{"input": "/*\n * Copyright 2021 NVIDIA Corporation\n *\n * Licensed under the Apache License, Version 2.0 with the LLVM exception\n * (the \"License\"); you may not use this file except in compliance with\n * the License.\n *\n * You may obtain a copy of the License at\n *\n * http://llvm.org/foundation/relicensing/LICENSE.txt\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \n\n#include \n#include \n\nnamespace\n{\n\n// Once launched, this kernel will block the stream until `flag` updates is\n// non-zero. If `timeout` seconds pass, an error will be printed and the stream\n// will unblock.\n__global__ void block_stream(const volatile nvbench::int32_t *flag,\n volatile nvbench::int32_t *timeout_flag,\n nvbench::float64_t timeout)\n{\n const auto start_point = cuda::std::chrono::high_resolution_clock::now();\n const auto timeout_ns =\n cuda::std::chrono::nanoseconds(static_cast(timeout * 1e9));\n const auto timeout_point = start_point + timeout_ns;\n\n const bool use_timeout = timeout >= 0.;\n auto now = cuda::std::chrono::high_resolution_clock::now();\n while (!(*flag) && (!use_timeout || now < timeout_point))\n {\n now = cuda::std::chrono::high_resolution_clock::now();\n }\n\n if (now >= timeout_point)\n {\n *timeout_flag = 1;\n __threadfence_system(); // Ensure timeout flag visibility on host.\n printf(\"\\n\"\n \"######################################################################\\n\"\n \"##################### Possible Deadlock Detected #####################\\n\"\n \"######################################################################\\n\"\n \"\\n\"\n \"Forcing unblock: The current measurement appears to have deadlocked\\n\"\n \"and the results cannot be trusted.\\n\"\n \"\\n\"\n \"This happens when the KernelLauncher synchronizes the CUDA device.\\n\"\n \"If this is the case, pass the `sync` exec_tag to the `exec` call:\\n\"\n \"\\n\"\n \" state.exec(); // Deadlock\\n\"\n \" state.exec(nvbench::exec_tag::sync, ); // Safe\\n\"\n \"\\n\"\n \"This tells NVBench about the sync so it can run the benchmark safely.\\n\"\n \"\\n\"\n \"If the KernelLauncher does not synchronize but has a very long \\n\"\n \"execution time, this may be a false positive. If so, disable this\\n\"\n \"check with:\\n\"\n \"\\n\"\n \" state.set_blocking_kernel_timeout(-1);\\n\"\n \"\\n\"\n \"The current timeout is set to %0.5g seconds.\\n\"\n \"\\n\"\n \"For more information, see the 'Benchmarks that sync' section of the\\n\"\n \"NVBench documentation.\\n\"\n \"\\n\"\n \"If this happens while profiling with an external tool,\\n\"\n \"pass the `--profile` flag to the executable to disable use of blocking kernel\\n\"\n \"and to also run the benchmark only once.\\n\"\n \"\\n\"\n \"For more information, see the 'Benchmark Properties' section of the\\n\"\n \"NVBench documentation.\\n\\n\",\n timeout);\n }\n}\n\n} // namespace\n\nnamespace nvbench\n{\n\nblocking_kernel::blocking_kernel()\n{\n NVBENCH_CUDA_CALL(cudaHostRegister(&m_host_flag, sizeof(m_host_flag), cudaHostRegisterMapped));\n NVBENCH_CUDA_CALL(cudaHostGetDevicePointer(&m_device_flag, &m_host_flag, 0));\n NVBENCH_CUDA_CALL(\n cudaHostRegister(&m_host_timeout_flag, sizeof(m_host_timeout_flag), cudaHostRegisterMapped));\n NVBENCH_CUDA_CALL(cudaHostGetDevicePointer(&m_device_timeout_flag, &m_host_timeout_flag, 0));\n}\n\nblocking_kernel::~blocking_kernel()\n{\n NVBENCH_CUDA_CALL_NOEXCEPT(cudaHostUnregister(&m_host_flag));\n NVBENCH_CUDA_CALL_NOEXCEPT(cudaHostUnregister(&m_host_timeout_flag));\n}\n\nvoid blocking_kernel::block(const nvbench::cuda_stream &stream, nvbench::float64_t timeout)\n{\n m_host_flag = 0;\n m_host_timeout_flag = 0;\n block_stream<<<1, 1, 0, stream>>>(m_device_flag, m_device_timeout_flag, timeout);\n}\n\nvoid blocking_kernel::timeout_detected()\n{\n NVBENCH_THROW(std::runtime_error,\n \"{}\",\n \"Deadlock detected -- missing nvbench::exec_tag::sync? \"\n \"See stdout for details.\");\n}\n\n} // namespace nvbench\n", "output": "/*\n * Copyright 2021 NVIDIA Corporation\n *\n * Licensed under the Apache License, Version 2.0 with the LLVM exception\n * (the \"License\"); you may not use this file except in compliance with\n * the License.\n *\n * You may obtain a copy of the License at\n *\n * http://llvm.org/foundation/relicensing/LICENSE.txt\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \n\n#include \n#include \n\nnamespace\n{\n\n// Once launched, this kernel will block the stream until `flag` updates is\n// non-zero. If `timeout` seconds pass, an error will be printed and the stream\n// will unblock.\n__global__ void block_stream(const volatile nvbench::int32_t *flag,\n volatile nvbench::int32_t *timeout_flag,\n nvbench::float64_t timeout)\n{\n const auto start_point = cuda::std::chrono::high_resolution_clock::now();\n const auto timeout_ns =\n cuda::std::chrono::nanoseconds(static_cast(timeout * 1e9));\n const auto timeout_point = start_point + timeout_ns;\n\n const bool use_timeout = timeout >= 0.;\n auto now = cuda::std::chrono::high_resolution_clock::now();\n while (!(*flag) && (!use_timeout || now < timeout_point))\n {\n now = cuda::std::chrono::high_resolution_clock::now();\n }\n\n if (now >= timeout_point)\n {\n *timeout_flag = 1;\n __threadfence_system(); // Ensure timeout flag visibility on host.\n printf(\"\\n\"\n \"######################################################################\\n\"\n \"##################### Possible Deadlock Detected #####################\\n\"\n \"######################################################################\\n\"\n \"\\n\"\n \"Forcing unblock: The current measurement appears to have deadlocked\\n\"\n \"and the results cannot be trusted.\\n\"\n \"\\n\"\n \"This happens when the KernelLauncher synchronizes the CUDA device.\\n\"\n \"If this is the case, pass the `sync` exec_tag to the `exec` call:\\n\"\n \"\\n\"\n \" state.exec(); // Deadlock\\n\"\n \" state.exec(nvbench::exec_tag::sync, ); // Safe\\n\"\n \"\\n\"\n \"This tells NVBench about the sync so it can run the benchmark safely.\\n\"\n \"\\n\"\n \"If the KernelLauncher does not synchronize but has a very long \\n\"\n \"execution time, this may be a false positive. If so, disable this\\n\"\n \"check with:\\n\"\n \"\\n\"\n \" state.set_blocking_kernel_timeout(-1);\\n\"\n \"\\n\"\n \"The current timeout is set to %0.5g seconds.\\n\"\n \"\\n\"\n \"For more information, see the 'Benchmarks that sync' section of the\\n\"\n \"NVBench documentation.\\n\"\n \"\\n\"\n \"If this happens while profiling with an external tool,\\n\"\n \"pass the `--profile` flag to the executable to disable use of blocking kernel\\n\"\n \"and to also run the benchmark only once.\\n\"\n \"\\n\"\n \"For more information, see the 'Benchmark Properties' section of the\\n\"\n \"NVBench documentation.\\n\\n\",\n timeout);\n }\n}\n\n} // namespace\n\nnamespace nvbench\n{\n\nblocking_kernel::blocking_kernel()\n{\n NVBENCH_CUDA_CALL(hipHostRegister(&m_host_flag, sizeof(m_host_flag), hipHostRegisterMapped));\n NVBENCH_CUDA_CALL(hipHostGetDevicePointer(&m_device_flag, &m_host_flag, 0));\n NVBENCH_CUDA_CALL(\n hipHostRegister(&m_host_timeout_flag, sizeof(m_host_timeout_flag), hipHostRegisterMapped));\n NVBENCH_CUDA_CALL(hipHostGetDevicePointer(&m_device_timeout_flag, &m_host_timeout_flag, 0));\n}\n\nblocking_kernel::~blocking_kernel()\n{\n NVBENCH_CUDA_CALL_NOEXCEPT(hipHostUnregister(&m_host_flag));\n NVBENCH_CUDA_CALL_NOEXCEPT(hipHostUnregister(&m_host_timeout_flag));\n}\n\nvoid blocking_kernel::block(const nvbench::cuda_stream &stream, nvbench::float64_t timeout)\n{\n m_host_flag = 0;\n m_host_timeout_flag = 0;\n block_stream<<<1, 1, 0, stream>>>(m_device_flag, m_device_timeout_flag, timeout);\n}\n\nvoid blocking_kernel::timeout_detected()\n{\n NVBENCH_THROW(std::runtime_error,\n \"{}\",\n \"Deadlock detected -- missing nvbench::exec_tag::sync? \"\n \"See stdout for details.\");\n}\n\n} // namespace nvbench\n", "source": "/kaggle/working/nvbench/nvbench/blocking_kernel.cu"} {"input": "/*\n * Copyright 2025 NVIDIA Corporation\n *\n * Licensed under the Apache License, Version 2.0 with the LLVM exception\n * (the \"License\"); you may not use this file except in compliance with\n * the License.\n *\n * You may obtain a copy of the License at\n *\n * http://llvm.org/foundation/relicensing/LICENSE.txt\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n#include \n#include \n#include \n#include \n\n#include \n\n#include \n#include \n\nnamespace\n{\n\n__global__ void get_timestamps_kernel(nvbench::uint64_t *global_timestamp,\n nvbench::uint64_t *sm0_timestamp)\n{\n nvbench::uint32_t smid;\n asm volatile(\"mov.u32 %0, %%smid;\" : \"=r\"(smid));\n if (smid == 0)\n {\n nvbench::uint64_t gts, lts;\n asm volatile(\"mov.u64 %0, %%globaltimer;\" : \"=l\"(gts));\n lts = clock64();\n\n *global_timestamp = gts;\n *sm0_timestamp = lts;\n }\n}\n\n} // namespace\n\nnamespace nvbench::detail\n{\n\ntimestamps_kernel::timestamps_kernel()\n{\n NVBENCH_CUDA_CALL(\n cudaHostRegister(&m_host_timestamps, sizeof(nvbench::uint64_t) * 2, cudaHostRegisterMapped));\n NVBENCH_CUDA_CALL(cudaHostGetDevicePointer(&m_device_timestamps, &m_host_timestamps, 0));\n}\n\ntimestamps_kernel::~timestamps_kernel()\n{\n NVBENCH_CUDA_CALL_NOEXCEPT(cudaHostUnregister(&m_host_timestamps));\n}\n\nvoid timestamps_kernel::record(const nvbench::cuda_stream &stream)\n{\n m_host_timestamps[0] = 0;\n m_host_timestamps[1] = 0;\n\n int device_id = 0;\n int num_sms = 0;\n\n NVBENCH_CUDA_CALL(cudaGetDevice(&device_id));\n NVBENCH_CUDA_CALL(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, device_id));\n\n get_timestamps_kernel<<(num_sms), 1, 0, stream.get_stream()>>>(\n m_device_timestamps,\n m_device_timestamps + 1);\n}\n\n} // namespace nvbench::detail\n", "output": "/*\n * Copyright 2025 NVIDIA Corporation\n *\n * Licensed under the Apache License, Version 2.0 with the LLVM exception\n * (the \"License\"); you may not use this file except in compliance with\n * the License.\n *\n * You may obtain a copy of the License at\n *\n * http://llvm.org/foundation/relicensing/LICENSE.txt\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n#include \n#include \n#include \n#include \n\n#include \n\n#include \n#include \n\nnamespace\n{\n\n__global__ void get_timestamps_kernel(nvbench::uint64_t *global_timestamp,\n nvbench::uint64_t *sm0_timestamp)\n{\n nvbench::uint32_t smid;\n asm volatile(\"mov.u32 %0, %%smid;\" : \"=r\"(smid));\n if (smid == 0)\n {\n nvbench::uint64_t gts, lts;\n asm volatile(\"mov.u64 %0, %%globaltimer;\" : \"=l\"(gts));\n lts = clock64();\n\n *global_timestamp = gts;\n *sm0_timestamp = lts;\n }\n}\n\n} // namespace\n\nnamespace nvbench::detail\n{\n\ntimestamps_kernel::timestamps_kernel()\n{\n NVBENCH_CUDA_CALL(\n hipHostRegister(&m_host_timestamps, sizeof(nvbench::uint64_t) * 2, hipHostRegisterMapped));\n NVBENCH_CUDA_CALL(hipHostGetDevicePointer(&m_device_timestamps, &m_host_timestamps, 0));\n}\n\ntimestamps_kernel::~timestamps_kernel()\n{\n NVBENCH_CUDA_CALL_NOEXCEPT(hipHostUnregister(&m_host_timestamps));\n}\n\nvoid timestamps_kernel::record(const nvbench::cuda_stream &stream)\n{\n m_host_timestamps[0] = 0;\n m_host_timestamps[1] = 0;\n\n int device_id = 0;\n int num_sms = 0;\n\n NVBENCH_CUDA_CALL(hipGetDevice(&device_id));\n NVBENCH_CUDA_CALL(hipDeviceGetAttribute(&num_sms, hipDeviceAttributeMultiprocessorCount, device_id));\n\n get_timestamps_kernel<<(num_sms), 1, 0, stream.get_stream()>>>(\n m_device_timestamps,\n m_device_timestamps + 1);\n}\n\n} // namespace nvbench::detail\n", "source": "/kaggle/working/nvbench/nvbench/detail/timestamps_kernel.cu"} {"input": "/*\n * Copyright 2021 NVIDIA Corporation\n *\n * Licensed under the Apache License, Version 2.0 with the LLVM exception\n * (the \"License\"); you may not use this file except in compliance with\n * the License.\n *\n * You may obtain a copy of the License at\n *\n * http://llvm.org/foundation/relicensing/LICENSE.txt\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n#include \n\n// Thrust vectors simplify memory management:\n#include \n\ntemplate \n__global__ void kernel(std::size_t stride,\n std::size_t elements,\n const nvbench::int32_t *__restrict__ in,\n nvbench::int32_t *__restrict__ out)\n{\n const std::size_t tid = threadIdx.x + blockIdx.x * blockDim.x;\n const std::size_t step = gridDim.x * blockDim.x;\n\n for (std::size_t i = stride * tid; i < stride * elements; i += stride * step)\n {\n for (int j = 0; j < ItemsPerThread; j++)\n {\n const auto read_id = (ItemsPerThread * i + j) % elements;\n const auto write_id = tid + j * elements;\n out[write_id] = in[read_id];\n }\n }\n}\n\n// `throughput_bench` copies a 128 MiB buffer of int32_t, and reports throughput\n// and cache hit rates.\n//\n// Calling state.collect_*() enables particular metric collection if nvbench\n// was build with CUPTI support (CMake option: -DNVBench_ENABLE_CUPTI=ON).\ntemplate \nvoid throughput_bench(nvbench::state &state, nvbench::type_list>)\n{\n // Allocate input data:\n const std::size_t stride = static_cast(state.get_int64(\"Stride\"));\n const std::size_t elements = 128 * 1024 * 1024 / sizeof(nvbench::int32_t);\n thrust::device_vector input(elements);\n thrust::device_vector output(elements * ItemsPerThread);\n\n // Provide throughput information:\n state.add_element_count(elements, \"Elements\");\n state.collect_dram_throughput();\n state.collect_l1_hit_rates();\n state.collect_l2_hit_rates();\n state.collect_loads_efficiency();\n state.collect_stores_efficiency();\n\n const auto threads_in_block = 256;\n const auto blocks_in_grid =\n static_cast((elements + threads_in_block - 1) / threads_in_block);\n\n state.exec([&](nvbench::launch &launch) {\n kernel<<>>(\n stride,\n elements,\n thrust::raw_pointer_cast(input.data()),\n thrust::raw_pointer_cast(output.data()));\n });\n}\n\nusing items_per_thread = nvbench::enum_type_list<1, 2>;\n\nNVBENCH_BENCH_TYPES(throughput_bench, NVBENCH_TYPE_AXES(items_per_thread))\n .set_type_axes_names({\"ItemsPerThread\"})\n .add_int64_axis(\"Stride\", nvbench::range(1, 4, 3));\n", "output": "/*\n * Copyright 2021 NVIDIA Corporation\n *\n * Licensed under the Apache License, Version 2.0 with the LLVM exception\n * (the \"License\"); you may not use this file except in compliance with\n * the License.\n *\n * You may obtain a copy of the License at\n *\n * http://llvm.org/foundation/relicensing/LICENSE.txt\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n\n// Thrust vectors simplify memory management:\n#include \n\ntemplate \n__global__ void kernel(std::size_t stride,\n std::size_t elements,\n const nvbench::int32_t *__restrict__ in,\n nvbench::int32_t *__restrict__ out)\n{\n const std::size_t tid = threadIdx.x + blockIdx.x * blockDim.x;\n const std::size_t step = gridDim.x * blockDim.x;\n\n for (std::size_t i = stride * tid; i < stride * elements; i += stride * step)\n {\n for (int j = 0; j < ItemsPerThread; j++)\n {\n const auto read_id = (ItemsPerThread * i + j) % elements;\n const auto write_id = tid + j * elements;\n out[write_id] = in[read_id];\n }\n }\n}\n\n// `throughput_bench` copies a 128 MiB buffer of int32_t, and reports throughput\n// and cache hit rates.\n//\n// Calling state.collect_*() enables particular metric collection if nvbench\n// was build with CUPTI support (CMake option: -DNVBench_ENABLE_CUPTI=ON).\ntemplate \nvoid throughput_bench(nvbench::state &state, nvbench::type_list>)\n{\n // Allocate input data:\n const std::size_t stride = static_cast(state.get_int64(\"Stride\"));\n const std::size_t elements = 128 * 1024 * 1024 / sizeof(nvbench::int32_t);\n thrust::device_vector input(elements);\n thrust::device_vector output(elements * ItemsPerThread);\n\n // Provide throughput information:\n state.add_element_count(elements, \"Elements\");\n state.collect_dram_throughput();\n state.collect_l1_hit_rates();\n state.collect_l2_hit_rates();\n state.collect_loads_efficiency();\n state.collect_stores_efficiency();\n\n const auto threads_in_block = 256;\n const auto blocks_in_grid =\n static_cast((elements + threads_in_block - 1) / threads_in_block);\n\n state.exec([&](nvbench::launch &launch) {\n kernel<<>>(\n stride,\n elements,\n thrust::raw_pointer_cast(input.data()),\n thrust::raw_pointer_cast(output.data()));\n });\n}\n\nusing items_per_thread = nvbench::enum_type_list<1, 2>;\n\nNVBENCH_BENCH_TYPES(throughput_bench, NVBENCH_TYPE_AXES(items_per_thread))\n .set_type_axes_names({\"ItemsPerThread\"})\n .add_int64_axis(\"Stride\", nvbench::range(1, 4, 3));\n", "source": "/kaggle/working/nvbench/examples/auto_throughput.cu"} {"input": "#include \n\n#include \"test_asserts.cuh\"\n\nnamespace\n{\n__global__ void multiply5(const int32_t *__restrict__ a, int32_t *__restrict__ b)\n{\n const auto id = blockIdx.x * blockDim.x + threadIdx.x;\n b[id] = 5 * a[id];\n}\n} // namespace\n\nint main()\n{\n multiply5<<<256, 256>>>(nullptr, nullptr);\n\n try\n {\n NVBENCH_CUDA_CALL(cudaStreamSynchronize(0));\n ASSERT(false);\n }\n catch (const std::runtime_error &)\n {\n ASSERT(cudaGetLastError() == cudaError_t::cudaSuccess);\n }\n\n return 0;\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n\n#include \"test_asserts.cuh\"\n\nnamespace\n{\n__global__ void multiply5(const int32_t *__restrict__ a, int32_t *__restrict__ b)\n{\n const auto id = blockIdx.x * blockDim.x + threadIdx.x;\n b[id] = 5 * a[id];\n}\n} // namespace\n\nint main()\n{\n multiply5<<<256, 256>>>(nullptr, nullptr);\n\n try\n {\n NVBENCH_CUDA_CALL(hipStreamSynchronize(0));\n ASSERT(false);\n }\n catch (const std::runtime_error &)\n {\n ASSERT(hipGetLastError() == hipError_t::hipSuccess);\n }\n\n return 0;\n}\n", "source": "/kaggle/working/nvbench/testing/reset_error.cu"} {"input": "/*\n * Copyright 2024 NVIDIA Corporation\n *\n * Licensed under the Apache License, Version 2.0 with the LLVM exception\n * (the \"License\"); you may not use this file except in compliance with\n * the License.\n *\n * You may obtain a copy of the License at\n *\n * http://llvm.org/foundation/relicensing/LICENSE.txt\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n#include \n#include \n\n/******************************************************************************\n * Install custom parser.\n * sSee for more details.\n ******************************************************************************/\n\n//\n// Step 1: Define a custom argument handler that accepts a vector of strings.\n// - This handler should modify the vector in place to remove any custom\n// arguments it handles. NVbench will then parse the remaining arguments.\n// - The handler should also update any application state needed to handle\n// the custom arguments.\n//\n\n// User code to handle a specific argument:\nvoid handle_my_custom_arg();\n\n// NVBench hook for modifying the command line arguments before parsing:\nvoid custom_arg_handler(std::vector &args)\n{\n // Handle and remove \"--my-custom-arg\"\n if (auto it = std::find(args.begin(), args.end(), \"--my-custom-arg\"); it != args.end())\n {\n handle_my_custom_arg();\n args.erase(it);\n }\n}\n\n//\n// Step 2: Install the custom argument handler.\n// - This is done by defining a macro that invokes the custom argument handler.\n//\n\n// Install the custom argument handler:\n// Either define this before any NVBench headers are included, or undefine and redefine:\n#undef NVBENCH_MAIN_CUSTOM_ARGS_HANDLER\n#define NVBENCH_MAIN_CUSTOM_ARGS_HANDLER(args) custom_arg_handler(args)\n\n// Step 3: Define `main`\n//\n// After installing the custom argument handler, define the main function using:\n//\n// ```\n// NVBENCH_MAIN\n// ```\n//\n// Here, this is done at the end of this file.\n\n/******************************************************************************\n * Unit test verification:\n ******************************************************************************/\n\n// Track whether the args are found / handled.\nbool h_custom_arg_found = false;\nbool h_handled_on_device = false;\n__device__ bool d_custom_arg_found = false;\n__device__ bool d_handled_on_device = false;\n\n// Copy host values to device:\nvoid copy_host_state_to_device()\n{\n NVBENCH_CUDA_CALL(cudaMemcpyToSymbol(d_custom_arg_found, &h_custom_arg_found, sizeof(bool)));\n NVBENCH_CUDA_CALL(cudaMemcpyToSymbol(d_handled_on_device, &h_handled_on_device, sizeof(bool)));\n}\n\n// Copy device values to host:\nvoid copy_device_state_to_host()\n{\n NVBENCH_CUDA_CALL(cudaMemcpyFromSymbol(&h_custom_arg_found, d_custom_arg_found, sizeof(bool)));\n NVBENCH_CUDA_CALL(cudaMemcpyFromSymbol(&h_handled_on_device, d_handled_on_device, sizeof(bool)));\n}\n\nvoid handle_my_custom_arg()\n{\n h_custom_arg_found = true;\n copy_host_state_to_device();\n}\n\nvoid verify()\n{\n copy_device_state_to_host();\n if (!h_custom_arg_found)\n {\n throw std::runtime_error(\"Custom argument not detected.\");\n }\n if (!h_handled_on_device)\n {\n throw std::runtime_error(\"Custom argument not handled on device.\");\n }\n}\n\n// Install a verification check to ensure the custom argument was handled.\n// Use the `PRE` finalize hook to ensure we check device state before resetting the context.\n#undef NVBENCH_MAIN_FINALIZE_CUSTOM_PRE\n#define NVBENCH_MAIN_FINALIZE_CUSTOM_PRE() verify()\n\n// Simple kernel/benchmark to make sure that the handler can successfully modify CUDA state:\n__global__ void kernel()\n{\n if (d_custom_arg_found)\n {\n d_handled_on_device = true;\n }\n}\nvoid bench(nvbench::state &state)\n{\n state.exec([](nvbench::launch &) { kernel<<<1, 1>>>(); });\n}\nNVBENCH_BENCH(bench);\n\n// Define the customized main function:\nNVBENCH_MAIN\n", "output": "/*\n * Copyright 2024 NVIDIA Corporation\n *\n * Licensed under the Apache License, Version 2.0 with the LLVM exception\n * (the \"License\"); you may not use this file except in compliance with\n * the License.\n *\n * You may obtain a copy of the License at\n *\n * http://llvm.org/foundation/relicensing/LICENSE.txt\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n/******************************************************************************\n * Install custom parser.\n * sSee for more details.\n ******************************************************************************/\n\n//\n// Step 1: Define a custom argument handler that accepts a vector of strings.\n// - This handler should modify the vector in place to remove any custom\n// arguments it handles. NVbench will then parse the remaining arguments.\n// - The handler should also update any application state needed to handle\n// the custom arguments.\n//\n\n// User code to handle a specific argument:\nvoid handle_my_custom_arg();\n\n// NVBench hook for modifying the command line arguments before parsing:\nvoid custom_arg_handler(std::vector &args)\n{\n // Handle and remove \"--my-custom-arg\"\n if (auto it = std::find(args.begin(), args.end(), \"--my-custom-arg\"); it != args.end())\n {\n handle_my_custom_arg();\n args.erase(it);\n }\n}\n\n//\n// Step 2: Install the custom argument handler.\n// - This is done by defining a macro that invokes the custom argument handler.\n//\n\n// Install the custom argument handler:\n// Either define this before any NVBench headers are included, or undefine and redefine:\n#undef NVBENCH_MAIN_CUSTOM_ARGS_HANDLER\n#define NVBENCH_MAIN_CUSTOM_ARGS_HANDLER(args) custom_arg_handler(args)\n\n// Step 3: Define `main`\n//\n// After installing the custom argument handler, define the main function using:\n//\n// ```\n// NVBENCH_MAIN\n// ```\n//\n// Here, this is done at the end of this file.\n\n/******************************************************************************\n * Unit test verification:\n ******************************************************************************/\n\n// Track whether the args are found / handled.\nbool h_custom_arg_found = false;\nbool h_handled_on_device = false;\n__device__ bool d_custom_arg_found = false;\n__device__ bool d_handled_on_device = false;\n\n// Copy host values to device:\nvoid copy_host_state_to_device()\n{\n NVBENCH_CUDA_CALL(hipMemcpyToSymbol(HIP_SYMBOL(d_custom_arg_found), &h_custom_arg_found, sizeof(bool)));\n NVBENCH_CUDA_CALL(hipMemcpyToSymbol(HIP_SYMBOL(d_handled_on_device), &h_handled_on_device, sizeof(bool)));\n}\n\n// Copy device values to host:\nvoid copy_device_state_to_host()\n{\n NVBENCH_CUDA_CALL(hipMemcpyFromSymbol(&h_custom_arg_found,HIP_SYMBOL(d_custom_arg_found), sizeof(bool)));\n NVBENCH_CUDA_CALL(hipMemcpyFromSymbol(&h_handled_on_device,HIP_SYMBOL(d_handled_on_device), sizeof(bool)));\n}\n\nvoid handle_my_custom_arg()\n{\n h_custom_arg_found = true;\n copy_host_state_to_device();\n}\n\nvoid verify()\n{\n copy_device_state_to_host();\n if (!h_custom_arg_found)\n {\n throw std::runtime_error(\"Custom argument not detected.\");\n }\n if (!h_handled_on_device)\n {\n throw std::runtime_error(\"Custom argument not handled on device.\");\n }\n}\n\n// Install a verification check to ensure the custom argument was handled.\n// Use the `PRE` finalize hook to ensure we check device state before resetting the context.\n#undef NVBENCH_MAIN_FINALIZE_CUSTOM_PRE\n#define NVBENCH_MAIN_FINALIZE_CUSTOM_PRE() verify()\n\n// Simple kernel/benchmark to make sure that the handler can successfully modify CUDA state:\n__global__ void kernel()\n{\n if (d_custom_arg_found)\n {\n d_handled_on_device = true;\n }\n}\nvoid bench(nvbench::state &state)\n{\n state.exec([](nvbench::launch &) { kernel<<<1, 1>>>(); });\n}\nNVBENCH_BENCH(bench);\n\n// Define the customized main function:\nNVBENCH_MAIN\n", "source": "/kaggle/working/nvbench/testing/custom_main_custom_args.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \"utils.cuh\"\n#include \n#include \n#include \n\n#include \n#include // For mkdir\n#include // For mode_t\n\nnamespace cg = cooperative_groups;\n\n// ------------------------------------------------------------------------------------------------\n// Linear layer\ntypedef struct {\n float* w; // OC, C\n float* b; // OC\n} LinearParams;\n\ntypedef struct {\n float* inp; // N, C\n float* out; // N, OC\n} LinearActs;\n\n\nvoid linear_set_param_ptrs(\n LinearParams* params,\n float* params_memory,\n int C, int OC\n) {\n params->w = params_memory;\n params_memory += OC * C;\n params->b = params_memory;\n}\n\ninline size_t linear_count_params(int C, int OC) {\n return OC * C + OC;\n}\n// ---------------------------------------------\n// kernels\n\n// copy the matmul forward and backward kernels from llm.c\n\n// is there no better way other than just adding bias with a whole separate kernel?\n// this is a highly memory-bound operation, should be fused into the matmul kernel\n// but i can't seem to find a cuBLAS function that does this\n__global__ void add_bias(float* out, const float* bias, int N, int OC) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = idx; i < N * OC; i += stride) {\n int col = i % OC;\n out[i] += bias[col];\n }\n}\n\n// kernel 2 calls cuBLAS, which should be very efficient\nvoid matmul_forward2(cublasHandle_t cublas_handle,\n float* out,\n const float* inp, const float* weight, const float* bias,\n int N, int C, int OC,\n const int block_size) {\n // for reference API is:\n // cublasStatus_t cublasSgemm(cublasHandle_t handle,\n // cublasOperation_t transa, cublasOperation_t transb,\n // int m, int n, int k,\n // const float *alpha,\n // const float *A, int lda,\n // const float *B, int ldb,\n // const float *beta,\n // float *C, int ldc)\n // for us, inp is (B*T, C), weight is (OC, C), out is (B*T, OC)\n // cuBLAS does C = alpha * A * B + beta * C\n // where A is mxk, B is kxn, C is mxn\n // now, because we use row-major storage, cuBLAS (which is column-major) sees our matrices transposed.\n // algorithmically / in e.g. PyTorch we want to do: out = inp @ weight.T\n // but because cuBLAS is column-major, we actually want to get it to calculate out.T . Mathematically, this is:\n // out.T = weight @ inp.T\n // but again, our variables look transposed, so using the actual weight/inp we have here in this function, this becomes\n // out.T = weight.T @ inp\n // so we need to get cuBLAS to calculate weight.T @ inp (the variables here are the actual ones in this function)\n // => need to call cuBLAS with A = weight, B = inp\n // => need to call cuBLAS with transa = CUBLAS_OP_T, transb = CUBLAS_OP_N\n\n const float alpha = 1.0f;\n const float beta = 0.0f;\n cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, OC, N, C, &alpha, weight, C, inp, C, &beta, out, OC));\n // and now we still have to add the bias... (ew)\n if (bias != NULL) {\n int grid_size = ceil_div(OC * N, block_size);\n add_bias<<>>(out, bias, N, OC);\n cudaCheck(cudaGetLastError());\n }\n}\n\n// use shared memory and coarsening + reductions\n__global__ void matmul_backward_bias_kernel_faster(float* dbias, const float* dout, int N, int OC) {\n extern __shared__ float shared[];\n int o = blockIdx.x; // range [0, OC)\n int tid = threadIdx.x; // range [0, block_size)\n int block_size = blockDim.x;\n const float* x = dout + o;\n // thread coarsening\n double sum = 0.0;\n for (int i = tid; i < N; i += block_size) {\n sum += x[i * OC];\n }\n shared[tid] = (float) sum;\n __syncthreads();\n // reductions\n for (int stride = block_size / 2; stride >= 1; stride /= 2) {\n __syncthreads();\n if (tid < stride) {\n shared[tid] += shared[tid + stride];\n }\n }\n // write the final result (at thread 0) to global memory\n if (tid == 0) {\n dbias[o] = shared[0];\n }\n}\n\n// version1: simple cuBLAS calls\nvoid matmul_backward1(cublasHandle_t cublas_handle,\n float* dinp, float* dweight, float* dbias,\n float* dout, float* inp, float* weight,\n int N, int C, int OC) {\n float alpha = 1.0f;\n float beta = 0.0f; // note we must use beta = 1.0 so that we do a +=, as we should, because gradients add\n\n // for reference the API is:\n // cublasStatus_t cublasSgemm(cublasHandle_t handle,\n // cublasOperation_t transa, cublasOperation_t transb,\n // int m, int n, int k,\n // const float *alpha,\n // const float *A, int lda,\n // const float *B, int ldb,\n // const float *beta,\n // float *C, int ldc)\n\n // recall the forward pass was calculated with alpha = 1.0f, beta = 0.0f as:\n // cublasSgemm(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, OC, B*T, C, &alpha, weight, C, inp, C, &beta, out, OC);\n\n // backward to input\n cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, C, N, OC, &alpha, weight, C, dout, OC, &beta, dinp, C));\n // backward to weight\n cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, C, OC, N, &alpha, inp, C, dout, OC, &beta, dweight, C));\n // backward to bias, if given\n if (dbias != NULL) {\n\n // sum over B,T using matrix vector multiplication with cuBLAS\n // for reference this API is:\n // cublasStatus_t cublasSgemv(cublasHandle_t handle, cublasOperation_t trans,\n // int m, int n,\n // const float *alpha,\n // const float *A, int lda,\n // const float *x, int incx,\n // const float *beta,\n // float *y, int incy)\n // dout is (B,T,OC), or in 2D terms (B*T, OC)\n // cublasCheck(cublasSgemv(cublas_handle, CUBLAS_OP_N, B*T, OC, &alpha, dout, B*T, ones, 1, &beta, dbias, 1));\n // cublasCheck(cublasSgemv(cublas_handle, CUBLAS_OP_T, OC, B*T, &alpha, dout, OC, ones, 1, &beta, dbias, 1));\n\n // ugh the above isn't working...\n // let's just do naive calculation for now, fix later\n // const int block_size=128;\n // const int grid_size=(OC + block_size - 1) / block_size;\n // matmul_backward_bias_kernel<<>>(dbias, dout, B, T, OC);\n\n // bit faster\n const int block_size=512;\n dim3 block_dim(block_size);\n dim3 grid_dim(OC);\n size_t shared_mem_size = block_size * sizeof(float);\n matmul_backward_bias_kernel_faster<<>>(dbias, dout, N, OC);\n }\n}\n\n// ------------------------------------------------------------------------------------------------\n// Broadcast tensor layer\n\n// kernel to broadcast a tensor from shape (B, C) to (B, C, H, W)\n// Assuming one thread per pixel\n__global__ void broadcast_last_dims_forward_kernel(\n const float* x,\n float* out,\n int N, int H, int W, int n_bk_per_img\n) {\n int n = blockIdx.x / n_bk_per_img;\n __shared__ float x_sh[1];\n if (threadIdx.x == 0) {\n x_sh[0] = x[n];\n }\n __syncthreads();\n\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < N * H * W) {\n out[idx] = x_sh[0];\n }\n}\n\nvoid broadcast_last_dims_forward(\n const float* x,\n float* out,\n int N, int H, int W, int block_size\n) {\n block_size = min_int(block_size, H * W);\n int n_bk_per_img = ceil_div(H * W, block_size);\n broadcast_last_dims_forward_kernel<<>>(x, out, N, H, W, n_bk_per_img);\n};\n\n// backwards for broadcast\n// use one warp per image\n__global__ void broadcast_last_dims_backward_kernel(\n const float* dout,\n float* dx,\n int N, int H, int W\n) {\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();\n if (idx >= N) { return; }\n\n int img_size = H * W;\n dout += idx * img_size;\n dx += idx;\n\n float sum = 0.0f;\n for (int i = warp.thread_rank(); i < img_size; i += warp.size()) {\n sum += dout[i];\n }\n sum = cg::reduce(warp, sum, cg::plus());\n if (warp.thread_rank() == 0) {\n dx[0] = sum;\n }\n}\n\nvoid broadcast_last_dims_backward(\n const float* dout,\n float* dx,\n int N, int H, int W, int block_size\n) {\n block_size = min_int(block_size, H * W);\n assert(block_size % 32 == 0);\n int warps_per_block = block_size / 32;\n int gridDim = ceil_div(N, warps_per_block);\n broadcast_last_dims_backward_kernel<<>>(dout, dx, N, H, W);\n};\n\n// ------------------------------------------------------------------------------------------------\n// Add layer\n__global__ void add_forward_kernel(\n const float* a, const float* b,\n float* out,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < N) {\n out[idx] = a[idx] + b[idx];\n }\n}\n\nvoid add_forward(\n const float* a, const float* b,\n float* out,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n add_forward_kernel<<>>(a, b, out, N);\n}\n\n// add A and B in place, and store result in B\n__global__ void add_inplace_forward_kernel(\n const float* a, float* b,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < N) {\n b[idx] += a[idx];\n }\n}\n\nvoid add_inplace_forward(\n const float* a, float* b,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n add_inplace_forward_kernel<<>>(a, b, N);\n}\n\n// ------------------------------------------------------------------------------------------------\n// SiLU layer\n\ntypedef struct {\n float* x;\n float* out;\n} SiluActs;\n\n// ----------------------------------------------\n// kernels\n\n__global__ void silu_forward_kernel(\n const float* x,\n float* out,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx >= N) {\n return;\n }\n\n float x_val = x[idx];\n out[idx] = x_val / (1.0f + expf(-x_val));\n}\n\nvoid silu_forward(\n const float* x,\n float* out,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n silu_forward_kernel<<>>(x, out, N);\n}\n\n__global__ void silu_backward_kernel(\n const float* dout, const float* x,\n float* dx,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx >= N) {\n return;\n }\n float out_val = dout[idx];\n float x_val = x[idx];\n float expx = expf(-x_val);\n float grad_silu = (1.0f + x_val * expx / (1.0f + expx)) / (1.0f + expx);\n dx[idx] = out_val * grad_silu;\n}\n\nvoid silu_backward(\n const float* dout, const float* x,\n float* dx,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n silu_backward_kernel<<>>(dout, x, dx, N);\n}\n\n// ------------------------------------------------------------------------------------------------\n// Upsample\ntypedef struct {\n float* x;\n float* out;\n} UpsampleActs;\n\n// one thread per pixel\n__global__ void upsample_forward_kernel(\n const float* x,\n float* out,\n int B, int C, int H, int W\n) {\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int img_size = H * W;\n if (flat_idx >= B * C * img_size) {\n return;\n }\n\n int H_out = H*2;\n int W_out = W*2;\n int img_out_size = H_out * W_out;\n\n int b = flat_idx / (C * img_size);\n int c = (flat_idx / img_size) % C;\n int i = (flat_idx / W) % H;\n int j = flat_idx % W;\n\n // move pointers\n x += b * C * img_size + c * img_size + i * W + j;\n out += b * C * img_out_size + c * img_out_size + 2 * i * W_out + 2 * j;\n\n float x_val = x[0];\n\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n out[ii * W_out + jj] = x_val;\n }\n }\n}\n\nvoid upsample_forward1(\n float* out,\n const float* x,\n int B, int C, int H, int W,\n int block_size\n) {\n int gridDim = ceil_div(B * C * H * W, block_size);\n upsample_forward_kernel<<>>(x, out, B, C, H, W);\n cudaCheck(cudaGetLastError());\n}\n\n// number of threads equals the number of pixels\n// in downsampled image\n__global__ void upsample_backward_kernel(\n const float* dout,\n float* dx,\n int B, int C, int H, int W\n) {\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int img_in_size = H * W;\n if (flat_idx >= B * C * img_in_size) {\n return;\n }\n\n int H_out = H*2;\n int W_out = W*2;\n int img_out_size = H_out * W_out;\n\n int b = flat_idx / (C * img_in_size);\n int c = (flat_idx / img_in_size) % C;\n int i = (flat_idx / W) % H;\n int j = flat_idx % W;\n\n // move pointers\n dx += b * C * img_in_size + c * img_in_size + i * W + j;\n dout += b * C * img_out_size + c * img_out_size + 2 * i * W_out + 2 * j;\n\n float dout_sum = 0.0f;\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n dout_sum += dout[ii * W_out + jj];\n }\n }\n dx[0] = dout_sum;\n}\n\nvoid upsample_backward1(\n float* dx,\n const float* dout,\n int B, int C, int H, int W,\n int block_size\n) {\n int gridDim = ceil_div(B * C * H * W, block_size);\n upsample_backward_kernel<<>>(dout, dx, B, C, H, W);\n cudaCheck(cudaGetLastError());\n}\n\n// ------------------------------------------------------------------------------------------------\n// Avgpool layer\n\ntypedef struct {\n float* x;\n float* out;\n} AvgpoolActs;\n\n__global__ void avgpool_2d_forward_kernel(\n const float* x, // (B, C, H, W)\n float* out, // (B, C, H/2, W/2)\n int B, int C, int H, int W\n) {\n //printf(\"started kernel\\n\");\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int H_2 = H / 2;\n int W_2 = W / 2;\n int img_out_size = H_2 * W_2;\n if (flat_idx >= B * C * img_out_size) {\n return;\n }\n int img_size = H * W;\n int b = flat_idx / (C * img_out_size);\n int c = (flat_idx / img_out_size) % C;\n int i = (flat_idx / W_2) % H_2;\n int j = flat_idx % W_2;\n\n // move pointers\n x += b * C * img_size + c * img_size + 2 * i * W + 2 * j;\n out += b * C * img_out_size + c * img_out_size + i * W_2 + j;\n\n float sum = 0.0f;\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n sum += x[ii * W + jj];\n }\n }\n out[0] = sum / 4.0f;\n}\n\nvoid avgpool_2d_forward1(\n float* out,\n const float* x,\n int B, int C, int H, int W,\n const int block_size\n) {\n int gridDim = ceil_div(B * C * (H / 2) * (W / 2), block_size);\n avgpool_2d_forward_kernel<<>>(x, out, B, C, H, W);\n cudaCheck(cudaGetLastError());\n}\n\n__global__ void avgpool_2d_backward_kernel(\n const float* dout,\n float* dx,\n int B, int C, int H_out, int W_out // H_out and W_out are the shape of the output from forward, i.e. half the original sizes\n) {\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int img_out_size = H_out * W_out;\n if (flat_idx >= B * C * img_out_size) {\n return;\n }\n int H_in = H_out * 2;\n int W_in = W_out * 2;\n int img_in_size = H_in * W_in;\n int b = flat_idx / (C * img_out_size);\n int c = (flat_idx / img_out_size) % C;\n int i = (flat_idx / W_out) % H_out;\n int j = flat_idx % W_out;\n\n // move pointers\n dout += b * C * img_out_size + c * img_out_size + i * W_out + j;\n dx += b * C * img_in_size + c * img_in_size + 2 * i * W_in + 2 * j;\n\n float d = dout[0] / 4.0f;\n\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n dx[ii * W_in + jj] = d;\n }\n }\n}\n\nvoid avgpool_2d_backward1(\n const float* dout,\n float* dx,\n int B, int C, int H, int W,\n const int block_size\n) {\n int H_out = H / 2;\n int W_out = W / 2;\n int gridDim = ceil_div(B * C * H_out * W_out, block_size);\n avgpool_2d_backward_kernel<<>>(dout, dx, B, C, H_out, W_out);\n cudaCheck(cudaGetLastError());\n}\n\n// ------------------------------------------------------------------------------------------------\n// Concat channel layer\n// used to concat unet skip connections in the channel dimension\ntypedef struct {\n float* x1;\n float* x2;\n float* out;\n} ConcatChannelActs;\n\n__global__ void concat_channel_forward_kernel(\n const float* x1, const float* x2,\n float* out,\n int B, int C1, int C2, int H, int W\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < B * C1 * H * W) {\n // copy input from x1\n int b = idx / (C1 * H * W);\n int c = (idx / H / W) % C1;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + c * H * W + h * W + w;\n out[out_idx] = x1[idx];\n }\n if (idx < B * C2 * H * W) {\n // copy input from x2\n // move over from x1\n int b = idx / (C2 * H * W);\n int c = (idx / H / W) % C2;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + (C1 + c) * H * W + h * W + w;\n \n out[out_idx] = x2[idx];\n }\n}\n\nvoid concat_channel_forward(\n const float* x1, const float* x2,\n float* out,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n int N = B * max_int(C1, C2) * H * W;\n int n_blk = ceil_div(N, block_size);\n concat_channel_forward_kernel<<>>(x1, x2, out, B, C1, C2, H, W);\n}\n\n__global__ void concat_channel_backward_kernel(\n const float* dout,\n float* dx1, float* dx2,\n int B, int C1, int C2, int H, int W\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < B * C1 * H * W) {\n int b = idx / (C1 * H * W);\n int c = (idx / H / W) % C1;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + c * H * W + h * W + w;\n\n dx1[idx] = dout[out_idx];\n }\n if (idx < B * C2 * H * W) {\n int b = idx / (C2 * H * W);\n int c = (idx / H / W) % C2;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + (C1 + c) * H * W + h * W + w;\n \n dx2[idx] = dout[out_idx];\n }\n}\n\nvoid concat_channel_backward(\n const float* dout,\n float* dx1, float* dx2,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n int N = B * max_int(C1, C2) * H * W;\n int n_blk = ceil_div(N, block_size);\n concat_channel_backward_kernel<<>>(dout, dx1, dx2, B, C1, C2, H, W);\n}\n\n// ------------------------------------------------------------------------------------------------\n// Conv2d 1x1 layer\n\ntemplate \n__global__ void conv2d_k1_forward_kernel(\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C, int H, int W, int O\n) {\n // constants\n const int BC = 16;\n assert(O % BO == 0);\n assert(C % BC == 0);\n assert(H * W % BH == 0);\n // n_threads per block = (BO * BH) / (TO * TH) = 256\n // n_blocks = B * (O / BO) * (H * W / BH)\n // set gridDim to (BxHxW / BH, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int block_col = blockIdx.x % blocks_per_col;\n const int batch = blockIdx.x / blocks_per_col;\n\n const int threads_per_col = BH / TH;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // shared mem for weight and x\n const int weight_s_sz = BO * BC;\n const int x_s_sz = BC * BH;\n __shared__ float w_s[weight_s_sz];\n __shared__ float x_s[x_s_sz];\n\n // local registers for results;\n float thread_results[TO * TH] = {0.0};\n float reg_w[TO] = {0.0};\n float reg_x[TH] = {0.0};\n\n // get offsets to x and weight\n int x_block_offset = batch * C * H * W + block_col * BH;\n int w_block_offset = block_row * BO * C;\n\n for (int c_idx = 0; c_idx < C; c_idx += BC) {\n int x_inner_offset = c_idx * H * W;\n int w_inner_offset = c_idx;\n // load weights to smem\n for (int i = threadIdx.x; i < weight_s_sz; i += blockDim.x) {\n int w_i = i / BC;\n int w_j = i % BC;\n int w_abs_idx = w_block_offset + w_inner_offset + w_i * C + w_j;\n w_s[w_i * BC + w_j] = weight[w_abs_idx];\n }\n for (int i = threadIdx.x; i < x_s_sz; i += blockDim.x) {\n int x_i = i / BH;\n int x_j = i % BH;\n int x_abs_idx = x_block_offset + x_inner_offset + x_i * H * W + x_j;\n x_s[x_i * BH + x_j] = x[x_abs_idx];\n }\n __syncthreads();\n\n\n // calculate thread results\n for (int dot_idx = 0; dot_idx < BC; dot_idx++) {\n for (int i = 0; i < TO; i++) {\n reg_w[i] = w_s[(thread_row * TO + i) * BC + dot_idx];\n }\n for (int j = 0; j < TH; j++) {\n reg_x[j] = x_s[dot_idx * BH + thread_col * TH + j];\n }\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n thread_results[i * TH + j] += reg_w[i] * reg_x[j];\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n // set out and bias directly to thread offset\n int out_block_offset = batch * O * H * W + block_row * BO * H * W + block_col * BH;\n int bias_offset = block_row * BO;\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n int bias_idx = bias_offset + thread_row * TO + i;\n float out_val = thread_results[i * TH + j] + bias[bias_idx];\n int out_idx = out_block_offset + thread_row * TO * H * W + thread_col * TH + i * H * W + j;\n atomicAdd(out + out_idx, out_val);\n }\n }\n}\n\nvoid conv2d_k1_forward2(\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C_in, int H, int W, int C_out\n) {\n cudaCheck(cudaMemset(out, 0, B * C_out * H * W * sizeof(float)));\n const int BO = 64;\n const int BH = 64;\n const int TO = 4;\n const int TH = 4;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k1_forward_kernel<<>>(x, weight, bias, out, B, C_in, H, W, C_out);\n}\n\n__global__ void permute_dout_x_kernel(\n const float* x, const float* dout,\n float* x_perm, float* dout_perm,\n int B, int C_in, int C_out, int H, int W\n) {\n int flat_idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n int h = (flat_idx / W) % H;\n int w = flat_idx % W;\n\n if (flat_idx < B * C_in * H * W) {\n // permute x\n int b = flat_idx / (C_in * img_size);\n int c = (flat_idx / img_size) % C_in;\n\n int out_ind = c * B * img_size + b * img_size + h * W + w;\n x_perm[out_ind] = x[flat_idx];\n }\n\n if (flat_idx < B * C_out * H * W) {\n // permute dout\n int b = flat_idx / (C_out * img_size);\n int c = (flat_idx / img_size) % C_out;\n\n int out_ind = c * B * img_size + b * img_size + h * W + w;\n dout_perm[out_ind] = dout[flat_idx];\n }\n}\n\n\n__global__ void permute_dx_kernel(\n const float* dx_perm,\n float* dx,\n int B, int C_in, int H, int W\n) {\n int flat_idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n if (flat_idx < B * C_in * img_size) {\n int c = flat_idx / (B * img_size);\n int b = (flat_idx / img_size) % B;\n int h = (flat_idx / W) % H;\n int w = flat_idx % W;\n int out_ind = b * C_in * img_size + c * img_size + h * W + w;\n dx[out_ind] = dx_perm[flat_idx];\n }\n}\n\n\n// we want to sum the B, H, W dimensions of dout\n// since dout_perm is shaped as (C_out, B, H, W)\n// we will have each block take care of one entry of C_out\n// and the threads in each block will use coorperative_groups to do the sum\n__global__ void dbias_kernel(\n const float* dout_perm, // shape (C_out, B, H, W)\n float* dbias,\n int B, int C_out, int H, int W\n) {\n int bhw = B * H * W;\n extern __shared__ float shared[];\n\n // move dout_perm to the C_out index\n if (threadIdx.x < bhw) {\n dout_perm += blockIdx.x * B * H * W;\n float th_sum = 0.0f;\n for (int i = threadIdx.x; i < bhw; i += blockDim.x) {\n th_sum += dout_perm[i];\n }\n shared[threadIdx.x] = th_sum;\n }\n __syncthreads();\n\n // reduce the shared memory\n for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {\n if (threadIdx.x < stride) {\n shared[threadIdx.x] += shared[threadIdx.x + stride];\n }\n __syncthreads();\n }\n\n if (threadIdx.x == 0) {\n dbias[blockIdx.x] = shared[0];\n }\n}\n\n// uses cublas\nvoid conv2d_k1_backward1(\n cublasHandle_t cublas_handle,\n const float* dout, const float* x, const float* weight,\n float* dx, float* dweight, float* dbias,\n int B, int C_in, int C_out, int H, int W,\n int block_size\n) {\n // shapes:\n // x, dx: (B, C_in, H, W)\n // weight: (C_out, C_in)\n // bias: (C_out,)\n // dout: (B, C_out, H, W)\n\n // plan (the shapes are raw array shapes, note that cublas interprets the arrays as transposed):\n // 1. transpose x to x_perm of shape (C_in, B*H*W), and dout to dout_perm of shape (C_out, B*H*W)\n // 2. do one sgemm call to get dw = x.T @ dout_perm, and another call to get dx_perm = dout @ w.T\n // 3. dx_perm is of shape (C_in, B*H*W), so call a kernel to permute it\n // 4. reduce to get dbias\n float *x_perm, *dout_perm, *dx_perm;\n cudaCheck(cudaMalloc(&x_perm, C_in * B * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&dout_perm, C_out * B * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&dx_perm, C_in * B * H * W * sizeof(float)));\n\n // 1. permute x and dout\n int total_th1 = (C_in > C_out) ? C_in * B * H * W : C_out * B * H * W;\n int n_blk1 = ceil_div(total_th1, block_size);\n permute_dout_x_kernel<<>>(x, dout, x_perm, dout_perm, B, C_in, C_out, H, W);\n\n // 2.a sgemm call for dw = x_perm.T @ dout_perm\n const float alpha = 1.0f;\n const float beta = 0.0f;\n int non_C_dim = B * H * W;\n cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, C_in, C_out, non_C_dim, &alpha, x_perm, non_C_dim, dout_perm, non_C_dim, &beta, dweight, C_in));\n\n // 2.b sgemm call for dx_perm = dout_perm @ weight.T\n cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, non_C_dim, C_in, C_out, &alpha, dout_perm, non_C_dim, weight, C_in, &beta, dx_perm, non_C_dim));\n\n // 3. permute dx_perm from (C_out, B, H, W) to (B, C_out, H, W)\n int total_th2 = C_in * B * H * W;\n int n_blk2 = ceil_div(total_th2, block_size);\n permute_dx_kernel<<>>(dx_perm, dx, B, C_in, H, W);\n\n // 4. reduce dout to get dbias\n dbias_kernel<<>>(dout_perm, dbias, B, C_out, H, W);\n\n // free memory\n cudaCheck(cudaFree(x_perm));\n cudaCheck(cudaFree(dout_perm));\n cudaCheck(cudaFree(dx_perm));\n}\n\n// ------------------------------------------------------------------------------------------------\n// Conv2d 3x3 layer\ntypedef struct {\n float* w;\n float* b;\n} ConvK3Params;\n\ninline void convk3_set_param_ptrs(\n ConvK3Params* params,\n float* params_memory,\n int C_in, int C_out\n) {\n params->w = params_memory;\n params->b = params->w + C_in * C_out * 9;\n}\n\ninline size_t convk3_count_params(int C_in, int C_out) {\n return C_in * C_out * 9 + C_out;\n}\n\ntypedef struct {\n float* inp;\n float* out;\n} ConvK3Acts;\n\n// ---------------------------------------------\n// kernels\n\n// this kernel will roughly follow kernel 5 here: https://github.com/siboehm/SGEMM_CUDA, with a few differences\n// the main operation is convolving the input X with the weights\n// weights are shaped (O, Cx9), and X is shaped (B, C, HxW)\n// We are morally doing a matmul with dimensions O, C, HxW\ntemplate \n__global__ void conv2d_k3_forward_kernel2(\n const float* x, const float* weight, const float* bias,\n float* out,\n const int B, const int C, const int O, const int H, const int W\n) {\n // constants\n const int C9 = C * 9;\n const int weight_size = O * C * 9;\n const int x_size = B * C * H * W;\n // n_threads / block = (BO * BH) / (TO * TH) = 256\n // n_blocks = B * (O / BO) * (H * W / BH)\n // we will set gridDim to (BxHxW / BH, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // for compute step, to determine which of the size TOxTH block the thread is responsible for\n const int threads_per_col = ceil_div(BH, TH);\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // shared mem for loading weights and X\n const int w_s_sz = BO * BC * 9;\n const int w_s_width = BC * 9;\n const int x_s_sz = BC * BH;\n __shared__ float w_s[w_s_sz];\n __shared__ float x_s[x_s_sz];\n\n // registers to store local results\n float thread_results[TO * TH * 9] = {0.0};\n // registers for weights and x\n float reg_w[TO] = {0.0};\n float reg_x[TH] = {0.0};\n\n // calculate offset in x for bounds checking later\n const int x_block_offset = batch * C * H * W + block_col * BH;\n // calculate weight offset for bounds checking\n const int w_block_offset = block_row * BO * C9;\n\n // outer loop over BC size blocks of channel C\n for (int c_idx = 0; c_idx < C; c_idx += BC) {\n // x and weight offsets for this block\n int x_inner_offset = c_idx * H * W;\n int w_inner_offset = c_idx * 9;\n // load weight into smem\n int w_abs_idx;\n for (int i = threadIdx.x; i < w_s_sz; i += blockDim.x) {\n int w_i = i / w_s_width;\n int w_j = i % w_s_width;\n w_abs_idx = w_block_offset + w_inner_offset + w_i * C9 + w_j;\n if (w_abs_idx < weight_size) {\n w_s[w_j * BO + w_i] = weight[w_abs_idx];\n }\n }\n // load x into smem\n int x_abs_idx;\n for (int i = threadIdx.x; i < x_s_sz; i += blockDim.x) {\n int x_i = i / BH;\n int x_j = i % BH;\n x_abs_idx = x_block_offset + x_inner_offset + x_i * H * W + x_j;\n if (x_abs_idx < x_size) {\n x_s[x_i * BH + x_j] = x[x_abs_idx];\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int dot_idx = 0; dot_idx < BC; dot_idx++) {\n // load x_s into registers\n for (int j = 0; j < TH; j++) {\n reg_x[j] = x_s[dot_idx * BH + thread_col * TH + j];\n }\n for (int conv_idx = 0; conv_idx < 9; conv_idx++) {\n for (int i = 0; i < TO; i++) {\n reg_w[i] = w_s[(dot_idx * 9 + conv_idx) * BO + thread_row * TO + i];\n }\n\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n float val = reg_w[i] * reg_x[j];\n thread_results[conv_idx * TO * TH + i * TH + j] += val;\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n // need to be careful here: for each pixel at (h_abs, w_abs) and its product with kernel k (k < 9)\n // There are 2 options\n // 1. the product should be added to some pixel, then do an atomiAdd\n // 2. the product should not contribute to any pixel because it would be out of bound, in this case just skip\n int out_batch_offset = batch * O * H * W;\n for (int j = 0; j < TH; j++) {\n int hw_abs = block_col * BH + thread_col * TH + j;\n int h_abs = hw_abs / W;\n int w_abs = hw_abs % W;\n for (int i = 0; i < TO; i++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target_abs = h_abs + 1 - k_i;\n int w_target_abs = w_abs + 1 - k_j;\n if (h_target_abs >= 0 && h_target_abs < H && w_target_abs >= 0 && w_target_abs < W) {\n int out_target_abs = out_batch_offset + o_abs * H * W + h_target_abs * W + w_target_abs;\n int th_res_idx = k * TO * TH + i * TH + j;\n atomicAdd(out + out_target_abs, thread_results[th_res_idx]);\n }\n }\n int out_abs = out_batch_offset + o_abs * H * W + h_abs * W + w_abs;\n atomicAdd(out + out_abs, bias[o_abs]);\n }\n }\n}\n\n// almost same as 2, but we vectorize the loads and use float4\ntemplate \n__global__ void conv2d_k3_forward_kernel3(\n float* x, float* weight, float* bias,\n float* out,\n const int B, const int C, const int O, const int H, const int W\n) {\n // constants\n assert(TH == 4); // needed for reg_x loads later, so we can do one vectorized load for each thread\n const int C9 = C * 9;\n // n_threads / block = (BO * BH) / (TO * TH) = 256\n // n_blocks = B * (O / BO) * (H * W / BH)\n // we will set gridDim to (BxHxW / BH, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // for compute step, to determine which of the size TOxTH block the thread is responsible for\n const int threads_per_col = ceil_div(BH, TH);\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // shared mem for loading weights and X\n const int w_s_sz = BO * BC * 9;\n const int w_s_width = BC * 9;\n const int x_s_sz = BC * BH;\n __shared__ float w_s[w_s_sz];\n __shared__ float x_s[x_s_sz];\n\n // registers to store local results\n float thread_results[TO * TH * 9] = {0.0};\n // registers for weights and x\n float reg_w[TO] = {0.0};\n float reg_x[TH] = {0.0};\n\n // calculate offset in x for bounds checking later\n const int x_block_offset = batch * C * H * W + block_col * BH;\n // calculate weight offset for bounds checking\n const int w_block_offset = block_row * BO * C9;\n\n // outer loop over BC size blocks of channel C\n for (int c_idx = 0; c_idx < C; c_idx += BC) {\n // x and weight offsets for this block\n int x_inner_offset = c_idx * H * W;\n int w_inner_offset = c_idx * 9;\n // load weight into smem\n for (int i = threadIdx.x; i < w_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (w_s_width / 4);\n int inner_row = i / (w_s_width / 4);\n float4 tmp = \n reinterpret_cast(&weight[w_block_offset + w_inner_offset + inner_row * C9 + inner_col * 4])[0];\n w_s[(inner_col * 4 + 0) * BO + inner_row] = tmp.x;\n w_s[(inner_col * 4 + 1) * BO + inner_row] = tmp.y;\n w_s[(inner_col * 4 + 2) * BO + inner_row] = tmp.z;\n w_s[(inner_col * 4 + 3) * BO + inner_row] = tmp.w;\n }\n // load x into smem\n for (int i = threadIdx.x; i < x_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (BH / 4);\n int inner_row = i / (BH / 4);\n reinterpret_cast(&x_s[inner_row * BH + inner_col * 4])[0] = \n reinterpret_cast(&x[x_block_offset + x_inner_offset + inner_row * H * W + inner_col * 4])[0];\n }\n __syncthreads();\n\n // calculate per thread results\n for (int dot_idx = 0; dot_idx < BC; dot_idx++) {\n // load x_s into registers\n // NOTE: relying on the fact that TH = 4\n reinterpret_cast(reg_x)[0] = reinterpret_cast(&x_s[dot_idx * BH + thread_col * TH])[0];\n for (int conv_idx = 0; conv_idx < 9; conv_idx++) {\n for (int i = 0; i < TO; i++) {\n reg_w[i] = w_s[(dot_idx * 9 + conv_idx) * BO + thread_row * TO + i];\n }\n\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n float val = reg_w[i] * reg_x[j];\n thread_results[conv_idx * TO * TH + i * TH + j] += val;\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n // need to be careful here: for each pixel at (h_abs, w_abs) and its product with kernel k (k < 9)\n // There are 2 options\n // 1. the product should be added to some pixel, then do an atomiAdd\n // 2. the product should not contribute to any pixel because it would be out of bound, in this case just skip\n int out_batch_offset = batch * O * H * W;\n for (int j = 0; j < TH; j++) {\n int hw_abs = block_col * BH + thread_col * TH + j;\n int h_abs = hw_abs / W;\n int w_abs = hw_abs % W;\n for (int i = 0; i < TO; i++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target_abs = h_abs + 1 - k_i;\n int w_target_abs = w_abs + 1 - k_j;\n if (h_target_abs >= 0 && h_target_abs < H && w_target_abs >= 0 && w_target_abs < W) {\n int out_target_abs = out_batch_offset + o_abs * H * W + h_target_abs * W + w_target_abs;\n int th_res_idx = k * TO * TH + i * TH + j;\n atomicAdd(out + out_target_abs, thread_results[th_res_idx]);\n }\n }\n int out_abs = out_batch_offset + o_abs * H * W + h_abs * W + w_abs;\n atomicAdd(out + out_abs, bias[o_abs]);\n }\n }\n}\n\nvoid conv2d_k3_forward3(\n float* x, float* weight, float* bias,\n float* out,\n const int B, const int C_in, const int C_out, const int H, const int W\n) {\n // since the kernel will add to the output, we need to zero it first\n cudaCheck(cudaMemset(out, 0, B * C_out * H * W * sizeof(float)));\n\n if (H * W == 64) {\n const int BH = 64;\n const int TH = 4;\n assert(C_out % 64 == 0 && C_in % 8 == 0);\n const int BC = 16;\n const int BO = 64;\n const int TO = 4;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel3<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n } else if (C_in == 3) {\n const int BH = 64;\n assert(H * W % BH == 0);\n const int TH = 4;\n const int BC = 3;\n assert(C_out % 64 == 0);\n const int BO = 64;\n const int TO = 4;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel2<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n } else if (C_out == 3) {\n const int BH = 64 * 2;\n assert(H * W % BH == 0);\n const int TH = 4;\n assert(C_in % 8 == 0);\n const int BC = 8;\n const int BO = 3;\n const int TO = 3;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel2<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n } else {\n const int BH = 64 * 4;\n assert(H * W % BH == 0);\n const int TH = 4;\n assert(C_out % 64 == 0 && C_in % 8 == 0);\n const int BC = 16;\n const int BO = 8;\n const int TO = 2;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel3<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n }\n cudaCheck(cudaGetLastError());\n}\n\n// almost the same as conv2d_k3_forward_kernel, except that the compute and the writes\n// are different, because now we want to treat the C dimension as the inner product dimension\ntemplate \n__global__ void dx_backward_kernel1(\n float* dout, float* weight,\n float* dx,\n const int B, const int C, const int O, const int H, const int W\n) {\n const int C9 = C * 9;\n // n_threads per block = (BC * BH) / (TC * TH)\n // set gridDim to (BxHxW / BH, C / TC)\n // find the output block that this thread block is responsible for\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // index for the TCxTH block that the thread is responsible\n const int threads_per_col = BH / TH;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for weights and dout\n const int w_s_width = BC * 9;\n const int w_s_sz = BO * w_s_width;\n const int dout_s_sz = BO * BH;\n __shared__ float w_s[w_s_sz];\n __shared__ float dout_s[dout_s_sz];\n\n // registers for local results\n float thread_results[TC * TH * 9] = {0.0};\n // regs for weights and dout\n float reg_w[TC * 9] = {0.0};\n float reg_dout[TH] = {0.0};\n\n // calculate the block offsets\n const int dout_block_offset = batch * O * H * W + block_col * BH;\n const int w_block_offset = block_row * w_s_width;\n\n // outer loop over BO sizes of channel O\n for (int o_idx = 0; o_idx < O; o_idx += BO) {\n // offsets for this inner block\n int w_inner_offset = o_idx * C9;\n int dout_inner_offset = o_idx * H * W;\n\n // load weights to smem\n if (BC % 4 == 0) {\n // vectorize dram load\n for (int i = threadIdx.x; i < w_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (w_s_width / 4);\n int inner_row = i / (w_s_width / 4);\n reinterpret_cast(&w_s[inner_row * w_s_width + inner_col * 4])[0] = \n reinterpret_cast(&weight[w_block_offset + w_inner_offset + inner_row * C9 + inner_col * 4])[0];\n }\n } else {\n for (int i = threadIdx.x; i < w_s_sz; i += blockDim.x) {\n int inner_col = i % w_s_width;\n int inner_row = i / w_s_width;\n w_s[inner_row * w_s_width + inner_col] = weight[w_block_offset + w_inner_offset + inner_row * C9 + inner_col];\n }\n }\n // load dout\n for (int i = threadIdx.x; i < dout_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (BH / 4);\n int inner_row = i / (BH / 4);\n reinterpret_cast(&dout_s[inner_row * BH + inner_col * 4])[0] = \n reinterpret_cast(&dout[dout_block_offset + dout_inner_offset + inner_row * H * W + inner_col * 4])[0];\n }\n __syncthreads();\n \n\n // calculate per thread results\n for (int dot_idx = 0; dot_idx < BO; dot_idx++) {\n // load dout and w into registers\n for (int j = 0; j < (TH / 4); j++) {\n reinterpret_cast(®_dout[j * 4])[0] = \n reinterpret_cast(&dout_s[dot_idx * BH + thread_col * TH + j * 4])[0];\n }\n // non vectorized load\n //for (int j = 0; j < TH; j++) {\n // reg_dout[j] = dout_s[dot_idx * BH + thread_col * TH + j];\n //}\n if (TC % 4 == 0) {\n // vectorize smem load\n for (int i = 0; i < (TC * 9 / 4); i++) {\n reinterpret_cast(®_w[i * 4])[0] = \n reinterpret_cast(&w_s[dot_idx * w_s_width + thread_row * TC * 9 + i * 4])[0];\n }\n } else {\n for (int i = 0; i < TC * 9; i++) {\n reg_w[i] = w_s[dot_idx * w_s_width + thread_row * TC * 9 + i];\n }\n }\n\n for (int k = 0; k < 9; k++) {\n for (int i = 0; i < TC; i++) {\n for (int j = 0; j < TH; j++) {\n thread_results[k * TC * TH + i * TH + j] += \n reg_w[i * 9 + k] * reg_dout[j];\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n int out_batch_offset = batch * C * H * W;\n for (int j = 0; j < TH; j++) {\n int hw_abs = block_col * BH + thread_col * TH + j;\n int h_abs = hw_abs / W;\n int w_abs = hw_abs % W;\n for (int i = 0; i < TC; i++) {\n int c_abs = block_row * BC + thread_row * TC + i;\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target_abs = h_abs - 1 + k_i;\n int w_target_abs = w_abs - 1 + k_j;\n if (h_target_abs >= 0 && h_target_abs < H && w_target_abs >= 0 && w_target_abs < W) {\n int out_target_abs = out_batch_offset + c_abs * H * W + h_target_abs * W + w_target_abs;\n int th_res_idx = k * TC * TH + i * TH + j;\n atomicAdd(dx + out_target_abs, thread_results[th_res_idx]);\n }\n }\n }\n }\n}\n\nvoid dx_backward(\n float* dout, float* weight,\n float* dx,\n const int B, const int C_in, const int C_out, const int H, const int W\n) {\n if (H < 16) {\n const int BH = 64;\n const int TH = 4;\n const int BC = 16;\n const int TC = 4;\n const int BO = 16;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n } else if (C_in == 3) {\n const int BH = 64 * 4;\n const int TH = 4;\n const int BC = 3;\n const int TC = 3;\n const int BO = 16;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n } else if (C_out == 3) {\n const int BH = 64 * 4;\n const int TH = 4;\n const int BC = 16;\n const int TC = 4;\n const int BO = 3;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n } else {\n const int BH = 64 * 4;\n const int TH = 4;\n const int BC = 16;\n const int TC = 4;\n const int BO = 16;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n }\n}\n\n// same as dweight backward 4, except we also do some of the reduce for dbias\n// so it doesn't have to reduce an entire BxHxW length for each O channel\ntemplate \n__global__ void dweight_dbias_backward_kernel1(\n float* dout, float* x,\n float* dweight_buf, float* dbias_buf,\n const int B, const int C, const int O, const int H, const int W\n) {\n // number of heights we need to load into smem in one go \n const int NH = BW / WSZ;\n // n_threads per block = (BC * BO) / (TO * TC)\n // gridDim will be (B * (W / (WSZ - 1)) * (C / BC), O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_chan = C / BC;\n const int blocks_per_width = W / (WSZ - 1);\n const int block_w = blockIdx.x % blocks_per_width;\n const int block_col = (blockIdx.x / blocks_per_width) % blocks_per_chan;\n const int batch = blockIdx.x / (blocks_per_chan * blocks_per_width);\n const int w_offset = block_w * (WSZ - 1);\n\n // warp indices for loading from global mem\n const int n_warps = ceil_div((int)blockDim.x, 32);\n const int warp_id = threadIdx.x / 32;\n const int lane_id = threadIdx.x % 32;\n // make sure that the width sections can be loaded by a single warp\n assert(WSZ <= 32);\n\n // index for the TO*TC thread block\n const int threads_per_col = BC / TC;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for dout and x\n __shared__ float dout_s[BO * BW];\n __shared__ float x_s[BC * BW];\n\n // registers for local mem and local results\n float thread_results[TO * TC * 9] = {0.0};\n float reg_dout[TO] = {0.0};\n float reg_x[TC] = {0.0};\n float thread_bias_results[TO] = {0.0};\n\n // block offsets\n const int dout_block_offset = batch * O * H * W + block_row * BO * H * W;\n const int x_block_offset = batch * C * H * W + block_col * BC * H * W;\n\n //int debug_idx = 0;\n\n // we will fuse the calculation for all 9 kernel filters\n // which means we will have to reload some of the image widths\n // we can load NW blocks into smem in one go, which means we will be stepping\n // over the H dimension in (NW - 1) steps\n // this is the outer loop over H\n int h_step = (NH < H) ? NH - 1 : H;\n // if NH == H, then we are loading the entire image in one go, so step by H\n for (int h = 0; h < H; h += h_step) {\n // load dout into smem\n int h_offset = h * W;\n int nh = (h < H - (NH - 1)) ? NH : NH - 1;\n int wsz = (w_offset < W - (WSZ - 1)) ? WSZ : WSZ - 1;\n int o_per_warp = ceil_div(BO, n_warps);\n if (lane_id < wsz) {\n for (int warp_o = 0; warp_o < o_per_warp; warp_o++) {\n int inner_o = warp_id * o_per_warp + warp_o;\n if (inner_o < BO) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n dout_s[(inner_h * WSZ + inner_w) * BO + inner_o] = \n dout[dout_block_offset + inner_o * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n int c_per_warp = ceil_div(BC, n_warps);\n if (lane_id < wsz) {\n for (int warp_c = 0; warp_c < c_per_warp; warp_c++) {\n int inner_c = warp_id * c_per_warp + warp_c;\n if (inner_c < BC) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n x_s[(inner_h * WSZ + inner_w) * BC + inner_c] = \n x[x_block_offset + inner_c * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n __syncthreads();\n \n\n // if loading the last block, fill in the last block's buffer to be 0\n if (nh == NH - 1) {\n for (int i = threadIdx.x; i < WSZ * BO; i++) {\n int inner_w = i % WSZ;\n int inner_o = i / WSZ;\n dout_s[(inner_w + nh * WSZ) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < WSZ * BC; i++) {\n int inner_w = i % WSZ;\n int inner_c = i / WSZ;\n\n x_s[(inner_w + nh * WSZ) * BC + inner_c] = 0.0f;\n } \n }\n // if loading the last width section, fill in the final value in the buffer to 0\n if (wsz == WSZ - 1) {\n for (int i = threadIdx.x; i < nh * BO; i++) {\n int inner_o = i % BO;\n int inner_h = i / BO;\n dout_s[(inner_h * WSZ + wsz) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < nh * BC; i++) {\n int inner_c = i % BC;\n int inner_h = i / BC;\n x_s[(inner_h * WSZ + wsz) * BC + inner_c] = 0.0f;\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int x_dot_idx = 0; x_dot_idx < WSZ * nh; x_dot_idx++) {\n int x_h = x_dot_idx / WSZ;\n int x_w = x_dot_idx % WSZ;\n // load TC of the C channels of one image pixel into reg_w\n for (int j = 0; j < TC; j++) {\n reg_x[j] = x_s[x_dot_idx * BC + thread_col * TC + j];\n }\n\n // for each conv filter, load a different row and calculate the outer product\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n // calculate the corresponding (h, w) that should be loaded from dout\n int dout_h = x_h + 1 - k_i;\n int dout_w = x_w + 1 - k_j;\n int dout_h_abs = h + dout_h;\n int dout_w_abs = w_offset + dout_w;\n int dout_dot_idx = -1;\n if (dout_h >= 0 && dout_h < NH && dout_w >= 0 && dout_w < WSZ && dout_w_abs < W && dout_h_abs < H) {\n if (x_h == NH - 1 && k_i == 1 && H > NH) {\n // this logic is a bit tricky: for the last row in smem\n // because it will be loaded as part of the next block as well, and it will be the first row in the next block\n // for kernels 4, 5, 6 we need to only compute them once\n // this statement ensures that we don't compute the 4, 5, 6 kernels for the last row in smem\n // so it will be computed as part of the first row next time\n dout_dot_idx = -1;\n } else if (x_w == WSZ - 1 && k_j == 1 && W > WSZ) {\n // if x_w is at the last pixel, and we have more than one width sections, and we are not are the last section\n // then don't compute 1, 4, 7 kernels and wait for the next section to compute them\n dout_dot_idx = -1;\n } else {\n dout_dot_idx = dout_h * WSZ + dout_w;\n }\n }\n\n // if in bounds of the dout section in smem, load the corresponding pixel from smem\n // and add the products to the thread results\n // otherwise the contribution is 0\n if (dout_dot_idx != -1) {\n for (int i = 0; i < TO; i++) {\n reg_dout[i] = dout_s[dout_dot_idx * BO + thread_row * TO + i];\n }\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n float new_val = reg_dout[i] * reg_x[j];\n thread_results[k * TO * TC + i * TC + j] += new_val;\n }\n }\n // update dbias when k = 4, where dout_h and dout_w match x_h and x_w\n if (k == 4 && block_col == 0) {\n for (int i = 0; i < TO; i++) {\n thread_bias_results[i] += reg_dout[i];\n }\n }\n }\n }\n __syncthreads();\n }\n }\n\n // write out dweight results\n for (int i = 0; i < TO; i++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n for (int j = 0; j < TC; j++) {\n int c_abs = block_col * BC + thread_col * TC + j;\n for (int k = 0; k < 9; k++) {\n dweight_buf[o_abs * C * 9 * B * blocks_per_width + c_abs * 9 * B * blocks_per_width + k * B * blocks_per_width + batch * blocks_per_width + block_w] = thread_results[k * TO * TC + i * TC + j];\n }\n }\n // write out dbias results\n if (block_col == 0) {\n dbias_buf[o_abs * B * blocks_per_width + batch * blocks_per_width + block_w] = thread_bias_results[i];\n }\n }\n}\n\n// C * O blocks\n// 288 (9 * 32) threads per block, use one warp to reduce each batch dim\n__global__ void dweight_reduce_kernel(\n float* dweight_buf, float* dweight,\n const int B, const int C, const int O\n) {\n assert(gridDim.x == O * C);\n int o = blockIdx.x / C;\n int c = blockIdx.x % C;\n // move input pointer\n dweight_buf += o * C * 9 * B + c * 9 * B;\n assert(blockDim.x == 32 * 9);\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n\n float thread_sum = 0.0f;\n for (int i = lane_id; i < B; i += 32) {\n thread_sum += dweight_buf[warp_id * B + i];\n }\n\n // warp reduce\n float warp_sum = cg::reduce(warp, thread_sum, cg::plus{});\n\n if (lane_id == 0) {\n dweight[o * C * 9 + c * 9 + warp_id] = warp_sum;\n }\n}\n\n__global__ void dbias_reduce_kernel1(\n float* dbias_buf, float* dbias,\n const int B\n) {\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n\n dbias_buf += blockIdx.x * B;\n\n int lane_id = threadIdx.x % 32;\n\n float db_thread_sum = 0.0f;\n for (int i = lane_id; i < B; i += 32) {\n db_thread_sum += dbias_buf[i];\n }\n\n float db_warp_sum = cg::reduce(warp, db_thread_sum, cg::plus{});\n\n if (lane_id == 0) {\n dbias[blockIdx.x] = db_warp_sum;\n }\n}\n\nvoid dweight_dbias_backward1(\n float* dout, float* x,\n float* dweight_buf, float* dweight,\n float* dbias_buf, float* dbias,\n const int B, const int C, const int O, const int H, const int W\n) {\n if (H * W == 64) {\n const int BC = 64;\n const int BO = 64;\n const int WSZ = 8;\n const int BW = WSZ * 8;\n const int TO = 4;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(cudaDeviceSynchronize());\n cudaCheck(cudaGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B, C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B);\n } else if (C == 3) {\n const int BC = 3;\n const int BO = 64;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 3;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(cudaDeviceSynchronize());\n cudaCheck(cudaGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B * W / (WSZ - 1));\n } else if (O == 3) {\n const int BC = 64;\n const int BO = 3;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 3;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(cudaDeviceSynchronize());\n cudaCheck(cudaGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B * W / (WSZ - 1));\n } else {\n const int BC = 64;\n const int BO = 64;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(cudaDeviceSynchronize());\n cudaCheck(cudaGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B * W / (WSZ - 1));\n }\n}\n\nvoid conv2d_k3_backward2(\n float* dout, float* x, float* weight,\n float* dweight_buf, float* dbias_buf,\n float* dx, float* dweight, float* dbias,\n const int B, const int C_in, const int C_out, const int H, const int W\n) {\n // dx currently requires the output buffer to be zeroed\n cudaCheck(cudaMemset(dx, 0, B * C_in * H * W * sizeof(float)));\n dx_backward(dout, weight, dx, B, C_in, C_out, H, W);\n dweight_dbias_backward1(dout, x, dweight_buf, dweight, dbias_buf, dbias, B, C_in, C_out, H, W);\n}\n\n// ------------------------------------------------------------------------------------------------\n// GroupNorm\n\ntypedef struct {\n float* w;\n float* b;\n} GroupNormParams;\n\ninline void gn_set_param_ptrs(\n GroupNormParams* params,\n float* params_memory,\n int C\n) {\n params->w = params_memory;\n params->b = params->w + C;\n}\n\ntypedef struct {\n float* x;\n float* out;\n float* mean;\n float* rstd;\n} GroupNormActs;\n\ninline void gn_set_act_ptrs(\n GroupNormActs* acts,\n float* acts_memory,\n int B, int C, int H, int W, int n_groups\n) {\n acts->out = acts_memory;\n acts->mean = acts->out + B * C * H * W;\n acts->rstd = acts->mean + B * n_groups;\n}\n\ntypedef struct {\n float* dx;\n float* dout;\n} GroupNormBackActs;\n\n// ---------------------------------------------\n// kernels\n\n// Essentially taken from llm.c's kernel 5\n// using kernel 5 because for images, each \"channel\" is effectively\n// H * W * group_size, which is quite large\n// we will have one block per group, which means there are B * C / group_size blocks\n__global__ void groupnorm_forward_kernel(\n const float* x, const float* weight, const float* bias,\n float* out, float* mean, float* rstd,\n int B, int C, int img_size, int group_size, int n_groups\n) {\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n __shared__ float shared_sum[32]; // block_size max is 1024 = 32 * 32 warps\n __shared__ float shared_sum2[32]; // warps will be writing into shared memeory after warp-reduce\n int num_warps = blockDim.x / 32;\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n int block_pixels = img_size * group_size;\n // group index\n int g = blockIdx.x % n_groups;\n\n // move pointers\n x += blockIdx.x * img_size * group_size;\n out += blockIdx.x * img_size * group_size;\n // each block will only every acces group_size channels\n weight += g * group_size;\n bias += g * group_size;\n\n float thread_sum = 0.0f;\n float thread_sum2 = 0.0f;\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n float val = x[i];\n thread_sum += val;\n thread_sum2 += val * val;\n }\n\n // warp reduce\n float warp_sum = cg::reduce(warp, thread_sum, cg::plus{});\n float warp_sum2 = cg:: reduce(warp, thread_sum2, cg::plus{});\n // store warp sum into shared memory\n shared_sum[warp_id] = warp_sum;\n shared_sum2[warp_id] = warp_sum2;\n __syncthreads();\n \n // load warp sums from shared memory\n warp_sum = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;\n warp_sum2 = (lane_id < num_warps) ? shared_sum2[lane_id] : 0.0f;\n float block_sum = cg::reduce(warp, warp_sum, cg::plus{});\n float block_sum2 = cg::reduce(warp, warp_sum2, cg::plus{});\n block_sum /= block_pixels;\n block_sum2 /= block_pixels;\n float m = block_sum;\n float var = block_sum2 - m * m;\n float s = rsqrtf(var + 1e-5f);\n if (threadIdx.x == 0 && mean != nullptr) {\n mean[blockIdx.x] = m;\n }\n if (threadIdx.x == 0 && rstd != nullptr) {\n rstd[blockIdx.x] = s;\n }\n\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n int c_mod_group = (i / img_size) % group_size;\n float n = s * (x[i] - m);\n out[i] = n * weight[c_mod_group] + bias[c_mod_group];\n }\n}\n\nvoid groupnorm_forward(\n const float* x, const float* weight, const float* bias,\n float* out, float* mean, float* rstd,\n int B, int C, int H, int W, int n_groups\n) {\n int img_size = H * W;\n int group_size = C / n_groups;\n int n_blocks = B * n_groups;\n int block_size = max_int(min_int(512, img_size * group_size), 32);\n groupnorm_forward_kernel<<>>(\n x, weight, bias, out, mean, rstd, B, C, img_size, group_size, n_groups\n );\n cudaCheck(cudaGetLastError());\n}\n\n// most similar to kernel 2 in llm.c\n// not doing any float16 optimizations yet\n// main change is to have each block allocated to a single group\n__global__ void groupnorm_backward_kernel(\n const float* dout, const float* x, const float* mean, const float* rstd, const float* weight,\n float* dx, float* dweight, float* dbias,\n int B, int C, int img_size, int group_size, int n_groups\n) {\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n __shared__ float shared_sum[32]; // block_size max is 1024 = 32 * 32 warps\n __shared__ float shared_sum2[32]; // warps will be writing into shared memeory after warp-reduce\n int num_warps = blockDim.x / 32;\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n int block_pixels = img_size * group_size;\n // group index\n int g = blockIdx.x % n_groups;\n\n // move pointers\n dout += blockIdx.x * img_size * group_size;\n x += blockIdx.x * img_size * group_size;\n dx += blockIdx.x * img_size * group_size;\n weight += g * group_size;\n dweight += g * group_size;\n dbias += g * group_size;\n\n float m_val = mean[blockIdx.x];\n float rstd_val = rstd[blockIdx.x];\n\n\n // calculate the two mean terms in the group dimension\n // first is dout * weight, and second is dout * weight * norm\n // where norm = (x - mean) * rstd\n float w_dout_thread = 0.0f;\n float w_dout_norm_thread = 0.0f;\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n int c_mod_group = (i / img_size) % group_size;\n float cur_w_dout = weight[c_mod_group] * dout[i];\n w_dout_thread += cur_w_dout;\n float norm = (x[i] - m_val) * rstd_val;\n w_dout_norm_thread += cur_w_dout * norm;\n }\n // warp reduce\n float w_dout_warp = cg::reduce(warp, w_dout_thread, cg::plus{});\n float w_dout_norm_warp = cg::reduce(warp, w_dout_norm_thread, cg::plus{});\n // store warp sum in shared mem\n shared_sum[warp_id] = w_dout_warp;\n shared_sum2[warp_id] = w_dout_norm_warp;\n __syncthreads();\n\n // load warp sums from shared memory\n w_dout_warp = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;\n w_dout_norm_warp = (lane_id < num_warps) ? shared_sum2[lane_id] : 0.0f;\n float w_dout_block = cg::reduce(warp, w_dout_warp, cg::plus{});\n float w_dout_norm_block = cg::reduce(warp, w_dout_norm_warp, cg::plus{});\n w_dout_block /= block_pixels;\n w_dout_norm_block /= block_pixels;\n\n // update dx\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n // in bounds of image\n // accumulate dw and db\n float dout_val = dout[i];\n float norm = (x[i] - m_val) * rstd_val;\n\n // update dx\n int c_mod_group = (i / img_size) % group_size;\n float w_dout = weight[c_mod_group] * dout_val;\n dx[i] = (w_dout - w_dout_block - norm * w_dout_norm_block) * rstd_val;\n }\n // update dw and db\n // use different methods when the image size is large or small\n\n // if the image size is larger than the block size\n // loop over the channels and use the whole block on each channel\n // otherwise, assign each warp to a channel\n // in either case image size must be larger than the warp size\n //assert(img_size % warp.size() == 0);\n assert(blockDim.x % warp.size() == 0);\n if (img_size % blockDim.x == 0) {\n for (int c = 0; c < group_size; c++) {\n float dw_thread = 0.0f;\n float db_thread = 0.0f;\n for (int i = threadIdx.x; i < img_size; i += blockDim.x) {\n float dout_val = dout[i];\n db_thread += dout_val;\n float norm = (x[i] - m_val) * rstd_val;\n dw_thread += dout_val * norm;\n }\n\n // move pointers\n dout += img_size;\n x += img_size;\n\n // warp reduce\n float dw_warp = cg::reduce(warp, dw_thread, cg::plus{});\n float db_warp = cg::reduce(warp, db_thread, cg::plus{});\n ////// store warp sum in shared mem\n if (lane_id == 0) {\n shared_sum[warp_id] = dw_warp;\n shared_sum2[warp_id] = db_warp;\n }\n __syncthreads();\n // use the first thread to reduce the shared memory sums and save to global memory\n if (threadIdx.x == 0) {\n float dw_block = 0.0f;\n float db_block = 0.0f;\n for (int i = 0; i < num_warps; i++) {\n dw_block += shared_sum[i];\n db_block += shared_sum2[i];\n }\n atomicAdd(dweight + c, dw_block);\n atomicAdd(dbias + c, db_block);\n }\n }\n } else {\n // if group size is large, need to loop over the group channels with the whole block\n int block_reps = ceil_div(group_size, num_warps);\n for (int br = 0; br < block_reps; br++) {\n float dw_thread = 0.0f;\n float db_thread = 0.0f;\n\n int ch = br * num_warps + warp_id;\n if (ch < group_size) {\n const float* dout_ch = dout + ch * img_size;\n const float* x_ch = x + ch * img_size;\n for (int i = lane_id; i < img_size; i += warp.size()) {\n float dout_val = dout_ch[i];\n db_thread += dout_val;\n float norm = (x_ch[i] - m_val) * rstd_val;\n dw_thread += dout_val * norm;\n }\n \n // warp reduce\n float dw_warp = cg::reduce(warp, dw_thread, cg::plus{});\n float db_warp = cg::reduce(warp, db_thread, cg::plus{});\n // since each warp takes care of an entire image\n // directly store result\n if (lane_id == 0) {\n atomicAdd(dweight + ch, dw_warp);\n atomicAdd(dbias + ch, db_warp);\n }\n }\n }\n }\n}\n\nvoid groupnorm_backward(\n const float* dout, const float* x, const float* mean, const float* rstd, const float* weight,\n float* dx, float* dweight, float* dbias,\n int B, int C, int H, int W, int n_groups\n) {\n int img_size = H * W;\n int group_size = C / n_groups;\n int n_blocks = B * n_groups;\n int block_size = max_int(min_int(512, img_size * group_size), 32 * group_size);\n groupnorm_backward_kernel<<>>(\n dout, x, mean, rstd, weight, dx, dweight, dbias, B, C, img_size, group_size, n_groups\n );\n cudaCheck(cudaGetLastError());\n}\n// ------------------------------------------------------------------------------------------------\n// ResBlock\n\n#define NUM_RES_PARAM_TENSORS 12\ntypedef struct {\n float* gn1_w; // C\n float* gn1_b; // C\n float* cv3_1_w; // C_out, C, 3, 3\n float* cv3_1_b; // C_out\n float* l_emb_w; // C_out, C_emb\n float* l_emb_b; // C_out\n float* gn2_w; // C_out\n float* gn2_b; // C_out\n float* cv3_2_w; // C_out, C_out, 3, 3\n float* cv3_2_b; // C_out\n float* res_cv1_w; // C_out, C\n float* res_cv1_b; // C_out\n size_t param_sizes[NUM_RES_PARAM_TENSORS];\n size_t n_params;\n} ResBlockParameters;\n\n#define NUM_RES_ACT_TENSORS 18 // not counting input and emb, which will come from outs of previous layers\ntypedef struct {\n float* gn1; // B, C, H, W\n float* gn1_mean; // B, n_groups\n float* gn1_rstd; // B, n_groups\n float* silu1; // B, C, H, W\n float* ud_h; // B, C, H_out, W_out\n float* ud_x; // B, C, H_out, W_out\n float* cv3_1; // B, C_out, H_out, W_out\n float* silu_emb; // B, C_emb\n float* l_emb; // B, C_out\n float* broad_emb; // B, C_out, H_out, W_out\n float* add1; // B, C_out, H_out, W_out\n float* gn2; // B, C_out, H_out, W_out\n float* gn2_mean; // B, n_groups\n float* gn2_rstd; // B, n_groups\n float* silu2; // B, C_out, H_out, W_out\n float* cv3_2; // B, C_out, H_out, W_out\n float* res_cv1; // B, C_out, H_out, W_out\n float* add2; // B, C_out, H_out, W_out\n size_t act_sizes[NUM_RES_ACT_TENSORS];\n size_t n_acts;\n float* input; // B, C, H, W\n float* emb; // B, C_emb\n} ResBlockActivations;\n\n\n// the last two backward tensors, dx and demb\n// we don't need to allocate extra space for these\n// instead we will just overwrite forward pass activations acts->input and acts->emb\n#define NUM_RES_BACKWARD_TENSORS 7\ntypedef struct {\n float* buf_BCemb;\n float* buf_BCHoWo;\n float* buf1_BCHW;\n float* buf2_BCHW;\n float* dout; // B, C_out, H_out, W_out\n float* dweight_buf;\n float* dbias_buf;\n float* dx; // B, C, H, W\n float* demb; // B, C_emb\n size_t back_sizes[NUM_RES_BACKWARD_TENSORS];\n size_t n_backs;\n} ResBlockBackwardActivations;\n\nvoid resblock_count_params(\n ResBlockParameters* params,\n int C, int C_emb, int C_out, int B, int H, int W, int up, int down, int gn_n_groups\n) {\n int H_out = H;\n int W_out = W;\n if (up) {\n H_out *= 2;\n W_out *= 2;\n } else if (down) {\n H_out /= 2;\n W_out /= 2;\n }\n // fill and count param sizes\n size_t* param_sizes = params->param_sizes;\n param_sizes[0] = C; // gn1_w\n param_sizes[1] = C; // gn1_b\n param_sizes[2] = C_out * C * 9; // cv3_1_w\n param_sizes[3] = C_out; // cv3_1_b\n param_sizes[4] = C_emb * C_out; // l_emb_w\n param_sizes[5] = C_out; // l_emb_b\n param_sizes[6] = C_out; // gn2_w\n param_sizes[7] = C_out; // gn2_b\n param_sizes[8] = C_out * C_out * 9; // cv3_2_w\n param_sizes[9] = C_out; // cv3_2_b\n param_sizes[10] = C_out * C; // res_cv1_w\n param_sizes[11] = C_out; // res_cv1_b\n if (C_out != C) {\n param_sizes[10] = C_out * C; // res_cv1_w\n param_sizes[11] = C_out; // res_cv1_b\n } else {\n param_sizes[10] = 0;\n param_sizes[11] = 0;\n }\n\n // count number of params\n size_t num_parameters = 0;\n for (size_t i = 0; i < NUM_RES_PARAM_TENSORS; i++) {\n num_parameters += param_sizes[i];\n }\n params->n_params = num_parameters;\n}\n\nvoid resblock_count_acts(\n ResBlockActivations* acts,\n int C, int C_emb, int C_out, int B, int H, int W, int up, int down, int gn_n_groups\n) {\n int H_out = H;\n int W_out = W;\n if (up) {\n H_out *= 2;\n W_out *= 2;\n } else if (down) {\n H_out /= 2;\n W_out /= 2;\n }\n // fill and count activation sizes\n size_t* acts_sizes = acts->act_sizes;\n acts_sizes[0] = B * C * H * W; // gn1\n acts_sizes[1] = B * gn_n_groups; // gn1_mean\n acts_sizes[2] = B * gn_n_groups; // gn1_rstd\n acts_sizes[3] = B * C * H * W; // silu1\n acts_sizes[4] = B * C * H_out * W_out; // ud_h\n acts_sizes[5] = B * C * H_out * W_out; // ud_x\n acts_sizes[6] = B * C_out * H_out * W_out; // cv3_1\n acts_sizes[7] = B * C_emb; // silu_emb\n acts_sizes[8] = B * C_out; // l_emb\n acts_sizes[9] = B * C_out * H_out * W_out; // broad_emb\n acts_sizes[10] = B * C_out * H_out * W_out; // add1\n acts_sizes[11] = B * C_out * H_out * W_out; // gn2\n acts_sizes[12] = B * gn_n_groups; // gn2_mean\n acts_sizes[13] = B * gn_n_groups; // gn2_rstd\n acts_sizes[14] = B * C_out * H_out * W_out; // silu2\n acts_sizes[15] = B * C_out * H_out * W_out; // cv3_2\n acts_sizes[16] = (C_out == C) ? 0 : B * C_out * H_out * W_out; // res_cv1\n acts_sizes[17] = B * C_out * H_out * W_out; // add2\n\n size_t num_acts_params = 0;\n for (size_t i = 0; i < NUM_RES_ACT_TENSORS; i++) {\n num_acts_params += acts_sizes[i];\n }\n acts->n_acts = num_acts_params;\n}\n\nvoid set_resblock_params_ptrs(\n int C,\n int C_out,\n ResBlockParameters *params,\n float* params_memory_gpu\n) {\n float** ptrs[] = {\n ¶ms->gn1_w, ¶ms->gn1_b, ¶ms->cv3_1_w, ¶ms->cv3_1_b,\n ¶ms->l_emb_w, ¶ms->l_emb_b, ¶ms->gn2_w, ¶ms->gn2_b,\n ¶ms->cv3_2_w, ¶ms->cv3_2_b, ¶ms->res_cv1_w, ¶ms->res_cv1_b\n };\n int num_param_tensors = NUM_RES_PARAM_TENSORS;\n if (C_out == C) {\n num_param_tensors -= 2;\n }\n float* params_ptr = params_memory_gpu;\n for (int i = 0; i < num_param_tensors; i++) {\n *(ptrs[i]) = params_ptr;\n params_ptr += params->param_sizes[i];\n }\n if (C_out == C) {\n params->res_cv1_w = nullptr;\n params->res_cv1_b = nullptr;\n }\n}\n\nvoid set_resblock_acts_ptrs(\n int C,\n int C_out,\n ResBlockActivations *acts,\n float* acts_memory\n) {\n float** acts_ptrs[] = {\n &acts->gn1, &acts->gn1_mean, &acts->gn1_rstd, &acts->silu1,\n &acts->ud_h, &acts->ud_x, &acts->cv3_1, &acts->silu_emb, &acts->l_emb, &acts->broad_emb,\n &acts->add1, &acts->gn2, &acts->gn2_mean, &acts->gn2_rstd, &acts->silu2, &acts->cv3_2,\n &acts->res_cv1, &acts->add2\n };\n float* acts_ptr = acts_memory;\n for (int i = 0; i < NUM_RES_ACT_TENSORS - 2; i++) {\n *(acts_ptrs[i]) = acts_ptr;\n acts_ptr += acts->act_sizes[i];\n }\n\n if (C_out == C) {\n acts->res_cv1 = nullptr;\n acts->add2 = acts_ptr;\n } else {\n acts->res_cv1 = acts_ptr;\n acts_ptr += acts->act_sizes[NUM_RES_ACT_TENSORS - 2];\n acts->add2 = acts_ptr;\n }\n}\n\n// assume all memory is allocated, run the forward pass\nvoid resblock_forward(\n cublasHandle_t cublas_handle,\n int C,\n int C_emb,\n int C_out,\n int B,\n int H,\n int W,\n int block_size,\n int up,\n int down,\n int gn_n_groups,\n ResBlockParameters *params,\n ResBlockActivations *acts\n) {\n int H_out = H;\n int W_out = W;\n if (up) {\n H_out *= 2;\n W_out *= 2;\n } else if (down) {\n H_out /= 2;\n W_out /= 2;\n }\n // up to first 3x3 conv in BigGAN res block\n // note: residual stream needs a convolution at the end if C != C_out\n groupnorm_forward(acts->input, params->gn1_w, params->gn1_b, acts->gn1, acts->gn1_mean, acts->gn1_rstd,\n B, C, H, W, gn_n_groups);\n cudaCheck(cudaGetLastError());\n silu_forward(acts->gn1, acts->silu1, B * C * H * W, block_size);\n cudaCheck(cudaGetLastError());\n if (up) {\n upsample_forward1(acts->ud_h, acts->silu1, B, C, H, W, block_size);\n upsample_forward1(acts->ud_x, acts->input, B, C, H, W, block_size);\n } else if (down) {\n avgpool_2d_forward1(acts->ud_h, acts->silu1, B, C, H, W, block_size);\n avgpool_2d_forward1(acts->ud_x, acts->input, B, C, H, W, block_size);\n } else {\n cudaCheck(cudaMemcpy(acts->ud_h, acts->silu1, B * C * H * W * sizeof(float), cudaMemcpyDeviceToDevice));\n cudaCheck(cudaMemcpy(acts->ud_x, acts->input, B * C * H * W * sizeof(float), cudaMemcpyDeviceToDevice));\n }\n cudaCheck(cudaGetLastError());\n conv2d_k3_forward3(acts->ud_h, params->cv3_1_w, params->cv3_1_b, acts->cv3_1, B, C, C_out, H_out, W_out);\n cudaCheck(cudaGetLastError());\n\n // project and add embedding layer\n silu_forward(acts->emb, acts->silu_emb, B * C_emb, block_size);\n cudaCheck(cudaGetLastError());\n matmul_forward2(cublas_handle, acts->l_emb, acts->silu_emb, params->l_emb_w, params->l_emb_b, B, C_emb, C_out, block_size);\n cudaCheck(cudaGetLastError());\n broadcast_last_dims_forward(acts->l_emb, acts->broad_emb, B * C_out, H_out, W_out, block_size);\n cudaCheck(cudaGetLastError());\n add_forward(acts->cv3_1, acts->broad_emb, acts->add1, B * C_out * H_out * W_out, block_size);\n cudaCheck(cudaGetLastError());\n\n // second group of norm -> silu -> 3x3 conv\n groupnorm_forward(acts->add1, params->gn2_w, params->gn2_b, acts->gn2, acts->gn2_mean, acts->gn2_rstd,\n B, C_out, H_out, W_out, gn_n_groups);\n cudaCheck(cudaGetLastError());\n silu_forward(acts->gn2, acts->silu2, B * C_out * H_out * W_out, block_size);\n cudaCheck(cudaGetLastError());\n conv2d_k3_forward3(acts->silu2, params->cv3_2_w, params->cv3_2_b, acts->cv3_2, B, C_out, C_out, H_out, W_out);\n cudaCheck(cudaGetLastError());\n\n // add residual (change channels with 1x1 conv if neccessary)\n if (C_out == C) {\n add_forward(acts->cv3_2, acts->ud_x, acts->add2, B * C_out * H_out * W_out, block_size);\n cudaCheck(cudaGetLastError());\n } else {\n conv2d_k1_forward2(acts->ud_x, params->res_cv1_w, params->res_cv1_b, acts->res_cv1, B, C, H_out, W_out, C_out);\n cudaCheck(cudaGetLastError());\n add_forward(acts->cv3_2, acts->res_cv1, acts->add2, B * C_out * H_out * W_out, block_size);\n cudaCheck(cudaGetLastError());\n }\n}\n\nvoid resblock_backward(\n cublasHandle_t cublas_handle,\n int C,\n int C_emb,\n int C_out,\n int B,\n int H,\n int W,\n int block_size,\n int up,\n int down,\n int gn_n_groups,\n ResBlockParameters *params,\n ResBlockParameters *grads,\n ResBlockActivations *acts,\n ResBlockBackwardActivations *back_acts\n) {\n int H_out = H;\n int W_out = W;\n if (up) {\n H_out *= 2;\n W_out *= 2;\n } else if (down) {\n H_out /= 2;\n W_out /= 2;\n }\n // Note: we're using the activation of the final conv layer as a buffer to store\n // backward pass gradients, because we don't need it in the backward pass anymore.\n float *demb = back_acts->demb;\n float *dx = back_acts->dx;\n float* dout = back_acts->dout;\n float *res_dx = back_acts->buf_BCHoWo;\n float* buf1_BCoHoWo = acts->cv3_2;\n float* buf2_BCoHoWo = acts->add2;\n cudaCheck(cudaMemset(buf1_BCoHoWo, 0, B * C_out * H_out * W_out * sizeof(float)));\n cudaCheck(cudaMemset(buf2_BCoHoWo, 0, B * C_out * H_out * W_out * sizeof(float)));\n // Note: we don't need to run a separate kernel to backprop through an add operation\n // we can just share the dout buffer between the two inputs\n if (C_out != C) {\n conv2d_k1_backward1(cublas_handle, dout, acts->ud_x, params->res_cv1_w, res_dx, grads->res_cv1_w, grads->res_cv1_b, B, C, C_out, H_out, W_out, block_size);\n cudaCheck(cudaGetLastError());\n } else {\n // point res_dx to dout\n res_dx = dout;\n }\n // NOTE: after backprop through cv1 above, we can use ud_x as a (B, C, H_out, W_out) buffer\n float* buf1_BCHoWo = acts->ud_x;\n\n // second group of norm -> silu -> 3x3 conv\n conv2d_k3_backward2(dout, acts->silu2, params->cv3_2_w, back_acts->dweight_buf, back_acts->dbias_buf, buf1_BCoHoWo, grads->cv3_2_w, grads->cv3_2_b, B, C_out, C_out, H_out, W_out);\n cudaCheck(cudaGetLastError());\n silu_backward(buf1_BCoHoWo, acts->gn2, buf2_BCoHoWo, B * C_out * H_out * W_out, block_size);\n cudaCheck(cudaGetLastError());\n groupnorm_backward(buf2_BCoHoWo, acts->add1, acts->gn2_mean, acts->gn2_rstd, params->gn2_w, buf1_BCoHoWo, grads->gn2_w, grads->gn2_b, B, C_out, H_out, W_out, gn_n_groups);\n cudaCheck(cudaGetLastError());\n\n // emb layer\n // note at this point we can also use the last layer activations as buffers\n broadcast_last_dims_backward(buf1_BCoHoWo, acts->l_emb, B * C_out, H_out, W_out, block_size);\n cudaCheck(cudaGetLastError());\n matmul_backward1(cublas_handle, back_acts->buf_BCemb, grads->l_emb_w, grads->l_emb_b, acts->l_emb, acts->silu_emb, params->l_emb_w, B, C_emb, C_out);\n cudaCheck(cudaGetLastError());\n silu_backward(back_acts->buf_BCemb, acts->emb, demb, B * C_emb, block_size);\n cudaCheck(cudaGetLastError());\n\n // first layers\n conv2d_k3_backward2(buf1_BCoHoWo, acts->ud_h, params->cv3_1_w, back_acts->dweight_buf, back_acts->dbias_buf, buf1_BCHoWo, grads->cv3_1_w, grads->cv3_1_b, B, C, C_out, H_out, W_out);\n cudaCheck(cudaGetLastError());\n // use acts->ud_x as a buffer for dx\n float* buf1_BCHW = back_acts->buf1_BCHW;\n float* buf2_BCHW = back_acts->buf2_BCHW;\n if (up) {\n upsample_backward1(buf1_BCHW, buf1_BCHoWo, B, C, H, W, block_size);\n cudaCheck(cudaGetLastError());\n upsample_backward1(buf2_BCHW, res_dx, B, C, H, W, block_size);\n cudaCheck(cudaGetLastError());\n } else if (down) {\n avgpool_2d_backward1(buf1_BCHoWo, buf1_BCHW, B, C, H, W, block_size);\n cudaCheck(cudaGetLastError());\n avgpool_2d_backward1(res_dx, buf2_BCHW, B, C, H, W, block_size);\n cudaCheck(cudaGetLastError());\n } else {\n cudaCheck(cudaMemcpy(buf1_BCHW, buf1_BCHoWo, B * C * H * W * sizeof(float), cudaMemcpyDeviceToDevice));\n cudaCheck(cudaGetLastError());\n cudaCheck(cudaMemcpy(buf2_BCHW, res_dx, B * C * H * W * sizeof(float), cudaMemcpyDeviceToDevice));\n cudaCheck(cudaGetLastError());\n }\n // use silu as the dx buffer\n silu_backward(buf1_BCHW, acts->gn1, acts->silu1, B * C * H * W, block_size);\n cudaCheck(cudaGetLastError());\n groupnorm_backward(acts->silu1, acts->input, acts->gn1_mean, acts->gn1_rstd, params->gn1_w, buf1_BCHW, grads->gn1_w, grads->gn1_b, B, C, H, W, gn_n_groups);\n cudaCheck(cudaGetLastError());\n\n add_forward(buf1_BCHW, buf2_BCHW, dx, B * C * H * W, block_size);\n cudaCheck(cudaGetLastError());\n}\n\n// ------------------------------------------------------------------------------------------------\n// self attention block\n// assumes you have Q, K, V\n__global__ void permute_kernel(float* q, float* k, float* v,\n const float* inp,\n int B, int N, int NH, int d) {\n // okay so now, this kernel wants Q,K,V to all be of shape (B, NH, N, d)\n // but instead, we have a single tensor QKV (inp) of shape (B, N, 3, NH, d)\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Q[b][nh_][n][d_] = inp[b][n][0][nh_][d_]\n\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n\n int inp_idx = \\\n (b * N * 3 * NH * d)\n + (n * 3 * NH * d)\n + (0 * NH * d)\n + (nh_ * d)\n + d_;\n\n q[idx] = inp[inp_idx];\n k[idx] = inp[inp_idx + NH * d];\n v[idx] = inp[inp_idx + 2 * (NH * d)];\n }\n}\n\n__global__ void scale_kernel(float* inp, float scale, int B, int NH, int T) {\n // scales the pre-softmax attention scores by scale\n // we don't want causal attention, so we don't set any values to infinity\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * NH * T * T) {\n inp[idx] *= scale;\n }\n}\n\n// warp-level reduction for finding the maximum value\n__device__ float warpReduceMax(float val) {\n for (int offset = 16; offset > 0; offset /= 2) {\n val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));\n }\n return val;\n}\n\n// warp-level reduction for summing values\n__device__ float warpReduceSum(float val) {\n for (int offset = 16; offset > 0; offset /= 2) {\n val += __shfl_down_sync(0xFFFFFFFF, val, offset);\n }\n return val;\n}\n\n__global__ void softmax_forward_kernel4(float* out, const float* inp, int N, int C) {\n // out is (N, C) just like inp. Each row of inp will get softmaxed.\n // same as kernel3, but can handle any block size (multiple of 32)\n // each row of C elements is handled by block_size threads\n // furthermore, each block_size threads get executed in warps of 32 threads\n\n // special reduction operations warpReduceMax/warpReduceSum are used for intra-warp reductions\n // shared memory is used for inter-warp reduction\n extern __shared__ float shared[];\n int idx = blockIdx.x;\n int tid = threadIdx.x;\n int warpId = threadIdx.x / 32; // warp index within a block\n int laneId = threadIdx.x % 32; // thread index within a warp\n\n // the number of warps per block. recall that blockDim.x is block_size\n int warpsPerBlock = blockDim.x / 32;\n\n // shared[] must be allocated to have 2 * warpsPerBlock elements\n // first half for max values, the second half for sum values\n float* maxvals = shared;\n float* sumvals = &shared[warpsPerBlock];\n\n // one row of inp, i.e. inp[idx, :] of shape (C,)\n const float* x = inp + idx * C;\n\n // first, thread coarsening by directly accessing global memory in series\n float maxval = -INFINITY;\n for (int i = tid; i < C; i += blockDim.x) {\n maxval = fmaxf(maxval, x[i]);\n }\n // now within-warp reductions for maxval\n maxval = warpReduceMax(maxval);\n\n // the 0th thread of each warp writes the maxval of that warp to shared memory\n if (laneId == 0) maxvals[warpId] = maxval;\n __syncthreads();\n\n // now the 0th thread reduces the maxvals in shared memory, i.e. across warps\n if (tid == 0) {\n float val = maxvals[tid];\n for (int i = 1; i < warpsPerBlock; i++) {\n val = fmaxf(val, maxvals[i]);\n }\n // store the final max in the first position\n maxvals[0] = val;\n }\n __syncthreads();\n // broadcast the max to all threads\n float offset = maxvals[0];\n\n // compute expf and write the result to global memory\n for (int i = tid; i < C; i += blockDim.x) {\n // subtract max for numerical stability\n out[idx * C + i] = expf(x[i] - offset);\n }\n\n // okay now we calculated exp(x - max(x))\n // step 2: sum all the values and divide by the sum\n\n // thread coarsening for sum\n x = out + idx * C;\n float sumval = 0.0f;\n for (int i = tid; i < C; i += blockDim.x) {\n sumval += x[i];\n }\n // within-warp reduction for sumval\n sumval = warpReduceSum(sumval);\n\n // write sumval to shared memory\n if (laneId == 0) sumvals[warpId] = sumval;\n __syncthreads();\n\n // inter-thread reduction of sum\n if (tid == 0) {\n float val = sumvals[tid];\n for (int i = 1; i < warpsPerBlock; ++i) {\n val += sumvals[i];\n }\n sumvals[0] = val;\n }\n __syncthreads();\n // broadcast the sum to all threads\n float sum = sumvals[0];\n\n // divide the whole row by the sum\n for (int i = tid; i < C; i += blockDim.x) {\n out[idx * C + i] = x[i] / sum;\n }\n}\n\n__global__ void unpermute_kernel(const float* inp, float *out, int B, int N, int NH, int d) {\n // out has shape (B, nh, N, d) but we need to unpermute it to (B, N, nh, d)\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n // out[b][n][nh_][d_] <- inp[b][nh_][n][d_]\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n\n int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_;\n out[other_idx] = inp[idx];\n }\n}\n\n// a simple attention forward function taken directly from llm.c kernel 3\nvoid attention_forward1\n(\n cublasHandle_t cublas_handle,\n float* out, float* qkvr, float* preatt, float* att,\n float* inp,\n int B, int T, int C, int NH,\n const int block_size\n) {\n // note: resure inp as a scratch buffer because it is not used in backward\n // inp is (B, T, 3C) QKV\n // preatt, att are (B, NH, T, T)\n // output is (B, T, C)\n int HS = C / NH; // head size\n\n // permute and separate inp from (B, T, 3, NH, HS) to 3X (B, NH, T, HS)\n float *q, *k, *v;\n q = qkvr + 0 * B * T * C;\n k = qkvr + 1 * B * T * C;\n v = qkvr + 2 * B * T * C;\n int total_threads = B * NH * T * HS;\n int num_blocks = ceil_div(total_threads, block_size);\n permute_kernel<<>>(q, k, v, inp, B, T, NH, HS);\n\n // batched matrix multiply with cuBLAS\n const float alpha = 1.0f;\n const float beta = 0.0f;\n cublasCheck(cublasSgemmStridedBatched(cublas_handle,\n CUBLAS_OP_T, CUBLAS_OP_N,\n T, T, HS,\n &alpha,\n k, HS, T * HS,\n q, HS, T * HS,\n &beta,\n preatt, T, T * T,\n B * NH));\n\n // multiply all elements of preatt elementwise by scale\n float scale = 1.0f / sqrtf(HS);\n total_threads = B * NH * T * T;\n num_blocks = ceil_div(total_threads, block_size);\n scale_kernel<<>>(preatt, scale, B, NH, T);\n\n // softmax. preatt is (B, NH, T, T) but we view it as (B * NH * T, T) and use the softmax kernel\n int softmax_block_size = 256;\n int grid_size = B * NH * T;\n size_t shared_mem_size = 2 * softmax_block_size / 32 * sizeof(float);\n softmax_forward_kernel4<<>>(att, preatt, B * NH * T, T);\n\n // new approach: first cuBLAS another batched matmul\n // y = att @ v # (B, nh, T, T) @ (B, nh, T, hs) -> (B, nh, T, hs)\n float* vaccum = inp;\n cublasCheck(cublasSgemmStridedBatched(cublas_handle,\n CUBLAS_OP_N, CUBLAS_OP_N,\n HS, T, T,\n &alpha,\n v, HS, T * HS,\n att, T, T * T,\n &beta,\n vaccum, HS, T * HS,\n B * NH));\n\n // now unpermute\n // y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side\n num_blocks = ceil_div(B * T * C, block_size);\n unpermute_kernel<<>>(vaccum, out, B, T, NH, HS);\n}\n\n__global__ void permute_kernel_backward(float* dinp,\n const float* dq, const float* dk, const float* dv,\n int B, int N, int NH, int d) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n\n int inp_idx = (b * N * 3 * NH * d) + (n * 3 * NH * d) + (0 * NH * d) + (nh_ * d) + d_;\n dinp[inp_idx] = dq[idx];\n dinp[inp_idx + NH * d] = dk[idx];\n dinp[inp_idx + 2 * (NH * d)] = dv[idx];\n }\n}\n\n__global__ void unpermute_kernel_backward(float* dinp, const float *dout, int B, int N, int NH, int d) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_;\n dinp[idx] = dout[other_idx];\n }\n}\n\n// slight edit from the softmax_autoregressive_backward_kernel from train_gpt2_fp32\n// our attention is no longer autoregressive\n__global__ void softmax_backward_kernel(float* dpreatt, const float* datt, const float* att,\n int B, int T, int C, float scale) {\n constexpr const int BlockSize = 256;\n constexpr int T_per_block = 4;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n __shared__ float block_acc[32];\n\n int idx = blockIdx.y;\n // go through blocks in reverse order, so the slowest block starts first\n int t0 = T - 1 - T_per_block*blockIdx.x;\n\n att += idx * T * T;\n datt += idx * T * T;\n dpreatt += idx * T * T;\n\n if (warp.meta_group_rank() == 0) {\n block_acc[warp.thread_rank()] = 0;\n }\n\n for(int to = 0; to < T_per_block; ++to) {\n int t = t0 - to;\n if(t < 0) return;\n const float* att_bth = att + t * T;\n const float* datt_bth = datt + t * T;\n float* dpreatt_bth = dpreatt + t * T;\n\n float local_sum = 0.0f;\n for (int t2 = block.thread_rank(); t2 < T; t2 += BlockSize) {\n local_sum += att_bth[t2] * datt_bth[t2];\n }\n\n block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus{});\n block.sync();\n local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus{});\n\n for (int t3 = block.thread_rank(); t3 < T; t3 += BlockSize) {\n // don't touch the cache. Some parts will still be here from the previous loop, and\n // we want to exploit those.\n float acc = __ldcs(att_bth + t3) * (__ldcs(datt_bth + t3) - local_sum);\n __stcs(dpreatt_bth + t3, scale * acc);\n }\n }\n}\n\n// the sequence of transformations in this compound op is:\n// inp (B,T,3C) -> qkvr (B,T,3C) -> preatt (B,NH,T,T) -> att (B,NH,T,T) -> vaccum (B,T,C) -> out (B,T,C)\nvoid attention_backward\n(\n cublasHandle_t cublas_handle,\n float* dinp, float* dqkvr, float* dpreatt, float* datt, float* scratch,\n const float* dout,\n const float* qkvr, const float* att,\n int B, int T, int C, int NH\n) {\n const int block_size = 256;\n int HS = C / NH; // head size\n const float one = 1.0f;\n const float zero = 0.0f; // note beta = 1.0f so that we accumulate gradients (+=)\n // unpack convenience pointers into q, k, v\n const float *q, *k, *v;\n q = qkvr + 0 * B * T * C;\n k = qkvr + 1 * B * T * C;\n v = qkvr + 2 * B * T * C;\n float *dq, *dk, *dv;\n dq = dqkvr + 0 * B * T * C;\n dk = dqkvr + 1 * B * T * C;\n dv = dqkvr + 2 * B * T * C;\n // backward through the unpermute operation\n int num_blocks = ceil_div(B * T * C, block_size);\n unpermute_kernel_backward<<>>(scratch, dout, B, T, NH, HS);\n cudaCheck(cudaGetLastError());\n // backward into datt\n cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, &one, v, HS, T * HS, scratch, HS, T * HS, &zero, datt, T, T * T, B * NH));\n // backward into dv\n cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &one, scratch, HS, T * HS, att, T, T * T, &zero, dv, HS, T * HS, B * NH));\n // backward into preatt\n int hs = C / NH; // head size\n float scale = 1.0f / sqrtf(hs);\n softmax_backward_kernel<<>>(dpreatt, datt, att, B, T, C, scale);\n cudaCheck(cudaGetLastError());\n // backward into q\n cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, &one, k, HS, T * HS, dpreatt, T, T * T, &zero, dq, HS, T * HS, B * NH));\n // backward into k\n cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &one, q, HS, T * HS, dpreatt, T, T * T, &zero, dk, HS, T * HS, B * NH));\n // backward into inp\n num_blocks = ceil_div(B * NH * T * HS, block_size);\n permute_kernel_backward<<>>(dinp, dq, dk, dv, B, T, NH, HS);\n cudaCheck(cudaGetLastError());\n}\n\n// ------------------------------------------------------------------------------------------------\n// Attention layer, including groupnorm, qkv projection, and residual stream\n\n#define NUM_ATT_PARAM_TENSORS 6\ntypedef struct {\n float* gn_w; // C\n float* gn_b; // C\n float* qkv_w; // 3xC, C\n float* qkv_b; // 3xC\n float* proj_w; // C, C\n float* proj_b; // C\n size_t param_sizes[NUM_ATT_PARAM_TENSORS];\n size_t n_params;\n} AttentionParams;\n\n#define NUM_ATT_ACT_TENSORS 12 // not counting input, which will reuse output memory from previous layer\ntypedef struct {\n float* gn; // B, C, H, W\n float* gn_mean; // B, n_groups\n float* gn_rstd; // B, n_groups\n float* perm1; // B, HxW, C\n float* qkv1; // B, HxW, 3xC\n float* qkv2; // B, HxW, 3xC\n float* preatt; // B, NH, T, T\n float* att; // B, NH, T, T\n float* att_out; // B, HxW, C\n float* proj; // B, HxW, C\n float* perm2; // B, C, H, W\n float* add; // B, C, H, W\n size_t act_sizes[NUM_ATT_ACT_TENSORS];\n size_t n_acts;\n float* input; // B, C, H, W\n} AttentionActs;\n\n#define NUM_ATT_BACKWARD_ACTS_TENSORS 7 // dinp will use input memory\ntypedef struct {\n float* buf1_BCHW; // B, C, H, W\n float* buf2_BCHW; // B, C, H, W\n float* buf_B3CHW; // B, 3xC, H, W\n float* dqkvr; // B, HxW, 3xC\n float* dpreatt; // B, NH, T, T\n float* datt; // B, NH, T, T\n float* dout; // B, C, H, W\n float* dinp; // B, C, H, W\n size_t back_sizes[NUM_ATT_BACKWARD_ACTS_TENSORS];\n size_t n_backs;\n} AttentionBackwardActs;\n\nvoid attention_block_count_params(AttentionParams* params, int C) {\n size_t *param_sizes = params->param_sizes;\n param_sizes[0] = C; // gn_w\n param_sizes[1] = C; // gn_b\n param_sizes[2] = 3 * C * C; // qkv_w\n param_sizes[3] = 3 * C; // qkv_b\n param_sizes[4] = C * C; // proj_w\n param_sizes[5] = C; // proj_b\n\n size_t num_parameters = 0;\n for (int i = 0; i < NUM_ATT_PARAM_TENSORS; i++) {\n num_parameters += param_sizes[i];\n }\n params->n_params = num_parameters;\n}\nvoid attention_block_count_acts(\n AttentionActs* acts,\n int B, int C, int H, int W, int gn_n_groups, int HS\n) {\n int C3 = 3 * C;\n int T = H * W;\n int NH = C / HS;\n size_t *act_sizes = acts->act_sizes;\n act_sizes[0] = B * C * H * W; // gn\n act_sizes[1] = B * gn_n_groups; // gn_mean\n act_sizes[2] = B * gn_n_groups; // gn_rstd\n act_sizes[3] = B * T * C; // perm1\n act_sizes[4] = B * T * C3; // qkv1\n act_sizes[5] = B * T * C3; // qkv2\n act_sizes[6] = B * NH * T * T; // preatt\n act_sizes[7] = B * NH * T * T; // att\n act_sizes[8] = B * T * C; // att_out\n act_sizes[9] = B * T * C; // proj\n act_sizes[10] = B * C * H * W; // perm2\n act_sizes[11] = B * C * H * W; // add\n\n size_t num_acts_params = 0;\n for (int i = 0; i < NUM_ATT_ACT_TENSORS; i++) {\n num_acts_params += act_sizes[i];\n }\n acts->n_acts = num_acts_params;\n}\n\nvoid set_attention_params_pointers(\n AttentionParams* params,\n float* params_memory\n) {\n float** ptrs[] = {\n ¶ms->gn_w,\n ¶ms->gn_b,\n ¶ms->qkv_w,\n ¶ms->qkv_b,\n ¶ms->proj_w,\n ¶ms->proj_b\n };\n for (int i = 0; i < NUM_ATT_PARAM_TENSORS; i++) {\n *(ptrs[i]) = params_memory;\n params_memory += params->param_sizes[i];\n }\n}\n\nvoid set_attention_acts_pointers(\n AttentionActs* acts,\n float* acts_memory\n) {\n float** ptrs[] = {\n &acts->gn,\n &acts->gn_mean,\n &acts->gn_rstd,\n &acts->perm1,\n &acts->qkv1,\n &acts->qkv2,\n &acts->preatt,\n &acts->att,\n &acts->att_out,\n &acts->proj,\n &acts->perm2,\n &acts->add\n };\n for (int i = 0; i < NUM_ATT_ACT_TENSORS; i++) {\n *(ptrs[i]) = acts_memory;\n acts_memory += acts->act_sizes[i];\n }\n}\n\n// ----------------------------------------------------\n// kernels\n\n// this is an opportunity for improvement:\n// rewrite the attention layer so we don't need to permute the input\n__global__ void attention_permute_kernel(\n const float* inp, float* out,\n int B, int C, int H, int W\n) {\n // from B, C, H, W to B, HxW, C\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n if (idx >= B * C * img_size) { return; }\n\n int b = idx / (img_size * C);\n int c = idx / img_size % C;\n int h = idx / W % H;\n int w = idx % W;\n\n out[b * C * img_size + h * C * W + w * C + c] = inp[idx];\n}\n\nvoid attention_permute(\n const float* inp, float* out,\n int B, int C, int H, int W, int block_size\n) {\n int num_threads = B * C * H * W;\n attention_permute_kernel<<>>(inp, out, B, C, H, W);\n}\n\n__global__ void attention_unpermute_kernel(\n const float* inp, float* out,\n int B, int C, int H, int W\n) {\n // from B, HxW, C to B, C, H, W\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n if (idx >= B * C * img_size) { return; }\n\n int b = idx / (img_size * C);\n int c = idx / img_size % C;\n int h = idx / W % H;\n int w = idx % W;\n\n out[idx] = inp[b * C * img_size + h * C * W + w * C + c];\n}\nvoid attention_unpermute(\n const float* inp, float* out,\n int B, int C, int H, int W, int block_size\n) {\n int num_threads = B * C * H * W;\n attention_unpermute_kernel<<>>(inp, out, B, C, H, W);\n}\n\nvoid attention_block_forward\n(\n cublasHandle_t cublas_handle,\n int B, int C, int H, int W, int HS, int gn_n_groups, int block_size,\n AttentionParams* params, AttentionActs* acts\n) {\n // apply groupnorm\n groupnorm_forward(acts->input, params->gn_w, params->gn_b, acts->gn, acts->gn_mean, acts->gn_rstd, B, C, H, W, gn_n_groups);\n // permute from B, C, H, W to B, HxW, C\n attention_permute(acts->gn, acts->perm1, B, C, H, W, block_size);\n // get qkv\n matmul_forward2(cublas_handle, acts->qkv1, acts->perm1, params->qkv_w, params->qkv_b, B * H * W, C, 3 * C, block_size);\n // apply attention\n attention_forward1(cublas_handle, acts->att_out, acts->qkv2, acts->preatt, acts->att, acts->qkv1, B, H * W, C, C / HS, block_size);\n // attention projection\n matmul_forward2(cublas_handle, acts->proj, acts->att_out, params->proj_w, params->proj_b, B * H * W, C, C, block_size);\n // permute from B, HxW, C to B, C, H, W\n attention_unpermute(acts->proj, acts->perm2, B, C, H, W, block_size);\n // add back residual stream\n add_forward(acts->input, acts->perm2, acts->add, B * C * H * W, block_size);\n}\n\nvoid attention_block_backward\n(\n cublasHandle_t cublas_handle,\n int B, int C, int H, int W, int HS, int gn_n_groups, int block_size,\n AttentionParams* params, AttentionActs* acts,\n AttentionBackwardActs* back_acts, AttentionParams* grads\n) {\n // permute from B, C, H, W to B, HxW, C\n attention_permute(back_acts->dout, back_acts->buf1_BCHW, B, C, H, W, block_size);\n // backward through projection\n matmul_backward1(cublas_handle, back_acts->buf2_BCHW, grads->proj_w, grads->proj_b, back_acts->buf1_BCHW, acts->att_out, params->proj_w, B * H * W, C, C);\n // backward through attention\n attention_backward(cublas_handle, back_acts->buf_B3CHW, back_acts->dqkvr, back_acts->dpreatt, back_acts->datt, back_acts->buf1_BCHW, back_acts->buf2_BCHW, acts->qkv2, acts->att, B, H * W, C, C / HS);\n // backward through qkv\n matmul_backward1(cublas_handle, back_acts->buf1_BCHW, grads->qkv_w, grads->qkv_b, back_acts->buf_B3CHW, acts->perm1, params->qkv_w, B * H * W, C, 3 * C);\n // permute from B, HxW, C to B, C, H, W\n attention_unpermute(back_acts->buf1_BCHW, back_acts->buf2_BCHW, B, C, H, W, block_size);\n // backward through groupnorm\n groupnorm_backward(back_acts->buf2_BCHW, acts->input, acts->gn_mean, acts->gn_rstd, params->gn_w, back_acts->buf1_BCHW, grads->gn_w, grads->gn_b, B, C, H, W, gn_n_groups);\n // add residual gradient\n add_forward(back_acts->buf1_BCHW, back_acts->dout, back_acts->dinp, B * C * H * W, block_size);\n}\n\n// ------------------------------------------------------------------------------------------------\n// MSE loss\n// this will only be called with one block\n__global__ void mse_forward_kernel(\n const float* inp, const float* y, float* loss, int N\n) {\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n\n __shared__ float shared_sum[32]; // max 32 warps\n int num_warps = blockDim.x / 32;\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n\n float thread_sum = 0.0f;\n for (int i = threadIdx.x; i < N; i += blockDim.x) {\n float diff = inp[i] - y[i];\n thread_sum += diff * diff;\n }\n\n // warp reduce\n float warp_sum = cg::reduce(warp, thread_sum, cg::plus{});\n shared_sum[warp_id] = warp_sum;\n __syncthreads();\n \n warp_sum = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;\n float block_sum = cg::reduce(warp, warp_sum, cg::plus{});\n\n if (threadIdx.x == 0) {\n loss[0] = block_sum / N;\n }\n}\n\nvoid mse_forward(const float* inp, const float* y, float* loss, int N, int block_size) {\n mse_forward_kernel<<<1, block_size>>>(inp, y, loss, N);\n}\n\n// not calculating dy in backward\n// note that backward for MSE doesn't require the loss value\n__global__ void mse_backward_kernel(\n const float* inp, const float* y, float* dinp, int N\n) {\n int idx = threadIdx.x + blockIdx.x * blockDim.x;\n if (idx < N) {\n float x = inp[idx];\n dinp[idx] = 2.0f * (x - y[idx]) / N;\n }\n}\n\nvoid mse_backward(const float* inp, const float* y, float* dinp, int N, int block_size) {\n int n_blocks = ceil_div(N, block_size);\n mse_backward_kernel<<>>(inp, y, dinp, N);\n}\n\n// ------------------------------------------------------------------------------------------------\n// DataLoader\n\ntypedef struct {\n int B;\n int C_in;\n int H;\n int W;\n int img_size;\n int n_imgs;\n FILE* img_file;\n long cur_position;\n long file_size;\n long header_offset;\n long num_batches;\n float* input;\n} DataLoader;\n\nvoid dataloader_init(DataLoader* loader, const char* filename, int B) {\n loader->B = B;\n\n loader->img_file = fopenCheck(filename, \"rb\");\n int data_header[256];\n loader->header_offset = 256 * sizeof(int);\n freadCheck(data_header, sizeof(int), 256, loader->img_file);\n if (data_header[0] != 20240620) {\n fprintf(stderr, \"Invalid magic data file\\n\");\n exit(1);\n }\n loader->n_imgs = data_header[1];\n loader->C_in = data_header[2];\n loader->H = data_header[3];\n loader->W = data_header[4];\n loader->img_size = loader->C_in * loader->H * loader->W;\n\n // determine file size\n fseekCheck(loader->img_file, 0, SEEK_SET);\n fseekCheck(loader->img_file, 0, SEEK_END);\n loader->file_size = ftell(loader->img_file);\n fseekCheck(loader->img_file, 0, SEEK_SET);\n assert((loader->file_size - loader->header_offset) / (loader->img_size * sizeof(float)) == loader->n_imgs);\n \n loader->cur_position = loader->header_offset;\n loader->num_batches = loader->n_imgs / B;\n\n // allocate cuda CPU pinned memory for inputs\n cudaMallocHost((void**)&loader->input, B * loader->img_size * sizeof(float));\n}\n\nvoid dataloader_reset(DataLoader* loader) {\n loader->cur_position = loader->header_offset;\n}\n\nvoid dataloader_next_batch(DataLoader* loader) {\n int B = loader->B;\n int img_size = loader->img_size;\n if (loader->cur_position + (B * img_size * sizeof(float)) > loader->file_size) {\n dataloader_reset(loader);\n }\n fseekCheck(loader->img_file, loader->cur_position, SEEK_SET);\n freadCheck(loader->input, sizeof(float), B * img_size, loader->img_file);\n loader->cur_position += B * img_size * sizeof(float);\n}\n\nvoid dataloader_free(DataLoader* loader) {\n fcloseCheck(loader->img_file);\n cudaFreeHost(loader->input);\n}\n\n\n// ------------------------------------------------------------------------------------------------\n// Diffusion\ntypedef struct {\n int B;\n int img_size;\n int max_period;\n float* betas; // (max_period,)\n float* sqrt_alphas_cumprod; // (max_period,)\n float* sqrt_one_minus_alphas_cumprod; // (max_period,)\n int block_size;\n curandState* curand_states;\n} GaussianDiffusion;\n\n__global__ void init_curand_states(curandState *state, unsigned long seed, int N) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < N) {\n curand_init(seed, idx, 0, &state[idx]);\n }\n}\n\nvoid diffusion_init(\n GaussianDiffusion* diffusion, int B, int C_in, int H, int W, int max_period, int block_size = 512,\n unsigned long seed = 0, float beta_start = 1e-4, float beta_end = 0.02f\n) {\n diffusion->B = B;\n diffusion->max_period = max_period;\n diffusion->block_size = block_size;\n diffusion->img_size = C_in * H * W;\n\n float scale = 1000.0f / max_period;\n beta_start *= scale;\n beta_end *= scale;\n\n // calculate betas and alphas on host then move to device\n float* betas = (float*)mallocCheck(max_period * sizeof(float));\n float* sqrt_alphas_cumprod = (float*)mallocCheck(max_period * sizeof(float));\n float* sqrt_one_minus_alphas_cumprod = (float*)mallocCheck(max_period * sizeof(float));\n\n float alpha_cumprod = 1.0f;\n for (int t = 1; t <= max_period; t++) {\n // linear interpolate for beta\n betas[t - 1] = (beta_start * (max_period - t) + beta_end * (t - 1)) / (max_period - 1);\n alpha_cumprod *= 1.0f - betas[t - 1];\n sqrt_alphas_cumprod[t - 1] = sqrtf(alpha_cumprod);\n sqrt_one_minus_alphas_cumprod[t - 1] = sqrtf(1.0f - alpha_cumprod);\n }\n\n float *d_betas, *d_sqrt_alphas_cumprod, *d_sqrt_one_minus_alphas_cumprod;\n cudaCheck(cudaMalloc(&d_betas, max_period * sizeof(float)));\n cudaCheck(cudaMalloc(&d_sqrt_alphas_cumprod, max_period * sizeof(float)));\n cudaCheck(cudaMalloc(&d_sqrt_one_minus_alphas_cumprod, max_period * sizeof(float)));\n\n cudaCheck(cudaMemcpy(d_betas, betas, max_period * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_sqrt_alphas_cumprod, sqrt_alphas_cumprod, max_period * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_sqrt_one_minus_alphas_cumprod, sqrt_one_minus_alphas_cumprod, max_period * sizeof(float), cudaMemcpyHostToDevice));\n\n free(betas);\n free(sqrt_alphas_cumprod);\n free(sqrt_one_minus_alphas_cumprod);\n\n diffusion->betas = d_betas;\n diffusion->sqrt_alphas_cumprod = d_sqrt_alphas_cumprod;\n diffusion->sqrt_one_minus_alphas_cumprod = d_sqrt_one_minus_alphas_cumprod;\n\n // init curand states for sampling noise of the size of one batch of inputs\n if (diffusion->curand_states == nullptr) {\n printf(\"initializing curand states\\n\");\n int total_states = B * C_in * H * W;\n int n_blocks = ceil_div(total_states, block_size);\n cudaCheck(cudaMalloc(&diffusion->curand_states, total_states * sizeof(curandState)));\n init_curand_states<<>>(diffusion->curand_states, seed, total_states);\n }\n}\n\n// fill BxN noise array with standard gaussians\n__global__ void diffusion_draw_normal_kernel(\n float* noise, int B, int N, curandState* curand_states\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * N) {\n noise[idx] = curand_normal(&curand_states[idx]);\n }\n}\n\nvoid diffusion_draw_normal(\n GaussianDiffusion *diffusion,\n float* noise\n) {\n int B = diffusion->B;\n int img_size = diffusion->img_size;\n int total_states = B * img_size;\n int n_blocks = ceil_div(total_states, diffusion->block_size);\n diffusion_draw_normal_kernel<<block_size>>>(noise, B, img_size, diffusion->curand_states);\n}\n\n// given input x of shape (B, N) and timesteps of shape (B,)\n// replace each x entry by the result of flowing forward by diffusion for time t\n// with randomness coming from noise\n__global__ void diffusion_forward_by_t_kernel(\n float* x, float* timesteps, float* noise,\n float* sqrt_alphas_cumprod, float* sqrt_one_minus_alphas_cumprod,\n int B, int N, curandState* curand_states\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * N) {\n int b = idx / N;\n float x_val = x[idx];\n float noise_val = noise[idx];\n int t = (int)timesteps[b];\n float sqrt_alpha_t = sqrt_alphas_cumprod[t];\n float sqrt_one_minus_alpha_t = sqrt_one_minus_alphas_cumprod[t];\n x[idx] = x_val * sqrt_alpha_t + noise_val * sqrt_one_minus_alpha_t;\n }\n}\n\nvoid diffusion_forward_by_t(\n GaussianDiffusion* diffusion,\n float* x, float* timesteps, float* noise\n) {\n int B = diffusion->B;\n int img_size = diffusion->img_size;\n int block_size = diffusion->block_size;\n int total_states = B * img_size;\n int n_blocks = ceil_div(total_states, block_size);\n diffusion_forward_by_t_kernel<<>>(\n x, timesteps, noise, diffusion->sqrt_alphas_cumprod, diffusion->sqrt_one_minus_alphas_cumprod, B, img_size, diffusion->curand_states\n );\n}\n\nvoid diffusion_free(GaussianDiffusion* diffusion) {\n cudaCheck(cudaFree(diffusion->betas));\n cudaCheck(cudaFree(diffusion->sqrt_alphas_cumprod));\n cudaCheck(cudaFree(diffusion->sqrt_one_minus_alphas_cumprod));\n cudaCheck(cudaFree(diffusion->curand_states));\n}\n\n// one block with B threads\n__global__ void sample_timestep_kernel(\n float* timesteps, int B,\n curandState* curand_states, int max_period\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B) {\n int t = curand(&curand_states[idx]) % max_period;\n timesteps[idx] = (float)t;\n }\n}\n\n// samples B timesteps in range [0, max_period), convert them to floats and\n// store in timesteps\nvoid sample_timesteps(GaussianDiffusion *diffusion, float* timesteps) {\n sample_timestep_kernel<<<1, diffusion->B>>>(timesteps, diffusion->B, diffusion->curand_states, diffusion->max_period);\n}\n\n// ------------------------------------------------------------------------------------------------\n// timestep embedding states\ntypedef struct {\n float* freqs;\n int half_dim;\n int B;\n int max_period;\n} TimestepEmbedding;\n\n__global__ void fill_freqs_kernel(float* freqs, int half_dim, int max_period) {\n int idx = threadIdx.x + blockIdx.x * blockDim.x;\n if (idx < half_dim) {\n float log_period = logf(max_period);\n float logit = -log_period * idx / half_dim;\n freqs[idx] = expf(logit);\n }\n}\n\nvoid init_timestep_embedding(TimestepEmbedding* emb, int dim, int B, int max_period) {\n assert(dim % 2 == 0);\n int half = dim / 2;\n emb->half_dim = half;\n emb->max_period = max_period;\n emb->B = B;\n cudaCheck(cudaMalloc(&emb->freqs, half * sizeof(float)));\n // fill freq values\n int block_size = 512;\n int n_blocks = ceil_div(half, block_size);\n fill_freqs_kernel<<>>(emb->freqs, half, max_period);\n}\n\nvoid free_timestep_embedding(TimestepEmbedding* emb) {\n cudaCheck(cudaFree(emb->freqs));\n}\n\n__global__ void fill_embeddings_kernel(\n const float* timesteps, const float* freqs, float* embeddings, int half, int B\n) {\n int idx = threadIdx.x + blockIdx.x * blockDim.x;\n if (idx < B * half) {\n int i = idx / half;\n int j = idx % half;\n float arg = timesteps[i] * freqs[j];\n embeddings[i * 2 * half + j] = cosf(arg);\n embeddings[i * 2 * half + j + half] = sinf(arg);\n }\n}\n\n// timesteps is a (B,) array of timesteps that are integers but saved as floats\n// out is a (B, dim) array\nvoid get_timestep_embeddings(\n TimestepEmbedding* emb, const float* timesteps, float* out\n) {\n int block_size = 512;\n int n_blocks = ceil_div(emb->half_dim * emb->B, block_size);\n fill_embeddings_kernel<<>>(timesteps, emb->freqs, out, emb->half_dim, emb->B);\n cudaCheck(cudaGetLastError());\n}\n\n// ------------------------------------------------------------------------------------------------\n// Unet init\n\ntypedef struct {\n int C_in;\n int C_model;\n int C_time;\n int C_out;\n int B;\n int H;\n int W;\n int gn_n_groups;\n int HS;\n int block_size;\n int n_layers;\n int n_levels;\n int n_res;\n int att_start_level;\n int* channel_mult;\n int max_period;\n} UnetConfig;\n\n// we just need a single set of buffers that are shared for all the backward pass activations\n#define NUM_SKIP_CONNECTIONS 12 // (n_res + 1) * n_levels\n#define NUM_BACK_BUFS 10\n#define UNET_NUM_BACK_ACTS 22\ntypedef struct {\n float* dinp_buf1; // used to store the input gradient, which will be reused as dout for the previous layer\n float* dinp_buf2;\n float* buf1_BCHW; // NOTE: first 4 buffers will have max B * C * H * W size\n float* buf2_BCHW;\n float* buf3_BCHW;\n float* buf4_BCHW;\n float* buf1_B3CHW; // NOTE: last 2 buffers will have max B * 3 * C * H * W size for attention layers\n float* buf2_B3CHW; // so max of H * W is only over dimensions that have attention\n float* demb_buf1; // used to store cumulative gradients w.r.t. timestep embeddings\n float* demb_buf2;\n float* skip_dinps[NUM_SKIP_CONNECTIONS];\n} UnetBackActs;\n\n// This is used as a block to store all the params / grads / acts for all the layers\ntypedef struct {\n float* mem_ptr;\n size_t total_size;\n int n_layers;\n void** layer_ptrs;\n size_t* layer_sizes;\n} UnetBlock;\n\ntypedef struct {\n float* back_mem;\n UnetBackActs back_acts;\n size_t back_sizes[UNET_NUM_BACK_ACTS];\n size_t n_backs;\n} UnetBackBlock;\n\ntypedef struct {\n UnetConfig config;\n float* input;\n float* timestep_emb; // this is the already computed embeddings of shape (B, C_model)\n UnetBlock params;\n UnetBlock grads;\n UnetBlock acts;\n UnetBackBlock back;\n float* skip_acts[NUM_SKIP_CONNECTIONS]; // pointers to skip connetion activations\n // the next two are for debug purposes only, just to keep track of the buffer that contains the grad w.r.t. inputs\n float* dinp;\n // model output\n float* output;\n // targets and losses\n float* target;\n float mean_loss;\n float* gpu_loss;\n // buffers for AdamW\n float* m_memory;\n float* v_memory;\n TimestepEmbedding time_emb_state;\n} Unet;\n\nvoid unet_init(Unet* unet) {\n unet->input = nullptr;\n unet->timestep_emb = nullptr;\n unet->params.mem_ptr = nullptr;\n unet->grads.mem_ptr = nullptr;\n unet->acts.mem_ptr = nullptr;\n unet->back.back_mem = nullptr;\n unet->dinp = nullptr;\n unet->output = nullptr;\n unet->target = nullptr;\n unet->mean_loss = -1.0f;\n unet->gpu_loss = nullptr;\n unet->m_memory = nullptr;\n unet->v_memory = nullptr;\n\n init_timestep_embedding(&unet->time_emb_state, unet->config.C_model, unet->config.B, unet->config.max_period);\n}\n\n\ntemplate \nvoid free_not_null(T** ptr) {\n if (*ptr != nullptr) {\n free(*ptr);\n *ptr = nullptr;\n }\n}\n\nvoid free_unet_block(UnetBlock* params) {\n free_not_null(¶ms->layer_sizes);\n for (int l = 0; l < params->n_layers; l++) {\n free_not_null(&(params->layer_ptrs[l]));\n }\n free_not_null(¶ms->layer_ptrs);\n cudaCheck(cudaFree(params->mem_ptr));\n}\n\nvoid free_unet(Unet* unet) {\n //free(unet->config.channel_mult);\n\n cudaCheck(cudaFree(unet->input));\n cudaCheck(cudaFree(unet->timestep_emb));\n cudaCheck(cudaFree(unet->gpu_loss));\n cudaCheck(cudaFree(unet->m_memory));\n cudaCheck(cudaFree(unet->v_memory));\n\n free_unet_block(&(unet->params));\n free_unet_block(&(unet->grads));\n free_unet_block(&(unet->acts));\n\n cudaCheck(cudaFree(unet->back.back_mem));\n free_timestep_embedding(&unet->time_emb_state);\n}\n\nvoid unet_count_layers(UnetConfig* config) {\n int n_layers = 0;\n int n_res = config->n_res;\n int n_levels = config->n_levels;\n int att_start_level = config->att_start_level;\n\n // timestep emb\n n_layers += 3;\n\n // initial conv\n n_layers += 1;\n\n // add downsample resblock layers\n n_layers += n_res * n_levels;\n\n // add downsample attention\n n_layers += n_res * (n_levels - att_start_level);\n\n // add avgpool layers\n n_layers += n_levels - 1;\n\n // add middle layers\n n_layers += 3;\n\n // add upsample concat layers\n n_layers += (n_res + 1) * n_levels;\n\n // add upsample resblock layers\n n_layers += (n_res + 1) * n_levels;\n\n // add upsample attention layers\n n_layers += (n_res + 1) * (n_levels - att_start_level);\n\n // add upsample layers\n n_layers += n_levels - 1;\n\n // output layers\n n_layers += 3;\n\n config->n_layers = n_layers;\n}\n\n// ------------------------------------------------------------------------------------------------\n// allocate memory for structs with info for each layer and count params\n\nint make_linear_params_ptrs(\n UnetBlock* params,\n int layer,\n int C_in, int C_out\n) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(LinearParams));\n params->layer_sizes[layer] = linear_count_params(C_in, C_out);\n return layer + 1;\n}\n\nint make_silu_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_convk3_params_ptrs(\n UnetBlock* params, int layer,\n int C_in, int OC, int B, int H, int W\n) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ConvK3Params));\n params->layer_sizes[layer] = convk3_count_params(C_in, OC);\n return layer + 1;\n}\n\nint make_resblock_params_ptrs(\n UnetBlock* params, int layer,\n int C, int C_emb, int C_out, int B, int H, int W, int gn_n_groups\n) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ResBlockParameters));\n ResBlockParameters* resblock_params = (ResBlockParameters*)params->layer_ptrs[layer];\n resblock_count_params(resblock_params, C, C_emb, C_out, B, H, W, 0, 0, gn_n_groups);\n params->layer_sizes[layer] = resblock_params->n_params;\n return layer + 1;\n}\n\nint make_avgpool_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_attention_block_params_ptrs(UnetBlock* params, int layer, int C) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(AttentionParams));\n AttentionParams* attention_params = (AttentionParams*)params->layer_ptrs[layer];\n attention_block_count_params(attention_params, C);\n params->layer_sizes[layer] = attention_params->n_params;\n return layer + 1;\n}\n\nint make_concat_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_upsample_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_groupnorm_params_ptrs(UnetBlock* params, int layer, int C) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(GroupNormParams));\n params->layer_sizes[layer] = C * 2;\n return layer + 1;\n}\n\nvoid unet_setup_and_count_params(UnetConfig* config, UnetBlock* params) {\n assert (config->n_layers > 0);\n\n params->n_layers = config->n_layers;\n params->layer_ptrs = (void**)mallocCheck(config->n_layers * sizeof(void*));\n params->layer_sizes = (size_t*)mallocCheck(config->n_layers * sizeof(size_t));\n\n int B = config->B;\n int C_in = config->C_in;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int C_out = config->C_out;\n int H = config->H;\n int W = config->W;\n int gn_n_groups = config->gn_n_groups;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n int layer = 0;\n layer = make_linear_params_ptrs(params, layer, C_model, C_time);\n layer = make_silu_params_ptrs(params, layer);\n layer = make_linear_params_ptrs(params, layer, C_time, C_time);\n int OC = C_model;\n layer = make_convk3_params_ptrs(params, layer, C_in, OC, B, H, W);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n layer = make_resblock_params_ptrs(params, layer, C, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_params_ptrs(params, layer, C);\n }\n }\n if (l < n_levels - 1) {\n layer = make_avgpool_params_ptrs(params, layer);\n H /= 2;\n W /= 2;\n }\n }\n\n // middle layers\n layer = make_resblock_params_ptrs(params, layer, C, C_time, C, B, H, W, gn_n_groups);\n layer = make_attention_block_params_ptrs(params, layer, C);\n layer = make_resblock_params_ptrs(params, layer, C, C_time, C, B, H, W, gn_n_groups);\n\n // upwards: concat then resblock\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = C_model * mult;\n for (int r = 0; r < n_res + 1; r++) {\n if (r == n_res && l > 0) {\n int prev_mult = config->channel_mult[l - 1];\n C_skip = C_model * prev_mult;\n }\n OC = C_model * mult;\n layer = make_concat_params_ptrs(params, layer);\n layer = make_resblock_params_ptrs(params, layer, C + C_skip, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_params_ptrs(params, layer, C);\n }\n }\n if (l > 0) {\n layer = make_upsample_params_ptrs(params, layer);\n H *= 2;\n W *= 2;\n }\n }\n\n layer = make_groupnorm_params_ptrs(params, layer, C);\n layer = make_silu_params_ptrs(params, layer);\n layer = make_convk3_params_ptrs(params, layer, C, C_out, B, H, W);\n\n // count total size\n size_t total_size = 0;\n for (int i = 0; i < config->n_layers; i++) {\n total_size += params->layer_sizes[i];\n }\n params->total_size = total_size;\n}\n\n// ------------------------------------------------------------------------------------------------\n// setup Unet activations pointers\nint make_linear_acts_ptrs(\n UnetBlock* acts,\n int layer,\n int B, int C_out\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(LinearActs));\n // acts buffers only hold outputs\n acts->layer_sizes[layer] = B * C_out;\n return layer + 1;\n}\n\nint make_silu_acts_ptrs(UnetBlock* acts, int layer, int N) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(SiluActs));\n acts->layer_sizes[layer] = N;\n return layer + 1;\n}\n\nint make_convk3_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int OC, int H, int W\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ConvK3Acts));\n acts->layer_sizes[layer] = B * OC * H * W;\n return layer + 1;\n}\n\nint make_resblock_acts_ptrs(\n UnetBlock* acts, int layer,\n int C, int C_emb, int C_out, int B, int H, int W, int gn_n_groups\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ResBlockActivations));\n ResBlockActivations* resblock_acts = (ResBlockActivations*)acts->layer_ptrs[layer];\n resblock_count_acts(resblock_acts, C, C_emb, C_out, B, H, W, 0, 0, gn_n_groups);\n acts->layer_sizes[layer] = resblock_acts->n_acts;\n return layer + 1;\n}\n\nint make_avgpool_acts_ptrs(UnetBlock* acts, int layer, int B, int C, int H, int W) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(AvgpoolActs));\n acts->layer_sizes[layer] = B * C * H/2 * W/2;\n return layer + 1;\n}\n\nint make_attention_block_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C, int H, int W, int gn_n_groups, int HS\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(AttentionActs));\n AttentionActs* att_acts = (AttentionActs*)acts->layer_ptrs[layer];\n attention_block_count_acts(att_acts, B, C, H, W, gn_n_groups, HS);\n acts->layer_sizes[layer] = att_acts->n_acts;\n return layer + 1;\n}\n\nint make_concat_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C1, int C2, int H, int W\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ConcatChannelActs));\n acts->layer_sizes[layer] = B * (C1 + C2) * H * W;\n return layer + 1;\n}\n\nint make_upsample_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C, int H, int W\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(UpsampleActs));\n acts->layer_sizes[layer] = B * C * H * 2 * W * 2;\n return layer + 1;\n}\n\nint make_groupnorm_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C, int H, int W, int gn_n_groups\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(GroupNormActs));\n // out (B, C, H, W), mean and rstd (B, gn_n_groups)\n acts->layer_sizes[layer] = B * C * H * W + B * gn_n_groups * 2;\n return layer + 1;\n}\n\nvoid unet_setup_and_count_acts(UnetConfig* config, UnetBlock* acts) {\n assert (config->n_layers > 0);\n\n acts->n_layers = config->n_layers;\n acts->layer_ptrs = (void**)mallocCheck(config->n_layers * sizeof(void*));\n acts->layer_sizes = (size_t*)mallocCheck(config->n_layers * sizeof(size_t));\n\n int B = config->B;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int C_out = config->C_out;\n int H = config->H;\n int W = config->W;\n int gn_n_groups = config->gn_n_groups;\n int HS = config->HS;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n int layer = 0;\n layer = make_linear_acts_ptrs(acts, layer, B, C_time);\n layer = make_silu_acts_ptrs(acts, layer, B * C_time);\n layer = make_linear_acts_ptrs(acts, layer, B, C_time);\n int OC = C_model;\n layer = make_convk3_acts_ptrs(acts, layer, B, OC, H, W);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n layer = make_resblock_acts_ptrs(acts, layer, C, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups, HS);\n }\n }\n if (l < n_levels - 1) {\n layer = make_avgpool_acts_ptrs(acts, layer, B, C, H, W);\n H /= 2;\n W /= 2;\n }\n }\n\n layer = make_resblock_acts_ptrs(acts, layer, C, C_time, C, B, H, W, gn_n_groups);\n layer = make_attention_block_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups, HS);\n layer = make_resblock_acts_ptrs(acts, layer, C, C_time, C, B, H, W, gn_n_groups);\n\n // upwards\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = mult * C_model;\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n int prev_mult = config->channel_mult[l - 1];\n C_skip = prev_mult * C_model;\n }\n layer = make_concat_acts_ptrs(acts, layer, B, C, C_skip, H, W);\n OC = mult * C_model;\n layer = make_resblock_acts_ptrs(acts, layer, C + C_skip, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups, HS);\n }\n }\n if (l > 0) {\n layer = make_upsample_acts_ptrs(acts, layer, B, C, H, W);\n H *= 2;\n W *= 2;\n }\n }\n layer = make_groupnorm_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups);\n layer = make_silu_acts_ptrs(acts, layer, B * C * H * W);\n layer = make_convk3_acts_ptrs(acts, layer, B, C_out, H, W);\n\n // count total size\n size_t total_size = 0;\n for (int i = 0; i < config->n_layers; i++) {\n total_size += acts->layer_sizes[i];\n }\n acts->total_size = total_size;\n}\n\nvoid unet_setup_and_count_backs(UnetConfig* config, UnetBackBlock* back) {\n int B = config->B;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int H = config->H;\n int W = config->W;\n\n size_t max_BCHW = 3 * B * C_model * H * W;\n size_t max_B3CHW = 2 * B * 3 * C_model * H * W;\n back->back_sizes[0] = max_BCHW;\n back->back_sizes[1] = max_BCHW;\n back->back_sizes[2] = max_BCHW;\n back->back_sizes[3] = max_BCHW;\n back->back_sizes[4] = max_BCHW;\n back->back_sizes[5] = max_BCHW;\n back->back_sizes[6] = max_B3CHW;\n back->back_sizes[7] = max_B3CHW;\n back->back_sizes[8] = B * C_time;\n back->back_sizes[9] = B * C_time;\n\n // fill the sizes of the skip connection backward acts\n int C = C_model;\n for (int l = 0; l < config->n_levels; l++) {\n for (int r = 0; r < (config->n_res + 1); r++) {\n int s = l * 3 + r + NUM_BACK_BUFS;\n assert(s < UNET_NUM_BACK_ACTS);\n back->back_sizes[s] = B * C * H * W;\n int mult = config->channel_mult[l];\n C = C_model * mult;\n if (l > 0 && r == 0) {\n H /= 2;\n W /= 2;\n }\n }\n }\n\n size_t n_backs = 0;\n for (int i = 0; i < UNET_NUM_BACK_ACTS; i++) {\n n_backs += back->back_sizes[i];\n }\n back->n_backs = n_backs;\n}\n\n\nvoid unet_make_ptrs_and_count_memory(Unet* unet)\n{\n unet_setup_and_count_params(&(unet->config), &(unet->params));\n unet_setup_and_count_params(&(unet->config), &(unet->grads));\n unet_setup_and_count_acts(&(unet->config), &(unet->acts));\n unet_setup_and_count_backs(&(unet->config), &(unet->back));\n}\n\n// ------------------------------------------------------------------------------------------------\n// malloc unet memory and set pointers\ntypedef struct {\n float* input;\n float* prev_out;\n int layer;\n float* emb;\n float* emb_cum; // cumulative demb\n float* mem_ptr;\n int skip_idx;\n} LayerInfo;\n\n\nvoid set_linear_params_ptrs(\n UnetBlock* params,\n LayerInfo* l_info,\n int C_in, int C_out\n) {\n int layer = l_info->layer;\n LinearParams* linear_params = (LinearParams*)params->layer_ptrs[layer];\n linear_set_param_ptrs(linear_params, l_info->mem_ptr, C_in, C_out);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_convk3_params_ptrs(UnetBlock* params, LayerInfo* l_info, int C_in, int OC) {\n int layer = l_info->layer;\n ConvK3Params* convk3_params = (ConvK3Params*)params->layer_ptrs[layer];\n convk3_set_param_ptrs(convk3_params, l_info->mem_ptr, C_in, OC);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_resblock_layer_params_ptrs(UnetBlock* params, LayerInfo* l_info, int C, int OC) {\n int layer = l_info->layer;\n ResBlockParameters* resblock_params = (ResBlockParameters*)params->layer_ptrs[layer];\n set_resblock_params_ptrs(C, OC, resblock_params, l_info->mem_ptr);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_attention_layer_params_ptrs(UnetBlock* params, LayerInfo* l_info) {\n int layer = l_info->layer;\n AttentionParams* att_params = (AttentionParams*)params->layer_ptrs[layer];\n set_attention_params_pointers(att_params, l_info->mem_ptr);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_groupnorm_layer_params_ptrs(UnetBlock* params, LayerInfo* l_info, int C) {\n int layer = l_info->layer;\n GroupNormParams* gn_params = (GroupNormParams*)params->layer_ptrs[layer];\n gn_set_param_ptrs(gn_params, l_info->mem_ptr, C);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\n\nvoid point_unet_params(UnetConfig* config, UnetBlock* params) {\n // common constants\n int C_in = config->C_in;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int C_out = config->C_out;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n LayerInfo l_info;\n l_info.layer = 0;\n l_info.mem_ptr = params->mem_ptr;\n set_linear_params_ptrs(params, &l_info, C_model, C_time);\n l_info.layer++; // skip silu layer\n set_linear_params_ptrs(params, &l_info, C_time, C_time);\n int OC = C_model;\n set_convk3_params_ptrs(params, &l_info, C_in, OC);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n set_resblock_layer_params_ptrs(params, &l_info, C, OC);\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_params_ptrs(params, &l_info);\n }\n }\n if (l < n_levels - 1) {\n l_info.layer++; // skip avgpool layer\n }\n }\n\n set_resblock_layer_params_ptrs(params, &l_info, C, C);\n set_attention_layer_params_ptrs(params, &l_info);\n set_resblock_layer_params_ptrs(params, &l_info, C, C);\n\n // upward: skip one layer for concat\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = mult * C_model;\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config->channel_mult[l - 1];\n }\n l_info.layer++;\n OC = mult * C_model;\n set_resblock_layer_params_ptrs(params, &l_info, C + C_skip, OC);\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_params_ptrs(params, &l_info);\n }\n }\n if (l > 0) {\n l_info.layer++; // upsample\n }\n }\n\n set_groupnorm_layer_params_ptrs(params, &l_info, C);\n l_info.layer++; // silu\n set_convk3_params_ptrs(params, &l_info, C, C_out);\n}\n\nvoid set_linear_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n LinearActs* linear_acts = (LinearActs*)acts->layer_ptrs[layer];\n linear_acts->out = l_info->mem_ptr;\n l_info->mem_ptr += acts->layer_sizes[layer];\n linear_acts->inp = l_info->input;\n l_info->input = linear_acts->out;\n l_info->layer++;\n}\n\nvoid set_silu_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n SiluActs* silu_acts = (SiluActs*)acts->layer_ptrs[layer];\n silu_acts->out = l_info->mem_ptr;\n l_info->mem_ptr += acts->layer_sizes[layer];\n silu_acts->x = l_info->input;\n l_info->input = silu_acts->out;\n l_info->layer++;\n}\n\nvoid set_convk3_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n ConvK3Acts* convk3_acts = (ConvK3Acts*)acts->layer_ptrs[layer];\n convk3_acts->out = l_info->mem_ptr;\n l_info->mem_ptr += acts->layer_sizes[layer];\n convk3_acts->inp = l_info->input;\n l_info->input = convk3_acts->out;\n l_info->layer++;\n}\n\nvoid set_resblock_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info, int C, int OC) {\n int layer = l_info->layer;\n ResBlockActivations* resblock_acts = (ResBlockActivations*)acts->layer_ptrs[layer];\n set_resblock_acts_ptrs(C, OC, resblock_acts, l_info->mem_ptr);\n resblock_acts->input = l_info->input;\n resblock_acts->emb = l_info->emb;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = resblock_acts->add2;\n l_info->layer++;\n}\n\nvoid set_avgpool_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n AvgpoolActs* avgpool_acts = (AvgpoolActs*)acts->layer_ptrs[layer];\n avgpool_acts->out = l_info->mem_ptr;\n avgpool_acts->x = l_info->input;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = avgpool_acts->out;\n l_info->layer++;\n}\n\nvoid set_attention_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n AttentionActs* att_acts = (AttentionActs*)acts->layer_ptrs[layer];\n set_attention_acts_pointers(att_acts, l_info->mem_ptr);\n att_acts->input = l_info->input;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = att_acts->add;\n l_info->layer++;\n}\n\nvoid set_concat_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info, float** skip_acts) {\n int layer = l_info->layer;\n ConcatChannelActs* cat_acts = (ConcatChannelActs*)acts->layer_ptrs[layer];\n cat_acts->out = l_info->mem_ptr;\n cat_acts->x1 = l_info->input;\n cat_acts->x2 = skip_acts[l_info->skip_idx];\n\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = cat_acts->out;\n l_info->layer++;\n l_info->skip_idx--;\n}\n\nvoid set_upsample_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n UpsampleActs* up_acts = (UpsampleActs*)acts->layer_ptrs[layer];\n up_acts->x = l_info->input;\n up_acts->out = l_info->mem_ptr;\n\n l_info->mem_ptr += acts-> layer_sizes[layer];\n l_info->input = up_acts->out;\n l_info->layer++;\n}\n\nvoid set_groupnorm_layer_acts_ptrs(\n UnetBlock* acts, LayerInfo* l_info,\n int B, int C, int H, int W, int gn_n_groups\n) {\n int layer = l_info->layer;\n GroupNormActs* gn_a = (GroupNormActs*)acts->layer_ptrs[layer];\n gn_set_act_ptrs(gn_a, l_info->mem_ptr, B, C, H, W, gn_n_groups);\n gn_a->x = l_info->input;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = gn_a->out;\n l_info->layer++;\n}\n\nvoid point_unet_acts(UnetConfig* config, UnetBlock* acts, float** skip_acts, float* input, float* timestep_emb) {\n // common constants\n int B = config->B;\n int C_model = config->C_model;\n int H = config->H;\n int W = config->W;\n int gn_n_groups = config->gn_n_groups;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n LayerInfo l_info;\n l_info.layer = 0;\n l_info.mem_ptr = acts->mem_ptr;\n\n // timestep emb\n l_info.input = timestep_emb;\n set_linear_acts_ptrs(acts, &l_info);\n set_silu_acts_ptrs(acts, &l_info);\n set_linear_acts_ptrs(acts, &l_info);\n // point input to model input, and set outputu of the last linear layer as the embedding input\n l_info.emb = l_info.input;\n l_info.input = input;\n\n int OC = C_model;\n set_convk3_acts_ptrs(acts, &l_info);\n int skip_idx = 0;\n skip_acts[skip_idx] = l_info.input; // set the output of the previouse layer as input to the skip connection\n skip_idx++;\n\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n set_resblock_layer_acts_ptrs(acts, &l_info, C, OC);\n if (l < att_start_level) {\n // for the first 2 layers, the skip connection comes from the resblocks\n skip_acts[skip_idx] = l_info.input;\n skip_idx++;\n }\n\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_acts_ptrs(acts, &l_info);\n skip_acts[skip_idx] = l_info.input;\n skip_idx++;\n }\n }\n if (l < n_levels - 1) {\n set_avgpool_layer_acts_ptrs(acts, &l_info);\n skip_acts[skip_idx] = l_info.input;\n skip_idx++;\n }\n }\n\n set_resblock_layer_acts_ptrs(acts, &l_info, C, C);\n set_attention_layer_acts_ptrs(acts, &l_info);\n set_resblock_layer_acts_ptrs(acts, &l_info, C, C);\n\n // upward: set concat\n l_info.skip_idx = NUM_SKIP_CONNECTIONS - 1;\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = mult * C_model;\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config->channel_mult[l - 1];\n }\n set_concat_layer_acts_ptrs(acts, &l_info, skip_acts);\n OC = mult * C_model;\n set_resblock_layer_acts_ptrs(acts, &l_info, C + C_skip, OC);\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_acts_ptrs(acts, &l_info);\n }\n }\n if (l > 0) {\n set_upsample_layer_acts_ptrs(acts, &l_info);\n }\n }\n set_groupnorm_layer_acts_ptrs(acts, &l_info, B, C, H, W, gn_n_groups);\n set_silu_acts_ptrs(acts, &l_info);\n set_convk3_acts_ptrs(acts, &l_info);\n}\n\nvoid set_unet_back_pointers(UnetBackBlock* back) {\n UnetBackActs* back_acts = &(back->back_acts);\n float* back_mem_ptr = back->back_mem;\n\n float** ptrs[] = {\n &(back_acts->dinp_buf1),\n &(back_acts->dinp_buf2),\n &(back_acts->buf1_BCHW),\n &(back_acts->buf2_BCHW),\n &(back_acts->buf3_BCHW),\n &(back_acts->buf4_BCHW),\n &(back_acts->buf1_B3CHW),\n &(back_acts->buf2_B3CHW),\n &(back_acts->demb_buf1),\n &(back_acts->demb_buf2)\n };\n\n for (int i = 0; i < NUM_BACK_BUFS; i++) {\n *(ptrs[i]) = back_mem_ptr;\n back_mem_ptr += back->back_sizes[i];\n }\n\n // point the back acts for skip connections\n for (int s = 0; s < NUM_SKIP_CONNECTIONS; s++) {\n back_acts->skip_dinps[s] = back_mem_ptr;\n back_mem_ptr += back->back_sizes[NUM_BACK_BUFS + s];\n }\n}\n\nvoid malloc_and_point(Unet* unet) {\n float *param_mem, *grad_mem, *act_mem, *back_mem;\n cudaCheck(cudaMalloc(¶m_mem, unet->params.total_size * sizeof(float)));\n cudaCheck(cudaMalloc(&grad_mem, unet->params.total_size * sizeof(float)));\n cudaCheck(cudaMalloc(&act_mem, unet->acts.total_size * sizeof(float)));\n cudaCheck(cudaMalloc(&back_mem, unet->back.n_backs * sizeof(float)));\n\n cudaCheck(cudaMemset(param_mem, 0, unet->params.total_size * sizeof(float)));\n cudaCheck(cudaMemset(grad_mem, 0, unet->params.total_size * sizeof(float)));\n cudaCheck(cudaMemset(act_mem, 0, unet->acts.total_size * sizeof(float)));\n cudaCheck(cudaMemset(back_mem, 0, unet->back.n_backs * sizeof(float)));\n\n // constants\n int B = unet->config.B;\n int C_in = unet->config.C_in;\n int C_model = unet->config.C_model;\n int H = unet->config.H;\n int W = unet->config.W;\n\n // malloc and point input and timestep emb\n cudaCheck(cudaMalloc(&(unet->input), B * C_in * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&(unet->timestep_emb), B * C_model * sizeof(float)));\n // malloc a nicely aligned buffer for the loss\n cudaCheck(cudaMalloc(&(unet->gpu_loss), 64 * sizeof(float)));\n\n unet->params.mem_ptr = param_mem;\n unet->grads.mem_ptr = grad_mem;\n unet->acts.mem_ptr = act_mem;\n unet->back.back_mem = back_mem;\n\n point_unet_params(&(unet->config), &(unet->params));\n point_unet_params(&(unet->config), &(unet->grads));\n point_unet_acts(&(unet->config), &(unet->acts), unet->skip_acts, unet->input, unet->timestep_emb);\n set_unet_back_pointers(&(unet->back));\n}\n\n// ------------------------------------------------------------------------------------------------\n// forward pass\n\nvoid linear_layer_forward(\n Unet* unet,\n LayerInfo* l_info,\n cublasHandle_t cublas_handle, int N, int C, int OC, int block_size\n) {\n int layer = l_info->layer;\n LinearParams* linear_params = (LinearParams*)unet->params.layer_ptrs[layer];\n LinearActs* linear_acts = (LinearActs*)unet->acts.layer_ptrs[layer];\n matmul_forward2(cublas_handle, linear_acts->out, linear_acts->inp, linear_params->w, linear_params->b, N, C, OC, block_size);\n l_info->layer++;\n}\n\nvoid silu_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int N, int block_size\n) {\n SiluActs* acts = (SiluActs*)unet->acts.layer_ptrs[l_info->layer];\n silu_forward(acts->x, acts->out, N, block_size);\n l_info->layer++;\n}\n\nvoid convk3_layer_forward(\n Unet* unet, LayerInfo* l_info,\n cublasHandle_t cublas_handle, int B, int C, int OC, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n ConvK3Params* conv_params = (ConvK3Params*)unet->params.layer_ptrs[layer];\n ConvK3Acts* conv_acts = (ConvK3Acts*)unet->acts.layer_ptrs[layer];\n conv2d_k3_forward3(conv_acts->inp, conv_params->w, conv_params->b, conv_acts->out, B, C, OC, H, W);\n l_info->layer++;\n}\nvoid resblock_layer_forward(\n Unet* unet, LayerInfo* l_info,\n cublasHandle_t cublas_handle, int C_in, int C_emb, int C_out, int B, int H, int W, int gn_n_groups, int block_size\n) {\n int layer = l_info->layer;\n ResBlockParameters* resblock_params = (ResBlockParameters*)unet->params.layer_ptrs[layer];\n ResBlockActivations* resblock_acts = (ResBlockActivations*)unet->acts.layer_ptrs[layer];\n resblock_forward(cublas_handle, C_in, C_emb, C_out, B, H, W, block_size, 0, 0, gn_n_groups, resblock_params, resblock_acts);\n l_info->layer++;\n}\n\nvoid avgpool_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n AvgpoolActs* avgpool_acts = (AvgpoolActs*)unet->acts.layer_ptrs[layer];\n avgpool_2d_forward1(avgpool_acts->out, avgpool_acts->x, B, C, H, W, block_size);\n l_info->layer++;\n}\n\nvoid attention_layer_forward(\n Unet* unet, LayerInfo* l_info,\n cublasHandle_t cublas_handle, int B, int C, int H, int W, int HS, int gn_n_groups, int block_size\n) {\n int layer = l_info->layer;\n AttentionParams* att_params = (AttentionParams*)unet->params.layer_ptrs[layer];\n AttentionActs* att_acts = (AttentionActs*)unet->acts.layer_ptrs[layer];\n attention_block_forward(cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, att_params, att_acts);\n l_info->layer++;\n}\n\nvoid concat_layer_forward(\n Unet* unet,\n LayerInfo* l_info,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n ConcatChannelActs* cat_acts = (ConcatChannelActs*)unet->acts.layer_ptrs[layer];\n concat_channel_forward(cat_acts->x1, cat_acts->x2, cat_acts->out, B, C1, C2, H, W, block_size);\n l_info->layer++;\n}\n\nvoid upsample_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n UpsampleActs* up_acts = (UpsampleActs*)unet->acts.layer_ptrs[layer];\n upsample_forward1(up_acts->out, up_acts->x, B, C, H, W, block_size);\n l_info->layer++;\n}\n\nvoid groupnorm_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int gn_n_groups\n) {\n int layer = l_info->layer;\n GroupNormParams* gn_p = (GroupNormParams*)unet->params.layer_ptrs[layer];\n GroupNormActs* gn_a = (GroupNormActs*)unet->acts.layer_ptrs[layer];\n groupnorm_forward(gn_a->x, gn_p->w, gn_p->b, gn_a->out, gn_a->mean, gn_a->rstd, B, C, H, W, gn_n_groups);\n l_info->layer++;\n}\n\nvoid unet_forward(\n cublasHandle_t cublas_handle,\n Unet* unet\n) {\n // common constants\n UnetConfig config = unet->config;\n int B = config.B;\n int C_in = config.C_in;\n int C_model = config.C_model;\n int C_time = config.C_time;\n int C_out = config.C_out;\n int H = config.H;\n int W = config.W;\n int block_size = config.block_size;\n int gn_n_groups = config.gn_n_groups;\n int HS = config.HS;\n int n_levels = config.n_levels;\n int n_res = config.n_res;\n int att_start_level = config.att_start_level;\n\n\n LayerInfo l_info;\n l_info.layer = 0;\n linear_layer_forward(unet, &l_info, cublas_handle, B, C_model, C_time, block_size);\n silu_layer_forward(unet, &l_info, B * C_time, block_size);\n linear_layer_forward(unet, &l_info, cublas_handle, B, C_time, C_time, block_size);\n int OC = C_model;\n convk3_layer_forward(unet, &l_info, cublas_handle, B, C_in, OC, H, W, block_size);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config.channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n resblock_layer_forward(unet, &l_info, cublas_handle, C, C_time, OC, B, H, W, gn_n_groups, block_size);\n C = OC;\n if (l >= att_start_level) {\n attention_layer_forward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n }\n }\n if (l < n_levels - 1) {\n avgpool_layer_forward(unet, &l_info, B, C, H, W, block_size);\n H /= 2;\n W /= 2;\n }\n }\n\n resblock_layer_forward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n attention_layer_forward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n resblock_layer_forward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n\n // upward\n for (int l = n_levels - 1; l >= 0; l--) {\n int C_skip = C_model * config.channel_mult[l];\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config.channel_mult[l - 1];\n }\n concat_layer_forward(unet, &l_info, B, C, C_skip, H, W, block_size);\n OC = C_model * config.channel_mult[l];\n resblock_layer_forward(unet, &l_info, cublas_handle, C + C_skip, C_time, OC, B, H, W, gn_n_groups, block_size);\n C = OC;\n if (l >= att_start_level) {\n attention_layer_forward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n }\n }\n if (l > 0) {\n upsample_layer_forward(unet, &l_info, B, C, H, W, block_size);\n H *= 2;\n W *= 2;\n }\n }\n\n groupnorm_layer_forward(unet, &l_info, B, C, H, W, gn_n_groups);\n silu_layer_forward(unet, &l_info, B * C * H * W, block_size);\n convk3_layer_forward(unet, &l_info, cublas_handle, B, C, C_out, H, W, block_size);\n\n ConvK3Acts* last_acts = (ConvK3Acts*)unet->acts.layer_ptrs[l_info.layer - 1];\n unet->output = last_acts->out;\n if (unet->target == nullptr) {\n // don't have targets, not computing loss\n unet->mean_loss = -1.0f;\n } else {\n mse_forward(unet->output, unet->target, unet->gpu_loss, B * C_out * H * W, block_size);\n cudaCheck(cudaMemcpy(&(unet->mean_loss), unet->gpu_loss, sizeof(float), cudaMemcpyDeviceToHost));\n }\n}\n\n// ------------------------------------------------------------------------------------------------\n// Unet backward\n\nvoid swap_layer_in_out_ptrs(LayerInfo* l_info) {\n float* temp = l_info->input;\n l_info->input = l_info->prev_out;\n l_info->prev_out = temp;\n}\n\nvoid linear_layer_backward(\n Unet* unet,\n LayerInfo* l_info,\n cublasHandle_t cublas_handle, int N, int C, int OC\n) {\n int layer = l_info->layer;\n LinearParams* linear_params = (LinearParams*)unet->params.layer_ptrs[layer];\n LinearParams* linear_grads = (LinearParams*)unet->grads.layer_ptrs[layer];\n LinearActs* linear_acts = (LinearActs*)unet->acts.layer_ptrs[layer];\n matmul_backward1(cublas_handle, l_info->prev_out, linear_grads->w, linear_grads->b, l_info->input, linear_acts->inp, linear_params->w, N, C, OC);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid silu_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int N, int block_size\n) {\n SiluActs* acts = (SiluActs*)unet->acts.layer_ptrs[l_info->layer];\n silu_backward(l_info->input, acts->x, l_info->prev_out, N, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid convk3_layer_backward(\n Unet* unet, LayerInfo* l_info,\n cublasHandle_t cublas_handle, int B, int C, int OC, int H, int W, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n add_inplace_forward(skip_dinp, l_info->input, B * OC * H * W, block_size);\n l_info->skip_idx--;\n }\n int layer = l_info->layer;\n ConvK3Params* conv_params = (ConvK3Params*)unet->params.layer_ptrs[layer];\n ConvK3Params* conv_grads = (ConvK3Params*)unet->grads.layer_ptrs[layer];\n ConvK3Acts* conv_acts = (ConvK3Acts*)unet->acts.layer_ptrs[layer];\n UnetBackActs* unet_back_acts = &(unet->back.back_acts);\n conv2d_k3_backward2(l_info->input, conv_acts->inp, conv_params->w, unet_back_acts->buf1_BCHW, unet_back_acts->buf2_BCHW, l_info->prev_out, conv_grads->w, conv_grads->b, B, C, OC, H, W);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid resblock_layer_backward(\n Unet* unet, LayerInfo* l_info,\n cublasHandle_t cublas_handle, int C_in, int C_emb, int C_out, int B, int H, int W, int gn_n_groups, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n add_inplace_forward(skip_dinp, l_info->input, B * C_out * H * W, block_size);\n l_info->skip_idx--;\n }\n int layer = l_info->layer;\n ResBlockParameters* resblock_params = (ResBlockParameters*)unet->params.layer_ptrs[layer];\n ResBlockParameters* resblock_grads = (ResBlockParameters*)unet->grads.layer_ptrs[layer];\n ResBlockActivations* resblock_acts = (ResBlockActivations*)unet->acts.layer_ptrs[layer];\n ResBlockBackwardActivations back_acts;\n UnetBackActs* unet_back_acts = &(unet->back.back_acts);\n\n // point back_acts to buffers\n back_acts.buf_BCemb = unet_back_acts->buf1_BCHW;\n back_acts.buf_BCHoWo = unet_back_acts->buf2_BCHW;\n back_acts.buf1_BCHW = unet_back_acts->buf3_BCHW;\n back_acts.buf2_BCHW = unet_back_acts->buf4_BCHW;\n back_acts.dweight_buf = unet_back_acts->buf1_B3CHW;\n back_acts.dbias_buf = unet_back_acts->buf2_B3CHW;\n back_acts.dout = l_info->input;\n back_acts.dx = l_info->prev_out;\n back_acts.demb = l_info->emb;\n\n resblock_backward(cublas_handle, C_in, C_emb, C_out, B, H, W, block_size, 0, 0, gn_n_groups, resblock_params, resblock_grads, resblock_acts, &back_acts);\n swap_layer_in_out_ptrs(l_info);\n\n // add demb to cumulative demb\n add_inplace_forward(l_info->emb, l_info->emb_cum, B * C_emb, block_size);\n l_info->layer--;\n}\n\nvoid attention_layer_backward(\n Unet* unet, LayerInfo* l_info,\n cublasHandle_t cublas_handle, int B, int C, int H, int W, int HS, int gn_n_groups, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n //printf(\"att adding skip %d\\n\", l_info->skip_idx);\n add_inplace_forward(skip_dinp, l_info->input, B * C * H * W, block_size);\n l_info->skip_idx--;\n }\n int layer = l_info->layer;\n AttentionParams* att_params = (AttentionParams*)unet->params.layer_ptrs[layer];\n AttentionParams* att_grads = (AttentionParams*)unet->grads.layer_ptrs[layer];\n AttentionActs* att_acts = (AttentionActs*)unet->acts.layer_ptrs[layer];\n AttentionBackwardActs att_backs;\n UnetBackActs* unet_backs = &(unet->back.back_acts);\n att_backs.buf1_BCHW = unet_backs->buf1_BCHW;\n att_backs.buf2_BCHW = unet_backs->buf2_BCHW;\n att_backs.buf_B3CHW = unet_backs->buf1_B3CHW;\n att_backs.dqkvr = unet_backs->buf2_B3CHW;\n att_backs.dpreatt = unet_backs->buf3_BCHW;\n att_backs.datt = unet_backs->buf4_BCHW;\n att_backs.dout = l_info->input;\n att_backs.dinp = l_info->prev_out;\n\n attention_block_backward(cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, att_params, att_acts, &att_backs, att_grads);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid concat_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n //printf(\"concat skip %d\\n\", l_info->skip_idx);\n concat_channel_backward(l_info->input, l_info->prev_out, skip_dinp, B, C1, C2, H, W, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n l_info->skip_idx++;\n}\n\n\nvoid avgpool_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n //printf(\"avgpool adding skip %d\\n\", l_info->skip_idx);\n add_inplace_forward(skip_dinp, l_info->input, B * C * H/2 * W/2, block_size);\n l_info->skip_idx--;\n }\n avgpool_2d_backward1(l_info->input, l_info->prev_out, B, C, H, W, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid upsample_layer_backward(\n Unet* unet, LayerInfo *l_info,\n int B, int C, int H, int W, int block_size\n) {\n upsample_backward1(l_info->prev_out, l_info->input, B, C, H, W, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid groupnorm_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int gn_n_groups\n) {\n int layer = l_info->layer;\n GroupNormParams* gn_p = (GroupNormParams*)unet->params.layer_ptrs[layer];\n GroupNormParams* gn_g = (GroupNormParams*)unet->grads.layer_ptrs[layer];\n GroupNormActs* gn_a = (GroupNormActs*)unet->acts.layer_ptrs[layer];\n groupnorm_backward(l_info->input, gn_a->x, gn_a->mean, gn_a->rstd, gn_p->w, l_info->prev_out, gn_g->w, gn_g->b, B, C, H, W, gn_n_groups);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid unet_backward(\n cublasHandle_t cublas_handle,\n Unet* unet\n \n) {\n if (unet->mean_loss == -1.0f or unet->target == nullptr) {\n printf(\"Error: must run forward with targets before calling backward\\n\");\n exit(EXIT_FAILURE);\n }\n // common constants\n UnetConfig config = unet->config;\n int B = config.B;\n int C_in = config.C_in;\n int C_model = config.C_model;\n int C_time = config.C_time;\n int C_out = config.C_out;\n int H = config.H;\n int W = config.W;\n int block_size = config.block_size;\n int gn_n_groups = config.gn_n_groups;\n int HS = config.HS;\n int n_levels = config.n_levels;\n int n_res = config.n_res;\n int att_start_level = config.att_start_level;\n\n UnetBackActs* back_acts = &(unet->back.back_acts);\n mse_backward(unet->output, unet->target, back_acts->dinp_buf1, B * C_out * H * W, block_size);\n\n LayerInfo l_info;\n l_info.layer = config.n_layers - 1;\n l_info.input = back_acts->dinp_buf1;\n l_info.prev_out = back_acts->dinp_buf2;\n l_info.emb = back_acts->demb_buf1;\n l_info.emb_cum = back_acts->demb_buf2;\n l_info.skip_idx = 0;\n\n int C;\n int OC;\n int C_skip;\n C = C_model;\n OC = C_model;\n convk3_layer_backward(unet, &l_info, cublas_handle, B, C, C_out, H, W, block_size);\n silu_layer_backward(unet, &l_info, B * C * H * W, block_size);\n groupnorm_layer_backward(unet, &l_info, B, C, H, W, gn_n_groups);\n\n for (int l = 0; l < n_levels; l++) {\n for (int r = n_res; r >= 0; r--) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config.channel_mult[l - 1];\n } else {\n C_skip = C_model * config.channel_mult[l];\n }\n if (l >= att_start_level) {\n attention_layer_backward(unet, &l_info, cublas_handle, B, OC, H, W, HS, gn_n_groups, block_size);\n }\n if (r == 0 && l < n_levels - 1) {\n C = C_model * config.channel_mult[l + 1];\n }\n resblock_layer_backward(unet, &l_info, cublas_handle, C + C_skip, C_time, OC, B, H, W, gn_n_groups, block_size);\n concat_layer_backward(unet, &l_info, B, C, C_skip, H, W, block_size);\n }\n if (l < n_levels - 1) {\n H /= 2;\n W /= 2;\n upsample_layer_backward(unet, &l_info, B, C, H, W, block_size);\n OC = C_model * config.channel_mult[l + 1];\n }\n }\n\n resblock_layer_backward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n attention_layer_backward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n resblock_layer_backward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n\n // we will need to change the backward for all the other blocks below to include the dout contribution from the skip connection.\n l_info.skip_idx = NUM_SKIP_CONNECTIONS - 1;\n for (int l = n_levels - 1; l >= 0; l--) {\n if (l < n_levels - 1) {\n H *= 2;\n W *= 2;\n avgpool_layer_backward(unet, &l_info, B, C, H, W, block_size, 1);\n }\n for (int r = 0; r < n_res; r++) {\n if (l >= att_start_level) {\n attention_layer_backward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, 1);\n }\n int res_add_skip = (l < att_start_level) ? 1 : 0;\n if (r == n_res - 1 && l > 0) {\n int prev_mult = config.channel_mult[l - 1];\n C = prev_mult * C_model;\n }\n resblock_layer_backward(unet, &l_info, cublas_handle, C, C_time, OC, B, H, W, gn_n_groups, block_size, res_add_skip);\n OC = C;\n }\n }\n\n convk3_layer_backward(unet, &l_info, cublas_handle, B, C_in, OC, H, W, block_size, 1);\n unet->dinp = l_info.input;\n\n // reset dout to point to emb_cum\n l_info.input = l_info.emb_cum;\n linear_layer_backward(unet, &l_info, cublas_handle, B, C_time, C_time);\n silu_layer_backward(unet, &l_info, B * C_time, block_size);\n linear_layer_backward(unet, &l_info, cublas_handle, B, C_model, C_time);\n}\n\n// ------------------------------------------------------------------------------------------------\n// zero grad and ADAMW\n\nvoid unet_zero_grad(Unet* unet) {\n cudaCheck(cudaMemset(unet->grads.mem_ptr, 0, unet->grads.total_size * sizeof(float)));\n cudaCheck(cudaMemset(unet->back.back_mem, 0, unet->back.n_backs * sizeof(float)));\n //cudaCheck(cudaMemset(unet->acts.mem_ptr, 0, unet->acts.total_size * sizeof(float)));\n}\n\n// below copied from llm.c\n\n// Implements linear interpolation using only two floating-point operations (as opposed to three in a naive implementation).\n// Reference: https://developer.nvidia.com/blog/lerp-faster-cuda\n__device__ inline float lerp(float start, float end, float weight) {\n return fma(weight, end, fma(-weight, start, start));\n}\n\n__global__ void adamw_kernel2(float* params_memory, float* grads_memory, float* m_memory, float* v_memory, long num_parameters,\n float learning_rate, float beta1, float beta2, float beta1_correction, float beta2_correction, float eps, float weight_decay) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i >= num_parameters) return; // guard\n float grad = grads_memory[i];\n float m = m_memory[i];\n float v = v_memory[i];\n // update the first moment (momentum)\n m = lerp(grad, m, beta1);\n m_memory[i] = m;\n // update the second moment (RMSprop)\n v = lerp(grad * grad, v, beta2);\n v_memory[i] = v;\n m /= beta1_correction; // m_hat\n v /= beta2_correction; // v_hat\n params_memory[i] -= learning_rate * (m / (sqrtf(v) + eps) + weight_decay * params_memory[i]);\n}\n\nvoid unet_update(Unet* unet, float learning_rate, float beta1, float beta2, float eps, float weight_decay, int t) {\n // lazy allocation\n if (unet->m_memory == nullptr) {\n size_t param_size = unet->params.total_size * sizeof(float);\n cudaCheck(cudaMalloc(&(unet->m_memory), param_size));\n cudaCheck(cudaMalloc(&(unet->v_memory), param_size));\n cudaCheck(cudaMemset(unet->m_memory, 0, param_size));\n cudaCheck(cudaMemset(unet->v_memory, 0, param_size));\n printf(\"allocated %zu MiB for AdamW state m\\n\", param_size / 1024 / 1024);\n printf(\"allocated %zu MiB for AdamW state v\\n\", param_size / 1024 / 1024);\n }\n\n int block_size = unet->config.block_size;\n int n_blocks = ceil_div(unet->params.total_size, (size_t)block_size);\n float beta1_correction = 1.0f - powf(beta1, t);\n float beta2_correction = 1.0f - powf(beta2, t);\n adamw_kernel2<<>>(unet->params.mem_ptr, unet->grads.mem_ptr, unet->m_memory, unet->v_memory, unet->params.total_size,\n learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay);\n cudaCheck(cudaGetLastError());\n}\n\n// ------------------------------------------------------------------------------------------------\n// save model states\n\nvoid save_unet_states(Unet* unet, GaussianDiffusion* diffusion, const char* filename) {\n // create header\n int model_header[256];\n model_header[0] = 12345678; // magic number\n model_header[1] = unet->config.B;\n model_header[2] = unet->config.C_in;\n model_header[3] = unet->config.C_model;\n model_header[4] = unet->config.C_out;\n model_header[5] = unet->config.H;\n model_header[6] = unet->config.W;\n model_header[7] = unet->config.max_period;\n model_header[8] = 1; // whether to save the AdamW states, this will be 0 if the model is saved from pytorch\n model_header[9] = 1; // whether to save rng states\n\n FILE* model_file = fopenCheck(filename, \"wb\");\n printf(\"\\nsaving model details\");\n int result = fwrite(model_header, sizeof(int), 256, model_file);\n assert(result == 256);\n\n // save params\n float* params_cpu_buf = (float*)mallocCheck(unet->params.total_size * sizeof(float));\n cudaCheck(cudaMemcpy(params_cpu_buf, unet->params.mem_ptr, unet->params.total_size * sizeof(float), cudaMemcpyDeviceToHost));\n printf(\"saving model params\\n\");\n result = fwrite(params_cpu_buf, sizeof(float), unet->params.total_size, model_file);\n assert(result == unet->params.total_size);\n\n if (model_header[8]) {\n // save AdamW states\n cudaCheck(cudaMemcpy(params_cpu_buf, unet->m_memory, unet->params.total_size * sizeof(float), cudaMemcpyDeviceToHost));\n printf(\"saving AdamW state m\\n\");\n result = fwrite(params_cpu_buf, sizeof(float), unet->params.total_size, model_file);\n assert(result == unet->params.total_size);\n cudaCheck(cudaMemcpy(params_cpu_buf, unet->v_memory, unet->params.total_size * sizeof(float), cudaMemcpyDeviceToHost));\n printf(\"saving AdamW state v\\n\");\n result = fwrite(params_cpu_buf, sizeof(float), unet->params.total_size, model_file);\n assert(result == unet->params.total_size);\n }\n\n if (model_header[9]) {\n // save rngs\n int n_states = diffusion->B * diffusion->img_size;\n void *curand_state_buf = mallocCheck(n_states * sizeof(curandState_t));\n cudaCheck(cudaMemcpy(curand_state_buf, diffusion->curand_states, diffusion->B * diffusion->img_size * sizeof(curandState_t), cudaMemcpyDeviceToHost));\n printf(\"saving rng states\\n\");\n result = fwrite(curand_state_buf, sizeof(curandState_t), n_states, model_file);\n assert(result == n_states);\n free(curand_state_buf);\n }\n\n fcloseCheck(model_file);\n free(params_cpu_buf);\n printf(\"saved model and diffusion states\\n\\n\");\n}\n\n// ------------------------------------------------------------------------------------------------\n// load unet and diffusion states\n\nvoid load_unet_and_diffusion_states(\n Unet* unet, GaussianDiffusion* diffusion, const char* filename\n) {\n FILE *model_file = fopenCheck(filename, \"rb\");\n // read header\n printf(\"\\nloading model header\\n\");\n int model_header[256];\n freadCheck(model_header, sizeof(int), 256, model_file);\n if (model_header[0] != 12345678) {\n fprintf(stderr, \"Invalid magic model file\\n\");\n exit(1);\n }\n int B = model_header[1];\n int C_in = model_header[2];\n int C_model = model_header[3];\n int C_out = model_header[4];\n int H = model_header[5];\n int W = model_header[6];\n int max_period = model_header[7];\n int save_adamw = model_header[8];\n int save_rng = model_header[9];\n\n // TODO: cleanup this code\n int* channel_mult = (int*)mallocCheck(4 * sizeof(int));\n channel_mult[0] = 1;\n channel_mult[1] = 2;\n channel_mult[2] = 3;\n channel_mult[3] = 4;\n\n // setup unet config\n UnetConfig* config = &(unet->config);\n config->B = B;\n config->C_in = C_in;\n config->C_model = C_model;\n config->C_time = C_model * 4;\n config->C_out = C_out;\n config->H = H;\n config->W = W;\n config->block_size = 512;\n config->gn_n_groups = 32;\n config->n_levels = 4;\n config->channel_mult = channel_mult;\n config->HS = 32;\n config->n_res = 2;\n config->max_period = max_period;\n config->att_start_level = 2;\n\n unet_init(unet);\n unet_count_layers(config);\n unet_make_ptrs_and_count_memory(unet);\n malloc_and_point(unet);\n\n // read params\n float* params_mem_cpu = (float*)mallocCheck(unet->params.total_size * sizeof(float));\n printf(\"loading model params\\n\");\n freadCheck(params_mem_cpu, sizeof(float), unet->params.total_size, model_file);\n cudaCheck(cudaMemcpy(unet->params.mem_ptr, params_mem_cpu, unet->params.total_size * sizeof(float), cudaMemcpyHostToDevice));\n\n if (save_adamw) {\n // read AdamW states\n printf(\"loading AdamW states\\n\");\n freadCheck(params_mem_cpu, sizeof(float), unet->params.total_size, model_file);\n cudaCheck(cudaMalloc(&(unet->m_memory), unet->params.total_size * sizeof(float)));\n cudaCheck(cudaMemcpy(unet->m_memory, params_mem_cpu, unet->params.total_size * sizeof(float), cudaMemcpyHostToDevice));\n freadCheck(params_mem_cpu, sizeof(float), unet->params.total_size, model_file);\n cudaCheck(cudaMalloc(&(unet->v_memory), unet->params.total_size * sizeof(float)));\n cudaCheck(cudaMemcpy(unet->v_memory, params_mem_cpu, unet->params.total_size * sizeof(float), cudaMemcpyHostToDevice));\n } else {\n printf(\"not loading AdamW states\\n\");\n unet->m_memory = nullptr;\n unet->v_memory = nullptr;\n }\n\n if (save_rng) {\n // read curand rng states\n printf(\"loading rng states\\n\");\n int n_states = B * C_in * H * W;\n void *curand_state_buf = mallocCheck(n_states * sizeof(curandState_t));\n freadCheck(curand_state_buf, sizeof(curandState_t), n_states, model_file);\n // initialize diffusion curant buffer\n cudaCheck(cudaMalloc(&(diffusion->curand_states), n_states * sizeof(curandState_t)));\n cudaCheck(cudaMemcpy(diffusion->curand_states, curand_state_buf, n_states * sizeof(curandState_t), cudaMemcpyHostToDevice));\n free(curand_state_buf);\n } else {\n printf(\"not loading rng states\\n\");\n diffusion->curand_states = nullptr;\n }\n diffusion_init(diffusion, B, C_in, H, W, 1000);\n\n free(params_mem_cpu);\n fcloseCheck(model_file);\n printf(\"loaded model and diffusion states \\n\\n\");\n}\n\n\n\n// ------------------------------------------------------------------------------------------------\n\ntypedef struct {\n int count;\n float mean_loss;\n} LossCounter;\n\nvoid update_loss_counter(LossCounter* counter, float loss) {\n float old_mean = counter->mean_loss;\n float old_count = (float)counter->count;\n counter->mean_loss = old_mean * old_count / (old_count + 1) + loss / (old_count + 1);\n counter->count++;\n}\n\nvoid clear_loss_counter(LossCounter* counter) {\n counter->count = 0;\n counter->mean_loss = 0.0f;\n}\n\nvoid create_directory(const char *dir) {\n struct stat st = {0};\n if (stat(dir, &st) == -1) {\n mkdir(dir, 0700); // Create directory with permissions 0700\n }\n}\n\nint main(int argc, char **argv) {\n // Default values\n const char *model_weights_file = \"unet_init.bin\";\n const char *data_file = \"data/elephant_train.bin\";\n const char *log_filename = \"log.txt\";\n\n // Parse command line arguments\n for (int i = 1; i < argc; ++i) {\n if (strcmp(argv[i], \"--model_weights\") == 0 && i + 1 < argc) {\n model_weights_file = argv[++i];\n } else if (strcmp(argv[i], \"--data_file\") == 0 && i + 1 < argc) {\n data_file = argv[++i];\n } else if (strcmp(argv[i], \"--log_file\") == 0 && i + 1 < argc) {\n log_filename = argv[++i];\n }\n }\n printf(\"loading model weights from %s\\n\", model_weights_file);\n printf(\"loading data from %s\\n\", data_file);\n\n // set up the device\n int deviceIdx = 0;\n cudaCheck(cudaSetDevice(deviceIdx));\n cudaDeviceProp deviceProp;\n cudaGetDeviceProperties(&deviceProp, deviceIdx);\n printf(\"Device %d: %s\\n\", deviceIdx, deviceProp.name);\n printf(\"device registers per block: %d\\n\", deviceProp.regsPerBlock);\n printf(\"shared mem per block: %ld\\n\", deviceProp.sharedMemPerBlock);\n\n // setup cublas\n cublasHandle_t cublas_handle;\n cublasCheck(cublasCreate(&cublas_handle));\n cublasCheck(cublasSetMathMode(cublas_handle, CUBLAS_DEFAULT_MATH));\n\n Unet unet;\n GaussianDiffusion diffusion;\n load_unet_and_diffusion_states(&unet, &diffusion, model_weights_file);\n\n int B = unet.config.B;\n int C_in = unet.config.C_in;\n int C_out = unet.config.C_out;\n int H = unet.config.H;\n int W = unet.config.W;\n\n DataLoader loader;\n dataloader_init(&loader, data_file, B);\n\n // allocate memory for target separately, since unet can still be run without targets for inference\n cudaCheck(cudaMalloc(&(unet.target), B * C_out * H * W * sizeof(float)));\n float *d_timesteps_buf;\n cudaCheck(cudaMalloc(&d_timesteps_buf, B * sizeof(float)));\n\n\n int n_iters = 100000;\n int n_log_iter = 100;\n int n_save_iter = 10000;\n //int n_iters = 100;\n //int n_log_iter = 5;\n //int n_save_iter = 50;\n\n // wipe existing logs\n FILE* log_file = fopen(log_filename, \"w\");\n fprintf(log_file, \"training unet,\\nmodel file %s,\\ndata filename %s\\n\", model_weights_file, data_file);\n fprintf(log_file, \"starting training for %d iterations\\n\", n_iters);\n fcloseCheck(log_file);\n\n const char *model_save_dir = \"./models/\";\n create_directory(model_save_dir);\n\n LossCounter loss_counter;\n clear_loss_counter(&loss_counter);\n\n printf(\"Training for %d iterations\\n\", n_iters);\n float time_elapsed = 0.0;\n cudaEvent_t start, end;\n cudaCheck(cudaEventCreate(&start));\n cudaCheck(cudaEventCreate(&end));\n cudaCheck(cudaEventRecord(start));\n cudaCheck(cudaEventSynchronize(start));\n for (int iter = 1; iter <= n_iters; iter++) {\n unet_zero_grad(&unet);\n dataloader_next_batch(&loader);\n // copy loaded data into input\n cudaCheck(cudaMemcpy(unet.input, loader.input, B * C_in * H * W * sizeof(float), cudaMemcpyHostToDevice));\n\n // sample timesteps and get timestep embeddings\n sample_timesteps(&diffusion, d_timesteps_buf);\n get_timestep_embeddings(&unet.time_emb_state, d_timesteps_buf, unet.timestep_emb);\n \n // target will be injected random noise\n diffusion_draw_normal(&diffusion, unet.target);\n\n // push input through forward diffusion process by t to get x_t\n diffusion_forward_by_t(&diffusion, unet.input, d_timesteps_buf, unet.target);\n\n unet_forward(cublas_handle, &unet);\n unet_backward(cublas_handle, &unet);\n unet_update(&unet, 0.0001f, 0.9f, 0.999f, 1e-8f, 0.0f, iter);\n\n update_loss_counter(&loss_counter, unet.mean_loss);\n\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(end));\n cudaCheck(cudaEventElapsedTime(&time_elapsed, start, end));\n\n if (iter % n_log_iter == 0) {\n printf(\"step %4d/%d | loss %7.6f | mean loss %7.6f | cur time %.4f s\\n\", iter, n_iters, unet.mean_loss, loss_counter.mean_loss, time_elapsed / 1000);\n FILE *log_file = fopen(log_filename, \"a\");\n fprintf(log_file, \"step %4d/%d | loss %7.6f | mean loss %7.6f | cur time %.4f s\\n\", iter, n_iters, unet.mean_loss, loss_counter.mean_loss, time_elapsed / 1000);\n fclose(log_file);\n clear_loss_counter(&loss_counter);\n }\n\n if (iter % n_save_iter == 0) {\n char save_filename[256];\n sprintf(save_filename, \"./models/model_%d.bin\", iter);\n save_unet_states(&unet, &diffusion, save_filename);\n }\n }\n printf(\"average time per iteration: %f ms\\n\", time_elapsed / n_iters);\n\n cudaCheck(cudaEventDestroy(start));\n cudaCheck(cudaEventDestroy(end));\n cublasCheck(cublasDestroy(cublas_handle));\n cudaCheck(cudaFree(d_timesteps_buf));\n cudaCheck(cudaFree(unet.target));\n\n dataloader_free(&loader);\n diffusion_free(&diffusion);\n free_unet(&unet);\n\n printf(\"success\\n\");\n}", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \"utils.cuh\"\n#include \n#include \n#include \n\n#include \n#include // For mkdir\n#include // For mode_t\n\nnamespace cg = cooperative_groups;\n\n// ------------------------------------------------------------------------------------------------\n// Linear layer\ntypedef struct {\n float* w; // OC, C\n float* b; // OC\n} LinearParams;\n\ntypedef struct {\n float* inp; // N, C\n float* out; // N, OC\n} LinearActs;\n\n\nvoid linear_set_param_ptrs(\n LinearParams* params,\n float* params_memory,\n int C, int OC\n) {\n params->w = params_memory;\n params_memory += OC * C;\n params->b = params_memory;\n}\n\ninline size_t linear_count_params(int C, int OC) {\n return OC * C + OC;\n}\n// ---------------------------------------------\n// kernels\n\n// copy the matmul forward and backward kernels from llm.c\n\n// is there no better way other than just adding bias with a whole separate kernel?\n// this is a highly memory-bound operation, should be fused into the matmul kernel\n// but i can't seem to find a cuBLAS function that does this\n__global__ void add_bias(float* out, const float* bias, int N, int OC) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = idx; i < N * OC; i += stride) {\n int col = i % OC;\n out[i] += bias[col];\n }\n}\n\n// kernel 2 calls cuBLAS, which should be very efficient\nvoid matmul_forward2(hipblasHandle_t cublas_handle,\n float* out,\n const float* inp, const float* weight, const float* bias,\n int N, int C, int OC,\n const int block_size) {\n // for reference API is:\n // cublasStatus_t cublasSgemm(cublasHandle_t handle,\n // cublasOperation_t transa, cublasOperation_t transb,\n // int m, int n, int k,\n // const float *alpha,\n // const float *A, int lda,\n // const float *B, int ldb,\n // const float *beta,\n // float *C, int ldc)\n // for us, inp is (B*T, C), weight is (OC, C), out is (B*T, OC)\n // cuBLAS does C = alpha * A * B + beta * C\n // where A is mxk, B is kxn, C is mxn\n // now, because we use row-major storage, cuBLAS (which is column-major) sees our matrices transposed.\n // algorithmically / in e.g. PyTorch we want to do: out = inp @ weight.T\n // but because cuBLAS is column-major, we actually want to get it to calculate out.T . Mathematically, this is:\n // out.T = weight @ inp.T\n // but again, our variables look transposed, so using the actual weight/inp we have here in this function, this becomes\n // out.T = weight.T @ inp\n // so we need to get cuBLAS to calculate weight.T @ inp (the variables here are the actual ones in this function)\n // => need to call cuBLAS with A = weight, B = inp\n // => need to call cuBLAS with transa = CUBLAS_OP_T, transb = CUBLAS_OP_N\n\n const float alpha = 1.0f;\n const float beta = 0.0f;\n cublasCheck(hipblasSgemm(cublas_handle, HIPBLAS_OP_T, HIPBLAS_OP_N, OC, N, C, &alpha, weight, C, inp, C, &beta, out, OC));\n // and now we still have to add the bias... (ew)\n if (bias != NULL) {\n int grid_size = ceil_div(OC * N, block_size);\n add_bias<<>>(out, bias, N, OC);\n cudaCheck(hipGetLastError());\n }\n}\n\n// use shared memory and coarsening + reductions\n__global__ void matmul_backward_bias_kernel_faster(float* dbias, const float* dout, int N, int OC) {\n extern __shared__ float shared[];\n int o = blockIdx.x; // range [0, OC)\n int tid = threadIdx.x; // range [0, block_size)\n int block_size = blockDim.x;\n const float* x = dout + o;\n // thread coarsening\n double sum = 0.0;\n for (int i = tid; i < N; i += block_size) {\n sum += x[i * OC];\n }\n shared[tid] = (float) sum;\n __syncthreads();\n // reductions\n for (int stride = block_size / 2; stride >= 1; stride /= 2) {\n __syncthreads();\n if (tid < stride) {\n shared[tid] += shared[tid + stride];\n }\n }\n // write the final result (at thread 0) to global memory\n if (tid == 0) {\n dbias[o] = shared[0];\n }\n}\n\n// version1: simple cuBLAS calls\nvoid matmul_backward1(hipblasHandle_t cublas_handle,\n float* dinp, float* dweight, float* dbias,\n float* dout, float* inp, float* weight,\n int N, int C, int OC) {\n float alpha = 1.0f;\n float beta = 0.0f; // note we must use beta = 1.0 so that we do a +=, as we should, because gradients add\n\n // for reference the API is:\n // cublasStatus_t cublasSgemm(cublasHandle_t handle,\n // cublasOperation_t transa, cublasOperation_t transb,\n // int m, int n, int k,\n // const float *alpha,\n // const float *A, int lda,\n // const float *B, int ldb,\n // const float *beta,\n // float *C, int ldc)\n\n // recall the forward pass was calculated with alpha = 1.0f, beta = 0.0f as:\n // cublasSgemm(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, OC, B*T, C, &alpha, weight, C, inp, C, &beta, out, OC);\n\n // backward to input\n cublasCheck(hipblasSgemm(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_N, C, N, OC, &alpha, weight, C, dout, OC, &beta, dinp, C));\n // backward to weight\n cublasCheck(hipblasSgemm(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_T, C, OC, N, &alpha, inp, C, dout, OC, &beta, dweight, C));\n // backward to bias, if given\n if (dbias != NULL) {\n\n // sum over B,T using matrix vector multiplication with cuBLAS\n // for reference this API is:\n // cublasStatus_t cublasSgemv(cublasHandle_t handle, cublasOperation_t trans,\n // int m, int n,\n // const float *alpha,\n // const float *A, int lda,\n // const float *x, int incx,\n // const float *beta,\n // float *y, int incy)\n // dout is (B,T,OC), or in 2D terms (B*T, OC)\n // cublasCheck(cublasSgemv(cublas_handle, CUBLAS_OP_N, B*T, OC, &alpha, dout, B*T, ones, 1, &beta, dbias, 1));\n // cublasCheck(cublasSgemv(cublas_handle, CUBLAS_OP_T, OC, B*T, &alpha, dout, OC, ones, 1, &beta, dbias, 1));\n\n // ugh the above isn't working...\n // let's just do naive calculation for now, fix later\n // const int block_size=128;\n // const int grid_size=(OC + block_size - 1) / block_size;\n // matmul_backward_bias_kernel<<>>(dbias, dout, B, T, OC);\n\n // bit faster\n const int block_size=512;\n dim3 block_dim(block_size);\n dim3 grid_dim(OC);\n size_t shared_mem_size = block_size * sizeof(float);\n matmul_backward_bias_kernel_faster<<>>(dbias, dout, N, OC);\n }\n}\n\n// ------------------------------------------------------------------------------------------------\n// Broadcast tensor layer\n\n// kernel to broadcast a tensor from shape (B, C) to (B, C, H, W)\n// Assuming one thread per pixel\n__global__ void broadcast_last_dims_forward_kernel(\n const float* x,\n float* out,\n int N, int H, int W, int n_bk_per_img\n) {\n int n = blockIdx.x / n_bk_per_img;\n __shared__ float x_sh[1];\n if (threadIdx.x == 0) {\n x_sh[0] = x[n];\n }\n __syncthreads();\n\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < N * H * W) {\n out[idx] = x_sh[0];\n }\n}\n\nvoid broadcast_last_dims_forward(\n const float* x,\n float* out,\n int N, int H, int W, int block_size\n) {\n block_size = min_int(block_size, H * W);\n int n_bk_per_img = ceil_div(H * W, block_size);\n broadcast_last_dims_forward_kernel<<>>(x, out, N, H, W, n_bk_per_img);\n};\n\n// backwards for broadcast\n// use one warp per image\n__global__ void broadcast_last_dims_backward_kernel(\n const float* dout,\n float* dx,\n int N, int H, int W\n) {\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();\n if (idx >= N) { return; }\n\n int img_size = H * W;\n dout += idx * img_size;\n dx += idx;\n\n float sum = 0.0f;\n for (int i = warp.thread_rank(); i < img_size; i += warp.size()) {\n sum += dout[i];\n }\n sum = cg::reduce(warp, sum, cg::plus());\n if (warp.thread_rank() == 0) {\n dx[0] = sum;\n }\n}\n\nvoid broadcast_last_dims_backward(\n const float* dout,\n float* dx,\n int N, int H, int W, int block_size\n) {\n block_size = min_int(block_size, H * W);\n assert(block_size % 32 == 0);\n int warps_per_block = block_size / 32;\n int gridDim = ceil_div(N, warps_per_block);\n broadcast_last_dims_backward_kernel<<>>(dout, dx, N, H, W);\n};\n\n// ------------------------------------------------------------------------------------------------\n// Add layer\n__global__ void add_forward_kernel(\n const float* a, const float* b,\n float* out,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < N) {\n out[idx] = a[idx] + b[idx];\n }\n}\n\nvoid add_forward(\n const float* a, const float* b,\n float* out,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n add_forward_kernel<<>>(a, b, out, N);\n}\n\n// add A and B in place, and store result in B\n__global__ void add_inplace_forward_kernel(\n const float* a, float* b,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < N) {\n b[idx] += a[idx];\n }\n}\n\nvoid add_inplace_forward(\n const float* a, float* b,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n add_inplace_forward_kernel<<>>(a, b, N);\n}\n\n// ------------------------------------------------------------------------------------------------\n// SiLU layer\n\ntypedef struct {\n float* x;\n float* out;\n} SiluActs;\n\n// ----------------------------------------------\n// kernels\n\n__global__ void silu_forward_kernel(\n const float* x,\n float* out,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx >= N) {\n return;\n }\n\n float x_val = x[idx];\n out[idx] = x_val / (1.0f + expf(-x_val));\n}\n\nvoid silu_forward(\n const float* x,\n float* out,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n silu_forward_kernel<<>>(x, out, N);\n}\n\n__global__ void silu_backward_kernel(\n const float* dout, const float* x,\n float* dx,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx >= N) {\n return;\n }\n float out_val = dout[idx];\n float x_val = x[idx];\n float expx = expf(-x_val);\n float grad_silu = (1.0f + x_val * expx / (1.0f + expx)) / (1.0f + expx);\n dx[idx] = out_val * grad_silu;\n}\n\nvoid silu_backward(\n const float* dout, const float* x,\n float* dx,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n silu_backward_kernel<<>>(dout, x, dx, N);\n}\n\n// ------------------------------------------------------------------------------------------------\n// Upsample\ntypedef struct {\n float* x;\n float* out;\n} UpsampleActs;\n\n// one thread per pixel\n__global__ void upsample_forward_kernel(\n const float* x,\n float* out,\n int B, int C, int H, int W\n) {\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int img_size = H * W;\n if (flat_idx >= B * C * img_size) {\n return;\n }\n\n int H_out = H*2;\n int W_out = W*2;\n int img_out_size = H_out * W_out;\n\n int b = flat_idx / (C * img_size);\n int c = (flat_idx / img_size) % C;\n int i = (flat_idx / W) % H;\n int j = flat_idx % W;\n\n // move pointers\n x += b * C * img_size + c * img_size + i * W + j;\n out += b * C * img_out_size + c * img_out_size + 2 * i * W_out + 2 * j;\n\n float x_val = x[0];\n\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n out[ii * W_out + jj] = x_val;\n }\n }\n}\n\nvoid upsample_forward1(\n float* out,\n const float* x,\n int B, int C, int H, int W,\n int block_size\n) {\n int gridDim = ceil_div(B * C * H * W, block_size);\n upsample_forward_kernel<<>>(x, out, B, C, H, W);\n cudaCheck(hipGetLastError());\n}\n\n// number of threads equals the number of pixels\n// in downsampled image\n__global__ void upsample_backward_kernel(\n const float* dout,\n float* dx,\n int B, int C, int H, int W\n) {\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int img_in_size = H * W;\n if (flat_idx >= B * C * img_in_size) {\n return;\n }\n\n int H_out = H*2;\n int W_out = W*2;\n int img_out_size = H_out * W_out;\n\n int b = flat_idx / (C * img_in_size);\n int c = (flat_idx / img_in_size) % C;\n int i = (flat_idx / W) % H;\n int j = flat_idx % W;\n\n // move pointers\n dx += b * C * img_in_size + c * img_in_size + i * W + j;\n dout += b * C * img_out_size + c * img_out_size + 2 * i * W_out + 2 * j;\n\n float dout_sum = 0.0f;\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n dout_sum += dout[ii * W_out + jj];\n }\n }\n dx[0] = dout_sum;\n}\n\nvoid upsample_backward1(\n float* dx,\n const float* dout,\n int B, int C, int H, int W,\n int block_size\n) {\n int gridDim = ceil_div(B * C * H * W, block_size);\n upsample_backward_kernel<<>>(dout, dx, B, C, H, W);\n cudaCheck(hipGetLastError());\n}\n\n// ------------------------------------------------------------------------------------------------\n// Avgpool layer\n\ntypedef struct {\n float* x;\n float* out;\n} AvgpoolActs;\n\n__global__ void avgpool_2d_forward_kernel(\n const float* x, // (B, C, H, W)\n float* out, // (B, C, H/2, W/2)\n int B, int C, int H, int W\n) {\n //printf(\"started kernel\\n\");\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int H_2 = H / 2;\n int W_2 = W / 2;\n int img_out_size = H_2 * W_2;\n if (flat_idx >= B * C * img_out_size) {\n return;\n }\n int img_size = H * W;\n int b = flat_idx / (C * img_out_size);\n int c = (flat_idx / img_out_size) % C;\n int i = (flat_idx / W_2) % H_2;\n int j = flat_idx % W_2;\n\n // move pointers\n x += b * C * img_size + c * img_size + 2 * i * W + 2 * j;\n out += b * C * img_out_size + c * img_out_size + i * W_2 + j;\n\n float sum = 0.0f;\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n sum += x[ii * W + jj];\n }\n }\n out[0] = sum / 4.0f;\n}\n\nvoid avgpool_2d_forward1(\n float* out,\n const float* x,\n int B, int C, int H, int W,\n const int block_size\n) {\n int gridDim = ceil_div(B * C * (H / 2) * (W / 2), block_size);\n avgpool_2d_forward_kernel<<>>(x, out, B, C, H, W);\n cudaCheck(hipGetLastError());\n}\n\n__global__ void avgpool_2d_backward_kernel(\n const float* dout,\n float* dx,\n int B, int C, int H_out, int W_out // H_out and W_out are the shape of the output from forward, i.e. half the original sizes\n) {\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int img_out_size = H_out * W_out;\n if (flat_idx >= B * C * img_out_size) {\n return;\n }\n int H_in = H_out * 2;\n int W_in = W_out * 2;\n int img_in_size = H_in * W_in;\n int b = flat_idx / (C * img_out_size);\n int c = (flat_idx / img_out_size) % C;\n int i = (flat_idx / W_out) % H_out;\n int j = flat_idx % W_out;\n\n // move pointers\n dout += b * C * img_out_size + c * img_out_size + i * W_out + j;\n dx += b * C * img_in_size + c * img_in_size + 2 * i * W_in + 2 * j;\n\n float d = dout[0] / 4.0f;\n\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n dx[ii * W_in + jj] = d;\n }\n }\n}\n\nvoid avgpool_2d_backward1(\n const float* dout,\n float* dx,\n int B, int C, int H, int W,\n const int block_size\n) {\n int H_out = H / 2;\n int W_out = W / 2;\n int gridDim = ceil_div(B * C * H_out * W_out, block_size);\n avgpool_2d_backward_kernel<<>>(dout, dx, B, C, H_out, W_out);\n cudaCheck(hipGetLastError());\n}\n\n// ------------------------------------------------------------------------------------------------\n// Concat channel layer\n// used to concat unet skip connections in the channel dimension\ntypedef struct {\n float* x1;\n float* x2;\n float* out;\n} ConcatChannelActs;\n\n__global__ void concat_channel_forward_kernel(\n const float* x1, const float* x2,\n float* out,\n int B, int C1, int C2, int H, int W\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < B * C1 * H * W) {\n // copy input from x1\n int b = idx / (C1 * H * W);\n int c = (idx / H / W) % C1;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + c * H * W + h * W + w;\n out[out_idx] = x1[idx];\n }\n if (idx < B * C2 * H * W) {\n // copy input from x2\n // move over from x1\n int b = idx / (C2 * H * W);\n int c = (idx / H / W) % C2;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + (C1 + c) * H * W + h * W + w;\n \n out[out_idx] = x2[idx];\n }\n}\n\nvoid concat_channel_forward(\n const float* x1, const float* x2,\n float* out,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n int N = B * max_int(C1, C2) * H * W;\n int n_blk = ceil_div(N, block_size);\n concat_channel_forward_kernel<<>>(x1, x2, out, B, C1, C2, H, W);\n}\n\n__global__ void concat_channel_backward_kernel(\n const float* dout,\n float* dx1, float* dx2,\n int B, int C1, int C2, int H, int W\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < B * C1 * H * W) {\n int b = idx / (C1 * H * W);\n int c = (idx / H / W) % C1;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + c * H * W + h * W + w;\n\n dx1[idx] = dout[out_idx];\n }\n if (idx < B * C2 * H * W) {\n int b = idx / (C2 * H * W);\n int c = (idx / H / W) % C2;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + (C1 + c) * H * W + h * W + w;\n \n dx2[idx] = dout[out_idx];\n }\n}\n\nvoid concat_channel_backward(\n const float* dout,\n float* dx1, float* dx2,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n int N = B * max_int(C1, C2) * H * W;\n int n_blk = ceil_div(N, block_size);\n concat_channel_backward_kernel<<>>(dout, dx1, dx2, B, C1, C2, H, W);\n}\n\n// ------------------------------------------------------------------------------------------------\n// Conv2d 1x1 layer\n\ntemplate \n__global__ void conv2d_k1_forward_kernel(\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C, int H, int W, int O\n) {\n // constants\n const int BC = 16;\n assert(O % BO == 0);\n assert(C % BC == 0);\n assert(H * W % BH == 0);\n // n_threads per block = (BO * BH) / (TO * TH) = 256\n // n_blocks = B * (O / BO) * (H * W / BH)\n // set gridDim to (BxHxW / BH, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int block_col = blockIdx.x % blocks_per_col;\n const int batch = blockIdx.x / blocks_per_col;\n\n const int threads_per_col = BH / TH;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // shared mem for weight and x\n const int weight_s_sz = BO * BC;\n const int x_s_sz = BC * BH;\n __shared__ float w_s[weight_s_sz];\n __shared__ float x_s[x_s_sz];\n\n // local registers for results;\n float thread_results[TO * TH] = {0.0};\n float reg_w[TO] = {0.0};\n float reg_x[TH] = {0.0};\n\n // get offsets to x and weight\n int x_block_offset = batch * C * H * W + block_col * BH;\n int w_block_offset = block_row * BO * C;\n\n for (int c_idx = 0; c_idx < C; c_idx += BC) {\n int x_inner_offset = c_idx * H * W;\n int w_inner_offset = c_idx;\n // load weights to smem\n for (int i = threadIdx.x; i < weight_s_sz; i += blockDim.x) {\n int w_i = i / BC;\n int w_j = i % BC;\n int w_abs_idx = w_block_offset + w_inner_offset + w_i * C + w_j;\n w_s[w_i * BC + w_j] = weight[w_abs_idx];\n }\n for (int i = threadIdx.x; i < x_s_sz; i += blockDim.x) {\n int x_i = i / BH;\n int x_j = i % BH;\n int x_abs_idx = x_block_offset + x_inner_offset + x_i * H * W + x_j;\n x_s[x_i * BH + x_j] = x[x_abs_idx];\n }\n __syncthreads();\n\n\n // calculate thread results\n for (int dot_idx = 0; dot_idx < BC; dot_idx++) {\n for (int i = 0; i < TO; i++) {\n reg_w[i] = w_s[(thread_row * TO + i) * BC + dot_idx];\n }\n for (int j = 0; j < TH; j++) {\n reg_x[j] = x_s[dot_idx * BH + thread_col * TH + j];\n }\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n thread_results[i * TH + j] += reg_w[i] * reg_x[j];\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n // set out and bias directly to thread offset\n int out_block_offset = batch * O * H * W + block_row * BO * H * W + block_col * BH;\n int bias_offset = block_row * BO;\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n int bias_idx = bias_offset + thread_row * TO + i;\n float out_val = thread_results[i * TH + j] + bias[bias_idx];\n int out_idx = out_block_offset + thread_row * TO * H * W + thread_col * TH + i * H * W + j;\n atomicAdd(out + out_idx, out_val);\n }\n }\n}\n\nvoid conv2d_k1_forward2(\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C_in, int H, int W, int C_out\n) {\n cudaCheck(hipMemset(out, 0, B * C_out * H * W * sizeof(float)));\n const int BO = 64;\n const int BH = 64;\n const int TO = 4;\n const int TH = 4;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k1_forward_kernel<<>>(x, weight, bias, out, B, C_in, H, W, C_out);\n}\n\n__global__ void permute_dout_x_kernel(\n const float* x, const float* dout,\n float* x_perm, float* dout_perm,\n int B, int C_in, int C_out, int H, int W\n) {\n int flat_idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n int h = (flat_idx / W) % H;\n int w = flat_idx % W;\n\n if (flat_idx < B * C_in * H * W) {\n // permute x\n int b = flat_idx / (C_in * img_size);\n int c = (flat_idx / img_size) % C_in;\n\n int out_ind = c * B * img_size + b * img_size + h * W + w;\n x_perm[out_ind] = x[flat_idx];\n }\n\n if (flat_idx < B * C_out * H * W) {\n // permute dout\n int b = flat_idx / (C_out * img_size);\n int c = (flat_idx / img_size) % C_out;\n\n int out_ind = c * B * img_size + b * img_size + h * W + w;\n dout_perm[out_ind] = dout[flat_idx];\n }\n}\n\n\n__global__ void permute_dx_kernel(\n const float* dx_perm,\n float* dx,\n int B, int C_in, int H, int W\n) {\n int flat_idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n if (flat_idx < B * C_in * img_size) {\n int c = flat_idx / (B * img_size);\n int b = (flat_idx / img_size) % B;\n int h = (flat_idx / W) % H;\n int w = flat_idx % W;\n int out_ind = b * C_in * img_size + c * img_size + h * W + w;\n dx[out_ind] = dx_perm[flat_idx];\n }\n}\n\n\n// we want to sum the B, H, W dimensions of dout\n// since dout_perm is shaped as (C_out, B, H, W)\n// we will have each block take care of one entry of C_out\n// and the threads in each block will use coorperative_groups to do the sum\n__global__ void dbias_kernel(\n const float* dout_perm, // shape (C_out, B, H, W)\n float* dbias,\n int B, int C_out, int H, int W\n) {\n int bhw = B * H * W;\n extern __shared__ float shared[];\n\n // move dout_perm to the C_out index\n if (threadIdx.x < bhw) {\n dout_perm += blockIdx.x * B * H * W;\n float th_sum = 0.0f;\n for (int i = threadIdx.x; i < bhw; i += blockDim.x) {\n th_sum += dout_perm[i];\n }\n shared[threadIdx.x] = th_sum;\n }\n __syncthreads();\n\n // reduce the shared memory\n for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {\n if (threadIdx.x < stride) {\n shared[threadIdx.x] += shared[threadIdx.x + stride];\n }\n __syncthreads();\n }\n\n if (threadIdx.x == 0) {\n dbias[blockIdx.x] = shared[0];\n }\n}\n\n// uses cublas\nvoid conv2d_k1_backward1(\n hipblasHandle_t cublas_handle,\n const float* dout, const float* x, const float* weight,\n float* dx, float* dweight, float* dbias,\n int B, int C_in, int C_out, int H, int W,\n int block_size\n) {\n // shapes:\n // x, dx: (B, C_in, H, W)\n // weight: (C_out, C_in)\n // bias: (C_out,)\n // dout: (B, C_out, H, W)\n\n // plan (the shapes are raw array shapes, note that cublas interprets the arrays as transposed):\n // 1. transpose x to x_perm of shape (C_in, B*H*W), and dout to dout_perm of shape (C_out, B*H*W)\n // 2. do one sgemm call to get dw = x.T @ dout_perm, and another call to get dx_perm = dout @ w.T\n // 3. dx_perm is of shape (C_in, B*H*W), so call a kernel to permute it\n // 4. reduce to get dbias\n float *x_perm, *dout_perm, *dx_perm;\n cudaCheck(hipMalloc(&x_perm, C_in * B * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&dout_perm, C_out * B * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&dx_perm, C_in * B * H * W * sizeof(float)));\n\n // 1. permute x and dout\n int total_th1 = (C_in > C_out) ? C_in * B * H * W : C_out * B * H * W;\n int n_blk1 = ceil_div(total_th1, block_size);\n permute_dout_x_kernel<<>>(x, dout, x_perm, dout_perm, B, C_in, C_out, H, W);\n\n // 2.a sgemm call for dw = x_perm.T @ dout_perm\n const float alpha = 1.0f;\n const float beta = 0.0f;\n int non_C_dim = B * H * W;\n cublasCheck(hipblasSgemm(cublas_handle, HIPBLAS_OP_T, HIPBLAS_OP_N, C_in, C_out, non_C_dim, &alpha, x_perm, non_C_dim, dout_perm, non_C_dim, &beta, dweight, C_in));\n\n // 2.b sgemm call for dx_perm = dout_perm @ weight.T\n cublasCheck(hipblasSgemm(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_T, non_C_dim, C_in, C_out, &alpha, dout_perm, non_C_dim, weight, C_in, &beta, dx_perm, non_C_dim));\n\n // 3. permute dx_perm from (C_out, B, H, W) to (B, C_out, H, W)\n int total_th2 = C_in * B * H * W;\n int n_blk2 = ceil_div(total_th2, block_size);\n permute_dx_kernel<<>>(dx_perm, dx, B, C_in, H, W);\n\n // 4. reduce dout to get dbias\n dbias_kernel<<>>(dout_perm, dbias, B, C_out, H, W);\n\n // free memory\n cudaCheck(hipFree(x_perm));\n cudaCheck(hipFree(dout_perm));\n cudaCheck(hipFree(dx_perm));\n}\n\n// ------------------------------------------------------------------------------------------------\n// Conv2d 3x3 layer\ntypedef struct {\n float* w;\n float* b;\n} ConvK3Params;\n\ninline void convk3_set_param_ptrs(\n ConvK3Params* params,\n float* params_memory,\n int C_in, int C_out\n) {\n params->w = params_memory;\n params->b = params->w + C_in * C_out * 9;\n}\n\ninline size_t convk3_count_params(int C_in, int C_out) {\n return C_in * C_out * 9 + C_out;\n}\n\ntypedef struct {\n float* inp;\n float* out;\n} ConvK3Acts;\n\n// ---------------------------------------------\n// kernels\n\n// this kernel will roughly follow kernel 5 here: https://github.com/siboehm/SGEMM_CUDA, with a few differences\n// the main operation is convolving the input X with the weights\n// weights are shaped (O, Cx9), and X is shaped (B, C, HxW)\n// We are morally doing a matmul with dimensions O, C, HxW\ntemplate \n__global__ void conv2d_k3_forward_kernel2(\n const float* x, const float* weight, const float* bias,\n float* out,\n const int B, const int C, const int O, const int H, const int W\n) {\n // constants\n const int C9 = C * 9;\n const int weight_size = O * C * 9;\n const int x_size = B * C * H * W;\n // n_threads / block = (BO * BH) / (TO * TH) = 256\n // n_blocks = B * (O / BO) * (H * W / BH)\n // we will set gridDim to (BxHxW / BH, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // for compute step, to determine which of the size TOxTH block the thread is responsible for\n const int threads_per_col = ceil_div(BH, TH);\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // shared mem for loading weights and X\n const int w_s_sz = BO * BC * 9;\n const int w_s_width = BC * 9;\n const int x_s_sz = BC * BH;\n __shared__ float w_s[w_s_sz];\n __shared__ float x_s[x_s_sz];\n\n // registers to store local results\n float thread_results[TO * TH * 9] = {0.0};\n // registers for weights and x\n float reg_w[TO] = {0.0};\n float reg_x[TH] = {0.0};\n\n // calculate offset in x for bounds checking later\n const int x_block_offset = batch * C * H * W + block_col * BH;\n // calculate weight offset for bounds checking\n const int w_block_offset = block_row * BO * C9;\n\n // outer loop over BC size blocks of channel C\n for (int c_idx = 0; c_idx < C; c_idx += BC) {\n // x and weight offsets for this block\n int x_inner_offset = c_idx * H * W;\n int w_inner_offset = c_idx * 9;\n // load weight into smem\n int w_abs_idx;\n for (int i = threadIdx.x; i < w_s_sz; i += blockDim.x) {\n int w_i = i / w_s_width;\n int w_j = i % w_s_width;\n w_abs_idx = w_block_offset + w_inner_offset + w_i * C9 + w_j;\n if (w_abs_idx < weight_size) {\n w_s[w_j * BO + w_i] = weight[w_abs_idx];\n }\n }\n // load x into smem\n int x_abs_idx;\n for (int i = threadIdx.x; i < x_s_sz; i += blockDim.x) {\n int x_i = i / BH;\n int x_j = i % BH;\n x_abs_idx = x_block_offset + x_inner_offset + x_i * H * W + x_j;\n if (x_abs_idx < x_size) {\n x_s[x_i * BH + x_j] = x[x_abs_idx];\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int dot_idx = 0; dot_idx < BC; dot_idx++) {\n // load x_s into registers\n for (int j = 0; j < TH; j++) {\n reg_x[j] = x_s[dot_idx * BH + thread_col * TH + j];\n }\n for (int conv_idx = 0; conv_idx < 9; conv_idx++) {\n for (int i = 0; i < TO; i++) {\n reg_w[i] = w_s[(dot_idx * 9 + conv_idx) * BO + thread_row * TO + i];\n }\n\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n float val = reg_w[i] * reg_x[j];\n thread_results[conv_idx * TO * TH + i * TH + j] += val;\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n // need to be careful here: for each pixel at (h_abs, w_abs) and its product with kernel k (k < 9)\n // There are 2 options\n // 1. the product should be added to some pixel, then do an atomiAdd\n // 2. the product should not contribute to any pixel because it would be out of bound, in this case just skip\n int out_batch_offset = batch * O * H * W;\n for (int j = 0; j < TH; j++) {\n int hw_abs = block_col * BH + thread_col * TH + j;\n int h_abs = hw_abs / W;\n int w_abs = hw_abs % W;\n for (int i = 0; i < TO; i++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target_abs = h_abs + 1 - k_i;\n int w_target_abs = w_abs + 1 - k_j;\n if (h_target_abs >= 0 && h_target_abs < H && w_target_abs >= 0 && w_target_abs < W) {\n int out_target_abs = out_batch_offset + o_abs * H * W + h_target_abs * W + w_target_abs;\n int th_res_idx = k * TO * TH + i * TH + j;\n atomicAdd(out + out_target_abs, thread_results[th_res_idx]);\n }\n }\n int out_abs = out_batch_offset + o_abs * H * W + h_abs * W + w_abs;\n atomicAdd(out + out_abs, bias[o_abs]);\n }\n }\n}\n\n// almost same as 2, but we vectorize the loads and use float4\ntemplate \n__global__ void conv2d_k3_forward_kernel3(\n float* x, float* weight, float* bias,\n float* out,\n const int B, const int C, const int O, const int H, const int W\n) {\n // constants\n assert(TH == 4); // needed for reg_x loads later, so we can do one vectorized load for each thread\n const int C9 = C * 9;\n // n_threads / block = (BO * BH) / (TO * TH) = 256\n // n_blocks = B * (O / BO) * (H * W / BH)\n // we will set gridDim to (BxHxW / BH, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // for compute step, to determine which of the size TOxTH block the thread is responsible for\n const int threads_per_col = ceil_div(BH, TH);\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // shared mem for loading weights and X\n const int w_s_sz = BO * BC * 9;\n const int w_s_width = BC * 9;\n const int x_s_sz = BC * BH;\n __shared__ float w_s[w_s_sz];\n __shared__ float x_s[x_s_sz];\n\n // registers to store local results\n float thread_results[TO * TH * 9] = {0.0};\n // registers for weights and x\n float reg_w[TO] = {0.0};\n float reg_x[TH] = {0.0};\n\n // calculate offset in x for bounds checking later\n const int x_block_offset = batch * C * H * W + block_col * BH;\n // calculate weight offset for bounds checking\n const int w_block_offset = block_row * BO * C9;\n\n // outer loop over BC size blocks of channel C\n for (int c_idx = 0; c_idx < C; c_idx += BC) {\n // x and weight offsets for this block\n int x_inner_offset = c_idx * H * W;\n int w_inner_offset = c_idx * 9;\n // load weight into smem\n for (int i = threadIdx.x; i < w_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (w_s_width / 4);\n int inner_row = i / (w_s_width / 4);\n float4 tmp = \n reinterpret_cast(&weight[w_block_offset + w_inner_offset + inner_row * C9 + inner_col * 4])[0];\n w_s[(inner_col * 4 + 0) * BO + inner_row] = tmp.x;\n w_s[(inner_col * 4 + 1) * BO + inner_row] = tmp.y;\n w_s[(inner_col * 4 + 2) * BO + inner_row] = tmp.z;\n w_s[(inner_col * 4 + 3) * BO + inner_row] = tmp.w;\n }\n // load x into smem\n for (int i = threadIdx.x; i < x_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (BH / 4);\n int inner_row = i / (BH / 4);\n reinterpret_cast(&x_s[inner_row * BH + inner_col * 4])[0] = \n reinterpret_cast(&x[x_block_offset + x_inner_offset + inner_row * H * W + inner_col * 4])[0];\n }\n __syncthreads();\n\n // calculate per thread results\n for (int dot_idx = 0; dot_idx < BC; dot_idx++) {\n // load x_s into registers\n // NOTE: relying on the fact that TH = 4\n reinterpret_cast(reg_x)[0] = reinterpret_cast(&x_s[dot_idx * BH + thread_col * TH])[0];\n for (int conv_idx = 0; conv_idx < 9; conv_idx++) {\n for (int i = 0; i < TO; i++) {\n reg_w[i] = w_s[(dot_idx * 9 + conv_idx) * BO + thread_row * TO + i];\n }\n\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n float val = reg_w[i] * reg_x[j];\n thread_results[conv_idx * TO * TH + i * TH + j] += val;\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n // need to be careful here: for each pixel at (h_abs, w_abs) and its product with kernel k (k < 9)\n // There are 2 options\n // 1. the product should be added to some pixel, then do an atomiAdd\n // 2. the product should not contribute to any pixel because it would be out of bound, in this case just skip\n int out_batch_offset = batch * O * H * W;\n for (int j = 0; j < TH; j++) {\n int hw_abs = block_col * BH + thread_col * TH + j;\n int h_abs = hw_abs / W;\n int w_abs = hw_abs % W;\n for (int i = 0; i < TO; i++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target_abs = h_abs + 1 - k_i;\n int w_target_abs = w_abs + 1 - k_j;\n if (h_target_abs >= 0 && h_target_abs < H && w_target_abs >= 0 && w_target_abs < W) {\n int out_target_abs = out_batch_offset + o_abs * H * W + h_target_abs * W + w_target_abs;\n int th_res_idx = k * TO * TH + i * TH + j;\n atomicAdd(out + out_target_abs, thread_results[th_res_idx]);\n }\n }\n int out_abs = out_batch_offset + o_abs * H * W + h_abs * W + w_abs;\n atomicAdd(out + out_abs, bias[o_abs]);\n }\n }\n}\n\nvoid conv2d_k3_forward3(\n float* x, float* weight, float* bias,\n float* out,\n const int B, const int C_in, const int C_out, const int H, const int W\n) {\n // since the kernel will add to the output, we need to zero it first\n cudaCheck(hipMemset(out, 0, B * C_out * H * W * sizeof(float)));\n\n if (H * W == 64) {\n const int BH = 64;\n const int TH = 4;\n assert(C_out % 64 == 0 && C_in % 8 == 0);\n const int BC = 16;\n const int BO = 64;\n const int TO = 4;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel3<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n } else if (C_in == 3) {\n const int BH = 64;\n assert(H * W % BH == 0);\n const int TH = 4;\n const int BC = 3;\n assert(C_out % 64 == 0);\n const int BO = 64;\n const int TO = 4;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel2<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n } else if (C_out == 3) {\n const int BH = 64 * 2;\n assert(H * W % BH == 0);\n const int TH = 4;\n assert(C_in % 8 == 0);\n const int BC = 8;\n const int BO = 3;\n const int TO = 3;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel2<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n } else {\n const int BH = 64 * 4;\n assert(H * W % BH == 0);\n const int TH = 4;\n assert(C_out % 64 == 0 && C_in % 8 == 0);\n const int BC = 16;\n const int BO = 8;\n const int TO = 2;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel3<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n }\n cudaCheck(hipGetLastError());\n}\n\n// almost the same as conv2d_k3_forward_kernel, except that the compute and the writes\n// are different, because now we want to treat the C dimension as the inner product dimension\ntemplate \n__global__ void dx_backward_kernel1(\n float* dout, float* weight,\n float* dx,\n const int B, const int C, const int O, const int H, const int W\n) {\n const int C9 = C * 9;\n // n_threads per block = (BC * BH) / (TC * TH)\n // set gridDim to (BxHxW / BH, C / TC)\n // find the output block that this thread block is responsible for\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // index for the TCxTH block that the thread is responsible\n const int threads_per_col = BH / TH;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for weights and dout\n const int w_s_width = BC * 9;\n const int w_s_sz = BO * w_s_width;\n const int dout_s_sz = BO * BH;\n __shared__ float w_s[w_s_sz];\n __shared__ float dout_s[dout_s_sz];\n\n // registers for local results\n float thread_results[TC * TH * 9] = {0.0};\n // regs for weights and dout\n float reg_w[TC * 9] = {0.0};\n float reg_dout[TH] = {0.0};\n\n // calculate the block offsets\n const int dout_block_offset = batch * O * H * W + block_col * BH;\n const int w_block_offset = block_row * w_s_width;\n\n // outer loop over BO sizes of channel O\n for (int o_idx = 0; o_idx < O; o_idx += BO) {\n // offsets for this inner block\n int w_inner_offset = o_idx * C9;\n int dout_inner_offset = o_idx * H * W;\n\n // load weights to smem\n if (BC % 4 == 0) {\n // vectorize dram load\n for (int i = threadIdx.x; i < w_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (w_s_width / 4);\n int inner_row = i / (w_s_width / 4);\n reinterpret_cast(&w_s[inner_row * w_s_width + inner_col * 4])[0] = \n reinterpret_cast(&weight[w_block_offset + w_inner_offset + inner_row * C9 + inner_col * 4])[0];\n }\n } else {\n for (int i = threadIdx.x; i < w_s_sz; i += blockDim.x) {\n int inner_col = i % w_s_width;\n int inner_row = i / w_s_width;\n w_s[inner_row * w_s_width + inner_col] = weight[w_block_offset + w_inner_offset + inner_row * C9 + inner_col];\n }\n }\n // load dout\n for (int i = threadIdx.x; i < dout_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (BH / 4);\n int inner_row = i / (BH / 4);\n reinterpret_cast(&dout_s[inner_row * BH + inner_col * 4])[0] = \n reinterpret_cast(&dout[dout_block_offset + dout_inner_offset + inner_row * H * W + inner_col * 4])[0];\n }\n __syncthreads();\n \n\n // calculate per thread results\n for (int dot_idx = 0; dot_idx < BO; dot_idx++) {\n // load dout and w into registers\n for (int j = 0; j < (TH / 4); j++) {\n reinterpret_cast(®_dout[j * 4])[0] = \n reinterpret_cast(&dout_s[dot_idx * BH + thread_col * TH + j * 4])[0];\n }\n // non vectorized load\n //for (int j = 0; j < TH; j++) {\n // reg_dout[j] = dout_s[dot_idx * BH + thread_col * TH + j];\n //}\n if (TC % 4 == 0) {\n // vectorize smem load\n for (int i = 0; i < (TC * 9 / 4); i++) {\n reinterpret_cast(®_w[i * 4])[0] = \n reinterpret_cast(&w_s[dot_idx * w_s_width + thread_row * TC * 9 + i * 4])[0];\n }\n } else {\n for (int i = 0; i < TC * 9; i++) {\n reg_w[i] = w_s[dot_idx * w_s_width + thread_row * TC * 9 + i];\n }\n }\n\n for (int k = 0; k < 9; k++) {\n for (int i = 0; i < TC; i++) {\n for (int j = 0; j < TH; j++) {\n thread_results[k * TC * TH + i * TH + j] += \n reg_w[i * 9 + k] * reg_dout[j];\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n int out_batch_offset = batch * C * H * W;\n for (int j = 0; j < TH; j++) {\n int hw_abs = block_col * BH + thread_col * TH + j;\n int h_abs = hw_abs / W;\n int w_abs = hw_abs % W;\n for (int i = 0; i < TC; i++) {\n int c_abs = block_row * BC + thread_row * TC + i;\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target_abs = h_abs - 1 + k_i;\n int w_target_abs = w_abs - 1 + k_j;\n if (h_target_abs >= 0 && h_target_abs < H && w_target_abs >= 0 && w_target_abs < W) {\n int out_target_abs = out_batch_offset + c_abs * H * W + h_target_abs * W + w_target_abs;\n int th_res_idx = k * TC * TH + i * TH + j;\n atomicAdd(dx + out_target_abs, thread_results[th_res_idx]);\n }\n }\n }\n }\n}\n\nvoid dx_backward(\n float* dout, float* weight,\n float* dx,\n const int B, const int C_in, const int C_out, const int H, const int W\n) {\n if (H < 16) {\n const int BH = 64;\n const int TH = 4;\n const int BC = 16;\n const int TC = 4;\n const int BO = 16;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n } else if (C_in == 3) {\n const int BH = 64 * 4;\n const int TH = 4;\n const int BC = 3;\n const int TC = 3;\n const int BO = 16;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n } else if (C_out == 3) {\n const int BH = 64 * 4;\n const int TH = 4;\n const int BC = 16;\n const int TC = 4;\n const int BO = 3;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n } else {\n const int BH = 64 * 4;\n const int TH = 4;\n const int BC = 16;\n const int TC = 4;\n const int BO = 16;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n }\n}\n\n// same as dweight backward 4, except we also do some of the reduce for dbias\n// so it doesn't have to reduce an entire BxHxW length for each O channel\ntemplate \n__global__ void dweight_dbias_backward_kernel1(\n float* dout, float* x,\n float* dweight_buf, float* dbias_buf,\n const int B, const int C, const int O, const int H, const int W\n) {\n // number of heights we need to load into smem in one go \n const int NH = BW / WSZ;\n // n_threads per block = (BC * BO) / (TO * TC)\n // gridDim will be (B * (W / (WSZ - 1)) * (C / BC), O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_chan = C / BC;\n const int blocks_per_width = W / (WSZ - 1);\n const int block_w = blockIdx.x % blocks_per_width;\n const int block_col = (blockIdx.x / blocks_per_width) % blocks_per_chan;\n const int batch = blockIdx.x / (blocks_per_chan * blocks_per_width);\n const int w_offset = block_w * (WSZ - 1);\n\n // warp indices for loading from global mem\n const int n_warps = ceil_div((int)blockDim.x, 32);\n const int warp_id = threadIdx.x / 32;\n const int lane_id = threadIdx.x % 32;\n // make sure that the width sections can be loaded by a single warp\n assert(WSZ <= 32);\n\n // index for the TO*TC thread block\n const int threads_per_col = BC / TC;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for dout and x\n __shared__ float dout_s[BO * BW];\n __shared__ float x_s[BC * BW];\n\n // registers for local mem and local results\n float thread_results[TO * TC * 9] = {0.0};\n float reg_dout[TO] = {0.0};\n float reg_x[TC] = {0.0};\n float thread_bias_results[TO] = {0.0};\n\n // block offsets\n const int dout_block_offset = batch * O * H * W + block_row * BO * H * W;\n const int x_block_offset = batch * C * H * W + block_col * BC * H * W;\n\n //int debug_idx = 0;\n\n // we will fuse the calculation for all 9 kernel filters\n // which means we will have to reload some of the image widths\n // we can load NW blocks into smem in one go, which means we will be stepping\n // over the H dimension in (NW - 1) steps\n // this is the outer loop over H\n int h_step = (NH < H) ? NH - 1 : H;\n // if NH == H, then we are loading the entire image in one go, so step by H\n for (int h = 0; h < H; h += h_step) {\n // load dout into smem\n int h_offset = h * W;\n int nh = (h < H - (NH - 1)) ? NH : NH - 1;\n int wsz = (w_offset < W - (WSZ - 1)) ? WSZ : WSZ - 1;\n int o_per_warp = ceil_div(BO, n_warps);\n if (lane_id < wsz) {\n for (int warp_o = 0; warp_o < o_per_warp; warp_o++) {\n int inner_o = warp_id * o_per_warp + warp_o;\n if (inner_o < BO) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n dout_s[(inner_h * WSZ + inner_w) * BO + inner_o] = \n dout[dout_block_offset + inner_o * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n int c_per_warp = ceil_div(BC, n_warps);\n if (lane_id < wsz) {\n for (int warp_c = 0; warp_c < c_per_warp; warp_c++) {\n int inner_c = warp_id * c_per_warp + warp_c;\n if (inner_c < BC) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n x_s[(inner_h * WSZ + inner_w) * BC + inner_c] = \n x[x_block_offset + inner_c * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n __syncthreads();\n \n\n // if loading the last block, fill in the last block's buffer to be 0\n if (nh == NH - 1) {\n for (int i = threadIdx.x; i < WSZ * BO; i++) {\n int inner_w = i % WSZ;\n int inner_o = i / WSZ;\n dout_s[(inner_w + nh * WSZ) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < WSZ * BC; i++) {\n int inner_w = i % WSZ;\n int inner_c = i / WSZ;\n\n x_s[(inner_w + nh * WSZ) * BC + inner_c] = 0.0f;\n } \n }\n // if loading the last width section, fill in the final value in the buffer to 0\n if (wsz == WSZ - 1) {\n for (int i = threadIdx.x; i < nh * BO; i++) {\n int inner_o = i % BO;\n int inner_h = i / BO;\n dout_s[(inner_h * WSZ + wsz) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < nh * BC; i++) {\n int inner_c = i % BC;\n int inner_h = i / BC;\n x_s[(inner_h * WSZ + wsz) * BC + inner_c] = 0.0f;\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int x_dot_idx = 0; x_dot_idx < WSZ * nh; x_dot_idx++) {\n int x_h = x_dot_idx / WSZ;\n int x_w = x_dot_idx % WSZ;\n // load TC of the C channels of one image pixel into reg_w\n for (int j = 0; j < TC; j++) {\n reg_x[j] = x_s[x_dot_idx * BC + thread_col * TC + j];\n }\n\n // for each conv filter, load a different row and calculate the outer product\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n // calculate the corresponding (h, w) that should be loaded from dout\n int dout_h = x_h + 1 - k_i;\n int dout_w = x_w + 1 - k_j;\n int dout_h_abs = h + dout_h;\n int dout_w_abs = w_offset + dout_w;\n int dout_dot_idx = -1;\n if (dout_h >= 0 && dout_h < NH && dout_w >= 0 && dout_w < WSZ && dout_w_abs < W && dout_h_abs < H) {\n if (x_h == NH - 1 && k_i == 1 && H > NH) {\n // this logic is a bit tricky: for the last row in smem\n // because it will be loaded as part of the next block as well, and it will be the first row in the next block\n // for kernels 4, 5, 6 we need to only compute them once\n // this statement ensures that we don't compute the 4, 5, 6 kernels for the last row in smem\n // so it will be computed as part of the first row next time\n dout_dot_idx = -1;\n } else if (x_w == WSZ - 1 && k_j == 1 && W > WSZ) {\n // if x_w is at the last pixel, and we have more than one width sections, and we are not are the last section\n // then don't compute 1, 4, 7 kernels and wait for the next section to compute them\n dout_dot_idx = -1;\n } else {\n dout_dot_idx = dout_h * WSZ + dout_w;\n }\n }\n\n // if in bounds of the dout section in smem, load the corresponding pixel from smem\n // and add the products to the thread results\n // otherwise the contribution is 0\n if (dout_dot_idx != -1) {\n for (int i = 0; i < TO; i++) {\n reg_dout[i] = dout_s[dout_dot_idx * BO + thread_row * TO + i];\n }\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n float new_val = reg_dout[i] * reg_x[j];\n thread_results[k * TO * TC + i * TC + j] += new_val;\n }\n }\n // update dbias when k = 4, where dout_h and dout_w match x_h and x_w\n if (k == 4 && block_col == 0) {\n for (int i = 0; i < TO; i++) {\n thread_bias_results[i] += reg_dout[i];\n }\n }\n }\n }\n __syncthreads();\n }\n }\n\n // write out dweight results\n for (int i = 0; i < TO; i++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n for (int j = 0; j < TC; j++) {\n int c_abs = block_col * BC + thread_col * TC + j;\n for (int k = 0; k < 9; k++) {\n dweight_buf[o_abs * C * 9 * B * blocks_per_width + c_abs * 9 * B * blocks_per_width + k * B * blocks_per_width + batch * blocks_per_width + block_w] = thread_results[k * TO * TC + i * TC + j];\n }\n }\n // write out dbias results\n if (block_col == 0) {\n dbias_buf[o_abs * B * blocks_per_width + batch * blocks_per_width + block_w] = thread_bias_results[i];\n }\n }\n}\n\n// C * O blocks\n// 288 (9 * 32) threads per block, use one warp to reduce each batch dim\n__global__ void dweight_reduce_kernel(\n float* dweight_buf, float* dweight,\n const int B, const int C, const int O\n) {\n assert(gridDim.x == O * C);\n int o = blockIdx.x / C;\n int c = blockIdx.x % C;\n // move input pointer\n dweight_buf += o * C * 9 * B + c * 9 * B;\n assert(blockDim.x == 32 * 9);\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n\n float thread_sum = 0.0f;\n for (int i = lane_id; i < B; i += 32) {\n thread_sum += dweight_buf[warp_id * B + i];\n }\n\n // warp reduce\n float warp_sum = cg::reduce(warp, thread_sum, cg::plus{});\n\n if (lane_id == 0) {\n dweight[o * C * 9 + c * 9 + warp_id] = warp_sum;\n }\n}\n\n__global__ void dbias_reduce_kernel1(\n float* dbias_buf, float* dbias,\n const int B\n) {\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n\n dbias_buf += blockIdx.x * B;\n\n int lane_id = threadIdx.x % 32;\n\n float db_thread_sum = 0.0f;\n for (int i = lane_id; i < B; i += 32) {\n db_thread_sum += dbias_buf[i];\n }\n\n float db_warp_sum = cg::reduce(warp, db_thread_sum, cg::plus{});\n\n if (lane_id == 0) {\n dbias[blockIdx.x] = db_warp_sum;\n }\n}\n\nvoid dweight_dbias_backward1(\n float* dout, float* x,\n float* dweight_buf, float* dweight,\n float* dbias_buf, float* dbias,\n const int B, const int C, const int O, const int H, const int W\n) {\n if (H * W == 64) {\n const int BC = 64;\n const int BO = 64;\n const int WSZ = 8;\n const int BW = WSZ * 8;\n const int TO = 4;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(hipDeviceSynchronize());\n cudaCheck(hipGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B, C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B);\n } else if (C == 3) {\n const int BC = 3;\n const int BO = 64;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 3;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(hipDeviceSynchronize());\n cudaCheck(hipGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B * W / (WSZ - 1));\n } else if (O == 3) {\n const int BC = 64;\n const int BO = 3;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 3;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(hipDeviceSynchronize());\n cudaCheck(hipGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B * W / (WSZ - 1));\n } else {\n const int BC = 64;\n const int BO = 64;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(hipDeviceSynchronize());\n cudaCheck(hipGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B * W / (WSZ - 1));\n }\n}\n\nvoid conv2d_k3_backward2(\n float* dout, float* x, float* weight,\n float* dweight_buf, float* dbias_buf,\n float* dx, float* dweight, float* dbias,\n const int B, const int C_in, const int C_out, const int H, const int W\n) {\n // dx currently requires the output buffer to be zeroed\n cudaCheck(hipMemset(dx, 0, B * C_in * H * W * sizeof(float)));\n dx_backward(dout, weight, dx, B, C_in, C_out, H, W);\n dweight_dbias_backward1(dout, x, dweight_buf, dweight, dbias_buf, dbias, B, C_in, C_out, H, W);\n}\n\n// ------------------------------------------------------------------------------------------------\n// GroupNorm\n\ntypedef struct {\n float* w;\n float* b;\n} GroupNormParams;\n\ninline void gn_set_param_ptrs(\n GroupNormParams* params,\n float* params_memory,\n int C\n) {\n params->w = params_memory;\n params->b = params->w + C;\n}\n\ntypedef struct {\n float* x;\n float* out;\n float* mean;\n float* rstd;\n} GroupNormActs;\n\ninline void gn_set_act_ptrs(\n GroupNormActs* acts,\n float* acts_memory,\n int B, int C, int H, int W, int n_groups\n) {\n acts->out = acts_memory;\n acts->mean = acts->out + B * C * H * W;\n acts->rstd = acts->mean + B * n_groups;\n}\n\ntypedef struct {\n float* dx;\n float* dout;\n} GroupNormBackActs;\n\n// ---------------------------------------------\n// kernels\n\n// Essentially taken from llm.c's kernel 5\n// using kernel 5 because for images, each \"channel\" is effectively\n// H * W * group_size, which is quite large\n// we will have one block per group, which means there are B * C / group_size blocks\n__global__ void groupnorm_forward_kernel(\n const float* x, const float* weight, const float* bias,\n float* out, float* mean, float* rstd,\n int B, int C, int img_size, int group_size, int n_groups\n) {\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n __shared__ float shared_sum[32]; // block_size max is 1024 = 32 * 32 warps\n __shared__ float shared_sum2[32]; // warps will be writing into shared memeory after warp-reduce\n int num_warps = blockDim.x / 32;\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n int block_pixels = img_size * group_size;\n // group index\n int g = blockIdx.x % n_groups;\n\n // move pointers\n x += blockIdx.x * img_size * group_size;\n out += blockIdx.x * img_size * group_size;\n // each block will only every acces group_size channels\n weight += g * group_size;\n bias += g * group_size;\n\n float thread_sum = 0.0f;\n float thread_sum2 = 0.0f;\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n float val = x[i];\n thread_sum += val;\n thread_sum2 += val * val;\n }\n\n // warp reduce\n float warp_sum = cg::reduce(warp, thread_sum, cg::plus{});\n float warp_sum2 = cg:: reduce(warp, thread_sum2, cg::plus{});\n // store warp sum into shared memory\n shared_sum[warp_id] = warp_sum;\n shared_sum2[warp_id] = warp_sum2;\n __syncthreads();\n \n // load warp sums from shared memory\n warp_sum = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;\n warp_sum2 = (lane_id < num_warps) ? shared_sum2[lane_id] : 0.0f;\n float block_sum = cg::reduce(warp, warp_sum, cg::plus{});\n float block_sum2 = cg::reduce(warp, warp_sum2, cg::plus{});\n block_sum /= block_pixels;\n block_sum2 /= block_pixels;\n float m = block_sum;\n float var = block_sum2 - m * m;\n float s = rsqrtf(var + 1e-5f);\n if (threadIdx.x == 0 && mean != nullptr) {\n mean[blockIdx.x] = m;\n }\n if (threadIdx.x == 0 && rstd != nullptr) {\n rstd[blockIdx.x] = s;\n }\n\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n int c_mod_group = (i / img_size) % group_size;\n float n = s * (x[i] - m);\n out[i] = n * weight[c_mod_group] + bias[c_mod_group];\n }\n}\n\nvoid groupnorm_forward(\n const float* x, const float* weight, const float* bias,\n float* out, float* mean, float* rstd,\n int B, int C, int H, int W, int n_groups\n) {\n int img_size = H * W;\n int group_size = C / n_groups;\n int n_blocks = B * n_groups;\n int block_size = max_int(min_int(512, img_size * group_size), 32);\n groupnorm_forward_kernel<<>>(\n x, weight, bias, out, mean, rstd, B, C, img_size, group_size, n_groups\n );\n cudaCheck(hipGetLastError());\n}\n\n// most similar to kernel 2 in llm.c\n// not doing any float16 optimizations yet\n// main change is to have each block allocated to a single group\n__global__ void groupnorm_backward_kernel(\n const float* dout, const float* x, const float* mean, const float* rstd, const float* weight,\n float* dx, float* dweight, float* dbias,\n int B, int C, int img_size, int group_size, int n_groups\n) {\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n __shared__ float shared_sum[32]; // block_size max is 1024 = 32 * 32 warps\n __shared__ float shared_sum2[32]; // warps will be writing into shared memeory after warp-reduce\n int num_warps = blockDim.x / 32;\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n int block_pixels = img_size * group_size;\n // group index\n int g = blockIdx.x % n_groups;\n\n // move pointers\n dout += blockIdx.x * img_size * group_size;\n x += blockIdx.x * img_size * group_size;\n dx += blockIdx.x * img_size * group_size;\n weight += g * group_size;\n dweight += g * group_size;\n dbias += g * group_size;\n\n float m_val = mean[blockIdx.x];\n float rstd_val = rstd[blockIdx.x];\n\n\n // calculate the two mean terms in the group dimension\n // first is dout * weight, and second is dout * weight * norm\n // where norm = (x - mean) * rstd\n float w_dout_thread = 0.0f;\n float w_dout_norm_thread = 0.0f;\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n int c_mod_group = (i / img_size) % group_size;\n float cur_w_dout = weight[c_mod_group] * dout[i];\n w_dout_thread += cur_w_dout;\n float norm = (x[i] - m_val) * rstd_val;\n w_dout_norm_thread += cur_w_dout * norm;\n }\n // warp reduce\n float w_dout_warp = cg::reduce(warp, w_dout_thread, cg::plus{});\n float w_dout_norm_warp = cg::reduce(warp, w_dout_norm_thread, cg::plus{});\n // store warp sum in shared mem\n shared_sum[warp_id] = w_dout_warp;\n shared_sum2[warp_id] = w_dout_norm_warp;\n __syncthreads();\n\n // load warp sums from shared memory\n w_dout_warp = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;\n w_dout_norm_warp = (lane_id < num_warps) ? shared_sum2[lane_id] : 0.0f;\n float w_dout_block = cg::reduce(warp, w_dout_warp, cg::plus{});\n float w_dout_norm_block = cg::reduce(warp, w_dout_norm_warp, cg::plus{});\n w_dout_block /= block_pixels;\n w_dout_norm_block /= block_pixels;\n\n // update dx\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n // in bounds of image\n // accumulate dw and db\n float dout_val = dout[i];\n float norm = (x[i] - m_val) * rstd_val;\n\n // update dx\n int c_mod_group = (i / img_size) % group_size;\n float w_dout = weight[c_mod_group] * dout_val;\n dx[i] = (w_dout - w_dout_block - norm * w_dout_norm_block) * rstd_val;\n }\n // update dw and db\n // use different methods when the image size is large or small\n\n // if the image size is larger than the block size\n // loop over the channels and use the whole block on each channel\n // otherwise, assign each warp to a channel\n // in either case image size must be larger than the warp size\n //assert(img_size % warp.size() == 0);\n assert(blockDim.x % warp.size() == 0);\n if (img_size % blockDim.x == 0) {\n for (int c = 0; c < group_size; c++) {\n float dw_thread = 0.0f;\n float db_thread = 0.0f;\n for (int i = threadIdx.x; i < img_size; i += blockDim.x) {\n float dout_val = dout[i];\n db_thread += dout_val;\n float norm = (x[i] - m_val) * rstd_val;\n dw_thread += dout_val * norm;\n }\n\n // move pointers\n dout += img_size;\n x += img_size;\n\n // warp reduce\n float dw_warp = cg::reduce(warp, dw_thread, cg::plus{});\n float db_warp = cg::reduce(warp, db_thread, cg::plus{});\n ////// store warp sum in shared mem\n if (lane_id == 0) {\n shared_sum[warp_id] = dw_warp;\n shared_sum2[warp_id] = db_warp;\n }\n __syncthreads();\n // use the first thread to reduce the shared memory sums and save to global memory\n if (threadIdx.x == 0) {\n float dw_block = 0.0f;\n float db_block = 0.0f;\n for (int i = 0; i < num_warps; i++) {\n dw_block += shared_sum[i];\n db_block += shared_sum2[i];\n }\n atomicAdd(dweight + c, dw_block);\n atomicAdd(dbias + c, db_block);\n }\n }\n } else {\n // if group size is large, need to loop over the group channels with the whole block\n int block_reps = ceil_div(group_size, num_warps);\n for (int br = 0; br < block_reps; br++) {\n float dw_thread = 0.0f;\n float db_thread = 0.0f;\n\n int ch = br * num_warps + warp_id;\n if (ch < group_size) {\n const float* dout_ch = dout + ch * img_size;\n const float* x_ch = x + ch * img_size;\n for (int i = lane_id; i < img_size; i += warp.size()) {\n float dout_val = dout_ch[i];\n db_thread += dout_val;\n float norm = (x_ch[i] - m_val) * rstd_val;\n dw_thread += dout_val * norm;\n }\n \n // warp reduce\n float dw_warp = cg::reduce(warp, dw_thread, cg::plus{});\n float db_warp = cg::reduce(warp, db_thread, cg::plus{});\n // since each warp takes care of an entire image\n // directly store result\n if (lane_id == 0) {\n atomicAdd(dweight + ch, dw_warp);\n atomicAdd(dbias + ch, db_warp);\n }\n }\n }\n }\n}\n\nvoid groupnorm_backward(\n const float* dout, const float* x, const float* mean, const float* rstd, const float* weight,\n float* dx, float* dweight, float* dbias,\n int B, int C, int H, int W, int n_groups\n) {\n int img_size = H * W;\n int group_size = C / n_groups;\n int n_blocks = B * n_groups;\n int block_size = max_int(min_int(512, img_size * group_size), 32 * group_size);\n groupnorm_backward_kernel<<>>(\n dout, x, mean, rstd, weight, dx, dweight, dbias, B, C, img_size, group_size, n_groups\n );\n cudaCheck(hipGetLastError());\n}\n// ------------------------------------------------------------------------------------------------\n// ResBlock\n\n#define NUM_RES_PARAM_TENSORS 12\ntypedef struct {\n float* gn1_w; // C\n float* gn1_b; // C\n float* cv3_1_w; // C_out, C, 3, 3\n float* cv3_1_b; // C_out\n float* l_emb_w; // C_out, C_emb\n float* l_emb_b; // C_out\n float* gn2_w; // C_out\n float* gn2_b; // C_out\n float* cv3_2_w; // C_out, C_out, 3, 3\n float* cv3_2_b; // C_out\n float* res_cv1_w; // C_out, C\n float* res_cv1_b; // C_out\n size_t param_sizes[NUM_RES_PARAM_TENSORS];\n size_t n_params;\n} ResBlockParameters;\n\n#define NUM_RES_ACT_TENSORS 18 // not counting input and emb, which will come from outs of previous layers\ntypedef struct {\n float* gn1; // B, C, H, W\n float* gn1_mean; // B, n_groups\n float* gn1_rstd; // B, n_groups\n float* silu1; // B, C, H, W\n float* ud_h; // B, C, H_out, W_out\n float* ud_x; // B, C, H_out, W_out\n float* cv3_1; // B, C_out, H_out, W_out\n float* silu_emb; // B, C_emb\n float* l_emb; // B, C_out\n float* broad_emb; // B, C_out, H_out, W_out\n float* add1; // B, C_out, H_out, W_out\n float* gn2; // B, C_out, H_out, W_out\n float* gn2_mean; // B, n_groups\n float* gn2_rstd; // B, n_groups\n float* silu2; // B, C_out, H_out, W_out\n float* cv3_2; // B, C_out, H_out, W_out\n float* res_cv1; // B, C_out, H_out, W_out\n float* add2; // B, C_out, H_out, W_out\n size_t act_sizes[NUM_RES_ACT_TENSORS];\n size_t n_acts;\n float* input; // B, C, H, W\n float* emb; // B, C_emb\n} ResBlockActivations;\n\n\n// the last two backward tensors, dx and demb\n// we don't need to allocate extra space for these\n// instead we will just overwrite forward pass activations acts->input and acts->emb\n#define NUM_RES_BACKWARD_TENSORS 7\ntypedef struct {\n float* buf_BCemb;\n float* buf_BCHoWo;\n float* buf1_BCHW;\n float* buf2_BCHW;\n float* dout; // B, C_out, H_out, W_out\n float* dweight_buf;\n float* dbias_buf;\n float* dx; // B, C, H, W\n float* demb; // B, C_emb\n size_t back_sizes[NUM_RES_BACKWARD_TENSORS];\n size_t n_backs;\n} ResBlockBackwardActivations;\n\nvoid resblock_count_params(\n ResBlockParameters* params,\n int C, int C_emb, int C_out, int B, int H, int W, int up, int down, int gn_n_groups\n) {\n int H_out = H;\n int W_out = W;\n if (up) {\n H_out *= 2;\n W_out *= 2;\n } else if (down) {\n H_out /= 2;\n W_out /= 2;\n }\n // fill and count param sizes\n size_t* param_sizes = params->param_sizes;\n param_sizes[0] = C; // gn1_w\n param_sizes[1] = C; // gn1_b\n param_sizes[2] = C_out * C * 9; // cv3_1_w\n param_sizes[3] = C_out; // cv3_1_b\n param_sizes[4] = C_emb * C_out; // l_emb_w\n param_sizes[5] = C_out; // l_emb_b\n param_sizes[6] = C_out; // gn2_w\n param_sizes[7] = C_out; // gn2_b\n param_sizes[8] = C_out * C_out * 9; // cv3_2_w\n param_sizes[9] = C_out; // cv3_2_b\n param_sizes[10] = C_out * C; // res_cv1_w\n param_sizes[11] = C_out; // res_cv1_b\n if (C_out != C) {\n param_sizes[10] = C_out * C; // res_cv1_w\n param_sizes[11] = C_out; // res_cv1_b\n } else {\n param_sizes[10] = 0;\n param_sizes[11] = 0;\n }\n\n // count number of params\n size_t num_parameters = 0;\n for (size_t i = 0; i < NUM_RES_PARAM_TENSORS; i++) {\n num_parameters += param_sizes[i];\n }\n params->n_params = num_parameters;\n}\n\nvoid resblock_count_acts(\n ResBlockActivations* acts,\n int C, int C_emb, int C_out, int B, int H, int W, int up, int down, int gn_n_groups\n) {\n int H_out = H;\n int W_out = W;\n if (up) {\n H_out *= 2;\n W_out *= 2;\n } else if (down) {\n H_out /= 2;\n W_out /= 2;\n }\n // fill and count activation sizes\n size_t* acts_sizes = acts->act_sizes;\n acts_sizes[0] = B * C * H * W; // gn1\n acts_sizes[1] = B * gn_n_groups; // gn1_mean\n acts_sizes[2] = B * gn_n_groups; // gn1_rstd\n acts_sizes[3] = B * C * H * W; // silu1\n acts_sizes[4] = B * C * H_out * W_out; // ud_h\n acts_sizes[5] = B * C * H_out * W_out; // ud_x\n acts_sizes[6] = B * C_out * H_out * W_out; // cv3_1\n acts_sizes[7] = B * C_emb; // silu_emb\n acts_sizes[8] = B * C_out; // l_emb\n acts_sizes[9] = B * C_out * H_out * W_out; // broad_emb\n acts_sizes[10] = B * C_out * H_out * W_out; // add1\n acts_sizes[11] = B * C_out * H_out * W_out; // gn2\n acts_sizes[12] = B * gn_n_groups; // gn2_mean\n acts_sizes[13] = B * gn_n_groups; // gn2_rstd\n acts_sizes[14] = B * C_out * H_out * W_out; // silu2\n acts_sizes[15] = B * C_out * H_out * W_out; // cv3_2\n acts_sizes[16] = (C_out == C) ? 0 : B * C_out * H_out * W_out; // res_cv1\n acts_sizes[17] = B * C_out * H_out * W_out; // add2\n\n size_t num_acts_params = 0;\n for (size_t i = 0; i < NUM_RES_ACT_TENSORS; i++) {\n num_acts_params += acts_sizes[i];\n }\n acts->n_acts = num_acts_params;\n}\n\nvoid set_resblock_params_ptrs(\n int C,\n int C_out,\n ResBlockParameters *params,\n float* params_memory_gpu\n) {\n float** ptrs[] = {\n ¶ms->gn1_w, ¶ms->gn1_b, ¶ms->cv3_1_w, ¶ms->cv3_1_b,\n ¶ms->l_emb_w, ¶ms->l_emb_b, ¶ms->gn2_w, ¶ms->gn2_b,\n ¶ms->cv3_2_w, ¶ms->cv3_2_b, ¶ms->res_cv1_w, ¶ms->res_cv1_b\n };\n int num_param_tensors = NUM_RES_PARAM_TENSORS;\n if (C_out == C) {\n num_param_tensors -= 2;\n }\n float* params_ptr = params_memory_gpu;\n for (int i = 0; i < num_param_tensors; i++) {\n *(ptrs[i]) = params_ptr;\n params_ptr += params->param_sizes[i];\n }\n if (C_out == C) {\n params->res_cv1_w = nullptr;\n params->res_cv1_b = nullptr;\n }\n}\n\nvoid set_resblock_acts_ptrs(\n int C,\n int C_out,\n ResBlockActivations *acts,\n float* acts_memory\n) {\n float** acts_ptrs[] = {\n &acts->gn1, &acts->gn1_mean, &acts->gn1_rstd, &acts->silu1,\n &acts->ud_h, &acts->ud_x, &acts->cv3_1, &acts->silu_emb, &acts->l_emb, &acts->broad_emb,\n &acts->add1, &acts->gn2, &acts->gn2_mean, &acts->gn2_rstd, &acts->silu2, &acts->cv3_2,\n &acts->res_cv1, &acts->add2\n };\n float* acts_ptr = acts_memory;\n for (int i = 0; i < NUM_RES_ACT_TENSORS - 2; i++) {\n *(acts_ptrs[i]) = acts_ptr;\n acts_ptr += acts->act_sizes[i];\n }\n\n if (C_out == C) {\n acts->res_cv1 = nullptr;\n acts->add2 = acts_ptr;\n } else {\n acts->res_cv1 = acts_ptr;\n acts_ptr += acts->act_sizes[NUM_RES_ACT_TENSORS - 2];\n acts->add2 = acts_ptr;\n }\n}\n\n// assume all memory is allocated, run the forward pass\nvoid resblock_forward(\n hipblasHandle_t cublas_handle,\n int C,\n int C_emb,\n int C_out,\n int B,\n int H,\n int W,\n int block_size,\n int up,\n int down,\n int gn_n_groups,\n ResBlockParameters *params,\n ResBlockActivations *acts\n) {\n int H_out = H;\n int W_out = W;\n if (up) {\n H_out *= 2;\n W_out *= 2;\n } else if (down) {\n H_out /= 2;\n W_out /= 2;\n }\n // up to first 3x3 conv in BigGAN res block\n // note: residual stream needs a convolution at the end if C != C_out\n groupnorm_forward(acts->input, params->gn1_w, params->gn1_b, acts->gn1, acts->gn1_mean, acts->gn1_rstd,\n B, C, H, W, gn_n_groups);\n cudaCheck(hipGetLastError());\n silu_forward(acts->gn1, acts->silu1, B * C * H * W, block_size);\n cudaCheck(hipGetLastError());\n if (up) {\n upsample_forward1(acts->ud_h, acts->silu1, B, C, H, W, block_size);\n upsample_forward1(acts->ud_x, acts->input, B, C, H, W, block_size);\n } else if (down) {\n avgpool_2d_forward1(acts->ud_h, acts->silu1, B, C, H, W, block_size);\n avgpool_2d_forward1(acts->ud_x, acts->input, B, C, H, W, block_size);\n } else {\n cudaCheck(hipMemcpy(acts->ud_h, acts->silu1, B * C * H * W * sizeof(float), hipMemcpyDeviceToDevice));\n cudaCheck(hipMemcpy(acts->ud_x, acts->input, B * C * H * W * sizeof(float), hipMemcpyDeviceToDevice));\n }\n cudaCheck(hipGetLastError());\n conv2d_k3_forward3(acts->ud_h, params->cv3_1_w, params->cv3_1_b, acts->cv3_1, B, C, C_out, H_out, W_out);\n cudaCheck(hipGetLastError());\n\n // project and add embedding layer\n silu_forward(acts->emb, acts->silu_emb, B * C_emb, block_size);\n cudaCheck(hipGetLastError());\n matmul_forward2(cublas_handle, acts->l_emb, acts->silu_emb, params->l_emb_w, params->l_emb_b, B, C_emb, C_out, block_size);\n cudaCheck(hipGetLastError());\n broadcast_last_dims_forward(acts->l_emb, acts->broad_emb, B * C_out, H_out, W_out, block_size);\n cudaCheck(hipGetLastError());\n add_forward(acts->cv3_1, acts->broad_emb, acts->add1, B * C_out * H_out * W_out, block_size);\n cudaCheck(hipGetLastError());\n\n // second group of norm -> silu -> 3x3 conv\n groupnorm_forward(acts->add1, params->gn2_w, params->gn2_b, acts->gn2, acts->gn2_mean, acts->gn2_rstd,\n B, C_out, H_out, W_out, gn_n_groups);\n cudaCheck(hipGetLastError());\n silu_forward(acts->gn2, acts->silu2, B * C_out * H_out * W_out, block_size);\n cudaCheck(hipGetLastError());\n conv2d_k3_forward3(acts->silu2, params->cv3_2_w, params->cv3_2_b, acts->cv3_2, B, C_out, C_out, H_out, W_out);\n cudaCheck(hipGetLastError());\n\n // add residual (change channels with 1x1 conv if neccessary)\n if (C_out == C) {\n add_forward(acts->cv3_2, acts->ud_x, acts->add2, B * C_out * H_out * W_out, block_size);\n cudaCheck(hipGetLastError());\n } else {\n conv2d_k1_forward2(acts->ud_x, params->res_cv1_w, params->res_cv1_b, acts->res_cv1, B, C, H_out, W_out, C_out);\n cudaCheck(hipGetLastError());\n add_forward(acts->cv3_2, acts->res_cv1, acts->add2, B * C_out * H_out * W_out, block_size);\n cudaCheck(hipGetLastError());\n }\n}\n\nvoid resblock_backward(\n hipblasHandle_t cublas_handle,\n int C,\n int C_emb,\n int C_out,\n int B,\n int H,\n int W,\n int block_size,\n int up,\n int down,\n int gn_n_groups,\n ResBlockParameters *params,\n ResBlockParameters *grads,\n ResBlockActivations *acts,\n ResBlockBackwardActivations *back_acts\n) {\n int H_out = H;\n int W_out = W;\n if (up) {\n H_out *= 2;\n W_out *= 2;\n } else if (down) {\n H_out /= 2;\n W_out /= 2;\n }\n // Note: we're using the activation of the final conv layer as a buffer to store\n // backward pass gradients, because we don't need it in the backward pass anymore.\n float *demb = back_acts->demb;\n float *dx = back_acts->dx;\n float* dout = back_acts->dout;\n float *res_dx = back_acts->buf_BCHoWo;\n float* buf1_BCoHoWo = acts->cv3_2;\n float* buf2_BCoHoWo = acts->add2;\n cudaCheck(hipMemset(buf1_BCoHoWo, 0, B * C_out * H_out * W_out * sizeof(float)));\n cudaCheck(hipMemset(buf2_BCoHoWo, 0, B * C_out * H_out * W_out * sizeof(float)));\n // Note: we don't need to run a separate kernel to backprop through an add operation\n // we can just share the dout buffer between the two inputs\n if (C_out != C) {\n conv2d_k1_backward1(cublas_handle, dout, acts->ud_x, params->res_cv1_w, res_dx, grads->res_cv1_w, grads->res_cv1_b, B, C, C_out, H_out, W_out, block_size);\n cudaCheck(hipGetLastError());\n } else {\n // point res_dx to dout\n res_dx = dout;\n }\n // NOTE: after backprop through cv1 above, we can use ud_x as a (B, C, H_out, W_out) buffer\n float* buf1_BCHoWo = acts->ud_x;\n\n // second group of norm -> silu -> 3x3 conv\n conv2d_k3_backward2(dout, acts->silu2, params->cv3_2_w, back_acts->dweight_buf, back_acts->dbias_buf, buf1_BCoHoWo, grads->cv3_2_w, grads->cv3_2_b, B, C_out, C_out, H_out, W_out);\n cudaCheck(hipGetLastError());\n silu_backward(buf1_BCoHoWo, acts->gn2, buf2_BCoHoWo, B * C_out * H_out * W_out, block_size);\n cudaCheck(hipGetLastError());\n groupnorm_backward(buf2_BCoHoWo, acts->add1, acts->gn2_mean, acts->gn2_rstd, params->gn2_w, buf1_BCoHoWo, grads->gn2_w, grads->gn2_b, B, C_out, H_out, W_out, gn_n_groups);\n cudaCheck(hipGetLastError());\n\n // emb layer\n // note at this point we can also use the last layer activations as buffers\n broadcast_last_dims_backward(buf1_BCoHoWo, acts->l_emb, B * C_out, H_out, W_out, block_size);\n cudaCheck(hipGetLastError());\n matmul_backward1(cublas_handle, back_acts->buf_BCemb, grads->l_emb_w, grads->l_emb_b, acts->l_emb, acts->silu_emb, params->l_emb_w, B, C_emb, C_out);\n cudaCheck(hipGetLastError());\n silu_backward(back_acts->buf_BCemb, acts->emb, demb, B * C_emb, block_size);\n cudaCheck(hipGetLastError());\n\n // first layers\n conv2d_k3_backward2(buf1_BCoHoWo, acts->ud_h, params->cv3_1_w, back_acts->dweight_buf, back_acts->dbias_buf, buf1_BCHoWo, grads->cv3_1_w, grads->cv3_1_b, B, C, C_out, H_out, W_out);\n cudaCheck(hipGetLastError());\n // use acts->ud_x as a buffer for dx\n float* buf1_BCHW = back_acts->buf1_BCHW;\n float* buf2_BCHW = back_acts->buf2_BCHW;\n if (up) {\n upsample_backward1(buf1_BCHW, buf1_BCHoWo, B, C, H, W, block_size);\n cudaCheck(hipGetLastError());\n upsample_backward1(buf2_BCHW, res_dx, B, C, H, W, block_size);\n cudaCheck(hipGetLastError());\n } else if (down) {\n avgpool_2d_backward1(buf1_BCHoWo, buf1_BCHW, B, C, H, W, block_size);\n cudaCheck(hipGetLastError());\n avgpool_2d_backward1(res_dx, buf2_BCHW, B, C, H, W, block_size);\n cudaCheck(hipGetLastError());\n } else {\n cudaCheck(hipMemcpy(buf1_BCHW, buf1_BCHoWo, B * C * H * W * sizeof(float), hipMemcpyDeviceToDevice));\n cudaCheck(hipGetLastError());\n cudaCheck(hipMemcpy(buf2_BCHW, res_dx, B * C * H * W * sizeof(float), hipMemcpyDeviceToDevice));\n cudaCheck(hipGetLastError());\n }\n // use silu as the dx buffer\n silu_backward(buf1_BCHW, acts->gn1, acts->silu1, B * C * H * W, block_size);\n cudaCheck(hipGetLastError());\n groupnorm_backward(acts->silu1, acts->input, acts->gn1_mean, acts->gn1_rstd, params->gn1_w, buf1_BCHW, grads->gn1_w, grads->gn1_b, B, C, H, W, gn_n_groups);\n cudaCheck(hipGetLastError());\n\n add_forward(buf1_BCHW, buf2_BCHW, dx, B * C * H * W, block_size);\n cudaCheck(hipGetLastError());\n}\n\n// ------------------------------------------------------------------------------------------------\n// self attention block\n// assumes you have Q, K, V\n__global__ void permute_kernel(float* q, float* k, float* v,\n const float* inp,\n int B, int N, int NH, int d) {\n // okay so now, this kernel wants Q,K,V to all be of shape (B, NH, N, d)\n // but instead, we have a single tensor QKV (inp) of shape (B, N, 3, NH, d)\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Q[b][nh_][n][d_] = inp[b][n][0][nh_][d_]\n\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n\n int inp_idx = \\\n (b * N * 3 * NH * d)\n + (n * 3 * NH * d)\n + (0 * NH * d)\n + (nh_ * d)\n + d_;\n\n q[idx] = inp[inp_idx];\n k[idx] = inp[inp_idx + NH * d];\n v[idx] = inp[inp_idx + 2 * (NH * d)];\n }\n}\n\n__global__ void scale_kernel(float* inp, float scale, int B, int NH, int T) {\n // scales the pre-softmax attention scores by scale\n // we don't want causal attention, so we don't set any values to infinity\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * NH * T * T) {\n inp[idx] *= scale;\n }\n}\n\n// warp-level reduction for finding the maximum value\n__device__ float warpReduceMax(float val) {\n for (int offset = 16; offset > 0; offset /= 2) {\n val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));\n }\n return val;\n}\n\n// warp-level reduction for summing values\n__device__ float warpReduceSum(float val) {\n for (int offset = 16; offset > 0; offset /= 2) {\n val += __shfl_down_sync(0xFFFFFFFF, val, offset);\n }\n return val;\n}\n\n__global__ void softmax_forward_kernel4(float* out, const float* inp, int N, int C) {\n // out is (N, C) just like inp. Each row of inp will get softmaxed.\n // same as kernel3, but can handle any block size (multiple of 32)\n // each row of C elements is handled by block_size threads\n // furthermore, each block_size threads get executed in warps of 32 threads\n\n // special reduction operations warpReduceMax/warpReduceSum are used for intra-warp reductions\n // shared memory is used for inter-warp reduction\n extern __shared__ float shared[];\n int idx = blockIdx.x;\n int tid = threadIdx.x;\n int warpId = threadIdx.x / 32; // warp index within a block\n int laneId = threadIdx.x % 32; // thread index within a warp\n\n // the number of warps per block. recall that blockDim.x is block_size\n int warpsPerBlock = blockDim.x / 32;\n\n // shared[] must be allocated to have 2 * warpsPerBlock elements\n // first half for max values, the second half for sum values\n float* maxvals = shared;\n float* sumvals = &shared[warpsPerBlock];\n\n // one row of inp, i.e. inp[idx, :] of shape (C,)\n const float* x = inp + idx * C;\n\n // first, thread coarsening by directly accessing global memory in series\n float maxval = -INFINITY;\n for (int i = tid; i < C; i += blockDim.x) {\n maxval = fmaxf(maxval, x[i]);\n }\n // now within-warp reductions for maxval\n maxval = warpReduceMax(maxval);\n\n // the 0th thread of each warp writes the maxval of that warp to shared memory\n if (laneId == 0) maxvals[warpId] = maxval;\n __syncthreads();\n\n // now the 0th thread reduces the maxvals in shared memory, i.e. across warps\n if (tid == 0) {\n float val = maxvals[tid];\n for (int i = 1; i < warpsPerBlock; i++) {\n val = fmaxf(val, maxvals[i]);\n }\n // store the final max in the first position\n maxvals[0] = val;\n }\n __syncthreads();\n // broadcast the max to all threads\n float offset = maxvals[0];\n\n // compute expf and write the result to global memory\n for (int i = tid; i < C; i += blockDim.x) {\n // subtract max for numerical stability\n out[idx * C + i] = expf(x[i] - offset);\n }\n\n // okay now we calculated exp(x - max(x))\n // step 2: sum all the values and divide by the sum\n\n // thread coarsening for sum\n x = out + idx * C;\n float sumval = 0.0f;\n for (int i = tid; i < C; i += blockDim.x) {\n sumval += x[i];\n }\n // within-warp reduction for sumval\n sumval = warpReduceSum(sumval);\n\n // write sumval to shared memory\n if (laneId == 0) sumvals[warpId] = sumval;\n __syncthreads();\n\n // inter-thread reduction of sum\n if (tid == 0) {\n float val = sumvals[tid];\n for (int i = 1; i < warpsPerBlock; ++i) {\n val += sumvals[i];\n }\n sumvals[0] = val;\n }\n __syncthreads();\n // broadcast the sum to all threads\n float sum = sumvals[0];\n\n // divide the whole row by the sum\n for (int i = tid; i < C; i += blockDim.x) {\n out[idx * C + i] = x[i] / sum;\n }\n}\n\n__global__ void unpermute_kernel(const float* inp, float *out, int B, int N, int NH, int d) {\n // out has shape (B, nh, N, d) but we need to unpermute it to (B, N, nh, d)\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n // out[b][n][nh_][d_] <- inp[b][nh_][n][d_]\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n\n int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_;\n out[other_idx] = inp[idx];\n }\n}\n\n// a simple attention forward function taken directly from llm.c kernel 3\nvoid attention_forward1\n(\n hipblasHandle_t cublas_handle,\n float* out, float* qkvr, float* preatt, float* att,\n float* inp,\n int B, int T, int C, int NH,\n const int block_size\n) {\n // note: resure inp as a scratch buffer because it is not used in backward\n // inp is (B, T, 3C) QKV\n // preatt, att are (B, NH, T, T)\n // output is (B, T, C)\n int HS = C / NH; // head size\n\n // permute and separate inp from (B, T, 3, NH, HS) to 3X (B, NH, T, HS)\n float *q, *k, *v;\n q = qkvr + 0 * B * T * C;\n k = qkvr + 1 * B * T * C;\n v = qkvr + 2 * B * T * C;\n int total_threads = B * NH * T * HS;\n int num_blocks = ceil_div(total_threads, block_size);\n permute_kernel<<>>(q, k, v, inp, B, T, NH, HS);\n\n // batched matrix multiply with cuBLAS\n const float alpha = 1.0f;\n const float beta = 0.0f;\n cublasCheck(hipblasSgemmStridedBatched(cublas_handle,\n HIPBLAS_OP_T, HIPBLAS_OP_N,\n T, T, HS,\n &alpha,\n k, HS, T * HS,\n q, HS, T * HS,\n &beta,\n preatt, T, T * T,\n B * NH));\n\n // multiply all elements of preatt elementwise by scale\n float scale = 1.0f / sqrtf(HS);\n total_threads = B * NH * T * T;\n num_blocks = ceil_div(total_threads, block_size);\n scale_kernel<<>>(preatt, scale, B, NH, T);\n\n // softmax. preatt is (B, NH, T, T) but we view it as (B * NH * T, T) and use the softmax kernel\n int softmax_block_size = 256;\n int grid_size = B * NH * T;\n size_t shared_mem_size = 2 * softmax_block_size / 32 * sizeof(float);\n softmax_forward_kernel4<<>>(att, preatt, B * NH * T, T);\n\n // new approach: first cuBLAS another batched matmul\n // y = att @ v # (B, nh, T, T) @ (B, nh, T, hs) -> (B, nh, T, hs)\n float* vaccum = inp;\n cublasCheck(hipblasSgemmStridedBatched(cublas_handle,\n HIPBLAS_OP_N, HIPBLAS_OP_N,\n HS, T, T,\n &alpha,\n v, HS, T * HS,\n att, T, T * T,\n &beta,\n vaccum, HS, T * HS,\n B * NH));\n\n // now unpermute\n // y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side\n num_blocks = ceil_div(B * T * C, block_size);\n unpermute_kernel<<>>(vaccum, out, B, T, NH, HS);\n}\n\n__global__ void permute_kernel_backward(float* dinp,\n const float* dq, const float* dk, const float* dv,\n int B, int N, int NH, int d) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n\n int inp_idx = (b * N * 3 * NH * d) + (n * 3 * NH * d) + (0 * NH * d) + (nh_ * d) + d_;\n dinp[inp_idx] = dq[idx];\n dinp[inp_idx + NH * d] = dk[idx];\n dinp[inp_idx + 2 * (NH * d)] = dv[idx];\n }\n}\n\n__global__ void unpermute_kernel_backward(float* dinp, const float *dout, int B, int N, int NH, int d) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_;\n dinp[idx] = dout[other_idx];\n }\n}\n\n// slight edit from the softmax_autoregressive_backward_kernel from train_gpt2_fp32\n// our attention is no longer autoregressive\n__global__ void softmax_backward_kernel(float* dpreatt, const float* datt, const float* att,\n int B, int T, int C, float scale) {\n constexpr const int BlockSize = 256;\n constexpr int T_per_block = 4;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n __shared__ float block_acc[32];\n\n int idx = blockIdx.y;\n // go through blocks in reverse order, so the slowest block starts first\n int t0 = T - 1 - T_per_block*blockIdx.x;\n\n att += idx * T * T;\n datt += idx * T * T;\n dpreatt += idx * T * T;\n\n if (warp.meta_group_rank() == 0) {\n block_acc[warp.thread_rank()] = 0;\n }\n\n for(int to = 0; to < T_per_block; ++to) {\n int t = t0 - to;\n if(t < 0) return;\n const float* att_bth = att + t * T;\n const float* datt_bth = datt + t * T;\n float* dpreatt_bth = dpreatt + t * T;\n\n float local_sum = 0.0f;\n for (int t2 = block.thread_rank(); t2 < T; t2 += BlockSize) {\n local_sum += att_bth[t2] * datt_bth[t2];\n }\n\n block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus{});\n block.sync();\n local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus{});\n\n for (int t3 = block.thread_rank(); t3 < T; t3 += BlockSize) {\n // don't touch the cache. Some parts will still be here from the previous loop, and\n // we want to exploit those.\n float acc = __ldcs(att_bth + t3) * (__ldcs(datt_bth + t3) - local_sum);\n __stcs(dpreatt_bth + t3, scale * acc);\n }\n }\n}\n\n// the sequence of transformations in this compound op is:\n// inp (B,T,3C) -> qkvr (B,T,3C) -> preatt (B,NH,T,T) -> att (B,NH,T,T) -> vaccum (B,T,C) -> out (B,T,C)\nvoid attention_backward\n(\n hipblasHandle_t cublas_handle,\n float* dinp, float* dqkvr, float* dpreatt, float* datt, float* scratch,\n const float* dout,\n const float* qkvr, const float* att,\n int B, int T, int C, int NH\n) {\n const int block_size = 256;\n int HS = C / NH; // head size\n const float one = 1.0f;\n const float zero = 0.0f; // note beta = 1.0f so that we accumulate gradients (+=)\n // unpack convenience pointers into q, k, v\n const float *q, *k, *v;\n q = qkvr + 0 * B * T * C;\n k = qkvr + 1 * B * T * C;\n v = qkvr + 2 * B * T * C;\n float *dq, *dk, *dv;\n dq = dqkvr + 0 * B * T * C;\n dk = dqkvr + 1 * B * T * C;\n dv = dqkvr + 2 * B * T * C;\n // backward through the unpermute operation\n int num_blocks = ceil_div(B * T * C, block_size);\n unpermute_kernel_backward<<>>(scratch, dout, B, T, NH, HS);\n cudaCheck(hipGetLastError());\n // backward into datt\n cublasCheck(hipblasSgemmStridedBatched(cublas_handle, HIPBLAS_OP_T, HIPBLAS_OP_N, T, T, HS, &one, v, HS, T * HS, scratch, HS, T * HS, &zero, datt, T, T * T, B * NH));\n // backward into dv\n cublasCheck(hipblasSgemmStridedBatched(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_T, HS, T, T, &one, scratch, HS, T * HS, att, T, T * T, &zero, dv, HS, T * HS, B * NH));\n // backward into preatt\n int hs = C / NH; // head size\n float scale = 1.0f / sqrtf(hs);\n softmax_backward_kernel<<>>(dpreatt, datt, att, B, T, C, scale);\n cudaCheck(hipGetLastError());\n // backward into q\n cublasCheck(hipblasSgemmStridedBatched(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_N, HS, T, T, &one, k, HS, T * HS, dpreatt, T, T * T, &zero, dq, HS, T * HS, B * NH));\n // backward into k\n cublasCheck(hipblasSgemmStridedBatched(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_T, HS, T, T, &one, q, HS, T * HS, dpreatt, T, T * T, &zero, dk, HS, T * HS, B * NH));\n // backward into inp\n num_blocks = ceil_div(B * NH * T * HS, block_size);\n permute_kernel_backward<<>>(dinp, dq, dk, dv, B, T, NH, HS);\n cudaCheck(hipGetLastError());\n}\n\n// ------------------------------------------------------------------------------------------------\n// Attention layer, including groupnorm, qkv projection, and residual stream\n\n#define NUM_ATT_PARAM_TENSORS 6\ntypedef struct {\n float* gn_w; // C\n float* gn_b; // C\n float* qkv_w; // 3xC, C\n float* qkv_b; // 3xC\n float* proj_w; // C, C\n float* proj_b; // C\n size_t param_sizes[NUM_ATT_PARAM_TENSORS];\n size_t n_params;\n} AttentionParams;\n\n#define NUM_ATT_ACT_TENSORS 12 // not counting input, which will reuse output memory from previous layer\ntypedef struct {\n float* gn; // B, C, H, W\n float* gn_mean; // B, n_groups\n float* gn_rstd; // B, n_groups\n float* perm1; // B, HxW, C\n float* qkv1; // B, HxW, 3xC\n float* qkv2; // B, HxW, 3xC\n float* preatt; // B, NH, T, T\n float* att; // B, NH, T, T\n float* att_out; // B, HxW, C\n float* proj; // B, HxW, C\n float* perm2; // B, C, H, W\n float* add; // B, C, H, W\n size_t act_sizes[NUM_ATT_ACT_TENSORS];\n size_t n_acts;\n float* input; // B, C, H, W\n} AttentionActs;\n\n#define NUM_ATT_BACKWARD_ACTS_TENSORS 7 // dinp will use input memory\ntypedef struct {\n float* buf1_BCHW; // B, C, H, W\n float* buf2_BCHW; // B, C, H, W\n float* buf_B3CHW; // B, 3xC, H, W\n float* dqkvr; // B, HxW, 3xC\n float* dpreatt; // B, NH, T, T\n float* datt; // B, NH, T, T\n float* dout; // B, C, H, W\n float* dinp; // B, C, H, W\n size_t back_sizes[NUM_ATT_BACKWARD_ACTS_TENSORS];\n size_t n_backs;\n} AttentionBackwardActs;\n\nvoid attention_block_count_params(AttentionParams* params, int C) {\n size_t *param_sizes = params->param_sizes;\n param_sizes[0] = C; // gn_w\n param_sizes[1] = C; // gn_b\n param_sizes[2] = 3 * C * C; // qkv_w\n param_sizes[3] = 3 * C; // qkv_b\n param_sizes[4] = C * C; // proj_w\n param_sizes[5] = C; // proj_b\n\n size_t num_parameters = 0;\n for (int i = 0; i < NUM_ATT_PARAM_TENSORS; i++) {\n num_parameters += param_sizes[i];\n }\n params->n_params = num_parameters;\n}\nvoid attention_block_count_acts(\n AttentionActs* acts,\n int B, int C, int H, int W, int gn_n_groups, int HS\n) {\n int C3 = 3 * C;\n int T = H * W;\n int NH = C / HS;\n size_t *act_sizes = acts->act_sizes;\n act_sizes[0] = B * C * H * W; // gn\n act_sizes[1] = B * gn_n_groups; // gn_mean\n act_sizes[2] = B * gn_n_groups; // gn_rstd\n act_sizes[3] = B * T * C; // perm1\n act_sizes[4] = B * T * C3; // qkv1\n act_sizes[5] = B * T * C3; // qkv2\n act_sizes[6] = B * NH * T * T; // preatt\n act_sizes[7] = B * NH * T * T; // att\n act_sizes[8] = B * T * C; // att_out\n act_sizes[9] = B * T * C; // proj\n act_sizes[10] = B * C * H * W; // perm2\n act_sizes[11] = B * C * H * W; // add\n\n size_t num_acts_params = 0;\n for (int i = 0; i < NUM_ATT_ACT_TENSORS; i++) {\n num_acts_params += act_sizes[i];\n }\n acts->n_acts = num_acts_params;\n}\n\nvoid set_attention_params_pointers(\n AttentionParams* params,\n float* params_memory\n) {\n float** ptrs[] = {\n ¶ms->gn_w,\n ¶ms->gn_b,\n ¶ms->qkv_w,\n ¶ms->qkv_b,\n ¶ms->proj_w,\n ¶ms->proj_b\n };\n for (int i = 0; i < NUM_ATT_PARAM_TENSORS; i++) {\n *(ptrs[i]) = params_memory;\n params_memory += params->param_sizes[i];\n }\n}\n\nvoid set_attention_acts_pointers(\n AttentionActs* acts,\n float* acts_memory\n) {\n float** ptrs[] = {\n &acts->gn,\n &acts->gn_mean,\n &acts->gn_rstd,\n &acts->perm1,\n &acts->qkv1,\n &acts->qkv2,\n &acts->preatt,\n &acts->att,\n &acts->att_out,\n &acts->proj,\n &acts->perm2,\n &acts->add\n };\n for (int i = 0; i < NUM_ATT_ACT_TENSORS; i++) {\n *(ptrs[i]) = acts_memory;\n acts_memory += acts->act_sizes[i];\n }\n}\n\n// ----------------------------------------------------\n// kernels\n\n// this is an opportunity for improvement:\n// rewrite the attention layer so we don't need to permute the input\n__global__ void attention_permute_kernel(\n const float* inp, float* out,\n int B, int C, int H, int W\n) {\n // from B, C, H, W to B, HxW, C\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n if (idx >= B * C * img_size) { return; }\n\n int b = idx / (img_size * C);\n int c = idx / img_size % C;\n int h = idx / W % H;\n int w = idx % W;\n\n out[b * C * img_size + h * C * W + w * C + c] = inp[idx];\n}\n\nvoid attention_permute(\n const float* inp, float* out,\n int B, int C, int H, int W, int block_size\n) {\n int num_threads = B * C * H * W;\n attention_permute_kernel<<>>(inp, out, B, C, H, W);\n}\n\n__global__ void attention_unpermute_kernel(\n const float* inp, float* out,\n int B, int C, int H, int W\n) {\n // from B, HxW, C to B, C, H, W\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n if (idx >= B * C * img_size) { return; }\n\n int b = idx / (img_size * C);\n int c = idx / img_size % C;\n int h = idx / W % H;\n int w = idx % W;\n\n out[idx] = inp[b * C * img_size + h * C * W + w * C + c];\n}\nvoid attention_unpermute(\n const float* inp, float* out,\n int B, int C, int H, int W, int block_size\n) {\n int num_threads = B * C * H * W;\n attention_unpermute_kernel<<>>(inp, out, B, C, H, W);\n}\n\nvoid attention_block_forward\n(\n hipblasHandle_t cublas_handle,\n int B, int C, int H, int W, int HS, int gn_n_groups, int block_size,\n AttentionParams* params, AttentionActs* acts\n) {\n // apply groupnorm\n groupnorm_forward(acts->input, params->gn_w, params->gn_b, acts->gn, acts->gn_mean, acts->gn_rstd, B, C, H, W, gn_n_groups);\n // permute from B, C, H, W to B, HxW, C\n attention_permute(acts->gn, acts->perm1, B, C, H, W, block_size);\n // get qkv\n matmul_forward2(cublas_handle, acts->qkv1, acts->perm1, params->qkv_w, params->qkv_b, B * H * W, C, 3 * C, block_size);\n // apply attention\n attention_forward1(cublas_handle, acts->att_out, acts->qkv2, acts->preatt, acts->att, acts->qkv1, B, H * W, C, C / HS, block_size);\n // attention projection\n matmul_forward2(cublas_handle, acts->proj, acts->att_out, params->proj_w, params->proj_b, B * H * W, C, C, block_size);\n // permute from B, HxW, C to B, C, H, W\n attention_unpermute(acts->proj, acts->perm2, B, C, H, W, block_size);\n // add back residual stream\n add_forward(acts->input, acts->perm2, acts->add, B * C * H * W, block_size);\n}\n\nvoid attention_block_backward\n(\n hipblasHandle_t cublas_handle,\n int B, int C, int H, int W, int HS, int gn_n_groups, int block_size,\n AttentionParams* params, AttentionActs* acts,\n AttentionBackwardActs* back_acts, AttentionParams* grads\n) {\n // permute from B, C, H, W to B, HxW, C\n attention_permute(back_acts->dout, back_acts->buf1_BCHW, B, C, H, W, block_size);\n // backward through projection\n matmul_backward1(cublas_handle, back_acts->buf2_BCHW, grads->proj_w, grads->proj_b, back_acts->buf1_BCHW, acts->att_out, params->proj_w, B * H * W, C, C);\n // backward through attention\n attention_backward(cublas_handle, back_acts->buf_B3CHW, back_acts->dqkvr, back_acts->dpreatt, back_acts->datt, back_acts->buf1_BCHW, back_acts->buf2_BCHW, acts->qkv2, acts->att, B, H * W, C, C / HS);\n // backward through qkv\n matmul_backward1(cublas_handle, back_acts->buf1_BCHW, grads->qkv_w, grads->qkv_b, back_acts->buf_B3CHW, acts->perm1, params->qkv_w, B * H * W, C, 3 * C);\n // permute from B, HxW, C to B, C, H, W\n attention_unpermute(back_acts->buf1_BCHW, back_acts->buf2_BCHW, B, C, H, W, block_size);\n // backward through groupnorm\n groupnorm_backward(back_acts->buf2_BCHW, acts->input, acts->gn_mean, acts->gn_rstd, params->gn_w, back_acts->buf1_BCHW, grads->gn_w, grads->gn_b, B, C, H, W, gn_n_groups);\n // add residual gradient\n add_forward(back_acts->buf1_BCHW, back_acts->dout, back_acts->dinp, B * C * H * W, block_size);\n}\n\n// ------------------------------------------------------------------------------------------------\n// MSE loss\n// this will only be called with one block\n__global__ void mse_forward_kernel(\n const float* inp, const float* y, float* loss, int N\n) {\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n\n __shared__ float shared_sum[32]; // max 32 warps\n int num_warps = blockDim.x / 32;\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n\n float thread_sum = 0.0f;\n for (int i = threadIdx.x; i < N; i += blockDim.x) {\n float diff = inp[i] - y[i];\n thread_sum += diff * diff;\n }\n\n // warp reduce\n float warp_sum = cg::reduce(warp, thread_sum, cg::plus{});\n shared_sum[warp_id] = warp_sum;\n __syncthreads();\n \n warp_sum = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;\n float block_sum = cg::reduce(warp, warp_sum, cg::plus{});\n\n if (threadIdx.x == 0) {\n loss[0] = block_sum / N;\n }\n}\n\nvoid mse_forward(const float* inp, const float* y, float* loss, int N, int block_size) {\n mse_forward_kernel<<<1, block_size>>>(inp, y, loss, N);\n}\n\n// not calculating dy in backward\n// note that backward for MSE doesn't require the loss value\n__global__ void mse_backward_kernel(\n const float* inp, const float* y, float* dinp, int N\n) {\n int idx = threadIdx.x + blockIdx.x * blockDim.x;\n if (idx < N) {\n float x = inp[idx];\n dinp[idx] = 2.0f * (x - y[idx]) / N;\n }\n}\n\nvoid mse_backward(const float* inp, const float* y, float* dinp, int N, int block_size) {\n int n_blocks = ceil_div(N, block_size);\n mse_backward_kernel<<>>(inp, y, dinp, N);\n}\n\n// ------------------------------------------------------------------------------------------------\n// DataLoader\n\ntypedef struct {\n int B;\n int C_in;\n int H;\n int W;\n int img_size;\n int n_imgs;\n FILE* img_file;\n long cur_position;\n long file_size;\n long header_offset;\n long num_batches;\n float* input;\n} DataLoader;\n\nvoid dataloader_init(DataLoader* loader, const char* filename, int B) {\n loader->B = B;\n\n loader->img_file = fopenCheck(filename, \"rb\");\n int data_header[256];\n loader->header_offset = 256 * sizeof(int);\n freadCheck(data_header, sizeof(int), 256, loader->img_file);\n if (data_header[0] != 20240620) {\n fprintf(stderr, \"Invalid magic data file\\n\");\n exit(1);\n }\n loader->n_imgs = data_header[1];\n loader->C_in = data_header[2];\n loader->H = data_header[3];\n loader->W = data_header[4];\n loader->img_size = loader->C_in * loader->H * loader->W;\n\n // determine file size\n fseekCheck(loader->img_file, 0, SEEK_SET);\n fseekCheck(loader->img_file, 0, SEEK_END);\n loader->file_size = ftell(loader->img_file);\n fseekCheck(loader->img_file, 0, SEEK_SET);\n assert((loader->file_size - loader->header_offset) / (loader->img_size * sizeof(float)) == loader->n_imgs);\n \n loader->cur_position = loader->header_offset;\n loader->num_batches = loader->n_imgs / B;\n\n // allocate cuda CPU pinned memory for inputs\n hipHostMalloc((void**)&loader->input, B * loader->img_size * sizeof(float));\n}\n\nvoid dataloader_reset(DataLoader* loader) {\n loader->cur_position = loader->header_offset;\n}\n\nvoid dataloader_next_batch(DataLoader* loader) {\n int B = loader->B;\n int img_size = loader->img_size;\n if (loader->cur_position + (B * img_size * sizeof(float)) > loader->file_size) {\n dataloader_reset(loader);\n }\n fseekCheck(loader->img_file, loader->cur_position, SEEK_SET);\n freadCheck(loader->input, sizeof(float), B * img_size, loader->img_file);\n loader->cur_position += B * img_size * sizeof(float);\n}\n\nvoid dataloader_free(DataLoader* loader) {\n fcloseCheck(loader->img_file);\n hipHostFree(loader->input);\n}\n\n\n// ------------------------------------------------------------------------------------------------\n// Diffusion\ntypedef struct {\n int B;\n int img_size;\n int max_period;\n float* betas; // (max_period,)\n float* sqrt_alphas_cumprod; // (max_period,)\n float* sqrt_one_minus_alphas_cumprod; // (max_period,)\n int block_size;\n hiprandState* curand_states;\n} GaussianDiffusion;\n\n__global__ void init_curand_states(hiprandState *state, unsigned long seed, int N) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < N) {\n hiprand_init(seed, idx, 0, &state[idx]);\n }\n}\n\nvoid diffusion_init(\n GaussianDiffusion* diffusion, int B, int C_in, int H, int W, int max_period, int block_size = 512,\n unsigned long seed = 0, float beta_start = 1e-4, float beta_end = 0.02f\n) {\n diffusion->B = B;\n diffusion->max_period = max_period;\n diffusion->block_size = block_size;\n diffusion->img_size = C_in * H * W;\n\n float scale = 1000.0f / max_period;\n beta_start *= scale;\n beta_end *= scale;\n\n // calculate betas and alphas on host then move to device\n float* betas = (float*)mallocCheck(max_period * sizeof(float));\n float* sqrt_alphas_cumprod = (float*)mallocCheck(max_period * sizeof(float));\n float* sqrt_one_minus_alphas_cumprod = (float*)mallocCheck(max_period * sizeof(float));\n\n float alpha_cumprod = 1.0f;\n for (int t = 1; t <= max_period; t++) {\n // linear interpolate for beta\n betas[t - 1] = (beta_start * (max_period - t) + beta_end * (t - 1)) / (max_period - 1);\n alpha_cumprod *= 1.0f - betas[t - 1];\n sqrt_alphas_cumprod[t - 1] = sqrtf(alpha_cumprod);\n sqrt_one_minus_alphas_cumprod[t - 1] = sqrtf(1.0f - alpha_cumprod);\n }\n\n float *d_betas, *d_sqrt_alphas_cumprod, *d_sqrt_one_minus_alphas_cumprod;\n cudaCheck(hipMalloc(&d_betas, max_period * sizeof(float)));\n cudaCheck(hipMalloc(&d_sqrt_alphas_cumprod, max_period * sizeof(float)));\n cudaCheck(hipMalloc(&d_sqrt_one_minus_alphas_cumprod, max_period * sizeof(float)));\n\n cudaCheck(hipMemcpy(d_betas, betas, max_period * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_sqrt_alphas_cumprod, sqrt_alphas_cumprod, max_period * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_sqrt_one_minus_alphas_cumprod, sqrt_one_minus_alphas_cumprod, max_period * sizeof(float), hipMemcpyHostToDevice));\n\n free(betas);\n free(sqrt_alphas_cumprod);\n free(sqrt_one_minus_alphas_cumprod);\n\n diffusion->betas = d_betas;\n diffusion->sqrt_alphas_cumprod = d_sqrt_alphas_cumprod;\n diffusion->sqrt_one_minus_alphas_cumprod = d_sqrt_one_minus_alphas_cumprod;\n\n // init curand states for sampling noise of the size of one batch of inputs\n if (diffusion->curand_states == nullptr) {\n printf(\"initializing hiprand states\\n\");\n int total_states = B * C_in * H * W;\n int n_blocks = ceil_div(total_states, block_size);\n cudaCheck(hipMalloc(&diffusion->curand_states, total_states * sizeof(hiprandState)));\n init_curand_states<<>>(diffusion->curand_states, seed, total_states);\n }\n}\n\n// fill BxN noise array with standard gaussians\n__global__ void diffusion_draw_normal_kernel(\n float* noise, int B, int N, hiprandState* curand_states\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * N) {\n noise[idx] = hiprand_normal(&curand_states[idx]);\n }\n}\n\nvoid diffusion_draw_normal(\n GaussianDiffusion *diffusion,\n float* noise\n) {\n int B = diffusion->B;\n int img_size = diffusion->img_size;\n int total_states = B * img_size;\n int n_blocks = ceil_div(total_states, diffusion->block_size);\n diffusion_draw_normal_kernel<<block_size>>>(noise, B, img_size, diffusion->curand_states);\n}\n\n// given input x of shape (B, N) and timesteps of shape (B,)\n// replace each x entry by the result of flowing forward by diffusion for time t\n// with randomness coming from noise\n__global__ void diffusion_forward_by_t_kernel(\n float* x, float* timesteps, float* noise,\n float* sqrt_alphas_cumprod, float* sqrt_one_minus_alphas_cumprod,\n int B, int N, hiprandState* curand_states\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * N) {\n int b = idx / N;\n float x_val = x[idx];\n float noise_val = noise[idx];\n int t = (int)timesteps[b];\n float sqrt_alpha_t = sqrt_alphas_cumprod[t];\n float sqrt_one_minus_alpha_t = sqrt_one_minus_alphas_cumprod[t];\n x[idx] = x_val * sqrt_alpha_t + noise_val * sqrt_one_minus_alpha_t;\n }\n}\n\nvoid diffusion_forward_by_t(\n GaussianDiffusion* diffusion,\n float* x, float* timesteps, float* noise\n) {\n int B = diffusion->B;\n int img_size = diffusion->img_size;\n int block_size = diffusion->block_size;\n int total_states = B * img_size;\n int n_blocks = ceil_div(total_states, block_size);\n diffusion_forward_by_t_kernel<<>>(\n x, timesteps, noise, diffusion->sqrt_alphas_cumprod, diffusion->sqrt_one_minus_alphas_cumprod, B, img_size, diffusion->curand_states\n );\n}\n\nvoid diffusion_free(GaussianDiffusion* diffusion) {\n cudaCheck(hipFree(diffusion->betas));\n cudaCheck(hipFree(diffusion->sqrt_alphas_cumprod));\n cudaCheck(hipFree(diffusion->sqrt_one_minus_alphas_cumprod));\n cudaCheck(hipFree(diffusion->curand_states));\n}\n\n// one block with B threads\n__global__ void sample_timestep_kernel(\n float* timesteps, int B,\n hiprandState* curand_states, int max_period\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B) {\n int t = hiprand(&curand_states[idx]) % max_period;\n timesteps[idx] = (float)t;\n }\n}\n\n// samples B timesteps in range [0, max_period), convert them to floats and\n// store in timesteps\nvoid sample_timesteps(GaussianDiffusion *diffusion, float* timesteps) {\n sample_timestep_kernel<<<1, diffusion->B>>>(timesteps, diffusion->B, diffusion->curand_states, diffusion->max_period);\n}\n\n// ------------------------------------------------------------------------------------------------\n// timestep embedding states\ntypedef struct {\n float* freqs;\n int half_dim;\n int B;\n int max_period;\n} TimestepEmbedding;\n\n__global__ void fill_freqs_kernel(float* freqs, int half_dim, int max_period) {\n int idx = threadIdx.x + blockIdx.x * blockDim.x;\n if (idx < half_dim) {\n float log_period = logf(max_period);\n float logit = -log_period * idx / half_dim;\n freqs[idx] = expf(logit);\n }\n}\n\nvoid init_timestep_embedding(TimestepEmbedding* emb, int dim, int B, int max_period) {\n assert(dim % 2 == 0);\n int half = dim / 2;\n emb->half_dim = half;\n emb->max_period = max_period;\n emb->B = B;\n cudaCheck(hipMalloc(&emb->freqs, half * sizeof(float)));\n // fill freq values\n int block_size = 512;\n int n_blocks = ceil_div(half, block_size);\n fill_freqs_kernel<<>>(emb->freqs, half, max_period);\n}\n\nvoid free_timestep_embedding(TimestepEmbedding* emb) {\n cudaCheck(hipFree(emb->freqs));\n}\n\n__global__ void fill_embeddings_kernel(\n const float* timesteps, const float* freqs, float* embeddings, int half, int B\n) {\n int idx = threadIdx.x + blockIdx.x * blockDim.x;\n if (idx < B * half) {\n int i = idx / half;\n int j = idx % half;\n float arg = timesteps[i] * freqs[j];\n embeddings[i * 2 * half + j] = cosf(arg);\n embeddings[i * 2 * half + j + half] = sinf(arg);\n }\n}\n\n// timesteps is a (B,) array of timesteps that are integers but saved as floats\n// out is a (B, dim) array\nvoid get_timestep_embeddings(\n TimestepEmbedding* emb, const float* timesteps, float* out\n) {\n int block_size = 512;\n int n_blocks = ceil_div(emb->half_dim * emb->B, block_size);\n fill_embeddings_kernel<<>>(timesteps, emb->freqs, out, emb->half_dim, emb->B);\n cudaCheck(hipGetLastError());\n}\n\n// ------------------------------------------------------------------------------------------------\n// Unet init\n\ntypedef struct {\n int C_in;\n int C_model;\n int C_time;\n int C_out;\n int B;\n int H;\n int W;\n int gn_n_groups;\n int HS;\n int block_size;\n int n_layers;\n int n_levels;\n int n_res;\n int att_start_level;\n int* channel_mult;\n int max_period;\n} UnetConfig;\n\n// we just need a single set of buffers that are shared for all the backward pass activations\n#define NUM_SKIP_CONNECTIONS 12 // (n_res + 1) * n_levels\n#define NUM_BACK_BUFS 10\n#define UNET_NUM_BACK_ACTS 22\ntypedef struct {\n float* dinp_buf1; // used to store the input gradient, which will be reused as dout for the previous layer\n float* dinp_buf2;\n float* buf1_BCHW; // NOTE: first 4 buffers will have max B * C * H * W size\n float* buf2_BCHW;\n float* buf3_BCHW;\n float* buf4_BCHW;\n float* buf1_B3CHW; // NOTE: last 2 buffers will have max B * 3 * C * H * W size for attention layers\n float* buf2_B3CHW; // so max of H * W is only over dimensions that have attention\n float* demb_buf1; // used to store cumulative gradients w.r.t. timestep embeddings\n float* demb_buf2;\n float* skip_dinps[NUM_SKIP_CONNECTIONS];\n} UnetBackActs;\n\n// This is used as a block to store all the params / grads / acts for all the layers\ntypedef struct {\n float* mem_ptr;\n size_t total_size;\n int n_layers;\n void** layer_ptrs;\n size_t* layer_sizes;\n} UnetBlock;\n\ntypedef struct {\n float* back_mem;\n UnetBackActs back_acts;\n size_t back_sizes[UNET_NUM_BACK_ACTS];\n size_t n_backs;\n} UnetBackBlock;\n\ntypedef struct {\n UnetConfig config;\n float* input;\n float* timestep_emb; // this is the already computed embeddings of shape (B, C_model)\n UnetBlock params;\n UnetBlock grads;\n UnetBlock acts;\n UnetBackBlock back;\n float* skip_acts[NUM_SKIP_CONNECTIONS]; // pointers to skip connetion activations\n // the next two are for debug purposes only, just to keep track of the buffer that contains the grad w.r.t. inputs\n float* dinp;\n // model output\n float* output;\n // targets and losses\n float* target;\n float mean_loss;\n float* gpu_loss;\n // buffers for AdamW\n float* m_memory;\n float* v_memory;\n TimestepEmbedding time_emb_state;\n} Unet;\n\nvoid unet_init(Unet* unet) {\n unet->input = nullptr;\n unet->timestep_emb = nullptr;\n unet->params.mem_ptr = nullptr;\n unet->grads.mem_ptr = nullptr;\n unet->acts.mem_ptr = nullptr;\n unet->back.back_mem = nullptr;\n unet->dinp = nullptr;\n unet->output = nullptr;\n unet->target = nullptr;\n unet->mean_loss = -1.0f;\n unet->gpu_loss = nullptr;\n unet->m_memory = nullptr;\n unet->v_memory = nullptr;\n\n init_timestep_embedding(&unet->time_emb_state, unet->config.C_model, unet->config.B, unet->config.max_period);\n}\n\n\ntemplate \nvoid free_not_null(T** ptr) {\n if (*ptr != nullptr) {\n free(*ptr);\n *ptr = nullptr;\n }\n}\n\nvoid free_unet_block(UnetBlock* params) {\n free_not_null(¶ms->layer_sizes);\n for (int l = 0; l < params->n_layers; l++) {\n free_not_null(&(params->layer_ptrs[l]));\n }\n free_not_null(¶ms->layer_ptrs);\n cudaCheck(hipFree(params->mem_ptr));\n}\n\nvoid free_unet(Unet* unet) {\n //free(unet->config.channel_mult);\n\n cudaCheck(hipFree(unet->input));\n cudaCheck(hipFree(unet->timestep_emb));\n cudaCheck(hipFree(unet->gpu_loss));\n cudaCheck(hipFree(unet->m_memory));\n cudaCheck(hipFree(unet->v_memory));\n\n free_unet_block(&(unet->params));\n free_unet_block(&(unet->grads));\n free_unet_block(&(unet->acts));\n\n cudaCheck(hipFree(unet->back.back_mem));\n free_timestep_embedding(&unet->time_emb_state);\n}\n\nvoid unet_count_layers(UnetConfig* config) {\n int n_layers = 0;\n int n_res = config->n_res;\n int n_levels = config->n_levels;\n int att_start_level = config->att_start_level;\n\n // timestep emb\n n_layers += 3;\n\n // initial conv\n n_layers += 1;\n\n // add downsample resblock layers\n n_layers += n_res * n_levels;\n\n // add downsample attention\n n_layers += n_res * (n_levels - att_start_level);\n\n // add avgpool layers\n n_layers += n_levels - 1;\n\n // add middle layers\n n_layers += 3;\n\n // add upsample concat layers\n n_layers += (n_res + 1) * n_levels;\n\n // add upsample resblock layers\n n_layers += (n_res + 1) * n_levels;\n\n // add upsample attention layers\n n_layers += (n_res + 1) * (n_levels - att_start_level);\n\n // add upsample layers\n n_layers += n_levels - 1;\n\n // output layers\n n_layers += 3;\n\n config->n_layers = n_layers;\n}\n\n// ------------------------------------------------------------------------------------------------\n// allocate memory for structs with info for each layer and count params\n\nint make_linear_params_ptrs(\n UnetBlock* params,\n int layer,\n int C_in, int C_out\n) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(LinearParams));\n params->layer_sizes[layer] = linear_count_params(C_in, C_out);\n return layer + 1;\n}\n\nint make_silu_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_convk3_params_ptrs(\n UnetBlock* params, int layer,\n int C_in, int OC, int B, int H, int W\n) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ConvK3Params));\n params->layer_sizes[layer] = convk3_count_params(C_in, OC);\n return layer + 1;\n}\n\nint make_resblock_params_ptrs(\n UnetBlock* params, int layer,\n int C, int C_emb, int C_out, int B, int H, int W, int gn_n_groups\n) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ResBlockParameters));\n ResBlockParameters* resblock_params = (ResBlockParameters*)params->layer_ptrs[layer];\n resblock_count_params(resblock_params, C, C_emb, C_out, B, H, W, 0, 0, gn_n_groups);\n params->layer_sizes[layer] = resblock_params->n_params;\n return layer + 1;\n}\n\nint make_avgpool_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_attention_block_params_ptrs(UnetBlock* params, int layer, int C) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(AttentionParams));\n AttentionParams* attention_params = (AttentionParams*)params->layer_ptrs[layer];\n attention_block_count_params(attention_params, C);\n params->layer_sizes[layer] = attention_params->n_params;\n return layer + 1;\n}\n\nint make_concat_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_upsample_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_groupnorm_params_ptrs(UnetBlock* params, int layer, int C) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(GroupNormParams));\n params->layer_sizes[layer] = C * 2;\n return layer + 1;\n}\n\nvoid unet_setup_and_count_params(UnetConfig* config, UnetBlock* params) {\n assert (config->n_layers > 0);\n\n params->n_layers = config->n_layers;\n params->layer_ptrs = (void**)mallocCheck(config->n_layers * sizeof(void*));\n params->layer_sizes = (size_t*)mallocCheck(config->n_layers * sizeof(size_t));\n\n int B = config->B;\n int C_in = config->C_in;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int C_out = config->C_out;\n int H = config->H;\n int W = config->W;\n int gn_n_groups = config->gn_n_groups;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n int layer = 0;\n layer = make_linear_params_ptrs(params, layer, C_model, C_time);\n layer = make_silu_params_ptrs(params, layer);\n layer = make_linear_params_ptrs(params, layer, C_time, C_time);\n int OC = C_model;\n layer = make_convk3_params_ptrs(params, layer, C_in, OC, B, H, W);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n layer = make_resblock_params_ptrs(params, layer, C, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_params_ptrs(params, layer, C);\n }\n }\n if (l < n_levels - 1) {\n layer = make_avgpool_params_ptrs(params, layer);\n H /= 2;\n W /= 2;\n }\n }\n\n // middle layers\n layer = make_resblock_params_ptrs(params, layer, C, C_time, C, B, H, W, gn_n_groups);\n layer = make_attention_block_params_ptrs(params, layer, C);\n layer = make_resblock_params_ptrs(params, layer, C, C_time, C, B, H, W, gn_n_groups);\n\n // upwards: concat then resblock\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = C_model * mult;\n for (int r = 0; r < n_res + 1; r++) {\n if (r == n_res && l > 0) {\n int prev_mult = config->channel_mult[l - 1];\n C_skip = C_model * prev_mult;\n }\n OC = C_model * mult;\n layer = make_concat_params_ptrs(params, layer);\n layer = make_resblock_params_ptrs(params, layer, C + C_skip, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_params_ptrs(params, layer, C);\n }\n }\n if (l > 0) {\n layer = make_upsample_params_ptrs(params, layer);\n H *= 2;\n W *= 2;\n }\n }\n\n layer = make_groupnorm_params_ptrs(params, layer, C);\n layer = make_silu_params_ptrs(params, layer);\n layer = make_convk3_params_ptrs(params, layer, C, C_out, B, H, W);\n\n // count total size\n size_t total_size = 0;\n for (int i = 0; i < config->n_layers; i++) {\n total_size += params->layer_sizes[i];\n }\n params->total_size = total_size;\n}\n\n// ------------------------------------------------------------------------------------------------\n// setup Unet activations pointers\nint make_linear_acts_ptrs(\n UnetBlock* acts,\n int layer,\n int B, int C_out\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(LinearActs));\n // acts buffers only hold outputs\n acts->layer_sizes[layer] = B * C_out;\n return layer + 1;\n}\n\nint make_silu_acts_ptrs(UnetBlock* acts, int layer, int N) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(SiluActs));\n acts->layer_sizes[layer] = N;\n return layer + 1;\n}\n\nint make_convk3_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int OC, int H, int W\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ConvK3Acts));\n acts->layer_sizes[layer] = B * OC * H * W;\n return layer + 1;\n}\n\nint make_resblock_acts_ptrs(\n UnetBlock* acts, int layer,\n int C, int C_emb, int C_out, int B, int H, int W, int gn_n_groups\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ResBlockActivations));\n ResBlockActivations* resblock_acts = (ResBlockActivations*)acts->layer_ptrs[layer];\n resblock_count_acts(resblock_acts, C, C_emb, C_out, B, H, W, 0, 0, gn_n_groups);\n acts->layer_sizes[layer] = resblock_acts->n_acts;\n return layer + 1;\n}\n\nint make_avgpool_acts_ptrs(UnetBlock* acts, int layer, int B, int C, int H, int W) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(AvgpoolActs));\n acts->layer_sizes[layer] = B * C * H/2 * W/2;\n return layer + 1;\n}\n\nint make_attention_block_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C, int H, int W, int gn_n_groups, int HS\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(AttentionActs));\n AttentionActs* att_acts = (AttentionActs*)acts->layer_ptrs[layer];\n attention_block_count_acts(att_acts, B, C, H, W, gn_n_groups, HS);\n acts->layer_sizes[layer] = att_acts->n_acts;\n return layer + 1;\n}\n\nint make_concat_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C1, int C2, int H, int W\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ConcatChannelActs));\n acts->layer_sizes[layer] = B * (C1 + C2) * H * W;\n return layer + 1;\n}\n\nint make_upsample_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C, int H, int W\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(UpsampleActs));\n acts->layer_sizes[layer] = B * C * H * 2 * W * 2;\n return layer + 1;\n}\n\nint make_groupnorm_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C, int H, int W, int gn_n_groups\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(GroupNormActs));\n // out (B, C, H, W), mean and rstd (B, gn_n_groups)\n acts->layer_sizes[layer] = B * C * H * W + B * gn_n_groups * 2;\n return layer + 1;\n}\n\nvoid unet_setup_and_count_acts(UnetConfig* config, UnetBlock* acts) {\n assert (config->n_layers > 0);\n\n acts->n_layers = config->n_layers;\n acts->layer_ptrs = (void**)mallocCheck(config->n_layers * sizeof(void*));\n acts->layer_sizes = (size_t*)mallocCheck(config->n_layers * sizeof(size_t));\n\n int B = config->B;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int C_out = config->C_out;\n int H = config->H;\n int W = config->W;\n int gn_n_groups = config->gn_n_groups;\n int HS = config->HS;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n int layer = 0;\n layer = make_linear_acts_ptrs(acts, layer, B, C_time);\n layer = make_silu_acts_ptrs(acts, layer, B * C_time);\n layer = make_linear_acts_ptrs(acts, layer, B, C_time);\n int OC = C_model;\n layer = make_convk3_acts_ptrs(acts, layer, B, OC, H, W);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n layer = make_resblock_acts_ptrs(acts, layer, C, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups, HS);\n }\n }\n if (l < n_levels - 1) {\n layer = make_avgpool_acts_ptrs(acts, layer, B, C, H, W);\n H /= 2;\n W /= 2;\n }\n }\n\n layer = make_resblock_acts_ptrs(acts, layer, C, C_time, C, B, H, W, gn_n_groups);\n layer = make_attention_block_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups, HS);\n layer = make_resblock_acts_ptrs(acts, layer, C, C_time, C, B, H, W, gn_n_groups);\n\n // upwards\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = mult * C_model;\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n int prev_mult = config->channel_mult[l - 1];\n C_skip = prev_mult * C_model;\n }\n layer = make_concat_acts_ptrs(acts, layer, B, C, C_skip, H, W);\n OC = mult * C_model;\n layer = make_resblock_acts_ptrs(acts, layer, C + C_skip, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups, HS);\n }\n }\n if (l > 0) {\n layer = make_upsample_acts_ptrs(acts, layer, B, C, H, W);\n H *= 2;\n W *= 2;\n }\n }\n layer = make_groupnorm_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups);\n layer = make_silu_acts_ptrs(acts, layer, B * C * H * W);\n layer = make_convk3_acts_ptrs(acts, layer, B, C_out, H, W);\n\n // count total size\n size_t total_size = 0;\n for (int i = 0; i < config->n_layers; i++) {\n total_size += acts->layer_sizes[i];\n }\n acts->total_size = total_size;\n}\n\nvoid unet_setup_and_count_backs(UnetConfig* config, UnetBackBlock* back) {\n int B = config->B;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int H = config->H;\n int W = config->W;\n\n size_t max_BCHW = 3 * B * C_model * H * W;\n size_t max_B3CHW = 2 * B * 3 * C_model * H * W;\n back->back_sizes[0] = max_BCHW;\n back->back_sizes[1] = max_BCHW;\n back->back_sizes[2] = max_BCHW;\n back->back_sizes[3] = max_BCHW;\n back->back_sizes[4] = max_BCHW;\n back->back_sizes[5] = max_BCHW;\n back->back_sizes[6] = max_B3CHW;\n back->back_sizes[7] = max_B3CHW;\n back->back_sizes[8] = B * C_time;\n back->back_sizes[9] = B * C_time;\n\n // fill the sizes of the skip connection backward acts\n int C = C_model;\n for (int l = 0; l < config->n_levels; l++) {\n for (int r = 0; r < (config->n_res + 1); r++) {\n int s = l * 3 + r + NUM_BACK_BUFS;\n assert(s < UNET_NUM_BACK_ACTS);\n back->back_sizes[s] = B * C * H * W;\n int mult = config->channel_mult[l];\n C = C_model * mult;\n if (l > 0 && r == 0) {\n H /= 2;\n W /= 2;\n }\n }\n }\n\n size_t n_backs = 0;\n for (int i = 0; i < UNET_NUM_BACK_ACTS; i++) {\n n_backs += back->back_sizes[i];\n }\n back->n_backs = n_backs;\n}\n\n\nvoid unet_make_ptrs_and_count_memory(Unet* unet)\n{\n unet_setup_and_count_params(&(unet->config), &(unet->params));\n unet_setup_and_count_params(&(unet->config), &(unet->grads));\n unet_setup_and_count_acts(&(unet->config), &(unet->acts));\n unet_setup_and_count_backs(&(unet->config), &(unet->back));\n}\n\n// ------------------------------------------------------------------------------------------------\n// malloc unet memory and set pointers\ntypedef struct {\n float* input;\n float* prev_out;\n int layer;\n float* emb;\n float* emb_cum; // cumulative demb\n float* mem_ptr;\n int skip_idx;\n} LayerInfo;\n\n\nvoid set_linear_params_ptrs(\n UnetBlock* params,\n LayerInfo* l_info,\n int C_in, int C_out\n) {\n int layer = l_info->layer;\n LinearParams* linear_params = (LinearParams*)params->layer_ptrs[layer];\n linear_set_param_ptrs(linear_params, l_info->mem_ptr, C_in, C_out);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_convk3_params_ptrs(UnetBlock* params, LayerInfo* l_info, int C_in, int OC) {\n int layer = l_info->layer;\n ConvK3Params* convk3_params = (ConvK3Params*)params->layer_ptrs[layer];\n convk3_set_param_ptrs(convk3_params, l_info->mem_ptr, C_in, OC);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_resblock_layer_params_ptrs(UnetBlock* params, LayerInfo* l_info, int C, int OC) {\n int layer = l_info->layer;\n ResBlockParameters* resblock_params = (ResBlockParameters*)params->layer_ptrs[layer];\n set_resblock_params_ptrs(C, OC, resblock_params, l_info->mem_ptr);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_attention_layer_params_ptrs(UnetBlock* params, LayerInfo* l_info) {\n int layer = l_info->layer;\n AttentionParams* att_params = (AttentionParams*)params->layer_ptrs[layer];\n set_attention_params_pointers(att_params, l_info->mem_ptr);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_groupnorm_layer_params_ptrs(UnetBlock* params, LayerInfo* l_info, int C) {\n int layer = l_info->layer;\n GroupNormParams* gn_params = (GroupNormParams*)params->layer_ptrs[layer];\n gn_set_param_ptrs(gn_params, l_info->mem_ptr, C);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\n\nvoid point_unet_params(UnetConfig* config, UnetBlock* params) {\n // common constants\n int C_in = config->C_in;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int C_out = config->C_out;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n LayerInfo l_info;\n l_info.layer = 0;\n l_info.mem_ptr = params->mem_ptr;\n set_linear_params_ptrs(params, &l_info, C_model, C_time);\n l_info.layer++; // skip silu layer\n set_linear_params_ptrs(params, &l_info, C_time, C_time);\n int OC = C_model;\n set_convk3_params_ptrs(params, &l_info, C_in, OC);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n set_resblock_layer_params_ptrs(params, &l_info, C, OC);\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_params_ptrs(params, &l_info);\n }\n }\n if (l < n_levels - 1) {\n l_info.layer++; // skip avgpool layer\n }\n }\n\n set_resblock_layer_params_ptrs(params, &l_info, C, C);\n set_attention_layer_params_ptrs(params, &l_info);\n set_resblock_layer_params_ptrs(params, &l_info, C, C);\n\n // upward: skip one layer for concat\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = mult * C_model;\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config->channel_mult[l - 1];\n }\n l_info.layer++;\n OC = mult * C_model;\n set_resblock_layer_params_ptrs(params, &l_info, C + C_skip, OC);\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_params_ptrs(params, &l_info);\n }\n }\n if (l > 0) {\n l_info.layer++; // upsample\n }\n }\n\n set_groupnorm_layer_params_ptrs(params, &l_info, C);\n l_info.layer++; // silu\n set_convk3_params_ptrs(params, &l_info, C, C_out);\n}\n\nvoid set_linear_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n LinearActs* linear_acts = (LinearActs*)acts->layer_ptrs[layer];\n linear_acts->out = l_info->mem_ptr;\n l_info->mem_ptr += acts->layer_sizes[layer];\n linear_acts->inp = l_info->input;\n l_info->input = linear_acts->out;\n l_info->layer++;\n}\n\nvoid set_silu_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n SiluActs* silu_acts = (SiluActs*)acts->layer_ptrs[layer];\n silu_acts->out = l_info->mem_ptr;\n l_info->mem_ptr += acts->layer_sizes[layer];\n silu_acts->x = l_info->input;\n l_info->input = silu_acts->out;\n l_info->layer++;\n}\n\nvoid set_convk3_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n ConvK3Acts* convk3_acts = (ConvK3Acts*)acts->layer_ptrs[layer];\n convk3_acts->out = l_info->mem_ptr;\n l_info->mem_ptr += acts->layer_sizes[layer];\n convk3_acts->inp = l_info->input;\n l_info->input = convk3_acts->out;\n l_info->layer++;\n}\n\nvoid set_resblock_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info, int C, int OC) {\n int layer = l_info->layer;\n ResBlockActivations* resblock_acts = (ResBlockActivations*)acts->layer_ptrs[layer];\n set_resblock_acts_ptrs(C, OC, resblock_acts, l_info->mem_ptr);\n resblock_acts->input = l_info->input;\n resblock_acts->emb = l_info->emb;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = resblock_acts->add2;\n l_info->layer++;\n}\n\nvoid set_avgpool_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n AvgpoolActs* avgpool_acts = (AvgpoolActs*)acts->layer_ptrs[layer];\n avgpool_acts->out = l_info->mem_ptr;\n avgpool_acts->x = l_info->input;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = avgpool_acts->out;\n l_info->layer++;\n}\n\nvoid set_attention_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n AttentionActs* att_acts = (AttentionActs*)acts->layer_ptrs[layer];\n set_attention_acts_pointers(att_acts, l_info->mem_ptr);\n att_acts->input = l_info->input;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = att_acts->add;\n l_info->layer++;\n}\n\nvoid set_concat_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info, float** skip_acts) {\n int layer = l_info->layer;\n ConcatChannelActs* cat_acts = (ConcatChannelActs*)acts->layer_ptrs[layer];\n cat_acts->out = l_info->mem_ptr;\n cat_acts->x1 = l_info->input;\n cat_acts->x2 = skip_acts[l_info->skip_idx];\n\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = cat_acts->out;\n l_info->layer++;\n l_info->skip_idx--;\n}\n\nvoid set_upsample_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n UpsampleActs* up_acts = (UpsampleActs*)acts->layer_ptrs[layer];\n up_acts->x = l_info->input;\n up_acts->out = l_info->mem_ptr;\n\n l_info->mem_ptr += acts-> layer_sizes[layer];\n l_info->input = up_acts->out;\n l_info->layer++;\n}\n\nvoid set_groupnorm_layer_acts_ptrs(\n UnetBlock* acts, LayerInfo* l_info,\n int B, int C, int H, int W, int gn_n_groups\n) {\n int layer = l_info->layer;\n GroupNormActs* gn_a = (GroupNormActs*)acts->layer_ptrs[layer];\n gn_set_act_ptrs(gn_a, l_info->mem_ptr, B, C, H, W, gn_n_groups);\n gn_a->x = l_info->input;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = gn_a->out;\n l_info->layer++;\n}\n\nvoid point_unet_acts(UnetConfig* config, UnetBlock* acts, float** skip_acts, float* input, float* timestep_emb) {\n // common constants\n int B = config->B;\n int C_model = config->C_model;\n int H = config->H;\n int W = config->W;\n int gn_n_groups = config->gn_n_groups;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n LayerInfo l_info;\n l_info.layer = 0;\n l_info.mem_ptr = acts->mem_ptr;\n\n // timestep emb\n l_info.input = timestep_emb;\n set_linear_acts_ptrs(acts, &l_info);\n set_silu_acts_ptrs(acts, &l_info);\n set_linear_acts_ptrs(acts, &l_info);\n // point input to model input, and set outputu of the last linear layer as the embedding input\n l_info.emb = l_info.input;\n l_info.input = input;\n\n int OC = C_model;\n set_convk3_acts_ptrs(acts, &l_info);\n int skip_idx = 0;\n skip_acts[skip_idx] = l_info.input; // set the output of the previouse layer as input to the skip connection\n skip_idx++;\n\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n set_resblock_layer_acts_ptrs(acts, &l_info, C, OC);\n if (l < att_start_level) {\n // for the first 2 layers, the skip connection comes from the resblocks\n skip_acts[skip_idx] = l_info.input;\n skip_idx++;\n }\n\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_acts_ptrs(acts, &l_info);\n skip_acts[skip_idx] = l_info.input;\n skip_idx++;\n }\n }\n if (l < n_levels - 1) {\n set_avgpool_layer_acts_ptrs(acts, &l_info);\n skip_acts[skip_idx] = l_info.input;\n skip_idx++;\n }\n }\n\n set_resblock_layer_acts_ptrs(acts, &l_info, C, C);\n set_attention_layer_acts_ptrs(acts, &l_info);\n set_resblock_layer_acts_ptrs(acts, &l_info, C, C);\n\n // upward: set concat\n l_info.skip_idx = NUM_SKIP_CONNECTIONS - 1;\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = mult * C_model;\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config->channel_mult[l - 1];\n }\n set_concat_layer_acts_ptrs(acts, &l_info, skip_acts);\n OC = mult * C_model;\n set_resblock_layer_acts_ptrs(acts, &l_info, C + C_skip, OC);\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_acts_ptrs(acts, &l_info);\n }\n }\n if (l > 0) {\n set_upsample_layer_acts_ptrs(acts, &l_info);\n }\n }\n set_groupnorm_layer_acts_ptrs(acts, &l_info, B, C, H, W, gn_n_groups);\n set_silu_acts_ptrs(acts, &l_info);\n set_convk3_acts_ptrs(acts, &l_info);\n}\n\nvoid set_unet_back_pointers(UnetBackBlock* back) {\n UnetBackActs* back_acts = &(back->back_acts);\n float* back_mem_ptr = back->back_mem;\n\n float** ptrs[] = {\n &(back_acts->dinp_buf1),\n &(back_acts->dinp_buf2),\n &(back_acts->buf1_BCHW),\n &(back_acts->buf2_BCHW),\n &(back_acts->buf3_BCHW),\n &(back_acts->buf4_BCHW),\n &(back_acts->buf1_B3CHW),\n &(back_acts->buf2_B3CHW),\n &(back_acts->demb_buf1),\n &(back_acts->demb_buf2)\n };\n\n for (int i = 0; i < NUM_BACK_BUFS; i++) {\n *(ptrs[i]) = back_mem_ptr;\n back_mem_ptr += back->back_sizes[i];\n }\n\n // point the back acts for skip connections\n for (int s = 0; s < NUM_SKIP_CONNECTIONS; s++) {\n back_acts->skip_dinps[s] = back_mem_ptr;\n back_mem_ptr += back->back_sizes[NUM_BACK_BUFS + s];\n }\n}\n\nvoid malloc_and_point(Unet* unet) {\n float *param_mem, *grad_mem, *act_mem, *back_mem;\n cudaCheck(hipMalloc(¶m_mem, unet->params.total_size * sizeof(float)));\n cudaCheck(hipMalloc(&grad_mem, unet->params.total_size * sizeof(float)));\n cudaCheck(hipMalloc(&act_mem, unet->acts.total_size * sizeof(float)));\n cudaCheck(hipMalloc(&back_mem, unet->back.n_backs * sizeof(float)));\n\n cudaCheck(hipMemset(param_mem, 0, unet->params.total_size * sizeof(float)));\n cudaCheck(hipMemset(grad_mem, 0, unet->params.total_size * sizeof(float)));\n cudaCheck(hipMemset(act_mem, 0, unet->acts.total_size * sizeof(float)));\n cudaCheck(hipMemset(back_mem, 0, unet->back.n_backs * sizeof(float)));\n\n // constants\n int B = unet->config.B;\n int C_in = unet->config.C_in;\n int C_model = unet->config.C_model;\n int H = unet->config.H;\n int W = unet->config.W;\n\n // malloc and point input and timestep emb\n cudaCheck(hipMalloc(&(unet->input), B * C_in * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&(unet->timestep_emb), B * C_model * sizeof(float)));\n // malloc a nicely aligned buffer for the loss\n cudaCheck(hipMalloc(&(unet->gpu_loss), 64 * sizeof(float)));\n\n unet->params.mem_ptr = param_mem;\n unet->grads.mem_ptr = grad_mem;\n unet->acts.mem_ptr = act_mem;\n unet->back.back_mem = back_mem;\n\n point_unet_params(&(unet->config), &(unet->params));\n point_unet_params(&(unet->config), &(unet->grads));\n point_unet_acts(&(unet->config), &(unet->acts), unet->skip_acts, unet->input, unet->timestep_emb);\n set_unet_back_pointers(&(unet->back));\n}\n\n// ------------------------------------------------------------------------------------------------\n// forward pass\n\nvoid linear_layer_forward(\n Unet* unet,\n LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int N, int C, int OC, int block_size\n) {\n int layer = l_info->layer;\n LinearParams* linear_params = (LinearParams*)unet->params.layer_ptrs[layer];\n LinearActs* linear_acts = (LinearActs*)unet->acts.layer_ptrs[layer];\n matmul_forward2(cublas_handle, linear_acts->out, linear_acts->inp, linear_params->w, linear_params->b, N, C, OC, block_size);\n l_info->layer++;\n}\n\nvoid silu_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int N, int block_size\n) {\n SiluActs* acts = (SiluActs*)unet->acts.layer_ptrs[l_info->layer];\n silu_forward(acts->x, acts->out, N, block_size);\n l_info->layer++;\n}\n\nvoid convk3_layer_forward(\n Unet* unet, LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int B, int C, int OC, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n ConvK3Params* conv_params = (ConvK3Params*)unet->params.layer_ptrs[layer];\n ConvK3Acts* conv_acts = (ConvK3Acts*)unet->acts.layer_ptrs[layer];\n conv2d_k3_forward3(conv_acts->inp, conv_params->w, conv_params->b, conv_acts->out, B, C, OC, H, W);\n l_info->layer++;\n}\nvoid resblock_layer_forward(\n Unet* unet, LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int C_in, int C_emb, int C_out, int B, int H, int W, int gn_n_groups, int block_size\n) {\n int layer = l_info->layer;\n ResBlockParameters* resblock_params = (ResBlockParameters*)unet->params.layer_ptrs[layer];\n ResBlockActivations* resblock_acts = (ResBlockActivations*)unet->acts.layer_ptrs[layer];\n resblock_forward(cublas_handle, C_in, C_emb, C_out, B, H, W, block_size, 0, 0, gn_n_groups, resblock_params, resblock_acts);\n l_info->layer++;\n}\n\nvoid avgpool_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n AvgpoolActs* avgpool_acts = (AvgpoolActs*)unet->acts.layer_ptrs[layer];\n avgpool_2d_forward1(avgpool_acts->out, avgpool_acts->x, B, C, H, W, block_size);\n l_info->layer++;\n}\n\nvoid attention_layer_forward(\n Unet* unet, LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int B, int C, int H, int W, int HS, int gn_n_groups, int block_size\n) {\n int layer = l_info->layer;\n AttentionParams* att_params = (AttentionParams*)unet->params.layer_ptrs[layer];\n AttentionActs* att_acts = (AttentionActs*)unet->acts.layer_ptrs[layer];\n attention_block_forward(cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, att_params, att_acts);\n l_info->layer++;\n}\n\nvoid concat_layer_forward(\n Unet* unet,\n LayerInfo* l_info,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n ConcatChannelActs* cat_acts = (ConcatChannelActs*)unet->acts.layer_ptrs[layer];\n concat_channel_forward(cat_acts->x1, cat_acts->x2, cat_acts->out, B, C1, C2, H, W, block_size);\n l_info->layer++;\n}\n\nvoid upsample_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n UpsampleActs* up_acts = (UpsampleActs*)unet->acts.layer_ptrs[layer];\n upsample_forward1(up_acts->out, up_acts->x, B, C, H, W, block_size);\n l_info->layer++;\n}\n\nvoid groupnorm_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int gn_n_groups\n) {\n int layer = l_info->layer;\n GroupNormParams* gn_p = (GroupNormParams*)unet->params.layer_ptrs[layer];\n GroupNormActs* gn_a = (GroupNormActs*)unet->acts.layer_ptrs[layer];\n groupnorm_forward(gn_a->x, gn_p->w, gn_p->b, gn_a->out, gn_a->mean, gn_a->rstd, B, C, H, W, gn_n_groups);\n l_info->layer++;\n}\n\nvoid unet_forward(\n hipblasHandle_t cublas_handle,\n Unet* unet\n) {\n // common constants\n UnetConfig config = unet->config;\n int B = config.B;\n int C_in = config.C_in;\n int C_model = config.C_model;\n int C_time = config.C_time;\n int C_out = config.C_out;\n int H = config.H;\n int W = config.W;\n int block_size = config.block_size;\n int gn_n_groups = config.gn_n_groups;\n int HS = config.HS;\n int n_levels = config.n_levels;\n int n_res = config.n_res;\n int att_start_level = config.att_start_level;\n\n\n LayerInfo l_info;\n l_info.layer = 0;\n linear_layer_forward(unet, &l_info, cublas_handle, B, C_model, C_time, block_size);\n silu_layer_forward(unet, &l_info, B * C_time, block_size);\n linear_layer_forward(unet, &l_info, cublas_handle, B, C_time, C_time, block_size);\n int OC = C_model;\n convk3_layer_forward(unet, &l_info, cublas_handle, B, C_in, OC, H, W, block_size);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config.channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n resblock_layer_forward(unet, &l_info, cublas_handle, C, C_time, OC, B, H, W, gn_n_groups, block_size);\n C = OC;\n if (l >= att_start_level) {\n attention_layer_forward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n }\n }\n if (l < n_levels - 1) {\n avgpool_layer_forward(unet, &l_info, B, C, H, W, block_size);\n H /= 2;\n W /= 2;\n }\n }\n\n resblock_layer_forward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n attention_layer_forward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n resblock_layer_forward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n\n // upward\n for (int l = n_levels - 1; l >= 0; l--) {\n int C_skip = C_model * config.channel_mult[l];\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config.channel_mult[l - 1];\n }\n concat_layer_forward(unet, &l_info, B, C, C_skip, H, W, block_size);\n OC = C_model * config.channel_mult[l];\n resblock_layer_forward(unet, &l_info, cublas_handle, C + C_skip, C_time, OC, B, H, W, gn_n_groups, block_size);\n C = OC;\n if (l >= att_start_level) {\n attention_layer_forward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n }\n }\n if (l > 0) {\n upsample_layer_forward(unet, &l_info, B, C, H, W, block_size);\n H *= 2;\n W *= 2;\n }\n }\n\n groupnorm_layer_forward(unet, &l_info, B, C, H, W, gn_n_groups);\n silu_layer_forward(unet, &l_info, B * C * H * W, block_size);\n convk3_layer_forward(unet, &l_info, cublas_handle, B, C, C_out, H, W, block_size);\n\n ConvK3Acts* last_acts = (ConvK3Acts*)unet->acts.layer_ptrs[l_info.layer - 1];\n unet->output = last_acts->out;\n if (unet->target == nullptr) {\n // don't have targets, not computing loss\n unet->mean_loss = -1.0f;\n } else {\n mse_forward(unet->output, unet->target, unet->gpu_loss, B * C_out * H * W, block_size);\n cudaCheck(hipMemcpy(&(unet->mean_loss), unet->gpu_loss, sizeof(float), hipMemcpyDeviceToHost));\n }\n}\n\n// ------------------------------------------------------------------------------------------------\n// Unet backward\n\nvoid swap_layer_in_out_ptrs(LayerInfo* l_info) {\n float* temp = l_info->input;\n l_info->input = l_info->prev_out;\n l_info->prev_out = temp;\n}\n\nvoid linear_layer_backward(\n Unet* unet,\n LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int N, int C, int OC\n) {\n int layer = l_info->layer;\n LinearParams* linear_params = (LinearParams*)unet->params.layer_ptrs[layer];\n LinearParams* linear_grads = (LinearParams*)unet->grads.layer_ptrs[layer];\n LinearActs* linear_acts = (LinearActs*)unet->acts.layer_ptrs[layer];\n matmul_backward1(cublas_handle, l_info->prev_out, linear_grads->w, linear_grads->b, l_info->input, linear_acts->inp, linear_params->w, N, C, OC);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid silu_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int N, int block_size\n) {\n SiluActs* acts = (SiluActs*)unet->acts.layer_ptrs[l_info->layer];\n silu_backward(l_info->input, acts->x, l_info->prev_out, N, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid convk3_layer_backward(\n Unet* unet, LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int B, int C, int OC, int H, int W, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n add_inplace_forward(skip_dinp, l_info->input, B * OC * H * W, block_size);\n l_info->skip_idx--;\n }\n int layer = l_info->layer;\n ConvK3Params* conv_params = (ConvK3Params*)unet->params.layer_ptrs[layer];\n ConvK3Params* conv_grads = (ConvK3Params*)unet->grads.layer_ptrs[layer];\n ConvK3Acts* conv_acts = (ConvK3Acts*)unet->acts.layer_ptrs[layer];\n UnetBackActs* unet_back_acts = &(unet->back.back_acts);\n conv2d_k3_backward2(l_info->input, conv_acts->inp, conv_params->w, unet_back_acts->buf1_BCHW, unet_back_acts->buf2_BCHW, l_info->prev_out, conv_grads->w, conv_grads->b, B, C, OC, H, W);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid resblock_layer_backward(\n Unet* unet, LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int C_in, int C_emb, int C_out, int B, int H, int W, int gn_n_groups, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n add_inplace_forward(skip_dinp, l_info->input, B * C_out * H * W, block_size);\n l_info->skip_idx--;\n }\n int layer = l_info->layer;\n ResBlockParameters* resblock_params = (ResBlockParameters*)unet->params.layer_ptrs[layer];\n ResBlockParameters* resblock_grads = (ResBlockParameters*)unet->grads.layer_ptrs[layer];\n ResBlockActivations* resblock_acts = (ResBlockActivations*)unet->acts.layer_ptrs[layer];\n ResBlockBackwardActivations back_acts;\n UnetBackActs* unet_back_acts = &(unet->back.back_acts);\n\n // point back_acts to buffers\n back_acts.buf_BCemb = unet_back_acts->buf1_BCHW;\n back_acts.buf_BCHoWo = unet_back_acts->buf2_BCHW;\n back_acts.buf1_BCHW = unet_back_acts->buf3_BCHW;\n back_acts.buf2_BCHW = unet_back_acts->buf4_BCHW;\n back_acts.dweight_buf = unet_back_acts->buf1_B3CHW;\n back_acts.dbias_buf = unet_back_acts->buf2_B3CHW;\n back_acts.dout = l_info->input;\n back_acts.dx = l_info->prev_out;\n back_acts.demb = l_info->emb;\n\n resblock_backward(cublas_handle, C_in, C_emb, C_out, B, H, W, block_size, 0, 0, gn_n_groups, resblock_params, resblock_grads, resblock_acts, &back_acts);\n swap_layer_in_out_ptrs(l_info);\n\n // add demb to cumulative demb\n add_inplace_forward(l_info->emb, l_info->emb_cum, B * C_emb, block_size);\n l_info->layer--;\n}\n\nvoid attention_layer_backward(\n Unet* unet, LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int B, int C, int H, int W, int HS, int gn_n_groups, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n //printf(\"att adding skip %d\\n\", l_info->skip_idx);\n add_inplace_forward(skip_dinp, l_info->input, B * C * H * W, block_size);\n l_info->skip_idx--;\n }\n int layer = l_info->layer;\n AttentionParams* att_params = (AttentionParams*)unet->params.layer_ptrs[layer];\n AttentionParams* att_grads = (AttentionParams*)unet->grads.layer_ptrs[layer];\n AttentionActs* att_acts = (AttentionActs*)unet->acts.layer_ptrs[layer];\n AttentionBackwardActs att_backs;\n UnetBackActs* unet_backs = &(unet->back.back_acts);\n att_backs.buf1_BCHW = unet_backs->buf1_BCHW;\n att_backs.buf2_BCHW = unet_backs->buf2_BCHW;\n att_backs.buf_B3CHW = unet_backs->buf1_B3CHW;\n att_backs.dqkvr = unet_backs->buf2_B3CHW;\n att_backs.dpreatt = unet_backs->buf3_BCHW;\n att_backs.datt = unet_backs->buf4_BCHW;\n att_backs.dout = l_info->input;\n att_backs.dinp = l_info->prev_out;\n\n attention_block_backward(cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, att_params, att_acts, &att_backs, att_grads);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid concat_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n //printf(\"concat skip %d\\n\", l_info->skip_idx);\n concat_channel_backward(l_info->input, l_info->prev_out, skip_dinp, B, C1, C2, H, W, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n l_info->skip_idx++;\n}\n\n\nvoid avgpool_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n //printf(\"avgpool adding skip %d\\n\", l_info->skip_idx);\n add_inplace_forward(skip_dinp, l_info->input, B * C * H/2 * W/2, block_size);\n l_info->skip_idx--;\n }\n avgpool_2d_backward1(l_info->input, l_info->prev_out, B, C, H, W, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid upsample_layer_backward(\n Unet* unet, LayerInfo *l_info,\n int B, int C, int H, int W, int block_size\n) {\n upsample_backward1(l_info->prev_out, l_info->input, B, C, H, W, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid groupnorm_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int gn_n_groups\n) {\n int layer = l_info->layer;\n GroupNormParams* gn_p = (GroupNormParams*)unet->params.layer_ptrs[layer];\n GroupNormParams* gn_g = (GroupNormParams*)unet->grads.layer_ptrs[layer];\n GroupNormActs* gn_a = (GroupNormActs*)unet->acts.layer_ptrs[layer];\n groupnorm_backward(l_info->input, gn_a->x, gn_a->mean, gn_a->rstd, gn_p->w, l_info->prev_out, gn_g->w, gn_g->b, B, C, H, W, gn_n_groups);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid unet_backward(\n hipblasHandle_t cublas_handle,\n Unet* unet\n \n) {\n if (unet->mean_loss == -1.0f or unet->target == nullptr) {\n printf(\"Error: must run forward with targets before calling backward\\n\");\n exit(EXIT_FAILURE);\n }\n // common constants\n UnetConfig config = unet->config;\n int B = config.B;\n int C_in = config.C_in;\n int C_model = config.C_model;\n int C_time = config.C_time;\n int C_out = config.C_out;\n int H = config.H;\n int W = config.W;\n int block_size = config.block_size;\n int gn_n_groups = config.gn_n_groups;\n int HS = config.HS;\n int n_levels = config.n_levels;\n int n_res = config.n_res;\n int att_start_level = config.att_start_level;\n\n UnetBackActs* back_acts = &(unet->back.back_acts);\n mse_backward(unet->output, unet->target, back_acts->dinp_buf1, B * C_out * H * W, block_size);\n\n LayerInfo l_info;\n l_info.layer = config.n_layers - 1;\n l_info.input = back_acts->dinp_buf1;\n l_info.prev_out = back_acts->dinp_buf2;\n l_info.emb = back_acts->demb_buf1;\n l_info.emb_cum = back_acts->demb_buf2;\n l_info.skip_idx = 0;\n\n int C;\n int OC;\n int C_skip;\n C = C_model;\n OC = C_model;\n convk3_layer_backward(unet, &l_info, cublas_handle, B, C, C_out, H, W, block_size);\n silu_layer_backward(unet, &l_info, B * C * H * W, block_size);\n groupnorm_layer_backward(unet, &l_info, B, C, H, W, gn_n_groups);\n\n for (int l = 0; l < n_levels; l++) {\n for (int r = n_res; r >= 0; r--) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config.channel_mult[l - 1];\n } else {\n C_skip = C_model * config.channel_mult[l];\n }\n if (l >= att_start_level) {\n attention_layer_backward(unet, &l_info, cublas_handle, B, OC, H, W, HS, gn_n_groups, block_size);\n }\n if (r == 0 && l < n_levels - 1) {\n C = C_model * config.channel_mult[l + 1];\n }\n resblock_layer_backward(unet, &l_info, cublas_handle, C + C_skip, C_time, OC, B, H, W, gn_n_groups, block_size);\n concat_layer_backward(unet, &l_info, B, C, C_skip, H, W, block_size);\n }\n if (l < n_levels - 1) {\n H /= 2;\n W /= 2;\n upsample_layer_backward(unet, &l_info, B, C, H, W, block_size);\n OC = C_model * config.channel_mult[l + 1];\n }\n }\n\n resblock_layer_backward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n attention_layer_backward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n resblock_layer_backward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n\n // we will need to change the backward for all the other blocks below to include the dout contribution from the skip connection.\n l_info.skip_idx = NUM_SKIP_CONNECTIONS - 1;\n for (int l = n_levels - 1; l >= 0; l--) {\n if (l < n_levels - 1) {\n H *= 2;\n W *= 2;\n avgpool_layer_backward(unet, &l_info, B, C, H, W, block_size, 1);\n }\n for (int r = 0; r < n_res; r++) {\n if (l >= att_start_level) {\n attention_layer_backward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, 1);\n }\n int res_add_skip = (l < att_start_level) ? 1 : 0;\n if (r == n_res - 1 && l > 0) {\n int prev_mult = config.channel_mult[l - 1];\n C = prev_mult * C_model;\n }\n resblock_layer_backward(unet, &l_info, cublas_handle, C, C_time, OC, B, H, W, gn_n_groups, block_size, res_add_skip);\n OC = C;\n }\n }\n\n convk3_layer_backward(unet, &l_info, cublas_handle, B, C_in, OC, H, W, block_size, 1);\n unet->dinp = l_info.input;\n\n // reset dout to point to emb_cum\n l_info.input = l_info.emb_cum;\n linear_layer_backward(unet, &l_info, cublas_handle, B, C_time, C_time);\n silu_layer_backward(unet, &l_info, B * C_time, block_size);\n linear_layer_backward(unet, &l_info, cublas_handle, B, C_model, C_time);\n}\n\n// ------------------------------------------------------------------------------------------------\n// zero grad and ADAMW\n\nvoid unet_zero_grad(Unet* unet) {\n cudaCheck(hipMemset(unet->grads.mem_ptr, 0, unet->grads.total_size * sizeof(float)));\n cudaCheck(hipMemset(unet->back.back_mem, 0, unet->back.n_backs * sizeof(float)));\n //cudaCheck(cudaMemset(unet->acts.mem_ptr, 0, unet->acts.total_size * sizeof(float)));\n}\n\n// below copied from llm.c\n\n// Implements linear interpolation using only two floating-point operations (as opposed to three in a naive implementation).\n// Reference: https://developer.nvidia.com/blog/lerp-faster-cuda\n__device__ inline float lerp(float start, float end, float weight) {\n return fma(weight, end, fma(-weight, start, start));\n}\n\n__global__ void adamw_kernel2(float* params_memory, float* grads_memory, float* m_memory, float* v_memory, long num_parameters,\n float learning_rate, float beta1, float beta2, float beta1_correction, float beta2_correction, float eps, float weight_decay) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i >= num_parameters) return; // guard\n float grad = grads_memory[i];\n float m = m_memory[i];\n float v = v_memory[i];\n // update the first moment (momentum)\n m = lerp(grad, m, beta1);\n m_memory[i] = m;\n // update the second moment (RMSprop)\n v = lerp(grad * grad, v, beta2);\n v_memory[i] = v;\n m /= beta1_correction; // m_hat\n v /= beta2_correction; // v_hat\n params_memory[i] -= learning_rate * (m / (sqrtf(v) + eps) + weight_decay * params_memory[i]);\n}\n\nvoid unet_update(Unet* unet, float learning_rate, float beta1, float beta2, float eps, float weight_decay, int t) {\n // lazy allocation\n if (unet->m_memory == nullptr) {\n size_t param_size = unet->params.total_size * sizeof(float);\n cudaCheck(hipMalloc(&(unet->m_memory), param_size));\n cudaCheck(hipMalloc(&(unet->v_memory), param_size));\n cudaCheck(hipMemset(unet->m_memory, 0, param_size));\n cudaCheck(hipMemset(unet->v_memory, 0, param_size));\n printf(\"allocated %zu MiB for AdamW state m\\n\", param_size / 1024 / 1024);\n printf(\"allocated %zu MiB for AdamW state v\\n\", param_size / 1024 / 1024);\n }\n\n int block_size = unet->config.block_size;\n int n_blocks = ceil_div(unet->params.total_size, (size_t)block_size);\n float beta1_correction = 1.0f - powf(beta1, t);\n float beta2_correction = 1.0f - powf(beta2, t);\n adamw_kernel2<<>>(unet->params.mem_ptr, unet->grads.mem_ptr, unet->m_memory, unet->v_memory, unet->params.total_size,\n learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay);\n cudaCheck(hipGetLastError());\n}\n\n// ------------------------------------------------------------------------------------------------\n// save model states\n\nvoid save_unet_states(Unet* unet, GaussianDiffusion* diffusion, const char* filename) {\n // create header\n int model_header[256];\n model_header[0] = 12345678; // magic number\n model_header[1] = unet->config.B;\n model_header[2] = unet->config.C_in;\n model_header[3] = unet->config.C_model;\n model_header[4] = unet->config.C_out;\n model_header[5] = unet->config.H;\n model_header[6] = unet->config.W;\n model_header[7] = unet->config.max_period;\n model_header[8] = 1; // whether to save the AdamW states, this will be 0 if the model is saved from pytorch\n model_header[9] = 1; // whether to save rng states\n\n FILE* model_file = fopenCheck(filename, \"wb\");\n printf(\"\\nsaving model details\");\n int result = fwrite(model_header, sizeof(int), 256, model_file);\n assert(result == 256);\n\n // save params\n float* params_cpu_buf = (float*)mallocCheck(unet->params.total_size * sizeof(float));\n cudaCheck(hipMemcpy(params_cpu_buf, unet->params.mem_ptr, unet->params.total_size * sizeof(float), hipMemcpyDeviceToHost));\n printf(\"saving model params\\n\");\n result = fwrite(params_cpu_buf, sizeof(float), unet->params.total_size, model_file);\n assert(result == unet->params.total_size);\n\n if (model_header[8]) {\n // save AdamW states\n cudaCheck(hipMemcpy(params_cpu_buf, unet->m_memory, unet->params.total_size * sizeof(float), hipMemcpyDeviceToHost));\n printf(\"saving AdamW state m\\n\");\n result = fwrite(params_cpu_buf, sizeof(float), unet->params.total_size, model_file);\n assert(result == unet->params.total_size);\n cudaCheck(hipMemcpy(params_cpu_buf, unet->v_memory, unet->params.total_size * sizeof(float), hipMemcpyDeviceToHost));\n printf(\"saving AdamW state v\\n\");\n result = fwrite(params_cpu_buf, sizeof(float), unet->params.total_size, model_file);\n assert(result == unet->params.total_size);\n }\n\n if (model_header[9]) {\n // save rngs\n int n_states = diffusion->B * diffusion->img_size;\n void *curand_state_buf = mallocCheck(n_states * sizeof(hiprandState_t));\n cudaCheck(hipMemcpy(curand_state_buf, diffusion->curand_states, diffusion->B * diffusion->img_size * sizeof(hiprandState_t), hipMemcpyDeviceToHost));\n printf(\"saving rng states\\n\");\n result = fwrite(curand_state_buf, sizeof(hiprandState_t), n_states, model_file);\n assert(result == n_states);\n free(curand_state_buf);\n }\n\n fcloseCheck(model_file);\n free(params_cpu_buf);\n printf(\"saved model and diffusion states\\n\\n\");\n}\n\n// ------------------------------------------------------------------------------------------------\n// load unet and diffusion states\n\nvoid load_unet_and_diffusion_states(\n Unet* unet, GaussianDiffusion* diffusion, const char* filename\n) {\n FILE *model_file = fopenCheck(filename, \"rb\");\n // read header\n printf(\"\\nloading model header\\n\");\n int model_header[256];\n freadCheck(model_header, sizeof(int), 256, model_file);\n if (model_header[0] != 12345678) {\n fprintf(stderr, \"Invalid magic model file\\n\");\n exit(1);\n }\n int B = model_header[1];\n int C_in = model_header[2];\n int C_model = model_header[3];\n int C_out = model_header[4];\n int H = model_header[5];\n int W = model_header[6];\n int max_period = model_header[7];\n int save_adamw = model_header[8];\n int save_rng = model_header[9];\n\n // TODO: cleanup this code\n int* channel_mult = (int*)mallocCheck(4 * sizeof(int));\n channel_mult[0] = 1;\n channel_mult[1] = 2;\n channel_mult[2] = 3;\n channel_mult[3] = 4;\n\n // setup unet config\n UnetConfig* config = &(unet->config);\n config->B = B;\n config->C_in = C_in;\n config->C_model = C_model;\n config->C_time = C_model * 4;\n config->C_out = C_out;\n config->H = H;\n config->W = W;\n config->block_size = 512;\n config->gn_n_groups = 32;\n config->n_levels = 4;\n config->channel_mult = channel_mult;\n config->HS = 32;\n config->n_res = 2;\n config->max_period = max_period;\n config->att_start_level = 2;\n\n unet_init(unet);\n unet_count_layers(config);\n unet_make_ptrs_and_count_memory(unet);\n malloc_and_point(unet);\n\n // read params\n float* params_mem_cpu = (float*)mallocCheck(unet->params.total_size * sizeof(float));\n printf(\"loading model params\\n\");\n freadCheck(params_mem_cpu, sizeof(float), unet->params.total_size, model_file);\n cudaCheck(hipMemcpy(unet->params.mem_ptr, params_mem_cpu, unet->params.total_size * sizeof(float), hipMemcpyHostToDevice));\n\n if (save_adamw) {\n // read AdamW states\n printf(\"loading AdamW states\\n\");\n freadCheck(params_mem_cpu, sizeof(float), unet->params.total_size, model_file);\n cudaCheck(hipMalloc(&(unet->m_memory), unet->params.total_size * sizeof(float)));\n cudaCheck(hipMemcpy(unet->m_memory, params_mem_cpu, unet->params.total_size * sizeof(float), hipMemcpyHostToDevice));\n freadCheck(params_mem_cpu, sizeof(float), unet->params.total_size, model_file);\n cudaCheck(hipMalloc(&(unet->v_memory), unet->params.total_size * sizeof(float)));\n cudaCheck(hipMemcpy(unet->v_memory, params_mem_cpu, unet->params.total_size * sizeof(float), hipMemcpyHostToDevice));\n } else {\n printf(\"not loading AdamW states\\n\");\n unet->m_memory = nullptr;\n unet->v_memory = nullptr;\n }\n\n if (save_rng) {\n // read curand rng states\n printf(\"loading rng states\\n\");\n int n_states = B * C_in * H * W;\n void *curand_state_buf = mallocCheck(n_states * sizeof(hiprandState_t));\n freadCheck(curand_state_buf, sizeof(hiprandState_t), n_states, model_file);\n // initialize diffusion curant buffer\n cudaCheck(hipMalloc(&(diffusion->curand_states), n_states * sizeof(hiprandState_t)));\n cudaCheck(hipMemcpy(diffusion->curand_states, curand_state_buf, n_states * sizeof(hiprandState_t), hipMemcpyHostToDevice));\n free(curand_state_buf);\n } else {\n printf(\"not loading rng states\\n\");\n diffusion->curand_states = nullptr;\n }\n diffusion_init(diffusion, B, C_in, H, W, 1000);\n\n free(params_mem_cpu);\n fcloseCheck(model_file);\n printf(\"loaded model and diffusion states \\n\\n\");\n}\n\n\n\n// ------------------------------------------------------------------------------------------------\n\ntypedef struct {\n int count;\n float mean_loss;\n} LossCounter;\n\nvoid update_loss_counter(LossCounter* counter, float loss) {\n float old_mean = counter->mean_loss;\n float old_count = (float)counter->count;\n counter->mean_loss = old_mean * old_count / (old_count + 1) + loss / (old_count + 1);\n counter->count++;\n}\n\nvoid clear_loss_counter(LossCounter* counter) {\n counter->count = 0;\n counter->mean_loss = 0.0f;\n}\n\nvoid create_directory(const char *dir) {\n struct stat st = {0};\n if (stat(dir, &st) == -1) {\n mkdir(dir, 0700); // Create directory with permissions 0700\n }\n}\n\nint main(int argc, char **argv) {\n // Default values\n const char *model_weights_file = \"unet_init.bin\";\n const char *data_file = \"data/elephant_train.bin\";\n const char *log_filename = \"log.txt\";\n\n // Parse command line arguments\n for (int i = 1; i < argc; ++i) {\n if (strcmp(argv[i], \"--model_weights\") == 0 && i + 1 < argc) {\n model_weights_file = argv[++i];\n } else if (strcmp(argv[i], \"--data_file\") == 0 && i + 1 < argc) {\n data_file = argv[++i];\n } else if (strcmp(argv[i], \"--log_file\") == 0 && i + 1 < argc) {\n log_filename = argv[++i];\n }\n }\n printf(\"loading model weights from %s\\n\", model_weights_file);\n printf(\"loading data from %s\\n\", data_file);\n\n // set up the device\n int deviceIdx = 0;\n cudaCheck(hipSetDevice(deviceIdx));\n hipDeviceProp_t deviceProp;\n hipGetDeviceProperties(&deviceProp, deviceIdx);\n printf(\"Device %d: %s\\n\", deviceIdx, deviceProp.name);\n printf(\"device registers per block: %d\\n\", deviceProp.regsPerBlock);\n printf(\"shared mem per block: %ld\\n\", deviceProp.sharedMemPerBlock);\n\n // setup cublas\n hipblasHandle_t cublas_handle;\n cublasCheck(hipblasCreate(&cublas_handle));\n cublasCheck(hipblasSetMathMode(cublas_handle, HIPBLAS_DEFAULT_MATH));\n\n Unet unet;\n GaussianDiffusion diffusion;\n load_unet_and_diffusion_states(&unet, &diffusion, model_weights_file);\n\n int B = unet.config.B;\n int C_in = unet.config.C_in;\n int C_out = unet.config.C_out;\n int H = unet.config.H;\n int W = unet.config.W;\n\n DataLoader loader;\n dataloader_init(&loader, data_file, B);\n\n // allocate memory for target separately, since unet can still be run without targets for inference\n cudaCheck(hipMalloc(&(unet.target), B * C_out * H * W * sizeof(float)));\n float *d_timesteps_buf;\n cudaCheck(hipMalloc(&d_timesteps_buf, B * sizeof(float)));\n\n\n int n_iters = 100000;\n int n_log_iter = 100;\n int n_save_iter = 10000;\n //int n_iters = 100;\n //int n_log_iter = 5;\n //int n_save_iter = 50;\n\n // wipe existing logs\n FILE* log_file = fopen(log_filename, \"w\");\n fprintf(log_file, \"training unet,\\nmodel file %s,\\ndata filename %s\\n\", model_weights_file, data_file);\n fprintf(log_file, \"starting training for %d iterations\\n\", n_iters);\n fcloseCheck(log_file);\n\n const char *model_save_dir = \"./models/\";\n create_directory(model_save_dir);\n\n LossCounter loss_counter;\n clear_loss_counter(&loss_counter);\n\n printf(\"Training for %d iterations\\n\", n_iters);\n float time_elapsed = 0.0;\n hipEvent_t start, end;\n cudaCheck(hipEventCreate(&start));\n cudaCheck(hipEventCreate(&end));\n cudaCheck(hipEventRecord(start));\n cudaCheck(hipEventSynchronize(start));\n for (int iter = 1; iter <= n_iters; iter++) {\n unet_zero_grad(&unet);\n dataloader_next_batch(&loader);\n // copy loaded data into input\n cudaCheck(hipMemcpy(unet.input, loader.input, B * C_in * H * W * sizeof(float), hipMemcpyHostToDevice));\n\n // sample timesteps and get timestep embeddings\n sample_timesteps(&diffusion, d_timesteps_buf);\n get_timestep_embeddings(&unet.time_emb_state, d_timesteps_buf, unet.timestep_emb);\n \n // target will be injected random noise\n diffusion_draw_normal(&diffusion, unet.target);\n\n // push input through forward diffusion process by t to get x_t\n diffusion_forward_by_t(&diffusion, unet.input, d_timesteps_buf, unet.target);\n\n unet_forward(cublas_handle, &unet);\n unet_backward(cublas_handle, &unet);\n unet_update(&unet, 0.0001f, 0.9f, 0.999f, 1e-8f, 0.0f, iter);\n\n update_loss_counter(&loss_counter, unet.mean_loss);\n\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(end));\n cudaCheck(hipEventElapsedTime(&time_elapsed, start, end));\n\n if (iter % n_log_iter == 0) {\n printf(\"step %4d/%d | loss %7.6f | mean loss %7.6f | cur time %.4f s\\n\", iter, n_iters, unet.mean_loss, loss_counter.mean_loss, time_elapsed / 1000);\n FILE *log_file = fopen(log_filename, \"a\");\n fprintf(log_file, \"step %4d/%d | loss %7.6f | mean loss %7.6f | cur time %.4f s\\n\", iter, n_iters, unet.mean_loss, loss_counter.mean_loss, time_elapsed / 1000);\n fclose(log_file);\n clear_loss_counter(&loss_counter);\n }\n\n if (iter % n_save_iter == 0) {\n char save_filename[256];\n sprintf(save_filename, \"./models/model_%d.bin\", iter);\n save_unet_states(&unet, &diffusion, save_filename);\n }\n }\n printf(\"average time per iteration: %f ms\\n\", time_elapsed / n_iters);\n\n cudaCheck(hipEventDestroy(start));\n cudaCheck(hipEventDestroy(end));\n cublasCheck(hipblasDestroy(cublas_handle));\n cudaCheck(hipFree(d_timesteps_buf));\n cudaCheck(hipFree(unet.target));\n\n dataloader_free(&loader);\n diffusion_free(&diffusion);\n free_unet(&unet);\n\n printf(\"success\\n\");\n}", "source": "/kaggle/working/unet.cu/train_unet.cu"} {"input": "#include \n#include \n#include \n#include \"common.h\"\n#include \"avgpool.cuh\"\n\n// ----------------------------------------------------------------------------\n// GPU kernels\n\n__global__ void avgpool_2d_forward_kernel(\n const float* x, // (B, C, H, W)\n float* out, // (B, C, H/2, W/2)\n int B, int C, int H, int W\n) {\n //printf(\"started kernel\\n\");\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int H_2 = H / 2;\n int W_2 = W / 2;\n int img_out_size = H_2 * W_2;\n if (flat_idx >= B * C * img_out_size) {\n return;\n }\n int img_size = H * W;\n int b = flat_idx / (C * img_out_size);\n int c = (flat_idx / img_out_size) % C;\n int i = (flat_idx / W_2) % H_2;\n int j = flat_idx % W_2;\n\n // move pointers\n x += b * C * img_size + c * img_size + 2 * i * W + 2 * j;\n out += b * C * img_out_size + c * img_out_size + i * W_2 + j;\n\n float sum = 0.0f;\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n sum += x[ii * W + jj];\n }\n }\n out[0] = sum / 4.0f;\n}\n\n\n__global__ void avgpool_2d_backward_kernel(\n const float* dout,\n float* dx,\n int B, int C, int H_out, int W_out // H_out and W_out are the shape of the output from forward, i.e. half the original sizes\n) {\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int img_out_size = H_out * W_out;\n if (flat_idx >= B * C * img_out_size) {\n return;\n }\n int H_in = H_out * 2;\n int W_in = W_out * 2;\n int img_in_size = H_in * W_in;\n int b = flat_idx / (C * img_out_size);\n int c = (flat_idx / img_out_size) % C;\n int i = (flat_idx / W_out) % H_out;\n int j = flat_idx % W_out;\n\n // move pointers\n dout += b * C * img_out_size + c * img_out_size + i * W_out + j;\n dx += b * C * img_in_size + c * img_in_size + 2 * i * W_in + 2 * j;\n\n float d = dout[0] / 4.0f;\n\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n dx[ii * W_in + jj] = d;\n }\n }\n}\n\n// ----------------------------------------------------------------------------\n// kernel launcher\n\nvoid avgpool_2d_forward1(\n float* out,\n const float* x,\n int B, int C, int H, int W,\n const int block_size\n) {\n int gridDim = ceil_div(B * C * (H / 2) * (W / 2), block_size);\n avgpool_2d_forward_kernel<<>>(x, out, B, C, H, W);\n cudaCheck(cudaGetLastError());\n}\n\nvoid avgpool_2d_backward1(\n const float* dout,\n float* dx,\n int B, int C, int H, int W,\n const int block_size\n) {\n int H_out = H / 2;\n int W_out = W / 2;\n int gridDim = ceil_div(B * C * H_out * W_out, block_size);\n avgpool_2d_backward_kernel<<>>(dout, dx, B, C, H_out, W_out);\n cudaCheck(cudaGetLastError());\n}\n\n// ----------------------------------------------------------------------------\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n srand(0);\n int B = 1;\n int C = 64;\n int H = 32;\n int W = 32;\n\n // set up the device\n int deviceIdx = 0;\n cudaCheck(cudaSetDevice(deviceIdx));\n cudaDeviceProp deviceProp;\n cudaGetDeviceProperties(&deviceProp, deviceIdx);\n printf(\"Device %d: %s\\n\", deviceIdx, deviceProp.name);\n\n // create host memory to load data\n float* x = (float*)malloc(B * C * H * W * sizeof(float));\n float* out = (float*)malloc(B * C * H/2 * W/2 * sizeof(float));\n float* dx = (float*)malloc(B * C * H * W * sizeof(float));\n float* dout = (float*)malloc(B * C * H/2 * W/2 * sizeof(float));\n\n // read saved output\n FILE *file = fopen(\"down.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(x, sizeof(float), B * C * H * W, file);\n freadCheck(out, sizeof(float), B * C * H/2 * W/2, file);\n freadCheck(dout, sizeof(float), B * C * H/2 * W/2, file);\n freadCheck(dx, sizeof(float), B * C * H * W, file);\n fclose(file);\n\n\n // allocate device memory\n float *d_x, *d_out, *d_dout, *d_dx;\n cudaCheck(cudaMalloc(&d_x, B * C * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&d_out, B * C * H/2 * W/2 * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dout, B * C * H/2 * W/2 * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dx, B * C * H * W * sizeof(float)));\n cudaCheck(cudaMemcpy(d_x, x, B * C * H * W * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_dout, dout, B * C * H/2 * W/2 * sizeof(float), cudaMemcpyHostToDevice));\n\n printf(\"Checking forward pass\\n\");\n int block_sizes[] = {32, 64, 128};\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"Checking block size %d\\n\", block_size);\n avgpool_2d_forward1(d_out, d_x, B, C, H, W, block_size);\n validate_result(d_out, out, \"out\", B * C * (H/2) * (W/2));\n }\n printf(\"Forward pass: all results match\\n\\n\");\n\n printf(\"Checking backward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"Checking block size %d\\n\", block_size);\n avgpool_2d_backward1(d_dout, d_dx, B, C, H, W, block_size);\n validate_result(d_dx, dx, \"dx\", B * C * H * W);\n }\n printf(\"Backward pass: all results match\\n\\n\");\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, avgpool_2d_forward1,\n d_out, d_x,\n B, C, H, W, block_size);\n\n float tflops = (float)B * H * W * C / elapsed_time * 1e3f / 1e12f;\n printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n }\n\n printf(\"\\nBackward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, avgpool_2d_backward1,\n d_dout, d_dx,\n B, C, H, W, block_size);\n\n float tflops = (float)B * H * W * C / elapsed_time * 1e3f / 1e12f;\n printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n }\n\n // free memory\n free(x);\n free(out);\n free(dout);\n free(dx);\n cudaFree(d_x);\n cudaFree(d_out);\n cudaFree(d_dout);\n cudaFree(d_dx);\n\n return 0;\n}\n#endif", "output": "#include \n#include \n#include \n#include \"common.h\"\n#include \"avgpool.cuh\"\n\n// ----------------------------------------------------------------------------\n// GPU kernels\n\n__global__ void avgpool_2d_forward_kernel(\n const float* x, // (B, C, H, W)\n float* out, // (B, C, H/2, W/2)\n int B, int C, int H, int W\n) {\n //printf(\"started kernel\\n\");\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int H_2 = H / 2;\n int W_2 = W / 2;\n int img_out_size = H_2 * W_2;\n if (flat_idx >= B * C * img_out_size) {\n return;\n }\n int img_size = H * W;\n int b = flat_idx / (C * img_out_size);\n int c = (flat_idx / img_out_size) % C;\n int i = (flat_idx / W_2) % H_2;\n int j = flat_idx % W_2;\n\n // move pointers\n x += b * C * img_size + c * img_size + 2 * i * W + 2 * j;\n out += b * C * img_out_size + c * img_out_size + i * W_2 + j;\n\n float sum = 0.0f;\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n sum += x[ii * W + jj];\n }\n }\n out[0] = sum / 4.0f;\n}\n\n\n__global__ void avgpool_2d_backward_kernel(\n const float* dout,\n float* dx,\n int B, int C, int H_out, int W_out // H_out and W_out are the shape of the output from forward, i.e. half the original sizes\n) {\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int img_out_size = H_out * W_out;\n if (flat_idx >= B * C * img_out_size) {\n return;\n }\n int H_in = H_out * 2;\n int W_in = W_out * 2;\n int img_in_size = H_in * W_in;\n int b = flat_idx / (C * img_out_size);\n int c = (flat_idx / img_out_size) % C;\n int i = (flat_idx / W_out) % H_out;\n int j = flat_idx % W_out;\n\n // move pointers\n dout += b * C * img_out_size + c * img_out_size + i * W_out + j;\n dx += b * C * img_in_size + c * img_in_size + 2 * i * W_in + 2 * j;\n\n float d = dout[0] / 4.0f;\n\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n dx[ii * W_in + jj] = d;\n }\n }\n}\n\n// ----------------------------------------------------------------------------\n// kernel launcher\n\nvoid avgpool_2d_forward1(\n float* out,\n const float* x,\n int B, int C, int H, int W,\n const int block_size\n) {\n int gridDim = ceil_div(B * C * (H / 2) * (W / 2), block_size);\n avgpool_2d_forward_kernel<<>>(x, out, B, C, H, W);\n cudaCheck(hipGetLastError());\n}\n\nvoid avgpool_2d_backward1(\n const float* dout,\n float* dx,\n int B, int C, int H, int W,\n const int block_size\n) {\n int H_out = H / 2;\n int W_out = W / 2;\n int gridDim = ceil_div(B * C * H_out * W_out, block_size);\n avgpool_2d_backward_kernel<<>>(dout, dx, B, C, H_out, W_out);\n cudaCheck(hipGetLastError());\n}\n\n// ----------------------------------------------------------------------------\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n srand(0);\n int B = 1;\n int C = 64;\n int H = 32;\n int W = 32;\n\n // set up the device\n int deviceIdx = 0;\n cudaCheck(hipSetDevice(deviceIdx));\n hipDeviceProp_t deviceProp;\n hipGetDeviceProperties(&deviceProp, deviceIdx);\n printf(\"Device %d: %s\\n\", deviceIdx, deviceProp.name);\n\n // create host memory to load data\n float* x = (float*)malloc(B * C * H * W * sizeof(float));\n float* out = (float*)malloc(B * C * H/2 * W/2 * sizeof(float));\n float* dx = (float*)malloc(B * C * H * W * sizeof(float));\n float* dout = (float*)malloc(B * C * H/2 * W/2 * sizeof(float));\n\n // read saved output\n FILE *file = fopen(\"down.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(x, sizeof(float), B * C * H * W, file);\n freadCheck(out, sizeof(float), B * C * H/2 * W/2, file);\n freadCheck(dout, sizeof(float), B * C * H/2 * W/2, file);\n freadCheck(dx, sizeof(float), B * C * H * W, file);\n fclose(file);\n\n\n // allocate device memory\n float *d_x, *d_out, *d_dout, *d_dx;\n cudaCheck(hipMalloc(&d_x, B * C * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&d_out, B * C * H/2 * W/2 * sizeof(float)));\n cudaCheck(hipMalloc(&d_dout, B * C * H/2 * W/2 * sizeof(float)));\n cudaCheck(hipMalloc(&d_dx, B * C * H * W * sizeof(float)));\n cudaCheck(hipMemcpy(d_x, x, B * C * H * W * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_dout, dout, B * C * H/2 * W/2 * sizeof(float), hipMemcpyHostToDevice));\n\n printf(\"Checking forward pass\\n\");\n int block_sizes[] = {32, 64, 128};\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"Checking block size %d\\n\", block_size);\n avgpool_2d_forward1(d_out, d_x, B, C, H, W, block_size);\n validate_result(d_out, out, \"out\", B * C * (H/2) * (W/2));\n }\n printf(\"Forward pass: all results match\\n\\n\");\n\n printf(\"Checking backward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"Checking block size %d\\n\", block_size);\n avgpool_2d_backward1(d_dout, d_dx, B, C, H, W, block_size);\n validate_result(d_dx, dx, \"dx\", B * C * H * W);\n }\n printf(\"Backward pass: all results match\\n\\n\");\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, avgpool_2d_forward1,\n d_out, d_x,\n B, C, H, W, block_size);\n\n float tflops = (float)B * H * W * C / elapsed_time * 1e3f / 1e12f;\n printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n }\n\n printf(\"\\nBackward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, avgpool_2d_backward1,\n d_dout, d_dx,\n B, C, H, W, block_size);\n\n float tflops = (float)B * H * W * C / elapsed_time * 1e3f / 1e12f;\n printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n }\n\n // free memory\n free(x);\n free(out);\n free(dout);\n free(dx);\n hipFree(d_x);\n hipFree(d_out);\n hipFree(d_dout);\n hipFree(d_dx);\n\n return 0;\n}\n#endif", "source": "/kaggle/working/unet.cu/dev/avgpool.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \"conv2d_k3.cuh\"\n\n\n__global__ void permute_w_tile_x_kernel1(\n const float* x, const float* weight,\n float* x_tiled, float* weight_perm,\n int B, int C_in, int C_out, int H, int W\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx < C_out * C_in * 9) {\n // permute weight\n int c_out = idx / (C_in * 9);\n int c_in = (idx / 9) % C_in;\n int k = idx % 9;\n int out_ind = k * C_out * C_in + c_out * C_in + c_in;\n weight_perm[out_ind] = weight[idx];\n }\n\n if (idx < B * C_in * H * W) {\n // tile x to shape (9, C_in, B, H, W)\n // for a given first dimension k, the entry at an index should be\n // what the k-th kernel is multiplied with in the convolution for that pixel\n // (not a good explanation... but obvious with a picture)\n // for example, a value at (H, W) = (0, 0) will never be convolved with kernels 5, 7 and 8\n // because it is on a corner\n int x_offset = B * C_in * H * W;\n int img_size = H * W;\n int b = idx / (C_in * img_size);\n int c = (idx / img_size) % C_in;\n int h_start = (idx / W) % H;\n int w_start = idx % W;\n float x_val = x[idx];\n for (int i = 0; i < 3; i++) {\n for (int j = 0; j < 3; j++) {\n int h_end = h_start + 1 - i;\n int w_end = w_start + 1 - j;\n if (h_end >= 0 && h_end < H && w_end >= 0 && w_end < W) {\n x_tiled[c * B * img_size + b * img_size + h_end * W + w_end] = x_val;\n }\n // shift pointer to next kernel\n x_tiled += x_offset;\n }\n }\n }\n}\n\n\n__global__ void reduce_transpose_add_bias_kernel1(\n const float* out_tmp, const float* bias,\n float* out,\n int B, int C_out, int H, int W\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n // make sure we're in bound\n assert(idx < B * C_out * H * W);\n // make sure every block lies in the same channel\n assert((B * H * W) % blockDim.x == 0);\n\n int img_size = H * W;\n int c = idx / (B * img_size);\n int b = (idx / img_size) % B;\n int h = (idx / W) % H;\n int w = idx % W;\n\n __shared__ float bias_sh[1];\n if (threadIdx.x == 0) {\n bias_sh[0] = bias[c];\n }\n __syncthreads();\n\n float out_val = bias_sh[0];\n\n for (int k = 0; k < 9; k++) {\n float out_tmp_val = out_tmp[k * C_out * B * img_size + c * B * img_size + b * img_size + h * W + w];\n out_val += out_tmp_val;\n }\n\n out[b * C_out * img_size + c * img_size + h * W + w] = out_val;\n}\n\nvoid conv2d_k3_forward1(\n cublasHandle_t cublas_handle,\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C_in, int C_out, int H, int W,\n int block_size\n) {\n // implements Conv2D with kernel size 1 and padding 1\n // shapes:\n // x (B, C_in, H, W)\n // weight (C_out, C_in, 3, 3)\n // bias (C_out,)\n\n // plan:\n // 1. transpose and tile x into x_tiled of shape (9, C_in, B, H, W)\n // 2. transpose weight into shape (9, C_out, C_in) (we will do this in the same kernel as one above)\n // 3. call cublasSgemmStridedBatched with 9 batches\n // 4. reduce over the 9 kernel dimensions, add bias, and transpose output to shape (B, C_out, H, W)\n\n // 1. and 2.\n // transpose weight and tile x\n float *x_tiled, *weight_perm;\n cudaCheck(cudaMalloc(&x_tiled, 9 * C_in * B * H * W * sizeof(float)));\n cudaCheck(cudaMemset(x_tiled, 0, 9 * C_in * B * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&weight_perm, 9 * C_out * C_in * sizeof(float)));\n int total_th_1 = max_int(B * C_in * H * W, C_in * C_out * 9);\n int n_blk_1 = ceil_div(total_th_1, block_size);\n permute_w_tile_x_kernel1<<>>(x, weight, x_tiled, weight_perm, B, C_in, C_out, H, W);\n\n // 3. matmul via cublasSgemmStridedBatched\n float* out_tmp;\n cudaCheck(cudaMalloc(&out_tmp, 9 * C_out * B * H * W * sizeof(float)));\n float alpha = 1.0f;\n float beta = 0.0f;\n int bhw = B * H * W;\n cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, bhw, C_out, C_in, &alpha, x_tiled, bhw, C_in * bhw, weight_perm, C_in, C_out * C_in, &beta, out_tmp, bhw, C_out * bhw, 9));\n\n // reduce, add bias, and transpose\n int total_th_2 = B * C_out * H * W;\n int n_blk_2 = ceil_div(total_th_2, block_size);\n reduce_transpose_add_bias_kernel1<<>>(out_tmp, bias, out, B, C_out, H, W);\n\n // free memory\n cudaCheck(cudaFree(x_tiled));\n cudaCheck(cudaFree(weight_perm));\n cudaCheck(cudaFree(out_tmp));\n}\n\n\n__global__ void transpose_tile_x_kernel2(\n const float* x,\n float* x_tiled,\n int B, int C_in, int H, int W\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * C_in * H * W) {\n // similar to permute_w_tile_x_kernel1, except now the target shape is (C_in, 9, B, H, W)\n int img_size = H * W;\n int b = idx / (C_in * img_size);\n int c = (idx / img_size) % C_in;\n int h_start = (idx / W) % H;\n int w_start = idx % W;\n float x_val = x[idx];\n for (int i = 0; i < 3; i++) {\n for (int j = 0; j < 3; j++) {\n int h_end = h_start + 1 - i;\n int w_end = w_start + 1 - j;\n if (h_end >= 0 && h_end < H && w_end >= 0 && w_end < W) {\n x_tiled[c * 9 * B * img_size + (i * 3 + j) * B * img_size + b * img_size + h_end * W + w_end] = x_val;\n }\n }\n }\n }\n}\n\n__global__ void transpose_out_add_bias_kernel2(\n const float* out_tmp, const float* bias,\n float* out,\n int B, int C_out, int H, int W\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n // make sure we're in bound\n assert(idx < B * C_out * H * W);\n // make sure every block lies in the same channel\n assert((B * H * W) % blockDim.x == 0);\n\n int img_size = H * W;\n int c = idx / (B * img_size);\n int b = (idx / img_size) % B;\n int h = (idx / W) % H;\n int w = idx % W;\n\n __shared__ float bias_sh[1];\n if (threadIdx.x == 0) {\n bias_sh[0] = bias[c];\n }\n __syncthreads();\n \n out[b * C_out * img_size + c * img_size + h * W + w] = bias_sh[0] + out_tmp[idx];\n}\n\nvoid conv2d_k3_forward2(\n cublasHandle_t cublas_handle,\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C_in, int C_out, int H, int W,\n int block_size,\n float* t1, float* t2, float* t3\n) {\n // similar to kernel 1, but transpose into a different shape so that we can just make one big matmul call\n // and no need to reduce afterwards\n // plan:\n // 1. transpose and tile x into x_tiled of shape (C_in, 9, B, H, W)\n // 2. call cublasSgemm on x_tiled and weight, no need for transposes!\n // 3. add bias\n\n // 1. transpose x to x_tiled\n cudaEvent_t start, end1, end2, end3;\n if (t1 != nullptr) {\n cudaCheck(cudaEventCreate(&start));\n cudaCheck(cudaEventCreate(&end1));\n cudaCheck(cudaEventCreate(&end2));\n cudaCheck(cudaEventCreate(&end3));\n cudaCheck(cudaEventRecord(start, nullptr));\n }\n float* x_tiled;\n cudaCheck(cudaMalloc(&x_tiled, C_in * 9 * B * H * W * sizeof(float)));\n cudaCheck(cudaMemset(x_tiled, 0, C_in * 9 * B * H * W * sizeof(float)));\n if (t1 != nullptr) {\n cudaCheck(cudaEventRecord(end1, nullptr));\n }\n int total_th_1 = B * C_in * H * W;\n int n_blk_1 = ceil_div(total_th_1, block_size);\n transpose_tile_x_kernel2<<>>(x, x_tiled, B, C_in, H, W);\n\n // 2. matmul with cublasSgemm\n float* out_tmp;\n cudaCheck(cudaMalloc(&out_tmp, C_out * B * H * W * sizeof(float)));\n float alpha = 1.0f;\n float beta = 0.0f;\n int bhw = B * H * W;\n cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, bhw, C_out, C_in * 9, &alpha, x_tiled, bhw, weight, C_in * 9, &beta, out_tmp, bhw));\n\n\n // 3. transpose and add bias\n int total_th_2 = B * C_out * H * W;\n int n_blk_2 = ceil_div(total_th_2, block_size);\n transpose_out_add_bias_kernel2<<>>(out_tmp, bias, out, B, C_out, H, W);\n if (t2 != nullptr) {\n cudaCheck(cudaEventRecord(end2, nullptr));\n }\n\n\n // free memory\n cudaCheck(cudaFree(x_tiled));\n cudaCheck(cudaFree(out_tmp));\n if (t3 != nullptr) {\n cudaCheck(cudaEventRecord(end3, nullptr));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end1));\n cudaCheck(cudaEventSynchronize(end2));\n cudaCheck(cudaEventSynchronize(end3));\n float elapsed_time1, elapsed_time2, elapsed_time3;\n cudaCheck(cudaEventElapsedTime(&elapsed_time1, start, end1));\n cudaCheck(cudaEventElapsedTime(&elapsed_time2, end1, end2));\n cudaCheck(cudaEventElapsedTime(&elapsed_time3, end2, end3));\n *t1 += elapsed_time1;\n *t2 += elapsed_time2;\n *t3 += elapsed_time3;\n }\n}\n\n// this kernel will roughly follow kernel 5 here: https://github.com/siboehm/SGEMM_CUDA, with a few differences\n// the main operation is convolving the input X with the weights\n// weights are shaped (O, Cx9), and X is shaped (B, C, HxW)\n// We are morally doing a matmul with dimensions O, C, HxW\ntemplate \n__global__ void conv2d_k3_forward_kernel1(\n const float* x, const float* weight, const float* bias,\n float* out,\n const int B, const int C, const int O, const int H, const int W\n) {\n // constants\n const int C9 = C * 9;\n const int weight_size = O * C * 9;\n const int x_size = B * C * H * W;\n // n_threads / block = (BO * BH) / (TO * TH) = 256\n // n_blocks = B * (O / BO) * (H * W / BH)\n // we will set gridDim to (BxHxW / BH, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // for compute step, to determine which of the size TOxTH block the thread is responsible for\n const int threads_per_col = ceil_div(BH, TH);\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // shared mem for loading weights and X\n const int w_s_sz = BO * BC * 9;\n const int w_s_width = BC * 9;\n const int x_s_sz = BC * BH;\n __shared__ float w_s[w_s_sz];\n __shared__ float x_s[x_s_sz];\n\n // registers to store local results\n float thread_results[TO * TH * 9] = {0.0};\n // registers for weights and x\n float reg_w[TO] = {0.0};\n float reg_x[TH] = {0.0};\n\n // calculate offset in x for bounds checking later\n const int x_block_offset = batch * C * H * W + block_col * BH;\n // calculate weight offset for bounds checking\n const int w_block_offset = block_row * BO * C9;\n\n // outer loop over BC size blocks of channel C\n for (int c_idx = 0; c_idx < C; c_idx += BC) {\n // x and weight offsets for this block\n int x_inner_offset = c_idx * H * W;\n int w_inner_offset = c_idx * 9;\n // load weight into smem\n int w_abs_idx;\n for (int i = threadIdx.x; i < w_s_sz; i += blockDim.x) {\n int w_i = i / w_s_width;\n int w_j = i % w_s_width;\n w_abs_idx = w_block_offset + w_inner_offset + w_i * C9 + w_j;\n if (w_abs_idx < weight_size) {\n w_s[w_i * w_s_width + w_j] = weight[w_abs_idx];\n }\n }\n // load x into smem\n int x_abs_idx;\n for (int i = threadIdx.x; i < x_s_sz; i += blockDim.x) {\n int x_i = i / BH;\n int x_j = i % BH;\n x_abs_idx = x_block_offset + x_inner_offset + x_i * H * W + x_j;\n if (x_abs_idx < x_size) {\n x_s[x_i * BH + x_j] = x[x_abs_idx];\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int dot_idx = 0; dot_idx < BC; dot_idx++) {\n // load x_s into registers\n for (int j = 0; j < TH; j++) {\n reg_x[j] = x_s[dot_idx * BH + thread_col * TH + j];\n }\n for (int conv_idx = 0; conv_idx < 9; conv_idx++) {\n for (int i = 0; i < TO; i++) {\n reg_w[i] = w_s[(thread_row * TO + i) * w_s_width + dot_idx * 9 + conv_idx];\n }\n\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n float val = reg_w[i] * reg_x[j];\n thread_results[conv_idx * TO * TH + i * TH + j] += val;\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n // need to be careful here: for each pixel at (h_abs, w_abs) and its product with kernel k (k < 9)\n // There are 2 options\n // 1. the product should be added to some pixel, then do an atomiAdd\n // 2. the product should not contribute to any pixel because it would be out of bound, in this case just skip\n int out_batch_offset = batch * O * H * W;\n for (int j = 0; j < TH; j++) {\n int hw_abs = block_col * BH + thread_col * TH + j;\n int h_abs = hw_abs / W;\n int w_abs = hw_abs % W;\n for (int i = 0; i < TO; i++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target_abs = h_abs + 1 - k_i;\n int w_target_abs = w_abs + 1 - k_j;\n if (h_target_abs >= 0 && h_target_abs < H && w_target_abs >= 0 && w_target_abs < W) {\n int out_target_abs = out_batch_offset + o_abs * H * W + h_target_abs * W + w_target_abs;\n int th_res_idx = k * TO * TH + i * TH + j;\n atomicAdd(out + out_target_abs, thread_results[th_res_idx]);\n }\n }\n int out_abs = out_batch_offset + o_abs * H * W + h_abs * W + w_abs;\n atomicAdd(out + out_abs, bias[o_abs]);\n }\n }\n}\n\n\n// almost same as 1, except we load A as transposed\ntemplate \n__global__ void conv2d_k3_forward_kernel2(\n const float* x, const float* weight, const float* bias,\n float* out,\n const int B, const int C, const int O, const int H, const int W\n) {\n // constants\n const int C9 = C * 9;\n const int weight_size = O * C * 9;\n const int x_size = B * C * H * W;\n // n_threads / block = (BO * BH) / (TO * TH) = 256\n // n_blocks = B * (O / BO) * (H * W / BH)\n // we will set gridDim to (BxHxW / BH, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // for compute step, to determine which of the size TOxTH block the thread is responsible for\n const int threads_per_col = ceil_div(BH, TH);\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // shared mem for loading weights and X\n const int w_s_sz = BO * BC * 9;\n const int w_s_width = BC * 9;\n const int x_s_sz = BC * BH;\n __shared__ float w_s[w_s_sz];\n __shared__ float x_s[x_s_sz];\n\n // registers to store local results\n float thread_results[TO * TH * 9] = {0.0};\n // registers for weights and x\n float reg_w[TO] = {0.0};\n float reg_x[TH] = {0.0};\n\n // calculate offset in x for bounds checking later\n const int x_block_offset = batch * C * H * W + block_col * BH;\n // calculate weight offset for bounds checking\n const int w_block_offset = block_row * BO * C9;\n\n // outer loop over BC size blocks of channel C\n for (int c_idx = 0; c_idx < C; c_idx += BC) {\n // x and weight offsets for this block\n int x_inner_offset = c_idx * H * W;\n int w_inner_offset = c_idx * 9;\n // load weight into smem\n int w_abs_idx;\n for (int i = threadIdx.x; i < w_s_sz; i += blockDim.x) {\n int w_i = i / w_s_width;\n int w_j = i % w_s_width;\n w_abs_idx = w_block_offset + w_inner_offset + w_i * C9 + w_j;\n if (w_abs_idx < weight_size) {\n w_s[w_j * BO + w_i] = weight[w_abs_idx];\n }\n }\n // load x into smem\n int x_abs_idx;\n for (int i = threadIdx.x; i < x_s_sz; i += blockDim.x) {\n int x_i = i / BH;\n int x_j = i % BH;\n x_abs_idx = x_block_offset + x_inner_offset + x_i * H * W + x_j;\n if (x_abs_idx < x_size) {\n x_s[x_i * BH + x_j] = x[x_abs_idx];\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int dot_idx = 0; dot_idx < BC; dot_idx++) {\n // load x_s into registers\n for (int j = 0; j < TH; j++) {\n reg_x[j] = x_s[dot_idx * BH + thread_col * TH + j];\n }\n for (int conv_idx = 0; conv_idx < 9; conv_idx++) {\n for (int i = 0; i < TO; i++) {\n reg_w[i] = w_s[(dot_idx * 9 + conv_idx) * BO + thread_row * TO + i];\n }\n\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n float val = reg_w[i] * reg_x[j];\n thread_results[conv_idx * TO * TH + i * TH + j] += val;\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n // need to be careful here: for each pixel at (h_abs, w_abs) and its product with kernel k (k < 9)\n // There are 2 options\n // 1. the product should be added to some pixel, then do an atomiAdd\n // 2. the product should not contribute to any pixel because it would be out of bound, in this case just skip\n int out_batch_offset = batch * O * H * W;\n for (int j = 0; j < TH; j++) {\n int hw_abs = block_col * BH + thread_col * TH + j;\n int h_abs = hw_abs / W;\n int w_abs = hw_abs % W;\n for (int i = 0; i < TO; i++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target_abs = h_abs + 1 - k_i;\n int w_target_abs = w_abs + 1 - k_j;\n if (h_target_abs >= 0 && h_target_abs < H && w_target_abs >= 0 && w_target_abs < W) {\n int out_target_abs = out_batch_offset + o_abs * H * W + h_target_abs * W + w_target_abs;\n int th_res_idx = k * TO * TH + i * TH + j;\n atomicAdd(out + out_target_abs, thread_results[th_res_idx]);\n }\n }\n int out_abs = out_batch_offset + o_abs * H * W + h_abs * W + w_abs;\n atomicAdd(out + out_abs, bias[o_abs]);\n }\n }\n}\n\n\n// almost same as 2, but we vectorize the loads and use float4\ntemplate \n__global__ void conv2d_k3_forward_kernel3(\n float* x, float* weight, float* bias,\n float* out,\n const int B, const int C, const int O, const int H, const int W\n) {\n // constants\n assert(TH == 4); // needed for reg_x loads later, so we can do one vectorized load for each thread\n const int C9 = C * 9;\n // n_threads / block = (BO * BH) / (TO * TH) = 256\n // n_blocks = B * (O / BO) * (H * W / BH)\n // we will set gridDim to (BxHxW / BH, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // for compute step, to determine which of the size TOxTH block the thread is responsible for\n const int threads_per_col = ceil_div(BH, TH);\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // shared mem for loading weights and X\n const int w_s_sz = BO * BC * 9;\n const int w_s_width = BC * 9;\n const int x_s_sz = BC * BH;\n __shared__ float w_s[w_s_sz];\n __shared__ float x_s[x_s_sz];\n\n // registers to store local results\n float thread_results[TO * TH * 9] = {0.0};\n // registers for weights and x\n float reg_w[TO] = {0.0};\n float reg_x[TH] = {0.0};\n\n // calculate offset in x for bounds checking later\n const int x_block_offset = batch * C * H * W + block_col * BH;\n // calculate weight offset for bounds checking\n const int w_block_offset = block_row * BO * C9;\n\n // outer loop over BC size blocks of channel C\n for (int c_idx = 0; c_idx < C; c_idx += BC) {\n // x and weight offsets for this block\n int x_inner_offset = c_idx * H * W;\n int w_inner_offset = c_idx * 9;\n // load weight into smem\n for (int i = threadIdx.x; i < w_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (w_s_width / 4);\n int inner_row = i / (w_s_width / 4);\n float4 tmp = \n reinterpret_cast(&weight[w_block_offset + w_inner_offset + inner_row * C9 + inner_col * 4])[0];\n w_s[(inner_col * 4 + 0) * BO + inner_row] = tmp.x;\n w_s[(inner_col * 4 + 1) * BO + inner_row] = tmp.y;\n w_s[(inner_col * 4 + 2) * BO + inner_row] = tmp.z;\n w_s[(inner_col * 4 + 3) * BO + inner_row] = tmp.w;\n }\n // load x into smem\n for (int i = threadIdx.x; i < x_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (BH / 4);\n int inner_row = i / (BH / 4);\n reinterpret_cast(&x_s[inner_row * BH + inner_col * 4])[0] = \n reinterpret_cast(&x[x_block_offset + x_inner_offset + inner_row * H * W + inner_col * 4])[0];\n }\n __syncthreads();\n\n // calculate per thread results\n for (int dot_idx = 0; dot_idx < BC; dot_idx++) {\n // load x_s into registers\n // NOTE: relying on the fact that TH = 4\n reinterpret_cast(reg_x)[0] = reinterpret_cast(&x_s[dot_idx * BH + thread_col * TH])[0];\n for (int conv_idx = 0; conv_idx < 9; conv_idx++) {\n for (int i = 0; i < TO; i++) {\n reg_w[i] = w_s[(dot_idx * 9 + conv_idx) * BO + thread_row * TO + i];\n }\n\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n float val = reg_w[i] * reg_x[j];\n thread_results[conv_idx * TO * TH + i * TH + j] += val;\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n // need to be careful here: for each pixel at (h_abs, w_abs) and its product with kernel k (k < 9)\n // There are 2 options\n // 1. the product should be added to some pixel, then do an atomiAdd\n // 2. the product should not contribute to any pixel because it would be out of bound, in this case just skip\n int out_batch_offset = batch * O * H * W;\n for (int j = 0; j < TH; j++) {\n int hw_abs = block_col * BH + thread_col * TH + j;\n int h_abs = hw_abs / W;\n int w_abs = hw_abs % W;\n for (int i = 0; i < TO; i++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target_abs = h_abs + 1 - k_i;\n int w_target_abs = w_abs + 1 - k_j;\n if (h_target_abs >= 0 && h_target_abs < H && w_target_abs >= 0 && w_target_abs < W) {\n int out_target_abs = out_batch_offset + o_abs * H * W + h_target_abs * W + w_target_abs;\n int th_res_idx = k * TO * TH + i * TH + j;\n atomicAdd(out + out_target_abs, thread_results[th_res_idx]);\n }\n }\n int out_abs = out_batch_offset + o_abs * H * W + h_abs * W + w_abs;\n atomicAdd(out + out_abs, bias[o_abs]);\n }\n }\n}\n\n// CPU version for debugging\nvoid conv_k3_forward_cpu(\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C, int OC, int H, int W\n) {\n for (int b = 0; b < B; b++) {\n for (int oc = 0; oc < OC; oc++) {\n for (int h = 0; h < H; h++) {\n for (int w = 0; w < W; w++) {\n for (int k = 0; k < 9; k++) {\n // check if the output will actually be used, if not skip\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target = h + 1 - k_i;\n int w_target = w + 1 - k_j;\n if (h_target < 0 || h_target >= H || w_target < 0 || w_target >= W) {\n continue;\n }\n float out_bohw = 0.0f;\n for (int c = 0; c < C; c++) {\n float new_val = weight[oc * C * 9 + c * 9 + k] * x[b * C * H * W + c * H * W + h * W + w];\n out_bohw += new_val;\n //if (b == 0 && h == 0 && w == 1 && k == 5 && oc == 0) {\n // printf(\"(0, 1) OC 0 kern 5: contribution from %d: %f, out val %f\\n\", c, new_val, out_bohw);\n //}\n }\n int out_idx = b * OC * H * W + oc * H * W + h_target * W + w_target;\n //if (h_target == 0 && w_target == 0 && oc == 0 && b == 0) {\n // printf(\"(0, 0) getting contribution from (%d, %d) at k %d: %f, prev_val %f\\n\", h, w, k, out_bohw, out[out_idx]);\n //}\n out[out_idx] += out_bohw;\n //if (h_target == 0 && w_target == 0 && oc == 0 && b == 0) {\n // printf(\"(0, 0) getting contribution from (%d, %d) at k %d: %f, new val %f\\n\", h, w, k, out_bohw, out[out_idx]);\n //}\n }\n \n //if (h == 0 && w == 0 && oc == 0 && b == 0) {\n // printf(\"(0, 0) adding bias %f\\n\", bias[oc]);\n //}\n out[b * OC * H * W + oc * H * W + h * W + w] += bias[oc];\n }\n }\n printf(\"(0, 0) final value: %f\\n\", out[0]);\n return;\n }\n }\n}\n\nvoid conv2d_k3_forward3(\n float* x, float* weight, float* bias,\n float* out,\n const int B, const int C_in, const int C_out, const int H, const int W\n) {\n // since the kernel will add to the output, we need to zero it first\n cudaCheck(cudaMemset(out, 0, B * C_out * H * W * sizeof(float)));\n\n if (H * W == 64) {\n const int BH = 64;\n const int TH = 4;\n assert(C_out % 64 == 0 && C_in % 8 == 0);\n const int BC = 16;\n const int BO = 64;\n const int TO = 4;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel3<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n } else if (C_in == 3) {\n const int BH = 64;\n assert(H * W % BH == 0);\n const int TH = 4;\n const int BC = 3;\n assert(C_out % 64 == 0);\n const int BO = 64;\n const int TO = 4;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel2<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n } else if (C_out == 3) {\n const int BH = 64 * 2;\n assert(H * W % BH == 0);\n const int TH = 4;\n assert(C_in % 8 == 0);\n const int BC = 8;\n const int BO = 3;\n const int TO = 3;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel2<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n } else {\n const int BH = 64 * 4;\n assert(H * W % BH == 0);\n const int TH = 4;\n assert(C_out % 64 == 0 && C_in % 8 == 0);\n const int BC = 16;\n const int BO = 8;\n const int TO = 2;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel3<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n }\n cudaCheck(cudaGetLastError());\n}\n\n// ----------------------------------------------------------------------------\n// backward pass\n\n__global__ void permute_dout_tile_x_kernel(\n const float* dout, const float* x,\n float* dout_perm, float* x_tiled,\n int B, int C_in, int C_out, int H, int W\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n if (idx < B * C_out * H * W) {\n // permute dout\n int b = idx / (C_out * img_size);\n int c = (idx / img_size) % C_out;\n int h = (idx / W) % H;\n int w = idx % W;\n dout_perm[c * B * img_size + b * img_size + h * W + w] = dout[idx];\n }\n\n if (idx < B * C_in * H * W) {\n // tile x_tiled, see permute_w_tile_x_kernel1 for explanation\n // except that we tile into a shape of (B, H, W, C_in, 9), the transpose of the tiling there\n int b = idx / (C_in * img_size);\n int c = (idx / img_size) % C_in;\n int h_start = (idx / W) % H;\n int w_start = idx % W;\n for (int i = 0; i < 3; i++) {\n for (int j = 0; j < 3; j++) {\n int h_end = h_start + 1 - i;\n int w_end = w_start + 1 - j;\n if (h_end >= 0 && h_end < H && w_end >= 0 && w_end < W) {\n x_tiled[b * img_size * C_in * 9 + h_end * W * C_in * 9 + w_end * C_in * 9 + c * 9 + i * 3 + j]\n = x[idx];\n }\n }\n }\n }\n}\n\n__global__ void reduce_dx_kernel(\n const float* dx_tiled,\n float* dx,\n int B, int C_in, int H, int W\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * C_in * H * W) {\n int img_size = H * W;\n int b = idx / (C_in * img_size);\n int c = (idx / img_size) % C_in;\n int h_end = (idx / W) % H;\n int w_end = idx % W;\n float sum = 0.0f;\n for (int i = 0; i < 3; i++) {\n for (int j = 0; j < 3; j++) {\n int h_start = h_end + 1 - i;\n int w_start = w_end + 1 - j;\n if (h_start >= 0 && h_start < H && w_start >= 0 && w_start < W) {\n sum += dx_tiled[c * 9 * B * img_size + (i * 3 + j) * B * img_size + b * img_size + h_start * W + w_start];\n }\n }\n }\n dx[idx] = sum;\n }\n}\n\n\n__global__ void reduce_dbias_kernel1(\n const float* dout_perm,\n float* dbias,\n int B, int C_out, int H, int W\n) {\n // dout_perm has shape (C_out, B, H, W)\n // each block will reduce over one channel\n dout_perm += blockIdx.x * B * H * W;\n int bhw = B * H * W;\n assert(blockDim.x <= bhw);\n extern __shared__ float sdata[];\n \n float sum = 0.0f;\n for (int i = threadIdx.x; i < bhw; i += blockDim.x) {\n sum += dout_perm[i];\n }\n sdata[threadIdx.x] = sum;\n\n __syncthreads();\n\n for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {\n if (threadIdx.x < stride) {\n sdata[threadIdx.x] += sdata[threadIdx.x + stride];\n }\n __syncthreads();\n }\n\n if (threadIdx.x == 0) {\n dbias[blockIdx.x] = sdata[0];\n }\n}\n\n// use cg to do warp reduce\n// doesn't really seem faster\n__global__ void reduce_dbias_kernel2(\n const float* dout_perm,\n float* dbias,\n int B, int C_out, int H, int W\n) {\n int bhw = B * H * W;\n assert(blockDim.x <= bhw);\n // move dout_perm to the right block\n dout_perm += blockIdx.x * bhw;\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n __shared__ float shared_sum[32];\n\n int num_warps = blockDim.x / 32;\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n\n float thread_sum = 0.0f;\n for (int i = threadIdx.x; i < bhw; i += blockDim.x) {\n thread_sum += dout_perm[i];\n }\n \n // warp reduce\n float warp_sum = cg::reduce(warp, thread_sum, cg::plus{});\n shared_sum[warp_id] = warp_sum;\n __syncthreads();\n\n // warp reduce\n warp_sum = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;\n float block_sum = cg::reduce(warp, warp_sum, cg::plus{});\n\n if (threadIdx.x == 0) {\n dbias[blockIdx.x] = block_sum;\n }\n}\n\nvoid conv2d_k3_backward1(\n cublasHandle_t cublas_handle,\n const float* dout, const float* x, const float* weight,\n float* dx, float* dweight, float* dbias,\n int B, int C_in, int C_out, int H, int W,\n int block_size,\n float* t1, float* t2, float* t3, float* t4, float* t5, float* t6\n) {\n // shapes:\n // dout (B, C_out, H, W)\n // x (B, C_in, H, W)\n // weight (C_out, C_in, 3, 3)\n // bias (C_out,)\n\n // we need to calculate dx, dweight, dbias separately\n // dbias is a single reduce\n // for dweight, we need to tile x into x_tile of shape (B * H * W, C_in * 9)\n // then do a matmul with dout_perm (shape (C_out, B * H * W))\n // for dx, we do dout_perm @ weight.T (store it in x_tile), and then reduce to get dx\n // If we don't want to do a reduce when computing dx, we need to make 9 copies of dout\n // which might be ok but increases memory usage. We can try it later.\n cudaEvent_t start, end1, end2, end3, end4, end5, end6;\n if (t1 != nullptr) {\n cudaCheck(cudaEventCreate(&start));\n cudaCheck(cudaEventCreate(&end1));\n cudaCheck(cudaEventCreate(&end2));\n cudaCheck(cudaEventCreate(&end3));\n cudaCheck(cudaEventCreate(&end4));\n cudaCheck(cudaEventCreate(&end5));\n cudaCheck(cudaEventCreate(&end6));\n cudaCheck(cudaEventRecord(start, nullptr));\n }\n\n // 1. transpose to dout_perm and tile x_tile\n float *dout_perm, *x_tiled;\n cudaCheck(cudaMalloc(&x_tiled, B * H * W * C_in * 9 * sizeof(float)));\n cudaCheck(cudaMemset(x_tiled, 0, B * H * W * C_in * 9 * sizeof(float)));\n cudaCheck(cudaMalloc(&dout_perm, B * C_out * H * W * sizeof(float)));\n if (t1 != nullptr) {\n cudaCheck(cudaEventRecord(end1, nullptr));\n }\n int total_th_1 = max_int(C_in, C_out) * B * H * W;\n int n_blk_1 = ceil_div(total_th_1, block_size);\n permute_dout_tile_x_kernel<<>>(dout, x, dout_perm, x_tiled, B, C_in, C_out, H, W);\n cudaCheck(cudaGetLastError());\n if (t2 != nullptr) {\n cudaCheck(cudaEventRecord(end2, nullptr));\n }\n\n // 2. get dweight by single matmul between x_tiled and dout_perm\n float alpha = 1.0f;\n float beta = 0.0f;\n int bhw = B * H * W;\n cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, C_in * 9, C_out, bhw, &alpha, x_tiled, C_in * 9, dout_perm, bhw, &beta, dweight, C_in * 9));\n\n if (t2 != nullptr) {\n cudaCheck(cudaEventRecord(end3, nullptr));\n }\n\n // 3. get dx_tile by matmul between dout_perm and weight.T, and save output in x_tiled\n // note: beta = 0 overwrites existing values in x_tiled\n cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, bhw, C_in * 9, C_out, &alpha, dout_perm, bhw, weight, C_in * 9, &beta, x_tiled, bhw));\n\n\n // 4. reduce x_tiled to get dx\n int total_th_2 = B * C_in * H * W;\n int n_blk_2 = ceil_div(total_th_2, block_size);\n reduce_dx_kernel<<>>(x_tiled, dx, B, C_in, H, W);\n\n if (t2 != nullptr) {\n cudaCheck(cudaEventRecord(end4, nullptr));\n }\n\n // 5. reduce dbias\n //reduce_dbias_kernel1<<>>(dout_perm, dbias, B, C_out, H, W);\n reduce_dbias_kernel2<<>>(dout_perm, dbias, B, C_out, H, W);\n\n if (t5 != nullptr) {\n cudaCheck(cudaEventRecord(end5, nullptr));\n }\n\n // free memory\n cudaCheck(cudaFree(dout_perm));\n cudaCheck(cudaFree(x_tiled));\n if (t5 != nullptr) {\n cudaCheck(cudaEventRecord(end6, nullptr));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end1));\n cudaCheck(cudaEventSynchronize(end2));\n cudaCheck(cudaEventSynchronize(end3));\n cudaCheck(cudaEventSynchronize(end4));\n cudaCheck(cudaEventSynchronize(end5));\n cudaCheck(cudaEventSynchronize(end6));\n float elapsed1, elapsed2, elapsed3, elapsed4, elapsed5, elapsed6;\n cudaCheck(cudaEventElapsedTime(&elapsed1, start, end1));\n cudaCheck(cudaEventElapsedTime(&elapsed2, end1, end2));\n cudaCheck(cudaEventElapsedTime(&elapsed3, end2, end3));\n cudaCheck(cudaEventElapsedTime(&elapsed4, end3, end4));\n cudaCheck(cudaEventElapsedTime(&elapsed5, end4, end5));\n cudaCheck(cudaEventElapsedTime(&elapsed6, end5, end6));\n *t1 += elapsed1;\n *t2 += elapsed2;\n *t3 += elapsed3;\n *t4 += elapsed4;\n *t5 += elapsed5;\n *t6 += elapsed6;\n }\n}\n\n// almost the same as conv2d_k3_forward_kernel, except that the compute and the writes\n// are different, because now we want to treat the C dimension as the inner product dimension\ntemplate \n__global__ void dx_backward_kernel1(\n float* dout, float* weight,\n float* dx,\n const int B, const int C, const int O, const int H, const int W\n) {\n const int C9 = C * 9;\n // n_threads per block = (BC * BH) / (TC * TH)\n // set gridDim to (BxHxW / BH, C / TC)\n // find the output block that this thread block is responsible for\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // index for the TCxTH block that the thread is responsible\n const int threads_per_col = BH / TH;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for weights and dout\n const int w_s_width = BC * 9;\n const int w_s_sz = BO * w_s_width;\n const int dout_s_sz = BO * BH;\n __shared__ float w_s[w_s_sz];\n __shared__ float dout_s[dout_s_sz];\n\n // registers for local results\n float thread_results[TC * TH * 9] = {0.0};\n // regs for weights and dout\n float reg_w[TC * 9] = {0.0};\n float reg_dout[TH] = {0.0};\n\n // calculate the block offsets\n const int dout_block_offset = batch * O * H * W + block_col * BH;\n const int w_block_offset = block_row * w_s_width;\n\n // outer loop over BO sizes of channel O\n for (int o_idx = 0; o_idx < O; o_idx += BO) {\n // offsets for this inner block\n int w_inner_offset = o_idx * C9;\n int dout_inner_offset = o_idx * H * W;\n\n // load weights to smem\n if (BC % 4 == 0) {\n // vectorize dram load\n for (int i = threadIdx.x; i < w_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (w_s_width / 4);\n int inner_row = i / (w_s_width / 4);\n reinterpret_cast(&w_s[inner_row * w_s_width + inner_col * 4])[0] = \n reinterpret_cast(&weight[w_block_offset + w_inner_offset + inner_row * C9 + inner_col * 4])[0];\n }\n } else {\n for (int i = threadIdx.x; i < w_s_sz; i += blockDim.x) {\n int inner_col = i % w_s_width;\n int inner_row = i / w_s_width;\n w_s[inner_row * w_s_width + inner_col] = weight[w_block_offset + w_inner_offset + inner_row * C9 + inner_col];\n }\n }\n // load dout\n for (int i = threadIdx.x; i < dout_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (BH / 4);\n int inner_row = i / (BH / 4);\n reinterpret_cast(&dout_s[inner_row * BH + inner_col * 4])[0] = \n reinterpret_cast(&dout[dout_block_offset + dout_inner_offset + inner_row * H * W + inner_col * 4])[0];\n }\n __syncthreads();\n \n\n // calculate per thread results\n for (int dot_idx = 0; dot_idx < BO; dot_idx++) {\n // load dout and w into registers\n for (int j = 0; j < (TH / 4); j++) {\n reinterpret_cast(®_dout[j * 4])[0] = \n reinterpret_cast(&dout_s[dot_idx * BH + thread_col * TH + j * 4])[0];\n }\n // non vectorized load\n //for (int j = 0; j < TH; j++) {\n // reg_dout[j] = dout_s[dot_idx * BH + thread_col * TH + j];\n //}\n if (TC % 4 == 0) {\n // vectorize smem load\n for (int i = 0; i < (TC * 9 / 4); i++) {\n reinterpret_cast(®_w[i * 4])[0] = \n reinterpret_cast(&w_s[dot_idx * w_s_width + thread_row * TC * 9 + i * 4])[0];\n }\n } else {\n for (int i = 0; i < TC * 9; i++) {\n reg_w[i] = w_s[dot_idx * w_s_width + thread_row * TC * 9 + i];\n }\n }\n\n for (int k = 0; k < 9; k++) {\n for (int i = 0; i < TC; i++) {\n for (int j = 0; j < TH; j++) {\n thread_results[k * TC * TH + i * TH + j] += \n reg_w[i * 9 + k] * reg_dout[j];\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n int out_batch_offset = batch * C * H * W;\n for (int j = 0; j < TH; j++) {\n int hw_abs = block_col * BH + thread_col * TH + j;\n int h_abs = hw_abs / W;\n int w_abs = hw_abs % W;\n for (int i = 0; i < TC; i++) {\n int c_abs = block_row * BC + thread_row * TC + i;\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target_abs = h_abs - 1 + k_i;\n int w_target_abs = w_abs - 1 + k_j;\n if (h_target_abs >= 0 && h_target_abs < H && w_target_abs >= 0 && w_target_abs < W) {\n int out_target_abs = out_batch_offset + c_abs * H * W + h_target_abs * W + w_target_abs;\n int th_res_idx = k * TC * TH + i * TH + j;\n atomicAdd(dx + out_target_abs, thread_results[th_res_idx]);\n }\n }\n }\n }\n}\n\n// I wonder how fast does a transposition kernel take for the weights\n__global__ void transpose_weight_kernel(\n float* weight, float* out,\n int C, int O\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < C * O * 9) {\n int o = idx / (C * 9);\n int c = (idx % C * 9) / 9;\n int k = idx % 9;\n out[c * O * 9 + o * 9 + k] = weight[idx];\n }\n}\n\nvoid dx_backward(\n float* dout, float* weight,\n float* dx,\n const int B, const int C_in, const int C_out, const int H, const int W\n) {\n if (H < 16) {\n const int BH = 64;\n const int TH = 4;\n const int BC = 16;\n const int TC = 4;\n const int BO = 16;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n } else if (C_in == 3) {\n const int BH = 64 * 4;\n const int TH = 4;\n const int BC = 3;\n const int TC = 3;\n const int BO = 16;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n } else if (C_out == 3) {\n const int BH = 64 * 4;\n const int TH = 4;\n const int BC = 16;\n const int TC = 4;\n const int BO = 3;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n } else {\n const int BH = 64 * 4;\n const int TH = 4;\n const int BC = 16;\n const int TC = 4;\n const int BO = 16;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n }\n}\n\nvoid dx_backward_cpu(\n float* dout, float* weight, float* dx,\n int B, int C, int O, int H, int W\n) {\n printf(\"dx[0] before zeroing: %f\\n\", dx[0]);\n for (int b = 0; b < B; b++) {\n for (int c = 0; c < C; c++) {\n for (int h = 0; h < H; h++) {\n for (int w = 0; w < W; w++) {\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target = h - 1 + k_i;\n int w_target = w - 1 + k_j;\n if (h_target < 0 || h_target >= H || w_target < 0 || w_target >= W) {\n continue;\n }\n float out_bchw = 0.0f;\n for (int o = 0; o < O; o++) {\n float new_val = weight[o * C * 9 + c * 9 + k] * dout[b * O * H * W + o * H * W + h * W + w];\n out_bchw += new_val;\n }\n int out_idx = b * C * H * W + c * H * W + h_target * W + w_target;\n dx[out_idx] += out_bchw;\n }\n }\n }\n if (b == 0 && c == 0) {\n printf(\"dx[0]: %f\\n\", dx[0]);\n return;\n }\n }\n }\n}\n\n// dweight_buf will have shape (O, C, 9, B)\n// and then we have to reduce over the last dimension to get the actual weight\ntemplate \n__global__ void dweight_backward_kernel1(\n float* dout, float* x,\n float* dweight_buf,\n const int B, const int C, const int O, const int H, const int W\n) {\n // ensure that we load entire image widths into smem\n // so we can do all 9 kernels at once\n assert(BW % W == 0);\n\n // n_threads per block = (BC * BO) / (TO * TC)\n // gridDim will be (B * C / BC, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = C / BC;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // index for the TO*TC thread block\n const int threads_per_col = BC / TC;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for dout and x\n __shared__ float dout_s[BO * BW];\n __shared__ float x_s[BC * BW];\n\n // registers for local mem and local results\n float thread_results[TO * TC * 9] = {0.0};\n float reg_dout[TO] = {0.0};\n float reg_x[TC] = {0.0};\n\n // block offsets\n const int dout_block_offset = batch * O * H * W + block_row * BO * H * W;\n const int x_block_offset = batch * C * H * W + block_col * BC * H * W;\n\n // number of image widths we will load to smem in one round\n const int NW = BW / W;\n\n // we will fuse the calculation for all 9 kernel filters\n // which means we will have to reload some of the image widths\n // we can load NW blocks into smem in one go, which means we will be stepping\n // over the H dimension in (NW - 1) steps\n // this is the outer loop over H\n for (int h = 0; h < H; h += (NW - 1)) {\n // load dout into smem\n int dout_inner_offset = h * W;\n int nh = (h < H - (NW - 1)) ? NW : NW - 1;\n int dout_load_sz = nh * W * BO;\n for (int i = threadIdx.x; i < dout_load_sz / 4; i++) {\n int inner_col = i % (nh * W / 4);\n int inner_row = i / (nh * W / 4);\n float4 tmp = \n reinterpret_cast(&dout[dout_block_offset + dout_inner_offset + inner_row * H * W + inner_col * 4])[0];\n dout_s[(inner_col * 4 + 0) * BO + inner_row] = tmp.x;\n dout_s[(inner_col * 4 + 1) * BO + inner_row] = tmp.y;\n dout_s[(inner_col * 4 + 2) * BO + inner_row] = tmp.z;\n dout_s[(inner_col * 4 + 3) * BO + inner_row] = tmp.w;\n }\n \n // load x to smem\n int x_inner_offset = h * W;\n int x_load_sz = nh * W * BC;\n for (int i = threadIdx.x; i < x_load_sz / 4; i++) {\n int inner_col = i % (nh * W / 4);\n int inner_row = i / (nh * W / 4);\n float4 tmp = \n reinterpret_cast(&x[x_block_offset + x_inner_offset + inner_row * H * W + inner_col * 4])[0];\n x_s[(inner_col * 4 + 0) * BC + inner_row] = tmp.x;\n x_s[(inner_col * 4 + 1) * BC + inner_row] = tmp.y;\n x_s[(inner_col * 4 + 2) * BC + inner_row] = tmp.z;\n x_s[(inner_col * 4 + 3) * BC + inner_row] = tmp.w;\n }\n __syncthreads();\n\n // if loading the last block, fill in the last block's buffer to be 0\n if (nh == NW - 1) {\n for (int i = threadIdx.x; i < W * BO; i++) {\n int inner_col = i % W;\n int inner_row = i / W;\n dout_s[(inner_col + nh * W) * BO + inner_row] = 0.0f;\n }\n for (int i = threadIdx.x; i < W * BC; i++) {\n int inner_col = i % W;\n int inner_row = i / W;\n x_s[(inner_col + nh * W) * BC + inner_row] = 0.0f;\n } \n }\n __syncthreads();\n\n // calculate per thread results\n for (int x_dot_idx = 0; x_dot_idx < BW; x_dot_idx++) {\n int x_h = x_dot_idx / W;\n int x_w = x_dot_idx % W;\n // load one pixel with TC of the C channels into reg_w\n for (int j = 0; j < TC; j++) {\n reg_x[j] = x_s[x_dot_idx * BC + thread_col * TC + j];\n }\n\n // for each conv filter, load a different row and calculate the outer product\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n // calculate the corresponding (h, w) that should be loaded from dout\n int dout_h = x_h + 1 - k_i;\n int dout_w = x_w + 1 - k_j;\n int dout_dot_idx = -1;\n if (dout_h >= 0 && dout_h < NW && dout_w >= 0 && dout_w < W) {\n // this logic is a bit tricky: for the last row in smem\n // because it will be loaded as part of the next block as well, and it will be the first row in the next block\n // for kernels 4, 5, 6 we need to only compute them once\n // this statement ensures that we don't compute the 4, 5, 6 kernels for the last row in smem\n // so it will be computed as part of the first row next time\n if (x_h == NW - 1 && k_i == 1) {\n dout_dot_idx = -1;\n } else {\n dout_dot_idx = dout_h * W + dout_w;\n }\n }\n // if in bounds of the dout section in smem, load the corresponding pixel from smem\n // and add the products to the thread results\n // otherwise the contribution is 0\n if (dout_dot_idx != -1) {\n for (int i = 0; i < TO; i++) {\n reg_dout[i] = dout_s[dout_dot_idx * BO + thread_row * TO + i];\n }\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n thread_results[k * TO * TC + i * TC + j] += reg_dout[i] * reg_x[j];\n }\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n for (int k = 0; k < 9; k++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n int c_abs = block_col * BC + thread_col * TC + j;\n dweight_buf[o_abs * C * 9 * B + c_abs * 9 * B + k * B + batch] = thread_results[k * TO * TC + i * TC + j];\n }\n }\n }\n}\n\n\n// C * O blocks\n// 288 (9 * 32) threads per block, use one warp to reduce each batch dim\n__global__ void dweight_reduce_kernel(\n float* dweight_buf, float* dweight,\n const int B, const int C, const int O\n) {\n assert(gridDim.x == O * C);\n int o = blockIdx.x / C;\n int c = blockIdx.x % C;\n // move input pointer\n dweight_buf += o * C * 9 * B + c * 9 * B;\n assert(blockDim.x == 32 * 9);\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n\n float thread_sum = 0.0f;\n for (int i = lane_id; i < B; i += 32) {\n thread_sum += dweight_buf[warp_id * B + i];\n }\n\n // warp reduce\n float warp_sum = cg::reduce(warp, thread_sum, cg::plus{});\n\n if (lane_id == 0) {\n dweight[o * C * 9 + c * 9 + warp_id] = warp_sum;\n }\n}\n\nvoid dweight_backward1(\n float* dout, float* x,\n float* dweight_buf, float* dweight,\n const int B, const int C, const int O, const int H, const int W\n) {\n const int BC = 32;\n const int BO = 32;\n const int BW = 2 * 64;\n const int TO = 4;\n const int TC = 4;\n dim3 gridDim(B * C / BC, O / BO);\n dim3 blockDim(BC * BO / (TO * TC));\n dweight_backward_kernel1<<>>(dout, x, dweight_buf, B, C, O, H, W);\n cudaCheck(cudaGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B, C, O);\n}\n\n\n// instead of having to load multiples of entire widths of the image into smem\n// we now only load fractions of the width each time\n// so the outer loop needs to loop over the height multiples (NH - 1) as well as the width multiples (WS - 1)\n// results: This did seem to help with memory access: compute intensity is higher\n// right now this one is much slower than kernel 1\n// I think this is because the smem loads are now not aligned, which causes a lot of uncoalesced loads\ntemplate \n__global__ void dweight_backward_kernel2(\n float* dout, float* x,\n float* dweight_buf,\n const int B, const int C, const int O, const int H, const int W\n) {\n // number of heights we need to load into smem in one go \n assert(BW % WSZ == 0);\n const int NH = BW / WSZ;\n // n_threads per block = (BC * BO) / (TO * TC)\n // gridDim will be (B * C / BC, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = C / BC;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // index for the TO*TC thread block\n const int threads_per_col = BC / TC;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for dout and x\n __shared__ float dout_s[BO * BW];\n __shared__ float x_s[BC * BW];\n\n // registers for local mem and local results\n float thread_results[TO * TC * 9] = {0.0};\n float reg_dout[TO] = {0.0};\n float reg_x[TC] = {0.0};\n\n // block offsets\n const int dout_block_offset = batch * O * H * W + block_row * BO * H * W;\n const int x_block_offset = batch * C * H * W + block_col * BC * H * W;\n\n //int debug_idx = 0;\n\n // we will fuse the calculation for all 9 kernel filters\n // which means we will have to reload some of the image widths\n // we can load NW blocks into smem in one go, which means we will be stepping\n // over the H dimension in (NW - 1) steps\n // this is the outer loop over H\n for (int h = 0; h < H; h += (NH - 1)) {\n for (int w_offset = 0; w_offset < W; w_offset += (WSZ - 1)) {\n // load dout into smem\n int h_offset = h * W;\n int nh = (h < H - (NH - 1)) ? NH : NH - 1;\n int wsz = (w_offset < W - (WSZ - 1)) ? WSZ : WSZ - 1;\n int dout_load_sz = nh * wsz * BO;\n for (int i = threadIdx.x; i < dout_load_sz; i++) {\n int inner_w = i % wsz;\n int inner_h = (i / wsz) % nh;\n int inner_o = i / (wsz * nh);\n dout_s[(inner_h * WSZ + inner_w) * BO + inner_o] = \n dout[dout_block_offset + h_offset + inner_h * W + w_offset + inner_w + inner_o * H * W];\n }\n \n // load x to smem\n int x_load_sz = nh * wsz * BC;\n for (int i = threadIdx.x; i < x_load_sz; i++) {\n int inner_w = i % wsz;\n int inner_h = (i / wsz) % nh;\n int inner_c = i / (wsz * nh);\n x_s[(inner_h * WSZ + inner_w) * BC + inner_c] = \n x[x_block_offset + h_offset + inner_h * W + w_offset + inner_w + inner_c * H * W];\n }\n\n __syncthreads();\n\n // if loading the last block, fill in the last block's buffer to be 0\n if (nh == NH - 1) {\n for (int i = threadIdx.x; i < WSZ * BO; i++) {\n int inner_w = i % WSZ;\n int inner_o = i / WSZ;\n dout_s[(inner_w + nh * WSZ) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < WSZ * BC; i++) {\n int inner_w = i % WSZ;\n int inner_c = i / WSZ;\n\n x_s[(inner_w + nh * WSZ) * BC + inner_c] = 0.0f;\n } \n }\n // if loading the last width section, fill in the final value in the buffer to 0\n if (wsz == WSZ - 1) {\n for (int i = threadIdx.x; i < nh * BO; i++) {\n int inner_o = i % BO;\n int inner_h = i / BO;\n dout_s[(inner_h * WSZ + wsz) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < nh * BC; i++) {\n int inner_c = i % BC;\n int inner_h = i / BC;\n x_s[(inner_h * WSZ + wsz) * BC + inner_c] = 0.0f;\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int x_dot_idx = 0; x_dot_idx < WSZ * nh; x_dot_idx++) {\n int x_h = x_dot_idx / WSZ;\n int x_w = x_dot_idx % WSZ;\n // load TC of the C channels of one image pixel into reg_w\n for (int j = 0; j < TC; j++) {\n reg_x[j] = x_s[x_dot_idx * BC + thread_col * TC + j];\n }\n\n // for each conv filter, load a different row and calculate the outer product\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n // calculate the corresponding (h, w) that should be loaded from dout\n int dout_h = x_h + 1 - k_i;\n int dout_w = x_w + 1 - k_j;\n int dout_h_abs = h + dout_h;\n int dout_w_abs = w_offset + dout_w;\n int dout_dot_idx = -1;\n if (dout_h >= 0 && dout_h < NH && dout_w >= 0 && dout_w < WSZ && dout_w_abs < W && dout_h_abs < H) {\n if (x_h == NH - 1 && k_i == 1) {\n // this logic is a bit tricky: for the last row in smem\n // because it will be loaded as part of the next block as well, and it will be the first row in the next block\n // for kernels 4, 5, 6 we need to only compute them once\n // this statement ensures that we don't compute the 4, 5, 6 kernels for the last row in smem\n // so it will be computed as part of the first row next time\n dout_dot_idx = -1;\n } else if (x_w == WSZ - 1 && k_j == 1) {\n dout_dot_idx = -1;\n } else {\n dout_dot_idx = dout_h * WSZ + dout_w;\n }\n }\n\n // if in bounds of the dout section in smem, load the corresponding pixel from smem\n // and add the products to the thread results\n // otherwise the contribution is 0\n if (dout_dot_idx != -1) {\n for (int i = 0; i < TO; i++) {\n assert(dout_dot_idx * BO + thread_row * TO + i < BO * BW);\n reg_dout[i] = dout_s[dout_dot_idx * BO + thread_row * TO + i];\n }\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n float new_val = reg_dout[i] * reg_x[j];\n thread_results[k * TO * TC + i * TC + j] += new_val;\n }\n }\n }\n }\n }\n __syncthreads();\n }\n }\n\n // write out results\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n for (int k = 0; k < 9; k++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n int c_abs = block_col * BC + thread_col * TC + j;\n dweight_buf[o_abs * C * 9 * B + c_abs * 9 * B + k * B + batch] = thread_results[k * TO * TC + i * TC + j];\n }\n }\n }\n}\n\nvoid dweight_backward2(\n float* dout, float* x,\n float* dweight_buf, float* dweight,\n const int B, const int C, const int O, const int H, const int W\n) {\n const int BC = 32;\n const int BO = 32;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 4;\n dim3 gridDim(B * C / BC, O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_backward_kernel2<<>>(dout, x, dweight_buf, B, C, O, H, W);\n cudaCheck(cudaDeviceSynchronize());\n cudaCheck(cudaGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B, C, O);\n}\n\n\n// kernel 2 does something silly: because are now loading WSZ sections of widths,\n// and WSZ is definitely not a nice power of 2, we need to be careful\n// which thread we are using when loading global memory.\n// This kernel only changes the way loading is done compared to kernel 2\n// effectively we do some warp tiling when loading from global mem\n// kernel 3 is now quite fast (2.6ms)\n// profiling tells us that we are now using very few blocks, so seems like a good idea\n// to spread out the different width sections on different blocks and do a bigger reduce afterwards\ntemplate \n__global__ void dweight_backward_kernel3(\n float* dout, float* x,\n float* dweight_buf,\n const int B, const int C, const int O, const int H, const int W\n) {\n // number of heights we need to load into smem in one go \n const int NH = BW / WSZ;\n // n_threads per block = (BC * BO) / (TO * TC)\n // gridDim will be (B * C / BC, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = C / BC;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // warp indices for loading from global mem\n const int n_warps = blockDim.x / 32;\n const int warp_id = threadIdx.x / 32;\n const int lane_id = threadIdx.x % 32;\n // make sure that the width sections can be loaded by a single warp\n assert(WSZ <= 32);\n\n // index for the TO*TC thread block\n const int threads_per_col = BC / TC;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for dout and x\n __shared__ float dout_s[BO * BW];\n __shared__ float x_s[BC * BW];\n\n // registers for local mem and local results\n float thread_results[TO * TC * 9] = {0.0};\n float reg_dout[TO] = {0.0};\n float reg_x[TC] = {0.0};\n\n // block offsets\n const int dout_block_offset = batch * O * H * W + block_row * BO * H * W;\n const int x_block_offset = batch * C * H * W + block_col * BC * H * W;\n\n //int debug_idx = 0;\n\n // we will fuse the calculation for all 9 kernel filters\n // which means we will have to reload some of the image widths\n // we can load NW blocks into smem in one go, which means we will be stepping\n // over the H dimension in (NW - 1) steps\n // this is the outer loop over H\n for (int h = 0; h < H; h += (NH - 1)) {\n for (int w_offset = 0; w_offset < W; w_offset += (WSZ - 1)) {\n // load dout into smem\n int h_offset = h * W;\n int nh = (h < H - (NH - 1)) ? NH : NH - 1;\n int wsz = (w_offset < W - (WSZ - 1)) ? WSZ : WSZ - 1;\n int o_per_warp = ceil_div(BO, n_warps);\n if (lane_id < wsz) {\n for (int warp_o = 0; warp_o < o_per_warp; warp_o++) {\n int inner_o = warp_id * o_per_warp + warp_o;\n if (inner_o < BO) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n dout_s[(inner_h * WSZ + inner_w) * BO + inner_o] = \n dout[dout_block_offset + inner_o * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n int c_per_warp = ceil_div(BC, n_warps);\n if (lane_id < wsz) {\n for (int warp_c = 0; warp_c < c_per_warp; warp_c++) {\n int inner_c = warp_id * c_per_warp + warp_c;\n if (inner_c < BC) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n x_s[(inner_h * WSZ + inner_w) * BC + inner_c] = \n x[x_block_offset + inner_c * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n __syncthreads();\n \n\n // if loading the last block, fill in the last block's buffer to be 0\n if (nh == NH - 1) {\n for (int i = threadIdx.x; i < WSZ * BO; i++) {\n int inner_w = i % WSZ;\n int inner_o = i / WSZ;\n dout_s[(inner_w + nh * WSZ) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < WSZ * BC; i++) {\n int inner_w = i % WSZ;\n int inner_c = i / WSZ;\n\n x_s[(inner_w + nh * WSZ) * BC + inner_c] = 0.0f;\n } \n }\n // if loading the last width section, fill in the final value in the buffer to 0\n if (wsz == WSZ - 1) {\n for (int i = threadIdx.x; i < nh * BO; i++) {\n int inner_o = i % BO;\n int inner_h = i / BO;\n dout_s[(inner_h * WSZ + wsz) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < nh * BC; i++) {\n int inner_c = i % BC;\n int inner_h = i / BC;\n x_s[(inner_h * WSZ + wsz) * BC + inner_c] = 0.0f;\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int x_dot_idx = 0; x_dot_idx < WSZ * nh; x_dot_idx++) {\n int x_h = x_dot_idx / WSZ;\n int x_w = x_dot_idx % WSZ;\n // load TC of the C channels of one image pixel into reg_w\n for (int j = 0; j < TC; j++) {\n reg_x[j] = x_s[x_dot_idx * BC + thread_col * TC + j];\n }\n\n // for each conv filter, load a different row and calculate the outer product\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n // calculate the corresponding (h, w) that should be loaded from dout\n int dout_h = x_h + 1 - k_i;\n int dout_w = x_w + 1 - k_j;\n int dout_h_abs = h + dout_h;\n int dout_w_abs = w_offset + dout_w;\n int dout_dot_idx = -1;\n if (dout_h >= 0 && dout_h < NH && dout_w >= 0 && dout_w < WSZ && dout_w_abs < W && dout_h_abs < H) {\n if (x_h == NH - 1 && k_i == 1) {\n // this logic is a bit tricky: for the last row in smem\n // because it will be loaded as part of the next block as well, and it will be the first row in the next block\n // for kernels 4, 5, 6 we need to only compute them once\n // this statement ensures that we don't compute the 4, 5, 6 kernels for the last row in smem\n // so it will be computed as part of the first row next time\n dout_dot_idx = -1;\n } else if (x_w == WSZ - 1 && k_j == 1) {\n dout_dot_idx = -1;\n } else {\n dout_dot_idx = dout_h * WSZ + dout_w;\n }\n }\n\n // if in bounds of the dout section in smem, load the corresponding pixel from smem\n // and add the products to the thread results\n // otherwise the contribution is 0\n if (dout_dot_idx != -1) {\n for (int i = 0; i < TO; i++) {\n reg_dout[i] = dout_s[dout_dot_idx * BO + thread_row * TO + i];\n }\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n float new_val = reg_dout[i] * reg_x[j];\n thread_results[k * TO * TC + i * TC + j] += new_val;\n }\n }\n }\n }\n }\n __syncthreads();\n }\n }\n\n // write out results\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n for (int k = 0; k < 9; k++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n int c_abs = block_col * BC + thread_col * TC + j;\n dweight_buf[o_abs * C * 9 * B + c_abs * 9 * B + k * B + batch] = thread_results[k * TO * TC + i * TC + j];\n }\n }\n }\n}\n\nvoid dweight_backward3(\n float* dout, float* x,\n float* dweight_buf, float* dweight,\n const int B, const int C, const int O, const int H, const int W\n) {\n const int BC = 64;\n const int BO = 64;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * C / BC, O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_backward_kernel3<<>>(dout, x, dweight_buf, B, C, O, H, W);\n cudaCheck(cudaDeviceSynchronize());\n cudaCheck(cudaGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B, C, O);\n}\n\n\n// trying to improve on kernel 3 by splitting the widths of each image\n// to be calculated on different blocks\n// Output will now be of shape (O, C, 9, B, W / (WSZ - 1))\n// and we will need to reduce over the last two dimensions\ntemplate \n__global__ void dweight_backward_kernel4(\n float* dout, float* x,\n float* dweight_buf,\n const int B, const int C, const int O, const int H, const int W\n) {\n // number of heights we need to load into smem in one go \n const int NH = BW / WSZ;\n // n_threads per block = (BC * BO) / (TO * TC)\n // gridDim will be (B * (W / (WSZ - 1)) * (C / BC), O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_chan = C / BC;\n const int blocks_per_width = W / (WSZ - 1);\n const int block_w = blockIdx.x % blocks_per_width;\n const int block_col = (blockIdx.x / blocks_per_width) % blocks_per_chan;\n const int batch = blockIdx.x / (blocks_per_chan * blocks_per_width);\n const int w_offset = block_w * (WSZ - 1);\n\n // warp indices for loading from global mem\n const int n_warps = blockDim.x / 32;\n const int warp_id = threadIdx.x / 32;\n const int lane_id = threadIdx.x % 32;\n // make sure that the width sections can be loaded by a single warp\n assert(WSZ <= 32);\n\n // index for the TO*TC thread block\n const int threads_per_col = BC / TC;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for dout and x\n __shared__ float dout_s[BO * BW];\n __shared__ float x_s[BC * BW];\n\n // registers for local mem and local results\n float thread_results[TO * TC * 9] = {0.0};\n float reg_dout[TO] = {0.0};\n float reg_x[TC] = {0.0};\n\n // block offsets\n const int dout_block_offset = batch * O * H * W + block_row * BO * H * W;\n const int x_block_offset = batch * C * H * W + block_col * BC * H * W;\n\n //int debug_idx = 0;\n\n // we will fuse the calculation for all 9 kernel filters\n // which means we will have to reload some of the image widths\n // we can load NW blocks into smem in one go, which means we will be stepping\n // over the H dimension in (NW - 1) steps\n // this is the outer loop over H\n for (int h = 0; h < H; h += (NH - 1)) {\n // load dout into smem\n int h_offset = h * W;\n int nh = (h < H - (NH - 1)) ? NH : NH - 1;\n int wsz = (w_offset < W - (WSZ - 1)) ? WSZ : WSZ - 1;\n int o_per_warp = ceil_div(BO, n_warps);\n if (lane_id < wsz) {\n for (int warp_o = 0; warp_o < o_per_warp; warp_o++) {\n int inner_o = warp_id * o_per_warp + warp_o;\n if (inner_o < BO) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n dout_s[(inner_h * WSZ + inner_w) * BO + inner_o] = \n dout[dout_block_offset + inner_o * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n int c_per_warp = ceil_div(BC, n_warps);\n if (lane_id < wsz) {\n for (int warp_c = 0; warp_c < c_per_warp; warp_c++) {\n int inner_c = warp_id * c_per_warp + warp_c;\n if (inner_c < BC) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n x_s[(inner_h * WSZ + inner_w) * BC + inner_c] = \n x[x_block_offset + inner_c * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n __syncthreads();\n \n\n // if loading the last block, fill in the last block's buffer to be 0\n if (nh == NH - 1) {\n for (int i = threadIdx.x; i < WSZ * BO; i++) {\n int inner_w = i % WSZ;\n int inner_o = i / WSZ;\n dout_s[(inner_w + nh * WSZ) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < WSZ * BC; i++) {\n int inner_w = i % WSZ;\n int inner_c = i / WSZ;\n\n x_s[(inner_w + nh * WSZ) * BC + inner_c] = 0.0f;\n } \n }\n // if loading the last width section, fill in the final value in the buffer to 0\n if (wsz == WSZ - 1) {\n for (int i = threadIdx.x; i < nh * BO; i++) {\n int inner_o = i % BO;\n int inner_h = i / BO;\n dout_s[(inner_h * WSZ + wsz) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < nh * BC; i++) {\n int inner_c = i % BC;\n int inner_h = i / BC;\n x_s[(inner_h * WSZ + wsz) * BC + inner_c] = 0.0f;\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int x_dot_idx = 0; x_dot_idx < WSZ * nh; x_dot_idx++) {\n int x_h = x_dot_idx / WSZ;\n int x_w = x_dot_idx % WSZ;\n // load TC of the C channels of one image pixel into reg_w\n for (int j = 0; j < TC; j++) {\n reg_x[j] = x_s[x_dot_idx * BC + thread_col * TC + j];\n }\n\n // for each conv filter, load a different row and calculate the outer product\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n // calculate the corresponding (h, w) that should be loaded from dout\n int dout_h = x_h + 1 - k_i;\n int dout_w = x_w + 1 - k_j;\n int dout_h_abs = h + dout_h;\n int dout_w_abs = w_offset + dout_w;\n int dout_dot_idx = -1;\n if (dout_h >= 0 && dout_h < NH && dout_w >= 0 && dout_w < WSZ && dout_w_abs < W && dout_h_abs < H) {\n if (x_h == NH - 1 && k_i == 1) {\n // this logic is a bit tricky: for the last row in smem\n // because it will be loaded as part of the next block as well, and it will be the first row in the next block\n // for kernels 4, 5, 6 we need to only compute them once\n // this statement ensures that we don't compute the 4, 5, 6 kernels for the last row in smem\n // so it will be computed as part of the first row next time\n dout_dot_idx = -1;\n } else if (x_w == WSZ - 1 && k_j == 1) {\n dout_dot_idx = -1;\n } else {\n dout_dot_idx = dout_h * WSZ + dout_w;\n }\n }\n\n // if in bounds of the dout section in smem, load the corresponding pixel from smem\n // and add the products to the thread results\n // otherwise the contribution is 0\n if (dout_dot_idx != -1) {\n for (int i = 0; i < TO; i++) {\n reg_dout[i] = dout_s[dout_dot_idx * BO + thread_row * TO + i];\n }\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n float new_val = reg_dout[i] * reg_x[j];\n thread_results[k * TO * TC + i * TC + j] += new_val;\n }\n }\n }\n }\n __syncthreads();\n }\n }\n\n // write out results\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n for (int k = 0; k < 9; k++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n int c_abs = block_col * BC + thread_col * TC + j;\n dweight_buf[o_abs * C * 9 * B * blocks_per_width + c_abs * 9 * B * blocks_per_width + k * B * blocks_per_width + batch * blocks_per_width + block_w] = thread_results[k * TO * TC + i * TC + j];\n }\n }\n }\n}\n\nvoid dweight_backward4(\n float* dout, float* x,\n float* dweight_buf, float* dweight,\n const int B, const int C, const int O, const int H, const int W\n) {\n const int BC = 64;\n const int BO = 64;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_backward_kernel4<<>>(dout, x, dweight_buf, B, C, O, H, W);\n cudaCheck(cudaDeviceSynchronize());\n cudaCheck(cudaGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n}\n\n// same as kernel 4, but vectorize loads to speed up smem load\n// unfortunately as is this doesn't give any speedup\n// probably because we are now using too few threads to load\n// I could try to shuffle and make each warp load more heights as well, but that doesn't seem first order right now\ntemplate \n__global__ void dweight_backward_kernel5(\n float* dout, float* x,\n float* dweight_buf,\n const int B, const int C, const int O, const int H, const int W\n) {\n // number of heights we need to load into smem in one go \n const int NH = BW / WSZ;\n // n_threads per block = (BC * BO) / (TO * TC)\n // gridDim will be (B * (W / (WSZ - 1)) * (C / BC), O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_chan = C / BC;\n const int blocks_per_width = W / (WSZ - 1);\n const int block_w = blockIdx.x % blocks_per_width;\n const int block_col = (blockIdx.x / blocks_per_width) % blocks_per_chan;\n const int batch = blockIdx.x / (blocks_per_chan * blocks_per_width);\n const int w_offset = block_w * (WSZ - 1);\n\n // warp indices for loading from global mem\n const int n_warps = blockDim.x / 32;\n const int warp_id = threadIdx.x / 32;\n const int lane_id = threadIdx.x % 32;\n // make sure that the width sections can be loaded by a single warp\n assert(WSZ <= 32);\n\n // index for the TO*TC thread block\n const int threads_per_col = BC / TC;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for dout and x\n __shared__ float dout_s[BO * BW];\n __shared__ float x_s[BC * BW];\n\n // registers for local mem and local results\n float thread_results[TO * TC * 9] = {0.0};\n float reg_dout[TO] = {0.0};\n float reg_x[TC] = {0.0};\n\n // block offsets\n const int dout_block_offset = batch * O * H * W + block_row * BO * H * W;\n const int x_block_offset = batch * C * H * W + block_col * BC * H * W;\n\n //int debug_idx = 0;\n\n // we will fuse the calculation for all 9 kernel filters\n // which means we will have to reload some of the image widths\n // we can load NW blocks into smem in one go, which means we will be stepping\n // over the H dimension in (NW - 1) steps\n // this is the outer loop over H\n for (int h = 0; h < H; h += (NH - 1)) {\n // load dout into smem\n int h_offset = h * W;\n int nh = (h < H - (NH - 1)) ? NH : NH - 1;\n int wsz = (w_offset < W - (WSZ - 1)) ? WSZ : WSZ - 1;\n int o_per_warp = ceil_div(BO, n_warps);\n if (lane_id < ceil_div(wsz, 4)) {\n for (int warp_o = 0; warp_o < o_per_warp; warp_o++) {\n int inner_o = warp_id * o_per_warp + warp_o;\n if (inner_o < BO) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n float4 tmp = \n reinterpret_cast(&dout[dout_block_offset + inner_o * H * W + h_offset + inner_h * W + w_offset + lane_id * 4])[0];\n dout_s[(inner_h * WSZ + lane_id * 4 + 0) * BO + inner_o] = tmp.x;\n if (lane_id < wsz / 4) {\n dout_s[(inner_h * WSZ + lane_id * 4 + 1) * BO + inner_o] = tmp.y;\n dout_s[(inner_h * WSZ + lane_id * 4 + 2) * BO + inner_o] = tmp.z;\n dout_s[(inner_h * WSZ + lane_id * 4 + 3) * BO + inner_o] = tmp.w;\n }\n }\n }\n }\n }\n\n int c_per_warp = ceil_div(BC, n_warps);\n if (lane_id < ceil_div(wsz, 4)) {\n for (int warp_c = 0; warp_c < c_per_warp; warp_c++) {\n int inner_c = warp_id * c_per_warp + warp_c;\n if (inner_c < BC) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n float4 tmp = \n reinterpret_cast(&x[x_block_offset + inner_c * H * W + h_offset + inner_h * W + w_offset + lane_id * 4])[0];\n x_s[(inner_h * WSZ + lane_id * 4 + 0) * BC + inner_c] = tmp.x;\n if (lane_id < wsz / 4) {\n x_s[(inner_h * WSZ + lane_id * 4 + 1) * BC + inner_c] = tmp.y;\n x_s[(inner_h * WSZ + lane_id * 4 + 2) * BC + inner_c] = tmp.z;\n x_s[(inner_h * WSZ + lane_id * 4 + 3) * BC + inner_c] = tmp.w;\n }\n }\n }\n }\n }\n\n __syncthreads();\n \n\n // if loading the last block, fill in the last block's buffer to be 0\n if (nh == NH - 1) {\n for (int i = threadIdx.x; i < WSZ * BO; i++) {\n int inner_w = i % WSZ;\n int inner_o = i / WSZ;\n dout_s[(inner_w + nh * WSZ) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < WSZ * BC; i++) {\n int inner_w = i % WSZ;\n int inner_c = i / WSZ;\n\n x_s[(inner_w + nh * WSZ) * BC + inner_c] = 0.0f;\n } \n }\n // if loading the last width section, fill in the final value in the buffer to 0\n if (wsz == WSZ - 1) {\n for (int i = threadIdx.x; i < nh * BO; i++) {\n int inner_o = i % BO;\n int inner_h = i / BO;\n dout_s[(inner_h * WSZ + wsz) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < nh * BC; i++) {\n int inner_c = i % BC;\n int inner_h = i / BC;\n x_s[(inner_h * WSZ + wsz) * BC + inner_c] = 0.0f;\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int x_dot_idx = 0; x_dot_idx < WSZ * nh; x_dot_idx++) {\n int x_h = x_dot_idx / WSZ;\n int x_w = x_dot_idx % WSZ;\n // load TC of the C channels of one image pixel into reg_w\n for (int j = 0; j < TC; j++) {\n reg_x[j] = x_s[x_dot_idx * BC + thread_col * TC + j];\n }\n\n // for each conv filter, load a different row and calculate the outer product\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n // calculate the corresponding (h, w) that should be loaded from dout\n int dout_h = x_h + 1 - k_i;\n int dout_w = x_w + 1 - k_j;\n int dout_h_abs = h + dout_h;\n int dout_w_abs = w_offset + dout_w;\n int dout_dot_idx = -1;\n if (dout_h >= 0 && dout_h < NH && dout_w >= 0 && dout_w < WSZ && dout_w_abs < W && dout_h_abs < H) {\n if (x_h == NH - 1 && k_i == 1) {\n // this logic is a bit tricky: for the last row in smem\n // because it will be loaded as part of the next block as well, and it will be the first row in the next block\n // for kernels 4, 5, 6 we need to only compute them once\n // this statement ensures that we don't compute the 4, 5, 6 kernels for the last row in smem\n // so it will be computed as part of the first row next time\n dout_dot_idx = -1;\n } else if (x_w == WSZ - 1 && k_j == 1) {\n dout_dot_idx = -1;\n } else {\n dout_dot_idx = dout_h * WSZ + dout_w;\n }\n }\n\n // if in bounds of the dout section in smem, load the corresponding pixel from smem\n // and add the products to the thread results\n // otherwise the contribution is 0\n if (dout_dot_idx != -1) {\n for (int i = 0; i < TO; i++) {\n reg_dout[i] = dout_s[dout_dot_idx * BO + thread_row * TO + i];\n }\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n float new_val = reg_dout[i] * reg_x[j];\n thread_results[k * TO * TC + i * TC + j] += new_val;\n }\n }\n }\n }\n __syncthreads();\n }\n }\n\n // write out results\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n for (int k = 0; k < 9; k++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n int c_abs = block_col * BC + thread_col * TC + j;\n dweight_buf[o_abs * C * 9 * B * blocks_per_width + c_abs * 9 * B * blocks_per_width + k * B * blocks_per_width + batch * blocks_per_width + block_w] = thread_results[k * TO * TC + i * TC + j];\n }\n }\n }\n}\n\nvoid dweight_backward5(\n float* dout, float* x,\n float* dweight_buf, float* dweight,\n const int B, const int C, const int O, const int H, const int W\n) {\n const int BC = 64;\n const int BO = 64;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_backward_kernel5<<>>(dout, x, dweight_buf, B, C, O, H, W);\n cudaCheck(cudaDeviceSynchronize());\n cudaCheck(cudaGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n}\n\n// same as dweight backward 4, except we also do some of the reduce for dbias\n// so it doesn't have to reduce an entire BxHxW length for each O channel\ntemplate \n__global__ void dweight_dbias_backward_kernel1(\n float* dout, float* x,\n float* dweight_buf, float* dbias_buf,\n const int B, const int C, const int O, const int H, const int W\n) {\n // number of heights we need to load into smem in one go \n const int NH = BW / WSZ;\n // n_threads per block = (BC * BO) / (TO * TC)\n // gridDim will be (B * (W / (WSZ - 1)) * (C / BC), O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_chan = C / BC;\n const int blocks_per_width = W / (WSZ - 1);\n const int block_w = blockIdx.x % blocks_per_width;\n const int block_col = (blockIdx.x / blocks_per_width) % blocks_per_chan;\n const int batch = blockIdx.x / (blocks_per_chan * blocks_per_width);\n const int w_offset = block_w * (WSZ - 1);\n\n // warp indices for loading from global mem\n const int n_warps = ceil_div((int)blockDim.x, 32);\n const int warp_id = threadIdx.x / 32;\n const int lane_id = threadIdx.x % 32;\n // make sure that the width sections can be loaded by a single warp\n assert(WSZ <= 32);\n\n // index for the TO*TC thread block\n const int threads_per_col = BC / TC;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for dout and x\n __shared__ float dout_s[BO * BW];\n __shared__ float x_s[BC * BW];\n\n // registers for local mem and local results\n float thread_results[TO * TC * 9] = {0.0};\n float reg_dout[TO] = {0.0};\n float reg_x[TC] = {0.0};\n float thread_bias_results[TO] = {0.0};\n\n // block offsets\n const int dout_block_offset = batch * O * H * W + block_row * BO * H * W;\n const int x_block_offset = batch * C * H * W + block_col * BC * H * W;\n\n //int debug_idx = 0;\n\n // we will fuse the calculation for all 9 kernel filters\n // which means we will have to reload some of the image widths\n // we can load NW blocks into smem in one go, which means we will be stepping\n // over the H dimension in (NW - 1) steps\n // this is the outer loop over H\n int h_step = (NH < H) ? NH - 1 : H;\n // if NH == H, then we are loading the entire image in one go, so step by H\n for (int h = 0; h < H; h += h_step) {\n // load dout into smem\n int h_offset = h * W;\n int nh = (h < H - (NH - 1)) ? NH : NH - 1;\n int wsz = (w_offset < W - (WSZ - 1)) ? WSZ : WSZ - 1;\n int o_per_warp = ceil_div(BO, n_warps);\n if (lane_id < wsz) {\n for (int warp_o = 0; warp_o < o_per_warp; warp_o++) {\n int inner_o = warp_id * o_per_warp + warp_o;\n if (inner_o < BO) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n dout_s[(inner_h * WSZ + inner_w) * BO + inner_o] = \n dout[dout_block_offset + inner_o * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n int c_per_warp = ceil_div(BC, n_warps);\n if (lane_id < wsz) {\n for (int warp_c = 0; warp_c < c_per_warp; warp_c++) {\n int inner_c = warp_id * c_per_warp + warp_c;\n if (inner_c < BC) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n x_s[(inner_h * WSZ + inner_w) * BC + inner_c] = \n x[x_block_offset + inner_c * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n __syncthreads();\n \n\n // if loading the last block, fill in the last block's buffer to be 0\n if (nh == NH - 1) {\n for (int i = threadIdx.x; i < WSZ * BO; i++) {\n int inner_w = i % WSZ;\n int inner_o = i / WSZ;\n dout_s[(inner_w + nh * WSZ) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < WSZ * BC; i++) {\n int inner_w = i % WSZ;\n int inner_c = i / WSZ;\n\n x_s[(inner_w + nh * WSZ) * BC + inner_c] = 0.0f;\n } \n }\n // if loading the last width section, fill in the final value in the buffer to 0\n if (wsz == WSZ - 1) {\n for (int i = threadIdx.x; i < nh * BO; i++) {\n int inner_o = i % BO;\n int inner_h = i / BO;\n dout_s[(inner_h * WSZ + wsz) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < nh * BC; i++) {\n int inner_c = i % BC;\n int inner_h = i / BC;\n x_s[(inner_h * WSZ + wsz) * BC + inner_c] = 0.0f;\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int x_dot_idx = 0; x_dot_idx < WSZ * nh; x_dot_idx++) {\n int x_h = x_dot_idx / WSZ;\n int x_w = x_dot_idx % WSZ;\n // load TC of the C channels of one image pixel into reg_w\n for (int j = 0; j < TC; j++) {\n reg_x[j] = x_s[x_dot_idx * BC + thread_col * TC + j];\n }\n\n // for each conv filter, load a different row and calculate the outer product\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n // calculate the corresponding (h, w) that should be loaded from dout\n int dout_h = x_h + 1 - k_i;\n int dout_w = x_w + 1 - k_j;\n int dout_h_abs = h + dout_h;\n int dout_w_abs = w_offset + dout_w;\n int dout_dot_idx = -1;\n if (dout_h >= 0 && dout_h < NH && dout_w >= 0 && dout_w < WSZ && dout_w_abs < W && dout_h_abs < H) {\n if (x_h == NH - 1 && k_i == 1 && H > NH) {\n // this logic is a bit tricky: for the last row in smem\n // because it will be loaded as part of the next block as well, and it will be the first row in the next block\n // for kernels 4, 5, 6 we need to only compute them once\n // this statement ensures that we don't compute the 4, 5, 6 kernels for the last row in smem\n // so it will be computed as part of the first row next time\n dout_dot_idx = -1;\n } else if (x_w == WSZ - 1 && k_j == 1 && W > WSZ) {\n // if x_w is at the last pixel, and we have more than one width sections, and we are not are the last section\n // then don't compute 1, 4, 7 kernels and wait for the next section to compute them\n dout_dot_idx = -1;\n } else {\n dout_dot_idx = dout_h * WSZ + dout_w;\n }\n }\n\n // if in bounds of the dout section in smem, load the corresponding pixel from smem\n // and add the products to the thread results\n // otherwise the contribution is 0\n if (dout_dot_idx != -1) {\n for (int i = 0; i < TO; i++) {\n reg_dout[i] = dout_s[dout_dot_idx * BO + thread_row * TO + i];\n }\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n float new_val = reg_dout[i] * reg_x[j];\n thread_results[k * TO * TC + i * TC + j] += new_val;\n }\n }\n // update dbias when k = 4, where dout_h and dout_w match x_h and x_w\n if (k == 4 && block_col == 0) {\n for (int i = 0; i < TO; i++) {\n thread_bias_results[i] += reg_dout[i];\n }\n }\n }\n }\n __syncthreads();\n }\n }\n\n // write out dweight results\n for (int i = 0; i < TO; i++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n for (int j = 0; j < TC; j++) {\n int c_abs = block_col * BC + thread_col * TC + j;\n for (int k = 0; k < 9; k++) {\n dweight_buf[o_abs * C * 9 * B * blocks_per_width + c_abs * 9 * B * blocks_per_width + k * B * blocks_per_width + batch * blocks_per_width + block_w] = thread_results[k * TO * TC + i * TC + j];\n }\n }\n // write out dbias results\n if (block_col == 0) {\n dbias_buf[o_abs * B * blocks_per_width + batch * blocks_per_width + block_w] = thread_bias_results[i];\n }\n }\n}\n\n\n__global__ void dbias_reduce_kernel1(\n float* dbias_buf, float* dbias,\n const int B\n) {\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n\n dbias_buf += blockIdx.x * B;\n\n int lane_id = threadIdx.x % 32;\n\n float db_thread_sum = 0.0f;\n for (int i = lane_id; i < B; i += 32) {\n db_thread_sum += dbias_buf[i];\n }\n\n float db_warp_sum = cg::reduce(warp, db_thread_sum, cg::plus{});\n\n if (lane_id == 0) {\n dbias[blockIdx.x] = db_warp_sum;\n }\n}\n\nvoid dweight_dbias_backward1(\n float* dout, float* x,\n float* dweight_buf, float* dweight,\n float* dbias_buf, float* dbias,\n const int B, const int C, const int O, const int H, const int W\n) {\n if (H * W == 64) {\n const int BC = 64;\n const int BO = 64;\n const int WSZ = 8;\n const int BW = WSZ * 8;\n const int TO = 4;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(cudaDeviceSynchronize());\n cudaCheck(cudaGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B, C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B);\n } else if (C == 3) {\n const int BC = 3;\n const int BO = 64;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 3;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(cudaDeviceSynchronize());\n cudaCheck(cudaGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B * W / (WSZ - 1));\n } else if (O == 3) {\n const int BC = 64;\n const int BO = 3;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 3;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(cudaDeviceSynchronize());\n cudaCheck(cudaGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B * W / (WSZ - 1));\n } else {\n const int BC = 64;\n const int BO = 64;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(cudaDeviceSynchronize());\n cudaCheck(cudaGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B * W / (WSZ - 1));\n }\n}\n\nvoid conv2d_k3_backward2(\n float* dout, float* x, float* weight,\n float* dweight_buf, float* dbias_buf,\n float* dx, float* dweight, float* dbias,\n const int B, const int C_in, const int C_out, const int H, const int W\n) {\n // dx currently requires the output buffer to be zeroed\n cudaCheck(cudaMemset(dx, 0, B * C_in * H * W * sizeof(float)));\n dx_backward(dout, weight, dx, B, C_in, C_out, H, W);\n dweight_dbias_backward1(dout, x, dweight_buf, dweight, dbias_buf, dbias, B, C_in, C_out, H, W);\n}\n// ----------------------------------------------------------------------------\n// kernel launcher\nvoid conv2d_k3_forward(\n cublasHandle_t cublas_handle,\n int kernel_num,\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C_in, int C_out, int H, int W,\n int block_size\n) {\n switch (kernel_num) {\n case 1:\n conv2d_k3_forward1(cublas_handle, x, weight, bias, out, B, C_in, C_out, H, W, block_size);\n break;\n case 2:\n conv2d_k3_forward2(cublas_handle, x, weight, bias, out, B, C_in, C_out, H, W, block_size);\n break;\n default:\n printf(\"Invalid kernel number\\n\");\n break;\n }\n}\n\n\n#ifndef LINKING\nint main( int argc, char **argv) {\n const int B = 32;\n const int C_in = 192;\n const int C_out = 64;\n const int H = 64;\n const int W = 64;\n\n // set up the device\n int deviceIdx = 0;\n cudaCheck(cudaSetDevice(deviceIdx));\n cudaDeviceProp deviceProp;\n cudaGetDeviceProperties(&deviceProp, deviceIdx);\n printf(\"Device %d: %s\\n\", deviceIdx, deviceProp.name);\n printf(\"device registers per block: %d\\n\", deviceProp.regsPerBlock);\n printf(\"shared mem per block: %ld\\n\", deviceProp.sharedMemPerBlock);\n\n // setup cublas\n cublasHandle_t cublas_handle;\n cublasCheck(cublasCreate(&cublas_handle));\n cublasCheck(cublasSetMathMode(cublas_handle, CUBLAS_DEFAULT_MATH));\n\n // create host memory to load data\n float* x = (float*)malloc(B * C_in * H * W * sizeof(float));\n float* weight = (float*)malloc(C_out * C_in * 3 * 3 * sizeof(float));\n float* bias = (float*)malloc(C_out * sizeof(float));\n float* out = (float*)malloc(B * C_out * H * W * sizeof(float));\n float* dout = (float*)malloc(B * C_out * H * W * sizeof(float));\n float* dx = (float*)malloc(B * C_in * H * W * sizeof(float));\n float* dweight = (float*)malloc(C_out * C_in * 3 * 3 * sizeof(float));\n float* dbias = (float*)malloc(C_out * sizeof(float));\n // debug\n float* dout_perm = (float*)malloc(B * C_out * H * W * sizeof(float));\n float* x_tiled = (float*)malloc(B * H * W * C_in * 9 * sizeof(float));\n\n // read saved output\n FILE *file = fopen(\"conv2d_k3.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(x, sizeof(float), B * C_in * H * W, file);\n freadCheck(weight, sizeof(float), C_out * C_in * 3 * 3, file);\n freadCheck(bias, sizeof(float), C_out, file);\n freadCheck(out, sizeof(float), B * C_out * H * W, file);\n freadCheck(dout, sizeof(float), B * C_out * H * W, file);\n freadCheck(dx, sizeof(float), B * C_in * H * W, file);\n freadCheck(dweight, sizeof(float), C_out * C_in * 3 * 3, file);\n freadCheck(dbias, sizeof(float), C_out, file);\n freadCheck(dout_perm, sizeof(float), B * C_out * H * W, file);\n freadCheck(x_tiled, sizeof(float), B * H * W * C_in * 9, file);\n fclose(file);\n\n // allocate device memory\n float *d_x, *d_weight, *d_bias, *d_out, *d_dout, *d_dx, *d_dweight, *d_dbias;\n cudaCheck(cudaMalloc(&d_x, B * C_in * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&d_weight, C_out * C_in * 3 * 3 * sizeof(float)));\n cudaCheck(cudaMalloc(&d_bias, C_out * sizeof(float)));\n cudaCheck(cudaMalloc(&d_out, B * C_out * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dout, B * C_out * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dx, B * C_in * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dweight, C_out * C_in * 3 * 3 * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dbias, C_out * sizeof(float)));\n\n cudaCheck(cudaMemset(d_out, 0, B * C_out * H * W * sizeof(float)));\n cudaCheck(cudaMemcpy(d_x, x, B * C_in * H * W * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_weight, weight, C_out * C_in * 3 * 3 * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_bias, bias, C_out * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_dout, dout, B * C_out * H * W * sizeof(float), cudaMemcpyHostToDevice));\n\n // read kernel number from command line\n int kernel_num = 2;\n if (argc > 1) {\n kernel_num = atoi(argv[1]);\n }\n printf(\"Running kernel %d\\n\", kernel_num);\n\n printf(\"Checking forward pass\\n\");\n float forward_acc = 1e-1;\n\n //for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n // int block_size = block_sizes[j];\n // printf(\"\\nBlock size: %d\\n\", block_size);\n // conv2d_k3_forward(cublas_handle, kernel_num, d_x, d_weight, d_bias, d_out, B, C_in, C_out, H, W, block_size);\n // validate_result(d_out, out, \"out\", B * C_out * H * W, forward_acc);\n //}\n //printf(\"\\nChecking cpu forward\\n\");\n //float* out_cpu_buf = (float*)malloc(B * C_out * H * W * sizeof(float));\n //conv_k3_forward_cpu(x, weight, bias, out_cpu_buf, B, C_in, C_out, H, W);\n //float* out_gpu_buf;\n //cudaCheck(cudaMalloc(&out_gpu_buf, B * C_out * H * W * sizeof(float)));\n //cudaCheck(cudaMemcpy(out_gpu_buf, out, B * C_out * H * W * sizeof(float), cudaMemcpyHostToDevice));\n //validate_result(out_gpu_buf, out, \"out cpu\", B * C_out * H * W);\n //cudaCheck(cudaFree(out_gpu_buf));\n //free(out_cpu_buf);\n\n printf(\"\\nChecking new kernel\\n\");\n conv2d_k3_forward3(d_x, d_weight, d_bias, d_out, B, C_in, C_out, H, W);\n cudaCheck(cudaGetLastError());\n printf(\"\\nChecking new kernel results\\n\");\n validate_result(d_out, out, \"out\", B * C_out * H * W, forward_acc);\n printf(\"Forward pass successful\\n\\n\");\n\n printf(\"\\nmanually benchmark new forward kernel\\n\");\n void* flush_buffer;\n cudaCheck(cudaMalloc(&flush_buffer, deviceProp.l2CacheSize));\n int repeat_times = 1000;\n float new_time = 0.0f;\n for (int i = 0; i < repeat_times; i++) {\n // clear L2\n cudaCheck(cudaMemset(flush_buffer, 0, deviceProp.l2CacheSize));\n cudaEvent_t start, end;\n cudaCheck(cudaEventCreate(&start));\n cudaCheck(cudaEventCreate(&end));\n cudaCheck(cudaEventRecord(start, nullptr));\n conv2d_k3_forward3(d_x, d_weight, d_bias, d_out, B, C_in, C_out, H, W);\n cudaCheck(cudaEventRecord(end, nullptr));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n new_time += elapsed_time;\n }\n new_time /= repeat_times;\n printf(\"new kernel time: %.4f ms\\n\", new_time);\n\n printf(\"Checking backward pass\\n\");\n\n float dw_acc = 1e-2;\n float* d_dweight_buf;\n cudaCheck(cudaMalloc(&d_dweight_buf, C_out * C_in * 3 * 3 * B * 32 * sizeof(float)));\n float* d_dbias_buf;\n cudaCheck(cudaMalloc(&d_dbias_buf, C_out * 256 * sizeof(float)));\n conv2d_k3_backward2(d_dout, d_x, d_weight, d_dweight_buf, d_dbias_buf, d_dx, d_dweight, d_dbias, B, C_in, C_out, H, W);\n\n printf(\"\\nChecking dweight\\n\");\n validate_result(d_dweight, dweight, \"dweight\", C_out * C_in * 3 * 3, dw_acc);\n\n printf(\"\\nChecking dbias\\n\");\n validate_result(d_dbias, dbias, \"dbias\", C_out, dw_acc);\n\n printf(\"\\nChecking dx\\n\");\n validate_result(d_dx, dx, \"dx\", B * C_in * H * W, dw_acc);\n\n new_time = 0.0f;\n printf(\"\\nManually benchmarking new dweight kernel\\n\");\n for (int i = 0; i < repeat_times; i++) {\n // clear L2\n cudaCheck(cudaMemset(flush_buffer, 0, deviceProp.l2CacheSize));\n cudaEvent_t start, end;\n cudaCheck(cudaEventCreate(&start));\n cudaCheck(cudaEventCreate(&end));\n cudaCheck(cudaEventRecord(start, nullptr));\n conv2d_k3_backward2(d_dout, d_x, d_weight, d_dweight_buf, d_dbias_buf, d_dx, d_dweight, d_dbias, B, C_in, C_out, H, W);\n cudaCheck(cudaEventRecord(end, nullptr));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n new_time += elapsed_time;\n }\n new_time /= repeat_times;\n printf(\"new kernel time: %.4f ms\\n\", new_time);\n cudaCheck(cudaFree(d_dweight_buf));\n cudaCheck(cudaFree(d_dbias_buf));\n\n\n\n printf(\"Backward pass successful\\n\");\n\n // old benchmarking code\n //printf(\"All results match. Starting benchmarks.\\n\\n\");\n //printf(\"Forward pass benchmarks:\\n\");\n //for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n // int block_size = block_sizes[j];\n\n // int repeat_times = 100;\n // float elapsed_time = benchmark_kernel(repeat_times, conv2d_k3_forward,\n // cublas_handle, kernel_num, d_x, d_weight, d_bias, d_out,\n // B, C_in, C_out, H, W, block_size);\n\n // float tflops = (float)B * H * W * C_in * C_out * 9 * 2 / elapsed_time * 1e3f / 1e12f;\n // printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n //}\n printf(\"\\n\\n Old benchmarks\\n\");\n printf(\"Running our own forward benchmarks\\n\");\n float t1 = 0.0f;\n float t2 = 0.0f;\n float t3 = 0.0f;\n float t_total = 0.0f;\n for (int i = 0; i < repeat_times; i++) {\n // clear L2\n cudaCheck(cudaMemset(flush_buffer, 0, deviceProp.l2CacheSize));\n conv2d_k3_forward2(cublas_handle, d_x, d_weight, d_bias, d_out, B, C_in, C_out, H, W, 512, &t1, &t2, &t3);\n }\n t1 /= repeat_times;\n t2 /= repeat_times;\n t3 /= repeat_times;\n t_total = t1 + t2 + t3;\n\n printf(\"t1: %.4f ms\\n\", t1);\n printf(\"t2: %.4f ms\\n\", t2);\n printf(\"t3: %.4f ms\\n\", t3);\n printf(\"t_total: %.4f ms\\n\", t_total);\n\n //printf(\"\\nBackward pass benchmarks:\\n\");\n //for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n // int block_size = block_sizes[j];\n\n // int repeat_times = 100;\n // float elapsed_time = benchmark_kernel(repeat_times, conv2d_k3_backward1,\n // cublas_handle, d_dout, d_x, d_weight, d_dx, d_dweight, d_dbias,\n // B, C_in, C_out, H, W, block_size);\n\n // float tflops = (float)B * H * W * C_in * C_out * (9 + 1) * 2 / elapsed_time * 1e3f / 1e12f;\n // printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n //}\n printf(\"\\nRunning our own backward benchmarks\\n\");\n t1 = 0.0f;\n t2 = 0.0f;\n t3 = 0.0f;\n float t4 = 0.0f;\n float t5 = 0.0f;\n float t6 = 0.0f;\n t_total = 0.0f;\n for (int i = 0; i < repeat_times; i++) {\n // clear L2\n cudaCheck(cudaMemset(flush_buffer, 0, deviceProp.l2CacheSize));\n conv2d_k3_backward1(cublas_handle, d_dout, d_x, d_weight, d_dx, d_dweight, d_dbias, B, C_in, C_out, H, W, 512, &t1, &t2, &t3, &t4, &t5, &t6);\n }\n cudaCheck(cudaFree(flush_buffer));\n\n t1 /= repeat_times;\n t2 /= repeat_times;\n t3 /= repeat_times;\n t4 /= repeat_times;\n t5 /= repeat_times;\n t6 /= repeat_times;\n t_total = t1 + t2 + t3 + t4 + t5 + t6;\n\n printf(\"t1: %.4f ms\\n\", t1);\n printf(\"t2: %.4f ms\\n\", t2);\n printf(\"t3: %.4f ms\\n\", t3);\n printf(\"t4: %.4f ms\\n\", t4);\n printf(\"t5: %.4f ms\\n\", t5);\n printf(\"t6: %.4f ms\\n\", t6);\n printf(\"t_total: %.4f ms\\n\", t_total);\n\n // free memory\n free(x);\n free(weight);\n free(bias);\n free(out);\n free(dout);\n free(dx);\n free(dweight);\n free(dbias);\n free(dout_perm);\n free(x_tiled);\n cudaCheck(cudaFree(d_x));\n cudaCheck(cudaFree(d_weight));\n cudaCheck(cudaFree(d_bias));\n cudaCheck(cudaFree(d_out));\n cudaCheck(cudaFree(d_dout));\n cudaCheck(cudaFree(d_dx));\n cudaCheck(cudaFree(d_dweight));\n cudaCheck(cudaFree(d_dbias));\n cublasCheck(cublasDestroy(cublas_handle));\n return 0;\n}\n#endif", "output": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \"conv2d_k3.cuh\"\n\n\n__global__ void permute_w_tile_x_kernel1(\n const float* x, const float* weight,\n float* x_tiled, float* weight_perm,\n int B, int C_in, int C_out, int H, int W\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx < C_out * C_in * 9) {\n // permute weight\n int c_out = idx / (C_in * 9);\n int c_in = (idx / 9) % C_in;\n int k = idx % 9;\n int out_ind = k * C_out * C_in + c_out * C_in + c_in;\n weight_perm[out_ind] = weight[idx];\n }\n\n if (idx < B * C_in * H * W) {\n // tile x to shape (9, C_in, B, H, W)\n // for a given first dimension k, the entry at an index should be\n // what the k-th kernel is multiplied with in the convolution for that pixel\n // (not a good explanation... but obvious with a picture)\n // for example, a value at (H, W) = (0, 0) will never be convolved with kernels 5, 7 and 8\n // because it is on a corner\n int x_offset = B * C_in * H * W;\n int img_size = H * W;\n int b = idx / (C_in * img_size);\n int c = (idx / img_size) % C_in;\n int h_start = (idx / W) % H;\n int w_start = idx % W;\n float x_val = x[idx];\n for (int i = 0; i < 3; i++) {\n for (int j = 0; j < 3; j++) {\n int h_end = h_start + 1 - i;\n int w_end = w_start + 1 - j;\n if (h_end >= 0 && h_end < H && w_end >= 0 && w_end < W) {\n x_tiled[c * B * img_size + b * img_size + h_end * W + w_end] = x_val;\n }\n // shift pointer to next kernel\n x_tiled += x_offset;\n }\n }\n }\n}\n\n\n__global__ void reduce_transpose_add_bias_kernel1(\n const float* out_tmp, const float* bias,\n float* out,\n int B, int C_out, int H, int W\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n // make sure we're in bound\n assert(idx < B * C_out * H * W);\n // make sure every block lies in the same channel\n assert((B * H * W) % blockDim.x == 0);\n\n int img_size = H * W;\n int c = idx / (B * img_size);\n int b = (idx / img_size) % B;\n int h = (idx / W) % H;\n int w = idx % W;\n\n __shared__ float bias_sh[1];\n if (threadIdx.x == 0) {\n bias_sh[0] = bias[c];\n }\n __syncthreads();\n\n float out_val = bias_sh[0];\n\n for (int k = 0; k < 9; k++) {\n float out_tmp_val = out_tmp[k * C_out * B * img_size + c * B * img_size + b * img_size + h * W + w];\n out_val += out_tmp_val;\n }\n\n out[b * C_out * img_size + c * img_size + h * W + w] = out_val;\n}\n\nvoid conv2d_k3_forward1(\n hipblasHandle_t cublas_handle,\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C_in, int C_out, int H, int W,\n int block_size\n) {\n // implements Conv2D with kernel size 1 and padding 1\n // shapes:\n // x (B, C_in, H, W)\n // weight (C_out, C_in, 3, 3)\n // bias (C_out,)\n\n // plan:\n // 1. transpose and tile x into x_tiled of shape (9, C_in, B, H, W)\n // 2. transpose weight into shape (9, C_out, C_in) (we will do this in the same kernel as one above)\n // 3. call cublasSgemmStridedBatched with 9 batches\n // 4. reduce over the 9 kernel dimensions, add bias, and transpose output to shape (B, C_out, H, W)\n\n // 1. and 2.\n // transpose weight and tile x\n float *x_tiled, *weight_perm;\n cudaCheck(hipMalloc(&x_tiled, 9 * C_in * B * H * W * sizeof(float)));\n cudaCheck(hipMemset(x_tiled, 0, 9 * C_in * B * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&weight_perm, 9 * C_out * C_in * sizeof(float)));\n int total_th_1 = max_int(B * C_in * H * W, C_in * C_out * 9);\n int n_blk_1 = ceil_div(total_th_1, block_size);\n permute_w_tile_x_kernel1<<>>(x, weight, x_tiled, weight_perm, B, C_in, C_out, H, W);\n\n // 3. matmul via cublasSgemmStridedBatched\n float* out_tmp;\n cudaCheck(hipMalloc(&out_tmp, 9 * C_out * B * H * W * sizeof(float)));\n float alpha = 1.0f;\n float beta = 0.0f;\n int bhw = B * H * W;\n cublasCheck(hipblasSgemmStridedBatched(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_N, bhw, C_out, C_in, &alpha, x_tiled, bhw, C_in * bhw, weight_perm, C_in, C_out * C_in, &beta, out_tmp, bhw, C_out * bhw, 9));\n\n // reduce, add bias, and transpose\n int total_th_2 = B * C_out * H * W;\n int n_blk_2 = ceil_div(total_th_2, block_size);\n reduce_transpose_add_bias_kernel1<<>>(out_tmp, bias, out, B, C_out, H, W);\n\n // free memory\n cudaCheck(hipFree(x_tiled));\n cudaCheck(hipFree(weight_perm));\n cudaCheck(hipFree(out_tmp));\n}\n\n\n__global__ void transpose_tile_x_kernel2(\n const float* x,\n float* x_tiled,\n int B, int C_in, int H, int W\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * C_in * H * W) {\n // similar to permute_w_tile_x_kernel1, except now the target shape is (C_in, 9, B, H, W)\n int img_size = H * W;\n int b = idx / (C_in * img_size);\n int c = (idx / img_size) % C_in;\n int h_start = (idx / W) % H;\n int w_start = idx % W;\n float x_val = x[idx];\n for (int i = 0; i < 3; i++) {\n for (int j = 0; j < 3; j++) {\n int h_end = h_start + 1 - i;\n int w_end = w_start + 1 - j;\n if (h_end >= 0 && h_end < H && w_end >= 0 && w_end < W) {\n x_tiled[c * 9 * B * img_size + (i * 3 + j) * B * img_size + b * img_size + h_end * W + w_end] = x_val;\n }\n }\n }\n }\n}\n\n__global__ void transpose_out_add_bias_kernel2(\n const float* out_tmp, const float* bias,\n float* out,\n int B, int C_out, int H, int W\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n // make sure we're in bound\n assert(idx < B * C_out * H * W);\n // make sure every block lies in the same channel\n assert((B * H * W) % blockDim.x == 0);\n\n int img_size = H * W;\n int c = idx / (B * img_size);\n int b = (idx / img_size) % B;\n int h = (idx / W) % H;\n int w = idx % W;\n\n __shared__ float bias_sh[1];\n if (threadIdx.x == 0) {\n bias_sh[0] = bias[c];\n }\n __syncthreads();\n \n out[b * C_out * img_size + c * img_size + h * W + w] = bias_sh[0] + out_tmp[idx];\n}\n\nvoid conv2d_k3_forward2(\n hipblasHandle_t cublas_handle,\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C_in, int C_out, int H, int W,\n int block_size,\n float* t1, float* t2, float* t3\n) {\n // similar to kernel 1, but transpose into a different shape so that we can just make one big matmul call\n // and no need to reduce afterwards\n // plan:\n // 1. transpose and tile x into x_tiled of shape (C_in, 9, B, H, W)\n // 2. call cublasSgemm on x_tiled and weight, no need for transposes!\n // 3. add bias\n\n // 1. transpose x to x_tiled\n hipEvent_t start, end1, end2, end3;\n if (t1 != nullptr) {\n cudaCheck(hipEventCreate(&start));\n cudaCheck(hipEventCreate(&end1));\n cudaCheck(hipEventCreate(&end2));\n cudaCheck(hipEventCreate(&end3));\n cudaCheck(hipEventRecord(start, nullptr));\n }\n float* x_tiled;\n cudaCheck(hipMalloc(&x_tiled, C_in * 9 * B * H * W * sizeof(float)));\n cudaCheck(hipMemset(x_tiled, 0, C_in * 9 * B * H * W * sizeof(float)));\n if (t1 != nullptr) {\n cudaCheck(hipEventRecord(end1, nullptr));\n }\n int total_th_1 = B * C_in * H * W;\n int n_blk_1 = ceil_div(total_th_1, block_size);\n transpose_tile_x_kernel2<<>>(x, x_tiled, B, C_in, H, W);\n\n // 2. matmul with cublasSgemm\n float* out_tmp;\n cudaCheck(hipMalloc(&out_tmp, C_out * B * H * W * sizeof(float)));\n float alpha = 1.0f;\n float beta = 0.0f;\n int bhw = B * H * W;\n cublasCheck(hipblasSgemm(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_N, bhw, C_out, C_in * 9, &alpha, x_tiled, bhw, weight, C_in * 9, &beta, out_tmp, bhw));\n\n\n // 3. transpose and add bias\n int total_th_2 = B * C_out * H * W;\n int n_blk_2 = ceil_div(total_th_2, block_size);\n transpose_out_add_bias_kernel2<<>>(out_tmp, bias, out, B, C_out, H, W);\n if (t2 != nullptr) {\n cudaCheck(hipEventRecord(end2, nullptr));\n }\n\n\n // free memory\n cudaCheck(hipFree(x_tiled));\n cudaCheck(hipFree(out_tmp));\n if (t3 != nullptr) {\n cudaCheck(hipEventRecord(end3, nullptr));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end1));\n cudaCheck(hipEventSynchronize(end2));\n cudaCheck(hipEventSynchronize(end3));\n float elapsed_time1, elapsed_time2, elapsed_time3;\n cudaCheck(hipEventElapsedTime(&elapsed_time1, start, end1));\n cudaCheck(hipEventElapsedTime(&elapsed_time2, end1, end2));\n cudaCheck(hipEventElapsedTime(&elapsed_time3, end2, end3));\n *t1 += elapsed_time1;\n *t2 += elapsed_time2;\n *t3 += elapsed_time3;\n }\n}\n\n// this kernel will roughly follow kernel 5 here: https://github.com/siboehm/SGEMM_CUDA, with a few differences\n// the main operation is convolving the input X with the weights\n// weights are shaped (O, Cx9), and X is shaped (B, C, HxW)\n// We are morally doing a matmul with dimensions O, C, HxW\ntemplate \n__global__ void conv2d_k3_forward_kernel1(\n const float* x, const float* weight, const float* bias,\n float* out,\n const int B, const int C, const int O, const int H, const int W\n) {\n // constants\n const int C9 = C * 9;\n const int weight_size = O * C * 9;\n const int x_size = B * C * H * W;\n // n_threads / block = (BO * BH) / (TO * TH) = 256\n // n_blocks = B * (O / BO) * (H * W / BH)\n // we will set gridDim to (BxHxW / BH, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // for compute step, to determine which of the size TOxTH block the thread is responsible for\n const int threads_per_col = ceil_div(BH, TH);\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // shared mem for loading weights and X\n const int w_s_sz = BO * BC * 9;\n const int w_s_width = BC * 9;\n const int x_s_sz = BC * BH;\n __shared__ float w_s[w_s_sz];\n __shared__ float x_s[x_s_sz];\n\n // registers to store local results\n float thread_results[TO * TH * 9] = {0.0};\n // registers for weights and x\n float reg_w[TO] = {0.0};\n float reg_x[TH] = {0.0};\n\n // calculate offset in x for bounds checking later\n const int x_block_offset = batch * C * H * W + block_col * BH;\n // calculate weight offset for bounds checking\n const int w_block_offset = block_row * BO * C9;\n\n // outer loop over BC size blocks of channel C\n for (int c_idx = 0; c_idx < C; c_idx += BC) {\n // x and weight offsets for this block\n int x_inner_offset = c_idx * H * W;\n int w_inner_offset = c_idx * 9;\n // load weight into smem\n int w_abs_idx;\n for (int i = threadIdx.x; i < w_s_sz; i += blockDim.x) {\n int w_i = i / w_s_width;\n int w_j = i % w_s_width;\n w_abs_idx = w_block_offset + w_inner_offset + w_i * C9 + w_j;\n if (w_abs_idx < weight_size) {\n w_s[w_i * w_s_width + w_j] = weight[w_abs_idx];\n }\n }\n // load x into smem\n int x_abs_idx;\n for (int i = threadIdx.x; i < x_s_sz; i += blockDim.x) {\n int x_i = i / BH;\n int x_j = i % BH;\n x_abs_idx = x_block_offset + x_inner_offset + x_i * H * W + x_j;\n if (x_abs_idx < x_size) {\n x_s[x_i * BH + x_j] = x[x_abs_idx];\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int dot_idx = 0; dot_idx < BC; dot_idx++) {\n // load x_s into registers\n for (int j = 0; j < TH; j++) {\n reg_x[j] = x_s[dot_idx * BH + thread_col * TH + j];\n }\n for (int conv_idx = 0; conv_idx < 9; conv_idx++) {\n for (int i = 0; i < TO; i++) {\n reg_w[i] = w_s[(thread_row * TO + i) * w_s_width + dot_idx * 9 + conv_idx];\n }\n\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n float val = reg_w[i] * reg_x[j];\n thread_results[conv_idx * TO * TH + i * TH + j] += val;\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n // need to be careful here: for each pixel at (h_abs, w_abs) and its product with kernel k (k < 9)\n // There are 2 options\n // 1. the product should be added to some pixel, then do an atomiAdd\n // 2. the product should not contribute to any pixel because it would be out of bound, in this case just skip\n int out_batch_offset = batch * O * H * W;\n for (int j = 0; j < TH; j++) {\n int hw_abs = block_col * BH + thread_col * TH + j;\n int h_abs = hw_abs / W;\n int w_abs = hw_abs % W;\n for (int i = 0; i < TO; i++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target_abs = h_abs + 1 - k_i;\n int w_target_abs = w_abs + 1 - k_j;\n if (h_target_abs >= 0 && h_target_abs < H && w_target_abs >= 0 && w_target_abs < W) {\n int out_target_abs = out_batch_offset + o_abs * H * W + h_target_abs * W + w_target_abs;\n int th_res_idx = k * TO * TH + i * TH + j;\n atomicAdd(out + out_target_abs, thread_results[th_res_idx]);\n }\n }\n int out_abs = out_batch_offset + o_abs * H * W + h_abs * W + w_abs;\n atomicAdd(out + out_abs, bias[o_abs]);\n }\n }\n}\n\n\n// almost same as 1, except we load A as transposed\ntemplate \n__global__ void conv2d_k3_forward_kernel2(\n const float* x, const float* weight, const float* bias,\n float* out,\n const int B, const int C, const int O, const int H, const int W\n) {\n // constants\n const int C9 = C * 9;\n const int weight_size = O * C * 9;\n const int x_size = B * C * H * W;\n // n_threads / block = (BO * BH) / (TO * TH) = 256\n // n_blocks = B * (O / BO) * (H * W / BH)\n // we will set gridDim to (BxHxW / BH, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // for compute step, to determine which of the size TOxTH block the thread is responsible for\n const int threads_per_col = ceil_div(BH, TH);\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // shared mem for loading weights and X\n const int w_s_sz = BO * BC * 9;\n const int w_s_width = BC * 9;\n const int x_s_sz = BC * BH;\n __shared__ float w_s[w_s_sz];\n __shared__ float x_s[x_s_sz];\n\n // registers to store local results\n float thread_results[TO * TH * 9] = {0.0};\n // registers for weights and x\n float reg_w[TO] = {0.0};\n float reg_x[TH] = {0.0};\n\n // calculate offset in x for bounds checking later\n const int x_block_offset = batch * C * H * W + block_col * BH;\n // calculate weight offset for bounds checking\n const int w_block_offset = block_row * BO * C9;\n\n // outer loop over BC size blocks of channel C\n for (int c_idx = 0; c_idx < C; c_idx += BC) {\n // x and weight offsets for this block\n int x_inner_offset = c_idx * H * W;\n int w_inner_offset = c_idx * 9;\n // load weight into smem\n int w_abs_idx;\n for (int i = threadIdx.x; i < w_s_sz; i += blockDim.x) {\n int w_i = i / w_s_width;\n int w_j = i % w_s_width;\n w_abs_idx = w_block_offset + w_inner_offset + w_i * C9 + w_j;\n if (w_abs_idx < weight_size) {\n w_s[w_j * BO + w_i] = weight[w_abs_idx];\n }\n }\n // load x into smem\n int x_abs_idx;\n for (int i = threadIdx.x; i < x_s_sz; i += blockDim.x) {\n int x_i = i / BH;\n int x_j = i % BH;\n x_abs_idx = x_block_offset + x_inner_offset + x_i * H * W + x_j;\n if (x_abs_idx < x_size) {\n x_s[x_i * BH + x_j] = x[x_abs_idx];\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int dot_idx = 0; dot_idx < BC; dot_idx++) {\n // load x_s into registers\n for (int j = 0; j < TH; j++) {\n reg_x[j] = x_s[dot_idx * BH + thread_col * TH + j];\n }\n for (int conv_idx = 0; conv_idx < 9; conv_idx++) {\n for (int i = 0; i < TO; i++) {\n reg_w[i] = w_s[(dot_idx * 9 + conv_idx) * BO + thread_row * TO + i];\n }\n\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n float val = reg_w[i] * reg_x[j];\n thread_results[conv_idx * TO * TH + i * TH + j] += val;\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n // need to be careful here: for each pixel at (h_abs, w_abs) and its product with kernel k (k < 9)\n // There are 2 options\n // 1. the product should be added to some pixel, then do an atomiAdd\n // 2. the product should not contribute to any pixel because it would be out of bound, in this case just skip\n int out_batch_offset = batch * O * H * W;\n for (int j = 0; j < TH; j++) {\n int hw_abs = block_col * BH + thread_col * TH + j;\n int h_abs = hw_abs / W;\n int w_abs = hw_abs % W;\n for (int i = 0; i < TO; i++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target_abs = h_abs + 1 - k_i;\n int w_target_abs = w_abs + 1 - k_j;\n if (h_target_abs >= 0 && h_target_abs < H && w_target_abs >= 0 && w_target_abs < W) {\n int out_target_abs = out_batch_offset + o_abs * H * W + h_target_abs * W + w_target_abs;\n int th_res_idx = k * TO * TH + i * TH + j;\n atomicAdd(out + out_target_abs, thread_results[th_res_idx]);\n }\n }\n int out_abs = out_batch_offset + o_abs * H * W + h_abs * W + w_abs;\n atomicAdd(out + out_abs, bias[o_abs]);\n }\n }\n}\n\n\n// almost same as 2, but we vectorize the loads and use float4\ntemplate \n__global__ void conv2d_k3_forward_kernel3(\n float* x, float* weight, float* bias,\n float* out,\n const int B, const int C, const int O, const int H, const int W\n) {\n // constants\n assert(TH == 4); // needed for reg_x loads later, so we can do one vectorized load for each thread\n const int C9 = C * 9;\n // n_threads / block = (BO * BH) / (TO * TH) = 256\n // n_blocks = B * (O / BO) * (H * W / BH)\n // we will set gridDim to (BxHxW / BH, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // for compute step, to determine which of the size TOxTH block the thread is responsible for\n const int threads_per_col = ceil_div(BH, TH);\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // shared mem for loading weights and X\n const int w_s_sz = BO * BC * 9;\n const int w_s_width = BC * 9;\n const int x_s_sz = BC * BH;\n __shared__ float w_s[w_s_sz];\n __shared__ float x_s[x_s_sz];\n\n // registers to store local results\n float thread_results[TO * TH * 9] = {0.0};\n // registers for weights and x\n float reg_w[TO] = {0.0};\n float reg_x[TH] = {0.0};\n\n // calculate offset in x for bounds checking later\n const int x_block_offset = batch * C * H * W + block_col * BH;\n // calculate weight offset for bounds checking\n const int w_block_offset = block_row * BO * C9;\n\n // outer loop over BC size blocks of channel C\n for (int c_idx = 0; c_idx < C; c_idx += BC) {\n // x and weight offsets for this block\n int x_inner_offset = c_idx * H * W;\n int w_inner_offset = c_idx * 9;\n // load weight into smem\n for (int i = threadIdx.x; i < w_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (w_s_width / 4);\n int inner_row = i / (w_s_width / 4);\n float4 tmp = \n reinterpret_cast(&weight[w_block_offset + w_inner_offset + inner_row * C9 + inner_col * 4])[0];\n w_s[(inner_col * 4 + 0) * BO + inner_row] = tmp.x;\n w_s[(inner_col * 4 + 1) * BO + inner_row] = tmp.y;\n w_s[(inner_col * 4 + 2) * BO + inner_row] = tmp.z;\n w_s[(inner_col * 4 + 3) * BO + inner_row] = tmp.w;\n }\n // load x into smem\n for (int i = threadIdx.x; i < x_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (BH / 4);\n int inner_row = i / (BH / 4);\n reinterpret_cast(&x_s[inner_row * BH + inner_col * 4])[0] = \n reinterpret_cast(&x[x_block_offset + x_inner_offset + inner_row * H * W + inner_col * 4])[0];\n }\n __syncthreads();\n\n // calculate per thread results\n for (int dot_idx = 0; dot_idx < BC; dot_idx++) {\n // load x_s into registers\n // NOTE: relying on the fact that TH = 4\n reinterpret_cast(reg_x)[0] = reinterpret_cast(&x_s[dot_idx * BH + thread_col * TH])[0];\n for (int conv_idx = 0; conv_idx < 9; conv_idx++) {\n for (int i = 0; i < TO; i++) {\n reg_w[i] = w_s[(dot_idx * 9 + conv_idx) * BO + thread_row * TO + i];\n }\n\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n float val = reg_w[i] * reg_x[j];\n thread_results[conv_idx * TO * TH + i * TH + j] += val;\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n // need to be careful here: for each pixel at (h_abs, w_abs) and its product with kernel k (k < 9)\n // There are 2 options\n // 1. the product should be added to some pixel, then do an atomiAdd\n // 2. the product should not contribute to any pixel because it would be out of bound, in this case just skip\n int out_batch_offset = batch * O * H * W;\n for (int j = 0; j < TH; j++) {\n int hw_abs = block_col * BH + thread_col * TH + j;\n int h_abs = hw_abs / W;\n int w_abs = hw_abs % W;\n for (int i = 0; i < TO; i++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target_abs = h_abs + 1 - k_i;\n int w_target_abs = w_abs + 1 - k_j;\n if (h_target_abs >= 0 && h_target_abs < H && w_target_abs >= 0 && w_target_abs < W) {\n int out_target_abs = out_batch_offset + o_abs * H * W + h_target_abs * W + w_target_abs;\n int th_res_idx = k * TO * TH + i * TH + j;\n atomicAdd(out + out_target_abs, thread_results[th_res_idx]);\n }\n }\n int out_abs = out_batch_offset + o_abs * H * W + h_abs * W + w_abs;\n atomicAdd(out + out_abs, bias[o_abs]);\n }\n }\n}\n\n// CPU version for debugging\nvoid conv_k3_forward_cpu(\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C, int OC, int H, int W\n) {\n for (int b = 0; b < B; b++) {\n for (int oc = 0; oc < OC; oc++) {\n for (int h = 0; h < H; h++) {\n for (int w = 0; w < W; w++) {\n for (int k = 0; k < 9; k++) {\n // check if the output will actually be used, if not skip\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target = h + 1 - k_i;\n int w_target = w + 1 - k_j;\n if (h_target < 0 || h_target >= H || w_target < 0 || w_target >= W) {\n continue;\n }\n float out_bohw = 0.0f;\n for (int c = 0; c < C; c++) {\n float new_val = weight[oc * C * 9 + c * 9 + k] * x[b * C * H * W + c * H * W + h * W + w];\n out_bohw += new_val;\n //if (b == 0 && h == 0 && w == 1 && k == 5 && oc == 0) {\n // printf(\"(0, 1) OC 0 kern 5: contribution from %d: %f, out val %f\\n\", c, new_val, out_bohw);\n //}\n }\n int out_idx = b * OC * H * W + oc * H * W + h_target * W + w_target;\n //if (h_target == 0 && w_target == 0 && oc == 0 && b == 0) {\n // printf(\"(0, 0) getting contribution from (%d, %d) at k %d: %f, prev_val %f\\n\", h, w, k, out_bohw, out[out_idx]);\n //}\n out[out_idx] += out_bohw;\n //if (h_target == 0 && w_target == 0 && oc == 0 && b == 0) {\n // printf(\"(0, 0) getting contribution from (%d, %d) at k %d: %f, new val %f\\n\", h, w, k, out_bohw, out[out_idx]);\n //}\n }\n \n //if (h == 0 && w == 0 && oc == 0 && b == 0) {\n // printf(\"(0, 0) adding bias %f\\n\", bias[oc]);\n //}\n out[b * OC * H * W + oc * H * W + h * W + w] += bias[oc];\n }\n }\n printf(\"(0, 0) final value: %f\\n\", out[0]);\n return;\n }\n }\n}\n\nvoid conv2d_k3_forward3(\n float* x, float* weight, float* bias,\n float* out,\n const int B, const int C_in, const int C_out, const int H, const int W\n) {\n // since the kernel will add to the output, we need to zero it first\n cudaCheck(hipMemset(out, 0, B * C_out * H * W * sizeof(float)));\n\n if (H * W == 64) {\n const int BH = 64;\n const int TH = 4;\n assert(C_out % 64 == 0 && C_in % 8 == 0);\n const int BC = 16;\n const int BO = 64;\n const int TO = 4;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel3<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n } else if (C_in == 3) {\n const int BH = 64;\n assert(H * W % BH == 0);\n const int TH = 4;\n const int BC = 3;\n assert(C_out % 64 == 0);\n const int BO = 64;\n const int TO = 4;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel2<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n } else if (C_out == 3) {\n const int BH = 64 * 2;\n assert(H * W % BH == 0);\n const int TH = 4;\n assert(C_in % 8 == 0);\n const int BC = 8;\n const int BO = 3;\n const int TO = 3;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel2<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n } else {\n const int BH = 64 * 4;\n assert(H * W % BH == 0);\n const int TH = 4;\n assert(C_out % 64 == 0 && C_in % 8 == 0);\n const int BC = 16;\n const int BO = 8;\n const int TO = 2;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k3_forward_kernel3<<>>(\n x, weight, bias, out, B, C_in, C_out, H, W\n );\n }\n cudaCheck(hipGetLastError());\n}\n\n// ----------------------------------------------------------------------------\n// backward pass\n\n__global__ void permute_dout_tile_x_kernel(\n const float* dout, const float* x,\n float* dout_perm, float* x_tiled,\n int B, int C_in, int C_out, int H, int W\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n if (idx < B * C_out * H * W) {\n // permute dout\n int b = idx / (C_out * img_size);\n int c = (idx / img_size) % C_out;\n int h = (idx / W) % H;\n int w = idx % W;\n dout_perm[c * B * img_size + b * img_size + h * W + w] = dout[idx];\n }\n\n if (idx < B * C_in * H * W) {\n // tile x_tiled, see permute_w_tile_x_kernel1 for explanation\n // except that we tile into a shape of (B, H, W, C_in, 9), the transpose of the tiling there\n int b = idx / (C_in * img_size);\n int c = (idx / img_size) % C_in;\n int h_start = (idx / W) % H;\n int w_start = idx % W;\n for (int i = 0; i < 3; i++) {\n for (int j = 0; j < 3; j++) {\n int h_end = h_start + 1 - i;\n int w_end = w_start + 1 - j;\n if (h_end >= 0 && h_end < H && w_end >= 0 && w_end < W) {\n x_tiled[b * img_size * C_in * 9 + h_end * W * C_in * 9 + w_end * C_in * 9 + c * 9 + i * 3 + j]\n = x[idx];\n }\n }\n }\n }\n}\n\n__global__ void reduce_dx_kernel(\n const float* dx_tiled,\n float* dx,\n int B, int C_in, int H, int W\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * C_in * H * W) {\n int img_size = H * W;\n int b = idx / (C_in * img_size);\n int c = (idx / img_size) % C_in;\n int h_end = (idx / W) % H;\n int w_end = idx % W;\n float sum = 0.0f;\n for (int i = 0; i < 3; i++) {\n for (int j = 0; j < 3; j++) {\n int h_start = h_end + 1 - i;\n int w_start = w_end + 1 - j;\n if (h_start >= 0 && h_start < H && w_start >= 0 && w_start < W) {\n sum += dx_tiled[c * 9 * B * img_size + (i * 3 + j) * B * img_size + b * img_size + h_start * W + w_start];\n }\n }\n }\n dx[idx] = sum;\n }\n}\n\n\n__global__ void reduce_dbias_kernel1(\n const float* dout_perm,\n float* dbias,\n int B, int C_out, int H, int W\n) {\n // dout_perm has shape (C_out, B, H, W)\n // each block will reduce over one channel\n dout_perm += blockIdx.x * B * H * W;\n int bhw = B * H * W;\n assert(blockDim.x <= bhw);\n extern __shared__ float sdata[];\n \n float sum = 0.0f;\n for (int i = threadIdx.x; i < bhw; i += blockDim.x) {\n sum += dout_perm[i];\n }\n sdata[threadIdx.x] = sum;\n\n __syncthreads();\n\n for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {\n if (threadIdx.x < stride) {\n sdata[threadIdx.x] += sdata[threadIdx.x + stride];\n }\n __syncthreads();\n }\n\n if (threadIdx.x == 0) {\n dbias[blockIdx.x] = sdata[0];\n }\n}\n\n// use cg to do warp reduce\n// doesn't really seem faster\n__global__ void reduce_dbias_kernel2(\n const float* dout_perm,\n float* dbias,\n int B, int C_out, int H, int W\n) {\n int bhw = B * H * W;\n assert(blockDim.x <= bhw);\n // move dout_perm to the right block\n dout_perm += blockIdx.x * bhw;\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n __shared__ float shared_sum[32];\n\n int num_warps = blockDim.x / 32;\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n\n float thread_sum = 0.0f;\n for (int i = threadIdx.x; i < bhw; i += blockDim.x) {\n thread_sum += dout_perm[i];\n }\n \n // warp reduce\n float warp_sum = cg::reduce(warp, thread_sum, cg::plus{});\n shared_sum[warp_id] = warp_sum;\n __syncthreads();\n\n // warp reduce\n warp_sum = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;\n float block_sum = cg::reduce(warp, warp_sum, cg::plus{});\n\n if (threadIdx.x == 0) {\n dbias[blockIdx.x] = block_sum;\n }\n}\n\nvoid conv2d_k3_backward1(\n hipblasHandle_t cublas_handle,\n const float* dout, const float* x, const float* weight,\n float* dx, float* dweight, float* dbias,\n int B, int C_in, int C_out, int H, int W,\n int block_size,\n float* t1, float* t2, float* t3, float* t4, float* t5, float* t6\n) {\n // shapes:\n // dout (B, C_out, H, W)\n // x (B, C_in, H, W)\n // weight (C_out, C_in, 3, 3)\n // bias (C_out,)\n\n // we need to calculate dx, dweight, dbias separately\n // dbias is a single reduce\n // for dweight, we need to tile x into x_tile of shape (B * H * W, C_in * 9)\n // then do a matmul with dout_perm (shape (C_out, B * H * W))\n // for dx, we do dout_perm @ weight.T (store it in x_tile), and then reduce to get dx\n // If we don't want to do a reduce when computing dx, we need to make 9 copies of dout\n // which might be ok but increases memory usage. We can try it later.\n hipEvent_t start, end1, end2, end3, end4, end5, end6;\n if (t1 != nullptr) {\n cudaCheck(hipEventCreate(&start));\n cudaCheck(hipEventCreate(&end1));\n cudaCheck(hipEventCreate(&end2));\n cudaCheck(hipEventCreate(&end3));\n cudaCheck(hipEventCreate(&end4));\n cudaCheck(hipEventCreate(&end5));\n cudaCheck(hipEventCreate(&end6));\n cudaCheck(hipEventRecord(start, nullptr));\n }\n\n // 1. transpose to dout_perm and tile x_tile\n float *dout_perm, *x_tiled;\n cudaCheck(hipMalloc(&x_tiled, B * H * W * C_in * 9 * sizeof(float)));\n cudaCheck(hipMemset(x_tiled, 0, B * H * W * C_in * 9 * sizeof(float)));\n cudaCheck(hipMalloc(&dout_perm, B * C_out * H * W * sizeof(float)));\n if (t1 != nullptr) {\n cudaCheck(hipEventRecord(end1, nullptr));\n }\n int total_th_1 = max_int(C_in, C_out) * B * H * W;\n int n_blk_1 = ceil_div(total_th_1, block_size);\n permute_dout_tile_x_kernel<<>>(dout, x, dout_perm, x_tiled, B, C_in, C_out, H, W);\n cudaCheck(hipGetLastError());\n if (t2 != nullptr) {\n cudaCheck(hipEventRecord(end2, nullptr));\n }\n\n // 2. get dweight by single matmul between x_tiled and dout_perm\n float alpha = 1.0f;\n float beta = 0.0f;\n int bhw = B * H * W;\n cublasCheck(hipblasSgemm(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_N, C_in * 9, C_out, bhw, &alpha, x_tiled, C_in * 9, dout_perm, bhw, &beta, dweight, C_in * 9));\n\n if (t2 != nullptr) {\n cudaCheck(hipEventRecord(end3, nullptr));\n }\n\n // 3. get dx_tile by matmul between dout_perm and weight.T, and save output in x_tiled\n // note: beta = 0 overwrites existing values in x_tiled\n cublasCheck(hipblasSgemm(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_T, bhw, C_in * 9, C_out, &alpha, dout_perm, bhw, weight, C_in * 9, &beta, x_tiled, bhw));\n\n\n // 4. reduce x_tiled to get dx\n int total_th_2 = B * C_in * H * W;\n int n_blk_2 = ceil_div(total_th_2, block_size);\n reduce_dx_kernel<<>>(x_tiled, dx, B, C_in, H, W);\n\n if (t2 != nullptr) {\n cudaCheck(hipEventRecord(end4, nullptr));\n }\n\n // 5. reduce dbias\n //reduce_dbias_kernel1<<>>(dout_perm, dbias, B, C_out, H, W);\n reduce_dbias_kernel2<<>>(dout_perm, dbias, B, C_out, H, W);\n\n if (t5 != nullptr) {\n cudaCheck(hipEventRecord(end5, nullptr));\n }\n\n // free memory\n cudaCheck(hipFree(dout_perm));\n cudaCheck(hipFree(x_tiled));\n if (t5 != nullptr) {\n cudaCheck(hipEventRecord(end6, nullptr));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end1));\n cudaCheck(hipEventSynchronize(end2));\n cudaCheck(hipEventSynchronize(end3));\n cudaCheck(hipEventSynchronize(end4));\n cudaCheck(hipEventSynchronize(end5));\n cudaCheck(hipEventSynchronize(end6));\n float elapsed1, elapsed2, elapsed3, elapsed4, elapsed5, elapsed6;\n cudaCheck(hipEventElapsedTime(&elapsed1, start, end1));\n cudaCheck(hipEventElapsedTime(&elapsed2, end1, end2));\n cudaCheck(hipEventElapsedTime(&elapsed3, end2, end3));\n cudaCheck(hipEventElapsedTime(&elapsed4, end3, end4));\n cudaCheck(hipEventElapsedTime(&elapsed5, end4, end5));\n cudaCheck(hipEventElapsedTime(&elapsed6, end5, end6));\n *t1 += elapsed1;\n *t2 += elapsed2;\n *t3 += elapsed3;\n *t4 += elapsed4;\n *t5 += elapsed5;\n *t6 += elapsed6;\n }\n}\n\n// almost the same as conv2d_k3_forward_kernel, except that the compute and the writes\n// are different, because now we want to treat the C dimension as the inner product dimension\ntemplate \n__global__ void dx_backward_kernel1(\n float* dout, float* weight,\n float* dx,\n const int B, const int C, const int O, const int H, const int W\n) {\n const int C9 = C * 9;\n // n_threads per block = (BC * BH) / (TC * TH)\n // set gridDim to (BxHxW / BH, C / TC)\n // find the output block that this thread block is responsible for\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // index for the TCxTH block that the thread is responsible\n const int threads_per_col = BH / TH;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for weights and dout\n const int w_s_width = BC * 9;\n const int w_s_sz = BO * w_s_width;\n const int dout_s_sz = BO * BH;\n __shared__ float w_s[w_s_sz];\n __shared__ float dout_s[dout_s_sz];\n\n // registers for local results\n float thread_results[TC * TH * 9] = {0.0};\n // regs for weights and dout\n float reg_w[TC * 9] = {0.0};\n float reg_dout[TH] = {0.0};\n\n // calculate the block offsets\n const int dout_block_offset = batch * O * H * W + block_col * BH;\n const int w_block_offset = block_row * w_s_width;\n\n // outer loop over BO sizes of channel O\n for (int o_idx = 0; o_idx < O; o_idx += BO) {\n // offsets for this inner block\n int w_inner_offset = o_idx * C9;\n int dout_inner_offset = o_idx * H * W;\n\n // load weights to smem\n if (BC % 4 == 0) {\n // vectorize dram load\n for (int i = threadIdx.x; i < w_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (w_s_width / 4);\n int inner_row = i / (w_s_width / 4);\n reinterpret_cast(&w_s[inner_row * w_s_width + inner_col * 4])[0] = \n reinterpret_cast(&weight[w_block_offset + w_inner_offset + inner_row * C9 + inner_col * 4])[0];\n }\n } else {\n for (int i = threadIdx.x; i < w_s_sz; i += blockDim.x) {\n int inner_col = i % w_s_width;\n int inner_row = i / w_s_width;\n w_s[inner_row * w_s_width + inner_col] = weight[w_block_offset + w_inner_offset + inner_row * C9 + inner_col];\n }\n }\n // load dout\n for (int i = threadIdx.x; i < dout_s_sz / 4; i += blockDim.x) {\n int inner_col = i % (BH / 4);\n int inner_row = i / (BH / 4);\n reinterpret_cast(&dout_s[inner_row * BH + inner_col * 4])[0] = \n reinterpret_cast(&dout[dout_block_offset + dout_inner_offset + inner_row * H * W + inner_col * 4])[0];\n }\n __syncthreads();\n \n\n // calculate per thread results\n for (int dot_idx = 0; dot_idx < BO; dot_idx++) {\n // load dout and w into registers\n for (int j = 0; j < (TH / 4); j++) {\n reinterpret_cast(®_dout[j * 4])[0] = \n reinterpret_cast(&dout_s[dot_idx * BH + thread_col * TH + j * 4])[0];\n }\n // non vectorized load\n //for (int j = 0; j < TH; j++) {\n // reg_dout[j] = dout_s[dot_idx * BH + thread_col * TH + j];\n //}\n if (TC % 4 == 0) {\n // vectorize smem load\n for (int i = 0; i < (TC * 9 / 4); i++) {\n reinterpret_cast(®_w[i * 4])[0] = \n reinterpret_cast(&w_s[dot_idx * w_s_width + thread_row * TC * 9 + i * 4])[0];\n }\n } else {\n for (int i = 0; i < TC * 9; i++) {\n reg_w[i] = w_s[dot_idx * w_s_width + thread_row * TC * 9 + i];\n }\n }\n\n for (int k = 0; k < 9; k++) {\n for (int i = 0; i < TC; i++) {\n for (int j = 0; j < TH; j++) {\n thread_results[k * TC * TH + i * TH + j] += \n reg_w[i * 9 + k] * reg_dout[j];\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n int out_batch_offset = batch * C * H * W;\n for (int j = 0; j < TH; j++) {\n int hw_abs = block_col * BH + thread_col * TH + j;\n int h_abs = hw_abs / W;\n int w_abs = hw_abs % W;\n for (int i = 0; i < TC; i++) {\n int c_abs = block_row * BC + thread_row * TC + i;\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target_abs = h_abs - 1 + k_i;\n int w_target_abs = w_abs - 1 + k_j;\n if (h_target_abs >= 0 && h_target_abs < H && w_target_abs >= 0 && w_target_abs < W) {\n int out_target_abs = out_batch_offset + c_abs * H * W + h_target_abs * W + w_target_abs;\n int th_res_idx = k * TC * TH + i * TH + j;\n atomicAdd(dx + out_target_abs, thread_results[th_res_idx]);\n }\n }\n }\n }\n}\n\n// I wonder how fast does a transposition kernel take for the weights\n__global__ void transpose_weight_kernel(\n float* weight, float* out,\n int C, int O\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < C * O * 9) {\n int o = idx / (C * 9);\n int c = (idx % C * 9) / 9;\n int k = idx % 9;\n out[c * O * 9 + o * 9 + k] = weight[idx];\n }\n}\n\nvoid dx_backward(\n float* dout, float* weight,\n float* dx,\n const int B, const int C_in, const int C_out, const int H, const int W\n) {\n if (H < 16) {\n const int BH = 64;\n const int TH = 4;\n const int BC = 16;\n const int TC = 4;\n const int BO = 16;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n } else if (C_in == 3) {\n const int BH = 64 * 4;\n const int TH = 4;\n const int BC = 3;\n const int TC = 3;\n const int BO = 16;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n } else if (C_out == 3) {\n const int BH = 64 * 4;\n const int TH = 4;\n const int BC = 16;\n const int TC = 4;\n const int BO = 3;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n } else {\n const int BH = 64 * 4;\n const int TH = 4;\n const int BC = 16;\n const int TC = 4;\n const int BO = 16;\n dim3 gridDim(B * H * W / BH, C_in / BC);\n dim3 blockDim((BC * BH) / (TC * TH));\n dx_backward_kernel1<<>>(dout, weight, dx, B, C_in, C_out, H, W);\n }\n}\n\nvoid dx_backward_cpu(\n float* dout, float* weight, float* dx,\n int B, int C, int O, int H, int W\n) {\n printf(\"dx[0] before zeroing: %f\\n\", dx[0]);\n for (int b = 0; b < B; b++) {\n for (int c = 0; c < C; c++) {\n for (int h = 0; h < H; h++) {\n for (int w = 0; w < W; w++) {\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n int h_target = h - 1 + k_i;\n int w_target = w - 1 + k_j;\n if (h_target < 0 || h_target >= H || w_target < 0 || w_target >= W) {\n continue;\n }\n float out_bchw = 0.0f;\n for (int o = 0; o < O; o++) {\n float new_val = weight[o * C * 9 + c * 9 + k] * dout[b * O * H * W + o * H * W + h * W + w];\n out_bchw += new_val;\n }\n int out_idx = b * C * H * W + c * H * W + h_target * W + w_target;\n dx[out_idx] += out_bchw;\n }\n }\n }\n if (b == 0 && c == 0) {\n printf(\"dx[0]: %f\\n\", dx[0]);\n return;\n }\n }\n }\n}\n\n// dweight_buf will have shape (O, C, 9, B)\n// and then we have to reduce over the last dimension to get the actual weight\ntemplate \n__global__ void dweight_backward_kernel1(\n float* dout, float* x,\n float* dweight_buf,\n const int B, const int C, const int O, const int H, const int W\n) {\n // ensure that we load entire image widths into smem\n // so we can do all 9 kernels at once\n assert(BW % W == 0);\n\n // n_threads per block = (BC * BO) / (TO * TC)\n // gridDim will be (B * C / BC, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = C / BC;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // index for the TO*TC thread block\n const int threads_per_col = BC / TC;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for dout and x\n __shared__ float dout_s[BO * BW];\n __shared__ float x_s[BC * BW];\n\n // registers for local mem and local results\n float thread_results[TO * TC * 9] = {0.0};\n float reg_dout[TO] = {0.0};\n float reg_x[TC] = {0.0};\n\n // block offsets\n const int dout_block_offset = batch * O * H * W + block_row * BO * H * W;\n const int x_block_offset = batch * C * H * W + block_col * BC * H * W;\n\n // number of image widths we will load to smem in one round\n const int NW = BW / W;\n\n // we will fuse the calculation for all 9 kernel filters\n // which means we will have to reload some of the image widths\n // we can load NW blocks into smem in one go, which means we will be stepping\n // over the H dimension in (NW - 1) steps\n // this is the outer loop over H\n for (int h = 0; h < H; h += (NW - 1)) {\n // load dout into smem\n int dout_inner_offset = h * W;\n int nh = (h < H - (NW - 1)) ? NW : NW - 1;\n int dout_load_sz = nh * W * BO;\n for (int i = threadIdx.x; i < dout_load_sz / 4; i++) {\n int inner_col = i % (nh * W / 4);\n int inner_row = i / (nh * W / 4);\n float4 tmp = \n reinterpret_cast(&dout[dout_block_offset + dout_inner_offset + inner_row * H * W + inner_col * 4])[0];\n dout_s[(inner_col * 4 + 0) * BO + inner_row] = tmp.x;\n dout_s[(inner_col * 4 + 1) * BO + inner_row] = tmp.y;\n dout_s[(inner_col * 4 + 2) * BO + inner_row] = tmp.z;\n dout_s[(inner_col * 4 + 3) * BO + inner_row] = tmp.w;\n }\n \n // load x to smem\n int x_inner_offset = h * W;\n int x_load_sz = nh * W * BC;\n for (int i = threadIdx.x; i < x_load_sz / 4; i++) {\n int inner_col = i % (nh * W / 4);\n int inner_row = i / (nh * W / 4);\n float4 tmp = \n reinterpret_cast(&x[x_block_offset + x_inner_offset + inner_row * H * W + inner_col * 4])[0];\n x_s[(inner_col * 4 + 0) * BC + inner_row] = tmp.x;\n x_s[(inner_col * 4 + 1) * BC + inner_row] = tmp.y;\n x_s[(inner_col * 4 + 2) * BC + inner_row] = tmp.z;\n x_s[(inner_col * 4 + 3) * BC + inner_row] = tmp.w;\n }\n __syncthreads();\n\n // if loading the last block, fill in the last block's buffer to be 0\n if (nh == NW - 1) {\n for (int i = threadIdx.x; i < W * BO; i++) {\n int inner_col = i % W;\n int inner_row = i / W;\n dout_s[(inner_col + nh * W) * BO + inner_row] = 0.0f;\n }\n for (int i = threadIdx.x; i < W * BC; i++) {\n int inner_col = i % W;\n int inner_row = i / W;\n x_s[(inner_col + nh * W) * BC + inner_row] = 0.0f;\n } \n }\n __syncthreads();\n\n // calculate per thread results\n for (int x_dot_idx = 0; x_dot_idx < BW; x_dot_idx++) {\n int x_h = x_dot_idx / W;\n int x_w = x_dot_idx % W;\n // load one pixel with TC of the C channels into reg_w\n for (int j = 0; j < TC; j++) {\n reg_x[j] = x_s[x_dot_idx * BC + thread_col * TC + j];\n }\n\n // for each conv filter, load a different row and calculate the outer product\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n // calculate the corresponding (h, w) that should be loaded from dout\n int dout_h = x_h + 1 - k_i;\n int dout_w = x_w + 1 - k_j;\n int dout_dot_idx = -1;\n if (dout_h >= 0 && dout_h < NW && dout_w >= 0 && dout_w < W) {\n // this logic is a bit tricky: for the last row in smem\n // because it will be loaded as part of the next block as well, and it will be the first row in the next block\n // for kernels 4, 5, 6 we need to only compute them once\n // this statement ensures that we don't compute the 4, 5, 6 kernels for the last row in smem\n // so it will be computed as part of the first row next time\n if (x_h == NW - 1 && k_i == 1) {\n dout_dot_idx = -1;\n } else {\n dout_dot_idx = dout_h * W + dout_w;\n }\n }\n // if in bounds of the dout section in smem, load the corresponding pixel from smem\n // and add the products to the thread results\n // otherwise the contribution is 0\n if (dout_dot_idx != -1) {\n for (int i = 0; i < TO; i++) {\n reg_dout[i] = dout_s[dout_dot_idx * BO + thread_row * TO + i];\n }\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n thread_results[k * TO * TC + i * TC + j] += reg_dout[i] * reg_x[j];\n }\n }\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n for (int k = 0; k < 9; k++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n int c_abs = block_col * BC + thread_col * TC + j;\n dweight_buf[o_abs * C * 9 * B + c_abs * 9 * B + k * B + batch] = thread_results[k * TO * TC + i * TC + j];\n }\n }\n }\n}\n\n\n// C * O blocks\n// 288 (9 * 32) threads per block, use one warp to reduce each batch dim\n__global__ void dweight_reduce_kernel(\n float* dweight_buf, float* dweight,\n const int B, const int C, const int O\n) {\n assert(gridDim.x == O * C);\n int o = blockIdx.x / C;\n int c = blockIdx.x % C;\n // move input pointer\n dweight_buf += o * C * 9 * B + c * 9 * B;\n assert(blockDim.x == 32 * 9);\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n\n float thread_sum = 0.0f;\n for (int i = lane_id; i < B; i += 32) {\n thread_sum += dweight_buf[warp_id * B + i];\n }\n\n // warp reduce\n float warp_sum = cg::reduce(warp, thread_sum, cg::plus{});\n\n if (lane_id == 0) {\n dweight[o * C * 9 + c * 9 + warp_id] = warp_sum;\n }\n}\n\nvoid dweight_backward1(\n float* dout, float* x,\n float* dweight_buf, float* dweight,\n const int B, const int C, const int O, const int H, const int W\n) {\n const int BC = 32;\n const int BO = 32;\n const int BW = 2 * 64;\n const int TO = 4;\n const int TC = 4;\n dim3 gridDim(B * C / BC, O / BO);\n dim3 blockDim(BC * BO / (TO * TC));\n dweight_backward_kernel1<<>>(dout, x, dweight_buf, B, C, O, H, W);\n cudaCheck(hipGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B, C, O);\n}\n\n\n// instead of having to load multiples of entire widths of the image into smem\n// we now only load fractions of the width each time\n// so the outer loop needs to loop over the height multiples (NH - 1) as well as the width multiples (WS - 1)\n// results: This did seem to help with memory access: compute intensity is higher\n// right now this one is much slower than kernel 1\n// I think this is because the smem loads are now not aligned, which causes a lot of uncoalesced loads\ntemplate \n__global__ void dweight_backward_kernel2(\n float* dout, float* x,\n float* dweight_buf,\n const int B, const int C, const int O, const int H, const int W\n) {\n // number of heights we need to load into smem in one go \n assert(BW % WSZ == 0);\n const int NH = BW / WSZ;\n // n_threads per block = (BC * BO) / (TO * TC)\n // gridDim will be (B * C / BC, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = C / BC;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // index for the TO*TC thread block\n const int threads_per_col = BC / TC;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for dout and x\n __shared__ float dout_s[BO * BW];\n __shared__ float x_s[BC * BW];\n\n // registers for local mem and local results\n float thread_results[TO * TC * 9] = {0.0};\n float reg_dout[TO] = {0.0};\n float reg_x[TC] = {0.0};\n\n // block offsets\n const int dout_block_offset = batch * O * H * W + block_row * BO * H * W;\n const int x_block_offset = batch * C * H * W + block_col * BC * H * W;\n\n //int debug_idx = 0;\n\n // we will fuse the calculation for all 9 kernel filters\n // which means we will have to reload some of the image widths\n // we can load NW blocks into smem in one go, which means we will be stepping\n // over the H dimension in (NW - 1) steps\n // this is the outer loop over H\n for (int h = 0; h < H; h += (NH - 1)) {\n for (int w_offset = 0; w_offset < W; w_offset += (WSZ - 1)) {\n // load dout into smem\n int h_offset = h * W;\n int nh = (h < H - (NH - 1)) ? NH : NH - 1;\n int wsz = (w_offset < W - (WSZ - 1)) ? WSZ : WSZ - 1;\n int dout_load_sz = nh * wsz * BO;\n for (int i = threadIdx.x; i < dout_load_sz; i++) {\n int inner_w = i % wsz;\n int inner_h = (i / wsz) % nh;\n int inner_o = i / (wsz * nh);\n dout_s[(inner_h * WSZ + inner_w) * BO + inner_o] = \n dout[dout_block_offset + h_offset + inner_h * W + w_offset + inner_w + inner_o * H * W];\n }\n \n // load x to smem\n int x_load_sz = nh * wsz * BC;\n for (int i = threadIdx.x; i < x_load_sz; i++) {\n int inner_w = i % wsz;\n int inner_h = (i / wsz) % nh;\n int inner_c = i / (wsz * nh);\n x_s[(inner_h * WSZ + inner_w) * BC + inner_c] = \n x[x_block_offset + h_offset + inner_h * W + w_offset + inner_w + inner_c * H * W];\n }\n\n __syncthreads();\n\n // if loading the last block, fill in the last block's buffer to be 0\n if (nh == NH - 1) {\n for (int i = threadIdx.x; i < WSZ * BO; i++) {\n int inner_w = i % WSZ;\n int inner_o = i / WSZ;\n dout_s[(inner_w + nh * WSZ) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < WSZ * BC; i++) {\n int inner_w = i % WSZ;\n int inner_c = i / WSZ;\n\n x_s[(inner_w + nh * WSZ) * BC + inner_c] = 0.0f;\n } \n }\n // if loading the last width section, fill in the final value in the buffer to 0\n if (wsz == WSZ - 1) {\n for (int i = threadIdx.x; i < nh * BO; i++) {\n int inner_o = i % BO;\n int inner_h = i / BO;\n dout_s[(inner_h * WSZ + wsz) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < nh * BC; i++) {\n int inner_c = i % BC;\n int inner_h = i / BC;\n x_s[(inner_h * WSZ + wsz) * BC + inner_c] = 0.0f;\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int x_dot_idx = 0; x_dot_idx < WSZ * nh; x_dot_idx++) {\n int x_h = x_dot_idx / WSZ;\n int x_w = x_dot_idx % WSZ;\n // load TC of the C channels of one image pixel into reg_w\n for (int j = 0; j < TC; j++) {\n reg_x[j] = x_s[x_dot_idx * BC + thread_col * TC + j];\n }\n\n // for each conv filter, load a different row and calculate the outer product\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n // calculate the corresponding (h, w) that should be loaded from dout\n int dout_h = x_h + 1 - k_i;\n int dout_w = x_w + 1 - k_j;\n int dout_h_abs = h + dout_h;\n int dout_w_abs = w_offset + dout_w;\n int dout_dot_idx = -1;\n if (dout_h >= 0 && dout_h < NH && dout_w >= 0 && dout_w < WSZ && dout_w_abs < W && dout_h_abs < H) {\n if (x_h == NH - 1 && k_i == 1) {\n // this logic is a bit tricky: for the last row in smem\n // because it will be loaded as part of the next block as well, and it will be the first row in the next block\n // for kernels 4, 5, 6 we need to only compute them once\n // this statement ensures that we don't compute the 4, 5, 6 kernels for the last row in smem\n // so it will be computed as part of the first row next time\n dout_dot_idx = -1;\n } else if (x_w == WSZ - 1 && k_j == 1) {\n dout_dot_idx = -1;\n } else {\n dout_dot_idx = dout_h * WSZ + dout_w;\n }\n }\n\n // if in bounds of the dout section in smem, load the corresponding pixel from smem\n // and add the products to the thread results\n // otherwise the contribution is 0\n if (dout_dot_idx != -1) {\n for (int i = 0; i < TO; i++) {\n assert(dout_dot_idx * BO + thread_row * TO + i < BO * BW);\n reg_dout[i] = dout_s[dout_dot_idx * BO + thread_row * TO + i];\n }\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n float new_val = reg_dout[i] * reg_x[j];\n thread_results[k * TO * TC + i * TC + j] += new_val;\n }\n }\n }\n }\n }\n __syncthreads();\n }\n }\n\n // write out results\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n for (int k = 0; k < 9; k++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n int c_abs = block_col * BC + thread_col * TC + j;\n dweight_buf[o_abs * C * 9 * B + c_abs * 9 * B + k * B + batch] = thread_results[k * TO * TC + i * TC + j];\n }\n }\n }\n}\n\nvoid dweight_backward2(\n float* dout, float* x,\n float* dweight_buf, float* dweight,\n const int B, const int C, const int O, const int H, const int W\n) {\n const int BC = 32;\n const int BO = 32;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 4;\n dim3 gridDim(B * C / BC, O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_backward_kernel2<<>>(dout, x, dweight_buf, B, C, O, H, W);\n cudaCheck(hipDeviceSynchronize());\n cudaCheck(hipGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B, C, O);\n}\n\n\n// kernel 2 does something silly: because are now loading WSZ sections of widths,\n// and WSZ is definitely not a nice power of 2, we need to be careful\n// which thread we are using when loading global memory.\n// This kernel only changes the way loading is done compared to kernel 2\n// effectively we do some warp tiling when loading from global mem\n// kernel 3 is now quite fast (2.6ms)\n// profiling tells us that we are now using very few blocks, so seems like a good idea\n// to spread out the different width sections on different blocks and do a bigger reduce afterwards\ntemplate \n__global__ void dweight_backward_kernel3(\n float* dout, float* x,\n float* dweight_buf,\n const int B, const int C, const int O, const int H, const int W\n) {\n // number of heights we need to load into smem in one go \n const int NH = BW / WSZ;\n // n_threads per block = (BC * BO) / (TO * TC)\n // gridDim will be (B * C / BC, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = C / BC;\n assert(gridDim.x == B * blocks_per_col);\n const int batch = blockIdx.x / blocks_per_col;\n const int block_col = blockIdx.x % blocks_per_col;\n\n // warp indices for loading from global mem\n const int n_warps = blockDim.x / 32;\n const int warp_id = threadIdx.x / 32;\n const int lane_id = threadIdx.x % 32;\n // make sure that the width sections can be loaded by a single warp\n assert(WSZ <= 32);\n\n // index for the TO*TC thread block\n const int threads_per_col = BC / TC;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for dout and x\n __shared__ float dout_s[BO * BW];\n __shared__ float x_s[BC * BW];\n\n // registers for local mem and local results\n float thread_results[TO * TC * 9] = {0.0};\n float reg_dout[TO] = {0.0};\n float reg_x[TC] = {0.0};\n\n // block offsets\n const int dout_block_offset = batch * O * H * W + block_row * BO * H * W;\n const int x_block_offset = batch * C * H * W + block_col * BC * H * W;\n\n //int debug_idx = 0;\n\n // we will fuse the calculation for all 9 kernel filters\n // which means we will have to reload some of the image widths\n // we can load NW blocks into smem in one go, which means we will be stepping\n // over the H dimension in (NW - 1) steps\n // this is the outer loop over H\n for (int h = 0; h < H; h += (NH - 1)) {\n for (int w_offset = 0; w_offset < W; w_offset += (WSZ - 1)) {\n // load dout into smem\n int h_offset = h * W;\n int nh = (h < H - (NH - 1)) ? NH : NH - 1;\n int wsz = (w_offset < W - (WSZ - 1)) ? WSZ : WSZ - 1;\n int o_per_warp = ceil_div(BO, n_warps);\n if (lane_id < wsz) {\n for (int warp_o = 0; warp_o < o_per_warp; warp_o++) {\n int inner_o = warp_id * o_per_warp + warp_o;\n if (inner_o < BO) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n dout_s[(inner_h * WSZ + inner_w) * BO + inner_o] = \n dout[dout_block_offset + inner_o * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n int c_per_warp = ceil_div(BC, n_warps);\n if (lane_id < wsz) {\n for (int warp_c = 0; warp_c < c_per_warp; warp_c++) {\n int inner_c = warp_id * c_per_warp + warp_c;\n if (inner_c < BC) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n x_s[(inner_h * WSZ + inner_w) * BC + inner_c] = \n x[x_block_offset + inner_c * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n __syncthreads();\n \n\n // if loading the last block, fill in the last block's buffer to be 0\n if (nh == NH - 1) {\n for (int i = threadIdx.x; i < WSZ * BO; i++) {\n int inner_w = i % WSZ;\n int inner_o = i / WSZ;\n dout_s[(inner_w + nh * WSZ) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < WSZ * BC; i++) {\n int inner_w = i % WSZ;\n int inner_c = i / WSZ;\n\n x_s[(inner_w + nh * WSZ) * BC + inner_c] = 0.0f;\n } \n }\n // if loading the last width section, fill in the final value in the buffer to 0\n if (wsz == WSZ - 1) {\n for (int i = threadIdx.x; i < nh * BO; i++) {\n int inner_o = i % BO;\n int inner_h = i / BO;\n dout_s[(inner_h * WSZ + wsz) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < nh * BC; i++) {\n int inner_c = i % BC;\n int inner_h = i / BC;\n x_s[(inner_h * WSZ + wsz) * BC + inner_c] = 0.0f;\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int x_dot_idx = 0; x_dot_idx < WSZ * nh; x_dot_idx++) {\n int x_h = x_dot_idx / WSZ;\n int x_w = x_dot_idx % WSZ;\n // load TC of the C channels of one image pixel into reg_w\n for (int j = 0; j < TC; j++) {\n reg_x[j] = x_s[x_dot_idx * BC + thread_col * TC + j];\n }\n\n // for each conv filter, load a different row and calculate the outer product\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n // calculate the corresponding (h, w) that should be loaded from dout\n int dout_h = x_h + 1 - k_i;\n int dout_w = x_w + 1 - k_j;\n int dout_h_abs = h + dout_h;\n int dout_w_abs = w_offset + dout_w;\n int dout_dot_idx = -1;\n if (dout_h >= 0 && dout_h < NH && dout_w >= 0 && dout_w < WSZ && dout_w_abs < W && dout_h_abs < H) {\n if (x_h == NH - 1 && k_i == 1) {\n // this logic is a bit tricky: for the last row in smem\n // because it will be loaded as part of the next block as well, and it will be the first row in the next block\n // for kernels 4, 5, 6 we need to only compute them once\n // this statement ensures that we don't compute the 4, 5, 6 kernels for the last row in smem\n // so it will be computed as part of the first row next time\n dout_dot_idx = -1;\n } else if (x_w == WSZ - 1 && k_j == 1) {\n dout_dot_idx = -1;\n } else {\n dout_dot_idx = dout_h * WSZ + dout_w;\n }\n }\n\n // if in bounds of the dout section in smem, load the corresponding pixel from smem\n // and add the products to the thread results\n // otherwise the contribution is 0\n if (dout_dot_idx != -1) {\n for (int i = 0; i < TO; i++) {\n reg_dout[i] = dout_s[dout_dot_idx * BO + thread_row * TO + i];\n }\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n float new_val = reg_dout[i] * reg_x[j];\n thread_results[k * TO * TC + i * TC + j] += new_val;\n }\n }\n }\n }\n }\n __syncthreads();\n }\n }\n\n // write out results\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n for (int k = 0; k < 9; k++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n int c_abs = block_col * BC + thread_col * TC + j;\n dweight_buf[o_abs * C * 9 * B + c_abs * 9 * B + k * B + batch] = thread_results[k * TO * TC + i * TC + j];\n }\n }\n }\n}\n\nvoid dweight_backward3(\n float* dout, float* x,\n float* dweight_buf, float* dweight,\n const int B, const int C, const int O, const int H, const int W\n) {\n const int BC = 64;\n const int BO = 64;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * C / BC, O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_backward_kernel3<<>>(dout, x, dweight_buf, B, C, O, H, W);\n cudaCheck(hipDeviceSynchronize());\n cudaCheck(hipGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B, C, O);\n}\n\n\n// trying to improve on kernel 3 by splitting the widths of each image\n// to be calculated on different blocks\n// Output will now be of shape (O, C, 9, B, W / (WSZ - 1))\n// and we will need to reduce over the last two dimensions\ntemplate \n__global__ void dweight_backward_kernel4(\n float* dout, float* x,\n float* dweight_buf,\n const int B, const int C, const int O, const int H, const int W\n) {\n // number of heights we need to load into smem in one go \n const int NH = BW / WSZ;\n // n_threads per block = (BC * BO) / (TO * TC)\n // gridDim will be (B * (W / (WSZ - 1)) * (C / BC), O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_chan = C / BC;\n const int blocks_per_width = W / (WSZ - 1);\n const int block_w = blockIdx.x % blocks_per_width;\n const int block_col = (blockIdx.x / blocks_per_width) % blocks_per_chan;\n const int batch = blockIdx.x / (blocks_per_chan * blocks_per_width);\n const int w_offset = block_w * (WSZ - 1);\n\n // warp indices for loading from global mem\n const int n_warps = blockDim.x / 32;\n const int warp_id = threadIdx.x / 32;\n const int lane_id = threadIdx.x % 32;\n // make sure that the width sections can be loaded by a single warp\n assert(WSZ <= 32);\n\n // index for the TO*TC thread block\n const int threads_per_col = BC / TC;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for dout and x\n __shared__ float dout_s[BO * BW];\n __shared__ float x_s[BC * BW];\n\n // registers for local mem and local results\n float thread_results[TO * TC * 9] = {0.0};\n float reg_dout[TO] = {0.0};\n float reg_x[TC] = {0.0};\n\n // block offsets\n const int dout_block_offset = batch * O * H * W + block_row * BO * H * W;\n const int x_block_offset = batch * C * H * W + block_col * BC * H * W;\n\n //int debug_idx = 0;\n\n // we will fuse the calculation for all 9 kernel filters\n // which means we will have to reload some of the image widths\n // we can load NW blocks into smem in one go, which means we will be stepping\n // over the H dimension in (NW - 1) steps\n // this is the outer loop over H\n for (int h = 0; h < H; h += (NH - 1)) {\n // load dout into smem\n int h_offset = h * W;\n int nh = (h < H - (NH - 1)) ? NH : NH - 1;\n int wsz = (w_offset < W - (WSZ - 1)) ? WSZ : WSZ - 1;\n int o_per_warp = ceil_div(BO, n_warps);\n if (lane_id < wsz) {\n for (int warp_o = 0; warp_o < o_per_warp; warp_o++) {\n int inner_o = warp_id * o_per_warp + warp_o;\n if (inner_o < BO) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n dout_s[(inner_h * WSZ + inner_w) * BO + inner_o] = \n dout[dout_block_offset + inner_o * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n int c_per_warp = ceil_div(BC, n_warps);\n if (lane_id < wsz) {\n for (int warp_c = 0; warp_c < c_per_warp; warp_c++) {\n int inner_c = warp_id * c_per_warp + warp_c;\n if (inner_c < BC) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n x_s[(inner_h * WSZ + inner_w) * BC + inner_c] = \n x[x_block_offset + inner_c * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n __syncthreads();\n \n\n // if loading the last block, fill in the last block's buffer to be 0\n if (nh == NH - 1) {\n for (int i = threadIdx.x; i < WSZ * BO; i++) {\n int inner_w = i % WSZ;\n int inner_o = i / WSZ;\n dout_s[(inner_w + nh * WSZ) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < WSZ * BC; i++) {\n int inner_w = i % WSZ;\n int inner_c = i / WSZ;\n\n x_s[(inner_w + nh * WSZ) * BC + inner_c] = 0.0f;\n } \n }\n // if loading the last width section, fill in the final value in the buffer to 0\n if (wsz == WSZ - 1) {\n for (int i = threadIdx.x; i < nh * BO; i++) {\n int inner_o = i % BO;\n int inner_h = i / BO;\n dout_s[(inner_h * WSZ + wsz) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < nh * BC; i++) {\n int inner_c = i % BC;\n int inner_h = i / BC;\n x_s[(inner_h * WSZ + wsz) * BC + inner_c] = 0.0f;\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int x_dot_idx = 0; x_dot_idx < WSZ * nh; x_dot_idx++) {\n int x_h = x_dot_idx / WSZ;\n int x_w = x_dot_idx % WSZ;\n // load TC of the C channels of one image pixel into reg_w\n for (int j = 0; j < TC; j++) {\n reg_x[j] = x_s[x_dot_idx * BC + thread_col * TC + j];\n }\n\n // for each conv filter, load a different row and calculate the outer product\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n // calculate the corresponding (h, w) that should be loaded from dout\n int dout_h = x_h + 1 - k_i;\n int dout_w = x_w + 1 - k_j;\n int dout_h_abs = h + dout_h;\n int dout_w_abs = w_offset + dout_w;\n int dout_dot_idx = -1;\n if (dout_h >= 0 && dout_h < NH && dout_w >= 0 && dout_w < WSZ && dout_w_abs < W && dout_h_abs < H) {\n if (x_h == NH - 1 && k_i == 1) {\n // this logic is a bit tricky: for the last row in smem\n // because it will be loaded as part of the next block as well, and it will be the first row in the next block\n // for kernels 4, 5, 6 we need to only compute them once\n // this statement ensures that we don't compute the 4, 5, 6 kernels for the last row in smem\n // so it will be computed as part of the first row next time\n dout_dot_idx = -1;\n } else if (x_w == WSZ - 1 && k_j == 1) {\n dout_dot_idx = -1;\n } else {\n dout_dot_idx = dout_h * WSZ + dout_w;\n }\n }\n\n // if in bounds of the dout section in smem, load the corresponding pixel from smem\n // and add the products to the thread results\n // otherwise the contribution is 0\n if (dout_dot_idx != -1) {\n for (int i = 0; i < TO; i++) {\n reg_dout[i] = dout_s[dout_dot_idx * BO + thread_row * TO + i];\n }\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n float new_val = reg_dout[i] * reg_x[j];\n thread_results[k * TO * TC + i * TC + j] += new_val;\n }\n }\n }\n }\n __syncthreads();\n }\n }\n\n // write out results\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n for (int k = 0; k < 9; k++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n int c_abs = block_col * BC + thread_col * TC + j;\n dweight_buf[o_abs * C * 9 * B * blocks_per_width + c_abs * 9 * B * blocks_per_width + k * B * blocks_per_width + batch * blocks_per_width + block_w] = thread_results[k * TO * TC + i * TC + j];\n }\n }\n }\n}\n\nvoid dweight_backward4(\n float* dout, float* x,\n float* dweight_buf, float* dweight,\n const int B, const int C, const int O, const int H, const int W\n) {\n const int BC = 64;\n const int BO = 64;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_backward_kernel4<<>>(dout, x, dweight_buf, B, C, O, H, W);\n cudaCheck(hipDeviceSynchronize());\n cudaCheck(hipGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n}\n\n// same as kernel 4, but vectorize loads to speed up smem load\n// unfortunately as is this doesn't give any speedup\n// probably because we are now using too few threads to load\n// I could try to shuffle and make each warp load more heights as well, but that doesn't seem first order right now\ntemplate \n__global__ void dweight_backward_kernel5(\n float* dout, float* x,\n float* dweight_buf,\n const int B, const int C, const int O, const int H, const int W\n) {\n // number of heights we need to load into smem in one go \n const int NH = BW / WSZ;\n // n_threads per block = (BC * BO) / (TO * TC)\n // gridDim will be (B * (W / (WSZ - 1)) * (C / BC), O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_chan = C / BC;\n const int blocks_per_width = W / (WSZ - 1);\n const int block_w = blockIdx.x % blocks_per_width;\n const int block_col = (blockIdx.x / blocks_per_width) % blocks_per_chan;\n const int batch = blockIdx.x / (blocks_per_chan * blocks_per_width);\n const int w_offset = block_w * (WSZ - 1);\n\n // warp indices for loading from global mem\n const int n_warps = blockDim.x / 32;\n const int warp_id = threadIdx.x / 32;\n const int lane_id = threadIdx.x % 32;\n // make sure that the width sections can be loaded by a single warp\n assert(WSZ <= 32);\n\n // index for the TO*TC thread block\n const int threads_per_col = BC / TC;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for dout and x\n __shared__ float dout_s[BO * BW];\n __shared__ float x_s[BC * BW];\n\n // registers for local mem and local results\n float thread_results[TO * TC * 9] = {0.0};\n float reg_dout[TO] = {0.0};\n float reg_x[TC] = {0.0};\n\n // block offsets\n const int dout_block_offset = batch * O * H * W + block_row * BO * H * W;\n const int x_block_offset = batch * C * H * W + block_col * BC * H * W;\n\n //int debug_idx = 0;\n\n // we will fuse the calculation for all 9 kernel filters\n // which means we will have to reload some of the image widths\n // we can load NW blocks into smem in one go, which means we will be stepping\n // over the H dimension in (NW - 1) steps\n // this is the outer loop over H\n for (int h = 0; h < H; h += (NH - 1)) {\n // load dout into smem\n int h_offset = h * W;\n int nh = (h < H - (NH - 1)) ? NH : NH - 1;\n int wsz = (w_offset < W - (WSZ - 1)) ? WSZ : WSZ - 1;\n int o_per_warp = ceil_div(BO, n_warps);\n if (lane_id < ceil_div(wsz, 4)) {\n for (int warp_o = 0; warp_o < o_per_warp; warp_o++) {\n int inner_o = warp_id * o_per_warp + warp_o;\n if (inner_o < BO) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n float4 tmp = \n reinterpret_cast(&dout[dout_block_offset + inner_o * H * W + h_offset + inner_h * W + w_offset + lane_id * 4])[0];\n dout_s[(inner_h * WSZ + lane_id * 4 + 0) * BO + inner_o] = tmp.x;\n if (lane_id < wsz / 4) {\n dout_s[(inner_h * WSZ + lane_id * 4 + 1) * BO + inner_o] = tmp.y;\n dout_s[(inner_h * WSZ + lane_id * 4 + 2) * BO + inner_o] = tmp.z;\n dout_s[(inner_h * WSZ + lane_id * 4 + 3) * BO + inner_o] = tmp.w;\n }\n }\n }\n }\n }\n\n int c_per_warp = ceil_div(BC, n_warps);\n if (lane_id < ceil_div(wsz, 4)) {\n for (int warp_c = 0; warp_c < c_per_warp; warp_c++) {\n int inner_c = warp_id * c_per_warp + warp_c;\n if (inner_c < BC) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n float4 tmp = \n reinterpret_cast(&x[x_block_offset + inner_c * H * W + h_offset + inner_h * W + w_offset + lane_id * 4])[0];\n x_s[(inner_h * WSZ + lane_id * 4 + 0) * BC + inner_c] = tmp.x;\n if (lane_id < wsz / 4) {\n x_s[(inner_h * WSZ + lane_id * 4 + 1) * BC + inner_c] = tmp.y;\n x_s[(inner_h * WSZ + lane_id * 4 + 2) * BC + inner_c] = tmp.z;\n x_s[(inner_h * WSZ + lane_id * 4 + 3) * BC + inner_c] = tmp.w;\n }\n }\n }\n }\n }\n\n __syncthreads();\n \n\n // if loading the last block, fill in the last block's buffer to be 0\n if (nh == NH - 1) {\n for (int i = threadIdx.x; i < WSZ * BO; i++) {\n int inner_w = i % WSZ;\n int inner_o = i / WSZ;\n dout_s[(inner_w + nh * WSZ) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < WSZ * BC; i++) {\n int inner_w = i % WSZ;\n int inner_c = i / WSZ;\n\n x_s[(inner_w + nh * WSZ) * BC + inner_c] = 0.0f;\n } \n }\n // if loading the last width section, fill in the final value in the buffer to 0\n if (wsz == WSZ - 1) {\n for (int i = threadIdx.x; i < nh * BO; i++) {\n int inner_o = i % BO;\n int inner_h = i / BO;\n dout_s[(inner_h * WSZ + wsz) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < nh * BC; i++) {\n int inner_c = i % BC;\n int inner_h = i / BC;\n x_s[(inner_h * WSZ + wsz) * BC + inner_c] = 0.0f;\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int x_dot_idx = 0; x_dot_idx < WSZ * nh; x_dot_idx++) {\n int x_h = x_dot_idx / WSZ;\n int x_w = x_dot_idx % WSZ;\n // load TC of the C channels of one image pixel into reg_w\n for (int j = 0; j < TC; j++) {\n reg_x[j] = x_s[x_dot_idx * BC + thread_col * TC + j];\n }\n\n // for each conv filter, load a different row and calculate the outer product\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n // calculate the corresponding (h, w) that should be loaded from dout\n int dout_h = x_h + 1 - k_i;\n int dout_w = x_w + 1 - k_j;\n int dout_h_abs = h + dout_h;\n int dout_w_abs = w_offset + dout_w;\n int dout_dot_idx = -1;\n if (dout_h >= 0 && dout_h < NH && dout_w >= 0 && dout_w < WSZ && dout_w_abs < W && dout_h_abs < H) {\n if (x_h == NH - 1 && k_i == 1) {\n // this logic is a bit tricky: for the last row in smem\n // because it will be loaded as part of the next block as well, and it will be the first row in the next block\n // for kernels 4, 5, 6 we need to only compute them once\n // this statement ensures that we don't compute the 4, 5, 6 kernels for the last row in smem\n // so it will be computed as part of the first row next time\n dout_dot_idx = -1;\n } else if (x_w == WSZ - 1 && k_j == 1) {\n dout_dot_idx = -1;\n } else {\n dout_dot_idx = dout_h * WSZ + dout_w;\n }\n }\n\n // if in bounds of the dout section in smem, load the corresponding pixel from smem\n // and add the products to the thread results\n // otherwise the contribution is 0\n if (dout_dot_idx != -1) {\n for (int i = 0; i < TO; i++) {\n reg_dout[i] = dout_s[dout_dot_idx * BO + thread_row * TO + i];\n }\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n float new_val = reg_dout[i] * reg_x[j];\n thread_results[k * TO * TC + i * TC + j] += new_val;\n }\n }\n }\n }\n __syncthreads();\n }\n }\n\n // write out results\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n for (int k = 0; k < 9; k++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n int c_abs = block_col * BC + thread_col * TC + j;\n dweight_buf[o_abs * C * 9 * B * blocks_per_width + c_abs * 9 * B * blocks_per_width + k * B * blocks_per_width + batch * blocks_per_width + block_w] = thread_results[k * TO * TC + i * TC + j];\n }\n }\n }\n}\n\nvoid dweight_backward5(\n float* dout, float* x,\n float* dweight_buf, float* dweight,\n const int B, const int C, const int O, const int H, const int W\n) {\n const int BC = 64;\n const int BO = 64;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_backward_kernel5<<>>(dout, x, dweight_buf, B, C, O, H, W);\n cudaCheck(hipDeviceSynchronize());\n cudaCheck(hipGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n}\n\n// same as dweight backward 4, except we also do some of the reduce for dbias\n// so it doesn't have to reduce an entire BxHxW length for each O channel\ntemplate \n__global__ void dweight_dbias_backward_kernel1(\n float* dout, float* x,\n float* dweight_buf, float* dbias_buf,\n const int B, const int C, const int O, const int H, const int W\n) {\n // number of heights we need to load into smem in one go \n const int NH = BW / WSZ;\n // n_threads per block = (BC * BO) / (TO * TC)\n // gridDim will be (B * (W / (WSZ - 1)) * (C / BC), O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_chan = C / BC;\n const int blocks_per_width = W / (WSZ - 1);\n const int block_w = blockIdx.x % blocks_per_width;\n const int block_col = (blockIdx.x / blocks_per_width) % blocks_per_chan;\n const int batch = blockIdx.x / (blocks_per_chan * blocks_per_width);\n const int w_offset = block_w * (WSZ - 1);\n\n // warp indices for loading from global mem\n const int n_warps = ceil_div((int)blockDim.x, 32);\n const int warp_id = threadIdx.x / 32;\n const int lane_id = threadIdx.x % 32;\n // make sure that the width sections can be loaded by a single warp\n assert(WSZ <= 32);\n\n // index for the TO*TC thread block\n const int threads_per_col = BC / TC;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // smem for dout and x\n __shared__ float dout_s[BO * BW];\n __shared__ float x_s[BC * BW];\n\n // registers for local mem and local results\n float thread_results[TO * TC * 9] = {0.0};\n float reg_dout[TO] = {0.0};\n float reg_x[TC] = {0.0};\n float thread_bias_results[TO] = {0.0};\n\n // block offsets\n const int dout_block_offset = batch * O * H * W + block_row * BO * H * W;\n const int x_block_offset = batch * C * H * W + block_col * BC * H * W;\n\n //int debug_idx = 0;\n\n // we will fuse the calculation for all 9 kernel filters\n // which means we will have to reload some of the image widths\n // we can load NW blocks into smem in one go, which means we will be stepping\n // over the H dimension in (NW - 1) steps\n // this is the outer loop over H\n int h_step = (NH < H) ? NH - 1 : H;\n // if NH == H, then we are loading the entire image in one go, so step by H\n for (int h = 0; h < H; h += h_step) {\n // load dout into smem\n int h_offset = h * W;\n int nh = (h < H - (NH - 1)) ? NH : NH - 1;\n int wsz = (w_offset < W - (WSZ - 1)) ? WSZ : WSZ - 1;\n int o_per_warp = ceil_div(BO, n_warps);\n if (lane_id < wsz) {\n for (int warp_o = 0; warp_o < o_per_warp; warp_o++) {\n int inner_o = warp_id * o_per_warp + warp_o;\n if (inner_o < BO) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n dout_s[(inner_h * WSZ + inner_w) * BO + inner_o] = \n dout[dout_block_offset + inner_o * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n int c_per_warp = ceil_div(BC, n_warps);\n if (lane_id < wsz) {\n for (int warp_c = 0; warp_c < c_per_warp; warp_c++) {\n int inner_c = warp_id * c_per_warp + warp_c;\n if (inner_c < BC) {\n for (int inner_h = 0; inner_h < nh; inner_h++) {\n int inner_w = lane_id;\n x_s[(inner_h * WSZ + inner_w) * BC + inner_c] = \n x[x_block_offset + inner_c * H * W + h_offset + inner_h * W + w_offset + inner_w];\n }\n }\n }\n }\n\n __syncthreads();\n \n\n // if loading the last block, fill in the last block's buffer to be 0\n if (nh == NH - 1) {\n for (int i = threadIdx.x; i < WSZ * BO; i++) {\n int inner_w = i % WSZ;\n int inner_o = i / WSZ;\n dout_s[(inner_w + nh * WSZ) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < WSZ * BC; i++) {\n int inner_w = i % WSZ;\n int inner_c = i / WSZ;\n\n x_s[(inner_w + nh * WSZ) * BC + inner_c] = 0.0f;\n } \n }\n // if loading the last width section, fill in the final value in the buffer to 0\n if (wsz == WSZ - 1) {\n for (int i = threadIdx.x; i < nh * BO; i++) {\n int inner_o = i % BO;\n int inner_h = i / BO;\n dout_s[(inner_h * WSZ + wsz) * BO + inner_o] = 0.0f;\n }\n for (int i = threadIdx.x; i < nh * BC; i++) {\n int inner_c = i % BC;\n int inner_h = i / BC;\n x_s[(inner_h * WSZ + wsz) * BC + inner_c] = 0.0f;\n }\n }\n __syncthreads();\n\n // calculate per thread results\n for (int x_dot_idx = 0; x_dot_idx < WSZ * nh; x_dot_idx++) {\n int x_h = x_dot_idx / WSZ;\n int x_w = x_dot_idx % WSZ;\n // load TC of the C channels of one image pixel into reg_w\n for (int j = 0; j < TC; j++) {\n reg_x[j] = x_s[x_dot_idx * BC + thread_col * TC + j];\n }\n\n // for each conv filter, load a different row and calculate the outer product\n for (int k = 0; k < 9; k++) {\n int k_i = k / 3;\n int k_j = k % 3;\n // calculate the corresponding (h, w) that should be loaded from dout\n int dout_h = x_h + 1 - k_i;\n int dout_w = x_w + 1 - k_j;\n int dout_h_abs = h + dout_h;\n int dout_w_abs = w_offset + dout_w;\n int dout_dot_idx = -1;\n if (dout_h >= 0 && dout_h < NH && dout_w >= 0 && dout_w < WSZ && dout_w_abs < W && dout_h_abs < H) {\n if (x_h == NH - 1 && k_i == 1 && H > NH) {\n // this logic is a bit tricky: for the last row in smem\n // because it will be loaded as part of the next block as well, and it will be the first row in the next block\n // for kernels 4, 5, 6 we need to only compute them once\n // this statement ensures that we don't compute the 4, 5, 6 kernels for the last row in smem\n // so it will be computed as part of the first row next time\n dout_dot_idx = -1;\n } else if (x_w == WSZ - 1 && k_j == 1 && W > WSZ) {\n // if x_w is at the last pixel, and we have more than one width sections, and we are not are the last section\n // then don't compute 1, 4, 7 kernels and wait for the next section to compute them\n dout_dot_idx = -1;\n } else {\n dout_dot_idx = dout_h * WSZ + dout_w;\n }\n }\n\n // if in bounds of the dout section in smem, load the corresponding pixel from smem\n // and add the products to the thread results\n // otherwise the contribution is 0\n if (dout_dot_idx != -1) {\n for (int i = 0; i < TO; i++) {\n reg_dout[i] = dout_s[dout_dot_idx * BO + thread_row * TO + i];\n }\n for (int j = 0; j < TC; j++) {\n for (int i = 0; i < TO; i++) {\n float new_val = reg_dout[i] * reg_x[j];\n thread_results[k * TO * TC + i * TC + j] += new_val;\n }\n }\n // update dbias when k = 4, where dout_h and dout_w match x_h and x_w\n if (k == 4 && block_col == 0) {\n for (int i = 0; i < TO; i++) {\n thread_bias_results[i] += reg_dout[i];\n }\n }\n }\n }\n __syncthreads();\n }\n }\n\n // write out dweight results\n for (int i = 0; i < TO; i++) {\n int o_abs = block_row * BO + thread_row * TO + i;\n for (int j = 0; j < TC; j++) {\n int c_abs = block_col * BC + thread_col * TC + j;\n for (int k = 0; k < 9; k++) {\n dweight_buf[o_abs * C * 9 * B * blocks_per_width + c_abs * 9 * B * blocks_per_width + k * B * blocks_per_width + batch * blocks_per_width + block_w] = thread_results[k * TO * TC + i * TC + j];\n }\n }\n // write out dbias results\n if (block_col == 0) {\n dbias_buf[o_abs * B * blocks_per_width + batch * blocks_per_width + block_w] = thread_bias_results[i];\n }\n }\n}\n\n\n__global__ void dbias_reduce_kernel1(\n float* dbias_buf, float* dbias,\n const int B\n) {\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n\n dbias_buf += blockIdx.x * B;\n\n int lane_id = threadIdx.x % 32;\n\n float db_thread_sum = 0.0f;\n for (int i = lane_id; i < B; i += 32) {\n db_thread_sum += dbias_buf[i];\n }\n\n float db_warp_sum = cg::reduce(warp, db_thread_sum, cg::plus{});\n\n if (lane_id == 0) {\n dbias[blockIdx.x] = db_warp_sum;\n }\n}\n\nvoid dweight_dbias_backward1(\n float* dout, float* x,\n float* dweight_buf, float* dweight,\n float* dbias_buf, float* dbias,\n const int B, const int C, const int O, const int H, const int W\n) {\n if (H * W == 64) {\n const int BC = 64;\n const int BO = 64;\n const int WSZ = 8;\n const int BW = WSZ * 8;\n const int TO = 4;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(hipDeviceSynchronize());\n cudaCheck(hipGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B, C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B);\n } else if (C == 3) {\n const int BC = 3;\n const int BO = 64;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 3;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(hipDeviceSynchronize());\n cudaCheck(hipGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B * W / (WSZ - 1));\n } else if (O == 3) {\n const int BC = 64;\n const int BO = 3;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 3;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(hipDeviceSynchronize());\n cudaCheck(hipGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B * W / (WSZ - 1));\n } else {\n const int BC = 64;\n const int BO = 64;\n const int WSZ = 8 + 1;\n const int BW = WSZ * 9;\n const int TO = 4;\n const int TC = 4;\n assert(BW % WSZ == 0);\n dim3 gridDim(B * (W / (WSZ - 1)) * (C / BC), O / BO);\n dim3 blockDim(BC * BO / (TC * TO));\n dweight_dbias_backward_kernel1<<>>(dout, x, dweight_buf, dbias_buf, B, C, O, H, W);\n cudaCheck(hipDeviceSynchronize());\n cudaCheck(hipGetLastError());\n\n int block_size = 32 * 9;\n int n_blocks = O * C;\n dweight_reduce_kernel<<>>(dweight_buf, dweight, B * W / (WSZ - 1), C, O);\n dbias_reduce_kernel1<<>>(dbias_buf, dbias, B * W / (WSZ - 1));\n }\n}\n\nvoid conv2d_k3_backward2(\n float* dout, float* x, float* weight,\n float* dweight_buf, float* dbias_buf,\n float* dx, float* dweight, float* dbias,\n const int B, const int C_in, const int C_out, const int H, const int W\n) {\n // dx currently requires the output buffer to be zeroed\n cudaCheck(hipMemset(dx, 0, B * C_in * H * W * sizeof(float)));\n dx_backward(dout, weight, dx, B, C_in, C_out, H, W);\n dweight_dbias_backward1(dout, x, dweight_buf, dweight, dbias_buf, dbias, B, C_in, C_out, H, W);\n}\n// ----------------------------------------------------------------------------\n// kernel launcher\nvoid conv2d_k3_forward(\n hipblasHandle_t cublas_handle,\n int kernel_num,\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C_in, int C_out, int H, int W,\n int block_size\n) {\n switch (kernel_num) {\n case 1:\n conv2d_k3_forward1(cublas_handle, x, weight, bias, out, B, C_in, C_out, H, W, block_size);\n break;\n case 2:\n conv2d_k3_forward2(cublas_handle, x, weight, bias, out, B, C_in, C_out, H, W, block_size);\n break;\n default:\n printf(\"Invalid kernel number\\n\");\n break;\n }\n}\n\n\n#ifndef LINKING\nint main( int argc, char **argv) {\n const int B = 32;\n const int C_in = 192;\n const int C_out = 64;\n const int H = 64;\n const int W = 64;\n\n // set up the device\n int deviceIdx = 0;\n cudaCheck(hipSetDevice(deviceIdx));\n hipDeviceProp_t deviceProp;\n hipGetDeviceProperties(&deviceProp, deviceIdx);\n printf(\"Device %d: %s\\n\", deviceIdx, deviceProp.name);\n printf(\"device registers per block: %d\\n\", deviceProp.regsPerBlock);\n printf(\"shared mem per block: %ld\\n\", deviceProp.sharedMemPerBlock);\n\n // setup cublas\n hipblasHandle_t cublas_handle;\n cublasCheck(hipblasCreate(&cublas_handle));\n cublasCheck(hipblasSetMathMode(cublas_handle, HIPBLAS_DEFAULT_MATH));\n\n // create host memory to load data\n float* x = (float*)malloc(B * C_in * H * W * sizeof(float));\n float* weight = (float*)malloc(C_out * C_in * 3 * 3 * sizeof(float));\n float* bias = (float*)malloc(C_out * sizeof(float));\n float* out = (float*)malloc(B * C_out * H * W * sizeof(float));\n float* dout = (float*)malloc(B * C_out * H * W * sizeof(float));\n float* dx = (float*)malloc(B * C_in * H * W * sizeof(float));\n float* dweight = (float*)malloc(C_out * C_in * 3 * 3 * sizeof(float));\n float* dbias = (float*)malloc(C_out * sizeof(float));\n // debug\n float* dout_perm = (float*)malloc(B * C_out * H * W * sizeof(float));\n float* x_tiled = (float*)malloc(B * H * W * C_in * 9 * sizeof(float));\n\n // read saved output\n FILE *file = fopen(\"conv2d_k3.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(x, sizeof(float), B * C_in * H * W, file);\n freadCheck(weight, sizeof(float), C_out * C_in * 3 * 3, file);\n freadCheck(bias, sizeof(float), C_out, file);\n freadCheck(out, sizeof(float), B * C_out * H * W, file);\n freadCheck(dout, sizeof(float), B * C_out * H * W, file);\n freadCheck(dx, sizeof(float), B * C_in * H * W, file);\n freadCheck(dweight, sizeof(float), C_out * C_in * 3 * 3, file);\n freadCheck(dbias, sizeof(float), C_out, file);\n freadCheck(dout_perm, sizeof(float), B * C_out * H * W, file);\n freadCheck(x_tiled, sizeof(float), B * H * W * C_in * 9, file);\n fclose(file);\n\n // allocate device memory\n float *d_x, *d_weight, *d_bias, *d_out, *d_dout, *d_dx, *d_dweight, *d_dbias;\n cudaCheck(hipMalloc(&d_x, B * C_in * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&d_weight, C_out * C_in * 3 * 3 * sizeof(float)));\n cudaCheck(hipMalloc(&d_bias, C_out * sizeof(float)));\n cudaCheck(hipMalloc(&d_out, B * C_out * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&d_dout, B * C_out * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&d_dx, B * C_in * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&d_dweight, C_out * C_in * 3 * 3 * sizeof(float)));\n cudaCheck(hipMalloc(&d_dbias, C_out * sizeof(float)));\n\n cudaCheck(hipMemset(d_out, 0, B * C_out * H * W * sizeof(float)));\n cudaCheck(hipMemcpy(d_x, x, B * C_in * H * W * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_weight, weight, C_out * C_in * 3 * 3 * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_bias, bias, C_out * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_dout, dout, B * C_out * H * W * sizeof(float), hipMemcpyHostToDevice));\n\n // read kernel number from command line\n int kernel_num = 2;\n if (argc > 1) {\n kernel_num = atoi(argv[1]);\n }\n printf(\"Running kernel %d\\n\", kernel_num);\n\n printf(\"Checking forward pass\\n\");\n float forward_acc = 1e-1;\n\n //for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n // int block_size = block_sizes[j];\n // printf(\"\\nBlock size: %d\\n\", block_size);\n // conv2d_k3_forward(cublas_handle, kernel_num, d_x, d_weight, d_bias, d_out, B, C_in, C_out, H, W, block_size);\n // validate_result(d_out, out, \"out\", B * C_out * H * W, forward_acc);\n //}\n //printf(\"\\nChecking cpu forward\\n\");\n //float* out_cpu_buf = (float*)malloc(B * C_out * H * W * sizeof(float));\n //conv_k3_forward_cpu(x, weight, bias, out_cpu_buf, B, C_in, C_out, H, W);\n //float* out_gpu_buf;\n //cudaCheck(cudaMalloc(&out_gpu_buf, B * C_out * H * W * sizeof(float)));\n //cudaCheck(cudaMemcpy(out_gpu_buf, out, B * C_out * H * W * sizeof(float), cudaMemcpyHostToDevice));\n //validate_result(out_gpu_buf, out, \"out cpu\", B * C_out * H * W);\n //cudaCheck(cudaFree(out_gpu_buf));\n //free(out_cpu_buf);\n\n printf(\"\\nChecking new kernel\\n\");\n conv2d_k3_forward3(d_x, d_weight, d_bias, d_out, B, C_in, C_out, H, W);\n cudaCheck(hipGetLastError());\n printf(\"\\nChecking new kernel results\\n\");\n validate_result(d_out, out, \"out\", B * C_out * H * W, forward_acc);\n printf(\"Forward pass successful\\n\\n\");\n\n printf(\"\\nmanually benchmark new forward kernel\\n\");\n void* flush_buffer;\n cudaCheck(hipMalloc(&flush_buffer, deviceProp.l2CacheSize));\n int repeat_times = 1000;\n float new_time = 0.0f;\n for (int i = 0; i < repeat_times; i++) {\n // clear L2\n cudaCheck(hipMemset(flush_buffer, 0, deviceProp.l2CacheSize));\n hipEvent_t start, end;\n cudaCheck(hipEventCreate(&start));\n cudaCheck(hipEventCreate(&end));\n cudaCheck(hipEventRecord(start, nullptr));\n conv2d_k3_forward3(d_x, d_weight, d_bias, d_out, B, C_in, C_out, H, W);\n cudaCheck(hipEventRecord(end, nullptr));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n new_time += elapsed_time;\n }\n new_time /= repeat_times;\n printf(\"new kernel time: %.4f ms\\n\", new_time);\n\n printf(\"Checking backward pass\\n\");\n\n float dw_acc = 1e-2;\n float* d_dweight_buf;\n cudaCheck(hipMalloc(&d_dweight_buf, C_out * C_in * 3 * 3 * B * 32 * sizeof(float)));\n float* d_dbias_buf;\n cudaCheck(hipMalloc(&d_dbias_buf, C_out * 256 * sizeof(float)));\n conv2d_k3_backward2(d_dout, d_x, d_weight, d_dweight_buf, d_dbias_buf, d_dx, d_dweight, d_dbias, B, C_in, C_out, H, W);\n\n printf(\"\\nChecking dweight\\n\");\n validate_result(d_dweight, dweight, \"dweight\", C_out * C_in * 3 * 3, dw_acc);\n\n printf(\"\\nChecking dbias\\n\");\n validate_result(d_dbias, dbias, \"dbias\", C_out, dw_acc);\n\n printf(\"\\nChecking dx\\n\");\n validate_result(d_dx, dx, \"dx\", B * C_in * H * W, dw_acc);\n\n new_time = 0.0f;\n printf(\"\\nManually benchmarking new dweight kernel\\n\");\n for (int i = 0; i < repeat_times; i++) {\n // clear L2\n cudaCheck(hipMemset(flush_buffer, 0, deviceProp.l2CacheSize));\n hipEvent_t start, end;\n cudaCheck(hipEventCreate(&start));\n cudaCheck(hipEventCreate(&end));\n cudaCheck(hipEventRecord(start, nullptr));\n conv2d_k3_backward2(d_dout, d_x, d_weight, d_dweight_buf, d_dbias_buf, d_dx, d_dweight, d_dbias, B, C_in, C_out, H, W);\n cudaCheck(hipEventRecord(end, nullptr));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n new_time += elapsed_time;\n }\n new_time /= repeat_times;\n printf(\"new kernel time: %.4f ms\\n\", new_time);\n cudaCheck(hipFree(d_dweight_buf));\n cudaCheck(hipFree(d_dbias_buf));\n\n\n\n printf(\"Backward pass successful\\n\");\n\n // old benchmarking code\n //printf(\"All results match. Starting benchmarks.\\n\\n\");\n //printf(\"Forward pass benchmarks:\\n\");\n //for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n // int block_size = block_sizes[j];\n\n // int repeat_times = 100;\n // float elapsed_time = benchmark_kernel(repeat_times, conv2d_k3_forward,\n // cublas_handle, kernel_num, d_x, d_weight, d_bias, d_out,\n // B, C_in, C_out, H, W, block_size);\n\n // float tflops = (float)B * H * W * C_in * C_out * 9 * 2 / elapsed_time * 1e3f / 1e12f;\n // printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n //}\n printf(\"\\n\\n Old benchmarks\\n\");\n printf(\"Running our own forward benchmarks\\n\");\n float t1 = 0.0f;\n float t2 = 0.0f;\n float t3 = 0.0f;\n float t_total = 0.0f;\n for (int i = 0; i < repeat_times; i++) {\n // clear L2\n cudaCheck(hipMemset(flush_buffer, 0, deviceProp.l2CacheSize));\n conv2d_k3_forward2(cublas_handle, d_x, d_weight, d_bias, d_out, B, C_in, C_out, H, W, 512, &t1, &t2, &t3);\n }\n t1 /= repeat_times;\n t2 /= repeat_times;\n t3 /= repeat_times;\n t_total = t1 + t2 + t3;\n\n printf(\"t1: %.4f ms\\n\", t1);\n printf(\"t2: %.4f ms\\n\", t2);\n printf(\"t3: %.4f ms\\n\", t3);\n printf(\"t_total: %.4f ms\\n\", t_total);\n\n //printf(\"\\nBackward pass benchmarks:\\n\");\n //for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n // int block_size = block_sizes[j];\n\n // int repeat_times = 100;\n // float elapsed_time = benchmark_kernel(repeat_times, conv2d_k3_backward1,\n // cublas_handle, d_dout, d_x, d_weight, d_dx, d_dweight, d_dbias,\n // B, C_in, C_out, H, W, block_size);\n\n // float tflops = (float)B * H * W * C_in * C_out * (9 + 1) * 2 / elapsed_time * 1e3f / 1e12f;\n // printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n //}\n printf(\"\\nRunning our own backward benchmarks\\n\");\n t1 = 0.0f;\n t2 = 0.0f;\n t3 = 0.0f;\n float t4 = 0.0f;\n float t5 = 0.0f;\n float t6 = 0.0f;\n t_total = 0.0f;\n for (int i = 0; i < repeat_times; i++) {\n // clear L2\n cudaCheck(hipMemset(flush_buffer, 0, deviceProp.l2CacheSize));\n conv2d_k3_backward1(cublas_handle, d_dout, d_x, d_weight, d_dx, d_dweight, d_dbias, B, C_in, C_out, H, W, 512, &t1, &t2, &t3, &t4, &t5, &t6);\n }\n cudaCheck(hipFree(flush_buffer));\n\n t1 /= repeat_times;\n t2 /= repeat_times;\n t3 /= repeat_times;\n t4 /= repeat_times;\n t5 /= repeat_times;\n t6 /= repeat_times;\n t_total = t1 + t2 + t3 + t4 + t5 + t6;\n\n printf(\"t1: %.4f ms\\n\", t1);\n printf(\"t2: %.4f ms\\n\", t2);\n printf(\"t3: %.4f ms\\n\", t3);\n printf(\"t4: %.4f ms\\n\", t4);\n printf(\"t5: %.4f ms\\n\", t5);\n printf(\"t6: %.4f ms\\n\", t6);\n printf(\"t_total: %.4f ms\\n\", t_total);\n\n // free memory\n free(x);\n free(weight);\n free(bias);\n free(out);\n free(dout);\n free(dx);\n free(dweight);\n free(dbias);\n free(dout_perm);\n free(x_tiled);\n cudaCheck(hipFree(d_x));\n cudaCheck(hipFree(d_weight));\n cudaCheck(hipFree(d_bias));\n cudaCheck(hipFree(d_out));\n cudaCheck(hipFree(d_dout));\n cudaCheck(hipFree(d_dx));\n cudaCheck(hipFree(d_dweight));\n cudaCheck(hipFree(d_dbias));\n cublasCheck(hipblasDestroy(cublas_handle));\n return 0;\n}\n#endif", "source": "/kaggle/working/unet.cu/dev/conv2d_k3.cu"} {"input": "#include \n#include \"common.h\"\n#include \n#include \n#include \"timestep_embedding.cuh\"\n\n\n\n\n__global__ void fill_freqs_kernel(float* freqs, int half_dim, int max_period) {\n int idx = threadIdx.x + blockIdx.x * blockDim.x;\n if (idx < half_dim) {\n float log_period = logf(max_period);\n float logit = -log_period * idx / half_dim;\n freqs[idx] = expf(logit);\n }\n}\n\nvoid init_timestep_embedding(TimestepEmbedding* emb, int dim, int B, int max_period) {\n assert(dim % 2 == 0);\n int half = dim / 2;\n emb->half_dim = half;\n emb->max_period = max_period;\n emb->B = B;\n cudaCheck(cudaMalloc(&emb->freqs, half * sizeof(float)));\n // fill freq values\n int block_size = 512;\n int n_blocks = ceil_div(half, block_size);\n fill_freqs_kernel<<>>(emb->freqs, half, max_period);\n}\n\nvoid free_timestep_embedding(TimestepEmbedding* emb) {\n cudaCheck(cudaFree(emb->freqs));\n}\n\n__global__ void fill_embeddings_kernel(\n const float* timesteps, const float* freqs, float* embeddings, int half, int B\n) {\n int idx = threadIdx.x + blockIdx.x * blockDim.x;\n if (idx < B * half) {\n int i = idx / half;\n int j = idx % half;\n float arg = timesteps[i] * freqs[j];\n embeddings[i * 2 * half + j] = cosf(arg);\n embeddings[i * 2 * half + j + half] = sinf(arg);\n }\n}\n\n\n// timesteps is a (B,) array of timesteps that are integers but saved as floats\n// out is a (B, dim) array\nvoid get_timestep_embeddings(\n TimestepEmbedding* emb, const float* timesteps, float* out\n) {\n int block_size = 512;\n int n_blocks = ceil_div(emb->half_dim * emb->B, block_size);\n fill_embeddings_kernel<<>>(timesteps, emb->freqs, out, emb->half_dim, emb->B);\n cudaCheck(cudaGetLastError());\n}\n\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n int B = 32;\n int dim = 64;\n int max_period = 1000;\n\n // read ground truth\n float* timesteps = (float*)malloc(B * sizeof(float));\n float* embeddings = (float*)malloc(B * dim * sizeof(float));\n float* freqs = (float*)malloc(dim / 2 * sizeof(float));\n FILE *file = fopenCheck(\"time_emb.bin\", \"rb\");\n freadCheck(timesteps, sizeof(float), B, file);\n freadCheck(embeddings, sizeof(float), B * dim, file);\n freadCheck(freqs, sizeof(float), dim / 2, file);\n fcloseCheck(file);\n\n TimestepEmbedding t_emb;\n init_timestep_embedding(&t_emb, dim, B, max_period);\n float *d_timesteps, *d_embeddings;\n cudaCheck(cudaMalloc(&d_timesteps, B * sizeof(float)));\n cudaCheck(cudaMalloc(&d_embeddings, B * dim * sizeof(float)));\n\n cudaCheck(cudaMemcpy(d_timesteps, timesteps, B * sizeof(float), cudaMemcpyHostToDevice));\n\n\n printf(\"Checking freqs\\n\");\n validate_result(t_emb.freqs, freqs, \"freqs\", dim / 2);\n\n\n get_timestep_embeddings(&t_emb, d_timesteps, d_embeddings);\n\n printf(\"Checking timestep embeddings\\n\");\n float emb_acc = 1e-3;\n validate_result(d_embeddings, embeddings, \"embeddings\", B * dim, emb_acc);\n\n free_timestep_embedding(&t_emb);\n free(timesteps);\n free(embeddings);\n free(freqs);\n cudaFree(d_timesteps);\n}\n#endif", "output": "#include \n#include \"common.h\"\n#include \n#include \n#include \"timestep_embedding.cuh\"\n\n\n\n\n__global__ void fill_freqs_kernel(float* freqs, int half_dim, int max_period) {\n int idx = threadIdx.x + blockIdx.x * blockDim.x;\n if (idx < half_dim) {\n float log_period = logf(max_period);\n float logit = -log_period * idx / half_dim;\n freqs[idx] = expf(logit);\n }\n}\n\nvoid init_timestep_embedding(TimestepEmbedding* emb, int dim, int B, int max_period) {\n assert(dim % 2 == 0);\n int half = dim / 2;\n emb->half_dim = half;\n emb->max_period = max_period;\n emb->B = B;\n cudaCheck(hipMalloc(&emb->freqs, half * sizeof(float)));\n // fill freq values\n int block_size = 512;\n int n_blocks = ceil_div(half, block_size);\n fill_freqs_kernel<<>>(emb->freqs, half, max_period);\n}\n\nvoid free_timestep_embedding(TimestepEmbedding* emb) {\n cudaCheck(hipFree(emb->freqs));\n}\n\n__global__ void fill_embeddings_kernel(\n const float* timesteps, const float* freqs, float* embeddings, int half, int B\n) {\n int idx = threadIdx.x + blockIdx.x * blockDim.x;\n if (idx < B * half) {\n int i = idx / half;\n int j = idx % half;\n float arg = timesteps[i] * freqs[j];\n embeddings[i * 2 * half + j] = cosf(arg);\n embeddings[i * 2 * half + j + half] = sinf(arg);\n }\n}\n\n\n// timesteps is a (B,) array of timesteps that are integers but saved as floats\n// out is a (B, dim) array\nvoid get_timestep_embeddings(\n TimestepEmbedding* emb, const float* timesteps, float* out\n) {\n int block_size = 512;\n int n_blocks = ceil_div(emb->half_dim * emb->B, block_size);\n fill_embeddings_kernel<<>>(timesteps, emb->freqs, out, emb->half_dim, emb->B);\n cudaCheck(hipGetLastError());\n}\n\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n int B = 32;\n int dim = 64;\n int max_period = 1000;\n\n // read ground truth\n float* timesteps = (float*)malloc(B * sizeof(float));\n float* embeddings = (float*)malloc(B * dim * sizeof(float));\n float* freqs = (float*)malloc(dim / 2 * sizeof(float));\n FILE *file = fopenCheck(\"time_emb.bin\", \"rb\");\n freadCheck(timesteps, sizeof(float), B, file);\n freadCheck(embeddings, sizeof(float), B * dim, file);\n freadCheck(freqs, sizeof(float), dim / 2, file);\n fcloseCheck(file);\n\n TimestepEmbedding t_emb;\n init_timestep_embedding(&t_emb, dim, B, max_period);\n float *d_timesteps, *d_embeddings;\n cudaCheck(hipMalloc(&d_timesteps, B * sizeof(float)));\n cudaCheck(hipMalloc(&d_embeddings, B * dim * sizeof(float)));\n\n cudaCheck(hipMemcpy(d_timesteps, timesteps, B * sizeof(float), hipMemcpyHostToDevice));\n\n\n printf(\"Checking freqs\\n\");\n validate_result(t_emb.freqs, freqs, \"freqs\", dim / 2);\n\n\n get_timestep_embeddings(&t_emb, d_timesteps, d_embeddings);\n\n printf(\"Checking timestep embeddings\\n\");\n float emb_acc = 1e-3;\n validate_result(d_embeddings, embeddings, \"embeddings\", B * dim, emb_acc);\n\n free_timestep_embedding(&t_emb);\n free(timesteps);\n free(embeddings);\n free(freqs);\n hipFree(d_timesteps);\n}\n#endif", "source": "/kaggle/working/unet.cu/dev/timestep_embedding.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n#include \"common.h\"\n\n\n__global__ void permute_kernel(float* q, float* k, float* v,\n const float* inp,\n int B, int N, int NH, int d) {\n // okay so now, this kernel wants Q,K,V to all be of shape (B, NH, N, d)\n // but instead, we have a single tensor QKV (inp) of shape (B, N, 3, NH, d)\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Q[b][nh_][n][d_] = inp[b][n][0][nh_][d_]\n\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n\n int inp_idx = \\\n (b * N * 3 * NH * d)\n + (n * 3 * NH * d)\n + (0 * NH * d)\n + (nh_ * d)\n + d_;\n\n q[idx] = inp[inp_idx];\n k[idx] = inp[inp_idx + NH * d];\n v[idx] = inp[inp_idx + 2 * (NH * d)];\n }\n}\n\n__global__ void scale_kernel(float* inp, float scale, int B, int NH, int T) {\n // scales the pre-softmax attention scores by scale\n // we don't want causal attention, so we don't set any values to infinity\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * NH * T * T) {\n inp[idx] *= scale;\n }\n}\n\n// warp-level reduction for finding the maximum value\n__device__ float warpReduceMax(float val) {\n for (int offset = 16; offset > 0; offset /= 2) {\n val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));\n }\n return val;\n}\n\n// warp-level reduction for summing values\n__device__ float warpReduceSum(float val) {\n for (int offset = 16; offset > 0; offset /= 2) {\n val += __shfl_down_sync(0xFFFFFFFF, val, offset);\n }\n return val;\n}\n\n__global__ void softmax_forward_kernel4(float* out, const float* inp, int N, int C) {\n // out is (N, C) just like inp. Each row of inp will get softmaxed.\n // same as kernel3, but can handle any block size (multiple of 32)\n // each row of C elements is handled by block_size threads\n // furthermore, each block_size threads get executed in warps of 32 threads\n\n // special reduction operations warpReduceMax/warpReduceSum are used for intra-warp reductions\n // shared memory is used for inter-warp reduction\n extern __shared__ float shared[];\n int idx = blockIdx.x;\n int tid = threadIdx.x;\n int warpId = threadIdx.x / 32; // warp index within a block\n int laneId = threadIdx.x % 32; // thread index within a warp\n\n // the number of warps per block. recall that blockDim.x is block_size\n int warpsPerBlock = blockDim.x / 32;\n\n // shared[] must be allocated to have 2 * warpsPerBlock elements\n // first half for max values, the second half for sum values\n float* maxvals = shared;\n float* sumvals = &shared[warpsPerBlock];\n\n // one row of inp, i.e. inp[idx, :] of shape (C,)\n const float* x = inp + idx * C;\n\n // first, thread coarsening by directly accessing global memory in series\n float maxval = -INFINITY;\n for (int i = tid; i < C; i += blockDim.x) {\n maxval = fmaxf(maxval, x[i]);\n }\n // now within-warp reductions for maxval\n maxval = warpReduceMax(maxval);\n\n // the 0th thread of each warp writes the maxval of that warp to shared memory\n if (laneId == 0) maxvals[warpId] = maxval;\n __syncthreads();\n\n // now the 0th thread reduces the maxvals in shared memory, i.e. across warps\n if (tid == 0) {\n float val = maxvals[tid];\n for (int i = 1; i < warpsPerBlock; i++) {\n val = fmaxf(val, maxvals[i]);\n }\n // store the final max in the first position\n maxvals[0] = val;\n }\n __syncthreads();\n // broadcast the max to all threads\n float offset = maxvals[0];\n\n // compute expf and write the result to global memory\n for (int i = tid; i < C; i += blockDim.x) {\n // subtract max for numerical stability\n out[idx * C + i] = expf(x[i] - offset);\n }\n\n // okay now we calculated exp(x - max(x))\n // step 2: sum all the values and divide by the sum\n\n // thread coarsening for sum\n x = out + idx * C;\n float sumval = 0.0f;\n for (int i = tid; i < C; i += blockDim.x) {\n sumval += x[i];\n }\n // within-warp reduction for sumval\n sumval = warpReduceSum(sumval);\n\n // write sumval to shared memory\n if (laneId == 0) sumvals[warpId] = sumval;\n __syncthreads();\n\n // inter-thread reduction of sum\n if (tid == 0) {\n float val = sumvals[tid];\n for (int i = 1; i < warpsPerBlock; ++i) {\n val += sumvals[i];\n }\n sumvals[0] = val;\n }\n __syncthreads();\n // broadcast the sum to all threads\n float sum = sumvals[0];\n\n // divide the whole row by the sum\n for (int i = tid; i < C; i += blockDim.x) {\n out[idx * C + i] = x[i] / sum;\n }\n}\n\n__global__ void unpermute_kernel(const float* inp, float *out, int B, int N, int NH, int d) {\n // out has shape (B, nh, N, d) but we need to unpermute it to (B, N, nh, d)\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n // out[b][n][nh_][d_] <- inp[b][nh_][n][d_]\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n\n int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_;\n out[other_idx] = inp[idx];\n }\n}\n\n// a simple attention forward function taken directly from llm.c kernel 3\nvoid attention_forward1\n(\n cublasHandle_t cublas_handle,\n float* out, float* qkvr, float* preatt, float* att,\n float* inp,\n int B, int T, int C, int NH,\n const int block_size\n) {\n // note: resure inp as a scratch buffer because it is not used in backward\n // inp is (B, T, 3C) QKV\n // preatt, att are (B, NH, T, T)\n // output is (B, T, C)\n int HS = C / NH; // head size\n\n // permute and separate inp from (B, T, 3, NH, HS) to 3X (B, NH, T, HS)\n float *q, *k, *v;\n q = qkvr + 0 * B * T * C;\n k = qkvr + 1 * B * T * C;\n v = qkvr + 2 * B * T * C;\n int total_threads = B * NH * T * HS;\n int num_blocks = ceil_div(total_threads, block_size);\n permute_kernel<<>>(q, k, v, inp, B, T, NH, HS);\n\n // batched matrix multiply with cuBLAS\n const float alpha = 1.0f;\n const float beta = 0.0f;\n cublasCheck(cublasSgemmStridedBatched(cublas_handle,\n CUBLAS_OP_T, CUBLAS_OP_N,\n T, T, HS,\n &alpha,\n k, HS, T * HS,\n q, HS, T * HS,\n &beta,\n preatt, T, T * T,\n B * NH));\n\n // multiply all elements of preatt elementwise by scale\n float scale = 1.0f / sqrtf(HS);\n total_threads = B * NH * T * T;\n num_blocks = ceil_div(total_threads, block_size);\n scale_kernel<<>>(preatt, scale, B, NH, T);\n\n // softmax. preatt is (B, NH, T, T) but we view it as (B * NH * T, T) and use the softmax kernel\n int softmax_block_size = 256;\n int grid_size = B * NH * T;\n size_t shared_mem_size = 2 * softmax_block_size / 32 * sizeof(float);\n softmax_forward_kernel4<<>>(att, preatt, B * NH * T, T);\n\n // new approach: first cuBLAS another batched matmul\n // y = att @ v # (B, nh, T, T) @ (B, nh, T, hs) -> (B, nh, T, hs)\n float* vaccum = inp;\n cublasCheck(cublasSgemmStridedBatched(cublas_handle,\n CUBLAS_OP_N, CUBLAS_OP_N,\n HS, T, T,\n &alpha,\n v, HS, T * HS,\n att, T, T * T,\n &beta,\n vaccum, HS, T * HS,\n B * NH));\n\n // now unpermute\n // y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side\n num_blocks = ceil_div(B * T * C, block_size);\n unpermute_kernel<<>>(vaccum, out, B, T, NH, HS);\n}\n\n__global__ void permute_kernel_backward(float* dinp,\n const float* dq, const float* dk, const float* dv,\n int B, int N, int NH, int d) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n\n int inp_idx = (b * N * 3 * NH * d) + (n * 3 * NH * d) + (0 * NH * d) + (nh_ * d) + d_;\n dinp[inp_idx] = dq[idx];\n dinp[inp_idx + NH * d] = dk[idx];\n dinp[inp_idx + 2 * (NH * d)] = dv[idx];\n }\n}\n\n__global__ void unpermute_kernel_backward(float* dinp, const float *dout, int B, int N, int NH, int d) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_;\n dinp[idx] = dout[other_idx];\n }\n}\n\n// slight edit from the softmax_autoregressive_backward_kernel from train_gpt2_fp32\n// our attention is no longer autoregressive\n__global__ void softmax_backward_kernel(float* dpreatt, const float* datt, const float* att,\n int B, int T, int C, float scale) {\n constexpr const int BlockSize = 256;\n constexpr int T_per_block = 4;\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n __shared__ float block_acc[32];\n\n int idx = blockIdx.y;\n // go through blocks in reverse order, so the slowest block starts first\n int t0 = T - 1 - T_per_block*blockIdx.x;\n\n att += idx * T * T;\n datt += idx * T * T;\n dpreatt += idx * T * T;\n\n if (warp.meta_group_rank() == 0) {\n block_acc[warp.thread_rank()] = 0;\n }\n\n for(int to = 0; to < T_per_block; ++to) {\n int t = t0 - to;\n if(t < 0) return;\n const float* att_bth = att + t * T;\n const float* datt_bth = datt + t * T;\n float* dpreatt_bth = dpreatt + t * T;\n\n float local_sum = 0.0f;\n for (int t2 = block.thread_rank(); t2 < T; t2 += BlockSize) {\n local_sum += att_bth[t2] * datt_bth[t2];\n }\n\n block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus{});\n block.sync();\n local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus{});\n\n for (int t3 = block.thread_rank(); t3 < T; t3 += BlockSize) {\n // don't touch the cache. Some parts will still be here from the previous loop, and\n // we want to exploit those.\n float acc = __ldcs(att_bth + t3) * (__ldcs(datt_bth + t3) - local_sum);\n __stcs(dpreatt_bth + t3, scale * acc);\n }\n }\n}\n\n// the sequence of transformations in this compound op is:\n// inp (B,T,3C) -> qkvr (B,T,3C) -> preatt (B,NH,T,T) -> att (B,NH,T,T) -> vaccum (B,T,C) -> out (B,T,C)\nvoid attention_backward\n(\n cublasHandle_t cublas_handle,\n float* dinp, float* dqkvr, float* dpreatt, float* datt, float* scratch,\n const float* dout,\n const float* qkvr, const float* att,\n int B, int T, int C, int NH\n) {\n const int block_size = 256;\n int HS = C / NH; // head size\n const float one = 1.0f;\n const float zero = 0.0f; // note beta = 1.0f so that we accumulate gradients (+=)\n // unpack convenience pointers into q, k, v\n const float *q, *k, *v;\n q = qkvr + 0 * B * T * C;\n k = qkvr + 1 * B * T * C;\n v = qkvr + 2 * B * T * C;\n float *dq, *dk, *dv;\n dq = dqkvr + 0 * B * T * C;\n dk = dqkvr + 1 * B * T * C;\n dv = dqkvr + 2 * B * T * C;\n // backward through the unpermute operation\n int num_blocks = ceil_div(B * T * C, block_size);\n unpermute_kernel_backward<<>>(scratch, dout, B, T, NH, HS);\n cudaCheck(cudaGetLastError());\n // backward into datt\n cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, &one, v, HS, T * HS, scratch, HS, T * HS, &zero, datt, T, T * T, B * NH));\n // backward into dv\n cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &one, scratch, HS, T * HS, att, T, T * T, &zero, dv, HS, T * HS, B * NH));\n // backward into preatt\n int hs = C / NH; // head size\n float scale = 1.0f / sqrtf(hs);\n softmax_backward_kernel<<>>(dpreatt, datt, att, B, T, C, scale);\n cudaCheck(cudaGetLastError());\n // backward into q\n cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, &one, k, HS, T * HS, dpreatt, T, T * T, &zero, dq, HS, T * HS, B * NH));\n // backward into k\n cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &one, q, HS, T * HS, dpreatt, T, T * T, &zero, dk, HS, T * HS, B * NH));\n // backward into inp\n num_blocks = ceil_div(B * NH * T * HS, block_size);\n permute_kernel_backward<<>>(dinp, dq, dk, dv, B, T, NH, HS);\n cudaCheck(cudaGetLastError());\n}\n\n\n#ifndef LINKING\nint main(int argc, char** argv) {\n int B = 4;\n int T = 1024;\n int C = 256;\n int HS = 32;\n int NH = C / HS;\n\n // setup cublas\n cublasHandle_t cublas_handle;\n cublasCheck(cublasCreate(&cublas_handle));\n cublasCheck(cublasSetMathMode(cublas_handle, CUBLAS_DEFAULT_MATH));\n\n // create host memory to load data\n float* qkv = (float*)malloc(B * T * 3 * C * sizeof(float));\n float* out = (float*)malloc(B * T * C * sizeof(float));\n float* dout = (float*)malloc(B * T * C * sizeof(float));\n float* dqkv = (float*)malloc(B * T * 3 * C * sizeof(float));\n\n // read saved output\n FILE * file = fopen(\"attention.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(qkv, sizeof(float), B * T * 3 * C, file);\n freadCheck(out, sizeof(float), B * T * C, file);\n freadCheck(dout, sizeof(float), B * T * C, file);\n freadCheck(dqkv, sizeof(float), B * T * 3 * C, file);\n fclose(file);\n\n // allocate device memory\n float *d_inp, *d_qkvr, *d_out, *d_preatt, *d_att, *d_dout, *d_dinp, *d_dqkvr, *d_dpreatt, *d_datt, *d_scratch;\n cudaCheck(cudaMalloc(&d_inp, B * T * 3 * C * sizeof(float)));\n cudaCheck(cudaMemcpy(d_inp, qkv, B * T * 3 * C * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMalloc(&d_qkvr, B * T * 3 * C * sizeof(float)));\n cudaCheck(cudaMemset(d_qkvr, 0, B * T * 3 * C * sizeof(float)));\n cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(float)));\n cudaCheck(cudaMemset(d_out, 0, B * T * C * sizeof(float)));\n cudaCheck(cudaMalloc(&d_preatt, B * NH * T * T * sizeof(float)));\n cudaCheck(cudaMemset(d_preatt, 0, B * NH * T * T * sizeof(float)));\n cudaCheck(cudaMalloc(&d_att, B * NH * T * T * sizeof(float)));\n cudaCheck(cudaMemset(d_att, 0, B * NH * T * T * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dout, B * T * C * sizeof(float)));\n cudaCheck(cudaMemcpy(d_dout, dout, B * T * C * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMalloc(&d_dinp, B * T * 3 * C * sizeof(float)));\n cudaCheck(cudaMemset(d_dinp, 0, B * T * 3 * C * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dqkvr, B * T * 3 * C * sizeof(float)));\n cudaCheck(cudaMemset(d_dqkvr, 0, B * T * 3 * C * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dpreatt, B * NH * T * T * sizeof(float)));\n cudaCheck(cudaMemset(d_dpreatt, 0, B * NH * T * T * sizeof(float)));\n cudaCheck(cudaMalloc(&d_datt, B * NH * T * T * sizeof(float)));\n cudaCheck(cudaMemset(d_datt, 0, B * NH * T * T * sizeof(float)));\n cudaCheck(cudaMalloc(&d_scratch, B * T * C * sizeof(float)));\n cudaCheck(cudaMemset(d_scratch, 0, B * T * C * sizeof(float)));\n\n // forward pass\n int block_sizes[] = {256, 512};\n printf(\"Checking forward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"\\nBlock size: %d\\n\", block_size);\n attention_forward1(cublas_handle, d_out, d_qkvr, d_preatt, d_att, d_inp, B, T, C, NH, block_size);\n validate_result(d_out, out, \"out\", B * T * C);\n // reset d_inp, which is used as a buffer\n cudaCheck(cudaMemcpy(d_inp, qkv, B * T * 3 * C * sizeof(float), cudaMemcpyHostToDevice));\n }\n\n printf(\"\\nChecking backward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"\\nBlock size: %d\\n\", block_size);\n attention_backward(cublas_handle, d_dinp, d_dqkvr, d_dpreatt, d_datt, d_scratch, d_dout, d_qkvr, d_att, B, T, C, NH);\n validate_result(d_dinp, dqkv, \"dinp\", B * T * 3 * C);\n }\n\n printf(\"\\nAll results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmarks:\\n\");\n int repeat_times = 100;\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j ++) {\n int block_size = block_sizes[j];\n float elapsed_time = benchmark_kernel(\n repeat_times, attention_forward1,\n cublas_handle, d_out, d_qkvr, d_preatt, d_att, d_inp, B, T, C, NH, block_size\n );\n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n printf(\"\\nBackward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j ++) {\n int block_size = block_sizes[j];\n float elapsed_time = benchmark_kernel(\n repeat_times, attention_backward,\n cublas_handle, d_dinp, d_dqkvr, d_dpreatt, d_datt, d_scratch, d_dout, d_qkvr, d_att, B, T, C, NH\n );\n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n\n // free memory\n free(qkv);\n free(out);\n free(dout);\n free(dqkv);\n cudaCheck(cudaFree(d_inp));\n cudaCheck(cudaFree(d_qkvr));\n cudaCheck(cudaFree(d_out));\n cudaCheck(cudaFree(d_preatt));\n cudaCheck(cudaFree(d_att));\n cudaCheck(cudaFree(d_dout));\n cudaCheck(cudaFree(d_dinp));\n cudaCheck(cudaFree(d_dqkvr));\n cudaCheck(cudaFree(d_dpreatt));\n cudaCheck(cudaFree(d_datt));\n cudaCheck(cudaFree(d_scratch));\n\n // destroy cublas handle\n cublasCheck(cublasDestroy(cublas_handle));\n}\n#endif", "output": "#include \n#include \n#include \n#include \n#include \n#include \n#include \"common.h\"\n\n\n__global__ void permute_kernel(float* q, float* k, float* v,\n const float* inp,\n int B, int N, int NH, int d) {\n // okay so now, this kernel wants Q,K,V to all be of shape (B, NH, N, d)\n // but instead, we have a single tensor QKV (inp) of shape (B, N, 3, NH, d)\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Q[b][nh_][n][d_] = inp[b][n][0][nh_][d_]\n\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n\n int inp_idx = \\\n (b * N * 3 * NH * d)\n + (n * 3 * NH * d)\n + (0 * NH * d)\n + (nh_ * d)\n + d_;\n\n q[idx] = inp[inp_idx];\n k[idx] = inp[inp_idx + NH * d];\n v[idx] = inp[inp_idx + 2 * (NH * d)];\n }\n}\n\n__global__ void scale_kernel(float* inp, float scale, int B, int NH, int T) {\n // scales the pre-softmax attention scores by scale\n // we don't want causal attention, so we don't set any values to infinity\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * NH * T * T) {\n inp[idx] *= scale;\n }\n}\n\n// warp-level reduction for finding the maximum value\n__device__ float warpReduceMax(float val) {\n for (int offset = 16; offset > 0; offset /= 2) {\n val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));\n }\n return val;\n}\n\n// warp-level reduction for summing values\n__device__ float warpReduceSum(float val) {\n for (int offset = 16; offset > 0; offset /= 2) {\n val += __shfl_down_sync(0xFFFFFFFF, val, offset);\n }\n return val;\n}\n\n__global__ void softmax_forward_kernel4(float* out, const float* inp, int N, int C) {\n // out is (N, C) just like inp. Each row of inp will get softmaxed.\n // same as kernel3, but can handle any block size (multiple of 32)\n // each row of C elements is handled by block_size threads\n // furthermore, each block_size threads get executed in warps of 32 threads\n\n // special reduction operations warpReduceMax/warpReduceSum are used for intra-warp reductions\n // shared memory is used for inter-warp reduction\n extern __shared__ float shared[];\n int idx = blockIdx.x;\n int tid = threadIdx.x;\n int warpId = threadIdx.x / 32; // warp index within a block\n int laneId = threadIdx.x % 32; // thread index within a warp\n\n // the number of warps per block. recall that blockDim.x is block_size\n int warpsPerBlock = blockDim.x / 32;\n\n // shared[] must be allocated to have 2 * warpsPerBlock elements\n // first half for max values, the second half for sum values\n float* maxvals = shared;\n float* sumvals = &shared[warpsPerBlock];\n\n // one row of inp, i.e. inp[idx, :] of shape (C,)\n const float* x = inp + idx * C;\n\n // first, thread coarsening by directly accessing global memory in series\n float maxval = -INFINITY;\n for (int i = tid; i < C; i += blockDim.x) {\n maxval = fmaxf(maxval, x[i]);\n }\n // now within-warp reductions for maxval\n maxval = warpReduceMax(maxval);\n\n // the 0th thread of each warp writes the maxval of that warp to shared memory\n if (laneId == 0) maxvals[warpId] = maxval;\n __syncthreads();\n\n // now the 0th thread reduces the maxvals in shared memory, i.e. across warps\n if (tid == 0) {\n float val = maxvals[tid];\n for (int i = 1; i < warpsPerBlock; i++) {\n val = fmaxf(val, maxvals[i]);\n }\n // store the final max in the first position\n maxvals[0] = val;\n }\n __syncthreads();\n // broadcast the max to all threads\n float offset = maxvals[0];\n\n // compute expf and write the result to global memory\n for (int i = tid; i < C; i += blockDim.x) {\n // subtract max for numerical stability\n out[idx * C + i] = expf(x[i] - offset);\n }\n\n // okay now we calculated exp(x - max(x))\n // step 2: sum all the values and divide by the sum\n\n // thread coarsening for sum\n x = out + idx * C;\n float sumval = 0.0f;\n for (int i = tid; i < C; i += blockDim.x) {\n sumval += x[i];\n }\n // within-warp reduction for sumval\n sumval = warpReduceSum(sumval);\n\n // write sumval to shared memory\n if (laneId == 0) sumvals[warpId] = sumval;\n __syncthreads();\n\n // inter-thread reduction of sum\n if (tid == 0) {\n float val = sumvals[tid];\n for (int i = 1; i < warpsPerBlock; ++i) {\n val += sumvals[i];\n }\n sumvals[0] = val;\n }\n __syncthreads();\n // broadcast the sum to all threads\n float sum = sumvals[0];\n\n // divide the whole row by the sum\n for (int i = tid; i < C; i += blockDim.x) {\n out[idx * C + i] = x[i] / sum;\n }\n}\n\n__global__ void unpermute_kernel(const float* inp, float *out, int B, int N, int NH, int d) {\n // out has shape (B, nh, N, d) but we need to unpermute it to (B, N, nh, d)\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n // out[b][n][nh_][d_] <- inp[b][nh_][n][d_]\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n\n int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_;\n out[other_idx] = inp[idx];\n }\n}\n\n// a simple attention forward function taken directly from llm.c kernel 3\nvoid attention_forward1\n(\n hipblasHandle_t cublas_handle,\n float* out, float* qkvr, float* preatt, float* att,\n float* inp,\n int B, int T, int C, int NH,\n const int block_size\n) {\n // note: resure inp as a scratch buffer because it is not used in backward\n // inp is (B, T, 3C) QKV\n // preatt, att are (B, NH, T, T)\n // output is (B, T, C)\n int HS = C / NH; // head size\n\n // permute and separate inp from (B, T, 3, NH, HS) to 3X (B, NH, T, HS)\n float *q, *k, *v;\n q = qkvr + 0 * B * T * C;\n k = qkvr + 1 * B * T * C;\n v = qkvr + 2 * B * T * C;\n int total_threads = B * NH * T * HS;\n int num_blocks = ceil_div(total_threads, block_size);\n permute_kernel<<>>(q, k, v, inp, B, T, NH, HS);\n\n // batched matrix multiply with cuBLAS\n const float alpha = 1.0f;\n const float beta = 0.0f;\n cublasCheck(hipblasSgemmStridedBatched(cublas_handle,\n HIPBLAS_OP_T, HIPBLAS_OP_N,\n T, T, HS,\n &alpha,\n k, HS, T * HS,\n q, HS, T * HS,\n &beta,\n preatt, T, T * T,\n B * NH));\n\n // multiply all elements of preatt elementwise by scale\n float scale = 1.0f / sqrtf(HS);\n total_threads = B * NH * T * T;\n num_blocks = ceil_div(total_threads, block_size);\n scale_kernel<<>>(preatt, scale, B, NH, T);\n\n // softmax. preatt is (B, NH, T, T) but we view it as (B * NH * T, T) and use the softmax kernel\n int softmax_block_size = 256;\n int grid_size = B * NH * T;\n size_t shared_mem_size = 2 * softmax_block_size / 32 * sizeof(float);\n softmax_forward_kernel4<<>>(att, preatt, B * NH * T, T);\n\n // new approach: first cuBLAS another batched matmul\n // y = att @ v # (B, nh, T, T) @ (B, nh, T, hs) -> (B, nh, T, hs)\n float* vaccum = inp;\n cublasCheck(hipblasSgemmStridedBatched(cublas_handle,\n HIPBLAS_OP_N, HIPBLAS_OP_N,\n HS, T, T,\n &alpha,\n v, HS, T * HS,\n att, T, T * T,\n &beta,\n vaccum, HS, T * HS,\n B * NH));\n\n // now unpermute\n // y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side\n num_blocks = ceil_div(B * T * C, block_size);\n unpermute_kernel<<>>(vaccum, out, B, T, NH, HS);\n}\n\n__global__ void permute_kernel_backward(float* dinp,\n const float* dq, const float* dk, const float* dv,\n int B, int N, int NH, int d) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n\n int inp_idx = (b * N * 3 * NH * d) + (n * 3 * NH * d) + (0 * NH * d) + (nh_ * d) + d_;\n dinp[inp_idx] = dq[idx];\n dinp[inp_idx + NH * d] = dk[idx];\n dinp[inp_idx + 2 * (NH * d)] = dv[idx];\n }\n}\n\n__global__ void unpermute_kernel_backward(float* dinp, const float *dout, int B, int N, int NH, int d) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * NH * N * d) {\n int b = idx / (NH * N * d);\n int rest = idx % (NH * N * d);\n int nh_ = rest / (N * d);\n rest = rest % (N * d);\n int n = rest / d;\n int d_ = rest % d;\n int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_;\n dinp[idx] = dout[other_idx];\n }\n}\n\n// slight edit from the softmax_autoregressive_backward_kernel from train_gpt2_fp32\n// our attention is no longer autoregressive\n__global__ void softmax_backward_kernel(float* dpreatt, const float* datt, const float* att,\n int B, int T, int C, float scale) {\n constexpr const int BlockSize = 256;\n constexpr int T_per_block = 4;\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n __shared__ float block_acc[32];\n\n int idx = blockIdx.y;\n // go through blocks in reverse order, so the slowest block starts first\n int t0 = T - 1 - T_per_block*blockIdx.x;\n\n att += idx * T * T;\n datt += idx * T * T;\n dpreatt += idx * T * T;\n\n if (warp.meta_group_rank() == 0) {\n block_acc[warp.thread_rank()] = 0;\n }\n\n for(int to = 0; to < T_per_block; ++to) {\n int t = t0 - to;\n if(t < 0) return;\n const float* att_bth = att + t * T;\n const float* datt_bth = datt + t * T;\n float* dpreatt_bth = dpreatt + t * T;\n\n float local_sum = 0.0f;\n for (int t2 = block.thread_rank(); t2 < T; t2 += BlockSize) {\n local_sum += att_bth[t2] * datt_bth[t2];\n }\n\n block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus{});\n block.sync();\n local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus{});\n\n for (int t3 = block.thread_rank(); t3 < T; t3 += BlockSize) {\n // don't touch the cache. Some parts will still be here from the previous loop, and\n // we want to exploit those.\n float acc = __ldcs(att_bth + t3) * (__ldcs(datt_bth + t3) - local_sum);\n __stcs(dpreatt_bth + t3, scale * acc);\n }\n }\n}\n\n// the sequence of transformations in this compound op is:\n// inp (B,T,3C) -> qkvr (B,T,3C) -> preatt (B,NH,T,T) -> att (B,NH,T,T) -> vaccum (B,T,C) -> out (B,T,C)\nvoid attention_backward\n(\n hipblasHandle_t cublas_handle,\n float* dinp, float* dqkvr, float* dpreatt, float* datt, float* scratch,\n const float* dout,\n const float* qkvr, const float* att,\n int B, int T, int C, int NH\n) {\n const int block_size = 256;\n int HS = C / NH; // head size\n const float one = 1.0f;\n const float zero = 0.0f; // note beta = 1.0f so that we accumulate gradients (+=)\n // unpack convenience pointers into q, k, v\n const float *q, *k, *v;\n q = qkvr + 0 * B * T * C;\n k = qkvr + 1 * B * T * C;\n v = qkvr + 2 * B * T * C;\n float *dq, *dk, *dv;\n dq = dqkvr + 0 * B * T * C;\n dk = dqkvr + 1 * B * T * C;\n dv = dqkvr + 2 * B * T * C;\n // backward through the unpermute operation\n int num_blocks = ceil_div(B * T * C, block_size);\n unpermute_kernel_backward<<>>(scratch, dout, B, T, NH, HS);\n cudaCheck(hipGetLastError());\n // backward into datt\n cublasCheck(hipblasSgemmStridedBatched(cublas_handle, HIPBLAS_OP_T, HIPBLAS_OP_N, T, T, HS, &one, v, HS, T * HS, scratch, HS, T * HS, &zero, datt, T, T * T, B * NH));\n // backward into dv\n cublasCheck(hipblasSgemmStridedBatched(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_T, HS, T, T, &one, scratch, HS, T * HS, att, T, T * T, &zero, dv, HS, T * HS, B * NH));\n // backward into preatt\n int hs = C / NH; // head size\n float scale = 1.0f / sqrtf(hs);\n softmax_backward_kernel<<>>(dpreatt, datt, att, B, T, C, scale);\n cudaCheck(hipGetLastError());\n // backward into q\n cublasCheck(hipblasSgemmStridedBatched(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_N, HS, T, T, &one, k, HS, T * HS, dpreatt, T, T * T, &zero, dq, HS, T * HS, B * NH));\n // backward into k\n cublasCheck(hipblasSgemmStridedBatched(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_T, HS, T, T, &one, q, HS, T * HS, dpreatt, T, T * T, &zero, dk, HS, T * HS, B * NH));\n // backward into inp\n num_blocks = ceil_div(B * NH * T * HS, block_size);\n permute_kernel_backward<<>>(dinp, dq, dk, dv, B, T, NH, HS);\n cudaCheck(hipGetLastError());\n}\n\n\n#ifndef LINKING\nint main(int argc, char** argv) {\n int B = 4;\n int T = 1024;\n int C = 256;\n int HS = 32;\n int NH = C / HS;\n\n // setup cublas\n hipblasHandle_t cublas_handle;\n cublasCheck(hipblasCreate(&cublas_handle));\n cublasCheck(hipblasSetMathMode(cublas_handle, HIPBLAS_DEFAULT_MATH));\n\n // create host memory to load data\n float* qkv = (float*)malloc(B * T * 3 * C * sizeof(float));\n float* out = (float*)malloc(B * T * C * sizeof(float));\n float* dout = (float*)malloc(B * T * C * sizeof(float));\n float* dqkv = (float*)malloc(B * T * 3 * C * sizeof(float));\n\n // read saved output\n FILE * file = fopen(\"attention.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(qkv, sizeof(float), B * T * 3 * C, file);\n freadCheck(out, sizeof(float), B * T * C, file);\n freadCheck(dout, sizeof(float), B * T * C, file);\n freadCheck(dqkv, sizeof(float), B * T * 3 * C, file);\n fclose(file);\n\n // allocate device memory\n float *d_inp, *d_qkvr, *d_out, *d_preatt, *d_att, *d_dout, *d_dinp, *d_dqkvr, *d_dpreatt, *d_datt, *d_scratch;\n cudaCheck(hipMalloc(&d_inp, B * T * 3 * C * sizeof(float)));\n cudaCheck(hipMemcpy(d_inp, qkv, B * T * 3 * C * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMalloc(&d_qkvr, B * T * 3 * C * sizeof(float)));\n cudaCheck(hipMemset(d_qkvr, 0, B * T * 3 * C * sizeof(float)));\n cudaCheck(hipMalloc(&d_out, B * T * C * sizeof(float)));\n cudaCheck(hipMemset(d_out, 0, B * T * C * sizeof(float)));\n cudaCheck(hipMalloc(&d_preatt, B * NH * T * T * sizeof(float)));\n cudaCheck(hipMemset(d_preatt, 0, B * NH * T * T * sizeof(float)));\n cudaCheck(hipMalloc(&d_att, B * NH * T * T * sizeof(float)));\n cudaCheck(hipMemset(d_att, 0, B * NH * T * T * sizeof(float)));\n cudaCheck(hipMalloc(&d_dout, B * T * C * sizeof(float)));\n cudaCheck(hipMemcpy(d_dout, dout, B * T * C * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMalloc(&d_dinp, B * T * 3 * C * sizeof(float)));\n cudaCheck(hipMemset(d_dinp, 0, B * T * 3 * C * sizeof(float)));\n cudaCheck(hipMalloc(&d_dqkvr, B * T * 3 * C * sizeof(float)));\n cudaCheck(hipMemset(d_dqkvr, 0, B * T * 3 * C * sizeof(float)));\n cudaCheck(hipMalloc(&d_dpreatt, B * NH * T * T * sizeof(float)));\n cudaCheck(hipMemset(d_dpreatt, 0, B * NH * T * T * sizeof(float)));\n cudaCheck(hipMalloc(&d_datt, B * NH * T * T * sizeof(float)));\n cudaCheck(hipMemset(d_datt, 0, B * NH * T * T * sizeof(float)));\n cudaCheck(hipMalloc(&d_scratch, B * T * C * sizeof(float)));\n cudaCheck(hipMemset(d_scratch, 0, B * T * C * sizeof(float)));\n\n // forward pass\n int block_sizes[] = {256, 512};\n printf(\"Checking forward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"\\nBlock size: %d\\n\", block_size);\n attention_forward1(cublas_handle, d_out, d_qkvr, d_preatt, d_att, d_inp, B, T, C, NH, block_size);\n validate_result(d_out, out, \"out\", B * T * C);\n // reset d_inp, which is used as a buffer\n cudaCheck(hipMemcpy(d_inp, qkv, B * T * 3 * C * sizeof(float), hipMemcpyHostToDevice));\n }\n\n printf(\"\\nChecking backward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"\\nBlock size: %d\\n\", block_size);\n attention_backward(cublas_handle, d_dinp, d_dqkvr, d_dpreatt, d_datt, d_scratch, d_dout, d_qkvr, d_att, B, T, C, NH);\n validate_result(d_dinp, dqkv, \"dinp\", B * T * 3 * C);\n }\n\n printf(\"\\nAll results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmarks:\\n\");\n int repeat_times = 100;\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j ++) {\n int block_size = block_sizes[j];\n float elapsed_time = benchmark_kernel(\n repeat_times, attention_forward1,\n cublas_handle, d_out, d_qkvr, d_preatt, d_att, d_inp, B, T, C, NH, block_size\n );\n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n printf(\"\\nBackward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j ++) {\n int block_size = block_sizes[j];\n float elapsed_time = benchmark_kernel(\n repeat_times, attention_backward,\n cublas_handle, d_dinp, d_dqkvr, d_dpreatt, d_datt, d_scratch, d_dout, d_qkvr, d_att, B, T, C, NH\n );\n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n\n // free memory\n free(qkv);\n free(out);\n free(dout);\n free(dqkv);\n cudaCheck(hipFree(d_inp));\n cudaCheck(hipFree(d_qkvr));\n cudaCheck(hipFree(d_out));\n cudaCheck(hipFree(d_preatt));\n cudaCheck(hipFree(d_att));\n cudaCheck(hipFree(d_dout));\n cudaCheck(hipFree(d_dinp));\n cudaCheck(hipFree(d_dqkvr));\n cudaCheck(hipFree(d_dpreatt));\n cudaCheck(hipFree(d_datt));\n cudaCheck(hipFree(d_scratch));\n\n // destroy cublas handle\n cublasCheck(hipblasDestroy(cublas_handle));\n}\n#endif", "source": "/kaggle/working/unet.cu/dev/attention.cu"} {"input": "#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \"silu.cuh\"\n\n\n\n// ----------------------------------------------------------------------------\n// GPU kernels\n\n\n__global__ void silu_forward_kernel(\n const float* x,\n float* out,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx >= N) {\n return;\n }\n\n float x_val = x[idx];\n out[idx] = x_val / (1.0f + expf(-x_val));\n}\n\n\n__global__ void silu_backward_kernel(\n const float* dout, const float* x,\n float* dx,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx >= N) {\n return;\n }\n float out_val = dout[idx];\n float x_val = x[idx];\n float expx = expf(-x_val);\n float grad_silu = (1.0f + x_val * expx / (1.0f + expx)) / (1.0f + expx);\n dx[idx] = out_val * grad_silu;\n}\n\n\n// ----------------------------------------------------------------------------\n// CUDA kernel launcher\n\nvoid silu_forward(\n const float* x,\n float* out,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n silu_forward_kernel<<>>(x, out, N);\n}\n\nvoid silu_backward(\n const float* dout, const float* x,\n float* dx,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n silu_backward_kernel<<>>(dout, x, dx, N);\n}\n\n// ----------------------------------------------------------------------------\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n int B = 1;\n int C = 32;\n int H = 32;\n int W = 32;\n int N = B * C * H * W;\n\n // set up the device\n int deviceIdx = 0;\n cudaCheck(cudaSetDevice(deviceIdx));\n cudaDeviceProp deviceProp;\n cudaGetDeviceProperties(&deviceProp, deviceIdx);\n printf(\"Device %d: %s\\n\", deviceIdx, deviceProp.name);\n\n // create host memory and load data\n float* x = (float*)malloc(N * sizeof(float));\n float* out = (float*)malloc(N * sizeof(float));\n float* dout = (float*)malloc(N * sizeof(float));\n float* dx = (float*)malloc(N * sizeof(float));\n\n // read saved output\n FILE *file = fopenCheck(\"silu.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n\n freadCheck(x, sizeof(float), N, file);\n freadCheck(out, sizeof(float), N, file);\n freadCheck(dout, sizeof(float), N, file);\n freadCheck(dx, sizeof(float), N, file);\n fcloseCheck(file);\n\n\n // allocate device memory\n float *d_x, *d_out, *d_dout, *d_dx;\n cudaCheck(cudaMalloc(&d_x, N * sizeof(float)));\n cudaCheck(cudaMalloc(&d_out, N * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dout, N * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dx, N * sizeof(float)));\n cudaCheck(cudaMemcpy(d_x, x, N * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_dout, dout, N * sizeof(float), cudaMemcpyHostToDevice));\n\n int block_sizes[] = {128, 256, 512, 1024};\n printf(\"Checking forward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"Checking block size %d\\n\", block_size);\n silu_forward(d_x, d_out, N, block_size);\n validate_result(d_out, out, \"out\", N);\n }\n printf(\"Forward pass: all results match\\n\\n\");\n\n printf(\"Checking backward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"Checking block size %d\\n\", block_size);\n silu_backward(d_dout, d_x, d_dx, N, block_size);\n validate_result(d_dx, dx, \"dx\", N);\n }\n printf(\"Backward pass: all results match\\n\\n\");\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, silu_forward,\n d_x, d_out, N, block_size);\n\n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n\n printf(\"\\nBackward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, silu_backward,\n d_dout, d_x, d_dx,\n N, block_size);\n\n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n\n // free memory\n free(x);\n free(out);\n free(dout);\n free(dx);\n\n cudaCheck(cudaFree(d_x));\n cudaCheck(cudaFree(d_out));\n cudaCheck(cudaFree(d_dout));\n cudaCheck(cudaFree(d_dx));\n}\n#endif", "output": "#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \"silu.cuh\"\n\n\n\n// ----------------------------------------------------------------------------\n// GPU kernels\n\n\n__global__ void silu_forward_kernel(\n const float* x,\n float* out,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx >= N) {\n return;\n }\n\n float x_val = x[idx];\n out[idx] = x_val / (1.0f + expf(-x_val));\n}\n\n\n__global__ void silu_backward_kernel(\n const float* dout, const float* x,\n float* dx,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx >= N) {\n return;\n }\n float out_val = dout[idx];\n float x_val = x[idx];\n float expx = expf(-x_val);\n float grad_silu = (1.0f + x_val * expx / (1.0f + expx)) / (1.0f + expx);\n dx[idx] = out_val * grad_silu;\n}\n\n\n// ----------------------------------------------------------------------------\n// CUDA kernel launcher\n\nvoid silu_forward(\n const float* x,\n float* out,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n silu_forward_kernel<<>>(x, out, N);\n}\n\nvoid silu_backward(\n const float* dout, const float* x,\n float* dx,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n silu_backward_kernel<<>>(dout, x, dx, N);\n}\n\n// ----------------------------------------------------------------------------\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n int B = 1;\n int C = 32;\n int H = 32;\n int W = 32;\n int N = B * C * H * W;\n\n // set up the device\n int deviceIdx = 0;\n cudaCheck(hipSetDevice(deviceIdx));\n hipDeviceProp_t deviceProp;\n hipGetDeviceProperties(&deviceProp, deviceIdx);\n printf(\"Device %d: %s\\n\", deviceIdx, deviceProp.name);\n\n // create host memory and load data\n float* x = (float*)malloc(N * sizeof(float));\n float* out = (float*)malloc(N * sizeof(float));\n float* dout = (float*)malloc(N * sizeof(float));\n float* dx = (float*)malloc(N * sizeof(float));\n\n // read saved output\n FILE *file = fopenCheck(\"silu.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n\n freadCheck(x, sizeof(float), N, file);\n freadCheck(out, sizeof(float), N, file);\n freadCheck(dout, sizeof(float), N, file);\n freadCheck(dx, sizeof(float), N, file);\n fcloseCheck(file);\n\n\n // allocate device memory\n float *d_x, *d_out, *d_dout, *d_dx;\n cudaCheck(hipMalloc(&d_x, N * sizeof(float)));\n cudaCheck(hipMalloc(&d_out, N * sizeof(float)));\n cudaCheck(hipMalloc(&d_dout, N * sizeof(float)));\n cudaCheck(hipMalloc(&d_dx, N * sizeof(float)));\n cudaCheck(hipMemcpy(d_x, x, N * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_dout, dout, N * sizeof(float), hipMemcpyHostToDevice));\n\n int block_sizes[] = {128, 256, 512, 1024};\n printf(\"Checking forward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"Checking block size %d\\n\", block_size);\n silu_forward(d_x, d_out, N, block_size);\n validate_result(d_out, out, \"out\", N);\n }\n printf(\"Forward pass: all results match\\n\\n\");\n\n printf(\"Checking backward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"Checking block size %d\\n\", block_size);\n silu_backward(d_dout, d_x, d_dx, N, block_size);\n validate_result(d_dx, dx, \"dx\", N);\n }\n printf(\"Backward pass: all results match\\n\\n\");\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, silu_forward,\n d_x, d_out, N, block_size);\n\n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n\n printf(\"\\nBackward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, silu_backward,\n d_dout, d_x, d_dx,\n N, block_size);\n\n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n\n // free memory\n free(x);\n free(out);\n free(dout);\n free(dx);\n\n cudaCheck(hipFree(d_x));\n cudaCheck(hipFree(d_out));\n cudaCheck(hipFree(d_dout));\n cudaCheck(hipFree(d_dx));\n}\n#endif", "source": "/kaggle/working/unet.cu/dev/silu.cu"} {"input": "#include \n#include \"add.cuh\"\n#include \"common.h\"\n\n// simple kernel to add two tensors\n__global__ void add_forward_kernel(\n const float* a, const float* b,\n float* out,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < N) {\n out[idx] = a[idx] + b[idx];\n }\n}\n\nvoid add_forward(\n const float* a, const float* b,\n float* out,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n add_forward_kernel<<>>(a, b, out, N);\n}\n\n\n// add A and B in place, and store result in B\n__global__ void add_inplace_forward_kernel(\n const float* a, float* b,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < N) {\n b[idx] += a[idx];\n }\n}\n\nvoid add_inplace_forward(\n const float* a, float* b,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n add_inplace_forward_kernel<<>>(a, b, N);\n}", "output": "#include \n#include \"add.cuh\"\n#include \"common.h\"\n\n// simple kernel to add two tensors\n__global__ void add_forward_kernel(\n const float* a, const float* b,\n float* out,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < N) {\n out[idx] = a[idx] + b[idx];\n }\n}\n\nvoid add_forward(\n const float* a, const float* b,\n float* out,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n add_forward_kernel<<>>(a, b, out, N);\n}\n\n\n// add A and B in place, and store result in B\n__global__ void add_inplace_forward_kernel(\n const float* a, float* b,\n int N\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < N) {\n b[idx] += a[idx];\n }\n}\n\nvoid add_inplace_forward(\n const float* a, float* b,\n int N, int block_size\n) {\n int n_blk = ceil_div(N, block_size);\n add_inplace_forward_kernel<<>>(a, b, N);\n}", "source": "/kaggle/working/unet.cu/dev/add.cu"} {"input": "#include \n#include \n#include \n#include // For POSIX systems to use getcwd\n#include // For PATH_MAX\n#include \n#include \n#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \"linear.cuh\"\n#include \"silu.cuh\"\n#include \"conv2d_k3.cuh\"\n#include \"resblock.cuh\"\n#include \"add.cuh\"\n#include \"avgpool.cuh\"\n#include \"attention_block.cuh\"\n#include \"concat_channel.cuh\"\n#include \"add.cuh\"\n#include \"upsample.cuh\"\n#include \"groupnorm.cuh\"\n#include \"mse.cuh\"\n#include \"timestep_embedding.cuh\"\n#include \"rand.h\"\n#include \n\n\n\ntypedef struct {\n int C_in;\n int C_model;\n int C_time;\n int C_out;\n int B;\n int H;\n int W;\n int gn_n_groups;\n int HS;\n int block_size;\n int n_layers;\n int n_levels;\n int n_res;\n int att_start_level;\n int* channel_mult;\n int max_period;\n} UnetConfig;\n\n// ------------------------------------------------------------------------------------------------\n// DataLoader\n\nconst char* DATASET_FILE = \"../datasets/elephant_train.bin\";\ntypedef struct {\n int B;\n int C_in;\n int H;\n int W;\n int img_size;\n int n_imgs;\n FILE* img_file;\n long cur_position;\n long file_size;\n long header_offset;\n long num_batches;\n float* input;\n} DataLoader;\n\nvoid dataloader_init(DataLoader* loader, const char* filename, int B) {\n loader->B = B;\n\n loader->img_file = fopenCheck(filename, \"rb\");\n int data_header[256];\n loader->header_offset = 256 * sizeof(int);\n freadCheck(data_header, sizeof(int), 256, loader->img_file);\n if (data_header[0] != 20240620) {\n fprintf(stderr, \"Invalid magic data file\\n\");\n exit(1);\n }\n loader->n_imgs = data_header[1];\n loader->C_in = data_header[2];\n loader->H = data_header[3];\n loader->W = data_header[4];\n loader->img_size = loader->C_in * loader->H * loader->W;\n\n // determine file size\n fseekCheck(loader->img_file, 0, SEEK_SET);\n fseekCheck(loader->img_file, 0, SEEK_END);\n loader->file_size = ftell(loader->img_file);\n fseekCheck(loader->img_file, 0, SEEK_SET);\n assert((loader->file_size - loader->header_offset) / (loader->img_size * sizeof(float)) == loader->n_imgs);\n \n loader->cur_position = loader->header_offset;\n loader->num_batches = loader->n_imgs / B;\n\n // allocate cuda CPU pinned memory for inputs\n cudaMallocHost((void**)&loader->input, B * loader->img_size * sizeof(float));\n}\n\nvoid dataloader_reset(DataLoader* loader) {\n loader->cur_position = loader->header_offset;\n}\n\nvoid dataloader_next_batch(DataLoader* loader) {\n int B = loader->B;\n int img_size = loader->img_size;\n if (loader->cur_position + (B * img_size * sizeof(float)) > loader->file_size) {\n dataloader_reset(loader);\n }\n fseekCheck(loader->img_file, loader->cur_position, SEEK_SET);\n freadCheck(loader->input, sizeof(float), B * img_size, loader->img_file);\n loader->cur_position += B * img_size * sizeof(float);\n}\n\nvoid dataloader_free(DataLoader* loader) {\n fcloseCheck(loader->img_file);\n cudaFreeHost(loader->input);\n}\n\n// ------------------------------------------------------------------------------------------------\n// Diffusion\ntypedef struct {\n int B;\n int img_size;\n int max_period;\n float* betas; // (max_period,)\n float* sqrt_alphas_cumprod; // (max_period,)\n float* sqrt_one_minus_alphas_cumprod; // (max_period,)\n int block_size;\n curandState* curand_states;\n mt19937_state host_rng;\n} GaussianDiffusion;\n\n__global__ void init_curand_states(curandState *state, unsigned long seed, int N) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < N) {\n curand_init(seed, idx, 0, &state[idx]);\n }\n}\n\nvoid diffusion_init(\n GaussianDiffusion* diffusion, int B, int C_in, int H, int W, int max_period, int block_size = 512,\n unsigned long seed = 0, float beta_start = 1e-4, float beta_end = 0.02f\n) {\n diffusion->B = B;\n diffusion->max_period = max_period;\n diffusion->block_size = block_size;\n diffusion->img_size = C_in * H * W;\n\n float scale = 1000.0f / max_period;\n beta_start *= scale;\n beta_end *= scale;\n\n // calculate betas and alphas on host then move to device\n float* betas = (float*)mallocCheck(max_period * sizeof(float));\n float* sqrt_alphas_cumprod = (float*)mallocCheck(max_period * sizeof(float));\n float* sqrt_one_minus_alphas_cumprod = (float*)mallocCheck(max_period * sizeof(float));\n\n float alpha_cumprod = 1.0f;\n for (int t = 1; t <= max_period; t++) {\n // linear interpolate for beta\n betas[t - 1] = (beta_start * (max_period - t) + beta_end * (t - 1)) / (max_period - 1);\n alpha_cumprod *= 1.0f - betas[t - 1];\n sqrt_alphas_cumprod[t - 1] = sqrtf(alpha_cumprod);\n sqrt_one_minus_alphas_cumprod[t - 1] = sqrtf(1.0f - alpha_cumprod);\n }\n\n float *d_betas, *d_sqrt_alphas_cumprod, *d_sqrt_one_minus_alphas_cumprod;\n cudaCheck(cudaMalloc(&d_betas, max_period * sizeof(float)));\n cudaCheck(cudaMalloc(&d_sqrt_alphas_cumprod, max_period * sizeof(float)));\n cudaCheck(cudaMalloc(&d_sqrt_one_minus_alphas_cumprod, max_period * sizeof(float)));\n\n cudaCheck(cudaMemcpy(d_betas, betas, max_period * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_sqrt_alphas_cumprod, sqrt_alphas_cumprod, max_period * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_sqrt_one_minus_alphas_cumprod, sqrt_one_minus_alphas_cumprod, max_period * sizeof(float), cudaMemcpyHostToDevice));\n\n free(betas);\n free(sqrt_alphas_cumprod);\n free(sqrt_one_minus_alphas_cumprod);\n\n diffusion->betas = d_betas;\n diffusion->sqrt_alphas_cumprod = d_sqrt_alphas_cumprod;\n diffusion->sqrt_one_minus_alphas_cumprod = d_sqrt_one_minus_alphas_cumprod;\n\n // init curand states for sampling noise of the size of one batch of inputs\n int total_states = B * C_in * H * W;\n int n_blocks = ceil_div(total_states, block_size);\n cudaCheck(cudaMalloc(&diffusion->curand_states, total_states * sizeof(curandState)));\n init_curand_states<<>>(diffusion->curand_states, seed, total_states);\n\n // init host rng for drawing random timesteps\n manual_seed(&diffusion->host_rng, seed);\n}\n\n// gridDim will be of shape (N / block_size, B)\n// expect x and x_t to be of shape (B, N)\n// and timesteps to be of shape (B,)\n// expect x_t to be filled with standard normals\n__global__ void diffusion_sample_kernel(\n float* timesteps, float* x, float* x_t,\n float* sqrt_alphas_cumprod, float* sqrt_one_minus_alphas_cumprod,\n int B, int N\n) {\n int b = blockIdx.y;\n int n = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (n < N && b < B) {\n int t = (int)timesteps[b];\n float sqrt_alpha_t = sqrt_alphas_cumprod[t];\n float sqrt_one_minus_alpha_t = sqrt_one_minus_alphas_cumprod[t];\n float x_t_val = x_t[b * N + n]; // x_t starts out as standard normals\n float x_val = x[b * N + n];\n x_t[b * N + n] = x_val * sqrt_alpha_t + x_t_val * sqrt_one_minus_alpha_t;\n }\n}\n\nvoid diffusion_sample(\n float* timesteps, float* x, float* x_t,\n GaussianDiffusion *diffusion,\n int B, int N\n) {\n dim3 gridDim(ceil_div(N, diffusion->block_size), B);\n diffusion_sample_kernel<<block_size>>>(timesteps, x, x_t, diffusion->sqrt_alphas_cumprod, diffusion->sqrt_one_minus_alphas_cumprod, B, N);\n cudaCheck(cudaDeviceSynchronize());\n}\n\n// a kernel that creates x_t = sqrt_alpha * x + sqrt(1 - alpha) * noise in one go\n// that draws the noise from standard normals, gets alpha from the timesteps and the precomputed alpha cumprod\n// and saves output to x_t\n// if the input x, timesteps and the alpha buffers are null, then it will just fill x_t with standard normals\n// to be called with B * N total threads\n__global__ void diffusion_draw_and_sample_xt_kernel(\n float* x_t, int B, int N, curandState* curand_states,\n float* sqrt_alphas_cumprod, float* sqrt_one_minus_alphas_cumprod,\n float* timesteps = nullptr, float* x = nullptr\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * N) {\n float out_val = curand_normal(&curand_states[idx]);\n if (timesteps != nullptr && x != nullptr) {\n int b = idx / N;\n int n = idx % N;\n float x_val = x[b * N + n];\n int t = (int)timesteps[b];\n float sqrt_alpha_t = sqrt_alphas_cumprod[t];\n float sqrt_one_minus_alpha_t = sqrt_one_minus_alphas_cumprod[t];\n out_val = x_val * sqrt_alpha_t + out_val * sqrt_one_minus_alpha_t;\n }\n x_t[idx] = out_val;\n }\n}\n\nvoid diffusion_draw_and_sample_xt(\n GaussianDiffusion *diffusion,\n float* x_t,\n float* timesteps = nullptr, float* x = nullptr\n) {\n int total_states = diffusion->B * diffusion->img_size;\n int n_blocks = ceil_div(total_states, diffusion->block_size);\n diffusion_draw_and_sample_xt_kernel<<block_size>>>(\n x_t, diffusion->B, diffusion->img_size, diffusion->curand_states, diffusion->sqrt_alphas_cumprod, diffusion->sqrt_one_minus_alphas_cumprod, timesteps, x\n );\n}\n\nvoid diffusion_free(GaussianDiffusion* diffusion) {\n cudaCheck(cudaFree(diffusion->betas));\n cudaCheck(cudaFree(diffusion->sqrt_alphas_cumprod));\n cudaCheck(cudaFree(diffusion->sqrt_one_minus_alphas_cumprod));\n cudaCheck(cudaFree(diffusion->curand_states));\n}\n\n// samples B timesteps in range [0, max_period), convert them to floats and\n// store in timesteps\nvoid sample_timesteps(GaussianDiffusion *diffusion, float* timesteps, int N = -1) {\n mt19937_state* rng = &diffusion->host_rng;\n if (N < 0) {\n N = diffusion->B;\n }\n int max_period = diffusion->max_period;\n for (int i = 0; i < N; i++) {\n float rand_float = randfloat32(rng);\n timesteps[i] = (float)(int)(rand_float * max_period);\n }\n}\n\n// ------------------------------------------------------------------------------------------------\n// Unet init\n\n// we just need a single set of buffers that are shared for all the backward pass activations\n#define NUM_SKIP_CONNECTIONS 12 // (n_res + 1) * n_levels\n#define NUM_BACK_BUFS 10\n#define UNET_NUM_BACK_ACTS 22\ntypedef struct {\n float* dinp_buf1; // used to store the input gradient, which will be reused as dout for the previous layer\n float* dinp_buf2;\n float* buf1_BCHW; // NOTE: first 4 buffers will have max B * C * H * W size\n float* buf2_BCHW;\n float* buf3_BCHW;\n float* buf4_BCHW;\n float* buf1_B3CHW; // NOTE: last 2 buffers will have max B * 3 * C * H * W size for attention layers\n float* buf2_B3CHW; // so max of H * W is only over dimensions that have attention\n float* demb_buf1; // used to store cumulative gradients w.r.t. timestep embeddings\n float* demb_buf2;\n float* skip_dinps[NUM_SKIP_CONNECTIONS];\n} UnetBackActs;\n\n// This is used as a block to store all the params / grads / acts for all the layers\ntypedef struct {\n float* mem_ptr;\n size_t total_size;\n int n_layers;\n void** layer_ptrs;\n size_t* layer_sizes;\n} UnetBlock;\n\ntypedef struct {\n float* back_mem;\n UnetBackActs back_acts;\n size_t back_sizes[UNET_NUM_BACK_ACTS];\n size_t n_backs;\n} UnetBackBlock;\n\ntypedef struct {\n UnetConfig config;\n float* input;\n float* timestep_emb; // this is the already computed embeddings of shape (B, C_model)\n UnetBlock params;\n UnetBlock grads;\n UnetBlock acts;\n UnetBackBlock back;\n float* skip_acts[NUM_SKIP_CONNECTIONS]; // pointers to skip connetion activations\n // the next two are for debug purposes only, just to keep track of the buffer that contains the grad w.r.t. inputs\n float* dinp;\n // model output\n float* output;\n // targets and losses\n float* target;\n float mean_loss;\n float* gpu_loss;\n // buffers for AdamW\n float* m_memory;\n float* v_memory;\n TimestepEmbedding time_emb_state;\n} Unet;\n\nvoid unet_init(Unet* unet) {\n unet->input = nullptr;\n unet->timestep_emb = nullptr;\n unet->params.mem_ptr = nullptr;\n unet->grads.mem_ptr = nullptr;\n unet->acts.mem_ptr = nullptr;\n unet->back.back_mem = nullptr;\n unet->dinp = nullptr;\n unet->output = nullptr;\n unet->target = nullptr;\n unet->mean_loss = -1.0f;\n unet->gpu_loss = nullptr;\n unet->m_memory = nullptr;\n unet->v_memory = nullptr;\n\n init_timestep_embedding(&unet->time_emb_state, unet->config.C_model, unet->config.B, unet->config.max_period);\n}\n\n\ntemplate \nvoid free_not_null(T** ptr) {\n if (*ptr != nullptr) {\n free(*ptr);\n *ptr = nullptr;\n }\n}\n\nvoid free_unet_block(UnetBlock* params) {\n free_not_null(¶ms->layer_sizes);\n for (int l = 0; l < params->n_layers; l++) {\n free_not_null(&(params->layer_ptrs[l]));\n }\n free_not_null(¶ms->layer_ptrs);\n cudaCheck(cudaFree(params->mem_ptr));\n}\n\nvoid free_unet(Unet* unet) {\n cudaCheck(cudaFree(unet->input));\n cudaCheck(cudaFree(unet->timestep_emb));\n cudaCheck(cudaFree(unet->gpu_loss));\n cudaCheck(cudaFree(unet->m_memory));\n cudaCheck(cudaFree(unet->v_memory));\n\n free_unet_block(&(unet->params));\n free_unet_block(&(unet->grads));\n free_unet_block(&(unet->acts));\n\n cudaCheck(cudaFree(unet->back.back_mem));\n free_timestep_embedding(&unet->time_emb_state);\n}\n\nvoid unet_count_layers(UnetConfig* config) {\n int n_layers = 0;\n int n_res = config->n_res;\n int n_levels = config->n_levels;\n int att_start_level = config->att_start_level;\n\n // timestep emb\n n_layers += 3;\n\n // initial conv\n n_layers += 1;\n\n // add downsample resblock layers\n n_layers += n_res * n_levels;\n\n // add downsample attention\n n_layers += n_res * (n_levels - att_start_level);\n\n // add avgpool layers\n n_layers += n_levels - 1;\n\n // add middle layers\n n_layers += 3;\n\n // add upsample concat layers\n n_layers += (n_res + 1) * n_levels;\n\n // add upsample resblock layers\n n_layers += (n_res + 1) * n_levels;\n\n // add upsample attention layers\n n_layers += (n_res + 1) * (n_levels - att_start_level);\n\n // add upsample layers\n n_layers += n_levels - 1;\n\n // output layers\n n_layers += 3;\n\n config->n_layers = n_layers;\n}\n\n\n// ------------------------------------------------------------------------------------------------\n// allocate memory for structs with info for each layer and count params\n\n\nint make_linear_params_ptrs(\n UnetBlock* params,\n int layer,\n int C_in, int C_out\n) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(LinearParams));\n params->layer_sizes[layer] = linear_count_params(C_in, C_out);\n return layer + 1;\n}\n\nint make_silu_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_convk3_params_ptrs(\n UnetBlock* params, int layer,\n int C_in, int OC, int B, int H, int W\n) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ConvK3Params));\n params->layer_sizes[layer] = convk3_count_params(C_in, OC);\n return layer + 1;\n}\n\nint make_resblock_params_ptrs(\n UnetBlock* params, int layer,\n int C, int C_emb, int C_out, int B, int H, int W, int gn_n_groups\n) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ResBlockParameters));\n ResBlockParameters* resblock_params = (ResBlockParameters*)params->layer_ptrs[layer];\n resblock_count_params(resblock_params, C, C_emb, C_out, B, H, W, 0, 0, gn_n_groups);\n params->layer_sizes[layer] = resblock_params->n_params;\n return layer + 1;\n}\n\nint make_avgpool_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_attention_block_params_ptrs(UnetBlock* params, int layer, int C) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(AttentionParams));\n AttentionParams* attention_params = (AttentionParams*)params->layer_ptrs[layer];\n attention_block_count_params(attention_params, C);\n params->layer_sizes[layer] = attention_params->n_params;\n return layer + 1;\n}\n\nint make_concat_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_upsample_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_groupnorm_params_ptrs(UnetBlock* params, int layer, int C) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(GroupNormParams));\n params->layer_sizes[layer] = C * 2;\n return layer + 1;\n}\n\n\nvoid unet_setup_and_count_params(UnetConfig* config, UnetBlock* params) {\n assert (config->n_layers > 0);\n\n params->n_layers = config->n_layers;\n params->layer_ptrs = (void**)mallocCheck(config->n_layers * sizeof(void*));\n params->layer_sizes = (size_t*)mallocCheck(config->n_layers * sizeof(size_t));\n\n int B = config->B;\n int C_in = config->C_in;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int C_out = config->C_out;\n int H = config->H;\n int W = config->W;\n int gn_n_groups = config->gn_n_groups;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n int layer = 0;\n layer = make_linear_params_ptrs(params, layer, C_model, C_time);\n layer = make_silu_params_ptrs(params, layer);\n layer = make_linear_params_ptrs(params, layer, C_time, C_time);\n int OC = C_model;\n layer = make_convk3_params_ptrs(params, layer, C_in, OC, B, H, W);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n layer = make_resblock_params_ptrs(params, layer, C, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_params_ptrs(params, layer, C);\n }\n }\n if (l < n_levels - 1) {\n layer = make_avgpool_params_ptrs(params, layer);\n H /= 2;\n W /= 2;\n }\n }\n\n // middle layers\n layer = make_resblock_params_ptrs(params, layer, C, C_time, C, B, H, W, gn_n_groups);\n layer = make_attention_block_params_ptrs(params, layer, C);\n layer = make_resblock_params_ptrs(params, layer, C, C_time, C, B, H, W, gn_n_groups);\n\n // upwards: concat then resblock\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = C_model * mult;\n for (int r = 0; r < n_res + 1; r++) {\n if (r == n_res && l > 0) {\n int prev_mult = config->channel_mult[l - 1];\n C_skip = C_model * prev_mult;\n }\n OC = C_model * mult;\n layer = make_concat_params_ptrs(params, layer);\n layer = make_resblock_params_ptrs(params, layer, C + C_skip, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_params_ptrs(params, layer, C);\n }\n }\n if (l > 0) {\n layer = make_upsample_params_ptrs(params, layer);\n H *= 2;\n W *= 2;\n }\n }\n\n layer = make_groupnorm_params_ptrs(params, layer, C);\n layer = make_silu_params_ptrs(params, layer);\n layer = make_convk3_params_ptrs(params, layer, C, C_out, B, H, W);\n\n // count total size\n size_t total_size = 0;\n for (int i = 0; i < config->n_layers; i++) {\n total_size += params->layer_sizes[i];\n }\n params->total_size = total_size;\n}\n\nint make_linear_acts_ptrs(\n UnetBlock* acts,\n int layer,\n int B, int C_out\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(LinearActs));\n // acts buffers only hold outputs\n acts->layer_sizes[layer] = B * C_out;\n return layer + 1;\n}\n\nint make_silu_acts_ptrs(UnetBlock* acts, int layer, int N) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(SiluActs));\n acts->layer_sizes[layer] = N;\n return layer + 1;\n}\n\nint make_convk3_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int OC, int H, int W\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ConvK3Acts));\n acts->layer_sizes[layer] = B * OC * H * W;\n return layer + 1;\n}\n\nint make_resblock_acts_ptrs(\n UnetBlock* acts, int layer,\n int C, int C_emb, int C_out, int B, int H, int W, int gn_n_groups\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ResBlockActivations));\n ResBlockActivations* resblock_acts = (ResBlockActivations*)acts->layer_ptrs[layer];\n resblock_count_acts(resblock_acts, C, C_emb, C_out, B, H, W, 0, 0, gn_n_groups);\n acts->layer_sizes[layer] = resblock_acts->n_acts;\n return layer + 1;\n}\n\nint make_avgpool_acts_ptrs(UnetBlock* acts, int layer, int B, int C, int H, int W) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(AvgpoolActs));\n acts->layer_sizes[layer] = B * C * H/2 * W/2;\n return layer + 1;\n}\n\nint make_attention_block_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C, int H, int W, int gn_n_groups, int HS\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(AttentionActs));\n AttentionActs* att_acts = (AttentionActs*)acts->layer_ptrs[layer];\n attention_block_count_acts(att_acts, B, C, H, W, gn_n_groups, HS);\n acts->layer_sizes[layer] = att_acts->n_acts;\n return layer + 1;\n}\n\nint make_concat_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C1, int C2, int H, int W\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ConcatChannelActs));\n acts->layer_sizes[layer] = B * (C1 + C2) * H * W;\n return layer + 1;\n}\n\nint make_upsample_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C, int H, int W\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(UpsampleActs));\n acts->layer_sizes[layer] = B * C * H * 2 * W * 2;\n return layer + 1;\n}\n\nint make_groupnorm_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C, int H, int W, int gn_n_groups\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(GroupNormActs));\n // out (B, C, H, W), mean and rstd (B, gn_n_groups)\n acts->layer_sizes[layer] = B * C * H * W + B * gn_n_groups * 2;\n return layer + 1;\n}\n\nvoid unet_setup_and_count_acts(UnetConfig* config, UnetBlock* acts) {\n assert (config->n_layers > 0);\n\n acts->n_layers = config->n_layers;\n acts->layer_ptrs = (void**)mallocCheck(config->n_layers * sizeof(void*));\n acts->layer_sizes = (size_t*)mallocCheck(config->n_layers * sizeof(size_t));\n\n int B = config->B;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int C_out = config->C_out;\n int H = config->H;\n int W = config->W;\n int gn_n_groups = config->gn_n_groups;\n int HS = config->HS;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n int layer = 0;\n layer = make_linear_acts_ptrs(acts, layer, B, C_time);\n layer = make_silu_acts_ptrs(acts, layer, B * C_time);\n layer = make_linear_acts_ptrs(acts, layer, B, C_time);\n int OC = C_model;\n layer = make_convk3_acts_ptrs(acts, layer, B, OC, H, W);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n layer = make_resblock_acts_ptrs(acts, layer, C, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups, HS);\n }\n }\n if (l < n_levels - 1) {\n layer = make_avgpool_acts_ptrs(acts, layer, B, C, H, W);\n H /= 2;\n W /= 2;\n }\n }\n\n layer = make_resblock_acts_ptrs(acts, layer, C, C_time, C, B, H, W, gn_n_groups);\n layer = make_attention_block_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups, HS);\n layer = make_resblock_acts_ptrs(acts, layer, C, C_time, C, B, H, W, gn_n_groups);\n\n // upwards\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = mult * C_model;\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n int prev_mult = config->channel_mult[l - 1];\n C_skip = prev_mult * C_model;\n }\n layer = make_concat_acts_ptrs(acts, layer, B, C, C_skip, H, W);\n OC = mult * C_model;\n layer = make_resblock_acts_ptrs(acts, layer, C + C_skip, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups, HS);\n }\n }\n if (l > 0) {\n layer = make_upsample_acts_ptrs(acts, layer, B, C, H, W);\n H *= 2;\n W *= 2;\n }\n }\n layer = make_groupnorm_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups);\n layer = make_silu_acts_ptrs(acts, layer, B * C * H * W);\n layer = make_convk3_acts_ptrs(acts, layer, B, C_out, H, W);\n\n // count total size\n size_t total_size = 0;\n for (int i = 0; i < config->n_layers; i++) {\n total_size += acts->layer_sizes[i];\n }\n acts->total_size = total_size;\n}\n\nvoid unet_setup_and_count_backs(UnetConfig* config, UnetBackBlock* back) {\n int B = config->B;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int H = config->H;\n int W = config->W;\n\n size_t max_BCHW = 3 * B * C_model * H * W;\n size_t max_B3CHW = 2 * B * 3 * C_model * H * W;\n back->back_sizes[0] = max_BCHW;\n back->back_sizes[1] = max_BCHW;\n back->back_sizes[2] = max_BCHW;\n back->back_sizes[3] = max_BCHW;\n back->back_sizes[4] = max_BCHW;\n back->back_sizes[5] = max_BCHW;\n back->back_sizes[6] = max_B3CHW;\n back->back_sizes[7] = max_B3CHW;\n back->back_sizes[8] = B * C_time;\n back->back_sizes[9] = B * C_time;\n\n // fill the sizes of the skip connection backward acts\n int C = C_model;\n for (int l = 0; l < config->n_levels; l++) {\n for (int r = 0; r < (config->n_res + 1); r++) {\n int s = l * 3 + r + NUM_BACK_BUFS;\n assert(s < UNET_NUM_BACK_ACTS);\n back->back_sizes[s] = B * C * H * W;\n int mult = config->channel_mult[l];\n C = C_model * mult;\n if (l > 0 && r == 0) {\n H /= 2;\n W /= 2;\n }\n }\n }\n\n size_t n_backs = 0;\n for (int i = 0; i < UNET_NUM_BACK_ACTS; i++) {\n n_backs += back->back_sizes[i];\n }\n back->n_backs = n_backs;\n}\n\n\nvoid unet_make_ptrs_and_count_memory(Unet* unet)\n{\n unet_setup_and_count_params(&(unet->config), &(unet->params));\n unet_setup_and_count_params(&(unet->config), &(unet->grads));\n unet_setup_and_count_acts(&(unet->config), &(unet->acts));\n unet_setup_and_count_backs(&(unet->config), &(unet->back));\n}\n\n// ------------------------------------------------------------------------------------------------\n// malloc unet memory and set pointers\n\ntypedef struct {\n float* input;\n float* prev_out;\n int layer;\n float* emb;\n float* emb_cum; // cumulative demb\n float* mem_ptr;\n int skip_idx;\n} LayerInfo;\n\n\nvoid set_linear_params_ptrs(\n UnetBlock* params,\n LayerInfo* l_info,\n int C_in, int C_out\n) {\n int layer = l_info->layer;\n LinearParams* linear_params = (LinearParams*)params->layer_ptrs[layer];\n linear_set_param_ptrs(linear_params, l_info->mem_ptr, C_in, C_out);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_convk3_params_ptrs(UnetBlock* params, LayerInfo* l_info, int C_in, int OC) {\n int layer = l_info->layer;\n ConvK3Params* convk3_params = (ConvK3Params*)params->layer_ptrs[layer];\n convk3_set_param_ptrs(convk3_params, l_info->mem_ptr, C_in, OC);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_resblock_layer_params_ptrs(UnetBlock* params, LayerInfo* l_info, int C, int OC) {\n int layer = l_info->layer;\n ResBlockParameters* resblock_params = (ResBlockParameters*)params->layer_ptrs[layer];\n set_resblock_params_ptrs(C, OC, resblock_params, l_info->mem_ptr);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_attention_layer_params_ptrs(UnetBlock* params, LayerInfo* l_info) {\n int layer = l_info->layer;\n AttentionParams* att_params = (AttentionParams*)params->layer_ptrs[layer];\n set_attention_params_pointers(att_params, l_info->mem_ptr);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_groupnorm_layer_params_ptrs(UnetBlock* params, LayerInfo* l_info, int C) {\n int layer = l_info->layer;\n GroupNormParams* gn_params = (GroupNormParams*)params->layer_ptrs[layer];\n gn_set_param_ptrs(gn_params, l_info->mem_ptr, C);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\n\nvoid point_unet_params(UnetConfig* config, UnetBlock* params) {\n // common constants\n int C_in = config->C_in;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int C_out = config->C_out;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n LayerInfo l_info;\n l_info.layer = 0;\n l_info.mem_ptr = params->mem_ptr;\n set_linear_params_ptrs(params, &l_info, C_model, C_time);\n l_info.layer++; // skip silu layer\n set_linear_params_ptrs(params, &l_info, C_time, C_time);\n int OC = C_model;\n set_convk3_params_ptrs(params, &l_info, C_in, OC);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n set_resblock_layer_params_ptrs(params, &l_info, C, OC);\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_params_ptrs(params, &l_info);\n }\n }\n if (l < n_levels - 1) {\n l_info.layer++; // skip avgpool layer\n }\n }\n\n set_resblock_layer_params_ptrs(params, &l_info, C, C);\n set_attention_layer_params_ptrs(params, &l_info);\n set_resblock_layer_params_ptrs(params, &l_info, C, C);\n\n // upward: skip one layer for concat\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = mult * C_model;\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config->channel_mult[l - 1];\n }\n l_info.layer++;\n OC = mult * C_model;\n set_resblock_layer_params_ptrs(params, &l_info, C + C_skip, OC);\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_params_ptrs(params, &l_info);\n }\n }\n if (l > 0) {\n l_info.layer++; // upsample\n }\n }\n\n set_groupnorm_layer_params_ptrs(params, &l_info, C);\n l_info.layer++; // silu\n set_convk3_params_ptrs(params, &l_info, C, C_out);\n}\n\nvoid set_linear_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n LinearActs* linear_acts = (LinearActs*)acts->layer_ptrs[layer];\n linear_acts->out = l_info->mem_ptr;\n l_info->mem_ptr += acts->layer_sizes[layer];\n linear_acts->inp = l_info->input;\n l_info->input = linear_acts->out;\n l_info->layer++;\n}\n\nvoid set_silu_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n SiluActs* silu_acts = (SiluActs*)acts->layer_ptrs[layer];\n silu_acts->out = l_info->mem_ptr;\n l_info->mem_ptr += acts->layer_sizes[layer];\n silu_acts->x = l_info->input;\n l_info->input = silu_acts->out;\n l_info->layer++;\n}\n\nvoid set_convk3_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n ConvK3Acts* convk3_acts = (ConvK3Acts*)acts->layer_ptrs[layer];\n convk3_acts->out = l_info->mem_ptr;\n l_info->mem_ptr += acts->layer_sizes[layer];\n convk3_acts->inp = l_info->input;\n l_info->input = convk3_acts->out;\n l_info->layer++;\n}\n\nvoid set_resblock_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info, int C, int OC) {\n int layer = l_info->layer;\n ResBlockActivations* resblock_acts = (ResBlockActivations*)acts->layer_ptrs[layer];\n set_resblock_acts_ptrs(C, OC, resblock_acts, l_info->mem_ptr);\n resblock_acts->input = l_info->input;\n resblock_acts->emb = l_info->emb;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = resblock_acts->add2;\n l_info->layer++;\n}\n\nvoid set_avgpool_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n AvgpoolActs* avgpool_acts = (AvgpoolActs*)acts->layer_ptrs[layer];\n avgpool_acts->out = l_info->mem_ptr;\n avgpool_acts->x = l_info->input;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = avgpool_acts->out;\n l_info->layer++;\n}\n\nvoid set_attention_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n AttentionActs* att_acts = (AttentionActs*)acts->layer_ptrs[layer];\n set_attention_acts_pointers(att_acts, l_info->mem_ptr);\n att_acts->input = l_info->input;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = att_acts->add;\n l_info->layer++;\n}\n\nvoid set_concat_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info, float** skip_acts) {\n int layer = l_info->layer;\n ConcatChannelActs* cat_acts = (ConcatChannelActs*)acts->layer_ptrs[layer];\n cat_acts->out = l_info->mem_ptr;\n cat_acts->x1 = l_info->input;\n cat_acts->x2 = skip_acts[l_info->skip_idx];\n\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = cat_acts->out;\n l_info->layer++;\n l_info->skip_idx--;\n}\n\nvoid set_upsample_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n UpsampleActs* up_acts = (UpsampleActs*)acts->layer_ptrs[layer];\n up_acts->x = l_info->input;\n up_acts->out = l_info->mem_ptr;\n\n l_info->mem_ptr += acts-> layer_sizes[layer];\n l_info->input = up_acts->out;\n l_info->layer++;\n}\n\nvoid set_groupnorm_layer_acts_ptrs(\n UnetBlock* acts, LayerInfo* l_info,\n int B, int C, int H, int W, int gn_n_groups\n) {\n int layer = l_info->layer;\n GroupNormActs* gn_a = (GroupNormActs*)acts->layer_ptrs[layer];\n gn_set_act_ptrs(gn_a, l_info->mem_ptr, B, C, H, W, gn_n_groups);\n gn_a->x = l_info->input;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = gn_a->out;\n l_info->layer++;\n}\n\nvoid point_unet_acts(UnetConfig* config, UnetBlock* acts, float** skip_acts, float* input, float* timestep_emb) {\n // common constants\n int B = config->B;\n int C_model = config->C_model;\n int H = config->H;\n int W = config->W;\n int gn_n_groups = config->gn_n_groups;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n LayerInfo l_info;\n l_info.layer = 0;\n l_info.mem_ptr = acts->mem_ptr;\n\n // timestep emb\n l_info.input = timestep_emb;\n set_linear_acts_ptrs(acts, &l_info);\n set_silu_acts_ptrs(acts, &l_info);\n set_linear_acts_ptrs(acts, &l_info);\n // point input to model input, and set outputu of the last linear layer as the embedding input\n l_info.emb = l_info.input;\n l_info.input = input;\n\n int OC = C_model;\n set_convk3_acts_ptrs(acts, &l_info);\n int skip_idx = 0;\n skip_acts[skip_idx] = l_info.input; // set the output of the previouse layer as input to the skip connection\n skip_idx++;\n\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n set_resblock_layer_acts_ptrs(acts, &l_info, C, OC);\n if (l < att_start_level) {\n // for the first 2 layers, the skip connection comes from the resblocks\n skip_acts[skip_idx] = l_info.input;\n skip_idx++;\n }\n\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_acts_ptrs(acts, &l_info);\n skip_acts[skip_idx] = l_info.input;\n skip_idx++;\n }\n }\n if (l < n_levels - 1) {\n set_avgpool_layer_acts_ptrs(acts, &l_info);\n skip_acts[skip_idx] = l_info.input;\n skip_idx++;\n }\n }\n\n set_resblock_layer_acts_ptrs(acts, &l_info, C, C);\n set_attention_layer_acts_ptrs(acts, &l_info);\n set_resblock_layer_acts_ptrs(acts, &l_info, C, C);\n\n // upward: set concat\n l_info.skip_idx = NUM_SKIP_CONNECTIONS - 1;\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = mult * C_model;\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config->channel_mult[l - 1];\n }\n set_concat_layer_acts_ptrs(acts, &l_info, skip_acts);\n OC = mult * C_model;\n set_resblock_layer_acts_ptrs(acts, &l_info, C + C_skip, OC);\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_acts_ptrs(acts, &l_info);\n }\n }\n if (l > 0) {\n set_upsample_layer_acts_ptrs(acts, &l_info);\n }\n }\n set_groupnorm_layer_acts_ptrs(acts, &l_info, B, C, H, W, gn_n_groups);\n set_silu_acts_ptrs(acts, &l_info);\n set_convk3_acts_ptrs(acts, &l_info);\n}\n\nvoid set_unet_back_pointers(UnetBackBlock* back) {\n UnetBackActs* back_acts = &(back->back_acts);\n float* back_mem_ptr = back->back_mem;\n\n float** ptrs[] = {\n &(back_acts->dinp_buf1),\n &(back_acts->dinp_buf2),\n &(back_acts->buf1_BCHW),\n &(back_acts->buf2_BCHW),\n &(back_acts->buf3_BCHW),\n &(back_acts->buf4_BCHW),\n &(back_acts->buf1_B3CHW),\n &(back_acts->buf2_B3CHW),\n &(back_acts->demb_buf1),\n &(back_acts->demb_buf2)\n };\n\n for (int i = 0; i < NUM_BACK_BUFS; i++) {\n *(ptrs[i]) = back_mem_ptr;\n back_mem_ptr += back->back_sizes[i];\n }\n\n // point the back acts for skip connections\n for (int s = 0; s < NUM_SKIP_CONNECTIONS; s++) {\n back_acts->skip_dinps[s] = back_mem_ptr;\n back_mem_ptr += back->back_sizes[NUM_BACK_BUFS + s];\n }\n}\n\nvoid malloc_and_point(Unet* unet) {\n float *param_mem, *grad_mem, *act_mem, *back_mem;\n cudaCheck(cudaMalloc(¶m_mem, unet->params.total_size * sizeof(float)));\n cudaCheck(cudaMalloc(&grad_mem, unet->params.total_size * sizeof(float)));\n cudaCheck(cudaMalloc(&act_mem, unet->acts.total_size * sizeof(float)));\n cudaCheck(cudaMalloc(&back_mem, unet->back.n_backs * sizeof(float)));\n\n cudaCheck(cudaMemset(param_mem, 0, unet->params.total_size * sizeof(float)));\n cudaCheck(cudaMemset(grad_mem, 0, unet->params.total_size * sizeof(float)));\n cudaCheck(cudaMemset(act_mem, 0, unet->acts.total_size * sizeof(float)));\n cudaCheck(cudaMemset(back_mem, 0, unet->back.n_backs * sizeof(float)));\n\n // constants\n int B = unet->config.B;\n int C_in = unet->config.C_in;\n int C_model = unet->config.C_model;\n int H = unet->config.H;\n int W = unet->config.W;\n\n // malloc and point input and timestep emb\n cudaCheck(cudaMalloc(&(unet->input), B * C_in * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&(unet->timestep_emb), B * C_model * sizeof(float)));\n // malloc a nicely aligned buffer for the loss\n cudaCheck(cudaMalloc(&(unet->gpu_loss), 64 * sizeof(float)));\n\n unet->params.mem_ptr = param_mem;\n unet->grads.mem_ptr = grad_mem;\n unet->acts.mem_ptr = act_mem;\n unet->back.back_mem = back_mem;\n\n point_unet_params(&(unet->config), &(unet->params));\n point_unet_params(&(unet->config), &(unet->grads));\n point_unet_acts(&(unet->config), &(unet->acts), unet->skip_acts, unet->input, unet->timestep_emb);\n set_unet_back_pointers(&(unet->back));\n}\n\n// ------------------------------------------------------------------------------------------------\n// forward pass\n\nvoid linear_layer_forward(\n Unet* unet,\n LayerInfo* l_info,\n cublasHandle_t cublas_handle, int N, int C, int OC, int block_size\n) {\n int layer = l_info->layer;\n LinearParams* linear_params = (LinearParams*)unet->params.layer_ptrs[layer];\n LinearActs* linear_acts = (LinearActs*)unet->acts.layer_ptrs[layer];\n matmul_forward2(cublas_handle, linear_acts->out, linear_acts->inp, linear_params->w, linear_params->b, N, C, OC, block_size);\n l_info->layer++;\n}\n\nvoid silu_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int N, int block_size\n) {\n SiluActs* acts = (SiluActs*)unet->acts.layer_ptrs[l_info->layer];\n silu_forward(acts->x, acts->out, N, block_size);\n l_info->layer++;\n}\n\nvoid convk3_layer_forward(\n Unet* unet, LayerInfo* l_info,\n cublasHandle_t cublas_handle, int B, int C, int OC, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n ConvK3Params* conv_params = (ConvK3Params*)unet->params.layer_ptrs[layer];\n ConvK3Acts* conv_acts = (ConvK3Acts*)unet->acts.layer_ptrs[layer];\n conv2d_k3_forward3(conv_acts->inp, conv_params->w, conv_params->b, conv_acts->out, B, C, OC, H, W);\n l_info->layer++;\n}\nvoid resblock_layer_forward(\n Unet* unet, LayerInfo* l_info,\n cublasHandle_t cublas_handle, int C_in, int C_emb, int C_out, int B, int H, int W, int gn_n_groups, int block_size\n) {\n int layer = l_info->layer;\n ResBlockParameters* resblock_params = (ResBlockParameters*)unet->params.layer_ptrs[layer];\n ResBlockActivations* resblock_acts = (ResBlockActivations*)unet->acts.layer_ptrs[layer];\n resblock_forward(cublas_handle, C_in, C_emb, C_out, B, H, W, block_size, 0, 0, gn_n_groups, resblock_params, resblock_acts);\n l_info->layer++;\n}\n\nvoid avgpool_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n AvgpoolActs* avgpool_acts = (AvgpoolActs*)unet->acts.layer_ptrs[layer];\n avgpool_2d_forward1(avgpool_acts->out, avgpool_acts->x, B, C, H, W, block_size);\n l_info->layer++;\n}\n\nvoid attention_layer_forward(\n Unet* unet, LayerInfo* l_info,\n cublasHandle_t cublas_handle, int B, int C, int H, int W, int HS, int gn_n_groups, int block_size\n) {\n int layer = l_info->layer;\n AttentionParams* att_params = (AttentionParams*)unet->params.layer_ptrs[layer];\n AttentionActs* att_acts = (AttentionActs*)unet->acts.layer_ptrs[layer];\n attention_block_forward(cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, att_params, att_acts);\n l_info->layer++;\n}\n\nvoid concat_layer_forward(\n Unet* unet,\n LayerInfo* l_info,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n ConcatChannelActs* cat_acts = (ConcatChannelActs*)unet->acts.layer_ptrs[layer];\n concat_channel_forward(cat_acts->x1, cat_acts->x2, cat_acts->out, B, C1, C2, H, W, block_size);\n l_info->layer++;\n}\n\nvoid upsample_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n UpsampleActs* up_acts = (UpsampleActs*)unet->acts.layer_ptrs[layer];\n upsample_forward1(up_acts->out, up_acts->x, B, C, H, W, block_size);\n l_info->layer++;\n}\n\nvoid groupnorm_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int gn_n_groups\n) {\n int layer = l_info->layer;\n GroupNormParams* gn_p = (GroupNormParams*)unet->params.layer_ptrs[layer];\n GroupNormActs* gn_a = (GroupNormActs*)unet->acts.layer_ptrs[layer];\n groupnorm_forward(gn_a->x, gn_p->w, gn_p->b, gn_a->out, gn_a->mean, gn_a->rstd, B, C, H, W, gn_n_groups);\n l_info->layer++;\n}\n\n\nvoid unet_forward(\n cublasHandle_t cublas_handle,\n Unet* unet,\n float* t_res = nullptr, float* t_att = nullptr\n) {\n // common constants\n UnetConfig config = unet->config;\n int B = config.B;\n int C_in = config.C_in;\n int C_model = config.C_model;\n int C_time = config.C_time;\n int C_out = config.C_out;\n int H = config.H;\n int W = config.W;\n int block_size = config.block_size;\n int gn_n_groups = config.gn_n_groups;\n int HS = config.HS;\n int n_levels = config.n_levels;\n int n_res = config.n_res;\n int att_start_level = config.att_start_level;\n\n cudaEvent_t start, end;\n if (t_res != nullptr) {\n cudaCheck(cudaEventCreate(&start));\n cudaCheck(cudaEventCreate(&end));\n }\n\n LayerInfo l_info;\n l_info.layer = 0;\n linear_layer_forward(unet, &l_info, cublas_handle, B, C_model, C_time, block_size);\n silu_layer_forward(unet, &l_info, B * C_time, block_size);\n linear_layer_forward(unet, &l_info, cublas_handle, B, C_time, C_time, block_size);\n int OC = C_model;\n convk3_layer_forward(unet, &l_info, cublas_handle, B, C_in, OC, H, W, block_size);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config.channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(start));\n }\n resblock_layer_forward(unet, &l_info, cublas_handle, C, C_time, OC, B, H, W, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n *t_res += elapsed_time;\n printf(\"down resblock: l %d, r %d, Cin %d, Cout %d, H %d, W %d, time %f\\n\", l, r, C, OC, H, W, elapsed_time);\n }\n C = OC;\n if (l >= att_start_level) {\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(start));\n }\n attention_layer_forward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n *t_att += elapsed_time;\n printf(\"down att: l %d, r %d, C %d, H %d, W %d, time %f\\n\", l, r, C, H, W, elapsed_time);\n }\n }\n }\n if (l < n_levels - 1) {\n avgpool_layer_forward(unet, &l_info, B, C, H, W, block_size);\n H /= 2;\n W /= 2;\n }\n }\n\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(start));\n }\n resblock_layer_forward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n *t_res += elapsed_time;\n printf(\"middle resblock 1: Cin %d, Cout %d, H %d, W %d, time %f\\n\", C, OC, H, W, elapsed_time);\n cudaCheck(cudaEventRecord(start));\n }\n attention_layer_forward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n *t_att += elapsed_time;\n printf(\"middle att: Cin %d, H %d, W %d, time %f\\n\", C, H, W, elapsed_time);\n cudaCheck(cudaEventRecord(start));\n }\n resblock_layer_forward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n printf(\"middle resblock 2: Cin %d, Cout %d, H %d, W %d, time %f\\n\", C, OC, H, W, elapsed_time);\n *t_res += elapsed_time;\n }\n\n // upward\n for (int l = n_levels - 1; l >= 0; l--) {\n int C_skip = C_model * config.channel_mult[l];\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config.channel_mult[l - 1];\n }\n concat_layer_forward(unet, &l_info, B, C, C_skip, H, W, block_size);\n OC = C_model * config.channel_mult[l];\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(start));\n }\n resblock_layer_forward(unet, &l_info, cublas_handle, C + C_skip, C_time, OC, B, H, W, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n printf(\"up resblock: l %d, r %d, Cin %d, Cskip %d, Cout %d, H %d, W %d, time %f\\n\", l, r, C, C_skip, OC, H, W, elapsed_time);\n *t_res += elapsed_time;\n }\n C = OC;\n if (l >= att_start_level) {\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(start));\n }\n attention_layer_forward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n printf(\"up att: l %d, r %d, C %d, H %d, W %d, time %f\\n\", l, r, C, H, W, elapsed_time);\n *t_att += elapsed_time;\n }\n }\n }\n if (l > 0) {\n upsample_layer_forward(unet, &l_info, B, C, H, W, block_size);\n H *= 2;\n W *= 2;\n }\n }\n\n groupnorm_layer_forward(unet, &l_info, B, C, H, W, gn_n_groups);\n silu_layer_forward(unet, &l_info, B * C * H * W, block_size);\n convk3_layer_forward(unet, &l_info, cublas_handle, B, C, C_out, H, W, block_size);\n\n ConvK3Acts* last_acts = (ConvK3Acts*)unet->acts.layer_ptrs[l_info.layer - 1];\n unet->output = last_acts->out;\n if (unet->target == nullptr) {\n // don't have targets, not computing loss\n unet->mean_loss = -1.0f;\n } else {\n mse_forward(unet->output, unet->target, unet->gpu_loss, B * C_out * H * W, block_size);\n cudaCheck(cudaMemcpy(&(unet->mean_loss), unet->gpu_loss, sizeof(float), cudaMemcpyDeviceToHost));\n }\n}\n\n// ------------------------------------------------------------------------------------------------\n// backward pass\nvoid swap_layer_in_out_ptrs(LayerInfo* l_info) {\n float* temp = l_info->input;\n l_info->input = l_info->prev_out;\n l_info->prev_out = temp;\n}\n\nvoid linear_layer_backward(\n Unet* unet,\n LayerInfo* l_info,\n cublasHandle_t cublas_handle, int N, int C, int OC\n) {\n int layer = l_info->layer;\n LinearParams* linear_params = (LinearParams*)unet->params.layer_ptrs[layer];\n LinearParams* linear_grads = (LinearParams*)unet->grads.layer_ptrs[layer];\n LinearActs* linear_acts = (LinearActs*)unet->acts.layer_ptrs[layer];\n matmul_backward1(cublas_handle, l_info->prev_out, linear_grads->w, linear_grads->b, l_info->input, linear_acts->inp, linear_params->w, N, C, OC);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid silu_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int N, int block_size\n) {\n SiluActs* acts = (SiluActs*)unet->acts.layer_ptrs[l_info->layer];\n silu_backward(l_info->input, acts->x, l_info->prev_out, N, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid convk3_layer_backward(\n Unet* unet, LayerInfo* l_info,\n cublasHandle_t cublas_handle, int B, int C, int OC, int H, int W, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n add_inplace_forward(skip_dinp, l_info->input, B * OC * H * W, block_size);\n l_info->skip_idx--;\n }\n int layer = l_info->layer;\n ConvK3Params* conv_params = (ConvK3Params*)unet->params.layer_ptrs[layer];\n ConvK3Params* conv_grads = (ConvK3Params*)unet->grads.layer_ptrs[layer];\n ConvK3Acts* conv_acts = (ConvK3Acts*)unet->acts.layer_ptrs[layer];\n UnetBackActs* unet_back_acts = &(unet->back.back_acts);\n conv2d_k3_backward2(l_info->input, conv_acts->inp, conv_params->w, unet_back_acts->buf1_BCHW, unet_back_acts->buf2_BCHW, l_info->prev_out, conv_grads->w, conv_grads->b, B, C, OC, H, W);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid resblock_layer_backward(\n Unet* unet, LayerInfo* l_info,\n cublasHandle_t cublas_handle, int C_in, int C_emb, int C_out, int B, int H, int W, int gn_n_groups, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n add_inplace_forward(skip_dinp, l_info->input, B * C_out * H * W, block_size);\n l_info->skip_idx--;\n }\n int layer = l_info->layer;\n ResBlockParameters* resblock_params = (ResBlockParameters*)unet->params.layer_ptrs[layer];\n ResBlockParameters* resblock_grads = (ResBlockParameters*)unet->grads.layer_ptrs[layer];\n ResBlockActivations* resblock_acts = (ResBlockActivations*)unet->acts.layer_ptrs[layer];\n ResBlockBackwardActivations back_acts;\n UnetBackActs* unet_back_acts = &(unet->back.back_acts);\n\n // point back_acts to buffers\n back_acts.buf_BCemb = unet_back_acts->buf1_BCHW;\n back_acts.buf_BCHoWo = unet_back_acts->buf2_BCHW;\n back_acts.buf1_BCHW = unet_back_acts->buf3_BCHW;\n back_acts.buf2_BCHW = unet_back_acts->buf4_BCHW;\n back_acts.dweight_buf = unet_back_acts->buf1_B3CHW;\n back_acts.dbias_buf = unet_back_acts->buf2_B3CHW;\n back_acts.dout = l_info->input;\n back_acts.dx = l_info->prev_out;\n back_acts.demb = l_info->emb;\n\n resblock_backward(cublas_handle, C_in, C_emb, C_out, B, H, W, block_size, 0, 0, gn_n_groups, resblock_params, resblock_grads, resblock_acts, &back_acts);\n swap_layer_in_out_ptrs(l_info);\n\n // add demb to cumulative demb\n add_inplace_forward(l_info->emb, l_info->emb_cum, B * C_emb, block_size);\n l_info->layer--;\n}\n\nvoid attention_layer_backward(\n Unet* unet, LayerInfo* l_info,\n cublasHandle_t cublas_handle, int B, int C, int H, int W, int HS, int gn_n_groups, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n //printf(\"att adding skip %d\\n\", l_info->skip_idx);\n add_inplace_forward(skip_dinp, l_info->input, B * C * H * W, block_size);\n l_info->skip_idx--;\n }\n int layer = l_info->layer;\n AttentionParams* att_params = (AttentionParams*)unet->params.layer_ptrs[layer];\n AttentionParams* att_grads = (AttentionParams*)unet->grads.layer_ptrs[layer];\n AttentionActs* att_acts = (AttentionActs*)unet->acts.layer_ptrs[layer];\n AttentionBackwardActs att_backs;\n UnetBackActs* unet_backs = &(unet->back.back_acts);\n att_backs.buf1_BCHW = unet_backs->buf1_BCHW;\n att_backs.buf2_BCHW = unet_backs->buf2_BCHW;\n att_backs.buf_B3CHW = unet_backs->buf1_B3CHW;\n att_backs.dqkvr = unet_backs->buf2_B3CHW;\n att_backs.dpreatt = unet_backs->buf3_BCHW;\n att_backs.datt = unet_backs->buf4_BCHW;\n att_backs.dout = l_info->input;\n att_backs.dinp = l_info->prev_out;\n\n attention_block_backward(cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, att_params, att_acts, &att_backs, att_grads);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid concat_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n //printf(\"concat skip %d\\n\", l_info->skip_idx);\n concat_channel_backward(l_info->input, l_info->prev_out, skip_dinp, B, C1, C2, H, W, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n l_info->skip_idx++;\n}\n\n\nvoid avgpool_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n //printf(\"avgpool adding skip %d\\n\", l_info->skip_idx);\n add_inplace_forward(skip_dinp, l_info->input, B * C * H/2 * W/2, block_size);\n l_info->skip_idx--;\n }\n avgpool_2d_backward1(l_info->input, l_info->prev_out, B, C, H, W, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid upsample_layer_backward(\n Unet* unet, LayerInfo *l_info,\n int B, int C, int H, int W, int block_size\n) {\n upsample_backward1(l_info->prev_out, l_info->input, B, C, H, W, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid groupnorm_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int gn_n_groups\n) {\n int layer = l_info->layer;\n GroupNormParams* gn_p = (GroupNormParams*)unet->params.layer_ptrs[layer];\n GroupNormParams* gn_g = (GroupNormParams*)unet->grads.layer_ptrs[layer];\n GroupNormActs* gn_a = (GroupNormActs*)unet->acts.layer_ptrs[layer];\n groupnorm_backward(l_info->input, gn_a->x, gn_a->mean, gn_a->rstd, gn_p->w, l_info->prev_out, gn_g->w, gn_g->b, B, C, H, W, gn_n_groups);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid unet_backward(\n cublasHandle_t cublas_handle,\n Unet* unet,\n float* t_res = nullptr, float* t_att = nullptr\n) {\n if (unet->mean_loss == -1.0f or unet->target == nullptr) {\n printf(\"Error: must run forward with targets before calling backward\\n\");\n exit(EXIT_FAILURE);\n }\n cudaEvent_t start, end;\n if (t_res != nullptr) {\n cudaCheck(cudaEventCreate(&start));\n cudaCheck(cudaEventCreate(&end));\n }\n // common constants\n UnetConfig config = unet->config;\n int B = config.B;\n int C_in = config.C_in;\n int C_model = config.C_model;\n int C_time = config.C_time;\n int C_out = config.C_out;\n int H = config.H;\n int W = config.W;\n int block_size = config.block_size;\n int gn_n_groups = config.gn_n_groups;\n int HS = config.HS;\n int n_levels = config.n_levels;\n int n_res = config.n_res;\n int att_start_level = config.att_start_level;\n\n UnetBackActs* back_acts = &(unet->back.back_acts);\n mse_backward(unet->output, unet->target, back_acts->dinp_buf1, B * C_out * H * W, block_size);\n\n LayerInfo l_info;\n l_info.layer = config.n_layers - 1;\n l_info.input = back_acts->dinp_buf1;\n l_info.prev_out = back_acts->dinp_buf2;\n l_info.emb = back_acts->demb_buf1;\n l_info.emb_cum = back_acts->demb_buf2;\n l_info.skip_idx = 0;\n\n int C;\n int OC;\n int C_skip;\n C = C_model;\n OC = C_model;\n convk3_layer_backward(unet, &l_info, cublas_handle, B, C, C_out, H, W, block_size);\n silu_layer_backward(unet, &l_info, B * C * H * W, block_size);\n groupnorm_layer_backward(unet, &l_info, B, C, H, W, gn_n_groups);\n\n for (int l = 0; l < n_levels; l++) {\n for (int r = n_res; r >= 0; r--) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config.channel_mult[l - 1];\n } else {\n C_skip = C_model * config.channel_mult[l];\n }\n if (l >= att_start_level) {\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(start));\n }\n attention_layer_backward(unet, &l_info, cublas_handle, B, OC, H, W, HS, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n *t_att += elapsed_time;\n printf(\"up att back: l %d, r %d, Cout %d, H %d, W %d, time %f\\n\", l, r, OC, H, W, elapsed_time);\n }\n }\n if (r == 0 && l < n_levels - 1) {\n C = C_model * config.channel_mult[l + 1];\n }\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(start));\n }\n resblock_layer_backward(unet, &l_info, cublas_handle, C + C_skip, C_time, OC, B, H, W, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n *t_res += elapsed_time;\n printf(\"up res back: l %d, r %d, Cin %d, Cskip %d, Cout %d, H %d, W %d, time %f\\n\", l, r, C, C_skip, OC, H, W, elapsed_time);\n }\n concat_layer_backward(unet, &l_info, B, C, C_skip, H, W, block_size);\n }\n if (l < n_levels - 1) {\n H /= 2;\n W /= 2;\n upsample_layer_backward(unet, &l_info, B, C, H, W, block_size);\n OC = C_model * config.channel_mult[l + 1];\n }\n }\n\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(start));\n }\n resblock_layer_backward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n *t_res += elapsed_time;\n printf(\"middle resblock 2: Cin %d, Cout %d, H %d, W %d, time %f\\n\", C, OC, H, W, elapsed_time);\n cudaCheck(cudaEventRecord(start));\n }\n attention_layer_backward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n *t_att += elapsed_time;\n printf(\"middle att: Cin %d, H %d, W %d, time %f\\n\", C, H, W, elapsed_time);\n cudaCheck(cudaEventRecord(start));\n }\n resblock_layer_backward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n printf(\"middle resblock 1: Cin %d, Cout %d, H %d, W %d, time %f\\n\", C, OC, H, W, elapsed_time);\n *t_res += elapsed_time;\n }\n\n // we will need to change the backward for all the other blocks below to include the dout contribution from the skip connection.\n l_info.skip_idx = NUM_SKIP_CONNECTIONS - 1;\n for (int l = n_levels - 1; l >= 0; l--) {\n if (l < n_levels - 1) {\n H *= 2;\n W *= 2;\n avgpool_layer_backward(unet, &l_info, B, C, H, W, block_size, 1);\n }\n for (int r = 0; r < n_res; r++) {\n if (l >= att_start_level) {\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(start));\n }\n attention_layer_backward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, 1);\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n printf(\"down att back: l %d, r %d, C %d, H %d, W %d, time %f\\n\", l, r, C, H, W, elapsed_time);\n *t_att += elapsed_time;\n }\n }\n int res_add_skip = (l < att_start_level) ? 1 : 0;\n if (r == n_res - 1 && l > 0) {\n int prev_mult = config.channel_mult[l - 1];\n C = prev_mult * C_model;\n }\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(start));\n }\n resblock_layer_backward(unet, &l_info, cublas_handle, C, C_time, OC, B, H, W, gn_n_groups, block_size, res_add_skip);\n if (t_res != nullptr) {\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n printf(\"down res back: l %d, r %d, Cin %d, Cout %d, H %d, W %d, time %f\\n\", l, r, C, OC, H, W, elapsed_time);\n *t_res += elapsed_time;\n }\n OC = C;\n }\n }\n\n convk3_layer_backward(unet, &l_info, cublas_handle, B, C_in, OC, H, W, block_size, 1);\n unet->dinp = l_info.input;\n\n // reset dout to point to emb_cum\n l_info.input = l_info.emb_cum;\n linear_layer_backward(unet, &l_info, cublas_handle, B, C_time, C_time);\n silu_layer_backward(unet, &l_info, B * C_time, block_size);\n linear_layer_backward(unet, &l_info, cublas_handle, B, C_model, C_time);\n}\n\n\n// ------------------------------------------------------------------------------------------------\n// functions to update gradient\n\nvoid unet_zero_grad(Unet* unet) {\n cudaCheck(cudaMemset(unet->grads.mem_ptr, 0, unet->grads.total_size * sizeof(float)));\n cudaCheck(cudaMemset(unet->back.back_mem, 0, unet->back.n_backs * sizeof(float)));\n //cudaCheck(cudaMemset(unet->acts.mem_ptr, 0, unet->acts.total_size * sizeof(float)));\n}\n\n// below copied from llm.c\n\n// Implements linear interpolation using only two floating-point operations (as opposed to three in a naive implementation).\n// Reference: https://developer.nvidia.com/blog/lerp-faster-cuda\n__device__ inline float lerp(float start, float end, float weight) {\n return fma(weight, end, fma(-weight, start, start));\n}\n\n__global__ void adamw_kernel2(float* params_memory, float* grads_memory, float* m_memory, float* v_memory, long num_parameters,\n float learning_rate, float beta1, float beta2, float beta1_correction, float beta2_correction, float eps, float weight_decay) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i >= num_parameters) return; // guard\n float grad = grads_memory[i];\n float m = m_memory[i];\n float v = v_memory[i];\n // update the first moment (momentum)\n m = lerp(grad, m, beta1);\n m_memory[i] = m;\n // update the second moment (RMSprop)\n v = lerp(grad * grad, v, beta2);\n v_memory[i] = v;\n m /= beta1_correction; // m_hat\n v /= beta2_correction; // v_hat\n params_memory[i] -= learning_rate * (m / (sqrtf(v) + eps) + weight_decay * params_memory[i]);\n}\n\nvoid unet_update(Unet* unet, float learning_rate, float beta1, float beta2, float eps, float weight_decay, int t) {\n // lazy allocation\n if (unet->m_memory == nullptr) {\n size_t param_size = unet->params.total_size * sizeof(float);\n cudaCheck(cudaMalloc(&(unet->m_memory), param_size));\n cudaCheck(cudaMalloc(&(unet->v_memory), param_size));\n cudaCheck(cudaMemset(unet->m_memory, 0, param_size));\n cudaCheck(cudaMemset(unet->v_memory, 0, param_size));\n printf(\"allocated %zu MiB for AdamW state m\\n\", param_size / 1024 / 1024);\n printf(\"allocated %zu MiB for AdamW state v\\n\", param_size / 1024 / 1024);\n }\n\n int block_size = unet->config.block_size;\n int n_blocks = ceil_div(unet->params.total_size, (size_t)block_size);\n float beta1_correction = 1.0f - powf(beta1, t);\n float beta2_correction = 1.0f - powf(beta2, t);\n adamw_kernel2<<>>(unet->params.mem_ptr, unet->grads.mem_ptr, unet->m_memory, unet->v_memory, unet->params.total_size,\n learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay);\n cudaCheck(cudaGetLastError());\n}\n\n\n// ------------------------------------------------------------------------------------------------\n// run tests\n\n#define NUM_DEBUG_STATES 5\ntypedef struct {\n float* input;\n float* timestep_emb;\n float* out;\n float* dinp;\n float* target;\n} DebugStates;\n\n\nint main(int argc, char **argv) {\n // setup cublas\n cublasHandle_t cublas_handle;\n cublasCheck(cublasCreate(&cublas_handle));\n cublasCheck(cublasSetMathMode(cublas_handle, CUBLAS_DEFAULT_MATH));\n\n char cwd[PATH_MAX];\n if (getcwd(cwd, sizeof(cwd)) != NULL) {\n //printf(\"Current working dir: %s\\n\", cwd);\n } else {\n perror(\"getcwd() error\");\n return 1;\n }\n\n // read header\n FILE *model_file = fopenCheck(\"unet_test_params.bin\", \"rb\");\n int model_header[256];\n freadCheck(model_header, sizeof(int), 256, model_file);\n if (model_header[0] != 12345678) {\n fprintf(stderr, \"Invalid magic model file\\n\");\n exit(1);\n }\n int B = model_header[1];\n int C_in = model_header[2];\n int C_model = model_header[3];\n int C_out = model_header[4];\n int H = model_header[5];\n int W = model_header[6];\n int max_period = model_header[7];\n int C_time = C_model * 4;\n const int n_levels = 4;\n int channel_mult[n_levels] = {1, 2, 3, 4};\n\n // dataloader\n DataLoader loader;\n dataloader_init(&loader, DATASET_FILE, B);\n\n\n Unet unet;\n UnetConfig* config = &(unet.config);\n config->B = B;\n config->C_in = C_in;\n config->C_model = C_model;\n config->C_time = C_time;\n config->C_out = C_out;\n config->H = H;\n config->W = W;\n config->block_size = 512;\n config->gn_n_groups = 32;\n config->n_levels = n_levels;\n config->channel_mult = channel_mult;\n config->HS = 32;\n config->n_res = 2;\n config->max_period = max_period;\n config->att_start_level = 2;\n\n unet_init(&unet);\n\n unet_count_layers(config);\n unet_make_ptrs_and_count_memory(&unet);\n malloc_and_point(&unet);\n\n // allocate host memory for params\n float* params_mem_cpu = (float*)mallocCheck(unet.params.total_size * sizeof(float));\n freadCheck(params_mem_cpu, sizeof(float), unet.params.total_size, model_file);\n fcloseCheck(model_file);\n UnetBlock params_cpu;\n unet_setup_and_count_params(config, ¶ms_cpu);\n params_cpu.mem_ptr = params_mem_cpu;\n point_unet_params(config, ¶ms_cpu);\n\n // copy params to device\n cudaCheck(cudaMemcpy(unet.params.mem_ptr, params_mem_cpu, unet.params.total_size * sizeof(float), cudaMemcpyHostToDevice));\n\n // allocate host memory for debug states\n FILE *debug_file = fopenCheck(\"unet_test_acts.bin\", \"rb\");\n size_t state_sizes[NUM_DEBUG_STATES];\n int out_size = B * C_out * H * W;\n state_sizes[0] = B * C_in * H * W; // input\n state_sizes[1] = B; // timestep emb\n state_sizes[2] = out_size; // out\n state_sizes[3] = B * C_in * H * W; // dinp\n state_sizes[4] = out_size; // target\n\n size_t num_debug_params = 0;\n for (int i = 0; i < NUM_DEBUG_STATES; i++) {\n num_debug_params += state_sizes[i];\n }\n\n float* debug_states_cpu = (float*)mallocCheck(num_debug_params * sizeof(float));\n freadCheck(debug_states_cpu, sizeof(float), num_debug_params, debug_file);\n // point to the debug states\n DebugStates debug_states;\n float** debug_ptrs[] = {\n &(debug_states.input),\n &(debug_states.timestep_emb),\n &(debug_states.out),\n &(debug_states.dinp),\n &(debug_states.target),\n };\n\n float* debug_mem_ptr = debug_states_cpu;\n for (int i = 0; i < NUM_DEBUG_STATES; i++) {\n *(debug_ptrs[i]) = debug_mem_ptr;\n debug_mem_ptr += state_sizes[i];\n }\n\n // read the groud truth grads\n float* grads_mem_cpu = (float*)mallocCheck(unet.grads.total_size * sizeof(float));\n freadCheck(grads_mem_cpu, sizeof(float), unet.grads.total_size, debug_file);\n fcloseCheck(debug_file);\n\n\n // read the losses\n float losses[10];\n FILE *loss_file = fopenCheck(\"unet_test_losses.bin\", \"rb\");\n freadCheck(losses, sizeof(float), 10, loss_file);\n fcloseCheck(loss_file);\n\n // create a grads struct and point to cpu grads for easier validation\n UnetBlock grads_cpu;\n unet_setup_and_count_params(config, &grads_cpu);\n grads_cpu.mem_ptr = grads_mem_cpu;\n point_unet_params(config, &grads_cpu);\n\n //// copy input to device\n //cudaCheck(cudaMemcpy(unet.input, debug_states.input, B * C_in * H * W * sizeof(float), cudaMemcpyHostToDevice));\n // load input\n dataloader_next_batch(&loader);\n cudaCheck(cudaMemcpy(unet.input, loader.input, B * C_in * H * W * sizeof(float), cudaMemcpyHostToDevice));\n\n // try out the random number generator\n GaussianDiffusion diffusion;\n diffusion_init(&diffusion, B, C_in, H, W, 1000);\n printf(\"successful gaussian init\\n\");\n\n float rand_nums[B * 10];\n sample_timesteps(&diffusion, rand_nums, B * 10);\n\n // compute timestep emb\n float* d_timesteps;\n cudaCheck(cudaMalloc(&d_timesteps, B * sizeof(float)));\n cudaMemcpy(d_timesteps, rand_nums, B * sizeof(float), cudaMemcpyHostToDevice);\n get_timestep_embeddings(&unet.time_emb_state, d_timesteps, unet.timestep_emb);\n\n // generate all the gaussian noise\n float *all_noise = (float*)malloc(10 * B * C_in * H * W * sizeof(float));\n normal_(all_noise, 10 * B * C_in * H * W, 0.0f, 1.0f, &diffusion.host_rng);\n\n\n cudaCheck(cudaMalloc(&(unet.target), out_size * sizeof(float)));\n\n //printf(\"drawing and saving x_t for verification\\n\");\n //float* x_t_cpu = (float*)mallocCheck(B * C_in * H * W * sizeof(float));\n //diffusion_draw_and_sample_xt(&diffusion, unet.target, d_timesteps, unet.input);\n //print_gpu_first_5(\"x_t\", unet.target);\n //cudaCheck(cudaMemcpy(x_t_cpu, unet.target, B * C_in * H * W * sizeof(float), cudaMemcpyDeviceToHost));\n //FILE *sample_debug_file = fopenCheck(\"./sample_debug.bin\", \"wb\");\n //size_t result = fwrite(rand_nums, sizeof(float), B, sample_debug_file);\n //if (result != B) {\n // fprintf(stderr, \"Error writing to file\\n\");\n // exit(1);\n //}\n //result = fwrite(loader.input, sizeof(float), B * C_in * H * W, sample_debug_file);\n //if (result != B * C_in * H * W) {\n // fprintf(stderr, \"Error writing to file\\n\");\n // exit(1);\n //}\n //result = fwrite(x_t_cpu, sizeof(float), B * C_in * H * W, sample_debug_file);\n //if (result != B * C_in * H * W) {\n // fprintf(stderr, \"Error writing to file\\n\");\n // exit(1);\n //}\n //fcloseCheck(sample_debug_file);\n\n\n // copy the targets to gpu\n //cudaCheck(cudaMemcpy(unet.target, debug_states.target, out_size * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(unet.target, all_noise, out_size * sizeof(float), cudaMemcpyHostToDevice));\n diffusion_sample(d_timesteps, unet.input, unet.target, &diffusion, B, C_in * H * W);\n //diffusion_draw_and_sample_xt(&diffusion, unet.target, d_timesteps, unet.input);\n\n\n\n printf(\"\\nCheck that unet targets match the debug states\\n\");\n validate_result(unet.target, debug_states.target, \"target\", out_size);\n\n\n float acc = 1e-4;\n float med_acc = 1e-2;\n printf(\"Running 10 update loops\\n\");\n // loop 1\n printf(\"Checking forward pass\\n\");\n float t_res = 0.0f;\n float t_att = 0.0f;\n unet_forward(cublas_handle, &unet, &t_res, &t_att);\n printf(\"resblock time: %f ms, attention time: %f ms\\n\", t_res, t_att);\n printf(\"\\nChecking out\\n\");\n validate_result(unet.output, debug_states.out, \"out\", out_size, med_acc);\n printf(\"cpu loss: %f, gpu loss: %f\\n\", losses[0], unet.mean_loss);\n printf(\"\\nForward pass successful\\n\");\n\n printf(\"\\nChecking backward pass\\n\");\n t_res = 0.0f;\n t_att = 0.0f;\n unet_backward(cublas_handle, &unet, &t_res, &t_att);\n printf(\"resblock back time: %f ms, attention back time: %f ms\\n\", t_res, t_att);\n printf(\"\\nChecking grads\\n\");\n validate_result(unet.grads.mem_ptr, grads_mem_cpu, \"grads\", unet.grads.total_size, med_acc);\n printf(\"\\nChecking dinp\\n\");\n validate_result(unet.dinp, debug_states.dinp, \"dinp\", B * C_in * H * W, acc);\n unet_update(&unet, 1e-3f, 0.9f, 0.999f, 1e-8f, 0.01f, 1);\n unet_zero_grad(&unet);\n\n // loop 2-9\n cudaEvent_t start, end;\n cudaCheck(cudaEventCreate(&start));\n cudaCheck(cudaEventCreate(&end));\n cudaCheck(cudaProfilerStart());\n cudaCheck(cudaEventRecord(start));\n for (int i = 1; i < 10; i++) {\n dataloader_next_batch(&loader);\n cudaCheck(cudaMemcpy(unet.input, loader.input, B * C_in * H * W * sizeof(float), cudaMemcpyHostToDevice));\n cudaMemcpy(d_timesteps, rand_nums + i * B, B * sizeof(float), cudaMemcpyHostToDevice);\n get_timestep_embeddings(&unet.time_emb_state, d_timesteps, unet.timestep_emb);\n\n cudaCheck(cudaMemcpy(unet.target, all_noise + i * out_size, out_size * sizeof(float), cudaMemcpyHostToDevice));\n diffusion_sample(d_timesteps, unet.input, unet.target, &diffusion, B, C_in * H * W);\n //diffusion_draw_and_sample_xt(&diffusion, unet.target, d_timesteps, unet.input);\n\n unet_forward(cublas_handle, &unet);\n unet_backward(cublas_handle, &unet);\n printf(\"step %d, cpu loss: %f, gpu loss: %f\\n\", i+1, losses[i], unet.mean_loss);\n unet_update(&unet, 1e-3f, 0.9f, 0.999f, 1e-8f, 0.01f, i+1);\n unet_zero_grad(&unet);\n }\n cudaCheck(cudaEventRecord(end));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float time_elapsed;\n cudaCheck(cudaEventElapsedTime(&time_elapsed, start, end));\n time_elapsed /= 9;\n printf(\"Average time per step: %f ms\\n\", time_elapsed);\n\n\n\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmark\\n\");\n int repeat_times = 50;\n float elapsed_time = benchmark_kernel(repeat_times, unet_forward, cublas_handle, &unet, nullptr, nullptr);\n printf(\"Forward pass time: %f ms\\n\", elapsed_time);\n\n printf(\"\\nBackward pass benchmark\\n\");\n elapsed_time = benchmark_kernel(repeat_times, unet_backward, cublas_handle, &unet, nullptr, nullptr);\n printf(\"Backward pass time: %f ms\\n\", elapsed_time);\n\n // free memory\n free(params_mem_cpu);\n free(debug_states_cpu);\n free_unet(&unet);\n cudaCheck(cudaFree(unet.target));\n cudaCheck(cudaFree(d_timesteps));\n dataloader_free(&loader);\n diffusion_free(&diffusion);\n free(all_noise);\n\n // free cublas\n cublasCheck(cublasDestroy(cublas_handle));\n}", "output": "#include \n#include \n#include // For POSIX systems to use getcwd\n#include // For PATH_MAX\n#include \n#include \n#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \"linear.cuh\"\n#include \"silu.cuh\"\n#include \"conv2d_k3.cuh\"\n#include \"resblock.cuh\"\n#include \"add.cuh\"\n#include \"avgpool.cuh\"\n#include \"attention_block.cuh\"\n#include \"concat_channel.cuh\"\n#include \"add.cuh\"\n#include \"upsample.cuh\"\n#include \"groupnorm.cuh\"\n#include \"mse.cuh\"\n#include \"timestep_embedding.cuh\"\n#include \"rand.h\"\n#include \n\n\n\ntypedef struct {\n int C_in;\n int C_model;\n int C_time;\n int C_out;\n int B;\n int H;\n int W;\n int gn_n_groups;\n int HS;\n int block_size;\n int n_layers;\n int n_levels;\n int n_res;\n int att_start_level;\n int* channel_mult;\n int max_period;\n} UnetConfig;\n\n// ------------------------------------------------------------------------------------------------\n// DataLoader\n\nconst char* DATASET_FILE = \"../datasets/elephant_train.bin\";\ntypedef struct {\n int B;\n int C_in;\n int H;\n int W;\n int img_size;\n int n_imgs;\n FILE* img_file;\n long cur_position;\n long file_size;\n long header_offset;\n long num_batches;\n float* input;\n} DataLoader;\n\nvoid dataloader_init(DataLoader* loader, const char* filename, int B) {\n loader->B = B;\n\n loader->img_file = fopenCheck(filename, \"rb\");\n int data_header[256];\n loader->header_offset = 256 * sizeof(int);\n freadCheck(data_header, sizeof(int), 256, loader->img_file);\n if (data_header[0] != 20240620) {\n fprintf(stderr, \"Invalid magic data file\\n\");\n exit(1);\n }\n loader->n_imgs = data_header[1];\n loader->C_in = data_header[2];\n loader->H = data_header[3];\n loader->W = data_header[4];\n loader->img_size = loader->C_in * loader->H * loader->W;\n\n // determine file size\n fseekCheck(loader->img_file, 0, SEEK_SET);\n fseekCheck(loader->img_file, 0, SEEK_END);\n loader->file_size = ftell(loader->img_file);\n fseekCheck(loader->img_file, 0, SEEK_SET);\n assert((loader->file_size - loader->header_offset) / (loader->img_size * sizeof(float)) == loader->n_imgs);\n \n loader->cur_position = loader->header_offset;\n loader->num_batches = loader->n_imgs / B;\n\n // allocate cuda CPU pinned memory for inputs\n hipHostMalloc((void**)&loader->input, B * loader->img_size * sizeof(float));\n}\n\nvoid dataloader_reset(DataLoader* loader) {\n loader->cur_position = loader->header_offset;\n}\n\nvoid dataloader_next_batch(DataLoader* loader) {\n int B = loader->B;\n int img_size = loader->img_size;\n if (loader->cur_position + (B * img_size * sizeof(float)) > loader->file_size) {\n dataloader_reset(loader);\n }\n fseekCheck(loader->img_file, loader->cur_position, SEEK_SET);\n freadCheck(loader->input, sizeof(float), B * img_size, loader->img_file);\n loader->cur_position += B * img_size * sizeof(float);\n}\n\nvoid dataloader_free(DataLoader* loader) {\n fcloseCheck(loader->img_file);\n hipHostFree(loader->input);\n}\n\n// ------------------------------------------------------------------------------------------------\n// Diffusion\ntypedef struct {\n int B;\n int img_size;\n int max_period;\n float* betas; // (max_period,)\n float* sqrt_alphas_cumprod; // (max_period,)\n float* sqrt_one_minus_alphas_cumprod; // (max_period,)\n int block_size;\n hiprandState* curand_states;\n mt19937_state host_rng;\n} GaussianDiffusion;\n\n__global__ void init_curand_states(hiprandState *state, unsigned long seed, int N) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < N) {\n hiprand_init(seed, idx, 0, &state[idx]);\n }\n}\n\nvoid diffusion_init(\n GaussianDiffusion* diffusion, int B, int C_in, int H, int W, int max_period, int block_size = 512,\n unsigned long seed = 0, float beta_start = 1e-4, float beta_end = 0.02f\n) {\n diffusion->B = B;\n diffusion->max_period = max_period;\n diffusion->block_size = block_size;\n diffusion->img_size = C_in * H * W;\n\n float scale = 1000.0f / max_period;\n beta_start *= scale;\n beta_end *= scale;\n\n // calculate betas and alphas on host then move to device\n float* betas = (float*)mallocCheck(max_period * sizeof(float));\n float* sqrt_alphas_cumprod = (float*)mallocCheck(max_period * sizeof(float));\n float* sqrt_one_minus_alphas_cumprod = (float*)mallocCheck(max_period * sizeof(float));\n\n float alpha_cumprod = 1.0f;\n for (int t = 1; t <= max_period; t++) {\n // linear interpolate for beta\n betas[t - 1] = (beta_start * (max_period - t) + beta_end * (t - 1)) / (max_period - 1);\n alpha_cumprod *= 1.0f - betas[t - 1];\n sqrt_alphas_cumprod[t - 1] = sqrtf(alpha_cumprod);\n sqrt_one_minus_alphas_cumprod[t - 1] = sqrtf(1.0f - alpha_cumprod);\n }\n\n float *d_betas, *d_sqrt_alphas_cumprod, *d_sqrt_one_minus_alphas_cumprod;\n cudaCheck(hipMalloc(&d_betas, max_period * sizeof(float)));\n cudaCheck(hipMalloc(&d_sqrt_alphas_cumprod, max_period * sizeof(float)));\n cudaCheck(hipMalloc(&d_sqrt_one_minus_alphas_cumprod, max_period * sizeof(float)));\n\n cudaCheck(hipMemcpy(d_betas, betas, max_period * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_sqrt_alphas_cumprod, sqrt_alphas_cumprod, max_period * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_sqrt_one_minus_alphas_cumprod, sqrt_one_minus_alphas_cumprod, max_period * sizeof(float), hipMemcpyHostToDevice));\n\n free(betas);\n free(sqrt_alphas_cumprod);\n free(sqrt_one_minus_alphas_cumprod);\n\n diffusion->betas = d_betas;\n diffusion->sqrt_alphas_cumprod = d_sqrt_alphas_cumprod;\n diffusion->sqrt_one_minus_alphas_cumprod = d_sqrt_one_minus_alphas_cumprod;\n\n // init curand states for sampling noise of the size of one batch of inputs\n int total_states = B * C_in * H * W;\n int n_blocks = ceil_div(total_states, block_size);\n cudaCheck(hipMalloc(&diffusion->curand_states, total_states * sizeof(hiprandState)));\n init_curand_states<<>>(diffusion->curand_states, seed, total_states);\n\n // init host rng for drawing random timesteps\n manual_seed(&diffusion->host_rng, seed);\n}\n\n// gridDim will be of shape (N / block_size, B)\n// expect x and x_t to be of shape (B, N)\n// and timesteps to be of shape (B,)\n// expect x_t to be filled with standard normals\n__global__ void diffusion_sample_kernel(\n float* timesteps, float* x, float* x_t,\n float* sqrt_alphas_cumprod, float* sqrt_one_minus_alphas_cumprod,\n int B, int N\n) {\n int b = blockIdx.y;\n int n = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (n < N && b < B) {\n int t = (int)timesteps[b];\n float sqrt_alpha_t = sqrt_alphas_cumprod[t];\n float sqrt_one_minus_alpha_t = sqrt_one_minus_alphas_cumprod[t];\n float x_t_val = x_t[b * N + n]; // x_t starts out as standard normals\n float x_val = x[b * N + n];\n x_t[b * N + n] = x_val * sqrt_alpha_t + x_t_val * sqrt_one_minus_alpha_t;\n }\n}\n\nvoid diffusion_sample(\n float* timesteps, float* x, float* x_t,\n GaussianDiffusion *diffusion,\n int B, int N\n) {\n dim3 gridDim(ceil_div(N, diffusion->block_size), B);\n diffusion_sample_kernel<<block_size>>>(timesteps, x, x_t, diffusion->sqrt_alphas_cumprod, diffusion->sqrt_one_minus_alphas_cumprod, B, N);\n cudaCheck(hipDeviceSynchronize());\n}\n\n// a kernel that creates x_t = sqrt_alpha * x + sqrt(1 - alpha) * noise in one go\n// that draws the noise from standard normals, gets alpha from the timesteps and the precomputed alpha cumprod\n// and saves output to x_t\n// if the input x, timesteps and the alpha buffers are null, then it will just fill x_t with standard normals\n// to be called with B * N total threads\n__global__ void diffusion_draw_and_sample_xt_kernel(\n float* x_t, int B, int N, hiprandState* curand_states,\n float* sqrt_alphas_cumprod, float* sqrt_one_minus_alphas_cumprod,\n float* timesteps = nullptr, float* x = nullptr\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * N) {\n float out_val = hiprand_normal(&curand_states[idx]);\n if (timesteps != nullptr && x != nullptr) {\n int b = idx / N;\n int n = idx % N;\n float x_val = x[b * N + n];\n int t = (int)timesteps[b];\n float sqrt_alpha_t = sqrt_alphas_cumprod[t];\n float sqrt_one_minus_alpha_t = sqrt_one_minus_alphas_cumprod[t];\n out_val = x_val * sqrt_alpha_t + out_val * sqrt_one_minus_alpha_t;\n }\n x_t[idx] = out_val;\n }\n}\n\nvoid diffusion_draw_and_sample_xt(\n GaussianDiffusion *diffusion,\n float* x_t,\n float* timesteps = nullptr, float* x = nullptr\n) {\n int total_states = diffusion->B * diffusion->img_size;\n int n_blocks = ceil_div(total_states, diffusion->block_size);\n diffusion_draw_and_sample_xt_kernel<<block_size>>>(\n x_t, diffusion->B, diffusion->img_size, diffusion->curand_states, diffusion->sqrt_alphas_cumprod, diffusion->sqrt_one_minus_alphas_cumprod, timesteps, x\n );\n}\n\nvoid diffusion_free(GaussianDiffusion* diffusion) {\n cudaCheck(hipFree(diffusion->betas));\n cudaCheck(hipFree(diffusion->sqrt_alphas_cumprod));\n cudaCheck(hipFree(diffusion->sqrt_one_minus_alphas_cumprod));\n cudaCheck(hipFree(diffusion->curand_states));\n}\n\n// samples B timesteps in range [0, max_period), convert them to floats and\n// store in timesteps\nvoid sample_timesteps(GaussianDiffusion *diffusion, float* timesteps, int N = -1) {\n mt19937_state* rng = &diffusion->host_rng;\n if (N < 0) {\n N = diffusion->B;\n }\n int max_period = diffusion->max_period;\n for (int i = 0; i < N; i++) {\n float rand_float = randfloat32(rng);\n timesteps[i] = (float)(int)(rand_float * max_period);\n }\n}\n\n// ------------------------------------------------------------------------------------------------\n// Unet init\n\n// we just need a single set of buffers that are shared for all the backward pass activations\n#define NUM_SKIP_CONNECTIONS 12 // (n_res + 1) * n_levels\n#define NUM_BACK_BUFS 10\n#define UNET_NUM_BACK_ACTS 22\ntypedef struct {\n float* dinp_buf1; // used to store the input gradient, which will be reused as dout for the previous layer\n float* dinp_buf2;\n float* buf1_BCHW; // NOTE: first 4 buffers will have max B * C * H * W size\n float* buf2_BCHW;\n float* buf3_BCHW;\n float* buf4_BCHW;\n float* buf1_B3CHW; // NOTE: last 2 buffers will have max B * 3 * C * H * W size for attention layers\n float* buf2_B3CHW; // so max of H * W is only over dimensions that have attention\n float* demb_buf1; // used to store cumulative gradients w.r.t. timestep embeddings\n float* demb_buf2;\n float* skip_dinps[NUM_SKIP_CONNECTIONS];\n} UnetBackActs;\n\n// This is used as a block to store all the params / grads / acts for all the layers\ntypedef struct {\n float* mem_ptr;\n size_t total_size;\n int n_layers;\n void** layer_ptrs;\n size_t* layer_sizes;\n} UnetBlock;\n\ntypedef struct {\n float* back_mem;\n UnetBackActs back_acts;\n size_t back_sizes[UNET_NUM_BACK_ACTS];\n size_t n_backs;\n} UnetBackBlock;\n\ntypedef struct {\n UnetConfig config;\n float* input;\n float* timestep_emb; // this is the already computed embeddings of shape (B, C_model)\n UnetBlock params;\n UnetBlock grads;\n UnetBlock acts;\n UnetBackBlock back;\n float* skip_acts[NUM_SKIP_CONNECTIONS]; // pointers to skip connetion activations\n // the next two are for debug purposes only, just to keep track of the buffer that contains the grad w.r.t. inputs\n float* dinp;\n // model output\n float* output;\n // targets and losses\n float* target;\n float mean_loss;\n float* gpu_loss;\n // buffers for AdamW\n float* m_memory;\n float* v_memory;\n TimestepEmbedding time_emb_state;\n} Unet;\n\nvoid unet_init(Unet* unet) {\n unet->input = nullptr;\n unet->timestep_emb = nullptr;\n unet->params.mem_ptr = nullptr;\n unet->grads.mem_ptr = nullptr;\n unet->acts.mem_ptr = nullptr;\n unet->back.back_mem = nullptr;\n unet->dinp = nullptr;\n unet->output = nullptr;\n unet->target = nullptr;\n unet->mean_loss = -1.0f;\n unet->gpu_loss = nullptr;\n unet->m_memory = nullptr;\n unet->v_memory = nullptr;\n\n init_timestep_embedding(&unet->time_emb_state, unet->config.C_model, unet->config.B, unet->config.max_period);\n}\n\n\ntemplate \nvoid free_not_null(T** ptr) {\n if (*ptr != nullptr) {\n free(*ptr);\n *ptr = nullptr;\n }\n}\n\nvoid free_unet_block(UnetBlock* params) {\n free_not_null(¶ms->layer_sizes);\n for (int l = 0; l < params->n_layers; l++) {\n free_not_null(&(params->layer_ptrs[l]));\n }\n free_not_null(¶ms->layer_ptrs);\n cudaCheck(hipFree(params->mem_ptr));\n}\n\nvoid free_unet(Unet* unet) {\n cudaCheck(hipFree(unet->input));\n cudaCheck(hipFree(unet->timestep_emb));\n cudaCheck(hipFree(unet->gpu_loss));\n cudaCheck(hipFree(unet->m_memory));\n cudaCheck(hipFree(unet->v_memory));\n\n free_unet_block(&(unet->params));\n free_unet_block(&(unet->grads));\n free_unet_block(&(unet->acts));\n\n cudaCheck(hipFree(unet->back.back_mem));\n free_timestep_embedding(&unet->time_emb_state);\n}\n\nvoid unet_count_layers(UnetConfig* config) {\n int n_layers = 0;\n int n_res = config->n_res;\n int n_levels = config->n_levels;\n int att_start_level = config->att_start_level;\n\n // timestep emb\n n_layers += 3;\n\n // initial conv\n n_layers += 1;\n\n // add downsample resblock layers\n n_layers += n_res * n_levels;\n\n // add downsample attention\n n_layers += n_res * (n_levels - att_start_level);\n\n // add avgpool layers\n n_layers += n_levels - 1;\n\n // add middle layers\n n_layers += 3;\n\n // add upsample concat layers\n n_layers += (n_res + 1) * n_levels;\n\n // add upsample resblock layers\n n_layers += (n_res + 1) * n_levels;\n\n // add upsample attention layers\n n_layers += (n_res + 1) * (n_levels - att_start_level);\n\n // add upsample layers\n n_layers += n_levels - 1;\n\n // output layers\n n_layers += 3;\n\n config->n_layers = n_layers;\n}\n\n\n// ------------------------------------------------------------------------------------------------\n// allocate memory for structs with info for each layer and count params\n\n\nint make_linear_params_ptrs(\n UnetBlock* params,\n int layer,\n int C_in, int C_out\n) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(LinearParams));\n params->layer_sizes[layer] = linear_count_params(C_in, C_out);\n return layer + 1;\n}\n\nint make_silu_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_convk3_params_ptrs(\n UnetBlock* params, int layer,\n int C_in, int OC, int B, int H, int W\n) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ConvK3Params));\n params->layer_sizes[layer] = convk3_count_params(C_in, OC);\n return layer + 1;\n}\n\nint make_resblock_params_ptrs(\n UnetBlock* params, int layer,\n int C, int C_emb, int C_out, int B, int H, int W, int gn_n_groups\n) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ResBlockParameters));\n ResBlockParameters* resblock_params = (ResBlockParameters*)params->layer_ptrs[layer];\n resblock_count_params(resblock_params, C, C_emb, C_out, B, H, W, 0, 0, gn_n_groups);\n params->layer_sizes[layer] = resblock_params->n_params;\n return layer + 1;\n}\n\nint make_avgpool_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_attention_block_params_ptrs(UnetBlock* params, int layer, int C) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(AttentionParams));\n AttentionParams* attention_params = (AttentionParams*)params->layer_ptrs[layer];\n attention_block_count_params(attention_params, C);\n params->layer_sizes[layer] = attention_params->n_params;\n return layer + 1;\n}\n\nint make_concat_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_upsample_params_ptrs(UnetBlock* params, int layer) {\n params->layer_ptrs[layer] = nullptr;\n params->layer_sizes[layer] = 0;\n return layer + 1;\n}\n\nint make_groupnorm_params_ptrs(UnetBlock* params, int layer, int C) {\n params->layer_ptrs[layer] = (void*)mallocCheck(sizeof(GroupNormParams));\n params->layer_sizes[layer] = C * 2;\n return layer + 1;\n}\n\n\nvoid unet_setup_and_count_params(UnetConfig* config, UnetBlock* params) {\n assert (config->n_layers > 0);\n\n params->n_layers = config->n_layers;\n params->layer_ptrs = (void**)mallocCheck(config->n_layers * sizeof(void*));\n params->layer_sizes = (size_t*)mallocCheck(config->n_layers * sizeof(size_t));\n\n int B = config->B;\n int C_in = config->C_in;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int C_out = config->C_out;\n int H = config->H;\n int W = config->W;\n int gn_n_groups = config->gn_n_groups;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n int layer = 0;\n layer = make_linear_params_ptrs(params, layer, C_model, C_time);\n layer = make_silu_params_ptrs(params, layer);\n layer = make_linear_params_ptrs(params, layer, C_time, C_time);\n int OC = C_model;\n layer = make_convk3_params_ptrs(params, layer, C_in, OC, B, H, W);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n layer = make_resblock_params_ptrs(params, layer, C, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_params_ptrs(params, layer, C);\n }\n }\n if (l < n_levels - 1) {\n layer = make_avgpool_params_ptrs(params, layer);\n H /= 2;\n W /= 2;\n }\n }\n\n // middle layers\n layer = make_resblock_params_ptrs(params, layer, C, C_time, C, B, H, W, gn_n_groups);\n layer = make_attention_block_params_ptrs(params, layer, C);\n layer = make_resblock_params_ptrs(params, layer, C, C_time, C, B, H, W, gn_n_groups);\n\n // upwards: concat then resblock\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = C_model * mult;\n for (int r = 0; r < n_res + 1; r++) {\n if (r == n_res && l > 0) {\n int prev_mult = config->channel_mult[l - 1];\n C_skip = C_model * prev_mult;\n }\n OC = C_model * mult;\n layer = make_concat_params_ptrs(params, layer);\n layer = make_resblock_params_ptrs(params, layer, C + C_skip, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_params_ptrs(params, layer, C);\n }\n }\n if (l > 0) {\n layer = make_upsample_params_ptrs(params, layer);\n H *= 2;\n W *= 2;\n }\n }\n\n layer = make_groupnorm_params_ptrs(params, layer, C);\n layer = make_silu_params_ptrs(params, layer);\n layer = make_convk3_params_ptrs(params, layer, C, C_out, B, H, W);\n\n // count total size\n size_t total_size = 0;\n for (int i = 0; i < config->n_layers; i++) {\n total_size += params->layer_sizes[i];\n }\n params->total_size = total_size;\n}\n\nint make_linear_acts_ptrs(\n UnetBlock* acts,\n int layer,\n int B, int C_out\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(LinearActs));\n // acts buffers only hold outputs\n acts->layer_sizes[layer] = B * C_out;\n return layer + 1;\n}\n\nint make_silu_acts_ptrs(UnetBlock* acts, int layer, int N) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(SiluActs));\n acts->layer_sizes[layer] = N;\n return layer + 1;\n}\n\nint make_convk3_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int OC, int H, int W\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ConvK3Acts));\n acts->layer_sizes[layer] = B * OC * H * W;\n return layer + 1;\n}\n\nint make_resblock_acts_ptrs(\n UnetBlock* acts, int layer,\n int C, int C_emb, int C_out, int B, int H, int W, int gn_n_groups\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ResBlockActivations));\n ResBlockActivations* resblock_acts = (ResBlockActivations*)acts->layer_ptrs[layer];\n resblock_count_acts(resblock_acts, C, C_emb, C_out, B, H, W, 0, 0, gn_n_groups);\n acts->layer_sizes[layer] = resblock_acts->n_acts;\n return layer + 1;\n}\n\nint make_avgpool_acts_ptrs(UnetBlock* acts, int layer, int B, int C, int H, int W) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(AvgpoolActs));\n acts->layer_sizes[layer] = B * C * H/2 * W/2;\n return layer + 1;\n}\n\nint make_attention_block_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C, int H, int W, int gn_n_groups, int HS\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(AttentionActs));\n AttentionActs* att_acts = (AttentionActs*)acts->layer_ptrs[layer];\n attention_block_count_acts(att_acts, B, C, H, W, gn_n_groups, HS);\n acts->layer_sizes[layer] = att_acts->n_acts;\n return layer + 1;\n}\n\nint make_concat_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C1, int C2, int H, int W\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(ConcatChannelActs));\n acts->layer_sizes[layer] = B * (C1 + C2) * H * W;\n return layer + 1;\n}\n\nint make_upsample_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C, int H, int W\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(UpsampleActs));\n acts->layer_sizes[layer] = B * C * H * 2 * W * 2;\n return layer + 1;\n}\n\nint make_groupnorm_acts_ptrs(\n UnetBlock* acts, int layer,\n int B, int C, int H, int W, int gn_n_groups\n) {\n acts->layer_ptrs[layer] = (void*)mallocCheck(sizeof(GroupNormActs));\n // out (B, C, H, W), mean and rstd (B, gn_n_groups)\n acts->layer_sizes[layer] = B * C * H * W + B * gn_n_groups * 2;\n return layer + 1;\n}\n\nvoid unet_setup_and_count_acts(UnetConfig* config, UnetBlock* acts) {\n assert (config->n_layers > 0);\n\n acts->n_layers = config->n_layers;\n acts->layer_ptrs = (void**)mallocCheck(config->n_layers * sizeof(void*));\n acts->layer_sizes = (size_t*)mallocCheck(config->n_layers * sizeof(size_t));\n\n int B = config->B;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int C_out = config->C_out;\n int H = config->H;\n int W = config->W;\n int gn_n_groups = config->gn_n_groups;\n int HS = config->HS;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n int layer = 0;\n layer = make_linear_acts_ptrs(acts, layer, B, C_time);\n layer = make_silu_acts_ptrs(acts, layer, B * C_time);\n layer = make_linear_acts_ptrs(acts, layer, B, C_time);\n int OC = C_model;\n layer = make_convk3_acts_ptrs(acts, layer, B, OC, H, W);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n layer = make_resblock_acts_ptrs(acts, layer, C, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups, HS);\n }\n }\n if (l < n_levels - 1) {\n layer = make_avgpool_acts_ptrs(acts, layer, B, C, H, W);\n H /= 2;\n W /= 2;\n }\n }\n\n layer = make_resblock_acts_ptrs(acts, layer, C, C_time, C, B, H, W, gn_n_groups);\n layer = make_attention_block_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups, HS);\n layer = make_resblock_acts_ptrs(acts, layer, C, C_time, C, B, H, W, gn_n_groups);\n\n // upwards\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = mult * C_model;\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n int prev_mult = config->channel_mult[l - 1];\n C_skip = prev_mult * C_model;\n }\n layer = make_concat_acts_ptrs(acts, layer, B, C, C_skip, H, W);\n OC = mult * C_model;\n layer = make_resblock_acts_ptrs(acts, layer, C + C_skip, C_time, OC, B, H, W, gn_n_groups);\n C = OC;\n if (l >= att_start_level) {\n layer = make_attention_block_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups, HS);\n }\n }\n if (l > 0) {\n layer = make_upsample_acts_ptrs(acts, layer, B, C, H, W);\n H *= 2;\n W *= 2;\n }\n }\n layer = make_groupnorm_acts_ptrs(acts, layer, B, C, H, W, gn_n_groups);\n layer = make_silu_acts_ptrs(acts, layer, B * C * H * W);\n layer = make_convk3_acts_ptrs(acts, layer, B, C_out, H, W);\n\n // count total size\n size_t total_size = 0;\n for (int i = 0; i < config->n_layers; i++) {\n total_size += acts->layer_sizes[i];\n }\n acts->total_size = total_size;\n}\n\nvoid unet_setup_and_count_backs(UnetConfig* config, UnetBackBlock* back) {\n int B = config->B;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int H = config->H;\n int W = config->W;\n\n size_t max_BCHW = 3 * B * C_model * H * W;\n size_t max_B3CHW = 2 * B * 3 * C_model * H * W;\n back->back_sizes[0] = max_BCHW;\n back->back_sizes[1] = max_BCHW;\n back->back_sizes[2] = max_BCHW;\n back->back_sizes[3] = max_BCHW;\n back->back_sizes[4] = max_BCHW;\n back->back_sizes[5] = max_BCHW;\n back->back_sizes[6] = max_B3CHW;\n back->back_sizes[7] = max_B3CHW;\n back->back_sizes[8] = B * C_time;\n back->back_sizes[9] = B * C_time;\n\n // fill the sizes of the skip connection backward acts\n int C = C_model;\n for (int l = 0; l < config->n_levels; l++) {\n for (int r = 0; r < (config->n_res + 1); r++) {\n int s = l * 3 + r + NUM_BACK_BUFS;\n assert(s < UNET_NUM_BACK_ACTS);\n back->back_sizes[s] = B * C * H * W;\n int mult = config->channel_mult[l];\n C = C_model * mult;\n if (l > 0 && r == 0) {\n H /= 2;\n W /= 2;\n }\n }\n }\n\n size_t n_backs = 0;\n for (int i = 0; i < UNET_NUM_BACK_ACTS; i++) {\n n_backs += back->back_sizes[i];\n }\n back->n_backs = n_backs;\n}\n\n\nvoid unet_make_ptrs_and_count_memory(Unet* unet)\n{\n unet_setup_and_count_params(&(unet->config), &(unet->params));\n unet_setup_and_count_params(&(unet->config), &(unet->grads));\n unet_setup_and_count_acts(&(unet->config), &(unet->acts));\n unet_setup_and_count_backs(&(unet->config), &(unet->back));\n}\n\n// ------------------------------------------------------------------------------------------------\n// malloc unet memory and set pointers\n\ntypedef struct {\n float* input;\n float* prev_out;\n int layer;\n float* emb;\n float* emb_cum; // cumulative demb\n float* mem_ptr;\n int skip_idx;\n} LayerInfo;\n\n\nvoid set_linear_params_ptrs(\n UnetBlock* params,\n LayerInfo* l_info,\n int C_in, int C_out\n) {\n int layer = l_info->layer;\n LinearParams* linear_params = (LinearParams*)params->layer_ptrs[layer];\n linear_set_param_ptrs(linear_params, l_info->mem_ptr, C_in, C_out);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_convk3_params_ptrs(UnetBlock* params, LayerInfo* l_info, int C_in, int OC) {\n int layer = l_info->layer;\n ConvK3Params* convk3_params = (ConvK3Params*)params->layer_ptrs[layer];\n convk3_set_param_ptrs(convk3_params, l_info->mem_ptr, C_in, OC);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_resblock_layer_params_ptrs(UnetBlock* params, LayerInfo* l_info, int C, int OC) {\n int layer = l_info->layer;\n ResBlockParameters* resblock_params = (ResBlockParameters*)params->layer_ptrs[layer];\n set_resblock_params_ptrs(C, OC, resblock_params, l_info->mem_ptr);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_attention_layer_params_ptrs(UnetBlock* params, LayerInfo* l_info) {\n int layer = l_info->layer;\n AttentionParams* att_params = (AttentionParams*)params->layer_ptrs[layer];\n set_attention_params_pointers(att_params, l_info->mem_ptr);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\nvoid set_groupnorm_layer_params_ptrs(UnetBlock* params, LayerInfo* l_info, int C) {\n int layer = l_info->layer;\n GroupNormParams* gn_params = (GroupNormParams*)params->layer_ptrs[layer];\n gn_set_param_ptrs(gn_params, l_info->mem_ptr, C);\n l_info->mem_ptr += params->layer_sizes[layer];\n l_info->layer++;\n}\n\n\nvoid point_unet_params(UnetConfig* config, UnetBlock* params) {\n // common constants\n int C_in = config->C_in;\n int C_model = config->C_model;\n int C_time = config->C_time;\n int C_out = config->C_out;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n LayerInfo l_info;\n l_info.layer = 0;\n l_info.mem_ptr = params->mem_ptr;\n set_linear_params_ptrs(params, &l_info, C_model, C_time);\n l_info.layer++; // skip silu layer\n set_linear_params_ptrs(params, &l_info, C_time, C_time);\n int OC = C_model;\n set_convk3_params_ptrs(params, &l_info, C_in, OC);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n set_resblock_layer_params_ptrs(params, &l_info, C, OC);\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_params_ptrs(params, &l_info);\n }\n }\n if (l < n_levels - 1) {\n l_info.layer++; // skip avgpool layer\n }\n }\n\n set_resblock_layer_params_ptrs(params, &l_info, C, C);\n set_attention_layer_params_ptrs(params, &l_info);\n set_resblock_layer_params_ptrs(params, &l_info, C, C);\n\n // upward: skip one layer for concat\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = mult * C_model;\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config->channel_mult[l - 1];\n }\n l_info.layer++;\n OC = mult * C_model;\n set_resblock_layer_params_ptrs(params, &l_info, C + C_skip, OC);\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_params_ptrs(params, &l_info);\n }\n }\n if (l > 0) {\n l_info.layer++; // upsample\n }\n }\n\n set_groupnorm_layer_params_ptrs(params, &l_info, C);\n l_info.layer++; // silu\n set_convk3_params_ptrs(params, &l_info, C, C_out);\n}\n\nvoid set_linear_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n LinearActs* linear_acts = (LinearActs*)acts->layer_ptrs[layer];\n linear_acts->out = l_info->mem_ptr;\n l_info->mem_ptr += acts->layer_sizes[layer];\n linear_acts->inp = l_info->input;\n l_info->input = linear_acts->out;\n l_info->layer++;\n}\n\nvoid set_silu_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n SiluActs* silu_acts = (SiluActs*)acts->layer_ptrs[layer];\n silu_acts->out = l_info->mem_ptr;\n l_info->mem_ptr += acts->layer_sizes[layer];\n silu_acts->x = l_info->input;\n l_info->input = silu_acts->out;\n l_info->layer++;\n}\n\nvoid set_convk3_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n ConvK3Acts* convk3_acts = (ConvK3Acts*)acts->layer_ptrs[layer];\n convk3_acts->out = l_info->mem_ptr;\n l_info->mem_ptr += acts->layer_sizes[layer];\n convk3_acts->inp = l_info->input;\n l_info->input = convk3_acts->out;\n l_info->layer++;\n}\n\nvoid set_resblock_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info, int C, int OC) {\n int layer = l_info->layer;\n ResBlockActivations* resblock_acts = (ResBlockActivations*)acts->layer_ptrs[layer];\n set_resblock_acts_ptrs(C, OC, resblock_acts, l_info->mem_ptr);\n resblock_acts->input = l_info->input;\n resblock_acts->emb = l_info->emb;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = resblock_acts->add2;\n l_info->layer++;\n}\n\nvoid set_avgpool_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n AvgpoolActs* avgpool_acts = (AvgpoolActs*)acts->layer_ptrs[layer];\n avgpool_acts->out = l_info->mem_ptr;\n avgpool_acts->x = l_info->input;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = avgpool_acts->out;\n l_info->layer++;\n}\n\nvoid set_attention_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n AttentionActs* att_acts = (AttentionActs*)acts->layer_ptrs[layer];\n set_attention_acts_pointers(att_acts, l_info->mem_ptr);\n att_acts->input = l_info->input;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = att_acts->add;\n l_info->layer++;\n}\n\nvoid set_concat_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info, float** skip_acts) {\n int layer = l_info->layer;\n ConcatChannelActs* cat_acts = (ConcatChannelActs*)acts->layer_ptrs[layer];\n cat_acts->out = l_info->mem_ptr;\n cat_acts->x1 = l_info->input;\n cat_acts->x2 = skip_acts[l_info->skip_idx];\n\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = cat_acts->out;\n l_info->layer++;\n l_info->skip_idx--;\n}\n\nvoid set_upsample_layer_acts_ptrs(UnetBlock* acts, LayerInfo* l_info) {\n int layer = l_info->layer;\n UpsampleActs* up_acts = (UpsampleActs*)acts->layer_ptrs[layer];\n up_acts->x = l_info->input;\n up_acts->out = l_info->mem_ptr;\n\n l_info->mem_ptr += acts-> layer_sizes[layer];\n l_info->input = up_acts->out;\n l_info->layer++;\n}\n\nvoid set_groupnorm_layer_acts_ptrs(\n UnetBlock* acts, LayerInfo* l_info,\n int B, int C, int H, int W, int gn_n_groups\n) {\n int layer = l_info->layer;\n GroupNormActs* gn_a = (GroupNormActs*)acts->layer_ptrs[layer];\n gn_set_act_ptrs(gn_a, l_info->mem_ptr, B, C, H, W, gn_n_groups);\n gn_a->x = l_info->input;\n l_info->mem_ptr += acts->layer_sizes[layer];\n l_info->input = gn_a->out;\n l_info->layer++;\n}\n\nvoid point_unet_acts(UnetConfig* config, UnetBlock* acts, float** skip_acts, float* input, float* timestep_emb) {\n // common constants\n int B = config->B;\n int C_model = config->C_model;\n int H = config->H;\n int W = config->W;\n int gn_n_groups = config->gn_n_groups;\n int n_levels = config->n_levels;\n int n_res = config->n_res;\n int att_start_level = config->att_start_level;\n\n LayerInfo l_info;\n l_info.layer = 0;\n l_info.mem_ptr = acts->mem_ptr;\n\n // timestep emb\n l_info.input = timestep_emb;\n set_linear_acts_ptrs(acts, &l_info);\n set_silu_acts_ptrs(acts, &l_info);\n set_linear_acts_ptrs(acts, &l_info);\n // point input to model input, and set outputu of the last linear layer as the embedding input\n l_info.emb = l_info.input;\n l_info.input = input;\n\n int OC = C_model;\n set_convk3_acts_ptrs(acts, &l_info);\n int skip_idx = 0;\n skip_acts[skip_idx] = l_info.input; // set the output of the previouse layer as input to the skip connection\n skip_idx++;\n\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config->channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n set_resblock_layer_acts_ptrs(acts, &l_info, C, OC);\n if (l < att_start_level) {\n // for the first 2 layers, the skip connection comes from the resblocks\n skip_acts[skip_idx] = l_info.input;\n skip_idx++;\n }\n\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_acts_ptrs(acts, &l_info);\n skip_acts[skip_idx] = l_info.input;\n skip_idx++;\n }\n }\n if (l < n_levels - 1) {\n set_avgpool_layer_acts_ptrs(acts, &l_info);\n skip_acts[skip_idx] = l_info.input;\n skip_idx++;\n }\n }\n\n set_resblock_layer_acts_ptrs(acts, &l_info, C, C);\n set_attention_layer_acts_ptrs(acts, &l_info);\n set_resblock_layer_acts_ptrs(acts, &l_info, C, C);\n\n // upward: set concat\n l_info.skip_idx = NUM_SKIP_CONNECTIONS - 1;\n for (int l = n_levels - 1; l >= 0; l--) {\n int mult = config->channel_mult[l];\n int C_skip = mult * C_model;\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config->channel_mult[l - 1];\n }\n set_concat_layer_acts_ptrs(acts, &l_info, skip_acts);\n OC = mult * C_model;\n set_resblock_layer_acts_ptrs(acts, &l_info, C + C_skip, OC);\n C = OC;\n if (l >= att_start_level) {\n set_attention_layer_acts_ptrs(acts, &l_info);\n }\n }\n if (l > 0) {\n set_upsample_layer_acts_ptrs(acts, &l_info);\n }\n }\n set_groupnorm_layer_acts_ptrs(acts, &l_info, B, C, H, W, gn_n_groups);\n set_silu_acts_ptrs(acts, &l_info);\n set_convk3_acts_ptrs(acts, &l_info);\n}\n\nvoid set_unet_back_pointers(UnetBackBlock* back) {\n UnetBackActs* back_acts = &(back->back_acts);\n float* back_mem_ptr = back->back_mem;\n\n float** ptrs[] = {\n &(back_acts->dinp_buf1),\n &(back_acts->dinp_buf2),\n &(back_acts->buf1_BCHW),\n &(back_acts->buf2_BCHW),\n &(back_acts->buf3_BCHW),\n &(back_acts->buf4_BCHW),\n &(back_acts->buf1_B3CHW),\n &(back_acts->buf2_B3CHW),\n &(back_acts->demb_buf1),\n &(back_acts->demb_buf2)\n };\n\n for (int i = 0; i < NUM_BACK_BUFS; i++) {\n *(ptrs[i]) = back_mem_ptr;\n back_mem_ptr += back->back_sizes[i];\n }\n\n // point the back acts for skip connections\n for (int s = 0; s < NUM_SKIP_CONNECTIONS; s++) {\n back_acts->skip_dinps[s] = back_mem_ptr;\n back_mem_ptr += back->back_sizes[NUM_BACK_BUFS + s];\n }\n}\n\nvoid malloc_and_point(Unet* unet) {\n float *param_mem, *grad_mem, *act_mem, *back_mem;\n cudaCheck(hipMalloc(¶m_mem, unet->params.total_size * sizeof(float)));\n cudaCheck(hipMalloc(&grad_mem, unet->params.total_size * sizeof(float)));\n cudaCheck(hipMalloc(&act_mem, unet->acts.total_size * sizeof(float)));\n cudaCheck(hipMalloc(&back_mem, unet->back.n_backs * sizeof(float)));\n\n cudaCheck(hipMemset(param_mem, 0, unet->params.total_size * sizeof(float)));\n cudaCheck(hipMemset(grad_mem, 0, unet->params.total_size * sizeof(float)));\n cudaCheck(hipMemset(act_mem, 0, unet->acts.total_size * sizeof(float)));\n cudaCheck(hipMemset(back_mem, 0, unet->back.n_backs * sizeof(float)));\n\n // constants\n int B = unet->config.B;\n int C_in = unet->config.C_in;\n int C_model = unet->config.C_model;\n int H = unet->config.H;\n int W = unet->config.W;\n\n // malloc and point input and timestep emb\n cudaCheck(hipMalloc(&(unet->input), B * C_in * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&(unet->timestep_emb), B * C_model * sizeof(float)));\n // malloc a nicely aligned buffer for the loss\n cudaCheck(hipMalloc(&(unet->gpu_loss), 64 * sizeof(float)));\n\n unet->params.mem_ptr = param_mem;\n unet->grads.mem_ptr = grad_mem;\n unet->acts.mem_ptr = act_mem;\n unet->back.back_mem = back_mem;\n\n point_unet_params(&(unet->config), &(unet->params));\n point_unet_params(&(unet->config), &(unet->grads));\n point_unet_acts(&(unet->config), &(unet->acts), unet->skip_acts, unet->input, unet->timestep_emb);\n set_unet_back_pointers(&(unet->back));\n}\n\n// ------------------------------------------------------------------------------------------------\n// forward pass\n\nvoid linear_layer_forward(\n Unet* unet,\n LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int N, int C, int OC, int block_size\n) {\n int layer = l_info->layer;\n LinearParams* linear_params = (LinearParams*)unet->params.layer_ptrs[layer];\n LinearActs* linear_acts = (LinearActs*)unet->acts.layer_ptrs[layer];\n matmul_forward2(cublas_handle, linear_acts->out, linear_acts->inp, linear_params->w, linear_params->b, N, C, OC, block_size);\n l_info->layer++;\n}\n\nvoid silu_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int N, int block_size\n) {\n SiluActs* acts = (SiluActs*)unet->acts.layer_ptrs[l_info->layer];\n silu_forward(acts->x, acts->out, N, block_size);\n l_info->layer++;\n}\n\nvoid convk3_layer_forward(\n Unet* unet, LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int B, int C, int OC, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n ConvK3Params* conv_params = (ConvK3Params*)unet->params.layer_ptrs[layer];\n ConvK3Acts* conv_acts = (ConvK3Acts*)unet->acts.layer_ptrs[layer];\n conv2d_k3_forward3(conv_acts->inp, conv_params->w, conv_params->b, conv_acts->out, B, C, OC, H, W);\n l_info->layer++;\n}\nvoid resblock_layer_forward(\n Unet* unet, LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int C_in, int C_emb, int C_out, int B, int H, int W, int gn_n_groups, int block_size\n) {\n int layer = l_info->layer;\n ResBlockParameters* resblock_params = (ResBlockParameters*)unet->params.layer_ptrs[layer];\n ResBlockActivations* resblock_acts = (ResBlockActivations*)unet->acts.layer_ptrs[layer];\n resblock_forward(cublas_handle, C_in, C_emb, C_out, B, H, W, block_size, 0, 0, gn_n_groups, resblock_params, resblock_acts);\n l_info->layer++;\n}\n\nvoid avgpool_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n AvgpoolActs* avgpool_acts = (AvgpoolActs*)unet->acts.layer_ptrs[layer];\n avgpool_2d_forward1(avgpool_acts->out, avgpool_acts->x, B, C, H, W, block_size);\n l_info->layer++;\n}\n\nvoid attention_layer_forward(\n Unet* unet, LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int B, int C, int H, int W, int HS, int gn_n_groups, int block_size\n) {\n int layer = l_info->layer;\n AttentionParams* att_params = (AttentionParams*)unet->params.layer_ptrs[layer];\n AttentionActs* att_acts = (AttentionActs*)unet->acts.layer_ptrs[layer];\n attention_block_forward(cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, att_params, att_acts);\n l_info->layer++;\n}\n\nvoid concat_layer_forward(\n Unet* unet,\n LayerInfo* l_info,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n ConcatChannelActs* cat_acts = (ConcatChannelActs*)unet->acts.layer_ptrs[layer];\n concat_channel_forward(cat_acts->x1, cat_acts->x2, cat_acts->out, B, C1, C2, H, W, block_size);\n l_info->layer++;\n}\n\nvoid upsample_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int block_size\n) {\n int layer = l_info->layer;\n UpsampleActs* up_acts = (UpsampleActs*)unet->acts.layer_ptrs[layer];\n upsample_forward1(up_acts->out, up_acts->x, B, C, H, W, block_size);\n l_info->layer++;\n}\n\nvoid groupnorm_layer_forward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int gn_n_groups\n) {\n int layer = l_info->layer;\n GroupNormParams* gn_p = (GroupNormParams*)unet->params.layer_ptrs[layer];\n GroupNormActs* gn_a = (GroupNormActs*)unet->acts.layer_ptrs[layer];\n groupnorm_forward(gn_a->x, gn_p->w, gn_p->b, gn_a->out, gn_a->mean, gn_a->rstd, B, C, H, W, gn_n_groups);\n l_info->layer++;\n}\n\n\nvoid unet_forward(\n hipblasHandle_t cublas_handle,\n Unet* unet,\n float* t_res = nullptr, float* t_att = nullptr\n) {\n // common constants\n UnetConfig config = unet->config;\n int B = config.B;\n int C_in = config.C_in;\n int C_model = config.C_model;\n int C_time = config.C_time;\n int C_out = config.C_out;\n int H = config.H;\n int W = config.W;\n int block_size = config.block_size;\n int gn_n_groups = config.gn_n_groups;\n int HS = config.HS;\n int n_levels = config.n_levels;\n int n_res = config.n_res;\n int att_start_level = config.att_start_level;\n\n hipEvent_t start, end;\n if (t_res != nullptr) {\n cudaCheck(hipEventCreate(&start));\n cudaCheck(hipEventCreate(&end));\n }\n\n LayerInfo l_info;\n l_info.layer = 0;\n linear_layer_forward(unet, &l_info, cublas_handle, B, C_model, C_time, block_size);\n silu_layer_forward(unet, &l_info, B * C_time, block_size);\n linear_layer_forward(unet, &l_info, cublas_handle, B, C_time, C_time, block_size);\n int OC = C_model;\n convk3_layer_forward(unet, &l_info, cublas_handle, B, C_in, OC, H, W, block_size);\n\n int C = OC;\n for (int l = 0; l < n_levels; l++) {\n int mult = config.channel_mult[l];\n OC = C_model * mult;\n for (int r = 0; r < n_res; r++) {\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(start));\n }\n resblock_layer_forward(unet, &l_info, cublas_handle, C, C_time, OC, B, H, W, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n *t_res += elapsed_time;\n printf(\"down resblock: l %d, r %d, Cin %d, Cout %d, H %d, W %d, time %f\\n\", l, r, C, OC, H, W, elapsed_time);\n }\n C = OC;\n if (l >= att_start_level) {\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(start));\n }\n attention_layer_forward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n *t_att += elapsed_time;\n printf(\"down att: l %d, r %d, C %d, H %d, W %d, time %f\\n\", l, r, C, H, W, elapsed_time);\n }\n }\n }\n if (l < n_levels - 1) {\n avgpool_layer_forward(unet, &l_info, B, C, H, W, block_size);\n H /= 2;\n W /= 2;\n }\n }\n\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(start));\n }\n resblock_layer_forward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n *t_res += elapsed_time;\n printf(\"middle resblock 1: Cin %d, Cout %d, H %d, W %d, time %f\\n\", C, OC, H, W, elapsed_time);\n cudaCheck(hipEventRecord(start));\n }\n attention_layer_forward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n *t_att += elapsed_time;\n printf(\"middle att: Cin %d, H %d, W %d, time %f\\n\", C, H, W, elapsed_time);\n cudaCheck(hipEventRecord(start));\n }\n resblock_layer_forward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n printf(\"middle resblock 2: Cin %d, Cout %d, H %d, W %d, time %f\\n\", C, OC, H, W, elapsed_time);\n *t_res += elapsed_time;\n }\n\n // upward\n for (int l = n_levels - 1; l >= 0; l--) {\n int C_skip = C_model * config.channel_mult[l];\n for (int r = 0; r <= n_res; r++) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config.channel_mult[l - 1];\n }\n concat_layer_forward(unet, &l_info, B, C, C_skip, H, W, block_size);\n OC = C_model * config.channel_mult[l];\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(start));\n }\n resblock_layer_forward(unet, &l_info, cublas_handle, C + C_skip, C_time, OC, B, H, W, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n printf(\"up resblock: l %d, r %d, Cin %d, Cskip %d, Cout %d, H %d, W %d, time %f\\n\", l, r, C, C_skip, OC, H, W, elapsed_time);\n *t_res += elapsed_time;\n }\n C = OC;\n if (l >= att_start_level) {\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(start));\n }\n attention_layer_forward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n printf(\"up att: l %d, r %d, C %d, H %d, W %d, time %f\\n\", l, r, C, H, W, elapsed_time);\n *t_att += elapsed_time;\n }\n }\n }\n if (l > 0) {\n upsample_layer_forward(unet, &l_info, B, C, H, W, block_size);\n H *= 2;\n W *= 2;\n }\n }\n\n groupnorm_layer_forward(unet, &l_info, B, C, H, W, gn_n_groups);\n silu_layer_forward(unet, &l_info, B * C * H * W, block_size);\n convk3_layer_forward(unet, &l_info, cublas_handle, B, C, C_out, H, W, block_size);\n\n ConvK3Acts* last_acts = (ConvK3Acts*)unet->acts.layer_ptrs[l_info.layer - 1];\n unet->output = last_acts->out;\n if (unet->target == nullptr) {\n // don't have targets, not computing loss\n unet->mean_loss = -1.0f;\n } else {\n mse_forward(unet->output, unet->target, unet->gpu_loss, B * C_out * H * W, block_size);\n cudaCheck(hipMemcpy(&(unet->mean_loss), unet->gpu_loss, sizeof(float), hipMemcpyDeviceToHost));\n }\n}\n\n// ------------------------------------------------------------------------------------------------\n// backward pass\nvoid swap_layer_in_out_ptrs(LayerInfo* l_info) {\n float* temp = l_info->input;\n l_info->input = l_info->prev_out;\n l_info->prev_out = temp;\n}\n\nvoid linear_layer_backward(\n Unet* unet,\n LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int N, int C, int OC\n) {\n int layer = l_info->layer;\n LinearParams* linear_params = (LinearParams*)unet->params.layer_ptrs[layer];\n LinearParams* linear_grads = (LinearParams*)unet->grads.layer_ptrs[layer];\n LinearActs* linear_acts = (LinearActs*)unet->acts.layer_ptrs[layer];\n matmul_backward1(cublas_handle, l_info->prev_out, linear_grads->w, linear_grads->b, l_info->input, linear_acts->inp, linear_params->w, N, C, OC);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid silu_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int N, int block_size\n) {\n SiluActs* acts = (SiluActs*)unet->acts.layer_ptrs[l_info->layer];\n silu_backward(l_info->input, acts->x, l_info->prev_out, N, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid convk3_layer_backward(\n Unet* unet, LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int B, int C, int OC, int H, int W, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n add_inplace_forward(skip_dinp, l_info->input, B * OC * H * W, block_size);\n l_info->skip_idx--;\n }\n int layer = l_info->layer;\n ConvK3Params* conv_params = (ConvK3Params*)unet->params.layer_ptrs[layer];\n ConvK3Params* conv_grads = (ConvK3Params*)unet->grads.layer_ptrs[layer];\n ConvK3Acts* conv_acts = (ConvK3Acts*)unet->acts.layer_ptrs[layer];\n UnetBackActs* unet_back_acts = &(unet->back.back_acts);\n conv2d_k3_backward2(l_info->input, conv_acts->inp, conv_params->w, unet_back_acts->buf1_BCHW, unet_back_acts->buf2_BCHW, l_info->prev_out, conv_grads->w, conv_grads->b, B, C, OC, H, W);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid resblock_layer_backward(\n Unet* unet, LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int C_in, int C_emb, int C_out, int B, int H, int W, int gn_n_groups, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n add_inplace_forward(skip_dinp, l_info->input, B * C_out * H * W, block_size);\n l_info->skip_idx--;\n }\n int layer = l_info->layer;\n ResBlockParameters* resblock_params = (ResBlockParameters*)unet->params.layer_ptrs[layer];\n ResBlockParameters* resblock_grads = (ResBlockParameters*)unet->grads.layer_ptrs[layer];\n ResBlockActivations* resblock_acts = (ResBlockActivations*)unet->acts.layer_ptrs[layer];\n ResBlockBackwardActivations back_acts;\n UnetBackActs* unet_back_acts = &(unet->back.back_acts);\n\n // point back_acts to buffers\n back_acts.buf_BCemb = unet_back_acts->buf1_BCHW;\n back_acts.buf_BCHoWo = unet_back_acts->buf2_BCHW;\n back_acts.buf1_BCHW = unet_back_acts->buf3_BCHW;\n back_acts.buf2_BCHW = unet_back_acts->buf4_BCHW;\n back_acts.dweight_buf = unet_back_acts->buf1_B3CHW;\n back_acts.dbias_buf = unet_back_acts->buf2_B3CHW;\n back_acts.dout = l_info->input;\n back_acts.dx = l_info->prev_out;\n back_acts.demb = l_info->emb;\n\n resblock_backward(cublas_handle, C_in, C_emb, C_out, B, H, W, block_size, 0, 0, gn_n_groups, resblock_params, resblock_grads, resblock_acts, &back_acts);\n swap_layer_in_out_ptrs(l_info);\n\n // add demb to cumulative demb\n add_inplace_forward(l_info->emb, l_info->emb_cum, B * C_emb, block_size);\n l_info->layer--;\n}\n\nvoid attention_layer_backward(\n Unet* unet, LayerInfo* l_info,\n hipblasHandle_t cublas_handle, int B, int C, int H, int W, int HS, int gn_n_groups, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n //printf(\"att adding skip %d\\n\", l_info->skip_idx);\n add_inplace_forward(skip_dinp, l_info->input, B * C * H * W, block_size);\n l_info->skip_idx--;\n }\n int layer = l_info->layer;\n AttentionParams* att_params = (AttentionParams*)unet->params.layer_ptrs[layer];\n AttentionParams* att_grads = (AttentionParams*)unet->grads.layer_ptrs[layer];\n AttentionActs* att_acts = (AttentionActs*)unet->acts.layer_ptrs[layer];\n AttentionBackwardActs att_backs;\n UnetBackActs* unet_backs = &(unet->back.back_acts);\n att_backs.buf1_BCHW = unet_backs->buf1_BCHW;\n att_backs.buf2_BCHW = unet_backs->buf2_BCHW;\n att_backs.buf_B3CHW = unet_backs->buf1_B3CHW;\n att_backs.dqkvr = unet_backs->buf2_B3CHW;\n att_backs.dpreatt = unet_backs->buf3_BCHW;\n att_backs.datt = unet_backs->buf4_BCHW;\n att_backs.dout = l_info->input;\n att_backs.dinp = l_info->prev_out;\n\n attention_block_backward(cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, att_params, att_acts, &att_backs, att_grads);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid concat_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n //printf(\"concat skip %d\\n\", l_info->skip_idx);\n concat_channel_backward(l_info->input, l_info->prev_out, skip_dinp, B, C1, C2, H, W, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n l_info->skip_idx++;\n}\n\n\nvoid avgpool_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int block_size,\n int add_skip = 0\n) {\n if (add_skip) {\n // need to add a skip connection component to dout, only used in the input downward layers\n float* skip_dinp = unet->back.back_acts.skip_dinps[l_info->skip_idx];\n //printf(\"avgpool adding skip %d\\n\", l_info->skip_idx);\n add_inplace_forward(skip_dinp, l_info->input, B * C * H/2 * W/2, block_size);\n l_info->skip_idx--;\n }\n avgpool_2d_backward1(l_info->input, l_info->prev_out, B, C, H, W, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid upsample_layer_backward(\n Unet* unet, LayerInfo *l_info,\n int B, int C, int H, int W, int block_size\n) {\n upsample_backward1(l_info->prev_out, l_info->input, B, C, H, W, block_size);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid groupnorm_layer_backward(\n Unet* unet, LayerInfo* l_info,\n int B, int C, int H, int W, int gn_n_groups\n) {\n int layer = l_info->layer;\n GroupNormParams* gn_p = (GroupNormParams*)unet->params.layer_ptrs[layer];\n GroupNormParams* gn_g = (GroupNormParams*)unet->grads.layer_ptrs[layer];\n GroupNormActs* gn_a = (GroupNormActs*)unet->acts.layer_ptrs[layer];\n groupnorm_backward(l_info->input, gn_a->x, gn_a->mean, gn_a->rstd, gn_p->w, l_info->prev_out, gn_g->w, gn_g->b, B, C, H, W, gn_n_groups);\n swap_layer_in_out_ptrs(l_info);\n l_info->layer--;\n}\n\nvoid unet_backward(\n hipblasHandle_t cublas_handle,\n Unet* unet,\n float* t_res = nullptr, float* t_att = nullptr\n) {\n if (unet->mean_loss == -1.0f or unet->target == nullptr) {\n printf(\"Error: must run forward with targets before calling backward\\n\");\n exit(EXIT_FAILURE);\n }\n hipEvent_t start, end;\n if (t_res != nullptr) {\n cudaCheck(hipEventCreate(&start));\n cudaCheck(hipEventCreate(&end));\n }\n // common constants\n UnetConfig config = unet->config;\n int B = config.B;\n int C_in = config.C_in;\n int C_model = config.C_model;\n int C_time = config.C_time;\n int C_out = config.C_out;\n int H = config.H;\n int W = config.W;\n int block_size = config.block_size;\n int gn_n_groups = config.gn_n_groups;\n int HS = config.HS;\n int n_levels = config.n_levels;\n int n_res = config.n_res;\n int att_start_level = config.att_start_level;\n\n UnetBackActs* back_acts = &(unet->back.back_acts);\n mse_backward(unet->output, unet->target, back_acts->dinp_buf1, B * C_out * H * W, block_size);\n\n LayerInfo l_info;\n l_info.layer = config.n_layers - 1;\n l_info.input = back_acts->dinp_buf1;\n l_info.prev_out = back_acts->dinp_buf2;\n l_info.emb = back_acts->demb_buf1;\n l_info.emb_cum = back_acts->demb_buf2;\n l_info.skip_idx = 0;\n\n int C;\n int OC;\n int C_skip;\n C = C_model;\n OC = C_model;\n convk3_layer_backward(unet, &l_info, cublas_handle, B, C, C_out, H, W, block_size);\n silu_layer_backward(unet, &l_info, B * C * H * W, block_size);\n groupnorm_layer_backward(unet, &l_info, B, C, H, W, gn_n_groups);\n\n for (int l = 0; l < n_levels; l++) {\n for (int r = n_res; r >= 0; r--) {\n if (r == n_res && l > 0) {\n C_skip = C_model * config.channel_mult[l - 1];\n } else {\n C_skip = C_model * config.channel_mult[l];\n }\n if (l >= att_start_level) {\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(start));\n }\n attention_layer_backward(unet, &l_info, cublas_handle, B, OC, H, W, HS, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n *t_att += elapsed_time;\n printf(\"up att back: l %d, r %d, Cout %d, H %d, W %d, time %f\\n\", l, r, OC, H, W, elapsed_time);\n }\n }\n if (r == 0 && l < n_levels - 1) {\n C = C_model * config.channel_mult[l + 1];\n }\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(start));\n }\n resblock_layer_backward(unet, &l_info, cublas_handle, C + C_skip, C_time, OC, B, H, W, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n *t_res += elapsed_time;\n printf(\"up res back: l %d, r %d, Cin %d, Cskip %d, Cout %d, H %d, W %d, time %f\\n\", l, r, C, C_skip, OC, H, W, elapsed_time);\n }\n concat_layer_backward(unet, &l_info, B, C, C_skip, H, W, block_size);\n }\n if (l < n_levels - 1) {\n H /= 2;\n W /= 2;\n upsample_layer_backward(unet, &l_info, B, C, H, W, block_size);\n OC = C_model * config.channel_mult[l + 1];\n }\n }\n\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(start));\n }\n resblock_layer_backward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n *t_res += elapsed_time;\n printf(\"middle resblock 2: Cin %d, Cout %d, H %d, W %d, time %f\\n\", C, OC, H, W, elapsed_time);\n cudaCheck(hipEventRecord(start));\n }\n attention_layer_backward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n *t_att += elapsed_time;\n printf(\"middle att: Cin %d, H %d, W %d, time %f\\n\", C, H, W, elapsed_time);\n cudaCheck(hipEventRecord(start));\n }\n resblock_layer_backward(unet, &l_info, cublas_handle, C, C_time, C, B, H, W, gn_n_groups, block_size);\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n printf(\"middle resblock 1: Cin %d, Cout %d, H %d, W %d, time %f\\n\", C, OC, H, W, elapsed_time);\n *t_res += elapsed_time;\n }\n\n // we will need to change the backward for all the other blocks below to include the dout contribution from the skip connection.\n l_info.skip_idx = NUM_SKIP_CONNECTIONS - 1;\n for (int l = n_levels - 1; l >= 0; l--) {\n if (l < n_levels - 1) {\n H *= 2;\n W *= 2;\n avgpool_layer_backward(unet, &l_info, B, C, H, W, block_size, 1);\n }\n for (int r = 0; r < n_res; r++) {\n if (l >= att_start_level) {\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(start));\n }\n attention_layer_backward(unet, &l_info, cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, 1);\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n printf(\"down att back: l %d, r %d, C %d, H %d, W %d, time %f\\n\", l, r, C, H, W, elapsed_time);\n *t_att += elapsed_time;\n }\n }\n int res_add_skip = (l < att_start_level) ? 1 : 0;\n if (r == n_res - 1 && l > 0) {\n int prev_mult = config.channel_mult[l - 1];\n C = prev_mult * C_model;\n }\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(start));\n }\n resblock_layer_backward(unet, &l_info, cublas_handle, C, C_time, OC, B, H, W, gn_n_groups, block_size, res_add_skip);\n if (t_res != nullptr) {\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n printf(\"down res back: l %d, r %d, Cin %d, Cout %d, H %d, W %d, time %f\\n\", l, r, C, OC, H, W, elapsed_time);\n *t_res += elapsed_time;\n }\n OC = C;\n }\n }\n\n convk3_layer_backward(unet, &l_info, cublas_handle, B, C_in, OC, H, W, block_size, 1);\n unet->dinp = l_info.input;\n\n // reset dout to point to emb_cum\n l_info.input = l_info.emb_cum;\n linear_layer_backward(unet, &l_info, cublas_handle, B, C_time, C_time);\n silu_layer_backward(unet, &l_info, B * C_time, block_size);\n linear_layer_backward(unet, &l_info, cublas_handle, B, C_model, C_time);\n}\n\n\n// ------------------------------------------------------------------------------------------------\n// functions to update gradient\n\nvoid unet_zero_grad(Unet* unet) {\n cudaCheck(hipMemset(unet->grads.mem_ptr, 0, unet->grads.total_size * sizeof(float)));\n cudaCheck(hipMemset(unet->back.back_mem, 0, unet->back.n_backs * sizeof(float)));\n //cudaCheck(cudaMemset(unet->acts.mem_ptr, 0, unet->acts.total_size * sizeof(float)));\n}\n\n// below copied from llm.c\n\n// Implements linear interpolation using only two floating-point operations (as opposed to three in a naive implementation).\n// Reference: https://developer.nvidia.com/blog/lerp-faster-cuda\n__device__ inline float lerp(float start, float end, float weight) {\n return fma(weight, end, fma(-weight, start, start));\n}\n\n__global__ void adamw_kernel2(float* params_memory, float* grads_memory, float* m_memory, float* v_memory, long num_parameters,\n float learning_rate, float beta1, float beta2, float beta1_correction, float beta2_correction, float eps, float weight_decay) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i >= num_parameters) return; // guard\n float grad = grads_memory[i];\n float m = m_memory[i];\n float v = v_memory[i];\n // update the first moment (momentum)\n m = lerp(grad, m, beta1);\n m_memory[i] = m;\n // update the second moment (RMSprop)\n v = lerp(grad * grad, v, beta2);\n v_memory[i] = v;\n m /= beta1_correction; // m_hat\n v /= beta2_correction; // v_hat\n params_memory[i] -= learning_rate * (m / (sqrtf(v) + eps) + weight_decay * params_memory[i]);\n}\n\nvoid unet_update(Unet* unet, float learning_rate, float beta1, float beta2, float eps, float weight_decay, int t) {\n // lazy allocation\n if (unet->m_memory == nullptr) {\n size_t param_size = unet->params.total_size * sizeof(float);\n cudaCheck(hipMalloc(&(unet->m_memory), param_size));\n cudaCheck(hipMalloc(&(unet->v_memory), param_size));\n cudaCheck(hipMemset(unet->m_memory, 0, param_size));\n cudaCheck(hipMemset(unet->v_memory, 0, param_size));\n printf(\"allocated %zu MiB for AdamW state m\\n\", param_size / 1024 / 1024);\n printf(\"allocated %zu MiB for AdamW state v\\n\", param_size / 1024 / 1024);\n }\n\n int block_size = unet->config.block_size;\n int n_blocks = ceil_div(unet->params.total_size, (size_t)block_size);\n float beta1_correction = 1.0f - powf(beta1, t);\n float beta2_correction = 1.0f - powf(beta2, t);\n adamw_kernel2<<>>(unet->params.mem_ptr, unet->grads.mem_ptr, unet->m_memory, unet->v_memory, unet->params.total_size,\n learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay);\n cudaCheck(hipGetLastError());\n}\n\n\n// ------------------------------------------------------------------------------------------------\n// run tests\n\n#define NUM_DEBUG_STATES 5\ntypedef struct {\n float* input;\n float* timestep_emb;\n float* out;\n float* dinp;\n float* target;\n} DebugStates;\n\n\nint main(int argc, char **argv) {\n // setup cublas\n hipblasHandle_t cublas_handle;\n cublasCheck(hipblasCreate(&cublas_handle));\n cublasCheck(hipblasSetMathMode(cublas_handle, HIPBLAS_DEFAULT_MATH));\n\n char cwd[PATH_MAX];\n if (getcwd(cwd, sizeof(cwd)) != NULL) {\n //printf(\"Current working dir: %s\\n\", cwd);\n } else {\n perror(\"getcwd() error\");\n return 1;\n }\n\n // read header\n FILE *model_file = fopenCheck(\"unet_test_params.bin\", \"rb\");\n int model_header[256];\n freadCheck(model_header, sizeof(int), 256, model_file);\n if (model_header[0] != 12345678) {\n fprintf(stderr, \"Invalid magic model file\\n\");\n exit(1);\n }\n int B = model_header[1];\n int C_in = model_header[2];\n int C_model = model_header[3];\n int C_out = model_header[4];\n int H = model_header[5];\n int W = model_header[6];\n int max_period = model_header[7];\n int C_time = C_model * 4;\n const int n_levels = 4;\n int channel_mult[n_levels] = {1, 2, 3, 4};\n\n // dataloader\n DataLoader loader;\n dataloader_init(&loader, DATASET_FILE, B);\n\n\n Unet unet;\n UnetConfig* config = &(unet.config);\n config->B = B;\n config->C_in = C_in;\n config->C_model = C_model;\n config->C_time = C_time;\n config->C_out = C_out;\n config->H = H;\n config->W = W;\n config->block_size = 512;\n config->gn_n_groups = 32;\n config->n_levels = n_levels;\n config->channel_mult = channel_mult;\n config->HS = 32;\n config->n_res = 2;\n config->max_period = max_period;\n config->att_start_level = 2;\n\n unet_init(&unet);\n\n unet_count_layers(config);\n unet_make_ptrs_and_count_memory(&unet);\n malloc_and_point(&unet);\n\n // allocate host memory for params\n float* params_mem_cpu = (float*)mallocCheck(unet.params.total_size * sizeof(float));\n freadCheck(params_mem_cpu, sizeof(float), unet.params.total_size, model_file);\n fcloseCheck(model_file);\n UnetBlock params_cpu;\n unet_setup_and_count_params(config, ¶ms_cpu);\n params_cpu.mem_ptr = params_mem_cpu;\n point_unet_params(config, ¶ms_cpu);\n\n // copy params to device\n cudaCheck(hipMemcpy(unet.params.mem_ptr, params_mem_cpu, unet.params.total_size * sizeof(float), hipMemcpyHostToDevice));\n\n // allocate host memory for debug states\n FILE *debug_file = fopenCheck(\"unet_test_acts.bin\", \"rb\");\n size_t state_sizes[NUM_DEBUG_STATES];\n int out_size = B * C_out * H * W;\n state_sizes[0] = B * C_in * H * W; // input\n state_sizes[1] = B; // timestep emb\n state_sizes[2] = out_size; // out\n state_sizes[3] = B * C_in * H * W; // dinp\n state_sizes[4] = out_size; // target\n\n size_t num_debug_params = 0;\n for (int i = 0; i < NUM_DEBUG_STATES; i++) {\n num_debug_params += state_sizes[i];\n }\n\n float* debug_states_cpu = (float*)mallocCheck(num_debug_params * sizeof(float));\n freadCheck(debug_states_cpu, sizeof(float), num_debug_params, debug_file);\n // point to the debug states\n DebugStates debug_states;\n float** debug_ptrs[] = {\n &(debug_states.input),\n &(debug_states.timestep_emb),\n &(debug_states.out),\n &(debug_states.dinp),\n &(debug_states.target),\n };\n\n float* debug_mem_ptr = debug_states_cpu;\n for (int i = 0; i < NUM_DEBUG_STATES; i++) {\n *(debug_ptrs[i]) = debug_mem_ptr;\n debug_mem_ptr += state_sizes[i];\n }\n\n // read the groud truth grads\n float* grads_mem_cpu = (float*)mallocCheck(unet.grads.total_size * sizeof(float));\n freadCheck(grads_mem_cpu, sizeof(float), unet.grads.total_size, debug_file);\n fcloseCheck(debug_file);\n\n\n // read the losses\n float losses[10];\n FILE *loss_file = fopenCheck(\"unet_test_losses.bin\", \"rb\");\n freadCheck(losses, sizeof(float), 10, loss_file);\n fcloseCheck(loss_file);\n\n // create a grads struct and point to cpu grads for easier validation\n UnetBlock grads_cpu;\n unet_setup_and_count_params(config, &grads_cpu);\n grads_cpu.mem_ptr = grads_mem_cpu;\n point_unet_params(config, &grads_cpu);\n\n //// copy input to device\n //cudaCheck(cudaMemcpy(unet.input, debug_states.input, B * C_in * H * W * sizeof(float), cudaMemcpyHostToDevice));\n // load input\n dataloader_next_batch(&loader);\n cudaCheck(hipMemcpy(unet.input, loader.input, B * C_in * H * W * sizeof(float), hipMemcpyHostToDevice));\n\n // try out the random number generator\n GaussianDiffusion diffusion;\n diffusion_init(&diffusion, B, C_in, H, W, 1000);\n printf(\"successful gaussian init\\n\");\n\n float rand_nums[B * 10];\n sample_timesteps(&diffusion, rand_nums, B * 10);\n\n // compute timestep emb\n float* d_timesteps;\n cudaCheck(hipMalloc(&d_timesteps, B * sizeof(float)));\n hipMemcpy(d_timesteps, rand_nums, B * sizeof(float), hipMemcpyHostToDevice);\n get_timestep_embeddings(&unet.time_emb_state, d_timesteps, unet.timestep_emb);\n\n // generate all the gaussian noise\n float *all_noise = (float*)malloc(10 * B * C_in * H * W * sizeof(float));\n normal_(all_noise, 10 * B * C_in * H * W, 0.0f, 1.0f, &diffusion.host_rng);\n\n\n cudaCheck(hipMalloc(&(unet.target), out_size * sizeof(float)));\n\n //printf(\"drawing and saving x_t for verification\\n\");\n //float* x_t_cpu = (float*)mallocCheck(B * C_in * H * W * sizeof(float));\n //diffusion_draw_and_sample_xt(&diffusion, unet.target, d_timesteps, unet.input);\n //print_gpu_first_5(\"x_t\", unet.target);\n //cudaCheck(cudaMemcpy(x_t_cpu, unet.target, B * C_in * H * W * sizeof(float), cudaMemcpyDeviceToHost));\n //FILE *sample_debug_file = fopenCheck(\"./sample_debug.bin\", \"wb\");\n //size_t result = fwrite(rand_nums, sizeof(float), B, sample_debug_file);\n //if (result != B) {\n // fprintf(stderr, \"Error writing to file\\n\");\n // exit(1);\n //}\n //result = fwrite(loader.input, sizeof(float), B * C_in * H * W, sample_debug_file);\n //if (result != B * C_in * H * W) {\n // fprintf(stderr, \"Error writing to file\\n\");\n // exit(1);\n //}\n //result = fwrite(x_t_cpu, sizeof(float), B * C_in * H * W, sample_debug_file);\n //if (result != B * C_in * H * W) {\n // fprintf(stderr, \"Error writing to file\\n\");\n // exit(1);\n //}\n //fcloseCheck(sample_debug_file);\n\n\n // copy the targets to gpu\n //cudaCheck(cudaMemcpy(unet.target, debug_states.target, out_size * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(hipMemcpy(unet.target, all_noise, out_size * sizeof(float), hipMemcpyHostToDevice));\n diffusion_sample(d_timesteps, unet.input, unet.target, &diffusion, B, C_in * H * W);\n //diffusion_draw_and_sample_xt(&diffusion, unet.target, d_timesteps, unet.input);\n\n\n\n printf(\"\\nCheck that unet targets match the debug states\\n\");\n validate_result(unet.target, debug_states.target, \"target\", out_size);\n\n\n float acc = 1e-4;\n float med_acc = 1e-2;\n printf(\"Running 10 update loops\\n\");\n // loop 1\n printf(\"Checking forward pass\\n\");\n float t_res = 0.0f;\n float t_att = 0.0f;\n unet_forward(cublas_handle, &unet, &t_res, &t_att);\n printf(\"resblock time: %f ms, attention time: %f ms\\n\", t_res, t_att);\n printf(\"\\nChecking out\\n\");\n validate_result(unet.output, debug_states.out, \"out\", out_size, med_acc);\n printf(\"cpu loss: %f, gpu loss: %f\\n\", losses[0], unet.mean_loss);\n printf(\"\\nForward pass successful\\n\");\n\n printf(\"\\nChecking backward pass\\n\");\n t_res = 0.0f;\n t_att = 0.0f;\n unet_backward(cublas_handle, &unet, &t_res, &t_att);\n printf(\"resblock back time: %f ms, attention back time: %f ms\\n\", t_res, t_att);\n printf(\"\\nChecking grads\\n\");\n validate_result(unet.grads.mem_ptr, grads_mem_cpu, \"grads\", unet.grads.total_size, med_acc);\n printf(\"\\nChecking dinp\\n\");\n validate_result(unet.dinp, debug_states.dinp, \"dinp\", B * C_in * H * W, acc);\n unet_update(&unet, 1e-3f, 0.9f, 0.999f, 1e-8f, 0.01f, 1);\n unet_zero_grad(&unet);\n\n // loop 2-9\n hipEvent_t start, end;\n cudaCheck(hipEventCreate(&start));\n cudaCheck(hipEventCreate(&end));\n cudaCheck(hipProfilerStart());\n cudaCheck(hipEventRecord(start));\n for (int i = 1; i < 10; i++) {\n dataloader_next_batch(&loader);\n cudaCheck(hipMemcpy(unet.input, loader.input, B * C_in * H * W * sizeof(float), hipMemcpyHostToDevice));\n hipMemcpy(d_timesteps, rand_nums + i * B, B * sizeof(float), hipMemcpyHostToDevice);\n get_timestep_embeddings(&unet.time_emb_state, d_timesteps, unet.timestep_emb);\n\n cudaCheck(hipMemcpy(unet.target, all_noise + i * out_size, out_size * sizeof(float), hipMemcpyHostToDevice));\n diffusion_sample(d_timesteps, unet.input, unet.target, &diffusion, B, C_in * H * W);\n //diffusion_draw_and_sample_xt(&diffusion, unet.target, d_timesteps, unet.input);\n\n unet_forward(cublas_handle, &unet);\n unet_backward(cublas_handle, &unet);\n printf(\"step %d, cpu loss: %f, gpu loss: %f\\n\", i+1, losses[i], unet.mean_loss);\n unet_update(&unet, 1e-3f, 0.9f, 0.999f, 1e-8f, 0.01f, i+1);\n unet_zero_grad(&unet);\n }\n cudaCheck(hipEventRecord(end));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float time_elapsed;\n cudaCheck(hipEventElapsedTime(&time_elapsed, start, end));\n time_elapsed /= 9;\n printf(\"Average time per step: %f ms\\n\", time_elapsed);\n\n\n\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmark\\n\");\n int repeat_times = 50;\n float elapsed_time = benchmark_kernel(repeat_times, unet_forward, cublas_handle, &unet, nullptr, nullptr);\n printf(\"Forward pass time: %f ms\\n\", elapsed_time);\n\n printf(\"\\nBackward pass benchmark\\n\");\n elapsed_time = benchmark_kernel(repeat_times, unet_backward, cublas_handle, &unet, nullptr, nullptr);\n printf(\"Backward pass time: %f ms\\n\", elapsed_time);\n\n // free memory\n free(params_mem_cpu);\n free(debug_states_cpu);\n free_unet(&unet);\n cudaCheck(hipFree(unet.target));\n cudaCheck(hipFree(d_timesteps));\n dataloader_free(&loader);\n diffusion_free(&diffusion);\n free(all_noise);\n\n // free cublas\n cublasCheck(hipblasDestroy(cublas_handle));\n}", "source": "/kaggle/working/unet.cu/dev/unet_test.cu"} {"input": "#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \"attention.cuh\"\n#include \"groupnorm.cuh\"\n#include \"linear.cuh\"\n#include \"add.cuh\"\n#include \"attention_block.cuh\"\n\n\n__global__ void attention_permute_kernel(\n const float* inp, float* out,\n int B, int C, int H, int W\n) {\n // from B, C, H, W to B, HxW, C\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n if (idx >= B * C * img_size) { return; }\n\n int b = idx / (img_size * C);\n int c = idx / img_size % C;\n int h = idx / W % H;\n int w = idx % W;\n\n out[b * C * img_size + h * C * W + w * C + c] = inp[idx];\n}\n\nvoid attention_permute(\n const float* inp, float* out,\n int B, int C, int H, int W, int block_size\n) {\n int num_threads = B * C * H * W;\n attention_permute_kernel<<>>(inp, out, B, C, H, W);\n}\n\n__global__ void attention_unpermute_kernel(\n const float* inp, float* out,\n int B, int C, int H, int W\n) {\n // from B, HxW, C to B, C, H, W\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n if (idx >= B * C * img_size) { return; }\n\n int b = idx / (img_size * C);\n int c = idx / img_size % C;\n int h = idx / W % H;\n int w = idx % W;\n\n out[idx] = inp[b * C * img_size + h * C * W + w * C + c];\n}\n\nvoid attention_unpermute(\n const float* inp, float* out,\n int B, int C, int H, int W, int block_size\n) {\n int num_threads = B * C * H * W;\n attention_unpermute_kernel<<>>(inp, out, B, C, H, W);\n}\n\n\nvoid attention_block_forward\n(\n cublasHandle_t cublas_handle,\n int B, int C, int H, int W, int HS, int gn_n_groups, int block_size,\n AttentionParams* params, AttentionActs* acts\n) {\n // apply groupnorm\n groupnorm_forward(acts->input, params->gn_w, params->gn_b, acts->gn, acts->gn_mean, acts->gn_rstd, B, C, H, W, gn_n_groups);\n // permute from B, C, H, W to B, HxW, C\n attention_permute(acts->gn, acts->perm1, B, C, H, W, block_size);\n // get qkv\n matmul_forward2(cublas_handle, acts->qkv1, acts->perm1, params->qkv_w, params->qkv_b, B * H * W, C, 3 * C, block_size);\n // apply attention\n attention_forward1(cublas_handle, acts->att_out, acts->qkv2, acts->preatt, acts->att, acts->qkv1, B, H * W, C, C / HS, block_size);\n // attention projection\n matmul_forward2(cublas_handle, acts->proj, acts->att_out, params->proj_w, params->proj_b, B * H * W, C, C, block_size);\n // permute from B, HxW, C to B, C, H, W\n attention_unpermute(acts->proj, acts->perm2, B, C, H, W, block_size);\n // add back residual stream\n add_forward(acts->input, acts->perm2, acts->add, B * C * H * W, block_size);\n}\n\n\nvoid attention_block_backward\n(\n cublasHandle_t cublas_handle,\n int B, int C, int H, int W, int HS, int gn_n_groups, int block_size,\n AttentionParams* params, AttentionActs* acts,\n AttentionBackwardActs* back_acts, AttentionParams* grads\n) {\n // permute from B, C, H, W to B, HxW, C\n attention_permute(back_acts->dout, back_acts->buf1_BCHW, B, C, H, W, block_size);\n // backward through projection\n matmul_backward1(cublas_handle, back_acts->buf2_BCHW, grads->proj_w, grads->proj_b, back_acts->buf1_BCHW, acts->att_out, params->proj_w, B * H * W, C, C);\n // backward through attention\n attention_backward(cublas_handle, back_acts->buf_B3CHW, back_acts->dqkvr, back_acts->dpreatt, back_acts->datt, back_acts->buf1_BCHW, back_acts->buf2_BCHW, acts->qkv2, acts->att, B, H * W, C, C / HS);\n // backward through qkv\n matmul_backward1(cublas_handle, back_acts->buf1_BCHW, grads->qkv_w, grads->qkv_b, back_acts->buf_B3CHW, acts->perm1, params->qkv_w, B * H * W, C, 3 * C);\n // permute from B, HxW, C to B, C, H, W\n attention_unpermute(back_acts->buf1_BCHW, back_acts->buf2_BCHW, B, C, H, W, block_size);\n // backward through groupnorm\n groupnorm_backward(back_acts->buf2_BCHW, acts->input, acts->gn_mean, acts->gn_rstd, params->gn_w, back_acts->buf1_BCHW, grads->gn_w, grads->gn_b, B, C, H, W, gn_n_groups);\n // add residual gradient\n add_forward(back_acts->buf1_BCHW, back_acts->dout, back_acts->dinp, B * C * H * W, block_size);\n}\n\n\nvoid set_attention_params_pointers(\n AttentionParams* params,\n float* params_memory\n) {\n float** ptrs[] = {\n ¶ms->gn_w,\n ¶ms->gn_b,\n ¶ms->qkv_w,\n ¶ms->qkv_b,\n ¶ms->proj_w,\n ¶ms->proj_b\n };\n for (int i = 0; i < NUM_ATT_PARAM_TENSORS; i++) {\n *(ptrs[i]) = params_memory;\n params_memory += params->param_sizes[i];\n }\n}\n\nvoid set_attention_acts_pointers(\n AttentionActs* acts,\n float* acts_memory\n) {\n float** ptrs[] = {\n &acts->gn,\n &acts->gn_mean,\n &acts->gn_rstd,\n &acts->perm1,\n &acts->qkv1,\n &acts->qkv2,\n &acts->preatt,\n &acts->att,\n &acts->att_out,\n &acts->proj,\n &acts->perm2,\n &acts->add\n };\n for (int i = 0; i < NUM_ATT_ACT_TENSORS; i++) {\n *(ptrs[i]) = acts_memory;\n acts_memory += acts->act_sizes[i];\n }\n}\n\nvoid set_attention_back_acts_pointers(\n AttentionBackwardActs* back_acts,\n float* back_acts_memory,\n size_t* back_act_sizes\n) {\n float** ptrs[] = {\n &back_acts->buf1_BCHW,\n &back_acts->buf2_BCHW,\n &back_acts->buf_B3CHW,\n &back_acts->dqkvr,\n &back_acts->dpreatt,\n &back_acts->datt,\n &back_acts->dout,\n };\n for (int i = 0; i < NUM_ATT_BACKWARD_ACTS_TENSORS; i++) {\n *(ptrs[i]) = back_acts_memory;\n back_acts_memory += back_act_sizes[i];\n }\n}\n\n\nvoid attention_block_count_backs(\n AttentionBackwardActs* back_acts,\n int B, int C, int H, int W, int HS, int gn_n_groups\n){\n int NH = C / HS;\n int C3 = 3 * C;\n int T = H * W;\n\n // count back acts\n size_t *back_act_sizes = back_acts->back_sizes;\n back_act_sizes[0] = B * C * H * W; // buf1_BCHW\n back_act_sizes[1] = B * C * H * W; // buf2_BCHW\n back_act_sizes[2] = B * C3 * H * W; // buf_B3CHW\n back_act_sizes[3] = B * T * C3; // dqkvr\n back_act_sizes[4] = B * NH * T * T; // dpreatt\n back_act_sizes[5] = B * NH * T * T; // datt\n back_act_sizes[6] = B * C * H * W; // dout\n \n size_t num_back_act_params = 0;\n for (int i = 0; i < NUM_ATT_BACKWARD_ACTS_TENSORS; i++) {\n num_back_act_params += back_act_sizes[i];\n }\n back_acts->n_backs = num_back_act_params;\n}\n\nvoid attention_block_count_params(AttentionParams* params, int C) {\n size_t *param_sizes = params->param_sizes;\n param_sizes[0] = C; // gn_w\n param_sizes[1] = C; // gn_b\n param_sizes[2] = 3 * C * C; // qkv_w\n param_sizes[3] = 3 * C; // qkv_b\n param_sizes[4] = C * C; // proj_w\n param_sizes[5] = C; // proj_b\n\n size_t num_parameters = 0;\n for (int i = 0; i < NUM_ATT_PARAM_TENSORS; i++) {\n num_parameters += param_sizes[i];\n }\n params->n_params = num_parameters;\n}\nvoid attention_block_count_acts(\n AttentionActs* acts,\n int B, int C, int H, int W, int gn_n_groups, int HS\n) {\n int C3 = 3 * C;\n int T = H * W;\n int NH = C / HS;\n size_t *act_sizes = acts->act_sizes;\n act_sizes[0] = B * C * H * W; // gn\n act_sizes[1] = B * gn_n_groups; // gn_mean\n act_sizes[2] = B * gn_n_groups; // gn_rstd\n act_sizes[3] = B * T * C; // perm1\n act_sizes[4] = B * T * C3; // qkv1\n act_sizes[5] = B * T * C3; // qkv2\n act_sizes[6] = B * NH * T * T; // preatt\n act_sizes[7] = B * NH * T * T; // att\n act_sizes[8] = B * T * C; // att_out\n act_sizes[9] = B * T * C; // proj\n act_sizes[10] = B * C * H * W; // perm2\n act_sizes[11] = B * C * H * W; // add\n\n size_t num_acts_params = 0;\n for (int i = 0; i < NUM_ATT_ACT_TENSORS; i++) {\n num_acts_params += act_sizes[i];\n }\n acts->n_acts = num_acts_params;\n}\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n // setup cublas\n cublasHandle_t cublas_handle;\n cublasCheck(cublasCreate(&cublas_handle));\n cublasCheck(cublasSetMathMode(cublas_handle, CUBLAS_DEFAULT_MATH));\n\n // read header\n FILE *model_file = fopenCheck(\"attention_block_params.bin\", \"rb\");\n int model_header[256];\n freadCheck(model_header, sizeof(int), 256, model_file);\n if (model_header[0] != 12345678) {\n fprintf(stderr, \"Invalid magic model file, magic number %d\\n\", model_header[0]);\n exit(1);\n }\n int B = model_header[1];\n int C = model_header[2];\n int H = model_header[3];\n int W = model_header[4];\n int HS = model_header[5];\n int gn_n_groups = model_header[6];\n\n AttentionParams params, grads, params_cpu;\n attention_block_count_params(¶ms, C);\n attention_block_count_params(&grads, C);\n attention_block_count_params(¶ms_cpu, C);\n\n size_t num_parameters = params.n_params;\n // read params to host\n float* params_memory_cpu = (float*)mallocCheck(num_parameters * sizeof(float));\n freadCheck(params_memory_cpu, sizeof(float), num_parameters, model_file);\n fcloseCheck(model_file);\n\n // copy params to device\n float* params_memory_gpu;\n cudaCheck(cudaMalloc(¶ms_memory_gpu, num_parameters * sizeof(float)));\n cudaCheck(cudaMemcpy(params_memory_gpu, params_memory_cpu, num_parameters * sizeof(float), cudaMemcpyHostToDevice));\n\n // create device memory for grads\n float* grads_memory_gpu;\n cudaCheck(cudaMalloc(&grads_memory_gpu, num_parameters * sizeof(float)));\n cudaCheck(cudaMemset(grads_memory_gpu, 0, num_parameters * sizeof(float)));\n\n // set param pointers\n set_attention_params_pointers(¶ms, params_memory_gpu);\n set_attention_params_pointers(&grads, grads_memory_gpu);\n set_attention_params_pointers(¶ms_cpu, params_memory_cpu);\n\n // calculate activation sizes\n int C3 = 3 * C;\n int T = H * W;\n AttentionActs acts;\n attention_block_count_acts(&acts, B, C, H, W, gn_n_groups, HS);\n size_t num_acts_params = acts.n_acts;\n\n // allocate device act memory\n float* acts_memory_gpu;\n cudaCheck(cudaMalloc(&acts_memory_gpu, num_acts_params * sizeof(float)));\n cudaCheck(cudaMemset(acts_memory_gpu, 0, num_acts_params * sizeof(float)));\n\n // set acts pointers\n set_attention_acts_pointers(&acts, acts_memory_gpu);\n\n // malloc input, which is not part of the large block of acts memory\n cudaCheck(cudaMalloc(&acts.input, B * C * H * W * sizeof(float)));\n\n // create debug states\n size_t debug_state_sizes[NUM_ATT_DEBUG_STATES];\n debug_state_sizes[0] = B * C * H * W; // inp\n debug_state_sizes[1] = B * C * H * W; // gn\n debug_state_sizes[2] = B * C * H * W; // perm1\n debug_state_sizes[3] = B * C3 * H * W; // qkv\n debug_state_sizes[4] = B * C * H * W; // att\n debug_state_sizes[5] = B * C * H * W; // proj\n debug_state_sizes[6] = B * C * H * W; // out\n debug_state_sizes[7] = B * C * H * W; // dout\n debug_state_sizes[8] = B * C * H * W; // dinp\n\n size_t num_debug_states = 0;\n for (int i = 0; i < NUM_ATT_DEBUG_STATES; i++) {\n num_debug_states += debug_state_sizes[i];\n }\n\n // allocate host debug state memory and host grads memory\n float* debug_states_memory_cpu = (float*)mallocCheck(num_debug_states * sizeof(float));\n float* grads_memory_cpu = (float*)mallocCheck(num_parameters * sizeof(float));\n AttentionParams grads_cpu;\n attention_block_count_params(&grads_cpu, C);\n set_attention_params_pointers(&grads_cpu, grads_memory_cpu);\n FILE *debug_states_file = fopenCheck(\"attention_block_states.bin\", \"rb\");\n printf(\"num debug states: %ld\\n\", num_debug_states);\n freadCheck(debug_states_memory_cpu, sizeof(float), num_debug_states, debug_states_file);\n freadCheck(grads_memory_cpu, sizeof(float), num_parameters, debug_states_file);\n fcloseCheck(debug_states_file);\n\n // set debug states pointers\n AttentionDebugStates debug_states;\n float* debug_states_ptr = debug_states_memory_cpu;\n float** debug_states_ptrs[] = {\n &debug_states.inp,\n &debug_states.gn,\n &debug_states.perm1,\n &debug_states.qkv,\n &debug_states.att,\n &debug_states.proj,\n &debug_states.out,\n &debug_states.dout,\n &debug_states.dinp\n };\n for (int i = 0; i < NUM_ATT_DEBUG_STATES; i++) {\n *(debug_states_ptrs[i]) = debug_states_ptr;\n debug_states_ptr += debug_state_sizes[i];\n }\n\n // create attention backward acts\n AttentionBackwardActs back_acts;\n attention_block_count_backs(&back_acts, B, C, H, W, HS, gn_n_groups);\n size_t num_back_act_params = back_acts.n_backs;\n\n float* back_acts_memory_gpu;\n cudaCheck(cudaMalloc(&back_acts_memory_gpu, num_back_act_params * sizeof(float)));\n cudaCheck(cudaMemset(back_acts_memory_gpu, 0, num_back_act_params * sizeof(float)));\n\n // set pointers\n set_attention_back_acts_pointers(&back_acts, back_acts_memory_gpu, back_acts.back_sizes);\n back_acts.dinp = acts.input;\n\n // copy input and dout to device\n cudaCheck(cudaMemcpy(acts.input, debug_states.inp, B * C * H * W * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(back_acts.dout, debug_states.dout, B * C * H * W * sizeof(float), cudaMemcpyHostToDevice));\n\n int block_size = 512;\n\n printf(\"Checking forward pass\\n\");\n attention_block_forward(cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, ¶ms, &acts);\n\n // validate results\n float acc = 1e-4;\n printf(\"\\nChecking gn\\n\");\n validate_result(acts.gn, debug_states.gn, \"gn\", B * C * H * W, acc);\n printf(\"\\nChecking perm1\\n\");\n validate_result(acts.perm1, debug_states.perm1, \"perm1\", B * T * C, acc);\n // don't check qkv because it is used as a buffer in attention step so values are wrong\n //printf(\"\\nChecking qkv\\n\");\n //validate_result(acts.qkv1, debug_states.qkv, \"qkv\", B * C3 * H * W, acc);\n printf(\"\\nChecking att\\n\");\n validate_result(acts.att_out, debug_states.att, \"att_out\", B * T * C, acc);\n printf(\"\\nChecking proj\\n\");\n validate_result(acts.proj, debug_states.proj, \"proj\", B * T * C, acc);\n printf(\"\\nChecking out\\n\");\n validate_result(acts.add, debug_states.out, \"out\", B * C * H * W, acc);\n\n printf(\"\\nForward successful\\n\");\n\n printf(\"Checking backward pass\\n\");\n attention_block_backward(cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, ¶ms, &acts, &back_acts, &grads);\n\n printf(\"\\nChecking dproj_w\\n\");\n validate_result(grads.proj_w, grads_cpu.proj_w, \"proj_w\", C * C, acc);\n printf(\"\\nChecking dproj_b\\n\");\n validate_result(grads.proj_b, grads_cpu.proj_b, \"proj_b\", C, acc);\n printf(\"\\nChecking dqkv_w\\n\");\n validate_result(grads.qkv_w, grads_cpu.qkv_w, \"qkv_w\", 3 * C * C, acc);\n printf(\"\\nChecking dqkv_b\\n\");\n validate_result(grads.qkv_b, grads_cpu.qkv_b, \"qkv_b\", 3 * C, acc);\n printf(\"\\nChecking dgn_w\\n\");\n validate_result(grads.gn_w, grads_cpu.gn_w, \"gn_w\", C, acc);\n printf(\"\\nChecking dgn_b\\n\");\n validate_result(grads.gn_b, grads_cpu.gn_b, \"gn_b\", C, acc);\n printf(\"\\nChecking dinp\\n\");\n validate_result(back_acts.dinp, debug_states.dinp, \"dinp\", B * C * H * W, acc);\n\n printf(\"\\nAll results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmark\\n\");\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(\n repeat_times, attention_block_forward,\n cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, ¶ms, &acts\n );\n printf(\"forward pass | time %.4f ms\\n\", elapsed_time);\n\n printf(\"Backward pass benchmark\\n\");\n elapsed_time = benchmark_kernel(\n repeat_times, attention_block_backward,\n cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, ¶ms, &acts, &back_acts, &grads\n );\n printf(\"backward pass | time %.4f ms\\n\", elapsed_time);\n\n\n cublasCheck(cublasDestroy(cublas_handle));\n // free memory\n cudaCheck(cudaFree(params_memory_gpu));\n cudaCheck(cudaFree(acts_memory_gpu));\n cudaCheck(cudaFree(back_acts_memory_gpu));\n cudaCheck(cudaFree(grads_memory_gpu));\n cudaCheck(cudaFree(acts.input));\n free(params_memory_cpu);\n free(debug_states_memory_cpu);\n free(grads_memory_cpu);\n\n}\n#endif", "output": "#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \"attention.cuh\"\n#include \"groupnorm.cuh\"\n#include \"linear.cuh\"\n#include \"add.cuh\"\n#include \"attention_block.cuh\"\n\n\n__global__ void attention_permute_kernel(\n const float* inp, float* out,\n int B, int C, int H, int W\n) {\n // from B, C, H, W to B, HxW, C\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n if (idx >= B * C * img_size) { return; }\n\n int b = idx / (img_size * C);\n int c = idx / img_size % C;\n int h = idx / W % H;\n int w = idx % W;\n\n out[b * C * img_size + h * C * W + w * C + c] = inp[idx];\n}\n\nvoid attention_permute(\n const float* inp, float* out,\n int B, int C, int H, int W, int block_size\n) {\n int num_threads = B * C * H * W;\n attention_permute_kernel<<>>(inp, out, B, C, H, W);\n}\n\n__global__ void attention_unpermute_kernel(\n const float* inp, float* out,\n int B, int C, int H, int W\n) {\n // from B, HxW, C to B, C, H, W\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n if (idx >= B * C * img_size) { return; }\n\n int b = idx / (img_size * C);\n int c = idx / img_size % C;\n int h = idx / W % H;\n int w = idx % W;\n\n out[idx] = inp[b * C * img_size + h * C * W + w * C + c];\n}\n\nvoid attention_unpermute(\n const float* inp, float* out,\n int B, int C, int H, int W, int block_size\n) {\n int num_threads = B * C * H * W;\n attention_unpermute_kernel<<>>(inp, out, B, C, H, W);\n}\n\n\nvoid attention_block_forward\n(\n hipblasHandle_t cublas_handle,\n int B, int C, int H, int W, int HS, int gn_n_groups, int block_size,\n AttentionParams* params, AttentionActs* acts\n) {\n // apply groupnorm\n groupnorm_forward(acts->input, params->gn_w, params->gn_b, acts->gn, acts->gn_mean, acts->gn_rstd, B, C, H, W, gn_n_groups);\n // permute from B, C, H, W to B, HxW, C\n attention_permute(acts->gn, acts->perm1, B, C, H, W, block_size);\n // get qkv\n matmul_forward2(cublas_handle, acts->qkv1, acts->perm1, params->qkv_w, params->qkv_b, B * H * W, C, 3 * C, block_size);\n // apply attention\n attention_forward1(cublas_handle, acts->att_out, acts->qkv2, acts->preatt, acts->att, acts->qkv1, B, H * W, C, C / HS, block_size);\n // attention projection\n matmul_forward2(cublas_handle, acts->proj, acts->att_out, params->proj_w, params->proj_b, B * H * W, C, C, block_size);\n // permute from B, HxW, C to B, C, H, W\n attention_unpermute(acts->proj, acts->perm2, B, C, H, W, block_size);\n // add back residual stream\n add_forward(acts->input, acts->perm2, acts->add, B * C * H * W, block_size);\n}\n\n\nvoid attention_block_backward\n(\n hipblasHandle_t cublas_handle,\n int B, int C, int H, int W, int HS, int gn_n_groups, int block_size,\n AttentionParams* params, AttentionActs* acts,\n AttentionBackwardActs* back_acts, AttentionParams* grads\n) {\n // permute from B, C, H, W to B, HxW, C\n attention_permute(back_acts->dout, back_acts->buf1_BCHW, B, C, H, W, block_size);\n // backward through projection\n matmul_backward1(cublas_handle, back_acts->buf2_BCHW, grads->proj_w, grads->proj_b, back_acts->buf1_BCHW, acts->att_out, params->proj_w, B * H * W, C, C);\n // backward through attention\n attention_backward(cublas_handle, back_acts->buf_B3CHW, back_acts->dqkvr, back_acts->dpreatt, back_acts->datt, back_acts->buf1_BCHW, back_acts->buf2_BCHW, acts->qkv2, acts->att, B, H * W, C, C / HS);\n // backward through qkv\n matmul_backward1(cublas_handle, back_acts->buf1_BCHW, grads->qkv_w, grads->qkv_b, back_acts->buf_B3CHW, acts->perm1, params->qkv_w, B * H * W, C, 3 * C);\n // permute from B, HxW, C to B, C, H, W\n attention_unpermute(back_acts->buf1_BCHW, back_acts->buf2_BCHW, B, C, H, W, block_size);\n // backward through groupnorm\n groupnorm_backward(back_acts->buf2_BCHW, acts->input, acts->gn_mean, acts->gn_rstd, params->gn_w, back_acts->buf1_BCHW, grads->gn_w, grads->gn_b, B, C, H, W, gn_n_groups);\n // add residual gradient\n add_forward(back_acts->buf1_BCHW, back_acts->dout, back_acts->dinp, B * C * H * W, block_size);\n}\n\n\nvoid set_attention_params_pointers(\n AttentionParams* params,\n float* params_memory\n) {\n float** ptrs[] = {\n ¶ms->gn_w,\n ¶ms->gn_b,\n ¶ms->qkv_w,\n ¶ms->qkv_b,\n ¶ms->proj_w,\n ¶ms->proj_b\n };\n for (int i = 0; i < NUM_ATT_PARAM_TENSORS; i++) {\n *(ptrs[i]) = params_memory;\n params_memory += params->param_sizes[i];\n }\n}\n\nvoid set_attention_acts_pointers(\n AttentionActs* acts,\n float* acts_memory\n) {\n float** ptrs[] = {\n &acts->gn,\n &acts->gn_mean,\n &acts->gn_rstd,\n &acts->perm1,\n &acts->qkv1,\n &acts->qkv2,\n &acts->preatt,\n &acts->att,\n &acts->att_out,\n &acts->proj,\n &acts->perm2,\n &acts->add\n };\n for (int i = 0; i < NUM_ATT_ACT_TENSORS; i++) {\n *(ptrs[i]) = acts_memory;\n acts_memory += acts->act_sizes[i];\n }\n}\n\nvoid set_attention_back_acts_pointers(\n AttentionBackwardActs* back_acts,\n float* back_acts_memory,\n size_t* back_act_sizes\n) {\n float** ptrs[] = {\n &back_acts->buf1_BCHW,\n &back_acts->buf2_BCHW,\n &back_acts->buf_B3CHW,\n &back_acts->dqkvr,\n &back_acts->dpreatt,\n &back_acts->datt,\n &back_acts->dout,\n };\n for (int i = 0; i < NUM_ATT_BACKWARD_ACTS_TENSORS; i++) {\n *(ptrs[i]) = back_acts_memory;\n back_acts_memory += back_act_sizes[i];\n }\n}\n\n\nvoid attention_block_count_backs(\n AttentionBackwardActs* back_acts,\n int B, int C, int H, int W, int HS, int gn_n_groups\n){\n int NH = C / HS;\n int C3 = 3 * C;\n int T = H * W;\n\n // count back acts\n size_t *back_act_sizes = back_acts->back_sizes;\n back_act_sizes[0] = B * C * H * W; // buf1_BCHW\n back_act_sizes[1] = B * C * H * W; // buf2_BCHW\n back_act_sizes[2] = B * C3 * H * W; // buf_B3CHW\n back_act_sizes[3] = B * T * C3; // dqkvr\n back_act_sizes[4] = B * NH * T * T; // dpreatt\n back_act_sizes[5] = B * NH * T * T; // datt\n back_act_sizes[6] = B * C * H * W; // dout\n \n size_t num_back_act_params = 0;\n for (int i = 0; i < NUM_ATT_BACKWARD_ACTS_TENSORS; i++) {\n num_back_act_params += back_act_sizes[i];\n }\n back_acts->n_backs = num_back_act_params;\n}\n\nvoid attention_block_count_params(AttentionParams* params, int C) {\n size_t *param_sizes = params->param_sizes;\n param_sizes[0] = C; // gn_w\n param_sizes[1] = C; // gn_b\n param_sizes[2] = 3 * C * C; // qkv_w\n param_sizes[3] = 3 * C; // qkv_b\n param_sizes[4] = C * C; // proj_w\n param_sizes[5] = C; // proj_b\n\n size_t num_parameters = 0;\n for (int i = 0; i < NUM_ATT_PARAM_TENSORS; i++) {\n num_parameters += param_sizes[i];\n }\n params->n_params = num_parameters;\n}\nvoid attention_block_count_acts(\n AttentionActs* acts,\n int B, int C, int H, int W, int gn_n_groups, int HS\n) {\n int C3 = 3 * C;\n int T = H * W;\n int NH = C / HS;\n size_t *act_sizes = acts->act_sizes;\n act_sizes[0] = B * C * H * W; // gn\n act_sizes[1] = B * gn_n_groups; // gn_mean\n act_sizes[2] = B * gn_n_groups; // gn_rstd\n act_sizes[3] = B * T * C; // perm1\n act_sizes[4] = B * T * C3; // qkv1\n act_sizes[5] = B * T * C3; // qkv2\n act_sizes[6] = B * NH * T * T; // preatt\n act_sizes[7] = B * NH * T * T; // att\n act_sizes[8] = B * T * C; // att_out\n act_sizes[9] = B * T * C; // proj\n act_sizes[10] = B * C * H * W; // perm2\n act_sizes[11] = B * C * H * W; // add\n\n size_t num_acts_params = 0;\n for (int i = 0; i < NUM_ATT_ACT_TENSORS; i++) {\n num_acts_params += act_sizes[i];\n }\n acts->n_acts = num_acts_params;\n}\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n // setup cublas\n hipblasHandle_t cublas_handle;\n cublasCheck(hipblasCreate(&cublas_handle));\n cublasCheck(hipblasSetMathMode(cublas_handle, HIPBLAS_DEFAULT_MATH));\n\n // read header\n FILE *model_file = fopenCheck(\"attention_block_params.bin\", \"rb\");\n int model_header[256];\n freadCheck(model_header, sizeof(int), 256, model_file);\n if (model_header[0] != 12345678) {\n fprintf(stderr, \"Invalid magic model file, magic number %d\\n\", model_header[0]);\n exit(1);\n }\n int B = model_header[1];\n int C = model_header[2];\n int H = model_header[3];\n int W = model_header[4];\n int HS = model_header[5];\n int gn_n_groups = model_header[6];\n\n AttentionParams params, grads, params_cpu;\n attention_block_count_params(¶ms, C);\n attention_block_count_params(&grads, C);\n attention_block_count_params(¶ms_cpu, C);\n\n size_t num_parameters = params.n_params;\n // read params to host\n float* params_memory_cpu = (float*)mallocCheck(num_parameters * sizeof(float));\n freadCheck(params_memory_cpu, sizeof(float), num_parameters, model_file);\n fcloseCheck(model_file);\n\n // copy params to device\n float* params_memory_gpu;\n cudaCheck(hipMalloc(¶ms_memory_gpu, num_parameters * sizeof(float)));\n cudaCheck(hipMemcpy(params_memory_gpu, params_memory_cpu, num_parameters * sizeof(float), hipMemcpyHostToDevice));\n\n // create device memory for grads\n float* grads_memory_gpu;\n cudaCheck(hipMalloc(&grads_memory_gpu, num_parameters * sizeof(float)));\n cudaCheck(hipMemset(grads_memory_gpu, 0, num_parameters * sizeof(float)));\n\n // set param pointers\n set_attention_params_pointers(¶ms, params_memory_gpu);\n set_attention_params_pointers(&grads, grads_memory_gpu);\n set_attention_params_pointers(¶ms_cpu, params_memory_cpu);\n\n // calculate activation sizes\n int C3 = 3 * C;\n int T = H * W;\n AttentionActs acts;\n attention_block_count_acts(&acts, B, C, H, W, gn_n_groups, HS);\n size_t num_acts_params = acts.n_acts;\n\n // allocate device act memory\n float* acts_memory_gpu;\n cudaCheck(hipMalloc(&acts_memory_gpu, num_acts_params * sizeof(float)));\n cudaCheck(hipMemset(acts_memory_gpu, 0, num_acts_params * sizeof(float)));\n\n // set acts pointers\n set_attention_acts_pointers(&acts, acts_memory_gpu);\n\n // malloc input, which is not part of the large block of acts memory\n cudaCheck(hipMalloc(&acts.input, B * C * H * W * sizeof(float)));\n\n // create debug states\n size_t debug_state_sizes[NUM_ATT_DEBUG_STATES];\n debug_state_sizes[0] = B * C * H * W; // inp\n debug_state_sizes[1] = B * C * H * W; // gn\n debug_state_sizes[2] = B * C * H * W; // perm1\n debug_state_sizes[3] = B * C3 * H * W; // qkv\n debug_state_sizes[4] = B * C * H * W; // att\n debug_state_sizes[5] = B * C * H * W; // proj\n debug_state_sizes[6] = B * C * H * W; // out\n debug_state_sizes[7] = B * C * H * W; // dout\n debug_state_sizes[8] = B * C * H * W; // dinp\n\n size_t num_debug_states = 0;\n for (int i = 0; i < NUM_ATT_DEBUG_STATES; i++) {\n num_debug_states += debug_state_sizes[i];\n }\n\n // allocate host debug state memory and host grads memory\n float* debug_states_memory_cpu = (float*)mallocCheck(num_debug_states * sizeof(float));\n float* grads_memory_cpu = (float*)mallocCheck(num_parameters * sizeof(float));\n AttentionParams grads_cpu;\n attention_block_count_params(&grads_cpu, C);\n set_attention_params_pointers(&grads_cpu, grads_memory_cpu);\n FILE *debug_states_file = fopenCheck(\"attention_block_states.bin\", \"rb\");\n printf(\"num debug states: %ld\\n\", num_debug_states);\n freadCheck(debug_states_memory_cpu, sizeof(float), num_debug_states, debug_states_file);\n freadCheck(grads_memory_cpu, sizeof(float), num_parameters, debug_states_file);\n fcloseCheck(debug_states_file);\n\n // set debug states pointers\n AttentionDebugStates debug_states;\n float* debug_states_ptr = debug_states_memory_cpu;\n float** debug_states_ptrs[] = {\n &debug_states.inp,\n &debug_states.gn,\n &debug_states.perm1,\n &debug_states.qkv,\n &debug_states.att,\n &debug_states.proj,\n &debug_states.out,\n &debug_states.dout,\n &debug_states.dinp\n };\n for (int i = 0; i < NUM_ATT_DEBUG_STATES; i++) {\n *(debug_states_ptrs[i]) = debug_states_ptr;\n debug_states_ptr += debug_state_sizes[i];\n }\n\n // create attention backward acts\n AttentionBackwardActs back_acts;\n attention_block_count_backs(&back_acts, B, C, H, W, HS, gn_n_groups);\n size_t num_back_act_params = back_acts.n_backs;\n\n float* back_acts_memory_gpu;\n cudaCheck(hipMalloc(&back_acts_memory_gpu, num_back_act_params * sizeof(float)));\n cudaCheck(hipMemset(back_acts_memory_gpu, 0, num_back_act_params * sizeof(float)));\n\n // set pointers\n set_attention_back_acts_pointers(&back_acts, back_acts_memory_gpu, back_acts.back_sizes);\n back_acts.dinp = acts.input;\n\n // copy input and dout to device\n cudaCheck(hipMemcpy(acts.input, debug_states.inp, B * C * H * W * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(back_acts.dout, debug_states.dout, B * C * H * W * sizeof(float), hipMemcpyHostToDevice));\n\n int block_size = 512;\n\n printf(\"Checking forward pass\\n\");\n attention_block_forward(cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, ¶ms, &acts);\n\n // validate results\n float acc = 1e-4;\n printf(\"\\nChecking gn\\n\");\n validate_result(acts.gn, debug_states.gn, \"gn\", B * C * H * W, acc);\n printf(\"\\nChecking perm1\\n\");\n validate_result(acts.perm1, debug_states.perm1, \"perm1\", B * T * C, acc);\n // don't check qkv because it is used as a buffer in attention step so values are wrong\n //printf(\"\\nChecking qkv\\n\");\n //validate_result(acts.qkv1, debug_states.qkv, \"qkv\", B * C3 * H * W, acc);\n printf(\"\\nChecking att\\n\");\n validate_result(acts.att_out, debug_states.att, \"att_out\", B * T * C, acc);\n printf(\"\\nChecking proj\\n\");\n validate_result(acts.proj, debug_states.proj, \"proj\", B * T * C, acc);\n printf(\"\\nChecking out\\n\");\n validate_result(acts.add, debug_states.out, \"out\", B * C * H * W, acc);\n\n printf(\"\\nForward successful\\n\");\n\n printf(\"Checking backward pass\\n\");\n attention_block_backward(cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, ¶ms, &acts, &back_acts, &grads);\n\n printf(\"\\nChecking dproj_w\\n\");\n validate_result(grads.proj_w, grads_cpu.proj_w, \"proj_w\", C * C, acc);\n printf(\"\\nChecking dproj_b\\n\");\n validate_result(grads.proj_b, grads_cpu.proj_b, \"proj_b\", C, acc);\n printf(\"\\nChecking dqkv_w\\n\");\n validate_result(grads.qkv_w, grads_cpu.qkv_w, \"qkv_w\", 3 * C * C, acc);\n printf(\"\\nChecking dqkv_b\\n\");\n validate_result(grads.qkv_b, grads_cpu.qkv_b, \"qkv_b\", 3 * C, acc);\n printf(\"\\nChecking dgn_w\\n\");\n validate_result(grads.gn_w, grads_cpu.gn_w, \"gn_w\", C, acc);\n printf(\"\\nChecking dgn_b\\n\");\n validate_result(grads.gn_b, grads_cpu.gn_b, \"gn_b\", C, acc);\n printf(\"\\nChecking dinp\\n\");\n validate_result(back_acts.dinp, debug_states.dinp, \"dinp\", B * C * H * W, acc);\n\n printf(\"\\nAll results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmark\\n\");\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(\n repeat_times, attention_block_forward,\n cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, ¶ms, &acts\n );\n printf(\"forward pass | time %.4f ms\\n\", elapsed_time);\n\n printf(\"Backward pass benchmark\\n\");\n elapsed_time = benchmark_kernel(\n repeat_times, attention_block_backward,\n cublas_handle, B, C, H, W, HS, gn_n_groups, block_size, ¶ms, &acts, &back_acts, &grads\n );\n printf(\"backward pass | time %.4f ms\\n\", elapsed_time);\n\n\n cublasCheck(hipblasDestroy(cublas_handle));\n // free memory\n cudaCheck(hipFree(params_memory_gpu));\n cudaCheck(hipFree(acts_memory_gpu));\n cudaCheck(hipFree(back_acts_memory_gpu));\n cudaCheck(hipFree(grads_memory_gpu));\n cudaCheck(hipFree(acts.input));\n free(params_memory_cpu);\n free(debug_states_memory_cpu);\n free(grads_memory_cpu);\n\n}\n#endif", "source": "/kaggle/working/unet.cu/dev/attention_block.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \"conv2d_k1.cuh\"\n\n\n\n__global__ void permute_x_kernel(\n const float* x_in,\n float* x_out,\n int B, int C_in, int H, int W\n) {\n int flat_idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (flat_idx < B * C_in * H * W) {\n int img_size = H * W;\n int b = flat_idx / (C_in * img_size);\n int c = (flat_idx / img_size) % C_in;\n int h = (flat_idx / W) % H;\n int w = flat_idx % W;\n\n // load from adjacent entries in x_in to enable coalescing\n int out_ind = c * B * img_size + b * img_size + h * W + w;\n x_out[out_ind] = x_in[flat_idx]; \n }\n}\n\n\n__global__ void permute_out_add_bias_kernel(\n const float* out_perm, const float* bias,\n float* out,\n int B, int C_out, int H, int W\n) {\n int flat_idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (flat_idx < B * C_out * H * W) {\n int img_size = H * W;\n int c = flat_idx / (B * img_size);\n int b = (flat_idx / img_size) % B;\n int h = (flat_idx / W) % H;\n int w = flat_idx % W;\n\n int out_ind = b * C_out * img_size + c * img_size + h * W + w;\n float val = out_perm[flat_idx];\n if (bias != NULL) {\n val += bias[c];\n }\n out[out_ind] = val;\n }\n}\n\n\n// kernel that uses cuBLAS to do the 1D conv, then adds bias by hand\nvoid conv2d_k1_forward1(\n cublasHandle_t cublas_handle,\n float* out,\n const float* x, const float* weight, const float* bias,\n int B, int C_in, int H, int W, int C_out,\n const int block_size,\n float* t1, float* t2, float* t3\n) {\n // x shape (B, C_in, H, W)\n // weight (C_out, C_in, 1, 1), which is essentially (C_out, C_in)\n // bias (C_out,)\n // out (B, C_out, H, W)\n\n // cublas stores matrices in column-major order\n // so plan is as follows:\n // 1. permute x to shape (C_in, B, H, W)\n // 2. do x @ weight -> out with cublas, where out ends up with shape (C_out, B, H, W)\n // 3. permute out to shape (B, C_out, H, W), and add bias at the same time.\n\n // 1. allocate x_perm and permute x\n cudaEvent_t start, end1, end2, end3;\n if (t1 != nullptr) {\n cudaCheck(cudaEventCreate(&start));\n cudaCheck(cudaEventCreate(&end1));\n cudaCheck(cudaEventCreate(&end2));\n cudaCheck(cudaEventCreate(&end3));\n cudaCheck(cudaEventRecord(start, nullptr));\n }\n float *x_perm, *out_perm;\n cudaCheck(cudaMalloc(&x_perm, C_in * B * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&out_perm, C_out * B * H * W * sizeof(float)));\n int total_threads1 = B * C_in * H * W;\n int num_blocks1 = ceil_div(total_threads1, block_size);\n permute_x_kernel<<>>(x, x_perm, B, C_in, H, W);\n if (t1 != nullptr) {\n cudaCheck(cudaEventRecord(end1, nullptr));\n }\n\n // 2. x @ weight with cublas\n const float alpha = 1.0f;\n const float beta = 0.0f;\n int non_C_dim = B * H * W;\n cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, non_C_dim, C_out, C_in, &alpha, x_perm, non_C_dim, weight, C_in, &beta, out_perm, non_C_dim));\n if (t2 != nullptr) {\n cudaCheck(cudaEventRecord(end2, nullptr));\n }\n\n // 3. permute back and add bias\n int total_threads2 = B * C_out * H * W;\n int num_blocks2 = ceil_div(total_threads2, block_size);\n permute_out_add_bias_kernel<<>>(out_perm, bias, out, B, C_out, H, W);\n\n // free memory\n cudaCheck(cudaFree(x_perm));\n cudaCheck(cudaFree(out_perm));\n if (t3 != nullptr) {\n cudaCheck(cudaEventRecord(end3, nullptr));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end1));\n cudaCheck(cudaEventSynchronize(end2));\n cudaCheck(cudaEventSynchronize(end3));\n float elapsed_time1, elapsed_time2, elapsed_time3;\n cudaCheck(cudaEventElapsedTime(&elapsed_time1, start, end1));\n cudaCheck(cudaEventElapsedTime(&elapsed_time2, end1, end2));\n cudaCheck(cudaEventElapsedTime(&elapsed_time3, end2, end3));\n *t1 += elapsed_time1;\n *t2 += elapsed_time2;\n *t3 += elapsed_time3;\n }\n}\n\ntemplate \n__global__ void conv2d_k1_forward_kernel(\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C, int H, int W, int O\n) {\n // constants\n const int BC = 16;\n assert(O % BO == 0);\n assert(C % BC == 0);\n assert(H * W % BH == 0);\n // n_threads per block = (BO * BH) / (TO * TH) = 256\n // n_blocks = B * (O / BO) * (H * W / BH)\n // set gridDim to (BxHxW / BH, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int block_col = blockIdx.x % blocks_per_col;\n const int batch = blockIdx.x / blocks_per_col;\n\n const int threads_per_col = BH / TH;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // shared mem for weight and x\n const int weight_s_sz = BO * BC;\n const int x_s_sz = BC * BH;\n __shared__ float w_s[weight_s_sz];\n __shared__ float x_s[x_s_sz];\n\n // local registers for results;\n float thread_results[TO * TH] = {0.0};\n float reg_w[TO] = {0.0};\n float reg_x[TH] = {0.0};\n\n // get offsets to x and weight\n int x_block_offset = batch * C * H * W + block_col * BH;\n int w_block_offset = block_row * BO * C;\n\n for (int c_idx = 0; c_idx < C; c_idx += BC) {\n int x_inner_offset = c_idx * H * W;\n int w_inner_offset = c_idx;\n // load weights to smem\n for (int i = threadIdx.x; i < weight_s_sz; i += blockDim.x) {\n int w_i = i / BC;\n int w_j = i % BC;\n int w_abs_idx = w_block_offset + w_inner_offset + w_i * C + w_j;\n w_s[w_i * BC + w_j] = weight[w_abs_idx];\n }\n for (int i = threadIdx.x; i < x_s_sz; i += blockDim.x) {\n int x_i = i / BH;\n int x_j = i % BH;\n int x_abs_idx = x_block_offset + x_inner_offset + x_i * H * W + x_j;\n x_s[x_i * BH + x_j] = x[x_abs_idx];\n }\n __syncthreads();\n\n\n // calculate thread results\n for (int dot_idx = 0; dot_idx < BC; dot_idx++) {\n for (int i = 0; i < TO; i++) {\n reg_w[i] = w_s[(thread_row * TO + i) * BC + dot_idx];\n }\n for (int j = 0; j < TH; j++) {\n reg_x[j] = x_s[dot_idx * BH + thread_col * TH + j];\n }\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n thread_results[i * TH + j] += reg_w[i] * reg_x[j];\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n // set out and bias directly to thread offset\n int out_block_offset = batch * O * H * W + block_row * BO * H * W + block_col * BH;\n int bias_offset = block_row * BO;\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n int bias_idx = bias_offset + thread_row * TO + i;\n float out_val = thread_results[i * TH + j] + bias[bias_idx];\n int out_idx = out_block_offset + thread_row * TO * H * W + thread_col * TH + i * H * W + j;\n atomicAdd(out + out_idx, out_val);\n }\n }\n}\n\nvoid conv2d_k1_forward2(\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C_in, int H, int W, int C_out\n) {\n cudaCheck(cudaMemset(out, 0, B * C_out * H * W * sizeof(float)));\n const int BO = 64;\n const int BH = 64;\n const int TO = 4;\n const int TH = 4;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k1_forward_kernel<<>>(x, weight, bias, out, B, C_in, H, W, C_out);\n}\n\n\n\n// --------------------------------------------------------------------------------\n// backward pass\n\n\n__global__ void permute_dout_x_kernel(\n const float* x, const float* dout,\n float* x_perm, float* dout_perm,\n int B, int C_in, int C_out, int H, int W\n) {\n int flat_idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n int h = (flat_idx / W) % H;\n int w = flat_idx % W;\n\n if (flat_idx < B * C_in * H * W) {\n // permute x\n int b = flat_idx / (C_in * img_size);\n int c = (flat_idx / img_size) % C_in;\n\n int out_ind = c * B * img_size + b * img_size + h * W + w;\n x_perm[out_ind] = x[flat_idx];\n }\n\n if (flat_idx < B * C_out * H * W) {\n // permute dout\n int b = flat_idx / (C_out * img_size);\n int c = (flat_idx / img_size) % C_out;\n\n int out_ind = c * B * img_size + b * img_size + h * W + w;\n dout_perm[out_ind] = dout[flat_idx];\n }\n}\n\n\n__global__ void permute_dx_kernel(\n const float* dx_perm,\n float* dx,\n int B, int C_in, int H, int W\n) {\n int flat_idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n if (flat_idx < B * C_in * img_size) {\n int c = flat_idx / (B * img_size);\n int b = (flat_idx / img_size) % B;\n int h = (flat_idx / W) % H;\n int w = flat_idx % W;\n int out_ind = b * C_in * img_size + c * img_size + h * W + w;\n dx[out_ind] = dx_perm[flat_idx];\n }\n}\n\n\n// we want to sum the B, H, W dimensions of dout\n// since dout_perm is shaped as (C_out, B, H, W)\n// we will have each block take care of one entry of C_out\n// and the threads in each block will use coorperative_groups to do the sum\n__global__ void dbias_kernel(\n const float* dout_perm, // shape (C_out, B, H, W)\n float* dbias,\n int B, int C_out, int H, int W\n) {\n int bhw = B * H * W;\n extern __shared__ float shared[];\n\n // move dout_perm to the C_out index\n if (threadIdx.x < bhw) {\n dout_perm += blockIdx.x * B * H * W;\n float th_sum = 0.0f;\n for (int i = threadIdx.x; i < bhw; i += blockDim.x) {\n th_sum += dout_perm[i];\n }\n shared[threadIdx.x] = th_sum;\n }\n __syncthreads();\n\n // reduce the shared memory\n for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {\n if (threadIdx.x < stride) {\n shared[threadIdx.x] += shared[threadIdx.x + stride];\n }\n __syncthreads();\n }\n\n if (threadIdx.x == 0) {\n dbias[blockIdx.x] = shared[0];\n }\n}\n\n\n// uses cublas\nvoid conv2d_k1_backward1(\n cublasHandle_t cublas_handle,\n const float* dout, const float* x, const float* weight,\n float* dx, float* dweight, float* dbias,\n int B, int C_in, int C_out, int H, int W,\n int block_size\n) {\n // shapes:\n // x, dx: (B, C_in, H, W)\n // weight: (C_out, C_in)\n // bias: (C_out,)\n // dout: (B, C_out, H, W)\n\n // plan (the shapes are raw array shapes, note that cublas interprets the arrays as transposed):\n // 1. transpose x to x_perm of shape (C_in, B*H*W), and dout to dout_perm of shape (C_out, B*H*W)\n // 2. do one sgemm call to get dw = x.T @ dout_perm, and another call to get dx_perm = dout @ w.T\n // 3. dx_perm is of shape (C_in, B*H*W), so call a kernel to permute it\n // 4. reduce to get dbias\n float *x_perm, *dout_perm, *dx_perm;\n cudaCheck(cudaMalloc(&x_perm, C_in * B * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&dout_perm, C_out * B * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&dx_perm, C_in * B * H * W * sizeof(float)));\n\n // 1. permute x and dout\n int total_th1 = (C_in > C_out) ? C_in * B * H * W : C_out * B * H * W;\n int n_blk1 = ceil_div(total_th1, block_size);\n permute_dout_x_kernel<<>>(x, dout, x_perm, dout_perm, B, C_in, C_out, H, W);\n\n // 2.a sgemm call for dw = x_perm.T @ dout_perm\n const float alpha = 1.0f;\n const float beta = 0.0f;\n int non_C_dim = B * H * W;\n cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, C_in, C_out, non_C_dim, &alpha, x_perm, non_C_dim, dout_perm, non_C_dim, &beta, dweight, C_in));\n\n // 2.b sgemm call for dx_perm = dout_perm @ weight.T\n cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, non_C_dim, C_in, C_out, &alpha, dout_perm, non_C_dim, weight, C_in, &beta, dx_perm, non_C_dim));\n\n // 3. permute dx_perm from (C_out, B, H, W) to (B, C_out, H, W)\n int total_th2 = C_in * B * H * W;\n int n_blk2 = ceil_div(total_th2, block_size);\n permute_dx_kernel<<>>(dx_perm, dx, B, C_in, H, W);\n\n // 4. reduce dout to get dbias\n dbias_kernel<<>>(dout_perm, dbias, B, C_out, H, W);\n\n // free memory\n cudaCheck(cudaFree(x_perm));\n cudaCheck(cudaFree(dout_perm));\n cudaCheck(cudaFree(dx_perm));\n}\n\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n srand(0);\n int B = 32;\n int C_in = 64;\n int C_out = 128;\n int H = 64;\n int W = 64;\n\n // set up the device\n int deviceIdx = 0;\n cudaCheck(cudaSetDevice(deviceIdx));\n cudaDeviceProp deviceProp;\n cudaGetDeviceProperties(&deviceProp, deviceIdx);\n printf(\"Device %d: %s\\n\", deviceIdx, deviceProp.name);\n\n // setup cublas\n cublasHandle_t cublas_handle;\n cublasCheck(cublasCreate(&cublas_handle));\n cublasCheck(cublasSetMathMode(cublas_handle, CUBLAS_DEFAULT_MATH));\n\n // create host memory to load data\n float* x = (float*)malloc(B * C_in * H * W * sizeof(float));\n float* weight = (float*)malloc(C_out * C_in * sizeof(float));\n float* bias = (float*)malloc(C_out * sizeof(float));\n float* out = (float*)malloc(B * C_out * H * W * sizeof(float));\n float* dout = (float*)malloc(B * C_out * H * W * sizeof(float));\n float* dx = (float*)malloc(B * C_in * H * W * sizeof(float));\n float* dweight = (float*)malloc(C_out * C_in * sizeof(float));\n float* dbias = (float*)malloc(C_out * sizeof(float));\n \n // read saved output\n FILE *file = fopen(\"conv2d_k1.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(x, sizeof(float), B * C_in * H * W, file);\n freadCheck(weight, sizeof(float), C_out * C_in, file);\n freadCheck(bias, sizeof(float), C_out, file);\n freadCheck(out, sizeof(float), B * C_out * H * W, file);\n freadCheck(dout, sizeof(float), B * C_out * H * W, file);\n freadCheck(dx, sizeof(float), B * C_in * H * W, file);\n freadCheck(dweight, sizeof(float), C_out * C_in, file);\n freadCheck(dbias, sizeof(float), C_out, file);\n fclose(file);\n\n // allocate device memory\n float *d_x, *d_weight, *d_bias, *d_out, *d_dout, *d_dx, *d_dweight, *d_dbias;\n cudaCheck(cudaMalloc(&d_x, B * C_in * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&d_weight, C_out * C_in * sizeof(float)));\n cudaCheck(cudaMalloc(&d_bias, C_out * sizeof(float)));\n cudaCheck(cudaMalloc(&d_out, B * C_out * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dout, B * C_out * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dx, B * C_in * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dweight, C_out * C_in * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dbias, C_out * sizeof(float)));\n cudaCheck(cudaMemcpy(d_x, x, B * C_in * H * W * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_weight, weight, C_out * C_in * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_bias, bias, C_out * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_dout, dout, B * C_out * H * W * sizeof(float), cudaMemcpyHostToDevice));\n\n //int block_sizes[] = {128, 256, 512, 1024};\n int block_sizes[] = {512};\n printf(\"Checking forward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n printf(\"Checking block size %d\\n\", block_sizes[j]);\n int block_size = block_sizes[j];\n conv2d_k1_forward1(cublas_handle, d_out, d_x, d_weight, d_bias, B, C_in, H, W, C_out, block_size);\n validate_result(d_out, out, \"out\", B * C_out * H * W);\n }\n printf(\"\\nChecking new kernel\\n\");\n conv2d_k1_forward2(d_x, d_weight, d_bias, d_out, B, C_in, H, W, C_out);\n validate_result(d_out, out, \"new_out\", B * C_out * H * W);\n printf(\"Running our own forward benchmark for new kernel\\n\");\n void* flush_buffer;\n cudaCheck(cudaMalloc(&flush_buffer, deviceProp.l2CacheSize));\n int repeat_times = 1000;\n float new_time = 0.0f;\n for (int i = 0; i < repeat_times; i++) {\n // clear L2\n cudaCheck(cudaMemset(flush_buffer, 0, deviceProp.l2CacheSize));\n cudaEvent_t start, end;\n cudaCheck(cudaEventCreate(&start));\n cudaCheck(cudaEventCreate(&end));\n cudaCheck(cudaEventRecord(start, nullptr));\n conv2d_k1_forward2(d_x, d_weight, d_bias, d_out, B, C_in, H, W, C_out);\n cudaCheck(cudaEventRecord(end, nullptr));\n cudaCheck(cudaEventSynchronize(start));\n cudaCheck(cudaEventSynchronize(end));\n float elapsed_time;\n cudaCheck(cudaEventElapsedTime(&elapsed_time, start, end));\n new_time += elapsed_time;\n }\n new_time /= repeat_times;\n printf(\"new kernel time: %.4f ms\\n\", new_time);\n \n printf(\"Forward pass successful\\n\\n\");\n\n printf(\"Checking backward pass\\n\");\n float accuracy = 1e-1;\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"\\nChecking block size %d\\n\", block_size);\n conv2d_k1_backward1(cublas_handle, d_dout, d_x, d_weight, d_dx, d_dweight, d_dbias, B, C_in, C_out, H, W, block_size);\n printf(\"\\nChecking dweight\\n\");\n validate_result(d_dweight, dweight, \"dweight\", C_out * C_in, accuracy);\n printf(\"\\nchecking dbias\\n\");\n validate_result(d_dbias, dbias, \"dbias\", C_out, accuracy);\n printf(\"\\nchecking dx\\n\");\n validate_result(d_dx, dx, \"dx\", B * C_in * H * W);\n }\n printf(\"Backward pass successful\\n\\n\");\n\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n //printf(\"Forward pass benchmarks:\\n\");\n //for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n // int block_size = block_sizes[j];\n\n // int repeat_times = 100;\n // float elapsed_time = benchmark_kernel(repeat_times, conv2d_k1_forward1,\n // cublas_handle, d_out, d_x, d_weight, d_bias,\n // B, C_in, H, W, C_out, block_size);\n\n // // napkin math: estimate the flops achieved\n // // e.g. A100 40GB PCIe is advertised at 19.5 TFLOPS fp32\n // float tflops = (float)B * H * W * C_in * C_out * 2 / elapsed_time * 1e3f / 1e12f;\n // printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n //}\n printf(\"Running our own forward benchmarks\\n\");\n float t1 = 0.0f;\n float t2 = 0.0f;\n float t3 = 0.0f;\n float t_total = 0.0f;\n for (int i = 0; i < repeat_times; i++) {\n // clear L2\n cudaCheck(cudaMemset(flush_buffer, 0, deviceProp.l2CacheSize));\n conv2d_k1_forward1(cublas_handle, d_out, d_x, d_weight, d_bias, B, C_in, H, W, C_out, 512, &t1, &t2, &t3);\n }\n cudaCheck(cudaFree(flush_buffer));\n t1 /= repeat_times;\n t2 /= repeat_times;\n t3 /= repeat_times;\n t_total = t1 + t2 + t3;\n\n printf(\"t1: %.4f ms\\n\", t1);\n printf(\"t2: %.4f ms\\n\", t2);\n printf(\"t3: %.4f ms\\n\", t3);\n printf(\"t_total: %.4f ms\\n\", t_total);\n\n printf(\"\\nBackward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, conv2d_k1_backward1,\n cublas_handle, d_dout, d_x, d_weight, d_dx, d_dweight, d_dbias,\n B, C_in, C_out, H, W, block_size);\n\n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n\n\n // free memory\n free(x);\n free(weight);\n free(bias);\n free(out);\n free(dout);\n free(dx);\n free(dweight);\n free(dbias);\n cudaCheck(cudaFree(d_x));\n cudaCheck(cudaFree(d_weight));\n cudaCheck(cudaFree(d_bias));\n cudaCheck(cudaFree(d_out));\n cudaCheck(cudaFree(d_dout));\n cudaCheck(cudaFree(d_dx));\n cudaCheck(cudaFree(d_dweight));\n cudaCheck(cudaFree(d_dbias));\n cublasCheck(cublasDestroy(cublas_handle));\n return 0;\n}\n#endif", "output": "#include \n#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \"conv2d_k1.cuh\"\n\n\n\n__global__ void permute_x_kernel(\n const float* x_in,\n float* x_out,\n int B, int C_in, int H, int W\n) {\n int flat_idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (flat_idx < B * C_in * H * W) {\n int img_size = H * W;\n int b = flat_idx / (C_in * img_size);\n int c = (flat_idx / img_size) % C_in;\n int h = (flat_idx / W) % H;\n int w = flat_idx % W;\n\n // load from adjacent entries in x_in to enable coalescing\n int out_ind = c * B * img_size + b * img_size + h * W + w;\n x_out[out_ind] = x_in[flat_idx]; \n }\n}\n\n\n__global__ void permute_out_add_bias_kernel(\n const float* out_perm, const float* bias,\n float* out,\n int B, int C_out, int H, int W\n) {\n int flat_idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (flat_idx < B * C_out * H * W) {\n int img_size = H * W;\n int c = flat_idx / (B * img_size);\n int b = (flat_idx / img_size) % B;\n int h = (flat_idx / W) % H;\n int w = flat_idx % W;\n\n int out_ind = b * C_out * img_size + c * img_size + h * W + w;\n float val = out_perm[flat_idx];\n if (bias != NULL) {\n val += bias[c];\n }\n out[out_ind] = val;\n }\n}\n\n\n// kernel that uses cuBLAS to do the 1D conv, then adds bias by hand\nvoid conv2d_k1_forward1(\n hipblasHandle_t cublas_handle,\n float* out,\n const float* x, const float* weight, const float* bias,\n int B, int C_in, int H, int W, int C_out,\n const int block_size,\n float* t1, float* t2, float* t3\n) {\n // x shape (B, C_in, H, W)\n // weight (C_out, C_in, 1, 1), which is essentially (C_out, C_in)\n // bias (C_out,)\n // out (B, C_out, H, W)\n\n // cublas stores matrices in column-major order\n // so plan is as follows:\n // 1. permute x to shape (C_in, B, H, W)\n // 2. do x @ weight -> out with cublas, where out ends up with shape (C_out, B, H, W)\n // 3. permute out to shape (B, C_out, H, W), and add bias at the same time.\n\n // 1. allocate x_perm and permute x\n hipEvent_t start, end1, end2, end3;\n if (t1 != nullptr) {\n cudaCheck(hipEventCreate(&start));\n cudaCheck(hipEventCreate(&end1));\n cudaCheck(hipEventCreate(&end2));\n cudaCheck(hipEventCreate(&end3));\n cudaCheck(hipEventRecord(start, nullptr));\n }\n float *x_perm, *out_perm;\n cudaCheck(hipMalloc(&x_perm, C_in * B * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&out_perm, C_out * B * H * W * sizeof(float)));\n int total_threads1 = B * C_in * H * W;\n int num_blocks1 = ceil_div(total_threads1, block_size);\n permute_x_kernel<<>>(x, x_perm, B, C_in, H, W);\n if (t1 != nullptr) {\n cudaCheck(hipEventRecord(end1, nullptr));\n }\n\n // 2. x @ weight with cublas\n const float alpha = 1.0f;\n const float beta = 0.0f;\n int non_C_dim = B * H * W;\n cublasCheck(hipblasSgemm(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_N, non_C_dim, C_out, C_in, &alpha, x_perm, non_C_dim, weight, C_in, &beta, out_perm, non_C_dim));\n if (t2 != nullptr) {\n cudaCheck(hipEventRecord(end2, nullptr));\n }\n\n // 3. permute back and add bias\n int total_threads2 = B * C_out * H * W;\n int num_blocks2 = ceil_div(total_threads2, block_size);\n permute_out_add_bias_kernel<<>>(out_perm, bias, out, B, C_out, H, W);\n\n // free memory\n cudaCheck(hipFree(x_perm));\n cudaCheck(hipFree(out_perm));\n if (t3 != nullptr) {\n cudaCheck(hipEventRecord(end3, nullptr));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end1));\n cudaCheck(hipEventSynchronize(end2));\n cudaCheck(hipEventSynchronize(end3));\n float elapsed_time1, elapsed_time2, elapsed_time3;\n cudaCheck(hipEventElapsedTime(&elapsed_time1, start, end1));\n cudaCheck(hipEventElapsedTime(&elapsed_time2, end1, end2));\n cudaCheck(hipEventElapsedTime(&elapsed_time3, end2, end3));\n *t1 += elapsed_time1;\n *t2 += elapsed_time2;\n *t3 += elapsed_time3;\n }\n}\n\ntemplate \n__global__ void conv2d_k1_forward_kernel(\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C, int H, int W, int O\n) {\n // constants\n const int BC = 16;\n assert(O % BO == 0);\n assert(C % BC == 0);\n assert(H * W % BH == 0);\n // n_threads per block = (BO * BH) / (TO * TH) = 256\n // n_blocks = B * (O / BO) * (H * W / BH)\n // set gridDim to (BxHxW / BH, O / BO)\n const int block_row = blockIdx.y;\n const int blocks_per_col = H * W / BH;\n assert(gridDim.x == B * blocks_per_col);\n const int block_col = blockIdx.x % blocks_per_col;\n const int batch = blockIdx.x / blocks_per_col;\n\n const int threads_per_col = BH / TH;\n const int thread_col = threadIdx.x % threads_per_col;\n const int thread_row = threadIdx.x / threads_per_col;\n\n // shared mem for weight and x\n const int weight_s_sz = BO * BC;\n const int x_s_sz = BC * BH;\n __shared__ float w_s[weight_s_sz];\n __shared__ float x_s[x_s_sz];\n\n // local registers for results;\n float thread_results[TO * TH] = {0.0};\n float reg_w[TO] = {0.0};\n float reg_x[TH] = {0.0};\n\n // get offsets to x and weight\n int x_block_offset = batch * C * H * W + block_col * BH;\n int w_block_offset = block_row * BO * C;\n\n for (int c_idx = 0; c_idx < C; c_idx += BC) {\n int x_inner_offset = c_idx * H * W;\n int w_inner_offset = c_idx;\n // load weights to smem\n for (int i = threadIdx.x; i < weight_s_sz; i += blockDim.x) {\n int w_i = i / BC;\n int w_j = i % BC;\n int w_abs_idx = w_block_offset + w_inner_offset + w_i * C + w_j;\n w_s[w_i * BC + w_j] = weight[w_abs_idx];\n }\n for (int i = threadIdx.x; i < x_s_sz; i += blockDim.x) {\n int x_i = i / BH;\n int x_j = i % BH;\n int x_abs_idx = x_block_offset + x_inner_offset + x_i * H * W + x_j;\n x_s[x_i * BH + x_j] = x[x_abs_idx];\n }\n __syncthreads();\n\n\n // calculate thread results\n for (int dot_idx = 0; dot_idx < BC; dot_idx++) {\n for (int i = 0; i < TO; i++) {\n reg_w[i] = w_s[(thread_row * TO + i) * BC + dot_idx];\n }\n for (int j = 0; j < TH; j++) {\n reg_x[j] = x_s[dot_idx * BH + thread_col * TH + j];\n }\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n thread_results[i * TH + j] += reg_w[i] * reg_x[j];\n }\n }\n }\n __syncthreads();\n }\n\n // write out results\n // set out and bias directly to thread offset\n int out_block_offset = batch * O * H * W + block_row * BO * H * W + block_col * BH;\n int bias_offset = block_row * BO;\n for (int i = 0; i < TO; i++) {\n for (int j = 0; j < TH; j++) {\n int bias_idx = bias_offset + thread_row * TO + i;\n float out_val = thread_results[i * TH + j] + bias[bias_idx];\n int out_idx = out_block_offset + thread_row * TO * H * W + thread_col * TH + i * H * W + j;\n atomicAdd(out + out_idx, out_val);\n }\n }\n}\n\nvoid conv2d_k1_forward2(\n const float* x, const float* weight, const float* bias,\n float* out,\n int B, int C_in, int H, int W, int C_out\n) {\n cudaCheck(hipMemset(out, 0, B * C_out * H * W * sizeof(float)));\n const int BO = 64;\n const int BH = 64;\n const int TO = 4;\n const int TH = 4;\n dim3 gridDim(B * H * W / BH, C_out / BO);\n dim3 blockDim((BO * BH) / (TO * TH));\n conv2d_k1_forward_kernel<<>>(x, weight, bias, out, B, C_in, H, W, C_out);\n}\n\n\n\n// --------------------------------------------------------------------------------\n// backward pass\n\n\n__global__ void permute_dout_x_kernel(\n const float* x, const float* dout,\n float* x_perm, float* dout_perm,\n int B, int C_in, int C_out, int H, int W\n) {\n int flat_idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n int h = (flat_idx / W) % H;\n int w = flat_idx % W;\n\n if (flat_idx < B * C_in * H * W) {\n // permute x\n int b = flat_idx / (C_in * img_size);\n int c = (flat_idx / img_size) % C_in;\n\n int out_ind = c * B * img_size + b * img_size + h * W + w;\n x_perm[out_ind] = x[flat_idx];\n }\n\n if (flat_idx < B * C_out * H * W) {\n // permute dout\n int b = flat_idx / (C_out * img_size);\n int c = (flat_idx / img_size) % C_out;\n\n int out_ind = c * B * img_size + b * img_size + h * W + w;\n dout_perm[out_ind] = dout[flat_idx];\n }\n}\n\n\n__global__ void permute_dx_kernel(\n const float* dx_perm,\n float* dx,\n int B, int C_in, int H, int W\n) {\n int flat_idx = blockIdx.x * blockDim.x + threadIdx.x;\n int img_size = H * W;\n if (flat_idx < B * C_in * img_size) {\n int c = flat_idx / (B * img_size);\n int b = (flat_idx / img_size) % B;\n int h = (flat_idx / W) % H;\n int w = flat_idx % W;\n int out_ind = b * C_in * img_size + c * img_size + h * W + w;\n dx[out_ind] = dx_perm[flat_idx];\n }\n}\n\n\n// we want to sum the B, H, W dimensions of dout\n// since dout_perm is shaped as (C_out, B, H, W)\n// we will have each block take care of one entry of C_out\n// and the threads in each block will use coorperative_groups to do the sum\n__global__ void dbias_kernel(\n const float* dout_perm, // shape (C_out, B, H, W)\n float* dbias,\n int B, int C_out, int H, int W\n) {\n int bhw = B * H * W;\n extern __shared__ float shared[];\n\n // move dout_perm to the C_out index\n if (threadIdx.x < bhw) {\n dout_perm += blockIdx.x * B * H * W;\n float th_sum = 0.0f;\n for (int i = threadIdx.x; i < bhw; i += blockDim.x) {\n th_sum += dout_perm[i];\n }\n shared[threadIdx.x] = th_sum;\n }\n __syncthreads();\n\n // reduce the shared memory\n for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {\n if (threadIdx.x < stride) {\n shared[threadIdx.x] += shared[threadIdx.x + stride];\n }\n __syncthreads();\n }\n\n if (threadIdx.x == 0) {\n dbias[blockIdx.x] = shared[0];\n }\n}\n\n\n// uses cublas\nvoid conv2d_k1_backward1(\n hipblasHandle_t cublas_handle,\n const float* dout, const float* x, const float* weight,\n float* dx, float* dweight, float* dbias,\n int B, int C_in, int C_out, int H, int W,\n int block_size\n) {\n // shapes:\n // x, dx: (B, C_in, H, W)\n // weight: (C_out, C_in)\n // bias: (C_out,)\n // dout: (B, C_out, H, W)\n\n // plan (the shapes are raw array shapes, note that cublas interprets the arrays as transposed):\n // 1. transpose x to x_perm of shape (C_in, B*H*W), and dout to dout_perm of shape (C_out, B*H*W)\n // 2. do one sgemm call to get dw = x.T @ dout_perm, and another call to get dx_perm = dout @ w.T\n // 3. dx_perm is of shape (C_in, B*H*W), so call a kernel to permute it\n // 4. reduce to get dbias\n float *x_perm, *dout_perm, *dx_perm;\n cudaCheck(hipMalloc(&x_perm, C_in * B * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&dout_perm, C_out * B * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&dx_perm, C_in * B * H * W * sizeof(float)));\n\n // 1. permute x and dout\n int total_th1 = (C_in > C_out) ? C_in * B * H * W : C_out * B * H * W;\n int n_blk1 = ceil_div(total_th1, block_size);\n permute_dout_x_kernel<<>>(x, dout, x_perm, dout_perm, B, C_in, C_out, H, W);\n\n // 2.a sgemm call for dw = x_perm.T @ dout_perm\n const float alpha = 1.0f;\n const float beta = 0.0f;\n int non_C_dim = B * H * W;\n cublasCheck(hipblasSgemm(cublas_handle, HIPBLAS_OP_T, HIPBLAS_OP_N, C_in, C_out, non_C_dim, &alpha, x_perm, non_C_dim, dout_perm, non_C_dim, &beta, dweight, C_in));\n\n // 2.b sgemm call for dx_perm = dout_perm @ weight.T\n cublasCheck(hipblasSgemm(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_T, non_C_dim, C_in, C_out, &alpha, dout_perm, non_C_dim, weight, C_in, &beta, dx_perm, non_C_dim));\n\n // 3. permute dx_perm from (C_out, B, H, W) to (B, C_out, H, W)\n int total_th2 = C_in * B * H * W;\n int n_blk2 = ceil_div(total_th2, block_size);\n permute_dx_kernel<<>>(dx_perm, dx, B, C_in, H, W);\n\n // 4. reduce dout to get dbias\n dbias_kernel<<>>(dout_perm, dbias, B, C_out, H, W);\n\n // free memory\n cudaCheck(hipFree(x_perm));\n cudaCheck(hipFree(dout_perm));\n cudaCheck(hipFree(dx_perm));\n}\n\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n srand(0);\n int B = 32;\n int C_in = 64;\n int C_out = 128;\n int H = 64;\n int W = 64;\n\n // set up the device\n int deviceIdx = 0;\n cudaCheck(hipSetDevice(deviceIdx));\n hipDeviceProp_t deviceProp;\n hipGetDeviceProperties(&deviceProp, deviceIdx);\n printf(\"Device %d: %s\\n\", deviceIdx, deviceProp.name);\n\n // setup cublas\n hipblasHandle_t cublas_handle;\n cublasCheck(hipblasCreate(&cublas_handle));\n cublasCheck(hipblasSetMathMode(cublas_handle, HIPBLAS_DEFAULT_MATH));\n\n // create host memory to load data\n float* x = (float*)malloc(B * C_in * H * W * sizeof(float));\n float* weight = (float*)malloc(C_out * C_in * sizeof(float));\n float* bias = (float*)malloc(C_out * sizeof(float));\n float* out = (float*)malloc(B * C_out * H * W * sizeof(float));\n float* dout = (float*)malloc(B * C_out * H * W * sizeof(float));\n float* dx = (float*)malloc(B * C_in * H * W * sizeof(float));\n float* dweight = (float*)malloc(C_out * C_in * sizeof(float));\n float* dbias = (float*)malloc(C_out * sizeof(float));\n \n // read saved output\n FILE *file = fopen(\"conv2d_k1.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(x, sizeof(float), B * C_in * H * W, file);\n freadCheck(weight, sizeof(float), C_out * C_in, file);\n freadCheck(bias, sizeof(float), C_out, file);\n freadCheck(out, sizeof(float), B * C_out * H * W, file);\n freadCheck(dout, sizeof(float), B * C_out * H * W, file);\n freadCheck(dx, sizeof(float), B * C_in * H * W, file);\n freadCheck(dweight, sizeof(float), C_out * C_in, file);\n freadCheck(dbias, sizeof(float), C_out, file);\n fclose(file);\n\n // allocate device memory\n float *d_x, *d_weight, *d_bias, *d_out, *d_dout, *d_dx, *d_dweight, *d_dbias;\n cudaCheck(hipMalloc(&d_x, B * C_in * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&d_weight, C_out * C_in * sizeof(float)));\n cudaCheck(hipMalloc(&d_bias, C_out * sizeof(float)));\n cudaCheck(hipMalloc(&d_out, B * C_out * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&d_dout, B * C_out * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&d_dx, B * C_in * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&d_dweight, C_out * C_in * sizeof(float)));\n cudaCheck(hipMalloc(&d_dbias, C_out * sizeof(float)));\n cudaCheck(hipMemcpy(d_x, x, B * C_in * H * W * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_weight, weight, C_out * C_in * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_bias, bias, C_out * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_dout, dout, B * C_out * H * W * sizeof(float), hipMemcpyHostToDevice));\n\n //int block_sizes[] = {128, 256, 512, 1024};\n int block_sizes[] = {512};\n printf(\"Checking forward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n printf(\"Checking block size %d\\n\", block_sizes[j]);\n int block_size = block_sizes[j];\n conv2d_k1_forward1(cublas_handle, d_out, d_x, d_weight, d_bias, B, C_in, H, W, C_out, block_size);\n validate_result(d_out, out, \"out\", B * C_out * H * W);\n }\n printf(\"\\nChecking new kernel\\n\");\n conv2d_k1_forward2(d_x, d_weight, d_bias, d_out, B, C_in, H, W, C_out);\n validate_result(d_out, out, \"new_out\", B * C_out * H * W);\n printf(\"Running our own forward benchmark for new kernel\\n\");\n void* flush_buffer;\n cudaCheck(hipMalloc(&flush_buffer, deviceProp.l2CacheSize));\n int repeat_times = 1000;\n float new_time = 0.0f;\n for (int i = 0; i < repeat_times; i++) {\n // clear L2\n cudaCheck(hipMemset(flush_buffer, 0, deviceProp.l2CacheSize));\n hipEvent_t start, end;\n cudaCheck(hipEventCreate(&start));\n cudaCheck(hipEventCreate(&end));\n cudaCheck(hipEventRecord(start, nullptr));\n conv2d_k1_forward2(d_x, d_weight, d_bias, d_out, B, C_in, H, W, C_out);\n cudaCheck(hipEventRecord(end, nullptr));\n cudaCheck(hipEventSynchronize(start));\n cudaCheck(hipEventSynchronize(end));\n float elapsed_time;\n cudaCheck(hipEventElapsedTime(&elapsed_time, start, end));\n new_time += elapsed_time;\n }\n new_time /= repeat_times;\n printf(\"new kernel time: %.4f ms\\n\", new_time);\n \n printf(\"Forward pass successful\\n\\n\");\n\n printf(\"Checking backward pass\\n\");\n float accuracy = 1e-1;\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"\\nChecking block size %d\\n\", block_size);\n conv2d_k1_backward1(cublas_handle, d_dout, d_x, d_weight, d_dx, d_dweight, d_dbias, B, C_in, C_out, H, W, block_size);\n printf(\"\\nChecking dweight\\n\");\n validate_result(d_dweight, dweight, \"dweight\", C_out * C_in, accuracy);\n printf(\"\\nchecking dbias\\n\");\n validate_result(d_dbias, dbias, \"dbias\", C_out, accuracy);\n printf(\"\\nchecking dx\\n\");\n validate_result(d_dx, dx, \"dx\", B * C_in * H * W);\n }\n printf(\"Backward pass successful\\n\\n\");\n\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n //printf(\"Forward pass benchmarks:\\n\");\n //for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n // int block_size = block_sizes[j];\n\n // int repeat_times = 100;\n // float elapsed_time = benchmark_kernel(repeat_times, conv2d_k1_forward1,\n // cublas_handle, d_out, d_x, d_weight, d_bias,\n // B, C_in, H, W, C_out, block_size);\n\n // // napkin math: estimate the flops achieved\n // // e.g. A100 40GB PCIe is advertised at 19.5 TFLOPS fp32\n // float tflops = (float)B * H * W * C_in * C_out * 2 / elapsed_time * 1e3f / 1e12f;\n // printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n //}\n printf(\"Running our own forward benchmarks\\n\");\n float t1 = 0.0f;\n float t2 = 0.0f;\n float t3 = 0.0f;\n float t_total = 0.0f;\n for (int i = 0; i < repeat_times; i++) {\n // clear L2\n cudaCheck(hipMemset(flush_buffer, 0, deviceProp.l2CacheSize));\n conv2d_k1_forward1(cublas_handle, d_out, d_x, d_weight, d_bias, B, C_in, H, W, C_out, 512, &t1, &t2, &t3);\n }\n cudaCheck(hipFree(flush_buffer));\n t1 /= repeat_times;\n t2 /= repeat_times;\n t3 /= repeat_times;\n t_total = t1 + t2 + t3;\n\n printf(\"t1: %.4f ms\\n\", t1);\n printf(\"t2: %.4f ms\\n\", t2);\n printf(\"t3: %.4f ms\\n\", t3);\n printf(\"t_total: %.4f ms\\n\", t_total);\n\n printf(\"\\nBackward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, conv2d_k1_backward1,\n cublas_handle, d_dout, d_x, d_weight, d_dx, d_dweight, d_dbias,\n B, C_in, C_out, H, W, block_size);\n\n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n\n\n // free memory\n free(x);\n free(weight);\n free(bias);\n free(out);\n free(dout);\n free(dx);\n free(dweight);\n free(dbias);\n cudaCheck(hipFree(d_x));\n cudaCheck(hipFree(d_weight));\n cudaCheck(hipFree(d_bias));\n cudaCheck(hipFree(d_out));\n cudaCheck(hipFree(d_dout));\n cudaCheck(hipFree(d_dx));\n cudaCheck(hipFree(d_dweight));\n cudaCheck(hipFree(d_dbias));\n cublasCheck(hipblasDestroy(cublas_handle));\n return 0;\n}\n#endif", "source": "/kaggle/working/unet.cu/dev/conv2d_k1.cu"} {"input": "#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \"linear.cuh\"\n\n// ----------------------------------------------------------------------------\n// GPU kernels\n\n// copy the matmul forward and backward kernels from llm.c\n\n// is there no better way other than just adding bias with a whole separate kernel?\n// this is a highly memory-bound operation, should be fused into the matmul kernel\n// but i can't seem to find a cuBLAS function that does this\n__global__ void add_bias(float* out, const float* bias, int N, int OC) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = idx; i < N * OC; i += stride) {\n int col = i % OC;\n out[i] += bias[col];\n }\n}\n\n// kernel 2 calls cuBLAS, which should be very efficient\nvoid matmul_forward2(cublasHandle_t cublas_handle,\n float* out,\n const float* inp, const float* weight, const float* bias,\n int N, int C, int OC,\n const int block_size) {\n // for reference API is:\n // cublasStatus_t cublasSgemm(cublasHandle_t handle,\n // cublasOperation_t transa, cublasOperation_t transb,\n // int m, int n, int k,\n // const float *alpha,\n // const float *A, int lda,\n // const float *B, int ldb,\n // const float *beta,\n // float *C, int ldc)\n // for us, inp is (B*T, C), weight is (OC, C), out is (B*T, OC)\n // cuBLAS does C = alpha * A * B + beta * C\n // where A is mxk, B is kxn, C is mxn\n // now, because we use row-major storage, cuBLAS (which is column-major) sees our matrices transposed.\n // algorithmically / in e.g. PyTorch we want to do: out = inp @ weight.T\n // but because cuBLAS is column-major, we actually want to get it to calculate out.T . Mathematically, this is:\n // out.T = weight @ inp.T\n // but again, our variables look transposed, so using the actual weight/inp we have here in this function, this becomes\n // out.T = weight.T @ inp\n // so we need to get cuBLAS to calculate weight.T @ inp (the variables here are the actual ones in this function)\n // => need to call cuBLAS with A = weight, B = inp\n // => need to call cuBLAS with transa = CUBLAS_OP_T, transb = CUBLAS_OP_N\n\n const float alpha = 1.0f;\n const float beta = 0.0f;\n cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, OC, N, C, &alpha, weight, C, inp, C, &beta, out, OC));\n // and now we still have to add the bias... (ew)\n if (bias != NULL) {\n int grid_size = ceil_div(OC * N, block_size);\n add_bias<<>>(out, bias, N, OC);\n cudaCheck(cudaGetLastError());\n }\n}\n\n// use shared memory and coarsening + reductions\n__global__ void matmul_backward_bias_kernel_faster(float* dbias, const float* dout, int N, int OC) {\n extern __shared__ float shared[];\n int o = blockIdx.x; // range [0, OC)\n int tid = threadIdx.x; // range [0, block_size)\n int block_size = blockDim.x;\n const float* x = dout + o;\n // thread coarsening\n double sum = 0.0;\n for (int i = tid; i < N; i += block_size) {\n sum += x[i * OC];\n }\n shared[tid] = (float) sum;\n __syncthreads();\n // reductions\n for (int stride = block_size / 2; stride >= 1; stride /= 2) {\n __syncthreads();\n if (tid < stride) {\n shared[tid] += shared[tid + stride];\n }\n }\n // write the final result (at thread 0) to global memory\n if (tid == 0) {\n dbias[o] = shared[0];\n }\n}\n\n// version1: simple cuBLAS calls\nvoid matmul_backward1(cublasHandle_t cublas_handle,\n float* dinp, float* dweight, float* dbias,\n float* dout, float* inp, float* weight,\n int N, int C, int OC) {\n float alpha = 1.0f;\n float beta = 0.0f; // note we must use beta = 1.0 so that we do a +=, as we should, because gradients add\n\n // for reference the API is:\n // cublasStatus_t cublasSgemm(cublasHandle_t handle,\n // cublasOperation_t transa, cublasOperation_t transb,\n // int m, int n, int k,\n // const float *alpha,\n // const float *A, int lda,\n // const float *B, int ldb,\n // const float *beta,\n // float *C, int ldc)\n\n // recall the forward pass was calculated with alpha = 1.0f, beta = 0.0f as:\n // cublasSgemm(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, OC, B*T, C, &alpha, weight, C, inp, C, &beta, out, OC);\n\n // backward to input\n cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, C, N, OC, &alpha, weight, C, dout, OC, &beta, dinp, C));\n // backward to weight\n cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, C, OC, N, &alpha, inp, C, dout, OC, &beta, dweight, C));\n // backward to bias, if given\n if (dbias != NULL) {\n\n // sum over B,T using matrix vector multiplication with cuBLAS\n // for reference this API is:\n // cublasStatus_t cublasSgemv(cublasHandle_t handle, cublasOperation_t trans,\n // int m, int n,\n // const float *alpha,\n // const float *A, int lda,\n // const float *x, int incx,\n // const float *beta,\n // float *y, int incy)\n // dout is (B,T,OC), or in 2D terms (B*T, OC)\n // cublasCheck(cublasSgemv(cublas_handle, CUBLAS_OP_N, B*T, OC, &alpha, dout, B*T, ones, 1, &beta, dbias, 1));\n // cublasCheck(cublasSgemv(cublas_handle, CUBLAS_OP_T, OC, B*T, &alpha, dout, OC, ones, 1, &beta, dbias, 1));\n\n // ugh the above isn't working...\n // let's just do naive calculation for now, fix later\n // const int block_size=128;\n // const int grid_size=(OC + block_size - 1) / block_size;\n // matmul_backward_bias_kernel<<>>(dbias, dout, B, T, OC);\n\n // bit faster\n const int block_size=512;\n dim3 block_dim(block_size);\n dim3 grid_dim(OC);\n size_t shared_mem_size = block_size * sizeof(float);\n matmul_backward_bias_kernel_faster<<>>(dbias, dout, N, OC);\n }\n}\n\n// ----------------------------------------------------------------------------\n\nvoid linear_set_param_ptrs(\n LinearParams* params,\n float* params_memory,\n int C, int OC\n) {\n params->w = params_memory;\n params_memory += OC * C;\n params->b = params_memory;\n}\n\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n int N = 32;\n int C = 64;\n int OC = 128;\n\n // setup cublas\n cublasHandle_t cublas_handle;\n cublasCheck(cublasCreate(&cublas_handle));\n cublasCheck(cublasSetMathMode(cublas_handle, CUBLAS_DEFAULT_MATH));\n\n\n // create host memory\n float *inp = (float*)malloc(N * C * sizeof(float));\n float *weight = (float*)malloc(OC * C * sizeof(float));\n float *bias = (float*)malloc(OC * sizeof(float));\n float *out = (float*)malloc(N * OC * sizeof(float));\n float *dout = (float*)malloc(N * OC * sizeof(float));\n float *dinp = (float*)malloc(N * C * sizeof(float));\n float *dweight = (float*)malloc(OC * C * sizeof(float));\n float *dbias = (float*)malloc(OC * sizeof(float));\n\n FILE *file = fopen(\"linear.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(inp, sizeof(float), N * C, file);\n freadCheck(weight, sizeof(float), OC * C, file);\n freadCheck(bias, sizeof(float), OC, file);\n freadCheck(out, sizeof(float), N * OC, file);\n freadCheck(dout, sizeof(float), N * OC, file);\n freadCheck(dinp, sizeof(float), N * C, file);\n freadCheck(dweight, sizeof(float), OC * C, file);\n freadCheck(dbias, sizeof(float), OC, file);\n fclose(file);\n\n\n // allocate device memory\n float *d_inp, *d_weight, *d_bias, *d_out, *d_dout, *d_dinp, *d_dweight, *d_dbias;\n cudaCheck(cudaMalloc(&d_inp, N * C * sizeof(float)));\n cudaCheck(cudaMalloc(&d_weight, OC * C * sizeof(float)));\n cudaCheck(cudaMalloc(&d_bias, OC * sizeof(float)));\n cudaCheck(cudaMalloc(&d_out, N * OC * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dout, N * OC * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dinp, N * C * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dweight, OC * C * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dbias, OC * sizeof(float)));\n\n cudaCheck(cudaMemcpy(d_inp, inp, N * C * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_weight, weight, OC * C * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_bias, bias, OC * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_dout, dout, N * OC * sizeof(float), cudaMemcpyHostToDevice));\n\n int block_sizes[] = {128, 256, 512, 1024};\n\n printf(\"Checking forward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"\\nBlock size: %d\\n\", block_size);\n matmul_forward2(cublas_handle, d_out, d_inp, d_weight, d_bias, N, C, OC, block_size);\n validate_result(d_out, out, \"out\", N * OC);\n }\n\n printf(\"Forward pass successful\\n\");\n\n printf(\"Checking backward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"\\nBlock size: %d\\n\", block_size);\n matmul_backward1(cublas_handle, d_dinp, d_dweight, d_dbias, d_dout, d_inp, d_weight, N, C, OC);\n printf(\"Checking dbias\\n\");\n validate_result(d_dbias, dbias, \"dbias\", OC);\n printf(\"Checking dinp\\n\");\n validate_result(d_dinp, dinp, \"dinp\", N * C);\n printf(\"Checking dweight\\n\");\n validate_result(d_dweight, dweight, \"dweight\", OC * C);\n }\n printf(\"Backward pass successful\\n\");\n\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, matmul_forward2,\n cublas_handle, d_out, d_inp, d_weight, d_bias,\n N, C, OC, block_size);\n\n float tflops = (float)N * C * OC * 2 / elapsed_time * 1e3f / 1e12f;\n printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n }\n\n printf(\"\\nBackward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, matmul_backward1,\n cublas_handle, d_dinp, d_dweight, d_dbias, d_dout, d_inp, d_weight,\n N, C, OC);\n\n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n\n // free memory\n free(inp);\n free(weight);\n free(bias);\n free(out);\n free(dout);\n free(dinp);\n free(dweight);\n free(dbias);\n\n cudaCheck(cudaFree(d_inp));\n cudaCheck(cudaFree(d_weight));\n cudaCheck(cudaFree(d_bias));\n cudaCheck(cudaFree(d_out));\n cudaCheck(cudaFree(d_dout));\n cudaCheck(cudaFree(d_dinp));\n cudaCheck(cudaFree(d_dweight));\n cudaCheck(cudaFree(d_dbias));\n\n cublasCheck(cublasDestroy(cublas_handle));\n return 0;\n}\n#endif", "output": "#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \"linear.cuh\"\n\n// ----------------------------------------------------------------------------\n// GPU kernels\n\n// copy the matmul forward and backward kernels from llm.c\n\n// is there no better way other than just adding bias with a whole separate kernel?\n// this is a highly memory-bound operation, should be fused into the matmul kernel\n// but i can't seem to find a cuBLAS function that does this\n__global__ void add_bias(float* out, const float* bias, int N, int OC) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = idx; i < N * OC; i += stride) {\n int col = i % OC;\n out[i] += bias[col];\n }\n}\n\n// kernel 2 calls cuBLAS, which should be very efficient\nvoid matmul_forward2(hipblasHandle_t cublas_handle,\n float* out,\n const float* inp, const float* weight, const float* bias,\n int N, int C, int OC,\n const int block_size) {\n // for reference API is:\n // cublasStatus_t cublasSgemm(cublasHandle_t handle,\n // cublasOperation_t transa, cublasOperation_t transb,\n // int m, int n, int k,\n // const float *alpha,\n // const float *A, int lda,\n // const float *B, int ldb,\n // const float *beta,\n // float *C, int ldc)\n // for us, inp is (B*T, C), weight is (OC, C), out is (B*T, OC)\n // cuBLAS does C = alpha * A * B + beta * C\n // where A is mxk, B is kxn, C is mxn\n // now, because we use row-major storage, cuBLAS (which is column-major) sees our matrices transposed.\n // algorithmically / in e.g. PyTorch we want to do: out = inp @ weight.T\n // but because cuBLAS is column-major, we actually want to get it to calculate out.T . Mathematically, this is:\n // out.T = weight @ inp.T\n // but again, our variables look transposed, so using the actual weight/inp we have here in this function, this becomes\n // out.T = weight.T @ inp\n // so we need to get cuBLAS to calculate weight.T @ inp (the variables here are the actual ones in this function)\n // => need to call cuBLAS with A = weight, B = inp\n // => need to call cuBLAS with transa = CUBLAS_OP_T, transb = CUBLAS_OP_N\n\n const float alpha = 1.0f;\n const float beta = 0.0f;\n cublasCheck(hipblasSgemm(cublas_handle, HIPBLAS_OP_T, HIPBLAS_OP_N, OC, N, C, &alpha, weight, C, inp, C, &beta, out, OC));\n // and now we still have to add the bias... (ew)\n if (bias != NULL) {\n int grid_size = ceil_div(OC * N, block_size);\n add_bias<<>>(out, bias, N, OC);\n cudaCheck(hipGetLastError());\n }\n}\n\n// use shared memory and coarsening + reductions\n__global__ void matmul_backward_bias_kernel_faster(float* dbias, const float* dout, int N, int OC) {\n extern __shared__ float shared[];\n int o = blockIdx.x; // range [0, OC)\n int tid = threadIdx.x; // range [0, block_size)\n int block_size = blockDim.x;\n const float* x = dout + o;\n // thread coarsening\n double sum = 0.0;\n for (int i = tid; i < N; i += block_size) {\n sum += x[i * OC];\n }\n shared[tid] = (float) sum;\n __syncthreads();\n // reductions\n for (int stride = block_size / 2; stride >= 1; stride /= 2) {\n __syncthreads();\n if (tid < stride) {\n shared[tid] += shared[tid + stride];\n }\n }\n // write the final result (at thread 0) to global memory\n if (tid == 0) {\n dbias[o] = shared[0];\n }\n}\n\n// version1: simple cuBLAS calls\nvoid matmul_backward1(hipblasHandle_t cublas_handle,\n float* dinp, float* dweight, float* dbias,\n float* dout, float* inp, float* weight,\n int N, int C, int OC) {\n float alpha = 1.0f;\n float beta = 0.0f; // note we must use beta = 1.0 so that we do a +=, as we should, because gradients add\n\n // for reference the API is:\n // cublasStatus_t cublasSgemm(cublasHandle_t handle,\n // cublasOperation_t transa, cublasOperation_t transb,\n // int m, int n, int k,\n // const float *alpha,\n // const float *A, int lda,\n // const float *B, int ldb,\n // const float *beta,\n // float *C, int ldc)\n\n // recall the forward pass was calculated with alpha = 1.0f, beta = 0.0f as:\n // cublasSgemm(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, OC, B*T, C, &alpha, weight, C, inp, C, &beta, out, OC);\n\n // backward to input\n cublasCheck(hipblasSgemm(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_N, C, N, OC, &alpha, weight, C, dout, OC, &beta, dinp, C));\n // backward to weight\n cublasCheck(hipblasSgemm(cublas_handle, HIPBLAS_OP_N, HIPBLAS_OP_T, C, OC, N, &alpha, inp, C, dout, OC, &beta, dweight, C));\n // backward to bias, if given\n if (dbias != NULL) {\n\n // sum over B,T using matrix vector multiplication with cuBLAS\n // for reference this API is:\n // cublasStatus_t cublasSgemv(cublasHandle_t handle, cublasOperation_t trans,\n // int m, int n,\n // const float *alpha,\n // const float *A, int lda,\n // const float *x, int incx,\n // const float *beta,\n // float *y, int incy)\n // dout is (B,T,OC), or in 2D terms (B*T, OC)\n // cublasCheck(cublasSgemv(cublas_handle, CUBLAS_OP_N, B*T, OC, &alpha, dout, B*T, ones, 1, &beta, dbias, 1));\n // cublasCheck(cublasSgemv(cublas_handle, CUBLAS_OP_T, OC, B*T, &alpha, dout, OC, ones, 1, &beta, dbias, 1));\n\n // ugh the above isn't working...\n // let's just do naive calculation for now, fix later\n // const int block_size=128;\n // const int grid_size=(OC + block_size - 1) / block_size;\n // matmul_backward_bias_kernel<<>>(dbias, dout, B, T, OC);\n\n // bit faster\n const int block_size=512;\n dim3 block_dim(block_size);\n dim3 grid_dim(OC);\n size_t shared_mem_size = block_size * sizeof(float);\n matmul_backward_bias_kernel_faster<<>>(dbias, dout, N, OC);\n }\n}\n\n// ----------------------------------------------------------------------------\n\nvoid linear_set_param_ptrs(\n LinearParams* params,\n float* params_memory,\n int C, int OC\n) {\n params->w = params_memory;\n params_memory += OC * C;\n params->b = params_memory;\n}\n\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n int N = 32;\n int C = 64;\n int OC = 128;\n\n // setup cublas\n hipblasHandle_t cublas_handle;\n cublasCheck(hipblasCreate(&cublas_handle));\n cublasCheck(hipblasSetMathMode(cublas_handle, HIPBLAS_DEFAULT_MATH));\n\n\n // create host memory\n float *inp = (float*)malloc(N * C * sizeof(float));\n float *weight = (float*)malloc(OC * C * sizeof(float));\n float *bias = (float*)malloc(OC * sizeof(float));\n float *out = (float*)malloc(N * OC * sizeof(float));\n float *dout = (float*)malloc(N * OC * sizeof(float));\n float *dinp = (float*)malloc(N * C * sizeof(float));\n float *dweight = (float*)malloc(OC * C * sizeof(float));\n float *dbias = (float*)malloc(OC * sizeof(float));\n\n FILE *file = fopen(\"linear.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(inp, sizeof(float), N * C, file);\n freadCheck(weight, sizeof(float), OC * C, file);\n freadCheck(bias, sizeof(float), OC, file);\n freadCheck(out, sizeof(float), N * OC, file);\n freadCheck(dout, sizeof(float), N * OC, file);\n freadCheck(dinp, sizeof(float), N * C, file);\n freadCheck(dweight, sizeof(float), OC * C, file);\n freadCheck(dbias, sizeof(float), OC, file);\n fclose(file);\n\n\n // allocate device memory\n float *d_inp, *d_weight, *d_bias, *d_out, *d_dout, *d_dinp, *d_dweight, *d_dbias;\n cudaCheck(hipMalloc(&d_inp, N * C * sizeof(float)));\n cudaCheck(hipMalloc(&d_weight, OC * C * sizeof(float)));\n cudaCheck(hipMalloc(&d_bias, OC * sizeof(float)));\n cudaCheck(hipMalloc(&d_out, N * OC * sizeof(float)));\n cudaCheck(hipMalloc(&d_dout, N * OC * sizeof(float)));\n cudaCheck(hipMalloc(&d_dinp, N * C * sizeof(float)));\n cudaCheck(hipMalloc(&d_dweight, OC * C * sizeof(float)));\n cudaCheck(hipMalloc(&d_dbias, OC * sizeof(float)));\n\n cudaCheck(hipMemcpy(d_inp, inp, N * C * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_weight, weight, OC * C * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_bias, bias, OC * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_dout, dout, N * OC * sizeof(float), hipMemcpyHostToDevice));\n\n int block_sizes[] = {128, 256, 512, 1024};\n\n printf(\"Checking forward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"\\nBlock size: %d\\n\", block_size);\n matmul_forward2(cublas_handle, d_out, d_inp, d_weight, d_bias, N, C, OC, block_size);\n validate_result(d_out, out, \"out\", N * OC);\n }\n\n printf(\"Forward pass successful\\n\");\n\n printf(\"Checking backward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"\\nBlock size: %d\\n\", block_size);\n matmul_backward1(cublas_handle, d_dinp, d_dweight, d_dbias, d_dout, d_inp, d_weight, N, C, OC);\n printf(\"Checking dbias\\n\");\n validate_result(d_dbias, dbias, \"dbias\", OC);\n printf(\"Checking dinp\\n\");\n validate_result(d_dinp, dinp, \"dinp\", N * C);\n printf(\"Checking dweight\\n\");\n validate_result(d_dweight, dweight, \"dweight\", OC * C);\n }\n printf(\"Backward pass successful\\n\");\n\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, matmul_forward2,\n cublas_handle, d_out, d_inp, d_weight, d_bias,\n N, C, OC, block_size);\n\n float tflops = (float)N * C * OC * 2 / elapsed_time * 1e3f / 1e12f;\n printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n }\n\n printf(\"\\nBackward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, matmul_backward1,\n cublas_handle, d_dinp, d_dweight, d_dbias, d_dout, d_inp, d_weight,\n N, C, OC);\n\n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n\n // free memory\n free(inp);\n free(weight);\n free(bias);\n free(out);\n free(dout);\n free(dinp);\n free(dweight);\n free(dbias);\n\n cudaCheck(hipFree(d_inp));\n cudaCheck(hipFree(d_weight));\n cudaCheck(hipFree(d_bias));\n cudaCheck(hipFree(d_out));\n cudaCheck(hipFree(d_dout));\n cudaCheck(hipFree(d_dinp));\n cudaCheck(hipFree(d_dweight));\n cudaCheck(hipFree(d_dbias));\n\n cublasCheck(hipblasDestroy(cublas_handle));\n return 0;\n}\n#endif", "source": "/kaggle/working/unet.cu/dev/linear.cu"} {"input": "#include \n#include \n#include \n#include \"common.h\"\n#include \"concat_channel.cuh\"\n\n\n__global__ void concat_channel_forward_kernel(\n const float* x1, const float* x2,\n float* out,\n int B, int C1, int C2, int H, int W\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < B * C1 * H * W) {\n // copy input from x1\n int b = idx / (C1 * H * W);\n int c = (idx / H / W) % C1;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + c * H * W + h * W + w;\n out[out_idx] = x1[idx];\n }\n if (idx < B * C2 * H * W) {\n // copy input from x2\n // move over from x1\n int b = idx / (C2 * H * W);\n int c = (idx / H / W) % C2;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + (C1 + c) * H * W + h * W + w;\n \n out[out_idx] = x2[idx];\n }\n}\n\n__global__ void concat_channel_backward_kernel(\n const float* dout,\n float* dx1, float* dx2,\n int B, int C1, int C2, int H, int W\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < B * C1 * H * W) {\n int b = idx / (C1 * H * W);\n int c = (idx / H / W) % C1;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + c * H * W + h * W + w;\n\n dx1[idx] = dout[out_idx];\n }\n if (idx < B * C2 * H * W) {\n int b = idx / (C2 * H * W);\n int c = (idx / H / W) % C2;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + (C1 + c) * H * W + h * W + w;\n \n dx2[idx] = dout[out_idx];\n }\n}\n\nvoid concat_channel_forward(\n const float* x1, const float* x2,\n float* out,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n int N = B * max_int(C1, C2) * H * W;\n int n_blk = ceil_div(N, block_size);\n concat_channel_forward_kernel<<>>(x1, x2, out, B, C1, C2, H, W);\n}\n\nvoid concat_channel_backward(\n const float* dout,\n float* dx1, float* dx2,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n int N = B * max_int(C1, C2) * H * W;\n int n_blk = ceil_div(N, block_size);\n concat_channel_backward_kernel<<>>(dout, dx1, dx2, B, C1, C2, H, W);\n}", "output": "#include \n#include \n#include \n#include \"common.h\"\n#include \"concat_channel.cuh\"\n\n\n__global__ void concat_channel_forward_kernel(\n const float* x1, const float* x2,\n float* out,\n int B, int C1, int C2, int H, int W\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < B * C1 * H * W) {\n // copy input from x1\n int b = idx / (C1 * H * W);\n int c = (idx / H / W) % C1;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + c * H * W + h * W + w;\n out[out_idx] = x1[idx];\n }\n if (idx < B * C2 * H * W) {\n // copy input from x2\n // move over from x1\n int b = idx / (C2 * H * W);\n int c = (idx / H / W) % C2;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + (C1 + c) * H * W + h * W + w;\n \n out[out_idx] = x2[idx];\n }\n}\n\n__global__ void concat_channel_backward_kernel(\n const float* dout,\n float* dx1, float* dx2,\n int B, int C1, int C2, int H, int W\n) {\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n if (idx < B * C1 * H * W) {\n int b = idx / (C1 * H * W);\n int c = (idx / H / W) % C1;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + c * H * W + h * W + w;\n\n dx1[idx] = dout[out_idx];\n }\n if (idx < B * C2 * H * W) {\n int b = idx / (C2 * H * W);\n int c = (idx / H / W) % C2;\n int h = (idx / W) % H;\n int w = idx % W;\n int out_idx = b * (C1 + C2) * H * W + (C1 + c) * H * W + h * W + w;\n \n dx2[idx] = dout[out_idx];\n }\n}\n\nvoid concat_channel_forward(\n const float* x1, const float* x2,\n float* out,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n int N = B * max_int(C1, C2) * H * W;\n int n_blk = ceil_div(N, block_size);\n concat_channel_forward_kernel<<>>(x1, x2, out, B, C1, C2, H, W);\n}\n\nvoid concat_channel_backward(\n const float* dout,\n float* dx1, float* dx2,\n int B, int C1, int C2, int H, int W, int block_size\n) {\n int N = B * max_int(C1, C2) * H * W;\n int n_blk = ceil_div(N, block_size);\n concat_channel_backward_kernel<<>>(dout, dx1, dx2, B, C1, C2, H, W);\n}", "source": "/kaggle/working/unet.cu/dev/concat_channel.cu"} {"input": "#include \n#include \n#include \n#include \"common.h\"\n\n\n// this will only be called with one block\n__global__ void mse_forward_kernel(\n const float* inp, const float* y, float* loss, int N\n) {\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n\n __shared__ float shared_sum[32]; // max 32 warps\n int num_warps = blockDim.x / 32;\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n\n float thread_sum = 0.0f;\n for (int i = threadIdx.x; i < N; i += blockDim.x) {\n float diff = inp[i] - y[i];\n thread_sum += diff * diff;\n }\n\n // warp reduce\n float warp_sum = cg::reduce(warp, thread_sum, cg::plus{});\n shared_sum[warp_id] = warp_sum;\n __syncthreads();\n \n warp_sum = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;\n float block_sum = cg::reduce(warp, warp_sum, cg::plus{});\n\n if (threadIdx.x == 0) {\n loss[0] = block_sum / N;\n }\n}\n\nvoid mse_forward(const float* inp, const float* y, float* loss, int N, int block_size) {\n mse_forward_kernel<<<1, block_size>>>(inp, y, loss, N);\n}\n\n\n\n// not calculating dy in backward\n// note that backward for MSE doesn't require the loss value\n__global__ void mse_backward_kernel(\n const float* inp, const float* y, float* dinp, int N\n) {\n int idx = threadIdx.x + blockIdx.x * blockDim.x;\n if (idx < N) {\n float x = inp[idx];\n dinp[idx] = 2.0f * (x - y[idx]) / N;\n }\n}\n\nvoid mse_backward(const float* inp, const float* y, float* dinp, int N, int block_size) {\n int n_blocks = ceil_div(N, block_size);\n mse_backward_kernel<<>>(inp, y, dinp, N);\n}\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n int N = 32 * 3 * 64 * 64;\n int block_size = 512;\n\n // host memory\n float* x = (float*)malloc(N * sizeof(float));\n float* y = (float*)malloc(N * sizeof(float));\n float* dx = (float*)malloc(N * sizeof(float));\n float* loss = (float*)malloc(sizeof(float));\n\n FILE *file = fopenCheck(\"mse.bin\", \"rb\");\n if (!file) {\n return -1;\n }\n freadCheck(x, sizeof(float), N, file);\n freadCheck(y, sizeof(float), N, file);\n freadCheck(dx, sizeof(float), N, file);\n freadCheck(loss, sizeof(float), 1, file);\n fcloseCheck(file);\n\n // device memory\n float *d_x, *d_y, *d_dx, *d_loss;\n cudaCheck(cudaMalloc(&d_x, N * sizeof(float)));\n cudaCheck(cudaMalloc(&d_y, N * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dx, N * sizeof(float)));\n cudaCheck(cudaMalloc(&d_loss, sizeof(float)));\n\n cudaCheck(cudaMemcpy(d_x, x, N * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_y, y, N * sizeof(float), cudaMemcpyHostToDevice));\n\n printf(\"Checking forward pass\\n\");\n mse_forward(d_x, d_y, d_loss, N, block_size);\n validate_result(d_loss, loss, \"loss\", 1);\n printf(\"Forward pass successful\\n\");\n\n printf(\"Checking backward pass\\n\");\n mse_backward(d_x, d_y, d_dx, N, block_size);\n validate_result(d_dx, dx, \"dinp\", N);\n printf(\"Backward pass successful\\n\");\n}\n#endif", "output": "#include \n#include \n#include \n#include \"common.h\"\n\n\n// this will only be called with one block\n__global__ void mse_forward_kernel(\n const float* inp, const float* y, float* loss, int N\n) {\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n\n __shared__ float shared_sum[32]; // max 32 warps\n int num_warps = blockDim.x / 32;\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n\n float thread_sum = 0.0f;\n for (int i = threadIdx.x; i < N; i += blockDim.x) {\n float diff = inp[i] - y[i];\n thread_sum += diff * diff;\n }\n\n // warp reduce\n float warp_sum = cg::reduce(warp, thread_sum, cg::plus{});\n shared_sum[warp_id] = warp_sum;\n __syncthreads();\n \n warp_sum = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;\n float block_sum = cg::reduce(warp, warp_sum, cg::plus{});\n\n if (threadIdx.x == 0) {\n loss[0] = block_sum / N;\n }\n}\n\nvoid mse_forward(const float* inp, const float* y, float* loss, int N, int block_size) {\n mse_forward_kernel<<<1, block_size>>>(inp, y, loss, N);\n}\n\n\n\n// not calculating dy in backward\n// note that backward for MSE doesn't require the loss value\n__global__ void mse_backward_kernel(\n const float* inp, const float* y, float* dinp, int N\n) {\n int idx = threadIdx.x + blockIdx.x * blockDim.x;\n if (idx < N) {\n float x = inp[idx];\n dinp[idx] = 2.0f * (x - y[idx]) / N;\n }\n}\n\nvoid mse_backward(const float* inp, const float* y, float* dinp, int N, int block_size) {\n int n_blocks = ceil_div(N, block_size);\n mse_backward_kernel<<>>(inp, y, dinp, N);\n}\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n int N = 32 * 3 * 64 * 64;\n int block_size = 512;\n\n // host memory\n float* x = (float*)malloc(N * sizeof(float));\n float* y = (float*)malloc(N * sizeof(float));\n float* dx = (float*)malloc(N * sizeof(float));\n float* loss = (float*)malloc(sizeof(float));\n\n FILE *file = fopenCheck(\"mse.bin\", \"rb\");\n if (!file) {\n return -1;\n }\n freadCheck(x, sizeof(float), N, file);\n freadCheck(y, sizeof(float), N, file);\n freadCheck(dx, sizeof(float), N, file);\n freadCheck(loss, sizeof(float), 1, file);\n fcloseCheck(file);\n\n // device memory\n float *d_x, *d_y, *d_dx, *d_loss;\n cudaCheck(hipMalloc(&d_x, N * sizeof(float)));\n cudaCheck(hipMalloc(&d_y, N * sizeof(float)));\n cudaCheck(hipMalloc(&d_dx, N * sizeof(float)));\n cudaCheck(hipMalloc(&d_loss, sizeof(float)));\n\n cudaCheck(hipMemcpy(d_x, x, N * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_y, y, N * sizeof(float), hipMemcpyHostToDevice));\n\n printf(\"Checking forward pass\\n\");\n mse_forward(d_x, d_y, d_loss, N, block_size);\n validate_result(d_loss, loss, \"loss\", 1);\n printf(\"Forward pass successful\\n\");\n\n printf(\"Checking backward pass\\n\");\n mse_backward(d_x, d_y, d_dx, N, block_size);\n validate_result(d_dx, dx, \"dinp\", N);\n printf(\"Backward pass successful\\n\");\n}\n#endif", "source": "/kaggle/working/unet.cu/dev/mse.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \n#include \"groupnorm.cuh\"\n\n\n// -----------------------------------------------------------------------------------------------\n// GPU kernels\n\n// Essentially taken from llm.c's kernel 5\n// using kernel 5 because for images, each \"channel\" is effectively\n// H * W * group_size, which is quite large\n// we will have one block per group, which means there are B * C / group_size blocks\n__global__ void groupnorm_forward_kernel(\n const float* x, const float* weight, const float* bias,\n float* out, float* mean, float* rstd,\n int B, int C, int img_size, int group_size, int n_groups\n) {\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n __shared__ float shared_sum[32]; // block_size max is 1024 = 32 * 32 warps\n __shared__ float shared_sum2[32]; // warps will be writing into shared memeory after warp-reduce\n int num_warps = blockDim.x / 32;\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n int block_pixels = img_size * group_size;\n // group index\n int g = blockIdx.x % n_groups;\n\n // move pointers\n x += blockIdx.x * img_size * group_size;\n out += blockIdx.x * img_size * group_size;\n // each block will only every acces group_size channels\n weight += g * group_size;\n bias += g * group_size;\n\n float thread_sum = 0.0f;\n float thread_sum2 = 0.0f;\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n float val = x[i];\n thread_sum += val;\n thread_sum2 += val * val;\n }\n\n // warp reduce\n float warp_sum = cg::reduce(warp, thread_sum, cg::plus{});\n float warp_sum2 = cg:: reduce(warp, thread_sum2, cg::plus{});\n // store warp sum into shared memory\n shared_sum[warp_id] = warp_sum;\n shared_sum2[warp_id] = warp_sum2;\n __syncthreads();\n \n // load warp sums from shared memory\n warp_sum = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;\n warp_sum2 = (lane_id < num_warps) ? shared_sum2[lane_id] : 0.0f;\n float block_sum = cg::reduce(warp, warp_sum, cg::plus{});\n float block_sum2 = cg::reduce(warp, warp_sum2, cg::plus{});\n block_sum /= block_pixels;\n block_sum2 /= block_pixels;\n float m = block_sum;\n float var = block_sum2 - m * m;\n float s = rsqrtf(var + 1e-5f);\n if (threadIdx.x == 0 && mean != nullptr) {\n mean[blockIdx.x] = m;\n }\n if (threadIdx.x == 0 && rstd != nullptr) {\n rstd[blockIdx.x] = s;\n }\n\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n int c_mod_group = (i / img_size) % group_size;\n float n = s * (x[i] - m);\n out[i] = n * weight[c_mod_group] + bias[c_mod_group];\n }\n}\n\n// most similar to kernel 2 in llm.c\n// not doing any float16 optimizations yet\n// main change is to have each block allocated to a single group\n__global__ void groupnorm_backward_kernel(\n const float* dout, const float* x, const float* mean, const float* rstd, const float* weight,\n float* dx, float* dweight, float* dbias,\n int B, int C, int img_size, int group_size, int n_groups\n) {\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n __shared__ float shared_sum[32]; // block_size max is 1024 = 32 * 32 warps\n __shared__ float shared_sum2[32]; // warps will be writing into shared memeory after warp-reduce\n int num_warps = blockDim.x / 32;\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n int block_pixels = img_size * group_size;\n // group index\n int g = blockIdx.x % n_groups;\n\n // move pointers\n dout += blockIdx.x * img_size * group_size;\n x += blockIdx.x * img_size * group_size;\n dx += blockIdx.x * img_size * group_size;\n weight += g * group_size;\n dweight += g * group_size;\n dbias += g * group_size;\n\n float m_val = mean[blockIdx.x];\n float rstd_val = rstd[blockIdx.x];\n\n\n // calculate the two mean terms in the group dimension\n // first is dout * weight, and second is dout * weight * norm\n // where norm = (x - mean) * rstd\n float w_dout_thread = 0.0f;\n float w_dout_norm_thread = 0.0f;\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n int c_mod_group = (i / img_size) % group_size;\n float cur_w_dout = weight[c_mod_group] * dout[i];\n w_dout_thread += cur_w_dout;\n float norm = (x[i] - m_val) * rstd_val;\n w_dout_norm_thread += cur_w_dout * norm;\n }\n // warp reduce\n float w_dout_warp = cg::reduce(warp, w_dout_thread, cg::plus{});\n float w_dout_norm_warp = cg::reduce(warp, w_dout_norm_thread, cg::plus{});\n // store warp sum in shared mem\n shared_sum[warp_id] = w_dout_warp;\n shared_sum2[warp_id] = w_dout_norm_warp;\n __syncthreads();\n\n // load warp sums from shared memory\n w_dout_warp = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;\n w_dout_norm_warp = (lane_id < num_warps) ? shared_sum2[lane_id] : 0.0f;\n float w_dout_block = cg::reduce(warp, w_dout_warp, cg::plus{});\n float w_dout_norm_block = cg::reduce(warp, w_dout_norm_warp, cg::plus{});\n w_dout_block /= block_pixels;\n w_dout_norm_block /= block_pixels;\n\n // update dx\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n // in bounds of image\n // accumulate dw and db\n float dout_val = dout[i];\n float norm = (x[i] - m_val) * rstd_val;\n\n // update dx\n int c_mod_group = (i / img_size) % group_size;\n float w_dout = weight[c_mod_group] * dout_val;\n dx[i] = (w_dout - w_dout_block - norm * w_dout_norm_block) * rstd_val;\n }\n // update dw and db\n // use different methods when the image size is large or small\n\n // if the image size is larger than the block size\n // loop over the channels and use the whole block on each channel\n // otherwise, assign each warp to a channel\n // in either case image size must be larger than the warp size\n //assert(img_size % warp.size() == 0);\n assert(blockDim.x % warp.size() == 0);\n if (img_size % blockDim.x == 0) {\n for (int c = 0; c < group_size; c++) {\n float dw_thread = 0.0f;\n float db_thread = 0.0f;\n for (int i = threadIdx.x; i < img_size; i += blockDim.x) {\n float dout_val = dout[i];\n db_thread += dout_val;\n float norm = (x[i] - m_val) * rstd_val;\n dw_thread += dout_val * norm;\n }\n\n // move pointers\n dout += img_size;\n x += img_size;\n\n // warp reduce\n float dw_warp = cg::reduce(warp, dw_thread, cg::plus{});\n float db_warp = cg::reduce(warp, db_thread, cg::plus{});\n ////// store warp sum in shared mem\n if (lane_id == 0) {\n shared_sum[warp_id] = dw_warp;\n shared_sum2[warp_id] = db_warp;\n }\n __syncthreads();\n // use the first thread to reduce the shared memory sums and save to global memory\n if (threadIdx.x == 0) {\n float dw_block = 0.0f;\n float db_block = 0.0f;\n for (int i = 0; i < num_warps; i++) {\n dw_block += shared_sum[i];\n db_block += shared_sum2[i];\n }\n atomicAdd(dweight + c, dw_block);\n atomicAdd(dbias + c, db_block);\n }\n }\n } else {\n // if group size is large, need to loop over the group channels with the whole block\n int block_reps = ceil_div(group_size, num_warps);\n for (int br = 0; br < block_reps; br++) {\n float dw_thread = 0.0f;\n float db_thread = 0.0f;\n\n int ch = br * num_warps + warp_id;\n if (ch < group_size) {\n const float* dout_ch = dout + ch * img_size;\n const float* x_ch = x + ch * img_size;\n for (int i = lane_id; i < img_size; i += warp.size()) {\n float dout_val = dout_ch[i];\n db_thread += dout_val;\n float norm = (x_ch[i] - m_val) * rstd_val;\n dw_thread += dout_val * norm;\n }\n \n // warp reduce\n float dw_warp = cg::reduce(warp, dw_thread, cg::plus{});\n float db_warp = cg::reduce(warp, db_thread, cg::plus{});\n // since each warp takes care of an entire image\n // directly store result\n if (lane_id == 0) {\n atomicAdd(dweight + ch, dw_warp);\n atomicAdd(dbias + ch, db_warp);\n }\n }\n }\n }\n}\n\n// -----------------------------------------------------------------------------------------------\n// CUDA kernel launcher\n\nvoid groupnorm_forward(\n const float* x, const float* weight, const float* bias,\n float* out, float* mean, float* rstd,\n int B, int C, int H, int W, int n_groups\n) {\n int img_size = H * W;\n int group_size = C / n_groups;\n int n_blocks = B * n_groups;\n int block_size = max_int(min_int(512, img_size * group_size), 32);\n groupnorm_forward_kernel<<>>(\n x, weight, bias, out, mean, rstd, B, C, img_size, group_size, n_groups\n );\n cudaCheck(cudaGetLastError());\n}\n\nvoid groupnorm_backward(\n const float* dout, const float* x, const float* mean, const float* rstd, const float* weight,\n float* dx, float* dweight, float* dbias,\n int B, int C, int H, int W, int n_groups\n) {\n int img_size = H * W;\n int group_size = C / n_groups;\n int n_blocks = B * n_groups;\n int block_size = max_int(min_int(512, img_size * group_size), 32 * group_size);\n groupnorm_backward_kernel<<>>(\n dout, x, mean, rstd, weight, dx, dweight, dbias, B, C, img_size, group_size, n_groups\n );\n cudaCheck(cudaGetLastError());\n}\n\n\n// -----------------------------------------------------------------------------------------------\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n //int B = 32;\n //int C = 64;\n //int H = 32;\n //int W = 32;\n srand(0);\n int B = 16;\n int C = 128;\n int H = 4;\n int W = 8;\n\n int img_size = H * W;\n int n_groups = 32;\n\n\n // create host memory to load data\n float *x = (float*)malloc(B * C * img_size * sizeof(float));\n float *weight = (float*)malloc(C * sizeof(float));\n float *bias = (float*)malloc(C * sizeof(float));\n float *out = (float*)malloc(B * C * img_size * sizeof(float));\n float *dout = (float*)malloc(B * C * img_size * sizeof(float));\n float *dx = (float*)malloc(B * C * img_size * sizeof(float));\n float *dweight = (float*)malloc(C * sizeof(float));\n float *dbias = (float*)malloc(C * sizeof(float));\n\n\n // read saved output\n FILE *file = fopenCheck(\"groupnorm.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(x, sizeof(float), B * C * img_size, file);\n freadCheck(weight, sizeof(float), C, file);\n freadCheck(bias, sizeof(float), C, file);\n freadCheck(out, sizeof(float), B * C * img_size, file);\n freadCheck(dout, sizeof(float), B * C * img_size, file);\n freadCheck(dx, sizeof(float), B * C * img_size, file);\n freadCheck(dweight, sizeof(float), C, file);\n freadCheck(dbias, sizeof(float), C, file);\n fcloseCheck(file);\n\n\n // device memory\n float *d_x, *d_weight, *d_bias, *d_out, *d_mean, *d_rstd, *d_dout, *d_dx, *d_dweight, *d_dbias;\n cudaCheck(cudaMalloc(&d_x, B * C * img_size * sizeof(float)));\n cudaCheck(cudaMalloc(&d_weight, C * sizeof(float)));\n cudaCheck(cudaMalloc(&d_bias, C * sizeof(float)));\n cudaCheck(cudaMalloc(&d_out, B * C * img_size * sizeof(float)));\n cudaCheck(cudaMalloc(&d_mean, B * n_groups * sizeof(float)));\n cudaCheck(cudaMalloc(&d_rstd, B * n_groups * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dout, B * C * img_size * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dx, B * C * img_size * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dweight, C * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dbias, C * sizeof(float)));\n cudaCheck(cudaMemcpy(d_x, x, B * C * img_size * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_weight, weight, C * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_bias, bias, C * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_dout, dout, B * C * img_size * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemset(d_dx, 0, B * C * img_size * sizeof(float)));\n cudaCheck(cudaMemset(d_dweight, 0, C * sizeof(float)));\n cudaCheck(cudaMemset(d_dbias, 0, C * sizeof(float)));\n\n printf(\"Checking forward pass\\n\");\n float fwd_acc = 1e-1;\n groupnorm_forward(d_x, d_weight, d_bias, d_out, d_mean, d_rstd, B, C, H, W, n_groups);\n validate_result(d_out, out, \"out\", B * C * img_size, fwd_acc);\n printf(\"Forward pass successful\\n\");\n\n // check the values of rstd\n float *rstd = (float*)malloc(B * n_groups * sizeof(float));\n\n printf(\"Checking backward pass\\n\");\n groupnorm_backward(d_dout, d_x, d_mean, d_rstd, d_weight, d_dx, d_dweight, d_dbias, B, C, H, W, n_groups);\n float acc = 1e-3;\n printf(\"Checking dbias\\n\");\n validate_result(d_dbias, dbias, \"dbias\", C, acc);\n printf(\"Checking dweight\\n\");\n validate_result(d_dweight, dweight, \"dweight\", C, acc);\n printf(\"Checking dx\\n\");\n validate_result(d_dx, dx, \"dx\", B * C * img_size, 1.0f);\n printf(\"Backward pass successful\\n\");\n\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmarks\\n\");\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, groupnorm_forward,\n d_x, d_weight, d_bias, d_out, d_mean, d_rstd,\n B, C, H, W, n_groups);\n printf(\"time %.4f ms\\n\", elapsed_time);\n\n printf(\"Backward pass benchmarks\\n\");\n elapsed_time = benchmark_kernel(repeat_times, groupnorm_backward,\n d_dout, d_x, d_mean, d_rstd, d_weight, d_dx, d_dweight, d_dbias,\n B, C, H, W, n_groups);\n printf(\"time %.4f ms\\n\", elapsed_time);\n}\n#endif", "output": "#include \n#include \n#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \n#include \"groupnorm.cuh\"\n\n\n// -----------------------------------------------------------------------------------------------\n// GPU kernels\n\n// Essentially taken from llm.c's kernel 5\n// using kernel 5 because for images, each \"channel\" is effectively\n// H * W * group_size, which is quite large\n// we will have one block per group, which means there are B * C / group_size blocks\n__global__ void groupnorm_forward_kernel(\n const float* x, const float* weight, const float* bias,\n float* out, float* mean, float* rstd,\n int B, int C, int img_size, int group_size, int n_groups\n) {\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n __shared__ float shared_sum[32]; // block_size max is 1024 = 32 * 32 warps\n __shared__ float shared_sum2[32]; // warps will be writing into shared memeory after warp-reduce\n int num_warps = blockDim.x / 32;\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n int block_pixels = img_size * group_size;\n // group index\n int g = blockIdx.x % n_groups;\n\n // move pointers\n x += blockIdx.x * img_size * group_size;\n out += blockIdx.x * img_size * group_size;\n // each block will only every acces group_size channels\n weight += g * group_size;\n bias += g * group_size;\n\n float thread_sum = 0.0f;\n float thread_sum2 = 0.0f;\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n float val = x[i];\n thread_sum += val;\n thread_sum2 += val * val;\n }\n\n // warp reduce\n float warp_sum = cg::reduce(warp, thread_sum, cg::plus{});\n float warp_sum2 = cg:: reduce(warp, thread_sum2, cg::plus{});\n // store warp sum into shared memory\n shared_sum[warp_id] = warp_sum;\n shared_sum2[warp_id] = warp_sum2;\n __syncthreads();\n \n // load warp sums from shared memory\n warp_sum = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;\n warp_sum2 = (lane_id < num_warps) ? shared_sum2[lane_id] : 0.0f;\n float block_sum = cg::reduce(warp, warp_sum, cg::plus{});\n float block_sum2 = cg::reduce(warp, warp_sum2, cg::plus{});\n block_sum /= block_pixels;\n block_sum2 /= block_pixels;\n float m = block_sum;\n float var = block_sum2 - m * m;\n float s = rsqrtf(var + 1e-5f);\n if (threadIdx.x == 0 && mean != nullptr) {\n mean[blockIdx.x] = m;\n }\n if (threadIdx.x == 0 && rstd != nullptr) {\n rstd[blockIdx.x] = s;\n }\n\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n int c_mod_group = (i / img_size) % group_size;\n float n = s * (x[i] - m);\n out[i] = n * weight[c_mod_group] + bias[c_mod_group];\n }\n}\n\n// most similar to kernel 2 in llm.c\n// not doing any float16 optimizations yet\n// main change is to have each block allocated to a single group\n__global__ void groupnorm_backward_kernel(\n const float* dout, const float* x, const float* mean, const float* rstd, const float* weight,\n float* dx, float* dweight, float* dbias,\n int B, int C, int img_size, int group_size, int n_groups\n) {\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n __shared__ float shared_sum[32]; // block_size max is 1024 = 32 * 32 warps\n __shared__ float shared_sum2[32]; // warps will be writing into shared memeory after warp-reduce\n int num_warps = blockDim.x / 32;\n int warp_id = threadIdx.x / 32;\n int lane_id = threadIdx.x % 32;\n int block_pixels = img_size * group_size;\n // group index\n int g = blockIdx.x % n_groups;\n\n // move pointers\n dout += blockIdx.x * img_size * group_size;\n x += blockIdx.x * img_size * group_size;\n dx += blockIdx.x * img_size * group_size;\n weight += g * group_size;\n dweight += g * group_size;\n dbias += g * group_size;\n\n float m_val = mean[blockIdx.x];\n float rstd_val = rstd[blockIdx.x];\n\n\n // calculate the two mean terms in the group dimension\n // first is dout * weight, and second is dout * weight * norm\n // where norm = (x - mean) * rstd\n float w_dout_thread = 0.0f;\n float w_dout_norm_thread = 0.0f;\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n int c_mod_group = (i / img_size) % group_size;\n float cur_w_dout = weight[c_mod_group] * dout[i];\n w_dout_thread += cur_w_dout;\n float norm = (x[i] - m_val) * rstd_val;\n w_dout_norm_thread += cur_w_dout * norm;\n }\n // warp reduce\n float w_dout_warp = cg::reduce(warp, w_dout_thread, cg::plus{});\n float w_dout_norm_warp = cg::reduce(warp, w_dout_norm_thread, cg::plus{});\n // store warp sum in shared mem\n shared_sum[warp_id] = w_dout_warp;\n shared_sum2[warp_id] = w_dout_norm_warp;\n __syncthreads();\n\n // load warp sums from shared memory\n w_dout_warp = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;\n w_dout_norm_warp = (lane_id < num_warps) ? shared_sum2[lane_id] : 0.0f;\n float w_dout_block = cg::reduce(warp, w_dout_warp, cg::plus{});\n float w_dout_norm_block = cg::reduce(warp, w_dout_norm_warp, cg::plus{});\n w_dout_block /= block_pixels;\n w_dout_norm_block /= block_pixels;\n\n // update dx\n for (int i = threadIdx.x; i < block_pixels; i += blockDim.x) {\n // in bounds of image\n // accumulate dw and db\n float dout_val = dout[i];\n float norm = (x[i] - m_val) * rstd_val;\n\n // update dx\n int c_mod_group = (i / img_size) % group_size;\n float w_dout = weight[c_mod_group] * dout_val;\n dx[i] = (w_dout - w_dout_block - norm * w_dout_norm_block) * rstd_val;\n }\n // update dw and db\n // use different methods when the image size is large or small\n\n // if the image size is larger than the block size\n // loop over the channels and use the whole block on each channel\n // otherwise, assign each warp to a channel\n // in either case image size must be larger than the warp size\n //assert(img_size % warp.size() == 0);\n assert(blockDim.x % warp.size() == 0);\n if (img_size % blockDim.x == 0) {\n for (int c = 0; c < group_size; c++) {\n float dw_thread = 0.0f;\n float db_thread = 0.0f;\n for (int i = threadIdx.x; i < img_size; i += blockDim.x) {\n float dout_val = dout[i];\n db_thread += dout_val;\n float norm = (x[i] - m_val) * rstd_val;\n dw_thread += dout_val * norm;\n }\n\n // move pointers\n dout += img_size;\n x += img_size;\n\n // warp reduce\n float dw_warp = cg::reduce(warp, dw_thread, cg::plus{});\n float db_warp = cg::reduce(warp, db_thread, cg::plus{});\n ////// store warp sum in shared mem\n if (lane_id == 0) {\n shared_sum[warp_id] = dw_warp;\n shared_sum2[warp_id] = db_warp;\n }\n __syncthreads();\n // use the first thread to reduce the shared memory sums and save to global memory\n if (threadIdx.x == 0) {\n float dw_block = 0.0f;\n float db_block = 0.0f;\n for (int i = 0; i < num_warps; i++) {\n dw_block += shared_sum[i];\n db_block += shared_sum2[i];\n }\n atomicAdd(dweight + c, dw_block);\n atomicAdd(dbias + c, db_block);\n }\n }\n } else {\n // if group size is large, need to loop over the group channels with the whole block\n int block_reps = ceil_div(group_size, num_warps);\n for (int br = 0; br < block_reps; br++) {\n float dw_thread = 0.0f;\n float db_thread = 0.0f;\n\n int ch = br * num_warps + warp_id;\n if (ch < group_size) {\n const float* dout_ch = dout + ch * img_size;\n const float* x_ch = x + ch * img_size;\n for (int i = lane_id; i < img_size; i += warp.size()) {\n float dout_val = dout_ch[i];\n db_thread += dout_val;\n float norm = (x_ch[i] - m_val) * rstd_val;\n dw_thread += dout_val * norm;\n }\n \n // warp reduce\n float dw_warp = cg::reduce(warp, dw_thread, cg::plus{});\n float db_warp = cg::reduce(warp, db_thread, cg::plus{});\n // since each warp takes care of an entire image\n // directly store result\n if (lane_id == 0) {\n atomicAdd(dweight + ch, dw_warp);\n atomicAdd(dbias + ch, db_warp);\n }\n }\n }\n }\n}\n\n// -----------------------------------------------------------------------------------------------\n// CUDA kernel launcher\n\nvoid groupnorm_forward(\n const float* x, const float* weight, const float* bias,\n float* out, float* mean, float* rstd,\n int B, int C, int H, int W, int n_groups\n) {\n int img_size = H * W;\n int group_size = C / n_groups;\n int n_blocks = B * n_groups;\n int block_size = max_int(min_int(512, img_size * group_size), 32);\n groupnorm_forward_kernel<<>>(\n x, weight, bias, out, mean, rstd, B, C, img_size, group_size, n_groups\n );\n cudaCheck(hipGetLastError());\n}\n\nvoid groupnorm_backward(\n const float* dout, const float* x, const float* mean, const float* rstd, const float* weight,\n float* dx, float* dweight, float* dbias,\n int B, int C, int H, int W, int n_groups\n) {\n int img_size = H * W;\n int group_size = C / n_groups;\n int n_blocks = B * n_groups;\n int block_size = max_int(min_int(512, img_size * group_size), 32 * group_size);\n groupnorm_backward_kernel<<>>(\n dout, x, mean, rstd, weight, dx, dweight, dbias, B, C, img_size, group_size, n_groups\n );\n cudaCheck(hipGetLastError());\n}\n\n\n// -----------------------------------------------------------------------------------------------\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n //int B = 32;\n //int C = 64;\n //int H = 32;\n //int W = 32;\n srand(0);\n int B = 16;\n int C = 128;\n int H = 4;\n int W = 8;\n\n int img_size = H * W;\n int n_groups = 32;\n\n\n // create host memory to load data\n float *x = (float*)malloc(B * C * img_size * sizeof(float));\n float *weight = (float*)malloc(C * sizeof(float));\n float *bias = (float*)malloc(C * sizeof(float));\n float *out = (float*)malloc(B * C * img_size * sizeof(float));\n float *dout = (float*)malloc(B * C * img_size * sizeof(float));\n float *dx = (float*)malloc(B * C * img_size * sizeof(float));\n float *dweight = (float*)malloc(C * sizeof(float));\n float *dbias = (float*)malloc(C * sizeof(float));\n\n\n // read saved output\n FILE *file = fopenCheck(\"groupnorm.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(x, sizeof(float), B * C * img_size, file);\n freadCheck(weight, sizeof(float), C, file);\n freadCheck(bias, sizeof(float), C, file);\n freadCheck(out, sizeof(float), B * C * img_size, file);\n freadCheck(dout, sizeof(float), B * C * img_size, file);\n freadCheck(dx, sizeof(float), B * C * img_size, file);\n freadCheck(dweight, sizeof(float), C, file);\n freadCheck(dbias, sizeof(float), C, file);\n fcloseCheck(file);\n\n\n // device memory\n float *d_x, *d_weight, *d_bias, *d_out, *d_mean, *d_rstd, *d_dout, *d_dx, *d_dweight, *d_dbias;\n cudaCheck(hipMalloc(&d_x, B * C * img_size * sizeof(float)));\n cudaCheck(hipMalloc(&d_weight, C * sizeof(float)));\n cudaCheck(hipMalloc(&d_bias, C * sizeof(float)));\n cudaCheck(hipMalloc(&d_out, B * C * img_size * sizeof(float)));\n cudaCheck(hipMalloc(&d_mean, B * n_groups * sizeof(float)));\n cudaCheck(hipMalloc(&d_rstd, B * n_groups * sizeof(float)));\n cudaCheck(hipMalloc(&d_dout, B * C * img_size * sizeof(float)));\n cudaCheck(hipMalloc(&d_dx, B * C * img_size * sizeof(float)));\n cudaCheck(hipMalloc(&d_dweight, C * sizeof(float)));\n cudaCheck(hipMalloc(&d_dbias, C * sizeof(float)));\n cudaCheck(hipMemcpy(d_x, x, B * C * img_size * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_weight, weight, C * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_bias, bias, C * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_dout, dout, B * C * img_size * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemset(d_dx, 0, B * C * img_size * sizeof(float)));\n cudaCheck(hipMemset(d_dweight, 0, C * sizeof(float)));\n cudaCheck(hipMemset(d_dbias, 0, C * sizeof(float)));\n\n printf(\"Checking forward pass\\n\");\n float fwd_acc = 1e-1;\n groupnorm_forward(d_x, d_weight, d_bias, d_out, d_mean, d_rstd, B, C, H, W, n_groups);\n validate_result(d_out, out, \"out\", B * C * img_size, fwd_acc);\n printf(\"Forward pass successful\\n\");\n\n // check the values of rstd\n float *rstd = (float*)malloc(B * n_groups * sizeof(float));\n\n printf(\"Checking backward pass\\n\");\n groupnorm_backward(d_dout, d_x, d_mean, d_rstd, d_weight, d_dx, d_dweight, d_dbias, B, C, H, W, n_groups);\n float acc = 1e-3;\n printf(\"Checking dbias\\n\");\n validate_result(d_dbias, dbias, \"dbias\", C, acc);\n printf(\"Checking dweight\\n\");\n validate_result(d_dweight, dweight, \"dweight\", C, acc);\n printf(\"Checking dx\\n\");\n validate_result(d_dx, dx, \"dx\", B * C * img_size, 1.0f);\n printf(\"Backward pass successful\\n\");\n\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmarks\\n\");\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, groupnorm_forward,\n d_x, d_weight, d_bias, d_out, d_mean, d_rstd,\n B, C, H, W, n_groups);\n printf(\"time %.4f ms\\n\", elapsed_time);\n\n printf(\"Backward pass benchmarks\\n\");\n elapsed_time = benchmark_kernel(repeat_times, groupnorm_backward,\n d_dout, d_x, d_mean, d_rstd, d_weight, d_dx, d_dweight, d_dbias,\n B, C, H, W, n_groups);\n printf(\"time %.4f ms\\n\", elapsed_time);\n}\n#endif", "source": "/kaggle/working/unet.cu/dev/groupnorm.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \"broadcast.cuh\"\n\n// kernel to broadcast a tensor from shape (B, C) to (B, C, H, W)\n// Assuming one thread per pixel\n__global__ void broadcast_last_dims_forward_kernel(\n const float* x,\n float* out,\n int N, int H, int W, int n_bk_per_img\n) {\n int n = blockIdx.x / n_bk_per_img;\n __shared__ float x_sh[1];\n if (threadIdx.x == 0) {\n x_sh[0] = x[n];\n }\n __syncthreads();\n\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < N * H * W) {\n out[idx] = x_sh[0];\n }\n}\n\nvoid broadcast_last_dims_forward(\n const float* x,\n float* out,\n int N, int H, int W, int block_size\n) {\n block_size = min_int(block_size, H * W);\n int n_bk_per_img = ceil_div(H * W, block_size);\n broadcast_last_dims_forward_kernel<<>>(x, out, N, H, W, n_bk_per_img);\n};\n\n// backwards for broadcast\n// use one warp per image\n__global__ void broadcast_last_dims_backward_kernel(\n const float* dout,\n float* dx,\n int N, int H, int W\n) {\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();\n if (idx >= N) { return; }\n\n int img_size = H * W;\n dout += idx * img_size;\n dx += idx;\n\n float sum = 0.0f;\n for (int i = warp.thread_rank(); i < img_size; i += warp.size()) {\n sum += dout[i];\n }\n sum = cg::reduce(warp, sum, cg::plus());\n if (warp.thread_rank() == 0) {\n dx[0] = sum;\n }\n}\n\nvoid broadcast_last_dims_backward(\n const float* dout,\n float* dx,\n int N, int H, int W, int block_size\n) {\n block_size = min_int(block_size, H * W);\n assert(block_size % 32 == 0);\n int warps_per_block = block_size / 32;\n int gridDim = ceil_div(N, warps_per_block);\n broadcast_last_dims_backward_kernel<<>>(dout, dx, N, H, W);\n};\n\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n int N = 8 * 64;\n int H = 16;\n int W = 16;\n\n // create host memory\n float *x = (float*)malloc(N * sizeof(float));\n float *out = (float*)malloc(N * H * W * sizeof(float));\n float *dout = (float*)malloc(N * H * W * sizeof(float));\n float *dx = (float*)malloc(N * sizeof(float));\n\n // read saved output\n FILE *file = fopen(\"broadcast.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(x, sizeof(float), N, file);\n freadCheck(out, sizeof(float), N * H * W, file);\n freadCheck(dout, sizeof(float), N * H * W, file);\n freadCheck(dx, sizeof(float), N, file);\n fclose(file);\n\n // allocate device memory\n float *d_x, *d_out, *d_dout, *d_dx;\n cudaCheck(cudaMalloc(&d_x, N * sizeof(float)));\n cudaCheck(cudaMalloc(&d_out, N * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dout, N * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dx, N * sizeof(float)));\n cudaCheck(cudaMemcpy(d_x, x, N * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_dout, dout, N * H * W * sizeof(float), cudaMemcpyHostToDevice));\n\n int block_sizes[] = {512};\n printf(\"Checking forward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"\\nChecking block size %d\\n\", block_size);\n broadcast_last_dims_forward(d_x, d_out, N, H, W, block_size);\n validate_result(d_out, out, \"out\", N * H * W);\n }\n printf(\"Checking backward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"\\nChecking block size %d\\n\", block_size);\n broadcast_last_dims_backward(d_dout, d_dx, N, H, W, block_size);\n validate_result(d_dx, dx, \"dx\", N);\n }\n\n\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, broadcast_last_dims_forward,\n d_out, d_x, N, H, W, block_size);\n \n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n printf(\"Backward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, broadcast_last_dims_backward,\n d_dout, d_dx, N, H, W, block_size);\n \n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n}\n#endif", "output": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"common.h\"\n#include \"broadcast.cuh\"\n\n// kernel to broadcast a tensor from shape (B, C) to (B, C, H, W)\n// Assuming one thread per pixel\n__global__ void broadcast_last_dims_forward_kernel(\n const float* x,\n float* out,\n int N, int H, int W, int n_bk_per_img\n) {\n int n = blockIdx.x / n_bk_per_img;\n __shared__ float x_sh[1];\n if (threadIdx.x == 0) {\n x_sh[0] = x[n];\n }\n __syncthreads();\n\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < N * H * W) {\n out[idx] = x_sh[0];\n }\n}\n\nvoid broadcast_last_dims_forward(\n const float* x,\n float* out,\n int N, int H, int W, int block_size\n) {\n block_size = min_int(block_size, H * W);\n int n_bk_per_img = ceil_div(H * W, block_size);\n broadcast_last_dims_forward_kernel<<>>(x, out, N, H, W, n_bk_per_img);\n};\n\n// backwards for broadcast\n// use one warp per image\n__global__ void broadcast_last_dims_backward_kernel(\n const float* dout,\n float* dx,\n int N, int H, int W\n) {\n namespace cg = cooperative_groups;\n cg::thread_block block = cg::this_thread_block();\n cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);\n int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();\n if (idx >= N) { return; }\n\n int img_size = H * W;\n dout += idx * img_size;\n dx += idx;\n\n float sum = 0.0f;\n for (int i = warp.thread_rank(); i < img_size; i += warp.size()) {\n sum += dout[i];\n }\n sum = cg::reduce(warp, sum, cg::plus());\n if (warp.thread_rank() == 0) {\n dx[0] = sum;\n }\n}\n\nvoid broadcast_last_dims_backward(\n const float* dout,\n float* dx,\n int N, int H, int W, int block_size\n) {\n block_size = min_int(block_size, H * W);\n assert(block_size % 32 == 0);\n int warps_per_block = block_size / 32;\n int gridDim = ceil_div(N, warps_per_block);\n broadcast_last_dims_backward_kernel<<>>(dout, dx, N, H, W);\n};\n\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n int N = 8 * 64;\n int H = 16;\n int W = 16;\n\n // create host memory\n float *x = (float*)malloc(N * sizeof(float));\n float *out = (float*)malloc(N * H * W * sizeof(float));\n float *dout = (float*)malloc(N * H * W * sizeof(float));\n float *dx = (float*)malloc(N * sizeof(float));\n\n // read saved output\n FILE *file = fopen(\"broadcast.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(x, sizeof(float), N, file);\n freadCheck(out, sizeof(float), N * H * W, file);\n freadCheck(dout, sizeof(float), N * H * W, file);\n freadCheck(dx, sizeof(float), N, file);\n fclose(file);\n\n // allocate device memory\n float *d_x, *d_out, *d_dout, *d_dx;\n cudaCheck(hipMalloc(&d_x, N * sizeof(float)));\n cudaCheck(hipMalloc(&d_out, N * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&d_dout, N * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&d_dx, N * sizeof(float)));\n cudaCheck(hipMemcpy(d_x, x, N * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_dout, dout, N * H * W * sizeof(float), hipMemcpyHostToDevice));\n\n int block_sizes[] = {512};\n printf(\"Checking forward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"\\nChecking block size %d\\n\", block_size);\n broadcast_last_dims_forward(d_x, d_out, N, H, W, block_size);\n validate_result(d_out, out, \"out\", N * H * W);\n }\n printf(\"Checking backward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"\\nChecking block size %d\\n\", block_size);\n broadcast_last_dims_backward(d_dout, d_dx, N, H, W, block_size);\n validate_result(d_dx, dx, \"dx\", N);\n }\n\n\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, broadcast_last_dims_forward,\n d_out, d_x, N, H, W, block_size);\n \n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n printf(\"Backward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, broadcast_last_dims_backward,\n d_dout, d_dx, N, H, W, block_size);\n \n printf(\"block_size %4d | time %.4f ms\\n\", block_size, elapsed_time);\n }\n}\n#endif", "source": "/kaggle/working/unet.cu/dev/broadcast.cu"} {"input": "#include \n#include \n#include \n#include \"common.h\"\n#include \"upsample.cuh\"\n\n// ----------------------------------------------------------------------------\n// GPU kernels\n\n// one thread per pixel\n__global__ void upsample_forward_kernel(\n const float* x,\n float* out,\n int B, int C, int H, int W\n) {\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int img_size = H * W;\n if (flat_idx >= B * C * img_size) {\n return;\n }\n\n int H_out = H*2;\n int W_out = W*2;\n int img_out_size = H_out * W_out;\n\n int b = flat_idx / (C * img_size);\n int c = (flat_idx / img_size) % C;\n int i = (flat_idx / W) % H;\n int j = flat_idx % W;\n\n // move pointers\n x += b * C * img_size + c * img_size + i * W + j;\n out += b * C * img_out_size + c * img_out_size + 2 * i * W_out + 2 * j;\n\n float x_val = x[0];\n\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n out[ii * W_out + jj] = x_val;\n }\n }\n}\n\n\n// number of threads equals the number of pixels\n// in downsampled image\n__global__ void upsample_backward_kernel(\n const float* dout,\n float* dx,\n int B, int C, int H, int W\n) {\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int img_in_size = H * W;\n if (flat_idx >= B * C * img_in_size) {\n return;\n }\n\n int H_out = H*2;\n int W_out = W*2;\n int img_out_size = H_out * W_out;\n\n int b = flat_idx / (C * img_in_size);\n int c = (flat_idx / img_in_size) % C;\n int i = (flat_idx / W) % H;\n int j = flat_idx % W;\n\n // move pointers\n dx += b * C * img_in_size + c * img_in_size + i * W + j;\n dout += b * C * img_out_size + c * img_out_size + 2 * i * W_out + 2 * j;\n\n float dout_sum = 0.0f;\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n dout_sum += dout[ii * W_out + jj];\n }\n }\n dx[0] = dout_sum;\n}\n\n// ----------------------------------------------------------------------------\n// kernel launcher\nvoid upsample_forward1(\n float* out,\n const float* x,\n int B, int C, int H, int W,\n int block_size\n) {\n int gridDim = ceil_div(B * C * H * W, block_size);\n upsample_forward_kernel<<>>(x, out, B, C, H, W);\n cudaCheck(cudaGetLastError());\n}\n\nvoid upsample_backward1(\n float* dx,\n const float* dout,\n int B, int C, int H, int W,\n int block_size\n) {\n int gridDim = ceil_div(B * C * H * W, block_size);\n upsample_backward_kernel<<>>(dout, dx, B, C, H, W);\n cudaCheck(cudaGetLastError());\n}\n\nvoid upsample_forward(\n int kernel_num,\n float* out,\n const float* x,\n int B, int C, int H, int W,\n int block_size\n) {\n switch (kernel_num) {\n case 1:\n upsample_forward1(out, x, B, C, H, W, block_size);\n break;\n default:\n printf(\"Invalid kernel number\\n\");\n exit(1);\n }\n}\n\nvoid upsample_backward(\n int kernel_num,\n float* dx,\n const float* dout,\n int B, int C, int H, int W,\n int block_size\n) {\n switch (kernel_num) {\n case 1:\n upsample_backward1(dx, dout, B, C, H, W, block_size);\n break;\n default:\n printf(\"Invalid kernel number\\n\");\n exit(1);\n }\n}\n// ----------------------------------------------------------------------------\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n srand(0);\n int B = 1;\n int C = 64;\n int H = 32;\n int W = 32;\n\n // set up the device\n int deviceIdx = 0;\n cudaCheck(cudaSetDevice(deviceIdx));\n cudaDeviceProp deviceProp;\n cudaGetDeviceProperties(&deviceProp, deviceIdx);\n printf(\"Device %d: %s\\n\", deviceIdx, deviceProp.name);\n\n // create host memory to load data\n float* x = (float*)malloc(B * C * H * W * sizeof(float));\n float* out = (float*)malloc(B * C * H*2 * W*2 * sizeof(float));\n float* dx = (float*)malloc(B * C * H * W * sizeof(float));\n float* dout = (float*)malloc(B * C * H*2 * W*2 * sizeof(float));\n\n // read saved output\n FILE *file = fopen(\"upsample.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(x, sizeof(float), B * C * H * W, file);\n freadCheck(out, sizeof(float), B * C * H*2 * W*2, file);\n freadCheck(dout, sizeof(float), B * C * H*2 * W*2, file);\n freadCheck(dx, sizeof(float), B * C * H * W, file);\n fclose(file);\n\n\n // allocate device memory\n float *d_x, *d_out, *d_dout, *d_dx;\n cudaCheck(cudaMalloc(&d_x, B * C * H * W * sizeof(float)));\n cudaCheck(cudaMalloc(&d_out, B * C * H*2 * W*2 * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dout, B * C * H*2 * W*2 * sizeof(float)));\n cudaCheck(cudaMalloc(&d_dx, B * C * H * W * sizeof(float)));\n cudaCheck(cudaMemcpy(d_x, x, B * C * H * W * sizeof(float), cudaMemcpyHostToDevice));\n cudaCheck(cudaMemcpy(d_dout, dout, B * C * H*2 * W*2 * sizeof(float), cudaMemcpyHostToDevice));\n\n // read kernel number from command line\n int kernel_num = 1;\n if (argc > 1) {\n kernel_num = atoi(argv[1]);\n }\n printf(\"Running kernel %d\\n\", kernel_num);\n\n int block_sizes[] = {32, 64, 128};\n printf(\"Checking forward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"Checking block size %d\\n\", block_size);\n upsample_forward(kernel_num, d_out, d_x, B, C, H, W, block_size);\n validate_result(d_out, out, \"out\", B * C * (H/2) * (W/2));\n }\n printf(\"Forward pass: all results match\\n\\n\");\n\n printf(\"Checking backward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"Checking block size %d\\n\", block_size);\n upsample_backward(kernel_num, d_dx, d_dout, B, C, H, W, block_size);\n validate_result(d_dx, dx, \"dx\", B * C * H * W);\n }\n printf(\"Backward pass: all results match\\n\\n\");\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, upsample_forward,\n kernel_num, d_out, d_x,\n B, C, H, W, block_size);\n\n float tflops = (float)B * H * W * C / elapsed_time * 1e3f / 1e12f;\n printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n }\n\n printf(\"\\nBackward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, upsample_backward,\n kernel_num, d_dx, d_dout,\n B, C, H, W, block_size);\n\n float tflops = (float)B * H * W * C / elapsed_time * 1e3f / 1e12f;\n printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n }\n\n // free memory\n free(x);\n free(out);\n free(dout);\n free(dx);\n\n cudaFree(d_x);\n cudaFree(d_out);\n cudaFree(d_dout);\n cudaFree(d_dx);\n\n return 0;\n}\n#endif", "output": "#include \n#include \n#include \n#include \"common.h\"\n#include \"upsample.cuh\"\n\n// ----------------------------------------------------------------------------\n// GPU kernels\n\n// one thread per pixel\n__global__ void upsample_forward_kernel(\n const float* x,\n float* out,\n int B, int C, int H, int W\n) {\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int img_size = H * W;\n if (flat_idx >= B * C * img_size) {\n return;\n }\n\n int H_out = H*2;\n int W_out = W*2;\n int img_out_size = H_out * W_out;\n\n int b = flat_idx / (C * img_size);\n int c = (flat_idx / img_size) % C;\n int i = (flat_idx / W) % H;\n int j = flat_idx % W;\n\n // move pointers\n x += b * C * img_size + c * img_size + i * W + j;\n out += b * C * img_out_size + c * img_out_size + 2 * i * W_out + 2 * j;\n\n float x_val = x[0];\n\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n out[ii * W_out + jj] = x_val;\n }\n }\n}\n\n\n// number of threads equals the number of pixels\n// in downsampled image\n__global__ void upsample_backward_kernel(\n const float* dout,\n float* dx,\n int B, int C, int H, int W\n) {\n int flat_idx = blockDim.x * blockIdx.x + threadIdx.x;\n int img_in_size = H * W;\n if (flat_idx >= B * C * img_in_size) {\n return;\n }\n\n int H_out = H*2;\n int W_out = W*2;\n int img_out_size = H_out * W_out;\n\n int b = flat_idx / (C * img_in_size);\n int c = (flat_idx / img_in_size) % C;\n int i = (flat_idx / W) % H;\n int j = flat_idx % W;\n\n // move pointers\n dx += b * C * img_in_size + c * img_in_size + i * W + j;\n dout += b * C * img_out_size + c * img_out_size + 2 * i * W_out + 2 * j;\n\n float dout_sum = 0.0f;\n for (int ii = 0; ii < 2; ii++) {\n for (int jj = 0; jj < 2; jj++) {\n dout_sum += dout[ii * W_out + jj];\n }\n }\n dx[0] = dout_sum;\n}\n\n// ----------------------------------------------------------------------------\n// kernel launcher\nvoid upsample_forward1(\n float* out,\n const float* x,\n int B, int C, int H, int W,\n int block_size\n) {\n int gridDim = ceil_div(B * C * H * W, block_size);\n upsample_forward_kernel<<>>(x, out, B, C, H, W);\n cudaCheck(hipGetLastError());\n}\n\nvoid upsample_backward1(\n float* dx,\n const float* dout,\n int B, int C, int H, int W,\n int block_size\n) {\n int gridDim = ceil_div(B * C * H * W, block_size);\n upsample_backward_kernel<<>>(dout, dx, B, C, H, W);\n cudaCheck(hipGetLastError());\n}\n\nvoid upsample_forward(\n int kernel_num,\n float* out,\n const float* x,\n int B, int C, int H, int W,\n int block_size\n) {\n switch (kernel_num) {\n case 1:\n upsample_forward1(out, x, B, C, H, W, block_size);\n break;\n default:\n printf(\"Invalid kernel number\\n\");\n exit(1);\n }\n}\n\nvoid upsample_backward(\n int kernel_num,\n float* dx,\n const float* dout,\n int B, int C, int H, int W,\n int block_size\n) {\n switch (kernel_num) {\n case 1:\n upsample_backward1(dx, dout, B, C, H, W, block_size);\n break;\n default:\n printf(\"Invalid kernel number\\n\");\n exit(1);\n }\n}\n// ----------------------------------------------------------------------------\n\n#ifndef LINKING\nint main(int argc, char **argv) {\n srand(0);\n int B = 1;\n int C = 64;\n int H = 32;\n int W = 32;\n\n // set up the device\n int deviceIdx = 0;\n cudaCheck(hipSetDevice(deviceIdx));\n hipDeviceProp_t deviceProp;\n hipGetDeviceProperties(&deviceProp, deviceIdx);\n printf(\"Device %d: %s\\n\", deviceIdx, deviceProp.name);\n\n // create host memory to load data\n float* x = (float*)malloc(B * C * H * W * sizeof(float));\n float* out = (float*)malloc(B * C * H*2 * W*2 * sizeof(float));\n float* dx = (float*)malloc(B * C * H * W * sizeof(float));\n float* dout = (float*)malloc(B * C * H*2 * W*2 * sizeof(float));\n\n // read saved output\n FILE *file = fopen(\"upsample.bin\", \"rb\");\n if (!file) {\n perror(\"Failed to load data\");\n return -1;\n }\n freadCheck(x, sizeof(float), B * C * H * W, file);\n freadCheck(out, sizeof(float), B * C * H*2 * W*2, file);\n freadCheck(dout, sizeof(float), B * C * H*2 * W*2, file);\n freadCheck(dx, sizeof(float), B * C * H * W, file);\n fclose(file);\n\n\n // allocate device memory\n float *d_x, *d_out, *d_dout, *d_dx;\n cudaCheck(hipMalloc(&d_x, B * C * H * W * sizeof(float)));\n cudaCheck(hipMalloc(&d_out, B * C * H*2 * W*2 * sizeof(float)));\n cudaCheck(hipMalloc(&d_dout, B * C * H*2 * W*2 * sizeof(float)));\n cudaCheck(hipMalloc(&d_dx, B * C * H * W * sizeof(float)));\n cudaCheck(hipMemcpy(d_x, x, B * C * H * W * sizeof(float), hipMemcpyHostToDevice));\n cudaCheck(hipMemcpy(d_dout, dout, B * C * H*2 * W*2 * sizeof(float), hipMemcpyHostToDevice));\n\n // read kernel number from command line\n int kernel_num = 1;\n if (argc > 1) {\n kernel_num = atoi(argv[1]);\n }\n printf(\"Running kernel %d\\n\", kernel_num);\n\n int block_sizes[] = {32, 64, 128};\n printf(\"Checking forward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"Checking block size %d\\n\", block_size);\n upsample_forward(kernel_num, d_out, d_x, B, C, H, W, block_size);\n validate_result(d_out, out, \"out\", B * C * (H/2) * (W/2));\n }\n printf(\"Forward pass: all results match\\n\\n\");\n\n printf(\"Checking backward pass\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n printf(\"Checking block size %d\\n\", block_size);\n upsample_backward(kernel_num, d_dx, d_dout, B, C, H, W, block_size);\n validate_result(d_dx, dx, \"dx\", B * C * H * W);\n }\n printf(\"Backward pass: all results match\\n\\n\");\n printf(\"All results match. Starting benchmarks.\\n\\n\");\n printf(\"Forward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, upsample_forward,\n kernel_num, d_out, d_x,\n B, C, H, W, block_size);\n\n float tflops = (float)B * H * W * C / elapsed_time * 1e3f / 1e12f;\n printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n }\n\n printf(\"\\nBackward pass benchmarks:\\n\");\n for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {\n int block_size = block_sizes[j];\n\n int repeat_times = 100;\n float elapsed_time = benchmark_kernel(repeat_times, upsample_backward,\n kernel_num, d_dx, d_dout,\n B, C, H, W, block_size);\n\n float tflops = (float)B * H * W * C / elapsed_time * 1e3f / 1e12f;\n printf(\"block_size %4d | time %.4f ms | tflops %.2f\\n\", block_size, elapsed_time, tflops);\n }\n\n // free memory\n free(x);\n free(out);\n free(dout);\n free(dx);\n\n hipFree(d_x);\n hipFree(d_out);\n hipFree(d_dout);\n hipFree(d_dx);\n\n return 0;\n}\n#endif", "source": "/kaggle/working/unet.cu/dev/upsample.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 3 --experimental --amap --skip-excluded-preprocessor-conditional-blocks %clang_args\n\n// CHECK: #include \n#include \n\n#define NUM_WORK_ELEMENTS (512)\n#define WORK_BUFFER_SIZE (NUM_WORK_ELEMENTS * sizeof(float))\n#define NUM_NODES NUM_WORK_ELEMENTS\n#define NUM_FRAMES 1000\n\ntypedef struct {\n int index;\n float adder;\n} node_t;\n\ntypedef struct {\n node_t *vector_table[NUM_NODES];\n node_t nodes[NUM_NODES];\n} graph_control_t;\n\nfloat *host_in_p = nullptr;\nfloat *host_out_p = nullptr;\ngraph_control_t *host_gc_p = nullptr;\nfloat *dev_in_p = nullptr;\nfloat *dev_out_p = nullptr;\ngraph_control_t *dev_gc_p = nullptr;\nint global_index[NUM_NODES] = {0};\n\n// CHECK: hipGraph_t graph;\n// CHECK-NEXT: hipGraphExec_t graphExec;\n// CHECK-NEXT: hipGraphExec_t instance;\n// CHECK-NEXT: hipStream_t stream;\n// CHECK-NEXT: hipGraphNode_t node[NUM_NODES] = { 0 };\ncudaGraph_t graph;\ncudaGraphExec_t graphExec;\ncudaGraphExec_t instance;\ncudaStream_t stream;\ncudaGraphNode_t node[NUM_NODES] = { 0 };\n\n__global__\nvoid add(const void *index,\n const graph_control_t *gc_p,\n const float *a,\n float *b) {\n const int *node_index_p = static_cast(index);\n int node_index = *node_index_p;\n float adder = 0.0f;\n int i = threadIdx.x;\n node_t *node_p = gc_p->vector_table[node_index];\n node_index = node_p->index;\n adder = node_p->adder;\n if (i == node_index) {\n for (int j = 0; j < 100; ++j)\n b[i] = a[i] + adder + node_index;\n }\n}\n\nvoid init(void) {\n int i = 0;\n // CHECK: hipStreamCreate(&stream);\n cudaStreamCreate(&stream);\n host_in_p = (float*)malloc(sizeof(float)*WORK_BUFFER_SIZE);\n host_out_p = (float*)malloc(sizeof(float)*WORK_BUFFER_SIZE);\n host_gc_p = (graph_control_t*)malloc(sizeof(graph_control_t));\n for (i = 0; i < WORK_BUFFER_SIZE; ++i) {\n host_in_p[i] = 1;\n }\n for (i = 0; i < WORK_BUFFER_SIZE; ++i) {\n host_out_p[i] = 42;\n }\n // CHECK: hipMalloc(&dev_in_p, WORK_BUFFER_SIZE);\n // CHECK-NEXT: hipMalloc(&dev_out_p, WORK_BUFFER_SIZE);\n // CHECK-NEXT: hipMalloc(&dev_gc_p, sizeof(graph_control_t));\n cudaMalloc(&dev_in_p, WORK_BUFFER_SIZE);\n cudaMalloc(&dev_out_p, WORK_BUFFER_SIZE);\n cudaMalloc(&dev_gc_p, sizeof(graph_control_t));\n for (i = 0; i < NUM_NODES; ++i) {\n host_gc_p->nodes[i].adder = 0;\n host_gc_p->nodes[i].index = i;\n host_gc_p->vector_table[i] = &(dev_gc_p->nodes[i]);\n global_index[i] = i;\n }\n // CHECK: if (hipGraphCreate(&graph, 0))\n if (cudaGraphCreate(&graph, 0))\n printf(\"Failed to create graph\\n\");\n\n /* Create one long vertical graph with dependency between each element\n Would have parallel nodes and multiple instances of \n it running in parallel if this was a real use-case,\n but since that is not the point here, lets skip that.\n */\n for (i = 0; i < NUM_NODES; ++i) {\n void *kargs[] = { &global_index[i],\n &dev_gc_p,\n &dev_in_p,\n &dev_out_p };\n // CHECK: hipKernelNodeParams params = { .func = (void*)add,\n cudaKernelNodeParams params = { .func = (void*)add,\n .gridDim = dim3(1, 1, 1),\n .blockDim = dim3(NUM_WORK_ELEMENTS, 1 ,1),\n .sharedMemBytes = 0,\n .kernelParams = kargs,\n .extra = NULL };\n // CHECK: if (hipGraphAddKernelNode(&node[i],\n if (cudaGraphAddKernelNode(&node[i],\n graph,\n 0,\n 0,\n ¶ms))\n printf(\"Failed to create kernel node\\n\");\n }\n for (i = 0; i < (NUM_NODES - 1); ++i)\n // [TODO][#2062] Rename all DO-NOT-CHECK back\n#if CUDA_VERSION < 13000\n // DO-NOT-CHECK: hipGraphAddDependencies(graph,\n cudaGraphAddDependencies(graph,\n &node[i],\n &node[i+1],\n 1);\n#endif\n // CHECK: hipGraphInstantiate(&graphExec,\n cudaGraphInstantiate(&graphExec,\n graph,\n NULL,\n NULL,\n 0);\n /* Same input for dataplane every frame */\n // CHECK: hipMemcpy(dev_in_p,\n cudaMemcpy(dev_in_p,\n host_in_p,\n WORK_BUFFER_SIZE,\n // CHECK: hipMemcpyHostToDevice);\n cudaMemcpyHostToDevice);\n}\n\nvoid clean_up() {\n free(host_in_p);\n free(host_out_p);\n free(host_gc_p);\n // CHECK: hipFree(dev_in_p);\n // CHECK-NEXT: hipFree(dev_out_p);\n // CHECK-NEXT: hipFree(dev_gc_p);\n // CHECK-NEXT: hipDeviceReset();\n cudaFree(dev_in_p);\n cudaFree(dev_out_p);\n cudaFree(dev_gc_p);\n cudaDeviceReset();\n}\n\nvoid graph_launch() {\n // CHECK: hipGraphLaunch(graphExec, stream);\n cudaGraphLaunch(graphExec, stream);\n}\n\n/* Set up graphs, update input, execute and read out output once per frame */\nint main (void) {\n int i = 0;\n /* Set up kernels (including parameters) and graphs once at init stage */\n init();\n /* Execute a bunch of frames */\n for (int framenbr = 0; framenbr < NUM_FRAMES; ++framenbr) {\n /* Set up control structure for one frame */\n for(i = 0; i < NUM_NODES; ++i) {\n host_gc_p->nodes[i].adder = framenbr * 1.0;\n }\n /* One copy for all data needed to execute one frame, queued in same stream as the rest */\n // CHECK: hipMemcpyAsync(dev_gc_p,\n cudaMemcpyAsync(dev_gc_p,\n host_gc_p,\n sizeof(graph_control_t),\n // CHECK: hipMemcpyHostToDevice,\n cudaMemcpyHostToDevice,\n stream);\n /* Kick graph */\n graph_launch();\n /* One read-out for all data produced this frame, queued in same stream as the rest. */\n // CHECK: hipMemcpyAsync(host_out_p,\n cudaMemcpyAsync(host_out_p,\n dev_out_p,\n WORK_BUFFER_SIZE,\n // CHECK: hipMemcpyDeviceToHost,\n cudaMemcpyDeviceToHost,\n stream);\n\n /*****************************************************************************************/\n /* Here, use the offloaded CPU core to do things involved in setting up the control and */ \n /* data plane info for next frame */\n /*****************************************************************************************/\n\n /* If still not done, wait for current frame to finish */\n // CHECK: hipStreamSynchronize(stream);\n cudaStreamSynchronize(stream);\n /* Just some printouts to uncomment for quick sanity check of functional behaviour */\n printf(\"\\n\\nOutput buffer, frame %d:\\n\", framenbr);\n for(i = 0; i < NUM_WORK_ELEMENTS; ++i) {\n printf(\"%f \", host_out_p[i]);\n if (i && i % 16 == 0)\n printf(\"\\n\");\n }\n }\n clean_up();\n return 0;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 3 --experimental --amap --skip-excluded-preprocessor-conditional-blocks %clang_args\n\n// CHECK: #include \n#include \"hip/hip_runtime.h\"\n#include \n\n#define NUM_WORK_ELEMENTS (512)\n#define WORK_BUFFER_SIZE (NUM_WORK_ELEMENTS * sizeof(float))\n#define NUM_NODES NUM_WORK_ELEMENTS\n#define NUM_FRAMES 1000\n\ntypedef struct {\n int index;\n float adder;\n} node_t;\n\ntypedef struct {\n node_t *vector_table[NUM_NODES];\n node_t nodes[NUM_NODES];\n} graph_control_t;\n\nfloat *host_in_p = nullptr;\nfloat *host_out_p = nullptr;\ngraph_control_t *host_gc_p = nullptr;\nfloat *dev_in_p = nullptr;\nfloat *dev_out_p = nullptr;\ngraph_control_t *dev_gc_p = nullptr;\nint global_index[NUM_NODES] = {0};\n\n// CHECK: hipGraph_t graph;\n// CHECK-NEXT: hipGraphExec_t graphExec;\n// CHECK-NEXT: hipGraphExec_t instance;\n// CHECK-NEXT: hipStream_t stream;\n// CHECK-NEXT: hipGraphNode_t node[NUM_NODES] = { 0 };\nhipGraph_t graph;\nhipGraphExec_t graphExec;\nhipGraphExec_t instance;\nhipStream_t stream;\nhipGraphNode_t node[NUM_NODES] = { 0 };\n\n__global__\nvoid add(const void *index,\n const graph_control_t *gc_p,\n const float *a,\n float *b) {\n const int *node_index_p = static_cast(index);\n int node_index = *node_index_p;\n float adder = 0.0f;\n int i = threadIdx.x;\n node_t *node_p = gc_p->vector_table[node_index];\n node_index = node_p->index;\n adder = node_p->adder;\n if (i == node_index) {\n for (int j = 0; j < 100; ++j)\n b[i] = a[i] + adder + node_index;\n }\n}\n\nvoid init(void) {\n int i = 0;\n // CHECK: hipStreamCreate(&stream);\n hipStreamCreate(&stream);\n host_in_p = (float*)malloc(sizeof(float)*WORK_BUFFER_SIZE);\n host_out_p = (float*)malloc(sizeof(float)*WORK_BUFFER_SIZE);\n host_gc_p = (graph_control_t*)malloc(sizeof(graph_control_t));\n for (i = 0; i < WORK_BUFFER_SIZE; ++i) {\n host_in_p[i] = 1;\n }\n for (i = 0; i < WORK_BUFFER_SIZE; ++i) {\n host_out_p[i] = 42;\n }\n // CHECK: hipMalloc(&dev_in_p, WORK_BUFFER_SIZE);\n // CHECK-NEXT: hipMalloc(&dev_out_p, WORK_BUFFER_SIZE);\n // CHECK-NEXT: hipMalloc(&dev_gc_p, sizeof(graph_control_t));\n hipMalloc(&dev_in_p, WORK_BUFFER_SIZE);\n hipMalloc(&dev_out_p, WORK_BUFFER_SIZE);\n hipMalloc(&dev_gc_p, sizeof(graph_control_t));\n for (i = 0; i < NUM_NODES; ++i) {\n host_gc_p->nodes[i].adder = 0;\n host_gc_p->nodes[i].index = i;\n host_gc_p->vector_table[i] = &(dev_gc_p->nodes[i]);\n global_index[i] = i;\n }\n // CHECK: if (hipGraphCreate(&graph, 0))\n if (hipGraphCreate(&graph, 0))\n printf(\"Failed to create graph\\n\");\n\n /* Create one long vertical graph with dependency between each element\n Would have parallel nodes and multiple instances of \n it running in parallel if this was a real use-case,\n but since that is not the point here, lets skip that.\n */\n for (i = 0; i < NUM_NODES; ++i) {\n void *kargs[] = { &global_index[i],\n &dev_gc_p,\n &dev_in_p,\n &dev_out_p };\n // CHECK: hipKernelNodeParams params = { .func = (void*)add,\n hipKernelNodeParams params = { .func = (void*)add,\n .gridDim = dim3(1, 1, 1),\n .blockDim = dim3(NUM_WORK_ELEMENTS, 1 ,1),\n .sharedMemBytes = 0,\n .kernelParams = kargs,\n .extra = NULL };\n // CHECK: if (hipGraphAddKernelNode(&node[i],\n if (hipGraphAddKernelNode(&node[i],\n graph,\n 0,\n 0,\n ¶ms))\n printf(\"Failed to create kernel node\\n\");\n }\n for (i = 0; i < (NUM_NODES - 1); ++i)\n // [TODO][#2062] Rename all DO-NOT-CHECK back\n#if CUDA_VERSION < 13000\n // DO-NOT-CHECK: hipGraphAddDependencies(graph,\n hipGraphAddDependencies(graph,\n &node[i],\n &node[i+1],\n 1);\n#endif\n // CHECK: hipGraphInstantiate(&graphExec,\n hipGraphInstantiate(&graphExec,\n graph,\n NULL,\n NULL,\n 0);\n /* Same input for dataplane every frame */\n // CHECK: hipMemcpy(dev_in_p,\n hipMemcpy(dev_in_p,\n host_in_p,\n WORK_BUFFER_SIZE,\n // CHECK: hipMemcpyHostToDevice);\n hipMemcpyHostToDevice);\n}\n\nvoid clean_up() {\n free(host_in_p);\n free(host_out_p);\n free(host_gc_p);\n // CHECK: hipFree(dev_in_p);\n // CHECK-NEXT: hipFree(dev_out_p);\n // CHECK-NEXT: hipFree(dev_gc_p);\n // CHECK-NEXT: hipDeviceReset();\n hipFree(dev_in_p);\n hipFree(dev_out_p);\n hipFree(dev_gc_p);\n hipDeviceReset();\n}\n\nvoid graph_launch() {\n // CHECK: hipGraphLaunch(graphExec, stream);\n hipGraphLaunch(graphExec, stream);\n}\n\n/* Set up graphs, update input, execute and read out output once per frame */\nint main (void) {\n int i = 0;\n /* Set up kernels (including parameters) and graphs once at init stage */\n init();\n /* Execute a bunch of frames */\n for (int framenbr = 0; framenbr < NUM_FRAMES; ++framenbr) {\n /* Set up control structure for one frame */\n for(i = 0; i < NUM_NODES; ++i) {\n host_gc_p->nodes[i].adder = framenbr * 1.0;\n }\n /* One copy for all data needed to execute one frame, queued in same stream as the rest */\n // CHECK: hipMemcpyAsync(dev_gc_p,\n hipMemcpyAsync(dev_gc_p,\n host_gc_p,\n sizeof(graph_control_t),\n // CHECK: hipMemcpyHostToDevice,\n hipMemcpyHostToDevice,\n stream);\n /* Kick graph */\n graph_launch();\n /* One read-out for all data produced this frame, queued in same stream as the rest. */\n // CHECK: hipMemcpyAsync(host_out_p,\n hipMemcpyAsync(host_out_p,\n dev_out_p,\n WORK_BUFFER_SIZE,\n // CHECK: hipMemcpyDeviceToHost,\n hipMemcpyDeviceToHost,\n stream);\n\n /*****************************************************************************************/\n /* Here, use the offloaded CPU core to do things involved in setting up the control and */ \n /* data plane info for next frame */\n /*****************************************************************************************/\n\n /* If still not done, wait for current frame to finish */\n // CHECK: hipStreamSynchronize(stream);\n hipStreamSynchronize(stream);\n /* Just some printouts to uncomment for quick sanity check of functional behaviour */\n printf(\"\\n\\nOutput buffer, frame %d:\\n\", framenbr);\n for(i = 0; i < NUM_WORK_ELEMENTS; ++i) {\n printf(\"%f \", host_out_p[i]);\n if (i && i % 16 == 0)\n printf(\"\\n\");\n }\n }\n clean_up();\n return 0;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/graph/simple_mechs.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args %clang_args\n// CHECK: #include \n#include \n\n__global__\nvoid add(int n, float *x, float *y)\n{\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = index; i < n; i += stride)\n y[i] = x[i] + y[i];\n}\n\nint main(int argc, char *argv[])\n{\n int numElements = 10;\n bool testResult = true;\n float *A, *B;\n // CHECK: hipMallocManaged(&A, numElements * sizeof(float));\n cudaMallocManaged(&A, numElements * sizeof(float));\n // CHECK: hipMallocManaged(&B, numElements * sizeof(float));\n cudaMallocManaged(&B, numElements * sizeof(float));\n for (int i = 0; i < numElements; i++) {\n A[i] = 1.0f;\n B[i] = 2.0f;\n }\n int blockSize = 256;\n int numBlocks = (numElements + blockSize - 1) / blockSize;\n dim3 dimGrid(numBlocks, 1, 1);\n dim3 dimBlock(blockSize, 1, 1);\n // CHECK: add<<>>(numElements, A, B);\n add<<>>(numElements, A, B);\n // CHECK: hipDeviceSynchronize();\n cudaDeviceSynchronize();\n float maxError = 0.0f;\n for (int i = 0; i < numElements; i++)\n maxError = fmax(maxError, fabs(B[i]-3.0f));\n // CHECK: hipFree(A);\n cudaFree(A);\n // CHECK: hipFree(B);\n cudaFree(B);\n if(maxError == 0.0f)\n return 0;\n return -1;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args %clang_args\n// CHECK: #include \n#include \"hip/hip_runtime.h\"\n#include \n\n__global__\nvoid add(int n, float *x, float *y)\n{\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = index; i < n; i += stride)\n y[i] = x[i] + y[i];\n}\n\nint main(int argc, char *argv[])\n{\n int numElements = 10;\n bool testResult = true;\n float *A, *B;\n // CHECK: hipMallocManaged(&A, numElements * sizeof(float));\n hipMallocManaged(&A, numElements * sizeof(float));\n // CHECK: hipMallocManaged(&B, numElements * sizeof(float));\n hipMallocManaged(&B, numElements * sizeof(float));\n for (int i = 0; i < numElements; i++) {\n A[i] = 1.0f;\n B[i] = 2.0f;\n }\n int blockSize = 256;\n int numBlocks = (numElements + blockSize - 1) / blockSize;\n dim3 dimGrid(numBlocks, 1, 1);\n dim3 dimBlock(blockSize, 1, 1);\n // CHECK: add<<>>(numElements, A, B);\n add<<>>(numElements, A, B);\n // CHECK: hipDeviceSynchronize();\n hipDeviceSynchronize();\n float maxError = 0.0f;\n for (int i = 0; i < numElements; i++)\n maxError = fmax(maxError, fabs(B[i]-3.0f));\n // CHECK: hipFree(A);\n hipFree(A);\n // CHECK: hipFree(B);\n hipFree(B);\n if(maxError == 0.0f)\n return 0;\n return -1;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/options/kernel-execution-syntax/none-kernel-execution-syntax.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n#include \n\n__global__\nvoid add(int n, float *x, float *y)\n{\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = index; i < n; i += stride)\n y[i] = x[i] + y[i];\n}\n\nint main(int argc, char *argv[])\n{\n int numElements = 10;\n bool testResult = true;\n float *A, *B;\n // CHECK: hipMallocManaged(&A, numElements * sizeof(float));\n cudaMallocManaged(&A, numElements * sizeof(float));\n // CHECK: hipMallocManaged(&B, numElements * sizeof(float));\n cudaMallocManaged(&B, numElements * sizeof(float));\n for (int i = 0; i < numElements; i++) {\n A[i] = 1.0f;\n B[i] = 2.0f;\n }\n int blockSize = 256;\n int numBlocks = (numElements + blockSize - 1) / blockSize;\n dim3 dimGrid(numBlocks, 1, 1);\n dim3 dimBlock(blockSize, 1, 1);\n // CHECK: hipLaunchKernelGGL(add, dim3(dimGrid), dim3(dimBlock), 0, 0, numElements, A, B);\n add<<>>(numElements, A, B);\n // CHECK: hipDeviceSynchronize();\n cudaDeviceSynchronize();\n float maxError = 0.0f;\n for (int i = 0; i < numElements; i++)\n maxError = fmax(maxError, fabs(B[i]-3.0f));\n // CHECK: hipFree(A);\n cudaFree(A);\n // CHECK: hipFree(B);\n cudaFree(B);\n if(maxError == 0.0f)\n return 0;\n return -1;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n#include \"hip/hip_runtime.h\"\n#include \n\n__global__\nvoid add(int n, float *x, float *y)\n{\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = index; i < n; i += stride)\n y[i] = x[i] + y[i];\n}\n\nint main(int argc, char *argv[])\n{\n int numElements = 10;\n bool testResult = true;\n float *A, *B;\n // CHECK: hipMallocManaged(&A, numElements * sizeof(float));\n hipMallocManaged(&A, numElements * sizeof(float));\n // CHECK: hipMallocManaged(&B, numElements * sizeof(float));\n hipMallocManaged(&B, numElements * sizeof(float));\n for (int i = 0; i < numElements; i++) {\n A[i] = 1.0f;\n B[i] = 2.0f;\n }\n int blockSize = 256;\n int numBlocks = (numElements + blockSize - 1) / blockSize;\n dim3 dimGrid(numBlocks, 1, 1);\n dim3 dimBlock(blockSize, 1, 1);\n // CHECK: hipLaunchKernelGGL(add, dim3(dimGrid), dim3(dimBlock), 0, 0, numElements, A, B);\n add<<>>(numElements, A, B);\n // CHECK: hipDeviceSynchronize();\n hipDeviceSynchronize();\n float maxError = 0.0f;\n for (int i = 0; i < numElements; i++)\n maxError = fmax(maxError, fabs(B[i]-3.0f));\n // CHECK: hipFree(A);\n hipFree(A);\n // CHECK: hipFree(B);\n hipFree(B);\n if(maxError == 0.0f)\n return 0;\n return -1;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/options/kernel-execution-syntax/hip-kernel-execution-syntax.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --cuda-kernel-execution-syntax %clang_args\n// CHECK: #include \n#include \n\n__global__\nvoid add(int n, float *x, float *y)\n{\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = index; i < n; i += stride)\n y[i] = x[i] + y[i];\n}\n\nint main(int argc, char *argv[])\n{\n int numElements = 10;\n bool testResult = true;\n float *A, *B;\n // CHECK: hipMallocManaged(&A, numElements * sizeof(float));\n cudaMallocManaged(&A, numElements * sizeof(float));\n // CHECK: hipMallocManaged(&B, numElements * sizeof(float));\n cudaMallocManaged(&B, numElements * sizeof(float));\n for (int i = 0; i < numElements; i++) {\n A[i] = 1.0f;\n B[i] = 2.0f;\n }\n int blockSize = 256;\n int numBlocks = (numElements + blockSize - 1) / blockSize;\n dim3 dimGrid(numBlocks, 1, 1);\n dim3 dimBlock(blockSize, 1, 1);\n // CHECK: add<<>>(numElements, A, B);\n add<<>>(numElements, A, B);\n // CHECK: hipDeviceSynchronize();\n cudaDeviceSynchronize();\n float maxError = 0.0f;\n for (int i = 0; i < numElements; i++)\n maxError = fmax(maxError, fabs(B[i]-3.0f));\n // CHECK: hipFree(A);\n cudaFree(A);\n // CHECK: hipFree(B);\n cudaFree(B);\n if(maxError == 0.0f)\n return 0;\n return -1;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --cuda-kernel-execution-syntax %clang_args\n// CHECK: #include \n#include \"hip/hip_runtime.h\"\n#include \n\n__global__\nvoid add(int n, float *x, float *y)\n{\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = index; i < n; i += stride)\n y[i] = x[i] + y[i];\n}\n\nint main(int argc, char *argv[])\n{\n int numElements = 10;\n bool testResult = true;\n float *A, *B;\n // CHECK: hipMallocManaged(&A, numElements * sizeof(float));\n hipMallocManaged(&A, numElements * sizeof(float));\n // CHECK: hipMallocManaged(&B, numElements * sizeof(float));\n hipMallocManaged(&B, numElements * sizeof(float));\n for (int i = 0; i < numElements; i++) {\n A[i] = 1.0f;\n B[i] = 2.0f;\n }\n int blockSize = 256;\n int numBlocks = (numElements + blockSize - 1) / blockSize;\n dim3 dimGrid(numBlocks, 1, 1);\n dim3 dimBlock(blockSize, 1, 1);\n // CHECK: add<<>>(numElements, A, B);\n add<<>>(numElements, A, B);\n // CHECK: hipDeviceSynchronize();\n hipDeviceSynchronize();\n float maxError = 0.0f;\n for (int i = 0; i < numElements; i++)\n maxError = fmax(maxError, fabs(B[i]-3.0f));\n // CHECK: hipFree(A);\n hipFree(A);\n // CHECK: hipFree(B);\n hipFree(B);\n if(maxError == 0.0f)\n return 0;\n return -1;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/options/kernel-execution-syntax/cuda-kernel-execution-syntax.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 2 --cuda-kernel-execution-syntax --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n#include \n\n__global__\nvoid add(int n, float *x, float *y)\n{\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = index; i < n; i += stride)\n y[i] = x[i] + y[i];\n}\n\nint main(int argc, char *argv[])\n{\n int numElements = 10;\n bool testResult = true;\n float *A, *B;\n // CHECK: hipMallocManaged(&A, numElements * sizeof(float));\n cudaMallocManaged(&A, numElements * sizeof(float));\n // CHECK: hipMallocManaged(&B, numElements * sizeof(float));\n cudaMallocManaged(&B, numElements * sizeof(float));\n for (int i = 0; i < numElements; i++) {\n A[i] = 1.0f;\n B[i] = 2.0f;\n }\n int blockSize = 256;\n int numBlocks = (numElements + blockSize - 1) / blockSize;\n dim3 dimGrid(numBlocks, 1, 1);\n dim3 dimBlock(blockSize, 1, 1);\n // CHECK: hipLaunchKernelGGL(add, dim3(dimGrid), dim3(dimBlock), 0, 0, numElements, A, B);\n add<<>>(numElements, A, B);\n // CHECK: hipDeviceSynchronize();\n cudaDeviceSynchronize();\n float maxError = 0.0f;\n for (int i = 0; i < numElements; i++)\n maxError = fmax(maxError, fabs(B[i]-3.0f));\n // CHECK: hipFree(A);\n cudaFree(A);\n // CHECK: hipFree(B);\n cudaFree(B);\n if(maxError == 0.0f)\n return 0;\n return -1;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 2 --cuda-kernel-execution-syntax --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n#include \"hip/hip_runtime.h\"\n#include \n\n__global__\nvoid add(int n, float *x, float *y)\n{\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = index; i < n; i += stride)\n y[i] = x[i] + y[i];\n}\n\nint main(int argc, char *argv[])\n{\n int numElements = 10;\n bool testResult = true;\n float *A, *B;\n // CHECK: hipMallocManaged(&A, numElements * sizeof(float));\n hipMallocManaged(&A, numElements * sizeof(float));\n // CHECK: hipMallocManaged(&B, numElements * sizeof(float));\n hipMallocManaged(&B, numElements * sizeof(float));\n for (int i = 0; i < numElements; i++) {\n A[i] = 1.0f;\n B[i] = 2.0f;\n }\n int blockSize = 256;\n int numBlocks = (numElements + blockSize - 1) / blockSize;\n dim3 dimGrid(numBlocks, 1, 1);\n dim3 dimBlock(blockSize, 1, 1);\n // CHECK: hipLaunchKernelGGL(add, dim3(dimGrid), dim3(dimBlock), 0, 0, numElements, A, B);\n add<<>>(numElements, A, B);\n // CHECK: hipDeviceSynchronize();\n hipDeviceSynchronize();\n float maxError = 0.0f;\n for (int i = 0; i < numElements; i++)\n maxError = fmax(maxError, fabs(B[i]-3.0f));\n // CHECK: hipFree(A);\n hipFree(A);\n // CHECK: hipFree(B);\n hipFree(B);\n if(maxError == 0.0f)\n return 0;\n return -1;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/options/kernel-execution-syntax/both-kernel-execution-syntax.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 3 --amap --skip-excluded-preprocessor-conditional-blocks --experimental %clang_args -D__CUDA_API_VERSION_INTERNAL\n\n// CHECK: #include \n// CHECK-NEXT: #include \n// CHECK-NEXT: #include \n#include \n#include \n#include \n#include \n\n __global__ __constant__ float INF_F;\n __global__ __constant__ float NAN_F;\n __global__ __constant__ float MIN_DENORM_F;\n __global__ __constant__ float MAX_NORMAL_F;\n __global__ __constant__ float NEG_ZERO_F;\n __global__ __constant__ float ZERO_F;\n __global__ __constant__ float ONE_F;\n __global__ __constant__ float SQRT_HALF_F;\n __global__ __constant__ float SQRT_HALF_HI_F;\n __global__ __constant__ float SQRT_HALF_LO_F;\n __global__ __constant__ float SQRT_TWO_F;\n __global__ __constant__ float THIRD_F;\n __global__ __constant__ float PIO4_F;\n __global__ __constant__ float PIO2_F;\n __global__ __constant__ float _3PIO4_F;\n __global__ __constant__ float _2_OVER_PI_F;\n __global__ __constant__ float SQRT_2_OVER_PI_F;\n __global__ __constant__ float PI_F;\n __global__ __constant__ float L2E_F;\n __global__ __constant__ float L2T_F;\n __global__ __constant__ float LG2_F;\n __global__ __constant__ float LGE_F;\n __global__ __constant__ float LN2_F;\n __global__ __constant__ float LNT_F;\n __global__ __constant__ float LNPI_F;\n __global__ __constant__ float TWO_TO_M126_F;\n __global__ __constant__ float TWO_TO_126_F;\n __global__ __constant__ float NORM_HUGE_F;\n __global__ __constant__ float TWO_TO_23_F;\n __global__ __constant__ float TWO_TO_24_F;\n __global__ __constant__ float TWO_TO_31_F;\n __global__ __constant__ float TWO_TO_32_F;\n __global__ __constant__ float REMQUO_BITS_F;\n __global__ __constant__ float REMQUO_MASK_F;\n __global__ __constant__ double TRIG_PLOSS_F;\n __global__ __constant__ double INF;\n __global__ __constant__ double NAN_;\n __global__ __constant__ double NEG_ZERO;\n __global__ __constant__ double MIN_DENORM;\n __global__ __constant__ double ZERO;\n __global__ __constant__ double ONE;\n __global__ __constant__ double SQRT_TWO;\n __global__ __constant__ double SQRT_HALF;\n __global__ __constant__ double SQRT_HALF_HI;\n __global__ __constant__ double SQRT_HALF_LO;\n __global__ __constant__ double THIRD;\n __global__ __constant__ double TWOTHIRD;\n __global__ __constant__ double PIO4;\n __global__ __constant__ double PIO4_HI;\n __global__ __constant__ double PIO4_LO;\n __global__ __constant__ double PIO2;\n __global__ __constant__ double PIO2_HI;\n __global__ __constant__ double PIO2_LO;\n __global__ __constant__ double _3PIO4;\n __global__ __constant__ double _2_OVER_PI;\n __global__ __constant__ double _PI;\n __global__ __constant__ double PI_HI;\n __global__ __constant__ double PI_LO;\n __global__ __constant__ double SQRT_2PI;\n __global__ __constant__ double SQRT_2PI_HI;\n __global__ __constant__ double SQRT_2PI_LO;\n __global__ __constant__ double SQRT_PIO2;\n __global__ __constant__ double SQRT_PIO2_HI;\n __global__ __constant__ double SQRT_PIO2_LO;\n __global__ __constant__ double SQRT_2OPI;\n __global__ __constant__ double L2E;\n __global__ __constant__ double L2E_HI;\n __global__ __constant__ double L2E_LO;\n __global__ __constant__ double L2T;\n __global__ __constant__ double LG2;\n __global__ __constant__ double LG2_HI;\n __global__ __constant__ double LG2_LO;\n __global__ __constant__ double LGE;\n __global__ __constant__ double LGE_HI;\n __global__ __constant__ double LGE_LO;\n __global__ __constant__ double LN2;\n __global__ __constant__ double LN2_HI;\n __global__ __constant__ double LN2_LO;\n __global__ __constant__ double LNT;\n __global__ __constant__ double LNT_HI;\n __global__ __constant__ double LNT_LO;\n __global__ __constant__ double LNPI;\n __global__ __constant__ double LN2_X_1024;\n __global__ __constant__ double LN2_X_1025;\n __global__ __constant__ double LN2_X_1075;\n __global__ __constant__ double LG2_X_1024;\n __global__ __constant__ double LG2_X_1075;\n __global__ __constant__ double TWO_TO_23;\n __global__ __constant__ double TWO_TO_52;\n __global__ __constant__ double TWO_TO_53;\n __global__ __constant__ double TWO_TO_54;\n __global__ __constant__ double TWO_TO_M54;\n __global__ __constant__ double TWO_TO_M1022;\n __global__ __constant__ double TRIG_PLOSS;\n __global__ __constant__ double DBL2INT_CVT;\n\n__global__ void init() {\n // CHECK: INF_F = HIP_INF_F;\n // CHECK-NEXT: NAN_F = HIP_NAN_F;\n // CHECK-NEXT: MIN_DENORM_F = HIP_MIN_DENORM_F;\n // CHECK-NEXT: MAX_NORMAL_F = HIP_MAX_NORMAL_F;\n // CHECK-NEXT: NEG_ZERO_F = HIP_NEG_ZERO_F;\n // CHECK-NEXT: ZERO_F = HIP_ZERO_F;\n // CHECK-NEXT: ONE_F = HIP_ONE_F;\n // CHECK-NEXT: SQRT_HALF_F = HIP_SQRT_HALF_F;\n // CHECK-NEXT: SQRT_HALF_HI_F = HIP_SQRT_HALF_HI_F;\n // CHECK-NEXT: SQRT_HALF_LO_F = HIP_SQRT_HALF_LO_F;\n // CHECK-NEXT: SQRT_TWO_F = HIP_SQRT_TWO_F;\n // CHECK-NEXT: THIRD_F = HIP_THIRD_F;\n // CHECK-NEXT: PIO4_F = HIP_PIO4_F;\n // CHECK-NEXT: PIO2_F = HIP_PIO2_F;\n // CHECK-NEXT: _3PIO4_F = HIP_3PIO4_F;\n // CHECK-NEXT: _2_OVER_PI_F = HIP_2_OVER_PI_F;\n // CHECK-NEXT: SQRT_2_OVER_PI_F = HIP_SQRT_2_OVER_PI_F;\n // CHECK-NEXT: PI_F = HIP_PI_F;\n // CHECK-NEXT: L2E_F = HIP_L2E_F;\n // CHECK-NEXT: L2T_F = HIP_L2T_F;\n // CHECK-NEXT: LG2_F = HIP_LG2_F;\n // CHECK-NEXT: LGE_F = HIP_LGE_F;\n // CHECK-NEXT: LN2_F = HIP_LN2_F;\n // CHECK-NEXT: LNT_F = HIP_LNT_F;\n // CHECK-NEXT: LNPI_F = HIP_LNPI_F;\n // CHECK-NEXT: TWO_TO_M126_F = HIP_TWO_TO_M126_F;\n // CHECK-NEXT: TWO_TO_126_F = HIP_TWO_TO_126_F;\n // CHECK-NEXT: NORM_HUGE_F = HIP_NORM_HUGE_F;\n // CHECK-NEXT: TWO_TO_23_F = HIP_TWO_TO_23_F;\n // CHECK-NEXT: TWO_TO_24_F = HIP_TWO_TO_24_F;\n // CHECK-NEXT: TWO_TO_31_F = HIP_TWO_TO_31_F;\n // CHECK-NEXT: TWO_TO_32_F = HIP_TWO_TO_32_F;\n // CHECK-NEXT: REMQUO_BITS_F = HIP_REMQUO_BITS_F;\n // CHECK-NEXT: REMQUO_MASK_F = HIP_REMQUO_MASK_F;\n // CHECK-NEXT: TRIG_PLOSS_F = HIP_TRIG_PLOSS_F;\n // CHECK-NEXT: INF = HIP_INF;\n // CHECK-NEXT: NAN_ = HIP_NAN;\n // CHECK-NEXT: NEG_ZERO = HIP_NEG_ZERO;\n // CHECK-NEXT: MIN_DENORM = HIP_MIN_DENORM;\n // CHECK-NEXT: ZERO = HIP_ZERO;\n // CHECK-NEXT: ONE = HIP_ONE;\n // CHECK-NEXT: SQRT_TWO = HIP_SQRT_TWO;\n // CHECK-NEXT: SQRT_HALF = HIP_SQRT_HALF;\n // CHECK-NEXT: SQRT_HALF_HI = HIP_SQRT_HALF_HI;\n // CHECK-NEXT: SQRT_HALF_LO = HIP_SQRT_HALF_LO;\n // CHECK-NEXT: THIRD = HIP_THIRD;\n // CHECK-NEXT: TWOTHIRD = HIP_TWOTHIRD;\n // CHECK-NEXT: PIO4 = HIP_PIO4;\n // CHECK-NEXT: PIO4_HI = HIP_PIO4_HI;\n // CHECK-NEXT: PIO4_LO = HIP_PIO4_LO;\n // CHECK-NEXT: PIO2 = HIP_PIO2;\n // CHECK-NEXT: PIO2_HI = HIP_PIO2_HI;\n // CHECK-NEXT: PIO2_LO = HIP_PIO2_LO;\n // CHECK-NEXT: _3PIO4 = HIP_3PIO4;\n // CHECK-NEXT: _2_OVER_PI = HIP_2_OVER_PI;\n // CHECK-NEXT: _PI = HIP_PI;\n // CHECK-NEXT: PI_HI = HIP_PI_HI;\n // CHECK-NEXT: PI_LO = HIP_PI_LO;\n // CHECK-NEXT: SQRT_2PI = HIP_SQRT_2PI;\n // CHECK-NEXT: SQRT_2PI_HI = HIP_SQRT_2PI_HI;\n // CHECK-NEXT: SQRT_2PI_LO = HIP_SQRT_2PI_LO;\n // CHECK-NEXT: SQRT_PIO2 = HIP_SQRT_PIO2;\n // CHECK-NEXT: SQRT_PIO2_HI = HIP_SQRT_PIO2_HI;\n // CHECK-NEXT: SQRT_PIO2_LO = HIP_SQRT_PIO2_LO;\n // CHECK-NEXT: SQRT_2OPI = HIP_SQRT_2OPI;\n // CHECK-NEXT: L2E = HIP_L2E;\n // CHECK-NEXT: L2E_HI = HIP_L2E_HI;\n // CHECK-NEXT: L2E_LO = HIP_L2E_LO;\n // CHECK-NEXT: L2T = HIP_L2T;\n // CHECK-NEXT: LG2 = HIP_LG2;\n // CHECK-NEXT: LG2_HI = HIP_LG2_HI;\n // CHECK-NEXT: LG2_LO = HIP_LG2_LO;\n // CHECK-NEXT: LGE = HIP_LGE;\n // CHECK-NEXT: LGE_HI = HIP_LGE_HI;\n // CHECK-NEXT: LGE_LO = HIP_LGE_LO;\n // CHECK-NEXT: LN2 = HIP_LN2;\n // CHECK-NEXT: LN2_HI = HIP_LN2_HI;\n // CHECK-NEXT: LN2_LO = HIP_LN2_LO;\n // CHECK-NEXT: LNT = HIP_LNT;\n // CHECK-NEXT: LNT_HI = HIP_LNT_HI;\n // CHECK-NEXT: LNT_LO = HIP_LNT_LO;\n // CHECK-NEXT: LNPI = HIP_LNPI;\n // CHECK-NEXT: LN2_X_1024 = HIP_LN2_X_1024;\n // CHECK-NEXT: LN2_X_1025 = HIP_LN2_X_1025;\n // CHECK-NEXT: LN2_X_1075 = HIP_LN2_X_1075;\n // CHECK-NEXT: LG2_X_1024 = HIP_LG2_X_1024;\n // CHECK-NEXT: LG2_X_1075 = HIP_LG2_X_1075;\n // CHECK-NEXT: TWO_TO_23 = HIP_TWO_TO_23;\n // CHECK-NEXT: TWO_TO_52 = HIP_TWO_TO_52;\n // CHECK-NEXT: TWO_TO_53 = HIP_TWO_TO_53;\n // CHECK-NEXT: TWO_TO_54 = HIP_TWO_TO_54;\n // CHECK-NEXT: TWO_TO_M54 = HIP_TWO_TO_M54;\n // CHECK-NEXT: TWO_TO_M1022 = HIP_TWO_TO_M1022;\n // CHECK-NEXT: TRIG_PLOSS = HIP_TRIG_PLOSS;\n // CHECK-NEXT: DBL2INT_CVT = HIP_DBL2INT_CVT;\n INF_F = CUDART_INF_F;\n NAN_F = CUDART_NAN_F;\n MIN_DENORM_F = CUDART_MIN_DENORM_F;\n MAX_NORMAL_F = CUDART_MAX_NORMAL_F;\n NEG_ZERO_F = CUDART_NEG_ZERO_F;\n ZERO_F = CUDART_ZERO_F;\n ONE_F = CUDART_ONE_F;\n SQRT_HALF_F = CUDART_SQRT_HALF_F;\n SQRT_HALF_HI_F = CUDART_SQRT_HALF_HI_F;\n SQRT_HALF_LO_F = CUDART_SQRT_HALF_LO_F;\n SQRT_TWO_F = CUDART_SQRT_TWO_F;\n THIRD_F = CUDART_THIRD_F;\n PIO4_F = CUDART_PIO4_F;\n PIO2_F = CUDART_PIO2_F;\n _3PIO4_F = CUDART_3PIO4_F;\n _2_OVER_PI_F = CUDART_2_OVER_PI_F;\n SQRT_2_OVER_PI_F = CUDART_SQRT_2_OVER_PI_F;\n PI_F = CUDART_PI_F;\n L2E_F = CUDART_L2E_F;\n L2T_F = CUDART_L2T_F;\n LG2_F = CUDART_LG2_F;\n LGE_F = CUDART_LGE_F;\n LN2_F = CUDART_LN2_F;\n LNT_F = CUDART_LNT_F;\n LNPI_F = CUDART_LNPI_F;\n TWO_TO_M126_F = CUDART_TWO_TO_M126_F;\n TWO_TO_126_F = CUDART_TWO_TO_126_F;\n NORM_HUGE_F = CUDART_NORM_HUGE_F;\n TWO_TO_23_F = CUDART_TWO_TO_23_F;\n TWO_TO_24_F = CUDART_TWO_TO_24_F;\n TWO_TO_31_F = CUDART_TWO_TO_31_F;\n TWO_TO_32_F = CUDART_TWO_TO_32_F;\n REMQUO_BITS_F = CUDART_REMQUO_BITS_F;\n REMQUO_MASK_F = CUDART_REMQUO_MASK_F;\n TRIG_PLOSS_F = CUDART_TRIG_PLOSS_F;\n INF = CUDART_INF;\n NAN_ = CUDART_NAN;\n NEG_ZERO = CUDART_NEG_ZERO;\n MIN_DENORM = CUDART_MIN_DENORM;\n ZERO = CUDART_ZERO;\n ONE = CUDART_ONE;\n SQRT_TWO = CUDART_SQRT_TWO;\n SQRT_HALF = CUDART_SQRT_HALF;\n SQRT_HALF_HI = CUDART_SQRT_HALF_HI;\n SQRT_HALF_LO = CUDART_SQRT_HALF_LO;\n THIRD = CUDART_THIRD;\n TWOTHIRD = CUDART_TWOTHIRD;\n PIO4 = CUDART_PIO4;\n PIO4_HI = CUDART_PIO4_HI;\n PIO4_LO = CUDART_PIO4_LO;\n PIO2 = CUDART_PIO2;\n PIO2_HI = CUDART_PIO2_HI;\n PIO2_LO = CUDART_PIO2_LO;\n _3PIO4 = CUDART_3PIO4;\n _2_OVER_PI = CUDART_2_OVER_PI;\n _PI = CUDART_PI;\n PI_HI = CUDART_PI_HI;\n PI_LO = CUDART_PI_LO;\n SQRT_2PI = CUDART_SQRT_2PI;\n SQRT_2PI_HI = CUDART_SQRT_2PI_HI;\n SQRT_2PI_LO = CUDART_SQRT_2PI_LO;\n SQRT_PIO2 = CUDART_SQRT_PIO2;\n SQRT_PIO2_HI = CUDART_SQRT_PIO2_HI;\n SQRT_PIO2_LO = CUDART_SQRT_PIO2_LO;\n SQRT_2OPI = CUDART_SQRT_2OPI;\n L2E = CUDART_L2E;\n L2E_HI = CUDART_L2E_HI;\n L2E_LO = CUDART_L2E_LO;\n L2T = CUDART_L2T;\n LG2 = CUDART_LG2;\n LG2_HI = CUDART_LG2_HI;\n LG2_LO = CUDART_LG2_LO;\n LGE = CUDART_LGE;\n LGE_HI = CUDART_LGE_HI;\n LGE_LO = CUDART_LGE_LO;\n LN2 = CUDART_LN2;\n LN2_HI = CUDART_LN2_HI;\n LN2_LO = CUDART_LN2_LO;\n LNT = CUDART_LNT;\n LNT_HI = CUDART_LNT_HI;\n LNT_LO = CUDART_LNT_LO;\n LNPI = CUDART_LNPI;\n LN2_X_1024 = CUDART_LN2_X_1024;\n LN2_X_1025 = CUDART_LN2_X_1025;\n LN2_X_1075 = CUDART_LN2_X_1075;\n LG2_X_1024 = CUDART_LG2_X_1024;\n LG2_X_1075 = CUDART_LG2_X_1075;\n TWO_TO_23 = CUDART_TWO_TO_23;\n TWO_TO_52 = CUDART_TWO_TO_52;\n TWO_TO_53 = CUDART_TWO_TO_53;\n TWO_TO_54 = CUDART_TWO_TO_54;\n TWO_TO_M54 = CUDART_TWO_TO_M54;\n TWO_TO_M1022 = CUDART_TWO_TO_M1022;\n TRIG_PLOSS = CUDART_TRIG_PLOSS;\n DBL2INT_CVT = CUDART_DBL2INT_CVT;\n}\n\nint main() {\n printf(\"08. CUDA Runtime API Defines synthetic test\\n\");\n\n // CHECK: int IPC_HANDLE_SIZE = HIP_IPC_HANDLE_SIZE;\n // CHECK-NEXT: int ArrayDefault = hipArrayDefault;\n // CHECK-NEXT: int ArrayLayered = hipArrayLayered;\n // CHECK-NEXT: int ArraySurfaceLoadStore = hipArraySurfaceLoadStore;\n // CHECK-NEXT: int ArrayCubemap = hipArrayCubemap;\n // CHECK-NEXT: int ArrayTextureGather = hipArrayTextureGather;\n // CHECK-NEXT: int DeviceBlockingSync = hipDeviceScheduleBlockingSync;\n // CHECK-NEXT: int DeviceLmemResizeToMax = hipDeviceLmemResizeToMax;\n // CHECK-NEXT: int DeviceMapHost = hipDeviceMapHost;\n // CHECK-NEXT: int DeviceScheduleAuto = hipDeviceScheduleAuto;\n // CHECK-NEXT: int DeviceScheduleSpin = hipDeviceScheduleSpin;\n // CHECK-NEXT: int DeviceScheduleYield = hipDeviceScheduleYield;\n // CHECK-NEXT: int DeviceScheduleBlockingSync = hipDeviceScheduleBlockingSync;\n // CHECK-NEXT: int DeviceScheduleMask = hipDeviceScheduleMask;\n // CHECK-NEXT: int EventDefault = hipEventDefault;\n // CHECK-NEXT: int EventBlockingSync = hipEventBlockingSync;\n // CHECK-NEXT: int EventDisableTiming = hipEventDisableTiming;\n // CHECK-NEXT: int EventInterprocess = hipEventInterprocess;\n // CHECK-NEXT: int HostAllocDefault = hipHostMallocDefault;\n // CHECK-NEXT: int HostAllocPortable = hipHostMallocPortable;\n // CHECK-NEXT: int HostAllocMapped = hipHostMallocMapped;\n // CHECK-NEXT: int HostAllocWriteCombined = hipHostMallocWriteCombined;\n // CHECK-NEXT: int HostRegisterDefault = hipHostRegisterDefault;\n // CHECK-NEXT: int HostRegisterPortable = hipHostRegisterPortable;\n // CHECK-NEXT: int HostRegisterMapped = hipHostRegisterMapped;\n // CHECK-NEXT: int IpcMemLazyEnablePeerAccess = hipIpcMemLazyEnablePeerAccess;\n // CHECK-NEXT: int MemAttachGlobal = hipMemAttachGlobal;\n // CHECK-NEXT: int MemAttachHost = hipMemAttachHost;\n // CHECK-NEXT: int MemAttachSingle = hipMemAttachSingle;\n // CHECK-NEXT: int TextureType1D = hipTextureType1D;\n // CHECK-NEXT: int TextureType2D = hipTextureType2D;\n // CHECK-NEXT: int TextureType3D = hipTextureType3D;\n // CHECK-NEXT: int TextureTypeCubemap = hipTextureTypeCubemap;\n // CHECK-NEXT: int TextureType1DLayered = hipTextureType1DLayered;\n // CHECK-NEXT: int TextureType2DLayered = hipTextureType2DLayered;\n // CHECK-NEXT: int TextureTypeCubemapLayered = hipTextureTypeCubemapLayered;\n // CHECK-NEXT: int OccupancyDefault = hipOccupancyDefault;\n // CHECK-NEXT: int OccupancyDisableCachingOverride = hipOccupancyDisableCachingOverride;\n // CHECK-NEXT: int StreamDefault = hipStreamDefault;\n // CHECK-NEXT: int StreamNonBlocking = hipStreamNonBlocking;\n // CHECK-NEXT: hipStream_t StreamPerThread = hipStreamPerThread;\n int IPC_HANDLE_SIZE = CUDA_IPC_HANDLE_SIZE;\n int ArrayDefault = cudaArrayDefault;\n int ArrayLayered = cudaArrayLayered;\n int ArraySurfaceLoadStore = cudaArraySurfaceLoadStore;\n int ArrayCubemap = cudaArrayCubemap;\n int ArrayTextureGather = cudaArrayTextureGather;\n int DeviceBlockingSync = cudaDeviceBlockingSync;\n int DeviceLmemResizeToMax = cudaDeviceLmemResizeToMax;\n int DeviceMapHost = cudaDeviceMapHost;\n int DeviceScheduleAuto = cudaDeviceScheduleAuto;\n int DeviceScheduleSpin = cudaDeviceScheduleSpin;\n int DeviceScheduleYield = cudaDeviceScheduleYield;\n int DeviceScheduleBlockingSync = cudaDeviceScheduleBlockingSync;\n int DeviceScheduleMask = cudaDeviceScheduleMask;\n int EventDefault = cudaEventDefault;\n int EventBlockingSync = cudaEventBlockingSync;\n int EventDisableTiming = cudaEventDisableTiming;\n int EventInterprocess = cudaEventInterprocess;\n int HostAllocDefault = cudaHostAllocDefault;\n int HostAllocPortable = cudaHostAllocPortable;\n int HostAllocMapped = cudaHostAllocMapped;\n int HostAllocWriteCombined = cudaHostAllocWriteCombined;\n int HostRegisterDefault = cudaHostRegisterDefault;\n int HostRegisterPortable = cudaHostRegisterPortable;\n int HostRegisterMapped = cudaHostRegisterMapped;\n int IpcMemLazyEnablePeerAccess = cudaIpcMemLazyEnablePeerAccess;\n int MemAttachGlobal = cudaMemAttachGlobal;\n int MemAttachHost = cudaMemAttachHost;\n int MemAttachSingle = cudaMemAttachSingle;\n int TextureType1D = cudaTextureType1D;\n int TextureType2D = cudaTextureType2D;\n int TextureType3D = cudaTextureType3D;\n int TextureTypeCubemap = cudaTextureTypeCubemap;\n int TextureType1DLayered = cudaTextureType1DLayered;\n int TextureType2DLayered = cudaTextureType2DLayered;\n int TextureTypeCubemapLayered = cudaTextureTypeCubemapLayered;\n int OccupancyDefault = cudaOccupancyDefault;\n int OccupancyDisableCachingOverride = cudaOccupancyDisableCachingOverride;\n int StreamDefault = cudaStreamDefault;\n int StreamNonBlocking = cudaStreamNonBlocking;\n cudaStream_t StreamPerThread = cudaStreamPerThread;\n\n#if CUDA_VERSION >= 7050\n // CHECK: int HostRegisterIoMemory = hipHostRegisterIoMemory;\n int HostRegisterIoMemory = cudaHostRegisterIoMemory;\n#endif\n\n#if CUDA_VERSION >= 8000\n // CHECK: int CpuDeviceId = hipCpuDeviceId;\n // CHECK-NEXT: int InvalidDeviceId = hipInvalidDeviceId;\n int CpuDeviceId = cudaCpuDeviceId;\n int InvalidDeviceId = cudaInvalidDeviceId;\n#endif\n\n#if CUDA_VERSION >= 9000\n#if CUDA_VERSION < 13000\n // CHECK: int CooperativeLaunchMultiDeviceNoPreSync = hipCooperativeLaunchMultiDeviceNoPreSync;\n // CHECK-NEXT: int CooperativeLaunchMultiDeviceNoPostSync = hipCooperativeLaunchMultiDeviceNoPostSync;\n int CooperativeLaunchMultiDeviceNoPreSync = cudaCooperativeLaunchMultiDeviceNoPreSync;\n int CooperativeLaunchMultiDeviceNoPostSync = cudaCooperativeLaunchMultiDeviceNoPostSync;\n#endif\n // CHECK: hipStream_t StreamLegacy = hipStreamLegacy;\n cudaStream_t StreamLegacy = cudaStreamLegacy;\n#endif\n\n#if CUDA_VERSION >= 10000\n // CHECK: int EXTERNAL_MEMORY_DEDICATED = hipExternalMemoryDedicated;\n int EXTERNAL_MEMORY_DEDICATED = cudaExternalMemoryDedicated;\n#endif\n\n#if CUDA_VERSION >= 11010\n // CHECK: int HostRegisterReadOnly = hipHostRegisterReadOnly;\n int HostRegisterReadOnly = cudaHostRegisterReadOnly;\n#endif\n\n#if CUDA_VERSION >= 12020\n // CHECK: int DevAttrHostNumaId = hipDeviceAttributeHostNumaId;\n int DevAttrHostNumaId = cudaDevAttrHostNumaId;\n#endif\n\n#if CUDA_VERSION >= 12030\n // CHECK: int GRAPH_KERNEL_NODE_PORT_DEFAULT = hipGraphKernelNodePortDefault;\n int GRAPH_KERNEL_NODE_PORT_DEFAULT = cudaGraphKernelNodePortDefault;\n\n // CHECK: int GRAPH_KERNEL_NODE_PORT_LAUNCH_ORDER = hipGraphKernelNodePortLaunchCompletion;\n int GRAPH_KERNEL_NODE_PORT_LAUNCH_ORDER = cudaGraphKernelNodePortLaunchCompletion;\n\n // CHECK: int GRAPH_KERNEL_NODE_PORT_PROGRAMMATIC = hipGraphKernelNodePortProgrammatic;\n int GRAPH_KERNEL_NODE_PORT_PROGRAMMATIC = cudaGraphKernelNodePortProgrammatic;\n#endif\n\n return 0;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 3 --amap --skip-excluded-preprocessor-conditional-blocks --experimental %clang_args -D__CUDA_API_VERSION_INTERNAL\n\n// CHECK: #include \n// CHECK-NEXT: #include \n// CHECK-NEXT: #include \n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n __global__ __constant__ float INF_F;\n __global__ __constant__ float NAN_F;\n __global__ __constant__ float MIN_DENORM_F;\n __global__ __constant__ float MAX_NORMAL_F;\n __global__ __constant__ float NEG_ZERO_F;\n __global__ __constant__ float ZERO_F;\n __global__ __constant__ float ONE_F;\n __global__ __constant__ float SQRT_HALF_F;\n __global__ __constant__ float SQRT_HALF_HI_F;\n __global__ __constant__ float SQRT_HALF_LO_F;\n __global__ __constant__ float SQRT_TWO_F;\n __global__ __constant__ float THIRD_F;\n __global__ __constant__ float PIO4_F;\n __global__ __constant__ float PIO2_F;\n __global__ __constant__ float _3PIO4_F;\n __global__ __constant__ float _2_OVER_PI_F;\n __global__ __constant__ float SQRT_2_OVER_PI_F;\n __global__ __constant__ float PI_F;\n __global__ __constant__ float L2E_F;\n __global__ __constant__ float L2T_F;\n __global__ __constant__ float LG2_F;\n __global__ __constant__ float LGE_F;\n __global__ __constant__ float LN2_F;\n __global__ __constant__ float LNT_F;\n __global__ __constant__ float LNPI_F;\n __global__ __constant__ float TWO_TO_M126_F;\n __global__ __constant__ float TWO_TO_126_F;\n __global__ __constant__ float NORM_HUGE_F;\n __global__ __constant__ float TWO_TO_23_F;\n __global__ __constant__ float TWO_TO_24_F;\n __global__ __constant__ float TWO_TO_31_F;\n __global__ __constant__ float TWO_TO_32_F;\n __global__ __constant__ float REMQUO_BITS_F;\n __global__ __constant__ float REMQUO_MASK_F;\n __global__ __constant__ double TRIG_PLOSS_F;\n __global__ __constant__ double INF;\n __global__ __constant__ double NAN_;\n __global__ __constant__ double NEG_ZERO;\n __global__ __constant__ double MIN_DENORM;\n __global__ __constant__ double ZERO;\n __global__ __constant__ double ONE;\n __global__ __constant__ double SQRT_TWO;\n __global__ __constant__ double SQRT_HALF;\n __global__ __constant__ double SQRT_HALF_HI;\n __global__ __constant__ double SQRT_HALF_LO;\n __global__ __constant__ double THIRD;\n __global__ __constant__ double TWOTHIRD;\n __global__ __constant__ double PIO4;\n __global__ __constant__ double PIO4_HI;\n __global__ __constant__ double PIO4_LO;\n __global__ __constant__ double PIO2;\n __global__ __constant__ double PIO2_HI;\n __global__ __constant__ double PIO2_LO;\n __global__ __constant__ double _3PIO4;\n __global__ __constant__ double _2_OVER_PI;\n __global__ __constant__ double _PI;\n __global__ __constant__ double PI_HI;\n __global__ __constant__ double PI_LO;\n __global__ __constant__ double SQRT_2PI;\n __global__ __constant__ double SQRT_2PI_HI;\n __global__ __constant__ double SQRT_2PI_LO;\n __global__ __constant__ double SQRT_PIO2;\n __global__ __constant__ double SQRT_PIO2_HI;\n __global__ __constant__ double SQRT_PIO2_LO;\n __global__ __constant__ double SQRT_2OPI;\n __global__ __constant__ double L2E;\n __global__ __constant__ double L2E_HI;\n __global__ __constant__ double L2E_LO;\n __global__ __constant__ double L2T;\n __global__ __constant__ double LG2;\n __global__ __constant__ double LG2_HI;\n __global__ __constant__ double LG2_LO;\n __global__ __constant__ double LGE;\n __global__ __constant__ double LGE_HI;\n __global__ __constant__ double LGE_LO;\n __global__ __constant__ double LN2;\n __global__ __constant__ double LN2_HI;\n __global__ __constant__ double LN2_LO;\n __global__ __constant__ double LNT;\n __global__ __constant__ double LNT_HI;\n __global__ __constant__ double LNT_LO;\n __global__ __constant__ double LNPI;\n __global__ __constant__ double LN2_X_1024;\n __global__ __constant__ double LN2_X_1025;\n __global__ __constant__ double LN2_X_1075;\n __global__ __constant__ double LG2_X_1024;\n __global__ __constant__ double LG2_X_1075;\n __global__ __constant__ double TWO_TO_23;\n __global__ __constant__ double TWO_TO_52;\n __global__ __constant__ double TWO_TO_53;\n __global__ __constant__ double TWO_TO_54;\n __global__ __constant__ double TWO_TO_M54;\n __global__ __constant__ double TWO_TO_M1022;\n __global__ __constant__ double TRIG_PLOSS;\n __global__ __constant__ double DBL2INT_CVT;\n\n__global__ void init() {\n // CHECK: INF_F = HIP_INF_F;\n // CHECK-NEXT: NAN_F = HIP_NAN_F;\n // CHECK-NEXT: MIN_DENORM_F = HIP_MIN_DENORM_F;\n // CHECK-NEXT: MAX_NORMAL_F = HIP_MAX_NORMAL_F;\n // CHECK-NEXT: NEG_ZERO_F = HIP_NEG_ZERO_F;\n // CHECK-NEXT: ZERO_F = HIP_ZERO_F;\n // CHECK-NEXT: ONE_F = HIP_ONE_F;\n // CHECK-NEXT: SQRT_HALF_F = HIP_SQRT_HALF_F;\n // CHECK-NEXT: SQRT_HALF_HI_F = HIP_SQRT_HALF_HI_F;\n // CHECK-NEXT: SQRT_HALF_LO_F = HIP_SQRT_HALF_LO_F;\n // CHECK-NEXT: SQRT_TWO_F = HIP_SQRT_TWO_F;\n // CHECK-NEXT: THIRD_F = HIP_THIRD_F;\n // CHECK-NEXT: PIO4_F = HIP_PIO4_F;\n // CHECK-NEXT: PIO2_F = HIP_PIO2_F;\n // CHECK-NEXT: _3PIO4_F = HIP_3PIO4_F;\n // CHECK-NEXT: _2_OVER_PI_F = HIP_2_OVER_PI_F;\n // CHECK-NEXT: SQRT_2_OVER_PI_F = HIP_SQRT_2_OVER_PI_F;\n // CHECK-NEXT: PI_F = HIP_PI_F;\n // CHECK-NEXT: L2E_F = HIP_L2E_F;\n // CHECK-NEXT: L2T_F = HIP_L2T_F;\n // CHECK-NEXT: LG2_F = HIP_LG2_F;\n // CHECK-NEXT: LGE_F = HIP_LGE_F;\n // CHECK-NEXT: LN2_F = HIP_LN2_F;\n // CHECK-NEXT: LNT_F = HIP_LNT_F;\n // CHECK-NEXT: LNPI_F = HIP_LNPI_F;\n // CHECK-NEXT: TWO_TO_M126_F = HIP_TWO_TO_M126_F;\n // CHECK-NEXT: TWO_TO_126_F = HIP_TWO_TO_126_F;\n // CHECK-NEXT: NORM_HUGE_F = HIP_NORM_HUGE_F;\n // CHECK-NEXT: TWO_TO_23_F = HIP_TWO_TO_23_F;\n // CHECK-NEXT: TWO_TO_24_F = HIP_TWO_TO_24_F;\n // CHECK-NEXT: TWO_TO_31_F = HIP_TWO_TO_31_F;\n // CHECK-NEXT: TWO_TO_32_F = HIP_TWO_TO_32_F;\n // CHECK-NEXT: REMQUO_BITS_F = HIP_REMQUO_BITS_F;\n // CHECK-NEXT: REMQUO_MASK_F = HIP_REMQUO_MASK_F;\n // CHECK-NEXT: TRIG_PLOSS_F = HIP_TRIG_PLOSS_F;\n // CHECK-NEXT: INF = HIP_INF;\n // CHECK-NEXT: NAN_ = HIP_NAN;\n // CHECK-NEXT: NEG_ZERO = HIP_NEG_ZERO;\n // CHECK-NEXT: MIN_DENORM = HIP_MIN_DENORM;\n // CHECK-NEXT: ZERO = HIP_ZERO;\n // CHECK-NEXT: ONE = HIP_ONE;\n // CHECK-NEXT: SQRT_TWO = HIP_SQRT_TWO;\n // CHECK-NEXT: SQRT_HALF = HIP_SQRT_HALF;\n // CHECK-NEXT: SQRT_HALF_HI = HIP_SQRT_HALF_HI;\n // CHECK-NEXT: SQRT_HALF_LO = HIP_SQRT_HALF_LO;\n // CHECK-NEXT: THIRD = HIP_THIRD;\n // CHECK-NEXT: TWOTHIRD = HIP_TWOTHIRD;\n // CHECK-NEXT: PIO4 = HIP_PIO4;\n // CHECK-NEXT: PIO4_HI = HIP_PIO4_HI;\n // CHECK-NEXT: PIO4_LO = HIP_PIO4_LO;\n // CHECK-NEXT: PIO2 = HIP_PIO2;\n // CHECK-NEXT: PIO2_HI = HIP_PIO2_HI;\n // CHECK-NEXT: PIO2_LO = HIP_PIO2_LO;\n // CHECK-NEXT: _3PIO4 = HIP_3PIO4;\n // CHECK-NEXT: _2_OVER_PI = HIP_2_OVER_PI;\n // CHECK-NEXT: _PI = HIP_PI;\n // CHECK-NEXT: PI_HI = HIP_PI_HI;\n // CHECK-NEXT: PI_LO = HIP_PI_LO;\n // CHECK-NEXT: SQRT_2PI = HIP_SQRT_2PI;\n // CHECK-NEXT: SQRT_2PI_HI = HIP_SQRT_2PI_HI;\n // CHECK-NEXT: SQRT_2PI_LO = HIP_SQRT_2PI_LO;\n // CHECK-NEXT: SQRT_PIO2 = HIP_SQRT_PIO2;\n // CHECK-NEXT: SQRT_PIO2_HI = HIP_SQRT_PIO2_HI;\n // CHECK-NEXT: SQRT_PIO2_LO = HIP_SQRT_PIO2_LO;\n // CHECK-NEXT: SQRT_2OPI = HIP_SQRT_2OPI;\n // CHECK-NEXT: L2E = HIP_L2E;\n // CHECK-NEXT: L2E_HI = HIP_L2E_HI;\n // CHECK-NEXT: L2E_LO = HIP_L2E_LO;\n // CHECK-NEXT: L2T = HIP_L2T;\n // CHECK-NEXT: LG2 = HIP_LG2;\n // CHECK-NEXT: LG2_HI = HIP_LG2_HI;\n // CHECK-NEXT: LG2_LO = HIP_LG2_LO;\n // CHECK-NEXT: LGE = HIP_LGE;\n // CHECK-NEXT: LGE_HI = HIP_LGE_HI;\n // CHECK-NEXT: LGE_LO = HIP_LGE_LO;\n // CHECK-NEXT: LN2 = HIP_LN2;\n // CHECK-NEXT: LN2_HI = HIP_LN2_HI;\n // CHECK-NEXT: LN2_LO = HIP_LN2_LO;\n // CHECK-NEXT: LNT = HIP_LNT;\n // CHECK-NEXT: LNT_HI = HIP_LNT_HI;\n // CHECK-NEXT: LNT_LO = HIP_LNT_LO;\n // CHECK-NEXT: LNPI = HIP_LNPI;\n // CHECK-NEXT: LN2_X_1024 = HIP_LN2_X_1024;\n // CHECK-NEXT: LN2_X_1025 = HIP_LN2_X_1025;\n // CHECK-NEXT: LN2_X_1075 = HIP_LN2_X_1075;\n // CHECK-NEXT: LG2_X_1024 = HIP_LG2_X_1024;\n // CHECK-NEXT: LG2_X_1075 = HIP_LG2_X_1075;\n // CHECK-NEXT: TWO_TO_23 = HIP_TWO_TO_23;\n // CHECK-NEXT: TWO_TO_52 = HIP_TWO_TO_52;\n // CHECK-NEXT: TWO_TO_53 = HIP_TWO_TO_53;\n // CHECK-NEXT: TWO_TO_54 = HIP_TWO_TO_54;\n // CHECK-NEXT: TWO_TO_M54 = HIP_TWO_TO_M54;\n // CHECK-NEXT: TWO_TO_M1022 = HIP_TWO_TO_M1022;\n // CHECK-NEXT: TRIG_PLOSS = HIP_TRIG_PLOSS;\n // CHECK-NEXT: DBL2INT_CVT = HIP_DBL2INT_CVT;\n INF_F = HIP_INF_F;\n NAN_F = HIP_NAN_F;\n MIN_DENORM_F = HIP_MIN_DENORM_F;\n MAX_NORMAL_F = HIP_MAX_NORMAL_F;\n NEG_ZERO_F = HIP_NEG_ZERO_F;\n ZERO_F = HIP_ZERO_F;\n ONE_F = HIP_ONE_F;\n SQRT_HALF_F = HIP_SQRT_HALF_F;\n SQRT_HALF_HI_F = HIP_SQRT_HALF_HI_F;\n SQRT_HALF_LO_F = HIP_SQRT_HALF_LO_F;\n SQRT_TWO_F = HIP_SQRT_TWO_F;\n THIRD_F = HIP_THIRD_F;\n PIO4_F = HIP_PIO4_F;\n PIO2_F = HIP_PIO2_F;\n _3PIO4_F = HIP_3PIO4_F;\n _2_OVER_PI_F = HIP_2_OVER_PI_F;\n SQRT_2_OVER_PI_F = HIP_SQRT_2_OVER_PI_F;\n PI_F = HIP_PI_F;\n L2E_F = HIP_L2E_F;\n L2T_F = HIP_L2T_F;\n LG2_F = HIP_LG2_F;\n LGE_F = HIP_LGE_F;\n LN2_F = HIP_LN2_F;\n LNT_F = HIP_LNT_F;\n LNPI_F = HIP_LNPI_F;\n TWO_TO_M126_F = HIP_TWO_TO_M126_F;\n TWO_TO_126_F = HIP_TWO_TO_126_F;\n NORM_HUGE_F = HIP_NORM_HUGE_F;\n TWO_TO_23_F = HIP_TWO_TO_23_F;\n TWO_TO_24_F = HIP_TWO_TO_24_F;\n TWO_TO_31_F = HIP_TWO_TO_31_F;\n TWO_TO_32_F = HIP_TWO_TO_32_F;\n REMQUO_BITS_F = HIP_REMQUO_BITS_F;\n REMQUO_MASK_F = HIP_REMQUO_MASK_F;\n TRIG_PLOSS_F = HIP_TRIG_PLOSS_F;\n INF = HIP_INF;\n NAN_ = HIP_NAN;\n NEG_ZERO = HIP_NEG_ZERO;\n MIN_DENORM = HIP_MIN_DENORM;\n ZERO = HIP_ZERO;\n ONE = HIP_ONE;\n SQRT_TWO = HIP_SQRT_TWO;\n SQRT_HALF = HIP_SQRT_HALF;\n SQRT_HALF_HI = HIP_SQRT_HALF_HI;\n SQRT_HALF_LO = HIP_SQRT_HALF_LO;\n THIRD = HIP_THIRD;\n TWOTHIRD = HIP_TWOTHIRD;\n PIO4 = HIP_PIO4;\n PIO4_HI = HIP_PIO4_HI;\n PIO4_LO = HIP_PIO4_LO;\n PIO2 = HIP_PIO2;\n PIO2_HI = HIP_PIO2_HI;\n PIO2_LO = HIP_PIO2_LO;\n _3PIO4 = HIP_3PIO4;\n _2_OVER_PI = HIP_2_OVER_PI;\n _PI = HIP_PI;\n PI_HI = HIP_PI_HI;\n PI_LO = HIP_PI_LO;\n SQRT_2PI = HIP_SQRT_2PI;\n SQRT_2PI_HI = HIP_SQRT_2PI_HI;\n SQRT_2PI_LO = HIP_SQRT_2PI_LO;\n SQRT_PIO2 = HIP_SQRT_PIO2;\n SQRT_PIO2_HI = HIP_SQRT_PIO2_HI;\n SQRT_PIO2_LO = HIP_SQRT_PIO2_LO;\n SQRT_2OPI = HIP_SQRT_2OPI;\n L2E = HIP_L2E;\n L2E_HI = HIP_L2E_HI;\n L2E_LO = HIP_L2E_LO;\n L2T = HIP_L2T;\n LG2 = HIP_LG2;\n LG2_HI = HIP_LG2_HI;\n LG2_LO = HIP_LG2_LO;\n LGE = HIP_LGE;\n LGE_HI = HIP_LGE_HI;\n LGE_LO = HIP_LGE_LO;\n LN2 = HIP_LN2;\n LN2_HI = HIP_LN2_HI;\n LN2_LO = HIP_LN2_LO;\n LNT = HIP_LNT;\n LNT_HI = HIP_LNT_HI;\n LNT_LO = HIP_LNT_LO;\n LNPI = HIP_LNPI;\n LN2_X_1024 = HIP_LN2_X_1024;\n LN2_X_1025 = HIP_LN2_X_1025;\n LN2_X_1075 = HIP_LN2_X_1075;\n LG2_X_1024 = HIP_LG2_X_1024;\n LG2_X_1075 = HIP_LG2_X_1075;\n TWO_TO_23 = HIP_TWO_TO_23;\n TWO_TO_52 = HIP_TWO_TO_52;\n TWO_TO_53 = HIP_TWO_TO_53;\n TWO_TO_54 = HIP_TWO_TO_54;\n TWO_TO_M54 = HIP_TWO_TO_M54;\n TWO_TO_M1022 = HIP_TWO_TO_M1022;\n TRIG_PLOSS = HIP_TRIG_PLOSS;\n DBL2INT_CVT = HIP_DBL2INT_CVT;\n}\n\nint main() {\n printf(\"08. CUDA Runtime API Defines synthetic test\\n\");\n\n // CHECK: int IPC_HANDLE_SIZE = HIP_IPC_HANDLE_SIZE;\n // CHECK-NEXT: int ArrayDefault = hipArrayDefault;\n // CHECK-NEXT: int ArrayLayered = hipArrayLayered;\n // CHECK-NEXT: int ArraySurfaceLoadStore = hipArraySurfaceLoadStore;\n // CHECK-NEXT: int ArrayCubemap = hipArrayCubemap;\n // CHECK-NEXT: int ArrayTextureGather = hipArrayTextureGather;\n // CHECK-NEXT: int DeviceBlockingSync = hipDeviceScheduleBlockingSync;\n // CHECK-NEXT: int DeviceLmemResizeToMax = hipDeviceLmemResizeToMax;\n // CHECK-NEXT: int DeviceMapHost = hipDeviceMapHost;\n // CHECK-NEXT: int DeviceScheduleAuto = hipDeviceScheduleAuto;\n // CHECK-NEXT: int DeviceScheduleSpin = hipDeviceScheduleSpin;\n // CHECK-NEXT: int DeviceScheduleYield = hipDeviceScheduleYield;\n // CHECK-NEXT: int DeviceScheduleBlockingSync = hipDeviceScheduleBlockingSync;\n // CHECK-NEXT: int DeviceScheduleMask = hipDeviceScheduleMask;\n // CHECK-NEXT: int EventDefault = hipEventDefault;\n // CHECK-NEXT: int EventBlockingSync = hipEventBlockingSync;\n // CHECK-NEXT: int EventDisableTiming = hipEventDisableTiming;\n // CHECK-NEXT: int EventInterprocess = hipEventInterprocess;\n // CHECK-NEXT: int HostAllocDefault = hipHostMallocDefault;\n // CHECK-NEXT: int HostAllocPortable = hipHostMallocPortable;\n // CHECK-NEXT: int HostAllocMapped = hipHostMallocMapped;\n // CHECK-NEXT: int HostAllocWriteCombined = hipHostMallocWriteCombined;\n // CHECK-NEXT: int HostRegisterDefault = hipHostRegisterDefault;\n // CHECK-NEXT: int HostRegisterPortable = hipHostRegisterPortable;\n // CHECK-NEXT: int HostRegisterMapped = hipHostRegisterMapped;\n // CHECK-NEXT: int IpcMemLazyEnablePeerAccess = hipIpcMemLazyEnablePeerAccess;\n // CHECK-NEXT: int MemAttachGlobal = hipMemAttachGlobal;\n // CHECK-NEXT: int MemAttachHost = hipMemAttachHost;\n // CHECK-NEXT: int MemAttachSingle = hipMemAttachSingle;\n // CHECK-NEXT: int TextureType1D = hipTextureType1D;\n // CHECK-NEXT: int TextureType2D = hipTextureType2D;\n // CHECK-NEXT: int TextureType3D = hipTextureType3D;\n // CHECK-NEXT: int TextureTypeCubemap = hipTextureTypeCubemap;\n // CHECK-NEXT: int TextureType1DLayered = hipTextureType1DLayered;\n // CHECK-NEXT: int TextureType2DLayered = hipTextureType2DLayered;\n // CHECK-NEXT: int TextureTypeCubemapLayered = hipTextureTypeCubemapLayered;\n // CHECK-NEXT: int OccupancyDefault = hipOccupancyDefault;\n // CHECK-NEXT: int OccupancyDisableCachingOverride = hipOccupancyDisableCachingOverride;\n // CHECK-NEXT: int StreamDefault = hipStreamDefault;\n // CHECK-NEXT: int StreamNonBlocking = hipStreamNonBlocking;\n // CHECK-NEXT: hipStream_t StreamPerThread = hipStreamPerThread;\n int IPC_HANDLE_SIZE = HIP_IPC_HANDLE_SIZE;\n int ArrayDefault = hipArrayDefault;\n int ArrayLayered = hipArrayLayered;\n int ArraySurfaceLoadStore = hipArraySurfaceLoadStore;\n int ArrayCubemap = hipArrayCubemap;\n int ArrayTextureGather = hipArrayTextureGather;\n int DeviceBlockingSync = hipDeviceScheduleBlockingSync;\n int DeviceLmemResizeToMax = hipDeviceLmemResizeToMax;\n int DeviceMapHost = hipDeviceMapHost;\n int DeviceScheduleAuto = hipDeviceScheduleAuto;\n int DeviceScheduleSpin = hipDeviceScheduleSpin;\n int DeviceScheduleYield = hipDeviceScheduleYield;\n int DeviceScheduleBlockingSync = hipDeviceScheduleBlockingSync;\n int DeviceScheduleMask = hipDeviceScheduleMask;\n int EventDefault = hipEventDefault;\n int EventBlockingSync = hipEventBlockingSync;\n int EventDisableTiming = hipEventDisableTiming;\n int EventInterprocess = hipEventInterprocess;\n int HostAllocDefault = hipHostMallocDefault;\n int HostAllocPortable = hipHostMallocPortable;\n int HostAllocMapped = hipHostMallocMapped;\n int HostAllocWriteCombined = hipHostMallocWriteCombined;\n int HostRegisterDefault = hipHostRegisterDefault;\n int HostRegisterPortable = hipHostRegisterPortable;\n int HostRegisterMapped = hipHostRegisterMapped;\n int IpcMemLazyEnablePeerAccess = hipIpcMemLazyEnablePeerAccess;\n int MemAttachGlobal = hipMemAttachGlobal;\n int MemAttachHost = hipMemAttachHost;\n int MemAttachSingle = hipMemAttachSingle;\n int TextureType1D = hipTextureType1D;\n int TextureType2D = hipTextureType2D;\n int TextureType3D = hipTextureType3D;\n int TextureTypeCubemap = hipTextureTypeCubemap;\n int TextureType1DLayered = hipTextureType1DLayered;\n int TextureType2DLayered = hipTextureType2DLayered;\n int TextureTypeCubemapLayered = hipTextureTypeCubemapLayered;\n int OccupancyDefault = hipOccupancyDefault;\n int OccupancyDisableCachingOverride = hipOccupancyDisableCachingOverride;\n int StreamDefault = hipStreamDefault;\n int StreamNonBlocking = hipStreamNonBlocking;\n hipStream_t StreamPerThread = hipStreamPerThread;\n\n#if CUDA_VERSION >= 7050\n // CHECK: int HostRegisterIoMemory = hipHostRegisterIoMemory;\n int HostRegisterIoMemory = hipHostRegisterIoMemory;\n#endif\n\n#if CUDA_VERSION >= 8000\n // CHECK: int CpuDeviceId = hipCpuDeviceId;\n // CHECK-NEXT: int InvalidDeviceId = hipInvalidDeviceId;\n int CpuDeviceId = hipCpuDeviceId;\n int InvalidDeviceId = hipInvalidDeviceId;\n#endif\n\n#if CUDA_VERSION >= 9000\n#if CUDA_VERSION < 13000\n // CHECK: int CooperativeLaunchMultiDeviceNoPreSync = hipCooperativeLaunchMultiDeviceNoPreSync;\n // CHECK-NEXT: int CooperativeLaunchMultiDeviceNoPostSync = hipCooperativeLaunchMultiDeviceNoPostSync;\n int CooperativeLaunchMultiDeviceNoPreSync = hipCooperativeLaunchMultiDeviceNoPreSync;\n int CooperativeLaunchMultiDeviceNoPostSync = hipCooperativeLaunchMultiDeviceNoPostSync;\n#endif\n // CHECK: hipStream_t StreamLegacy = hipStreamLegacy;\n hipStream_t StreamLegacy = hipStreamLegacy;\n#endif\n\n#if CUDA_VERSION >= 10000\n // CHECK: int EXTERNAL_MEMORY_DEDICATED = hipExternalMemoryDedicated;\n int EXTERNAL_MEMORY_DEDICATED = hipExternalMemoryDedicated;\n#endif\n\n#if CUDA_VERSION >= 11010\n // CHECK: int HostRegisterReadOnly = hipHostRegisterReadOnly;\n int HostRegisterReadOnly = hipHostRegisterReadOnly;\n#endif\n\n#if CUDA_VERSION >= 12020\n // CHECK: int DevAttrHostNumaId = hipDeviceAttributeHostNumaId;\n int DevAttrHostNumaId = hipDeviceAttributeHostNumaId;\n#endif\n\n#if CUDA_VERSION >= 12030\n // CHECK: int GRAPH_KERNEL_NODE_PORT_DEFAULT = hipGraphKernelNodePortDefault;\n int GRAPH_KERNEL_NODE_PORT_DEFAULT = hipGraphKernelNodePortDefault;\n\n // CHECK: int GRAPH_KERNEL_NODE_PORT_LAUNCH_ORDER = hipGraphKernelNodePortLaunchCompletion;\n int GRAPH_KERNEL_NODE_PORT_LAUNCH_ORDER = hipGraphKernelNodePortLaunchCompletion;\n\n // CHECK: int GRAPH_KERNEL_NODE_PORT_PROGRAMMATIC = hipGraphKernelNodePortProgrammatic;\n int GRAPH_KERNEL_NODE_PORT_PROGRAMMATIC = hipGraphKernelNodePortProgrammatic;\n#endif\n\n return 0;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/synthetic/runtime_defines.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 2 --skip-excluded-preprocessor-conditional-blocks --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n\n__global__ void axpy_kernel(float a, float* x, float* y) {\n y[threadIdx.x] = a * x[threadIdx.x];\n}\n\nvoid axpy(float a, float* x, float* y) {\nfloat* y_new = nullptr;\n#ifdef SOME_MACRO\n y_new = x;\n // CHECK: axpy_kernel <<<1, 1>>> (a, y_new, x);\n axpy_kernel <<<1, 1>>> (a, y_new, x);\n#endif\n\n#ifndef SOME_MACRO\n y_new = y;\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(2), 0, 0, a, y_new, x);\n axpy_kernel <<<1, 2>>> (a, y_new, x);\n#endif\n\n#ifdef SOME_MACRO\n // CHECK: axpy_kernel <<<1, 3>>> (a, y, x);\n axpy_kernel <<<1, 3>>> (a, y, x);\n#else\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(4), 0, 0, a, x, y);\n axpy_kernel <<<1, 4>>> (a, x, y);\n#endif\n\n#ifdef SOME_MACRO\n // CHECK: axpy_kernel <<<1, 5>>> (a, y, x);\n axpy_kernel <<<1, 5>>> (a, y, x);\n#elif defined SOME_MACRO_1\n // CHECK: axpy_kernel <<<1, 6>>> (a, x, y);\n axpy_kernel <<<1, 6>>> (a, x, y);\n#else\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(7), 0, 0, a, x, y);\n axpy_kernel <<<1, 7>>> (a, x, y);\n#endif\n\n#ifndef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(8), 0, 0, a, y, x);\n axpy_kernel <<<1, 8>>> (a, y, x);\n#elif !defined(SOME_MACRO_1)\n // CHECK: axpy_kernel <<<1, 9>>> (a, x, y);\n axpy_kernel <<<1, 9>>> (a, x, y);\n#else\n // CHECK: axpy_kernel <<<1, 10>>> (a, x, y);\n axpy_kernel <<<1, 10>>> (a, x, y);\n#endif\n\n}", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 2 --skip-excluded-preprocessor-conditional-blocks --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n\n#include \"hip/hip_runtime.h\"\n__global__ void axpy_kernel(float a, float* x, float* y) {\n y[threadIdx.x] = a * x[threadIdx.x];\n}\n\nvoid axpy(float a, float* x, float* y) {\nfloat* y_new = nullptr;\n#ifdef SOME_MACRO\n y_new = x;\n // CHECK: axpy_kernel <<<1, 1>>> (a, y_new, x);\n axpy_kernel <<<1, 1>>> (a, y_new, x);\n#endif\n\n#ifndef SOME_MACRO\n y_new = y;\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(2), 0, 0, a, y_new, x);\n axpy_kernel <<<1, 2>>> (a, y_new, x);\n#endif\n\n#ifdef SOME_MACRO\n // CHECK: axpy_kernel <<<1, 3>>> (a, y, x);\n axpy_kernel <<<1, 3>>> (a, y, x);\n#else\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(4), 0, 0, a, x, y);\n axpy_kernel <<<1, 4>>> (a, x, y);\n#endif\n\n#ifdef SOME_MACRO\n // CHECK: axpy_kernel <<<1, 5>>> (a, y, x);\n axpy_kernel <<<1, 5>>> (a, y, x);\n#elif defined SOME_MACRO_1\n // CHECK: axpy_kernel <<<1, 6>>> (a, x, y);\n axpy_kernel <<<1, 6>>> (a, x, y);\n#else\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(7), 0, 0, a, x, y);\n axpy_kernel <<<1, 7>>> (a, x, y);\n#endif\n\n#ifndef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(8), 0, 0, a, y, x);\n axpy_kernel <<<1, 8>>> (a, y, x);\n#elif !defined(SOME_MACRO_1)\n // CHECK: axpy_kernel <<<1, 9>>> (a, x, y);\n axpy_kernel <<<1, 9>>> (a, x, y);\n#else\n // CHECK: axpy_kernel <<<1, 10>>> (a, x, y);\n axpy_kernel <<<1, 10>>> (a, x, y);\n#endif\n\n}", "source": "/kaggle/working/HIPIFY/tests/unit_tests/pp/pp_if_else_conditionals_01.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n\n__global__ void axpy_kernel(float a, float* x, float* y) {\n y[threadIdx.x] = a * x[threadIdx.x];\n}\n\nvoid axpy(float a, float* x, float* y) {\nfloat* y_new = nullptr;\n#ifdef SOME_MACRO\n y_new = x;\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(1), 0, 0, a, y_new, x);\n axpy_kernel <<<1, 1>>> (a, y_new, x);\n#endif\n\n#ifndef SOME_MACRO\n y_new = y;\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(2), 0, 0, a, y_new, x);\n axpy_kernel <<<1, 2>>> (a, y_new, x);\n#endif\n\n#ifdef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(3), 0, 0, a, y, x);\n axpy_kernel <<<1, 3>>> (a, y, x);\n#else\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(4), 0, 0, a, x, y);\n axpy_kernel <<<1, 4>>> (a, x, y);\n#endif\n\n#ifdef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(5), 0, 0, a, y, x);\n axpy_kernel <<<1, 5>>> (a, y, x);\n#elif defined SOME_MACRO_1\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(6), 0, 0, a, x, y);\n axpy_kernel <<<1, 6>>> (a, x, y);\n#else\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(7), 0, 0, a, x, y);\n axpy_kernel <<<1, 7>>> (a, x, y);\n#endif\n\n#ifndef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(8), 0, 0, a, y, x);\n axpy_kernel <<<1, 8>>> (a, y, x);\n#elif !defined(SOME_MACRO_1)\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(9), 0, 0, a, x, y);\n axpy_kernel <<<1, 9>>> (a, x, y);\n#else\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(10), 0, 0, a, x, y);\n axpy_kernel <<<1, 10>>> (a, x, y);\n#endif\n\n}", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n\n#include \"hip/hip_runtime.h\"\n__global__ void axpy_kernel(float a, float* x, float* y) {\n y[threadIdx.x] = a * x[threadIdx.x];\n}\n\nvoid axpy(float a, float* x, float* y) {\nfloat* y_new = nullptr;\n#ifdef SOME_MACRO\n y_new = x;\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(1), 0, 0, a, y_new, x);\n axpy_kernel <<<1, 1>>> (a, y_new, x);\n#endif\n\n#ifndef SOME_MACRO\n y_new = y;\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(2), 0, 0, a, y_new, x);\n axpy_kernel <<<1, 2>>> (a, y_new, x);\n#endif\n\n#ifdef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(3), 0, 0, a, y, x);\n axpy_kernel <<<1, 3>>> (a, y, x);\n#else\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(4), 0, 0, a, x, y);\n axpy_kernel <<<1, 4>>> (a, x, y);\n#endif\n\n#ifdef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(5), 0, 0, a, y, x);\n axpy_kernel <<<1, 5>>> (a, y, x);\n#elif defined SOME_MACRO_1\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(6), 0, 0, a, x, y);\n axpy_kernel <<<1, 6>>> (a, x, y);\n#else\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(7), 0, 0, a, x, y);\n axpy_kernel <<<1, 7>>> (a, x, y);\n#endif\n\n#ifndef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(8), 0, 0, a, y, x);\n axpy_kernel <<<1, 8>>> (a, y, x);\n#elif !defined(SOME_MACRO_1)\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(9), 0, 0, a, x, y);\n axpy_kernel <<<1, 9>>> (a, x, y);\n#else\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(10), 0, 0, a, x, y);\n axpy_kernel <<<1, 10>>> (a, x, y);\n#endif\n\n}", "source": "/kaggle/working/HIPIFY/tests/unit_tests/pp/pp_if_else_conditionals_01_LLVM_10.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n\n#include \n\n__global__ void axpy_kernel(float a, float* x, float* y) {\n y[threadIdx.x] = a * x[threadIdx.x];\n}\n\nvoid axpy(float a, float* x, float* y) {\n\n#ifdef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(1), 0, 0, a, y, x);\n axpy_kernel <<<1, 1>>> (a, y, x);\n#endif\n\n#ifndef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(2), 0, 0, a, y, x);\n axpy_kernel <<<1, 2>>> (a, y, x);\n#endif\n\n#ifdef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(3), 0, 0, a, y, x);\n axpy_kernel <<<1, 3>>> (a, y, x);\n#else\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(4), 0, 0, a, x, y);\n axpy_kernel <<<1, 4>>> (a, x, y);\n#endif\n\n}", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n\n#include \n\n__global__ void axpy_kernel(float a, float* x, float* y) {\n y[threadIdx.x] = a * x[threadIdx.x];\n}\n\nvoid axpy(float a, float* x, float* y) {\n\n#ifdef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(1), 0, 0, a, y, x);\n axpy_kernel <<<1, 1>>> (a, y, x);\n#endif\n\n#ifndef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(2), 0, 0, a, y, x);\n axpy_kernel <<<1, 2>>> (a, y, x);\n#endif\n\n#ifdef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(3), 0, 0, a, y, x);\n axpy_kernel <<<1, 3>>> (a, y, x);\n#else\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(4), 0, 0, a, x, y);\n axpy_kernel <<<1, 4>>> (a, x, y);\n#endif\n\n}", "source": "/kaggle/working/HIPIFY/tests/unit_tests/pp/pp_if_else_conditionals_LLVM_10.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 2 --skip-excluded-preprocessor-conditional-blocks --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n\n#include \n\n__global__ void axpy_kernel(float a, float* x, float* y) {\n y[threadIdx.x] = a * x[threadIdx.x];\n}\n\nvoid axpy(float a, float* x, float* y) {\n\n#ifdef SOME_MACRO\n // CHECK: axpy_kernel <<<1, 1>>> (a, y, x);\n axpy_kernel <<<1, 1>>> (a, y, x);\n#endif\n\n#ifndef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(2), 0, 0, a, y, x);\n axpy_kernel <<<1, 2>>> (a, y, x);\n#endif\n\n#ifdef SOME_MACRO\n // CHECK: axpy_kernel <<<1, 3>>> (a, y, x);\n axpy_kernel <<<1, 3>>> (a, y, x);\n#else\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(4), 0, 0, a, x, y);\n axpy_kernel <<<1, 4>>> (a, x, y);\n#endif\n\n}", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 2 --skip-excluded-preprocessor-conditional-blocks --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n\n#include \n\n__global__ void axpy_kernel(float a, float* x, float* y) {\n y[threadIdx.x] = a * x[threadIdx.x];\n}\n\nvoid axpy(float a, float* x, float* y) {\n\n#ifdef SOME_MACRO\n // CHECK: axpy_kernel <<<1, 1>>> (a, y, x);\n axpy_kernel <<<1, 1>>> (a, y, x);\n#endif\n\n#ifndef SOME_MACRO\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(2), 0, 0, a, y, x);\n axpy_kernel <<<1, 2>>> (a, y, x);\n#endif\n\n#ifdef SOME_MACRO\n // CHECK: axpy_kernel <<<1, 3>>> (a, y, x);\n axpy_kernel <<<1, 3>>> (a, y, x);\n#else\n // CHECK: hipLaunchKernelGGL(axpy_kernel, dim3(1), dim3(4), 0, 0, a, x, y);\n axpy_kernel <<<1, 4>>> (a, x, y);\n#endif\n\n}", "source": "/kaggle/working/HIPIFY/tests/unit_tests/pp/pp_if_else_conditionals.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#include \n#include \n#include \n#include \"inc.h\"\n// CHECK: #include \n#include \n\n// CHECK: hipEvent_t t##_start, t##_end; \\\n// CHECK: hipEventCreate(&t##_start); \\\n// CHECK: hipEventCreate(&t##_end);\n#define TIMER_CREATE(t) \\\n cudaEvent_t t##_start, t##_end; \\\n cudaEventCreate(&t##_start); \\\n cudaEventCreate(&t##_end);\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start);\n#define TIMER_START(t) \\\n cudaEventRecord(t##_start); \\\n cudaEventSynchronize(t##_start); \\\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start); \\\n// CHECK: hipEventRecord(t##_end); \\\n// CHECK: hipEventSynchronize(t##_end); \\\n// CHECK: hipEventElapsedTime(&t, t##_start, t##_end);\n#define TIMER_END(t) \\\n cudaEventRecord(t##_start); \\\n cudaEventSynchronize(t##_start); \\\n cudaEventRecord(t##_end); \\\n cudaEventSynchronize(t##_end); \\\n cudaEventElapsedTime(&t, t##_start, t##_end); \n\n\ninline void cErrorCheck(const char *file, int line) {\n// CHECK: hipDeviceSynchronize();\n// CHECK: hipError_t err = hipGetLastError();\n// CHECK: if (err != hipSuccess) {\n// CHECK: printf(\"Error: %s\\n\", hipGetErrorString(err));\n cudaThreadSynchronize();\n cudaError_t err = cudaGetLastError();\n if (err != cudaSuccess) {\n printf(\"Error: %s\\n\", cudaGetErrorString(err));\n printf(\" @ %s: %d\\n\", file, line);\n exit(-1);\n }\n}\n\ninline dim3 K_GRID(int n, int threads = K_THREADS) {\n int blocks = (int)ceilf(sqrtf((float)n/threads));\n dim3 grid(blocks, blocks);\n return grid;\n}\n\ntypedef struct data {\n int n;\n float4 *r, *v, *f;\n} data;\n\ndata cpu, gpu;\n\n#define N 20\n\n__global__ void repulsion(data gpu);\n__global__ void integration(data gpu);\n\n\nint main() {\n printf(\"Cuda Test 1\\n\");\n\n int count = 0;\n // CHECK: hipGetDeviceCount(&count);\n cudaGetDeviceCount(&count);\n printf(\" %d CUDA devices found\\n\", count);\n if(!count) {\n ::exit(EXIT_FAILURE);\n }\n // CHECK: hipFree(0);\n cudaFree(0);\n\n cpu.n = N;\n\n cpu.r = (float4*)malloc(N * sizeof(float4));\n cpu.v = (float4*)malloc(N * sizeof(float4));\n cpu.f = (float4*)malloc(N * sizeof(float4));\n\n for(int i = 0; i < N; ++i) {\n cpu.v[i] = make_float4(0,0,0,0);\n cpu.r[i] = make_float4(RND(), RND(), RND(), 0);\n cpu.f[i] = make_float4(0,0.01,0,0);\n }\n\n gpu = cpu;\n // CHECK: hipMalloc(&gpu.r, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.v, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.f, N * sizeof(float4));\n cudaMalloc(&gpu.r, N * sizeof(float4));\n cudaMalloc(&gpu.v, N * sizeof(float4));\n cudaMalloc(&gpu.f, N * sizeof(float4));\n // CHECK: hipMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n cudaMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), cudaMemcpyHostToDevice);\n cudaMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), cudaMemcpyHostToDevice);\n cudaMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), cudaMemcpyHostToDevice);\n\n ERRORCHECK();\n float rep;\n TIMER_CREATE(rep);\n TIMER_START(rep);\n // CHECK: hipLaunchKernelGGL(integration, dim3(K_GRID(cpu.n)), dim3(K_THREADS), 0, 0, gpu);\n integration <<< K_GRID(cpu.n), K_THREADS >>>(gpu);\n\n TIMER_END(rep);\n printf(\"Took: %f ms\\n\", rep);\n ERRORCHECK();\n // CHECK: hipMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n cudaMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), cudaMemcpyDeviceToHost);\n cudaMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), cudaMemcpyDeviceToHost);\n cudaMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), cudaMemcpyDeviceToHost);\n // CHECK: hipHostFree(cpu.r);\n // CHECK: hipHostFree(cpu.v);\n // CHECK: hipHostFree(cpu.f);\n cudaFreeHost(cpu.r);\n cudaFreeHost(cpu.v);\n cudaFreeHost(cpu.f);\n // CHECK: hipFree(gpu.r);\n // CHECK: hipFree(gpu.v);\n // CHECK: hipFree(gpu.f);\n cudaFree(gpu.r);\n cudaFree(gpu.v);\n cudaFree(gpu.f);\n // CHECK: hipDeviceReset();\n cudaDeviceReset();\n\n printf(\"Results: \\n\");\n for(int i = 0; i < N; ++i) {\n printf(\"%f, %f, %f \\n\", cpu.r[i].x, cpu.r[i].y, cpu.r[i].z);\n }\n\n printf(\"Ready...\\n\");\n return 0;\n}\n\n__global__ void repulsion(data gpu) {\n int idx = K_INDEX();\n if(idx < N) {\n gpu.r[idx].x = 1;\n gpu.r[idx].y = 1;\n gpu.r[idx].z = 1;\n }\n}\n\n#define MULT4(v, s) v.x *= s; v.y *= s; v.z *= s; v.w *= s;\n#define ADD4(v1, v2) v1.x += v2.x; v1.y += v2.y; v1.z += v2.z; v1.w += v2.w;\n\n__global__ void integration(data gpu) {\n int i = K_INDEX();\n if(i < N) {\n MULT4(gpu.f[i], 0.01);\n MULT4(gpu.v[i], 0.01);\n ADD4(gpu.v[i], gpu.f[i]);\n ADD4(gpu.r[i], gpu.v[i]);\n gpu.f[i] = make_float4(0,0,0,0);\n }\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#include \n#include \n#include \n#include \"inc.h\"\n// CHECK: #include \n#include \n\n// CHECK: hipEvent_t t##_start, t##_end; \\\n// CHECK: hipEventCreate(&t##_start); \\\n// CHECK: hipEventCreate(&t##_end);\n#define TIMER_CREATE(t) \\\n hipEvent_t t##_start, t##_end; \\\n hipEventCreate(&t##_start); \\\n hipEventCreate(&t##_end);\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start);\n#define TIMER_START(t) \\\n hipEventRecord(t##_start); \\\n hipEventSynchronize(t##_start); \\\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start); \\\n// CHECK: hipEventRecord(t##_end); \\\n// CHECK: hipEventSynchronize(t##_end); \\\n// CHECK: hipEventElapsedTime(&t, t##_start, t##_end);\n#define TIMER_END(t) \\\n hipEventRecord(t##_start); \\\n hipEventSynchronize(t##_start); \\\n hipEventRecord(t##_end); \\\n hipEventSynchronize(t##_end); \\\n hipEventElapsedTime(&t, t##_start, t##_end); \n\n\ninline void cErrorCheck(const char *file, int line) {\n// CHECK: hipDeviceSynchronize();\n// CHECK: hipError_t err = hipGetLastError();\n// CHECK: if (err != hipSuccess) {\n// CHECK: printf(\"Error: %s\\n\", hipGetErrorString(err));\n hipDeviceSynchronize();\n hipError_t err = hipGetLastError();\n if (err != hipSuccess) {\n printf(\"Error: %s\\n\", hipGetErrorString(err));\n printf(\" @ %s: %d\\n\", file, line);\n exit(-1);\n }\n}\n\ninline dim3 K_GRID(int n, int threads = K_THREADS) {\n int blocks = (int)ceilf(sqrtf((float)n/threads));\n dim3 grid(blocks, blocks);\n return grid;\n}\n\ntypedef struct data {\n int n;\n float4 *r, *v, *f;\n} data;\n\ndata cpu, gpu;\n\n#define N 20\n\n__global__ void repulsion(data gpu);\n__global__ void integration(data gpu);\n\n\nint main() {\n printf(\"Cuda Test 1\\n\");\n\n int count = 0;\n // CHECK: hipGetDeviceCount(&count);\n hipGetDeviceCount(&count);\n printf(\" %d CUDA devices found\\n\", count);\n if(!count) {\n ::exit(EXIT_FAILURE);\n }\n // CHECK: hipFree(0);\n hipFree(0);\n\n cpu.n = N;\n\n cpu.r = (float4*)malloc(N * sizeof(float4));\n cpu.v = (float4*)malloc(N * sizeof(float4));\n cpu.f = (float4*)malloc(N * sizeof(float4));\n\n for(int i = 0; i < N; ++i) {\n cpu.v[i] = make_float4(0,0,0,0);\n cpu.r[i] = make_float4(RND(), RND(), RND(), 0);\n cpu.f[i] = make_float4(0,0.01,0,0);\n }\n\n gpu = cpu;\n // CHECK: hipMalloc(&gpu.r, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.v, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.f, N * sizeof(float4));\n hipMalloc(&gpu.r, N * sizeof(float4));\n hipMalloc(&gpu.v, N * sizeof(float4));\n hipMalloc(&gpu.f, N * sizeof(float4));\n // CHECK: hipMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n hipMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n hipMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n hipMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n\n ERRORCHECK();\n float rep;\n TIMER_CREATE(rep);\n TIMER_START(rep);\n // CHECK: hipLaunchKernelGGL(integration, dim3(K_GRID(cpu.n)), dim3(K_THREADS), 0, 0, gpu);\n integration <<< K_GRID(cpu.n), K_THREADS >>>(gpu);\n\n TIMER_END(rep);\n printf(\"Took: %f ms\\n\", rep);\n ERRORCHECK();\n // CHECK: hipMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n hipMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n hipMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n hipMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipHostFree(cpu.r);\n // CHECK: hipHostFree(cpu.v);\n // CHECK: hipHostFree(cpu.f);\n hipHostFree(cpu.r);\n hipHostFree(cpu.v);\n hipHostFree(cpu.f);\n // CHECK: hipFree(gpu.r);\n // CHECK: hipFree(gpu.v);\n // CHECK: hipFree(gpu.f);\n hipFree(gpu.r);\n hipFree(gpu.v);\n hipFree(gpu.f);\n // CHECK: hipDeviceReset();\n hipDeviceReset();\n\n printf(\"Results: \\n\");\n for(int i = 0; i < N; ++i) {\n printf(\"%f, %f, %f \\n\", cpu.r[i].x, cpu.r[i].y, cpu.r[i].z);\n }\n\n printf(\"Ready...\\n\");\n return 0;\n}\n\n__global__ void repulsion(data gpu) {\n int idx = K_INDEX();\n if(idx < N) {\n gpu.r[idx].x = 1;\n gpu.r[idx].y = 1;\n gpu.r[idx].z = 1;\n }\n}\n\n#define MULT4(v, s) v.x *= s; v.y *= s; v.z *= s; v.w *= s;\n#define ADD4(v1, v2) v1.x += v2.x; v1.y += v2.y; v1.z += v2.z; v1.w += v2.w;\n\n__global__ void integration(data gpu) {\n int i = K_INDEX();\n if(i < N) {\n MULT4(gpu.f[i], 0.01);\n MULT4(gpu.v[i], 0.01);\n ADD4(gpu.v[i], gpu.f[i]);\n ADD4(gpu.r[i], gpu.v[i]);\n gpu.f[i] = make_float4(0,0,0,0);\n }\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/compilation_database/compilation_database_before_13000/cd_intro_before_13000.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#include \n#include \n#include \n#include \"inc.h\"\n// CHECK: #include \n#include \n\n// CHECK: hipEvent_t t##_start, t##_end; \\\n// CHECK: hipEventCreate(&t##_start); \\\n// CHECK: hipEventCreate(&t##_end);\n#define TIMER_CREATE(t) \\\n cudaEvent_t t##_start, t##_end; \\\n cudaEventCreate(&t##_start); \\\n cudaEventCreate(&t##_end);\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start);\n#define TIMER_START(t) \\\n cudaEventRecord(t##_start); \\\n cudaEventSynchronize(t##_start); \\\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start); \\\n// CHECK: hipEventRecord(t##_end); \\\n// CHECK: hipEventSynchronize(t##_end); \\\n// CHECK: hipEventElapsedTime(&t, t##_start, t##_end);\n#define TIMER_END(t) \\\n cudaEventRecord(t##_start); \\\n cudaEventSynchronize(t##_start); \\\n cudaEventRecord(t##_end); \\\n cudaEventSynchronize(t##_end); \\\n cudaEventElapsedTime(&t, t##_start, t##_end); \n\n\ninline void cErrorCheck(const char *file, int line) {\n// CHECK: hipError_t err = hipGetLastError();\n// CHECK: if (err != hipSuccess) {\n// CHECK: printf(\"Error: %s\\n\", hipGetErrorString(err));\n cudaError_t err = cudaGetLastError();\n if (err != cudaSuccess) {\n printf(\"Error: %s\\n\", cudaGetErrorString(err));\n printf(\" @ %s: %d\\n\", file, line);\n exit(-1);\n }\n}\n\ninline dim3 K_GRID(int n, int threads = K_THREADS) {\n int blocks = (int)ceilf(sqrtf((float)n/threads));\n dim3 grid(blocks, blocks);\n return grid;\n}\n\ntypedef struct data {\n int n;\n float4 *r, *v, *f;\n} data;\n\ndata cpu, gpu;\n\n#define N 20\n\n__global__ void repulsion(data gpu);\n__global__ void integration(data gpu);\n\n\nint main() {\n printf(\"Cuda Test 1\\n\");\n\n int count = 0;\n // CHECK: hipGetDeviceCount(&count);\n cudaGetDeviceCount(&count);\n printf(\" %d CUDA devices found\\n\", count);\n if(!count) {\n ::exit(EXIT_FAILURE);\n }\n // CHECK: hipFree(0);\n cudaFree(0);\n\n cpu.n = N;\n\n cpu.r = (float4*)malloc(N * sizeof(float4));\n cpu.v = (float4*)malloc(N * sizeof(float4));\n cpu.f = (float4*)malloc(N * sizeof(float4));\n\n for(int i = 0; i < N; ++i) {\n cpu.v[i] = make_float4(0,0,0,0);\n cpu.r[i] = make_float4(RND(), RND(), RND(), 0);\n cpu.f[i] = make_float4(0,0.01,0,0);\n }\n\n gpu = cpu;\n // CHECK: hipMalloc(&gpu.r, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.v, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.f, N * sizeof(float4));\n cudaMalloc(&gpu.r, N * sizeof(float4));\n cudaMalloc(&gpu.v, N * sizeof(float4));\n cudaMalloc(&gpu.f, N * sizeof(float4));\n // CHECK: hipMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n cudaMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), cudaMemcpyHostToDevice);\n cudaMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), cudaMemcpyHostToDevice);\n cudaMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), cudaMemcpyHostToDevice);\n\n ERRORCHECK();\n float rep;\n TIMER_CREATE(rep);\n TIMER_START(rep);\n // CHECK: hipLaunchKernelGGL(integration, dim3(K_GRID(cpu.n)), dim3(K_THREADS), 0, 0, gpu);\n integration <<< K_GRID(cpu.n), K_THREADS >>>(gpu);\n\n TIMER_END(rep);\n printf(\"Took: %f ms\\n\", rep);\n ERRORCHECK();\n // CHECK: hipMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n cudaMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), cudaMemcpyDeviceToHost);\n cudaMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), cudaMemcpyDeviceToHost);\n cudaMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), cudaMemcpyDeviceToHost);\n // CHECK: hipHostFree(cpu.r);\n // CHECK: hipHostFree(cpu.v);\n // CHECK: hipHostFree(cpu.f);\n cudaFreeHost(cpu.r);\n cudaFreeHost(cpu.v);\n cudaFreeHost(cpu.f);\n // CHECK: hipFree(gpu.r);\n // CHECK: hipFree(gpu.v);\n // CHECK: hipFree(gpu.f);\n cudaFree(gpu.r);\n cudaFree(gpu.v);\n cudaFree(gpu.f);\n // CHECK: hipDeviceReset();\n cudaDeviceReset();\n\n printf(\"Results: \\n\");\n for(int i = 0; i < N; ++i) {\n printf(\"%f, %f, %f \\n\", cpu.r[i].x, cpu.r[i].y, cpu.r[i].z);\n }\n\n printf(\"Ready...\\n\");\n return 0;\n}\n\n__global__ void repulsion(data gpu) {\n int idx = K_INDEX();\n if(idx < N) {\n gpu.r[idx].x = 1;\n gpu.r[idx].y = 1;\n gpu.r[idx].z = 1;\n }\n}\n\n#define MULT4(v, s) v.x *= s; v.y *= s; v.z *= s; v.w *= s;\n#define ADD4(v1, v2) v1.x += v2.x; v1.y += v2.y; v1.z += v2.z; v1.w += v2.w;\n\n__global__ void integration(data gpu) {\n int i = K_INDEX();\n if(i < N) {\n MULT4(gpu.f[i], 0.01);\n MULT4(gpu.v[i], 0.01);\n ADD4(gpu.v[i], gpu.f[i]);\n ADD4(gpu.r[i], gpu.v[i]);\n gpu.f[i] = make_float4(0,0,0,0);\n }\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#include \n#include \n#include \n#include \"inc.h\"\n// CHECK: #include \n#include \n\n// CHECK: hipEvent_t t##_start, t##_end; \\\n// CHECK: hipEventCreate(&t##_start); \\\n// CHECK: hipEventCreate(&t##_end);\n#define TIMER_CREATE(t) \\\n hipEvent_t t##_start, t##_end; \\\n hipEventCreate(&t##_start); \\\n hipEventCreate(&t##_end);\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start);\n#define TIMER_START(t) \\\n hipEventRecord(t##_start); \\\n hipEventSynchronize(t##_start); \\\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start); \\\n// CHECK: hipEventRecord(t##_end); \\\n// CHECK: hipEventSynchronize(t##_end); \\\n// CHECK: hipEventElapsedTime(&t, t##_start, t##_end);\n#define TIMER_END(t) \\\n hipEventRecord(t##_start); \\\n hipEventSynchronize(t##_start); \\\n hipEventRecord(t##_end); \\\n hipEventSynchronize(t##_end); \\\n hipEventElapsedTime(&t, t##_start, t##_end); \n\n\ninline void cErrorCheck(const char *file, int line) {\n// CHECK: hipError_t err = hipGetLastError();\n// CHECK: if (err != hipSuccess) {\n// CHECK: printf(\"Error: %s\\n\", hipGetErrorString(err));\n hipError_t err = hipGetLastError();\n if (err != hipSuccess) {\n printf(\"Error: %s\\n\", hipGetErrorString(err));\n printf(\" @ %s: %d\\n\", file, line);\n exit(-1);\n }\n}\n\ninline dim3 K_GRID(int n, int threads = K_THREADS) {\n int blocks = (int)ceilf(sqrtf((float)n/threads));\n dim3 grid(blocks, blocks);\n return grid;\n}\n\ntypedef struct data {\n int n;\n float4 *r, *v, *f;\n} data;\n\ndata cpu, gpu;\n\n#define N 20\n\n__global__ void repulsion(data gpu);\n__global__ void integration(data gpu);\n\n\nint main() {\n printf(\"Cuda Test 1\\n\");\n\n int count = 0;\n // CHECK: hipGetDeviceCount(&count);\n hipGetDeviceCount(&count);\n printf(\" %d CUDA devices found\\n\", count);\n if(!count) {\n ::exit(EXIT_FAILURE);\n }\n // CHECK: hipFree(0);\n hipFree(0);\n\n cpu.n = N;\n\n cpu.r = (float4*)malloc(N * sizeof(float4));\n cpu.v = (float4*)malloc(N * sizeof(float4));\n cpu.f = (float4*)malloc(N * sizeof(float4));\n\n for(int i = 0; i < N; ++i) {\n cpu.v[i] = make_float4(0,0,0,0);\n cpu.r[i] = make_float4(RND(), RND(), RND(), 0);\n cpu.f[i] = make_float4(0,0.01,0,0);\n }\n\n gpu = cpu;\n // CHECK: hipMalloc(&gpu.r, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.v, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.f, N * sizeof(float4));\n hipMalloc(&gpu.r, N * sizeof(float4));\n hipMalloc(&gpu.v, N * sizeof(float4));\n hipMalloc(&gpu.f, N * sizeof(float4));\n // CHECK: hipMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n hipMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n hipMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n hipMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n\n ERRORCHECK();\n float rep;\n TIMER_CREATE(rep);\n TIMER_START(rep);\n // CHECK: hipLaunchKernelGGL(integration, dim3(K_GRID(cpu.n)), dim3(K_THREADS), 0, 0, gpu);\n integration <<< K_GRID(cpu.n), K_THREADS >>>(gpu);\n\n TIMER_END(rep);\n printf(\"Took: %f ms\\n\", rep);\n ERRORCHECK();\n // CHECK: hipMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n hipMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n hipMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n hipMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipHostFree(cpu.r);\n // CHECK: hipHostFree(cpu.v);\n // CHECK: hipHostFree(cpu.f);\n hipHostFree(cpu.r);\n hipHostFree(cpu.v);\n hipHostFree(cpu.f);\n // CHECK: hipFree(gpu.r);\n // CHECK: hipFree(gpu.v);\n // CHECK: hipFree(gpu.f);\n hipFree(gpu.r);\n hipFree(gpu.v);\n hipFree(gpu.f);\n // CHECK: hipDeviceReset();\n hipDeviceReset();\n\n printf(\"Results: \\n\");\n for(int i = 0; i < N; ++i) {\n printf(\"%f, %f, %f \\n\", cpu.r[i].x, cpu.r[i].y, cpu.r[i].z);\n }\n\n printf(\"Ready...\\n\");\n return 0;\n}\n\n__global__ void repulsion(data gpu) {\n int idx = K_INDEX();\n if(idx < N) {\n gpu.r[idx].x = 1;\n gpu.r[idx].y = 1;\n gpu.r[idx].z = 1;\n }\n}\n\n#define MULT4(v, s) v.x *= s; v.y *= s; v.z *= s; v.w *= s;\n#define ADD4(v1, v2) v1.x += v2.x; v1.y += v2.y; v1.z += v2.z; v1.w += v2.w;\n\n__global__ void integration(data gpu) {\n int i = K_INDEX();\n if(i < N) {\n MULT4(gpu.f[i], 0.01);\n MULT4(gpu.v[i], 0.01);\n ADD4(gpu.v[i], gpu.f[i]);\n ADD4(gpu.r[i], gpu.v[i]);\n gpu.f[i] = make_float4(0,0,0,0);\n }\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/compilation_database/compilation_database_13000/cd_intro.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// Synthetic test to warn only on device functions umin and umax as unsupported, but not on user defined ones.\n// ToDo: change lit testing in order to parse the output.\n\n#define LEN 1024\n#define SIZE LEN * sizeof(float)\n// CHECK: #include \n#include \n\nnamespace my {\n // user defined function\n unsigned int umin(unsigned int arg1, unsigned int arg2) {\n return (arg1 < arg2) ? arg1 : arg2;\n }\n // user defined function\n unsigned int umax(unsigned int arg1, unsigned int arg2) {\n return (arg1 > arg2) ? arg1 : arg2;\n }\n}\n\n__global__ void uint_arithm(float* A, float* B, float* C, unsigned int u1, unsigned int u2)\n{\n // device function call (warn if unsupported)\n unsigned int _umin = umin ( u1, u2 );\n // device function call (warn if unsupported)\n unsigned int _umax = umax ( u1, u2 );\n // device function call (warn if unsupported)\n unsigned int _umin_global = ::umin ( u1, u2 );\n // device function call (warn if unsupported)\n unsigned int _umax_global = ::umax(u1, u2);\n if (_umin != _umin_global) return;\n if (_umax != _umax_global) return;\n int i = threadIdx.x;\n A[i] = i + _umin;\n B[i] = i + _umax;\n C[i] = A[i] + B[i];\n}\n\nint main() {\n unsigned int u1 = 33;\n unsigned int u2 = 34;\n // user defined function call\n unsigned int _min = my::umin(u1, u2);\n // user defined function call\n unsigned int _max = my::umax(u1, u2);\n float *A, *B, *C;\n // CHECK: hipMalloc((void**)&A, SIZE);\n cudaMalloc((void**)&A, SIZE);\n // CHECK: hipMalloc((void**)&B, SIZE);\n cudaMalloc((void**)&B, SIZE);\n // CHECK: hipMalloc((void**)&C, SIZE);\n cudaMalloc((void**)&C, SIZE);\n dim3 dimGrid(LEN / 512, 1, 1);\n dim3 dimBlock(512, 1, 1);\n // CHECK: hipLaunchKernelGGL(uint_arithm, dim3(dimGrid), dim3(dimBlock), 0, 0, A, B, C, u1, u2);\n uint_arithm<<>>(A, B, C, u1, u2);\n return _min < _max;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// Synthetic test to warn only on device functions umin and umax as unsupported, but not on user defined ones.\n// ToDo: change lit testing in order to parse the output.\n\n#include \"hip/hip_runtime.h\"\n#define LEN 1024\n#define SIZE LEN * sizeof(float)\n// CHECK: #include \n#include \n\nnamespace my {\n // user defined function\n unsigned int umin(unsigned int arg1, unsigned int arg2) {\n return (arg1 < arg2) ? arg1 : arg2;\n }\n // user defined function\n unsigned int umax(unsigned int arg1, unsigned int arg2) {\n return (arg1 > arg2) ? arg1 : arg2;\n }\n}\n\n__global__ void uint_arithm(float* A, float* B, float* C, unsigned int u1, unsigned int u2)\n{\n // device function call (warn if unsupported)\n unsigned int _umin = umin ( u1, u2 );\n // device function call (warn if unsupported)\n unsigned int _umax = umax ( u1, u2 );\n // device function call (warn if unsupported)\n unsigned int _umin_global = ::umin ( u1, u2 );\n // device function call (warn if unsupported)\n unsigned int _umax_global = ::umax(u1, u2);\n if (_umin != _umin_global) return;\n if (_umax != _umax_global) return;\n int i = threadIdx.x;\n A[i] = i + _umin;\n B[i] = i + _umax;\n C[i] = A[i] + B[i];\n}\n\nint main() {\n unsigned int u1 = 33;\n unsigned int u2 = 34;\n // user defined function call\n unsigned int _min = my::umin(u1, u2);\n // user defined function call\n unsigned int _max = my::umax(u1, u2);\n float *A, *B, *C;\n // CHECK: hipMalloc((void**)&A, SIZE);\n hipMalloc((void**)&A, SIZE);\n // CHECK: hipMalloc((void**)&B, SIZE);\n hipMalloc((void**)&B, SIZE);\n // CHECK: hipMalloc((void**)&C, SIZE);\n hipMalloc((void**)&C, SIZE);\n dim3 dimGrid(LEN / 512, 1, 1);\n dim3 dimBlock(512, 1, 1);\n // CHECK: hipLaunchKernelGGL(uint_arithm, dim3(dimGrid), dim3(dimBlock), 0, 0, A, B, C, u1, u2);\n uint_arithm<<>>(A, B, C, u1, u2);\n return _min < _max;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/device/math_functions.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n/*\nCopyright (c) 2015-present Advanced Micro Devices, Inc. All rights reserved.\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in\nall copies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\nTHE SOFTWARE.\n*/\n\n// CHECK: #include \n#include \n#include \n#include \n\n#define NUM 1024\n#define SIZE 1024 * 4\n\n__device__ int globalIn[NUM];\n__device__ int globalOut[NUM];\n\n__global__ void Assign(int* Out) {\n int tid = threadIdx.x + blockIdx.x * blockDim.x;\n Out[tid] = globalIn[tid];\n globalOut[tid] = globalIn[tid];\n}\n\n__device__ __constant__ int globalConst[NUM];\n\n__global__ void checkAddress(int* addr, bool* out) {\n *out = (globalConst == addr);\n}\n\nint main() {\n int *A, *Am, *B, *Ad, *C, *Cm;\n A = new int[NUM];\n B = new int[NUM];\n C = new int[NUM];\n for (int i = 0; i < NUM; ++i) {\n A[i] = -1 * i;\n B[i] = 0;\n C[i] = 0;\n }\n // CHECK: hipMalloc((void**)&Ad, SIZE);\n cudaMalloc((void**)&Ad, SIZE);\n // CHECK: hipHostMalloc((void**)&Am, SIZE, hipHostMallocDefault);\n cudaMallocHost((void**)&Am, SIZE);\n // CHECK: hipHostMalloc((void**)&Cm, SIZE, hipHostMallocDefault);\n cudaMallocHost((void**)&Cm, SIZE);\n for (int i = 0; i < NUM; ++i) {\n Am[i] = -1 * i;\n Cm[i] = 0;\n }\n // CHECK: hipStream_t stream = NULL;\n cudaStream_t stream = NULL;\n // CHECK: hipStreamCreate(&stream);\n cudaStreamCreate(&stream);\n // CHECK: hipMemcpyToSymbolAsync(HIP_SYMBOL(globalIn), Am, SIZE, 0, hipMemcpyHostToDevice, stream);\n cudaMemcpyToSymbolAsync(globalIn, Am, SIZE, 0, cudaMemcpyHostToDevice, stream);\n // CHECK: hipStreamSynchronize(stream);\n cudaStreamSynchronize(stream);\n // CHECK: hipLaunchKernelGGL(Assign, dim3(1, 1, 1), dim3(NUM, 1, 1), 0, 0, Ad);\n Assign<<>>(Ad);\n // CHECK: hipMemcpy(B, Ad, SIZE, hipMemcpyDeviceToHost);\n cudaMemcpy(B, Ad, SIZE, cudaMemcpyDeviceToHost);\n // CHECK: hipMemcpyFromSymbolAsync(Cm, HIP_SYMBOL(globalOut), SIZE, 0, hipMemcpyDeviceToHost, stream);\n cudaMemcpyFromSymbolAsync(Cm, globalOut, SIZE, 0, cudaMemcpyDeviceToHost, stream);\n // CHECK: hipStreamSynchronize(stream);\n cudaStreamSynchronize(stream);\n for (int i = 0; i < NUM; ++i) {\n assert(Am[i] == B[i]);\n assert(Am[i] == Cm[i]);\n }\n for (int i = 0; i < NUM; ++i) {\n A[i] = -2 * i;\n B[i] = 0;\n }\n // CHECK: hipMemcpyToSymbol(HIP_SYMBOL(globalIn), A, SIZE, 0, hipMemcpyHostToDevice);\n cudaMemcpyToSymbol(globalIn, A, SIZE, 0, cudaMemcpyHostToDevice);\n // CHECK: hipLaunchKernelGGL(Assign, dim3(1, 1, 1), dim3(NUM, 1, 1), 0, 0, Ad);\n Assign<<>>(Ad);\n // CHECK: hipMemcpy(B, Ad, SIZE, hipMemcpyDeviceToHost);\n cudaMemcpy(B, Ad, SIZE, cudaMemcpyDeviceToHost);\n // CHECK: hipMemcpyFromSymbol(C, HIP_SYMBOL(globalOut), SIZE, 0, hipMemcpyDeviceToHost);\n cudaMemcpyFromSymbol(C, globalOut, SIZE, 0, cudaMemcpyDeviceToHost);\n for (int i = 0; i < NUM; ++i) {\n assert(A[i] == B[i]);\n assert(A[i] == C[i]);\n }\n for (int i = 0; i < NUM; ++i) {\n A[i] = -3 * i;\n B[i] = 0;\n }\n // CHECK: hipMemcpyToSymbolAsync(HIP_SYMBOL(globalIn), A, SIZE, 0, hipMemcpyHostToDevice, stream);\n cudaMemcpyToSymbolAsync(globalIn, A, SIZE, 0, cudaMemcpyHostToDevice, stream);\n // CHECK: hipStreamSynchronize(stream);\n cudaStreamSynchronize(stream);\n // CHECK: hipLaunchKernelGGL(Assign, dim3(1, 1, 1), dim3(NUM, 1, 1), 0, 0, Ad);\n Assign<<>>(Ad);\n // CHECK: hipMemcpy(B, Ad, SIZE, hipMemcpyDeviceToHost);\n cudaMemcpy(B, Ad, SIZE, cudaMemcpyDeviceToHost);\n // CHECK: hipMemcpyFromSymbolAsync(C, HIP_SYMBOL(globalOut), SIZE, 0, hipMemcpyDeviceToHost, stream);\n cudaMemcpyFromSymbolAsync(C, globalOut, SIZE, 0, cudaMemcpyDeviceToHost, stream);\n // CHECK: hipStreamSynchronize(stream);\n cudaStreamSynchronize(stream);\n for (int i = 0; i < NUM; ++i) {\n assert(A[i] == B[i]);\n assert(A[i] == C[i]);\n }\n bool *checkOkD;\n bool checkOk = false;\n size_t symbolSize = 0;\n int *symbolAddress;\n // CHECK: hipGetSymbolSize(&symbolSize, HIP_SYMBOL(globalConst));\n cudaGetSymbolSize(&symbolSize, globalConst);\n // CHECK: hipGetSymbolAddress((void**) &symbolAddress, HIP_SYMBOL(globalConst));\n cudaGetSymbolAddress((void**) &symbolAddress, globalConst);\n // CHECK: hipMalloc((void**)&checkOkD, sizeof(bool));\n cudaMalloc((void**)&checkOkD, sizeof(bool));\n // CHECK: hipLaunchKernelGGL(checkAddress, dim3(1, 1, 1), dim3(1, 1, 1), 0, 0, symbolAddress, checkOkD);\n checkAddress<<>>(symbolAddress, checkOkD);\n // CHECK: hipMemcpy(&checkOk, checkOkD, sizeof(bool), hipMemcpyDeviceToHost);\n cudaMemcpy(&checkOk, checkOkD, sizeof(bool), cudaMemcpyDeviceToHost);\n // CHECK: hipFree(checkOkD);\n cudaFree(checkOkD);\n assert(checkOk);\n assert(symbolSize == SIZE);\n // CHECK: hipHostFree(Am);\n cudaFreeHost(Am);\n // CHECK: hipHostFree(Cm);\n cudaFreeHost(Cm);\n // CHECK: hipFree(Ad);\n cudaFree(Ad);\n delete[] A;\n delete[] B;\n delete[] C;\n return 0;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n/*\nCopyright (c) 2015-present Advanced Micro Devices, Inc. All rights reserved.\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in\nall copies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\nTHE SOFTWARE.\n*/\n\n// CHECK: #include \n#include \n#include \n#include \n\n#define NUM 1024\n#define SIZE 1024 * 4\n\n__device__ int globalIn[NUM];\n__device__ int globalOut[NUM];\n\n__global__ void Assign(int* Out) {\n int tid = threadIdx.x + blockIdx.x * blockDim.x;\n Out[tid] = globalIn[tid];\n globalOut[tid] = globalIn[tid];\n}\n\n__device__ __constant__ int globalConst[NUM];\n\n__global__ void checkAddress(int* addr, bool* out) {\n *out = (globalConst == addr);\n}\n\nint main() {\n int *A, *Am, *B, *Ad, *C, *Cm;\n A = new int[NUM];\n B = new int[NUM];\n C = new int[NUM];\n for (int i = 0; i < NUM; ++i) {\n A[i] = -1 * i;\n B[i] = 0;\n C[i] = 0;\n }\n // CHECK: hipMalloc((void**)&Ad, SIZE);\n hipMalloc((void**)&Ad, SIZE);\n // CHECK: hipHostMalloc((void**)&Am, SIZE, hipHostMallocDefault);\n hipHostMalloc((void**)&Am, SIZE);\n // CHECK: hipHostMalloc((void**)&Cm, SIZE, hipHostMallocDefault);\n hipHostMalloc((void**)&Cm, SIZE);\n for (int i = 0; i < NUM; ++i) {\n Am[i] = -1 * i;\n Cm[i] = 0;\n }\n // CHECK: hipStream_t stream = NULL;\n hipStream_t stream = NULL;\n // CHECK: hipStreamCreate(&stream);\n hipStreamCreate(&stream);\n // CHECK: hipMemcpyToSymbolAsync(HIP_SYMBOL(globalIn), Am, SIZE, 0, hipMemcpyHostToDevice, stream);\n hipMemcpyToSymbolAsync(HIP_SYMBOL(globalIn), Am, SIZE, 0, hipMemcpyHostToDevice, stream);\n // CHECK: hipStreamSynchronize(stream);\n hipStreamSynchronize(stream);\n // CHECK: hipLaunchKernelGGL(Assign, dim3(1, 1, 1), dim3(NUM, 1, 1), 0, 0, Ad);\n Assign<<>>(Ad);\n // CHECK: hipMemcpy(B, Ad, SIZE, hipMemcpyDeviceToHost);\n hipMemcpy(B, Ad, SIZE, hipMemcpyDeviceToHost);\n // CHECK: hipMemcpyFromSymbolAsync(Cm, HIP_SYMBOL(globalOut), SIZE, 0, hipMemcpyDeviceToHost, stream);\n hipMemcpyFromSymbolAsync(Cm,HIP_SYMBOL(globalOut), SIZE, 0, hipMemcpyDeviceToHost, stream);\n // CHECK: hipStreamSynchronize(stream);\n hipStreamSynchronize(stream);\n for (int i = 0; i < NUM; ++i) {\n assert(Am[i] == B[i]);\n assert(Am[i] == Cm[i]);\n }\n for (int i = 0; i < NUM; ++i) {\n A[i] = -2 * i;\n B[i] = 0;\n }\n // CHECK: hipMemcpyToSymbol(HIP_SYMBOL(globalIn), A, SIZE, 0, hipMemcpyHostToDevice);\n hipMemcpyToSymbol(HIP_SYMBOL(globalIn), A, SIZE, 0, hipMemcpyHostToDevice);\n // CHECK: hipLaunchKernelGGL(Assign, dim3(1, 1, 1), dim3(NUM, 1, 1), 0, 0, Ad);\n Assign<<>>(Ad);\n // CHECK: hipMemcpy(B, Ad, SIZE, hipMemcpyDeviceToHost);\n hipMemcpy(B, Ad, SIZE, hipMemcpyDeviceToHost);\n // CHECK: hipMemcpyFromSymbol(C, HIP_SYMBOL(globalOut), SIZE, 0, hipMemcpyDeviceToHost);\n hipMemcpyFromSymbol(C,HIP_SYMBOL(globalOut), SIZE, 0, hipMemcpyDeviceToHost);\n for (int i = 0; i < NUM; ++i) {\n assert(A[i] == B[i]);\n assert(A[i] == C[i]);\n }\n for (int i = 0; i < NUM; ++i) {\n A[i] = -3 * i;\n B[i] = 0;\n }\n // CHECK: hipMemcpyToSymbolAsync(HIP_SYMBOL(globalIn), A, SIZE, 0, hipMemcpyHostToDevice, stream);\n hipMemcpyToSymbolAsync(HIP_SYMBOL(globalIn), A, SIZE, 0, hipMemcpyHostToDevice, stream);\n // CHECK: hipStreamSynchronize(stream);\n hipStreamSynchronize(stream);\n // CHECK: hipLaunchKernelGGL(Assign, dim3(1, 1, 1), dim3(NUM, 1, 1), 0, 0, Ad);\n Assign<<>>(Ad);\n // CHECK: hipMemcpy(B, Ad, SIZE, hipMemcpyDeviceToHost);\n hipMemcpy(B, Ad, SIZE, hipMemcpyDeviceToHost);\n // CHECK: hipMemcpyFromSymbolAsync(C, HIP_SYMBOL(globalOut), SIZE, 0, hipMemcpyDeviceToHost, stream);\n hipMemcpyFromSymbolAsync(C,HIP_SYMBOL(globalOut), SIZE, 0, hipMemcpyDeviceToHost, stream);\n // CHECK: hipStreamSynchronize(stream);\n hipStreamSynchronize(stream);\n for (int i = 0; i < NUM; ++i) {\n assert(A[i] == B[i]);\n assert(A[i] == C[i]);\n }\n bool *checkOkD;\n bool checkOk = false;\n size_t symbolSize = 0;\n int *symbolAddress;\n // CHECK: hipGetSymbolSize(&symbolSize, HIP_SYMBOL(globalConst));\n hipGetSymbolSize(&symbolSize,HIP_SYMBOL(globalConst));\n // CHECK: hipGetSymbolAddress((void**) &symbolAddress, HIP_SYMBOL(globalConst));\n hipGetSymbolAddress((void**) &symbolAddress,HIP_SYMBOL(globalConst));\n // CHECK: hipMalloc((void**)&checkOkD, sizeof(bool));\n hipMalloc((void**)&checkOkD, sizeof(bool));\n // CHECK: hipLaunchKernelGGL(checkAddress, dim3(1, 1, 1), dim3(1, 1, 1), 0, 0, symbolAddress, checkOkD);\n checkAddress<<>>(symbolAddress, checkOkD);\n // CHECK: hipMemcpy(&checkOk, checkOkD, sizeof(bool), hipMemcpyDeviceToHost);\n hipMemcpy(&checkOk, checkOkD, sizeof(bool), hipMemcpyDeviceToHost);\n // CHECK: hipFree(checkOkD);\n hipFree(checkOkD);\n assert(checkOk);\n assert(symbolSize == SIZE);\n // CHECK: hipHostFree(Am);\n hipHostFree(Am);\n // CHECK: hipHostFree(Cm);\n hipHostFree(Cm);\n // CHECK: hipFree(Ad);\n hipFree(Ad);\n delete[] A;\n delete[] B;\n delete[] C;\n return 0;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/device/device_symbols.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 2 --skip-excluded-preprocessor-conditional-blocks --hip-kernel-execution-syntax %clang_args \"-Xclang\" \"-fcuda-allow-variadic-functions\"\n\n/*\nCopyright (c) 2015-present Advanced Micro Devices, Inc. All rights reserved.\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\nThe above copyright notice and this permission notice shall be included in\nall copies or substantial portions of the Software.\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\nTHE SOFTWARE.\n*/\n\n// CHECK: #include \"hip/hip_runtime.h\"\n#include \"cuda_runtime.h\"\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#define EXIT_WAIVED 2\n\nconst char* sampleName = \"hipSimpleAtomicsTest\";\n\nusing namespace std;\n// Auto-Verification Code\nbool testResult = true;\n\nbool computeGoldBitwise(...) {\n return true;\n}\n\ntemplate{}>::type* = nullptr>\nbool computeGoldBitwise(T* gpuData, int len) {\n T val = 0xff;\n for (int i = 0; i < len; ++i) {\n // 9th element should be 1\n val &= (2 * i + 7);\n }\n if (val != gpuData[8]) {\n printf(\"atomicAnd failed\\n\");\n return false;\n }\n\n val = 0;\n for (int i = 0; i < len; ++i) {\n // 10th element should be 0xff\n val |= (1 << i);\n }\n if (val != gpuData[9]) {\n printf(\"atomicOr failed\\n\");\n return false;\n }\n\n val = 0xff;\n for (int i = 0; i < len; ++i) {\n // 11th element should be 0xff\n val ^= i;\n }\n if (val != gpuData[10]) {\n printf(\"atomicXor failed\\n\");\n return false;\n }\n\n return true;\n}\n\ntemplate\nbool computeGold(T* gpuData, int len) {\n T val = 0;\n for (int i = 0; i < len; ++i) {\n val += 10;\n }\n if (val != gpuData[0]) {\n printf(\"atomicAdd failed\\n\");\n return false;\n }\n\n val = 0;\n for (int i = 0; i < len; ++i) {\n val -= 10;\n }\n if (val != gpuData[1]) {\n printf(\"atomicSub failed\\n\");\n return false;\n }\n\n bool found = false;\n for (T i = 0; i < len; ++i) {\n // third element should be a member of [0, len)\n if (i == gpuData[2]) {\n found = true;\n break;\n }\n }\n if (!found) {\n printf(\"atomicExch failed\\n\");\n return false;\n }\n\n val = -(1 << 8);\n for (T i = 0; i < len; ++i) {\n // fourth element should be len-1\n val = max(val, i);\n }\n if (val != gpuData[3]) {\n printf(\"atomicMax failed\\n\");\n return false;\n }\n\n val = 1 << 8;\n for (T i = 0; i < len; ++i) {\n val = min(val, i);\n }\n if (val != gpuData[4]) {\n printf(\"atomicMin failed\\n\");\n return false;\n }\n\n int limit = 17;\n val = 0;\n for (int i = 0; i < len; ++i) {\n val = (val >= limit) ? 0 : val + 1;\n }\n if (val != gpuData[5]) {\n printf(\"atomicInc failed\\n\");\n return false;\n }\n\n limit = 137;\n val = 0;\n for (int i = 0; i < len; ++i) {\n val = ((val == 0) || (val > limit)) ? limit : val - 1;\n }\n if (val != gpuData[6]) {\n printf(\"atomicDec failed\\n\");\n return false;\n }\n\n found = false;\n for (T i = 0; i < len; ++i) {\n // eighth element should be a member of [0, len)\n if (i == gpuData[7]) {\n found = true;\n break;\n }\n }\n if (!found) {\n printf(\"atomicCAS failed\\n\");\n return false;\n }\n\n return computeGoldBitwise(gpuData, len);\n}\n\n__device__\nvoid testKernelExch(...) {}\n\ntemplate{}>::type* = nullptr>\n__device__\nvoid testKernelExch(T* g_odata) {\n // access thread id\n const T tid = blockDim.x * blockIdx.x + threadIdx.x;\n // Atomic exchange\n atomicExch(&g_odata[2], tid);\n}\n\n__device__\nvoid testKernelSub(...) {}\n\ntemplate<\n typename T,\n typename enable_if<\n is_same{} || is_same{}>::type* = nullptr>\n__device__\nvoid testKernelSub(T* g_odata) {\n // Atomic subtraction (final should be 0)\n atomicSub(&g_odata[1], 10);\n}\n\n__device__\nvoid testKernelIntegral(...) {}\n\ntemplate{}>::type* = nullptr>\n__device__\nvoid testKernelIntegral(T* g_odata) {\n // access thread id\n const T tid = blockDim.x * blockIdx.x + threadIdx.x;\n // Atomic maximum\n atomicMax(&g_odata[3], tid);\n // Atomic minimum\n atomicMin(&g_odata[4], tid);\n // Atomic increment (modulo 17+1)\n atomicInc((unsigned int*)&g_odata[5], 17);\n // Atomic decrement\n atomicDec((unsigned int*)&g_odata[6], 137);\n // Atomic compare-and-swap\n atomicCAS(&g_odata[7], tid - 1, tid);\n // Atomic AND\n atomicAnd(&g_odata[8], 2 * tid + 7);\n // Atomic OR\n atomicOr(&g_odata[9], 1 << tid);\n // Atomic XOR\n atomicXor(&g_odata[10], tid);\n testKernelSub(g_odata);\n}\n\ntemplate\n__global__ void testKernel(T* g_odata) {\n // Atomic addition\n atomicAdd(&g_odata[0], 10);\n testKernelIntegral(g_odata);\n testKernelExch(g_odata);\n}\n\ntemplate\nvoid runTest() {\n // CHECK: hipDeviceProp_t deviceProp;\n cudaDeviceProp deviceProp;\n deviceProp.major = 0;\n deviceProp.minor = 0;\n int dev = 0;\n // CHECK: hipGetDeviceProperties(&deviceProp, dev);\n cudaGetDeviceProperties(&deviceProp, dev);\n // Statistics about the GPU device\n printf(\n \"> GPU device has %d Multi-Processors, \"\n \"SM %d.%d compute capabilities\\n\\n\",\n deviceProp.multiProcessorCount, deviceProp.major, deviceProp.minor);\n unsigned int numThreads = 256;\n unsigned int numBlocks = 64;\n unsigned int numData = 11;\n unsigned int memSize = sizeof(T) * numData;\n\n // Allocate mem for the result on host side\n T* hOData = (T*)malloc(memSize);\n // Initialize the memory\n for (unsigned int i = 0; i < numData; i++) hOData[i] = 0;\n // To make the AND and XOR tests generate something other than 0...\n hOData[8] = hOData[10] = 0xff;\n // Allocate device memory for result\n T* dOData;\n // CHECK: hipMalloc((void**)&dOData, memSize);\n cudaMalloc((void**)&dOData, memSize);\n // Copy host memory to device to initialize to zero\n // CHECK: hipMemcpy(dOData, hOData, memSize, hipMemcpyHostToDevice);\n cudaMemcpy(dOData, hOData, memSize, cudaMemcpyHostToDevice);\n // Execute the kernel\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(testKernel), dim3(numBlocks), dim3(numThreads), 0, 0, dOData);\n testKernel<<>>(dOData);\n // Copy result from device to host\n // CHECK: hipMemcpy(hOData, dOData, memSize, hipMemcpyDeviceToHost);\n cudaMemcpy(hOData, dOData, memSize, cudaMemcpyDeviceToHost);\n // Compute reference solution\n testResult = computeGold(hOData, numThreads * numBlocks);\n // Cleanup memory\n free(hOData);\n // CHECK: hipFree(dOData);\n cudaFree(dOData);\n}\n\nint main(int argc, char** argv) {\n printf(\"%s starting...\\n\", sampleName);\n runTest();\n runTest();\n runTest();\n runTest();\n#if CUDA_VERSION >= 8000 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 600\n runTest();\n#endif\n // CHECK: hipDeviceReset();\n cudaDeviceReset();\n printf(\"%s completed, returned %s\\n\", sampleName, testResult ? \"OK\" : \"ERROR!\");\n exit(testResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 2 --skip-excluded-preprocessor-conditional-blocks --hip-kernel-execution-syntax %clang_args \"-Xclang\" \"-fcuda-allow-variadic-functions\"\n\n/*\nCopyright (c) 2015-present Advanced Micro Devices, Inc. All rights reserved.\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\nThe above copyright notice and this permission notice shall be included in\nall copies or substantial portions of the Software.\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\nTHE SOFTWARE.\n*/\n\n// CHECK: #include \"hip/hip_runtime.h\"\n#include \"hip/hip_runtime.h\"\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#define EXIT_WAIVED 2\n\nconst char* sampleName = \"hipSimpleAtomicsTest\";\n\nusing namespace std;\n// Auto-Verification Code\nbool testResult = true;\n\nbool computeGoldBitwise(...) {\n return true;\n}\n\ntemplate{}>::type* = nullptr>\nbool computeGoldBitwise(T* gpuData, int len) {\n T val = 0xff;\n for (int i = 0; i < len; ++i) {\n // 9th element should be 1\n val &= (2 * i + 7);\n }\n if (val != gpuData[8]) {\n printf(\"atomicAnd failed\\n\");\n return false;\n }\n\n val = 0;\n for (int i = 0; i < len; ++i) {\n // 10th element should be 0xff\n val |= (1 << i);\n }\n if (val != gpuData[9]) {\n printf(\"atomicOr failed\\n\");\n return false;\n }\n\n val = 0xff;\n for (int i = 0; i < len; ++i) {\n // 11th element should be 0xff\n val ^= i;\n }\n if (val != gpuData[10]) {\n printf(\"atomicXor failed\\n\");\n return false;\n }\n\n return true;\n}\n\ntemplate\nbool computeGold(T* gpuData, int len) {\n T val = 0;\n for (int i = 0; i < len; ++i) {\n val += 10;\n }\n if (val != gpuData[0]) {\n printf(\"atomicAdd failed\\n\");\n return false;\n }\n\n val = 0;\n for (int i = 0; i < len; ++i) {\n val -= 10;\n }\n if (val != gpuData[1]) {\n printf(\"atomicSub failed\\n\");\n return false;\n }\n\n bool found = false;\n for (T i = 0; i < len; ++i) {\n // third element should be a member of [0, len)\n if (i == gpuData[2]) {\n found = true;\n break;\n }\n }\n if (!found) {\n printf(\"atomicExch failed\\n\");\n return false;\n }\n\n val = -(1 << 8);\n for (T i = 0; i < len; ++i) {\n // fourth element should be len-1\n val = max(val, i);\n }\n if (val != gpuData[3]) {\n printf(\"atomicMax failed\\n\");\n return false;\n }\n\n val = 1 << 8;\n for (T i = 0; i < len; ++i) {\n val = min(val, i);\n }\n if (val != gpuData[4]) {\n printf(\"atomicMin failed\\n\");\n return false;\n }\n\n int limit = 17;\n val = 0;\n for (int i = 0; i < len; ++i) {\n val = (val >= limit) ? 0 : val + 1;\n }\n if (val != gpuData[5]) {\n printf(\"atomicInc failed\\n\");\n return false;\n }\n\n limit = 137;\n val = 0;\n for (int i = 0; i < len; ++i) {\n val = ((val == 0) || (val > limit)) ? limit : val - 1;\n }\n if (val != gpuData[6]) {\n printf(\"atomicDec failed\\n\");\n return false;\n }\n\n found = false;\n for (T i = 0; i < len; ++i) {\n // eighth element should be a member of [0, len)\n if (i == gpuData[7]) {\n found = true;\n break;\n }\n }\n if (!found) {\n printf(\"atomicCAS failed\\n\");\n return false;\n }\n\n return computeGoldBitwise(gpuData, len);\n}\n\n__device__\nvoid testKernelExch(...) {}\n\ntemplate{}>::type* = nullptr>\n__device__\nvoid testKernelExch(T* g_odata) {\n // access thread id\n const T tid = blockDim.x * blockIdx.x + threadIdx.x;\n // Atomic exchange\n atomicExch(&g_odata[2], tid);\n}\n\n__device__\nvoid testKernelSub(...) {}\n\ntemplate<\n typename T,\n typename enable_if<\n is_same{} || is_same{}>::type* = nullptr>\n__device__\nvoid testKernelSub(T* g_odata) {\n // Atomic subtraction (final should be 0)\n atomicSub(&g_odata[1], 10);\n}\n\n__device__\nvoid testKernelIntegral(...) {}\n\ntemplate{}>::type* = nullptr>\n__device__\nvoid testKernelIntegral(T* g_odata) {\n // access thread id\n const T tid = blockDim.x * blockIdx.x + threadIdx.x;\n // Atomic maximum\n atomicMax(&g_odata[3], tid);\n // Atomic minimum\n atomicMin(&g_odata[4], tid);\n // Atomic increment (modulo 17+1)\n atomicInc((unsigned int*)&g_odata[5], 17);\n // Atomic decrement\n atomicDec((unsigned int*)&g_odata[6], 137);\n // Atomic compare-and-swap\n atomicCAS(&g_odata[7], tid - 1, tid);\n // Atomic AND\n atomicAnd(&g_odata[8], 2 * tid + 7);\n // Atomic OR\n atomicOr(&g_odata[9], 1 << tid);\n // Atomic XOR\n atomicXor(&g_odata[10], tid);\n testKernelSub(g_odata);\n}\n\ntemplate\n__global__ void testKernel(T* g_odata) {\n // Atomic addition\n atomicAdd(&g_odata[0], 10);\n testKernelIntegral(g_odata);\n testKernelExch(g_odata);\n}\n\ntemplate\nvoid runTest() {\n // CHECK: hipDeviceProp_t deviceProp;\n hipDeviceProp_t deviceProp;\n deviceProp.major = 0;\n deviceProp.minor = 0;\n int dev = 0;\n // CHECK: hipGetDeviceProperties(&deviceProp, dev);\n hipGetDeviceProperties(&deviceProp, dev);\n // Statistics about the GPU device\n printf(\n \"> GPU device has %d Multi-Processors, \"\n \"SM %d.%d compute capabilities\\n\\n\",\n deviceProp.multiProcessorCount, deviceProp.major, deviceProp.minor);\n unsigned int numThreads = 256;\n unsigned int numBlocks = 64;\n unsigned int numData = 11;\n unsigned int memSize = sizeof(T) * numData;\n\n // Allocate mem for the result on host side\n T* hOData = (T*)malloc(memSize);\n // Initialize the memory\n for (unsigned int i = 0; i < numData; i++) hOData[i] = 0;\n // To make the AND and XOR tests generate something other than 0...\n hOData[8] = hOData[10] = 0xff;\n // Allocate device memory for result\n T* dOData;\n // CHECK: hipMalloc((void**)&dOData, memSize);\n hipMalloc((void**)&dOData, memSize);\n // Copy host memory to device to initialize to zero\n // CHECK: hipMemcpy(dOData, hOData, memSize, hipMemcpyHostToDevice);\n hipMemcpy(dOData, hOData, memSize, hipMemcpyHostToDevice);\n // Execute the kernel\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(testKernel), dim3(numBlocks), dim3(numThreads), 0, 0, dOData);\n testKernel<<>>(dOData);\n // Copy result from device to host\n // CHECK: hipMemcpy(hOData, dOData, memSize, hipMemcpyDeviceToHost);\n hipMemcpy(hOData, dOData, memSize, hipMemcpyDeviceToHost);\n // Compute reference solution\n testResult = computeGold(hOData, numThreads * numBlocks);\n // Cleanup memory\n free(hOData);\n // CHECK: hipFree(dOData);\n hipFree(dOData);\n}\n\nint main(int argc, char** argv) {\n printf(\"%s starting...\\n\", sampleName);\n runTest();\n runTest();\n runTest();\n runTest();\n#if CUDA_VERSION >= 8000 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 600\n runTest();\n#endif\n // CHECK: hipDeviceReset();\n hipDeviceReset();\n printf(\"%s completed, returned %s\\n\", sampleName, testResult ? \"OK\" : \"ERROR!\");\n exit(testResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/device/atomics.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 2 --amap --skip-excluded-preprocessor-conditional-blocks %clang_args\n\n// CHECK: #include \n// CHECK: #include \"hip/hip_complex.h\"\n#include \"cuComplex.h\"\n\n#define TYPEFLOAT\n#define DIMX 100\n#define DIMY 40\n#define moveX 2\n#define moveY 1\n\n#define MAXITERATIONS 10\n\n#ifdef TYPEFLOAT\n#define TYPE float\n// CHECK: #define cTYPE hipFloatComplex\n#define cTYPE cuFloatComplex\n// CHECK: #define cMakecuComplex(re,i) make_hipFloatComplex(re,i)\n#define cMakecuComplex(re,i) make_cuFloatComplex(re,i)\n#endif\n#ifdef TYPEDOUBLE\n// CHECK: #define TYPE hipDoubleComplex\n#define TYPE cuDoubleComplex\n// CHECK: #define cMakecuComplex(re,i) make_hipDoubleComplex(re,i)\n#define cMakecuComplex(re,i) make_cuDoubleComplex(re,i)\n#endif\n\n__device__ cTYPE juliaFunctor(cTYPE p, cTYPE c) {\n // CHECK: return hipCaddf(hipCmulf(p, p), c);\n return cuCaddf(cuCmulf(p, p), c);\n}\n\n__device__ cTYPE convertToComplex(int x, int y, float zoom) {\n TYPE jx = 1.5 * (x - DIMX / 2) / (0.5 * zoom * DIMX) + moveX;\n TYPE jy = (y - DIMY / 2) / (0.5 * zoom * DIMY) + moveY;\n return cMakecuComplex(jx, jy);\n}\n\n__device__ int evolveComplexPoint(cTYPE p, cTYPE c) {\n int it = 1;\n // CHECK: while (it <= MAXITERATIONS && hipCabsf(p) <= 4) {\n while (it <= MAXITERATIONS && cuCabsf(p) <= 4) {\n p = juliaFunctor(p, c);\n it++;\n }\n return it;\n}\n\n__global__ void computeJulia(int* data, cTYPE c, float zoom) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n int j = blockIdx.y * blockDim.y + threadIdx.y;\n\n if (i\n// CHECK: #include \"hip/hip_complex.h\"\n#include \"hip/hip_runtime.h\"\n#include \"hip/hip_complex.h\"\n\n#define TYPEFLOAT\n#define DIMX 100\n#define DIMY 40\n#define moveX 2\n#define moveY 1\n\n#define MAXITERATIONS 10\n\n#ifdef TYPEFLOAT\n#define TYPE float\n// CHECK: #define cTYPE hipFloatComplex\n#define cTYPE hipFloatComplex\n// CHECK: #define cMakecuComplex(re,i) make_hipFloatComplex(re,i)\n#define cMakecuComplex(re,i) make_hipFloatComplex(re,i)\n#endif\n#ifdef TYPEDOUBLE\n// CHECK: #define TYPE hipDoubleComplex\n#define TYPE hipDoubleComplex\n// CHECK: #define cMakecuComplex(re,i) make_hipDoubleComplex(re,i)\n#define cMakecuComplex(re,i) make_hipDoubleComplex(re,i)\n#endif\n\n__device__ cTYPE juliaFunctor(cTYPE p, cTYPE c) {\n // CHECK: return hipCaddf(hipCmulf(p, p), c);\n return hipCaddf(hipCmulf(p, p), c);\n}\n\n__device__ cTYPE convertToComplex(int x, int y, float zoom) {\n TYPE jx = 1.5 * (x - DIMX / 2) / (0.5 * zoom * DIMX) + moveX;\n TYPE jy = (y - DIMY / 2) / (0.5 * zoom * DIMY) + moveY;\n return cMakecuComplex(jx, jy);\n}\n\n__device__ int evolveComplexPoint(cTYPE p, cTYPE c) {\n int it = 1;\n // CHECK: while (it <= MAXITERATIONS && hipCabsf(p) <= 4) {\n while (it <= MAXITERATIONS && hipCabsf(p) <= 4) {\n p = juliaFunctor(p, c);\n it++;\n }\n return it;\n}\n\n__global__ void computeJulia(int* data, cTYPE c, float zoom) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n int j = blockIdx.y * blockDim.y + threadIdx.y;\n\n if (i\n#include \n// CHECK: #include \n#include \n// CHECK: #include \n#include \n// CHECK: #include \n#include \n\n// CHECK: #define CUDA_CALL(x) do { if((x) != hipSuccess) {\n#define CUDA_CALL(x) do { if((x) != cudaSuccess) { \\\n printf(\"Error at %s:%d\\n\",__FILE__,__LINE__); \\\n return EXIT_FAILURE;}} while(0)\n// CHECK: #define CURAND_CALL(x) do { if((x)!=HIPRAND_STATUS_SUCCESS) {\n#define CURAND_CALL(x) do { if((x)!=CURAND_STATUS_SUCCESS) { \\\n printf(\"Error at %s:%d\\n\",__FILE__,__LINE__);\\\n return EXIT_FAILURE;}} while(0)\n\n\n#define HOURS 16\n#define OPENING_HOUR 7\n#define CLOSING_HOUR (OPENING_HOUR + HOURS)\n\n#define access_2D(type, ptr, row, column, pitch)\\\n *((type*)((char*)ptr + (row) * pitch) + column)\n\nenum API_TYPE {\n HOST_API = 0,\n SIMPLE_DEVICE_API = 1,\n ROBUST_DEVICE_API = 2,\n};\n\n/* global variables */\nAPI_TYPE api;\nint report_break;\nint cashiers_load_h[HOURS];\n__constant__ int cashiers_load[HOURS];\n// CHECK: __global__ void setup_kernel(hiprandState *state)\n__global__ void setup_kernel(curandState *state)\n{\n int id = threadIdx.x + blockIdx.x * blockDim.x;\n /* Each thread gets same seed, a different sequence \n number, no offset */\n // CHECK: hiprand_init(1234, id, 0, &state[id]);\n curand_init(1234, id, 0, &state[id]);\n}\n\n__inline__ __device__\nvoid update_queue(int id, int min, unsigned int new_customers,\n unsigned int &queue_length,\n unsigned int *queue_lengths, size_t pitch)\n{\n int balance;\n balance = new_customers - 2 * cashiers_load[(min-1)/60];\n if (balance + (int)queue_length <= 0){\n queue_length = 0;\n }else{\n queue_length += balance;\n }\n /* Store results */\n access_2D(unsigned int, queue_lengths, min-1, id, pitch)\n = queue_length;\n}\n\n// CHECK: __global__ void simple_device_API_kernel(hiprandState *state,\n__global__ void simple_device_API_kernel(curandState *state,\n unsigned int *queue_lengths, size_t pitch)\n{\n int id = threadIdx.x + blockIdx.x * blockDim.x;\n unsigned int new_customers;\n unsigned int queue_length = 0;\n /* Copy state to local memory for efficiency */\n // CHECK: hiprandState localState = state[id];\n curandState localState = state[id];\n /* Simulate queue in time */\n for(int min = 1; min <= 60 * HOURS; min++) {\n /* Draw number of new customers depending on API */\n // CHECK: new_customers = hiprand_poisson(&localState,\n new_customers = curand_poisson(&localState,\n 4*(sin((float)min/100.0)+1));\n /* Update queue */\n update_queue(id, min, new_customers, queue_length,\n queue_lengths, pitch); \n }\n /* Copy state back to global memory */\n state[id] = localState;\n}\n\n\n__global__ void host_API_kernel(unsigned int *poisson_numbers,\n unsigned int *queue_lengths, size_t pitch)\n{\n int id = threadIdx.x + blockIdx.x * blockDim.x;\n unsigned int new_customers;\n unsigned int queue_length = 0;\n /* Simulate queue in time */\n for(int min = 1; min <= 60 * HOURS; min++) {\n /* Get random number from global memory */\n new_customers = poisson_numbers\n [blockDim.x * gridDim.x * (min -1) + id];\n /* Update queue */\n update_queue(id, min, new_customers, queue_length,\n queue_lengths, pitch);\n }\n}\n// CHECK: __global__ void robust_device_API_kernel(hiprandState *state,\n// CHECK: hiprandDiscreteDistribution_t poisson_1,\n// CHECK: hiprandDiscreteDistribution_t poisson_2,\n// CHECK: hiprandDiscreteDistribution_t poisson_3,\n__global__ void robust_device_API_kernel(curandState *state,\n curandDiscreteDistribution_t poisson_1,\n curandDiscreteDistribution_t poisson_2,\n curandDiscreteDistribution_t poisson_3,\n unsigned int *queue_lengths, size_t pitch)\n{\n int id = threadIdx.x + blockIdx.x * 64;\n unsigned int new_customers;\n unsigned int queue_length = 0;\n /* Copy state to local memory for efficiency */\n // CHECK: hiprandState localState = state[id];\n curandState localState = state[id];\n /* Simulate queue in time */\n /* first 3 hours */\n for(int min = 1; min <= 60 * 3; min++) {\n /* draw number of new customers depending on API */\n new_customers =\n // CHECK: hiprand_discrete(&localState, poisson_2);\n curand_discrete(&localState, poisson_2);\n /* Update queue */\n update_queue(id, min, new_customers, queue_length,\n queue_lengths, pitch);\n }\n /* second 3 hours */\n for(int min = 60 * 3 + 1; min <= 60 * 6; min++) {\n /* draw number of new customers depending on API */\n new_customers =\n // CHECK: hiprand_discrete(&localState, poisson_1);\n curand_discrete(&localState, poisson_1);\n /* Update queue */\n update_queue(id, min, new_customers, queue_length,\n queue_lengths, pitch); \n }\n /* after 6 hours */\n for(int min = 60 * 6 + 1; min <= 60 * HOURS; min++) {\n /* draw number of new customers depending on API */\n new_customers =\n // CHECK: hiprand_discrete(&localState, poisson_3);\n curand_discrete(&localState, poisson_3);\n /* Update queue */\n update_queue(id, min, new_customers, queue_length,\n queue_lengths, pitch);\n }\n /* Copy state back to global memory */\n state[id] = localState;\n}\n\n/* Set time intervals between reports */\nvoid report_settings()\n{\n do{\n printf(\"Set time intervals between queue reports\");\n printf(\"(in minutes > 0)\\n\");\n if (scanf(\"%d\", &report_break) == 0) continue;\n }while(report_break <= 0);\n}\n\n\n/* Set number of cashiers each hour */\nvoid add_cachiers(int *cashiers_load)\n{\n int i, min, max, begin, end;\n printf(\"Cashier serves 2 customers per minute...\\n\");\n for (i = 0; i < HOURS; i++){\n cashiers_load_h[i] = 0;\n }\n while (true){\n printf(\"Adding cashier...\\n\");\n min = OPENING_HOUR;\n max = CLOSING_HOUR-1;\n do{\n printf(\"Set hour that cahier comes (%d-%d)\",\n min, max);\n printf(\" [type 0 to finish adding cashiers]\\n\");\n if (scanf(\"%d\", &begin) == 0) continue;\n }while (begin > max || (begin < min && begin != 0));\n if (begin == 0) break;\n min = begin+1;\n max = CLOSING_HOUR;\n do{\n printf(\"Set hour that cahier leaves (%d-%d)\",\n min, max);\n printf(\" [type 0 to finish adding cashiers]\\n\");\n if (scanf(\"%d\", &end) == 0) continue;\n }while (end > max || (end < min && end != 0));\n if (end == 0) break;\n for (i = begin - OPENING_HOUR;\n i < end - OPENING_HOUR; i++){\n cashiers_load_h[i]++;\n }\n }\n for (i = OPENING_HOUR; i < CLOSING_HOUR; i++){\n printf(\"\\n%2d:00 - %2d:00 %d cashier\",\n i, i+1, cashiers_load_h[i-OPENING_HOUR]);\n if (cashiers_load[i-OPENING_HOUR] != 1) printf(\"s\");\n }\n printf(\"\\n\");\n}\n\n/* Set API type */\nAPI_TYPE set_API_type()\n{\n printf(\"Choose API type:\\n\");\n int choose;\n do{\n printf(\"type 1 for HOST API\\n\");\n printf(\"type 2 for SIMPLE DEVICE API\\n\");\n printf(\"type 3 for ROBUST DEVICE API\\n\");\n if (scanf(\"%d\", &choose) == 0) continue;\n }while( choose < 1 || choose > 3);\n switch(choose){\n case 1: return HOST_API;\n case 2: return SIMPLE_DEVICE_API;\n case 3: return ROBUST_DEVICE_API;\n default:\n fprintf(stderr, \"wrong API\\n\");\n return HOST_API;\n }\n}\n\nvoid settings()\n{\n add_cachiers(cashiers_load);\n // CHECK: hipMemcpyToSymbol(HIP_SYMBOL(\"cashiers_load\"), cashiers_load_h,\n // CHECK: HOURS * sizeof(int), 0, hipMemcpyHostToDevice);\n cudaMemcpyToSymbol(\"cashiers_load\", cashiers_load_h,\n HOURS * sizeof(int), 0, cudaMemcpyHostToDevice);\n report_settings();\n api = set_API_type();\n}\n\nvoid print_statistics(unsigned int *hostResults, size_t pitch)\n{\n int min, i, hour, minute;\n unsigned int sum;\n for(min = report_break; min <= 60 * HOURS;\n min += report_break) {\n sum = 0;\n for(i = 0; i < 64 * 64; i++) {\n sum += access_2D(unsigned int, hostResults,\n min-1, i, pitch);\n }\n hour = OPENING_HOUR + min/60;\n minute = min%60;\n printf(\"%2d:%02d # of waiting customers = %10.4g |\",\n hour, minute, (float)sum/(64.0 * 64.0));\n printf(\" # of cashiers = %d | \",\n cashiers_load_h[(min-1)/60]);\n printf(\"# of new customers/min ~= \");\n switch (api){\n case HOST_API:\n printf(\"%2.2f\\n\", 4.0);\n break;\n case SIMPLE_DEVICE_API:\n printf(\"%2.2f\\n\",\n 4*(sin((float)min/100.0)+1));\n break;\n case ROBUST_DEVICE_API:\n if (min <= 3 * 60){\n printf(\"%2.2f\\n\", 2.0);\n }else{\n if (min <= 6 * 60){\n printf(\"%2.2f\\n\", 1.0);\n }else{\n printf(\"%2.2f\\n\", 3.0);\n }\n }\n break;\n default:\n fprintf(stderr, \"Wrong API\\n\");\n }\n }\n}\n\nint main(int argc, char *argv[])\n{\n int n;\n size_t pitch;\n // CHECK: hiprandState *devStates;\n curandState *devStates;\n unsigned int *devResults, *hostResults;\n unsigned int *poisson_numbers_d;\n // CHECK: hiprandDiscreteDistribution_t poisson_1, poisson_2;\n // CHECK: hiprandDiscreteDistribution_t poisson_3;\n // CHECK: hiprandGenerator_t gen;\n curandDiscreteDistribution_t poisson_1, poisson_2;\n curandDiscreteDistribution_t poisson_3;\n curandGenerator_t gen;\n\n /* Setting cashiers, report and API */\n settings();\n\n /* Allocate space for results on device */\n // CHECK: CUDA_CALL(hipMallocPitch((void **)&devResults, &pitch,\n CUDA_CALL(cudaMallocPitch((void **)&devResults, &pitch,\n 64 * 64 * sizeof(unsigned int), 60 * HOURS));\n\n /* Allocate space for results on host */\n hostResults = (unsigned int *)calloc(pitch * 60 * HOURS,\n sizeof(unsigned int));\n\n /* Allocate space for prng states on device */\n // CHECK: CUDA_CALL(hipMalloc((void **)&devStates, 64 * 64 *\n // CHECK: sizeof(hiprandState)));\n CUDA_CALL(cudaMalloc((void **)&devStates, 64 * 64 * \n sizeof(curandState)));\n\n /* Setup prng states */\n if (api != HOST_API){\n // CHECK: hipLaunchKernelGGL(setup_kernel, dim3(64), dim3(64), 0, 0, devStates);\n setup_kernel<<<64, 64>>>(devStates);\n }\n /* Simulate queue */\n switch (api){\n case HOST_API:\n /* Create pseudo-random number generator */\n // CHECK: CURAND_CALL(hiprandCreateGenerator(&gen,\n // CHECK: HIPRAND_RNG_PSEUDO_DEFAULT));\n CURAND_CALL(curandCreateGenerator(&gen,\n CURAND_RNG_PSEUDO_DEFAULT));\n /* Set seed */\n // CHECK: CURAND_CALL(hiprandSetPseudoRandomGeneratorSeed(\n CURAND_CALL(curandSetPseudoRandomGeneratorSeed(\n gen, 1234ULL));\n /* compute n */\n n = 64 * 64 * HOURS * 60;\n /* Allocate n unsigned ints on device */\n // CHECK: CUDA_CALL(hipMalloc((void **)&poisson_numbers_d,\n CUDA_CALL(cudaMalloc((void **)&poisson_numbers_d,\n n * sizeof(unsigned int)));\n /* Generate n unsigned ints on device */\n // CHECK: CURAND_CALL(hiprandGeneratePoisson(gen,\n CURAND_CALL(curandGeneratePoisson(gen,\n poisson_numbers_d, n, 4.0));\n // CHECK: hipLaunchKernelGGL(host_API_kernel, dim3(64), dim3(64), 0, 0, poisson_numbers_d,\n host_API_kernel<<<64, 64>>>(poisson_numbers_d,\n devResults, pitch);\n /* Cleanup */\n // CHECK: CURAND_CALL(hiprandDestroyGenerator(gen));\n CURAND_CALL(curandDestroyGenerator(gen));\n break;\n case SIMPLE_DEVICE_API:\n // CHECK: hipLaunchKernelGGL(simple_device_API_kernel, dim3(64), dim3(64), 0, 0, devStates,\n simple_device_API_kernel<<<64, 64>>>(devStates,\n devResults, pitch);\n break;\n case ROBUST_DEVICE_API:\n /* Create histograms for Poisson(1) */\n // CHECK: CURAND_CALL(hiprandCreatePoissonDistribution(1.0,\n CURAND_CALL(curandCreatePoissonDistribution(1.0,\n &poisson_1));\n /* Create histograms for Poisson(2) */\n // CHECK: CURAND_CALL(hiprandCreatePoissonDistribution(2.0,\n CURAND_CALL(curandCreatePoissonDistribution(2.0,\n &poisson_2));\n /* Create histograms for Poisson(3) */\n // CHECK: CURAND_CALL(hiprandCreatePoissonDistribution(3.0,\n CURAND_CALL(curandCreatePoissonDistribution(3.0,\n &poisson_3));\n // CHECK: hipLaunchKernelGGL(robust_device_API_kernel, dim3(64), dim3(64), 0, 0, devStates,\n robust_device_API_kernel<<<64, 64>>>(devStates,\n poisson_1, poisson_2, poisson_3,\n devResults, pitch);\n /* Cleanup */\n // CHECK: CURAND_CALL(hiprandDestroyDistribution(poisson_1));\n // CHECK: CURAND_CALL(hiprandDestroyDistribution(poisson_2));\n // CHECK: CURAND_CALL(hiprandDestroyDistribution(poisson_3));\n CURAND_CALL(curandDestroyDistribution(poisson_1));\n CURAND_CALL(curandDestroyDistribution(poisson_2));\n CURAND_CALL(curandDestroyDistribution(poisson_3));\n break;\n default:\n fprintf(stderr, \"Wrong API\\n\");\n }\n /* Copy device memory to host */\n // CHECK: CUDA_CALL(hipMemcpy2D(hostResults, pitch, devResults,\n // CHECK: 60 * HOURS, hipMemcpyDeviceToHost));\n CUDA_CALL(cudaMemcpy2D(hostResults, pitch, devResults,\n pitch, 64 * 64 * sizeof(unsigned int),\n 60 * HOURS, cudaMemcpyDeviceToHost));\n /* Show result */\n print_statistics(hostResults, pitch);\n /* Cleanup */\n // CHECK: CUDA_CALL(hipFree(devStates));\n // CHECK: CUDA_CALL(hipFree(devResults));\n CUDA_CALL(cudaFree(devStates));\n CUDA_CALL(cudaFree(devResults));\n free(hostResults);\n return EXIT_SUCCESS;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n// Taken from: http://docs.nvidia.com/cuda/curand/device-api-overview.html#poisson-api-example\n/*\n * This program uses CURAND library for Poisson distribution\n * to simulate queues in store for 16 hours. It shows the\n * difference of using 3 different APIs:\n * - HOST API -arrival of customers is described by Poisson(4)\n * - SIMPLE DEVICE API -arrival of customers is described by\n * Poisson(4*(sin(x/100)+1)), where x is number of minutes\n * from store opening time.\n * - ROBUST DEVICE API -arrival of customers is described by:\n * - Poisson(2) for first 3 hours.\n * - Poisson(1) for second 3 hours.\n * - Poisson(3) after 6 hours. \n */\n\n#include \n#include \n// CHECK: #include \n#include \n// CHECK: #include \n#include \n// CHECK: #include \n#include \n\n// CHECK: #define CUDA_CALL(x) do { if((x) != hipSuccess) {\n#define CUDA_CALL(x) do { if((x) != hipSuccess) { \\\n printf(\"Error at %s:%d\\n\",__FILE__,__LINE__); \\\n return EXIT_FAILURE;}} while(0)\n// CHECK: #define CURAND_CALL(x) do { if((x)!=HIPRAND_STATUS_SUCCESS) {\n#define CURAND_CALL(x) do { if((x)!=HIPRAND_STATUS_SUCCESS) { \\\n printf(\"Error at %s:%d\\n\",__FILE__,__LINE__);\\\n return EXIT_FAILURE;}} while(0)\n\n\n#define HOURS 16\n#define OPENING_HOUR 7\n#define CLOSING_HOUR (OPENING_HOUR + HOURS)\n\n#define access_2D(type, ptr, row, column, pitch)\\\n *((type*)((char*)ptr + (row) * pitch) + column)\n\nenum API_TYPE {\n HOST_API = 0,\n SIMPLE_DEVICE_API = 1,\n ROBUST_DEVICE_API = 2,\n};\n\n/* global variables */\nAPI_TYPE api;\nint report_break;\nint cashiers_load_h[HOURS];\n__constant__ int cashiers_load[HOURS];\n// CHECK: __global__ void setup_kernel(hiprandState *state)\n__global__ void setup_kernel(hiprandState *state)\n{\n int id = threadIdx.x + blockIdx.x * blockDim.x;\n /* Each thread gets same seed, a different sequence \n number, no offset */\n // CHECK: hiprand_init(1234, id, 0, &state[id]);\n hiprand_init(1234, id, 0, &state[id]);\n}\n\n__inline__ __device__\nvoid update_queue(int id, int min, unsigned int new_customers,\n unsigned int &queue_length,\n unsigned int *queue_lengths, size_t pitch)\n{\n int balance;\n balance = new_customers - 2 * cashiers_load[(min-1)/60];\n if (balance + (int)queue_length <= 0){\n queue_length = 0;\n }else{\n queue_length += balance;\n }\n /* Store results */\n access_2D(unsigned int, queue_lengths, min-1, id, pitch)\n = queue_length;\n}\n\n// CHECK: __global__ void simple_device_API_kernel(hiprandState *state,\n__global__ void simple_device_API_kernel(hiprandState *state,\n unsigned int *queue_lengths, size_t pitch)\n{\n int id = threadIdx.x + blockIdx.x * blockDim.x;\n unsigned int new_customers;\n unsigned int queue_length = 0;\n /* Copy state to local memory for efficiency */\n // CHECK: hiprandState localState = state[id];\n hiprandState localState = state[id];\n /* Simulate queue in time */\n for(int min = 1; min <= 60 * HOURS; min++) {\n /* Draw number of new customers depending on API */\n // CHECK: new_customers = hiprand_poisson(&localState,\n new_customers = hiprand_poisson(&localState,\n 4*(sin((float)min/100.0)+1));\n /* Update queue */\n update_queue(id, min, new_customers, queue_length,\n queue_lengths, pitch); \n }\n /* Copy state back to global memory */\n state[id] = localState;\n}\n\n\n__global__ void host_API_kernel(unsigned int *poisson_numbers,\n unsigned int *queue_lengths, size_t pitch)\n{\n int id = threadIdx.x + blockIdx.x * blockDim.x;\n unsigned int new_customers;\n unsigned int queue_length = 0;\n /* Simulate queue in time */\n for(int min = 1; min <= 60 * HOURS; min++) {\n /* Get random number from global memory */\n new_customers = poisson_numbers\n [blockDim.x * gridDim.x * (min -1) + id];\n /* Update queue */\n update_queue(id, min, new_customers, queue_length,\n queue_lengths, pitch);\n }\n}\n// CHECK: __global__ void robust_device_API_kernel(hiprandState *state,\n// CHECK: hiprandDiscreteDistribution_t poisson_1,\n// CHECK: hiprandDiscreteDistribution_t poisson_2,\n// CHECK: hiprandDiscreteDistribution_t poisson_3,\n__global__ void robust_device_API_kernel(hiprandState *state,\n hiprandDiscreteDistribution_t poisson_1,\n hiprandDiscreteDistribution_t poisson_2,\n hiprandDiscreteDistribution_t poisson_3,\n unsigned int *queue_lengths, size_t pitch)\n{\n int id = threadIdx.x + blockIdx.x * 64;\n unsigned int new_customers;\n unsigned int queue_length = 0;\n /* Copy state to local memory for efficiency */\n // CHECK: hiprandState localState = state[id];\n hiprandState localState = state[id];\n /* Simulate queue in time */\n /* first 3 hours */\n for(int min = 1; min <= 60 * 3; min++) {\n /* draw number of new customers depending on API */\n new_customers =\n // CHECK: hiprand_discrete(&localState, poisson_2);\n hiprand_discrete(&localState, poisson_2);\n /* Update queue */\n update_queue(id, min, new_customers, queue_length,\n queue_lengths, pitch);\n }\n /* second 3 hours */\n for(int min = 60 * 3 + 1; min <= 60 * 6; min++) {\n /* draw number of new customers depending on API */\n new_customers =\n // CHECK: hiprand_discrete(&localState, poisson_1);\n hiprand_discrete(&localState, poisson_1);\n /* Update queue */\n update_queue(id, min, new_customers, queue_length,\n queue_lengths, pitch); \n }\n /* after 6 hours */\n for(int min = 60 * 6 + 1; min <= 60 * HOURS; min++) {\n /* draw number of new customers depending on API */\n new_customers =\n // CHECK: hiprand_discrete(&localState, poisson_3);\n hiprand_discrete(&localState, poisson_3);\n /* Update queue */\n update_queue(id, min, new_customers, queue_length,\n queue_lengths, pitch);\n }\n /* Copy state back to global memory */\n state[id] = localState;\n}\n\n/* Set time intervals between reports */\nvoid report_settings()\n{\n do{\n printf(\"Set time intervals between queue reports\");\n printf(\"(in minutes > 0)\\n\");\n if (scanf(\"%d\", &report_break) == 0) continue;\n }while(report_break <= 0);\n}\n\n\n/* Set number of cashiers each hour */\nvoid add_cachiers(int *cashiers_load)\n{\n int i, min, max, begin, end;\n printf(\"Cashier serves 2 customers per minute...\\n\");\n for (i = 0; i < HOURS; i++){\n cashiers_load_h[i] = 0;\n }\n while (true){\n printf(\"Adding cashier...\\n\");\n min = OPENING_HOUR;\n max = CLOSING_HOUR-1;\n do{\n printf(\"Set hour that cahier comes (%d-%d)\",\n min, max);\n printf(\" [type 0 to finish adding cashiers]\\n\");\n if (scanf(\"%d\", &begin) == 0) continue;\n }while (begin > max || (begin < min && begin != 0));\n if (begin == 0) break;\n min = begin+1;\n max = CLOSING_HOUR;\n do{\n printf(\"Set hour that cahier leaves (%d-%d)\",\n min, max);\n printf(\" [type 0 to finish adding cashiers]\\n\");\n if (scanf(\"%d\", &end) == 0) continue;\n }while (end > max || (end < min && end != 0));\n if (end == 0) break;\n for (i = begin - OPENING_HOUR;\n i < end - OPENING_HOUR; i++){\n cashiers_load_h[i]++;\n }\n }\n for (i = OPENING_HOUR; i < CLOSING_HOUR; i++){\n printf(\"\\n%2d:00 - %2d:00 %d cashier\",\n i, i+1, cashiers_load_h[i-OPENING_HOUR]);\n if (cashiers_load[i-OPENING_HOUR] != 1) printf(\"s\");\n }\n printf(\"\\n\");\n}\n\n/* Set API type */\nAPI_TYPE set_API_type()\n{\n printf(\"Choose API type:\\n\");\n int choose;\n do{\n printf(\"type 1 for HOST API\\n\");\n printf(\"type 2 for SIMPLE DEVICE API\\n\");\n printf(\"type 3 for ROBUST DEVICE API\\n\");\n if (scanf(\"%d\", &choose) == 0) continue;\n }while( choose < 1 || choose > 3);\n switch(choose){\n case 1: return HOST_API;\n case 2: return SIMPLE_DEVICE_API;\n case 3: return ROBUST_DEVICE_API;\n default:\n fprintf(stderr, \"wrong API\\n\");\n return HOST_API;\n }\n}\n\nvoid settings()\n{\n add_cachiers(cashiers_load);\n // CHECK: hipMemcpyToSymbol(HIP_SYMBOL(\"cashiers_load\"), cashiers_load_h,\n // CHECK: HOURS * sizeof(int), 0, hipMemcpyHostToDevice);\n hipMemcpyToSymbol(HIP_SYMBOL(\"cashiers_load\"), cashiers_load_h,\n HOURS * sizeof(int), 0, hipMemcpyHostToDevice);\n report_settings();\n api = set_API_type();\n}\n\nvoid print_statistics(unsigned int *hostResults, size_t pitch)\n{\n int min, i, hour, minute;\n unsigned int sum;\n for(min = report_break; min <= 60 * HOURS;\n min += report_break) {\n sum = 0;\n for(i = 0; i < 64 * 64; i++) {\n sum += access_2D(unsigned int, hostResults,\n min-1, i, pitch);\n }\n hour = OPENING_HOUR + min/60;\n minute = min%60;\n printf(\"%2d:%02d # of waiting customers = %10.4g |\",\n hour, minute, (float)sum/(64.0 * 64.0));\n printf(\" # of cashiers = %d | \",\n cashiers_load_h[(min-1)/60]);\n printf(\"# of new customers/min ~= \");\n switch (api){\n case HOST_API:\n printf(\"%2.2f\\n\", 4.0);\n break;\n case SIMPLE_DEVICE_API:\n printf(\"%2.2f\\n\",\n 4*(sin((float)min/100.0)+1));\n break;\n case ROBUST_DEVICE_API:\n if (min <= 3 * 60){\n printf(\"%2.2f\\n\", 2.0);\n }else{\n if (min <= 6 * 60){\n printf(\"%2.2f\\n\", 1.0);\n }else{\n printf(\"%2.2f\\n\", 3.0);\n }\n }\n break;\n default:\n fprintf(stderr, \"Wrong API\\n\");\n }\n }\n}\n\nint main(int argc, char *argv[])\n{\n int n;\n size_t pitch;\n // CHECK: hiprandState *devStates;\n hiprandState *devStates;\n unsigned int *devResults, *hostResults;\n unsigned int *poisson_numbers_d;\n // CHECK: hiprandDiscreteDistribution_t poisson_1, poisson_2;\n // CHECK: hiprandDiscreteDistribution_t poisson_3;\n // CHECK: hiprandGenerator_t gen;\n hiprandDiscreteDistribution_t poisson_1, poisson_2;\n hiprandDiscreteDistribution_t poisson_3;\n hiprandGenerator_t gen;\n\n /* Setting cashiers, report and API */\n settings();\n\n /* Allocate space for results on device */\n // CHECK: CUDA_CALL(hipMallocPitch((void **)&devResults, &pitch,\n CUDA_CALL(hipMallocPitch((void **)&devResults, &pitch,\n 64 * 64 * sizeof(unsigned int), 60 * HOURS));\n\n /* Allocate space for results on host */\n hostResults = (unsigned int *)calloc(pitch * 60 * HOURS,\n sizeof(unsigned int));\n\n /* Allocate space for prng states on device */\n // CHECK: CUDA_CALL(hipMalloc((void **)&devStates, 64 * 64 *\n // CHECK: sizeof(hiprandState)));\n CUDA_CALL(hipMalloc((void **)&devStates, 64 * 64 * \n sizeof(hiprandState)));\n\n /* Setup prng states */\n if (api != HOST_API){\n // CHECK: hipLaunchKernelGGL(setup_kernel, dim3(64), dim3(64), 0, 0, devStates);\n setup_kernel<<<64, 64>>>(devStates);\n }\n /* Simulate queue */\n switch (api){\n case HOST_API:\n /* Create pseudo-random number generator */\n // CHECK: CURAND_CALL(hiprandCreateGenerator(&gen,\n // CHECK: HIPRAND_RNG_PSEUDO_DEFAULT));\n CURAND_CALL(hiprandCreateGenerator(&gen,\n HIPRAND_RNG_PSEUDO_DEFAULT));\n /* Set seed */\n // CHECK: CURAND_CALL(hiprandSetPseudoRandomGeneratorSeed(\n CURAND_CALL(hiprandSetPseudoRandomGeneratorSeed(\n gen, 1234ULL));\n /* compute n */\n n = 64 * 64 * HOURS * 60;\n /* Allocate n unsigned ints on device */\n // CHECK: CUDA_CALL(hipMalloc((void **)&poisson_numbers_d,\n CUDA_CALL(hipMalloc((void **)&poisson_numbers_d,\n n * sizeof(unsigned int)));\n /* Generate n unsigned ints on device */\n // CHECK: CURAND_CALL(hiprandGeneratePoisson(gen,\n CURAND_CALL(hiprandGeneratePoisson(gen,\n poisson_numbers_d, n, 4.0));\n // CHECK: hipLaunchKernelGGL(host_API_kernel, dim3(64), dim3(64), 0, 0, poisson_numbers_d,\n host_API_kernel<<<64, 64>>>(poisson_numbers_d,\n devResults, pitch);\n /* Cleanup */\n // CHECK: CURAND_CALL(hiprandDestroyGenerator(gen));\n CURAND_CALL(hiprandDestroyGenerator(gen));\n break;\n case SIMPLE_DEVICE_API:\n // CHECK: hipLaunchKernelGGL(simple_device_API_kernel, dim3(64), dim3(64), 0, 0, devStates,\n simple_device_API_kernel<<<64, 64>>>(devStates,\n devResults, pitch);\n break;\n case ROBUST_DEVICE_API:\n /* Create histograms for Poisson(1) */\n // CHECK: CURAND_CALL(hiprandCreatePoissonDistribution(1.0,\n CURAND_CALL(hiprandCreatePoissonDistribution(1.0,\n &poisson_1));\n /* Create histograms for Poisson(2) */\n // CHECK: CURAND_CALL(hiprandCreatePoissonDistribution(2.0,\n CURAND_CALL(hiprandCreatePoissonDistribution(2.0,\n &poisson_2));\n /* Create histograms for Poisson(3) */\n // CHECK: CURAND_CALL(hiprandCreatePoissonDistribution(3.0,\n CURAND_CALL(hiprandCreatePoissonDistribution(3.0,\n &poisson_3));\n // CHECK: hipLaunchKernelGGL(robust_device_API_kernel, dim3(64), dim3(64), 0, 0, devStates,\n robust_device_API_kernel<<<64, 64>>>(devStates,\n poisson_1, poisson_2, poisson_3,\n devResults, pitch);\n /* Cleanup */\n // CHECK: CURAND_CALL(hiprandDestroyDistribution(poisson_1));\n // CHECK: CURAND_CALL(hiprandDestroyDistribution(poisson_2));\n // CHECK: CURAND_CALL(hiprandDestroyDistribution(poisson_3));\n CURAND_CALL(hiprandDestroyDistribution(poisson_1));\n CURAND_CALL(hiprandDestroyDistribution(poisson_2));\n CURAND_CALL(hiprandDestroyDistribution(poisson_3));\n break;\n default:\n fprintf(stderr, \"Wrong API\\n\");\n }\n /* Copy device memory to host */\n // CHECK: CUDA_CALL(hipMemcpy2D(hostResults, pitch, devResults,\n // CHECK: 60 * HOURS, hipMemcpyDeviceToHost));\n CUDA_CALL(hipMemcpy2D(hostResults, pitch, devResults,\n pitch, 64 * 64 * sizeof(unsigned int),\n 60 * HOURS, hipMemcpyDeviceToHost));\n /* Show result */\n print_statistics(hostResults, pitch);\n /* Cleanup */\n // CHECK: CUDA_CALL(hipFree(devStates));\n // CHECK: CUDA_CALL(hipFree(devResults));\n CUDA_CALL(hipFree(devStates));\n CUDA_CALL(hipFree(devResults));\n free(hostResults);\n return EXIT_SUCCESS;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/libraries/cuRAND/poisson_api_example.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n#include \n// CHECK: #include \n#include \n#define THRUST_NS_QUALIFIER ::thrust\n// CHECK: #include \n#include \n\n#include \n\n// TODO:\n// using namespace cub;\n\ntemplate \n__global__ void sort(const T* data_in, T* data_out){\n // CHECK: typedef ::hipcub::BlockRadixSort BlockRadixSortT;\n typedef ::cub::BlockRadixSort BlockRadixSortT;\n __shared__ typename BlockRadixSortT::TempStorage tmp_sort;\n double items[4];\n int i0 = 4 * (blockIdx.x * blockDim.x + threadIdx.x);\n for (int i = 0; i < 4; ++i){\n items[i] = data_in[i0 + i];\n }\n BlockRadixSortT(tmp_sort).Sort(items);\n for (int i = 0; i < 4; ++i){\n data_out[i0 + i] = items[i];\n }\n}\n\nint main(){\n double* d_gpu = NULL;\n double* result_gpu = NULL;\n double* data_sorted = new double[4096];\n // Allocate memory on the GPU\n // CHECK: hipMalloc(&d_gpu, 4096 * sizeof(double));\n cudaMalloc(&d_gpu, 4096 * sizeof(double));\n // CHECK: hipMalloc(&result_gpu, 4096 * sizeof(double));\n cudaMalloc(&result_gpu, 4096 * sizeof(double));\n // CHECK: hiprandGenerator_t gen;\n curandGenerator_t gen;\n // Create generator\n // CHECK: hiprandCreateGenerator(&gen, HIPRAND_RNG_PSEUDO_DEFAULT);\n curandCreateGenerator(&gen, CURAND_RNG_PSEUDO_DEFAULT);\n // Fill array with random numbers\n // CHECK: hiprandGenerateNormalDouble(gen, d_gpu, 4096, 0.0, 1.0);\n curandGenerateNormalDouble(gen, d_gpu, 4096, 0.0, 1.0);\n // Destroy generator\n // CHECK: hiprandDestroyGenerator(gen);\n curandDestroyGenerator(gen);\n // Sort data\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(sort), dim3(1), dim3(1024), 0, 0, d_gpu, result_gpu);\n sort<<<1, 1024>>>(d_gpu, result_gpu);\n // CHECK: hipMemcpy(data_sorted, result_gpu, 4096 * sizeof(double), hipMemcpyDeviceToHost);\n cudaMemcpy(data_sorted, result_gpu, 4096 * sizeof(double), cudaMemcpyDeviceToHost);\n // Write the sorted data to standard out\n for (int i = 0; i < 4096; ++i){\n std::cout << data_sorted[i] << \", \";\n }\n std::cout << std::endl;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n#include \"hip/hip_runtime.h\"\n#include \n// CHECK: #include \n#include \n#define THRUST_NS_QUALIFIER ::thrust\n// CHECK: #include \n#include \n\n#include \n\n// TODO:\n// using namespace cub;\n\ntemplate \n__global__ void sort(const T* data_in, T* data_out){\n // CHECK: typedef ::hipcub::BlockRadixSort BlockRadixSortT;\n typedef ::hipcub::BlockRadixSort BlockRadixSortT;\n __shared__ typename BlockRadixSortT::TempStorage tmp_sort;\n double items[4];\n int i0 = 4 * (blockIdx.x * blockDim.x + threadIdx.x);\n for (int i = 0; i < 4; ++i){\n items[i] = data_in[i0 + i];\n }\n BlockRadixSortT(tmp_sort).Sort(items);\n for (int i = 0; i < 4; ++i){\n data_out[i0 + i] = items[i];\n }\n}\n\nint main(){\n double* d_gpu = NULL;\n double* result_gpu = NULL;\n double* data_sorted = new double[4096];\n // Allocate memory on the GPU\n // CHECK: hipMalloc(&d_gpu, 4096 * sizeof(double));\n hipMalloc(&d_gpu, 4096 * sizeof(double));\n // CHECK: hipMalloc(&result_gpu, 4096 * sizeof(double));\n hipMalloc(&result_gpu, 4096 * sizeof(double));\n // CHECK: hiprandGenerator_t gen;\n hiprandGenerator_t gen;\n // Create generator\n // CHECK: hiprandCreateGenerator(&gen, HIPRAND_RNG_PSEUDO_DEFAULT);\n hiprandCreateGenerator(&gen, HIPRAND_RNG_PSEUDO_DEFAULT);\n // Fill array with random numbers\n // CHECK: hiprandGenerateNormalDouble(gen, d_gpu, 4096, 0.0, 1.0);\n hiprandGenerateNormalDouble(gen, d_gpu, 4096, 0.0, 1.0);\n // Destroy generator\n // CHECK: hiprandDestroyGenerator(gen);\n hiprandDestroyGenerator(gen);\n // Sort data\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(sort), dim3(1), dim3(1024), 0, 0, d_gpu, result_gpu);\n sort<<<1, 1024>>>(d_gpu, result_gpu);\n // CHECK: hipMemcpy(data_sorted, result_gpu, 4096 * sizeof(double), hipMemcpyDeviceToHost);\n hipMemcpy(data_sorted, result_gpu, 4096 * sizeof(double), hipMemcpyDeviceToHost);\n // Write the sorted data to standard out\n for (int i = 0; i < 4096; ++i){\n std::cout << data_sorted[i] << \", \";\n }\n std::cout << std::endl;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/libraries/CUB/cub_01.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args %clang_args\n// CHECK: #include \n#include \n#define THRUST_NS_QUALIFIER ::thrust\n// CHECK: #include \n#include \n\n// using namespace hipcub;\nusing namespace cub;\n\n// Simple CUDA kernel for computing tiled partial sums\ntemplate \n cub::BlockScanAlgorithm SCAN_ALGO>\n__global__ void ScanTilesKernel(int *d_in, int *d_out) {\n // Specialize collective types for problem context\n // CHECK: typedef ::hipcub::BlockLoad BlockLoadT;\n typedef ::cub::BlockLoad BlockLoadT;\n typedef BlockScan BlockScanT;\n // Allocate on-chip temporary storage\n __shared__ union {\n typename BlockLoadT::TempStorage load;\n typename BlockScanT::TempStorage reduce;\n } temp_storage;\n // Load data per thread\n int thread_data[ITEMS_PER_THREAD];\n int offset = blockIdx.x * (BLOCK_THREADS * ITEMS_PER_THREAD);\n BlockLoadT(temp_storage.load).Load(d_in + offset, offset);\n __syncthreads();\n // Compute the block-wide prefix sum\n BlockScanT(temp_storage).Sum(thread_data);\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args %clang_args\n// CHECK: #include \n#include \"hip/hip_runtime.h\"\n#include \n#define THRUST_NS_QUALIFIER ::thrust\n// CHECK: #include \n#include \n\n// using namespace hipcub;\nusing namespace hipcub;\n\n// Simple CUDA kernel for computing tiled partial sums\ntemplate \n hipcub::BlockScanAlgorithm SCAN_ALGO>\n__global__ void ScanTilesKernel(int *d_in, int *d_out) {\n // Specialize collective types for problem context\n // CHECK: typedef ::hipcub::BlockLoad BlockLoadT;\n typedef ::hipcub::BlockLoad BlockLoadT;\n typedef BlockScan BlockScanT;\n // Allocate on-chip temporary storage\n __shared__ union {\n typename BlockLoadT::TempStorage load;\n typename BlockScanT::TempStorage reduce;\n } temp_storage;\n // Load data per thread\n int thread_data[ITEMS_PER_THREAD];\n int offset = blockIdx.x * (BLOCK_THREADS * ITEMS_PER_THREAD);\n BlockLoadT(temp_storage.load).Load(d_in + offset, offset);\n __syncthreads();\n // Compute the block-wide prefix sum\n BlockScanT(temp_storage).Sum(thread_data);\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/libraries/CUB/cub_03.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n#include \n// CHECK: #include \n#include \n#define THRUST_NS_QUALIFIER ::thrust\n// CHECK: #include \n#include \n\n#include \n\ntemplate \n__global__ void sort(const T* data_in, T* data_out){\n // CHECK: typedef hipcub::BlockLoad BlockLoadT;\n typedef cub::BlockLoad BlockLoadT;\n // CHECK: typedef hipcub::BlockRadixSort BlockRadixSortT;\n typedef cub::BlockRadixSort BlockRadixSortT;\n // CHECK: typedef hipcub::BlockStore BlockStoreT;\n typedef cub::BlockStore BlockStoreT;\n __shared__ union {\n typename BlockLoadT::TempStorage load;\n typename BlockRadixSortT::TempStorage sort;\n typename BlockStoreT::TempStorage store;\n } tmp_storage;\n T items[ITEMS_PER_THREAD];\n BlockLoadT(tmp_storage.load).Load(data_in + blockIdx.x * BLOCK_WIDTH * ITEMS_PER_THREAD, items);\n __syncthreads();\n BlockRadixSortT(tmp_storage.sort).Sort(items);\n __syncthreads();\n BlockStoreT(tmp_storage.store).Store(data_out + blockIdx.x * BLOCK_WIDTH * ITEMS_PER_THREAD, items);\n}\n\nint main() {\n double* d_gpu = NULL;\n double* result_gpu = NULL;\n double* data_sorted = new double[1000*4096];\n // Allocate memory on the GPU\n // CHECK: hipMalloc(&d_gpu, 1000*4096 * sizeof(double));\n cudaMalloc(&d_gpu, 1000*4096 * sizeof(double));\n // CHECK: hipMalloc(&result_gpu, 1000*4096 * sizeof(double));\n cudaMalloc(&result_gpu, 1000*4096 * sizeof(double));\n // CHECK: hiprandGenerator_t gen;\n curandGenerator_t gen;\n // Create generator\n // CHECK: hiprandCreateGenerator(&gen, HIPRAND_RNG_PSEUDO_DEFAULT);\n curandCreateGenerator(&gen, CURAND_RNG_PSEUDO_DEFAULT);\n // Fill array with random numbers\n // CHECK: hiprandGenerateNormalDouble(gen, d_gpu, 1000*4096, 0.0, 1.0);\n curandGenerateNormalDouble(gen, d_gpu, 1000*4096, 0.0, 1.0);\n // Destroy generator\n // CHECK: hiprandDestroyGenerator(gen);\n curandDestroyGenerator(gen);\n // Sort data\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(sort<512, 8, hipcub::BLOCK_LOAD_TRANSPOSE, hipcub::BLOCK_STORE_TRANSPOSE>), dim3(1000), dim3(512), 0, 0, d_gpu, result_gpu);\n sort<512, 8, cub::BLOCK_LOAD_TRANSPOSE, cub::BLOCK_STORE_TRANSPOSE><<<1000, 512>>>(d_gpu, result_gpu);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(sort<256, 16, hipcub::BLOCK_LOAD_DIRECT, hipcub::BLOCK_STORE_DIRECT>), dim3(1000), dim3(256), 0, 0, d_gpu, result_gpu);\n sort<256, 16, cub::BLOCK_LOAD_DIRECT, cub::BLOCK_STORE_DIRECT><<<1000, 256>>>(d_gpu, result_gpu);\n // CHECK: hipMemcpy(data_sorted, result_gpu, 1000*4096*sizeof(double), hipMemcpyDeviceToHost);\n cudaMemcpy(data_sorted, result_gpu, 1000*4096*sizeof(double), cudaMemcpyDeviceToHost);\n // Write the sorted data to standard out\n for (int i = 0; i < 4095; ++i) {\n std::cout << data_sorted[i] << \", \";\n }\n std::cout << data_sorted[4095] << std::endl;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n#include \"hip/hip_runtime.h\"\n#include \n// CHECK: #include \n#include \n#define THRUST_NS_QUALIFIER ::thrust\n// CHECK: #include \n#include \n\n#include \n\ntemplate \n__global__ void sort(const T* data_in, T* data_out){\n // CHECK: typedef hipcub::BlockLoad BlockLoadT;\n typedef hipcub::BlockLoad BlockLoadT;\n // CHECK: typedef hipcub::BlockRadixSort BlockRadixSortT;\n typedef hipcub::BlockRadixSort BlockRadixSortT;\n // CHECK: typedef hipcub::BlockStore BlockStoreT;\n typedef hipcub::BlockStore BlockStoreT;\n __shared__ union {\n typename BlockLoadT::TempStorage load;\n typename BlockRadixSortT::TempStorage sort;\n typename BlockStoreT::TempStorage store;\n } tmp_storage;\n T items[ITEMS_PER_THREAD];\n BlockLoadT(tmp_storage.load).Load(data_in + blockIdx.x * BLOCK_WIDTH * ITEMS_PER_THREAD, items);\n __syncthreads();\n BlockRadixSortT(tmp_storage.sort).Sort(items);\n __syncthreads();\n BlockStoreT(tmp_storage.store).Store(data_out + blockIdx.x * BLOCK_WIDTH * ITEMS_PER_THREAD, items);\n}\n\nint main() {\n double* d_gpu = NULL;\n double* result_gpu = NULL;\n double* data_sorted = new double[1000*4096];\n // Allocate memory on the GPU\n // CHECK: hipMalloc(&d_gpu, 1000*4096 * sizeof(double));\n hipMalloc(&d_gpu, 1000*4096 * sizeof(double));\n // CHECK: hipMalloc(&result_gpu, 1000*4096 * sizeof(double));\n hipMalloc(&result_gpu, 1000*4096 * sizeof(double));\n // CHECK: hiprandGenerator_t gen;\n hiprandGenerator_t gen;\n // Create generator\n // CHECK: hiprandCreateGenerator(&gen, HIPRAND_RNG_PSEUDO_DEFAULT);\n hiprandCreateGenerator(&gen, HIPRAND_RNG_PSEUDO_DEFAULT);\n // Fill array with random numbers\n // CHECK: hiprandGenerateNormalDouble(gen, d_gpu, 1000*4096, 0.0, 1.0);\n hiprandGenerateNormalDouble(gen, d_gpu, 1000*4096, 0.0, 1.0);\n // Destroy generator\n // CHECK: hiprandDestroyGenerator(gen);\n hiprandDestroyGenerator(gen);\n // Sort data\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(sort<512, 8, hipcub::BLOCK_LOAD_TRANSPOSE, hipcub::BLOCK_STORE_TRANSPOSE>), dim3(1000), dim3(512), 0, 0, d_gpu, result_gpu);\n sort<512, 8, hipcub::BLOCK_LOAD_TRANSPOSE, hipcub::BLOCK_STORE_TRANSPOSE><<<1000, 512>>>(d_gpu, result_gpu);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(sort<256, 16, hipcub::BLOCK_LOAD_DIRECT, hipcub::BLOCK_STORE_DIRECT>), dim3(1000), dim3(256), 0, 0, d_gpu, result_gpu);\n sort<256, 16, hipcub::BLOCK_LOAD_DIRECT, hipcub::BLOCK_STORE_DIRECT><<<1000, 256>>>(d_gpu, result_gpu);\n // CHECK: hipMemcpy(data_sorted, result_gpu, 1000*4096*sizeof(double), hipMemcpyDeviceToHost);\n hipMemcpy(data_sorted, result_gpu, 1000*4096*sizeof(double), hipMemcpyDeviceToHost);\n // Write the sorted data to standard out\n for (int i = 0; i < 4095; ++i) {\n std::cout << data_sorted[i] << \", \";\n }\n std::cout << data_sorted[4095] << std::endl;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/libraries/CUB/cub_02.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 4 --amap --hip-kernel-execution-syntax --skip-excluded-preprocessor-conditional-blocks --hipdnn %clang_args\n\n#include \n#include \n#include \n#include \n\n// CHECK: #include \n#include \n// CHECK: #include \"hipDNN.h\"\n#include \"cudnn.h\"\n\n// CHECK: hipError_t err = (f); \\\n// CHECK: if (err != hipSuccess) { \\\n\n#define CUDA_CALL(f) { \\\n cudaError_t err = (f); \\\n if (err != cudaSuccess) { \\\n std::cout \\\n << \" Error occurred: \" << err << std::endl; \\\n std::exit(1); \\\n } \\\n}\n// CHECK: hipdnnStatus_t err = (f); \\\n// CHECK: if (err != HIPDNN_STATUS_SUCCESS) { \\\n\n#define CUDNN_CALL(f) { \\\n cudnnStatus_t err = (f); \\\n if (err != CUDNN_STATUS_SUCCESS) { \\\n std::cout \\\n << \" Error occurred: \" << err << std::endl; \\\n std::exit(1); \\\n } \\\n}\n\n__global__ void dev_const(float *px, float k) {\n int tid = threadIdx.x + blockIdx.x * blockDim.x;\n px[tid] = k;\n}\n\n__global__ void dev_iota(float *px) {\n int tid = threadIdx.x + blockIdx.x * blockDim.x;\n px[tid] = tid;\n}\n\nvoid print(const float *data, int n, int c, int h, int w) {\n std::vector buffer(1 << 20);\n // CHECK: CUDA_CALL(hipMemcpy(\n CUDA_CALL(cudaMemcpy(\n buffer.data(), data,\n n * c * h * w * sizeof(float),\n // CHECK: hipMemcpyDeviceToHost));\n cudaMemcpyDeviceToHost));\n int a = 0;\n for (int i = 0; i < n; ++i) {\n for (int j = 0; j < c; ++j) {\n std::cout << \"n=\" << i << \", c=\" << j << \":\" << std::endl;\n for (int k = 0; k < h; ++k) {\n for (int l = 0; l < w; ++l) {\n std::cout << std::setw(4) << std::right << buffer[a];\n ++a;\n }\n std::cout << std::endl;\n }\n }\n }\n std::cout << std::endl;\n}\n\nint main() {\n // CHECK: hipdnnHandle_t cudnn;\n cudnnHandle_t cudnn;\n // CHECK: CUDNN_CALL(hipdnnCreate(&cudnn));\n CUDNN_CALL(cudnnCreate(&cudnn));\n\n // input\n const int in_n = 1;\n const int in_c = 1;\n const int in_h = 5;\n const int in_w = 5;\n std::cout << \"in_n: \" << in_n << std::endl;\n std::cout << \"in_c: \" << in_c << std::endl;\n std::cout << \"in_h: \" << in_h << std::endl;\n std::cout << \"in_w: \" << in_w << std::endl;\n std::cout << std::endl;\n // CHECK: hipdnnTensorDescriptor_t in_desc;\n cudnnTensorDescriptor_t in_desc;\n // CHECK: CUDNN_CALL(hipdnnCreateTensorDescriptor(&in_desc));\n CUDNN_CALL(cudnnCreateTensorDescriptor(&in_desc));\n // CHECK: CUDNN_CALL(hipdnnSetTensor4dDescriptor(\n CUDNN_CALL(cudnnSetTensor4dDescriptor(\n // CHECK: in_desc, HIPDNN_TENSOR_NCHW, HIPDNN_DATA_FLOAT,\n in_desc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT,\n in_n, in_c, in_h, in_w));\n\n float *in_data;\n // CHECK: CUDA_CALL(hipMalloc(\n CUDA_CALL(cudaMalloc(\n &in_data, in_n * in_c * in_h * in_w * sizeof(float)));\n\n // filter\n const int filt_k = 1;\n const int filt_c = 1;\n const int filt_h = 2;\n const int filt_w = 2;\n std::cout << \"filt_k: \" << filt_k << std::endl;\n std::cout << \"filt_c: \" << filt_c << std::endl;\n std::cout << \"filt_h: \" << filt_h << std::endl;\n std::cout << \"filt_w: \" << filt_w << std::endl;\n std::cout << std::endl;\n\n // CHECK: hipdnnFilterDescriptor_t filt_desc;\n cudnnFilterDescriptor_t filt_desc;\n // CHECK: CUDNN_CALL(hipdnnCreateFilterDescriptor(&filt_desc));\n CUDNN_CALL(cudnnCreateFilterDescriptor(&filt_desc));\n // CHECK: CUDNN_CALL(hipdnnSetFilter4dDescriptor(\n CUDNN_CALL(cudnnSetFilter4dDescriptor(\n // CHECK: filt_desc, HIPDNN_DATA_FLOAT, HIPDNN_TENSOR_NCHW,\n filt_desc, CUDNN_DATA_FLOAT, CUDNN_TENSOR_NCHW,\n filt_k, filt_c, filt_h, filt_w));\n\n float *filt_data;\n // CUDA_CALL(hipMalloc(\n CUDA_CALL(cudaMalloc(\n &filt_data, filt_k * filt_c * filt_h * filt_w * sizeof(float)));\n\n // convolution\n const int pad_h = 1;\n const int pad_w = 1;\n const int str_h = 1;\n const int str_w = 1;\n const int dil_h = 1;\n const int dil_w = 1;\n std::cout << \"pad_h: \" << pad_h << std::endl;\n std::cout << \"pad_w: \" << pad_w << std::endl;\n std::cout << \"str_h: \" << str_h << std::endl;\n std::cout << \"str_w: \" << str_w << std::endl;\n std::cout << \"dil_h: \" << dil_h << std::endl;\n std::cout << \"dil_w: \" << dil_w << std::endl;\n std::cout << std::endl;\n\n // CHECK: hipdnnConvolutionDescriptor_t conv_desc;\n cudnnConvolutionDescriptor_t conv_desc;\n // CUDNN_CALL(hipdnnCreateConvolutionDescriptor(&conv_desc));\n CUDNN_CALL(cudnnCreateConvolutionDescriptor(&conv_desc));\n // CHECK: CUDNN_CALL(hipdnnSetConvolution2dDescriptor(\n CUDNN_CALL(cudnnSetConvolution2dDescriptor(\n conv_desc,\n pad_h, pad_w, str_h, str_w, dil_h, dil_w,\n // CHECK: HIPDNN_CONVOLUTION, HIPDNN_DATA_FLOAT));\n CUDNN_CONVOLUTION, CUDNN_DATA_FLOAT));\n\n // output\n int out_n;\n int out_c;\n int out_h;\n int out_w;\n\n // CHECK: CUDNN_CALL(hipdnnGetConvolution2dForwardOutputDim(\n CUDNN_CALL(cudnnGetConvolution2dForwardOutputDim(\n conv_desc, in_desc, filt_desc,\n &out_n, &out_c, &out_h, &out_w));\n\n std::cout << \"out_n: \" << out_n << std::endl;\n std::cout << \"out_c: \" << out_c << std::endl;\n std::cout << \"out_h: \" << out_h << std::endl;\n std::cout << \"out_w: \" << out_w << std::endl;\n std::cout << std::endl;\n // CHECK: hipdnnTensorDescriptor_t out_desc;\n cudnnTensorDescriptor_t out_desc;\n // CHECK: CUDNN_CALL(hipdnnCreateTensorDescriptor(&out_desc));\n CUDNN_CALL(cudnnCreateTensorDescriptor(&out_desc));\n // CHECK: CUDNN_CALL(hipdnnSetTensor4dDescriptor(\n CUDNN_CALL(cudnnSetTensor4dDescriptor(\n // CHECK: out_desc, HIPDNN_TENSOR_NCHW, HIPDNN_DATA_FLOAT,\n out_desc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT,\n out_n, out_c, out_h, out_w));\n\n\n cudnnDataType_t *dataType = nullptr;\n cudnnTensorFormat_t *tensorFormat = nullptr;\n int *p_filt_k = nullptr;\n int *p_filt_c = nullptr;\n int *p_filt_h = nullptr;\n int *p_filt_w = nullptr;\n\n // CHECK: CUDNN_CALL(hipdnnGetFilter4dDescriptor(\n CUDNN_CALL(cudnnGetFilter4dDescriptor(\n filt_desc, dataType, tensorFormat,\n p_filt_k, p_filt_c, p_filt_h, p_filt_w));\n\n float *out_data;\n // CHECK: CUDA_CALL(hipMalloc(\n CUDA_CALL(cudaMalloc(\n &out_data, out_n * out_c * out_h * out_w * sizeof(float)));\n\n // algorithm\n // CHECK: hipdnnConvolutionFwdAlgo_t algo;\n cudnnConvolutionFwdAlgo_t algo;\n#if CUDNN_VERSION < 8001\n // CHECK: CUDNN_CALL(hipdnnGetConvolutionForwardAlgorithm(\n CUDNN_CALL(cudnnGetConvolutionForwardAlgorithm(\n cudnn,\n in_desc, filt_desc, conv_desc, out_desc,\n // CHECK: HIPDNN_CONVOLUTION_FWD_PREFER_FASTEST, 0, &algo));\n CUDNN_CONVOLUTION_FWD_PREFER_FASTEST, 0, &algo));\n#endif\n\n std::cout << \"Convolution algorithm: \" << algo << std::endl;\n std::cout << std::endl;\n\n // workspace\n size_t ws_size;\n // CHECK: CUDNN_CALL(hipdnnGetConvolutionForwardWorkspaceSize(\n CUDNN_CALL(cudnnGetConvolutionForwardWorkspaceSize(\n cudnn, in_desc, filt_desc, conv_desc, out_desc, algo, &ws_size));\n\n float *ws_data;\n // CHECK: CUDA_CALL(hipMalloc(&ws_data, ws_size));\n CUDA_CALL(cudaMalloc(&ws_data, ws_size));\n\n std::cout << \"Workspace size: \" << ws_size << std::endl;\n std::cout << std::endl;\n\n // perform\n float alpha = 1.f;\n float beta = 0.f;\n // CHECK: hipLaunchKernelGGL(dev_iota, dim3(in_w * in_h), dim3(in_n * in_c), 0, 0, in_data);\n // CHECK: hipLaunchKernelGGL(dev_const, dim3(filt_w * filt_h), dim3(filt_k * filt_c), 0, 0, filt_data, 1.f);\n dev_iota<<>>(in_data);\n dev_const<<>>(filt_data, 1.f);\n // CHECK: CUDNN_CALL(hipdnnConvolutionForward(\n CUDNN_CALL(cudnnConvolutionForward(\n cudnn,\n &alpha, in_desc, in_data, filt_desc, filt_data,\n conv_desc, algo, ws_data, ws_size,\n &beta, out_desc, out_data));\n\n // results\n std::cout << \"in_data:\" << std::endl;\n print(in_data, in_n, in_c, in_h, in_w);\n \n std::cout << \"filt_data:\" << std::endl;\n print(filt_data, filt_k, filt_c, filt_h, filt_w);\n \n std::cout << \"out_data:\" << std::endl;\n print(out_data, out_n, out_c, out_h, out_w);\n\n // finalizing\n // CHECK: CUDA_CALL(hipFree(ws_data));\n CUDA_CALL(cudaFree(ws_data));\n // CHECK: CUDA_CALL(hipFree(out_data));\n CUDA_CALL(cudaFree(out_data));\n // CHECK: CUDNN_CALL(hipdnnDestroyTensorDescriptor(out_desc));\n CUDNN_CALL(cudnnDestroyTensorDescriptor(out_desc));\n // CHECK: CUDNN_CALL(hipdnnDestroyConvolutionDescriptor(conv_desc));\n CUDNN_CALL(cudnnDestroyConvolutionDescriptor(conv_desc));\n // CHECK: CUDA_CALL(hipFree(filt_data));\n CUDA_CALL(cudaFree(filt_data));\n // CHECK: CUDNN_CALL(hipdnnDestroyFilterDescriptor(filt_desc));\n CUDNN_CALL(cudnnDestroyFilterDescriptor(filt_desc));\n // CHECK: CUDA_CALL(hipFree(in_data));\n CUDA_CALL(cudaFree(in_data));\n // CHECK: CUDNN_CALL(hipdnnDestroyTensorDescriptor(in_desc));\n CUDNN_CALL(cudnnDestroyTensorDescriptor(in_desc));\n // CHECK: CUDNN_CALL(hipdnnDestroy(cudnn));\n CUDNN_CALL(cudnnDestroy(cudnn));\n return 0;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 4 --amap --hip-kernel-execution-syntax --skip-excluded-preprocessor-conditional-blocks --hipdnn %clang_args\n\n#include \n#include \n#include \n#include \n\n// CHECK: #include \n#include \n// CHECK: #include \"hipDNN.h\"\n#include \"hipDNN.h\"\n\n// CHECK: hipError_t err = (f); \\\n// CHECK: if (err != hipSuccess) { \\\n\n#define CUDA_CALL(f) { \\\n hipError_t err = (f); \\\n if (err != hipSuccess) { \\\n std::cout \\\n << \" Error occurred: \" << err << std::endl; \\\n std::exit(1); \\\n } \\\n}\n// CHECK: hipdnnStatus_t err = (f); \\\n// CHECK: if (err != HIPDNN_STATUS_SUCCESS) { \\\n\n#define CUDNN_CALL(f) { \\\n hipdnnStatus_t err = (f); \\\n if (err != HIPDNN_STATUS_SUCCESS) { \\\n std::cout \\\n << \" Error occurred: \" << err << std::endl; \\\n std::exit(1); \\\n } \\\n}\n\n__global__ void dev_const(float *px, float k) {\n int tid = threadIdx.x + blockIdx.x * blockDim.x;\n px[tid] = k;\n}\n\n__global__ void dev_iota(float *px) {\n int tid = threadIdx.x + blockIdx.x * blockDim.x;\n px[tid] = tid;\n}\n\nvoid print(const float *data, int n, int c, int h, int w) {\n std::vector buffer(1 << 20);\n // CHECK: CUDA_CALL(hipMemcpy(\n CUDA_CALL(hipMemcpy(\n buffer.data(), data,\n n * c * h * w * sizeof(float),\n // CHECK: hipMemcpyDeviceToHost));\n hipMemcpyDeviceToHost));\n int a = 0;\n for (int i = 0; i < n; ++i) {\n for (int j = 0; j < c; ++j) {\n std::cout << \"n=\" << i << \", c=\" << j << \":\" << std::endl;\n for (int k = 0; k < h; ++k) {\n for (int l = 0; l < w; ++l) {\n std::cout << std::setw(4) << std::right << buffer[a];\n ++a;\n }\n std::cout << std::endl;\n }\n }\n }\n std::cout << std::endl;\n}\n\nint main() {\n // CHECK: hipdnnHandle_t cudnn;\n hipdnnHandle_t cudnn;\n // CHECK: CUDNN_CALL(hipdnnCreate(&cudnn));\n CUDNN_CALL(hipdnnCreate(&cudnn));\n\n // input\n const int in_n = 1;\n const int in_c = 1;\n const int in_h = 5;\n const int in_w = 5;\n std::cout << \"in_n: \" << in_n << std::endl;\n std::cout << \"in_c: \" << in_c << std::endl;\n std::cout << \"in_h: \" << in_h << std::endl;\n std::cout << \"in_w: \" << in_w << std::endl;\n std::cout << std::endl;\n // CHECK: hipdnnTensorDescriptor_t in_desc;\n hipdnnTensorDescriptor_t in_desc;\n // CHECK: CUDNN_CALL(hipdnnCreateTensorDescriptor(&in_desc));\n CUDNN_CALL(hipdnnCreateTensorDescriptor(&in_desc));\n // CHECK: CUDNN_CALL(hipdnnSetTensor4dDescriptor(\n CUDNN_CALL(hipdnnSetTensor4dDescriptor(\n // CHECK: in_desc, HIPDNN_TENSOR_NCHW, HIPDNN_DATA_FLOAT,\n in_desc, HIPDNN_TENSOR_NCHW, HIPDNN_DATA_FLOAT,\n in_n, in_c, in_h, in_w));\n\n float *in_data;\n // CHECK: CUDA_CALL(hipMalloc(\n CUDA_CALL(hipMalloc(\n &in_data, in_n * in_c * in_h * in_w * sizeof(float)));\n\n // filter\n const int filt_k = 1;\n const int filt_c = 1;\n const int filt_h = 2;\n const int filt_w = 2;\n std::cout << \"filt_k: \" << filt_k << std::endl;\n std::cout << \"filt_c: \" << filt_c << std::endl;\n std::cout << \"filt_h: \" << filt_h << std::endl;\n std::cout << \"filt_w: \" << filt_w << std::endl;\n std::cout << std::endl;\n\n // CHECK: hipdnnFilterDescriptor_t filt_desc;\n hipdnnFilterDescriptor_t filt_desc;\n // CHECK: CUDNN_CALL(hipdnnCreateFilterDescriptor(&filt_desc));\n CUDNN_CALL(hipdnnCreateFilterDescriptor(&filt_desc));\n // CHECK: CUDNN_CALL(hipdnnSetFilter4dDescriptor(\n CUDNN_CALL(hipdnnSetFilter4dDescriptor(\n // CHECK: filt_desc, HIPDNN_DATA_FLOAT, HIPDNN_TENSOR_NCHW,\n filt_desc, HIPDNN_DATA_FLOAT, HIPDNN_TENSOR_NCHW,\n filt_k, filt_c, filt_h, filt_w));\n\n float *filt_data;\n // CUDA_CALL(hipMalloc(\n CUDA_CALL(hipMalloc(\n &filt_data, filt_k * filt_c * filt_h * filt_w * sizeof(float)));\n\n // convolution\n const int pad_h = 1;\n const int pad_w = 1;\n const int str_h = 1;\n const int str_w = 1;\n const int dil_h = 1;\n const int dil_w = 1;\n std::cout << \"pad_h: \" << pad_h << std::endl;\n std::cout << \"pad_w: \" << pad_w << std::endl;\n std::cout << \"str_h: \" << str_h << std::endl;\n std::cout << \"str_w: \" << str_w << std::endl;\n std::cout << \"dil_h: \" << dil_h << std::endl;\n std::cout << \"dil_w: \" << dil_w << std::endl;\n std::cout << std::endl;\n\n // CHECK: hipdnnConvolutionDescriptor_t conv_desc;\n hipdnnConvolutionDescriptor_t conv_desc;\n // CUDNN_CALL(hipdnnCreateConvolutionDescriptor(&conv_desc));\n CUDNN_CALL(hipdnnCreateConvolutionDescriptor(&conv_desc));\n // CHECK: CUDNN_CALL(hipdnnSetConvolution2dDescriptor(\n CUDNN_CALL(hipdnnSetConvolution2dDescriptor(\n conv_desc,\n pad_h, pad_w, str_h, str_w, dil_h, dil_w,\n // CHECK: HIPDNN_CONVOLUTION, HIPDNN_DATA_FLOAT));\n HIPDNN_CONVOLUTION, HIPDNN_DATA_FLOAT));\n\n // output\n int out_n;\n int out_c;\n int out_h;\n int out_w;\n\n // CHECK: CUDNN_CALL(hipdnnGetConvolution2dForwardOutputDim(\n CUDNN_CALL(hipdnnGetConvolution2dForwardOutputDim(\n conv_desc, in_desc, filt_desc,\n &out_n, &out_c, &out_h, &out_w));\n\n std::cout << \"out_n: \" << out_n << std::endl;\n std::cout << \"out_c: \" << out_c << std::endl;\n std::cout << \"out_h: \" << out_h << std::endl;\n std::cout << \"out_w: \" << out_w << std::endl;\n std::cout << std::endl;\n // CHECK: hipdnnTensorDescriptor_t out_desc;\n hipdnnTensorDescriptor_t out_desc;\n // CHECK: CUDNN_CALL(hipdnnCreateTensorDescriptor(&out_desc));\n CUDNN_CALL(hipdnnCreateTensorDescriptor(&out_desc));\n // CHECK: CUDNN_CALL(hipdnnSetTensor4dDescriptor(\n CUDNN_CALL(hipdnnSetTensor4dDescriptor(\n // CHECK: out_desc, HIPDNN_TENSOR_NCHW, HIPDNN_DATA_FLOAT,\n out_desc, HIPDNN_TENSOR_NCHW, HIPDNN_DATA_FLOAT,\n out_n, out_c, out_h, out_w));\n\n\n hipdnnDataType_t *dataType = nullptr;\n hipdnnTensorFormat_t *tensorFormat = nullptr;\n int *p_filt_k = nullptr;\n int *p_filt_c = nullptr;\n int *p_filt_h = nullptr;\n int *p_filt_w = nullptr;\n\n // CHECK: CUDNN_CALL(hipdnnGetFilter4dDescriptor(\n CUDNN_CALL(hipdnnGetFilter4dDescriptor(\n filt_desc, dataType, tensorFormat,\n p_filt_k, p_filt_c, p_filt_h, p_filt_w));\n\n float *out_data;\n // CHECK: CUDA_CALL(hipMalloc(\n CUDA_CALL(hipMalloc(\n &out_data, out_n * out_c * out_h * out_w * sizeof(float)));\n\n // algorithm\n // CHECK: hipdnnConvolutionFwdAlgo_t algo;\n hipdnnConvolutionFwdAlgo_t algo;\n#if CUDNN_VERSION < 8001\n // CHECK: CUDNN_CALL(hipdnnGetConvolutionForwardAlgorithm(\n CUDNN_CALL(hipdnnGetConvolutionForwardAlgorithm(\n cudnn,\n in_desc, filt_desc, conv_desc, out_desc,\n // CHECK: HIPDNN_CONVOLUTION_FWD_PREFER_FASTEST, 0, &algo));\n HIPDNN_CONVOLUTION_FWD_PREFER_FASTEST, 0, &algo));\n#endif\n\n std::cout << \"Convolution algorithm: \" << algo << std::endl;\n std::cout << std::endl;\n\n // workspace\n size_t ws_size;\n // CHECK: CUDNN_CALL(hipdnnGetConvolutionForwardWorkspaceSize(\n CUDNN_CALL(hipdnnGetConvolutionForwardWorkspaceSize(\n cudnn, in_desc, filt_desc, conv_desc, out_desc, algo, &ws_size));\n\n float *ws_data;\n // CHECK: CUDA_CALL(hipMalloc(&ws_data, ws_size));\n CUDA_CALL(hipMalloc(&ws_data, ws_size));\n\n std::cout << \"Workspace size: \" << ws_size << std::endl;\n std::cout << std::endl;\n\n // perform\n float alpha = 1.f;\n float beta = 0.f;\n // CHECK: hipLaunchKernelGGL(dev_iota, dim3(in_w * in_h), dim3(in_n * in_c), 0, 0, in_data);\n // CHECK: hipLaunchKernelGGL(dev_const, dim3(filt_w * filt_h), dim3(filt_k * filt_c), 0, 0, filt_data, 1.f);\n dev_iota<<>>(in_data);\n dev_const<<>>(filt_data, 1.f);\n // CHECK: CUDNN_CALL(hipdnnConvolutionForward(\n CUDNN_CALL(hipdnnConvolutionForward(\n cudnn,\n &alpha, in_desc, in_data, filt_desc, filt_data,\n conv_desc, algo, ws_data, ws_size,\n &beta, out_desc, out_data));\n\n // results\n std::cout << \"in_data:\" << std::endl;\n print(in_data, in_n, in_c, in_h, in_w);\n \n std::cout << \"filt_data:\" << std::endl;\n print(filt_data, filt_k, filt_c, filt_h, filt_w);\n \n std::cout << \"out_data:\" << std::endl;\n print(out_data, out_n, out_c, out_h, out_w);\n\n // finalizing\n // CHECK: CUDA_CALL(hipFree(ws_data));\n CUDA_CALL(hipFree(ws_data));\n // CHECK: CUDA_CALL(hipFree(out_data));\n CUDA_CALL(hipFree(out_data));\n // CHECK: CUDNN_CALL(hipdnnDestroyTensorDescriptor(out_desc));\n CUDNN_CALL(hipdnnDestroyTensorDescriptor(out_desc));\n // CHECK: CUDNN_CALL(hipdnnDestroyConvolutionDescriptor(conv_desc));\n CUDNN_CALL(hipdnnDestroyConvolutionDescriptor(conv_desc));\n // CHECK: CUDA_CALL(hipFree(filt_data));\n CUDA_CALL(hipFree(filt_data));\n // CHECK: CUDNN_CALL(hipdnnDestroyFilterDescriptor(filt_desc));\n CUDNN_CALL(hipdnnDestroyFilterDescriptor(filt_desc));\n // CHECK: CUDA_CALL(hipFree(in_data));\n CUDA_CALL(hipFree(in_data));\n // CHECK: CUDNN_CALL(hipdnnDestroyTensorDescriptor(in_desc));\n CUDNN_CALL(hipdnnDestroyTensorDescriptor(in_desc));\n // CHECK: CUDNN_CALL(hipdnnDestroy(cudnn));\n CUDNN_CALL(hipdnnDestroy(cudnn));\n return 0;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/libraries/cuDNN/cudnn_convolution_forward.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args %clang_args\n\n/*\nCopyright (c) 2015-present Advanced Micro Devices, Inc. All rights reserved.\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in\nall copies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\nTHE SOFTWARE.\n*/\n\n// CHECK: #include \n#include \n#include \n\n// Random predefiend 32 and 64 bit values\nconstexpr uint32_t value32 = 0x70F0F0FF;\nconstexpr uint64_t value64 = 0x7FFF0000FFFF0000;\nconstexpr unsigned int writeFlag = 0;\n\n__global__\nvoid fn(float* px, float* py) {\n bool a[42];\n __shared__ double b[69];\n for (auto&& x : b) x = *py++;\n for (auto&& x : a) x = *px++ > 0.0;\n for (auto&& x : a) if (x) *--py = *--px;\n}\n\nvoid testWrite() {\n int64_t* signalPtr;\n // CHECK: hipStream_t stream;\n cudaStream_t stream;\n // CHECK: hipStreamCreate(&stream);\n cudaStreamCreate(&stream);\n int64_t* host_ptr64 = (int64_t*)malloc(sizeof(int64_t));\n int32_t* host_ptr32 = (int32_t*)malloc(sizeof(int32_t));\n // hipExtMallocWithFlags((void**)&signalPtr, 8, hipMallocSignalMemory);\n void* device_ptr64;\n void* device_ptr32;\n *host_ptr64 = 0x0;\n *host_ptr32 = 0x0;\n *signalPtr = 0x0;\n // CHECK: hipHostRegister(host_ptr64, sizeof(int64_t), 0);\n cudaHostRegister(host_ptr64, sizeof(int64_t), 0);\n // CHECK: hipHostRegister(host_ptr32, sizeof(int32_t), 0);\n cudaHostRegister(host_ptr32, sizeof(int32_t), 0);\n // CHECK: hipStreamWriteValue64(stream, hipDeviceptr_t(host_ptr64), value64, writeFlag);\n cuStreamWriteValue64(stream, CUdeviceptr(host_ptr64), value64, writeFlag);\n // CHECK: hipStreamWriteValue32(stream, hipDeviceptr_t(host_ptr32), value32, writeFlag);\n cuStreamWriteValue32(stream, CUdeviceptr(host_ptr32), value32, writeFlag);\n // CHECK: hipStreamSynchronize(stream);\n cudaStreamSynchronize(stream);\n // CHECK: hipHostGetDevicePointer((void**)&device_ptr64, host_ptr64, 0);\n cudaHostGetDevicePointer((void**)&device_ptr64, host_ptr64, 0);\n // CHECK: hipHostGetDevicePointer((void**)&device_ptr32, host_ptr32, 0);\n cudaHostGetDevicePointer((void**)&device_ptr32, host_ptr32, 0);\n // Reset values\n *host_ptr64 = 0x0;\n *host_ptr32 = 0x0;\n // CHECK: hipStreamWriteValue64(stream, hipDeviceptr_t(device_ptr64), value64, writeFlag);\n cuStreamWriteValue64(stream, CUdeviceptr(device_ptr64), value64, writeFlag);\n // CHECK: hipStreamWriteValue32(stream, hipDeviceptr_t(device_ptr32), value32, writeFlag);\n cuStreamWriteValue32(stream, CUdeviceptr(device_ptr32), value32, writeFlag);\n // CHECK: hipStreamSynchronize(stream);\n cudaStreamSynchronize(stream);\n // Test Writing to Signal Memory\n // CHECK: hipStreamWriteValue64(stream, hipDeviceptr_t(signalPtr), value64, writeFlag);\n cuStreamWriteValue64(stream, CUdeviceptr(signalPtr), value64, writeFlag);\n // CHECK: hipStreamSynchronize(stream);\n cudaStreamSynchronize(stream);\n // Cleanup\n // CHECK: hipStreamDestroy(stream);\n cudaStreamDestroy(stream);\n // CHECK: hipHostUnregister(host_ptr64);\n cudaHostUnregister(host_ptr64);\n // CHECK: hipHostUnregister(host_ptr32);\n cudaHostUnregister(host_ptr32);\n // CHECK: hipFree(signalPtr);\n cudaFree(signalPtr);\n free(host_ptr32);\n free(host_ptr64);\n}\n\nvoid testWait() {\n int64_t* signalPtr;\n // random data values\n int32_t DATA_INIT = 0x1234;\n int32_t DATA_UPDATE = 0X4321;\n\n struct TEST_WAIT {\n int compareOp;\n uint64_t mask;\n int64_t waitValue;\n int64_t signalValueFail;\n int64_t signalValuePass;\n };\n\n TEST_WAIT testCases[] = {\n {\n // mask will ignore few MSB bits\n // CHECK: hipStreamWaitValueGte,\n CU_STREAM_WAIT_VALUE_GEQ,\n 0x0000FFFFFFFFFFFF,\n 0x000000007FFF0001,\n 0x7FFF00007FFF0000,\n 0x000000007FFF0001\n },\n {\n // CHECK: hipStreamWaitValueGte,\n CU_STREAM_WAIT_VALUE_GEQ,\n 0xF,\n 0x4,\n 0x3,\n 0x6\n },\n {\n // mask will ignore few MSB bits\n // CHECK: hipStreamWaitValueEq,\n CU_STREAM_WAIT_VALUE_EQ,\n 0x0000FFFFFFFFFFFF,\n 0x000000000FFF0001,\n 0x7FFF00000FFF0000,\n 0x7F0000000FFF0001\n },\n {\n // CHECK: hipStreamWaitValueEq,\n CU_STREAM_WAIT_VALUE_EQ,\n 0xFF,\n 0x11,\n 0x25,\n 0x11\n },\n {\n // mask will discard bits 8 to 11\n // CHECK: hipStreamWaitValueAnd,\n CU_STREAM_WAIT_VALUE_AND,\n 0xFF,\n 0xF4A,\n 0xF35,\n 0X02\n },\n {\n // mask is set to ignore the sign bit.\n // CHECK: hipStreamWaitValueNor,\n CU_STREAM_WAIT_VALUE_NOR,\n 0x7FFFFFFFFFFFFFFF,\n 0x7FFFFFFFFFFFF247,\n 0x7FFFFFFFFFFFFdbd,\n 0x7FFFFFFFFFFFFdb5\n },\n {\n // mask is set to apply NOR for bits 0 to 3.\n // CHECK: hipStreamWaitValueNor,\n CU_STREAM_WAIT_VALUE_NOR,\n 0xF,\n 0x7E,\n 0x7D,\n 0x76\n }\n };\n\n struct TEST_WAIT32_NO_MASK {\n int compareOp;\n int32_t waitValue;\n int32_t signalValueFail;\n int32_t signalValuePass;\n };\n\n // default mask 0xFFFFFFFF will be used.\n TEST_WAIT32_NO_MASK testCasesNoMask32[] = {\n {\n // CHECK: hipStreamWaitValueGte,\n CU_STREAM_WAIT_VALUE_GEQ,\n 0x7FFF0001,\n 0x7FFF0000,\n 0x7FFF0010\n },\n {\n // CHECK: hipStreamWaitValueEq,\n CU_STREAM_WAIT_VALUE_EQ,\n 0x7FFFFFFF,\n 0x7FFF0000,\n 0x7FFFFFFF\n },\n {\n // CHECK: hipStreamWaitValueAnd,\n CU_STREAM_WAIT_VALUE_AND,\n 0x70F0F0F0,\n 0x0F0F0F0F,\n 0X1F0F0F0F\n },\n {\n // CHECK: hipStreamWaitValueNor,\n CU_STREAM_WAIT_VALUE_NOR,\n 0x7AAAAAAA,\n static_cast(0x85555555),\n static_cast(0x9AAAAAAA)\n }\n };\n\n struct TEST_WAIT64_NO_MASK {\n int compareOp;\n int64_t waitValue;\n int64_t signalValueFail;\n int64_t signalValuePass;\n };\n\n // default mask 0xFFFFFFFFFFFFFFFF will be used.\n TEST_WAIT64_NO_MASK testCasesNoMask64[] = {\n {\n // CHECK: hipStreamWaitValueGte,\n CU_STREAM_WAIT_VALUE_GEQ,\n 0x7FFFFFFFFFFF0001,\n 0x7FFFFFFFFFFF0000,\n 0x7FFFFFFFFFFF0001\n },\n {\n // CHECK: hipStreamWaitValueEq,\n CU_STREAM_WAIT_VALUE_EQ,\n 0x7FFFFFFFFFFFFFFF,\n 0x7FFFFFFF0FFF0000,\n 0x7FFFFFFFFFFFFFFF\n },\n {\n // CHECK: hipStreamWaitValueAnd,\n CU_STREAM_WAIT_VALUE_AND,\n 0x70F0F0F0F0F0F0F0,\n 0x0F0F0F0F0F0F0F0F,\n 0X1F0F0F0F0F0F0F0F\n },\n {\n // CHECK: hipStreamWaitValueNor,\n CU_STREAM_WAIT_VALUE_NOR,\n 0x4724724747247247,\n static_cast(0xbddbddbdbddbddbd),\n static_cast(0xbddbddbdbddbddb3)\n }\n };\n\n // CHECK: hipStream_t stream;\n cudaStream_t stream;\n // CHECK: hipStreamCreate(&stream);\n cudaStreamCreate(&stream);\n // hipExtMallocWithFlags((void**)&signalPtr, 8, hipMallocSignalMemory);\n int64_t* dataPtr64 = (int64_t*)malloc(sizeof(int64_t));\n int32_t* dataPtr32 = (int32_t*)malloc(sizeof(int32_t));\n // hipHostRegister(dataPtr64, sizeof(int64_t), 0);\n cudaHostRegister(dataPtr64, sizeof(int64_t), 0);\n // CHECK: hipHostRegister(dataPtr32, sizeof(int32_t), 0);\n cudaHostRegister(dataPtr32, sizeof(int32_t), 0);\n // Run-1: streamWait is blocking (wait conditions is false)\n // Run-2: streamWait is non-blocking (wait condition is true)\n for (int run = 0; run < 2; run++) {\n bool isBlocking = run == 0;\n for (const auto & tc : testCases) {\n *signalPtr = isBlocking ? tc.signalValueFail : tc.signalValuePass;\n *dataPtr64 = DATA_INIT;\n // CHECK: hipStreamWaitValue64(stream, hipDeviceptr_t(signalPtr), tc.waitValue, tc.compareOp);\n cuStreamWaitValue64(stream, CUdeviceptr(signalPtr), tc.waitValue, tc.compareOp);\n // CHECK: hipStreamWriteValue64(stream, hipDeviceptr_t(dataPtr64), DATA_UPDATE, writeFlag);\n cuStreamWriteValue64(stream, CUdeviceptr(dataPtr64), DATA_UPDATE, writeFlag);\n if (isBlocking) {\n // Trigger an implict flush and verify stream has pending work.\n // CHECK: if (hipStreamQuery(stream) != hipErrorNotReady) {}\n if (cudaStreamQuery(stream) != cudaErrorNotReady) {}\n // update signal to unblock the wait.\n *signalPtr = tc.signalValuePass;\n }\n // CHECK: if (hipStreamQuery(stream) != hipSuccess) {}\n if (cudaStreamQuery(stream) != cudaSuccess) {}\n // CHECK: hipStreamSynchronize(stream);\n cudaStreamSynchronize(stream);\n if (*dataPtr64 != DATA_UPDATE) {}\n // 32-bit API\n *signalPtr = isBlocking ? tc.signalValueFail : tc.signalValuePass;\n *dataPtr32 = DATA_INIT;\n // CHECK: hipStreamWaitValue32(stream, hipDeviceptr_t(signalPtr), tc.waitValue, tc.compareOp);\n cuStreamWaitValue32(stream, CUdeviceptr(signalPtr), tc.waitValue, tc.compareOp);\n // CHECK: hipStreamWriteValue32(stream, hipDeviceptr_t(dataPtr32), DATA_UPDATE, writeFlag);\n cuStreamWriteValue32(stream, CUdeviceptr(dataPtr32), DATA_UPDATE, writeFlag);\n if (isBlocking) {\n // Trigger an implict flush and verify stream has pending work.\n // CHECK: if (hipStreamQuery(stream) != hipErrorNotReady) {}\n if (cudaStreamQuery(stream) != cudaErrorNotReady) {}\n // update signal to unblock the wait.\n *signalPtr = static_cast(tc.signalValuePass);\n }\n // CHECK: hipStreamSynchronize(stream);\n cudaStreamSynchronize(stream);\n if (*dataPtr32 != DATA_UPDATE) {}\n }\n }\n // Run-1: streamWait is blocking (wait conditions is false)\n // Run-2: streamWait is non-blocking (wait condition is true)\n for (int run = 0; run < 2; run++) {\n bool isBlocking = run == 0;\n for (const auto& tc : testCasesNoMask32) {\n *signalPtr = isBlocking ? tc.signalValueFail : tc.signalValuePass;\n *dataPtr32 = DATA_INIT;\n // CHECK: hipStreamWaitValue32(stream, hipDeviceptr_t(signalPtr), tc.waitValue, tc.compareOp);\n cuStreamWaitValue32(stream, CUdeviceptr(signalPtr), tc.waitValue, tc.compareOp);\n // CHECK: hipStreamWriteValue32(stream, hipDeviceptr_t(dataPtr32), DATA_UPDATE, writeFlag);\n cuStreamWriteValue32(stream, CUdeviceptr(dataPtr32), DATA_UPDATE, writeFlag);\n if (isBlocking) {\n // Trigger an implict flush and verify stream has pending work.\n // CHECK: if (hipStreamQuery(stream) != hipErrorNotReady) {}\n if (cudaStreamQuery(stream) != cudaErrorNotReady) {}\n // update signal to unblock the wait.\n *signalPtr = tc.signalValuePass;\n }\n // CHECK: hipStreamSynchronize(stream);\n cudaStreamSynchronize(stream);\n if (*dataPtr32 != DATA_UPDATE) {}\n }\n }\n // Run-1: streamWait is blocking (wait conditions is false)\n // Run-2: streamWait is non-blocking (wait condition is true)\n for (int run = 0; run < 2; run++) {\n bool isBlocking = run == 0;\n for (const auto& tc : testCasesNoMask64) {\n *signalPtr = isBlocking ? tc.signalValueFail : tc.signalValuePass;\n *dataPtr64 = DATA_INIT;\n // CHECK: hipStreamWaitValue64(stream, hipDeviceptr_t(signalPtr), tc.waitValue, tc.compareOp);\n cuStreamWaitValue64(stream, CUdeviceptr(signalPtr), tc.waitValue, tc.compareOp);\n // CHECK: hipStreamWriteValue64(stream, hipDeviceptr_t(dataPtr64), DATA_UPDATE, writeFlag);\n cuStreamWriteValue64(stream, CUdeviceptr(dataPtr64), DATA_UPDATE, writeFlag);\n if (isBlocking) {\n // Trigger an implict flush and verify stream has pending work.\n // CHECK: if (hipStreamQuery(stream) != hipErrorNotReady) {}\n if (cudaStreamQuery(stream) != cudaErrorNotReady) {}\n // update signal to unblock the wait.\n *signalPtr = tc.signalValuePass;\n }\n // CHECK: hipStreamSynchronize(stream);\n cudaStreamSynchronize(stream);\n if (*dataPtr64 != DATA_UPDATE) {}\n }\n }\n // Cleanup\n // CHECK: hipFree(signalPtr);\n cudaFree(signalPtr);\n // CHECK: hipHostUnregister(dataPtr64);\n cudaHostUnregister(dataPtr64);\n // CHECK: hipHostUnregister(dataPtr32);\n cudaHostUnregister(dataPtr32);\n free(dataPtr64);\n free(dataPtr32);\n // CHECK: hipStreamDestroy(stream);\n cudaStreamDestroy(stream);\n}\n\nint main() {\n // CHECK: hipFuncCache_t cacheConfig;\n cudaFuncCache cacheConfig;\n void* func;\n // CHECK: hipFuncSetCacheConfig(reinterpret_cast(func), cacheConfig);\n cudaFuncSetCacheConfig(func, cacheConfig);\n // CHECK: hipFuncAttributes attr{};\n cudaFuncAttributes attr{};\n // CHECK: auto r = hipFuncGetAttributes(&attr, reinterpret_cast(&fn));\n auto r = cudaFuncGetAttributes(&attr, &fn);\n // CHECK: if (r != hipSuccess || attr.maxThreadsPerBlock == 0) {\n if (r != cudaSuccess || attr.maxThreadsPerBlock == 0) {\n return 1;\n }\n testWrite();\n testWait();\n return 0;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args %clang_args\n\n/*\nCopyright (c) 2015-present Advanced Micro Devices, Inc. All rights reserved.\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in\nall copies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\nTHE SOFTWARE.\n*/\n\n// CHECK: #include \n#include \n#include \n\n// Random predefiend 32 and 64 bit values\nconstexpr uint32_t value32 = 0x70F0F0FF;\nconstexpr uint64_t value64 = 0x7FFF0000FFFF0000;\nconstexpr unsigned int writeFlag = 0;\n\n__global__\nvoid fn(float* px, float* py) {\n bool a[42];\n __shared__ double b[69];\n for (auto&& x : b) x = *py++;\n for (auto&& x : a) x = *px++ > 0.0;\n for (auto&& x : a) if (x) *--py = *--px;\n}\n\nvoid testWrite() {\n int64_t* signalPtr;\n // CHECK: hipStream_t stream;\n hipStream_t stream;\n // CHECK: hipStreamCreate(&stream);\n hipStreamCreate(&stream);\n int64_t* host_ptr64 = (int64_t*)malloc(sizeof(int64_t));\n int32_t* host_ptr32 = (int32_t*)malloc(sizeof(int32_t));\n // hipExtMallocWithFlags((void**)&signalPtr, 8, hipMallocSignalMemory);\n void* device_ptr64;\n void* device_ptr32;\n *host_ptr64 = 0x0;\n *host_ptr32 = 0x0;\n *signalPtr = 0x0;\n // CHECK: hipHostRegister(host_ptr64, sizeof(int64_t), 0);\n hipHostRegister(host_ptr64, sizeof(int64_t), 0);\n // CHECK: hipHostRegister(host_ptr32, sizeof(int32_t), 0);\n hipHostRegister(host_ptr32, sizeof(int32_t), 0);\n // CHECK: hipStreamWriteValue64(stream, hipDeviceptr_t(host_ptr64), value64, writeFlag);\n hipStreamWriteValue64(stream, hipDeviceptr_t(host_ptr64), value64, writeFlag);\n // CHECK: hipStreamWriteValue32(stream, hipDeviceptr_t(host_ptr32), value32, writeFlag);\n hipStreamWriteValue32(stream, hipDeviceptr_t(host_ptr32), value32, writeFlag);\n // CHECK: hipStreamSynchronize(stream);\n hipStreamSynchronize(stream);\n // CHECK: hipHostGetDevicePointer((void**)&device_ptr64, host_ptr64, 0);\n hipHostGetDevicePointer((void**)&device_ptr64, host_ptr64, 0);\n // CHECK: hipHostGetDevicePointer((void**)&device_ptr32, host_ptr32, 0);\n hipHostGetDevicePointer((void**)&device_ptr32, host_ptr32, 0);\n // Reset values\n *host_ptr64 = 0x0;\n *host_ptr32 = 0x0;\n // CHECK: hipStreamWriteValue64(stream, hipDeviceptr_t(device_ptr64), value64, writeFlag);\n hipStreamWriteValue64(stream, hipDeviceptr_t(device_ptr64), value64, writeFlag);\n // CHECK: hipStreamWriteValue32(stream, hipDeviceptr_t(device_ptr32), value32, writeFlag);\n hipStreamWriteValue32(stream, hipDeviceptr_t(device_ptr32), value32, writeFlag);\n // CHECK: hipStreamSynchronize(stream);\n hipStreamSynchronize(stream);\n // Test Writing to Signal Memory\n // CHECK: hipStreamWriteValue64(stream, hipDeviceptr_t(signalPtr), value64, writeFlag);\n hipStreamWriteValue64(stream, hipDeviceptr_t(signalPtr), value64, writeFlag);\n // CHECK: hipStreamSynchronize(stream);\n hipStreamSynchronize(stream);\n // Cleanup\n // CHECK: hipStreamDestroy(stream);\n hipStreamDestroy(stream);\n // CHECK: hipHostUnregister(host_ptr64);\n hipHostUnregister(host_ptr64);\n // CHECK: hipHostUnregister(host_ptr32);\n hipHostUnregister(host_ptr32);\n // CHECK: hipFree(signalPtr);\n hipFree(signalPtr);\n free(host_ptr32);\n free(host_ptr64);\n}\n\nvoid testWait() {\n int64_t* signalPtr;\n // random data values\n int32_t DATA_INIT = 0x1234;\n int32_t DATA_UPDATE = 0X4321;\n\n struct TEST_WAIT {\n int compareOp;\n uint64_t mask;\n int64_t waitValue;\n int64_t signalValueFail;\n int64_t signalValuePass;\n };\n\n TEST_WAIT testCases[] = {\n {\n // mask will ignore few MSB bits\n // CHECK: hipStreamWaitValueGte,\n hipStreamWaitValueGte,\n 0x0000FFFFFFFFFFFF,\n 0x000000007FFF0001,\n 0x7FFF00007FFF0000,\n 0x000000007FFF0001\n },\n {\n // CHECK: hipStreamWaitValueGte,\n hipStreamWaitValueGte,\n 0xF,\n 0x4,\n 0x3,\n 0x6\n },\n {\n // mask will ignore few MSB bits\n // CHECK: hipStreamWaitValueEq,\n hipStreamWaitValueEq,\n 0x0000FFFFFFFFFFFF,\n 0x000000000FFF0001,\n 0x7FFF00000FFF0000,\n 0x7F0000000FFF0001\n },\n {\n // CHECK: hipStreamWaitValueEq,\n hipStreamWaitValueEq,\n 0xFF,\n 0x11,\n 0x25,\n 0x11\n },\n {\n // mask will discard bits 8 to 11\n // CHECK: hipStreamWaitValueAnd,\n hipStreamWaitValueAnd,\n 0xFF,\n 0xF4A,\n 0xF35,\n 0X02\n },\n {\n // mask is set to ignore the sign bit.\n // CHECK: hipStreamWaitValueNor,\n hipStreamWaitValueNor,\n 0x7FFFFFFFFFFFFFFF,\n 0x7FFFFFFFFFFFF247,\n 0x7FFFFFFFFFFFFdbd,\n 0x7FFFFFFFFFFFFdb5\n },\n {\n // mask is set to apply NOR for bits 0 to 3.\n // CHECK: hipStreamWaitValueNor,\n hipStreamWaitValueNor,\n 0xF,\n 0x7E,\n 0x7D,\n 0x76\n }\n };\n\n struct TEST_WAIT32_NO_MASK {\n int compareOp;\n int32_t waitValue;\n int32_t signalValueFail;\n int32_t signalValuePass;\n };\n\n // default mask 0xFFFFFFFF will be used.\n TEST_WAIT32_NO_MASK testCasesNoMask32[] = {\n {\n // CHECK: hipStreamWaitValueGte,\n hipStreamWaitValueGte,\n 0x7FFF0001,\n 0x7FFF0000,\n 0x7FFF0010\n },\n {\n // CHECK: hipStreamWaitValueEq,\n hipStreamWaitValueEq,\n 0x7FFFFFFF,\n 0x7FFF0000,\n 0x7FFFFFFF\n },\n {\n // CHECK: hipStreamWaitValueAnd,\n hipStreamWaitValueAnd,\n 0x70F0F0F0,\n 0x0F0F0F0F,\n 0X1F0F0F0F\n },\n {\n // CHECK: hipStreamWaitValueNor,\n hipStreamWaitValueNor,\n 0x7AAAAAAA,\n static_cast(0x85555555),\n static_cast(0x9AAAAAAA)\n }\n };\n\n struct TEST_WAIT64_NO_MASK {\n int compareOp;\n int64_t waitValue;\n int64_t signalValueFail;\n int64_t signalValuePass;\n };\n\n // default mask 0xFFFFFFFFFFFFFFFF will be used.\n TEST_WAIT64_NO_MASK testCasesNoMask64[] = {\n {\n // CHECK: hipStreamWaitValueGte,\n hipStreamWaitValueGte,\n 0x7FFFFFFFFFFF0001,\n 0x7FFFFFFFFFFF0000,\n 0x7FFFFFFFFFFF0001\n },\n {\n // CHECK: hipStreamWaitValueEq,\n hipStreamWaitValueEq,\n 0x7FFFFFFFFFFFFFFF,\n 0x7FFFFFFF0FFF0000,\n 0x7FFFFFFFFFFFFFFF\n },\n {\n // CHECK: hipStreamWaitValueAnd,\n hipStreamWaitValueAnd,\n 0x70F0F0F0F0F0F0F0,\n 0x0F0F0F0F0F0F0F0F,\n 0X1F0F0F0F0F0F0F0F\n },\n {\n // CHECK: hipStreamWaitValueNor,\n hipStreamWaitValueNor,\n 0x4724724747247247,\n static_cast(0xbddbddbdbddbddbd),\n static_cast(0xbddbddbdbddbddb3)\n }\n };\n\n // CHECK: hipStream_t stream;\n hipStream_t stream;\n // CHECK: hipStreamCreate(&stream);\n hipStreamCreate(&stream);\n // hipExtMallocWithFlags((void**)&signalPtr, 8, hipMallocSignalMemory);\n int64_t* dataPtr64 = (int64_t*)malloc(sizeof(int64_t));\n int32_t* dataPtr32 = (int32_t*)malloc(sizeof(int32_t));\n // hipHostRegister(dataPtr64, sizeof(int64_t), 0);\n hipHostRegister(dataPtr64, sizeof(int64_t), 0);\n // CHECK: hipHostRegister(dataPtr32, sizeof(int32_t), 0);\n hipHostRegister(dataPtr32, sizeof(int32_t), 0);\n // Run-1: streamWait is blocking (wait conditions is false)\n // Run-2: streamWait is non-blocking (wait condition is true)\n for (int run = 0; run < 2; run++) {\n bool isBlocking = run == 0;\n for (const auto & tc : testCases) {\n *signalPtr = isBlocking ? tc.signalValueFail : tc.signalValuePass;\n *dataPtr64 = DATA_INIT;\n // CHECK: hipStreamWaitValue64(stream, hipDeviceptr_t(signalPtr), tc.waitValue, tc.compareOp);\n hipStreamWaitValue64(stream, hipDeviceptr_t(signalPtr), tc.waitValue, tc.compareOp);\n // CHECK: hipStreamWriteValue64(stream, hipDeviceptr_t(dataPtr64), DATA_UPDATE, writeFlag);\n hipStreamWriteValue64(stream, hipDeviceptr_t(dataPtr64), DATA_UPDATE, writeFlag);\n if (isBlocking) {\n // Trigger an implict flush and verify stream has pending work.\n // CHECK: if (hipStreamQuery(stream) != hipErrorNotReady) {}\n if (hipStreamQuery(stream) != hipErrorNotReady) {}\n // update signal to unblock the wait.\n *signalPtr = tc.signalValuePass;\n }\n // CHECK: if (hipStreamQuery(stream) != hipSuccess) {}\n if (hipStreamQuery(stream) != hipSuccess) {}\n // CHECK: hipStreamSynchronize(stream);\n hipStreamSynchronize(stream);\n if (*dataPtr64 != DATA_UPDATE) {}\n // 32-bit API\n *signalPtr = isBlocking ? tc.signalValueFail : tc.signalValuePass;\n *dataPtr32 = DATA_INIT;\n // CHECK: hipStreamWaitValue32(stream, hipDeviceptr_t(signalPtr), tc.waitValue, tc.compareOp);\n hipStreamWaitValue32(stream, hipDeviceptr_t(signalPtr), tc.waitValue, tc.compareOp);\n // CHECK: hipStreamWriteValue32(stream, hipDeviceptr_t(dataPtr32), DATA_UPDATE, writeFlag);\n hipStreamWriteValue32(stream, hipDeviceptr_t(dataPtr32), DATA_UPDATE, writeFlag);\n if (isBlocking) {\n // Trigger an implict flush and verify stream has pending work.\n // CHECK: if (hipStreamQuery(stream) != hipErrorNotReady) {}\n if (hipStreamQuery(stream) != hipErrorNotReady) {}\n // update signal to unblock the wait.\n *signalPtr = static_cast(tc.signalValuePass);\n }\n // CHECK: hipStreamSynchronize(stream);\n hipStreamSynchronize(stream);\n if (*dataPtr32 != DATA_UPDATE) {}\n }\n }\n // Run-1: streamWait is blocking (wait conditions is false)\n // Run-2: streamWait is non-blocking (wait condition is true)\n for (int run = 0; run < 2; run++) {\n bool isBlocking = run == 0;\n for (const auto& tc : testCasesNoMask32) {\n *signalPtr = isBlocking ? tc.signalValueFail : tc.signalValuePass;\n *dataPtr32 = DATA_INIT;\n // CHECK: hipStreamWaitValue32(stream, hipDeviceptr_t(signalPtr), tc.waitValue, tc.compareOp);\n hipStreamWaitValue32(stream, hipDeviceptr_t(signalPtr), tc.waitValue, tc.compareOp);\n // CHECK: hipStreamWriteValue32(stream, hipDeviceptr_t(dataPtr32), DATA_UPDATE, writeFlag);\n hipStreamWriteValue32(stream, hipDeviceptr_t(dataPtr32), DATA_UPDATE, writeFlag);\n if (isBlocking) {\n // Trigger an implict flush and verify stream has pending work.\n // CHECK: if (hipStreamQuery(stream) != hipErrorNotReady) {}\n if (hipStreamQuery(stream) != hipErrorNotReady) {}\n // update signal to unblock the wait.\n *signalPtr = tc.signalValuePass;\n }\n // CHECK: hipStreamSynchronize(stream);\n hipStreamSynchronize(stream);\n if (*dataPtr32 != DATA_UPDATE) {}\n }\n }\n // Run-1: streamWait is blocking (wait conditions is false)\n // Run-2: streamWait is non-blocking (wait condition is true)\n for (int run = 0; run < 2; run++) {\n bool isBlocking = run == 0;\n for (const auto& tc : testCasesNoMask64) {\n *signalPtr = isBlocking ? tc.signalValueFail : tc.signalValuePass;\n *dataPtr64 = DATA_INIT;\n // CHECK: hipStreamWaitValue64(stream, hipDeviceptr_t(signalPtr), tc.waitValue, tc.compareOp);\n hipStreamWaitValue64(stream, hipDeviceptr_t(signalPtr), tc.waitValue, tc.compareOp);\n // CHECK: hipStreamWriteValue64(stream, hipDeviceptr_t(dataPtr64), DATA_UPDATE, writeFlag);\n hipStreamWriteValue64(stream, hipDeviceptr_t(dataPtr64), DATA_UPDATE, writeFlag);\n if (isBlocking) {\n // Trigger an implict flush and verify stream has pending work.\n // CHECK: if (hipStreamQuery(stream) != hipErrorNotReady) {}\n if (hipStreamQuery(stream) != hipErrorNotReady) {}\n // update signal to unblock the wait.\n *signalPtr = tc.signalValuePass;\n }\n // CHECK: hipStreamSynchronize(stream);\n hipStreamSynchronize(stream);\n if (*dataPtr64 != DATA_UPDATE) {}\n }\n }\n // Cleanup\n // CHECK: hipFree(signalPtr);\n hipFree(signalPtr);\n // CHECK: hipHostUnregister(dataPtr64);\n hipHostUnregister(dataPtr64);\n // CHECK: hipHostUnregister(dataPtr32);\n hipHostUnregister(dataPtr32);\n free(dataPtr64);\n free(dataPtr32);\n // CHECK: hipStreamDestroy(stream);\n hipStreamDestroy(stream);\n}\n\nint main() {\n // CHECK: hipFuncCache_t cacheConfig;\n hipFuncCache_t cacheConfig;\n void* func;\n // CHECK: hipFuncSetCacheConfig(reinterpret_cast(func), cacheConfig);\n hipFuncSetCacheConfig(reinterpret_cast(func), cacheConfig);\n // CHECK: hipFuncAttributes attr{};\n hipFuncAttributes attr{};\n // CHECK: auto r = hipFuncGetAttributes(&attr, reinterpret_cast(&fn));\n auto r = hipFuncGetAttributes(&attr,reinterpret_cast(&fn));\n // CHECK: if (r != hipSuccess || attr.maxThreadsPerBlock == 0) {\n if (r != hipSuccess || attr.maxThreadsPerBlock == 0) {\n return 1;\n }\n testWrite();\n testWait();\n return 0;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/casts/reinterpret_cast.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n#include \n\n__global__ void test_0() {\n int a = 10;\n}\n\nnamespace first {\n __global__ void test_1() {\n int b = 20;\n }\n namespace second {\n __global__ void test_2() {\n int c = 30;\n }\n }\n}\n\nint main() {\n // CHECK: hipLaunchKernelGGL(::test_0, dim3(1), dim3(1), 0, 0);\n ::test_0<<<1, 1>>>();\n // CHECK: hipLaunchKernelGGL(first::test_1, dim3(1), dim3(1), 0, 0);\n first::test_1<<<1, 1>>>();\n // CHECK: hipLaunchKernelGGL(first::second::test_2, dim3(1), dim3(1), 0, 0);\n first::second::test_2<<<1, 1>>>();\n return 0;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// CHECK: #include \n#include \n\n__global__ void test_0() {\n int a = 10;\n}\n\nnamespace first {\n __global__ void test_1() {\n int b = 20;\n }\n namespace second {\n __global__ void test_2() {\n int c = 30;\n }\n }\n}\n\nint main() {\n // CHECK: hipLaunchKernelGGL(::test_0, dim3(1), dim3(1), 0, 0);\n ::test_0<<<1, 1>>>();\n // CHECK: hipLaunchKernelGGL(first::test_1, dim3(1), dim3(1), 0, 0);\n first::test_1<<<1, 1>>>();\n // CHECK: hipLaunchKernelGGL(first::second::test_2, dim3(1), dim3(1), 0, 0);\n first::second::test_2<<<1, 1>>>();\n return 0;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/namespace/ns_kernel_launch.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#pragma once\n// CHECK: #include \n#include \n#include \n// CHECK: #include \"hip/hip_fp16.h\"\n#include \"cuda_fp16.h\"\n/**\n * Allocate GPU memory for `count` elements of type `T`.\n */\ntemplate\nstatic T *gpuMalloc(size_t count) {\n T *ret = nullptr;\n // CHECK: hipMalloc(&ret, count * sizeof(T));\n cudaMalloc(&ret, count * sizeof(T));\n return ret;\n}\n\ntemplate\n__global__ void add(int n, T *x, T *y) {\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n // CHECK: half2 tmp = __float2half2_rn(0.0f);\n half2 tmp = __float2half2_rn(0.0f);\n // CHECK: float max_val = fmax((float)reinterpret_cast(tmp.x), (float)reinterpret_cast(tmp.y));\n float max_val = fmax((float)tmp.x, (float)tmp.y);\n // CHECK: float min_val = fmin((float)reinterpret_cast(tmp.x), (float)reinterpret_cast(tmp.y));\n float min_val = fmin((float)reinterpret_cast(tmp.x), (float)reinterpret_cast(tmp.y));\n for (int i = index; i < n; i += stride)\n y[i] = max_val - min_val + x[i] + y[i];\n}\n\nint main(int argc, char* argv[]) {\n size_t numElements = 50;\n float *A = gpuMalloc(numElements);\n float* B = gpuMalloc(numElements);\n for (int i = 0; i < numElements; ++i) {\n A[i] = 1.0f;\n B[i] = 2.0f;\n }\n int blockSize = 512;\n int numBlocks = (numElements + blockSize - 1) / blockSize;\n dim3 dimGrid(numBlocks, 1, 1);\n dim3 dimBlock(blockSize, 1, 1);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(add), dim3(dimGrid), dim3(dimBlock), 0, 0, numElements, A, B);\n add<<>>(numElements, A, B);\n // CHECK: hipDeviceSynchronize();\n cudaDeviceSynchronize();\n float maxError = 0.0f;\n for (int i = 0; i < numElements; ++i)\n maxError = fmax(maxError, fabs(B[i] - 3.0f));\n // CHECK: hipFree(A);\n cudaFree(A);\n // CHECK: hipFree(B);\n cudaFree(B);\n if (maxError == 0.0f)\n return 0;\n return -1;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#pragma once\n// CHECK: #include \n#include \n#include \n// CHECK: #include \"hip/hip_fp16.h\"\n#include \"hip/hip_fp16.h\"\n/**\n * Allocate GPU memory for `count` elements of type `T`.\n */\ntemplate\nstatic T *gpuMalloc(size_t count) {\n T *ret = nullptr;\n // CHECK: hipMalloc(&ret, count * sizeof(T));\n hipMalloc(&ret, count * sizeof(T));\n return ret;\n}\n\ntemplate\n__global__ void add(int n, T *x, T *y) {\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n // CHECK: half2 tmp = __float2half2_rn(0.0f);\n half2 tmp = __float2half2_rn(0.0f);\n // CHECK: float max_val = fmax((float)reinterpret_cast(tmp.x), (float)reinterpret_cast(tmp.y));\n float max_val = fmax((float)tmp.x, (float)tmp.y);\n // CHECK: float min_val = fmin((float)reinterpret_cast(tmp.x), (float)reinterpret_cast(tmp.y));\n float min_val = fmin((float)reinterpret_cast(tmp.x), (float)reinterpret_cast(tmp.y));\n for (int i = index; i < n; i += stride)\n y[i] = max_val - min_val + x[i] + y[i];\n}\n\nint main(int argc, char* argv[]) {\n size_t numElements = 50;\n float *A = gpuMalloc(numElements);\n float* B = gpuMalloc(numElements);\n for (int i = 0; i < numElements; ++i) {\n A[i] = 1.0f;\n B[i] = 2.0f;\n }\n int blockSize = 512;\n int numBlocks = (numElements + blockSize - 1) / blockSize;\n dim3 dimGrid(numBlocks, 1, 1);\n dim3 dimBlock(blockSize, 1, 1);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(add), dim3(dimGrid), dim3(dimBlock), 0, 0, numElements, A, B);\n add<<>>(numElements, A, B);\n // CHECK: hipDeviceSynchronize();\n hipDeviceSynchronize();\n float maxError = 0.0f;\n for (int i = 0; i < numElements; ++i)\n maxError = fmax(maxError, fabs(B[i] - 3.0f));\n // CHECK: hipFree(A);\n hipFree(A);\n // CHECK: hipFree(B);\n hipFree(B);\n if (maxError == 0.0f)\n return 0;\n return -1;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/samples/half2_allocators.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n/*\nCopyright (c) 2015-2016 Advanced Micro Devices, Inc. All rights reserved.\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\nThe above copyright notice and this permission notice shall be included in\nall copies or substantial portions of the Software.\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\nTHE SOFTWARE.\n*/\n\n#include\n#include\n#include\n#include\n#include\n\n#define LEN 1024\n#define SIZE LEN * sizeof(float)\n#define ITER 1024*1024\n\n// CHECK: if(status != hipSuccess) {\n#define check(msg, status){ \\\nif(status != cudaSuccess) { \\\n printf(\"%s failed. \\n\", #msg); \\\n} \\\n}\n\n__global__ void Inc1(float *Ad, float *Bd){\n\tint tx = threadIdx.x + blockIdx.x * blockDim.x;\n\tif(tx < 1 ){\n\t\tfor(int i=0;i>>(Ad, Bd);\n\tA[0] = -(ITER*1.0f);\n\tstd::cout<<\"Same cache line before completion: \\t\"<< A[0]<>>(Ad, Bd);\n\tA[0] = -(ITER*1.0f);\n\tstd::cout<<\"Diff cache line before completion: \\t\"<\n#include\n#include\n#include\n#include\n\n#define LEN 1024\n#define SIZE LEN * sizeof(float)\n#define ITER 1024*1024\n\n// CHECK: if(status != hipSuccess) {\n#define check(msg, status){ \\\nif(status != hipSuccess) { \\\n printf(\"%s failed. \\n\", #msg); \\\n} \\\n}\n\n__global__ void Inc1(float *Ad, float *Bd){\n\tint tx = threadIdx.x + blockIdx.x * blockDim.x;\n\tif(tx < 1 ){\n\t\tfor(int i=0;i>>(Ad, Bd);\n\tA[0] = -(ITER*1.0f);\n\tstd::cout<<\"Same cache line before completion: \\t\"<< A[0]<>>(Ad, Bd);\n\tA[0] = -(ITER*1.0f);\n\tstd::cout<<\"Diff cache line before completion: \\t\"<\n#include \n#include \n// CHECK: #include \n#include \n\n#define K_THREADS 64\n#define K_INDEX() ((gridDim.x * blockIdx.y + blockIdx.x) * blockDim.x + threadIdx.x)\n#define RND() ((rand() & 0x7FFF) / float(0x8000))\n#define ERRORCHECK() cErrorCheck(__FILE__, __LINE__)\n\n// CHECK: hipEvent_t t##_start, t##_end; \\\n// CHECK: hipEventCreate(&t##_start); \\\n// CHECK: hipEventCreate(&t##_end);\n#define TIMER_CREATE(t) \\\n cudaEvent_t t##_start, t##_end; \\\n cudaEventCreate(&t##_start); \\\n cudaEventCreate(&t##_end);\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start);\n#define TIMER_START(t) \\\n cudaEventRecord(t##_start); \\\n cudaEventSynchronize(t##_start); \\\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start); \\\n// CHECK: hipEventRecord(t##_end); \\\n// CHECK: hipEventSynchronize(t##_end); \\\n// CHECK: hipEventElapsedTime(&t, t##_start, t##_end);\n#define TIMER_END(t) \\\n cudaEventRecord(t##_start); \\\n cudaEventSynchronize(t##_start); \\\n cudaEventRecord(t##_end); \\\n cudaEventSynchronize(t##_end); \\\n cudaEventElapsedTime(&t, t##_start, t##_end); \n\n\ninline void cErrorCheck(const char *file, int line) {\n// CHECK: hipError_t err = hipGetLastError();\n// CHECK: if (err != hipSuccess) {\n// CHECK: printf(\"Error: %s\\n\", hipGetErrorString(err));\n cudaError_t err = cudaGetLastError();\n if (err != cudaSuccess) {\n printf(\"Error: %s\\n\", cudaGetErrorString(err));\n printf(\" @ %s: %d\\n\", file, line);\n exit(-1);\n }\n}\n\ninline dim3 K_GRID(int n, int threads = K_THREADS) {\n int blocks = (int)ceilf(sqrtf((float)n/threads));\n dim3 grid(blocks, blocks);\n return grid;\n}\n\ntypedef struct data {\n int n;\n float4 *r, *v, *f;\n} data;\n\ndata cpu, gpu;\n\n#define N 20\n\n__global__ void repulsion(data gpu);\n__global__ void integration(data gpu);\n\n\nint main() {\n printf(\"Cuda Test 1\\n\");\n\n int count = 0;\n // CHECK: hipGetDeviceCount(&count);\n cudaGetDeviceCount(&count);\n printf(\" %d CUDA devices found\\n\", count);\n if(!count) {\n ::exit(EXIT_FAILURE);\n }\n // CHECK: hipFree(0);\n cudaFree(0);\n\n cpu.n = N;\n\n cpu.r = (float4*)malloc(N * sizeof(float4));\n cpu.v = (float4*)malloc(N * sizeof(float4));\n cpu.f = (float4*)malloc(N * sizeof(float4));\n\n for(int i = 0; i < N; ++i) {\n cpu.v[i] = make_float4(0,0,0,0);\n cpu.r[i] = make_float4(RND(), RND(), RND(), 0);\n cpu.f[i] = make_float4(0,0.01,0,0);\n }\n\n gpu = cpu;\n // CHECK: hipMalloc(&gpu.r, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.v, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.f, N * sizeof(float4));\n cudaMalloc(&gpu.r, N * sizeof(float4));\n cudaMalloc(&gpu.v, N * sizeof(float4));\n cudaMalloc(&gpu.f, N * sizeof(float4));\n // CHECK: hipMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n cudaMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), cudaMemcpyHostToDevice);\n cudaMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), cudaMemcpyHostToDevice);\n cudaMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), cudaMemcpyHostToDevice);\n\n ERRORCHECK();\n float rep;\n TIMER_CREATE(rep);\n TIMER_START(rep);\n // CHECK: hipLaunchKernelGGL(integration, dim3(K_GRID(cpu.n)), dim3(K_THREADS), 0, 0, gpu);\n integration <<< K_GRID(cpu.n), K_THREADS >>>(gpu);\n\n TIMER_END(rep);\n printf(\"Took: %f ms\\n\", rep);\n ERRORCHECK();\n // CHECK: hipMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n cudaMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), cudaMemcpyDeviceToHost);\n cudaMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), cudaMemcpyDeviceToHost);\n cudaMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), cudaMemcpyDeviceToHost);\n // CHECK: hipHostFree(cpu.r);\n // CHECK: hipHostFree(cpu.v);\n // CHECK: hipHostFree(cpu.f);\n cudaFreeHost(cpu.r);\n cudaFreeHost(cpu.v);\n cudaFreeHost(cpu.f);\n // CHECK: hipFree(gpu.r);\n // CHECK: hipFree(gpu.v);\n // CHECK: hipFree(gpu.f);\n cudaFree(gpu.r);\n cudaFree(gpu.v);\n cudaFree(gpu.f);\n // CHECK: hipDeviceReset();\n cudaDeviceReset();\n\n printf(\"Results: \\n\");\n for(int i = 0; i < N; ++i) {\n printf(\"%f, %f, %f \\n\", cpu.r[i].x, cpu.r[i].y, cpu.r[i].z);\n }\n\n printf(\"Ready...\\n\");\n return 0;\n}\n\n__global__ void repulsion(data gpu) {\n int idx = K_INDEX();\n if(idx < N) {\n gpu.r[idx].x = 1;\n gpu.r[idx].y = 1;\n gpu.r[idx].z = 1;\n }\n}\n\n#define MULT4(v, s) v.x *= s; v.y *= s; v.z *= s; v.w *= s;\n#define ADD4(v1, v2) v1.x += v2.x; v1.y += v2.y; v1.z += v2.z; v1.w += v2.w;\n\n__global__ void integration(data gpu) {\n int i = K_INDEX();\n if(i < N) {\n MULT4(gpu.f[i], 0.01);\n MULT4(gpu.v[i], 0.01);\n ADD4(gpu.v[i], gpu.f[i]);\n ADD4(gpu.r[i], gpu.v[i]);\n gpu.f[i] = make_float4(0,0,0,0);\n }\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#include \n#include \n#include \n// CHECK: #include \n#include \n\n#define K_THREADS 64\n#define K_INDEX() ((gridDim.x * blockIdx.y + blockIdx.x) * blockDim.x + threadIdx.x)\n#define RND() ((rand() & 0x7FFF) / float(0x8000))\n#define ERRORCHECK() cErrorCheck(__FILE__, __LINE__)\n\n// CHECK: hipEvent_t t##_start, t##_end; \\\n// CHECK: hipEventCreate(&t##_start); \\\n// CHECK: hipEventCreate(&t##_end);\n#define TIMER_CREATE(t) \\\n hipEvent_t t##_start, t##_end; \\\n hipEventCreate(&t##_start); \\\n hipEventCreate(&t##_end);\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start);\n#define TIMER_START(t) \\\n hipEventRecord(t##_start); \\\n hipEventSynchronize(t##_start); \\\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start); \\\n// CHECK: hipEventRecord(t##_end); \\\n// CHECK: hipEventSynchronize(t##_end); \\\n// CHECK: hipEventElapsedTime(&t, t##_start, t##_end);\n#define TIMER_END(t) \\\n hipEventRecord(t##_start); \\\n hipEventSynchronize(t##_start); \\\n hipEventRecord(t##_end); \\\n hipEventSynchronize(t##_end); \\\n hipEventElapsedTime(&t, t##_start, t##_end); \n\n\ninline void cErrorCheck(const char *file, int line) {\n// CHECK: hipError_t err = hipGetLastError();\n// CHECK: if (err != hipSuccess) {\n// CHECK: printf(\"Error: %s\\n\", hipGetErrorString(err));\n hipError_t err = hipGetLastError();\n if (err != hipSuccess) {\n printf(\"Error: %s\\n\", hipGetErrorString(err));\n printf(\" @ %s: %d\\n\", file, line);\n exit(-1);\n }\n}\n\ninline dim3 K_GRID(int n, int threads = K_THREADS) {\n int blocks = (int)ceilf(sqrtf((float)n/threads));\n dim3 grid(blocks, blocks);\n return grid;\n}\n\ntypedef struct data {\n int n;\n float4 *r, *v, *f;\n} data;\n\ndata cpu, gpu;\n\n#define N 20\n\n__global__ void repulsion(data gpu);\n__global__ void integration(data gpu);\n\n\nint main() {\n printf(\"Cuda Test 1\\n\");\n\n int count = 0;\n // CHECK: hipGetDeviceCount(&count);\n hipGetDeviceCount(&count);\n printf(\" %d CUDA devices found\\n\", count);\n if(!count) {\n ::exit(EXIT_FAILURE);\n }\n // CHECK: hipFree(0);\n hipFree(0);\n\n cpu.n = N;\n\n cpu.r = (float4*)malloc(N * sizeof(float4));\n cpu.v = (float4*)malloc(N * sizeof(float4));\n cpu.f = (float4*)malloc(N * sizeof(float4));\n\n for(int i = 0; i < N; ++i) {\n cpu.v[i] = make_float4(0,0,0,0);\n cpu.r[i] = make_float4(RND(), RND(), RND(), 0);\n cpu.f[i] = make_float4(0,0.01,0,0);\n }\n\n gpu = cpu;\n // CHECK: hipMalloc(&gpu.r, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.v, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.f, N * sizeof(float4));\n hipMalloc(&gpu.r, N * sizeof(float4));\n hipMalloc(&gpu.v, N * sizeof(float4));\n hipMalloc(&gpu.f, N * sizeof(float4));\n // CHECK: hipMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n hipMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n hipMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n hipMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n\n ERRORCHECK();\n float rep;\n TIMER_CREATE(rep);\n TIMER_START(rep);\n // CHECK: hipLaunchKernelGGL(integration, dim3(K_GRID(cpu.n)), dim3(K_THREADS), 0, 0, gpu);\n integration <<< K_GRID(cpu.n), K_THREADS >>>(gpu);\n\n TIMER_END(rep);\n printf(\"Took: %f ms\\n\", rep);\n ERRORCHECK();\n // CHECK: hipMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n hipMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n hipMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n hipMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipHostFree(cpu.r);\n // CHECK: hipHostFree(cpu.v);\n // CHECK: hipHostFree(cpu.f);\n hipHostFree(cpu.r);\n hipHostFree(cpu.v);\n hipHostFree(cpu.f);\n // CHECK: hipFree(gpu.r);\n // CHECK: hipFree(gpu.v);\n // CHECK: hipFree(gpu.f);\n hipFree(gpu.r);\n hipFree(gpu.v);\n hipFree(gpu.f);\n // CHECK: hipDeviceReset();\n hipDeviceReset();\n\n printf(\"Results: \\n\");\n for(int i = 0; i < N; ++i) {\n printf(\"%f, %f, %f \\n\", cpu.r[i].x, cpu.r[i].y, cpu.r[i].z);\n }\n\n printf(\"Ready...\\n\");\n return 0;\n}\n\n__global__ void repulsion(data gpu) {\n int idx = K_INDEX();\n if(idx < N) {\n gpu.r[idx].x = 1;\n gpu.r[idx].y = 1;\n gpu.r[idx].z = 1;\n }\n}\n\n#define MULT4(v, s) v.x *= s; v.y *= s; v.z *= s; v.w *= s;\n#define ADD4(v1, v2) v1.x += v2.x; v1.y += v2.y; v1.z += v2.z; v1.w += v2.w;\n\n__global__ void integration(data gpu) {\n int i = K_INDEX();\n if(i < N) {\n MULT4(gpu.f[i], 0.01);\n MULT4(gpu.v[i], 0.01);\n ADD4(gpu.v[i], gpu.f[i]);\n ADD4(gpu.r[i], gpu.v[i]);\n gpu.f[i] = make_float4(0,0,0,0);\n }\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/samples/intro.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#include \n\n// CHECK: #include \n#include \n\n#define TOKEN_PASTE(X, Y) X ## Y\n#define ARG_LIST_AS_MACRO a, device_x, device_y\n#define KERNEL_CALL_AS_MACRO axpy<<<1, kDataLen>>>\n#define KERNEL_NAME_MACRO axpy\n\n// CHECK: #define COMPLETE_LAUNCH hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, device_x, device_y)\n#define COMPLETE_LAUNCH axpy<<<1, kDataLen>>>(a, device_x, device_y)\n\n\ntemplate\n__global__ void axpy(T a, T *x, T *y) {\n y[threadIdx.x] = a * x[threadIdx.x];\n}\n\n__global__ void empty() {\n}\n\nint main(int argc, char* argv[]) {\n const int kDataLen = 4;\n\n float a = 2.0f;\n float host_x[kDataLen] = {1.0f, 2.0f, 3.0f, 4.0f};\n float host_y[kDataLen];\n\n // Copy input data to device.\n float* device_x;\n float* device_y;\n\n // CHECK: hipMalloc(&device_x, kDataLen * sizeof(float));\n cudaMalloc(&device_x, kDataLen * sizeof(float));\n\n#ifdef HERRING\n // CHECK: hipMalloc(&device_y, kDataLen * sizeof(float));\n cudaMalloc(&device_y, kDataLen * sizeof(float));\n#else\n // CHECK: hipMalloc(&device_y, kDataLen * sizeof(double));\n cudaMalloc(&device_y, kDataLen * sizeof(double));\n#endif\n\n // CHECK: hipMemcpy(device_x, host_x, kDataLen * sizeof(float), hipMemcpyHostToDevice);\n cudaMemcpy(device_x, host_x, kDataLen * sizeof(float), cudaMemcpyHostToDevice);\n\n // Launch the kernel in numerous different strange ways to exercise the prerocessor.\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, device_x, device_y);\n axpy<<<1, kDataLen>>>(a, device_x, device_y);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, device_x, device_y);\n axpy<<<1, kDataLen>>>(a, device_x, device_y);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, TOKEN_PASTE(device, _x), device_y);\n axpy<<<1, kDataLen>>>(a, TOKEN_PASTE(device, _x), device_y);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, ARG_LIST_AS_MACRO);\n axpy<<<1, kDataLen>>>(ARG_LIST_AS_MACRO);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(KERNEL_NAME_MACRO), dim3(1), dim3(kDataLen), 0, 0, ARG_LIST_AS_MACRO);\n KERNEL_NAME_MACRO<<<1, kDataLen>>>(ARG_LIST_AS_MACRO);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, ARG_LIST_AS_MACRO);\n KERNEL_CALL_AS_MACRO(ARG_LIST_AS_MACRO);\n\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), 0, 0);\n empty<<<1, kDataLen>>> ( );\n\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), 0, 0);\n empty<<<1, kDataLen, 0>>>();\n\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), 0, 0);\n empty<<<1, kDataLen, 0, 0>>>();\n\n // CHECK: COMPLETE_LAUNCH;\n COMPLETE_LAUNCH;\n\n\n // Copy output data to host.\n // CHECK: hipDeviceSynchronize();\n cudaDeviceSynchronize();\n\n // CHECK: hipMemcpy(host_y, device_y, kDataLen * sizeof(float), hipMemcpyDeviceToHost);\n cudaMemcpy(host_y, device_y, kDataLen * sizeof(float), cudaMemcpyDeviceToHost);\n\n // Print the results.\n for (int i = 0; i < kDataLen; ++i) {\n std::cout << \"y[\" << i << \"] = \" << host_y[i] << \"\\n\";\n }\n\n // CHECK: hipDeviceReset();\n cudaDeviceReset();\n return 0;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#include \n\n// CHECK: #include \n#include \n\n#define TOKEN_PASTE(X, Y) X ## Y\n#define ARG_LIST_AS_MACRO a, device_x, device_y\n#define KERNEL_CALL_AS_MACRO axpy<<<1, kDataLen>>>\n#define KERNEL_NAME_MACRO axpy\n\n// CHECK: #define COMPLETE_LAUNCH hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, device_x, device_y)\n#define COMPLETE_LAUNCH axpy<<<1, kDataLen>>>(a, device_x, device_y)\n\n\ntemplate\n__global__ void axpy(T a, T *x, T *y) {\n y[threadIdx.x] = a * x[threadIdx.x];\n}\n\n__global__ void empty() {\n}\n\nint main(int argc, char* argv[]) {\n const int kDataLen = 4;\n\n float a = 2.0f;\n float host_x[kDataLen] = {1.0f, 2.0f, 3.0f, 4.0f};\n float host_y[kDataLen];\n\n // Copy input data to device.\n float* device_x;\n float* device_y;\n\n // CHECK: hipMalloc(&device_x, kDataLen * sizeof(float));\n hipMalloc(&device_x, kDataLen * sizeof(float));\n\n#ifdef HERRING\n // CHECK: hipMalloc(&device_y, kDataLen * sizeof(float));\n hipMalloc(&device_y, kDataLen * sizeof(float));\n#else\n // CHECK: hipMalloc(&device_y, kDataLen * sizeof(double));\n hipMalloc(&device_y, kDataLen * sizeof(double));\n#endif\n\n // CHECK: hipMemcpy(device_x, host_x, kDataLen * sizeof(float), hipMemcpyHostToDevice);\n hipMemcpy(device_x, host_x, kDataLen * sizeof(float), hipMemcpyHostToDevice);\n\n // Launch the kernel in numerous different strange ways to exercise the prerocessor.\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, device_x, device_y);\n axpy<<<1, kDataLen>>>(a, device_x, device_y);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, device_x, device_y);\n axpy<<<1, kDataLen>>>(a, device_x, device_y);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, TOKEN_PASTE(device, _x), device_y);\n axpy<<<1, kDataLen>>>(a, TOKEN_PASTE(device, _x), device_y);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, ARG_LIST_AS_MACRO);\n axpy<<<1, kDataLen>>>(ARG_LIST_AS_MACRO);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(KERNEL_NAME_MACRO), dim3(1), dim3(kDataLen), 0, 0, ARG_LIST_AS_MACRO);\n KERNEL_NAME_MACRO<<<1, kDataLen>>>(ARG_LIST_AS_MACRO);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, ARG_LIST_AS_MACRO);\n KERNEL_CALL_AS_MACRO(ARG_LIST_AS_MACRO);\n\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), 0, 0);\n empty<<<1, kDataLen>>> ( );\n\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), 0, 0);\n empty<<<1, kDataLen, 0>>>();\n\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), 0, 0);\n empty<<<1, kDataLen, 0, 0>>>();\n\n // CHECK: COMPLETE_LAUNCH;\n COMPLETE_LAUNCH;\n\n\n // Copy output data to host.\n // CHECK: hipDeviceSynchronize();\n hipDeviceSynchronize();\n\n // CHECK: hipMemcpy(host_y, device_y, kDataLen * sizeof(float), hipMemcpyDeviceToHost);\n hipMemcpy(host_y, device_y, kDataLen * sizeof(float), hipMemcpyDeviceToHost);\n\n // Print the results.\n for (int i = 0; i < kDataLen; ++i) {\n std::cout << \"y[\" << i << \"] = \" << host_y[i] << \"\\n\";\n }\n\n // CHECK: hipDeviceReset();\n hipDeviceReset();\n return 0;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/samples/axpy.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n/*\nCopyright (c) 2015-2016 Advanced Micro Devices, Inc. All rights reserved.\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in\nall copies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\nTHE SOFTWARE.\n*/\n#include \n#include \n\n#define CHECK(cmd) \\\n{\\\n cudaError_t error = cmd;\\\n if (error != cudaSuccess) { \\\n fprintf(stderr, \"error: '%s'(%d) at %s:%d\\n\", cudaGetErrorString(error), error,__FILE__, __LINE__); \\\n exit(EXIT_FAILURE);\\\n\t }\\\n}\n\n\n/*\n * Square each element in the array A and write to array C.\n */\ntemplate \n__global__ void\nvector_square(T *C_d, const T *A_d, size_t N)\n{\n size_t offset = (blockIdx.x * blockDim.x + threadIdx.x);\n size_t stride = blockDim.x * gridDim.x;\n\n for (size_t i=offset; i>> (C_d, A_d, N);\n\n printf (\"info: copy Device2Host\\n\");\n // CHECK: CHECK ( hipMemcpy(C_h, C_d, Nbytes, hipMemcpyDeviceToHost));\n CHECK ( cudaMemcpy(C_h, C_d, Nbytes, cudaMemcpyDeviceToHost));\n\n printf (\"info: check result\\n\");\n for (size_t i=0; i\n#include \n\n#define CHECK(cmd) \\\n{\\\n hipError_t error = cmd;\\\n if (error != hipSuccess) { \\\n fprintf(stderr, \"error: '%s'(%d) at %s:%d\\n\", hipGetErrorString(error), error,__FILE__, __LINE__); \\\n exit(EXIT_FAILURE);\\\n\t }\\\n}\n\n\n/*\n * Square each element in the array A and write to array C.\n */\ntemplate \n__global__ void\nvector_square(T *C_d, const T *A_d, size_t N)\n{\n size_t offset = (blockIdx.x * blockDim.x + threadIdx.x);\n size_t stride = blockDim.x * gridDim.x;\n\n for (size_t i=offset; i>> (C_d, A_d, N);\n\n printf (\"info: copy Device2Host\\n\");\n // CHECK: CHECK ( hipMemcpy(C_h, C_d, Nbytes, hipMemcpyDeviceToHost));\n CHECK ( hipMemcpy(C_h, C_d, Nbytes, hipMemcpyDeviceToHost));\n\n printf (\"info: check result\\n\");\n for (size_t i=0; i\n#include \n#include \n// CHECK: #include \n#include \n\n#define BlockSize 16 // Size of blocks, 32 x 32 threads, fixed, used globally\n\n__global__ void gpu_Comput (int *h, int N, int T) {\n\n// Array loaded with global thread ID that acesses that location\n\n\tint col = threadIdx.x + blockDim.x * blockIdx.x; \n\tint row = threadIdx.y + blockDim.y * blockIdx.y;\n\n\tint threadID = col + row * N;\n\tint index = row + col * N; // sequentially down each row\n\n\tfor (int t = 0; t < T; t++) // loop to repeat to reduce other time effects\n\t h[index] = threadID; // load array with flattened global thread ID\n}\n\nvoid printArray(int *h, int N) {\n\n\tprintf(\"Results of computation, every N/8 numbers, eight numbers\\n\");\n\n\tfor (int row = 0; row < N; row += N/8) {\n\t for (int col = 0; col < N; col += N/8)\n\t printf(\"%6d \", h[col + row * N]);\n\t printf(\"\\n\");\n\t}\n}\n\nint main(int argc, char *argv[]) {\n\n\tint T = 100; // number of iterations, entered at keyboard\n\tint B = 1; // number of blocks, entered at keyboard\n\tchar key;\n\n\tint *h, *dev_h; // ptr to array holding numbers on host and device\n // CHECK: hipEvent_t start, stop;\n\tcudaEvent_t start, stop; // cuda events to measure time\n\tfloat elapsed_time_ms1;\n // CHECK: hipEventCreate( &start );\n // CHECK: hipEventCreate( &stop );\n\tcudaEventCreate( &start );\n\tcudaEventCreate( &stop );\n\n/* ------------------------- Keyboard input -----------------------------------*/\n\ndo { // loop to repeat complete program\n\n\tprintf(\"Grid Structure 2-D grid, 2-D blocks\\n\");\n\tprintf(\"Blocks fixed at 16 x 16 threads, 512 threads, max for compute cap. 1.x\\n\");\n\tprintf(\"Enter number of blocks in grid, each dimension, currently %d\\n\",B);\n\tscanf(\"%d\",&B);\n\tprintf(\"Enter number of iterations, currently %d\\n\",T);\n\tscanf(\"%d\",&T);\n\n\tint N = B * BlockSize; // size of data array, given input data\n\n\tprintf(\"Array size (and total grid-block size) %d x %d\\n\", N, N);\n\n\tdim3 Block(BlockSize, BlockSize); //Block structure, 32 x 32 max\n\tdim3 Grid(B, B); //Grid structure, B x B\n\n/* ------------------------- Allocate Memory-----------------------------------*/\n\n\tint size = N * N * sizeof(int); // number of bytes in total in array\n\th = (int*) malloc(size); // Array on host\n // CHECK: hipMalloc((void**)&dev_h, size);\n\tcudaMalloc((void**)&dev_h, size); // allocate device memory\n\n/* ------------------------- GPU Computation -----------------------------------*/\n\n // CHECK: hipEventRecord( start, 0 );\n\tcudaEventRecord( start, 0 );\n // CHECK: hipLaunchKernelGGL(gpu_Comput, dim3(Grid), dim3(Block), 0, 0, dev_h, N, T);\n\tgpu_Comput<<< Grid, Block >>>(dev_h, N, T);\n // CHECK: hipEventRecord( stop, 0 );\n // CHECK: hipEventSynchronize( stop );\n // CHECK: hipEventElapsedTime( &elapsed_time_ms1, start, stop );\n\tcudaEventRecord( stop, 0 ); // instrument code to measue end time\n\tcudaEventSynchronize( stop ); // wait for all work done by threads\n\tcudaEventElapsedTime( &elapsed_time_ms1, start, stop );\n // CHECK: hipMemcpy(h,dev_h, size ,hipMemcpyDeviceToHost);\n\tcudaMemcpy(h,dev_h, size ,cudaMemcpyDeviceToHost); //Get results to check\n\n\tprintArray(h,N);\n\tprintf(\"\\nTime to calculate results on GPU: %f ms.\\n\", elapsed_time_ms1);\n\n/* -------------------------REPEAT PROGRAM INPUT-----------------------------------*/\n\n\tprintf(\"\\nEnter c to repeat, return to terminate\\n\");\n\n\tscanf(\"%c\",&key);\n\tscanf(\"%c\",&key);\n\n} while (key == 'c'); // loop of complete program\n\n/* -------------- clean up ---------------------------------------*/\n\nfree(h);\n // CHECK: hipFree(dev_h);\n\tcudaFree(dev_h);\n // CHECK: hipEventDestroy(start);\n // CHECK: hipEventDestroy(stop);\n\tcudaEventDestroy(start);\n\tcudaEventDestroy(stop);\n\n\treturn 0;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n// To measure effects of memory coalescing. Coalescing.cu\n// B. Wilkinson Jan 30, 2011\n\n#include \n#include \n#include \n// CHECK: #include \n#include \n\n#define BlockSize 16 // Size of blocks, 32 x 32 threads, fixed, used globally\n\n__global__ void gpu_Comput (int *h, int N, int T) {\n\n// Array loaded with global thread ID that acesses that location\n\n\tint col = threadIdx.x + blockDim.x * blockIdx.x; \n\tint row = threadIdx.y + blockDim.y * blockIdx.y;\n\n\tint threadID = col + row * N;\n\tint index = row + col * N; // sequentially down each row\n\n\tfor (int t = 0; t < T; t++) // loop to repeat to reduce other time effects\n\t h[index] = threadID; // load array with flattened global thread ID\n}\n\nvoid printArray(int *h, int N) {\n\n\tprintf(\"Results of computation, every N/8 numbers, eight numbers\\n\");\n\n\tfor (int row = 0; row < N; row += N/8) {\n\t for (int col = 0; col < N; col += N/8)\n\t printf(\"%6d \", h[col + row * N]);\n\t printf(\"\\n\");\n\t}\n}\n\nint main(int argc, char *argv[]) {\n\n\tint T = 100; // number of iterations, entered at keyboard\n\tint B = 1; // number of blocks, entered at keyboard\n\tchar key;\n\n\tint *h, *dev_h; // ptr to array holding numbers on host and device\n // CHECK: hipEvent_t start, stop;\n\thipEvent_t start, stop; // cuda events to measure time\n\tfloat elapsed_time_ms1;\n // CHECK: hipEventCreate( &start );\n // CHECK: hipEventCreate( &stop );\n\thipEventCreate( &start );\n\thipEventCreate( &stop );\n\n/* ------------------------- Keyboard input -----------------------------------*/\n\ndo { // loop to repeat complete program\n\n\tprintf(\"Grid Structure 2-D grid, 2-D blocks\\n\");\n\tprintf(\"Blocks fixed at 16 x 16 threads, 512 threads, max for compute cap. 1.x\\n\");\n\tprintf(\"Enter number of blocks in grid, each dimension, currently %d\\n\",B);\n\tscanf(\"%d\",&B);\n\tprintf(\"Enter number of iterations, currently %d\\n\",T);\n\tscanf(\"%d\",&T);\n\n\tint N = B * BlockSize; // size of data array, given input data\n\n\tprintf(\"Array size (and total grid-block size) %d x %d\\n\", N, N);\n\n\tdim3 Block(BlockSize, BlockSize); //Block structure, 32 x 32 max\n\tdim3 Grid(B, B); //Grid structure, B x B\n\n/* ------------------------- Allocate Memory-----------------------------------*/\n\n\tint size = N * N * sizeof(int); // number of bytes in total in array\n\th = (int*) malloc(size); // Array on host\n // CHECK: hipMalloc((void**)&dev_h, size);\n\thipMalloc((void**)&dev_h, size); // allocate device memory\n\n/* ------------------------- GPU Computation -----------------------------------*/\n\n // CHECK: hipEventRecord( start, 0 );\n\thipEventRecord( start, 0 );\n // CHECK: hipLaunchKernelGGL(gpu_Comput, dim3(Grid), dim3(Block), 0, 0, dev_h, N, T);\n\tgpu_Comput<<< Grid, Block >>>(dev_h, N, T);\n // CHECK: hipEventRecord( stop, 0 );\n // CHECK: hipEventSynchronize( stop );\n // CHECK: hipEventElapsedTime( &elapsed_time_ms1, start, stop );\n\thipEventRecord( stop, 0 ); // instrument code to measue end time\n\thipEventSynchronize( stop ); // wait for all work done by threads\n\thipEventElapsedTime( &elapsed_time_ms1, start, stop );\n // CHECK: hipMemcpy(h,dev_h, size ,hipMemcpyDeviceToHost);\n\thipMemcpy(h,dev_h, size ,hipMemcpyDeviceToHost); //Get results to check\n\n\tprintArray(h,N);\n\tprintf(\"\\nTime to calculate results on GPU: %f ms.\\n\", elapsed_time_ms1);\n\n/* -------------------------REPEAT PROGRAM INPUT-----------------------------------*/\n\n\tprintf(\"\\nEnter c to repeat, return to terminate\\n\");\n\n\tscanf(\"%c\",&key);\n\tscanf(\"%c\",&key);\n\n} while (key == 'c'); // loop of complete program\n\n/* -------------- clean up ---------------------------------------*/\n\nfree(h);\n // CHECK: hipFree(dev_h);\n\thipFree(dev_h);\n // CHECK: hipEventDestroy(start);\n // CHECK: hipEventDestroy(stop);\n\thipEventDestroy(start);\n\thipEventDestroy(stop);\n\n\treturn 0;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/samples/coalescing.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#include \n#include \n#include \n// CHECK: #include \n#include \n\n#define K_THREADS 64\n#define K_INDEX() ((gridDim.x * blockIdx.y + blockIdx.x) * blockDim.x + threadIdx.x)\n#define RND() ((rand() & 0x7FFF) / float(0x8000))\n#define ERRORCHECK() cErrorCheck(__FILE__, __LINE__)\n\n// CHECK: hipEvent_t t##_start, t##_end; \\\n// CHECK: hipEventCreate(&t##_start); \\\n// CHECK: hipEventCreate(&t##_end);\n#define TIMER_CREATE(t) \\\n cudaEvent_t t##_start, t##_end; \\\n cudaEventCreate(&t##_start); \\\n cudaEventCreate(&t##_end);\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start);\n#define TIMER_START(t) \\\n cudaEventRecord(t##_start); \\\n cudaEventSynchronize(t##_start); \\\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start); \\\n// CHECK: hipEventRecord(t##_end); \\\n// CHECK: hipEventSynchronize(t##_end); \\\n// CHECK: hipEventElapsedTime(&t, t##_start, t##_end);\n#define TIMER_END(t) \\\n cudaEventRecord(t##_start); \\\n cudaEventSynchronize(t##_start); \\\n cudaEventRecord(t##_end); \\\n cudaEventSynchronize(t##_end); \\\n cudaEventElapsedTime(&t, t##_start, t##_end); \n\n\ninline void cErrorCheck(const char *file, int line) {\n// CHECK: hipDeviceSynchronize();\n// CHECK: hipError_t err = hipGetLastError();\n// CHECK: if (err != hipSuccess) {\n// CHECK: printf(\"Error: %s\\n\", hipGetErrorString(err));\n cudaThreadSynchronize();\n cudaError_t err = cudaGetLastError();\n if (err != cudaSuccess) {\n printf(\"Error: %s\\n\", cudaGetErrorString(err));\n printf(\" @ %s: %d\\n\", file, line);\n exit(-1);\n }\n}\n\ninline dim3 K_GRID(int n, int threads = K_THREADS) {\n int blocks = (int)ceilf(sqrtf((float)n/threads));\n dim3 grid(blocks, blocks);\n return grid;\n}\n\ntypedef struct data {\n int n;\n float4 *r, *v, *f;\n} data;\n\ndata cpu, gpu;\n\n#define N 20\n\n__global__ void repulsion(data gpu);\n__global__ void integration(data gpu);\n\n\nint main() {\n printf(\"Cuda Test 1\\n\");\n\n int count = 0;\n // CHECK: hipGetDeviceCount(&count);\n cudaGetDeviceCount(&count);\n printf(\" %d CUDA devices found\\n\", count);\n if(!count) {\n ::exit(EXIT_FAILURE);\n }\n // CHECK: hipFree(0);\n cudaFree(0);\n\n cpu.n = N;\n\n cpu.r = (float4*)malloc(N * sizeof(float4));\n cpu.v = (float4*)malloc(N * sizeof(float4));\n cpu.f = (float4*)malloc(N * sizeof(float4));\n\n for(int i = 0; i < N; ++i) {\n cpu.v[i] = make_float4(0,0,0,0);\n cpu.r[i] = make_float4(RND(), RND(), RND(), 0);\n cpu.f[i] = make_float4(0,0.01,0,0);\n }\n\n gpu = cpu;\n // CHECK: hipMalloc(&gpu.r, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.v, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.f, N * sizeof(float4));\n cudaMalloc(&gpu.r, N * sizeof(float4));\n cudaMalloc(&gpu.v, N * sizeof(float4));\n cudaMalloc(&gpu.f, N * sizeof(float4));\n // CHECK: hipMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n cudaMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), cudaMemcpyHostToDevice);\n cudaMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), cudaMemcpyHostToDevice);\n cudaMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), cudaMemcpyHostToDevice);\n\n ERRORCHECK();\n float rep;\n TIMER_CREATE(rep);\n TIMER_START(rep);\n // CHECK: hipLaunchKernelGGL(integration, dim3(K_GRID(cpu.n)), dim3(K_THREADS), 0, 0, gpu);\n integration <<< K_GRID(cpu.n), K_THREADS >>>(gpu);\n\n TIMER_END(rep);\n printf(\"Took: %f ms\\n\", rep);\n ERRORCHECK();\n // CHECK: hipMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n cudaMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), cudaMemcpyDeviceToHost);\n cudaMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), cudaMemcpyDeviceToHost);\n cudaMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), cudaMemcpyDeviceToHost);\n // CHECK: hipHostFree(cpu.r);\n // CHECK: hipHostFree(cpu.v);\n // CHECK: hipHostFree(cpu.f);\n cudaFreeHost(cpu.r);\n cudaFreeHost(cpu.v);\n cudaFreeHost(cpu.f);\n // CHECK: hipFree(gpu.r);\n // CHECK: hipFree(gpu.v);\n // CHECK: hipFree(gpu.f);\n cudaFree(gpu.r);\n cudaFree(gpu.v);\n cudaFree(gpu.f);\n // CHECK: hipDeviceReset();\n cudaDeviceReset();\n\n printf(\"Results: \\n\");\n for(int i = 0; i < N; ++i) {\n printf(\"%f, %f, %f \\n\", cpu.r[i].x, cpu.r[i].y, cpu.r[i].z);\n }\n\n printf(\"Ready...\\n\");\n return 0;\n}\n\n__global__ void repulsion(data gpu) {\n int idx = K_INDEX();\n if(idx < N) {\n gpu.r[idx].x = 1;\n gpu.r[idx].y = 1;\n gpu.r[idx].z = 1;\n }\n}\n\n#define MULT4(v, s) v.x *= s; v.y *= s; v.z *= s; v.w *= s;\n#define ADD4(v1, v2) v1.x += v2.x; v1.y += v2.y; v1.z += v2.z; v1.w += v2.w;\n\n__global__ void integration(data gpu) {\n int i = K_INDEX();\n if(i < N) {\n MULT4(gpu.f[i], 0.01);\n MULT4(gpu.v[i], 0.01);\n ADD4(gpu.v[i], gpu.f[i]);\n ADD4(gpu.r[i], gpu.v[i]);\n gpu.f[i] = make_float4(0,0,0,0);\n }\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#include \n#include \n#include \n// CHECK: #include \n#include \n\n#define K_THREADS 64\n#define K_INDEX() ((gridDim.x * blockIdx.y + blockIdx.x) * blockDim.x + threadIdx.x)\n#define RND() ((rand() & 0x7FFF) / float(0x8000))\n#define ERRORCHECK() cErrorCheck(__FILE__, __LINE__)\n\n// CHECK: hipEvent_t t##_start, t##_end; \\\n// CHECK: hipEventCreate(&t##_start); \\\n// CHECK: hipEventCreate(&t##_end);\n#define TIMER_CREATE(t) \\\n hipEvent_t t##_start, t##_end; \\\n hipEventCreate(&t##_start); \\\n hipEventCreate(&t##_end);\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start);\n#define TIMER_START(t) \\\n hipEventRecord(t##_start); \\\n hipEventSynchronize(t##_start); \\\n\n// CHECK: hipEventRecord(t##_start); \\\n// CHECK: hipEventSynchronize(t##_start); \\\n// CHECK: hipEventRecord(t##_end); \\\n// CHECK: hipEventSynchronize(t##_end); \\\n// CHECK: hipEventElapsedTime(&t, t##_start, t##_end);\n#define TIMER_END(t) \\\n hipEventRecord(t##_start); \\\n hipEventSynchronize(t##_start); \\\n hipEventRecord(t##_end); \\\n hipEventSynchronize(t##_end); \\\n hipEventElapsedTime(&t, t##_start, t##_end); \n\n\ninline void cErrorCheck(const char *file, int line) {\n// CHECK: hipDeviceSynchronize();\n// CHECK: hipError_t err = hipGetLastError();\n// CHECK: if (err != hipSuccess) {\n// CHECK: printf(\"Error: %s\\n\", hipGetErrorString(err));\n hipDeviceSynchronize();\n hipError_t err = hipGetLastError();\n if (err != hipSuccess) {\n printf(\"Error: %s\\n\", hipGetErrorString(err));\n printf(\" @ %s: %d\\n\", file, line);\n exit(-1);\n }\n}\n\ninline dim3 K_GRID(int n, int threads = K_THREADS) {\n int blocks = (int)ceilf(sqrtf((float)n/threads));\n dim3 grid(blocks, blocks);\n return grid;\n}\n\ntypedef struct data {\n int n;\n float4 *r, *v, *f;\n} data;\n\ndata cpu, gpu;\n\n#define N 20\n\n__global__ void repulsion(data gpu);\n__global__ void integration(data gpu);\n\n\nint main() {\n printf(\"Cuda Test 1\\n\");\n\n int count = 0;\n // CHECK: hipGetDeviceCount(&count);\n hipGetDeviceCount(&count);\n printf(\" %d CUDA devices found\\n\", count);\n if(!count) {\n ::exit(EXIT_FAILURE);\n }\n // CHECK: hipFree(0);\n hipFree(0);\n\n cpu.n = N;\n\n cpu.r = (float4*)malloc(N * sizeof(float4));\n cpu.v = (float4*)malloc(N * sizeof(float4));\n cpu.f = (float4*)malloc(N * sizeof(float4));\n\n for(int i = 0; i < N; ++i) {\n cpu.v[i] = make_float4(0,0,0,0);\n cpu.r[i] = make_float4(RND(), RND(), RND(), 0);\n cpu.f[i] = make_float4(0,0.01,0,0);\n }\n\n gpu = cpu;\n // CHECK: hipMalloc(&gpu.r, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.v, N * sizeof(float4));\n // CHECK: hipMalloc(&gpu.f, N * sizeof(float4));\n hipMalloc(&gpu.r, N * sizeof(float4));\n hipMalloc(&gpu.v, N * sizeof(float4));\n hipMalloc(&gpu.f, N * sizeof(float4));\n // CHECK: hipMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n hipMemcpy(gpu.r, cpu.r, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n hipMemcpy(gpu.v, cpu.v, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n hipMemcpy(gpu.f, cpu.f, cpu.n * sizeof(float4), hipMemcpyHostToDevice);\n\n ERRORCHECK();\n float rep;\n TIMER_CREATE(rep);\n TIMER_START(rep);\n // CHECK: hipLaunchKernelGGL(integration, dim3(K_GRID(cpu.n)), dim3(K_THREADS), 0, 0, gpu);\n integration <<< K_GRID(cpu.n), K_THREADS >>>(gpu);\n\n TIMER_END(rep);\n printf(\"Took: %f ms\\n\", rep);\n ERRORCHECK();\n // CHECK: hipMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n hipMemcpy(cpu.r, gpu.r, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n hipMemcpy(cpu.v, gpu.v, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n hipMemcpy(cpu.f, gpu.f, cpu.n * sizeof(float4), hipMemcpyDeviceToHost);\n // CHECK: hipHostFree(cpu.r);\n // CHECK: hipHostFree(cpu.v);\n // CHECK: hipHostFree(cpu.f);\n hipHostFree(cpu.r);\n hipHostFree(cpu.v);\n hipHostFree(cpu.f);\n // CHECK: hipFree(gpu.r);\n // CHECK: hipFree(gpu.v);\n // CHECK: hipFree(gpu.f);\n hipFree(gpu.r);\n hipFree(gpu.v);\n hipFree(gpu.f);\n // CHECK: hipDeviceReset();\n hipDeviceReset();\n\n printf(\"Results: \\n\");\n for(int i = 0; i < N; ++i) {\n printf(\"%f, %f, %f \\n\", cpu.r[i].x, cpu.r[i].y, cpu.r[i].z);\n }\n\n printf(\"Ready...\\n\");\n return 0;\n}\n\n__global__ void repulsion(data gpu) {\n int idx = K_INDEX();\n if(idx < N) {\n gpu.r[idx].x = 1;\n gpu.r[idx].y = 1;\n gpu.r[idx].z = 1;\n }\n}\n\n#define MULT4(v, s) v.x *= s; v.y *= s; v.z *= s; v.w *= s;\n#define ADD4(v1, v2) v1.x += v2.x; v1.y += v2.y; v1.z += v2.z; v1.w += v2.w;\n\n__global__ void integration(data gpu) {\n int i = K_INDEX();\n if(i < N) {\n MULT4(gpu.f[i], 0.01);\n MULT4(gpu.v[i], 0.01);\n ADD4(gpu.v[i], gpu.f[i]);\n ADD4(gpu.r[i], gpu.v[i]);\n gpu.f[i] = make_float4(0,0,0,0);\n }\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/samples/intro_before_13000.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n// Kernel definition\n__global__ void vecAdd(float* A, float* B, float* C)\n{\n int i = threadIdx.x;\n A[i] = 0;\n B[i] = i;\n C[i] = A[i] + B[i];\n}\n// CHECK: #include \n#include \n#define SIZE 10\n#define KERNELINVOKES 5000000\nint vecadd(int gpudevice, int rank)\n{\n int devcheck(int, int);\n devcheck(gpudevice, rank);\n float A[SIZE], B[SIZE], C[SIZE];\n // Kernel invocation\n float *devPtrA;\n float *devPtrB;\n float *devPtrC;\n int memsize = SIZE * sizeof(float);\n // CHECK: hipMalloc((void**)&devPtrA, memsize);\n // CHECK: hipMalloc((void**)&devPtrB, memsize);\n // CHECK: hipMalloc((void**)&devPtrC, memsize);\n cudaMalloc((void**)&devPtrA, memsize);\n cudaMalloc((void**)&devPtrB, memsize);\n cudaMalloc((void**)&devPtrC, memsize);\n // CHECK: hipMemcpy(devPtrA, A, memsize, hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(devPtrB, B, memsize, hipMemcpyHostToDevice);\n cudaMemcpy(devPtrA, A, memsize, cudaMemcpyHostToDevice);\n cudaMemcpy(devPtrB, B, memsize, cudaMemcpyHostToDevice);\n for (int i = 0; i>>(devPtrA, devPtrB, devPtrC);\n }\n // CHECK: hipMemcpy(C, devPtrC, memsize, hipMemcpyDeviceToHost);\n cudaMemcpy(C, devPtrC, memsize, cudaMemcpyDeviceToHost);\n // calculate only up to gpudevice to show the unique output\n // of each rank's kernel launch\n for (int i = 0; i= device_count)\n {\n printf(\"gpudevice >= device_count ... exiting\\n\");\n exit(1);\n }\n // CHECK: hipError_t cudareturn;\n // CHECK: hipDeviceProp_t deviceProp;\n // CHECK: hipGetDeviceProperties(&deviceProp, gpudevice);\n cudaError_t cudareturn;\n cudaDeviceProp deviceProp;\n cudaGetDeviceProperties(&deviceProp, gpudevice);\n if (deviceProp.warpSize <= 1)\n {\n printf(\"rank %d: warning, CUDA Device Emulation (CPU) detected, exiting\\n\", rank);\n exit(1);\n }\n // CHECK: cudareturn = hipSetDevice(gpudevice);\n cudareturn = cudaSetDevice(gpudevice);\n // CHECK: if (cudareturn == hipErrorInvalidDevice)\n if (cudareturn == cudaErrorInvalidDevice)\n {\n // CHECK: perror(\"hipSetDevice returned hipErrorInvalidDevice\");\n perror(\"cudaSetDevice returned cudaErrorInvalidDevice\");\n }\n else\n {\n // CHECK: hipGetDevice(&device);\n cudaGetDevice(&device);\n printf(\"rank %d: cudaGetDevice()=%d\\n\", rank, device);\n }\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n// Kernel definition\n#include \"hip/hip_runtime.h\"\n__global__ void vecAdd(float* A, float* B, float* C)\n{\n int i = threadIdx.x;\n A[i] = 0;\n B[i] = i;\n C[i] = A[i] + B[i];\n}\n// CHECK: #include \n#include \n#define SIZE 10\n#define KERNELINVOKES 5000000\nint vecadd(int gpudevice, int rank)\n{\n int devcheck(int, int);\n devcheck(gpudevice, rank);\n float A[SIZE], B[SIZE], C[SIZE];\n // Kernel invocation\n float *devPtrA;\n float *devPtrB;\n float *devPtrC;\n int memsize = SIZE * sizeof(float);\n // CHECK: hipMalloc((void**)&devPtrA, memsize);\n // CHECK: hipMalloc((void**)&devPtrB, memsize);\n // CHECK: hipMalloc((void**)&devPtrC, memsize);\n hipMalloc((void**)&devPtrA, memsize);\n hipMalloc((void**)&devPtrB, memsize);\n hipMalloc((void**)&devPtrC, memsize);\n // CHECK: hipMemcpy(devPtrA, A, memsize, hipMemcpyHostToDevice);\n // CHECK: hipMemcpy(devPtrB, B, memsize, hipMemcpyHostToDevice);\n hipMemcpy(devPtrA, A, memsize, hipMemcpyHostToDevice);\n hipMemcpy(devPtrB, B, memsize, hipMemcpyHostToDevice);\n for (int i = 0; i>>(devPtrA, devPtrB, devPtrC);\n }\n // CHECK: hipMemcpy(C, devPtrC, memsize, hipMemcpyDeviceToHost);\n hipMemcpy(C, devPtrC, memsize, hipMemcpyDeviceToHost);\n // calculate only up to gpudevice to show the unique output\n // of each rank's kernel launch\n for (int i = 0; i= device_count)\n {\n printf(\"gpudevice >= device_count ... exiting\\n\");\n exit(1);\n }\n // CHECK: hipError_t cudareturn;\n // CHECK: hipDeviceProp_t deviceProp;\n // CHECK: hipGetDeviceProperties(&deviceProp, gpudevice);\n hipError_t cudareturn;\n hipDeviceProp_t deviceProp;\n hipGetDeviceProperties(&deviceProp, gpudevice);\n if (deviceProp.warpSize <= 1)\n {\n printf(\"rank %d: warning, CUDA Device Emulation (CPU) detected, exiting\\n\", rank);\n exit(1);\n }\n // CHECK: cudareturn = hipSetDevice(gpudevice);\n cudareturn = hipSetDevice(gpudevice);\n // CHECK: if (cudareturn == hipErrorInvalidDevice)\n if (cudareturn == hipErrorInvalidDevice)\n {\n // CHECK: perror(\"hipSetDevice returned hipErrorInvalidDevice\");\n perror(\"hipSetDevice returned hipErrorInvalidDevice\");\n }\n else\n {\n // CHECK: hipGetDevice(&device);\n hipGetDevice(&device);\n printf(\"rank %d: hipGetDevice()=%d\\n\", rank, device);\n }\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/samples/vec_add.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args %clang_args\n/*\nCopyright (c) 2015-present Advanced Micro Devices, Inc. All rights reserved.\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in\nall copies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\nTHE SOFTWARE.\n*/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n// CHECK: #include \n#include \n\nvoid check_error(void) {\n // CHECK: hipError_t err = hipGetLastError();\n cudaError_t err = cudaGetLastError();\n // CHECK: if (err != hipSuccess) {\n if (err != cudaSuccess) {\n std::cerr\n << \"Error: \"\n // CHECK: << hipGetErrorString(err)\n << cudaGetErrorString(err)\n << std::endl;\n exit(err);\n }\n}\n\n__global__ void atomic_reduction_kernel(int *in, int *out, int ARRAYSIZE) {\n int sum = int(0);\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n for (int i = idx; i < ARRAYSIZE; i += blockDim.x * gridDim.x) {\n sum += in[i];\n }\n // CHECK: atomicAdd(out, sum);\n atomicAdd(out, sum);\n}\n\n__global__ void atomic_reduction_kernel2(int *in, int *out, int ARRAYSIZE) {\n int sum = int(0);\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n for(int i = idx * 16; i < ARRAYSIZE; i += blockDim.x * gridDim.x * 16) {\n sum += in[i] + in[i+1] + in[i+2] + in[i+3] + in[i+4] + in[i+5] + in[i+6] + in[i+7] + in[i+8] + in[i+9] + in[i+10]\n + in[i+11] + in[i+12] + in[i+13] + in[i+14] + in[i+15];\n }\n // CHECK: atomicAdd(out, sum);\n atomicAdd(out, sum);\n}\n\n__global__ void atomic_reduction_kernel3(int *in, int *out, int ARRAYSIZE) {\n int sum = int(0);\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n for(int i = idx * 4; i < ARRAYSIZE; i += blockDim.x * gridDim.x * 4) {\n sum += in[i] + in[i+1] + in[i+2] + in[i+3];\n }\n // CHECK: atomicAdd(out, sum);\n atomicAdd(out, sum);\n}\n\nint main(int argc, char **argv) {\n unsigned int ARRAYSIZE = 52428800;\n if (argc < 2) {\n printf(\"Usage: ./reduction num_of_elems\\n\");\n printf(\"using default value: %d\\n\", ARRAYSIZE);\n } else\n ARRAYSIZE = atoi(argv[1]);\n int N = 10;\n printf(\"ARRAYSIZE: %d\\n\", ARRAYSIZE);\n std::cout << \"Array size: \" << ARRAYSIZE * sizeof(int) / 1024.0 / 1024.0 << \" MB\" << std::endl;\n int *Array = (int*)malloc(ARRAYSIZE * sizeof(int));\n int checksum = 0; \n for (int i = 0; i < ARRAYSIZE; ++i) {\n Array[i] = rand()%2;\n checksum += Array[i];\n }\n int *in, *out;\n std::chrono::high_resolution_clock::time_point t1, t2;\n long long size = sizeof(int) * ARRAYSIZE;\n // CHECK: hipMalloc(&in, size);\n cudaMalloc(&in, size);\n // CHECK: hipMalloc(&out, sizeof(int));\n cudaMalloc(&out, sizeof(int));\n check_error();\n // CHECK: hipMemcpy(in, Array, ARRAYSIZE * sizeof(int), hipMemcpyHostToDevice);\n cudaMemcpy(in, Array, ARRAYSIZE * sizeof(int), cudaMemcpyHostToDevice);\n // CHECK: hipDeviceSynchronize();\n cudaDeviceSynchronize();\n check_error();\n // CHECK: hipDeviceProp_t props;\n cudaDeviceProp props;\n // CHECK: hipGetDeviceProperties(&props, 0);\n cudaGetDeviceProperties(&props, 0);\n int threads = 256;\n int blocks = std::min((ARRAYSIZE + threads - 1) / threads, 2048u);\n t1 = std::chrono::high_resolution_clock::now();\n for(int i = 0; i < N; ++i) {\n // CHECK: hipMemsetAsync(out, 0, sizeof(int));\n cudaMemsetAsync(out, 0, sizeof(int));\n // CHECK: atomic_reduction_kernel<<>>(in, out, ARRAYSIZE);\n atomic_reduction_kernel<<>>(in, out, ARRAYSIZE);\n // CHECK: atomic_reduction_kernel2<<>>(in, out, ARRAYSIZE);\n atomic_reduction_kernel2<<>>(in, out, ARRAYSIZE);\n // CHECK: atomic_reduction_kernel3<<>>(in, out, ARRAYSIZE);\n atomic_reduction_kernel3<<>>(in, out, ARRAYSIZE);\n check_error();\n // CHECK: hipDeviceSynchronize();\n cudaDeviceSynchronize();\n check_error();\n }\n t2 = std::chrono::high_resolution_clock::now();\n double times = std::chrono::duration_cast>(t2 - t1).count();\n float GB = (float)ARRAYSIZE * sizeof(int) * N;\n std::cout << \"The average performance of reduction is \" << 1.0E-09 * GB / times << \" GBytes/sec\" << std::endl;\n int sum;\n // CHECK: hipMemcpy(&sum, out, sizeof(int), hipMemcpyDeviceToHost);\n cudaMemcpy(&sum, out, sizeof(int), cudaMemcpyDeviceToHost);\n check_error();\n if(sum == checksum)\n std::cout << \"VERIFICATION: result is CORRECT\" << std::endl << std::endl;\n else\n std::cout << \"VERIFICATION: result is INCORRECT!!\" << std::endl << std::endl;\n // CHECK: hipFree(in);\n cudaFree(in);\n // CHECK: hipFree(out);\n cudaFree(out);\n check_error();\n free(Array);\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args %clang_args\n/*\nCopyright (c) 2015-present Advanced Micro Devices, Inc. All rights reserved.\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in\nall copies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\nTHE SOFTWARE.\n*/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n// CHECK: #include \n#include \n\nvoid check_error(void) {\n // CHECK: hipError_t err = hipGetLastError();\n hipError_t err = hipGetLastError();\n // CHECK: if (err != hipSuccess) {\n if (err != hipSuccess) {\n std::cerr\n << \"Error: \"\n // CHECK: << hipGetErrorString(err)\n << hipGetErrorString(err)\n << std::endl;\n exit(err);\n }\n}\n\n__global__ void atomic_reduction_kernel(int *in, int *out, int ARRAYSIZE) {\n int sum = int(0);\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n for (int i = idx; i < ARRAYSIZE; i += blockDim.x * gridDim.x) {\n sum += in[i];\n }\n // CHECK: atomicAdd(out, sum);\n atomicAdd(out, sum);\n}\n\n__global__ void atomic_reduction_kernel2(int *in, int *out, int ARRAYSIZE) {\n int sum = int(0);\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n for(int i = idx * 16; i < ARRAYSIZE; i += blockDim.x * gridDim.x * 16) {\n sum += in[i] + in[i+1] + in[i+2] + in[i+3] + in[i+4] + in[i+5] + in[i+6] + in[i+7] + in[i+8] + in[i+9] + in[i+10]\n + in[i+11] + in[i+12] + in[i+13] + in[i+14] + in[i+15];\n }\n // CHECK: atomicAdd(out, sum);\n atomicAdd(out, sum);\n}\n\n__global__ void atomic_reduction_kernel3(int *in, int *out, int ARRAYSIZE) {\n int sum = int(0);\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n for(int i = idx * 4; i < ARRAYSIZE; i += blockDim.x * gridDim.x * 4) {\n sum += in[i] + in[i+1] + in[i+2] + in[i+3];\n }\n // CHECK: atomicAdd(out, sum);\n atomicAdd(out, sum);\n}\n\nint main(int argc, char **argv) {\n unsigned int ARRAYSIZE = 52428800;\n if (argc < 2) {\n printf(\"Usage: ./reduction num_of_elems\\n\");\n printf(\"using default value: %d\\n\", ARRAYSIZE);\n } else\n ARRAYSIZE = atoi(argv[1]);\n int N = 10;\n printf(\"ARRAYSIZE: %d\\n\", ARRAYSIZE);\n std::cout << \"Array size: \" << ARRAYSIZE * sizeof(int) / 1024.0 / 1024.0 << \" MB\" << std::endl;\n int *Array = (int*)malloc(ARRAYSIZE * sizeof(int));\n int checksum = 0; \n for (int i = 0; i < ARRAYSIZE; ++i) {\n Array[i] = rand()%2;\n checksum += Array[i];\n }\n int *in, *out;\n std::chrono::high_resolution_clock::time_point t1, t2;\n long long size = sizeof(int) * ARRAYSIZE;\n // CHECK: hipMalloc(&in, size);\n hipMalloc(&in, size);\n // CHECK: hipMalloc(&out, sizeof(int));\n hipMalloc(&out, sizeof(int));\n check_error();\n // CHECK: hipMemcpy(in, Array, ARRAYSIZE * sizeof(int), hipMemcpyHostToDevice);\n hipMemcpy(in, Array, ARRAYSIZE * sizeof(int), hipMemcpyHostToDevice);\n // CHECK: hipDeviceSynchronize();\n hipDeviceSynchronize();\n check_error();\n // CHECK: hipDeviceProp_t props;\n hipDeviceProp_t props;\n // CHECK: hipGetDeviceProperties(&props, 0);\n hipGetDeviceProperties(&props, 0);\n int threads = 256;\n int blocks = std::min((ARRAYSIZE + threads - 1) / threads, 2048u);\n t1 = std::chrono::high_resolution_clock::now();\n for(int i = 0; i < N; ++i) {\n // CHECK: hipMemsetAsync(out, 0, sizeof(int));\n hipMemsetAsync(out, 0, sizeof(int));\n // CHECK: atomic_reduction_kernel<<>>(in, out, ARRAYSIZE);\n atomic_reduction_kernel<<>>(in, out, ARRAYSIZE);\n // CHECK: atomic_reduction_kernel2<<>>(in, out, ARRAYSIZE);\n atomic_reduction_kernel2<<>>(in, out, ARRAYSIZE);\n // CHECK: atomic_reduction_kernel3<<>>(in, out, ARRAYSIZE);\n atomic_reduction_kernel3<<>>(in, out, ARRAYSIZE);\n check_error();\n // CHECK: hipDeviceSynchronize();\n hipDeviceSynchronize();\n check_error();\n }\n t2 = std::chrono::high_resolution_clock::now();\n double times = std::chrono::duration_cast>(t2 - t1).count();\n float GB = (float)ARRAYSIZE * sizeof(int) * N;\n std::cout << \"The average performance of reduction is \" << 1.0E-09 * GB / times << \" GBytes/sec\" << std::endl;\n int sum;\n // CHECK: hipMemcpy(&sum, out, sizeof(int), hipMemcpyDeviceToHost);\n hipMemcpy(&sum, out, sizeof(int), hipMemcpyDeviceToHost);\n check_error();\n if(sum == checksum)\n std::cout << \"VERIFICATION: result is CORRECT\" << std::endl << std::endl;\n else\n std::cout << \"VERIFICATION: result is INCORRECT!!\" << std::endl << std::endl;\n // CHECK: hipFree(in);\n hipFree(in);\n // CHECK: hipFree(out);\n hipFree(out);\n check_error();\n free(Array);\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/samples/reduction.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#pragma once\n// CHECK: #include \n#include \n#include \n\n/**\n * Allocate GPU memory for `count` elements of type `T`.\n */\ntemplate\nstatic T* gpuMalloc(size_t count) {\n T* ret = nullptr;\n // CHECK: hipMalloc(&ret, count * sizeof(T));\n cudaMalloc(&ret, count * sizeof(T));\n return ret;\n}\n\ntemplate\n__global__ void add(int n, T* x, T* y) {\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = index; i < n; i += stride)\n y[i] = x[i] + y[i];\n}\n\nint main(int argc, char* argv[]) {\n size_t numElements = 50;\n float *A = gpuMalloc(numElements);\n float* B = gpuMalloc(numElements);\n for (int i = 0; i < numElements; ++i) {\n A[i] = 1.0f;\n B[i] = 2.0f;\n }\n int blockSize = 512;\n int numBlocks = (numElements + blockSize - 1) / blockSize;\n dim3 dimGrid(numBlocks, 1, 1);\n dim3 dimBlock(blockSize, 1, 1);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(add), dim3(dimGrid), dim3(dimBlock), 0, 0, numElements, A, B);\n add<<>>(numElements, A, B);\n // CHECK: hipDeviceSynchronize();\n cudaDeviceSynchronize();\n float maxError = 0.0f;\n for (int i = 0; i < numElements; ++i)\n maxError = fmax(maxError, fabs(B[i] - 3.0f));\n // CHECK: hipFree(A);\n cudaFree(A);\n // CHECK: hipFree(B);\n cudaFree(B);\n if (maxError == 0.0f)\n return 0;\n return -1;\n }\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#pragma once\n// CHECK: #include \n#include \n#include \n\n/**\n * Allocate GPU memory for `count` elements of type `T`.\n */\ntemplate\nstatic T* gpuMalloc(size_t count) {\n T* ret = nullptr;\n // CHECK: hipMalloc(&ret, count * sizeof(T));\n hipMalloc(&ret, count * sizeof(T));\n return ret;\n}\n\ntemplate\n__global__ void add(int n, T* x, T* y) {\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = index; i < n; i += stride)\n y[i] = x[i] + y[i];\n}\n\nint main(int argc, char* argv[]) {\n size_t numElements = 50;\n float *A = gpuMalloc(numElements);\n float* B = gpuMalloc(numElements);\n for (int i = 0; i < numElements; ++i) {\n A[i] = 1.0f;\n B[i] = 2.0f;\n }\n int blockSize = 512;\n int numBlocks = (numElements + blockSize - 1) / blockSize;\n dim3 dimGrid(numBlocks, 1, 1);\n dim3 dimBlock(blockSize, 1, 1);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(add), dim3(dimGrid), dim3(dimBlock), 0, 0, numElements, A, B);\n add<<>>(numElements, A, B);\n // CHECK: hipDeviceSynchronize();\n hipDeviceSynchronize();\n float maxError = 0.0f;\n for (int i = 0; i < numElements; ++i)\n maxError = fmax(maxError, fabs(B[i] - 3.0f));\n // CHECK: hipFree(A);\n hipFree(A);\n // CHECK: hipFree(B);\n hipFree(B);\n if (maxError == 0.0f)\n return 0;\n return -1;\n }\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/samples/allocators.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#include \n#include \n\n// CHECK: #include \n#include \n\ntemplate\n__global__ void axpy(T a, T *x, T *y) {\n y[threadIdx.x] = a * x[threadIdx.x];\n}\n\ntemplate\n__global__ void axpy_2(T1 a, T2 *x, T2 *y) {\n y[threadIdx.x] = a * x[threadIdx.x];\n}\n\ntemplate\n__global__ void axpy_empty() {\n}\n\n__global__ void empty() {\n}\n\n__global__ void nonempty(int x, int y, int z) {\n}\n\nint main(int argc, char* argv[]) {\n const int kDataLen = 4;\n\n float a = 2.0f;\n float host_x[kDataLen] = {1.0f, 2.0f, 3.0f, 4.0f};\n float host_y[kDataLen];\n\n // Copy input data to device.\n float* device_x;\n float* device_y;\n\n // CHECK: hipMalloc(&device_x, kDataLen * sizeof(float));\n cudaMalloc(&device_x, kDataLen * sizeof(float));\n // CHECK: hipMalloc(&device_y, kDataLen * sizeof(float));\n cudaMalloc(&device_y, kDataLen * sizeof(float));\n // CHECK: hipMemcpy(device_x, host_x, kDataLen * sizeof(float), hipMemcpyHostToDevice);\n cudaMemcpy(device_x, host_x, kDataLen * sizeof(float), cudaMemcpyHostToDevice);\n\n int x = 1, y = 2, z = 3;\n size_t N = 32;\n // CHECK: hipStream_t stream = NULL;\n cudaStream_t stream = NULL;\n // CHECK: hipStreamCreate(&stream);\n cudaStreamCreate(&stream);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, device_x, device_y);\n axpy<<<1, kDataLen>>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, device_x, device_y);\n axpy<<>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, device_x, device_y);\n axpy<<<1, dim3(kDataLen)>>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, device_x, device_y);\n axpy<<>>(a, device_x, device_y);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, 0, a, device_x, device_y);\n axpy<<<1, kDataLen, N>>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, 0, a, device_x, device_y);\n axpy<<>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, 0, a, device_x, device_y);\n axpy<<<1, dim3(kDataLen), N>>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, 0, a, device_x, device_y);\n axpy<<>>(a, device_x, device_y);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, stream, a, device_x, device_y);\n axpy<<<1, kDataLen, N, stream>>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, stream, a, device_x, device_y);\n axpy<<>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, stream, a, device_x, device_y);\n axpy<<<1, dim3(kDataLen), N, stream>>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, stream, a, device_x, device_y);\n axpy<<>>(a, device_x, device_y);\n\n double h_x[kDataLen] = {1.0f, 2.0f, 3.0f, 4.0f};\n double h_y[kDataLen];\n\n // Copy input data to device.\n double* d_x;\n double* d_y;\n\n // CHECK: hipMalloc(&d_x, kDataLen * sizeof(double));\n cudaMalloc(&d_x, kDataLen * sizeof(double));\n // CHECK: hipMalloc(&d_y, kDataLen * sizeof(double));\n cudaMalloc(&d_y, kDataLen * sizeof(double));\n // CHECK: hipMemcpy(d_x, h_x, kDataLen * sizeof(double), hipMemcpyHostToDevice);\n cudaMemcpy(d_x, h_x, kDataLen * sizeof(double), cudaMemcpyHostToDevice);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_2), dim3(1), dim3(kDataLen*2+10), N*N, stream, a, d_x, d_y);\n axpy_2<<<1, kDataLen*2+10, N*N, stream>>>(a, d_x, d_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_2), dim3(1,1,1), dim3(kDataLen*2+10), N*N, stream, a, d_x, d_y);\n axpy_2<<>>(a, d_x, d_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_2), dim3(1), dim3(kDataLen*2+10), N*N, stream, a, d_x, d_y);\n axpy_2<<<1, dim3(kDataLen*2+10), N*N, stream>>>(a, d_x, d_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_2), dim3(1,1,1), dim3(kDataLen*2+10), N*N, stream, a, d_x, d_y);\n axpy_2<<>>(a, d_x, d_y);\n\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), 0, 0);\n axpy_empty<<<1, kDataLen>>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), 0, 0);\n axpy_empty<<>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), 0, 0);\n axpy_empty<<<1, dim3(kDataLen)>>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), 0, 0);\n axpy_empty<<>>();\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, 0);\n axpy_empty<<<1, kDataLen, N>>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, 0);\n axpy_empty<<>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, 0);\n axpy_empty<<<1, dim3(kDataLen), N>>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, 0);\n axpy_empty<<>>();\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, stream);\n axpy_empty<<<1, kDataLen, N, stream>>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, stream);\n axpy_empty<<>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, stream);\n axpy_empty<<<1, dim3(kDataLen), N, stream>>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, stream);\n axpy_empty<<>>();\n\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), 0, 0);\n empty<<<1, kDataLen>>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), 0, 0);\n empty<<>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), 0, 0);\n empty<<<1, dim3(kDataLen)>>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), 0, 0);\n empty<<>> ( );\n\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, 0);\n empty<<<1, kDataLen, N>>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, 0);\n empty<<>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, 0);\n empty<<<1, dim3(kDataLen), N>>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, 0);\n empty<<>> ( );\n\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, stream);\n empty<<<1, kDataLen, N, stream>>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, stream);\n empty<<>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, stream);\n empty<<<1, dim3(kDataLen), N, stream>>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, stream);\n empty<<>> ( );\n\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), 0, 0, x, y, z);\n nonempty<<<1, kDataLen>>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), 0, 0, x, y, z);\n nonempty<<>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), 0, 0, x, y, z);\n nonempty<<<1, dim3(kDataLen)>>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), 0, 0, x, y, z);\n nonempty<<>> (x, y, z);\n\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, 0, x, y, z);\n nonempty<<<1, kDataLen, N>>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, 0, x, y, z);\n nonempty<<>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, 0, x, y, z);\n nonempty<<<1, dim3(kDataLen), N>>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, 0, x, y, z);\n nonempty<<>> (x, y, z);\n\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, stream, x, y, z);\n nonempty<<<1, kDataLen, N, stream>>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, stream, x, y, z);\n nonempty<<>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, stream, x, y, z);\n nonempty<<<1, dim3(kDataLen), N, stream>>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, stream, x, y, z);\n nonempty<<>> (x, y, z);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_2), dim3(x,y,z), dim3(std::min(kDataLen*2+10,x)), std::min(x,y), stream, a, std::min(d_x,d_y), std::max(d_x,d_y));\n axpy_2<<>>(a, std::min(d_x,d_y), std::max(d_x,d_y));\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_2), dim3(x,y,z), dim3(std::min(kDataLen*2+10,x)), std::min(x,y), 0, a, std::min(d_x,d_y), std::max(d_x,d_y));\n axpy_2<<>>(a, std::min(d_x,d_y), std::max(d_x,d_y));\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_2), dim3(x,y,z), dim3(std::min(kDataLen*2+10,x)), 0, 0, a, std::min(d_x,d_y), std::max(d_x,d_y));\n axpy_2<<>>(a, std::min(d_x,d_y), std::max(d_x,d_y));\n\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(x,y,z), dim3(x,y,std::min(y,z)), 0, 0, x, y, z);\n nonempty<<>>(x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(x,y,z), dim3(x,y,std::min(std::max(x,y),z)), 0, 0, x, y, z);\n nonempty<<>>(x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(x,y,z), dim3(x,y,std::min(std::max(x,int(N)),z)), 0, 0, x, y, z);\n nonempty<<>>(x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(x,y,z), dim3(x,y,std::min(std::max(x,int(N+N -x/y + y*1)),z)), 0, 0, x, y, z);\n nonempty<<>>(x, y, z);\n\n // Copy output data to host.\n // CHECK: hipDeviceSynchronize();\n cudaDeviceSynchronize();\n\n // CHECK: hipMemcpy(host_y, device_y, kDataLen * sizeof(float), hipMemcpyDeviceToHost);\n cudaMemcpy(host_y, device_y, kDataLen * sizeof(float), cudaMemcpyDeviceToHost);\n\n // Print the results.\n for (int i = 0; i < kDataLen; ++i) {\n std::cout << \"y[\" << i << \"] = \" << host_y[i] << \"\\n\";\n }\n\n // CHECK: hipDeviceReset();\n cudaDeviceReset();\n return 0;\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n#include \n#include \n\n// CHECK: #include \n#include \n\ntemplate\n__global__ void axpy(T a, T *x, T *y) {\n y[threadIdx.x] = a * x[threadIdx.x];\n}\n\ntemplate\n__global__ void axpy_2(T1 a, T2 *x, T2 *y) {\n y[threadIdx.x] = a * x[threadIdx.x];\n}\n\ntemplate\n__global__ void axpy_empty() {\n}\n\n__global__ void empty() {\n}\n\n__global__ void nonempty(int x, int y, int z) {\n}\n\nint main(int argc, char* argv[]) {\n const int kDataLen = 4;\n\n float a = 2.0f;\n float host_x[kDataLen] = {1.0f, 2.0f, 3.0f, 4.0f};\n float host_y[kDataLen];\n\n // Copy input data to device.\n float* device_x;\n float* device_y;\n\n // CHECK: hipMalloc(&device_x, kDataLen * sizeof(float));\n hipMalloc(&device_x, kDataLen * sizeof(float));\n // CHECK: hipMalloc(&device_y, kDataLen * sizeof(float));\n hipMalloc(&device_y, kDataLen * sizeof(float));\n // CHECK: hipMemcpy(device_x, host_x, kDataLen * sizeof(float), hipMemcpyHostToDevice);\n hipMemcpy(device_x, host_x, kDataLen * sizeof(float), hipMemcpyHostToDevice);\n\n int x = 1, y = 2, z = 3;\n size_t N = 32;\n // CHECK: hipStream_t stream = NULL;\n hipStream_t stream = NULL;\n // CHECK: hipStreamCreate(&stream);\n hipStreamCreate(&stream);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, device_x, device_y);\n axpy<<<1, kDataLen>>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, device_x, device_y);\n axpy<<>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, device_x, device_y);\n axpy<<<1, dim3(kDataLen)>>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), 0, 0, a, device_x, device_y);\n axpy<<>>(a, device_x, device_y);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, 0, a, device_x, device_y);\n axpy<<<1, kDataLen, N>>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, 0, a, device_x, device_y);\n axpy<<>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, 0, a, device_x, device_y);\n axpy<<<1, dim3(kDataLen), N>>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, 0, a, device_x, device_y);\n axpy<<>>(a, device_x, device_y);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, stream, a, device_x, device_y);\n axpy<<<1, kDataLen, N, stream>>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, stream, a, device_x, device_y);\n axpy<<>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, stream, a, device_x, device_y);\n axpy<<<1, dim3(kDataLen), N, stream>>>(a, device_x, device_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy), dim3(1), dim3(kDataLen), N, stream, a, device_x, device_y);\n axpy<<>>(a, device_x, device_y);\n\n double h_x[kDataLen] = {1.0f, 2.0f, 3.0f, 4.0f};\n double h_y[kDataLen];\n\n // Copy input data to device.\n double* d_x;\n double* d_y;\n\n // CHECK: hipMalloc(&d_x, kDataLen * sizeof(double));\n hipMalloc(&d_x, kDataLen * sizeof(double));\n // CHECK: hipMalloc(&d_y, kDataLen * sizeof(double));\n hipMalloc(&d_y, kDataLen * sizeof(double));\n // CHECK: hipMemcpy(d_x, h_x, kDataLen * sizeof(double), hipMemcpyHostToDevice);\n hipMemcpy(d_x, h_x, kDataLen * sizeof(double), hipMemcpyHostToDevice);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_2), dim3(1), dim3(kDataLen*2+10), N*N, stream, a, d_x, d_y);\n axpy_2<<<1, kDataLen*2+10, N*N, stream>>>(a, d_x, d_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_2), dim3(1,1,1), dim3(kDataLen*2+10), N*N, stream, a, d_x, d_y);\n axpy_2<<>>(a, d_x, d_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_2), dim3(1), dim3(kDataLen*2+10), N*N, stream, a, d_x, d_y);\n axpy_2<<<1, dim3(kDataLen*2+10), N*N, stream>>>(a, d_x, d_y);\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_2), dim3(1,1,1), dim3(kDataLen*2+10), N*N, stream, a, d_x, d_y);\n axpy_2<<>>(a, d_x, d_y);\n\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), 0, 0);\n axpy_empty<<<1, kDataLen>>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), 0, 0);\n axpy_empty<<>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), 0, 0);\n axpy_empty<<<1, dim3(kDataLen)>>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), 0, 0);\n axpy_empty<<>>();\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, 0);\n axpy_empty<<<1, kDataLen, N>>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, 0);\n axpy_empty<<>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, 0);\n axpy_empty<<<1, dim3(kDataLen), N>>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, 0);\n axpy_empty<<>>();\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, stream);\n axpy_empty<<<1, kDataLen, N, stream>>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, stream);\n axpy_empty<<>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, stream);\n axpy_empty<<<1, dim3(kDataLen), N, stream>>>();\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_empty), dim3(1), dim3(kDataLen), N, stream);\n axpy_empty<<>>();\n\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), 0, 0);\n empty<<<1, kDataLen>>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), 0, 0);\n empty<<>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), 0, 0);\n empty<<<1, dim3(kDataLen)>>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), 0, 0);\n empty<<>> ( );\n\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, 0);\n empty<<<1, kDataLen, N>>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, 0);\n empty<<>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, 0);\n empty<<<1, dim3(kDataLen), N>>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, 0);\n empty<<>> ( );\n\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, stream);\n empty<<<1, kDataLen, N, stream>>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, stream);\n empty<<>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, stream);\n empty<<<1, dim3(kDataLen), N, stream>>> ( );\n // CHECK: hipLaunchKernelGGL(empty, dim3(1), dim3(kDataLen), N, stream);\n empty<<>> ( );\n\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), 0, 0, x, y, z);\n nonempty<<<1, kDataLen>>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), 0, 0, x, y, z);\n nonempty<<>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), 0, 0, x, y, z);\n nonempty<<<1, dim3(kDataLen)>>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), 0, 0, x, y, z);\n nonempty<<>> (x, y, z);\n\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, 0, x, y, z);\n nonempty<<<1, kDataLen, N>>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, 0, x, y, z);\n nonempty<<>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, 0, x, y, z);\n nonempty<<<1, dim3(kDataLen), N>>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, 0, x, y, z);\n nonempty<<>> (x, y, z);\n\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, stream, x, y, z);\n nonempty<<<1, kDataLen, N, stream>>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, stream, x, y, z);\n nonempty<<>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, stream, x, y, z);\n nonempty<<<1, dim3(kDataLen), N, stream>>> (x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(1), dim3(kDataLen), N, stream, x, y, z);\n nonempty<<>> (x, y, z);\n\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_2), dim3(x,y,z), dim3(std::min(kDataLen*2+10,x)), std::min(x,y), stream, a, std::min(d_x,d_y), std::max(d_x,d_y));\n axpy_2<<>>(a, std::min(d_x,d_y), std::max(d_x,d_y));\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_2), dim3(x,y,z), dim3(std::min(kDataLen*2+10,x)), std::min(x,y), 0, a, std::min(d_x,d_y), std::max(d_x,d_y));\n axpy_2<<>>(a, std::min(d_x,d_y), std::max(d_x,d_y));\n // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(axpy_2), dim3(x,y,z), dim3(std::min(kDataLen*2+10,x)), 0, 0, a, std::min(d_x,d_y), std::max(d_x,d_y));\n axpy_2<<>>(a, std::min(d_x,d_y), std::max(d_x,d_y));\n\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(x,y,z), dim3(x,y,std::min(y,z)), 0, 0, x, y, z);\n nonempty<<>>(x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(x,y,z), dim3(x,y,std::min(std::max(x,y),z)), 0, 0, x, y, z);\n nonempty<<>>(x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(x,y,z), dim3(x,y,std::min(std::max(x,int(N)),z)), 0, 0, x, y, z);\n nonempty<<>>(x, y, z);\n // CHECK: hipLaunchKernelGGL(nonempty, dim3(x,y,z), dim3(x,y,std::min(std::max(x,int(N+N -x/y + y*1)),z)), 0, 0, x, y, z);\n nonempty<<>>(x, y, z);\n\n // Copy output data to host.\n // CHECK: hipDeviceSynchronize();\n hipDeviceSynchronize();\n\n // CHECK: hipMemcpy(host_y, device_y, kDataLen * sizeof(float), hipMemcpyDeviceToHost);\n hipMemcpy(host_y, device_y, kDataLen * sizeof(float), hipMemcpyDeviceToHost);\n\n // Print the results.\n for (int i = 0; i < kDataLen; ++i) {\n std::cout << \"y[\" << i << \"] = \" << host_y[i] << \"\\n\";\n }\n\n // CHECK: hipDeviceReset();\n hipDeviceReset();\n return 0;\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/kernel_launch/kernel_launch_syntax.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// Synthetic test to warn only on device functions umin and umax as unsupported, but not on user defined ones.\n// ToDo: change lit testing in order to parse the output.\n\n#define LEN 1024\n#define SIZE LEN * sizeof(float)\n#define ITER 1024*1024\n\n// CHECK: #include \n#include \n\n#define CUDA_LAUNCH(cuda_call,dimGrid,dimBlock, ...) \\\n cuda_call<<>>(__VA_ARGS__);\n\n__global__ void Inc1(float *Ad, float *Bd) {\n int tx = threadIdx.x + blockIdx.x * blockDim.x;\n if (tx < 1) {\n for (int i = 0; i < ITER; ++i) {\n Ad[tx] = Ad[tx] + 1.0f;\n for (int j = 0; j < 256; ++j) {\n Bd[tx] = Ad[tx];\n }\n }\n }\n}\n\nint main() {\n float *A, *Ad, *Bd;\n A = new float[LEN];\n for (int i = 0; i < LEN; ++i) {\n A[i] = 0.0f;\n }\n // CHECK: hipError_t status;\n cudaError_t status;\n // CHECK: status = hipHostRegister(A, SIZE, hipHostRegisterMapped);\n status = cudaHostRegister(A, SIZE, cudaHostRegisterMapped);\n // CHECK: hipHostGetDevicePointer(&Ad, A, 0);\n cudaHostGetDevicePointer(&Ad, A, 0);\n // CHECK: hipMalloc((void**)&Bd, SIZE);\n cudaMalloc((void**)&Bd, SIZE);\n dim3 dimGrid(LEN / 512, 1, 1);\n dim3 dimBlock(512, 1, 1);\n\n // CHECK: hipLaunchKernelGGL(Inc1, dim3(dimGrid), dim3(dimBlock), 0, 0, Ad, Bd);\n CUDA_LAUNCH(Inc1, dimGrid, dimBlock, Ad, Bd);\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n// Synthetic test to warn only on device functions umin and umax as unsupported, but not on user defined ones.\n// ToDo: change lit testing in order to parse the output.\n\n#include \"hip/hip_runtime.h\"\n#define LEN 1024\n#define SIZE LEN * sizeof(float)\n#define ITER 1024*1024\n\n// CHECK: #include \n#include \n\n#define CUDA_LAUNCH(cuda_call,dimGrid,dimBlock, ...) \\\n cuda_call<<>>(__VA_ARGS__);\n\n__global__ void Inc1(float *Ad, float *Bd) {\n int tx = threadIdx.x + blockIdx.x * blockDim.x;\n if (tx < 1) {\n for (int i = 0; i < ITER; ++i) {\n Ad[tx] = Ad[tx] + 1.0f;\n for (int j = 0; j < 256; ++j) {\n Bd[tx] = Ad[tx];\n }\n }\n }\n}\n\nint main() {\n float *A, *Ad, *Bd;\n A = new float[LEN];\n for (int i = 0; i < LEN; ++i) {\n A[i] = 0.0f;\n }\n // CHECK: hipError_t status;\n hipError_t status;\n // CHECK: status = hipHostRegister(A, SIZE, hipHostRegisterMapped);\n status = hipHostRegister(A, SIZE, hipHostRegisterMapped);\n // CHECK: hipHostGetDevicePointer(&Ad, A, 0);\n hipHostGetDevicePointer(&Ad, A, 0);\n // CHECK: hipMalloc((void**)&Bd, SIZE);\n hipMalloc((void**)&Bd, SIZE);\n dim3 dimGrid(LEN / 512, 1, 1);\n dim3 dimBlock(512, 1, 1);\n\n // CHECK: hipLaunchKernelGGL(Inc1, dim3(dimGrid), dim3(dimBlock), 0, 0, Ad, Bd);\n CUDA_LAUNCH(Inc1, dimGrid, dimBlock, Ad, Bd);\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/kernel_launch/kernel_launch_01.cu"} {"input": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n// CHECK: #include \n#include \n#include \n\n__global__ void my_kernel(int arg) {}\n\nvoid test_launches() {\n struct { int x, y; } config = {1, 2};\n dim3 grid_arr[2] = {dim3(1), dim3(2)};\n int shared_mem = 0;\n\n // CHECK: hipStream_t stream = 0;\n cudaStream_t stream = 0;\n\n // 1. Array indexing\n // CHECK: hipLaunchKernelGGL(my_kernel, dim3(grid_arr[0]), dim3(grid_arr[1]), 0, 0, 0);\n my_kernel<<>>(0);\n\n // 2. Member access\n // CHECK: hipLaunchKernelGGL(my_kernel, dim3(config.x), dim3(config.y), 0, 0, 0);\n my_kernel<<>>(0);\n\n // 3. Math, Logical, and Ternary operators\n // CHECK: hipLaunchKernelGGL(my_kernel, dim3(config.x == 1 ? 2 : 1), dim3(256), 0, 0, 0);\n my_kernel<<>>(0);\n\n // 4. Address-of operator\n // CHECK: hipLaunchKernelGGL(my_kernel, dim3(1), dim3(256), shared_mem, &stream, 0);\n my_kernel<<<1, 256, shared_mem, &stream>>>(0);\n\n // 5. Std calls\n // CHECK: hipLaunchKernelGGL(my_kernel, dim3(1), dim3(std::max(config.x, config.y)), shared_mem, 0, std::min(config.x, config.y));\n my_kernel<<<1, std::max(config.x, config.y), shared_mem>>>(std::min(config.x, config.y));\n}\n", "output": "// RUN: %run_test hipify \"%s\" \"%t\" %hipify_args 1 --hip-kernel-execution-syntax %clang_args\n\n// CHECK: #include \n#include \n#include \n\n__global__ void my_kernel(int arg) {}\n\nvoid test_launches() {\n struct { int x, y; } config = {1, 2};\n dim3 grid_arr[2] = {dim3(1), dim3(2)};\n int shared_mem = 0;\n\n // CHECK: hipStream_t stream = 0;\n hipStream_t stream = 0;\n\n // 1. Array indexing\n // CHECK: hipLaunchKernelGGL(my_kernel, dim3(grid_arr[0]), dim3(grid_arr[1]), 0, 0, 0);\n my_kernel<<>>(0);\n\n // 2. Member access\n // CHECK: hipLaunchKernelGGL(my_kernel, dim3(config.x), dim3(config.y), 0, 0, 0);\n my_kernel<<>>(0);\n\n // 3. Math, Logical, and Ternary operators\n // CHECK: hipLaunchKernelGGL(my_kernel, dim3(config.x == 1 ? 2 : 1), dim3(256), 0, 0, 0);\n my_kernel<<>>(0);\n\n // 4. Address-of operator\n // CHECK: hipLaunchKernelGGL(my_kernel, dim3(1), dim3(256), shared_mem, &stream, 0);\n my_kernel<<<1, 256, shared_mem, &stream>>>(0);\n\n // 5. Std calls\n // CHECK: hipLaunchKernelGGL(my_kernel, dim3(1), dim3(std::max(config.x, config.y)), shared_mem, 0, std::min(config.x, config.y));\n my_kernel<<<1, std::max(config.x, config.y), shared_mem>>>(std::min(config.x, config.y));\n}\n", "source": "/kaggle/working/HIPIFY/tests/unit_tests/kernel_launch/complex_kernel_launches.cu"} {"input": "#include \n#include \n#include \n\n#include \"../../utils/utils.cuh\"\n\n#define M 4096 // Rows of A, Rows of C\n#define N 8192 // Columns of B, Columns of C\n#define K 1024 // Columns of A, Rows of B\n#define NGPUS 2 // Number of GPUs to use for computation\n\n// Kernel to perform matrix multiplication on a portion of the matrices\n__global__ void matMulKernelTP(int *A, int *B, int *C, int m, int n, int k, int col_start, int col_size) {\n int row = blockIdx.y * blockDim.y + threadIdx.y;\n int local_col = blockIdx.x * blockDim.x + threadIdx.x;\n int col = col_start + local_col;\n \n if (row < m && local_col < col_size) {\n int value = 0;\n for (int i = 0; i < k; i++) {\n value += A[row * k + i] * B[i * n + col];\n }\n C[row * col_size + local_col] = value;\n }\n}\n\n// Host function for matrix multiplication\nvoid matMulHost(const int* A, const int* B, int* C, int m, int n, int k) {\n for (int i = 0; i < m; i++) {\n for (int j = 0; j < n; j++) {\n int sum = 0;\n for (int p = 0; p < k; p++) {\n sum += A[i * k + p] * B[p * n + j];\n }\n C[i * n + j] = sum;\n }\n }\n}\n\nint main() {\n auto total_start = std::chrono::high_resolution_clock::now();\n // Host matrices\n int *A = new int[M * K];\n int *B = new int[K * N];\n int *C = new int[M * N](); // GPU result\n int *C_host = new int[M * N](); // CPU result\n\n // Initialize matrices with sample values\n for (int i = 0; i < M; i++) {\n for (int j = 0; j < K; j++) {\n A[i * K + j] = (i + j) % 10; // Simple pattern to avoid large numbers in multiplication\n }\n }\n \n for (int i = 0; i < K; i++) {\n for (int j = 0; j < N; j++) {\n B[i * N + j] = (i + j) % 10;\n }\n }\n\n // Calculate split for tensor parallelism across N dimension\n const int cols_per_gpu = N / NGPUS;\n \n // Arrays for each GPU\n int *d_A[NGPUS], *d_B[NGPUS], *d_C[NGPUS];\n\n // Create streams and events per GPU\n cudaStream_t streams[NGPUS];\n cudaEvent_t startEvents[NGPUS], stopEvents[NGPUS];\n\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::cudaSetDevice(gpu));\n cudaCheckError(::cudaStreamCreate(&streams[gpu]));\n cudaCheckError(::cudaEventCreate(&startEvents[gpu]));\n cudaCheckError(::cudaEventCreate(&stopEvents[gpu]));\n }\n\n // Set up each GPU\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::cudaSetDevice(gpu));\n \n // Allocate memory on current GPU\n cudaCheckError(::cudaMalloc(&d_A[gpu], M * K * sizeof(int)));\n cudaCheckError(::cudaMalloc(&d_B[gpu], K * N * sizeof(int)));\n cudaCheckError(::cudaMalloc(&d_C[gpu], M * cols_per_gpu * sizeof(int)));\n \n // Copy input matrices to current GPU\n cudaCheckError(::cudaMemcpyAsync(d_A[gpu], A, M * K * sizeof(int), cudaMemcpyHostToDevice, streams[gpu]));\n cudaCheckError(::cudaMemcpyAsync(d_B[gpu], B, K * N * sizeof(int), cudaMemcpyHostToDevice, streams[gpu]));\n \n // Set grid and block dimensions for this GPU's portion\n const dim3 threadsPerBlock(16, 16); // Using 16x16 thread blocks for better occupancy\n const dim3 numBlocks(\n (cols_per_gpu + threadsPerBlock.x - 1) / threadsPerBlock.x,\n (M + threadsPerBlock.y - 1) / threadsPerBlock.y\n );\n \n const int col_start = gpu * cols_per_gpu;\n\n cudaCheckError(::cudaEventRecord(startEvents[gpu], streams[gpu]));\n\n matMulKernelTP<<>>(\n d_A[gpu], d_B[gpu], d_C[gpu],\n M, N, K, col_start, cols_per_gpu\n );\n\n // Record stop event for this GPU\n cudaCheckError(::cudaEventRecord(stopEvents[gpu], streams[gpu]));\n }\n \n // Copy results back from each GPU and combine\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::cudaSetDevice(gpu));\n // Copy each row's portion separately to maintain correct layout\n for (int row = 0; row < M; row++) {\n cudaCheckError(::cudaMemcpyAsync(\n &C[row * N + gpu * cols_per_gpu], // Destination in host matrix\n &d_C[gpu][row * cols_per_gpu], // Source from GPU\n cols_per_gpu * sizeof(int), // Size of this GPU's portion\n cudaMemcpyDeviceToHost,\n streams[gpu]\n ));\n }\n // Ensure all work in stream is completed before measuring timing.\n cudaCheckError(::cudaStreamSynchronize(streams[gpu]));\n }\n\n // Elapsed time per GPU.\n for (int gpu = 0; gpu < NGPUS; ++gpu) {\n float gpu_time = 0;\n cudaCheckError(::cudaEventElapsedTime(&gpu_time, startEvents[gpu], stopEvents[gpu]));\n std::cout << \"GPU \" << gpu << \" execution time: \" << gpu_time << \"ms\" << std::endl;\n }\n\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::cudaEventDestroy(startEvents[gpu]));\n cudaCheckError(::cudaEventDestroy(stopEvents[gpu]));\n cudaCheckError(::cudaStreamDestroy(streams[gpu]));\n }\n\n // Start CPU timing\n auto cpu_start = std::chrono::high_resolution_clock::now();\n \n // Compute result on CPU for verification\n std::cout << \"Computing on CPU for verification...\" << std::endl;\n matMulHost(A, B, C_host, M, N, K);\n \n // Stop CPU timing\n auto cpu_end = std::chrono::high_resolution_clock::now();\n auto cpu_duration = std::chrono::duration_cast(cpu_end - cpu_start);\n \n // Print timing results\n std::cout << \"\\nTiming Results:\" << std::endl;\n std::cout << \"CPU Time: \" << cpu_duration.count() << \" ms\" << std::endl;\n \n // Compare results\n bool match = true;\n int mismatch_count = 0;\n for (int i = 0; i < M * N; i++) {\n if (C[i] != C_host[i]) {\n match = false;\n mismatch_count++;\n }\n }\n \n // Output comparison results\n std::cout << \"\\nResult comparison:\" << std::endl;\n if (match) {\n std::cout << \"GPU and CPU results match perfectly!\" << std::endl;\n } else {\n std::cout << \"Found \" << mismatch_count << \" mismatches between GPU and CPU results.\" << std::endl;\n }\n \n // Function to print a 4x4 region of both matrices\n auto print_region = [&](const char* region_name, int start_row, int start_col) {\n std::cout << \"\\n\" << region_name << \" (4x4):\" << std::endl;\n std::cout << \"GPU Result:\" << std::endl;\n for (int i = 0; i < 4; i++) {\n for (int j = 0; j < 4; j++) {\n std::cout << C[(start_row + i) * N + (start_col + j)] << \" \";\n }\n std::cout << std::endl;\n }\n \n std::cout << \"CPU Result:\" << std::endl;\n for (int i = 0; i < 4; i++) {\n for (int j = 0; j < 4; j++) {\n std::cout << C_host[(start_row + i) * N + (start_col + j)] << \" \";\n }\n std::cout << std::endl;\n }\n };\n \n // Print different regions of the matrices\n print_region(\"Top-left corner\", 0, 0);\n print_region(\"Middle region\", M/2 - 2, N/2 - 2);\n print_region(\"Bottom-right corner\", M - 4, N - 4);\n \n // Free device memory on each GPU\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::cudaSetDevice(gpu));\n cudaCheckError(::cudaFree(d_A[gpu]));\n cudaCheckError(::cudaFree(d_B[gpu]));\n cudaCheckError(::cudaFree(d_C[gpu]));\n }\n\n // Free host memory\n delete[] A;\n delete[] B;\n delete[] C;\n delete[] C_host;\n \n auto total_end = std::chrono::high_resolution_clock::now();\n auto total_duration = std::chrono::duration_cast(total_end - total_start);\n \n std::cout << \"Total Time: \" << cpu_duration.count() << \" ms\" << std::endl;\n return 0;\n}\n", "output": "#include \n#include \n#include \n\n#include \"../../utils/utils.cuh\"\n\n#define M 4096 // Rows of A, Rows of C\n#define N 8192 // Columns of B, Columns of C\n#define K 1024 // Columns of A, Rows of B\n#define NGPUS 2 // Number of GPUs to use for computation\n\n// Kernel to perform matrix multiplication on a portion of the matrices\n__global__ void matMulKernelTP(int *A, int *B, int *C, int m, int n, int k, int col_start, int col_size) {\n int row = blockIdx.y * blockDim.y + threadIdx.y;\n int local_col = blockIdx.x * blockDim.x + threadIdx.x;\n int col = col_start + local_col;\n \n if (row < m && local_col < col_size) {\n int value = 0;\n for (int i = 0; i < k; i++) {\n value += A[row * k + i] * B[i * n + col];\n }\n C[row * col_size + local_col] = value;\n }\n}\n\n// Host function for matrix multiplication\nvoid matMulHost(const int* A, const int* B, int* C, int m, int n, int k) {\n for (int i = 0; i < m; i++) {\n for (int j = 0; j < n; j++) {\n int sum = 0;\n for (int p = 0; p < k; p++) {\n sum += A[i * k + p] * B[p * n + j];\n }\n C[i * n + j] = sum;\n }\n }\n}\n\nint main() {\n auto total_start = std::chrono::high_resolution_clock::now();\n // Host matrices\n int *A = new int[M * K];\n int *B = new int[K * N];\n int *C = new int[M * N](); // GPU result\n int *C_host = new int[M * N](); // CPU result\n\n // Initialize matrices with sample values\n for (int i = 0; i < M; i++) {\n for (int j = 0; j < K; j++) {\n A[i * K + j] = (i + j) % 10; // Simple pattern to avoid large numbers in multiplication\n }\n }\n \n for (int i = 0; i < K; i++) {\n for (int j = 0; j < N; j++) {\n B[i * N + j] = (i + j) % 10;\n }\n }\n\n // Calculate split for tensor parallelism across N dimension\n const int cols_per_gpu = N / NGPUS;\n \n // Arrays for each GPU\n int *d_A[NGPUS], *d_B[NGPUS], *d_C[NGPUS];\n\n // Create streams and events per GPU\n hipStream_t streams[NGPUS];\n hipEvent_t startEvents[NGPUS], stopEvents[NGPUS];\n\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::hipSetDevice(gpu));\n cudaCheckError(::hipStreamCreate(&streams[gpu]));\n cudaCheckError(::hipEventCreate(&startEvents[gpu]));\n cudaCheckError(::hipEventCreate(&stopEvents[gpu]));\n }\n\n // Set up each GPU\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::hipSetDevice(gpu));\n \n // Allocate memory on current GPU\n cudaCheckError(::hipMalloc(&d_A[gpu], M * K * sizeof(int)));\n cudaCheckError(::hipMalloc(&d_B[gpu], K * N * sizeof(int)));\n cudaCheckError(::hipMalloc(&d_C[gpu], M * cols_per_gpu * sizeof(int)));\n \n // Copy input matrices to current GPU\n cudaCheckError(::hipMemcpyAsync(d_A[gpu], A, M * K * sizeof(int), hipMemcpyHostToDevice, streams[gpu]));\n cudaCheckError(::hipMemcpyAsync(d_B[gpu], B, K * N * sizeof(int), hipMemcpyHostToDevice, streams[gpu]));\n \n // Set grid and block dimensions for this GPU's portion\n const dim3 threadsPerBlock(16, 16); // Using 16x16 thread blocks for better occupancy\n const dim3 numBlocks(\n (cols_per_gpu + threadsPerBlock.x - 1) / threadsPerBlock.x,\n (M + threadsPerBlock.y - 1) / threadsPerBlock.y\n );\n \n const int col_start = gpu * cols_per_gpu;\n\n cudaCheckError(::hipEventRecord(startEvents[gpu], streams[gpu]));\n\n matMulKernelTP<<>>(\n d_A[gpu], d_B[gpu], d_C[gpu],\n M, N, K, col_start, cols_per_gpu\n );\n\n // Record stop event for this GPU\n cudaCheckError(::hipEventRecord(stopEvents[gpu], streams[gpu]));\n }\n \n // Copy results back from each GPU and combine\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::hipSetDevice(gpu));\n // Copy each row's portion separately to maintain correct layout\n for (int row = 0; row < M; row++) {\n cudaCheckError(::hipMemcpyAsync(\n &C[row * N + gpu * cols_per_gpu], // Destination in host matrix\n &d_C[gpu][row * cols_per_gpu], // Source from GPU\n cols_per_gpu * sizeof(int), // Size of this GPU's portion\n hipMemcpyDeviceToHost,\n streams[gpu]\n ));\n }\n // Ensure all work in stream is completed before measuring timing.\n cudaCheckError(::hipStreamSynchronize(streams[gpu]));\n }\n\n // Elapsed time per GPU.\n for (int gpu = 0; gpu < NGPUS; ++gpu) {\n float gpu_time = 0;\n cudaCheckError(::hipEventElapsedTime(&gpu_time, startEvents[gpu], stopEvents[gpu]));\n std::cout << \"GPU \" << gpu << \" execution time: \" << gpu_time << \"ms\" << std::endl;\n }\n\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::hipEventDestroy(startEvents[gpu]));\n cudaCheckError(::hipEventDestroy(stopEvents[gpu]));\n cudaCheckError(::hipStreamDestroy(streams[gpu]));\n }\n\n // Start CPU timing\n auto cpu_start = std::chrono::high_resolution_clock::now();\n \n // Compute result on CPU for verification\n std::cout << \"Computing on CPU for verification...\" << std::endl;\n matMulHost(A, B, C_host, M, N, K);\n \n // Stop CPU timing\n auto cpu_end = std::chrono::high_resolution_clock::now();\n auto cpu_duration = std::chrono::duration_cast(cpu_end - cpu_start);\n \n // Print timing results\n std::cout << \"\\nTiming Results:\" << std::endl;\n std::cout << \"CPU Time: \" << cpu_duration.count() << \" ms\" << std::endl;\n \n // Compare results\n bool match = true;\n int mismatch_count = 0;\n for (int i = 0; i < M * N; i++) {\n if (C[i] != C_host[i]) {\n match = false;\n mismatch_count++;\n }\n }\n \n // Output comparison results\n std::cout << \"\\nResult comparison:\" << std::endl;\n if (match) {\n std::cout << \"GPU and CPU results match perfectly!\" << std::endl;\n } else {\n std::cout << \"Found \" << mismatch_count << \" mismatches between GPU and CPU results.\" << std::endl;\n }\n \n // Function to print a 4x4 region of both matrices\n auto print_region = [&](const char* region_name, int start_row, int start_col) {\n std::cout << \"\\n\" << region_name << \" (4x4):\" << std::endl;\n std::cout << \"GPU Result:\" << std::endl;\n for (int i = 0; i < 4; i++) {\n for (int j = 0; j < 4; j++) {\n std::cout << C[(start_row + i) * N + (start_col + j)] << \" \";\n }\n std::cout << std::endl;\n }\n \n std::cout << \"CPU Result:\" << std::endl;\n for (int i = 0; i < 4; i++) {\n for (int j = 0; j < 4; j++) {\n std::cout << C_host[(start_row + i) * N + (start_col + j)] << \" \";\n }\n std::cout << std::endl;\n }\n };\n \n // Print different regions of the matrices\n print_region(\"Top-left corner\", 0, 0);\n print_region(\"Middle region\", M/2 - 2, N/2 - 2);\n print_region(\"Bottom-right corner\", M - 4, N - 4);\n \n // Free device memory on each GPU\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::hipSetDevice(gpu));\n cudaCheckError(::hipFree(d_A[gpu]));\n cudaCheckError(::hipFree(d_B[gpu]));\n cudaCheckError(::hipFree(d_C[gpu]));\n }\n\n // Free host memory\n delete[] A;\n delete[] B;\n delete[] C;\n delete[] C_host;\n \n auto total_end = std::chrono::high_resolution_clock::now();\n auto total_duration = std::chrono::duration_cast(total_end - total_start);\n \n std::cout << \"Total Time: \" << cpu_duration.count() << \" ms\" << std::endl;\n return 0;\n}\n", "source": "/kaggle/working/cuda_examples/TensorParallelFromScratch/02_matmul_tp/matmul_tp_big.cu"} {"input": "#include \n#include \n\n#include \"../../utils/utils.cuh\"\n\n#define M 2 // Rows of A, Rows of C\n#define N 2 // Columns of B, Columns of C\n#define K 3 // Columns of A, Rows of B\n\n// Kernel to perform matrix multiplication on a portion of the matrices\n__global__ void matMulKernelTP(int *A, int *B, int *C, int m, int n, int k, int col_start, int col_size) {\n int row = blockIdx.y * blockDim.y + threadIdx.y;\n int local_col = blockIdx.x * blockDim.x + threadIdx.x;\n int col = col_start + local_col;\n \n if (row < m && local_col < col_size) {\n int value = 0;\n for (int i = 0; i < k; i++) {\n value += A[row * k + i] * B[i * n + col];\n }\n C[row * col_size + local_col] = value;\n }\n}\n\nint main() {\n // Host matrices\n int A[M][K] = { {1, 2, 3}, {4, 5, 6} };\n int B[K][N] = { {100, 200}, {300, 400}, {500, 600} };\n int C[M][N] = {0};\n\n // Calculate split for tensor parallelism across N dimension\n int cols_per_gpu = N / 2;\n \n // Arrays for each GPU\n int *d_A[2], *d_B[2], *d_C[2];\n \n // Set up each GPU\n for (int gpu = 0; gpu < 2; gpu++) {\n cudaCheckError(::cudaSetDevice(gpu));\n \n // Allocate memory on current GPU\n cudaCheckError(::cudaMalloc(&d_A[gpu], M * K * sizeof(int)));\n cudaCheckError(::cudaMalloc(&d_B[gpu], K * N * sizeof(int)));\n cudaCheckError(::cudaMalloc(&d_C[gpu], M * cols_per_gpu * sizeof(int)));\n \n // Copy input matrices to current GPU\n cudaCheckError(::cudaMemcpy(d_A[gpu], A, M * K * sizeof(int), cudaMemcpyHostToDevice));\n cudaCheckError(::cudaMemcpy(d_B[gpu], B, K * N * sizeof(int), cudaMemcpyHostToDevice));\n \n // Set grid and block dimensions for this GPU's portion\n dim3 threadsPerBlock(cols_per_gpu, M);\n dim3 numBlocks(\n (cols_per_gpu + threadsPerBlock.x - 1) / threadsPerBlock.x,\n (M + threadsPerBlock.y - 1) / threadsPerBlock.y\n );\n \n // Launch kernel for this GPU's portion\n int col_start = gpu * cols_per_gpu;\n matMulKernelTP<<>>(\n d_A[gpu], d_B[gpu], d_C[gpu],\n M, N, K, col_start, cols_per_gpu\n );\n }\n \n // Copy results back from each GPU and combine\n for (int gpu = 0; gpu < 2; gpu++) {\n cudaCheckError(::cudaSetDevice(gpu));\n // Copy each row's portion separately to maintain correct layout\n for (int row = 0; row < M; row++) {\n cudaCheckError(::cudaMemcpy(\n &C[row][gpu * cols_per_gpu], // Destination in host matrix\n &d_C[gpu][row * cols_per_gpu], // Source from GPU\n cols_per_gpu * sizeof(int), // Size of this GPU's portion\n cudaMemcpyDeviceToHost\n ));\n }\n }\n \n // Output the result matrix\n std::cout << \"Result Matrix C:\" << std::endl;\n for (int i = 0; i < M; i++) {\n for (int j = 0; j < N; j++) {\n std::cout << C[i][j] << \" \";\n }\n std::cout << std::endl;\n }\n \n // Free device memory on each GPU\n for (int gpu = 0; gpu < 2; gpu++) {\n cudaCheckError(::cudaSetDevice(gpu));\n cudaCheckError(::cudaFree(d_A[gpu]));\n cudaCheckError(::cudaFree(d_B[gpu]));\n cudaCheckError(::cudaFree(d_C[gpu]));\n }\n \n return 0;\n}\n", "output": "#include \n#include \n\n#include \"../../utils/utils.cuh\"\n\n#define M 2 // Rows of A, Rows of C\n#define N 2 // Columns of B, Columns of C\n#define K 3 // Columns of A, Rows of B\n\n// Kernel to perform matrix multiplication on a portion of the matrices\n__global__ void matMulKernelTP(int *A, int *B, int *C, int m, int n, int k, int col_start, int col_size) {\n int row = blockIdx.y * blockDim.y + threadIdx.y;\n int local_col = blockIdx.x * blockDim.x + threadIdx.x;\n int col = col_start + local_col;\n \n if (row < m && local_col < col_size) {\n int value = 0;\n for (int i = 0; i < k; i++) {\n value += A[row * k + i] * B[i * n + col];\n }\n C[row * col_size + local_col] = value;\n }\n}\n\nint main() {\n // Host matrices\n int A[M][K] = { {1, 2, 3}, {4, 5, 6} };\n int B[K][N] = { {100, 200}, {300, 400}, {500, 600} };\n int C[M][N] = {0};\n\n // Calculate split for tensor parallelism across N dimension\n int cols_per_gpu = N / 2;\n \n // Arrays for each GPU\n int *d_A[2], *d_B[2], *d_C[2];\n \n // Set up each GPU\n for (int gpu = 0; gpu < 2; gpu++) {\n cudaCheckError(::hipSetDevice(gpu));\n \n // Allocate memory on current GPU\n cudaCheckError(::hipMalloc(&d_A[gpu], M * K * sizeof(int)));\n cudaCheckError(::hipMalloc(&d_B[gpu], K * N * sizeof(int)));\n cudaCheckError(::hipMalloc(&d_C[gpu], M * cols_per_gpu * sizeof(int)));\n \n // Copy input matrices to current GPU\n cudaCheckError(::hipMemcpy(d_A[gpu], A, M * K * sizeof(int), hipMemcpyHostToDevice));\n cudaCheckError(::hipMemcpy(d_B[gpu], B, K * N * sizeof(int), hipMemcpyHostToDevice));\n \n // Set grid and block dimensions for this GPU's portion\n dim3 threadsPerBlock(cols_per_gpu, M);\n dim3 numBlocks(\n (cols_per_gpu + threadsPerBlock.x - 1) / threadsPerBlock.x,\n (M + threadsPerBlock.y - 1) / threadsPerBlock.y\n );\n \n // Launch kernel for this GPU's portion\n int col_start = gpu * cols_per_gpu;\n matMulKernelTP<<>>(\n d_A[gpu], d_B[gpu], d_C[gpu],\n M, N, K, col_start, cols_per_gpu\n );\n }\n \n // Copy results back from each GPU and combine\n for (int gpu = 0; gpu < 2; gpu++) {\n cudaCheckError(::hipSetDevice(gpu));\n // Copy each row's portion separately to maintain correct layout\n for (int row = 0; row < M; row++) {\n cudaCheckError(::hipMemcpy(\n &C[row][gpu * cols_per_gpu], // Destination in host matrix\n &d_C[gpu][row * cols_per_gpu], // Source from GPU\n cols_per_gpu * sizeof(int), // Size of this GPU's portion\n hipMemcpyDeviceToHost\n ));\n }\n }\n \n // Output the result matrix\n std::cout << \"Result Matrix C:\" << std::endl;\n for (int i = 0; i < M; i++) {\n for (int j = 0; j < N; j++) {\n std::cout << C[i][j] << \" \";\n }\n std::cout << std::endl;\n }\n \n // Free device memory on each GPU\n for (int gpu = 0; gpu < 2; gpu++) {\n cudaCheckError(::hipSetDevice(gpu));\n cudaCheckError(::hipFree(d_A[gpu]));\n cudaCheckError(::hipFree(d_B[gpu]));\n cudaCheckError(::hipFree(d_C[gpu]));\n }\n \n return 0;\n}\n", "source": "/kaggle/working/cuda_examples/TensorParallelFromScratch/02_matmul_tp/matmul_tp_small.cu"} {"input": "#include \n#include \n\n#include \"../../utils/utils.cuh\"\n\n#define M 2 // Rows of A, Rows of C\n#define N 2 // Columns of B, Columns of C\n#define K 3 // Columns of A, Rows of B\n\n// Kernel to perform our simple matrix multiplication\n__global__ void matMulKernel(int *A, int *B, int *C, int m, int n, int k) {\n int row = blockIdx.y * blockDim.y + threadIdx.y;\n int col = blockIdx.x * blockDim.x + threadIdx.x;\n printf(\"row: %d, col: %d\\n\", row, col);\n if (row < m && col < n) {\n int value = 0;\n for (int i = 0; i < k; i++) {\n value += A[row * k + i] * B[i * n + col];\n }\n C[row * n + col] = value;\n }\n}\n\nint main() {\n // Host matrices\n int A[M][K] = { {1, 2, 3}, {4, 5, 6} };\n int B[K][N] = { {100, 200}, {300, 400}, {500, 600} };\n int C[N][M] = {0};\n\n int *d_A, *d_B, *d_C;\n\n // Allocate device memory\n cudaCheckError(::cudaMalloc(&d_A, M * K * sizeof(int)));\n cudaCheckError(::cudaMalloc(&d_B, K * N * sizeof(int)));\n cudaCheckError(::cudaMalloc(&d_C, M * N * sizeof(int)));\n\n // Copy matrices from host to device\n cudaCheckError(::cudaMemcpy(d_A, A, M * K * sizeof(int), cudaMemcpyHostToDevice));\n cudaCheckError(::cudaMemcpy(d_B, B, K * N * sizeof(int), cudaMemcpyHostToDevice));\n\n // Set our grid and block size\n dim3 threadsPerBlock(M, N);\n dim3 numBlocks((N + threadsPerBlock.x - 1) / threadsPerBlock.x, (M + threadsPerBlock.y - 1) / threadsPerBlock.y);\n\n // Launch the kernel\n matMulKernel<<>>(d_A, d_B, d_C, M, N, K);\n\n // Copy result back to host\n cudaCheckError(::cudaMemcpy(C, d_C, M * N * sizeof(int), cudaMemcpyDeviceToHost));\n\n // Output the result matrix\n std::cout << \"Result Matrix C:\" << std::endl;\n for (int i = 0; i < M; i++) {\n for (int j = 0; j < N; j++) {\n std::cout << C[i][j] << \" \";\n }\n std::cout << std::endl;\n }\n\n // Free device memory\n cudaCheckError(::cudaFree(d_A));\n cudaCheckError(::cudaFree(d_B));\n cudaCheckError(::cudaFree(d_C));\n\n return 0;\n}\n", "output": "#include \n#include \n\n#include \"../../utils/utils.cuh\"\n\n#define M 2 // Rows of A, Rows of C\n#define N 2 // Columns of B, Columns of C\n#define K 3 // Columns of A, Rows of B\n\n// Kernel to perform our simple matrix multiplication\n__global__ void matMulKernel(int *A, int *B, int *C, int m, int n, int k) {\n int row = blockIdx.y * blockDim.y + threadIdx.y;\n int col = blockIdx.x * blockDim.x + threadIdx.x;\n printf(\"row: %d, col: %d\\n\", row, col);\n if (row < m && col < n) {\n int value = 0;\n for (int i = 0; i < k; i++) {\n value += A[row * k + i] * B[i * n + col];\n }\n C[row * n + col] = value;\n }\n}\n\nint main() {\n // Host matrices\n int A[M][K] = { {1, 2, 3}, {4, 5, 6} };\n int B[K][N] = { {100, 200}, {300, 400}, {500, 600} };\n int C[N][M] = {0};\n\n int *d_A, *d_B, *d_C;\n\n // Allocate device memory\n cudaCheckError(::hipMalloc(&d_A, M * K * sizeof(int)));\n cudaCheckError(::hipMalloc(&d_B, K * N * sizeof(int)));\n cudaCheckError(::hipMalloc(&d_C, M * N * sizeof(int)));\n\n // Copy matrices from host to device\n cudaCheckError(::hipMemcpy(d_A, A, M * K * sizeof(int), hipMemcpyHostToDevice));\n cudaCheckError(::hipMemcpy(d_B, B, K * N * sizeof(int), hipMemcpyHostToDevice));\n\n // Set our grid and block size\n dim3 threadsPerBlock(M, N);\n dim3 numBlocks((N + threadsPerBlock.x - 1) / threadsPerBlock.x, (M + threadsPerBlock.y - 1) / threadsPerBlock.y);\n\n // Launch the kernel\n matMulKernel<<>>(d_A, d_B, d_C, M, N, K);\n\n // Copy result back to host\n cudaCheckError(::hipMemcpy(C, d_C, M * N * sizeof(int), hipMemcpyDeviceToHost));\n\n // Output the result matrix\n std::cout << \"Result Matrix C:\" << std::endl;\n for (int i = 0; i < M; i++) {\n for (int j = 0; j < N; j++) {\n std::cout << C[i][j] << \" \";\n }\n std::cout << std::endl;\n }\n\n // Free device memory\n cudaCheckError(::hipFree(d_A));\n cudaCheckError(::hipFree(d_B));\n cudaCheckError(::hipFree(d_C));\n\n return 0;\n}\n", "source": "/kaggle/working/cuda_examples/TensorParallelFromScratch/01_simple_matmul_no_tp/matmul.cu"} {"input": "#include \n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n#define W 1024\n\n__global__ void matMul(float* out, const float* in1, const float* in2, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n float sum = 0.0f;\n if (row < width && col < width) {\n for (int i = 0; i < width; ++i) {\n sum += in1[row * width + i] * in2[i * width + col];\n }\n out[row * width + col] = sum;\n }\n}\n\nvoid doMatMul(float* out, const float* in1, const float* in2, int width) {\n float *dIn1, *dIn2, *dOut;\n const size_t size = width * width * sizeof(float);\n\n printf(\"Allocating device memory...\\n\");\n cudaCheckError(::cudaMalloc(&dIn1, size));\n cudaCheckError(::cudaMalloc(&dIn2, size));\n cudaCheckError(::cudaMalloc(&dOut, size));\n\n printf(\"Copying data to device...\\n\");\n cudaCheckError(::cudaMemcpy(dIn1, in1, size, cudaMemcpyHostToDevice));\n cudaCheckError(::cudaMemcpy(dIn2, in2, size, cudaMemcpyHostToDevice));\n\n dim3 threadsPerBlock(16, 16);\n dim3 blocksPerGrid((width + threadsPerBlock.x - 1) / threadsPerBlock.x,\n (width + threadsPerBlock.y - 1) / threadsPerBlock.y);\n printf(\"Launching kernel...\\n\");\n for (int j = 0; j < 100; ++j)\n {\n nvtxRangePush(\"Matrix multiplication kernel\");\n matMul<<>>(dOut, dIn1, dIn2, width);\n nvtxRangePop();\n }\n cudaCheckError(::cudaDeviceSynchronize());\n\n printf(\"Copying result to host...\\n\");\n cudaCheckError(::cudaMemcpy(out, dOut, size, cudaMemcpyHostToDevice));\n\n printf(\"Freeing device memory...\\n\");\n cudaCheckError(::cudaFree(dIn1));\n cudaCheckError(::cudaFree(dIn2));\n cudaCheckError(::cudaFree(dOut));\n}\n\nint main() {\n float *in1, *in2, *out;\n const size_t width = W * W * sizeof(float);\n cudaCheckError(::cudaMallocHost(&in1, width));\n cudaCheckError(::cudaMallocHost(&in2, width));\n cudaCheckError(::cudaMallocHost(&out, width));\n for (int i = 0; i < W * W; ++i) {\n in1[i] = static_cast(rand()) / RAND_MAX;\n in2[i] = static_cast(rand()) / RAND_MAX;\n }\n printf(\"Starting matrix multiplication...\\n\");\n doMatMul(out, in1, in2, width);\n printf(\"Matrix multiplication completed.\\n\");\n printf(\"Freeing host memory...\\n\");\n cudaCheckError(::cudaFreeHost(in1));\n cudaCheckError(::cudaFreeHost(in2));\n cudaCheckError(::cudaFreeHost(out));\n return 0;\n}\n", "output": "#include \n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n#define W 1024\n\n__global__ void matMul(float* out, const float* in1, const float* in2, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n float sum = 0.0f;\n if (row < width && col < width) {\n for (int i = 0; i < width; ++i) {\n sum += in1[row * width + i] * in2[i * width + col];\n }\n out[row * width + col] = sum;\n }\n}\n\nvoid doMatMul(float* out, const float* in1, const float* in2, int width) {\n float *dIn1, *dIn2, *dOut;\n const size_t size = width * width * sizeof(float);\n\n printf(\"Allocating device memory...\\n\");\n cudaCheckError(::hipMalloc(&dIn1, size));\n cudaCheckError(::hipMalloc(&dIn2, size));\n cudaCheckError(::hipMalloc(&dOut, size));\n\n printf(\"Copying data to device...\\n\");\n cudaCheckError(::hipMemcpy(dIn1, in1, size, hipMemcpyHostToDevice));\n cudaCheckError(::hipMemcpy(dIn2, in2, size, hipMemcpyHostToDevice));\n\n dim3 threadsPerBlock(16, 16);\n dim3 blocksPerGrid((width + threadsPerBlock.x - 1) / threadsPerBlock.x,\n (width + threadsPerBlock.y - 1) / threadsPerBlock.y);\n printf(\"Launching kernel...\\n\");\n for (int j = 0; j < 100; ++j)\n {\n nvtxRangePush(\"Matrix multiplication kernel\");\n matMul<<>>(dOut, dIn1, dIn2, width);\n nvtxRangePop();\n }\n cudaCheckError(::hipDeviceSynchronize());\n\n printf(\"Copying result to host...\\n\");\n cudaCheckError(::hipMemcpy(out, dOut, size, hipMemcpyHostToDevice));\n\n printf(\"Freeing device memory...\\n\");\n cudaCheckError(::hipFree(dIn1));\n cudaCheckError(::hipFree(dIn2));\n cudaCheckError(::hipFree(dOut));\n}\n\nint main() {\n float *in1, *in2, *out;\n const size_t width = W * W * sizeof(float);\n cudaCheckError(::hipHostMalloc(&in1, width));\n cudaCheckError(::hipHostMalloc(&in2, width));\n cudaCheckError(::hipHostMalloc(&out, width));\n for (int i = 0; i < W * W; ++i) {\n in1[i] = static_cast(rand()) / RAND_MAX;\n in2[i] = static_cast(rand()) / RAND_MAX;\n }\n printf(\"Starting matrix multiplication...\\n\");\n doMatMul(out, in1, in2, width);\n printf(\"Matrix multiplication completed.\\n\");\n printf(\"Freeing host memory...\\n\");\n cudaCheckError(::hipHostFree(in1));\n cudaCheckError(::hipHostFree(in2));\n cudaCheckError(::hipHostFree(out));\n return 0;\n}\n", "source": "/kaggle/working/cuda_examples/ProfilingExamples/nvtx/nvtx_example.cu"} {"input": "#include \n#include \n#include \"../../utils/utils.cuh\"\n\n#define N (1 << 20) // 1M elements\n#define THREADS 1024 // Threads per block\n#define BLOCKS (N + THREADS - 1) / THREADS // Number of blocks\n\n// Naive Reduction (with bank conflicts)\n__global__ void reduce_naive(float *input, float *output) {\n __shared__ float shared_mem[THREADS];\n\n int tid = threadIdx.x;\n int global_idx = blockIdx.x * blockDim.x + tid;\n\n // Load global memory into shared memory\n shared_mem[tid] = (global_idx < N) ? input[global_idx] : 0.0f;\n __syncthreads();\n\n // Reduction in shared memory\n for (int stride = 1; stride < blockDim.x; stride *= 2) {\n if (tid % (2 * stride) == 0) {\n shared_mem[tid] += shared_mem[tid + stride];\n }\n __syncthreads();\n }\n\n // Write block sum to output\n if (tid == 0) output[blockIdx.x] = shared_mem[0];\n}\n\n// Optimized Reduction (Avoiding Bank Conflicts)\n__global__ void reduce_optimized(float *input, float *output) {\n __shared__ float shared_mem[THREADS + (THREADS / 32)];\n\n int tid = threadIdx.x;\n int global_idx = blockIdx.x * blockDim.x + tid;\n\n // Load data into shared memory (coalesced)\n shared_mem[tid] = (global_idx < N) ? input[global_idx] : 0.0f;\n __syncthreads();\n\n for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {\n if (tid < stride) {\n shared_mem[tid] += shared_mem[tid + stride];\n }\n __syncthreads();\n }\n\n // Write block sum to output\n if (tid == 0) output[blockIdx.x] = shared_mem[0];\n}\n\n// CPU Reduction (Baseline for Benchmark)\nfloat reduce_cpu(float *input, int size) {\n float sum = 0.0f;\n for (int i = 0; i < size; i++) {\n sum += input[i];\n }\n return sum;\n}\n\n// Benchmarking Function\nvoid benchmark() {\n float *h_input, *h_output;\n float *d_input, *d_output;\n cudaCheckError(::cudaMallocHost(&h_input, N * sizeof(float)));\n cudaCheckError(::cudaMallocHost(&h_output, BLOCKS * sizeof(float)));\n cudaCheckError(::cudaMalloc(&d_input, N * sizeof(float)));\n cudaCheckError(::cudaMalloc(&d_output, BLOCKS * sizeof(float)));\n\n // Initialize input\n for (int i = 0; i < N; i++) {\n h_input[i] = 1.0f; // Simple case where sum = N\n }\n\n cudaCheckError(::cudaMemcpy(d_input, h_input, N * sizeof(float), cudaMemcpyHostToDevice));\n\n cudaEvent_t start, stop;\n cudaCheckError(::cudaEventCreate((&start)));\n cudaCheckError(::cudaEventCreate(&stop));\n // Benchmark Na\u00efve Kernel\n cudaCheckError(::cudaEventRecord(start));\n reduce_naive<<>>(d_input, d_output);\n cudaCheckError(::cudaEventRecord(stop));\n cudaCheckError(::cudaEventSynchronize(stop));\n float time_naive;\n cudaCheckError(::cudaEventElapsedTime(&time_naive, start, stop));\n printf(\"Na\u00efve Kernel Time: %f ms\\n\", time_naive);\n cudaCheckError(::cudaMemcpy(h_output, d_output, BLOCKS * sizeof(float), cudaMemcpyDeviceToHost));\n\n // Final sum on CPU\n float sum_naive = reduce_cpu(h_output, BLOCKS);\n\n // Benchmark Optimized Kernel\n cudaCheckError(::cudaEventRecord(start));\n reduce_optimized<<>>(d_input, d_output);\n cudaCheckError(::cudaEventRecord(stop));\n cudaCheckError(::cudaEventSynchronize(stop));\n float time_optimized;\n cudaCheckError(::cudaEventElapsedTime(&time_optimized, start, stop));\n printf(\"Optimized Kernel Time: %f ms\\n\", time_optimized);\n cudaCheckError(::cudaMemcpy(h_output, d_output, BLOCKS * sizeof(float), cudaMemcpyDeviceToHost));\n \n float sum_optimized = reduce_cpu(h_output, BLOCKS);\n\n // CPU Baseline\n float sum_cpu = reduce_cpu(h_input, N);\n\n printf(\"CPU Sum: %f\\n\", sum_cpu);\n printf(\"Na\u00efve GPU Sum: %f\\n\", sum_naive);\n printf(\"Optimized GPU Sum: %f\\n\", sum_optimized);\n\n // Cleanup\n cudaCheckError(::cudaFreeHost(h_input));\n cudaCheckError(::cudaFreeHost(h_output));\n cudaCheckError(::cudaFree(d_input));\n cudaCheckError(::cudaFree(d_output));\n}\n\nint main() {\n benchmark();\n return 0;\n}\n", "output": "#include \n#include \n#include \"../../utils/utils.cuh\"\n\n#define N (1 << 20) // 1M elements\n#define THREADS 1024 // Threads per block\n#define BLOCKS (N + THREADS - 1) / THREADS // Number of blocks\n\n// Naive Reduction (with bank conflicts)\n__global__ void reduce_naive(float *input, float *output) {\n __shared__ float shared_mem[THREADS];\n\n int tid = threadIdx.x;\n int global_idx = blockIdx.x * blockDim.x + tid;\n\n // Load global memory into shared memory\n shared_mem[tid] = (global_idx < N) ? input[global_idx] : 0.0f;\n __syncthreads();\n\n // Reduction in shared memory\n for (int stride = 1; stride < blockDim.x; stride *= 2) {\n if (tid % (2 * stride) == 0) {\n shared_mem[tid] += shared_mem[tid + stride];\n }\n __syncthreads();\n }\n\n // Write block sum to output\n if (tid == 0) output[blockIdx.x] = shared_mem[0];\n}\n\n// Optimized Reduction (Avoiding Bank Conflicts)\n__global__ void reduce_optimized(float *input, float *output) {\n __shared__ float shared_mem[THREADS + (THREADS / 32)];\n\n int tid = threadIdx.x;\n int global_idx = blockIdx.x * blockDim.x + tid;\n\n // Load data into shared memory (coalesced)\n shared_mem[tid] = (global_idx < N) ? input[global_idx] : 0.0f;\n __syncthreads();\n\n for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {\n if (tid < stride) {\n shared_mem[tid] += shared_mem[tid + stride];\n }\n __syncthreads();\n }\n\n // Write block sum to output\n if (tid == 0) output[blockIdx.x] = shared_mem[0];\n}\n\n// CPU Reduction (Baseline for Benchmark)\nfloat reduce_cpu(float *input, int size) {\n float sum = 0.0f;\n for (int i = 0; i < size; i++) {\n sum += input[i];\n }\n return sum;\n}\n\n// Benchmarking Function\nvoid benchmark() {\n float *h_input, *h_output;\n float *d_input, *d_output;\n cudaCheckError(::hipHostMalloc(&h_input, N * sizeof(float)));\n cudaCheckError(::hipHostMalloc(&h_output, BLOCKS * sizeof(float)));\n cudaCheckError(::hipMalloc(&d_input, N * sizeof(float)));\n cudaCheckError(::hipMalloc(&d_output, BLOCKS * sizeof(float)));\n\n // Initialize input\n for (int i = 0; i < N; i++) {\n h_input[i] = 1.0f; // Simple case where sum = N\n }\n\n cudaCheckError(::hipMemcpy(d_input, h_input, N * sizeof(float), hipMemcpyHostToDevice));\n\n hipEvent_t start, stop;\n cudaCheckError(::hipEventCreate((&start)));\n cudaCheckError(::hipEventCreate(&stop));\n // Benchmark Na\u00efve Kernel\n cudaCheckError(::hipEventRecord(start));\n reduce_naive<<>>(d_input, d_output);\n cudaCheckError(::hipEventRecord(stop));\n cudaCheckError(::hipEventSynchronize(stop));\n float time_naive;\n cudaCheckError(::hipEventElapsedTime(&time_naive, start, stop));\n printf(\"Na\u00efve Kernel Time: %f ms\\n\", time_naive);\n cudaCheckError(::hipMemcpy(h_output, d_output, BLOCKS * sizeof(float), hipMemcpyDeviceToHost));\n\n // Final sum on CPU\n float sum_naive = reduce_cpu(h_output, BLOCKS);\n\n // Benchmark Optimized Kernel\n cudaCheckError(::hipEventRecord(start));\n reduce_optimized<<>>(d_input, d_output);\n cudaCheckError(::hipEventRecord(stop));\n cudaCheckError(::hipEventSynchronize(stop));\n float time_optimized;\n cudaCheckError(::hipEventElapsedTime(&time_optimized, start, stop));\n printf(\"Optimized Kernel Time: %f ms\\n\", time_optimized);\n cudaCheckError(::hipMemcpy(h_output, d_output, BLOCKS * sizeof(float), hipMemcpyDeviceToHost));\n \n float sum_optimized = reduce_cpu(h_output, BLOCKS);\n\n // CPU Baseline\n float sum_cpu = reduce_cpu(h_input, N);\n\n printf(\"CPU Sum: %f\\n\", sum_cpu);\n printf(\"Na\u00efve GPU Sum: %f\\n\", sum_naive);\n printf(\"Optimized GPU Sum: %f\\n\", sum_optimized);\n\n // Cleanup\n cudaCheckError(::hipHostFree(h_input));\n cudaCheckError(::hipHostFree(h_output));\n cudaCheckError(::hipFree(d_input));\n cudaCheckError(::hipFree(d_output));\n}\n\nint main() {\n benchmark();\n return 0;\n}\n", "source": "/kaggle/working/cuda_examples/ProfilingExamples/memory_bank_conflict/bank_conflict.cu"} {"input": "#include \n#include \n#include \"../../utils/utils.cuh\"\n\n#define N (1 << 20) // 1M elements\n#define THREADS 1024 // Threads per block\n#define BLOCKS (N + THREADS - 1) / THREADS // Number of blocks\n\n__global__ void reduce_global(float *input, float *output) {\n int tid = threadIdx.x;\n int global_idx = blockIdx.x * blockDim.x * 2 + tid; // Unroll by 2 for better memory efficiency\n\n float sum = 0.0f;\n if (global_idx < N) sum += input[global_idx];\n if (global_idx + blockDim.x < N) sum += input[global_idx + blockDim.x];\n\n // In-register reduction (warp-level)\n for (int stride = warpSize / 2; stride > 0; stride >>= 1) {\n sum += __shfl_down_sync(0xFFFFFFFF, sum, stride);\n }\n\n // Store only the first thread's result per warp\n if (tid % warpSize == 0) {\n atomicAdd(&output[blockIdx.x], sum);\n }\n}\n\n\n\n// Optimized Reduction with Shared Memory\n__global__ void reduce_shared(float *input, float *output) {\n __shared__ float shared_mem[THREADS + (THREADS / 32)];\n\n int tid = threadIdx.x;\n int global_idx = blockIdx.x * blockDim.x + tid;\n\n // Load data into shared memory (coalesced)\n shared_mem[tid] = (global_idx < N) ? input[global_idx] : 0.0f;\n __syncthreads();\n\n for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {\n if (tid < stride) {\n shared_mem[tid] += shared_mem[tid + stride];\n }\n __syncthreads();\n }\n\n // Write block sum to output\n if (tid == 0) output[blockIdx.x] = shared_mem[0];\n}\n\n// CPU Reduction (Baseline for Benchmark)\nfloat reduce_cpu(float *input, int size) {\n float sum = 0.0f;\n for (int i = 0; i < size; i++) {\n sum += input[i];\n }\n return sum;\n}\n\n// Benchmarking Function\nvoid benchmark() {\n float *h_input, *h_output;\n float *d_input, *d_output;\n cudaCheckError(::cudaMallocHost(&h_input, N * sizeof(float)));\n cudaCheckError(::cudaMallocHost(&h_output, BLOCKS * sizeof(float)));\n cudaCheckError(::cudaMalloc(&d_input, N * sizeof(float)));\n cudaCheckError(::cudaMalloc(&d_output, BLOCKS * sizeof(float)));\n\n // Initialize input\n for (int i = 0; i < N; i++) {\n h_input[i] = 1.0f; // Simple case where sum = N\n }\n\n cudaCheckError(::cudaMemcpy(d_input, h_input, N * sizeof(float), cudaMemcpyHostToDevice));\n\n cudaEvent_t start, stop;\n cudaCheckError(::cudaEventCreate((&start)));\n cudaCheckError(::cudaEventCreate(&stop));\n // Benchmark Global Kernel\n cudaCheckError(::cudaEventRecord(start));\n reduce_global<<>>(d_input, d_output);\n cudaCheckError(::cudaEventRecord(stop));\n cudaCheckError(::cudaEventSynchronize(stop));\n float time_naive;\n cudaCheckError(::cudaEventElapsedTime(&time_naive, start, stop));\n printf(\"Global Kernel Time: %f ms\\n\", time_naive);\n cudaCheckError(::cudaMemcpy(h_output, d_output, BLOCKS * sizeof(float), cudaMemcpyDeviceToHost));\n\n // Final sum on CPU\n float sum_naive = reduce_cpu(h_output, BLOCKS);\n\n // Benchmark Optimized Kernel\n cudaCheckError(::cudaEventRecord(start));\n reduce_shared<<>>(d_input, d_output);\n cudaCheckError(::cudaEventRecord(stop));\n cudaCheckError(::cudaEventSynchronize(stop));\n float time_optimized;\n cudaCheckError(::cudaEventElapsedTime(&time_optimized, start, stop));\n printf(\"Shmem Kernel Time: %f ms\\n\", time_optimized);\n cudaCheckError(::cudaMemcpy(h_output, d_output, BLOCKS * sizeof(float), cudaMemcpyDeviceToHost));\n \n float sum_optimized = reduce_cpu(h_output, BLOCKS);\n\n // CPU Baseline\n float sum_cpu = reduce_cpu(h_input, N);\n\n printf(\"CPU Sum: %f\\n\", sum_cpu);\n printf(\"GPU Global memory Sum: %f\\n\", sum_naive);\n printf(\"GPU Shared memory Sum: %f\\n\", sum_optimized);\n\n // Cleanup\n cudaCheckError(::cudaFreeHost(h_input));\n cudaCheckError(::cudaFreeHost(h_output));\n cudaCheckError(::cudaFree(d_input));\n cudaCheckError(::cudaFree(d_output));\n}\n\nint main() {\n benchmark();\n return 0;\n}\n", "output": "#include \n#include \n#include \"../../utils/utils.cuh\"\n\n#define N (1 << 20) // 1M elements\n#define THREADS 1024 // Threads per block\n#define BLOCKS (N + THREADS - 1) / THREADS // Number of blocks\n\n__global__ void reduce_global(float *input, float *output) {\n int tid = threadIdx.x;\n int global_idx = blockIdx.x * blockDim.x * 2 + tid; // Unroll by 2 for better memory efficiency\n\n float sum = 0.0f;\n if (global_idx < N) sum += input[global_idx];\n if (global_idx + blockDim.x < N) sum += input[global_idx + blockDim.x];\n\n // In-register reduction (warp-level)\n for (int stride = warpSize / 2; stride > 0; stride >>= 1) {\n sum += __shfl_down_sync(0xFFFFFFFF, sum, stride);\n }\n\n // Store only the first thread's result per warp\n if (tid % warpSize == 0) {\n atomicAdd(&output[blockIdx.x], sum);\n }\n}\n\n\n\n// Optimized Reduction with Shared Memory\n__global__ void reduce_shared(float *input, float *output) {\n __shared__ float shared_mem[THREADS + (THREADS / 32)];\n\n int tid = threadIdx.x;\n int global_idx = blockIdx.x * blockDim.x + tid;\n\n // Load data into shared memory (coalesced)\n shared_mem[tid] = (global_idx < N) ? input[global_idx] : 0.0f;\n __syncthreads();\n\n for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {\n if (tid < stride) {\n shared_mem[tid] += shared_mem[tid + stride];\n }\n __syncthreads();\n }\n\n // Write block sum to output\n if (tid == 0) output[blockIdx.x] = shared_mem[0];\n}\n\n// CPU Reduction (Baseline for Benchmark)\nfloat reduce_cpu(float *input, int size) {\n float sum = 0.0f;\n for (int i = 0; i < size; i++) {\n sum += input[i];\n }\n return sum;\n}\n\n// Benchmarking Function\nvoid benchmark() {\n float *h_input, *h_output;\n float *d_input, *d_output;\n cudaCheckError(::hipHostMalloc(&h_input, N * sizeof(float)));\n cudaCheckError(::hipHostMalloc(&h_output, BLOCKS * sizeof(float)));\n cudaCheckError(::hipMalloc(&d_input, N * sizeof(float)));\n cudaCheckError(::hipMalloc(&d_output, BLOCKS * sizeof(float)));\n\n // Initialize input\n for (int i = 0; i < N; i++) {\n h_input[i] = 1.0f; // Simple case where sum = N\n }\n\n cudaCheckError(::hipMemcpy(d_input, h_input, N * sizeof(float), hipMemcpyHostToDevice));\n\n hipEvent_t start, stop;\n cudaCheckError(::hipEventCreate((&start)));\n cudaCheckError(::hipEventCreate(&stop));\n // Benchmark Global Kernel\n cudaCheckError(::hipEventRecord(start));\n reduce_global<<>>(d_input, d_output);\n cudaCheckError(::hipEventRecord(stop));\n cudaCheckError(::hipEventSynchronize(stop));\n float time_naive;\n cudaCheckError(::hipEventElapsedTime(&time_naive, start, stop));\n printf(\"Global Kernel Time: %f ms\\n\", time_naive);\n cudaCheckError(::hipMemcpy(h_output, d_output, BLOCKS * sizeof(float), hipMemcpyDeviceToHost));\n\n // Final sum on CPU\n float sum_naive = reduce_cpu(h_output, BLOCKS);\n\n // Benchmark Optimized Kernel\n cudaCheckError(::hipEventRecord(start));\n reduce_shared<<>>(d_input, d_output);\n cudaCheckError(::hipEventRecord(stop));\n cudaCheckError(::hipEventSynchronize(stop));\n float time_optimized;\n cudaCheckError(::hipEventElapsedTime(&time_optimized, start, stop));\n printf(\"Shmem Kernel Time: %f ms\\n\", time_optimized);\n cudaCheckError(::hipMemcpy(h_output, d_output, BLOCKS * sizeof(float), hipMemcpyDeviceToHost));\n \n float sum_optimized = reduce_cpu(h_output, BLOCKS);\n\n // CPU Baseline\n float sum_cpu = reduce_cpu(h_input, N);\n\n printf(\"CPU Sum: %f\\n\", sum_cpu);\n printf(\"GPU Global memory Sum: %f\\n\", sum_naive);\n printf(\"GPU Shared memory Sum: %f\\n\", sum_optimized);\n\n // Cleanup\n cudaCheckError(::hipHostFree(h_input));\n cudaCheckError(::hipHostFree(h_output));\n cudaCheckError(::hipFree(d_input));\n cudaCheckError(::hipFree(d_output));\n}\n\nint main() {\n benchmark();\n return 0;\n}\n", "source": "/kaggle/working/cuda_examples/ProfilingExamples/global_vs_shared_mem/global_vs_shared.cu"} {"input": "#include \n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n// Simple matmul kernel.\n__global__ void matMul(float *out, const float *in1, const float *in2, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n float sum = 0.0f;\n if (row < width && col < width) {\n for (int i = 0; i < width; ++i) {\n sum += in1[row * width + i] * in2[i * width + col];\n }\n out[row * width + col] = sum;\n }\n}\n\n// Add one kernel to emphasize serial execution within a stream.\n__global__ void addOne(float *out, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n if (row < width && col < width) {\n int idx = row * width + col;\n out[idx] += 1.0f;\n }\n}\n\n// Minus one kernel to emphasize serial execution within a stream.\n__global__ void minusOne(float *out, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n if (row < width && col < width) {\n int idx = row * width + col;\n out[idx] -= 1.0f;\n }\n}\n\n// Add matrices to a final result for comparison.\n__global__ void matAdd(float *out, const float *in1, const float *in2, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n if (row < width && col < width) {\n out[row * width + col] = in1[row * width + col] + in2[row * width + col];\n }\n}\n\n// Perform a matmul on separate streams for a million elements, then add them.\nint main() {\n // Set dimensionality.\n const int width = 1024;\n const int size = width * width * sizeof(float);\n\n // Create separate streams to run matmuls asynchronously.\n cudaStream_t stream1, stream2;\n cudaCheckError(::cudaStreamCreate(&stream1));\n cudaCheckError(::cudaStreamCreate(&stream2));\n\n // Allocations and initializations.\n float *hA, *hB, *hC, *hD, *hResult;\n float *dA, *dB, *dC, *dD, *dResultAB, *dResultCD, *dFinal;\n\n // Host.\n cudaCheckError(::cudaMallocHost(&hA, size));\n cudaCheckError(::cudaMallocHost(&hB, size));\n cudaCheckError(::cudaMallocHost(&hC, size));\n cudaCheckError(::cudaMallocHost(&hD, size));\n cudaCheckError(::cudaMallocHost(&hResult, size));\n\n for (int i = 0; i < width * width; ++i) {\n hA[i] = 1.0f;\n hB[i] = 2.0f;\n hC[i] = 3.0f;\n hD[i] = 4.0f;\n }\n\n // Device.\n cudaCheckError(::cudaMalloc(&dA, size));\n cudaCheckError(::cudaMalloc(&dB, size));\n cudaCheckError(::cudaMalloc(&dC, size));\n cudaCheckError(::cudaMalloc(&dD, size));\n cudaCheckError(::cudaMalloc(&dResultAB, size));\n cudaCheckError(::cudaMalloc(&dResultCD, size));\n cudaCheckError(::cudaMalloc(&dFinal, size));\n\n // Copy data using separate streams. This is not necessary, per se, but it will\n // allow for overlap of copy and execution in a batch mode case.\n cudaCheckError(::cudaMemcpyAsync(dA, hA, size, cudaMemcpyHostToDevice, stream1));\n cudaCheckError(::cudaMemcpyAsync(dB, hB, size, cudaMemcpyHostToDevice, stream1));\n cudaCheckError(::cudaMemcpyAsync(dC, hC, size, cudaMemcpyHostToDevice, stream2));\n cudaCheckError(::cudaMemcpyAsync(dD, hD, size, cudaMemcpyHostToDevice, stream2));\n \n // Launch matmul kernels on different streams.\n dim3 threadsPerBlock(32, 32);\n dim3 numBlocks((width + threadsPerBlock.x - 1) / threadsPerBlock.x, \n (width + threadsPerBlock.y - 1) / threadsPerBlock.y);\n matMul<<>>(dResultAB, dA, dB, width);\n matMul<<>>(dResultCD, dC, dD, width);\n\n // Add 1.0f to each element in the result matrices\n addOne<<>>(dResultAB, width);\n addOne<<>>(dResultCD, width);\n\n // Subtract 1.0f to each element in the result matrices\n minusOne<<>>(dResultAB, width);\n minusOne<<>>(dResultCD, width);\n \n // Synchronize work prior to computing on each stream's result array.\n cudaCheckError(::cudaStreamSynchronize(stream1));\n cudaCheckError(::cudaStreamSynchronize(stream2));\n\n matAdd<<>>(dFinal, dResultAB, dResultCD, width);\n cudaCheckError(::cudaMemcpy(hResult, dFinal, size, cudaMemcpyDeviceToHost));\n\n // Verify results.\n float expected = (1.0f * 2.0f + 3.0f * 4.0f) * width;\n for (int i = 0; i < width * width; ++i) {\n assert(fabs(hResult[i] - expected) < 1e-5);\n }\n\n std::cout << \"Test passed!\" << std::endl;\n\n // Clean up.\n cudaCheckError(::cudaFree(dA));\n cudaCheckError(::cudaFree(dB));\n cudaCheckError(::cudaFree(dC));\n cudaCheckError(::cudaFree(dD));\n cudaCheckError(::cudaFree(dResultAB));\n cudaCheckError(::cudaFree(dResultCD));\n cudaCheckError(::cudaFree(dFinal));\n cudaCheckError(::cudaStreamDestroy(stream1));\n cudaCheckError(::cudaStreamDestroy(stream2));\n cudaCheckError(::cudaFreeHost(hA));\n cudaCheckError(::cudaFreeHost(hB));\n cudaCheckError(::cudaFreeHost(hC));\n cudaCheckError(::cudaFreeHost(hD));\n cudaCheckError(::cudaFreeHost(hResult));\n\n return 0;\n}\n\n", "output": "#include \n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n// Simple matmul kernel.\n__global__ void matMul(float *out, const float *in1, const float *in2, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n float sum = 0.0f;\n if (row < width && col < width) {\n for (int i = 0; i < width; ++i) {\n sum += in1[row * width + i] * in2[i * width + col];\n }\n out[row * width + col] = sum;\n }\n}\n\n// Add one kernel to emphasize serial execution within a stream.\n__global__ void addOne(float *out, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n if (row < width && col < width) {\n int idx = row * width + col;\n out[idx] += 1.0f;\n }\n}\n\n// Minus one kernel to emphasize serial execution within a stream.\n__global__ void minusOne(float *out, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n if (row < width && col < width) {\n int idx = row * width + col;\n out[idx] -= 1.0f;\n }\n}\n\n// Add matrices to a final result for comparison.\n__global__ void matAdd(float *out, const float *in1, const float *in2, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n if (row < width && col < width) {\n out[row * width + col] = in1[row * width + col] + in2[row * width + col];\n }\n}\n\n// Perform a matmul on separate streams for a million elements, then add them.\nint main() {\n // Set dimensionality.\n const int width = 1024;\n const int size = width * width * sizeof(float);\n\n // Create separate streams to run matmuls asynchronously.\n hipStream_t stream1, stream2;\n cudaCheckError(::hipStreamCreate(&stream1));\n cudaCheckError(::hipStreamCreate(&stream2));\n\n // Allocations and initializations.\n float *hA, *hB, *hC, *hD, *hResult;\n float *dA, *dB, *dC, *dD, *dResultAB, *dResultCD, *dFinal;\n\n // Host.\n cudaCheckError(::hipHostMalloc(&hA, size));\n cudaCheckError(::hipHostMalloc(&hB, size));\n cudaCheckError(::hipHostMalloc(&hC, size));\n cudaCheckError(::hipHostMalloc(&hD, size));\n cudaCheckError(::hipHostMalloc(&hResult, size));\n\n for (int i = 0; i < width * width; ++i) {\n hA[i] = 1.0f;\n hB[i] = 2.0f;\n hC[i] = 3.0f;\n hD[i] = 4.0f;\n }\n\n // Device.\n cudaCheckError(::hipMalloc(&dA, size));\n cudaCheckError(::hipMalloc(&dB, size));\n cudaCheckError(::hipMalloc(&dC, size));\n cudaCheckError(::hipMalloc(&dD, size));\n cudaCheckError(::hipMalloc(&dResultAB, size));\n cudaCheckError(::hipMalloc(&dResultCD, size));\n cudaCheckError(::hipMalloc(&dFinal, size));\n\n // Copy data using separate streams. This is not necessary, per se, but it will\n // allow for overlap of copy and execution in a batch mode case.\n cudaCheckError(::hipMemcpyAsync(dA, hA, size, hipMemcpyHostToDevice, stream1));\n cudaCheckError(::hipMemcpyAsync(dB, hB, size, hipMemcpyHostToDevice, stream1));\n cudaCheckError(::hipMemcpyAsync(dC, hC, size, hipMemcpyHostToDevice, stream2));\n cudaCheckError(::hipMemcpyAsync(dD, hD, size, hipMemcpyHostToDevice, stream2));\n \n // Launch matmul kernels on different streams.\n dim3 threadsPerBlock(32, 32);\n dim3 numBlocks((width + threadsPerBlock.x - 1) / threadsPerBlock.x, \n (width + threadsPerBlock.y - 1) / threadsPerBlock.y);\n matMul<<>>(dResultAB, dA, dB, width);\n matMul<<>>(dResultCD, dC, dD, width);\n\n // Add 1.0f to each element in the result matrices\n addOne<<>>(dResultAB, width);\n addOne<<>>(dResultCD, width);\n\n // Subtract 1.0f to each element in the result matrices\n minusOne<<>>(dResultAB, width);\n minusOne<<>>(dResultCD, width);\n \n // Synchronize work prior to computing on each stream's result array.\n cudaCheckError(::hipStreamSynchronize(stream1));\n cudaCheckError(::hipStreamSynchronize(stream2));\n\n matAdd<<>>(dFinal, dResultAB, dResultCD, width);\n cudaCheckError(::hipMemcpy(hResult, dFinal, size, hipMemcpyDeviceToHost));\n\n // Verify results.\n float expected = (1.0f * 2.0f + 3.0f * 4.0f) * width;\n for (int i = 0; i < width * width; ++i) {\n assert(fabs(hResult[i] - expected) < 1e-5);\n }\n\n std::cout << \"Test passed!\" << std::endl;\n\n // Clean up.\n cudaCheckError(::hipFree(dA));\n cudaCheckError(::hipFree(dB));\n cudaCheckError(::hipFree(dC));\n cudaCheckError(::hipFree(dD));\n cudaCheckError(::hipFree(dResultAB));\n cudaCheckError(::hipFree(dResultCD));\n cudaCheckError(::hipFree(dFinal));\n cudaCheckError(::hipStreamDestroy(stream1));\n cudaCheckError(::hipStreamDestroy(stream2));\n cudaCheckError(::hipHostFree(hA));\n cudaCheckError(::hipHostFree(hB));\n cudaCheckError(::hipHostFree(hC));\n cudaCheckError(::hipHostFree(hD));\n cudaCheckError(::hipHostFree(hResult));\n\n return 0;\n}\n\n", "source": "/kaggle/working/cuda_examples/PerformanceChecklistExamples/cuda_streams/cuda_streams.cu"} {"input": "#include \n#include \n#include \"../../utils/utils.cuh\"\n\n/**\n * @file vector_add.cu\n * \n * @brief We demonstrate adding 2 vectors of size 2048 using CUDA\n */\n\n/**\n * @brief CUDA kernel to perform vector addition\n * \n * @param d_v1 Device vector holding a input vector\n * @param d_v2 Device vector holding another input vector\n * @param d_v3 Device vector holding the sum of d_v1 and d_v2\n */\n__global__ void vector_add_kernel(float *d_v1, float *d_v2, float *d_v3, int n) {\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n if (i < n) {\n d_v3[i] = d_v1[i] + d_v2[i];\n }\n}\n\nint main() {\n int N = 2048;\n size_t size = sizeof(float) * N;\n\n // Allocate memory for the host vectors\n float *v1, *v2, *v3;\n cudaCheckError(::cudaMallocHost(&v1, size));\n cudaCheckError(::cudaMallocHost(&v2, size));\n cudaCheckError(::cudaMallocHost(&v3, size));\n\n // Fill the host vectors with data\n for (size_t i = 0; i < N; i++) {\n v1[i] = 6.4;\n v2[i] = 3.5;\n }\n\n // Create 3 device vectors for holding N float elements\n float *d_v1, *d_v2, *d_v3;\n cudaCheckError(::cudaMalloc(&d_v1, size));\n cudaCheckError(::cudaMalloc(&d_v2, size));\n cudaCheckError(::cudaMalloc(&d_v3, size));\n\n // Copy the host vectors into the device vectors\n cudaCheckError(::cudaMemcpy(d_v1, v1, size, cudaMemcpyHostToDevice));\n cudaCheckError(::cudaMemcpy(d_v2, v2, size, cudaMemcpyHostToDevice));\n\n // Launch the kernel using 8 thread blocks and 256 threads / block\n vector_add_kernel<<<8, 256>>>(d_v1, d_v2, d_v3, N);\n\n // Copy the result device vector into the result host vector\n cudaCheckError(::cudaMemcpy(v3, d_v3, size, cudaMemcpyDeviceToHost));\n\n cudaCheckError(::cudaDeviceSynchronize());\n\n // Verify the result host vector\n for (int i = 0; i < N; i++) {\n assert(v3[i] == (float)9.9);\n }\n\n // Free the device vector\n cudaCheckError(::cudaFree(d_v1));\n cudaCheckError(::cudaFree(d_v2));\n cudaCheckError(::cudaFree(d_v3));\n\n // Free the host vectors\n cudaCheckError(::cudaFreeHost(v1));\n cudaCheckError(::cudaFreeHost(v2));\n cudaCheckError(::cudaFreeHost(v3));\n\n return 0;\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n/**\n * @file vector_add.cu\n * \n * @brief We demonstrate adding 2 vectors of size 2048 using CUDA\n */\n\n/**\n * @brief CUDA kernel to perform vector addition\n * \n * @param d_v1 Device vector holding a input vector\n * @param d_v2 Device vector holding another input vector\n * @param d_v3 Device vector holding the sum of d_v1 and d_v2\n */\n__global__ void vector_add_kernel(float *d_v1, float *d_v2, float *d_v3, int n) {\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n if (i < n) {\n d_v3[i] = d_v1[i] + d_v2[i];\n }\n}\n\nint main() {\n int N = 2048;\n size_t size = sizeof(float) * N;\n\n // Allocate memory for the host vectors\n float *v1, *v2, *v3;\n cudaCheckError(::hipHostMalloc(&v1, size));\n cudaCheckError(::hipHostMalloc(&v2, size));\n cudaCheckError(::hipHostMalloc(&v3, size));\n\n // Fill the host vectors with data\n for (size_t i = 0; i < N; i++) {\n v1[i] = 6.4;\n v2[i] = 3.5;\n }\n\n // Create 3 device vectors for holding N float elements\n float *d_v1, *d_v2, *d_v3;\n cudaCheckError(::hipMalloc(&d_v1, size));\n cudaCheckError(::hipMalloc(&d_v2, size));\n cudaCheckError(::hipMalloc(&d_v3, size));\n\n // Copy the host vectors into the device vectors\n cudaCheckError(::hipMemcpy(d_v1, v1, size, hipMemcpyHostToDevice));\n cudaCheckError(::hipMemcpy(d_v2, v2, size, hipMemcpyHostToDevice));\n\n // Launch the kernel using 8 thread blocks and 256 threads / block\n vector_add_kernel<<<8, 256>>>(d_v1, d_v2, d_v3, N);\n\n // Copy the result device vector into the result host vector\n cudaCheckError(::hipMemcpy(v3, d_v3, size, hipMemcpyDeviceToHost));\n\n cudaCheckError(::hipDeviceSynchronize());\n\n // Verify the result host vector\n for (int i = 0; i < N; i++) {\n assert(v3[i] == (float)9.9);\n }\n\n // Free the device vector\n cudaCheckError(::hipFree(d_v1));\n cudaCheckError(::hipFree(d_v2));\n cudaCheckError(::hipFree(d_v3));\n\n // Free the host vectors\n cudaCheckError(::hipHostFree(v1));\n cudaCheckError(::hipHostFree(v2));\n cudaCheckError(::hipHostFree(v3));\n\n return 0;\n}\n", "source": "/kaggle/working/cuda_examples/KernelAndLibExamples/vector_add/vector_add.cu"} {"input": "#include \n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n// Device (GPU) code to print thread index from the GPU.\n__global__ void helloFromGPU()\n{\n int threadId = blockIdx.x * blockDim.x + threadIdx.x;\n printf(\"Hello from GPU thread: %d!\\n\", threadId);\n}\n\n// Host code\nint main()\n{\n int nDevices;\n cudaGetDeviceCount(&nDevices);\n for (int i = 0; i < nDevices; i++) {\n cudaDeviceProp prop;\n cudaGetDeviceProperties(&prop, i);\n printf(\"Device Number: %d\\n\", i);\n printf(\" Device name: %s\\n\", prop.name);\n printf(\" major-minor: %d-%d\\n\", prop.major, prop.minor);\n printf(\" Concurrent kernels: %s\\n\", prop.concurrentKernels ? \"yes\" : \"no\");\n printf(\" Concurrent computation/communication: %s\\n\\n\",prop.deviceOverlap ? \"yes\" : \"no\");\n printf(\" Managed memory supported: %s\\n\", prop.managedMemory ? \"yes\" : \"no\");\n printf(\" Memory Clock Rate (MHz): %d\\n\",\n prop.memoryClockRate/1024);\n printf(\" Memory Bus Width (bits): %d\\n\",\n prop.memoryBusWidth);\n printf(\" Peak Memory Bandwidth (GB/s): %.1f\\n\",\n 2.0*prop.memoryClockRate*(prop.memoryBusWidth/8)/1.0e6);\n printf(\" Total global memory (Gbytes): %.1f\\n\",(float)(prop.totalGlobalMem)/1024.0/1024.0/1024.0);\n printf(\" Shared memory per block (Kbytes): %.1f\\n\",(float)(prop.sharedMemPerBlock)/1024.0);\n printf(\" L2 cache size (Kbytes): %.1f\\n\",(float)(prop.l2CacheSize)/1024);\n printf(\" Warp-size: %d\\n\", prop.warpSize);\n }\n printf(\"\\n\");\n printf(\"Hello from CPU!\\n\");\n helloFromGPU<<<1,4>>>();\n cudaCheckError(::cudaDeviceSynchronize());\n return 0;\n}\n\n", "output": "#include \n#include \n#include <>\n#include \"../../utils/utils.cuh\"\n\n// Device (GPU) code to print thread index from the GPU.\n__global__ void helloFromGPU()\n{\n int threadId = blockIdx.x * blockDim.x + threadIdx.x;\n printf(\"Hello from GPU thread: %d!\\n\", threadId);\n}\n\n// Host code\nint main()\n{\n int nDevices;\n hipGetDeviceCount(&nDevices);\n for (int i = 0; i < nDevices; i++) {\n hipDeviceProp_t prop;\n hipGetDeviceProperties(&prop, i);\n printf(\"Device Number: %d\\n\", i);\n printf(\" Device name: %s\\n\", prop.name);\n printf(\" major-minor: %d-%d\\n\", prop.major, prop.minor);\n printf(\" Concurrent kernels: %s\\n\", prop.concurrentKernels ? \"yes\" : \"no\");\n printf(\" Concurrent computation/communication: %s\\n\\n\",prop.deviceOverlap ? \"yes\" : \"no\");\n printf(\" Managed memory supported: %s\\n\", prop.managedMemory ? \"yes\" : \"no\");\n printf(\" Memory Clock Rate (MHz): %d\\n\",\n prop.memoryClockRate/1024);\n printf(\" Memory Bus Width (bits): %d\\n\",\n prop.memoryBusWidth);\n printf(\" Peak Memory Bandwidth (GB/s): %.1f\\n\",\n 2.0*prop.memoryClockRate*(prop.memoryBusWidth/8)/1.0e6);\n printf(\" Total global memory (Gbytes): %.1f\\n\",(float)(prop.totalGlobalMem)/1024.0/1024.0/1024.0);\n printf(\" Shared memory per block (Kbytes): %.1f\\n\",(float)(prop.sharedMemPerBlock)/1024.0);\n printf(\" L2 cache size (Kbytes): %.1f\\n\",(float)(prop.l2CacheSize)/1024);\n printf(\" Warp-size: %d\\n\", prop.warpSize);\n }\n printf(\"\\n\");\n printf(\"Hello from CPU!\\n\");\n helloFromGPU<<<1,4>>>();\n cudaCheckError(::hipDeviceSynchronize());\n return 0;\n}\n\n", "source": "/kaggle/working/cuda_examples/SetupAndInitExamples/setup_check/hello_world.cu"} {"input": "#include \n#include \n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n// Simple kernel for use.\n__global__ void vector_add_kernel(float *d_v1, float *d_v2, float *d_v3, int n) {\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n if (i < n) {\n d_v3[i] = d_v1[i] + d_v2[i];\n }\n}\n\n// Structure to pass around host memory pointers.\nstruct Memory {\n float *v1;\n float *v2;\n float *v3;\n};\n\nvoid benchmark(int N, bool usePinnedMemory) {\n size_t size = N * sizeof(float);\n Memory mem;\n \n // Allocate host-pinned memory.\n if (usePinnedMemory) {\n cudaCheckError(::cudaMallocHost(&mem.v1, size));\n cudaCheckError(::cudaMallocHost(&mem.v2, size));\n cudaCheckError(::cudaMallocHost(&mem.v3, size));\n } else {\n // Allocate pageable memory.\n mem.v1 = (float*)malloc(size);\n mem.v2 = (float*)malloc(size);\n mem.v3 = (float*)malloc(size);\n }\n\n // Initialize data\n for (int i = 0; i < N; i++) {\n mem.v1[i] = 1.0f;\n mem.v2[i] = 2.0f;\n }\n\n // Allocate device memory\n float *d_v1, *d_v2, *d_v3;\n cudaCheckError(::cudaMalloc(&d_v1, size));\n cudaCheckError(::cudaMalloc(&d_v2, size));\n cudaCheckError(::cudaMalloc(&d_v3, size));\n\n // Create CUDA stream\n cudaStream_t stream;\n cudaCheckError(::cudaStreamCreate(&stream));\n\n // Create CUDA events for timing\n cudaEvent_t start, stop;\n cudaCheckError(::cudaEventCreate(&start));\n cudaCheckError(::cudaEventCreate(&stop));\n\n // Record start time\n cudaCheckError(::cudaEventRecord(start, stream));\n\n // Asynchronous memory transfers\n cudaCheckError(::cudaMemcpyAsync(d_v1, mem.v1, size, cudaMemcpyHostToDevice, stream));\n cudaCheckError(::cudaMemcpyAsync(d_v2, mem.v2, size, cudaMemcpyHostToDevice, stream));\n\n // Launch kernel\n int threadsPerBlock = 256;\n int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;\n vector_add_kernel<<>>(d_v1, d_v2, d_v3, N);\n\n // Asynchronous memory transfer back to host\n cudaCheckError(::cudaMemcpyAsync(mem.v3, d_v3, size, cudaMemcpyDeviceToHost, stream));\n\n // Record stop time\n cudaCheckError(::cudaEventRecord(stop, stream));\n cudaCheckError(::cudaEventSynchronize(stop));\n\n float milliseconds = 0;\n cudaCheckError(::cudaEventElapsedTime(&milliseconds, start, stop));\n\n // Calculate bandwidth\n float bandwidth = 3 * size / (milliseconds / 1000) / 1e9; // GB/s\n\n std::cout << (usePinnedMemory ? \"Pinned\" : \"Pageable\") << \" memory, N = \" << N \n << \", Time: \" << milliseconds << \" ms, Bandwidth: \" << bandwidth << \" GB/s\" << std::endl;\n\n // Clean up\n cudaCheckError(::cudaFree(d_v1));\n cudaCheckError(::cudaFree(d_v2));\n cudaCheckError(::cudaFree(d_v3));\n cudaCheckError(::cudaStreamDestroy(stream));\n cudaCheckError(::cudaEventDestroy(start));\n cudaCheckError(::cudaEventDestroy(stop));\n\n if (usePinnedMemory) {\n cudaCheckError(::cudaFreeHost(mem.v1));\n cudaCheckError(::cudaFreeHost(mem.v2));\n cudaCheckError(::cudaFreeHost(mem.v3));\n } else {\n free(mem.v1);\n free(mem.v2);\n free(mem.v3);\n }\n}\n\nint main() {\n // Small, to fairly large memory pushes:\n // 1024 * 4 bytes = 4KB\n // 1024 * 1024 * 4 bytes = 4MB\n // 10 * 1024 * 1024 * 4 bytes = 40MB\n std::vector sizes = {1024, 1024 * 1024, 10 * 1024 * 1024};\n \n for (int N : sizes) {\n benchmark(N, false); // Pageable memory\n benchmark(N, true); // Pinned memory\n std::cout << std::endl;\n }\n\n return 0;\n}\n", "output": "#include \n#include \n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n// Simple kernel for use.\n__global__ void vector_add_kernel(float *d_v1, float *d_v2, float *d_v3, int n) {\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n if (i < n) {\n d_v3[i] = d_v1[i] + d_v2[i];\n }\n}\n\n// Structure to pass around host memory pointers.\nstruct Memory {\n float *v1;\n float *v2;\n float *v3;\n};\n\nvoid benchmark(int N, bool usePinnedMemory) {\n size_t size = N * sizeof(float);\n Memory mem;\n \n // Allocate host-pinned memory.\n if (usePinnedMemory) {\n cudaCheckError(::hipHostMalloc(&mem.v1, size));\n cudaCheckError(::hipHostMalloc(&mem.v2, size));\n cudaCheckError(::hipHostMalloc(&mem.v3, size));\n } else {\n // Allocate pageable memory.\n mem.v1 = (float*)malloc(size);\n mem.v2 = (float*)malloc(size);\n mem.v3 = (float*)malloc(size);\n }\n\n // Initialize data\n for (int i = 0; i < N; i++) {\n mem.v1[i] = 1.0f;\n mem.v2[i] = 2.0f;\n }\n\n // Allocate device memory\n float *d_v1, *d_v2, *d_v3;\n cudaCheckError(::hipMalloc(&d_v1, size));\n cudaCheckError(::hipMalloc(&d_v2, size));\n cudaCheckError(::hipMalloc(&d_v3, size));\n\n // Create CUDA stream\n hipStream_t stream;\n cudaCheckError(::hipStreamCreate(&stream));\n\n // Create CUDA events for timing\n hipEvent_t start, stop;\n cudaCheckError(::hipEventCreate(&start));\n cudaCheckError(::hipEventCreate(&stop));\n\n // Record start time\n cudaCheckError(::hipEventRecord(start, stream));\n\n // Asynchronous memory transfers\n cudaCheckError(::hipMemcpyAsync(d_v1, mem.v1, size, hipMemcpyHostToDevice, stream));\n cudaCheckError(::hipMemcpyAsync(d_v2, mem.v2, size, hipMemcpyHostToDevice, stream));\n\n // Launch kernel\n int threadsPerBlock = 256;\n int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;\n vector_add_kernel<<>>(d_v1, d_v2, d_v3, N);\n\n // Asynchronous memory transfer back to host\n cudaCheckError(::hipMemcpyAsync(mem.v3, d_v3, size, hipMemcpyDeviceToHost, stream));\n\n // Record stop time\n cudaCheckError(::hipEventRecord(stop, stream));\n cudaCheckError(::hipEventSynchronize(stop));\n\n float milliseconds = 0;\n cudaCheckError(::hipEventElapsedTime(&milliseconds, start, stop));\n\n // Calculate bandwidth\n float bandwidth = 3 * size / (milliseconds / 1000) / 1e9; // GB/s\n\n std::cout << (usePinnedMemory ? \"Pinned\" : \"Pageable\") << \" memory, N = \" << N \n << \", Time: \" << milliseconds << \" ms, Bandwidth: \" << bandwidth << \" GB/s\" << std::endl;\n\n // Clean up\n cudaCheckError(::hipFree(d_v1));\n cudaCheckError(::hipFree(d_v2));\n cudaCheckError(::hipFree(d_v3));\n cudaCheckError(::hipStreamDestroy(stream));\n cudaCheckError(::hipEventDestroy(start));\n cudaCheckError(::hipEventDestroy(stop));\n\n if (usePinnedMemory) {\n cudaCheckError(::hipHostFree(mem.v1));\n cudaCheckError(::hipHostFree(mem.v2));\n cudaCheckError(::hipHostFree(mem.v3));\n } else {\n free(mem.v1);\n free(mem.v2);\n free(mem.v3);\n }\n}\n\nint main() {\n // Small, to fairly large memory pushes:\n // 1024 * 4 bytes = 4KB\n // 1024 * 1024 * 4 bytes = 4MB\n // 10 * 1024 * 1024 * 4 bytes = 40MB\n std::vector sizes = {1024, 1024 * 1024, 10 * 1024 * 1024};\n \n for (int N : sizes) {\n benchmark(N, false); // Pageable memory\n benchmark(N, true); // Pinned memory\n std::cout << std::endl;\n }\n\n return 0;\n}\n", "source": "/kaggle/working/cuda_examples/MemoryAndStructureExamples/pinned_vs_pageable/pinned_vs_pageable.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n// Kernel to intialize weights with random values.\n__global__ void initializeWeightsKernel(float* weights, int size, unsigned long seed) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < size) {\n curandStateXORWOW state;\n curand_init(seed, idx, 0, &state);\n weights[idx] = curand_normal(&state);\n }\n}\n\n// Function to initialize weights using cudaMalloc\nvoid initializeWeightsCudaMalloc(int size) {\n\n float* d_weights;\n // Allocate memory for distribution, but don't initialize values until kernel\n cudaCheckError(::cudaMalloc(&d_weights, size * sizeof(float)));\n\n // Launch kernel to initialize weights\n initializeWeightsKernel<<<(size + 255) / 256, 256>>>(d_weights, size, time(NULL));\n cudaCheckError(::cudaDeviceSynchronize());\n\n cudaCheckError(::cudaFree(d_weights));\n\n std::cout << \"Initialized weights using cudaMalloc.\" << std::endl;\n}\n\n// Function to initialize weights using thrust::device_vector\nvoid initializeWeightsThrust(int size) {\n\n // Create a device vector and initialize with zeros\n thrust::device_vector weights(size);\n\n // Launch kernel to initialize weights\n initializeWeightsKernel<<<(size + 255) / 256, 256>>>(thrust::raw_pointer_cast(weights.data()), size, time(NULL));\n\n std::cout << \"Initialized weights using thrust::device_vector.\" << std::endl;\n}\n\nint main() {\n const int size = 1000000; // Large number of weights\n\n // Initialize weights using thrust::device_vector\n initializeWeightsThrust(size);\n\n // Initialize weights using cudaMalloc\n initializeWeightsCudaMalloc(size);\n\n return 0;\n}\n", "output": "#include \n#include \n#include \n#include \n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n// Kernel to intialize weights with random values.\n__global__ void initializeWeightsKernel(float* weights, int size, unsigned long seed) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < size) {\n hiprandStateXORWOW state;\n hiprand_init(seed, idx, 0, &state);\n weights[idx] = hiprand_normal(&state);\n }\n}\n\n// Function to initialize weights using cudaMalloc\nvoid initializeWeightsCudaMalloc(int size) {\n\n float* d_weights;\n // Allocate memory for distribution, but don't initialize values until kernel\n cudaCheckError(::hipMalloc(&d_weights, size * sizeof(float)));\n\n // Launch kernel to initialize weights\n initializeWeightsKernel<<<(size + 255) / 256, 256>>>(d_weights, size, time(NULL));\n cudaCheckError(::hipDeviceSynchronize());\n\n cudaCheckError(::hipFree(d_weights));\n\n std::cout << \"Initialized weights using hipMalloc.\" << std::endl;\n}\n\n// Function to initialize weights using thrust::device_vector\nvoid initializeWeightsThrust(int size) {\n\n // Create a device vector and initialize with zeros\n thrust::device_vector weights(size);\n\n // Launch kernel to initialize weights\n initializeWeightsKernel<<<(size + 255) / 256, 256>>>(thrust::raw_pointer_cast(weights.data()), size, time(NULL));\n\n std::cout << \"Initialized weights using thrust::device_vector.\" << std::endl;\n}\n\nint main() {\n const int size = 1000000; // Large number of weights\n\n // Initialize weights using thrust::device_vector\n initializeWeightsThrust(size);\n\n // Initialize weights using cudaMalloc\n initializeWeightsCudaMalloc(size);\n\n return 0;\n}\n", "source": "/kaggle/working/cuda_examples/MemoryAndStructureExamples/alloc_init_vs_alloc_uninit/alloc_init_vs_alloc_uninit.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Simple demonstration of cub::BlockRadixSort\n *\n * To compile using the command line:\n * nvcc -arch=sm_XX example_block_radix_sort.cu -I../.. -lcudart -O3\n *\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console (define before including cub.h)\n#define CUB_STDERR\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../../test/test_util.h\"\n\nusing namespace cub;\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\n/// Verbose output\nbool g_verbose = false;\n\n/// Timing iterations\nint g_timing_iterations = 100;\n\n/// Default grid size\nint g_grid_size = 1;\n\n/// Uniform key samples\nbool g_uniform_keys;\n\n\n//---------------------------------------------------------------------\n// Kernels\n//---------------------------------------------------------------------\n\n/**\n * Simple kernel for performing a block-wide sorting over integers\n */\ntemplate <\n typename Key,\n int BLOCK_THREADS,\n int ITEMS_PER_THREAD>\n__launch_bounds__ (BLOCK_THREADS)\n__global__ void BlockSortKernel(\n Key *d_in, // Tile of input\n Key *d_out, // Tile of output\n clock_t *d_elapsed) // Elapsed cycle count of block scan\n{\n enum { TILE_SIZE = BLOCK_THREADS * ITEMS_PER_THREAD };\n\n // Specialize BlockLoad type for our thread block (uses warp-striped loads for coalescing, then transposes in shared memory to a blocked arrangement)\n typedef BlockLoad BlockLoadT;\n\n // Specialize BlockRadixSort type for our thread block\n typedef BlockRadixSort BlockRadixSortT;\n\n // Shared memory\n __shared__ union TempStorage\n {\n typename BlockLoadT::TempStorage load;\n typename BlockRadixSortT::TempStorage sort;\n } temp_storage;\n\n // Per-thread tile items\n Key items[ITEMS_PER_THREAD];\n\n // Our current block's offset\n int block_offset = blockIdx.x * TILE_SIZE;\n\n // Load items into a blocked arrangement\n BlockLoadT(temp_storage.load).Load(d_in + block_offset, items);\n\n // Barrier for smem reuse\n __syncthreads();\n\n // Start cycle timer\n clock_t start = clock();\n\n // Sort keys\n BlockRadixSortT(temp_storage.sort).SortBlockedToStriped(items);\n\n // Stop cycle timer\n clock_t stop = clock();\n\n // Store output in striped fashion\n StoreDirectStriped(threadIdx.x, d_out + block_offset, items);\n\n // Store elapsed clocks\n if (threadIdx.x == 0)\n {\n d_elapsed[blockIdx.x] = (start > stop) ? start - stop : stop - start;\n }\n}\n\n\n\n//---------------------------------------------------------------------\n// Host utilities\n//---------------------------------------------------------------------\n\n\n/**\n * Initialize sorting problem (and solution).\n */\ntemplate \nvoid Initialize(\n Key *h_in,\n Key *h_reference,\n int num_items,\n int tile_size)\n{\n for (int i = 0; i < num_items; ++i)\n {\n if (g_uniform_keys)\n {\n h_in[i] = 0;\n }\n else\n {\n RandomBits(h_in[i]);\n }\n h_reference[i] = h_in[i];\n }\n\n // Only sort the first tile\n std::sort(h_reference, h_reference + tile_size);\n}\n\n\n/**\n * Test BlockScan\n */\ntemplate <\n typename Key,\n int BLOCK_THREADS,\n int ITEMS_PER_THREAD>\nvoid Test()\n{\n const int TILE_SIZE = BLOCK_THREADS * ITEMS_PER_THREAD;\n\n // Allocate host arrays\n Key *h_in = new Key[TILE_SIZE * g_grid_size];\n Key *h_reference = new Key[TILE_SIZE * g_grid_size];\n clock_t *h_elapsed = new clock_t[g_grid_size];\n\n // Initialize problem and reference output on host\n Initialize(h_in, h_reference, TILE_SIZE * g_grid_size, TILE_SIZE);\n\n // Initialize device arrays\n Key *d_in = NULL;\n Key *d_out = NULL;\n clock_t *d_elapsed = NULL;\n CubDebugExit(cudaMalloc((void**)&d_in, sizeof(Key) * TILE_SIZE * g_grid_size));\n CubDebugExit(cudaMalloc((void**)&d_out, sizeof(Key) * TILE_SIZE * g_grid_size));\n CubDebugExit(cudaMalloc((void**)&d_elapsed, sizeof(clock_t) * g_grid_size));\n\n // Display input problem data\n if (g_verbose)\n {\n printf(\"Input data: \");\n for (int i = 0; i < TILE_SIZE; i++)\n std::cout << h_in[i] << \", \";\n printf(\"\\n\\n\");\n }\n\n // Kernel props\n int max_sm_occupancy;\n CubDebugExit(MaxSmOccupancy(max_sm_occupancy, BlockSortKernel, BLOCK_THREADS));\n\n // Copy problem to device\n CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(Key) * TILE_SIZE * g_grid_size, cudaMemcpyHostToDevice));\n\n printf(\"BlockRadixSort %d items (%d timing iterations, %d blocks, %d threads, %d items per thread, %d SM occupancy):\\n\",\n TILE_SIZE * g_grid_size, g_timing_iterations, g_grid_size, BLOCK_THREADS, ITEMS_PER_THREAD, max_sm_occupancy);\n fflush(stdout);\n\n // Run kernel once to prime caches and check result\n BlockSortKernel<<>>(\n d_in,\n d_out,\n d_elapsed);\n\n // Check for kernel errors and STDIO from the kernel, if any\n CubDebugExit(cudaPeekAtLastError());\n CubDebugExit(cudaDeviceSynchronize());\n\n // Check results\n printf(\"\\tOutput items: \");\n int compare = CompareDeviceResults(h_reference, d_out, TILE_SIZE, g_verbose, g_verbose);\n printf(\"%s\\n\", compare ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n fflush(stdout);\n\n // Run this several times and average the performance results\n GpuTimer timer;\n float elapsed_millis = 0.0;\n unsigned long long elapsed_clocks = 0;\n\n for (int i = 0; i < g_timing_iterations; ++i)\n {\n timer.Start();\n\n // Run kernel\n BlockSortKernel<<>>(\n d_in,\n d_out,\n d_elapsed);\n\n timer.Stop();\n elapsed_millis += timer.ElapsedMillis();\n\n // Copy clocks from device\n CubDebugExit(cudaMemcpy(h_elapsed, d_elapsed, sizeof(clock_t) * g_grid_size, cudaMemcpyDeviceToHost));\n for (int j = 0; j < g_grid_size; j++)\n {\n elapsed_clocks += h_elapsed[j];\n }\n }\n\n // Check for kernel errors and STDIO from the kernel, if any\n CubDebugExit(cudaDeviceSynchronize());\n\n // Display timing results\n float avg_millis = elapsed_millis / g_timing_iterations;\n float avg_items_per_sec = float(TILE_SIZE * g_grid_size) / avg_millis / 1000.0f;\n double avg_clocks = double(elapsed_clocks) / g_timing_iterations / g_grid_size;\n double avg_clocks_per_item = avg_clocks / TILE_SIZE;\n\n printf(\"\\tAverage BlockRadixSort::SortBlocked clocks: %.3f\\n\", avg_clocks);\n printf(\"\\tAverage BlockRadixSort::SortBlocked clocks per item: %.3f\\n\", avg_clocks_per_item);\n printf(\"\\tAverage kernel millis: %.4f\\n\", avg_millis);\n printf(\"\\tAverage million items / sec: %.4f\\n\", avg_items_per_sec);\n fflush(stdout);\n\n // Cleanup\n if (h_in) delete[] h_in;\n if (h_reference) delete[] h_reference;\n if (h_elapsed) delete[] h_elapsed;\n if (d_in) CubDebugExit(cudaFree(d_in));\n if (d_out) CubDebugExit(cudaFree(d_out));\n if (d_elapsed) CubDebugExit(cudaFree(d_elapsed));\n}\n\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n g_uniform_keys = args.CheckCmdLineFlag(\"uniform\");\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n args.GetCmdLineArgument(\"grid-size\", g_grid_size);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=] \"\n \"[--i=]\"\n \"[--grid-size=]\"\n \"[--v] \"\n \"\\n\", argv[0], g_timing_iterations, g_grid_size);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n fflush(stdout);\n\n // Run tests\n printf(\"\\nuint32:\\n\"); fflush(stdout);\n Test();\n printf(\"\\n\"); fflush(stdout);\n\n printf(\"\\nfp32:\\n\"); fflush(stdout);\n Test();\n printf(\"\\n\"); fflush(stdout);\n\n printf(\"\\nuint8:\\n\"); fflush(stdout);\n Test();\n printf(\"\\n\"); fflush(stdout);\n\n return 0;\n}\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Simple demonstration of cub::BlockRadixSort\n *\n * To compile using the command line:\n * nvcc -arch=sm_XX example_block_radix_sort.cu -I../.. -lcudart -O3\n *\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console (define before including cub.h)\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../../test/test_util.h\"\n\nusing namespace hipcub;\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\n/// Verbose output\nbool g_verbose = false;\n\n/// Timing iterations\nint g_timing_iterations = 100;\n\n/// Default grid size\nint g_grid_size = 1;\n\n/// Uniform key samples\nbool g_uniform_keys;\n\n\n//---------------------------------------------------------------------\n// Kernels\n//---------------------------------------------------------------------\n\n/**\n * Simple kernel for performing a block-wide sorting over integers\n */\ntemplate <\n typename Key,\n int BLOCK_THREADS,\n int ITEMS_PER_THREAD>\n__launch_bounds__ (BLOCK_THREADS)\n__global__ void BlockSortKernel(\n Key *d_in, // Tile of input\n Key *d_out, // Tile of output\n clock_t *d_elapsed) // Elapsed cycle count of block scan\n{\n enum { TILE_SIZE = BLOCK_THREADS * ITEMS_PER_THREAD };\n\n // Specialize BlockLoad type for our thread block (uses warp-striped loads for coalescing, then transposes in shared memory to a blocked arrangement)\n typedef BlockLoad BlockLoadT;\n\n // Specialize BlockRadixSort type for our thread block\n typedef BlockRadixSort BlockRadixSortT;\n\n // Shared memory\n __shared__ union TempStorage\n {\n typename BlockLoadT::TempStorage load;\n typename BlockRadixSortT::TempStorage sort;\n } temp_storage;\n\n // Per-thread tile items\n Key items[ITEMS_PER_THREAD];\n\n // Our current block's offset\n int block_offset = blockIdx.x * TILE_SIZE;\n\n // Load items into a blocked arrangement\n BlockLoadT(temp_storage.load).Load(d_in + block_offset, items);\n\n // Barrier for smem reuse\n __syncthreads();\n\n // Start cycle timer\n clock_t start = clock();\n\n // Sort keys\n BlockRadixSortT(temp_storage.sort).SortBlockedToStriped(items);\n\n // Stop cycle timer\n clock_t stop = clock();\n\n // Store output in striped fashion\n StoreDirectStriped(threadIdx.x, d_out + block_offset, items);\n\n // Store elapsed clocks\n if (threadIdx.x == 0)\n {\n d_elapsed[blockIdx.x] = (start > stop) ? start - stop : stop - start;\n }\n}\n\n\n\n//---------------------------------------------------------------------\n// Host utilities\n//---------------------------------------------------------------------\n\n\n/**\n * Initialize sorting problem (and solution).\n */\ntemplate \nvoid Initialize(\n Key *h_in,\n Key *h_reference,\n int num_items,\n int tile_size)\n{\n for (int i = 0; i < num_items; ++i)\n {\n if (g_uniform_keys)\n {\n h_in[i] = 0;\n }\n else\n {\n RandomBits(h_in[i]);\n }\n h_reference[i] = h_in[i];\n }\n\n // Only sort the first tile\n std::sort(h_reference, h_reference + tile_size);\n}\n\n\n/**\n * Test BlockScan\n */\ntemplate <\n typename Key,\n int BLOCK_THREADS,\n int ITEMS_PER_THREAD>\nvoid Test()\n{\n const int TILE_SIZE = BLOCK_THREADS * ITEMS_PER_THREAD;\n\n // Allocate host arrays\n Key *h_in = new Key[TILE_SIZE * g_grid_size];\n Key *h_reference = new Key[TILE_SIZE * g_grid_size];\n clock_t *h_elapsed = new clock_t[g_grid_size];\n\n // Initialize problem and reference output on host\n Initialize(h_in, h_reference, TILE_SIZE * g_grid_size, TILE_SIZE);\n\n // Initialize device arrays\n Key *d_in = NULL;\n Key *d_out = NULL;\n clock_t *d_elapsed = NULL;\n CubDebugExit(hipMalloc((void**)&d_in, sizeof(Key) * TILE_SIZE * g_grid_size));\n CubDebugExit(hipMalloc((void**)&d_out, sizeof(Key) * TILE_SIZE * g_grid_size));\n CubDebugExit(hipMalloc((void**)&d_elapsed, sizeof(clock_t) * g_grid_size));\n\n // Display input problem data\n if (g_verbose)\n {\n printf(\"Input data: \");\n for (int i = 0; i < TILE_SIZE; i++)\n std::cout << h_in[i] << \", \";\n printf(\"\\n\\n\");\n }\n\n // Kernel props\n int max_sm_occupancy;\n CubDebugExit(MaxSmOccupancy(max_sm_occupancy, BlockSortKernel, BLOCK_THREADS));\n\n // Copy problem to device\n CubDebugExit(hipMemcpy(d_in, h_in, sizeof(Key) * TILE_SIZE * g_grid_size, hipMemcpyHostToDevice));\n\n printf(\"BlockRadixSort %d items (%d timing iterations, %d blocks, %d threads, %d items per thread, %d SM occupancy):\\n\",\n TILE_SIZE * g_grid_size, g_timing_iterations, g_grid_size, BLOCK_THREADS, ITEMS_PER_THREAD, max_sm_occupancy);\n fflush(stdout);\n\n // Run kernel once to prime caches and check result\n BlockSortKernel<<>>(\n d_in,\n d_out,\n d_elapsed);\n\n // Check for kernel errors and STDIO from the kernel, if any\n CubDebugExit(hipPeekAtLastError());\n CubDebugExit(hipDeviceSynchronize());\n\n // Check results\n printf(\"\\tOutput items: \");\n int compare = CompareDeviceResults(h_reference, d_out, TILE_SIZE, g_verbose, g_verbose);\n printf(\"%s\\n\", compare ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n fflush(stdout);\n\n // Run this several times and average the performance results\n GpuTimer timer;\n float elapsed_millis = 0.0;\n unsigned long long elapsed_clocks = 0;\n\n for (int i = 0; i < g_timing_iterations; ++i)\n {\n timer.Start();\n\n // Run kernel\n BlockSortKernel<<>>(\n d_in,\n d_out,\n d_elapsed);\n\n timer.Stop();\n elapsed_millis += timer.ElapsedMillis();\n\n // Copy clocks from device\n CubDebugExit(hipMemcpy(h_elapsed, d_elapsed, sizeof(clock_t) * g_grid_size, hipMemcpyDeviceToHost));\n for (int j = 0; j < g_grid_size; j++)\n {\n elapsed_clocks += h_elapsed[j];\n }\n }\n\n // Check for kernel errors and STDIO from the kernel, if any\n CubDebugExit(hipDeviceSynchronize());\n\n // Display timing results\n float avg_millis = elapsed_millis / g_timing_iterations;\n float avg_items_per_sec = float(TILE_SIZE * g_grid_size) / avg_millis / 1000.0f;\n double avg_clocks = double(elapsed_clocks) / g_timing_iterations / g_grid_size;\n double avg_clocks_per_item = avg_clocks / TILE_SIZE;\n\n printf(\"\\tAverage BlockRadixSort::SortBlocked clocks: %.3f\\n\", avg_clocks);\n printf(\"\\tAverage BlockRadixSort::SortBlocked clocks per item: %.3f\\n\", avg_clocks_per_item);\n printf(\"\\tAverage kernel millis: %.4f\\n\", avg_millis);\n printf(\"\\tAverage million items / sec: %.4f\\n\", avg_items_per_sec);\n fflush(stdout);\n\n // Cleanup\n if (h_in) delete[] h_in;\n if (h_reference) delete[] h_reference;\n if (h_elapsed) delete[] h_elapsed;\n if (d_in) CubDebugExit(hipFree(d_in));\n if (d_out) CubDebugExit(hipFree(d_out));\n if (d_elapsed) CubDebugExit(hipFree(d_elapsed));\n}\n\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n g_uniform_keys = args.CheckCmdLineFlag(\"uniform\");\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n args.GetCmdLineArgument(\"grid-size\", g_grid_size);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=] \"\n \"[--i=]\"\n \"[--grid-size=]\"\n \"[--v] \"\n \"\\n\", argv[0], g_timing_iterations, g_grid_size);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n fflush(stdout);\n\n // Run tests\n printf(\"\\nuint32:\\n\"); fflush(stdout);\n Test();\n printf(\"\\n\"); fflush(stdout);\n\n printf(\"\\nfp32:\\n\"); fflush(stdout);\n Test();\n printf(\"\\n\"); fflush(stdout);\n\n printf(\"\\nuint8:\\n\"); fflush(stdout);\n Test();\n printf(\"\\n\"); fflush(stdout);\n\n return 0;\n}\n\n", "source": "/kaggle/working/cub/examples/block/example_block_radix_sort.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Simple demonstration of cub::BlockReduce\n *\n * To compile using the command line:\n * nvcc -arch=sm_XX example_block_reduce.cu -I../.. -lcudart -O3\n *\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console (define before including cub.h)\n#define CUB_STDERR\n\n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../../test/test_util.h\"\n\nusing namespace cub;\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\n/// Verbose output\nbool g_verbose = false;\n\n/// Timing iterations\nint g_timing_iterations = 100;\n\n/// Default grid size\nint g_grid_size = 1;\n\n\n\n//---------------------------------------------------------------------\n// Kernels\n//---------------------------------------------------------------------\n\n/**\n * Simple kernel for performing a block-wide reduction.\n */\ntemplate <\n int BLOCK_THREADS,\n int ITEMS_PER_THREAD,\n BlockReduceAlgorithm ALGORITHM>\n__global__ void BlockReduceKernel(\n int *d_in, // Tile of input\n int *d_out, // Tile aggregate\n clock_t *d_elapsed) // Elapsed cycle count of block reduction\n{\n // Specialize BlockReduce type for our thread block\n typedef BlockReduce BlockReduceT;\n\n // Shared memory\n __shared__ typename BlockReduceT::TempStorage temp_storage;\n\n // Per-thread tile data\n int data[ITEMS_PER_THREAD];\n LoadDirectStriped(threadIdx.x, d_in, data);\n\n // Start cycle timer\n clock_t start = clock();\n\n // Compute sum\n int aggregate = BlockReduceT(temp_storage).Sum(data);\n\n // Stop cycle timer\n clock_t stop = clock();\n\n // Store aggregate and elapsed clocks\n if (threadIdx.x == 0)\n {\n *d_elapsed = (start > stop) ? start - stop : stop - start;\n *d_out = aggregate;\n }\n}\n\n\n\n//---------------------------------------------------------------------\n// Host utilities\n//---------------------------------------------------------------------\n\n/**\n * Initialize reduction problem (and solution).\n * Returns the aggregate\n */\nint Initialize(int *h_in, int num_items)\n{\n int inclusive = 0;\n\n for (int i = 0; i < num_items; ++i)\n {\n h_in[i] = i % 17;\n inclusive += h_in[i];\n }\n\n return inclusive;\n}\n\n\n/**\n * Test thread block reduction\n */\ntemplate <\n int BLOCK_THREADS,\n int ITEMS_PER_THREAD,\n BlockReduceAlgorithm ALGORITHM>\nvoid Test()\n{\n const int TILE_SIZE = BLOCK_THREADS * ITEMS_PER_THREAD;\n\n // Allocate host arrays\n int *h_in = new int[TILE_SIZE];\n int *h_gpu = new int[TILE_SIZE + 1];\n\n // Initialize problem and reference output on host\n int h_aggregate = Initialize(h_in, TILE_SIZE);\n\n // Initialize device arrays\n int *d_in = NULL;\n int *d_out = NULL;\n clock_t *d_elapsed = NULL;\n cudaMalloc((void**)&d_in, sizeof(int) * TILE_SIZE);\n cudaMalloc((void**)&d_out, sizeof(int) * 1);\n cudaMalloc((void**)&d_elapsed, sizeof(clock_t));\n\n // Display input problem data\n if (g_verbose)\n {\n printf(\"Input data: \");\n for (int i = 0; i < TILE_SIZE; i++)\n printf(\"%d, \", h_in[i]);\n printf(\"\\n\\n\");\n }\n\n // Kernel props\n int max_sm_occupancy;\n CubDebugExit(MaxSmOccupancy(max_sm_occupancy, BlockReduceKernel, BLOCK_THREADS));\n\n // Copy problem to device\n cudaMemcpy(d_in, h_in, sizeof(int) * TILE_SIZE, cudaMemcpyHostToDevice);\n\n printf(\"BlockReduce algorithm %s on %d items (%d timing iterations, %d blocks, %d threads, %d items per thread, %d SM occupancy):\\n\",\n (ALGORITHM == BLOCK_REDUCE_RAKING) ? \"BLOCK_REDUCE_RAKING\" : \"BLOCK_REDUCE_WARP_REDUCTIONS\",\n TILE_SIZE, g_timing_iterations, g_grid_size, BLOCK_THREADS, ITEMS_PER_THREAD, max_sm_occupancy);\n\n // Run kernel\n BlockReduceKernel<<>>(\n d_in,\n d_out,\n d_elapsed);\n\n // Check total aggregate\n printf(\"\\tAggregate: \");\n int compare = CompareDeviceResults(&h_aggregate, d_out, 1, g_verbose, g_verbose);\n printf(\"%s\\n\", compare ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Run this several times and average the performance results\n GpuTimer timer;\n float elapsed_millis = 0.0;\n clock_t elapsed_clocks = 0;\n\n for (int i = 0; i < g_timing_iterations; ++i)\n {\n // Copy problem to device\n cudaMemcpy(d_in, h_in, sizeof(int) * TILE_SIZE, cudaMemcpyHostToDevice);\n\n timer.Start();\n\n // Run kernel\n BlockReduceKernel<<>>(\n d_in,\n d_out,\n d_elapsed);\n\n timer.Stop();\n elapsed_millis += timer.ElapsedMillis();\n\n // Copy clocks from device\n clock_t clocks;\n CubDebugExit(cudaMemcpy(&clocks, d_elapsed, sizeof(clock_t), cudaMemcpyDeviceToHost));\n elapsed_clocks += clocks;\n\n }\n\n // Check for kernel errors and STDIO from the kernel, if any\n CubDebugExit(cudaPeekAtLastError());\n CubDebugExit(cudaDeviceSynchronize());\n\n // Display timing results\n float avg_millis = elapsed_millis / g_timing_iterations;\n float avg_items_per_sec = float(TILE_SIZE * g_grid_size) / avg_millis / 1000.0f;\n float avg_clocks = float(elapsed_clocks) / g_timing_iterations;\n float avg_clocks_per_item = avg_clocks / TILE_SIZE;\n\n printf(\"\\tAverage BlockReduce::Sum clocks: %.3f\\n\", avg_clocks);\n printf(\"\\tAverage BlockReduce::Sum clocks per item: %.3f\\n\", avg_clocks_per_item);\n printf(\"\\tAverage kernel millis: %.4f\\n\", avg_millis);\n printf(\"\\tAverage million items / sec: %.4f\\n\", avg_items_per_sec);\n\n // Cleanup\n if (h_in) delete[] h_in;\n if (h_gpu) delete[] h_gpu;\n if (d_in) cudaFree(d_in);\n if (d_out) cudaFree(d_out);\n if (d_elapsed) cudaFree(d_elapsed);\n}\n\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n args.GetCmdLineArgument(\"grid-size\", g_grid_size);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=] \"\n \"[--i=] \"\n \"[--grid-size=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // Run tests\n Test<1024, 1, BLOCK_REDUCE_RAKING>();\n Test<512, 2, BLOCK_REDUCE_RAKING>();\n Test<256, 4, BLOCK_REDUCE_RAKING>();\n Test<128, 8, BLOCK_REDUCE_RAKING>();\n Test<64, 16, BLOCK_REDUCE_RAKING>();\n Test<32, 32, BLOCK_REDUCE_RAKING>();\n Test<16, 64, BLOCK_REDUCE_RAKING>();\n\n printf(\"-------------\\n\");\n\n Test<1024, 1, BLOCK_REDUCE_WARP_REDUCTIONS>();\n Test<512, 2, BLOCK_REDUCE_WARP_REDUCTIONS>();\n Test<256, 4, BLOCK_REDUCE_WARP_REDUCTIONS>();\n Test<128, 8, BLOCK_REDUCE_WARP_REDUCTIONS>();\n Test<64, 16, BLOCK_REDUCE_WARP_REDUCTIONS>();\n Test<32, 32, BLOCK_REDUCE_WARP_REDUCTIONS>();\n Test<16, 64, BLOCK_REDUCE_WARP_REDUCTIONS>();\n\n return 0;\n}\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Simple demonstration of cub::BlockReduce\n *\n * To compile using the command line:\n * nvcc -arch=sm_XX example_block_reduce.cu -I../.. -lcudart -O3\n *\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console (define before including cub.h)\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../../test/test_util.h\"\n\nusing namespace hipcub;\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\n/// Verbose output\nbool g_verbose = false;\n\n/// Timing iterations\nint g_timing_iterations = 100;\n\n/// Default grid size\nint g_grid_size = 1;\n\n\n\n//---------------------------------------------------------------------\n// Kernels\n//---------------------------------------------------------------------\n\n/**\n * Simple kernel for performing a block-wide reduction.\n */\ntemplate <\n int BLOCK_THREADS,\n int ITEMS_PER_THREAD,\n BlockReduceAlgorithm ALGORITHM>\n__global__ void BlockReduceKernel(\n int *d_in, // Tile of input\n int *d_out, // Tile aggregate\n clock_t *d_elapsed) // Elapsed cycle count of block reduction\n{\n // Specialize BlockReduce type for our thread block\n typedef BlockReduce BlockReduceT;\n\n // Shared memory\n __shared__ typename BlockReduceT::TempStorage temp_storage;\n\n // Per-thread tile data\n int data[ITEMS_PER_THREAD];\n LoadDirectStriped(threadIdx.x, d_in, data);\n\n // Start cycle timer\n clock_t start = clock();\n\n // Compute sum\n int aggregate = BlockReduceT(temp_storage).Sum(data);\n\n // Stop cycle timer\n clock_t stop = clock();\n\n // Store aggregate and elapsed clocks\n if (threadIdx.x == 0)\n {\n *d_elapsed = (start > stop) ? start - stop : stop - start;\n *d_out = aggregate;\n }\n}\n\n\n\n//---------------------------------------------------------------------\n// Host utilities\n//---------------------------------------------------------------------\n\n/**\n * Initialize reduction problem (and solution).\n * Returns the aggregate\n */\nint Initialize(int *h_in, int num_items)\n{\n int inclusive = 0;\n\n for (int i = 0; i < num_items; ++i)\n {\n h_in[i] = i % 17;\n inclusive += h_in[i];\n }\n\n return inclusive;\n}\n\n\n/**\n * Test thread block reduction\n */\ntemplate <\n int BLOCK_THREADS,\n int ITEMS_PER_THREAD,\n BlockReduceAlgorithm ALGORITHM>\nvoid Test()\n{\n const int TILE_SIZE = BLOCK_THREADS * ITEMS_PER_THREAD;\n\n // Allocate host arrays\n int *h_in = new int[TILE_SIZE];\n int *h_gpu = new int[TILE_SIZE + 1];\n\n // Initialize problem and reference output on host\n int h_aggregate = Initialize(h_in, TILE_SIZE);\n\n // Initialize device arrays\n int *d_in = NULL;\n int *d_out = NULL;\n clock_t *d_elapsed = NULL;\n hipMalloc((void**)&d_in, sizeof(int) * TILE_SIZE);\n hipMalloc((void**)&d_out, sizeof(int) * 1);\n hipMalloc((void**)&d_elapsed, sizeof(clock_t));\n\n // Display input problem data\n if (g_verbose)\n {\n printf(\"Input data: \");\n for (int i = 0; i < TILE_SIZE; i++)\n printf(\"%d, \", h_in[i]);\n printf(\"\\n\\n\");\n }\n\n // Kernel props\n int max_sm_occupancy;\n CubDebugExit(MaxSmOccupancy(max_sm_occupancy, BlockReduceKernel, BLOCK_THREADS));\n\n // Copy problem to device\n hipMemcpy(d_in, h_in, sizeof(int) * TILE_SIZE, hipMemcpyHostToDevice);\n\n printf(\"BlockReduce algorithm %s on %d items (%d timing iterations, %d blocks, %d threads, %d items per thread, %d SM occupancy):\\n\",\n (ALGORITHM == BLOCK_REDUCE_RAKING) ? \"BLOCK_REDUCE_RAKING\" : \"BLOCK_REDUCE_WARP_REDUCTIONS\",\n TILE_SIZE, g_timing_iterations, g_grid_size, BLOCK_THREADS, ITEMS_PER_THREAD, max_sm_occupancy);\n\n // Run kernel\n BlockReduceKernel<<>>(\n d_in,\n d_out,\n d_elapsed);\n\n // Check total aggregate\n printf(\"\\tAggregate: \");\n int compare = CompareDeviceResults(&h_aggregate, d_out, 1, g_verbose, g_verbose);\n printf(\"%s\\n\", compare ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Run this several times and average the performance results\n GpuTimer timer;\n float elapsed_millis = 0.0;\n clock_t elapsed_clocks = 0;\n\n for (int i = 0; i < g_timing_iterations; ++i)\n {\n // Copy problem to device\n hipMemcpy(d_in, h_in, sizeof(int) * TILE_SIZE, hipMemcpyHostToDevice);\n\n timer.Start();\n\n // Run kernel\n BlockReduceKernel<<>>(\n d_in,\n d_out,\n d_elapsed);\n\n timer.Stop();\n elapsed_millis += timer.ElapsedMillis();\n\n // Copy clocks from device\n clock_t clocks;\n CubDebugExit(hipMemcpy(&clocks, d_elapsed, sizeof(clock_t), hipMemcpyDeviceToHost));\n elapsed_clocks += clocks;\n\n }\n\n // Check for kernel errors and STDIO from the kernel, if any\n CubDebugExit(hipPeekAtLastError());\n CubDebugExit(hipDeviceSynchronize());\n\n // Display timing results\n float avg_millis = elapsed_millis / g_timing_iterations;\n float avg_items_per_sec = float(TILE_SIZE * g_grid_size) / avg_millis / 1000.0f;\n float avg_clocks = float(elapsed_clocks) / g_timing_iterations;\n float avg_clocks_per_item = avg_clocks / TILE_SIZE;\n\n printf(\"\\tAverage BlockReduce::Sum clocks: %.3f\\n\", avg_clocks);\n printf(\"\\tAverage BlockReduce::Sum clocks per item: %.3f\\n\", avg_clocks_per_item);\n printf(\"\\tAverage kernel millis: %.4f\\n\", avg_millis);\n printf(\"\\tAverage million items / sec: %.4f\\n\", avg_items_per_sec);\n\n // Cleanup\n if (h_in) delete[] h_in;\n if (h_gpu) delete[] h_gpu;\n if (d_in) hipFree(d_in);\n if (d_out) hipFree(d_out);\n if (d_elapsed) hipFree(d_elapsed);\n}\n\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n args.GetCmdLineArgument(\"grid-size\", g_grid_size);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=] \"\n \"[--i=] \"\n \"[--grid-size=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // Run tests\n Test<1024, 1, BLOCK_REDUCE_RAKING>();\n Test<512, 2, BLOCK_REDUCE_RAKING>();\n Test<256, 4, BLOCK_REDUCE_RAKING>();\n Test<128, 8, BLOCK_REDUCE_RAKING>();\n Test<64, 16, BLOCK_REDUCE_RAKING>();\n Test<32, 32, BLOCK_REDUCE_RAKING>();\n Test<16, 64, BLOCK_REDUCE_RAKING>();\n\n printf(\"-------------\\n\");\n\n Test<1024, 1, BLOCK_REDUCE_WARP_REDUCTIONS>();\n Test<512, 2, BLOCK_REDUCE_WARP_REDUCTIONS>();\n Test<256, 4, BLOCK_REDUCE_WARP_REDUCTIONS>();\n Test<128, 8, BLOCK_REDUCE_WARP_REDUCTIONS>();\n Test<64, 16, BLOCK_REDUCE_WARP_REDUCTIONS>();\n Test<32, 32, BLOCK_REDUCE_WARP_REDUCTIONS>();\n Test<16, 64, BLOCK_REDUCE_WARP_REDUCTIONS>();\n\n return 0;\n}\n\n", "source": "/kaggle/working/cub/examples/block/example_block_reduce.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Simple demonstration of cub::BlockScan\n *\n * To compile using the command line:\n * nvcc -arch=sm_XX example_block_scan.cu -I../.. -lcudart -O3\n *\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console (define before including cub.h)\n#define CUB_STDERR\n\n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../../test/test_util.h\"\n\nusing namespace cub;\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\n/// Verbose output\nbool g_verbose = false;\n\n/// Timing iterations\nint g_timing_iterations = 100;\n\n/// Default grid size\nint g_grid_size = 1;\n\n\n\n//---------------------------------------------------------------------\n// Kernels\n//---------------------------------------------------------------------\n\n/**\n * Simple kernel for performing a block-wide exclusive prefix sum over integers\n */\ntemplate <\n int BLOCK_THREADS,\n int ITEMS_PER_THREAD,\n BlockScanAlgorithm ALGORITHM>\n__global__ void BlockPrefixSumKernel(\n int *d_in, // Tile of input\n int *d_out, // Tile of output\n clock_t *d_elapsed) // Elapsed cycle count of block scan\n{\n // Specialize BlockLoad type for our thread block (uses warp-striped loads for coalescing, then transposes in shared memory to a blocked arrangement)\n typedef BlockLoad BlockLoadT;\n\n // Specialize BlockStore type for our thread block (uses warp-striped loads for coalescing, then transposes in shared memory to a blocked arrangement)\n typedef BlockStore BlockStoreT;\n\n // Specialize BlockScan type for our thread block\n typedef BlockScan BlockScanT;\n\n // Shared memory\n __shared__ union TempStorage\n {\n typename BlockLoadT::TempStorage load;\n typename BlockStoreT::TempStorage store;\n typename BlockScanT::TempStorage scan;\n } temp_storage;\n\n // Per-thread tile data\n int data[ITEMS_PER_THREAD];\n\n // Load items into a blocked arrangement\n BlockLoadT(temp_storage.load).Load(d_in, data);\n\n // Barrier for smem reuse\n __syncthreads();\n\n // Start cycle timer\n clock_t start = clock();\n\n // Compute exclusive prefix sum\n int aggregate;\n BlockScanT(temp_storage.scan).ExclusiveSum(data, data, aggregate);\n\n // Stop cycle timer\n clock_t stop = clock();\n\n // Barrier for smem reuse\n __syncthreads();\n\n // Store items from a blocked arrangement\n BlockStoreT(temp_storage.store).Store(d_out, data);\n\n // Store aggregate and elapsed clocks\n if (threadIdx.x == 0)\n {\n *d_elapsed = (start > stop) ? start - stop : stop - start;\n d_out[BLOCK_THREADS * ITEMS_PER_THREAD] = aggregate;\n }\n}\n\n\n\n//---------------------------------------------------------------------\n// Host utilities\n//---------------------------------------------------------------------\n\n/**\n * Initialize exclusive prefix sum problem (and solution).\n * Returns the aggregate\n */\nint Initialize(\n int *h_in,\n int *h_reference,\n int num_items)\n{\n int inclusive = 0;\n\n for (int i = 0; i < num_items; ++i)\n {\n h_in[i] = i % 17;\n\n h_reference[i] = inclusive;\n inclusive += h_in[i];\n }\n\n return inclusive;\n}\n\n\n/**\n * Test thread block scan\n */\ntemplate <\n int BLOCK_THREADS,\n int ITEMS_PER_THREAD,\n BlockScanAlgorithm ALGORITHM>\nvoid Test()\n{\n const int TILE_SIZE = BLOCK_THREADS * ITEMS_PER_THREAD;\n\n // Allocate host arrays\n int *h_in = new int[TILE_SIZE];\n int *h_reference = new int[TILE_SIZE];\n int *h_gpu = new int[TILE_SIZE + 1];\n\n // Initialize problem and reference output on host\n int h_aggregate = Initialize(h_in, h_reference, TILE_SIZE);\n\n // Initialize device arrays\n int *d_in = NULL;\n int *d_out = NULL;\n clock_t *d_elapsed = NULL;\n cudaMalloc((void**)&d_in, sizeof(int) * TILE_SIZE);\n cudaMalloc((void**)&d_out, sizeof(int) * (TILE_SIZE + 1));\n cudaMalloc((void**)&d_elapsed, sizeof(clock_t));\n\n // Display input problem data\n if (g_verbose)\n {\n printf(\"Input data: \");\n for (int i = 0; i < TILE_SIZE; i++)\n printf(\"%d, \", h_in[i]);\n printf(\"\\n\\n\");\n }\n\n // Kernel props\n int max_sm_occupancy;\n CubDebugExit(MaxSmOccupancy(max_sm_occupancy, BlockPrefixSumKernel, BLOCK_THREADS));\n\n // Copy problem to device\n cudaMemcpy(d_in, h_in, sizeof(int) * TILE_SIZE, cudaMemcpyHostToDevice);\n\n printf(\"BlockScan algorithm %s on %d items (%d timing iterations, %d blocks, %d threads, %d items per thread, %d SM occupancy):\\n\",\n (ALGORITHM == BLOCK_SCAN_RAKING) ? \"BLOCK_SCAN_RAKING\" : (ALGORITHM == BLOCK_SCAN_RAKING_MEMOIZE) ? \"BLOCK_SCAN_RAKING_MEMOIZE\" : \"BLOCK_SCAN_WARP_SCANS\",\n TILE_SIZE, g_timing_iterations, g_grid_size, BLOCK_THREADS, ITEMS_PER_THREAD, max_sm_occupancy);\n\n // Run aggregate/prefix kernel\n BlockPrefixSumKernel<<>>(\n d_in,\n d_out,\n d_elapsed);\n\n // Check results\n printf(\"\\tOutput items: \");\n int compare = CompareDeviceResults(h_reference, d_out, TILE_SIZE, g_verbose, g_verbose);\n printf(\"%s\\n\", compare ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Check total aggregate\n printf(\"\\tAggregate: \");\n compare = CompareDeviceResults(&h_aggregate, d_out + TILE_SIZE, 1, g_verbose, g_verbose);\n printf(\"%s\\n\", compare ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Run this several times and average the performance results\n GpuTimer timer;\n float elapsed_millis = 0.0;\n clock_t elapsed_clocks = 0;\n\n for (int i = 0; i < g_timing_iterations; ++i)\n {\n // Copy problem to device\n cudaMemcpy(d_in, h_in, sizeof(int) * TILE_SIZE, cudaMemcpyHostToDevice);\n\n timer.Start();\n\n // Run aggregate/prefix kernel\n BlockPrefixSumKernel<<>>(\n d_in,\n d_out,\n d_elapsed);\n\n timer.Stop();\n elapsed_millis += timer.ElapsedMillis();\n\n // Copy clocks from device\n clock_t clocks;\n CubDebugExit(cudaMemcpy(&clocks, d_elapsed, sizeof(clock_t), cudaMemcpyDeviceToHost));\n elapsed_clocks += clocks;\n\n }\n\n // Check for kernel errors and STDIO from the kernel, if any\n CubDebugExit(cudaPeekAtLastError());\n CubDebugExit(cudaDeviceSynchronize());\n\n // Display timing results\n float avg_millis = elapsed_millis / g_timing_iterations;\n float avg_items_per_sec = float(TILE_SIZE * g_grid_size) / avg_millis / 1000.0f;\n float avg_clocks = float(elapsed_clocks) / g_timing_iterations;\n float avg_clocks_per_item = avg_clocks / TILE_SIZE;\n\n printf(\"\\tAverage BlockScan::Sum clocks: %.3f\\n\", avg_clocks);\n printf(\"\\tAverage BlockScan::Sum clocks per item: %.3f\\n\", avg_clocks_per_item);\n printf(\"\\tAverage kernel millis: %.4f\\n\", avg_millis);\n printf(\"\\tAverage million items / sec: %.4f\\n\", avg_items_per_sec);\n\n // Cleanup\n if (h_in) delete[] h_in;\n if (h_reference) delete[] h_reference;\n if (h_gpu) delete[] h_gpu;\n if (d_in) cudaFree(d_in);\n if (d_out) cudaFree(d_out);\n if (d_elapsed) cudaFree(d_elapsed);\n}\n\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n args.GetCmdLineArgument(\"grid-size\", g_grid_size);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=] \"\n \"[--i=]\"\n \"[--grid-size=]\"\n \"[--v] \"\n \"\\n\", argv[0], g_timing_iterations, g_grid_size);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // Run tests\n Test<1024, 1, BLOCK_SCAN_RAKING>();\n Test<512, 2, BLOCK_SCAN_RAKING>();\n Test<256, 4, BLOCK_SCAN_RAKING>();\n Test<128, 8, BLOCK_SCAN_RAKING>();\n Test<64, 16, BLOCK_SCAN_RAKING>();\n Test<32, 32, BLOCK_SCAN_RAKING>();\n\n printf(\"-------------\\n\");\n\n Test<1024, 1, BLOCK_SCAN_RAKING_MEMOIZE>();\n Test<512, 2, BLOCK_SCAN_RAKING_MEMOIZE>();\n Test<256, 4, BLOCK_SCAN_RAKING_MEMOIZE>();\n Test<128, 8, BLOCK_SCAN_RAKING_MEMOIZE>();\n Test<64, 16, BLOCK_SCAN_RAKING_MEMOIZE>();\n Test<32, 32, BLOCK_SCAN_RAKING_MEMOIZE>();\n\n printf(\"-------------\\n\");\n\n Test<1024, 1, BLOCK_SCAN_WARP_SCANS>();\n Test<512, 2, BLOCK_SCAN_WARP_SCANS>();\n Test<256, 4, BLOCK_SCAN_WARP_SCANS>();\n Test<128, 8, BLOCK_SCAN_WARP_SCANS>();\n Test<64, 16, BLOCK_SCAN_WARP_SCANS>();\n Test<32, 32, BLOCK_SCAN_WARP_SCANS>();\n\n\n return 0;\n}\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Simple demonstration of cub::BlockScan\n *\n * To compile using the command line:\n * nvcc -arch=sm_XX example_block_scan.cu -I../.. -lcudart -O3\n *\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console (define before including cub.h)\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../../test/test_util.h\"\n\nusing namespace hipcub;\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\n/// Verbose output\nbool g_verbose = false;\n\n/// Timing iterations\nint g_timing_iterations = 100;\n\n/// Default grid size\nint g_grid_size = 1;\n\n\n\n//---------------------------------------------------------------------\n// Kernels\n//---------------------------------------------------------------------\n\n/**\n * Simple kernel for performing a block-wide exclusive prefix sum over integers\n */\ntemplate <\n int BLOCK_THREADS,\n int ITEMS_PER_THREAD,\n BlockScanAlgorithm ALGORITHM>\n__global__ void BlockPrefixSumKernel(\n int *d_in, // Tile of input\n int *d_out, // Tile of output\n clock_t *d_elapsed) // Elapsed cycle count of block scan\n{\n // Specialize BlockLoad type for our thread block (uses warp-striped loads for coalescing, then transposes in shared memory to a blocked arrangement)\n typedef BlockLoad BlockLoadT;\n\n // Specialize BlockStore type for our thread block (uses warp-striped loads for coalescing, then transposes in shared memory to a blocked arrangement)\n typedef BlockStore BlockStoreT;\n\n // Specialize BlockScan type for our thread block\n typedef BlockScan BlockScanT;\n\n // Shared memory\n __shared__ union TempStorage\n {\n typename BlockLoadT::TempStorage load;\n typename BlockStoreT::TempStorage store;\n typename BlockScanT::TempStorage scan;\n } temp_storage;\n\n // Per-thread tile data\n int data[ITEMS_PER_THREAD];\n\n // Load items into a blocked arrangement\n BlockLoadT(temp_storage.load).Load(d_in, data);\n\n // Barrier for smem reuse\n __syncthreads();\n\n // Start cycle timer\n clock_t start = clock();\n\n // Compute exclusive prefix sum\n int aggregate;\n BlockScanT(temp_storage.scan).ExclusiveSum(data, data, aggregate);\n\n // Stop cycle timer\n clock_t stop = clock();\n\n // Barrier for smem reuse\n __syncthreads();\n\n // Store items from a blocked arrangement\n BlockStoreT(temp_storage.store).Store(d_out, data);\n\n // Store aggregate and elapsed clocks\n if (threadIdx.x == 0)\n {\n *d_elapsed = (start > stop) ? start - stop : stop - start;\n d_out[BLOCK_THREADS * ITEMS_PER_THREAD] = aggregate;\n }\n}\n\n\n\n//---------------------------------------------------------------------\n// Host utilities\n//---------------------------------------------------------------------\n\n/**\n * Initialize exclusive prefix sum problem (and solution).\n * Returns the aggregate\n */\nint Initialize(\n int *h_in,\n int *h_reference,\n int num_items)\n{\n int inclusive = 0;\n\n for (int i = 0; i < num_items; ++i)\n {\n h_in[i] = i % 17;\n\n h_reference[i] = inclusive;\n inclusive += h_in[i];\n }\n\n return inclusive;\n}\n\n\n/**\n * Test thread block scan\n */\ntemplate <\n int BLOCK_THREADS,\n int ITEMS_PER_THREAD,\n BlockScanAlgorithm ALGORITHM>\nvoid Test()\n{\n const int TILE_SIZE = BLOCK_THREADS * ITEMS_PER_THREAD;\n\n // Allocate host arrays\n int *h_in = new int[TILE_SIZE];\n int *h_reference = new int[TILE_SIZE];\n int *h_gpu = new int[TILE_SIZE + 1];\n\n // Initialize problem and reference output on host\n int h_aggregate = Initialize(h_in, h_reference, TILE_SIZE);\n\n // Initialize device arrays\n int *d_in = NULL;\n int *d_out = NULL;\n clock_t *d_elapsed = NULL;\n hipMalloc((void**)&d_in, sizeof(int) * TILE_SIZE);\n hipMalloc((void**)&d_out, sizeof(int) * (TILE_SIZE + 1));\n hipMalloc((void**)&d_elapsed, sizeof(clock_t));\n\n // Display input problem data\n if (g_verbose)\n {\n printf(\"Input data: \");\n for (int i = 0; i < TILE_SIZE; i++)\n printf(\"%d, \", h_in[i]);\n printf(\"\\n\\n\");\n }\n\n // Kernel props\n int max_sm_occupancy;\n CubDebugExit(MaxSmOccupancy(max_sm_occupancy, BlockPrefixSumKernel, BLOCK_THREADS));\n\n // Copy problem to device\n hipMemcpy(d_in, h_in, sizeof(int) * TILE_SIZE, hipMemcpyHostToDevice);\n\n printf(\"BlockScan algorithm %s on %d items (%d timing iterations, %d blocks, %d threads, %d items per thread, %d SM occupancy):\\n\",\n (ALGORITHM == BLOCK_SCAN_RAKING) ? \"BLOCK_SCAN_RAKING\" : (ALGORITHM == BLOCK_SCAN_RAKING_MEMOIZE) ? \"BLOCK_SCAN_RAKING_MEMOIZE\" : \"BLOCK_SCAN_WARP_SCANS\",\n TILE_SIZE, g_timing_iterations, g_grid_size, BLOCK_THREADS, ITEMS_PER_THREAD, max_sm_occupancy);\n\n // Run aggregate/prefix kernel\n BlockPrefixSumKernel<<>>(\n d_in,\n d_out,\n d_elapsed);\n\n // Check results\n printf(\"\\tOutput items: \");\n int compare = CompareDeviceResults(h_reference, d_out, TILE_SIZE, g_verbose, g_verbose);\n printf(\"%s\\n\", compare ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Check total aggregate\n printf(\"\\tAggregate: \");\n compare = CompareDeviceResults(&h_aggregate, d_out + TILE_SIZE, 1, g_verbose, g_verbose);\n printf(\"%s\\n\", compare ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Run this several times and average the performance results\n GpuTimer timer;\n float elapsed_millis = 0.0;\n clock_t elapsed_clocks = 0;\n\n for (int i = 0; i < g_timing_iterations; ++i)\n {\n // Copy problem to device\n hipMemcpy(d_in, h_in, sizeof(int) * TILE_SIZE, hipMemcpyHostToDevice);\n\n timer.Start();\n\n // Run aggregate/prefix kernel\n BlockPrefixSumKernel<<>>(\n d_in,\n d_out,\n d_elapsed);\n\n timer.Stop();\n elapsed_millis += timer.ElapsedMillis();\n\n // Copy clocks from device\n clock_t clocks;\n CubDebugExit(hipMemcpy(&clocks, d_elapsed, sizeof(clock_t), hipMemcpyDeviceToHost));\n elapsed_clocks += clocks;\n\n }\n\n // Check for kernel errors and STDIO from the kernel, if any\n CubDebugExit(hipPeekAtLastError());\n CubDebugExit(hipDeviceSynchronize());\n\n // Display timing results\n float avg_millis = elapsed_millis / g_timing_iterations;\n float avg_items_per_sec = float(TILE_SIZE * g_grid_size) / avg_millis / 1000.0f;\n float avg_clocks = float(elapsed_clocks) / g_timing_iterations;\n float avg_clocks_per_item = avg_clocks / TILE_SIZE;\n\n printf(\"\\tAverage BlockScan::Sum clocks: %.3f\\n\", avg_clocks);\n printf(\"\\tAverage BlockScan::Sum clocks per item: %.3f\\n\", avg_clocks_per_item);\n printf(\"\\tAverage kernel millis: %.4f\\n\", avg_millis);\n printf(\"\\tAverage million items / sec: %.4f\\n\", avg_items_per_sec);\n\n // Cleanup\n if (h_in) delete[] h_in;\n if (h_reference) delete[] h_reference;\n if (h_gpu) delete[] h_gpu;\n if (d_in) hipFree(d_in);\n if (d_out) hipFree(d_out);\n if (d_elapsed) hipFree(d_elapsed);\n}\n\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n args.GetCmdLineArgument(\"grid-size\", g_grid_size);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=] \"\n \"[--i=]\"\n \"[--grid-size=]\"\n \"[--v] \"\n \"\\n\", argv[0], g_timing_iterations, g_grid_size);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // Run tests\n Test<1024, 1, BLOCK_SCAN_RAKING>();\n Test<512, 2, BLOCK_SCAN_RAKING>();\n Test<256, 4, BLOCK_SCAN_RAKING>();\n Test<128, 8, BLOCK_SCAN_RAKING>();\n Test<64, 16, BLOCK_SCAN_RAKING>();\n Test<32, 32, BLOCK_SCAN_RAKING>();\n\n printf(\"-------------\\n\");\n\n Test<1024, 1, BLOCK_SCAN_RAKING_MEMOIZE>();\n Test<512, 2, BLOCK_SCAN_RAKING_MEMOIZE>();\n Test<256, 4, BLOCK_SCAN_RAKING_MEMOIZE>();\n Test<128, 8, BLOCK_SCAN_RAKING_MEMOIZE>();\n Test<64, 16, BLOCK_SCAN_RAKING_MEMOIZE>();\n Test<32, 32, BLOCK_SCAN_RAKING_MEMOIZE>();\n\n printf(\"-------------\\n\");\n\n Test<1024, 1, BLOCK_SCAN_WARP_SCANS>();\n Test<512, 2, BLOCK_SCAN_WARP_SCANS>();\n Test<256, 4, BLOCK_SCAN_WARP_SCANS>();\n Test<128, 8, BLOCK_SCAN_WARP_SCANS>();\n Test<64, 16, BLOCK_SCAN_WARP_SCANS>();\n Test<32, 32, BLOCK_SCAN_WARP_SCANS>();\n\n\n return 0;\n}\n\n", "source": "/kaggle/working/cub/examples/block/example_block_scan.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Simple demonstration of cub::BlockReduce with dynamic shared memory\n *\n * To compile using the command line:\n * nvcc -arch=sm_XX example_block_reduce_dyn_smem.cu -I../.. -lcudart -O3 -std=c++14\n *\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console (define before including cub.h)\n#define CUB_STDERR\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../../test/test_util.h\"\n\n// Some implementation details rely on c++14\n#if CUB_CPP_DIALECT >= 2014\n\nusing namespace cub;\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\n/// Verbose output\nbool g_verbose = false;\n\n/// Default grid size\nint g_grid_size = 1;\n\n//---------------------------------------------------------------------\n// Kernels\n//---------------------------------------------------------------------\n\n/**\n * Simple kernel for performing a block-wide reduction.\n */\ntemplate \n__global__ void BlockReduceKernel(\n int *d_in, // Tile of input\n int *d_out // Tile aggregate\n )\n{\n // Specialize BlockReduce type for our thread block\n using BlockReduceT = cub::BlockReduce;\n using TempStorageT = typename BlockReduceT::TempStorage;\n\n union ShmemLayout\n {\n TempStorageT reduce;\n int aggregate;\n };\n\n // shared memory byte-array\n extern __shared__ __align__(alignof(ShmemLayout)) char smem[];\n \n // cast to lvalue reference of expected type\n auto& temp_storage = reinterpret_cast(smem);\n\n int data = d_in[threadIdx.x];\n\n // Compute sum\n int aggregate = BlockReduceT(temp_storage).Sum(data);\n\n // block-wide sync barrier necessary to re-use shared mem safely\n __syncthreads();\n int* smem_integers = reinterpret_cast(smem);\n if (threadIdx.x == 0) smem_integers[0] = aggregate;\n\n // sync to make new shared value available to all threads\n __syncthreads();\n aggregate = smem_integers[0];\n\n // all threads write the aggregate to output\n d_out[threadIdx.x] = aggregate;\n}\n\n\n\n//---------------------------------------------------------------------\n// Host utilities\n//---------------------------------------------------------------------\n\n/**\n * Initialize reduction problem (and solution).\n * Returns the aggregate\n */\nint Initialize(int *h_in, int num_items)\n{\n int inclusive = 0;\n\n for (int i = 0; i < num_items; ++i)\n {\n h_in[i] = i % 17;\n inclusive += h_in[i];\n }\n\n return inclusive;\n}\n\n/**\n * Test thread block reduction\n */\ntemplate \nvoid Test()\n{\n // Allocate host arrays\n int *h_in = new int[BLOCK_THREADS];\n\n // Initialize problem and reference output on host\n int h_aggregate = Initialize(h_in, BLOCK_THREADS);\n\n // Initialize device arrays\n int *d_in = NULL;\n int *d_out = NULL;\n cudaMalloc((void**)&d_in, sizeof(int) * BLOCK_THREADS);\n cudaMalloc((void**)&d_out, sizeof(int) * BLOCK_THREADS);\n\n // Display input problem data\n if (g_verbose)\n {\n printf(\"Input data: \");\n for (int i = 0; i < BLOCK_THREADS; i++)\n printf(\"%d, \", h_in[i]);\n printf(\"\\n\\n\");\n }\n\n // Copy problem to device\n cudaMemcpy(d_in, h_in, sizeof(int) * BLOCK_THREADS, cudaMemcpyHostToDevice);\n\n // determine necessary storage size:\n auto block_reduce_temp_bytes =\n sizeof(typename cub::BlockReduce::TempStorage);\n // finally, we need to make sure that we can hold at least one integer\n // needed in the kernel to exchange data after reduction\n auto smem_size = (std::max)(1 * sizeof(int), block_reduce_temp_bytes);\n\n // use default stream\n cudaStream_t stream = NULL;\n\n // Run reduction kernel\n BlockReduceKernel\n <<>>(\n d_in,\n d_out);\n\n // Check total aggregate\n printf(\"\\tAggregate: \");\n int compare = 0;\n for (int i = 0; i < BLOCK_THREADS; i++) {\n compare = compare || CompareDeviceResults(\n &h_aggregate, d_out + i, 1, g_verbose, g_verbose);\n }\n printf(\"%s\\n\", compare ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Check for kernel errors and STDIO from the kernel, if any\n CubDebugExit(cudaPeekAtLastError());\n CubDebugExit(cudaDeviceSynchronize());\n\n // Cleanup\n if (h_in) delete[] h_in;\n if (d_in) cudaFree(d_in);\n if (d_out) cudaFree(d_out);\n}\n\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"grid-size\", g_grid_size);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=] \"\n \"[--grid-size=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // Run tests\n Test<1024>();\n Test<512>();\n Test<256>();\n Test<128>();\n Test<64>();\n Test<32>();\n Test<16>();\n\n return 0;\n}\n\n#else // < C++14\n\nint main() {}\n\n#endif\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Simple demonstration of cub::BlockReduce with dynamic shared memory\n *\n * To compile using the command line:\n * nvcc -arch=sm_XX example_block_reduce_dyn_smem.cu -I../.. -lcudart -O3 -std=c++14\n *\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console (define before including cub.h)\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../../test/test_util.h\"\n\n// Some implementation details rely on c++14\n#if CUB_CPP_DIALECT >= 2014\n\nusing namespace hipcub;\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\n/// Verbose output\nbool g_verbose = false;\n\n/// Default grid size\nint g_grid_size = 1;\n\n//---------------------------------------------------------------------\n// Kernels\n//---------------------------------------------------------------------\n\n/**\n * Simple kernel for performing a block-wide reduction.\n */\ntemplate \n__global__ void BlockReduceKernel(\n int *d_in, // Tile of input\n int *d_out // Tile aggregate\n )\n{\n // Specialize BlockReduce type for our thread block\n using BlockReduceT = hipcub::BlockReduce;\n using TempStorageT = typename BlockReduceT::TempStorage;\n\n union ShmemLayout\n {\n TempStorageT reduce;\n int aggregate;\n };\n\n // shared memory byte-array\n extern __shared__ __align__(alignof(ShmemLayout)) char smem[];\n \n // cast to lvalue reference of expected type\n auto& temp_storage = reinterpret_cast(smem);\n\n int data = d_in[threadIdx.x];\n\n // Compute sum\n int aggregate = BlockReduceT(temp_storage).Sum(data);\n\n // block-wide sync barrier necessary to re-use shared mem safely\n __syncthreads();\n int* smem_integers = reinterpret_cast(smem);\n if (threadIdx.x == 0) smem_integers[0] = aggregate;\n\n // sync to make new shared value available to all threads\n __syncthreads();\n aggregate = smem_integers[0];\n\n // all threads write the aggregate to output\n d_out[threadIdx.x] = aggregate;\n}\n\n\n\n//---------------------------------------------------------------------\n// Host utilities\n//---------------------------------------------------------------------\n\n/**\n * Initialize reduction problem (and solution).\n * Returns the aggregate\n */\nint Initialize(int *h_in, int num_items)\n{\n int inclusive = 0;\n\n for (int i = 0; i < num_items; ++i)\n {\n h_in[i] = i % 17;\n inclusive += h_in[i];\n }\n\n return inclusive;\n}\n\n/**\n * Test thread block reduction\n */\ntemplate \nvoid Test()\n{\n // Allocate host arrays\n int *h_in = new int[BLOCK_THREADS];\n\n // Initialize problem and reference output on host\n int h_aggregate = Initialize(h_in, BLOCK_THREADS);\n\n // Initialize device arrays\n int *d_in = NULL;\n int *d_out = NULL;\n hipMalloc((void**)&d_in, sizeof(int) * BLOCK_THREADS);\n hipMalloc((void**)&d_out, sizeof(int) * BLOCK_THREADS);\n\n // Display input problem data\n if (g_verbose)\n {\n printf(\"Input data: \");\n for (int i = 0; i < BLOCK_THREADS; i++)\n printf(\"%d, \", h_in[i]);\n printf(\"\\n\\n\");\n }\n\n // Copy problem to device\n hipMemcpy(d_in, h_in, sizeof(int) * BLOCK_THREADS, hipMemcpyHostToDevice);\n\n // determine necessary storage size:\n auto block_reduce_temp_bytes =\n sizeof(typename hipcub::BlockReduce::TempStorage);\n // finally, we need to make sure that we can hold at least one integer\n // needed in the kernel to exchange data after reduction\n auto smem_size = (std::max)(1 * sizeof(int), block_reduce_temp_bytes);\n\n // use default stream\n hipStream_t stream = NULL;\n\n // Run reduction kernel\n BlockReduceKernel\n <<>>(\n d_in,\n d_out);\n\n // Check total aggregate\n printf(\"\\tAggregate: \");\n int compare = 0;\n for (int i = 0; i < BLOCK_THREADS; i++) {\n compare = compare || CompareDeviceResults(\n &h_aggregate, d_out + i, 1, g_verbose, g_verbose);\n }\n printf(\"%s\\n\", compare ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Check for kernel errors and STDIO from the kernel, if any\n CubDebugExit(hipPeekAtLastError());\n CubDebugExit(hipDeviceSynchronize());\n\n // Cleanup\n if (h_in) delete[] h_in;\n if (d_in) hipFree(d_in);\n if (d_out) hipFree(d_out);\n}\n\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"grid-size\", g_grid_size);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=] \"\n \"[--grid-size=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // Run tests\n Test<1024>();\n Test<512>();\n Test<256>();\n Test<128>();\n Test<64>();\n Test<32>();\n Test<16>();\n\n return 0;\n}\n\n#else // < C++14\n\nint main() {}\n\n#endif\n", "source": "/kaggle/working/cub/examples/block/example_block_reduce_dyn_smem.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n\n#include \n\n#include \n\ntemplate \n__global__ void init_kernel(ScanTileStateT tile_state, int blocks_in_grid)\n{\n tile_state.InitializeStatus(blocks_in_grid);\n}\n\ntemplate \n__global__ void decoupled_look_back_kernel(cub::ScanTileState tile_state)\n{\n using scan_op_t = cub::Sum;\n using scan_tile_state_t = cub::ScanTileState;\n using tile_prefix_op = cub::TilePrefixCallbackOp;\n using temp_storage_t = typename tile_prefix_op::TempStorage;\n\n // Allocate temp storage in shared memory\n __shared__ temp_storage_t temp_storage;\n\n scan_op_t scan_op{};\n const unsigned int threads_in_warp = 32;\n const unsigned int tid = threadIdx.x;\n\n // Construct prefix op\n tile_prefix_op prefix(tile_state, temp_storage, scan_op);\n const unsigned int tile_idx = prefix.GetTileIdx();\n\n // Compute block aggregate\n MessageT block_aggregate = blockIdx.x;\n\n if (tile_idx == 0)\n {\n // There are no blocks to look back to, immediately set the inclusive state\n if (tid == 0)\n {\n tile_state.SetInclusive(tile_idx, block_aggregate);\n printf(\"tile %d: inclusive = %d\\n\", tile_idx, block_aggregate);\n }\n }\n else\n {\n // Only the first warp in the block can perform the look back\n const unsigned int warp_id = tid / threads_in_warp;\n\n if (warp_id == 0)\n {\n // Perform the decoupled look-back\n // Invocation of the prefix will block until the look-back is complete.\n MessageT exclusive_prefix = prefix(block_aggregate);\n\n if (tid == 0)\n {\n MessageT inclusive_prefix = scan_op(exclusive_prefix, block_aggregate);\n printf(\"tile %d: exclusive = %d inclusive = %d\\n\",\n tile_idx,\n exclusive_prefix,\n inclusive_prefix);\n }\n }\n }\n}\n\ntemplate \nvoid decoupled_look_back_example(int blocks_in_grid)\n{\n using scan_tile_state_t = cub::ScanTileState;\n\n // Query temporary storage requirements\n std::size_t temp_storage_bytes{};\n scan_tile_state_t::AllocationSize(blocks_in_grid, temp_storage_bytes);\n\n // Allocate temporary storage\n thrust::device_vector temp_storage(temp_storage_bytes);\n std::uint8_t *d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n // Initialize temporary storage\n scan_tile_state_t tile_status;\n tile_status.Init(blocks_in_grid, d_temp_storage, temp_storage_bytes);\n const unsigned int threads_in_init_block = 256;\n const unsigned int blocks_in_init_grid = cub::DivideAndRoundUp(blocks_in_grid,\n threads_in_init_block);\n init_kernel<<>>(tile_status, blocks_in_grid);\n\n // Launch decoupled look-back\n const unsigned int threads_in_block = 256;\n decoupled_look_back_kernel<<>>(tile_status);\n\n // Wait for kernel to finish\n cudaDeviceSynchronize();\n}\n\nint main() { decoupled_look_back_example(14); }\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \n\n#include \n\ntemplate \n__global__ void init_kernel(ScanTileStateT tile_state, int blocks_in_grid)\n{\n tile_state.InitializeStatus(blocks_in_grid);\n}\n\ntemplate \n__global__ void decoupled_look_back_kernel(hipcub::ScanTileState tile_state)\n{\n using scan_op_t = hipcub::Sum;\n using scan_tile_state_t = hipcub::ScanTileState;\n using tile_prefix_op = hipcub::TilePrefixCallbackOp;\n using temp_storage_t = typename tile_prefix_op::TempStorage;\n\n // Allocate temp storage in shared memory\n __shared__ temp_storage_t temp_storage;\n\n scan_op_t scan_op{};\n const unsigned int threads_in_warp = 32;\n const unsigned int tid = threadIdx.x;\n\n // Construct prefix op\n tile_prefix_op prefix(tile_state, temp_storage, scan_op);\n const unsigned int tile_idx = prefix.GetTileIdx();\n\n // Compute block aggregate\n MessageT block_aggregate = blockIdx.x;\n\n if (tile_idx == 0)\n {\n // There are no blocks to look back to, immediately set the inclusive state\n if (tid == 0)\n {\n tile_state.SetInclusive(tile_idx, block_aggregate);\n printf(\"tile %d: inclusive = %d\\n\", tile_idx, block_aggregate);\n }\n }\n else\n {\n // Only the first warp in the block can perform the look back\n const unsigned int warp_id = tid / threads_in_warp;\n\n if (warp_id == 0)\n {\n // Perform the decoupled look-back\n // Invocation of the prefix will block until the look-back is complete.\n MessageT exclusive_prefix = prefix(block_aggregate);\n\n if (tid == 0)\n {\n MessageT inclusive_prefix = scan_op(exclusive_prefix, block_aggregate);\n printf(\"tile %d: exclusive = %d inclusive = %d\\n\",\n tile_idx,\n exclusive_prefix,\n inclusive_prefix);\n }\n }\n }\n}\n\ntemplate \nvoid decoupled_look_back_example(int blocks_in_grid)\n{\n using scan_tile_state_t = hipcub::ScanTileState;\n\n // Query temporary storage requirements\n std::size_t temp_storage_bytes{};\n scan_tile_state_t::AllocationSize(blocks_in_grid, temp_storage_bytes);\n\n // Allocate temporary storage\n thrust::device_vector temp_storage(temp_storage_bytes);\n std::uint8_t *d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n // Initialize temporary storage\n scan_tile_state_t tile_status;\n tile_status.Init(blocks_in_grid, d_temp_storage, temp_storage_bytes);\n const unsigned int threads_in_init_block = 256;\n const unsigned int blocks_in_init_grid = hipcub::DivideAndRoundUp(blocks_in_grid,\n threads_in_init_block);\n init_kernel<<>>(tile_status, blocks_in_grid);\n\n // Launch decoupled look-back\n const unsigned int threads_in_block = 256;\n decoupled_look_back_kernel<<>>(tile_status);\n\n // Wait for kernel to finish\n hipDeviceSynchronize();\n}\n\nint main() { decoupled_look_back_example(14); }\n", "source": "/kaggle/working/cub/examples/device/example_device_decoupled_look_back.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \n#include \n\n#include \"thrust/scan.h\"\n#include \n#include \n\nclass generator_t\n{\npublic:\n generator_t();\n ~generator_t();\n\n template \n void operator()(seed_t seed,\n thrust::device_vector &data,\n bit_entropy entropy,\n T min = std::numeric_limits::lowest(),\n T max = std::numeric_limits::max());\n\n template \n thrust::device_vector power_law_segment_offsets(seed_t seed,\n std::size_t total_elements,\n std::size_t total_segments);\n\n double *distribution();\n curandGenerator_t &gen() { return m_gen; }\n\n double *prepare_random_generator(seed_t seed, std::size_t num_items);\n double *prepare_lognormal_random_generator(seed_t seed, std::size_t num_items);\n\nprivate:\n curandGenerator_t m_gen;\n thrust::device_vector m_distribution;\n};\n\ntemplate \nstruct random_to_item_t\n{\n double m_min;\n double m_max;\n\n __host__ __device__ random_to_item_t(T min, T max)\n : m_min(static_cast(min))\n , m_max(static_cast(max))\n {}\n\n __host__ __device__ T operator()(double random_value)\n {\n return static_cast((m_max - m_min) * random_value + m_min);\n }\n};\n\ngenerator_t::generator_t() { curandCreateGenerator(&m_gen, CURAND_RNG_PSEUDO_DEFAULT); }\n\ngenerator_t::~generator_t() { curandDestroyGenerator(m_gen); }\n\ndouble *generator_t::distribution() { return thrust::raw_pointer_cast(m_distribution.data()); }\n\ndouble *generator_t::prepare_random_generator(seed_t seed, std::size_t num_items)\n{\n curandSetPseudoRandomGeneratorSeed(m_gen, seed.get());\n\n m_distribution.resize(num_items);\n curandGenerateUniformDouble(m_gen, this->distribution(), num_items);\n\n return this->distribution();\n}\n\ndouble *generator_t::prepare_lognormal_random_generator(seed_t seed, std::size_t num_segments)\n{\n curandSetPseudoRandomGeneratorSeed(m_gen, seed.get());\n\n m_distribution.resize(num_segments);\n\n const double mean = 3.0;\n const double sigma = 1.2;\n curandGenerateLogNormalDouble(m_gen, this->distribution(), num_segments, mean, sigma);\n\n return this->distribution();\n}\n\ntemplate \n__global__ void and_kernel(T *d_in, T *d_tmp, std::size_t n)\n{\n const std::size_t i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < n)\n {\n d_in[i] = d_in[i] & d_tmp[i];\n }\n}\n\n__global__ void and_kernel(float *d_in, float *d_tmp, std::size_t n)\n{\n const std::size_t i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < n)\n {\n std::uint32_t result = reinterpret_cast(d_in[i]) &\n reinterpret_cast(d_tmp[i]);\n d_in[i] = reinterpret_cast(result);\n }\n}\n\n__global__ void and_kernel(double *d_in, double *d_tmp, std::size_t n)\n{\n const std::size_t i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < n)\n {\n std::uint64_t result = reinterpret_cast(d_in[i]) &\n reinterpret_cast(d_tmp[i]);\n d_in[i] = reinterpret_cast(result);\n }\n}\n\n__global__ void and_kernel(complex *d_in, complex *d_tmp, std::size_t n)\n{\n const std::size_t i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < n)\n {\n double in_real = d_in[i].real();\n double in_imag = d_in[i].imag();\n\n double tmp_real = d_tmp[i].real();\n double tmp_imag = d_tmp[i].imag();\n\n std::uint64_t result_real = reinterpret_cast(in_real) &\n reinterpret_cast(tmp_real);\n\n std::uint64_t result_imag = reinterpret_cast(in_imag) &\n reinterpret_cast(tmp_imag);\n\n d_in[i].real(reinterpret_cast(result_real));\n d_in[i].imag(reinterpret_cast(result_imag));\n }\n}\n\n__global__ void set_real_kernel(complex *d_in, complex min, complex max, double *d_tmp, std::size_t n)\n{\n const std::size_t i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < n)\n {\n d_in[i].real(random_to_item_t{min.real(), max.real()}(d_tmp[i]));\n }\n}\n\n__global__ void set_imag_kernel(complex *d_in, complex min, complex max, double *d_tmp, std::size_t n)\n{\n const std::size_t i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < n)\n {\n d_in[i].imag(random_to_item_t{min.imag(), max.imag()}(d_tmp[i]));\n }\n}\n\ntemplate \nstruct lognormal_transformer_t\n{\n std::size_t total_elements;\n double sum;\n\n __device__ T operator()(double val)\n {\n return floor(val * total_elements / sum);\n }\n};\n\ntemplate \n__global__ void lognormal_adjust_kernel(T *segment_sizes, std::size_t diff)\n{\n const unsigned tid = threadIdx.x + blockIdx.x * blockDim.x;\n\n if (tid < diff)\n {\n segment_sizes[tid]++;\n }\n}\n\ntemplate \nthrust::device_vector generator_t::power_law_segment_offsets(seed_t seed,\n std::size_t total_elements,\n std::size_t total_segments)\n{\n thrust::device_vector segment_sizes(total_segments + 1);\n prepare_lognormal_random_generator(seed, total_segments);\n\n if (thrust::count(m_distribution.begin(), m_distribution.end(), 0.0) == total_segments)\n {\n thrust::fill_n(m_distribution.begin(), total_segments, 1.0);\n }\n\n double sum = thrust::reduce(m_distribution.begin(), m_distribution.end());\n thrust::transform(m_distribution.begin(),\n m_distribution.end(),\n segment_sizes.begin(),\n lognormal_transformer_t{total_elements, sum});\n\n const int diff = total_elements - thrust::reduce(segment_sizes.begin(), segment_sizes.end());\n const int block_size = 256;\n const int grid_size = (std::abs(diff) + block_size - 1) / block_size;\n\n T *d_segment_sizes = thrust::raw_pointer_cast(segment_sizes.data());\n lognormal_adjust_kernel<<>>(d_segment_sizes, diff);\n\n thrust::exclusive_scan(segment_sizes.begin(), segment_sizes.end(), segment_sizes.begin());\n\n return segment_sizes;\n}\n\ntemplate \nvoid generator_t::operator()(seed_t seed,\n thrust::device_vector &data,\n bit_entropy entropy,\n T min,\n T max)\n{\n switch (entropy)\n {\n case bit_entropy::_1_000: {\n prepare_random_generator(seed, data.size());\n\n thrust::transform(m_distribution.begin(),\n m_distribution.end(),\n data.begin(),\n random_to_item_t(min, max));\n return;\n }\n case bit_entropy::_0_000: {\n std::mt19937 rng;\n rng.seed(static_cast(seed.get()));\n std::uniform_real_distribution dist(0.0f, 1.0f);\n T random_value = random_to_item_t(min, max)(dist(rng));\n thrust::fill(data.begin(), data.end(), random_value);\n return;\n }\n default: {\n prepare_random_generator(seed, data.size());\n\n thrust::transform(m_distribution.begin(),\n m_distribution.end(),\n data.begin(),\n random_to_item_t(min, max));\n\n const int number_of_steps = static_cast(entropy);\n thrust::device_vector tmp(data.size());\n const int threads_in_block = 256;\n const int blocks_in_grid = (data.size() + threads_in_block - 1) / threads_in_block;\n\n for (int i = 0; i < number_of_steps; i++, ++seed)\n {\n (*this)(seed, tmp, bit_entropy::_1_000, min, max);\n and_kernel<<>>(thrust::raw_pointer_cast(data.data()),\n thrust::raw_pointer_cast(tmp.data()),\n data.size());\n cudaStreamSynchronize(0);\n }\n return;\n }\n };\n}\n\ntemplate <>\nvoid generator_t::operator()(seed_t seed,\n thrust::device_vector &data,\n bit_entropy entropy,\n complex min,\n complex max)\n{\n const int threads_in_block = 256;\n const int blocks_in_grid = (data.size() + threads_in_block - 1) / threads_in_block;\n\n switch (entropy)\n {\n case bit_entropy::_1_000: {\n prepare_random_generator(seed, data.size()); ++seed;\n set_real_kernel<<>>(thrust::raw_pointer_cast(data.data()),\n min,\n max,\n thrust::raw_pointer_cast(m_distribution.data()),\n data.size());\n prepare_random_generator(seed, data.size()); ++seed;\n set_imag_kernel<<>>(thrust::raw_pointer_cast(data.data()),\n min,\n max,\n thrust::raw_pointer_cast(m_distribution.data()),\n data.size());\n return;\n }\n case bit_entropy::_0_000: {\n std::mt19937 rng;\n rng.seed(static_cast(seed.get()));\n std::uniform_real_distribution dist(0.0f, 1.0f);\n double random_imag = random_to_item_t(min.imag(), max.imag())(dist(rng));\n double random_real = random_to_item_t(min.imag(), max.imag())(dist(rng));\n complex random_value(random_real, random_imag);\n thrust::fill(data.begin(), data.end(), random_value);\n return;\n }\n default: {\n prepare_random_generator(seed, data.size());\n\n prepare_random_generator(seed, data.size()); ++seed;\n set_real_kernel<<>>(thrust::raw_pointer_cast(data.data()),\n min,\n max,\n thrust::raw_pointer_cast(m_distribution.data()),\n data.size());\n prepare_random_generator(seed, data.size()); ++seed;\n set_imag_kernel<<>>(thrust::raw_pointer_cast(data.data()),\n min,\n max,\n thrust::raw_pointer_cast(m_distribution.data()),\n data.size());\n\n const int number_of_steps = static_cast(entropy);\n thrust::device_vector tmp(data.size());\n\n for (int i = 0; i < number_of_steps; i++, ++seed)\n {\n (*this)(seed, tmp, bit_entropy::_1_000, min, max);\n and_kernel<<>>(thrust::raw_pointer_cast(data.data()),\n thrust::raw_pointer_cast(tmp.data()),\n data.size());\n cudaStreamSynchronize(0);\n }\n return;\n }\n };\n}\n\nstruct random_to_probability_t\n{\n double m_probability;\n\n __host__ __device__ bool operator()(double random_value)\n {\n return random_value < m_probability;\n }\n};\n\ntemplate <>\nvoid generator_t::operator()(seed_t seed,\n thrust::device_vector &data,\n bit_entropy entropy,\n bool /* min */,\n bool /* max */)\n{\n if (entropy == bit_entropy::_0_000)\n {\n thrust::fill(data.begin(), data.end(), false);\n }\n else if (entropy == bit_entropy::_1_000)\n {\n thrust::fill(data.begin(), data.end(), true);\n }\n else\n {\n prepare_random_generator(seed, data.size());\n thrust::transform(m_distribution.begin(),\n m_distribution.end(),\n data.begin(),\n random_to_probability_t{entropy_to_probability(entropy)});\n }\n}\n\ntemplate \nvoid gen(seed_t seed, thrust::device_vector &data, bit_entropy entropy, T min, T max)\n{\n generator_t{}(seed, data, entropy, min, max);\n}\n\n#define INSTANTIATE_RND(TYPE) \\\n template void gen(seed_t, \\\n thrust::device_vector & data, \\\n bit_entropy, \\\n TYPE min, \\\n TYPE max)\n\n#define INSTANTIATE(TYPE) INSTANTIATE_RND(TYPE);\n\nINSTANTIATE(bool);\n\nINSTANTIATE(uint8_t);\nINSTANTIATE(uint16_t);\nINSTANTIATE(uint32_t);\nINSTANTIATE(uint64_t);\nINSTANTIATE(uint128_t);\n\nINSTANTIATE(int8_t);\nINSTANTIATE(int16_t);\nINSTANTIATE(int32_t);\nINSTANTIATE(int64_t);\nINSTANTIATE(int128_t);\n\nINSTANTIATE(float);\nINSTANTIATE(double);\nINSTANTIATE(complex);\n\n#undef INSTANTIATE\n#undef INSTANTIATE_RND\n\n\ntemplate \nthrust::device_vector gen_power_law_offsets(seed_t seed,\n std::size_t total_elements,\n std::size_t total_segments)\n{\n return generator_t{}.power_law_segment_offsets(seed, total_elements, total_segments);\n}\n\n#define INSTANTIATE(TYPE) \\\n template thrust::device_vector gen_power_law_offsets(seed_t, std::size_t, std::size_t)\n\nINSTANTIATE(uint32_t);\nINSTANTIATE(uint64_t);\n#undef INSTANTIATE\n\n\ntemplate \nstruct offset_to_iterator_t\n{\n T *base_it;\n\n __host__ __device__ __forceinline__ T*\n operator()(std::size_t offset) const\n {\n return base_it + offset;\n }\n};\n\ntemplate \nstruct repeat_index_t\n{\n __host__ __device__ __forceinline__ thrust::constant_iterator operator()(std::size_t i)\n {\n return thrust::constant_iterator(static_cast(i));\n }\n};\n\nstruct offset_to_size_t\n{\n std::size_t *offsets = nullptr;\n\n __host__ __device__ __forceinline__ std::size_t operator()(std::size_t i)\n {\n return offsets[i + 1] - offsets[i];\n }\n};\n\ntemplate \nthrust::device_vector gen_power_law_key_segments(seed_t seed,\n std::size_t total_elements,\n thrust::device_vector &segment_offsets)\n{\n std::size_t total_segments = segment_offsets.size() - 1;\n thrust::device_vector out(total_elements);\n std::size_t *d_offsets = thrust::raw_pointer_cast(segment_offsets.data());\n T *d_out = thrust::raw_pointer_cast(out.data());\n\n thrust::counting_iterator iota(0);\n offset_to_iterator_t dst_transform_op{d_out};\n\n auto d_range_srcs = thrust::make_transform_iterator(iota, repeat_index_t{});\n auto d_range_dsts = thrust::make_transform_iterator(d_offsets, dst_transform_op);\n auto d_range_sizes = thrust::make_transform_iterator(iota, offset_to_size_t{d_offsets});\n\n std::uint8_t *d_temp_storage = nullptr;\n std::size_t temp_storage_bytes = 0;\n cub::DeviceCopy::Batched(d_temp_storage,\n temp_storage_bytes,\n d_range_srcs,\n d_range_dsts,\n d_range_sizes,\n total_segments);\n \n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n cub::DeviceCopy::Batched(d_temp_storage,\n temp_storage_bytes,\n d_range_srcs,\n d_range_dsts,\n d_range_sizes,\n total_segments);\n cudaDeviceSynchronize();\n\n return out;\n}\n\ntemplate \nthrust::device_vector gen_power_law_key_segments(seed_t seed,\n std::size_t total_elements,\n std::size_t total_segments)\n{\n thrust::device_vector segment_offsets =\n gen_power_law_offsets(seed, total_elements, total_segments);\n return gen_power_law_key_segments(seed, total_elements, segment_offsets);\n}\n\n#define INSTANTIATE(TYPE) \\\n template thrust::device_vector gen_power_law_key_segments(seed_t, \\\n std::size_t, \\\n std::size_t)\n\nINSTANTIATE(bool);\n\nINSTANTIATE(uint8_t);\nINSTANTIATE(uint16_t);\nINSTANTIATE(uint32_t);\nINSTANTIATE(uint64_t);\nINSTANTIATE(uint128_t);\n\nINSTANTIATE(int8_t);\nINSTANTIATE(int16_t);\nINSTANTIATE(int32_t);\nINSTANTIATE(int64_t);\nINSTANTIATE(int128_t);\n\nINSTANTIATE(float);\nINSTANTIATE(double);\nINSTANTIATE(complex);\n#undef INSTANTIATE\n\ntemplate \nstruct gt_t\n{\n T val;\n\n __device__ bool operator()(T x)\n {\n return x > val;\n }\n};\n\ntemplate \nthrust::device_vector gen_uniform_offsets(seed_t seed,\n T total_elements,\n T min_segment_size,\n T max_segment_size)\n{\n thrust::device_vector segment_offsets(total_elements + 2);\n gen(seed, segment_offsets, bit_entropy::_1_000, min_segment_size, max_segment_size);\n segment_offsets[total_elements] = total_elements + 1;\n thrust::exclusive_scan(segment_offsets.begin(), segment_offsets.end(), segment_offsets.begin());\n typename thrust::device_vector::iterator iter =\n thrust::find_if(segment_offsets.begin(), segment_offsets.end(), gt_t{total_elements});\n *iter = total_elements;\n segment_offsets.erase(iter + 1, segment_offsets.end());\n return segment_offsets;\n}\n\n#define INSTANTIATE(TYPE) \\\n template thrust::device_vector gen_uniform_offsets(seed_t, TYPE, TYPE, TYPE)\n\nINSTANTIATE(uint32_t);\nINSTANTIATE(uint64_t);\n\n#undef INSTANTIATE\n\n/**\n * @brief Generates a vector of random key segments.\n *\n * Not all parameter combinations can be satisfied. For instance, if the total\n * elements is less than the minimal segment size, the function will return a\n * vector with a single element that is outside of the requested range. \n * At most one segment can be out of the requested range.\n */\ntemplate \nthrust::device_vector gen_uniform_key_segments(seed_t seed,\n std::size_t total_elements,\n std::size_t min_segment_size,\n std::size_t max_segment_size)\n{\n thrust::device_vector segment_offsets =\n gen_uniform_offsets(seed, total_elements, min_segment_size, max_segment_size);\n return gen_power_law_key_segments(seed, total_elements, segment_offsets);\n}\n\n#define INSTANTIATE(TYPE) \\\n template thrust::device_vector gen_uniform_key_segments(seed_t, \\\n std::size_t, \\\n std::size_t, \\\n std::size_t)\n\nINSTANTIATE(bool);\n\nINSTANTIATE(uint8_t);\nINSTANTIATE(uint16_t);\nINSTANTIATE(uint32_t);\nINSTANTIATE(uint64_t);\nINSTANTIATE(uint128_t);\n\nINSTANTIATE(int8_t);\nINSTANTIATE(int16_t);\nINSTANTIATE(int32_t);\nINSTANTIATE(int64_t);\nINSTANTIATE(int128_t);\n\nINSTANTIATE(float);\nINSTANTIATE(double);\nINSTANTIATE(complex);\n#undef INSTANTIATE\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \n#include \n\n#include \"thrust/scan.h\"\n#include \n#include \n\nclass generator_t\n{\npublic:\n generator_t();\n ~generator_t();\n\n template \n void operator()(seed_t seed,\n thrust::device_vector &data,\n bit_entropy entropy,\n T min = std::numeric_limits::lowest(),\n T max = std::numeric_limits::max());\n\n template \n thrust::device_vector power_law_segment_offsets(seed_t seed,\n std::size_t total_elements,\n std::size_t total_segments);\n\n double *distribution();\n hiprandGenerator_t &gen() { return m_gen; }\n\n double *prepare_random_generator(seed_t seed, std::size_t num_items);\n double *prepare_lognormal_random_generator(seed_t seed, std::size_t num_items);\n\nprivate:\n hiprandGenerator_t m_gen;\n thrust::device_vector m_distribution;\n};\n\ntemplate \nstruct random_to_item_t\n{\n double m_min;\n double m_max;\n\n __host__ __device__ random_to_item_t(T min, T max)\n : m_min(static_cast(min))\n , m_max(static_cast(max))\n {}\n\n __host__ __device__ T operator()(double random_value)\n {\n return static_cast((m_max - m_min) * random_value + m_min);\n }\n};\n\ngenerator_t::generator_t() { hiprandCreateGenerator(&m_gen, HIPRAND_RNG_PSEUDO_DEFAULT); }\n\ngenerator_t::~generator_t() { hiprandDestroyGenerator(m_gen); }\n\ndouble *generator_t::distribution() { return thrust::raw_pointer_cast(m_distribution.data()); }\n\ndouble *generator_t::prepare_random_generator(seed_t seed, std::size_t num_items)\n{\n hiprandSetPseudoRandomGeneratorSeed(m_gen, seed.get());\n\n m_distribution.resize(num_items);\n hiprandGenerateUniformDouble(m_gen, this->distribution(), num_items);\n\n return this->distribution();\n}\n\ndouble *generator_t::prepare_lognormal_random_generator(seed_t seed, std::size_t num_segments)\n{\n hiprandSetPseudoRandomGeneratorSeed(m_gen, seed.get());\n\n m_distribution.resize(num_segments);\n\n const double mean = 3.0;\n const double sigma = 1.2;\n hiprandGenerateLogNormalDouble(m_gen, this->distribution(), num_segments, mean, sigma);\n\n return this->distribution();\n}\n\ntemplate \n__global__ void and_kernel(T *d_in, T *d_tmp, std::size_t n)\n{\n const std::size_t i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < n)\n {\n d_in[i] = d_in[i] & d_tmp[i];\n }\n}\n\n__global__ void and_kernel(float *d_in, float *d_tmp, std::size_t n)\n{\n const std::size_t i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < n)\n {\n std::uint32_t result = reinterpret_cast(d_in[i]) &\n reinterpret_cast(d_tmp[i]);\n d_in[i] = reinterpret_cast(result);\n }\n}\n\n__global__ void and_kernel(double *d_in, double *d_tmp, std::size_t n)\n{\n const std::size_t i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < n)\n {\n std::uint64_t result = reinterpret_cast(d_in[i]) &\n reinterpret_cast(d_tmp[i]);\n d_in[i] = reinterpret_cast(result);\n }\n}\n\n__global__ void and_kernel(complex *d_in, complex *d_tmp, std::size_t n)\n{\n const std::size_t i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < n)\n {\n double in_real = d_in[i].real();\n double in_imag = d_in[i].imag();\n\n double tmp_real = d_tmp[i].real();\n double tmp_imag = d_tmp[i].imag();\n\n std::uint64_t result_real = reinterpret_cast(in_real) &\n reinterpret_cast(tmp_real);\n\n std::uint64_t result_imag = reinterpret_cast(in_imag) &\n reinterpret_cast(tmp_imag);\n\n d_in[i].real(reinterpret_cast(result_real));\n d_in[i].imag(reinterpret_cast(result_imag));\n }\n}\n\n__global__ void set_real_kernel(complex *d_in, complex min, complex max, double *d_tmp, std::size_t n)\n{\n const std::size_t i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < n)\n {\n d_in[i].real(random_to_item_t{min.real(), max.real()}(d_tmp[i]));\n }\n}\n\n__global__ void set_imag_kernel(complex *d_in, complex min, complex max, double *d_tmp, std::size_t n)\n{\n const std::size_t i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < n)\n {\n d_in[i].imag(random_to_item_t{min.imag(), max.imag()}(d_tmp[i]));\n }\n}\n\ntemplate \nstruct lognormal_transformer_t\n{\n std::size_t total_elements;\n double sum;\n\n __device__ T operator()(double val)\n {\n return floor(val * total_elements / sum);\n }\n};\n\ntemplate \n__global__ void lognormal_adjust_kernel(T *segment_sizes, std::size_t diff)\n{\n const unsigned tid = threadIdx.x + blockIdx.x * blockDim.x;\n\n if (tid < diff)\n {\n segment_sizes[tid]++;\n }\n}\n\ntemplate \nthrust::device_vector generator_t::power_law_segment_offsets(seed_t seed,\n std::size_t total_elements,\n std::size_t total_segments)\n{\n thrust::device_vector segment_sizes(total_segments + 1);\n prepare_lognormal_random_generator(seed, total_segments);\n\n if (thrust::count(m_distribution.begin(), m_distribution.end(), 0.0) == total_segments)\n {\n thrust::fill_n(m_distribution.begin(), total_segments, 1.0);\n }\n\n double sum = thrust::reduce(m_distribution.begin(), m_distribution.end());\n thrust::transform(m_distribution.begin(),\n m_distribution.end(),\n segment_sizes.begin(),\n lognormal_transformer_t{total_elements, sum});\n\n const int diff = total_elements - thrust::reduce(segment_sizes.begin(), segment_sizes.end());\n const int block_size = 256;\n const int grid_size = (std::abs(diff) + block_size - 1) / block_size;\n\n T *d_segment_sizes = thrust::raw_pointer_cast(segment_sizes.data());\n lognormal_adjust_kernel<<>>(d_segment_sizes, diff);\n\n thrust::exclusive_scan(segment_sizes.begin(), segment_sizes.end(), segment_sizes.begin());\n\n return segment_sizes;\n}\n\ntemplate \nvoid generator_t::operator()(seed_t seed,\n thrust::device_vector &data,\n bit_entropy entropy,\n T min,\n T max)\n{\n switch (entropy)\n {\n case bit_entropy::_1_000: {\n prepare_random_generator(seed, data.size());\n\n thrust::transform(m_distribution.begin(),\n m_distribution.end(),\n data.begin(),\n random_to_item_t(min, max));\n return;\n }\n case bit_entropy::_0_000: {\n std::mt19937 rng;\n rng.seed(static_cast(seed.get()));\n std::uniform_real_distribution dist(0.0f, 1.0f);\n T random_value = random_to_item_t(min, max)(dist(rng));\n thrust::fill(data.begin(), data.end(), random_value);\n return;\n }\n default: {\n prepare_random_generator(seed, data.size());\n\n thrust::transform(m_distribution.begin(),\n m_distribution.end(),\n data.begin(),\n random_to_item_t(min, max));\n\n const int number_of_steps = static_cast(entropy);\n thrust::device_vector tmp(data.size());\n const int threads_in_block = 256;\n const int blocks_in_grid = (data.size() + threads_in_block - 1) / threads_in_block;\n\n for (int i = 0; i < number_of_steps; i++, ++seed)\n {\n (*this)(seed, tmp, bit_entropy::_1_000, min, max);\n and_kernel<<>>(thrust::raw_pointer_cast(data.data()),\n thrust::raw_pointer_cast(tmp.data()),\n data.size());\n hipStreamSynchronize(0);\n }\n return;\n }\n };\n}\n\ntemplate <>\nvoid generator_t::operator()(seed_t seed,\n thrust::device_vector &data,\n bit_entropy entropy,\n complex min,\n complex max)\n{\n const int threads_in_block = 256;\n const int blocks_in_grid = (data.size() + threads_in_block - 1) / threads_in_block;\n\n switch (entropy)\n {\n case bit_entropy::_1_000: {\n prepare_random_generator(seed, data.size()); ++seed;\n set_real_kernel<<>>(thrust::raw_pointer_cast(data.data()),\n min,\n max,\n thrust::raw_pointer_cast(m_distribution.data()),\n data.size());\n prepare_random_generator(seed, data.size()); ++seed;\n set_imag_kernel<<>>(thrust::raw_pointer_cast(data.data()),\n min,\n max,\n thrust::raw_pointer_cast(m_distribution.data()),\n data.size());\n return;\n }\n case bit_entropy::_0_000: {\n std::mt19937 rng;\n rng.seed(static_cast(seed.get()));\n std::uniform_real_distribution dist(0.0f, 1.0f);\n double random_imag = random_to_item_t(min.imag(), max.imag())(dist(rng));\n double random_real = random_to_item_t(min.imag(), max.imag())(dist(rng));\n complex random_value(random_real, random_imag);\n thrust::fill(data.begin(), data.end(), random_value);\n return;\n }\n default: {\n prepare_random_generator(seed, data.size());\n\n prepare_random_generator(seed, data.size()); ++seed;\n set_real_kernel<<>>(thrust::raw_pointer_cast(data.data()),\n min,\n max,\n thrust::raw_pointer_cast(m_distribution.data()),\n data.size());\n prepare_random_generator(seed, data.size()); ++seed;\n set_imag_kernel<<>>(thrust::raw_pointer_cast(data.data()),\n min,\n max,\n thrust::raw_pointer_cast(m_distribution.data()),\n data.size());\n\n const int number_of_steps = static_cast(entropy);\n thrust::device_vector tmp(data.size());\n\n for (int i = 0; i < number_of_steps; i++, ++seed)\n {\n (*this)(seed, tmp, bit_entropy::_1_000, min, max);\n and_kernel<<>>(thrust::raw_pointer_cast(data.data()),\n thrust::raw_pointer_cast(tmp.data()),\n data.size());\n hipStreamSynchronize(0);\n }\n return;\n }\n };\n}\n\nstruct random_to_probability_t\n{\n double m_probability;\n\n __host__ __device__ bool operator()(double random_value)\n {\n return random_value < m_probability;\n }\n};\n\ntemplate <>\nvoid generator_t::operator()(seed_t seed,\n thrust::device_vector &data,\n bit_entropy entropy,\n bool /* min */,\n bool /* max */)\n{\n if (entropy == bit_entropy::_0_000)\n {\n thrust::fill(data.begin(), data.end(), false);\n }\n else if (entropy == bit_entropy::_1_000)\n {\n thrust::fill(data.begin(), data.end(), true);\n }\n else\n {\n prepare_random_generator(seed, data.size());\n thrust::transform(m_distribution.begin(),\n m_distribution.end(),\n data.begin(),\n random_to_probability_t{entropy_to_probability(entropy)});\n }\n}\n\ntemplate \nvoid gen(seed_t seed, thrust::device_vector &data, bit_entropy entropy, T min, T max)\n{\n generator_t{}(seed, data, entropy, min, max);\n}\n\n#define INSTANTIATE_RND(TYPE) \\\n template void gen(seed_t, \\\n thrust::device_vector & data, \\\n bit_entropy, \\\n TYPE min, \\\n TYPE max)\n\n#define INSTANTIATE(TYPE) INSTANTIATE_RND(TYPE);\n\nINSTANTIATE(bool);\n\nINSTANTIATE(uint8_t);\nINSTANTIATE(uint16_t);\nINSTANTIATE(uint32_t);\nINSTANTIATE(uint64_t);\nINSTANTIATE(uint128_t);\n\nINSTANTIATE(int8_t);\nINSTANTIATE(int16_t);\nINSTANTIATE(int32_t);\nINSTANTIATE(int64_t);\nINSTANTIATE(int128_t);\n\nINSTANTIATE(float);\nINSTANTIATE(double);\nINSTANTIATE(complex);\n\n#undef INSTANTIATE\n#undef INSTANTIATE_RND\n\n\ntemplate \nthrust::device_vector gen_power_law_offsets(seed_t seed,\n std::size_t total_elements,\n std::size_t total_segments)\n{\n return generator_t{}.power_law_segment_offsets(seed, total_elements, total_segments);\n}\n\n#define INSTANTIATE(TYPE) \\\n template thrust::device_vector gen_power_law_offsets(seed_t, std::size_t, std::size_t)\n\nINSTANTIATE(uint32_t);\nINSTANTIATE(uint64_t);\n#undef INSTANTIATE\n\n\ntemplate \nstruct offset_to_iterator_t\n{\n T *base_it;\n\n __host__ __device__ __forceinline__ T*\n operator()(std::size_t offset) const\n {\n return base_it + offset;\n }\n};\n\ntemplate \nstruct repeat_index_t\n{\n __host__ __device__ __forceinline__ thrust::constant_iterator operator()(std::size_t i)\n {\n return thrust::constant_iterator(static_cast(i));\n }\n};\n\nstruct offset_to_size_t\n{\n std::size_t *offsets = nullptr;\n\n __host__ __device__ __forceinline__ std::size_t operator()(std::size_t i)\n {\n return offsets[i + 1] - offsets[i];\n }\n};\n\ntemplate \nthrust::device_vector gen_power_law_key_segments(seed_t seed,\n std::size_t total_elements,\n thrust::device_vector &segment_offsets)\n{\n std::size_t total_segments = segment_offsets.size() - 1;\n thrust::device_vector out(total_elements);\n std::size_t *d_offsets = thrust::raw_pointer_cast(segment_offsets.data());\n T *d_out = thrust::raw_pointer_cast(out.data());\n\n thrust::counting_iterator iota(0);\n offset_to_iterator_t dst_transform_op{d_out};\n\n auto d_range_srcs = thrust::make_transform_iterator(iota, repeat_index_t{});\n auto d_range_dsts = thrust::make_transform_iterator(d_offsets, dst_transform_op);\n auto d_range_sizes = thrust::make_transform_iterator(iota, offset_to_size_t{d_offsets});\n\n std::uint8_t *d_temp_storage = nullptr;\n std::size_t temp_storage_bytes = 0;\n hipcub::DeviceCopy::Batched(d_temp_storage,\n temp_storage_bytes,\n d_range_srcs,\n d_range_dsts,\n d_range_sizes,\n total_segments);\n \n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n hipcub::DeviceCopy::Batched(d_temp_storage,\n temp_storage_bytes,\n d_range_srcs,\n d_range_dsts,\n d_range_sizes,\n total_segments);\n hipDeviceSynchronize();\n\n return out;\n}\n\ntemplate \nthrust::device_vector gen_power_law_key_segments(seed_t seed,\n std::size_t total_elements,\n std::size_t total_segments)\n{\n thrust::device_vector segment_offsets =\n gen_power_law_offsets(seed, total_elements, total_segments);\n return gen_power_law_key_segments(seed, total_elements, segment_offsets);\n}\n\n#define INSTANTIATE(TYPE) \\\n template thrust::device_vector gen_power_law_key_segments(seed_t, \\\n std::size_t, \\\n std::size_t)\n\nINSTANTIATE(bool);\n\nINSTANTIATE(uint8_t);\nINSTANTIATE(uint16_t);\nINSTANTIATE(uint32_t);\nINSTANTIATE(uint64_t);\nINSTANTIATE(uint128_t);\n\nINSTANTIATE(int8_t);\nINSTANTIATE(int16_t);\nINSTANTIATE(int32_t);\nINSTANTIATE(int64_t);\nINSTANTIATE(int128_t);\n\nINSTANTIATE(float);\nINSTANTIATE(double);\nINSTANTIATE(complex);\n#undef INSTANTIATE\n\ntemplate \nstruct gt_t\n{\n T val;\n\n __device__ bool operator()(T x)\n {\n return x > val;\n }\n};\n\ntemplate \nthrust::device_vector gen_uniform_offsets(seed_t seed,\n T total_elements,\n T min_segment_size,\n T max_segment_size)\n{\n thrust::device_vector segment_offsets(total_elements + 2);\n gen(seed, segment_offsets, bit_entropy::_1_000, min_segment_size, max_segment_size);\n segment_offsets[total_elements] = total_elements + 1;\n thrust::exclusive_scan(segment_offsets.begin(), segment_offsets.end(), segment_offsets.begin());\n typename thrust::device_vector::iterator iter =\n thrust::find_if(segment_offsets.begin(), segment_offsets.end(), gt_t{total_elements});\n *iter = total_elements;\n segment_offsets.erase(iter + 1, segment_offsets.end());\n return segment_offsets;\n}\n\n#define INSTANTIATE(TYPE) \\\n template thrust::device_vector gen_uniform_offsets(seed_t, TYPE, TYPE, TYPE)\n\nINSTANTIATE(uint32_t);\nINSTANTIATE(uint64_t);\n\n#undef INSTANTIATE\n\n/**\n * @brief Generates a vector of random key segments.\n *\n * Not all parameter combinations can be satisfied. For instance, if the total\n * elements is less than the minimal segment size, the function will return a\n * vector with a single element that is outside of the requested range. \n * At most one segment can be out of the requested range.\n */\ntemplate \nthrust::device_vector gen_uniform_key_segments(seed_t seed,\n std::size_t total_elements,\n std::size_t min_segment_size,\n std::size_t max_segment_size)\n{\n thrust::device_vector segment_offsets =\n gen_uniform_offsets(seed, total_elements, min_segment_size, max_segment_size);\n return gen_power_law_key_segments(seed, total_elements, segment_offsets);\n}\n\n#define INSTANTIATE(TYPE) \\\n template thrust::device_vector gen_uniform_key_segments(seed_t, \\\n std::size_t, \\\n std::size_t, \\\n std::size_t)\n\nINSTANTIATE(bool);\n\nINSTANTIATE(uint8_t);\nINSTANTIATE(uint16_t);\nINSTANTIATE(uint32_t);\nINSTANTIATE(uint64_t);\nINSTANTIATE(uint128_t);\n\nINSTANTIATE(int8_t);\nINSTANTIATE(int16_t);\nINSTANTIATE(int32_t);\nINSTANTIATE(int64_t);\nINSTANTIATE(int128_t);\n\nINSTANTIATE(float);\nINSTANTIATE(double);\nINSTANTIATE(complex);\n#undef INSTANTIATE\n", "source": "/kaggle/working/cub/benchmarks/nvbench_helper/nvbench_helper.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceSelect::Unique utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"test_util.h\"\n\n#include \n#include \n\nusing namespace cub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\nfloat g_device_giga_bandwidth;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB DeviceSelect entrypoints\n//---------------------------------------------------------------------\n\n\n/**\n * Dispatch to unique entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceSelect::Unique(d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items);\n }\n return error;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceSelect\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type cub_backend,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items)\n{\n *d_cdp_error = Dispatch(cub_backend,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n d_num_selected_out,\n num_items);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Dispatch to CDP kernel\n */\ntemplate \ncudaError_t Dispatch(Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items)\n{\n // Invoke kernel to invoke device-side dispatch\n cudaError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n Int2Type{},\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n d_num_selected_out,\n num_items);\n CubDebugExit(retval);\n\n // Copy out temp_storage_bytes\n CubDebugExit(cudaMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n cudaMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(cudaMemcpy(&retval,\n d_cdp_error,\n sizeof(cudaError_t) * 1,\n cudaMemcpyDeviceToHost));\n return retval;\n}\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\n\n/**\n * Initialize problem\n */\ntemplate \nvoid Initialize(\n int entropy_reduction,\n T *h_in,\n int num_items,\n int max_segment)\n{\n unsigned int max_int = (unsigned int) -1;\n\n int key = 0;\n int i = 0;\n while (i < num_items)\n {\n // Select number of repeating occurrences for the current run\n int repeat;\n if (max_segment < 0)\n {\n repeat = num_items;\n }\n else if (max_segment < 2)\n {\n repeat = 1;\n }\n else\n {\n RandomBits(repeat, entropy_reduction);\n repeat = (int) ((double(repeat) * double(max_segment)) / double(max_int));\n repeat = CUB_MAX(1, repeat);\n }\n\n int j = i;\n while (j < CUB_MIN(i + repeat, num_items))\n {\n InitValue(INTEGER_SEED, h_in[j], key);\n j++;\n }\n\n i = j;\n key++;\n }\n\n if (g_verbose)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n\n/**\n * Solve unique problem\n */\ntemplate <\n typename InputIteratorT,\n typename T>\nint Solve(\n InputIteratorT h_in,\n T *h_reference,\n int num_items)\n{\n int num_selected = 0;\n if (num_items > 0)\n {\n h_reference[num_selected] = h_in[0];\n num_selected++;\n }\n\n for (int i = 1; i < num_items; ++i)\n {\n if (h_in[i] != h_in[i - 1])\n {\n h_reference[num_selected] = h_in[i];\n num_selected++;\n }\n }\n\n return num_selected;\n}\n\n\n\n/**\n * Test DeviceSelect for a given problem input\n */\ntemplate <\n Backend BACKEND,\n typename DeviceInputIteratorT,\n typename T>\nvoid Test(\n DeviceInputIteratorT d_in,\n T *h_reference,\n int num_selected,\n int num_items)\n{\n // Allocate device output array and num selected\n T *d_out = NULL;\n int *d_num_selected_out = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(T) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_num_selected_out, sizeof(int)));\n\n // Allocate CDP device arrays\n size_t *d_temp_storage_bytes = NULL;\n cudaError_t *d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(cudaError_t) * 1));\n\n // Allocate temporary storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(Int2Type(), 1, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items));\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n // Clear device output array\n CubDebugExit(cudaMemset(d_out, 0, sizeof(T) * num_items));\n CubDebugExit(cudaMemset(d_num_selected_out, 0, sizeof(int)));\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(Int2Type(), 1, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items));\n\n // Check for correctness (and display results, if specified)\n int compare1 = CompareDeviceResults(h_reference, d_out, num_selected, true, g_verbose);\n printf(\"\\t Data %s \", compare1 ? \"FAIL\" : \"PASS\");\n\n int compare2 = CompareDeviceResults(&num_selected, d_num_selected_out, 1, true, g_verbose);\n printf(\"\\t Count %s \", compare2 ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Performance\n GpuTimer gpu_timer;\n gpu_timer.Start();\n CubDebugExit(Dispatch(Int2Type(), g_timing_iterations, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items));\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n\n // Display performance\n if (g_timing_iterations > 0)\n {\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n float giga_bandwidth = float((num_items + num_selected) * sizeof(T)) / avg_millis / 1000.0f / 1000.0f;\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s, %.1f%% peak\", avg_millis, giga_rate, giga_bandwidth, giga_bandwidth / g_device_giga_bandwidth * 100.0);\n }\n printf(\"\\n\\n\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Cleanup\n if (d_out) CubDebugExit(g_allocator.DeviceFree(d_out));\n if (d_num_selected_out) CubDebugExit(g_allocator.DeviceFree(d_num_selected_out));\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare1 | compare2);\n}\n\n\n/**\n * Test DeviceSelect on pointer type\n */\ntemplate <\n Backend BACKEND,\n typename T>\nvoid TestPointer(\n int num_items,\n int entropy_reduction,\n int max_segment)\n{\n // Allocate host arrays\n T* h_in = new T[num_items];\n T* h_reference = new T[num_items];\n\n // Initialize problem and solution\n Initialize(entropy_reduction, h_in, num_items, max_segment);\n int num_selected = Solve(h_in, h_reference, num_items);\n\n printf(\"\\nPointer %s cub::DeviceSelect::Unique %d items, %d selected (avg run length %.3f), %s %d-byte elements, entropy_reduction %d\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items, num_selected, float(num_items) / num_selected,\n typeid(T).name(),\n (int) sizeof(T),\n entropy_reduction);\n fflush(stdout);\n\n // Allocate problem device arrays\n T *d_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_in, sizeof(T) * num_items));\n\n // Initialize device input\n CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(T) * num_items, cudaMemcpyHostToDevice));\n\n // Run Test\n Test(d_in, h_reference, num_selected, num_items);\n\n // Cleanup\n if (h_in) delete[] h_in;\n if (h_reference) delete[] h_reference;\n if (d_in) CubDebugExit(g_allocator.DeviceFree(d_in));\n}\n\n\n/**\n * Test DeviceSelect on iterator type\n */\ntemplate <\n Backend BACKEND,\n typename T>\nvoid TestIterator(\n int num_items)\n{\n // Use a counting iterator as the input\n CountingInputIterator h_in(0);\n\n // Allocate host arrays\n T* h_reference = new T[num_items];\n\n // Initialize problem and solution\n int num_selected = Solve(h_in, h_reference, num_items);\n\n printf(\"\\nIterator %s cub::DeviceSelect::Unique %d items, %d selected (avg run length %.3f), %s %d-byte elements\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items, num_selected, float(num_items) / num_selected,\n typeid(T).name(),\n (int) sizeof(T));\n fflush(stdout);\n\n // Run Test\n Test(h_in, h_reference, num_selected, num_items);\n\n // Cleanup\n if (h_reference) delete[] h_reference;\n}\n\n\n/**\n * Test different gen modes\n */\ntemplate <\n Backend BACKEND,\n typename T>\nvoid Test(\n int num_items)\n{\n for (int max_segment = 1; ((max_segment > 0) && (max_segment < num_items)); max_segment *= 11)\n {\n TestPointer(num_items, 0, max_segment);\n TestPointer(num_items, 2, max_segment);\n TestPointer(num_items, 7, max_segment);\n }\n}\n\n\n/**\n * Test different dispatch\n */\ntemplate <\n typename T>\nvoid TestOp(\n int num_items)\n{\n#if TEST_CDP == 0\n Test(num_items);\n#elif TEST_CDP == 1\n Test(num_items);\n#endif // TEST_CDP\n}\n\n\n/**\n * Test different input sizes\n */\ntemplate \nvoid Test(\n int num_items)\n{\n if (num_items < 0)\n {\n TestOp(0);\n TestOp(1);\n TestOp(100);\n TestOp(10000);\n TestOp(1000000);\n }\n else\n {\n TestOp(num_items);\n }\n}\n\ntemplate \nvoid TestIteratorOp(int num_items)\n{\n void *d_temp_storage{};\n std::size_t temp_storage_size{};\n\n thrust::device_vector num_selected(1);\n\n auto in = thrust::make_counting_iterator(static_cast(0));\n auto out = thrust::make_discard_iterator();\n\n CubDebugExit(cub::DeviceSelect::Unique(d_temp_storage,\n temp_storage_size,\n in,\n out,\n num_selected.begin(),\n num_items));\n\n thrust::device_vector temp_storage(temp_storage_size);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n CubDebugExit(cub::DeviceSelect::Unique(d_temp_storage,\n temp_storage_size,\n in,\n out,\n num_selected.begin(),\n num_items));\n\n AssertEquals(num_selected[0], num_items);\n}\n\ntemplate \nvoid TestIterator(int num_items)\n{\n if (num_items < 0)\n {\n TestIteratorOp(0);\n TestIteratorOp(1);\n TestIteratorOp(100);\n TestIteratorOp(10000);\n TestIteratorOp(1000000);\n }\n else\n {\n TestIteratorOp(num_items);\n }\n}\n\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n int num_items = -1;\n int entropy_reduction = 0;\n int maxseg = 1000;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n args.GetCmdLineArgument(\"maxseg\", maxseg);\n args.GetCmdLineArgument(\"entropy\", entropy_reduction);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--maxseg=]\"\n \"[--entropy=]\"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n g_device_giga_bandwidth = args.device_giga_bandwidth;\n printf(\"\\n\");\n\n // %PARAM% TEST_CDP cdp 0:1\n\n // Test different input types\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n\n TestIterator(num_items);\n\n return 0;\n}\n\n\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceSelect::Unique utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"test_util.h\"\n\n#include \n#include \n\nusing namespace hipcub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\nfloat g_device_giga_bandwidth;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB DeviceSelect entrypoints\n//---------------------------------------------------------------------\n\n\n/**\n * Dispatch to unique entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceSelect::Unique(d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items);\n }\n return error;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceSelect\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type cub_backend,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items)\n{\n *d_cdp_error = Dispatch(cub_backend,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n d_num_selected_out,\n num_items);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Dispatch to CDP kernel\n */\ntemplate \nhipError_t Dispatch(Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items)\n{\n // Invoke kernel to invoke device-side dispatch\n hipError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n Int2Type{},\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n d_num_selected_out,\n num_items);\n CubDebugExit(retval);\n\n // Copy out temp_storage_bytes\n CubDebugExit(hipMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n hipMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(hipMemcpy(&retval,\n d_cdp_error,\n sizeof(hipError_t) * 1,\n hipMemcpyDeviceToHost));\n return retval;\n}\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\n\n/**\n * Initialize problem\n */\ntemplate \nvoid Initialize(\n int entropy_reduction,\n T *h_in,\n int num_items,\n int max_segment)\n{\n unsigned int max_int = (unsigned int) -1;\n\n int key = 0;\n int i = 0;\n while (i < num_items)\n {\n // Select number of repeating occurrences for the current run\n int repeat;\n if (max_segment < 0)\n {\n repeat = num_items;\n }\n else if (max_segment < 2)\n {\n repeat = 1;\n }\n else\n {\n RandomBits(repeat, entropy_reduction);\n repeat = (int) ((double(repeat) * double(max_segment)) / double(max_int));\n repeat = CUB_MAX(1, repeat);\n }\n\n int j = i;\n while (j < CUB_MIN(i + repeat, num_items))\n {\n InitValue(INTEGER_SEED, h_in[j], key);\n j++;\n }\n\n i = j;\n key++;\n }\n\n if (g_verbose)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n\n/**\n * Solve unique problem\n */\ntemplate <\n typename InputIteratorT,\n typename T>\nint Solve(\n InputIteratorT h_in,\n T *h_reference,\n int num_items)\n{\n int num_selected = 0;\n if (num_items > 0)\n {\n h_reference[num_selected] = h_in[0];\n num_selected++;\n }\n\n for (int i = 1; i < num_items; ++i)\n {\n if (h_in[i] != h_in[i - 1])\n {\n h_reference[num_selected] = h_in[i];\n num_selected++;\n }\n }\n\n return num_selected;\n}\n\n\n\n/**\n * Test DeviceSelect for a given problem input\n */\ntemplate <\n Backend BACKEND,\n typename DeviceInputIteratorT,\n typename T>\nvoid Test(\n DeviceInputIteratorT d_in,\n T *h_reference,\n int num_selected,\n int num_items)\n{\n // Allocate device output array and num selected\n T *d_out = NULL;\n int *d_num_selected_out = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(T) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_num_selected_out, sizeof(int)));\n\n // Allocate CDP device arrays\n size_t *d_temp_storage_bytes = NULL;\n hipError_t *d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(hipError_t) * 1));\n\n // Allocate temporary storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(Int2Type(), 1, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items));\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n // Clear device output array\n CubDebugExit(hipMemset(d_out, 0, sizeof(T) * num_items));\n CubDebugExit(hipMemset(d_num_selected_out, 0, sizeof(int)));\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(Int2Type(), 1, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items));\n\n // Check for correctness (and display results, if specified)\n int compare1 = CompareDeviceResults(h_reference, d_out, num_selected, true, g_verbose);\n printf(\"\\t Data %s \", compare1 ? \"FAIL\" : \"PASS\");\n\n int compare2 = CompareDeviceResults(&num_selected, d_num_selected_out, 1, true, g_verbose);\n printf(\"\\t Count %s \", compare2 ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Performance\n GpuTimer gpu_timer;\n gpu_timer.Start();\n CubDebugExit(Dispatch(Int2Type(), g_timing_iterations, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items));\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n\n // Display performance\n if (g_timing_iterations > 0)\n {\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n float giga_bandwidth = float((num_items + num_selected) * sizeof(T)) / avg_millis / 1000.0f / 1000.0f;\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s, %.1f%% peak\", avg_millis, giga_rate, giga_bandwidth, giga_bandwidth / g_device_giga_bandwidth * 100.0);\n }\n printf(\"\\n\\n\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Cleanup\n if (d_out) CubDebugExit(g_allocator.DeviceFree(d_out));\n if (d_num_selected_out) CubDebugExit(g_allocator.DeviceFree(d_num_selected_out));\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare1 | compare2);\n}\n\n\n/**\n * Test DeviceSelect on pointer type\n */\ntemplate <\n Backend BACKEND,\n typename T>\nvoid TestPointer(\n int num_items,\n int entropy_reduction,\n int max_segment)\n{\n // Allocate host arrays\n T* h_in = new T[num_items];\n T* h_reference = new T[num_items];\n\n // Initialize problem and solution\n Initialize(entropy_reduction, h_in, num_items, max_segment);\n int num_selected = Solve(h_in, h_reference, num_items);\n\n printf(\"\\nPointer %s hipcub::DeviceSelect::Unique %d items, %d selected (avg run length %.3f), %s %d-byte elements, entropy_reduction %d\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items, num_selected, float(num_items) / num_selected,\n typeid(T).name(),\n (int) sizeof(T),\n entropy_reduction);\n fflush(stdout);\n\n // Allocate problem device arrays\n T *d_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_in, sizeof(T) * num_items));\n\n // Initialize device input\n CubDebugExit(hipMemcpy(d_in, h_in, sizeof(T) * num_items, hipMemcpyHostToDevice));\n\n // Run Test\n Test(d_in, h_reference, num_selected, num_items);\n\n // Cleanup\n if (h_in) delete[] h_in;\n if (h_reference) delete[] h_reference;\n if (d_in) CubDebugExit(g_allocator.DeviceFree(d_in));\n}\n\n\n/**\n * Test DeviceSelect on iterator type\n */\ntemplate <\n Backend BACKEND,\n typename T>\nvoid TestIterator(\n int num_items)\n{\n // Use a counting iterator as the input\n CountingInputIterator h_in(0);\n\n // Allocate host arrays\n T* h_reference = new T[num_items];\n\n // Initialize problem and solution\n int num_selected = Solve(h_in, h_reference, num_items);\n\n printf(\"\\nIterator %s hipcub::DeviceSelect::Unique %d items, %d selected (avg run length %.3f), %s %d-byte elements\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items, num_selected, float(num_items) / num_selected,\n typeid(T).name(),\n (int) sizeof(T));\n fflush(stdout);\n\n // Run Test\n Test(h_in, h_reference, num_selected, num_items);\n\n // Cleanup\n if (h_reference) delete[] h_reference;\n}\n\n\n/**\n * Test different gen modes\n */\ntemplate <\n Backend BACKEND,\n typename T>\nvoid Test(\n int num_items)\n{\n for (int max_segment = 1; ((max_segment > 0) && (max_segment < num_items)); max_segment *= 11)\n {\n TestPointer(num_items, 0, max_segment);\n TestPointer(num_items, 2, max_segment);\n TestPointer(num_items, 7, max_segment);\n }\n}\n\n\n/**\n * Test different dispatch\n */\ntemplate <\n typename T>\nvoid TestOp(\n int num_items)\n{\n#if TEST_CDP == 0\n Test(num_items);\n#elif TEST_CDP == 1\n Test(num_items);\n#endif // TEST_CDP\n}\n\n\n/**\n * Test different input sizes\n */\ntemplate \nvoid Test(\n int num_items)\n{\n if (num_items < 0)\n {\n TestOp(0);\n TestOp(1);\n TestOp(100);\n TestOp(10000);\n TestOp(1000000);\n }\n else\n {\n TestOp(num_items);\n }\n}\n\ntemplate \nvoid TestIteratorOp(int num_items)\n{\n void *d_temp_storage{};\n std::size_t temp_storage_size{};\n\n thrust::device_vector num_selected(1);\n\n auto in = thrust::make_counting_iterator(static_cast(0));\n auto out = thrust::make_discard_iterator();\n\n CubDebugExit(hipcub::DeviceSelect::Unique(d_temp_storage,\n temp_storage_size,\n in,\n out,\n num_selected.begin(),\n num_items));\n\n thrust::device_vector temp_storage(temp_storage_size);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n CubDebugExit(hipcub::DeviceSelect::Unique(d_temp_storage,\n temp_storage_size,\n in,\n out,\n num_selected.begin(),\n num_items));\n\n AssertEquals(num_selected[0], num_items);\n}\n\ntemplate \nvoid TestIterator(int num_items)\n{\n if (num_items < 0)\n {\n TestIteratorOp(0);\n TestIteratorOp(1);\n TestIteratorOp(100);\n TestIteratorOp(10000);\n TestIteratorOp(1000000);\n }\n else\n {\n TestIteratorOp(num_items);\n }\n}\n\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n int num_items = -1;\n int entropy_reduction = 0;\n int maxseg = 1000;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n args.GetCmdLineArgument(\"maxseg\", maxseg);\n args.GetCmdLineArgument(\"entropy\", entropy_reduction);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--maxseg=]\"\n \"[--entropy=]\"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n g_device_giga_bandwidth = args.device_giga_bandwidth;\n printf(\"\\n\");\n\n // %PARAM% TEST_CDP cdp 0:1\n\n // Test different input types\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n\n TestIterator(num_items);\n\n return 0;\n}\n\n\n\n", "source": "/kaggle/working/cub/test/test_device_select_unique.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceReduce::ReduceByKey utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \"test_util.h\"\n\n#include \n#include \n\nusing namespace cub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB entrypoints\n//---------------------------------------------------------------------\n\n/**\n * Dispatch to reduce-by-key entrypoint\n */\ntemplate <\n typename KeyInputIteratorT,\n typename KeyOutputIteratorT,\n typename ValueInputIteratorT,\n typename ValueOutputIteratorT,\n typename NumRunsIteratorT,\n typename EqualityOpT,\n typename ReductionOpT,\n typename OffsetT>\nCUB_RUNTIME_FUNCTION __forceinline__\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n KeyInputIteratorT d_keys_in,\n KeyOutputIteratorT d_keys_out,\n ValueInputIteratorT d_values_in,\n ValueOutputIteratorT d_values_out,\n NumRunsIteratorT d_num_runs,\n EqualityOpT /*equality_op*/,\n ReductionOpT reduction_op,\n OffsetT num_items)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceReduce::ReduceByKey(\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_keys_out,\n d_values_in,\n d_values_out,\n d_num_runs,\n reduction_op,\n num_items);\n }\n return error;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceSelect\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type cub_backend,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n KeyInputIteratorT d_keys_in,\n KeyOutputIteratorT d_keys_out,\n ValueInputIteratorT d_values_in,\n ValueOutputIteratorT d_values_out,\n NumRunsIteratorT d_num_runs,\n EqualityOpT equality_op,\n ReductionOpT reduction_op,\n OffsetT num_items)\n{\n\n *d_cdp_error = Dispatch(cub_backend,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_keys_out,\n d_values_in,\n d_values_out,\n d_num_runs,\n equality_op,\n reduction_op,\n num_items);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Dispatch to CDP kernel\n */\ntemplate \n__forceinline__ cudaError_t\nDispatch(Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n KeyInputIteratorT d_keys_in,\n KeyOutputIteratorT d_keys_out,\n ValueInputIteratorT d_values_in,\n ValueOutputIteratorT d_values_out,\n NumRunsIteratorT d_num_runs,\n EqualityOpT equality_op,\n ReductionOpT reduction_op,\n OffsetT num_items)\n{\n // Invoke kernel to invoke device-side dispatch\n cudaError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n Int2Type{},\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_keys_out,\n d_values_in,\n d_values_out,\n d_num_runs,\n equality_op,\n reduction_op,\n num_items);\n CubDebugExit(retval);\n\n // Copy out temp_storage_bytes\n CubDebugExit(cudaMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n cudaMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(cudaMemcpy(&retval,\n d_cdp_error,\n sizeof(cudaError_t) * 1,\n cudaMemcpyDeviceToHost));\n return retval;\n}\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\n\n/**\n * Initialize problem\n */\ntemplate \nvoid Initialize(\n int entropy_reduction,\n T *h_in,\n int num_items,\n int max_segment)\n{\n unsigned int max_int = (unsigned int) -1;\n\n int key = 0;\n int i = 0;\n while (i < num_items)\n {\n // Select number of repeating occurrences\n\n int repeat;\n\n if (max_segment < 0)\n {\n repeat = num_items;\n }\n else if (max_segment < 2)\n {\n repeat = 1;\n }\n else\n {\n RandomBits(repeat, entropy_reduction);\n repeat = (int) ((double(repeat) * double(max_segment)) / double(max_int));\n repeat = CUB_MAX(1, repeat);\n }\n\n int j = i;\n while (j < CUB_MIN(i + repeat, num_items))\n {\n InitValue(INTEGER_SEED, h_in[j], key);\n j++;\n }\n\n i = j;\n key++;\n }\n\n if (g_verbose)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n\n/**\n * Solve problem. Returns total number of segments identified\n */\ntemplate <\n typename KeyInputIteratorT,\n typename ValueInputIteratorT,\n typename KeyT,\n typename ValueT,\n typename EqualityOpT,\n typename ReductionOpT>\nint Solve(\n KeyInputIteratorT h_keys_in,\n KeyT *h_keys_reference,\n ValueInputIteratorT h_values_in,\n ValueT *h_values_reference,\n EqualityOpT equality_op,\n ReductionOpT reduction_op,\n int num_items)\n{\n using AccumT = cub::detail::accumulator_t;\n\n // First item\n KeyT previous = h_keys_in[0];\n AccumT aggregate = h_values_in[0];\n int num_segments = 0;\n\n // Subsequent items\n for (int i = 1; i < num_items; ++i)\n {\n if (!equality_op(previous, h_keys_in[i]))\n {\n h_keys_reference[num_segments] = previous;\n h_values_reference[num_segments] = static_cast(aggregate);\n num_segments++;\n aggregate = h_values_in[i];\n }\n else\n {\n aggregate = static_cast(reduction_op(aggregate, h_values_in[i]));\n }\n previous = h_keys_in[i];\n }\n\n h_keys_reference[num_segments] = previous;\n h_values_reference[num_segments] = static_cast(aggregate);\n num_segments++;\n\n return num_segments;\n}\n\n\n\n/**\n * Test DeviceSelect for a given problem input\n */\ntemplate <\n Backend BACKEND,\n typename DeviceKeyInputIteratorT,\n typename DeviceValueInputIteratorT,\n typename KeyT,\n typename ValueT,\n typename EqualityOpT,\n typename ReductionOpT>\nvoid Test(\n DeviceKeyInputIteratorT d_keys_in,\n DeviceValueInputIteratorT d_values_in,\n KeyT* h_keys_reference,\n ValueT* h_values_reference,\n EqualityOpT equality_op,\n ReductionOpT reduction_op,\n int num_segments,\n int num_items)\n{\n // Allocate device output arrays and number of segments\n KeyT* d_keys_out = NULL;\n ValueT* d_values_out = NULL;\n int* d_num_runs = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys_out, sizeof(KeyT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_values_out, sizeof(ValueT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_num_runs, sizeof(int)));\n\n // Allocate CDP device arrays\n size_t *d_temp_storage_bytes = NULL;\n cudaError_t *d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(cudaError_t) * 1));\n\n // Allocate temporary storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(Int2Type(), 1, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, d_num_runs, equality_op, reduction_op, num_items));\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n // Clear device output arrays\n CubDebugExit(cudaMemset(d_keys_out, 0, sizeof(KeyT) * num_items));\n CubDebugExit(cudaMemset(d_values_out, 0, sizeof(ValueT) * num_items));\n CubDebugExit(cudaMemset(d_num_runs, 0, sizeof(int)));\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(Int2Type(), 1, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, d_num_runs, equality_op, reduction_op, num_items));\n\n // Check for correctness (and display results, if specified)\n int compare1 = CompareDeviceResults(h_keys_reference, d_keys_out, num_segments, true, g_verbose);\n printf(\"\\t Keys %s \", compare1 ? \"FAIL\" : \"PASS\");\n\n int compare2 = CompareDeviceResults(h_values_reference, d_values_out, num_segments, true, g_verbose);\n printf(\"\\t Values %s \", compare2 ? \"FAIL\" : \"PASS\");\n\n int compare3 = CompareDeviceResults(&num_segments, d_num_runs, 1, true, g_verbose);\n printf(\"\\t Count %s \", compare3 ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Performance\n GpuTimer gpu_timer;\n gpu_timer.Start();\n CubDebugExit(Dispatch(Int2Type(), g_timing_iterations, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, d_num_runs, equality_op, reduction_op, num_items));\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n\n // Display performance\n if (g_timing_iterations > 0)\n {\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n int bytes_moved = ((num_items + num_segments) * sizeof(KeyT)) + ((num_items + num_segments) * sizeof(ValueT));\n float giga_bandwidth = float(bytes_moved) / avg_millis / 1000.0f / 1000.0f;\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s\", avg_millis, giga_rate, giga_bandwidth);\n }\n printf(\"\\n\\n\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Cleanup\n if (d_keys_out) CubDebugExit(g_allocator.DeviceFree(d_keys_out));\n if (d_values_out) CubDebugExit(g_allocator.DeviceFree(d_values_out));\n if (d_num_runs) CubDebugExit(g_allocator.DeviceFree(d_num_runs));\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare1 | compare2 | compare3);\n}\n\n\n/**\n * Test DeviceSelect on pointer type\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT,\n typename ReductionOpT>\nvoid TestPointer(\n int num_items,\n int entropy_reduction,\n int max_segment,\n ReductionOpT reduction_op)\n{\n // Allocate host arrays\n KeyT* h_keys_in = new KeyT[num_items];\n KeyT* h_keys_reference = new KeyT[num_items];\n\n ValueT* h_values_in = new ValueT[num_items];\n ValueT* h_values_reference = new ValueT[num_items];\n\n for (int i = 0; i < num_items; ++i)\n InitValue(INTEGER_SEED, h_values_in[i], 1);\n\n // Initialize problem and solution\n Equality equality_op;\n Initialize(entropy_reduction, h_keys_in, num_items, max_segment);\n int num_segments = Solve(h_keys_in, h_keys_reference, h_values_in, h_values_reference, equality_op, reduction_op, num_items);\n\n printf(\"\\nPointer %s cub::DeviceReduce::ReduceByKey %s reduction of %d items, %d segments (avg run length %.3f), {%s,%s} key value pairs, max_segment %d, entropy_reduction %d\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n (std::is_same::value) ? \"Sum\" : \"Max\",\n num_items, num_segments, float(num_items) / num_segments,\n typeid(KeyT).name(), typeid(ValueT).name(),\n max_segment, entropy_reduction);\n fflush(stdout);\n\n // Allocate problem device arrays\n KeyT *d_keys_in = NULL;\n ValueT *d_values_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys_in, sizeof(KeyT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_values_in, sizeof(ValueT) * num_items));\n\n // Initialize device input\n CubDebugExit(cudaMemcpy(d_keys_in, h_keys_in, sizeof(KeyT) * num_items, cudaMemcpyHostToDevice));\n CubDebugExit(cudaMemcpy(d_values_in, h_values_in, sizeof(ValueT) * num_items, cudaMemcpyHostToDevice));\n\n // Run Test\n Test(d_keys_in, d_values_in, h_keys_reference, h_values_reference, equality_op, reduction_op, num_segments, num_items);\n\n // Cleanup\n if (h_keys_in) delete[] h_keys_in;\n if (h_values_in) delete[] h_values_in;\n if (h_keys_reference) delete[] h_keys_reference;\n if (h_values_reference) delete[] h_values_reference;\n if (d_keys_in) CubDebugExit(g_allocator.DeviceFree(d_keys_in));\n if (d_values_in) CubDebugExit(g_allocator.DeviceFree(d_values_in));\n}\n\n\n/**\n * Test on iterator type\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT,\n typename ReductionOpT>\nvoid TestIterator(\n int num_items,\n int entropy_reduction,\n int max_segment,\n ReductionOpT reduction_op)\n{\n // Allocate host arrays\n KeyT* h_keys_in = new KeyT[num_items];\n KeyT* h_keys_reference = new KeyT[num_items];\n\n ValueT one_val;\n InitValue(INTEGER_SEED, one_val, 1);\n ConstantInputIterator h_values_in(one_val);\n ValueT* h_values_reference = new ValueT[num_items];\n\n // Initialize problem and solution\n Equality equality_op;\n Initialize(entropy_reduction, h_keys_in, num_items, max_segment);\n int num_segments = Solve(h_keys_in, h_keys_reference, h_values_in, h_values_reference, equality_op, reduction_op, num_items);\n\n printf(\"\\nIterator %s cub::DeviceReduce::ReduceByKey %s reduction of %d items, %d segments (avg run length %.3f), {%s,%s} key value pairs, max_segment %d, entropy_reduction %d\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n (std::is_same::value) ? \"Sum\" : \"Max\",\n num_items, num_segments, float(num_items) / num_segments,\n typeid(KeyT).name(), typeid(ValueT).name(),\n max_segment, entropy_reduction);\n fflush(stdout);\n\n // Allocate problem device arrays\n KeyT *d_keys_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys_in, sizeof(KeyT) * num_items));\n\n // Initialize device input\n CubDebugExit(cudaMemcpy(d_keys_in, h_keys_in, sizeof(KeyT) * num_items, cudaMemcpyHostToDevice));\n\n // Run Test\n Test(d_keys_in, h_values_in, h_keys_reference, h_values_reference, equality_op, reduction_op, num_segments, num_items);\n\n // Cleanup\n if (h_keys_in) delete[] h_keys_in;\n if (h_keys_reference) delete[] h_keys_reference;\n if (h_values_reference) delete[] h_values_reference;\n if (d_keys_in) CubDebugExit(g_allocator.DeviceFree(d_keys_in));\n}\n\n\n/**\n * Test different gen modes\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT,\n typename ReductionOpT>\nvoid Test(\n int num_items,\n ReductionOpT reduction_op,\n int max_segment)\n{\n // 0 key-bit entropy reduction rounds\n TestPointer(num_items, 0, max_segment, reduction_op);\n\n if (max_segment > 1)\n {\n // 2 key-bit entropy reduction rounds\n TestPointer(num_items, 2, max_segment, reduction_op);\n\n // 7 key-bit entropy reduction rounds\n TestPointer(num_items, 7, max_segment, reduction_op);\n }\n}\n\n\n/**\n * Test different avg segment lengths modes\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT,\n typename ReductionOpT>\nvoid Test(\n int num_items,\n ReductionOpT reduction_op)\n{\n Test(num_items, reduction_op, -1);\n Test(num_items, reduction_op, 1);\n\n // Evaluate different max-segment lengths\n for (int max_segment = 3; max_segment < CUB_MIN(num_items, (unsigned short) -1); max_segment *= 11)\n {\n Test(num_items, reduction_op, max_segment);\n }\n}\n\n\n\n/**\n * Test different dispatch\n */\ntemplate <\n typename KeyT,\n typename ValueT,\n typename ReductionOpT>\nvoid TestDispatch(\n int num_items,\n ReductionOpT reduction_op)\n{\n#if TEST_CDP == 0\n Test(num_items, reduction_op);\n#elif TEST_CDP == 1\n Test(num_items, reduction_op);\n#endif // TEST_CDP\n}\n\n\n/**\n * Test different input sizes\n */\ntemplate <\n typename KeyT,\n typename ValueT,\n typename ReductionOpT>\nvoid TestSize(\n int num_items,\n ReductionOpT reduction_op)\n{\n if (num_items < 0)\n {\n TestDispatch(1, reduction_op);\n TestDispatch(100, reduction_op);\n TestDispatch(10000, reduction_op);\n TestDispatch(1000000, reduction_op);\n }\n else\n {\n TestDispatch(num_items, reduction_op);\n }\n\n}\n\n\ntemplate <\n typename KeyT,\n typename ValueT>\nvoid TestOp(\n int num_items)\n{\n TestSize(num_items, cub::Sum());\n TestSize(num_items, cub::Max());\n}\n\n\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n int num_items = -1;\n int entropy_reduction = 0;\n int maxseg = 1000;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n args.GetCmdLineArgument(\"maxseg\", maxseg);\n args.GetCmdLineArgument(\"entropy\", entropy_reduction);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--maxseg=]\"\n \"[--entropy=]\"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n printf(\"\\n\");\n\n // Get ptx version\n int ptx_version = 0;\n CubDebugExit(PtxVersion(ptx_version));\n\n // %PARAM% TEST_CDP cdp 0:1\n\n // Test different input types\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n\n return 0;\n}\n\n\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceReduce::ReduceByKey utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \"test_util.h\"\n\n#include \n#include \n\nusing namespace hipcub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB entrypoints\n//---------------------------------------------------------------------\n\n/**\n * Dispatch to reduce-by-key entrypoint\n */\ntemplate <\n typename KeyInputIteratorT,\n typename KeyOutputIteratorT,\n typename ValueInputIteratorT,\n typename ValueOutputIteratorT,\n typename NumRunsIteratorT,\n typename EqualityOpT,\n typename ReductionOpT,\n typename OffsetT>\nHIPCUB_RUNTIME_FUNCTION __forceinline__\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n KeyInputIteratorT d_keys_in,\n KeyOutputIteratorT d_keys_out,\n ValueInputIteratorT d_values_in,\n ValueOutputIteratorT d_values_out,\n NumRunsIteratorT d_num_runs,\n EqualityOpT /*equality_op*/,\n ReductionOpT reduction_op,\n OffsetT num_items)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceReduce::ReduceByKey(\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_keys_out,\n d_values_in,\n d_values_out,\n d_num_runs,\n reduction_op,\n num_items);\n }\n return error;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceSelect\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type cub_backend,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n KeyInputIteratorT d_keys_in,\n KeyOutputIteratorT d_keys_out,\n ValueInputIteratorT d_values_in,\n ValueOutputIteratorT d_values_out,\n NumRunsIteratorT d_num_runs,\n EqualityOpT equality_op,\n ReductionOpT reduction_op,\n OffsetT num_items)\n{\n\n *d_cdp_error = Dispatch(cub_backend,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_keys_out,\n d_values_in,\n d_values_out,\n d_num_runs,\n equality_op,\n reduction_op,\n num_items);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Dispatch to CDP kernel\n */\ntemplate \n__forceinline__ hipError_t\nDispatch(Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n KeyInputIteratorT d_keys_in,\n KeyOutputIteratorT d_keys_out,\n ValueInputIteratorT d_values_in,\n ValueOutputIteratorT d_values_out,\n NumRunsIteratorT d_num_runs,\n EqualityOpT equality_op,\n ReductionOpT reduction_op,\n OffsetT num_items)\n{\n // Invoke kernel to invoke device-side dispatch\n hipError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n Int2Type{},\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_keys_out,\n d_values_in,\n d_values_out,\n d_num_runs,\n equality_op,\n reduction_op,\n num_items);\n CubDebugExit(retval);\n\n // Copy out temp_storage_bytes\n CubDebugExit(hipMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n hipMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(hipMemcpy(&retval,\n d_cdp_error,\n sizeof(hipError_t) * 1,\n hipMemcpyDeviceToHost));\n return retval;\n}\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\n\n/**\n * Initialize problem\n */\ntemplate \nvoid Initialize(\n int entropy_reduction,\n T *h_in,\n int num_items,\n int max_segment)\n{\n unsigned int max_int = (unsigned int) -1;\n\n int key = 0;\n int i = 0;\n while (i < num_items)\n {\n // Select number of repeating occurrences\n\n int repeat;\n\n if (max_segment < 0)\n {\n repeat = num_items;\n }\n else if (max_segment < 2)\n {\n repeat = 1;\n }\n else\n {\n RandomBits(repeat, entropy_reduction);\n repeat = (int) ((double(repeat) * double(max_segment)) / double(max_int));\n repeat = CUB_MAX(1, repeat);\n }\n\n int j = i;\n while (j < CUB_MIN(i + repeat, num_items))\n {\n InitValue(INTEGER_SEED, h_in[j], key);\n j++;\n }\n\n i = j;\n key++;\n }\n\n if (g_verbose)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n\n/**\n * Solve problem. Returns total number of segments identified\n */\ntemplate <\n typename KeyInputIteratorT,\n typename ValueInputIteratorT,\n typename KeyT,\n typename ValueT,\n typename EqualityOpT,\n typename ReductionOpT>\nint Solve(\n KeyInputIteratorT h_keys_in,\n KeyT *h_keys_reference,\n ValueInputIteratorT h_values_in,\n ValueT *h_values_reference,\n EqualityOpT equality_op,\n ReductionOpT reduction_op,\n int num_items)\n{\n using AccumT = hipcub::detail::accumulator_t;\n\n // First item\n KeyT previous = h_keys_in[0];\n AccumT aggregate = h_values_in[0];\n int num_segments = 0;\n\n // Subsequent items\n for (int i = 1; i < num_items; ++i)\n {\n if (!equality_op(previous, h_keys_in[i]))\n {\n h_keys_reference[num_segments] = previous;\n h_values_reference[num_segments] = static_cast(aggregate);\n num_segments++;\n aggregate = h_values_in[i];\n }\n else\n {\n aggregate = static_cast(reduction_op(aggregate, h_values_in[i]));\n }\n previous = h_keys_in[i];\n }\n\n h_keys_reference[num_segments] = previous;\n h_values_reference[num_segments] = static_cast(aggregate);\n num_segments++;\n\n return num_segments;\n}\n\n\n\n/**\n * Test DeviceSelect for a given problem input\n */\ntemplate <\n Backend BACKEND,\n typename DeviceKeyInputIteratorT,\n typename DeviceValueInputIteratorT,\n typename KeyT,\n typename ValueT,\n typename EqualityOpT,\n typename ReductionOpT>\nvoid Test(\n DeviceKeyInputIteratorT d_keys_in,\n DeviceValueInputIteratorT d_values_in,\n KeyT* h_keys_reference,\n ValueT* h_values_reference,\n EqualityOpT equality_op,\n ReductionOpT reduction_op,\n int num_segments,\n int num_items)\n{\n // Allocate device output arrays and number of segments\n KeyT* d_keys_out = NULL;\n ValueT* d_values_out = NULL;\n int* d_num_runs = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys_out, sizeof(KeyT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_values_out, sizeof(ValueT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_num_runs, sizeof(int)));\n\n // Allocate CDP device arrays\n size_t *d_temp_storage_bytes = NULL;\n hipError_t *d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(hipError_t) * 1));\n\n // Allocate temporary storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(Int2Type(), 1, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, d_num_runs, equality_op, reduction_op, num_items));\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n // Clear device output arrays\n CubDebugExit(hipMemset(d_keys_out, 0, sizeof(KeyT) * num_items));\n CubDebugExit(hipMemset(d_values_out, 0, sizeof(ValueT) * num_items));\n CubDebugExit(hipMemset(d_num_runs, 0, sizeof(int)));\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(Int2Type(), 1, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, d_num_runs, equality_op, reduction_op, num_items));\n\n // Check for correctness (and display results, if specified)\n int compare1 = CompareDeviceResults(h_keys_reference, d_keys_out, num_segments, true, g_verbose);\n printf(\"\\t Keys %s \", compare1 ? \"FAIL\" : \"PASS\");\n\n int compare2 = CompareDeviceResults(h_values_reference, d_values_out, num_segments, true, g_verbose);\n printf(\"\\t Values %s \", compare2 ? \"FAIL\" : \"PASS\");\n\n int compare3 = CompareDeviceResults(&num_segments, d_num_runs, 1, true, g_verbose);\n printf(\"\\t Count %s \", compare3 ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Performance\n GpuTimer gpu_timer;\n gpu_timer.Start();\n CubDebugExit(Dispatch(Int2Type(), g_timing_iterations, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, d_num_runs, equality_op, reduction_op, num_items));\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n\n // Display performance\n if (g_timing_iterations > 0)\n {\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n int bytes_moved = ((num_items + num_segments) * sizeof(KeyT)) + ((num_items + num_segments) * sizeof(ValueT));\n float giga_bandwidth = float(bytes_moved) / avg_millis / 1000.0f / 1000.0f;\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s\", avg_millis, giga_rate, giga_bandwidth);\n }\n printf(\"\\n\\n\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Cleanup\n if (d_keys_out) CubDebugExit(g_allocator.DeviceFree(d_keys_out));\n if (d_values_out) CubDebugExit(g_allocator.DeviceFree(d_values_out));\n if (d_num_runs) CubDebugExit(g_allocator.DeviceFree(d_num_runs));\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare1 | compare2 | compare3);\n}\n\n\n/**\n * Test DeviceSelect on pointer type\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT,\n typename ReductionOpT>\nvoid TestPointer(\n int num_items,\n int entropy_reduction,\n int max_segment,\n ReductionOpT reduction_op)\n{\n // Allocate host arrays\n KeyT* h_keys_in = new KeyT[num_items];\n KeyT* h_keys_reference = new KeyT[num_items];\n\n ValueT* h_values_in = new ValueT[num_items];\n ValueT* h_values_reference = new ValueT[num_items];\n\n for (int i = 0; i < num_items; ++i)\n InitValue(INTEGER_SEED, h_values_in[i], 1);\n\n // Initialize problem and solution\n Equality equality_op;\n Initialize(entropy_reduction, h_keys_in, num_items, max_segment);\n int num_segments = Solve(h_keys_in, h_keys_reference, h_values_in, h_values_reference, equality_op, reduction_op, num_items);\n\n printf(\"\\nPointer %s hipcub::DeviceReduce::ReduceByKey %s reduction of %d items, %d segments (avg run length %.3f), {%s,%s} key value pairs, max_segment %d, entropy_reduction %d\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n (std::is_same::value) ? \"Sum\" : \"Max\",\n num_items, num_segments, float(num_items) / num_segments,\n typeid(KeyT).name(), typeid(ValueT).name(),\n max_segment, entropy_reduction);\n fflush(stdout);\n\n // Allocate problem device arrays\n KeyT *d_keys_in = NULL;\n ValueT *d_values_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys_in, sizeof(KeyT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_values_in, sizeof(ValueT) * num_items));\n\n // Initialize device input\n CubDebugExit(hipMemcpy(d_keys_in, h_keys_in, sizeof(KeyT) * num_items, hipMemcpyHostToDevice));\n CubDebugExit(hipMemcpy(d_values_in, h_values_in, sizeof(ValueT) * num_items, hipMemcpyHostToDevice));\n\n // Run Test\n Test(d_keys_in, d_values_in, h_keys_reference, h_values_reference, equality_op, reduction_op, num_segments, num_items);\n\n // Cleanup\n if (h_keys_in) delete[] h_keys_in;\n if (h_values_in) delete[] h_values_in;\n if (h_keys_reference) delete[] h_keys_reference;\n if (h_values_reference) delete[] h_values_reference;\n if (d_keys_in) CubDebugExit(g_allocator.DeviceFree(d_keys_in));\n if (d_values_in) CubDebugExit(g_allocator.DeviceFree(d_values_in));\n}\n\n\n/**\n * Test on iterator type\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT,\n typename ReductionOpT>\nvoid TestIterator(\n int num_items,\n int entropy_reduction,\n int max_segment,\n ReductionOpT reduction_op)\n{\n // Allocate host arrays\n KeyT* h_keys_in = new KeyT[num_items];\n KeyT* h_keys_reference = new KeyT[num_items];\n\n ValueT one_val;\n InitValue(INTEGER_SEED, one_val, 1);\n ConstantInputIterator h_values_in(one_val);\n ValueT* h_values_reference = new ValueT[num_items];\n\n // Initialize problem and solution\n Equality equality_op;\n Initialize(entropy_reduction, h_keys_in, num_items, max_segment);\n int num_segments = Solve(h_keys_in, h_keys_reference, h_values_in, h_values_reference, equality_op, reduction_op, num_items);\n\n printf(\"\\nIterator %s hipcub::DeviceReduce::ReduceByKey %s reduction of %d items, %d segments (avg run length %.3f), {%s,%s} key value pairs, max_segment %d, entropy_reduction %d\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n (std::is_same::value) ? \"Sum\" : \"Max\",\n num_items, num_segments, float(num_items) / num_segments,\n typeid(KeyT).name(), typeid(ValueT).name(),\n max_segment, entropy_reduction);\n fflush(stdout);\n\n // Allocate problem device arrays\n KeyT *d_keys_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys_in, sizeof(KeyT) * num_items));\n\n // Initialize device input\n CubDebugExit(hipMemcpy(d_keys_in, h_keys_in, sizeof(KeyT) * num_items, hipMemcpyHostToDevice));\n\n // Run Test\n Test(d_keys_in, h_values_in, h_keys_reference, h_values_reference, equality_op, reduction_op, num_segments, num_items);\n\n // Cleanup\n if (h_keys_in) delete[] h_keys_in;\n if (h_keys_reference) delete[] h_keys_reference;\n if (h_values_reference) delete[] h_values_reference;\n if (d_keys_in) CubDebugExit(g_allocator.DeviceFree(d_keys_in));\n}\n\n\n/**\n * Test different gen modes\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT,\n typename ReductionOpT>\nvoid Test(\n int num_items,\n ReductionOpT reduction_op,\n int max_segment)\n{\n // 0 key-bit entropy reduction rounds\n TestPointer(num_items, 0, max_segment, reduction_op);\n\n if (max_segment > 1)\n {\n // 2 key-bit entropy reduction rounds\n TestPointer(num_items, 2, max_segment, reduction_op);\n\n // 7 key-bit entropy reduction rounds\n TestPointer(num_items, 7, max_segment, reduction_op);\n }\n}\n\n\n/**\n * Test different avg segment lengths modes\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT,\n typename ReductionOpT>\nvoid Test(\n int num_items,\n ReductionOpT reduction_op)\n{\n Test(num_items, reduction_op, -1);\n Test(num_items, reduction_op, 1);\n\n // Evaluate different max-segment lengths\n for (int max_segment = 3; max_segment < CUB_MIN(num_items, (unsigned short) -1); max_segment *= 11)\n {\n Test(num_items, reduction_op, max_segment);\n }\n}\n\n\n\n/**\n * Test different dispatch\n */\ntemplate <\n typename KeyT,\n typename ValueT,\n typename ReductionOpT>\nvoid TestDispatch(\n int num_items,\n ReductionOpT reduction_op)\n{\n#if TEST_CDP == 0\n Test(num_items, reduction_op);\n#elif TEST_CDP == 1\n Test(num_items, reduction_op);\n#endif // TEST_CDP\n}\n\n\n/**\n * Test different input sizes\n */\ntemplate <\n typename KeyT,\n typename ValueT,\n typename ReductionOpT>\nvoid TestSize(\n int num_items,\n ReductionOpT reduction_op)\n{\n if (num_items < 0)\n {\n TestDispatch(1, reduction_op);\n TestDispatch(100, reduction_op);\n TestDispatch(10000, reduction_op);\n TestDispatch(1000000, reduction_op);\n }\n else\n {\n TestDispatch(num_items, reduction_op);\n }\n\n}\n\n\ntemplate <\n typename KeyT,\n typename ValueT>\nvoid TestOp(\n int num_items)\n{\n TestSize(num_items, hipcub::Sum());\n TestSize(num_items, hipcub::Max());\n}\n\n\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n int num_items = -1;\n int entropy_reduction = 0;\n int maxseg = 1000;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n args.GetCmdLineArgument(\"maxseg\", maxseg);\n args.GetCmdLineArgument(\"entropy\", entropy_reduction);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--maxseg=]\"\n \"[--entropy=]\"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n printf(\"\\n\");\n\n // Get ptx version\n int ptx_version = 0;\n CubDebugExit(PtxVersion(ptx_version));\n\n // %PARAM% TEST_CDP cdp 0:1\n\n // Test different input types\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n TestOp(num_items);\n\n return 0;\n}\n\n\n\n", "source": "/kaggle/working/cub/test/test_device_reduce_by_key.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n#include \n\n#include \n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void warp_combine_scan_kernel(T *in, T *inclusive_out, T *exclusive_out, ActionT action)\n{\n using warp_scan_t = cub::WarpScan;\n using storage_t = typename warp_scan_t::TempStorage;\n\n __shared__ storage_t storage[TOTAL_WARPS];\n\n const int tid = cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n\n // Get warp index\n int warp_id = tid / LOGICAL_WARP_THREADS;\n\n T inc_out, exc_out;\n T thread_data = in[tid];\n\n warp_scan_t scan(storage[warp_id]);\n\n action(scan, thread_data, inc_out, exc_out);\n\n inclusive_out[tid] = inc_out;\n exclusive_out[tid] = exc_out;\n}\n\ntemplate \nvoid warp_combine_scan(thrust::device_vector &in,\n thrust::device_vector &inclusive_out,\n thrust::device_vector &exclusive_out,\n ActionT action)\n{\n warp_combine_scan_kernel\n <<<1, LOGICAL_WARP_THREADS * TOTAL_WARPS>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(inclusive_out.data()),\n thrust::raw_pointer_cast(exclusive_out.data()),\n action);\n\n REQUIRE(cudaSuccess == cudaPeekAtLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n}\n\ntemplate \n__global__ void warp_scan_kernel(T *in, T *out, ActionT action)\n{\n using warp_scan_t = cub::WarpScan;\n using storage_t = typename warp_scan_t::TempStorage;\n\n __shared__ storage_t storage[TOTAL_WARPS];\n\n const int tid = cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n\n // Get warp index\n int warp_id = tid / LOGICAL_WARP_THREADS;\n\n T thread_data = in[tid];\n\n warp_scan_t scan(storage[warp_id]);\n\n action(scan, thread_data);\n\n out[tid] = thread_data;\n}\n\ntemplate \nvoid warp_scan(thrust::device_vector &in, thrust::device_vector &out, ActionT action)\n{\n warp_scan_kernel\n <<<1, LOGICAL_WARP_THREADS * TOTAL_WARPS>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(out.data()),\n action);\n\n REQUIRE(cudaSuccess == cudaPeekAtLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n}\n\nenum class scan_mode\n{\n exclusive,\n inclusive\n};\n\ntemplate \nstruct sum_op_t\n{\n template \n __device__ void operator()(WarpScanT &scan, T &thread_data) const\n {\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveSum(thread_data, thread_data);\n }\n else\n {\n scan.InclusiveSum(thread_data, thread_data);\n }\n }\n};\n\ntemplate \nstruct sum_aggregate_op_t\n{\n int m_target_thread_id;\n T *m_d_warp_aggregate;\n\n template \n __device__ void operator()(cub::WarpScan &scan, T &thread_data) const\n {\n T warp_aggregate{};\n\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveSum(thread_data, thread_data, warp_aggregate);\n }\n else\n {\n scan.InclusiveSum(thread_data, thread_data, warp_aggregate);\n }\n\n const int tid = cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n\n if (tid % LOGICAL_WARP_THREADS == m_target_thread_id)\n {\n m_d_warp_aggregate[tid / LOGICAL_WARP_THREADS] = warp_aggregate;\n }\n }\n};\n\ntemplate \nstruct min_op_t\n{\n template \n __device__ void operator()(WarpScanT &scan, T &thread_data) const\n {\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveScan(thread_data, thread_data, cub::Min{});\n }\n else\n {\n scan.InclusiveScan(thread_data, thread_data, cub::Min{});\n }\n }\n};\n\ntemplate \nstruct min_aggregate_op_t\n{\n int m_target_thread_id;\n T *m_d_warp_aggregate;\n\n template \n __device__ void operator()(cub::WarpScan &scan, T &thread_data) const\n {\n T warp_aggregate{};\n\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveScan(thread_data, thread_data, cub::Min{}, warp_aggregate);\n }\n else\n {\n scan.InclusiveScan(thread_data, thread_data, cub::Min{}, warp_aggregate);\n }\n\n const int tid = cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n\n if (tid % LOGICAL_WARP_THREADS == m_target_thread_id)\n {\n m_d_warp_aggregate[tid / LOGICAL_WARP_THREADS] = warp_aggregate;\n }\n }\n};\n\ntemplate \nstruct min_init_value_op_t\n{\n T initial_value;\n template \n __device__ void operator()(WarpScanT &scan, T &thread_data) const\n {\n scan.ExclusiveScan(thread_data, thread_data, initial_value, cub::Min{});\n }\n};\n\ntemplate \nstruct min_init_value_aggregate_op_t\n{\n int m_target_thread_id;\n T initial_value;\n T *m_d_warp_aggregate;\n\n template \n __device__ void operator()(cub::WarpScan &scan, T &thread_data) const\n {\n T warp_aggregate{};\n\n scan.ExclusiveScan(thread_data, thread_data, initial_value, cub::Min{}, warp_aggregate);\n\n const int tid = cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n\n if (tid % LOGICAL_WARP_THREADS == m_target_thread_id)\n {\n m_d_warp_aggregate[tid / LOGICAL_WARP_THREADS] = warp_aggregate;\n }\n }\n};\n\nstruct min_scan_op_t\n{\n template \n __device__ void\n operator()(WarpScanT &scan, T &thread_data, T &inclusive_output, T &exclusive_output) const\n {\n scan.Scan(thread_data, inclusive_output, exclusive_output, cub::Min{});\n }\n};\n\ntemplate \nstruct min_init_value_scan_op_t\n{\n T initial_value;\n template \n __device__ void\n operator()(WarpScanT &scan, T &thread_data, T &inclusive_output, T &exclusive_output) const\n {\n scan.Scan(thread_data, inclusive_output, exclusive_output, initial_value, cub::Min{});\n }\n};\n\ntemplate \nthrust::host_vector compute_host_reference(scan_mode mode,\n thrust::host_vector &result,\n int logical_warp_threads,\n ScanOpT scan_op,\n T initial_value = T{})\n{\n if (result.empty())\n {\n return thrust::host_vector{};\n }\n // TODO : assert result.size() % logical_warp_threads == 0\n\n int num_warps = CUB_QUOTIENT_CEILING(static_cast(result.size()), logical_warp_threads);\n thrust::host_vector warp_accumulator(num_warps);\n if (mode == scan_mode::exclusive)\n {\n for (int w = 0; w < num_warps; ++w)\n {\n T *output = result.data() + w * logical_warp_threads;\n T accumulator = output[0];\n T current = static_cast(scan_op(initial_value, output[0]));\n output[0] = initial_value;\n for (int i = 1; i < logical_warp_threads; i++)\n {\n accumulator = static_cast(scan_op(accumulator, output[i]));\n T tmp = output[i];\n output[i] = current;\n current = static_cast(scan_op(current, tmp));\n }\n warp_accumulator[w] = accumulator;\n }\n }\n else\n {\n for (int w = 0; w < num_warps; ++w)\n {\n T *output = result.data() + w * logical_warp_threads;\n T current = initial_value;\n for (int i = 0; i < logical_warp_threads; i++)\n {\n current = static_cast(scan_op(current, output[i]));\n output[i] = current;\n }\n warp_accumulator[w] = current;\n }\n }\n\n return warp_accumulator;\n}\n\nusing types = c2h::type_list;\nusing logical_warp_threads = c2h::enum_type_list;\nusing modes = c2h::enum_type_list;\n\nusing vec_types = c2h::type_list;\n\nusing warp_combine_type = int;\n\ntemplate \nstruct total_warps_t\n{\nprivate:\n static constexpr int max_warps = 2;\n static constexpr bool is_arch_warp = (logical_warp_threads == CUB_WARP_THREADS(0));\n static constexpr bool is_pow_of_two = ((logical_warp_threads & (logical_warp_threads - 1)) == 0);\n static constexpr int total_warps = (is_arch_warp || is_pow_of_two) ? max_warps : 1;\n\npublic:\n static constexpr int value() { return total_warps; }\n};\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int logical_warp_threads = c2h::get<1, TestType>::value;\n static constexpr scan_mode mode = c2h::get<2, TestType>::value;\n static constexpr int total_warps = total_warps_t::value();\n static constexpr int tile_size = total_warps * logical_warp_threads;\n};\n\nCUB_TEST(\"Warp scan works with sum\", \"[scan][warp]\", types, logical_warp_threads, modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n warp_scan(d_in,\n d_out,\n sum_op_t{});\n\n thrust::host_vector h_out = d_in;\n\n compute_host_reference(params::mode, h_out, params::logical_warp_threads, std::plus{});\n REQUIRE_APPROX_EQ(h_out, d_out);\n}\n\nCUB_TEST(\"Warp scan works with vec_types\", \"[scan][warp]\", vec_types, logical_warp_threads, modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n warp_scan(d_in,\n d_out,\n sum_op_t{});\n\n thrust::host_vector h_out = d_in;\n\n compute_host_reference(params::mode, h_out, params::logical_warp_threads, std::plus{});\n REQUIRE(h_out == d_out);\n}\n\nCUB_TEST(\"Warp scan works with custom types\",\n \"[scan][warp]\",\n c2h::type_list>,\n logical_warp_threads,\n modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n warp_scan(d_in,\n d_out,\n sum_op_t{});\n\n thrust::host_vector h_out = d_in;\n\n compute_host_reference(params::mode, h_out, params::logical_warp_threads, std::plus{});\n REQUIRE(h_out == d_out);\n}\n\nCUB_TEST(\"Warp scan returns valid warp aggregate\",\n \"[scan][warp]\",\n c2h::type_list>,\n logical_warp_threads,\n modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_warp_aggregates(params::total_warps);\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n const int target_thread_id = GENERATE_COPY(take(2, random(0, params::logical_warp_threads - 1)));\n\n warp_scan(\n d_in,\n d_out,\n sum_aggregate_op_t{target_thread_id,\n thrust::raw_pointer_cast(d_warp_aggregates.data())});\n\n thrust::host_vector h_out = d_in;\n\n auto h_warp_aggregates =\n compute_host_reference(params::mode, h_out, params::logical_warp_threads, std::plus{});\n REQUIRE(h_out == d_out);\n REQUIRE(h_warp_aggregates == d_warp_aggregates);\n}\n\n// TODO : Do we need all the types?\nCUB_TEST(\"Warp scan works with custom scan op\", \"[scan][warp]\", types, logical_warp_threads, modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n warp_scan(d_in,\n d_out,\n min_op_t{});\n\n thrust::host_vector h_out = d_in;\n\n compute_host_reference(\n params::mode,\n h_out,\n params::logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n std::numeric_limits::max());\n\n // From the documentation -\n // Computes an exclusive prefix scan using the specified binary scan functor\n // across the calling warp. Because no initial value is supplied, the output\n // computed for warp-lane0 is undefined.\n\n // When comparing device output, the corresponding undefined data points need\n // to be fixed\n\n if (params::mode == scan_mode::exclusive)\n {\n for (size_t i = 0; i < h_out.size(); i += params::logical_warp_threads)\n {\n d_out[i] = h_out[i];\n }\n }\n REQUIRE_APPROX_EQ(h_out, d_out);\n}\n\nCUB_TEST(\"Warp custom op scan returns valid warp aggregate\",\n \"[scan][warp]\",\n types,\n logical_warp_threads,\n modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_warp_aggregates(params::total_warps);\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n const int target_thread_id = GENERATE_COPY(take(2, random(0, params::logical_warp_threads - 1)));\n\n warp_scan(\n d_in,\n d_out,\n min_aggregate_op_t{target_thread_id,\n thrust::raw_pointer_cast(d_warp_aggregates.data())});\n\n thrust::host_vector h_out = d_in;\n\n auto h_warp_aggregates = compute_host_reference(\n params::mode,\n h_out,\n params::logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n std::numeric_limits::max());\n\n // From the documentation -\n // Computes an exclusive prefix scan using the specified binary scan functor\n // across the calling warp. Because no initial value is supplied, the output\n // computed for warp-lane0 is undefined.\n\n // When comparing device output, the corresponding undefined data points need\n // to be fixed\n\n if (params::mode == scan_mode::exclusive)\n {\n for (size_t i = 0; i < h_out.size(); i += params::logical_warp_threads)\n {\n d_out[i] = h_out[i];\n }\n }\n REQUIRE(h_out == d_out);\n REQUIRE(h_warp_aggregates == d_warp_aggregates);\n}\n\nCUB_TEST(\"Warp custom op scan works with initial value\",\n \"[scan][warp]\",\n types,\n logical_warp_threads,\n c2h::enum_type_list)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n const type initial_value = static_cast(GENERATE_COPY(take(2, random(0, params::tile_size))));\n\n warp_scan(d_in,\n d_out,\n min_init_value_op_t{\n initial_value});\n\n thrust::host_vector h_out = d_in;\n\n compute_host_reference(\n params::mode,\n h_out,\n params::logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n initial_value);\n\n REQUIRE_APPROX_EQ(h_out, d_out);\n}\n\nCUB_TEST(\"Warp custom op scan with initial value returns valid warp aggregate\",\n \"[scan][warp]\",\n types,\n logical_warp_threads,\n c2h::enum_type_list)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_warp_aggregates(params::total_warps);\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n const int target_thread_id = GENERATE_COPY(take(2, random(0, params::logical_warp_threads - 1)));\n const type initial_value = static_cast(GENERATE_COPY(take(2, random(0, params::tile_size))));\n\n warp_scan(\n d_in,\n d_out,\n min_init_value_aggregate_op_t{target_thread_id,\n initial_value,\n thrust::raw_pointer_cast(d_warp_aggregates.data())});\n\n thrust::host_vector h_out = d_in;\n\n auto h_warp_aggregates = compute_host_reference(\n params::mode,\n h_out,\n params::logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n initial_value);\n\n REQUIRE(h_out == d_out);\n REQUIRE(h_warp_aggregates == d_warp_aggregates);\n}\n\nCUB_TEST(\"Warp combination scan works with custom scan op\", \"[scan][warp]\", logical_warp_threads)\n{\n constexpr int logical_warp_threads = c2h::get<0, TestType>();\n constexpr int total_warps = total_warps_t::value();\n using type = int;\n\n thrust::device_vector d_inclusive_out(total_warps * logical_warp_threads);\n thrust::device_vector d_exclusive_out(total_warps * logical_warp_threads);\n thrust::device_vector d_in(total_warps * logical_warp_threads);\n c2h::gen(CUB_SEED(10), d_in);\n\n warp_combine_scan(d_in,\n d_inclusive_out,\n d_exclusive_out,\n min_scan_op_t{});\n\n thrust::host_vector h_exclusive_out = d_in;\n thrust::host_vector h_inclusive_out = d_in;\n\n compute_host_reference(\n scan_mode::exclusive,\n h_exclusive_out,\n logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n std::numeric_limits::max());\n\n compute_host_reference(\n scan_mode::inclusive,\n h_inclusive_out,\n logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n std::numeric_limits::max());\n\n // According to WarpScan::Scan documentation -\n // Because no initial value is supplied, the exclusive_output computed for warp-lane0 is\n // undefined.\n\n // When comparing device output, the corresponding undefined data points need\n // to be fixed\n\n for (size_t i = 0; i < h_exclusive_out.size(); i += logical_warp_threads)\n {\n d_exclusive_out[i] = h_exclusive_out[i];\n }\n\n REQUIRE(h_inclusive_out == d_inclusive_out);\n REQUIRE(h_exclusive_out == d_exclusive_out);\n}\n\nCUB_TEST(\"Warp combination custom scan works with initial value\",\n \"[scan][warp]\",\n logical_warp_threads)\n{\n constexpr int logical_warp_threads = c2h::get<0, TestType>();\n constexpr int total_warps = total_warps_t::value();\n using type = int;\n\n thrust::device_vector d_inclusive_out(total_warps * logical_warp_threads);\n thrust::device_vector d_exclusive_out(total_warps * logical_warp_threads);\n thrust::device_vector d_in(total_warps * logical_warp_threads);\n c2h::gen(CUB_SEED(10), d_in);\n\n const type initial_value = GENERATE_COPY(take(2, random(0, total_warps * logical_warp_threads)));\n\n warp_combine_scan(d_in,\n d_inclusive_out,\n d_exclusive_out,\n min_init_value_scan_op_t{\n initial_value});\n\n thrust::host_vector h_exclusive_out = d_in;\n thrust::host_vector h_inclusive_out = d_in;\n\n compute_host_reference(\n scan_mode::exclusive,\n h_exclusive_out,\n logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n initial_value);\n\n compute_host_reference(\n scan_mode::inclusive,\n h_inclusive_out,\n logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n initial_value);\n\n REQUIRE(h_inclusive_out == d_inclusive_out);\n REQUIRE(h_exclusive_out == d_exclusive_out);\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n#include \n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void warp_combine_scan_kernel(T *in, T *inclusive_out, T *exclusive_out, ActionT action)\n{\n using warp_scan_t = hipcub::WarpScan;\n using storage_t = typename warp_scan_t::TempStorage;\n\n __shared__ storage_t storage[TOTAL_WARPS];\n\n const int tid = hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n\n // Get warp index\n int warp_id = tid / LOGICAL_WARP_THREADS;\n\n T inc_out, exc_out;\n T thread_data = in[tid];\n\n warp_scan_t scan(storage[warp_id]);\n\n action(scan, thread_data, inc_out, exc_out);\n\n inclusive_out[tid] = inc_out;\n exclusive_out[tid] = exc_out;\n}\n\ntemplate \nvoid warp_combine_scan(thrust::device_vector &in,\n thrust::device_vector &inclusive_out,\n thrust::device_vector &exclusive_out,\n ActionT action)\n{\n warp_combine_scan_kernel\n <<<1, LOGICAL_WARP_THREADS * TOTAL_WARPS>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(inclusive_out.data()),\n thrust::raw_pointer_cast(exclusive_out.data()),\n action);\n\n REQUIRE(hipSuccess == hipPeekAtLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n}\n\ntemplate \n__global__ void warp_scan_kernel(T *in, T *out, ActionT action)\n{\n using warp_scan_t = hipcub::WarpScan;\n using storage_t = typename warp_scan_t::TempStorage;\n\n __shared__ storage_t storage[TOTAL_WARPS];\n\n const int tid = hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n\n // Get warp index\n int warp_id = tid / LOGICAL_WARP_THREADS;\n\n T thread_data = in[tid];\n\n warp_scan_t scan(storage[warp_id]);\n\n action(scan, thread_data);\n\n out[tid] = thread_data;\n}\n\ntemplate \nvoid warp_scan(thrust::device_vector &in, thrust::device_vector &out, ActionT action)\n{\n warp_scan_kernel\n <<<1, LOGICAL_WARP_THREADS * TOTAL_WARPS>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(out.data()),\n action);\n\n REQUIRE(hipSuccess == hipPeekAtLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n}\n\nenum class scan_mode\n{\n exclusive,\n inclusive\n};\n\ntemplate \nstruct sum_op_t\n{\n template \n __device__ void operator()(WarpScanT &scan, T &thread_data) const\n {\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveSum(thread_data, thread_data);\n }\n else\n {\n scan.InclusiveSum(thread_data, thread_data);\n }\n }\n};\n\ntemplate \nstruct sum_aggregate_op_t\n{\n int m_target_thread_id;\n T *m_d_warp_aggregate;\n\n template \n __device__ void operator()(hipcub::WarpScan &scan, T &thread_data) const\n {\n T warp_aggregate{};\n\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveSum(thread_data, thread_data, warp_aggregate);\n }\n else\n {\n scan.InclusiveSum(thread_data, thread_data, warp_aggregate);\n }\n\n const int tid = hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n\n if (tid % LOGICAL_WARP_THREADS == m_target_thread_id)\n {\n m_d_warp_aggregate[tid / LOGICAL_WARP_THREADS] = warp_aggregate;\n }\n }\n};\n\ntemplate \nstruct min_op_t\n{\n template \n __device__ void operator()(WarpScanT &scan, T &thread_data) const\n {\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveScan(thread_data, thread_data, hipcub::Min{});\n }\n else\n {\n scan.InclusiveScan(thread_data, thread_data, hipcub::Min{});\n }\n }\n};\n\ntemplate \nstruct min_aggregate_op_t\n{\n int m_target_thread_id;\n T *m_d_warp_aggregate;\n\n template \n __device__ void operator()(hipcub::WarpScan &scan, T &thread_data) const\n {\n T warp_aggregate{};\n\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveScan(thread_data, thread_data, hipcub::Min{}, warp_aggregate);\n }\n else\n {\n scan.InclusiveScan(thread_data, thread_data, hipcub::Min{}, warp_aggregate);\n }\n\n const int tid = hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n\n if (tid % LOGICAL_WARP_THREADS == m_target_thread_id)\n {\n m_d_warp_aggregate[tid / LOGICAL_WARP_THREADS] = warp_aggregate;\n }\n }\n};\n\ntemplate \nstruct min_init_value_op_t\n{\n T initial_value;\n template \n __device__ void operator()(WarpScanT &scan, T &thread_data) const\n {\n scan.ExclusiveScan(thread_data, thread_data, initial_value, hipcub::Min{});\n }\n};\n\ntemplate \nstruct min_init_value_aggregate_op_t\n{\n int m_target_thread_id;\n T initial_value;\n T *m_d_warp_aggregate;\n\n template \n __device__ void operator()(hipcub::WarpScan &scan, T &thread_data) const\n {\n T warp_aggregate{};\n\n scan.ExclusiveScan(thread_data, thread_data, initial_value, hipcub::Min{}, warp_aggregate);\n\n const int tid = hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n\n if (tid % LOGICAL_WARP_THREADS == m_target_thread_id)\n {\n m_d_warp_aggregate[tid / LOGICAL_WARP_THREADS] = warp_aggregate;\n }\n }\n};\n\nstruct min_scan_op_t\n{\n template \n __device__ void\n operator()(WarpScanT &scan, T &thread_data, T &inclusive_output, T &exclusive_output) const\n {\n scan.Scan(thread_data, inclusive_output, exclusive_output, hipcub::Min{});\n }\n};\n\ntemplate \nstruct min_init_value_scan_op_t\n{\n T initial_value;\n template \n __device__ void\n operator()(WarpScanT &scan, T &thread_data, T &inclusive_output, T &exclusive_output) const\n {\n scan.Scan(thread_data, inclusive_output, exclusive_output, initial_value, hipcub::Min{});\n }\n};\n\ntemplate \nthrust::host_vector compute_host_reference(scan_mode mode,\n thrust::host_vector &result,\n int logical_warp_threads,\n ScanOpT scan_op,\n T initial_value = T{})\n{\n if (result.empty())\n {\n return thrust::host_vector{};\n }\n // TODO : assert result.size() % logical_warp_threads == 0\n\n int num_warps = CUB_QUOTIENT_CEILING(static_cast(result.size()), logical_warp_threads);\n thrust::host_vector warp_accumulator(num_warps);\n if (mode == scan_mode::exclusive)\n {\n for (int w = 0; w < num_warps; ++w)\n {\n T *output = result.data() + w * logical_warp_threads;\n T accumulator = output[0];\n T current = static_cast(scan_op(initial_value, output[0]));\n output[0] = initial_value;\n for (int i = 1; i < logical_warp_threads; i++)\n {\n accumulator = static_cast(scan_op(accumulator, output[i]));\n T tmp = output[i];\n output[i] = current;\n current = static_cast(scan_op(current, tmp));\n }\n warp_accumulator[w] = accumulator;\n }\n }\n else\n {\n for (int w = 0; w < num_warps; ++w)\n {\n T *output = result.data() + w * logical_warp_threads;\n T current = initial_value;\n for (int i = 0; i < logical_warp_threads; i++)\n {\n current = static_cast(scan_op(current, output[i]));\n output[i] = current;\n }\n warp_accumulator[w] = current;\n }\n }\n\n return warp_accumulator;\n}\n\nusing types = c2h::type_list;\nusing logical_warp_threads = c2h::enum_type_list;\nusing modes = c2h::enum_type_list;\n\nusing vec_types = c2h::type_list;\n\nusing warp_combine_type = int;\n\ntemplate \nstruct total_warps_t\n{\nprivate:\n static constexpr int max_warps = 2;\n static constexpr bool is_arch_warp = (logical_warp_threads == CUB_WARP_THREADS(0));\n static constexpr bool is_pow_of_two = ((logical_warp_threads & (logical_warp_threads - 1)) == 0);\n static constexpr int total_warps = (is_arch_warp || is_pow_of_two) ? max_warps : 1;\n\npublic:\n static constexpr int value() { return total_warps; }\n};\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int logical_warp_threads = c2h::get<1, TestType>::value;\n static constexpr scan_mode mode = c2h::get<2, TestType>::value;\n static constexpr int total_warps = total_warps_t::value();\n static constexpr int tile_size = total_warps * logical_warp_threads;\n};\n\nCUB_TEST(\"Warp scan works with sum\", \"[scan][warp]\", types, logical_warp_threads, modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n warp_scan(d_in,\n d_out,\n sum_op_t{});\n\n thrust::host_vector h_out = d_in;\n\n compute_host_reference(params::mode, h_out, params::logical_warp_threads, std::plus{});\n REQUIRE_APPROX_EQ(h_out, d_out);\n}\n\nCUB_TEST(\"Warp scan works with vec_types\", \"[scan][warp]\", vec_types, logical_warp_threads, modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n warp_scan(d_in,\n d_out,\n sum_op_t{});\n\n thrust::host_vector h_out = d_in;\n\n compute_host_reference(params::mode, h_out, params::logical_warp_threads, std::plus{});\n REQUIRE(h_out == d_out);\n}\n\nCUB_TEST(\"Warp scan works with custom types\",\n \"[scan][warp]\",\n c2h::type_list>,\n logical_warp_threads,\n modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n warp_scan(d_in,\n d_out,\n sum_op_t{});\n\n thrust::host_vector h_out = d_in;\n\n compute_host_reference(params::mode, h_out, params::logical_warp_threads, std::plus{});\n REQUIRE(h_out == d_out);\n}\n\nCUB_TEST(\"Warp scan returns valid warp aggregate\",\n \"[scan][warp]\",\n c2h::type_list>,\n logical_warp_threads,\n modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_warp_aggregates(params::total_warps);\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n const int target_thread_id = GENERATE_COPY(take(2, random(0, params::logical_warp_threads - 1)));\n\n warp_scan(\n d_in,\n d_out,\n sum_aggregate_op_t{target_thread_id,\n thrust::raw_pointer_cast(d_warp_aggregates.data())});\n\n thrust::host_vector h_out = d_in;\n\n auto h_warp_aggregates =\n compute_host_reference(params::mode, h_out, params::logical_warp_threads, std::plus{});\n REQUIRE(h_out == d_out);\n REQUIRE(h_warp_aggregates == d_warp_aggregates);\n}\n\n// TODO : Do we need all the types?\nCUB_TEST(\"Warp scan works with custom scan op\", \"[scan][warp]\", types, logical_warp_threads, modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n warp_scan(d_in,\n d_out,\n min_op_t{});\n\n thrust::host_vector h_out = d_in;\n\n compute_host_reference(\n params::mode,\n h_out,\n params::logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n std::numeric_limits::max());\n\n // From the documentation -\n // Computes an exclusive prefix scan using the specified binary scan functor\n // across the calling warp. Because no initial value is supplied, the output\n // computed for warp-lane0 is undefined.\n\n // When comparing device output, the corresponding undefined data points need\n // to be fixed\n\n if (params::mode == scan_mode::exclusive)\n {\n for (size_t i = 0; i < h_out.size(); i += params::logical_warp_threads)\n {\n d_out[i] = h_out[i];\n }\n }\n REQUIRE_APPROX_EQ(h_out, d_out);\n}\n\nCUB_TEST(\"Warp custom op scan returns valid warp aggregate\",\n \"[scan][warp]\",\n types,\n logical_warp_threads,\n modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_warp_aggregates(params::total_warps);\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n const int target_thread_id = GENERATE_COPY(take(2, random(0, params::logical_warp_threads - 1)));\n\n warp_scan(\n d_in,\n d_out,\n min_aggregate_op_t{target_thread_id,\n thrust::raw_pointer_cast(d_warp_aggregates.data())});\n\n thrust::host_vector h_out = d_in;\n\n auto h_warp_aggregates = compute_host_reference(\n params::mode,\n h_out,\n params::logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n std::numeric_limits::max());\n\n // From the documentation -\n // Computes an exclusive prefix scan using the specified binary scan functor\n // across the calling warp. Because no initial value is supplied, the output\n // computed for warp-lane0 is undefined.\n\n // When comparing device output, the corresponding undefined data points need\n // to be fixed\n\n if (params::mode == scan_mode::exclusive)\n {\n for (size_t i = 0; i < h_out.size(); i += params::logical_warp_threads)\n {\n d_out[i] = h_out[i];\n }\n }\n REQUIRE(h_out == d_out);\n REQUIRE(h_warp_aggregates == d_warp_aggregates);\n}\n\nCUB_TEST(\"Warp custom op scan works with initial value\",\n \"[scan][warp]\",\n types,\n logical_warp_threads,\n c2h::enum_type_list)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n const type initial_value = static_cast(GENERATE_COPY(take(2, random(0, params::tile_size))));\n\n warp_scan(d_in,\n d_out,\n min_init_value_op_t{\n initial_value});\n\n thrust::host_vector h_out = d_in;\n\n compute_host_reference(\n params::mode,\n h_out,\n params::logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n initial_value);\n\n REQUIRE_APPROX_EQ(h_out, d_out);\n}\n\nCUB_TEST(\"Warp custom op scan with initial value returns valid warp aggregate\",\n \"[scan][warp]\",\n types,\n logical_warp_threads,\n c2h::enum_type_list)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_warp_aggregates(params::total_warps);\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n const int target_thread_id = GENERATE_COPY(take(2, random(0, params::logical_warp_threads - 1)));\n const type initial_value = static_cast(GENERATE_COPY(take(2, random(0, params::tile_size))));\n\n warp_scan(\n d_in,\n d_out,\n min_init_value_aggregate_op_t{target_thread_id,\n initial_value,\n thrust::raw_pointer_cast(d_warp_aggregates.data())});\n\n thrust::host_vector h_out = d_in;\n\n auto h_warp_aggregates = compute_host_reference(\n params::mode,\n h_out,\n params::logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n initial_value);\n\n REQUIRE(h_out == d_out);\n REQUIRE(h_warp_aggregates == d_warp_aggregates);\n}\n\nCUB_TEST(\"Warp combination scan works with custom scan op\", \"[scan][warp]\", logical_warp_threads)\n{\n constexpr int logical_warp_threads = c2h::get<0, TestType>();\n constexpr int total_warps = total_warps_t::value();\n using type = int;\n\n thrust::device_vector d_inclusive_out(total_warps * logical_warp_threads);\n thrust::device_vector d_exclusive_out(total_warps * logical_warp_threads);\n thrust::device_vector d_in(total_warps * logical_warp_threads);\n c2h::gen(CUB_SEED(10), d_in);\n\n warp_combine_scan(d_in,\n d_inclusive_out,\n d_exclusive_out,\n min_scan_op_t{});\n\n thrust::host_vector h_exclusive_out = d_in;\n thrust::host_vector h_inclusive_out = d_in;\n\n compute_host_reference(\n scan_mode::exclusive,\n h_exclusive_out,\n logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n std::numeric_limits::max());\n\n compute_host_reference(\n scan_mode::inclusive,\n h_inclusive_out,\n logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n std::numeric_limits::max());\n\n // According to WarpScan::Scan documentation -\n // Because no initial value is supplied, the exclusive_output computed for warp-lane0 is\n // undefined.\n\n // When comparing device output, the corresponding undefined data points need\n // to be fixed\n\n for (size_t i = 0; i < h_exclusive_out.size(); i += logical_warp_threads)\n {\n d_exclusive_out[i] = h_exclusive_out[i];\n }\n\n REQUIRE(h_inclusive_out == d_inclusive_out);\n REQUIRE(h_exclusive_out == d_exclusive_out);\n}\n\nCUB_TEST(\"Warp combination custom scan works with initial value\",\n \"[scan][warp]\",\n logical_warp_threads)\n{\n constexpr int logical_warp_threads = c2h::get<0, TestType>();\n constexpr int total_warps = total_warps_t::value();\n using type = int;\n\n thrust::device_vector d_inclusive_out(total_warps * logical_warp_threads);\n thrust::device_vector d_exclusive_out(total_warps * logical_warp_threads);\n thrust::device_vector d_in(total_warps * logical_warp_threads);\n c2h::gen(CUB_SEED(10), d_in);\n\n const type initial_value = GENERATE_COPY(take(2, random(0, total_warps * logical_warp_threads)));\n\n warp_combine_scan(d_in,\n d_inclusive_out,\n d_exclusive_out,\n min_init_value_scan_op_t{\n initial_value});\n\n thrust::host_vector h_exclusive_out = d_in;\n thrust::host_vector h_inclusive_out = d_in;\n\n compute_host_reference(\n scan_mode::exclusive,\n h_exclusive_out,\n logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n initial_value);\n\n compute_host_reference(\n scan_mode::inclusive,\n h_inclusive_out,\n logical_warp_threads,\n [](type l, type r) { return std::min(l, r); },\n initial_value);\n\n REQUIRE(h_inclusive_out == d_inclusive_out);\n REQUIRE(h_exclusive_out == d_exclusive_out);\n}\n", "source": "/kaggle/working/cub/test/catch2_test_warp_scan.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of BlockMergeSort utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n#include \n\n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\nstruct CustomLess\n{\n template \n __device__ __host__ bool operator()(const DataType &lhs, const DataType &rhs)\n {\n return lhs < rhs;\n }\n};\n\ntemplate \n__global__ void block_merge_sort_kernel(\n KeyT *data, \n int valid_items,\n KeyT oob_default,\n ActionT action)\n{\n using BlockMergeSort =\n cub::BlockMergeSort;\n\n __shared__ typename BlockMergeSort::TempStorage temp_storage_shuffle;\n\n KeyT thread_data[ItemsPerThread];\n\n const int thread_offset = static_cast(threadIdx.x) * ItemsPerThread;\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n const int idx = thread_offset + item;\n thread_data[item] = idx < valid_items ? data[idx] : KeyT();\n }\n __syncthreads();\n\n BlockMergeSort sort(temp_storage_shuffle);\n\n action(sort, thread_data, valid_items, oob_default);\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n const int idx = thread_offset + item;\n\n if (idx >= valid_items)\n {\n break;\n }\n\n data[idx] = thread_data[item];\n }\n}\n\ntemplate \n__global__ void block_merge_sort_kernel(\n KeyT *keys, \n ValueT *vals, \n int valid_items,\n KeyT oob_default,\n ActionT action)\n{\n using BlockMergeSort =\n cub::BlockMergeSort;\n\n __shared__ typename BlockMergeSort::TempStorage temp_storage_shuffle;\n\n KeyT thread_keys[ItemsPerThread];\n ValueT thread_vals[ItemsPerThread];\n\n const int thread_offset = static_cast(threadIdx.x) * ItemsPerThread;\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n const int idx = thread_offset + item;\n thread_keys[item] = idx < valid_items ? keys[idx] : KeyT{};\n thread_vals[item] = idx < valid_items ? vals[idx] : ValueT{};\n }\n __syncthreads();\n\n BlockMergeSort sort(temp_storage_shuffle);\n\n action(sort, thread_keys, thread_vals, valid_items, oob_default);\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n const int idx = thread_offset + item;\n\n if (idx >= valid_items)\n {\n break;\n }\n\n keys[idx] = thread_keys[item];\n vals[idx] = thread_vals[item];\n }\n}\n\nstruct stable_sort_keys_partial_tile_t\n{\n template \n __device__ void operator()(BlockMergeSortT &sort,\n KeyT &thread_data,\n int valid_items,\n DefaultT oob_default) const\n {\n sort.StableSort(thread_data, CustomLess{}, valid_items, oob_default);\n }\n};\n\nstruct stable_sort_pairs_partial_tile_t\n{\n template \n __device__ void operator()(BlockMergeSortT &sort,\n KeyT &thread_keys,\n ValueT &thread_vals,\n int valid_items,\n DefaultT oob_default) const\n {\n sort.StableSort(thread_keys,\n thread_vals,\n CustomLess{},\n valid_items,\n oob_default);\n }\n};\n\nstruct stable_sort_pairs_full_tile_t\n{\n template \n __device__ void operator()(BlockMergeSortT &sort,\n KeyT &thread_keys,\n ValueT &thread_vals,\n int /* valid_items */,\n DefaultT /* oob_default */) const\n {\n sort.StableSort(thread_keys, thread_vals, CustomLess());\n }\n};\n\nstruct stable_sort_keys_full_tile_t\n{\n template \n __device__ void operator()(BlockMergeSortT &sort,\n KeyT &thread_keys,\n int /* valid_items */,\n DefaultT /* oob_default */) const\n {\n sort.StableSort(thread_keys, CustomLess());\n }\n};\n\ntemplate \nvoid block_merge_sort(\n thrust::device_vector &keys, \n ActionT action)\n{\n block_merge_sort_kernel\n <<<1, ThreadsInBlock>>>(\n thrust::raw_pointer_cast(keys.data()),\n static_cast(keys.size()),\n std::numeric_limits::max(),\n action);\n\n REQUIRE( cudaSuccess == cudaPeekAtLastError() );\n REQUIRE( cudaSuccess == cudaDeviceSynchronize() );\n}\n\ntemplate \nvoid block_merge_sort(\n thrust::device_vector &keys, \n thrust::device_vector &vals, \n ActionT action)\n{\n block_merge_sort_kernel\n <<<1, ThreadsInBlock>>>(\n thrust::raw_pointer_cast(keys.data()),\n thrust::raw_pointer_cast(vals.data()),\n static_cast(keys.size()),\n std::numeric_limits::max(),\n action);\n\n REQUIRE( cudaSuccess == cudaPeekAtLastError() );\n REQUIRE( cudaSuccess == cudaDeviceSynchronize() );\n}\n\n// %PARAM% THREADS_IN_BLOCK bs 64:256\n\nusing key_types = c2h::type_list;\nusing threads_in_block = c2h::enum_type_list;\nusing items_per_thread = c2h::enum_type_list;\n\n\ntemplate \nstruct params_t\n{\n using key_t = typename c2h::get<0, TestType>;\n\n static constexpr int items_per_thread = c2h::get<1, TestType>::value;\n static constexpr int threads_in_block = c2h::get<2, TestType>::value;\n static constexpr int tile_size = items_per_thread * threads_in_block;\n};\n\n\nCUB_TEST(\"Block merge sort can sort keys in partial tiles\",\n \"[merge sort][block]\",\n key_types,\n items_per_thread,\n threads_in_block)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n\n thrust::device_vector d_keys(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n\n c2h::gen(CUB_SEED(10), d_keys);\n\n thrust::host_vector h_reference = d_keys;\n std::stable_sort(\n thrust::raw_pointer_cast(h_reference.data()), \n thrust::raw_pointer_cast(h_reference.data()) + h_reference.size(), \n CustomLess{});\n\n block_merge_sort(\n d_keys,\n stable_sort_keys_partial_tile_t{});\n\n REQUIRE( h_reference == d_keys );\n}\n\nCUB_TEST(\"Block merge sort can sort keys in full tiles\",\n \"[merge sort][block]\",\n key_types,\n items_per_thread,\n threads_in_block)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n\n thrust::device_vector d_keys(params::tile_size);\n\n c2h::gen(CUB_SEED(10), d_keys);\n\n thrust::host_vector h_reference = d_keys;\n std::stable_sort(\n thrust::raw_pointer_cast(h_reference.data()), \n thrust::raw_pointer_cast(h_reference.data()) + h_reference.size(), \n CustomLess{});\n\n block_merge_sort(\n d_keys,\n stable_sort_keys_full_tile_t{});\n\n REQUIRE( h_reference == d_keys );\n}\n\nCUB_TEST(\"Block merge sort can sort pairs in partial tiles\",\n \"[merge sort][block]\",\n key_types,\n items_per_thread,\n threads_in_block)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n using value_t = key_t;\n using pair_t = std::pair;\n\n thrust::device_vector d_keys(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n thrust::device_vector d_vals(d_keys.size());\n\n c2h::gen(CUB_SEED(5), d_keys);\n c2h::gen(CUB_SEED(5), d_vals);\n\n thrust::host_vector h_keys = d_keys;\n thrust::host_vector h_vals = d_vals;\n\n thrust::host_vector h_ref(d_keys.size());\n\n for (std::size_t idx = 0; idx < h_ref.size(); idx++)\n {\n h_ref[idx] = std::make_pair(h_keys[idx], h_vals[idx]);\n }\n\n std::stable_sort(\n thrust::raw_pointer_cast(h_ref.data()), \n thrust::raw_pointer_cast(h_ref.data()) + h_ref.size(), \n [](pair_t l, pair_t r) -> bool {\n return l.first < r.first;\n });\n\n for (std::size_t idx = 0; idx < h_ref.size(); idx++)\n {\n h_keys[idx] = h_ref[idx].first;\n h_vals[idx] = h_ref[idx].second;\n }\n\n block_merge_sort(\n d_keys,\n d_vals,\n stable_sort_pairs_partial_tile_t{});\n\n REQUIRE( h_keys == d_keys );\n REQUIRE( h_vals == d_vals );\n}\n\nCUB_TEST(\"Block merge sort can sort pairs in full tiles\",\n \"[merge sort][block]\",\n key_types,\n items_per_thread,\n threads_in_block)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n using value_t = key_t;\n using pair_t = std::pair;\n\n thrust::device_vector d_keys(params::tile_size);\n thrust::device_vector d_vals(d_keys.size());\n\n c2h::gen(CUB_SEED(5), d_keys);\n c2h::gen(CUB_SEED(5), d_vals);\n\n thrust::host_vector h_keys = d_keys;\n thrust::host_vector h_vals = d_vals;\n\n thrust::host_vector h_ref(d_keys.size());\n\n for (std::size_t idx = 0; idx < h_ref.size(); idx++)\n {\n h_ref[idx] = std::make_pair(h_keys[idx], h_vals[idx]);\n }\n\n std::stable_sort(\n thrust::raw_pointer_cast(h_ref.data()), \n thrust::raw_pointer_cast(h_ref.data()) + h_ref.size(), \n [](pair_t l, pair_t r) -> bool {\n return l.first < r.first;\n });\n\n for (std::size_t idx = 0; idx < h_ref.size(); idx++)\n {\n h_keys[idx] = h_ref[idx].first;\n h_vals[idx] = h_ref[idx].second;\n }\n\n block_merge_sort(\n d_keys,\n d_vals,\n stable_sort_pairs_full_tile_t{});\n\n REQUIRE( h_keys == d_keys );\n REQUIRE( h_vals == d_vals );\n}\n\n\nCUB_TEST(\"Block merge sort can sort pairs with mixed types\",\n \"[merge sort][block]\",\n threads_in_block)\n{\n using key_t = std::int32_t;\n using value_t = std::int64_t;\n using pair_t = std::pair;\n\n constexpr int items_per_thread = 2;\n constexpr int threads_in_block = c2h::get<0, TestType>::value;\n constexpr int tile_size = items_per_thread * threads_in_block;\n\n thrust::device_vector d_keys(tile_size);\n thrust::device_vector d_vals(d_keys.size());\n\n c2h::gen(CUB_SEED(5), d_keys);\n c2h::gen(CUB_SEED(5), d_vals);\n\n thrust::host_vector h_keys = d_keys;\n thrust::host_vector h_vals = d_vals;\n\n thrust::host_vector h_ref(d_keys.size());\n\n for (std::size_t idx = 0; idx < h_ref.size(); idx++)\n {\n h_ref[idx] = std::make_pair(h_keys[idx], h_vals[idx]);\n }\n\n std::stable_sort(\n thrust::raw_pointer_cast(h_ref.data()), \n thrust::raw_pointer_cast(h_ref.data()) + h_ref.size(), \n [](pair_t l, pair_t r) -> bool {\n return l.first < r.first;\n });\n\n for (std::size_t idx = 0; idx < h_ref.size(); idx++)\n {\n h_keys[idx] = h_ref[idx].first;\n h_vals[idx] = h_ref[idx].second;\n }\n\n block_merge_sort(\n d_keys,\n d_vals,\n stable_sort_pairs_full_tile_t{});\n\n REQUIRE( h_keys == d_keys );\n REQUIRE( h_vals == d_vals );\n}\n\nCUB_TEST(\"Block merge sort can sort large tiles\",\n \"[merge sort][block]\",\n threads_in_block)\n{\n using key_t = std::uint16_t;\n\n constexpr int items_per_thread = 2;\n\n // Repurpose block sizes\n constexpr int cmake_threads_in_block = c2h::get<0, TestType>::value;\n constexpr int threads_in_block = cmake_threads_in_block < 256\n ? 512\n : 1024;\n\n constexpr int tile_size = threads_in_block * items_per_thread;\n\n thrust::device_vector d_keys(tile_size);\n c2h::gen(CUB_SEED(10), d_keys);\n\n thrust::host_vector h_reference = d_keys;\n std::stable_sort(\n thrust::raw_pointer_cast(h_reference.data()), \n thrust::raw_pointer_cast(h_reference.data()) + h_reference.size(), \n CustomLess{});\n\n block_merge_sort(\n d_keys,\n stable_sort_keys_full_tile_t{});\n\n REQUIRE( h_reference == d_keys );\n}\n\nCUB_TEST(\"Block merge sort is stable\",\n \"[merge sort][block]\",\n threads_in_block)\n{\n using key_t = c2h::custom_type_t;\n\n constexpr int items_per_thread = 2;\n constexpr int threads_in_block = c2h::get<0, TestType>::value;\n constexpr int tile_size = threads_in_block * items_per_thread;\n\n thrust::device_vector d_keys(tile_size);\n c2h::gen(CUB_SEED(10), d_keys);\n\n thrust::host_vector h_reference = d_keys;\n std::stable_sort(\n thrust::raw_pointer_cast(h_reference.data()), \n thrust::raw_pointer_cast(h_reference.data()) + h_reference.size(), \n CustomLess{});\n\n block_merge_sort(\n d_keys,\n stable_sort_keys_full_tile_t{});\n\n REQUIRE( h_reference == d_keys );\n}\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of BlockMergeSort utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n\n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\nstruct CustomLess\n{\n template \n __device__ __host__ bool operator()(const DataType &lhs, const DataType &rhs)\n {\n return lhs < rhs;\n }\n};\n\ntemplate \n__global__ void block_merge_sort_kernel(\n KeyT *data, \n int valid_items,\n KeyT oob_default,\n ActionT action)\n{\n using BlockMergeSort =\n hipcub::BlockMergeSort;\n\n __shared__ typename BlockMergeSort::TempStorage temp_storage_shuffle;\n\n KeyT thread_data[ItemsPerThread];\n\n const int thread_offset = static_cast(threadIdx.x) * ItemsPerThread;\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n const int idx = thread_offset + item;\n thread_data[item] = idx < valid_items ? data[idx] : KeyT();\n }\n __syncthreads();\n\n BlockMergeSort sort(temp_storage_shuffle);\n\n action(sort, thread_data, valid_items, oob_default);\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n const int idx = thread_offset + item;\n\n if (idx >= valid_items)\n {\n break;\n }\n\n data[idx] = thread_data[item];\n }\n}\n\ntemplate \n__global__ void block_merge_sort_kernel(\n KeyT *keys, \n ValueT *vals, \n int valid_items,\n KeyT oob_default,\n ActionT action)\n{\n using BlockMergeSort =\n hipcub::BlockMergeSort;\n\n __shared__ typename BlockMergeSort::TempStorage temp_storage_shuffle;\n\n KeyT thread_keys[ItemsPerThread];\n ValueT thread_vals[ItemsPerThread];\n\n const int thread_offset = static_cast(threadIdx.x) * ItemsPerThread;\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n const int idx = thread_offset + item;\n thread_keys[item] = idx < valid_items ? keys[idx] : KeyT{};\n thread_vals[item] = idx < valid_items ? vals[idx] : ValueT{};\n }\n __syncthreads();\n\n BlockMergeSort sort(temp_storage_shuffle);\n\n action(sort, thread_keys, thread_vals, valid_items, oob_default);\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n const int idx = thread_offset + item;\n\n if (idx >= valid_items)\n {\n break;\n }\n\n keys[idx] = thread_keys[item];\n vals[idx] = thread_vals[item];\n }\n}\n\nstruct stable_sort_keys_partial_tile_t\n{\n template \n __device__ void operator()(BlockMergeSortT &sort,\n KeyT &thread_data,\n int valid_items,\n DefaultT oob_default) const\n {\n sort.StableSort(thread_data, CustomLess{}, valid_items, oob_default);\n }\n};\n\nstruct stable_sort_pairs_partial_tile_t\n{\n template \n __device__ void operator()(BlockMergeSortT &sort,\n KeyT &thread_keys,\n ValueT &thread_vals,\n int valid_items,\n DefaultT oob_default) const\n {\n sort.StableSort(thread_keys,\n thread_vals,\n CustomLess{},\n valid_items,\n oob_default);\n }\n};\n\nstruct stable_sort_pairs_full_tile_t\n{\n template \n __device__ void operator()(BlockMergeSortT &sort,\n KeyT &thread_keys,\n ValueT &thread_vals,\n int /* valid_items */,\n DefaultT /* oob_default */) const\n {\n sort.StableSort(thread_keys, thread_vals, CustomLess());\n }\n};\n\nstruct stable_sort_keys_full_tile_t\n{\n template \n __device__ void operator()(BlockMergeSortT &sort,\n KeyT &thread_keys,\n int /* valid_items */,\n DefaultT /* oob_default */) const\n {\n sort.StableSort(thread_keys, CustomLess());\n }\n};\n\ntemplate \nvoid block_merge_sort(\n thrust::device_vector &keys, \n ActionT action)\n{\n block_merge_sort_kernel\n <<<1, ThreadsInBlock>>>(\n thrust::raw_pointer_cast(keys.data()),\n static_cast(keys.size()),\n std::numeric_limits::max(),\n action);\n\n REQUIRE( hipSuccess == hipPeekAtLastError() );\n REQUIRE( hipSuccess == hipDeviceSynchronize() );\n}\n\ntemplate \nvoid block_merge_sort(\n thrust::device_vector &keys, \n thrust::device_vector &vals, \n ActionT action)\n{\n block_merge_sort_kernel\n <<<1, ThreadsInBlock>>>(\n thrust::raw_pointer_cast(keys.data()),\n thrust::raw_pointer_cast(vals.data()),\n static_cast(keys.size()),\n std::numeric_limits::max(),\n action);\n\n REQUIRE( hipSuccess == hipPeekAtLastError() );\n REQUIRE( hipSuccess == hipDeviceSynchronize() );\n}\n\n// %PARAM% THREADS_IN_BLOCK bs 64:256\n\nusing key_types = c2h::type_list;\nusing threads_in_block = c2h::enum_type_list;\nusing items_per_thread = c2h::enum_type_list;\n\n\ntemplate \nstruct params_t\n{\n using key_t = typename c2h::get<0, TestType>;\n\n static constexpr int items_per_thread = c2h::get<1, TestType>::value;\n static constexpr int threads_in_block = c2h::get<2, TestType>::value;\n static constexpr int tile_size = items_per_thread * threads_in_block;\n};\n\n\nCUB_TEST(\"Block merge sort can sort keys in partial tiles\",\n \"[merge sort][block]\",\n key_types,\n items_per_thread,\n threads_in_block)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n\n thrust::device_vector d_keys(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n\n c2h::gen(CUB_SEED(10), d_keys);\n\n thrust::host_vector h_reference = d_keys;\n std::stable_sort(\n thrust::raw_pointer_cast(h_reference.data()), \n thrust::raw_pointer_cast(h_reference.data()) + h_reference.size(), \n CustomLess{});\n\n block_merge_sort(\n d_keys,\n stable_sort_keys_partial_tile_t{});\n\n REQUIRE( h_reference == d_keys );\n}\n\nCUB_TEST(\"Block merge sort can sort keys in full tiles\",\n \"[merge sort][block]\",\n key_types,\n items_per_thread,\n threads_in_block)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n\n thrust::device_vector d_keys(params::tile_size);\n\n c2h::gen(CUB_SEED(10), d_keys);\n\n thrust::host_vector h_reference = d_keys;\n std::stable_sort(\n thrust::raw_pointer_cast(h_reference.data()), \n thrust::raw_pointer_cast(h_reference.data()) + h_reference.size(), \n CustomLess{});\n\n block_merge_sort(\n d_keys,\n stable_sort_keys_full_tile_t{});\n\n REQUIRE( h_reference == d_keys );\n}\n\nCUB_TEST(\"Block merge sort can sort pairs in partial tiles\",\n \"[merge sort][block]\",\n key_types,\n items_per_thread,\n threads_in_block)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n using value_t = key_t;\n using pair_t = std::pair;\n\n thrust::device_vector d_keys(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n thrust::device_vector d_vals(d_keys.size());\n\n c2h::gen(CUB_SEED(5), d_keys);\n c2h::gen(CUB_SEED(5), d_vals);\n\n thrust::host_vector h_keys = d_keys;\n thrust::host_vector h_vals = d_vals;\n\n thrust::host_vector h_ref(d_keys.size());\n\n for (std::size_t idx = 0; idx < h_ref.size(); idx++)\n {\n h_ref[idx] = std::make_pair(h_keys[idx], h_vals[idx]);\n }\n\n std::stable_sort(\n thrust::raw_pointer_cast(h_ref.data()), \n thrust::raw_pointer_cast(h_ref.data()) + h_ref.size(), \n [](pair_t l, pair_t r) -> bool {\n return l.first < r.first;\n });\n\n for (std::size_t idx = 0; idx < h_ref.size(); idx++)\n {\n h_keys[idx] = h_ref[idx].first;\n h_vals[idx] = h_ref[idx].second;\n }\n\n block_merge_sort(\n d_keys,\n d_vals,\n stable_sort_pairs_partial_tile_t{});\n\n REQUIRE( h_keys == d_keys );\n REQUIRE( h_vals == d_vals );\n}\n\nCUB_TEST(\"Block merge sort can sort pairs in full tiles\",\n \"[merge sort][block]\",\n key_types,\n items_per_thread,\n threads_in_block)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n using value_t = key_t;\n using pair_t = std::pair;\n\n thrust::device_vector d_keys(params::tile_size);\n thrust::device_vector d_vals(d_keys.size());\n\n c2h::gen(CUB_SEED(5), d_keys);\n c2h::gen(CUB_SEED(5), d_vals);\n\n thrust::host_vector h_keys = d_keys;\n thrust::host_vector h_vals = d_vals;\n\n thrust::host_vector h_ref(d_keys.size());\n\n for (std::size_t idx = 0; idx < h_ref.size(); idx++)\n {\n h_ref[idx] = std::make_pair(h_keys[idx], h_vals[idx]);\n }\n\n std::stable_sort(\n thrust::raw_pointer_cast(h_ref.data()), \n thrust::raw_pointer_cast(h_ref.data()) + h_ref.size(), \n [](pair_t l, pair_t r) -> bool {\n return l.first < r.first;\n });\n\n for (std::size_t idx = 0; idx < h_ref.size(); idx++)\n {\n h_keys[idx] = h_ref[idx].first;\n h_vals[idx] = h_ref[idx].second;\n }\n\n block_merge_sort(\n d_keys,\n d_vals,\n stable_sort_pairs_full_tile_t{});\n\n REQUIRE( h_keys == d_keys );\n REQUIRE( h_vals == d_vals );\n}\n\n\nCUB_TEST(\"Block merge sort can sort pairs with mixed types\",\n \"[merge sort][block]\",\n threads_in_block)\n{\n using key_t = std::int32_t;\n using value_t = std::int64_t;\n using pair_t = std::pair;\n\n constexpr int items_per_thread = 2;\n constexpr int threads_in_block = c2h::get<0, TestType>::value;\n constexpr int tile_size = items_per_thread * threads_in_block;\n\n thrust::device_vector d_keys(tile_size);\n thrust::device_vector d_vals(d_keys.size());\n\n c2h::gen(CUB_SEED(5), d_keys);\n c2h::gen(CUB_SEED(5), d_vals);\n\n thrust::host_vector h_keys = d_keys;\n thrust::host_vector h_vals = d_vals;\n\n thrust::host_vector h_ref(d_keys.size());\n\n for (std::size_t idx = 0; idx < h_ref.size(); idx++)\n {\n h_ref[idx] = std::make_pair(h_keys[idx], h_vals[idx]);\n }\n\n std::stable_sort(\n thrust::raw_pointer_cast(h_ref.data()), \n thrust::raw_pointer_cast(h_ref.data()) + h_ref.size(), \n [](pair_t l, pair_t r) -> bool {\n return l.first < r.first;\n });\n\n for (std::size_t idx = 0; idx < h_ref.size(); idx++)\n {\n h_keys[idx] = h_ref[idx].first;\n h_vals[idx] = h_ref[idx].second;\n }\n\n block_merge_sort(\n d_keys,\n d_vals,\n stable_sort_pairs_full_tile_t{});\n\n REQUIRE( h_keys == d_keys );\n REQUIRE( h_vals == d_vals );\n}\n\nCUB_TEST(\"Block merge sort can sort large tiles\",\n \"[merge sort][block]\",\n threads_in_block)\n{\n using key_t = std::uint16_t;\n\n constexpr int items_per_thread = 2;\n\n // Repurpose block sizes\n constexpr int cmake_threads_in_block = c2h::get<0, TestType>::value;\n constexpr int threads_in_block = cmake_threads_in_block < 256\n ? 512\n : 1024;\n\n constexpr int tile_size = threads_in_block * items_per_thread;\n\n thrust::device_vector d_keys(tile_size);\n c2h::gen(CUB_SEED(10), d_keys);\n\n thrust::host_vector h_reference = d_keys;\n std::stable_sort(\n thrust::raw_pointer_cast(h_reference.data()), \n thrust::raw_pointer_cast(h_reference.data()) + h_reference.size(), \n CustomLess{});\n\n block_merge_sort(\n d_keys,\n stable_sort_keys_full_tile_t{});\n\n REQUIRE( h_reference == d_keys );\n}\n\nCUB_TEST(\"Block merge sort is stable\",\n \"[merge sort][block]\",\n threads_in_block)\n{\n using key_t = c2h::custom_type_t;\n\n constexpr int items_per_thread = 2;\n constexpr int threads_in_block = c2h::get<0, TestType>::value;\n constexpr int tile_size = threads_in_block * items_per_thread;\n\n thrust::device_vector d_keys(tile_size);\n c2h::gen(CUB_SEED(10), d_keys);\n\n thrust::host_vector h_reference = d_keys;\n std::stable_sort(\n thrust::raw_pointer_cast(h_reference.data()), \n thrust::raw_pointer_cast(h_reference.data()) + h_reference.size(), \n CustomLess{});\n\n block_merge_sort(\n d_keys,\n stable_sort_keys_full_tile_t{});\n\n REQUIRE( h_reference == d_keys );\n}\n\n", "source": "/kaggle/working/cub/test/catch2_test_block_merge_sort.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n#include \n#include \n\n#include \n#include \n\n#include \"fill_striped.cuh\"\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void warp_store_kernel(OutputIteratorT output_iterator, ActionT action)\n{\n using warp_store_t = cub::WarpStore;\n using storage_t = typename warp_store_t::TempStorage;\n\n constexpr int tile_size = ITEMS_PER_THREAD * LOGICAL_WARP_THREADS;\n __shared__ storage_t storage[TOTAL_WARPS];\n\n const int tid = cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n T reg[ITEMS_PER_THREAD];\n\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n reg[item] = static_cast(tid * ITEMS_PER_THREAD + item);\n }\n\n const int warp_id = tid / LOGICAL_WARP_THREADS;\n warp_store_t store(storage[warp_id]);\n\n action(store, output_iterator + (warp_id * tile_size), reg);\n}\n\ntemplate \nvoid warp_store(OutputIteratorT output_iterator, ActionT action)\n{\n warp_store_kernel<<<1, TOTAL_WARPS * LOGICAL_WARP_THREADS>>>(output_iterator, action);\n}\n\nstruct guarded_store_t\n{\n int valid_items;\n template \n __device__ void\n operator()(cub::WarpStore store,\n OutputIteratorT output,\n T (®)[ITEMS_PER_THREAD])\n {\n store.Store(output, reg, valid_items);\n }\n};\n\nstruct unguarded_store_t\n{\n template \n __device__ void\n operator()(cub::WarpStore store,\n OutputIteratorT output,\n T (®)[ITEMS_PER_THREAD])\n {\n store.Store(output, reg);\n }\n};\n\ntemplate \nthrust::device_vector compute_reference(int valid_items)\n{\n const int tile_size = LOGICAL_WARP_THREADS * ITEMS_PER_THREAD;\n const int total_item_count = TOTAL_WARPS * tile_size;\n thrust::device_vector d_input(total_item_count);\n\n CUB_IF_CONSTEXPR(StoreAlgorithm == cub::WarpStoreAlgorithm::WARP_STORE_STRIPED)\n {\n thrust::host_vector input(total_item_count);\n fill_striped(\n input.begin());\n d_input = input;\n }\n else\n {\n c2h::gen(c2h::modulo_t{d_input.size()}, d_input);\n }\n if (valid_items != total_item_count)\n {\n for (int warp_id = 0; warp_id < TOTAL_WARPS; warp_id++)\n {\n thrust::fill(d_input.begin() + warp_id * tile_size + valid_items,\n d_input.begin() + (warp_id + 1) * tile_size,\n T{});\n }\n }\n return d_input;\n}\n\n// %PARAM% LWT lwt 4:16:32\n// %PARAM% ALGO_TYPE alg 0:1:2:3\n\nusing types = c2h::type_list;\nusing items_per_thread = c2h::enum_type_list;\nusing logical_warp_threads = c2h::enum_type_list;\nusing algorithms = c2h::enum_type_list;\nusing algorithm =\n c2h::enum_type_list::value>;\n\nusing cache_store_modifier = c2h::enum_type_list;\n\nconstexpr int guarded_store_tests_count = 30;\n\ntemplate \nstruct total_warps_t\n{\nprivate:\n static constexpr int max_warps = 2;\n static constexpr bool is_arch_warp = (logical_warp_threads == CUB_WARP_THREADS(0));\n static constexpr bool is_pow_of_two = ((logical_warp_threads & (logical_warp_threads - 1)) == 0);\n static constexpr int total_warps = (is_arch_warp || is_pow_of_two) ? max_warps : 1;\n\npublic:\n static constexpr int value() { return total_warps; }\n};\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int logical_warp_threads = c2h::get<1, TestType>::value;\n static constexpr int items_per_thread = c2h::get<2, TestType>::value;\n static constexpr cub::WarpStoreAlgorithm algorithm = c2h::get<3, TestType>::value;\n static constexpr int total_warps = total_warps_t::value();\n static constexpr int tile_size = logical_warp_threads * items_per_thread;\n static constexpr int total_item_count = total_warps * tile_size;\n};\n\nCUB_TEST(\"Warp store guarded range works with pointer\",\n \"[store][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::total_item_count, type{});\n const int valid_items =\n GENERATE_COPY(take(guarded_store_tests_count, random(0, params::tile_size - 1)));\n auto out = thrust::raw_pointer_cast(d_out.data());\n warp_store(out, guarded_store_t{valid_items});\n auto d_expected_output = compute_reference(valid_items);\n REQUIRE(d_expected_output == d_out);\n}\n\nCUB_TEST(\"Warp store guarded range works with cache modified iterator\",\n \"[store][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm,\n cache_store_modifier)\n{\n using params = params_t;\n using type = typename params::type;\n constexpr cub::CacheStoreModifier store_modifier = c2h::get<4, TestType>::value;\n\n thrust::device_vector d_out(params::total_item_count, type{});\n const int valid_items =\n GENERATE_COPY(take(guarded_store_tests_count, random(0, params::tile_size - 1)));\n auto out =\n cub::CacheModifiedOutputIterator(thrust::raw_pointer_cast(d_out.data()));\n warp_store(out, guarded_store_t{valid_items});\n auto d_expected_output = compute_reference(valid_items);\n REQUIRE(d_expected_output == d_out);\n}\n\nCUB_TEST(\"Warp store unguarded range works with pointer\",\n \"[store][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::total_item_count, type{});\n const int valid_items = params::tile_size;\n auto out = thrust::raw_pointer_cast(d_out.data());\n warp_store(out, unguarded_store_t{});\n auto d_expected_output = compute_reference(valid_items);\n REQUIRE(d_expected_output == d_out);\n}\n\nCUB_TEST(\"Warp store unguarded range works with cache modified iterator\",\n \"[store][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm,\n cache_store_modifier)\n{\n using params = params_t;\n using type = typename params::type;\n constexpr cub::CacheStoreModifier store_modifier = c2h::get<4, TestType>::value;\n\n thrust::device_vector d_out(params::total_item_count, type{});\n const int valid_items = params::tile_size;\n auto out =\n cub::CacheModifiedOutputIterator(thrust::raw_pointer_cast(d_out.data()));\n warp_store(out, unguarded_store_t{});\n auto d_expected_output = compute_reference(valid_items);\n REQUIRE(d_expected_output == d_out);\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n#include \n#include \n\n#include \"fill_striped.cuh\"\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void warp_store_kernel(OutputIteratorT output_iterator, ActionT action)\n{\n using warp_store_t = hipcub::WarpStore;\n using storage_t = typename warp_store_t::TempStorage;\n\n constexpr int tile_size = ITEMS_PER_THREAD * LOGICAL_WARP_THREADS;\n __shared__ storage_t storage[TOTAL_WARPS];\n\n const int tid = hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n T reg[ITEMS_PER_THREAD];\n\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n reg[item] = static_cast(tid * ITEMS_PER_THREAD + item);\n }\n\n const int warp_id = tid / LOGICAL_WARP_THREADS;\n warp_store_t store(storage[warp_id]);\n\n action(store, output_iterator + (warp_id * tile_size), reg);\n}\n\ntemplate \nvoid warp_store(OutputIteratorT output_iterator, ActionT action)\n{\n warp_store_kernel<<<1, TOTAL_WARPS * LOGICAL_WARP_THREADS>>>(output_iterator, action);\n}\n\nstruct guarded_store_t\n{\n int valid_items;\n template \n __device__ void\n operator()(hipcub::WarpStore store,\n OutputIteratorT output,\n T (®)[ITEMS_PER_THREAD])\n {\n store.Store(output, reg, valid_items);\n }\n};\n\nstruct unguarded_store_t\n{\n template \n __device__ void\n operator()(hipcub::WarpStore store,\n OutputIteratorT output,\n T (®)[ITEMS_PER_THREAD])\n {\n store.Store(output, reg);\n }\n};\n\ntemplate \nthrust::device_vector compute_reference(int valid_items)\n{\n const int tile_size = LOGICAL_WARP_THREADS * ITEMS_PER_THREAD;\n const int total_item_count = TOTAL_WARPS * tile_size;\n thrust::device_vector d_input(total_item_count);\n\n CUB_IF_CONSTEXPR(StoreAlgorithm == hipcub::WarpStoreAlgorithm::WARP_STORE_STRIPED)\n {\n thrust::host_vector input(total_item_count);\n fill_striped(\n input.begin());\n d_input = input;\n }\n else\n {\n c2h::gen(c2h::modulo_t{d_input.size()}, d_input);\n }\n if (valid_items != total_item_count)\n {\n for (int warp_id = 0; warp_id < TOTAL_WARPS; warp_id++)\n {\n thrust::fill(d_input.begin() + warp_id * tile_size + valid_items,\n d_input.begin() + (warp_id + 1) * tile_size,\n T{});\n }\n }\n return d_input;\n}\n\n// %PARAM% LWT lwt 4:16:32\n// %PARAM% ALGO_TYPE alg 0:1:2:3\n\nusing types = c2h::type_list;\nusing items_per_thread = c2h::enum_type_list;\nusing logical_warp_threads = c2h::enum_type_list;\nusing algorithms = c2h::enum_type_list;\nusing algorithm =\n c2h::enum_type_list::value>;\n\nusing cache_store_modifier = c2h::enum_type_list;\n\nconstexpr int guarded_store_tests_count = 30;\n\ntemplate \nstruct total_warps_t\n{\nprivate:\n static constexpr int max_warps = 2;\n static constexpr bool is_arch_warp = (logical_warp_threads == CUB_WARP_THREADS(0));\n static constexpr bool is_pow_of_two = ((logical_warp_threads & (logical_warp_threads - 1)) == 0);\n static constexpr int total_warps = (is_arch_warp || is_pow_of_two) ? max_warps : 1;\n\npublic:\n static constexpr int value() { return total_warps; }\n};\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int logical_warp_threads = c2h::get<1, TestType>::value;\n static constexpr int items_per_thread = c2h::get<2, TestType>::value;\n static constexpr hipcub::WarpStoreAlgorithm algorithm = c2h::get<3, TestType>::value;\n static constexpr int total_warps = total_warps_t::value();\n static constexpr int tile_size = logical_warp_threads * items_per_thread;\n static constexpr int total_item_count = total_warps * tile_size;\n};\n\nCUB_TEST(\"Warp store guarded range works with pointer\",\n \"[store][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::total_item_count, type{});\n const int valid_items =\n GENERATE_COPY(take(guarded_store_tests_count, random(0, params::tile_size - 1)));\n auto out = thrust::raw_pointer_cast(d_out.data());\n warp_store(out, guarded_store_t{valid_items});\n auto d_expected_output = compute_reference(valid_items);\n REQUIRE(d_expected_output == d_out);\n}\n\nCUB_TEST(\"Warp store guarded range works with cache modified iterator\",\n \"[store][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm,\n cache_store_modifier)\n{\n using params = params_t;\n using type = typename params::type;\n constexpr hipcub::CacheStoreModifier store_modifier = c2h::get<4, TestType>::value;\n\n thrust::device_vector d_out(params::total_item_count, type{});\n const int valid_items =\n GENERATE_COPY(take(guarded_store_tests_count, random(0, params::tile_size - 1)));\n auto out =\n hipcub::CacheModifiedOutputIterator(thrust::raw_pointer_cast(d_out.data()));\n warp_store(out, guarded_store_t{valid_items});\n auto d_expected_output = compute_reference(valid_items);\n REQUIRE(d_expected_output == d_out);\n}\n\nCUB_TEST(\"Warp store unguarded range works with pointer\",\n \"[store][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::total_item_count, type{});\n const int valid_items = params::tile_size;\n auto out = thrust::raw_pointer_cast(d_out.data());\n warp_store(out, unguarded_store_t{});\n auto d_expected_output = compute_reference(valid_items);\n REQUIRE(d_expected_output == d_out);\n}\n\nCUB_TEST(\"Warp store unguarded range works with cache modified iterator\",\n \"[store][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm,\n cache_store_modifier)\n{\n using params = params_t;\n using type = typename params::type;\n constexpr hipcub::CacheStoreModifier store_modifier = c2h::get<4, TestType>::value;\n\n thrust::device_vector d_out(params::total_item_count, type{});\n const int valid_items = params::tile_size;\n auto out =\n hipcub::CacheModifiedOutputIterator(thrust::raw_pointer_cast(d_out.data()));\n warp_store(out, unguarded_store_t{});\n auto d_expected_output = compute_reference(valid_items);\n REQUIRE(d_expected_output == d_out);\n}\n", "source": "/kaggle/working/cub/test/catch2_test_warp_store.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#undef NDEBUG\n#include \n\n#include \n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void init_kernel(ScanTileStateT tile_state, int blocks_in_grid)\n{\n tile_state.InitializeStatus(blocks_in_grid);\n}\n\ntemplate \n__global__ void decoupled_look_back_kernel(cub::ScanTileState tile_state,\n MessageT *tile_data)\n{\n using scan_op_t = cub::Sum;\n using scan_tile_state_t = cub::ScanTileState;\n using tile_prefix_op = cub::TilePrefixCallbackOp;\n using temp_storage_t = typename tile_prefix_op::TempStorage;\n\n // Allocate temp storage in shared memory\n __shared__ temp_storage_t temp_storage;\n\n scan_op_t scan_op{};\n const unsigned int threads_in_warp = 32;\n const unsigned int tid = threadIdx.x;\n\n // Construct prefix op\n tile_prefix_op prefix(tile_state, temp_storage, scan_op);\n const unsigned int tile_idx = prefix.GetTileIdx();\n\n // \"Compute\" tile aggregate\n MessageT tile_aggregate = tile_data[tile_idx];\n\n if (tile_idx == 0)\n {\n // There are no blocks to look back to, immediately set the inclusive state\n if (tid == 0)\n {\n tile_state.SetInclusive(tile_idx, tile_aggregate);\n tile_data[tile_idx] = tile_aggregate;\n }\n }\n else\n {\n // Only the first warp in the block can perform the look back\n const unsigned int warp_id = tid / threads_in_warp;\n\n if (warp_id == 0)\n {\n // Perform the decoupled look-back\n // Invocation of the prefix will block until the look-back is complete.\n MessageT exclusive_prefix = prefix(tile_aggregate);\n\n if (tid == 0)\n {\n MessageT inclusive_prefix = scan_op(exclusive_prefix, tile_aggregate);\n tile_data[tile_idx] = inclusive_prefix;\n }\n }\n __syncthreads();\n\n assert(tile_data[tile_idx] == prefix.GetInclusivePrefix());\n assert(tile_aggregate == prefix.GetBlockAggregate());\n }\n}\n\nusing message_types = c2h::type_list;\n\ntemplate \nthrust::host_vector\ncompute_reference(const thrust::device_vector &tile_aggregates)\n{\n if (tile_aggregates.empty())\n {\n return {};\n }\n\n thrust::host_vector reference = tile_aggregates;\n MessageT *h_reference = thrust::raw_pointer_cast(reference.data());\n\n MessageT aggregate = h_reference[0];\n for (std::size_t i = 1; i < reference.size(); i++)\n {\n aggregate += h_reference[i];\n h_reference[i] = aggregate;\n }\n\n return reference;\n}\n\nCUB_TEST(\"Decoupled look-back works with various message types\",\n \"[decoupled look-back][device]\",\n message_types)\n{\n using message_t = typename c2h::get<0, TestType>;\n using scan_tile_state_t = cub::ScanTileState;\n\n const int max_tiles = 1024 * 1024;\n const int num_tiles = GENERATE_COPY(take(10, random(1, max_tiles)));\n\n thrust::device_vector tile_data(num_tiles);\n message_t *d_tile_data = thrust::raw_pointer_cast(tile_data.data());\n\n c2h::gen(CUB_SEED(2), tile_data);\n thrust::host_vector reference = compute_reference(tile_data);\n\n // Query temporary storage requirements\n std::size_t temp_storage_bytes{};\n scan_tile_state_t::AllocationSize(num_tiles, temp_storage_bytes);\n\n // Allocate temporary storage\n thrust::device_vector temp_storage(temp_storage_bytes);\n std::uint8_t *d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n // Initialize temporary storage\n scan_tile_state_t tile_status;\n cudaError_t status = tile_status.Init(num_tiles, d_temp_storage, temp_storage_bytes);\n REQUIRE(status == cudaSuccess);\n\n const unsigned int threads_in_init_block = 256;\n const unsigned int blocks_in_init_grid = cub::DivideAndRoundUp(num_tiles, threads_in_init_block);\n init_kernel<<>>(tile_status, num_tiles);\n REQUIRE(cudaSuccess == cudaPeekAtLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n // Launch decoupled look-back\n const unsigned int threads_in_block = 256;\n decoupled_look_back_kernel<<>>(tile_status, d_tile_data);\n REQUIRE(cudaSuccess == cudaPeekAtLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n REQUIRE(reference == tile_data);\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#undef NDEBUG\n#include \n\n#include \n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void init_kernel(ScanTileStateT tile_state, int blocks_in_grid)\n{\n tile_state.InitializeStatus(blocks_in_grid);\n}\n\ntemplate \n__global__ void decoupled_look_back_kernel(hipcub::ScanTileState tile_state,\n MessageT *tile_data)\n{\n using scan_op_t = hipcub::Sum;\n using scan_tile_state_t = hipcub::ScanTileState;\n using tile_prefix_op = hipcub::TilePrefixCallbackOp;\n using temp_storage_t = typename tile_prefix_op::TempStorage;\n\n // Allocate temp storage in shared memory\n __shared__ temp_storage_t temp_storage;\n\n scan_op_t scan_op{};\n const unsigned int threads_in_warp = 32;\n const unsigned int tid = threadIdx.x;\n\n // Construct prefix op\n tile_prefix_op prefix(tile_state, temp_storage, scan_op);\n const unsigned int tile_idx = prefix.GetTileIdx();\n\n // \"Compute\" tile aggregate\n MessageT tile_aggregate = tile_data[tile_idx];\n\n if (tile_idx == 0)\n {\n // There are no blocks to look back to, immediately set the inclusive state\n if (tid == 0)\n {\n tile_state.SetInclusive(tile_idx, tile_aggregate);\n tile_data[tile_idx] = tile_aggregate;\n }\n }\n else\n {\n // Only the first warp in the block can perform the look back\n const unsigned int warp_id = tid / threads_in_warp;\n\n if (warp_id == 0)\n {\n // Perform the decoupled look-back\n // Invocation of the prefix will block until the look-back is complete.\n MessageT exclusive_prefix = prefix(tile_aggregate);\n\n if (tid == 0)\n {\n MessageT inclusive_prefix = scan_op(exclusive_prefix, tile_aggregate);\n tile_data[tile_idx] = inclusive_prefix;\n }\n }\n __syncthreads();\n\n assert(tile_data[tile_idx] == prefix.GetInclusivePrefix());\n assert(tile_aggregate == prefix.GetBlockAggregate());\n }\n}\n\nusing message_types = c2h::type_list;\n\ntemplate \nthrust::host_vector\ncompute_reference(const thrust::device_vector &tile_aggregates)\n{\n if (tile_aggregates.empty())\n {\n return {};\n }\n\n thrust::host_vector reference = tile_aggregates;\n MessageT *h_reference = thrust::raw_pointer_cast(reference.data());\n\n MessageT aggregate = h_reference[0];\n for (std::size_t i = 1; i < reference.size(); i++)\n {\n aggregate += h_reference[i];\n h_reference[i] = aggregate;\n }\n\n return reference;\n}\n\nCUB_TEST(\"Decoupled look-back works with various message types\",\n \"[decoupled look-back][device]\",\n message_types)\n{\n using message_t = typename c2h::get<0, TestType>;\n using scan_tile_state_t = hipcub::ScanTileState;\n\n const int max_tiles = 1024 * 1024;\n const int num_tiles = GENERATE_COPY(take(10, random(1, max_tiles)));\n\n thrust::device_vector tile_data(num_tiles);\n message_t *d_tile_data = thrust::raw_pointer_cast(tile_data.data());\n\n c2h::gen(CUB_SEED(2), tile_data);\n thrust::host_vector reference = compute_reference(tile_data);\n\n // Query temporary storage requirements\n std::size_t temp_storage_bytes{};\n scan_tile_state_t::AllocationSize(num_tiles, temp_storage_bytes);\n\n // Allocate temporary storage\n thrust::device_vector temp_storage(temp_storage_bytes);\n std::uint8_t *d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n // Initialize temporary storage\n scan_tile_state_t tile_status;\n hipError_t status = tile_status.Init(num_tiles, d_temp_storage, temp_storage_bytes);\n REQUIRE(status == hipSuccess);\n\n const unsigned int threads_in_init_block = 256;\n const unsigned int blocks_in_init_grid = hipcub::DivideAndRoundUp(num_tiles, threads_in_init_block);\n init_kernel<<>>(tile_status, num_tiles);\n REQUIRE(hipSuccess == hipPeekAtLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n // Launch decoupled look-back\n const unsigned int threads_in_block = 256;\n decoupled_look_back_kernel<<>>(tile_status, d_tile_data);\n REQUIRE(hipSuccess == hipPeekAtLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n REQUIRE(reference == tile_data);\n}\n", "source": "/kaggle/working/cub/test/catch2_test_device_decoupled_look_back.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n\n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void block_reduce_kernel(T *in, T *out, int valid_items, ActionT action)\n{\n using block_reduce_t = cub::BlockReduce;\n using storage_t = typename block_reduce_t::TempStorage;\n\n __shared__ storage_t storage;\n\n T thread_data[ItemsPerThread];\n\n const int tid = static_cast(cub::RowMajorTid(BlockDimX, BlockDimY, BlockDimZ));\n const int thread_offset = tid * ItemsPerThread;\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n const int idx = thread_offset + item;\n thread_data[item] = idx < valid_items ? in[idx] : T();\n }\n __syncthreads();\n\n block_reduce_t reduce(storage);\n\n T aggregate = action(reduce, thread_data, valid_items);\n\n if (tid == 0)\n {\n out[0] = aggregate;\n }\n}\n\ntemplate \nvoid block_reduce(thrust::device_vector &in, thrust::device_vector &out, ActionT action)\n{\n dim3 block_dims(BlockDimX, BlockDimY, BlockDimZ);\n\n block_reduce_kernel\n <<<1, block_dims>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(out.data()),\n static_cast(in.size()),\n action);\n\n REQUIRE(cudaSuccess == cudaPeekAtLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n}\n\nstruct sum_partial_tile_op_t\n{\n template \n __device__ T operator()(BlockReduceT &reduce,\n T (&thread_data)[ItemsPerThread],\n int valid_items) const\n {\n return reduce.Sum(thread_data[0], valid_items);\n }\n};\n\nstruct sum_full_tile_op_t\n{\n template \n __device__ T operator()(BlockReduceT &reduce,\n T (&thread_data)[ItemsPerThread],\n int /* valid_items */) const\n {\n return reduce.Sum(thread_data);\n }\n};\n\nstruct max_partial_tile_op_t\n{\n template \n __device__ T operator()(BlockReduceT &reduce, T (&thread_data)[ItemsPerThread], int valid_items) const\n {\n return reduce.Reduce(thread_data[0], cub::Max{}, valid_items);\n }\n};\n\nstruct max_full_tile_op_t\n{\n template \n __device__ T operator()(BlockReduceT &reduce,\n T (&thread_data)[ItemsPerThread],\n int /* valid_items */) const\n {\n return reduce.Reduce(thread_data, cub::Max{});\n }\n};\n\nusing types = c2h::type_list;\nusing vec_types = c2h::type_list;\n\n// %PARAM% TEST_DIM_X dimx 7:32:65:128\n// %PARAM% TEST_DIM_YZ dimyz 1:2\n\nusing block_dim_xs = c2h::enum_type_list;\nusing block_dim_yzs = c2h::enum_type_list;\nusing items_per_thread = c2h::enum_type_list;\nusing single_item_per_thread = c2h::enum_type_list;\nusing algorithm =\n c2h::enum_type_list;\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int items_per_thread = c2h::get<1, TestType>::value;\n static constexpr int block_dim_x = c2h::get<2, TestType>::value;\n static constexpr int block_dim_y = c2h::get<3, TestType>::value;\n static constexpr int block_dim_z = block_dim_y;\n static constexpr int tile_size = items_per_thread * block_dim_x * block_dim_y * block_dim_z;\n\n static constexpr cub::BlockReduceAlgorithm algorithm = c2h::get<4, TestType>::value;\n};\n\nCUB_TEST(\"Block reduce works with sum\",\n \"[reduce][block]\",\n types,\n items_per_thread,\n block_dim_xs,\n block_dim_yzs,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(1);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_reference(1, std::accumulate(h_in.begin() + 1, h_in.end(), h_in[0], [](const type &lhs, const type &rhs) {\n return static_cast(lhs + rhs);\n }));\n\n block_reduce(d_in, d_out, sum_full_tile_op_t{});\n\n REQUIRE_APPROX_EQ(h_reference, d_out);\n}\n\nCUB_TEST(\"Block reduce works with sum in partial tiles\",\n \"[reduce][block]\",\n types,\n single_item_per_thread,\n block_dim_xs,\n block_dim_yzs,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(1);\n thrust::device_vector d_in(GENERATE_COPY(take(2, random(1, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_in);\n\n thrust::host_vector h_in = d_in;\n std::vector h_reference(1, std::accumulate(h_in.begin() + 1, h_in.end(), h_in[0], [](const type &lhs, const type &rhs) {\n return static_cast(lhs + rhs);\n }));\n\n block_reduce(d_in, d_out, sum_partial_tile_op_t{});\n\n REQUIRE_APPROX_EQ(h_reference, d_out);\n}\n\nCUB_TEST(\"Block reduce works with custom op\",\n \"[reduce][block]\",\n types,\n items_per_thread,\n block_dim_xs,\n block_dim_yzs,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(1);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_reference(\n 1,\n std::accumulate(h_in.begin() + 1, h_in.end(), h_in[0], [](const type &lhs, const type &rhs) {\n return std::max(lhs, rhs);\n }));\n\n block_reduce(d_in, d_out, max_full_tile_op_t{});\n\n REQUIRE_APPROX_EQ(h_reference, d_out);\n}\n\nCUB_TEST(\"Block reduce works with custom op in partial tiles\",\n \"[reduce][block]\",\n types,\n single_item_per_thread,\n block_dim_xs,\n block_dim_yzs,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(1);\n thrust::device_vector d_in(GENERATE_COPY(take(2, random(1, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_in);\n\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_reference(\n 1,\n std::accumulate(h_in.begin() + 1, h_in.end(), h_in[0], [](const type &lhs, const type &rhs) {\n return std::max(lhs, rhs);\n }));\n\n block_reduce(d_in, d_out, max_partial_tile_op_t{});\n\n REQUIRE_APPROX_EQ(h_reference, d_out);\n}\n\nCUB_TEST(\"Block reduce works with custom types\",\n \"[reduce][block]\",\n block_dim_xs,\n block_dim_yzs,\n algorithm)\n{\n using type = c2h::custom_type_t;\n\n constexpr int items_per_thread = 1;\n constexpr int block_dim_x = c2h::get<0, TestType>::value;\n constexpr int block_dim_y = c2h::get<1, TestType>::value;\n constexpr int block_dim_z = block_dim_y;\n constexpr cub::BlockReduceAlgorithm algorithm = c2h::get<2, TestType>::value;\n\n constexpr int tile_size = block_dim_x * block_dim_y * block_dim_z * items_per_thread;\n\n thrust::device_vector d_out(1);\n thrust::device_vector d_in(GENERATE_COPY(take(2, random(1, tile_size))));\n c2h::gen(CUB_SEED(10), d_in);\n\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_reference(1, std::accumulate(h_in.begin() + 1, h_in.end(), h_in[0], [](const type &lhs, const type &rhs) {\n return static_cast(lhs + rhs);\n }));\n\n block_reduce(d_in, d_out, sum_partial_tile_op_t{});\n\n REQUIRE(h_reference == d_out);\n}\n\nCUB_TEST(\"Block reduce works with vec types\",\n \"[reduce][block]\",\n vec_types,\n block_dim_xs,\n block_dim_yzs,\n algorithm)\n{\n using type = c2h::get<0, TestType>;\n\n constexpr int items_per_thread = 1;\n constexpr int block_dim_x = c2h::get<1, TestType>::value;\n constexpr int block_dim_y = c2h::get<2, TestType>::value;\n constexpr int block_dim_z = block_dim_y;\n constexpr cub::BlockReduceAlgorithm algorithm = c2h::get<3, TestType>::value;\n\n constexpr int tile_size = block_dim_x * block_dim_y * block_dim_z * items_per_thread;\n\n thrust::device_vector d_out(1);\n thrust::device_vector d_in(GENERATE_COPY(take(2, random(1, tile_size))));\n c2h::gen(CUB_SEED(10), d_in);\n\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_reference(1, std::accumulate(h_in.begin() + 1, h_in.end(), h_in[0], [](const type &lhs, const type &rhs) {\n return static_cast(lhs + rhs);\n }));\n\n block_reduce(d_in, d_out, sum_partial_tile_op_t{});\n\n REQUIRE(h_reference == d_out);\n}\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void block_reduce_kernel(T *in, T *out, int valid_items, ActionT action)\n{\n using block_reduce_t = hipcub::BlockReduce;\n using storage_t = typename block_reduce_t::TempStorage;\n\n __shared__ storage_t storage;\n\n T thread_data[ItemsPerThread];\n\n const int tid = static_cast(hipcub::RowMajorTid(BlockDimX, BlockDimY, BlockDimZ));\n const int thread_offset = tid * ItemsPerThread;\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n const int idx = thread_offset + item;\n thread_data[item] = idx < valid_items ? in[idx] : T();\n }\n __syncthreads();\n\n block_reduce_t reduce(storage);\n\n T aggregate = action(reduce, thread_data, valid_items);\n\n if (tid == 0)\n {\n out[0] = aggregate;\n }\n}\n\ntemplate \nvoid block_reduce(thrust::device_vector &in, thrust::device_vector &out, ActionT action)\n{\n dim3 block_dims(BlockDimX, BlockDimY, BlockDimZ);\n\n block_reduce_kernel\n <<<1, block_dims>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(out.data()),\n static_cast(in.size()),\n action);\n\n REQUIRE(hipSuccess == hipPeekAtLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n}\n\nstruct sum_partial_tile_op_t\n{\n template \n __device__ T operator()(BlockReduceT &reduce,\n T (&thread_data)[ItemsPerThread],\n int valid_items) const\n {\n return reduce.Sum(thread_data[0], valid_items);\n }\n};\n\nstruct sum_full_tile_op_t\n{\n template \n __device__ T operator()(BlockReduceT &reduce,\n T (&thread_data)[ItemsPerThread],\n int /* valid_items */) const\n {\n return reduce.Sum(thread_data);\n }\n};\n\nstruct max_partial_tile_op_t\n{\n template \n __device__ T operator()(BlockReduceT &reduce, T (&thread_data)[ItemsPerThread], int valid_items) const\n {\n return reduce.Reduce(thread_data[0], hipcub::Max{}, valid_items);\n }\n};\n\nstruct max_full_tile_op_t\n{\n template \n __device__ T operator()(BlockReduceT &reduce,\n T (&thread_data)[ItemsPerThread],\n int /* valid_items */) const\n {\n return reduce.Reduce(thread_data, hipcub::Max{});\n }\n};\n\nusing types = c2h::type_list;\nusing vec_types = c2h::type_list;\n\n// %PARAM% TEST_DIM_X dimx 7:32:65:128\n// %PARAM% TEST_DIM_YZ dimyz 1:2\n\nusing block_dim_xs = c2h::enum_type_list;\nusing block_dim_yzs = c2h::enum_type_list;\nusing items_per_thread = c2h::enum_type_list;\nusing single_item_per_thread = c2h::enum_type_list;\nusing algorithm =\n c2h::enum_type_list;\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int items_per_thread = c2h::get<1, TestType>::value;\n static constexpr int block_dim_x = c2h::get<2, TestType>::value;\n static constexpr int block_dim_y = c2h::get<3, TestType>::value;\n static constexpr int block_dim_z = block_dim_y;\n static constexpr int tile_size = items_per_thread * block_dim_x * block_dim_y * block_dim_z;\n\n static constexpr hipcub::BlockReduceAlgorithm algorithm = c2h::get<4, TestType>::value;\n};\n\nCUB_TEST(\"Block reduce works with sum\",\n \"[reduce][block]\",\n types,\n items_per_thread,\n block_dim_xs,\n block_dim_yzs,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(1);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_reference(1, std::accumulate(h_in.begin() + 1, h_in.end(), h_in[0], [](const type &lhs, const type &rhs) {\n return static_cast(lhs + rhs);\n }));\n\n block_reduce(d_in, d_out, sum_full_tile_op_t{});\n\n REQUIRE_APPROX_EQ(h_reference, d_out);\n}\n\nCUB_TEST(\"Block reduce works with sum in partial tiles\",\n \"[reduce][block]\",\n types,\n single_item_per_thread,\n block_dim_xs,\n block_dim_yzs,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(1);\n thrust::device_vector d_in(GENERATE_COPY(take(2, random(1, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_in);\n\n thrust::host_vector h_in = d_in;\n std::vector h_reference(1, std::accumulate(h_in.begin() + 1, h_in.end(), h_in[0], [](const type &lhs, const type &rhs) {\n return static_cast(lhs + rhs);\n }));\n\n block_reduce(d_in, d_out, sum_partial_tile_op_t{});\n\n REQUIRE_APPROX_EQ(h_reference, d_out);\n}\n\nCUB_TEST(\"Block reduce works with custom op\",\n \"[reduce][block]\",\n types,\n items_per_thread,\n block_dim_xs,\n block_dim_yzs,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(1);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_reference(\n 1,\n std::accumulate(h_in.begin() + 1, h_in.end(), h_in[0], [](const type &lhs, const type &rhs) {\n return std::max(lhs, rhs);\n }));\n\n block_reduce(d_in, d_out, max_full_tile_op_t{});\n\n REQUIRE_APPROX_EQ(h_reference, d_out);\n}\n\nCUB_TEST(\"Block reduce works with custom op in partial tiles\",\n \"[reduce][block]\",\n types,\n single_item_per_thread,\n block_dim_xs,\n block_dim_yzs,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(1);\n thrust::device_vector d_in(GENERATE_COPY(take(2, random(1, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_in);\n\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_reference(\n 1,\n std::accumulate(h_in.begin() + 1, h_in.end(), h_in[0], [](const type &lhs, const type &rhs) {\n return std::max(lhs, rhs);\n }));\n\n block_reduce(d_in, d_out, max_partial_tile_op_t{});\n\n REQUIRE_APPROX_EQ(h_reference, d_out);\n}\n\nCUB_TEST(\"Block reduce works with custom types\",\n \"[reduce][block]\",\n block_dim_xs,\n block_dim_yzs,\n algorithm)\n{\n using type = c2h::custom_type_t;\n\n constexpr int items_per_thread = 1;\n constexpr int block_dim_x = c2h::get<0, TestType>::value;\n constexpr int block_dim_y = c2h::get<1, TestType>::value;\n constexpr int block_dim_z = block_dim_y;\n constexpr hipcub::BlockReduceAlgorithm algorithm = c2h::get<2, TestType>::value;\n\n constexpr int tile_size = block_dim_x * block_dim_y * block_dim_z * items_per_thread;\n\n thrust::device_vector d_out(1);\n thrust::device_vector d_in(GENERATE_COPY(take(2, random(1, tile_size))));\n c2h::gen(CUB_SEED(10), d_in);\n\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_reference(1, std::accumulate(h_in.begin() + 1, h_in.end(), h_in[0], [](const type &lhs, const type &rhs) {\n return static_cast(lhs + rhs);\n }));\n\n block_reduce(d_in, d_out, sum_partial_tile_op_t{});\n\n REQUIRE(h_reference == d_out);\n}\n\nCUB_TEST(\"Block reduce works with vec types\",\n \"[reduce][block]\",\n vec_types,\n block_dim_xs,\n block_dim_yzs,\n algorithm)\n{\n using type = c2h::get<0, TestType>;\n\n constexpr int items_per_thread = 1;\n constexpr int block_dim_x = c2h::get<1, TestType>::value;\n constexpr int block_dim_y = c2h::get<2, TestType>::value;\n constexpr int block_dim_z = block_dim_y;\n constexpr hipcub::BlockReduceAlgorithm algorithm = c2h::get<3, TestType>::value;\n\n constexpr int tile_size = block_dim_x * block_dim_y * block_dim_z * items_per_thread;\n\n thrust::device_vector d_out(1);\n thrust::device_vector d_in(GENERATE_COPY(take(2, random(1, tile_size))));\n c2h::gen(CUB_SEED(10), d_in);\n\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_reference(1, std::accumulate(h_in.begin() + 1, h_in.end(), h_in[0], [](const type &lhs, const type &rhs) {\n return static_cast(lhs + rhs);\n }));\n\n block_reduce(d_in, d_out, sum_partial_tile_op_t{});\n\n REQUIRE(h_reference == d_out);\n}\n\n", "source": "/kaggle/working/cub/test/catch2_test_block_reduce.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of BlockMergeSort utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n#include \n\n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void block_shuffle_kernel(T *data, ActionT action)\n{\n using block_shuffle_t = cub::BlockShuffle;\n using temp_storage_t = typename block_shuffle_t::TempStorage;\n\n __shared__ temp_storage_t temp_storage;\n\n T thread_data[ItemsPerThread];\n\n data += cub::RowMajorTid(BlockDimX, BlockDimY, BlockDimZ) * ItemsPerThread;\n for (int item = 0; item < ItemsPerThread; item++)\n {\n thread_data[item] = data[item];\n }\n __syncthreads();\n\n block_shuffle_t block_shuffle(temp_storage);\n action(block_shuffle, thread_data);\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n data[item] = thread_data[item];\n }\n}\n\nstruct up_op_t \n{\n template \n __device__ void\n operator()(BlockShuffleT &block_shuffle,\n T (&thread_data)[ItemsPerThread]) const\n {\n block_shuffle.Up(thread_data, thread_data);\n }\n};\n\nstruct offset_op_t \n{\n int m_distance;\n\n __host__ offset_op_t(int distance)\n : m_distance(distance)\n {}\n\n template \n __device__ void\n operator()(BlockShuffleT &block_shuffle,\n T (&thread_data)[ItemsPerThread]) const\n {\n block_shuffle.Offset(thread_data[0], thread_data[0], m_distance);\n }\n};\n\nstruct rotate_op_t \n{\n unsigned int m_distance;\n\n __host__ rotate_op_t(unsigned int distance)\n : m_distance(distance)\n {}\n\n template \n __device__ void\n operator()(BlockShuffleT &block_shuffle,\n T (&thread_data)[ItemsPerThread]) const\n {\n block_shuffle.Rotate(thread_data[0], thread_data[0], m_distance);\n }\n};\n\ntemplate \nstruct up_with_suffix_op_t \n{\n int m_target_thread_id;\n T * m_d_suffix_ptr;\n\n __host__ up_with_suffix_op_t(\n int target_thread_id,\n T *d_suffix_ptr)\n : m_target_thread_id(target_thread_id)\n , m_d_suffix_ptr(d_suffix_ptr)\n {}\n\n template \n __device__ void\n operator()(BlockShuffleT &block_shuffle,\n T (&thread_data)[ItemsPerThread]) const\n {\n T suffix{};\n\n block_shuffle.Up(thread_data, thread_data, suffix);\n\n if (cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z) == m_target_thread_id)\n {\n m_d_suffix_ptr[0] = suffix;\n }\n }\n};\n\nstruct down_op_t \n{\n template \n __device__ void\n operator()(BlockShuffleT &block_shuffle,\n T (&thread_data)[ItemsPerThread]) const\n {\n block_shuffle.Down(thread_data, thread_data);\n }\n};\n\ntemplate \nstruct down_with_prefix_op_t \n{\n int m_target_thread_id;\n T * m_d_prefix_ptr;\n\n __host__ down_with_prefix_op_t(\n int target_thread_id,\n T *d_prefix_ptr)\n : m_target_thread_id(target_thread_id)\n , m_d_prefix_ptr(d_prefix_ptr)\n {}\n\n template \n __device__ void\n operator()(BlockShuffleT &block_shuffle,\n T (&thread_data)[ItemsPerThread]) const\n {\n T prefix{};\n\n block_shuffle.Down(thread_data, thread_data, prefix);\n\n if (cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z) == m_target_thread_id)\n {\n m_d_prefix_ptr[0] = prefix;\n }\n }\n};\n\ntemplate \nvoid block_shuffle(\n thrust::device_vector &data, \n ActionT action)\n{\n dim3 block(BlockDimX, BlockDimY, BlockDimZ);\n block_shuffle_kernel\n <<<1, block>>>(\n thrust::raw_pointer_cast(data.data()),\n action);\n\n REQUIRE( cudaSuccess == cudaPeekAtLastError() );\n REQUIRE( cudaSuccess == cudaDeviceSynchronize() );\n}\n\n// %PARAM% MULTI_DIM mdim 0:1\n// %PARAM% DIM_IDX dim_idx 0:1:2\n\n#if MULTI_DIM\nusing block_dim_xs = c2h::enum_type_list;\nusing block_dim_yz = c2h::enum_type_list;\n#else\nusing block_dim_xs = c2h::enum_type_list;\nusing block_dim_yz = c2h::enum_type_list;\n#endif\n\nusing block_dim_x = c2h::enum_type_list::value>;\n\nusing types = c2h::type_list;\nusing items_per_thread = c2h::enum_type_list;\nusing single_item_per_thread = c2h::enum_type_list;\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int items_per_thread = c2h::get<1, TestType>::value;\n static constexpr int block_dim_x = c2h::get<2, TestType>::value;\n static constexpr int block_dim_y = c2h::get<3, TestType>::value;\n static constexpr int block_dim_z = block_dim_y;\n static constexpr int threads_in_block = block_dim_x * block_dim_y * block_dim_z;\n static constexpr int tile_size = items_per_thread * threads_in_block;\n};\n\n\nCUB_TEST(\"Block shuffle offset works\",\n \"[shuffle][block]\",\n types,\n single_item_per_thread,\n block_dim_x,\n block_dim_yz)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n const int distance = \n GENERATE_COPY(take(4, random(1 - params::tile_size, \n params::tile_size - 1)));\n\n thrust::host_vector h_data = d_data;\n thrust::host_vector h_ref(params::tile_size);\n\n for (int i = 0; i < static_cast(h_data.size()); i++)\n {\n const int source = i + distance;\n h_ref[i] = (source >= 0) && (source < params::tile_size) \n ? h_data[source]\n : h_data[i];\n }\n\n block_shuffle(d_data, offset_op_t{distance});\n\n REQUIRE(h_ref == d_data);\n}\n\nCUB_TEST(\"Block shuffle rotate works\",\n \"[shuffle][block]\",\n types,\n single_item_per_thread,\n block_dim_x,\n block_dim_yz)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n thrust::device_vector d_ref = d_data;\n\n const unsigned int distance = \n GENERATE_COPY(take(4, random(0, params::tile_size - 1)));\n\n thrust::host_vector h_ref = d_data;\n std::rotate(h_ref.begin(), h_ref.begin() + distance, h_ref.end());\n\n block_shuffle(d_data, rotate_op_t{distance});\n\n REQUIRE(h_ref == d_data);\n}\n\nCUB_TEST(\"Block shuffle up works\",\n \"[shuffle][block]\",\n types,\n items_per_thread,\n block_dim_x,\n block_dim_yz)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n thrust::device_vector d_ref(params::tile_size);\n thrust::copy(d_data.begin(), d_data.end() - 1, d_ref.begin() + 1);\n thrust::copy(d_data.begin(), d_data.begin() + 1, d_ref.begin());\n\n block_shuffle(d_data, up_op_t{});\n\n REQUIRE(d_ref == d_data);\n}\n\nCUB_TEST(\"Block shuffle up works when suffix is required\",\n \"[shuffle][block]\",\n types,\n items_per_thread,\n block_dim_x,\n block_dim_yz)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n const int target_thread_id = \n GENERATE_COPY(take(2, random(0, params::threads_in_block - 1)));\n\n thrust::device_vector d_ref(params::tile_size);\n thrust::copy(d_data.begin(), d_data.end() - 1, d_ref.begin() + 1);\n thrust::copy(d_data.begin(), d_data.begin() + 1, d_ref.begin());\n\n thrust::device_vector d_suffix(1);\n thrust::device_vector d_suffix_ref(1);\n thrust::copy(d_data.end() - 1, d_data.end(), d_suffix_ref.begin());\n\n block_shuffle(\n d_data,\n up_with_suffix_op_t{target_thread_id,\n thrust::raw_pointer_cast(d_suffix.data())});\n\n REQUIRE(d_ref == d_data);\n REQUIRE(d_suffix_ref == d_suffix);\n}\n\nCUB_TEST(\"Block shuffle down works\",\n \"[shuffle][block]\",\n types,\n items_per_thread,\n block_dim_x,\n block_dim_yz)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n thrust::device_vector d_ref(params::tile_size);\n thrust::copy(d_data.begin() + 1, d_data.end(), d_ref.begin());\n thrust::copy(d_data.end() - 1, d_data.end(), d_ref.end() - 1);\n\n block_shuffle(d_data, down_op_t{});\n\n REQUIRE(d_ref == d_data);\n}\n\nCUB_TEST(\"Block shuffle down works when prefix is required\",\n \"[shuffle][block]\",\n types,\n items_per_thread,\n block_dim_x,\n block_dim_yz)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n const int target_thread_id = \n GENERATE_COPY(take(2, random(0, params::threads_in_block - 1)));\n\n thrust::device_vector d_ref(params::tile_size);\n thrust::copy(d_data.begin() + 1, d_data.end(), d_ref.begin());\n thrust::copy(d_data.end() - 1, d_data.end(), d_ref.end() - 1);\n\n thrust::device_vector d_prefix(1);\n thrust::device_vector d_prefix_ref(1);\n thrust::copy(d_data.begin(), d_data.begin() + 1, d_prefix_ref.begin());\n\n block_shuffle(\n d_data,\n down_with_prefix_op_t{target_thread_id,\n thrust::raw_pointer_cast(d_prefix.data())});\n\n REQUIRE(d_ref == d_data);\n REQUIRE(d_prefix_ref == d_prefix);\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of BlockMergeSort utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n\n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void block_shuffle_kernel(T *data, ActionT action)\n{\n using block_shuffle_t = hipcub::BlockShuffle;\n using temp_storage_t = typename block_shuffle_t::TempStorage;\n\n __shared__ temp_storage_t temp_storage;\n\n T thread_data[ItemsPerThread];\n\n data += hipcub::RowMajorTid(BlockDimX, BlockDimY, BlockDimZ) * ItemsPerThread;\n for (int item = 0; item < ItemsPerThread; item++)\n {\n thread_data[item] = data[item];\n }\n __syncthreads();\n\n block_shuffle_t block_shuffle(temp_storage);\n action(block_shuffle, thread_data);\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n data[item] = thread_data[item];\n }\n}\n\nstruct up_op_t \n{\n template \n __device__ void\n operator()(BlockShuffleT &block_shuffle,\n T (&thread_data)[ItemsPerThread]) const\n {\n block_shuffle.Up(thread_data, thread_data);\n }\n};\n\nstruct offset_op_t \n{\n int m_distance;\n\n __host__ offset_op_t(int distance)\n : m_distance(distance)\n {}\n\n template \n __device__ void\n operator()(BlockShuffleT &block_shuffle,\n T (&thread_data)[ItemsPerThread]) const\n {\n block_shuffle.Offset(thread_data[0], thread_data[0], m_distance);\n }\n};\n\nstruct rotate_op_t \n{\n unsigned int m_distance;\n\n __host__ rotate_op_t(unsigned int distance)\n : m_distance(distance)\n {}\n\n template \n __device__ void\n operator()(BlockShuffleT &block_shuffle,\n T (&thread_data)[ItemsPerThread]) const\n {\n block_shuffle.Rotate(thread_data[0], thread_data[0], m_distance);\n }\n};\n\ntemplate \nstruct up_with_suffix_op_t \n{\n int m_target_thread_id;\n T * m_d_suffix_ptr;\n\n __host__ up_with_suffix_op_t(\n int target_thread_id,\n T *d_suffix_ptr)\n : m_target_thread_id(target_thread_id)\n , m_d_suffix_ptr(d_suffix_ptr)\n {}\n\n template \n __device__ void\n operator()(BlockShuffleT &block_shuffle,\n T (&thread_data)[ItemsPerThread]) const\n {\n T suffix{};\n\n block_shuffle.Up(thread_data, thread_data, suffix);\n\n if (hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z) == m_target_thread_id)\n {\n m_d_suffix_ptr[0] = suffix;\n }\n }\n};\n\nstruct down_op_t \n{\n template \n __device__ void\n operator()(BlockShuffleT &block_shuffle,\n T (&thread_data)[ItemsPerThread]) const\n {\n block_shuffle.Down(thread_data, thread_data);\n }\n};\n\ntemplate \nstruct down_with_prefix_op_t \n{\n int m_target_thread_id;\n T * m_d_prefix_ptr;\n\n __host__ down_with_prefix_op_t(\n int target_thread_id,\n T *d_prefix_ptr)\n : m_target_thread_id(target_thread_id)\n , m_d_prefix_ptr(d_prefix_ptr)\n {}\n\n template \n __device__ void\n operator()(BlockShuffleT &block_shuffle,\n T (&thread_data)[ItemsPerThread]) const\n {\n T prefix{};\n\n block_shuffle.Down(thread_data, thread_data, prefix);\n\n if (hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z) == m_target_thread_id)\n {\n m_d_prefix_ptr[0] = prefix;\n }\n }\n};\n\ntemplate \nvoid block_shuffle(\n thrust::device_vector &data, \n ActionT action)\n{\n dim3 block(BlockDimX, BlockDimY, BlockDimZ);\n block_shuffle_kernel\n <<<1, block>>>(\n thrust::raw_pointer_cast(data.data()),\n action);\n\n REQUIRE( hipSuccess == hipPeekAtLastError() );\n REQUIRE( hipSuccess == hipDeviceSynchronize() );\n}\n\n// %PARAM% MULTI_DIM mdim 0:1\n// %PARAM% DIM_IDX dim_idx 0:1:2\n\n#if MULTI_DIM\nusing block_dim_xs = c2h::enum_type_list;\nusing block_dim_yz = c2h::enum_type_list;\n#else\nusing block_dim_xs = c2h::enum_type_list;\nusing block_dim_yz = c2h::enum_type_list;\n#endif\n\nusing block_dim_x = c2h::enum_type_list::value>;\n\nusing types = c2h::type_list;\nusing items_per_thread = c2h::enum_type_list;\nusing single_item_per_thread = c2h::enum_type_list;\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int items_per_thread = c2h::get<1, TestType>::value;\n static constexpr int block_dim_x = c2h::get<2, TestType>::value;\n static constexpr int block_dim_y = c2h::get<3, TestType>::value;\n static constexpr int block_dim_z = block_dim_y;\n static constexpr int threads_in_block = block_dim_x * block_dim_y * block_dim_z;\n static constexpr int tile_size = items_per_thread * threads_in_block;\n};\n\n\nCUB_TEST(\"Block shuffle offset works\",\n \"[shuffle][block]\",\n types,\n single_item_per_thread,\n block_dim_x,\n block_dim_yz)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n const int distance = \n GENERATE_COPY(take(4, random(1 - params::tile_size, \n params::tile_size - 1)));\n\n thrust::host_vector h_data = d_data;\n thrust::host_vector h_ref(params::tile_size);\n\n for (int i = 0; i < static_cast(h_data.size()); i++)\n {\n const int source = i + distance;\n h_ref[i] = (source >= 0) && (source < params::tile_size) \n ? h_data[source]\n : h_data[i];\n }\n\n block_shuffle(d_data, offset_op_t{distance});\n\n REQUIRE(h_ref == d_data);\n}\n\nCUB_TEST(\"Block shuffle rotate works\",\n \"[shuffle][block]\",\n types,\n single_item_per_thread,\n block_dim_x,\n block_dim_yz)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n thrust::device_vector d_ref = d_data;\n\n const unsigned int distance = \n GENERATE_COPY(take(4, random(0, params::tile_size - 1)));\n\n thrust::host_vector h_ref = d_data;\n std::rotate(h_ref.begin(), h_ref.begin() + distance, h_ref.end());\n\n block_shuffle(d_data, rotate_op_t{distance});\n\n REQUIRE(h_ref == d_data);\n}\n\nCUB_TEST(\"Block shuffle up works\",\n \"[shuffle][block]\",\n types,\n items_per_thread,\n block_dim_x,\n block_dim_yz)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n thrust::device_vector d_ref(params::tile_size);\n thrust::copy(d_data.begin(), d_data.end() - 1, d_ref.begin() + 1);\n thrust::copy(d_data.begin(), d_data.begin() + 1, d_ref.begin());\n\n block_shuffle(d_data, up_op_t{});\n\n REQUIRE(d_ref == d_data);\n}\n\nCUB_TEST(\"Block shuffle up works when suffix is required\",\n \"[shuffle][block]\",\n types,\n items_per_thread,\n block_dim_x,\n block_dim_yz)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n const int target_thread_id = \n GENERATE_COPY(take(2, random(0, params::threads_in_block - 1)));\n\n thrust::device_vector d_ref(params::tile_size);\n thrust::copy(d_data.begin(), d_data.end() - 1, d_ref.begin() + 1);\n thrust::copy(d_data.begin(), d_data.begin() + 1, d_ref.begin());\n\n thrust::device_vector d_suffix(1);\n thrust::device_vector d_suffix_ref(1);\n thrust::copy(d_data.end() - 1, d_data.end(), d_suffix_ref.begin());\n\n block_shuffle(\n d_data,\n up_with_suffix_op_t{target_thread_id,\n thrust::raw_pointer_cast(d_suffix.data())});\n\n REQUIRE(d_ref == d_data);\n REQUIRE(d_suffix_ref == d_suffix);\n}\n\nCUB_TEST(\"Block shuffle down works\",\n \"[shuffle][block]\",\n types,\n items_per_thread,\n block_dim_x,\n block_dim_yz)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n thrust::device_vector d_ref(params::tile_size);\n thrust::copy(d_data.begin() + 1, d_data.end(), d_ref.begin());\n thrust::copy(d_data.end() - 1, d_data.end(), d_ref.end() - 1);\n\n block_shuffle(d_data, down_op_t{});\n\n REQUIRE(d_ref == d_data);\n}\n\nCUB_TEST(\"Block shuffle down works when prefix is required\",\n \"[shuffle][block]\",\n types,\n items_per_thread,\n block_dim_x,\n block_dim_yz)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n const int target_thread_id = \n GENERATE_COPY(take(2, random(0, params::threads_in_block - 1)));\n\n thrust::device_vector d_ref(params::tile_size);\n thrust::copy(d_data.begin() + 1, d_data.end(), d_ref.begin());\n thrust::copy(d_data.end() - 1, d_data.end(), d_ref.end() - 1);\n\n thrust::device_vector d_prefix(1);\n thrust::device_vector d_prefix_ref(1);\n thrust::copy(d_data.begin(), d_data.begin() + 1, d_prefix_ref.begin());\n\n block_shuffle(\n d_data,\n down_with_prefix_op_t{target_thread_id,\n thrust::raw_pointer_cast(d_prefix.data())});\n\n REQUIRE(d_ref == d_data);\n REQUIRE(d_prefix_ref == d_prefix);\n}\n", "source": "/kaggle/working/cub/test/catch2_test_block_shuffle.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of iterator utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"test_util.h\"\n\nusing namespace cub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n\ntemplate \nstruct TransformOp\n{\n // Increment transform\n __host__ __device__ __forceinline__ T operator()(T input) const\n {\n T addend;\n InitValue(INTEGER_SEED, addend, 1);\n return input + addend;\n }\n};\n\nstruct SelectOp\n{\n template \n __host__ __device__ __forceinline__ bool operator()(T input)\n {\n return true;\n }\n};\n\n\n//---------------------------------------------------------------------\n// Test kernels\n//---------------------------------------------------------------------\n\n/**\n * Test random access input iterator\n */\ntemplate <\n typename InputIteratorT,\n typename T>\n__global__ void Kernel(\n InputIteratorT d_in,\n T *d_out,\n InputIteratorT *d_itrs)\n{\n d_out[0] = *d_in; // Value at offset 0\n d_out[1] = d_in[100]; // Value at offset 100\n d_out[2] = *(d_in + 1000); // Value at offset 1000\n d_out[3] = *(d_in + 10000); // Value at offset 10000\n\n d_in++;\n d_out[4] = d_in[0]; // Value at offset 1\n\n d_in += 20;\n d_out[5] = d_in[0]; // Value at offset 21\n d_itrs[0] = d_in; // Iterator at offset 21\n\n d_in -= 10;\n d_out[6] = d_in[0]; // Value at offset 11;\n\n d_in -= 11;\n d_out[7] = d_in[0]; // Value at offset 0\n d_itrs[1] = d_in; // Iterator at offset 0\n}\n\n\n\n//---------------------------------------------------------------------\n// Host testing subroutines\n//---------------------------------------------------------------------\n\n\n/**\n * Run iterator test on device\n */\ntemplate <\n typename InputIteratorT,\n typename T,\n int TEST_VALUES>\nvoid Test(\n InputIteratorT d_in,\n T (&h_reference)[TEST_VALUES])\n{\n // Allocate device arrays\n T *d_out = NULL;\n InputIteratorT *d_itrs = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(T) * TEST_VALUES));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_itrs, sizeof(InputIteratorT) * 2));\n\n int compare;\n\n // Run unguarded kernel\n Kernel<<<1, 1>>>(d_in, d_out, d_itrs);\n\n CubDebugExit(cudaPeekAtLastError());\n CubDebugExit(cudaDeviceSynchronize());\n\n // Check results\n compare = CompareDeviceResults(h_reference, d_out, TEST_VALUES, g_verbose, g_verbose);\n printf(\"\\tValues: %s\\n\", (compare) ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Check iterator at offset 21\n InputIteratorT h_itr = d_in + 21;\n compare = CompareDeviceResults(&h_itr, d_itrs, 1, g_verbose, g_verbose);\n printf(\"\\tIterators: %s\\n\", (compare) ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Check iterator at offset 0\n compare = CompareDeviceResults(&d_in, d_itrs + 1, 1, g_verbose, g_verbose);\n printf(\"\\tIterators: %s\\n\", (compare) ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Cleanup\n if (d_out) CubDebugExit(g_allocator.DeviceFree(d_out));\n if (d_itrs) CubDebugExit(g_allocator.DeviceFree(d_itrs));\n}\n\n\n/**\n * Test constant iterator\n */\ntemplate \nvoid TestConstant(T base)\n{\n printf(\"\\nTesting constant iterator on type %s (base: %lld)\\n\", typeid(T).name(), (unsigned long long) (base)); fflush(stdout);\n\n //\n // Test iterator manipulation in kernel\n //\n\n T h_reference[8] = {base, base, base, base, base, base, base, base};\n ConstantInputIterator d_itr(base);\n Test(d_itr, h_reference);\n}\n\n\n/**\n * Test counting iterator\n */\ntemplate \nvoid TestCounting(T base)\n{\n printf(\"\\nTesting counting iterator on type %s (base: %d) \\n\", typeid(T).name(), int(base)); fflush(stdout);\n\n //\n // Test iterator manipulation in kernel\n //\n\n // Initialize reference data\n T h_reference[8];\n h_reference[0] = static_cast(base + 0); // Value at offset 0\n h_reference[1] = static_cast(base + 100); // Value at offset 100\n h_reference[2] = static_cast(base + 1000); // Value at offset 1000\n h_reference[3] = static_cast(base + 10000); // Value at offset 10000\n h_reference[4] = static_cast(base + 1); // Value at offset 1\n h_reference[5] = static_cast(base + 21); // Value at offset 21\n h_reference[6] = static_cast(base + 11); // Value at offset 11\n h_reference[7] = static_cast(base + 0); // Value at offset 0;\n\n CountingInputIterator d_itr(base);\n Test(d_itr, h_reference);\n}\n\n\n/**\n * Test modified iterator\n */\ntemplate \nvoid TestModified()\n{\n printf(\"\\nTesting cache-modified iterator on type %s\\n\", typeid(T).name()); fflush(stdout);\n\n //\n // Test iterator manipulation in kernel\n //\n\n constexpr int TEST_VALUES = 11000;\n\n T *h_data = new T[TEST_VALUES];\n for (int i = 0; i < TEST_VALUES; ++i)\n {\n RandomBits(h_data[i]);\n }\n\n // Allocate device arrays\n T *d_data = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_data, sizeof(T) * TEST_VALUES));\n CubDebugExit(cudaMemcpy(d_data, h_data, sizeof(T) * TEST_VALUES, cudaMemcpyHostToDevice));\n\n // Initialize reference data\n T h_reference[8];\n h_reference[0] = h_data[0]; // Value at offset 0\n h_reference[1] = h_data[100]; // Value at offset 100\n h_reference[2] = h_data[1000]; // Value at offset 1000\n h_reference[3] = h_data[10000]; // Value at offset 10000\n h_reference[4] = h_data[1]; // Value at offset 1\n h_reference[5] = h_data[21]; // Value at offset 21\n h_reference[6] = h_data[11]; // Value at offset 11\n h_reference[7] = h_data[0]; // Value at offset 0;\n\n Test(CacheModifiedInputIterator((CastT*) d_data), h_reference);\n Test(CacheModifiedInputIterator((CastT*) d_data), h_reference);\n Test(CacheModifiedInputIterator((CastT*) d_data), h_reference);\n Test(CacheModifiedInputIterator((CastT*) d_data), h_reference);\n Test(CacheModifiedInputIterator((CastT*) d_data), h_reference);\n Test(CacheModifiedInputIterator((CastT*) d_data), h_reference);\n Test(CacheModifiedInputIterator((CastT*) d_data), h_reference);\n\n if (h_data) delete[] h_data;\n if (d_data) CubDebugExit(g_allocator.DeviceFree(d_data));\n}\n\n\n/**\n * Test transform iterator\n */\ntemplate \nvoid TestTransform()\n{\n printf(\"\\nTesting transform iterator on type %s\\n\", typeid(T).name()); fflush(stdout);\n\n //\n // Test iterator manipulation in kernel\n //\n\n constexpr int TEST_VALUES = 11000;\n\n T *h_data = new T[TEST_VALUES];\n for (int i = 0; i < TEST_VALUES; ++i)\n {\n InitValue(INTEGER_SEED, h_data[i], i);\n }\n\n // Allocate device arrays\n T *d_data = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_data, sizeof(T) * TEST_VALUES));\n CubDebugExit(cudaMemcpy(d_data, h_data, sizeof(T) * TEST_VALUES, cudaMemcpyHostToDevice));\n\n TransformOp op;\n\n // Initialize reference data\n T h_reference[8];\n h_reference[0] = op(h_data[0]); // Value at offset 0\n h_reference[1] = op(h_data[100]); // Value at offset 100\n h_reference[2] = op(h_data[1000]); // Value at offset 1000\n h_reference[3] = op(h_data[10000]); // Value at offset 10000\n h_reference[4] = op(h_data[1]); // Value at offset 1\n h_reference[5] = op(h_data[21]); // Value at offset 21\n h_reference[6] = op(h_data[11]); // Value at offset 11\n h_reference[7] = op(h_data[0]); // Value at offset 0;\n\n TransformInputIterator, CastT*> d_itr((CastT*) d_data, op);\n Test(d_itr, h_reference);\n\n if (h_data) delete[] h_data;\n if (d_data) CubDebugExit(g_allocator.DeviceFree(d_data));\n}\n\n\n/**\n * Test tex-obj texture iterator\n */\ntemplate \nvoid TestTexObj()\n{\n printf(\"\\nTesting tex-obj iterator on type %s\\n\", typeid(T).name()); fflush(stdout);\n\n //\n // Test iterator manipulation in kernel\n //\n\n const unsigned int TEST_VALUES = 11000;\n const unsigned int DUMMY_OFFSET = 500;\n const unsigned int DUMMY_TEST_VALUES = TEST_VALUES - DUMMY_OFFSET;\n\n T *h_data = new T[TEST_VALUES];\n for (unsigned int i = 0; i < TEST_VALUES; ++i)\n {\n RandomBits(h_data[i]);\n }\n\n // Allocate device arrays\n T *d_data = NULL;\n T *d_dummy = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_data, sizeof(T) * TEST_VALUES));\n CubDebugExit(cudaMemcpy(d_data, h_data, sizeof(T) * TEST_VALUES, cudaMemcpyHostToDevice));\n\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_dummy, sizeof(T) * DUMMY_TEST_VALUES));\n CubDebugExit(cudaMemcpy(d_dummy, h_data + DUMMY_OFFSET, sizeof(T) * DUMMY_TEST_VALUES, cudaMemcpyHostToDevice));\n\n // Initialize reference data\n T h_reference[8];\n h_reference[0] = h_data[0]; // Value at offset 0\n h_reference[1] = h_data[100]; // Value at offset 100\n h_reference[2] = h_data[1000]; // Value at offset 1000\n h_reference[3] = h_data[10000]; // Value at offset 10000\n h_reference[4] = h_data[1]; // Value at offset 1\n h_reference[5] = h_data[21]; // Value at offset 21\n h_reference[6] = h_data[11]; // Value at offset 11\n h_reference[7] = h_data[0]; // Value at offset 0;\n\n // Create and bind obj-based test iterator\n TexObjInputIterator d_obj_itr;\n CubDebugExit(d_obj_itr.BindTexture((CastT*) d_data, sizeof(T) * TEST_VALUES));\n\n Test(d_obj_itr, h_reference);\n\n if (h_data) delete[] h_data;\n if (d_data) CubDebugExit(g_allocator.DeviceFree(d_data));\n if (d_dummy) CubDebugExit(g_allocator.DeviceFree(d_dummy));\n}\n\n/**\n * Test texture transform iterator\n */\ntemplate \nvoid TestTexTransform()\n{\n printf(\"\\nTesting tex-transform iterator on type %s\\n\", typeid(T).name()); fflush(stdout);\n\n //\n // Test iterator manipulation in kernel\n //\n\n constexpr int TEST_VALUES = 11000;\n\n T *h_data = new T[TEST_VALUES];\n for (int i = 0; i < TEST_VALUES; ++i)\n {\n InitValue(INTEGER_SEED, h_data[i], i);\n }\n\n // Allocate device arrays\n T *d_data = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_data, sizeof(T) * TEST_VALUES));\n CubDebugExit(cudaMemcpy(d_data, h_data, sizeof(T) * TEST_VALUES, cudaMemcpyHostToDevice));\n\n TransformOp op;\n\n // Initialize reference data\n T h_reference[8];\n h_reference[0] = op(h_data[0]); // Value at offset 0\n h_reference[1] = op(h_data[100]); // Value at offset 100\n h_reference[2] = op(h_data[1000]); // Value at offset 1000\n h_reference[3] = op(h_data[10000]); // Value at offset 10000\n h_reference[4] = op(h_data[1]); // Value at offset 1\n h_reference[5] = op(h_data[21]); // Value at offset 21\n h_reference[6] = op(h_data[11]); // Value at offset 11\n h_reference[7] = op(h_data[0]); // Value at offset 0;\n\n // Create and bind texture iterator\n typedef TexObjInputIterator TextureIterator;\n\n TextureIterator d_tex_itr;\n CubDebugExit(d_tex_itr.BindTexture((CastT*) d_data, sizeof(T) * TEST_VALUES));\n\n // Create transform iterator\n TransformInputIterator, TextureIterator> xform_itr(d_tex_itr, op);\n\n Test(xform_itr, h_reference);\n\n CubDebugExit(d_tex_itr.UnbindTexture());\n if (h_data) delete[] h_data;\n if (d_data) CubDebugExit(g_allocator.DeviceFree(d_data));\n}\n\n/**\n * Run non-integer tests\n */\ntemplate \nvoid Test(Int2Type /* is_integer */)\n{\n TestModified();\n TestTransform();\n TestTexObj();\n TestTexTransform();\n}\n\n/**\n * Run integer tests\n */\ntemplate \nvoid Test(Int2Type /* is_integer */)\n{\n TestConstant(0);\n TestConstant(99);\n\n TestCounting(0);\n TestCounting(99);\n\n // Run non-integer tests\n Test(Int2Type());\n}\n\n/**\n * Run tests\n */\ntemplate \nvoid Test()\n{\n enum {\n IS_INTEGER = (Traits::CATEGORY == SIGNED_INTEGER) || (Traits::CATEGORY == UNSIGNED_INTEGER)\n };\n\n // Test non-const type\n Test(Int2Type());\n\n // Test non-const type\n Test(Int2Type());\n}\n\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // Get ptx version\n int ptx_version = 0;\n CubDebugExit(PtxVersion(ptx_version));\n\n // Evaluate different data types\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n\n printf(\"\\nTest complete\\n\"); fflush(stdout);\n\n return 0;\n}\n\n\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of iterator utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"test_util.h\"\n\nusing namespace hipcub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n\ntemplate \nstruct TransformOp\n{\n // Increment transform\n __host__ __device__ __forceinline__ T operator()(T input) const\n {\n T addend;\n InitValue(INTEGER_SEED, addend, 1);\n return input + addend;\n }\n};\n\nstruct SelectOp\n{\n template \n __host__ __device__ __forceinline__ bool operator()(T input)\n {\n return true;\n }\n};\n\n\n//---------------------------------------------------------------------\n// Test kernels\n//---------------------------------------------------------------------\n\n/**\n * Test random access input iterator\n */\ntemplate <\n typename InputIteratorT,\n typename T>\n__global__ void Kernel(\n InputIteratorT d_in,\n T *d_out,\n InputIteratorT *d_itrs)\n{\n d_out[0] = *d_in; // Value at offset 0\n d_out[1] = d_in[100]; // Value at offset 100\n d_out[2] = *(d_in + 1000); // Value at offset 1000\n d_out[3] = *(d_in + 10000); // Value at offset 10000\n\n d_in++;\n d_out[4] = d_in[0]; // Value at offset 1\n\n d_in += 20;\n d_out[5] = d_in[0]; // Value at offset 21\n d_itrs[0] = d_in; // Iterator at offset 21\n\n d_in -= 10;\n d_out[6] = d_in[0]; // Value at offset 11;\n\n d_in -= 11;\n d_out[7] = d_in[0]; // Value at offset 0\n d_itrs[1] = d_in; // Iterator at offset 0\n}\n\n\n\n//---------------------------------------------------------------------\n// Host testing subroutines\n//---------------------------------------------------------------------\n\n\n/**\n * Run iterator test on device\n */\ntemplate <\n typename InputIteratorT,\n typename T,\n int TEST_VALUES>\nvoid Test(\n InputIteratorT d_in,\n T (&h_reference)[TEST_VALUES])\n{\n // Allocate device arrays\n T *d_out = NULL;\n InputIteratorT *d_itrs = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(T) * TEST_VALUES));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_itrs, sizeof(InputIteratorT) * 2));\n\n int compare;\n\n // Run unguarded kernel\n Kernel<<<1, 1>>>(d_in, d_out, d_itrs);\n\n CubDebugExit(hipPeekAtLastError());\n CubDebugExit(hipDeviceSynchronize());\n\n // Check results\n compare = CompareDeviceResults(h_reference, d_out, TEST_VALUES, g_verbose, g_verbose);\n printf(\"\\tValues: %s\\n\", (compare) ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Check iterator at offset 21\n InputIteratorT h_itr = d_in + 21;\n compare = CompareDeviceResults(&h_itr, d_itrs, 1, g_verbose, g_verbose);\n printf(\"\\tIterators: %s\\n\", (compare) ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Check iterator at offset 0\n compare = CompareDeviceResults(&d_in, d_itrs + 1, 1, g_verbose, g_verbose);\n printf(\"\\tIterators: %s\\n\", (compare) ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Cleanup\n if (d_out) CubDebugExit(g_allocator.DeviceFree(d_out));\n if (d_itrs) CubDebugExit(g_allocator.DeviceFree(d_itrs));\n}\n\n\n/**\n * Test constant iterator\n */\ntemplate \nvoid TestConstant(T base)\n{\n printf(\"\\nTesting constant iterator on type %s (base: %lld)\\n\", typeid(T).name(), (unsigned long long) (base)); fflush(stdout);\n\n //\n // Test iterator manipulation in kernel\n //\n\n T h_reference[8] = {base, base, base, base, base, base, base, base};\n ConstantInputIterator d_itr(base);\n Test(d_itr, h_reference);\n}\n\n\n/**\n * Test counting iterator\n */\ntemplate \nvoid TestCounting(T base)\n{\n printf(\"\\nTesting counting iterator on type %s (base: %d) \\n\", typeid(T).name(), int(base)); fflush(stdout);\n\n //\n // Test iterator manipulation in kernel\n //\n\n // Initialize reference data\n T h_reference[8];\n h_reference[0] = static_cast(base + 0); // Value at offset 0\n h_reference[1] = static_cast(base + 100); // Value at offset 100\n h_reference[2] = static_cast(base + 1000); // Value at offset 1000\n h_reference[3] = static_cast(base + 10000); // Value at offset 10000\n h_reference[4] = static_cast(base + 1); // Value at offset 1\n h_reference[5] = static_cast(base + 21); // Value at offset 21\n h_reference[6] = static_cast(base + 11); // Value at offset 11\n h_reference[7] = static_cast(base + 0); // Value at offset 0;\n\n CountingInputIterator d_itr(base);\n Test(d_itr, h_reference);\n}\n\n\n/**\n * Test modified iterator\n */\ntemplate \nvoid TestModified()\n{\n printf(\"\\nTesting cache-modified iterator on type %s\\n\", typeid(T).name()); fflush(stdout);\n\n //\n // Test iterator manipulation in kernel\n //\n\n constexpr int TEST_VALUES = 11000;\n\n T *h_data = new T[TEST_VALUES];\n for (int i = 0; i < TEST_VALUES; ++i)\n {\n RandomBits(h_data[i]);\n }\n\n // Allocate device arrays\n T *d_data = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_data, sizeof(T) * TEST_VALUES));\n CubDebugExit(hipMemcpy(d_data, h_data, sizeof(T) * TEST_VALUES, hipMemcpyHostToDevice));\n\n // Initialize reference data\n T h_reference[8];\n h_reference[0] = h_data[0]; // Value at offset 0\n h_reference[1] = h_data[100]; // Value at offset 100\n h_reference[2] = h_data[1000]; // Value at offset 1000\n h_reference[3] = h_data[10000]; // Value at offset 10000\n h_reference[4] = h_data[1]; // Value at offset 1\n h_reference[5] = h_data[21]; // Value at offset 21\n h_reference[6] = h_data[11]; // Value at offset 11\n h_reference[7] = h_data[0]; // Value at offset 0;\n\n Test(CacheModifiedInputIterator((CastT*) d_data), h_reference);\n Test(CacheModifiedInputIterator((CastT*) d_data), h_reference);\n Test(CacheModifiedInputIterator((CastT*) d_data), h_reference);\n Test(CacheModifiedInputIterator((CastT*) d_data), h_reference);\n Test(CacheModifiedInputIterator((CastT*) d_data), h_reference);\n Test(CacheModifiedInputIterator((CastT*) d_data), h_reference);\n Test(CacheModifiedInputIterator((CastT*) d_data), h_reference);\n\n if (h_data) delete[] h_data;\n if (d_data) CubDebugExit(g_allocator.DeviceFree(d_data));\n}\n\n\n/**\n * Test transform iterator\n */\ntemplate \nvoid TestTransform()\n{\n printf(\"\\nTesting transform iterator on type %s\\n\", typeid(T).name()); fflush(stdout);\n\n //\n // Test iterator manipulation in kernel\n //\n\n constexpr int TEST_VALUES = 11000;\n\n T *h_data = new T[TEST_VALUES];\n for (int i = 0; i < TEST_VALUES; ++i)\n {\n InitValue(INTEGER_SEED, h_data[i], i);\n }\n\n // Allocate device arrays\n T *d_data = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_data, sizeof(T) * TEST_VALUES));\n CubDebugExit(hipMemcpy(d_data, h_data, sizeof(T) * TEST_VALUES, hipMemcpyHostToDevice));\n\n TransformOp op;\n\n // Initialize reference data\n T h_reference[8];\n h_reference[0] = op(h_data[0]); // Value at offset 0\n h_reference[1] = op(h_data[100]); // Value at offset 100\n h_reference[2] = op(h_data[1000]); // Value at offset 1000\n h_reference[3] = op(h_data[10000]); // Value at offset 10000\n h_reference[4] = op(h_data[1]); // Value at offset 1\n h_reference[5] = op(h_data[21]); // Value at offset 21\n h_reference[6] = op(h_data[11]); // Value at offset 11\n h_reference[7] = op(h_data[0]); // Value at offset 0;\n\n TransformInputIterator, CastT*> d_itr((CastT*) d_data, op);\n Test(d_itr, h_reference);\n\n if (h_data) delete[] h_data;\n if (d_data) CubDebugExit(g_allocator.DeviceFree(d_data));\n}\n\n\n/**\n * Test tex-obj texture iterator\n */\ntemplate \nvoid TestTexObj()\n{\n printf(\"\\nTesting tex-obj iterator on type %s\\n\", typeid(T).name()); fflush(stdout);\n\n //\n // Test iterator manipulation in kernel\n //\n\n const unsigned int TEST_VALUES = 11000;\n const unsigned int DUMMY_OFFSET = 500;\n const unsigned int DUMMY_TEST_VALUES = TEST_VALUES - DUMMY_OFFSET;\n\n T *h_data = new T[TEST_VALUES];\n for (unsigned int i = 0; i < TEST_VALUES; ++i)\n {\n RandomBits(h_data[i]);\n }\n\n // Allocate device arrays\n T *d_data = NULL;\n T *d_dummy = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_data, sizeof(T) * TEST_VALUES));\n CubDebugExit(hipMemcpy(d_data, h_data, sizeof(T) * TEST_VALUES, hipMemcpyHostToDevice));\n\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_dummy, sizeof(T) * DUMMY_TEST_VALUES));\n CubDebugExit(hipMemcpy(d_dummy, h_data + DUMMY_OFFSET, sizeof(T) * DUMMY_TEST_VALUES, hipMemcpyHostToDevice));\n\n // Initialize reference data\n T h_reference[8];\n h_reference[0] = h_data[0]; // Value at offset 0\n h_reference[1] = h_data[100]; // Value at offset 100\n h_reference[2] = h_data[1000]; // Value at offset 1000\n h_reference[3] = h_data[10000]; // Value at offset 10000\n h_reference[4] = h_data[1]; // Value at offset 1\n h_reference[5] = h_data[21]; // Value at offset 21\n h_reference[6] = h_data[11]; // Value at offset 11\n h_reference[7] = h_data[0]; // Value at offset 0;\n\n // Create and bind obj-based test iterator\n TexObjInputIterator d_obj_itr;\n CubDebugExit(d_obj_itr.BindTexture((CastT*) d_data, sizeof(T) * TEST_VALUES));\n\n Test(d_obj_itr, h_reference);\n\n if (h_data) delete[] h_data;\n if (d_data) CubDebugExit(g_allocator.DeviceFree(d_data));\n if (d_dummy) CubDebugExit(g_allocator.DeviceFree(d_dummy));\n}\n\n/**\n * Test texture transform iterator\n */\ntemplate \nvoid TestTexTransform()\n{\n printf(\"\\nTesting tex-transform iterator on type %s\\n\", typeid(T).name()); fflush(stdout);\n\n //\n // Test iterator manipulation in kernel\n //\n\n constexpr int TEST_VALUES = 11000;\n\n T *h_data = new T[TEST_VALUES];\n for (int i = 0; i < TEST_VALUES; ++i)\n {\n InitValue(INTEGER_SEED, h_data[i], i);\n }\n\n // Allocate device arrays\n T *d_data = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_data, sizeof(T) * TEST_VALUES));\n CubDebugExit(hipMemcpy(d_data, h_data, sizeof(T) * TEST_VALUES, hipMemcpyHostToDevice));\n\n TransformOp op;\n\n // Initialize reference data\n T h_reference[8];\n h_reference[0] = op(h_data[0]); // Value at offset 0\n h_reference[1] = op(h_data[100]); // Value at offset 100\n h_reference[2] = op(h_data[1000]); // Value at offset 1000\n h_reference[3] = op(h_data[10000]); // Value at offset 10000\n h_reference[4] = op(h_data[1]); // Value at offset 1\n h_reference[5] = op(h_data[21]); // Value at offset 21\n h_reference[6] = op(h_data[11]); // Value at offset 11\n h_reference[7] = op(h_data[0]); // Value at offset 0;\n\n // Create and bind texture iterator\n typedef TexObjInputIterator TextureIterator;\n\n TextureIterator d_tex_itr;\n CubDebugExit(d_tex_itr.BindTexture((CastT*) d_data, sizeof(T) * TEST_VALUES));\n\n // Create transform iterator\n TransformInputIterator, TextureIterator> xform_itr(d_tex_itr, op);\n\n Test(xform_itr, h_reference);\n\n CubDebugExit(d_tex_itr.UnbindTexture());\n if (h_data) delete[] h_data;\n if (d_data) CubDebugExit(g_allocator.DeviceFree(d_data));\n}\n\n/**\n * Run non-integer tests\n */\ntemplate \nvoid Test(Int2Type /* is_integer */)\n{\n TestModified();\n TestTransform();\n TestTexObj();\n TestTexTransform();\n}\n\n/**\n * Run integer tests\n */\ntemplate \nvoid Test(Int2Type /* is_integer */)\n{\n TestConstant(0);\n TestConstant(99);\n\n TestCounting(0);\n TestCounting(99);\n\n // Run non-integer tests\n Test(Int2Type());\n}\n\n/**\n * Run tests\n */\ntemplate \nvoid Test()\n{\n enum {\n IS_INTEGER = (Traits::CATEGORY == SIGNED_INTEGER) || (Traits::CATEGORY == UNSIGNED_INTEGER)\n };\n\n // Test non-const type\n Test(Int2Type());\n\n // Test non-const type\n Test(Int2Type());\n}\n\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // Get ptx version\n int ptx_version = 0;\n CubDebugExit(PtxVersion(ptx_version));\n\n // Evaluate different data types\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n\n printf(\"\\nTest complete\\n\"); fflush(stdout);\n\n return 0;\n}\n\n\n\n", "source": "/kaggle/working/cub/test/test_iterator.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceReduce utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"test_util.h\"\n#include \n\n#define TEST_HALF_T !_NVHPC_CUDA\n\n#define TEST_BF_T !_NVHPC_CUDA\n\n#if TEST_HALF_T\n#include \n\n// Half support is provided by SM53+. We currently test against a few older architectures. \n// The specializations below can be removed once we drop these architectures. \nnamespace cub {\n\ntemplate <>\n__host__ __device__ __forceinline__ //\n__half Min::operator()(__half& a, __half& b) const \n{\n NV_IF_TARGET(NV_PROVIDES_SM_53, \n (return CUB_MIN(a, b);),\n (return CUB_MIN(__half2float(a), __half2float(b));));\n}\n\ntemplate <>\n__host__ __device__ __forceinline__ //\nKeyValuePair \nArgMin::operator()(const KeyValuePair &a, \n const KeyValuePair &b) const \n{\n const float av = __half2float(a.value);\n const float bv = __half2float(b.value);\n\n if ((bv < av) || ((av == bv) && (b.key < a.key)))\n {\n return b;\n }\n\n return a;\n}\n\ntemplate <>\n__host__ __device__ __forceinline__ //\n__half Max::operator()(__half& a, __half& b) const \n{\n NV_IF_TARGET(NV_PROVIDES_SM_53, \n (return CUB_MAX(a, b);),\n (return CUB_MAX(__half2float(a), __half2float(b));));\n}\n\ntemplate <>\n__host__ __device__ __forceinline__ //\nKeyValuePair \nArgMax::operator()(const KeyValuePair &a, \n const KeyValuePair &b) const \n{\n const float av = __half2float(a.value);\n const float bv = __half2float(b.value);\n\n if ((bv > av) || ((av == bv) && (b.key < a.key)))\n {\n return b;\n }\n\n return a;\n}\n\n} // namespace cub\n#endif\n\n#if TEST_BF_T\n#include \n#endif\n\nusing namespace cub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nint g_ptx_version;\nint g_sm_count;\ndouble g_device_giga_bandwidth;\nbool g_verbose = false;\nbool g_verbose_input = false;\nint g_timing_iterations = 0;\nCachingDeviceAllocator g_allocator(true);\n\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CUB_SEGMENTED, // CUB segmented method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n CDP_SEGMENTED, // GPU-based segmented method\n};\n\ninline const char* BackendToString(Backend b)\n{\n switch (b)\n {\n case CUB:\n return \"CUB\";\n case CUB_SEGMENTED:\n return \"CUB_SEGMENTED\";\n case CDP:\n return \"CDP\";\n case CDP_SEGMENTED:\n return \"CDP_SEGMENTED\";\n default:\n break;\n }\n\n return \"\";\n}\n\n// Custom max functor\nstruct CustomMax\n{\n /// Boolean max operator, returns (a > b) ? a : b\n template \n __host__ __device__ auto operator()(T&& a, C&& b) \n -> cub::detail::accumulator_t\n {\n return CUB_MAX(a, b);\n }\n\n#if TEST_HALF_T\n __host__ __device__ __half operator()(__half& a, __half& b) \n {\n return cub::Max{}(a, b);\n }\n#endif\n};\n\n// Comparing results computed on CPU and GPU for extended floating point types is impossible. \n// For instance, when used with a constant iterator of two, the accumulator in sequential reference \n// computation (CPU) bumps into the 4096 limits, which will never change (`4096 + 2 = 4096`). \n// Meanwhile, per-thread aggregates (`2 * 16 = 32`) are accumulated within and among thread blocks, \n// yielding `inf` as a result. No reasonable epsilon can be selected to compare `inf` with `4096`. \n// To make `__half` and `__nv_bfloat16` arithmetic associative, the function object below raises \n// extended floating points to the area of unsigned short integers. This allows us to test large \n// inputs with few code-path differences in device algorithms. \nstruct ExtendedFloatSum\n{\n template \n __host__ __device__ T operator()(T a, T b) const\n {\n T result{};\n result.__x = a.raw() + b.raw();\n return result;\n }\n\n#if TEST_HALF_T\n __host__ __device__ __half operator()(__half a, __half b) const\n {\n uint16_t result = this->operator()(half_t{a}, half_t(b)).raw();\n return reinterpret_cast<__half &>(result);\n }\n#endif\n\n#if TEST_BF_T\n __device__ __nv_bfloat16 operator()(__nv_bfloat16 a, __nv_bfloat16 b) const\n {\n uint16_t result = this->operator()(bfloat16_t{a}, bfloat16_t(b)).raw();\n return reinterpret_cast<__nv_bfloat16 &>(result);\n }\n#endif\n};\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB DeviceReduce entrypoints\n//---------------------------------------------------------------------\n\n/**\n * Dispatch to reduce entrypoint (custom-max)\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int /*max_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n ReductionOpT reduction_op)\n{\n using InputT = cub::detail::value_t;\n\n // The output value type\n using OutputT = cub::detail::non_void_value_t;\n\n // Max-identity\n OutputT identity = Traits::Lowest(); // replace with std::numeric_limits::lowest() when C++ support is more prevalent\n\n // Invoke kernel to device reduction directly\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceReduce::Reduce(d_temp_storage, temp_storage_bytes,\n d_in, d_out, num_items, reduction_op, identity);\n }\n\n return error;\n}\n\n/**\n * Dispatch to sum entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n cudaError_t * /*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int /*max_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n cub::Sum /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items);\n }\n\n return error;\n}\n\n/**\n * Dispatch to extended fp sum entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION cudaError_t Dispatch(Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n cudaError_t * /*d_cdp_error*/,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int /*max_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n ExtendedFloatSum reduction_op)\n{\n using InputT = cub::detail::value_t;\n using OutputT = cub::detail::non_void_value_t;\n\n OutputT identity{};\n\n // Invoke kernel to device reduction directly\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceReduce::Reduce(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items,\n reduction_op,\n identity);\n }\n\n return error;\n}\n\n/**\n * Dispatch to min entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int /*max_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n cub::Min /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceReduce::Min(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items);\n }\n\n return error;\n}\n\n/**\n * Dispatch to max entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int /*max_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n cub::Max /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceReduce::Max(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items);\n }\n\n return error;\n}\n\n/**\n * Dispatch to argmin entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int /*max_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n cub::ArgMin /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceReduce::ArgMin(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items);\n }\n\n return error;\n}\n\n/**\n * Dispatch to argmax entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int /*max_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n cub::ArgMax /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceReduce::ArgMax(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items);\n }\n\n return error;\n}\n\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB DeviceSegmentedReduce entrypoints\n//---------------------------------------------------------------------\n\n/**\n * Dispatch to reduce entrypoint (custom-max)\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int /*num_items*/,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n ReductionOpT reduction_op)\n{\n // The input value type\n using InputT = cub::detail::value_t;\n\n // The output value type\n using OutputT = cub::detail::non_void_value_t;\n\n // Max-identity\n OutputT identity = Traits::Lowest(); // replace with std::numeric_limits::lowest() when C++ support is more prevalent\n\n // Invoke kernel to device reduction directly\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceSegmentedReduce::Reduce(d_temp_storage, temp_storage_bytes,\n d_in, d_out, max_segments, d_segment_begin_offsets, d_segment_end_offsets, reduction_op, identity);\n }\n return error;\n}\n\n/**\n * Dispatch to sum entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int /*num_items*/,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n cub::Sum /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceSegmentedReduce::Sum(d_temp_storage, temp_storage_bytes,\n d_in, d_out, max_segments, d_segment_begin_offsets, d_segment_end_offsets);\n }\n return error;\n}\n\n/**\n * Dispatch to extended fp sum entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION cudaError_t Dispatch(Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n cudaError_t * /*d_cdp_error*/,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int /*num_items*/,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n ExtendedFloatSum reduction_op)\n{\n using InputT = cub::detail::value_t;\n using OutputT = cub::detail::non_void_value_t;\n\n OutputT identity{};\n\n // Invoke kernel to device reduction directly\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceSegmentedReduce::Reduce(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n max_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n reduction_op,\n identity);\n }\n\n return error;\n}\n\n/**\n * Dispatch to min entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int /*num_items*/,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n cub::Min /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceSegmentedReduce::Min(d_temp_storage, temp_storage_bytes,\n d_in, d_out, max_segments, d_segment_begin_offsets, d_segment_end_offsets);\n }\n return error;\n}\n\n/**\n * Dispatch to max entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int /*num_items*/,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n cub::Max /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceSegmentedReduce::Max(d_temp_storage, temp_storage_bytes,\n d_in, d_out, max_segments, d_segment_begin_offsets, d_segment_end_offsets);\n }\n return error;\n}\n\n/**\n * Dispatch to argmin entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int /*num_items*/,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n cub::ArgMin /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceSegmentedReduce::ArgMin(d_temp_storage, temp_storage_bytes,\n d_in, d_out, max_segments, d_segment_begin_offsets, d_segment_end_offsets);\n }\n return error;\n}\n\n/**\n * Dispatch to argmax entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int /*num_items*/,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n cub::ArgMax /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceSegmentedReduce::ArgMax(d_temp_storage, temp_storage_bytes,\n d_in, d_out, max_segments, d_segment_begin_offsets, d_segment_end_offsets);\n }\n return error;\n}\n\n\n//---------------------------------------------------------------------\n// CUDA nested-parallelism test kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceReduce\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type cub_backend,\n int timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n ReductionOpT reduction_op)\n{\n *d_cdp_error = Dispatch(cub_backend,\n timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items,\n max_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n reduction_op);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Launch kernel and dispatch on device. Should only be called from host code.\n * The CubBackend should be one of the non-CDP CUB backends to invoke from the\n * device.\n */\ntemplate \ncudaError_t LaunchCDPKernel(Int2Type cub_backend,\n int timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n ReductionOpT reduction_op)\n{\n cudaError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n cub_backend,\n timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items,\n max_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n reduction_op);\n CubDebugExit(retval);\n CubDebugExit(cub::detail::device_synchronize());\n\n // Copy out temp_storage_bytes\n CubDebugExit(cudaMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n cudaMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(cudaMemcpy(&retval,\n d_cdp_error,\n sizeof(cudaError_t) * 1,\n cudaMemcpyDeviceToHost));\n\n return retval;\n}\n\n// Specializations of Dispatch that translate the CDP backend to the appropriate\n// CUB backend, and uses the CUB backend to launch the CDP kernel.\n#define DEFINE_CDP_DISPATCHER(CdpBackend, CubBackend) \\\n template \\\n cudaError_t Dispatch(Int2Type, \\\n int timing_iterations, \\\n size_t *d_temp_storage_bytes, \\\n cudaError_t *d_cdp_error, \\\n \\\n void *d_temp_storage, \\\n size_t &temp_storage_bytes, \\\n InputIteratorT d_in, \\\n OutputIteratorT d_out, \\\n int num_items, \\\n int max_segments, \\\n BeginOffsetIteratorT d_segment_begin_offsets, \\\n EndOffsetIteratorT d_segment_end_offsets, \\\n ReductionOpT reduction_op) \\\n { \\\n Int2Type cub_backend{}; \\\n return LaunchCDPKernel(cub_backend, \\\n timing_iterations, \\\n d_temp_storage_bytes, \\\n d_cdp_error, \\\n d_temp_storage, \\\n temp_storage_bytes, \\\n d_in, \\\n d_out, \\\n num_items, \\\n max_segments, \\\n d_segment_begin_offsets, \\\n d_segment_end_offsets, \\\n reduction_op); \\\n }\n\nDEFINE_CDP_DISPATCHER(CDP, CUB)\nDEFINE_CDP_DISPATCHER(CDP_SEGMENTED, CUB_SEGMENTED)\n\n#undef DEFINE_CDP_DISPATCHER\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Problem generation\n//---------------------------------------------------------------------\n/// Initialize problem\ntemplate \nvoid Initialize(\n GenMode gen_mode,\n InputT *h_in,\n int num_items)\n{\n for (int i = 0; i < num_items; ++i)\n {\n InitValue(gen_mode, h_in[i], i);\n }\n\n if (g_verbose_input)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n\n/// Solve problem (max/custom-max functor)\ntemplate \nstruct Solution\n{\n using OutputT = _OutputT;\n using InitT = OutputT;\n using AccumT = cub::detail::accumulator_t;\n\n template \n static void Solve(HostInputIteratorT h_in, OutputT *h_reference, OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets, EndOffsetIteratorT h_segment_end_offsets, ReductionOpT reduction_op)\n {\n for (int i = 0; i < num_segments; ++i)\n {\n AccumT aggregate = Traits::Lowest(); // replace with std::numeric_limits::lowest() when C++ support is more prevalent\n for (int j = h_segment_begin_offsets[i]; j < h_segment_end_offsets[i]; ++j)\n aggregate = reduction_op(aggregate, OutputT(h_in[j]));\n h_reference[i] = aggregate;\n }\n }\n};\n\n/// Solve problem (min functor)\ntemplate \nstruct Solution\n{\n using OutputT = _OutputT;\n using InitT = OutputT;\n using AccumT = cub::detail::accumulator_t;\n\n template \n static void Solve(HostInputIteratorT h_in, OutputT *h_reference, OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets, EndOffsetIteratorT h_segment_end_offsets, cub::Min reduction_op)\n {\n for (int i = 0; i < num_segments; ++i)\n {\n AccumT aggregate = Traits::Max(); // replace with std::numeric_limits::max() when C++ support is more prevalent\n for (int j = h_segment_begin_offsets[i]; j < h_segment_end_offsets[i]; ++j)\n aggregate = reduction_op(aggregate, OutputT(h_in[j]));\n h_reference[i] = aggregate;\n }\n }\n};\n\n\n/// Solve problem (sum functor)\ntemplate \nstruct Solution\n{\n using OutputT = _OutputT;\n using InitT = OutputT;\n using AccumT = cub::detail::accumulator_t;\n\n template \n static void Solve(HostInputIteratorT h_in, OutputT *h_reference, OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets, EndOffsetIteratorT h_segment_end_offsets, ReductionOpT reduction_op)\n {\n for (int i = 0; i < num_segments; ++i)\n {\n AccumT aggregate{};\n for (int j = h_segment_begin_offsets[i]; j < h_segment_end_offsets[i]; ++j)\n aggregate = reduction_op(aggregate, h_in[j]);\n h_reference[i] = static_cast(aggregate);\n }\n }\n};\n\ntemplate \nstruct Solution\n{\n using OutputT = _OutputT;\n using InitT = OutputT;\n using AccumT = cub::detail::accumulator_t;\n\n template \n static void Solve(HostInputIteratorT h_in,\n OutputT *h_reference,\n OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets,\n EndOffsetIteratorT h_segment_end_offsets,\n ReductionOpT reduction_op)\n {\n for (int i = 0; i < num_segments; ++i)\n {\n AccumT aggregate{};\n for (int j = h_segment_begin_offsets[i]; j < h_segment_end_offsets[i]; ++j)\n aggregate = reduction_op(aggregate, h_in[j]);\n h_reference[i] = static_cast(aggregate);\n }\n }\n};\n\n/// Solve problem (argmin functor)\ntemplate \nstruct Solution\n{\n typedef KeyValuePair OutputT;\n\n template \n static void Solve(HostInputIteratorT h_in, OutputT *h_reference, OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets, EndOffsetIteratorT h_segment_end_offsets, cub::ArgMin reduction_op)\n {\n for (int i = 0; i < num_segments; ++i)\n {\n const auto segment_begin = h_segment_begin_offsets[i];\n const auto segment_end = h_segment_end_offsets[i];\n\n if (segment_begin < segment_end) \n {\n OutputT aggregate(0, OutputValueT(h_in[segment_begin]));\n for (int j = segment_begin + 1; j < segment_end; ++j)\n {\n OutputT item(j - segment_begin, OutputValueT(h_in[j]));\n aggregate = reduction_op(aggregate, item);\n }\n h_reference[i] = aggregate;\n }\n else \n {\n // Guaranteed output for empty segments\n OutputT aggregate(1, Traits::Max()); // replace with std::numeric_limits::max() when C++ support is more prevalent\n h_reference[i] = aggregate;\n }\n }\n }\n};\n\n\n/// Solve problem (argmax functor)\ntemplate \nstruct Solution\n{\n typedef KeyValuePair OutputT;\n\n template \n static void Solve(HostInputIteratorT h_in, OutputT *h_reference, OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets, EndOffsetIteratorT h_segment_end_offsets, cub::ArgMax reduction_op)\n {\n for (int i = 0; i < num_segments; ++i)\n {\n const auto segment_begin = h_segment_begin_offsets[i];\n const auto segment_end = h_segment_end_offsets[i];\n\n if (segment_begin < segment_end) \n {\n OutputT aggregate(0, OutputValueT(h_in[segment_begin])); \n for (int j = segment_begin + 1; j < segment_end; ++j)\n {\n OutputT item(j - segment_begin, OutputValueT(h_in[j]));\n aggregate = reduction_op(aggregate, item);\n }\n h_reference[i] = aggregate;\n }\n else \n {\n OutputT aggregate(1, Traits::Lowest());\n h_reference[i] = aggregate;\n }\n }\n }\n};\n\n\n//---------------------------------------------------------------------\n// Problem generation\n//---------------------------------------------------------------------\n\ntemplate \nIt unwrap_it(It it) {\n return it;\n}\n\n#if TEST_HALF_T\n__half* unwrap_it(half_t *it) {\n return reinterpret_cast<__half*>(it);\n}\n\ntemplate \nConstantInputIterator<__half, OffsetT> unwrap_it(ConstantInputIterator it) {\n half_t wrapped_val = *it;\n __half val = wrapped_val.operator __half();\n return ConstantInputIterator<__half, OffsetT>(val);\n}\n#endif\n\n#if TEST_BF_T\n__nv_bfloat16* unwrap_it(bfloat16_t* it) {\n return reinterpret_cast<__nv_bfloat16*>(it);\n}\n\ntemplate \nConstantInputIterator<__nv_bfloat16, OffsetT> unwrap_it(ConstantInputIterator it) {\n bfloat16_t wrapped_val = *it;\n __nv_bfloat16 val = wrapped_val.operator __nv_bfloat16();\n return ConstantInputIterator<__nv_bfloat16, OffsetT>(val);\n}\n#endif\n\ntemplate ()))>\nstd::integral_constant::value> //\nreference_extended_fp(WrappedItT)\n{\n return {};\n}\n\nExtendedFloatSum unwrap_op(std::true_type /* extended float */, cub::Sum) //\n{\n return {};\n}\n\ntemplate \nOpT unwrap_op(std::integral_constant /* base case */, OpT op)\n{\n return op;\n}\n\n/// Test DeviceReduce for a given problem input\ntemplate <\n typename BackendT,\n typename DeviceWrappedInputIteratorT,\n typename DeviceWrappedOutputIteratorT,\n typename HostReferenceIteratorT,\n typename OffsetT,\n typename BeginOffsetIteratorT,\n typename EndOffsetIteratorT,\n typename ReductionOpT>\nvoid Test(\n BackendT backend,\n DeviceWrappedInputIteratorT d_wrapped_in,\n DeviceWrappedOutputIteratorT d_wrapped_out,\n OffsetT num_items,\n OffsetT num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n ReductionOpT reduction_op,\n HostReferenceIteratorT h_reference)\n{\n // Input data types\n auto d_in = unwrap_it(d_wrapped_in);\n auto d_out = unwrap_it(d_wrapped_out);\n\n // Allocate CDP device arrays for temp storage size and error\n size_t *d_temp_storage_bytes = NULL;\n cudaError_t *d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(cudaError_t) * 1));\n\n // Inquire temp device storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(backend, 1,\n d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes,\n d_in, d_out, num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets,\n reduction_op));\n\n // Allocate temp device storage\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(backend, 1,\n d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes,\n d_in, d_out, num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets,\n reduction_op));\n\n // Check for correctness (and display results, if specified)\n int compare = CompareDeviceResults(h_reference, d_wrapped_out, num_segments, g_verbose, g_verbose);\n\n printf(\"\\t%s\", compare ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Performance\n if (g_timing_iterations > 0)\n {\n GpuTimer gpu_timer;\n gpu_timer.Start();\n\n CubDebugExit(Dispatch(backend, g_timing_iterations,\n d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes,\n d_in, d_out, num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets,\n reduction_op));\n\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n\n // Display performance\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n float giga_bandwidth = giga_rate * sizeof(h_reference[0]);\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s, %.1f%% peak\",\n avg_millis, giga_rate, giga_bandwidth, giga_bandwidth / g_device_giga_bandwidth * 100.0);\n\n }\n\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare);\n}\n\ntemplate \nvoid SolveAndTest(HostInputIteratorT h_in,\n DeviceInputIteratorT d_in,\n OffsetT num_items,\n OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets,\n EndOffsetIteratorT h_segment_end_offsets,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n ReductionOpT wrapped_reduction_op)\n{\n auto reduction_op = unwrap_op(reference_extended_fp(d_in), wrapped_reduction_op);\n\n using InputValueT = cub::detail::value_t;\n using SolutionT = Solution;\n using OutputT = typename SolutionT::OutputT;\n\n printf(\"\\n\\n%s cub::DeviceReduce<%s> %d items (%s), %d segments\\n\",\n BackendToString(BACKEND),\n typeid(ReductionOpT).name(),\n num_items,\n typeid(HostInputIteratorT).name(),\n num_segments);\n fflush(stdout);\n\n // Allocate and solve solution\n OutputT *h_reference = new OutputT[num_segments];\n SolutionT::Solve(h_in, h_reference, num_segments, h_segment_begin_offsets, h_segment_end_offsets, reduction_op);\n\n // Run with discard iterator\n DiscardOutputIterator discard_itr;\n Test(Int2Type(), d_in, discard_itr, num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets, reduction_op, h_reference);\n\n // Run with output data\n OutputT *d_out = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(OutputT) * num_segments));\n CubDebugExit(cudaMemset(d_out, 0, sizeof(OutputT) * num_segments));\n Test(Int2Type(), d_in, d_out, num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets, reduction_op, h_reference);\n\n // Cleanup\n if (d_out) CubDebugExit(g_allocator.DeviceFree(d_out));\n if (h_reference) delete[] h_reference;\n}\n\n/// Test specific problem type\ntemplate <\n Backend BACKEND,\n typename InputT,\n typename OutputT,\n typename OffsetT,\n typename ReductionOpT>\nvoid TestProblem(\n OffsetT num_items,\n OffsetT num_segments,\n GenMode gen_mode,\n ReductionOpT reduction_op)\n{\n printf(\"\\n\\nInitializing %d %s->%s (gen mode %d)... \", num_items, typeid(InputT).name(), typeid(OutputT).name(), gen_mode); fflush(stdout);\n fflush(stdout);\n\n // Initialize value data\n InputT* h_in = new InputT[num_items];\n Initialize(gen_mode, h_in, num_items);\n\n // Initialize segment data\n OffsetT *h_segment_offsets = new OffsetT[num_segments + 1];\n InitializeSegments(num_items, num_segments, h_segment_offsets, g_verbose_input);\n\n // Initialize device data\n OffsetT *d_segment_offsets = NULL;\n InputT *d_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_in, sizeof(InputT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_segment_offsets, sizeof(OffsetT) * (num_segments + 1)));\n CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(InputT) * num_items, cudaMemcpyHostToDevice));\n CubDebugExit(cudaMemcpy(d_segment_offsets, h_segment_offsets, sizeof(OffsetT) * (num_segments + 1), cudaMemcpyHostToDevice));\n\n SolveAndTest(h_in, d_in, num_items, num_segments,\n h_segment_offsets, h_segment_offsets + 1, d_segment_offsets, d_segment_offsets + 1, reduction_op);\n\n if (h_segment_offsets) delete[] h_segment_offsets;\n if (d_segment_offsets) CubDebugExit(g_allocator.DeviceFree(d_segment_offsets));\n if (h_in) delete[] h_in;\n if (d_in) CubDebugExit(g_allocator.DeviceFree(d_in));\n}\n\n\n/// Test different operators\ntemplate <\n Backend BACKEND,\n typename OutputT,\n typename HostInputIteratorT,\n typename DeviceInputIteratorT,\n typename OffsetT,\n typename BeginOffsetIteratorT,\n typename EndOffsetIteratorT>\nvoid TestByOp(\n HostInputIteratorT h_in,\n DeviceInputIteratorT d_in,\n OffsetT num_items,\n OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets,\n EndOffsetIteratorT h_segment_end_offsets,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets)\n{\n SolveAndTest(h_in, d_in, num_items, num_segments,\n h_segment_begin_offsets, h_segment_end_offsets, d_segment_begin_offsets, d_segment_end_offsets, CustomMax());\n SolveAndTest(h_in, d_in, num_items, num_segments,\n h_segment_begin_offsets, h_segment_end_offsets, d_segment_begin_offsets, d_segment_end_offsets, Sum());\n SolveAndTest(h_in, d_in, num_items, num_segments,\n h_segment_begin_offsets, h_segment_end_offsets, d_segment_begin_offsets, d_segment_end_offsets, Min());\n SolveAndTest(h_in, d_in, num_items, num_segments,\n h_segment_begin_offsets, h_segment_end_offsets, d_segment_begin_offsets, d_segment_end_offsets, ArgMin());\n SolveAndTest(h_in, d_in, num_items, num_segments,\n h_segment_begin_offsets, h_segment_end_offsets, d_segment_begin_offsets, d_segment_end_offsets, Max());\n SolveAndTest(h_in, d_in, num_items, num_segments,\n h_segment_begin_offsets, h_segment_end_offsets, d_segment_begin_offsets, d_segment_end_offsets, ArgMax());\n}\n\ntemplate\nstruct TransformFunctor1\n{\n __host__ __device__ __forceinline__ OffsetT operator()(OffsetT offset) const\n {\n return offset;\n }\n};\n\ntemplate\nstruct TransformFunctor2\n{\n __host__ __device__ __forceinline__ OffsetT operator()(OffsetT offset) const\n {\n return offset;\n }\n};\n\n/// Test different backends\ntemplate <\n typename InputT,\n typename OutputT,\n typename OffsetT>\nvoid TestByBackend(\n OffsetT num_items,\n OffsetT max_segments,\n GenMode gen_mode)\n{\n#if TEST_CDP == 0\n constexpr auto NonSegmentedBackend = CUB;\n constexpr auto SegmentedBackend = CUB_SEGMENTED;\n#else // TEST_CDP\n constexpr auto NonSegmentedBackend = CDP;\n constexpr auto SegmentedBackend = CDP_SEGMENTED;\n#endif // TEST_CDP\n\n // Initialize host data\n printf(\"\\n\\nInitializing %d %s -> %s (gen mode %d)... \",\n num_items, typeid(InputT).name(), typeid(OutputT).name(), gen_mode); fflush(stdout);\n\n InputT *h_in = new InputT[num_items];\n OffsetT *h_segment_offsets = new OffsetT[max_segments + 1];\n Initialize(gen_mode, h_in, num_items);\n\n // Initialize device data\n InputT *d_in = NULL;\n OffsetT *d_segment_offsets = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_in, sizeof(InputT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_segment_offsets, sizeof(OffsetT) * (max_segments + 1)));\n CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(InputT) * num_items, cudaMemcpyHostToDevice));\n\n //\n // Test single-segment implementations\n //\n\n InitializeSegments(num_items, 1, h_segment_offsets, g_verbose_input);\n\n // Page-aligned-input tests\n TestByOp(h_in, d_in, num_items, 1,\n h_segment_offsets, h_segment_offsets + 1, (OffsetT*) NULL, (OffsetT*)NULL);\n\n // Non-page-aligned-input tests\n if (num_items > 1)\n {\n InitializeSegments(num_items - 1, 1, h_segment_offsets, g_verbose_input);\n TestByOp(h_in + 1, d_in + 1, num_items - 1, 1,\n h_segment_offsets, h_segment_offsets + 1, (OffsetT*) NULL, (OffsetT*)NULL);\n }\n\n //\n // Test segmented implementation\n //\n\n // Right now we assign a single thread block to each segment, so lets keep it to under 128K items per segment\n int max_items_per_segment = 128000;\n\n for (int num_segments = cub::DivideAndRoundUp(num_items, max_items_per_segment);\n num_segments < max_segments;\n num_segments = (num_segments * 32) + 1)\n {\n // Test with segment pointer\n InitializeSegments(num_items, num_segments, h_segment_offsets, g_verbose_input);\n CubDebugExit(cudaMemcpy(d_segment_offsets, h_segment_offsets, sizeof(OffsetT) * (num_segments + 1), cudaMemcpyHostToDevice));\n TestByOp(h_in, d_in, num_items, num_segments,\n h_segment_offsets, h_segment_offsets + 1, d_segment_offsets, d_segment_offsets + 1);\n\n // Test with segment iterator\n typedef CastOp IdentityOpT;\n IdentityOpT identity_op;\n TransformInputIterator h_segment_offsets_itr(\n h_segment_offsets,\n identity_op);\n TransformInputIterator d_segment_offsets_itr(\n d_segment_offsets,\n identity_op);\n\n TestByOp(h_in, d_in, num_items, num_segments,\n h_segment_offsets_itr, h_segment_offsets_itr + 1, d_segment_offsets_itr, d_segment_offsets_itr + 1);\n\n // Test with transform iterators of different types\n\n typedef TransformFunctor1 TransformFunctor1T;\n typedef TransformFunctor2 TransformFunctor2T;\n\n TransformInputIterator h_segment_begin_offsets_itr(h_segment_offsets, TransformFunctor1T());\n TransformInputIterator h_segment_end_offsets_itr(h_segment_offsets + 1, TransformFunctor2T());\n\n TransformInputIterator d_segment_begin_offsets_itr(d_segment_offsets, TransformFunctor1T());\n TransformInputIterator d_segment_end_offsets_itr(d_segment_offsets + 1, TransformFunctor2T());\n\n TestByOp(h_in, d_in, num_items, num_segments,\n h_segment_begin_offsets_itr, h_segment_end_offsets_itr,\n d_segment_begin_offsets_itr, d_segment_end_offsets_itr);\n }\n\n if (h_in) delete[] h_in;\n if (h_segment_offsets) delete[] h_segment_offsets;\n if (d_in) CubDebugExit(g_allocator.DeviceFree(d_in));\n if (d_segment_offsets) CubDebugExit(g_allocator.DeviceFree(d_segment_offsets));\n}\n\n\n/// Test different input-generation modes\ntemplate <\n typename InputT,\n typename OutputT,\n typename OffsetT>\nvoid TestByGenMode(\n OffsetT num_items,\n OffsetT max_segments)\n{\n //\n // Test pointer support using different input-generation modes\n //\n\n TestByBackend(num_items, max_segments, UNIFORM);\n TestByBackend(num_items, max_segments, INTEGER_SEED);\n TestByBackend(num_items, max_segments, RANDOM);\n\n //\n // Test iterator support using a constant-iterator and SUM\n //\n\n InputT val;\n InitValue(UNIFORM, val, 0);\n ConstantInputIterator in(val);\n\n OffsetT *h_segment_offsets = new OffsetT[1 + 1];\n InitializeSegments(num_items, 1, h_segment_offsets, g_verbose_input);\n\n#if TEST_CDP == 0\n constexpr auto Backend = CUB;\n#else // TEST_CDP\n constexpr auto Backend = CDP;\n#endif // TEST_CDP\n\n SolveAndTest(in, in, num_items, 1,\n h_segment_offsets, h_segment_offsets + 1, (OffsetT*) NULL, (OffsetT*)NULL, Sum());\n\n if (h_segment_offsets) delete[] h_segment_offsets;\n}\n\n/// Test different problem sizes\ntemplate \nvoid TestBySize(OffsetT max_items, OffsetT max_segments, OffsetT tile_size)\n{\n // Test 0, 1, many\n TestByGenMode(0, max_segments);\n TestByGenMode(1, max_segments);\n TestByGenMode(max_items, max_segments);\n\n // Test random problem sizes from a log-distribution [8, max_items-ish)\n int num_iterations = 8;\n double max_exp = log(double(max_items)) / log(double(2.0));\n for (int i = 0; i < num_iterations; ++i)\n {\n OffsetT num_items = (OffsetT)pow(2.0, RandomValue(max_exp - 3.0) + 3.0);\n TestByGenMode(num_items, max_segments);\n }\n\n //\n // White-box testing of single-segment problems around specific sizes\n //\n\n#if TEST_CDP == 0\n constexpr auto Backend = CUB;\n#else // TEST_CDP\n constexpr auto Backend = CDP;\n#endif // TEST_CDP\n\n // Tile-boundaries: multiple blocks, one tile per block\n TestProblem(tile_size * 4, 1, RANDOM, Sum());\n TestProblem(tile_size * 4 + 1, 1, RANDOM, Sum());\n TestProblem(tile_size * 4 - 1, 1, RANDOM, Sum());\n\n // Tile-boundaries: multiple blocks, multiple tiles per block\n OffsetT sm_occupancy = 32;\n OffsetT occupancy = tile_size * sm_occupancy * g_sm_count;\n TestProblem(occupancy, 1, RANDOM, Sum());\n TestProblem(occupancy + 1, 1, RANDOM, Sum());\n TestProblem(occupancy - 1, 1, RANDOM, Sum());\n};\n\nclass CustomInputT\n{\n char m_val{};\n\npublic:\n __host__ __device__ explicit CustomInputT(char val)\n : m_val(val)\n {}\n\n __host__ __device__ int get() const { return static_cast(m_val); }\n};\n\nclass CustomAccumulatorT\n{\n int m_val{0};\n int m_magic_value{42};\n\n __host__ __device__ CustomAccumulatorT(int val)\n : m_val(val)\n {}\n\npublic:\n __host__ __device__ CustomAccumulatorT()\n {}\n\n __host__ __device__ CustomAccumulatorT(const CustomAccumulatorT &in)\n : m_val(in.is_valid() * in.get())\n , m_magic_value(in.is_valid() * 42)\n {}\n\n __host__ __device__ void operator=(const CustomInputT &in)\n {\n if (this->is_valid())\n {\n m_val = in.get();\n }\n }\n\n __host__ __device__ void operator=(const CustomAccumulatorT &in)\n {\n if (this->is_valid() && in.is_valid())\n {\n m_val = in.get();\n }\n }\n\n __host__ __device__ CustomAccumulatorT \n operator+(const CustomInputT &in) const\n {\n const int multiplier = this->is_valid();\n return {(m_val + in.get()) * multiplier};\n }\n\n __host__ __device__ CustomAccumulatorT\n operator+(const CustomAccumulatorT &in) const\n {\n const int multiplier = this->is_valid() && in.is_valid();\n return {(m_val + in.get()) * multiplier};\n }\n\n __host__ __device__ int get() const { return m_val; }\n\n __host__ __device__ bool is_valid() const { return m_magic_value == 42; }\n};\n\nclass CustomOutputT\n{\n bool *m_d_flag{};\n int m_expected{};\n\npublic:\n __host__ __device__ CustomOutputT(bool *d_flag, int expected)\n : m_d_flag(d_flag)\n , m_expected(expected)\n {}\n\n __host__ __device__ void operator=(const CustomAccumulatorT &accum) const\n {\n *m_d_flag = accum.is_valid() && (accum.get() == m_expected);\n }\n};\n\n__global__ void InitializeTestAccumulatorTypes(int num_items,\n int expected,\n bool *d_flag,\n CustomInputT *d_in,\n CustomOutputT *d_out)\n{\n const int idx = static_cast(threadIdx.x + blockIdx.x * blockDim.x);\n\n if (idx < num_items)\n {\n d_in[idx] = CustomInputT(1);\n }\n\n if (idx == 0)\n {\n *d_out = CustomOutputT{d_flag, expected};\n }\n}\n\ntemplate \nvoid TestBigIndicesHelper(OffsetT num_items)\n{\n thrust::constant_iterator const_iter(T{1});\n thrust::device_vector out(1);\n std::size_t* d_out = thrust::raw_pointer_cast(out.data());\n\n std::uint8_t *d_temp_storage{};\n std::size_t temp_storage_bytes{};\n\n CubDebugExit(\n cub::DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, const_iter, d_out, num_items));\n\n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n CubDebugExit(\n cub::DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, const_iter, d_out, num_items));\n std::size_t result = out[0];\n\n AssertEquals(result, num_items);\n}\n\ntemplate \nvoid TestBigIndices()\n{\n TestBigIndicesHelper(1ull << 30);\n TestBigIndicesHelper(1ull << 31);\n TestBigIndicesHelper((1ull << 32) - 1);\n TestBigIndicesHelper(1ull << 33);\n}\n\n#if TEST_TYPES == 3\nvoid TestAccumulatorTypes()\n{\n const int num_items = 2 * 1024 * 1024;\n const int expected = num_items;\n const int block_size = 256;\n const int grid_size = (num_items + block_size - 1) / block_size;\n\n CustomInputT *d_in{};\n CustomOutputT *d_out{};\n CustomAccumulatorT init{};\n bool *d_flag{};\n\n CubDebugExit(\n g_allocator.DeviceAllocate((void **)&d_out, sizeof(CustomOutputT)));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_flag, sizeof(bool)));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_in,\n sizeof(CustomInputT) * num_items));\n\n InitializeTestAccumulatorTypes<<>>(num_items,\n expected,\n d_flag,\n d_in,\n d_out);\n\n std::uint8_t *d_temp_storage{};\n std::size_t temp_storage_bytes{};\n\n CubDebugExit(cub::DeviceReduce::Reduce(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items,\n cub::Sum{},\n init));\n\n CubDebugExit(\n g_allocator.DeviceAllocate((void **)&d_temp_storage, temp_storage_bytes));\n CubDebugExit(cudaMemset(d_temp_storage, 1, temp_storage_bytes));\n\n CubDebugExit(cub::DeviceReduce::Reduce(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items,\n cub::Sum{},\n init));\n\n bool ok{};\n CubDebugExit(cudaMemcpy(&ok, d_flag, sizeof(bool), cudaMemcpyDeviceToHost));\n\n AssertTrue(ok);\n\n CubDebugExit(g_allocator.DeviceFree(d_out));\n CubDebugExit(g_allocator.DeviceFree(d_in));\n}\n\n/**\n * ArgMin should return max value for empty input. This interferes with\n * input data containing infinity values. This test checks that ArgMin\n * works correctly with infinity values.\n */\nvoid TestFloatInfInArgMin()\n{\n using in_t = float;\n using offset_t = int;\n using out_t = cub::KeyValuePair;\n\n const int n = 10;\n const float inf = ::cuda::std::numeric_limits::infinity();\n\n thrust::device_vector in(n, inf);\n thrust::device_vector out(1);\n\n const in_t *d_in = thrust::raw_pointer_cast(in.data());\n out_t *d_out = thrust::raw_pointer_cast(out.data());\n\n std::uint8_t *d_temp_storage{};\n std::size_t temp_storage_bytes{};\n\n CubDebugExit(cub::DeviceReduce::ArgMin(d_temp_storage, temp_storage_bytes, d_in, d_out, n));\n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n CubDebugExit(cub::DeviceReduce::ArgMin(d_temp_storage, temp_storage_bytes, d_in, d_out, n));\n\n const out_t result = out[0];\n AssertEquals(result.key, 0);\n AssertEquals(result.value, inf);\n}\n\n/**\n * ArgMax should return lowest value for empty input. This interferes with\n * input data containing infinity values. This test checks that ArgMax\n * works correctly with infinity values.\n */\nvoid TestFloatInfInArgMax()\n{\n using in_t = float;\n using offset_t = int;\n using out_t = cub::KeyValuePair;\n\n const int n = 10;\n const float inf = ::cuda::std::numeric_limits::infinity();\n \n thrust::device_vector in(n, -inf);\n thrust::device_vector out(1);\n\n const in_t *d_in = thrust::raw_pointer_cast(in.data());\n out_t *d_out = thrust::raw_pointer_cast(out.data());\n\n std::uint8_t *d_temp_storage{};\n std::size_t temp_storage_bytes{}; \n\n CubDebugExit(cub::DeviceReduce::ArgMax(d_temp_storage, temp_storage_bytes, d_in, d_out, n));\n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n CubDebugExit(cub::DeviceReduce::ArgMax(d_temp_storage, temp_storage_bytes, d_in, d_out, n));\n\n const out_t result = out[0];\n AssertEquals(result.key, 0);\n AssertEquals(result.value, -inf);\n}\n\nvoid TestFloatInfInArg()\n{\n TestFloatInfInArgMin();\n TestFloatInfInArgMax();\n}\n#endif\n\ntemplate \nstruct GetTileSize\n{\n OffsetT max_items{};\n OffsetT max_segments{};\n OffsetT tile_size{};\n\n GetTileSize(OffsetT max_items, OffsetT max_segments)\n : max_items(max_items)\n , max_segments(max_segments)\n {}\n\n template \n CUB_RUNTIME_FUNCTION cudaError_t Invoke()\n {\n this->tile_size = ActivePolicyT::ReducePolicy::BLOCK_THREADS *\n ActivePolicyT::ReducePolicy::ITEMS_PER_THREAD;\n return cudaSuccess;\n }\n};\n\n/// Test problem type\ntemplate \nvoid TestType(OffsetT max_items, OffsetT max_segments)\n{\n // Inspect the tuning policies to determine this arch's tile size:\n using MaxPolicyT =\n typename DeviceReducePolicy::MaxPolicy;\n GetTileSize dispatch(max_items, max_segments);\n CubDebugExit(MaxPolicyT::Invoke(g_ptx_version, dispatch));\n\n TestBySize(max_items, max_segments, dispatch.tile_size);\n}\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n typedef int OffsetT;\n\n OffsetT max_items = 27000000;\n OffsetT max_segments = 34000;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n g_verbose_input = args.CheckCmdLineFlag(\"v2\");\n args.GetCmdLineArgument(\"n\", max_items);\n args.GetCmdLineArgument(\"s\", max_segments);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--s= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n g_device_giga_bandwidth = args.device_giga_bandwidth;\n\n // Get ptx version\n CubDebugExit(PtxVersion(g_ptx_version));\n\n // Get SM count\n g_sm_count = args.deviceProp.multiProcessorCount;\n\n // %PARAM% TEST_CDP cdp 0:1\n // %PARAM% TEST_TYPES types 0:1:2:3:4\n\n#if TEST_TYPES == 0\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n#elif TEST_TYPES == 1\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n#elif TEST_TYPES == 2\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n#elif TEST_TYPES == 3\n TestType(max_items, max_segments);\n TestAccumulatorTypes();\n TestFloatInfInArg();\n\n#if TEST_HALF_T\n TestType(max_items, max_segments);\n#endif\n#else // TEST_TYPES == 4\n TestType(max_items, max_segments);\n TestBigIndices();\n\n#if TEST_BF_T\n TestType(max_items, max_segments);\n#endif\n#endif\n printf(\"\\n\");\n return 0;\n}\n\n\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceReduce utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"test_util.h\"\n#include \n\n#define TEST_HALF_T !_NVHPC_CUDA\n\n#define TEST_BF_T !_NVHPC_CUDA\n\n#if TEST_HALF_T\n#include \n\n// Half support is provided by SM53+. We currently test against a few older architectures. \n// The specializations below can be removed once we drop these architectures. \nnamespace cub {\n\ntemplate <>\n__host__ __device__ __forceinline__ //\n__half Min::operator()(__half& a, __half& b) const \n{\n NV_IF_TARGET(NV_PROVIDES_SM_53, \n (return CUB_MIN(a, b);),\n (return CUB_MIN(__half2float(a), __half2float(b));));\n}\n\ntemplate <>\n__host__ __device__ __forceinline__ //\nKeyValuePair \nArgMin::operator()(const KeyValuePair &a, \n const KeyValuePair &b) const \n{\n const float av = __half2float(a.value);\n const float bv = __half2float(b.value);\n\n if ((bv < av) || ((av == bv) && (b.key < a.key)))\n {\n return b;\n }\n\n return a;\n}\n\ntemplate <>\n__host__ __device__ __forceinline__ //\n__half Max::operator()(__half& a, __half& b) const \n{\n NV_IF_TARGET(NV_PROVIDES_SM_53, \n (return CUB_MAX(a, b);),\n (return CUB_MAX(__half2float(a), __half2float(b));));\n}\n\ntemplate <>\n__host__ __device__ __forceinline__ //\nKeyValuePair \nArgMax::operator()(const KeyValuePair &a, \n const KeyValuePair &b) const \n{\n const float av = __half2float(a.value);\n const float bv = __half2float(b.value);\n\n if ((bv > av) || ((av == bv) && (b.key < a.key)))\n {\n return b;\n }\n\n return a;\n}\n\n} // namespace cub\n#endif\n\n#if TEST_BF_T\n#include \n#endif\n\nusing namespace hipcub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nint g_ptx_version;\nint g_sm_count;\ndouble g_device_giga_bandwidth;\nbool g_verbose = false;\nbool g_verbose_input = false;\nint g_timing_iterations = 0;\nCachingDeviceAllocator g_allocator(true);\n\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CUB_SEGMENTED, // CUB segmented method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n CDP_SEGMENTED, // GPU-based segmented method\n};\n\ninline const char* BackendToString(Backend b)\n{\n switch (b)\n {\n case CUB:\n return \"CUB\";\n case CUB_SEGMENTED:\n return \"CUB_SEGMENTED\";\n case CDP:\n return \"CDP\";\n case CDP_SEGMENTED:\n return \"CDP_SEGMENTED\";\n default:\n break;\n }\n\n return \"\";\n}\n\n// Custom max functor\nstruct CustomMax\n{\n /// Boolean max operator, returns (a > b) ? a : b\n template \n __host__ __device__ auto operator()(T&& a, C&& b) \n -> hipcub::detail::accumulator_t\n {\n return CUB_MAX(a, b);\n }\n\n#if TEST_HALF_T\n __host__ __device__ __half operator()(__half& a, __half& b) \n {\n return hipcub::Max{}(a, b);\n }\n#endif\n};\n\n// Comparing results computed on CPU and GPU for extended floating point types is impossible. \n// For instance, when used with a constant iterator of two, the accumulator in sequential reference \n// computation (CPU) bumps into the 4096 limits, which will never change (`4096 + 2 = 4096`). \n// Meanwhile, per-thread aggregates (`2 * 16 = 32`) are accumulated within and among thread blocks, \n// yielding `inf` as a result. No reasonable epsilon can be selected to compare `inf` with `4096`. \n// To make `__half` and `__nv_bfloat16` arithmetic associative, the function object below raises \n// extended floating points to the area of unsigned short integers. This allows us to test large \n// inputs with few code-path differences in device algorithms. \nstruct ExtendedFloatSum\n{\n template \n __host__ __device__ T operator()(T a, T b) const\n {\n T result{};\n result.__x = a.raw() + b.raw();\n return result;\n }\n\n#if TEST_HALF_T\n __host__ __device__ __half operator()(__half a, __half b) const\n {\n uint16_t result = this->operator()(half_t{a}, half_t(b)).raw();\n return reinterpret_cast<__half &>(result);\n }\n#endif\n\n#if TEST_BF_T\n __device__ __hip_bfloat16 operator()(__hip_bfloat16 a, __hip_bfloat16 b) const\n {\n uint16_t result = this->operator()(bfloat16_t{a}, bfloat16_t(b)).raw();\n return reinterpret_cast<__hip_bfloat16 &>(result);\n }\n#endif\n};\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB DeviceReduce entrypoints\n//---------------------------------------------------------------------\n\n/**\n * Dispatch to reduce entrypoint (custom-max)\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int /*max_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n ReductionOpT reduction_op)\n{\n using InputT = hipcub::detail::value_t;\n\n // The output value type\n using OutputT = hipcub::detail::non_void_value_t;\n\n // Max-identity\n OutputT identity = Traits::Lowest(); // replace with std::numeric_limits::lowest() when C++ support is more prevalent\n\n // Invoke kernel to device reduction directly\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceReduce::Reduce(d_temp_storage, temp_storage_bytes,\n d_in, d_out, num_items, reduction_op, identity);\n }\n\n return error;\n}\n\n/**\n * Dispatch to sum entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n hipError_t * /*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int /*max_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n hipcub::Sum /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items);\n }\n\n return error;\n}\n\n/**\n * Dispatch to extended fp sum entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION hipError_t Dispatch(Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n hipError_t * /*d_cdp_error*/,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int /*max_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n ExtendedFloatSum reduction_op)\n{\n using InputT = hipcub::detail::value_t;\n using OutputT = hipcub::detail::non_void_value_t;\n\n OutputT identity{};\n\n // Invoke kernel to device reduction directly\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceReduce::Reduce(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items,\n reduction_op,\n identity);\n }\n\n return error;\n}\n\n/**\n * Dispatch to min entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int /*max_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n hipcub::Min /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceReduce::Min(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items);\n }\n\n return error;\n}\n\n/**\n * Dispatch to max entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int /*max_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n hipcub::Max /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceReduce::Max(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items);\n }\n\n return error;\n}\n\n/**\n * Dispatch to argmin entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int /*max_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n hipcub::ArgMin /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceReduce::ArgMin(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items);\n }\n\n return error;\n}\n\n/**\n * Dispatch to argmax entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int /*max_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n hipcub::ArgMax /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceReduce::ArgMax(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items);\n }\n\n return error;\n}\n\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB DeviceSegmentedReduce entrypoints\n//---------------------------------------------------------------------\n\n/**\n * Dispatch to reduce entrypoint (custom-max)\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int /*num_items*/,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n ReductionOpT reduction_op)\n{\n // The input value type\n using InputT = hipcub::detail::value_t;\n\n // The output value type\n using OutputT = hipcub::detail::non_void_value_t;\n\n // Max-identity\n OutputT identity = Traits::Lowest(); // replace with std::numeric_limits::lowest() when C++ support is more prevalent\n\n // Invoke kernel to device reduction directly\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceSegmentedReduce::Reduce(d_temp_storage, temp_storage_bytes,\n d_in, d_out, max_segments, d_segment_begin_offsets, d_segment_end_offsets, reduction_op, identity);\n }\n return error;\n}\n\n/**\n * Dispatch to sum entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int /*num_items*/,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n hipcub::Sum /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceSegmentedReduce::Sum(d_temp_storage, temp_storage_bytes,\n d_in, d_out, max_segments, d_segment_begin_offsets, d_segment_end_offsets);\n }\n return error;\n}\n\n/**\n * Dispatch to extended fp sum entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION hipError_t Dispatch(Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n hipError_t * /*d_cdp_error*/,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int /*num_items*/,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n ExtendedFloatSum reduction_op)\n{\n using InputT = hipcub::detail::value_t;\n using OutputT = hipcub::detail::non_void_value_t;\n\n OutputT identity{};\n\n // Invoke kernel to device reduction directly\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceSegmentedReduce::Reduce(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n max_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n reduction_op,\n identity);\n }\n\n return error;\n}\n\n/**\n * Dispatch to min entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int /*num_items*/,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n hipcub::Min /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceSegmentedReduce::Min(d_temp_storage, temp_storage_bytes,\n d_in, d_out, max_segments, d_segment_begin_offsets, d_segment_end_offsets);\n }\n return error;\n}\n\n/**\n * Dispatch to max entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int /*num_items*/,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n hipcub::Max /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceSegmentedReduce::Max(d_temp_storage, temp_storage_bytes,\n d_in, d_out, max_segments, d_segment_begin_offsets, d_segment_end_offsets);\n }\n return error;\n}\n\n/**\n * Dispatch to argmin entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int /*num_items*/,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n hipcub::ArgMin /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceSegmentedReduce::ArgMin(d_temp_storage, temp_storage_bytes,\n d_in, d_out, max_segments, d_segment_begin_offsets, d_segment_end_offsets);\n }\n return error;\n}\n\n/**\n * Dispatch to argmax entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int /*num_items*/,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n hipcub::ArgMax /*reduction_op*/)\n{\n // Invoke kernel to device reduction directly\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_iterations; ++i)\n {\n error = DeviceSegmentedReduce::ArgMax(d_temp_storage, temp_storage_bytes,\n d_in, d_out, max_segments, d_segment_begin_offsets, d_segment_end_offsets);\n }\n return error;\n}\n\n\n//---------------------------------------------------------------------\n// CUDA nested-parallelism test kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceReduce\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type cub_backend,\n int timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n ReductionOpT reduction_op)\n{\n *d_cdp_error = Dispatch(cub_backend,\n timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items,\n max_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n reduction_op);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Launch kernel and dispatch on device. Should only be called from host code.\n * The CubBackend should be one of the non-CDP CUB backends to invoke from the\n * device.\n */\ntemplate \nhipError_t LaunchCDPKernel(Int2Type cub_backend,\n int timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n int num_items,\n int max_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n ReductionOpT reduction_op)\n{\n hipError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n cub_backend,\n timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items,\n max_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n reduction_op);\n CubDebugExit(retval);\n CubDebugExit(hipcub::detail::device_synchronize());\n\n // Copy out temp_storage_bytes\n CubDebugExit(hipMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n hipMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(hipMemcpy(&retval,\n d_cdp_error,\n sizeof(hipError_t) * 1,\n hipMemcpyDeviceToHost));\n\n return retval;\n}\n\n// Specializations of Dispatch that translate the CDP backend to the appropriate\n// CUB backend, and uses the CUB backend to launch the CDP kernel.\n#define DEFINE_CDP_DISPATCHER(CdpBackend, CubBackend) \\\n template \\\n hipError_t Dispatch(Int2Type, \\\n int timing_iterations, \\\n size_t *d_temp_storage_bytes, \\\n hipError_t *d_cdp_error, \\\n \\\n void *d_temp_storage, \\\n size_t &temp_storage_bytes, \\\n InputIteratorT d_in, \\\n OutputIteratorT d_out, \\\n int num_items, \\\n int max_segments, \\\n BeginOffsetIteratorT d_segment_begin_offsets, \\\n EndOffsetIteratorT d_segment_end_offsets, \\\n ReductionOpT reduction_op) \\\n { \\\n Int2Type cub_backend{}; \\\n return LaunchCDPKernel(cub_backend, \\\n timing_iterations, \\\n d_temp_storage_bytes, \\\n d_cdp_error, \\\n d_temp_storage, \\\n temp_storage_bytes, \\\n d_in, \\\n d_out, \\\n num_items, \\\n max_segments, \\\n d_segment_begin_offsets, \\\n d_segment_end_offsets, \\\n reduction_op); \\\n }\n\nDEFINE_CDP_DISPATCHER(CDP, CUB)\nDEFINE_CDP_DISPATCHER(CDP_SEGMENTED, CUB_SEGMENTED)\n\n#undef DEFINE_CDP_DISPATCHER\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Problem generation\n//---------------------------------------------------------------------\n/// Initialize problem\ntemplate \nvoid Initialize(\n GenMode gen_mode,\n InputT *h_in,\n int num_items)\n{\n for (int i = 0; i < num_items; ++i)\n {\n InitValue(gen_mode, h_in[i], i);\n }\n\n if (g_verbose_input)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n\n/// Solve problem (max/custom-max functor)\ntemplate \nstruct Solution\n{\n using OutputT = _OutputT;\n using InitT = OutputT;\n using AccumT = hipcub::detail::accumulator_t;\n\n template \n static void Solve(HostInputIteratorT h_in, OutputT *h_reference, OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets, EndOffsetIteratorT h_segment_end_offsets, ReductionOpT reduction_op)\n {\n for (int i = 0; i < num_segments; ++i)\n {\n AccumT aggregate = Traits::Lowest(); // replace with std::numeric_limits::lowest() when C++ support is more prevalent\n for (int j = h_segment_begin_offsets[i]; j < h_segment_end_offsets[i]; ++j)\n aggregate = reduction_op(aggregate, OutputT(h_in[j]));\n h_reference[i] = aggregate;\n }\n }\n};\n\n/// Solve problem (min functor)\ntemplate \nstruct Solution\n{\n using OutputT = _OutputT;\n using InitT = OutputT;\n using AccumT = hipcub::detail::accumulator_t;\n\n template \n static void Solve(HostInputIteratorT h_in, OutputT *h_reference, OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets, EndOffsetIteratorT h_segment_end_offsets, hipcub::Min reduction_op)\n {\n for (int i = 0; i < num_segments; ++i)\n {\n AccumT aggregate = Traits::Max(); // replace with std::numeric_limits::max() when C++ support is more prevalent\n for (int j = h_segment_begin_offsets[i]; j < h_segment_end_offsets[i]; ++j)\n aggregate = reduction_op(aggregate, OutputT(h_in[j]));\n h_reference[i] = aggregate;\n }\n }\n};\n\n\n/// Solve problem (sum functor)\ntemplate \nstruct Solution\n{\n using OutputT = _OutputT;\n using InitT = OutputT;\n using AccumT = hipcub::detail::accumulator_t;\n\n template \n static void Solve(HostInputIteratorT h_in, OutputT *h_reference, OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets, EndOffsetIteratorT h_segment_end_offsets, ReductionOpT reduction_op)\n {\n for (int i = 0; i < num_segments; ++i)\n {\n AccumT aggregate{};\n for (int j = h_segment_begin_offsets[i]; j < h_segment_end_offsets[i]; ++j)\n aggregate = reduction_op(aggregate, h_in[j]);\n h_reference[i] = static_cast(aggregate);\n }\n }\n};\n\ntemplate \nstruct Solution\n{\n using OutputT = _OutputT;\n using InitT = OutputT;\n using AccumT = hipcub::detail::accumulator_t;\n\n template \n static void Solve(HostInputIteratorT h_in,\n OutputT *h_reference,\n OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets,\n EndOffsetIteratorT h_segment_end_offsets,\n ReductionOpT reduction_op)\n {\n for (int i = 0; i < num_segments; ++i)\n {\n AccumT aggregate{};\n for (int j = h_segment_begin_offsets[i]; j < h_segment_end_offsets[i]; ++j)\n aggregate = reduction_op(aggregate, h_in[j]);\n h_reference[i] = static_cast(aggregate);\n }\n }\n};\n\n/// Solve problem (argmin functor)\ntemplate \nstruct Solution\n{\n typedef KeyValuePair OutputT;\n\n template \n static void Solve(HostInputIteratorT h_in, OutputT *h_reference, OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets, EndOffsetIteratorT h_segment_end_offsets, hipcub::ArgMin reduction_op)\n {\n for (int i = 0; i < num_segments; ++i)\n {\n const auto segment_begin = h_segment_begin_offsets[i];\n const auto segment_end = h_segment_end_offsets[i];\n\n if (segment_begin < segment_end) \n {\n OutputT aggregate(0, OutputValueT(h_in[segment_begin]));\n for (int j = segment_begin + 1; j < segment_end; ++j)\n {\n OutputT item(j - segment_begin, OutputValueT(h_in[j]));\n aggregate = reduction_op(aggregate, item);\n }\n h_reference[i] = aggregate;\n }\n else \n {\n // Guaranteed output for empty segments\n OutputT aggregate(1, Traits::Max()); // replace with std::numeric_limits::max() when C++ support is more prevalent\n h_reference[i] = aggregate;\n }\n }\n }\n};\n\n\n/// Solve problem (argmax functor)\ntemplate \nstruct Solution\n{\n typedef KeyValuePair OutputT;\n\n template \n static void Solve(HostInputIteratorT h_in, OutputT *h_reference, OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets, EndOffsetIteratorT h_segment_end_offsets, hipcub::ArgMax reduction_op)\n {\n for (int i = 0; i < num_segments; ++i)\n {\n const auto segment_begin = h_segment_begin_offsets[i];\n const auto segment_end = h_segment_end_offsets[i];\n\n if (segment_begin < segment_end) \n {\n OutputT aggregate(0, OutputValueT(h_in[segment_begin])); \n for (int j = segment_begin + 1; j < segment_end; ++j)\n {\n OutputT item(j - segment_begin, OutputValueT(h_in[j]));\n aggregate = reduction_op(aggregate, item);\n }\n h_reference[i] = aggregate;\n }\n else \n {\n OutputT aggregate(1, Traits::Lowest());\n h_reference[i] = aggregate;\n }\n }\n }\n};\n\n\n//---------------------------------------------------------------------\n// Problem generation\n//---------------------------------------------------------------------\n\ntemplate \nIt unwrap_it(It it) {\n return it;\n}\n\n#if TEST_HALF_T\n__half* unwrap_it(half_t *it) {\n return reinterpret_cast<__half*>(it);\n}\n\ntemplate \nConstantInputIterator<__half, OffsetT> unwrap_it(ConstantInputIterator it) {\n half_t wrapped_val = *it;\n __half val = wrapped_val.operator __half();\n return ConstantInputIterator<__half, OffsetT>(val);\n}\n#endif\n\n#if TEST_BF_T\n__hip_bfloat16* unwrap_it(bfloat16_t* it) {\n return reinterpret_cast<__hip_bfloat16*>(it);\n}\n\ntemplate \nConstantInputIterator<__hip_bfloat16, OffsetT> unwrap_it(ConstantInputIterator it) {\n bfloat16_t wrapped_val = *it;\n __hip_bfloat16 val = wrapped_val.operator __hip_bfloat16();\n return ConstantInputIterator<__hip_bfloat16, OffsetT>(val);\n}\n#endif\n\ntemplate ()))>\nstd::integral_constant::value> //\nreference_extended_fp(WrappedItT)\n{\n return {};\n}\n\nExtendedFloatSum unwrap_op(std::true_type /* extended float */, hipcub::Sum) //\n{\n return {};\n}\n\ntemplate \nOpT unwrap_op(std::integral_constant /* base case */, OpT op)\n{\n return op;\n}\n\n/// Test DeviceReduce for a given problem input\ntemplate <\n typename BackendT,\n typename DeviceWrappedInputIteratorT,\n typename DeviceWrappedOutputIteratorT,\n typename HostReferenceIteratorT,\n typename OffsetT,\n typename BeginOffsetIteratorT,\n typename EndOffsetIteratorT,\n typename ReductionOpT>\nvoid Test(\n BackendT backend,\n DeviceWrappedInputIteratorT d_wrapped_in,\n DeviceWrappedOutputIteratorT d_wrapped_out,\n OffsetT num_items,\n OffsetT num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n ReductionOpT reduction_op,\n HostReferenceIteratorT h_reference)\n{\n // Input data types\n auto d_in = unwrap_it(d_wrapped_in);\n auto d_out = unwrap_it(d_wrapped_out);\n\n // Allocate CDP device arrays for temp storage size and error\n size_t *d_temp_storage_bytes = NULL;\n hipError_t *d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(hipError_t) * 1));\n\n // Inquire temp device storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(backend, 1,\n d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes,\n d_in, d_out, num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets,\n reduction_op));\n\n // Allocate temp device storage\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(backend, 1,\n d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes,\n d_in, d_out, num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets,\n reduction_op));\n\n // Check for correctness (and display results, if specified)\n int compare = CompareDeviceResults(h_reference, d_wrapped_out, num_segments, g_verbose, g_verbose);\n\n printf(\"\\t%s\", compare ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Performance\n if (g_timing_iterations > 0)\n {\n GpuTimer gpu_timer;\n gpu_timer.Start();\n\n CubDebugExit(Dispatch(backend, g_timing_iterations,\n d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes,\n d_in, d_out, num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets,\n reduction_op));\n\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n\n // Display performance\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n float giga_bandwidth = giga_rate * sizeof(h_reference[0]);\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s, %.1f%% peak\",\n avg_millis, giga_rate, giga_bandwidth, giga_bandwidth / g_device_giga_bandwidth * 100.0);\n\n }\n\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare);\n}\n\ntemplate \nvoid SolveAndTest(HostInputIteratorT h_in,\n DeviceInputIteratorT d_in,\n OffsetT num_items,\n OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets,\n EndOffsetIteratorT h_segment_end_offsets,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n ReductionOpT wrapped_reduction_op)\n{\n auto reduction_op = unwrap_op(reference_extended_fp(d_in), wrapped_reduction_op);\n\n using InputValueT = hipcub::detail::value_t;\n using SolutionT = Solution;\n using OutputT = typename SolutionT::OutputT;\n\n printf(\"\\n\\n%s hipcub::DeviceReduce<%s> %d items (%s), %d segments\\n\",\n BackendToString(BACKEND),\n typeid(ReductionOpT).name(),\n num_items,\n typeid(HostInputIteratorT).name(),\n num_segments);\n fflush(stdout);\n\n // Allocate and solve solution\n OutputT *h_reference = new OutputT[num_segments];\n SolutionT::Solve(h_in, h_reference, num_segments, h_segment_begin_offsets, h_segment_end_offsets, reduction_op);\n\n // Run with discard iterator\n DiscardOutputIterator discard_itr;\n Test(Int2Type(), d_in, discard_itr, num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets, reduction_op, h_reference);\n\n // Run with output data\n OutputT *d_out = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(OutputT) * num_segments));\n CubDebugExit(hipMemset(d_out, 0, sizeof(OutputT) * num_segments));\n Test(Int2Type(), d_in, d_out, num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets, reduction_op, h_reference);\n\n // Cleanup\n if (d_out) CubDebugExit(g_allocator.DeviceFree(d_out));\n if (h_reference) delete[] h_reference;\n}\n\n/// Test specific problem type\ntemplate <\n Backend BACKEND,\n typename InputT,\n typename OutputT,\n typename OffsetT,\n typename ReductionOpT>\nvoid TestProblem(\n OffsetT num_items,\n OffsetT num_segments,\n GenMode gen_mode,\n ReductionOpT reduction_op)\n{\n printf(\"\\n\\nInitializing %d %s->%s (gen mode %d)... \", num_items, typeid(InputT).name(), typeid(OutputT).name(), gen_mode); fflush(stdout);\n fflush(stdout);\n\n // Initialize value data\n InputT* h_in = new InputT[num_items];\n Initialize(gen_mode, h_in, num_items);\n\n // Initialize segment data\n OffsetT *h_segment_offsets = new OffsetT[num_segments + 1];\n InitializeSegments(num_items, num_segments, h_segment_offsets, g_verbose_input);\n\n // Initialize device data\n OffsetT *d_segment_offsets = NULL;\n InputT *d_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_in, sizeof(InputT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_segment_offsets, sizeof(OffsetT) * (num_segments + 1)));\n CubDebugExit(hipMemcpy(d_in, h_in, sizeof(InputT) * num_items, hipMemcpyHostToDevice));\n CubDebugExit(hipMemcpy(d_segment_offsets, h_segment_offsets, sizeof(OffsetT) * (num_segments + 1), hipMemcpyHostToDevice));\n\n SolveAndTest(h_in, d_in, num_items, num_segments,\n h_segment_offsets, h_segment_offsets + 1, d_segment_offsets, d_segment_offsets + 1, reduction_op);\n\n if (h_segment_offsets) delete[] h_segment_offsets;\n if (d_segment_offsets) CubDebugExit(g_allocator.DeviceFree(d_segment_offsets));\n if (h_in) delete[] h_in;\n if (d_in) CubDebugExit(g_allocator.DeviceFree(d_in));\n}\n\n\n/// Test different operators\ntemplate <\n Backend BACKEND,\n typename OutputT,\n typename HostInputIteratorT,\n typename DeviceInputIteratorT,\n typename OffsetT,\n typename BeginOffsetIteratorT,\n typename EndOffsetIteratorT>\nvoid TestByOp(\n HostInputIteratorT h_in,\n DeviceInputIteratorT d_in,\n OffsetT num_items,\n OffsetT num_segments,\n BeginOffsetIteratorT h_segment_begin_offsets,\n EndOffsetIteratorT h_segment_end_offsets,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets)\n{\n SolveAndTest(h_in, d_in, num_items, num_segments,\n h_segment_begin_offsets, h_segment_end_offsets, d_segment_begin_offsets, d_segment_end_offsets, CustomMax());\n SolveAndTest(h_in, d_in, num_items, num_segments,\n h_segment_begin_offsets, h_segment_end_offsets, d_segment_begin_offsets, d_segment_end_offsets, Sum());\n SolveAndTest(h_in, d_in, num_items, num_segments,\n h_segment_begin_offsets, h_segment_end_offsets, d_segment_begin_offsets, d_segment_end_offsets, Min());\n SolveAndTest(h_in, d_in, num_items, num_segments,\n h_segment_begin_offsets, h_segment_end_offsets, d_segment_begin_offsets, d_segment_end_offsets, ArgMin());\n SolveAndTest(h_in, d_in, num_items, num_segments,\n h_segment_begin_offsets, h_segment_end_offsets, d_segment_begin_offsets, d_segment_end_offsets, Max());\n SolveAndTest(h_in, d_in, num_items, num_segments,\n h_segment_begin_offsets, h_segment_end_offsets, d_segment_begin_offsets, d_segment_end_offsets, ArgMax());\n}\n\ntemplate\nstruct TransformFunctor1\n{\n __host__ __device__ __forceinline__ OffsetT operator()(OffsetT offset) const\n {\n return offset;\n }\n};\n\ntemplate\nstruct TransformFunctor2\n{\n __host__ __device__ __forceinline__ OffsetT operator()(OffsetT offset) const\n {\n return offset;\n }\n};\n\n/// Test different backends\ntemplate <\n typename InputT,\n typename OutputT,\n typename OffsetT>\nvoid TestByBackend(\n OffsetT num_items,\n OffsetT max_segments,\n GenMode gen_mode)\n{\n#if TEST_CDP == 0\n constexpr auto NonSegmentedBackend = CUB;\n constexpr auto SegmentedBackend = CUB_SEGMENTED;\n#else // TEST_CDP\n constexpr auto NonSegmentedBackend = CDP;\n constexpr auto SegmentedBackend = CDP_SEGMENTED;\n#endif // TEST_CDP\n\n // Initialize host data\n printf(\"\\n\\nInitializing %d %s -> %s (gen mode %d)... \",\n num_items, typeid(InputT).name(), typeid(OutputT).name(), gen_mode); fflush(stdout);\n\n InputT *h_in = new InputT[num_items];\n OffsetT *h_segment_offsets = new OffsetT[max_segments + 1];\n Initialize(gen_mode, h_in, num_items);\n\n // Initialize device data\n InputT *d_in = NULL;\n OffsetT *d_segment_offsets = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_in, sizeof(InputT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_segment_offsets, sizeof(OffsetT) * (max_segments + 1)));\n CubDebugExit(hipMemcpy(d_in, h_in, sizeof(InputT) * num_items, hipMemcpyHostToDevice));\n\n //\n // Test single-segment implementations\n //\n\n InitializeSegments(num_items, 1, h_segment_offsets, g_verbose_input);\n\n // Page-aligned-input tests\n TestByOp(h_in, d_in, num_items, 1,\n h_segment_offsets, h_segment_offsets + 1, (OffsetT*) NULL, (OffsetT*)NULL);\n\n // Non-page-aligned-input tests\n if (num_items > 1)\n {\n InitializeSegments(num_items - 1, 1, h_segment_offsets, g_verbose_input);\n TestByOp(h_in + 1, d_in + 1, num_items - 1, 1,\n h_segment_offsets, h_segment_offsets + 1, (OffsetT*) NULL, (OffsetT*)NULL);\n }\n\n //\n // Test segmented implementation\n //\n\n // Right now we assign a single thread block to each segment, so lets keep it to under 128K items per segment\n int max_items_per_segment = 128000;\n\n for (int num_segments = hipcub::DivideAndRoundUp(num_items, max_items_per_segment);\n num_segments < max_segments;\n num_segments = (num_segments * 32) + 1)\n {\n // Test with segment pointer\n InitializeSegments(num_items, num_segments, h_segment_offsets, g_verbose_input);\n CubDebugExit(hipMemcpy(d_segment_offsets, h_segment_offsets, sizeof(OffsetT) * (num_segments + 1), hipMemcpyHostToDevice));\n TestByOp(h_in, d_in, num_items, num_segments,\n h_segment_offsets, h_segment_offsets + 1, d_segment_offsets, d_segment_offsets + 1);\n\n // Test with segment iterator\n typedef CastOp IdentityOpT;\n IdentityOpT identity_op;\n TransformInputIterator h_segment_offsets_itr(\n h_segment_offsets,\n identity_op);\n TransformInputIterator d_segment_offsets_itr(\n d_segment_offsets,\n identity_op);\n\n TestByOp(h_in, d_in, num_items, num_segments,\n h_segment_offsets_itr, h_segment_offsets_itr + 1, d_segment_offsets_itr, d_segment_offsets_itr + 1);\n\n // Test with transform iterators of different types\n\n typedef TransformFunctor1 TransformFunctor1T;\n typedef TransformFunctor2 TransformFunctor2T;\n\n TransformInputIterator h_segment_begin_offsets_itr(h_segment_offsets, TransformFunctor1T());\n TransformInputIterator h_segment_end_offsets_itr(h_segment_offsets + 1, TransformFunctor2T());\n\n TransformInputIterator d_segment_begin_offsets_itr(d_segment_offsets, TransformFunctor1T());\n TransformInputIterator d_segment_end_offsets_itr(d_segment_offsets + 1, TransformFunctor2T());\n\n TestByOp(h_in, d_in, num_items, num_segments,\n h_segment_begin_offsets_itr, h_segment_end_offsets_itr,\n d_segment_begin_offsets_itr, d_segment_end_offsets_itr);\n }\n\n if (h_in) delete[] h_in;\n if (h_segment_offsets) delete[] h_segment_offsets;\n if (d_in) CubDebugExit(g_allocator.DeviceFree(d_in));\n if (d_segment_offsets) CubDebugExit(g_allocator.DeviceFree(d_segment_offsets));\n}\n\n\n/// Test different input-generation modes\ntemplate <\n typename InputT,\n typename OutputT,\n typename OffsetT>\nvoid TestByGenMode(\n OffsetT num_items,\n OffsetT max_segments)\n{\n //\n // Test pointer support using different input-generation modes\n //\n\n TestByBackend(num_items, max_segments, UNIFORM);\n TestByBackend(num_items, max_segments, INTEGER_SEED);\n TestByBackend(num_items, max_segments, RANDOM);\n\n //\n // Test iterator support using a constant-iterator and SUM\n //\n\n InputT val;\n InitValue(UNIFORM, val, 0);\n ConstantInputIterator in(val);\n\n OffsetT *h_segment_offsets = new OffsetT[1 + 1];\n InitializeSegments(num_items, 1, h_segment_offsets, g_verbose_input);\n\n#if TEST_CDP == 0\n constexpr auto Backend = CUB;\n#else // TEST_CDP\n constexpr auto Backend = CDP;\n#endif // TEST_CDP\n\n SolveAndTest(in, in, num_items, 1,\n h_segment_offsets, h_segment_offsets + 1, (OffsetT*) NULL, (OffsetT*)NULL, Sum());\n\n if (h_segment_offsets) delete[] h_segment_offsets;\n}\n\n/// Test different problem sizes\ntemplate \nvoid TestBySize(OffsetT max_items, OffsetT max_segments, OffsetT tile_size)\n{\n // Test 0, 1, many\n TestByGenMode(0, max_segments);\n TestByGenMode(1, max_segments);\n TestByGenMode(max_items, max_segments);\n\n // Test random problem sizes from a log-distribution [8, max_items-ish)\n int num_iterations = 8;\n double max_exp = log(double(max_items)) / log(double(2.0));\n for (int i = 0; i < num_iterations; ++i)\n {\n OffsetT num_items = (OffsetT)pow(2.0, RandomValue(max_exp - 3.0) + 3.0);\n TestByGenMode(num_items, max_segments);\n }\n\n //\n // White-box testing of single-segment problems around specific sizes\n //\n\n#if TEST_CDP == 0\n constexpr auto Backend = CUB;\n#else // TEST_CDP\n constexpr auto Backend = CDP;\n#endif // TEST_CDP\n\n // Tile-boundaries: multiple blocks, one tile per block\n TestProblem(tile_size * 4, 1, RANDOM, Sum());\n TestProblem(tile_size * 4 + 1, 1, RANDOM, Sum());\n TestProblem(tile_size * 4 - 1, 1, RANDOM, Sum());\n\n // Tile-boundaries: multiple blocks, multiple tiles per block\n OffsetT sm_occupancy = 32;\n OffsetT occupancy = tile_size * sm_occupancy * g_sm_count;\n TestProblem(occupancy, 1, RANDOM, Sum());\n TestProblem(occupancy + 1, 1, RANDOM, Sum());\n TestProblem(occupancy - 1, 1, RANDOM, Sum());\n};\n\nclass CustomInputT\n{\n char m_val{};\n\npublic:\n __host__ __device__ explicit CustomInputT(char val)\n : m_val(val)\n {}\n\n __host__ __device__ int get() const { return static_cast(m_val); }\n};\n\nclass CustomAccumulatorT\n{\n int m_val{0};\n int m_magic_value{42};\n\n __host__ __device__ CustomAccumulatorT(int val)\n : m_val(val)\n {}\n\npublic:\n __host__ __device__ CustomAccumulatorT()\n {}\n\n __host__ __device__ CustomAccumulatorT(const CustomAccumulatorT &in)\n : m_val(in.is_valid() * in.get())\n , m_magic_value(in.is_valid() * 42)\n {}\n\n __host__ __device__ void operator=(const CustomInputT &in)\n {\n if (this->is_valid())\n {\n m_val = in.get();\n }\n }\n\n __host__ __device__ void operator=(const CustomAccumulatorT &in)\n {\n if (this->is_valid() && in.is_valid())\n {\n m_val = in.get();\n }\n }\n\n __host__ __device__ CustomAccumulatorT \n operator+(const CustomInputT &in) const\n {\n const int multiplier = this->is_valid();\n return {(m_val + in.get()) * multiplier};\n }\n\n __host__ __device__ CustomAccumulatorT\n operator+(const CustomAccumulatorT &in) const\n {\n const int multiplier = this->is_valid() && in.is_valid();\n return {(m_val + in.get()) * multiplier};\n }\n\n __host__ __device__ int get() const { return m_val; }\n\n __host__ __device__ bool is_valid() const { return m_magic_value == 42; }\n};\n\nclass CustomOutputT\n{\n bool *m_d_flag{};\n int m_expected{};\n\npublic:\n __host__ __device__ CustomOutputT(bool *d_flag, int expected)\n : m_d_flag(d_flag)\n , m_expected(expected)\n {}\n\n __host__ __device__ void operator=(const CustomAccumulatorT &accum) const\n {\n *m_d_flag = accum.is_valid() && (accum.get() == m_expected);\n }\n};\n\n__global__ void InitializeTestAccumulatorTypes(int num_items,\n int expected,\n bool *d_flag,\n CustomInputT *d_in,\n CustomOutputT *d_out)\n{\n const int idx = static_cast(threadIdx.x + blockIdx.x * blockDim.x);\n\n if (idx < num_items)\n {\n d_in[idx] = CustomInputT(1);\n }\n\n if (idx == 0)\n {\n *d_out = CustomOutputT{d_flag, expected};\n }\n}\n\ntemplate \nvoid TestBigIndicesHelper(OffsetT num_items)\n{\n thrust::constant_iterator const_iter(T{1});\n thrust::device_vector out(1);\n std::size_t* d_out = thrust::raw_pointer_cast(out.data());\n\n std::uint8_t *d_temp_storage{};\n std::size_t temp_storage_bytes{};\n\n CubDebugExit(\n hipcub::DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, const_iter, d_out, num_items));\n\n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n CubDebugExit(\n hipcub::DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, const_iter, d_out, num_items));\n std::size_t result = out[0];\n\n AssertEquals(result, num_items);\n}\n\ntemplate \nvoid TestBigIndices()\n{\n TestBigIndicesHelper(1ull << 30);\n TestBigIndicesHelper(1ull << 31);\n TestBigIndicesHelper((1ull << 32) - 1);\n TestBigIndicesHelper(1ull << 33);\n}\n\n#if TEST_TYPES == 3\nvoid TestAccumulatorTypes()\n{\n const int num_items = 2 * 1024 * 1024;\n const int expected = num_items;\n const int block_size = 256;\n const int grid_size = (num_items + block_size - 1) / block_size;\n\n CustomInputT *d_in{};\n CustomOutputT *d_out{};\n CustomAccumulatorT init{};\n bool *d_flag{};\n\n CubDebugExit(\n g_allocator.DeviceAllocate((void **)&d_out, sizeof(CustomOutputT)));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_flag, sizeof(bool)));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_in,\n sizeof(CustomInputT) * num_items));\n\n InitializeTestAccumulatorTypes<<>>(num_items,\n expected,\n d_flag,\n d_in,\n d_out);\n\n std::uint8_t *d_temp_storage{};\n std::size_t temp_storage_bytes{};\n\n CubDebugExit(hipcub::DeviceReduce::Reduce(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items,\n hipcub::Sum{},\n init));\n\n CubDebugExit(\n g_allocator.DeviceAllocate((void **)&d_temp_storage, temp_storage_bytes));\n CubDebugExit(hipMemset(d_temp_storage, 1, temp_storage_bytes));\n\n CubDebugExit(hipcub::DeviceReduce::Reduce(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items,\n hipcub::Sum{},\n init));\n\n bool ok{};\n CubDebugExit(hipMemcpy(&ok, d_flag, sizeof(bool), hipMemcpyDeviceToHost));\n\n AssertTrue(ok);\n\n CubDebugExit(g_allocator.DeviceFree(d_out));\n CubDebugExit(g_allocator.DeviceFree(d_in));\n}\n\n/**\n * ArgMin should return max value for empty input. This interferes with\n * input data containing infinity values. This test checks that ArgMin\n * works correctly with infinity values.\n */\nvoid TestFloatInfInArgMin()\n{\n using in_t = float;\n using offset_t = int;\n using out_t = hipcub::KeyValuePair;\n\n const int n = 10;\n const float inf = ::cuda::std::numeric_limits::infinity();\n\n thrust::device_vector in(n, inf);\n thrust::device_vector out(1);\n\n const in_t *d_in = thrust::raw_pointer_cast(in.data());\n out_t *d_out = thrust::raw_pointer_cast(out.data());\n\n std::uint8_t *d_temp_storage{};\n std::size_t temp_storage_bytes{};\n\n CubDebugExit(hipcub::DeviceReduce::ArgMin(d_temp_storage, temp_storage_bytes, d_in, d_out, n));\n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n CubDebugExit(hipcub::DeviceReduce::ArgMin(d_temp_storage, temp_storage_bytes, d_in, d_out, n));\n\n const out_t result = out[0];\n AssertEquals(result.key, 0);\n AssertEquals(result.value, inf);\n}\n\n/**\n * ArgMax should return lowest value for empty input. This interferes with\n * input data containing infinity values. This test checks that ArgMax\n * works correctly with infinity values.\n */\nvoid TestFloatInfInArgMax()\n{\n using in_t = float;\n using offset_t = int;\n using out_t = hipcub::KeyValuePair;\n\n const int n = 10;\n const float inf = ::cuda::std::numeric_limits::infinity();\n \n thrust::device_vector in(n, -inf);\n thrust::device_vector out(1);\n\n const in_t *d_in = thrust::raw_pointer_cast(in.data());\n out_t *d_out = thrust::raw_pointer_cast(out.data());\n\n std::uint8_t *d_temp_storage{};\n std::size_t temp_storage_bytes{}; \n\n CubDebugExit(hipcub::DeviceReduce::ArgMax(d_temp_storage, temp_storage_bytes, d_in, d_out, n));\n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n CubDebugExit(hipcub::DeviceReduce::ArgMax(d_temp_storage, temp_storage_bytes, d_in, d_out, n));\n\n const out_t result = out[0];\n AssertEquals(result.key, 0);\n AssertEquals(result.value, -inf);\n}\n\nvoid TestFloatInfInArg()\n{\n TestFloatInfInArgMin();\n TestFloatInfInArgMax();\n}\n#endif\n\ntemplate \nstruct GetTileSize\n{\n OffsetT max_items{};\n OffsetT max_segments{};\n OffsetT tile_size{};\n\n GetTileSize(OffsetT max_items, OffsetT max_segments)\n : max_items(max_items)\n , max_segments(max_segments)\n {}\n\n template \n HIPCUB_RUNTIME_FUNCTION hipError_t Invoke()\n {\n this->tile_size = ActivePolicyT::ReducePolicy::BLOCK_THREADS *\n ActivePolicyT::ReducePolicy::ITEMS_PER_THREAD;\n return hipSuccess;\n }\n};\n\n/// Test problem type\ntemplate \nvoid TestType(OffsetT max_items, OffsetT max_segments)\n{\n // Inspect the tuning policies to determine this arch's tile size:\n using MaxPolicyT =\n typename DeviceReducePolicy::MaxPolicy;\n GetTileSize dispatch(max_items, max_segments);\n CubDebugExit(MaxPolicyT::Invoke(g_ptx_version, dispatch));\n\n TestBySize(max_items, max_segments, dispatch.tile_size);\n}\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n typedef int OffsetT;\n\n OffsetT max_items = 27000000;\n OffsetT max_segments = 34000;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n g_verbose_input = args.CheckCmdLineFlag(\"v2\");\n args.GetCmdLineArgument(\"n\", max_items);\n args.GetCmdLineArgument(\"s\", max_segments);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--s= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n g_device_giga_bandwidth = args.device_giga_bandwidth;\n\n // Get ptx version\n CubDebugExit(PtxVersion(g_ptx_version));\n\n // Get SM count\n g_sm_count = args.deviceProp.multiProcessorCount;\n\n // %PARAM% TEST_CDP cdp 0:1\n // %PARAM% TEST_TYPES types 0:1:2:3:4\n\n#if TEST_TYPES == 0\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n#elif TEST_TYPES == 1\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n#elif TEST_TYPES == 2\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n TestType(max_items, max_segments);\n#elif TEST_TYPES == 3\n TestType(max_items, max_segments);\n TestAccumulatorTypes();\n TestFloatInfInArg();\n\n#if TEST_HALF_T\n TestType(max_items, max_segments);\n#endif\n#else // TEST_TYPES == 4\n TestType(max_items, max_segments);\n TestBigIndices();\n\n#if TEST_BF_T\n TestType(max_items, max_segments);\n#endif\n#endif\n printf(\"\\n\");\n return 0;\n}\n\n\n\n", "source": "/kaggle/working/cub/test/test_device_reduce.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceSelect::If and DevicePartition::If utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"test_util.h\"\n\n#include \n#include \n\nusing namespace cub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\nfloat g_device_giga_bandwidth;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n\n// Selection functor type\ntemplate \nstruct LessThan\n{\n T compare;\n\n __host__ __device__ __forceinline__\n LessThan(T compare) : compare(compare) {}\n\n __host__ __device__ __forceinline__\n bool operator()(const T &a) const {\n return (a < compare);\n }\n};\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB DeviceSelect entrypoints\n//---------------------------------------------------------------------\n\n\n/**\n * Dispatch to select if entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n Int2Type /*is_flagged*/,\n Int2Type /*is_partition*/,\n int timing_timing_iterations,\n size_t* /*d_temp_storage_bytes*/,\n cudaError_t* /*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n FlagIteratorT /*d_flags*/,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items,\n SelectOpT select_op)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceSelect::If(d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items, select_op);\n }\n return error;\n}\n\n\n/**\n * Dispatch to partition if entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n Int2Type /*is_flagged*/,\n Int2Type /*is_partition*/,\n int timing_timing_iterations,\n size_t* /*d_temp_storage_bytes*/,\n cudaError_t* /*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n FlagIteratorT /*d_flags*/,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items,\n SelectOpT select_op)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DevicePartition::If(d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items, select_op);\n }\n return error;\n}\n\n\n/**\n * Dispatch to select flagged entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n Int2Type /*is_flagged*/,\n Int2Type /*partition*/,\n int timing_timing_iterations,\n size_t* /*d_temp_storage_bytes*/,\n cudaError_t* /*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n FlagIteratorT d_flags,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items,\n SelectOpT /*select_op*/)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceSelect::Flagged(d_temp_storage, temp_storage_bytes, d_in, d_flags, d_out, d_num_selected_out, num_items);\n }\n return error;\n}\n\n\n/**\n * Dispatch to partition flagged entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n Int2Type /*is_flagged*/,\n Int2Type /*partition*/,\n int timing_timing_iterations,\n size_t* /*d_temp_storage_bytes*/,\n cudaError_t* /*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n FlagIteratorT d_flags,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items,\n SelectOpT /*select_op*/)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DevicePartition::Flagged(d_temp_storage, temp_storage_bytes, d_in, d_flags, d_out, d_num_selected_out, num_items);\n }\n return error;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceSelect\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type cub_backend,\n IsFlaggedTag is_flagged,\n IsPartitionTag is_partition,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n InputIteratorT d_in,\n FlagIteratorT d_flags,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items,\n SelectOpT select_op)\n{\n *d_cdp_error = Dispatch(cub_backend,\n is_flagged,\n is_partition,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_flags,\n d_out,\n d_num_selected_out,\n num_items,\n select_op);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Dispatch to CDP kernel\n */\ntemplate \ncudaError_t Dispatch(Int2Type /*dispatch_to*/,\n IsFlaggedTag is_flagged,\n IsPartitionTag is_partition,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n FlagIteratorT d_flags,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items,\n SelectOpT select_op)\n{\n // Invoke kernel to invoke device-side dispatch\n cudaError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n Int2Type{},\n is_flagged,\n is_partition,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_flags,\n d_out,\n d_num_selected_out,\n num_items,\n select_op);\n CubDebugExit(retval);\n\n // Copy out temp_storage_bytes\n CubDebugExit(cudaMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n cudaMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(cudaMemcpy(&retval,\n d_cdp_error,\n sizeof(cudaError_t) * 1,\n cudaMemcpyDeviceToHost));\n return retval;\n}\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\n\n/**\n * Initialize problem\n */\ntemplate \nvoid Initialize(\n T* h_in,\n int num_items)\n{\n for (int i = 0; i < num_items; ++i)\n {\n // Initialize each item to a randomly selected value from [0..126]\n unsigned int value;\n RandomBits(value, 0, 0, 7);\n if (value == 127)\n value = 126;\n InitValue(INTEGER_SEED, h_in[i], value);\n }\n\n if (g_verbose)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n\n/**\n * Solve selection problem (and set corresponding flags)\n */\ntemplate <\n typename InputIteratorT,\n typename FlagIteratorT,\n typename SelectOpT,\n typename T>\nint Solve(\n InputIteratorT h_in,\n SelectOpT select_op,\n T* h_reference,\n FlagIteratorT h_flags,\n int num_items)\n{\n int num_selected = 0;\n for (int i = 0; i < num_items; ++i)\n {\n if ((h_flags[i] = select_op(h_in[i])))\n {\n h_reference[num_selected] = h_in[i];\n num_selected++;\n }\n else\n {\n h_reference[num_items - (i - num_selected) - 1] = h_in[i];\n }\n }\n\n return num_selected;\n}\n\n\n\n/**\n * Test DeviceSelect for a given problem input\n */\ntemplate <\n Backend BACKEND,\n bool IS_FLAGGED,\n bool IS_PARTITION,\n typename DeviceInputIteratorT,\n typename FlagT,\n typename SelectOpT,\n typename T>\nvoid Test(\n DeviceInputIteratorT d_in,\n FlagT* h_flags,\n SelectOpT select_op,\n T* h_reference,\n int num_selected,\n int num_items)\n{\n // Allocate device flags, output, and num-selected\n FlagT* d_flags = NULL;\n T* d_out = NULL;\n int* d_num_selected_out = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_flags, sizeof(FlagT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(T) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_num_selected_out, sizeof(int)));\n\n // Allocate CDP device arrays\n size_t* d_temp_storage_bytes = NULL;\n cudaError_t* d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(cudaError_t) * 1));\n\n // Allocate temporary storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(Int2Type(), Int2Type(), Int2Type(), 1, d_temp_storage_bytes, d_cdp_error,\n d_temp_storage, temp_storage_bytes, d_in, d_flags, d_out, d_num_selected_out, num_items, select_op));\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n // Copy flags and clear device output array\n CubDebugExit(cudaMemcpy(d_flags, h_flags, sizeof(FlagT) * num_items, cudaMemcpyHostToDevice));\n CubDebugExit(cudaMemset(d_out, 0, sizeof(T) * num_items));\n CubDebugExit(cudaMemset(d_num_selected_out, 0, sizeof(int)));\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(Int2Type(), Int2Type(), Int2Type(), 1, d_temp_storage_bytes, d_cdp_error,\n d_temp_storage, temp_storage_bytes, d_in, d_flags, d_out, d_num_selected_out, num_items, select_op));\n\n // Check for correctness (and display results, if specified)\n int compare1 = (IS_PARTITION) ?\n CompareDeviceResults(h_reference, d_out, num_items, true, g_verbose) :\n CompareDeviceResults(h_reference, d_out, num_selected, true, g_verbose);\n printf(\"\\t Data %s\\n\", compare1 ? \"FAIL\" : \"PASS\");\n\n int compare2 = CompareDeviceResults(&num_selected, d_num_selected_out, 1, true, g_verbose);\n printf(\"\\t Count %s\\n\", compare2 ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Performance\n GpuTimer gpu_timer;\n gpu_timer.Start();\n CubDebugExit(Dispatch(Int2Type(), Int2Type(), Int2Type(), g_timing_iterations, d_temp_storage_bytes, d_cdp_error,\n d_temp_storage, temp_storage_bytes, d_in, d_flags, d_out, d_num_selected_out, num_items, select_op));\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n\n // Display performance\n if (g_timing_iterations > 0)\n {\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n int num_output_items = (IS_PARTITION) ? num_items : num_selected;\n int num_flag_items = (IS_FLAGGED) ? num_items : 0;\n size_t num_bytes = sizeof(T) * (num_items + num_output_items) + sizeof(FlagT) * num_flag_items;\n float giga_bandwidth = float(num_bytes) / avg_millis / 1000.0f / 1000.0f;\n\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s, %.1f%% peak\", avg_millis, giga_rate, giga_bandwidth, giga_bandwidth / g_device_giga_bandwidth * 100.0);\n }\n printf(\"\\n\\n\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Cleanup\n if (d_flags) CubDebugExit(g_allocator.DeviceFree(d_flags));\n if (d_out) CubDebugExit(g_allocator.DeviceFree(d_out));\n if (d_num_selected_out) CubDebugExit(g_allocator.DeviceFree(d_num_selected_out));\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare1 | compare2);\n}\n\n\n/**\n * Test on pointer type\n */\ntemplate <\n Backend BACKEND,\n bool IS_FLAGGED,\n bool IS_PARTITION,\n typename T>\nvoid TestPointer(\n int num_items,\n float select_ratio)\n{\n typedef char FlagT;\n\n // Allocate host arrays\n T* h_in = new T[num_items];\n FlagT* h_flags = new FlagT[num_items];\n T* h_reference = new T[num_items];\n\n // Initialize input\n Initialize(h_in, num_items);\n\n // Select a comparison value that is select_ratio through the space of [0,127]\n T compare;\n if (select_ratio <= 0.0)\n InitValue(INTEGER_SEED, compare, 0); // select none\n else if (select_ratio >= 1.0)\n InitValue(INTEGER_SEED, compare, 127); // select all\n else\n InitValue(INTEGER_SEED, compare, int(double(double(127) * select_ratio)));\n\n LessThan select_op(compare);\n int num_selected = Solve(h_in, select_op, h_reference, h_flags, num_items);\n\n if (g_verbose) std::cout << \"\\nComparison item: \" << compare << \"\\n\";\n printf(\"\\nPointer %s cub::%s::%s %d items, %d selected (select ratio %.3f), %s %d-byte elements\\n\",\n (IS_PARTITION) ? \"DevicePartition\" : \"DeviceSelect\",\n (IS_FLAGGED) ? \"Flagged\" : \"If\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items, num_selected, float(num_selected) / num_items, typeid(T).name(), (int) sizeof(T));\n fflush(stdout);\n\n // Allocate problem device arrays\n T *d_in = NULL;\n\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_in, sizeof(T) * num_items));\n\n // Initialize device input\n CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(T) * num_items, cudaMemcpyHostToDevice));\n\n // Run Test\n Test(d_in, h_flags, select_op, h_reference, num_selected, num_items);\n\n // Cleanup\n if (h_in) delete[] h_in;\n if (h_reference) delete[] h_reference;\n if (h_flags) delete[] h_flags;\n if (d_in) CubDebugExit(g_allocator.DeviceFree(d_in));\n}\n\n\n/**\n * Test on iterator type\n */\ntemplate <\n Backend BACKEND,\n bool IS_FLAGGED,\n bool IS_PARTITION,\n typename T>\nvoid TestIterator(\n int num_items,\n float select_ratio)\n{\n typedef char FlagT;\n\n // Allocate host arrays\n T* h_reference = new T[num_items];\n FlagT* h_flags = new FlagT[num_items];\n\n // Use counting iterator as the input\n CountingInputIterator h_in(0);\n\n // Select a comparison value that is select_ratio through the space of [0,127]\n T compare;\n if (select_ratio <= 0.0)\n InitValue(INTEGER_SEED, compare, 0); // select none\n else if (select_ratio >= 1.0)\n InitValue(INTEGER_SEED, compare, 127); // select all\n else\n InitValue(INTEGER_SEED, compare, int(double(double(127) * select_ratio)));\n\n LessThan select_op(compare);\n int num_selected = Solve(h_in, select_op, h_reference, h_flags, num_items);\n\n if (g_verbose) std::cout << \"\\nComparison item: \" << compare << \"\\n\";\n printf(\"\\nIterator %s cub::%s::%s %d items, %d selected (select ratio %.3f), %s %d-byte elements\\n\",\n (IS_PARTITION) ? \"DevicePartition\" : \"DeviceSelect\",\n (IS_FLAGGED) ? \"Flagged\" : \"If\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items, num_selected, float(num_selected) / num_items, typeid(T).name(), (int) sizeof(T));\n fflush(stdout);\n\n // Run Test\n Test(h_in, h_flags, select_op, h_reference, num_selected, num_items);\n\n // Cleanup\n if (h_reference) delete[] h_reference;\n if (h_flags) delete[] h_flags;\n}\n\n\n/**\n * Test different selection ratios\n */\ntemplate <\n Backend BACKEND,\n bool IS_FLAGGED,\n bool IS_PARTITION,\n typename T>\nvoid Test(\n int num_items)\n{\n for (float select_ratio = 0.0f; select_ratio <= 1.0f; select_ratio += 0.2f)\n {\n TestPointer(num_items, select_ratio);\n }\n}\n\n\n/**\n * Test (select vs. partition) and (flagged vs. functor)\n */\ntemplate <\n Backend BACKEND,\n typename T>\nvoid TestMethod(\n int num_items)\n{\n // Functor\n Test(num_items);\n Test(num_items);\n\n // Flagged\n Test(num_items);\n Test(num_items);\n}\n\n\n/**\n * Test different dispatch\n */\ntemplate <\n typename T>\nvoid TestOp(\n int num_items)\n{\n#if TEST_CDP == 0\n TestMethod(num_items);\n#elif TEST_CDP == 1\n TestMethod(num_items);\n#endif // TEST_CDP\n}\n\n\n/**\n * Test different input sizes\n */\ntemplate \nvoid Test(\n int num_items)\n{\n if (num_items < 0)\n {\n TestOp(0);\n TestOp(1);\n TestOp(100);\n TestOp(10000);\n TestOp(1000000);\n }\n else\n {\n TestOp(num_items);\n }\n}\n\ntemplate\nstruct pair_to_col_t\n{\n __host__ __device__ T0 operator()(const thrust::tuple &in)\n {\n return thrust::get<0>(in);\n }\n};\n\ntemplate\nstruct select_t {\n __host__ __device__ bool operator()(const thrust::tuple &in) {\n return static_cast(thrust::get<0>(in)) > thrust::get<1>(in);\n }\n};\n\ntemplate \nvoid TestMixedOp(int num_items)\n{\n const T0 target_value = static_cast(42);\n thrust::device_vector col_a(num_items, target_value);\n thrust::device_vector col_b(num_items, static_cast(4.2));\n\n thrust::device_vector result(num_items);\n\n auto in = thrust::make_zip_iterator(col_a.begin(), col_b.begin());\n auto out = thrust::make_transform_output_iterator(result.begin(), pair_to_col_t{});\n\n void *d_tmp_storage {};\n std::size_t tmp_storage_size{};\n cub::DeviceSelect::If(\n d_tmp_storage, tmp_storage_size,\n in, out, thrust::make_discard_iterator(),\n num_items, select_t{});\n\n thrust::device_vector tmp_storage(tmp_storage_size);\n d_tmp_storage = thrust::raw_pointer_cast(tmp_storage.data());\n\n cub::DeviceSelect::If(\n d_tmp_storage, tmp_storage_size,\n in, out, thrust::make_discard_iterator(),\n num_items, select_t{});\n\n AssertEquals(num_items, thrust::count(result.begin(), result.end(), target_value));\n}\n\n/**\n * Test different input sizes\n */\ntemplate \nvoid TestMixed(int num_items)\n{\n if (num_items < 0)\n {\n TestMixedOp(0);\n TestMixedOp(1);\n TestMixedOp(100);\n TestMixedOp(10000);\n TestMixedOp(1000000);\n }\n else\n {\n TestMixedOp(num_items);\n }\n}\n\nvoid TestFlagsNormalization()\n{\n const int num_items = 1024 * 1024;\n thrust::device_vector result(num_items);\n\n void *d_tmp_storage{};\n std::size_t tmp_storage_size{};\n CubDebugExit(\n cub::DeviceSelect::Flagged(d_tmp_storage,\n tmp_storage_size,\n cub::CountingInputIterator(0), // in\n cub::CountingInputIterator(1), // flags\n thrust::raw_pointer_cast(result.data()), // out\n thrust::make_discard_iterator(), // num_out\n num_items));\n\n thrust::device_vector tmp_storage(tmp_storage_size);\n d_tmp_storage = thrust::raw_pointer_cast(tmp_storage.data());\n\n CubDebugExit(\n cub::DeviceSelect::Flagged(d_tmp_storage,\n tmp_storage_size,\n cub::CountingInputIterator(0), // in\n cub::CountingInputIterator(1), // flags\n thrust::raw_pointer_cast(result.data()), // out\n thrust::make_discard_iterator(), // num_out\n num_items));\n\n AssertTrue(thrust::equal(result.begin(),\n result.end(),\n thrust::make_counting_iterator(0)));\n}\n\nvoid TestFlagsAliasingInPartition()\n{\n int h_items[]{0, 1, 0, 2, 0, 3, 0, 4, 0, 5};\n constexpr int num_items = sizeof(h_items) / sizeof(h_items[0]);\n\n int *d_in{};\n int *d_out{};\n\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_in, sizeof(h_items)));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_out, sizeof(h_items)));\n\n CubDebugExit(\n cudaMemcpy(d_in, h_items, sizeof(h_items), cudaMemcpyHostToDevice));\n\n // alias flags and keys\n int *d_flags = d_in;\n\n void *d_tmp_storage{};\n std::size_t tmp_storage_size{};\n\n CubDebugExit(\n cub::DevicePartition::Flagged(d_tmp_storage,\n tmp_storage_size,\n d_in,\n d_flags,\n d_out,\n thrust::make_discard_iterator(), // num_out\n num_items));\n\n thrust::device_vector tmp_storage(tmp_storage_size);\n d_tmp_storage = thrust::raw_pointer_cast(tmp_storage.data());\n\n CubDebugExit(\n cub::DevicePartition::Flagged(d_tmp_storage,\n tmp_storage_size,\n d_in,\n d_flags,\n d_out,\n thrust::make_discard_iterator(), // num_out\n num_items));\n\n AssertTrue(thrust::equal(thrust::device,\n d_out,\n d_out + num_items / 2,\n thrust::make_counting_iterator(1)));\n\n AssertEquals(\n thrust::count(thrust::device, d_out + num_items / 2, d_out + num_items, 0),\n num_items / 2);\n\n CubDebugExit(g_allocator.DeviceFree(d_out));\n CubDebugExit(g_allocator.DeviceFree(d_in));\n}\n\nstruct Odd\n{\n __host__ __device__ bool operator()(int v) const { return v % 2; }\n};\n\nvoid TestIfInPlace()\n{\n const int num_items = 4 * 1024 * 1024;\n const int num_iters = 42;\n\n thrust::device_vector num_out(1);\n thrust::device_vector data(num_items);\n thrust::device_vector reference(num_items);\n thrust::device_vector reference_out(1);\n\n thrust::sequence(data.begin(), data.end());\n\n Odd op{};\n\n int *d_num_out = thrust::raw_pointer_cast(num_out.data());\n int *d_data = thrust::raw_pointer_cast(data.data());\n int *d_reference = thrust::raw_pointer_cast(reference.data());\n int *d_reference_out = thrust::raw_pointer_cast(reference_out.data());\n\n void *d_tmp_storage{};\n std::size_t tmp_storage_size{};\n\n CubDebugExit(\n cub::DeviceSelect::If(d_tmp_storage,\n tmp_storage_size,\n d_data,\n d_num_out,\n num_items,\n op));\n\n thrust::device_vector tmp_storage(tmp_storage_size);\n d_tmp_storage = thrust::raw_pointer_cast(tmp_storage.data());\n\n thrust::default_random_engine g{};\n\n for (int iter = 0; iter < num_iters; iter++)\n {\n thrust::shuffle(data.begin(), data.end(), g);\n\n CubDebugExit(\n cub::DeviceSelect::If(d_tmp_storage,\n tmp_storage_size,\n d_data,\n d_reference,\n d_reference_out,\n num_items,\n op));\n\n CubDebugExit(\n cub::DeviceSelect::If(d_tmp_storage,\n tmp_storage_size,\n d_data,\n d_num_out,\n num_items,\n op));\n\n AssertEquals(num_out, reference_out);\n const int num_selected = num_out[0];\n\n const bool match_reference = thrust::equal(reference.begin(),\n reference.begin() + num_selected,\n data.begin());\n AssertTrue(match_reference);\n }\n}\n\nvoid TestFlaggedInPlace()\n{\n const int num_items = 4 * 1024 * 1024;\n const int num_iters = 42;\n\n thrust::device_vector num_out(1);\n thrust::device_vector data(num_items);\n thrust::device_vector flags(num_items);\n\n int h_num_out{};\n int *d_num_out = thrust::raw_pointer_cast(num_out.data());\n int *d_data = thrust::raw_pointer_cast(data.data());\n bool *d_flags = thrust::raw_pointer_cast(flags.data());\n\n void *d_tmp_storage{};\n std::size_t tmp_storage_size{};\n\n CubDebugExit(\n cub::DeviceSelect::Flagged(d_tmp_storage,\n tmp_storage_size,\n d_data,\n d_flags,\n d_num_out,\n num_items));\n\n thrust::device_vector tmp_storage(tmp_storage_size);\n d_tmp_storage = thrust::raw_pointer_cast(tmp_storage.data());\n\n thrust::default_random_engine g{};\n\n for (int iter = 0; iter < num_iters; iter++)\n {\n const int num_selected = RandomValue(num_items);\n\n thrust::sequence(data.begin(), data.end());\n thrust::fill(flags.begin(), flags.begin() + num_selected, true);\n thrust::fill(flags.begin() + num_selected, flags.end(), false);\n thrust::shuffle(flags.begin(), flags.end(), g);\n\n CubDebugExit(\n cub::DeviceSelect::Flagged(d_tmp_storage,\n tmp_storage_size,\n d_data,\n d_flags,\n d_num_out,\n num_items));\n\n cudaMemcpy(&h_num_out, d_num_out, sizeof(int), cudaMemcpyDeviceToHost);\n\n AssertEquals(num_selected, h_num_out);\n\n auto selection_perm_begin = thrust::make_permutation_iterator(flags.begin(),\n data.begin());\n auto selection_perm_end = selection_perm_begin + num_selected;\n\n AssertEquals(num_selected,\n thrust::count(selection_perm_begin, selection_perm_end, true));\n }\n}\n\nvoid TestFlaggedInPlaceWithAliasedFlags()\n{\n const int num_items = 1024 * 1024;\n const int num_iters = 42;\n\n thrust::device_vector num_out(1);\n thrust::device_vector data(num_items);\n thrust::device_vector reference(num_items);\n thrust::device_vector flags(num_items);\n\n int h_num_out{};\n int *d_num_out = thrust::raw_pointer_cast(num_out.data());\n int *d_data = thrust::raw_pointer_cast(data.data());\n int *d_flags = d_data; // alias\n int *d_allocated_flags = thrust::raw_pointer_cast(data.data());\n int *d_reference = thrust::raw_pointer_cast(reference.data());\n\n void *d_tmp_storage{};\n std::size_t tmp_storage_size{};\n\n CubDebugExit(\n cub::DeviceSelect::Flagged(d_tmp_storage,\n tmp_storage_size,\n d_data,\n d_flags,\n d_num_out,\n num_items));\n\n thrust::device_vector tmp_storage(tmp_storage_size);\n d_tmp_storage = thrust::raw_pointer_cast(tmp_storage.data());\n\n thrust::default_random_engine g{};\n\n for (int iter = 0; iter < num_iters; iter++)\n {\n const int num_selected = RandomValue(num_items);\n\n thrust::sequence(data.begin(), data.begin() + num_selected, 1);\n thrust::fill(data.begin() + num_selected, data.end(), 0);\n thrust::shuffle(data.begin(), data.end(), g);\n\n CubDebugExit(\n cub::DeviceSelect::Flagged(d_tmp_storage,\n tmp_storage_size,\n d_data, // in\n d_allocated_flags,\n d_reference, // out\n d_num_out,\n num_items));\n\n CubDebugExit(\n cub::DeviceSelect::Flagged(d_tmp_storage,\n tmp_storage_size,\n d_data,\n d_flags,\n d_num_out,\n num_items));\n\n cudaMemcpy(&h_num_out, d_num_out, sizeof(int), cudaMemcpyDeviceToHost);\n\n AssertEquals(num_selected, h_num_out);\n\n const bool match_reference = thrust::equal(reference.begin(),\n reference.begin() + num_selected,\n data.begin());\n AssertTrue(match_reference);\n }\n}\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n int num_items = -1;\n float select_ratio = 0.5;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n args.GetCmdLineArgument(\"ratio\", select_ratio);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--ratio=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n g_device_giga_bandwidth = args.device_giga_bandwidth;\n printf(\"\\n\");\n\n // %PARAM% TEST_CDP cdp 0:1\n\n TestFlagsAliasingInPartition();\n\n TestFlaggedInPlace();\n TestFlaggedInPlaceWithAliasedFlags();\n TestIfInPlace();\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n\n TestMixed(num_items);\n TestFlagsNormalization();\n\n return 0;\n}\n\n\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceSelect::If and DevicePartition::If utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"test_util.h\"\n\n#include \n#include \n\nusing namespace hipcub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\nfloat g_device_giga_bandwidth;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n\n// Selection functor type\ntemplate \nstruct LessThan\n{\n T compare;\n\n __host__ __device__ __forceinline__\n LessThan(T compare) : compare(compare) {}\n\n __host__ __device__ __forceinline__\n bool operator()(const T &a) const {\n return (a < compare);\n }\n};\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB DeviceSelect entrypoints\n//---------------------------------------------------------------------\n\n\n/**\n * Dispatch to select if entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n Int2Type /*is_flagged*/,\n Int2Type /*is_partition*/,\n int timing_timing_iterations,\n size_t* /*d_temp_storage_bytes*/,\n hipError_t* /*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n FlagIteratorT /*d_flags*/,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items,\n SelectOpT select_op)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceSelect::If(d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items, select_op);\n }\n return error;\n}\n\n\n/**\n * Dispatch to partition if entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n Int2Type /*is_flagged*/,\n Int2Type /*is_partition*/,\n int timing_timing_iterations,\n size_t* /*d_temp_storage_bytes*/,\n hipError_t* /*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n FlagIteratorT /*d_flags*/,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items,\n SelectOpT select_op)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DevicePartition::If(d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items, select_op);\n }\n return error;\n}\n\n\n/**\n * Dispatch to select flagged entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n Int2Type /*is_flagged*/,\n Int2Type /*partition*/,\n int timing_timing_iterations,\n size_t* /*d_temp_storage_bytes*/,\n hipError_t* /*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n FlagIteratorT d_flags,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items,\n SelectOpT /*select_op*/)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceSelect::Flagged(d_temp_storage, temp_storage_bytes, d_in, d_flags, d_out, d_num_selected_out, num_items);\n }\n return error;\n}\n\n\n/**\n * Dispatch to partition flagged entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n Int2Type /*is_flagged*/,\n Int2Type /*partition*/,\n int timing_timing_iterations,\n size_t* /*d_temp_storage_bytes*/,\n hipError_t* /*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n InputIteratorT d_in,\n FlagIteratorT d_flags,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items,\n SelectOpT /*select_op*/)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DevicePartition::Flagged(d_temp_storage, temp_storage_bytes, d_in, d_flags, d_out, d_num_selected_out, num_items);\n }\n return error;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceSelect\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type cub_backend,\n IsFlaggedTag is_flagged,\n IsPartitionTag is_partition,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n InputIteratorT d_in,\n FlagIteratorT d_flags,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items,\n SelectOpT select_op)\n{\n *d_cdp_error = Dispatch(cub_backend,\n is_flagged,\n is_partition,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_flags,\n d_out,\n d_num_selected_out,\n num_items,\n select_op);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Dispatch to CDP kernel\n */\ntemplate \nhipError_t Dispatch(Int2Type /*dispatch_to*/,\n IsFlaggedTag is_flagged,\n IsPartitionTag is_partition,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n FlagIteratorT d_flags,\n OutputIteratorT d_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items,\n SelectOpT select_op)\n{\n // Invoke kernel to invoke device-side dispatch\n hipError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n Int2Type{},\n is_flagged,\n is_partition,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_flags,\n d_out,\n d_num_selected_out,\n num_items,\n select_op);\n CubDebugExit(retval);\n\n // Copy out temp_storage_bytes\n CubDebugExit(hipMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n hipMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(hipMemcpy(&retval,\n d_cdp_error,\n sizeof(hipError_t) * 1,\n hipMemcpyDeviceToHost));\n return retval;\n}\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\n\n/**\n * Initialize problem\n */\ntemplate \nvoid Initialize(\n T* h_in,\n int num_items)\n{\n for (int i = 0; i < num_items; ++i)\n {\n // Initialize each item to a randomly selected value from [0..126]\n unsigned int value;\n RandomBits(value, 0, 0, 7);\n if (value == 127)\n value = 126;\n InitValue(INTEGER_SEED, h_in[i], value);\n }\n\n if (g_verbose)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n\n/**\n * Solve selection problem (and set corresponding flags)\n */\ntemplate <\n typename InputIteratorT,\n typename FlagIteratorT,\n typename SelectOpT,\n typename T>\nint Solve(\n InputIteratorT h_in,\n SelectOpT select_op,\n T* h_reference,\n FlagIteratorT h_flags,\n int num_items)\n{\n int num_selected = 0;\n for (int i = 0; i < num_items; ++i)\n {\n if ((h_flags[i] = select_op(h_in[i])))\n {\n h_reference[num_selected] = h_in[i];\n num_selected++;\n }\n else\n {\n h_reference[num_items - (i - num_selected) - 1] = h_in[i];\n }\n }\n\n return num_selected;\n}\n\n\n\n/**\n * Test DeviceSelect for a given problem input\n */\ntemplate <\n Backend BACKEND,\n bool IS_FLAGGED,\n bool IS_PARTITION,\n typename DeviceInputIteratorT,\n typename FlagT,\n typename SelectOpT,\n typename T>\nvoid Test(\n DeviceInputIteratorT d_in,\n FlagT* h_flags,\n SelectOpT select_op,\n T* h_reference,\n int num_selected,\n int num_items)\n{\n // Allocate device flags, output, and num-selected\n FlagT* d_flags = NULL;\n T* d_out = NULL;\n int* d_num_selected_out = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_flags, sizeof(FlagT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(T) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_num_selected_out, sizeof(int)));\n\n // Allocate CDP device arrays\n size_t* d_temp_storage_bytes = NULL;\n hipError_t* d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(hipError_t) * 1));\n\n // Allocate temporary storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(Int2Type(), Int2Type(), Int2Type(), 1, d_temp_storage_bytes, d_cdp_error,\n d_temp_storage, temp_storage_bytes, d_in, d_flags, d_out, d_num_selected_out, num_items, select_op));\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n // Copy flags and clear device output array\n CubDebugExit(hipMemcpy(d_flags, h_flags, sizeof(FlagT) * num_items, hipMemcpyHostToDevice));\n CubDebugExit(hipMemset(d_out, 0, sizeof(T) * num_items));\n CubDebugExit(hipMemset(d_num_selected_out, 0, sizeof(int)));\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(Int2Type(), Int2Type(), Int2Type(), 1, d_temp_storage_bytes, d_cdp_error,\n d_temp_storage, temp_storage_bytes, d_in, d_flags, d_out, d_num_selected_out, num_items, select_op));\n\n // Check for correctness (and display results, if specified)\n int compare1 = (IS_PARTITION) ?\n CompareDeviceResults(h_reference, d_out, num_items, true, g_verbose) :\n CompareDeviceResults(h_reference, d_out, num_selected, true, g_verbose);\n printf(\"\\t Data %s\\n\", compare1 ? \"FAIL\" : \"PASS\");\n\n int compare2 = CompareDeviceResults(&num_selected, d_num_selected_out, 1, true, g_verbose);\n printf(\"\\t Count %s\\n\", compare2 ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Performance\n GpuTimer gpu_timer;\n gpu_timer.Start();\n CubDebugExit(Dispatch(Int2Type(), Int2Type(), Int2Type(), g_timing_iterations, d_temp_storage_bytes, d_cdp_error,\n d_temp_storage, temp_storage_bytes, d_in, d_flags, d_out, d_num_selected_out, num_items, select_op));\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n\n // Display performance\n if (g_timing_iterations > 0)\n {\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n int num_output_items = (IS_PARTITION) ? num_items : num_selected;\n int num_flag_items = (IS_FLAGGED) ? num_items : 0;\n size_t num_bytes = sizeof(T) * (num_items + num_output_items) + sizeof(FlagT) * num_flag_items;\n float giga_bandwidth = float(num_bytes) / avg_millis / 1000.0f / 1000.0f;\n\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s, %.1f%% peak\", avg_millis, giga_rate, giga_bandwidth, giga_bandwidth / g_device_giga_bandwidth * 100.0);\n }\n printf(\"\\n\\n\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Cleanup\n if (d_flags) CubDebugExit(g_allocator.DeviceFree(d_flags));\n if (d_out) CubDebugExit(g_allocator.DeviceFree(d_out));\n if (d_num_selected_out) CubDebugExit(g_allocator.DeviceFree(d_num_selected_out));\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare1 | compare2);\n}\n\n\n/**\n * Test on pointer type\n */\ntemplate <\n Backend BACKEND,\n bool IS_FLAGGED,\n bool IS_PARTITION,\n typename T>\nvoid TestPointer(\n int num_items,\n float select_ratio)\n{\n typedef char FlagT;\n\n // Allocate host arrays\n T* h_in = new T[num_items];\n FlagT* h_flags = new FlagT[num_items];\n T* h_reference = new T[num_items];\n\n // Initialize input\n Initialize(h_in, num_items);\n\n // Select a comparison value that is select_ratio through the space of [0,127]\n T compare;\n if (select_ratio <= 0.0)\n InitValue(INTEGER_SEED, compare, 0); // select none\n else if (select_ratio >= 1.0)\n InitValue(INTEGER_SEED, compare, 127); // select all\n else\n InitValue(INTEGER_SEED, compare, int(double(double(127) * select_ratio)));\n\n LessThan select_op(compare);\n int num_selected = Solve(h_in, select_op, h_reference, h_flags, num_items);\n\n if (g_verbose) std::cout << \"\\nComparison item: \" << compare << \"\\n\";\n printf(\"\\nPointer %s cub::%s::%s %d items, %d selected (select ratio %.3f), %s %d-byte elements\\n\",\n (IS_PARTITION) ? \"DevicePartition\" : \"DeviceSelect\",\n (IS_FLAGGED) ? \"Flagged\" : \"If\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items, num_selected, float(num_selected) / num_items, typeid(T).name(), (int) sizeof(T));\n fflush(stdout);\n\n // Allocate problem device arrays\n T *d_in = NULL;\n\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_in, sizeof(T) * num_items));\n\n // Initialize device input\n CubDebugExit(hipMemcpy(d_in, h_in, sizeof(T) * num_items, hipMemcpyHostToDevice));\n\n // Run Test\n Test(d_in, h_flags, select_op, h_reference, num_selected, num_items);\n\n // Cleanup\n if (h_in) delete[] h_in;\n if (h_reference) delete[] h_reference;\n if (h_flags) delete[] h_flags;\n if (d_in) CubDebugExit(g_allocator.DeviceFree(d_in));\n}\n\n\n/**\n * Test on iterator type\n */\ntemplate <\n Backend BACKEND,\n bool IS_FLAGGED,\n bool IS_PARTITION,\n typename T>\nvoid TestIterator(\n int num_items,\n float select_ratio)\n{\n typedef char FlagT;\n\n // Allocate host arrays\n T* h_reference = new T[num_items];\n FlagT* h_flags = new FlagT[num_items];\n\n // Use counting iterator as the input\n CountingInputIterator h_in(0);\n\n // Select a comparison value that is select_ratio through the space of [0,127]\n T compare;\n if (select_ratio <= 0.0)\n InitValue(INTEGER_SEED, compare, 0); // select none\n else if (select_ratio >= 1.0)\n InitValue(INTEGER_SEED, compare, 127); // select all\n else\n InitValue(INTEGER_SEED, compare, int(double(double(127) * select_ratio)));\n\n LessThan select_op(compare);\n int num_selected = Solve(h_in, select_op, h_reference, h_flags, num_items);\n\n if (g_verbose) std::cout << \"\\nComparison item: \" << compare << \"\\n\";\n printf(\"\\nIterator %s cub::%s::%s %d items, %d selected (select ratio %.3f), %s %d-byte elements\\n\",\n (IS_PARTITION) ? \"DevicePartition\" : \"DeviceSelect\",\n (IS_FLAGGED) ? \"Flagged\" : \"If\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items, num_selected, float(num_selected) / num_items, typeid(T).name(), (int) sizeof(T));\n fflush(stdout);\n\n // Run Test\n Test(h_in, h_flags, select_op, h_reference, num_selected, num_items);\n\n // Cleanup\n if (h_reference) delete[] h_reference;\n if (h_flags) delete[] h_flags;\n}\n\n\n/**\n * Test different selection ratios\n */\ntemplate <\n Backend BACKEND,\n bool IS_FLAGGED,\n bool IS_PARTITION,\n typename T>\nvoid Test(\n int num_items)\n{\n for (float select_ratio = 0.0f; select_ratio <= 1.0f; select_ratio += 0.2f)\n {\n TestPointer(num_items, select_ratio);\n }\n}\n\n\n/**\n * Test (select vs. partition) and (flagged vs. functor)\n */\ntemplate <\n Backend BACKEND,\n typename T>\nvoid TestMethod(\n int num_items)\n{\n // Functor\n Test(num_items);\n Test(num_items);\n\n // Flagged\n Test(num_items);\n Test(num_items);\n}\n\n\n/**\n * Test different dispatch\n */\ntemplate <\n typename T>\nvoid TestOp(\n int num_items)\n{\n#if TEST_CDP == 0\n TestMethod(num_items);\n#elif TEST_CDP == 1\n TestMethod(num_items);\n#endif // TEST_CDP\n}\n\n\n/**\n * Test different input sizes\n */\ntemplate \nvoid Test(\n int num_items)\n{\n if (num_items < 0)\n {\n TestOp(0);\n TestOp(1);\n TestOp(100);\n TestOp(10000);\n TestOp(1000000);\n }\n else\n {\n TestOp(num_items);\n }\n}\n\ntemplate\nstruct pair_to_col_t\n{\n __host__ __device__ T0 operator()(const thrust::tuple &in)\n {\n return thrust::get<0>(in);\n }\n};\n\ntemplate\nstruct select_t {\n __host__ __device__ bool operator()(const thrust::tuple &in) {\n return static_cast(thrust::get<0>(in)) > thrust::get<1>(in);\n }\n};\n\ntemplate \nvoid TestMixedOp(int num_items)\n{\n const T0 target_value = static_cast(42);\n thrust::device_vector col_a(num_items, target_value);\n thrust::device_vector col_b(num_items, static_cast(4.2));\n\n thrust::device_vector result(num_items);\n\n auto in = thrust::make_zip_iterator(col_a.begin(), col_b.begin());\n auto out = thrust::make_transform_output_iterator(result.begin(), pair_to_col_t{});\n\n void *d_tmp_storage {};\n std::size_t tmp_storage_size{};\n hipcub::DeviceSelect::If(\n d_tmp_storage, tmp_storage_size,\n in, out, thrust::make_discard_iterator(),\n num_items, select_t{});\n\n thrust::device_vector tmp_storage(tmp_storage_size);\n d_tmp_storage = thrust::raw_pointer_cast(tmp_storage.data());\n\n hipcub::DeviceSelect::If(\n d_tmp_storage, tmp_storage_size,\n in, out, thrust::make_discard_iterator(),\n num_items, select_t{});\n\n AssertEquals(num_items, thrust::count(result.begin(), result.end(), target_value));\n}\n\n/**\n * Test different input sizes\n */\ntemplate \nvoid TestMixed(int num_items)\n{\n if (num_items < 0)\n {\n TestMixedOp(0);\n TestMixedOp(1);\n TestMixedOp(100);\n TestMixedOp(10000);\n TestMixedOp(1000000);\n }\n else\n {\n TestMixedOp(num_items);\n }\n}\n\nvoid TestFlagsNormalization()\n{\n const int num_items = 1024 * 1024;\n thrust::device_vector result(num_items);\n\n void *d_tmp_storage{};\n std::size_t tmp_storage_size{};\n CubDebugExit(\n hipcub::DeviceSelect::Flagged(d_tmp_storage,\n tmp_storage_size,\n hipcub::CountingInputIterator(0), // in\n hipcub::CountingInputIterator(1), // flags\n thrust::raw_pointer_cast(result.data()), // out\n thrust::make_discard_iterator(), // num_out\n num_items));\n\n thrust::device_vector tmp_storage(tmp_storage_size);\n d_tmp_storage = thrust::raw_pointer_cast(tmp_storage.data());\n\n CubDebugExit(\n hipcub::DeviceSelect::Flagged(d_tmp_storage,\n tmp_storage_size,\n hipcub::CountingInputIterator(0), // in\n hipcub::CountingInputIterator(1), // flags\n thrust::raw_pointer_cast(result.data()), // out\n thrust::make_discard_iterator(), // num_out\n num_items));\n\n AssertTrue(thrust::equal(result.begin(),\n result.end(),\n thrust::make_counting_iterator(0)));\n}\n\nvoid TestFlagsAliasingInPartition()\n{\n int h_items[]{0, 1, 0, 2, 0, 3, 0, 4, 0, 5};\n constexpr int num_items = sizeof(h_items) / sizeof(h_items[0]);\n\n int *d_in{};\n int *d_out{};\n\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_in, sizeof(h_items)));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_out, sizeof(h_items)));\n\n CubDebugExit(\n hipMemcpy(d_in, h_items, sizeof(h_items), hipMemcpyHostToDevice));\n\n // alias flags and keys\n int *d_flags = d_in;\n\n void *d_tmp_storage{};\n std::size_t tmp_storage_size{};\n\n CubDebugExit(\n hipcub::DevicePartition::Flagged(d_tmp_storage,\n tmp_storage_size,\n d_in,\n d_flags,\n d_out,\n thrust::make_discard_iterator(), // num_out\n num_items));\n\n thrust::device_vector tmp_storage(tmp_storage_size);\n d_tmp_storage = thrust::raw_pointer_cast(tmp_storage.data());\n\n CubDebugExit(\n hipcub::DevicePartition::Flagged(d_tmp_storage,\n tmp_storage_size,\n d_in,\n d_flags,\n d_out,\n thrust::make_discard_iterator(), // num_out\n num_items));\n\n AssertTrue(thrust::equal(thrust::device,\n d_out,\n d_out + num_items / 2,\n thrust::make_counting_iterator(1)));\n\n AssertEquals(\n thrust::count(thrust::device, d_out + num_items / 2, d_out + num_items, 0),\n num_items / 2);\n\n CubDebugExit(g_allocator.DeviceFree(d_out));\n CubDebugExit(g_allocator.DeviceFree(d_in));\n}\n\nstruct Odd\n{\n __host__ __device__ bool operator()(int v) const { return v % 2; }\n};\n\nvoid TestIfInPlace()\n{\n const int num_items = 4 * 1024 * 1024;\n const int num_iters = 42;\n\n thrust::device_vector num_out(1);\n thrust::device_vector data(num_items);\n thrust::device_vector reference(num_items);\n thrust::device_vector reference_out(1);\n\n thrust::sequence(data.begin(), data.end());\n\n Odd op{};\n\n int *d_num_out = thrust::raw_pointer_cast(num_out.data());\n int *d_data = thrust::raw_pointer_cast(data.data());\n int *d_reference = thrust::raw_pointer_cast(reference.data());\n int *d_reference_out = thrust::raw_pointer_cast(reference_out.data());\n\n void *d_tmp_storage{};\n std::size_t tmp_storage_size{};\n\n CubDebugExit(\n hipcub::DeviceSelect::If(d_tmp_storage,\n tmp_storage_size,\n d_data,\n d_num_out,\n num_items,\n op));\n\n thrust::device_vector tmp_storage(tmp_storage_size);\n d_tmp_storage = thrust::raw_pointer_cast(tmp_storage.data());\n\n thrust::default_random_engine g{};\n\n for (int iter = 0; iter < num_iters; iter++)\n {\n thrust::shuffle(data.begin(), data.end(), g);\n\n CubDebugExit(\n hipcub::DeviceSelect::If(d_tmp_storage,\n tmp_storage_size,\n d_data,\n d_reference,\n d_reference_out,\n num_items,\n op));\n\n CubDebugExit(\n hipcub::DeviceSelect::If(d_tmp_storage,\n tmp_storage_size,\n d_data,\n d_num_out,\n num_items,\n op));\n\n AssertEquals(num_out, reference_out);\n const int num_selected = num_out[0];\n\n const bool match_reference = thrust::equal(reference.begin(),\n reference.begin() + num_selected,\n data.begin());\n AssertTrue(match_reference);\n }\n}\n\nvoid TestFlaggedInPlace()\n{\n const int num_items = 4 * 1024 * 1024;\n const int num_iters = 42;\n\n thrust::device_vector num_out(1);\n thrust::device_vector data(num_items);\n thrust::device_vector flags(num_items);\n\n int h_num_out{};\n int *d_num_out = thrust::raw_pointer_cast(num_out.data());\n int *d_data = thrust::raw_pointer_cast(data.data());\n bool *d_flags = thrust::raw_pointer_cast(flags.data());\n\n void *d_tmp_storage{};\n std::size_t tmp_storage_size{};\n\n CubDebugExit(\n hipcub::DeviceSelect::Flagged(d_tmp_storage,\n tmp_storage_size,\n d_data,\n d_flags,\n d_num_out,\n num_items));\n\n thrust::device_vector tmp_storage(tmp_storage_size);\n d_tmp_storage = thrust::raw_pointer_cast(tmp_storage.data());\n\n thrust::default_random_engine g{};\n\n for (int iter = 0; iter < num_iters; iter++)\n {\n const int num_selected = RandomValue(num_items);\n\n thrust::sequence(data.begin(), data.end());\n thrust::fill(flags.begin(), flags.begin() + num_selected, true);\n thrust::fill(flags.begin() + num_selected, flags.end(), false);\n thrust::shuffle(flags.begin(), flags.end(), g);\n\n CubDebugExit(\n hipcub::DeviceSelect::Flagged(d_tmp_storage,\n tmp_storage_size,\n d_data,\n d_flags,\n d_num_out,\n num_items));\n\n hipMemcpy(&h_num_out, d_num_out, sizeof(int), hipMemcpyDeviceToHost);\n\n AssertEquals(num_selected, h_num_out);\n\n auto selection_perm_begin = thrust::make_permutation_iterator(flags.begin(),\n data.begin());\n auto selection_perm_end = selection_perm_begin + num_selected;\n\n AssertEquals(num_selected,\n thrust::count(selection_perm_begin, selection_perm_end, true));\n }\n}\n\nvoid TestFlaggedInPlaceWithAliasedFlags()\n{\n const int num_items = 1024 * 1024;\n const int num_iters = 42;\n\n thrust::device_vector num_out(1);\n thrust::device_vector data(num_items);\n thrust::device_vector reference(num_items);\n thrust::device_vector flags(num_items);\n\n int h_num_out{};\n int *d_num_out = thrust::raw_pointer_cast(num_out.data());\n int *d_data = thrust::raw_pointer_cast(data.data());\n int *d_flags = d_data; // alias\n int *d_allocated_flags = thrust::raw_pointer_cast(data.data());\n int *d_reference = thrust::raw_pointer_cast(reference.data());\n\n void *d_tmp_storage{};\n std::size_t tmp_storage_size{};\n\n CubDebugExit(\n hipcub::DeviceSelect::Flagged(d_tmp_storage,\n tmp_storage_size,\n d_data,\n d_flags,\n d_num_out,\n num_items));\n\n thrust::device_vector tmp_storage(tmp_storage_size);\n d_tmp_storage = thrust::raw_pointer_cast(tmp_storage.data());\n\n thrust::default_random_engine g{};\n\n for (int iter = 0; iter < num_iters; iter++)\n {\n const int num_selected = RandomValue(num_items);\n\n thrust::sequence(data.begin(), data.begin() + num_selected, 1);\n thrust::fill(data.begin() + num_selected, data.end(), 0);\n thrust::shuffle(data.begin(), data.end(), g);\n\n CubDebugExit(\n hipcub::DeviceSelect::Flagged(d_tmp_storage,\n tmp_storage_size,\n d_data, // in\n d_allocated_flags,\n d_reference, // out\n d_num_out,\n num_items));\n\n CubDebugExit(\n hipcub::DeviceSelect::Flagged(d_tmp_storage,\n tmp_storage_size,\n d_data,\n d_flags,\n d_num_out,\n num_items));\n\n hipMemcpy(&h_num_out, d_num_out, sizeof(int), hipMemcpyDeviceToHost);\n\n AssertEquals(num_selected, h_num_out);\n\n const bool match_reference = thrust::equal(reference.begin(),\n reference.begin() + num_selected,\n data.begin());\n AssertTrue(match_reference);\n }\n}\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n int num_items = -1;\n float select_ratio = 0.5;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n args.GetCmdLineArgument(\"ratio\", select_ratio);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--ratio=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n g_device_giga_bandwidth = args.device_giga_bandwidth;\n printf(\"\\n\");\n\n // %PARAM% TEST_CDP cdp 0:1\n\n TestFlagsAliasingInPartition();\n\n TestFlaggedInPlace();\n TestFlaggedInPlaceWithAliasedFlags();\n TestIfInPlace();\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n\n TestMixed(num_items);\n TestFlagsNormalization();\n\n return 0;\n}\n\n\n\n", "source": "/kaggle/working/cub/test/test_device_select_if.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2021 NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceScan utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \"test_util.h\"\n\n#include \n#include \n#include \n\nusing namespace cub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\ndouble g_device_giga_bandwidth;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n\nenum AliasMode\n{\n AliasNone, // output is allocated\n AliasKeys, // output is an alias of input keys\n AliasValues // output is an alias of input values\n};\n\n\n/**\n * \\brief WrapperFunctor (for precluding test-specialized dispatch to *Sum variants)\n */\ntemplate\nstruct WrapperFunctor\n{\n OpT op;\n\n WrapperFunctor(OpT op) : op(op) {}\n\n template \n __host__ __device__ __forceinline__ auto operator()(const T &a, const U &b) const\n -> decltype(op(a, b))\n {\n return static_cast(op(a, b));\n }\n};\n\n/**\n * \\brief DivideByFiveFunctor (used by TestIterator)\n */\ntemplate\nstruct DivideByFiveFunctor\n{\n template \n __host__ __device__ __forceinline__ OutputT operator()(const T &a) const\n {\n return static_cast(a / 5);\n }\n};\n\n/**\n * \\brief Mod2Equality (used for non-bool keys to make keys more likely to equal each other)\n */\nstruct Mod2Equality\n{\n template \n __host__ __device__ __forceinline__ T operator()(const T &a, const T &b) const\n {\n return (a % 2) == (b % 2);\n }\n};\n\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB DeviceScan entrypoints\n//---------------------------------------------------------------------\n\n/**\n * Dispatch to exclusive scan entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n IsPrimitiveT /*is_primitive*/,\n int timing_timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n KeysInputIteratorT d_keys_in,\n ValuesInputIteratorT d_values_in,\n ValuesOutputIteratorT d_values_out,\n ScanOpT scan_op,\n InitialValueT initial_value,\n OffsetT num_items,\n EqualityOpT equality_op)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::ExclusiveScanByKey(d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_values_out, scan_op, initial_value, num_items, equality_op);\n }\n return error;\n}\n\n\n/**\n * Dispatch to exclusive sum entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n Int2Type /*is_primitive*/,\n int timing_timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n KeysInputIteratorT d_keys_in,\n ValuesInputIteratorT d_values_in,\n ValuesOutputIteratorT d_values_out,\n Sum /*scan_op*/,\n InitialValueT /*initial_value*/,\n OffsetT num_items,\n EqualityOpT equality_op)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::ExclusiveSumByKey(d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_values_out, num_items, equality_op);\n }\n return error;\n}\n\n\n/**\n * Dispatch to inclusive scan entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n IsPrimitiveT /*is_primitive*/,\n int timing_timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n KeysInputIteratorT d_keys_in,\n ValuesInputIteratorT d_values_in,\n ValuesOutputIteratorT d_values_out,\n ScanOpT scan_op,\n NullType /*initial_value*/,\n OffsetT num_items,\n EqualityOpT equality_op)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::InclusiveScanByKey(d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_values_out, scan_op, num_items, equality_op);\n }\n return error;\n}\n\n/**\n * Dispatch to inclusive sum entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n Int2Type /*is_primitive*/,\n int timing_timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n KeysInputIteratorT d_keys_in,\n ValuesInputIteratorT d_values_in,\n ValuesOutputIteratorT d_values_out,\n Sum /*scan_op*/,\n NullType /*initial_value*/,\n OffsetT num_items,\n EqualityOpT equality_op)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::InclusiveSumByKey(d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_values_out, num_items, equality_op);\n }\n return error;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceScan\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type cub_backend,\n IsPrimitiveT is_primitive,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n KeysInputIteratorT d_keys_in,\n ValuesInputIteratorT d_values_in,\n ValuesOutputIteratorT d_values_out,\n ScanOpT scan_op,\n InitialValueT initial_value,\n OffsetT num_items,\n EqualityOpT equality_op)\n{\n *d_cdp_error = Dispatch(cub_backend,\n is_primitive,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_values_in,\n d_values_out,\n scan_op,\n initial_value,\n num_items,\n equality_op);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Dispatch to CDP kernel\n */\ntemplate \ncudaError_t Dispatch(Int2Type /*dispatch_to*/,\n IsPrimitiveT is_primitive,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n KeysInputIteratorT d_keys_in,\n ValuesInputIteratorT d_values_in,\n ValuesOutputIteratorT d_values_out,\n ScanOpT scan_op,\n InitialValueT initial_value,\n OffsetT num_items,\n EqualityOpT equality_op)\n{\n // Invoke kernel to invoke device-side dispatch\n cudaError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n Int2Type{},\n is_primitive,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_values_in,\n d_values_out,\n scan_op,\n initial_value,\n num_items,\n equality_op);\n CubDebugExit(retval);\n\n // Copy out temp_storage_bytes\n CubDebugExit(cudaMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n cudaMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(cudaMemcpy(&retval,\n d_cdp_error,\n sizeof(cudaError_t) * 1,\n cudaMemcpyDeviceToHost));\n return retval;\n}\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\n\n/**\n * Initialize problem\n */\ntemplate \nvoid Initialize(\n GenMode gen_mode,\n T *h_in,\n int num_items)\n{\n for (int i = 0; i < num_items; ++i)\n {\n InitValue(gen_mode, h_in[i], i);\n }\n\n if (g_verbose)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n/**\n * Solve exclusive-scan problem\n */\ntemplate <\n typename KeysInputIteratorT,\n typename ValuesInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT,\n typename EqualityOpT>\nvoid Solve(\n KeysInputIteratorT h_keys_in,\n ValuesInputIteratorT h_values_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n using ValueT = cub::detail::value_t;\n using AccumT = cub::detail::accumulator_t;\n\n if (num_items > 0)\n {\n for (int i = 0; i < num_items;) {\n AccumT val = static_cast(h_values_in[i]);\n h_reference[i] = initial_value;\n AccumT inclusive = static_cast(scan_op(initial_value, val));\n\n ++i;\n\n for (; i < num_items && equality_op(h_keys_in[i - 1], h_keys_in[i]); ++i)\n {\n val = static_cast(h_values_in[i]);\n h_reference[i] = static_cast(inclusive);\n inclusive = static_cast(scan_op(inclusive, val));\n }\n }\n }\n}\n\n\n/**\n * Solve inclusive-scan problem\n */\ntemplate <\n typename KeysInputIteratorT,\n typename ValuesInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename EqualityOpT>\nvoid Solve(\n KeysInputIteratorT h_keys_in,\n ValuesInputIteratorT h_values_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n NullType /*initial_value*/,\n EqualityOpT equality_op)\n{\n using ValueT = cub::detail::value_t;\n using AccumT = cub::detail::accumulator_t;\n\n if (num_items > 0)\n {\n for (int i = 0; i < num_items;) {\n AccumT inclusive = h_values_in[i];\n h_reference[i] = static_cast(inclusive);\n\n ++i;\n\n for (; i < num_items && equality_op(h_keys_in[i - 1], h_keys_in[i]); ++i)\n {\n AccumT val = h_values_in[i];\n inclusive = static_cast(scan_op(inclusive, val));\n h_reference[i] = static_cast(inclusive);\n }\n }\n }\n}\n\ntemplate\nstruct AllocateOutput {\n static void run(OutputT *&d_out, DeviceInputIteratorT, int num_items) {\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(OutputT) * num_items));\n }\n};\n\ntemplate\nstruct AllocateOutput {\n static void run(OutputT *&d_out, OutputT *d_in, int /* num_items */) {\n d_out = d_in;\n }\n};\n\n/**\n * Test DeviceScan for a given problem input\n */\ntemplate <\n Backend BACKEND,\n typename KeysInputIteratorT,\n typename ValuesInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT,\n typename EqualityOpT,\n AliasMode Mode=AliasNone>\nvoid Test(\n KeysInputIteratorT d_keys_in,\n ValuesInputIteratorT d_values_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n using KeyT = cub::detail::value_t;\n using InputT = cub::detail::value_t;\n\n // Allocate device output array\n OutputT *d_values_out = NULL;\n\n if (Mode == AliasKeys)\n {\n AllocateOutput::run(\n d_values_out,\n d_keys_in,\n num_items);\n }\n else\n {\n AllocateOutput::run(\n d_values_out,\n d_values_in,\n num_items);\n }\n\n // Allocate CDP device arrays\n size_t *d_temp_storage_bytes = NULL;\n cudaError_t *d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(cudaError_t) * 1));\n\n // Allocate temporary storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(\n Int2Type(),\n Int2Type::PRIMITIVE>(),\n 1,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_values_in,\n d_values_out,\n scan_op,\n initial_value,\n num_items,\n equality_op));\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n // Clear device output array\n if (Mode == AliasNone)\n {\n CubDebugExit(cudaMemset(d_values_out, 0, sizeof(OutputT) * num_items));\n }\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(\n Int2Type(),\n Int2Type::PRIMITIVE>(),\n 1,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_values_in,\n d_values_out,\n scan_op,\n initial_value,\n num_items,\n equality_op));\n\n // Check for correctness (and display results, if specified)\n const int compare = CompareDeviceResults(h_reference,\n d_values_out,\n num_items,\n true,\n g_verbose);\n\n printf(\"\\t%s\", compare ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Display performance\n if (g_timing_iterations > 0)\n {\n // Performance\n GpuTimer gpu_timer;\n gpu_timer.Start();\n CubDebugExit(Dispatch(Int2Type(),\n Int2Type::PRIMITIVE>(),\n g_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_values_in,\n d_values_out,\n scan_op,\n initial_value,\n num_items,\n equality_op));\n\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n float giga_bandwidth = giga_rate * (sizeof(InputT) + sizeof(OutputT));\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s, %.1f%% peak\",\n avg_millis, giga_rate, giga_bandwidth, giga_bandwidth / g_device_giga_bandwidth * 100.0);\n }\n\n printf(\"\\n\\n\");\n\n // Cleanup\n if (Mode == AliasNone)\n {\n if (d_values_out) \n {\n CubDebugExit(g_allocator.DeviceFree(d_values_out));\n }\n }\n\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare);\n}\n\ntemplate \nvoid TestInplaceValues(KeysInputIteratorT d_keys_in,\n OutputT *d_values_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n Test(d_keys_in,\n d_values_in,\n h_reference,\n num_items,\n scan_op,\n initial_value,\n equality_op);\n}\n\ntemplate \nvoid TestInplaceValues(KeysInputIteratorT,\n ValuesInputIteratorT,\n OutputT *,\n int,\n ScanOpT,\n InitialValueT,\n EqualityOpT)\n{}\n\ntemplate \nvoid TestInplaceKeys(T *d_keys_in,\n ValuesInputIteratorT d_values_in,\n T *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n Test(d_keys_in,\n d_values_in,\n h_reference,\n num_items,\n scan_op,\n initial_value,\n equality_op);\n}\n\ntemplate \nvoid TestInplaceKeys(KeysInputIteratorT,\n ValuesInputIteratorT,\n OutputT *,\n int,\n ScanOpT,\n InitialValueT,\n EqualityOpT)\n{}\n\n/**\n * Test DeviceScan on pointer type\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT,\n typename EqualityOpT>\nvoid TestPointer(\n int num_items,\n GenMode gen_mode,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n printf(\"\\nPointer %s %s cub::DeviceScan::%s %d items, %s->%s (%d->%d bytes) , gen-mode %s\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n (std::is_same::value) ? \"Inclusive\" : \"Exclusive\",\n (std::is_same::value) ? \"Sum\" : \"Scan\",\n num_items,\n typeid(InputT).name(), typeid(OutputT).name(), (int) sizeof(InputT), (int) sizeof(OutputT),\n (gen_mode == RANDOM) ? \"RANDOM\" : (gen_mode == INTEGER_SEED) ? \"SEQUENTIAL\" : \"HOMOGENOUS\");\n fflush(stdout);\n\n // Allocate host arrays\n KeyT* h_keys_in = new KeyT[num_items];\n InputT* h_values_in = new InputT[num_items];\n OutputT* h_reference = new OutputT[num_items];\n\n // Initialize problem and solution\n Initialize(gen_mode, h_keys_in, num_items);\n Initialize(gen_mode, h_values_in, num_items);\n\n // If the output type is primitive and the operator is cub::Sum, the test\n // dispatcher throws away scan_op and initial_value for exclusive scan.\n // Without an initial_value arg, the accumulator switches to the input value\n // type.\n // Do the same thing here:\n if (Traits::PRIMITIVE &&\n std::is_same::value &&\n !std::is_same::value)\n {\n Solve(h_keys_in, h_values_in, h_reference, num_items, cub::Sum{}, InputT{}, equality_op);\n }\n else\n {\n Solve(h_keys_in, h_values_in, h_reference, num_items, scan_op, initial_value, equality_op);\n }\n\n // Allocate problem device arrays\n KeyT *d_keys_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys_in, sizeof(KeyT) * num_items));\n InputT *d_values_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_values_in, sizeof(InputT) * num_items));\n\n // Initialize device input\n CubDebugExit(cudaMemcpy(d_keys_in, h_keys_in, sizeof(KeyT) * num_items, cudaMemcpyHostToDevice));\n CubDebugExit(cudaMemcpy(d_values_in, h_values_in, sizeof(InputT) * num_items, cudaMemcpyHostToDevice));\n\n // Run Test\n Test(d_keys_in,\n d_values_in,\n h_reference,\n num_items,\n scan_op,\n initial_value,\n equality_op);\n\n // Test in/out values aliasing\n TestInplaceValues(d_keys_in,\n d_values_in, \n h_reference,\n num_items,\n scan_op,\n initial_value,\n equality_op);\n\n CubDebugExit(cudaMemcpy(d_values_in, h_values_in, sizeof(InputT) * num_items, cudaMemcpyHostToDevice));\n\n // Test keys/values aliasing (should go last, changes keys)\n TestInplaceKeys(d_keys_in,\n d_values_in,\n h_reference,\n num_items,\n scan_op,\n initial_value,\n equality_op);\n\n // Cleanup\n if (h_keys_in) delete[] h_keys_in;\n if (h_values_in) delete[] h_values_in;\n if (h_reference) delete[] h_reference;\n if (d_keys_in) CubDebugExit(g_allocator.DeviceFree(d_keys_in));\n if (d_values_in) CubDebugExit(g_allocator.DeviceFree(d_values_in));\n}\n\n\n/**\n * Test DeviceScan on iterator type\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT,\n typename EqualityOpT>\nvoid TestIterator(\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n printf(\"\\nIterator %s %s cub::DeviceScan::%s %d items, %s->%s (%d->%d bytes)\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n (std::is_same::value) ? \"Inclusive\" : \"Exclusive\",\n (std::is_same::value) ? \"Sum\" : \"Scan\",\n num_items,\n typeid(InputT).name(), typeid(OutputT).name(), (int) sizeof(InputT), (int) sizeof(OutputT));\n fflush(stdout);\n\n // Use a counting iterator followed by div as the keys\n using CountingIterT = CountingInputIterator;\n CountingIterT h_keys_in_helper(0);\n TransformInputIterator, CountingIterT> h_keys_in(h_keys_in_helper, DivideByFiveFunctor());\n\n // Use a constant iterator as the input\n InputT val = InputT();\n ConstantInputIterator h_values_in(val);\n\n // Allocate host arrays\n OutputT* h_reference = new OutputT[num_items];\n\n // Initialize problem and solution\n Solve(h_keys_in, h_values_in, h_reference, num_items, scan_op, initial_value, equality_op);\n\n // Run Test\n Test(h_keys_in, h_values_in, h_reference, num_items, scan_op, initial_value, equality_op);\n\n // Cleanup\n if (h_reference) delete[] h_reference;\n}\n\n\n/**\n * Test different gen modes\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT,\n typename EqualityOpT>\nvoid Test(\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n TestPointer( num_items, UNIFORM, scan_op, initial_value, equality_op);\n TestPointer( num_items, RANDOM, scan_op, initial_value, equality_op);\n TestIterator( num_items, scan_op, initial_value, equality_op);\n}\n\n\n/**\n * Test different dispatch\n */\ntemplate <\n typename KeyT,\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT,\n typename EqualityOpT>\nvoid Test(\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n#if TEST_CDP == 0\n Test(num_items, scan_op, initial_value, equality_op);\n#elif TEST_CDP == 1\n Test(num_items, scan_op, initial_value, equality_op);\n#endif // TEST_CDP\n}\n\n\n/**\n * Test different operators\n */\ntemplate \nvoid TestOp(\n int num_items,\n OutputT identity,\n OutputT initial_value,\n EqualityOpT equality_op)\n{\n // Exclusive (use identity as initial value because it will dispatch to *Sum variants that don't take initial values)\n Test(num_items, cub::Sum(), identity, equality_op);\n Test(num_items, cub::Max(), identity, equality_op);\n\n // Exclusive (non-specialized, so we can test initial-value)\n Test(num_items, WrapperFunctor(cub::Sum()), initial_value, equality_op);\n Test(num_items, WrapperFunctor(cub::Max()), initial_value, equality_op);\n\n // Inclusive (no initial value)\n Test(num_items, cub::Sum(), NullType(), equality_op);\n Test(num_items, cub::Max(), NullType(), equality_op);\n}\n\n/**\n * Test different key type and equality operator\n */\ntemplate \nvoid TestKeyTAndEqualityOp(\n int num_items,\n OutputT identity,\n OutputT initial_value)\n{\n TestOp(num_items, identity, initial_value, Equality());\n TestOp( num_items, identity, initial_value, Mod2Equality());\n}\n\n/**\n * Test different input sizes\n */\ntemplate <\n typename InputT,\n typename OutputT>\nvoid TestSize(\n int num_items,\n OutputT identity,\n OutputT initial_value)\n{\n if (num_items < 0)\n {\n TestKeyTAndEqualityOp(0, identity, initial_value);\n TestKeyTAndEqualityOp(1, identity, initial_value);\n TestKeyTAndEqualityOp(100, identity, initial_value);\n TestKeyTAndEqualityOp(10000, identity, initial_value);\n TestKeyTAndEqualityOp(1000000, identity, initial_value);\n }\n else\n {\n TestKeyTAndEqualityOp(num_items, identity, initial_value);\n }\n}\n\n\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n int num_items = -1;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n g_device_giga_bandwidth = args.device_giga_bandwidth;\n printf(\"\\n\");\n\n // %PARAM% TEST_CDP cdp 0:1\n // %PARAM% TEST_VALUE_TYPES types 0:1:2:3:4:5\n\n#if TEST_VALUE_TYPES == 0\n\n // Test different input+output data types\n TestSize(num_items, (int)0, (int)99);\n\n // Test same input+output data types\n TestSize(num_items, (unsigned char)0, (unsigned char)99);\n TestSize(num_items, (char)0, (char)99);\n\n#elif TEST_VALUE_TYPES == 1\n\n TestSize(num_items, (unsigned short)0, (unsigned short)99);\n TestSize(num_items, (unsigned int)0, (unsigned int)99);\n TestSize(num_items,\n (unsigned long long)0,\n (unsigned long long)99);\n#elif TEST_VALUE_TYPES == 2\n\n TestSize(num_items, make_uchar2(0, 0), make_uchar2(17, 21));\n TestSize(num_items, make_char2(0, 0), make_char2(17, 21));\n TestSize(num_items, make_ushort2(0, 0), make_ushort2(17, 21));\n\n#elif TEST_VALUE_TYPES == 3\n\n TestSize(num_items, make_uint2(0, 0), make_uint2(17, 21));\n TestSize(num_items,\n make_ulonglong2(0, 0),\n make_ulonglong2(17, 21));\n TestSize(num_items,\n make_uchar4(0, 0, 0, 0),\n make_uchar4(17, 21, 32, 85));\n\n#elif TEST_VALUE_TYPES == 4\n\n TestSize(num_items,\n make_char4(0, 0, 0, 0),\n make_char4(17, 21, 32, 85));\n\n TestSize(num_items,\n make_ushort4(0, 0, 0, 0),\n make_ushort4(17, 21, 32, 85));\n TestSize(num_items,\n make_uint4(0, 0, 0, 0),\n make_uint4(17, 21, 32, 85));\n\n#elif TEST_VALUE_TYPES == 5\n\n TestSize(num_items,\n make_ulonglong4(0, 0, 0, 0),\n make_ulonglong4(17, 21, 32, 85));\n\n TestSize(num_items,\n TestFoo::MakeTestFoo(0, 0, 0, 0),\n TestFoo::MakeTestFoo(std::numeric_limits::max(),\n std::numeric_limits::max(),\n std::numeric_limits::max(),\n std::numeric_limits::max()));\n\n TestSize(num_items, \n TestBar(0, 0), \n TestBar(std::numeric_limits::max(), \n std::numeric_limits::max()));\n\n#endif\n\n return 0;\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2021 NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceScan utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \"test_util.h\"\n\n#include \n#include \n#include \n\nusing namespace hipcub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\ndouble g_device_giga_bandwidth;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n\nenum AliasMode\n{\n AliasNone, // output is allocated\n AliasKeys, // output is an alias of input keys\n AliasValues // output is an alias of input values\n};\n\n\n/**\n * \\brief WrapperFunctor (for precluding test-specialized dispatch to *Sum variants)\n */\ntemplate\nstruct WrapperFunctor\n{\n OpT op;\n\n WrapperFunctor(OpT op) : op(op) {}\n\n template \n __host__ __device__ __forceinline__ auto operator()(const T &a, const U &b) const\n -> decltype(op(a, b))\n {\n return static_cast(op(a, b));\n }\n};\n\n/**\n * \\brief DivideByFiveFunctor (used by TestIterator)\n */\ntemplate\nstruct DivideByFiveFunctor\n{\n template \n __host__ __device__ __forceinline__ OutputT operator()(const T &a) const\n {\n return static_cast(a / 5);\n }\n};\n\n/**\n * \\brief Mod2Equality (used for non-bool keys to make keys more likely to equal each other)\n */\nstruct Mod2Equality\n{\n template \n __host__ __device__ __forceinline__ T operator()(const T &a, const T &b) const\n {\n return (a % 2) == (b % 2);\n }\n};\n\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB DeviceScan entrypoints\n//---------------------------------------------------------------------\n\n/**\n * Dispatch to exclusive scan entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n IsPrimitiveT /*is_primitive*/,\n int timing_timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n KeysInputIteratorT d_keys_in,\n ValuesInputIteratorT d_values_in,\n ValuesOutputIteratorT d_values_out,\n ScanOpT scan_op,\n InitialValueT initial_value,\n OffsetT num_items,\n EqualityOpT equality_op)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::ExclusiveScanByKey(d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_values_out, scan_op, initial_value, num_items, equality_op);\n }\n return error;\n}\n\n\n/**\n * Dispatch to exclusive sum entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n Int2Type /*is_primitive*/,\n int timing_timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n KeysInputIteratorT d_keys_in,\n ValuesInputIteratorT d_values_in,\n ValuesOutputIteratorT d_values_out,\n Sum /*scan_op*/,\n InitialValueT /*initial_value*/,\n OffsetT num_items,\n EqualityOpT equality_op)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::ExclusiveSumByKey(d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_values_out, num_items, equality_op);\n }\n return error;\n}\n\n\n/**\n * Dispatch to inclusive scan entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n IsPrimitiveT /*is_primitive*/,\n int timing_timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n KeysInputIteratorT d_keys_in,\n ValuesInputIteratorT d_values_in,\n ValuesOutputIteratorT d_values_out,\n ScanOpT scan_op,\n NullType /*initial_value*/,\n OffsetT num_items,\n EqualityOpT equality_op)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::InclusiveScanByKey(d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_values_out, scan_op, num_items, equality_op);\n }\n return error;\n}\n\n/**\n * Dispatch to inclusive sum entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n Int2Type /*is_primitive*/,\n int timing_timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n KeysInputIteratorT d_keys_in,\n ValuesInputIteratorT d_values_in,\n ValuesOutputIteratorT d_values_out,\n Sum /*scan_op*/,\n NullType /*initial_value*/,\n OffsetT num_items,\n EqualityOpT equality_op)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::InclusiveSumByKey(d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_values_out, num_items, equality_op);\n }\n return error;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceScan\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type cub_backend,\n IsPrimitiveT is_primitive,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n KeysInputIteratorT d_keys_in,\n ValuesInputIteratorT d_values_in,\n ValuesOutputIteratorT d_values_out,\n ScanOpT scan_op,\n InitialValueT initial_value,\n OffsetT num_items,\n EqualityOpT equality_op)\n{\n *d_cdp_error = Dispatch(cub_backend,\n is_primitive,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_values_in,\n d_values_out,\n scan_op,\n initial_value,\n num_items,\n equality_op);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Dispatch to CDP kernel\n */\ntemplate \nhipError_t Dispatch(Int2Type /*dispatch_to*/,\n IsPrimitiveT is_primitive,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n KeysInputIteratorT d_keys_in,\n ValuesInputIteratorT d_values_in,\n ValuesOutputIteratorT d_values_out,\n ScanOpT scan_op,\n InitialValueT initial_value,\n OffsetT num_items,\n EqualityOpT equality_op)\n{\n // Invoke kernel to invoke device-side dispatch\n hipError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n Int2Type{},\n is_primitive,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_values_in,\n d_values_out,\n scan_op,\n initial_value,\n num_items,\n equality_op);\n CubDebugExit(retval);\n\n // Copy out temp_storage_bytes\n CubDebugExit(hipMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n hipMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(hipMemcpy(&retval,\n d_cdp_error,\n sizeof(hipError_t) * 1,\n hipMemcpyDeviceToHost));\n return retval;\n}\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\n\n/**\n * Initialize problem\n */\ntemplate \nvoid Initialize(\n GenMode gen_mode,\n T *h_in,\n int num_items)\n{\n for (int i = 0; i < num_items; ++i)\n {\n InitValue(gen_mode, h_in[i], i);\n }\n\n if (g_verbose)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n/**\n * Solve exclusive-scan problem\n */\ntemplate <\n typename KeysInputIteratorT,\n typename ValuesInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT,\n typename EqualityOpT>\nvoid Solve(\n KeysInputIteratorT h_keys_in,\n ValuesInputIteratorT h_values_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n using ValueT = hipcub::detail::value_t;\n using AccumT = hipcub::detail::accumulator_t;\n\n if (num_items > 0)\n {\n for (int i = 0; i < num_items;) {\n AccumT val = static_cast(h_values_in[i]);\n h_reference[i] = initial_value;\n AccumT inclusive = static_cast(scan_op(initial_value, val));\n\n ++i;\n\n for (; i < num_items && equality_op(h_keys_in[i - 1], h_keys_in[i]); ++i)\n {\n val = static_cast(h_values_in[i]);\n h_reference[i] = static_cast(inclusive);\n inclusive = static_cast(scan_op(inclusive, val));\n }\n }\n }\n}\n\n\n/**\n * Solve inclusive-scan problem\n */\ntemplate <\n typename KeysInputIteratorT,\n typename ValuesInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename EqualityOpT>\nvoid Solve(\n KeysInputIteratorT h_keys_in,\n ValuesInputIteratorT h_values_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n NullType /*initial_value*/,\n EqualityOpT equality_op)\n{\n using ValueT = hipcub::detail::value_t;\n using AccumT = hipcub::detail::accumulator_t;\n\n if (num_items > 0)\n {\n for (int i = 0; i < num_items;) {\n AccumT inclusive = h_values_in[i];\n h_reference[i] = static_cast(inclusive);\n\n ++i;\n\n for (; i < num_items && equality_op(h_keys_in[i - 1], h_keys_in[i]); ++i)\n {\n AccumT val = h_values_in[i];\n inclusive = static_cast(scan_op(inclusive, val));\n h_reference[i] = static_cast(inclusive);\n }\n }\n }\n}\n\ntemplate\nstruct AllocateOutput {\n static void run(OutputT *&d_out, DeviceInputIteratorT, int num_items) {\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(OutputT) * num_items));\n }\n};\n\ntemplate\nstruct AllocateOutput {\n static void run(OutputT *&d_out, OutputT *d_in, int /* num_items */) {\n d_out = d_in;\n }\n};\n\n/**\n * Test DeviceScan for a given problem input\n */\ntemplate <\n Backend BACKEND,\n typename KeysInputIteratorT,\n typename ValuesInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT,\n typename EqualityOpT,\n AliasMode Mode=AliasNone>\nvoid Test(\n KeysInputIteratorT d_keys_in,\n ValuesInputIteratorT d_values_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n using KeyT = hipcub::detail::value_t;\n using InputT = hipcub::detail::value_t;\n\n // Allocate device output array\n OutputT *d_values_out = NULL;\n\n if (Mode == AliasKeys)\n {\n AllocateOutput::run(\n d_values_out,\n d_keys_in,\n num_items);\n }\n else\n {\n AllocateOutput::run(\n d_values_out,\n d_values_in,\n num_items);\n }\n\n // Allocate CDP device arrays\n size_t *d_temp_storage_bytes = NULL;\n hipError_t *d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(hipError_t) * 1));\n\n // Allocate temporary storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(\n Int2Type(),\n Int2Type::PRIMITIVE>(),\n 1,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_values_in,\n d_values_out,\n scan_op,\n initial_value,\n num_items,\n equality_op));\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n // Clear device output array\n if (Mode == AliasNone)\n {\n CubDebugExit(hipMemset(d_values_out, 0, sizeof(OutputT) * num_items));\n }\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(\n Int2Type(),\n Int2Type::PRIMITIVE>(),\n 1,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_values_in,\n d_values_out,\n scan_op,\n initial_value,\n num_items,\n equality_op));\n\n // Check for correctness (and display results, if specified)\n const int compare = CompareDeviceResults(h_reference,\n d_values_out,\n num_items,\n true,\n g_verbose);\n\n printf(\"\\t%s\", compare ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Display performance\n if (g_timing_iterations > 0)\n {\n // Performance\n GpuTimer gpu_timer;\n gpu_timer.Start();\n CubDebugExit(Dispatch(Int2Type(),\n Int2Type::PRIMITIVE>(),\n g_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_values_in,\n d_values_out,\n scan_op,\n initial_value,\n num_items,\n equality_op));\n\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n float giga_bandwidth = giga_rate * (sizeof(InputT) + sizeof(OutputT));\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s, %.1f%% peak\",\n avg_millis, giga_rate, giga_bandwidth, giga_bandwidth / g_device_giga_bandwidth * 100.0);\n }\n\n printf(\"\\n\\n\");\n\n // Cleanup\n if (Mode == AliasNone)\n {\n if (d_values_out) \n {\n CubDebugExit(g_allocator.DeviceFree(d_values_out));\n }\n }\n\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare);\n}\n\ntemplate \nvoid TestInplaceValues(KeysInputIteratorT d_keys_in,\n OutputT *d_values_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n Test(d_keys_in,\n d_values_in,\n h_reference,\n num_items,\n scan_op,\n initial_value,\n equality_op);\n}\n\ntemplate \nvoid TestInplaceValues(KeysInputIteratorT,\n ValuesInputIteratorT,\n OutputT *,\n int,\n ScanOpT,\n InitialValueT,\n EqualityOpT)\n{}\n\ntemplate \nvoid TestInplaceKeys(T *d_keys_in,\n ValuesInputIteratorT d_values_in,\n T *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n Test(d_keys_in,\n d_values_in,\n h_reference,\n num_items,\n scan_op,\n initial_value,\n equality_op);\n}\n\ntemplate \nvoid TestInplaceKeys(KeysInputIteratorT,\n ValuesInputIteratorT,\n OutputT *,\n int,\n ScanOpT,\n InitialValueT,\n EqualityOpT)\n{}\n\n/**\n * Test DeviceScan on pointer type\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT,\n typename EqualityOpT>\nvoid TestPointer(\n int num_items,\n GenMode gen_mode,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n printf(\"\\nPointer %s %s hipcub::DeviceScan::%s %d items, %s->%s (%d->%d bytes) , gen-mode %s\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n (std::is_same::value) ? \"Inclusive\" : \"Exclusive\",\n (std::is_same::value) ? \"Sum\" : \"Scan\",\n num_items,\n typeid(InputT).name(), typeid(OutputT).name(), (int) sizeof(InputT), (int) sizeof(OutputT),\n (gen_mode == RANDOM) ? \"RANDOM\" : (gen_mode == INTEGER_SEED) ? \"SEQUENTIAL\" : \"HOMOGENOUS\");\n fflush(stdout);\n\n // Allocate host arrays\n KeyT* h_keys_in = new KeyT[num_items];\n InputT* h_values_in = new InputT[num_items];\n OutputT* h_reference = new OutputT[num_items];\n\n // Initialize problem and solution\n Initialize(gen_mode, h_keys_in, num_items);\n Initialize(gen_mode, h_values_in, num_items);\n\n // If the output type is primitive and the operator is cub::Sum, the test\n // dispatcher throws away scan_op and initial_value for exclusive scan.\n // Without an initial_value arg, the accumulator switches to the input value\n // type.\n // Do the same thing here:\n if (Traits::PRIMITIVE &&\n std::is_same::value &&\n !std::is_same::value)\n {\n Solve(h_keys_in, h_values_in, h_reference, num_items, hipcub::Sum{}, InputT{}, equality_op);\n }\n else\n {\n Solve(h_keys_in, h_values_in, h_reference, num_items, scan_op, initial_value, equality_op);\n }\n\n // Allocate problem device arrays\n KeyT *d_keys_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys_in, sizeof(KeyT) * num_items));\n InputT *d_values_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_values_in, sizeof(InputT) * num_items));\n\n // Initialize device input\n CubDebugExit(hipMemcpy(d_keys_in, h_keys_in, sizeof(KeyT) * num_items, hipMemcpyHostToDevice));\n CubDebugExit(hipMemcpy(d_values_in, h_values_in, sizeof(InputT) * num_items, hipMemcpyHostToDevice));\n\n // Run Test\n Test(d_keys_in,\n d_values_in,\n h_reference,\n num_items,\n scan_op,\n initial_value,\n equality_op);\n\n // Test in/out values aliasing\n TestInplaceValues(d_keys_in,\n d_values_in, \n h_reference,\n num_items,\n scan_op,\n initial_value,\n equality_op);\n\n CubDebugExit(hipMemcpy(d_values_in, h_values_in, sizeof(InputT) * num_items, hipMemcpyHostToDevice));\n\n // Test keys/values aliasing (should go last, changes keys)\n TestInplaceKeys(d_keys_in,\n d_values_in,\n h_reference,\n num_items,\n scan_op,\n initial_value,\n equality_op);\n\n // Cleanup\n if (h_keys_in) delete[] h_keys_in;\n if (h_values_in) delete[] h_values_in;\n if (h_reference) delete[] h_reference;\n if (d_keys_in) CubDebugExit(g_allocator.DeviceFree(d_keys_in));\n if (d_values_in) CubDebugExit(g_allocator.DeviceFree(d_values_in));\n}\n\n\n/**\n * Test DeviceScan on iterator type\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT,\n typename EqualityOpT>\nvoid TestIterator(\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n printf(\"\\nIterator %s %s hipcub::DeviceScan::%s %d items, %s->%s (%d->%d bytes)\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n (std::is_same::value) ? \"Inclusive\" : \"Exclusive\",\n (std::is_same::value) ? \"Sum\" : \"Scan\",\n num_items,\n typeid(InputT).name(), typeid(OutputT).name(), (int) sizeof(InputT), (int) sizeof(OutputT));\n fflush(stdout);\n\n // Use a counting iterator followed by div as the keys\n using CountingIterT = CountingInputIterator;\n CountingIterT h_keys_in_helper(0);\n TransformInputIterator, CountingIterT> h_keys_in(h_keys_in_helper, DivideByFiveFunctor());\n\n // Use a constant iterator as the input\n InputT val = InputT();\n ConstantInputIterator h_values_in(val);\n\n // Allocate host arrays\n OutputT* h_reference = new OutputT[num_items];\n\n // Initialize problem and solution\n Solve(h_keys_in, h_values_in, h_reference, num_items, scan_op, initial_value, equality_op);\n\n // Run Test\n Test(h_keys_in, h_values_in, h_reference, num_items, scan_op, initial_value, equality_op);\n\n // Cleanup\n if (h_reference) delete[] h_reference;\n}\n\n\n/**\n * Test different gen modes\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT,\n typename EqualityOpT>\nvoid Test(\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n TestPointer( num_items, UNIFORM, scan_op, initial_value, equality_op);\n TestPointer( num_items, RANDOM, scan_op, initial_value, equality_op);\n TestIterator( num_items, scan_op, initial_value, equality_op);\n}\n\n\n/**\n * Test different dispatch\n */\ntemplate <\n typename KeyT,\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT,\n typename EqualityOpT>\nvoid Test(\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value,\n EqualityOpT equality_op)\n{\n#if TEST_CDP == 0\n Test(num_items, scan_op, initial_value, equality_op);\n#elif TEST_CDP == 1\n Test(num_items, scan_op, initial_value, equality_op);\n#endif // TEST_CDP\n}\n\n\n/**\n * Test different operators\n */\ntemplate \nvoid TestOp(\n int num_items,\n OutputT identity,\n OutputT initial_value,\n EqualityOpT equality_op)\n{\n // Exclusive (use identity as initial value because it will dispatch to *Sum variants that don't take initial values)\n Test(num_items, hipcub::Sum(), identity, equality_op);\n Test(num_items, hipcub::Max(), identity, equality_op);\n\n // Exclusive (non-specialized, so we can test initial-value)\n Test(num_items, WrapperFunctor(hipcub::Sum()), initial_value, equality_op);\n Test(num_items, WrapperFunctor(hipcub::Max()), initial_value, equality_op);\n\n // Inclusive (no initial value)\n Test(num_items, hipcub::Sum(), NullType(), equality_op);\n Test(num_items, hipcub::Max(), NullType(), equality_op);\n}\n\n/**\n * Test different key type and equality operator\n */\ntemplate \nvoid TestKeyTAndEqualityOp(\n int num_items,\n OutputT identity,\n OutputT initial_value)\n{\n TestOp(num_items, identity, initial_value, Equality());\n TestOp( num_items, identity, initial_value, Mod2Equality());\n}\n\n/**\n * Test different input sizes\n */\ntemplate <\n typename InputT,\n typename OutputT>\nvoid TestSize(\n int num_items,\n OutputT identity,\n OutputT initial_value)\n{\n if (num_items < 0)\n {\n TestKeyTAndEqualityOp(0, identity, initial_value);\n TestKeyTAndEqualityOp(1, identity, initial_value);\n TestKeyTAndEqualityOp(100, identity, initial_value);\n TestKeyTAndEqualityOp(10000, identity, initial_value);\n TestKeyTAndEqualityOp(1000000, identity, initial_value);\n }\n else\n {\n TestKeyTAndEqualityOp(num_items, identity, initial_value);\n }\n}\n\n\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n int num_items = -1;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n g_device_giga_bandwidth = args.device_giga_bandwidth;\n printf(\"\\n\");\n\n // %PARAM% TEST_CDP cdp 0:1\n // %PARAM% TEST_VALUE_TYPES types 0:1:2:3:4:5\n\n#if TEST_VALUE_TYPES == 0\n\n // Test different input+output data types\n TestSize(num_items, (int)0, (int)99);\n\n // Test same input+output data types\n TestSize(num_items, (unsigned char)0, (unsigned char)99);\n TestSize(num_items, (char)0, (char)99);\n\n#elif TEST_VALUE_TYPES == 1\n\n TestSize(num_items, (unsigned short)0, (unsigned short)99);\n TestSize(num_items, (unsigned int)0, (unsigned int)99);\n TestSize(num_items,\n (unsigned long long)0,\n (unsigned long long)99);\n#elif TEST_VALUE_TYPES == 2\n\n TestSize(num_items, make_uchar2(0, 0), make_uchar2(17, 21));\n TestSize(num_items, make_char2(0, 0), make_char2(17, 21));\n TestSize(num_items, make_ushort2(0, 0), make_ushort2(17, 21));\n\n#elif TEST_VALUE_TYPES == 3\n\n TestSize(num_items, make_uint2(0, 0), make_uint2(17, 21));\n TestSize(num_items,\n make_ulonglong2(0, 0),\n make_ulonglong2(17, 21));\n TestSize(num_items,\n make_uchar4(0, 0, 0, 0),\n make_uchar4(17, 21, 32, 85));\n\n#elif TEST_VALUE_TYPES == 4\n\n TestSize(num_items,\n make_char4(0, 0, 0, 0),\n make_char4(17, 21, 32, 85));\n\n TestSize(num_items,\n make_ushort4(0, 0, 0, 0),\n make_ushort4(17, 21, 32, 85));\n TestSize(num_items,\n make_uint4(0, 0, 0, 0),\n make_uint4(17, 21, 32, 85));\n\n#elif TEST_VALUE_TYPES == 5\n\n TestSize(num_items,\n make_ulonglong4(0, 0, 0, 0),\n make_ulonglong4(17, 21, 32, 85));\n\n TestSize(num_items,\n TestFoo::MakeTestFoo(0, 0, 0, 0),\n TestFoo::MakeTestFoo(std::numeric_limits::max(),\n std::numeric_limits::max(),\n std::numeric_limits::max(),\n std::numeric_limits::max()));\n\n TestSize(num_items, \n TestBar(0, 0), \n TestBar(std::numeric_limits::max(), \n std::numeric_limits::max()));\n\n#endif\n\n return 0;\n}\n", "source": "/kaggle/working/cub/test/test_device_scan_by_key.cu"} {"input": "/******************************************************************************\n* Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n*\n* Redistribution and use in source and binary forms, with or without\n* modification, are permitted provided that the following conditions are met:\n* * Redistributions of source code must retain the above copyright\n* notice, this list of conditions and the following disclaimer.\n* * Redistributions in binary form must reproduce the above copyright\n* notice, this list of conditions and the following disclaimer in the\n* documentation and/or other materials provided with the distribution.\n* * Neither the name of the NVIDIA CORPORATION nor the\n* names of its contributors may be used to endorse or promote products\n* derived from this software without specific prior written permission.\n*\n* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n* ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n* DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n* DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n* (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n* ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n* SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n*\n******************************************************************************/\n\n/******************************************************************************\n * Test of BlockHistogram utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n#include \n\n#include \n\n#include \n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void block_histogram_kernel(T *d_samples, HistoCounter *d_histogram)\n{\n // Parameterize BlockHistogram type for our thread block\n using block_histogram_t =\n cub::BlockHistogram;\n\n // Allocate temp storage in shared memory\n __shared__ typename block_histogram_t::TempStorage temp_storage;\n\n // Per-thread tile data\n T data[ITEMS_PER_THREAD];\n cub::LoadDirectStriped(threadIdx.x, d_samples, data);\n\n // Test histo (writing directly to histogram buffer in global)\n block_histogram_t(temp_storage).Histogram(data, d_histogram);\n}\n\ntemplate \nvoid block_histogram(thrust::device_vector &d_samples,\n thrust::device_vector &d_histogram)\n{\n block_histogram_kernel\n <<<1, ThreadsInBlock>>>(thrust::raw_pointer_cast(d_samples.data()),\n thrust::raw_pointer_cast(d_histogram.data()));\n\n REQUIRE(cudaSuccess == cudaPeekAtLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n}\n\n// %PARAM% TEST_BINS bins 32:256:1024\n\nusing types = c2h::type_list;\nusing threads_in_block = c2h::enum_type_list;\nusing items_per_thread = c2h::enum_type_list;\nusing bins = c2h::enum_type_list;\nusing algorithms = c2h::enum_type_list;\n\ntemplate \nstruct params_t\n{\n using sample_t = typename c2h::get<0, TestType>;\n\n static constexpr int items_per_thread = c2h::get<1, TestType>::value;\n static constexpr int threads_in_block = c2h::get<2, TestType>::value;\n static constexpr int bins = c2h::get<3, TestType>::value;\n static constexpr int num_samples = threads_in_block * items_per_thread;\n static constexpr cub::BlockHistogramAlgorithm algorithm =\n c2h::get<4, TestType>::value;\n};\n\nCUB_TEST(\"Block histogram can be computed with uniform input\",\n \"[histogram][block]\",\n types,\n items_per_thread,\n threads_in_block,\n bins,\n algorithms)\n{\n using params = params_t;\n using sample_t = typename params::sample_t;\n\n const sample_t uniform_value =\n static_cast(GENERATE_COPY(take(10, random(0, params::bins - 1))));\n\n thrust::host_vector h_samples(params::num_samples, uniform_value);\n thrust::host_vector h_reference(params::bins);\n h_reference[static_cast(uniform_value)] = params::num_samples;\n\n // Allocate problem device arrays\n thrust::device_vector d_samples = h_samples;\n thrust::device_vector d_histogram(params::bins);\n\n // Run kernel\n block_histogram(d_samples, d_histogram);\n\n REQUIRE(h_reference == d_histogram);\n}\n\ntemplate \nthrust::host_vector\ncompute_host_reference(int bins, const thrust::host_vector &h_samples)\n{\n thrust::host_vector h_reference(bins);\n for (const SampleT &sample : h_samples)\n {\n h_reference[sample]++;\n }\n\n return h_reference;\n}\n\nCUB_TEST(\"Block histogram can be computed with modulo input\",\n \"[histogram][block]\",\n types,\n items_per_thread,\n threads_in_block,\n bins,\n algorithms)\n{\n using params = params_t;\n using sample_t = typename params::sample_t;\n\n // Allocate problem device arrays\n thrust::device_vector d_histogram(params::bins);\n thrust::device_vector d_samples(params::num_samples);\n\n c2h::gen(c2h::modulo_t{params::bins}, d_samples);\n\n thrust::host_vector h_samples = d_samples;\n auto h_reference = compute_host_reference(params::bins, h_samples);\n\n // Run kernel\n block_histogram(d_samples, d_histogram);\n\n REQUIRE(h_reference == d_histogram);\n}\n\nCUB_TEST(\"Block histogram can be computed with random input\",\n \"[histogram][block]\",\n types,\n items_per_thread,\n threads_in_block,\n bins,\n algorithms)\n{\n using params = params_t;\n using sample_t = typename params::sample_t;\n\n // Allocate problem device arrays\n thrust::device_vector d_histogram(params::bins);\n thrust::device_vector d_samples(params::num_samples);\n\n const sample_t min_bin = static_cast(0);\n const sample_t max_bin = \n static_cast(\n std::min(static_cast(std::numeric_limits::max()),\n static_cast(params::bins - 1)));\n\n c2h::gen(CUB_SEED(10), d_samples, min_bin, max_bin);\n\n thrust::host_vector h_samples = d_samples;\n auto h_reference = compute_host_reference(params::bins, h_samples);\n\n // Run kernel\n block_histogram(d_samples, d_histogram);\n\n REQUIRE(h_reference == d_histogram);\n}\n\n", "output": "/******************************************************************************\n* Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n*\n* Redistribution and use in source and binary forms, with or without\n* modification, are permitted provided that the following conditions are met:\n* * Redistributions of source code must retain the above copyright\n* notice, this list of conditions and the following disclaimer.\n* * Redistributions in binary form must reproduce the above copyright\n* notice, this list of conditions and the following disclaimer in the\n* documentation and/or other materials provided with the distribution.\n* * Neither the name of the NVIDIA CORPORATION nor the\n* names of its contributors may be used to endorse or promote products\n* derived from this software without specific prior written permission.\n*\n* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n* ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n* DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n* DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n* (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n* ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n* SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n*\n******************************************************************************/\n\n/******************************************************************************\n * Test of BlockHistogram utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n\n#include \n\n#include \n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void block_histogram_kernel(T *d_samples, HistoCounter *d_histogram)\n{\n // Parameterize BlockHistogram type for our thread block\n using block_histogram_t =\n hipcub::BlockHistogram;\n\n // Allocate temp storage in shared memory\n __shared__ typename block_histogram_t::TempStorage temp_storage;\n\n // Per-thread tile data\n T data[ITEMS_PER_THREAD];\n hipcub::LoadDirectStriped(threadIdx.x, d_samples, data);\n\n // Test histo (writing directly to histogram buffer in global)\n block_histogram_t(temp_storage).Histogram(data, d_histogram);\n}\n\ntemplate \nvoid block_histogram(thrust::device_vector &d_samples,\n thrust::device_vector &d_histogram)\n{\n block_histogram_kernel\n <<<1, ThreadsInBlock>>>(thrust::raw_pointer_cast(d_samples.data()),\n thrust::raw_pointer_cast(d_histogram.data()));\n\n REQUIRE(hipSuccess == hipPeekAtLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n}\n\n// %PARAM% TEST_BINS bins 32:256:1024\n\nusing types = c2h::type_list;\nusing threads_in_block = c2h::enum_type_list;\nusing items_per_thread = c2h::enum_type_list;\nusing bins = c2h::enum_type_list;\nusing algorithms = c2h::enum_type_list;\n\ntemplate \nstruct params_t\n{\n using sample_t = typename c2h::get<0, TestType>;\n\n static constexpr int items_per_thread = c2h::get<1, TestType>::value;\n static constexpr int threads_in_block = c2h::get<2, TestType>::value;\n static constexpr int bins = c2h::get<3, TestType>::value;\n static constexpr int num_samples = threads_in_block * items_per_thread;\n static constexpr hipcub::BlockHistogramAlgorithm algorithm =\n c2h::get<4, TestType>::value;\n};\n\nCUB_TEST(\"Block histogram can be computed with uniform input\",\n \"[histogram][block]\",\n types,\n items_per_thread,\n threads_in_block,\n bins,\n algorithms)\n{\n using params = params_t;\n using sample_t = typename params::sample_t;\n\n const sample_t uniform_value =\n static_cast(GENERATE_COPY(take(10, random(0, params::bins - 1))));\n\n thrust::host_vector h_samples(params::num_samples, uniform_value);\n thrust::host_vector h_reference(params::bins);\n h_reference[static_cast(uniform_value)] = params::num_samples;\n\n // Allocate problem device arrays\n thrust::device_vector d_samples = h_samples;\n thrust::device_vector d_histogram(params::bins);\n\n // Run kernel\n block_histogram(d_samples, d_histogram);\n\n REQUIRE(h_reference == d_histogram);\n}\n\ntemplate \nthrust::host_vector\ncompute_host_reference(int bins, const thrust::host_vector &h_samples)\n{\n thrust::host_vector h_reference(bins);\n for (const SampleT &sample : h_samples)\n {\n h_reference[sample]++;\n }\n\n return h_reference;\n}\n\nCUB_TEST(\"Block histogram can be computed with modulo input\",\n \"[histogram][block]\",\n types,\n items_per_thread,\n threads_in_block,\n bins,\n algorithms)\n{\n using params = params_t;\n using sample_t = typename params::sample_t;\n\n // Allocate problem device arrays\n thrust::device_vector d_histogram(params::bins);\n thrust::device_vector d_samples(params::num_samples);\n\n c2h::gen(c2h::modulo_t{params::bins}, d_samples);\n\n thrust::host_vector h_samples = d_samples;\n auto h_reference = compute_host_reference(params::bins, h_samples);\n\n // Run kernel\n block_histogram(d_samples, d_histogram);\n\n REQUIRE(h_reference == d_histogram);\n}\n\nCUB_TEST(\"Block histogram can be computed with random input\",\n \"[histogram][block]\",\n types,\n items_per_thread,\n threads_in_block,\n bins,\n algorithms)\n{\n using params = params_t;\n using sample_t = typename params::sample_t;\n\n // Allocate problem device arrays\n thrust::device_vector d_histogram(params::bins);\n thrust::device_vector d_samples(params::num_samples);\n\n const sample_t min_bin = static_cast(0);\n const sample_t max_bin = \n static_cast(\n std::min(static_cast(std::numeric_limits::max()),\n static_cast(params::bins - 1)));\n\n c2h::gen(CUB_SEED(10), d_samples, min_bin, max_bin);\n\n thrust::host_vector h_samples = d_samples;\n auto h_reference = compute_host_reference(params::bins, h_samples);\n\n // Run kernel\n block_histogram(d_samples, d_histogram);\n\n REQUIRE(h_reference == d_histogram);\n}\n\n", "source": "/kaggle/working/cub/test/catch2_test_block_histogram.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceRunLengthEncode utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \n#include \n\n#include \"test_util.h\"\n\nusing namespace cub;\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n// Operation types\nenum RleMethod\n{\n RLE, // Run length encode\n NON_TRIVIAL,\n};\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB entrypoints\n//---------------------------------------------------------------------\n\n/**\n * Dispatch to run-length encode entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__ cudaError_t Dispatch(Int2Type /*method*/,\n Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n cudaError_t * /*d_cdp_error*/,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n UniqueOutputIteratorT d_unique_out,\n OffsetsOutputIteratorT /*d_offsets_out*/,\n LengthsOutputIteratorT d_lengths_out,\n NumRunsIterator d_num_runs,\n cub::Equality /*equality_op*/,\n OffsetT num_items)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceRunLengthEncode::Encode(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_unique_out,\n d_lengths_out,\n d_num_runs,\n num_items);\n }\n return error;\n}\n\n/**\n * Dispatch to non-trivial runs entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__ cudaError_t Dispatch(Int2Type /*method*/,\n Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n cudaError_t * /*d_cdp_error*/,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n UniqueOutputIteratorT /*d_unique_out*/,\n OffsetsOutputIteratorT d_offsets_out,\n LengthsOutputIteratorT d_lengths_out,\n NumRunsIterator d_num_runs,\n cub::Equality /*equality_op*/,\n OffsetT num_items)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceRunLengthEncode::NonTrivialRuns(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_offsets_out,\n d_lengths_out,\n d_num_runs,\n num_items);\n }\n return error;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceRunLengthEncode\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type method,\n Int2Type cub_backend,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n InputIteratorT d_in,\n UniqueOutputIteratorT d_unique_out,\n OffsetsOutputIteratorT d_offsets_out,\n LengthsOutputIteratorT d_lengths_out,\n NumRunsIterator d_num_runs,\n cub::Equality equality_op,\n OffsetT num_items)\n{\n *d_cdp_error = Dispatch(method,\n cub_backend,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_unique_out,\n d_offsets_out,\n d_lengths_out,\n d_num_runs,\n equality_op,\n num_items);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Dispatch to CDP kernel\n */\ntemplate \n__forceinline__ cudaError_t Dispatch(Int2Type method,\n Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n UniqueOutputIteratorT d_unique_out,\n OffsetsOutputIteratorT d_offsets_out,\n LengthsOutputIteratorT d_lengths_out,\n NumRunsIterator d_num_runs,\n EqualityOp equality_op,\n OffsetT num_items)\n{\n // Invoke kernel to invoke device-side dispatch\n cudaError_t retval = thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n method,\n Int2Type{},\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_unique_out,\n d_offsets_out,\n d_lengths_out,\n d_num_runs,\n equality_op,\n num_items);\n CubDebugExit(retval);\n\n // Copy out temp_storage_bytes\n CubDebugExit(cudaMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n cudaMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(cudaMemcpy(&retval, d_cdp_error, sizeof(cudaError_t) * 1, cudaMemcpyDeviceToHost));\n return retval;\n}\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\n/**\n * Initialize problem\n */\ntemplate \nvoid Initialize(int entropy_reduction, T *h_in, int num_items, int max_segment)\n{\n unsigned int max_int = (unsigned int)-1;\n\n int key = 0;\n int i = 0;\n while (i < num_items)\n {\n // Select number of repeating occurrences for the current run\n int repeat;\n if (max_segment < 0)\n {\n repeat = num_items;\n }\n else if (max_segment < 2)\n {\n repeat = 1;\n }\n else\n {\n RandomBits(repeat, entropy_reduction);\n repeat = (int)((double(repeat) * double(max_segment)) / double(max_int));\n repeat = CUB_MAX(1, repeat);\n }\n\n int j = i;\n while (j < CUB_MIN(i + repeat, num_items))\n {\n InitValue(INTEGER_SEED, h_in[j], key);\n j++;\n }\n\n i = j;\n key++;\n }\n\n if (g_verbose)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n/**\n * Solve problem. Returns total number of segments identified\n */\ntemplate \nint Solve(InputIteratorT h_in,\n T *h_unique_reference,\n OffsetT *h_offsets_reference,\n LengthT *h_lengths_reference,\n EqualityOp equality_op,\n int num_items)\n{\n if (num_items == 0)\n return 0;\n\n // First item\n T previous = h_in[0];\n LengthT length = 1;\n int num_runs = 0;\n int run_begin = 0;\n\n // Subsequent items\n for (int i = 1; i < num_items; ++i)\n {\n if (!equality_op(previous, h_in[i]))\n {\n if ((RLE_METHOD != NON_TRIVIAL) || (length > 1))\n {\n h_unique_reference[num_runs] = previous;\n h_offsets_reference[num_runs] = run_begin;\n h_lengths_reference[num_runs] = length;\n num_runs++;\n }\n length = 1;\n run_begin = i;\n }\n else\n {\n length++;\n }\n previous = h_in[i];\n }\n\n if ((RLE_METHOD != NON_TRIVIAL) || (length > 1))\n {\n h_unique_reference[num_runs] = previous;\n h_offsets_reference[num_runs] = run_begin;\n h_lengths_reference[num_runs] = length;\n num_runs++;\n }\n\n return num_runs;\n}\n\n/**\n * Test DeviceRunLengthEncode for a given problem input\n */\ntemplate \nvoid Test(DeviceInputIteratorT d_in,\n T *h_unique_reference,\n OffsetT *h_offsets_reference,\n LengthT *h_lengths_reference,\n EqualityOp equality_op,\n int num_runs,\n int num_items)\n{\n // Allocate device output arrays and number of segments\n T *d_unique_out = NULL;\n LengthT *d_offsets_out = NULL;\n OffsetT *d_lengths_out = NULL;\n int *d_num_runs = NULL;\n\n if (RLE_METHOD == RLE)\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_unique_out, sizeof(T) * num_items));\n if (RLE_METHOD == NON_TRIVIAL)\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_offsets_out, sizeof(OffsetT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_lengths_out, sizeof(LengthT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_num_runs, sizeof(int)));\n\n // Allocate CDP device arrays\n size_t *d_temp_storage_bytes = NULL;\n cudaError_t *d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_cdp_error, sizeof(cudaError_t) * 1));\n\n // Allocate temporary storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(Int2Type(),\n Int2Type(),\n 1,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_unique_out,\n d_offsets_out,\n d_lengths_out,\n d_num_runs,\n equality_op,\n num_items));\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n // Clear device output arrays\n if (RLE_METHOD == RLE)\n CubDebugExit(cudaMemset(d_unique_out, 0, sizeof(T) * num_items));\n if (RLE_METHOD == NON_TRIVIAL)\n CubDebugExit(cudaMemset(d_offsets_out, 0, sizeof(OffsetT) * num_items));\n CubDebugExit(cudaMemset(d_lengths_out, 0, sizeof(LengthT) * num_items));\n CubDebugExit(cudaMemset(d_num_runs, 0, sizeof(int)));\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(Int2Type(),\n Int2Type(),\n 1,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_unique_out,\n d_offsets_out,\n d_lengths_out,\n d_num_runs,\n equality_op,\n num_items));\n\n // Check for correctness (and display results, if specified)\n int compare0 = 0;\n int compare1 = 0;\n int compare2 = 0;\n int compare3 = 0;\n\n if (RLE_METHOD == RLE)\n {\n compare0 = CompareDeviceResults(h_unique_reference, d_unique_out, num_runs, true, g_verbose);\n printf(\"\\t Keys %s\\n\", compare0 ? \"FAIL\" : \"PASS\");\n }\n\n if (RLE_METHOD != RLE)\n {\n compare1 = CompareDeviceResults(h_offsets_reference, d_offsets_out, num_runs, true, g_verbose);\n printf(\"\\t Offsets %s\\n\", compare1 ? \"FAIL\" : \"PASS\");\n }\n\n compare2 = CompareDeviceResults(h_lengths_reference, d_lengths_out, num_runs, true, g_verbose);\n printf(\"\\t Lengths %s\\n\", compare2 ? \"FAIL\" : \"PASS\");\n\n compare3 = CompareDeviceResults(&num_runs, d_num_runs, 1, true, g_verbose);\n printf(\"\\t Count %s\\n\", compare3 ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Performance\n GpuTimer gpu_timer;\n gpu_timer.Start();\n CubDebugExit(Dispatch(Int2Type(),\n Int2Type(),\n g_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_unique_out,\n d_offsets_out,\n d_lengths_out,\n d_num_runs,\n equality_op,\n num_items));\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n\n // Display performance\n if (g_timing_iterations > 0)\n {\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n int bytes_moved = (num_items * sizeof(T)) + (num_runs * (sizeof(OffsetT) + sizeof(LengthT)));\n float giga_bandwidth = float(bytes_moved) / avg_millis / 1000.0f / 1000.0f;\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s\",\n avg_millis,\n giga_rate,\n giga_bandwidth);\n }\n printf(\"\\n\\n\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Cleanup\n if (d_unique_out)\n CubDebugExit(g_allocator.DeviceFree(d_unique_out));\n if (d_offsets_out)\n CubDebugExit(g_allocator.DeviceFree(d_offsets_out));\n if (d_lengths_out)\n CubDebugExit(g_allocator.DeviceFree(d_lengths_out));\n if (d_num_runs)\n CubDebugExit(g_allocator.DeviceFree(d_num_runs));\n if (d_temp_storage_bytes)\n CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error)\n CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage)\n CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare0 | compare1 | compare2 | compare3);\n}\n\n/**\n * Test DeviceRunLengthEncode on pointer type\n */\ntemplate \nvoid TestPointer(int num_items, int entropy_reduction, int max_segment)\n{\n // Allocate host arrays\n T *h_in = new T[num_items];\n T *h_unique_reference = new T[num_items];\n OffsetT *h_offsets_reference = new OffsetT[num_items];\n LengthT *h_lengths_reference = new LengthT[num_items];\n\n for (int i = 0; i < num_items; ++i)\n InitValue(INTEGER_SEED, h_offsets_reference[i], 1);\n\n // Initialize problem and solution\n Equality equality_op;\n Initialize(entropy_reduction, h_in, num_items, max_segment);\n\n int num_runs = Solve(h_in,\n h_unique_reference,\n h_offsets_reference,\n h_lengths_reference,\n equality_op,\n num_items);\n\n printf(\"\\nPointer %s cub::%s on %d items, %d segments (avg run length %.3f), {%s key, %s offset, \"\n \"%s length}, max_segment %d, entropy_reduction %d\\n\",\n (RLE_METHOD == RLE) ? \"DeviceRunLengthEncode::Encode\"\n : (RLE_METHOD == NON_TRIVIAL) ? \"DeviceRunLengthEncode::NonTrivialRuns\"\n : \"Other\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items,\n num_runs,\n float(num_items) / num_runs,\n typeid(T).name(),\n typeid(OffsetT).name(),\n typeid(LengthT).name(),\n max_segment,\n entropy_reduction);\n fflush(stdout);\n\n // Allocate problem device arrays\n T *d_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_in, sizeof(T) * num_items));\n\n // Initialize device input\n CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(T) * num_items, cudaMemcpyHostToDevice));\n\n // Run Test\n Test(d_in,\n h_unique_reference,\n h_offsets_reference,\n h_lengths_reference,\n equality_op,\n num_runs,\n num_items);\n\n // Cleanup\n if (h_in)\n delete[] h_in;\n if (h_unique_reference)\n delete[] h_unique_reference;\n if (h_offsets_reference)\n delete[] h_offsets_reference;\n if (h_lengths_reference)\n delete[] h_lengths_reference;\n if (d_in)\n CubDebugExit(g_allocator.DeviceFree(d_in));\n}\n\n/**\n * Test on iterator type\n */\ntemplate \nvoid TestIterator(int num_items, Int2Type /*is_primitive*/)\n{\n // Allocate host arrays\n T *h_unique_reference = new T[num_items];\n OffsetT *h_offsets_reference = new OffsetT[num_items];\n LengthT *h_lengths_reference = new LengthT[num_items];\n\n T one_val;\n InitValue(INTEGER_SEED, one_val, 1);\n ConstantInputIterator h_in(one_val);\n\n // Initialize problem and solution\n Equality equality_op;\n int num_runs = Solve(h_in,\n h_unique_reference,\n h_offsets_reference,\n h_lengths_reference,\n equality_op,\n num_items);\n\n printf(\"\\nIterator %s cub::%s on %d items, %d segments (avg run length %.3f), {%s key, %s \"\n \"offset, %s length}\\n\",\n (RLE_METHOD == RLE) ? \"DeviceRunLengthEncode::Encode\"\n : (RLE_METHOD == NON_TRIVIAL) ? \"DeviceRunLengthEncode::NonTrivialRuns\"\n : \"Other\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items,\n num_runs,\n float(num_items) / num_runs,\n typeid(T).name(),\n typeid(OffsetT).name(),\n typeid(LengthT).name());\n fflush(stdout);\n\n // Run Test\n Test(h_in,\n h_unique_reference,\n h_offsets_reference,\n h_lengths_reference,\n equality_op,\n num_runs,\n num_items);\n\n // Cleanup\n if (h_unique_reference)\n delete[] h_unique_reference;\n if (h_offsets_reference)\n delete[] h_offsets_reference;\n if (h_lengths_reference)\n delete[] h_lengths_reference;\n}\n\ntemplate \nvoid TestIterator(int /*num_items*/, Int2Type /*is_primitive*/)\n{}\n\n/**\n * Test different gen modes\n */\ntemplate \nvoid Test(int num_items)\n{\n // Test iterator (one run)\n TestIterator(num_items,\n Int2Type::PRIMITIVE>());\n\n // Evaluate different run lengths / segment sizes\n const int max_seg_limit = CUB_MIN(num_items, 1 << 16);\n const int max_seg_inc = 4;\n for (int max_segment = 1, entropy_reduction = 0; max_segment <= max_seg_limit;\n max_segment <<= max_seg_inc, entropy_reduction++)\n {\n const int max_seg = CUB_MAX(1, max_segment);\n TestPointer(num_items, entropy_reduction, max_seg);\n }\n}\n\n/**\n * Test different dispatch\n */\ntemplate \nvoid TestDispatch(int num_items)\n{\n#if TEST_CDP == 0\n Test(num_items);\n Test(num_items);\n#elif TEST_CDP == 1\n Test(num_items);\n Test(num_items);\n#endif\n}\n\n/**\n * Test different input sizes\n */\ntemplate \nvoid TestSize(int num_items)\n{\n if (num_items < 0)\n {\n TestDispatch(0);\n TestDispatch(1);\n TestDispatch(100);\n TestDispatch(10000);\n TestDispatch(1000000);\n }\n else\n {\n TestDispatch(num_items);\n }\n}\n\n/**\n * @brief Test with NaNs as input to ensure we don't create an additional (invalid) run on the first\n * item when (item[0] == item[0]) is false.\n */\ntemplate \nvoid TestNaNs()\n{\n using OffsetT = int32_t;\n using LengthT = int32_t;\n\n const auto quiet_nan = std::numeric_limits::quiet_NaN();\n\n // Allocate host data\n std::vector h_in{quiet_nan, quiet_nan, quiet_nan};\n const OffsetT num_items = static_cast(h_in.size());\n T *h_unique_reference = new T[num_items];\n OffsetT *h_offsets_reference = new OffsetT[num_items];\n LengthT *h_lengths_reference = new LengthT[num_items];\n\n // Initialize problem and solution\n Equality equality_op;\n\n // Get host-side data for verification\n int num_runs = Solve(h_in.data(),\n h_unique_reference,\n h_offsets_reference,\n h_lengths_reference,\n equality_op,\n num_items);\n\n // Allocate problem device arrays\n thrust::device_vector d_in(num_items);\n d_in = h_in;\n\n // Run Test\n Test(thrust::raw_pointer_cast(d_in.data()),\n h_unique_reference,\n h_offsets_reference,\n h_lengths_reference,\n equality_op,\n num_runs,\n num_items);\n}\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char **argv)\n{\n int num_items = -1;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\",\n argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n printf(\"\\n\");\n\n // %PARAM% TEST_CDP cdp 0:1\n\n // Run tests with leading/trailing NaNs, where (NaN == NaN) is false\n TestNaNs();\n TestNaNs();\n\n // Test different input types\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n\n return 0;\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceRunLengthEncode utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \n#include \n\n#include \"test_util.h\"\n\nusing namespace hipcub;\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n// Operation types\nenum RleMethod\n{\n RLE, // Run length encode\n NON_TRIVIAL,\n};\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB entrypoints\n//---------------------------------------------------------------------\n\n/**\n * Dispatch to run-length encode entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__ hipError_t Dispatch(Int2Type /*method*/,\n Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n hipError_t * /*d_cdp_error*/,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n UniqueOutputIteratorT d_unique_out,\n OffsetsOutputIteratorT /*d_offsets_out*/,\n LengthsOutputIteratorT d_lengths_out,\n NumRunsIterator d_num_runs,\n hipcub::Equality /*equality_op*/,\n OffsetT num_items)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceRunLengthEncode::Encode(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_unique_out,\n d_lengths_out,\n d_num_runs,\n num_items);\n }\n return error;\n}\n\n/**\n * Dispatch to non-trivial runs entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__ hipError_t Dispatch(Int2Type /*method*/,\n Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n hipError_t * /*d_cdp_error*/,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n UniqueOutputIteratorT /*d_unique_out*/,\n OffsetsOutputIteratorT d_offsets_out,\n LengthsOutputIteratorT d_lengths_out,\n NumRunsIterator d_num_runs,\n hipcub::Equality /*equality_op*/,\n OffsetT num_items)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceRunLengthEncode::NonTrivialRuns(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_offsets_out,\n d_lengths_out,\n d_num_runs,\n num_items);\n }\n return error;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceRunLengthEncode\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type method,\n Int2Type cub_backend,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n InputIteratorT d_in,\n UniqueOutputIteratorT d_unique_out,\n OffsetsOutputIteratorT d_offsets_out,\n LengthsOutputIteratorT d_lengths_out,\n NumRunsIterator d_num_runs,\n hipcub::Equality equality_op,\n OffsetT num_items)\n{\n *d_cdp_error = Dispatch(method,\n cub_backend,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_unique_out,\n d_offsets_out,\n d_lengths_out,\n d_num_runs,\n equality_op,\n num_items);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Dispatch to CDP kernel\n */\ntemplate \n__forceinline__ hipError_t Dispatch(Int2Type method,\n Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n UniqueOutputIteratorT d_unique_out,\n OffsetsOutputIteratorT d_offsets_out,\n LengthsOutputIteratorT d_lengths_out,\n NumRunsIterator d_num_runs,\n EqualityOp equality_op,\n OffsetT num_items)\n{\n // Invoke kernel to invoke device-side dispatch\n hipError_t retval = thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n method,\n Int2Type{},\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_unique_out,\n d_offsets_out,\n d_lengths_out,\n d_num_runs,\n equality_op,\n num_items);\n CubDebugExit(retval);\n\n // Copy out temp_storage_bytes\n CubDebugExit(hipMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n hipMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(hipMemcpy(&retval, d_cdp_error, sizeof(hipError_t) * 1, hipMemcpyDeviceToHost));\n return retval;\n}\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\n/**\n * Initialize problem\n */\ntemplate \nvoid Initialize(int entropy_reduction, T *h_in, int num_items, int max_segment)\n{\n unsigned int max_int = (unsigned int)-1;\n\n int key = 0;\n int i = 0;\n while (i < num_items)\n {\n // Select number of repeating occurrences for the current run\n int repeat;\n if (max_segment < 0)\n {\n repeat = num_items;\n }\n else if (max_segment < 2)\n {\n repeat = 1;\n }\n else\n {\n RandomBits(repeat, entropy_reduction);\n repeat = (int)((double(repeat) * double(max_segment)) / double(max_int));\n repeat = CUB_MAX(1, repeat);\n }\n\n int j = i;\n while (j < CUB_MIN(i + repeat, num_items))\n {\n InitValue(INTEGER_SEED, h_in[j], key);\n j++;\n }\n\n i = j;\n key++;\n }\n\n if (g_verbose)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n/**\n * Solve problem. Returns total number of segments identified\n */\ntemplate \nint Solve(InputIteratorT h_in,\n T *h_unique_reference,\n OffsetT *h_offsets_reference,\n LengthT *h_lengths_reference,\n EqualityOp equality_op,\n int num_items)\n{\n if (num_items == 0)\n return 0;\n\n // First item\n T previous = h_in[0];\n LengthT length = 1;\n int num_runs = 0;\n int run_begin = 0;\n\n // Subsequent items\n for (int i = 1; i < num_items; ++i)\n {\n if (!equality_op(previous, h_in[i]))\n {\n if ((RLE_METHOD != NON_TRIVIAL) || (length > 1))\n {\n h_unique_reference[num_runs] = previous;\n h_offsets_reference[num_runs] = run_begin;\n h_lengths_reference[num_runs] = length;\n num_runs++;\n }\n length = 1;\n run_begin = i;\n }\n else\n {\n length++;\n }\n previous = h_in[i];\n }\n\n if ((RLE_METHOD != NON_TRIVIAL) || (length > 1))\n {\n h_unique_reference[num_runs] = previous;\n h_offsets_reference[num_runs] = run_begin;\n h_lengths_reference[num_runs] = length;\n num_runs++;\n }\n\n return num_runs;\n}\n\n/**\n * Test DeviceRunLengthEncode for a given problem input\n */\ntemplate \nvoid Test(DeviceInputIteratorT d_in,\n T *h_unique_reference,\n OffsetT *h_offsets_reference,\n LengthT *h_lengths_reference,\n EqualityOp equality_op,\n int num_runs,\n int num_items)\n{\n // Allocate device output arrays and number of segments\n T *d_unique_out = NULL;\n LengthT *d_offsets_out = NULL;\n OffsetT *d_lengths_out = NULL;\n int *d_num_runs = NULL;\n\n if (RLE_METHOD == RLE)\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_unique_out, sizeof(T) * num_items));\n if (RLE_METHOD == NON_TRIVIAL)\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_offsets_out, sizeof(OffsetT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_lengths_out, sizeof(LengthT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_num_runs, sizeof(int)));\n\n // Allocate CDP device arrays\n size_t *d_temp_storage_bytes = NULL;\n hipError_t *d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_cdp_error, sizeof(hipError_t) * 1));\n\n // Allocate temporary storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(Int2Type(),\n Int2Type(),\n 1,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_unique_out,\n d_offsets_out,\n d_lengths_out,\n d_num_runs,\n equality_op,\n num_items));\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n // Clear device output arrays\n if (RLE_METHOD == RLE)\n CubDebugExit(hipMemset(d_unique_out, 0, sizeof(T) * num_items));\n if (RLE_METHOD == NON_TRIVIAL)\n CubDebugExit(hipMemset(d_offsets_out, 0, sizeof(OffsetT) * num_items));\n CubDebugExit(hipMemset(d_lengths_out, 0, sizeof(LengthT) * num_items));\n CubDebugExit(hipMemset(d_num_runs, 0, sizeof(int)));\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(Int2Type(),\n Int2Type(),\n 1,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_unique_out,\n d_offsets_out,\n d_lengths_out,\n d_num_runs,\n equality_op,\n num_items));\n\n // Check for correctness (and display results, if specified)\n int compare0 = 0;\n int compare1 = 0;\n int compare2 = 0;\n int compare3 = 0;\n\n if (RLE_METHOD == RLE)\n {\n compare0 = CompareDeviceResults(h_unique_reference, d_unique_out, num_runs, true, g_verbose);\n printf(\"\\t Keys %s\\n\", compare0 ? \"FAIL\" : \"PASS\");\n }\n\n if (RLE_METHOD != RLE)\n {\n compare1 = CompareDeviceResults(h_offsets_reference, d_offsets_out, num_runs, true, g_verbose);\n printf(\"\\t Offsets %s\\n\", compare1 ? \"FAIL\" : \"PASS\");\n }\n\n compare2 = CompareDeviceResults(h_lengths_reference, d_lengths_out, num_runs, true, g_verbose);\n printf(\"\\t Lengths %s\\n\", compare2 ? \"FAIL\" : \"PASS\");\n\n compare3 = CompareDeviceResults(&num_runs, d_num_runs, 1, true, g_verbose);\n printf(\"\\t Count %s\\n\", compare3 ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Performance\n GpuTimer gpu_timer;\n gpu_timer.Start();\n CubDebugExit(Dispatch(Int2Type(),\n Int2Type(),\n g_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_unique_out,\n d_offsets_out,\n d_lengths_out,\n d_num_runs,\n equality_op,\n num_items));\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n\n // Display performance\n if (g_timing_iterations > 0)\n {\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n int bytes_moved = (num_items * sizeof(T)) + (num_runs * (sizeof(OffsetT) + sizeof(LengthT)));\n float giga_bandwidth = float(bytes_moved) / avg_millis / 1000.0f / 1000.0f;\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s\",\n avg_millis,\n giga_rate,\n giga_bandwidth);\n }\n printf(\"\\n\\n\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Cleanup\n if (d_unique_out)\n CubDebugExit(g_allocator.DeviceFree(d_unique_out));\n if (d_offsets_out)\n CubDebugExit(g_allocator.DeviceFree(d_offsets_out));\n if (d_lengths_out)\n CubDebugExit(g_allocator.DeviceFree(d_lengths_out));\n if (d_num_runs)\n CubDebugExit(g_allocator.DeviceFree(d_num_runs));\n if (d_temp_storage_bytes)\n CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error)\n CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage)\n CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare0 | compare1 | compare2 | compare3);\n}\n\n/**\n * Test DeviceRunLengthEncode on pointer type\n */\ntemplate \nvoid TestPointer(int num_items, int entropy_reduction, int max_segment)\n{\n // Allocate host arrays\n T *h_in = new T[num_items];\n T *h_unique_reference = new T[num_items];\n OffsetT *h_offsets_reference = new OffsetT[num_items];\n LengthT *h_lengths_reference = new LengthT[num_items];\n\n for (int i = 0; i < num_items; ++i)\n InitValue(INTEGER_SEED, h_offsets_reference[i], 1);\n\n // Initialize problem and solution\n Equality equality_op;\n Initialize(entropy_reduction, h_in, num_items, max_segment);\n\n int num_runs = Solve(h_in,\n h_unique_reference,\n h_offsets_reference,\n h_lengths_reference,\n equality_op,\n num_items);\n\n printf(\"\\nPointer %s cub::%s on %d items, %d segments (avg run length %.3f), {%s key, %s offset, \"\n \"%s length}, max_segment %d, entropy_reduction %d\\n\",\n (RLE_METHOD == RLE) ? \"DeviceRunLengthEncode::Encode\"\n : (RLE_METHOD == NON_TRIVIAL) ? \"DeviceRunLengthEncode::NonTrivialRuns\"\n : \"Other\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items,\n num_runs,\n float(num_items) / num_runs,\n typeid(T).name(),\n typeid(OffsetT).name(),\n typeid(LengthT).name(),\n max_segment,\n entropy_reduction);\n fflush(stdout);\n\n // Allocate problem device arrays\n T *d_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_in, sizeof(T) * num_items));\n\n // Initialize device input\n CubDebugExit(hipMemcpy(d_in, h_in, sizeof(T) * num_items, hipMemcpyHostToDevice));\n\n // Run Test\n Test(d_in,\n h_unique_reference,\n h_offsets_reference,\n h_lengths_reference,\n equality_op,\n num_runs,\n num_items);\n\n // Cleanup\n if (h_in)\n delete[] h_in;\n if (h_unique_reference)\n delete[] h_unique_reference;\n if (h_offsets_reference)\n delete[] h_offsets_reference;\n if (h_lengths_reference)\n delete[] h_lengths_reference;\n if (d_in)\n CubDebugExit(g_allocator.DeviceFree(d_in));\n}\n\n/**\n * Test on iterator type\n */\ntemplate \nvoid TestIterator(int num_items, Int2Type /*is_primitive*/)\n{\n // Allocate host arrays\n T *h_unique_reference = new T[num_items];\n OffsetT *h_offsets_reference = new OffsetT[num_items];\n LengthT *h_lengths_reference = new LengthT[num_items];\n\n T one_val;\n InitValue(INTEGER_SEED, one_val, 1);\n ConstantInputIterator h_in(one_val);\n\n // Initialize problem and solution\n Equality equality_op;\n int num_runs = Solve(h_in,\n h_unique_reference,\n h_offsets_reference,\n h_lengths_reference,\n equality_op,\n num_items);\n\n printf(\"\\nIterator %s cub::%s on %d items, %d segments (avg run length %.3f), {%s key, %s \"\n \"offset, %s length}\\n\",\n (RLE_METHOD == RLE) ? \"DeviceRunLengthEncode::Encode\"\n : (RLE_METHOD == NON_TRIVIAL) ? \"DeviceRunLengthEncode::NonTrivialRuns\"\n : \"Other\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items,\n num_runs,\n float(num_items) / num_runs,\n typeid(T).name(),\n typeid(OffsetT).name(),\n typeid(LengthT).name());\n fflush(stdout);\n\n // Run Test\n Test(h_in,\n h_unique_reference,\n h_offsets_reference,\n h_lengths_reference,\n equality_op,\n num_runs,\n num_items);\n\n // Cleanup\n if (h_unique_reference)\n delete[] h_unique_reference;\n if (h_offsets_reference)\n delete[] h_offsets_reference;\n if (h_lengths_reference)\n delete[] h_lengths_reference;\n}\n\ntemplate \nvoid TestIterator(int /*num_items*/, Int2Type /*is_primitive*/)\n{}\n\n/**\n * Test different gen modes\n */\ntemplate \nvoid Test(int num_items)\n{\n // Test iterator (one run)\n TestIterator(num_items,\n Int2Type::PRIMITIVE>());\n\n // Evaluate different run lengths / segment sizes\n const int max_seg_limit = CUB_MIN(num_items, 1 << 16);\n const int max_seg_inc = 4;\n for (int max_segment = 1, entropy_reduction = 0; max_segment <= max_seg_limit;\n max_segment <<= max_seg_inc, entropy_reduction++)\n {\n const int max_seg = CUB_MAX(1, max_segment);\n TestPointer(num_items, entropy_reduction, max_seg);\n }\n}\n\n/**\n * Test different dispatch\n */\ntemplate \nvoid TestDispatch(int num_items)\n{\n#if TEST_CDP == 0\n Test(num_items);\n Test(num_items);\n#elif TEST_CDP == 1\n Test(num_items);\n Test(num_items);\n#endif\n}\n\n/**\n * Test different input sizes\n */\ntemplate \nvoid TestSize(int num_items)\n{\n if (num_items < 0)\n {\n TestDispatch(0);\n TestDispatch(1);\n TestDispatch(100);\n TestDispatch(10000);\n TestDispatch(1000000);\n }\n else\n {\n TestDispatch(num_items);\n }\n}\n\n/**\n * @brief Test with NaNs as input to ensure we don't create an additional (invalid) run on the first\n * item when (item[0] == item[0]) is false.\n */\ntemplate \nvoid TestNaNs()\n{\n using OffsetT = int32_t;\n using LengthT = int32_t;\n\n const auto quiet_nan = std::numeric_limits::quiet_NaN();\n\n // Allocate host data\n std::vector h_in{quiet_nan, quiet_nan, quiet_nan};\n const OffsetT num_items = static_cast(h_in.size());\n T *h_unique_reference = new T[num_items];\n OffsetT *h_offsets_reference = new OffsetT[num_items];\n LengthT *h_lengths_reference = new LengthT[num_items];\n\n // Initialize problem and solution\n Equality equality_op;\n\n // Get host-side data for verification\n int num_runs = Solve(h_in.data(),\n h_unique_reference,\n h_offsets_reference,\n h_lengths_reference,\n equality_op,\n num_items);\n\n // Allocate problem device arrays\n thrust::device_vector d_in(num_items);\n d_in = h_in;\n\n // Run Test\n Test(thrust::raw_pointer_cast(d_in.data()),\n h_unique_reference,\n h_offsets_reference,\n h_lengths_reference,\n equality_op,\n num_runs,\n num_items);\n}\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char **argv)\n{\n int num_items = -1;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\",\n argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n printf(\"\\n\");\n\n // %PARAM% TEST_CDP cdp 0:1\n\n // Run tests with leading/trailing NaNs, where (NaN == NaN) is false\n TestNaNs();\n TestNaNs();\n\n // Test different input types\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n TestSize(num_items);\n\n return 0;\n}\n", "source": "/kaggle/working/cub/test/test_device_run_length_encode.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n\n#include \n\n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#define TEST_HALF_T !_NVHPC_CUDA\n\n#define TEST_BF_T !_NVHPC_CUDA\n\n#if TEST_HALF_T\n#include \n#endif\n\n#if TEST_BF_T\n#include \n#endif\n\nusing namespace cub;\n\ntemplate \nstruct UnwrapHalfAndBfloat16\n{\n using Type = T;\n};\n\n#if TEST_HALF_T\ntemplate <>\nstruct UnwrapHalfAndBfloat16\n{\n using Type = __half;\n};\n#endif\n\n#if TEST_BF_T\ntemplate <>\nstruct UnwrapHalfAndBfloat16\n{\n using Type = __nv_bfloat16;\n};\n#endif\n\nconstexpr static int MAX_ITERATIONS = 2;\n\n\nclass SizeGroupDescription\n{\npublic:\n SizeGroupDescription(const int segments,\n const int segment_size)\n : segments(segments)\n , segment_size(segment_size)\n {}\n\n int segments {};\n int segment_size {};\n};\n\ntemplate \nstruct SegmentChecker\n{\n const KeyT *sorted_keys {};\n const int *offsets {};\n\n SegmentChecker(const KeyT *sorted_keys,\n const int *offsets)\n : sorted_keys(sorted_keys)\n , offsets(offsets)\n {}\n\n bool operator()(int segment_id)\n {\n const int segment_begin = offsets[segment_id];\n const int segment_end = offsets[segment_id + 1];\n\n int counter = 0;\n for (int i = segment_begin; i < segment_end; i++)\n {\n if (sorted_keys[i] != static_cast(counter++))\n {\n return false;\n }\n }\n\n return true;\n }\n};\n\ntemplate \nstruct DescendingSegmentChecker\n{\n const KeyT *sorted_keys{};\n const int *offsets{};\n\n DescendingSegmentChecker(const KeyT *sorted_keys,\n const int *offsets)\n : sorted_keys(sorted_keys)\n , offsets(offsets)\n {}\n\n bool operator()(int segment_id)\n {\n const int segment_begin = offsets[segment_id];\n const int segment_end = offsets[segment_id + 1];\n\n int counter = 0;\n for (int i = segment_end - 1; i >= segment_begin; i--)\n {\n if (sorted_keys[i] != static_cast(counter++))\n {\n return false;\n }\n }\n\n return true;\n }\n};\n\ntemplate \nstruct ReversedIota\n{\n KeyT *data {};\n const int *offsets {};\n\n ReversedIota(KeyT *data,\n const int *offsets)\n : data(data)\n , offsets(offsets)\n {}\n\n void operator()(int segment_id) const\n {\n const int segment_begin = offsets[segment_id];\n const int segment_end = offsets[segment_id + 1];\n const int segment_size = segment_end - segment_begin;\n\n int count = 0;\n for (int i = segment_begin; i < segment_end; i++)\n {\n data[i] = static_cast(segment_size - 1 - count++);\n }\n }\n};\n\n\ntemplate \nstruct Iota\n{\n KeyT *data{};\n const int *offsets{};\n\n Iota(KeyT *data, const int *offsets)\n : data(data)\n , offsets(offsets)\n {}\n\n void operator()(int segment_id) const\n {\n const int segment_begin = offsets[segment_id];\n const int segment_end = offsets[segment_id + 1];\n\n int count = 0;\n for (int i = segment_begin; i < segment_end; i++)\n {\n data[i] = static_cast(count++);\n }\n }\n};\n\n\ntemplate \nclass Input\n{\n thrust::default_random_engine random_engine;\n thrust::device_vector d_segment_sizes;\n thrust::device_vector d_offsets;\n thrust::host_vector h_offsets;\n\n using MaskedValueT = cub::detail::conditional_t<\n std::is_same::value, KeyT, ValueT>;\n\n bool reverse {};\n int num_items {};\n thrust::device_vector d_keys;\n thrust::device_vector d_values;\n thrust::host_vector h_keys;\n thrust::host_vector h_values;\n\npublic:\n Input(bool reverse, const thrust::host_vector &h_segment_sizes)\n : d_segment_sizes(h_segment_sizes)\n , d_offsets(d_segment_sizes.size() + 1)\n , h_offsets(d_segment_sizes.size() + 1)\n , reverse(reverse)\n , num_items(static_cast(\n thrust::reduce(d_segment_sizes.begin(), d_segment_sizes.end())))\n , d_keys(num_items)\n , d_values(num_items)\n , h_keys(num_items)\n , h_values(num_items)\n {\n update();\n }\n\n Input(thrust::host_vector &h_offsets)\n : d_offsets(h_offsets)\n , h_offsets(h_offsets)\n , reverse(false)\n , num_items(h_offsets.back())\n , d_keys(num_items)\n , d_values(num_items)\n {\n }\n\n void shuffle()\n {\n thrust::shuffle(d_segment_sizes.begin(), d_segment_sizes.end(), random_engine);\n\n update();\n }\n\n int get_num_items() const\n {\n return num_items;\n }\n\n int get_num_segments() const\n {\n return static_cast(d_segment_sizes.size());\n }\n\n const KeyT *get_d_keys() const\n {\n return thrust::raw_pointer_cast(d_keys.data());\n }\n\n thrust::device_vector &get_d_keys_vec()\n {\n return d_keys;\n }\n\n thrust::device_vector &get_d_values_vec()\n {\n return d_values;\n }\n\n KeyT *get_d_keys()\n {\n return thrust::raw_pointer_cast(d_keys.data());\n }\n\n const thrust::host_vector& get_h_offsets()\n {\n return h_offsets;\n }\n\n MaskedValueT *get_d_values()\n {\n return thrust::raw_pointer_cast(d_values.data());\n }\n\n const int *get_d_offsets() const\n {\n return thrust::raw_pointer_cast(d_offsets.data());\n }\n\n template \n bool check_output_implementation(const T *keys_output)\n {\n const int *offsets = thrust::raw_pointer_cast(h_offsets.data());\n\n if (reverse)\n {\n DescendingSegmentChecker checker{keys_output, offsets};\n\n for (int i = 0; i < get_num_segments(); i++)\n {\n if (!checker(i))\n {\n return false;\n }\n }\n }\n else\n {\n SegmentChecker checker{keys_output, offsets};\n\n for (int i = 0; i < get_num_segments(); i++)\n {\n if (!checker(i))\n {\n return false;\n }\n }\n }\n\n return true;\n }\n\n bool check_output(const KeyT *d_keys_output,\n const MaskedValueT *d_values_output = nullptr)\n {\n KeyT *keys_output = thrust::raw_pointer_cast(h_keys.data());\n MaskedValueT *values_output = thrust::raw_pointer_cast(h_values.data());\n\n cudaMemcpy(keys_output,\n d_keys_output,\n sizeof(KeyT) * num_items,\n cudaMemcpyDeviceToHost);\n\n const bool keys_ok = check_output_implementation(keys_output);\n\n if (std::is_same::value || d_values_output == nullptr)\n {\n return keys_ok;\n }\n\n cudaMemcpy(values_output,\n d_values_output,\n sizeof(ValueT) * num_items,\n cudaMemcpyDeviceToHost);\n\n const bool values_ok = check_output_implementation(values_output);\n\n return keys_ok && values_ok;\n }\n\nprivate:\n void update()\n {\n fill_offsets();\n gen_keys();\n }\n\n void fill_offsets()\n {\n thrust::copy(d_segment_sizes.begin(), d_segment_sizes.end(), d_offsets.begin());\n thrust::exclusive_scan(d_offsets.begin(), d_offsets.end(), d_offsets.begin(), 0u);\n thrust::copy(d_offsets.begin(), d_offsets.end(), h_offsets.begin());\n }\n\n void gen_keys()\n {\n KeyT *keys_output = thrust::raw_pointer_cast(h_keys.data());\n const int *offsets = thrust::raw_pointer_cast(h_offsets.data());\n\n if (reverse)\n {\n Iota generator{keys_output, offsets};\n\n for (int i = 0; i < get_num_segments(); i++)\n {\n generator(i);\n }\n }\n else\n {\n ReversedIota generator{keys_output, offsets};\n\n for (int i = 0; i < get_num_segments(); i++)\n {\n generator(i);\n }\n }\n\n d_keys = h_keys;\n d_values = d_keys;\n }\n};\n\ntemplate ::value>\nclass InputDescription\n{\n thrust::host_vector segment_sizes;\n\npublic:\n InputDescription& add(const SizeGroupDescription &group)\n {\n if (static_cast(group.segment_size) <\n static_cast((std::numeric_limits::max)()))\n {\n for (int i = 0; i < group.segments; i++)\n {\n segment_sizes.push_back(group.segment_size);\n }\n }\n\n return *this;\n }\n\n template \n Input gen(bool reverse)\n {\n return Input(reverse, segment_sizes);\n }\n};\n\ntemplate \nclass InputDescription\n{\n thrust::host_vector segment_sizes;\n\npublic:\n InputDescription& add(const SizeGroupDescription &group)\n {\n for (int i = 0; i < group.segments; i++)\n {\n segment_sizes.push_back(group.segment_size);\n }\n\n return *this;\n }\n\n template \n Input gen(bool reverse)\n {\n return Input(reverse, segment_sizes);\n }\n};\n\n\ntemplate \nvoid Sort(bool pairs,\n bool descending,\n bool double_buffer,\n bool stable_sort,\n\n void *tmp_storage,\n std::size_t &temp_storage_bytes,\n\n WrappedKeyT *wrapped_input_keys,\n WrappedKeyT *wrapped_output_keys,\n\n ValueT *input_values,\n ValueT *output_values,\n\n int num_items,\n int num_segments,\n const int *d_offsets,\n\n int *keys_selector = nullptr,\n int *values_selector = nullptr)\n{\n using KeyT = typename UnwrapHalfAndBfloat16::Type;\n\n auto input_keys = reinterpret_cast(wrapped_input_keys);\n auto output_keys = reinterpret_cast(wrapped_output_keys);\n\n if (stable_sort)\n {\n if (pairs)\n {\n if (descending)\n {\n if (double_buffer)\n {\n cub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n cub::DoubleBuffer values_buffer(input_values, output_values);\n values_buffer.selector = *values_selector;\n\n CubDebugExit(cub::DeviceSegmentedSort::StableSortPairsDescending(\n tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n values_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n *values_selector = values_buffer.selector;\n }\n else\n {\n CubDebugExit(cub::DeviceSegmentedSort::StableSortPairsDescending(\n tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n input_values,\n output_values,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n else\n {\n if (double_buffer)\n {\n cub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n cub::DoubleBuffer values_buffer(input_values, output_values);\n values_buffer.selector = *values_selector;\n\n CubDebugExit(\n cub::DeviceSegmentedSort::StableSortPairs(tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n values_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n *values_selector = values_buffer.selector;\n }\n else\n {\n CubDebugExit(\n cub::DeviceSegmentedSort::StableSortPairs(tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n input_values,\n output_values,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n }\n else\n {\n if (descending)\n {\n if (double_buffer)\n {\n cub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n CubDebugExit(cub::DeviceSegmentedSort::StableSortKeysDescending(\n tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n }\n else\n {\n CubDebugExit(cub::DeviceSegmentedSort::StableSortKeysDescending(\n tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n else\n {\n if (double_buffer)\n {\n cub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n CubDebugExit(\n cub::DeviceSegmentedSort::StableSortKeys(tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n }\n else\n {\n CubDebugExit(\n cub::DeviceSegmentedSort::StableSortKeys(tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n }\n }\n else\n {\n if (pairs)\n {\n if (descending)\n {\n if (double_buffer)\n {\n cub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n cub::DoubleBuffer values_buffer(input_values, output_values);\n values_buffer.selector = *values_selector;\n\n CubDebugExit(\n cub::DeviceSegmentedSort::SortPairsDescending(tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n values_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n *values_selector = values_buffer.selector;\n }\n else\n {\n CubDebugExit(\n cub::DeviceSegmentedSort::SortPairsDescending(tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n input_values,\n output_values,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n else\n {\n if (double_buffer)\n {\n cub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n cub::DoubleBuffer values_buffer(input_values, output_values);\n values_buffer.selector = *values_selector;\n\n CubDebugExit(cub::DeviceSegmentedSort::SortPairs(tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n values_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n *values_selector = values_buffer.selector;\n }\n else\n {\n CubDebugExit(cub::DeviceSegmentedSort::SortPairs(tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n input_values,\n output_values,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n }\n else\n {\n if (descending)\n {\n if (double_buffer)\n {\n cub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n CubDebugExit(\n cub::DeviceSegmentedSort::SortKeysDescending(tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n }\n else\n {\n CubDebugExit(\n cub::DeviceSegmentedSort::SortKeysDescending(tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n else\n {\n if (double_buffer)\n {\n cub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n CubDebugExit(cub::DeviceSegmentedSort::SortKeys(tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n }\n else\n {\n CubDebugExit(cub::DeviceSegmentedSort::SortKeys(tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n }\n }\n}\n\ntemplate \nstd::size_t Sort(bool pairs,\n bool descending,\n bool double_buffer,\n bool stable_sort,\n\n KeyT *input_keys,\n KeyT *output_keys,\n\n ValueT *input_values,\n ValueT *output_values,\n\n int num_items,\n int num_segments,\n const int *d_offsets,\n\n int *keys_selector = nullptr,\n int *values_selector = nullptr)\n{\n std::size_t temp_storage_bytes = 42ul;\n\n Sort(pairs,\n descending,\n double_buffer,\n stable_sort,\n nullptr,\n temp_storage_bytes,\n input_keys,\n output_keys,\n input_values,\n output_values,\n num_items,\n num_segments,\n d_offsets,\n keys_selector,\n values_selector);\n\n thrust::device_vector temp_storage(temp_storage_bytes);\n std::uint8_t *d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n Sort(pairs,\n descending,\n double_buffer,\n stable_sort,\n d_temp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n input_values,\n output_values,\n num_items,\n num_segments,\n d_offsets,\n keys_selector,\n values_selector);\n\n return temp_storage_bytes;\n}\n\n\nconstexpr bool keys_only = false;\nconstexpr bool pairs = true;\n\nconstexpr bool ascending = false;\nconstexpr bool descending = true;\n\nconstexpr bool pointers = false;\nconstexpr bool double_buffer = true;\n\nconstexpr bool unstable = false;\nconstexpr bool stable = true;\n\n\nvoid TestZeroSegments()\n{\n // Type doesn't affect the escape logic, so it should be fine\n // to test only one set of types here.\n\n using KeyT = std::uint8_t;\n using ValueT = std::uint64_t;\n\n for (bool stable_sort: { unstable, stable })\n {\n for (bool sort_pairs: { keys_only, pairs })\n {\n for (bool sort_descending: { ascending, descending })\n {\n for (bool sort_buffer: { pointers, double_buffer })\n {\n cub::DoubleBuffer keys_buffer(nullptr, nullptr);\n cub::DoubleBuffer values_buffer(nullptr, nullptr);\n values_buffer.selector = 1;\n\n Sort(sort_pairs,\n sort_descending,\n sort_buffer,\n stable_sort,\n nullptr,\n nullptr,\n nullptr,\n nullptr,\n int{},\n int{},\n nullptr,\n &keys_buffer.selector,\n &values_buffer.selector);\n\n AssertEquals(keys_buffer.selector, 0);\n AssertEquals(values_buffer.selector, 1);\n }\n }\n }\n }\n}\n\n\nvoid TestEmptySegments(int segments)\n{\n // Type doesn't affect the escape logic, so it should be fine\n // to test only one set of types here.\n\n using KeyT = std::uint8_t;\n using ValueT = std::uint64_t;\n\n thrust::device_vector offsets(segments + 1, int{});\n const int *d_offsets = thrust::raw_pointer_cast(offsets.data());\n\n for (bool sort_stable: { unstable, stable })\n {\n for (bool sort_pairs: { keys_only, pairs })\n {\n for (bool sort_descending: { ascending, descending })\n {\n for (bool sort_buffer: { pointers, double_buffer })\n {\n cub::DoubleBuffer keys_buffer(nullptr, nullptr);\n cub::DoubleBuffer values_buffer(nullptr, nullptr);\n values_buffer.selector = 1;\n\n Sort(sort_pairs,\n sort_descending,\n sort_buffer,\n sort_stable,\n nullptr,\n nullptr,\n nullptr,\n nullptr,\n int{},\n segments,\n d_offsets,\n &keys_buffer.selector,\n &values_buffer.selector);\n\n AssertEquals(keys_buffer.selector, 0);\n AssertEquals(values_buffer.selector, 1);\n }\n }\n }\n }\n}\n\n\ntemplate \nvoid TestSameSizeSegments(int segment_size,\n int segments,\n bool skip_values = false)\n{\n const int num_items = segment_size * segments;\n\n thrust::device_vector offsets(segments + 1);\n thrust::sequence(offsets.begin(),\n offsets.end(),\n int{},\n segment_size);\n\n const int *d_offsets = thrust::raw_pointer_cast(offsets.data());\n\n const KeyT target_key {1};\n const ValueT target_value {42};\n\n thrust::device_vector keys_input(num_items);\n thrust::device_vector keys_output(num_items);\n\n KeyT *d_keys_input = thrust::raw_pointer_cast(keys_input.data());\n KeyT *d_keys_output = thrust::raw_pointer_cast(keys_output.data());\n\n thrust::device_vector values_input(num_items);\n thrust::device_vector values_output(num_items);\n\n thrust::host_vector host_keys(num_items);\n thrust::host_vector host_values(num_items);\n\n ValueT *d_values_input = thrust::raw_pointer_cast(values_input.data());\n ValueT *d_values_output = thrust::raw_pointer_cast(values_output.data());\n\n for (bool stable_sort: { unstable, stable })\n {\n for (bool sort_pairs: { keys_only, pairs })\n {\n if (sort_pairs)\n {\n if (skip_values)\n {\n continue;\n }\n }\n\n for (bool sort_descending: { ascending, descending })\n {\n for (bool sort_buffers: { pointers, double_buffer })\n {\n cub::DoubleBuffer keys_buffer(nullptr, nullptr);\n cub::DoubleBuffer values_buffer(nullptr, nullptr);\n values_buffer.selector = 1;\n\n thrust::fill(keys_input.begin(), keys_input.end(), target_key);\n thrust::fill(keys_output.begin(), keys_output.end(), KeyT{});\n\n if (sort_pairs)\n {\n if (sort_buffers)\n {\n thrust::fill(values_input.begin(), values_input.end(), ValueT{});\n thrust::fill(values_output.begin(), values_output.end(), target_value);\n }\n else\n {\n thrust::fill(values_input.begin(), values_input.end(), target_value);\n thrust::fill(values_output.begin(), values_output.end(), ValueT{});\n }\n }\n\n const std::size_t temp_storage_bytes =\n Sort(sort_pairs,\n sort_descending,\n sort_buffers,\n stable_sort,\n d_keys_input,\n d_keys_output,\n d_values_input,\n d_values_output,\n num_items,\n segments,\n d_offsets,\n &keys_buffer.selector,\n &values_buffer.selector);\n\n // If temporary storage size is defined by extra keys storage\n if (sort_buffers)\n {\n if (2 * segments * sizeof(unsigned int) < num_items * sizeof(KeyT))\n {\n std::size_t extra_temp_storage_bytes{};\n\n Sort(sort_pairs,\n sort_descending,\n pointers,\n stable_sort,\n nullptr,\n extra_temp_storage_bytes,\n d_keys_input,\n d_keys_output,\n d_values_input,\n d_values_output,\n num_items,\n segments,\n d_offsets,\n &keys_buffer.selector,\n &values_buffer.selector);\n\n AssertTrue(extra_temp_storage_bytes > temp_storage_bytes);\n }\n }\n\n {\n host_keys = keys_buffer.selector || !sort_buffers ? keys_output\n : keys_input;\n const std::size_t items_selected =\n thrust::count(host_keys.begin(), host_keys.end(), target_key);\n AssertEquals(static_cast(items_selected), num_items);\n }\n\n if (sort_pairs)\n {\n host_values = values_buffer.selector || !sort_buffers\n ? values_output\n : values_input;\n const std::size_t items_selected =\n thrust::count(host_values.begin(),\n host_values.end(),\n target_value);\n\n AssertEquals(static_cast(items_selected), num_items);\n }\n }\n }\n }\n }\n}\n\n\ntemplate \nvoid InputTest(bool sort_descending,\n Input &input)\n{\n thrust::device_vector keys_output(input.get_num_items());\n KeyT *d_keys_output = thrust::raw_pointer_cast(keys_output.data());\n\n thrust::device_vector values_output(input.get_num_items());\n ValueT *d_values_output = thrust::raw_pointer_cast(values_output.data());\n\n for (bool stable_sort: { unstable, stable })\n {\n for (bool sort_pairs : { keys_only, pairs })\n {\n for (bool sort_buffers : {pointers, double_buffer})\n {\n for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++)\n {\n thrust::fill(keys_output.begin(), keys_output.end(), KeyT{});\n thrust::fill(values_output.begin(), values_output.end(), ValueT{});\n\n cub::DoubleBuffer keys_buffer(input.get_d_keys(),\n d_keys_output);\n cub::DoubleBuffer values_buffer(input.get_d_values(),\n d_values_output);\n\n Sort(sort_pairs,\n sort_descending,\n sort_buffers,\n stable_sort,\n input.get_d_keys(),\n d_keys_output,\n input.get_d_values(),\n d_values_output,\n input.get_num_items(),\n input.get_num_segments(),\n input.get_d_offsets(),\n &keys_buffer.selector,\n &values_buffer.selector);\n\n if (sort_buffers)\n {\n if (sort_pairs)\n {\n AssertTrue(input.check_output(keys_buffer.Current(),\n values_buffer.Current()));\n }\n else\n {\n AssertTrue(input.check_output(keys_buffer.Current()));\n }\n }\n else\n {\n if (sort_pairs)\n {\n AssertTrue(input.check_output(d_keys_output, d_values_output));\n }\n else\n {\n AssertTrue(input.check_output(d_keys_output));\n }\n }\n\n input.shuffle();\n }\n }\n }\n }\n}\n\nstruct ComparisonPredicate\n{\n template \n __host__ __device__ bool operator()(const T &lhs, const T &rhs) const\n {\n return lhs == rhs;\n }\n\n __host__ __device__ bool operator()(const half_t &lhs, const half_t &rhs) const\n {\n return lhs.raw() == rhs.raw();\n }\n};\n\ntemplate \nbool compare_two_outputs(const thrust::host_vector &offsets,\n const thrust::host_vector &lhs,\n const thrust::host_vector &rhs)\n{\n const auto num_segments = static_cast(offsets.size() - 1);\n\n for (std::size_t segment_id = 0; segment_id < num_segments; segment_id++)\n {\n auto lhs_begin = lhs.cbegin() + offsets[segment_id];\n auto lhs_end = lhs.cbegin() + offsets[segment_id + 1];\n auto rhs_begin = rhs.cbegin() + offsets[segment_id];\n\n auto err = thrust::mismatch(lhs_begin, lhs_end, rhs_begin, ComparisonPredicate{});\n\n if (err.first != lhs_end)\n {\n const auto idx = thrust::distance(lhs_begin, err.first);\n const auto segment_size = std::distance(lhs_begin, lhs_end);\n\n std::cerr << \"Mismatch in segment \" << segment_id\n << \" at position \" << idx << \" / \" << segment_size\n << \": \"\n << static_cast(lhs_begin[idx]) << \" vs \"\n << static_cast(rhs_begin[idx]) << \" (\"\n << typeid(lhs_begin[idx]).name() << \")\" << std::endl;\n\n return false;\n }\n }\n\n return true;\n}\n\ntemplate \nvoid RandomizeInput(thrust::host_vector &h_keys,\n thrust::host_vector &h_values)\n{\n for (std::size_t i = 0; i < h_keys.size(); i++)\n {\n h_keys[i] = RandomValue((std::numeric_limits::max)()) > 128;\n h_values[i] = RandomValue((std::numeric_limits::max)());\n }\n}\n\ntemplate \nvoid RandomizeInput(thrust::host_vector &h_keys,\n thrust::host_vector &h_values)\n{\n for (std::size_t i = 0; i < h_keys.size(); i++)\n {\n h_keys[i] = RandomValue((std::numeric_limits::max)());\n h_values[i] = RandomValue((std::numeric_limits::max)());\n }\n}\n\n#if TEST_HALF_T\nvoid RandomizeInput(thrust::host_vector &h_keys,\n thrust::host_vector &h_values)\n{\n for (std::size_t i = 0; i < h_keys.size(); i++)\n {\n h_keys[i] = RandomValue((std::numeric_limits::max)());\n h_values[i] = RandomValue((std::numeric_limits::max)());\n }\n}\n#endif\n\n#if TEST_BF_T\nvoid RandomizeInput(thrust::host_vector &h_keys,\n thrust::host_vector &h_values)\n{\n for (std::size_t i = 0; i < h_keys.size(); i++)\n {\n h_keys[i] = RandomValue((std::numeric_limits::max)());\n h_values[i] = RandomValue((std::numeric_limits::max)());\n }\n}\n#endif\n\n\n\ntemplate \nvoid HostReferenceSort(bool sort_pairs,\n bool sort_descending,\n unsigned int num_segments,\n const thrust::host_vector &h_offsets,\n thrust::host_vector &h_keys,\n thrust::host_vector &h_values)\n{\n for (unsigned int segment_i = 0;\n segment_i < num_segments;\n segment_i++)\n {\n const int segment_begin = h_offsets[segment_i];\n const int segment_end = h_offsets[segment_i + 1];\n\n if (sort_pairs)\n {\n if (sort_descending)\n {\n thrust::stable_sort_by_key(h_keys.begin() + segment_begin,\n h_keys.begin() + segment_end,\n h_values.begin() + segment_begin,\n thrust::greater{});\n }\n else\n {\n thrust::stable_sort_by_key(h_keys.begin() + segment_begin,\n h_keys.begin() + segment_end,\n h_values.begin() + segment_begin);\n }\n }\n else\n {\n if (sort_descending)\n {\n thrust::stable_sort(h_keys.begin() + segment_begin,\n h_keys.begin() + segment_end,\n thrust::greater{});\n }\n else\n {\n thrust::stable_sort(h_keys.begin() + segment_begin,\n h_keys.begin() + segment_end);\n }\n }\n }\n}\n\n\n#if STORE_ON_FAILURE\ntemplate \nvoid DumpInput(bool sort_pairs,\n bool sort_descending,\n bool sort_buffers,\n Input &input,\n thrust::host_vector &h_keys,\n thrust::host_vector &h_values)\n{\n const thrust::host_vector &h_offsets = input.get_h_offsets();\n\n std::cout << \"sort pairs: \" << sort_pairs << \"\\n\";\n std::cout << \"sort descending: \" << sort_descending << \"\\n\";\n std::cout << \"sort buffers: \" << sort_buffers << \"\\n\";\n std::cout << \"num_items: \" << input.get_num_items() << \"\\n\";\n std::cout << \"num_segments: \" << input.get_num_segments() << \"\\n\";\n std::cout << \"key type: \" << typeid(h_keys[0]).name() << \"\\n\";\n std::cout << \"value type: \" << typeid(h_values[0]).name() << \"\\n\";\n std::cout << \"offset type: \" << typeid(h_offsets[0]).name() << \"\\n\";\n\n std::ofstream offsets_dump(\"offsets\", std::ios::binary);\n offsets_dump.write(reinterpret_cast(\n thrust::raw_pointer_cast(h_offsets.data())),\n sizeof(int) * h_offsets.size());\n\n std::ofstream keys_dump(\"keys\", std::ios::binary);\n keys_dump.write(reinterpret_cast(\n thrust::raw_pointer_cast(h_keys.data())),\n sizeof(KeyT) * h_keys.size());\n\n std::ofstream values_dump(\"values\", std::ios::binary);\n values_dump.write(reinterpret_cast(\n thrust::raw_pointer_cast(h_values.data())),\n sizeof(ValueT) * h_values.size());\n}\n#endif\n\n\ntemplate \nvoid InputTestRandom(Input &input)\n{\n thrust::host_vector h_keys_output(input.get_num_items());\n thrust::device_vector keys_output(input.get_num_items());\n\n thrust::host_vector h_values_output(input.get_num_items());\n thrust::device_vector values_output(input.get_num_items());\n\n KeyT *d_keys_output = thrust::raw_pointer_cast(keys_output.data());\n ValueT *d_values_output = thrust::raw_pointer_cast(values_output.data());\n\n thrust::host_vector h_keys(input.get_num_items());\n thrust::host_vector h_values(input.get_num_items());\n\n const thrust::host_vector &h_offsets = input.get_h_offsets();\n\n for (bool stable_sort: { unstable, stable })\n {\n for (bool sort_pairs: { keys_only, pairs })\n {\n for (bool sort_descending: { ascending, descending })\n {\n for (bool sort_buffers: { pointers, double_buffer })\n {\n for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++)\n {\n RandomizeInput(h_keys, h_values);\n\n#if STORE_ON_FAILURE\n auto h_keys_backup = h_keys;\n auto h_values_backup = h_values;\n#endif\n\n input.get_d_keys_vec() = h_keys;\n input.get_d_values_vec() = h_values;\n\n cub::DoubleBuffer keys_buffer(input.get_d_keys(), d_keys_output);\n cub::DoubleBuffer values_buffer(input.get_d_values(), d_values_output);\n\n Sort(sort_pairs,\n sort_descending,\n sort_buffers,\n stable_sort,\n input.get_d_keys(),\n d_keys_output,\n input.get_d_values(),\n d_values_output,\n input.get_num_items(),\n input.get_num_segments(),\n input.get_d_offsets(),\n &keys_buffer.selector,\n &values_buffer.selector);\n\n HostReferenceSort(sort_pairs,\n sort_descending,\n input.get_num_segments(),\n h_offsets,\n h_keys,\n h_values);\n\n if (sort_buffers)\n {\n if (keys_buffer.selector)\n {\n h_keys_output = keys_output;\n }\n else\n {\n h_keys_output = input.get_d_keys_vec();\n }\n\n if (values_buffer.selector)\n {\n h_values_output = values_output;\n }\n else\n {\n h_values_output = input.get_d_values_vec();\n }\n }\n else\n {\n h_keys_output = keys_output;\n h_values_output = values_output;\n }\n\n const bool keys_ok =\n compare_two_outputs(h_offsets, h_keys, h_keys_output);\n\n const bool values_ok =\n sort_pairs\n ? compare_two_outputs(h_offsets, h_values, h_values_output)\n : true;\n\n#if STORE_ON_FAILURE\n if (!keys_ok || !values_ok)\n {\n DumpInput(sort_pairs,\n sort_descending,\n sort_buffers,\n input,\n h_keys_backup,\n h_values_backup);\n }\n#endif\n\n AssertTrue(keys_ok);\n AssertTrue(values_ok);\n\n input.shuffle();\n }\n }\n }\n }\n }\n}\n\ntemplate ::value>\nstruct EdgeTestDispatch\n{\n // Edge cases that needs to be tested\n const int empty_short_circuit_segment_size = 0;\n const int copy_short_circuit_segment_size = 1;\n const int swap_short_circuit_segment_size = 2;\n\n const int a_few = 2;\n const int a_bunch_of = 42;\n const int a_lot_of = 420;\n\n template \n CUB_RUNTIME_FUNCTION __forceinline__ cudaError_t Invoke()\n {\n NV_IF_TARGET(NV_IS_HOST,\n (using SmallAndMediumPolicyT =\n typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT;\n using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy;\n\n const int small_segment_max_segment_size =\n SmallAndMediumPolicyT::SmallPolicyT::ITEMS_PER_TILE;\n\n const int items_per_small_segment =\n SmallAndMediumPolicyT::SmallPolicyT::ITEMS_PER_THREAD;\n\n const int medium_segment_max_segment_size =\n SmallAndMediumPolicyT::MediumPolicyT::ITEMS_PER_TILE;\n\n const int single_thread_segment_size = items_per_small_segment;\n\n const int large_cached_segment_max_segment_size =\n LargeSegmentPolicyT::BLOCK_THREADS *\n LargeSegmentPolicyT::ITEMS_PER_THREAD;\n\n for (bool sort_descending : {ascending, descending}) {\n Input edge_cases =\n InputDescription()\n .add({a_lot_of, empty_short_circuit_segment_size})\n .add({a_lot_of, copy_short_circuit_segment_size})\n .add({a_lot_of, swap_short_circuit_segment_size})\n .add({a_lot_of, swap_short_circuit_segment_size + 1})\n .add({a_lot_of, swap_short_circuit_segment_size + 1})\n .add({a_lot_of, single_thread_segment_size - 1})\n .add({a_lot_of, single_thread_segment_size})\n .add({a_lot_of, single_thread_segment_size + 1})\n .add({a_lot_of, single_thread_segment_size * 2 - 1})\n .add({a_lot_of, single_thread_segment_size * 2})\n .add({a_lot_of, single_thread_segment_size * 2 + 1})\n .add({a_bunch_of, small_segment_max_segment_size - 1})\n .add({a_bunch_of, small_segment_max_segment_size})\n .add({a_bunch_of, small_segment_max_segment_size + 1})\n .add({a_bunch_of, medium_segment_max_segment_size - 1})\n .add({a_bunch_of, medium_segment_max_segment_size})\n .add({a_bunch_of, medium_segment_max_segment_size + 1})\n .add({a_bunch_of, large_cached_segment_max_segment_size - 1})\n .add({a_bunch_of, large_cached_segment_max_segment_size})\n .add({a_bunch_of, large_cached_segment_max_segment_size + 1})\n .add({a_few, large_cached_segment_max_segment_size * 2})\n .add({a_few, large_cached_segment_max_segment_size * 3})\n .add({a_few, large_cached_segment_max_segment_size * 5})\n .template gen(sort_descending);\n\n InputTest(sort_descending, edge_cases);\n }));\n\n return cudaSuccess;\n }\n};\n\ntemplate \nstruct EdgeTestDispatch\n{\n template \n CUB_RUNTIME_FUNCTION __forceinline__ cudaError_t Invoke()\n {\n // Edge case test is using an optimized testing approach which is\n // incompatible with duplicates. RandomTest is used for other types.\n return cudaSuccess;\n }\n};\n\ntemplate \nvoid EdgePatternsTest()\n{\n int ptx_version = 0;\n if (CubDebug(PtxVersion(ptx_version)))\n {\n return;\n }\n\n using MaxPolicyT =\n typename cub::DeviceSegmentedSortPolicy::MaxPolicy;\n using EdgeTestDispatchT = EdgeTestDispatch;\n EdgeTestDispatchT dispatch;\n\n MaxPolicyT::Invoke(ptx_version, dispatch);\n\n}\n\ntemplate \nInput GenRandomInput(int max_items,\n int min_segments,\n int max_segments,\n bool descending)\n{\n int items_generated {};\n const int segments_num = RandomValue(max_segments) + min_segments;\n\n thrust::host_vector segment_sizes;\n segment_sizes.reserve(segments_num);\n\n const int max_segment_size = 6000;\n\n for (int segment_id = 0; segment_id < segments_num; segment_id++)\n {\n const int segment_size_raw = RandomValue(max_segment_size);\n const int segment_size = segment_size_raw > 0 ? segment_size_raw : 0;\n\n if (segment_size + items_generated > max_items)\n {\n break;\n }\n\n items_generated += segment_size;\n segment_sizes.push_back(segment_size);\n }\n\n return Input{descending, segment_sizes};\n}\n\ntemplate \nvoid RandomTest(int min_segments,\n int max_segments)\n{\n const int max_items = 10000000;\n\n for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++)\n {\n Input edge_cases = GenRandomInput(max_items,\n min_segments,\n max_segments,\n descending);\n\n InputTestRandom(edge_cases);\n }\n}\n\n\ntemplate \nvoid Test()\n{\n for (int segment_size: { 1, 1024, 24 * 1024 })\n {\n for (int segments: { 1, 1024 })\n {\n TestSameSizeSegments(segment_size, segments);\n }\n }\n\n RandomTest(1 << 2, 1 << 8);\n RandomTest(1 << 9, 1 << 19);\n EdgePatternsTest();\n}\n\n\n#if TEST_CDP == 1\ntemplate \n__global__ void LauncherKernel(\n void *tmp_storage,\n std::size_t temp_storage_bytes,\n const KeyT *in_keys,\n KeyT *out_keys,\n int num_items,\n int num_segments,\n const int *offsets)\n{\n CubDebug(cub::DeviceSegmentedSort::SortKeys(tmp_storage,\n temp_storage_bytes,\n in_keys,\n out_keys,\n num_items,\n num_segments,\n offsets,\n offsets + 1));\n}\n\ntemplate \nvoid TestDeviceSideLaunch(Input &input)\n{\n thrust::host_vector h_keys_output(input.get_num_items());\n thrust::device_vector keys_output(input.get_num_items());\n\n thrust::host_vector h_values_output(input.get_num_items());\n thrust::device_vector values_output(input.get_num_items());\n\n KeyT *d_keys_output = thrust::raw_pointer_cast(keys_output.data());\n\n thrust::host_vector h_keys(input.get_num_items());\n thrust::host_vector h_values(input.get_num_items());\n\n const thrust::host_vector &h_offsets = input.get_h_offsets();\n\n for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++)\n {\n RandomizeInput(h_keys, h_values);\n\n input.get_d_keys_vec() = h_keys;\n input.get_d_values_vec() = h_values;\n\n const KeyT *d_input = input.get_d_keys();\n\n std::size_t temp_storage_bytes{};\n cub::DeviceSegmentedSort::SortKeys(nullptr,\n temp_storage_bytes,\n d_input,\n d_keys_output,\n input.get_num_items(),\n input.get_num_segments(),\n input.get_d_offsets(),\n input.get_d_offsets() + 1);\n\n thrust::device_vector temp_storage(temp_storage_bytes);\n std::uint8_t *d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n LauncherKernel<<<1, 1>>>(\n d_temp_storage,\n temp_storage_bytes,\n d_input,\n d_keys_output,\n input.get_num_items(),\n input.get_num_segments(),\n input.get_d_offsets());\n CubDebugExit(cudaDeviceSynchronize());\n CubDebugExit(cudaPeekAtLastError());\n\n HostReferenceSort(false,\n false,\n input.get_num_segments(),\n h_offsets,\n h_keys,\n h_values);\n\n h_keys_output = keys_output;\n\n const bool keys_ok =\n compare_two_outputs(h_offsets, h_keys, h_keys_output);\n\n AssertTrue(keys_ok);\n\n input.shuffle();\n }\n}\n\ntemplate \nvoid TestDeviceSideLaunch(int min_segments, int max_segments)\n{\n const int max_items = 10000000;\n\n for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++)\n {\n Input edge_cases =\n GenRandomInput(max_items,\n min_segments,\n max_segments,\n descending);\n\n TestDeviceSideLaunch(edge_cases);\n }\n}\n\ntemplate \nvoid TestDeviceSideLaunch()\n{\n TestDeviceSideLaunch(1 << 2, 1 << 8);\n TestDeviceSideLaunch(1 << 9, 1 << 19);\n}\n\n#endif // TEST_CDP\n\nvoid TestUnspecifiedRanges()\n{\n const std::size_t num_items = 1024 * 1024;\n const std::size_t max_segments = 42;\n const std::size_t avg_segment_size = num_items / max_segments;\n\n for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++)\n {\n thrust::host_vector h_offsets_begin;\n thrust::host_vector h_offsets_end;\n\n h_offsets_begin.reserve(max_segments + 1);\n h_offsets_end.reserve(max_segments + 1);\n\n {\n int offset = 0;\n\n for (std::size_t sid = 0; sid < max_segments; sid++)\n {\n const int segment_size =\n RandomValue(static_cast(avg_segment_size));\n const bool segment_is_utilized = RandomValue(100) > 60;\n\n if (segment_is_utilized)\n {\n h_offsets_begin.push_back(offset);\n h_offsets_end.push_back(offset + segment_size);\n }\n\n offset += segment_size;\n }\n\n if (h_offsets_begin.empty())\n {\n h_offsets_begin.push_back(avg_segment_size);\n h_offsets_end.push_back(num_items);\n }\n }\n\n thrust::device_vector keys(num_items);\n thrust::device_vector values(num_items);\n\n thrust::sequence(keys.rbegin(), keys.rend());\n thrust::sequence(values.rbegin(), values.rend());\n\n thrust::device_vector d_offsets_begin = h_offsets_begin;\n thrust::device_vector d_offsets_end = h_offsets_end;\n\n thrust::device_vector expected_keys = keys;\n thrust::device_vector expected_values = values;\n\n const int num_segments = static_cast(h_offsets_begin.size());\n\n for (int sid = 0; sid < num_segments; sid++)\n {\n const int segment_begin = h_offsets_begin[sid];\n const int segment_end = h_offsets_end[sid];\n\n thrust::sort_by_key(expected_keys.begin() + segment_begin,\n expected_keys.begin() + segment_end,\n expected_values.begin() + segment_begin);\n }\n\n thrust::device_vector result_keys = keys;\n thrust::device_vector result_values = values;\n\n {\n cub::DoubleBuffer keys_buffer(\n thrust::raw_pointer_cast(keys.data()), \n thrust::raw_pointer_cast(result_keys.data()));\n\n cub::DoubleBuffer values_buffer(\n thrust::raw_pointer_cast(values.data()), \n thrust::raw_pointer_cast(result_values.data()));\n\n std::size_t temp_storage_bytes{};\n std::uint8_t *d_temp_storage{};\n\n CubDebugExit(cub::DeviceSegmentedSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n keys_buffer, values_buffer, \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data())));\n\n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n CubDebugExit(cub::DeviceSegmentedSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n keys_buffer, values_buffer, \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data())));\n\n for (int sid = 0; sid < num_segments; sid++)\n {\n const int segment_begin = h_offsets_begin[sid];\n const int segment_end = h_offsets_end[sid];\n\n if (keys_buffer.selector == 0)\n {\n thrust::copy(\n keys.begin() + segment_begin,\n keys.begin() + segment_end,\n result_keys.begin() + segment_begin);\n }\n \n if (values_buffer.selector == 0)\n {\n thrust::copy(\n values.begin() + segment_begin,\n values.begin() + segment_end,\n result_values.begin() + segment_begin);\n }\n }\n }\n\n AssertEquals(result_values, expected_values);\n AssertEquals(result_keys, expected_keys);\n\n thrust::sequence(keys.rbegin(), keys.rend());\n thrust::sequence(values.rbegin(), values.rend());\n\n result_keys = keys;\n result_values = values;\n\n {\n std::size_t temp_storage_bytes{};\n std::uint8_t *d_temp_storage{};\n\n CubDebugExit(cub::DeviceSegmentedSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n thrust::raw_pointer_cast(keys.data()), \n thrust::raw_pointer_cast(result_keys.data()), \n thrust::raw_pointer_cast(values.data()), \n thrust::raw_pointer_cast(result_values.data()), \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data())));\n\n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n CubDebugExit(cub::DeviceSegmentedSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n thrust::raw_pointer_cast(keys.data()), \n thrust::raw_pointer_cast(result_keys.data()), \n thrust::raw_pointer_cast(values.data()), \n thrust::raw_pointer_cast(result_values.data()), \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data())));\n }\n\n AssertEquals(result_values, expected_values);\n AssertEquals(result_keys, expected_keys);\n }\n}\n\nint main(int argc, char** argv)\n{\n CommandLineArgs args(argc, argv);\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // %PARAM% TEST_CDP cdp 0:1\n\n#if TEST_CDP == 0\n TestZeroSegments();\n TestEmptySegments(1 << 2);\n TestEmptySegments(1 << 22);\n\n#if TEST_HALF_T\n Test();\n#endif\n\n#if TEST_BF_T\n Test();\n#endif\n\n Test();\n Test();\n Test();\n\n#elif TEST_CDP == 1\n TestDeviceSideLaunch();\n#endif // TEST_CDP\n\n TestUnspecifiedRanges();\n\n return 0;\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n\n#include \n\n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#define TEST_HALF_T !_NVHPC_CUDA\n\n#define TEST_BF_T !_NVHPC_CUDA\n\n#if TEST_HALF_T\n#include \n#endif\n\n#if TEST_BF_T\n#include \n#endif\n\nusing namespace hipcub;\n\ntemplate \nstruct UnwrapHalfAndBfloat16\n{\n using Type = T;\n};\n\n#if TEST_HALF_T\ntemplate <>\nstruct UnwrapHalfAndBfloat16\n{\n using Type = __half;\n};\n#endif\n\n#if TEST_BF_T\ntemplate <>\nstruct UnwrapHalfAndBfloat16\n{\n using Type = __hip_bfloat16;\n};\n#endif\n\nconstexpr static int MAX_ITERATIONS = 2;\n\n\nclass SizeGroupDescription\n{\npublic:\n SizeGroupDescription(const int segments,\n const int segment_size)\n : segments(segments)\n , segment_size(segment_size)\n {}\n\n int segments {};\n int segment_size {};\n};\n\ntemplate \nstruct SegmentChecker\n{\n const KeyT *sorted_keys {};\n const int *offsets {};\n\n SegmentChecker(const KeyT *sorted_keys,\n const int *offsets)\n : sorted_keys(sorted_keys)\n , offsets(offsets)\n {}\n\n bool operator()(int segment_id)\n {\n const int segment_begin = offsets[segment_id];\n const int segment_end = offsets[segment_id + 1];\n\n int counter = 0;\n for (int i = segment_begin; i < segment_end; i++)\n {\n if (sorted_keys[i] != static_cast(counter++))\n {\n return false;\n }\n }\n\n return true;\n }\n};\n\ntemplate \nstruct DescendingSegmentChecker\n{\n const KeyT *sorted_keys{};\n const int *offsets{};\n\n DescendingSegmentChecker(const KeyT *sorted_keys,\n const int *offsets)\n : sorted_keys(sorted_keys)\n , offsets(offsets)\n {}\n\n bool operator()(int segment_id)\n {\n const int segment_begin = offsets[segment_id];\n const int segment_end = offsets[segment_id + 1];\n\n int counter = 0;\n for (int i = segment_end - 1; i >= segment_begin; i--)\n {\n if (sorted_keys[i] != static_cast(counter++))\n {\n return false;\n }\n }\n\n return true;\n }\n};\n\ntemplate \nstruct ReversedIota\n{\n KeyT *data {};\n const int *offsets {};\n\n ReversedIota(KeyT *data,\n const int *offsets)\n : data(data)\n , offsets(offsets)\n {}\n\n void operator()(int segment_id) const\n {\n const int segment_begin = offsets[segment_id];\n const int segment_end = offsets[segment_id + 1];\n const int segment_size = segment_end - segment_begin;\n\n int count = 0;\n for (int i = segment_begin; i < segment_end; i++)\n {\n data[i] = static_cast(segment_size - 1 - count++);\n }\n }\n};\n\n\ntemplate \nstruct Iota\n{\n KeyT *data{};\n const int *offsets{};\n\n Iota(KeyT *data, const int *offsets)\n : data(data)\n , offsets(offsets)\n {}\n\n void operator()(int segment_id) const\n {\n const int segment_begin = offsets[segment_id];\n const int segment_end = offsets[segment_id + 1];\n\n int count = 0;\n for (int i = segment_begin; i < segment_end; i++)\n {\n data[i] = static_cast(count++);\n }\n }\n};\n\n\ntemplate \nclass Input\n{\n thrust::default_random_engine random_engine;\n thrust::device_vector d_segment_sizes;\n thrust::device_vector d_offsets;\n thrust::host_vector h_offsets;\n\n using MaskedValueT = hipcub::detail::conditional_t<\n std::is_same::value, KeyT, ValueT>;\n\n bool reverse {};\n int num_items {};\n thrust::device_vector d_keys;\n thrust::device_vector d_values;\n thrust::host_vector h_keys;\n thrust::host_vector h_values;\n\npublic:\n Input(bool reverse, const thrust::host_vector &h_segment_sizes)\n : d_segment_sizes(h_segment_sizes)\n , d_offsets(d_segment_sizes.size() + 1)\n , h_offsets(d_segment_sizes.size() + 1)\n , reverse(reverse)\n , num_items(static_cast(\n thrust::reduce(d_segment_sizes.begin(), d_segment_sizes.end())))\n , d_keys(num_items)\n , d_values(num_items)\n , h_keys(num_items)\n , h_values(num_items)\n {\n update();\n }\n\n Input(thrust::host_vector &h_offsets)\n : d_offsets(h_offsets)\n , h_offsets(h_offsets)\n , reverse(false)\n , num_items(h_offsets.back())\n , d_keys(num_items)\n , d_values(num_items)\n {\n }\n\n void shuffle()\n {\n thrust::shuffle(d_segment_sizes.begin(), d_segment_sizes.end(), random_engine);\n\n update();\n }\n\n int get_num_items() const\n {\n return num_items;\n }\n\n int get_num_segments() const\n {\n return static_cast(d_segment_sizes.size());\n }\n\n const KeyT *get_d_keys() const\n {\n return thrust::raw_pointer_cast(d_keys.data());\n }\n\n thrust::device_vector &get_d_keys_vec()\n {\n return d_keys;\n }\n\n thrust::device_vector &get_d_values_vec()\n {\n return d_values;\n }\n\n KeyT *get_d_keys()\n {\n return thrust::raw_pointer_cast(d_keys.data());\n }\n\n const thrust::host_vector& get_h_offsets()\n {\n return h_offsets;\n }\n\n MaskedValueT *get_d_values()\n {\n return thrust::raw_pointer_cast(d_values.data());\n }\n\n const int *get_d_offsets() const\n {\n return thrust::raw_pointer_cast(d_offsets.data());\n }\n\n template \n bool check_output_implementation(const T *keys_output)\n {\n const int *offsets = thrust::raw_pointer_cast(h_offsets.data());\n\n if (reverse)\n {\n DescendingSegmentChecker checker{keys_output, offsets};\n\n for (int i = 0; i < get_num_segments(); i++)\n {\n if (!checker(i))\n {\n return false;\n }\n }\n }\n else\n {\n SegmentChecker checker{keys_output, offsets};\n\n for (int i = 0; i < get_num_segments(); i++)\n {\n if (!checker(i))\n {\n return false;\n }\n }\n }\n\n return true;\n }\n\n bool check_output(const KeyT *d_keys_output,\n const MaskedValueT *d_values_output = nullptr)\n {\n KeyT *keys_output = thrust::raw_pointer_cast(h_keys.data());\n MaskedValueT *values_output = thrust::raw_pointer_cast(h_values.data());\n\n hipMemcpy(keys_output,\n d_keys_output,\n sizeof(KeyT) * num_items,\n hipMemcpyDeviceToHost);\n\n const bool keys_ok = check_output_implementation(keys_output);\n\n if (std::is_same::value || d_values_output == nullptr)\n {\n return keys_ok;\n }\n\n hipMemcpy(values_output,\n d_values_output,\n sizeof(ValueT) * num_items,\n hipMemcpyDeviceToHost);\n\n const bool values_ok = check_output_implementation(values_output);\n\n return keys_ok && values_ok;\n }\n\nprivate:\n void update()\n {\n fill_offsets();\n gen_keys();\n }\n\n void fill_offsets()\n {\n thrust::copy(d_segment_sizes.begin(), d_segment_sizes.end(), d_offsets.begin());\n thrust::exclusive_scan(d_offsets.begin(), d_offsets.end(), d_offsets.begin(), 0u);\n thrust::copy(d_offsets.begin(), d_offsets.end(), h_offsets.begin());\n }\n\n void gen_keys()\n {\n KeyT *keys_output = thrust::raw_pointer_cast(h_keys.data());\n const int *offsets = thrust::raw_pointer_cast(h_offsets.data());\n\n if (reverse)\n {\n Iota generator{keys_output, offsets};\n\n for (int i = 0; i < get_num_segments(); i++)\n {\n generator(i);\n }\n }\n else\n {\n ReversedIota generator{keys_output, offsets};\n\n for (int i = 0; i < get_num_segments(); i++)\n {\n generator(i);\n }\n }\n\n d_keys = h_keys;\n d_values = d_keys;\n }\n};\n\ntemplate ::value>\nclass InputDescription\n{\n thrust::host_vector segment_sizes;\n\npublic:\n InputDescription& add(const SizeGroupDescription &group)\n {\n if (static_cast(group.segment_size) <\n static_cast((std::numeric_limits::max)()))\n {\n for (int i = 0; i < group.segments; i++)\n {\n segment_sizes.push_back(group.segment_size);\n }\n }\n\n return *this;\n }\n\n template \n Input gen(bool reverse)\n {\n return Input(reverse, segment_sizes);\n }\n};\n\ntemplate \nclass InputDescription\n{\n thrust::host_vector segment_sizes;\n\npublic:\n InputDescription& add(const SizeGroupDescription &group)\n {\n for (int i = 0; i < group.segments; i++)\n {\n segment_sizes.push_back(group.segment_size);\n }\n\n return *this;\n }\n\n template \n Input gen(bool reverse)\n {\n return Input(reverse, segment_sizes);\n }\n};\n\n\ntemplate \nvoid Sort(bool pairs,\n bool descending,\n bool double_buffer,\n bool stable_sort,\n\n void *tmp_storage,\n std::size_t &temp_storage_bytes,\n\n WrappedKeyT *wrapped_input_keys,\n WrappedKeyT *wrapped_output_keys,\n\n ValueT *input_values,\n ValueT *output_values,\n\n int num_items,\n int num_segments,\n const int *d_offsets,\n\n int *keys_selector = nullptr,\n int *values_selector = nullptr)\n{\n using KeyT = typename UnwrapHalfAndBfloat16::Type;\n\n auto input_keys = reinterpret_cast(wrapped_input_keys);\n auto output_keys = reinterpret_cast(wrapped_output_keys);\n\n if (stable_sort)\n {\n if (pairs)\n {\n if (descending)\n {\n if (double_buffer)\n {\n hipcub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n hipcub::DoubleBuffer values_buffer(input_values, output_values);\n values_buffer.selector = *values_selector;\n\n CubDebugExit(hipcub::DeviceSegmentedSort::StableSortPairsDescending(\n tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n values_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n *values_selector = values_buffer.selector;\n }\n else\n {\n CubDebugExit(hipcub::DeviceSegmentedSort::StableSortPairsDescending(\n tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n input_values,\n output_values,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n else\n {\n if (double_buffer)\n {\n hipcub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n hipcub::DoubleBuffer values_buffer(input_values, output_values);\n values_buffer.selector = *values_selector;\n\n CubDebugExit(\n hipcub::DeviceSegmentedSort::StableSortPairs(tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n values_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n *values_selector = values_buffer.selector;\n }\n else\n {\n CubDebugExit(\n hipcub::DeviceSegmentedSort::StableSortPairs(tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n input_values,\n output_values,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n }\n else\n {\n if (descending)\n {\n if (double_buffer)\n {\n hipcub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n CubDebugExit(hipcub::DeviceSegmentedSort::StableSortKeysDescending(\n tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n }\n else\n {\n CubDebugExit(hipcub::DeviceSegmentedSort::StableSortKeysDescending(\n tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n else\n {\n if (double_buffer)\n {\n hipcub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n CubDebugExit(\n hipcub::DeviceSegmentedSort::StableSortKeys(tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n }\n else\n {\n CubDebugExit(\n hipcub::DeviceSegmentedSort::StableSortKeys(tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n }\n }\n else\n {\n if (pairs)\n {\n if (descending)\n {\n if (double_buffer)\n {\n hipcub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n hipcub::DoubleBuffer values_buffer(input_values, output_values);\n values_buffer.selector = *values_selector;\n\n CubDebugExit(\n hipcub::DeviceSegmentedSort::SortPairsDescending(tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n values_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n *values_selector = values_buffer.selector;\n }\n else\n {\n CubDebugExit(\n hipcub::DeviceSegmentedSort::SortPairsDescending(tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n input_values,\n output_values,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n else\n {\n if (double_buffer)\n {\n hipcub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n hipcub::DoubleBuffer values_buffer(input_values, output_values);\n values_buffer.selector = *values_selector;\n\n CubDebugExit(hipcub::DeviceSegmentedSort::SortPairs(tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n values_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n *values_selector = values_buffer.selector;\n }\n else\n {\n CubDebugExit(hipcub::DeviceSegmentedSort::SortPairs(tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n input_values,\n output_values,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n }\n else\n {\n if (descending)\n {\n if (double_buffer)\n {\n hipcub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n CubDebugExit(\n hipcub::DeviceSegmentedSort::SortKeysDescending(tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n }\n else\n {\n CubDebugExit(\n hipcub::DeviceSegmentedSort::SortKeysDescending(tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n else\n {\n if (double_buffer)\n {\n hipcub::DoubleBuffer keys_buffer(input_keys, output_keys);\n keys_buffer.selector = *keys_selector;\n\n CubDebugExit(hipcub::DeviceSegmentedSort::SortKeys(tmp_storage,\n temp_storage_bytes,\n keys_buffer,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n\n *keys_selector = keys_buffer.selector;\n }\n else\n {\n CubDebugExit(hipcub::DeviceSegmentedSort::SortKeys(tmp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n num_items,\n num_segments,\n d_offsets,\n d_offsets + 1));\n }\n }\n }\n }\n}\n\ntemplate \nstd::size_t Sort(bool pairs,\n bool descending,\n bool double_buffer,\n bool stable_sort,\n\n KeyT *input_keys,\n KeyT *output_keys,\n\n ValueT *input_values,\n ValueT *output_values,\n\n int num_items,\n int num_segments,\n const int *d_offsets,\n\n int *keys_selector = nullptr,\n int *values_selector = nullptr)\n{\n std::size_t temp_storage_bytes = 42ul;\n\n Sort(pairs,\n descending,\n double_buffer,\n stable_sort,\n nullptr,\n temp_storage_bytes,\n input_keys,\n output_keys,\n input_values,\n output_values,\n num_items,\n num_segments,\n d_offsets,\n keys_selector,\n values_selector);\n\n thrust::device_vector temp_storage(temp_storage_bytes);\n std::uint8_t *d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n Sort(pairs,\n descending,\n double_buffer,\n stable_sort,\n d_temp_storage,\n temp_storage_bytes,\n input_keys,\n output_keys,\n input_values,\n output_values,\n num_items,\n num_segments,\n d_offsets,\n keys_selector,\n values_selector);\n\n return temp_storage_bytes;\n}\n\n\nconstexpr bool keys_only = false;\nconstexpr bool pairs = true;\n\nconstexpr bool ascending = false;\nconstexpr bool descending = true;\n\nconstexpr bool pointers = false;\nconstexpr bool double_buffer = true;\n\nconstexpr bool unstable = false;\nconstexpr bool stable = true;\n\n\nvoid TestZeroSegments()\n{\n // Type doesn't affect the escape logic, so it should be fine\n // to test only one set of types here.\n\n using KeyT = std::uint8_t;\n using ValueT = std::uint64_t;\n\n for (bool stable_sort: { unstable, stable })\n {\n for (bool sort_pairs: { keys_only, pairs })\n {\n for (bool sort_descending: { ascending, descending })\n {\n for (bool sort_buffer: { pointers, double_buffer })\n {\n hipcub::DoubleBuffer keys_buffer(nullptr, nullptr);\n hipcub::DoubleBuffer values_buffer(nullptr, nullptr);\n values_buffer.selector = 1;\n\n Sort(sort_pairs,\n sort_descending,\n sort_buffer,\n stable_sort,\n nullptr,\n nullptr,\n nullptr,\n nullptr,\n int{},\n int{},\n nullptr,\n &keys_buffer.selector,\n &values_buffer.selector);\n\n AssertEquals(keys_buffer.selector, 0);\n AssertEquals(values_buffer.selector, 1);\n }\n }\n }\n }\n}\n\n\nvoid TestEmptySegments(int segments)\n{\n // Type doesn't affect the escape logic, so it should be fine\n // to test only one set of types here.\n\n using KeyT = std::uint8_t;\n using ValueT = std::uint64_t;\n\n thrust::device_vector offsets(segments + 1, int{});\n const int *d_offsets = thrust::raw_pointer_cast(offsets.data());\n\n for (bool sort_stable: { unstable, stable })\n {\n for (bool sort_pairs: { keys_only, pairs })\n {\n for (bool sort_descending: { ascending, descending })\n {\n for (bool sort_buffer: { pointers, double_buffer })\n {\n hipcub::DoubleBuffer keys_buffer(nullptr, nullptr);\n hipcub::DoubleBuffer values_buffer(nullptr, nullptr);\n values_buffer.selector = 1;\n\n Sort(sort_pairs,\n sort_descending,\n sort_buffer,\n sort_stable,\n nullptr,\n nullptr,\n nullptr,\n nullptr,\n int{},\n segments,\n d_offsets,\n &keys_buffer.selector,\n &values_buffer.selector);\n\n AssertEquals(keys_buffer.selector, 0);\n AssertEquals(values_buffer.selector, 1);\n }\n }\n }\n }\n}\n\n\ntemplate \nvoid TestSameSizeSegments(int segment_size,\n int segments,\n bool skip_values = false)\n{\n const int num_items = segment_size * segments;\n\n thrust::device_vector offsets(segments + 1);\n thrust::sequence(offsets.begin(),\n offsets.end(),\n int{},\n segment_size);\n\n const int *d_offsets = thrust::raw_pointer_cast(offsets.data());\n\n const KeyT target_key {1};\n const ValueT target_value {42};\n\n thrust::device_vector keys_input(num_items);\n thrust::device_vector keys_output(num_items);\n\n KeyT *d_keys_input = thrust::raw_pointer_cast(keys_input.data());\n KeyT *d_keys_output = thrust::raw_pointer_cast(keys_output.data());\n\n thrust::device_vector values_input(num_items);\n thrust::device_vector values_output(num_items);\n\n thrust::host_vector host_keys(num_items);\n thrust::host_vector host_values(num_items);\n\n ValueT *d_values_input = thrust::raw_pointer_cast(values_input.data());\n ValueT *d_values_output = thrust::raw_pointer_cast(values_output.data());\n\n for (bool stable_sort: { unstable, stable })\n {\n for (bool sort_pairs: { keys_only, pairs })\n {\n if (sort_pairs)\n {\n if (skip_values)\n {\n continue;\n }\n }\n\n for (bool sort_descending: { ascending, descending })\n {\n for (bool sort_buffers: { pointers, double_buffer })\n {\n hipcub::DoubleBuffer keys_buffer(nullptr, nullptr);\n hipcub::DoubleBuffer values_buffer(nullptr, nullptr);\n values_buffer.selector = 1;\n\n thrust::fill(keys_input.begin(), keys_input.end(), target_key);\n thrust::fill(keys_output.begin(), keys_output.end(), KeyT{});\n\n if (sort_pairs)\n {\n if (sort_buffers)\n {\n thrust::fill(values_input.begin(), values_input.end(), ValueT{});\n thrust::fill(values_output.begin(), values_output.end(), target_value);\n }\n else\n {\n thrust::fill(values_input.begin(), values_input.end(), target_value);\n thrust::fill(values_output.begin(), values_output.end(), ValueT{});\n }\n }\n\n const std::size_t temp_storage_bytes =\n Sort(sort_pairs,\n sort_descending,\n sort_buffers,\n stable_sort,\n d_keys_input,\n d_keys_output,\n d_values_input,\n d_values_output,\n num_items,\n segments,\n d_offsets,\n &keys_buffer.selector,\n &values_buffer.selector);\n\n // If temporary storage size is defined by extra keys storage\n if (sort_buffers)\n {\n if (2 * segments * sizeof(unsigned int) < num_items * sizeof(KeyT))\n {\n std::size_t extra_temp_storage_bytes{};\n\n Sort(sort_pairs,\n sort_descending,\n pointers,\n stable_sort,\n nullptr,\n extra_temp_storage_bytes,\n d_keys_input,\n d_keys_output,\n d_values_input,\n d_values_output,\n num_items,\n segments,\n d_offsets,\n &keys_buffer.selector,\n &values_buffer.selector);\n\n AssertTrue(extra_temp_storage_bytes > temp_storage_bytes);\n }\n }\n\n {\n host_keys = keys_buffer.selector || !sort_buffers ? keys_output\n : keys_input;\n const std::size_t items_selected =\n thrust::count(host_keys.begin(), host_keys.end(), target_key);\n AssertEquals(static_cast(items_selected), num_items);\n }\n\n if (sort_pairs)\n {\n host_values = values_buffer.selector || !sort_buffers\n ? values_output\n : values_input;\n const std::size_t items_selected =\n thrust::count(host_values.begin(),\n host_values.end(),\n target_value);\n\n AssertEquals(static_cast(items_selected), num_items);\n }\n }\n }\n }\n }\n}\n\n\ntemplate \nvoid InputTest(bool sort_descending,\n Input &input)\n{\n thrust::device_vector keys_output(input.get_num_items());\n KeyT *d_keys_output = thrust::raw_pointer_cast(keys_output.data());\n\n thrust::device_vector values_output(input.get_num_items());\n ValueT *d_values_output = thrust::raw_pointer_cast(values_output.data());\n\n for (bool stable_sort: { unstable, stable })\n {\n for (bool sort_pairs : { keys_only, pairs })\n {\n for (bool sort_buffers : {pointers, double_buffer})\n {\n for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++)\n {\n thrust::fill(keys_output.begin(), keys_output.end(), KeyT{});\n thrust::fill(values_output.begin(), values_output.end(), ValueT{});\n\n hipcub::DoubleBuffer keys_buffer(input.get_d_keys(),\n d_keys_output);\n hipcub::DoubleBuffer values_buffer(input.get_d_values(),\n d_values_output);\n\n Sort(sort_pairs,\n sort_descending,\n sort_buffers,\n stable_sort,\n input.get_d_keys(),\n d_keys_output,\n input.get_d_values(),\n d_values_output,\n input.get_num_items(),\n input.get_num_segments(),\n input.get_d_offsets(),\n &keys_buffer.selector,\n &values_buffer.selector);\n\n if (sort_buffers)\n {\n if (sort_pairs)\n {\n AssertTrue(input.check_output(keys_buffer.Current(),\n values_buffer.Current()));\n }\n else\n {\n AssertTrue(input.check_output(keys_buffer.Current()));\n }\n }\n else\n {\n if (sort_pairs)\n {\n AssertTrue(input.check_output(d_keys_output, d_values_output));\n }\n else\n {\n AssertTrue(input.check_output(d_keys_output));\n }\n }\n\n input.shuffle();\n }\n }\n }\n }\n}\n\nstruct ComparisonPredicate\n{\n template \n __host__ __device__ bool operator()(const T &lhs, const T &rhs) const\n {\n return lhs == rhs;\n }\n\n __host__ __device__ bool operator()(const half_t &lhs, const half_t &rhs) const\n {\n return lhs.raw() == rhs.raw();\n }\n};\n\ntemplate \nbool compare_two_outputs(const thrust::host_vector &offsets,\n const thrust::host_vector &lhs,\n const thrust::host_vector &rhs)\n{\n const auto num_segments = static_cast(offsets.size() - 1);\n\n for (std::size_t segment_id = 0; segment_id < num_segments; segment_id++)\n {\n auto lhs_begin = lhs.cbegin() + offsets[segment_id];\n auto lhs_end = lhs.cbegin() + offsets[segment_id + 1];\n auto rhs_begin = rhs.cbegin() + offsets[segment_id];\n\n auto err = thrust::mismatch(lhs_begin, lhs_end, rhs_begin, ComparisonPredicate{});\n\n if (err.first != lhs_end)\n {\n const auto idx = thrust::distance(lhs_begin, err.first);\n const auto segment_size = std::distance(lhs_begin, lhs_end);\n\n std::cerr << \"Mismatch in segment \" << segment_id\n << \" at position \" << idx << \" / \" << segment_size\n << \": \"\n << static_cast(lhs_begin[idx]) << \" vs \"\n << static_cast(rhs_begin[idx]) << \" (\"\n << typeid(lhs_begin[idx]).name() << \")\" << std::endl;\n\n return false;\n }\n }\n\n return true;\n}\n\ntemplate \nvoid RandomizeInput(thrust::host_vector &h_keys,\n thrust::host_vector &h_values)\n{\n for (std::size_t i = 0; i < h_keys.size(); i++)\n {\n h_keys[i] = RandomValue((std::numeric_limits::max)()) > 128;\n h_values[i] = RandomValue((std::numeric_limits::max)());\n }\n}\n\ntemplate \nvoid RandomizeInput(thrust::host_vector &h_keys,\n thrust::host_vector &h_values)\n{\n for (std::size_t i = 0; i < h_keys.size(); i++)\n {\n h_keys[i] = RandomValue((std::numeric_limits::max)());\n h_values[i] = RandomValue((std::numeric_limits::max)());\n }\n}\n\n#if TEST_HALF_T\nvoid RandomizeInput(thrust::host_vector &h_keys,\n thrust::host_vector &h_values)\n{\n for (std::size_t i = 0; i < h_keys.size(); i++)\n {\n h_keys[i] = RandomValue((std::numeric_limits::max)());\n h_values[i] = RandomValue((std::numeric_limits::max)());\n }\n}\n#endif\n\n#if TEST_BF_T\nvoid RandomizeInput(thrust::host_vector &h_keys,\n thrust::host_vector &h_values)\n{\n for (std::size_t i = 0; i < h_keys.size(); i++)\n {\n h_keys[i] = RandomValue((std::numeric_limits::max)());\n h_values[i] = RandomValue((std::numeric_limits::max)());\n }\n}\n#endif\n\n\n\ntemplate \nvoid HostReferenceSort(bool sort_pairs,\n bool sort_descending,\n unsigned int num_segments,\n const thrust::host_vector &h_offsets,\n thrust::host_vector &h_keys,\n thrust::host_vector &h_values)\n{\n for (unsigned int segment_i = 0;\n segment_i < num_segments;\n segment_i++)\n {\n const int segment_begin = h_offsets[segment_i];\n const int segment_end = h_offsets[segment_i + 1];\n\n if (sort_pairs)\n {\n if (sort_descending)\n {\n thrust::stable_sort_by_key(h_keys.begin() + segment_begin,\n h_keys.begin() + segment_end,\n h_values.begin() + segment_begin,\n thrust::greater{});\n }\n else\n {\n thrust::stable_sort_by_key(h_keys.begin() + segment_begin,\n h_keys.begin() + segment_end,\n h_values.begin() + segment_begin);\n }\n }\n else\n {\n if (sort_descending)\n {\n thrust::stable_sort(h_keys.begin() + segment_begin,\n h_keys.begin() + segment_end,\n thrust::greater{});\n }\n else\n {\n thrust::stable_sort(h_keys.begin() + segment_begin,\n h_keys.begin() + segment_end);\n }\n }\n }\n}\n\n\n#if STORE_ON_FAILURE\ntemplate \nvoid DumpInput(bool sort_pairs,\n bool sort_descending,\n bool sort_buffers,\n Input &input,\n thrust::host_vector &h_keys,\n thrust::host_vector &h_values)\n{\n const thrust::host_vector &h_offsets = input.get_h_offsets();\n\n std::cout << \"sort pairs: \" << sort_pairs << \"\\n\";\n std::cout << \"sort descending: \" << sort_descending << \"\\n\";\n std::cout << \"sort buffers: \" << sort_buffers << \"\\n\";\n std::cout << \"num_items: \" << input.get_num_items() << \"\\n\";\n std::cout << \"num_segments: \" << input.get_num_segments() << \"\\n\";\n std::cout << \"key type: \" << typeid(h_keys[0]).name() << \"\\n\";\n std::cout << \"value type: \" << typeid(h_values[0]).name() << \"\\n\";\n std::cout << \"offset type: \" << typeid(h_offsets[0]).name() << \"\\n\";\n\n std::ofstream offsets_dump(\"offsets\", std::ios::binary);\n offsets_dump.write(reinterpret_cast(\n thrust::raw_pointer_cast(h_offsets.data())),\n sizeof(int) * h_offsets.size());\n\n std::ofstream keys_dump(\"keys\", std::ios::binary);\n keys_dump.write(reinterpret_cast(\n thrust::raw_pointer_cast(h_keys.data())),\n sizeof(KeyT) * h_keys.size());\n\n std::ofstream values_dump(\"values\", std::ios::binary);\n values_dump.write(reinterpret_cast(\n thrust::raw_pointer_cast(h_values.data())),\n sizeof(ValueT) * h_values.size());\n}\n#endif\n\n\ntemplate \nvoid InputTestRandom(Input &input)\n{\n thrust::host_vector h_keys_output(input.get_num_items());\n thrust::device_vector keys_output(input.get_num_items());\n\n thrust::host_vector h_values_output(input.get_num_items());\n thrust::device_vector values_output(input.get_num_items());\n\n KeyT *d_keys_output = thrust::raw_pointer_cast(keys_output.data());\n ValueT *d_values_output = thrust::raw_pointer_cast(values_output.data());\n\n thrust::host_vector h_keys(input.get_num_items());\n thrust::host_vector h_values(input.get_num_items());\n\n const thrust::host_vector &h_offsets = input.get_h_offsets();\n\n for (bool stable_sort: { unstable, stable })\n {\n for (bool sort_pairs: { keys_only, pairs })\n {\n for (bool sort_descending: { ascending, descending })\n {\n for (bool sort_buffers: { pointers, double_buffer })\n {\n for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++)\n {\n RandomizeInput(h_keys, h_values);\n\n#if STORE_ON_FAILURE\n auto h_keys_backup = h_keys;\n auto h_values_backup = h_values;\n#endif\n\n input.get_d_keys_vec() = h_keys;\n input.get_d_values_vec() = h_values;\n\n hipcub::DoubleBuffer keys_buffer(input.get_d_keys(), d_keys_output);\n hipcub::DoubleBuffer values_buffer(input.get_d_values(), d_values_output);\n\n Sort(sort_pairs,\n sort_descending,\n sort_buffers,\n stable_sort,\n input.get_d_keys(),\n d_keys_output,\n input.get_d_values(),\n d_values_output,\n input.get_num_items(),\n input.get_num_segments(),\n input.get_d_offsets(),\n &keys_buffer.selector,\n &values_buffer.selector);\n\n HostReferenceSort(sort_pairs,\n sort_descending,\n input.get_num_segments(),\n h_offsets,\n h_keys,\n h_values);\n\n if (sort_buffers)\n {\n if (keys_buffer.selector)\n {\n h_keys_output = keys_output;\n }\n else\n {\n h_keys_output = input.get_d_keys_vec();\n }\n\n if (values_buffer.selector)\n {\n h_values_output = values_output;\n }\n else\n {\n h_values_output = input.get_d_values_vec();\n }\n }\n else\n {\n h_keys_output = keys_output;\n h_values_output = values_output;\n }\n\n const bool keys_ok =\n compare_two_outputs(h_offsets, h_keys, h_keys_output);\n\n const bool values_ok =\n sort_pairs\n ? compare_two_outputs(h_offsets, h_values, h_values_output)\n : true;\n\n#if STORE_ON_FAILURE\n if (!keys_ok || !values_ok)\n {\n DumpInput(sort_pairs,\n sort_descending,\n sort_buffers,\n input,\n h_keys_backup,\n h_values_backup);\n }\n#endif\n\n AssertTrue(keys_ok);\n AssertTrue(values_ok);\n\n input.shuffle();\n }\n }\n }\n }\n }\n}\n\ntemplate ::value>\nstruct EdgeTestDispatch\n{\n // Edge cases that needs to be tested\n const int empty_short_circuit_segment_size = 0;\n const int copy_short_circuit_segment_size = 1;\n const int swap_short_circuit_segment_size = 2;\n\n const int a_few = 2;\n const int a_bunch_of = 42;\n const int a_lot_of = 420;\n\n template \n HIPCUB_RUNTIME_FUNCTION __forceinline__ hipError_t Invoke()\n {\n NV_IF_TARGET(NV_IS_HOST,\n (using SmallAndMediumPolicyT =\n typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT;\n using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy;\n\n const int small_segment_max_segment_size =\n SmallAndMediumPolicyT::SmallPolicyT::ITEMS_PER_TILE;\n\n const int items_per_small_segment =\n SmallAndMediumPolicyT::SmallPolicyT::ITEMS_PER_THREAD;\n\n const int medium_segment_max_segment_size =\n SmallAndMediumPolicyT::MediumPolicyT::ITEMS_PER_TILE;\n\n const int single_thread_segment_size = items_per_small_segment;\n\n const int large_cached_segment_max_segment_size =\n LargeSegmentPolicyT::BLOCK_THREADS *\n LargeSegmentPolicyT::ITEMS_PER_THREAD;\n\n for (bool sort_descending : {ascending, descending}) {\n Input edge_cases =\n InputDescription()\n .add({a_lot_of, empty_short_circuit_segment_size})\n .add({a_lot_of, copy_short_circuit_segment_size})\n .add({a_lot_of, swap_short_circuit_segment_size})\n .add({a_lot_of, swap_short_circuit_segment_size + 1})\n .add({a_lot_of, swap_short_circuit_segment_size + 1})\n .add({a_lot_of, single_thread_segment_size - 1})\n .add({a_lot_of, single_thread_segment_size})\n .add({a_lot_of, single_thread_segment_size + 1})\n .add({a_lot_of, single_thread_segment_size * 2 - 1})\n .add({a_lot_of, single_thread_segment_size * 2})\n .add({a_lot_of, single_thread_segment_size * 2 + 1})\n .add({a_bunch_of, small_segment_max_segment_size - 1})\n .add({a_bunch_of, small_segment_max_segment_size})\n .add({a_bunch_of, small_segment_max_segment_size + 1})\n .add({a_bunch_of, medium_segment_max_segment_size - 1})\n .add({a_bunch_of, medium_segment_max_segment_size})\n .add({a_bunch_of, medium_segment_max_segment_size + 1})\n .add({a_bunch_of, large_cached_segment_max_segment_size - 1})\n .add({a_bunch_of, large_cached_segment_max_segment_size})\n .add({a_bunch_of, large_cached_segment_max_segment_size + 1})\n .add({a_few, large_cached_segment_max_segment_size * 2})\n .add({a_few, large_cached_segment_max_segment_size * 3})\n .add({a_few, large_cached_segment_max_segment_size * 5})\n .template gen(sort_descending);\n\n InputTest(sort_descending, edge_cases);\n }));\n\n return hipSuccess;\n }\n};\n\ntemplate \nstruct EdgeTestDispatch\n{\n template \n HIPCUB_RUNTIME_FUNCTION __forceinline__ hipError_t Invoke()\n {\n // Edge case test is using an optimized testing approach which is\n // incompatible with duplicates. RandomTest is used for other types.\n return hipSuccess;\n }\n};\n\ntemplate \nvoid EdgePatternsTest()\n{\n int ptx_version = 0;\n if (HipcubDebug(PtxVersion(ptx_version)))\n {\n return;\n }\n\n using MaxPolicyT =\n typename hipcub::DeviceSegmentedSortPolicy::MaxPolicy;\n using EdgeTestDispatchT = EdgeTestDispatch;\n EdgeTestDispatchT dispatch;\n\n MaxPolicyT::Invoke(ptx_version, dispatch);\n\n}\n\ntemplate \nInput GenRandomInput(int max_items,\n int min_segments,\n int max_segments,\n bool descending)\n{\n int items_generated {};\n const int segments_num = RandomValue(max_segments) + min_segments;\n\n thrust::host_vector segment_sizes;\n segment_sizes.reserve(segments_num);\n\n const int max_segment_size = 6000;\n\n for (int segment_id = 0; segment_id < segments_num; segment_id++)\n {\n const int segment_size_raw = RandomValue(max_segment_size);\n const int segment_size = segment_size_raw > 0 ? segment_size_raw : 0;\n\n if (segment_size + items_generated > max_items)\n {\n break;\n }\n\n items_generated += segment_size;\n segment_sizes.push_back(segment_size);\n }\n\n return Input{descending, segment_sizes};\n}\n\ntemplate \nvoid RandomTest(int min_segments,\n int max_segments)\n{\n const int max_items = 10000000;\n\n for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++)\n {\n Input edge_cases = GenRandomInput(max_items,\n min_segments,\n max_segments,\n descending);\n\n InputTestRandom(edge_cases);\n }\n}\n\n\ntemplate \nvoid Test()\n{\n for (int segment_size: { 1, 1024, 24 * 1024 })\n {\n for (int segments: { 1, 1024 })\n {\n TestSameSizeSegments(segment_size, segments);\n }\n }\n\n RandomTest(1 << 2, 1 << 8);\n RandomTest(1 << 9, 1 << 19);\n EdgePatternsTest();\n}\n\n\n#if TEST_CDP == 1\ntemplate \n__global__ void LauncherKernel(\n void *tmp_storage,\n std::size_t temp_storage_bytes,\n const KeyT *in_keys,\n KeyT *out_keys,\n int num_items,\n int num_segments,\n const int *offsets)\n{\n HipcubDebug(hipcub::DeviceSegmentedSort::SortKeys(tmp_storage,\n temp_storage_bytes,\n in_keys,\n out_keys,\n num_items,\n num_segments,\n offsets,\n offsets + 1));\n}\n\ntemplate \nvoid TestDeviceSideLaunch(Input &input)\n{\n thrust::host_vector h_keys_output(input.get_num_items());\n thrust::device_vector keys_output(input.get_num_items());\n\n thrust::host_vector h_values_output(input.get_num_items());\n thrust::device_vector values_output(input.get_num_items());\n\n KeyT *d_keys_output = thrust::raw_pointer_cast(keys_output.data());\n\n thrust::host_vector h_keys(input.get_num_items());\n thrust::host_vector h_values(input.get_num_items());\n\n const thrust::host_vector &h_offsets = input.get_h_offsets();\n\n for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++)\n {\n RandomizeInput(h_keys, h_values);\n\n input.get_d_keys_vec() = h_keys;\n input.get_d_values_vec() = h_values;\n\n const KeyT *d_input = input.get_d_keys();\n\n std::size_t temp_storage_bytes{};\n hipcub::DeviceSegmentedSort::SortKeys(nullptr,\n temp_storage_bytes,\n d_input,\n d_keys_output,\n input.get_num_items(),\n input.get_num_segments(),\n input.get_d_offsets(),\n input.get_d_offsets() + 1);\n\n thrust::device_vector temp_storage(temp_storage_bytes);\n std::uint8_t *d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n LauncherKernel<<<1, 1>>>(\n d_temp_storage,\n temp_storage_bytes,\n d_input,\n d_keys_output,\n input.get_num_items(),\n input.get_num_segments(),\n input.get_d_offsets());\n CubDebugExit(hipDeviceSynchronize());\n CubDebugExit(hipPeekAtLastError());\n\n HostReferenceSort(false,\n false,\n input.get_num_segments(),\n h_offsets,\n h_keys,\n h_values);\n\n h_keys_output = keys_output;\n\n const bool keys_ok =\n compare_two_outputs(h_offsets, h_keys, h_keys_output);\n\n AssertTrue(keys_ok);\n\n input.shuffle();\n }\n}\n\ntemplate \nvoid TestDeviceSideLaunch(int min_segments, int max_segments)\n{\n const int max_items = 10000000;\n\n for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++)\n {\n Input edge_cases =\n GenRandomInput(max_items,\n min_segments,\n max_segments,\n descending);\n\n TestDeviceSideLaunch(edge_cases);\n }\n}\n\ntemplate \nvoid TestDeviceSideLaunch()\n{\n TestDeviceSideLaunch(1 << 2, 1 << 8);\n TestDeviceSideLaunch(1 << 9, 1 << 19);\n}\n\n#endif // TEST_CDP\n\nvoid TestUnspecifiedRanges()\n{\n const std::size_t num_items = 1024 * 1024;\n const std::size_t max_segments = 42;\n const std::size_t avg_segment_size = num_items / max_segments;\n\n for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++)\n {\n thrust::host_vector h_offsets_begin;\n thrust::host_vector h_offsets_end;\n\n h_offsets_begin.reserve(max_segments + 1);\n h_offsets_end.reserve(max_segments + 1);\n\n {\n int offset = 0;\n\n for (std::size_t sid = 0; sid < max_segments; sid++)\n {\n const int segment_size =\n RandomValue(static_cast(avg_segment_size));\n const bool segment_is_utilized = RandomValue(100) > 60;\n\n if (segment_is_utilized)\n {\n h_offsets_begin.push_back(offset);\n h_offsets_end.push_back(offset + segment_size);\n }\n\n offset += segment_size;\n }\n\n if (h_offsets_begin.empty())\n {\n h_offsets_begin.push_back(avg_segment_size);\n h_offsets_end.push_back(num_items);\n }\n }\n\n thrust::device_vector keys(num_items);\n thrust::device_vector values(num_items);\n\n thrust::sequence(keys.rbegin(), keys.rend());\n thrust::sequence(values.rbegin(), values.rend());\n\n thrust::device_vector d_offsets_begin = h_offsets_begin;\n thrust::device_vector d_offsets_end = h_offsets_end;\n\n thrust::device_vector expected_keys = keys;\n thrust::device_vector expected_values = values;\n\n const int num_segments = static_cast(h_offsets_begin.size());\n\n for (int sid = 0; sid < num_segments; sid++)\n {\n const int segment_begin = h_offsets_begin[sid];\n const int segment_end = h_offsets_end[sid];\n\n thrust::sort_by_key(expected_keys.begin() + segment_begin,\n expected_keys.begin() + segment_end,\n expected_values.begin() + segment_begin);\n }\n\n thrust::device_vector result_keys = keys;\n thrust::device_vector result_values = values;\n\n {\n hipcub::DoubleBuffer keys_buffer(\n thrust::raw_pointer_cast(keys.data()), \n thrust::raw_pointer_cast(result_keys.data()));\n\n hipcub::DoubleBuffer values_buffer(\n thrust::raw_pointer_cast(values.data()), \n thrust::raw_pointer_cast(result_values.data()));\n\n std::size_t temp_storage_bytes{};\n std::uint8_t *d_temp_storage{};\n\n CubDebugExit(hipcub::DeviceSegmentedSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n keys_buffer, values_buffer, \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data())));\n\n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n CubDebugExit(hipcub::DeviceSegmentedSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n keys_buffer, values_buffer, \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data())));\n\n for (int sid = 0; sid < num_segments; sid++)\n {\n const int segment_begin = h_offsets_begin[sid];\n const int segment_end = h_offsets_end[sid];\n\n if (keys_buffer.selector == 0)\n {\n thrust::copy(\n keys.begin() + segment_begin,\n keys.begin() + segment_end,\n result_keys.begin() + segment_begin);\n }\n \n if (values_buffer.selector == 0)\n {\n thrust::copy(\n values.begin() + segment_begin,\n values.begin() + segment_end,\n result_values.begin() + segment_begin);\n }\n }\n }\n\n AssertEquals(result_values, expected_values);\n AssertEquals(result_keys, expected_keys);\n\n thrust::sequence(keys.rbegin(), keys.rend());\n thrust::sequence(values.rbegin(), values.rend());\n\n result_keys = keys;\n result_values = values;\n\n {\n std::size_t temp_storage_bytes{};\n std::uint8_t *d_temp_storage{};\n\n CubDebugExit(hipcub::DeviceSegmentedSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n thrust::raw_pointer_cast(keys.data()), \n thrust::raw_pointer_cast(result_keys.data()), \n thrust::raw_pointer_cast(values.data()), \n thrust::raw_pointer_cast(result_values.data()), \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data())));\n\n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n CubDebugExit(hipcub::DeviceSegmentedSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n thrust::raw_pointer_cast(keys.data()), \n thrust::raw_pointer_cast(result_keys.data()), \n thrust::raw_pointer_cast(values.data()), \n thrust::raw_pointer_cast(result_values.data()), \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data())));\n }\n\n AssertEquals(result_values, expected_values);\n AssertEquals(result_keys, expected_keys);\n }\n}\n\nint main(int argc, char** argv)\n{\n CommandLineArgs args(argc, argv);\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // %PARAM% TEST_CDP cdp 0:1\n\n#if TEST_CDP == 0\n TestZeroSegments();\n TestEmptySegments(1 << 2);\n TestEmptySegments(1 << 22);\n\n#if TEST_HALF_T\n Test();\n#endif\n\n#if TEST_BF_T\n Test();\n#endif\n\n Test();\n Test();\n Test();\n\n#elif TEST_CDP == 1\n TestDeviceSideLaunch();\n#endif // TEST_CDP\n\n TestUnspecifiedRanges();\n\n return 0;\n}\n", "source": "/kaggle/working/cub/test/test_device_segmented_sort.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#undef NDEBUG\n#include \n\n#include \"catch2_test_block_radix_sort.cuh\"\n#include \"cub/block/radix_rank_sort_operations.cuh\"\n\n#include \n#include \n#include \n\n// example-begin custom-type\nstruct custom_t\n{\n float f;\n int unused;\n long long int lli;\n\n custom_t() = default;\n __device__ custom_t(float f, long long int lli)\n : f(f)\n , unused(42)\n , lli(lli)\n {}\n};\n\nstatic __device__ bool operator==(const custom_t &lhs, const custom_t &rhs)\n{\n return lhs.f == rhs.f && lhs.lli == rhs.lli;\n}\n\nstruct decomposer_t\n{\n __device__ ::cuda::std::tuple //\n operator()(custom_t &key) const\n {\n return {key.f, key.lli};\n }\n};\n// example-end custom-type\n\n__global__ void sort_keys()\n{\n // example-begin keys\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+2.5, 4}, //\n {-2.5, 0}, //\n {+1.1, 3}, //\n },\n {\n // thread 1 keys\n {+0.0, 1}, //\n {-0.0, 2}, //\n {+3.7, 5} //\n }};\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage).Sort(thread_keys[threadIdx.x], decomposer_t{});\n\n custom_t expected_output[2][3] = //\n {{\n // thread 0 expected keys\n {-2.5, 0}, //\n {+0.0, 1}, //\n {-0.0, 2} //\n },\n {\n // thread 1 expected keys\n {+1.1, 3}, //\n {+2.5, 4}, //\n {+3.7, 5} //\n }};\n // example-end keys\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_output[threadIdx.x][2]);\n}\n\n__global__ void sort_keys_bits()\n{\n // example-begin keys-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 1 key each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][1] = //\n {{\n {24.2, 1ll << 61} // thread 0 keys\n },\n {\n {42.4, 1ll << 60} // thread 1 keys\n }};\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000001110000011001100110011010 00100000000000...0000\n // decompose(in[1]) = 01000010001010011001100110011010 00010000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .Sort(thread_keys[threadIdx.x], decomposer_t{}, begin_bit, end_bit);\n\n custom_t expected_output[2][3] = //\n {{\n {42.4, 1ll << 60}, // thread 0 expected keys\n },\n {\n {24.2, 1ll << 61} // thread 1 expected keys\n }};\n // example-end keys-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n}\n\n__global__ void sort_keys_descending()\n{\n // example-begin keys-descending\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+1.1, 2}, //\n {+2.5, 1}, //\n {-0.0, 4}, //\n },\n {\n // thread 1 keys\n {+0.0, 3}, //\n {-2.5, 5}, //\n {+3.7, 0} //\n }};\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage).SortDescending(thread_keys[threadIdx.x], decomposer_t{});\n\n custom_t expected_output[2][3] = //\n {{\n // thread 0 expected keys\n {+3.7, 0}, //\n {+2.5, 1}, //\n {+1.1, 2}, //\n },\n {\n // thread 1 expected keys\n {-0.0, 4}, //\n {+0.0, 3}, //\n {-2.5, 5} //\n }};\n // example-end keys-descending\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_output[threadIdx.x][2]);\n}\n\n__global__ void sort_keys_descending_bits()\n{\n // example-begin keys-descending-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 1 key each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][1] = //\n {{\n {42.4, 1ll << 60} // thread 0 keys\n },\n {\n {24.2, 1ll << 61} // thread 1 keys\n }};\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000010001010011001100110011010 00010000000000...0000\n // decompose(in[1]) = 01000001110000011001100110011010 00100000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortDescending(thread_keys[threadIdx.x], decomposer_t{}, begin_bit, end_bit);\n\n custom_t expected_output[2][3] = //\n {{\n {24.2, 1ll << 61}, // thread 0 expected keys\n },\n {\n {42.4, 1ll << 60} // thread 1 expected keys\n }};\n // example-end keys-descending-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n}\n\n__global__ void sort_pairs()\n{\n // example-begin pairs\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys and values each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+2.5, 4}, //\n {-2.5, 0}, //\n {+1.1, 3}, //\n },\n {\n // thread 1 keys\n {+0.0, 1}, //\n {-0.0, 2}, //\n {+3.7, 5} //\n }};\n\n int thread_values[2][3] = //\n {{4, 0, 3}, // thread 0 values\n {1, 2, 5}}; // thread 1 values\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .Sort(thread_keys[threadIdx.x], thread_values[threadIdx.x], decomposer_t{});\n\n custom_t expected_keys[2][3] = //\n {{\n // thread 0 expected keys\n {-2.5, 0}, //\n {+0.0, 1}, //\n {-0.0, 2} //\n },\n {\n // thread 1 expected keys\n {+1.1, 3}, //\n {+2.5, 4}, //\n {+3.7, 5} //\n }};\n\n int expected_values[2][3] = //\n {{0, 1, 2}, // thread 0 expected values\n {3, 4, 5}}; // thread 1 expected values\n // example-end pairs\n\n assert(thread_keys[threadIdx.x][0] == expected_keys[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_keys[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_keys[threadIdx.x][2]);\n\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][1] == expected_values[threadIdx.x][1]);\n assert(thread_values[threadIdx.x][2] == expected_values[threadIdx.x][2]);\n}\n\n__global__ void sort_pairs_bits()\n{\n // example-begin pairs-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys and values each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][1] = //\n {{\n {24.2, 1ll << 61} // thread 0 keys\n },\n {\n {42.4, 1ll << 60} // thread 1 keys\n }};\n\n int thread_values[2][1] = //\n {{1}, // thread 0 values\n {0}}; // thread 1 values\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000001110000011001100110011010 00100000000000...0000\n // decompose(in[1]) = 01000010001010011001100110011010 00010000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .Sort(thread_keys[threadIdx.x], thread_values[threadIdx.x], decomposer_t{}, begin_bit, end_bit);\n\n custom_t expected_keys[2][3] = //\n {{\n {42.4, 1ll << 60}, // thread 0 expected keys\n },\n {\n {24.2, 1ll << 61} // thread 1 expected keys\n }};\n\n int expected_values[2][1] = //\n {{0}, // thread 0 values\n {1}}; // thread 1 values\n // example-end pairs-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_keys[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n}\n\n__global__ void sort_pairs_descending()\n{\n // example-begin pairs-descending\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys and values each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+1.1, 2}, //\n {+2.5, 1}, //\n {-0.0, 4}, //\n },\n {\n // thread 1 keys\n {+0.0, 3}, //\n {-2.5, 5}, //\n {+3.7, 0} //\n }};\n\n int thread_values[2][3] = //\n {{2, 1, 4}, // thread 0 values\n {3, 5, 0}}; // thread 1 values\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortDescending(thread_keys[threadIdx.x], thread_values[threadIdx.x], decomposer_t{});\n\n custom_t expected_keys[2][3] = //\n {{\n // thread 0 expected keys\n {+3.7, 0}, //\n {+2.5, 1}, //\n {+1.1, 2}, //\n },\n {\n // thread 1 expected keys\n {-0.0, 4}, //\n {+0.0, 3}, //\n {-2.5, 5} //\n }};\n\n int expected_values[2][3] = //\n {{0, 1, 2}, // thread 0 expected values\n {4, 3, 5}}; // thread 1 expected values\n // example-end pairs-descending\n\n assert(thread_keys[threadIdx.x][0] == expected_keys[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_keys[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_keys[threadIdx.x][2]);\n\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][1] == expected_values[threadIdx.x][1]);\n assert(thread_values[threadIdx.x][2] == expected_values[threadIdx.x][2]);\n}\n\n__global__ void sort_pairs_descending_bits()\n{\n // example-begin pairs-descending-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys and values each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][1] = //\n {{\n {42.4, 1ll << 60} // thread 0 keys\n },\n {\n {24.2, 1ll << 61} // thread 1 keys\n }};\n\n int thread_values[2][1] = //\n {{1}, // thread 0 values\n {0}}; // thread 1 values\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000010001010011001100110011010 00010000000000...0000\n // decompose(in[1]) = 01000001110000011001100110011010 00100000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortDescending(thread_keys[threadIdx.x],\n thread_values[threadIdx.x],\n decomposer_t{},\n begin_bit,\n end_bit);\n\n custom_t expected_output[2][3] = //\n {{\n {24.2, 1ll << 61}, // thread 0 expected keys\n },\n {\n {42.4, 1ll << 60} // thread 1 expected keys\n }};\n\n int expected_values[2][1] = //\n {{0}, // thread 0 expected values\n {1}}; // thread 1 expected values\n // example-end pairs-descending-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n}\n\n__global__ void sort_keys_blocked_to_striped()\n{\n // example-begin keys-striped\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+2.5, 4}, //\n {-2.5, 0}, //\n {+1.1, 3}, //\n },\n {\n // thread 1 keys\n {+0.0, 1}, //\n {-0.0, 2}, //\n {+3.7, 5} //\n }};\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage).SortBlockedToStriped(thread_keys[threadIdx.x], decomposer_t{});\n\n custom_t expected_output[2][3] = //\n {{\n // thread 0 expected keys\n {-2.5, 0}, //\n {-0.0, 2}, //\n {+2.5, 4} //\n },\n {\n // thread 1 expected keys\n {+0.0, 1}, //\n {+1.1, 3}, //\n {+3.7, 5} //\n }};\n // example-end keys-striped\n \n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_output[threadIdx.x][2]);\n}\n\n__global__ void sort_keys_blocked_to_striped_bits()\n{\n // example-begin keys-striped-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 2 keys each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][2] = //\n {{ // thread 0 keys\n {24.2, 1ll << 62}, \n {42.4, 1ll << 61} \n },\n { // thread 1 keys\n {42.4, 1ll << 60}, \n {24.2, 1ll << 59} \n }};\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000001110000011001100110011010 01000000000000...0000\n // decompose(in[1]) = 01000010001010011001100110011010 00100000000000...0000\n // decompose(in[2]) = 01000001110000011001100110011010 00010000000000...0000\n // decompose(in[3]) = 01000010001010011001100110011010 00001000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0100xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // decompose(in[2]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // decompose(in[3]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0000xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortBlockedToStriped(thread_keys[threadIdx.x], decomposer_t{}, begin_bit, end_bit);\n\n custom_t expected_output[2][3] = //\n {{ // thread 0 expected keys\n {24.2, 1ll << 59}, \n {42.4, 1ll << 61}\n },\n { // thread 1 expected keys\n {42.4, 1ll << 60},\n {24.2, 1ll << 62} \n }};\n // example-end keys-striped-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n}\n\n__global__ void sort_pairs_blocked_to_striped()\n{\n // example-begin pairs-striped\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys and values each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+2.5, 4}, //\n {-2.5, 0}, //\n {+1.1, 3}, //\n },\n {\n // thread 1 keys\n {+0.0, 1}, //\n {-0.0, 2}, //\n {+3.7, 5} //\n }};\n\n int thread_values[2][3] = //\n {{4, 0, 3}, // thread 0 values\n {1, 2, 5}}; // thread 1 values\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortBlockedToStriped(thread_keys[threadIdx.x], thread_values[threadIdx.x], decomposer_t{});\n\n custom_t expected_output[2][3] = //\n {{\n // thread 0 expected keys\n {-2.5, 0}, //\n {-0.0, 2}, //\n {+2.5, 4} //\n },\n {\n // thread 1 expected keys\n {+0.0, 1}, //\n {+1.1, 3}, //\n {+3.7, 5} //\n }};\n\n int expected_values[2][3] = //\n {{0, 2, 4}, // thread 0 values\n {1, 3, 5}}; // thread 1 values\n // example-end pairs-striped\n \n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_output[threadIdx.x][2]);\n\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][1] == expected_values[threadIdx.x][1]);\n assert(thread_values[threadIdx.x][2] == expected_values[threadIdx.x][2]);\n}\n\n__global__ void sort_pairs_blocked_to_striped_bits()\n{\n // example-begin pairs-striped-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 2 keys and values each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][2] = //\n {{ // thread 0 keys\n {24.2, 1ll << 62}, \n {42.4, 1ll << 61} \n },\n { // thread 1 keys\n {42.4, 1ll << 60}, \n {24.2, 1ll << 59} \n }};\n\n int thread_values[2][2] = //\n {{3, 2}, // thread 0 values\n {1, 0}}; // thread 1 values\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000001110000011001100110011010 01000000000000...0000\n // decompose(in[1]) = 01000010001010011001100110011010 00100000000000...0000\n // decompose(in[2]) = 01000001110000011001100110011010 00010000000000...0000\n // decompose(in[3]) = 01000010001010011001100110011010 00001000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0100xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // decompose(in[2]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // decompose(in[3]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0000xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortBlockedToStriped(thread_keys[threadIdx.x],\n thread_values[threadIdx.x],\n decomposer_t{},\n begin_bit,\n end_bit);\n\n custom_t expected_output[2][3] = //\n {{ // thread 0 expected keys\n {24.2, 1ll << 59}, \n {42.4, 1ll << 61}\n },\n { // thread 1 expected keys\n {42.4, 1ll << 60},\n {24.2, 1ll << 62} \n }};\n\n int expected_values[2][2] = //\n {{0, 2}, // thread 0 values\n {1, 3}}; // thread 1 values\n // example-end pairs-striped-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][1] == expected_values[threadIdx.x][1]);\n}\n\n__global__ void sort_keys_descending_blocked_to_striped()\n{\n // example-begin keys-striped-descending\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+1.1, 2}, //\n {+2.5, 1}, //\n {-0.0, 4}, //\n },\n {\n // thread 1 keys\n {+0.0, 3}, //\n {-2.5, 5}, //\n {+3.7, 0} //\n }};\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortDescendingBlockedToStriped(thread_keys[threadIdx.x], decomposer_t{});\n\n custom_t expected_output[2][3] = //\n {{\n // thread 0 expected keys\n {+3.7, 0}, //\n {+1.1, 2}, //\n {+0.0, 3} //\n },\n {\n // thread 1 expected keys\n {+2.5, 1}, //\n {-0.0, 4}, //\n {-2.5, 5} //\n }};\n // example-end keys-striped-descending\n \n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_output[threadIdx.x][2]);\n}\n\n__global__ void sort_keys_descending_blocked_to_striped_bits()\n{\n // example-begin keys-striped-descending-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 2 keys each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][2] = //\n {{ // thread 0 keys\n {24.2, 1ll << 62}, \n {42.4, 1ll << 61} \n },\n { // thread 1 keys\n {42.4, 1ll << 60}, \n {24.2, 1ll << 59} \n }};\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000001110000011001100110011010 01000000000000...0000\n // decompose(in[1]) = 01000010001010011001100110011010 00100000000000...0000\n // decompose(in[2]) = 01000001110000011001100110011010 00010000000000...0000\n // decompose(in[3]) = 01000010001010011001100110011010 00001000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0100xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // decompose(in[2]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // decompose(in[3]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0000xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortDescendingBlockedToStriped(thread_keys[threadIdx.x], decomposer_t{}, begin_bit, end_bit);\n\n custom_t expected_output[2][2] = //\n {{ // thread 0 expected keys\n {24.2, 1ll << 62}, //\n {42.4, 1ll << 60} //\n },\n { // thread 1 expected keys\n {42.4, 1ll << 61}, //\n {24.2, 1ll << 59} //\n }};\n // example-end keys-striped-descending-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n}\n\n__global__ void sort_pairs_descending_blocked_to_striped()\n{\n // example-begin pairs-striped-descending\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys and values each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+1.1, 2}, //\n {+2.5, 1}, //\n {-0.0, 4}, //\n },\n {\n // thread 1 keys\n {+0.0, 3}, //\n {-2.5, 5}, //\n {+3.7, 0} //\n }};\n\n int thread_values[2][3] = //\n {{2, 1, 4}, // thread 0 values\n {3, 5, 0}}; // thread 1 values\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortDescendingBlockedToStriped(thread_keys[threadIdx.x],\n thread_values[threadIdx.x],\n decomposer_t{});\n\n custom_t expected_output[2][3] = //\n {{\n // thread 0 expected keys\n {+3.7, 0}, //\n {+1.1, 2}, //\n {+0.0, 3} //\n },\n {\n // thread 1 expected keys\n {+2.5, 1}, //\n {-0.0, 4}, //\n {-2.5, 5} //\n }};\n\n int expected_values[2][3] = //\n {{0, 2, 3}, // thread 0 values\n {1, 4, 5}}; // thread 1 values\n // example-end pairs-striped-descending\n \n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_output[threadIdx.x][2]);\n\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][1] == expected_values[threadIdx.x][1]);\n assert(thread_values[threadIdx.x][2] == expected_values[threadIdx.x][2]);\n}\n\n__global__ void sort_pairs_descending_blocked_to_striped_bits()\n{\n // example-begin pairs-striped-descending-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 2 keys and values each\n using block_radix_sort_t = cub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][2] = //\n {{ // thread 0 keys\n {24.2, 1ll << 62}, \n {42.4, 1ll << 61} \n },\n { // thread 1 keys\n {42.4, 1ll << 60}, \n {24.2, 1ll << 59} \n }};\n\n int thread_values[2][2] = //\n {{3, 2}, // thread 0 values\n {1, 0}}; // thread 1 values\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000001110000011001100110011010 01000000000000...0000\n // decompose(in[1]) = 01000010001010011001100110011010 00100000000000...0000\n // decompose(in[2]) = 01000001110000011001100110011010 00010000000000...0000\n // decompose(in[3]) = 01000010001010011001100110011010 00001000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0100xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // decompose(in[2]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // decompose(in[3]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0000xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortDescendingBlockedToStriped(thread_keys[threadIdx.x],\n thread_values[threadIdx.x],\n decomposer_t{},\n begin_bit,\n end_bit);\n\n custom_t expected_output[2][2] = //\n {{ // thread 0 expected keys\n {24.2, 1ll << 62}, //\n {42.4, 1ll << 60} //\n },\n { // thread 1 expected keys\n {42.4, 1ll << 61}, //\n {24.2, 1ll << 59} //\n }};\n\n int expected_values[2][2] = //\n {{3, 1}, // thread 0 values\n {2, 0}}; // thread 1 values\n // example-end pairs-striped-descending-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][1] == expected_values[threadIdx.x][1]);\n}\n\nTEST_CASE(\"Block radix sort works in some corner cases\", \"[radix][sort][block]\")\n{\n sort_keys<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_keys_bits<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_keys_descending<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_keys_descending_bits<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_pairs<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_pairs_bits<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_pairs_descending<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_pairs_descending_bits<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_keys_blocked_to_striped<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_keys_blocked_to_striped_bits<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_pairs_blocked_to_striped<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_pairs_blocked_to_striped_bits<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_keys_descending_blocked_to_striped<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_keys_descending_blocked_to_striped_bits<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_pairs_descending_blocked_to_striped<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n\n sort_pairs_descending_blocked_to_striped_bits<<<1, 2>>>();\n REQUIRE(cudaSuccess == cudaGetLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#undef NDEBUG\n#include \n\n#include \"catch2_test_block_radix_sort.cuh\"\n#include \"cub/block/radix_rank_sort_operations.cuh\"\n\n#include \n#include \n#include \n\n// example-begin custom-type\nstruct custom_t\n{\n float f;\n int unused;\n long long int lli;\n\n custom_t() = default;\n __device__ custom_t(float f, long long int lli)\n : f(f)\n , unused(42)\n , lli(lli)\n {}\n};\n\nstatic __device__ bool operator==(const custom_t &lhs, const custom_t &rhs)\n{\n return lhs.f == rhs.f && lhs.lli == rhs.lli;\n}\n\nstruct decomposer_t\n{\n __device__ ::cuda::std::tuple //\n operator()(custom_t &key) const\n {\n return {key.f, key.lli};\n }\n};\n// example-end custom-type\n\n__global__ void sort_keys()\n{\n // example-begin keys\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+2.5, 4}, //\n {-2.5, 0}, //\n {+1.1, 3}, //\n },\n {\n // thread 1 keys\n {+0.0, 1}, //\n {-0.0, 2}, //\n {+3.7, 5} //\n }};\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage).Sort(thread_keys[threadIdx.x], decomposer_t{});\n\n custom_t expected_output[2][3] = //\n {{\n // thread 0 expected keys\n {-2.5, 0}, //\n {+0.0, 1}, //\n {-0.0, 2} //\n },\n {\n // thread 1 expected keys\n {+1.1, 3}, //\n {+2.5, 4}, //\n {+3.7, 5} //\n }};\n // example-end keys\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_output[threadIdx.x][2]);\n}\n\n__global__ void sort_keys_bits()\n{\n // example-begin keys-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 1 key each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][1] = //\n {{\n {24.2, 1ll << 61} // thread 0 keys\n },\n {\n {42.4, 1ll << 60} // thread 1 keys\n }};\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000001110000011001100110011010 00100000000000...0000\n // decompose(in[1]) = 01000010001010011001100110011010 00010000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .Sort(thread_keys[threadIdx.x], decomposer_t{}, begin_bit, end_bit);\n\n custom_t expected_output[2][3] = //\n {{\n {42.4, 1ll << 60}, // thread 0 expected keys\n },\n {\n {24.2, 1ll << 61} // thread 1 expected keys\n }};\n // example-end keys-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n}\n\n__global__ void sort_keys_descending()\n{\n // example-begin keys-descending\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+1.1, 2}, //\n {+2.5, 1}, //\n {-0.0, 4}, //\n },\n {\n // thread 1 keys\n {+0.0, 3}, //\n {-2.5, 5}, //\n {+3.7, 0} //\n }};\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage).SortDescending(thread_keys[threadIdx.x], decomposer_t{});\n\n custom_t expected_output[2][3] = //\n {{\n // thread 0 expected keys\n {+3.7, 0}, //\n {+2.5, 1}, //\n {+1.1, 2}, //\n },\n {\n // thread 1 expected keys\n {-0.0, 4}, //\n {+0.0, 3}, //\n {-2.5, 5} //\n }};\n // example-end keys-descending\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_output[threadIdx.x][2]);\n}\n\n__global__ void sort_keys_descending_bits()\n{\n // example-begin keys-descending-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 1 key each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][1] = //\n {{\n {42.4, 1ll << 60} // thread 0 keys\n },\n {\n {24.2, 1ll << 61} // thread 1 keys\n }};\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000010001010011001100110011010 00010000000000...0000\n // decompose(in[1]) = 01000001110000011001100110011010 00100000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortDescending(thread_keys[threadIdx.x], decomposer_t{}, begin_bit, end_bit);\n\n custom_t expected_output[2][3] = //\n {{\n {24.2, 1ll << 61}, // thread 0 expected keys\n },\n {\n {42.4, 1ll << 60} // thread 1 expected keys\n }};\n // example-end keys-descending-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n}\n\n__global__ void sort_pairs()\n{\n // example-begin pairs\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys and values each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+2.5, 4}, //\n {-2.5, 0}, //\n {+1.1, 3}, //\n },\n {\n // thread 1 keys\n {+0.0, 1}, //\n {-0.0, 2}, //\n {+3.7, 5} //\n }};\n\n int thread_values[2][3] = //\n {{4, 0, 3}, // thread 0 values\n {1, 2, 5}}; // thread 1 values\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .Sort(thread_keys[threadIdx.x], thread_values[threadIdx.x], decomposer_t{});\n\n custom_t expected_keys[2][3] = //\n {{\n // thread 0 expected keys\n {-2.5, 0}, //\n {+0.0, 1}, //\n {-0.0, 2} //\n },\n {\n // thread 1 expected keys\n {+1.1, 3}, //\n {+2.5, 4}, //\n {+3.7, 5} //\n }};\n\n int expected_values[2][3] = //\n {{0, 1, 2}, // thread 0 expected values\n {3, 4, 5}}; // thread 1 expected values\n // example-end pairs\n\n assert(thread_keys[threadIdx.x][0] == expected_keys[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_keys[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_keys[threadIdx.x][2]);\n\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][1] == expected_values[threadIdx.x][1]);\n assert(thread_values[threadIdx.x][2] == expected_values[threadIdx.x][2]);\n}\n\n__global__ void sort_pairs_bits()\n{\n // example-begin pairs-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys and values each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][1] = //\n {{\n {24.2, 1ll << 61} // thread 0 keys\n },\n {\n {42.4, 1ll << 60} // thread 1 keys\n }};\n\n int thread_values[2][1] = //\n {{1}, // thread 0 values\n {0}}; // thread 1 values\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000001110000011001100110011010 00100000000000...0000\n // decompose(in[1]) = 01000010001010011001100110011010 00010000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .Sort(thread_keys[threadIdx.x], thread_values[threadIdx.x], decomposer_t{}, begin_bit, end_bit);\n\n custom_t expected_keys[2][3] = //\n {{\n {42.4, 1ll << 60}, // thread 0 expected keys\n },\n {\n {24.2, 1ll << 61} // thread 1 expected keys\n }};\n\n int expected_values[2][1] = //\n {{0}, // thread 0 values\n {1}}; // thread 1 values\n // example-end pairs-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_keys[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n}\n\n__global__ void sort_pairs_descending()\n{\n // example-begin pairs-descending\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys and values each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+1.1, 2}, //\n {+2.5, 1}, //\n {-0.0, 4}, //\n },\n {\n // thread 1 keys\n {+0.0, 3}, //\n {-2.5, 5}, //\n {+3.7, 0} //\n }};\n\n int thread_values[2][3] = //\n {{2, 1, 4}, // thread 0 values\n {3, 5, 0}}; // thread 1 values\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortDescending(thread_keys[threadIdx.x], thread_values[threadIdx.x], decomposer_t{});\n\n custom_t expected_keys[2][3] = //\n {{\n // thread 0 expected keys\n {+3.7, 0}, //\n {+2.5, 1}, //\n {+1.1, 2}, //\n },\n {\n // thread 1 expected keys\n {-0.0, 4}, //\n {+0.0, 3}, //\n {-2.5, 5} //\n }};\n\n int expected_values[2][3] = //\n {{0, 1, 2}, // thread 0 expected values\n {4, 3, 5}}; // thread 1 expected values\n // example-end pairs-descending\n\n assert(thread_keys[threadIdx.x][0] == expected_keys[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_keys[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_keys[threadIdx.x][2]);\n\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][1] == expected_values[threadIdx.x][1]);\n assert(thread_values[threadIdx.x][2] == expected_values[threadIdx.x][2]);\n}\n\n__global__ void sort_pairs_descending_bits()\n{\n // example-begin pairs-descending-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys and values each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][1] = //\n {{\n {42.4, 1ll << 60} // thread 0 keys\n },\n {\n {24.2, 1ll << 61} // thread 1 keys\n }};\n\n int thread_values[2][1] = //\n {{1}, // thread 0 values\n {0}}; // thread 1 values\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000010001010011001100110011010 00010000000000...0000\n // decompose(in[1]) = 01000001110000011001100110011010 00100000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortDescending(thread_keys[threadIdx.x],\n thread_values[threadIdx.x],\n decomposer_t{},\n begin_bit,\n end_bit);\n\n custom_t expected_output[2][3] = //\n {{\n {24.2, 1ll << 61}, // thread 0 expected keys\n },\n {\n {42.4, 1ll << 60} // thread 1 expected keys\n }};\n\n int expected_values[2][1] = //\n {{0}, // thread 0 expected values\n {1}}; // thread 1 expected values\n // example-end pairs-descending-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n}\n\n__global__ void sort_keys_blocked_to_striped()\n{\n // example-begin keys-striped\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+2.5, 4}, //\n {-2.5, 0}, //\n {+1.1, 3}, //\n },\n {\n // thread 1 keys\n {+0.0, 1}, //\n {-0.0, 2}, //\n {+3.7, 5} //\n }};\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage).SortBlockedToStriped(thread_keys[threadIdx.x], decomposer_t{});\n\n custom_t expected_output[2][3] = //\n {{\n // thread 0 expected keys\n {-2.5, 0}, //\n {-0.0, 2}, //\n {+2.5, 4} //\n },\n {\n // thread 1 expected keys\n {+0.0, 1}, //\n {+1.1, 3}, //\n {+3.7, 5} //\n }};\n // example-end keys-striped\n \n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_output[threadIdx.x][2]);\n}\n\n__global__ void sort_keys_blocked_to_striped_bits()\n{\n // example-begin keys-striped-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 2 keys each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][2] = //\n {{ // thread 0 keys\n {24.2, 1ll << 62}, \n {42.4, 1ll << 61} \n },\n { // thread 1 keys\n {42.4, 1ll << 60}, \n {24.2, 1ll << 59} \n }};\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000001110000011001100110011010 01000000000000...0000\n // decompose(in[1]) = 01000010001010011001100110011010 00100000000000...0000\n // decompose(in[2]) = 01000001110000011001100110011010 00010000000000...0000\n // decompose(in[3]) = 01000010001010011001100110011010 00001000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0100xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // decompose(in[2]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // decompose(in[3]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0000xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortBlockedToStriped(thread_keys[threadIdx.x], decomposer_t{}, begin_bit, end_bit);\n\n custom_t expected_output[2][3] = //\n {{ // thread 0 expected keys\n {24.2, 1ll << 59}, \n {42.4, 1ll << 61}\n },\n { // thread 1 expected keys\n {42.4, 1ll << 60},\n {24.2, 1ll << 62} \n }};\n // example-end keys-striped-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n}\n\n__global__ void sort_pairs_blocked_to_striped()\n{\n // example-begin pairs-striped\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys and values each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+2.5, 4}, //\n {-2.5, 0}, //\n {+1.1, 3}, //\n },\n {\n // thread 1 keys\n {+0.0, 1}, //\n {-0.0, 2}, //\n {+3.7, 5} //\n }};\n\n int thread_values[2][3] = //\n {{4, 0, 3}, // thread 0 values\n {1, 2, 5}}; // thread 1 values\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortBlockedToStriped(thread_keys[threadIdx.x], thread_values[threadIdx.x], decomposer_t{});\n\n custom_t expected_output[2][3] = //\n {{\n // thread 0 expected keys\n {-2.5, 0}, //\n {-0.0, 2}, //\n {+2.5, 4} //\n },\n {\n // thread 1 expected keys\n {+0.0, 1}, //\n {+1.1, 3}, //\n {+3.7, 5} //\n }};\n\n int expected_values[2][3] = //\n {{0, 2, 4}, // thread 0 values\n {1, 3, 5}}; // thread 1 values\n // example-end pairs-striped\n \n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_output[threadIdx.x][2]);\n\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][1] == expected_values[threadIdx.x][1]);\n assert(thread_values[threadIdx.x][2] == expected_values[threadIdx.x][2]);\n}\n\n__global__ void sort_pairs_blocked_to_striped_bits()\n{\n // example-begin pairs-striped-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 2 keys and values each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][2] = //\n {{ // thread 0 keys\n {24.2, 1ll << 62}, \n {42.4, 1ll << 61} \n },\n { // thread 1 keys\n {42.4, 1ll << 60}, \n {24.2, 1ll << 59} \n }};\n\n int thread_values[2][2] = //\n {{3, 2}, // thread 0 values\n {1, 0}}; // thread 1 values\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000001110000011001100110011010 01000000000000...0000\n // decompose(in[1]) = 01000010001010011001100110011010 00100000000000...0000\n // decompose(in[2]) = 01000001110000011001100110011010 00010000000000...0000\n // decompose(in[3]) = 01000010001010011001100110011010 00001000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0100xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // decompose(in[2]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // decompose(in[3]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0000xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortBlockedToStriped(thread_keys[threadIdx.x],\n thread_values[threadIdx.x],\n decomposer_t{},\n begin_bit,\n end_bit);\n\n custom_t expected_output[2][3] = //\n {{ // thread 0 expected keys\n {24.2, 1ll << 59}, \n {42.4, 1ll << 61}\n },\n { // thread 1 expected keys\n {42.4, 1ll << 60},\n {24.2, 1ll << 62} \n }};\n\n int expected_values[2][2] = //\n {{0, 2}, // thread 0 values\n {1, 3}}; // thread 1 values\n // example-end pairs-striped-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][1] == expected_values[threadIdx.x][1]);\n}\n\n__global__ void sort_keys_descending_blocked_to_striped()\n{\n // example-begin keys-striped-descending\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+1.1, 2}, //\n {+2.5, 1}, //\n {-0.0, 4}, //\n },\n {\n // thread 1 keys\n {+0.0, 3}, //\n {-2.5, 5}, //\n {+3.7, 0} //\n }};\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortDescendingBlockedToStriped(thread_keys[threadIdx.x], decomposer_t{});\n\n custom_t expected_output[2][3] = //\n {{\n // thread 0 expected keys\n {+3.7, 0}, //\n {+1.1, 2}, //\n {+0.0, 3} //\n },\n {\n // thread 1 expected keys\n {+2.5, 1}, //\n {-0.0, 4}, //\n {-2.5, 5} //\n }};\n // example-end keys-striped-descending\n \n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_output[threadIdx.x][2]);\n}\n\n__global__ void sort_keys_descending_blocked_to_striped_bits()\n{\n // example-begin keys-striped-descending-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 2 keys each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][2] = //\n {{ // thread 0 keys\n {24.2, 1ll << 62}, \n {42.4, 1ll << 61} \n },\n { // thread 1 keys\n {42.4, 1ll << 60}, \n {24.2, 1ll << 59} \n }};\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000001110000011001100110011010 01000000000000...0000\n // decompose(in[1]) = 01000010001010011001100110011010 00100000000000...0000\n // decompose(in[2]) = 01000001110000011001100110011010 00010000000000...0000\n // decompose(in[3]) = 01000010001010011001100110011010 00001000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0100xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // decompose(in[2]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // decompose(in[3]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0000xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortDescendingBlockedToStriped(thread_keys[threadIdx.x], decomposer_t{}, begin_bit, end_bit);\n\n custom_t expected_output[2][2] = //\n {{ // thread 0 expected keys\n {24.2, 1ll << 62}, //\n {42.4, 1ll << 60} //\n },\n { // thread 1 expected keys\n {42.4, 1ll << 61}, //\n {24.2, 1ll << 59} //\n }};\n // example-end keys-striped-descending-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n}\n\n__global__ void sort_pairs_descending_blocked_to_striped()\n{\n // example-begin pairs-striped-descending\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 3 keys and values each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][3] = //\n {{\n // thread 0 keys\n {+1.1, 2}, //\n {+2.5, 1}, //\n {-0.0, 4}, //\n },\n {\n // thread 1 keys\n {+0.0, 3}, //\n {-2.5, 5}, //\n {+3.7, 0} //\n }};\n\n int thread_values[2][3] = //\n {{2, 1, 4}, // thread 0 values\n {3, 5, 0}}; // thread 1 values\n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortDescendingBlockedToStriped(thread_keys[threadIdx.x],\n thread_values[threadIdx.x],\n decomposer_t{});\n\n custom_t expected_output[2][3] = //\n {{\n // thread 0 expected keys\n {+3.7, 0}, //\n {+1.1, 2}, //\n {+0.0, 3} //\n },\n {\n // thread 1 expected keys\n {+2.5, 1}, //\n {-0.0, 4}, //\n {-2.5, 5} //\n }};\n\n int expected_values[2][3] = //\n {{0, 2, 3}, // thread 0 values\n {1, 4, 5}}; // thread 1 values\n // example-end pairs-striped-descending\n \n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n assert(thread_keys[threadIdx.x][2] == expected_output[threadIdx.x][2]);\n\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][1] == expected_values[threadIdx.x][1]);\n assert(thread_values[threadIdx.x][2] == expected_values[threadIdx.x][2]);\n}\n\n__global__ void sort_pairs_descending_blocked_to_striped_bits()\n{\n // example-begin pairs-striped-descending-bits\n // Specialize `cub::BlockRadixSort` for a 1D block of 2 threads owning 2 keys and values each\n using block_radix_sort_t = hipcub::BlockRadixSort;\n\n // Allocate shared memory for `cub::BlockRadixSort`\n __shared__ block_radix_sort_t::TempStorage temp_storage;\n\n // Obtain a segment of consecutive items that are blocked across threads\n custom_t thread_keys[2][2] = //\n {{ // thread 0 keys\n {24.2, 1ll << 62}, \n {42.4, 1ll << 61} \n },\n { // thread 1 keys\n {42.4, 1ll << 60}, \n {24.2, 1ll << 59} \n }};\n\n int thread_values[2][2] = //\n {{3, 2}, // thread 0 values\n {1, 0}}; // thread 1 values\n\n const int begin_bit = sizeof(long long int) * 8 - 4; // 60\n const int end_bit = sizeof(long long int) * 8 + 4; // 68\n\n // Decomposition orders the bits as follows:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = 01000001110000011001100110011010 01000000000000...0000\n // decompose(in[1]) = 01000010001010011001100110011010 00100000000000...0000\n // decompose(in[2]) = 01000001110000011001100110011010 00010000000000...0000\n // decompose(in[3]) = 01000010001010011001100110011010 00001000000000...0000\n // <----------- higher bits / lower bits -----------> \n //\n // The bit subrange `[60, 68)` specifies differentiating key bits:\n //\n // <------------- fp32 -----------> <------ int64 ------>\n // decompose(in[0]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0100xxxxxxxxxx...xxxx\n // decompose(in[1]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0010xxxxxxxxxx...xxxx\n // decompose(in[2]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0001xxxxxxxxxx...xxxx\n // decompose(in[3]) = xxxxxxxxxxxxxxxxxxxxxxxxxxxx1010 0000xxxxxxxxxx...xxxx\n // <----------- higher bits / lower bits -----------> \n\n // Collectively sort the keys\n block_radix_sort_t(temp_storage)\n .SortDescendingBlockedToStriped(thread_keys[threadIdx.x],\n thread_values[threadIdx.x],\n decomposer_t{},\n begin_bit,\n end_bit);\n\n custom_t expected_output[2][2] = //\n {{ // thread 0 expected keys\n {24.2, 1ll << 62}, //\n {42.4, 1ll << 60} //\n },\n { // thread 1 expected keys\n {42.4, 1ll << 61}, //\n {24.2, 1ll << 59} //\n }};\n\n int expected_values[2][2] = //\n {{3, 1}, // thread 0 values\n {2, 0}}; // thread 1 values\n // example-end pairs-striped-descending-bits\n\n assert(thread_keys[threadIdx.x][0] == expected_output[threadIdx.x][0]);\n assert(thread_keys[threadIdx.x][1] == expected_output[threadIdx.x][1]);\n\n assert(thread_values[threadIdx.x][0] == expected_values[threadIdx.x][0]);\n assert(thread_values[threadIdx.x][1] == expected_values[threadIdx.x][1]);\n}\n\nTEST_CASE(\"Block radix sort works in some corner cases\", \"[radix][sort][block]\")\n{\n sort_keys<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_keys_bits<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_keys_descending<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_keys_descending_bits<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_pairs<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_pairs_bits<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_pairs_descending<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_pairs_descending_bits<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_keys_blocked_to_striped<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_keys_blocked_to_striped_bits<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_pairs_blocked_to_striped<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_pairs_blocked_to_striped_bits<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_keys_descending_blocked_to_striped<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_keys_descending_blocked_to_striped_bits<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_pairs_descending_blocked_to_striped<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n\n sort_pairs_descending_blocked_to_striped_bits<<<1, 2>>>();\n REQUIRE(hipSuccess == hipGetLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n}\n", "source": "/kaggle/working/cub/test/catch2_test_block_radix_sort_custom.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"c2h/custom_type.cuh\"\n#include \"catch2_test_helper.h\"\n\nstruct CustomLess\n{\n template \n __device__ __host__ bool operator()(const T &lhs, const T &rhs)\n {\n return lhs < rhs;\n }\n};\n\n/**\n * @brief Kernel to dispatch to the appropriate WarpMergeSort member function, sorting keys-only.\n */\ntemplate \n__global__ void\nwarp_merge_sort_kernel(T *in, T *out, SegmentSizeItT segment_sizes, T oob_default, ActionT action)\n{\n using warp_merge_sort_t = cub::WarpMergeSort;\n using storage_t = typename warp_merge_sort_t::TempStorage;\n\n // Get linear thread and warp index\n const int tid = threadIdx.x;\n const int warp_id = tid / LOGICAL_WARP_THREADS;\n\n // Test case of partially finished CTA\n if (warp_id >= TOTAL_WARPS)\n {\n return;\n }\n\n // Thread-local storage & warp-scope temporary storage allocation\n T thread_data[ITEMS_PER_THREAD];\n __shared__ storage_t storage[TOTAL_WARPS];\n\n // Instantiate warp-scope algorithm\n warp_merge_sort_t warp_sort(storage[warp_id]);\n\n const int warp_offset = LOGICAL_WARP_THREADS * ITEMS_PER_THREAD * warp_id;\n const int thread_offset = warp_offset + warp_sort.get_linear_tid() * ITEMS_PER_THREAD;\n const int valid_items = segment_sizes[warp_id];\n\n // Load data\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n const int idx = thread_offset + item;\n thread_data[item] = in[idx];\n }\n cub::WARP_SYNC(warp_sort.get_member_mask());\n\n // Run merge sort test\n action(warp_sort, thread_data, valid_items, oob_default);\n\n // Store data\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n const int idx = thread_offset + item;\n out[idx] = (idx - warp_offset) >= valid_items ? oob_default : thread_data[item];\n }\n}\n\n/**\n * @brief Kernel to dispatch to the appropriate WarpMergeSort member function, sorting key-value\n * pairs.\n */\ntemplate \n__global__ void warp_merge_sort_kernel(KeyT *keys_in,\n KeyT *keys_out,\n ValueT *values_in,\n ValueT *values_out,\n SegmentSizeItT segment_sizes,\n KeyT oob_default,\n ActionT action)\n{\n using warp_merge_sort_t =\n cub::WarpMergeSort;\n using storage_t = typename warp_merge_sort_t::TempStorage;\n\n // Get linear thread and warp index\n const int tid = cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n const int warp_id = tid / LOGICAL_WARP_THREADS;\n\n // Test case of partially finished CTA\n if (warp_id >= TOTAL_WARPS)\n {\n return;\n }\n\n // Thread-local storage & warp-scope temporary storage allocation\n KeyT keys[ITEMS_PER_THREAD];\n ValueT values[ITEMS_PER_THREAD];\n __shared__ storage_t storage[TOTAL_WARPS];\n\n // Instantiate warp-scope algorithm\n warp_merge_sort_t warp_sort(storage[warp_id]);\n\n const int warp_offset = LOGICAL_WARP_THREADS * ITEMS_PER_THREAD * warp_id;\n const int thread_offset = warp_offset + warp_sort.get_linear_tid() * ITEMS_PER_THREAD;\n const int valid_items = segment_sizes[warp_id];\n\n // Load data\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n const int idx = thread_offset + item;\n keys[item] = keys_in[idx];\n values[item] = values_in[idx];\n }\n cub::WARP_SYNC(warp_sort.get_member_mask());\n\n // Run merge sort test\n action(warp_sort, keys, values, valid_items, oob_default);\n\n // Store data\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n const int idx = thread_offset + item;\n keys_out[idx] = (idx - warp_offset) >= valid_items ? oob_default : keys[item];\n values_out[idx] = (idx - warp_offset) >= valid_items ? ValueT{} : values[item];\n }\n}\n\n// -----------------------------------------------------------\n// Dimensions being instantiated:\n// {full,partial} x {stable, 'unstable'} x {keys, kv-pairs}\n// -----------------------------------------------------------\n\n/**\n * @brief Delegate wrapper for WarpMergeSort::StableSort on keys-only\n */\nstruct warp_stable_sort_keys_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n T (&thread_data)[ITEMS_PER_THREAD],\n int /*valid_items*/,\n T /*oob_default*/) const\n {\n warp_sort.StableSort(thread_data, CustomLess{});\n }\n};\n\n/**\n * @brief Delegate wrapper for partial WarpMergeSort::StableSort keys-only\n */\nstruct warp_partial_stable_sort_keys_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n T (&thread_data)[ITEMS_PER_THREAD],\n int valid_items,\n T oob_default) const\n {\n warp_sort.StableSort(thread_data, CustomLess{}, valid_items, oob_default);\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpMergeSort::Sort on keys-only\n */\nstruct warp_sort_keys_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n T (&thread_data)[ITEMS_PER_THREAD],\n int /*valid_items*/,\n T /*oob_default*/) const\n {\n warp_sort.Sort(thread_data, CustomLess{});\n }\n};\n\n/**\n * @brief Delegate wrapper for partial WarpMergeSort::StableSort keys-only\n */\nstruct warp_partial_sort_keys_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n T (&thread_data)[ITEMS_PER_THREAD],\n int valid_items,\n T oob_default) const\n {\n warp_sort.Sort(thread_data, CustomLess{}, valid_items, oob_default);\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpMergeSort::StableSort on key-value pairs\n */\nstruct warp_stable_sort_pairs_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n KeyT (&keys)[ITEMS_PER_THREAD],\n ValueT (&values)[ITEMS_PER_THREAD],\n int /*valid_items*/,\n KeyT /*oob_default*/) const\n {\n warp_sort.StableSort(keys, values, CustomLess{});\n }\n};\n\n/**\n * @brief Delegate wrapper for partial WarpMergeSort::StableSort key-value pairs\n */\nstruct warp_partial_stable_sort_pairs_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n KeyT (&keys)[ITEMS_PER_THREAD],\n ValueT (&values)[ITEMS_PER_THREAD],\n int valid_items,\n KeyT oob_default) const\n {\n warp_sort.StableSort(keys, values, CustomLess{}, valid_items, oob_default);\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpMergeSort::Sort on key-value pairs\n */\nstruct warp_sort_pairs_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n KeyT (&keys)[ITEMS_PER_THREAD],\n ValueT (&values)[ITEMS_PER_THREAD],\n int /*valid_items*/,\n KeyT /*oob_default*/) const\n {\n warp_sort.Sort(keys, values, CustomLess{});\n }\n};\n\n/**\n * @brief Delegate wrapper for partial WarpMergeSort::StableSort key-value pairs\n */\nstruct warp_partial_sort_pairs_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n KeyT (&keys)[ITEMS_PER_THREAD],\n ValueT (&values)[ITEMS_PER_THREAD],\n int valid_items,\n KeyT oob_default) const\n {\n warp_sort.Sort(keys, values, CustomLess{}, valid_items, oob_default);\n }\n};\n\n/**\n * @brief Dispatch helper function for sorting keys\n */\ntemplate \nvoid warp_merge_sort(thrust::device_vector &in,\n thrust::device_vector &out,\n SegmentSizesItT segment_sizes,\n T oob_default,\n ActionT action)\n{\n warp_merge_sort_kernel\n <<<1, LOGICAL_WARP_THREADS * TOTAL_WARPS>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(out.data()),\n segment_sizes,\n oob_default,\n action);\n\n REQUIRE(cudaSuccess == cudaPeekAtLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n}\n\n/**\n * @brief Dispatch helper function for sorting key-value pairs\n */\ntemplate \nvoid warp_merge_sort(thrust::device_vector &keys_in,\n thrust::device_vector &keys_out,\n thrust::device_vector &values_in,\n thrust::device_vector &values_out,\n SegmentSizesItT segment_sizes,\n KeyT oob_default,\n ActionT action)\n{\n warp_merge_sort_kernel\n <<<1, LOGICAL_WARP_THREADS * TOTAL_WARPS>>>(thrust::raw_pointer_cast(keys_in.data()),\n thrust::raw_pointer_cast(keys_out.data()),\n thrust::raw_pointer_cast(values_in.data()),\n thrust::raw_pointer_cast(values_out.data()),\n segment_sizes,\n oob_default,\n action);\n\n REQUIRE(cudaSuccess == cudaPeekAtLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n}\n\n/**\n * @brief Performs a stable sort on per-warp segments of data and assigns oob_default to items that\n * are out-of-bounds.\n */\ntemplate \nvoid compute_host_reference(RandomItT h_data,\n SegmentSizeItT segment_sizes,\n unsigned int num_segments,\n T oob_default,\n int logical_warp_items)\n{\n for (unsigned int segment_id = 0; segment_id < num_segments; segment_id++)\n {\n unsigned int segment_size = segment_sizes[segment_id];\n std::stable_sort(h_data, h_data + segment_size);\n std::fill(h_data + segment_size, h_data + logical_warp_items, oob_default);\n h_data += logical_warp_items;\n }\n}\n\n/**\n * @brief Stability requirement of the sorting algorithm\n */\nenum class stability\n{\n stable,\n unstable\n};\n\n// List of key types to test\nusing custom_t =\n c2h::custom_type_t;\nusing key_types = c2h::type_list;\n\n// List of value types\nusing value_types = c2h::type_list;\n\n// Logical warp sizes to test\nusing logical_warp_threads = c2h::enum_type_list;\n\n// Number of items per thread to test\nusing items_per_thread_list = c2h::enum_type_list;\n\n// Whether the sort is required to be stable or not\nusing stability_list = c2h::enum_type_list;\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int logical_warp_threads = c2h::get<1, TestType>::value;\n static constexpr int items_per_thread = c2h::get<2, TestType>::value;\n static constexpr int logical_warp_items = logical_warp_threads * items_per_thread;\n static constexpr int total_warps = 2;\n static constexpr int tile_size = items_per_thread * total_warps * logical_warp_threads;\n static constexpr bool is_stable = c2h::get<3, TestType>::value == stability::stable;\n};\n\nCUB_TEST(\"Warp sort on keys-only works\",\n \"[sort][warp]\",\n key_types,\n logical_warp_threads,\n items_per_thread_list,\n stability_list)\n{\n using params = params_t;\n using type = typename params::type;\n using warp_sort_delegate =\n cub::detail::conditional_t;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n auto segment_sizes = thrust::make_constant_iterator(params::logical_warp_items);\n const auto oob_default = std::numeric_limits::max();\n c2h::gen(CUB_SEED(10), d_in);\n\n // Run test\n warp_merge_sort(\n d_in,\n d_out,\n segment_sizes,\n oob_default,\n warp_sort_delegate{});\n\n // Prepare verification data\n thrust::host_vector h_in_out = d_in;\n compute_host_reference(h_in_out.begin(),\n segment_sizes,\n params::total_warps,\n oob_default,\n params::logical_warp_items);\n\n // Verify results\n REQUIRE(h_in_out == d_out);\n}\n\nCUB_TEST(\"Warp sort keys-only on partial warp-tile works\",\n \"[sort][warp]\",\n key_types,\n logical_warp_threads,\n items_per_thread_list,\n stability_list)\n{\n using params = params_t;\n using type = typename params::type;\n using warp_sort_delegate = cub::detail::\n conditional_t;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_segment_sizes(params::total_warps);\n const auto oob_default = std::numeric_limits::max();\n c2h::gen(CUB_SEED(5), d_in);\n c2h::gen(CUB_SEED(5), d_segment_sizes, 0, params::logical_warp_items);\n\n // Run test\n warp_merge_sort(\n d_in,\n d_out,\n d_segment_sizes.cbegin(),\n oob_default,\n warp_sort_delegate{});\n\n // Prepare verification data\n thrust::host_vector h_in_out = d_in;\n thrust::host_vector segment_sizes = d_segment_sizes;\n compute_host_reference(h_in_out.begin(),\n segment_sizes,\n params::total_warps,\n oob_default,\n params::logical_warp_items);\n\n // Verify results\n REQUIRE(h_in_out == d_out);\n}\n\nCUB_TEST(\"Warp sort on keys-value pairs works\",\n \"[sort][warp]\",\n key_types,\n logical_warp_threads,\n items_per_thread_list,\n stability_list,\n value_types)\n{\n using params = params_t;\n using key_type = typename params::type;\n using value_type = typename c2h::get<4, TestType>;\n using warp_sort_delegate =\n cub::detail::conditional_t;\n\n // Prepare test data\n thrust::device_vector d_keys_in(params::tile_size);\n thrust::device_vector d_keys_out(params::tile_size);\n thrust::device_vector d_values_in(params::tile_size);\n thrust::device_vector d_values_out(params::tile_size);\n auto segment_sizes = thrust::make_constant_iterator(params::logical_warp_items);\n const auto oob_default = std::numeric_limits::max();\n c2h::gen(CUB_SEED(10), d_keys_in);\n\n // Run test\n warp_merge_sort(\n d_keys_in,\n d_keys_out,\n d_values_in,\n d_values_out,\n segment_sizes,\n oob_default,\n warp_stable_sort_pairs_t{});\n\n // Prepare verification data\n thrust::host_vector h_keys_in_out = d_keys_in;\n thrust::host_vector h_values_in_out = d_values_in;\n auto cpu_kv_pairs = thrust::make_zip_iterator(h_keys_in_out.begin(), h_values_in_out.begin());\n compute_host_reference(cpu_kv_pairs,\n segment_sizes,\n params::total_warps,\n thrust::make_tuple(oob_default, value_type{}),\n params::logical_warp_items);\n\n // Verify results\n REQUIRE(h_keys_in_out == d_keys_out);\n REQUIRE(h_values_in_out == d_values_out);\n}\n\nCUB_TEST(\"Warp sort on key-value pairs of a partial warp-tile works\",\n \"[sort][warp]\",\n key_types,\n logical_warp_threads,\n items_per_thread_list,\n stability_list,\n value_types)\n{\n using params = params_t;\n using key_type = typename params::type;\n using value_type = typename c2h::get<4, TestType>;\n using warp_sort_delegate = cub::detail::\n conditional_t;\n\n // Prepare test data\n thrust::device_vector d_keys_in(params::tile_size);\n thrust::device_vector d_keys_out(params::tile_size);\n thrust::device_vector d_values_in(params::tile_size);\n thrust::device_vector d_values_out(params::tile_size);\n thrust::device_vector d_segment_sizes(params::total_warps);\n const auto oob_default = std::numeric_limits::max();\n c2h::gen(CUB_SEED(5), d_keys_in);\n c2h::gen(CUB_SEED(5), d_segment_sizes, 0, params::logical_warp_items);\n\n // Run test\n warp_merge_sort(\n d_keys_in,\n d_keys_out,\n d_values_in,\n d_values_out,\n d_segment_sizes.cbegin(),\n oob_default,\n warp_sort_delegate{});\n\n // Prepare verification data\n thrust::host_vector h_keys_in_out = d_keys_in;\n thrust::host_vector h_values_in_out = d_values_in;\n thrust::host_vector segment_sizes = d_segment_sizes;\n auto cpu_kv_pairs = thrust::make_zip_iterator(h_keys_in_out.begin(), h_values_in_out.begin());\n compute_host_reference(cpu_kv_pairs,\n segment_sizes,\n params::total_warps,\n thrust::make_tuple(oob_default, value_type{}),\n params::logical_warp_items);\n\n // Verify results\n REQUIRE(h_keys_in_out == d_keys_out);\n REQUIRE(h_values_in_out == d_values_out);\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"c2h/custom_type.cuh\"\n#include \"catch2_test_helper.h\"\n\nstruct CustomLess\n{\n template \n __device__ __host__ bool operator()(const T &lhs, const T &rhs)\n {\n return lhs < rhs;\n }\n};\n\n/**\n * @brief Kernel to dispatch to the appropriate WarpMergeSort member function, sorting keys-only.\n */\ntemplate \n__global__ void\nwarp_merge_sort_kernel(T *in, T *out, SegmentSizeItT segment_sizes, T oob_default, ActionT action)\n{\n using warp_merge_sort_t = hipcub::WarpMergeSort;\n using storage_t = typename warp_merge_sort_t::TempStorage;\n\n // Get linear thread and warp index\n const int tid = threadIdx.x;\n const int warp_id = tid / LOGICAL_WARP_THREADS;\n\n // Test case of partially finished CTA\n if (warp_id >= TOTAL_WARPS)\n {\n return;\n }\n\n // Thread-local storage & warp-scope temporary storage allocation\n T thread_data[ITEMS_PER_THREAD];\n __shared__ storage_t storage[TOTAL_WARPS];\n\n // Instantiate warp-scope algorithm\n warp_merge_sort_t warp_sort(storage[warp_id]);\n\n const int warp_offset = LOGICAL_WARP_THREADS * ITEMS_PER_THREAD * warp_id;\n const int thread_offset = warp_offset + warp_sort.get_linear_tid() * ITEMS_PER_THREAD;\n const int valid_items = segment_sizes[warp_id];\n\n // Load data\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n const int idx = thread_offset + item;\n thread_data[item] = in[idx];\n }\n hipcub::WARP_SYNC(warp_sort.get_member_mask());\n\n // Run merge sort test\n action(warp_sort, thread_data, valid_items, oob_default);\n\n // Store data\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n const int idx = thread_offset + item;\n out[idx] = (idx - warp_offset) >= valid_items ? oob_default : thread_data[item];\n }\n}\n\n/**\n * @brief Kernel to dispatch to the appropriate WarpMergeSort member function, sorting key-value\n * pairs.\n */\ntemplate \n__global__ void warp_merge_sort_kernel(KeyT *keys_in,\n KeyT *keys_out,\n ValueT *values_in,\n ValueT *values_out,\n SegmentSizeItT segment_sizes,\n KeyT oob_default,\n ActionT action)\n{\n using warp_merge_sort_t =\n hipcub::WarpMergeSort;\n using storage_t = typename warp_merge_sort_t::TempStorage;\n\n // Get linear thread and warp index\n const int tid = hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n const int warp_id = tid / LOGICAL_WARP_THREADS;\n\n // Test case of partially finished CTA\n if (warp_id >= TOTAL_WARPS)\n {\n return;\n }\n\n // Thread-local storage & warp-scope temporary storage allocation\n KeyT keys[ITEMS_PER_THREAD];\n ValueT values[ITEMS_PER_THREAD];\n __shared__ storage_t storage[TOTAL_WARPS];\n\n // Instantiate warp-scope algorithm\n warp_merge_sort_t warp_sort(storage[warp_id]);\n\n const int warp_offset = LOGICAL_WARP_THREADS * ITEMS_PER_THREAD * warp_id;\n const int thread_offset = warp_offset + warp_sort.get_linear_tid() * ITEMS_PER_THREAD;\n const int valid_items = segment_sizes[warp_id];\n\n // Load data\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n const int idx = thread_offset + item;\n keys[item] = keys_in[idx];\n values[item] = values_in[idx];\n }\n hipcub::WARP_SYNC(warp_sort.get_member_mask());\n\n // Run merge sort test\n action(warp_sort, keys, values, valid_items, oob_default);\n\n // Store data\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n const int idx = thread_offset + item;\n keys_out[idx] = (idx - warp_offset) >= valid_items ? oob_default : keys[item];\n values_out[idx] = (idx - warp_offset) >= valid_items ? ValueT{} : values[item];\n }\n}\n\n// -----------------------------------------------------------\n// Dimensions being instantiated:\n// {full,partial} x {stable, 'unstable'} x {keys, kv-pairs}\n// -----------------------------------------------------------\n\n/**\n * @brief Delegate wrapper for WarpMergeSort::StableSort on keys-only\n */\nstruct warp_stable_sort_keys_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n T (&thread_data)[ITEMS_PER_THREAD],\n int /*valid_items*/,\n T /*oob_default*/) const\n {\n warp_sort.StableSort(thread_data, CustomLess{});\n }\n};\n\n/**\n * @brief Delegate wrapper for partial WarpMergeSort::StableSort keys-only\n */\nstruct warp_partial_stable_sort_keys_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n T (&thread_data)[ITEMS_PER_THREAD],\n int valid_items,\n T oob_default) const\n {\n warp_sort.StableSort(thread_data, CustomLess{}, valid_items, oob_default);\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpMergeSort::Sort on keys-only\n */\nstruct warp_sort_keys_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n T (&thread_data)[ITEMS_PER_THREAD],\n int /*valid_items*/,\n T /*oob_default*/) const\n {\n warp_sort.Sort(thread_data, CustomLess{});\n }\n};\n\n/**\n * @brief Delegate wrapper for partial WarpMergeSort::StableSort keys-only\n */\nstruct warp_partial_sort_keys_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n T (&thread_data)[ITEMS_PER_THREAD],\n int valid_items,\n T oob_default) const\n {\n warp_sort.Sort(thread_data, CustomLess{}, valid_items, oob_default);\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpMergeSort::StableSort on key-value pairs\n */\nstruct warp_stable_sort_pairs_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n KeyT (&keys)[ITEMS_PER_THREAD],\n ValueT (&values)[ITEMS_PER_THREAD],\n int /*valid_items*/,\n KeyT /*oob_default*/) const\n {\n warp_sort.StableSort(keys, values, CustomLess{});\n }\n};\n\n/**\n * @brief Delegate wrapper for partial WarpMergeSort::StableSort key-value pairs\n */\nstruct warp_partial_stable_sort_pairs_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n KeyT (&keys)[ITEMS_PER_THREAD],\n ValueT (&values)[ITEMS_PER_THREAD],\n int valid_items,\n KeyT oob_default) const\n {\n warp_sort.StableSort(keys, values, CustomLess{}, valid_items, oob_default);\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpMergeSort::Sort on key-value pairs\n */\nstruct warp_sort_pairs_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n KeyT (&keys)[ITEMS_PER_THREAD],\n ValueT (&values)[ITEMS_PER_THREAD],\n int /*valid_items*/,\n KeyT /*oob_default*/) const\n {\n warp_sort.Sort(keys, values, CustomLess{});\n }\n};\n\n/**\n * @brief Delegate wrapper for partial WarpMergeSort::StableSort key-value pairs\n */\nstruct warp_partial_sort_pairs_t\n{\n template \n __device__ void operator()(WarpSortT &warp_sort,\n KeyT (&keys)[ITEMS_PER_THREAD],\n ValueT (&values)[ITEMS_PER_THREAD],\n int valid_items,\n KeyT oob_default) const\n {\n warp_sort.Sort(keys, values, CustomLess{}, valid_items, oob_default);\n }\n};\n\n/**\n * @brief Dispatch helper function for sorting keys\n */\ntemplate \nvoid warp_merge_sort(thrust::device_vector &in,\n thrust::device_vector &out,\n SegmentSizesItT segment_sizes,\n T oob_default,\n ActionT action)\n{\n warp_merge_sort_kernel\n <<<1, LOGICAL_WARP_THREADS * TOTAL_WARPS>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(out.data()),\n segment_sizes,\n oob_default,\n action);\n\n REQUIRE(hipSuccess == hipPeekAtLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n}\n\n/**\n * @brief Dispatch helper function for sorting key-value pairs\n */\ntemplate \nvoid warp_merge_sort(thrust::device_vector &keys_in,\n thrust::device_vector &keys_out,\n thrust::device_vector &values_in,\n thrust::device_vector &values_out,\n SegmentSizesItT segment_sizes,\n KeyT oob_default,\n ActionT action)\n{\n warp_merge_sort_kernel\n <<<1, LOGICAL_WARP_THREADS * TOTAL_WARPS>>>(thrust::raw_pointer_cast(keys_in.data()),\n thrust::raw_pointer_cast(keys_out.data()),\n thrust::raw_pointer_cast(values_in.data()),\n thrust::raw_pointer_cast(values_out.data()),\n segment_sizes,\n oob_default,\n action);\n\n REQUIRE(hipSuccess == hipPeekAtLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n}\n\n/**\n * @brief Performs a stable sort on per-warp segments of data and assigns oob_default to items that\n * are out-of-bounds.\n */\ntemplate \nvoid compute_host_reference(RandomItT h_data,\n SegmentSizeItT segment_sizes,\n unsigned int num_segments,\n T oob_default,\n int logical_warp_items)\n{\n for (unsigned int segment_id = 0; segment_id < num_segments; segment_id++)\n {\n unsigned int segment_size = segment_sizes[segment_id];\n std::stable_sort(h_data, h_data + segment_size);\n std::fill(h_data + segment_size, h_data + logical_warp_items, oob_default);\n h_data += logical_warp_items;\n }\n}\n\n/**\n * @brief Stability requirement of the sorting algorithm\n */\nenum class stability\n{\n stable,\n unstable\n};\n\n// List of key types to test\nusing custom_t =\n c2h::custom_type_t;\nusing key_types = c2h::type_list;\n\n// List of value types\nusing value_types = c2h::type_list;\n\n// Logical warp sizes to test\nusing logical_warp_threads = c2h::enum_type_list;\n\n// Number of items per thread to test\nusing items_per_thread_list = c2h::enum_type_list;\n\n// Whether the sort is required to be stable or not\nusing stability_list = c2h::enum_type_list;\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int logical_warp_threads = c2h::get<1, TestType>::value;\n static constexpr int items_per_thread = c2h::get<2, TestType>::value;\n static constexpr int logical_warp_items = logical_warp_threads * items_per_thread;\n static constexpr int total_warps = 2;\n static constexpr int tile_size = items_per_thread * total_warps * logical_warp_threads;\n static constexpr bool is_stable = c2h::get<3, TestType>::value == stability::stable;\n};\n\nCUB_TEST(\"Warp sort on keys-only works\",\n \"[sort][warp]\",\n key_types,\n logical_warp_threads,\n items_per_thread_list,\n stability_list)\n{\n using params = params_t;\n using type = typename params::type;\n using warp_sort_delegate =\n hipcub::detail::conditional_t;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n auto segment_sizes = thrust::make_constant_iterator(params::logical_warp_items);\n const auto oob_default = std::numeric_limits::max();\n c2h::gen(CUB_SEED(10), d_in);\n\n // Run test\n warp_merge_sort(\n d_in,\n d_out,\n segment_sizes,\n oob_default,\n warp_sort_delegate{});\n\n // Prepare verification data\n thrust::host_vector h_in_out = d_in;\n compute_host_reference(h_in_out.begin(),\n segment_sizes,\n params::total_warps,\n oob_default,\n params::logical_warp_items);\n\n // Verify results\n REQUIRE(h_in_out == d_out);\n}\n\nCUB_TEST(\"Warp sort keys-only on partial warp-tile works\",\n \"[sort][warp]\",\n key_types,\n logical_warp_threads,\n items_per_thread_list,\n stability_list)\n{\n using params = params_t;\n using type = typename params::type;\n using warp_sort_delegate = hipcub::detail::\n conditional_t;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_segment_sizes(params::total_warps);\n const auto oob_default = std::numeric_limits::max();\n c2h::gen(CUB_SEED(5), d_in);\n c2h::gen(CUB_SEED(5), d_segment_sizes, 0, params::logical_warp_items);\n\n // Run test\n warp_merge_sort(\n d_in,\n d_out,\n d_segment_sizes.cbegin(),\n oob_default,\n warp_sort_delegate{});\n\n // Prepare verification data\n thrust::host_vector h_in_out = d_in;\n thrust::host_vector segment_sizes = d_segment_sizes;\n compute_host_reference(h_in_out.begin(),\n segment_sizes,\n params::total_warps,\n oob_default,\n params::logical_warp_items);\n\n // Verify results\n REQUIRE(h_in_out == d_out);\n}\n\nCUB_TEST(\"Warp sort on keys-value pairs works\",\n \"[sort][warp]\",\n key_types,\n logical_warp_threads,\n items_per_thread_list,\n stability_list,\n value_types)\n{\n using params = params_t;\n using key_type = typename params::type;\n using value_type = typename c2h::get<4, TestType>;\n using warp_sort_delegate =\n hipcub::detail::conditional_t;\n\n // Prepare test data\n thrust::device_vector d_keys_in(params::tile_size);\n thrust::device_vector d_keys_out(params::tile_size);\n thrust::device_vector d_values_in(params::tile_size);\n thrust::device_vector d_values_out(params::tile_size);\n auto segment_sizes = thrust::make_constant_iterator(params::logical_warp_items);\n const auto oob_default = std::numeric_limits::max();\n c2h::gen(CUB_SEED(10), d_keys_in);\n\n // Run test\n warp_merge_sort(\n d_keys_in,\n d_keys_out,\n d_values_in,\n d_values_out,\n segment_sizes,\n oob_default,\n warp_stable_sort_pairs_t{});\n\n // Prepare verification data\n thrust::host_vector h_keys_in_out = d_keys_in;\n thrust::host_vector h_values_in_out = d_values_in;\n auto cpu_kv_pairs = thrust::make_zip_iterator(h_keys_in_out.begin(), h_values_in_out.begin());\n compute_host_reference(cpu_kv_pairs,\n segment_sizes,\n params::total_warps,\n thrust::make_tuple(oob_default, value_type{}),\n params::logical_warp_items);\n\n // Verify results\n REQUIRE(h_keys_in_out == d_keys_out);\n REQUIRE(h_values_in_out == d_values_out);\n}\n\nCUB_TEST(\"Warp sort on key-value pairs of a partial warp-tile works\",\n \"[sort][warp]\",\n key_types,\n logical_warp_threads,\n items_per_thread_list,\n stability_list,\n value_types)\n{\n using params = params_t;\n using key_type = typename params::type;\n using value_type = typename c2h::get<4, TestType>;\n using warp_sort_delegate = hipcub::detail::\n conditional_t;\n\n // Prepare test data\n thrust::device_vector d_keys_in(params::tile_size);\n thrust::device_vector d_keys_out(params::tile_size);\n thrust::device_vector d_values_in(params::tile_size);\n thrust::device_vector d_values_out(params::tile_size);\n thrust::device_vector d_segment_sizes(params::total_warps);\n const auto oob_default = std::numeric_limits::max();\n c2h::gen(CUB_SEED(5), d_keys_in);\n c2h::gen(CUB_SEED(5), d_segment_sizes, 0, params::logical_warp_items);\n\n // Run test\n warp_merge_sort(\n d_keys_in,\n d_keys_out,\n d_values_in,\n d_values_out,\n d_segment_sizes.cbegin(),\n oob_default,\n warp_sort_delegate{});\n\n // Prepare verification data\n thrust::host_vector h_keys_in_out = d_keys_in;\n thrust::host_vector h_values_in_out = d_values_in;\n thrust::host_vector segment_sizes = d_segment_sizes;\n auto cpu_kv_pairs = thrust::make_zip_iterator(h_keys_in_out.begin(), h_values_in_out.begin());\n compute_host_reference(cpu_kv_pairs,\n segment_sizes,\n params::total_warps,\n thrust::make_tuple(oob_default, value_type{}),\n params::logical_warp_items);\n\n // Verify results\n REQUIRE(h_keys_in_out == d_keys_out);\n REQUIRE(h_values_in_out == d_values_out);\n}\n", "source": "/kaggle/working/cub/test/catch2_test_warp_merge_sort.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"fill_striped.cuh\"\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void warp_load_kernel(InputIteratorT input_iterator, ActionT action, int *error_counter)\n{\n using warp_load_t = cub::WarpLoad;\n using storage_t = typename warp_load_t::TempStorage;\n\n constexpr int tile_size = ITEMS_PER_THREAD * LOGICAL_WARP_THREADS;\n\n __shared__ storage_t storage[TOTAL_WARPS];\n\n const int linear_tid = threadIdx.x;\n\n const int warp_id = linear_tid / LOGICAL_WARP_THREADS;\n warp_load_t load(storage[warp_id]);\n\n // Test WarpLoad specialization\n T reg[ITEMS_PER_THREAD];\n action.load(load, input_iterator + (warp_id * tile_size), reg);\n\n // Verify data was loaded as expected\n action.verify(reg, error_counter);\n}\n\ntemplate \nvoid warp_load(InputIteratorT input_iterator, ActionT action, int *error_counter)\n{\n warp_load_kernel\n <<<1, TOTAL_WARPS * LOGICAL_WARP_THREADS>>>(input_iterator, action, error_counter);\n REQUIRE(cudaSuccess == cudaPeekAtLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n}\n\n/**\n * @brief WarpLoad test specialisation for guarded loads\n */\ntemplate \nstruct guarded_load_t\n{\n int valid_items;\n T oob_default;\n\n template \n __device__ void load(cub::WarpLoad load,\n InputIteratorT input,\n T (®)[ITEMS_PER_THREAD])\n {\n load.Load(input, reg, valid_items, oob_default);\n }\n\n template \n __device__ void verify(T (®)[ITEMS_PER_THREAD], int *error_counter)\n {\n const auto linear_tid = cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n const auto lane_id = linear_tid % LOGICAL_WARP_THREADS;\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n const auto expected_value = static_cast(linear_tid * ITEMS_PER_THREAD + item);\n\n const bool is_oob = LoadAlgorithm == cub::WarpLoadAlgorithm::WARP_LOAD_STRIPED\n ? item * LOGICAL_WARP_THREADS + lane_id >= valid_items\n : lane_id * ITEMS_PER_THREAD + item >= valid_items;\n\n if (is_oob)\n {\n if (reg[item] != oob_default)\n {\n atomicAdd(error_counter, 1);\n }\n }\n else if (reg[item] != expected_value)\n {\n atomicAdd(error_counter, 1);\n }\n }\n }\n};\n\n/**\n * @brief WarpLoad test specialisation for unguarded loads\n */\nstruct unguarded_load_t\n{\n template \n __device__ void load(cub::WarpLoad load,\n InputIteratorT input,\n T (®)[ITEMS_PER_THREAD])\n {\n load.Load(input, reg);\n }\n\n template \n __device__ void verify(T (®)[ITEMS_PER_THREAD], int *error_counter)\n {\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n const auto expected_value = static_cast(threadIdx.x * ITEMS_PER_THREAD + item);\n\n if (reg[item] != expected_value)\n {\n atomicAdd(error_counter, 1);\n }\n }\n }\n};\n\ntemplate \nthrust::device_vector generate_input()\n{\n const int tile_size = LOGICAL_WARP_THREADS * ITEMS_PER_THREAD;\n const int num_items = TOTAL_WARPS * tile_size;\n\n thrust::device_vector d_input(num_items);\n\n if (LoadAlgorithm == cub::WarpLoadAlgorithm::WARP_LOAD_STRIPED)\n {\n thrust::host_vector h_input(num_items);\n\n // In this case we need different stripe pattern, so the\n // items/threads parameters are swapped\n\n constexpr int FAKE_BLOCK_SIZE = ITEMS_PER_THREAD * TOTAL_WARPS;\n\n fill_striped(h_input.begin());\n d_input = h_input;\n }\n else\n {\n c2h::gen(c2h::modulo_t{num_items}, d_input);\n }\n\n return d_input;\n}\n\n// %PARAM% LWT lwt 4:16:32\n// %PARAM% ALGO_TYPE alg 0:1:2:3\n\nusing types = c2h::type_list;\nusing items_per_thread = c2h::enum_type_list;\nusing logical_warp_threads = c2h::enum_type_list;\nusing algorithms = c2h::enum_type_list;\nusing algorithm =\n c2h::enum_type_list::value>;\n\nusing cache_load_modifier = c2h::enum_type_list;\n\nconstexpr int guarded_load_tests_count = 30;\n\ntemplate \nstruct total_warps_t\n{\nprivate:\n static constexpr int max_warps = 2;\n static constexpr bool is_arch_warp = (logical_warp_threads == CUB_WARP_THREADS(0));\n static constexpr bool is_pow_of_two = ((logical_warp_threads & (logical_warp_threads - 1)) == 0);\n static constexpr int total_warps = (is_arch_warp || is_pow_of_two) ? max_warps : 1;\n\npublic:\n static constexpr int value() { return total_warps; }\n};\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int logical_warp_threads = c2h::get<1, TestType>::value;\n static constexpr int items_per_thread = c2h::get<2, TestType>::value;\n static constexpr cub::WarpLoadAlgorithm algorithm = c2h::get<3, TestType>::value;\n static constexpr int total_warps = total_warps_t::value();\n static constexpr int tile_size = logical_warp_threads * items_per_thread;\n static constexpr int total_item_count = total_warps * tile_size;\n};\n\nCUB_TEST(\"Warp load guarded range works with pointer\",\n \"[load][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n using delegate_t = guarded_load_t;\n\n const int valid_items =\n GENERATE_COPY(take(guarded_load_tests_count, random(0, params::tile_size - 1)));\n const auto oob_default = static_cast(valid_items);\n\n auto d_in = generate_input();\n thrust::device_vector d_error_counter(1, 0);\n\n warp_load(thrust::raw_pointer_cast(d_in.data()),\n delegate_t{valid_items, oob_default},\n thrust::raw_pointer_cast(d_error_counter.data()));\n\n const int num_errors = d_error_counter[0];\n const int expected_error_count = 0;\n REQUIRE(num_errors == expected_error_count);\n}\n\nCUB_TEST(\"Warp load guarded range works with cache modified iterator\",\n \"[load][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm,\n cache_load_modifier)\n{\n using params = params_t;\n using type = typename params::type;\n using delegate_t = guarded_load_t;\n constexpr cub::CacheLoadModifier load_modifier = c2h::get<4, TestType>::value;\n\n const int valid_items =\n GENERATE_COPY(take(guarded_load_tests_count, random(0, params::tile_size - 1)));\n const auto oob_default = static_cast(valid_items);\n\n auto d_in = generate_input();\n auto in_it =\n cub::CacheModifiedInputIterator(thrust::raw_pointer_cast(d_in.data()));\n thrust::device_vector d_error_counter(1, 0);\n\n warp_load(in_it,\n delegate_t{valid_items, oob_default},\n thrust::raw_pointer_cast(d_error_counter.data()));\n\n const auto num_errors = d_error_counter[0];\n const int expected_error_count = 0;\n REQUIRE(num_errors == expected_error_count);\n}\n\nCUB_TEST(\"Warp load unguarded range works with pointer\",\n \"[load][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n using delegate_t = unguarded_load_t;\n\n auto d_in = generate_input();\n thrust::device_vector d_error_counter(1, 0);\n\n warp_load(thrust::raw_pointer_cast(d_in.data()),\n delegate_t{},\n thrust::raw_pointer_cast(d_error_counter.data()));\n\n const auto num_errors = d_error_counter[0];\n const int expected_error_count = 0;\n REQUIRE(num_errors == expected_error_count);\n}\n\nCUB_TEST(\"Warp load unguarded range works with cache modified iterator\",\n \"[load][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm,\n cache_load_modifier)\n{\n using params = params_t;\n using type = typename params::type;\n using delegate_t = unguarded_load_t;\n constexpr cub::CacheLoadModifier load_modifier = c2h::get<4, TestType>::value;\n\n auto d_in = generate_input();\n auto in_it =\n cub::CacheModifiedInputIterator(thrust::raw_pointer_cast(d_in.data()));\n thrust::device_vector d_error_counter(1, 0);\n\n warp_load(in_it, delegate_t{}, thrust::raw_pointer_cast(d_error_counter.data()));\n\n const auto num_errors = d_error_counter[0];\n const int expected_error_count = 0;\n REQUIRE(num_errors == expected_error_count);\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"fill_striped.cuh\"\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void warp_load_kernel(InputIteratorT input_iterator, ActionT action, int *error_counter)\n{\n using warp_load_t = hipcub::WarpLoad;\n using storage_t = typename warp_load_t::TempStorage;\n\n constexpr int tile_size = ITEMS_PER_THREAD * LOGICAL_WARP_THREADS;\n\n __shared__ storage_t storage[TOTAL_WARPS];\n\n const int linear_tid = threadIdx.x;\n\n const int warp_id = linear_tid / LOGICAL_WARP_THREADS;\n warp_load_t load(storage[warp_id]);\n\n // Test WarpLoad specialization\n T reg[ITEMS_PER_THREAD];\n action.load(load, input_iterator + (warp_id * tile_size), reg);\n\n // Verify data was loaded as expected\n action.verify(reg, error_counter);\n}\n\ntemplate \nvoid warp_load(InputIteratorT input_iterator, ActionT action, int *error_counter)\n{\n warp_load_kernel\n <<<1, TOTAL_WARPS * LOGICAL_WARP_THREADS>>>(input_iterator, action, error_counter);\n REQUIRE(hipSuccess == hipPeekAtLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n}\n\n/**\n * @brief WarpLoad test specialisation for guarded loads\n */\ntemplate \nstruct guarded_load_t\n{\n int valid_items;\n T oob_default;\n\n template \n __device__ void load(hipcub::WarpLoad load,\n InputIteratorT input,\n T (®)[ITEMS_PER_THREAD])\n {\n load.Load(input, reg, valid_items, oob_default);\n }\n\n template \n __device__ void verify(T (®)[ITEMS_PER_THREAD], int *error_counter)\n {\n const auto linear_tid = hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n const auto lane_id = linear_tid % LOGICAL_WARP_THREADS;\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n const auto expected_value = static_cast(linear_tid * ITEMS_PER_THREAD + item);\n\n const bool is_oob = LoadAlgorithm == hipcub::WarpLoadAlgorithm::WARP_LOAD_STRIPED\n ? item * LOGICAL_WARP_THREADS + lane_id >= valid_items\n : lane_id * ITEMS_PER_THREAD + item >= valid_items;\n\n if (is_oob)\n {\n if (reg[item] != oob_default)\n {\n atomicAdd(error_counter, 1);\n }\n }\n else if (reg[item] != expected_value)\n {\n atomicAdd(error_counter, 1);\n }\n }\n }\n};\n\n/**\n * @brief WarpLoad test specialisation for unguarded loads\n */\nstruct unguarded_load_t\n{\n template \n __device__ void load(hipcub::WarpLoad load,\n InputIteratorT input,\n T (®)[ITEMS_PER_THREAD])\n {\n load.Load(input, reg);\n }\n\n template \n __device__ void verify(T (®)[ITEMS_PER_THREAD], int *error_counter)\n {\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n const auto expected_value = static_cast(threadIdx.x * ITEMS_PER_THREAD + item);\n\n if (reg[item] != expected_value)\n {\n atomicAdd(error_counter, 1);\n }\n }\n }\n};\n\ntemplate \nthrust::device_vector generate_input()\n{\n const int tile_size = LOGICAL_WARP_THREADS * ITEMS_PER_THREAD;\n const int num_items = TOTAL_WARPS * tile_size;\n\n thrust::device_vector d_input(num_items);\n\n if (LoadAlgorithm == hipcub::WarpLoadAlgorithm::WARP_LOAD_STRIPED)\n {\n thrust::host_vector h_input(num_items);\n\n // In this case we need different stripe pattern, so the\n // items/threads parameters are swapped\n\n constexpr int FAKE_BLOCK_SIZE = ITEMS_PER_THREAD * TOTAL_WARPS;\n\n fill_striped(h_input.begin());\n d_input = h_input;\n }\n else\n {\n c2h::gen(c2h::modulo_t{num_items}, d_input);\n }\n\n return d_input;\n}\n\n// %PARAM% LWT lwt 4:16:32\n// %PARAM% ALGO_TYPE alg 0:1:2:3\n\nusing types = c2h::type_list;\nusing items_per_thread = c2h::enum_type_list;\nusing logical_warp_threads = c2h::enum_type_list;\nusing algorithms = c2h::enum_type_list;\nusing algorithm =\n c2h::enum_type_list::value>;\n\nusing cache_load_modifier = c2h::enum_type_list;\n\nconstexpr int guarded_load_tests_count = 30;\n\ntemplate \nstruct total_warps_t\n{\nprivate:\n static constexpr int max_warps = 2;\n static constexpr bool is_arch_warp = (logical_warp_threads == CUB_WARP_THREADS(0));\n static constexpr bool is_pow_of_two = ((logical_warp_threads & (logical_warp_threads - 1)) == 0);\n static constexpr int total_warps = (is_arch_warp || is_pow_of_two) ? max_warps : 1;\n\npublic:\n static constexpr int value() { return total_warps; }\n};\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int logical_warp_threads = c2h::get<1, TestType>::value;\n static constexpr int items_per_thread = c2h::get<2, TestType>::value;\n static constexpr hipcub::WarpLoadAlgorithm algorithm = c2h::get<3, TestType>::value;\n static constexpr int total_warps = total_warps_t::value();\n static constexpr int tile_size = logical_warp_threads * items_per_thread;\n static constexpr int total_item_count = total_warps * tile_size;\n};\n\nCUB_TEST(\"Warp load guarded range works with pointer\",\n \"[load][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n using delegate_t = guarded_load_t;\n\n const int valid_items =\n GENERATE_COPY(take(guarded_load_tests_count, random(0, params::tile_size - 1)));\n const auto oob_default = static_cast(valid_items);\n\n auto d_in = generate_input();\n thrust::device_vector d_error_counter(1, 0);\n\n warp_load(thrust::raw_pointer_cast(d_in.data()),\n delegate_t{valid_items, oob_default},\n thrust::raw_pointer_cast(d_error_counter.data()));\n\n const int num_errors = d_error_counter[0];\n const int expected_error_count = 0;\n REQUIRE(num_errors == expected_error_count);\n}\n\nCUB_TEST(\"Warp load guarded range works with cache modified iterator\",\n \"[load][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm,\n cache_load_modifier)\n{\n using params = params_t;\n using type = typename params::type;\n using delegate_t = guarded_load_t;\n constexpr hipcub::CacheLoadModifier load_modifier = c2h::get<4, TestType>::value;\n\n const int valid_items =\n GENERATE_COPY(take(guarded_load_tests_count, random(0, params::tile_size - 1)));\n const auto oob_default = static_cast(valid_items);\n\n auto d_in = generate_input();\n auto in_it =\n hipcub::CacheModifiedInputIterator(thrust::raw_pointer_cast(d_in.data()));\n thrust::device_vector d_error_counter(1, 0);\n\n warp_load(in_it,\n delegate_t{valid_items, oob_default},\n thrust::raw_pointer_cast(d_error_counter.data()));\n\n const auto num_errors = d_error_counter[0];\n const int expected_error_count = 0;\n REQUIRE(num_errors == expected_error_count);\n}\n\nCUB_TEST(\"Warp load unguarded range works with pointer\",\n \"[load][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n using delegate_t = unguarded_load_t;\n\n auto d_in = generate_input();\n thrust::device_vector d_error_counter(1, 0);\n\n warp_load(thrust::raw_pointer_cast(d_in.data()),\n delegate_t{},\n thrust::raw_pointer_cast(d_error_counter.data()));\n\n const auto num_errors = d_error_counter[0];\n const int expected_error_count = 0;\n REQUIRE(num_errors == expected_error_count);\n}\n\nCUB_TEST(\"Warp load unguarded range works with cache modified iterator\",\n \"[load][warp]\",\n types,\n logical_warp_threads,\n items_per_thread,\n algorithm,\n cache_load_modifier)\n{\n using params = params_t;\n using type = typename params::type;\n using delegate_t = unguarded_load_t;\n constexpr hipcub::CacheLoadModifier load_modifier = c2h::get<4, TestType>::value;\n\n auto d_in = generate_input();\n auto in_it =\n hipcub::CacheModifiedInputIterator(thrust::raw_pointer_cast(d_in.data()));\n thrust::device_vector d_error_counter(1, 0);\n\n warp_load(in_it, delegate_t{}, thrust::raw_pointer_cast(d_error_counter.data()));\n\n const auto num_errors = d_error_counter[0];\n const int expected_error_count = 0;\n REQUIRE(num_errors == expected_error_count);\n}\n", "source": "/kaggle/working/cub/test/catch2_test_warp_load.cu"} {"input": "/*******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"test_util.h\"\n#include \"cub/thread/thread_sort.cuh\"\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n\nstruct CustomLess\n{\n template \n __host__ __device__ bool operator()(DataType &lhs, DataType &rhs)\n {\n return lhs < rhs;\n }\n};\n\n\ntemplate \n__global__ void kernel(const KeyT *keys_in,\n KeyT *keys_out,\n const ValueT *values_in,\n ValueT *values_out)\n{\n KeyT thread_keys[ItemsPerThread];\n KeyT thread_values[ItemsPerThread];\n\n const auto thread_offset = ItemsPerThread * threadIdx.x;\n keys_in += thread_offset;\n keys_out += thread_offset;\n values_in += thread_offset;\n values_out += thread_offset;\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n thread_keys[item] = keys_in[item];\n thread_values[item] = values_in[item];\n }\n\n cub::StableOddEvenSort(thread_keys, thread_values, CustomLess{});\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n keys_out[item] = thread_keys[item];\n values_out[item] = thread_values[item];\n }\n}\n\n\ntemplate \nvoid Test()\n{\n const unsigned int threads_in_block = 1024;\n const unsigned int elements = threads_in_block * ItemsPerThread;\n\n thrust::default_random_engine re;\n thrust::device_vector data_source(elements);\n\n for (int iteration = 0; iteration < 10; iteration++)\n {\n thrust::sequence(data_source.begin(), data_source.end());\n thrust::shuffle(data_source.begin(), data_source.end(), re);\n thrust::device_vector in_keys(data_source);\n thrust::device_vector out_keys(elements);\n\n thrust::shuffle(data_source.begin(), data_source.end(), re);\n thrust::device_vector in_values(data_source);\n thrust::device_vector out_values(elements);\n\n thrust::host_vector host_keys(in_keys);\n thrust::host_vector host_values(in_values);\n\n kernel<<<1, threads_in_block>>>(\n thrust::raw_pointer_cast(in_keys.data()),\n thrust::raw_pointer_cast(out_keys.data()),\n thrust::raw_pointer_cast(in_values.data()),\n thrust::raw_pointer_cast(out_values.data()));\n\n for (unsigned int tid = 0; tid < threads_in_block; tid++)\n {\n const auto thread_begin = tid * ItemsPerThread;\n const auto thread_end = thread_begin + ItemsPerThread;\n\n thrust::sort_by_key(host_keys.begin() + thread_begin,\n host_keys.begin() + thread_end,\n host_values.begin() + thread_begin,\n CustomLess{});\n }\n\n AssertEquals(host_keys, out_keys);\n AssertEquals(host_values, out_values);\n }\n}\n\n\ntemplate \nvoid Test()\n{\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n}\n\nint main()\n{\n Test();\n Test();\n\n return 0;\n}\n", "output": "/*******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \"test_util.h\"\n#include \"cub/thread/thread_sort.cuh\"\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n\nstruct CustomLess\n{\n template \n __host__ __device__ bool operator()(DataType &lhs, DataType &rhs)\n {\n return lhs < rhs;\n }\n};\n\n\ntemplate \n__global__ void kernel(const KeyT *keys_in,\n KeyT *keys_out,\n const ValueT *values_in,\n ValueT *values_out)\n{\n KeyT thread_keys[ItemsPerThread];\n KeyT thread_values[ItemsPerThread];\n\n const auto thread_offset = ItemsPerThread * threadIdx.x;\n keys_in += thread_offset;\n keys_out += thread_offset;\n values_in += thread_offset;\n values_out += thread_offset;\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n thread_keys[item] = keys_in[item];\n thread_values[item] = values_in[item];\n }\n\n hipcub::StableOddEvenSort(thread_keys, thread_values, CustomLess{});\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n keys_out[item] = thread_keys[item];\n values_out[item] = thread_values[item];\n }\n}\n\n\ntemplate \nvoid Test()\n{\n const unsigned int threads_in_block = 1024;\n const unsigned int elements = threads_in_block * ItemsPerThread;\n\n thrust::default_random_engine re;\n thrust::device_vector data_source(elements);\n\n for (int iteration = 0; iteration < 10; iteration++)\n {\n thrust::sequence(data_source.begin(), data_source.end());\n thrust::shuffle(data_source.begin(), data_source.end(), re);\n thrust::device_vector in_keys(data_source);\n thrust::device_vector out_keys(elements);\n\n thrust::shuffle(data_source.begin(), data_source.end(), re);\n thrust::device_vector in_values(data_source);\n thrust::device_vector out_values(elements);\n\n thrust::host_vector host_keys(in_keys);\n thrust::host_vector host_values(in_values);\n\n kernel<<<1, threads_in_block>>>(\n thrust::raw_pointer_cast(in_keys.data()),\n thrust::raw_pointer_cast(out_keys.data()),\n thrust::raw_pointer_cast(in_values.data()),\n thrust::raw_pointer_cast(out_values.data()));\n\n for (unsigned int tid = 0; tid < threads_in_block; tid++)\n {\n const auto thread_begin = tid * ItemsPerThread;\n const auto thread_end = thread_begin + ItemsPerThread;\n\n thrust::sort_by_key(host_keys.begin() + thread_begin,\n host_keys.begin() + thread_end,\n host_values.begin() + thread_begin,\n CustomLess{});\n }\n\n AssertEquals(host_keys, out_keys);\n AssertEquals(host_values, out_values);\n }\n}\n\n\ntemplate \nvoid Test()\n{\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n}\n\nint main()\n{\n Test();\n Test();\n\n return 0;\n}\n", "source": "/kaggle/working/cub/test/test_thread_sort.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n#include \n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \nstruct output_idx\n{\n static __device__ int get(int item)\n {\n return static_cast(threadIdx.x) * ItemsPerThread + item;\n }\n};\n\ntemplate \nstruct output_idx\n{\n static __device__ int get(int item)\n {\n return static_cast(threadIdx.x) + ThreadsInBlock * item;\n }\n};\n\ntemplate \n__global__ void kernel(std::integral_constant,\n InputIteratorT input,\n OutputIteratorT output,\n int num_items)\n{\n using input_t = cub::detail::value_t;\n using block_load_t = cub::BlockLoad;\n using storage_t =\n typename block_load_t::TempStorage;\n\n __shared__ storage_t storage;\n block_load_t block_load(storage);\n\n input_t data[ItemsPerThread];\n\n if (ItemsPerThread * ThreadsInBlock == num_items)\n {\n block_load.Load(input, data);\n }\n else\n {\n block_load.Load(input, data, num_items);\n }\n\n for (int i = 0; i < ItemsPerThread; i++)\n {\n const int idx =\n output_idx::get(i);\n\n if (idx < num_items)\n {\n output[idx] = data[i];\n }\n }\n}\n\ntemplate \n__global__ void kernel(std::integral_constant,\n InputIteratorT input,\n OutputIteratorT output,\n int num_items)\n{\n for (int i = 0; i < ItemsPerThread; i++)\n {\n const int idx =\n output_idx::get(i);\n\n if (idx < num_items)\n {\n output[idx] = input[idx];\n }\n }\n}\n\ntemplate \nvoid block_load(\n InputIteratorT input,\n OutputIteratorT output,\n int num_items)\n{\n using input_t = cub::detail::value_t;\n using block_load_t = cub::BlockLoad;\n using storage_t = typename block_load_t::TempStorage;\n constexpr bool sufficient_resources = sizeof(storage_t) <= 1024 * 48;\n\n kernel<<<1, ThreadsInBlock>>>(\n std::integral_constant{},\n input, output, num_items);\n\n REQUIRE( cudaSuccess == cudaPeekAtLastError() );\n REQUIRE( cudaSuccess == cudaDeviceSynchronize() );\n}\n\n\n// %PARAM% IPT it 1:11\n\nusing types = c2h::type_list;\nusing vec_types = c2h::type_list;\n\nusing even_threads_in_block = c2h::enum_type_list;\nusing odd_threads_in_block = c2h::enum_type_list;\nusing a_block_size = c2h::enum_type_list;\n\nusing items_per_thread = c2h::enum_type_list;\nusing load_algorithm = c2h::enum_type_list<\n cub::BlockLoadAlgorithm,\n cub::BlockLoadAlgorithm::BLOCK_LOAD_DIRECT,\n cub::BlockLoadAlgorithm::BLOCK_LOAD_STRIPED,\n cub::BlockLoadAlgorithm::BLOCK_LOAD_VECTORIZE,\n cub::BlockLoadAlgorithm::BLOCK_LOAD_TRANSPOSE,\n cub::BlockLoadAlgorithm::BLOCK_LOAD_WARP_TRANSPOSE, \n cub::BlockLoadAlgorithm::BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED>;\n\nusing odd_load_algorithm = c2h::enum_type_list<\n cub::BlockLoadAlgorithm,\n cub::BlockLoadAlgorithm::BLOCK_LOAD_DIRECT,\n cub::BlockLoadAlgorithm::BLOCK_LOAD_STRIPED,\n cub::BlockLoadAlgorithm::BLOCK_LOAD_VECTORIZE,\n cub::BlockLoadAlgorithm::BLOCK_LOAD_TRANSPOSE>;\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int items_per_thread = c2h::get<1, TestType>::value;\n static constexpr int threads_in_block = c2h::get<2, TestType>::value;\n static constexpr int tile_size = items_per_thread * threads_in_block;\n static constexpr cub::BlockLoadAlgorithm load_algorithm = \n c2h::get<3, TestType>::value;\n};\n\nCUB_TEST(\"Block load works with even block sizes\",\n \"[load][block]\",\n types,\n items_per_thread,\n even_threads_in_block,\n load_algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_load(thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\nCUB_TEST(\"Block load works with even odd sizes\",\n \"[load][block]\",\n types,\n items_per_thread,\n odd_threads_in_block,\n odd_load_algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_load(thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\nCUB_TEST(\"Block load works with even vector types\",\n \"[load][block]\",\n vec_types,\n items_per_thread,\n a_block_size,\n load_algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_load(thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\nCUB_TEST(\"Block load works with custom types\",\n \"[load][block]\",\n items_per_thread,\n load_algorithm)\n{\n using type = c2h::custom_type_t;\n constexpr int items_per_thread = c2h::get<0, TestType>::value;\n constexpr int threads_in_block = 64;\n constexpr int tile_size = items_per_thread * threads_in_block;\n static constexpr cub::BlockLoadAlgorithm load_algorithm = \n c2h::get<1, TestType>::value;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_load(thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\nCUB_TEST(\"Block load works with caching iterators\",\n \"[load][block]\",\n items_per_thread,\n load_algorithm)\n{\n using type = int;\n constexpr int items_per_thread = c2h::get<0, TestType>::value;\n constexpr int threads_in_block = 64;\n constexpr int tile_size = items_per_thread * threads_in_block;\n static constexpr cub::BlockLoadAlgorithm load_algorithm = \n c2h::get<1, TestType>::value;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n cub::CacheModifiedInputIterator in(\n thrust::raw_pointer_cast(d_input.data()));\n\n thrust::device_vector d_output(d_input.size());\n\n block_load(in,\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \nstruct output_idx\n{\n static __device__ int get(int item)\n {\n return static_cast(threadIdx.x) * ItemsPerThread + item;\n }\n};\n\ntemplate \nstruct output_idx\n{\n static __device__ int get(int item)\n {\n return static_cast(threadIdx.x) + ThreadsInBlock * item;\n }\n};\n\ntemplate \n__global__ void kernel(std::integral_constant,\n InputIteratorT input,\n OutputIteratorT output,\n int num_items)\n{\n using input_t = hipcub::detail::value_t;\n using block_load_t = hipcub::BlockLoad;\n using storage_t =\n typename block_load_t::TempStorage;\n\n __shared__ storage_t storage;\n block_load_t block_load(storage);\n\n input_t data[ItemsPerThread];\n\n if (ItemsPerThread * ThreadsInBlock == num_items)\n {\n block_load.Load(input, data);\n }\n else\n {\n block_load.Load(input, data, num_items);\n }\n\n for (int i = 0; i < ItemsPerThread; i++)\n {\n const int idx =\n output_idx::get(i);\n\n if (idx < num_items)\n {\n output[idx] = data[i];\n }\n }\n}\n\ntemplate \n__global__ void kernel(std::integral_constant,\n InputIteratorT input,\n OutputIteratorT output,\n int num_items)\n{\n for (int i = 0; i < ItemsPerThread; i++)\n {\n const int idx =\n output_idx::get(i);\n\n if (idx < num_items)\n {\n output[idx] = input[idx];\n }\n }\n}\n\ntemplate \nvoid block_load(\n InputIteratorT input,\n OutputIteratorT output,\n int num_items)\n{\n using input_t = hipcub::detail::value_t;\n using block_load_t = hipcub::BlockLoad;\n using storage_t = typename block_load_t::TempStorage;\n constexpr bool sufficient_resources = sizeof(storage_t) <= 1024 * 48;\n\n kernel<<<1, ThreadsInBlock>>>(\n std::integral_constant{},\n input, output, num_items);\n\n REQUIRE( hipSuccess == hipPeekAtLastError() );\n REQUIRE( hipSuccess == hipDeviceSynchronize() );\n}\n\n\n// %PARAM% IPT it 1:11\n\nusing types = c2h::type_list;\nusing vec_types = c2h::type_list;\n\nusing even_threads_in_block = c2h::enum_type_list;\nusing odd_threads_in_block = c2h::enum_type_list;\nusing a_block_size = c2h::enum_type_list;\n\nusing items_per_thread = c2h::enum_type_list;\nusing load_algorithm = c2h::enum_type_list<\n hipcub::BlockLoadAlgorithm,\n hipcub::BlockLoadAlgorithm::BLOCK_LOAD_DIRECT,\n hipcub::BlockLoadAlgorithm::BLOCK_LOAD_STRIPED,\n hipcub::BlockLoadAlgorithm::BLOCK_LOAD_VECTORIZE,\n hipcub::BlockLoadAlgorithm::BLOCK_LOAD_TRANSPOSE,\n hipcub::BlockLoadAlgorithm::BLOCK_LOAD_WARP_TRANSPOSE, \n hipcub::BlockLoadAlgorithm::BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED>;\n\nusing odd_load_algorithm = c2h::enum_type_list<\n hipcub::BlockLoadAlgorithm,\n hipcub::BlockLoadAlgorithm::BLOCK_LOAD_DIRECT,\n hipcub::BlockLoadAlgorithm::BLOCK_LOAD_STRIPED,\n hipcub::BlockLoadAlgorithm::BLOCK_LOAD_VECTORIZE,\n hipcub::BlockLoadAlgorithm::BLOCK_LOAD_TRANSPOSE>;\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int items_per_thread = c2h::get<1, TestType>::value;\n static constexpr int threads_in_block = c2h::get<2, TestType>::value;\n static constexpr int tile_size = items_per_thread * threads_in_block;\n static constexpr hipcub::BlockLoadAlgorithm load_algorithm = \n c2h::get<3, TestType>::value;\n};\n\nCUB_TEST(\"Block load works with even block sizes\",\n \"[load][block]\",\n types,\n items_per_thread,\n even_threads_in_block,\n load_algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_load(thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\nCUB_TEST(\"Block load works with even odd sizes\",\n \"[load][block]\",\n types,\n items_per_thread,\n odd_threads_in_block,\n odd_load_algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_load(thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\nCUB_TEST(\"Block load works with even vector types\",\n \"[load][block]\",\n vec_types,\n items_per_thread,\n a_block_size,\n load_algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_load(thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\nCUB_TEST(\"Block load works with custom types\",\n \"[load][block]\",\n items_per_thread,\n load_algorithm)\n{\n using type = c2h::custom_type_t;\n constexpr int items_per_thread = c2h::get<0, TestType>::value;\n constexpr int threads_in_block = 64;\n constexpr int tile_size = items_per_thread * threads_in_block;\n static constexpr hipcub::BlockLoadAlgorithm load_algorithm = \n c2h::get<1, TestType>::value;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_load(thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\nCUB_TEST(\"Block load works with caching iterators\",\n \"[load][block]\",\n items_per_thread,\n load_algorithm)\n{\n using type = int;\n constexpr int items_per_thread = c2h::get<0, TestType>::value;\n constexpr int threads_in_block = 64;\n constexpr int tile_size = items_per_thread * threads_in_block;\n static constexpr hipcub::BlockLoadAlgorithm load_algorithm = \n c2h::get<1, TestType>::value;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n hipcub::CacheModifiedInputIterator in(\n thrust::raw_pointer_cast(d_input.data()));\n\n thrust::device_vector d_output(d_input.size());\n\n block_load(in,\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\n", "source": "/kaggle/working/cub/test/catch2_test_block_load.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void block_adj_diff_kernel(T *data, ActionT action, bool in_place)\n{\n using block_adjacent_differencet_t =\n cub::BlockAdjacentDifference;\n using temp_storage_t = typename block_adjacent_differencet_t::TempStorage;\n\n __shared__ temp_storage_t temp_storage;\n\n T thread_in[ItemsPerThread];\n T thread_out[ItemsPerThread];\n\n const int thread_offset = static_cast(threadIdx.x) * ItemsPerThread;\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n thread_in[item] = data[thread_offset + item];\n }\n __syncthreads();\n\n block_adjacent_differencet_t adj_diff(temp_storage);\n\n if (in_place)\n {\n action(adj_diff, thread_in, thread_in);\n\n for (unsigned int item = 0; item < ItemsPerThread; item++)\n {\n data[thread_offset + item] = thread_in[item];\n }\n }\n else\n {\n action(adj_diff, thread_in, thread_out);\n\n for (unsigned int item = 0; item < ItemsPerThread; item++)\n {\n data[thread_offset + item] = thread_out[item];\n }\n }\n}\n\ntemplate \nstruct custom_difference_t\n{\n __host__ __device__ T operator()(const T &lhs, const T &rhs)\n {\n return lhs - rhs;\n }\n};\n\ntemplate \nstruct base_op_t\n{\n template \n __device__ void operator()(BlockAdjDiff &block_adj_diff,\n T (&input)[ItemsPerThread],\n T (&output)[ItemsPerThread]) const\n {\n if (ReadLeft)\n {\n block_adj_diff.SubtractLeft(input, output, custom_difference_t{});\n }\n else\n {\n block_adj_diff.SubtractRight(input, output, custom_difference_t{});\n }\n }\n};\n\ntemplate \nstruct last_tile_op_t\n{\n int m_valid_items{};\n\n __host__ last_tile_op_t(int valid_items)\n : m_valid_items(valid_items)\n {}\n\n template \n __device__ void operator()(BlockAdjDiff &block_adj_diff,\n T (&input)[ITEMS_PER_THREAD],\n T (&output)[ITEMS_PER_THREAD]) const\n {\n custom_difference_t diff{};\n\n if (ReadLeft)\n {\n block_adj_diff.SubtractLeftPartialTile(input,\n output,\n diff,\n m_valid_items);\n }\n else\n {\n block_adj_diff.SubtractRightPartialTile(input,\n output,\n diff,\n m_valid_items);\n }\n }\n};\n\ntemplate \nstruct middle_tile_op_t\n{\n T m_neighbour_tile_value;\n\n __host__ middle_tile_op_t(T neighbour_tile_value)\n : m_neighbour_tile_value(neighbour_tile_value)\n {}\n\n template \n __device__ void operator()(BlockAdjDiff &block_adj_diff,\n T (&input)[ITEMS_PER_THREAD],\n T (&output)[ITEMS_PER_THREAD]) const\n {\n custom_difference_t diff{};\n\n if (ReadLeft)\n {\n block_adj_diff.SubtractLeft(input,\n output,\n diff,\n m_neighbour_tile_value);\n }\n else\n {\n block_adj_diff.SubtractRight(input,\n output,\n diff,\n m_neighbour_tile_value);\n }\n }\n};\n\ntemplate \nstruct last_tile_with_pred_op_t\n{\n int m_valid_items;\n T m_neighbour_tile_value;\n\n __host__ last_tile_with_pred_op_t(\n int valid_items,\n T neighbour_tile_value)\n : m_valid_items(valid_items)\n , m_neighbour_tile_value(neighbour_tile_value)\n {\n }\n\n template \n __device__ void operator()(BlockAdjDiff &block_adj_diff, \n T (&input)[ITEMS_PER_THREAD],\n T (&output)[ITEMS_PER_THREAD]) const\n {\n custom_difference_t diff{};\n block_adj_diff.SubtractLeftPartialTile(input,\n output,\n diff,\n m_valid_items,\n m_neighbour_tile_value);\n }\n};\n\ntemplate \nvoid block_adj_diff(thrust::device_vector &data,\n bool in_place,\n ActionT action)\n{\n block_adj_diff_kernel\n <<<1, ThreadsInBlock>>>(thrust::raw_pointer_cast(data.data()),\n action,\n in_place);\n\n REQUIRE(cudaSuccess == cudaPeekAtLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n}\n\ntemplate \nvoid host_adj_diff(thrust::host_vector &h_data, int valid_items)\n{\n custom_difference_t diff{};\n\n if (ReadLeft)\n {\n for (int i = valid_items - 1; i > 0; i--)\n {\n h_data[i] = diff(h_data[i], h_data[i - 1]);\n }\n }\n else\n {\n for (int i = 0; i < valid_items - 1; i++)\n {\n h_data[i] = diff(h_data[i], h_data[i + 1]);\n }\n }\n}\n\ntemplate \nvoid host_adj_diff(thrust::host_vector &h_data,\n int valid_items,\n T neighbour_value)\n{\n custom_difference_t diff{};\n\n host_adj_diff(h_data, valid_items);\n\n if (valid_items == 0)\n {\n return;\n }\n\n if (ReadLeft)\n {\n h_data[0] = diff(h_data[0], neighbour_value);\n }\n else\n {\n h_data[valid_items - 1] = diff(h_data[valid_items - 1], neighbour_value);\n }\n}\n\n// %PARAM% THREADS_IN_BLOCK bs 64:256\n\nusing key_types =\n c2h::type_list;\n\nusing threads_in_block = c2h::enum_type_list;\nusing items_per_thread = c2h::enum_type_list;\nusing directions = c2h::enum_type_list;\nusing left_only = c2h::enum_type_list;\n\ntemplate \nstruct params_t\n{\n using key_t = typename c2h::get<0, TestType>;\n\n static constexpr int items_per_thread = c2h::get<1, TestType>::value;\n static constexpr int threads_in_block = c2h::get<2, TestType>::value;\n static constexpr int tile_size = items_per_thread * threads_in_block;\n static constexpr bool read_left = c2h::get<3, TestType>::value;\n};\n\nCUB_TEST(\"Block adjacent difference works with full tiles\",\n \"[adjacent difference][block]\",\n key_types,\n items_per_thread,\n threads_in_block,\n directions)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n const bool in_place = GENERATE(false, true);\n\n thrust::host_vector h_data = d_data;\n host_adj_diff(h_data, params::tile_size);\n\n block_adj_diff(\n d_data,\n in_place,\n base_op_t{});\n\n REQUIRE(h_data == d_data);\n}\n\nCUB_TEST(\"Block adjacent difference works with last tiles\",\n \"[adjacent difference][block]\",\n key_types,\n items_per_thread,\n threads_in_block,\n directions)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n const bool in_place = GENERATE(false, true);\n const int valid_items = GENERATE_COPY(take(10, random(0, params::tile_size)));\n\n thrust::host_vector h_data = d_data;\n host_adj_diff(h_data, valid_items);\n\n block_adj_diff(\n d_data,\n in_place,\n last_tile_op_t{valid_items});\n\n REQUIRE(h_data == d_data);\n}\n\nCUB_TEST(\"Block adjacent difference works with single tiles\",\n \"[adjacent difference][block]\",\n key_types,\n items_per_thread,\n threads_in_block,\n left_only)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n const bool in_place = GENERATE(false, true);\n const int valid_items = GENERATE_COPY(take(10, random(0, params::tile_size)));\n constexpr bool read_left = true;\n\n thrust::host_vector h_data = d_data;\n key_t neighbour_value = h_data[h_data.size() / 2];\n\n host_adj_diff(h_data, valid_items, neighbour_value);\n\n block_adj_diff(\n d_data,\n in_place,\n last_tile_with_pred_op_t{valid_items, neighbour_value});\n\n REQUIRE(h_data == d_data);\n}\n\nCUB_TEST(\"Block adjacent difference works with middle tiles\",\n \"[adjacent difference][block]\",\n key_types,\n items_per_thread,\n threads_in_block,\n directions)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n const bool in_place = GENERATE(false, true);\n\n thrust::host_vector h_data = d_data;\n key_t neighbour_value = h_data[h_data.size() / 2]; \n\n host_adj_diff(h_data, params::tile_size, neighbour_value);\n\n block_adj_diff(\n d_data,\n in_place,\n middle_tile_op_t{neighbour_value});\n\n REQUIRE(h_data == d_data);\n}\n\nCUB_TEST(\"Block adjacent difference supports custom types\",\n \"[adjacent difference][block]\",\n threads_in_block)\n{\n using key_t = c2h::custom_type_t;\n\n constexpr int items_per_thread = 2;\n constexpr int threads_in_block = c2h::get<0, TestType>::value;\n constexpr int tile_size = threads_in_block * items_per_thread;\n constexpr bool read_left = true;\n const bool in_place = true;\n\n thrust::device_vector d_data(tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n thrust::host_vector h_data = d_data;\n host_adj_diff(h_data, tile_size);\n\n block_adj_diff(d_data,\n in_place,\n base_op_t{});\n\n REQUIRE(h_data == d_data);\n}\n\n// TODO Test different input/output types\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void block_adj_diff_kernel(T *data, ActionT action, bool in_place)\n{\n using block_adjacent_differencet_t =\n hipcub::BlockAdjacentDifference;\n using temp_storage_t = typename block_adjacent_differencet_t::TempStorage;\n\n __shared__ temp_storage_t temp_storage;\n\n T thread_in[ItemsPerThread];\n T thread_out[ItemsPerThread];\n\n const int thread_offset = static_cast(threadIdx.x) * ItemsPerThread;\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n thread_in[item] = data[thread_offset + item];\n }\n __syncthreads();\n\n block_adjacent_differencet_t adj_diff(temp_storage);\n\n if (in_place)\n {\n action(adj_diff, thread_in, thread_in);\n\n for (unsigned int item = 0; item < ItemsPerThread; item++)\n {\n data[thread_offset + item] = thread_in[item];\n }\n }\n else\n {\n action(adj_diff, thread_in, thread_out);\n\n for (unsigned int item = 0; item < ItemsPerThread; item++)\n {\n data[thread_offset + item] = thread_out[item];\n }\n }\n}\n\ntemplate \nstruct custom_difference_t\n{\n __host__ __device__ T operator()(const T &lhs, const T &rhs)\n {\n return lhs - rhs;\n }\n};\n\ntemplate \nstruct base_op_t\n{\n template \n __device__ void operator()(BlockAdjDiff &block_adj_diff,\n T (&input)[ItemsPerThread],\n T (&output)[ItemsPerThread]) const\n {\n if (ReadLeft)\n {\n block_adj_diff.SubtractLeft(input, output, custom_difference_t{});\n }\n else\n {\n block_adj_diff.SubtractRight(input, output, custom_difference_t{});\n }\n }\n};\n\ntemplate \nstruct last_tile_op_t\n{\n int m_valid_items{};\n\n __host__ last_tile_op_t(int valid_items)\n : m_valid_items(valid_items)\n {}\n\n template \n __device__ void operator()(BlockAdjDiff &block_adj_diff,\n T (&input)[ITEMS_PER_THREAD],\n T (&output)[ITEMS_PER_THREAD]) const\n {\n custom_difference_t diff{};\n\n if (ReadLeft)\n {\n block_adj_diff.SubtractLeftPartialTile(input,\n output,\n diff,\n m_valid_items);\n }\n else\n {\n block_adj_diff.SubtractRightPartialTile(input,\n output,\n diff,\n m_valid_items);\n }\n }\n};\n\ntemplate \nstruct middle_tile_op_t\n{\n T m_neighbour_tile_value;\n\n __host__ middle_tile_op_t(T neighbour_tile_value)\n : m_neighbour_tile_value(neighbour_tile_value)\n {}\n\n template \n __device__ void operator()(BlockAdjDiff &block_adj_diff,\n T (&input)[ITEMS_PER_THREAD],\n T (&output)[ITEMS_PER_THREAD]) const\n {\n custom_difference_t diff{};\n\n if (ReadLeft)\n {\n block_adj_diff.SubtractLeft(input,\n output,\n diff,\n m_neighbour_tile_value);\n }\n else\n {\n block_adj_diff.SubtractRight(input,\n output,\n diff,\n m_neighbour_tile_value);\n }\n }\n};\n\ntemplate \nstruct last_tile_with_pred_op_t\n{\n int m_valid_items;\n T m_neighbour_tile_value;\n\n __host__ last_tile_with_pred_op_t(\n int valid_items,\n T neighbour_tile_value)\n : m_valid_items(valid_items)\n , m_neighbour_tile_value(neighbour_tile_value)\n {\n }\n\n template \n __device__ void operator()(BlockAdjDiff &block_adj_diff, \n T (&input)[ITEMS_PER_THREAD],\n T (&output)[ITEMS_PER_THREAD]) const\n {\n custom_difference_t diff{};\n block_adj_diff.SubtractLeftPartialTile(input,\n output,\n diff,\n m_valid_items,\n m_neighbour_tile_value);\n }\n};\n\ntemplate \nvoid block_adj_diff(thrust::device_vector &data,\n bool in_place,\n ActionT action)\n{\n block_adj_diff_kernel\n <<<1, ThreadsInBlock>>>(thrust::raw_pointer_cast(data.data()),\n action,\n in_place);\n\n REQUIRE(hipSuccess == hipPeekAtLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n}\n\ntemplate \nvoid host_adj_diff(thrust::host_vector &h_data, int valid_items)\n{\n custom_difference_t diff{};\n\n if (ReadLeft)\n {\n for (int i = valid_items - 1; i > 0; i--)\n {\n h_data[i] = diff(h_data[i], h_data[i - 1]);\n }\n }\n else\n {\n for (int i = 0; i < valid_items - 1; i++)\n {\n h_data[i] = diff(h_data[i], h_data[i + 1]);\n }\n }\n}\n\ntemplate \nvoid host_adj_diff(thrust::host_vector &h_data,\n int valid_items,\n T neighbour_value)\n{\n custom_difference_t diff{};\n\n host_adj_diff(h_data, valid_items);\n\n if (valid_items == 0)\n {\n return;\n }\n\n if (ReadLeft)\n {\n h_data[0] = diff(h_data[0], neighbour_value);\n }\n else\n {\n h_data[valid_items - 1] = diff(h_data[valid_items - 1], neighbour_value);\n }\n}\n\n// %PARAM% THREADS_IN_BLOCK bs 64:256\n\nusing key_types =\n c2h::type_list;\n\nusing threads_in_block = c2h::enum_type_list;\nusing items_per_thread = c2h::enum_type_list;\nusing directions = c2h::enum_type_list;\nusing left_only = c2h::enum_type_list;\n\ntemplate \nstruct params_t\n{\n using key_t = typename c2h::get<0, TestType>;\n\n static constexpr int items_per_thread = c2h::get<1, TestType>::value;\n static constexpr int threads_in_block = c2h::get<2, TestType>::value;\n static constexpr int tile_size = items_per_thread * threads_in_block;\n static constexpr bool read_left = c2h::get<3, TestType>::value;\n};\n\nCUB_TEST(\"Block adjacent difference works with full tiles\",\n \"[adjacent difference][block]\",\n key_types,\n items_per_thread,\n threads_in_block,\n directions)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n const bool in_place = GENERATE(false, true);\n\n thrust::host_vector h_data = d_data;\n host_adj_diff(h_data, params::tile_size);\n\n block_adj_diff(\n d_data,\n in_place,\n base_op_t{});\n\n REQUIRE(h_data == d_data);\n}\n\nCUB_TEST(\"Block adjacent difference works with last tiles\",\n \"[adjacent difference][block]\",\n key_types,\n items_per_thread,\n threads_in_block,\n directions)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n const bool in_place = GENERATE(false, true);\n const int valid_items = GENERATE_COPY(take(10, random(0, params::tile_size)));\n\n thrust::host_vector h_data = d_data;\n host_adj_diff(h_data, valid_items);\n\n block_adj_diff(\n d_data,\n in_place,\n last_tile_op_t{valid_items});\n\n REQUIRE(h_data == d_data);\n}\n\nCUB_TEST(\"Block adjacent difference works with single tiles\",\n \"[adjacent difference][block]\",\n key_types,\n items_per_thread,\n threads_in_block,\n left_only)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n const bool in_place = GENERATE(false, true);\n const int valid_items = GENERATE_COPY(take(10, random(0, params::tile_size)));\n constexpr bool read_left = true;\n\n thrust::host_vector h_data = d_data;\n key_t neighbour_value = h_data[h_data.size() / 2];\n\n host_adj_diff(h_data, valid_items, neighbour_value);\n\n block_adj_diff(\n d_data,\n in_place,\n last_tile_with_pred_op_t{valid_items, neighbour_value});\n\n REQUIRE(h_data == d_data);\n}\n\nCUB_TEST(\"Block adjacent difference works with middle tiles\",\n \"[adjacent difference][block]\",\n key_types,\n items_per_thread,\n threads_in_block,\n directions)\n{\n using params = params_t;\n using key_t = typename params::key_t;\n\n thrust::device_vector d_data(params::tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n const bool in_place = GENERATE(false, true);\n\n thrust::host_vector h_data = d_data;\n key_t neighbour_value = h_data[h_data.size() / 2]; \n\n host_adj_diff(h_data, params::tile_size, neighbour_value);\n\n block_adj_diff(\n d_data,\n in_place,\n middle_tile_op_t{neighbour_value});\n\n REQUIRE(h_data == d_data);\n}\n\nCUB_TEST(\"Block adjacent difference supports custom types\",\n \"[adjacent difference][block]\",\n threads_in_block)\n{\n using key_t = c2h::custom_type_t;\n\n constexpr int items_per_thread = 2;\n constexpr int threads_in_block = c2h::get<0, TestType>::value;\n constexpr int tile_size = threads_in_block * items_per_thread;\n constexpr bool read_left = true;\n const bool in_place = true;\n\n thrust::device_vector d_data(tile_size);\n c2h::gen(CUB_SEED(10), d_data);\n\n thrust::host_vector h_data = d_data;\n host_adj_diff(h_data, tile_size);\n\n block_adj_diff(d_data,\n in_place,\n base_op_t{});\n\n REQUIRE(h_data == d_data);\n}\n\n// TODO Test different input/output types\n\n", "source": "/kaggle/working/cub/test/catch2_test_block_adjacent_difference.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceRadixSort utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#if !_NVHPC_CUDA\n #include \n#endif\n\n#if !_NVHPC_CUDA\n #include \n#endif\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"test_util.h\"\n\nusing namespace cub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\nstd::size_t g_smallest_pre_sorted_num_items = (std::size_t(1) << 32) - 42;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method (allows overwriting of input)\n CUB_NO_OVERWRITE, // CUB method (disallows overwriting of input)\n\n CUB_SEGMENTED, // CUB method (allows overwriting of input)\n CUB_SEGMENTED_NO_OVERWRITE, // CUB method (disallows overwriting of input)\n\n // Same as above, but launches kernels from device using CDP.\n CDP,\n CDP_NO_OVERWRITE,\n CDP_SEGMENTED,\n CDP_SEGMENTED_NO_OVERWRITE,\n};\n\nstatic const char* BackendToString(Backend b)\n{\n switch (b)\n {\n case CUB:\n return \"CUB\";\n case CUB_NO_OVERWRITE:\n return \"CUB_NO_OVERWRITE\";\n case CUB_SEGMENTED:\n return \"CUB_SEGMENTED\";\n case CUB_SEGMENTED_NO_OVERWRITE:\n return \"CUB_SEGMENTED_NO_OVERWRITE\";\n case CDP:\n return \"CDP\";\n case CDP_NO_OVERWRITE:\n return \"CDP_NO_OVERWRITE\";\n case CDP_SEGMENTED:\n return \"CDP_SEGMENTED\";\n case CDP_SEGMENTED_NO_OVERWRITE:\n return \"CDP_SEGMENTED_NO_OVERWRITE\";\n default:\n break;\n }\n\n return \"\";\n}\n\n//---------------------------------------------------------------------\n// Dispatch to different DeviceRadixSort entrypoints\n//---------------------------------------------------------------------\n\n/**\n * Dispatch to CUB sorting entrypoint (specialized for ascending)\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int /*num_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n int begin_bit,\n int end_bit)\n{\n return DeviceRadixSort::SortPairs(\n d_temp_storage, temp_storage_bytes,\n d_keys, d_values,\n num_items, begin_bit, end_bit);\n}\n\n/**\n * Dispatch to CUB_NO_OVERWRITE sorting entrypoint (specialized for ascending)\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int /*num_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n int begin_bit,\n int end_bit)\n{\n KeyT const *const_keys_itr = d_keys.Current();\n ValueT const *const_values_itr = d_values.Current();\n\n cudaError_t retval = DeviceRadixSort::SortPairs(\n d_temp_storage, temp_storage_bytes,\n const_keys_itr, d_keys.Alternate(), const_values_itr, d_values.Alternate(),\n num_items, begin_bit, end_bit);\n\n d_keys.selector ^= 1;\n d_values.selector ^= 1;\n return retval;\n}\n\n/**\n * Dispatch to CUB sorting entrypoint (specialized for descending)\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int /*num_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n int begin_bit,\n int end_bit)\n{\n return DeviceRadixSort::SortPairsDescending(\n d_temp_storage, temp_storage_bytes,\n d_keys, d_values, num_items,\n begin_bit, end_bit);\n}\n\n\n/**\n * Dispatch to CUB_NO_OVERWRITE sorting entrypoint (specialized for descending)\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int /*num_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n int begin_bit,\n int end_bit)\n{\n KeyT const *const_keys_itr = d_keys.Current();\n ValueT const *const_values_itr = d_values.Current();\n\n cudaError_t retval = DeviceRadixSort::SortPairsDescending(\n d_temp_storage, temp_storage_bytes,\n const_keys_itr, d_keys.Alternate(), const_values_itr, d_values.Alternate(),\n num_items, begin_bit, end_bit);\n\n d_keys.selector ^= 1;\n d_values.selector ^= 1;\n return retval;\n}\n\n//---------------------------------------------------------------------\n// Dispatch to different DeviceRadixSort entrypoints\n//---------------------------------------------------------------------\n\n// Validates that `num_items` fits into `int`\n// TODO(canonizer): remove this check once num_items is templated for segmented sort.\ntemplate \n__host__ __device__ bool ValidateNumItemsForSegmentedSort(NumItemsT num_items)\n{\n if (static_cast(num_items) <\n static_cast(INT_MAX))\n {\n return true;\n }\n else\n {\n printf(\"cub::DeviceSegmentedRadixSort is currently limited by %d items but \"\n \"%lld were provided\\n\",\n INT_MAX,\n static_cast(num_items));\n }\n\n return false;\n}\n\n/**\n * Dispatch to CUB_SEGMENTED sorting entrypoint (specialized for ascending)\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n if (ValidateNumItemsForSegmentedSort(num_items))\n {\n return DeviceSegmentedRadixSort::SortPairs(d_temp_storage,\n temp_storage_bytes,\n d_keys,\n d_values,\n static_cast(num_items),\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit);\n }\n\n return cudaErrorInvalidValue;\n}\n\n/**\n * Dispatch to CUB_SEGMENTED_NO_OVERWRITE sorting entrypoint (specialized for ascending)\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n if (ValidateNumItemsForSegmentedSort(num_items))\n {\n KeyT const *const_keys_itr = d_keys.Current();\n ValueT const *const_values_itr = d_values.Current();\n\n cudaError_t retval =\n DeviceSegmentedRadixSort::SortPairs(d_temp_storage,\n temp_storage_bytes,\n const_keys_itr,\n d_keys.Alternate(),\n const_values_itr,\n d_values.Alternate(),\n static_cast(num_items),\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit);\n\n d_keys.selector ^= 1;\n d_values.selector ^= 1;\n return retval;\n }\n\n return cudaErrorInvalidValue;\n}\n\n\n/**\n * Dispatch to CUB_SEGMENTED sorting entrypoint (specialized for descending)\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n if (ValidateNumItemsForSegmentedSort(num_items))\n {\n return DeviceSegmentedRadixSort::SortPairsDescending(\n d_temp_storage,\n temp_storage_bytes,\n d_keys,\n d_values,\n static_cast(num_items),\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit);\n }\n\n return cudaErrorInvalidValue;\n}\n\n/**\n * Dispatch to CUB_SEGMENTED_NO_OVERWRITE sorting entrypoint (specialized for descending)\n */\ntemplate \nCUB_RUNTIME_FUNCTION\ncudaError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n if (ValidateNumItemsForSegmentedSort(num_items))\n {\n KeyT const *const_keys_itr = d_keys.Current();\n ValueT const *const_values_itr = d_values.Current();\n\n cudaError_t retval =\n DeviceSegmentedRadixSort::SortPairsDescending(d_temp_storage,\n temp_storage_bytes,\n const_keys_itr,\n d_keys.Alternate(),\n const_values_itr,\n d_values.Alternate(),\n static_cast(num_items),\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit);\n\n d_keys.selector ^= 1;\n d_values.selector ^= 1;\n return retval;\n }\n\n return cudaErrorInvalidValue;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceRadixSort\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type is_descending,\n Int2Type cub_backend,\n int *d_selector,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n void *d_temp_storage,\n size_t temp_storage_bytes,\n DoubleBuffer d_keys,\n DoubleBuffer d_values,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n *d_cdp_error = Dispatch(is_descending,\n cub_backend,\n d_selector,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys,\n d_values,\n num_items,\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n *d_selector = d_keys.selector;\n}\n\n/**\n * Launch kernel and dispatch on device. Should only be called from host code.\n * The CubBackend should be one of the non-CDP CUB backends to invoke from the\n * device.\n */\ntemplate \ncudaError_t LaunchCDPKernel(Int2Type is_descending,\n Int2Type cub_backend,\n int *d_selector,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n // Invoke kernel to invoke device-side dispatch:\n cudaError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n is_descending,\n cub_backend,\n d_selector,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys,\n d_values,\n num_items,\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit);\n CubDebugExit(retval);\n CubDebugExit(cub::detail::device_synchronize());\n\n // Copy out selector\n CubDebugExit(cudaMemcpy(&d_keys.selector,\n d_selector,\n sizeof(int) * 1,\n cudaMemcpyDeviceToHost));\n d_values.selector = d_keys.selector;\n\n // Copy out temp_storage_bytes\n CubDebugExit(cudaMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n cudaMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(cudaMemcpy(&retval,\n d_cdp_error,\n sizeof(cudaError_t) * 1,\n cudaMemcpyDeviceToHost));\n\n return retval;\n}\n\n// Specializations of Dispatch that translate the CDP backend to the appropriate\n// CUB backend, and uses the CUB backend to launch the CDP kernel.\n#define DEFINE_CDP_DISPATCHER(CdpBackend, CubBackend) \\\n template \\\n cudaError_t Dispatch(Int2Type is_descending, \\\n Int2Type /*dispatch_to*/, \\\n int *d_selector, \\\n size_t *d_temp_storage_bytes, \\\n cudaError_t *d_cdp_error, \\\n \\\n void *d_temp_storage, \\\n size_t &temp_storage_bytes, \\\n DoubleBuffer &d_keys, \\\n DoubleBuffer &d_values, \\\n NumItemsT num_items, \\\n int num_segments, \\\n BeginOffsetIteratorT d_segment_begin_offsets, \\\n EndOffsetIteratorT d_segment_end_offsets, \\\n int begin_bit, \\\n int end_bit) \\\n { \\\n Int2Type cub_backend{}; \\\n return LaunchCDPKernel(is_descending, \\\n cub_backend, \\\n d_selector, \\\n d_temp_storage_bytes, \\\n d_cdp_error, \\\n d_temp_storage, \\\n temp_storage_bytes, \\\n d_keys, \\\n d_values, \\\n num_items, \\\n num_segments, \\\n d_segment_begin_offsets, \\\n d_segment_end_offsets, \\\n begin_bit, \\\n end_bit); \\\n }\n\nDEFINE_CDP_DISPATCHER(CDP, CUB)\nDEFINE_CDP_DISPATCHER(CDP_NO_OVERWRITE, CUB_NO_OVERWRITE)\nDEFINE_CDP_DISPATCHER(CDP_SEGMENTED, CUB_SEGMENTED)\nDEFINE_CDP_DISPATCHER(CDP_SEGMENTED_NO_OVERWRITE, CUB_SEGMENTED_NO_OVERWRITE)\n\n#undef DEFINE_CDP_DISPATCHER\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Problem generation\n//---------------------------------------------------------------------\n\n\n/**\n * Simple key-value pairing\n */\ntemplate <\n typename KeyT,\n typename ValueT>\nstruct Pair\n{\n KeyT key;\n ValueT value;\n\n bool operator<(const Pair &b) const\n {\n return (key < b.key);\n }\n};\n\n\n/**\n * Simple key-value pairing (specialized for bool types)\n */\ntemplate \nstruct Pair\n{\n bool key;\n ValueT value;\n\n bool operator<(const Pair &b) const\n {\n return (!key && b.key);\n }\n};\n\n\n/**\n * Initialize key data\n */\ntemplate \nvoid InitializeKeyBits(\n GenMode gen_mode,\n KeyT *h_keys,\n NumItemsT num_items,\n int /*entropy_reduction*/)\n{\n for (NumItemsT i = 0; i < num_items; ++i)\n InitValue(gen_mode, h_keys[i], i);\n}\n\ntemplate ::UnsignedBits>\nUnsignedBits KeyBits(KeyT key)\n{\n UnsignedBits bits;\n memcpy(&bits, &key, sizeof(KeyT));\n return bits;\n}\n\n/** Initialize the reference array monotonically. */\ntemplate \nvoid InitializeKeysSorted(\n KeyT *h_keys,\n NumItemsT num_items)\n{\n using TraitsT = cub::Traits;\n using UnsignedBits = typename TraitsT::UnsignedBits;\n\n // Numbers to generate random runs.\n UnsignedBits max_inc = 1 << (sizeof(UnsignedBits) < 4 ? 3 :\n (sizeof(UnsignedBits) < 8 ? 14 : 24));\n UnsignedBits min_bits = TraitsT::TwiddleIn(KeyBits(TraitsT::Lowest()));\n UnsignedBits max_bits = TraitsT::TwiddleIn(KeyBits(TraitsT::Max()));\n NumItemsT max_run = std::max(\n NumItemsT(double(num_items) * (max_inc + 1) / max_bits),\n NumItemsT(1 << 14));\n\n UnsignedBits *h_key_bits = reinterpret_cast(h_keys);\n NumItemsT i = 0;\n // Start with the minimum twiddled key.\n UnsignedBits twiddled_key = min_bits;\n while (i < num_items)\n {\n // Generate random increment (avoid overflow).\n UnsignedBits inc_bits = 0;\n RandomBits(inc_bits);\n // twiddled_key < max_bits at this point.\n UnsignedBits inc = static_cast(std::min(1 + inc_bits % max_inc, max_bits - twiddled_key));\n twiddled_key += inc;\n\n // Generate random run length (ensure there are enough values to fill the rest).\n NumItemsT run_bits = 0;\n RandomBits(run_bits);\n NumItemsT run_length = std::min(1 + run_bits % max_run, num_items - i);\n if (twiddled_key == max_bits) run_length = num_items - i;\n NumItemsT run_end = i + run_length;\n\n // Fill the array.\n UnsignedBits key = TraitsT::TwiddleOut(twiddled_key);\n // Avoid -0.0 for floating-point keys.\n UnsignedBits negative_zero = UnsignedBits(1) << UnsignedBits(sizeof(UnsignedBits) * 8 - 1);\n if (TraitsT::CATEGORY == cub::FLOATING_POINT && key == negative_zero)\n {\n key = 0;\n }\n\n for (; i < run_end; ++i)\n {\n h_key_bits[i] = key;\n }\n }\n}\n\n\n/**\n * Initialize solution\n */\ntemplate \nvoid InitializeSolution(\n KeyT *h_keys,\n NumItemsT num_items,\n int num_segments,\n bool pre_sorted,\n NumItemsT *h_segment_offsets,\n int begin_bit,\n int end_bit,\n NumItemsT *&h_reference_ranks,\n KeyT *&h_reference_keys)\n{\n if (num_items == 0)\n {\n h_reference_ranks = nullptr;\n h_reference_keys = nullptr;\n return;\n }\n\n if (pre_sorted)\n {\n printf(\"Shuffling reference solution on CPU\\n\");\n // Note: begin_bit and end_bit are ignored here, and assumed to have the\n // default values (begin_bit == 0, end_bit == 8 * sizeof(KeyT)).\n // Otherwise, pre-sorting won't work, as it doesn't necessarily\n // correspond to the order of keys sorted by a subrange of bits.\n // num_segments is also ignored as assumed to be 1, as pre-sorted tests\n // are currently not supported for multiple segments.\n //\n // Pre-sorted tests with non-default begin_bit, end_bit or num_segments\n // != 1 are skipped in TestBits() and TestSegments(), respectively.\n AssertEquals(begin_bit, 0);\n AssertEquals(end_bit, static_cast(8 * sizeof(KeyT)));\n AssertEquals(num_segments, 1);\n\n // Copy to the reference solution.\n h_reference_keys = new KeyT[num_items];\n if (IS_DESCENDING)\n {\n // Copy in reverse.\n for (NumItemsT i = 0; i < num_items; ++i)\n {\n h_reference_keys[i] = h_keys[num_items - 1 - i];\n }\n // Copy back.\n memcpy(h_keys, h_reference_keys, num_items * sizeof(KeyT));\n }\n else\n {\n memcpy(h_reference_keys, h_keys, num_items * sizeof(KeyT));\n }\n\n // Summarize the pre-sorted array (element, 1st position, count).\n struct Element\n {\n KeyT key;\n NumItemsT num;\n NumItemsT index;\n };\n\n std::vector summary;\n KeyT cur_key = h_reference_keys[0];\n summary.push_back(Element{cur_key, 1, 0});\n for (NumItemsT i = 1; i < num_items; ++i)\n {\n KeyT key = h_reference_keys[i];\n if (key == cur_key)\n {\n // Same key.\n summary.back().num++;\n continue;\n }\n\n // Different key.\n cur_key = key;\n summary.push_back(Element{cur_key, 1, i});\n }\n\n // Generate a random permutation from the summary. Such a complicated\n // approach is used to permute the array and compute ranks in a\n // cache-friendly way and in a short time.\n if (WANT_RANKS)\n {\n h_reference_ranks = new NumItemsT[num_items];\n }\n NumItemsT max_run = 32;\n NumItemsT run = 0;\n NumItemsT i = 0;\n while (summary.size() > 0)\n {\n // Pick up a random element and a run.\n NumItemsT bits = 0;\n RandomBits(bits);\n NumItemsT summary_id = bits % summary.size();\n Element& element = summary[summary_id];\n run = std::min(1 + bits % (max_run - 1), element.num);\n for (NumItemsT j = 0; j < run; ++j)\n {\n h_keys[i + j] = element.key;\n if (WANT_RANKS)\n {\n h_reference_ranks[element.index + j] = i + j;\n }\n }\n i += run;\n element.index += run;\n element.num -= run;\n if (element.num == 0)\n {\n // Remove the empty entry.\n std::swap(summary[summary_id], summary.back());\n summary.pop_back();\n }\n }\n printf(\" Done.\\n\");\n }\n else\n {\n typedef Pair PairT;\n\n PairT *h_pairs = new PairT[num_items];\n\n int num_bits = end_bit - begin_bit;\n for (NumItemsT i = 0; i < num_items; ++i)\n {\n\n // Mask off unwanted portions\n if (num_bits < static_cast(sizeof(KeyT) * 8))\n {\n using UnsignedBits = typename cub::Traits::UnsignedBits;\n\n UnsignedBits base = 0;\n memcpy(&base, &h_keys[i], sizeof(KeyT));\n base &= ((UnsignedBits{1} << num_bits) - 1) << begin_bit;\n memcpy(&h_pairs[i].key, &base, sizeof(KeyT));\n }\n else\n {\n h_pairs[i].key = h_keys[i];\n }\n\n h_pairs[i].value = i;\n }\n\n printf(\"\\nSorting reference solution on CPU \"\n \"(%zd items, %d segments, %zd items/seg)...\",\n static_cast(num_items),\n num_segments,\n static_cast(num_items / num_segments));\n fflush(stdout);\n\n for (int i = 0; i < num_segments; ++i)\n {\n if (IS_DESCENDING) std::reverse(h_pairs + h_segment_offsets[i], h_pairs + h_segment_offsets[i + 1]);\n std::stable_sort( h_pairs + h_segment_offsets[i], h_pairs + h_segment_offsets[i + 1]);\n if (IS_DESCENDING) std::reverse(h_pairs + h_segment_offsets[i], h_pairs + h_segment_offsets[i + 1]);\n }\n\n printf(\" Done.\\n\"); fflush(stdout);\n\n if (WANT_RANKS)\n {\n h_reference_ranks = new NumItemsT[num_items];\n }\n h_reference_keys = new KeyT[num_items];\n\n for (NumItemsT i = 0; i < num_items; ++i)\n {\n if (WANT_RANKS)\n {\n h_reference_ranks[i] = h_pairs[i].value;\n }\n h_reference_keys[i] = h_keys[h_pairs[i].value];\n }\n\n if (h_pairs) delete[] h_pairs;\n }\n}\n\ntemplate \nvoid ResetKeys(KeyT *h_keys, NumItemsT num_items, bool pre_sorted, KeyT *reference_keys)\n{\n if (!pre_sorted) return;\n\n // Copy the reference keys back.\n if (IS_DESCENDING)\n {\n // Keys need to be copied in reverse.\n for (NumItemsT i = 0; i < num_items; ++i)\n {\n h_keys[i] = reference_keys[num_items - 1 - i];\n }\n }\n else\n {\n memcpy(h_keys, reference_keys, num_items * sizeof(KeyT));\n }\n}\n\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\ntemplate \nstruct UnwrapHalfAndBfloat16 {\n using Type = T;\n};\n\n#if !_NVHPC_CUDA\ntemplate <>\nstruct UnwrapHalfAndBfloat16 {\n using Type = __half;\n};\n#endif\n\n#if !_NVHPC_CUDA\ntemplate <>\nstruct UnwrapHalfAndBfloat16 {\n using Type = __nv_bfloat16;\n};\n#endif\n\ntemplate \nint compare_device_arrays(T* host_reference, T* d_tmp_buffer, T* d_data, NumItemsT num_items)\n{\n CubDebugExit(cudaMemcpy(d_tmp_buffer, host_reference, sizeof(T) * num_items, cudaMemcpyHostToDevice));\n\n auto d_reference_begin = d_tmp_buffer;\n auto d_reference_end = d_tmp_buffer + num_items;\n auto err = thrust::mismatch(thrust::device, d_reference_begin, d_reference_end, d_data);\n\n if (err.first != d_reference_end)\n {\n const auto index = thrust::distance(d_reference_begin, err.first);\n return CompareDeviceResults(host_reference + index, d_data + index, 1, true, g_verbose);\n }\n\n return 0;\n}\n\ntemplate \nint compare_device_arrays(\n CUB_NS_QUALIFIER::NullType */* h_reference */,\n CUB_NS_QUALIFIER::NullType */* d_tmp_buffer */,\n CUB_NS_QUALIFIER::NullType */* d_data */,\n NumItemsT /* num_items */)\n{\n return 0;\n}\n\n/**\n * Test DeviceRadixSort\n */\ntemplate <\n Backend BACKEND,\n bool IS_DESCENDING,\n typename KeyT,\n typename ValueT,\n typename BeginOffsetIteratorT,\n typename EndOffsetIteratorT,\n typename NumItemsT>\nvoid Test(\n KeyT *h_keys,\n ValueT *h_values,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit,\n KeyT *h_reference_keys,\n ValueT *h_reference_values)\n{\n // Key alias type\n using KeyAliasT = typename UnwrapHalfAndBfloat16::Type;\n\n const bool KEYS_ONLY = std::is_same::value;\n\n printf(\"%s %s cub::DeviceRadixSort %zd items, %d segments, \"\n \"%d-byte keys (%s) %d-byte values (%s), %d-byte num_items (%s), \"\n \"descending %d, begin_bit %d, end_bit %d\\n\",\n BackendToString(BACKEND),\n (KEYS_ONLY) ? \"keys-only\" : \"key-value\",\n static_cast(num_items),\n num_segments,\n static_cast(sizeof(KeyT)),\n typeid(KeyT).name(),\n (KEYS_ONLY) ? 0 : static_cast(sizeof(ValueT)),\n typeid(ValueT).name(),\n static_cast(sizeof(NumItemsT)),\n typeid(NumItemsT).name(),\n IS_DESCENDING,\n begin_bit,\n end_bit);\n\n if (g_verbose)\n {\n printf(\"Input keys:\\n\");\n DisplayResults(h_keys, num_items);\n printf(\"\\n\\n\");\n }\n\n // Allocate device arrays\n DoubleBuffer d_keys;\n DoubleBuffer d_values;\n int *d_selector;\n size_t *d_temp_storage_bytes;\n cudaError_t *d_cdp_error;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys.d_buffers[0], sizeof(KeyT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys.d_buffers[1], sizeof(KeyT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_selector, sizeof(int) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(cudaError_t) * 1));\n if (!KEYS_ONLY)\n {\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_values.d_buffers[0], sizeof(ValueT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_values.d_buffers[1], sizeof(ValueT) * num_items));\n }\n\n // Allocate temporary storage (and make it un-aligned)\n size_t temp_storage_bytes = 0;\n void *d_temp_storage = NULL;\n CubDebugExit(Dispatch(\n Int2Type(), Int2Type(), d_selector, d_temp_storage_bytes, d_cdp_error,\n d_temp_storage, temp_storage_bytes, d_keys, d_values,\n num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets,\n begin_bit, end_bit));\n\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes + 1));\n void* mis_aligned_temp = static_cast(d_temp_storage) + 1;\n\n // Initialize/clear device arrays\n d_keys.selector = 0;\n CubDebugExit(cudaMemcpy(d_keys.d_buffers[0], h_keys, sizeof(KeyT) * num_items, cudaMemcpyHostToDevice));\n CubDebugExit(cudaMemset(d_keys.d_buffers[1], 0, sizeof(KeyT) * num_items));\n if (!KEYS_ONLY)\n {\n d_values.selector = 0;\n CubDebugExit(cudaMemcpy(d_values.d_buffers[0], h_values, sizeof(ValueT) * num_items, cudaMemcpyHostToDevice));\n CubDebugExit(cudaMemset(d_values.d_buffers[1], 0, sizeof(ValueT) * num_items));\n }\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(\n Int2Type(), Int2Type(), d_selector, d_temp_storage_bytes, d_cdp_error,\n mis_aligned_temp, temp_storage_bytes, d_keys, d_values,\n num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets,\n begin_bit, end_bit));\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Check for correctness (and display results, if specified)\n printf(\"Warmup done. Checking results:\\n\"); fflush(stdout);\n\n int compare = 0;\n\n // If in/out API is used, we are not allowed to overwrite the input. \n // Let's check that the input buffer is not overwritten by the algorithm.\n if (BACKEND == CUB_NO_OVERWRITE)\n {\n KeyT *d_input_keys = reinterpret_cast(d_keys.d_buffers[0]);\n\n if (temp_storage_bytes < sizeof(KeyT) * num_items)\n {\n // For small input sizes, temporary storage is not large enough to fit keys.\n compare = CompareDeviceResults(h_keys, d_input_keys, num_items, true, g_verbose);\n }\n else \n {\n // If overwrite is not allowed, temporary storage is large enough to fit keys.\n KeyT* temp_keys = reinterpret_cast(d_temp_storage);\n\n compare = compare_device_arrays(h_keys, temp_keys, d_input_keys, num_items);\n printf(\"\\t Compare input keys: %s \", compare ? \"FAIL\" : \"PASS\"); fflush(stdout);\n }\n }\n\n // After the previous check is done, we can safely reuse alternative buffer to store \n // the reference results and compare current output.\n compare |= compare_device_arrays(h_reference_keys,\n reinterpret_cast(d_keys.Alternate()),\n reinterpret_cast(d_keys.Current()),\n num_items);\n\n printf(\"\\t Compare keys (selector %d): %s \", d_keys.selector, compare ? \"FAIL\" : \"PASS\"); fflush(stdout);\n\n if (!KEYS_ONLY)\n {\n int values_compare = compare_device_arrays(h_reference_values,\n d_values.Alternate(),\n d_values.Current(),\n num_items);\n\n compare |= values_compare;\n printf(\"\\t Compare values (selector %d): %s \", d_values.selector, values_compare ? \"FAIL\" : \"PASS\"); fflush(stdout);\n }\n\n // Performance\n if (g_timing_iterations)\n printf(\"\\nPerforming timing iterations:\\n\"); fflush(stdout);\n\n GpuTimer gpu_timer;\n float elapsed_millis = 0.0f;\n for (int i = 0; i < g_timing_iterations; ++i)\n {\n // Initialize/clear device arrays\n CubDebugExit(cudaMemcpy(d_keys.d_buffers[d_keys.selector], h_keys, sizeof(KeyT) * num_items, cudaMemcpyHostToDevice));\n CubDebugExit(cudaMemset(d_keys.d_buffers[d_keys.selector ^ 1], 0, sizeof(KeyT) * num_items));\n if (!KEYS_ONLY)\n {\n CubDebugExit(cudaMemcpy(d_values.d_buffers[d_values.selector], h_values, sizeof(ValueT) * num_items, cudaMemcpyHostToDevice));\n CubDebugExit(cudaMemset(d_values.d_buffers[d_values.selector ^ 1], 0, sizeof(ValueT) * num_items));\n }\n\n gpu_timer.Start();\n CubDebugExit(Dispatch(\n Int2Type(), Int2Type(), d_selector, d_temp_storage_bytes, d_cdp_error,\n mis_aligned_temp, temp_storage_bytes, d_keys, d_values,\n num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets,\n begin_bit, end_bit));\n gpu_timer.Stop();\n elapsed_millis += gpu_timer.ElapsedMillis();\n }\n\n // Display performance\n if (g_timing_iterations > 0)\n {\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n float giga_bandwidth = (KEYS_ONLY) ?\n giga_rate * sizeof(KeyT) * 2 :\n giga_rate * (sizeof(KeyT) + sizeof(ValueT)) * 2;\n printf(\"\\n%.3f elapsed ms, %.3f avg ms, %.3f billion items/s, %.3f logical GB/s\", elapsed_millis, avg_millis, giga_rate, giga_bandwidth);\n }\n\n printf(\"\\n\\n\");\n\n // Cleanup\n if (d_keys.d_buffers[0]) CubDebugExit(g_allocator.DeviceFree(d_keys.d_buffers[0]));\n if (d_keys.d_buffers[1]) CubDebugExit(g_allocator.DeviceFree(d_keys.d_buffers[1]));\n if (d_values.d_buffers[0]) CubDebugExit(g_allocator.DeviceFree(d_values.d_buffers[0]));\n if (d_values.d_buffers[1]) CubDebugExit(g_allocator.DeviceFree(d_values.d_buffers[1]));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_selector) CubDebugExit(g_allocator.DeviceFree(d_selector));\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n\n // Correctness asserts\n AssertEquals(0, compare);\n}\n\n// Returns whether there is enough memory for the test.\ntemplate \nbool HasEnoughMemory(std::size_t num_items, bool overwrite)\n{\n std::size_t total_mem = TotalGlobalMem();\n std::size_t value_size = std::is_same::value\n ? 0\n : sizeof(ValueT);\n // A conservative estimate of the amount of memory required.\n double factor = overwrite ? 2.25 : 3.25;\n std::size_t test_mem = static_cast\n (num_items * (sizeof(KeyT) + value_size) * factor);\n return test_mem < total_mem;\n}\n\n\n/**\n * Test backend\n */\ntemplate \nvoid TestBackend(KeyT *h_keys,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit,\n KeyT *h_reference_keys,\n NumItemsT *h_reference_ranks)\n{\n#if TEST_CDP == 0\n constexpr auto NonSegmentedOverwrite = CUB;\n constexpr auto NonSegmentedNoOverwrite = CUB_NO_OVERWRITE;\n constexpr auto SegmentedOverwrite = CUB_SEGMENTED;\n constexpr auto SegmentedNoOverwrite = CUB_SEGMENTED_NO_OVERWRITE;\n#else // TEST_CDP\n constexpr auto NonSegmentedOverwrite = CDP;\n constexpr auto NonSegmentedNoOverwrite = CDP_NO_OVERWRITE;\n constexpr auto SegmentedOverwrite = CDP_SEGMENTED;\n constexpr auto SegmentedNoOverwrite = CDP_SEGMENTED_NO_OVERWRITE;\n#endif // TEST_CDP\n\n const bool KEYS_ONLY = std::is_same::value;\n\n // A conservative check assuming overwrite is allowed.\n if (!HasEnoughMemory(static_cast(num_items), true))\n {\n printf(\"Skipping the test due to insufficient device memory\\n\");\n return;\n }\n\n std::unique_ptr h_value_data{};\n\n ValueT *h_values = nullptr;\n ValueT *h_reference_values = nullptr;\n\n if (!KEYS_ONLY)\n {\n h_value_data.reset(new ValueT[2 * static_cast(num_items)]);\n h_values = h_value_data.get();\n h_reference_values = h_value_data.get() + num_items;\n\n for (NumItemsT i = 0; i < num_items; ++i)\n {\n InitValue(INTEGER_SEED, h_values[i], i);\n InitValue(INTEGER_SEED, h_reference_values[i], h_reference_ranks[i]);\n }\n }\n\n // Skip segmented sort if num_items isn't int.\n // TODO(64bit-seg-sort): re-enable these tests once num_items is templated for\n // segmented sort.\n if (std::is_same::value)\n {\n printf(\"Testing segmented sort with overwrite\\n\");\n Test(h_keys,\n h_values,\n num_items,\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit,\n h_reference_keys,\n h_reference_values);\n printf(\"Testing segmented sort with no overwrite\\n\");\n Test(h_keys,\n h_values,\n num_items,\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit,\n h_reference_keys,\n h_reference_values);\n }\n else\n {\n printf(\"Skipping segmented sort tests (NumItemsT != int)\\n\");\n }\n\n if (num_segments == 1)\n {\n printf(\"Testing non-segmented sort with overwrite\\n\");\n Test(h_keys,\n h_values,\n num_items,\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit,\n h_reference_keys,\n h_reference_values);\n if (HasEnoughMemory(static_cast(num_items),\n false))\n {\n printf(\"Testing non-segmented sort with no overwrite\\n\");\n Test(h_keys,\n h_values,\n num_items,\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit,\n h_reference_keys,\n h_reference_values);\n }\n else\n {\n printf(\"Skipping no-overwrite tests with %zd items due to \"\n \"insufficient memory\\n\",\n static_cast(num_items));\n }\n }\n}\n\n\n// Smallest value type for TEST_VALUE_TYPE.\n// Unless TEST_VALUE_TYPE == 3, this is the only value type tested.\n#if TEST_VALUE_TYPE == 0\n// Test keys-only\nusing SmallestValueT = NullType;\n#elif TEST_VALUE_TYPE == 1\n// Test with 8b value\nusing SmallestValueT = unsigned char;\n#elif TEST_VALUE_TYPE == 2\n// Test with 32b value\nusing SmallestValueT = unsigned int;\n// Test with 64b value\n#elif TEST_VALUE_TYPE == 3\nusing SmallestValueT = unsigned long long;\n#endif\n\n\n/**\n * Test value type\n */\ntemplate \nvoid TestValueTypes(\n KeyT *h_keys,\n NumItemsT num_items,\n int num_segments,\n bool pre_sorted,\n NumItemsT *h_segment_offsets,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n // Initialize the solution\n NumItemsT *h_reference_ranks = NULL;\n KeyT *h_reference_keys = NULL;\n // If TEST_VALUE_TYPE == 0, no values are sorted, only keys.\n // Since ranks are only necessary when checking for values,\n // they are not computed in this case.\n InitializeSolution(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, begin_bit, end_bit, h_reference_ranks, h_reference_keys);\n\n TestBackend (h_keys, num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets, begin_bit, end_bit, h_reference_keys, h_reference_ranks);\n\n#if TEST_VALUE_TYPE == 3\n // Test with non-trivially-constructable value\n // These are cheap to build, so lump them in with the 64b value tests.\n TestBackend (h_keys, num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets, begin_bit, end_bit, h_reference_keys, h_reference_ranks);\n#endif\n\n // Cleanup\n ResetKeys(h_keys, num_items, pre_sorted, h_reference_keys);\n if (h_reference_ranks) delete[] h_reference_ranks;\n if (h_reference_keys) delete[] h_reference_keys;\n}\n\n\n\n/**\n * Test ascending/descending\n */\ntemplate \nvoid TestDirection(\n KeyT *h_keys,\n NumItemsT num_items,\n int num_segments,\n bool pre_sorted,\n NumItemsT *h_segment_offsets,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n TestValueTypes(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_begin_offsets, d_segment_end_offsets, begin_bit, end_bit);\n TestValueTypes(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_begin_offsets, d_segment_end_offsets, begin_bit, end_bit);\n}\n\n\n/**\n * Test different bit ranges\n */\ntemplate \nvoid TestBits(\n KeyT *h_keys,\n NumItemsT num_items,\n int num_segments,\n bool pre_sorted,\n NumItemsT *h_segment_offsets,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets)\n{\n // Don't test partial-word sorting for boolean, fp, or signed types (the bit-flipping techniques get in the way) or pre-sorted keys\n if ((Traits::CATEGORY == UNSIGNED_INTEGER)\n && (!std::is_same::value)\n && !pre_sorted)\n {\n // Partial bits\n int begin_bit = 1;\n int end_bit = (sizeof(KeyT) * 8) - 1;\n printf(\"Testing key bits [%d,%d)\\n\", begin_bit, end_bit); fflush(stdout);\n TestDirection(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_begin_offsets, d_segment_end_offsets, begin_bit, end_bit);\n\n // Equal bits\n begin_bit = end_bit = 0;\n printf(\"Testing key bits [%d,%d)\\n\", begin_bit, end_bit); fflush(stdout);\n TestDirection(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_begin_offsets, d_segment_end_offsets, begin_bit, end_bit);\n\n // Across subword boundaries\n int mid_bit = sizeof(KeyT) * 4;\n printf(\"Testing key bits [%d,%d)\\n\", mid_bit - 1, mid_bit + 1); fflush(stdout);\n TestDirection(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_begin_offsets, d_segment_end_offsets, mid_bit - 1, mid_bit + 1);\n }\n\n printf(\"Testing key bits [%d,%d)\\n\", 0, int(sizeof(KeyT)) * 8); fflush(stdout);\n TestDirection(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_begin_offsets, d_segment_end_offsets, 0, sizeof(KeyT) * 8);\n}\n\n\ntemplate\nstruct TransformFunctor1\n{\n __host__ __device__ __forceinline__ OffsetT operator()(OffsetT offset) const\n {\n return offset;\n }\n};\n\ntemplate\nstruct TransformFunctor2\n{\n __host__ __device__ __forceinline__ OffsetT operator()(OffsetT offset) const\n {\n return offset;\n }\n};\n\n\n/**\n* Test different segment iterators\n*/\ntemplate \nvoid TestSegmentIterators(\n KeyT *h_keys,\n NumItemsT num_items,\n int num_segments,\n bool pre_sorted,\n NumItemsT *h_segment_offsets,\n NumItemsT *d_segment_offsets)\n{\n InitializeSegments(num_items, num_segments, h_segment_offsets);\n CubDebugExit(cudaMemcpy(d_segment_offsets, h_segment_offsets, sizeof(NumItemsT) * (num_segments + 1), cudaMemcpyHostToDevice));\n\n // Test with segment pointer.\n // This is also used to test non-segmented sort.\n TestBits(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_offsets, d_segment_offsets + 1);\n\n if (num_segments > 1)\n {\n // Test with transform iterators of different types\n typedef TransformFunctor1 TransformFunctor1T;\n typedef TransformFunctor2 TransformFunctor2T;\n\n TransformInputIterator d_segment_begin_offsets_itr(d_segment_offsets, TransformFunctor1T());\n TransformInputIterator d_segment_end_offsets_itr(d_segment_offsets + 1, TransformFunctor2T());\n\n TestBits(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_begin_offsets_itr, d_segment_end_offsets_itr);\n }\n}\n\n\n/**\n * Test different segment compositions\n */\ntemplate \nvoid TestSegments(\n KeyT *h_keys,\n NumItemsT num_items,\n int max_segments,\n bool pre_sorted)\n{\n max_segments = static_cast(CUB_MIN(num_items, static_cast(max_segments)));\n NumItemsT *h_segment_offsets = new NumItemsT[max_segments + 1];\n\n NumItemsT *d_segment_offsets = nullptr;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_segment_offsets, sizeof(NumItemsT) * (max_segments + 1)));\n\n for (int num_segments = max_segments; num_segments > 1; num_segments = cub::DivideAndRoundUp(num_segments, 64))\n {\n // Pre-sorted tests are not supported for segmented sort\n if (num_items / num_segments < 128 * 1000 && !pre_sorted) {\n // Right now we assign a single thread block to each segment, so lets keep it to under 128K items per segment\n TestSegmentIterators(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_offsets);\n }\n }\n\n // Test single segment\n if (num_items > 0)\n {\n if (num_items < 128 * 1000 || pre_sorted)\n {\n // Right now we assign a single thread block to each segment, so lets\n // keep it to under 128K items per segment\n TestSegmentIterators(h_keys,\n num_items,\n 1,\n pre_sorted,\n h_segment_offsets,\n d_segment_offsets);\n }\n }\n\n if (h_segment_offsets) delete[] h_segment_offsets;\n if (d_segment_offsets) CubDebugExit(g_allocator.DeviceFree(d_segment_offsets));\n}\n\n/** \n * Test different NumItemsT, i.e. types of num_items \n */\ntemplate \nvoid TestNumItems(KeyT *h_keys,\n std::size_t num_items,\n int max_segments,\n bool pre_sorted)\n{\n if (!pre_sorted && num_items <= std::size_t(std::numeric_limits::max()))\n {\n TestSegments(h_keys, static_cast(num_items), max_segments, pre_sorted);\n }\n if (pre_sorted && num_items <= std::size_t(std::numeric_limits::max()))\n {\n TestSegments(h_keys, static_cast(num_items), max_segments, pre_sorted);\n }\n TestSegments(h_keys, num_items, max_segments, pre_sorted);\n}\n\n\n/**\n * Test different (sub)lengths and number of segments\n */\ntemplate \nvoid TestSizes(KeyT* h_keys,\n std::size_t max_items,\n int max_segments,\n bool pre_sorted)\n{\n if (pre_sorted)\n {\n // run a specific list of sizes, up to max_items\n std::size_t sizes[] = {g_smallest_pre_sorted_num_items, 4350000007ull};\n for (std::size_t num_items : sizes)\n {\n if (num_items > max_items) break;\n TestNumItems(h_keys, num_items, max_segments, pre_sorted);\n }\n }\n else\n {\n for (std::size_t num_items = max_items;\n num_items > 1;\n num_items = cub::DivideAndRoundUp(num_items, 64))\n {\n TestNumItems(h_keys, num_items, max_segments, pre_sorted);\n }\n }\n}\n\n/**\n * Test key sampling distributions\n */\ntemplate \nvoid TestGen(\n std::size_t max_items,\n int max_segments)\n{\n if (max_items == ~std::size_t(0))\n {\n max_items = 8000003;\n }\n\n if (max_segments < 0)\n {\n max_segments = 5003;\n }\n\n std::unique_ptr h_keys(new KeyT[max_items]);\n\n // Test trivial problems sizes\n h_keys[0] = static_cast(42);\n TestNumItems(h_keys.get(), 0, 0, false);\n TestNumItems(h_keys.get(), 1, 1, false);\n\n for (int entropy_reduction = 0; entropy_reduction <= 6; entropy_reduction += 6)\n {\n printf(\"\\nTesting random %s keys with entropy reduction factor %d\\n\", typeid(KeyT).name(), entropy_reduction); fflush(stdout);\n InitializeKeyBits(RANDOM, h_keys.get(), max_items, entropy_reduction);\n TestSizes(h_keys.get(), max_items, max_segments, false);\n }\n\n if (cub::Traits::CATEGORY == cub::FLOATING_POINT)\n {\n printf(\"\\nTesting random %s keys with some replaced with -0.0 or +0.0 \\n\", typeid(KeyT).name());\n fflush(stdout);\n InitializeKeyBits(RANDOM_MINUS_PLUS_ZERO, h_keys.get(), max_items, 0);\n // This just tests +/- 0 handling -- don't need to test multiple sizes\n TestNumItems(h_keys.get(), max_items, max_segments, false);\n }\n\n printf(\"\\nTesting uniform %s keys\\n\", typeid(KeyT).name()); fflush(stdout);\n InitializeKeyBits(UNIFORM, h_keys.get(), max_items, 0);\n TestSizes(h_keys.get(), max_items, max_segments, false);\n\n printf(\"\\nTesting natural number %s keys\\n\", typeid(KeyT).name()); fflush(stdout);\n InitializeKeyBits(INTEGER_SEED, h_keys.get(), max_items, 0);\n TestSizes(h_keys.get(), max_items, max_segments, false);\n\n if (WITH_PRE_SORTED)\n {\n // Presorting is only used for testing large input arrays.\n const std::size_t large_num_items = std::size_t(4350000007ull);\n\n // A conservative check for memory, as we don't know ValueT or whether\n // the overwrite is allowed until later.\n // For ValueT, the check is actually exact unless TEST_VALUE_TYPE == 3.\n if (!HasEnoughMemory(large_num_items, true))\n {\n printf(\"Skipping the permutation-based test due to insufficient device memory\\n\");\n return;\n }\n\n h_keys.reset(nullptr); // Explicitly free old buffer before allocating.\n h_keys.reset(new KeyT[large_num_items]);\n\n printf(\"\\nTesting pre-sorted and randomly permuted %s keys\\n\", typeid(KeyT).name());\n fflush(stdout);\n InitializeKeysSorted(h_keys.get(), large_num_items);\n fflush(stdout);\n TestSizes(h_keys.get(), large_num_items, max_segments, true);\n fflush(stdout);\n }\n}\n\n\n//---------------------------------------------------------------------\n// Simple test\n//---------------------------------------------------------------------\n\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT,\n bool IS_DESCENDING>\nvoid Test(\n std::size_t num_items,\n int num_segments,\n GenMode gen_mode,\n int entropy_reduction,\n int begin_bit,\n int end_bit)\n{\n const bool KEYS_ONLY = std::is_same::value;\n\n KeyT *h_keys = new KeyT[num_items];\n std::size_t *h_reference_ranks = NULL;\n KeyT *h_reference_keys = NULL;\n ValueT *h_values = NULL;\n ValueT *h_reference_values = NULL;\n size_t *h_segment_offsets = new std::size_t[num_segments + 1];\n\n std::size_t* d_segment_offsets = nullptr;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_segment_offsets, sizeof(std::size_t) * (num_segments + 1)));\n\n if (end_bit < 0)\n end_bit = sizeof(KeyT) * 8;\n\n InitializeKeyBits(gen_mode, h_keys, num_items, entropy_reduction);\n InitializeSegments(num_items, num_segments, h_segment_offsets);\n CubDebugExit(cudaMemcpy(d_segment_offsets, h_segment_offsets, sizeof(std::size_t) * (num_segments + 1), cudaMemcpyHostToDevice));\n InitializeSolution(\n h_keys, num_items, num_segments, false, h_segment_offsets,\n begin_bit, end_bit, h_reference_ranks, h_reference_keys);\n\n if (!KEYS_ONLY)\n {\n h_values = new ValueT[num_items];\n h_reference_values = new ValueT[num_items];\n\n for (std::size_t i = 0; i < num_items; ++i)\n {\n InitValue(INTEGER_SEED, h_values[i], i);\n InitValue(INTEGER_SEED, h_reference_values[i], h_reference_ranks[i]);\n }\n }\n if (h_reference_ranks) delete[] h_reference_ranks;\n\n printf(\"\\nTesting bits [%d,%d) of %s keys with gen-mode %d\\n\", begin_bit, end_bit, typeid(KeyT).name(), gen_mode); fflush(stdout);\n Test(\n h_keys, h_values,\n num_items, num_segments, d_segment_offsets, d_segment_offsets + 1,\n begin_bit, end_bit, h_reference_keys, h_reference_values);\n\n if (h_keys) delete[] h_keys;\n if (h_reference_keys) delete[] h_reference_keys;\n if (h_values) delete[] h_values;\n if (h_reference_values) delete[] h_reference_values;\n if (h_segment_offsets) delete[] h_segment_offsets;\n if (d_segment_offsets) CubDebugExit(g_allocator.DeviceFree(d_segment_offsets));\n}\n\n#if TEST_VALUE_TYPE == 0\nvoid TestUnspecifiedRanges()\n{\n const std::size_t num_items = 1024 * 1024;\n const std::size_t max_segments = 42;\n const std::size_t avg_segment_size = num_items / max_segments;\n\n for (int iteration = 0; iteration < 4; iteration++)\n {\n thrust::host_vector h_offsets_begin;\n thrust::host_vector h_offsets_end;\n\n h_offsets_begin.reserve(max_segments + 1);\n h_offsets_end.reserve(max_segments + 1);\n\n {\n int offset = 0;\n\n for (std::size_t sid = 0; sid < max_segments; sid++)\n {\n const int segment_size = \n static_cast(RandomValue(avg_segment_size));\n\n const bool segment_is_utilized = segment_size > 0 \n && RandomValue(100) > 60;\n\n if (segment_is_utilized)\n {\n h_offsets_begin.push_back(offset);\n h_offsets_end.push_back(offset + segment_size);\n }\n\n offset += segment_size;\n }\n\n if (h_offsets_begin.empty())\n {\n h_offsets_begin.push_back(avg_segment_size);\n h_offsets_end.push_back(num_items);\n }\n }\n\n thrust::device_vector keys(num_items);\n thrust::device_vector values(num_items);\n\n thrust::sequence(keys.rbegin(), keys.rend());\n thrust::sequence(values.rbegin(), values.rend());\n\n thrust::device_vector d_offsets_begin = h_offsets_begin;\n thrust::device_vector d_offsets_end = h_offsets_end;\n\n thrust::device_vector expected_keys = keys;\n thrust::device_vector expected_values = values;\n\n const int num_segments = static_cast(h_offsets_begin.size());\n\n thrust::device_vector result_keys = keys;\n thrust::device_vector result_values = values;\n\n for (int sid = 0; sid < num_segments; sid++)\n {\n const int segment_begin = h_offsets_begin[sid];\n const int segment_end = h_offsets_end[sid];\n\n thrust::sort_by_key(expected_keys.begin() + segment_begin,\n expected_keys.begin() + segment_end,\n expected_values.begin() + segment_begin);\n }\n\n {\n cub::DoubleBuffer keys_buffer(\n thrust::raw_pointer_cast(keys.data()), \n thrust::raw_pointer_cast(result_keys.data()));\n\n cub::DoubleBuffer values_buffer(\n thrust::raw_pointer_cast(values.data()), \n thrust::raw_pointer_cast(result_values.data()));\n\n std::size_t temp_storage_bytes{};\n std::uint8_t *d_temp_storage{nullptr};\n\n CubDebugExit(cub::DeviceSegmentedRadixSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n keys_buffer, values_buffer, \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data()),\n 0, sizeof(int) * 8));\n\n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n CubDebugExit(cub::DeviceSegmentedRadixSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n keys_buffer, values_buffer, \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data()),\n 0, sizeof(int) * 8));\n\n for (int sid = 0; sid < num_segments; sid++)\n {\n const int segment_begin = h_offsets_begin[sid];\n const int segment_end = h_offsets_end[sid];\n\n if (keys_buffer.selector == 0)\n {\n thrust::copy(\n keys.begin() + segment_begin,\n keys.begin() + segment_end,\n result_keys.begin() + segment_begin);\n }\n \n if (values_buffer.selector == 0)\n {\n thrust::copy(\n values.begin() + segment_begin,\n values.begin() + segment_end,\n result_values.begin() + segment_begin);\n }\n }\n }\n\n AssertEquals(result_keys, expected_keys); \n AssertEquals(result_values, expected_values);\n\n thrust::sequence(keys.rbegin(), keys.rend());\n thrust::sequence(values.rbegin(), values.rend());\n\n result_keys = keys;\n result_values = values;\n\n {\n std::size_t temp_storage_bytes{};\n std::uint8_t *d_temp_storage{};\n\n CubDebugExit(cub::DeviceSegmentedRadixSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n thrust::raw_pointer_cast(keys.data()), \n thrust::raw_pointer_cast(result_keys.data()), \n thrust::raw_pointer_cast(values.data()), \n thrust::raw_pointer_cast(result_values.data()), \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data()),\n 0, sizeof(int) * 8));\n\n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n CubDebugExit(cub::DeviceSegmentedRadixSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n thrust::raw_pointer_cast(keys.data()), \n thrust::raw_pointer_cast(result_keys.data()), \n thrust::raw_pointer_cast(values.data()), \n thrust::raw_pointer_cast(result_values.data()), \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data()),\n 0, sizeof(int) * 8));\n }\n\n AssertEquals(result_values, expected_values);\n AssertEquals(result_keys, expected_keys);\n }\n}\n#endif\n\n#if TEST_KEY_BYTES == 4\n// Following tests check that new decomposer API doesn't break old API. \n// It's disabled because some compilers don't like implicit conversions, which\n// is required for the test. Once we figure out how to temporarily enable conversion, we can\n// re-enable the test.\n#define ENABLING_CONVERSION_IS_FIGURED_OUT 0\n#if ENABLING_CONVERSION_IS_FIGURED_OUT \nstruct bit_selector\n{\n int bit;\n\n operator int() const\n {\n return bit;\n }\n};\n\ntemplate \nvoid device_radix_sort_keys_allows_implicit_conversions_for_bits_helper(BeginBitT begin_bit, EndBitT end_bit, Ts... args)\n{\n const int num_items = 0;\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n cub::DeviceRadixSort::SortKeys(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit,\n end_bit);\n }\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n cub::DeviceRadixSort::SortKeys(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit);\n }\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n cub::DeviceRadixSort::SortKeysDescending(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit,\n end_bit);\n }\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n cub::DeviceRadixSort::SortKeysDescending(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit);\n }\n}\n\ntemplate \nvoid device_radix_sort_pairs_allows_implicit_conversions_for_bits_helper(BeginBitT begin_bit,\n EndBitT end_bit,\n Ts... args)\n{\n const int num_items = 0;\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n cub::DeviceRadixSort::SortPairs(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit,\n end_bit);\n }\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n cub::DeviceRadixSort::SortPairs(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit);\n }\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n cub::DeviceRadixSort::SortPairsDescending(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit,\n end_bit);\n }\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n cub::DeviceRadixSort::SortPairsDescending(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit);\n }\n}\n\ntemplate \nvoid device_radix_sort_allows_implicit_conversions_for_bits(BeginBitT begin_bit, EndBitT end_bit)\n{\n int *d_i_ptr = nullptr;\n\n device_radix_sort_keys_allows_implicit_conversions_for_bits_helper(begin_bit,\n end_bit,\n d_i_ptr,\n d_i_ptr);\n\n cub::DoubleBuffer keys(d_i_ptr, d_i_ptr);\n device_radix_sort_keys_allows_implicit_conversions_for_bits_helper(begin_bit, end_bit, keys);\n\n #if TEST_VALUE_TYPE == 2\n unsigned int *d_u_ptr = nullptr;\n\n device_radix_sort_pairs_allows_implicit_conversions_for_bits_helper(begin_bit,\n end_bit,\n d_i_ptr,\n d_i_ptr,\n d_u_ptr,\n d_u_ptr);\n\n cub::DoubleBuffer pairs(d_u_ptr, d_u_ptr);\n device_radix_sort_pairs_allows_implicit_conversions_for_bits_helper(begin_bit,\n end_bit,\n keys,\n pairs);\n #endif\n}\n\nvoid device_radix_sort_allows_implicit_conversions_for_bits()\n{\n int begin_i = 0;\n long long int begin_lli = 0;\n bit_selector begin_bs{0};\n\n int end_i = 2;\n long long int end_lli = 2;\n bit_selector end_bs{2};\n\n device_radix_sort_allows_implicit_conversions_for_bits(begin_i, end_i);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_i, end_lli);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_i, end_bs);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_lli, end_i);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_lli, end_lli);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_lli, end_bs);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_bs, end_i);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_bs, end_lli);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_bs, end_bs);\n}\n#endif // ENABLING_CONVERSION_IS_FIGURED_OUT \n#endif // TEST_KEY_BYTES == 4\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n std::size_t num_items = ~std::size_t(0);\n int num_segments = -1;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"s\", num_segments);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--s= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // %PARAM% TEST_CDP cdp 0:1\n // %PARAM% TEST_KEY_BYTES bytes 1:2:4:8:16\n // %PARAM% TEST_VALUE_TYPE pairs 0:1:2:3\n // 0->Keys only\n // 1->uchar\n // 2->uint\n // 3->[ull,TestBar] (TestBar is cheap to build, included here to\n // reduce total number of targets)\n\n // To reduce testing time, some key types are only tested when not\n // testing pairs:\n#if TEST_VALUE_TYPE == 0\n#define TEST_EXTENDED_KEY_TYPES\n#endif\n\n // Compile/run thorough tests\n#if TEST_KEY_BYTES == 1\n\n TestGen (num_items, num_segments);\n\n#ifdef TEST_EXTENDED_KEY_TYPES\n TestGen (num_items, num_segments);\n TestGen (num_items, num_segments);\n TestGen (num_items, num_segments);\n#endif // TEST_EXTENDED_KEY_TYPES\n\n#elif TEST_KEY_BYTES == 2\n TestGen (num_items, num_segments);\n\n#ifdef TEST_EXTENDED_KEY_TYPES\n TestGen (num_items, num_segments);\n\n#if !_NVHPC_CUDA\n TestGen (num_items, num_segments);\n#endif // CTK >= 9\n\n#if !_NVHPC_CUDA\n#if !defined(__ICC)\n // Fails with `-0 != 0` with ICC for unknown reasons. See #333.\n TestGen (num_items, num_segments);\n#endif // !ICC\n#endif // CTK >= 11\n\n#endif // TEST_EXTENDED_KEY_TYPES\n\n#elif TEST_KEY_BYTES == 4\n\n TestGen (num_items, num_segments);\n\n#if TEST_VALUE_TYPE == 0\n TestUnspecifiedRanges();\n#endif\n\n#if ENABLING_CONVERSION_IS_FIGURED_OUT \n device_radix_sort_allows_implicit_conversions_for_bits();\n#endif\n\n#ifdef TEST_EXTENDED_KEY_TYPES\n TestGen (num_items, num_segments);\n TestGen (num_items, num_segments);\n#endif // TEST_EXTENDED_KEY_TYPES\n\n#elif TEST_KEY_BYTES == 8\n\n TestGen (num_items, num_segments);\n\n#ifdef TEST_EXTENDED_KEY_TYPES\n TestGen (num_items, num_segments);\n TestGen(num_items, num_segments);\n#endif // TEST_EXTENDED_KEY_TYPES\n\n#elif TEST_KEY_BYTES == 16\n\n#if CUB_IS_INT128_ENABLED \n TestGen<__int128_t, false>(num_items, num_segments);\n TestGen<__uint128_t, false>(num_items, num_segments);\n#else\n // Fix unused static function for MSVC\n BackendToString(CUB);\n#endif\n\n#endif // TEST_KEY_BYTES switch\n\n return 0;\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceRadixSort utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#if !_NVHPC_CUDA\n #include \n#endif\n\n#if !_NVHPC_CUDA\n #include \n#endif\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"test_util.h\"\n\nusing namespace hipcub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\nstd::size_t g_smallest_pre_sorted_num_items = (std::size_t(1) << 32) - 42;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method (allows overwriting of input)\n CUB_NO_OVERWRITE, // CUB method (disallows overwriting of input)\n\n CUB_SEGMENTED, // CUB method (allows overwriting of input)\n CUB_SEGMENTED_NO_OVERWRITE, // CUB method (disallows overwriting of input)\n\n // Same as above, but launches kernels from device using CDP.\n CDP,\n CDP_NO_OVERWRITE,\n CDP_SEGMENTED,\n CDP_SEGMENTED_NO_OVERWRITE,\n};\n\nstatic const char* BackendToString(Backend b)\n{\n switch (b)\n {\n case CUB:\n return \"CUB\";\n case CUB_NO_OVERWRITE:\n return \"CUB_NO_OVERWRITE\";\n case CUB_SEGMENTED:\n return \"CUB_SEGMENTED\";\n case CUB_SEGMENTED_NO_OVERWRITE:\n return \"CUB_SEGMENTED_NO_OVERWRITE\";\n case CDP:\n return \"CDP\";\n case CDP_NO_OVERWRITE:\n return \"CDP_NO_OVERWRITE\";\n case CDP_SEGMENTED:\n return \"CDP_SEGMENTED\";\n case CDP_SEGMENTED_NO_OVERWRITE:\n return \"CDP_SEGMENTED_NO_OVERWRITE\";\n default:\n break;\n }\n\n return \"\";\n}\n\n//---------------------------------------------------------------------\n// Dispatch to different DeviceRadixSort entrypoints\n//---------------------------------------------------------------------\n\n/**\n * Dispatch to CUB sorting entrypoint (specialized for ascending)\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int /*num_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n int begin_bit,\n int end_bit)\n{\n return DeviceRadixSort::SortPairs(\n d_temp_storage, temp_storage_bytes,\n d_keys, d_values,\n num_items, begin_bit, end_bit);\n}\n\n/**\n * Dispatch to CUB_NO_OVERWRITE sorting entrypoint (specialized for ascending)\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int /*num_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n int begin_bit,\n int end_bit)\n{\n KeyT const *const_keys_itr = d_keys.Current();\n ValueT const *const_values_itr = d_values.Current();\n\n hipError_t retval = DeviceRadixSort::SortPairs(\n d_temp_storage, temp_storage_bytes,\n const_keys_itr, d_keys.Alternate(), const_values_itr, d_values.Alternate(),\n num_items, begin_bit, end_bit);\n\n d_keys.selector ^= 1;\n d_values.selector ^= 1;\n return retval;\n}\n\n/**\n * Dispatch to CUB sorting entrypoint (specialized for descending)\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int /*num_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n int begin_bit,\n int end_bit)\n{\n return DeviceRadixSort::SortPairsDescending(\n d_temp_storage, temp_storage_bytes,\n d_keys, d_values, num_items,\n begin_bit, end_bit);\n}\n\n\n/**\n * Dispatch to CUB_NO_OVERWRITE sorting entrypoint (specialized for descending)\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int /*num_segments*/,\n BeginOffsetIteratorT /*d_segment_begin_offsets*/,\n EndOffsetIteratorT /*d_segment_end_offsets*/,\n int begin_bit,\n int end_bit)\n{\n KeyT const *const_keys_itr = d_keys.Current();\n ValueT const *const_values_itr = d_values.Current();\n\n hipError_t retval = DeviceRadixSort::SortPairsDescending(\n d_temp_storage, temp_storage_bytes,\n const_keys_itr, d_keys.Alternate(), const_values_itr, d_values.Alternate(),\n num_items, begin_bit, end_bit);\n\n d_keys.selector ^= 1;\n d_values.selector ^= 1;\n return retval;\n}\n\n//---------------------------------------------------------------------\n// Dispatch to different DeviceRadixSort entrypoints\n//---------------------------------------------------------------------\n\n// Validates that `num_items` fits into `int`\n// TODO(canonizer): remove this check once num_items is templated for segmented sort.\ntemplate \n__host__ __device__ bool ValidateNumItemsForSegmentedSort(NumItemsT num_items)\n{\n if (static_cast(num_items) <\n static_cast(INT_MAX))\n {\n return true;\n }\n else\n {\n printf(\"hipcub::DeviceSegmentedRadixSort is currently limited by %d items but \"\n \"%lld were provided\\n\",\n INT_MAX,\n static_cast(num_items));\n }\n\n return false;\n}\n\n/**\n * Dispatch to CUB_SEGMENTED sorting entrypoint (specialized for ascending)\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n if (ValidateNumItemsForSegmentedSort(num_items))\n {\n return DeviceSegmentedRadixSort::SortPairs(d_temp_storage,\n temp_storage_bytes,\n d_keys,\n d_values,\n static_cast(num_items),\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit);\n }\n\n return hipErrorInvalidValue;\n}\n\n/**\n * Dispatch to CUB_SEGMENTED_NO_OVERWRITE sorting entrypoint (specialized for ascending)\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n if (ValidateNumItemsForSegmentedSort(num_items))\n {\n KeyT const *const_keys_itr = d_keys.Current();\n ValueT const *const_values_itr = d_values.Current();\n\n hipError_t retval =\n DeviceSegmentedRadixSort::SortPairs(d_temp_storage,\n temp_storage_bytes,\n const_keys_itr,\n d_keys.Alternate(),\n const_values_itr,\n d_values.Alternate(),\n static_cast(num_items),\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit);\n\n d_keys.selector ^= 1;\n d_values.selector ^= 1;\n return retval;\n }\n\n return hipErrorInvalidValue;\n}\n\n\n/**\n * Dispatch to CUB_SEGMENTED sorting entrypoint (specialized for descending)\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n if (ValidateNumItemsForSegmentedSort(num_items))\n {\n return DeviceSegmentedRadixSort::SortPairsDescending(\n d_temp_storage,\n temp_storage_bytes,\n d_keys,\n d_values,\n static_cast(num_items),\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit);\n }\n\n return hipErrorInvalidValue;\n}\n\n/**\n * Dispatch to CUB_SEGMENTED_NO_OVERWRITE sorting entrypoint (specialized for descending)\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION\nhipError_t Dispatch(\n Int2Type /*is_descending*/,\n Int2Type /*dispatch_to*/,\n int */*d_selector*/,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t& temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n if (ValidateNumItemsForSegmentedSort(num_items))\n {\n KeyT const *const_keys_itr = d_keys.Current();\n ValueT const *const_values_itr = d_values.Current();\n\n hipError_t retval =\n DeviceSegmentedRadixSort::SortPairsDescending(d_temp_storage,\n temp_storage_bytes,\n const_keys_itr,\n d_keys.Alternate(),\n const_values_itr,\n d_values.Alternate(),\n static_cast(num_items),\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit);\n\n d_keys.selector ^= 1;\n d_values.selector ^= 1;\n return retval;\n }\n\n return hipErrorInvalidValue;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceRadixSort\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type is_descending,\n Int2Type cub_backend,\n int *d_selector,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n void *d_temp_storage,\n size_t temp_storage_bytes,\n DoubleBuffer d_keys,\n DoubleBuffer d_values,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n *d_cdp_error = Dispatch(is_descending,\n cub_backend,\n d_selector,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys,\n d_values,\n num_items,\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n *d_selector = d_keys.selector;\n}\n\n/**\n * Launch kernel and dispatch on device. Should only be called from host code.\n * The CubBackend should be one of the non-CDP CUB backends to invoke from the\n * device.\n */\ntemplate \nhipError_t LaunchCDPKernel(Int2Type is_descending,\n Int2Type cub_backend,\n int *d_selector,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n DoubleBuffer &d_keys,\n DoubleBuffer &d_values,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n // Invoke kernel to invoke device-side dispatch:\n hipError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n is_descending,\n cub_backend,\n d_selector,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys,\n d_values,\n num_items,\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit);\n CubDebugExit(retval);\n CubDebugExit(hipcub::detail::device_synchronize());\n\n // Copy out selector\n CubDebugExit(hipMemcpy(&d_keys.selector,\n d_selector,\n sizeof(int) * 1,\n hipMemcpyDeviceToHost));\n d_values.selector = d_keys.selector;\n\n // Copy out temp_storage_bytes\n CubDebugExit(hipMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n hipMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(hipMemcpy(&retval,\n d_cdp_error,\n sizeof(hipError_t) * 1,\n hipMemcpyDeviceToHost));\n\n return retval;\n}\n\n// Specializations of Dispatch that translate the CDP backend to the appropriate\n// CUB backend, and uses the CUB backend to launch the CDP kernel.\n#define DEFINE_CDP_DISPATCHER(CdpBackend, CubBackend) \\\n template \\\n hipError_t Dispatch(Int2Type is_descending, \\\n Int2Type /*dispatch_to*/, \\\n int *d_selector, \\\n size_t *d_temp_storage_bytes, \\\n hipError_t *d_cdp_error, \\\n \\\n void *d_temp_storage, \\\n size_t &temp_storage_bytes, \\\n DoubleBuffer &d_keys, \\\n DoubleBuffer &d_values, \\\n NumItemsT num_items, \\\n int num_segments, \\\n BeginOffsetIteratorT d_segment_begin_offsets, \\\n EndOffsetIteratorT d_segment_end_offsets, \\\n int begin_bit, \\\n int end_bit) \\\n { \\\n Int2Type cub_backend{}; \\\n return LaunchCDPKernel(is_descending, \\\n cub_backend, \\\n d_selector, \\\n d_temp_storage_bytes, \\\n d_cdp_error, \\\n d_temp_storage, \\\n temp_storage_bytes, \\\n d_keys, \\\n d_values, \\\n num_items, \\\n num_segments, \\\n d_segment_begin_offsets, \\\n d_segment_end_offsets, \\\n begin_bit, \\\n end_bit); \\\n }\n\nDEFINE_CDP_DISPATCHER(CDP, CUB)\nDEFINE_CDP_DISPATCHER(CDP_NO_OVERWRITE, CUB_NO_OVERWRITE)\nDEFINE_CDP_DISPATCHER(CDP_SEGMENTED, CUB_SEGMENTED)\nDEFINE_CDP_DISPATCHER(CDP_SEGMENTED_NO_OVERWRITE, CUB_SEGMENTED_NO_OVERWRITE)\n\n#undef DEFINE_CDP_DISPATCHER\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Problem generation\n//---------------------------------------------------------------------\n\n\n/**\n * Simple key-value pairing\n */\ntemplate <\n typename KeyT,\n typename ValueT>\nstruct Pair\n{\n KeyT key;\n ValueT value;\n\n bool operator<(const Pair &b) const\n {\n return (key < b.key);\n }\n};\n\n\n/**\n * Simple key-value pairing (specialized for bool types)\n */\ntemplate \nstruct Pair\n{\n bool key;\n ValueT value;\n\n bool operator<(const Pair &b) const\n {\n return (!key && b.key);\n }\n};\n\n\n/**\n * Initialize key data\n */\ntemplate \nvoid InitializeKeyBits(\n GenMode gen_mode,\n KeyT *h_keys,\n NumItemsT num_items,\n int /*entropy_reduction*/)\n{\n for (NumItemsT i = 0; i < num_items; ++i)\n InitValue(gen_mode, h_keys[i], i);\n}\n\ntemplate ::UnsignedBits>\nUnsignedBits KeyBits(KeyT key)\n{\n UnsignedBits bits;\n memcpy(&bits, &key, sizeof(KeyT));\n return bits;\n}\n\n/** Initialize the reference array monotonically. */\ntemplate \nvoid InitializeKeysSorted(\n KeyT *h_keys,\n NumItemsT num_items)\n{\n using TraitsT = hipcub::Traits;\n using UnsignedBits = typename TraitsT::UnsignedBits;\n\n // Numbers to generate random runs.\n UnsignedBits max_inc = 1 << (sizeof(UnsignedBits) < 4 ? 3 :\n (sizeof(UnsignedBits) < 8 ? 14 : 24));\n UnsignedBits min_bits = TraitsT::TwiddleIn(KeyBits(TraitsT::Lowest()));\n UnsignedBits max_bits = TraitsT::TwiddleIn(KeyBits(TraitsT::Max()));\n NumItemsT max_run = std::max(\n NumItemsT(double(num_items) * (max_inc + 1) / max_bits),\n NumItemsT(1 << 14));\n\n UnsignedBits *h_key_bits = reinterpret_cast(h_keys);\n NumItemsT i = 0;\n // Start with the minimum twiddled key.\n UnsignedBits twiddled_key = min_bits;\n while (i < num_items)\n {\n // Generate random increment (avoid overflow).\n UnsignedBits inc_bits = 0;\n RandomBits(inc_bits);\n // twiddled_key < max_bits at this point.\n UnsignedBits inc = static_cast(std::min(1 + inc_bits % max_inc, max_bits - twiddled_key));\n twiddled_key += inc;\n\n // Generate random run length (ensure there are enough values to fill the rest).\n NumItemsT run_bits = 0;\n RandomBits(run_bits);\n NumItemsT run_length = std::min(1 + run_bits % max_run, num_items - i);\n if (twiddled_key == max_bits) run_length = num_items - i;\n NumItemsT run_end = i + run_length;\n\n // Fill the array.\n UnsignedBits key = TraitsT::TwiddleOut(twiddled_key);\n // Avoid -0.0 for floating-point keys.\n UnsignedBits negative_zero = UnsignedBits(1) << UnsignedBits(sizeof(UnsignedBits) * 8 - 1);\n if (TraitsT::CATEGORY == hipcub::FLOATING_POINT && key == negative_zero)\n {\n key = 0;\n }\n\n for (; i < run_end; ++i)\n {\n h_key_bits[i] = key;\n }\n }\n}\n\n\n/**\n * Initialize solution\n */\ntemplate \nvoid InitializeSolution(\n KeyT *h_keys,\n NumItemsT num_items,\n int num_segments,\n bool pre_sorted,\n NumItemsT *h_segment_offsets,\n int begin_bit,\n int end_bit,\n NumItemsT *&h_reference_ranks,\n KeyT *&h_reference_keys)\n{\n if (num_items == 0)\n {\n h_reference_ranks = nullptr;\n h_reference_keys = nullptr;\n return;\n }\n\n if (pre_sorted)\n {\n printf(\"Shuffling reference solution on CPU\\n\");\n // Note: begin_bit and end_bit are ignored here, and assumed to have the\n // default values (begin_bit == 0, end_bit == 8 * sizeof(KeyT)).\n // Otherwise, pre-sorting won't work, as it doesn't necessarily\n // correspond to the order of keys sorted by a subrange of bits.\n // num_segments is also ignored as assumed to be 1, as pre-sorted tests\n // are currently not supported for multiple segments.\n //\n // Pre-sorted tests with non-default begin_bit, end_bit or num_segments\n // != 1 are skipped in TestBits() and TestSegments(), respectively.\n AssertEquals(begin_bit, 0);\n AssertEquals(end_bit, static_cast(8 * sizeof(KeyT)));\n AssertEquals(num_segments, 1);\n\n // Copy to the reference solution.\n h_reference_keys = new KeyT[num_items];\n if (IS_DESCENDING)\n {\n // Copy in reverse.\n for (NumItemsT i = 0; i < num_items; ++i)\n {\n h_reference_keys[i] = h_keys[num_items - 1 - i];\n }\n // Copy back.\n memcpy(h_keys, h_reference_keys, num_items * sizeof(KeyT));\n }\n else\n {\n memcpy(h_reference_keys, h_keys, num_items * sizeof(KeyT));\n }\n\n // Summarize the pre-sorted array (element, 1st position, count).\n struct Element\n {\n KeyT key;\n NumItemsT num;\n NumItemsT index;\n };\n\n std::vector summary;\n KeyT cur_key = h_reference_keys[0];\n summary.push_back(Element{cur_key, 1, 0});\n for (NumItemsT i = 1; i < num_items; ++i)\n {\n KeyT key = h_reference_keys[i];\n if (key == cur_key)\n {\n // Same key.\n summary.back().num++;\n continue;\n }\n\n // Different key.\n cur_key = key;\n summary.push_back(Element{cur_key, 1, i});\n }\n\n // Generate a random permutation from the summary. Such a complicated\n // approach is used to permute the array and compute ranks in a\n // cache-friendly way and in a short time.\n if (WANT_RANKS)\n {\n h_reference_ranks = new NumItemsT[num_items];\n }\n NumItemsT max_run = 32;\n NumItemsT run = 0;\n NumItemsT i = 0;\n while (summary.size() > 0)\n {\n // Pick up a random element and a run.\n NumItemsT bits = 0;\n RandomBits(bits);\n NumItemsT summary_id = bits % summary.size();\n Element& element = summary[summary_id];\n run = std::min(1 + bits % (max_run - 1), element.num);\n for (NumItemsT j = 0; j < run; ++j)\n {\n h_keys[i + j] = element.key;\n if (WANT_RANKS)\n {\n h_reference_ranks[element.index + j] = i + j;\n }\n }\n i += run;\n element.index += run;\n element.num -= run;\n if (element.num == 0)\n {\n // Remove the empty entry.\n std::swap(summary[summary_id], summary.back());\n summary.pop_back();\n }\n }\n printf(\" Done.\\n\");\n }\n else\n {\n typedef Pair PairT;\n\n PairT *h_pairs = new PairT[num_items];\n\n int num_bits = end_bit - begin_bit;\n for (NumItemsT i = 0; i < num_items; ++i)\n {\n\n // Mask off unwanted portions\n if (num_bits < static_cast(sizeof(KeyT) * 8))\n {\n using UnsignedBits = typename hipcub::Traits::UnsignedBits;\n\n UnsignedBits base = 0;\n memcpy(&base, &h_keys[i], sizeof(KeyT));\n base &= ((UnsignedBits{1} << num_bits) - 1) << begin_bit;\n memcpy(&h_pairs[i].key, &base, sizeof(KeyT));\n }\n else\n {\n h_pairs[i].key = h_keys[i];\n }\n\n h_pairs[i].value = i;\n }\n\n printf(\"\\nSorting reference solution on CPU \"\n \"(%zd items, %d segments, %zd items/seg)...\",\n static_cast(num_items),\n num_segments,\n static_cast(num_items / num_segments));\n fflush(stdout);\n\n for (int i = 0; i < num_segments; ++i)\n {\n if (IS_DESCENDING) std::reverse(h_pairs + h_segment_offsets[i], h_pairs + h_segment_offsets[i + 1]);\n std::stable_sort( h_pairs + h_segment_offsets[i], h_pairs + h_segment_offsets[i + 1]);\n if (IS_DESCENDING) std::reverse(h_pairs + h_segment_offsets[i], h_pairs + h_segment_offsets[i + 1]);\n }\n\n printf(\" Done.\\n\"); fflush(stdout);\n\n if (WANT_RANKS)\n {\n h_reference_ranks = new NumItemsT[num_items];\n }\n h_reference_keys = new KeyT[num_items];\n\n for (NumItemsT i = 0; i < num_items; ++i)\n {\n if (WANT_RANKS)\n {\n h_reference_ranks[i] = h_pairs[i].value;\n }\n h_reference_keys[i] = h_keys[h_pairs[i].value];\n }\n\n if (h_pairs) delete[] h_pairs;\n }\n}\n\ntemplate \nvoid ResetKeys(KeyT *h_keys, NumItemsT num_items, bool pre_sorted, KeyT *reference_keys)\n{\n if (!pre_sorted) return;\n\n // Copy the reference keys back.\n if (IS_DESCENDING)\n {\n // Keys need to be copied in reverse.\n for (NumItemsT i = 0; i < num_items; ++i)\n {\n h_keys[i] = reference_keys[num_items - 1 - i];\n }\n }\n else\n {\n memcpy(h_keys, reference_keys, num_items * sizeof(KeyT));\n }\n}\n\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\ntemplate \nstruct UnwrapHalfAndBfloat16 {\n using Type = T;\n};\n\n#if !_NVHPC_CUDA\ntemplate <>\nstruct UnwrapHalfAndBfloat16 {\n using Type = __half;\n};\n#endif\n\n#if !_NVHPC_CUDA\ntemplate <>\nstruct UnwrapHalfAndBfloat16 {\n using Type = __hip_bfloat16;\n};\n#endif\n\ntemplate \nint compare_device_arrays(T* host_reference, T* d_tmp_buffer, T* d_data, NumItemsT num_items)\n{\n CubDebugExit(hipMemcpy(d_tmp_buffer, host_reference, sizeof(T) * num_items, hipMemcpyHostToDevice));\n\n auto d_reference_begin = d_tmp_buffer;\n auto d_reference_end = d_tmp_buffer + num_items;\n auto err = thrust::mismatch(thrust::device, d_reference_begin, d_reference_end, d_data);\n\n if (err.first != d_reference_end)\n {\n const auto index = thrust::distance(d_reference_begin, err.first);\n return CompareDeviceResults(host_reference + index, d_data + index, 1, true, g_verbose);\n }\n\n return 0;\n}\n\ntemplate \nint compare_device_arrays(\n CUB_NS_QUALIFIER::NullType */* h_reference */,\n CUB_NS_QUALIFIER::NullType */* d_tmp_buffer */,\n CUB_NS_QUALIFIER::NullType */* d_data */,\n NumItemsT /* num_items */)\n{\n return 0;\n}\n\n/**\n * Test DeviceRadixSort\n */\ntemplate <\n Backend BACKEND,\n bool IS_DESCENDING,\n typename KeyT,\n typename ValueT,\n typename BeginOffsetIteratorT,\n typename EndOffsetIteratorT,\n typename NumItemsT>\nvoid Test(\n KeyT *h_keys,\n ValueT *h_values,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit,\n KeyT *h_reference_keys,\n ValueT *h_reference_values)\n{\n // Key alias type\n using KeyAliasT = typename UnwrapHalfAndBfloat16::Type;\n\n const bool KEYS_ONLY = std::is_same::value;\n\n printf(\"%s %s hipcub::DeviceRadixSort %zd items, %d segments, \"\n \"%d-byte keys (%s) %d-byte values (%s), %d-byte num_items (%s), \"\n \"descending %d, begin_bit %d, end_bit %d\\n\",\n BackendToString(BACKEND),\n (KEYS_ONLY) ? \"keys-only\" : \"key-value\",\n static_cast(num_items),\n num_segments,\n static_cast(sizeof(KeyT)),\n typeid(KeyT).name(),\n (KEYS_ONLY) ? 0 : static_cast(sizeof(ValueT)),\n typeid(ValueT).name(),\n static_cast(sizeof(NumItemsT)),\n typeid(NumItemsT).name(),\n IS_DESCENDING,\n begin_bit,\n end_bit);\n\n if (g_verbose)\n {\n printf(\"Input keys:\\n\");\n DisplayResults(h_keys, num_items);\n printf(\"\\n\\n\");\n }\n\n // Allocate device arrays\n DoubleBuffer d_keys;\n DoubleBuffer d_values;\n int *d_selector;\n size_t *d_temp_storage_bytes;\n hipError_t *d_cdp_error;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys.d_buffers[0], sizeof(KeyT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys.d_buffers[1], sizeof(KeyT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_selector, sizeof(int) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(hipError_t) * 1));\n if (!KEYS_ONLY)\n {\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_values.d_buffers[0], sizeof(ValueT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_values.d_buffers[1], sizeof(ValueT) * num_items));\n }\n\n // Allocate temporary storage (and make it un-aligned)\n size_t temp_storage_bytes = 0;\n void *d_temp_storage = NULL;\n CubDebugExit(Dispatch(\n Int2Type(), Int2Type(), d_selector, d_temp_storage_bytes, d_cdp_error,\n d_temp_storage, temp_storage_bytes, d_keys, d_values,\n num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets,\n begin_bit, end_bit));\n\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes + 1));\n void* mis_aligned_temp = static_cast(d_temp_storage) + 1;\n\n // Initialize/clear device arrays\n d_keys.selector = 0;\n CubDebugExit(hipMemcpy(d_keys.d_buffers[0], h_keys, sizeof(KeyT) * num_items, hipMemcpyHostToDevice));\n CubDebugExit(hipMemset(d_keys.d_buffers[1], 0, sizeof(KeyT) * num_items));\n if (!KEYS_ONLY)\n {\n d_values.selector = 0;\n CubDebugExit(hipMemcpy(d_values.d_buffers[0], h_values, sizeof(ValueT) * num_items, hipMemcpyHostToDevice));\n CubDebugExit(hipMemset(d_values.d_buffers[1], 0, sizeof(ValueT) * num_items));\n }\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(\n Int2Type(), Int2Type(), d_selector, d_temp_storage_bytes, d_cdp_error,\n mis_aligned_temp, temp_storage_bytes, d_keys, d_values,\n num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets,\n begin_bit, end_bit));\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Check for correctness (and display results, if specified)\n printf(\"Warmup done. Checking results:\\n\"); fflush(stdout);\n\n int compare = 0;\n\n // If in/out API is used, we are not allowed to overwrite the input. \n // Let's check that the input buffer is not overwritten by the algorithm.\n if (BACKEND == CUB_NO_OVERWRITE)\n {\n KeyT *d_input_keys = reinterpret_cast(d_keys.d_buffers[0]);\n\n if (temp_storage_bytes < sizeof(KeyT) * num_items)\n {\n // For small input sizes, temporary storage is not large enough to fit keys.\n compare = CompareDeviceResults(h_keys, d_input_keys, num_items, true, g_verbose);\n }\n else \n {\n // If overwrite is not allowed, temporary storage is large enough to fit keys.\n KeyT* temp_keys = reinterpret_cast(d_temp_storage);\n\n compare = compare_device_arrays(h_keys, temp_keys, d_input_keys, num_items);\n printf(\"\\t Compare input keys: %s \", compare ? \"FAIL\" : \"PASS\"); fflush(stdout);\n }\n }\n\n // After the previous check is done, we can safely reuse alternative buffer to store \n // the reference results and compare current output.\n compare |= compare_device_arrays(h_reference_keys,\n reinterpret_cast(d_keys.Alternate()),\n reinterpret_cast(d_keys.Current()),\n num_items);\n\n printf(\"\\t Compare keys (selector %d): %s \", d_keys.selector, compare ? \"FAIL\" : \"PASS\"); fflush(stdout);\n\n if (!KEYS_ONLY)\n {\n int values_compare = compare_device_arrays(h_reference_values,\n d_values.Alternate(),\n d_values.Current(),\n num_items);\n\n compare |= values_compare;\n printf(\"\\t Compare values (selector %d): %s \", d_values.selector, values_compare ? \"FAIL\" : \"PASS\"); fflush(stdout);\n }\n\n // Performance\n if (g_timing_iterations)\n printf(\"\\nPerforming timing iterations:\\n\"); fflush(stdout);\n\n GpuTimer gpu_timer;\n float elapsed_millis = 0.0f;\n for (int i = 0; i < g_timing_iterations; ++i)\n {\n // Initialize/clear device arrays\n CubDebugExit(hipMemcpy(d_keys.d_buffers[d_keys.selector], h_keys, sizeof(KeyT) * num_items, hipMemcpyHostToDevice));\n CubDebugExit(hipMemset(d_keys.d_buffers[d_keys.selector ^ 1], 0, sizeof(KeyT) * num_items));\n if (!KEYS_ONLY)\n {\n CubDebugExit(hipMemcpy(d_values.d_buffers[d_values.selector], h_values, sizeof(ValueT) * num_items, hipMemcpyHostToDevice));\n CubDebugExit(hipMemset(d_values.d_buffers[d_values.selector ^ 1], 0, sizeof(ValueT) * num_items));\n }\n\n gpu_timer.Start();\n CubDebugExit(Dispatch(\n Int2Type(), Int2Type(), d_selector, d_temp_storage_bytes, d_cdp_error,\n mis_aligned_temp, temp_storage_bytes, d_keys, d_values,\n num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets,\n begin_bit, end_bit));\n gpu_timer.Stop();\n elapsed_millis += gpu_timer.ElapsedMillis();\n }\n\n // Display performance\n if (g_timing_iterations > 0)\n {\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n float giga_bandwidth = (KEYS_ONLY) ?\n giga_rate * sizeof(KeyT) * 2 :\n giga_rate * (sizeof(KeyT) + sizeof(ValueT)) * 2;\n printf(\"\\n%.3f elapsed ms, %.3f avg ms, %.3f billion items/s, %.3f logical GB/s\", elapsed_millis, avg_millis, giga_rate, giga_bandwidth);\n }\n\n printf(\"\\n\\n\");\n\n // Cleanup\n if (d_keys.d_buffers[0]) CubDebugExit(g_allocator.DeviceFree(d_keys.d_buffers[0]));\n if (d_keys.d_buffers[1]) CubDebugExit(g_allocator.DeviceFree(d_keys.d_buffers[1]));\n if (d_values.d_buffers[0]) CubDebugExit(g_allocator.DeviceFree(d_values.d_buffers[0]));\n if (d_values.d_buffers[1]) CubDebugExit(g_allocator.DeviceFree(d_values.d_buffers[1]));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_selector) CubDebugExit(g_allocator.DeviceFree(d_selector));\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n\n // Correctness asserts\n AssertEquals(0, compare);\n}\n\n// Returns whether there is enough memory for the test.\ntemplate \nbool HasEnoughMemory(std::size_t num_items, bool overwrite)\n{\n std::size_t total_mem = TotalGlobalMem();\n std::size_t value_size = std::is_same::value\n ? 0\n : sizeof(ValueT);\n // A conservative estimate of the amount of memory required.\n double factor = overwrite ? 2.25 : 3.25;\n std::size_t test_mem = static_cast\n (num_items * (sizeof(KeyT) + value_size) * factor);\n return test_mem < total_mem;\n}\n\n\n/**\n * Test backend\n */\ntemplate \nvoid TestBackend(KeyT *h_keys,\n NumItemsT num_items,\n int num_segments,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit,\n KeyT *h_reference_keys,\n NumItemsT *h_reference_ranks)\n{\n#if TEST_CDP == 0\n constexpr auto NonSegmentedOverwrite = CUB;\n constexpr auto NonSegmentedNoOverwrite = CUB_NO_OVERWRITE;\n constexpr auto SegmentedOverwrite = CUB_SEGMENTED;\n constexpr auto SegmentedNoOverwrite = CUB_SEGMENTED_NO_OVERWRITE;\n#else // TEST_CDP\n constexpr auto NonSegmentedOverwrite = CDP;\n constexpr auto NonSegmentedNoOverwrite = CDP_NO_OVERWRITE;\n constexpr auto SegmentedOverwrite = CDP_SEGMENTED;\n constexpr auto SegmentedNoOverwrite = CDP_SEGMENTED_NO_OVERWRITE;\n#endif // TEST_CDP\n\n const bool KEYS_ONLY = std::is_same::value;\n\n // A conservative check assuming overwrite is allowed.\n if (!HasEnoughMemory(static_cast(num_items), true))\n {\n printf(\"Skipping the test due to insufficient device memory\\n\");\n return;\n }\n\n std::unique_ptr h_value_data{};\n\n ValueT *h_values = nullptr;\n ValueT *h_reference_values = nullptr;\n\n if (!KEYS_ONLY)\n {\n h_value_data.reset(new ValueT[2 * static_cast(num_items)]);\n h_values = h_value_data.get();\n h_reference_values = h_value_data.get() + num_items;\n\n for (NumItemsT i = 0; i < num_items; ++i)\n {\n InitValue(INTEGER_SEED, h_values[i], i);\n InitValue(INTEGER_SEED, h_reference_values[i], h_reference_ranks[i]);\n }\n }\n\n // Skip segmented sort if num_items isn't int.\n // TODO(64bit-seg-sort): re-enable these tests once num_items is templated for\n // segmented sort.\n if (std::is_same::value)\n {\n printf(\"Testing segmented sort with overwrite\\n\");\n Test(h_keys,\n h_values,\n num_items,\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit,\n h_reference_keys,\n h_reference_values);\n printf(\"Testing segmented sort with no overwrite\\n\");\n Test(h_keys,\n h_values,\n num_items,\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit,\n h_reference_keys,\n h_reference_values);\n }\n else\n {\n printf(\"Skipping segmented sort tests (NumItemsT != int)\\n\");\n }\n\n if (num_segments == 1)\n {\n printf(\"Testing non-segmented sort with overwrite\\n\");\n Test(h_keys,\n h_values,\n num_items,\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit,\n h_reference_keys,\n h_reference_values);\n if (HasEnoughMemory(static_cast(num_items),\n false))\n {\n printf(\"Testing non-segmented sort with no overwrite\\n\");\n Test(h_keys,\n h_values,\n num_items,\n num_segments,\n d_segment_begin_offsets,\n d_segment_end_offsets,\n begin_bit,\n end_bit,\n h_reference_keys,\n h_reference_values);\n }\n else\n {\n printf(\"Skipping no-overwrite tests with %zd items due to \"\n \"insufficient memory\\n\",\n static_cast(num_items));\n }\n }\n}\n\n\n// Smallest value type for TEST_VALUE_TYPE.\n// Unless TEST_VALUE_TYPE == 3, this is the only value type tested.\n#if TEST_VALUE_TYPE == 0\n// Test keys-only\nusing SmallestValueT = NullType;\n#elif TEST_VALUE_TYPE == 1\n// Test with 8b value\nusing SmallestValueT = unsigned char;\n#elif TEST_VALUE_TYPE == 2\n// Test with 32b value\nusing SmallestValueT = unsigned int;\n// Test with 64b value\n#elif TEST_VALUE_TYPE == 3\nusing SmallestValueT = unsigned long long;\n#endif\n\n\n/**\n * Test value type\n */\ntemplate \nvoid TestValueTypes(\n KeyT *h_keys,\n NumItemsT num_items,\n int num_segments,\n bool pre_sorted,\n NumItemsT *h_segment_offsets,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n // Initialize the solution\n NumItemsT *h_reference_ranks = NULL;\n KeyT *h_reference_keys = NULL;\n // If TEST_VALUE_TYPE == 0, no values are sorted, only keys.\n // Since ranks are only necessary when checking for values,\n // they are not computed in this case.\n InitializeSolution(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, begin_bit, end_bit, h_reference_ranks, h_reference_keys);\n\n TestBackend (h_keys, num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets, begin_bit, end_bit, h_reference_keys, h_reference_ranks);\n\n#if TEST_VALUE_TYPE == 3\n // Test with non-trivially-constructable value\n // These are cheap to build, so lump them in with the 64b value tests.\n TestBackend (h_keys, num_items, num_segments, d_segment_begin_offsets, d_segment_end_offsets, begin_bit, end_bit, h_reference_keys, h_reference_ranks);\n#endif\n\n // Cleanup\n ResetKeys(h_keys, num_items, pre_sorted, h_reference_keys);\n if (h_reference_ranks) delete[] h_reference_ranks;\n if (h_reference_keys) delete[] h_reference_keys;\n}\n\n\n\n/**\n * Test ascending/descending\n */\ntemplate \nvoid TestDirection(\n KeyT *h_keys,\n NumItemsT num_items,\n int num_segments,\n bool pre_sorted,\n NumItemsT *h_segment_offsets,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets,\n int begin_bit,\n int end_bit)\n{\n TestValueTypes(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_begin_offsets, d_segment_end_offsets, begin_bit, end_bit);\n TestValueTypes(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_begin_offsets, d_segment_end_offsets, begin_bit, end_bit);\n}\n\n\n/**\n * Test different bit ranges\n */\ntemplate \nvoid TestBits(\n KeyT *h_keys,\n NumItemsT num_items,\n int num_segments,\n bool pre_sorted,\n NumItemsT *h_segment_offsets,\n BeginOffsetIteratorT d_segment_begin_offsets,\n EndOffsetIteratorT d_segment_end_offsets)\n{\n // Don't test partial-word sorting for boolean, fp, or signed types (the bit-flipping techniques get in the way) or pre-sorted keys\n if ((Traits::CATEGORY == UNSIGNED_INTEGER)\n && (!std::is_same::value)\n && !pre_sorted)\n {\n // Partial bits\n int begin_bit = 1;\n int end_bit = (sizeof(KeyT) * 8) - 1;\n printf(\"Testing key bits [%d,%d)\\n\", begin_bit, end_bit); fflush(stdout);\n TestDirection(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_begin_offsets, d_segment_end_offsets, begin_bit, end_bit);\n\n // Equal bits\n begin_bit = end_bit = 0;\n printf(\"Testing key bits [%d,%d)\\n\", begin_bit, end_bit); fflush(stdout);\n TestDirection(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_begin_offsets, d_segment_end_offsets, begin_bit, end_bit);\n\n // Across subword boundaries\n int mid_bit = sizeof(KeyT) * 4;\n printf(\"Testing key bits [%d,%d)\\n\", mid_bit - 1, mid_bit + 1); fflush(stdout);\n TestDirection(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_begin_offsets, d_segment_end_offsets, mid_bit - 1, mid_bit + 1);\n }\n\n printf(\"Testing key bits [%d,%d)\\n\", 0, int(sizeof(KeyT)) * 8); fflush(stdout);\n TestDirection(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_begin_offsets, d_segment_end_offsets, 0, sizeof(KeyT) * 8);\n}\n\n\ntemplate\nstruct TransformFunctor1\n{\n __host__ __device__ __forceinline__ OffsetT operator()(OffsetT offset) const\n {\n return offset;\n }\n};\n\ntemplate\nstruct TransformFunctor2\n{\n __host__ __device__ __forceinline__ OffsetT operator()(OffsetT offset) const\n {\n return offset;\n }\n};\n\n\n/**\n* Test different segment iterators\n*/\ntemplate \nvoid TestSegmentIterators(\n KeyT *h_keys,\n NumItemsT num_items,\n int num_segments,\n bool pre_sorted,\n NumItemsT *h_segment_offsets,\n NumItemsT *d_segment_offsets)\n{\n InitializeSegments(num_items, num_segments, h_segment_offsets);\n CubDebugExit(hipMemcpy(d_segment_offsets, h_segment_offsets, sizeof(NumItemsT) * (num_segments + 1), hipMemcpyHostToDevice));\n\n // Test with segment pointer.\n // This is also used to test non-segmented sort.\n TestBits(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_offsets, d_segment_offsets + 1);\n\n if (num_segments > 1)\n {\n // Test with transform iterators of different types\n typedef TransformFunctor1 TransformFunctor1T;\n typedef TransformFunctor2 TransformFunctor2T;\n\n TransformInputIterator d_segment_begin_offsets_itr(d_segment_offsets, TransformFunctor1T());\n TransformInputIterator d_segment_end_offsets_itr(d_segment_offsets + 1, TransformFunctor2T());\n\n TestBits(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_begin_offsets_itr, d_segment_end_offsets_itr);\n }\n}\n\n\n/**\n * Test different segment compositions\n */\ntemplate \nvoid TestSegments(\n KeyT *h_keys,\n NumItemsT num_items,\n int max_segments,\n bool pre_sorted)\n{\n max_segments = static_cast(CUB_MIN(num_items, static_cast(max_segments)));\n NumItemsT *h_segment_offsets = new NumItemsT[max_segments + 1];\n\n NumItemsT *d_segment_offsets = nullptr;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_segment_offsets, sizeof(NumItemsT) * (max_segments + 1)));\n\n for (int num_segments = max_segments; num_segments > 1; num_segments = hipcub::DivideAndRoundUp(num_segments, 64))\n {\n // Pre-sorted tests are not supported for segmented sort\n if (num_items / num_segments < 128 * 1000 && !pre_sorted) {\n // Right now we assign a single thread block to each segment, so lets keep it to under 128K items per segment\n TestSegmentIterators(h_keys, num_items, num_segments, pre_sorted, h_segment_offsets, d_segment_offsets);\n }\n }\n\n // Test single segment\n if (num_items > 0)\n {\n if (num_items < 128 * 1000 || pre_sorted)\n {\n // Right now we assign a single thread block to each segment, so lets\n // keep it to under 128K items per segment\n TestSegmentIterators(h_keys,\n num_items,\n 1,\n pre_sorted,\n h_segment_offsets,\n d_segment_offsets);\n }\n }\n\n if (h_segment_offsets) delete[] h_segment_offsets;\n if (d_segment_offsets) CubDebugExit(g_allocator.DeviceFree(d_segment_offsets));\n}\n\n/** \n * Test different NumItemsT, i.e. types of num_items \n */\ntemplate \nvoid TestNumItems(KeyT *h_keys,\n std::size_t num_items,\n int max_segments,\n bool pre_sorted)\n{\n if (!pre_sorted && num_items <= std::size_t(std::numeric_limits::max()))\n {\n TestSegments(h_keys, static_cast(num_items), max_segments, pre_sorted);\n }\n if (pre_sorted && num_items <= std::size_t(std::numeric_limits::max()))\n {\n TestSegments(h_keys, static_cast(num_items), max_segments, pre_sorted);\n }\n TestSegments(h_keys, num_items, max_segments, pre_sorted);\n}\n\n\n/**\n * Test different (sub)lengths and number of segments\n */\ntemplate \nvoid TestSizes(KeyT* h_keys,\n std::size_t max_items,\n int max_segments,\n bool pre_sorted)\n{\n if (pre_sorted)\n {\n // run a specific list of sizes, up to max_items\n std::size_t sizes[] = {g_smallest_pre_sorted_num_items, 4350000007ull};\n for (std::size_t num_items : sizes)\n {\n if (num_items > max_items) break;\n TestNumItems(h_keys, num_items, max_segments, pre_sorted);\n }\n }\n else\n {\n for (std::size_t num_items = max_items;\n num_items > 1;\n num_items = hipcub::DivideAndRoundUp(num_items, 64))\n {\n TestNumItems(h_keys, num_items, max_segments, pre_sorted);\n }\n }\n}\n\n/**\n * Test key sampling distributions\n */\ntemplate \nvoid TestGen(\n std::size_t max_items,\n int max_segments)\n{\n if (max_items == ~std::size_t(0))\n {\n max_items = 8000003;\n }\n\n if (max_segments < 0)\n {\n max_segments = 5003;\n }\n\n std::unique_ptr h_keys(new KeyT[max_items]);\n\n // Test trivial problems sizes\n h_keys[0] = static_cast(42);\n TestNumItems(h_keys.get(), 0, 0, false);\n TestNumItems(h_keys.get(), 1, 1, false);\n\n for (int entropy_reduction = 0; entropy_reduction <= 6; entropy_reduction += 6)\n {\n printf(\"\\nTesting random %s keys with entropy reduction factor %d\\n\", typeid(KeyT).name(), entropy_reduction); fflush(stdout);\n InitializeKeyBits(RANDOM, h_keys.get(), max_items, entropy_reduction);\n TestSizes(h_keys.get(), max_items, max_segments, false);\n }\n\n if (hipcub::Traits::CATEGORY == hipcub::FLOATING_POINT)\n {\n printf(\"\\nTesting random %s keys with some replaced with -0.0 or +0.0 \\n\", typeid(KeyT).name());\n fflush(stdout);\n InitializeKeyBits(RANDOM_MINUS_PLUS_ZERO, h_keys.get(), max_items, 0);\n // This just tests +/- 0 handling -- don't need to test multiple sizes\n TestNumItems(h_keys.get(), max_items, max_segments, false);\n }\n\n printf(\"\\nTesting uniform %s keys\\n\", typeid(KeyT).name()); fflush(stdout);\n InitializeKeyBits(UNIFORM, h_keys.get(), max_items, 0);\n TestSizes(h_keys.get(), max_items, max_segments, false);\n\n printf(\"\\nTesting natural number %s keys\\n\", typeid(KeyT).name()); fflush(stdout);\n InitializeKeyBits(INTEGER_SEED, h_keys.get(), max_items, 0);\n TestSizes(h_keys.get(), max_items, max_segments, false);\n\n if (WITH_PRE_SORTED)\n {\n // Presorting is only used for testing large input arrays.\n const std::size_t large_num_items = std::size_t(4350000007ull);\n\n // A conservative check for memory, as we don't know ValueT or whether\n // the overwrite is allowed until later.\n // For ValueT, the check is actually exact unless TEST_VALUE_TYPE == 3.\n if (!HasEnoughMemory(large_num_items, true))\n {\n printf(\"Skipping the permutation-based test due to insufficient device memory\\n\");\n return;\n }\n\n h_keys.reset(nullptr); // Explicitly free old buffer before allocating.\n h_keys.reset(new KeyT[large_num_items]);\n\n printf(\"\\nTesting pre-sorted and randomly permuted %s keys\\n\", typeid(KeyT).name());\n fflush(stdout);\n InitializeKeysSorted(h_keys.get(), large_num_items);\n fflush(stdout);\n TestSizes(h_keys.get(), large_num_items, max_segments, true);\n fflush(stdout);\n }\n}\n\n\n//---------------------------------------------------------------------\n// Simple test\n//---------------------------------------------------------------------\n\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT,\n bool IS_DESCENDING>\nvoid Test(\n std::size_t num_items,\n int num_segments,\n GenMode gen_mode,\n int entropy_reduction,\n int begin_bit,\n int end_bit)\n{\n const bool KEYS_ONLY = std::is_same::value;\n\n KeyT *h_keys = new KeyT[num_items];\n std::size_t *h_reference_ranks = NULL;\n KeyT *h_reference_keys = NULL;\n ValueT *h_values = NULL;\n ValueT *h_reference_values = NULL;\n size_t *h_segment_offsets = new std::size_t[num_segments + 1];\n\n std::size_t* d_segment_offsets = nullptr;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_segment_offsets, sizeof(std::size_t) * (num_segments + 1)));\n\n if (end_bit < 0)\n end_bit = sizeof(KeyT) * 8;\n\n InitializeKeyBits(gen_mode, h_keys, num_items, entropy_reduction);\n InitializeSegments(num_items, num_segments, h_segment_offsets);\n CubDebugExit(hipMemcpy(d_segment_offsets, h_segment_offsets, sizeof(std::size_t) * (num_segments + 1), hipMemcpyHostToDevice));\n InitializeSolution(\n h_keys, num_items, num_segments, false, h_segment_offsets,\n begin_bit, end_bit, h_reference_ranks, h_reference_keys);\n\n if (!KEYS_ONLY)\n {\n h_values = new ValueT[num_items];\n h_reference_values = new ValueT[num_items];\n\n for (std::size_t i = 0; i < num_items; ++i)\n {\n InitValue(INTEGER_SEED, h_values[i], i);\n InitValue(INTEGER_SEED, h_reference_values[i], h_reference_ranks[i]);\n }\n }\n if (h_reference_ranks) delete[] h_reference_ranks;\n\n printf(\"\\nTesting bits [%d,%d) of %s keys with gen-mode %d\\n\", begin_bit, end_bit, typeid(KeyT).name(), gen_mode); fflush(stdout);\n Test(\n h_keys, h_values,\n num_items, num_segments, d_segment_offsets, d_segment_offsets + 1,\n begin_bit, end_bit, h_reference_keys, h_reference_values);\n\n if (h_keys) delete[] h_keys;\n if (h_reference_keys) delete[] h_reference_keys;\n if (h_values) delete[] h_values;\n if (h_reference_values) delete[] h_reference_values;\n if (h_segment_offsets) delete[] h_segment_offsets;\n if (d_segment_offsets) CubDebugExit(g_allocator.DeviceFree(d_segment_offsets));\n}\n\n#if TEST_VALUE_TYPE == 0\nvoid TestUnspecifiedRanges()\n{\n const std::size_t num_items = 1024 * 1024;\n const std::size_t max_segments = 42;\n const std::size_t avg_segment_size = num_items / max_segments;\n\n for (int iteration = 0; iteration < 4; iteration++)\n {\n thrust::host_vector h_offsets_begin;\n thrust::host_vector h_offsets_end;\n\n h_offsets_begin.reserve(max_segments + 1);\n h_offsets_end.reserve(max_segments + 1);\n\n {\n int offset = 0;\n\n for (std::size_t sid = 0; sid < max_segments; sid++)\n {\n const int segment_size = \n static_cast(RandomValue(avg_segment_size));\n\n const bool segment_is_utilized = segment_size > 0 \n && RandomValue(100) > 60;\n\n if (segment_is_utilized)\n {\n h_offsets_begin.push_back(offset);\n h_offsets_end.push_back(offset + segment_size);\n }\n\n offset += segment_size;\n }\n\n if (h_offsets_begin.empty())\n {\n h_offsets_begin.push_back(avg_segment_size);\n h_offsets_end.push_back(num_items);\n }\n }\n\n thrust::device_vector keys(num_items);\n thrust::device_vector values(num_items);\n\n thrust::sequence(keys.rbegin(), keys.rend());\n thrust::sequence(values.rbegin(), values.rend());\n\n thrust::device_vector d_offsets_begin = h_offsets_begin;\n thrust::device_vector d_offsets_end = h_offsets_end;\n\n thrust::device_vector expected_keys = keys;\n thrust::device_vector expected_values = values;\n\n const int num_segments = static_cast(h_offsets_begin.size());\n\n thrust::device_vector result_keys = keys;\n thrust::device_vector result_values = values;\n\n for (int sid = 0; sid < num_segments; sid++)\n {\n const int segment_begin = h_offsets_begin[sid];\n const int segment_end = h_offsets_end[sid];\n\n thrust::sort_by_key(expected_keys.begin() + segment_begin,\n expected_keys.begin() + segment_end,\n expected_values.begin() + segment_begin);\n }\n\n {\n hipcub::DoubleBuffer keys_buffer(\n thrust::raw_pointer_cast(keys.data()), \n thrust::raw_pointer_cast(result_keys.data()));\n\n hipcub::DoubleBuffer values_buffer(\n thrust::raw_pointer_cast(values.data()), \n thrust::raw_pointer_cast(result_values.data()));\n\n std::size_t temp_storage_bytes{};\n std::uint8_t *d_temp_storage{nullptr};\n\n CubDebugExit(hipcub::DeviceSegmentedRadixSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n keys_buffer, values_buffer, \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data()),\n 0, sizeof(int) * 8));\n\n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n CubDebugExit(hipcub::DeviceSegmentedRadixSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n keys_buffer, values_buffer, \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data()),\n 0, sizeof(int) * 8));\n\n for (int sid = 0; sid < num_segments; sid++)\n {\n const int segment_begin = h_offsets_begin[sid];\n const int segment_end = h_offsets_end[sid];\n\n if (keys_buffer.selector == 0)\n {\n thrust::copy(\n keys.begin() + segment_begin,\n keys.begin() + segment_end,\n result_keys.begin() + segment_begin);\n }\n \n if (values_buffer.selector == 0)\n {\n thrust::copy(\n values.begin() + segment_begin,\n values.begin() + segment_end,\n result_values.begin() + segment_begin);\n }\n }\n }\n\n AssertEquals(result_keys, expected_keys); \n AssertEquals(result_values, expected_values);\n\n thrust::sequence(keys.rbegin(), keys.rend());\n thrust::sequence(values.rbegin(), values.rend());\n\n result_keys = keys;\n result_values = values;\n\n {\n std::size_t temp_storage_bytes{};\n std::uint8_t *d_temp_storage{};\n\n CubDebugExit(hipcub::DeviceSegmentedRadixSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n thrust::raw_pointer_cast(keys.data()), \n thrust::raw_pointer_cast(result_keys.data()), \n thrust::raw_pointer_cast(values.data()), \n thrust::raw_pointer_cast(result_values.data()), \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data()),\n 0, sizeof(int) * 8));\n\n thrust::device_vector temp_storage(temp_storage_bytes);\n d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());\n\n CubDebugExit(hipcub::DeviceSegmentedRadixSort::SortPairs(\n d_temp_storage, temp_storage_bytes, \n thrust::raw_pointer_cast(keys.data()), \n thrust::raw_pointer_cast(result_keys.data()), \n thrust::raw_pointer_cast(values.data()), \n thrust::raw_pointer_cast(result_values.data()), \n num_items, num_segments, \n thrust::raw_pointer_cast(d_offsets_begin.data()),\n thrust::raw_pointer_cast(d_offsets_end.data()),\n 0, sizeof(int) * 8));\n }\n\n AssertEquals(result_values, expected_values);\n AssertEquals(result_keys, expected_keys);\n }\n}\n#endif\n\n#if TEST_KEY_BYTES == 4\n// Following tests check that new decomposer API doesn't break old API. \n// It's disabled because some compilers don't like implicit conversions, which\n// is required for the test. Once we figure out how to temporarily enable conversion, we can\n// re-enable the test.\n#define ENABLING_CONVERSION_IS_FIGURED_OUT 0\n#if ENABLING_CONVERSION_IS_FIGURED_OUT \nstruct bit_selector\n{\n int bit;\n\n operator int() const\n {\n return bit;\n }\n};\n\ntemplate \nvoid device_radix_sort_keys_allows_implicit_conversions_for_bits_helper(BeginBitT begin_bit, EndBitT end_bit, Ts... args)\n{\n const int num_items = 0;\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n hipcub::DeviceRadixSort::SortKeys(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit,\n end_bit);\n }\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n hipcub::DeviceRadixSort::SortKeys(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit);\n }\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n hipcub::DeviceRadixSort::SortKeysDescending(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit,\n end_bit);\n }\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n hipcub::DeviceRadixSort::SortKeysDescending(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit);\n }\n}\n\ntemplate \nvoid device_radix_sort_pairs_allows_implicit_conversions_for_bits_helper(BeginBitT begin_bit,\n EndBitT end_bit,\n Ts... args)\n{\n const int num_items = 0;\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n hipcub::DeviceRadixSort::SortPairs(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit,\n end_bit);\n }\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n hipcub::DeviceRadixSort::SortPairs(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit);\n }\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n hipcub::DeviceRadixSort::SortPairsDescending(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit,\n end_bit);\n }\n\n {\n std::size_t temp_storage_bytes = 0;\n std::uint8_t *d_temp_storage = nullptr;\n\n hipcub::DeviceRadixSort::SortPairsDescending(d_temp_storage,\n temp_storage_bytes,\n args...,\n num_items,\n begin_bit);\n }\n}\n\ntemplate \nvoid device_radix_sort_allows_implicit_conversions_for_bits(BeginBitT begin_bit, EndBitT end_bit)\n{\n int *d_i_ptr = nullptr;\n\n device_radix_sort_keys_allows_implicit_conversions_for_bits_helper(begin_bit,\n end_bit,\n d_i_ptr,\n d_i_ptr);\n\n hipcub::DoubleBuffer keys(d_i_ptr, d_i_ptr);\n device_radix_sort_keys_allows_implicit_conversions_for_bits_helper(begin_bit, end_bit, keys);\n\n #if TEST_VALUE_TYPE == 2\n unsigned int *d_u_ptr = nullptr;\n\n device_radix_sort_pairs_allows_implicit_conversions_for_bits_helper(begin_bit,\n end_bit,\n d_i_ptr,\n d_i_ptr,\n d_u_ptr,\n d_u_ptr);\n\n hipcub::DoubleBuffer pairs(d_u_ptr, d_u_ptr);\n device_radix_sort_pairs_allows_implicit_conversions_for_bits_helper(begin_bit,\n end_bit,\n keys,\n pairs);\n #endif\n}\n\nvoid device_radix_sort_allows_implicit_conversions_for_bits()\n{\n int begin_i = 0;\n long long int begin_lli = 0;\n bit_selector begin_bs{0};\n\n int end_i = 2;\n long long int end_lli = 2;\n bit_selector end_bs{2};\n\n device_radix_sort_allows_implicit_conversions_for_bits(begin_i, end_i);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_i, end_lli);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_i, end_bs);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_lli, end_i);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_lli, end_lli);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_lli, end_bs);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_bs, end_i);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_bs, end_lli);\n device_radix_sort_allows_implicit_conversions_for_bits(begin_bs, end_bs);\n}\n#endif // ENABLING_CONVERSION_IS_FIGURED_OUT \n#endif // TEST_KEY_BYTES == 4\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n std::size_t num_items = ~std::size_t(0);\n int num_segments = -1;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"s\", num_segments);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--s= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // %PARAM% TEST_CDP cdp 0:1\n // %PARAM% TEST_KEY_BYTES bytes 1:2:4:8:16\n // %PARAM% TEST_VALUE_TYPE pairs 0:1:2:3\n // 0->Keys only\n // 1->uchar\n // 2->uint\n // 3->[ull,TestBar] (TestBar is cheap to build, included here to\n // reduce total number of targets)\n\n // To reduce testing time, some key types are only tested when not\n // testing pairs:\n#if TEST_VALUE_TYPE == 0\n#define TEST_EXTENDED_KEY_TYPES\n#endif\n\n // Compile/run thorough tests\n#if TEST_KEY_BYTES == 1\n\n TestGen (num_items, num_segments);\n\n#ifdef TEST_EXTENDED_KEY_TYPES\n TestGen (num_items, num_segments);\n TestGen (num_items, num_segments);\n TestGen (num_items, num_segments);\n#endif // TEST_EXTENDED_KEY_TYPES\n\n#elif TEST_KEY_BYTES == 2\n TestGen (num_items, num_segments);\n\n#ifdef TEST_EXTENDED_KEY_TYPES\n TestGen (num_items, num_segments);\n\n#if !_NVHPC_CUDA\n TestGen (num_items, num_segments);\n#endif // CTK >= 9\n\n#if !_NVHPC_CUDA\n#if !defined(__ICC)\n // Fails with `-0 != 0` with ICC for unknown reasons. See #333.\n TestGen (num_items, num_segments);\n#endif // !ICC\n#endif // CTK >= 11\n\n#endif // TEST_EXTENDED_KEY_TYPES\n\n#elif TEST_KEY_BYTES == 4\n\n TestGen (num_items, num_segments);\n\n#if TEST_VALUE_TYPE == 0\n TestUnspecifiedRanges();\n#endif\n\n#if ENABLING_CONVERSION_IS_FIGURED_OUT \n device_radix_sort_allows_implicit_conversions_for_bits();\n#endif\n\n#ifdef TEST_EXTENDED_KEY_TYPES\n TestGen (num_items, num_segments);\n TestGen (num_items, num_segments);\n#endif // TEST_EXTENDED_KEY_TYPES\n\n#elif TEST_KEY_BYTES == 8\n\n TestGen (num_items, num_segments);\n\n#ifdef TEST_EXTENDED_KEY_TYPES\n TestGen (num_items, num_segments);\n TestGen(num_items, num_segments);\n#endif // TEST_EXTENDED_KEY_TYPES\n\n#elif TEST_KEY_BYTES == 16\n\n#if CUB_IS_INT128_ENABLED \n TestGen<__int128_t, false>(num_items, num_segments);\n TestGen<__uint128_t, false>(num_items, num_segments);\n#else\n // Fix unused static function for MSVC\n BackendToString(CUB);\n#endif\n\n#endif // TEST_KEY_BYTES switch\n\n return 0;\n}\n", "source": "/kaggle/working/cub/test/test_device_radix_sort.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\n/******************************************************************************\n * HELPER CLASS FOR RUN-LENGTH DECODING TESTS\n ******************************************************************************/\n\n/**\n * \\brief Class template to facilitate testing the BlockRunLengthDecode algorithm for all its\n * template parameter specialisations.\n *\n * \\tparam ItemItT The item type being run-length decoded\n * \\tparam RunLengthsItT Iterator type providing the runs' lengths\n * \\tparam RUNS_PER_THREAD The number of runs that each thread is getting assigned to\n * \\tparam DECODED_ITEMS_PER_THREAD The number of run-length decoded items that each thread is \n * decoding \\tparam TEST_RELATIVE_OFFSETS_ Whether to also retrieve each decoded item's \n * relative offset within its run \\tparam TEST_RUN_OFFSETS_ Whether to pass in each run's \n * offset instead of each run's length \\tparam BLOCK_DIM_X The thread block length in \n * threads along the X dimension\n * \\tparam BLOCK_DIM_Y The thread block length in threads along the Y dimension\n * \\tparam BLOCK_DIM_Z The thread block length in threads along the Z dimension\n */\ntemplate \nclass AgentTestBlockRunLengthDecode\n{\npublic:\n constexpr static uint32_t BLOCK_THREADS = BLOCK_DIM_X * BLOCK_DIM_Y * BLOCK_DIM_Z;\n constexpr static uint32_t RUNS_PER_BLOCK = RUNS_PER_THREAD * BLOCK_THREADS;\n constexpr static bool TEST_RELATIVE_OFFSETS = TEST_RELATIVE_OFFSETS_;\n\nprivate:\n using RunItemT = cub::detail::value_t;\n using RunLengthT = cub::detail::value_t;\n\n using BlockRunOffsetScanT =\n cub::BlockScan;\n\n using BlockRunLengthDecodeT =\n cub::BlockRunLengthDecode;\n\n using BlockLoadRunItemT = cub::BlockLoad;\n\n using BlockLoadRunLengthsT = cub::BlockLoad;\n\n using BlockStoreDecodedItemT = cub::BlockStore;\n\n using BlockStoreRelativeOffsetT = cub::BlockStore;\n\n __device__ __forceinline__ BlockRunLengthDecodeT\n InitBlockRunLengthDecode(RunItemT (&unique_items)[RUNS_PER_THREAD],\n RunLengthT (&run_lengths)[RUNS_PER_THREAD],\n RunLengthT &decoded_size,\n cub::Int2Type /*test_run_offsets*/)\n {\n RunLengthT run_offsets[RUNS_PER_THREAD];\n BlockRunOffsetScanT(temp_storage.run_offsets_scan_storage)\n .ExclusiveSum(run_lengths, run_offsets, decoded_size);\n\n // Ensure temporary shared memory can be repurposed\n cub::CTA_SYNC();\n\n // Construct BlockRunLengthDecode and initialize with the run offsets\n return BlockRunLengthDecodeT(temp_storage.decode.run_length_decode_storage,\n unique_items,\n run_offsets);\n }\n\n __device__ __forceinline__ BlockRunLengthDecodeT\n InitBlockRunLengthDecode(RunItemT (&unique_items)[RUNS_PER_THREAD],\n RunLengthT (&run_lengths)[RUNS_PER_THREAD],\n RunLengthT &decoded_size,\n cub::Int2Type /*test_run_offsets*/)\n {\n // Construct BlockRunLengthDecode and initialize with the run lengths\n return BlockRunLengthDecodeT(temp_storage.decode.run_length_decode_storage,\n unique_items,\n run_lengths,\n decoded_size);\n }\n\n __device__ __forceinline__ void LoadRuns(ItemItT d_block_unique_items,\n RunLengthsItT d_block_run_lengths,\n RunItemT (&unique_items)[RUNS_PER_THREAD],\n RunLengthT (&run_lengths)[RUNS_PER_THREAD],\n size_t num_valid_items)\n {\n if (num_valid_items < RUNS_PER_BLOCK)\n {\n BlockLoadRunItemT(temp_storage.load_uniques_storage)\n .Load(d_block_unique_items, unique_items, num_valid_items);\n }\n else\n {\n BlockLoadRunItemT(temp_storage.load_uniques_storage).Load(d_block_unique_items, unique_items);\n }\n\n // Ensure BlockLoad's temporary shared memory can be repurposed\n cub::CTA_SYNC();\n\n // Load this block's tile of run lengths\n if (num_valid_items < RUNS_PER_BLOCK)\n BlockLoadRunLengthsT(temp_storage.load_run_lengths_storage)\n .Load(d_block_run_lengths, run_lengths, num_valid_items, static_cast(0));\n else\n BlockLoadRunLengthsT(temp_storage.load_run_lengths_storage)\n .Load(d_block_run_lengths, run_lengths);\n\n // Ensure temporary shared memory can be repurposed\n cub::CTA_SYNC();\n }\n\npublic:\n union TempStorage\n {\n typename BlockLoadRunItemT::TempStorage load_uniques_storage;\n typename BlockLoadRunLengthsT::TempStorage load_run_lengths_storage;\n cub::detail::\n conditional_t\n run_offsets_scan_storage;\n struct\n {\n typename BlockRunLengthDecodeT::TempStorage run_length_decode_storage;\n typename BlockStoreDecodedItemT::TempStorage store_decoded_runs_storage;\n typename BlockStoreRelativeOffsetT::TempStorage store_relative_offsets;\n } decode;\n };\n\n TempStorage &temp_storage;\n\n __device__ __forceinline__ AgentTestBlockRunLengthDecode(TempStorage &temp_storage)\n : temp_storage(temp_storage)\n {}\n\n /**\n * \\brief Loads the given block (or tile) of runs, and computes their \"decompressed\" (run-length\n * decoded) size.\n */\n __device__ __forceinline__ uint32_t GetDecodedSize(ItemItT d_block_unique_items,\n RunLengthsItT d_block_run_lengths,\n size_t num_valid_runs)\n {\n // Load this block's tile of encoded runs\n RunItemT unique_items[RUNS_PER_THREAD];\n RunLengthT run_lengths[RUNS_PER_THREAD];\n LoadRuns(d_block_unique_items, d_block_run_lengths, unique_items, run_lengths, num_valid_runs);\n\n // Init the BlockRunLengthDecode and get the total decoded size of this block's tile (i.e., the\n // \"decompressed\" size)\n uint32_t decoded_size = 0U;\n BlockRunLengthDecodeT run_length_decode =\n InitBlockRunLengthDecode(unique_items,\n run_lengths,\n decoded_size,\n cub::Int2Type());\n return decoded_size;\n }\n\n /**\n * \\brief Loads the given block (or tile) of runs, run-length decodes them, and writes the results\n * to \\p d_block_decoded_out.\n */\n template \n __device__ __forceinline__ uint32_t WriteDecodedRuns(ItemItT d_block_unique_items,\n RunLengthsItT d_block_run_lengths,\n UniqueItemOutItT d_block_decoded_out,\n RelativeOffsetOutItT d_block_rel_out,\n size_t num_valid_runs)\n {\n // Load this block's tile of encoded runs\n RunItemT unique_items[RUNS_PER_THREAD];\n RunLengthT run_lengths[RUNS_PER_THREAD];\n LoadRuns(d_block_unique_items, d_block_run_lengths, unique_items, run_lengths, num_valid_runs);\n\n // Init the BlockRunLengthDecode and get the total decoded size of this block's tile (i.e., the\n // \"decompressed\" size)\n uint32_t decoded_size = 0U;\n BlockRunLengthDecodeT run_length_decode =\n InitBlockRunLengthDecode(unique_items,\n run_lengths,\n decoded_size,\n cub::Int2Type());\n\n // Run-length decode (\"decompress\") the runs into a window buffer of limited size. This is\n // repeated until all runs have been decoded.\n uint32_t decoded_window_offset = 0U;\n while (decoded_window_offset < decoded_size)\n {\n RunLengthT relative_offsets[DECODED_ITEMS_PER_THREAD];\n RunItemT decoded_items[DECODED_ITEMS_PER_THREAD];\n\n // The number of decoded items that are valid within this window (aka pass) of run-length\n // decoding\n uint32_t num_valid_items = decoded_size - decoded_window_offset;\n run_length_decode.RunLengthDecode(decoded_items, relative_offsets, decoded_window_offset);\n BlockStoreDecodedItemT(temp_storage.decode.store_decoded_runs_storage)\n .Store(d_block_decoded_out + decoded_window_offset, decoded_items, num_valid_items);\n\n if (TEST_RELATIVE_OFFSETS)\n {\n BlockStoreRelativeOffsetT(temp_storage.decode.store_relative_offsets)\n .Store(d_block_rel_out + decoded_window_offset, relative_offsets, num_valid_items);\n }\n\n decoded_window_offset += DECODED_ITEMS_PER_THREAD * BLOCK_THREADS;\n }\n return decoded_size;\n }\n};\n\n/******************************************************************************\n * [STAGE 1] RUN-LENGTH DECODING TEST KERNEL\n ******************************************************************************/\ntemplate \n__launch_bounds__(AgentTestBlockRunLengthDecode::BLOCK_THREADS) __global__\n void BlockRunLengthDecodeGetSizeKernel(const ItemItT d_unique_items,\n const RunLengthsItT d_run_lengths,\n const OffsetT num_runs,\n DecodedSizesOutT d_decoded_sizes)\n{\n constexpr OffsetT RUNS_PER_BLOCK = AgentTestBlockRunLengthDecode::RUNS_PER_BLOCK;\n\n __shared__ typename AgentTestBlockRunLengthDecode::TempStorage temp_storage;\n\n OffsetT block_offset = blockIdx.x * RUNS_PER_BLOCK;\n OffsetT num_valid_runs = (block_offset + RUNS_PER_BLOCK >= num_runs) ? (num_runs - block_offset)\n : RUNS_PER_BLOCK;\n\n AgentTestBlockRunLengthDecode run_length_decode_agent(temp_storage);\n uint64_t num_decoded_items = run_length_decode_agent.GetDecodedSize(d_unique_items + block_offset,\n d_run_lengths + block_offset,\n num_valid_runs);\n\n d_decoded_sizes[blockIdx.x] = num_decoded_items;\n}\n\n/******************************************************************************\n * [STAGE 2] RUN-LENGTH DECODING TEST KERNEL\n ******************************************************************************/\ntemplate \n__launch_bounds__(AgentTestBlockRunLengthDecode::BLOCK_THREADS) __global__\n void BlockRunLengthDecodeTestKernel(const ItemItT d_unique_items,\n const RunLengthsItT d_run_lengths,\n const DecodedSizesOutT d_decoded_offsets,\n const OffsetT num_runs,\n DecodedItemsOutItT d_decoded_items,\n RelativeOffsetOutItT d_relative_offsets)\n\n{\n constexpr OffsetT RUNS_PER_BLOCK = AgentTestBlockRunLengthDecode::RUNS_PER_BLOCK;\n\n __shared__ typename AgentTestBlockRunLengthDecode::TempStorage temp_storage;\n\n OffsetT block_offset = blockIdx.x * RUNS_PER_BLOCK;\n OffsetT num_valid_runs = (block_offset + RUNS_PER_BLOCK >= num_runs) ? (num_runs - block_offset)\n : RUNS_PER_BLOCK;\n\n AgentTestBlockRunLengthDecode run_length_decode_agent(temp_storage);\n run_length_decode_agent.WriteDecodedRuns(d_unique_items + block_offset,\n d_run_lengths + block_offset,\n d_decoded_items + d_decoded_offsets[blockIdx.x],\n d_relative_offsets + d_decoded_offsets[blockIdx.x],\n num_valid_runs);\n}\n\nstruct ModOp\n{\n using T = uint32_t;\n __host__ __device__ __forceinline__ T operator()(const T &x) const { return 1 + (x % 100); }\n};\n\ntemplate \nvoid TestAlgorithmSpecialisation()\n{\n constexpr uint32_t THREADS_PER_BLOCK = BLOCK_DIM_X * BLOCK_DIM_Y * BLOCK_DIM_Z;\n constexpr uint32_t RUNS_PER_BLOCK = RUNS_PER_THREAD * THREADS_PER_BLOCK;\n\n using RunItemT = float;\n using RunLengthT = uint32_t;\n using ItemItT = cub::CountingInputIterator;\n using RunLengthsItT =\n cub::TransformInputIterator>;\n\n ItemItT d_unique_items(1000U);\n RunLengthsItT d_run_lengths(cub::CountingInputIterator(0), ModOp{});\n\n constexpr uint32_t num_runs = 10000;\n constexpr uint32_t num_blocks = (num_runs + (RUNS_PER_BLOCK - 1U)) / RUNS_PER_BLOCK;\n\n size_t temp_storage_bytes = 0ULL;\n void *temp_storage = nullptr;\n uint32_t *h_num_decoded_total = nullptr;\n uint32_t *d_decoded_sizes = nullptr;\n uint32_t *d_decoded_offsets = nullptr;\n RunItemT *d_decoded_out = nullptr;\n RunLengthT *d_relative_offsets = nullptr;\n RunItemT *h_decoded_out = nullptr;\n RunLengthT *h_relative_offsets = nullptr;\n\n using AgentTestBlockRunLengthDecodeT = AgentTestBlockRunLengthDecode;\n\n enum : uint32_t\n {\n TIMER_SIZE_BEGIN = 0,\n TIMER_SIZE_END,\n TIMER_DECODE_BEGIN,\n TIMER_DECODE_END,\n NUM_TIMERS,\n };\n\n // Get temporary storage requirements for the scan (for computing offsets for the per-block\n // run-length decoded items)\n cub::DeviceScan::InclusiveSum(nullptr,\n temp_storage_bytes,\n d_decoded_sizes,\n d_decoded_offsets,\n num_blocks);\n\n // Allocate device memory\n CubDebugExit(cudaMalloc(&temp_storage, temp_storage_bytes));\n CubDebugExit(cudaMalloc(&d_decoded_sizes, num_blocks * sizeof(*d_decoded_sizes)));\n // Allocate for the exclusive sum PLUS the overall aggregate\n CubDebugExit(cudaMalloc(&d_decoded_offsets, (num_blocks + 1) * sizeof(*d_decoded_offsets)));\n CubDebugExit(cudaMallocHost(&h_num_decoded_total, sizeof(*h_num_decoded_total)));\n\n // Get the per-block number of items being decoded (i-th thread block writing size to\n // d_decoded_sizes[i])\n BlockRunLengthDecodeGetSizeKernel\n <<>>(d_unique_items,\n d_run_lengths,\n num_runs,\n d_decoded_sizes);\n\n // Compute offsets for the runs decoded by each block (exclusive sum + aggregate)\n CubDebugExit(cudaMemsetAsync(d_decoded_offsets, 0, sizeof(d_decoded_offsets[0])));\n CubDebugExit(cub::DeviceScan::InclusiveSum(temp_storage,\n temp_storage_bytes,\n d_decoded_sizes,\n &d_decoded_offsets[1],\n num_blocks));\n\n // Copy the total decoded size to CPU in order to allocate just the right amount of device memory\n CubDebugExit(cudaMemcpy(h_num_decoded_total,\n &d_decoded_offsets[num_blocks],\n sizeof(*h_num_decoded_total),\n cudaMemcpyDeviceToHost));\n\n // Allocate device memory for the run-length decoded output\n CubDebugExit(cudaMallocHost(&h_decoded_out, (*h_num_decoded_total) * sizeof(RunItemT)));\n CubDebugExit(cudaMalloc(&d_decoded_out, (*h_num_decoded_total) * sizeof(RunItemT)));\n if (TEST_RELATIVE_OFFSETS)\n {\n CubDebugExit(cudaMalloc(&d_relative_offsets, (*h_num_decoded_total) * sizeof(RunLengthT)));\n CubDebugExit(cudaMallocHost(&h_relative_offsets, (*h_num_decoded_total) * sizeof(RunLengthT)));\n }\n\n // Perform the block-wise run-length decoding (each block taking its offset from\n // d_decoded_offsets)\n BlockRunLengthDecodeTestKernel\n <<>>(d_unique_items,\n d_run_lengths,\n d_decoded_offsets,\n num_runs,\n d_decoded_out,\n d_relative_offsets);\n\n // Copy back results for verification\n CubDebugExit(cudaMemcpy(h_decoded_out,\n d_decoded_out,\n (*h_num_decoded_total) * sizeof(*h_decoded_out),\n cudaMemcpyDeviceToHost));\n\n if (TEST_RELATIVE_OFFSETS)\n {\n // Copy back the relative offsets\n CubDebugExit(cudaMemcpy(h_relative_offsets,\n d_relative_offsets,\n (*h_num_decoded_total) * sizeof(*h_relative_offsets),\n cudaMemcpyDeviceToHost));\n }\n\n // Generate host-side run-length decoded data for verification\n std::vector> host_golden;\n host_golden.reserve(*h_num_decoded_total);\n for (uint32_t run = 0; run < num_runs; run++)\n {\n for (RunLengthT i = 0; i < d_run_lengths[run]; i++)\n {\n host_golden.push_back({d_unique_items[run], i});\n }\n }\n\n // Verify the total run-length decoded size is correct\n REQUIRE(host_golden.size() == h_num_decoded_total[0]);\n\n // Verify the run-length decoded data is correct\n bool cmp_eq = true;\n for (uint32_t i = 0; i < host_golden.size(); i++)\n {\n if (host_golden[i].first != h_decoded_out[i])\n {\n FAIL(\"Mismatch at #\" << i << \": CPU item: \" << host_golden[i].first\n << \", GPU: \" << h_decoded_out[i] << \"\\n\");\n cmp_eq = false;\n }\n if (TEST_RELATIVE_OFFSETS)\n {\n if (host_golden[i].second != h_relative_offsets[i])\n {\n FAIL(\"Mismatch of relative offset at #\" << i\n << \": CPU item: \" << host_golden[i].first << \", GPU: \" << h_decoded_out[i]\n << \"; relative offsets: CPU: \" << host_golden[i].second\n << \", GPU: \" << h_relative_offsets[i] << \"\\n\");\n cmp_eq = false;\n break;\n }\n }\n }\n\n REQUIRE(cmp_eq == true);\n\n // Clean up memory allocations\n CubDebugExit(cudaFree(temp_storage));\n CubDebugExit(cudaFree(d_decoded_sizes));\n CubDebugExit(cudaFree(d_decoded_offsets));\n CubDebugExit(cudaFree(d_decoded_out));\n CubDebugExit(cudaFreeHost(h_num_decoded_total));\n CubDebugExit(cudaFreeHost(h_decoded_out));\n\n if (TEST_RELATIVE_OFFSETS)\n {\n CubDebugExit(cudaFree(d_relative_offsets));\n CubDebugExit(cudaFreeHost(h_relative_offsets));\n }\n}\n\nconstexpr bool DO_TEST_RELATIVE_OFFSETS = true;\nconstexpr bool DO_NOT_TEST_RELATIVE_OFFSETS = false;\n\nconstexpr bool TEST_WITH_RUN_OFFSETS = true;\nconstexpr bool TEST_WITH_RUN_LENGTHS = false;\n\ntemplate \nstruct params_t\n{\n static constexpr int runs_per_thread = RunsPerThread;\n static constexpr int decoded_items_per_thread = DecodedItemsPerThread;\n static constexpr int block_dim_x = BlockDimX;\n static constexpr int block_dim_y = BlockDimY;\n static constexpr int block_dim_z = BlockDimZ;\n};\n\nCUB_TEST_LIST(\"Block Run Length Decode works with run lengths and offsets relative to each run\",\n \"[rld][block]\",\n params_t<1, 1, 64>,\n params_t<1, 3, 32, 2, 3>,\n params_t<1, 1, 128>,\n params_t<1, 8, 128>,\n params_t<3, 1, 256>,\n params_t<1, 8, 256>,\n params_t<8, 1, 256>,\n params_t<1, 1, 256>,\n params_t<2, 2, 384>)\n{\n using params = TestType;\n\n TestAlgorithmSpecialisation();\n}\n\nCUB_TEST_LIST(\"Block Run Length Decode works with run lengths and performs normal run-length \"\n \"decoding\",\n \"[rld][block]\",\n params_t<1, 1, 64>,\n params_t<1, 3, 32, 2, 3>,\n params_t<1, 1, 128>,\n params_t<1, 8, 128>,\n params_t<3, 1, 256>,\n params_t<1, 8, 256>,\n params_t<8, 1, 256>,\n params_t<1, 1, 256>,\n params_t<2, 2, 384>)\n{\n using params = TestType;\n\n TestAlgorithmSpecialisation();\n}\n\nCUB_TEST_LIST(\"Block Run Length Decode works with run offsets and generates offsets relative to \"\n \"each run\",\n \"[rld][block]\",\n params_t<1, 1, 64>,\n params_t<1, 3, 32, 2, 3>,\n params_t<1, 1, 128>,\n params_t<1, 8, 128>,\n params_t<3, 1, 256>,\n params_t<1, 8, 256>,\n params_t<8, 1, 256>,\n params_t<1, 1, 256>,\n params_t<2, 2, 384>)\n{\n using params = TestType;\n\n TestAlgorithmSpecialisation();\n}\n\nCUB_TEST_LIST(\"Block Run Length Decode works with run offsets and performs normal run-length \"\n \"decoding\",\n \"[rld][block]\",\n params_t<1, 1, 64>,\n params_t<1, 3, 32, 2, 3>,\n params_t<1, 1, 128>,\n params_t<1, 8, 128>,\n params_t<3, 1, 256>,\n params_t<1, 8, 256>,\n params_t<8, 1, 256>,\n params_t<1, 1, 256>,\n params_t<2, 2, 384>)\n{\n using params = TestType;\n\n TestAlgorithmSpecialisation();\n}\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\n/******************************************************************************\n * HELPER CLASS FOR RUN-LENGTH DECODING TESTS\n ******************************************************************************/\n\n/**\n * \\brief Class template to facilitate testing the BlockRunLengthDecode algorithm for all its\n * template parameter specialisations.\n *\n * \\tparam ItemItT The item type being run-length decoded\n * \\tparam RunLengthsItT Iterator type providing the runs' lengths\n * \\tparam RUNS_PER_THREAD The number of runs that each thread is getting assigned to\n * \\tparam DECODED_ITEMS_PER_THREAD The number of run-length decoded items that each thread is \n * decoding \\tparam TEST_RELATIVE_OFFSETS_ Whether to also retrieve each decoded item's \n * relative offset within its run \\tparam TEST_RUN_OFFSETS_ Whether to pass in each run's \n * offset instead of each run's length \\tparam BLOCK_DIM_X The thread block length in \n * threads along the X dimension\n * \\tparam BLOCK_DIM_Y The thread block length in threads along the Y dimension\n * \\tparam BLOCK_DIM_Z The thread block length in threads along the Z dimension\n */\ntemplate \nclass AgentTestBlockRunLengthDecode\n{\npublic:\n constexpr static uint32_t BLOCK_THREADS = BLOCK_DIM_X * BLOCK_DIM_Y * BLOCK_DIM_Z;\n constexpr static uint32_t RUNS_PER_BLOCK = RUNS_PER_THREAD * BLOCK_THREADS;\n constexpr static bool TEST_RELATIVE_OFFSETS = TEST_RELATIVE_OFFSETS_;\n\nprivate:\n using RunItemT = hipcub::detail::value_t;\n using RunLengthT = hipcub::detail::value_t;\n\n using BlockRunOffsetScanT =\n hipcub::BlockScan;\n\n using BlockRunLengthDecodeT =\n hipcub::BlockRunLengthDecode;\n\n using BlockLoadRunItemT = hipcub::BlockLoad;\n\n using BlockLoadRunLengthsT = hipcub::BlockLoad;\n\n using BlockStoreDecodedItemT = hipcub::BlockStore;\n\n using BlockStoreRelativeOffsetT = hipcub::BlockStore;\n\n __device__ __forceinline__ BlockRunLengthDecodeT\n InitBlockRunLengthDecode(RunItemT (&unique_items)[RUNS_PER_THREAD],\n RunLengthT (&run_lengths)[RUNS_PER_THREAD],\n RunLengthT &decoded_size,\n hipcub::Int2Type /*test_run_offsets*/)\n {\n RunLengthT run_offsets[RUNS_PER_THREAD];\n BlockRunOffsetScanT(temp_storage.run_offsets_scan_storage)\n .ExclusiveSum(run_lengths, run_offsets, decoded_size);\n\n // Ensure temporary shared memory can be repurposed\n hipcub::CTA_SYNC();\n\n // Construct BlockRunLengthDecode and initialize with the run offsets\n return BlockRunLengthDecodeT(temp_storage.decode.run_length_decode_storage,\n unique_items,\n run_offsets);\n }\n\n __device__ __forceinline__ BlockRunLengthDecodeT\n InitBlockRunLengthDecode(RunItemT (&unique_items)[RUNS_PER_THREAD],\n RunLengthT (&run_lengths)[RUNS_PER_THREAD],\n RunLengthT &decoded_size,\n hipcub::Int2Type /*test_run_offsets*/)\n {\n // Construct BlockRunLengthDecode and initialize with the run lengths\n return BlockRunLengthDecodeT(temp_storage.decode.run_length_decode_storage,\n unique_items,\n run_lengths,\n decoded_size);\n }\n\n __device__ __forceinline__ void LoadRuns(ItemItT d_block_unique_items,\n RunLengthsItT d_block_run_lengths,\n RunItemT (&unique_items)[RUNS_PER_THREAD],\n RunLengthT (&run_lengths)[RUNS_PER_THREAD],\n size_t num_valid_items)\n {\n if (num_valid_items < RUNS_PER_BLOCK)\n {\n BlockLoadRunItemT(temp_storage.load_uniques_storage)\n .Load(d_block_unique_items, unique_items, num_valid_items);\n }\n else\n {\n BlockLoadRunItemT(temp_storage.load_uniques_storage).Load(d_block_unique_items, unique_items);\n }\n\n // Ensure BlockLoad's temporary shared memory can be repurposed\n hipcub::CTA_SYNC();\n\n // Load this block's tile of run lengths\n if (num_valid_items < RUNS_PER_BLOCK)\n BlockLoadRunLengthsT(temp_storage.load_run_lengths_storage)\n .Load(d_block_run_lengths, run_lengths, num_valid_items, static_cast(0));\n else\n BlockLoadRunLengthsT(temp_storage.load_run_lengths_storage)\n .Load(d_block_run_lengths, run_lengths);\n\n // Ensure temporary shared memory can be repurposed\n hipcub::CTA_SYNC();\n }\n\npublic:\n union TempStorage\n {\n typename BlockLoadRunItemT::TempStorage load_uniques_storage;\n typename BlockLoadRunLengthsT::TempStorage load_run_lengths_storage;\n hipcub::detail::\n conditional_t\n run_offsets_scan_storage;\n struct\n {\n typename BlockRunLengthDecodeT::TempStorage run_length_decode_storage;\n typename BlockStoreDecodedItemT::TempStorage store_decoded_runs_storage;\n typename BlockStoreRelativeOffsetT::TempStorage store_relative_offsets;\n } decode;\n };\n\n TempStorage &temp_storage;\n\n __device__ __forceinline__ AgentTestBlockRunLengthDecode(TempStorage &temp_storage)\n : temp_storage(temp_storage)\n {}\n\n /**\n * \\brief Loads the given block (or tile) of runs, and computes their \"decompressed\" (run-length\n * decoded) size.\n */\n __device__ __forceinline__ uint32_t GetDecodedSize(ItemItT d_block_unique_items,\n RunLengthsItT d_block_run_lengths,\n size_t num_valid_runs)\n {\n // Load this block's tile of encoded runs\n RunItemT unique_items[RUNS_PER_THREAD];\n RunLengthT run_lengths[RUNS_PER_THREAD];\n LoadRuns(d_block_unique_items, d_block_run_lengths, unique_items, run_lengths, num_valid_runs);\n\n // Init the BlockRunLengthDecode and get the total decoded size of this block's tile (i.e., the\n // \"decompressed\" size)\n uint32_t decoded_size = 0U;\n BlockRunLengthDecodeT run_length_decode =\n InitBlockRunLengthDecode(unique_items,\n run_lengths,\n decoded_size,\n hipcub::Int2Type());\n return decoded_size;\n }\n\n /**\n * \\brief Loads the given block (or tile) of runs, run-length decodes them, and writes the results\n * to \\p d_block_decoded_out.\n */\n template \n __device__ __forceinline__ uint32_t WriteDecodedRuns(ItemItT d_block_unique_items,\n RunLengthsItT d_block_run_lengths,\n UniqueItemOutItT d_block_decoded_out,\n RelativeOffsetOutItT d_block_rel_out,\n size_t num_valid_runs)\n {\n // Load this block's tile of encoded runs\n RunItemT unique_items[RUNS_PER_THREAD];\n RunLengthT run_lengths[RUNS_PER_THREAD];\n LoadRuns(d_block_unique_items, d_block_run_lengths, unique_items, run_lengths, num_valid_runs);\n\n // Init the BlockRunLengthDecode and get the total decoded size of this block's tile (i.e., the\n // \"decompressed\" size)\n uint32_t decoded_size = 0U;\n BlockRunLengthDecodeT run_length_decode =\n InitBlockRunLengthDecode(unique_items,\n run_lengths,\n decoded_size,\n hipcub::Int2Type());\n\n // Run-length decode (\"decompress\") the runs into a window buffer of limited size. This is\n // repeated until all runs have been decoded.\n uint32_t decoded_window_offset = 0U;\n while (decoded_window_offset < decoded_size)\n {\n RunLengthT relative_offsets[DECODED_ITEMS_PER_THREAD];\n RunItemT decoded_items[DECODED_ITEMS_PER_THREAD];\n\n // The number of decoded items that are valid within this window (aka pass) of run-length\n // decoding\n uint32_t num_valid_items = decoded_size - decoded_window_offset;\n run_length_decode.RunLengthDecode(decoded_items, relative_offsets, decoded_window_offset);\n BlockStoreDecodedItemT(temp_storage.decode.store_decoded_runs_storage)\n .Store(d_block_decoded_out + decoded_window_offset, decoded_items, num_valid_items);\n\n if (TEST_RELATIVE_OFFSETS)\n {\n BlockStoreRelativeOffsetT(temp_storage.decode.store_relative_offsets)\n .Store(d_block_rel_out + decoded_window_offset, relative_offsets, num_valid_items);\n }\n\n decoded_window_offset += DECODED_ITEMS_PER_THREAD * BLOCK_THREADS;\n }\n return decoded_size;\n }\n};\n\n/******************************************************************************\n * [STAGE 1] RUN-LENGTH DECODING TEST KERNEL\n ******************************************************************************/\ntemplate \n__launch_bounds__(AgentTestBlockRunLengthDecode::BLOCK_THREADS) __global__\n void BlockRunLengthDecodeGetSizeKernel(const ItemItT d_unique_items,\n const RunLengthsItT d_run_lengths,\n const OffsetT num_runs,\n DecodedSizesOutT d_decoded_sizes)\n{\n constexpr OffsetT RUNS_PER_BLOCK = AgentTestBlockRunLengthDecode::RUNS_PER_BLOCK;\n\n __shared__ typename AgentTestBlockRunLengthDecode::TempStorage temp_storage;\n\n OffsetT block_offset = blockIdx.x * RUNS_PER_BLOCK;\n OffsetT num_valid_runs = (block_offset + RUNS_PER_BLOCK >= num_runs) ? (num_runs - block_offset)\n : RUNS_PER_BLOCK;\n\n AgentTestBlockRunLengthDecode run_length_decode_agent(temp_storage);\n uint64_t num_decoded_items = run_length_decode_agent.GetDecodedSize(d_unique_items + block_offset,\n d_run_lengths + block_offset,\n num_valid_runs);\n\n d_decoded_sizes[blockIdx.x] = num_decoded_items;\n}\n\n/******************************************************************************\n * [STAGE 2] RUN-LENGTH DECODING TEST KERNEL\n ******************************************************************************/\ntemplate \n__launch_bounds__(AgentTestBlockRunLengthDecode::BLOCK_THREADS) __global__\n void BlockRunLengthDecodeTestKernel(const ItemItT d_unique_items,\n const RunLengthsItT d_run_lengths,\n const DecodedSizesOutT d_decoded_offsets,\n const OffsetT num_runs,\n DecodedItemsOutItT d_decoded_items,\n RelativeOffsetOutItT d_relative_offsets)\n\n{\n constexpr OffsetT RUNS_PER_BLOCK = AgentTestBlockRunLengthDecode::RUNS_PER_BLOCK;\n\n __shared__ typename AgentTestBlockRunLengthDecode::TempStorage temp_storage;\n\n OffsetT block_offset = blockIdx.x * RUNS_PER_BLOCK;\n OffsetT num_valid_runs = (block_offset + RUNS_PER_BLOCK >= num_runs) ? (num_runs - block_offset)\n : RUNS_PER_BLOCK;\n\n AgentTestBlockRunLengthDecode run_length_decode_agent(temp_storage);\n run_length_decode_agent.WriteDecodedRuns(d_unique_items + block_offset,\n d_run_lengths + block_offset,\n d_decoded_items + d_decoded_offsets[blockIdx.x],\n d_relative_offsets + d_decoded_offsets[blockIdx.x],\n num_valid_runs);\n}\n\nstruct ModOp\n{\n using T = uint32_t;\n __host__ __device__ __forceinline__ T operator()(const T &x) const { return 1 + (x % 100); }\n};\n\ntemplate \nvoid TestAlgorithmSpecialisation()\n{\n constexpr uint32_t THREADS_PER_BLOCK = BLOCK_DIM_X * BLOCK_DIM_Y * BLOCK_DIM_Z;\n constexpr uint32_t RUNS_PER_BLOCK = RUNS_PER_THREAD * THREADS_PER_BLOCK;\n\n using RunItemT = float;\n using RunLengthT = uint32_t;\n using ItemItT = hipcub::CountingInputIterator;\n using RunLengthsItT =\n hipcub::TransformInputIterator>;\n\n ItemItT d_unique_items(1000U);\n RunLengthsItT d_run_lengths(hipcub::CountingInputIterator(0), ModOp{});\n\n constexpr uint32_t num_runs = 10000;\n constexpr uint32_t num_blocks = (num_runs + (RUNS_PER_BLOCK - 1U)) / RUNS_PER_BLOCK;\n\n size_t temp_storage_bytes = 0ULL;\n void *temp_storage = nullptr;\n uint32_t *h_num_decoded_total = nullptr;\n uint32_t *d_decoded_sizes = nullptr;\n uint32_t *d_decoded_offsets = nullptr;\n RunItemT *d_decoded_out = nullptr;\n RunLengthT *d_relative_offsets = nullptr;\n RunItemT *h_decoded_out = nullptr;\n RunLengthT *h_relative_offsets = nullptr;\n\n using AgentTestBlockRunLengthDecodeT = AgentTestBlockRunLengthDecode;\n\n enum : uint32_t\n {\n TIMER_SIZE_BEGIN = 0,\n TIMER_SIZE_END,\n TIMER_DECODE_BEGIN,\n TIMER_DECODE_END,\n NUM_TIMERS,\n };\n\n // Get temporary storage requirements for the scan (for computing offsets for the per-block\n // run-length decoded items)\n hipcub::DeviceScan::InclusiveSum(nullptr,\n temp_storage_bytes,\n d_decoded_sizes,\n d_decoded_offsets,\n num_blocks);\n\n // Allocate device memory\n CubDebugExit(hipMalloc(&temp_storage, temp_storage_bytes));\n CubDebugExit(hipMalloc(&d_decoded_sizes, num_blocks * sizeof(*d_decoded_sizes)));\n // Allocate for the exclusive sum PLUS the overall aggregate\n CubDebugExit(hipMalloc(&d_decoded_offsets, (num_blocks + 1) * sizeof(*d_decoded_offsets)));\n CubDebugExit(hipHostMalloc(&h_num_decoded_total, sizeof(*h_num_decoded_total)));\n\n // Get the per-block number of items being decoded (i-th thread block writing size to\n // d_decoded_sizes[i])\n BlockRunLengthDecodeGetSizeKernel\n <<>>(d_unique_items,\n d_run_lengths,\n num_runs,\n d_decoded_sizes);\n\n // Compute offsets for the runs decoded by each block (exclusive sum + aggregate)\n CubDebugExit(hipMemsetAsync(d_decoded_offsets, 0, sizeof(d_decoded_offsets[0])));\n CubDebugExit(hipcub::DeviceScan::InclusiveSum(temp_storage,\n temp_storage_bytes,\n d_decoded_sizes,\n &d_decoded_offsets[1],\n num_blocks));\n\n // Copy the total decoded size to CPU in order to allocate just the right amount of device memory\n CubDebugExit(hipMemcpy(h_num_decoded_total,\n &d_decoded_offsets[num_blocks],\n sizeof(*h_num_decoded_total),\n hipMemcpyDeviceToHost));\n\n // Allocate device memory for the run-length decoded output\n CubDebugExit(hipHostMalloc(&h_decoded_out, (*h_num_decoded_total) * sizeof(RunItemT)));\n CubDebugExit(hipMalloc(&d_decoded_out, (*h_num_decoded_total) * sizeof(RunItemT)));\n if (TEST_RELATIVE_OFFSETS)\n {\n CubDebugExit(hipMalloc(&d_relative_offsets, (*h_num_decoded_total) * sizeof(RunLengthT)));\n CubDebugExit(hipHostMalloc(&h_relative_offsets, (*h_num_decoded_total) * sizeof(RunLengthT)));\n }\n\n // Perform the block-wise run-length decoding (each block taking its offset from\n // d_decoded_offsets)\n BlockRunLengthDecodeTestKernel\n <<>>(d_unique_items,\n d_run_lengths,\n d_decoded_offsets,\n num_runs,\n d_decoded_out,\n d_relative_offsets);\n\n // Copy back results for verification\n CubDebugExit(hipMemcpy(h_decoded_out,\n d_decoded_out,\n (*h_num_decoded_total) * sizeof(*h_decoded_out),\n hipMemcpyDeviceToHost));\n\n if (TEST_RELATIVE_OFFSETS)\n {\n // Copy back the relative offsets\n CubDebugExit(hipMemcpy(h_relative_offsets,\n d_relative_offsets,\n (*h_num_decoded_total) * sizeof(*h_relative_offsets),\n hipMemcpyDeviceToHost));\n }\n\n // Generate host-side run-length decoded data for verification\n std::vector> host_golden;\n host_golden.reserve(*h_num_decoded_total);\n for (uint32_t run = 0; run < num_runs; run++)\n {\n for (RunLengthT i = 0; i < d_run_lengths[run]; i++)\n {\n host_golden.push_back({d_unique_items[run], i});\n }\n }\n\n // Verify the total run-length decoded size is correct\n REQUIRE(host_golden.size() == h_num_decoded_total[0]);\n\n // Verify the run-length decoded data is correct\n bool cmp_eq = true;\n for (uint32_t i = 0; i < host_golden.size(); i++)\n {\n if (host_golden[i].first != h_decoded_out[i])\n {\n FAIL(\"Mismatch at #\" << i << \": CPU item: \" << host_golden[i].first\n << \", GPU: \" << h_decoded_out[i] << \"\\n\");\n cmp_eq = false;\n }\n if (TEST_RELATIVE_OFFSETS)\n {\n if (host_golden[i].second != h_relative_offsets[i])\n {\n FAIL(\"Mismatch of relative offset at #\" << i\n << \": CPU item: \" << host_golden[i].first << \", GPU: \" << h_decoded_out[i]\n << \"; relative offsets: CPU: \" << host_golden[i].second\n << \", GPU: \" << h_relative_offsets[i] << \"\\n\");\n cmp_eq = false;\n break;\n }\n }\n }\n\n REQUIRE(cmp_eq == true);\n\n // Clean up memory allocations\n CubDebugExit(hipFree(temp_storage));\n CubDebugExit(hipFree(d_decoded_sizes));\n CubDebugExit(hipFree(d_decoded_offsets));\n CubDebugExit(hipFree(d_decoded_out));\n CubDebugExit(hipHostFree(h_num_decoded_total));\n CubDebugExit(hipHostFree(h_decoded_out));\n\n if (TEST_RELATIVE_OFFSETS)\n {\n CubDebugExit(hipFree(d_relative_offsets));\n CubDebugExit(hipHostFree(h_relative_offsets));\n }\n}\n\nconstexpr bool DO_TEST_RELATIVE_OFFSETS = true;\nconstexpr bool DO_NOT_TEST_RELATIVE_OFFSETS = false;\n\nconstexpr bool TEST_WITH_RUN_OFFSETS = true;\nconstexpr bool TEST_WITH_RUN_LENGTHS = false;\n\ntemplate \nstruct params_t\n{\n static constexpr int runs_per_thread = RunsPerThread;\n static constexpr int decoded_items_per_thread = DecodedItemsPerThread;\n static constexpr int block_dim_x = BlockDimX;\n static constexpr int block_dim_y = BlockDimY;\n static constexpr int block_dim_z = BlockDimZ;\n};\n\nCUB_TEST_LIST(\"Block Run Length Decode works with run lengths and offsets relative to each run\",\n \"[rld][block]\",\n params_t<1, 1, 64>,\n params_t<1, 3, 32, 2, 3>,\n params_t<1, 1, 128>,\n params_t<1, 8, 128>,\n params_t<3, 1, 256>,\n params_t<1, 8, 256>,\n params_t<8, 1, 256>,\n params_t<1, 1, 256>,\n params_t<2, 2, 384>)\n{\n using params = TestType;\n\n TestAlgorithmSpecialisation();\n}\n\nCUB_TEST_LIST(\"Block Run Length Decode works with run lengths and performs normal run-length \"\n \"decoding\",\n \"[rld][block]\",\n params_t<1, 1, 64>,\n params_t<1, 3, 32, 2, 3>,\n params_t<1, 1, 128>,\n params_t<1, 8, 128>,\n params_t<3, 1, 256>,\n params_t<1, 8, 256>,\n params_t<8, 1, 256>,\n params_t<1, 1, 256>,\n params_t<2, 2, 384>)\n{\n using params = TestType;\n\n TestAlgorithmSpecialisation();\n}\n\nCUB_TEST_LIST(\"Block Run Length Decode works with run offsets and generates offsets relative to \"\n \"each run\",\n \"[rld][block]\",\n params_t<1, 1, 64>,\n params_t<1, 3, 32, 2, 3>,\n params_t<1, 1, 128>,\n params_t<1, 8, 128>,\n params_t<3, 1, 256>,\n params_t<1, 8, 256>,\n params_t<8, 1, 256>,\n params_t<1, 1, 256>,\n params_t<2, 2, 384>)\n{\n using params = TestType;\n\n TestAlgorithmSpecialisation();\n}\n\nCUB_TEST_LIST(\"Block Run Length Decode works with run offsets and performs normal run-length \"\n \"decoding\",\n \"[rld][block]\",\n params_t<1, 1, 64>,\n params_t<1, 3, 32, 2, 3>,\n params_t<1, 1, 128>,\n params_t<1, 8, 128>,\n params_t<3, 1, 256>,\n params_t<1, 8, 256>,\n params_t<8, 1, 256>,\n params_t<1, 1, 256>,\n params_t<2, 2, 384>)\n{\n using params = TestType;\n\n TestAlgorithmSpecialisation();\n}\n\n", "source": "/kaggle/working/cub/test/catch2_test_block_run_length_decode.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void block_scan_kernel(T *in, T *out, ActionT action)\n{\n using block_scan_t = cub::BlockScan;\n using storage_t = typename block_scan_t::TempStorage;\n\n __shared__ storage_t storage;\n\n T thread_data[ItemsPerThread];\n\n const int tid = static_cast(cub::RowMajorTid(BlockDimX, BlockDimY, BlockDimZ));\n const int thread_offset = tid * ItemsPerThread;\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n const int idx = thread_offset + item;\n thread_data[item] = in[idx];\n }\n __syncthreads();\n\n block_scan_t scan(storage);\n\n action(scan, thread_data);\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n const int idx = thread_offset + item;\n out[idx] = thread_data[item];\n }\n}\n\ntemplate \nvoid block_scan(thrust::device_vector &in, thrust::device_vector &out, ActionT action)\n{\n dim3 block_dims(BlockDimX, BlockDimY, BlockDimZ);\n\n block_scan_kernel\n <<<1, block_dims>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(out.data()),\n action);\n\n REQUIRE(cudaSuccess == cudaPeekAtLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n}\n\nenum class scan_mode\n{\n exclusive,\n inclusive\n};\n\ntemplate \nstruct sum_op_t\n{\n template \n __device__ void operator()(BlockScanT &scan, T (&thread_data)[ItemsPerThread]) const\n {\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveSum(thread_data, thread_data);\n }\n else \n {\n scan.InclusiveSum(thread_data, thread_data);\n }\n }\n};\n\ntemplate \nstruct min_op_t\n{\n template \n __device__ void operator()(BlockScanT &scan, int (&thread_data)[ItemsPerThread]) const\n {\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveScan(thread_data, thread_data, cub::Min{});\n }\n else \n {\n scan.InclusiveScan(thread_data, thread_data, cub::Min{});\n }\n }\n};\n\ntemplate \nstruct sum_aggregate_op_t\n{\n int m_target_thread_id;\n T *m_d_block_aggregate;\n\n template \n __device__ void operator()(BlockScanT &scan, T (&thread_data)[ItemsPerThread]) const\n {\n T block_aggregate{};\n\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveSum(thread_data, thread_data, block_aggregate);\n }\n else \n {\n scan.InclusiveSum(thread_data, thread_data, block_aggregate);\n }\n\n const int tid = static_cast(cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z));\n\n if (tid == m_target_thread_id)\n {\n *m_d_block_aggregate = block_aggregate;\n }\n }\n};\n\ntemplate \nstruct sum_prefix_op_t\n{\n T m_prefix;\n\n struct block_prefix_op_t\n {\n int linear_tid;\n T prefix;\n\n __device__ block_prefix_op_t(int linear_tid, T prefix)\n : linear_tid(linear_tid)\n , prefix(prefix)\n {}\n\n __device__ T operator()(T block_aggregate)\n {\n T retval = (linear_tid == 0) ? prefix : T{};\n prefix = prefix + block_aggregate;\n return retval;\n }\n };\n\n\n template \n __device__ void operator()(BlockScanT &scan, T (&thread_data)[ItemsPerThread]) const\n {\n const int tid = static_cast(cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z));\n block_prefix_op_t prefix_op{tid, m_prefix};\n\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveSum(thread_data, thread_data, prefix_op);\n }\n else \n {\n scan.InclusiveSum(thread_data, thread_data, prefix_op);\n }\n }\n};\n\ntemplate \nstruct min_prefix_op_t\n{\n T m_prefix;\n static constexpr T min_identity = std::numeric_limits::max();\n\n struct block_prefix_op_t\n {\n int linear_tid;\n T prefix;\n\n __device__ block_prefix_op_t(int linear_tid, T prefix)\n : linear_tid(linear_tid)\n , prefix(prefix)\n {}\n\n __device__ T operator()(T block_aggregate)\n {\n T retval = (linear_tid == 0) ? prefix : min_identity;\n prefix = cub::Min{}(prefix, block_aggregate);\n return retval;\n }\n };\n\n\n template \n __device__ void operator()(BlockScanT &scan, T (&thread_data)[ItemsPerThread]) const\n {\n const int tid = static_cast(cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z));\n block_prefix_op_t prefix_op{tid, m_prefix};\n\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveScan(thread_data, thread_data, cub::Min{}, prefix_op);\n }\n else\n {\n scan.InclusiveScan(thread_data, thread_data, cub::Min{}, prefix_op);\n }\n }\n};\n\ntemplate \nT host_scan(scan_mode mode, thrust::host_vector &result, ScanOpT scan_op, T initial_value = T{})\n{\n if (result.empty()) \n {\n return {};\n }\n\n T accumulator = static_cast(scan_op(initial_value, result[0]));\n T block_accumulator = result[0];\n\n if (mode == scan_mode::exclusive)\n {\n result[0] = initial_value;\n\n for (std::size_t i = 1; i < result.size(); i++)\n {\n T tmp = result[i];\n result[i] = accumulator;\n accumulator = static_cast(scan_op(accumulator, tmp));\n block_accumulator = static_cast(scan_op(block_accumulator, tmp));\n }\n }\n else\n {\n result[0] = accumulator;\n\n for (std::size_t i = 1; i < result.size(); i++)\n {\n accumulator = static_cast(scan_op(accumulator, result[i]));\n block_accumulator = static_cast(scan_op(block_accumulator, result[i]));\n result[i] = accumulator;\n }\n }\n\n return block_accumulator;\n}\n\n// %PARAM% ALGO_TYPE alg 0:1:2\n// %PARAM% TEST_MODE mode 0:1\n\nusing types = c2h::type_list;\nusing vec_types = c2h::type_list;\nusing block_dim_x = c2h::enum_type_list;\nusing block_dim_yz = c2h::enum_type_list;\nusing items_per_thread = c2h::enum_type_list;\nusing algorithms = c2h::enum_type_list;\nusing algorithm = c2h::enum_type_list::value>;\n\n#if TEST_MODE == 0\nusing modes = c2h::enum_type_list;\n#else\nusing modes = c2h::enum_type_list;\n#endif\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int block_dim_x = c2h::get<1, TestType>::value;\n static constexpr int block_dim_y = c2h::get<2, TestType>::value;\n static constexpr int block_dim_z = block_dim_y;\n static constexpr int items_per_thread = c2h::get<3, TestType>::value;\n static constexpr int tile_size = items_per_thread * block_dim_x * block_dim_y * block_dim_z;\n\n static constexpr cub::BlockScanAlgorithm algorithm = c2h::get<4, TestType>::value;\n static constexpr scan_mode mode = c2h::get<5, TestType>::value;\n};\n\nCUB_TEST(\"Block scan works with sum\",\n \"[scan][block]\",\n types,\n block_dim_x,\n block_dim_yz,\n items_per_thread,\n algorithm,\n modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n block_scan(d_in, d_out, sum_op_t{});\n\n\n thrust::host_vector h_out = d_in;\n host_scan(params::mode, h_out, std::plus{});\n\n REQUIRE_APPROX_EQ(h_out, d_out);\n}\n\nCUB_TEST(\"Block scan works with vec types\", \"[scan][block]\", vec_types, algorithm, modes)\n{\n constexpr int items_per_thread = 3;\n constexpr int block_dim_x = 256;\n constexpr int block_dim_y = 1;\n constexpr int block_dim_z = 1;\n constexpr int tile_size = items_per_thread * block_dim_x * block_dim_y * block_dim_z;\n constexpr cub::BlockScanAlgorithm algorithm = c2h::get<1, TestType>::value;\n constexpr scan_mode mode = c2h::get<2, TestType>::value;\n\n using type = typename c2h::get<0, TestType>;\n\n thrust::device_vector d_out(tile_size);\n thrust::device_vector d_in(tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n block_scan(d_in,\n d_out,\n sum_op_t{});\n\n thrust::host_vector h_out = d_in;\n host_scan(mode, h_out, std::plus{});\n\n REQUIRE(h_out == d_out);\n}\n\nCUB_TEST(\"Block scan works with custom types\", \"[scan][block]\", algorithm, modes)\n{\n constexpr int items_per_thread = 3;\n constexpr int block_dim_x = 256;\n constexpr int block_dim_y = 1;\n constexpr int block_dim_z = 1;\n constexpr int tile_size = items_per_thread * block_dim_x * block_dim_y * block_dim_z;\n constexpr cub::BlockScanAlgorithm algorithm = c2h::get<0, TestType>::value;\n constexpr scan_mode mode = c2h::get<1, TestType>::value;\n\n using type = c2h::custom_type_t;\n\n thrust::device_vector d_out(tile_size);\n thrust::device_vector d_in(tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n block_scan(d_in,\n d_out,\n sum_op_t{});\n\n thrust::host_vector h_out = d_in;\n host_scan(mode, h_out, std::plus{});\n\n REQUIRE(h_out == d_out);\n}\n\nCUB_TEST(\"Block scan returns valid block aggregate\",\n \"[scan][block]\",\n algorithm,\n modes,\n block_dim_yz)\n{\n constexpr int items_per_thread = 3;\n constexpr int block_dim_x = 64;\n constexpr int block_dim_y = c2h::get<2, TestType>::value;\n constexpr int block_dim_z = block_dim_y;\n constexpr int threads_in_block = block_dim_x * block_dim_y * block_dim_z;\n constexpr int tile_size = items_per_thread * threads_in_block;\n constexpr cub::BlockScanAlgorithm algorithm = c2h::get<0, TestType>::value;\n constexpr scan_mode mode = c2h::get<1, TestType>::value;\n\n using type = c2h::custom_type_t;\n\n const int target_thread_id = GENERATE_COPY(take(2, random(0, threads_in_block - 1)));\n\n thrust::device_vector d_block_aggregate(1);\n thrust::device_vector d_out(tile_size);\n thrust::device_vector d_in(tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n block_scan(\n d_in,\n d_out,\n sum_aggregate_op_t{target_thread_id,\n thrust::raw_pointer_cast(d_block_aggregate.data())});\n\n thrust::host_vector h_out = d_in;\n type block_aggregate = host_scan(mode, h_out, std::plus{});\n\n REQUIRE(h_out == d_out);\n REQUIRE(block_aggregate == d_block_aggregate[0]);\n}\n\nCUB_TEST(\"Block scan supports prefix op\",\n \"[scan][block]\",\n algorithm,\n modes,\n block_dim_yz)\n{\n constexpr int items_per_thread = 3;\n constexpr int block_dim_x = 64;\n constexpr int block_dim_y = c2h::get<2, TestType>::value;\n constexpr int block_dim_z = block_dim_y;\n constexpr int threads_in_block = block_dim_x * block_dim_y * block_dim_z;\n constexpr int tile_size = items_per_thread * threads_in_block;\n constexpr cub::BlockScanAlgorithm algorithm = c2h::get<0, TestType>::value;\n constexpr scan_mode mode = c2h::get<1, TestType>::value;\n\n using type = int;\n\n const type prefix = GENERATE_COPY(take(2, random(0, tile_size)));\n\n thrust::device_vector d_out(tile_size);\n thrust::device_vector d_in(tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n block_scan(\n d_in,\n d_out,\n sum_prefix_op_t{prefix});\n\n thrust::host_vector h_out = d_in;\n host_scan(mode, h_out, std::plus{}, prefix);\n\n REQUIRE(h_out == d_out);\n}\n\nCUB_TEST(\"Block scan supports custom scan op\",\n \"[scan][block]\",\n algorithm,\n modes,\n block_dim_yz)\n{\n constexpr int items_per_thread = 3;\n constexpr int block_dim_x = 64;\n constexpr int block_dim_y = c2h::get<2, TestType>::value;\n constexpr int block_dim_z = block_dim_y;\n constexpr int threads_in_block = block_dim_x * block_dim_y * block_dim_z;\n constexpr int tile_size = items_per_thread * threads_in_block;\n constexpr cub::BlockScanAlgorithm algorithm = c2h::get<0, TestType>::value;\n constexpr scan_mode mode = c2h::get<1, TestType>::value;\n\n using type = int;\n\n thrust::device_vector d_out(tile_size);\n thrust::device_vector d_in(tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n d_in[0] = INT_MIN;\n\n block_scan(\n d_in,\n d_out,\n min_op_t{});\n\n thrust::host_vector h_out = d_in;\n host_scan(mode, h_out, [](type l, type r) { return std::min(l, r); }, INT_MIN );\n\n REQUIRE(h_out == d_out);\n}\n\nCUB_TEST(\"Block scan supports prefix op and custom scan op\",\n \"[scan][block]\",\n algorithm,\n modes,\n block_dim_yz)\n{\n constexpr int items_per_thread = 3;\n constexpr int block_dim_x = 64;\n constexpr int block_dim_y = c2h::get<2, TestType>::value;\n constexpr int block_dim_z = block_dim_y;\n constexpr int threads_in_block = block_dim_x * block_dim_y * block_dim_z;\n constexpr int tile_size = items_per_thread * threads_in_block;\n constexpr cub::BlockScanAlgorithm algorithm = c2h::get<0, TestType>::value;\n constexpr scan_mode mode = c2h::get<1, TestType>::value;\n\n using type = int;\n\n const type prefix = GENERATE_COPY(take(2, random(0, tile_size)));\n\n thrust::device_vector d_out(tile_size);\n thrust::device_vector d_in(tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n block_scan(\n d_in,\n d_out,\n min_prefix_op_t{prefix});\n\n thrust::host_vector h_out = d_in;\n host_scan(mode, h_out, [] (type a, type b) { return std::min(a, b); }, prefix);\n\n REQUIRE(h_out == d_out);\n}\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void block_scan_kernel(T *in, T *out, ActionT action)\n{\n using block_scan_t = hipcub::BlockScan;\n using storage_t = typename block_scan_t::TempStorage;\n\n __shared__ storage_t storage;\n\n T thread_data[ItemsPerThread];\n\n const int tid = static_cast(hipcub::RowMajorTid(BlockDimX, BlockDimY, BlockDimZ));\n const int thread_offset = tid * ItemsPerThread;\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n const int idx = thread_offset + item;\n thread_data[item] = in[idx];\n }\n __syncthreads();\n\n block_scan_t scan(storage);\n\n action(scan, thread_data);\n\n for (int item = 0; item < ItemsPerThread; item++)\n {\n const int idx = thread_offset + item;\n out[idx] = thread_data[item];\n }\n}\n\ntemplate \nvoid block_scan(thrust::device_vector &in, thrust::device_vector &out, ActionT action)\n{\n dim3 block_dims(BlockDimX, BlockDimY, BlockDimZ);\n\n block_scan_kernel\n <<<1, block_dims>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(out.data()),\n action);\n\n REQUIRE(hipSuccess == hipPeekAtLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n}\n\nenum class scan_mode\n{\n exclusive,\n inclusive\n};\n\ntemplate \nstruct sum_op_t\n{\n template \n __device__ void operator()(BlockScanT &scan, T (&thread_data)[ItemsPerThread]) const\n {\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveSum(thread_data, thread_data);\n }\n else \n {\n scan.InclusiveSum(thread_data, thread_data);\n }\n }\n};\n\ntemplate \nstruct min_op_t\n{\n template \n __device__ void operator()(BlockScanT &scan, int (&thread_data)[ItemsPerThread]) const\n {\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveScan(thread_data, thread_data, hipcub::Min{});\n }\n else \n {\n scan.InclusiveScan(thread_data, thread_data, hipcub::Min{});\n }\n }\n};\n\ntemplate \nstruct sum_aggregate_op_t\n{\n int m_target_thread_id;\n T *m_d_block_aggregate;\n\n template \n __device__ void operator()(BlockScanT &scan, T (&thread_data)[ItemsPerThread]) const\n {\n T block_aggregate{};\n\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveSum(thread_data, thread_data, block_aggregate);\n }\n else \n {\n scan.InclusiveSum(thread_data, thread_data, block_aggregate);\n }\n\n const int tid = static_cast(hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z));\n\n if (tid == m_target_thread_id)\n {\n *m_d_block_aggregate = block_aggregate;\n }\n }\n};\n\ntemplate \nstruct sum_prefix_op_t\n{\n T m_prefix;\n\n struct block_prefix_op_t\n {\n int linear_tid;\n T prefix;\n\n __device__ block_prefix_op_t(int linear_tid, T prefix)\n : linear_tid(linear_tid)\n , prefix(prefix)\n {}\n\n __device__ T operator()(T block_aggregate)\n {\n T retval = (linear_tid == 0) ? prefix : T{};\n prefix = prefix + block_aggregate;\n return retval;\n }\n };\n\n\n template \n __device__ void operator()(BlockScanT &scan, T (&thread_data)[ItemsPerThread]) const\n {\n const int tid = static_cast(hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z));\n block_prefix_op_t prefix_op{tid, m_prefix};\n\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveSum(thread_data, thread_data, prefix_op);\n }\n else \n {\n scan.InclusiveSum(thread_data, thread_data, prefix_op);\n }\n }\n};\n\ntemplate \nstruct min_prefix_op_t\n{\n T m_prefix;\n static constexpr T min_identity = std::numeric_limits::max();\n\n struct block_prefix_op_t\n {\n int linear_tid;\n T prefix;\n\n __device__ block_prefix_op_t(int linear_tid, T prefix)\n : linear_tid(linear_tid)\n , prefix(prefix)\n {}\n\n __device__ T operator()(T block_aggregate)\n {\n T retval = (linear_tid == 0) ? prefix : min_identity;\n prefix = hipcub::Min{}(prefix, block_aggregate);\n return retval;\n }\n };\n\n\n template \n __device__ void operator()(BlockScanT &scan, T (&thread_data)[ItemsPerThread]) const\n {\n const int tid = static_cast(hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z));\n block_prefix_op_t prefix_op{tid, m_prefix};\n\n if (Mode == scan_mode::exclusive)\n {\n scan.ExclusiveScan(thread_data, thread_data, hipcub::Min{}, prefix_op);\n }\n else\n {\n scan.InclusiveScan(thread_data, thread_data, hipcub::Min{}, prefix_op);\n }\n }\n};\n\ntemplate \nT host_scan(scan_mode mode, thrust::host_vector &result, ScanOpT scan_op, T initial_value = T{})\n{\n if (result.empty()) \n {\n return {};\n }\n\n T accumulator = static_cast(scan_op(initial_value, result[0]));\n T block_accumulator = result[0];\n\n if (mode == scan_mode::exclusive)\n {\n result[0] = initial_value;\n\n for (std::size_t i = 1; i < result.size(); i++)\n {\n T tmp = result[i];\n result[i] = accumulator;\n accumulator = static_cast(scan_op(accumulator, tmp));\n block_accumulator = static_cast(scan_op(block_accumulator, tmp));\n }\n }\n else\n {\n result[0] = accumulator;\n\n for (std::size_t i = 1; i < result.size(); i++)\n {\n accumulator = static_cast(scan_op(accumulator, result[i]));\n block_accumulator = static_cast(scan_op(block_accumulator, result[i]));\n result[i] = accumulator;\n }\n }\n\n return block_accumulator;\n}\n\n// %PARAM% ALGO_TYPE alg 0:1:2\n// %PARAM% TEST_MODE mode 0:1\n\nusing types = c2h::type_list;\nusing vec_types = c2h::type_list;\nusing block_dim_x = c2h::enum_type_list;\nusing block_dim_yz = c2h::enum_type_list;\nusing items_per_thread = c2h::enum_type_list;\nusing algorithms = c2h::enum_type_list;\nusing algorithm = c2h::enum_type_list::value>;\n\n#if TEST_MODE == 0\nusing modes = c2h::enum_type_list;\n#else\nusing modes = c2h::enum_type_list;\n#endif\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int block_dim_x = c2h::get<1, TestType>::value;\n static constexpr int block_dim_y = c2h::get<2, TestType>::value;\n static constexpr int block_dim_z = block_dim_y;\n static constexpr int items_per_thread = c2h::get<3, TestType>::value;\n static constexpr int tile_size = items_per_thread * block_dim_x * block_dim_y * block_dim_z;\n\n static constexpr hipcub::BlockScanAlgorithm algorithm = c2h::get<4, TestType>::value;\n static constexpr scan_mode mode = c2h::get<5, TestType>::value;\n};\n\nCUB_TEST(\"Block scan works with sum\",\n \"[scan][block]\",\n types,\n block_dim_x,\n block_dim_yz,\n items_per_thread,\n algorithm,\n modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_out(params::tile_size);\n thrust::device_vector d_in(params::tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n block_scan(d_in, d_out, sum_op_t{});\n\n\n thrust::host_vector h_out = d_in;\n host_scan(params::mode, h_out, std::plus{});\n\n REQUIRE_APPROX_EQ(h_out, d_out);\n}\n\nCUB_TEST(\"Block scan works with vec types\", \"[scan][block]\", vec_types, algorithm, modes)\n{\n constexpr int items_per_thread = 3;\n constexpr int block_dim_x = 256;\n constexpr int block_dim_y = 1;\n constexpr int block_dim_z = 1;\n constexpr int tile_size = items_per_thread * block_dim_x * block_dim_y * block_dim_z;\n constexpr hipcub::BlockScanAlgorithm algorithm = c2h::get<1, TestType>::value;\n constexpr scan_mode mode = c2h::get<2, TestType>::value;\n\n using type = typename c2h::get<0, TestType>;\n\n thrust::device_vector d_out(tile_size);\n thrust::device_vector d_in(tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n block_scan(d_in,\n d_out,\n sum_op_t{});\n\n thrust::host_vector h_out = d_in;\n host_scan(mode, h_out, std::plus{});\n\n REQUIRE(h_out == d_out);\n}\n\nCUB_TEST(\"Block scan works with custom types\", \"[scan][block]\", algorithm, modes)\n{\n constexpr int items_per_thread = 3;\n constexpr int block_dim_x = 256;\n constexpr int block_dim_y = 1;\n constexpr int block_dim_z = 1;\n constexpr int tile_size = items_per_thread * block_dim_x * block_dim_y * block_dim_z;\n constexpr hipcub::BlockScanAlgorithm algorithm = c2h::get<0, TestType>::value;\n constexpr scan_mode mode = c2h::get<1, TestType>::value;\n\n using type = c2h::custom_type_t;\n\n thrust::device_vector d_out(tile_size);\n thrust::device_vector d_in(tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n block_scan(d_in,\n d_out,\n sum_op_t{});\n\n thrust::host_vector h_out = d_in;\n host_scan(mode, h_out, std::plus{});\n\n REQUIRE(h_out == d_out);\n}\n\nCUB_TEST(\"Block scan returns valid block aggregate\",\n \"[scan][block]\",\n algorithm,\n modes,\n block_dim_yz)\n{\n constexpr int items_per_thread = 3;\n constexpr int block_dim_x = 64;\n constexpr int block_dim_y = c2h::get<2, TestType>::value;\n constexpr int block_dim_z = block_dim_y;\n constexpr int threads_in_block = block_dim_x * block_dim_y * block_dim_z;\n constexpr int tile_size = items_per_thread * threads_in_block;\n constexpr hipcub::BlockScanAlgorithm algorithm = c2h::get<0, TestType>::value;\n constexpr scan_mode mode = c2h::get<1, TestType>::value;\n\n using type = c2h::custom_type_t;\n\n const int target_thread_id = GENERATE_COPY(take(2, random(0, threads_in_block - 1)));\n\n thrust::device_vector d_block_aggregate(1);\n thrust::device_vector d_out(tile_size);\n thrust::device_vector d_in(tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n block_scan(\n d_in,\n d_out,\n sum_aggregate_op_t{target_thread_id,\n thrust::raw_pointer_cast(d_block_aggregate.data())});\n\n thrust::host_vector h_out = d_in;\n type block_aggregate = host_scan(mode, h_out, std::plus{});\n\n REQUIRE(h_out == d_out);\n REQUIRE(block_aggregate == d_block_aggregate[0]);\n}\n\nCUB_TEST(\"Block scan supports prefix op\",\n \"[scan][block]\",\n algorithm,\n modes,\n block_dim_yz)\n{\n constexpr int items_per_thread = 3;\n constexpr int block_dim_x = 64;\n constexpr int block_dim_y = c2h::get<2, TestType>::value;\n constexpr int block_dim_z = block_dim_y;\n constexpr int threads_in_block = block_dim_x * block_dim_y * block_dim_z;\n constexpr int tile_size = items_per_thread * threads_in_block;\n constexpr hipcub::BlockScanAlgorithm algorithm = c2h::get<0, TestType>::value;\n constexpr scan_mode mode = c2h::get<1, TestType>::value;\n\n using type = int;\n\n const type prefix = GENERATE_COPY(take(2, random(0, tile_size)));\n\n thrust::device_vector d_out(tile_size);\n thrust::device_vector d_in(tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n block_scan(\n d_in,\n d_out,\n sum_prefix_op_t{prefix});\n\n thrust::host_vector h_out = d_in;\n host_scan(mode, h_out, std::plus{}, prefix);\n\n REQUIRE(h_out == d_out);\n}\n\nCUB_TEST(\"Block scan supports custom scan op\",\n \"[scan][block]\",\n algorithm,\n modes,\n block_dim_yz)\n{\n constexpr int items_per_thread = 3;\n constexpr int block_dim_x = 64;\n constexpr int block_dim_y = c2h::get<2, TestType>::value;\n constexpr int block_dim_z = block_dim_y;\n constexpr int threads_in_block = block_dim_x * block_dim_y * block_dim_z;\n constexpr int tile_size = items_per_thread * threads_in_block;\n constexpr hipcub::BlockScanAlgorithm algorithm = c2h::get<0, TestType>::value;\n constexpr scan_mode mode = c2h::get<1, TestType>::value;\n\n using type = int;\n\n thrust::device_vector d_out(tile_size);\n thrust::device_vector d_in(tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n d_in[0] = INT_MIN;\n\n block_scan(\n d_in,\n d_out,\n min_op_t{});\n\n thrust::host_vector h_out = d_in;\n host_scan(mode, h_out, [](type l, type r) { return std::min(l, r); }, INT_MIN );\n\n REQUIRE(h_out == d_out);\n}\n\nCUB_TEST(\"Block scan supports prefix op and custom scan op\",\n \"[scan][block]\",\n algorithm,\n modes,\n block_dim_yz)\n{\n constexpr int items_per_thread = 3;\n constexpr int block_dim_x = 64;\n constexpr int block_dim_y = c2h::get<2, TestType>::value;\n constexpr int block_dim_z = block_dim_y;\n constexpr int threads_in_block = block_dim_x * block_dim_y * block_dim_z;\n constexpr int tile_size = items_per_thread * threads_in_block;\n constexpr hipcub::BlockScanAlgorithm algorithm = c2h::get<0, TestType>::value;\n constexpr scan_mode mode = c2h::get<1, TestType>::value;\n\n using type = int;\n\n const type prefix = GENERATE_COPY(take(2, random(0, tile_size)));\n\n thrust::device_vector d_out(tile_size);\n thrust::device_vector d_in(tile_size);\n c2h::gen(CUB_SEED(10), d_in);\n\n block_scan(\n d_in,\n d_out,\n min_prefix_op_t{prefix});\n\n thrust::host_vector h_out = d_in;\n host_scan(mode, h_out, [] (type a, type b) { return std::min(a, b); }, prefix);\n\n REQUIRE(h_out == d_out);\n}\n\n", "source": "/kaggle/working/cub/test/catch2_test_block_scan.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"cub/util_type.cuh\"\n#define CUB_STDERR\n\n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"test_util.h\"\n#include \n\nbool g_verbose = false;\ncub::CachingDeviceAllocator g_allocator(true);\n\ntemplate \n__launch_bounds__(BlockThreads, 1) __global__ void kernel(Key *d_keys, int *d_ranks)\n{\n using block_radix_rank = cub::detail::\n block_radix_rank_t;\n\n using storage_t = typename block_radix_rank::TempStorage;\n\n // Allocate temp storage in shared memory\n __shared__ storage_t temp_storage;\n\n // Items per thread\n Key keys[ItemsPerThread];\n int ranks[ItemsPerThread];\n\n constexpr bool uses_warp_striped_arrangement =\n RankAlgorithm == cub::RadixRankAlgorithm::RADIX_RANK_MATCH ||\n RankAlgorithm == cub::RadixRankAlgorithm::RADIX_RANK_MATCH_EARLY_COUNTS_ANY ||\n RankAlgorithm == cub::RadixRankAlgorithm::RADIX_RANK_MATCH_EARLY_COUNTS_ATOMIC_OR;\n\n if (uses_warp_striped_arrangement)\n {\n cub::LoadDirectWarpStriped(threadIdx.x, d_keys, keys);\n }\n else\n {\n cub::LoadDirectBlocked(threadIdx.x, d_keys, keys);\n }\n\n cub::BFEDigitExtractor extractor(0, RadixBits);\n block_radix_rank(temp_storage).RankKeys(keys, ranks, extractor);\n\n if (uses_warp_striped_arrangement)\n {\n cub::StoreDirectWarpStriped(threadIdx.x, d_ranks, ranks);\n }\n else\n {\n cub::StoreDirectBlocked(threadIdx.x, d_ranks, ranks);\n }\n}\n\n//---------------------------------------------------------------------\n// Host testing subroutines\n//---------------------------------------------------------------------\n\n/**\n * Simple key-value pairing\n */\ntemplate \nstruct pair_t\n{\n Key key;\n int value;\n\n bool operator<(const pair_t &b) const { return (key < b.key); }\n};\n\ntemplate \nvoid Initialize(GenMode gen_mode, Key *h_keys, int *h_reference_ranks, int num_items, int num_bits)\n{\n std::unique_ptr[]> h_pairs_storage(new pair_t[num_items]);\n pair_t *h_pairs = h_pairs_storage.get();\n\n for (int i = 0; i < num_items; ++i)\n {\n InitValue(gen_mode, h_keys[i], i);\n\n // Mask off unwanted portions\n std::uint64_t base = 0;\n memcpy(&base, &h_keys[i], sizeof(Key));\n base &= (1ull << num_bits) - 1;\n memcpy(&h_keys[i], &base, sizeof(Key));\n\n h_pairs[i].key = h_keys[i];\n h_pairs[i].value = i;\n }\n\n if (DESCENDING)\n {\n std::reverse(h_pairs, h_pairs + num_items);\n }\n\n std::stable_sort(h_pairs, h_pairs + num_items);\n\n if (DESCENDING)\n {\n std::reverse(h_pairs, h_pairs + num_items);\n }\n\n for (int i = 0; i < num_items; ++i)\n {\n h_reference_ranks[h_pairs[i].value] = i;\n }\n}\n\ntemplate \nvoid TestDriver(GenMode gen_mode)\n{\n constexpr int tile_size = BlockThreads * ItemsPerThread;\n\n // Allocate host arrays\n std::unique_ptr h_keys(new Key[tile_size]);\n std::unique_ptr h_ranks(new int[tile_size]);\n std::unique_ptr h_reference_ranks(new int[tile_size]);\n\n // Allocate device arrays\n Key *d_keys = nullptr;\n int *d_ranks = nullptr;\n\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_keys, sizeof(Key) * tile_size));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_ranks, sizeof(int) * tile_size));\n\n // Initialize problem and solution on host\n Initialize(gen_mode, h_keys.get(), h_reference_ranks.get(), tile_size, RadixBits);\n\n // Copy problem to device\n CubDebugExit(cudaMemcpy(d_keys, h_keys.get(), sizeof(Key) * tile_size, cudaMemcpyHostToDevice));\n\n // Run kernel\n kernel\n <<<1, BlockThreads>>>(d_keys, d_ranks);\n\n // Flush kernel output / errors\n CubDebugExit(cudaPeekAtLastError());\n CubDebugExit(cudaDeviceSynchronize());\n\n // Check keys results\n const bool compare =\n CompareDeviceResults(h_reference_ranks.get(), d_ranks, tile_size, g_verbose, g_verbose);\n AssertEquals(0, compare);\n\n if (d_keys)\n {\n CubDebugExit(g_allocator.DeviceFree(d_keys));\n }\n\n if (d_ranks)\n {\n CubDebugExit(g_allocator.DeviceFree(d_ranks));\n }\n}\n\ntemplate \nvoid TestValid(cub::Int2Type /*fits_smem_capacity*/)\n{\n TestDriver(\n UNIFORM);\n\n TestDriver(\n INTEGER_SEED);\n}\n\ntemplate \nvoid TestValid(cub::Int2Type fits_smem_capacity)\n{}\n\ntemplate \nvoid Test()\n{\n // Check size of smem storage for the target arch to make sure it will fit\n using block_radix_rank = cub::detail::\n block_radix_rank_t;\n using storage_t = typename block_radix_rank::TempStorage;\n\n cub::Int2Type<(sizeof(storage_t) <= 48 * 1024)> fits_smem_capacity;\n\n TestValid(\n fits_smem_capacity);\n}\n\ntemplate \nvoid Test()\n{\n Test();\n Test();\n}\n\ntemplate \nvoid Test()\n{\n Test();\n Test();\n}\n\ntemplate \nvoid Test()\n{\n Test();\n Test();\n}\n\ntemplate \nvoid Test()\n{\n Test();\n Test();\n}\n\ntemplate \nvoid Test()\n{\n Test();\n Test();\n}\n\ntemplate \nvoid Test(cub::Int2Type /* multiple of hw warp */)\n{\n Test();\n\n // TODO(senior-zero):\n // - RADIX_RANK_MATCH_EARLY_COUNTS_ANY\n // - RADIX_RANK_MATCH_EARLY_COUNTS_ATOMIC_OR\n}\n\ntemplate \nvoid Test(cub::Int2Type /* multiple of hw warp */)\n{}\n\ntemplate \nvoid Test()\n{\n Test();\n Test();\n\n Test(cub::Int2Type<(BlockThreads % 32) == 0>{});\n}\n\nint main(int argc, char **argv)\n{\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\",\n argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n Test<16>();\n Test<32>();\n Test<128>();\n Test<130>();\n\n return 0;\n}\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#include \"cub/util_type.cuh\"\n#define HIPCUB_STDERR\n\n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"test_util.h\"\n#include \n\nbool g_verbose = false;\nhipcub::CachingDeviceAllocator g_allocator(true);\n\ntemplate \n__launch_bounds__(BlockThreads, 1) __global__ void kernel(Key *d_keys, int *d_ranks)\n{\n using block_radix_rank = hipcub::detail::\n block_radix_rank_t;\n\n using storage_t = typename block_radix_rank::TempStorage;\n\n // Allocate temp storage in shared memory\n __shared__ storage_t temp_storage;\n\n // Items per thread\n Key keys[ItemsPerThread];\n int ranks[ItemsPerThread];\n\n constexpr bool uses_warp_striped_arrangement =\n RankAlgorithm == hipcub::RadixRankAlgorithm::RADIX_RANK_MATCH ||\n RankAlgorithm == hipcub::RadixRankAlgorithm::RADIX_RANK_MATCH_EARLY_COUNTS_ANY ||\n RankAlgorithm == hipcub::RadixRankAlgorithm::RADIX_RANK_MATCH_EARLY_COUNTS_ATOMIC_OR;\n\n if (uses_warp_striped_arrangement)\n {\n hipcub::LoadDirectWarpStriped(threadIdx.x, d_keys, keys);\n }\n else\n {\n hipcub::LoadDirectBlocked(threadIdx.x, d_keys, keys);\n }\n\n hipcub::BFEDigitExtractor extractor(0, RadixBits);\n block_radix_rank(temp_storage).RankKeys(keys, ranks, extractor);\n\n if (uses_warp_striped_arrangement)\n {\n hipcub::StoreDirectWarpStriped(threadIdx.x, d_ranks, ranks);\n }\n else\n {\n hipcub::StoreDirectBlocked(threadIdx.x, d_ranks, ranks);\n }\n}\n\n//---------------------------------------------------------------------\n// Host testing subroutines\n//---------------------------------------------------------------------\n\n/**\n * Simple key-value pairing\n */\ntemplate \nstruct pair_t\n{\n Key key;\n int value;\n\n bool operator<(const pair_t &b) const { return (key < b.key); }\n};\n\ntemplate \nvoid Initialize(GenMode gen_mode, Key *h_keys, int *h_reference_ranks, int num_items, int num_bits)\n{\n std::unique_ptr[]> h_pairs_storage(new pair_t[num_items]);\n pair_t *h_pairs = h_pairs_storage.get();\n\n for (int i = 0; i < num_items; ++i)\n {\n InitValue(gen_mode, h_keys[i], i);\n\n // Mask off unwanted portions\n std::uint64_t base = 0;\n memcpy(&base, &h_keys[i], sizeof(Key));\n base &= (1ull << num_bits) - 1;\n memcpy(&h_keys[i], &base, sizeof(Key));\n\n h_pairs[i].key = h_keys[i];\n h_pairs[i].value = i;\n }\n\n if (DESCENDING)\n {\n std::reverse(h_pairs, h_pairs + num_items);\n }\n\n std::stable_sort(h_pairs, h_pairs + num_items);\n\n if (DESCENDING)\n {\n std::reverse(h_pairs, h_pairs + num_items);\n }\n\n for (int i = 0; i < num_items; ++i)\n {\n h_reference_ranks[h_pairs[i].value] = i;\n }\n}\n\ntemplate \nvoid TestDriver(GenMode gen_mode)\n{\n constexpr int tile_size = BlockThreads * ItemsPerThread;\n\n // Allocate host arrays\n std::unique_ptr h_keys(new Key[tile_size]);\n std::unique_ptr h_ranks(new int[tile_size]);\n std::unique_ptr h_reference_ranks(new int[tile_size]);\n\n // Allocate device arrays\n Key *d_keys = nullptr;\n int *d_ranks = nullptr;\n\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_keys, sizeof(Key) * tile_size));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_ranks, sizeof(int) * tile_size));\n\n // Initialize problem and solution on host\n Initialize(gen_mode, h_keys.get(), h_reference_ranks.get(), tile_size, RadixBits);\n\n // Copy problem to device\n CubDebugExit(hipMemcpy(d_keys, h_keys.get(), sizeof(Key) * tile_size, hipMemcpyHostToDevice));\n\n // Run kernel\n kernel\n <<<1, BlockThreads>>>(d_keys, d_ranks);\n\n // Flush kernel output / errors\n CubDebugExit(hipPeekAtLastError());\n CubDebugExit(hipDeviceSynchronize());\n\n // Check keys results\n const bool compare =\n CompareDeviceResults(h_reference_ranks.get(), d_ranks, tile_size, g_verbose, g_verbose);\n AssertEquals(0, compare);\n\n if (d_keys)\n {\n CubDebugExit(g_allocator.DeviceFree(d_keys));\n }\n\n if (d_ranks)\n {\n CubDebugExit(g_allocator.DeviceFree(d_ranks));\n }\n}\n\ntemplate \nvoid TestValid(hipcub::Int2Type /*fits_smem_capacity*/)\n{\n TestDriver(\n UNIFORM);\n\n TestDriver(\n INTEGER_SEED);\n}\n\ntemplate \nvoid TestValid(hipcub::Int2Type fits_smem_capacity)\n{}\n\ntemplate \nvoid Test()\n{\n // Check size of smem storage for the target arch to make sure it will fit\n using block_radix_rank = hipcub::detail::\n block_radix_rank_t;\n using storage_t = typename block_radix_rank::TempStorage;\n\n hipcub::Int2Type<(sizeof(storage_t) <= 48 * 1024)> fits_smem_capacity;\n\n TestValid(\n fits_smem_capacity);\n}\n\ntemplate \nvoid Test()\n{\n Test();\n Test();\n}\n\ntemplate \nvoid Test()\n{\n Test();\n Test();\n}\n\ntemplate \nvoid Test()\n{\n Test();\n Test();\n}\n\ntemplate \nvoid Test()\n{\n Test();\n Test();\n}\n\ntemplate \nvoid Test()\n{\n Test();\n Test();\n}\n\ntemplate \nvoid Test(hipcub::Int2Type /* multiple of hw warp */)\n{\n Test();\n\n // TODO(senior-zero):\n // - RADIX_RANK_MATCH_EARLY_COUNTS_ANY\n // - RADIX_RANK_MATCH_EARLY_COUNTS_ATOMIC_OR\n}\n\ntemplate \nvoid Test(hipcub::Int2Type /* multiple of hw warp */)\n{}\n\ntemplate \nvoid Test()\n{\n Test();\n Test();\n\n Test(hipcub::Int2Type<(BlockThreads % 32) == 0>{});\n}\n\nint main(int argc, char **argv)\n{\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\",\n argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n Test<16>();\n Test<32>();\n Test<128>();\n Test<130>();\n\n return 0;\n}\n\n", "source": "/kaggle/working/cub/test/test_block_radix_rank.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_cdp_helper.h\"\n#include \"catch2_test_helper.h\"\n\n// %PARAM% TEST_CDP cdp 0:1\n\ntemplate \n__global__ void cub_api_example_x2_0_kernel(const T *d_in, T *d_out, int num_items)\n{\n const int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < num_items)\n {\n d_out[i] = d_in[i] * T{2};\n }\n}\n\ntemplate \n__global__ void cub_api_example_x0_5_kernel(const T *d_in, T *d_out, int num_items)\n{\n const int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < num_items)\n {\n d_out[i] = d_in[i] / T{2};\n }\n}\n\nstruct cub_api_example_t\n{\n static constexpr int threads_in_block = 256;\n\n template \n CUB_RUNTIME_FUNCTION static cudaError_t invoke(std::uint8_t *d_temp_storage,\n std::size_t &temp_storage_bytes,\n KernelT kernel,\n const T *d_in,\n T *d_out,\n int num_items,\n bool should_be_invoked_on_device)\n {\n NV_IF_TARGET(NV_IS_HOST,\n (if (should_be_invoked_on_device) { return cudaErrorLaunchFailure; }),\n (if (!should_be_invoked_on_device) { return cudaErrorLaunchFailure; }));\n\n if (d_temp_storage == nullptr)\n {\n temp_storage_bytes = static_cast(num_items);\n return cudaSuccess;\n }\n\n if (temp_storage_bytes != static_cast(num_items))\n {\n return cudaErrorInvalidValue;\n }\n\n const int blocks_in_grid = (num_items + threads_in_block - 1) / threads_in_block;\n\n return thrust::cuda_cub::launcher::triple_chevron(blocks_in_grid, threads_in_block, 0, 0)\n .doit(kernel, d_in, d_out, num_items);\n }\n\n template \n CUB_RUNTIME_FUNCTION static cudaError_t x2_0(std::uint8_t *d_temp_storage,\n std::size_t &temp_storage_bytes,\n const T *d_in,\n T *d_out,\n int num_items,\n bool should_be_invoked_on_device)\n {\n return invoke(d_temp_storage,\n temp_storage_bytes,\n cub_api_example_x2_0_kernel,\n d_in,\n d_out,\n num_items,\n should_be_invoked_on_device);\n }\n\n template \n CUB_RUNTIME_FUNCTION static cudaError_t x0_5(std::uint8_t *d_temp_storage,\n std::size_t &temp_storage_bytes,\n const T *d_in,\n T *d_out,\n int num_items,\n bool should_be_invoked_on_device)\n {\n return invoke(d_temp_storage,\n temp_storage_bytes,\n cub_api_example_x0_5_kernel,\n d_in,\n d_out,\n num_items,\n should_be_invoked_on_device);\n }\n};\n\nDECLARE_CDP_WRAPPER(cub_api_example_t::x2_0, x2_0);\nDECLARE_CDP_WRAPPER(cub_api_example_t::x0_5, x0_5);\n\nCUB_TEST(\"CDP wrapper works with predefined invocables\", \"[test][utils]\")\n{\n int n = 42;\n thrust::device_vector in(n, 21);\n thrust::device_vector out(n);\n\n int *d_in = thrust::raw_pointer_cast(in.data());\n int *d_out = thrust::raw_pointer_cast(out.data());\n\n constexpr bool should_be_invoked_on_device = TEST_CDP;\n\n {\n x2_0(d_in, d_out, n, should_be_invoked_on_device);\n\n const auto actual = static_cast(thrust::count(out.begin(), out.end(), 42));\n const auto expected = static_cast(n);\n\n REQUIRE(actual == expected);\n }\n\n {\n x0_5(d_out, d_out, n, should_be_invoked_on_device);\n\n const auto actual = static_cast(thrust::count(out.begin(), out.end(), 21));\n const auto expected = static_cast(n);\n\n REQUIRE(actual == expected);\n }\n}\n\nstruct custom_x2_0_invocable\n{\n template \n CUB_RUNTIME_FUNCTION cudaError_t operator()(std::uint8_t *d_temp_storage,\n std::size_t &temp_storage_bytes,\n const T *d_in,\n T *d_out,\n int num_items,\n bool should_be_invoked_on_device)\n {\n return cub_api_example_t::x2_0(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items,\n should_be_invoked_on_device);\n }\n};\n\nstruct custom_x0_5_invocable\n{\n template \n CUB_RUNTIME_FUNCTION cudaError_t operator()(std::uint8_t *d_temp_storage,\n std::size_t &temp_storage_bytes,\n const T *d_in,\n T *d_out,\n int num_items,\n bool should_be_invoked_on_device)\n {\n return cub_api_example_t::x0_5(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items,\n should_be_invoked_on_device);\n }\n};\n\nCUB_TEST(\"CDP wrapper works with custom invocables\", \"[test][utils]\")\n{\n int n = 42;\n thrust::device_vector in(n, 21);\n thrust::device_vector out(n);\n\n int *d_in = thrust::raw_pointer_cast(in.data());\n int *d_out = thrust::raw_pointer_cast(out.data());\n\n constexpr bool should_be_invoked_on_device = TEST_CDP;\n\n {\n cdp_launch(custom_x2_0_invocable{}, d_in, d_out, n, should_be_invoked_on_device);\n\n const auto actual = static_cast(thrust::count(out.begin(), out.end(), 42));\n const auto expected = static_cast(n);\n\n REQUIRE(actual == expected);\n }\n\n {\n cdp_launch(custom_x0_5_invocable{}, d_out, d_out, n, should_be_invoked_on_device);\n\n const auto actual = static_cast(thrust::count(out.begin(), out.end(), 21));\n const auto expected = static_cast(n);\n\n REQUIRE(actual == expected);\n }\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_cdp_helper.h\"\n#include \"catch2_test_helper.h\"\n\n// %PARAM% TEST_CDP cdp 0:1\n\ntemplate \n__global__ void cub_api_example_x2_0_kernel(const T *d_in, T *d_out, int num_items)\n{\n const int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < num_items)\n {\n d_out[i] = d_in[i] * T{2};\n }\n}\n\ntemplate \n__global__ void cub_api_example_x0_5_kernel(const T *d_in, T *d_out, int num_items)\n{\n const int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < num_items)\n {\n d_out[i] = d_in[i] / T{2};\n }\n}\n\nstruct cub_api_example_t\n{\n static constexpr int threads_in_block = 256;\n\n template \n HIPCUB_RUNTIME_FUNCTION static hipError_t invoke(std::uint8_t *d_temp_storage,\n std::size_t &temp_storage_bytes,\n KernelT kernel,\n const T *d_in,\n T *d_out,\n int num_items,\n bool should_be_invoked_on_device)\n {\n NV_IF_TARGET(NV_IS_HOST,\n (if (should_be_invoked_on_device) { return hipErrorLaunchFailure; }),\n (if (!should_be_invoked_on_device) { return hipErrorLaunchFailure; }));\n\n if (d_temp_storage == nullptr)\n {\n temp_storage_bytes = static_cast(num_items);\n return hipSuccess;\n }\n\n if (temp_storage_bytes != static_cast(num_items))\n {\n return hipErrorInvalidValue;\n }\n\n const int blocks_in_grid = (num_items + threads_in_block - 1) / threads_in_block;\n\n return thrust::cuda_cub::launcher::triple_chevron(blocks_in_grid, threads_in_block, 0, 0)\n .doit(kernel, d_in, d_out, num_items);\n }\n\n template \n HIPCUB_RUNTIME_FUNCTION static hipError_t x2_0(std::uint8_t *d_temp_storage,\n std::size_t &temp_storage_bytes,\n const T *d_in,\n T *d_out,\n int num_items,\n bool should_be_invoked_on_device)\n {\n return invoke(d_temp_storage,\n temp_storage_bytes,\n cub_api_example_x2_0_kernel,\n d_in,\n d_out,\n num_items,\n should_be_invoked_on_device);\n }\n\n template \n HIPCUB_RUNTIME_FUNCTION static hipError_t x0_5(std::uint8_t *d_temp_storage,\n std::size_t &temp_storage_bytes,\n const T *d_in,\n T *d_out,\n int num_items,\n bool should_be_invoked_on_device)\n {\n return invoke(d_temp_storage,\n temp_storage_bytes,\n cub_api_example_x0_5_kernel,\n d_in,\n d_out,\n num_items,\n should_be_invoked_on_device);\n }\n};\n\nDECLARE_CDP_WRAPPER(cub_api_example_t::x2_0, x2_0);\nDECLARE_CDP_WRAPPER(cub_api_example_t::x0_5, x0_5);\n\nCUB_TEST(\"CDP wrapper works with predefined invocables\", \"[test][utils]\")\n{\n int n = 42;\n thrust::device_vector in(n, 21);\n thrust::device_vector out(n);\n\n int *d_in = thrust::raw_pointer_cast(in.data());\n int *d_out = thrust::raw_pointer_cast(out.data());\n\n constexpr bool should_be_invoked_on_device = TEST_CDP;\n\n {\n x2_0(d_in, d_out, n, should_be_invoked_on_device);\n\n const auto actual = static_cast(thrust::count(out.begin(), out.end(), 42));\n const auto expected = static_cast(n);\n\n REQUIRE(actual == expected);\n }\n\n {\n x0_5(d_out, d_out, n, should_be_invoked_on_device);\n\n const auto actual = static_cast(thrust::count(out.begin(), out.end(), 21));\n const auto expected = static_cast(n);\n\n REQUIRE(actual == expected);\n }\n}\n\nstruct custom_x2_0_invocable\n{\n template \n HIPCUB_RUNTIME_FUNCTION hipError_t operator()(std::uint8_t *d_temp_storage,\n std::size_t &temp_storage_bytes,\n const T *d_in,\n T *d_out,\n int num_items,\n bool should_be_invoked_on_device)\n {\n return cub_api_example_t::x2_0(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items,\n should_be_invoked_on_device);\n }\n};\n\nstruct custom_x0_5_invocable\n{\n template \n HIPCUB_RUNTIME_FUNCTION hipError_t operator()(std::uint8_t *d_temp_storage,\n std::size_t &temp_storage_bytes,\n const T *d_in,\n T *d_out,\n int num_items,\n bool should_be_invoked_on_device)\n {\n return cub_api_example_t::x0_5(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items,\n should_be_invoked_on_device);\n }\n};\n\nCUB_TEST(\"CDP wrapper works with custom invocables\", \"[test][utils]\")\n{\n int n = 42;\n thrust::device_vector in(n, 21);\n thrust::device_vector out(n);\n\n int *d_in = thrust::raw_pointer_cast(in.data());\n int *d_out = thrust::raw_pointer_cast(out.data());\n\n constexpr bool should_be_invoked_on_device = TEST_CDP;\n\n {\n cdp_launch(custom_x2_0_invocable{}, d_in, d_out, n, should_be_invoked_on_device);\n\n const auto actual = static_cast(thrust::count(out.begin(), out.end(), 42));\n const auto expected = static_cast(n);\n\n REQUIRE(actual == expected);\n }\n\n {\n cdp_launch(custom_x0_5_invocable{}, d_out, d_out, n, should_be_invoked_on_device);\n\n const auto actual = static_cast(thrust::count(out.begin(), out.end(), 21));\n const auto expected = static_cast(n);\n\n REQUIRE(actual == expected);\n }\n}\n", "source": "/kaggle/working/cub/test/catch2_test_cdp_wrapper.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n#include \n\n#include \n#include \n#include \n\n#include \n#include \n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"c2h/custom_type.cuh\"\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void warp_reduce_kernel(T *in, T *out, ActionT action)\n{\n using warp_reduce_t = cub::WarpReduce;\n using storage_t = typename warp_reduce_t::TempStorage;\n\n __shared__ storage_t storage[TOTAL_WARPS];\n\n const int tid = threadIdx.x;\n\n // Get warp index\n int warp_id = tid / LOGICAL_WARP_THREADS;\n\n // Load data\n T thread_data = in[tid];\n // Instantiate and run warp reduction\n warp_reduce_t warp_reduce(storage[warp_id]);\n auto result = action(tid, warp_reduce, thread_data);\n\n // Write warp aggregate\n out[tid] = result;\n}\n\n/**\n * @brief Delegate wrapper for WarpReduce::Sum\n */\ntemplate \nstruct warp_sum_t\n{\n template \n __device__ T operator()(int linear_tid,\n cub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n auto result = warp_reduce.Sum(thread_data);\n return ((linear_tid % LOGICAL_WARP_THREADS) == 0) ? result : thread_data;\n }\n};\n\n/**\n * @brief Delegate wrapper for partial WarpReduce::Sum\n */\ntemplate \nstruct warp_sum_partial_t\n{\n int num_valid;\n template \n __device__ __forceinline__ T operator()(int linear_tid,\n cub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n auto result = warp_reduce.Sum(thread_data, num_valid);\n return ((linear_tid % LOGICAL_WARP_THREADS) == 0) ? result : thread_data;\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpReduce::Reduce\n */\ntemplate \nstruct warp_reduce_t\n{\n ReductionOpT reduction_op;\n template \n __device__ __forceinline__ T operator()(int linear_tid,\n cub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n auto result = warp_reduce.Reduce(thread_data, reduction_op);\n return ((linear_tid % LOGICAL_WARP_THREADS) == 0) ? result : thread_data;\n }\n};\n\n/**\n * @brief Delegate wrapper for partial WarpReduce::Reduce\n */\ntemplate \nstruct warp_reduce_partial_t\n{\n int num_valid;\n ReductionOpT reduction_op;\n template \n __device__ T operator()(int linear_tid,\n cub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n auto result = warp_reduce.Reduce(thread_data, reduction_op, num_valid);\n return ((linear_tid % LOGICAL_WARP_THREADS) == 0) ? result : thread_data;\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpReduce::TailSegmentedSum\n */\ntemplate \nstruct warp_seg_sum_tail_t\n{\n uint8_t *d_flags;\n template \n __device__ T operator()(int linear_tid,\n cub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n const bool has_agg = (linear_tid % LOGICAL_WARP_THREADS == 0) ||\n ((linear_tid == 0) ? 0 : d_flags[linear_tid - 1]);\n auto result = warp_reduce.TailSegmentedSum(thread_data, d_flags[linear_tid]);\n return has_agg ? result : thread_data;\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpReduce::HeadSegmentedSum\n */\ntemplate \nstruct warp_seg_sum_head_t\n{\n uint8_t *d_flags;\n template \n __device__ T operator()(int linear_tid,\n cub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n const bool has_agg = ((linear_tid % LOGICAL_WARP_THREADS == 0) || d_flags[linear_tid]);\n auto result = warp_reduce.HeadSegmentedSum(thread_data, d_flags[linear_tid]);\n return (has_agg) ? result : thread_data;\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpReduce::TailSegmentedReduce\n */\ntemplate \nstruct warp_seg_reduce_tail_t\n{\n uint8_t *d_flags;\n ReductionOpT reduction_op;\n template \n __device__ T operator()(int linear_tid,\n cub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n const bool has_agg = (linear_tid % LOGICAL_WARP_THREADS == 0) ||\n ((linear_tid == 0) ? 0 : d_flags[linear_tid - 1]);\n auto result = warp_reduce.TailSegmentedReduce(thread_data, d_flags[linear_tid], reduction_op);\n return has_agg ? result : thread_data;\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpReduce::HeadSegmentedReduce\n */\ntemplate \nstruct warp_seg_reduce_head_t\n{\n uint8_t *d_flags;\n ReductionOpT reduction_op;\n template \n __device__ T operator()(int linear_tid,\n cub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n const bool has_agg = ((linear_tid % LOGICAL_WARP_THREADS == 0) || d_flags[linear_tid]);\n auto result = warp_reduce.HeadSegmentedReduce(thread_data, d_flags[linear_tid], reduction_op);\n return (has_agg) ? result : thread_data;\n }\n};\n\n/**\n * @brief Dispatch helper function\n */\ntemplate \nvoid warp_reduce(thrust::device_vector &in, thrust::device_vector &out, ActionT action)\n{\n warp_reduce_kernel\n <<<1, LOGICAL_WARP_THREADS * TOTAL_WARPS>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(out.data()),\n action);\n\n REQUIRE(cudaSuccess == cudaPeekAtLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n}\n\n/**\n * @brief Compares the results returned from system under test against the expected results.\n */\ntemplate ::value, int>::type = 0>\nvoid verify_results(const thrust::host_vector &expected_data,\n const thrust::device_vector &test_results)\n{\n REQUIRE_APPROX_EQ(expected_data, test_results);\n}\n\n/**\n * @brief Compares the results returned from system under test against the expected results.\n */\ntemplate ::value, int>::type = 0>\nvoid verify_results(const thrust::host_vector &expected_data,\n const thrust::device_vector &test_results)\n{\n REQUIRE(expected_data == test_results);\n}\n\nenum class reduce_mode\n{\n all,\n partial,\n head_flags,\n tail_flags,\n};\n\ntemplate \nvoid compute_host_reference(reduce_mode mode,\n InputItT h_in,\n FlagInputItT h_flags,\n int warps,\n int logical_warp_threads,\n int valid_warp_threads,\n ReductionOp reduction_op,\n ResultOutItT h_data_out)\n{\n // Accumulate segments (lane 0 of each warp is implicitly a segment head)\n for (int warp = 0; warp < warps; ++warp)\n {\n int warp_offset = warp * logical_warp_threads;\n int item_offset = warp_offset + valid_warp_threads - 1;\n\n // Last item in warp\n auto head_aggregate = h_in[item_offset];\n auto tail_aggregate = h_in[item_offset];\n\n if (mode != reduce_mode::tail_flags && h_flags[item_offset])\n {\n h_data_out[item_offset] = head_aggregate;\n }\n item_offset--;\n\n // Work backwards\n while (item_offset >= warp_offset)\n {\n if (h_flags[item_offset + 1])\n {\n head_aggregate = h_in[item_offset];\n }\n else\n {\n head_aggregate = reduction_op(head_aggregate, h_in[item_offset]);\n }\n\n if (h_flags[item_offset])\n {\n if (mode == reduce_mode::head_flags)\n {\n h_data_out[item_offset] = head_aggregate;\n }\n else if (mode == reduce_mode::tail_flags)\n {\n h_data_out[item_offset + 1] = tail_aggregate;\n tail_aggregate = h_in[item_offset];\n }\n }\n else\n {\n tail_aggregate = reduction_op(tail_aggregate, h_in[item_offset]);\n }\n\n item_offset--;\n }\n\n // Record last segment aggregate\n if (mode == reduce_mode::tail_flags)\n {\n h_data_out[warp_offset] = tail_aggregate;\n }\n else\n {\n h_data_out[warp_offset] = head_aggregate;\n }\n }\n}\n\n// List of types to test\nusing custom_t = c2h::custom_type_t;\nusing full_type_list = c2h::type_list;\n\nusing builtin_type_list = c2h::type_list;\n\n// Logical warp sizes to test\nusing logical_warp_threads = c2h::enum_type_list;\n\nusing segmented_modes =\n c2h::enum_type_list;\n\ntemplate \nstruct total_warps_t\n{\nprivate:\n static constexpr int max_warps = 2;\n static constexpr bool is_arch_warp = (logical_warp_threads == CUB_WARP_THREADS(0));\n static constexpr bool is_pow_of_two = ((logical_warp_threads & (logical_warp_threads - 1)) == 0);\n static constexpr int total_warps = (is_arch_warp || is_pow_of_two) ? max_warps : 1;\n\npublic:\n static constexpr int value() { return total_warps; }\n};\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int logical_warp_threads = c2h::get<1, TestType>::value;\n static constexpr int total_warps = total_warps_t::value();\n static constexpr int tile_size = total_warps * logical_warp_threads;\n};\n\nCUB_TEST(\"Warp sum works\", \"[reduce][warp]\", full_type_list, logical_warp_threads)\n{\n using params = params_t;\n using type = typename params::type;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n const auto valid_items = params::logical_warp_threads;\n c2h::gen(CUB_SEED(10), d_in);\n\n // Run test\n warp_reduce(d_in, d_out, warp_sum_t{});\n\n // Prepare verification data\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_out = h_in;\n auto h_flags = thrust::make_constant_iterator(false);\n compute_host_reference(reduce_mode::all,\n h_in,\n h_flags,\n params::total_warps,\n params::logical_warp_threads,\n valid_items,\n ::cuda::std::plus{},\n h_out.begin());\n\n // Verify results\n verify_results(h_out, d_out);\n}\n\nCUB_TEST(\"Warp reduce works\", \"[reduce][warp]\", builtin_type_list, logical_warp_threads)\n{\n using params = params_t;\n using type = typename params::type;\n using red_op_t = cub::Min;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n const auto valid_items = params::logical_warp_threads;\n c2h::gen(CUB_SEED(10), d_in);\n\n // Run test\n warp_reduce(d_in,\n d_out,\n warp_reduce_t{\n red_op_t{}});\n\n // Prepare verification data\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_out = h_in;\n auto h_flags = thrust::make_constant_iterator(false);\n compute_host_reference(reduce_mode::all,\n h_in,\n h_flags,\n params::total_warps,\n params::logical_warp_threads,\n valid_items,\n red_op_t{},\n h_out.begin());\n\n // Verify results\n verify_results(h_out, d_out);\n}\n\nCUB_TEST(\"Warp sum on partial warp works\", \"[reduce][warp]\", full_type_list, logical_warp_threads)\n{\n using params = params_t;\n using type = typename params::type;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n const int valid_items = GENERATE_COPY(take(2, random(1, params::logical_warp_threads)));\n c2h::gen(CUB_SEED(10), d_in);\n\n // Run test\n warp_reduce(d_in,\n d_out,\n warp_sum_partial_t{\n valid_items});\n\n // Prepare verification data\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_out = h_in;\n auto h_flags = thrust::make_constant_iterator(false);\n compute_host_reference(reduce_mode::all,\n h_in,\n h_flags,\n params::total_warps,\n params::logical_warp_threads,\n valid_items,\n ::cuda::std::plus{},\n h_out.begin());\n\n // Verify results\n verify_results(h_out, d_out);\n}\n\nCUB_TEST(\"Warp reduce on partial warp works\",\n \"[reduce][warp]\",\n builtin_type_list,\n logical_warp_threads)\n{\n using params = params_t;\n using type = typename params::type;\n using red_op_t = cub::Min;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n const int valid_items = GENERATE_COPY(take(2, random(1, params::logical_warp_threads)));\n c2h::gen(CUB_SEED(10), d_in);\n\n // Run test\n warp_reduce(\n d_in,\n d_out,\n warp_reduce_partial_t{valid_items, red_op_t{}});\n\n // Prepare verification data\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_out = h_in;\n auto h_flags = thrust::make_constant_iterator(false);\n compute_host_reference(reduce_mode::all,\n h_in,\n h_flags,\n params::total_warps,\n params::logical_warp_threads,\n valid_items,\n red_op_t{},\n h_out.begin());\n\n // Verify results\n verify_results(h_out, d_out);\n}\n\nCUB_TEST(\"Warp segmented sum works\",\n \"[reduce][warp]\",\n full_type_list,\n logical_warp_threads,\n segmented_modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n constexpr auto segmented_mod = c2h::get<2, TestType>::value;\n static_assert(segmented_mod == reduce_mode::tail_flags ||\n segmented_mod == reduce_mode::head_flags,\n \"Segmented tests must either be head or tail flags\");\n using warp_seg_sum_t = cub::detail::conditional_t<(segmented_mod == reduce_mode::tail_flags),\n warp_seg_sum_tail_t,\n warp_seg_sum_head_t>;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_flags(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n const auto valid_items = params::logical_warp_threads;\n constexpr uint8_t min = 0;\n constexpr uint8_t max = 2;\n c2h::gen(CUB_SEED(5), d_in);\n c2h::gen(CUB_SEED(5), d_flags, min, max);\n\n // Run test\n warp_reduce(\n d_in,\n d_out,\n warp_seg_sum_t{thrust::raw_pointer_cast(d_flags.data())});\n\n // Prepare verification data\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_flags = d_flags;\n thrust::host_vector h_out = h_in;\n compute_host_reference(segmented_mod,\n h_in,\n h_flags,\n params::total_warps,\n params::logical_warp_threads,\n valid_items,\n ::cuda::std::plus{},\n h_out.begin());\n\n // Verify results\n verify_results(h_out, d_out);\n}\n\nCUB_TEST(\"Warp segmented reduction works\",\n \"[reduce][warp]\",\n builtin_type_list,\n logical_warp_threads,\n segmented_modes)\n{\n using params = params_t;\n using type = typename params::type;\n using red_op_t = cub::Min;\n\n constexpr auto segmented_mod = c2h::get<2, TestType>::value;\n static_assert(segmented_mod == reduce_mode::tail_flags ||\n segmented_mod == reduce_mode::head_flags,\n \"Segmented tests must either be head or tail flags\");\n using warp_seg_reduction_t =\n cub::detail::conditional_t<(segmented_mod == reduce_mode::tail_flags),\n warp_seg_reduce_tail_t,\n warp_seg_reduce_head_t>;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_flags(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n const auto valid_items = params::logical_warp_threads;\n constexpr uint8_t min = 0;\n constexpr uint8_t max = 2;\n c2h::gen(CUB_SEED(5), d_in);\n c2h::gen(CUB_SEED(5), d_flags, min, max);\n\n // Run test\n warp_reduce(\n d_in,\n d_out,\n warp_seg_reduction_t{thrust::raw_pointer_cast(d_flags.data()), red_op_t{}});\n\n // Prepare verification data\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_flags = d_flags;\n thrust::host_vector h_out = h_in;\n compute_host_reference(segmented_mod,\n h_in,\n h_flags,\n params::total_warps,\n params::logical_warp_threads,\n valid_items,\n red_op_t{},\n h_out.begin());\n\n // Verify results\n verify_results(h_out, d_out);\n}", "output": "/******************************************************************************\n * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n#include \n#include \n#include \n\n#include \n#include \n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"c2h/custom_type.cuh\"\n#include \"catch2_test_helper.h\"\n\ntemplate \n__global__ void warp_reduce_kernel(T *in, T *out, ActionT action)\n{\n using warp_reduce_t = hipcub::WarpReduce;\n using storage_t = typename warp_reduce_t::TempStorage;\n\n __shared__ storage_t storage[TOTAL_WARPS];\n\n const int tid = threadIdx.x;\n\n // Get warp index\n int warp_id = tid / LOGICAL_WARP_THREADS;\n\n // Load data\n T thread_data = in[tid];\n // Instantiate and run warp reduction\n warp_reduce_t warp_reduce(storage[warp_id]);\n auto result = action(tid, warp_reduce, thread_data);\n\n // Write warp aggregate\n out[tid] = result;\n}\n\n/**\n * @brief Delegate wrapper for WarpReduce::Sum\n */\ntemplate \nstruct warp_sum_t\n{\n template \n __device__ T operator()(int linear_tid,\n hipcub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n auto result = warp_reduce.Sum(thread_data);\n return ((linear_tid % LOGICAL_WARP_THREADS) == 0) ? result : thread_data;\n }\n};\n\n/**\n * @brief Delegate wrapper for partial WarpReduce::Sum\n */\ntemplate \nstruct warp_sum_partial_t\n{\n int num_valid;\n template \n __device__ __forceinline__ T operator()(int linear_tid,\n hipcub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n auto result = warp_reduce.Sum(thread_data, num_valid);\n return ((linear_tid % LOGICAL_WARP_THREADS) == 0) ? result : thread_data;\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpReduce::Reduce\n */\ntemplate \nstruct warp_reduce_t\n{\n ReductionOpT reduction_op;\n template \n __device__ __forceinline__ T operator()(int linear_tid,\n hipcub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n auto result = warp_reduce.Reduce(thread_data, reduction_op);\n return ((linear_tid % LOGICAL_WARP_THREADS) == 0) ? result : thread_data;\n }\n};\n\n/**\n * @brief Delegate wrapper for partial WarpReduce::Reduce\n */\ntemplate \nstruct warp_reduce_partial_t\n{\n int num_valid;\n ReductionOpT reduction_op;\n template \n __device__ T operator()(int linear_tid,\n hipcub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n auto result = warp_reduce.Reduce(thread_data, reduction_op, num_valid);\n return ((linear_tid % LOGICAL_WARP_THREADS) == 0) ? result : thread_data;\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpReduce::TailSegmentedSum\n */\ntemplate \nstruct warp_seg_sum_tail_t\n{\n uint8_t *d_flags;\n template \n __device__ T operator()(int linear_tid,\n hipcub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n const bool has_agg = (linear_tid % LOGICAL_WARP_THREADS == 0) ||\n ((linear_tid == 0) ? 0 : d_flags[linear_tid - 1]);\n auto result = warp_reduce.TailSegmentedSum(thread_data, d_flags[linear_tid]);\n return has_agg ? result : thread_data;\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpReduce::HeadSegmentedSum\n */\ntemplate \nstruct warp_seg_sum_head_t\n{\n uint8_t *d_flags;\n template \n __device__ T operator()(int linear_tid,\n hipcub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n const bool has_agg = ((linear_tid % LOGICAL_WARP_THREADS == 0) || d_flags[linear_tid]);\n auto result = warp_reduce.HeadSegmentedSum(thread_data, d_flags[linear_tid]);\n return (has_agg) ? result : thread_data;\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpReduce::TailSegmentedReduce\n */\ntemplate \nstruct warp_seg_reduce_tail_t\n{\n uint8_t *d_flags;\n ReductionOpT reduction_op;\n template \n __device__ T operator()(int linear_tid,\n hipcub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n const bool has_agg = (linear_tid % LOGICAL_WARP_THREADS == 0) ||\n ((linear_tid == 0) ? 0 : d_flags[linear_tid - 1]);\n auto result = warp_reduce.TailSegmentedReduce(thread_data, d_flags[linear_tid], reduction_op);\n return has_agg ? result : thread_data;\n }\n};\n\n/**\n * @brief Delegate wrapper for WarpReduce::HeadSegmentedReduce\n */\ntemplate \nstruct warp_seg_reduce_head_t\n{\n uint8_t *d_flags;\n ReductionOpT reduction_op;\n template \n __device__ T operator()(int linear_tid,\n hipcub::WarpReduce &warp_reduce,\n T &thread_data) const\n {\n const bool has_agg = ((linear_tid % LOGICAL_WARP_THREADS == 0) || d_flags[linear_tid]);\n auto result = warp_reduce.HeadSegmentedReduce(thread_data, d_flags[linear_tid], reduction_op);\n return (has_agg) ? result : thread_data;\n }\n};\n\n/**\n * @brief Dispatch helper function\n */\ntemplate \nvoid warp_reduce(thrust::device_vector &in, thrust::device_vector &out, ActionT action)\n{\n warp_reduce_kernel\n <<<1, LOGICAL_WARP_THREADS * TOTAL_WARPS>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(out.data()),\n action);\n\n REQUIRE(hipSuccess == hipPeekAtLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n}\n\n/**\n * @brief Compares the results returned from system under test against the expected results.\n */\ntemplate ::value, int>::type = 0>\nvoid verify_results(const thrust::host_vector &expected_data,\n const thrust::device_vector &test_results)\n{\n REQUIRE_APPROX_EQ(expected_data, test_results);\n}\n\n/**\n * @brief Compares the results returned from system under test against the expected results.\n */\ntemplate ::value, int>::type = 0>\nvoid verify_results(const thrust::host_vector &expected_data,\n const thrust::device_vector &test_results)\n{\n REQUIRE(expected_data == test_results);\n}\n\nenum class reduce_mode\n{\n all,\n partial,\n head_flags,\n tail_flags,\n};\n\ntemplate \nvoid compute_host_reference(reduce_mode mode,\n InputItT h_in,\n FlagInputItT h_flags,\n int warps,\n int logical_warp_threads,\n int valid_warp_threads,\n ReductionOp reduction_op,\n ResultOutItT h_data_out)\n{\n // Accumulate segments (lane 0 of each warp is implicitly a segment head)\n for (int warp = 0; warp < warps; ++warp)\n {\n int warp_offset = warp * logical_warp_threads;\n int item_offset = warp_offset + valid_warp_threads - 1;\n\n // Last item in warp\n auto head_aggregate = h_in[item_offset];\n auto tail_aggregate = h_in[item_offset];\n\n if (mode != reduce_mode::tail_flags && h_flags[item_offset])\n {\n h_data_out[item_offset] = head_aggregate;\n }\n item_offset--;\n\n // Work backwards\n while (item_offset >= warp_offset)\n {\n if (h_flags[item_offset + 1])\n {\n head_aggregate = h_in[item_offset];\n }\n else\n {\n head_aggregate = reduction_op(head_aggregate, h_in[item_offset]);\n }\n\n if (h_flags[item_offset])\n {\n if (mode == reduce_mode::head_flags)\n {\n h_data_out[item_offset] = head_aggregate;\n }\n else if (mode == reduce_mode::tail_flags)\n {\n h_data_out[item_offset + 1] = tail_aggregate;\n tail_aggregate = h_in[item_offset];\n }\n }\n else\n {\n tail_aggregate = reduction_op(tail_aggregate, h_in[item_offset]);\n }\n\n item_offset--;\n }\n\n // Record last segment aggregate\n if (mode == reduce_mode::tail_flags)\n {\n h_data_out[warp_offset] = tail_aggregate;\n }\n else\n {\n h_data_out[warp_offset] = head_aggregate;\n }\n }\n}\n\n// List of types to test\nusing custom_t = c2h::custom_type_t;\nusing full_type_list = c2h::type_list;\n\nusing builtin_type_list = c2h::type_list;\n\n// Logical warp sizes to test\nusing logical_warp_threads = c2h::enum_type_list;\n\nusing segmented_modes =\n c2h::enum_type_list;\n\ntemplate \nstruct total_warps_t\n{\nprivate:\n static constexpr int max_warps = 2;\n static constexpr bool is_arch_warp = (logical_warp_threads == CUB_WARP_THREADS(0));\n static constexpr bool is_pow_of_two = ((logical_warp_threads & (logical_warp_threads - 1)) == 0);\n static constexpr int total_warps = (is_arch_warp || is_pow_of_two) ? max_warps : 1;\n\npublic:\n static constexpr int value() { return total_warps; }\n};\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int logical_warp_threads = c2h::get<1, TestType>::value;\n static constexpr int total_warps = total_warps_t::value();\n static constexpr int tile_size = total_warps * logical_warp_threads;\n};\n\nCUB_TEST(\"Warp sum works\", \"[reduce][warp]\", full_type_list, logical_warp_threads)\n{\n using params = params_t;\n using type = typename params::type;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n const auto valid_items = params::logical_warp_threads;\n c2h::gen(CUB_SEED(10), d_in);\n\n // Run test\n warp_reduce(d_in, d_out, warp_sum_t{});\n\n // Prepare verification data\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_out = h_in;\n auto h_flags = thrust::make_constant_iterator(false);\n compute_host_reference(reduce_mode::all,\n h_in,\n h_flags,\n params::total_warps,\n params::logical_warp_threads,\n valid_items,\n ::cuda::std::plus{},\n h_out.begin());\n\n // Verify results\n verify_results(h_out, d_out);\n}\n\nCUB_TEST(\"Warp reduce works\", \"[reduce][warp]\", builtin_type_list, logical_warp_threads)\n{\n using params = params_t;\n using type = typename params::type;\n using red_op_t = hipcub::Min;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n const auto valid_items = params::logical_warp_threads;\n c2h::gen(CUB_SEED(10), d_in);\n\n // Run test\n warp_reduce(d_in,\n d_out,\n warp_reduce_t{\n red_op_t{}});\n\n // Prepare verification data\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_out = h_in;\n auto h_flags = thrust::make_constant_iterator(false);\n compute_host_reference(reduce_mode::all,\n h_in,\n h_flags,\n params::total_warps,\n params::logical_warp_threads,\n valid_items,\n red_op_t{},\n h_out.begin());\n\n // Verify results\n verify_results(h_out, d_out);\n}\n\nCUB_TEST(\"Warp sum on partial warp works\", \"[reduce][warp]\", full_type_list, logical_warp_threads)\n{\n using params = params_t;\n using type = typename params::type;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n const int valid_items = GENERATE_COPY(take(2, random(1, params::logical_warp_threads)));\n c2h::gen(CUB_SEED(10), d_in);\n\n // Run test\n warp_reduce(d_in,\n d_out,\n warp_sum_partial_t{\n valid_items});\n\n // Prepare verification data\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_out = h_in;\n auto h_flags = thrust::make_constant_iterator(false);\n compute_host_reference(reduce_mode::all,\n h_in,\n h_flags,\n params::total_warps,\n params::logical_warp_threads,\n valid_items,\n ::cuda::std::plus{},\n h_out.begin());\n\n // Verify results\n verify_results(h_out, d_out);\n}\n\nCUB_TEST(\"Warp reduce on partial warp works\",\n \"[reduce][warp]\",\n builtin_type_list,\n logical_warp_threads)\n{\n using params = params_t;\n using type = typename params::type;\n using red_op_t = hipcub::Min;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n const int valid_items = GENERATE_COPY(take(2, random(1, params::logical_warp_threads)));\n c2h::gen(CUB_SEED(10), d_in);\n\n // Run test\n warp_reduce(\n d_in,\n d_out,\n warp_reduce_partial_t{valid_items, red_op_t{}});\n\n // Prepare verification data\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_out = h_in;\n auto h_flags = thrust::make_constant_iterator(false);\n compute_host_reference(reduce_mode::all,\n h_in,\n h_flags,\n params::total_warps,\n params::logical_warp_threads,\n valid_items,\n red_op_t{},\n h_out.begin());\n\n // Verify results\n verify_results(h_out, d_out);\n}\n\nCUB_TEST(\"Warp segmented sum works\",\n \"[reduce][warp]\",\n full_type_list,\n logical_warp_threads,\n segmented_modes)\n{\n using params = params_t;\n using type = typename params::type;\n\n constexpr auto segmented_mod = c2h::get<2, TestType>::value;\n static_assert(segmented_mod == reduce_mode::tail_flags ||\n segmented_mod == reduce_mode::head_flags,\n \"Segmented tests must either be head or tail flags\");\n using warp_seg_sum_t = hipcub::detail::conditional_t<(segmented_mod == reduce_mode::tail_flags),\n warp_seg_sum_tail_t,\n warp_seg_sum_head_t>;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_flags(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n const auto valid_items = params::logical_warp_threads;\n constexpr uint8_t min = 0;\n constexpr uint8_t max = 2;\n c2h::gen(CUB_SEED(5), d_in);\n c2h::gen(CUB_SEED(5), d_flags, min, max);\n\n // Run test\n warp_reduce(\n d_in,\n d_out,\n warp_seg_sum_t{thrust::raw_pointer_cast(d_flags.data())});\n\n // Prepare verification data\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_flags = d_flags;\n thrust::host_vector h_out = h_in;\n compute_host_reference(segmented_mod,\n h_in,\n h_flags,\n params::total_warps,\n params::logical_warp_threads,\n valid_items,\n ::cuda::std::plus{},\n h_out.begin());\n\n // Verify results\n verify_results(h_out, d_out);\n}\n\nCUB_TEST(\"Warp segmented reduction works\",\n \"[reduce][warp]\",\n builtin_type_list,\n logical_warp_threads,\n segmented_modes)\n{\n using params = params_t;\n using type = typename params::type;\n using red_op_t = hipcub::Min;\n\n constexpr auto segmented_mod = c2h::get<2, TestType>::value;\n static_assert(segmented_mod == reduce_mode::tail_flags ||\n segmented_mod == reduce_mode::head_flags,\n \"Segmented tests must either be head or tail flags\");\n using warp_seg_reduction_t =\n hipcub::detail::conditional_t<(segmented_mod == reduce_mode::tail_flags),\n warp_seg_reduce_tail_t,\n warp_seg_reduce_head_t>;\n\n // Prepare test data\n thrust::device_vector d_in(params::tile_size);\n thrust::device_vector d_flags(params::tile_size);\n thrust::device_vector d_out(params::tile_size);\n const auto valid_items = params::logical_warp_threads;\n constexpr uint8_t min = 0;\n constexpr uint8_t max = 2;\n c2h::gen(CUB_SEED(5), d_in);\n c2h::gen(CUB_SEED(5), d_flags, min, max);\n\n // Run test\n warp_reduce(\n d_in,\n d_out,\n warp_seg_reduction_t{thrust::raw_pointer_cast(d_flags.data()), red_op_t{}});\n\n // Prepare verification data\n thrust::host_vector h_in = d_in;\n thrust::host_vector h_flags = d_flags;\n thrust::host_vector h_out = h_in;\n compute_host_reference(segmented_mod,\n h_in,\n h_flags,\n params::total_warps,\n params::logical_warp_threads,\n valid_items,\n red_op_t{},\n h_out.begin());\n\n // Verify results\n verify_results(h_out, d_out);\n}", "source": "/kaggle/working/cub/test/catch2_test_warp_reduce.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \n\n#include \"fill_striped.cuh\"\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \nstruct exchange_data_t;\n\ntemplate \nstruct exchange_data_t::value>::type>\n{\n InputT input[ItemsPerThread];\n OutputT (&output)[ItemsPerThread] = input;\n\n template \n inline __device__ void\n scatter(cub::WarpExchange &exchange,\n int (&ranks)[ItemsPerThread])\n {\n exchange.ScatterToStriped(input, ranks);\n }\n};\n\ntemplate \nstruct exchange_data_t::value>::type>\n{\n InputT input[ItemsPerThread];\n OutputT output[ItemsPerThread];\n\n template \n inline __device__ void\n scatter(cub::WarpExchange &exchange,\n int (&ranks)[ItemsPerThread])\n {\n exchange.ScatterToStriped(input, output, ranks);\n }\n};\n\ntemplate \n__global__ void scatter_kernel(const InputT *input_data, OutputT *output_data)\n{\n using warp_exchange_t = cub::WarpExchange;\n using storage_t = typename warp_exchange_t::TempStorage;\n\n constexpr int tile_size = ITEMS_PER_THREAD * LOGICAL_WARP_THREADS;\n __shared__ storage_t temp_storage[TOTAL_WARPS];\n\n const int tid = cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n\n // Get warp index\n const int warp_id = tid / LOGICAL_WARP_THREADS;\n const int lane_id = tid % LOGICAL_WARP_THREADS;\n\n warp_exchange_t exchange(temp_storage[warp_id]);\n\n exchange_data_t exchange_data;\n\n // Reverse data\n int ranks[ITEMS_PER_THREAD];\n\n input_data += warp_id * tile_size;\n output_data += warp_id * tile_size;\n\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n const auto item_idx = lane_id * ITEMS_PER_THREAD + item;\n exchange_data.input[item] = input_data[item_idx];\n ranks[item] = tile_size - 1 - item_idx;\n }\n\n exchange_data.scatter(exchange, ranks);\n\n // Striped to blocked\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n output_data[item * LOGICAL_WARP_THREADS + lane_id] = exchange_data.output[item];\n }\n}\n\ntemplate \nvoid warp_scatter_strided(thrust::device_vector &in, thrust::device_vector &out)\n{\n scatter_kernel\n <<<1, LOGICAL_WARP_THREADS * TOTAL_WARPS>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(out.data()));\n REQUIRE(cudaSuccess == cudaPeekAtLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n}\n\ntemplate \n__global__ void kernel(const InputT *input_data, OutputT *output_data, ActionT action)\n{\n using warp_exchange_t = cub::WarpExchange;\n using storage_t = typename warp_exchange_t::TempStorage;\n\n constexpr int tile_size = ITEMS_PER_THREAD * LOGICAL_WARP_THREADS;\n __shared__ storage_t temp_storage[TOTAL_WARPS];\n\n const int tid = cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n\n // Get warp index\n const int warp_id = tid / LOGICAL_WARP_THREADS;\n const int lane_id = tid % LOGICAL_WARP_THREADS;\n\n warp_exchange_t exchange(temp_storage[warp_id]);\n\n exchange_data_t exchange_data;\n\n input_data += warp_id * tile_size;\n output_data += warp_id * tile_size;\n\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n exchange_data.input[item] = input_data[lane_id * ITEMS_PER_THREAD + item];\n }\n\n action(exchange_data.input, exchange_data.output, exchange);\n\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n output_data[lane_id * ITEMS_PER_THREAD + item] = exchange_data.output[item];\n }\n}\n\ntemplate \nvoid warp_exchange(thrust::device_vector &in,\n thrust::device_vector &out,\n ActionT action)\n{\n kernel\n <<<1, LOGICAL_WARP_THREADS * TOTAL_WARPS>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(out.data()),\n action);\n REQUIRE(cudaSuccess == cudaPeekAtLastError());\n REQUIRE(cudaSuccess == cudaDeviceSynchronize());\n}\n\nstruct blocked_to_striped\n{\n template \n __device__ void operator()(InputT (&input)[ITEMS_PER_THREAD],\n OutputT (&output)[ITEMS_PER_THREAD],\n cub::WarpExchange &exchange)\n {\n exchange.BlockedToStriped(input, output);\n }\n};\n\nstruct striped_to_blocked\n{\n template \n __device__ void operator()(InputT (&input)[ITEMS_PER_THREAD],\n OutputT (&output)[ITEMS_PER_THREAD],\n cub::WarpExchange &exchange)\n {\n exchange.StripedToBlocked(input, output);\n }\n};\n\ntemplate \nthrust::host_vector compute_host_reference(const thrust::device_vector &d_input,\n int tile_size)\n{\n thrust::host_vector input = d_input;\n\n int num_warps = CUB_QUOTIENT_CEILING(static_cast(d_input.size()), tile_size);\n for (int warp_id = 0; warp_id < num_warps; warp_id++)\n {\n const int warp_data_begin = tile_size * warp_id;\n const int warp_data_end = warp_data_begin + tile_size;\n thrust::reverse(input.begin() + warp_data_begin, input.begin() + warp_data_end);\n }\n return input;\n}\n\nusing inout_types = c2h::type_list,\n c2h::pair,\n c2h::pair,\n c2h::pair,\n c2h::pair,\n c2h::pair>;\n\nusing logical_warp_threads = c2h::enum_type_list;\nusing items_per_thread = c2h::enum_type_list;\n\ntemplate \nstruct total_warps_t\n{\nprivate:\n static constexpr int max_warps = 2;\n static constexpr bool is_arch_warp = (logical_warp_threads == CUB_WARP_THREADS(0));\n static constexpr bool is_pow_of_two = ((logical_warp_threads & (logical_warp_threads - 1)) == 0);\n static constexpr int total_warps = (is_arch_warp || is_pow_of_two) ? max_warps : 1;\n\npublic:\n static constexpr int value() { return total_warps; }\n};\n\ntemplate \nstruct params_t\n{\n using in_type = typename c2h::first>;\n using out_type = typename c2h::second>;\n\n static constexpr int logical_warp_threads = c2h::get<1, TestType>::value;\n static constexpr int items_per_thread = c2h::get<2, TestType>::value;\n static constexpr int total_warps = total_warps_t::value();\n static constexpr int tile_size = logical_warp_threads * items_per_thread;\n static constexpr int total_item_count = total_warps * tile_size;\n};\n\nCUB_TEST(\"Scatter to striped works\",\n \"[exchange][warp]\",\n inout_types,\n logical_warp_threads,\n items_per_thread)\n{\n using params = params_t;\n using in_type = typename params::in_type;\n using out_type = typename params::out_type;\n thrust::device_vector d_out(params::total_item_count);\n thrust::device_vector d_in(params::total_item_count);\n\n c2h::gen(c2h::modulo_t{d_in.size()}, d_in);\n\n warp_scatter_strided(\n d_in,\n d_out);\n\n auto h_expected_output = compute_host_reference(d_in, params::tile_size);\n REQUIRE(h_expected_output == d_out);\n}\n\nCUB_TEST(\"Blocked to striped works\",\n \"[exchange][warp]\",\n inout_types,\n logical_warp_threads,\n items_per_thread)\n{\n using params = params_t;\n using in_type = typename params::in_type;\n using out_type = typename params::out_type;\n thrust::device_vector d_out(params::total_item_count, out_type{});\n thrust::device_vector d_in(params::total_item_count);\n\n c2h::gen(c2h::modulo_t{d_in.size()}, d_in);\n\n warp_exchange(\n d_in,\n d_out,\n blocked_to_striped{});\n thrust::host_vector h_expected_output(d_out.size());\n fill_striped(h_expected_output.begin());\n\n REQUIRE(h_expected_output == d_out);\n}\n\nCUB_TEST(\"Striped to blocked works\",\n \"[exchange][warp]\",\n inout_types,\n logical_warp_threads,\n items_per_thread)\n{\n using params = params_t;\n using in_type = typename params::in_type;\n using out_type = typename params::out_type;\n thrust::device_vector d_out(params::total_item_count, out_type{});\n\n thrust::host_vector h_in(params::total_item_count);\n fill_striped(h_in.begin());\n thrust::device_vector d_in = h_in;\n\n warp_exchange(\n d_in,\n d_out,\n striped_to_blocked{});\n thrust::device_vector d_expected_output(d_out.size());\n c2h::gen(c2h::modulo_t{d_out.size()}, d_expected_output);\n\n REQUIRE(d_expected_output == d_out);\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \n\n#include \"fill_striped.cuh\"\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \nstruct exchange_data_t;\n\ntemplate \nstruct exchange_data_t::value>::type>\n{\n InputT input[ItemsPerThread];\n OutputT (&output)[ItemsPerThread] = input;\n\n template \n inline __device__ void\n scatter(hipcub::WarpExchange &exchange,\n int (&ranks)[ItemsPerThread])\n {\n exchange.ScatterToStriped(input, ranks);\n }\n};\n\ntemplate \nstruct exchange_data_t::value>::type>\n{\n InputT input[ItemsPerThread];\n OutputT output[ItemsPerThread];\n\n template \n inline __device__ void\n scatter(hipcub::WarpExchange &exchange,\n int (&ranks)[ItemsPerThread])\n {\n exchange.ScatterToStriped(input, output, ranks);\n }\n};\n\ntemplate \n__global__ void scatter_kernel(const InputT *input_data, OutputT *output_data)\n{\n using warp_exchange_t = hipcub::WarpExchange;\n using storage_t = typename warp_exchange_t::TempStorage;\n\n constexpr int tile_size = ITEMS_PER_THREAD * LOGICAL_WARP_THREADS;\n __shared__ storage_t temp_storage[TOTAL_WARPS];\n\n const int tid = hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n\n // Get warp index\n const int warp_id = tid / LOGICAL_WARP_THREADS;\n const int lane_id = tid % LOGICAL_WARP_THREADS;\n\n warp_exchange_t exchange(temp_storage[warp_id]);\n\n exchange_data_t exchange_data;\n\n // Reverse data\n int ranks[ITEMS_PER_THREAD];\n\n input_data += warp_id * tile_size;\n output_data += warp_id * tile_size;\n\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n const auto item_idx = lane_id * ITEMS_PER_THREAD + item;\n exchange_data.input[item] = input_data[item_idx];\n ranks[item] = tile_size - 1 - item_idx;\n }\n\n exchange_data.scatter(exchange, ranks);\n\n // Striped to blocked\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n output_data[item * LOGICAL_WARP_THREADS + lane_id] = exchange_data.output[item];\n }\n}\n\ntemplate \nvoid warp_scatter_strided(thrust::device_vector &in, thrust::device_vector &out)\n{\n scatter_kernel\n <<<1, LOGICAL_WARP_THREADS * TOTAL_WARPS>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(out.data()));\n REQUIRE(hipSuccess == hipPeekAtLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n}\n\ntemplate \n__global__ void kernel(const InputT *input_data, OutputT *output_data, ActionT action)\n{\n using warp_exchange_t = hipcub::WarpExchange;\n using storage_t = typename warp_exchange_t::TempStorage;\n\n constexpr int tile_size = ITEMS_PER_THREAD * LOGICAL_WARP_THREADS;\n __shared__ storage_t temp_storage[TOTAL_WARPS];\n\n const int tid = hipcub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z);\n\n // Get warp index\n const int warp_id = tid / LOGICAL_WARP_THREADS;\n const int lane_id = tid % LOGICAL_WARP_THREADS;\n\n warp_exchange_t exchange(temp_storage[warp_id]);\n\n exchange_data_t exchange_data;\n\n input_data += warp_id * tile_size;\n output_data += warp_id * tile_size;\n\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n exchange_data.input[item] = input_data[lane_id * ITEMS_PER_THREAD + item];\n }\n\n action(exchange_data.input, exchange_data.output, exchange);\n\n for (int item = 0; item < ITEMS_PER_THREAD; item++)\n {\n output_data[lane_id * ITEMS_PER_THREAD + item] = exchange_data.output[item];\n }\n}\n\ntemplate \nvoid warp_exchange(thrust::device_vector &in,\n thrust::device_vector &out,\n ActionT action)\n{\n kernel\n <<<1, LOGICAL_WARP_THREADS * TOTAL_WARPS>>>(thrust::raw_pointer_cast(in.data()),\n thrust::raw_pointer_cast(out.data()),\n action);\n REQUIRE(hipSuccess == hipPeekAtLastError());\n REQUIRE(hipSuccess == hipDeviceSynchronize());\n}\n\nstruct blocked_to_striped\n{\n template \n __device__ void operator()(InputT (&input)[ITEMS_PER_THREAD],\n OutputT (&output)[ITEMS_PER_THREAD],\n hipcub::WarpExchange &exchange)\n {\n exchange.BlockedToStriped(input, output);\n }\n};\n\nstruct striped_to_blocked\n{\n template \n __device__ void operator()(InputT (&input)[ITEMS_PER_THREAD],\n OutputT (&output)[ITEMS_PER_THREAD],\n hipcub::WarpExchange &exchange)\n {\n exchange.StripedToBlocked(input, output);\n }\n};\n\ntemplate \nthrust::host_vector compute_host_reference(const thrust::device_vector &d_input,\n int tile_size)\n{\n thrust::host_vector input = d_input;\n\n int num_warps = CUB_QUOTIENT_CEILING(static_cast(d_input.size()), tile_size);\n for (int warp_id = 0; warp_id < num_warps; warp_id++)\n {\n const int warp_data_begin = tile_size * warp_id;\n const int warp_data_end = warp_data_begin + tile_size;\n thrust::reverse(input.begin() + warp_data_begin, input.begin() + warp_data_end);\n }\n return input;\n}\n\nusing inout_types = c2h::type_list,\n c2h::pair,\n c2h::pair,\n c2h::pair,\n c2h::pair,\n c2h::pair>;\n\nusing logical_warp_threads = c2h::enum_type_list;\nusing items_per_thread = c2h::enum_type_list;\n\ntemplate \nstruct total_warps_t\n{\nprivate:\n static constexpr int max_warps = 2;\n static constexpr bool is_arch_warp = (logical_warp_threads == CUB_WARP_THREADS(0));\n static constexpr bool is_pow_of_two = ((logical_warp_threads & (logical_warp_threads - 1)) == 0);\n static constexpr int total_warps = (is_arch_warp || is_pow_of_two) ? max_warps : 1;\n\npublic:\n static constexpr int value() { return total_warps; }\n};\n\ntemplate \nstruct params_t\n{\n using in_type = typename c2h::first>;\n using out_type = typename c2h::second>;\n\n static constexpr int logical_warp_threads = c2h::get<1, TestType>::value;\n static constexpr int items_per_thread = c2h::get<2, TestType>::value;\n static constexpr int total_warps = total_warps_t::value();\n static constexpr int tile_size = logical_warp_threads * items_per_thread;\n static constexpr int total_item_count = total_warps * tile_size;\n};\n\nCUB_TEST(\"Scatter to striped works\",\n \"[exchange][warp]\",\n inout_types,\n logical_warp_threads,\n items_per_thread)\n{\n using params = params_t;\n using in_type = typename params::in_type;\n using out_type = typename params::out_type;\n thrust::device_vector d_out(params::total_item_count);\n thrust::device_vector d_in(params::total_item_count);\n\n c2h::gen(c2h::modulo_t{d_in.size()}, d_in);\n\n warp_scatter_strided(\n d_in,\n d_out);\n\n auto h_expected_output = compute_host_reference(d_in, params::tile_size);\n REQUIRE(h_expected_output == d_out);\n}\n\nCUB_TEST(\"Blocked to striped works\",\n \"[exchange][warp]\",\n inout_types,\n logical_warp_threads,\n items_per_thread)\n{\n using params = params_t;\n using in_type = typename params::in_type;\n using out_type = typename params::out_type;\n thrust::device_vector d_out(params::total_item_count, out_type{});\n thrust::device_vector d_in(params::total_item_count);\n\n c2h::gen(c2h::modulo_t{d_in.size()}, d_in);\n\n warp_exchange(\n d_in,\n d_out,\n blocked_to_striped{});\n thrust::host_vector h_expected_output(d_out.size());\n fill_striped(h_expected_output.begin());\n\n REQUIRE(h_expected_output == d_out);\n}\n\nCUB_TEST(\"Striped to blocked works\",\n \"[exchange][warp]\",\n inout_types,\n logical_warp_threads,\n items_per_thread)\n{\n using params = params_t;\n using in_type = typename params::in_type;\n using out_type = typename params::out_type;\n thrust::device_vector d_out(params::total_item_count, out_type{});\n\n thrust::host_vector h_in(params::total_item_count);\n fill_striped(h_in.begin());\n thrust::device_vector d_in = h_in;\n\n warp_exchange(\n d_in,\n d_out,\n striped_to_blocked{});\n thrust::device_vector d_expected_output(d_out.size());\n c2h::gen(c2h::modulo_t{d_out.size()}, d_expected_output);\n\n REQUIRE(d_expected_output == d_out);\n}\n", "source": "/kaggle/working/cub/test/catch2_test_warp_exchange.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n#include \n#include \n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \nstruct output_idx\n{\n static __device__ int get(int item)\n {\n return static_cast(threadIdx.x) * ItemsPerThread + item;\n }\n};\n\ntemplate \nstruct output_idx\n{\n static __device__ int get(int item)\n {\n return static_cast(threadIdx.x) + ThreadsInBlock * item;\n }\n};\n\ntemplate \n__global__ void kernel(std::integral_constant,\n InputIteratorT input,\n OutputIteratorT output,\n int num_items)\n{\n using input_t = cub::detail::value_t;\n using block_store_t = cub::BlockStore;\n using storage_t = typename block_store_t::TempStorage;\n\n __shared__ storage_t storage;\n block_store_t block_store(storage);\n\n input_t data[ItemsPerThread];\n\n for (int i = 0; i < ItemsPerThread; i++)\n {\n const int idx =\n output_idx::get(i);\n\n if (idx < num_items)\n {\n data[i] = input[idx];\n }\n }\n\n if (ItemsPerThread * ThreadsInBlock == num_items)\n {\n block_store.Store(output, data);\n }\n else\n {\n block_store.Store(output, data, num_items);\n }\n}\n\ntemplate \n__global__ void kernel(std::integral_constant,\n InputIteratorT input,\n OutputIteratorT output,\n int num_items)\n{\n for (int i = 0; i < ItemsPerThread; i++)\n {\n const int idx =\n output_idx::get(i);\n\n if (idx < num_items)\n {\n output[idx] = input[idx];\n }\n }\n}\n\ntemplate \nvoid block_store(\n InputIteratorT input,\n OutputIteratorT output,\n int num_items)\n{\n using input_t = cub::detail::value_t;\n using block_store_t = cub::BlockStore;\n using storage_t = typename block_store_t::TempStorage;\n constexpr bool sufficient_resources = sizeof(storage_t) <= 1024 * 48;\n\n kernel<<<1, ThreadsInBlock>>>(\n std::integral_constant{},\n input,\n output,\n num_items);\n\n REQUIRE( cudaSuccess == cudaPeekAtLastError() );\n REQUIRE( cudaSuccess == cudaDeviceSynchronize() );\n}\n\n\n// %PARAM% IPT it 1:11\n\nusing types = c2h::type_list;\nusing vec_types = c2h::type_list;\n\nusing even_threads_in_block = c2h::enum_type_list;\nusing odd_threads_in_block = c2h::enum_type_list;\nusing a_block_size = c2h::enum_type_list;\n\nusing items_per_thread = c2h::enum_type_list;\nusing store_algorithm = c2h::enum_type_list<\n cub::BlockStoreAlgorithm,\n cub::BlockStoreAlgorithm::BLOCK_STORE_DIRECT,\n cub::BlockStoreAlgorithm::BLOCK_STORE_STRIPED,\n cub::BlockStoreAlgorithm::BLOCK_STORE_VECTORIZE,\n cub::BlockStoreAlgorithm::BLOCK_STORE_TRANSPOSE,\n cub::BlockStoreAlgorithm::BLOCK_STORE_WARP_TRANSPOSE, \n cub::BlockStoreAlgorithm::BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED>;\n\nusing odd_store_algorithm = c2h::enum_type_list<\n cub::BlockStoreAlgorithm,\n cub::BlockStoreAlgorithm::BLOCK_STORE_DIRECT,\n cub::BlockStoreAlgorithm::BLOCK_STORE_STRIPED,\n cub::BlockStoreAlgorithm::BLOCK_STORE_VECTORIZE,\n cub::BlockStoreAlgorithm::BLOCK_STORE_TRANSPOSE>;\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int items_per_thread = c2h::get<1, TestType>::value;\n static constexpr int threads_in_block = c2h::get<2, TestType>::value;\n static constexpr int tile_size = items_per_thread * threads_in_block;\n static constexpr cub::BlockStoreAlgorithm store_algorithm = \n c2h::get<3, TestType>::value;\n};\n\nCUB_TEST(\"Block store works with even block sizes\",\n \"[store][block]\",\n types,\n items_per_thread,\n even_threads_in_block,\n store_algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_store(thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\nCUB_TEST(\"Block store works with even odd sizes\",\n \"[store][block]\",\n types,\n items_per_thread,\n odd_threads_in_block,\n odd_store_algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_store(thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(\n d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\nCUB_TEST(\"Block store works with even vector types\",\n \"[store][block]\",\n vec_types,\n items_per_thread,\n a_block_size,\n store_algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_store(\n thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\nCUB_TEST(\"Block store works with custom types\",\n \"[store][block]\",\n items_per_thread,\n store_algorithm)\n{\n using type = c2h::custom_type_t;\n constexpr int items_per_thread = c2h::get<0, TestType>::value;\n constexpr int threads_in_block = 64;\n constexpr int tile_size = items_per_thread * threads_in_block;\n static constexpr cub::BlockStoreAlgorithm store_algorithm =\n c2h::get<1, TestType>::value;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_store(\n thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE(d_input == d_output);\n}\n\nCUB_TEST(\"Block store works with caching iterators\",\n \"[store][block]\",\n items_per_thread,\n store_algorithm)\n{\n using type = int;\n constexpr int items_per_thread = c2h::get<0, TestType>::value;\n constexpr int threads_in_block = 64;\n constexpr int tile_size = items_per_thread * threads_in_block;\n static constexpr cub::BlockStoreAlgorithm store_algorithm =\n c2h::get<1, TestType>::value;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n cub::CacheModifiedOutputIterator out(\n thrust::raw_pointer_cast(d_output.data()));\n\n block_store(\n thrust::raw_pointer_cast(d_input.data()),\n out,\n static_cast(d_input.size()));\n\n REQUIRE(d_input == d_output);\n}\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n// Has to go after all cub headers. Otherwise, this test won't catch unused\n// variables in cub kernels.\n#include \"catch2_test_helper.h\"\n\ntemplate \nstruct output_idx\n{\n static __device__ int get(int item)\n {\n return static_cast(threadIdx.x) * ItemsPerThread + item;\n }\n};\n\ntemplate \nstruct output_idx\n{\n static __device__ int get(int item)\n {\n return static_cast(threadIdx.x) + ThreadsInBlock * item;\n }\n};\n\ntemplate \n__global__ void kernel(std::integral_constant,\n InputIteratorT input,\n OutputIteratorT output,\n int num_items)\n{\n using input_t = hipcub::detail::value_t;\n using block_store_t = hipcub::BlockStore;\n using storage_t = typename block_store_t::TempStorage;\n\n __shared__ storage_t storage;\n block_store_t block_store(storage);\n\n input_t data[ItemsPerThread];\n\n for (int i = 0; i < ItemsPerThread; i++)\n {\n const int idx =\n output_idx::get(i);\n\n if (idx < num_items)\n {\n data[i] = input[idx];\n }\n }\n\n if (ItemsPerThread * ThreadsInBlock == num_items)\n {\n block_store.Store(output, data);\n }\n else\n {\n block_store.Store(output, data, num_items);\n }\n}\n\ntemplate \n__global__ void kernel(std::integral_constant,\n InputIteratorT input,\n OutputIteratorT output,\n int num_items)\n{\n for (int i = 0; i < ItemsPerThread; i++)\n {\n const int idx =\n output_idx::get(i);\n\n if (idx < num_items)\n {\n output[idx] = input[idx];\n }\n }\n}\n\ntemplate \nvoid block_store(\n InputIteratorT input,\n OutputIteratorT output,\n int num_items)\n{\n using input_t = hipcub::detail::value_t;\n using block_store_t = hipcub::BlockStore;\n using storage_t = typename block_store_t::TempStorage;\n constexpr bool sufficient_resources = sizeof(storage_t) <= 1024 * 48;\n\n kernel<<<1, ThreadsInBlock>>>(\n std::integral_constant{},\n input,\n output,\n num_items);\n\n REQUIRE( hipSuccess == hipPeekAtLastError() );\n REQUIRE( hipSuccess == hipDeviceSynchronize() );\n}\n\n\n// %PARAM% IPT it 1:11\n\nusing types = c2h::type_list;\nusing vec_types = c2h::type_list;\n\nusing even_threads_in_block = c2h::enum_type_list;\nusing odd_threads_in_block = c2h::enum_type_list;\nusing a_block_size = c2h::enum_type_list;\n\nusing items_per_thread = c2h::enum_type_list;\nusing store_algorithm = c2h::enum_type_list<\n hipcub::BlockStoreAlgorithm,\n hipcub::BlockStoreAlgorithm::BLOCK_STORE_DIRECT,\n hipcub::BlockStoreAlgorithm::BLOCK_STORE_STRIPED,\n hipcub::BlockStoreAlgorithm::BLOCK_STORE_VECTORIZE,\n hipcub::BlockStoreAlgorithm::BLOCK_STORE_TRANSPOSE,\n hipcub::BlockStoreAlgorithm::BLOCK_STORE_WARP_TRANSPOSE, \n hipcub::BlockStoreAlgorithm::BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED>;\n\nusing odd_store_algorithm = c2h::enum_type_list<\n hipcub::BlockStoreAlgorithm,\n hipcub::BlockStoreAlgorithm::BLOCK_STORE_DIRECT,\n hipcub::BlockStoreAlgorithm::BLOCK_STORE_STRIPED,\n hipcub::BlockStoreAlgorithm::BLOCK_STORE_VECTORIZE,\n hipcub::BlockStoreAlgorithm::BLOCK_STORE_TRANSPOSE>;\n\ntemplate \nstruct params_t\n{\n using type = typename c2h::get<0, TestType>;\n\n static constexpr int items_per_thread = c2h::get<1, TestType>::value;\n static constexpr int threads_in_block = c2h::get<2, TestType>::value;\n static constexpr int tile_size = items_per_thread * threads_in_block;\n static constexpr hipcub::BlockStoreAlgorithm store_algorithm = \n c2h::get<3, TestType>::value;\n};\n\nCUB_TEST(\"Block store works with even block sizes\",\n \"[store][block]\",\n types,\n items_per_thread,\n even_threads_in_block,\n store_algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_store(thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\nCUB_TEST(\"Block store works with even odd sizes\",\n \"[store][block]\",\n types,\n items_per_thread,\n odd_threads_in_block,\n odd_store_algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_store(thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(\n d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\nCUB_TEST(\"Block store works with even vector types\",\n \"[store][block]\",\n vec_types,\n items_per_thread,\n a_block_size,\n store_algorithm)\n{\n using params = params_t;\n using type = typename params::type;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, params::tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_store(\n thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE( d_input == d_output );\n}\n\nCUB_TEST(\"Block store works with custom types\",\n \"[store][block]\",\n items_per_thread,\n store_algorithm)\n{\n using type = c2h::custom_type_t;\n constexpr int items_per_thread = c2h::get<0, TestType>::value;\n constexpr int threads_in_block = 64;\n constexpr int tile_size = items_per_thread * threads_in_block;\n static constexpr hipcub::BlockStoreAlgorithm store_algorithm =\n c2h::get<1, TestType>::value;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n\n block_store(\n thrust::raw_pointer_cast(d_input.data()),\n thrust::raw_pointer_cast(d_output.data()),\n static_cast(d_input.size()));\n\n REQUIRE(d_input == d_output);\n}\n\nCUB_TEST(\"Block store works with caching iterators\",\n \"[store][block]\",\n items_per_thread,\n store_algorithm)\n{\n using type = int;\n constexpr int items_per_thread = c2h::get<0, TestType>::value;\n constexpr int threads_in_block = 64;\n constexpr int tile_size = items_per_thread * threads_in_block;\n static constexpr hipcub::BlockStoreAlgorithm store_algorithm =\n c2h::get<1, TestType>::value;\n\n thrust::device_vector d_input(\n GENERATE_COPY(take(10, random(0, tile_size))));\n c2h::gen(CUB_SEED(10), d_input);\n\n thrust::device_vector d_output(d_input.size());\n hipcub::CacheModifiedOutputIterator out(\n thrust::raw_pointer_cast(d_output.data()));\n\n block_store(\n thrust::raw_pointer_cast(d_input.data()),\n out,\n static_cast(d_input.size()));\n\n REQUIRE(d_input == d_output);\n}\n\n", "source": "/kaggle/working/cub/test/catch2_test_block_store.cu"} {"input": "/******************************************************************************\n * Copyright (c) NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceSelect::Unique utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n#include \n#include \n\n#include \n\n#include \"test_util.h\"\n\n#include \n#include \n\nusing namespace cub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\nint g_repeat = 0;\nfloat g_device_giga_bandwidth;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB DeviceSelect entrypoints\n//---------------------------------------------------------------------\n\n\n/**\n * Dispatch to unique entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__\ncudaError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n cudaError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t &temp_storage_bytes,\n KeyInputIteratorT d_keys_in,\n ValueInputIteratorT d_values_in,\n KeyOutputIteratorT d_keys_out,\n ValueOutputIteratorT d_values_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceSelect::UniqueByKey(d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_keys_out, d_values_out, d_num_selected_out, num_items);\n }\n return error;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceSelect\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type cub_backend,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n KeyInputIteratorT d_keys_in,\n ValueInputIteratorT d_values_in,\n KeyOutputIteratorT d_keys_out,\n ValueOutputIteratorT d_values_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items)\n{\n *d_cdp_error = Dispatch(cub_backend,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_values_in,\n d_keys_out,\n d_values_out,\n d_num_selected_out,\n num_items);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Dispatch to CDP kernel\n */\ntemplate \ncudaError_t Dispatch(Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n KeyInputIteratorT d_keys_in,\n ValueInputIteratorT d_values_in,\n KeyOutputIteratorT d_keys_out,\n ValueOutputIteratorT d_values_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items)\n{\n // Invoke kernel to invoke device-side dispatch\n cudaError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n Int2Type{},\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_values_in,\n d_keys_out,\n d_values_out,\n d_num_selected_out,\n num_items);\n CubDebugExit(retval);\n\n // Copy out temp_storage_bytes\n CubDebugExit(cudaMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n cudaMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(cudaMemcpy(&retval,\n d_cdp_error,\n sizeof(cudaError_t) * 1,\n cudaMemcpyDeviceToHost));\n return retval;\n}\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\n\n/**\n * Initialize problem\n */\ntemplate \nvoid Initialize(\n int entropy_reduction,\n T *h_in,\n int num_items,\n int max_segment)\n{\n unsigned int max_int = (unsigned int) -1;\n\n int key = 0;\n int i = 0;\n while (i < num_items)\n {\n // Select number of repeating occurrences for the current run\n int repeat;\n if (max_segment < 0)\n {\n repeat = num_items;\n }\n else if (max_segment < 2)\n {\n repeat = 1;\n }\n else\n {\n RandomBits(repeat, entropy_reduction);\n repeat = (int) ((double(repeat) * double(max_segment)) / double(max_int));\n repeat = CUB_MAX(1, repeat);\n }\n\n int j = i;\n while (j < CUB_MIN(i + repeat, num_items))\n {\n InitValue(INTEGER_SEED, h_in[j], key);\n j++;\n }\n\n i = j;\n key++;\n }\n\n if (g_verbose)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n\n/**\n * Solve unique problem\n */\ntemplate <\n typename KeyInputIteratorT,\n typename ValueInputIteratorT,\n typename KeyT,\n typename ValueT>\nint Solve(\n KeyInputIteratorT h_keys_in,\n ValueInputIteratorT h_values_in,\n KeyT *h_keys_reference,\n ValueT *h_values_reference,\n int num_items)\n{\n int num_selected = 0;\n if (num_items > 0)\n {\n h_keys_reference[num_selected] = h_keys_in[0];\n h_values_reference[num_selected] = h_values_in[0];\n num_selected++;\n }\n\n for (int i = 1; i < num_items; ++i)\n {\n if (h_keys_in[i] != h_keys_in[i - 1])\n {\n h_keys_reference[num_selected] = h_keys_in[i];\n h_values_reference[num_selected] = h_values_in[i];\n num_selected++;\n }\n }\n\n return num_selected;\n}\n\n\n\n/**\n * Test DeviceSelect for a given problem input\n */\ntemplate <\n Backend BACKEND,\n typename KeyInputIteratorT,\n typename ValueInputIteratorT,\n typename KeyT,\n typename ValueT>\nvoid Test(\n KeyInputIteratorT d_keys_in,\n ValueInputIteratorT d_values_in,\n KeyT *h_keys_reference,\n ValueT *h_values_reference,\n int num_selected,\n int num_items)\n{\n // Allocate device output array and num selected\n KeyT *d_keys_out = NULL;\n ValueT *d_values_out = NULL;\n int *d_num_selected_out = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys_out, sizeof(KeyT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_values_out, sizeof(ValueT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_num_selected_out, sizeof(int)));\n\n // Allocate CDP device arrays\n size_t *d_temp_storage_bytes = NULL;\n cudaError_t *d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(cudaError_t) * 1));\n\n // Allocate temporary storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(Int2Type(), 1, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_keys_out, d_values_out, d_num_selected_out, num_items));\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n // Clear device output array\n CubDebugExit(cudaMemset(d_keys_out, 0, sizeof(KeyT) * num_items));\n CubDebugExit(cudaMemset(d_values_out, 0, sizeof(ValueT) * num_items));\n CubDebugExit(cudaMemset(d_num_selected_out, 0, sizeof(int)));\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(Int2Type(), 1, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_keys_out, d_values_out, d_num_selected_out, num_items));\n\n // Check for correctness (and display results, if specified)\n int compare11 = CompareDeviceResults(h_keys_reference, d_keys_out, num_selected, true, g_verbose);\n int compare12 = CompareDeviceResults(h_values_reference, d_values_out, num_selected, true, g_verbose);\n int compare1 = compare11 && compare12;\n printf(\"\\t Data %s \", compare1 ? \"FAIL\" : \"PASS\");\n\n int compare2 = CompareDeviceResults(&num_selected, d_num_selected_out, 1, true, g_verbose);\n printf(\"\\t Count %s \", compare2 ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Performance\n GpuTimer gpu_timer;\n gpu_timer.Start();\n CubDebugExit(Dispatch(Int2Type(), g_timing_iterations, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_keys_out, d_values_out, d_num_selected_out, num_items));\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n\n // Display performance\n if (g_timing_iterations > 0)\n {\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n float giga_bandwidth = float((num_items + num_selected) * (sizeof(KeyT) + sizeof(ValueT))) / avg_millis / 1000.0f / 1000.0f;\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s, %.1f%% peak\", avg_millis, giga_rate, giga_bandwidth, giga_bandwidth / g_device_giga_bandwidth * 100.0);\n }\n printf(\"\\n\\n\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Cleanup\n if (d_keys_out) CubDebugExit(g_allocator.DeviceFree(d_keys_out));\n if (d_values_out) CubDebugExit(g_allocator.DeviceFree(d_values_out));\n if (d_num_selected_out) CubDebugExit(g_allocator.DeviceFree(d_num_selected_out));\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare1 | compare2);\n}\n\n\n/**\n * Test DeviceSelect on pointer type\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT>\nvoid TestPointer(\n int num_items,\n int entropy_reduction,\n int max_segment)\n{\n // Allocate host arrays\n KeyT* h_keys_in = new KeyT[num_items];\n ValueT* h_values_in = new ValueT[num_items];\n KeyT* h_keys_reference = new KeyT[num_items];\n ValueT* h_values_reference = new ValueT[num_items];\n\n // Initialize problem and solution\n Initialize(entropy_reduction, h_keys_in, num_items, max_segment);\n Initialize(entropy_reduction, h_values_in, num_items, max_segment);\n int num_selected = Solve(h_keys_in, h_values_in, h_keys_reference, h_values_reference, num_items);\n\n printf(\"\\nPointer %s cub::DeviceSelect::Unique %d items, %d selected (avg run length %.3f), %s %d-byte elements, entropy_reduction %d\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items, num_selected, float(num_items) / num_selected,\n typeid(KeyT).name(),\n (int) sizeof(KeyT),\n entropy_reduction);\n fflush(stdout);\n\n // Allocate problem device arrays\n KeyT *d_keys_in = NULL;\n ValueT *d_values_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys_in, sizeof(KeyT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_values_in, sizeof(ValueT) * num_items));\n\n // Initialize device input\n CubDebugExit(cudaMemcpy(d_keys_in, h_keys_in, sizeof(KeyT) * num_items, cudaMemcpyHostToDevice));\n CubDebugExit(cudaMemcpy(d_values_in, h_values_in, sizeof(ValueT) * num_items, cudaMemcpyHostToDevice));\n\n // Run Test\n Test(d_keys_in, d_values_in, h_keys_reference, h_values_reference, num_selected, num_items);\n\n // Cleanup\n if (h_keys_in) delete[] h_keys_in;\n if (h_values_in) delete[] h_values_in;\n if (h_keys_reference) delete[] h_keys_reference;\n if (h_values_reference) delete[] h_values_reference;\n if (d_keys_in) CubDebugExit(g_allocator.DeviceFree(d_keys_in));\n if (d_values_in) CubDebugExit(g_allocator.DeviceFree(d_values_in));\n}\n\n\n/**\n * Test DeviceSelect on iterator type\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT>\nvoid TestIterator(\n int num_items)\n{\n // Use a counting iterator as the input\n CountingInputIterator h_keys_in(0);\n CountingInputIterator h_values_in(0);\n\n // Allocate host arrays\n KeyT* h_keys_reference = new KeyT[num_items];\n ValueT* h_values_reference = new ValueT[num_items];\n\n // Initialize problem and solution\n int num_selected = Solve(h_keys_in, h_values_in, h_keys_reference, h_values_reference, num_items);\n\n printf(\"\\nIterator %s cub::DeviceSelect::Unique %d items, %d selected (avg run length %.3f), %s %d-byte elements\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items, num_selected, float(num_items) / num_selected,\n typeid(KeyT).name(),\n (int) sizeof(ValueT));\n fflush(stdout);\n\n // Run Test\n Test(h_keys_in, h_values_in, h_keys_reference, h_values_reference, num_selected, num_items);\n\n // Cleanup\n if (h_keys_reference) delete[] h_keys_reference;\n if (h_values_reference) delete[] h_values_reference;\n}\n\n\n/**\n * Test different gen modes\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT>\nvoid Test(\n int num_items)\n{\n for (int max_segment = 1; ((max_segment > 0) && (max_segment < num_items)); max_segment *= 11)\n {\n TestPointer(num_items, 0, max_segment);\n TestPointer(num_items, 2, max_segment);\n TestPointer(num_items, 7, max_segment);\n }\n}\n\n\n/**\n * Test different dispatch\n */\ntemplate <\n typename KeyT,\n typename ValueT>\nvoid TestOp(\n int num_items)\n{\n#if TEST_CDP == 0\n Test(num_items);\n#elif TEST_CDP == 1\n Test(num_items);\n#endif // TEST_CDP\n}\n\n\n/**\n * Test different input sizes\n */\ntemplate <\n typename KeyT,\n typename ValueT>\nvoid Test(\n int num_items)\n{\n if (num_items < 0)\n {\n TestOp(0);\n TestOp(1);\n TestOp(100);\n TestOp(10000);\n TestOp(1000000);\n }\n else\n {\n TestOp(num_items);\n }\n}\n\n\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n int num_items = -1;\n int entropy_reduction = 0;\n int maxseg = 1000;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n args.GetCmdLineArgument(\"repeat\", g_repeat);\n args.GetCmdLineArgument(\"maxseg\", maxseg);\n args.GetCmdLineArgument(\"entropy\", entropy_reduction);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--maxseg=]\"\n \"[--entropy=]\"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n g_device_giga_bandwidth = args.device_giga_bandwidth;\n printf(\"\\n\");\n\n // %PARAM% TEST_CDP cdp 0:1\n\n // Test different input types\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n return 0;\n}\n\n\n\n", "output": "/******************************************************************************\n * Copyright (c) NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceSelect::Unique utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n#include \n\n#include \n\n#include \"test_util.h\"\n\n#include \n#include \n\nusing namespace hipcub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\nint g_repeat = 0;\nfloat g_device_giga_bandwidth;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB DeviceSelect entrypoints\n//---------------------------------------------------------------------\n\n\n/**\n * Dispatch to unique entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__\nhipError_t Dispatch(\n Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t */*d_temp_storage_bytes*/,\n hipError_t */*d_cdp_error*/,\n\n void* d_temp_storage,\n size_t &temp_storage_bytes,\n KeyInputIteratorT d_keys_in,\n ValueInputIteratorT d_values_in,\n KeyOutputIteratorT d_keys_out,\n ValueOutputIteratorT d_values_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceSelect::UniqueByKey(d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_keys_out, d_values_out, d_num_selected_out, num_items);\n }\n return error;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceSelect\n */\ntemplate \n__global__ void CDPDispatchKernel(Int2Type cub_backend,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n KeyInputIteratorT d_keys_in,\n ValueInputIteratorT d_values_in,\n KeyOutputIteratorT d_keys_out,\n ValueOutputIteratorT d_values_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items)\n{\n *d_cdp_error = Dispatch(cub_backend,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_values_in,\n d_keys_out,\n d_values_out,\n d_num_selected_out,\n num_items);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Dispatch to CDP kernel\n */\ntemplate \nhipError_t Dispatch(Int2Type /*dispatch_to*/,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n KeyInputIteratorT d_keys_in,\n ValueInputIteratorT d_values_in,\n KeyOutputIteratorT d_keys_out,\n ValueOutputIteratorT d_values_out,\n NumSelectedIteratorT d_num_selected_out,\n OffsetT num_items)\n{\n // Invoke kernel to invoke device-side dispatch\n hipError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n Int2Type{},\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_keys_in,\n d_values_in,\n d_keys_out,\n d_values_out,\n d_num_selected_out,\n num_items);\n CubDebugExit(retval);\n\n // Copy out temp_storage_bytes\n CubDebugExit(hipMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n hipMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(hipMemcpy(&retval,\n d_cdp_error,\n sizeof(hipError_t) * 1,\n hipMemcpyDeviceToHost));\n return retval;\n}\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\n\n/**\n * Initialize problem\n */\ntemplate \nvoid Initialize(\n int entropy_reduction,\n T *h_in,\n int num_items,\n int max_segment)\n{\n unsigned int max_int = (unsigned int) -1;\n\n int key = 0;\n int i = 0;\n while (i < num_items)\n {\n // Select number of repeating occurrences for the current run\n int repeat;\n if (max_segment < 0)\n {\n repeat = num_items;\n }\n else if (max_segment < 2)\n {\n repeat = 1;\n }\n else\n {\n RandomBits(repeat, entropy_reduction);\n repeat = (int) ((double(repeat) * double(max_segment)) / double(max_int));\n repeat = CUB_MAX(1, repeat);\n }\n\n int j = i;\n while (j < CUB_MIN(i + repeat, num_items))\n {\n InitValue(INTEGER_SEED, h_in[j], key);\n j++;\n }\n\n i = j;\n key++;\n }\n\n if (g_verbose)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n\n/**\n * Solve unique problem\n */\ntemplate <\n typename KeyInputIteratorT,\n typename ValueInputIteratorT,\n typename KeyT,\n typename ValueT>\nint Solve(\n KeyInputIteratorT h_keys_in,\n ValueInputIteratorT h_values_in,\n KeyT *h_keys_reference,\n ValueT *h_values_reference,\n int num_items)\n{\n int num_selected = 0;\n if (num_items > 0)\n {\n h_keys_reference[num_selected] = h_keys_in[0];\n h_values_reference[num_selected] = h_values_in[0];\n num_selected++;\n }\n\n for (int i = 1; i < num_items; ++i)\n {\n if (h_keys_in[i] != h_keys_in[i - 1])\n {\n h_keys_reference[num_selected] = h_keys_in[i];\n h_values_reference[num_selected] = h_values_in[i];\n num_selected++;\n }\n }\n\n return num_selected;\n}\n\n\n\n/**\n * Test DeviceSelect for a given problem input\n */\ntemplate <\n Backend BACKEND,\n typename KeyInputIteratorT,\n typename ValueInputIteratorT,\n typename KeyT,\n typename ValueT>\nvoid Test(\n KeyInputIteratorT d_keys_in,\n ValueInputIteratorT d_values_in,\n KeyT *h_keys_reference,\n ValueT *h_values_reference,\n int num_selected,\n int num_items)\n{\n // Allocate device output array and num selected\n KeyT *d_keys_out = NULL;\n ValueT *d_values_out = NULL;\n int *d_num_selected_out = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys_out, sizeof(KeyT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_values_out, sizeof(ValueT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_num_selected_out, sizeof(int)));\n\n // Allocate CDP device arrays\n size_t *d_temp_storage_bytes = NULL;\n hipError_t *d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(hipError_t) * 1));\n\n // Allocate temporary storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(Int2Type(), 1, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_keys_out, d_values_out, d_num_selected_out, num_items));\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n // Clear device output array\n CubDebugExit(hipMemset(d_keys_out, 0, sizeof(KeyT) * num_items));\n CubDebugExit(hipMemset(d_values_out, 0, sizeof(ValueT) * num_items));\n CubDebugExit(hipMemset(d_num_selected_out, 0, sizeof(int)));\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(Int2Type(), 1, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_keys_out, d_values_out, d_num_selected_out, num_items));\n\n // Check for correctness (and display results, if specified)\n int compare11 = CompareDeviceResults(h_keys_reference, d_keys_out, num_selected, true, g_verbose);\n int compare12 = CompareDeviceResults(h_values_reference, d_values_out, num_selected, true, g_verbose);\n int compare1 = compare11 && compare12;\n printf(\"\\t Data %s \", compare1 ? \"FAIL\" : \"PASS\");\n\n int compare2 = CompareDeviceResults(&num_selected, d_num_selected_out, 1, true, g_verbose);\n printf(\"\\t Count %s \", compare2 ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Performance\n GpuTimer gpu_timer;\n gpu_timer.Start();\n CubDebugExit(Dispatch(Int2Type(), g_timing_iterations, d_temp_storage_bytes, d_cdp_error, d_temp_storage, temp_storage_bytes, d_keys_in, d_values_in, d_keys_out, d_values_out, d_num_selected_out, num_items));\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n\n // Display performance\n if (g_timing_iterations > 0)\n {\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n float giga_bandwidth = float((num_items + num_selected) * (sizeof(KeyT) + sizeof(ValueT))) / avg_millis / 1000.0f / 1000.0f;\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s, %.1f%% peak\", avg_millis, giga_rate, giga_bandwidth, giga_bandwidth / g_device_giga_bandwidth * 100.0);\n }\n printf(\"\\n\\n\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Cleanup\n if (d_keys_out) CubDebugExit(g_allocator.DeviceFree(d_keys_out));\n if (d_values_out) CubDebugExit(g_allocator.DeviceFree(d_values_out));\n if (d_num_selected_out) CubDebugExit(g_allocator.DeviceFree(d_num_selected_out));\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare1 | compare2);\n}\n\n\n/**\n * Test DeviceSelect on pointer type\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT>\nvoid TestPointer(\n int num_items,\n int entropy_reduction,\n int max_segment)\n{\n // Allocate host arrays\n KeyT* h_keys_in = new KeyT[num_items];\n ValueT* h_values_in = new ValueT[num_items];\n KeyT* h_keys_reference = new KeyT[num_items];\n ValueT* h_values_reference = new ValueT[num_items];\n\n // Initialize problem and solution\n Initialize(entropy_reduction, h_keys_in, num_items, max_segment);\n Initialize(entropy_reduction, h_values_in, num_items, max_segment);\n int num_selected = Solve(h_keys_in, h_values_in, h_keys_reference, h_values_reference, num_items);\n\n printf(\"\\nPointer %s hipcub::DeviceSelect::Unique %d items, %d selected (avg run length %.3f), %s %d-byte elements, entropy_reduction %d\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items, num_selected, float(num_items) / num_selected,\n typeid(KeyT).name(),\n (int) sizeof(KeyT),\n entropy_reduction);\n fflush(stdout);\n\n // Allocate problem device arrays\n KeyT *d_keys_in = NULL;\n ValueT *d_values_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_keys_in, sizeof(KeyT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_values_in, sizeof(ValueT) * num_items));\n\n // Initialize device input\n CubDebugExit(hipMemcpy(d_keys_in, h_keys_in, sizeof(KeyT) * num_items, hipMemcpyHostToDevice));\n CubDebugExit(hipMemcpy(d_values_in, h_values_in, sizeof(ValueT) * num_items, hipMemcpyHostToDevice));\n\n // Run Test\n Test(d_keys_in, d_values_in, h_keys_reference, h_values_reference, num_selected, num_items);\n\n // Cleanup\n if (h_keys_in) delete[] h_keys_in;\n if (h_values_in) delete[] h_values_in;\n if (h_keys_reference) delete[] h_keys_reference;\n if (h_values_reference) delete[] h_values_reference;\n if (d_keys_in) CubDebugExit(g_allocator.DeviceFree(d_keys_in));\n if (d_values_in) CubDebugExit(g_allocator.DeviceFree(d_values_in));\n}\n\n\n/**\n * Test DeviceSelect on iterator type\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT>\nvoid TestIterator(\n int num_items)\n{\n // Use a counting iterator as the input\n CountingInputIterator h_keys_in(0);\n CountingInputIterator h_values_in(0);\n\n // Allocate host arrays\n KeyT* h_keys_reference = new KeyT[num_items];\n ValueT* h_values_reference = new ValueT[num_items];\n\n // Initialize problem and solution\n int num_selected = Solve(h_keys_in, h_values_in, h_keys_reference, h_values_reference, num_items);\n\n printf(\"\\nIterator %s hipcub::DeviceSelect::Unique %d items, %d selected (avg run length %.3f), %s %d-byte elements\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n num_items, num_selected, float(num_items) / num_selected,\n typeid(KeyT).name(),\n (int) sizeof(ValueT));\n fflush(stdout);\n\n // Run Test\n Test(h_keys_in, h_values_in, h_keys_reference, h_values_reference, num_selected, num_items);\n\n // Cleanup\n if (h_keys_reference) delete[] h_keys_reference;\n if (h_values_reference) delete[] h_values_reference;\n}\n\n\n/**\n * Test different gen modes\n */\ntemplate <\n Backend BACKEND,\n typename KeyT,\n typename ValueT>\nvoid Test(\n int num_items)\n{\n for (int max_segment = 1; ((max_segment > 0) && (max_segment < num_items)); max_segment *= 11)\n {\n TestPointer(num_items, 0, max_segment);\n TestPointer(num_items, 2, max_segment);\n TestPointer(num_items, 7, max_segment);\n }\n}\n\n\n/**\n * Test different dispatch\n */\ntemplate <\n typename KeyT,\n typename ValueT>\nvoid TestOp(\n int num_items)\n{\n#if TEST_CDP == 0\n Test(num_items);\n#elif TEST_CDP == 1\n Test(num_items);\n#endif // TEST_CDP\n}\n\n\n/**\n * Test different input sizes\n */\ntemplate <\n typename KeyT,\n typename ValueT>\nvoid Test(\n int num_items)\n{\n if (num_items < 0)\n {\n TestOp(0);\n TestOp(1);\n TestOp(100);\n TestOp(10000);\n TestOp(1000000);\n }\n else\n {\n TestOp(num_items);\n }\n}\n\n\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n int num_items = -1;\n int entropy_reduction = 0;\n int maxseg = 1000;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n args.GetCmdLineArgument(\"repeat\", g_repeat);\n args.GetCmdLineArgument(\"maxseg\", maxseg);\n args.GetCmdLineArgument(\"entropy\", entropy_reduction);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--maxseg=]\"\n \"[--entropy=]\"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n g_device_giga_bandwidth = args.device_giga_bandwidth;\n printf(\"\\n\");\n\n // %PARAM% TEST_CDP cdp 0:1\n\n // Test different input types\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n Test(num_items);\n\n return 0;\n}\n\n\n\n", "source": "/kaggle/working/cub/test/test_device_select_unique_by_key.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"test_util.h\"\n\n/**\n * @brief Host-side random data generation\n */\ntemplate \nvoid GenerateRandomData(\n T *rand_out,\n const std::size_t num_items,\n const T min_rand_val = std::numeric_limits::min(),\n const T max_rand_val = std::numeric_limits::max(),\n const std::uint_fast32_t seed = 320981U,\n typename std::enable_if::value && (sizeof(T) >= 2)>::type * = nullptr)\n{\n // initialize random number generator\n std::mt19937 rng(seed);\n std::uniform_int_distribution uni_dist(min_rand_val, max_rand_val);\n\n // generate random numbers\n for (std::size_t i = 0; i < num_items; ++i)\n {\n rand_out[i] = uni_dist(rng);\n }\n}\n\ntemplate \nvoid __global__ BaselineBatchMemCpyKernel(InputBufferIt input_buffer_it,\n OutputBufferIt output_buffer_it,\n BufferSizeIteratorT buffer_sizes,\n BufferOffsetT num_buffers)\n{\n BufferOffsetT gtid = blockDim.x * blockIdx.x + threadIdx.x;\n if (gtid >= num_buffers)\n {\n return;\n }\n for (BufferOffsetT i = 0; i < buffer_sizes[gtid]; i++)\n {\n reinterpret_cast(output_buffer_it[gtid])[i] =\n reinterpret_cast(input_buffer_it[gtid])[i];\n }\n}\n\ntemplate \nvoid InvokeBaselineBatchMemcpy(InputBufferIt input_buffer_it,\n OutputBufferIt output_buffer_it,\n BufferSizeIteratorT buffer_sizes,\n uint32_t num_buffers)\n{\n constexpr uint32_t block_threads = 128U;\n uint32_t num_blocks = (num_buffers + block_threads - 1) / block_threads;\n BaselineBatchMemCpyKernel<<>>(input_buffer_it,\n output_buffer_it,\n buffer_sizes,\n num_buffers);\n}\n\ntemplate \nvoid __global__ BaselineBatchMemCpyPerBlockKernel(InputBufferIt input_buffer_it,\n OutputBufferIt output_buffer_it,\n BufferSizeIteratorT buffer_sizes,\n BufferOffsetT num_buffers)\n{\n BufferOffsetT gbid = blockIdx.x;\n if (gbid >= num_buffers)\n {\n return;\n }\n for (BufferOffsetT i = threadIdx.x; i < buffer_sizes[gbid] / 8; i += blockDim.x)\n {\n reinterpret_cast(output_buffer_it[gbid])[i] =\n reinterpret_cast(input_buffer_it[gbid])[i];\n }\n}\n\n/**\n * @brief Used for generating a shuffled but cohesive sequence of output-buffer offsets for the\n * sequence of input-buffers.\n */\ntemplate \nstd::vector GetShuffledBufferOffsets(const std::vector &buffer_sizes,\n const std::uint_fast32_t seed = 320981U)\n{\n BufferOffsetT num_buffers = static_cast(buffer_sizes.size());\n\n // We're remapping the i-th buffer to pmt_idxs[i]\n std::mt19937 rng(seed);\n std::vector pmt_idxs(num_buffers);\n std::iota(pmt_idxs.begin(), pmt_idxs.end(), static_cast(0));\n std::shuffle(std::begin(pmt_idxs), std::end(pmt_idxs), rng);\n\n // Compute the offsets using the new mapping\n ByteOffsetT running_offset = {};\n std::vector permuted_offsets;\n permuted_offsets.reserve(num_buffers);\n for (auto permuted_buffer_idx : pmt_idxs)\n {\n permuted_offsets.emplace_back(running_offset);\n running_offset += buffer_sizes[permuted_buffer_idx];\n }\n\n // Generate the scatter indexes that identify where each buffer was mapped to\n std::vector scatter_idxs(num_buffers);\n for (BufferOffsetT i = 0; i < num_buffers; i++)\n {\n scatter_idxs[pmt_idxs[i]] = i;\n }\n\n std::vector new_offsets(num_buffers);\n for (BufferOffsetT i = 0; i < num_buffers; i++)\n {\n new_offsets[i] = permuted_offsets[scatter_idxs[i]];\n }\n\n return new_offsets;\n}\n\n/**\n * @brief Function object class template that takes an offset and returns an iterator at the given\n * offset relative to a fixed base iterator.\n *\n * @tparam IteratorT The random-access iterator type to be returned\n */\ntemplate \nstruct OffsetToPtrOp\n{\n template \n __host__ __device__ __forceinline__ IteratorT operator()(T offset) const\n {\n return base_it + offset;\n }\n IteratorT base_it;\n};\n\nenum class TestDataGen\n{\n // Random offsets into a data segment\n RANDOM,\n\n // Buffers cohesively reside next to each other\n CONSECUTIVE\n};\n\n/**\n * @brief\n *\n * @tparam AtomicT The most granular type being copied. All source and destination pointers will be\n * aligned based on this type, the number of bytes being copied will be an integer multiple of this\n * type's size\n * @tparam BufferOffsetT Type used for indexing into the array of buffers\n * @tparam BufferSizeT Type used for indexing into individual bytes of a buffer (large enough to\n * cover the max buffer size)\n * @tparam ByteOffsetT Type used for indexing into bytes over *all* the buffers' sizes\n */\ntemplate \nvoid RunTest(BufferOffsetT num_buffers,\n BufferSizeT min_buffer_size,\n BufferSizeT max_buffer_size,\n TestDataGen input_gen,\n TestDataGen output_gen)\n{\n using SrcPtrT = uint8_t *;\n\n // Buffer segment data (their offsets and sizes)\n std::vector h_buffer_sizes(num_buffers);\n std::vector h_buffer_src_offsets(num_buffers);\n std::vector h_buffer_dst_offsets(num_buffers);\n\n // Device-side resources\n void *d_in = nullptr;\n void *d_out = nullptr;\n ByteOffsetT *d_buffer_src_offsets = nullptr;\n ByteOffsetT *d_buffer_dst_offsets = nullptr;\n BufferSizeT *d_buffer_sizes = nullptr;\n void *d_temp_storage = nullptr;\n size_t temp_storage_bytes = 0;\n\n // Generate the buffer sizes\n GenerateRandomData(h_buffer_sizes.data(), h_buffer_sizes.size(), min_buffer_size, max_buffer_size);\n\n // Make sure buffer sizes are a multiple of the most granular unit (one AtomicT) being copied\n // (round down)\n for (BufferOffsetT i = 0; i < num_buffers; i++)\n {\n h_buffer_sizes[i] = (h_buffer_sizes[i] / sizeof(AtomicT)) * sizeof(AtomicT);\n }\n\n // Compute the total bytes to be copied\n ByteOffsetT num_total_bytes = 0;\n for (BufferOffsetT i = 0; i < num_buffers; i++)\n {\n if (input_gen == TestDataGen::CONSECUTIVE)\n {\n h_buffer_src_offsets[i] = num_total_bytes;\n }\n if (output_gen == TestDataGen::CONSECUTIVE)\n {\n h_buffer_dst_offsets[i] = num_total_bytes;\n }\n num_total_bytes += h_buffer_sizes[i];\n }\n\n // Shuffle input buffer source-offsets\n std::uint_fast32_t shuffle_seed = 320981U;\n if (input_gen == TestDataGen::RANDOM)\n {\n h_buffer_src_offsets = GetShuffledBufferOffsets(h_buffer_sizes,\n shuffle_seed);\n shuffle_seed += 42;\n }\n\n // Shuffle input buffer source-offsets\n if (output_gen == TestDataGen::RANDOM)\n {\n h_buffer_dst_offsets = GetShuffledBufferOffsets(h_buffer_sizes,\n shuffle_seed);\n }\n\n // Get temporary storage requirements\n CubDebugExit(cub::DeviceMemcpy::Batched(d_temp_storage,\n temp_storage_bytes,\n static_cast(nullptr),\n static_cast(nullptr),\n d_buffer_sizes,\n num_buffers));\n\n // Check if there's sufficient device memory to run this test\n std::size_t total_required_mem = num_total_bytes + //\n num_total_bytes + //\n (num_buffers * sizeof(d_buffer_src_offsets[0])) + //\n (num_buffers * sizeof(d_buffer_dst_offsets[0])) + //\n (num_buffers * sizeof(d_buffer_sizes[0])) + //\n temp_storage_bytes; //\n if (TotalGlobalMem() < total_required_mem)\n {\n std::cout\n << \"Skipping the test due to insufficient device memory\\n\" //\n << \" - Required: \" << total_required_mem << \" B, available: \" << TotalGlobalMem() << \" B\\n\" //\n << \" - Skipped test instance: \" //\n << \" -> Min. buffer size: \" << min_buffer_size << \", max. buffer size: \" << max_buffer_size //\n << \", num_buffers: \" << num_buffers //\n << \", in_gen: \" << ((input_gen == TestDataGen::RANDOM) ? \"SHFL\" : \"CONSECUTIVE\") //\n << \", out_gen: \" << ((output_gen == TestDataGen::RANDOM) ? \"SHFL\" : \"CONSECUTIVE\");\n return;\n }\n\n cudaEvent_t events[2];\n cudaEventCreate(&events[0]);\n cudaEventCreate(&events[1]);\n\n cudaStream_t stream;\n cudaStreamCreate(&stream);\n\n // Allocate device memory\n CubDebugExit(cudaMalloc(&d_in, num_total_bytes));\n CubDebugExit(cudaMalloc(&d_out, num_total_bytes));\n CubDebugExit(cudaMalloc(&d_buffer_src_offsets, num_buffers * sizeof(d_buffer_src_offsets[0])));\n CubDebugExit(cudaMalloc(&d_buffer_dst_offsets, num_buffers * sizeof(d_buffer_dst_offsets[0])));\n CubDebugExit(cudaMalloc(&d_buffer_sizes, num_buffers * sizeof(d_buffer_sizes[0])));\n CubDebugExit(cudaMalloc(&d_temp_storage, temp_storage_bytes));\n\n // Populate the data source with random data\n using RandomInitAliasT = uint16_t;\n std::size_t num_aliased_factor = sizeof(RandomInitAliasT) / sizeof(uint8_t);\n std::size_t num_aliased_units = CUB_QUOTIENT_CEILING(num_total_bytes, num_aliased_factor);\n std::unique_ptr h_in(new uint8_t[num_aliased_units * num_aliased_factor]);\n std::unique_ptr h_out(new uint8_t[num_total_bytes]);\n std::unique_ptr h_gpu_results(new uint8_t[num_total_bytes]);\n\n // Generate random offsets into the random-bits data buffer\n GenerateRandomData(reinterpret_cast(h_in.get()), num_aliased_units);\n\n // Prepare d_buffer_srcs\n OffsetToPtrOp src_transform_op{static_cast(d_in)};\n cub::TransformInputIterator, ByteOffsetT *> d_buffer_srcs(\n d_buffer_src_offsets,\n src_transform_op);\n\n // Prepare d_buffer_dsts\n OffsetToPtrOp dst_transform_op{static_cast(d_out)};\n cub::TransformInputIterator, ByteOffsetT *> d_buffer_dsts(\n d_buffer_dst_offsets,\n dst_transform_op);\n\n // Prepare random data segment (which serves for the buffer sources)\n CubDebugExit(cudaMemcpyAsync(d_in, h_in.get(), num_total_bytes, cudaMemcpyHostToDevice, stream));\n\n // Prepare d_buffer_src_offsets\n CubDebugExit(cudaMemcpyAsync(d_buffer_src_offsets,\n h_buffer_src_offsets.data(),\n h_buffer_src_offsets.size() * sizeof(h_buffer_src_offsets[0]),\n cudaMemcpyHostToDevice,\n stream));\n\n // Prepare d_buffer_dst_offsets\n CubDebugExit(cudaMemcpyAsync(d_buffer_dst_offsets,\n h_buffer_dst_offsets.data(),\n h_buffer_dst_offsets.size() * sizeof(h_buffer_dst_offsets[0]),\n cudaMemcpyHostToDevice,\n stream));\n\n // Prepare d_buffer_sizes\n CubDebugExit(cudaMemcpyAsync(d_buffer_sizes,\n h_buffer_sizes.data(),\n h_buffer_sizes.size() * sizeof(h_buffer_sizes[0]),\n cudaMemcpyHostToDevice,\n stream));\n\n // Record event before algorithm\n cudaEventRecord(events[0], stream);\n\n // Invoke device-side algorithm being under test\n CubDebugExit(cub::DeviceMemcpy::Batched(d_temp_storage,\n temp_storage_bytes,\n d_buffer_srcs,\n d_buffer_dsts,\n d_buffer_sizes,\n num_buffers,\n stream));\n\n // Record event after algorithm\n cudaEventRecord(events[1], stream);\n\n // Copy back the output buffer\n CubDebugExit(\n cudaMemcpyAsync(h_gpu_results.get(), d_out, num_total_bytes, cudaMemcpyDeviceToHost, stream));\n\n // Make sure results have been copied back to the host\n CubDebugExit(cudaStreamSynchronize(stream));\n\n // CPU-side result generation for verification\n for (BufferOffsetT i = 0; i < num_buffers; i++)\n {\n std::memcpy(h_out.get() + h_buffer_dst_offsets[i],\n h_in.get() + h_buffer_src_offsets[i],\n h_buffer_sizes[i]);\n }\n\n float duration = 0;\n cudaEventElapsedTime(&duration, events[0], events[1]);\n\n#ifdef CUB_TEST_BENCHMARK\n size_t stats_src_offsets = sizeof(ByteOffsetT) * num_buffers;\n size_t stats_dst_offsets = sizeof(ByteOffsetT) * num_buffers;\n size_t stats_sizes = sizeof(BufferSizeT) * num_buffers;\n size_t stats_data_copied = 2 * num_total_bytes;\n\n std::cout\n << \"Min. buffer size: \" << min_buffer_size << \", max. buffer size: \" << max_buffer_size //\n << \", num_buffers: \" << num_buffers //\n << \", in_gen: \" << ((input_gen == TestDataGen::RANDOM) ? \"SHFL\" : \"CONSECUTIVE\") //\n << \", out_gen: \" << ((output_gen == TestDataGen::RANDOM) ? \"SHFL\" : \"CONSECUTIVE\") //\n << \", src size: \" << stats_src_offsets << \", dst size: \" << stats_dst_offsets //\n << \", sizes size: \" << stats_sizes << \", cpy_data_size: \" << stats_data_copied //\n << \", total: \" << (stats_src_offsets + stats_dst_offsets + stats_sizes + stats_data_copied) //\n << \", duration: \" << duration //\n << \", BW: \"\n << ((double)(stats_src_offsets + stats_dst_offsets + stats_sizes + stats_data_copied) /\n 1000000000.0) /\n (duration / 1000.0)\n << \"GB/s \\n\";\n#endif\n\n for (ByteOffsetT i = 0; i < num_total_bytes; i++)\n {\n if (h_gpu_results.get()[i] != h_out.get()[i])\n {\n std::cout << \"Mismatch at index \" << i\n << \", CPU vs. GPU: \" << static_cast(h_gpu_results.get()[i]) << \", \"\n << static_cast(h_out.get()[i]) << \"\\n\";\n }\n AssertEquals(h_out.get()[i], h_gpu_results.get()[i]);\n }\n\n CubDebugExit(cudaFree(d_in));\n CubDebugExit(cudaFree(d_out));\n CubDebugExit(cudaFree(d_buffer_src_offsets));\n CubDebugExit(cudaFree(d_buffer_dst_offsets));\n CubDebugExit(cudaFree(d_buffer_sizes));\n CubDebugExit(cudaFree(d_temp_storage));\n}\n\ntemplate \n__global__ void TestVectorizedCopyKernel(const void *d_in, void *d_out, ByteOffsetT copy_size)\n{\n cub::detail::VectorizedCopy(threadIdx.x, d_out, copy_size, d_in);\n}\n\nstruct TupleMemberEqualityOp\n{\n template \n __host__ __device__ __forceinline__ bool operator()(T tuple)\n {\n return thrust::get<0>(tuple) == thrust::get<1>(tuple);\n }\n};\n\n/**\n * @brief Tests the VectorizedCopy for various aligned and misaligned input and output pointers.\n * @tparam VectorT The vector type used for vectorized stores (i.e., one of uint4, uint2, uint32_t)\n */\ntemplate \nvoid TestVectorizedCopy()\n{\n\n constexpr uint32_t threads_per_block = 8;\n\n std::vector in_offsets{0, 1, sizeof(uint32_t) - 1};\n std::vector out_offsets{0, 1, sizeof(VectorT) - 1};\n std::vector copy_sizes{0,\n 1,\n sizeof(uint32_t),\n sizeof(VectorT),\n 2 * threads_per_block * sizeof(VectorT)};\n for (auto copy_sizes_it = std::begin(copy_sizes); copy_sizes_it < std::end(copy_sizes);\n copy_sizes_it++)\n {\n for (auto in_offsets_it = std::begin(in_offsets); in_offsets_it < std::end(in_offsets);\n in_offsets_it++)\n {\n for (auto out_offsets_it = std::begin(out_offsets); out_offsets_it < std::end(out_offsets);\n out_offsets_it++)\n {\n std::size_t in_offset = *in_offsets_it;\n std::size_t out_offset = *out_offsets_it;\n std::size_t copy_size = *copy_sizes_it;\n\n // Prepare data\n const std::size_t alloc_size_in = in_offset + copy_size;\n const std::size_t alloc_size_out = out_offset + copy_size;\n thrust::device_vector data_in(alloc_size_in);\n thrust::device_vector data_out(alloc_size_out);\n thrust::sequence(data_in.begin(), data_in.end(), static_cast(0));\n thrust::fill_n(data_out.begin(), alloc_size_out, static_cast(0x42));\n\n auto d_in = thrust::raw_pointer_cast(data_in.data());\n auto d_out = thrust::raw_pointer_cast(data_out.data());\n\n TestVectorizedCopyKernel\n <<<1, threads_per_block>>>(d_in + in_offset,\n d_out + out_offset,\n static_cast(copy_size));\n auto zip_it = thrust::make_zip_iterator(data_in.begin() + in_offset,\n data_out.begin() + out_offset);\n\n bool success = thrust::all_of(zip_it, zip_it + copy_size, TupleMemberEqualityOp{});\n AssertTrue(success);\n }\n }\n }\n}\n\ntemplate \n__global__ void TestBitPackedCounterKernel(uint32_t *bins,\n uint32_t *increments,\n uint32_t *counts_out,\n uint32_t num_items)\n{\n using BitPackedCounterT =\n cub::detail::BitPackedCounter;\n BitPackedCounterT counter{};\n for (uint32_t i = 0; i < num_items; i++)\n {\n counter.Add(bins[i], increments[i]);\n }\n\n for (uint32_t i = 0; i < NUM_ITEMS; i++)\n {\n counts_out[i] = counter.Get(i);\n }\n}\n\n/**\n * @brief Tests BitPackedCounter that's used for computing the histogram of buffer sizes (i.e.,\n * small, medium, large).\n */\ntemplate \nvoid TestBitPackedCounter(const std::uint_fast32_t seed = 320981U)\n{\n\n constexpr uint32_t min_increment = 0;\n constexpr uint32_t max_increment = 4;\n constexpr double avg_increment = static_cast(min_increment) +\n (static_cast(max_increment - min_increment) / 2.0);\n std::uint32_t num_increments = \n static_cast(static_cast(MAX_ITEM_VALUE * NUM_ITEMS) / avg_increment);\n\n // Test input data\n std::array reference_counters{};\n thrust::host_vector h_bins(num_increments);\n thrust::host_vector h_increments(num_increments);\n\n // Generate random test input data\n GenerateRandomData(thrust::raw_pointer_cast(h_bins.data()),\n num_increments,\n 0U,\n NUM_ITEMS - 1U,\n seed);\n GenerateRandomData(thrust::raw_pointer_cast(h_increments.data()),\n num_increments,\n min_increment,\n max_increment,\n (seed + 17));\n\n // Make sure test data does not overflow any of the counters\n for (std::size_t i = 0; i < num_increments; i++)\n {\n // New increment for this bin would overflow => zero this increment\n if (reference_counters[h_bins[i]] + h_increments[i] >= MAX_ITEM_VALUE)\n {\n h_increments[i] = 0;\n }\n else\n {\n reference_counters[h_bins[i]] += h_increments[i];\n }\n }\n\n // Device memory\n thrust::device_vector bins_in(num_increments);\n thrust::device_vector increments_in(num_increments);\n thrust::device_vector counts_out(NUM_ITEMS);\n\n // Initialize device-side test data\n bins_in = h_bins;\n increments_in = h_increments;\n\n // Memory for GPU-generated results\n thrust::host_vector host_counts(num_increments);\n\n // Reset counters to arbitrary random value\n thrust::fill(counts_out.begin(), counts_out.end(), 814920U);\n\n // Run tests with densely bit-packed counters\n TestBitPackedCounterKernel\n <<<1, 1>>>(thrust::raw_pointer_cast(bins_in.data()),\n thrust::raw_pointer_cast(increments_in.data()),\n thrust::raw_pointer_cast(counts_out.data()),\n num_increments);\n\n // Result verification\n host_counts = counts_out;\n for (uint32_t i = 0; i < NUM_ITEMS; i++)\n {\n AssertEquals(reference_counters[i], host_counts[i]);\n }\n\n // Reset counters to arbitrary random value\n thrust::fill(counts_out.begin(), counts_out.end(), 814920U);\n\n // Run tests with bit-packed counters, where bit-count is a power-of-two\n TestBitPackedCounterKernel\n <<<1, 1>>>(thrust::raw_pointer_cast(bins_in.data()),\n thrust::raw_pointer_cast(increments_in.data()),\n thrust::raw_pointer_cast(counts_out.data()),\n num_increments);\n\n // Result verification\n host_counts = counts_out;\n for (uint32_t i = 0; i < NUM_ITEMS; i++)\n {\n AssertEquals(reference_counters[i], host_counts[i]);\n }\n}\n\nint main(int argc, char **argv)\n{\n CommandLineArgs args(argc, argv);\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n //---------------------------------------------------------------------\n // VectorizedCopy tests\n //---------------------------------------------------------------------\n TestVectorizedCopy();\n TestVectorizedCopy();\n\n //---------------------------------------------------------------------\n // BitPackedCounter tests\n //---------------------------------------------------------------------\n TestBitPackedCounter<1, 1>();\n TestBitPackedCounter<1, (0x01U << 16)>();\n TestBitPackedCounter<4, 1>();\n TestBitPackedCounter<4, 2>();\n TestBitPackedCounter<4, 255>();\n TestBitPackedCounter<4, 256>();\n TestBitPackedCounter<8, 1024>();\n TestBitPackedCounter<32, 1>();\n TestBitPackedCounter<32, 256>();\n\n //---------------------------------------------------------------------\n // DeviceMemcpy::Batched tests\n //---------------------------------------------------------------------\n // The most granular type being copied. Buffer's will be aligned and their size be an integer\n // multiple of this type\n using AtomicCopyT = uint8_t;\n\n // Type used for indexing into the array of buffers\n using BufferOffsetT = uint32_t;\n\n // Type used for indexing into individual bytes of a buffer (large enough to cover the max buffer\n using BufferSizeT = uint32_t;\n\n // Type used for indexing into bytes over *all* the buffers' sizes\n using ByteOffsetT = uint32_t;\n\n // Total number of bytes that are targeted to be copied on each run\n const BufferOffsetT target_copy_size = 64U << 20;\n\n // The number of randomly\n constexpr std::size_t num_rnd_buffer_range_tests = 32;\n\n // Each buffer's size will be random within this interval\n std::vector> buffer_size_ranges = {{0, 1},\n {1, 2},\n {0, 16},\n {1, 32},\n {1, 1024},\n {1, 32 * 1024},\n {128 * 1024, 256 * 1024},\n {target_copy_size,\n target_copy_size}};\n\n std::mt19937 rng(0);\n std::uniform_int_distribution size_dist(1, 1000000);\n for (std::size_t i = 0; i < num_rnd_buffer_range_tests; i++)\n {\n auto range_begin = size_dist(rng);\n auto range_end = size_dist(rng);\n if (range_begin > range_end)\n {\n std::swap(range_begin, range_end);\n }\n buffer_size_ranges.push_back({range_begin, range_end});\n }\n\n for (const auto &buffer_size_range : buffer_size_ranges)\n {\n BufferSizeT min_buffer_size =\n static_cast(CUB_ROUND_UP_NEAREST(buffer_size_range.first, sizeof(AtomicCopyT)));\n BufferSizeT max_buffer_size =\n static_cast(CUB_ROUND_UP_NEAREST(buffer_size_range.second,\n static_cast(sizeof(AtomicCopyT))));\n double average_buffer_size = (min_buffer_size + max_buffer_size) / 2.0;\n BufferOffsetT target_num_buffers =\n static_cast(target_copy_size / average_buffer_size);\n\n // Run tests with input buffer being consecutive and output buffers being consecutive\n RunTest(target_num_buffers,\n min_buffer_size,\n max_buffer_size,\n TestDataGen::CONSECUTIVE,\n TestDataGen::CONSECUTIVE);\n\n // Run tests with input buffer being randomly shuffled and output buffers being randomly\n // shuffled\n RunTest(target_num_buffers,\n min_buffer_size,\n max_buffer_size,\n TestDataGen::RANDOM,\n TestDataGen::RANDOM);\n }\n\n //---------------------------------------------------------------------\n // DeviceMemcpy::Batched test with 64-bit offsets\n //---------------------------------------------------------------------\n using ByteOffset64T = uint64_t;\n using BufferSize64T = uint64_t;\n ByteOffset64T large_target_copy_size =\n static_cast(std::numeric_limits::max()) + (128ULL * 1024ULL * 1024ULL);\n // Make sure min_buffer_size is in fact smaller than max buffer size\n constexpr BufferOffsetT single_buffer = 1;\n\n // Run tests with input buffer being consecutive and output buffers being consecutive\n RunTest(single_buffer,\n large_target_copy_size,\n large_target_copy_size,\n TestDataGen::CONSECUTIVE,\n TestDataGen::CONSECUTIVE);\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"test_util.h\"\n\n/**\n * @brief Host-side random data generation\n */\ntemplate \nvoid GenerateRandomData(\n T *rand_out,\n const std::size_t num_items,\n const T min_rand_val = std::numeric_limits::min(),\n const T max_rand_val = std::numeric_limits::max(),\n const std::uint_fast32_t seed = 320981U,\n typename std::enable_if::value && (sizeof(T) >= 2)>::type * = nullptr)\n{\n // initialize random number generator\n std::mt19937 rng(seed);\n std::uniform_int_distribution uni_dist(min_rand_val, max_rand_val);\n\n // generate random numbers\n for (std::size_t i = 0; i < num_items; ++i)\n {\n rand_out[i] = uni_dist(rng);\n }\n}\n\ntemplate \nvoid __global__ BaselineBatchMemCpyKernel(InputBufferIt input_buffer_it,\n OutputBufferIt output_buffer_it,\n BufferSizeIteratorT buffer_sizes,\n BufferOffsetT num_buffers)\n{\n BufferOffsetT gtid = blockDim.x * blockIdx.x + threadIdx.x;\n if (gtid >= num_buffers)\n {\n return;\n }\n for (BufferOffsetT i = 0; i < buffer_sizes[gtid]; i++)\n {\n reinterpret_cast(output_buffer_it[gtid])[i] =\n reinterpret_cast(input_buffer_it[gtid])[i];\n }\n}\n\ntemplate \nvoid InvokeBaselineBatchMemcpy(InputBufferIt input_buffer_it,\n OutputBufferIt output_buffer_it,\n BufferSizeIteratorT buffer_sizes,\n uint32_t num_buffers)\n{\n constexpr uint32_t block_threads = 128U;\n uint32_t num_blocks = (num_buffers + block_threads - 1) / block_threads;\n BaselineBatchMemCpyKernel<<>>(input_buffer_it,\n output_buffer_it,\n buffer_sizes,\n num_buffers);\n}\n\ntemplate \nvoid __global__ BaselineBatchMemCpyPerBlockKernel(InputBufferIt input_buffer_it,\n OutputBufferIt output_buffer_it,\n BufferSizeIteratorT buffer_sizes,\n BufferOffsetT num_buffers)\n{\n BufferOffsetT gbid = blockIdx.x;\n if (gbid >= num_buffers)\n {\n return;\n }\n for (BufferOffsetT i = threadIdx.x; i < buffer_sizes[gbid] / 8; i += blockDim.x)\n {\n reinterpret_cast(output_buffer_it[gbid])[i] =\n reinterpret_cast(input_buffer_it[gbid])[i];\n }\n}\n\n/**\n * @brief Used for generating a shuffled but cohesive sequence of output-buffer offsets for the\n * sequence of input-buffers.\n */\ntemplate \nstd::vector GetShuffledBufferOffsets(const std::vector &buffer_sizes,\n const std::uint_fast32_t seed = 320981U)\n{\n BufferOffsetT num_buffers = static_cast(buffer_sizes.size());\n\n // We're remapping the i-th buffer to pmt_idxs[i]\n std::mt19937 rng(seed);\n std::vector pmt_idxs(num_buffers);\n std::iota(pmt_idxs.begin(), pmt_idxs.end(), static_cast(0));\n std::shuffle(std::begin(pmt_idxs), std::end(pmt_idxs), rng);\n\n // Compute the offsets using the new mapping\n ByteOffsetT running_offset = {};\n std::vector permuted_offsets;\n permuted_offsets.reserve(num_buffers);\n for (auto permuted_buffer_idx : pmt_idxs)\n {\n permuted_offsets.emplace_back(running_offset);\n running_offset += buffer_sizes[permuted_buffer_idx];\n }\n\n // Generate the scatter indexes that identify where each buffer was mapped to\n std::vector scatter_idxs(num_buffers);\n for (BufferOffsetT i = 0; i < num_buffers; i++)\n {\n scatter_idxs[pmt_idxs[i]] = i;\n }\n\n std::vector new_offsets(num_buffers);\n for (BufferOffsetT i = 0; i < num_buffers; i++)\n {\n new_offsets[i] = permuted_offsets[scatter_idxs[i]];\n }\n\n return new_offsets;\n}\n\n/**\n * @brief Function object class template that takes an offset and returns an iterator at the given\n * offset relative to a fixed base iterator.\n *\n * @tparam IteratorT The random-access iterator type to be returned\n */\ntemplate \nstruct OffsetToPtrOp\n{\n template \n __host__ __device__ __forceinline__ IteratorT operator()(T offset) const\n {\n return base_it + offset;\n }\n IteratorT base_it;\n};\n\nenum class TestDataGen\n{\n // Random offsets into a data segment\n RANDOM,\n\n // Buffers cohesively reside next to each other\n CONSECUTIVE\n};\n\n/**\n * @brief\n *\n * @tparam AtomicT The most granular type being copied. All source and destination pointers will be\n * aligned based on this type, the number of bytes being copied will be an integer multiple of this\n * type's size\n * @tparam BufferOffsetT Type used for indexing into the array of buffers\n * @tparam BufferSizeT Type used for indexing into individual bytes of a buffer (large enough to\n * cover the max buffer size)\n * @tparam ByteOffsetT Type used for indexing into bytes over *all* the buffers' sizes\n */\ntemplate \nvoid RunTest(BufferOffsetT num_buffers,\n BufferSizeT min_buffer_size,\n BufferSizeT max_buffer_size,\n TestDataGen input_gen,\n TestDataGen output_gen)\n{\n using SrcPtrT = uint8_t *;\n\n // Buffer segment data (their offsets and sizes)\n std::vector h_buffer_sizes(num_buffers);\n std::vector h_buffer_src_offsets(num_buffers);\n std::vector h_buffer_dst_offsets(num_buffers);\n\n // Device-side resources\n void *d_in = nullptr;\n void *d_out = nullptr;\n ByteOffsetT *d_buffer_src_offsets = nullptr;\n ByteOffsetT *d_buffer_dst_offsets = nullptr;\n BufferSizeT *d_buffer_sizes = nullptr;\n void *d_temp_storage = nullptr;\n size_t temp_storage_bytes = 0;\n\n // Generate the buffer sizes\n GenerateRandomData(h_buffer_sizes.data(), h_buffer_sizes.size(), min_buffer_size, max_buffer_size);\n\n // Make sure buffer sizes are a multiple of the most granular unit (one AtomicT) being copied\n // (round down)\n for (BufferOffsetT i = 0; i < num_buffers; i++)\n {\n h_buffer_sizes[i] = (h_buffer_sizes[i] / sizeof(AtomicT)) * sizeof(AtomicT);\n }\n\n // Compute the total bytes to be copied\n ByteOffsetT num_total_bytes = 0;\n for (BufferOffsetT i = 0; i < num_buffers; i++)\n {\n if (input_gen == TestDataGen::CONSECUTIVE)\n {\n h_buffer_src_offsets[i] = num_total_bytes;\n }\n if (output_gen == TestDataGen::CONSECUTIVE)\n {\n h_buffer_dst_offsets[i] = num_total_bytes;\n }\n num_total_bytes += h_buffer_sizes[i];\n }\n\n // Shuffle input buffer source-offsets\n std::uint_fast32_t shuffle_seed = 320981U;\n if (input_gen == TestDataGen::RANDOM)\n {\n h_buffer_src_offsets = GetShuffledBufferOffsets(h_buffer_sizes,\n shuffle_seed);\n shuffle_seed += 42;\n }\n\n // Shuffle input buffer source-offsets\n if (output_gen == TestDataGen::RANDOM)\n {\n h_buffer_dst_offsets = GetShuffledBufferOffsets(h_buffer_sizes,\n shuffle_seed);\n }\n\n // Get temporary storage requirements\n CubDebugExit(hipcub::DeviceMemcpy::Batched(d_temp_storage,\n temp_storage_bytes,\n static_cast(nullptr),\n static_cast(nullptr),\n d_buffer_sizes,\n num_buffers));\n\n // Check if there's sufficient device memory to run this test\n std::size_t total_required_mem = num_total_bytes + //\n num_total_bytes + //\n (num_buffers * sizeof(d_buffer_src_offsets[0])) + //\n (num_buffers * sizeof(d_buffer_dst_offsets[0])) + //\n (num_buffers * sizeof(d_buffer_sizes[0])) + //\n temp_storage_bytes; //\n if (TotalGlobalMem() < total_required_mem)\n {\n std::cout\n << \"Skipping the test due to insufficient device memory\\n\" //\n << \" - Required: \" << total_required_mem << \" B, available: \" << TotalGlobalMem() << \" B\\n\" //\n << \" - Skipped test instance: \" //\n << \" -> Min. buffer size: \" << min_buffer_size << \", max. buffer size: \" << max_buffer_size //\n << \", num_buffers: \" << num_buffers //\n << \", in_gen: \" << ((input_gen == TestDataGen::RANDOM) ? \"SHFL\" : \"CONSECUTIVE\") //\n << \", out_gen: \" << ((output_gen == TestDataGen::RANDOM) ? \"SHFL\" : \"CONSECUTIVE\");\n return;\n }\n\n hipEvent_t events[2];\n hipEventCreate(&events[0]);\n hipEventCreate(&events[1]);\n\n hipStream_t stream;\n hipStreamCreate(&stream);\n\n // Allocate device memory\n CubDebugExit(hipMalloc(&d_in, num_total_bytes));\n CubDebugExit(hipMalloc(&d_out, num_total_bytes));\n CubDebugExit(hipMalloc(&d_buffer_src_offsets, num_buffers * sizeof(d_buffer_src_offsets[0])));\n CubDebugExit(hipMalloc(&d_buffer_dst_offsets, num_buffers * sizeof(d_buffer_dst_offsets[0])));\n CubDebugExit(hipMalloc(&d_buffer_sizes, num_buffers * sizeof(d_buffer_sizes[0])));\n CubDebugExit(hipMalloc(&d_temp_storage, temp_storage_bytes));\n\n // Populate the data source with random data\n using RandomInitAliasT = uint16_t;\n std::size_t num_aliased_factor = sizeof(RandomInitAliasT) / sizeof(uint8_t);\n std::size_t num_aliased_units = CUB_QUOTIENT_CEILING(num_total_bytes, num_aliased_factor);\n std::unique_ptr h_in(new uint8_t[num_aliased_units * num_aliased_factor]);\n std::unique_ptr h_out(new uint8_t[num_total_bytes]);\n std::unique_ptr h_gpu_results(new uint8_t[num_total_bytes]);\n\n // Generate random offsets into the random-bits data buffer\n GenerateRandomData(reinterpret_cast(h_in.get()), num_aliased_units);\n\n // Prepare d_buffer_srcs\n OffsetToPtrOp src_transform_op{static_cast(d_in)};\n hipcub::TransformInputIterator, ByteOffsetT *> d_buffer_srcs(\n d_buffer_src_offsets,\n src_transform_op);\n\n // Prepare d_buffer_dsts\n OffsetToPtrOp dst_transform_op{static_cast(d_out)};\n hipcub::TransformInputIterator, ByteOffsetT *> d_buffer_dsts(\n d_buffer_dst_offsets,\n dst_transform_op);\n\n // Prepare random data segment (which serves for the buffer sources)\n CubDebugExit(hipMemcpyAsync(d_in, h_in.get(), num_total_bytes, hipMemcpyHostToDevice, stream));\n\n // Prepare d_buffer_src_offsets\n CubDebugExit(hipMemcpyAsync(d_buffer_src_offsets,\n h_buffer_src_offsets.data(),\n h_buffer_src_offsets.size() * sizeof(h_buffer_src_offsets[0]),\n hipMemcpyHostToDevice,\n stream));\n\n // Prepare d_buffer_dst_offsets\n CubDebugExit(hipMemcpyAsync(d_buffer_dst_offsets,\n h_buffer_dst_offsets.data(),\n h_buffer_dst_offsets.size() * sizeof(h_buffer_dst_offsets[0]),\n hipMemcpyHostToDevice,\n stream));\n\n // Prepare d_buffer_sizes\n CubDebugExit(hipMemcpyAsync(d_buffer_sizes,\n h_buffer_sizes.data(),\n h_buffer_sizes.size() * sizeof(h_buffer_sizes[0]),\n hipMemcpyHostToDevice,\n stream));\n\n // Record event before algorithm\n hipEventRecord(events[0], stream);\n\n // Invoke device-side algorithm being under test\n CubDebugExit(hipcub::DeviceMemcpy::Batched(d_temp_storage,\n temp_storage_bytes,\n d_buffer_srcs,\n d_buffer_dsts,\n d_buffer_sizes,\n num_buffers,\n stream));\n\n // Record event after algorithm\n hipEventRecord(events[1], stream);\n\n // Copy back the output buffer\n CubDebugExit(\n hipMemcpyAsync(h_gpu_results.get(), d_out, num_total_bytes, hipMemcpyDeviceToHost, stream));\n\n // Make sure results have been copied back to the host\n CubDebugExit(hipStreamSynchronize(stream));\n\n // CPU-side result generation for verification\n for (BufferOffsetT i = 0; i < num_buffers; i++)\n {\n std::memcpy(h_out.get() + h_buffer_dst_offsets[i],\n h_in.get() + h_buffer_src_offsets[i],\n h_buffer_sizes[i]);\n }\n\n float duration = 0;\n hipEventElapsedTime(&duration, events[0], events[1]);\n\n#ifdef CUB_TEST_BENCHMARK\n size_t stats_src_offsets = sizeof(ByteOffsetT) * num_buffers;\n size_t stats_dst_offsets = sizeof(ByteOffsetT) * num_buffers;\n size_t stats_sizes = sizeof(BufferSizeT) * num_buffers;\n size_t stats_data_copied = 2 * num_total_bytes;\n\n std::cout\n << \"Min. buffer size: \" << min_buffer_size << \", max. buffer size: \" << max_buffer_size //\n << \", num_buffers: \" << num_buffers //\n << \", in_gen: \" << ((input_gen == TestDataGen::RANDOM) ? \"SHFL\" : \"CONSECUTIVE\") //\n << \", out_gen: \" << ((output_gen == TestDataGen::RANDOM) ? \"SHFL\" : \"CONSECUTIVE\") //\n << \", src size: \" << stats_src_offsets << \", dst size: \" << stats_dst_offsets //\n << \", sizes size: \" << stats_sizes << \", cpy_data_size: \" << stats_data_copied //\n << \", total: \" << (stats_src_offsets + stats_dst_offsets + stats_sizes + stats_data_copied) //\n << \", duration: \" << duration //\n << \", BW: \"\n << ((double)(stats_src_offsets + stats_dst_offsets + stats_sizes + stats_data_copied) /\n 1000000000.0) /\n (duration / 1000.0)\n << \"GB/s \\n\";\n#endif\n\n for (ByteOffsetT i = 0; i < num_total_bytes; i++)\n {\n if (h_gpu_results.get()[i] != h_out.get()[i])\n {\n std::cout << \"Mismatch at index \" << i\n << \", CPU vs. GPU: \" << static_cast(h_gpu_results.get()[i]) << \", \"\n << static_cast(h_out.get()[i]) << \"\\n\";\n }\n AssertEquals(h_out.get()[i], h_gpu_results.get()[i]);\n }\n\n CubDebugExit(hipFree(d_in));\n CubDebugExit(hipFree(d_out));\n CubDebugExit(hipFree(d_buffer_src_offsets));\n CubDebugExit(hipFree(d_buffer_dst_offsets));\n CubDebugExit(hipFree(d_buffer_sizes));\n CubDebugExit(hipFree(d_temp_storage));\n}\n\ntemplate \n__global__ void TestVectorizedCopyKernel(const void *d_in, void *d_out, ByteOffsetT copy_size)\n{\n hipcub::detail::VectorizedCopy(threadIdx.x, d_out, copy_size, d_in);\n}\n\nstruct TupleMemberEqualityOp\n{\n template \n __host__ __device__ __forceinline__ bool operator()(T tuple)\n {\n return thrust::get<0>(tuple) == thrust::get<1>(tuple);\n }\n};\n\n/**\n * @brief Tests the VectorizedCopy for various aligned and misaligned input and output pointers.\n * @tparam VectorT The vector type used for vectorized stores (i.e., one of uint4, uint2, uint32_t)\n */\ntemplate \nvoid TestVectorizedCopy()\n{\n\n constexpr uint32_t threads_per_block = 8;\n\n std::vector in_offsets{0, 1, sizeof(uint32_t) - 1};\n std::vector out_offsets{0, 1, sizeof(VectorT) - 1};\n std::vector copy_sizes{0,\n 1,\n sizeof(uint32_t),\n sizeof(VectorT),\n 2 * threads_per_block * sizeof(VectorT)};\n for (auto copy_sizes_it = std::begin(copy_sizes); copy_sizes_it < std::end(copy_sizes);\n copy_sizes_it++)\n {\n for (auto in_offsets_it = std::begin(in_offsets); in_offsets_it < std::end(in_offsets);\n in_offsets_it++)\n {\n for (auto out_offsets_it = std::begin(out_offsets); out_offsets_it < std::end(out_offsets);\n out_offsets_it++)\n {\n std::size_t in_offset = *in_offsets_it;\n std::size_t out_offset = *out_offsets_it;\n std::size_t copy_size = *copy_sizes_it;\n\n // Prepare data\n const std::size_t alloc_size_in = in_offset + copy_size;\n const std::size_t alloc_size_out = out_offset + copy_size;\n thrust::device_vector data_in(alloc_size_in);\n thrust::device_vector data_out(alloc_size_out);\n thrust::sequence(data_in.begin(), data_in.end(), static_cast(0));\n thrust::fill_n(data_out.begin(), alloc_size_out, static_cast(0x42));\n\n auto d_in = thrust::raw_pointer_cast(data_in.data());\n auto d_out = thrust::raw_pointer_cast(data_out.data());\n\n TestVectorizedCopyKernel\n <<<1, threads_per_block>>>(d_in + in_offset,\n d_out + out_offset,\n static_cast(copy_size));\n auto zip_it = thrust::make_zip_iterator(data_in.begin() + in_offset,\n data_out.begin() + out_offset);\n\n bool success = thrust::all_of(zip_it, zip_it + copy_size, TupleMemberEqualityOp{});\n AssertTrue(success);\n }\n }\n }\n}\n\ntemplate \n__global__ void TestBitPackedCounterKernel(uint32_t *bins,\n uint32_t *increments,\n uint32_t *counts_out,\n uint32_t num_items)\n{\n using BitPackedCounterT =\n hipcub::detail::BitPackedCounter;\n BitPackedCounterT counter{};\n for (uint32_t i = 0; i < num_items; i++)\n {\n counter.Add(bins[i], increments[i]);\n }\n\n for (uint32_t i = 0; i < NUM_ITEMS; i++)\n {\n counts_out[i] = counter.Get(i);\n }\n}\n\n/**\n * @brief Tests BitPackedCounter that's used for computing the histogram of buffer sizes (i.e.,\n * small, medium, large).\n */\ntemplate \nvoid TestBitPackedCounter(const std::uint_fast32_t seed = 320981U)\n{\n\n constexpr uint32_t min_increment = 0;\n constexpr uint32_t max_increment = 4;\n constexpr double avg_increment = static_cast(min_increment) +\n (static_cast(max_increment - min_increment) / 2.0);\n std::uint32_t num_increments = \n static_cast(static_cast(MAX_ITEM_VALUE * NUM_ITEMS) / avg_increment);\n\n // Test input data\n std::array reference_counters{};\n thrust::host_vector h_bins(num_increments);\n thrust::host_vector h_increments(num_increments);\n\n // Generate random test input data\n GenerateRandomData(thrust::raw_pointer_cast(h_bins.data()),\n num_increments,\n 0U,\n NUM_ITEMS - 1U,\n seed);\n GenerateRandomData(thrust::raw_pointer_cast(h_increments.data()),\n num_increments,\n min_increment,\n max_increment,\n (seed + 17));\n\n // Make sure test data does not overflow any of the counters\n for (std::size_t i = 0; i < num_increments; i++)\n {\n // New increment for this bin would overflow => zero this increment\n if (reference_counters[h_bins[i]] + h_increments[i] >= MAX_ITEM_VALUE)\n {\n h_increments[i] = 0;\n }\n else\n {\n reference_counters[h_bins[i]] += h_increments[i];\n }\n }\n\n // Device memory\n thrust::device_vector bins_in(num_increments);\n thrust::device_vector increments_in(num_increments);\n thrust::device_vector counts_out(NUM_ITEMS);\n\n // Initialize device-side test data\n bins_in = h_bins;\n increments_in = h_increments;\n\n // Memory for GPU-generated results\n thrust::host_vector host_counts(num_increments);\n\n // Reset counters to arbitrary random value\n thrust::fill(counts_out.begin(), counts_out.end(), 814920U);\n\n // Run tests with densely bit-packed counters\n TestBitPackedCounterKernel\n <<<1, 1>>>(thrust::raw_pointer_cast(bins_in.data()),\n thrust::raw_pointer_cast(increments_in.data()),\n thrust::raw_pointer_cast(counts_out.data()),\n num_increments);\n\n // Result verification\n host_counts = counts_out;\n for (uint32_t i = 0; i < NUM_ITEMS; i++)\n {\n AssertEquals(reference_counters[i], host_counts[i]);\n }\n\n // Reset counters to arbitrary random value\n thrust::fill(counts_out.begin(), counts_out.end(), 814920U);\n\n // Run tests with bit-packed counters, where bit-count is a power-of-two\n TestBitPackedCounterKernel\n <<<1, 1>>>(thrust::raw_pointer_cast(bins_in.data()),\n thrust::raw_pointer_cast(increments_in.data()),\n thrust::raw_pointer_cast(counts_out.data()),\n num_increments);\n\n // Result verification\n host_counts = counts_out;\n for (uint32_t i = 0; i < NUM_ITEMS; i++)\n {\n AssertEquals(reference_counters[i], host_counts[i]);\n }\n}\n\nint main(int argc, char **argv)\n{\n CommandLineArgs args(argc, argv);\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n //---------------------------------------------------------------------\n // VectorizedCopy tests\n //---------------------------------------------------------------------\n TestVectorizedCopy();\n TestVectorizedCopy();\n\n //---------------------------------------------------------------------\n // BitPackedCounter tests\n //---------------------------------------------------------------------\n TestBitPackedCounter<1, 1>();\n TestBitPackedCounter<1, (0x01U << 16)>();\n TestBitPackedCounter<4, 1>();\n TestBitPackedCounter<4, 2>();\n TestBitPackedCounter<4, 255>();\n TestBitPackedCounter<4, 256>();\n TestBitPackedCounter<8, 1024>();\n TestBitPackedCounter<32, 1>();\n TestBitPackedCounter<32, 256>();\n\n //---------------------------------------------------------------------\n // DeviceMemcpy::Batched tests\n //---------------------------------------------------------------------\n // The most granular type being copied. Buffer's will be aligned and their size be an integer\n // multiple of this type\n using AtomicCopyT = uint8_t;\n\n // Type used for indexing into the array of buffers\n using BufferOffsetT = uint32_t;\n\n // Type used for indexing into individual bytes of a buffer (large enough to cover the max buffer\n using BufferSizeT = uint32_t;\n\n // Type used for indexing into bytes over *all* the buffers' sizes\n using ByteOffsetT = uint32_t;\n\n // Total number of bytes that are targeted to be copied on each run\n const BufferOffsetT target_copy_size = 64U << 20;\n\n // The number of randomly\n constexpr std::size_t num_rnd_buffer_range_tests = 32;\n\n // Each buffer's size will be random within this interval\n std::vector> buffer_size_ranges = {{0, 1},\n {1, 2},\n {0, 16},\n {1, 32},\n {1, 1024},\n {1, 32 * 1024},\n {128 * 1024, 256 * 1024},\n {target_copy_size,\n target_copy_size}};\n\n std::mt19937 rng(0);\n std::uniform_int_distribution size_dist(1, 1000000);\n for (std::size_t i = 0; i < num_rnd_buffer_range_tests; i++)\n {\n auto range_begin = size_dist(rng);\n auto range_end = size_dist(rng);\n if (range_begin > range_end)\n {\n std::swap(range_begin, range_end);\n }\n buffer_size_ranges.push_back({range_begin, range_end});\n }\n\n for (const auto &buffer_size_range : buffer_size_ranges)\n {\n BufferSizeT min_buffer_size =\n static_cast(CUB_ROUND_UP_NEAREST(buffer_size_range.first, sizeof(AtomicCopyT)));\n BufferSizeT max_buffer_size =\n static_cast(CUB_ROUND_UP_NEAREST(buffer_size_range.second,\n static_cast(sizeof(AtomicCopyT))));\n double average_buffer_size = (min_buffer_size + max_buffer_size) / 2.0;\n BufferOffsetT target_num_buffers =\n static_cast(target_copy_size / average_buffer_size);\n\n // Run tests with input buffer being consecutive and output buffers being consecutive\n RunTest(target_num_buffers,\n min_buffer_size,\n max_buffer_size,\n TestDataGen::CONSECUTIVE,\n TestDataGen::CONSECUTIVE);\n\n // Run tests with input buffer being randomly shuffled and output buffers being randomly\n // shuffled\n RunTest(target_num_buffers,\n min_buffer_size,\n max_buffer_size,\n TestDataGen::RANDOM,\n TestDataGen::RANDOM);\n }\n\n //---------------------------------------------------------------------\n // DeviceMemcpy::Batched test with 64-bit offsets\n //---------------------------------------------------------------------\n using ByteOffset64T = uint64_t;\n using BufferSize64T = uint64_t;\n ByteOffset64T large_target_copy_size =\n static_cast(std::numeric_limits::max()) + (128ULL * 1024ULL * 1024ULL);\n // Make sure min_buffer_size is in fact smaller than max buffer size\n constexpr BufferOffsetT single_buffer = 1;\n\n // Run tests with input buffer being consecutive and output buffers being consecutive\n RunTest(single_buffer,\n large_target_copy_size,\n large_target_copy_size,\n TestDataGen::CONSECUTIVE,\n TestDataGen::CONSECUTIVE);\n}\n", "source": "/kaggle/working/cub/test/test_device_batch_memcpy.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of iterator utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n// This file tests deprecated CUB APIs. Silence deprecation warnings:\n#define CUB_IGNORE_DEPRECATED_API\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"test_util.h\"\n\nusing namespace cub;\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nCachingDeviceAllocator g_allocator(true);\n\n//---------------------------------------------------------------------\n// Test kernels\n//---------------------------------------------------------------------\n\n/**\n * Test random access input iterator\n */\ntemplate <\n typename InputIteratorT,\n typename T>\n__global__ void Kernel(\n InputIteratorT d_in,\n T *d_out,\n InputIteratorT *d_itrs)\n{\n d_out[0] = *d_in; // Value at offset 0\n d_out[1] = d_in[100]; // Value at offset 100\n d_out[2] = *(d_in + 1000); // Value at offset 1000\n d_out[3] = *(d_in + 10000); // Value at offset 10000\n\n d_in++;\n d_out[4] = d_in[0]; // Value at offset 1\n\n d_in += 20;\n d_out[5] = d_in[0]; // Value at offset 21\n d_itrs[0] = d_in; // Iterator at offset 21\n\n d_in -= 10;\n d_out[6] = d_in[0]; // Value at offset 11;\n\n d_in -= 11;\n d_out[7] = d_in[0]; // Value at offset 0\n d_itrs[1] = d_in; // Iterator at offset 0\n}\n\n\n\n//---------------------------------------------------------------------\n// Host testing subroutines\n//---------------------------------------------------------------------\n\n\n/**\n * Run iterator test on device\n */\ntemplate <\n typename InputIteratorT,\n typename T,\n int TEST_VALUES>\nvoid Test(\n InputIteratorT d_in,\n T (&h_reference)[TEST_VALUES])\n{\n // Allocate device arrays\n T *d_out = NULL;\n InputIteratorT *d_itrs = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(T) * TEST_VALUES));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_itrs, sizeof(InputIteratorT) * 2));\n\n int compare;\n\n // Run unguarded kernel\n Kernel<<<1, 1>>>(d_in, d_out, d_itrs);\n\n CubDebugExit(cudaPeekAtLastError());\n CubDebugExit(cudaDeviceSynchronize());\n\n // Check results\n compare = CompareDeviceResults(h_reference, d_out, TEST_VALUES, g_verbose, g_verbose);\n printf(\"\\tValues: %s\\n\", (compare) ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Check iterator at offset 21\n InputIteratorT h_itr = d_in + 21;\n compare = CompareDeviceResults(&h_itr, d_itrs, 1, g_verbose, g_verbose);\n printf(\"\\tIterators: %s\\n\", (compare) ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Check iterator at offset 0\n compare = CompareDeviceResults(&d_in, d_itrs + 1, 1, g_verbose, g_verbose);\n printf(\"\\tIterators: %s\\n\", (compare) ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Cleanup\n if (d_out)\n {\n CubDebugExit(g_allocator.DeviceFree(d_out));\n }\n if (d_itrs)\n {\n CubDebugExit(g_allocator.DeviceFree(d_itrs));\n }\n}\n\n/**\n * Test tex-ref texture iterator\n */\ntemplate \nvoid TestTexRef()\n{\n printf(\"\\nTesting tex-ref iterator on type %s\\n\", typeid(T).name()); fflush(stdout);\n\n //\n // Test iterator manipulation in kernel\n //\n\n constexpr int TEST_VALUES = 11000;\n constexpr unsigned int DUMMY_OFFSET = 500;\n constexpr unsigned int DUMMY_TEST_VALUES = TEST_VALUES - DUMMY_OFFSET;\n\n T *h_data = new T[TEST_VALUES];\n for (int i = 0; i < TEST_VALUES; ++i)\n {\n RandomBits(h_data[i]);\n }\n\n // Allocate device arrays\n T *d_data = NULL;\n T *d_dummy = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_data, sizeof(T) * TEST_VALUES));\n CubDebugExit(cudaMemcpy(d_data, h_data, sizeof(T) * TEST_VALUES, cudaMemcpyHostToDevice));\n\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_dummy, sizeof(T) * DUMMY_TEST_VALUES));\n CubDebugExit(cudaMemcpy(d_dummy, h_data + DUMMY_OFFSET, sizeof(T) * DUMMY_TEST_VALUES, cudaMemcpyHostToDevice));\n\n // Initialize reference data\n T h_reference[8];\n h_reference[0] = h_data[0]; // Value at offset 0\n h_reference[1] = h_data[100]; // Value at offset 100\n h_reference[2] = h_data[1000]; // Value at offset 1000\n h_reference[3] = h_data[10000]; // Value at offset 10000\n h_reference[4] = h_data[1]; // Value at offset 1\n h_reference[5] = h_data[21]; // Value at offset 21\n h_reference[6] = h_data[11]; // Value at offset 11\n h_reference[7] = h_data[0]; // Value at offset 0;\n\n // Create and bind ref-based test iterator\n TexRefInputIterator d_ref_itr;\n CubDebugExit(d_ref_itr.BindTexture((CastT*) d_data, sizeof(T) * TEST_VALUES));\n\n // Create and bind dummy iterator of same type to check with interferance\n TexRefInputIterator d_ref_itr2;\n CubDebugExit(d_ref_itr2.BindTexture((CastT*) d_dummy, sizeof(T) * DUMMY_TEST_VALUES));\n\n Test(d_ref_itr, h_reference);\n\n CubDebugExit(d_ref_itr.UnbindTexture());\n CubDebugExit(d_ref_itr2.UnbindTexture());\n\n if (h_data)\n {\n delete[] h_data;\n }\n if (d_data)\n {\n CubDebugExit(g_allocator.DeviceFree(d_data));\n }\n if (d_dummy)\n {\n CubDebugExit(g_allocator.DeviceFree(d_dummy));\n }\n}\n\n/**\n * Run non-integer tests\n */\ntemplate \nvoid Test()\n{\n TestTexRef();\n}\n\n/**\n * Run tests\n */\ntemplate \nvoid Test()\n{\n // Test non-const type\n Test();\n\n // Test non-const type\n Test();\n}\n\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // Evaluate different data types\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n\n printf(\"\\nTest complete\\n\");\n fflush(stdout);\n\n return 0;\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of iterator utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n// This file tests deprecated CUB APIs. Silence deprecation warnings:\n#define CUB_IGNORE_DEPRECATED_API\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"test_util.h\"\n\nusing namespace hipcub;\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nCachingDeviceAllocator g_allocator(true);\n\n//---------------------------------------------------------------------\n// Test kernels\n//---------------------------------------------------------------------\n\n/**\n * Test random access input iterator\n */\ntemplate <\n typename InputIteratorT,\n typename T>\n__global__ void Kernel(\n InputIteratorT d_in,\n T *d_out,\n InputIteratorT *d_itrs)\n{\n d_out[0] = *d_in; // Value at offset 0\n d_out[1] = d_in[100]; // Value at offset 100\n d_out[2] = *(d_in + 1000); // Value at offset 1000\n d_out[3] = *(d_in + 10000); // Value at offset 10000\n\n d_in++;\n d_out[4] = d_in[0]; // Value at offset 1\n\n d_in += 20;\n d_out[5] = d_in[0]; // Value at offset 21\n d_itrs[0] = d_in; // Iterator at offset 21\n\n d_in -= 10;\n d_out[6] = d_in[0]; // Value at offset 11;\n\n d_in -= 11;\n d_out[7] = d_in[0]; // Value at offset 0\n d_itrs[1] = d_in; // Iterator at offset 0\n}\n\n\n\n//---------------------------------------------------------------------\n// Host testing subroutines\n//---------------------------------------------------------------------\n\n\n/**\n * Run iterator test on device\n */\ntemplate <\n typename InputIteratorT,\n typename T,\n int TEST_VALUES>\nvoid Test(\n InputIteratorT d_in,\n T (&h_reference)[TEST_VALUES])\n{\n // Allocate device arrays\n T *d_out = NULL;\n InputIteratorT *d_itrs = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(T) * TEST_VALUES));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_itrs, sizeof(InputIteratorT) * 2));\n\n int compare;\n\n // Run unguarded kernel\n Kernel<<<1, 1>>>(d_in, d_out, d_itrs);\n\n CubDebugExit(hipPeekAtLastError());\n CubDebugExit(hipDeviceSynchronize());\n\n // Check results\n compare = CompareDeviceResults(h_reference, d_out, TEST_VALUES, g_verbose, g_verbose);\n printf(\"\\tValues: %s\\n\", (compare) ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Check iterator at offset 21\n InputIteratorT h_itr = d_in + 21;\n compare = CompareDeviceResults(&h_itr, d_itrs, 1, g_verbose, g_verbose);\n printf(\"\\tIterators: %s\\n\", (compare) ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Check iterator at offset 0\n compare = CompareDeviceResults(&d_in, d_itrs + 1, 1, g_verbose, g_verbose);\n printf(\"\\tIterators: %s\\n\", (compare) ? \"FAIL\" : \"PASS\");\n AssertEquals(0, compare);\n\n // Cleanup\n if (d_out)\n {\n CubDebugExit(g_allocator.DeviceFree(d_out));\n }\n if (d_itrs)\n {\n CubDebugExit(g_allocator.DeviceFree(d_itrs));\n }\n}\n\n/**\n * Test tex-ref texture iterator\n */\ntemplate \nvoid TestTexRef()\n{\n printf(\"\\nTesting tex-ref iterator on type %s\\n\", typeid(T).name()); fflush(stdout);\n\n //\n // Test iterator manipulation in kernel\n //\n\n constexpr int TEST_VALUES = 11000;\n constexpr unsigned int DUMMY_OFFSET = 500;\n constexpr unsigned int DUMMY_TEST_VALUES = TEST_VALUES - DUMMY_OFFSET;\n\n T *h_data = new T[TEST_VALUES];\n for (int i = 0; i < TEST_VALUES; ++i)\n {\n RandomBits(h_data[i]);\n }\n\n // Allocate device arrays\n T *d_data = NULL;\n T *d_dummy = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_data, sizeof(T) * TEST_VALUES));\n CubDebugExit(hipMemcpy(d_data, h_data, sizeof(T) * TEST_VALUES, hipMemcpyHostToDevice));\n\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_dummy, sizeof(T) * DUMMY_TEST_VALUES));\n CubDebugExit(hipMemcpy(d_dummy, h_data + DUMMY_OFFSET, sizeof(T) * DUMMY_TEST_VALUES, hipMemcpyHostToDevice));\n\n // Initialize reference data\n T h_reference[8];\n h_reference[0] = h_data[0]; // Value at offset 0\n h_reference[1] = h_data[100]; // Value at offset 100\n h_reference[2] = h_data[1000]; // Value at offset 1000\n h_reference[3] = h_data[10000]; // Value at offset 10000\n h_reference[4] = h_data[1]; // Value at offset 1\n h_reference[5] = h_data[21]; // Value at offset 21\n h_reference[6] = h_data[11]; // Value at offset 11\n h_reference[7] = h_data[0]; // Value at offset 0;\n\n // Create and bind ref-based test iterator\n TexRefInputIterator d_ref_itr;\n CubDebugExit(d_ref_itr.BindTexture((CastT*) d_data, sizeof(T) * TEST_VALUES));\n\n // Create and bind dummy iterator of same type to check with interferance\n TexRefInputIterator d_ref_itr2;\n CubDebugExit(d_ref_itr2.BindTexture((CastT*) d_dummy, sizeof(T) * DUMMY_TEST_VALUES));\n\n Test(d_ref_itr, h_reference);\n\n CubDebugExit(d_ref_itr.UnbindTexture());\n CubDebugExit(d_ref_itr2.UnbindTexture());\n\n if (h_data)\n {\n delete[] h_data;\n }\n if (d_data)\n {\n CubDebugExit(g_allocator.DeviceFree(d_data));\n }\n if (d_dummy)\n {\n CubDebugExit(g_allocator.DeviceFree(d_dummy));\n }\n}\n\n/**\n * Run non-integer tests\n */\ntemplate \nvoid Test()\n{\n TestTexRef();\n}\n\n/**\n * Run tests\n */\ntemplate \nvoid Test()\n{\n // Test non-const type\n Test();\n\n // Test non-const type\n Test();\n}\n\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // Evaluate different data types\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n Test();\n\n Test();\n Test();\n\n printf(\"\\nTest complete\\n\");\n fflush(stdout);\n\n return 0;\n}\n", "source": "/kaggle/working/cub/test/test_iterator_deprecated.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test evaluation for software global barrier throughput\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n\n#include \n\n#include \"test_util.h\"\n\nusing namespace cub;\n\n\n//---------------------------------------------------------------------\n// Test kernels\n//---------------------------------------------------------------------\n\n/**\n * Kernel that iterates through the specified number of software global barriers\n */\n__global__ void Kernel(\n GridBarrier global_barrier,\n int iterations)\n{\n for (int i = 0; i < iterations; i++)\n {\n global_barrier.Sync();\n }\n}\n\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n cudaError_t retval = cudaSuccess;\n\n // Defaults\n int iterations = 10000;\n int block_size = 128;\n int grid_size = -1;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n\n // Get args\n args.GetCmdLineArgument(\"i\", iterations);\n args.GetCmdLineArgument(\"grid-size\", grid_size);\n args.GetCmdLineArgument(\"block-size\", block_size);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=]\"\n \"[--i=]\"\n \"[--grid-size]\"\n \"[--block-size]\"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // Get device ordinal\n int device_ordinal;\n CubDebugExit(cudaGetDevice(&device_ordinal));\n\n // Get device SM version\n int sm_version = 0;\n CubDebugExit(SmVersion(sm_version, device_ordinal));\n\n // Get SM properties\n int sm_count, max_block_threads, max_sm_occupancy;\n CubDebugExit(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device_ordinal));\n CubDebugExit(cudaDeviceGetAttribute(&max_block_threads, cudaDevAttrMaxThreadsPerBlock, device_ordinal));\n CubDebugExit(MaxSmOccupancy(max_sm_occupancy, EmptyKernel, 32));\n\n // Compute grid size and occupancy\n int occupancy = CUB_MIN((max_block_threads / block_size), max_sm_occupancy);\n\n if (grid_size == -1)\n {\n grid_size = occupancy * sm_count;\n }\n else\n {\n occupancy = grid_size / sm_count;\n }\n\n printf(\"Initializing software global barrier for Kernel<<<%d,%d>>> with %d occupancy\\n\",\n grid_size, block_size, occupancy);\n fflush(stdout);\n\n // Init global barrier\n GridBarrierLifetime global_barrier;\n global_barrier.Setup(grid_size);\n\n // Time kernel\n GpuTimer gpu_timer;\n gpu_timer.Start();\n Kernel<<>>(global_barrier, iterations);\n gpu_timer.Stop();\n\n retval = CubDebug(cudaDeviceSynchronize());\n\n // Output timing results\n float avg_elapsed = gpu_timer.ElapsedMillis() / float(iterations);\n printf(\"%d iterations, %f total elapsed millis, %f avg elapsed millis\\n\",\n iterations,\n gpu_timer.ElapsedMillis(),\n avg_elapsed);\n\n return retval;\n}\n", "output": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test evaluation for software global barrier throughput\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n\n#include \n\n#include \"test_util.h\"\n\nusing namespace hipcub;\n\n\n//---------------------------------------------------------------------\n// Test kernels\n//---------------------------------------------------------------------\n\n/**\n * Kernel that iterates through the specified number of software global barriers\n */\n__global__ void Kernel(\n GridBarrier global_barrier,\n int iterations)\n{\n for (int i = 0; i < iterations; i++)\n {\n global_barrier.Sync();\n }\n}\n\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n hipError_t retval = hipSuccess;\n\n // Defaults\n int iterations = 10000;\n int block_size = 128;\n int grid_size = -1;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n\n // Get args\n args.GetCmdLineArgument(\"i\", iterations);\n args.GetCmdLineArgument(\"grid-size\", grid_size);\n args.GetCmdLineArgument(\"block-size\", block_size);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--device=]\"\n \"[--i=]\"\n \"[--grid-size]\"\n \"[--block-size]\"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n\n // Get device ordinal\n int device_ordinal;\n CubDebugExit(hipGetDevice(&device_ordinal));\n\n // Get device SM version\n int sm_version = 0;\n CubDebugExit(SmVersion(sm_version, device_ordinal));\n\n // Get SM properties\n int sm_count, max_block_threads, max_sm_occupancy;\n CubDebugExit(hipDeviceGetAttribute(&sm_count, hipDeviceAttributeMultiprocessorCount, device_ordinal));\n CubDebugExit(hipDeviceGetAttribute(&max_block_threads, hipDeviceAttributeMaxThreadsPerBlock, device_ordinal));\n CubDebugExit(MaxSmOccupancy(max_sm_occupancy, EmptyKernel, 32));\n\n // Compute grid size and occupancy\n int occupancy = CUB_MIN((max_block_threads / block_size), max_sm_occupancy);\n\n if (grid_size == -1)\n {\n grid_size = occupancy * sm_count;\n }\n else\n {\n occupancy = grid_size / sm_count;\n }\n\n printf(\"Initializing software global barrier for Kernel<<<%d,%d>>> with %d occupancy\\n\",\n grid_size, block_size, occupancy);\n fflush(stdout);\n\n // Init global barrier\n GridBarrierLifetime global_barrier;\n global_barrier.Setup(grid_size);\n\n // Time kernel\n GpuTimer gpu_timer;\n gpu_timer.Start();\n Kernel<<>>(global_barrier, iterations);\n gpu_timer.Stop();\n\n retval = HipcubDebug(hipDeviceSynchronize());\n\n // Output timing results\n float avg_elapsed = gpu_timer.ElapsedMillis() / float(iterations);\n printf(\"%d iterations, %f total elapsed millis, %f avg elapsed millis\\n\",\n iterations,\n gpu_timer.ElapsedMillis(),\n avg_elapsed);\n\n return retval;\n}\n", "source": "/kaggle/working/cub/test/test_grid_barrier.cu"} {"input": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceScan utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#define CUB_STDERR\n\n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \"test_util.h\"\n\n#include \n#include \n\nusing namespace cub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\ndouble g_device_giga_bandwidth;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n\n/**\n * \\brief WrapperFunctor (for precluding test-specialized dispatch to *Sum variants)\n */\ntemplate\nstruct WrapperFunctor\n{\n OpT op;\n\n WrapperFunctor(OpT op) : op(op) {}\n\n template \n __host__ __device__ __forceinline__ auto operator()(const T &a, const U &b) const\n -> decltype(op(a, b))\n {\n return static_cast(op(a, b));\n }\n};\n\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB DeviceScan entrypoints\n//---------------------------------------------------------------------\n\n/**\n * Dispatch to exclusive scan entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__ cudaError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n IsPrimitiveT /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /* d_temp_storage_bytes */,\n cudaError_t * /* d_cdp_error */,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT /* d_out */,\n ScanOpT scan_op,\n InitialValueT initial_value,\n OffsetT num_items)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::ExclusiveScan(d_temp_storage,\n temp_storage_bytes,\n d_in,\n scan_op,\n initial_value,\n num_items);\n }\n return error;\n}\n\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__ cudaError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n IsPrimitiveT /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n cudaError_t * /*d_cdp_error*/,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n ScanOpT scan_op,\n InitialValueT initial_value,\n OffsetT num_items)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::ExclusiveScan(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n scan_op,\n initial_value,\n num_items);\n }\n return error;\n}\n\n/**\n * Dispatch to exclusive sum entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__ cudaError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n Int2Type /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n cudaError_t * /*d_cdp_error*/,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT /* d_out */,\n Sum /*scan_op*/,\n InitialValueT /*initial_value*/,\n OffsetT num_items)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::ExclusiveSum(d_temp_storage,\n temp_storage_bytes,\n d_in,\n num_items);\n }\n return error;\n}\n\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__ cudaError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n Int2Type /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n cudaError_t * /*d_cdp_error*/,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n Sum /*scan_op*/,\n InitialValueT /*initial_value*/,\n OffsetT num_items)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::ExclusiveSum(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items);\n }\n return error;\n}\n\n/**\n * Dispatch to inclusive scan entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__ cudaError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n IsPrimitiveT /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n cudaError_t * /*d_cdp_error*/,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT /* d_out */,\n ScanOpT scan_op,\n NullType /* initial_value */,\n OffsetT num_items)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::InclusiveScan(d_temp_storage,\n temp_storage_bytes,\n d_in,\n scan_op,\n num_items);\n }\n return error;\n}\n\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__ cudaError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n IsPrimitiveT /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n cudaError_t * /*d_cdp_error*/,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n ScanOpT scan_op,\n NullType /*initial_value*/,\n OffsetT num_items)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::InclusiveScan(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n scan_op,\n num_items);\n }\n return error;\n}\n\n/**\n * Dispatch to inclusive sum entrypoint\n */\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__ cudaError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n Int2Type /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n cudaError_t * /*d_cdp_error*/,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT /* d_out */,\n Sum /*scan_op*/,\n NullType /*initial_value*/,\n OffsetT num_items)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::InclusiveSum(d_temp_storage,\n temp_storage_bytes,\n d_in,\n num_items);\n }\n return error;\n}\n\ntemplate \nCUB_RUNTIME_FUNCTION __forceinline__ cudaError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n Int2Type /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n cudaError_t * /*d_cdp_error*/,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n Sum /*scan_op*/,\n NullType /*initial_value*/,\n OffsetT num_items)\n{\n cudaError_t error = cudaSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::InclusiveSum(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items);\n }\n return error;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceScan\n */\ntemplate \n__global__ void CDPDispatchKernel(InPlaceT in_place,\n CubBackendT cub_backend,\n IsPrimitiveT is_primitive,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n ScanOpT scan_op,\n InitialValueT initial_value,\n OffsetT num_items)\n{\n *d_cdp_error = Dispatch(in_place,\n cub_backend,\n is_primitive,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n scan_op,\n initial_value,\n num_items);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Dispatch to CDP kernel\n */\ntemplate \ncudaError_t Dispatch(InPlaceT in_place,\n Int2Type dispatch_to,\n IsPrimitiveT is_primitive,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n cudaError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n ScanOpT scan_op,\n InitialValueT initial_value,\n OffsetT num_items)\n{\n // Invoke kernel to invoke device-side dispatch to CUB backend:\n (void)dispatch_to;\n using CubBackendT = Int2Type;\n CubBackendT cub_backend;\n cudaError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n in_place,\n cub_backend,\n is_primitive,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n scan_op,\n initial_value,\n num_items);\n CubDebugExit(retval);\n\n // Copy out temp_storage_bytes\n CubDebugExit(cudaMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n cudaMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(cudaMemcpy(&retval,\n d_cdp_error,\n sizeof(cudaError_t) * 1,\n cudaMemcpyDeviceToHost));\n return retval;\n}\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\n\n/**\n * Initialize problem\n */\ntemplate \nvoid Initialize(\n GenMode gen_mode,\n T *h_in,\n int num_items)\n{\n for (int i = 0; i < num_items; ++i)\n {\n InitValue(gen_mode, h_in[i], i);\n }\n\n if (g_verbose)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n/**\n * Solve exclusive-scan problem\n */\ntemplate <\n typename InputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\nvoid Solve(\n InputIteratorT h_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n using AccumT = \n cub::detail::accumulator_t<\n ScanOpT, \n InitialValueT, \n cub::detail::value_t>;\n\n if (num_items > 0)\n {\n AccumT val = static_cast(h_in[0]);\n h_reference[0] = initial_value;\n AccumT inclusive = static_cast(scan_op(initial_value, val));\n\n for (int i = 1; i < num_items; ++i)\n {\n val = static_cast(h_in[i]);\n h_reference[i] = static_cast(inclusive);\n inclusive = static_cast(scan_op(inclusive, val));\n }\n }\n}\n\n\n/**\n * Solve inclusive-scan problem\n */\ntemplate <\n typename InputIteratorT,\n typename OutputT,\n typename ScanOpT>\nvoid Solve(\n InputIteratorT h_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n NullType)\n{\n using AccumT = \n cub::detail::accumulator_t<\n ScanOpT, \n cub::detail::value_t, \n cub::detail::value_t>;\n\n if (num_items > 0)\n {\n AccumT inclusive = h_in[0];\n h_reference[0] = static_cast(inclusive);\n\n for (int i = 1; i < num_items; ++i)\n {\n AccumT val = h_in[i];\n inclusive = static_cast(scan_op(inclusive, val));\n h_reference[i] = static_cast(inclusive);\n }\n }\n}\n\ntemplate\nstruct AllocateOutput {\n static void run(OutputT *&d_out, DeviceInputIteratorT, int num_items) {\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(OutputT) * num_items));\n }\n};\n\ntemplate\nstruct AllocateOutput {\n static void run(OutputT *&d_out, OutputT *d_in, int /* num_items */) {\n d_out = d_in;\n }\n};\n\n/**\n * Test DeviceScan for a given problem input\n */\ntemplate <\n Backend BACKEND,\n typename DeviceInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT,\n bool InPlace=false>\nvoid Test(\n DeviceInputIteratorT d_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n using InputT = cub::detail::value_t;\n\n // Allocate device output array\n OutputT *d_out = NULL;\n AllocateOutput::run(d_out, d_in, num_items);\n\n // Allocate CDP device arrays\n size_t *d_temp_storage_bytes = NULL;\n cudaError_t *d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(cudaError_t) * 1));\n\n // Allocate temporary storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(\n Int2Type(),\n Int2Type(),\n Int2Type::PRIMITIVE>(),\n 1,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n scan_op,\n initial_value,\n num_items));\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n if (!InPlace)\n {\n // Clear device output array\n CubDebugExit(cudaMemset(d_out, 0, sizeof(OutputT) * num_items));\n }\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(\n Int2Type(),\n Int2Type(),\n Int2Type::PRIMITIVE>(),\n 1,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n scan_op,\n initial_value,\n num_items));\n\n // Check for correctness (and display results, if specified)\n int compare = CompareDeviceResults(h_reference, d_out, num_items, true, g_verbose);\n printf(\"\\t%s\", compare ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Performance\n if (g_timing_iterations > 0)\n {\n GpuTimer gpu_timer;\n gpu_timer.Start();\n CubDebugExit(Dispatch(Int2Type(),\n Int2Type(),\n Int2Type::PRIMITIVE>(),\n g_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n scan_op,\n initial_value,\n num_items));\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n\n // Display performance\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n float giga_bandwidth = giga_rate * (sizeof(InputT) + sizeof(OutputT));\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s, %.1f%% \"\n \"peak\",\n avg_millis,\n giga_rate,\n giga_bandwidth,\n giga_bandwidth / g_device_giga_bandwidth * 100.0);\n }\n\n printf(\"\\n\\n\");\n\n // Cleanup\n if (!InPlace)\n {\n if (d_out)\n {\n CubDebugExit(g_allocator.DeviceFree(d_out));\n }\n }\n\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare);\n}\n\ntemplate \n__global__ void FillInitValue(InitialValueT *ptr, InitialValueT initial_value) {\n *ptr = initial_value;\n}\n\ntemplate <\n Backend BACKEND,\n typename DeviceInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\ntypename std::enable_if::value>::type\nTestFutureInitValue(\n DeviceInputIteratorT d_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n // Allocate device initial_value\n InitialValueT *d_initial_value = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_initial_value, sizeof(InitialValueT)));\n FillInitValue<<<1, 1>>>(d_initial_value, initial_value);\n\n // Run test\n auto future_init_value = cub::FutureValue(d_initial_value);\n Test(d_in, h_reference, num_items, scan_op, future_init_value);\n\n // Cleanup\n if (d_initial_value) CubDebugExit(g_allocator.DeviceFree(d_initial_value));\n}\n\ntemplate <\n Backend BACKEND,\n typename DeviceInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\ntypename std::enable_if::value>::type\nTestFutureInitValue(\n DeviceInputIteratorT,\n OutputT *,\n int,\n ScanOpT,\n InitialValueT)\n{\n // cub::NullType does not have device pointer, so nothing to do here\n}\n\ntemplate <\n Backend BACKEND,\n typename DeviceInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\ntypename std::enable_if::value>::type\nTestFutureInitValueIter(\n DeviceInputIteratorT d_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n using IterT = cub::ConstantInputIterator;\n IterT iter(initial_value);\n auto future_init_value = cub::FutureValue(iter);\n Test(d_in, h_reference, num_items, scan_op, future_init_value);\n}\n\ntemplate <\n Backend BACKEND,\n typename DeviceInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\ntypename std::enable_if::value>::type\nTestFutureInitValueIter(\n DeviceInputIteratorT,\n OutputT *,\n int,\n ScanOpT,\n InitialValueT)\n{\n // cub::NullType does not have device pointer, so nothing to do here\n}\n\ntemplate \nvoid TestInplace(OutputT *d_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n Test(d_in,\n h_reference,\n num_items,\n scan_op,\n initial_value);\n}\n\ntemplate \nvoid TestInplace(DeviceInputIteratorT, OutputT *, int, ScanOpT, InitialValueT)\n{}\n\n/**\n * Test DeviceScan on pointer type\n */\ntemplate <\n Backend BACKEND,\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\nvoid TestPointer(\n int num_items,\n GenMode gen_mode,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n printf(\"\\nPointer %s %s cub::DeviceScan::%s %d items, %s->%s (%d->%d bytes) , gen-mode %s\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n (std::is_same::value) ? \"Inclusive\" : \"Exclusive\",\n (std::is_same::value) ? \"Sum\" : \"Scan\",\n num_items,\n typeid(InputT).name(), typeid(OutputT).name(), (int) sizeof(InputT), (int) sizeof(OutputT),\n (gen_mode == RANDOM) ? \"RANDOM\" : (gen_mode == INTEGER_SEED) ? \"SEQUENTIAL\" : \"HOMOGENOUS\");\n fflush(stdout);\n\n // Allocate host arrays\n InputT* h_in = new InputT[num_items];\n OutputT* h_reference = new OutputT[num_items];\n\n // Initialize problem and solution\n Initialize(gen_mode, h_in, num_items);\n\n // If the output type is primitive and the operator is cub::Sum, the test\n // dispatcher throws away scan_op and initial_value for exclusive scan.\n // Without an initial_value arg, the accumulator switches to the input value\n // type.\n // Do the same thing here:\n if (Traits::PRIMITIVE &&\n std::is_same::value &&\n !std::is_same::value)\n {\n Solve(h_in, h_reference, num_items, cub::Sum{}, InputT{});\n }\n else\n {\n Solve(h_in, h_reference, num_items, scan_op, initial_value);\n }\n\n // Allocate problem device arrays\n InputT *d_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_in, sizeof(InputT) * num_items));\n\n // Initialize device input\n CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(InputT) * num_items, cudaMemcpyHostToDevice));\n\n // Run Test\n Test(d_in, h_reference, num_items, scan_op, initial_value);\n TestFutureInitValue(d_in, h_reference, num_items, scan_op, initial_value);\n TestInplace(d_in, h_reference, num_items, scan_op, initial_value);\n\n // Cleanup\n if (h_in) delete[] h_in;\n if (h_reference) delete[] h_reference;\n if (d_in) CubDebugExit(g_allocator.DeviceFree(d_in));\n}\n\n\n/**\n * Test DeviceScan on iterator type\n */\ntemplate <\n Backend BACKEND,\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\nvoid TestIterator(\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n printf(\"\\nIterator %s %s cub::DeviceScan::%s %d items, %s->%s (%d->%d bytes)\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n (std::is_same::value) ? \"Inclusive\" : \"Exclusive\",\n (std::is_same::value) ? \"Sum\" : \"Scan\",\n num_items,\n typeid(InputT).name(), typeid(OutputT).name(), (int) sizeof(InputT), (int) sizeof(OutputT));\n fflush(stdout);\n\n // Use a constant iterator as the input\n InputT val = InputT();\n ConstantInputIterator h_in(val);\n\n // Allocate host arrays\n OutputT* h_reference = new OutputT[num_items];\n\n // Initialize problem and solution\n Solve(h_in, h_reference, num_items, scan_op, initial_value);\n\n // Run Test\n Test(h_in, h_reference, num_items, scan_op, initial_value);\n TestFutureInitValueIter(h_in, h_reference, num_items, scan_op, initial_value);\n\n // Cleanup\n if (h_reference) delete[] h_reference;\n}\n\n\n/**\n * Test different gen modes\n */\ntemplate <\n Backend BACKEND,\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\nvoid Test(\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n TestPointer( num_items, UNIFORM, scan_op, initial_value);\n TestPointer( num_items, RANDOM, scan_op, initial_value);\n TestIterator( num_items, scan_op, initial_value);\n}\n\n\n/**\n * Test different dispatch\n */\ntemplate <\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\nvoid Test(\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n#if TEST_CDP == 0\n Test(num_items, scan_op, initial_value);\n#elif TEST_CDP == 1\n Test(num_items, scan_op, initial_value);\n#endif // TEST_CDP\n}\n\n\n/**\n * Test different operators\n */\ntemplate \nvoid TestOp(\n int num_items,\n OutputT identity,\n OutputT initial_value)\n{\n // Exclusive (use identity as initial value because it will dispatch to *Sum variants that don't take initial values)\n Test(num_items, cub::Sum(), identity);\n Test(num_items, cub::Max(), identity);\n\n // Exclusive (non-specialized, so we can test initial-value)\n Test(num_items, WrapperFunctor(cub::Sum()), initial_value);\n Test(num_items, WrapperFunctor(cub::Max()), initial_value);\n\n // Inclusive (no initial value)\n Test(num_items, cub::Sum(), NullType());\n Test(num_items, cub::Max(), NullType());\n}\n\n\n/**\n * Test different input sizes\n */\ntemplate <\n typename InputT,\n typename OutputT>\nvoid TestSize(\n int num_items,\n OutputT identity,\n OutputT initial_value)\n{\n if (num_items < 0)\n {\n TestOp(0, identity, initial_value);\n TestOp(1, identity, initial_value);\n TestOp(100, identity, initial_value);\n TestOp(10000, identity, initial_value);\n TestOp(1000000, identity, initial_value);\n }\n else\n {\n TestOp(num_items, identity, initial_value);\n }\n}\n\nclass CustomInputT\n{\n char m_val{};\n\npublic:\n __host__ __device__ explicit CustomInputT(char val)\n : m_val(val)\n {}\n\n __host__ __device__ int get() const { return static_cast(m_val); }\n};\n\nclass CustomAccumulatorT\n{\n int m_val{0};\n int m_magic_value{42};\n\n __host__ __device__ CustomAccumulatorT(int val)\n : m_val(val)\n {}\n\npublic:\n __host__ __device__ CustomAccumulatorT()\n {}\n\n __host__ __device__ CustomAccumulatorT(const CustomAccumulatorT &in)\n : m_val(in.is_valid() * in.get())\n , m_magic_value(in.is_valid() * 42)\n {}\n\n __host__ __device__ CustomAccumulatorT(const CustomInputT &in)\n : m_val(in.get())\n , m_magic_value(42)\n {}\n\n __host__ __device__ void operator=(const CustomInputT &in)\n {\n if (this->is_valid())\n {\n m_val = in.get();\n }\n }\n\n __host__ __device__ void operator=(const CustomAccumulatorT &in)\n {\n if (this->is_valid() && in.is_valid())\n {\n m_val = in.get();\n }\n }\n\n __host__ __device__ CustomAccumulatorT \n operator+(const CustomInputT &in) const\n {\n const int multiplier = this->is_valid();\n return {(m_val + in.get()) * multiplier};\n }\n\n __host__ __device__ CustomAccumulatorT\n operator+(const CustomAccumulatorT &in) const\n {\n const int multiplier = this->is_valid() && in.is_valid();\n return {(m_val + in.get()) * multiplier};\n }\n\n __host__ __device__ int get() const { return m_val; }\n\n __host__ __device__ bool is_valid() const { return m_magic_value == 42; }\n};\n\nclass CustomOutputT\n{\n int *m_d_ok_count{};\n int m_expected{};\n\npublic:\n __host__ __device__ CustomOutputT(int *d_ok_count, int expected)\n : m_d_ok_count(d_ok_count)\n , m_expected(expected)\n {}\n\n __device__ void operator=(const CustomAccumulatorT &accum) const\n {\n const int ok = accum.is_valid() && (accum.get() == m_expected);\n atomicAdd(m_d_ok_count, ok);\n }\n};\n\n__global__ void InitializeTestAccumulatorTypes(int num_items,\n int *d_ok_count,\n CustomInputT *d_in,\n CustomOutputT *d_out)\n{\n const int idx = static_cast(threadIdx.x + blockIdx.x * blockDim.x);\n\n if (idx < num_items)\n {\n d_in[idx] = CustomInputT(1);\n d_out[idx] = CustomOutputT{d_ok_count, idx};\n }\n}\n\nvoid TestAccumulatorTypes()\n{\n const int num_items = 2 * 1024 * 1024;\n const int block_size = 256;\n const int grid_size = (num_items + block_size - 1) / block_size;\n\n CustomInputT *d_in{};\n CustomOutputT *d_out{};\n CustomAccumulatorT init{};\n int *d_ok_count{};\n\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_ok_count, sizeof(int)));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_out,\n sizeof(CustomOutputT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_in,\n sizeof(CustomInputT) * num_items));\n\n InitializeTestAccumulatorTypes<<>>(num_items,\n d_ok_count,\n d_in,\n d_out);\n\n std::uint8_t *d_temp_storage{};\n std::size_t temp_storage_bytes{};\n\n CubDebugExit(cub::DeviceScan::ExclusiveScan(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n cub::Sum{},\n init,\n num_items));\n\n CubDebugExit(\n g_allocator.DeviceAllocate((void **)&d_temp_storage, temp_storage_bytes));\n CubDebugExit(cudaMemset(d_temp_storage, 1, temp_storage_bytes));\n\n CubDebugExit(cub::DeviceScan::ExclusiveScan(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n cub::Sum{},\n init,\n num_items));\n\n int ok{};\n CubDebugExit(cudaMemcpy(&ok, d_ok_count, sizeof(int), cudaMemcpyDeviceToHost));\n\n AssertEquals(ok, num_items);\n\n CubDebugExit(g_allocator.DeviceFree(d_out));\n CubDebugExit(g_allocator.DeviceFree(d_in));\n CubDebugExit(g_allocator.DeviceFree(d_ok_count));\n}\n\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n int num_items = -1;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n g_device_giga_bandwidth = args.device_giga_bandwidth;\n printf(\"\\n\");\n\n // %PARAM% TEST_CDP cdp 0:1\n // %PARAM% TEST_VALUE_TYPES types 0:1:2\n\n#if TEST_VALUE_TYPES == 0\n\n // Test different input+output data types\n TestSize(num_items, (int)0, (int)99);\n\n // Test same input+output data types\n TestSize(num_items, (unsigned char)0, (unsigned char)99);\n TestSize(num_items, (char)0, (char)99);\n TestSize(num_items, (unsigned short)0, (unsigned short)99);\n TestSize(num_items, (unsigned int)0, (unsigned int)99);\n TestSize(num_items,\n (unsigned long long)0,\n (unsigned long long)99);\n\n#elif TEST_VALUE_TYPES == 1\n\n TestSize(num_items, make_uchar2(0, 0), make_uchar2(17, 21));\n TestSize(num_items, make_char2(0, 0), make_char2(17, 21));\n TestSize(num_items, make_ushort2(0, 0), make_ushort2(17, 21));\n TestSize(num_items, make_uint2(0, 0), make_uint2(17, 21));\n TestSize(num_items,\n make_ulonglong2(0, 0),\n make_ulonglong2(17, 21));\n TestSize(num_items,\n make_uchar4(0, 0, 0, 0),\n make_uchar4(17, 21, 32, 85));\n#elif TEST_VALUE_TYPES == 2\n TestSize(num_items,\n make_char4(0, 0, 0, 0),\n make_char4(17, 21, 32, 85));\n\n TestSize(num_items,\n make_ushort4(0, 0, 0, 0),\n make_ushort4(17, 21, 32, 85));\n TestSize(num_items,\n make_uint4(0, 0, 0, 0),\n make_uint4(17, 21, 32, 85));\n TestSize(num_items,\n make_ulonglong4(0, 0, 0, 0),\n make_ulonglong4(17, 21, 32, 85));\n\n TestSize(num_items,\n TestFoo::MakeTestFoo(0, 0, 0, 0),\n TestFoo::MakeTestFoo(std::numeric_limits::max(),\n std::numeric_limits::max(),\n std::numeric_limits::max(),\n std::numeric_limits::max()));\n\n TestSize(num_items, \n TestBar(0, 0), \n TestBar(std::numeric_limits::max(), \n std::numeric_limits::max()));\n\n TestAccumulatorTypes();\n#endif\n\n return 0;\n}\n\n", "output": "/******************************************************************************\n * Copyright (c) 2011, Duane Merrill. All rights reserved.\n * Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n/******************************************************************************\n * Test of DeviceScan utilities\n ******************************************************************************/\n\n// Ensure printing of CUDA runtime errors to console\n#include \"hip/hip_runtime.h\"\n#define HIPCUB_STDERR\n\n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \"test_util.h\"\n\n#include \n#include \n\nusing namespace hipcub;\n\n\n//---------------------------------------------------------------------\n// Globals, constants and typedefs\n//---------------------------------------------------------------------\n\nbool g_verbose = false;\nint g_timing_iterations = 0;\ndouble g_device_giga_bandwidth;\nCachingDeviceAllocator g_allocator(true);\n\n// Dispatch types\nenum Backend\n{\n CUB, // CUB method\n CDP, // GPU-based (dynamic parallelism) dispatch to CUB method\n};\n\n\n/**\n * \\brief WrapperFunctor (for precluding test-specialized dispatch to *Sum variants)\n */\ntemplate\nstruct WrapperFunctor\n{\n OpT op;\n\n WrapperFunctor(OpT op) : op(op) {}\n\n template \n __host__ __device__ __forceinline__ auto operator()(const T &a, const U &b) const\n -> decltype(op(a, b))\n {\n return static_cast(op(a, b));\n }\n};\n\n\n//---------------------------------------------------------------------\n// Dispatch to different CUB DeviceScan entrypoints\n//---------------------------------------------------------------------\n\n/**\n * Dispatch to exclusive scan entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__ hipError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n IsPrimitiveT /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /* d_temp_storage_bytes */,\n hipError_t * /* d_cdp_error */,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT /* d_out */,\n ScanOpT scan_op,\n InitialValueT initial_value,\n OffsetT num_items)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::ExclusiveScan(d_temp_storage,\n temp_storage_bytes,\n d_in,\n scan_op,\n initial_value,\n num_items);\n }\n return error;\n}\n\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__ hipError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n IsPrimitiveT /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n hipError_t * /*d_cdp_error*/,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n ScanOpT scan_op,\n InitialValueT initial_value,\n OffsetT num_items)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::ExclusiveScan(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n scan_op,\n initial_value,\n num_items);\n }\n return error;\n}\n\n/**\n * Dispatch to exclusive sum entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__ hipError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n Int2Type /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n hipError_t * /*d_cdp_error*/,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT /* d_out */,\n Sum /*scan_op*/,\n InitialValueT /*initial_value*/,\n OffsetT num_items)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::ExclusiveSum(d_temp_storage,\n temp_storage_bytes,\n d_in,\n num_items);\n }\n return error;\n}\n\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__ hipError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n Int2Type /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n hipError_t * /*d_cdp_error*/,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n Sum /*scan_op*/,\n InitialValueT /*initial_value*/,\n OffsetT num_items)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::ExclusiveSum(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items);\n }\n return error;\n}\n\n/**\n * Dispatch to inclusive scan entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__ hipError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n IsPrimitiveT /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n hipError_t * /*d_cdp_error*/,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT /* d_out */,\n ScanOpT scan_op,\n NullType /* initial_value */,\n OffsetT num_items)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::InclusiveScan(d_temp_storage,\n temp_storage_bytes,\n d_in,\n scan_op,\n num_items);\n }\n return error;\n}\n\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__ hipError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n IsPrimitiveT /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n hipError_t * /*d_cdp_error*/,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n ScanOpT scan_op,\n NullType /*initial_value*/,\n OffsetT num_items)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::InclusiveScan(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n scan_op,\n num_items);\n }\n return error;\n}\n\n/**\n * Dispatch to inclusive sum entrypoint\n */\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__ hipError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n Int2Type /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n hipError_t * /*d_cdp_error*/,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT /* d_out */,\n Sum /*scan_op*/,\n NullType /*initial_value*/,\n OffsetT num_items)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::InclusiveSum(d_temp_storage,\n temp_storage_bytes,\n d_in,\n num_items);\n }\n return error;\n}\n\ntemplate \nHIPCUB_RUNTIME_FUNCTION __forceinline__ hipError_t\nDispatch(Int2Type /*in_place*/,\n Int2Type /*dispatch_to*/,\n Int2Type /*is_primitive*/,\n int timing_timing_iterations,\n size_t * /*d_temp_storage_bytes*/,\n hipError_t * /*d_cdp_error*/,\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n Sum /*scan_op*/,\n NullType /*initial_value*/,\n OffsetT num_items)\n{\n hipError_t error = hipSuccess;\n for (int i = 0; i < timing_timing_iterations; ++i)\n {\n error = DeviceScan::InclusiveSum(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n num_items);\n }\n return error;\n}\n\n//---------------------------------------------------------------------\n// CUDA Nested Parallelism Test Kernel\n//---------------------------------------------------------------------\n\n#if TEST_CDP == 1\n\n/**\n * Simple wrapper kernel to invoke DeviceScan\n */\ntemplate \n__global__ void CDPDispatchKernel(InPlaceT in_place,\n CubBackendT cub_backend,\n IsPrimitiveT is_primitive,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n ScanOpT scan_op,\n InitialValueT initial_value,\n OffsetT num_items)\n{\n *d_cdp_error = Dispatch(in_place,\n cub_backend,\n is_primitive,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n scan_op,\n initial_value,\n num_items);\n\n *d_temp_storage_bytes = temp_storage_bytes;\n}\n\n/**\n * Dispatch to CDP kernel\n */\ntemplate \nhipError_t Dispatch(InPlaceT in_place,\n Int2Type dispatch_to,\n IsPrimitiveT is_primitive,\n int timing_timing_iterations,\n size_t *d_temp_storage_bytes,\n hipError_t *d_cdp_error,\n\n void *d_temp_storage,\n size_t &temp_storage_bytes,\n InputIteratorT d_in,\n OutputIteratorT d_out,\n ScanOpT scan_op,\n InitialValueT initial_value,\n OffsetT num_items)\n{\n // Invoke kernel to invoke device-side dispatch to CUB backend:\n (void)dispatch_to;\n using CubBackendT = Int2Type;\n CubBackendT cub_backend;\n hipError_t retval =\n thrust::cuda_cub::launcher::triple_chevron(1, 1, 0, 0)\n .doit(CDPDispatchKernel,\n in_place,\n cub_backend,\n is_primitive,\n timing_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n scan_op,\n initial_value,\n num_items);\n CubDebugExit(retval);\n\n // Copy out temp_storage_bytes\n CubDebugExit(hipMemcpy(&temp_storage_bytes,\n d_temp_storage_bytes,\n sizeof(size_t) * 1,\n hipMemcpyDeviceToHost));\n\n // Copy out error\n CubDebugExit(hipMemcpy(&retval,\n d_cdp_error,\n sizeof(hipError_t) * 1,\n hipMemcpyDeviceToHost));\n return retval;\n}\n\n#endif // TEST_CDP\n\n//---------------------------------------------------------------------\n// Test generation\n//---------------------------------------------------------------------\n\n\n/**\n * Initialize problem\n */\ntemplate \nvoid Initialize(\n GenMode gen_mode,\n T *h_in,\n int num_items)\n{\n for (int i = 0; i < num_items; ++i)\n {\n InitValue(gen_mode, h_in[i], i);\n }\n\n if (g_verbose)\n {\n printf(\"Input:\\n\");\n DisplayResults(h_in, num_items);\n printf(\"\\n\\n\");\n }\n}\n\n/**\n * Solve exclusive-scan problem\n */\ntemplate <\n typename InputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\nvoid Solve(\n InputIteratorT h_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n using AccumT = \n hipcub::detail::accumulator_t<\n ScanOpT, \n InitialValueT, \n hipcub::detail::value_t>;\n\n if (num_items > 0)\n {\n AccumT val = static_cast(h_in[0]);\n h_reference[0] = initial_value;\n AccumT inclusive = static_cast(scan_op(initial_value, val));\n\n for (int i = 1; i < num_items; ++i)\n {\n val = static_cast(h_in[i]);\n h_reference[i] = static_cast(inclusive);\n inclusive = static_cast(scan_op(inclusive, val));\n }\n }\n}\n\n\n/**\n * Solve inclusive-scan problem\n */\ntemplate <\n typename InputIteratorT,\n typename OutputT,\n typename ScanOpT>\nvoid Solve(\n InputIteratorT h_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n NullType)\n{\n using AccumT = \n hipcub::detail::accumulator_t<\n ScanOpT, \n hipcub::detail::value_t, \n hipcub::detail::value_t>;\n\n if (num_items > 0)\n {\n AccumT inclusive = h_in[0];\n h_reference[0] = static_cast(inclusive);\n\n for (int i = 1; i < num_items; ++i)\n {\n AccumT val = h_in[i];\n inclusive = static_cast(scan_op(inclusive, val));\n h_reference[i] = static_cast(inclusive);\n }\n }\n}\n\ntemplate\nstruct AllocateOutput {\n static void run(OutputT *&d_out, DeviceInputIteratorT, int num_items) {\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(OutputT) * num_items));\n }\n};\n\ntemplate\nstruct AllocateOutput {\n static void run(OutputT *&d_out, OutputT *d_in, int /* num_items */) {\n d_out = d_in;\n }\n};\n\n/**\n * Test DeviceScan for a given problem input\n */\ntemplate <\n Backend BACKEND,\n typename DeviceInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT,\n bool InPlace=false>\nvoid Test(\n DeviceInputIteratorT d_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n using InputT = hipcub::detail::value_t;\n\n // Allocate device output array\n OutputT *d_out = NULL;\n AllocateOutput::run(d_out, d_in, num_items);\n\n // Allocate CDP device arrays\n size_t *d_temp_storage_bytes = NULL;\n hipError_t *d_cdp_error = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_temp_storage_bytes, sizeof(size_t) * 1));\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_cdp_error, sizeof(hipError_t) * 1));\n\n // Allocate temporary storage\n void *d_temp_storage = NULL;\n size_t temp_storage_bytes = 0;\n CubDebugExit(Dispatch(\n Int2Type(),\n Int2Type(),\n Int2Type::PRIMITIVE>(),\n 1,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n scan_op,\n initial_value,\n num_items));\n CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));\n\n if (!InPlace)\n {\n // Clear device output array\n CubDebugExit(hipMemset(d_out, 0, sizeof(OutputT) * num_items));\n }\n\n // Run warmup/correctness iteration\n CubDebugExit(Dispatch(\n Int2Type(),\n Int2Type(),\n Int2Type::PRIMITIVE>(),\n 1,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n scan_op,\n initial_value,\n num_items));\n\n // Check for correctness (and display results, if specified)\n int compare = CompareDeviceResults(h_reference, d_out, num_items, true, g_verbose);\n printf(\"\\t%s\", compare ? \"FAIL\" : \"PASS\");\n\n // Flush any stdout/stderr\n fflush(stdout);\n fflush(stderr);\n\n // Performance\n if (g_timing_iterations > 0)\n {\n GpuTimer gpu_timer;\n gpu_timer.Start();\n CubDebugExit(Dispatch(Int2Type(),\n Int2Type(),\n Int2Type::PRIMITIVE>(),\n g_timing_iterations,\n d_temp_storage_bytes,\n d_cdp_error,\n d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n scan_op,\n initial_value,\n num_items));\n gpu_timer.Stop();\n float elapsed_millis = gpu_timer.ElapsedMillis();\n\n // Display performance\n float avg_millis = elapsed_millis / g_timing_iterations;\n float giga_rate = float(num_items) / avg_millis / 1000.0f / 1000.0f;\n float giga_bandwidth = giga_rate * (sizeof(InputT) + sizeof(OutputT));\n printf(\", %.3f avg ms, %.3f billion items/s, %.3f logical GB/s, %.1f%% \"\n \"peak\",\n avg_millis,\n giga_rate,\n giga_bandwidth,\n giga_bandwidth / g_device_giga_bandwidth * 100.0);\n }\n\n printf(\"\\n\\n\");\n\n // Cleanup\n if (!InPlace)\n {\n if (d_out)\n {\n CubDebugExit(g_allocator.DeviceFree(d_out));\n }\n }\n\n if (d_temp_storage_bytes) CubDebugExit(g_allocator.DeviceFree(d_temp_storage_bytes));\n if (d_cdp_error) CubDebugExit(g_allocator.DeviceFree(d_cdp_error));\n if (d_temp_storage) CubDebugExit(g_allocator.DeviceFree(d_temp_storage));\n\n // Correctness asserts\n AssertEquals(0, compare);\n}\n\ntemplate \n__global__ void FillInitValue(InitialValueT *ptr, InitialValueT initial_value) {\n *ptr = initial_value;\n}\n\ntemplate <\n Backend BACKEND,\n typename DeviceInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\ntypename std::enable_if::value>::type\nTestFutureInitValue(\n DeviceInputIteratorT d_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n // Allocate device initial_value\n InitialValueT *d_initial_value = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_initial_value, sizeof(InitialValueT)));\n FillInitValue<<<1, 1>>>(d_initial_value, initial_value);\n\n // Run test\n auto future_init_value = hipcub::FutureValue(d_initial_value);\n Test(d_in, h_reference, num_items, scan_op, future_init_value);\n\n // Cleanup\n if (d_initial_value) CubDebugExit(g_allocator.DeviceFree(d_initial_value));\n}\n\ntemplate <\n Backend BACKEND,\n typename DeviceInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\ntypename std::enable_if::value>::type\nTestFutureInitValue(\n DeviceInputIteratorT,\n OutputT *,\n int,\n ScanOpT,\n InitialValueT)\n{\n // cub::NullType does not have device pointer, so nothing to do here\n}\n\ntemplate <\n Backend BACKEND,\n typename DeviceInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\ntypename std::enable_if::value>::type\nTestFutureInitValueIter(\n DeviceInputIteratorT d_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n using IterT = hipcub::ConstantInputIterator;\n IterT iter(initial_value);\n auto future_init_value = hipcub::FutureValue(iter);\n Test(d_in, h_reference, num_items, scan_op, future_init_value);\n}\n\ntemplate <\n Backend BACKEND,\n typename DeviceInputIteratorT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\ntypename std::enable_if::value>::type\nTestFutureInitValueIter(\n DeviceInputIteratorT,\n OutputT *,\n int,\n ScanOpT,\n InitialValueT)\n{\n // cub::NullType does not have device pointer, so nothing to do here\n}\n\ntemplate \nvoid TestInplace(OutputT *d_in,\n OutputT *h_reference,\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n Test(d_in,\n h_reference,\n num_items,\n scan_op,\n initial_value);\n}\n\ntemplate \nvoid TestInplace(DeviceInputIteratorT, OutputT *, int, ScanOpT, InitialValueT)\n{}\n\n/**\n * Test DeviceScan on pointer type\n */\ntemplate <\n Backend BACKEND,\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\nvoid TestPointer(\n int num_items,\n GenMode gen_mode,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n printf(\"\\nPointer %s %s hipcub::DeviceScan::%s %d items, %s->%s (%d->%d bytes) , gen-mode %s\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n (std::is_same::value) ? \"Inclusive\" : \"Exclusive\",\n (std::is_same::value) ? \"Sum\" : \"Scan\",\n num_items,\n typeid(InputT).name(), typeid(OutputT).name(), (int) sizeof(InputT), (int) sizeof(OutputT),\n (gen_mode == RANDOM) ? \"RANDOM\" : (gen_mode == INTEGER_SEED) ? \"SEQUENTIAL\" : \"HOMOGENOUS\");\n fflush(stdout);\n\n // Allocate host arrays\n InputT* h_in = new InputT[num_items];\n OutputT* h_reference = new OutputT[num_items];\n\n // Initialize problem and solution\n Initialize(gen_mode, h_in, num_items);\n\n // If the output type is primitive and the operator is cub::Sum, the test\n // dispatcher throws away scan_op and initial_value for exclusive scan.\n // Without an initial_value arg, the accumulator switches to the input value\n // type.\n // Do the same thing here:\n if (Traits::PRIMITIVE &&\n std::is_same::value &&\n !std::is_same::value)\n {\n Solve(h_in, h_reference, num_items, hipcub::Sum{}, InputT{});\n }\n else\n {\n Solve(h_in, h_reference, num_items, scan_op, initial_value);\n }\n\n // Allocate problem device arrays\n InputT *d_in = NULL;\n CubDebugExit(g_allocator.DeviceAllocate((void**)&d_in, sizeof(InputT) * num_items));\n\n // Initialize device input\n CubDebugExit(hipMemcpy(d_in, h_in, sizeof(InputT) * num_items, hipMemcpyHostToDevice));\n\n // Run Test\n Test(d_in, h_reference, num_items, scan_op, initial_value);\n TestFutureInitValue(d_in, h_reference, num_items, scan_op, initial_value);\n TestInplace(d_in, h_reference, num_items, scan_op, initial_value);\n\n // Cleanup\n if (h_in) delete[] h_in;\n if (h_reference) delete[] h_reference;\n if (d_in) CubDebugExit(g_allocator.DeviceFree(d_in));\n}\n\n\n/**\n * Test DeviceScan on iterator type\n */\ntemplate <\n Backend BACKEND,\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\nvoid TestIterator(\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n printf(\"\\nIterator %s %s hipcub::DeviceScan::%s %d items, %s->%s (%d->%d bytes)\\n\",\n (BACKEND == CDP) ? \"CDP CUB\" : \"CUB\",\n (std::is_same::value) ? \"Inclusive\" : \"Exclusive\",\n (std::is_same::value) ? \"Sum\" : \"Scan\",\n num_items,\n typeid(InputT).name(), typeid(OutputT).name(), (int) sizeof(InputT), (int) sizeof(OutputT));\n fflush(stdout);\n\n // Use a constant iterator as the input\n InputT val = InputT();\n ConstantInputIterator h_in(val);\n\n // Allocate host arrays\n OutputT* h_reference = new OutputT[num_items];\n\n // Initialize problem and solution\n Solve(h_in, h_reference, num_items, scan_op, initial_value);\n\n // Run Test\n Test(h_in, h_reference, num_items, scan_op, initial_value);\n TestFutureInitValueIter(h_in, h_reference, num_items, scan_op, initial_value);\n\n // Cleanup\n if (h_reference) delete[] h_reference;\n}\n\n\n/**\n * Test different gen modes\n */\ntemplate <\n Backend BACKEND,\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\nvoid Test(\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n TestPointer( num_items, UNIFORM, scan_op, initial_value);\n TestPointer( num_items, RANDOM, scan_op, initial_value);\n TestIterator( num_items, scan_op, initial_value);\n}\n\n\n/**\n * Test different dispatch\n */\ntemplate <\n typename InputT,\n typename OutputT,\n typename ScanOpT,\n typename InitialValueT>\nvoid Test(\n int num_items,\n ScanOpT scan_op,\n InitialValueT initial_value)\n{\n#if TEST_CDP == 0\n Test(num_items, scan_op, initial_value);\n#elif TEST_CDP == 1\n Test(num_items, scan_op, initial_value);\n#endif // TEST_CDP\n}\n\n\n/**\n * Test different operators\n */\ntemplate \nvoid TestOp(\n int num_items,\n OutputT identity,\n OutputT initial_value)\n{\n // Exclusive (use identity as initial value because it will dispatch to *Sum variants that don't take initial values)\n Test(num_items, hipcub::Sum(), identity);\n Test(num_items, hipcub::Max(), identity);\n\n // Exclusive (non-specialized, so we can test initial-value)\n Test(num_items, WrapperFunctor(hipcub::Sum()), initial_value);\n Test(num_items, WrapperFunctor(hipcub::Max()), initial_value);\n\n // Inclusive (no initial value)\n Test(num_items, hipcub::Sum(), NullType());\n Test(num_items, hipcub::Max(), NullType());\n}\n\n\n/**\n * Test different input sizes\n */\ntemplate <\n typename InputT,\n typename OutputT>\nvoid TestSize(\n int num_items,\n OutputT identity,\n OutputT initial_value)\n{\n if (num_items < 0)\n {\n TestOp(0, identity, initial_value);\n TestOp(1, identity, initial_value);\n TestOp(100, identity, initial_value);\n TestOp(10000, identity, initial_value);\n TestOp(1000000, identity, initial_value);\n }\n else\n {\n TestOp(num_items, identity, initial_value);\n }\n}\n\nclass CustomInputT\n{\n char m_val{};\n\npublic:\n __host__ __device__ explicit CustomInputT(char val)\n : m_val(val)\n {}\n\n __host__ __device__ int get() const { return static_cast(m_val); }\n};\n\nclass CustomAccumulatorT\n{\n int m_val{0};\n int m_magic_value{42};\n\n __host__ __device__ CustomAccumulatorT(int val)\n : m_val(val)\n {}\n\npublic:\n __host__ __device__ CustomAccumulatorT()\n {}\n\n __host__ __device__ CustomAccumulatorT(const CustomAccumulatorT &in)\n : m_val(in.is_valid() * in.get())\n , m_magic_value(in.is_valid() * 42)\n {}\n\n __host__ __device__ CustomAccumulatorT(const CustomInputT &in)\n : m_val(in.get())\n , m_magic_value(42)\n {}\n\n __host__ __device__ void operator=(const CustomInputT &in)\n {\n if (this->is_valid())\n {\n m_val = in.get();\n }\n }\n\n __host__ __device__ void operator=(const CustomAccumulatorT &in)\n {\n if (this->is_valid() && in.is_valid())\n {\n m_val = in.get();\n }\n }\n\n __host__ __device__ CustomAccumulatorT \n operator+(const CustomInputT &in) const\n {\n const int multiplier = this->is_valid();\n return {(m_val + in.get()) * multiplier};\n }\n\n __host__ __device__ CustomAccumulatorT\n operator+(const CustomAccumulatorT &in) const\n {\n const int multiplier = this->is_valid() && in.is_valid();\n return {(m_val + in.get()) * multiplier};\n }\n\n __host__ __device__ int get() const { return m_val; }\n\n __host__ __device__ bool is_valid() const { return m_magic_value == 42; }\n};\n\nclass CustomOutputT\n{\n int *m_d_ok_count{};\n int m_expected{};\n\npublic:\n __host__ __device__ CustomOutputT(int *d_ok_count, int expected)\n : m_d_ok_count(d_ok_count)\n , m_expected(expected)\n {}\n\n __device__ void operator=(const CustomAccumulatorT &accum) const\n {\n const int ok = accum.is_valid() && (accum.get() == m_expected);\n atomicAdd(m_d_ok_count, ok);\n }\n};\n\n__global__ void InitializeTestAccumulatorTypes(int num_items,\n int *d_ok_count,\n CustomInputT *d_in,\n CustomOutputT *d_out)\n{\n const int idx = static_cast(threadIdx.x + blockIdx.x * blockDim.x);\n\n if (idx < num_items)\n {\n d_in[idx] = CustomInputT(1);\n d_out[idx] = CustomOutputT{d_ok_count, idx};\n }\n}\n\nvoid TestAccumulatorTypes()\n{\n const int num_items = 2 * 1024 * 1024;\n const int block_size = 256;\n const int grid_size = (num_items + block_size - 1) / block_size;\n\n CustomInputT *d_in{};\n CustomOutputT *d_out{};\n CustomAccumulatorT init{};\n int *d_ok_count{};\n\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_ok_count, sizeof(int)));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_out,\n sizeof(CustomOutputT) * num_items));\n CubDebugExit(g_allocator.DeviceAllocate((void **)&d_in,\n sizeof(CustomInputT) * num_items));\n\n InitializeTestAccumulatorTypes<<>>(num_items,\n d_ok_count,\n d_in,\n d_out);\n\n std::uint8_t *d_temp_storage{};\n std::size_t temp_storage_bytes{};\n\n CubDebugExit(hipcub::DeviceScan::ExclusiveScan(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n hipcub::Sum{},\n init,\n num_items));\n\n CubDebugExit(\n g_allocator.DeviceAllocate((void **)&d_temp_storage, temp_storage_bytes));\n CubDebugExit(hipMemset(d_temp_storage, 1, temp_storage_bytes));\n\n CubDebugExit(hipcub::DeviceScan::ExclusiveScan(d_temp_storage,\n temp_storage_bytes,\n d_in,\n d_out,\n hipcub::Sum{},\n init,\n num_items));\n\n int ok{};\n CubDebugExit(hipMemcpy(&ok, d_ok_count, sizeof(int), hipMemcpyDeviceToHost));\n\n AssertEquals(ok, num_items);\n\n CubDebugExit(g_allocator.DeviceFree(d_out));\n CubDebugExit(g_allocator.DeviceFree(d_in));\n CubDebugExit(g_allocator.DeviceFree(d_ok_count));\n}\n\n\n//---------------------------------------------------------------------\n// Main\n//---------------------------------------------------------------------\n\n/**\n * Main\n */\nint main(int argc, char** argv)\n{\n int num_items = -1;\n\n // Initialize command line\n CommandLineArgs args(argc, argv);\n g_verbose = args.CheckCmdLineFlag(\"v\");\n args.GetCmdLineArgument(\"n\", num_items);\n args.GetCmdLineArgument(\"i\", g_timing_iterations);\n\n // Print usage\n if (args.CheckCmdLineFlag(\"help\"))\n {\n printf(\"%s \"\n \"[--n= \"\n \"[--i= \"\n \"[--device=] \"\n \"[--v] \"\n \"\\n\", argv[0]);\n exit(0);\n }\n\n // Initialize device\n CubDebugExit(args.DeviceInit());\n g_device_giga_bandwidth = args.device_giga_bandwidth;\n printf(\"\\n\");\n\n // %PARAM% TEST_CDP cdp 0:1\n // %PARAM% TEST_VALUE_TYPES types 0:1:2\n\n#if TEST_VALUE_TYPES == 0\n\n // Test different input+output data types\n TestSize(num_items, (int)0, (int)99);\n\n // Test same input+output data types\n TestSize(num_items, (unsigned char)0, (unsigned char)99);\n TestSize(num_items, (char)0, (char)99);\n TestSize(num_items, (unsigned short)0, (unsigned short)99);\n TestSize(num_items, (unsigned int)0, (unsigned int)99);\n TestSize(num_items,\n (unsigned long long)0,\n (unsigned long long)99);\n\n#elif TEST_VALUE_TYPES == 1\n\n TestSize(num_items, make_uchar2(0, 0), make_uchar2(17, 21));\n TestSize(num_items, make_char2(0, 0), make_char2(17, 21));\n TestSize(num_items, make_ushort2(0, 0), make_ushort2(17, 21));\n TestSize(num_items, make_uint2(0, 0), make_uint2(17, 21));\n TestSize(num_items,\n make_ulonglong2(0, 0),\n make_ulonglong2(17, 21));\n TestSize(num_items,\n make_uchar4(0, 0, 0, 0),\n make_uchar4(17, 21, 32, 85));\n#elif TEST_VALUE_TYPES == 2\n TestSize(num_items,\n make_char4(0, 0, 0, 0),\n make_char4(17, 21, 32, 85));\n\n TestSize(num_items,\n make_ushort4(0, 0, 0, 0),\n make_ushort4(17, 21, 32, 85));\n TestSize(num_items,\n make_uint4(0, 0, 0, 0),\n make_uint4(17, 21, 32, 85));\n TestSize(num_items,\n make_ulonglong4(0, 0, 0, 0),\n make_ulonglong4(17, 21, 32, 85));\n\n TestSize(num_items,\n TestFoo::MakeTestFoo(0, 0, 0, 0),\n TestFoo::MakeTestFoo(std::numeric_limits::max(),\n std::numeric_limits::max(),\n std::numeric_limits::max(),\n std::numeric_limits::max()));\n\n TestSize(num_items, \n TestBar(0, 0), \n TestBar(std::numeric_limits::max(), \n std::numeric_limits::max()));\n\n TestAccumulatorTypes();\n#endif\n\n return 0;\n}\n\n", "source": "/kaggle/working/cub/test/test_device_scan.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n\n#include \"cuda_fp16.h\"\n#include \"helper_cuda.h\"\n\n#define NUM_OF_BLOCKS 128\n#define NUM_OF_THREADS 128\n\n__forceinline__ __device__ void reduceInShared_intrinsics(half2 *const v)\n{\n if (threadIdx.x < 64)\n v[threadIdx.x] = __hadd2(v[threadIdx.x], v[threadIdx.x + 64]);\n __syncthreads();\n if (threadIdx.x < 32)\n v[threadIdx.x] = __hadd2(v[threadIdx.x], v[threadIdx.x + 32]);\n __syncthreads();\n if (threadIdx.x < 16)\n v[threadIdx.x] = __hadd2(v[threadIdx.x], v[threadIdx.x + 16]);\n __syncthreads();\n if (threadIdx.x < 8)\n v[threadIdx.x] = __hadd2(v[threadIdx.x], v[threadIdx.x + 8]);\n __syncthreads();\n if (threadIdx.x < 4)\n v[threadIdx.x] = __hadd2(v[threadIdx.x], v[threadIdx.x + 4]);\n __syncthreads();\n if (threadIdx.x < 2)\n v[threadIdx.x] = __hadd2(v[threadIdx.x], v[threadIdx.x + 2]);\n __syncthreads();\n if (threadIdx.x < 1)\n v[threadIdx.x] = __hadd2(v[threadIdx.x], v[threadIdx.x + 1]);\n __syncthreads();\n}\n\n__forceinline__ __device__ void reduceInShared_native(half2 *const v)\n{\n if (threadIdx.x < 64)\n v[threadIdx.x] = v[threadIdx.x] + v[threadIdx.x + 64];\n __syncthreads();\n if (threadIdx.x < 32)\n v[threadIdx.x] = v[threadIdx.x] + v[threadIdx.x + 32];\n __syncthreads();\n if (threadIdx.x < 16)\n v[threadIdx.x] = v[threadIdx.x] + v[threadIdx.x + 16];\n __syncthreads();\n if (threadIdx.x < 8)\n v[threadIdx.x] = v[threadIdx.x] + v[threadIdx.x + 8];\n __syncthreads();\n if (threadIdx.x < 4)\n v[threadIdx.x] = v[threadIdx.x] + v[threadIdx.x + 4];\n __syncthreads();\n if (threadIdx.x < 2)\n v[threadIdx.x] = v[threadIdx.x] + v[threadIdx.x + 2];\n __syncthreads();\n if (threadIdx.x < 1)\n v[threadIdx.x] = v[threadIdx.x] + v[threadIdx.x + 1];\n __syncthreads();\n}\n\n__global__ void\nscalarProductKernel_intrinsics(half2 const *const a, half2 const *const b, float *const results, size_t const size)\n{\n const int stride = gridDim.x * blockDim.x;\n __shared__ half2 shArray[NUM_OF_THREADS];\n\n shArray[threadIdx.x] = __float2half2_rn(0.f);\n half2 value = __float2half2_rn(0.f);\n\n for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < size; i += stride) {\n value = __hfma2(a[i], b[i], value);\n }\n\n shArray[threadIdx.x] = value;\n __syncthreads();\n reduceInShared_intrinsics(shArray);\n\n if (threadIdx.x == 0) {\n half2 result = shArray[0];\n float f_result = __low2float(result) + __high2float(result);\n results[blockIdx.x] = f_result;\n }\n}\n\n__global__ void\nscalarProductKernel_native(half2 const *const a, half2 const *const b, float *const results, size_t const size)\n{\n const int stride = gridDim.x * blockDim.x;\n __shared__ half2 shArray[NUM_OF_THREADS];\n\n half2 value(0.f, 0.f);\n shArray[threadIdx.x] = value;\n\n for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < size; i += stride) {\n value = a[i] * b[i] + value;\n }\n\n shArray[threadIdx.x] = value;\n __syncthreads();\n reduceInShared_native(shArray);\n\n if (threadIdx.x == 0) {\n half2 result = shArray[0];\n float f_result = (float)result.y + (float)result.x;\n results[blockIdx.x] = f_result;\n }\n}\n\nvoid generateInput(half2 *a, size_t size)\n{\n for (size_t i = 0; i < size; ++i) {\n half2 temp;\n temp.x = static_cast(rand() % 4);\n temp.y = static_cast(rand() % 2);\n a[i] = temp;\n }\n}\n\nint main(int argc, char *argv[])\n{\n srand((unsigned int)time(NULL));\n size_t size = NUM_OF_BLOCKS * NUM_OF_THREADS * 16;\n\n half2 *vec[2];\n half2 *devVec[2];\n\n float *results;\n float *devResults;\n\n int devID = findCudaDevice(argc, (const char **)argv);\n\n cudaDeviceProp devProp;\n checkCudaErrors(cudaGetDeviceProperties(&devProp, devID));\n\n if (devProp.major < 5 || (devProp.major == 5 && devProp.minor < 3)) {\n printf(\"ERROR: fp16ScalarProduct requires GPU devices with compute SM 5.3 or \"\n \"higher.\\n\");\n return EXIT_WAIVED;\n }\n\n for (int i = 0; i < 2; ++i) {\n checkCudaErrors(cudaMallocHost((void **)&vec[i], size * sizeof *vec[i]));\n checkCudaErrors(cudaMalloc((void **)&devVec[i], size * sizeof *devVec[i]));\n }\n\n checkCudaErrors(cudaMallocHost((void **)&results, NUM_OF_BLOCKS * sizeof *results));\n checkCudaErrors(cudaMalloc((void **)&devResults, NUM_OF_BLOCKS * sizeof *devResults));\n\n for (int i = 0; i < 2; ++i) {\n generateInput(vec[i], size);\n checkCudaErrors(cudaMemcpy(devVec[i], vec[i], size * sizeof *vec[i], cudaMemcpyHostToDevice));\n }\n\n scalarProductKernel_native<<>>(devVec[0], devVec[1], devResults, size);\n\n checkCudaErrors(cudaMemcpy(results, devResults, NUM_OF_BLOCKS * sizeof *results, cudaMemcpyDeviceToHost));\n\n float result_native = 0;\n for (int i = 0; i < NUM_OF_BLOCKS; ++i) {\n result_native += results[i];\n }\n printf(\"Result native operators\\t: %f \\n\", result_native);\n\n scalarProductKernel_intrinsics<<>>(devVec[0], devVec[1], devResults, size);\n\n checkCudaErrors(cudaMemcpy(results, devResults, NUM_OF_BLOCKS * sizeof *results, cudaMemcpyDeviceToHost));\n\n float result_intrinsics = 0;\n for (int i = 0; i < NUM_OF_BLOCKS; ++i) {\n result_intrinsics += results[i];\n }\n printf(\"Result intrinsics\\t: %f \\n\", result_intrinsics);\n\n printf(\"&&&& fp16ScalarProduct %s\\n\", (fabs(result_intrinsics - result_native) < 0.00001) ? \"PASSED\" : \"FAILED\");\n\n for (int i = 0; i < 2; ++i) {\n checkCudaErrors(cudaFree(devVec[i]));\n checkCudaErrors(cudaFreeHost(vec[i]));\n }\n\n checkCudaErrors(cudaFree(devResults));\n checkCudaErrors(cudaFreeHost(results));\n\n return EXIT_SUCCESS;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n#include \"hip/hip_fp16.h\"\n#include \"hip/hip_runtime_api.h\"\n\n#define NUM_OF_BLOCKS 128\n#define NUM_OF_THREADS 128\n\n__forceinline__ __device__ void reduceInShared_intrinsics(half2 *const v)\n{\n if (threadIdx.x < 64)\n v[threadIdx.x] = __hadd2(v[threadIdx.x], v[threadIdx.x + 64]);\n __syncthreads();\n if (threadIdx.x < 32)\n v[threadIdx.x] = __hadd2(v[threadIdx.x], v[threadIdx.x + 32]);\n __syncthreads();\n if (threadIdx.x < 16)\n v[threadIdx.x] = __hadd2(v[threadIdx.x], v[threadIdx.x + 16]);\n __syncthreads();\n if (threadIdx.x < 8)\n v[threadIdx.x] = __hadd2(v[threadIdx.x], v[threadIdx.x + 8]);\n __syncthreads();\n if (threadIdx.x < 4)\n v[threadIdx.x] = __hadd2(v[threadIdx.x], v[threadIdx.x + 4]);\n __syncthreads();\n if (threadIdx.x < 2)\n v[threadIdx.x] = __hadd2(v[threadIdx.x], v[threadIdx.x + 2]);\n __syncthreads();\n if (threadIdx.x < 1)\n v[threadIdx.x] = __hadd2(v[threadIdx.x], v[threadIdx.x + 1]);\n __syncthreads();\n}\n\n__forceinline__ __device__ void reduceInShared_native(half2 *const v)\n{\n if (threadIdx.x < 64)\n v[threadIdx.x] = v[threadIdx.x] + v[threadIdx.x + 64];\n __syncthreads();\n if (threadIdx.x < 32)\n v[threadIdx.x] = v[threadIdx.x] + v[threadIdx.x + 32];\n __syncthreads();\n if (threadIdx.x < 16)\n v[threadIdx.x] = v[threadIdx.x] + v[threadIdx.x + 16];\n __syncthreads();\n if (threadIdx.x < 8)\n v[threadIdx.x] = v[threadIdx.x] + v[threadIdx.x + 8];\n __syncthreads();\n if (threadIdx.x < 4)\n v[threadIdx.x] = v[threadIdx.x] + v[threadIdx.x + 4];\n __syncthreads();\n if (threadIdx.x < 2)\n v[threadIdx.x] = v[threadIdx.x] + v[threadIdx.x + 2];\n __syncthreads();\n if (threadIdx.x < 1)\n v[threadIdx.x] = v[threadIdx.x] + v[threadIdx.x + 1];\n __syncthreads();\n}\n\n__global__ void\nscalarProductKernel_intrinsics(half2 const *const a, half2 const *const b, float *const results, size_t const size)\n{\n const int stride = gridDim.x * blockDim.x;\n __shared__ half2 shArray[NUM_OF_THREADS];\n\n shArray[threadIdx.x] = __float2half2_rn(0.f);\n half2 value = __float2half2_rn(0.f);\n\n for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < size; i += stride) {\n value = __hfma2(a[i], b[i], value);\n }\n\n shArray[threadIdx.x] = value;\n __syncthreads();\n reduceInShared_intrinsics(shArray);\n\n if (threadIdx.x == 0) {\n half2 result = shArray[0];\n float f_result = __low2float(result) + __high2float(result);\n results[blockIdx.x] = f_result;\n }\n}\n\n__global__ void\nscalarProductKernel_native(half2 const *const a, half2 const *const b, float *const results, size_t const size)\n{\n const int stride = gridDim.x * blockDim.x;\n __shared__ half2 shArray[NUM_OF_THREADS];\n\n half2 value(0.f, 0.f);\n shArray[threadIdx.x] = value;\n\n for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < size; i += stride) {\n value = a[i] * b[i] + value;\n }\n\n shArray[threadIdx.x] = value;\n __syncthreads();\n reduceInShared_native(shArray);\n\n if (threadIdx.x == 0) {\n half2 result = shArray[0];\n float f_result = (float)result.y + (float)result.x;\n results[blockIdx.x] = f_result;\n }\n}\n\nvoid generateInput(half2 *a, size_t size)\n{\n for (size_t i = 0; i < size; ++i) {\n half2 temp;\n temp.x = static_cast(rand() % 4);\n temp.y = static_cast(rand() % 2);\n a[i] = temp;\n }\n}\n\nint main(int argc, char *argv[])\n{\n srand((unsigned int)time(NULL));\n size_t size = NUM_OF_BLOCKS * NUM_OF_THREADS * 16;\n\n half2 *vec[2];\n half2 *devVec[2];\n\n float *results;\n float *devResults;\n\n int devID = findCudaDevice(argc, (const char **)argv);\n\n hipDeviceProp_t devProp;\n checkCudaErrors(hipGetDeviceProperties(&devProp, devID));\n\n if (devProp.major < 5 || (devProp.major == 5 && devProp.minor < 3)) {\n printf(\"ERROR: fp16ScalarProduct requires GPU devices with compute SM 5.3 or \"\n \"higher.\\n\");\n return EXIT_WAIVED;\n }\n\n for (int i = 0; i < 2; ++i) {\n checkCudaErrors(hipHostMalloc((void **)&vec[i], size * sizeof *vec[i]));\n checkCudaErrors(hipMalloc((void **)&devVec[i], size * sizeof *devVec[i]));\n }\n\n checkCudaErrors(hipHostMalloc((void **)&results, NUM_OF_BLOCKS * sizeof *results));\n checkCudaErrors(hipMalloc((void **)&devResults, NUM_OF_BLOCKS * sizeof *devResults));\n\n for (int i = 0; i < 2; ++i) {\n generateInput(vec[i], size);\n checkCudaErrors(hipMemcpy(devVec[i], vec[i], size * sizeof *vec[i], hipMemcpyHostToDevice));\n }\n\n scalarProductKernel_native<<>>(devVec[0], devVec[1], devResults, size);\n\n checkCudaErrors(hipMemcpy(results, devResults, NUM_OF_BLOCKS * sizeof *results, hipMemcpyDeviceToHost));\n\n float result_native = 0;\n for (int i = 0; i < NUM_OF_BLOCKS; ++i) {\n result_native += results[i];\n }\n printf(\"Result native operators\\t: %f \\n\", result_native);\n\n scalarProductKernel_intrinsics<<>>(devVec[0], devVec[1], devResults, size);\n\n checkCudaErrors(hipMemcpy(results, devResults, NUM_OF_BLOCKS * sizeof *results, hipMemcpyDeviceToHost));\n\n float result_intrinsics = 0;\n for (int i = 0; i < NUM_OF_BLOCKS; ++i) {\n result_intrinsics += results[i];\n }\n printf(\"Result intrinsics\\t: %f \\n\", result_intrinsics);\n\n printf(\"&&&& fp16ScalarProduct %s\\n\", (fabs(result_intrinsics - result_native) < 0.00001) ? \"PASSED\" : \"FAILED\");\n\n for (int i = 0; i < 2; ++i) {\n checkCudaErrors(hipFree(devVec[i]));\n checkCudaErrors(hipHostFree(vec[i]));\n }\n\n checkCudaErrors(hipFree(devResults));\n checkCudaErrors(hipHostFree(results));\n\n return EXIT_SUCCESS;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/fp16ScalarProduct/fp16ScalarProduct.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/**\n * Matrix multiplication: C = A * B.\n * Host code.\n *\n * This sample implements matrix multiplication which makes use of shared memory\n * to ensure data reuse, the matrix multiplication is done using tiling approach.\n * It has been written for clarity of exposition to illustrate various CUDA programming\n * principles, not with the goal of providing the most performant generic kernel for matrix multiplication.\n * See also:\n * V. Volkov and J. Demmel, \"Benchmarking GPUs to tune dense linear algebra,\"\n * in Proc. 2008 ACM/IEEE Conf. on Supercomputing (SC '08),\n * Piscataway, NJ: IEEE Press, 2008, pp. Art. 31:1-11.\n */\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n#include \n\n// Helper functions and utilities to work with CUDA\n#include \n#include \n\n/**\n * Matrix multiplication (CUDA Kernel) on the device: C = A * B\n * wA is A's width and wB is B's width\n */\ntemplate __global__ void MatrixMulCUDA(float *C, float *A, float *B, int wA, int wB)\n{\n // Block index\n int bx = blockIdx.x;\n int by = blockIdx.y;\n\n // Thread index\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = wA * BLOCK_SIZE * by;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * bx;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n // Csub is used to store the element of the block sub-matrix\n // that is computed by the thread\n float Csub = 0;\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n // Load the matrices from device memory\n // to shared memory; each thread loads\n // one element of each matrix\n As[ty][tx] = A[a + wA * ty + tx];\n Bs[ty][tx] = B[b + wB * ty + tx];\n\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n // Multiply the two matrices together;\n // each thread computes one element\n // of the block sub-matrix\n#pragma unroll\n\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[ty][k] * Bs[k][tx];\n }\n\n // Synchronize to make sure that the preceding\n // computation is done before loading two new\n // sub-matrices of A and B in the next iteration\n __syncthreads();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes one element\n int c = wB * BLOCK_SIZE * by + BLOCK_SIZE * bx;\n C[c + wB * ty + tx] = Csub;\n}\n\nvoid ConstantInit(float *data, int size, float val)\n{\n for (int i = 0; i < size; ++i) {\n data[i] = val;\n }\n}\n\n/**\n * Run a simple test of matrix multiplication using CUDA\n */\nint MatrixMultiply(int argc, char **argv, int block_size, const dim3 &dimsA, const dim3 &dimsB)\n{\n // Allocate host memory for matrices A and B\n unsigned int size_A = dimsA.x * dimsA.y;\n unsigned int mem_size_A = sizeof(float) * size_A;\n float *h_A;\n checkCudaErrors(cudaMallocHost(&h_A, mem_size_A));\n unsigned int size_B = dimsB.x * dimsB.y;\n unsigned int mem_size_B = sizeof(float) * size_B;\n float *h_B;\n checkCudaErrors(cudaMallocHost(&h_B, mem_size_B));\n cudaStream_t stream;\n\n // Initialize host memory\n const float valB = 0.01f;\n ConstantInit(h_A, size_A, 1.0f);\n ConstantInit(h_B, size_B, valB);\n\n // Allocate device memory\n float *d_A, *d_B, *d_C;\n\n // Allocate host matrix C\n dim3 dimsC(dimsB.x, dimsA.y, 1);\n unsigned int mem_size_C = dimsC.x * dimsC.y * sizeof(float);\n float *h_C;\n checkCudaErrors(cudaMallocHost(&h_C, mem_size_C));\n\n if (h_C == NULL) {\n fprintf(stderr, \"Failed to allocate host matrix C!\\n\");\n exit(EXIT_FAILURE);\n }\n\n checkCudaErrors(cudaMalloc(reinterpret_cast(&d_A), mem_size_A));\n checkCudaErrors(cudaMalloc(reinterpret_cast(&d_B), mem_size_B));\n checkCudaErrors(cudaMalloc(reinterpret_cast(&d_C), mem_size_C));\n // Allocate CUDA events that we'll use for timing\n cudaEvent_t start, stop;\n checkCudaErrors(cudaEventCreate(&start));\n checkCudaErrors(cudaEventCreate(&stop));\n\n checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));\n\n // copy host memory to device\n checkCudaErrors(cudaMemcpyAsync(d_A, h_A, mem_size_A, cudaMemcpyHostToDevice, stream));\n checkCudaErrors(cudaMemcpyAsync(d_B, h_B, mem_size_B, cudaMemcpyHostToDevice, stream));\n\n // Setup execution parameters\n dim3 threads(block_size, block_size);\n dim3 grid(dimsB.x / threads.x, dimsA.y / threads.y);\n\n // Create and start timer\n printf(\"Computing result using CUDA Kernel...\\n\");\n\n // Performs warmup operation using matrixMul CUDA kernel\n if (block_size == 16) {\n MatrixMulCUDA<16><<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n }\n else {\n MatrixMulCUDA<32><<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n }\n\n printf(\"done\\n\");\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n // Record the start event\n checkCudaErrors(cudaEventRecord(start, stream));\n\n // Execute the kernel\n int nIter = 300;\n\n for (int j = 0; j < nIter; j++) {\n if (block_size == 16) {\n MatrixMulCUDA<16><<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n }\n else {\n MatrixMulCUDA<32><<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n }\n }\n\n // Record the stop event\n checkCudaErrors(cudaEventRecord(stop, stream));\n\n // Wait for the stop event to complete\n checkCudaErrors(cudaEventSynchronize(stop));\n\n float msecTotal = 0.0f;\n checkCudaErrors(cudaEventElapsedTime(&msecTotal, start, stop));\n\n // Compute and print the performance\n float msecPerMatrixMul = msecTotal / nIter;\n double flopsPerMatrixMul =\n 2.0 * static_cast(dimsA.x) * static_cast(dimsA.y) * static_cast(dimsB.x);\n double gigaFlops = (flopsPerMatrixMul * 1.0e-9f) / (msecPerMatrixMul / 1000.0f);\n printf(\"Performance= %.2f GFlop/s, Time= %.3f msec, Size= %.0f Ops,\"\n \" WorkgroupSize= %u threads/block\\n\",\n gigaFlops,\n msecPerMatrixMul,\n flopsPerMatrixMul,\n threads.x * threads.y);\n\n // Copy result from device to host\n checkCudaErrors(cudaMemcpyAsync(h_C, d_C, mem_size_C, cudaMemcpyDeviceToHost, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n printf(\"Checking computed result for correctness: \");\n bool correct = true;\n\n // test relative error by the formula\n // |_cpu - _gpu|/<|x|, |y|> < eps\n double eps = 1.e-6; // machine zero\n\n for (int i = 0; i < static_cast(dimsC.x * dimsC.y); i++) {\n double abs_err = fabs(h_C[i] - (dimsA.x * valB));\n double dot_length = dimsA.x;\n double abs_val = fabs(h_C[i]);\n double rel_err = abs_err / abs_val / dot_length;\n\n if (rel_err > eps) {\n printf(\"Error! Matrix[%05d]=%.8f, ref=%.8f error term is > %E\\n\", i, h_C[i], dimsA.x * valB, eps);\n correct = false;\n }\n }\n\n printf(\"%s\\n\", correct ? \"Result = PASS\" : \"Result = FAIL\");\n\n // Clean up memory\n checkCudaErrors(cudaFreeHost(h_A));\n checkCudaErrors(cudaFreeHost(h_B));\n checkCudaErrors(cudaFreeHost(h_C));\n checkCudaErrors(cudaFree(d_A));\n checkCudaErrors(cudaFree(d_B));\n checkCudaErrors(cudaFree(d_C));\n checkCudaErrors(cudaEventDestroy(start));\n checkCudaErrors(cudaEventDestroy(stop));\n printf(\"\\nNOTE: The CUDA Samples are not meant for performance \"\n \"measurements. Results may vary when GPU Boost is enabled.\\n\");\n\n if (correct) {\n return EXIT_SUCCESS;\n }\n else {\n return EXIT_FAILURE;\n }\n}\n\n\n/**\n * Program main\n */\nint main(int argc, char **argv)\n{\n printf(\"[Matrix Multiply Using CUDA] - Starting...\\n\");\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\") || checkCmdLineFlag(argc, (const char **)argv, \"?\")) {\n printf(\"Usage -device=n (n >= 0 for deviceID)\\n\");\n printf(\" -wA=WidthA -hA=HeightA (Width x Height of Matrix A)\\n\");\n printf(\" -wB=WidthB -hB=HeightB (Width x Height of Matrix B)\\n\");\n printf(\" Note: Outer matrix dimensions of A & B matrices\"\n \" must be equal.\\n\");\n\n exit(EXIT_SUCCESS);\n }\n\n // This will pick the best possible CUDA capable device, otherwise\n // override the device ID based on input provided at the command line\n int dev = findCudaDevice(argc, (const char **)argv);\n\n int block_size = 32;\n\n dim3 dimsA(5 * 2 * block_size, 5 * 2 * block_size, 1);\n dim3 dimsB(5 * 4 * block_size, 5 * 2 * block_size, 1);\n\n // width of Matrix A\n if (checkCmdLineFlag(argc, (const char **)argv, \"wA\")) {\n dimsA.x = getCmdLineArgumentInt(argc, (const char **)argv, \"wA\");\n }\n\n // height of Matrix A\n if (checkCmdLineFlag(argc, (const char **)argv, \"hA\")) {\n dimsA.y = getCmdLineArgumentInt(argc, (const char **)argv, \"hA\");\n }\n\n // width of Matrix B\n if (checkCmdLineFlag(argc, (const char **)argv, \"wB\")) {\n dimsB.x = getCmdLineArgumentInt(argc, (const char **)argv, \"wB\");\n }\n\n // height of Matrix B\n if (checkCmdLineFlag(argc, (const char **)argv, \"hB\")) {\n dimsB.y = getCmdLineArgumentInt(argc, (const char **)argv, \"hB\");\n }\n\n if (dimsA.x != dimsB.y) {\n printf(\"Error: outer matrix dimensions must be equal. (%d != %d)\\n\", dimsA.x, dimsB.y);\n exit(EXIT_FAILURE);\n }\n\n printf(\"MatrixA(%d,%d), MatrixB(%d,%d)\\n\", dimsA.x, dimsA.y, dimsB.x, dimsB.y);\n\n checkCudaErrors(cudaProfilerStart());\n int matrix_result = MatrixMultiply(argc, argv, block_size, dimsA, dimsB);\n checkCudaErrors(cudaProfilerStop());\n\n exit(matrix_result);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/**\n * Matrix multiplication: C = A * B.\n * Host code.\n *\n * This sample implements matrix multiplication which makes use of shared memory\n * to ensure data reuse, the matrix multiplication is done using tiling approach.\n * It has been written for clarity of exposition to illustrate various CUDA programming\n * principles, not with the goal of providing the most performant generic kernel for matrix multiplication.\n * See also:\n * V. Volkov and J. Demmel, \"Benchmarking GPUs to tune dense linear algebra,\"\n * in Proc. 2008 ACM/IEEE Conf. on Supercomputing (SC '08),\n * Piscataway, NJ: IEEE Press, 2008, pp. Art. 31:1-11.\n */\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n#include \n\n// Helper functions and utilities to work with CUDA\n#include \n\n/**\n * Matrix multiplication (CUDA Kernel) on the device: C = A * B\n * wA is A's width and wB is B's width\n */\ntemplate __global__ void MatrixMulCUDA(float *C, float *A, float *B, int wA, int wB)\n{\n // Block index\n int bx = blockIdx.x;\n int by = blockIdx.y;\n\n // Thread index\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = wA * BLOCK_SIZE * by;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * bx;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n // Csub is used to store the element of the block sub-matrix\n // that is computed by the thread\n float Csub = 0;\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n // Load the matrices from device memory\n // to shared memory; each thread loads\n // one element of each matrix\n As[ty][tx] = A[a + wA * ty + tx];\n Bs[ty][tx] = B[b + wB * ty + tx];\n\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n // Multiply the two matrices together;\n // each thread computes one element\n // of the block sub-matrix\n#pragma unroll\n\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[ty][k] * Bs[k][tx];\n }\n\n // Synchronize to make sure that the preceding\n // computation is done before loading two new\n // sub-matrices of A and B in the next iteration\n __syncthreads();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes one element\n int c = wB * BLOCK_SIZE * by + BLOCK_SIZE * bx;\n C[c + wB * ty + tx] = Csub;\n}\n\nvoid ConstantInit(float *data, int size, float val)\n{\n for (int i = 0; i < size; ++i) {\n data[i] = val;\n }\n}\n\n/**\n * Run a simple test of matrix multiplication using CUDA\n */\nint MatrixMultiply(int argc, char **argv, int block_size, const dim3 &dimsA, const dim3 &dimsB)\n{\n // Allocate host memory for matrices A and B\n unsigned int size_A = dimsA.x * dimsA.y;\n unsigned int mem_size_A = sizeof(float) * size_A;\n float *h_A;\n checkCudaErrors(hipHostMalloc(&h_A, mem_size_A));\n unsigned int size_B = dimsB.x * dimsB.y;\n unsigned int mem_size_B = sizeof(float) * size_B;\n float *h_B;\n checkCudaErrors(hipHostMalloc(&h_B, mem_size_B));\n hipStream_t stream;\n\n // Initialize host memory\n const float valB = 0.01f;\n ConstantInit(h_A, size_A, 1.0f);\n ConstantInit(h_B, size_B, valB);\n\n // Allocate device memory\n float *d_A, *d_B, *d_C;\n\n // Allocate host matrix C\n dim3 dimsC(dimsB.x, dimsA.y, 1);\n unsigned int mem_size_C = dimsC.x * dimsC.y * sizeof(float);\n float *h_C;\n checkCudaErrors(hipHostMalloc(&h_C, mem_size_C));\n\n if (h_C == NULL) {\n fprintf(stderr, \"Failed to allocate host matrix C!\\n\");\n exit(EXIT_FAILURE);\n }\n\n checkCudaErrors(hipMalloc(reinterpret_cast(&d_A), mem_size_A));\n checkCudaErrors(hipMalloc(reinterpret_cast(&d_B), mem_size_B));\n checkCudaErrors(hipMalloc(reinterpret_cast(&d_C), mem_size_C));\n // Allocate CUDA events that we'll use for timing\n hipEvent_t start, stop;\n checkCudaErrors(hipEventCreate(&start));\n checkCudaErrors(hipEventCreate(&stop));\n\n checkCudaErrors(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));\n\n // copy host memory to device\n checkCudaErrors(hipMemcpyAsync(d_A, h_A, mem_size_A, hipMemcpyHostToDevice, stream));\n checkCudaErrors(hipMemcpyAsync(d_B, h_B, mem_size_B, hipMemcpyHostToDevice, stream));\n\n // Setup execution parameters\n dim3 threads(block_size, block_size);\n dim3 grid(dimsB.x / threads.x, dimsA.y / threads.y);\n\n // Create and start timer\n printf(\"Computing result using CUDA Kernel...\\n\");\n\n // Performs warmup operation using matrixMul CUDA kernel\n if (block_size == 16) {\n MatrixMulCUDA<16><<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n }\n else {\n MatrixMulCUDA<32><<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n }\n\n printf(\"done\\n\");\n checkCudaErrors(hipStreamSynchronize(stream));\n\n // Record the start event\n checkCudaErrors(hipEventRecord(start, stream));\n\n // Execute the kernel\n int nIter = 300;\n\n for (int j = 0; j < nIter; j++) {\n if (block_size == 16) {\n MatrixMulCUDA<16><<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n }\n else {\n MatrixMulCUDA<32><<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n }\n }\n\n // Record the stop event\n checkCudaErrors(hipEventRecord(stop, stream));\n\n // Wait for the stop event to complete\n checkCudaErrors(hipEventSynchronize(stop));\n\n float msecTotal = 0.0f;\n checkCudaErrors(hipEventElapsedTime(&msecTotal, start, stop));\n\n // Compute and print the performance\n float msecPerMatrixMul = msecTotal / nIter;\n double flopsPerMatrixMul =\n 2.0 * static_cast(dimsA.x) * static_cast(dimsA.y) * static_cast(dimsB.x);\n double gigaFlops = (flopsPerMatrixMul * 1.0e-9f) / (msecPerMatrixMul / 1000.0f);\n printf(\"Performance= %.2f GFlop/s, Time= %.3f msec, Size= %.0f Ops,\"\n \" WorkgroupSize= %u threads/block\\n\",\n gigaFlops,\n msecPerMatrixMul,\n flopsPerMatrixMul,\n threads.x * threads.y);\n\n // Copy result from device to host\n checkCudaErrors(hipMemcpyAsync(h_C, d_C, mem_size_C, hipMemcpyDeviceToHost, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n\n printf(\"Checking computed result for correctness: \");\n bool correct = true;\n\n // test relative error by the formula\n // |_cpu - _gpu|/<|x|, |y|> < eps\n double eps = 1.e-6; // machine zero\n\n for (int i = 0; i < static_cast(dimsC.x * dimsC.y); i++) {\n double abs_err = fabs(h_C[i] - (dimsA.x * valB));\n double dot_length = dimsA.x;\n double abs_val = fabs(h_C[i]);\n double rel_err = abs_err / abs_val / dot_length;\n\n if (rel_err > eps) {\n printf(\"Error! Matrix[%05d]=%.8f, ref=%.8f error term is > %E\\n\", i, h_C[i], dimsA.x * valB, eps);\n correct = false;\n }\n }\n\n printf(\"%s\\n\", correct ? \"Result = PASS\" : \"Result = FAIL\");\n\n // Clean up memory\n checkCudaErrors(hipHostFree(h_A));\n checkCudaErrors(hipHostFree(h_B));\n checkCudaErrors(hipHostFree(h_C));\n checkCudaErrors(hipFree(d_A));\n checkCudaErrors(hipFree(d_B));\n checkCudaErrors(hipFree(d_C));\n checkCudaErrors(hipEventDestroy(start));\n checkCudaErrors(hipEventDestroy(stop));\n printf(\"\\nNOTE: The CUDA Samples are not meant for performance \"\n \"measurements. Results may vary when GPU Boost is enabled.\\n\");\n\n if (correct) {\n return EXIT_SUCCESS;\n }\n else {\n return EXIT_FAILURE;\n }\n}\n\n\n/**\n * Program main\n */\nint main(int argc, char **argv)\n{\n printf(\"[Matrix Multiply Using CUDA] - Starting...\\n\");\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\") || checkCmdLineFlag(argc, (const char **)argv, \"?\")) {\n printf(\"Usage -device=n (n >= 0 for deviceID)\\n\");\n printf(\" -wA=WidthA -hA=HeightA (Width x Height of Matrix A)\\n\");\n printf(\" -wB=WidthB -hB=HeightB (Width x Height of Matrix B)\\n\");\n printf(\" Note: Outer matrix dimensions of A & B matrices\"\n \" must be equal.\\n\");\n\n exit(EXIT_SUCCESS);\n }\n\n // This will pick the best possible CUDA capable device, otherwise\n // override the device ID based on input provided at the command line\n int dev = findCudaDevice(argc, (const char **)argv);\n\n int block_size = 32;\n\n dim3 dimsA(5 * 2 * block_size, 5 * 2 * block_size, 1);\n dim3 dimsB(5 * 4 * block_size, 5 * 2 * block_size, 1);\n\n // width of Matrix A\n if (checkCmdLineFlag(argc, (const char **)argv, \"wA\")) {\n dimsA.x = getCmdLineArgumentInt(argc, (const char **)argv, \"wA\");\n }\n\n // height of Matrix A\n if (checkCmdLineFlag(argc, (const char **)argv, \"hA\")) {\n dimsA.y = getCmdLineArgumentInt(argc, (const char **)argv, \"hA\");\n }\n\n // width of Matrix B\n if (checkCmdLineFlag(argc, (const char **)argv, \"wB\")) {\n dimsB.x = getCmdLineArgumentInt(argc, (const char **)argv, \"wB\");\n }\n\n // height of Matrix B\n if (checkCmdLineFlag(argc, (const char **)argv, \"hB\")) {\n dimsB.y = getCmdLineArgumentInt(argc, (const char **)argv, \"hB\");\n }\n\n if (dimsA.x != dimsB.y) {\n printf(\"Error: outer matrix dimensions must be equal. (%d != %d)\\n\", dimsA.x, dimsB.y);\n exit(EXIT_FAILURE);\n }\n\n printf(\"MatrixA(%d,%d), MatrixB(%d,%d)\\n\", dimsA.x, dimsA.y, dimsB.x, dimsB.y);\n\n checkCudaErrors(hipProfilerStart());\n int matrix_result = MatrixMultiply(argc, argv, block_size, dimsA, dimsB);\n checkCudaErrors(hipProfilerStop());\n\n exit(matrix_result);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/matrixMul/matrixMul.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* pitchLinearTexture\n *\n * This example demonstrates how to use textures bound to pitch linear memory.\n * It performs a shift of matrix elements using wrap addressing mode (aka\n * periodic boundary conditions) on two arrays, a pitch linear and a CUDA array,\n * in order to highlight the differences in using each.\n *\n * Textures binding to pitch linear memory is a new feature in CUDA 2.2,\n * and allows use of texture features such as wrap addressing mode and\n * filtering which are not possible with textures bound to regular linear memory\n */\n\n// includes, system\n#include \n\n#ifdef _WIN32\n#define WINDOWS_LEAN_AND_MEAN\n#define NOMINMAX\n#include \n#endif\n\n// Includes CUDA\n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n\n#define NUM_REPS 100 // number of repetitions performed\n#define TILE_DIM 16 // tile/block size\n\nconst char *sSDKsample = \"simplePitchLinearTexture\";\n\n// Auto-Verification Code\nbool bTestResult = true;\n\n////////////////////////////////////////////////////////////////////////////////\n// NB: (1) The second argument \"pitch\" is in elements, not bytes\n// (2) normalized coordinates are used (required for wrap address mode)\n////////////////////////////////////////////////////////////////////////////////\n//! Shifts matrix elements using pitch linear array\n//! @param odata output data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void\nshiftPitchLinear(float *odata, int pitch, int width, int height, int shiftX, int shiftY, cudaTextureObject_t texRefPL)\n{\n int xid = blockIdx.x * blockDim.x + threadIdx.x;\n int yid = blockIdx.y * blockDim.y + threadIdx.y;\n\n odata[yid * pitch + xid] = tex2D(texRefPL, (xid + shiftX) / (float)width, (yid + shiftY) / (float)height);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Shifts matrix elements using regular array\n//! @param odata output data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void\nshiftArray(float *odata, int pitch, int width, int height, int shiftX, int shiftY, cudaTextureObject_t texRefArray)\n{\n int xid = blockIdx.x * blockDim.x + threadIdx.x;\n int yid = blockIdx.y * blockDim.y + threadIdx.y;\n\n odata[yid * pitch + xid] = tex2D(texRefArray, (xid + shiftX) / (float)width, (yid + shiftY) / (float)height);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Declaration, forward\nvoid runTest(int argc, char **argv);\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"%s starting...\\n\\n\", sSDKsample);\n\n runTest(argc, argv);\n\n printf(\"%s completed, returned %s\\n\", sSDKsample, bTestResult ? \"OK\" : \"ERROR!\");\n exit(bTestResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nvoid runTest(int argc, char **argv)\n{\n // Set array size\n const int nx = 2048;\n const int ny = 2048;\n\n // Setup shifts applied to x and y data\n const int x_shift = 5;\n const int y_shift = 7;\n\n if ((nx % TILE_DIM != 0) || (ny % TILE_DIM != 0)) {\n printf(\"nx and ny must be multiples of TILE_DIM\\n\");\n exit(EXIT_FAILURE);\n }\n\n // Setup execution configuration parameters\n dim3 dimGrid(nx / TILE_DIM, ny / TILE_DIM), dimBlock(TILE_DIM, TILE_DIM);\n\n // This will pick the best possible CUDA capable device\n int devID = findCudaDevice(argc, (const char **)argv);\n\n // CUDA events for timing\n cudaEvent_t start, stop;\n cudaEventCreate(&start);\n cudaEventCreate(&stop);\n\n // Host allocation and initialization\n float *h_idata = (float *)malloc(sizeof(float) * nx * ny);\n float *h_odata = (float *)malloc(sizeof(float) * nx * ny);\n float *gold = (float *)malloc(sizeof(float) * nx * ny);\n\n for (int i = 0; i < nx * ny; ++i) {\n h_idata[i] = (float)i;\n }\n\n // Device memory allocation\n // Pitch linear input data\n float *d_idataPL;\n size_t d_pitchBytes;\n\n checkCudaErrors(cudaMallocPitch((void **)&d_idataPL, &d_pitchBytes, nx * sizeof(float), ny));\n\n // Array input data\n cudaArray *d_idataArray;\n cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc();\n\n checkCudaErrors(cudaMallocArray(&d_idataArray, &channelDesc, nx, ny));\n\n // Pitch linear output data\n float *d_odata;\n checkCudaErrors(cudaMallocPitch((void **)&d_odata, &d_pitchBytes, nx * sizeof(float), ny));\n\n // Copy host data to device\n // Pitch linear\n size_t h_pitchBytes = nx * sizeof(float);\n\n checkCudaErrors(\n cudaMemcpy2D(d_idataPL, d_pitchBytes, h_idata, h_pitchBytes, nx * sizeof(float), ny, cudaMemcpyHostToDevice));\n\n // Array\n checkCudaErrors(cudaMemcpyToArray(d_idataArray, 0, 0, h_idata, nx * ny * sizeof(float), cudaMemcpyHostToDevice));\n\n cudaTextureObject_t texRefPL;\n cudaTextureObject_t texRefArray;\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypePitch2D;\n texRes.res.pitch2D.devPtr = d_idataPL;\n texRes.res.pitch2D.desc = channelDesc;\n texRes.res.pitch2D.width = nx;\n texRes.res.pitch2D.height = ny;\n texRes.res.pitch2D.pitchInBytes = h_pitchBytes;\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = cudaFilterModePoint;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&texRefPL, &texRes, &texDescr, NULL));\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = d_idataArray;\n texDescr.normalizedCoords = true;\n texDescr.filterMode = cudaFilterModePoint;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeElementType;\n checkCudaErrors(cudaCreateTextureObject(&texRefArray, &texRes, &texDescr, NULL));\n\n // Reference calculation\n for (int j = 0; j < ny; ++j) {\n int jshift = (j + y_shift) % ny;\n\n for (int i = 0; i < nx; ++i) {\n int ishift = (i + x_shift) % nx;\n gold[j * nx + i] = h_idata[jshift * nx + ishift];\n }\n }\n\n // Run ShiftPitchLinear kernel\n checkCudaErrors(cudaMemset2D(d_odata, d_pitchBytes, 0, nx * sizeof(float), ny));\n\n checkCudaErrors(cudaEventRecord(start, 0));\n\n for (int i = 0; i < NUM_REPS; ++i) {\n shiftPitchLinear<<>>(\n d_odata, (int)(d_pitchBytes / sizeof(float)), nx, ny, x_shift, y_shift, texRefPL);\n }\n\n checkCudaErrors(cudaEventRecord(stop, 0));\n checkCudaErrors(cudaEventSynchronize(stop));\n float timePL;\n checkCudaErrors(cudaEventElapsedTime(&timePL, start, stop));\n\n // Check results\n checkCudaErrors(\n cudaMemcpy2D(h_odata, h_pitchBytes, d_odata, d_pitchBytes, nx * sizeof(float), ny, cudaMemcpyDeviceToHost));\n\n bool res = compareData(gold, h_odata, nx * ny, 0.0f, 0.15f);\n\n bTestResult = true;\n\n if (res == false) {\n printf(\"*** shiftPitchLinear failed ***\\n\");\n bTestResult = false;\n }\n\n // Run ShiftArray kernel\n checkCudaErrors(cudaMemset2D(d_odata, d_pitchBytes, 0, nx * sizeof(float), ny));\n checkCudaErrors(cudaEventRecord(start, 0));\n\n for (int i = 0; i < NUM_REPS; ++i) {\n shiftArray<<>>(\n d_odata, (int)(d_pitchBytes / sizeof(float)), nx, ny, x_shift, y_shift, texRefArray);\n }\n\n checkCudaErrors(cudaEventRecord(stop, 0));\n checkCudaErrors(cudaEventSynchronize(stop));\n float timeArray;\n checkCudaErrors(cudaEventElapsedTime(&timeArray, start, stop));\n\n // Check results\n checkCudaErrors(\n cudaMemcpy2D(h_odata, h_pitchBytes, d_odata, d_pitchBytes, nx * sizeof(float), ny, cudaMemcpyDeviceToHost));\n res = compareData(gold, h_odata, nx * ny, 0.0f, 0.15f);\n\n if (res == false) {\n printf(\"*** shiftArray failed ***\\n\");\n bTestResult = false;\n }\n\n float bandwidthPL = 2.f * 1000.f * nx * ny * sizeof(float) / (1.e+9f) / (timePL / NUM_REPS);\n float bandwidthArray = 2.f * 1000.f * nx * ny * sizeof(float) / (1.e+9f) / (timeArray / NUM_REPS);\n\n printf(\"\\nBandwidth (GB/s) for pitch linear: %.2e; for array: %.2e\\n\", bandwidthPL, bandwidthArray);\n\n float fetchRatePL = nx * ny / 1.e+6f / (timePL / (1000.0f * NUM_REPS));\n float fetchRateArray = nx * ny / 1.e+6f / (timeArray / (1000.0f * NUM_REPS));\n\n printf(\"\\nTexture fetch rate (Mpix/s) for pitch linear: \"\n \"%.2e; for array: %.2e\\n\\n\",\n fetchRatePL,\n fetchRateArray);\n\n // Cleanup\n free(h_idata);\n free(h_odata);\n free(gold);\n\n checkCudaErrors(cudaDestroyTextureObject(texRefPL));\n checkCudaErrors(cudaDestroyTextureObject(texRefArray));\n checkCudaErrors(cudaFree(d_idataPL));\n checkCudaErrors(cudaFreeArray(d_idataArray));\n checkCudaErrors(cudaFree(d_odata));\n\n checkCudaErrors(cudaEventDestroy(start));\n checkCudaErrors(cudaEventDestroy(stop));\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* pitchLinearTexture\n *\n * This example demonstrates how to use textures bound to pitch linear memory.\n * It performs a shift of matrix elements using wrap addressing mode (aka\n * periodic boundary conditions) on two arrays, a pitch linear and a CUDA array,\n * in order to highlight the differences in using each.\n *\n * Textures binding to pitch linear memory is a new feature in CUDA 2.2,\n * and allows use of texture features such as wrap addressing mode and\n * filtering which are not possible with textures bound to regular linear memory\n */\n\n// includes, system\n#include \n\n#ifdef _WIN32\n#define WINDOWS_LEAN_AND_MEAN\n#define NOMINMAX\n#include \n#endif\n\n// Includes CUDA\n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n\n#define NUM_REPS 100 // number of repetitions performed\n#define TILE_DIM 16 // tile/block size\n\nconst char *sSDKsample = \"simplePitchLinearTexture\";\n\n// Auto-Verification Code\nbool bTestResult = true;\n\n////////////////////////////////////////////////////////////////////////////////\n// NB: (1) The second argument \"pitch\" is in elements, not bytes\n// (2) normalized coordinates are used (required for wrap address mode)\n////////////////////////////////////////////////////////////////////////////////\n//! Shifts matrix elements using pitch linear array\n//! @param odata output data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void\nshiftPitchLinear(float *odata, int pitch, int width, int height, int shiftX, int shiftY, hipTextureObject_t texRefPL)\n{\n int xid = blockIdx.x * blockDim.x + threadIdx.x;\n int yid = blockIdx.y * blockDim.y + threadIdx.y;\n\n odata[yid * pitch + xid] = tex2D(texRefPL, (xid + shiftX) / (float)width, (yid + shiftY) / (float)height);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Shifts matrix elements using regular array\n//! @param odata output data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void\nshiftArray(float *odata, int pitch, int width, int height, int shiftX, int shiftY, hipTextureObject_t texRefArray)\n{\n int xid = blockIdx.x * blockDim.x + threadIdx.x;\n int yid = blockIdx.y * blockDim.y + threadIdx.y;\n\n odata[yid * pitch + xid] = tex2D(texRefArray, (xid + shiftX) / (float)width, (yid + shiftY) / (float)height);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Declaration, forward\nvoid runTest(int argc, char **argv);\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"%s starting...\\n\\n\", sSDKsample);\n\n runTest(argc, argv);\n\n printf(\"%s completed, returned %s\\n\", sSDKsample, bTestResult ? \"OK\" : \"ERROR!\");\n exit(bTestResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nvoid runTest(int argc, char **argv)\n{\n // Set array size\n const int nx = 2048;\n const int ny = 2048;\n\n // Setup shifts applied to x and y data\n const int x_shift = 5;\n const int y_shift = 7;\n\n if ((nx % TILE_DIM != 0) || (ny % TILE_DIM != 0)) {\n printf(\"nx and ny must be multiples of TILE_DIM\\n\");\n exit(EXIT_FAILURE);\n }\n\n // Setup execution configuration parameters\n dim3 dimGrid(nx / TILE_DIM, ny / TILE_DIM), dimBlock(TILE_DIM, TILE_DIM);\n\n // This will pick the best possible CUDA capable device\n int devID = findCudaDevice(argc, (const char **)argv);\n\n // CUDA events for timing\n hipEvent_t start, stop;\n hipEventCreate(&start);\n hipEventCreate(&stop);\n\n // Host allocation and initialization\n float *h_idata = (float *)malloc(sizeof(float) * nx * ny);\n float *h_odata = (float *)malloc(sizeof(float) * nx * ny);\n float *gold = (float *)malloc(sizeof(float) * nx * ny);\n\n for (int i = 0; i < nx * ny; ++i) {\n h_idata[i] = (float)i;\n }\n\n // Device memory allocation\n // Pitch linear input data\n float *d_idataPL;\n size_t d_pitchBytes;\n\n checkCudaErrors(hipMallocPitch((void **)&d_idataPL, &d_pitchBytes, nx * sizeof(float), ny));\n\n // Array input data\n hipArray *d_idataArray;\n hipChannelFormatDesc channelDesc = hipCreateChannelDesc();\n\n checkCudaErrors(hipMallocArray(&d_idataArray, &channelDesc, nx, ny));\n\n // Pitch linear output data\n float *d_odata;\n checkCudaErrors(hipMallocPitch((void **)&d_odata, &d_pitchBytes, nx * sizeof(float), ny));\n\n // Copy host data to device\n // Pitch linear\n size_t h_pitchBytes = nx * sizeof(float);\n\n checkCudaErrors(\n hipMemcpy2D(d_idataPL, d_pitchBytes, h_idata, h_pitchBytes, nx * sizeof(float), ny, hipMemcpyHostToDevice));\n\n // Array\n checkCudaErrors(hipMemcpyToArray(d_idataArray, 0, 0, h_idata, nx * ny * sizeof(float), hipMemcpyHostToDevice));\n\n hipTextureObject_t texRefPL;\n hipTextureObject_t texRefArray;\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypePitch2D;\n texRes.res.pitch2D.devPtr = d_idataPL;\n texRes.res.pitch2D.desc = channelDesc;\n texRes.res.pitch2D.width = nx;\n texRes.res.pitch2D.height = ny;\n texRes.res.pitch2D.pitchInBytes = h_pitchBytes;\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = hipFilterModePoint;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&texRefPL, &texRes, &texDescr, NULL));\n memset(&texRes, 0, sizeof(hipResourceDesc));\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = d_idataArray;\n texDescr.normalizedCoords = true;\n texDescr.filterMode = hipFilterModePoint;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeElementType;\n checkCudaErrors(hipCreateTextureObject(&texRefArray, &texRes, &texDescr, NULL));\n\n // Reference calculation\n for (int j = 0; j < ny; ++j) {\n int jshift = (j + y_shift) % ny;\n\n for (int i = 0; i < nx; ++i) {\n int ishift = (i + x_shift) % nx;\n gold[j * nx + i] = h_idata[jshift * nx + ishift];\n }\n }\n\n // Run ShiftPitchLinear kernel\n checkCudaErrors(hipMemset2D(d_odata, d_pitchBytes, 0, nx * sizeof(float), ny));\n\n checkCudaErrors(hipEventRecord(start, 0));\n\n for (int i = 0; i < NUM_REPS; ++i) {\n shiftPitchLinear<<>>(\n d_odata, (int)(d_pitchBytes / sizeof(float)), nx, ny, x_shift, y_shift, texRefPL);\n }\n\n checkCudaErrors(hipEventRecord(stop, 0));\n checkCudaErrors(hipEventSynchronize(stop));\n float timePL;\n checkCudaErrors(hipEventElapsedTime(&timePL, start, stop));\n\n // Check results\n checkCudaErrors(\n hipMemcpy2D(h_odata, h_pitchBytes, d_odata, d_pitchBytes, nx * sizeof(float), ny, hipMemcpyDeviceToHost));\n\n bool res = compareData(gold, h_odata, nx * ny, 0.0f, 0.15f);\n\n bTestResult = true;\n\n if (res == false) {\n printf(\"*** shiftPitchLinear failed ***\\n\");\n bTestResult = false;\n }\n\n // Run ShiftArray kernel\n checkCudaErrors(hipMemset2D(d_odata, d_pitchBytes, 0, nx * sizeof(float), ny));\n checkCudaErrors(hipEventRecord(start, 0));\n\n for (int i = 0; i < NUM_REPS; ++i) {\n shiftArray<<>>(\n d_odata, (int)(d_pitchBytes / sizeof(float)), nx, ny, x_shift, y_shift, texRefArray);\n }\n\n checkCudaErrors(hipEventRecord(stop, 0));\n checkCudaErrors(hipEventSynchronize(stop));\n float timeArray;\n checkCudaErrors(hipEventElapsedTime(&timeArray, start, stop));\n\n // Check results\n checkCudaErrors(\n hipMemcpy2D(h_odata, h_pitchBytes, d_odata, d_pitchBytes, nx * sizeof(float), ny, hipMemcpyDeviceToHost));\n res = compareData(gold, h_odata, nx * ny, 0.0f, 0.15f);\n\n if (res == false) {\n printf(\"*** shiftArray failed ***\\n\");\n bTestResult = false;\n }\n\n float bandwidthPL = 2.f * 1000.f * nx * ny * sizeof(float) / (1.e+9f) / (timePL / NUM_REPS);\n float bandwidthArray = 2.f * 1000.f * nx * ny * sizeof(float) / (1.e+9f) / (timeArray / NUM_REPS);\n\n printf(\"\\nBandwidth (GB/s) for pitch linear: %.2e; for array: %.2e\\n\", bandwidthPL, bandwidthArray);\n\n float fetchRatePL = nx * ny / 1.e+6f / (timePL / (1000.0f * NUM_REPS));\n float fetchRateArray = nx * ny / 1.e+6f / (timeArray / (1000.0f * NUM_REPS));\n\n printf(\"\\nTexture fetch rate (Mpix/s) for pitch linear: \"\n \"%.2e; for array: %.2e\\n\\n\",\n fetchRatePL,\n fetchRateArray);\n\n // Cleanup\n free(h_idata);\n free(h_odata);\n free(gold);\n\n checkCudaErrors(hipDestroyTextureObject(texRefPL));\n checkCudaErrors(hipDestroyTextureObject(texRefArray));\n checkCudaErrors(hipFree(d_idataPL));\n checkCudaErrors(hipFreeArray(d_idataArray));\n checkCudaErrors(hipFree(d_odata));\n\n checkCudaErrors(hipEventDestroy(start));\n checkCudaErrors(hipEventDestroy(stop));\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simplePitchLinearTexture/simplePitchLinearTexture.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n/* Vector addition: C = A + B.\n *\n * This sample is a very basic sample that implements element by element\n * vector addition. It is the same as the sample illustrating Chapter 3\n * of the programming guide with some additions like error checking.\n *\n */\n\n// Device code\nextern \"C\" __global__ void VecAdd_kernel(const float *A, const float *B, float *C, int N)\n{\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (i < N)\n C[i] = A[i] + B[i];\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n/* Vector addition: C = A + B.\n *\n * This sample is a very basic sample that implements element by element\n * vector addition. It is the same as the sample illustrating Chapter 3\n * of the programming guide with some additions like error checking.\n *\n */\n\n// Device code\n#include \"hip/hip_runtime.h\"\nextern \"C\" __global__ void VecAdd_kernel(const float *A, const float *B, float *C, int N)\n{\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (i < N)\n C[i] = A[i] + B[i];\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/vectorAddDrv/vectorAdd_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#ifndef MAX\n#define MAX(a, b) (a > b ? a : b)\n#endif\n\n/* Add two vectors on the GPU */\n__global__ void vectorAddGPU(float *a, float *b, float *c, int N)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx < N) {\n c[idx] = a[idx] + b[idx];\n }\n}\n\n// Allocate generic memory with malloc() and pin it laster instead of using\n// cudaHostAlloc()\nbool bPinGenericMemory = false;\n\n// Macro to aligned up to the memory size in question\n#define MEMORY_ALIGNMENT 4096\n#define ALIGN_UP(x, size) (((size_t)x + (size - 1)) & (~(size - 1)))\n\nint main(int argc, char **argv)\n{\n int n, nelem, deviceCount;\n int idev = 0; // use default device 0\n char *device = NULL;\n unsigned int flags;\n size_t bytes;\n float *a, *b, *c; // Pinned memory allocated on the CPU\n float *a_UA, *b_UA, *c_UA; // Non-4K Aligned Pinned memory on the CPU\n float *d_a, *d_b, *d_c; // Device pointers for mapped memory\n float errorNorm, refNorm, ref, diff;\n cudaDeviceProp deviceProp;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\")) {\n printf(\"Usage: simpleZeroCopy [OPTION]\\n\\n\");\n printf(\"Options:\\n\");\n printf(\" --device=[device #] Specify the device to be used\\n\");\n printf(\" --use_generic_memory (optional) use generic page-aligned for system \"\n \"memory\\n\");\n return EXIT_SUCCESS;\n }\n\n /* Get the device selected by the user or default to 0, and then set it. */\n if (getCmdLineArgumentString(argc, (const char **)argv, \"device\", &device)) {\n cudaGetDeviceCount(&deviceCount);\n idev = atoi(device);\n\n if (idev >= deviceCount || idev < 0) {\n fprintf(stderr, \"Device number %d is invalid, will use default CUDA device 0.\\n\", idev);\n idev = 0;\n }\n }\n\n // if GPU found supports SM 1.2, then continue, otherwise we exit\n if (!checkCudaCapabilities(1, 2)) {\n exit(EXIT_SUCCESS);\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"use_generic_memory\")) {\n#if defined(__APPLE__) || defined(MACOSX)\n bPinGenericMemory = false; // Generic Pinning of System Paged memory is not\n // currently supported on Mac OSX\n#else\n bPinGenericMemory = true;\n#endif\n }\n\n if (bPinGenericMemory) {\n printf(\"> Using Generic System Paged Memory (malloc)\\n\");\n }\n else {\n printf(\"> Using CUDA Host Allocated (cudaHostAlloc)\\n\");\n }\n\n checkCudaErrors(cudaSetDevice(idev));\n\n /* Verify the selected device supports mapped memory and set the device\n flags for mapping host memory. */\n\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, idev));\n\n#if CUDART_VERSION >= 2020\n\n if (!deviceProp.canMapHostMemory) {\n fprintf(stderr, \"Device %d does not support mapping CPU host memory!\\n\", idev);\n\n exit(EXIT_SUCCESS);\n }\n\n checkCudaErrors(cudaSetDeviceFlags(cudaDeviceMapHost));\n#else\n fprintf(stderr,\n \"CUDART version %d.%d does not support \"\n \" field\\n\",\n ,\n CUDART_VERSION / 1000,\n (CUDART_VERSION % 100) / 10);\n\n exit(EXIT_SUCCESS);\n#endif\n\n#if CUDART_VERSION < 4000\n\n if (bPinGenericMemory) {\n fprintf(stderr,\n \"CUDART version %d.%d does not support function\\n\",\n CUDART_VERSION / 1000,\n (CUDART_VERSION % 100) / 10);\n\n exit(EXIT_SUCCESS);\n }\n\n#endif\n\n /* Allocate mapped CPU memory. */\n\n nelem = 1048576;\n bytes = nelem * sizeof(float);\n\n if (bPinGenericMemory) {\n#if CUDART_VERSION >= 4000\n a_UA = (float *)malloc(bytes + MEMORY_ALIGNMENT);\n b_UA = (float *)malloc(bytes + MEMORY_ALIGNMENT);\n c_UA = (float *)malloc(bytes + MEMORY_ALIGNMENT);\n\n // We need to ensure memory is aligned to 4K (so we will need to padd memory\n // accordingly)\n a = (float *)ALIGN_UP(a_UA, MEMORY_ALIGNMENT);\n b = (float *)ALIGN_UP(b_UA, MEMORY_ALIGNMENT);\n c = (float *)ALIGN_UP(c_UA, MEMORY_ALIGNMENT);\n\n checkCudaErrors(cudaHostRegister(a, bytes, cudaHostRegisterMapped));\n checkCudaErrors(cudaHostRegister(b, bytes, cudaHostRegisterMapped));\n checkCudaErrors(cudaHostRegister(c, bytes, cudaHostRegisterMapped));\n#endif\n }\n else {\n#if CUDART_VERSION >= 2020\n flags = cudaHostAllocMapped;\n checkCudaErrors(cudaHostAlloc((void **)&a, bytes, flags));\n checkCudaErrors(cudaHostAlloc((void **)&b, bytes, flags));\n checkCudaErrors(cudaHostAlloc((void **)&c, bytes, flags));\n#endif\n }\n\n /* Initialize the vectors. */\n\n for (n = 0; n < nelem; n++) {\n a[n] = rand() / (float)RAND_MAX;\n b[n] = rand() / (float)RAND_MAX;\n }\n\n /* Get the device pointers for the pinned CPU memory mapped into the GPU\n memory space. */\n\n#if CUDART_VERSION >= 2020\n checkCudaErrors(cudaHostGetDevicePointer((void **)&d_a, (void *)a, 0));\n checkCudaErrors(cudaHostGetDevicePointer((void **)&d_b, (void *)b, 0));\n checkCudaErrors(cudaHostGetDevicePointer((void **)&d_c, (void *)c, 0));\n#endif\n\n /* Call the GPU kernel using the CPU pointers residing in CPU mapped memory.\n */\n printf(\"> vectorAddGPU kernel will add vectors using mapped CPU memory...\\n\");\n dim3 block(256);\n dim3 grid((unsigned int)ceil(nelem / (float)block.x));\n vectorAddGPU<<>>(d_a, d_b, d_c, nelem);\n checkCudaErrors(cudaDeviceSynchronize());\n getLastCudaError(\"vectorAddGPU() execution failed\");\n\n /* Compare the results */\n\n printf(\"> Checking the results from vectorAddGPU() ...\\n\");\n errorNorm = 0.f;\n refNorm = 0.f;\n\n for (n = 0; n < nelem; n++) {\n ref = a[n] + b[n];\n diff = c[n] - ref;\n errorNorm += diff * diff;\n refNorm += ref * ref;\n }\n\n errorNorm = (float)sqrt((double)errorNorm);\n refNorm = (float)sqrt((double)refNorm);\n\n /* Memory clean up */\n\n printf(\"> Releasing CPU memory...\\n\");\n\n if (bPinGenericMemory) {\n#if CUDART_VERSION >= 4000\n checkCudaErrors(cudaHostUnregister(a));\n checkCudaErrors(cudaHostUnregister(b));\n checkCudaErrors(cudaHostUnregister(c));\n free(a_UA);\n free(b_UA);\n free(c_UA);\n#endif\n }\n else {\n#if CUDART_VERSION >= 2020\n checkCudaErrors(cudaFreeHost(a));\n checkCudaErrors(cudaFreeHost(b));\n checkCudaErrors(cudaFreeHost(c));\n#endif\n }\n\n exit(errorNorm / refNorm < 1.e-6f ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#ifndef MAX\n#define MAX(a, b) (a > b ? a : b)\n#endif\n\n/* Add two vectors on the GPU */\n__global__ void vectorAddGPU(float *a, float *b, float *c, int N)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx < N) {\n c[idx] = a[idx] + b[idx];\n }\n}\n\n// Allocate generic memory with malloc() and pin it laster instead of using\n// cudaHostAlloc()\nbool bPinGenericMemory = false;\n\n// Macro to aligned up to the memory size in question\n#define MEMORY_ALIGNMENT 4096\n#define ALIGN_UP(x, size) (((size_t)x + (size - 1)) & (~(size - 1)))\n\nint main(int argc, char **argv)\n{\n int n, nelem, deviceCount;\n int idev = 0; // use default device 0\n char *device = NULL;\n unsigned int flags;\n size_t bytes;\n float *a, *b, *c; // Pinned memory allocated on the CPU\n float *a_UA, *b_UA, *c_UA; // Non-4K Aligned Pinned memory on the CPU\n float *d_a, *d_b, *d_c; // Device pointers for mapped memory\n float errorNorm, refNorm, ref, diff;\n hipDeviceProp_t deviceProp;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\")) {\n printf(\"Usage: simpleZeroCopy [OPTION]\\n\\n\");\n printf(\"Options:\\n\");\n printf(\" --device=[device #] Specify the device to be used\\n\");\n printf(\" --use_generic_memory (optional) use generic page-aligned for system \"\n \"memory\\n\");\n return EXIT_SUCCESS;\n }\n\n /* Get the device selected by the user or default to 0, and then set it. */\n if (getCmdLineArgumentString(argc, (const char **)argv, \"device\", &device)) {\n hipGetDeviceCount(&deviceCount);\n idev = atoi(device);\n\n if (idev >= deviceCount || idev < 0) {\n fprintf(stderr, \"Device number %d is invalid, will use default CUDA device 0.\\n\", idev);\n idev = 0;\n }\n }\n\n // if GPU found supports SM 1.2, then continue, otherwise we exit\n if (!checkCudaCapabilities(1, 2)) {\n exit(EXIT_SUCCESS);\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"use_generic_memory\")) {\n#if defined(__APPLE__) || defined(MACOSX)\n bPinGenericMemory = false; // Generic Pinning of System Paged memory is not\n // currently supported on Mac OSX\n#else\n bPinGenericMemory = true;\n#endif\n }\n\n if (bPinGenericMemory) {\n printf(\"> Using Generic System Paged Memory (malloc)\\n\");\n }\n else {\n printf(\"> Using CUDA Host Allocated (hipHostAlloc)\\n\");\n }\n\n checkCudaErrors(hipSetDevice(idev));\n\n /* Verify the selected device supports mapped memory and set the device\n flags for mapping host memory. */\n\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, idev));\n\n#if CUDART_VERSION >= 2020\n\n if (!deviceProp.canMapHostMemory) {\n fprintf(stderr, \"Device %d does not support mapping CPU host memory!\\n\", idev);\n\n exit(EXIT_SUCCESS);\n }\n\n checkCudaErrors(hipSetDeviceFlags(hipDeviceMapHost));\n#else\n fprintf(stderr,\n \"CUDART version %d.%d does not support \"\n \" field\\n\",\n ,\n CUDART_VERSION / 1000,\n (CUDART_VERSION % 100) / 10);\n\n exit(EXIT_SUCCESS);\n#endif\n\n#if CUDART_VERSION < 4000\n\n if (bPinGenericMemory) {\n fprintf(stderr,\n \"CUDART version %d.%d does not support function\\n\",\n CUDART_VERSION / 1000,\n (CUDART_VERSION % 100) / 10);\n\n exit(EXIT_SUCCESS);\n }\n\n#endif\n\n /* Allocate mapped CPU memory. */\n\n nelem = 1048576;\n bytes = nelem * sizeof(float);\n\n if (bPinGenericMemory) {\n#if CUDART_VERSION >= 4000\n a_UA = (float *)malloc(bytes + MEMORY_ALIGNMENT);\n b_UA = (float *)malloc(bytes + MEMORY_ALIGNMENT);\n c_UA = (float *)malloc(bytes + MEMORY_ALIGNMENT);\n\n // We need to ensure memory is aligned to 4K (so we will need to padd memory\n // accordingly)\n a = (float *)ALIGN_UP(a_UA, MEMORY_ALIGNMENT);\n b = (float *)ALIGN_UP(b_UA, MEMORY_ALIGNMENT);\n c = (float *)ALIGN_UP(c_UA, MEMORY_ALIGNMENT);\n\n checkCudaErrors(hipHostRegister(a, bytes, hipHostRegisterMapped));\n checkCudaErrors(hipHostRegister(b, bytes, hipHostRegisterMapped));\n checkCudaErrors(hipHostRegister(c, bytes, hipHostRegisterMapped));\n#endif\n }\n else {\n#if CUDART_VERSION >= 2020\n flags = hipHostMallocMapped;\n checkCudaErrors(hipHostAlloc((void **)&a, bytes, flags));\n checkCudaErrors(hipHostAlloc((void **)&b, bytes, flags));\n checkCudaErrors(hipHostAlloc((void **)&c, bytes, flags));\n#endif\n }\n\n /* Initialize the vectors. */\n\n for (n = 0; n < nelem; n++) {\n a[n] = rand() / (float)RAND_MAX;\n b[n] = rand() / (float)RAND_MAX;\n }\n\n /* Get the device pointers for the pinned CPU memory mapped into the GPU\n memory space. */\n\n#if CUDART_VERSION >= 2020\n checkCudaErrors(hipHostGetDevicePointer((void **)&d_a, (void *)a, 0));\n checkCudaErrors(hipHostGetDevicePointer((void **)&d_b, (void *)b, 0));\n checkCudaErrors(hipHostGetDevicePointer((void **)&d_c, (void *)c, 0));\n#endif\n\n /* Call the GPU kernel using the CPU pointers residing in CPU mapped memory.\n */\n printf(\"> vectorAddGPU kernel will add vectors using mapped CPU memory...\\n\");\n dim3 block(256);\n dim3 grid((unsigned int)ceil(nelem / (float)block.x));\n vectorAddGPU<<>>(d_a, d_b, d_c, nelem);\n checkCudaErrors(hipDeviceSynchronize());\n getLastCudaError(\"vectorAddGPU() execution failed\");\n\n /* Compare the results */\n\n printf(\"> Checking the results from vectorAddGPU() ...\\n\");\n errorNorm = 0.f;\n refNorm = 0.f;\n\n for (n = 0; n < nelem; n++) {\n ref = a[n] + b[n];\n diff = c[n] - ref;\n errorNorm += diff * diff;\n refNorm += ref * ref;\n }\n\n errorNorm = (float)sqrt((double)errorNorm);\n refNorm = (float)sqrt((double)refNorm);\n\n /* Memory clean up */\n\n printf(\"> Releasing CPU memory...\\n\");\n\n if (bPinGenericMemory) {\n#if CUDART_VERSION >= 4000\n checkCudaErrors(hipHostUnregister(a));\n checkCudaErrors(hipHostUnregister(b));\n checkCudaErrors(hipHostUnregister(c));\n free(a_UA);\n free(b_UA);\n free(c_UA);\n#endif\n }\n else {\n#if CUDART_VERSION >= 2020\n checkCudaErrors(hipHostFree(a));\n checkCudaErrors(hipHostFree(b));\n checkCudaErrors(hipHostFree(c));\n#endif\n }\n\n exit(errorNorm / refNorm < 1.e-6f ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleZeroCopy/simpleZeroCopy.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * Multi-GPU sample using OpenMP for threading on the CPU side\n * needs a compiler that supports OpenMP 2.0\n */\n\n#include \n#include \n#include // stdio functions are used since C++ streams aren't necessarily thread safe\n\nusing namespace std;\n\n// a simple kernel that simply increments each array element by b\n__global__ void kernelAddConstant(int *g_a, const int b)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n g_a[idx] += b;\n}\n\n// a predicate that checks whether each array element is set to its index plus b\nint correctResult(int *data, const int n, const int b)\n{\n for (int i = 0; i < n; i++)\n if (data[i] != i + b)\n return 0;\n\n return 1;\n}\n\nint main(int argc, char *argv[])\n{\n int num_gpus = 0; // number of CUDA GPUs\n\n printf(\"%s Starting...\\n\\n\", argv[0]);\n\n /////////////////////////////////////////////////////////////////\n // determine the number of CUDA capable GPUs\n //\n cudaGetDeviceCount(&num_gpus);\n\n if (num_gpus < 1) {\n printf(\"no CUDA capable devices were detected\\n\");\n return 1;\n }\n\n /////////////////////////////////////////////////////////////////\n // display CPU and GPU configuration\n //\n printf(\"number of host CPUs:\\t%d\\n\", omp_get_num_procs());\n printf(\"number of CUDA devices:\\t%d\\n\", num_gpus);\n\n for (int i = 0; i < num_gpus; i++) {\n cudaDeviceProp dprop;\n cudaGetDeviceProperties(&dprop, i);\n printf(\" %d: %s\\n\", i, dprop.name);\n }\n\n printf(\"---------------------------\\n\");\n\n /////////////////////////////////////////////////////////////////\n // initialize data\n //\n unsigned int n = num_gpus * 8192;\n unsigned int nbytes = n * sizeof(int);\n int *a = 0; // pointer to data on the CPU\n int b = 3; // value by which the array is incremented\n a = (int *)malloc(nbytes);\n\n if (0 == a) {\n printf(\"couldn't allocate CPU memory\\n\");\n return 1;\n }\n\n for (unsigned int i = 0; i < n; i++)\n a[i] = i;\n\n ////////////////////////////////////////////////////////////////\n // run as many CPU threads as there are CUDA devices\n // each CPU thread controls a different device, processing its\n // portion of the data. It's possible to use more CPU threads\n // than there are CUDA devices, in which case several CPU\n // threads will be allocating resources and launching kernels\n // on the same device. For example, try omp_set_num_threads(2*num_gpus);\n // Recall that all variables declared inside an \"omp parallel\" scope are\n // local to each CPU thread\n //\n omp_set_num_threads(num_gpus); // create as many CPU threads as there are CUDA devices\n// omp_set_num_threads(2*num_gpus);// create twice as many CPU threads as there\n// are CUDA devices\n#pragma omp parallel\n {\n unsigned int cpu_thread_id = omp_get_thread_num();\n unsigned int num_cpu_threads = omp_get_num_threads();\n\n // set and check the CUDA device for this CPU thread\n int gpu_id = -1;\n checkCudaErrors(\n cudaSetDevice(cpu_thread_id % num_gpus)); // \"% num_gpus\" allows more CPU threads than GPU devices\n checkCudaErrors(cudaGetDevice(&gpu_id));\n printf(\"CPU thread %d (of %d) uses CUDA device %d\\n\", cpu_thread_id, num_cpu_threads, gpu_id);\n\n int *d_a = 0; // pointer to memory on the device associated with this CPU thread\n int *sub_a = a + cpu_thread_id * n / num_cpu_threads; // pointer to this CPU thread's portion of data\n unsigned int nbytes_per_kernel = nbytes / num_cpu_threads;\n dim3 gpu_threads(128); // 128 threads per block\n dim3 gpu_blocks(n / (gpu_threads.x * num_cpu_threads));\n\n checkCudaErrors(cudaMalloc((void **)&d_a, nbytes_per_kernel));\n checkCudaErrors(cudaMemset(d_a, 0, nbytes_per_kernel));\n checkCudaErrors(cudaMemcpy(d_a, sub_a, nbytes_per_kernel, cudaMemcpyHostToDevice));\n kernelAddConstant<<>>(d_a, b);\n\n checkCudaErrors(cudaMemcpy(sub_a, d_a, nbytes_per_kernel, cudaMemcpyDeviceToHost));\n checkCudaErrors(cudaFree(d_a));\n }\n printf(\"---------------------------\\n\");\n\n if (cudaSuccess != cudaGetLastError())\n printf(\"%s\\n\", cudaGetErrorString(cudaGetLastError()));\n\n ////////////////////////////////////////////////////////////////\n // check the result\n //\n bool bResult = correctResult(a, n, b);\n\n if (a)\n free(a); // free CPU memory\n\n exit(bResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * Multi-GPU sample using OpenMP for threading on the CPU side\n * needs a compiler that supports OpenMP 2.0\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include // stdio functions are used since C++ streams aren't necessarily thread safe\n\nusing namespace std;\n\n// a simple kernel that simply increments each array element by b\n__global__ void kernelAddConstant(int *g_a, const int b)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n g_a[idx] += b;\n}\n\n// a predicate that checks whether each array element is set to its index plus b\nint correctResult(int *data, const int n, const int b)\n{\n for (int i = 0; i < n; i++)\n if (data[i] != i + b)\n return 0;\n\n return 1;\n}\n\nint main(int argc, char *argv[])\n{\n int num_gpus = 0; // number of CUDA GPUs\n\n printf(\"%s Starting...\\n\\n\", argv[0]);\n\n /////////////////////////////////////////////////////////////////\n // determine the number of CUDA capable GPUs\n //\n hipGetDeviceCount(&num_gpus);\n\n if (num_gpus < 1) {\n printf(\"no CUDA capable devices were detected\\n\");\n return 1;\n }\n\n /////////////////////////////////////////////////////////////////\n // display CPU and GPU configuration\n //\n printf(\"number of host CPUs:\\t%d\\n\", omp_get_num_procs());\n printf(\"number of CUDA devices:\\t%d\\n\", num_gpus);\n\n for (int i = 0; i < num_gpus; i++) {\n hipDeviceProp_t dprop;\n hipGetDeviceProperties(&dprop, i);\n printf(\" %d: %s\\n\", i, dprop.name);\n }\n\n printf(\"---------------------------\\n\");\n\n /////////////////////////////////////////////////////////////////\n // initialize data\n //\n unsigned int n = num_gpus * 8192;\n unsigned int nbytes = n * sizeof(int);\n int *a = 0; // pointer to data on the CPU\n int b = 3; // value by which the array is incremented\n a = (int *)malloc(nbytes);\n\n if (0 == a) {\n printf(\"couldn't allocate CPU memory\\n\");\n return 1;\n }\n\n for (unsigned int i = 0; i < n; i++)\n a[i] = i;\n\n ////////////////////////////////////////////////////////////////\n // run as many CPU threads as there are CUDA devices\n // each CPU thread controls a different device, processing its\n // portion of the data. It's possible to use more CPU threads\n // than there are CUDA devices, in which case several CPU\n // threads will be allocating resources and launching kernels\n // on the same device. For example, try omp_set_num_threads(2*num_gpus);\n // Recall that all variables declared inside an \"omp parallel\" scope are\n // local to each CPU thread\n //\n omp_set_num_threads(num_gpus); // create as many CPU threads as there are CUDA devices\n// omp_set_num_threads(2*num_gpus);// create twice as many CPU threads as there\n// are CUDA devices\n#pragma omp parallel\n {\n unsigned int cpu_thread_id = omp_get_thread_num();\n unsigned int num_cpu_threads = omp_get_num_threads();\n\n // set and check the CUDA device for this CPU thread\n int gpu_id = -1;\n checkCudaErrors(\n hipSetDevice(cpu_thread_id % num_gpus)); // \"% num_gpus\" allows more CPU threads than GPU devices\n checkCudaErrors(hipGetDevice(&gpu_id));\n printf(\"CPU thread %d (of %d) uses CUDA device %d\\n\", cpu_thread_id, num_cpu_threads, gpu_id);\n\n int *d_a = 0; // pointer to memory on the device associated with this CPU thread\n int *sub_a = a + cpu_thread_id * n / num_cpu_threads; // pointer to this CPU thread's portion of data\n unsigned int nbytes_per_kernel = nbytes / num_cpu_threads;\n dim3 gpu_threads(128); // 128 threads per block\n dim3 gpu_blocks(n / (gpu_threads.x * num_cpu_threads));\n\n checkCudaErrors(hipMalloc((void **)&d_a, nbytes_per_kernel));\n checkCudaErrors(hipMemset(d_a, 0, nbytes_per_kernel));\n checkCudaErrors(hipMemcpy(d_a, sub_a, nbytes_per_kernel, hipMemcpyHostToDevice));\n kernelAddConstant<<>>(d_a, b);\n\n checkCudaErrors(hipMemcpy(sub_a, d_a, nbytes_per_kernel, hipMemcpyDeviceToHost));\n checkCudaErrors(hipFree(d_a));\n }\n printf(\"---------------------------\\n\");\n\n if (hipSuccess != hipGetLastError())\n printf(\"%s\\n\", hipGetErrorString(hipGetLastError()));\n\n ////////////////////////////////////////////////////////////////\n // check the result\n //\n bool bResult = correctResult(a, n, b);\n\n if (a)\n free(a); // free CPU memory\n\n exit(bResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/cudaOpenMP/cudaOpenMP.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates Inter Process Communication\n * using one process per GPU for computation.\n */\n#include \n#include \n#include \n\n#include \"helper_cuda.h\"\n#include \"helper_multiprocess.h\"\nstatic const char shmName[] = \"simpleIPCshm\";\n// For direct NVLINK and PCI-E peers, at max 8 simultaneous peers are allowed\n// For NVSWITCH connected peers like DGX-2, simultaneous peers are not limited\n// in the same way.\n#define MAX_DEVICES (32)\n#define DATA_SIZE (64ULL << 20ULL) // 64MB\n\n#if defined(__linux__)\n#define cpu_atomic_add32(a, x) __sync_add_and_fetch(a, x)\n#elif defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)\n#define cpu_atomic_add32(a, x) InterlockedAdd((volatile LONG *)a, x)\n#else\n#error Unsupported system\n#endif\n\ntypedef struct shmStruct_st\n{\n size_t nprocesses;\n int barrier;\n int sense;\n int devices[MAX_DEVICES];\n cudaIpcMemHandle_t memHandle[MAX_DEVICES];\n cudaIpcEventHandle_t eventHandle[MAX_DEVICES];\n} shmStruct;\n\n__global__ void simpleKernel(char *ptr, int sz, char val)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n for (; idx < sz; idx += (gridDim.x * blockDim.x)) {\n ptr[idx] = val;\n }\n}\n\nstatic void barrierWait(volatile int *barrier, volatile int *sense, unsigned int n)\n{\n int count;\n\n // Check-in\n count = cpu_atomic_add32(barrier, 1);\n if (count == n) // Last one in\n *sense = 1;\n while (!*sense)\n ;\n\n // Check-out\n count = cpu_atomic_add32(barrier, -1);\n if (count == 0) // Last one out\n *sense = 0;\n while (*sense)\n ;\n}\n\nstatic void childProcess(int id)\n{\n volatile shmStruct *shm = NULL;\n cudaStream_t stream;\n sharedMemoryInfo info;\n size_t procCount, i;\n int blocks = 0;\n int threads = 128;\n cudaDeviceProp prop;\n std::vector ptrs;\n std::vector events;\n std::vector verification_buffer(DATA_SIZE);\n char pidString[20] = {0};\n char lshmName[40] = {0};\n\n // Use parent process ID to create a unique shared memory name for Linux multi-process\n#ifdef __linux__\n pid_t pid;\n pid = getppid();\n snprintf(pidString, sizeof(pidString), \"%d\", pid);\n#endif\n strcat(lshmName, shmName);\n strcat(lshmName, pidString);\n\n printf(\"CP: lshmName = %s\\n\", lshmName);\n\n if (sharedMemoryOpen(lshmName, sizeof(shmStruct), &info) != 0) {\n printf(\"Failed to create shared memory slab\\n\");\n exit(EXIT_FAILURE);\n }\n shm = (volatile shmStruct *)info.addr;\n procCount = shm->nprocesses;\n\n printf(\"Process %d: Starting on device %d...\\n\", id, shm->devices[id]);\n\n checkCudaErrors(cudaSetDevice(shm->devices[id]));\n checkCudaErrors(cudaGetDeviceProperties(&prop, shm->devices[id]));\n checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));\n checkCudaErrors(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, simpleKernel, threads, 0));\n blocks *= prop.multiProcessorCount;\n\n // Open and track all the allocations and events created in the master\n // process for use later\n for (i = 0; i < procCount; i++) {\n void *ptr = NULL;\n cudaEvent_t event;\n\n // Notice, we don't need to explicitly enable peer access for\n // allocations on other devices.\n checkCudaErrors(\n cudaIpcOpenMemHandle(&ptr, *(cudaIpcMemHandle_t *)&shm->memHandle[i], cudaIpcMemLazyEnablePeerAccess));\n checkCudaErrors(cudaIpcOpenEventHandle(&event, *(cudaIpcEventHandle_t *)&shm->eventHandle[i]));\n\n ptrs.push_back(ptr);\n events.push_back(event);\n }\n\n // At each iteration of the loop, each sibling process will push work on\n // their respective devices accessing the next peer mapped buffer allocated\n // by the master process (these can come from other sibling processes as\n // well). To coordinate each process' access, we force the stream to wait for\n // the work already accessing this buffer asynchronously through IPC events,\n // allowing the CPU processes to continue to queue more work.\n for (i = 0; i < procCount; i++) {\n size_t bufferId = (i + id) % procCount;\n // Wait for the buffer to be accessed to be ready\n checkCudaErrors(cudaStreamWaitEvent(stream, events[bufferId], 0));\n // Push a simple kernel on it\n simpleKernel<<>>((char *)ptrs[bufferId], DATA_SIZE, id);\n checkCudaErrors(cudaGetLastError());\n // Signal that this buffer is ready for the next consumer\n checkCudaErrors(cudaEventRecord(events[bufferId], stream));\n // Wait for all my sibling processes to push this stage of their work\n // before proceeding to the next. This prevents siblings from racing\n // ahead and clobbering the recorded event or waiting on the wrong\n // recorded event.\n barrierWait(&shm->barrier, &shm->sense, (unsigned int)procCount);\n if (id == 0) {\n printf(\"Step %lld done\\n\", (unsigned long long)i);\n }\n }\n\n // Now wait for my buffer to be ready so I can copy it locally and verify it\n checkCudaErrors(cudaStreamWaitEvent(stream, events[id], 0));\n checkCudaErrors(cudaMemcpyAsync(&verification_buffer[0], ptrs[id], DATA_SIZE, cudaMemcpyDeviceToHost, stream));\n // And wait for all the queued up work to complete\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n printf(\"Process %d: verifying...\\n\", id);\n\n // The contents should have the id of the sibling just after me\n char compareId = (char)((id + 1) % procCount);\n for (unsigned long long j = 0; j < DATA_SIZE; j++) {\n if (verification_buffer[j] != compareId) {\n printf(\"Process %d: Verification mismatch at %lld: %d != %d\\n\",\n id,\n j,\n (int)verification_buffer[j],\n (int)compareId);\n }\n }\n\n // Clean up!\n for (i = 0; i < procCount; i++) {\n checkCudaErrors(cudaIpcCloseMemHandle(ptrs[i]));\n checkCudaErrors(cudaEventDestroy(events[i]));\n }\n\n checkCudaErrors(cudaStreamDestroy(stream));\n\n printf(\"Process %d complete!\\n\", id);\n}\n\nstatic void parentProcess(char *app)\n{\n sharedMemoryInfo info;\n int devCount, i;\n volatile shmStruct *shm = NULL;\n std::vector ptrs;\n std::vector events;\n std::vector processes;\n char pidString[20] = {0};\n char lshmName[40] = {0};\n\n // Use current process ID to create a unique shared memory name for Linux multi-process\n#ifdef __linux__\n pid_t pid;\n pid = getpid();\n snprintf(pidString, sizeof(pidString), \"%d\", pid);\n#endif\n strcat(lshmName, shmName);\n strcat(lshmName, pidString);\n\n printf(\"PP: lshmName = %s\\n\", lshmName);\n\n checkCudaErrors(cudaGetDeviceCount(&devCount));\n\n if (sharedMemoryCreate(lshmName, sizeof(*shm), &info) != 0) {\n printf(\"Failed to create shared memory slab\\n\");\n exit(EXIT_FAILURE);\n }\n shm = (volatile shmStruct *)info.addr;\n memset((void *)shm, 0, sizeof(*shm));\n\n // Pick all the devices that can access each other's memory for this test\n // Keep in mind that CUDA has minimal support for fork() without a\n // corresponding exec() in the child process, but in this case our\n // spawnProcess will always exec, so no need to worry.\n for (i = 0; i < devCount; i++) {\n bool allPeers = true;\n cudaDeviceProp prop;\n checkCudaErrors(cudaGetDeviceProperties(&prop, i));\n\n // CUDA IPC is only supported on devices with unified addressing\n if (!prop.unifiedAddressing) {\n printf(\"Device %d does not support unified addressing, skipping...\\n\", i);\n continue;\n }\n // This sample requires two processes accessing each device, so we need\n // to ensure exclusive or prohibited mode is not set\n int computeMode;\n checkCudaErrors(cudaDeviceGetAttribute(&computeMode, cudaDevAttrComputeMode, i));\n if (computeMode != cudaComputeModeDefault) {\n printf(\"Device %d is in an unsupported compute mode for this sample\\n\", i);\n continue;\n }\n\n for (int j = 0; j < shm->nprocesses; j++) {\n int canAccessPeerIJ, canAccessPeerJI;\n checkCudaErrors(cudaDeviceCanAccessPeer(&canAccessPeerJI, shm->devices[j], i));\n checkCudaErrors(cudaDeviceCanAccessPeer(&canAccessPeerIJ, i, shm->devices[j]));\n if (!canAccessPeerIJ || !canAccessPeerJI) {\n allPeers = false;\n break;\n }\n }\n if (allPeers) {\n // Enable peers here. This isn't necessary for IPC, but it will\n // setup the peers for the device. For systems that only allow 8\n // peers per GPU at a time, this acts to remove devices from CanAccessPeer\n for (int j = 0; j < shm->nprocesses; j++) {\n checkCudaErrors(cudaSetDevice(i));\n checkCudaErrors(cudaDeviceEnablePeerAccess(shm->devices[j], 0));\n checkCudaErrors(cudaSetDevice(shm->devices[j]));\n checkCudaErrors(cudaDeviceEnablePeerAccess(i, 0));\n }\n shm->devices[shm->nprocesses++] = i;\n if (shm->nprocesses >= MAX_DEVICES)\n break;\n }\n else {\n printf(\"Device %d is not peer capable with some other selected peers, \"\n \"skipping\\n\",\n i);\n }\n }\n\n if (shm->nprocesses == 0) {\n printf(\"No CUDA devices support IPC\\n\");\n exit(EXIT_WAIVED);\n }\n\n // Now allocate memory and an event for each process and fill the shared\n // memory buffer with the IPC handles to communicate\n for (i = 0; i < shm->nprocesses; i++) {\n void *ptr = NULL;\n cudaEvent_t event;\n\n checkCudaErrors(cudaSetDevice(shm->devices[i]));\n checkCudaErrors(cudaMalloc(&ptr, DATA_SIZE));\n checkCudaErrors(cudaIpcGetMemHandle((cudaIpcMemHandle_t *)&shm->memHandle[i], ptr));\n checkCudaErrors(cudaEventCreate(&event, cudaEventDisableTiming | cudaEventInterprocess));\n checkCudaErrors(cudaIpcGetEventHandle((cudaIpcEventHandle_t *)&shm->eventHandle[i], event));\n\n ptrs.push_back(ptr);\n events.push_back(event);\n }\n\n // Launch the child processes!\n for (i = 0; i < shm->nprocesses; i++) {\n char devIdx[12]; // Increased size to ensure enough space for formatted integer\n char *const args[] = {app, devIdx, NULL};\n Process process;\n\n snprintf(devIdx, sizeof(devIdx), \"%d\", i);\n\n if (spawnProcess(&process, app, args)) {\n printf(\"Failed to create process\\n\");\n exit(EXIT_FAILURE);\n }\n\n processes.push_back(process);\n }\n\n // And wait for them to finish\n for (i = 0; i < processes.size(); i++) {\n if (waitProcess(&processes[i]) != EXIT_SUCCESS) {\n printf(\"Process %d failed!\\n\", i);\n exit(EXIT_FAILURE);\n }\n }\n\n // Clean up!\n for (i = 0; i < shm->nprocesses; i++) {\n checkCudaErrors(cudaSetDevice(shm->devices[i]));\n checkCudaErrors(cudaEventSynchronize(events[i]));\n checkCudaErrors(cudaEventDestroy(events[i]));\n checkCudaErrors(cudaFree(ptrs[i]));\n }\n\n sharedMemoryClose(&info);\n}\n\nint main(int argc, char **argv)\n{\n#if defined(__arm__) || defined(__aarch64__)\n printf(\"Not supported on ARM\\n\");\n return EXIT_WAIVED;\n#else\n if (argc == 1) {\n parentProcess(argv[0]);\n }\n else {\n childProcess(atoi(argv[1]));\n }\n return EXIT_SUCCESS;\n#endif\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates Inter Process Communication\n * using one process per GPU for computation.\n */\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n#include \"hip/hip_runtime_api.h\"\n#include \"helper_multiprocess.h\"\nstatic const char shmName[] = \"simpleIPCshm\";\n// For direct NVLINK and PCI-E peers, at max 8 simultaneous peers are allowed\n// For NVSWITCH connected peers like DGX-2, simultaneous peers are not limited\n// in the same way.\n#define MAX_DEVICES (32)\n#define DATA_SIZE (64ULL << 20ULL) // 64MB\n\n#if defined(__linux__)\n#define cpu_atomic_add32(a, x) __sync_add_and_fetch(a, x)\n#elif defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)\n#define cpu_atomic_add32(a, x) InterlockedAdd((volatile LONG *)a, x)\n#else\n#error Unsupported system\n#endif\n\ntypedef struct shmStruct_st\n{\n size_t nprocesses;\n int barrier;\n int sense;\n int devices[MAX_DEVICES];\n hipIpcMemHandle_t memHandle[MAX_DEVICES];\n hipIpcEventHandle_t eventHandle[MAX_DEVICES];\n} shmStruct;\n\n__global__ void simpleKernel(char *ptr, int sz, char val)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n for (; idx < sz; idx += (gridDim.x * blockDim.x)) {\n ptr[idx] = val;\n }\n}\n\nstatic void barrierWait(volatile int *barrier, volatile int *sense, unsigned int n)\n{\n int count;\n\n // Check-in\n count = cpu_atomic_add32(barrier, 1);\n if (count == n) // Last one in\n *sense = 1;\n while (!*sense)\n ;\n\n // Check-out\n count = cpu_atomic_add32(barrier, -1);\n if (count == 0) // Last one out\n *sense = 0;\n while (*sense)\n ;\n}\n\nstatic void childProcess(int id)\n{\n volatile shmStruct *shm = NULL;\n hipStream_t stream;\n sharedMemoryInfo info;\n size_t procCount, i;\n int blocks = 0;\n int threads = 128;\n hipDeviceProp_t prop;\n std::vector ptrs;\n std::vector events;\n std::vector verification_buffer(DATA_SIZE);\n char pidString[20] = {0};\n char lshmName[40] = {0};\n\n // Use parent process ID to create a unique shared memory name for Linux multi-process\n#ifdef __linux__\n pid_t pid;\n pid = getppid();\n snprintf(pidString, sizeof(pidString), \"%d\", pid);\n#endif\n strcat(lshmName, shmName);\n strcat(lshmName, pidString);\n\n printf(\"CP: lshmName = %s\\n\", lshmName);\n\n if (sharedMemoryOpen(lshmName, sizeof(shmStruct), &info) != 0) {\n printf(\"Failed to create shared memory slab\\n\");\n exit(EXIT_FAILURE);\n }\n shm = (volatile shmStruct *)info.addr;\n procCount = shm->nprocesses;\n\n printf(\"Process %d: Starting on device %d...\\n\", id, shm->devices[id]);\n\n checkCudaErrors(hipSetDevice(shm->devices[id]));\n checkCudaErrors(hipGetDeviceProperties(&prop, shm->devices[id]));\n checkCudaErrors(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));\n checkCudaErrors(hipOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, simpleKernel, threads, 0));\n blocks *= prop.multiProcessorCount;\n\n // Open and track all the allocations and events created in the master\n // process for use later\n for (i = 0; i < procCount; i++) {\n void *ptr = NULL;\n hipEvent_t event;\n\n // Notice, we don't need to explicitly enable peer access for\n // allocations on other devices.\n checkCudaErrors(\n hipIpcOpenMemHandle(&ptr, *(hipIpcMemHandle_t *)&shm->memHandle[i], hipIpcMemLazyEnablePeerAccess));\n checkCudaErrors(hipIpcOpenEventHandle(&event, *(hipIpcEventHandle_t *)&shm->eventHandle[i]));\n\n ptrs.push_back(ptr);\n events.push_back(event);\n }\n\n // At each iteration of the loop, each sibling process will push work on\n // their respective devices accessing the next peer mapped buffer allocated\n // by the master process (these can come from other sibling processes as\n // well). To coordinate each process' access, we force the stream to wait for\n // the work already accessing this buffer asynchronously through IPC events,\n // allowing the CPU processes to continue to queue more work.\n for (i = 0; i < procCount; i++) {\n size_t bufferId = (i + id) % procCount;\n // Wait for the buffer to be accessed to be ready\n checkCudaErrors(hipStreamWaitEvent(stream, events[bufferId], 0));\n // Push a simple kernel on it\n simpleKernel<<>>((char *)ptrs[bufferId], DATA_SIZE, id);\n checkCudaErrors(hipGetLastError());\n // Signal that this buffer is ready for the next consumer\n checkCudaErrors(hipEventRecord(events[bufferId], stream));\n // Wait for all my sibling processes to push this stage of their work\n // before proceeding to the next. This prevents siblings from racing\n // ahead and clobbering the recorded event or waiting on the wrong\n // recorded event.\n barrierWait(&shm->barrier, &shm->sense, (unsigned int)procCount);\n if (id == 0) {\n printf(\"Step %lld done\\n\", (unsigned long long)i);\n }\n }\n\n // Now wait for my buffer to be ready so I can copy it locally and verify it\n checkCudaErrors(hipStreamWaitEvent(stream, events[id], 0));\n checkCudaErrors(hipMemcpyAsync(&verification_buffer[0], ptrs[id], DATA_SIZE, hipMemcpyDeviceToHost, stream));\n // And wait for all the queued up work to complete\n checkCudaErrors(hipStreamSynchronize(stream));\n\n printf(\"Process %d: verifying...\\n\", id);\n\n // The contents should have the id of the sibling just after me\n char compareId = (char)((id + 1) % procCount);\n for (unsigned long long j = 0; j < DATA_SIZE; j++) {\n if (verification_buffer[j] != compareId) {\n printf(\"Process %d: Verification mismatch at %lld: %d != %d\\n\",\n id,\n j,\n (int)verification_buffer[j],\n (int)compareId);\n }\n }\n\n // Clean up!\n for (i = 0; i < procCount; i++) {\n checkCudaErrors(hipIpcCloseMemHandle(ptrs[i]));\n checkCudaErrors(hipEventDestroy(events[i]));\n }\n\n checkCudaErrors(hipStreamDestroy(stream));\n\n printf(\"Process %d complete!\\n\", id);\n}\n\nstatic void parentProcess(char *app)\n{\n sharedMemoryInfo info;\n int devCount, i;\n volatile shmStruct *shm = NULL;\n std::vector ptrs;\n std::vector events;\n std::vector processes;\n char pidString[20] = {0};\n char lshmName[40] = {0};\n\n // Use current process ID to create a unique shared memory name for Linux multi-process\n#ifdef __linux__\n pid_t pid;\n pid = getpid();\n snprintf(pidString, sizeof(pidString), \"%d\", pid);\n#endif\n strcat(lshmName, shmName);\n strcat(lshmName, pidString);\n\n printf(\"PP: lshmName = %s\\n\", lshmName);\n\n checkCudaErrors(hipGetDeviceCount(&devCount));\n\n if (sharedMemoryCreate(lshmName, sizeof(*shm), &info) != 0) {\n printf(\"Failed to create shared memory slab\\n\");\n exit(EXIT_FAILURE);\n }\n shm = (volatile shmStruct *)info.addr;\n memset((void *)shm, 0, sizeof(*shm));\n\n // Pick all the devices that can access each other's memory for this test\n // Keep in mind that CUDA has minimal support for fork() without a\n // corresponding exec() in the child process, but in this case our\n // spawnProcess will always exec, so no need to worry.\n for (i = 0; i < devCount; i++) {\n bool allPeers = true;\n hipDeviceProp_t prop;\n checkCudaErrors(hipGetDeviceProperties(&prop, i));\n\n // CUDA IPC is only supported on devices with unified addressing\n if (!prop.unifiedAddressing) {\n printf(\"Device %d does not support unified addressing, skipping...\\n\", i);\n continue;\n }\n // This sample requires two processes accessing each device, so we need\n // to ensure exclusive or prohibited mode is not set\n int computeMode;\n checkCudaErrors(hipDeviceGetAttribute(&computeMode, hipDeviceAttributeComputeMode, i));\n if (computeMode != hipComputeModeDefault) {\n printf(\"Device %d is in an unsupported compute mode for this sample\\n\", i);\n continue;\n }\n\n for (int j = 0; j < shm->nprocesses; j++) {\n int canAccessPeerIJ, canAccessPeerJI;\n checkCudaErrors(hipDeviceCanAccessPeer(&canAccessPeerJI, shm->devices[j], i));\n checkCudaErrors(hipDeviceCanAccessPeer(&canAccessPeerIJ, i, shm->devices[j]));\n if (!canAccessPeerIJ || !canAccessPeerJI) {\n allPeers = false;\n break;\n }\n }\n if (allPeers) {\n // Enable peers here. This isn't necessary for IPC, but it will\n // setup the peers for the device. For systems that only allow 8\n // peers per GPU at a time, this acts to remove devices from CanAccessPeer\n for (int j = 0; j < shm->nprocesses; j++) {\n checkCudaErrors(hipSetDevice(i));\n checkCudaErrors(hipDeviceEnablePeerAccess(shm->devices[j], 0));\n checkCudaErrors(hipSetDevice(shm->devices[j]));\n checkCudaErrors(hipDeviceEnablePeerAccess(i, 0));\n }\n shm->devices[shm->nprocesses++] = i;\n if (shm->nprocesses >= MAX_DEVICES)\n break;\n }\n else {\n printf(\"Device %d is not peer capable with some other selected peers, \"\n \"skipping\\n\",\n i);\n }\n }\n\n if (shm->nprocesses == 0) {\n printf(\"No CUDA devices support IPC\\n\");\n exit(EXIT_WAIVED);\n }\n\n // Now allocate memory and an event for each process and fill the shared\n // memory buffer with the IPC handles to communicate\n for (i = 0; i < shm->nprocesses; i++) {\n void *ptr = NULL;\n hipEvent_t event;\n\n checkCudaErrors(hipSetDevice(shm->devices[i]));\n checkCudaErrors(hipMalloc(&ptr, DATA_SIZE));\n checkCudaErrors(hipIpcGetMemHandle((hipIpcMemHandle_t *)&shm->memHandle[i], ptr));\n checkCudaErrors(hipEventCreate(&event, hipEventDisableTiming | hipEventInterprocess));\n checkCudaErrors(hipIpcGetEventHandle((hipIpcEventHandle_t *)&shm->eventHandle[i], event));\n\n ptrs.push_back(ptr);\n events.push_back(event);\n }\n\n // Launch the child processes!\n for (i = 0; i < shm->nprocesses; i++) {\n char devIdx[12]; // Increased size to ensure enough space for formatted integer\n char *const args[] = {app, devIdx, NULL};\n Process process;\n\n snprintf(devIdx, sizeof(devIdx), \"%d\", i);\n\n if (spawnProcess(&process, app, args)) {\n printf(\"Failed to create process\\n\");\n exit(EXIT_FAILURE);\n }\n\n processes.push_back(process);\n }\n\n // And wait for them to finish\n for (i = 0; i < processes.size(); i++) {\n if (waitProcess(&processes[i]) != EXIT_SUCCESS) {\n printf(\"Process %d failed!\\n\", i);\n exit(EXIT_FAILURE);\n }\n }\n\n // Clean up!\n for (i = 0; i < shm->nprocesses; i++) {\n checkCudaErrors(hipSetDevice(shm->devices[i]));\n checkCudaErrors(hipEventSynchronize(events[i]));\n checkCudaErrors(hipEventDestroy(events[i]));\n checkCudaErrors(hipFree(ptrs[i]));\n }\n\n sharedMemoryClose(&info);\n}\n\nint main(int argc, char **argv)\n{\n#if defined(__arm__) || defined(__aarch64__)\n printf(\"Not supported on ARM\\n\");\n return EXIT_WAIVED;\n#else\n if (argc == 1) {\n parentProcess(argv[0]);\n }\n else {\n childProcess(atoi(argv[1]));\n }\n return EXIT_SUCCESS;\n#endif\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleIPC/simpleIPC.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* Template project which demonstrates the basics on how to setup a project\n * example application.\n * Host code.\n */\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// includes CUDA\n#include \n\n// includes, project\n#include \n#include // helper functions for SDK examples\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\nvoid runTest(int argc, char **argv);\n\nextern \"C\" void computeGold(float *reference, float *idata, const unsigned int len);\n\n////////////////////////////////////////////////////////////////////////////////\n//! Simple test kernel for device functionality\n//! @param g_idata input data in global memory\n//! @param g_odata output data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void testKernel(float *g_idata, float *g_odata)\n{\n // shared memory\n // the size is determined by the host application\n extern __shared__ float sdata[];\n\n // access thread id\n const unsigned int tid = threadIdx.x;\n // access number of threads in this block\n const unsigned int num_threads = blockDim.x;\n\n // read in input data from global memory\n sdata[tid] = g_idata[tid];\n __syncthreads();\n\n // perform some computations\n sdata[tid] = (float)num_threads * sdata[tid];\n __syncthreads();\n\n // write data to global memory\n g_odata[tid] = sdata[tid];\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv) { runTest(argc, argv); }\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nvoid runTest(int argc, char **argv)\n{\n bool bTestResult = true;\n\n printf(\"%s Starting...\\n\\n\", argv[0]);\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n int devID = findCudaDevice(argc, (const char **)argv);\n\n StopWatchInterface *timer = 0;\n sdkCreateTimer(&timer);\n sdkStartTimer(&timer);\n\n unsigned int num_threads = 32;\n unsigned int mem_size = sizeof(float) * num_threads;\n\n // allocate host memory\n float *h_idata = (float *)malloc(mem_size);\n\n // initalize the memory\n for (unsigned int i = 0; i < num_threads; ++i) {\n h_idata[i] = (float)i;\n }\n\n // allocate device memory\n float *d_idata;\n checkCudaErrors(cudaMalloc((void **)&d_idata, mem_size));\n // copy host memory to device\n checkCudaErrors(cudaMemcpy(d_idata, h_idata, mem_size, cudaMemcpyHostToDevice));\n\n // allocate device memory for result\n float *d_odata;\n checkCudaErrors(cudaMalloc((void **)&d_odata, mem_size));\n\n // setup execution parameters\n dim3 grid(1, 1, 1);\n dim3 threads(num_threads, 1, 1);\n\n // execute the kernel\n testKernel<<>>(d_idata, d_odata);\n\n // check if kernel execution generated and error\n getLastCudaError(\"Kernel execution failed\");\n\n // allocate mem for the result on host side\n float *h_odata = (float *)malloc(mem_size);\n // copy result from device to host\n checkCudaErrors(cudaMemcpy(h_odata, d_odata, sizeof(float) * num_threads, cudaMemcpyDeviceToHost));\n\n sdkStopTimer(&timer);\n printf(\"Processing time: %f (ms)\\n\", sdkGetTimerValue(&timer));\n sdkDeleteTimer(&timer);\n\n // compute reference solution\n float *reference = (float *)malloc(mem_size);\n computeGold(reference, h_idata, num_threads);\n\n // check result\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // write file for regression test\n sdkWriteFile(\"./data/regression.dat\", h_odata, num_threads, 0.0f, false);\n }\n else {\n // custom output handling when no regression test running\n // in this case check if the result is equivalent to the expected solution\n bTestResult = compareData(reference, h_odata, num_threads, 0.0f, 0.0f);\n }\n\n // cleanup memory\n free(h_idata);\n free(h_odata);\n free(reference);\n checkCudaErrors(cudaFree(d_idata));\n checkCudaErrors(cudaFree(d_odata));\n\n exit(bTestResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* Template project which demonstrates the basics on how to setup a project\n * example application.\n * Host code.\n */\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// includes CUDA\n#include \n\n// includes, project\n#include \n#include // helper functions for SDK examples\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\nvoid runTest(int argc, char **argv);\n\nextern \"C\" void computeGold(float *reference, float *idata, const unsigned int len);\n\n////////////////////////////////////////////////////////////////////////////////\n//! Simple test kernel for device functionality\n//! @param g_idata input data in global memory\n//! @param g_odata output data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void testKernel(float *g_idata, float *g_odata)\n{\n // shared memory\n // the size is determined by the host application\n extern __shared__ float sdata[];\n\n // access thread id\n const unsigned int tid = threadIdx.x;\n // access number of threads in this block\n const unsigned int num_threads = blockDim.x;\n\n // read in input data from global memory\n sdata[tid] = g_idata[tid];\n __syncthreads();\n\n // perform some computations\n sdata[tid] = (float)num_threads * sdata[tid];\n __syncthreads();\n\n // write data to global memory\n g_odata[tid] = sdata[tid];\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv) { runTest(argc, argv); }\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nvoid runTest(int argc, char **argv)\n{\n bool bTestResult = true;\n\n printf(\"%s Starting...\\n\\n\", argv[0]);\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n int devID = findCudaDevice(argc, (const char **)argv);\n\n StopWatchInterface *timer = 0;\n sdkCreateTimer(&timer);\n sdkStartTimer(&timer);\n\n unsigned int num_threads = 32;\n unsigned int mem_size = sizeof(float) * num_threads;\n\n // allocate host memory\n float *h_idata = (float *)malloc(mem_size);\n\n // initalize the memory\n for (unsigned int i = 0; i < num_threads; ++i) {\n h_idata[i] = (float)i;\n }\n\n // allocate device memory\n float *d_idata;\n checkCudaErrors(hipMalloc((void **)&d_idata, mem_size));\n // copy host memory to device\n checkCudaErrors(hipMemcpy(d_idata, h_idata, mem_size, hipMemcpyHostToDevice));\n\n // allocate device memory for result\n float *d_odata;\n checkCudaErrors(hipMalloc((void **)&d_odata, mem_size));\n\n // setup execution parameters\n dim3 grid(1, 1, 1);\n dim3 threads(num_threads, 1, 1);\n\n // execute the kernel\n testKernel<<>>(d_idata, d_odata);\n\n // check if kernel execution generated and error\n getLastCudaError(\"Kernel execution failed\");\n\n // allocate mem for the result on host side\n float *h_odata = (float *)malloc(mem_size);\n // copy result from device to host\n checkCudaErrors(hipMemcpy(h_odata, d_odata, sizeof(float) * num_threads, hipMemcpyDeviceToHost));\n\n sdkStopTimer(&timer);\n printf(\"Processing time: %f (ms)\\n\", sdkGetTimerValue(&timer));\n sdkDeleteTimer(&timer);\n\n // compute reference solution\n float *reference = (float *)malloc(mem_size);\n computeGold(reference, h_idata, num_threads);\n\n // check result\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // write file for regression test\n sdkWriteFile(\"./data/regression.dat\", h_odata, num_threads, 0.0f, false);\n }\n else {\n // custom output handling when no regression test running\n // in this case check if the result is equivalent to the expected solution\n bTestResult = compareData(reference, h_odata, num_threads, 0.0f, 0.0f);\n }\n\n // cleanup memory\n free(h_idata);\n free(h_odata);\n free(reference);\n checkCudaErrors(hipFree(d_idata));\n checkCudaErrors(hipFree(d_odata));\n\n exit(bTestResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/template/template.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n////////////////////////////////////////////////////////////////////////////////\n// Kernels\n////////////////////////////////////////////////////////////////////////////////\n//! Tests assert function.\n//! Thread whose id > N will print assertion failed error message.\n////////////////////////////////////////////////////////////////////////////////\n\nextern \"C\" __global__ void testKernel(int N)\n{\n int gtid = blockIdx.x * blockDim.x + threadIdx.x;\n assert(gtid < N);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n////////////////////////////////////////////////////////////////////////////////\n// Kernels\n////////////////////////////////////////////////////////////////////////////////\n//! Tests assert function.\n//! Thread whose id > N will print assertion failed error message.\n////////////////////////////////////////////////////////////////////////////////\n\n#include \"hip/hip_runtime.h\"\nextern \"C\" __global__ void testKernel(int N)\n{\n int gtid = blockIdx.x * blockDim.x + threadIdx.x;\n assert(gtid < N);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleAssert_nvrtc/simpleAssert_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// includes CUDA\n#include \n\n// includes, project\n#include \n#include // helper functions for SDK examples\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\nvoid runTest(int argc, char **argv);\n\ncudaAccessPolicyWindow initAccessPolicyWindow(void)\n{\n cudaAccessPolicyWindow accessPolicyWindow = {0};\n accessPolicyWindow.base_ptr = (void *)0;\n accessPolicyWindow.num_bytes = 0;\n accessPolicyWindow.hitRatio = 0.f;\n accessPolicyWindow.hitProp = cudaAccessPropertyNormal;\n accessPolicyWindow.missProp = cudaAccessPropertyStreaming;\n return accessPolicyWindow;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Simple test kernel for device functionality\n//! @param data input data in global memory\n//! @param dataSize input data size\n//! @param bigData input bigData in global memory\n//! @param bigDataSize input bigData size\n//! @param hitcount how many data access are done within block\n////////////////////////////////////////////////////////////////////////////////\nstatic __global__ void kernCacheSegmentTest(int *data, int dataSize, int *trash, int bigDataSize, int hitCount)\n{\n __shared__ unsigned int hit;\n int row = blockIdx.y * blockDim.y + threadIdx.y;\n int col = blockIdx.x * blockDim.x + threadIdx.x;\n int tID = row * blockDim.y + col;\n uint32_t psRand = tID;\n\n atomicExch(&hit, 0);\n __syncthreads();\n while (hit < hitCount) {\n psRand ^= psRand << 13;\n psRand ^= psRand >> 17;\n psRand ^= psRand << 5;\n\n int idx = tID - psRand;\n if (idx < 0) {\n idx = -idx;\n }\n\n if ((tID % 2) == 0) {\n data[psRand % dataSize] = data[psRand % dataSize] + data[idx % dataSize];\n }\n else {\n trash[psRand % bigDataSize] = trash[psRand % bigDataSize] + trash[idx % bigDataSize];\n }\n\n atomicAdd(&hit, 1);\n }\n}\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv) { runTest(argc, argv); }\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nvoid runTest(int argc, char **argv)\n{\n bool bTestResult = true;\n cudaAccessPolicyWindow accessPolicyWindow;\n cudaDeviceProp deviceProp;\n cudaStreamAttrValue streamAttrValue;\n cudaStream_t stream;\n cudaStreamAttrID streamAttrID;\n dim3 threads(32, 32);\n int *dataDevicePointer;\n int *dataHostPointer;\n int dataSize;\n int *bigDataDevicePointer;\n int *bigDataHostPointer;\n int bigDataSize;\n StopWatchInterface *timer = 0;\n\n printf(\"%s Starting...\\n\\n\", argv[0]);\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n int devID = findCudaDevice(argc, (const char **)argv);\n sdkCreateTimer(&timer);\n sdkStartTimer(&timer);\n // Get device properties\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, devID));\n dim3 blocks(deviceProp.maxGridSize[1], 1);\n\n // Make sure device the l2 optimization\n if (deviceProp.persistingL2CacheMaxSize == 0) {\n printf(\"Waiving execution as device %d does not support persisting L2 \"\n \"Caching\\n\",\n devID);\n exit(EXIT_WAIVED);\n }\n\n // Create stream to assiocate with window\n checkCudaErrors(cudaStreamCreate(&stream));\n\n // Set the amount of l2 cache that will be persisting to maximum the device\n // can support\n checkCudaErrors(cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, deviceProp.persistingL2CacheMaxSize));\n\n // Stream attribute to set\n streamAttrID = cudaStreamAttributeAccessPolicyWindow;\n\n // Default window\n streamAttrValue.accessPolicyWindow = initAccessPolicyWindow();\n accessPolicyWindow = initAccessPolicyWindow();\n\n // Allocate size of both buffers\n bigDataSize = (deviceProp.l2CacheSize * 4) / sizeof(int);\n dataSize = (deviceProp.l2CacheSize / 4) / sizeof(int);\n\n // Allocate data\n checkCudaErrors(cudaMallocHost(&dataHostPointer, dataSize * sizeof(int)));\n checkCudaErrors(cudaMallocHost(&bigDataHostPointer, bigDataSize * sizeof(int)));\n\n for (int i = 0; i < bigDataSize; ++i) {\n if (i < dataSize) {\n dataHostPointer[i] = i;\n }\n\n bigDataHostPointer[bigDataSize - i - 1] = i;\n }\n\n checkCudaErrors(cudaMalloc((void **)&dataDevicePointer, dataSize * sizeof(int)));\n checkCudaErrors(cudaMalloc((void **)&bigDataDevicePointer, bigDataSize * sizeof(int)));\n checkCudaErrors(\n cudaMemcpyAsync(dataDevicePointer, dataHostPointer, dataSize * sizeof(int), cudaMemcpyHostToDevice, stream));\n checkCudaErrors(cudaMemcpyAsync(\n bigDataDevicePointer, bigDataHostPointer, bigDataSize * sizeof(int), cudaMemcpyHostToDevice, stream));\n\n // Make a window for the buffer of interest\n accessPolicyWindow.base_ptr = (void *)dataDevicePointer;\n accessPolicyWindow.num_bytes = dataSize * sizeof(int);\n accessPolicyWindow.hitRatio = 1.f;\n accessPolicyWindow.hitProp = cudaAccessPropertyPersisting;\n accessPolicyWindow.missProp = cudaAccessPropertyNormal;\n streamAttrValue.accessPolicyWindow = accessPolicyWindow;\n\n // Assign window to stream\n checkCudaErrors(cudaStreamSetAttribute(stream, streamAttrID, &streamAttrValue));\n\n // Demote any previous persisting lines\n checkCudaErrors(cudaCtxResetPersistingL2Cache());\n\n checkCudaErrors(cudaStreamSynchronize(stream));\n kernCacheSegmentTest<<>>(\n dataDevicePointer, dataSize, bigDataDevicePointer, bigDataSize, 0xAFFFF);\n\n checkCudaErrors(cudaStreamSynchronize(stream));\n // check if kernel execution generated and error\n getLastCudaError(\"Kernel execution failed\");\n\n // Free memory\n checkCudaErrors(cudaFreeHost(dataHostPointer));\n checkCudaErrors(cudaFreeHost(bigDataHostPointer));\n checkCudaErrors(cudaFree(dataDevicePointer));\n checkCudaErrors(cudaFree(bigDataDevicePointer));\n\n sdkStopTimer(&timer);\n printf(\"Processing time: %f (ms)\\n\", sdkGetTimerValue(&timer));\n sdkDeleteTimer(&timer);\n\n exit(bTestResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// includes CUDA\n#include \n\n// includes, project\n#include \n#include // helper functions for SDK examples\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\nvoid runTest(int argc, char **argv);\n\nhipAccessPolicyWindow initAccessPolicyWindow(void)\n{\n hipAccessPolicyWindow accessPolicyWindow = {0};\n accessPolicyWindow.base_ptr = (void *)0;\n accessPolicyWindow.num_bytes = 0;\n accessPolicyWindow.hitRatio = 0.f;\n accessPolicyWindow.hitProp = hipAccessPropertyNormal;\n accessPolicyWindow.missProp = hipAccessPropertyStreaming;\n return accessPolicyWindow;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Simple test kernel for device functionality\n//! @param data input data in global memory\n//! @param dataSize input data size\n//! @param bigData input bigData in global memory\n//! @param bigDataSize input bigData size\n//! @param hitcount how many data access are done within block\n////////////////////////////////////////////////////////////////////////////////\nstatic __global__ void kernCacheSegmentTest(int *data, int dataSize, int *trash, int bigDataSize, int hitCount)\n{\n __shared__ unsigned int hit;\n int row = blockIdx.y * blockDim.y + threadIdx.y;\n int col = blockIdx.x * blockDim.x + threadIdx.x;\n int tID = row * blockDim.y + col;\n uint32_t psRand = tID;\n\n atomicExch(&hit, 0);\n __syncthreads();\n while (hit < hitCount) {\n psRand ^= psRand << 13;\n psRand ^= psRand >> 17;\n psRand ^= psRand << 5;\n\n int idx = tID - psRand;\n if (idx < 0) {\n idx = -idx;\n }\n\n if ((tID % 2) == 0) {\n data[psRand % dataSize] = data[psRand % dataSize] + data[idx % dataSize];\n }\n else {\n trash[psRand % bigDataSize] = trash[psRand % bigDataSize] + trash[idx % bigDataSize];\n }\n\n atomicAdd(&hit, 1);\n }\n}\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv) { runTest(argc, argv); }\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nvoid runTest(int argc, char **argv)\n{\n bool bTestResult = true;\n hipAccessPolicyWindow accessPolicyWindow;\n hipDeviceProp_t deviceProp;\n hipLaunchAttributeValue streamAttrValue;\n hipStream_t stream;\n hipLaunchAttributeID streamAttrID;\n dim3 threads(32, 32);\n int *dataDevicePointer;\n int *dataHostPointer;\n int dataSize;\n int *bigDataDevicePointer;\n int *bigDataHostPointer;\n int bigDataSize;\n StopWatchInterface *timer = 0;\n\n printf(\"%s Starting...\\n\\n\", argv[0]);\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n int devID = findCudaDevice(argc, (const char **)argv);\n sdkCreateTimer(&timer);\n sdkStartTimer(&timer);\n // Get device properties\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, devID));\n dim3 blocks(deviceProp.maxGridSize[1], 1);\n\n // Make sure device the l2 optimization\n if (deviceProp.persistingL2CacheMaxSize == 0) {\n printf(\"Waiving execution as device %d does not support persisting L2 \"\n \"Caching\\n\",\n devID);\n exit(EXIT_WAIVED);\n }\n\n // Create stream to assiocate with window\n checkCudaErrors(hipStreamCreate(&stream));\n\n // Set the amount of l2 cache that will be persisting to maximum the device\n // can support\n checkCudaErrors(hipDeviceSetLimit(cudaLimitPersistingL2CacheSize, deviceProp.persistingL2CacheMaxSize));\n\n // Stream attribute to set\n streamAttrID = hipLaunchAttributeAccessPolicyWindow;\n\n // Default window\n streamAttrValue.accessPolicyWindow = initAccessPolicyWindow();\n accessPolicyWindow = initAccessPolicyWindow();\n\n // Allocate size of both buffers\n bigDataSize = (deviceProp.l2CacheSize * 4) / sizeof(int);\n dataSize = (deviceProp.l2CacheSize / 4) / sizeof(int);\n\n // Allocate data\n checkCudaErrors(hipHostMalloc(&dataHostPointer, dataSize * sizeof(int)));\n checkCudaErrors(hipHostMalloc(&bigDataHostPointer, bigDataSize * sizeof(int)));\n\n for (int i = 0; i < bigDataSize; ++i) {\n if (i < dataSize) {\n dataHostPointer[i] = i;\n }\n\n bigDataHostPointer[bigDataSize - i - 1] = i;\n }\n\n checkCudaErrors(hipMalloc((void **)&dataDevicePointer, dataSize * sizeof(int)));\n checkCudaErrors(hipMalloc((void **)&bigDataDevicePointer, bigDataSize * sizeof(int)));\n checkCudaErrors(\n hipMemcpyAsync(dataDevicePointer, dataHostPointer, dataSize * sizeof(int), hipMemcpyHostToDevice, stream));\n checkCudaErrors(hipMemcpyAsync(\n bigDataDevicePointer, bigDataHostPointer, bigDataSize * sizeof(int), hipMemcpyHostToDevice, stream));\n\n // Make a window for the buffer of interest\n accessPolicyWindow.base_ptr = (void *)dataDevicePointer;\n accessPolicyWindow.num_bytes = dataSize * sizeof(int);\n accessPolicyWindow.hitRatio = 1.f;\n accessPolicyWindow.hitProp = hipAccessPropertyPersisting;\n accessPolicyWindow.missProp = hipAccessPropertyNormal;\n streamAttrValue.accessPolicyWindow = accessPolicyWindow;\n\n // Assign window to stream\n checkCudaErrors(hipStreamSetAttribute(stream, streamAttrID, &streamAttrValue));\n\n // Demote any previous persisting lines\n checkCudaErrors(cudaCtxResetPersistingL2Cache());\n\n checkCudaErrors(hipStreamSynchronize(stream));\n kernCacheSegmentTest<<>>(\n dataDevicePointer, dataSize, bigDataDevicePointer, bigDataSize, 0xAFFFF);\n\n checkCudaErrors(hipStreamSynchronize(stream));\n // check if kernel execution generated and error\n getLastCudaError(\"Kernel execution failed\");\n\n // Free memory\n checkCudaErrors(hipHostFree(dataHostPointer));\n checkCudaErrors(hipHostFree(bigDataHostPointer));\n checkCudaErrors(hipFree(dataDevicePointer));\n checkCudaErrors(hipFree(bigDataDevicePointer));\n\n sdkStopTimer(&timer);\n printf(\"Processing time: %f (ms)\\n\", sdkGetTimerValue(&timer));\n sdkDeleteTimer(&timer);\n\n exit(bTestResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleAttributes/simpleAttributes.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n\nnamespace cg = cooperative_groups;\n#include \n#include \n\n#include \"mergeSort_common.h\"\n\ninline __device__ void Comparator(uint &keyA, uint &valA, uint &keyB, uint &valB, uint arrowDir)\n{\n uint t;\n\n if ((keyA > keyB) == arrowDir) {\n t = keyA;\n keyA = keyB;\n keyB = t;\n t = valA;\n valA = valB;\n valB = t;\n }\n}\n\n__global__ void\nbitonicSortSharedKernel(uint *d_DstKey, uint *d_DstVal, uint *d_SrcKey, uint *d_SrcVal, uint arrayLength, uint sortDir)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Shared memory storage for one or more short vectors\n __shared__ uint s_key[SHARED_SIZE_LIMIT];\n __shared__ uint s_val[SHARED_SIZE_LIMIT];\n\n // Offset to the beginning of subbatch and load data\n d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n s_key[threadIdx.x + 0] = d_SrcKey[0];\n s_val[threadIdx.x + 0] = d_SrcVal[0];\n s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];\n s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];\n\n for (uint size = 2; size < arrayLength; size <<= 1) {\n // Bitonic merge\n uint dir = (threadIdx.x & (size / 2)) != 0;\n\n for (uint stride = size / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], dir);\n }\n }\n\n // ddd == sortDir for the last bitonic merge step\n {\n for (uint stride = arrayLength / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], sortDir);\n }\n }\n\n cg::sync(cta);\n d_DstKey[0] = s_key[threadIdx.x + 0];\n d_DstVal[0] = s_val[threadIdx.x + 0];\n d_DstKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n d_DstVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n}\n\n// Helper function (also used by odd-even merge sort)\nextern \"C\" uint factorRadix2(uint *log2L, uint L)\n{\n if (!L) {\n *log2L = 0;\n return 0;\n }\n else {\n for (*log2L = 0; (L & 1) == 0; L >>= 1, *log2L++)\n ;\n\n return L;\n }\n}\n\nextern \"C\" void bitonicSortShared(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint batchSize,\n uint arrayLength,\n uint sortDir)\n{\n // Nothing to sort\n if (arrayLength < 2) {\n return;\n }\n\n // Only power-of-two array lengths are supported by this implementation\n uint log2L;\n uint factorizationRemainder = factorRadix2(&log2L, arrayLength);\n assert(factorizationRemainder == 1);\n\n uint blockCount = batchSize * arrayLength / SHARED_SIZE_LIMIT;\n uint threadCount = SHARED_SIZE_LIMIT / 2;\n\n assert(arrayLength <= SHARED_SIZE_LIMIT);\n assert((batchSize * arrayLength) % SHARED_SIZE_LIMIT == 0);\n\n bitonicSortSharedKernel<<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, arrayLength, sortDir);\n getLastCudaError(\"bitonicSortSharedKernel<<<>>> failed!\\n\");\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Merge step 3: merge elementary intervals\n////////////////////////////////////////////////////////////////////////////////\nstatic inline __host__ __device__ uint iDivUp(uint a, uint b) { return ((a % b) == 0) ? (a / b) : (a / b + 1); }\n\nstatic inline __host__ __device__ uint getSampleCount(uint dividend) { return iDivUp(dividend, SAMPLE_STRIDE); }\n\ntemplate \nstatic inline __device__ void\nComparatorExtended(uint &keyA, uint &valA, uint &flagA, uint &keyB, uint &valB, uint &flagB, uint arrowDir)\n{\n uint t;\n\n if ((!(flagA || flagB) && ((keyA > keyB) == arrowDir)) || ((arrowDir == sortDir) && (flagA == 1))\n || ((arrowDir != sortDir) && (flagB == 1))) {\n t = keyA;\n keyA = keyB;\n keyB = t;\n t = valA;\n valA = valB;\n valB = t;\n t = flagA;\n flagA = flagB;\n flagB = t;\n }\n}\n\ntemplate \n__global__ void bitonicMergeElementaryIntervalsKernel(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint *d_LimitsA,\n uint *d_LimitsB,\n uint stride,\n uint N)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ uint s_key[2 * SAMPLE_STRIDE];\n __shared__ uint s_val[2 * SAMPLE_STRIDE];\n __shared__ uint s_inf[2 * SAMPLE_STRIDE];\n\n const uint intervalI = blockIdx.x & ((2 * stride) / SAMPLE_STRIDE - 1);\n const uint segmentBase = (blockIdx.x - intervalI) * SAMPLE_STRIDE;\n d_SrcKey += segmentBase;\n d_SrcVal += segmentBase;\n d_DstKey += segmentBase;\n d_DstVal += segmentBase;\n\n // Set up threadblock-wide parameters\n __shared__ uint startSrcA, lenSrcA, startSrcB, lenSrcB, startDst;\n\n if (threadIdx.x == 0) {\n uint segmentElementsA = stride;\n uint segmentElementsB = umin(stride, N - segmentBase - stride);\n uint segmentSamplesA = stride / SAMPLE_STRIDE;\n uint segmentSamplesB = getSampleCount(segmentElementsB);\n uint segmentSamples = segmentSamplesA + segmentSamplesB;\n\n startSrcA = d_LimitsA[blockIdx.x];\n startSrcB = d_LimitsB[blockIdx.x];\n startDst = startSrcA + startSrcB;\n\n uint endSrcA = (intervalI + 1 < segmentSamples) ? d_LimitsA[blockIdx.x + 1] : segmentElementsA;\n uint endSrcB = (intervalI + 1 < segmentSamples) ? d_LimitsB[blockIdx.x + 1] : segmentElementsB;\n lenSrcA = endSrcA - startSrcA;\n lenSrcB = endSrcB - startSrcB;\n }\n\n s_inf[threadIdx.x + 0] = 1;\n s_inf[threadIdx.x + SAMPLE_STRIDE] = 1;\n\n // Load input data\n cg::sync(cta);\n\n if (threadIdx.x < lenSrcA) {\n s_key[threadIdx.x] = d_SrcKey[0 + startSrcA + threadIdx.x];\n s_val[threadIdx.x] = d_SrcVal[0 + startSrcA + threadIdx.x];\n s_inf[threadIdx.x] = 0;\n }\n\n // Prepare for bitonic merge by inversing the ordering\n if (threadIdx.x < lenSrcB) {\n s_key[2 * SAMPLE_STRIDE - 1 - threadIdx.x] = d_SrcKey[stride + startSrcB + threadIdx.x];\n s_val[2 * SAMPLE_STRIDE - 1 - threadIdx.x] = d_SrcVal[stride + startSrcB + threadIdx.x];\n s_inf[2 * SAMPLE_STRIDE - 1 - threadIdx.x] = 0;\n }\n\n //\"Extended\" bitonic merge\n for (uint stride = SAMPLE_STRIDE; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n ComparatorExtended(s_key[pos + 0],\n s_val[pos + 0],\n s_inf[pos + 0],\n s_key[pos + stride],\n s_val[pos + stride],\n s_inf[pos + stride],\n sortDir);\n }\n\n // Store sorted data\n cg::sync(cta);\n d_DstKey += startDst;\n d_DstVal += startDst;\n\n if (threadIdx.x < lenSrcA) {\n d_DstKey[threadIdx.x] = s_key[threadIdx.x];\n d_DstVal[threadIdx.x] = s_val[threadIdx.x];\n }\n\n if (threadIdx.x < lenSrcB) {\n d_DstKey[lenSrcA + threadIdx.x] = s_key[lenSrcA + threadIdx.x];\n d_DstVal[lenSrcA + threadIdx.x] = s_val[lenSrcA + threadIdx.x];\n }\n}\n\nextern \"C\" void bitonicMergeElementaryIntervals(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint *d_LimitsA,\n uint *d_LimitsB,\n uint stride,\n uint N,\n uint sortDir)\n{\n uint lastSegmentElements = N % (2 * stride);\n\n uint mergePairs = (lastSegmentElements > stride) ? getSampleCount(N) : (N - lastSegmentElements) / SAMPLE_STRIDE;\n\n if (sortDir) {\n bitonicMergeElementaryIntervalsKernel<1U>\n <<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, d_LimitsA, d_LimitsB, stride, N);\n getLastCudaError(\"mergeElementaryIntervalsKernel<1> failed\\n\");\n }\n else {\n bitonicMergeElementaryIntervalsKernel<0U>\n <<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, d_LimitsA, d_LimitsB, stride, N);\n getLastCudaError(\"mergeElementaryIntervalsKernel<0> failed\\n\");\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n\nnamespace cg = cooperative_groups;\n#include \n#include \n\n#include \"mergeSort_common.h\"\n\ninline __device__ void Comparator(uint &keyA, uint &valA, uint &keyB, uint &valB, uint arrowDir)\n{\n uint t;\n\n if ((keyA > keyB) == arrowDir) {\n t = keyA;\n keyA = keyB;\n keyB = t;\n t = valA;\n valA = valB;\n valB = t;\n }\n}\n\n__global__ void\nbitonicSortSharedKernel(uint *d_DstKey, uint *d_DstVal, uint *d_SrcKey, uint *d_SrcVal, uint arrayLength, uint sortDir)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Shared memory storage for one or more short vectors\n __shared__ uint s_key[SHARED_SIZE_LIMIT];\n __shared__ uint s_val[SHARED_SIZE_LIMIT];\n\n // Offset to the beginning of subbatch and load data\n d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n s_key[threadIdx.x + 0] = d_SrcKey[0];\n s_val[threadIdx.x + 0] = d_SrcVal[0];\n s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];\n s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];\n\n for (uint size = 2; size < arrayLength; size <<= 1) {\n // Bitonic merge\n uint dir = (threadIdx.x & (size / 2)) != 0;\n\n for (uint stride = size / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], dir);\n }\n }\n\n // ddd == sortDir for the last bitonic merge step\n {\n for (uint stride = arrayLength / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], sortDir);\n }\n }\n\n cg::sync(cta);\n d_DstKey[0] = s_key[threadIdx.x + 0];\n d_DstVal[0] = s_val[threadIdx.x + 0];\n d_DstKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n d_DstVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n}\n\n// Helper function (also used by odd-even merge sort)\nextern \"C\" uint factorRadix2(uint *log2L, uint L)\n{\n if (!L) {\n *log2L = 0;\n return 0;\n }\n else {\n for (*log2L = 0; (L & 1) == 0; L >>= 1, *log2L++)\n ;\n\n return L;\n }\n}\n\nextern \"C\" void bitonicSortShared(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint batchSize,\n uint arrayLength,\n uint sortDir)\n{\n // Nothing to sort\n if (arrayLength < 2) {\n return;\n }\n\n // Only power-of-two array lengths are supported by this implementation\n uint log2L;\n uint factorizationRemainder = factorRadix2(&log2L, arrayLength);\n assert(factorizationRemainder == 1);\n\n uint blockCount = batchSize * arrayLength / SHARED_SIZE_LIMIT;\n uint threadCount = SHARED_SIZE_LIMIT / 2;\n\n assert(arrayLength <= SHARED_SIZE_LIMIT);\n assert((batchSize * arrayLength) % SHARED_SIZE_LIMIT == 0);\n\n bitonicSortSharedKernel<<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, arrayLength, sortDir);\n getLastCudaError(\"bitonicSortSharedKernel<<<>>> failed!\\n\");\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Merge step 3: merge elementary intervals\n////////////////////////////////////////////////////////////////////////////////\nstatic inline __host__ __device__ uint iDivUp(uint a, uint b) { return ((a % b) == 0) ? (a / b) : (a / b + 1); }\n\nstatic inline __host__ __device__ uint getSampleCount(uint dividend) { return iDivUp(dividend, SAMPLE_STRIDE); }\n\ntemplate \nstatic inline __device__ void\nComparatorExtended(uint &keyA, uint &valA, uint &flagA, uint &keyB, uint &valB, uint &flagB, uint arrowDir)\n{\n uint t;\n\n if ((!(flagA || flagB) && ((keyA > keyB) == arrowDir)) || ((arrowDir == sortDir) && (flagA == 1))\n || ((arrowDir != sortDir) && (flagB == 1))) {\n t = keyA;\n keyA = keyB;\n keyB = t;\n t = valA;\n valA = valB;\n valB = t;\n t = flagA;\n flagA = flagB;\n flagB = t;\n }\n}\n\ntemplate \n__global__ void bitonicMergeElementaryIntervalsKernel(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint *d_LimitsA,\n uint *d_LimitsB,\n uint stride,\n uint N)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ uint s_key[2 * SAMPLE_STRIDE];\n __shared__ uint s_val[2 * SAMPLE_STRIDE];\n __shared__ uint s_inf[2 * SAMPLE_STRIDE];\n\n const uint intervalI = blockIdx.x & ((2 * stride) / SAMPLE_STRIDE - 1);\n const uint segmentBase = (blockIdx.x - intervalI) * SAMPLE_STRIDE;\n d_SrcKey += segmentBase;\n d_SrcVal += segmentBase;\n d_DstKey += segmentBase;\n d_DstVal += segmentBase;\n\n // Set up threadblock-wide parameters\n __shared__ uint startSrcA, lenSrcA, startSrcB, lenSrcB, startDst;\n\n if (threadIdx.x == 0) {\n uint segmentElementsA = stride;\n uint segmentElementsB = umin(stride, N - segmentBase - stride);\n uint segmentSamplesA = stride / SAMPLE_STRIDE;\n uint segmentSamplesB = getSampleCount(segmentElementsB);\n uint segmentSamples = segmentSamplesA + segmentSamplesB;\n\n startSrcA = d_LimitsA[blockIdx.x];\n startSrcB = d_LimitsB[blockIdx.x];\n startDst = startSrcA + startSrcB;\n\n uint endSrcA = (intervalI + 1 < segmentSamples) ? d_LimitsA[blockIdx.x + 1] : segmentElementsA;\n uint endSrcB = (intervalI + 1 < segmentSamples) ? d_LimitsB[blockIdx.x + 1] : segmentElementsB;\n lenSrcA = endSrcA - startSrcA;\n lenSrcB = endSrcB - startSrcB;\n }\n\n s_inf[threadIdx.x + 0] = 1;\n s_inf[threadIdx.x + SAMPLE_STRIDE] = 1;\n\n // Load input data\n cg::sync(cta);\n\n if (threadIdx.x < lenSrcA) {\n s_key[threadIdx.x] = d_SrcKey[0 + startSrcA + threadIdx.x];\n s_val[threadIdx.x] = d_SrcVal[0 + startSrcA + threadIdx.x];\n s_inf[threadIdx.x] = 0;\n }\n\n // Prepare for bitonic merge by inversing the ordering\n if (threadIdx.x < lenSrcB) {\n s_key[2 * SAMPLE_STRIDE - 1 - threadIdx.x] = d_SrcKey[stride + startSrcB + threadIdx.x];\n s_val[2 * SAMPLE_STRIDE - 1 - threadIdx.x] = d_SrcVal[stride + startSrcB + threadIdx.x];\n s_inf[2 * SAMPLE_STRIDE - 1 - threadIdx.x] = 0;\n }\n\n //\"Extended\" bitonic merge\n for (uint stride = SAMPLE_STRIDE; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n ComparatorExtended(s_key[pos + 0],\n s_val[pos + 0],\n s_inf[pos + 0],\n s_key[pos + stride],\n s_val[pos + stride],\n s_inf[pos + stride],\n sortDir);\n }\n\n // Store sorted data\n cg::sync(cta);\n d_DstKey += startDst;\n d_DstVal += startDst;\n\n if (threadIdx.x < lenSrcA) {\n d_DstKey[threadIdx.x] = s_key[threadIdx.x];\n d_DstVal[threadIdx.x] = s_val[threadIdx.x];\n }\n\n if (threadIdx.x < lenSrcB) {\n d_DstKey[lenSrcA + threadIdx.x] = s_key[lenSrcA + threadIdx.x];\n d_DstVal[lenSrcA + threadIdx.x] = s_val[lenSrcA + threadIdx.x];\n }\n}\n\nextern \"C\" void bitonicMergeElementaryIntervals(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint *d_LimitsA,\n uint *d_LimitsB,\n uint stride,\n uint N,\n uint sortDir)\n{\n uint lastSegmentElements = N % (2 * stride);\n\n uint mergePairs = (lastSegmentElements > stride) ? getSampleCount(N) : (N - lastSegmentElements) / SAMPLE_STRIDE;\n\n if (sortDir) {\n bitonicMergeElementaryIntervalsKernel<1U>\n <<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, d_LimitsA, d_LimitsB, stride, N);\n getLastCudaError(\"mergeElementaryIntervalsKernel<1> failed\\n\");\n }\n else {\n bitonicMergeElementaryIntervalsKernel<0U>\n <<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, d_LimitsA, d_LimitsB, stride, N);\n getLastCudaError(\"mergeElementaryIntervalsKernel<0> failed\\n\");\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/mergeSort/bitonic.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * Based on \"Designing efficient sorting algorithms for manycore GPUs\"\n * by Nadathur Satish, Mark Harris, and Michael Garland\n * http://mgarland.org/files/papers/gpusort-ipdps09.pdf\n *\n * Victor Podlozhnyuk 09/24/2009\n */\n\n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \n\n#include \"mergeSort_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Helper functions\n////////////////////////////////////////////////////////////////////////////////\nstatic inline __host__ __device__ uint iDivUp(uint a, uint b) { return ((a % b) == 0) ? (a / b) : (a / b + 1); }\n\nstatic inline __host__ __device__ uint getSampleCount(uint dividend) { return iDivUp(dividend, SAMPLE_STRIDE); }\n\n#define W (sizeof(uint) * 8)\nstatic inline __device__ uint nextPowerOfTwo(uint x)\n{\n /*\n --x;\n x |= x >> 1;\n x |= x >> 2;\n x |= x >> 4;\n x |= x >> 8;\n x |= x >> 16;\n return ++x;\n */\n return 1U << (W - __clz(x - 1));\n}\n\ntemplate static inline __device__ uint binarySearchInclusive(uint val, uint *data, uint L, uint stride)\n{\n if (L == 0) {\n return 0;\n }\n\n uint pos = 0;\n\n for (; stride > 0; stride >>= 1) {\n uint newPos = umin(pos + stride, L);\n\n if ((sortDir && (data[newPos - 1] <= val)) || (!sortDir && (data[newPos - 1] >= val))) {\n pos = newPos;\n }\n }\n\n return pos;\n}\n\ntemplate static inline __device__ uint binarySearchExclusive(uint val, uint *data, uint L, uint stride)\n{\n if (L == 0) {\n return 0;\n }\n\n uint pos = 0;\n\n for (; stride > 0; stride >>= 1) {\n uint newPos = umin(pos + stride, L);\n\n if ((sortDir && (data[newPos - 1] < val)) || (!sortDir && (data[newPos - 1] > val))) {\n pos = newPos;\n }\n }\n\n return pos;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Bottom-level merge sort (binary search-based)\n////////////////////////////////////////////////////////////////////////////////\ntemplate \n__global__ void mergeSortSharedKernel(uint *d_DstKey, uint *d_DstVal, uint *d_SrcKey, uint *d_SrcVal, uint arrayLength)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ uint s_key[SHARED_SIZE_LIMIT];\n __shared__ uint s_val[SHARED_SIZE_LIMIT];\n\n d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n s_key[threadIdx.x + 0] = d_SrcKey[0];\n s_val[threadIdx.x + 0] = d_SrcVal[0];\n s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];\n s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];\n\n for (uint stride = 1; stride < arrayLength; stride <<= 1) {\n uint lPos = threadIdx.x & (stride - 1);\n uint *baseKey = s_key + 2 * (threadIdx.x - lPos);\n uint *baseVal = s_val + 2 * (threadIdx.x - lPos);\n\n cg::sync(cta);\n uint keyA = baseKey[lPos + 0];\n uint valA = baseVal[lPos + 0];\n uint keyB = baseKey[lPos + stride];\n uint valB = baseVal[lPos + stride];\n uint posA = binarySearchExclusive(keyA, baseKey + stride, stride, stride) + lPos;\n uint posB = binarySearchInclusive(keyB, baseKey + 0, stride, stride) + lPos;\n\n cg::sync(cta);\n baseKey[posA] = keyA;\n baseVal[posA] = valA;\n baseKey[posB] = keyB;\n baseVal[posB] = valB;\n }\n\n cg::sync(cta);\n d_DstKey[0] = s_key[threadIdx.x + 0];\n d_DstVal[0] = s_val[threadIdx.x + 0];\n d_DstKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n d_DstVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n}\n\nstatic void mergeSortShared(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint batchSize,\n uint arrayLength,\n uint sortDir)\n{\n if (arrayLength < 2) {\n return;\n }\n\n assert(SHARED_SIZE_LIMIT % arrayLength == 0);\n assert(((batchSize * arrayLength) % SHARED_SIZE_LIMIT) == 0);\n uint blockCount = batchSize * arrayLength / SHARED_SIZE_LIMIT;\n uint threadCount = SHARED_SIZE_LIMIT / 2;\n\n if (sortDir) {\n mergeSortSharedKernel<1U><<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, arrayLength);\n getLastCudaError(\"mergeSortShared<1><<<>>> failed\\n\");\n }\n else {\n mergeSortSharedKernel<0U><<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, arrayLength);\n getLastCudaError(\"mergeSortShared<0><<<>>> failed\\n\");\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Merge step 1: generate sample ranks\n////////////////////////////////////////////////////////////////////////////////\ntemplate \n__global__ void\ngenerateSampleRanksKernel(uint *d_RanksA, uint *d_RanksB, uint *d_SrcKey, uint stride, uint N, uint threadCount)\n{\n uint pos = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (pos >= threadCount) {\n return;\n }\n\n const uint i = pos & ((stride / SAMPLE_STRIDE) - 1);\n const uint segmentBase = (pos - i) * (2 * SAMPLE_STRIDE);\n d_SrcKey += segmentBase;\n d_RanksA += segmentBase / SAMPLE_STRIDE;\n d_RanksB += segmentBase / SAMPLE_STRIDE;\n\n const uint segmentElementsA = stride;\n const uint segmentElementsB = umin(stride, N - segmentBase - stride);\n const uint segmentSamplesA = getSampleCount(segmentElementsA);\n const uint segmentSamplesB = getSampleCount(segmentElementsB);\n\n if (i < segmentSamplesA) {\n d_RanksA[i] = i * SAMPLE_STRIDE;\n d_RanksB[i] = binarySearchExclusive(\n d_SrcKey[i * SAMPLE_STRIDE], d_SrcKey + stride, segmentElementsB, nextPowerOfTwo(segmentElementsB));\n }\n\n if (i < segmentSamplesB) {\n d_RanksB[(stride / SAMPLE_STRIDE) + i] = i * SAMPLE_STRIDE;\n d_RanksA[(stride / SAMPLE_STRIDE) + i] = binarySearchInclusive(\n d_SrcKey[stride + i * SAMPLE_STRIDE], d_SrcKey + 0, segmentElementsA, nextPowerOfTwo(segmentElementsA));\n }\n}\n\nstatic void generateSampleRanks(uint *d_RanksA, uint *d_RanksB, uint *d_SrcKey, uint stride, uint N, uint sortDir)\n{\n uint lastSegmentElements = N % (2 * stride);\n uint threadCount = (lastSegmentElements > stride) ? (N + 2 * stride - lastSegmentElements) / (2 * SAMPLE_STRIDE)\n : (N - lastSegmentElements) / (2 * SAMPLE_STRIDE);\n\n if (sortDir) {\n generateSampleRanksKernel<1U>\n <<>>(d_RanksA, d_RanksB, d_SrcKey, stride, N, threadCount);\n getLastCudaError(\"generateSampleRanksKernel<1U><<<>>> failed\\n\");\n }\n else {\n generateSampleRanksKernel<0U>\n <<>>(d_RanksA, d_RanksB, d_SrcKey, stride, N, threadCount);\n getLastCudaError(\"generateSampleRanksKernel<0U><<<>>> failed\\n\");\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Merge step 2: generate sample ranks and indices\n////////////////////////////////////////////////////////////////////////////////\n__global__ void mergeRanksAndIndicesKernel(uint *d_Limits, uint *d_Ranks, uint stride, uint N, uint threadCount)\n{\n uint pos = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (pos >= threadCount) {\n return;\n }\n\n const uint i = pos & ((stride / SAMPLE_STRIDE) - 1);\n const uint segmentBase = (pos - i) * (2 * SAMPLE_STRIDE);\n d_Ranks += (pos - i) * 2;\n d_Limits += (pos - i) * 2;\n\n const uint segmentElementsA = stride;\n const uint segmentElementsB = umin(stride, N - segmentBase - stride);\n const uint segmentSamplesA = getSampleCount(segmentElementsA);\n const uint segmentSamplesB = getSampleCount(segmentElementsB);\n\n if (i < segmentSamplesA) {\n uint dstPos = binarySearchExclusive<1U>(\n d_Ranks[i], d_Ranks + segmentSamplesA, segmentSamplesB, nextPowerOfTwo(segmentSamplesB))\n + i;\n d_Limits[dstPos] = d_Ranks[i];\n }\n\n if (i < segmentSamplesB) {\n uint dstPos = binarySearchInclusive<1U>(\n d_Ranks[segmentSamplesA + i], d_Ranks, segmentSamplesA, nextPowerOfTwo(segmentSamplesA))\n + i;\n d_Limits[dstPos] = d_Ranks[segmentSamplesA + i];\n }\n}\n\nstatic void mergeRanksAndIndices(uint *d_LimitsA, uint *d_LimitsB, uint *d_RanksA, uint *d_RanksB, uint stride, uint N)\n{\n uint lastSegmentElements = N % (2 * stride);\n uint threadCount = (lastSegmentElements > stride) ? (N + 2 * stride - lastSegmentElements) / (2 * SAMPLE_STRIDE)\n : (N - lastSegmentElements) / (2 * SAMPLE_STRIDE);\n\n mergeRanksAndIndicesKernel<<>>(d_LimitsA, d_RanksA, stride, N, threadCount);\n getLastCudaError(\"mergeRanksAndIndicesKernel(A)<<<>>> failed\\n\");\n\n mergeRanksAndIndicesKernel<<>>(d_LimitsB, d_RanksB, stride, N, threadCount);\n getLastCudaError(\"mergeRanksAndIndicesKernel(B)<<<>>> failed\\n\");\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Merge step 3: merge elementary intervals\n////////////////////////////////////////////////////////////////////////////////\ntemplate \ninline __device__ void merge(uint *dstKey,\n uint *dstVal,\n uint *srcAKey,\n uint *srcAVal,\n uint *srcBKey,\n uint *srcBVal,\n uint lenA,\n uint nPowTwoLenA,\n uint lenB,\n uint nPowTwoLenB,\n cg::thread_block cta)\n{\n uint keyA, valA, keyB, valB, dstPosA, dstPosB;\n\n if (threadIdx.x < lenA) {\n keyA = srcAKey[threadIdx.x];\n valA = srcAVal[threadIdx.x];\n dstPosA = binarySearchExclusive(keyA, srcBKey, lenB, nPowTwoLenB) + threadIdx.x;\n }\n\n if (threadIdx.x < lenB) {\n keyB = srcBKey[threadIdx.x];\n valB = srcBVal[threadIdx.x];\n dstPosB = binarySearchInclusive(keyB, srcAKey, lenA, nPowTwoLenA) + threadIdx.x;\n }\n\n cg::sync(cta);\n\n if (threadIdx.x < lenA) {\n dstKey[dstPosA] = keyA;\n dstVal[dstPosA] = valA;\n }\n\n if (threadIdx.x < lenB) {\n dstKey[dstPosB] = keyB;\n dstVal[dstPosB] = valB;\n }\n}\n\ntemplate \n__global__ void mergeElementaryIntervalsKernel(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint *d_LimitsA,\n uint *d_LimitsB,\n uint stride,\n uint N)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ uint s_key[2 * SAMPLE_STRIDE];\n __shared__ uint s_val[2 * SAMPLE_STRIDE];\n\n const uint intervalI = blockIdx.x & ((2 * stride) / SAMPLE_STRIDE - 1);\n const uint segmentBase = (blockIdx.x - intervalI) * SAMPLE_STRIDE;\n d_SrcKey += segmentBase;\n d_SrcVal += segmentBase;\n d_DstKey += segmentBase;\n d_DstVal += segmentBase;\n\n // Set up threadblock-wide parameters\n __shared__ uint startSrcA, startSrcB, lenSrcA, lenSrcB, startDstA, startDstB;\n\n if (threadIdx.x == 0) {\n uint segmentElementsA = stride;\n uint segmentElementsB = umin(stride, N - segmentBase - stride);\n uint segmentSamplesA = getSampleCount(segmentElementsA);\n uint segmentSamplesB = getSampleCount(segmentElementsB);\n uint segmentSamples = segmentSamplesA + segmentSamplesB;\n\n startSrcA = d_LimitsA[blockIdx.x];\n startSrcB = d_LimitsB[blockIdx.x];\n uint endSrcA = (intervalI + 1 < segmentSamples) ? d_LimitsA[blockIdx.x + 1] : segmentElementsA;\n uint endSrcB = (intervalI + 1 < segmentSamples) ? d_LimitsB[blockIdx.x + 1] : segmentElementsB;\n lenSrcA = endSrcA - startSrcA;\n lenSrcB = endSrcB - startSrcB;\n startDstA = startSrcA + startSrcB;\n startDstB = startDstA + lenSrcA;\n }\n\n // Load main input data\n cg::sync(cta);\n\n if (threadIdx.x < lenSrcA) {\n s_key[threadIdx.x + 0] = d_SrcKey[0 + startSrcA + threadIdx.x];\n s_val[threadIdx.x + 0] = d_SrcVal[0 + startSrcA + threadIdx.x];\n }\n\n if (threadIdx.x < lenSrcB) {\n s_key[threadIdx.x + SAMPLE_STRIDE] = d_SrcKey[stride + startSrcB + threadIdx.x];\n s_val[threadIdx.x + SAMPLE_STRIDE] = d_SrcVal[stride + startSrcB + threadIdx.x];\n }\n\n // Merge data in shared memory\n cg::sync(cta);\n merge(s_key,\n s_val,\n s_key + 0,\n s_val + 0,\n s_key + SAMPLE_STRIDE,\n s_val + SAMPLE_STRIDE,\n lenSrcA,\n SAMPLE_STRIDE,\n lenSrcB,\n SAMPLE_STRIDE,\n cta);\n\n // Store merged data\n cg::sync(cta);\n\n if (threadIdx.x < lenSrcA) {\n d_DstKey[startDstA + threadIdx.x] = s_key[threadIdx.x];\n d_DstVal[startDstA + threadIdx.x] = s_val[threadIdx.x];\n }\n\n if (threadIdx.x < lenSrcB) {\n d_DstKey[startDstB + threadIdx.x] = s_key[lenSrcA + threadIdx.x];\n d_DstVal[startDstB + threadIdx.x] = s_val[lenSrcA + threadIdx.x];\n }\n}\n\nstatic void mergeElementaryIntervals(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint *d_LimitsA,\n uint *d_LimitsB,\n uint stride,\n uint N,\n uint sortDir)\n{\n uint lastSegmentElements = N % (2 * stride);\n uint mergePairs = (lastSegmentElements > stride) ? getSampleCount(N) : (N - lastSegmentElements) / SAMPLE_STRIDE;\n\n if (sortDir) {\n mergeElementaryIntervalsKernel<1U>\n <<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, d_LimitsA, d_LimitsB, stride, N);\n getLastCudaError(\"mergeElementaryIntervalsKernel<1> failed\\n\");\n }\n else {\n mergeElementaryIntervalsKernel<0U>\n <<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, d_LimitsA, d_LimitsB, stride, N);\n getLastCudaError(\"mergeElementaryIntervalsKernel<0> failed\\n\");\n }\n}\n\nextern \"C\" void bitonicSortShared(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint batchSize,\n uint arrayLength,\n uint sortDir);\n\nextern \"C\" void bitonicMergeElementaryIntervals(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint *d_LimitsA,\n uint *d_LimitsB,\n uint stride,\n uint N,\n uint sortDir);\n\nstatic uint *d_RanksA, *d_RanksB, *d_LimitsA, *d_LimitsB;\nstatic const uint MAX_SAMPLE_COUNT = 32768;\n\nextern \"C\" void initMergeSort(void)\n{\n checkCudaErrors(cudaMalloc((void **)&d_RanksA, MAX_SAMPLE_COUNT * sizeof(uint)));\n checkCudaErrors(cudaMalloc((void **)&d_RanksB, MAX_SAMPLE_COUNT * sizeof(uint)));\n checkCudaErrors(cudaMalloc((void **)&d_LimitsA, MAX_SAMPLE_COUNT * sizeof(uint)));\n checkCudaErrors(cudaMalloc((void **)&d_LimitsB, MAX_SAMPLE_COUNT * sizeof(uint)));\n}\n\nextern \"C\" void closeMergeSort(void)\n{\n checkCudaErrors(cudaFree(d_RanksA));\n checkCudaErrors(cudaFree(d_RanksB));\n checkCudaErrors(cudaFree(d_LimitsB));\n checkCudaErrors(cudaFree(d_LimitsA));\n}\n\nextern \"C\" void mergeSort(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_BufKey,\n uint *d_BufVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint N,\n uint sortDir)\n{\n uint stageCount = 0;\n\n for (uint stride = SHARED_SIZE_LIMIT; stride < N; stride <<= 1, stageCount++)\n ;\n\n uint *ikey, *ival, *okey, *oval;\n\n if (stageCount & 1) {\n ikey = d_BufKey;\n ival = d_BufVal;\n okey = d_DstKey;\n oval = d_DstVal;\n }\n else {\n ikey = d_DstKey;\n ival = d_DstVal;\n okey = d_BufKey;\n oval = d_BufVal;\n }\n\n assert(N <= (SAMPLE_STRIDE * MAX_SAMPLE_COUNT));\n assert(N % SHARED_SIZE_LIMIT == 0);\n mergeSortShared(ikey, ival, d_SrcKey, d_SrcVal, N / SHARED_SIZE_LIMIT, SHARED_SIZE_LIMIT, sortDir);\n\n for (uint stride = SHARED_SIZE_LIMIT; stride < N; stride <<= 1) {\n uint lastSegmentElements = N % (2 * stride);\n\n // Find sample ranks and prepare for limiters merge\n generateSampleRanks(d_RanksA, d_RanksB, ikey, stride, N, sortDir);\n\n // Merge ranks and indices\n mergeRanksAndIndices(d_LimitsA, d_LimitsB, d_RanksA, d_RanksB, stride, N);\n\n // Merge elementary intervals\n mergeElementaryIntervals(okey, oval, ikey, ival, d_LimitsA, d_LimitsB, stride, N, sortDir);\n\n if (lastSegmentElements <= stride) {\n // Last merge segment consists of a single array which just needs to be\n // passed through\n checkCudaErrors(cudaMemcpy(okey + (N - lastSegmentElements),\n ikey + (N - lastSegmentElements),\n lastSegmentElements * sizeof(uint),\n cudaMemcpyDeviceToDevice));\n checkCudaErrors(cudaMemcpy(oval + (N - lastSegmentElements),\n ival + (N - lastSegmentElements),\n lastSegmentElements * sizeof(uint),\n cudaMemcpyDeviceToDevice));\n }\n\n uint *t;\n t = ikey;\n ikey = okey;\n okey = t;\n t = ival;\n ival = oval;\n oval = t;\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * Based on \"Designing efficient sorting algorithms for manycore GPUs\"\n * by Nadathur Satish, Mark Harris, and Michael Garland\n * http://mgarland.org/files/papers/gpusort-ipdps09.pdf\n *\n * Victor Podlozhnyuk 09/24/2009\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \n\n#include \"mergeSort_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Helper functions\n////////////////////////////////////////////////////////////////////////////////\nstatic inline __host__ __device__ uint iDivUp(uint a, uint b) { return ((a % b) == 0) ? (a / b) : (a / b + 1); }\n\nstatic inline __host__ __device__ uint getSampleCount(uint dividend) { return iDivUp(dividend, SAMPLE_STRIDE); }\n\n#define W (sizeof(uint) * 8)\nstatic inline __device__ uint nextPowerOfTwo(uint x)\n{\n /*\n --x;\n x |= x >> 1;\n x |= x >> 2;\n x |= x >> 4;\n x |= x >> 8;\n x |= x >> 16;\n return ++x;\n */\n return 1U << (W - __clz(x - 1));\n}\n\ntemplate static inline __device__ uint binarySearchInclusive(uint val, uint *data, uint L, uint stride)\n{\n if (L == 0) {\n return 0;\n }\n\n uint pos = 0;\n\n for (; stride > 0; stride >>= 1) {\n uint newPos = umin(pos + stride, L);\n\n if ((sortDir && (data[newPos - 1] <= val)) || (!sortDir && (data[newPos - 1] >= val))) {\n pos = newPos;\n }\n }\n\n return pos;\n}\n\ntemplate static inline __device__ uint binarySearchExclusive(uint val, uint *data, uint L, uint stride)\n{\n if (L == 0) {\n return 0;\n }\n\n uint pos = 0;\n\n for (; stride > 0; stride >>= 1) {\n uint newPos = umin(pos + stride, L);\n\n if ((sortDir && (data[newPos - 1] < val)) || (!sortDir && (data[newPos - 1] > val))) {\n pos = newPos;\n }\n }\n\n return pos;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Bottom-level merge sort (binary search-based)\n////////////////////////////////////////////////////////////////////////////////\ntemplate \n__global__ void mergeSortSharedKernel(uint *d_DstKey, uint *d_DstVal, uint *d_SrcKey, uint *d_SrcVal, uint arrayLength)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ uint s_key[SHARED_SIZE_LIMIT];\n __shared__ uint s_val[SHARED_SIZE_LIMIT];\n\n d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n s_key[threadIdx.x + 0] = d_SrcKey[0];\n s_val[threadIdx.x + 0] = d_SrcVal[0];\n s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];\n s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];\n\n for (uint stride = 1; stride < arrayLength; stride <<= 1) {\n uint lPos = threadIdx.x & (stride - 1);\n uint *baseKey = s_key + 2 * (threadIdx.x - lPos);\n uint *baseVal = s_val + 2 * (threadIdx.x - lPos);\n\n cg::sync(cta);\n uint keyA = baseKey[lPos + 0];\n uint valA = baseVal[lPos + 0];\n uint keyB = baseKey[lPos + stride];\n uint valB = baseVal[lPos + stride];\n uint posA = binarySearchExclusive(keyA, baseKey + stride, stride, stride) + lPos;\n uint posB = binarySearchInclusive(keyB, baseKey + 0, stride, stride) + lPos;\n\n cg::sync(cta);\n baseKey[posA] = keyA;\n baseVal[posA] = valA;\n baseKey[posB] = keyB;\n baseVal[posB] = valB;\n }\n\n cg::sync(cta);\n d_DstKey[0] = s_key[threadIdx.x + 0];\n d_DstVal[0] = s_val[threadIdx.x + 0];\n d_DstKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n d_DstVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n}\n\nstatic void mergeSortShared(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint batchSize,\n uint arrayLength,\n uint sortDir)\n{\n if (arrayLength < 2) {\n return;\n }\n\n assert(SHARED_SIZE_LIMIT % arrayLength == 0);\n assert(((batchSize * arrayLength) % SHARED_SIZE_LIMIT) == 0);\n uint blockCount = batchSize * arrayLength / SHARED_SIZE_LIMIT;\n uint threadCount = SHARED_SIZE_LIMIT / 2;\n\n if (sortDir) {\n mergeSortSharedKernel<1U><<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, arrayLength);\n getLastCudaError(\"mergeSortShared<1><<<>>> failed\\n\");\n }\n else {\n mergeSortSharedKernel<0U><<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, arrayLength);\n getLastCudaError(\"mergeSortShared<0><<<>>> failed\\n\");\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Merge step 1: generate sample ranks\n////////////////////////////////////////////////////////////////////////////////\ntemplate \n__global__ void\ngenerateSampleRanksKernel(uint *d_RanksA, uint *d_RanksB, uint *d_SrcKey, uint stride, uint N, uint threadCount)\n{\n uint pos = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (pos >= threadCount) {\n return;\n }\n\n const uint i = pos & ((stride / SAMPLE_STRIDE) - 1);\n const uint segmentBase = (pos - i) * (2 * SAMPLE_STRIDE);\n d_SrcKey += segmentBase;\n d_RanksA += segmentBase / SAMPLE_STRIDE;\n d_RanksB += segmentBase / SAMPLE_STRIDE;\n\n const uint segmentElementsA = stride;\n const uint segmentElementsB = umin(stride, N - segmentBase - stride);\n const uint segmentSamplesA = getSampleCount(segmentElementsA);\n const uint segmentSamplesB = getSampleCount(segmentElementsB);\n\n if (i < segmentSamplesA) {\n d_RanksA[i] = i * SAMPLE_STRIDE;\n d_RanksB[i] = binarySearchExclusive(\n d_SrcKey[i * SAMPLE_STRIDE], d_SrcKey + stride, segmentElementsB, nextPowerOfTwo(segmentElementsB));\n }\n\n if (i < segmentSamplesB) {\n d_RanksB[(stride / SAMPLE_STRIDE) + i] = i * SAMPLE_STRIDE;\n d_RanksA[(stride / SAMPLE_STRIDE) + i] = binarySearchInclusive(\n d_SrcKey[stride + i * SAMPLE_STRIDE], d_SrcKey + 0, segmentElementsA, nextPowerOfTwo(segmentElementsA));\n }\n}\n\nstatic void generateSampleRanks(uint *d_RanksA, uint *d_RanksB, uint *d_SrcKey, uint stride, uint N, uint sortDir)\n{\n uint lastSegmentElements = N % (2 * stride);\n uint threadCount = (lastSegmentElements > stride) ? (N + 2 * stride - lastSegmentElements) / (2 * SAMPLE_STRIDE)\n : (N - lastSegmentElements) / (2 * SAMPLE_STRIDE);\n\n if (sortDir) {\n generateSampleRanksKernel<1U>\n <<>>(d_RanksA, d_RanksB, d_SrcKey, stride, N, threadCount);\n getLastCudaError(\"generateSampleRanksKernel<1U><<<>>> failed\\n\");\n }\n else {\n generateSampleRanksKernel<0U>\n <<>>(d_RanksA, d_RanksB, d_SrcKey, stride, N, threadCount);\n getLastCudaError(\"generateSampleRanksKernel<0U><<<>>> failed\\n\");\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Merge step 2: generate sample ranks and indices\n////////////////////////////////////////////////////////////////////////////////\n__global__ void mergeRanksAndIndicesKernel(uint *d_Limits, uint *d_Ranks, uint stride, uint N, uint threadCount)\n{\n uint pos = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (pos >= threadCount) {\n return;\n }\n\n const uint i = pos & ((stride / SAMPLE_STRIDE) - 1);\n const uint segmentBase = (pos - i) * (2 * SAMPLE_STRIDE);\n d_Ranks += (pos - i) * 2;\n d_Limits += (pos - i) * 2;\n\n const uint segmentElementsA = stride;\n const uint segmentElementsB = umin(stride, N - segmentBase - stride);\n const uint segmentSamplesA = getSampleCount(segmentElementsA);\n const uint segmentSamplesB = getSampleCount(segmentElementsB);\n\n if (i < segmentSamplesA) {\n uint dstPos = binarySearchExclusive<1U>(\n d_Ranks[i], d_Ranks + segmentSamplesA, segmentSamplesB, nextPowerOfTwo(segmentSamplesB))\n + i;\n d_Limits[dstPos] = d_Ranks[i];\n }\n\n if (i < segmentSamplesB) {\n uint dstPos = binarySearchInclusive<1U>(\n d_Ranks[segmentSamplesA + i], d_Ranks, segmentSamplesA, nextPowerOfTwo(segmentSamplesA))\n + i;\n d_Limits[dstPos] = d_Ranks[segmentSamplesA + i];\n }\n}\n\nstatic void mergeRanksAndIndices(uint *d_LimitsA, uint *d_LimitsB, uint *d_RanksA, uint *d_RanksB, uint stride, uint N)\n{\n uint lastSegmentElements = N % (2 * stride);\n uint threadCount = (lastSegmentElements > stride) ? (N + 2 * stride - lastSegmentElements) / (2 * SAMPLE_STRIDE)\n : (N - lastSegmentElements) / (2 * SAMPLE_STRIDE);\n\n mergeRanksAndIndicesKernel<<>>(d_LimitsA, d_RanksA, stride, N, threadCount);\n getLastCudaError(\"mergeRanksAndIndicesKernel(A)<<<>>> failed\\n\");\n\n mergeRanksAndIndicesKernel<<>>(d_LimitsB, d_RanksB, stride, N, threadCount);\n getLastCudaError(\"mergeRanksAndIndicesKernel(B)<<<>>> failed\\n\");\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Merge step 3: merge elementary intervals\n////////////////////////////////////////////////////////////////////////////////\ntemplate \ninline __device__ void merge(uint *dstKey,\n uint *dstVal,\n uint *srcAKey,\n uint *srcAVal,\n uint *srcBKey,\n uint *srcBVal,\n uint lenA,\n uint nPowTwoLenA,\n uint lenB,\n uint nPowTwoLenB,\n cg::thread_block cta)\n{\n uint keyA, valA, keyB, valB, dstPosA, dstPosB;\n\n if (threadIdx.x < lenA) {\n keyA = srcAKey[threadIdx.x];\n valA = srcAVal[threadIdx.x];\n dstPosA = binarySearchExclusive(keyA, srcBKey, lenB, nPowTwoLenB) + threadIdx.x;\n }\n\n if (threadIdx.x < lenB) {\n keyB = srcBKey[threadIdx.x];\n valB = srcBVal[threadIdx.x];\n dstPosB = binarySearchInclusive(keyB, srcAKey, lenA, nPowTwoLenA) + threadIdx.x;\n }\n\n cg::sync(cta);\n\n if (threadIdx.x < lenA) {\n dstKey[dstPosA] = keyA;\n dstVal[dstPosA] = valA;\n }\n\n if (threadIdx.x < lenB) {\n dstKey[dstPosB] = keyB;\n dstVal[dstPosB] = valB;\n }\n}\n\ntemplate \n__global__ void mergeElementaryIntervalsKernel(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint *d_LimitsA,\n uint *d_LimitsB,\n uint stride,\n uint N)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ uint s_key[2 * SAMPLE_STRIDE];\n __shared__ uint s_val[2 * SAMPLE_STRIDE];\n\n const uint intervalI = blockIdx.x & ((2 * stride) / SAMPLE_STRIDE - 1);\n const uint segmentBase = (blockIdx.x - intervalI) * SAMPLE_STRIDE;\n d_SrcKey += segmentBase;\n d_SrcVal += segmentBase;\n d_DstKey += segmentBase;\n d_DstVal += segmentBase;\n\n // Set up threadblock-wide parameters\n __shared__ uint startSrcA, startSrcB, lenSrcA, lenSrcB, startDstA, startDstB;\n\n if (threadIdx.x == 0) {\n uint segmentElementsA = stride;\n uint segmentElementsB = umin(stride, N - segmentBase - stride);\n uint segmentSamplesA = getSampleCount(segmentElementsA);\n uint segmentSamplesB = getSampleCount(segmentElementsB);\n uint segmentSamples = segmentSamplesA + segmentSamplesB;\n\n startSrcA = d_LimitsA[blockIdx.x];\n startSrcB = d_LimitsB[blockIdx.x];\n uint endSrcA = (intervalI + 1 < segmentSamples) ? d_LimitsA[blockIdx.x + 1] : segmentElementsA;\n uint endSrcB = (intervalI + 1 < segmentSamples) ? d_LimitsB[blockIdx.x + 1] : segmentElementsB;\n lenSrcA = endSrcA - startSrcA;\n lenSrcB = endSrcB - startSrcB;\n startDstA = startSrcA + startSrcB;\n startDstB = startDstA + lenSrcA;\n }\n\n // Load main input data\n cg::sync(cta);\n\n if (threadIdx.x < lenSrcA) {\n s_key[threadIdx.x + 0] = d_SrcKey[0 + startSrcA + threadIdx.x];\n s_val[threadIdx.x + 0] = d_SrcVal[0 + startSrcA + threadIdx.x];\n }\n\n if (threadIdx.x < lenSrcB) {\n s_key[threadIdx.x + SAMPLE_STRIDE] = d_SrcKey[stride + startSrcB + threadIdx.x];\n s_val[threadIdx.x + SAMPLE_STRIDE] = d_SrcVal[stride + startSrcB + threadIdx.x];\n }\n\n // Merge data in shared memory\n cg::sync(cta);\n merge(s_key,\n s_val,\n s_key + 0,\n s_val + 0,\n s_key + SAMPLE_STRIDE,\n s_val + SAMPLE_STRIDE,\n lenSrcA,\n SAMPLE_STRIDE,\n lenSrcB,\n SAMPLE_STRIDE,\n cta);\n\n // Store merged data\n cg::sync(cta);\n\n if (threadIdx.x < lenSrcA) {\n d_DstKey[startDstA + threadIdx.x] = s_key[threadIdx.x];\n d_DstVal[startDstA + threadIdx.x] = s_val[threadIdx.x];\n }\n\n if (threadIdx.x < lenSrcB) {\n d_DstKey[startDstB + threadIdx.x] = s_key[lenSrcA + threadIdx.x];\n d_DstVal[startDstB + threadIdx.x] = s_val[lenSrcA + threadIdx.x];\n }\n}\n\nstatic void mergeElementaryIntervals(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint *d_LimitsA,\n uint *d_LimitsB,\n uint stride,\n uint N,\n uint sortDir)\n{\n uint lastSegmentElements = N % (2 * stride);\n uint mergePairs = (lastSegmentElements > stride) ? getSampleCount(N) : (N - lastSegmentElements) / SAMPLE_STRIDE;\n\n if (sortDir) {\n mergeElementaryIntervalsKernel<1U>\n <<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, d_LimitsA, d_LimitsB, stride, N);\n getLastCudaError(\"mergeElementaryIntervalsKernel<1> failed\\n\");\n }\n else {\n mergeElementaryIntervalsKernel<0U>\n <<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, d_LimitsA, d_LimitsB, stride, N);\n getLastCudaError(\"mergeElementaryIntervalsKernel<0> failed\\n\");\n }\n}\n\nextern \"C\" void bitonicSortShared(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint batchSize,\n uint arrayLength,\n uint sortDir);\n\nextern \"C\" void bitonicMergeElementaryIntervals(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint *d_LimitsA,\n uint *d_LimitsB,\n uint stride,\n uint N,\n uint sortDir);\n\nstatic uint *d_RanksA, *d_RanksB, *d_LimitsA, *d_LimitsB;\nstatic const uint MAX_SAMPLE_COUNT = 32768;\n\nextern \"C\" void initMergeSort(void)\n{\n checkCudaErrors(hipMalloc((void **)&d_RanksA, MAX_SAMPLE_COUNT * sizeof(uint)));\n checkCudaErrors(hipMalloc((void **)&d_RanksB, MAX_SAMPLE_COUNT * sizeof(uint)));\n checkCudaErrors(hipMalloc((void **)&d_LimitsA, MAX_SAMPLE_COUNT * sizeof(uint)));\n checkCudaErrors(hipMalloc((void **)&d_LimitsB, MAX_SAMPLE_COUNT * sizeof(uint)));\n}\n\nextern \"C\" void closeMergeSort(void)\n{\n checkCudaErrors(hipFree(d_RanksA));\n checkCudaErrors(hipFree(d_RanksB));\n checkCudaErrors(hipFree(d_LimitsB));\n checkCudaErrors(hipFree(d_LimitsA));\n}\n\nextern \"C\" void mergeSort(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_BufKey,\n uint *d_BufVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint N,\n uint sortDir)\n{\n uint stageCount = 0;\n\n for (uint stride = SHARED_SIZE_LIMIT; stride < N; stride <<= 1, stageCount++)\n ;\n\n uint *ikey, *ival, *okey, *oval;\n\n if (stageCount & 1) {\n ikey = d_BufKey;\n ival = d_BufVal;\n okey = d_DstKey;\n oval = d_DstVal;\n }\n else {\n ikey = d_DstKey;\n ival = d_DstVal;\n okey = d_BufKey;\n oval = d_BufVal;\n }\n\n assert(N <= (SAMPLE_STRIDE * MAX_SAMPLE_COUNT));\n assert(N % SHARED_SIZE_LIMIT == 0);\n mergeSortShared(ikey, ival, d_SrcKey, d_SrcVal, N / SHARED_SIZE_LIMIT, SHARED_SIZE_LIMIT, sortDir);\n\n for (uint stride = SHARED_SIZE_LIMIT; stride < N; stride <<= 1) {\n uint lastSegmentElements = N % (2 * stride);\n\n // Find sample ranks and prepare for limiters merge\n generateSampleRanks(d_RanksA, d_RanksB, ikey, stride, N, sortDir);\n\n // Merge ranks and indices\n mergeRanksAndIndices(d_LimitsA, d_LimitsB, d_RanksA, d_RanksB, stride, N);\n\n // Merge elementary intervals\n mergeElementaryIntervals(okey, oval, ikey, ival, d_LimitsA, d_LimitsB, stride, N, sortDir);\n\n if (lastSegmentElements <= stride) {\n // Last merge segment consists of a single array which just needs to be\n // passed through\n checkCudaErrors(hipMemcpy(okey + (N - lastSegmentElements),\n ikey + (N - lastSegmentElements),\n lastSegmentElements * sizeof(uint),\n hipMemcpyDeviceToDevice));\n checkCudaErrors(hipMemcpy(oval + (N - lastSegmentElements),\n ival + (N - lastSegmentElements),\n lastSegmentElements * sizeof(uint),\n hipMemcpyDeviceToDevice));\n }\n\n uint *t;\n t = ikey;\n ikey = okey;\n okey = t;\n t = ival;\n ival = oval;\n oval = t;\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/mergeSort/mergeSort.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample illustrates the usage of CUDA streams for overlapping\n * kernel execution with device/host memcopies. The kernel is used to\n * initialize an array to a specific value, after which the array is\n * copied to the host (CPU) memory. To increase performance, multiple\n * kernel/memcopy pairs are launched asynchronously, each pair in its\n * own stream. Devices with Compute Capability 1.1 can overlap a kernel\n * and a memcopy as long as they are issued in different streams. Kernels\n * are serialized. Thus, if n pairs are launched, streamed approach\n * can reduce the memcopy cost to the (1/n)th of a single copy of the entire\n * data set.\n *\n * Additionally, this sample uses CUDA events to measure elapsed time for\n * CUDA calls. Events are a part of CUDA API and provide a system independent\n * way to measure execution times on CUDA devices with approximately 0.5\n * microsecond precision.\n *\n * Elapsed times are averaged over nreps repetitions (10 by default).\n *\n */\n\nconst char *sSDKsample = \"simpleStreams\";\n\nconst char *sEventSyncMethod[] = {\"cudaEventDefault\", \"cudaEventBlockingSync\", \"cudaEventDisableTiming\", NULL};\n\nconst char *sDeviceSyncMethod[] = {\"cudaDeviceScheduleAuto\",\n \"cudaDeviceScheduleSpin\",\n \"cudaDeviceScheduleYield\",\n \"INVALID\",\n \"cudaDeviceScheduleBlockingSync\",\n NULL};\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#ifndef WIN32\n#include // for mmap() / munmap()\n#endif\n\n// Macro to aligned up to the memory size in question\n#define MEMORY_ALIGNMENT 4096\n#define ALIGN_UP(x, size) (((size_t)x + (size - 1)) & (~(size - 1)))\n\n__global__ void init_array(int *g_data, int *factor, int num_iterations)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n for (int i = 0; i < num_iterations; i++) {\n g_data[idx] += *factor; // non-coalesced on purpose, to burn time\n }\n}\n\nbool correct_data(int *a, const int n, const int c)\n{\n for (int i = 0; i < n; i++) {\n if (a[i] != c) {\n printf(\"%d: %d %d\\n\", i, a[i], c);\n return false;\n }\n }\n\n return true;\n}\n\ninline void AllocateHostMemory(bool bPinGenericMemory, int **pp_a, int **ppAligned_a, int nbytes)\n{\n#if CUDART_VERSION >= 4000\n#if !defined(__arm__) && !defined(__aarch64__)\n if (bPinGenericMemory) {\n// allocate a generic page-aligned chunk of system memory\n#ifdef WIN32\n printf(\"> VirtualAlloc() allocating %4.2f Mbytes of (generic page-aligned \"\n \"system memory)\\n\",\n (float)nbytes / 1048576.0f);\n *pp_a = (int *)VirtualAlloc(NULL, (nbytes + MEMORY_ALIGNMENT), MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);\n#else\n printf(\"> mmap() allocating %4.2f Mbytes (generic page-aligned system \"\n \"memory)\\n\",\n (float)nbytes / 1048576.0f);\n *pp_a = (int *)mmap(NULL, (nbytes + MEMORY_ALIGNMENT), PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0);\n#endif\n\n *ppAligned_a = (int *)ALIGN_UP(*pp_a, MEMORY_ALIGNMENT);\n\n printf(\"> cudaHostRegister() registering %4.2f Mbytes of generic allocated \"\n \"system memory\\n\",\n (float)nbytes / 1048576.0f);\n // pin allocate memory\n checkCudaErrors(cudaHostRegister(*ppAligned_a, nbytes, cudaHostRegisterMapped));\n }\n else\n#endif\n#endif\n {\n printf(\"> cudaMallocHost() allocating %4.2f Mbytes of system memory\\n\", (float)nbytes / 1048576.0f);\n // allocate host memory (pinned is required for achieve asynchronicity)\n checkCudaErrors(cudaMallocHost((void **)pp_a, nbytes));\n *ppAligned_a = *pp_a;\n }\n}\n\ninline void FreeHostMemory(bool bPinGenericMemory, int **pp_a, int **ppAligned_a, int nbytes)\n{\n#if CUDART_VERSION >= 4000\n#if !defined(__arm__) && !defined(__aarch64__)\n // CUDA 4.0 support pinning of generic host memory\n if (bPinGenericMemory) {\n // unpin and delete host memory\n checkCudaErrors(cudaHostUnregister(*ppAligned_a));\n#ifdef WIN32\n VirtualFree(*pp_a, 0, MEM_RELEASE);\n#else\n munmap(*pp_a, nbytes);\n#endif\n }\n else\n#endif\n#endif\n {\n cudaFreeHost(*pp_a);\n }\n}\n\nstatic const char *sSyncMethod[] = {\"0 (Automatic Blocking)\",\n \"1 (Spin Blocking)\",\n \"2 (Yield Blocking)\",\n \"3 (Undefined Blocking Method)\",\n \"4 (Blocking Sync Event) = low CPU utilization\",\n NULL};\n\nvoid printHelp()\n{\n printf(\"Usage: %s [options below]\\n\", sSDKsample);\n printf(\"\\t--sync_method=n for CPU/GPU synchronization\\n\");\n printf(\"\\t n=%s\\n\", sSyncMethod[0]);\n printf(\"\\t n=%s\\n\", sSyncMethod[1]);\n printf(\"\\t n=%s\\n\", sSyncMethod[2]);\n printf(\"\\t n=%s\\n\", sSyncMethod[4]);\n printf(\"\\t--use_generic_memory (default) use generic page-aligned for system \"\n \"memory\\n\");\n printf(\"\\t--use_cuda_malloc_host (optional) use cudaMallocHost to allocate \"\n \"system memory\\n\");\n}\n\n#if defined(__APPLE__) || defined(MACOSX)\n#define DEFAULT_PINNED_GENERIC_MEMORY false\n#else\n#define DEFAULT_PINNED_GENERIC_MEMORY true\n#endif\n\nint main(int argc, char **argv)\n{\n int cuda_device = 0;\n int nstreams = 4; // number of streams for CUDA calls\n int nreps = 10; // number of times each experiment is repeated\n int n = 16 * 1024 * 1024; // number of ints in the data set\n int nbytes = n * sizeof(int); // number of data bytes\n dim3 threads, blocks; // kernel launch configuration\n float elapsed_time, time_memcpy, time_kernel; // timing variables\n float scale_factor = 1.0f;\n\n // allocate generic memory and pin it laster instead of using cudaHostAlloc()\n\n bool bPinGenericMemory = DEFAULT_PINNED_GENERIC_MEMORY; // we want this to be the default behavior\n int device_sync_method = cudaDeviceBlockingSync; // by default we use BlockingSync\n\n int niterations; // number of iterations for the loop inside the kernel\n\n printf(\"[ %s ]\\n\\n\", sSDKsample);\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\")) {\n printHelp();\n return EXIT_SUCCESS;\n }\n\n if ((device_sync_method = getCmdLineArgumentInt(argc, (const char **)argv, \"sync_method\")) >= 0) {\n if (device_sync_method == 0 || device_sync_method == 1 || device_sync_method == 2 || device_sync_method == 4) {\n printf(\"Device synchronization method set to = %s\\n\", sSyncMethod[device_sync_method]);\n printf(\"Setting reps to 100 to demonstrate steady state\\n\");\n nreps = 100;\n }\n else {\n printf(\"Invalid command line option sync_method=\\\"%d\\\"\\n\", device_sync_method);\n return EXIT_FAILURE;\n }\n }\n else {\n printHelp();\n return EXIT_SUCCESS;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"use_generic_memory\")) {\n#if defined(__APPLE__) || defined(MACOSX)\n bPinGenericMemory = false; // Generic Pinning of System Paged memory not\n // currently supported on Mac OSX\n#else\n bPinGenericMemory = true;\n#endif\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"use_cuda_malloc_host\")) {\n bPinGenericMemory = false;\n }\n\n printf(\"\\n> \");\n cuda_device = findCudaDevice(argc, (const char **)argv);\n\n // check the compute capability of the device\n int num_devices = 0;\n checkCudaErrors(cudaGetDeviceCount(&num_devices));\n\n if (0 == num_devices) {\n printf(\"your system does not have a CUDA capable device, waiving test...\\n\");\n return EXIT_WAIVED;\n }\n\n // check if the command-line chosen device ID is within range, exit if not\n if (cuda_device >= num_devices) {\n printf(\"cuda_device=%d is invalid, must choose device ID between 0 and %d\\n\", cuda_device, num_devices - 1);\n return EXIT_FAILURE;\n }\n\n checkCudaErrors(cudaSetDevice(cuda_device));\n\n // Checking for compute capabilities\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, cuda_device));\n\n niterations = 5;\n\n // Check if GPU can map host memory (Generic Method), if not then we override\n // bPinGenericMemory to be false\n if (bPinGenericMemory) {\n printf(\"Device: <%s> canMapHostMemory: %s\\n\", deviceProp.name, deviceProp.canMapHostMemory ? \"Yes\" : \"No\");\n\n if (deviceProp.canMapHostMemory == 0) {\n printf(\"Using cudaMallocHost, CUDA device does not support mapping of \"\n \"generic host memory\\n\");\n bPinGenericMemory = false;\n }\n }\n\n // Anything that is less than 32 Cores will have scaled down workload\n scale_factor =\n max((32.0f / (_ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * (float)deviceProp.multiProcessorCount)),\n 1.0f);\n n = (int)rint((float)n / scale_factor);\n\n printf(\"> CUDA Capable: SM %d.%d hardware\\n\", deviceProp.major, deviceProp.minor);\n printf(\"> %d Multiprocessor(s) x %d (Cores/Multiprocessor) = %d (Cores)\\n\",\n deviceProp.multiProcessorCount,\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor),\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * deviceProp.multiProcessorCount);\n\n printf(\"> scale_factor = %1.4f\\n\", 1.0f / scale_factor);\n printf(\"> array_size = %d\\n\\n\", n);\n\n // enable use of blocking sync, to reduce CPU usage\n printf(\"> Using CPU/GPU Device Synchronization method (%s)\\n\", sDeviceSyncMethod[device_sync_method]);\n checkCudaErrors(cudaSetDeviceFlags(device_sync_method | (bPinGenericMemory ? cudaDeviceMapHost : 0)));\n\n // allocate host memory\n int c = 5; // value to which the array will be initialized\n int *h_a = 0; // pointer to the array data in host memory\n int *hAligned_a = 0; // pointer to the array data in host memory (aligned to\n // MEMORY_ALIGNMENT)\n\n // Allocate Host memory (could be using cudaMallocHost or VirtualAlloc/mmap if\n // using the new CUDA 4.0 features\n AllocateHostMemory(bPinGenericMemory, &h_a, &hAligned_a, nbytes);\n\n // allocate device memory\n int *d_a = 0,\n *d_c = 0; // pointers to data and init value in the device memory\n checkCudaErrors(cudaMalloc((void **)&d_a, nbytes));\n checkCudaErrors(cudaMemset(d_a, 0x0, nbytes));\n checkCudaErrors(cudaMalloc((void **)&d_c, sizeof(int)));\n checkCudaErrors(cudaMemcpy(d_c, &c, sizeof(int), cudaMemcpyHostToDevice));\n\n printf(\"\\nStarting Test\\n\");\n\n // allocate and initialize an array of stream handles\n cudaStream_t *streams = (cudaStream_t *)malloc(nstreams * sizeof(cudaStream_t));\n\n for (int i = 0; i < nstreams; i++) {\n checkCudaErrors(cudaStreamCreate(&(streams[i])));\n }\n\n // create CUDA event handles\n // use blocking sync\n cudaEvent_t start_event, stop_event;\n int eventflags = ((device_sync_method == cudaDeviceBlockingSync) ? cudaEventBlockingSync : cudaEventDefault);\n\n checkCudaErrors(cudaEventCreateWithFlags(&start_event, eventflags));\n checkCudaErrors(cudaEventCreateWithFlags(&stop_event, eventflags));\n\n // time memcopy from device\n checkCudaErrors(cudaEventRecord(start_event, 0)); // record in stream-0, to\n // ensure that all previous\n // CUDA calls have\n // completed\n checkCudaErrors(cudaMemcpyAsync(hAligned_a, d_a, nbytes, cudaMemcpyDeviceToHost, streams[0]));\n checkCudaErrors(cudaEventRecord(stop_event, 0));\n checkCudaErrors(cudaEventSynchronize(stop_event)); // block until the event is actually recorded\n checkCudaErrors(cudaEventElapsedTime(&time_memcpy, start_event, stop_event));\n printf(\"memcopy:\\t%.2f\\n\", time_memcpy);\n\n // time kernel\n threads = dim3(512, 1);\n blocks = dim3(n / threads.x, 1);\n checkCudaErrors(cudaEventRecord(start_event, 0));\n init_array<<>>(d_a, d_c, niterations);\n checkCudaErrors(cudaEventRecord(stop_event, 0));\n checkCudaErrors(cudaEventSynchronize(stop_event));\n checkCudaErrors(cudaEventElapsedTime(&time_kernel, start_event, stop_event));\n printf(\"kernel:\\t\\t%.2f\\n\", time_kernel);\n\n //////////////////////////////////////////////////////////////////////\n // time non-streamed execution for reference\n threads = dim3(512, 1);\n blocks = dim3(n / threads.x, 1);\n checkCudaErrors(cudaEventRecord(start_event, 0));\n\n for (int k = 0; k < nreps; k++) {\n init_array<<>>(d_a, d_c, niterations);\n checkCudaErrors(cudaMemcpy(hAligned_a, d_a, nbytes, cudaMemcpyDeviceToHost));\n }\n\n checkCudaErrors(cudaEventRecord(stop_event, 0));\n checkCudaErrors(cudaEventSynchronize(stop_event));\n checkCudaErrors(cudaEventElapsedTime(&elapsed_time, start_event, stop_event));\n printf(\"non-streamed:\\t%.2f\\n\", elapsed_time / nreps);\n\n //////////////////////////////////////////////////////////////////////\n // time execution with nstreams streams\n threads = dim3(512, 1);\n blocks = dim3(n / (nstreams * threads.x), 1);\n memset(hAligned_a, 255,\n nbytes); // set host memory bits to all 1s, for testing correctness\n checkCudaErrors(cudaMemset(d_a, 0, nbytes)); // set device memory to all 0s, for testing correctness\n checkCudaErrors(cudaEventRecord(start_event, 0));\n\n for (int k = 0; k < nreps; k++) {\n // asynchronously launch nstreams kernels, each operating on its own portion\n // of data\n for (int i = 0; i < nstreams; i++) {\n init_array<<>>(d_a + i * n / nstreams, d_c, niterations);\n }\n\n // asynchronously launch nstreams memcopies. Note that memcopy in stream x\n // will only\n // commence executing when all previous CUDA calls in stream x have\n // completed\n for (int i = 0; i < nstreams; i++) {\n checkCudaErrors(cudaMemcpyAsync(hAligned_a + i * n / nstreams,\n d_a + i * n / nstreams,\n nbytes / nstreams,\n cudaMemcpyDeviceToHost,\n streams[i]));\n }\n }\n\n checkCudaErrors(cudaEventRecord(stop_event, 0));\n checkCudaErrors(cudaEventSynchronize(stop_event));\n checkCudaErrors(cudaEventElapsedTime(&elapsed_time, start_event, stop_event));\n printf(\"%d streams:\\t%.2f\\n\", nstreams, elapsed_time / nreps);\n\n // check whether the output is correct\n printf(\"-------------------------------\\n\");\n bool bResults = correct_data(hAligned_a, n, c * nreps * niterations);\n\n // release resources\n for (int i = 0; i < nstreams; i++) {\n checkCudaErrors(cudaStreamDestroy(streams[i]));\n }\n\n checkCudaErrors(cudaEventDestroy(start_event));\n checkCudaErrors(cudaEventDestroy(stop_event));\n\n // Free cudaMallocHost or Generic Host allocated memory (from CUDA 4.0)\n FreeHostMemory(bPinGenericMemory, &h_a, &hAligned_a, nbytes);\n\n checkCudaErrors(cudaFree(d_a));\n checkCudaErrors(cudaFree(d_c));\n\n return bResults ? EXIT_SUCCESS : EXIT_FAILURE;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample illustrates the usage of CUDA streams for overlapping\n * kernel execution with device/host memcopies. The kernel is used to\n * initialize an array to a specific value, after which the array is\n * copied to the host (CPU) memory. To increase performance, multiple\n * kernel/memcopy pairs are launched asynchronously, each pair in its\n * own stream. Devices with Compute Capability 1.1 can overlap a kernel\n * and a memcopy as long as they are issued in different streams. Kernels\n * are serialized. Thus, if n pairs are launched, streamed approach\n * can reduce the memcopy cost to the (1/n)th of a single copy of the entire\n * data set.\n *\n * Additionally, this sample uses CUDA events to measure elapsed time for\n * CUDA calls. Events are a part of CUDA API and provide a system independent\n * way to measure execution times on CUDA devices with approximately 0.5\n * microsecond precision.\n *\n * Elapsed times are averaged over nreps repetitions (10 by default).\n *\n */\n\nconst char *sSDKsample = \"simpleStreams\";\n\nconst char *sEventSyncMethod[] = {\"hipEventDefault\", \"hipEventBlockingSync\", \"hipEventDisableTiming\", NULL};\n\nconst char *sDeviceSyncMethod[] = {\"hipDeviceScheduleAuto\",\n \"hipDeviceScheduleSpin\",\n \"hipDeviceScheduleYield\",\n \"INVALID\",\n \"hipDeviceScheduleBlockingSync\",\n NULL};\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#ifndef WIN32\n#include // for mmap() / munmap()\n#endif\n\n// Macro to aligned up to the memory size in question\n#define MEMORY_ALIGNMENT 4096\n#define ALIGN_UP(x, size) (((size_t)x + (size - 1)) & (~(size - 1)))\n\n__global__ void init_array(int *g_data, int *factor, int num_iterations)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n for (int i = 0; i < num_iterations; i++) {\n g_data[idx] += *factor; // non-coalesced on purpose, to burn time\n }\n}\n\nbool correct_data(int *a, const int n, const int c)\n{\n for (int i = 0; i < n; i++) {\n if (a[i] != c) {\n printf(\"%d: %d %d\\n\", i, a[i], c);\n return false;\n }\n }\n\n return true;\n}\n\ninline void AllocateHostMemory(bool bPinGenericMemory, int **pp_a, int **ppAligned_a, int nbytes)\n{\n#if CUDART_VERSION >= 4000\n#if !defined(__arm__) && !defined(__aarch64__)\n if (bPinGenericMemory) {\n// allocate a generic page-aligned chunk of system memory\n#ifdef WIN32\n printf(\"> VirtualAlloc() allocating %4.2f Mbytes of (generic page-aligned \"\n \"system memory)\\n\",\n (float)nbytes / 1048576.0f);\n *pp_a = (int *)VirtualAlloc(NULL, (nbytes + MEMORY_ALIGNMENT), MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);\n#else\n printf(\"> mmap() allocating %4.2f Mbytes (generic page-aligned system \"\n \"memory)\\n\",\n (float)nbytes / 1048576.0f);\n *pp_a = (int *)mmap(NULL, (nbytes + MEMORY_ALIGNMENT), PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0);\n#endif\n\n *ppAligned_a = (int *)ALIGN_UP(*pp_a, MEMORY_ALIGNMENT);\n\n printf(\"> hipHostRegister() registering %4.2f Mbytes of generic allocated \"\n \"system memory\\n\",\n (float)nbytes / 1048576.0f);\n // pin allocate memory\n checkCudaErrors(hipHostRegister(*ppAligned_a, nbytes, hipHostRegisterMapped));\n }\n else\n#endif\n#endif\n {\n printf(\"> hipHostMalloc() allocating %4.2f Mbytes of system memory\\n\", (float)nbytes / 1048576.0f);\n // allocate host memory (pinned is required for achieve asynchronicity)\n checkCudaErrors(hipHostMalloc((void **)pp_a, nbytes));\n *ppAligned_a = *pp_a;\n }\n}\n\ninline void FreeHostMemory(bool bPinGenericMemory, int **pp_a, int **ppAligned_a, int nbytes)\n{\n#if CUDART_VERSION >= 4000\n#if !defined(__arm__) && !defined(__aarch64__)\n // CUDA 4.0 support pinning of generic host memory\n if (bPinGenericMemory) {\n // unpin and delete host memory\n checkCudaErrors(hipHostUnregister(*ppAligned_a));\n#ifdef WIN32\n VirtualFree(*pp_a, 0, MEM_RELEASE);\n#else\n munmap(*pp_a, nbytes);\n#endif\n }\n else\n#endif\n#endif\n {\n hipHostFree(*pp_a);\n }\n}\n\nstatic const char *sSyncMethod[] = {\"0 (Automatic Blocking)\",\n \"1 (Spin Blocking)\",\n \"2 (Yield Blocking)\",\n \"3 (Undefined Blocking Method)\",\n \"4 (Blocking Sync Event) = low CPU utilization\",\n NULL};\n\nvoid printHelp()\n{\n printf(\"Usage: %s [options below]\\n\", sSDKsample);\n printf(\"\\t--sync_method=n for CPU/GPU synchronization\\n\");\n printf(\"\\t n=%s\\n\", sSyncMethod[0]);\n printf(\"\\t n=%s\\n\", sSyncMethod[1]);\n printf(\"\\t n=%s\\n\", sSyncMethod[2]);\n printf(\"\\t n=%s\\n\", sSyncMethod[4]);\n printf(\"\\t--use_generic_memory (default) use generic page-aligned for system \"\n \"memory\\n\");\n printf(\"\\t--use_cuda_malloc_host (optional) use hipHostMalloc to allocate \"\n \"system memory\\n\");\n}\n\n#if defined(__APPLE__) || defined(MACOSX)\n#define DEFAULT_PINNED_GENERIC_MEMORY false\n#else\n#define DEFAULT_PINNED_GENERIC_MEMORY true\n#endif\n\nint main(int argc, char **argv)\n{\n int cuda_device = 0;\n int nstreams = 4; // number of streams for CUDA calls\n int nreps = 10; // number of times each experiment is repeated\n int n = 16 * 1024 * 1024; // number of ints in the data set\n int nbytes = n * sizeof(int); // number of data bytes\n dim3 threads, blocks; // kernel launch configuration\n float elapsed_time, time_memcpy, time_kernel; // timing variables\n float scale_factor = 1.0f;\n\n // allocate generic memory and pin it laster instead of using cudaHostAlloc()\n\n bool bPinGenericMemory = DEFAULT_PINNED_GENERIC_MEMORY; // we want this to be the default behavior\n int device_sync_method = hipDeviceScheduleBlockingSync; // by default we use BlockingSync\n\n int niterations; // number of iterations for the loop inside the kernel\n\n printf(\"[ %s ]\\n\\n\", sSDKsample);\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\")) {\n printHelp();\n return EXIT_SUCCESS;\n }\n\n if ((device_sync_method = getCmdLineArgumentInt(argc, (const char **)argv, \"sync_method\")) >= 0) {\n if (device_sync_method == 0 || device_sync_method == 1 || device_sync_method == 2 || device_sync_method == 4) {\n printf(\"Device synchronization method set to = %s\\n\", sSyncMethod[device_sync_method]);\n printf(\"Setting reps to 100 to demonstrate steady state\\n\");\n nreps = 100;\n }\n else {\n printf(\"Invalid command line option sync_method=\\\"%d\\\"\\n\", device_sync_method);\n return EXIT_FAILURE;\n }\n }\n else {\n printHelp();\n return EXIT_SUCCESS;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"use_generic_memory\")) {\n#if defined(__APPLE__) || defined(MACOSX)\n bPinGenericMemory = false; // Generic Pinning of System Paged memory not\n // currently supported on Mac OSX\n#else\n bPinGenericMemory = true;\n#endif\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"use_cuda_malloc_host\")) {\n bPinGenericMemory = false;\n }\n\n printf(\"\\n> \");\n cuda_device = findCudaDevice(argc, (const char **)argv);\n\n // check the compute capability of the device\n int num_devices = 0;\n checkCudaErrors(hipGetDeviceCount(&num_devices));\n\n if (0 == num_devices) {\n printf(\"your system does not have a CUDA capable device, waiving test...\\n\");\n return EXIT_WAIVED;\n }\n\n // check if the command-line chosen device ID is within range, exit if not\n if (cuda_device >= num_devices) {\n printf(\"cuda_device=%d is invalid, must choose device ID between 0 and %d\\n\", cuda_device, num_devices - 1);\n return EXIT_FAILURE;\n }\n\n checkCudaErrors(hipSetDevice(cuda_device));\n\n // Checking for compute capabilities\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, cuda_device));\n\n niterations = 5;\n\n // Check if GPU can map host memory (Generic Method), if not then we override\n // bPinGenericMemory to be false\n if (bPinGenericMemory) {\n printf(\"Device: <%s> canMapHostMemory: %s\\n\", deviceProp.name, deviceProp.canMapHostMemory ? \"Yes\" : \"No\");\n\n if (deviceProp.canMapHostMemory == 0) {\n printf(\"Using hipHostMalloc, CUDA device does not support mapping of \"\n \"generic host memory\\n\");\n bPinGenericMemory = false;\n }\n }\n\n // Anything that is less than 32 Cores will have scaled down workload\n scale_factor =\n max((32.0f / (_ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * (float)deviceProp.multiProcessorCount)),\n 1.0f);\n n = (int)rint((float)n / scale_factor);\n\n printf(\"> CUDA Capable: SM %d.%d hardware\\n\", deviceProp.major, deviceProp.minor);\n printf(\"> %d Multiprocessor(s) x %d (Cores/Multiprocessor) = %d (Cores)\\n\",\n deviceProp.multiProcessorCount,\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor),\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * deviceProp.multiProcessorCount);\n\n printf(\"> scale_factor = %1.4f\\n\", 1.0f / scale_factor);\n printf(\"> array_size = %d\\n\\n\", n);\n\n // enable use of blocking sync, to reduce CPU usage\n printf(\"> Using CPU/GPU Device Synchronization method (%s)\\n\", sDeviceSyncMethod[device_sync_method]);\n checkCudaErrors(hipSetDeviceFlags(device_sync_method | (bPinGenericMemory ? hipDeviceMapHost : 0)));\n\n // allocate host memory\n int c = 5; // value to which the array will be initialized\n int *h_a = 0; // pointer to the array data in host memory\n int *hAligned_a = 0; // pointer to the array data in host memory (aligned to\n // MEMORY_ALIGNMENT)\n\n // Allocate Host memory (could be using cudaMallocHost or VirtualAlloc/mmap if\n // using the new CUDA 4.0 features\n AllocateHostMemory(bPinGenericMemory, &h_a, &hAligned_a, nbytes);\n\n // allocate device memory\n int *d_a = 0,\n *d_c = 0; // pointers to data and init value in the device memory\n checkCudaErrors(hipMalloc((void **)&d_a, nbytes));\n checkCudaErrors(hipMemset(d_a, 0x0, nbytes));\n checkCudaErrors(hipMalloc((void **)&d_c, sizeof(int)));\n checkCudaErrors(hipMemcpy(d_c, &c, sizeof(int), hipMemcpyHostToDevice));\n\n printf(\"\\nStarting Test\\n\");\n\n // allocate and initialize an array of stream handles\n hipStream_t *streams = (hipStream_t *)malloc(nstreams * sizeof(hipStream_t));\n\n for (int i = 0; i < nstreams; i++) {\n checkCudaErrors(hipStreamCreate(&(streams[i])));\n }\n\n // create CUDA event handles\n // use blocking sync\n hipEvent_t start_event, stop_event;\n int eventflags = ((device_sync_method == hipDeviceScheduleBlockingSync) ? hipEventBlockingSync : hipEventDefault);\n\n checkCudaErrors(hipEventCreateWithFlags(&start_event, eventflags));\n checkCudaErrors(hipEventCreateWithFlags(&stop_event, eventflags));\n\n // time memcopy from device\n checkCudaErrors(hipEventRecord(start_event, 0)); // record in stream-0, to\n // ensure that all previous\n // CUDA calls have\n // completed\n checkCudaErrors(hipMemcpyAsync(hAligned_a, d_a, nbytes, hipMemcpyDeviceToHost, streams[0]));\n checkCudaErrors(hipEventRecord(stop_event, 0));\n checkCudaErrors(hipEventSynchronize(stop_event)); // block until the event is actually recorded\n checkCudaErrors(hipEventElapsedTime(&time_memcpy, start_event, stop_event));\n printf(\"memcopy:\\t%.2f\\n\", time_memcpy);\n\n // time kernel\n threads = dim3(512, 1);\n blocks = dim3(n / threads.x, 1);\n checkCudaErrors(hipEventRecord(start_event, 0));\n init_array<<>>(d_a, d_c, niterations);\n checkCudaErrors(hipEventRecord(stop_event, 0));\n checkCudaErrors(hipEventSynchronize(stop_event));\n checkCudaErrors(hipEventElapsedTime(&time_kernel, start_event, stop_event));\n printf(\"kernel:\\t\\t%.2f\\n\", time_kernel);\n\n //////////////////////////////////////////////////////////////////////\n // time non-streamed execution for reference\n threads = dim3(512, 1);\n blocks = dim3(n / threads.x, 1);\n checkCudaErrors(hipEventRecord(start_event, 0));\n\n for (int k = 0; k < nreps; k++) {\n init_array<<>>(d_a, d_c, niterations);\n checkCudaErrors(hipMemcpy(hAligned_a, d_a, nbytes, hipMemcpyDeviceToHost));\n }\n\n checkCudaErrors(hipEventRecord(stop_event, 0));\n checkCudaErrors(hipEventSynchronize(stop_event));\n checkCudaErrors(hipEventElapsedTime(&elapsed_time, start_event, stop_event));\n printf(\"non-streamed:\\t%.2f\\n\", elapsed_time / nreps);\n\n //////////////////////////////////////////////////////////////////////\n // time execution with nstreams streams\n threads = dim3(512, 1);\n blocks = dim3(n / (nstreams * threads.x), 1);\n memset(hAligned_a, 255,\n nbytes); // set host memory bits to all 1s, for testing correctness\n checkCudaErrors(hipMemset(d_a, 0, nbytes)); // set device memory to all 0s, for testing correctness\n checkCudaErrors(hipEventRecord(start_event, 0));\n\n for (int k = 0; k < nreps; k++) {\n // asynchronously launch nstreams kernels, each operating on its own portion\n // of data\n for (int i = 0; i < nstreams; i++) {\n init_array<<>>(d_a + i * n / nstreams, d_c, niterations);\n }\n\n // asynchronously launch nstreams memcopies. Note that memcopy in stream x\n // will only\n // commence executing when all previous CUDA calls in stream x have\n // completed\n for (int i = 0; i < nstreams; i++) {\n checkCudaErrors(hipMemcpyAsync(hAligned_a + i * n / nstreams,\n d_a + i * n / nstreams,\n nbytes / nstreams,\n hipMemcpyDeviceToHost,\n streams[i]));\n }\n }\n\n checkCudaErrors(hipEventRecord(stop_event, 0));\n checkCudaErrors(hipEventSynchronize(stop_event));\n checkCudaErrors(hipEventElapsedTime(&elapsed_time, start_event, stop_event));\n printf(\"%d streams:\\t%.2f\\n\", nstreams, elapsed_time / nreps);\n\n // check whether the output is correct\n printf(\"-------------------------------\\n\");\n bool bResults = correct_data(hAligned_a, n, c * nreps * niterations);\n\n // release resources\n for (int i = 0; i < nstreams; i++) {\n checkCudaErrors(hipStreamDestroy(streams[i]));\n }\n\n checkCudaErrors(hipEventDestroy(start_event));\n checkCudaErrors(hipEventDestroy(stop_event));\n\n // Free cudaMallocHost or Generic Host allocated memory (from CUDA 4.0)\n FreeHostMemory(bPinGenericMemory, &h_a, &hAligned_a, nbytes);\n\n checkCudaErrors(hipFree(d_a));\n checkCudaErrors(hipFree(d_c));\n\n return bResults ? EXIT_SUCCESS : EXIT_FAILURE;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleStreams/simpleStreams.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * Quadro and Tesla GPUs with compute capability >= 2.0 can overlap two\n * memcopies with kernel execution. This sample illustrates the usage of CUDA\n * streams to achieve overlapping of kernel execution with copying data to and\n * from the device.\n *\n * Additionally, this sample uses CUDA events to measure elapsed time for\n * CUDA calls. Events are a part of CUDA API and provide a system independent\n * way to measure execution times on CUDA devices with approximately 0.5\n * microsecond precision.\n *\n * Elapsed times are averaged over nreps repetitions (10 by default).\n *\n */\n\nconst char *sSDKname = \"simpleMultiCopy\";\n\n// includes, system\n#include \n\n// include CUDA\n#include \n\n// includes, project\n#include \n#include // helper for shared that are common to CUDA Samples\n\n// includes, kernels\n// Declare the CUDA kernels here and main() code that is needed to launch\n// Compute workload on the system\n__global__ void incKernel(int *g_out, int *g_in, int N, int inner_reps)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx < N) {\n for (int i = 0; i < inner_reps; ++i) {\n g_out[idx] = g_in[idx] + 1;\n }\n }\n}\n\n#define STREAM_COUNT 4\n\n// Uncomment to simulate data source/sink IO times\n// #define SIMULATE_IO\n\nint *h_data_source;\nint *h_data_sink;\n\nint *h_data_in[STREAM_COUNT];\nint *d_data_in[STREAM_COUNT];\n\nint *h_data_out[STREAM_COUNT];\nint *d_data_out[STREAM_COUNT];\n\ncudaEvent_t cycleDone[STREAM_COUNT];\ncudaStream_t stream[STREAM_COUNT];\n\ncudaEvent_t start, stop;\n\nint N = 1 << 22;\nint nreps = 10; // number of times each experiment is repeated\nint inner_reps = 5;\n\nint memsize;\n\ndim3 block(512);\ndim3 grid;\n\nint thread_blocks;\n\nfloat processWithStreams(int streams_used);\nvoid init();\nbool test();\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char *argv[])\n{\n int cuda_device = 0;\n float scale_factor;\n cudaDeviceProp deviceProp;\n\n printf(\"[%s] - Starting...\\n\", sSDKname);\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"device\")) {\n cuda_device = getCmdLineArgumentInt(argc, (const char **)argv, \"device=\");\n\n if (cuda_device < 0) {\n printf(\"Invalid command line parameters\\n\");\n exit(EXIT_FAILURE);\n }\n else {\n printf(\"cuda_device = %d\\n\", cuda_device);\n cuda_device = gpuDeviceInit(cuda_device);\n\n if (cuda_device < 0) {\n printf(\"No CUDA Capable devices found, exiting...\\n\");\n exit(EXIT_SUCCESS);\n }\n }\n }\n else {\n // Otherwise pick the device with the highest Gflops/s\n cuda_device = gpuGetMaxGflopsDeviceId();\n checkCudaErrors(cudaSetDevice(cuda_device));\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, cuda_device));\n printf(\"> Using CUDA device [%d]: %s\\n\", cuda_device, deviceProp.name);\n }\n\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, cuda_device));\n printf(\"[%s] has %d MP(s) x %d (Cores/MP) = %d (Cores)\\n\",\n deviceProp.name,\n deviceProp.multiProcessorCount,\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor),\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * deviceProp.multiProcessorCount);\n\n // Anything that is less than 32 Cores will have scaled down workload\n scale_factor =\n max((32.0f / (_ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * (float)deviceProp.multiProcessorCount)),\n 1.0f);\n N = (int)((float)N / scale_factor);\n\n printf(\"> Device name: %s\\n\", deviceProp.name);\n printf(\"> CUDA Capability %d.%d hardware with %d multi-processors\\n\",\n deviceProp.major,\n deviceProp.minor,\n deviceProp.multiProcessorCount);\n printf(\"> scale_factor = %.2f\\n\", 1.0f / scale_factor);\n printf(\"> array_size = %d\\n\\n\", N);\n\n memsize = N * sizeof(int);\n\n thread_blocks = N / block.x;\n\n grid.x = thread_blocks % 65535;\n grid.y = (thread_blocks / 65535 + 1);\n\n // Allocate resources\n\n h_data_source = (int *)malloc(memsize);\n h_data_sink = (int *)malloc(memsize);\n\n for (int i = 0; i < STREAM_COUNT; ++i) {\n checkCudaErrors(cudaHostAlloc(&h_data_in[i], memsize, cudaHostAllocDefault));\n checkCudaErrors(cudaMalloc(&d_data_in[i], memsize));\n checkCudaErrors(cudaMemset(d_data_in[i], 0, memsize));\n\n checkCudaErrors(cudaHostAlloc(&h_data_out[i], memsize, cudaHostAllocDefault));\n checkCudaErrors(cudaMalloc(&d_data_out[i], memsize));\n\n checkCudaErrors(cudaStreamCreate(&stream[i]));\n checkCudaErrors(cudaEventCreate(&cycleDone[i]));\n\n cudaEventRecord(cycleDone[i], stream[i]);\n }\n\n cudaEventCreate(&start);\n cudaEventCreate(&stop);\n\n init();\n\n // Kernel warmup\n incKernel<<>>(d_data_out[0], d_data_in[0], N, inner_reps);\n\n // Time copies and kernel\n cudaEventRecord(start, 0);\n checkCudaErrors(cudaMemcpyAsync(d_data_in[0], h_data_in[0], memsize, cudaMemcpyHostToDevice, 0));\n cudaEventRecord(stop, 0);\n cudaEventSynchronize(stop);\n\n float memcpy_h2d_time;\n cudaEventElapsedTime(&memcpy_h2d_time, start, stop);\n\n cudaEventRecord(start, 0);\n checkCudaErrors(cudaMemcpyAsync(h_data_out[0], d_data_out[0], memsize, cudaMemcpyDeviceToHost, 0));\n cudaEventRecord(stop, 0);\n cudaEventSynchronize(stop);\n\n float memcpy_d2h_time;\n cudaEventElapsedTime(&memcpy_d2h_time, start, stop);\n\n cudaEventRecord(start, 0);\n incKernel<<>>(d_data_out[0], d_data_in[0], N, inner_reps);\n cudaEventRecord(stop, 0);\n cudaEventSynchronize(stop);\n\n float kernel_time;\n cudaEventElapsedTime(&kernel_time, start, stop);\n\n printf(\"\\n\");\n printf(\"Relevant properties of this CUDA device\\n\");\n int canOverlap;\n checkCudaErrors(cudaDeviceGetAttribute(&canOverlap, cudaDevAttrGpuOverlap, cuda_device));\n printf(\"(%s) Can overlap one CPU<>GPU data transfer with GPU kernel execution \"\n \"(device property \\\"cudaDevAttrGpuOverlap\\\")\\n\",\n canOverlap ? \"X\" : \" \");\n // printf(\"(%s) Can execute several GPU kernels simultaneously (compute\n // capability >= 2.0)\\n\", deviceProp.major >= 2 ? \"X\": \" \");\n printf(\"(%s) Can overlap two CPU<>GPU data transfers with GPU kernel execution\\n\"\n \" (Compute Capability >= 2.0 AND (Tesla product OR Quadro \"\n \"4000/5000/6000/K5000)\\n\",\n (deviceProp.major >= 2 && deviceProp.asyncEngineCount > 1) ? \"X\" : \" \");\n\n printf(\"\\n\");\n printf(\"Measured timings (throughput):\\n\");\n printf(\" Memcpy host to device\\t: %f ms (%f GB/s)\\n\", memcpy_h2d_time, (memsize * 1e-6) / memcpy_h2d_time);\n printf(\" Memcpy device to host\\t: %f ms (%f GB/s)\\n\", memcpy_d2h_time, (memsize * 1e-6) / memcpy_d2h_time);\n printf(\" Kernel\\t\\t\\t: %f ms (%f GB/s)\\n\", kernel_time, (inner_reps * memsize * 2e-6) / kernel_time);\n\n printf(\"\\n\");\n printf(\"Theoretical limits for speedup gained from overlapped data \"\n \"transfers:\\n\");\n printf(\"No overlap at all (transfer-kernel-transfer): %f ms \\n\", memcpy_h2d_time + memcpy_d2h_time + kernel_time);\n printf(\"Compute can overlap with one transfer: %f ms\\n\", max((memcpy_h2d_time + memcpy_d2h_time), kernel_time));\n printf(\"Compute can overlap with both data transfers: %f ms\\n\",\n max(max(memcpy_h2d_time, memcpy_d2h_time), kernel_time));\n\n // Process pipelined work\n float serial_time = processWithStreams(1);\n float overlap_time = processWithStreams(STREAM_COUNT);\n\n printf(\"\\nAverage measured timings over %d repetitions:\\n\", nreps);\n printf(\" Avg. time when execution fully serialized\\t: %f ms\\n\", serial_time / nreps);\n printf(\" Avg. time when overlapped using %d streams\\t: %f ms\\n\", STREAM_COUNT, overlap_time / nreps);\n printf(\" Avg. speedup gained (serialized - overlapped)\\t: %f ms\\n\", (serial_time - overlap_time) / nreps);\n\n printf(\"\\nMeasured throughput:\\n\");\n printf(\" Fully serialized execution\\t\\t: %f GB/s\\n\", (nreps * (memsize * 2e-6)) / serial_time);\n printf(\" Overlapped using %d streams\\t\\t: %f GB/s\\n\", STREAM_COUNT, (nreps * (memsize * 2e-6)) / overlap_time);\n\n // Verify the results, we will use the results for final output\n bool bResults = test();\n\n // Free resources\n\n free(h_data_source);\n free(h_data_sink);\n\n for (int i = 0; i < STREAM_COUNT; ++i) {\n cudaFreeHost(h_data_in[i]);\n cudaFree(d_data_in[i]);\n\n cudaFreeHost(h_data_out[i]);\n cudaFree(d_data_out[i]);\n\n cudaStreamDestroy(stream[i]);\n cudaEventDestroy(cycleDone[i]);\n }\n\n cudaEventDestroy(start);\n cudaEventDestroy(stop);\n\n // Test result\n exit(bResults ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\nfloat processWithStreams(int streams_used)\n{\n int current_stream = 0;\n\n float time;\n\n // Do processing in a loop\n //\n // Note: All memory commands are processed in the order they are issued,\n // independent of the stream they are enqueued in. Hence the pattern by\n // which the copy and kernel commands are enqueued in the stream\n // has an influence on the achieved overlap.\n\n cudaEventRecord(start, 0);\n\n for (int i = 0; i < nreps; ++i) {\n int next_stream = (current_stream + 1) % streams_used;\n\n#ifdef SIMULATE_IO\n // Store the result\n memcpy(h_data_sink, h_data_out[current_stream], memsize);\n\n // Read new input\n memcpy(h_data_in[next_stream], h_data_source, memsize);\n#endif\n\n // Ensure that processing and copying of the last cycle has finished\n cudaEventSynchronize(cycleDone[next_stream]);\n\n // Process current frame\n incKernel<<>>(\n d_data_out[current_stream], d_data_in[current_stream], N, inner_reps);\n\n // Upload next frame\n checkCudaErrors(cudaMemcpyAsync(\n d_data_in[next_stream], h_data_in[next_stream], memsize, cudaMemcpyHostToDevice, stream[next_stream]));\n\n // Download current frame\n checkCudaErrors(cudaMemcpyAsync(h_data_out[current_stream],\n d_data_out[current_stream],\n memsize,\n cudaMemcpyDeviceToHost,\n stream[current_stream]));\n\n checkCudaErrors(cudaEventRecord(cycleDone[current_stream], stream[current_stream]));\n\n current_stream = next_stream;\n }\n\n cudaEventRecord(stop, 0);\n\n cudaDeviceSynchronize();\n\n cudaEventElapsedTime(&time, start, stop);\n\n return time;\n}\n\nvoid init()\n{\n for (int i = 0; i < N; ++i) {\n h_data_source[i] = 0;\n }\n\n for (int i = 0; i < STREAM_COUNT; ++i) {\n memcpy(h_data_in[i], h_data_source, memsize);\n }\n}\n\nbool test()\n{\n bool passed = true;\n\n for (int j = 0; j < STREAM_COUNT; ++j) {\n for (int i = 0; i < N; ++i) {\n passed &= (h_data_out[j][i] == 1);\n }\n }\n\n return passed;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * Quadro and Tesla GPUs with compute capability >= 2.0 can overlap two\n * memcopies with kernel execution. This sample illustrates the usage of CUDA\n * streams to achieve overlapping of kernel execution with copying data to and\n * from the device.\n *\n * Additionally, this sample uses CUDA events to measure elapsed time for\n * CUDA calls. Events are a part of CUDA API and provide a system independent\n * way to measure execution times on CUDA devices with approximately 0.5\n * microsecond precision.\n *\n * Elapsed times are averaged over nreps repetitions (10 by default).\n *\n */\n\nconst char *sSDKname = \"simpleMultiCopy\";\n\n// includes, system\n#include \n\n// include CUDA\n#include \n\n// includes, project\n#include \n#include // helper for shared that are common to CUDA Samples\n\n// includes, kernels\n// Declare the CUDA kernels here and main() code that is needed to launch\n// Compute workload on the system\n__global__ void incKernel(int *g_out, int *g_in, int N, int inner_reps)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx < N) {\n for (int i = 0; i < inner_reps; ++i) {\n g_out[idx] = g_in[idx] + 1;\n }\n }\n}\n\n#define STREAM_COUNT 4\n\n// Uncomment to simulate data source/sink IO times\n// #define SIMULATE_IO\n\nint *h_data_source;\nint *h_data_sink;\n\nint *h_data_in[STREAM_COUNT];\nint *d_data_in[STREAM_COUNT];\n\nint *h_data_out[STREAM_COUNT];\nint *d_data_out[STREAM_COUNT];\n\nhipEvent_t cycleDone[STREAM_COUNT];\nhipStream_t stream[STREAM_COUNT];\n\nhipEvent_t start, stop;\n\nint N = 1 << 22;\nint nreps = 10; // number of times each experiment is repeated\nint inner_reps = 5;\n\nint memsize;\n\ndim3 block(512);\ndim3 grid;\n\nint thread_blocks;\n\nfloat processWithStreams(int streams_used);\nvoid init();\nbool test();\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char *argv[])\n{\n int cuda_device = 0;\n float scale_factor;\n hipDeviceProp_t deviceProp;\n\n printf(\"[%s] - Starting...\\n\", sSDKname);\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"device\")) {\n cuda_device = getCmdLineArgumentInt(argc, (const char **)argv, \"device=\");\n\n if (cuda_device < 0) {\n printf(\"Invalid command line parameters\\n\");\n exit(EXIT_FAILURE);\n }\n else {\n printf(\"cuda_device = %d\\n\", cuda_device);\n cuda_device = gpuDeviceInit(cuda_device);\n\n if (cuda_device < 0) {\n printf(\"No CUDA Capable devices found, exiting...\\n\");\n exit(EXIT_SUCCESS);\n }\n }\n }\n else {\n // Otherwise pick the device with the highest Gflops/s\n cuda_device = gpuGetMaxGflopsDeviceId();\n checkCudaErrors(hipSetDevice(cuda_device));\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, cuda_device));\n printf(\"> Using CUDA device [%d]: %s\\n\", cuda_device, deviceProp.name);\n }\n\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, cuda_device));\n printf(\"[%s] has %d MP(s) x %d (Cores/MP) = %d (Cores)\\n\",\n deviceProp.name,\n deviceProp.multiProcessorCount,\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor),\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * deviceProp.multiProcessorCount);\n\n // Anything that is less than 32 Cores will have scaled down workload\n scale_factor =\n max((32.0f / (_ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * (float)deviceProp.multiProcessorCount)),\n 1.0f);\n N = (int)((float)N / scale_factor);\n\n printf(\"> Device name: %s\\n\", deviceProp.name);\n printf(\"> CUDA Capability %d.%d hardware with %d multi-processors\\n\",\n deviceProp.major,\n deviceProp.minor,\n deviceProp.multiProcessorCount);\n printf(\"> scale_factor = %.2f\\n\", 1.0f / scale_factor);\n printf(\"> array_size = %d\\n\\n\", N);\n\n memsize = N * sizeof(int);\n\n thread_blocks = N / block.x;\n\n grid.x = thread_blocks % 65535;\n grid.y = (thread_blocks / 65535 + 1);\n\n // Allocate resources\n\n h_data_source = (int *)malloc(memsize);\n h_data_sink = (int *)malloc(memsize);\n\n for (int i = 0; i < STREAM_COUNT; ++i) {\n checkCudaErrors(hipHostAlloc(&h_data_in[i], memsize, hipHostMallocDefault));\n checkCudaErrors(hipMalloc(&d_data_in[i], memsize));\n checkCudaErrors(hipMemset(d_data_in[i], 0, memsize));\n\n checkCudaErrors(hipHostAlloc(&h_data_out[i], memsize, hipHostMallocDefault));\n checkCudaErrors(hipMalloc(&d_data_out[i], memsize));\n\n checkCudaErrors(hipStreamCreate(&stream[i]));\n checkCudaErrors(hipEventCreate(&cycleDone[i]));\n\n hipEventRecord(cycleDone[i], stream[i]);\n }\n\n hipEventCreate(&start);\n hipEventCreate(&stop);\n\n init();\n\n // Kernel warmup\n incKernel<<>>(d_data_out[0], d_data_in[0], N, inner_reps);\n\n // Time copies and kernel\n hipEventRecord(start, 0);\n checkCudaErrors(hipMemcpyAsync(d_data_in[0], h_data_in[0], memsize, hipMemcpyHostToDevice, 0));\n hipEventRecord(stop, 0);\n hipEventSynchronize(stop);\n\n float memcpy_h2d_time;\n hipEventElapsedTime(&memcpy_h2d_time, start, stop);\n\n hipEventRecord(start, 0);\n checkCudaErrors(hipMemcpyAsync(h_data_out[0], d_data_out[0], memsize, hipMemcpyDeviceToHost, 0));\n hipEventRecord(stop, 0);\n hipEventSynchronize(stop);\n\n float memcpy_d2h_time;\n hipEventElapsedTime(&memcpy_d2h_time, start, stop);\n\n hipEventRecord(start, 0);\n incKernel<<>>(d_data_out[0], d_data_in[0], N, inner_reps);\n hipEventRecord(stop, 0);\n hipEventSynchronize(stop);\n\n float kernel_time;\n hipEventElapsedTime(&kernel_time, start, stop);\n\n printf(\"\\n\");\n printf(\"Relevant properties of this CUDA device\\n\");\n int canOverlap;\n checkCudaErrors(hipDeviceGetAttribute(&canOverlap, hipDeviceAttributeAsyncEngineCount, cuda_device));\n printf(\"(%s) Can overlap one CPU<>GPU data transfer with GPU kernel execution \"\n \"(device property \\\"hipDeviceAttributeAsyncEngineCount\\\")\\n\",\n canOverlap ? \"X\" : \" \");\n // printf(\"(%s) Can execute several GPU kernels simultaneously (compute\n // capability >= 2.0)\\n\", deviceProp.major >= 2 ? \"X\": \" \");\n printf(\"(%s) Can overlap two CPU<>GPU data transfers with GPU kernel execution\\n\"\n \" (Compute Capability >= 2.0 AND (Tesla product OR Quadro \"\n \"4000/5000/6000/K5000)\\n\",\n (deviceProp.major >= 2 && deviceProp.asyncEngineCount > 1) ? \"X\" : \" \");\n\n printf(\"\\n\");\n printf(\"Measured timings (throughput):\\n\");\n printf(\" Memcpy host to device\\t: %f ms (%f GB/s)\\n\", memcpy_h2d_time, (memsize * 1e-6) / memcpy_h2d_time);\n printf(\" Memcpy device to host\\t: %f ms (%f GB/s)\\n\", memcpy_d2h_time, (memsize * 1e-6) / memcpy_d2h_time);\n printf(\" Kernel\\t\\t\\t: %f ms (%f GB/s)\\n\", kernel_time, (inner_reps * memsize * 2e-6) / kernel_time);\n\n printf(\"\\n\");\n printf(\"Theoretical limits for speedup gained from overlapped data \"\n \"transfers:\\n\");\n printf(\"No overlap at all (transfer-kernel-transfer): %f ms \\n\", memcpy_h2d_time + memcpy_d2h_time + kernel_time);\n printf(\"Compute can overlap with one transfer: %f ms\\n\", max((memcpy_h2d_time + memcpy_d2h_time), kernel_time));\n printf(\"Compute can overlap with both data transfers: %f ms\\n\",\n max(max(memcpy_h2d_time, memcpy_d2h_time), kernel_time));\n\n // Process pipelined work\n float serial_time = processWithStreams(1);\n float overlap_time = processWithStreams(STREAM_COUNT);\n\n printf(\"\\nAverage measured timings over %d repetitions:\\n\", nreps);\n printf(\" Avg. time when execution fully serialized\\t: %f ms\\n\", serial_time / nreps);\n printf(\" Avg. time when overlapped using %d streams\\t: %f ms\\n\", STREAM_COUNT, overlap_time / nreps);\n printf(\" Avg. speedup gained (serialized - overlapped)\\t: %f ms\\n\", (serial_time - overlap_time) / nreps);\n\n printf(\"\\nMeasured throughput:\\n\");\n printf(\" Fully serialized execution\\t\\t: %f GB/s\\n\", (nreps * (memsize * 2e-6)) / serial_time);\n printf(\" Overlapped using %d streams\\t\\t: %f GB/s\\n\", STREAM_COUNT, (nreps * (memsize * 2e-6)) / overlap_time);\n\n // Verify the results, we will use the results for final output\n bool bResults = test();\n\n // Free resources\n\n free(h_data_source);\n free(h_data_sink);\n\n for (int i = 0; i < STREAM_COUNT; ++i) {\n hipHostFree(h_data_in[i]);\n hipFree(d_data_in[i]);\n\n hipHostFree(h_data_out[i]);\n hipFree(d_data_out[i]);\n\n hipStreamDestroy(stream[i]);\n hipEventDestroy(cycleDone[i]);\n }\n\n hipEventDestroy(start);\n hipEventDestroy(stop);\n\n // Test result\n exit(bResults ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\nfloat processWithStreams(int streams_used)\n{\n int current_stream = 0;\n\n float time;\n\n // Do processing in a loop\n //\n // Note: All memory commands are processed in the order they are issued,\n // independent of the stream they are enqueued in. Hence the pattern by\n // which the copy and kernel commands are enqueued in the stream\n // has an influence on the achieved overlap.\n\n hipEventRecord(start, 0);\n\n for (int i = 0; i < nreps; ++i) {\n int next_stream = (current_stream + 1) % streams_used;\n\n#ifdef SIMULATE_IO\n // Store the result\n memcpy(h_data_sink, h_data_out[current_stream], memsize);\n\n // Read new input\n memcpy(h_data_in[next_stream], h_data_source, memsize);\n#endif\n\n // Ensure that processing and copying of the last cycle has finished\n hipEventSynchronize(cycleDone[next_stream]);\n\n // Process current frame\n incKernel<<>>(\n d_data_out[current_stream], d_data_in[current_stream], N, inner_reps);\n\n // Upload next frame\n checkCudaErrors(hipMemcpyAsync(\n d_data_in[next_stream], h_data_in[next_stream], memsize, hipMemcpyHostToDevice, stream[next_stream]));\n\n // Download current frame\n checkCudaErrors(hipMemcpyAsync(h_data_out[current_stream],\n d_data_out[current_stream],\n memsize,\n hipMemcpyDeviceToHost,\n stream[current_stream]));\n\n checkCudaErrors(hipEventRecord(cycleDone[current_stream], stream[current_stream]));\n\n current_stream = next_stream;\n }\n\n hipEventRecord(stop, 0);\n\n hipDeviceSynchronize();\n\n hipEventElapsedTime(&time, start, stop);\n\n return time;\n}\n\nvoid init()\n{\n for (int i = 0; i < N; ++i) {\n h_data_source[i] = 0;\n }\n\n for (int i = 0; i < STREAM_COUNT; ++i) {\n memcpy(h_data_in[i], h_data_source, memsize);\n }\n}\n\nbool test()\n{\n bool passed = true;\n\n for (int j = 0; j < STREAM_COUNT; ++j) {\n for (int i = 0; i < N; ++i) {\n passed &= (h_data_out[j][i] == 1);\n }\n }\n\n return passed;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleMultiCopy/simpleMultiCopy.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* Vector addition: C = A + B.\n *\n * This sample is a very basic sample that implements element by element\n * vector addition. It is the same as the sample illustrating Chapter 3\n * of the programming guide with some additions like error checking.\n *\n */\n\n// Device code\nextern \"C\" __global__ void VecAdd_kernel(const float *A, const float *B, float *C, int N)\n{\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (i < N)\n C[i] = A[i] + B[i];\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* Vector addition: C = A + B.\n *\n * This sample is a very basic sample that implements element by element\n * vector addition. It is the same as the sample illustrating Chapter 3\n * of the programming guide with some additions like error checking.\n *\n */\n\n// Device code\n#include \"hip/hip_runtime.h\"\nextern \"C\" __global__ void VecAdd_kernel(const float *A, const float *B, float *C, int N)\n{\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (i < N)\n C[i] = A[i] + B[i];\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleDrvRuntime/vectorAdd_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates how use texture fetches in CUDA\n *\n * This sample takes an input PGM image (imageFilename) and generates\n * an output PGM image (imageFilename_out). This CUDA kernel performs\n * a simple 2D transform (rotation) on the texture coordinates (u,v).\n */\n\n// Includes, system\n#include \n#include \n#include \n#include \n\n#ifdef _WIN32\n#define WINDOWS_LEAN_AND_MEAN\n#define NOMINMAX\n#include \n#endif\n\n// Includes CUDA\n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n\n#define MIN_EPSILON_ERROR 5e-3f\n\n////////////////////////////////////////////////////////////////////////////////\n// Define the files that are to be save and the reference images for validation\nconst char *imageFilename = \"teapot512.pgm\";\nconst char *refFilename = \"ref_rotated.pgm\";\nfloat angle = 0.5f; // angle to rotate image by (in radians)\n\n// Auto-Verification Code\nbool testResult = true;\n\nstatic const char *sampleName = \"simpleSurfaceWrite\";\n\n////////////////////////////////////////////////////////////////////////////////\n// Kernels\n////////////////////////////////////////////////////////////////////////////////\n//! Write to a cuArray (texture data source) using surface writes\n//! @param gIData input data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void surfaceWriteKernel(float *gIData, int width, int height, cudaSurfaceObject_t outputSurface)\n{\n // calculate surface coordinates\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // read from global memory and write to cuarray (via surface reference)\n surf2Dwrite(gIData[y * width + x], outputSurface, x * 4, y, cudaBoundaryModeTrap);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Transform an image using texture lookups\n//! @param gOData output data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void transformKernel(float *gOData, int width, int height, float theta, cudaTextureObject_t tex)\n{\n // calculate normalized texture coordinates\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n float u = x / (float)width;\n float v = y / (float)height;\n\n // transform coordinates\n u -= 0.5f;\n v -= 0.5f;\n float tu = u * cosf(theta) - v * sinf(theta) + 0.5f;\n float tv = v * cosf(theta) + u * sinf(theta) + 0.5f;\n\n // read from texture and write to global memory\n gOData[y * width + x] = tex2D(tex, tu, tv);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Declaration, forward\nvoid runTest(int argc, char **argv);\n\nextern \"C\" void computeGold(float *reference, float *idata, const unsigned int len);\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"%s starting...\\n\", sampleName);\n\n // Process command-line arguments\n if (argc > 1) {\n if (checkCmdLineFlag(argc, (const char **)argv, \"input\")) {\n getCmdLineArgumentString(argc, (const char **)argv, \"input\", (char **)&imageFilename);\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"reference\")) {\n getCmdLineArgumentString(argc, (const char **)argv, \"reference\", (char **)&refFilename);\n }\n else {\n printf(\"-input flag should be used with -reference flag\");\n exit(EXIT_FAILURE);\n }\n }\n else if (checkCmdLineFlag(argc, (const char **)argv, \"reference\")) {\n printf(\"-reference flag should be used with -input flag\");\n exit(EXIT_FAILURE);\n }\n }\n\n runTest(argc, argv);\n\n printf(\"%s completed, returned %s\\n\", sampleName, testResult ? \"OK\" : \"ERROR!\");\n exit(testResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nvoid runTest(int argc, char **argv)\n{\n // Use command-line specified CUDA device,\n // otherwise use device with highest Gflops/s\n int devID = findCudaDevice(argc, (const char **)argv);\n\n // Get number of SMs on this GPU\n cudaDeviceProp deviceProps;\n\n checkCudaErrors(cudaGetDeviceProperties(&deviceProps, devID));\n printf(\"CUDA device [%s] has %d Multi-Processors, SM %d.%d\\n\",\n deviceProps.name,\n deviceProps.multiProcessorCount,\n deviceProps.major,\n deviceProps.minor);\n\n // Load image from disk\n float *hData = NULL;\n unsigned int width, height;\n char *imagePath = sdkFindFilePath(imageFilename, argv[0]);\n\n if (imagePath == NULL) {\n printf(\"Unable to source image input file: %s\\n\", imageFilename);\n exit(EXIT_FAILURE);\n }\n\n sdkLoadPGM(imagePath, &hData, &width, &height);\n\n unsigned int size = width * height * sizeof(float);\n printf(\"Loaded '%s', %d x %d pixels\\n\", imageFilename, width, height);\n\n // Load reference image from image (output)\n float *hDataRef = (float *)malloc(size);\n char *refPath = sdkFindFilePath(refFilename, argv[0]);\n\n if (refPath == NULL) {\n printf(\"Unable to find reference image file: %s\\n\", refFilename);\n exit(EXIT_FAILURE);\n }\n\n sdkLoadPGM(refPath, &hDataRef, &width, &height);\n\n // Allocate device memory for result\n float *dData = NULL;\n checkCudaErrors(cudaMalloc((void **)&dData, size));\n\n // Allocate array and copy image data\n cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc(32, 0, 0, 0, cudaChannelFormatKindFloat);\n cudaArray *cuArray;\n checkCudaErrors(cudaMallocArray(&cuArray, &channelDesc, width, height, cudaArraySurfaceLoadStore));\n\n dim3 dimBlock(8, 8, 1);\n dim3 dimGrid(width / dimBlock.x, height / dimBlock.y, 1);\n\n cudaSurfaceObject_t outputSurface;\n cudaResourceDesc surfRes;\n memset(&surfRes, 0, sizeof(cudaResourceDesc));\n surfRes.resType = cudaResourceTypeArray;\n surfRes.res.array.array = cuArray;\n\n checkCudaErrors(cudaCreateSurfaceObject(&outputSurface, &surfRes));\n#if 1\n checkCudaErrors(cudaMemcpy(dData, hData, size, cudaMemcpyHostToDevice));\n surfaceWriteKernel<<>>(dData, width, height, outputSurface);\n#else // This is what differs from the example simpleTexture\n checkCudaErrors(cudaMemcpyToArray(cuArray, 0, 0, hData, size, cudaMemcpyHostToDevice));\n#endif\n\n cudaTextureObject_t tex;\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = cuArray;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = cudaFilterModeLinear;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n\n // Warmup\n transformKernel<<>>(dData, width, height, angle, tex);\n\n checkCudaErrors(cudaDeviceSynchronize());\n\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n sdkStartTimer(&timer);\n\n // Execute the kernel\n transformKernel<<>>(dData, width, height, angle, tex);\n\n // Check if kernel execution generated an error\n getLastCudaError(\"Kernel execution failed\");\n\n cudaDeviceSynchronize();\n sdkStopTimer(&timer);\n printf(\"Processing time: %f (ms)\\n\", sdkGetTimerValue(&timer));\n printf(\"%.2f Mpixels/sec\\n\", (width * height / (sdkGetTimerValue(&timer) / 1000.0f)) / 1e6);\n sdkDeleteTimer(&timer);\n\n // Allocate mem for the result on host side\n float *hOData = (float *)malloc(size);\n // copy result from device to host\n checkCudaErrors(cudaMemcpy(hOData, dData, size, cudaMemcpyDeviceToHost));\n\n // Write result to file\n char outputFilename[1024];\n strcpy(outputFilename, \"output.pgm\");\n sdkSavePGM(\"output.pgm\", hOData, width, height);\n printf(\"Wrote '%s'\\n\", outputFilename);\n\n // Write regression file if necessary\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // Write file for regression test\n sdkWriteFile(\"./data/regression.dat\", hOData, width * height, 0.0f, false);\n }\n else {\n // We need to reload the data from disk,\n // because it is inverted upon output\n sdkLoadPGM(outputFilename, &hOData, &width, &height);\n\n printf(\"Comparing files\\n\");\n printf(\"\\toutput: <%s>\\n\", outputFilename);\n printf(\"\\treference: <%s>\\n\", refPath);\n testResult = compareData(hOData, hDataRef, width * height, MIN_EPSILON_ERROR, 0.0f);\n }\n\n checkCudaErrors(cudaDestroySurfaceObject(outputSurface));\n checkCudaErrors(cudaDestroyTextureObject(tex));\n checkCudaErrors(cudaFree(dData));\n checkCudaErrors(cudaFreeArray(cuArray));\n free(imagePath);\n free(refPath);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates how use texture fetches in CUDA\n *\n * This sample takes an input PGM image (imageFilename) and generates\n * an output PGM image (imageFilename_out). This CUDA kernel performs\n * a simple 2D transform (rotation) on the texture coordinates (u,v).\n */\n\n// Includes, system\n#include \n#include \n#include \n#include \n\n#ifdef _WIN32\n#define WINDOWS_LEAN_AND_MEAN\n#define NOMINMAX\n#include \n#endif\n\n// Includes CUDA\n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n\n#define MIN_EPSILON_ERROR 5e-3f\n\n////////////////////////////////////////////////////////////////////////////////\n// Define the files that are to be save and the reference images for validation\nconst char *imageFilename = \"teapot512.pgm\";\nconst char *refFilename = \"ref_rotated.pgm\";\nfloat angle = 0.5f; // angle to rotate image by (in radians)\n\n// Auto-Verification Code\nbool testResult = true;\n\nstatic const char *sampleName = \"simpleSurfaceWrite\";\n\n////////////////////////////////////////////////////////////////////////////////\n// Kernels\n////////////////////////////////////////////////////////////////////////////////\n//! Write to a cuArray (texture data source) using surface writes\n//! @param gIData input data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void surfaceWriteKernel(float *gIData, int width, int height, hipSurfaceObject_t outputSurface)\n{\n // calculate surface coordinates\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // read from global memory and write to cuarray (via surface reference)\n surf2Dwrite(gIData[y * width + x], outputSurface, x * 4, y, hipBoundaryModeTrap);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Transform an image using texture lookups\n//! @param gOData output data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void transformKernel(float *gOData, int width, int height, float theta, hipTextureObject_t tex)\n{\n // calculate normalized texture coordinates\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n float u = x / (float)width;\n float v = y / (float)height;\n\n // transform coordinates\n u -= 0.5f;\n v -= 0.5f;\n float tu = u * cosf(theta) - v * sinf(theta) + 0.5f;\n float tv = v * cosf(theta) + u * sinf(theta) + 0.5f;\n\n // read from texture and write to global memory\n gOData[y * width + x] = tex2D(tex, tu, tv);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Declaration, forward\nvoid runTest(int argc, char **argv);\n\nextern \"C\" void computeGold(float *reference, float *idata, const unsigned int len);\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"%s starting...\\n\", sampleName);\n\n // Process command-line arguments\n if (argc > 1) {\n if (checkCmdLineFlag(argc, (const char **)argv, \"input\")) {\n getCmdLineArgumentString(argc, (const char **)argv, \"input\", (char **)&imageFilename);\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"reference\")) {\n getCmdLineArgumentString(argc, (const char **)argv, \"reference\", (char **)&refFilename);\n }\n else {\n printf(\"-input flag should be used with -reference flag\");\n exit(EXIT_FAILURE);\n }\n }\n else if (checkCmdLineFlag(argc, (const char **)argv, \"reference\")) {\n printf(\"-reference flag should be used with -input flag\");\n exit(EXIT_FAILURE);\n }\n }\n\n runTest(argc, argv);\n\n printf(\"%s completed, returned %s\\n\", sampleName, testResult ? \"OK\" : \"ERROR!\");\n exit(testResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nvoid runTest(int argc, char **argv)\n{\n // Use command-line specified CUDA device,\n // otherwise use device with highest Gflops/s\n int devID = findCudaDevice(argc, (const char **)argv);\n\n // Get number of SMs on this GPU\n hipDeviceProp_t deviceProps;\n\n checkCudaErrors(hipGetDeviceProperties(&deviceProps, devID));\n printf(\"CUDA device [%s] has %d Multi-Processors, SM %d.%d\\n\",\n deviceProps.name,\n deviceProps.multiProcessorCount,\n deviceProps.major,\n deviceProps.minor);\n\n // Load image from disk\n float *hData = NULL;\n unsigned int width, height;\n char *imagePath = sdkFindFilePath(imageFilename, argv[0]);\n\n if (imagePath == NULL) {\n printf(\"Unable to source image input file: %s\\n\", imageFilename);\n exit(EXIT_FAILURE);\n }\n\n sdkLoadPGM(imagePath, &hData, &width, &height);\n\n unsigned int size = width * height * sizeof(float);\n printf(\"Loaded '%s', %d x %d pixels\\n\", imageFilename, width, height);\n\n // Load reference image from image (output)\n float *hDataRef = (float *)malloc(size);\n char *refPath = sdkFindFilePath(refFilename, argv[0]);\n\n if (refPath == NULL) {\n printf(\"Unable to find reference image file: %s\\n\", refFilename);\n exit(EXIT_FAILURE);\n }\n\n sdkLoadPGM(refPath, &hDataRef, &width, &height);\n\n // Allocate device memory for result\n float *dData = NULL;\n checkCudaErrors(hipMalloc((void **)&dData, size));\n\n // Allocate array and copy image data\n hipChannelFormatDesc channelDesc = hipCreateChannelDesc(32, 0, 0, 0, hipChannelFormatKindFloat);\n hipArray *cuArray;\n checkCudaErrors(hipMallocArray(&cuArray, &channelDesc, width, height, hipArraySurfaceLoadStore));\n\n dim3 dimBlock(8, 8, 1);\n dim3 dimGrid(width / dimBlock.x, height / dimBlock.y, 1);\n\n hipSurfaceObject_t outputSurface;\n hipResourceDesc surfRes;\n memset(&surfRes, 0, sizeof(hipResourceDesc));\n surfRes.resType = hipResourceTypeArray;\n surfRes.res.array.array = cuArray;\n\n checkCudaErrors(hipCreateSurfaceObject(&outputSurface, &surfRes));\n#if 1\n checkCudaErrors(hipMemcpy(dData, hData, size, hipMemcpyHostToDevice));\n surfaceWriteKernel<<>>(dData, width, height, outputSurface);\n#else // This is what differs from the example simpleTexture\n checkCudaErrors(hipMemcpyToArray(cuArray, 0, 0, hData, size, hipMemcpyHostToDevice));\n#endif\n\n hipTextureObject_t tex;\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = cuArray;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = hipFilterModeLinear;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n\n // Warmup\n transformKernel<<>>(dData, width, height, angle, tex);\n\n checkCudaErrors(hipDeviceSynchronize());\n\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n sdkStartTimer(&timer);\n\n // Execute the kernel\n transformKernel<<>>(dData, width, height, angle, tex);\n\n // Check if kernel execution generated an error\n getLastCudaError(\"Kernel execution failed\");\n\n hipDeviceSynchronize();\n sdkStopTimer(&timer);\n printf(\"Processing time: %f (ms)\\n\", sdkGetTimerValue(&timer));\n printf(\"%.2f Mpixels/sec\\n\", (width * height / (sdkGetTimerValue(&timer) / 1000.0f)) / 1e6);\n sdkDeleteTimer(&timer);\n\n // Allocate mem for the result on host side\n float *hOData = (float *)malloc(size);\n // copy result from device to host\n checkCudaErrors(hipMemcpy(hOData, dData, size, hipMemcpyDeviceToHost));\n\n // Write result to file\n char outputFilename[1024];\n strcpy(outputFilename, \"output.pgm\");\n sdkSavePGM(\"output.pgm\", hOData, width, height);\n printf(\"Wrote '%s'\\n\", outputFilename);\n\n // Write regression file if necessary\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // Write file for regression test\n sdkWriteFile(\"./data/regression.dat\", hOData, width * height, 0.0f, false);\n }\n else {\n // We need to reload the data from disk,\n // because it is inverted upon output\n sdkLoadPGM(outputFilename, &hOData, &width, &height);\n\n printf(\"Comparing files\\n\");\n printf(\"\\toutput: <%s>\\n\", outputFilename);\n printf(\"\\treference: <%s>\\n\", refPath);\n testResult = compareData(hOData, hDataRef, width * height, MIN_EPSILON_ERROR, 0.0f);\n }\n\n checkCudaErrors(hipDestroySurfaceObject(outputSurface));\n checkCudaErrors(hipDestroyTextureObject(tex));\n checkCudaErrors(hipFree(dData));\n checkCudaErrors(hipFreeArray(cuArray));\n free(imagePath);\n free(refPath);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleSurfaceWrite/simpleSurfaceWrite.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n//\n// This sample demonstrates how HyperQ allows supporting devices to avoid false\n// dependencies between kernels in different streams.\n//\n// - Devices without HyperQ will run a maximum of two kernels at a time (one\n// kernel_A and one kernel_B).\n// - Devices with HyperQ will run up to 32 kernels simultaneously.\n\n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \n#include \n\nconst char *sSDKsample = \"hyperQ\";\n\n// This subroutine does no real work but runs for at least the specified number\n// of clock ticks.\n__device__ void clock_block(clock_t *d_o, clock_t clock_count)\n{\n unsigned int start_clock = (unsigned int)clock();\n\n clock_t clock_offset = 0;\n\n while (clock_offset < clock_count) {\n unsigned int end_clock = (unsigned int)clock();\n\n // The code below should work like\n // this (thanks to modular arithmetics):\n //\n // clock_offset = (clock_t) (end_clock > start_clock ?\n // end_clock - start_clock :\n // end_clock + (0xffffffffu - start_clock));\n //\n // Indeed, let m = 2^32 then\n // end - start = end + m - start (mod m).\n\n clock_offset = (clock_t)(end_clock - start_clock);\n }\n\n d_o[0] = clock_offset;\n}\n\n// We create two identical kernels calling clock_block(), we create two so that\n// we can identify dependencies in the profile timeline (\"kernel_B\" is always\n// dependent on \"kernel_A\" in the same stream).\n__global__ void kernel_A(clock_t *d_o, clock_t clock_count) { clock_block(d_o, clock_count); }\n__global__ void kernel_B(clock_t *d_o, clock_t clock_count) { clock_block(d_o, clock_count); }\n\n// Single-warp reduction kernel (note: this is not optimized for simplicity)\n__global__ void sum(clock_t *d_clocks, int N)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ clock_t s_clocks[32];\n\n clock_t my_sum = 0;\n\n for (int i = threadIdx.x; i < N; i += blockDim.x) {\n my_sum += d_clocks[i];\n }\n\n s_clocks[threadIdx.x] = my_sum;\n cg::sync(cta);\n\n for (int i = warpSize / 2; i > 0; i /= 2) {\n if (threadIdx.x < i) {\n s_clocks[threadIdx.x] += s_clocks[threadIdx.x + i];\n }\n\n cg::sync(cta);\n }\n\n if (threadIdx.x == 0) {\n d_clocks[0] = s_clocks[0];\n }\n}\n\nint main(int argc, char **argv)\n{\n int nstreams = 32; // One stream for each pair of kernels\n float kernel_time = 10; // Time each kernel should run in ms\n float elapsed_time;\n int cuda_device = 0;\n\n printf(\"starting %s...\\n\", sSDKsample);\n\n // Get number of streams (if overridden on the command line)\n if (checkCmdLineFlag(argc, (const char **)argv, \"nstreams\")) {\n nstreams = getCmdLineArgumentInt(argc, (const char **)argv, \"nstreams\");\n }\n\n // Use command-line specified CUDA device, otherwise use device with\n // highest Gflops/s\n cuda_device = findCudaDevice(argc, (const char **)argv);\n\n // Get device properties\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDevice(&cuda_device));\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, cuda_device));\n\n // Get device clock rate\n int clockRate;\n checkCudaErrors(cudaDeviceGetAttribute(&clockRate, cudaDevAttrClockRate, cuda_device));\n\n // HyperQ is available in devices of Compute Capability 3.5 and higher\n if (deviceProp.major < 3 || (deviceProp.major == 3 && deviceProp.minor < 5)) {\n if (deviceProp.concurrentKernels == 0) {\n printf(\"> GPU does not support concurrent kernel execution (SM 3.5 or \"\n \"higher required)\\n\");\n printf(\" CUDA kernel runs will be serialized\\n\");\n }\n else {\n printf(\"> GPU does not support HyperQ\\n\");\n printf(\" CUDA kernel runs will have limited concurrency\\n\");\n }\n }\n\n printf(\"> Detected Compute SM %d.%d hardware with %d multi-processors\\n\",\n deviceProp.major,\n deviceProp.minor,\n deviceProp.multiProcessorCount);\n\n // Allocate host memory for the output (reduced to a single value)\n clock_t *a = 0;\n checkCudaErrors(cudaMallocHost((void **)&a, sizeof(clock_t)));\n\n // Allocate device memory for the output (one value for each kernel)\n clock_t *d_a = 0;\n checkCudaErrors(cudaMalloc((void **)&d_a, 2 * nstreams * sizeof(clock_t)));\n\n // Allocate and initialize an array of stream handles\n cudaStream_t *streams = (cudaStream_t *)malloc(nstreams * sizeof(cudaStream_t));\n\n for (int i = 0; i < nstreams; i++) {\n checkCudaErrors(cudaStreamCreate(&(streams[i])));\n }\n\n // Create CUDA event handles\n cudaEvent_t start_event, stop_event;\n checkCudaErrors(cudaEventCreate(&start_event));\n checkCudaErrors(cudaEventCreate(&stop_event));\n\n // Target time per kernel is kernel_time ms, clockRate is in KHz\n // Target number of clocks = target time * clock frequency\n#if defined(__arm__) || defined(__aarch64__)\n // the kernel takes more time than the channel reset time on arm archs, so to\n // prevent hangs reduce time_clocks.\n clock_t time_clocks = (clock_t)(kernel_time * (clockRate / 100));\n#else\n clock_t time_clocks = (clock_t)(kernel_time * clockRate);\n#endif\n clock_t total_clocks = 0;\n\n // Start the clock\n checkCudaErrors(cudaEventRecord(start_event, 0));\n\n // Queue pairs of {kernel_A, kernel_B} in separate streams\n for (int i = 0; i < nstreams; ++i) {\n kernel_A<<<1, 1, 0, streams[i]>>>(&d_a[2 * i], time_clocks);\n total_clocks += time_clocks;\n kernel_B<<<1, 1, 0, streams[i]>>>(&d_a[2 * i + 1], time_clocks);\n total_clocks += time_clocks;\n }\n\n // Stop the clock in stream 0 (i.e. all previous kernels will be complete)\n checkCudaErrors(cudaEventRecord(stop_event, 0));\n\n // At this point the CPU has dispatched all work for the GPU and can\n // continue processing other tasks in parallel. In this sample we just want\n // to wait until all work is done so we use a blocking cudaMemcpy below.\n\n // Run the sum kernel and copy the result back to host\n sum<<<1, 32>>>(d_a, 2 * nstreams);\n checkCudaErrors(cudaMemcpy(a, d_a, sizeof(clock_t), cudaMemcpyDeviceToHost));\n\n // stop_event will have been recorded but including the synchronize here to\n // prevent copy/paste errors!\n checkCudaErrors(cudaEventSynchronize(stop_event));\n checkCudaErrors(cudaEventElapsedTime(&elapsed_time, start_event, stop_event));\n\n printf(\"Expected time for serial execution of %d sets of kernels is between \"\n \"approx. %.3fs and %.3fs\\n\",\n nstreams,\n (nstreams + 1) * kernel_time / 1000.0f,\n 2 * nstreams * kernel_time / 1000.0f);\n printf(\"Expected time for fully concurrent execution of %d sets of kernels is \"\n \"approx. %.3fs\\n\",\n nstreams,\n 2 * kernel_time / 1000.0f);\n printf(\"Measured time for sample = %.3fs\\n\", elapsed_time / 1000.0f);\n\n bool bTestResult = (a[0] >= total_clocks);\n\n // Release resources\n for (int i = 0; i < nstreams; i++) {\n cudaStreamDestroy(streams[i]);\n }\n\n free(streams);\n cudaEventDestroy(start_event);\n cudaEventDestroy(stop_event);\n cudaFreeHost(a);\n cudaFree(d_a);\n\n exit(bTestResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n//\n// This sample demonstrates how HyperQ allows supporting devices to avoid false\n// dependencies between kernels in different streams.\n//\n// - Devices without HyperQ will run a maximum of two kernels at a time (one\n// kernel_A and one kernel_B).\n// - Devices with HyperQ will run up to 32 kernels simultaneously.\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \n#include \n\nconst char *sSDKsample = \"hyperQ\";\n\n// This subroutine does no real work but runs for at least the specified number\n// of clock ticks.\n__device__ void clock_block(clock_t *d_o, clock_t clock_count)\n{\n unsigned int start_clock = (unsigned int)clock();\n\n clock_t clock_offset = 0;\n\n while (clock_offset < clock_count) {\n unsigned int end_clock = (unsigned int)clock();\n\n // The code below should work like\n // this (thanks to modular arithmetics):\n //\n // clock_offset = (clock_t) (end_clock > start_clock ?\n // end_clock - start_clock :\n // end_clock + (0xffffffffu - start_clock));\n //\n // Indeed, let m = 2^32 then\n // end - start = end + m - start (mod m).\n\n clock_offset = (clock_t)(end_clock - start_clock);\n }\n\n d_o[0] = clock_offset;\n}\n\n// We create two identical kernels calling clock_block(), we create two so that\n// we can identify dependencies in the profile timeline (\"kernel_B\" is always\n// dependent on \"kernel_A\" in the same stream).\n__global__ void kernel_A(clock_t *d_o, clock_t clock_count) { clock_block(d_o, clock_count); }\n__global__ void kernel_B(clock_t *d_o, clock_t clock_count) { clock_block(d_o, clock_count); }\n\n// Single-warp reduction kernel (note: this is not optimized for simplicity)\n__global__ void sum(clock_t *d_clocks, int N)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ clock_t s_clocks[32];\n\n clock_t my_sum = 0;\n\n for (int i = threadIdx.x; i < N; i += blockDim.x) {\n my_sum += d_clocks[i];\n }\n\n s_clocks[threadIdx.x] = my_sum;\n cg::sync(cta);\n\n for (int i = warpSize / 2; i > 0; i /= 2) {\n if (threadIdx.x < i) {\n s_clocks[threadIdx.x] += s_clocks[threadIdx.x + i];\n }\n\n cg::sync(cta);\n }\n\n if (threadIdx.x == 0) {\n d_clocks[0] = s_clocks[0];\n }\n}\n\nint main(int argc, char **argv)\n{\n int nstreams = 32; // One stream for each pair of kernels\n float kernel_time = 10; // Time each kernel should run in ms\n float elapsed_time;\n int cuda_device = 0;\n\n printf(\"starting %s...\\n\", sSDKsample);\n\n // Get number of streams (if overridden on the command line)\n if (checkCmdLineFlag(argc, (const char **)argv, \"nstreams\")) {\n nstreams = getCmdLineArgumentInt(argc, (const char **)argv, \"nstreams\");\n }\n\n // Use command-line specified CUDA device, otherwise use device with\n // highest Gflops/s\n cuda_device = findCudaDevice(argc, (const char **)argv);\n\n // Get device properties\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDevice(&cuda_device));\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, cuda_device));\n\n // Get device clock rate\n int clockRate;\n checkCudaErrors(hipDeviceGetAttribute(&clockRate, hipDeviceAttributeClockRate, cuda_device));\n\n // HyperQ is available in devices of Compute Capability 3.5 and higher\n if (deviceProp.major < 3 || (deviceProp.major == 3 && deviceProp.minor < 5)) {\n if (deviceProp.concurrentKernels == 0) {\n printf(\"> GPU does not support concurrent kernel execution (SM 3.5 or \"\n \"higher required)\\n\");\n printf(\" CUDA kernel runs will be serialized\\n\");\n }\n else {\n printf(\"> GPU does not support HyperQ\\n\");\n printf(\" CUDA kernel runs will have limited concurrency\\n\");\n }\n }\n\n printf(\"> Detected Compute SM %d.%d hardware with %d multi-processors\\n\",\n deviceProp.major,\n deviceProp.minor,\n deviceProp.multiProcessorCount);\n\n // Allocate host memory for the output (reduced to a single value)\n clock_t *a = 0;\n checkCudaErrors(hipHostMalloc((void **)&a, sizeof(clock_t)));\n\n // Allocate device memory for the output (one value for each kernel)\n clock_t *d_a = 0;\n checkCudaErrors(hipMalloc((void **)&d_a, 2 * nstreams * sizeof(clock_t)));\n\n // Allocate and initialize an array of stream handles\n hipStream_t *streams = (hipStream_t *)malloc(nstreams * sizeof(hipStream_t));\n\n for (int i = 0; i < nstreams; i++) {\n checkCudaErrors(hipStreamCreate(&(streams[i])));\n }\n\n // Create CUDA event handles\n hipEvent_t start_event, stop_event;\n checkCudaErrors(hipEventCreate(&start_event));\n checkCudaErrors(hipEventCreate(&stop_event));\n\n // Target time per kernel is kernel_time ms, clockRate is in KHz\n // Target number of clocks = target time * clock frequency\n#if defined(__arm__) || defined(__aarch64__)\n // the kernel takes more time than the channel reset time on arm archs, so to\n // prevent hangs reduce time_clocks.\n clock_t time_clocks = (clock_t)(kernel_time * (clockRate / 100));\n#else\n clock_t time_clocks = (clock_t)(kernel_time * clockRate);\n#endif\n clock_t total_clocks = 0;\n\n // Start the clock\n checkCudaErrors(hipEventRecord(start_event, 0));\n\n // Queue pairs of {kernel_A, kernel_B} in separate streams\n for (int i = 0; i < nstreams; ++i) {\n kernel_A<<<1, 1, 0, streams[i]>>>(&d_a[2 * i], time_clocks);\n total_clocks += time_clocks;\n kernel_B<<<1, 1, 0, streams[i]>>>(&d_a[2 * i + 1], time_clocks);\n total_clocks += time_clocks;\n }\n\n // Stop the clock in stream 0 (i.e. all previous kernels will be complete)\n checkCudaErrors(hipEventRecord(stop_event, 0));\n\n // At this point the CPU has dispatched all work for the GPU and can\n // continue processing other tasks in parallel. In this sample we just want\n // to wait until all work is done so we use a blocking cudaMemcpy below.\n\n // Run the sum kernel and copy the result back to host\n sum<<<1, 32>>>(d_a, 2 * nstreams);\n checkCudaErrors(hipMemcpy(a, d_a, sizeof(clock_t), hipMemcpyDeviceToHost));\n\n // stop_event will have been recorded but including the synchronize here to\n // prevent copy/paste errors!\n checkCudaErrors(hipEventSynchronize(stop_event));\n checkCudaErrors(hipEventElapsedTime(&elapsed_time, start_event, stop_event));\n\n printf(\"Expected time for serial execution of %d sets of kernels is between \"\n \"approx. %.3fs and %.3fs\\n\",\n nstreams,\n (nstreams + 1) * kernel_time / 1000.0f,\n 2 * nstreams * kernel_time / 1000.0f);\n printf(\"Expected time for fully concurrent execution of %d sets of kernels is \"\n \"approx. %.3fs\\n\",\n nstreams,\n 2 * kernel_time / 1000.0f);\n printf(\"Measured time for sample = %.3fs\\n\", elapsed_time / 1000.0f);\n\n bool bTestResult = (a[0] >= total_clocks);\n\n // Release resources\n for (int i = 0; i < nstreams; i++) {\n hipStreamDestroy(streams[i]);\n }\n\n free(streams);\n hipEventDestroy(start_event);\n hipEventDestroy(stop_event);\n hipHostFree(a);\n hipFree(d_a);\n\n exit(bTestResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleHyperQ/simpleHyperQ.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates how use texture fetches in CUDA\n *\n * This sample takes an input PGM image (image_filename) and generates\n * an output PGM image (image_filename_out). This CUDA kernel performs\n * a simple 2D transform (rotation) on the texture coordinates (u,v).\n */\n\n// Includes, system\n#include \n#include \n#include \n#include \n\n#ifdef _WIN32\n#define WINDOWS_LEAN_AND_MEAN\n#define NOMINMAX\n#include \n#endif\n\n// Includes CUDA\n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n\n#define MAX_EPSILON_ERROR 5e-3f\n\n// Define the files that are to be save and the reference images for validation\nconst char *imageFilename = \"teapot512.pgm\";\nconst char *refFilename = \"ref_rotated.pgm\";\n\nconst char *sampleName = \"simpleTexture\";\n\n////////////////////////////////////////////////////////////////////////////////\n// Constants\nconst float angle = 0.5f; // angle to rotate image by (in radians)\n\n// Auto-Verification Code\nbool testResult = true;\n\n////////////////////////////////////////////////////////////////////////////////\n//! Transform an image using texture lookups\n//! @param outputData output data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void transformKernel(float *outputData, int width, int height, float theta, cudaTextureObject_t tex)\n{\n // calculate normalized texture coordinates\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n float u = (float)x - (float)width / 2;\n float v = (float)y - (float)height / 2;\n float tu = u * cosf(theta) - v * sinf(theta);\n float tv = v * cosf(theta) + u * sinf(theta);\n\n tu /= (float)width;\n tv /= (float)height;\n\n // read from texture and write to global memory\n outputData[y * width + x] = tex2D(tex, tu + 0.5f, tv + 0.5f);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Declaration, forward\nvoid runTest(int argc, char **argv);\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"%s starting...\\n\", sampleName);\n\n // Process command-line arguments\n if (argc > 1) {\n if (checkCmdLineFlag(argc, (const char **)argv, \"input\")) {\n getCmdLineArgumentString(argc, (const char **)argv, \"input\", (char **)&imageFilename);\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"reference\")) {\n getCmdLineArgumentString(argc, (const char **)argv, \"reference\", (char **)&refFilename);\n }\n else {\n printf(\"-input flag should be used with -reference flag\");\n exit(EXIT_FAILURE);\n }\n }\n else if (checkCmdLineFlag(argc, (const char **)argv, \"reference\")) {\n printf(\"-reference flag should be used with -input flag\");\n exit(EXIT_FAILURE);\n }\n }\n\n runTest(argc, argv);\n\n printf(\"%s completed, returned %s\\n\", sampleName, testResult ? \"OK\" : \"ERROR!\");\n exit(testResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nvoid runTest(int argc, char **argv)\n{\n int devID = findCudaDevice(argc, (const char **)argv);\n\n // load image from disk\n float *hData = NULL;\n unsigned int width, height;\n char *imagePath = sdkFindFilePath(imageFilename, argv[0]);\n\n if (imagePath == NULL) {\n printf(\"Unable to source image file: %s\\n\", imageFilename);\n exit(EXIT_FAILURE);\n }\n\n sdkLoadPGM(imagePath, &hData, &width, &height);\n\n unsigned int size = width * height * sizeof(float);\n printf(\"Loaded '%s', %d x %d pixels\\n\", imageFilename, width, height);\n\n // Load reference image from image (output)\n float *hDataRef = (float *)malloc(size);\n char *refPath = sdkFindFilePath(refFilename, argv[0]);\n\n if (refPath == NULL) {\n printf(\"Unable to find reference image file: %s\\n\", refFilename);\n exit(EXIT_FAILURE);\n }\n\n sdkLoadPGM(refPath, &hDataRef, &width, &height);\n\n // Allocate device memory for result\n float *dData = NULL;\n checkCudaErrors(cudaMalloc((void **)&dData, size));\n\n // Allocate array and copy image data\n cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc(32, 0, 0, 0, cudaChannelFormatKindFloat);\n cudaArray *cuArray;\n checkCudaErrors(cudaMallocArray(&cuArray, &channelDesc, width, height));\n checkCudaErrors(cudaMemcpyToArray(cuArray, 0, 0, hData, size, cudaMemcpyHostToDevice));\n\n cudaTextureObject_t tex;\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = cuArray;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = cudaFilterModeLinear;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n\n dim3 dimBlock(8, 8, 1);\n dim3 dimGrid(width / dimBlock.x, height / dimBlock.y, 1);\n\n // Warmup\n transformKernel<<>>(dData, width, height, angle, tex);\n\n checkCudaErrors(cudaDeviceSynchronize());\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n sdkStartTimer(&timer);\n\n // Execute the kernel\n transformKernel<<>>(dData, width, height, angle, tex);\n\n // Check if kernel execution generated an error\n getLastCudaError(\"Kernel execution failed\");\n\n checkCudaErrors(cudaDeviceSynchronize());\n sdkStopTimer(&timer);\n printf(\"Processing time: %f (ms)\\n\", sdkGetTimerValue(&timer));\n printf(\"%.2f Mpixels/sec\\n\", (width * height / (sdkGetTimerValue(&timer) / 1000.0f)) / 1e6);\n sdkDeleteTimer(&timer);\n\n // Allocate mem for the result on host side\n float *hOutputData = (float *)malloc(size);\n // copy result from device to host\n checkCudaErrors(cudaMemcpy(hOutputData, dData, size, cudaMemcpyDeviceToHost));\n\n // Write result to file\n char outputFilename[1024];\n strcpy(outputFilename, imagePath);\n strcpy(outputFilename + strlen(imagePath) - 4, \"_out.pgm\");\n sdkSavePGM(outputFilename, hOutputData, width, height);\n printf(\"Wrote '%s'\\n\", outputFilename);\n\n // Write regression file if necessary\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // Write file for regression test\n sdkWriteFile(\"./data/regression.dat\", hOutputData, width * height, 0.0f, false);\n }\n else {\n // We need to reload the data from disk,\n // because it is inverted upon output\n sdkLoadPGM(outputFilename, &hOutputData, &width, &height);\n\n printf(\"Comparing files\\n\");\n printf(\"\\toutput: <%s>\\n\", outputFilename);\n printf(\"\\treference: <%s>\\n\", refPath);\n\n testResult = compareData(hOutputData, hDataRef, width * height, MAX_EPSILON_ERROR, 0.15f);\n }\n\n checkCudaErrors(cudaDestroyTextureObject(tex));\n checkCudaErrors(cudaFree(dData));\n checkCudaErrors(cudaFreeArray(cuArray));\n free(imagePath);\n free(refPath);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates how use texture fetches in CUDA\n *\n * This sample takes an input PGM image (image_filename) and generates\n * an output PGM image (image_filename_out). This CUDA kernel performs\n * a simple 2D transform (rotation) on the texture coordinates (u,v).\n */\n\n// Includes, system\n#include \n#include \n#include \n#include \n\n#ifdef _WIN32\n#define WINDOWS_LEAN_AND_MEAN\n#define NOMINMAX\n#include \n#endif\n\n// Includes CUDA\n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n\n#define MAX_EPSILON_ERROR 5e-3f\n\n// Define the files that are to be save and the reference images for validation\nconst char *imageFilename = \"teapot512.pgm\";\nconst char *refFilename = \"ref_rotated.pgm\";\n\nconst char *sampleName = \"simpleTexture\";\n\n////////////////////////////////////////////////////////////////////////////////\n// Constants\nconst float angle = 0.5f; // angle to rotate image by (in radians)\n\n// Auto-Verification Code\nbool testResult = true;\n\n////////////////////////////////////////////////////////////////////////////////\n//! Transform an image using texture lookups\n//! @param outputData output data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void transformKernel(float *outputData, int width, int height, float theta, hipTextureObject_t tex)\n{\n // calculate normalized texture coordinates\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n float u = (float)x - (float)width / 2;\n float v = (float)y - (float)height / 2;\n float tu = u * cosf(theta) - v * sinf(theta);\n float tv = v * cosf(theta) + u * sinf(theta);\n\n tu /= (float)width;\n tv /= (float)height;\n\n // read from texture and write to global memory\n outputData[y * width + x] = tex2D(tex, tu + 0.5f, tv + 0.5f);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Declaration, forward\nvoid runTest(int argc, char **argv);\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"%s starting...\\n\", sampleName);\n\n // Process command-line arguments\n if (argc > 1) {\n if (checkCmdLineFlag(argc, (const char **)argv, \"input\")) {\n getCmdLineArgumentString(argc, (const char **)argv, \"input\", (char **)&imageFilename);\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"reference\")) {\n getCmdLineArgumentString(argc, (const char **)argv, \"reference\", (char **)&refFilename);\n }\n else {\n printf(\"-input flag should be used with -reference flag\");\n exit(EXIT_FAILURE);\n }\n }\n else if (checkCmdLineFlag(argc, (const char **)argv, \"reference\")) {\n printf(\"-reference flag should be used with -input flag\");\n exit(EXIT_FAILURE);\n }\n }\n\n runTest(argc, argv);\n\n printf(\"%s completed, returned %s\\n\", sampleName, testResult ? \"OK\" : \"ERROR!\");\n exit(testResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nvoid runTest(int argc, char **argv)\n{\n int devID = findCudaDevice(argc, (const char **)argv);\n\n // load image from disk\n float *hData = NULL;\n unsigned int width, height;\n char *imagePath = sdkFindFilePath(imageFilename, argv[0]);\n\n if (imagePath == NULL) {\n printf(\"Unable to source image file: %s\\n\", imageFilename);\n exit(EXIT_FAILURE);\n }\n\n sdkLoadPGM(imagePath, &hData, &width, &height);\n\n unsigned int size = width * height * sizeof(float);\n printf(\"Loaded '%s', %d x %d pixels\\n\", imageFilename, width, height);\n\n // Load reference image from image (output)\n float *hDataRef = (float *)malloc(size);\n char *refPath = sdkFindFilePath(refFilename, argv[0]);\n\n if (refPath == NULL) {\n printf(\"Unable to find reference image file: %s\\n\", refFilename);\n exit(EXIT_FAILURE);\n }\n\n sdkLoadPGM(refPath, &hDataRef, &width, &height);\n\n // Allocate device memory for result\n float *dData = NULL;\n checkCudaErrors(hipMalloc((void **)&dData, size));\n\n // Allocate array and copy image data\n hipChannelFormatDesc channelDesc = hipCreateChannelDesc(32, 0, 0, 0, hipChannelFormatKindFloat);\n hipArray *cuArray;\n checkCudaErrors(hipMallocArray(&cuArray, &channelDesc, width, height));\n checkCudaErrors(hipMemcpyToArray(cuArray, 0, 0, hData, size, hipMemcpyHostToDevice));\n\n hipTextureObject_t tex;\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = cuArray;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = hipFilterModeLinear;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n\n dim3 dimBlock(8, 8, 1);\n dim3 dimGrid(width / dimBlock.x, height / dimBlock.y, 1);\n\n // Warmup\n transformKernel<<>>(dData, width, height, angle, tex);\n\n checkCudaErrors(hipDeviceSynchronize());\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n sdkStartTimer(&timer);\n\n // Execute the kernel\n transformKernel<<>>(dData, width, height, angle, tex);\n\n // Check if kernel execution generated an error\n getLastCudaError(\"Kernel execution failed\");\n\n checkCudaErrors(hipDeviceSynchronize());\n sdkStopTimer(&timer);\n printf(\"Processing time: %f (ms)\\n\", sdkGetTimerValue(&timer));\n printf(\"%.2f Mpixels/sec\\n\", (width * height / (sdkGetTimerValue(&timer) / 1000.0f)) / 1e6);\n sdkDeleteTimer(&timer);\n\n // Allocate mem for the result on host side\n float *hOutputData = (float *)malloc(size);\n // copy result from device to host\n checkCudaErrors(hipMemcpy(hOutputData, dData, size, hipMemcpyDeviceToHost));\n\n // Write result to file\n char outputFilename[1024];\n strcpy(outputFilename, imagePath);\n strcpy(outputFilename + strlen(imagePath) - 4, \"_out.pgm\");\n sdkSavePGM(outputFilename, hOutputData, width, height);\n printf(\"Wrote '%s'\\n\", outputFilename);\n\n // Write regression file if necessary\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // Write file for regression test\n sdkWriteFile(\"./data/regression.dat\", hOutputData, width * height, 0.0f, false);\n }\n else {\n // We need to reload the data from disk,\n // because it is inverted upon output\n sdkLoadPGM(outputFilename, &hOutputData, &width, &height);\n\n printf(\"Comparing files\\n\");\n printf(\"\\toutput: <%s>\\n\", outputFilename);\n printf(\"\\treference: <%s>\\n\", refPath);\n\n testResult = compareData(hOutputData, hDataRef, width * height, MAX_EPSILON_ERROR, 0.15f);\n }\n\n checkCudaErrors(hipDestroyTextureObject(tex));\n checkCudaErrors(hipFree(dData));\n checkCudaErrors(hipFreeArray(cuArray));\n free(imagePath);\n free(refPath);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleTexture/simpleTexture.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#ifndef _SIMPLETEXTURE_KERNEL_H_\n#define _SIMPLETEXTURE_KERNEL_H_\n#include \n\n////////////////////////////////////////////////////////////////////////////////\n//! Transform an image using texture lookups\n//! @param g_odata output data in global memory\n////////////////////////////////////////////////////////////////////////////////\nextern \"C\" __global__ void transformKernel(float *g_odata, int width, int height, float theta, CUtexObject tex)\n{\n // calculate normalized texture coordinates\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n float u = (float)x - (float)width / 2;\n float v = (float)y - (float)height / 2;\n float tu = u * cosf(theta) - v * sinf(theta);\n float tv = v * cosf(theta) + u * sinf(theta);\n\n tu /= (float)width;\n tv /= (float)height;\n\n // read from texture and write to global memory\n g_odata[y * width + x] = tex2D(tex, tu + 0.5f, tv + 0.5f);\n}\n\n#endif // #ifndef _SIMPLETEXTURE_KERNEL_H_\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#ifndef _SIMPLETEXTURE_KERNEL_H_\n#define _SIMPLETEXTURE_KERNEL_H_\n#include \n\n////////////////////////////////////////////////////////////////////////////////\n//! Transform an image using texture lookups\n//! @param g_odata output data in global memory\n////////////////////////////////////////////////////////////////////////////////\nextern \"C\" __global__ void transformKernel(float *g_odata, int width, int height, float theta, hipTextureObject_t tex)\n{\n // calculate normalized texture coordinates\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n float u = (float)x - (float)width / 2;\n float v = (float)y - (float)height / 2;\n float tu = u * cosf(theta) - v * sinf(theta);\n float tv = v * cosf(theta) + u * sinf(theta);\n\n tu /= (float)width;\n tv /= (float)height;\n\n // read from texture and write to global memory\n g_odata[y * width + x] = tex2D(tex, tu + 0.5f, tv + 0.5f);\n}\n\n#endif // #ifndef _SIMPLETEXTURE_KERNEL_H_\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleTextureDrv/simpleTexture_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include // helper functions for CUDA error check\n#include \n\nconst int manualBlockSize = 32;\n\n////////////////////////////////////////////////////////////////////////////////\n// Test kernel\n//\n// This kernel squares each array element. Each thread addresses\n// himself with threadIdx and blockIdx, so that it can handle any\n// execution configuration, including anything the launch configurator\n// API suggests.\n////////////////////////////////////////////////////////////////////////////////\n__global__ void square(uint32_t *array, int arrayCount)\n{\n extern __shared__ int dynamicSmem[];\n int idx = threadIdx.x + blockIdx.x * blockDim.x;\n\n if (idx < arrayCount) {\n array[idx] *= array[idx];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Potential occupancy calculator\n//\n// The potential occupancy is calculated according to the kernel and\n// execution configuration the user desires. Occupancy is defined in\n// terms of active blocks per multiprocessor, and the user can convert\n// it to other metrics.\n//\n// This wrapper routine computes the occupancy of kernel, and reports\n// it in terms of active warps / maximum warps per SM.\n////////////////////////////////////////////////////////////////////////////////\nstatic double reportPotentialOccupancy(void *kernel, int blockSize, size_t dynamicSMem)\n{\n int device;\n cudaDeviceProp prop;\n\n int numBlocks;\n int activeWarps;\n int maxWarps;\n\n double occupancy;\n\n checkCudaErrors(cudaGetDevice(&device));\n checkCudaErrors(cudaGetDeviceProperties(&prop, device));\n\n checkCudaErrors(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&numBlocks, kernel, blockSize, dynamicSMem));\n\n activeWarps = numBlocks * blockSize / prop.warpSize;\n maxWarps = prop.maxThreadsPerMultiProcessor / prop.warpSize;\n\n occupancy = (double)activeWarps / maxWarps;\n\n return occupancy;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Occupancy-based launch configurator\n//\n// The launch configurator, cudaOccupancyMaxPotentialBlockSize and\n// cudaOccupancyMaxPotentialBlockSizeVariableSMem, suggests a block\n// size that achieves the best theoretical occupancy. It also returns\n// the minimum number of blocks needed to achieve the occupancy on the\n// whole device.\n//\n// This launch configurator is purely occupancy-based. It doesn't\n// translate directly to performance, but the suggestion should\n// nevertheless be a good starting point for further optimizations.\n//\n// This function configures the launch based on the \"automatic\"\n// argument, records the runtime, and reports occupancy and runtime.\n////////////////////////////////////////////////////////////////////////////////\nstatic int launchConfig(uint32_t *array, int arrayCount, bool automatic)\n{\n int blockSize;\n int minGridSize;\n int gridSize;\n size_t dynamicSMemUsage = 0;\n\n cudaEvent_t start;\n cudaEvent_t end;\n\n float elapsedTime;\n\n double potentialOccupancy;\n\n checkCudaErrors(cudaEventCreate(&start));\n checkCudaErrors(cudaEventCreate(&end));\n\n if (automatic) {\n checkCudaErrors(\n cudaOccupancyMaxPotentialBlockSize(&minGridSize, &blockSize, (void *)square, dynamicSMemUsage, arrayCount));\n\n std::cout << \"Suggested block size: \" << blockSize << std::endl\n << \"Minimum grid size for maximum occupancy: \" << minGridSize << std::endl;\n }\n else {\n // This block size is too small. Given limited number of\n // active blocks per multiprocessor, the number of active\n // threads will be limited, and thus unable to achieve maximum\n // occupancy.\n //\n blockSize = manualBlockSize;\n }\n\n // Round up\n //\n gridSize = (arrayCount + blockSize - 1) / blockSize;\n\n // Launch and profile\n //\n checkCudaErrors(cudaEventRecord(start));\n square<<>>(array, arrayCount);\n checkCudaErrors(cudaEventRecord(end));\n\n checkCudaErrors(cudaDeviceSynchronize());\n\n // Calculate occupancy\n //\n potentialOccupancy = reportPotentialOccupancy((void *)square, blockSize, dynamicSMemUsage);\n\n std::cout << \"Potential occupancy: \" << potentialOccupancy * 100 << \"%\" << std::endl;\n\n // Report elapsed time\n //\n checkCudaErrors(cudaEventElapsedTime(&elapsedTime, start, end));\n std::cout << \"Elapsed time: \" << elapsedTime << \"ms\" << std::endl;\n\n return 0;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// The test\n//\n// The test generates an array and squares it with a CUDA kernel, then\n// verifies the result.\n////////////////////////////////////////////////////////////////////////////////\nstatic int test(bool automaticLaunchConfig, const int count = 1000000)\n{\n uint32_t *array;\n uint32_t *dArray;\n int size = count * sizeof(uint32_t);\n\n array = new uint32_t[count];\n\n for (uint32_t i = 0; i < count; i += 1) {\n array[i] = i;\n }\n\n checkCudaErrors(cudaMalloc(&dArray, size));\n checkCudaErrors(cudaMemcpy(dArray, array, size, cudaMemcpyHostToDevice));\n\n for (uint32_t i = 0; i < count; i += 1) {\n array[i] = 0;\n }\n\n launchConfig(dArray, count, automaticLaunchConfig);\n\n checkCudaErrors(cudaMemcpy(array, dArray, size, cudaMemcpyDeviceToHost));\n checkCudaErrors(cudaFree(dArray));\n\n // Verify the return data\n // Both GPU and CPU use uint32_t * uint32_t, which has well-defined overflow behavior (modulo 2^32)\n //\n for (uint32_t i = 0; i < count; i += 1) {\n if (array[i] != i * i) {\n std::cout << \"element \" << i << \" expected \" << i * i << \" actual \" << array[i] << std::endl;\n return 1;\n }\n }\n delete[] array;\n\n return 0;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Sample Main\n//\n// The sample runs the test with manually configured launch and\n// automatically configured launch, and reports the occupancy and\n// performance.\n////////////////////////////////////////////////////////////////////////////////\nint main()\n{\n int status;\n\n std::cout << \"starting Simple Occupancy\" << std::endl << std::endl;\n\n std::cout << \"[ Manual configuration with \" << manualBlockSize << \" threads per block ]\" << std::endl;\n\n status = test(false);\n if (status) {\n std::cerr << \"Test failed\\n\" << std::endl;\n return -1;\n }\n\n std::cout << std::endl;\n\n std::cout << \"[ Automatic, occupancy-based configuration ]\" << std::endl;\n status = test(true);\n if (status) {\n std::cerr << \"Test failed\\n\" << std::endl;\n return -1;\n }\n\n std::cout << std::endl;\n std::cout << \"Test PASSED\\n\" << std::endl;\n\n return 0;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include // helper functions for CUDA error check\n#include \n\nconst int manualBlockSize = 32;\n\n////////////////////////////////////////////////////////////////////////////////\n// Test kernel\n//\n// This kernel squares each array element. Each thread addresses\n// himself with threadIdx and blockIdx, so that it can handle any\n// execution configuration, including anything the launch configurator\n// API suggests.\n////////////////////////////////////////////////////////////////////////////////\n__global__ void square(uint32_t *array, int arrayCount)\n{\n extern __shared__ int dynamicSmem[];\n int idx = threadIdx.x + blockIdx.x * blockDim.x;\n\n if (idx < arrayCount) {\n array[idx] *= array[idx];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Potential occupancy calculator\n//\n// The potential occupancy is calculated according to the kernel and\n// execution configuration the user desires. Occupancy is defined in\n// terms of active blocks per multiprocessor, and the user can convert\n// it to other metrics.\n//\n// This wrapper routine computes the occupancy of kernel, and reports\n// it in terms of active warps / maximum warps per SM.\n////////////////////////////////////////////////////////////////////////////////\nstatic double reportPotentialOccupancy(void *kernel, int blockSize, size_t dynamicSMem)\n{\n int device;\n hipDeviceProp_t prop;\n\n int numBlocks;\n int activeWarps;\n int maxWarps;\n\n double occupancy;\n\n checkCudaErrors(hipGetDevice(&device));\n checkCudaErrors(hipGetDeviceProperties(&prop, device));\n\n checkCudaErrors(hipOccupancyMaxActiveBlocksPerMultiprocessor(&numBlocks, kernel, blockSize, dynamicSMem));\n\n activeWarps = numBlocks * blockSize / prop.warpSize;\n maxWarps = prop.maxThreadsPerMultiProcessor / prop.warpSize;\n\n occupancy = (double)activeWarps / maxWarps;\n\n return occupancy;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Occupancy-based launch configurator\n//\n// The launch configurator, cudaOccupancyMaxPotentialBlockSize and\n// cudaOccupancyMaxPotentialBlockSizeVariableSMem, suggests a block\n// size that achieves the best theoretical occupancy. It also returns\n// the minimum number of blocks needed to achieve the occupancy on the\n// whole device.\n//\n// This launch configurator is purely occupancy-based. It doesn't\n// translate directly to performance, but the suggestion should\n// nevertheless be a good starting point for further optimizations.\n//\n// This function configures the launch based on the \"automatic\"\n// argument, records the runtime, and reports occupancy and runtime.\n////////////////////////////////////////////////////////////////////////////////\nstatic int launchConfig(uint32_t *array, int arrayCount, bool automatic)\n{\n int blockSize;\n int minGridSize;\n int gridSize;\n size_t dynamicSMemUsage = 0;\n\n hipEvent_t start;\n hipEvent_t end;\n\n float elapsedTime;\n\n double potentialOccupancy;\n\n checkCudaErrors(hipEventCreate(&start));\n checkCudaErrors(hipEventCreate(&end));\n\n if (automatic) {\n checkCudaErrors(\n hipOccupancyMaxPotentialBlockSize(&minGridSize, &blockSize, (void *)square, dynamicSMemUsage, arrayCount));\n\n std::cout << \"Suggested block size: \" << blockSize << std::endl\n << \"Minimum grid size for maximum occupancy: \" << minGridSize << std::endl;\n }\n else {\n // This block size is too small. Given limited number of\n // active blocks per multiprocessor, the number of active\n // threads will be limited, and thus unable to achieve maximum\n // occupancy.\n //\n blockSize = manualBlockSize;\n }\n\n // Round up\n //\n gridSize = (arrayCount + blockSize - 1) / blockSize;\n\n // Launch and profile\n //\n checkCudaErrors(hipEventRecord(start));\n square<<>>(array, arrayCount);\n checkCudaErrors(hipEventRecord(end));\n\n checkCudaErrors(hipDeviceSynchronize());\n\n // Calculate occupancy\n //\n potentialOccupancy = reportPotentialOccupancy((void *)square, blockSize, dynamicSMemUsage);\n\n std::cout << \"Potential occupancy: \" << potentialOccupancy * 100 << \"%\" << std::endl;\n\n // Report elapsed time\n //\n checkCudaErrors(hipEventElapsedTime(&elapsedTime, start, end));\n std::cout << \"Elapsed time: \" << elapsedTime << \"ms\" << std::endl;\n\n return 0;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// The test\n//\n// The test generates an array and squares it with a CUDA kernel, then\n// verifies the result.\n////////////////////////////////////////////////////////////////////////////////\nstatic int test(bool automaticLaunchConfig, const int count = 1000000)\n{\n uint32_t *array;\n uint32_t *dArray;\n int size = count * sizeof(uint32_t);\n\n array = new uint32_t[count];\n\n for (uint32_t i = 0; i < count; i += 1) {\n array[i] = i;\n }\n\n checkCudaErrors(hipMalloc(&dArray, size));\n checkCudaErrors(hipMemcpy(dArray, array, size, hipMemcpyHostToDevice));\n\n for (uint32_t i = 0; i < count; i += 1) {\n array[i] = 0;\n }\n\n launchConfig(dArray, count, automaticLaunchConfig);\n\n checkCudaErrors(hipMemcpy(array, dArray, size, hipMemcpyDeviceToHost));\n checkCudaErrors(hipFree(dArray));\n\n // Verify the return data\n // Both GPU and CPU use uint32_t * uint32_t, which has well-defined overflow behavior (modulo 2^32)\n //\n for (uint32_t i = 0; i < count; i += 1) {\n if (array[i] != i * i) {\n std::cout << \"element \" << i << \" expected \" << i * i << \" actual \" << array[i] << std::endl;\n return 1;\n }\n }\n delete[] array;\n\n return 0;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Sample Main\n//\n// The sample runs the test with manually configured launch and\n// automatically configured launch, and reports the occupancy and\n// performance.\n////////////////////////////////////////////////////////////////////////////////\nint main()\n{\n int status;\n\n std::cout << \"starting Simple Occupancy\" << std::endl << std::endl;\n\n std::cout << \"[ Manual configuration with \" << manualBlockSize << \" threads per block ]\" << std::endl;\n\n status = test(false);\n if (status) {\n std::cerr << \"Test failed\\n\" << std::endl;\n return -1;\n }\n\n std::cout << std::endl;\n\n std::cout << \"[ Automatic, occupancy-based configuration ]\" << std::endl;\n status = test(true);\n if (status) {\n std::cerr << \"Test failed\\n\" << std::endl;\n return -1;\n }\n\n std::cout << std::endl;\n std::cout << \"Test PASSED\\n\" << std::endl;\n\n return 0;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleOccupancy/simpleOccupancy.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* A program demonstrating trivial use of system-wide atomics on migratable\n * memory.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#define min(a, b) (a) < (b) ? (a) : (b)\n#define max(a, b) (a) > (b) ? (a) : (b)\n\n#define LOOP_NUM 50\n\n__global__ void atomicKernel(int *atom_arr)\n{\n unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int i = 0; i < LOOP_NUM; i++) {\n // Atomic addition\n atomicAdd_system(&atom_arr[0], 10);\n\n // Atomic exchange\n atomicExch_system(&atom_arr[1], tid);\n\n // Atomic maximum\n atomicMax_system(&atom_arr[2], tid);\n\n // Atomic minimum\n atomicMin_system(&atom_arr[3], tid);\n\n // Atomic increment (modulo 17+1)\n atomicInc_system((unsigned int *)&atom_arr[4], 17);\n\n // Atomic decrement\n atomicDec_system((unsigned int *)&atom_arr[5], 137);\n\n // Atomic compare-and-swap\n atomicCAS_system(&atom_arr[6], tid - 1, tid);\n\n // Bitwise atomic instructions\n\n // Atomic AND\n atomicAnd_system(&atom_arr[7], 2 * tid + 7);\n\n // Atomic OR\n atomicOr_system(&atom_arr[8], 1 << tid);\n\n // Atomic XOR\n atomicXor_system(&atom_arr[9], tid);\n }\n}\n\nvoid atomicKernel_CPU(int *atom_arr, int no_of_threads)\n{\n for (int i = no_of_threads; i < 2 * no_of_threads; i++) {\n for (int j = 0; j < LOOP_NUM; j++) {\n // Atomic addition\n __sync_fetch_and_add(&atom_arr[0], 10);\n\n // Atomic exchange\n __sync_lock_test_and_set(&atom_arr[1], i);\n\n // Atomic maximum\n int old, expected;\n do {\n expected = atom_arr[2];\n old = __sync_val_compare_and_swap(&atom_arr[2], expected, max(expected, i));\n } while (old != expected);\n\n // Atomic minimum\n do {\n expected = atom_arr[3];\n old = __sync_val_compare_and_swap(&atom_arr[3], expected, min(expected, i));\n } while (old != expected);\n\n // Atomic increment (modulo 17+1)\n int limit = 17;\n do {\n expected = atom_arr[4];\n old = __sync_val_compare_and_swap(&atom_arr[4], expected, (expected >= limit) ? 0 : expected + 1);\n } while (old != expected);\n\n // Atomic decrement\n limit = 137;\n do {\n expected = atom_arr[5];\n old = __sync_val_compare_and_swap(\n &atom_arr[5], expected, ((expected == 0) || (expected > limit)) ? limit : expected - 1);\n } while (old != expected);\n\n // Atomic compare-and-swap\n __sync_val_compare_and_swap(&atom_arr[6], i - 1, i);\n\n // Bitwise atomic instructions\n\n // Atomic AND\n __sync_fetch_and_and(&atom_arr[7], 2 * i + 7);\n\n // Atomic OR\n __sync_fetch_and_or(&atom_arr[8], 1 << i);\n\n // Atomic XOR\n // 11th element should be 0xff\n __sync_fetch_and_xor(&atom_arr[9], i);\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Compute reference data set\n//! Each element is multiplied with the number of threads / array length\n//! @param reference reference data, computed but preallocated\n//! @param idata input data as provided to device\n//! @param len number of elements in reference / idata\n////////////////////////////////////////////////////////////////////////////////\nint verify(int *testData, const int len)\n{\n int val = 0;\n\n for (int i = 0; i < len * LOOP_NUM; ++i) {\n val += 10;\n }\n\n if (val != testData[0]) {\n printf(\"atomicAdd failed val = %d testData = %d\\n\", val, testData[0]);\n return false;\n }\n\n val = 0;\n\n bool found = false;\n\n for (int i = 0; i < len; ++i) {\n // second element should be a member of [0, len)\n if (i == testData[1]) {\n found = true;\n break;\n }\n }\n\n if (!found) {\n printf(\"atomicExch failed\\n\");\n return false;\n }\n\n val = -(1 << 8);\n\n for (int i = 0; i < len; ++i) {\n // third element should be len-1\n val = max(val, i);\n }\n\n if (val != testData[2]) {\n printf(\"atomicMax failed\\n\");\n return false;\n }\n\n val = 1 << 8;\n\n for (int i = 0; i < len; ++i) {\n val = min(val, i);\n }\n\n if (val != testData[3]) {\n printf(\"atomicMin failed\\n\");\n return false;\n }\n\n int limit = 17;\n val = 0;\n\n for (int i = 0; i < len * LOOP_NUM; ++i) {\n val = (val >= limit) ? 0 : val + 1;\n }\n\n if (val != testData[4]) {\n printf(\"atomicInc failed\\n\");\n return false;\n }\n\n limit = 137;\n val = 0;\n\n for (int i = 0; i < len * LOOP_NUM; ++i) {\n val = ((val == 0) || (val > limit)) ? limit : val - 1;\n }\n\n if (val != testData[5]) {\n printf(\"atomicDec failed\\n\");\n return false;\n }\n\n found = false;\n\n for (int i = 0; i < len; ++i) {\n // seventh element should be a member of [0, len)\n if (i == testData[6]) {\n found = true;\n break;\n }\n }\n\n if (!found) {\n printf(\"atomicCAS failed\\n\");\n return false;\n }\n\n val = 0xff;\n\n for (int i = 0; i < len; ++i) {\n // 8th element should be 1\n val &= (2 * i + 7);\n }\n\n if (val != testData[7]) {\n printf(\"atomicAnd failed\\n\");\n return false;\n }\n\n val = 0;\n\n for (int i = 0; i < len; ++i) {\n // 9th element should be 0xff\n val |= (1 << i);\n }\n\n if (val != testData[8]) {\n printf(\"atomicOr failed\\n\");\n return false;\n }\n\n val = 0xff;\n\n for (int i = 0; i < len; ++i) {\n // 11th element should be 0xff\n val ^= i;\n }\n\n if (val != testData[9]) {\n printf(\"atomicXor failed\\n\");\n return false;\n }\n\n return true;\n}\n\nint main(int argc, char **argv)\n{\n // set device\n cudaDeviceProp device_prop;\n int dev_id = findCudaDevice(argc, (const char **)argv);\n checkCudaErrors(cudaGetDeviceProperties(&device_prop, dev_id));\n\n if (!device_prop.managedMemory) {\n // This samples requires being run on a device that supports Unified Memory\n fprintf(stderr, \"Unified Memory not supported on this device\\n\");\n exit(EXIT_WAIVED);\n }\n\n int computeMode;\n checkCudaErrors(cudaDeviceGetAttribute(&computeMode, cudaDevAttrComputeMode, dev_id));\n if (computeMode == cudaComputeModeProhibited) {\n // This sample requires being run with a default or process exclusive mode\n fprintf(stderr,\n \"This sample requires a device in either default or process \"\n \"exclusive mode\\n\");\n exit(EXIT_WAIVED);\n }\n\n int concurrentManagedAccess;\n checkCudaErrors(cudaDeviceGetAttribute(&concurrentManagedAccess, cudaDevAttrConcurrentManagedAccess, dev_id));\n\n if (device_prop.major < 6) {\n printf(\"%s: requires a minimum CUDA compute 6.0 capability, waiving \"\n \"testing.\\n\",\n argv[0]);\n exit(EXIT_WAIVED);\n }\n\n unsigned int numThreads = 256;\n unsigned int numBlocks = 64;\n unsigned int numData = 10;\n\n int *atom_arr;\n\n if (device_prop.pageableMemoryAccess) {\n printf(\"CAN access pageable memory\\n\");\n atom_arr = (int *)malloc(sizeof(int) * numData);\n }\n else {\n printf(\"CANNOT access pageable memory\\n\");\n checkCudaErrors(cudaMallocManaged(&atom_arr, sizeof(int) * numData));\n }\n\n for (unsigned int i = 0; i < numData; i++)\n atom_arr[i] = 0;\n\n // To make the AND and XOR tests generate something other than 0...\n atom_arr[7] = atom_arr[9] = 0xff;\n\n atomicKernel<<>>(atom_arr);\n\n // Do device synchronize for devices that don't support concurrent managed access such as WSL.\n if (!concurrentManagedAccess) {\n checkCudaErrors(cudaDeviceSynchronize());\n }\n\n atomicKernel_CPU(atom_arr, numBlocks * numThreads);\n\n checkCudaErrors(cudaDeviceSynchronize());\n\n // Compute & verify reference solution\n int testResult = verify(atom_arr, 2 * numThreads * numBlocks);\n\n if (device_prop.pageableMemoryAccess) {\n free(atom_arr);\n }\n else {\n cudaFree(atom_arr);\n }\n\n printf(\"systemWideAtomics completed, returned %s \\n\", testResult ? \"OK\" : \"ERROR!\");\n exit(testResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* A program demonstrating trivial use of system-wide atomics on migratable\n * memory.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#define min(a, b) (a) < (b) ? (a) : (b)\n#define max(a, b) (a) > (b) ? (a) : (b)\n\n#define LOOP_NUM 50\n\n__global__ void atomicKernel(int *atom_arr)\n{\n unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int i = 0; i < LOOP_NUM; i++) {\n // Atomic addition\n atomicAdd_system(&atom_arr[0], 10);\n\n // Atomic exchange\n atomicExch_system(&atom_arr[1], tid);\n\n // Atomic maximum\n atomicMax_system(&atom_arr[2], tid);\n\n // Atomic minimum\n atomicMin_system(&atom_arr[3], tid);\n\n // Atomic increment (modulo 17+1)\n atomicInc_system((unsigned int *)&atom_arr[4], 17);\n\n // Atomic decrement\n atomicDec_system((unsigned int *)&atom_arr[5], 137);\n\n // Atomic compare-and-swap\n atomicCAS_system(&atom_arr[6], tid - 1, tid);\n\n // Bitwise atomic instructions\n\n // Atomic AND\n atomicAnd_system(&atom_arr[7], 2 * tid + 7);\n\n // Atomic OR\n atomicOr_system(&atom_arr[8], 1 << tid);\n\n // Atomic XOR\n atomicXor_system(&atom_arr[9], tid);\n }\n}\n\nvoid atomicKernel_CPU(int *atom_arr, int no_of_threads)\n{\n for (int i = no_of_threads; i < 2 * no_of_threads; i++) {\n for (int j = 0; j < LOOP_NUM; j++) {\n // Atomic addition\n __sync_fetch_and_add(&atom_arr[0], 10);\n\n // Atomic exchange\n __sync_lock_test_and_set(&atom_arr[1], i);\n\n // Atomic maximum\n int old, expected;\n do {\n expected = atom_arr[2];\n old = __sync_val_compare_and_swap(&atom_arr[2], expected, max(expected, i));\n } while (old != expected);\n\n // Atomic minimum\n do {\n expected = atom_arr[3];\n old = __sync_val_compare_and_swap(&atom_arr[3], expected, min(expected, i));\n } while (old != expected);\n\n // Atomic increment (modulo 17+1)\n int limit = 17;\n do {\n expected = atom_arr[4];\n old = __sync_val_compare_and_swap(&atom_arr[4], expected, (expected >= limit) ? 0 : expected + 1);\n } while (old != expected);\n\n // Atomic decrement\n limit = 137;\n do {\n expected = atom_arr[5];\n old = __sync_val_compare_and_swap(\n &atom_arr[5], expected, ((expected == 0) || (expected > limit)) ? limit : expected - 1);\n } while (old != expected);\n\n // Atomic compare-and-swap\n __sync_val_compare_and_swap(&atom_arr[6], i - 1, i);\n\n // Bitwise atomic instructions\n\n // Atomic AND\n __sync_fetch_and_and(&atom_arr[7], 2 * i + 7);\n\n // Atomic OR\n __sync_fetch_and_or(&atom_arr[8], 1 << i);\n\n // Atomic XOR\n // 11th element should be 0xff\n __sync_fetch_and_xor(&atom_arr[9], i);\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Compute reference data set\n//! Each element is multiplied with the number of threads / array length\n//! @param reference reference data, computed but preallocated\n//! @param idata input data as provided to device\n//! @param len number of elements in reference / idata\n////////////////////////////////////////////////////////////////////////////////\nint verify(int *testData, const int len)\n{\n int val = 0;\n\n for (int i = 0; i < len * LOOP_NUM; ++i) {\n val += 10;\n }\n\n if (val != testData[0]) {\n printf(\"atomicAdd failed val = %d testData = %d\\n\", val, testData[0]);\n return false;\n }\n\n val = 0;\n\n bool found = false;\n\n for (int i = 0; i < len; ++i) {\n // second element should be a member of [0, len)\n if (i == testData[1]) {\n found = true;\n break;\n }\n }\n\n if (!found) {\n printf(\"atomicExch failed\\n\");\n return false;\n }\n\n val = -(1 << 8);\n\n for (int i = 0; i < len; ++i) {\n // third element should be len-1\n val = max(val, i);\n }\n\n if (val != testData[2]) {\n printf(\"atomicMax failed\\n\");\n return false;\n }\n\n val = 1 << 8;\n\n for (int i = 0; i < len; ++i) {\n val = min(val, i);\n }\n\n if (val != testData[3]) {\n printf(\"atomicMin failed\\n\");\n return false;\n }\n\n int limit = 17;\n val = 0;\n\n for (int i = 0; i < len * LOOP_NUM; ++i) {\n val = (val >= limit) ? 0 : val + 1;\n }\n\n if (val != testData[4]) {\n printf(\"atomicInc failed\\n\");\n return false;\n }\n\n limit = 137;\n val = 0;\n\n for (int i = 0; i < len * LOOP_NUM; ++i) {\n val = ((val == 0) || (val > limit)) ? limit : val - 1;\n }\n\n if (val != testData[5]) {\n printf(\"atomicDec failed\\n\");\n return false;\n }\n\n found = false;\n\n for (int i = 0; i < len; ++i) {\n // seventh element should be a member of [0, len)\n if (i == testData[6]) {\n found = true;\n break;\n }\n }\n\n if (!found) {\n printf(\"atomicCAS failed\\n\");\n return false;\n }\n\n val = 0xff;\n\n for (int i = 0; i < len; ++i) {\n // 8th element should be 1\n val &= (2 * i + 7);\n }\n\n if (val != testData[7]) {\n printf(\"atomicAnd failed\\n\");\n return false;\n }\n\n val = 0;\n\n for (int i = 0; i < len; ++i) {\n // 9th element should be 0xff\n val |= (1 << i);\n }\n\n if (val != testData[8]) {\n printf(\"atomicOr failed\\n\");\n return false;\n }\n\n val = 0xff;\n\n for (int i = 0; i < len; ++i) {\n // 11th element should be 0xff\n val ^= i;\n }\n\n if (val != testData[9]) {\n printf(\"atomicXor failed\\n\");\n return false;\n }\n\n return true;\n}\n\nint main(int argc, char **argv)\n{\n // set device\n hipDeviceProp_t device_prop;\n int dev_id = findCudaDevice(argc, (const char **)argv);\n checkCudaErrors(hipGetDeviceProperties(&device_prop, dev_id));\n\n if (!device_prop.managedMemory) {\n // This samples requires being run on a device that supports Unified Memory\n fprintf(stderr, \"Unified Memory not supported on this device\\n\");\n exit(EXIT_WAIVED);\n }\n\n int computeMode;\n checkCudaErrors(hipDeviceGetAttribute(&computeMode, hipDeviceAttributeComputeMode, dev_id));\n if (computeMode == hipComputeModeProhibited) {\n // This sample requires being run with a default or process exclusive mode\n fprintf(stderr,\n \"This sample requires a device in either default or process \"\n \"exclusive mode\\n\");\n exit(EXIT_WAIVED);\n }\n\n int concurrentManagedAccess;\n checkCudaErrors(hipDeviceGetAttribute(&concurrentManagedAccess, hipDeviceAttributeConcurrentManagedAccess, dev_id));\n\n if (device_prop.major < 6) {\n printf(\"%s: requires a minimum CUDA compute 6.0 capability, waiving \"\n \"testing.\\n\",\n argv[0]);\n exit(EXIT_WAIVED);\n }\n\n unsigned int numThreads = 256;\n unsigned int numBlocks = 64;\n unsigned int numData = 10;\n\n int *atom_arr;\n\n if (device_prop.pageableMemoryAccess) {\n printf(\"CAN access pageable memory\\n\");\n atom_arr = (int *)malloc(sizeof(int) * numData);\n }\n else {\n printf(\"CANNOT access pageable memory\\n\");\n checkCudaErrors(hipMallocManaged(&atom_arr, sizeof(int) * numData));\n }\n\n for (unsigned int i = 0; i < numData; i++)\n atom_arr[i] = 0;\n\n // To make the AND and XOR tests generate something other than 0...\n atom_arr[7] = atom_arr[9] = 0xff;\n\n atomicKernel<<>>(atom_arr);\n\n // Do device synchronize for devices that don't support concurrent managed access such as WSL.\n if (!concurrentManagedAccess) {\n checkCudaErrors(hipDeviceSynchronize());\n }\n\n atomicKernel_CPU(atom_arr, numBlocks * numThreads);\n\n checkCudaErrors(hipDeviceSynchronize());\n\n // Compute & verify reference solution\n int testResult = verify(atom_arr, 2 * numThreads * numBlocks);\n\n if (device_prop.pageableMemoryAccess) {\n free(atom_arr);\n }\n else {\n hipFree(atom_arr);\n }\n\n printf(\"systemWideAtomics completed, returned %s \\n\", testResult ? \"OK\" : \"ERROR!\");\n exit(testResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/systemWideAtomics/systemWideAtomics.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/**\n * Matrix multiplication: C = A * B.\n * Host code.\n *\n * This sample implements matrix multiplication as described in Chapter 3\n * of the programming guide.\n * It has been written for clarity of exposition to illustrate various CUDA\n * programming principles, not with the goal of providing the most\n * performant generic kernel for matrix multiplication.\n *\n * See also:\n * V. Volkov and J. Demmel, \"Benchmarking GPUs to tune dense linear algebra,\"\n * in Proc. 2008 ACM/IEEE Conf. on Supercomputing (SC '08),\n * Piscataway, NJ: IEEE Press, 2008, pp. Art. 31:1-11.\n */\n\n/**\n * Matrix multiplication (CUDA Kernel) on the device: C = A * B\n * wA is A's width and wB is B's width\n */\n\n#include \n\ntemplate __device__ void matrixMulCUDA(float *C, float *A, float *B, int wA, int wB)\n{\n // Handle to thread block group\n cooperative_groups::thread_block cta = cooperative_groups::this_thread_block();\n // Block index\n int bx = blockIdx.x;\n int by = blockIdx.y;\n\n // Thread index\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = wA * BLOCK_SIZE * by;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * bx;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n // Csub is used to store the element of the block sub-matrix\n // that is computed by the thread\n float Csub = 0;\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n // Load the matrices from device memory\n // to shared memory; each thread loads\n // one element of each matrix\n As[ty][tx] = A[a + wA * ty + tx];\n Bs[ty][tx] = B[b + wB * ty + tx];\n\n // Synchronize to make sure the matrices are loaded\n cooperative_groups::sync(cta);\n\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[ty][k] * Bs[k][tx];\n }\n\n // Synchronize to make sure that the preceding\n // computation is done before loading two new\n // sub-matrices of A and B in the next iteration\n cooperative_groups::sync(cta);\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes one element\n int c = wB * BLOCK_SIZE * by + BLOCK_SIZE * bx;\n C[c + wB * ty + tx] = Csub;\n}\n\nextern \"C\" __global__ void matrixMulCUDA_block16(float *C, float *A, float *B, int wA, int wB)\n{\n matrixMulCUDA<16>(C, A, B, wA, wB);\n}\n\nextern \"C\" __global__ void matrixMulCUDA_block32(float *C, float *A, float *B, int wA, int wB)\n{\n matrixMulCUDA<32>(C, A, B, wA, wB);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/**\n * Matrix multiplication: C = A * B.\n * Host code.\n *\n * This sample implements matrix multiplication as described in Chapter 3\n * of the programming guide.\n * It has been written for clarity of exposition to illustrate various CUDA\n * programming principles, not with the goal of providing the most\n * performant generic kernel for matrix multiplication.\n *\n * See also:\n * V. Volkov and J. Demmel, \"Benchmarking GPUs to tune dense linear algebra,\"\n * in Proc. 2008 ACM/IEEE Conf. on Supercomputing (SC '08),\n * Piscataway, NJ: IEEE Press, 2008, pp. Art. 31:1-11.\n */\n\n/**\n * Matrix multiplication (CUDA Kernel) on the device: C = A * B\n * wA is A's width and wB is B's width\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n\ntemplate __device__ void matrixMulCUDA(float *C, float *A, float *B, int wA, int wB)\n{\n // Handle to thread block group\n cooperative_groups::thread_block cta = cooperative_groups::this_thread_block();\n // Block index\n int bx = blockIdx.x;\n int by = blockIdx.y;\n\n // Thread index\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = wA * BLOCK_SIZE * by;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * bx;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n // Csub is used to store the element of the block sub-matrix\n // that is computed by the thread\n float Csub = 0;\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n // Load the matrices from device memory\n // to shared memory; each thread loads\n // one element of each matrix\n As[ty][tx] = A[a + wA * ty + tx];\n Bs[ty][tx] = B[b + wB * ty + tx];\n\n // Synchronize to make sure the matrices are loaded\n cooperative_groups::sync(cta);\n\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[ty][k] * Bs[k][tx];\n }\n\n // Synchronize to make sure that the preceding\n // computation is done before loading two new\n // sub-matrices of A and B in the next iteration\n cooperative_groups::sync(cta);\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes one element\n int c = wB * BLOCK_SIZE * by + BLOCK_SIZE * bx;\n C[c + wB * ty + tx] = Csub;\n}\n\nextern \"C\" __global__ void matrixMulCUDA_block16(float *C, float *A, float *B, int wA, int wB)\n{\n matrixMulCUDA<16>(C, A, B, wA, wB);\n}\n\nextern \"C\" __global__ void matrixMulCUDA_block32(float *C, float *A, float *B, int wA, int wB)\n{\n matrixMulCUDA<32>(C, A, B, wA, wB);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/matrixMul_nvrtc/matrixMul_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This example shows how to use the clock function to measure the performance\n * of block of threads of a kernel accurately. Blocks are executed in parallel\n * and out of order. Since there's no synchronization mechanism between blocks,\n * we measure the clock once for each block. The clock samples are written to\n * device memory.\n */\n\n// System includes\n#include \n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// This kernel computes a standard parallel reduction and evaluates the\n// time it takes to do that for each block. The timing results are stored\n// in device memory.\n__global__ static void timedReduction(const float *input, float *output, clock_t *timer)\n{\n // __shared__ float shared[2 * blockDim.x];\n extern __shared__ float shared[];\n\n const int tid = threadIdx.x;\n const int bid = blockIdx.x;\n\n if (tid == 0)\n timer[bid] = clock();\n\n // Copy input.\n shared[tid] = input[tid];\n shared[tid + blockDim.x] = input[tid + blockDim.x];\n\n // Perform reduction to find minimum.\n for (int d = blockDim.x; d > 0; d /= 2) {\n __syncthreads();\n\n if (tid < d) {\n float f0 = shared[tid];\n float f1 = shared[tid + d];\n\n if (f1 < f0) {\n shared[tid] = f1;\n }\n }\n }\n\n // Write result.\n if (tid == 0)\n output[bid] = shared[0];\n\n __syncthreads();\n\n if (tid == 0)\n timer[bid + gridDim.x] = clock();\n}\n\n#define NUM_BLOCKS 64\n#define NUM_THREADS 256\n\n// It's interesting to change the number of blocks and the number of threads to\n// understand how to keep the hardware busy.\n//\n// Here are some numbers I get on my G80:\n// blocks - clocks\n// 1 - 3096\n// 8 - 3232\n// 16 - 3364\n// 32 - 4615\n// 64 - 9981\n//\n// With less than 16 blocks some of the multiprocessors of the device are idle.\n// With more than 16 you are using all the multiprocessors, but there's only one\n// block per multiprocessor and that doesn't allow you to hide the latency of\n// the memory. With more than 32 the speed scales linearly.\n\n// Start the main CUDA Sample here\nint main(int argc, char **argv)\n{\n printf(\"CUDA Clock sample\\n\");\n\n // This will pick the best possible CUDA capable device\n int dev = findCudaDevice(argc, (const char **)argv);\n\n float *dinput = NULL;\n float *doutput = NULL;\n clock_t *dtimer = NULL;\n\n clock_t timer[NUM_BLOCKS * 2];\n float input[NUM_THREADS * 2];\n\n for (int i = 0; i < NUM_THREADS * 2; i++) {\n input[i] = (float)i;\n }\n\n checkCudaErrors(cudaMalloc((void **)&dinput, sizeof(float) * NUM_THREADS * 2));\n checkCudaErrors(cudaMalloc((void **)&doutput, sizeof(float) * NUM_BLOCKS));\n checkCudaErrors(cudaMalloc((void **)&dtimer, sizeof(clock_t) * NUM_BLOCKS * 2));\n\n checkCudaErrors(cudaMemcpy(dinput, input, sizeof(float) * NUM_THREADS * 2, cudaMemcpyHostToDevice));\n\n timedReduction<<>>(dinput, doutput, dtimer);\n\n checkCudaErrors(cudaMemcpy(timer, dtimer, sizeof(clock_t) * NUM_BLOCKS * 2, cudaMemcpyDeviceToHost));\n\n checkCudaErrors(cudaFree(dinput));\n checkCudaErrors(cudaFree(doutput));\n checkCudaErrors(cudaFree(dtimer));\n\n long double avgElapsedClocks = 0;\n\n for (int i = 0; i < NUM_BLOCKS; i++) {\n avgElapsedClocks += (long double)(timer[i + NUM_BLOCKS] - timer[i]);\n }\n\n avgElapsedClocks = avgElapsedClocks / NUM_BLOCKS;\n printf(\"Average clocks/block = %Lf\\n\", avgElapsedClocks);\n\n return EXIT_SUCCESS;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This example shows how to use the clock function to measure the performance\n * of block of threads of a kernel accurately. Blocks are executed in parallel\n * and out of order. Since there's no synchronization mechanism between blocks,\n * we measure the clock once for each block. The clock samples are written to\n * device memory.\n */\n\n// System includes\n#include \n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// This kernel computes a standard parallel reduction and evaluates the\n// time it takes to do that for each block. The timing results are stored\n// in device memory.\n__global__ static void timedReduction(const float *input, float *output, clock_t *timer)\n{\n // __shared__ float shared[2 * blockDim.x];\n extern __shared__ float shared[];\n\n const int tid = threadIdx.x;\n const int bid = blockIdx.x;\n\n if (tid == 0)\n timer[bid] = clock();\n\n // Copy input.\n shared[tid] = input[tid];\n shared[tid + blockDim.x] = input[tid + blockDim.x];\n\n // Perform reduction to find minimum.\n for (int d = blockDim.x; d > 0; d /= 2) {\n __syncthreads();\n\n if (tid < d) {\n float f0 = shared[tid];\n float f1 = shared[tid + d];\n\n if (f1 < f0) {\n shared[tid] = f1;\n }\n }\n }\n\n // Write result.\n if (tid == 0)\n output[bid] = shared[0];\n\n __syncthreads();\n\n if (tid == 0)\n timer[bid + gridDim.x] = clock();\n}\n\n#define NUM_BLOCKS 64\n#define NUM_THREADS 256\n\n// It's interesting to change the number of blocks and the number of threads to\n// understand how to keep the hardware busy.\n//\n// Here are some numbers I get on my G80:\n// blocks - clocks\n// 1 - 3096\n// 8 - 3232\n// 16 - 3364\n// 32 - 4615\n// 64 - 9981\n//\n// With less than 16 blocks some of the multiprocessors of the device are idle.\n// With more than 16 you are using all the multiprocessors, but there's only one\n// block per multiprocessor and that doesn't allow you to hide the latency of\n// the memory. With more than 32 the speed scales linearly.\n\n// Start the main CUDA Sample here\nint main(int argc, char **argv)\n{\n printf(\"CUDA Clock sample\\n\");\n\n // This will pick the best possible CUDA capable device\n int dev = findCudaDevice(argc, (const char **)argv);\n\n float *dinput = NULL;\n float *doutput = NULL;\n clock_t *dtimer = NULL;\n\n clock_t timer[NUM_BLOCKS * 2];\n float input[NUM_THREADS * 2];\n\n for (int i = 0; i < NUM_THREADS * 2; i++) {\n input[i] = (float)i;\n }\n\n checkCudaErrors(hipMalloc((void **)&dinput, sizeof(float) * NUM_THREADS * 2));\n checkCudaErrors(hipMalloc((void **)&doutput, sizeof(float) * NUM_BLOCKS));\n checkCudaErrors(hipMalloc((void **)&dtimer, sizeof(clock_t) * NUM_BLOCKS * 2));\n\n checkCudaErrors(hipMemcpy(dinput, input, sizeof(float) * NUM_THREADS * 2, hipMemcpyHostToDevice));\n\n timedReduction<<>>(dinput, doutput, dtimer);\n\n checkCudaErrors(hipMemcpy(timer, dtimer, sizeof(clock_t) * NUM_BLOCKS * 2, hipMemcpyDeviceToHost));\n\n checkCudaErrors(hipFree(dinput));\n checkCudaErrors(hipFree(doutput));\n checkCudaErrors(hipFree(dtimer));\n\n long double avgElapsedClocks = 0;\n\n for (int i = 0; i < NUM_BLOCKS; i++) {\n avgElapsedClocks += (long double)(timer[i + NUM_BLOCKS] - timer[i]);\n }\n\n avgElapsedClocks = avgElapsedClocks / NUM_BLOCKS;\n printf(\"Average clocks/block = %Lf\\n\", avgElapsedClocks);\n\n return EXIT_SUCCESS;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/clock/clock.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#ifdef _WIN32\n#define WINDOWS_LEAN_AND_MEAN\n#define NOMINMAX\n#include \n#else\n#include \n#endif\n\n// Includes, system\n#include \n#include \n\n// Includes CUDA\n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n\nconst char *sampleName = \"simpleAssert\";\n\n////////////////////////////////////////////////////////////////////////////////\n// Auto-Verification Code\nbool testResult = true;\n\n////////////////////////////////////////////////////////////////////////////////\n// Kernels\n////////////////////////////////////////////////////////////////////////////////\n//! Tests assert function.\n//! Thread whose id > N will print assertion failed error message.\n////////////////////////////////////////////////////////////////////////////////\n__global__ void testKernel(int N)\n{\n int gtid = blockIdx.x * blockDim.x + threadIdx.x;\n assert(gtid < N);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Declaration, forward\nvoid runTest(int argc, char **argv);\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"%s starting...\\n\", sampleName);\n\n runTest(argc, argv);\n\n printf(\"%s completed, returned %s\\n\", sampleName, testResult ? \"OK\" : \"ERROR!\");\n exit(testResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\nvoid runTest(int argc, char **argv)\n{\n int Nblocks = 2;\n int Nthreads = 32;\n cudaError_t error;\n\n#ifndef _WIN32\n utsname OS_System_Type;\n uname(&OS_System_Type);\n\n printf(\"OS_System_Type.release = %s\\n\", OS_System_Type.release);\n\n if (!strcasecmp(OS_System_Type.sysname, \"Darwin\")) {\n printf(\"simpleAssert is not current supported on Mac OSX\\n\\n\");\n exit(EXIT_SUCCESS);\n }\n else {\n printf(\"OS Info: <%s>\\n\\n\", OS_System_Type.version);\n }\n\n#endif\n\n // This will pick the best possible CUDA capable device\n findCudaDevice(argc, (const char **)argv);\n\n // Kernel configuration, where a one-dimensional\n // grid and one-dimensional blocks are configured.\n dim3 dimGrid(Nblocks);\n dim3 dimBlock(Nthreads);\n\n printf(\"Launch kernel to generate assertion failures\\n\");\n testKernel<<>>(60);\n\n // Synchronize (flushes assert output).\n printf(\"\\n-- Begin assert output\\n\\n\");\n error = cudaDeviceSynchronize();\n printf(\"\\n-- End assert output\\n\\n\");\n\n // Check for errors and failed asserts in asynchronous kernel launch.\n if (error == cudaErrorAssert) {\n printf(\"Device assert failed as expected, \"\n \"CUDA error message is: %s\\n\\n\",\n cudaGetErrorString(error));\n }\n\n testResult = error == cudaErrorAssert;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#ifdef _WIN32\n#define WINDOWS_LEAN_AND_MEAN\n#define NOMINMAX\n#include \n#else\n#include \n#endif\n\n// Includes, system\n#include \n#include \n\n// Includes CUDA\n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n\nconst char *sampleName = \"simpleAssert\";\n\n////////////////////////////////////////////////////////////////////////////////\n// Auto-Verification Code\nbool testResult = true;\n\n////////////////////////////////////////////////////////////////////////////////\n// Kernels\n////////////////////////////////////////////////////////////////////////////////\n//! Tests assert function.\n//! Thread whose id > N will print assertion failed error message.\n////////////////////////////////////////////////////////////////////////////////\n__global__ void testKernel(int N)\n{\n int gtid = blockIdx.x * blockDim.x + threadIdx.x;\n assert(gtid < N);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Declaration, forward\nvoid runTest(int argc, char **argv);\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"%s starting...\\n\", sampleName);\n\n runTest(argc, argv);\n\n printf(\"%s completed, returned %s\\n\", sampleName, testResult ? \"OK\" : \"ERROR!\");\n exit(testResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\nvoid runTest(int argc, char **argv)\n{\n int Nblocks = 2;\n int Nthreads = 32;\n hipError_t error;\n\n#ifndef _WIN32\n utsname OS_System_Type;\n uname(&OS_System_Type);\n\n printf(\"OS_System_Type.release = %s\\n\", OS_System_Type.release);\n\n if (!strcasecmp(OS_System_Type.sysname, \"Darwin\")) {\n printf(\"simpleAssert is not current supported on Mac OSX\\n\\n\");\n exit(EXIT_SUCCESS);\n }\n else {\n printf(\"OS Info: <%s>\\n\\n\", OS_System_Type.version);\n }\n\n#endif\n\n // This will pick the best possible CUDA capable device\n findCudaDevice(argc, (const char **)argv);\n\n // Kernel configuration, where a one-dimensional\n // grid and one-dimensional blocks are configured.\n dim3 dimGrid(Nblocks);\n dim3 dimBlock(Nthreads);\n\n printf(\"Launch kernel to generate assertion failures\\n\");\n testKernel<<>>(60);\n\n // Synchronize (flushes assert output).\n printf(\"\\n-- Begin assert output\\n\\n\");\n error = hipDeviceSynchronize();\n printf(\"\\n-- End assert output\\n\\n\");\n\n // Check for errors and failed asserts in asynchronous kernel launch.\n if (error == hipErrorAssert) {\n printf(\"Device assert failed as expected, \"\n \"CUDA error message is: %s\\n\\n\",\n hipGetErrorString(error));\n }\n\n testResult = error == hipErrorAssert;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleAssert/simpleAssert.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample implements multi-threaded heterogeneous computing workloads with\n * the new CPU callbacks for CUDA streams and events introduced with CUDA 5.0.\n * Together with the thread safety of the CUDA API implementing heterogeneous\n * workloads that float between CPU threads and GPUs has become simple and\n * efficient.\n *\n * The workloads in the sample follow the form CPU preprocess -> GPU process ->\n * CPU postprocess.\n * Each CPU processing step is handled by its own dedicated thread. GPU\n * workloads are sent to all available GPUs in the system.\n *\n */\n\n// System includes\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#include \"multithreading.h\"\n\nconst int N_workloads = 8;\nconst int N_elements_per_workload = 100000;\n\nCUTBarrier thread_barrier;\n\nvoid CUDART_CB myStreamCallback(cudaStream_t event, cudaError_t status, void *data);\n\nstruct heterogeneous_workload\n{\n int id;\n int cudaDeviceID;\n\n int *h_data;\n int *d_data;\n cudaStream_t stream;\n\n bool success;\n};\n\n__global__ void incKernel(int *data, int N)\n{\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < N)\n data[i]++;\n}\n\nCUT_THREADPROC launch(void *void_arg)\n{\n heterogeneous_workload *workload = (heterogeneous_workload *)void_arg;\n\n // Select GPU for this CPU thread\n checkCudaErrors(cudaSetDevice(workload->cudaDeviceID));\n\n // Allocate Resources\n checkCudaErrors(cudaStreamCreate(&workload->stream));\n checkCudaErrors(cudaMalloc(&workload->d_data, N_elements_per_workload * sizeof(int)));\n checkCudaErrors(cudaHostAlloc(&workload->h_data, N_elements_per_workload * sizeof(int), cudaHostAllocPortable));\n\n // CPU thread generates data\n for (int i = 0; i < N_elements_per_workload; ++i) {\n workload->h_data[i] = workload->id + i;\n }\n\n // Schedule work for GPU in CUDA stream without blocking the CPU thread\n // Note: Dedicated streams enable concurrent execution of workloads on the GPU\n dim3 block(512);\n dim3 grid((N_elements_per_workload + block.x - 1) / block.x);\n\n checkCudaErrors(cudaMemcpyAsync(workload->d_data,\n workload->h_data,\n N_elements_per_workload * sizeof(int),\n cudaMemcpyHostToDevice,\n workload->stream));\n incKernel<<stream>>>(workload->d_data, N_elements_per_workload);\n checkCudaErrors(cudaMemcpyAsync(workload->h_data,\n workload->d_data,\n N_elements_per_workload * sizeof(int),\n cudaMemcpyDeviceToHost,\n workload->stream));\n\n // New in CUDA 5.0: Add a CPU callback which is called once all currently\n // pending operations in the CUDA stream have finished\n checkCudaErrors(cudaStreamAddCallback(workload->stream, myStreamCallback, workload, 0));\n\n CUT_THREADEND;\n // CPU thread end of life, GPU continues to process data...\n}\n\nCUT_THREADPROC postprocess(void *void_arg)\n{\n heterogeneous_workload *workload = (heterogeneous_workload *)void_arg;\n // ... GPU is done with processing, continue on new CPU thread...\n\n // Select GPU for this CPU thread\n checkCudaErrors(cudaSetDevice(workload->cudaDeviceID));\n\n // CPU thread consumes results from GPU\n workload->success = true;\n\n for (int i = 0; i < N_workloads; ++i) {\n workload->success &= workload->h_data[i] == i + workload->id + 1;\n }\n\n // Free Resources\n checkCudaErrors(cudaFree(workload->d_data));\n checkCudaErrors(cudaFreeHost(workload->h_data));\n checkCudaErrors(cudaStreamDestroy(workload->stream));\n\n // Signal the end of the heterogeneous workload to main thread\n cutIncrementBarrier(&thread_barrier);\n\n CUT_THREADEND;\n}\n\nvoid CUDART_CB myStreamCallback(cudaStream_t stream, cudaError_t status, void *data)\n{\n // Check status of GPU after stream operations are done\n checkCudaErrors(status);\n\n // Spawn new CPU worker thread and continue processing on the CPU\n cutStartThread(postprocess, data);\n}\n\nint main(int argc, char **argv)\n{\n int N_gpus, max_gpus = 0;\n int gpuInfo[32]; // assume a maximum of 32 GPUs in a system configuration\n\n printf(\"Starting simpleCallback\\n\");\n\n checkCudaErrors(cudaGetDeviceCount(&N_gpus));\n printf(\"Found %d CUDA capable GPUs\\n\", N_gpus);\n\n if (N_gpus > 32) {\n printf(\"simpleCallback only supports 32 GPU(s)\\n\");\n }\n\n for (int devid = 0; devid < N_gpus; devid++) {\n int SMversion;\n cudaDeviceProp deviceProp;\n cudaSetDevice(devid);\n cudaGetDeviceProperties(&deviceProp, devid);\n SMversion = deviceProp.major << 4 + deviceProp.minor;\n printf(\"GPU[%d] %s supports SM %d.%d\", devid, deviceProp.name, deviceProp.major, deviceProp.minor);\n printf(\", %s GPU Callback Functions\\n\", (SMversion >= 0x11) ? \"capable\" : \"NOT capable\");\n\n if (SMversion >= 0x11) {\n gpuInfo[max_gpus++] = devid;\n }\n }\n\n printf(\"%d GPUs available to run Callback Functions\\n\", max_gpus);\n\n heterogeneous_workload *workloads;\n workloads = (heterogeneous_workload *)malloc(N_workloads * sizeof(heterogeneous_workload));\n ;\n thread_barrier = cutCreateBarrier(N_workloads);\n\n // Main thread spawns a CPU worker thread for each heterogeneous workload\n printf(\"Starting %d heterogeneous computing workloads\\n\", N_workloads);\n\n for (int i = 0; i < N_workloads; ++i) {\n workloads[i].id = i;\n workloads[i].cudaDeviceID = gpuInfo[i % max_gpus]; // i % N_gpus;\n\n cutStartThread(launch, &workloads[i]);\n }\n\n // Sleep until all workloads have finished\n cutWaitForBarrier(&thread_barrier);\n printf(\"Total of %d workloads finished:\\n\", N_workloads);\n\n bool success = true;\n\n for (int i = 0; i < N_workloads; ++i) {\n success &= workloads[i].success;\n }\n\n printf(\"%s\\n\", success ? \"Success\" : \"Failure\");\n\n free(workloads);\n\n exit(success ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample implements multi-threaded heterogeneous computing workloads with\n * the new CPU callbacks for CUDA streams and events introduced with CUDA 5.0.\n * Together with the thread safety of the CUDA API implementing heterogeneous\n * workloads that float between CPU threads and GPUs has become simple and\n * efficient.\n *\n * The workloads in the sample follow the form CPU preprocess -> GPU process ->\n * CPU postprocess.\n * Each CPU processing step is handled by its own dedicated thread. GPU\n * workloads are sent to all available GPUs in the system.\n *\n */\n\n// System includes\n#include \"hip/hip_runtime.h\"\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#include \"multithreading.h\"\n\nconst int N_workloads = 8;\nconst int N_elements_per_workload = 100000;\n\nCUTBarrier thread_barrier;\n\nvoid CUDART_CB myStreamCallback(hipStream_t event, hipError_t status, void *data);\n\nstruct heterogeneous_workload\n{\n int id;\n int cudaDeviceID;\n\n int *h_data;\n int *d_data;\n hipStream_t stream;\n\n bool success;\n};\n\n__global__ void incKernel(int *data, int N)\n{\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < N)\n data[i]++;\n}\n\nCUT_THREADPROC launch(void *void_arg)\n{\n heterogeneous_workload *workload = (heterogeneous_workload *)void_arg;\n\n // Select GPU for this CPU thread\n checkCudaErrors(hipSetDevice(workload->cudaDeviceID));\n\n // Allocate Resources\n checkCudaErrors(hipStreamCreate(&workload->stream));\n checkCudaErrors(hipMalloc(&workload->d_data, N_elements_per_workload * sizeof(int)));\n checkCudaErrors(hipHostAlloc(&workload->h_data, N_elements_per_workload * sizeof(int), hipHostMallocPortable));\n\n // CPU thread generates data\n for (int i = 0; i < N_elements_per_workload; ++i) {\n workload->h_data[i] = workload->id + i;\n }\n\n // Schedule work for GPU in CUDA stream without blocking the CPU thread\n // Note: Dedicated streams enable concurrent execution of workloads on the GPU\n dim3 block(512);\n dim3 grid((N_elements_per_workload + block.x - 1) / block.x);\n\n checkCudaErrors(hipMemcpyAsync(workload->d_data,\n workload->h_data,\n N_elements_per_workload * sizeof(int),\n hipMemcpyHostToDevice,\n workload->stream));\n incKernel<<stream>>>(workload->d_data, N_elements_per_workload);\n checkCudaErrors(hipMemcpyAsync(workload->h_data,\n workload->d_data,\n N_elements_per_workload * sizeof(int),\n hipMemcpyDeviceToHost,\n workload->stream));\n\n // New in CUDA 5.0: Add a CPU callback which is called once all currently\n // pending operations in the CUDA stream have finished\n checkCudaErrors(hipStreamAddCallback(workload->stream, myStreamCallback, workload, 0));\n\n CUT_THREADEND;\n // CPU thread end of life, GPU continues to process data...\n}\n\nCUT_THREADPROC postprocess(void *void_arg)\n{\n heterogeneous_workload *workload = (heterogeneous_workload *)void_arg;\n // ... GPU is done with processing, continue on new CPU thread...\n\n // Select GPU for this CPU thread\n checkCudaErrors(hipSetDevice(workload->cudaDeviceID));\n\n // CPU thread consumes results from GPU\n workload->success = true;\n\n for (int i = 0; i < N_workloads; ++i) {\n workload->success &= workload->h_data[i] == i + workload->id + 1;\n }\n\n // Free Resources\n checkCudaErrors(hipFree(workload->d_data));\n checkCudaErrors(hipHostFree(workload->h_data));\n checkCudaErrors(hipStreamDestroy(workload->stream));\n\n // Signal the end of the heterogeneous workload to main thread\n cutIncrementBarrier(&thread_barrier);\n\n CUT_THREADEND;\n}\n\nvoid CUDART_CB myStreamCallback(hipStream_t stream, hipError_t status, void *data)\n{\n // Check status of GPU after stream operations are done\n checkCudaErrors(status);\n\n // Spawn new CPU worker thread and continue processing on the CPU\n cutStartThread(postprocess, data);\n}\n\nint main(int argc, char **argv)\n{\n int N_gpus, max_gpus = 0;\n int gpuInfo[32]; // assume a maximum of 32 GPUs in a system configuration\n\n printf(\"Starting simpleCallback\\n\");\n\n checkCudaErrors(hipGetDeviceCount(&N_gpus));\n printf(\"Found %d CUDA capable GPUs\\n\", N_gpus);\n\n if (N_gpus > 32) {\n printf(\"simpleCallback only supports 32 GPU(s)\\n\");\n }\n\n for (int devid = 0; devid < N_gpus; devid++) {\n int SMversion;\n hipDeviceProp_t deviceProp;\n hipSetDevice(devid);\n hipGetDeviceProperties(&deviceProp, devid);\n SMversion = deviceProp.major << 4 + deviceProp.minor;\n printf(\"GPU[%d] %s supports SM %d.%d\", devid, deviceProp.name, deviceProp.major, deviceProp.minor);\n printf(\", %s GPU Callback Functions\\n\", (SMversion >= 0x11) ? \"capable\" : \"NOT capable\");\n\n if (SMversion >= 0x11) {\n gpuInfo[max_gpus++] = devid;\n }\n }\n\n printf(\"%d GPUs available to run Callback Functions\\n\", max_gpus);\n\n heterogeneous_workload *workloads;\n workloads = (heterogeneous_workload *)malloc(N_workloads * sizeof(heterogeneous_workload));\n ;\n thread_barrier = cutCreateBarrier(N_workloads);\n\n // Main thread spawns a CPU worker thread for each heterogeneous workload\n printf(\"Starting %d heterogeneous computing workloads\\n\", N_workloads);\n\n for (int i = 0; i < N_workloads; ++i) {\n workloads[i].id = i;\n workloads[i].cudaDeviceID = gpuInfo[i % max_gpus]; // i % N_gpus;\n\n cutStartThread(launch, &workloads[i]);\n }\n\n // Sleep until all workloads have finished\n cutWaitForBarrier(&thread_barrier);\n printf(\"Total of %d workloads finished:\\n\", N_workloads);\n\n bool success = true;\n\n for (int i = 0; i < N_workloads; ++i) {\n success &= workloads[i].success;\n }\n\n printf(\"%s\\n\", success ? \"Success\" : \"Failure\");\n\n free(workloads);\n\n exit(success ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleCallback/simpleCallback.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Includes, system\n#include \n\n// Includes CUDA\n#include \n#include \n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n\nnamespace cg = cooperative_groups;\n\n#if __CUDA_ARCH__ >= 700\ntemplate \n__device__ void reduceBlockData(cuda::barrier &barrier,\n cg::thread_block_tile<32> &tile32,\n double &threadSum,\n double *result)\n{\n extern __shared__ double tmp[];\n\n#pragma unroll\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n threadSum += tile32.shfl_down(threadSum, offset);\n }\n if (tile32.thread_rank() == 0) {\n tmp[tile32.meta_group_rank()] = threadSum;\n }\n\n auto token = barrier.arrive();\n\n barrier.wait(std::move(token));\n\n // The warp 0 will perform last round of reduction\n if (tile32.meta_group_rank() == 0) {\n double beta = tile32.thread_rank() < tile32.meta_group_size() ? tmp[tile32.thread_rank()] : 0.0;\n\n#pragma unroll\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n beta += tile32.shfl_down(beta, offset);\n }\n\n if (tile32.thread_rank() == 0) {\n if (writeSquareRoot)\n *result = sqrt(beta);\n else\n *result = beta;\n }\n }\n}\n#endif\n\n__global__ void normVecByDotProductAWBarrier(float *vecA, float *vecB, double *partialResults, int size)\n{\n#if __CUDA_ARCH__ >= 700\n#pragma diag_suppress static_var_with_dynamic_init\n cg::thread_block cta = cg::this_thread_block();\n cg::grid_group grid = cg::this_grid();\n ;\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n __shared__ cuda::barrier barrier;\n\n if (threadIdx.x == 0) {\n init(&barrier, blockDim.x);\n }\n\n cg::sync(cta);\n\n double threadSum = 0.0;\n for (int i = grid.thread_rank(); i < size; i += grid.size()) {\n threadSum += (double)(vecA[i] * vecB[i]);\n }\n\n // Each thread block performs reduction of partial dotProducts and writes to\n // global mem.\n reduceBlockData(barrier, tile32, threadSum, &partialResults[blockIdx.x]);\n\n cg::sync(grid);\n\n // One block performs the final summation of partial dot products\n // of all the thread blocks and writes the sqrt of final dot product.\n if (blockIdx.x == 0) {\n threadSum = 0.0;\n for (int i = cta.thread_rank(); i < gridDim.x; i += cta.size()) {\n threadSum += partialResults[i];\n }\n reduceBlockData(barrier, tile32, threadSum, &partialResults[0]);\n }\n\n cg::sync(grid);\n\n const double finalValue = partialResults[0];\n\n // Perform normalization of vecA & vecB.\n for (int i = grid.thread_rank(); i < size; i += grid.size()) {\n vecA[i] = (float)vecA[i] / finalValue;\n vecB[i] = (float)vecB[i] / finalValue;\n }\n#endif\n}\n\nint runNormVecByDotProductAWBarrier(int argc, char **argv, int deviceId);\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"%s starting...\\n\", argv[0]);\n\n // This will pick the best possible CUDA capable device\n int dev = findCudaDevice(argc, (const char **)argv);\n\n int major = 0;\n checkCudaErrors(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev));\n\n // Arrive-Wait Barrier require a GPU of Volta (SM7X) architecture or higher.\n if (major < 7) {\n printf(\"simpleAWBarrier requires SM 7.0 or higher. Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n int supportsCooperativeLaunch = 0;\n checkCudaErrors(cudaDeviceGetAttribute(&supportsCooperativeLaunch, cudaDevAttrCooperativeLaunch, dev));\n\n if (!supportsCooperativeLaunch) {\n printf(\"\\nSelected GPU (%d) does not support Cooperative Kernel Launch, \"\n \"Waiving the run\\n\",\n dev);\n exit(EXIT_WAIVED);\n }\n\n int testResult = runNormVecByDotProductAWBarrier(argc, argv, dev);\n\n printf(\"%s completed, returned %s\\n\", argv[0], testResult ? \"OK\" : \"ERROR!\");\n exit(testResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\nint runNormVecByDotProductAWBarrier(int argc, char **argv, int deviceId)\n{\n float *vecA, *d_vecA;\n float *vecB, *d_vecB;\n double *d_partialResults;\n int size = 10000000;\n\n checkCudaErrors(cudaMallocHost(&vecA, sizeof(float) * size));\n checkCudaErrors(cudaMallocHost(&vecB, sizeof(float) * size));\n\n checkCudaErrors(cudaMalloc(&d_vecA, sizeof(float) * size));\n checkCudaErrors(cudaMalloc(&d_vecB, sizeof(float) * size));\n\n float baseVal = 2.0;\n for (int i = 0; i < size; i++) {\n vecA[i] = vecB[i] = baseVal;\n }\n\n cudaStream_t stream;\n checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));\n\n checkCudaErrors(cudaMemcpyAsync(d_vecA, vecA, sizeof(float) * size, cudaMemcpyHostToDevice, stream));\n checkCudaErrors(cudaMemcpyAsync(d_vecB, vecB, sizeof(float) * size, cudaMemcpyHostToDevice, stream));\n\n // Kernel configuration, where a one-dimensional\n // grid and one-dimensional blocks are configured.\n int minGridSize = 0, blockSize = 0;\n checkCudaErrors(\n cudaOccupancyMaxPotentialBlockSize(&minGridSize, &blockSize, (void *)normVecByDotProductAWBarrier, 0, size));\n\n int smemSize = ((blockSize / 32) + 1) * sizeof(double);\n\n int numBlocksPerSm = 0;\n checkCudaErrors(cudaOccupancyMaxActiveBlocksPerMultiprocessor(\n &numBlocksPerSm, normVecByDotProductAWBarrier, blockSize, smemSize));\n\n int multiProcessorCount = 0;\n checkCudaErrors(cudaDeviceGetAttribute(&multiProcessorCount, cudaDevAttrMultiProcessorCount, deviceId));\n\n minGridSize = multiProcessorCount * numBlocksPerSm;\n checkCudaErrors(cudaMalloc(&d_partialResults, minGridSize * sizeof(double)));\n\n printf(\"Launching normVecByDotProductAWBarrier kernel with numBlocks = %d \"\n \"blockSize = %d\\n\",\n minGridSize,\n blockSize);\n\n dim3 dimGrid(minGridSize, 1, 1), dimBlock(blockSize, 1, 1);\n\n void *kernelArgs[] = {(void *)&d_vecA, (void *)&d_vecB, (void *)&d_partialResults, (void *)&size};\n\n checkCudaErrors(cudaLaunchCooperativeKernel(\n (void *)normVecByDotProductAWBarrier, dimGrid, dimBlock, kernelArgs, smemSize, stream));\n\n checkCudaErrors(cudaMemcpyAsync(vecA, d_vecA, sizeof(float) * size, cudaMemcpyDeviceToHost, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n float expectedResult = (baseVal / sqrt(size * baseVal * baseVal));\n unsigned int matches = 0;\n for (int i = 0; i < size; i++) {\n if ((vecA[i] - expectedResult) > 0.00001) {\n printf(\"mismatch at i = %d\\n\", i);\n break;\n }\n else {\n matches++;\n }\n }\n\n printf(\"Result = %s\\n\", matches == size ? \"PASSED\" : \"FAILED\");\n checkCudaErrors(cudaFree(d_vecA));\n checkCudaErrors(cudaFree(d_vecB));\n checkCudaErrors(cudaFree(d_partialResults));\n\n checkCudaErrors(cudaFreeHost(vecA));\n checkCudaErrors(cudaFreeHost(vecB));\n return matches == size;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Includes, system\n#include \n\n// Includes CUDA\n#include \n#include \n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n\nnamespace cg = cooperative_groups;\n\n#if __CUDA_ARCH__ >= 700\ntemplate \n__device__ void reduceBlockData(cuda::barrier &barrier,\n cg::thread_block_tile<32> &tile32,\n double &threadSum,\n double *result)\n{\n extern __shared__ double tmp[];\n\n#pragma unroll\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n threadSum += tile32.shfl_down(threadSum, offset);\n }\n if (tile32.thread_rank() == 0) {\n tmp[tile32.meta_group_rank()] = threadSum;\n }\n\n auto token = barrier.arrive();\n\n barrier.wait(std::move(token));\n\n // The warp 0 will perform last round of reduction\n if (tile32.meta_group_rank() == 0) {\n double beta = tile32.thread_rank() < tile32.meta_group_size() ? tmp[tile32.thread_rank()] : 0.0;\n\n#pragma unroll\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n beta += tile32.shfl_down(beta, offset);\n }\n\n if (tile32.thread_rank() == 0) {\n if (writeSquareRoot)\n *result = sqrt(beta);\n else\n *result = beta;\n }\n }\n}\n#endif\n\n__global__ void normVecByDotProductAWBarrier(float *vecA, float *vecB, double *partialResults, int size)\n{\n#if __CUDA_ARCH__ >= 700\n#pragma diag_suppress static_var_with_dynamic_init\n cg::thread_block cta = cg::this_thread_block();\n cg::grid_group grid = cg::this_grid();\n ;\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n __shared__ cuda::barrier barrier;\n\n if (threadIdx.x == 0) {\n init(&barrier, blockDim.x);\n }\n\n cg::sync(cta);\n\n double threadSum = 0.0;\n for (int i = grid.thread_rank(); i < size; i += grid.size()) {\n threadSum += (double)(vecA[i] * vecB[i]);\n }\n\n // Each thread block performs reduction of partial dotProducts and writes to\n // global mem.\n reduceBlockData(barrier, tile32, threadSum, &partialResults[blockIdx.x]);\n\n cg::sync(grid);\n\n // One block performs the final summation of partial dot products\n // of all the thread blocks and writes the sqrt of final dot product.\n if (blockIdx.x == 0) {\n threadSum = 0.0;\n for (int i = cta.thread_rank(); i < gridDim.x; i += cta.size()) {\n threadSum += partialResults[i];\n }\n reduceBlockData(barrier, tile32, threadSum, &partialResults[0]);\n }\n\n cg::sync(grid);\n\n const double finalValue = partialResults[0];\n\n // Perform normalization of vecA & vecB.\n for (int i = grid.thread_rank(); i < size; i += grid.size()) {\n vecA[i] = (float)vecA[i] / finalValue;\n vecB[i] = (float)vecB[i] / finalValue;\n }\n#endif\n}\n\nint runNormVecByDotProductAWBarrier(int argc, char **argv, int deviceId);\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"%s starting...\\n\", argv[0]);\n\n // This will pick the best possible CUDA capable device\n int dev = findCudaDevice(argc, (const char **)argv);\n\n int major = 0;\n checkCudaErrors(hipDeviceGetAttribute(&major, hipDeviceAttributeComputeCapabilityMajor, dev));\n\n // Arrive-Wait Barrier require a GPU of Volta (SM7X) architecture or higher.\n if (major < 7) {\n printf(\"simpleAWBarrier requires SM 7.0 or higher. Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n int supportsCooperativeLaunch = 0;\n checkCudaErrors(hipDeviceGetAttribute(&supportsCooperativeLaunch, hipDeviceAttributeCooperativeLaunch, dev));\n\n if (!supportsCooperativeLaunch) {\n printf(\"\\nSelected GPU (%d) does not support Cooperative Kernel Launch, \"\n \"Waiving the run\\n\",\n dev);\n exit(EXIT_WAIVED);\n }\n\n int testResult = runNormVecByDotProductAWBarrier(argc, argv, dev);\n\n printf(\"%s completed, returned %s\\n\", argv[0], testResult ? \"OK\" : \"ERROR!\");\n exit(testResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\nint runNormVecByDotProductAWBarrier(int argc, char **argv, int deviceId)\n{\n float *vecA, *d_vecA;\n float *vecB, *d_vecB;\n double *d_partialResults;\n int size = 10000000;\n\n checkCudaErrors(hipHostMalloc(&vecA, sizeof(float) * size));\n checkCudaErrors(hipHostMalloc(&vecB, sizeof(float) * size));\n\n checkCudaErrors(hipMalloc(&d_vecA, sizeof(float) * size));\n checkCudaErrors(hipMalloc(&d_vecB, sizeof(float) * size));\n\n float baseVal = 2.0;\n for (int i = 0; i < size; i++) {\n vecA[i] = vecB[i] = baseVal;\n }\n\n hipStream_t stream;\n checkCudaErrors(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));\n\n checkCudaErrors(hipMemcpyAsync(d_vecA, vecA, sizeof(float) * size, hipMemcpyHostToDevice, stream));\n checkCudaErrors(hipMemcpyAsync(d_vecB, vecB, sizeof(float) * size, hipMemcpyHostToDevice, stream));\n\n // Kernel configuration, where a one-dimensional\n // grid and one-dimensional blocks are configured.\n int minGridSize = 0, blockSize = 0;\n checkCudaErrors(\n hipOccupancyMaxPotentialBlockSize(&minGridSize, &blockSize, (void *)normVecByDotProductAWBarrier, 0, size));\n\n int smemSize = ((blockSize / 32) + 1) * sizeof(double);\n\n int numBlocksPerSm = 0;\n checkCudaErrors(hipOccupancyMaxActiveBlocksPerMultiprocessor(\n &numBlocksPerSm, normVecByDotProductAWBarrier, blockSize, smemSize));\n\n int multiProcessorCount = 0;\n checkCudaErrors(hipDeviceGetAttribute(&multiProcessorCount, hipDeviceAttributeMultiprocessorCount, deviceId));\n\n minGridSize = multiProcessorCount * numBlocksPerSm;\n checkCudaErrors(hipMalloc(&d_partialResults, minGridSize * sizeof(double)));\n\n printf(\"Launching normVecByDotProductAWBarrier kernel with numBlocks = %d \"\n \"blockSize = %d\\n\",\n minGridSize,\n blockSize);\n\n dim3 dimGrid(minGridSize, 1, 1), dimBlock(blockSize, 1, 1);\n\n void *kernelArgs[] = {(void *)&d_vecA, (void *)&d_vecB, (void *)&d_partialResults, (void *)&size};\n\n checkCudaErrors(hipLaunchCooperativeKernel(reinterpret_cast((void *)normVecByDotProductAWBarrier), dimGrid, dimBlock, kernelArgs, smemSize, stream));\n\n checkCudaErrors(hipMemcpyAsync(vecA, d_vecA, sizeof(float) * size, hipMemcpyDeviceToHost, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n\n float expectedResult = (baseVal / sqrt(size * baseVal * baseVal));\n unsigned int matches = 0;\n for (int i = 0; i < size; i++) {\n if ((vecA[i] - expectedResult) > 0.00001) {\n printf(\"mismatch at i = %d\\n\", i);\n break;\n }\n else {\n matches++;\n }\n }\n\n printf(\"Result = %s\\n\", matches == size ? \"PASSED\" : \"FAILED\");\n checkCudaErrors(hipFree(d_vecA));\n checkCudaErrors(hipFree(d_vecB));\n checkCudaErrors(hipFree(d_partialResults));\n\n checkCudaErrors(hipHostFree(vecA));\n checkCudaErrors(hipHostFree(vecB));\n return matches == size;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleAWBarrier/simpleAWBarrier.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* Vector addition: C = A + B.\n *\n * This sample is a very basic sample that implements element by element\n * vector addition. It is the same as the sample illustrating Chapter 3\n * of the programming guide with some additions like error checking.\n *\n */\n\n// Device code\nextern \"C\" __global__ void VecAdd_kernel(const float *A, const float *B, float *C, int N)\n{\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (i < N)\n C[i] = A[i] + B[i];\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* Vector addition: C = A + B.\n *\n * This sample is a very basic sample that implements element by element\n * vector addition. It is the same as the sample illustrating Chapter 3\n * of the programming guide with some additions like error checking.\n *\n */\n\n// Device code\n#include \"hip/hip_runtime.h\"\nextern \"C\" __global__ void VecAdd_kernel(const float *A, const float *B, float *C, int N)\n{\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (i < N)\n C[i] = A[i] + B[i];\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/vectorAddMMAP/vectorAdd_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This example shows how to use the clock function to measure the performance\n * of block of threads of a kernel accurately. Blocks are executed in parallel\n * and out of order. Since there's no synchronization mechanism between blocks,\n * we measure the clock once for each block. The clock samples are written to\n * device memory.\n */\n\n// This kernel computes a standard parallel reduction and evaluates the\n// time it takes to do that for each block. The timing results are stored\n// in device memory.\n\nextern \"C\" __global__ void timedReduction(const float *input, float *output, clock_t *timer)\n{\n // __shared__ float shared[2 * blockDim.x];\n extern __shared__ float shared[];\n\n const int tid = threadIdx.x;\n const int bid = blockIdx.x;\n\n if (tid == 0)\n timer[bid] = clock();\n\n // Copy input.\n shared[tid] = input[tid];\n shared[tid + blockDim.x] = input[tid + blockDim.x];\n\n // Perform reduction to find minimum.\n for (int d = blockDim.x; d > 0; d /= 2) {\n __syncthreads();\n\n if (tid < d) {\n float f0 = shared[tid];\n float f1 = shared[tid + d];\n\n if (f1 < f0) {\n shared[tid] = f1;\n }\n }\n }\n\n // Write result.\n if (tid == 0)\n output[bid] = shared[0];\n\n __syncthreads();\n\n if (tid == 0)\n timer[bid + gridDim.x] = clock();\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This example shows how to use the clock function to measure the performance\n * of block of threads of a kernel accurately. Blocks are executed in parallel\n * and out of order. Since there's no synchronization mechanism between blocks,\n * we measure the clock once for each block. The clock samples are written to\n * device memory.\n */\n\n// This kernel computes a standard parallel reduction and evaluates the\n// time it takes to do that for each block. The timing results are stored\n// in device memory.\n\n#include \"hip/hip_runtime.h\"\nextern \"C\" __global__ void timedReduction(const float *input, float *output, clock_t *timer)\n{\n // __shared__ float shared[2 * blockDim.x];\n extern __shared__ float shared[];\n\n const int tid = threadIdx.x;\n const int bid = blockIdx.x;\n\n if (tid == 0)\n timer[bid] = clock();\n\n // Copy input.\n shared[tid] = input[tid];\n shared[tid + blockDim.x] = input[tid + blockDim.x];\n\n // Perform reduction to find minimum.\n for (int d = blockDim.x; d > 0; d /= 2) {\n __syncthreads();\n\n if (tid < d) {\n float f0 = shared[tid];\n float f1 = shared[tid + d];\n\n if (f1 < f0) {\n shared[tid] = f1;\n }\n }\n }\n\n // Write result.\n if (tid == 0)\n output[bid] = shared[0];\n\n __syncthreads();\n\n if (tid == 0)\n timer[bid + gridDim.x] = clock();\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/clock_nvrtc/clock_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/**\n * CUDA Kernel Device code\n *\n * Computes the vector addition of A and B into C. The 3 vectors have the same\n * number of elements numElements.\n */\n\nextern \"C\" __global__ void vectorAdd(const float *A, const float *B, float *C, int numElements)\n{\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (i < numElements) {\n C[i] = A[i] + B[i];\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/**\n * CUDA Kernel Device code\n *\n * Computes the vector addition of A and B into C. The 3 vectors have the same\n * number of elements numElements.\n */\n\n#include \"hip/hip_runtime.h\"\nextern \"C\" __global__ void vectorAdd(const float *A, const float *B, float *C, int numElements)\n{\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (i < numElements) {\n C[i] = A[i] + B[i];\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/vectorAdd_nvrtc/vectorAdd_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates how to use texture fetches from layered 2D textures\n * in CUDA C\n *\n * This sample first generates a 3D input data array for the layered texture\n * and the expected output. Then it starts CUDA C kernels, one for each layer,\n * which fetch their layer's texture data (using normalized texture coordinates)\n * transform it to the expected output, and write it to a 3D output data array.\n */\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// includes, kernels\n#include \n\n// includes, project\n#include \n#include // helper for shared that are common to CUDA Samples\n\nstatic const char *sSDKname = \"simpleLayeredTexture\";\n\n////////////////////////////////////////////////////////////////////////////////\n//! Transform a layer of a layered 2D texture using texture lookups\n//! @param g_odata output data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void transformKernel(float *g_odata, int width, int height, int layer, cudaTextureObject_t tex)\n{\n // calculate this thread's data point\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // 0.5f offset and division are necessary to access the original data points\n // in the texture (such that bilinear interpolation will not be activated).\n // For details, see also CUDA Programming Guide, Appendix D\n float u = (x + 0.5f) / (float)width;\n float v = (y + 0.5f) / (float)height;\n\n // read from texture, do expected transformation and write to global memory\n g_odata[layer * width * height + y * width + x] = -tex2DLayered(tex, u, v, layer) + layer;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"[%s] - Starting...\\n\", sSDKname);\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n int devID = findCudaDevice(argc, (const char **)argv);\n\n bool bResult = true;\n\n // get number of SMs on this GPU\n cudaDeviceProp deviceProps;\n\n checkCudaErrors(cudaGetDeviceProperties(&deviceProps, devID));\n printf(\"CUDA device [%s] has %d Multi-Processors \", deviceProps.name, deviceProps.multiProcessorCount);\n printf(\"SM %d.%d\\n\", deviceProps.major, deviceProps.minor);\n\n // generate input data for layered texture\n unsigned int width = 512, height = 512, num_layers = 5;\n unsigned int size = width * height * num_layers * sizeof(float);\n float *h_data = (float *)malloc(size);\n\n for (unsigned int layer = 0; layer < num_layers; layer++)\n for (int i = 0; i < (int)(width * height); i++) {\n h_data[layer * width * height + i] = (float)i;\n }\n\n // this is the expected transformation of the input data (the expected output)\n float *h_data_ref = (float *)malloc(size);\n\n for (unsigned int layer = 0; layer < num_layers; layer++)\n for (int i = 0; i < (int)(width * height); i++) {\n h_data_ref[layer * width * height + i] = -h_data[layer * width * height + i] + layer;\n }\n\n // allocate device memory for result\n float *d_data = NULL;\n checkCudaErrors(cudaMalloc((void **)&d_data, size));\n\n // allocate array and copy image data\n cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc(32, 0, 0, 0, cudaChannelFormatKindFloat);\n cudaArray *cu_3darray;\n checkCudaErrors(\n cudaMalloc3DArray(&cu_3darray, &channelDesc, make_cudaExtent(width, height, num_layers), cudaArrayLayered));\n cudaMemcpy3DParms myparms = {0};\n myparms.srcPos = make_cudaPos(0, 0, 0);\n myparms.dstPos = make_cudaPos(0, 0, 0);\n myparms.srcPtr = make_cudaPitchedPtr(h_data, width * sizeof(float), width, height);\n myparms.dstArray = cu_3darray;\n myparms.extent = make_cudaExtent(width, height, num_layers);\n myparms.kind = cudaMemcpyHostToDevice;\n checkCudaErrors(cudaMemcpy3D(&myparms));\n\n cudaTextureObject_t tex;\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = cu_3darray;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = cudaFilterModeLinear;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n\n dim3 dimBlock(8, 8, 1);\n dim3 dimGrid(width / dimBlock.x, height / dimBlock.y, 1);\n\n printf(\"Covering 2D data array of %d x %d: Grid size is %d x %d, each block has \"\n \"8 x 8 threads\\n\",\n width,\n height,\n dimGrid.x,\n dimGrid.y);\n\n transformKernel<<>>(d_data, width, height, 0,\n tex); // warmup (for better timing)\n\n // check if kernel execution generated an error\n getLastCudaError(\"warmup Kernel execution failed\");\n\n checkCudaErrors(cudaDeviceSynchronize());\n\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n sdkStartTimer(&timer);\n\n // execute the kernel\n for (unsigned int layer = 0; layer < num_layers; layer++)\n transformKernel<<>>(d_data, width, height, layer, tex);\n\n // check if kernel execution generated an error\n getLastCudaError(\"Kernel execution failed\");\n\n checkCudaErrors(cudaDeviceSynchronize());\n sdkStopTimer(&timer);\n printf(\"Processing time: %.3f msec\\n\", sdkGetTimerValue(&timer));\n printf(\"%.2f Mtexlookups/sec\\n\", (width * height * num_layers / (sdkGetTimerValue(&timer) / 1000.0f) / 1e6));\n sdkDeleteTimer(&timer);\n\n // allocate mem for the result on host side\n float *h_odata = (float *)malloc(size);\n // copy result from device to host\n checkCudaErrors(cudaMemcpy(h_odata, d_data, size, cudaMemcpyDeviceToHost));\n\n // write regression file if necessary\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // write file for regression test\n sdkWriteFile(\"./data/regression.dat\", h_odata, width * height, 0.0f, false);\n }\n else {\n printf(\"Comparing kernel output to expected data\\n\");\n\n#define MIN_EPSILON_ERROR 5e-3f\n bResult = compareData(h_odata, h_data_ref, width * height * num_layers, MIN_EPSILON_ERROR, 0.0f);\n }\n\n // cleanup memory\n free(h_data);\n free(h_data_ref);\n free(h_odata);\n\n checkCudaErrors(cudaDestroyTextureObject(tex));\n checkCudaErrors(cudaFree(d_data));\n checkCudaErrors(cudaFreeArray(cu_3darray));\n\n exit(bResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates how to use texture fetches from layered 2D textures\n * in CUDA C\n *\n * This sample first generates a 3D input data array for the layered texture\n * and the expected output. Then it starts CUDA C kernels, one for each layer,\n * which fetch their layer's texture data (using normalized texture coordinates)\n * transform it to the expected output, and write it to a 3D output data array.\n */\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// includes, kernels\n#include \n\n// includes, project\n#include \n#include // helper for shared that are common to CUDA Samples\n\nstatic const char *sSDKname = \"simpleLayeredTexture\";\n\n////////////////////////////////////////////////////////////////////////////////\n//! Transform a layer of a layered 2D texture using texture lookups\n//! @param g_odata output data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void transformKernel(float *g_odata, int width, int height, int layer, hipTextureObject_t tex)\n{\n // calculate this thread's data point\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // 0.5f offset and division are necessary to access the original data points\n // in the texture (such that bilinear interpolation will not be activated).\n // For details, see also CUDA Programming Guide, Appendix D\n float u = (x + 0.5f) / (float)width;\n float v = (y + 0.5f) / (float)height;\n\n // read from texture, do expected transformation and write to global memory\n g_odata[layer * width * height + y * width + x] = -tex2DLayered(tex, u, v, layer) + layer;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"[%s] - Starting...\\n\", sSDKname);\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n int devID = findCudaDevice(argc, (const char **)argv);\n\n bool bResult = true;\n\n // get number of SMs on this GPU\n hipDeviceProp_t deviceProps;\n\n checkCudaErrors(hipGetDeviceProperties(&deviceProps, devID));\n printf(\"CUDA device [%s] has %d Multi-Processors \", deviceProps.name, deviceProps.multiProcessorCount);\n printf(\"SM %d.%d\\n\", deviceProps.major, deviceProps.minor);\n\n // generate input data for layered texture\n unsigned int width = 512, height = 512, num_layers = 5;\n unsigned int size = width * height * num_layers * sizeof(float);\n float *h_data = (float *)malloc(size);\n\n for (unsigned int layer = 0; layer < num_layers; layer++)\n for (int i = 0; i < (int)(width * height); i++) {\n h_data[layer * width * height + i] = (float)i;\n }\n\n // this is the expected transformation of the input data (the expected output)\n float *h_data_ref = (float *)malloc(size);\n\n for (unsigned int layer = 0; layer < num_layers; layer++)\n for (int i = 0; i < (int)(width * height); i++) {\n h_data_ref[layer * width * height + i] = -h_data[layer * width * height + i] + layer;\n }\n\n // allocate device memory for result\n float *d_data = NULL;\n checkCudaErrors(hipMalloc((void **)&d_data, size));\n\n // allocate array and copy image data\n hipChannelFormatDesc channelDesc = hipCreateChannelDesc(32, 0, 0, 0, hipChannelFormatKindFloat);\n hipArray *cu_3darray;\n checkCudaErrors(\n hipMalloc3DArray(&cu_3darray, &channelDesc, make_hipExtent(width, height, num_layers), hipArrayLayered));\n hipMemcpy3DParms myparms = {0};\n myparms.srcPos = make_hipPos(0, 0, 0);\n myparms.dstPos = make_hipPos(0, 0, 0);\n myparms.srcPtr = make_hipPitchedPtr(h_data, width * sizeof(float), width, height);\n myparms.dstArray = cu_3darray;\n myparms.extent = make_hipExtent(width, height, num_layers);\n myparms.kind = hipMemcpyHostToDevice;\n checkCudaErrors(hipMemcpy3D(&myparms));\n\n hipTextureObject_t tex;\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = cu_3darray;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = hipFilterModeLinear;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n\n dim3 dimBlock(8, 8, 1);\n dim3 dimGrid(width / dimBlock.x, height / dimBlock.y, 1);\n\n printf(\"Covering 2D data array of %d x %d: Grid size is %d x %d, each block has \"\n \"8 x 8 threads\\n\",\n width,\n height,\n dimGrid.x,\n dimGrid.y);\n\n transformKernel<<>>(d_data, width, height, 0,\n tex); // warmup (for better timing)\n\n // check if kernel execution generated an error\n getLastCudaError(\"warmup Kernel execution failed\");\n\n checkCudaErrors(hipDeviceSynchronize());\n\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n sdkStartTimer(&timer);\n\n // execute the kernel\n for (unsigned int layer = 0; layer < num_layers; layer++)\n transformKernel<<>>(d_data, width, height, layer, tex);\n\n // check if kernel execution generated an error\n getLastCudaError(\"Kernel execution failed\");\n\n checkCudaErrors(hipDeviceSynchronize());\n sdkStopTimer(&timer);\n printf(\"Processing time: %.3f msec\\n\", sdkGetTimerValue(&timer));\n printf(\"%.2f Mtexlookups/sec\\n\", (width * height * num_layers / (sdkGetTimerValue(&timer) / 1000.0f) / 1e6));\n sdkDeleteTimer(&timer);\n\n // allocate mem for the result on host side\n float *h_odata = (float *)malloc(size);\n // copy result from device to host\n checkCudaErrors(hipMemcpy(h_odata, d_data, size, hipMemcpyDeviceToHost));\n\n // write regression file if necessary\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // write file for regression test\n sdkWriteFile(\"./data/regression.dat\", h_odata, width * height, 0.0f, false);\n }\n else {\n printf(\"Comparing kernel output to expected data\\n\");\n\n#define MIN_EPSILON_ERROR 5e-3f\n bResult = compareData(h_odata, h_data_ref, width * height * num_layers, MIN_EPSILON_ERROR, 0.0f);\n }\n\n // cleanup memory\n free(h_data);\n free(h_data_ref);\n free(h_odata);\n\n checkCudaErrors(hipDestroyTextureObject(tex));\n checkCudaErrors(hipFree(d_data));\n checkCudaErrors(hipFreeArray(cu_3darray));\n\n exit(bResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleLayeredTexture/simpleLayeredTexture.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/**\n *\n * This sample is a simple code that illustrates basic usage of\n * cooperative groups within the thread block. The code launches a single\n * thread block, creates a cooperative group of all threads in the block,\n * and a set of tiled partition cooperative groups. For each, it uses a\n * generic reduction function to calculate the sum of all the ranks in\n * that group. In each case the result is printed, together with the\n * expected answer (which is calculated using the analytical formula\n * (n-1)*n)/2, noting that the ranks start at zero).\n *\n */\n\n#include \n#include \n\nusing namespace cooperative_groups;\n\n/**\n * CUDA device function\n *\n * calculates the sum of val across the group g. The workspace array, x,\n * must be large enough to contain g.size() integers.\n */\n__device__ int sumReduction(thread_group g, int *x, int val)\n{\n // rank of this thread in the group\n int lane = g.thread_rank();\n\n // for each iteration of this loop, the number of threads active in the\n // reduction, i, is halved, and each active thread (with index [lane])\n // performs a single summation of it's own value with that\n // of a \"partner\" (with index [lane+i]).\n for (int i = g.size() / 2; i > 0; i /= 2) {\n // store value for this thread in temporary array\n x[lane] = val;\n\n // synchronize all threads in group\n g.sync();\n\n if (lane < i)\n // active threads perform summation of their value with\n // their partner's value\n val += x[lane + i];\n\n // synchronize all threads in group\n g.sync();\n }\n\n // master thread in group returns result, and others return -1.\n if (g.thread_rank() == 0)\n return val;\n else\n return -1;\n}\n\n/**\n * CUDA kernel device code\n *\n * Creates cooperative groups and performs reductions\n */\n__global__ void cgkernel()\n{\n // threadBlockGroup includes all threads in the block\n thread_block threadBlockGroup = this_thread_block();\n int threadBlockGroupSize = threadBlockGroup.size();\n\n // workspace array in shared memory required for reduction\n extern __shared__ int workspace[];\n\n int input, output, expectedOutput;\n\n // input to reduction, for each thread, is its' rank in the group\n input = threadBlockGroup.thread_rank();\n\n // expected output from analytical formula (n-1)(n)/2\n // (noting that indexing starts at 0 rather than 1)\n expectedOutput = (threadBlockGroupSize - 1) * threadBlockGroupSize / 2;\n\n // perform reduction\n output = sumReduction(threadBlockGroup, workspace, input);\n\n // master thread in group prints out result\n if (threadBlockGroup.thread_rank() == 0) {\n printf(\" Sum of all ranks 0..%d in threadBlockGroup is %d (expected %d)\\n\\n\",\n (int)threadBlockGroup.size() - 1,\n output,\n expectedOutput);\n\n printf(\" Now creating %d groups, each of size 16 threads:\\n\\n\", (int)threadBlockGroup.size() / 16);\n }\n\n threadBlockGroup.sync();\n\n // each tiledPartition16 group includes 16 threads\n thread_block_tile<16> tiledPartition16 = tiled_partition<16>(threadBlockGroup);\n\n // This offset allows each group to have its own unique area in the workspace\n // array\n int workspaceOffset = threadBlockGroup.thread_rank() - tiledPartition16.thread_rank();\n\n // input to reduction, for each thread, is its' rank in the group\n input = tiledPartition16.thread_rank();\n\n // expected output from analytical formula (n-1)(n)/2\n // (noting that indexing starts at 0 rather than 1)\n expectedOutput = 15 * 16 / 2;\n\n // Perform reduction\n output = sumReduction(tiledPartition16, workspace + workspaceOffset, input);\n\n // each master thread prints out result\n if (tiledPartition16.thread_rank() == 0)\n printf(\" Sum of all ranks 0..15 in this tiledPartition16 group is %d \"\n \"(expected %d)\\n\",\n output,\n expectedOutput);\n\n return;\n}\n\n/**\n * Host main routine\n */\nint main()\n{\n // Error code to check return values for CUDA calls\n cudaError_t err;\n\n // Launch the kernel\n\n int blocksPerGrid = 1;\n int threadsPerBlock = 64;\n\n printf(\"\\nLaunching a single block with %d threads...\\n\\n\", threadsPerBlock);\n\n // we use the optional third argument to specify the size\n // of shared memory required in the kernel\n cgkernel<<>>();\n err = cudaDeviceSynchronize();\n\n if (err != cudaSuccess) {\n fprintf(stderr, \"Failed to launch kernel (error code %s)!\\n\", cudaGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n printf(\"\\n...Done.\\n\\n\");\n\n return 0;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/**\n *\n * This sample is a simple code that illustrates basic usage of\n * cooperative groups within the thread block. The code launches a single\n * thread block, creates a cooperative group of all threads in the block,\n * and a set of tiled partition cooperative groups. For each, it uses a\n * generic reduction function to calculate the sum of all the ranks in\n * that group. In each case the result is printed, together with the\n * expected answer (which is calculated using the analytical formula\n * (n-1)*n)/2, noting that the ranks start at zero).\n *\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\nusing namespace cooperative_groups;\n\n/**\n * CUDA device function\n *\n * calculates the sum of val across the group g. The workspace array, x,\n * must be large enough to contain g.size() integers.\n */\n__device__ int sumReduction(thread_group g, int *x, int val)\n{\n // rank of this thread in the group\n int lane = g.thread_rank();\n\n // for each iteration of this loop, the number of threads active in the\n // reduction, i, is halved, and each active thread (with index [lane])\n // performs a single summation of it's own value with that\n // of a \"partner\" (with index [lane+i]).\n for (int i = g.size() / 2; i > 0; i /= 2) {\n // store value for this thread in temporary array\n x[lane] = val;\n\n // synchronize all threads in group\n g.sync();\n\n if (lane < i)\n // active threads perform summation of their value with\n // their partner's value\n val += x[lane + i];\n\n // synchronize all threads in group\n g.sync();\n }\n\n // master thread in group returns result, and others return -1.\n if (g.thread_rank() == 0)\n return val;\n else\n return -1;\n}\n\n/**\n * CUDA kernel device code\n *\n * Creates cooperative groups and performs reductions\n */\n__global__ void cgkernel()\n{\n // threadBlockGroup includes all threads in the block\n thread_block threadBlockGroup = this_thread_block();\n int threadBlockGroupSize = threadBlockGroup.size();\n\n // workspace array in shared memory required for reduction\n extern __shared__ int workspace[];\n\n int input, output, expectedOutput;\n\n // input to reduction, for each thread, is its' rank in the group\n input = threadBlockGroup.thread_rank();\n\n // expected output from analytical formula (n-1)(n)/2\n // (noting that indexing starts at 0 rather than 1)\n expectedOutput = (threadBlockGroupSize - 1) * threadBlockGroupSize / 2;\n\n // perform reduction\n output = sumReduction(threadBlockGroup, workspace, input);\n\n // master thread in group prints out result\n if (threadBlockGroup.thread_rank() == 0) {\n printf(\" Sum of all ranks 0..%d in threadBlockGroup is %d (expected %d)\\n\\n\",\n (int)threadBlockGroup.size() - 1,\n output,\n expectedOutput);\n\n printf(\" Now creating %d groups, each of size 16 threads:\\n\\n\", (int)threadBlockGroup.size() / 16);\n }\n\n threadBlockGroup.sync();\n\n // each tiledPartition16 group includes 16 threads\n thread_block_tile<16> tiledPartition16 = tiled_partition<16>(threadBlockGroup);\n\n // This offset allows each group to have its own unique area in the workspace\n // array\n int workspaceOffset = threadBlockGroup.thread_rank() - tiledPartition16.thread_rank();\n\n // input to reduction, for each thread, is its' rank in the group\n input = tiledPartition16.thread_rank();\n\n // expected output from analytical formula (n-1)(n)/2\n // (noting that indexing starts at 0 rather than 1)\n expectedOutput = 15 * 16 / 2;\n\n // Perform reduction\n output = sumReduction(tiledPartition16, workspace + workspaceOffset, input);\n\n // each master thread prints out result\n if (tiledPartition16.thread_rank() == 0)\n printf(\" Sum of all ranks 0..15 in this tiledPartition16 group is %d \"\n \"(expected %d)\\n\",\n output,\n expectedOutput);\n\n return;\n}\n\n/**\n * Host main routine\n */\nint main()\n{\n // Error code to check return values for CUDA calls\n hipError_t err;\n\n // Launch the kernel\n\n int blocksPerGrid = 1;\n int threadsPerBlock = 64;\n\n printf(\"\\nLaunching a single block with %d threads...\\n\\n\", threadsPerBlock);\n\n // we use the optional third argument to specify the size\n // of shared memory required in the kernel\n cgkernel<<>>();\n err = hipDeviceSynchronize();\n\n if (err != hipSuccess) {\n fprintf(stderr, \"Failed to launch kernel (error code %s)!\\n\", hipGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n printf(\"\\n...Done.\\n\\n\");\n\n return 0;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleCooperativeGroups/simpleCooperativeGroups.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* Simple example demonstrating how to use MPI with CUDA\n *\n * Generate some random numbers on one node.\n * Dispatch them to all nodes.\n * Compute their square root on each node's GPU.\n * Compute the average of the results using MPI.\n *\n * simpleMPI.cu: GPU part, compiled with nvcc\n */\n\n#include \nusing std::cerr;\nusing std::endl;\n\n#include \"simpleMPI.h\"\n\n// Error handling macro\n#define CUDA_CHECK(call) \\\n if ((call) != cudaSuccess) { \\\n cudaError_t err = cudaGetLastError(); \\\n cerr << \"CUDA error calling \\\"\" #call \"\\\", code is \" << err << endl; \\\n my_abort(err); \\\n }\n\n// Device code\n// Very simple GPU Kernel that computes square roots of input numbers\n__global__ void simpleMPIKernel(float *input, float *output)\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n output[tid] = sqrt(input[tid]);\n}\n\n// Initialize an array with random data (between 0 and 1)\nvoid initData(float *data, int dataSize)\n{\n for (int i = 0; i < dataSize; i++) {\n data[i] = (float)rand() / RAND_MAX;\n }\n}\n\n// CUDA computation on each node\n// No MPI here, only CUDA\nvoid computeGPU(float *hostData, int blockSize, int gridSize)\n{\n int dataSize = blockSize * gridSize;\n\n // Allocate data on GPU memory\n float *deviceInputData = NULL;\n CUDA_CHECK(cudaMalloc((void **)&deviceInputData, dataSize * sizeof(float)));\n\n float *deviceOutputData = NULL;\n CUDA_CHECK(cudaMalloc((void **)&deviceOutputData, dataSize * sizeof(float)));\n\n // Copy to GPU memory\n CUDA_CHECK(cudaMemcpy(deviceInputData, hostData, dataSize * sizeof(float), cudaMemcpyHostToDevice));\n\n // Run kernel\n simpleMPIKernel<<>>(deviceInputData, deviceOutputData);\n\n // Copy data back to CPU memory\n CUDA_CHECK(cudaMemcpy(hostData, deviceOutputData, dataSize * sizeof(float), cudaMemcpyDeviceToHost));\n\n // Free GPU memory\n CUDA_CHECK(cudaFree(deviceInputData));\n CUDA_CHECK(cudaFree(deviceOutputData));\n}\n\nfloat sum(float *data, int size)\n{\n float accum = 0.f;\n\n for (int i = 0; i < size; i++) {\n accum += data[i];\n }\n\n return accum;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* Simple example demonstrating how to use MPI with CUDA\n *\n * Generate some random numbers on one node.\n * Dispatch them to all nodes.\n * Compute their square root on each node's GPU.\n * Compute the average of the results using MPI.\n *\n * simpleMPI.cu: GPU part, compiled with nvcc\n */\n\n#include \"hip/hip_runtime.h\"\n#include \nusing std::cerr;\nusing std::endl;\n\n#include \"simpleMPI.h\"\n\n// Error handling macro\n#define CUDA_CHECK(call) \\\n if ((call) != hipSuccess) { \\\n hipError_t err = hipGetLastError(); \\\n cerr << \"CUDA error calling \\\"\" #call \"\\\", code is \" << err << endl; \\\n my_abort(err); \\\n }\n\n// Device code\n// Very simple GPU Kernel that computes square roots of input numbers\n__global__ void simpleMPIKernel(float *input, float *output)\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n output[tid] = sqrt(input[tid]);\n}\n\n// Initialize an array with random data (between 0 and 1)\nvoid initData(float *data, int dataSize)\n{\n for (int i = 0; i < dataSize; i++) {\n data[i] = (float)rand() / RAND_MAX;\n }\n}\n\n// CUDA computation on each node\n// No MPI here, only CUDA\nvoid computeGPU(float *hostData, int blockSize, int gridSize)\n{\n int dataSize = blockSize * gridSize;\n\n // Allocate data on GPU memory\n float *deviceInputData = NULL;\n CUDA_CHECK(hipMalloc((void **)&deviceInputData, dataSize * sizeof(float)));\n\n float *deviceOutputData = NULL;\n CUDA_CHECK(hipMalloc((void **)&deviceOutputData, dataSize * sizeof(float)));\n\n // Copy to GPU memory\n CUDA_CHECK(hipMemcpy(deviceInputData, hostData, dataSize * sizeof(float), hipMemcpyHostToDevice));\n\n // Run kernel\n simpleMPIKernel<<>>(deviceInputData, deviceOutputData);\n\n // Copy data back to CPU memory\n CUDA_CHECK(hipMemcpy(hostData, deviceOutputData, dataSize * sizeof(float), hipMemcpyDeviceToHost));\n\n // Free GPU memory\n CUDA_CHECK(hipFree(deviceInputData));\n CUDA_CHECK(hipFree(deviceOutputData));\n}\n\nfloat sum(float *data, int size)\n{\n float accum = 0.f;\n\n for (int i = 0; i < size; i++) {\n accum += data[i];\n }\n\n return accum;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleMPI/simpleMPI.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#ifndef _SIMPLETEXTURE3D_KERNEL_CU_\n#define _SIMPLETEXTURE3D_KERNEL_CU_\n\n#include \n#include \n#include \n#include \n#include \n#include \n\ntypedef unsigned int uint;\ntypedef unsigned char uchar;\n\ncudaArray *d_volumeArray = 0;\ncudaTextureObject_t tex; // 3D texture\n\n__global__ void d_render(uint *d_output, uint imageW, uint imageH, float w, cudaTextureObject_t texObj)\n{\n uint x = __umul24(blockIdx.x, blockDim.x) + threadIdx.x;\n uint y = __umul24(blockIdx.y, blockDim.y) + threadIdx.y;\n\n float u = x / (float)imageW;\n float v = y / (float)imageH;\n // read from 3D texture\n float voxel = tex3D(texObj, u, v, w);\n\n if ((x < imageW) && (y < imageH)) {\n // write output color\n uint i = __umul24(y, imageW) + x;\n d_output[i] = voxel * 255;\n }\n}\n\nextern \"C\" void setTextureFilterMode(bool bLinearFilter)\n{\n if (tex) {\n checkCudaErrors(cudaDestroyTextureObject(tex));\n }\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = d_volumeArray;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = bLinearFilter ? cudaFilterModeLinear : cudaFilterModePoint;\n ;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.addressMode[2] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeNormalizedFloat;\n\n checkCudaErrors(cudaCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void initCuda(const uchar *h_volume, cudaExtent volumeSize)\n{\n // create 3D array\n cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc();\n checkCudaErrors(cudaMalloc3DArray(&d_volumeArray, &channelDesc, volumeSize));\n\n // copy data to 3D array\n cudaMemcpy3DParms copyParams = {0};\n copyParams.srcPtr =\n make_cudaPitchedPtr((void *)h_volume, volumeSize.width * sizeof(uchar), volumeSize.width, volumeSize.height);\n copyParams.dstArray = d_volumeArray;\n copyParams.extent = volumeSize;\n copyParams.kind = cudaMemcpyHostToDevice;\n checkCudaErrors(cudaMemcpy3D(©Params));\n\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = d_volumeArray;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n // access with normalized texture coordinates\n texDescr.normalizedCoords = true;\n // linear interpolation\n texDescr.filterMode = cudaFilterModeLinear;\n // wrap texture coordinates\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.addressMode[2] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeNormalizedFloat;\n\n checkCudaErrors(cudaCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void render_kernel(dim3 gridSize, dim3 blockSize, uint *d_output, uint imageW, uint imageH, float w)\n{\n d_render<<>>(d_output, imageW, imageH, w, tex);\n}\n\nvoid cleanupCuda()\n{\n if (tex) {\n checkCudaErrors(cudaDestroyTextureObject(tex));\n }\n if (d_volumeArray) {\n checkCudaErrors(cudaFreeArray(d_volumeArray));\n }\n}\n\n#endif // #ifndef _SIMPLETEXTURE3D_KERNEL_CU_\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#ifndef _SIMPLETEXTURE3D_KERNEL_CU_\n#define _SIMPLETEXTURE3D_KERNEL_CU_\n\n#include \n#include \n#include \n#include \n#include \n#include \n\ntypedef unsigned int uint;\ntypedef unsigned char uchar;\n\nhipArray *d_volumeArray = 0;\nhipTextureObject_t tex; // 3D texture\n\n__global__ void d_render(uint *d_output, uint imageW, uint imageH, float w, hipTextureObject_t texObj)\n{\n uint x = __umul24(blockIdx.x, blockDim.x) + threadIdx.x;\n uint y = __umul24(blockIdx.y, blockDim.y) + threadIdx.y;\n\n float u = x / (float)imageW;\n float v = y / (float)imageH;\n // read from 3D texture\n float voxel = tex3D(texObj, u, v, w);\n\n if ((x < imageW) && (y < imageH)) {\n // write output color\n uint i = __umul24(y, imageW) + x;\n d_output[i] = voxel * 255;\n }\n}\n\nextern \"C\" void setTextureFilterMode(bool bLinearFilter)\n{\n if (tex) {\n checkCudaErrors(hipDestroyTextureObject(tex));\n }\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = d_volumeArray;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = bLinearFilter ? hipFilterModeLinear : hipFilterModePoint;\n ;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.addressMode[2] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeNormalizedFloat;\n\n checkCudaErrors(hipCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void initCuda(const uchar *h_volume, hipExtent volumeSize)\n{\n // create 3D array\n hipChannelFormatDesc channelDesc = hipCreateChannelDesc();\n checkCudaErrors(hipMalloc3DArray(&d_volumeArray, &channelDesc, volumeSize));\n\n // copy data to 3D array\n hipMemcpy3DParms copyParams = {0};\n copyParams.srcPtr =\n make_hipPitchedPtr((void *)h_volume, volumeSize.width * sizeof(uchar), volumeSize.width, volumeSize.height);\n copyParams.dstArray = d_volumeArray;\n copyParams.extent = volumeSize;\n copyParams.kind = hipMemcpyHostToDevice;\n checkCudaErrors(hipMemcpy3D(©Params));\n\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = d_volumeArray;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n // access with normalized texture coordinates\n texDescr.normalizedCoords = true;\n // linear interpolation\n texDescr.filterMode = hipFilterModeLinear;\n // wrap texture coordinates\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.addressMode[2] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeNormalizedFloat;\n\n checkCudaErrors(hipCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void render_kernel(dim3 gridSize, dim3 blockSize, uint *d_output, uint imageW, uint imageH, float w)\n{\n d_render<<>>(d_output, imageW, imageH, w, tex);\n}\n\nvoid cleanupCuda()\n{\n if (tex) {\n checkCudaErrors(hipDestroyTextureObject(tex));\n }\n if (d_volumeArray) {\n checkCudaErrors(hipFreeArray(d_volumeArray));\n }\n}\n\n#endif // #ifndef _SIMPLETEXTURE3D_KERNEL_CU_\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleTexture3D/simpleTexture3D_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates a combination of Peer-to-Peer (P2P) and\n * Unified Virtual Address Space (UVA) features new to SDK 4.0\n */\n\n// includes, system\n#include \n#include \n\n// CUDA includes\n#include \n\n// includes, project\n#include \n#include // helper for shared that are common to CUDA Samples\n\n__global__ void SimpleKernel(float *src, float *dst)\n{\n // Just a dummy kernel, doing enough for us to verify that everything\n // worked\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n dst[idx] = src[idx] * 2.0f;\n}\n\ninline bool IsAppBuiltAs64() { return sizeof(void *) == 8; }\n\nint main(int argc, char **argv)\n{\n printf(\"[%s] - Starting...\\n\", argv[0]);\n\n if (!IsAppBuiltAs64()) {\n printf(\"%s is only supported with on 64-bit OSs and the application must be \"\n \"built as a 64-bit target. Test is being waived.\\n\",\n argv[0]);\n exit(EXIT_WAIVED);\n }\n\n // Number of GPUs\n printf(\"Checking for multiple GPUs...\\n\");\n int gpu_n;\n checkCudaErrors(cudaGetDeviceCount(&gpu_n));\n printf(\"CUDA-capable device count: %i\\n\", gpu_n);\n\n if (gpu_n < 2) {\n printf(\"Two or more GPUs with Peer-to-Peer access capability are required for \"\n \"%s.\\n\",\n argv[0]);\n printf(\"Waiving test.\\n\");\n exit(EXIT_WAIVED);\n }\n\n // Query device properties\n cudaDeviceProp prop[64];\n int gpuid[2]; // we want to find the first two GPU's that can support P2P\n\n for (int i = 0; i < gpu_n; i++) {\n checkCudaErrors(cudaGetDeviceProperties(&prop[i], i));\n }\n // Check possibility for peer access\n printf(\"\\nChecking GPU(s) for support of peer to peer memory access...\\n\");\n\n int can_access_peer;\n int p2pCapableGPUs[2]; // We take only 1 pair of P2P capable GPUs\n p2pCapableGPUs[0] = p2pCapableGPUs[1] = -1;\n\n // Show all the combinations of supported P2P GPUs\n for (int i = 0; i < gpu_n; i++) {\n for (int j = 0; j < gpu_n; j++) {\n if (i == j) {\n continue;\n }\n checkCudaErrors(cudaDeviceCanAccessPeer(&can_access_peer, i, j));\n printf(\"> Peer access from %s (GPU%d) -> %s (GPU%d) : %s\\n\",\n prop[i].name,\n i,\n prop[j].name,\n j,\n can_access_peer ? \"Yes\" : \"No\");\n if (can_access_peer && p2pCapableGPUs[0] == -1) {\n p2pCapableGPUs[0] = i;\n p2pCapableGPUs[1] = j;\n }\n }\n }\n\n if (p2pCapableGPUs[0] == -1 || p2pCapableGPUs[1] == -1) {\n printf(\"Two or more GPUs with Peer-to-Peer access capability are required for \"\n \"%s.\\n\",\n argv[0]);\n printf(\"Peer to Peer access is not available amongst GPUs in the system, \"\n \"waiving test.\\n\");\n\n exit(EXIT_WAIVED);\n }\n\n // Use first pair of p2p capable GPUs detected.\n gpuid[0] = p2pCapableGPUs[0];\n gpuid[1] = p2pCapableGPUs[1];\n\n // Enable peer access\n printf(\"Enabling peer access between GPU%d and GPU%d...\\n\", gpuid[0], gpuid[1]);\n checkCudaErrors(cudaSetDevice(gpuid[0]));\n checkCudaErrors(cudaDeviceEnablePeerAccess(gpuid[1], 0));\n checkCudaErrors(cudaSetDevice(gpuid[1]));\n checkCudaErrors(cudaDeviceEnablePeerAccess(gpuid[0], 0));\n\n // Allocate buffers\n const size_t buf_size = 1024 * 1024 * 16 * sizeof(float);\n printf(\n \"Allocating buffers (%iMB on GPU%d, GPU%d and CPU Host)...\\n\", int(buf_size / 1024 / 1024), gpuid[0], gpuid[1]);\n checkCudaErrors(cudaSetDevice(gpuid[0]));\n float *g0;\n checkCudaErrors(cudaMalloc(&g0, buf_size));\n checkCudaErrors(cudaSetDevice(gpuid[1]));\n float *g1;\n checkCudaErrors(cudaMalloc(&g1, buf_size));\n float *h0;\n checkCudaErrors(cudaMallocHost(&h0, buf_size)); // Automatically portable with UVA\n\n // Create CUDA event handles\n printf(\"Creating event handles...\\n\");\n cudaEvent_t start_event, stop_event;\n float time_memcpy;\n int eventflags = cudaEventBlockingSync;\n checkCudaErrors(cudaEventCreateWithFlags(&start_event, eventflags));\n checkCudaErrors(cudaEventCreateWithFlags(&stop_event, eventflags));\n\n // P2P memcopy() benchmark\n checkCudaErrors(cudaEventRecord(start_event, 0));\n\n for (int i = 0; i < 100; i++) {\n // With UVA we don't need to specify source and target devices, the\n // runtime figures this out by itself from the pointers\n // Ping-pong copy between GPUs\n if (i % 2 == 0) {\n checkCudaErrors(cudaMemcpy(g1, g0, buf_size, cudaMemcpyDefault));\n }\n else {\n checkCudaErrors(cudaMemcpy(g0, g1, buf_size, cudaMemcpyDefault));\n }\n }\n\n checkCudaErrors(cudaEventRecord(stop_event, 0));\n checkCudaErrors(cudaEventSynchronize(stop_event));\n checkCudaErrors(cudaEventElapsedTime(&time_memcpy, start_event, stop_event));\n printf(\"cudaMemcpyPeer / cudaMemcpy between GPU%d and GPU%d: %.2fGB/s\\n\",\n gpuid[0],\n gpuid[1],\n (1.0f / (time_memcpy / 1000.0f)) * ((100.0f * buf_size)) / 1024.0f / 1024.0f / 1024.0f);\n\n // Prepare host buffer and copy to GPU 0\n printf(\"Preparing host buffer and memcpy to GPU%d...\\n\", gpuid[0]);\n\n for (int i = 0; i < buf_size / sizeof(float); i++) {\n h0[i] = float(i % 4096);\n }\n\n checkCudaErrors(cudaSetDevice(gpuid[0]));\n checkCudaErrors(cudaMemcpy(g0, h0, buf_size, cudaMemcpyDefault));\n\n // Kernel launch configuration\n const dim3 threads(512, 1);\n const dim3 blocks((buf_size / sizeof(float)) / threads.x, 1);\n\n // Run kernel on GPU 1, reading input from the GPU 0 buffer, writing\n // output to the GPU 1 buffer\n printf(\"Run kernel on GPU%d, taking source data from GPU%d and writing to \"\n \"GPU%d...\\n\",\n gpuid[1],\n gpuid[0],\n gpuid[1]);\n checkCudaErrors(cudaSetDevice(gpuid[1]));\n SimpleKernel<<>>(g0, g1);\n\n checkCudaErrors(cudaDeviceSynchronize());\n\n // Run kernel on GPU 0, reading input from the GPU 1 buffer, writing\n // output to the GPU 0 buffer\n printf(\"Run kernel on GPU%d, taking source data from GPU%d and writing to \"\n \"GPU%d...\\n\",\n gpuid[0],\n gpuid[1],\n gpuid[0]);\n checkCudaErrors(cudaSetDevice(gpuid[0]));\n SimpleKernel<<>>(g1, g0);\n\n checkCudaErrors(cudaDeviceSynchronize());\n\n // Copy data back to host and verify\n printf(\"Copy data back to host from GPU%d and verify results...\\n\", gpuid[0]);\n checkCudaErrors(cudaMemcpy(h0, g0, buf_size, cudaMemcpyDefault));\n\n int error_count = 0;\n\n for (int i = 0; i < buf_size / sizeof(float); i++) {\n // Re-generate input data and apply 2x '* 2.0f' computation of both\n // kernel runs\n if (h0[i] != float(i % 4096) * 2.0f * 2.0f) {\n printf(\"Verification error @ element %i: val = %f, ref = %f\\n\", i, h0[i], (float(i % 4096) * 2.0f * 2.0f));\n\n if (error_count++ > 10) {\n break;\n }\n }\n }\n\n // Disable peer access (also unregisters memory for non-UVA cases)\n printf(\"Disabling peer access...\\n\");\n checkCudaErrors(cudaSetDevice(gpuid[0]));\n checkCudaErrors(cudaDeviceDisablePeerAccess(gpuid[1]));\n checkCudaErrors(cudaSetDevice(gpuid[1]));\n checkCudaErrors(cudaDeviceDisablePeerAccess(gpuid[0]));\n\n // Cleanup and shutdown\n printf(\"Shutting down...\\n\");\n checkCudaErrors(cudaEventDestroy(start_event));\n checkCudaErrors(cudaEventDestroy(stop_event));\n checkCudaErrors(cudaSetDevice(gpuid[0]));\n checkCudaErrors(cudaFree(g0));\n checkCudaErrors(cudaSetDevice(gpuid[1]));\n checkCudaErrors(cudaFree(g1));\n checkCudaErrors(cudaFreeHost(h0));\n\n for (int i = 0; i < gpu_n; i++) {\n checkCudaErrors(cudaSetDevice(i));\n }\n\n if (error_count != 0) {\n printf(\"Test failed!\\n\");\n exit(EXIT_FAILURE);\n }\n else {\n printf(\"Test passed\\n\");\n exit(EXIT_SUCCESS);\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates a combination of Peer-to-Peer (P2P) and\n * Unified Virtual Address Space (UVA) features new to SDK 4.0\n */\n\n// includes, system\n#include \n#include \n\n// CUDA includes\n#include \n\n// includes, project\n#include \n#include // helper for shared that are common to CUDA Samples\n\n__global__ void SimpleKernel(float *src, float *dst)\n{\n // Just a dummy kernel, doing enough for us to verify that everything\n // worked\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n dst[idx] = src[idx] * 2.0f;\n}\n\ninline bool IsAppBuiltAs64() { return sizeof(void *) == 8; }\n\nint main(int argc, char **argv)\n{\n printf(\"[%s] - Starting...\\n\", argv[0]);\n\n if (!IsAppBuiltAs64()) {\n printf(\"%s is only supported with on 64-bit OSs and the application must be \"\n \"built as a 64-bit target. Test is being waived.\\n\",\n argv[0]);\n exit(EXIT_WAIVED);\n }\n\n // Number of GPUs\n printf(\"Checking for multiple GPUs...\\n\");\n int gpu_n;\n checkCudaErrors(hipGetDeviceCount(&gpu_n));\n printf(\"CUDA-capable device count: %i\\n\", gpu_n);\n\n if (gpu_n < 2) {\n printf(\"Two or more GPUs with Peer-to-Peer access capability are required for \"\n \"%s.\\n\",\n argv[0]);\n printf(\"Waiving test.\\n\");\n exit(EXIT_WAIVED);\n }\n\n // Query device properties\n hipDeviceProp_t prop[64];\n int gpuid[2]; // we want to find the first two GPU's that can support P2P\n\n for (int i = 0; i < gpu_n; i++) {\n checkCudaErrors(hipGetDeviceProperties(&prop[i], i));\n }\n // Check possibility for peer access\n printf(\"\\nChecking GPU(s) for support of peer to peer memory access...\\n\");\n\n int can_access_peer;\n int p2pCapableGPUs[2]; // We take only 1 pair of P2P capable GPUs\n p2pCapableGPUs[0] = p2pCapableGPUs[1] = -1;\n\n // Show all the combinations of supported P2P GPUs\n for (int i = 0; i < gpu_n; i++) {\n for (int j = 0; j < gpu_n; j++) {\n if (i == j) {\n continue;\n }\n checkCudaErrors(hipDeviceCanAccessPeer(&can_access_peer, i, j));\n printf(\"> Peer access from %s (GPU%d) -> %s (GPU%d) : %s\\n\",\n prop[i].name,\n i,\n prop[j].name,\n j,\n can_access_peer ? \"Yes\" : \"No\");\n if (can_access_peer && p2pCapableGPUs[0] == -1) {\n p2pCapableGPUs[0] = i;\n p2pCapableGPUs[1] = j;\n }\n }\n }\n\n if (p2pCapableGPUs[0] == -1 || p2pCapableGPUs[1] == -1) {\n printf(\"Two or more GPUs with Peer-to-Peer access capability are required for \"\n \"%s.\\n\",\n argv[0]);\n printf(\"Peer to Peer access is not available amongst GPUs in the system, \"\n \"waiving test.\\n\");\n\n exit(EXIT_WAIVED);\n }\n\n // Use first pair of p2p capable GPUs detected.\n gpuid[0] = p2pCapableGPUs[0];\n gpuid[1] = p2pCapableGPUs[1];\n\n // Enable peer access\n printf(\"Enabling peer access between GPU%d and GPU%d...\\n\", gpuid[0], gpuid[1]);\n checkCudaErrors(hipSetDevice(gpuid[0]));\n checkCudaErrors(hipDeviceEnablePeerAccess(gpuid[1], 0));\n checkCudaErrors(hipSetDevice(gpuid[1]));\n checkCudaErrors(hipDeviceEnablePeerAccess(gpuid[0], 0));\n\n // Allocate buffers\n const size_t buf_size = 1024 * 1024 * 16 * sizeof(float);\n printf(\n \"Allocating buffers (%iMB on GPU%d, GPU%d and CPU Host)...\\n\", int(buf_size / 1024 / 1024), gpuid[0], gpuid[1]);\n checkCudaErrors(hipSetDevice(gpuid[0]));\n float *g0;\n checkCudaErrors(hipMalloc(&g0, buf_size));\n checkCudaErrors(hipSetDevice(gpuid[1]));\n float *g1;\n checkCudaErrors(hipMalloc(&g1, buf_size));\n float *h0;\n checkCudaErrors(hipHostMalloc(&h0, buf_size)); // Automatically portable with UVA\n\n // Create CUDA event handles\n printf(\"Creating event handles...\\n\");\n hipEvent_t start_event, stop_event;\n float time_memcpy;\n int eventflags = hipEventBlockingSync;\n checkCudaErrors(hipEventCreateWithFlags(&start_event, eventflags));\n checkCudaErrors(hipEventCreateWithFlags(&stop_event, eventflags));\n\n // P2P memcopy() benchmark\n checkCudaErrors(hipEventRecord(start_event, 0));\n\n for (int i = 0; i < 100; i++) {\n // With UVA we don't need to specify source and target devices, the\n // runtime figures this out by itself from the pointers\n // Ping-pong copy between GPUs\n if (i % 2 == 0) {\n checkCudaErrors(hipMemcpy(g1, g0, buf_size, hipMemcpyDefault));\n }\n else {\n checkCudaErrors(hipMemcpy(g0, g1, buf_size, hipMemcpyDefault));\n }\n }\n\n checkCudaErrors(hipEventRecord(stop_event, 0));\n checkCudaErrors(hipEventSynchronize(stop_event));\n checkCudaErrors(hipEventElapsedTime(&time_memcpy, start_event, stop_event));\n printf(\"hipMemcpyPeer / hipMemcpy between GPU%d and GPU%d: %.2fGB/s\\n\",\n gpuid[0],\n gpuid[1],\n (1.0f / (time_memcpy / 1000.0f)) * ((100.0f * buf_size)) / 1024.0f / 1024.0f / 1024.0f);\n\n // Prepare host buffer and copy to GPU 0\n printf(\"Preparing host buffer and memcpy to GPU%d...\\n\", gpuid[0]);\n\n for (int i = 0; i < buf_size / sizeof(float); i++) {\n h0[i] = float(i % 4096);\n }\n\n checkCudaErrors(hipSetDevice(gpuid[0]));\n checkCudaErrors(hipMemcpy(g0, h0, buf_size, hipMemcpyDefault));\n\n // Kernel launch configuration\n const dim3 threads(512, 1);\n const dim3 blocks((buf_size / sizeof(float)) / threads.x, 1);\n\n // Run kernel on GPU 1, reading input from the GPU 0 buffer, writing\n // output to the GPU 1 buffer\n printf(\"Run kernel on GPU%d, taking source data from GPU%d and writing to \"\n \"GPU%d...\\n\",\n gpuid[1],\n gpuid[0],\n gpuid[1]);\n checkCudaErrors(hipSetDevice(gpuid[1]));\n SimpleKernel<<>>(g0, g1);\n\n checkCudaErrors(hipDeviceSynchronize());\n\n // Run kernel on GPU 0, reading input from the GPU 1 buffer, writing\n // output to the GPU 0 buffer\n printf(\"Run kernel on GPU%d, taking source data from GPU%d and writing to \"\n \"GPU%d...\\n\",\n gpuid[0],\n gpuid[1],\n gpuid[0]);\n checkCudaErrors(hipSetDevice(gpuid[0]));\n SimpleKernel<<>>(g1, g0);\n\n checkCudaErrors(hipDeviceSynchronize());\n\n // Copy data back to host and verify\n printf(\"Copy data back to host from GPU%d and verify results...\\n\", gpuid[0]);\n checkCudaErrors(hipMemcpy(h0, g0, buf_size, hipMemcpyDefault));\n\n int error_count = 0;\n\n for (int i = 0; i < buf_size / sizeof(float); i++) {\n // Re-generate input data and apply 2x '* 2.0f' computation of both\n // kernel runs\n if (h0[i] != float(i % 4096) * 2.0f * 2.0f) {\n printf(\"Verification error @ element %i: val = %f, ref = %f\\n\", i, h0[i], (float(i % 4096) * 2.0f * 2.0f));\n\n if (error_count++ > 10) {\n break;\n }\n }\n }\n\n // Disable peer access (also unregisters memory for non-UVA cases)\n printf(\"Disabling peer access...\\n\");\n checkCudaErrors(hipSetDevice(gpuid[0]));\n checkCudaErrors(hipDeviceDisablePeerAccess(gpuid[1]));\n checkCudaErrors(hipSetDevice(gpuid[1]));\n checkCudaErrors(hipDeviceDisablePeerAccess(gpuid[0]));\n\n // Cleanup and shutdown\n printf(\"Shutting down...\\n\");\n checkCudaErrors(hipEventDestroy(start_event));\n checkCudaErrors(hipEventDestroy(stop_event));\n checkCudaErrors(hipSetDevice(gpuid[0]));\n checkCudaErrors(hipFree(g0));\n checkCudaErrors(hipSetDevice(gpuid[1]));\n checkCudaErrors(hipFree(g1));\n checkCudaErrors(hipHostFree(h0));\n\n for (int i = 0; i < gpu_n; i++) {\n checkCudaErrors(hipSetDevice(i));\n }\n\n if (error_count != 0) {\n printf(\"Test failed!\\n\");\n exit(EXIT_FAILURE);\n }\n else {\n printf(\"Test passed\\n\");\n exit(EXIT_SUCCESS);\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleP2P/simpleP2P.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* Matrix multiplication: C = A * B.\n * Device code.\n */\n\n#ifndef _MATRIXMUL_KERNEL_H_\n#define _MATRIXMUL_KERNEL_H_\n\n#include \n\n#define AS(i, j) As[i][j]\n#define BS(i, j) Bs[i][j]\n\n////////////////////////////////////////////////////////////////////////////////\n//! Matrix multiplication on the device: C = A * B\n//! wA is A's width and wB is B's width\n////////////////////////////////////////////////////////////////////////////////\ntemplate \n__device__ void matrixMul(float *C, float *A, float *B, size_type wA, size_type wB)\n{\n // Block index\n size_type bx = blockIdx.x;\n size_type by = blockIdx.y;\n\n // Thread index\n size_type tx = threadIdx.x;\n size_type ty = threadIdx.y;\n\n // Index of the first sub-matrix of A processed by the block\n size_type aBegin = wA * block_size * by;\n\n // Index of the last sub-matrix of A processed by the block\n size_type aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n size_type aStep = block_size;\n\n // Index of the first sub-matrix of B processed by the block\n size_type bBegin = block_size * bx;\n\n // Step size used to iterate through the sub-matrices of B\n size_type bStep = block_size * wB;\n\n // Csub is used to store the element of the block sub-matrix\n // that is computed by the thread\n float Csub = 0;\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (size_type a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ float As[block_size][block_size];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ float Bs[block_size][block_size];\n\n // Load the matrices from device memory\n // to shared memory; each thread loads\n // one element of each matrix\n AS(ty, tx) = A[a + wA * ty + tx];\n BS(ty, tx) = B[b + wB * ty + tx];\n\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n // Multiply the two matrices together;\n // each thread computes one element\n // of the block sub-matrix\n#pragma unroll\n\n for (size_type k = 0; k < block_size; ++k)\n Csub += AS(ty, k) * BS(k, tx);\n\n // Synchronize to make sure that the preceding\n // computation is done before loading two new\n // sub-matrices of A and B in the next iteration\n __syncthreads();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes one element\n size_type c = wB * block_size * by + block_size * bx;\n C[c + wB * ty + tx] = Csub;\n}\n\n// C wrappers around our template kernel\nextern \"C\" __global__ void matrixMul_bs8_64bit(float *C, float *A, float *B, size_t wA, size_t wB)\n{\n matrixMul<8, size_t>(C, A, B, wA, wB);\n}\nextern \"C\" __global__ void matrixMul_bs16_64bit(float *C, float *A, float *B, size_t wA, size_t wB)\n{\n matrixMul<16, size_t>(C, A, B, wA, wB);\n}\nextern \"C\" __global__ void matrixMul_bs32_64bit(float *C, float *A, float *B, size_t wA, size_t wB)\n{\n matrixMul<32, size_t>(C, A, B, wA, wB);\n}\n\n#endif // #ifndef _MATRIXMUL_KERNEL_H_\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* Matrix multiplication: C = A * B.\n * Device code.\n */\n\n#include \"hip/hip_runtime.h\"\n#ifndef _MATRIXMUL_KERNEL_H_\n#define _MATRIXMUL_KERNEL_H_\n\n#include \n\n#define AS(i, j) As[i][j]\n#define BS(i, j) Bs[i][j]\n\n////////////////////////////////////////////////////////////////////////////////\n//! Matrix multiplication on the device: C = A * B\n//! wA is A's width and wB is B's width\n////////////////////////////////////////////////////////////////////////////////\ntemplate \n__device__ void matrixMul(float *C, float *A, float *B, size_type wA, size_type wB)\n{\n // Block index\n size_type bx = blockIdx.x;\n size_type by = blockIdx.y;\n\n // Thread index\n size_type tx = threadIdx.x;\n size_type ty = threadIdx.y;\n\n // Index of the first sub-matrix of A processed by the block\n size_type aBegin = wA * block_size * by;\n\n // Index of the last sub-matrix of A processed by the block\n size_type aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n size_type aStep = block_size;\n\n // Index of the first sub-matrix of B processed by the block\n size_type bBegin = block_size * bx;\n\n // Step size used to iterate through the sub-matrices of B\n size_type bStep = block_size * wB;\n\n // Csub is used to store the element of the block sub-matrix\n // that is computed by the thread\n float Csub = 0;\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (size_type a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ float As[block_size][block_size];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ float Bs[block_size][block_size];\n\n // Load the matrices from device memory\n // to shared memory; each thread loads\n // one element of each matrix\n AS(ty, tx) = A[a + wA * ty + tx];\n BS(ty, tx) = B[b + wB * ty + tx];\n\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n // Multiply the two matrices together;\n // each thread computes one element\n // of the block sub-matrix\n#pragma unroll\n\n for (size_type k = 0; k < block_size; ++k)\n Csub += AS(ty, k) * BS(k, tx);\n\n // Synchronize to make sure that the preceding\n // computation is done before loading two new\n // sub-matrices of A and B in the next iteration\n __syncthreads();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes one element\n size_type c = wB * block_size * by + block_size * bx;\n C[c + wB * ty + tx] = Csub;\n}\n\n// C wrappers around our template kernel\nextern \"C\" __global__ void matrixMul_bs8_64bit(float *C, float *A, float *B, size_t wA, size_t wB)\n{\n matrixMul<8, size_t>(C, A, B, wA, wB);\n}\nextern \"C\" __global__ void matrixMul_bs16_64bit(float *C, float *A, float *B, size_t wA, size_t wB)\n{\n matrixMul<16, size_t>(C, A, B, wA, wB);\n}\nextern \"C\" __global__ void matrixMul_bs32_64bit(float *C, float *A, float *B, size_t wA, size_t wB)\n{\n matrixMul<32, size_t>(C, A, B, wA, wB);\n}\n\n#endif // #ifndef _MATRIXMUL_KERNEL_H_\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/matrixMulDrv/matrixMul_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This application demonstrates how to use the CUDA API to use multiple GPUs,\n * with an emphasis on simple illustration of the techniques (not on\n * performance).\n *\n * Note that in order to detect multiple GPUs in your system you have to disable\n * SLI in the nvidia control panel. Otherwise only one GPU is visible to the\n * application. On the other side, you can still extend your desktop to screens\n * attached to both GPUs.\n */\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#ifndef MAX\n#define MAX(a, b) (a > b ? a : b)\n#endif\n\n#include \"simpleMultiGPU.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Data configuration\n////////////////////////////////////////////////////////////////////////////////\nconst int MAX_GPU_COUNT = 32;\nconst int DATA_N = 1048576 * 32;\n\n////////////////////////////////////////////////////////////////////////////////\n// Simple reduction kernel.\n// Refer to the 'reduction' CUDA Sample describing\n// reduction optimization strategies\n////////////////////////////////////////////////////////////////////////////////\n__global__ static void reduceKernel(float *d_Result, float *d_Input, int N)\n{\n const int tid = blockIdx.x * blockDim.x + threadIdx.x;\n const int threadN = gridDim.x * blockDim.x;\n float sum = 0;\n\n for (int pos = tid; pos < N; pos += threadN)\n sum += d_Input[pos];\n\n d_Result[tid] = sum;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n // Solver config\n TGPUplan plan[MAX_GPU_COUNT];\n\n // GPU reduction results\n float h_SumGPU[MAX_GPU_COUNT];\n\n float sumGPU;\n double sumCPU, diff;\n\n int i, j, gpuBase, GPU_N;\n\n const int BLOCK_N = 32;\n const int THREAD_N = 256;\n const int ACCUM_N = BLOCK_N * THREAD_N;\n\n printf(\"Starting simpleMultiGPU\\n\");\n checkCudaErrors(cudaGetDeviceCount(&GPU_N));\n\n if (GPU_N > MAX_GPU_COUNT) {\n GPU_N = MAX_GPU_COUNT;\n }\n\n printf(\"CUDA-capable device count: %i\\n\", GPU_N);\n\n printf(\"Generating input data...\\n\\n\");\n\n // Subdividing input data across GPUs\n // Get data sizes for each GPU\n for (i = 0; i < GPU_N; i++) {\n plan[i].dataN = DATA_N / GPU_N;\n }\n\n // Take into account \"odd\" data sizes\n for (i = 0; i < DATA_N % GPU_N; i++) {\n plan[i].dataN++;\n }\n\n // Assign data ranges to GPUs\n gpuBase = 0;\n\n for (i = 0; i < GPU_N; i++) {\n plan[i].h_Sum = h_SumGPU + i;\n gpuBase += plan[i].dataN;\n }\n\n // Create streams for issuing GPU command asynchronously and allocate memory\n // (GPU and System page-locked)\n for (i = 0; i < GPU_N; i++) {\n checkCudaErrors(cudaSetDevice(i));\n checkCudaErrors(cudaStreamCreate(&plan[i].stream));\n // Allocate memory\n checkCudaErrors(cudaMalloc((void **)&plan[i].d_Data, plan[i].dataN * sizeof(float)));\n checkCudaErrors(cudaMalloc((void **)&plan[i].d_Sum, ACCUM_N * sizeof(float)));\n checkCudaErrors(cudaMallocHost((void **)&plan[i].h_Sum_from_device, ACCUM_N * sizeof(float)));\n checkCudaErrors(cudaMallocHost((void **)&plan[i].h_Data, plan[i].dataN * sizeof(float)));\n\n for (j = 0; j < plan[i].dataN; j++) {\n plan[i].h_Data[j] = (float)rand() / (float)RAND_MAX;\n }\n }\n\n // Start timing and compute on GPU(s)\n printf(\"Computing with %d GPUs...\\n\", GPU_N);\n // create and start timer\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n\n // start the timer\n sdkStartTimer(&timer);\n\n // Copy data to GPU, launch the kernel and copy data back. All asynchronously\n for (i = 0; i < GPU_N; i++) {\n // Set device\n checkCudaErrors(cudaSetDevice(i));\n\n // Copy input data from CPU\n checkCudaErrors(cudaMemcpyAsync(\n plan[i].d_Data, plan[i].h_Data, plan[i].dataN * sizeof(float), cudaMemcpyHostToDevice, plan[i].stream));\n\n // Perform GPU computations\n reduceKernel<<>>(plan[i].d_Sum, plan[i].d_Data, plan[i].dataN);\n getLastCudaError(\"reduceKernel() execution failed.\\n\");\n\n // Read back GPU results\n checkCudaErrors(cudaMemcpyAsync(\n plan[i].h_Sum_from_device, plan[i].d_Sum, ACCUM_N * sizeof(float), cudaMemcpyDeviceToHost, plan[i].stream));\n }\n\n // Process GPU results\n for (i = 0; i < GPU_N; i++) {\n float sum;\n\n // Set device\n checkCudaErrors(cudaSetDevice(i));\n\n // Wait for all operations to finish\n cudaStreamSynchronize(plan[i].stream);\n\n // Finalize GPU reduction for current subvector\n sum = 0;\n\n for (j = 0; j < ACCUM_N; j++) {\n sum += plan[i].h_Sum_from_device[j];\n }\n\n *(plan[i].h_Sum) = (float)sum;\n\n // Shut down this GPU\n checkCudaErrors(cudaFreeHost(plan[i].h_Sum_from_device));\n checkCudaErrors(cudaFree(plan[i].d_Sum));\n checkCudaErrors(cudaFree(plan[i].d_Data));\n checkCudaErrors(cudaStreamDestroy(plan[i].stream));\n }\n\n sumGPU = 0;\n\n for (i = 0; i < GPU_N; i++) {\n sumGPU += h_SumGPU[i];\n }\n\n sdkStopTimer(&timer);\n printf(\" GPU Processing time: %f (ms)\\n\\n\", sdkGetTimerValue(&timer));\n sdkDeleteTimer(&timer);\n\n // Compute on Host CPU\n printf(\"Computing with Host CPU...\\n\\n\");\n\n sumCPU = 0;\n\n for (i = 0; i < GPU_N; i++) {\n for (j = 0; j < plan[i].dataN; j++) {\n sumCPU += plan[i].h_Data[j];\n }\n }\n\n // Compare GPU and CPU results\n printf(\"Comparing GPU and Host CPU results...\\n\");\n diff = fabs(sumCPU - sumGPU) / fabs(sumCPU);\n printf(\" GPU sum: %f\\n CPU sum: %f\\n\", sumGPU, sumCPU);\n printf(\" Relative difference: %E \\n\\n\", diff);\n\n // Cleanup and shutdown\n for (i = 0; i < GPU_N; i++) {\n checkCudaErrors(cudaSetDevice(i));\n checkCudaErrors(cudaFreeHost(plan[i].h_Data));\n }\n\n exit((diff < 1e-5) ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This application demonstrates how to use the CUDA API to use multiple GPUs,\n * with an emphasis on simple illustration of the techniques (not on\n * performance).\n *\n * Note that in order to detect multiple GPUs in your system you have to disable\n * SLI in the nvidia control panel. Otherwise only one GPU is visible to the\n * application. On the other side, you can still extend your desktop to screens\n * attached to both GPUs.\n */\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#ifndef MAX\n#define MAX(a, b) (a > b ? a : b)\n#endif\n\n#include \"simpleMultiGPU.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Data configuration\n////////////////////////////////////////////////////////////////////////////////\nconst int MAX_GPU_COUNT = 32;\nconst int DATA_N = 1048576 * 32;\n\n////////////////////////////////////////////////////////////////////////////////\n// Simple reduction kernel.\n// Refer to the 'reduction' CUDA Sample describing\n// reduction optimization strategies\n////////////////////////////////////////////////////////////////////////////////\n__global__ static void reduceKernel(float *d_Result, float *d_Input, int N)\n{\n const int tid = blockIdx.x * blockDim.x + threadIdx.x;\n const int threadN = gridDim.x * blockDim.x;\n float sum = 0;\n\n for (int pos = tid; pos < N; pos += threadN)\n sum += d_Input[pos];\n\n d_Result[tid] = sum;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n // Solver config\n TGPUplan plan[MAX_GPU_COUNT];\n\n // GPU reduction results\n float h_SumGPU[MAX_GPU_COUNT];\n\n float sumGPU;\n double sumCPU, diff;\n\n int i, j, gpuBase, GPU_N;\n\n const int BLOCK_N = 32;\n const int THREAD_N = 256;\n const int ACCUM_N = BLOCK_N * THREAD_N;\n\n printf(\"Starting simpleMultiGPU\\n\");\n checkCudaErrors(hipGetDeviceCount(&GPU_N));\n\n if (GPU_N > MAX_GPU_COUNT) {\n GPU_N = MAX_GPU_COUNT;\n }\n\n printf(\"CUDA-capable device count: %i\\n\", GPU_N);\n\n printf(\"Generating input data...\\n\\n\");\n\n // Subdividing input data across GPUs\n // Get data sizes for each GPU\n for (i = 0; i < GPU_N; i++) {\n plan[i].dataN = DATA_N / GPU_N;\n }\n\n // Take into account \"odd\" data sizes\n for (i = 0; i < DATA_N % GPU_N; i++) {\n plan[i].dataN++;\n }\n\n // Assign data ranges to GPUs\n gpuBase = 0;\n\n for (i = 0; i < GPU_N; i++) {\n plan[i].h_Sum = h_SumGPU + i;\n gpuBase += plan[i].dataN;\n }\n\n // Create streams for issuing GPU command asynchronously and allocate memory\n // (GPU and System page-locked)\n for (i = 0; i < GPU_N; i++) {\n checkCudaErrors(hipSetDevice(i));\n checkCudaErrors(hipStreamCreate(&plan[i].stream));\n // Allocate memory\n checkCudaErrors(hipMalloc((void **)&plan[i].d_Data, plan[i].dataN * sizeof(float)));\n checkCudaErrors(hipMalloc((void **)&plan[i].d_Sum, ACCUM_N * sizeof(float)));\n checkCudaErrors(hipHostMalloc((void **)&plan[i].h_Sum_from_device, ACCUM_N * sizeof(float)));\n checkCudaErrors(hipHostMalloc((void **)&plan[i].h_Data, plan[i].dataN * sizeof(float)));\n\n for (j = 0; j < plan[i].dataN; j++) {\n plan[i].h_Data[j] = (float)rand() / (float)RAND_MAX;\n }\n }\n\n // Start timing and compute on GPU(s)\n printf(\"Computing with %d GPUs...\\n\", GPU_N);\n // create and start timer\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n\n // start the timer\n sdkStartTimer(&timer);\n\n // Copy data to GPU, launch the kernel and copy data back. All asynchronously\n for (i = 0; i < GPU_N; i++) {\n // Set device\n checkCudaErrors(hipSetDevice(i));\n\n // Copy input data from CPU\n checkCudaErrors(hipMemcpyAsync(\n plan[i].d_Data, plan[i].h_Data, plan[i].dataN * sizeof(float), hipMemcpyHostToDevice, plan[i].stream));\n\n // Perform GPU computations\n reduceKernel<<>>(plan[i].d_Sum, plan[i].d_Data, plan[i].dataN);\n getLastCudaError(\"reduceKernel() execution failed.\\n\");\n\n // Read back GPU results\n checkCudaErrors(hipMemcpyAsync(\n plan[i].h_Sum_from_device, plan[i].d_Sum, ACCUM_N * sizeof(float), hipMemcpyDeviceToHost, plan[i].stream));\n }\n\n // Process GPU results\n for (i = 0; i < GPU_N; i++) {\n float sum;\n\n // Set device\n checkCudaErrors(hipSetDevice(i));\n\n // Wait for all operations to finish\n hipStreamSynchronize(plan[i].stream);\n\n // Finalize GPU reduction for current subvector\n sum = 0;\n\n for (j = 0; j < ACCUM_N; j++) {\n sum += plan[i].h_Sum_from_device[j];\n }\n\n *(plan[i].h_Sum) = (float)sum;\n\n // Shut down this GPU\n checkCudaErrors(hipHostFree(plan[i].h_Sum_from_device));\n checkCudaErrors(hipFree(plan[i].d_Sum));\n checkCudaErrors(hipFree(plan[i].d_Data));\n checkCudaErrors(hipStreamDestroy(plan[i].stream));\n }\n\n sumGPU = 0;\n\n for (i = 0; i < GPU_N; i++) {\n sumGPU += h_SumGPU[i];\n }\n\n sdkStopTimer(&timer);\n printf(\" GPU Processing time: %f (ms)\\n\\n\", sdkGetTimerValue(&timer));\n sdkDeleteTimer(&timer);\n\n // Compute on Host CPU\n printf(\"Computing with Host CPU...\\n\\n\");\n\n sumCPU = 0;\n\n for (i = 0; i < GPU_N; i++) {\n for (j = 0; j < plan[i].dataN; j++) {\n sumCPU += plan[i].h_Data[j];\n }\n }\n\n // Compare GPU and CPU results\n printf(\"Comparing GPU and Host CPU results...\\n\");\n diff = fabs(sumCPU - sumGPU) / fabs(sumCPU);\n printf(\" GPU sum: %f\\n CPU sum: %f\\n\", sumGPU, sumCPU);\n printf(\" Relative difference: %E \\n\\n\", diff);\n\n // Cleanup and shutdown\n for (i = 0; i < GPU_N; i++) {\n checkCudaErrors(hipSetDevice(i));\n checkCudaErrors(hipHostFree(plan[i].h_Data));\n }\n\n exit((diff < 1e-5) ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleMultiGPU/simpleMultiGPU.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#ifndef MAX\n#define MAX(a, b) (a > b ? a : b)\n#endif\n\n__global__ void testKernel(int val)\n{\n printf(\"[%d, %d]:\\t\\tValue is:%d\\n\",\n blockIdx.y * gridDim.x + blockIdx.x,\n threadIdx.z * blockDim.x * blockDim.y + threadIdx.y * blockDim.x + threadIdx.x,\n val);\n}\n\nint main(int argc, char **argv)\n{\n int devID;\n cudaDeviceProp props;\n\n // This will pick the best possible CUDA capable device\n devID = findCudaDevice(argc, (const char **)argv);\n\n // Get GPU information\n checkCudaErrors(cudaGetDevice(&devID));\n checkCudaErrors(cudaGetDeviceProperties(&props, devID));\n printf(\"Device %d: \\\"%s\\\" with Compute capability %d.%d\\n\", devID, props.name, props.major, props.minor);\n\n printf(\"printf() is called. Output:\\n\\n\");\n\n // Kernel configuration, where a two-dimensional grid and\n // three-dimensional blocks are configured.\n dim3 dimGrid(2, 2);\n dim3 dimBlock(2, 2, 2);\n testKernel<<>>(10);\n cudaDeviceSynchronize();\n\n return EXIT_SUCCESS;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#ifndef MAX\n#define MAX(a, b) (a > b ? a : b)\n#endif\n\n__global__ void testKernel(int val)\n{\n printf(\"[%d, %d]:\\t\\tValue is:%d\\n\",\n blockIdx.y * gridDim.x + blockIdx.x,\n threadIdx.z * blockDim.x * blockDim.y + threadIdx.y * blockDim.x + threadIdx.x,\n val);\n}\n\nint main(int argc, char **argv)\n{\n int devID;\n hipDeviceProp_t props;\n\n // This will pick the best possible CUDA capable device\n devID = findCudaDevice(argc, (const char **)argv);\n\n // Get GPU information\n checkCudaErrors(hipGetDevice(&devID));\n checkCudaErrors(hipGetDeviceProperties(&props, devID));\n printf(\"Device %d: \\\"%s\\\" with Compute capability %d.%d\\n\", devID, props.name, props.major, props.minor);\n\n printf(\"printf() is called. Output:\\n\\n\");\n\n // Kernel configuration, where a two-dimensional grid and\n // three-dimensional blocks are configured.\n dim3 dimGrid(2, 2);\n dim3 dimBlock(2, 2, 2);\n testKernel<<>>(10);\n hipDeviceSynchronize();\n\n return EXIT_SUCCESS;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simplePrintf/simplePrintf.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates how to use texture fetches from layered 2D textures\n * in CUDA C\n *\n * This sample first generates a 3D input data array for the layered texture\n * and the expected output. Then it starts CUDA C kernels, one for each layer,\n * which fetch their layer's texture data (using normalized texture coordinates)\n * transform it to the expected output, and write it to a 3D output data array.\n */\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// includes CUDA\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\nstatic const char *sSDKname = \"simpleCubemapTexture\";\n\n// includes, kernels\n\n////////////////////////////////////////////////////////////////////////////////\n//! Transform a cubemap face of a linear buffe using cubemap texture lookups\n//! @param g_odata output data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void transformKernel(float *g_odata, int width, cudaTextureObject_t tex)\n{\n // calculate this thread's data point\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // 0.5f offset and division are necessary to access the original data points\n // in the texture (such that bilinear interpolation will not be activated).\n // For details, see also CUDA Programming Guide, Appendix D\n\n float u = ((x + 0.5f) / (float)width) * 2.f - 1.f;\n float v = ((y + 0.5f) / (float)width) * 2.f - 1.f;\n\n float cx, cy, cz;\n\n for (unsigned int face = 0; face < 6; face++) {\n // Layer 0 is positive X face\n if (face == 0) {\n cx = 1;\n cy = -v;\n cz = -u;\n }\n // Layer 1 is negative X face\n else if (face == 1) {\n cx = -1;\n cy = -v;\n cz = u;\n }\n // Layer 2 is positive Y face\n else if (face == 2) {\n cx = u;\n cy = 1;\n cz = v;\n }\n // Layer 3 is negative Y face\n else if (face == 3) {\n cx = u;\n cy = -1;\n cz = -v;\n }\n // Layer 4 is positive Z face\n else if (face == 4) {\n cx = u;\n cy = -v;\n cz = 1;\n }\n // Layer 4 is negative Z face\n else if (face == 5) {\n cx = -u;\n cy = -v;\n cz = -1;\n }\n\n // read from texture, do expected transformation and write to global memory\n g_odata[face * width * width + y * width + x] = -texCubemap(tex, cx, cy, cz);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n int devID = findCudaDevice(argc, (const char **)argv);\n\n bool bResult = true;\n\n // get number of SMs on this GPU\n cudaDeviceProp deviceProps;\n\n checkCudaErrors(cudaGetDeviceProperties(&deviceProps, devID));\n printf(\"CUDA device [%s] has %d Multi-Processors \", deviceProps.name, deviceProps.multiProcessorCount);\n printf(\"SM %d.%d\\n\", deviceProps.major, deviceProps.minor);\n\n if (deviceProps.major < 2) {\n printf(\"%s requires SM 2.0 or higher for support of Texture Arrays. Test \"\n \"will exit... \\n\",\n sSDKname);\n\n exit(EXIT_WAIVED);\n }\n\n // generate input data for layered texture\n unsigned int width = 64, num_faces = 6, num_layers = 1;\n unsigned int cubemap_size = width * width * num_faces;\n unsigned int size = cubemap_size * num_layers * sizeof(float);\n float *h_data = (float *)malloc(size);\n\n for (int i = 0; i < (int)(cubemap_size * num_layers); i++) {\n h_data[i] = (float)i;\n }\n\n // this is the expected transformation of the input data (the expected output)\n float *h_data_ref = (float *)malloc(size);\n\n for (unsigned int layer = 0; layer < num_layers; layer++) {\n for (int i = 0; i < (int)(cubemap_size); i++) {\n h_data_ref[layer * cubemap_size + i] = -h_data[layer * cubemap_size + i] + layer;\n }\n }\n\n // allocate device memory for result\n float *d_data = NULL;\n checkCudaErrors(cudaMalloc((void **)&d_data, size));\n\n // allocate array and copy image data\n cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc(32, 0, 0, 0, cudaChannelFormatKindFloat);\n cudaArray *cu_3darray;\n // checkCudaErrors(cudaMalloc3DArray( &cu_3darray, &channelDesc,\n // make_cudaExtent(width, height, num_layers), cudaArrayLayered ));\n checkCudaErrors(\n cudaMalloc3DArray(&cu_3darray, &channelDesc, make_cudaExtent(width, width, num_faces), cudaArrayCubemap));\n cudaMemcpy3DParms myparms = {0};\n myparms.srcPos = make_cudaPos(0, 0, 0);\n myparms.dstPos = make_cudaPos(0, 0, 0);\n myparms.srcPtr = make_cudaPitchedPtr(h_data, width * sizeof(float), width, width);\n myparms.dstArray = cu_3darray;\n myparms.extent = make_cudaExtent(width, width, num_faces);\n myparms.kind = cudaMemcpyHostToDevice;\n checkCudaErrors(cudaMemcpy3D(&myparms));\n\n cudaTextureObject_t tex;\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = cu_3darray;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = cudaFilterModeLinear;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.addressMode[2] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n\n dim3 dimBlock(8, 8, 1);\n dim3 dimGrid(width / dimBlock.x, width / dimBlock.y, 1);\n\n printf(\"Covering Cubemap data array of %d~3 x %d: Grid size is %d x %d, each \"\n \"block has 8 x 8 threads\\n\",\n width,\n num_layers,\n dimGrid.x,\n dimGrid.y);\n\n transformKernel<<>>(d_data, width,\n tex); // warmup (for better timing)\n\n // check if kernel execution generated an error\n getLastCudaError(\"warmup Kernel execution failed\");\n\n checkCudaErrors(cudaDeviceSynchronize());\n\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n sdkStartTimer(&timer);\n\n // execute the kernel\n transformKernel<<>>(d_data, width, tex);\n\n // check if kernel execution generated an error\n getLastCudaError(\"Kernel execution failed\");\n\n checkCudaErrors(cudaDeviceSynchronize());\n sdkStopTimer(&timer);\n printf(\"Processing time: %.3f msec\\n\", sdkGetTimerValue(&timer));\n printf(\"%.2f Mtexlookups/sec\\n\", (cubemap_size / (sdkGetTimerValue(&timer) / 1000.0f) / 1e6));\n sdkDeleteTimer(&timer);\n\n // allocate mem for the result on host side\n float *h_odata = (float *)malloc(size);\n // copy result from device to host\n checkCudaErrors(cudaMemcpy(h_odata, d_data, size, cudaMemcpyDeviceToHost));\n\n // write regression file if necessary\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // write file for regression test\n sdkWriteFile(\"./data/regression.dat\", h_odata, width * width, 0.0f, false);\n }\n else {\n printf(\"Comparing kernel output to expected data\\n\");\n\n#define MIN_EPSILON_ERROR 5e-3f\n bResult = compareData(h_odata, h_data_ref, cubemap_size, MIN_EPSILON_ERROR, 0.0f);\n }\n\n // cleanup memory\n free(h_data);\n free(h_data_ref);\n free(h_odata);\n\n checkCudaErrors(cudaDestroyTextureObject(tex));\n checkCudaErrors(cudaFree(d_data));\n checkCudaErrors(cudaFreeArray(cu_3darray));\n\n exit(bResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates how to use texture fetches from layered 2D textures\n * in CUDA C\n *\n * This sample first generates a 3D input data array for the layered texture\n * and the expected output. Then it starts CUDA C kernels, one for each layer,\n * which fetch their layer's texture data (using normalized texture coordinates)\n * transform it to the expected output, and write it to a 3D output data array.\n */\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// includes CUDA\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\nstatic const char *sSDKname = \"simpleCubemapTexture\";\n\n// includes, kernels\n\n////////////////////////////////////////////////////////////////////////////////\n//! Transform a cubemap face of a linear buffe using cubemap texture lookups\n//! @param g_odata output data in global memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void transformKernel(float *g_odata, int width, hipTextureObject_t tex)\n{\n // calculate this thread's data point\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // 0.5f offset and division are necessary to access the original data points\n // in the texture (such that bilinear interpolation will not be activated).\n // For details, see also CUDA Programming Guide, Appendix D\n\n float u = ((x + 0.5f) / (float)width) * 2.f - 1.f;\n float v = ((y + 0.5f) / (float)width) * 2.f - 1.f;\n\n float cx, cy, cz;\n\n for (unsigned int face = 0; face < 6; face++) {\n // Layer 0 is positive X face\n if (face == 0) {\n cx = 1;\n cy = -v;\n cz = -u;\n }\n // Layer 1 is negative X face\n else if (face == 1) {\n cx = -1;\n cy = -v;\n cz = u;\n }\n // Layer 2 is positive Y face\n else if (face == 2) {\n cx = u;\n cy = 1;\n cz = v;\n }\n // Layer 3 is negative Y face\n else if (face == 3) {\n cx = u;\n cy = -1;\n cz = -v;\n }\n // Layer 4 is positive Z face\n else if (face == 4) {\n cx = u;\n cy = -v;\n cz = 1;\n }\n // Layer 4 is negative Z face\n else if (face == 5) {\n cx = -u;\n cy = -v;\n cz = -1;\n }\n\n // read from texture, do expected transformation and write to global memory\n g_odata[face * width * width + y * width + x] = -texCubemap(tex, cx, cy, cz);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n int devID = findCudaDevice(argc, (const char **)argv);\n\n bool bResult = true;\n\n // get number of SMs on this GPU\n hipDeviceProp_t deviceProps;\n\n checkCudaErrors(hipGetDeviceProperties(&deviceProps, devID));\n printf(\"CUDA device [%s] has %d Multi-Processors \", deviceProps.name, deviceProps.multiProcessorCount);\n printf(\"SM %d.%d\\n\", deviceProps.major, deviceProps.minor);\n\n if (deviceProps.major < 2) {\n printf(\"%s requires SM 2.0 or higher for support of Texture Arrays. Test \"\n \"will exit... \\n\",\n sSDKname);\n\n exit(EXIT_WAIVED);\n }\n\n // generate input data for layered texture\n unsigned int width = 64, num_faces = 6, num_layers = 1;\n unsigned int cubemap_size = width * width * num_faces;\n unsigned int size = cubemap_size * num_layers * sizeof(float);\n float *h_data = (float *)malloc(size);\n\n for (int i = 0; i < (int)(cubemap_size * num_layers); i++) {\n h_data[i] = (float)i;\n }\n\n // this is the expected transformation of the input data (the expected output)\n float *h_data_ref = (float *)malloc(size);\n\n for (unsigned int layer = 0; layer < num_layers; layer++) {\n for (int i = 0; i < (int)(cubemap_size); i++) {\n h_data_ref[layer * cubemap_size + i] = -h_data[layer * cubemap_size + i] + layer;\n }\n }\n\n // allocate device memory for result\n float *d_data = NULL;\n checkCudaErrors(hipMalloc((void **)&d_data, size));\n\n // allocate array and copy image data\n hipChannelFormatDesc channelDesc = hipCreateChannelDesc(32, 0, 0, 0, hipChannelFormatKindFloat);\n hipArray *cu_3darray;\n // checkCudaErrors(cudaMalloc3DArray( &cu_3darray, &channelDesc,\n // make_cudaExtent(width, height, num_layers), cudaArrayLayered ));\n checkCudaErrors(\n hipMalloc3DArray(&cu_3darray, &channelDesc, make_hipExtent(width, width, num_faces), hipArrayCubemap));\n hipMemcpy3DParms myparms = {0};\n myparms.srcPos = make_hipPos(0, 0, 0);\n myparms.dstPos = make_hipPos(0, 0, 0);\n myparms.srcPtr = make_hipPitchedPtr(h_data, width * sizeof(float), width, width);\n myparms.dstArray = cu_3darray;\n myparms.extent = make_hipExtent(width, width, num_faces);\n myparms.kind = hipMemcpyHostToDevice;\n checkCudaErrors(hipMemcpy3D(&myparms));\n\n hipTextureObject_t tex;\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = cu_3darray;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = hipFilterModeLinear;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.addressMode[2] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n\n dim3 dimBlock(8, 8, 1);\n dim3 dimGrid(width / dimBlock.x, width / dimBlock.y, 1);\n\n printf(\"Covering Cubemap data array of %d~3 x %d: Grid size is %d x %d, each \"\n \"block has 8 x 8 threads\\n\",\n width,\n num_layers,\n dimGrid.x,\n dimGrid.y);\n\n transformKernel<<>>(d_data, width,\n tex); // warmup (for better timing)\n\n // check if kernel execution generated an error\n getLastCudaError(\"warmup Kernel execution failed\");\n\n checkCudaErrors(hipDeviceSynchronize());\n\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n sdkStartTimer(&timer);\n\n // execute the kernel\n transformKernel<<>>(d_data, width, tex);\n\n // check if kernel execution generated an error\n getLastCudaError(\"Kernel execution failed\");\n\n checkCudaErrors(hipDeviceSynchronize());\n sdkStopTimer(&timer);\n printf(\"Processing time: %.3f msec\\n\", sdkGetTimerValue(&timer));\n printf(\"%.2f Mtexlookups/sec\\n\", (cubemap_size / (sdkGetTimerValue(&timer) / 1000.0f) / 1e6));\n sdkDeleteTimer(&timer);\n\n // allocate mem for the result on host side\n float *h_odata = (float *)malloc(size);\n // copy result from device to host\n checkCudaErrors(hipMemcpy(h_odata, d_data, size, hipMemcpyDeviceToHost));\n\n // write regression file if necessary\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // write file for regression test\n sdkWriteFile(\"./data/regression.dat\", h_odata, width * width, 0.0f, false);\n }\n else {\n printf(\"Comparing kernel output to expected data\\n\");\n\n#define MIN_EPSILON_ERROR 5e-3f\n bResult = compareData(h_odata, h_data_ref, cubemap_size, MIN_EPSILON_ERROR, 0.0f);\n }\n\n // cleanup memory\n free(h_data);\n free(h_data_ref);\n free(h_odata);\n\n checkCudaErrors(hipDestroyTextureObject(tex));\n checkCudaErrors(hipFree(d_data));\n checkCudaErrors(hipFreeArray(cu_3darray));\n\n exit(bResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleCubemapTexture/simpleCubemapTexture.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Utilities and system includes\n\n#include \n\n// clamp x to range [a, b]\n__device__ float clamp(float x, float a, float b) { return max(a, min(b, x)); }\n\n__device__ int clamp(int x, int a, int b) { return max(a, min(b, x)); }\n\n// convert floating point rgb color to 8-bit integer\n__device__ int rgbToInt(float r, float g, float b)\n{\n r = clamp(r, 0.0f, 255.0f);\n g = clamp(g, 0.0f, 255.0f);\n b = clamp(b, 0.0f, 255.0f);\n return (int(b) << 16) | (int(g) << 8) | int(r);\n}\n\n__global__ void cudaProcess(unsigned int *g_odata, int imgw)\n{\n extern __shared__ uchar4 sdata[];\n\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n int bw = blockDim.x;\n int bh = blockDim.y;\n int x = blockIdx.x * bw + tx;\n int y = blockIdx.y * bh + ty;\n\n uchar4 c4 = make_uchar4((x & 0x20) ? 100 : 0, 0, (y & 0x20) ? 100 : 0, 0);\n g_odata[y * imgw + x] = rgbToInt(c4.z, c4.y, c4.x);\n}\n\nextern \"C\" void launch_cudaProcess(dim3 grid, dim3 block, int sbytes, unsigned int *g_odata, int imgw)\n{\n cudaProcess<<>>(g_odata, imgw);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Utilities and system includes\n\n#include \"hip/hip_runtime.h\"\n#include \n\n// clamp x to range [a, b]\n__device__ float clamp(float x, float a, float b) { return max(a, min(b, x)); }\n\n__device__ int clamp(int x, int a, int b) { return max(a, min(b, x)); }\n\n// convert floating point rgb color to 8-bit integer\n__device__ int rgbToInt(float r, float g, float b)\n{\n r = clamp(r, 0.0f, 255.0f);\n g = clamp(g, 0.0f, 255.0f);\n b = clamp(b, 0.0f, 255.0f);\n return (int(b) << 16) | (int(g) << 8) | int(r);\n}\n\n__global__ void cudaProcess(unsigned int *g_odata, int imgw)\n{\n extern __shared__ uchar4 sdata[];\n\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n int bw = blockDim.x;\n int bh = blockDim.y;\n int x = blockIdx.x * bw + tx;\n int y = blockIdx.y * bh + ty;\n\n uchar4 c4 = make_uchar4((x & 0x20) ? 100 : 0, 0, (y & 0x20) ? 100 : 0, 0);\n g_odata[y * imgw + x] = rgbToInt(c4.z, c4.y, c4.x);\n}\n\nextern \"C\" void launch_cudaProcess(dim3 grid, dim3 block, int sbytes, unsigned int *g_odata, int imgw)\n{\n cudaProcess<<>>(g_odata, imgw);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleCUDA2GL/simpleCUDA2GL.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* This sample is a templatized version of the template project.\n * It also shows how to correctly templatize dynamically allocated shared\n * memory arrays.\n * Host code.\n */\n\n// System includes\n#include \n#include \n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#ifndef MAX\n#define MAX(a, b) (a > b ? a : b)\n#endif\n\n// includes, kernels\n#include \"sharedmem.cuh\"\n\nint g_TotalFailures = 0;\n\n////////////////////////////////////////////////////////////////////////////////\n//! Simple test kernel for device functionality\n//! @param g_idata input data in global memory\n//! @param g_odata output data in global memory\n////////////////////////////////////////////////////////////////////////////////\ntemplate __global__ void testKernel(T *g_idata, T *g_odata)\n{\n // Shared mem size is determined by the host app at run time\n SharedMemory smem;\n T *sdata = smem.getPointer();\n\n // access thread id\n const unsigned int tid = threadIdx.x;\n // access number of threads in this block\n const unsigned int num_threads = blockDim.x;\n\n // read in input data from global memory\n sdata[tid] = g_idata[tid];\n __syncthreads();\n\n // perform some computations\n sdata[tid] = (T)num_threads * sdata[tid];\n __syncthreads();\n\n // write data to global memory\n g_odata[tid] = sdata[tid];\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\ntemplate void runTest(int argc, char **argv, int len);\n\ntemplate void computeGold(T *reference, T *idata, const unsigned int len)\n{\n const T T_len = static_cast(len);\n\n for (unsigned int i = 0; i < len; ++i) {\n reference[i] = idata[i] * T_len;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"> runTest\\n\");\n runTest(argc, argv, 32);\n printf(\"> runTest\\n\");\n runTest(argc, argv, 64);\n\n printf(\"\\n[simpleTemplates] -> Test Results: %d Failures\\n\", g_TotalFailures);\n\n exit(g_TotalFailures == 0 ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\n// To completely templatize runTest (below) with cutil, we need to use\n// template specialization to wrap up CUTIL's array comparison and file writing\n// functions for different types.\n\n// Here's the generic wrapper for cutCompare*\ntemplate class ArrayComparator\n{\npublic:\n bool compare(const T *reference, T *data, unsigned int len)\n {\n fprintf(stderr, \"Error: no comparison function implemented for this type\\n\");\n return false;\n }\n};\n\n// Here's the specialization for ints:\ntemplate <> class ArrayComparator\n{\npublic:\n bool compare(const int *reference, int *data, unsigned int len)\n {\n return compareData(reference, data, len, 0.15f, 0.0f);\n }\n};\n\n// Here's the specialization for floats:\ntemplate <> class ArrayComparator\n{\npublic:\n bool compare(const float *reference, float *data, unsigned int len)\n {\n return compareData(reference, data, len, 0.15f, 0.15f);\n }\n};\n\n// Here's the generic wrapper for cutWriteFile*\ntemplate class ArrayFileWriter\n{\npublic:\n bool write(const char *filename, T *data, unsigned int len, float epsilon)\n {\n fprintf(stderr, \"Error: no file write function implemented for this type\\n\");\n return false;\n }\n};\n\n// Here's the specialization for ints:\ntemplate <> class ArrayFileWriter\n{\npublic:\n bool write(const char *filename, int *data, unsigned int len, float epsilon)\n {\n return sdkWriteFile(filename, data, len, epsilon, false);\n }\n};\n\n// Here's the specialization for floats:\ntemplate <> class ArrayFileWriter\n{\npublic:\n bool write(const char *filename, float *data, unsigned int len, float epsilon)\n {\n return sdkWriteFile(filename, data, len, epsilon, false);\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\ntemplate void runTest(int argc, char **argv, int len)\n{\n int devID;\n cudaDeviceProp deviceProps;\n\n devID = findCudaDevice(argc, (const char **)argv);\n\n // get number of SMs on this GPU\n checkCudaErrors(cudaGetDeviceProperties(&deviceProps, devID));\n printf(\"CUDA device [%s] has %d Multi-Processors\\n\", deviceProps.name, deviceProps.multiProcessorCount);\n\n // create and start timer\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n\n // start the timer\n sdkStartTimer(&timer);\n\n unsigned int num_threads = len;\n unsigned int mem_size = sizeof(float) * num_threads;\n\n // allocate host memory\n T *h_idata = (T *)malloc(mem_size);\n\n // initialize the memory\n for (unsigned int i = 0; i < num_threads; ++i) {\n h_idata[i] = (T)i;\n }\n\n // allocate device memory\n T *d_idata;\n checkCudaErrors(cudaMalloc((void **)&d_idata, mem_size));\n // copy host memory to device\n checkCudaErrors(cudaMemcpy(d_idata, h_idata, mem_size, cudaMemcpyHostToDevice));\n\n // allocate device memory for result\n T *d_odata;\n checkCudaErrors(cudaMalloc((void **)&d_odata, mem_size));\n\n // setup execution parameters\n dim3 grid(1, 1, 1);\n dim3 threads(num_threads, 1, 1);\n\n // execute the kernel\n testKernel<<>>(d_idata, d_odata);\n\n // check if kernel execution generated and error\n getLastCudaError(\"Kernel execution failed\");\n\n // allocate mem for the result on host side\n T *h_odata = (T *)malloc(mem_size);\n // copy result from device to host\n checkCudaErrors(cudaMemcpy(h_odata, d_odata, sizeof(T) * num_threads, cudaMemcpyDeviceToHost));\n\n sdkStopTimer(&timer);\n printf(\"Processing time: %f (ms)\\n\", sdkGetTimerValue(&timer));\n sdkDeleteTimer(&timer);\n\n // compute reference solution\n T *reference = (T *)malloc(mem_size);\n computeGold(reference, h_idata, num_threads);\n\n ArrayComparator comparator;\n ArrayFileWriter writer;\n\n // check result\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // write file for regression test\n writer.write(\"./data/regression.dat\", h_odata, num_threads, 0.0f);\n }\n else {\n // custom output handling when no regression test running\n // in this case check if the result is equivalent to the expected solution\n bool res = comparator.compare(reference, h_odata, num_threads);\n printf(\"Compare %s\\n\\n\", (1 == res) ? \"OK\" : \"MISMATCH\");\n g_TotalFailures += (1 != res);\n }\n\n // cleanup memory\n free(h_idata);\n free(h_odata);\n free(reference);\n checkCudaErrors(cudaFree(d_idata));\n checkCudaErrors(cudaFree(d_odata));\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* This sample is a templatized version of the template project.\n * It also shows how to correctly templatize dynamically allocated shared\n * memory arrays.\n * Host code.\n */\n\n// System includes\n#include \n#include \n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#ifndef MAX\n#define MAX(a, b) (a > b ? a : b)\n#endif\n\n// includes, kernels\n#include \"sharedmem.cuh\"\n\nint g_TotalFailures = 0;\n\n////////////////////////////////////////////////////////////////////////////////\n//! Simple test kernel for device functionality\n//! @param g_idata input data in global memory\n//! @param g_odata output data in global memory\n////////////////////////////////////////////////////////////////////////////////\ntemplate __global__ void testKernel(T *g_idata, T *g_odata)\n{\n // Shared mem size is determined by the host app at run time\n SharedMemory smem;\n T *sdata = smem.getPointer();\n\n // access thread id\n const unsigned int tid = threadIdx.x;\n // access number of threads in this block\n const unsigned int num_threads = blockDim.x;\n\n // read in input data from global memory\n sdata[tid] = g_idata[tid];\n __syncthreads();\n\n // perform some computations\n sdata[tid] = (T)num_threads * sdata[tid];\n __syncthreads();\n\n // write data to global memory\n g_odata[tid] = sdata[tid];\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\ntemplate void runTest(int argc, char **argv, int len);\n\ntemplate void computeGold(T *reference, T *idata, const unsigned int len)\n{\n const T T_len = static_cast(len);\n\n for (unsigned int i = 0; i < len; ++i) {\n reference[i] = idata[i] * T_len;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"> runTest\\n\");\n runTest(argc, argv, 32);\n printf(\"> runTest\\n\");\n runTest(argc, argv, 64);\n\n printf(\"\\n[simpleTemplates] -> Test Results: %d Failures\\n\", g_TotalFailures);\n\n exit(g_TotalFailures == 0 ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\n// To completely templatize runTest (below) with cutil, we need to use\n// template specialization to wrap up CUTIL's array comparison and file writing\n// functions for different types.\n\n// Here's the generic wrapper for cutCompare*\ntemplate class ArrayComparator\n{\npublic:\n bool compare(const T *reference, T *data, unsigned int len)\n {\n fprintf(stderr, \"Error: no comparison function implemented for this type\\n\");\n return false;\n }\n};\n\n// Here's the specialization for ints:\ntemplate <> class ArrayComparator\n{\npublic:\n bool compare(const int *reference, int *data, unsigned int len)\n {\n return compareData(reference, data, len, 0.15f, 0.0f);\n }\n};\n\n// Here's the specialization for floats:\ntemplate <> class ArrayComparator\n{\npublic:\n bool compare(const float *reference, float *data, unsigned int len)\n {\n return compareData(reference, data, len, 0.15f, 0.15f);\n }\n};\n\n// Here's the generic wrapper for cutWriteFile*\ntemplate class ArrayFileWriter\n{\npublic:\n bool write(const char *filename, T *data, unsigned int len, float epsilon)\n {\n fprintf(stderr, \"Error: no file write function implemented for this type\\n\");\n return false;\n }\n};\n\n// Here's the specialization for ints:\ntemplate <> class ArrayFileWriter\n{\npublic:\n bool write(const char *filename, int *data, unsigned int len, float epsilon)\n {\n return sdkWriteFile(filename, data, len, epsilon, false);\n }\n};\n\n// Here's the specialization for floats:\ntemplate <> class ArrayFileWriter\n{\npublic:\n bool write(const char *filename, float *data, unsigned int len, float epsilon)\n {\n return sdkWriteFile(filename, data, len, epsilon, false);\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\ntemplate void runTest(int argc, char **argv, int len)\n{\n int devID;\n hipDeviceProp_t deviceProps;\n\n devID = findCudaDevice(argc, (const char **)argv);\n\n // get number of SMs on this GPU\n checkCudaErrors(hipGetDeviceProperties(&deviceProps, devID));\n printf(\"CUDA device [%s] has %d Multi-Processors\\n\", deviceProps.name, deviceProps.multiProcessorCount);\n\n // create and start timer\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n\n // start the timer\n sdkStartTimer(&timer);\n\n unsigned int num_threads = len;\n unsigned int mem_size = sizeof(float) * num_threads;\n\n // allocate host memory\n T *h_idata = (T *)malloc(mem_size);\n\n // initialize the memory\n for (unsigned int i = 0; i < num_threads; ++i) {\n h_idata[i] = (T)i;\n }\n\n // allocate device memory\n T *d_idata;\n checkCudaErrors(hipMalloc((void **)&d_idata, mem_size));\n // copy host memory to device\n checkCudaErrors(hipMemcpy(d_idata, h_idata, mem_size, hipMemcpyHostToDevice));\n\n // allocate device memory for result\n T *d_odata;\n checkCudaErrors(hipMalloc((void **)&d_odata, mem_size));\n\n // setup execution parameters\n dim3 grid(1, 1, 1);\n dim3 threads(num_threads, 1, 1);\n\n // execute the kernel\n testKernel<<>>(d_idata, d_odata);\n\n // check if kernel execution generated and error\n getLastCudaError(\"Kernel execution failed\");\n\n // allocate mem for the result on host side\n T *h_odata = (T *)malloc(mem_size);\n // copy result from device to host\n checkCudaErrors(hipMemcpy(h_odata, d_odata, sizeof(T) * num_threads, hipMemcpyDeviceToHost));\n\n sdkStopTimer(&timer);\n printf(\"Processing time: %f (ms)\\n\", sdkGetTimerValue(&timer));\n sdkDeleteTimer(&timer);\n\n // compute reference solution\n T *reference = (T *)malloc(mem_size);\n computeGold(reference, h_idata, num_threads);\n\n ArrayComparator comparator;\n ArrayFileWriter writer;\n\n // check result\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // write file for regression test\n writer.write(\"./data/regression.dat\", h_odata, num_threads, 0.0f);\n }\n else {\n // custom output handling when no regression test running\n // in this case check if the result is equivalent to the expected solution\n bool res = comparator.compare(reference, h_odata, num_threads);\n printf(\"Compare %s\\n\\n\", (1 == res) ? \"OK\" : \"MISMATCH\");\n g_TotalFailures += (1 != res);\n }\n\n // cleanup memory\n free(h_idata);\n free(h_odata);\n free(reference);\n checkCudaErrors(hipFree(d_idata));\n checkCudaErrors(hipFree(d_odata));\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/simpleTemplates/simpleTemplates.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/**\n * Vector addition: C = A + B.\n *\n * This sample is a very basic sample that implements element by element\n * vector addition. It is the same as the sample illustrating Chapter 2\n * of the programming guide with some additions like error checking.\n */\n\n#include \n\n// For the CUDA runtime routines (prefixed with \"cuda_\")\n#include \n#include \n/**\n * CUDA Kernel Device code\n *\n * Computes the vector addition of A and B into C. The 3 vectors have the same\n * number of elements numElements.\n */\n__global__ void vectorAdd(const float *A, const float *B, float *C, int numElements)\n{\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (i < numElements) {\n C[i] = A[i] + B[i] + 0.0f;\n }\n}\n\n/**\n * Host main routine\n */\nint main(void)\n{\n // Error code to check return values for CUDA calls\n cudaError_t err = cudaSuccess;\n\n // Print the vector length to be used, and compute its size\n int numElements = 50000;\n size_t size = numElements * sizeof(float);\n printf(\"[Vector addition of %d elements]\\n\", numElements);\n\n // Allocate the host input vector A\n float *h_A = (float *)malloc(size);\n\n // Allocate the host input vector B\n float *h_B = (float *)malloc(size);\n\n // Allocate the host output vector C\n float *h_C = (float *)malloc(size);\n\n // Verify that allocations succeeded\n if (h_A == NULL || h_B == NULL || h_C == NULL) {\n fprintf(stderr, \"Failed to allocate host vectors!\\n\");\n exit(EXIT_FAILURE);\n }\n\n // Initialize the host input vectors\n for (int i = 0; i < numElements; ++i) {\n h_A[i] = rand() / (float)RAND_MAX;\n h_B[i] = rand() / (float)RAND_MAX;\n }\n\n // Allocate the device input vector A\n float *d_A = NULL;\n err = cudaMalloc((void **)&d_A, size);\n\n if (err != cudaSuccess) {\n fprintf(stderr, \"Failed to allocate device vector A (error code %s)!\\n\", cudaGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n // Allocate the device input vector B\n float *d_B = NULL;\n err = cudaMalloc((void **)&d_B, size);\n\n if (err != cudaSuccess) {\n fprintf(stderr, \"Failed to allocate device vector B (error code %s)!\\n\", cudaGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n // Allocate the device output vector C\n float *d_C = NULL;\n err = cudaMalloc((void **)&d_C, size);\n\n if (err != cudaSuccess) {\n fprintf(stderr, \"Failed to allocate device vector C (error code %s)!\\n\", cudaGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n // Copy the host input vectors A and B in host memory to the device input\n // vectors in\n // device memory\n printf(\"Copy input data from the host memory to the CUDA device\\n\");\n err = cudaMemcpy(d_A, h_A, size, cudaMemcpyHostToDevice);\n\n if (err != cudaSuccess) {\n fprintf(stderr, \"Failed to copy vector A from host to device (error code %s)!\\n\", cudaGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n err = cudaMemcpy(d_B, h_B, size, cudaMemcpyHostToDevice);\n\n if (err != cudaSuccess) {\n fprintf(stderr, \"Failed to copy vector B from host to device (error code %s)!\\n\", cudaGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n // Launch the Vector Add CUDA Kernel\n int threadsPerBlock = 256;\n int blocksPerGrid = (numElements + threadsPerBlock - 1) / threadsPerBlock;\n printf(\"CUDA kernel launch with %d blocks of %d threads\\n\", blocksPerGrid, threadsPerBlock);\n vectorAdd<<>>(d_A, d_B, d_C, numElements);\n err = cudaGetLastError();\n\n if (err != cudaSuccess) {\n fprintf(stderr, \"Failed to launch vectorAdd kernel (error code %s)!\\n\", cudaGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n // Copy the device result vector in device memory to the host result vector\n // in host memory.\n printf(\"Copy output data from the CUDA device to the host memory\\n\");\n err = cudaMemcpy(h_C, d_C, size, cudaMemcpyDeviceToHost);\n\n if (err != cudaSuccess) {\n fprintf(stderr, \"Failed to copy vector C from device to host (error code %s)!\\n\", cudaGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n // Verify that the result vector is correct\n for (int i = 0; i < numElements; ++i) {\n if (fabs(h_A[i] + h_B[i] - h_C[i]) > 1e-5) {\n fprintf(stderr, \"Result verification failed at element %d!\\n\", i);\n exit(EXIT_FAILURE);\n }\n }\n\n printf(\"Test PASSED\\n\");\n\n // Free device global memory\n err = cudaFree(d_A);\n\n if (err != cudaSuccess) {\n fprintf(stderr, \"Failed to free device vector A (error code %s)!\\n\", cudaGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n err = cudaFree(d_B);\n\n if (err != cudaSuccess) {\n fprintf(stderr, \"Failed to free device vector B (error code %s)!\\n\", cudaGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n err = cudaFree(d_C);\n\n if (err != cudaSuccess) {\n fprintf(stderr, \"Failed to free device vector C (error code %s)!\\n\", cudaGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n // Free host memory\n free(h_A);\n free(h_B);\n free(h_C);\n\n printf(\"Done\\n\");\n return 0;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/**\n * Vector addition: C = A + B.\n *\n * This sample is a very basic sample that implements element by element\n * vector addition. It is the same as the sample illustrating Chapter 2\n * of the programming guide with some additions like error checking.\n */\n\n#include \n\n// For the CUDA runtime routines (prefixed with \"cuda_\")\n#include \n#include \n/**\n * CUDA Kernel Device code\n *\n * Computes the vector addition of A and B into C. The 3 vectors have the same\n * number of elements numElements.\n */\n__global__ void vectorAdd(const float *A, const float *B, float *C, int numElements)\n{\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (i < numElements) {\n C[i] = A[i] + B[i] + 0.0f;\n }\n}\n\n/**\n * Host main routine\n */\nint main(void)\n{\n // Error code to check return values for CUDA calls\n hipError_t err = hipSuccess;\n\n // Print the vector length to be used, and compute its size\n int numElements = 50000;\n size_t size = numElements * sizeof(float);\n printf(\"[Vector addition of %d elements]\\n\", numElements);\n\n // Allocate the host input vector A\n float *h_A = (float *)malloc(size);\n\n // Allocate the host input vector B\n float *h_B = (float *)malloc(size);\n\n // Allocate the host output vector C\n float *h_C = (float *)malloc(size);\n\n // Verify that allocations succeeded\n if (h_A == NULL || h_B == NULL || h_C == NULL) {\n fprintf(stderr, \"Failed to allocate host vectors!\\n\");\n exit(EXIT_FAILURE);\n }\n\n // Initialize the host input vectors\n for (int i = 0; i < numElements; ++i) {\n h_A[i] = rand() / (float)RAND_MAX;\n h_B[i] = rand() / (float)RAND_MAX;\n }\n\n // Allocate the device input vector A\n float *d_A = NULL;\n err = hipMalloc((void **)&d_A, size);\n\n if (err != hipSuccess) {\n fprintf(stderr, \"Failed to allocate device vector A (error code %s)!\\n\", hipGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n // Allocate the device input vector B\n float *d_B = NULL;\n err = hipMalloc((void **)&d_B, size);\n\n if (err != hipSuccess) {\n fprintf(stderr, \"Failed to allocate device vector B (error code %s)!\\n\", hipGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n // Allocate the device output vector C\n float *d_C = NULL;\n err = hipMalloc((void **)&d_C, size);\n\n if (err != hipSuccess) {\n fprintf(stderr, \"Failed to allocate device vector C (error code %s)!\\n\", hipGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n // Copy the host input vectors A and B in host memory to the device input\n // vectors in\n // device memory\n printf(\"Copy input data from the host memory to the CUDA device\\n\");\n err = hipMemcpy(d_A, h_A, size, hipMemcpyHostToDevice);\n\n if (err != hipSuccess) {\n fprintf(stderr, \"Failed to copy vector A from host to device (error code %s)!\\n\", hipGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n err = hipMemcpy(d_B, h_B, size, hipMemcpyHostToDevice);\n\n if (err != hipSuccess) {\n fprintf(stderr, \"Failed to copy vector B from host to device (error code %s)!\\n\", hipGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n // Launch the Vector Add CUDA Kernel\n int threadsPerBlock = 256;\n int blocksPerGrid = (numElements + threadsPerBlock - 1) / threadsPerBlock;\n printf(\"CUDA kernel launch with %d blocks of %d threads\\n\", blocksPerGrid, threadsPerBlock);\n vectorAdd<<>>(d_A, d_B, d_C, numElements);\n err = hipGetLastError();\n\n if (err != hipSuccess) {\n fprintf(stderr, \"Failed to launch vectorAdd kernel (error code %s)!\\n\", hipGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n // Copy the device result vector in device memory to the host result vector\n // in host memory.\n printf(\"Copy output data from the CUDA device to the host memory\\n\");\n err = hipMemcpy(h_C, d_C, size, hipMemcpyDeviceToHost);\n\n if (err != hipSuccess) {\n fprintf(stderr, \"Failed to copy vector C from device to host (error code %s)!\\n\", hipGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n // Verify that the result vector is correct\n for (int i = 0; i < numElements; ++i) {\n if (fabs(h_A[i] + h_B[i] - h_C[i]) > 1e-5) {\n fprintf(stderr, \"Result verification failed at element %d!\\n\", i);\n exit(EXIT_FAILURE);\n }\n }\n\n printf(\"Test PASSED\\n\");\n\n // Free device global memory\n err = hipFree(d_A);\n\n if (err != hipSuccess) {\n fprintf(stderr, \"Failed to free device vector A (error code %s)!\\n\", hipGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n err = hipFree(d_B);\n\n if (err != hipSuccess) {\n fprintf(stderr, \"Failed to free device vector B (error code %s)!\\n\", hipGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n err = hipFree(d_C);\n\n if (err != hipSuccess) {\n fprintf(stderr, \"Failed to free device vector C (error code %s)!\\n\", hipGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n\n // Free host memory\n free(h_A);\n free(h_B);\n free(h_C);\n\n printf(\"Done\\n\");\n return 0;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/vectorAdd/vectorAdd.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample illustrates the usage of CUDA events for both GPU timing and\n * overlapping CPU and GPU execution. Events are inserted into a stream\n * of CUDA calls. Since CUDA stream calls are asynchronous, the CPU can\n * perform computations while GPU is executing (including DMA memcopies\n * between the host and device). CPU can query CUDA events to determine\n * whether GPU has completed tasks.\n */\n\n// includes, system\n#include \n\n// includes CUDA Runtime\n#include \n#include \n\n// includes, project\n#include \n#include // helper utility functions\n\n__global__ void increment_kernel(int *g_data, int inc_value)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n g_data[idx] = g_data[idx] + inc_value;\n}\n\nbool correct_output(int *data, const int n, const int x)\n{\n for (int i = 0; i < n; i++)\n if (data[i] != x) {\n printf(\"Error! data[%d] = %d, ref = %d\\n\", i, data[i], x);\n return false;\n }\n\n return true;\n}\n\nint main(int argc, char *argv[])\n{\n int devID;\n cudaDeviceProp deviceProps;\n\n printf(\"[%s] - Starting...\\n\", argv[0]);\n\n // This will pick the best possible CUDA capable device\n devID = findCudaDevice(argc, (const char **)argv);\n\n // get device name\n checkCudaErrors(cudaGetDeviceProperties(&deviceProps, devID));\n printf(\"CUDA device [%s]\\n\", deviceProps.name);\n\n int n = 16 * 1024 * 1024;\n int nbytes = n * sizeof(int);\n int value = 26;\n\n // allocate host memory\n int *a = 0;\n checkCudaErrors(cudaMallocHost((void **)&a, nbytes));\n memset(a, 0, nbytes);\n\n // allocate device memory\n int *d_a = 0;\n checkCudaErrors(cudaMalloc((void **)&d_a, nbytes));\n checkCudaErrors(cudaMemset(d_a, 255, nbytes));\n\n // set kernel launch configuration\n dim3 threads = dim3(512, 1);\n dim3 blocks = dim3(n / threads.x, 1);\n\n // create cuda event handles\n cudaEvent_t start, stop;\n checkCudaErrors(cudaEventCreate(&start));\n checkCudaErrors(cudaEventCreate(&stop));\n\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n sdkResetTimer(&timer);\n\n checkCudaErrors(cudaDeviceSynchronize());\n float gpu_time = 0.0f;\n\n // asynchronously issue work to the GPU (all to stream 0)\n checkCudaErrors(cudaProfilerStart());\n sdkStartTimer(&timer);\n cudaEventRecord(start, 0);\n cudaMemcpyAsync(d_a, a, nbytes, cudaMemcpyHostToDevice, 0);\n increment_kernel<<>>(d_a, value);\n cudaMemcpyAsync(a, d_a, nbytes, cudaMemcpyDeviceToHost, 0);\n cudaEventRecord(stop, 0);\n sdkStopTimer(&timer);\n checkCudaErrors(cudaProfilerStop());\n\n // have CPU do some work while waiting for stage 1 to finish\n unsigned long int counter = 0;\n\n while (cudaEventQuery(stop) == cudaErrorNotReady) {\n counter++;\n }\n\n checkCudaErrors(cudaEventElapsedTime(&gpu_time, start, stop));\n\n // print the cpu and gpu times\n printf(\"time spent executing by the GPU: %.2f\\n\", gpu_time);\n printf(\"time spent by CPU in CUDA calls: %.2f\\n\", sdkGetTimerValue(&timer));\n printf(\"CPU executed %lu iterations while waiting for GPU to finish\\n\", counter);\n\n // check the output for correctness\n bool bFinalResults = correct_output(a, n, value);\n\n // release resources\n checkCudaErrors(cudaEventDestroy(start));\n checkCudaErrors(cudaEventDestroy(stop));\n checkCudaErrors(cudaFreeHost(a));\n checkCudaErrors(cudaFree(d_a));\n\n exit(bFinalResults ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample illustrates the usage of CUDA events for both GPU timing and\n * overlapping CPU and GPU execution. Events are inserted into a stream\n * of CUDA calls. Since CUDA stream calls are asynchronous, the CPU can\n * perform computations while GPU is executing (including DMA memcopies\n * between the host and device). CPU can query CUDA events to determine\n * whether GPU has completed tasks.\n */\n\n// includes, system\n#include \n\n// includes CUDA Runtime\n#include \n#include \n\n// includes, project\n#include // helper utility functions\n\n__global__ void increment_kernel(int *g_data, int inc_value)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n g_data[idx] = g_data[idx] + inc_value;\n}\n\nbool correct_output(int *data, const int n, const int x)\n{\n for (int i = 0; i < n; i++)\n if (data[i] != x) {\n printf(\"Error! data[%d] = %d, ref = %d\\n\", i, data[i], x);\n return false;\n }\n\n return true;\n}\n\nint main(int argc, char *argv[])\n{\n int devID;\n hipDeviceProp_t deviceProps;\n\n printf(\"[%s] - Starting...\\n\", argv[0]);\n\n // This will pick the best possible CUDA capable device\n devID = findCudaDevice(argc, (const char **)argv);\n\n // get device name\n checkCudaErrors(hipGetDeviceProperties(&deviceProps, devID));\n printf(\"CUDA device [%s]\\n\", deviceProps.name);\n\n int n = 16 * 1024 * 1024;\n int nbytes = n * sizeof(int);\n int value = 26;\n\n // allocate host memory\n int *a = 0;\n checkCudaErrors(hipHostMalloc((void **)&a, nbytes));\n memset(a, 0, nbytes);\n\n // allocate device memory\n int *d_a = 0;\n checkCudaErrors(hipMalloc((void **)&d_a, nbytes));\n checkCudaErrors(hipMemset(d_a, 255, nbytes));\n\n // set kernel launch configuration\n dim3 threads = dim3(512, 1);\n dim3 blocks = dim3(n / threads.x, 1);\n\n // create cuda event handles\n hipEvent_t start, stop;\n checkCudaErrors(hipEventCreate(&start));\n checkCudaErrors(hipEventCreate(&stop));\n\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n sdkResetTimer(&timer);\n\n checkCudaErrors(hipDeviceSynchronize());\n float gpu_time = 0.0f;\n\n // asynchronously issue work to the GPU (all to stream 0)\n checkCudaErrors(hipProfilerStart());\n sdkStartTimer(&timer);\n hipEventRecord(start, 0);\n hipMemcpyAsync(d_a, a, nbytes, hipMemcpyHostToDevice, 0);\n increment_kernel<<>>(d_a, value);\n hipMemcpyAsync(a, d_a, nbytes, hipMemcpyDeviceToHost, 0);\n hipEventRecord(stop, 0);\n sdkStopTimer(&timer);\n checkCudaErrors(hipProfilerStop());\n\n // have CPU do some work while waiting for stage 1 to finish\n unsigned long int counter = 0;\n\n while (hipEventQuery(stop) == hipErrorNotReady) {\n counter++;\n }\n\n checkCudaErrors(hipEventElapsedTime(&gpu_time, start, stop));\n\n // print the cpu and gpu times\n printf(\"time spent executing by the GPU: %.2f\\n\", gpu_time);\n printf(\"time spent by CPU in CUDA calls: %.2f\\n\", sdkGetTimerValue(&timer));\n printf(\"CPU executed %lu iterations while waiting for GPU to finish\\n\", counter);\n\n // check the output for correctness\n bool bFinalResults = correct_output(a, n, value);\n\n // release resources\n checkCudaErrors(hipEventDestroy(start));\n checkCudaErrors(hipEventDestroy(stop));\n checkCudaErrors(hipHostFree(a));\n checkCudaErrors(hipFree(d_a));\n\n exit(bFinalResults ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/0_Introduction/asyncAPI/asyncAPI.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#ifndef QUASIRANDOMGENERATOR_KERNEL_CUH\n#define QUASIRANDOMGENERATOR_KERNEL_CUH\n\n#include \n#include \n#include \n\n#include \"quasirandomGenerator_common.h\"\n\n// Fast integer multiplication\n#define MUL(a, b) __umul24(a, b)\n\n////////////////////////////////////////////////////////////////////////////////\n// Niederreiter quasirandom number generation kernel\n////////////////////////////////////////////////////////////////////////////////\nstatic __constant__ unsigned int c_Table[QRNG_DIMENSIONS][QRNG_RESOLUTION];\n\nstatic __global__ void quasirandomGeneratorKernel(float *d_Output, unsigned int seed, unsigned int N)\n{\n unsigned int *dimBase = &c_Table[threadIdx.y][0];\n unsigned int tid = MUL(blockDim.x, blockIdx.x) + threadIdx.x;\n unsigned int threadN = MUL(blockDim.x, gridDim.x);\n\n for (unsigned int pos = tid; pos < N; pos += threadN) {\n unsigned int result = 0;\n unsigned int data = seed + pos;\n\n for (int bit = 0; bit < QRNG_RESOLUTION; bit++, data >>= 1)\n if (data & 1) {\n result ^= dimBase[bit];\n }\n\n d_Output[MUL(threadIdx.y, N) + pos] = (float)(result + 1) * INT_SCALE;\n }\n}\n\n// Table initialization routine\nextern \"C\" void initTableGPU(unsigned int tableCPU[QRNG_DIMENSIONS][QRNG_RESOLUTION])\n{\n checkCudaErrors(cudaMemcpyToSymbol(c_Table, tableCPU, QRNG_DIMENSIONS * QRNG_RESOLUTION * sizeof(unsigned int)));\n}\n\n// Host-side interface\nextern \"C\" void quasirandomGeneratorGPU(float *d_Output, unsigned int seed, unsigned int N)\n{\n dim3 threads(128, QRNG_DIMENSIONS);\n quasirandomGeneratorKernel<<<128, threads>>>(d_Output, seed, N);\n getLastCudaError(\"quasirandomGeneratorKernel() execution failed.\\n\");\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Moro's Inverse Cumulative Normal Distribution function approximation\n////////////////////////////////////////////////////////////////////////////////\n__device__ inline float MoroInvCNDgpu(unsigned int x)\n{\n const float a1 = 2.50662823884f;\n const float a2 = -18.61500062529f;\n const float a3 = 41.39119773534f;\n const float a4 = -25.44106049637f;\n const float b1 = -8.4735109309f;\n const float b2 = 23.08336743743f;\n const float b3 = -21.06224101826f;\n const float b4 = 3.13082909833f;\n const float c1 = 0.337475482272615f;\n const float c2 = 0.976169019091719f;\n const float c3 = 0.160797971491821f;\n const float c4 = 2.76438810333863E-02f;\n const float c5 = 3.8405729373609E-03f;\n const float c6 = 3.951896511919E-04f;\n const float c7 = 3.21767881768E-05f;\n const float c8 = 2.888167364E-07f;\n const float c9 = 3.960315187E-07f;\n\n float z;\n\n bool negate = false;\n\n // Ensure the conversion to floating point will give a value in the\n // range (0,0.5] by restricting the input to the bottom half of the\n // input domain. We will later reflect the result if the input was\n // originally in the top half of the input domain\n if (x >= 0x80000000UL) {\n x = 0xffffffffUL - x;\n negate = true;\n }\n\n // x is now in the range [0,0x80000000) (i.e. [0,0x7fffffff])\n // Convert to floating point in (0,0.5]\n const float x1 = 1.0f / static_cast(0xffffffffUL);\n const float x2 = x1 / 2.0f;\n float p1 = x * x1 + x2;\n // Convert to floating point in (-0.5,0]\n float p2 = p1 - 0.5f;\n\n // The input to the Moro inversion is p2 which is in the range\n // (-0.5,0]. This means that our output will be the negative side\n // of the bell curve (which we will reflect if \"negate\" is true).\n\n // Main body of the bell curve for |p| < 0.42\n if (p2 > -0.42f) {\n z = p2 * p2;\n z = p2 * (((a4 * z + a3) * z + a2) * z + a1) / ((((b4 * z + b3) * z + b2) * z + b1) * z + 1.0f);\n }\n // Special case (Chebychev) for tail\n else {\n z = __logf(-__logf(p1));\n z = -(c1 + z * (c2 + z * (c3 + z * (c4 + z * (c5 + z * (c6 + z * (c7 + z * (c8 + z * c9))))))));\n }\n\n // If the original input (x) was in the top half of the range, reflect\n // to get the positive side of the bell curve\n return negate ? -z : z;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Main kernel. Choose between transforming\n// input sequence and uniform ascending (0, 1) sequence\n////////////////////////////////////////////////////////////////////////////////\nstatic __global__ void inverseCNDKernel(float *d_Output, unsigned int *d_Input, unsigned int pathN)\n{\n unsigned int distance = ((unsigned int)-1) / (pathN + 1);\n unsigned int tid = MUL(blockDim.x, blockIdx.x) + threadIdx.x;\n unsigned int threadN = MUL(blockDim.x, gridDim.x);\n\n // Transform input number sequence if it's supplied\n if (d_Input) {\n for (unsigned int pos = tid; pos < pathN; pos += threadN) {\n unsigned int d = d_Input[pos];\n d_Output[pos] = (float)MoroInvCNDgpu(d);\n }\n }\n // Else generate input uniformly placed samples on the fly\n // and write to destination\n else {\n for (unsigned int pos = tid; pos < pathN; pos += threadN) {\n unsigned int d = (pos + 1) * distance;\n d_Output[pos] = (float)MoroInvCNDgpu(d);\n }\n }\n}\n\nextern \"C\" void inverseCNDgpu(float *d_Output, unsigned int *d_Input, unsigned int N)\n{\n inverseCNDKernel<<<128, 128>>>(d_Output, d_Input, N);\n getLastCudaError(\"inverseCNDKernel() execution failed.\\n\");\n}\n\n#endif\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#ifndef QUASIRANDOMGENERATOR_KERNEL_CUH\n#define QUASIRANDOMGENERATOR_KERNEL_CUH\n\n#include \n#include \n#include \n\n#include \"quasirandomGenerator_common.h\"\n\n// Fast integer multiplication\n#define MUL(a, b) __umul24(a, b)\n\n////////////////////////////////////////////////////////////////////////////////\n// Niederreiter quasirandom number generation kernel\n////////////////////////////////////////////////////////////////////////////////\nstatic __constant__ unsigned int c_Table[QRNG_DIMENSIONS][QRNG_RESOLUTION];\n\nstatic __global__ void quasirandomGeneratorKernel(float *d_Output, unsigned int seed, unsigned int N)\n{\n unsigned int *dimBase = &c_Table[threadIdx.y][0];\n unsigned int tid = MUL(blockDim.x, blockIdx.x) + threadIdx.x;\n unsigned int threadN = MUL(blockDim.x, gridDim.x);\n\n for (unsigned int pos = tid; pos < N; pos += threadN) {\n unsigned int result = 0;\n unsigned int data = seed + pos;\n\n for (int bit = 0; bit < QRNG_RESOLUTION; bit++, data >>= 1)\n if (data & 1) {\n result ^= dimBase[bit];\n }\n\n d_Output[MUL(threadIdx.y, N) + pos] = (float)(result + 1) * INT_SCALE;\n }\n}\n\n// Table initialization routine\nextern \"C\" void initTableGPU(unsigned int tableCPU[QRNG_DIMENSIONS][QRNG_RESOLUTION])\n{\n checkCudaErrors(hipMemcpyToSymbol(HIP_SYMBOL(c_Table), tableCPU, QRNG_DIMENSIONS * QRNG_RESOLUTION * sizeof(unsigned int)));\n}\n\n// Host-side interface\nextern \"C\" void quasirandomGeneratorGPU(float *d_Output, unsigned int seed, unsigned int N)\n{\n dim3 threads(128, QRNG_DIMENSIONS);\n quasirandomGeneratorKernel<<<128, threads>>>(d_Output, seed, N);\n getLastCudaError(\"quasirandomGeneratorKernel() execution failed.\\n\");\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Moro's Inverse Cumulative Normal Distribution function approximation\n////////////////////////////////////////////////////////////////////////////////\n__device__ inline float MoroInvCNDgpu(unsigned int x)\n{\n const float a1 = 2.50662823884f;\n const float a2 = -18.61500062529f;\n const float a3 = 41.39119773534f;\n const float a4 = -25.44106049637f;\n const float b1 = -8.4735109309f;\n const float b2 = 23.08336743743f;\n const float b3 = -21.06224101826f;\n const float b4 = 3.13082909833f;\n const float c1 = 0.337475482272615f;\n const float c2 = 0.976169019091719f;\n const float c3 = 0.160797971491821f;\n const float c4 = 2.76438810333863E-02f;\n const float c5 = 3.8405729373609E-03f;\n const float c6 = 3.951896511919E-04f;\n const float c7 = 3.21767881768E-05f;\n const float c8 = 2.888167364E-07f;\n const float c9 = 3.960315187E-07f;\n\n float z;\n\n bool negate = false;\n\n // Ensure the conversion to floating point will give a value in the\n // range (0,0.5] by restricting the input to the bottom half of the\n // input domain. We will later reflect the result if the input was\n // originally in the top half of the input domain\n if (x >= 0x80000000UL) {\n x = 0xffffffffUL - x;\n negate = true;\n }\n\n // x is now in the range [0,0x80000000) (i.e. [0,0x7fffffff])\n // Convert to floating point in (0,0.5]\n const float x1 = 1.0f / static_cast(0xffffffffUL);\n const float x2 = x1 / 2.0f;\n float p1 = x * x1 + x2;\n // Convert to floating point in (-0.5,0]\n float p2 = p1 - 0.5f;\n\n // The input to the Moro inversion is p2 which is in the range\n // (-0.5,0]. This means that our output will be the negative side\n // of the bell curve (which we will reflect if \"negate\" is true).\n\n // Main body of the bell curve for |p| < 0.42\n if (p2 > -0.42f) {\n z = p2 * p2;\n z = p2 * (((a4 * z + a3) * z + a2) * z + a1) / ((((b4 * z + b3) * z + b2) * z + b1) * z + 1.0f);\n }\n // Special case (Chebychev) for tail\n else {\n z = __logf(-__logf(p1));\n z = -(c1 + z * (c2 + z * (c3 + z * (c4 + z * (c5 + z * (c6 + z * (c7 + z * (c8 + z * c9))))))));\n }\n\n // If the original input (x) was in the top half of the range, reflect\n // to get the positive side of the bell curve\n return negate ? -z : z;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Main kernel. Choose between transforming\n// input sequence and uniform ascending (0, 1) sequence\n////////////////////////////////////////////////////////////////////////////////\nstatic __global__ void inverseCNDKernel(float *d_Output, unsigned int *d_Input, unsigned int pathN)\n{\n unsigned int distance = ((unsigned int)-1) / (pathN + 1);\n unsigned int tid = MUL(blockDim.x, blockIdx.x) + threadIdx.x;\n unsigned int threadN = MUL(blockDim.x, gridDim.x);\n\n // Transform input number sequence if it's supplied\n if (d_Input) {\n for (unsigned int pos = tid; pos < pathN; pos += threadN) {\n unsigned int d = d_Input[pos];\n d_Output[pos] = (float)MoroInvCNDgpu(d);\n }\n }\n // Else generate input uniformly placed samples on the fly\n // and write to destination\n else {\n for (unsigned int pos = tid; pos < pathN; pos += threadN) {\n unsigned int d = (pos + 1) * distance;\n d_Output[pos] = (float)MoroInvCNDgpu(d);\n }\n }\n}\n\nextern \"C\" void inverseCNDgpu(float *d_Output, unsigned int *d_Input, unsigned int N)\n{\n inverseCNDKernel<<<128, 128>>>(d_Output, d_Input, N);\n getLastCudaError(\"inverseCNDKernel() execution failed.\\n\");\n}\n\n#endif\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/quasirandomGenerator/quasirandomGenerator_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n////////////////////////////////////////////////////////////////////////////////\n// Global types and parameters\n////////////////////////////////////////////////////////////////////////////////\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \n\n#include \"binomialOptions_common.h\"\n#include \"realtype.h\"\n\n// Preprocessed input option data\ntypedef struct\n{\n real S;\n real X;\n real vDt;\n real puByDf;\n real pdByDf;\n} __TOptionData;\nstatic __constant__ __TOptionData d_OptionData[MAX_OPTIONS];\nstatic __device__ real d_CallValue[MAX_OPTIONS];\n\n////////////////////////////////////////////////////////////////////////////////\n// Overloaded shortcut functions for different precision modes\n////////////////////////////////////////////////////////////////////////////////\n#ifndef DOUBLE_PRECISION\n__device__ inline float expiryCallValue(float S, float X, float vDt, int i)\n{\n float d = S * __expf(vDt * (2.0f * i - NUM_STEPS)) - X;\n return (d > 0.0F) ? d : 0.0F;\n}\n#else\n__device__ inline double expiryCallValue(double S, double X, double vDt, int i)\n{\n double d = S * exp(vDt * (2.0 * i - NUM_STEPS)) - X;\n return (d > 0.0) ? d : 0.0;\n}\n#endif\n\n////////////////////////////////////////////////////////////////////////////////\n// GPU kernel\n////////////////////////////////////////////////////////////////////////////////\n#define THREADBLOCK_SIZE 128\n#define ELEMS_PER_THREAD (NUM_STEPS / THREADBLOCK_SIZE)\n#if NUM_STEPS % THREADBLOCK_SIZE\n#error Bad constants\n#endif\n\n__global__ void binomialOptionsKernel()\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ real call_exchange[THREADBLOCK_SIZE + 1];\n\n const int tid = threadIdx.x;\n const real S = d_OptionData[blockIdx.x].S;\n const real X = d_OptionData[blockIdx.x].X;\n const real vDt = d_OptionData[blockIdx.x].vDt;\n const real puByDf = d_OptionData[blockIdx.x].puByDf;\n const real pdByDf = d_OptionData[blockIdx.x].pdByDf;\n\n real call[ELEMS_PER_THREAD + 1];\n#pragma unroll\n for (int i = 0; i < ELEMS_PER_THREAD; ++i)\n call[i] = expiryCallValue(S, X, vDt, tid * ELEMS_PER_THREAD + i);\n\n if (tid == 0)\n call_exchange[THREADBLOCK_SIZE] = expiryCallValue(S, X, vDt, NUM_STEPS);\n\n int final_it = max(0, tid * ELEMS_PER_THREAD - 1);\n\n#pragma unroll 16\n for (int i = NUM_STEPS; i > 0; --i) {\n call_exchange[tid] = call[0];\n cg::sync(cta);\n call[ELEMS_PER_THREAD] = call_exchange[tid + 1];\n cg::sync(cta);\n\n if (i > final_it) {\n#pragma unroll\n for (int j = 0; j < ELEMS_PER_THREAD; ++j)\n call[j] = puByDf * call[j + 1] + pdByDf * call[j];\n }\n }\n\n if (tid == 0) {\n d_CallValue[blockIdx.x] = call[0];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Host-side interface to GPU binomialOptions\n////////////////////////////////////////////////////////////////////////////////\nextern \"C\" void binomialOptionsGPU(real *callValue, TOptionData *optionData, int optN)\n{\n __TOptionData h_OptionData[MAX_OPTIONS];\n\n for (int i = 0; i < optN; i++) {\n const real T = optionData[i].T;\n const real R = optionData[i].R;\n const real V = optionData[i].V;\n\n const real dt = T / (real)NUM_STEPS;\n const real vDt = V * sqrt(dt);\n const real rDt = R * dt;\n // Per-step interest and discount factors\n const real If = exp(rDt);\n const real Df = exp(-rDt);\n // Values and pseudoprobabilities of upward and downward moves\n const real u = exp(vDt);\n const real d = exp(-vDt);\n const real pu = (If - d) / (u - d);\n const real pd = (real)1.0 - pu;\n const real puByDf = pu * Df;\n const real pdByDf = pd * Df;\n\n h_OptionData[i].S = (real)optionData[i].S;\n h_OptionData[i].X = (real)optionData[i].X;\n h_OptionData[i].vDt = (real)vDt;\n h_OptionData[i].puByDf = (real)puByDf;\n h_OptionData[i].pdByDf = (real)pdByDf;\n }\n\n checkCudaErrors(cudaMemcpyToSymbol(d_OptionData, h_OptionData, optN * sizeof(__TOptionData)));\n binomialOptionsKernel<<>>();\n getLastCudaError(\"binomialOptionsKernel() execution failed.\\n\");\n checkCudaErrors(cudaMemcpyFromSymbol(callValue, d_CallValue, optN * sizeof(real)));\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n////////////////////////////////////////////////////////////////////////////////\n// Global types and parameters\n////////////////////////////////////////////////////////////////////////////////\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \n\n#include \"binomialOptions_common.h\"\n#include \"realtype.h\"\n\n// Preprocessed input option data\ntypedef struct\n{\n real S;\n real X;\n real vDt;\n real puByDf;\n real pdByDf;\n} __TOptionData;\nstatic __constant__ __TOptionData d_OptionData[MAX_OPTIONS];\nstatic __device__ real d_CallValue[MAX_OPTIONS];\n\n////////////////////////////////////////////////////////////////////////////////\n// Overloaded shortcut functions for different precision modes\n////////////////////////////////////////////////////////////////////////////////\n#ifndef DOUBLE_PRECISION\n__device__ inline float expiryCallValue(float S, float X, float vDt, int i)\n{\n float d = S * __expf(vDt * (2.0f * i - NUM_STEPS)) - X;\n return (d > 0.0F) ? d : 0.0F;\n}\n#else\n__device__ inline double expiryCallValue(double S, double X, double vDt, int i)\n{\n double d = S * exp(vDt * (2.0 * i - NUM_STEPS)) - X;\n return (d > 0.0) ? d : 0.0;\n}\n#endif\n\n////////////////////////////////////////////////////////////////////////////////\n// GPU kernel\n////////////////////////////////////////////////////////////////////////////////\n#define THREADBLOCK_SIZE 128\n#define ELEMS_PER_THREAD (NUM_STEPS / THREADBLOCK_SIZE)\n#if NUM_STEPS % THREADBLOCK_SIZE\n#error Bad constants\n#endif\n\n__global__ void binomialOptionsKernel()\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ real call_exchange[THREADBLOCK_SIZE + 1];\n\n const int tid = threadIdx.x;\n const real S = d_OptionData[blockIdx.x].S;\n const real X = d_OptionData[blockIdx.x].X;\n const real vDt = d_OptionData[blockIdx.x].vDt;\n const real puByDf = d_OptionData[blockIdx.x].puByDf;\n const real pdByDf = d_OptionData[blockIdx.x].pdByDf;\n\n real call[ELEMS_PER_THREAD + 1];\n#pragma unroll\n for (int i = 0; i < ELEMS_PER_THREAD; ++i)\n call[i] = expiryCallValue(S, X, vDt, tid * ELEMS_PER_THREAD + i);\n\n if (tid == 0)\n call_exchange[THREADBLOCK_SIZE] = expiryCallValue(S, X, vDt, NUM_STEPS);\n\n int final_it = max(0, tid * ELEMS_PER_THREAD - 1);\n\n#pragma unroll 16\n for (int i = NUM_STEPS; i > 0; --i) {\n call_exchange[tid] = call[0];\n cg::sync(cta);\n call[ELEMS_PER_THREAD] = call_exchange[tid + 1];\n cg::sync(cta);\n\n if (i > final_it) {\n#pragma unroll\n for (int j = 0; j < ELEMS_PER_THREAD; ++j)\n call[j] = puByDf * call[j + 1] + pdByDf * call[j];\n }\n }\n\n if (tid == 0) {\n d_CallValue[blockIdx.x] = call[0];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Host-side interface to GPU binomialOptions\n////////////////////////////////////////////////////////////////////////////////\nextern \"C\" void binomialOptionsGPU(real *callValue, TOptionData *optionData, int optN)\n{\n __TOptionData h_OptionData[MAX_OPTIONS];\n\n for (int i = 0; i < optN; i++) {\n const real T = optionData[i].T;\n const real R = optionData[i].R;\n const real V = optionData[i].V;\n\n const real dt = T / (real)NUM_STEPS;\n const real vDt = V * sqrt(dt);\n const real rDt = R * dt;\n // Per-step interest and discount factors\n const real If = exp(rDt);\n const real Df = exp(-rDt);\n // Values and pseudoprobabilities of upward and downward moves\n const real u = exp(vDt);\n const real d = exp(-vDt);\n const real pu = (If - d) / (u - d);\n const real pd = (real)1.0 - pu;\n const real puByDf = pu * Df;\n const real pdByDf = pd * Df;\n\n h_OptionData[i].S = (real)optionData[i].S;\n h_OptionData[i].X = (real)optionData[i].X;\n h_OptionData[i].vDt = (real)vDt;\n h_OptionData[i].puByDf = (real)puByDf;\n h_OptionData[i].pdByDf = (real)pdByDf;\n }\n\n checkCudaErrors(hipMemcpyToSymbol(HIP_SYMBOL(d_OptionData), h_OptionData, optN * sizeof(__TOptionData)));\n binomialOptionsKernel<<>>();\n getLastCudaError(\"binomialOptionsKernel() execution failed.\\n\");\n checkCudaErrors(hipMemcpyFromSymbol(callValue,HIP_SYMBOL(d_CallValue), optN * sizeof(real)));\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/binomialOptions/binomialOptions_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n This example demonstrates how to use the Cuda OpenGL bindings to\n dynamically modify a vertex buffer using a Cuda kernel.\n\n The steps are:\n 1. Create an empty vertex buffer object (VBO)\n 2. Register the VBO with Cuda\n 3. Map the VBO for writing from Cuda\n 4. Run Cuda kernel to modify the vertex positions\n 5. Unmap the VBO\n 6. Render the results using OpenGL\n\n Host code\n*/\n\n// includes, system\n#include \n#include \n#include \n#include \n\n#ifdef _WIN32\n#define WINDOWS_LEAN_AND_MEAN\n#define NOMINMAX\n#include \n#endif\n\n// OpenGL Graphics includes\n#include \n#if defined(__APPLE__) || defined(MACOSX)\n#pragma clang diagnostic ignored \"-Wdeprecated-declarations\"\n#include \n#ifndef glutCloseFunc\n#define glutCloseFunc glutWMCloseFunc\n#endif\n#else\n#include \n#endif\n\n// includes, cuda\n#include \n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n#include \n\n#define MAX_EPSILON_ERROR 10.0f\n#define THRESHOLD 0.30f\n#define REFRESH_DELAY 10 // ms\n\n////////////////////////////////////////////////////////////////////////////////\n// constants\nconst unsigned int window_width = 512;\nconst unsigned int window_height = 512;\n\nconst unsigned int mesh_width = 256;\nconst unsigned int mesh_height = 256;\n\n// vbo variables\nGLuint vbo;\nstruct cudaGraphicsResource *cuda_vbo_resource;\nvoid *d_vbo_buffer = NULL;\n\nfloat g_fAnim = 0.0;\n\n// mouse controls\nint mouse_old_x, mouse_old_y;\nint mouse_buttons = 0;\nfloat rotate_x = 0.0, rotate_y = 0.0;\nfloat translate_z = -3.0;\n\nStopWatchInterface *timer = NULL;\n\n// Auto-Verification Code\nint fpsCount = 0; // FPS count for averaging\nint fpsLimit = 1; // FPS limit for sampling\nint g_Index = 0;\nfloat avgFPS = 0.0f;\nunsigned int frameCount = 0;\nunsigned int g_TotalErrors = 0;\nbool g_bQAReadback = false;\n\nint *pArgc = NULL;\nchar **pArgv = NULL;\n\n#define MAX(a, b) ((a > b) ? a : b)\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\nbool runTest(int argc, char **argv, char *ref_file);\nvoid cleanup();\n\n// GL functionality\nbool initGL(int *argc, char **argv);\nvoid createVBO(GLuint *vbo, struct cudaGraphicsResource **vbo_res, unsigned int vbo_res_flags);\nvoid deleteVBO(GLuint *vbo, struct cudaGraphicsResource *vbo_res);\n\n// rendering callbacks\nvoid display();\nvoid keyboard(unsigned char key, int x, int y);\nvoid mouse(int button, int state, int x, int y);\nvoid motion(int x, int y);\nvoid timerEvent(int value);\n\n// Cuda functionality\nvoid runCuda(struct cudaGraphicsResource **vbo_resource);\nvoid runAutoTest(int devID, char **argv, char *ref_file);\nvoid checkResultCuda(int argc, char **argv, const GLuint &vbo);\n\nconst char *sSDKsample = \"simpleGL (VBO)\";\n\n///////////////////////////////////////////////////////////////////////////////\n//! Simple kernel to modify vertex positions in sine wave pattern\n//! @param data data in global memory\n///////////////////////////////////////////////////////////////////////////////\n__global__ void simple_vbo_kernel(float4 *pos, unsigned int width, unsigned int height, float time)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // calculate uv coordinates\n float u = x / (float)width;\n float v = y / (float)height;\n u = u * 2.0f - 1.0f;\n v = v * 2.0f - 1.0f;\n\n // calculate simple sine wave pattern\n float freq = 4.0f;\n float w = sinf(u * freq + time) * cosf(v * freq + time) * 0.5f;\n\n // write output vertex\n pos[y * width + x] = make_float4(u, w, v, 1.0f);\n}\n\n\nvoid launch_kernel(float4 *pos, unsigned int mesh_width, unsigned int mesh_height, float time)\n{\n // execute the kernel\n dim3 block(8, 8, 1);\n dim3 grid(mesh_width / block.x, mesh_height / block.y, 1);\n simple_vbo_kernel<<>>(pos, mesh_width, mesh_height, time);\n}\n\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n char *ref_file = NULL;\n\n pArgc = &argc;\n pArgv = argv;\n\n#if defined(__linux__)\n setenv(\"DISPLAY\", \":0\", 0);\n#endif\n\n printf(\"%s starting...\\n\", sSDKsample);\n\n if (argc > 1) {\n if (checkCmdLineFlag(argc, (const char **)argv, \"file\")) {\n // In this mode, we are running non-OpenGL and doing a compare of the VBO was generated correctly\n getCmdLineArgumentString(argc, (const char **)argv, \"file\", (char **)&ref_file);\n }\n }\n\n printf(\"\\n\");\n\n runTest(argc, argv, ref_file);\n\n printf(\"%s completed, returned %s\\n\", sSDKsample, (g_TotalErrors == 0) ? \"OK\" : \"ERROR!\");\n exit(g_TotalErrors == 0 ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\nvoid computeFPS()\n{\n frameCount++;\n fpsCount++;\n\n if (fpsCount == fpsLimit) {\n avgFPS = 1.f / (sdkGetAverageTimerValue(&timer) / 1000.f);\n fpsCount = 0;\n fpsLimit = (int)MAX(avgFPS, 1.f);\n\n sdkResetTimer(&timer);\n }\n\n char fps[256];\n sprintf(fps, \"Cuda GL Interop (VBO): %3.1f fps (Max 100Hz)\", avgFPS);\n glutSetWindowTitle(fps);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Initialize GL\n////////////////////////////////////////////////////////////////////////////////\nbool initGL(int *argc, char **argv)\n{\n glutInit(argc, argv);\n glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE);\n glutInitWindowSize(window_width, window_height);\n glutCreateWindow(\"Cuda GL Interop (VBO)\");\n glutDisplayFunc(display);\n glutKeyboardFunc(keyboard);\n glutMotionFunc(motion);\n glutTimerFunc(REFRESH_DELAY, timerEvent, 0);\n\n // initialize necessary OpenGL extensions\n if (!isGLVersionSupported(2, 0)) {\n fprintf(stderr, \"ERROR: Support for necessary OpenGL extensions missing.\");\n fflush(stderr);\n return false;\n }\n\n // default initialization\n glClearColor(0.0, 0.0, 0.0, 1.0);\n glDisable(GL_DEPTH_TEST);\n\n // viewport\n glViewport(0, 0, window_width, window_height);\n\n // projection\n glMatrixMode(GL_PROJECTION);\n glLoadIdentity();\n gluPerspective(60.0, (GLfloat)window_width / (GLfloat)window_height, 0.1, 10.0);\n\n SDK_CHECK_ERROR_GL();\n\n return true;\n}\n\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nbool runTest(int argc, char **argv, char *ref_file)\n{\n // Create the CUTIL timer\n sdkCreateTimer(&timer);\n\n // use command-line specified CUDA device, otherwise use device with highest Gflops/s\n int devID = findCudaDevice(argc, (const char **)argv);\n\n // command line mode only\n if (ref_file != NULL) {\n // create VBO\n checkCudaErrors(cudaMalloc((void **)&d_vbo_buffer, mesh_width * mesh_height * 4 * sizeof(float)));\n\n // run the cuda part\n runAutoTest(devID, argv, ref_file);\n\n // check result of Cuda step\n checkResultCuda(argc, argv, vbo);\n\n cudaFree(d_vbo_buffer);\n d_vbo_buffer = NULL;\n }\n else {\n // First initialize OpenGL context, so we can properly set the GL for CUDA.\n // This is necessary in order to achieve optimal performance with OpenGL/CUDA interop.\n if (false == initGL(&argc, argv)) {\n return false;\n }\n\n // register callbacks\n glutDisplayFunc(display);\n glutKeyboardFunc(keyboard);\n glutMouseFunc(mouse);\n glutMotionFunc(motion);\n#if defined(__APPLE__) || defined(MACOSX)\n atexit(cleanup);\n#else\n glutCloseFunc(cleanup);\n#endif\n\n // create VBO\n createVBO(&vbo, &cuda_vbo_resource, cudaGraphicsMapFlagsWriteDiscard);\n\n // run the cuda part\n runCuda(&cuda_vbo_resource);\n\n // start rendering mainloop\n glutMainLoop();\n }\n\n return true;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run the Cuda part of the computation\n////////////////////////////////////////////////////////////////////////////////\nvoid runCuda(struct cudaGraphicsResource **vbo_resource)\n{\n // map OpenGL buffer object for writing from CUDA\n float4 *dptr;\n checkCudaErrors(cudaGraphicsMapResources(1, vbo_resource, 0));\n size_t num_bytes;\n checkCudaErrors(cudaGraphicsResourceGetMappedPointer((void **)&dptr, &num_bytes, *vbo_resource));\n // printf(\"CUDA mapped VBO: May access %ld bytes\\n\", num_bytes);\n\n // execute the kernel\n // dim3 block(8, 8, 1);\n // dim3 grid(mesh_width / block.x, mesh_height / block.y, 1);\n // kernel<<< grid, block>>>(dptr, mesh_width, mesh_height, g_fAnim);\n\n launch_kernel(dptr, mesh_width, mesh_height, g_fAnim);\n\n // unmap buffer object\n checkCudaErrors(cudaGraphicsUnmapResources(1, vbo_resource, 0));\n}\n\n#ifdef _WIN32\n#ifndef FOPEN\n#define FOPEN(fHandle, filename, mode) fopen_s(&fHandle, filename, mode)\n#endif\n#else\n#ifndef FOPEN\n#define FOPEN(fHandle, filename, mode) (fHandle = fopen(filename, mode))\n#endif\n#endif\n\nvoid sdkDumpBin2(void *data, unsigned int bytes, const char *filename)\n{\n printf(\"sdkDumpBin: <%s>\\n\", filename);\n FILE *fp;\n FOPEN(fp, filename, \"wb\");\n fwrite(data, bytes, 1, fp);\n fflush(fp);\n fclose(fp);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run the Cuda part of the computation\n////////////////////////////////////////////////////////////////////////////////\nvoid runAutoTest(int devID, char **argv, char *ref_file)\n{\n char *reference_file = NULL;\n void *imageData = malloc(mesh_width * mesh_height * sizeof(float));\n\n // execute the kernel\n launch_kernel((float4 *)d_vbo_buffer, mesh_width, mesh_height, g_fAnim);\n\n cudaDeviceSynchronize();\n getLastCudaError(\"launch_kernel failed\");\n\n checkCudaErrors(\n cudaMemcpy(imageData, d_vbo_buffer, mesh_width * mesh_height * sizeof(float), cudaMemcpyDeviceToHost));\n\n sdkDumpBin2(imageData, mesh_width * mesh_height * sizeof(float), \"simpleGL.bin\");\n reference_file = sdkFindFilePath(ref_file, argv[0]);\n\n if (reference_file\n && !sdkCompareBin2BinFloat(\"simpleGL.bin\",\n reference_file,\n mesh_width * mesh_height * sizeof(float),\n MAX_EPSILON_ERROR,\n THRESHOLD,\n pArgv[0])) {\n g_TotalErrors++;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Create VBO\n////////////////////////////////////////////////////////////////////////////////\nvoid createVBO(GLuint *vbo, struct cudaGraphicsResource **vbo_res, unsigned int vbo_res_flags)\n{\n assert(vbo);\n\n // create buffer object\n glGenBuffers(1, vbo);\n glBindBuffer(GL_ARRAY_BUFFER, *vbo);\n\n // initialize buffer object\n unsigned int size = mesh_width * mesh_height * 4 * sizeof(float);\n glBufferData(GL_ARRAY_BUFFER, size, 0, GL_DYNAMIC_DRAW);\n\n glBindBuffer(GL_ARRAY_BUFFER, 0);\n\n // register this buffer object with CUDA\n checkCudaErrors(cudaGraphicsGLRegisterBuffer(vbo_res, *vbo, vbo_res_flags));\n\n SDK_CHECK_ERROR_GL();\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Delete VBO\n////////////////////////////////////////////////////////////////////////////////\nvoid deleteVBO(GLuint *vbo, struct cudaGraphicsResource *vbo_res)\n{\n\n // unregister this buffer object with CUDA\n checkCudaErrors(cudaGraphicsUnregisterResource(vbo_res));\n\n glBindBuffer(1, *vbo);\n glDeleteBuffers(1, vbo);\n\n *vbo = 0;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Display callback\n////////////////////////////////////////////////////////////////////////////////\nvoid display()\n{\n sdkStartTimer(&timer);\n\n // run CUDA kernel to generate vertex positions\n runCuda(&cuda_vbo_resource);\n\n glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);\n\n // set view matrix\n glMatrixMode(GL_MODELVIEW);\n glLoadIdentity();\n glTranslatef(0.0, 0.0, translate_z);\n glRotatef(rotate_x, 1.0, 0.0, 0.0);\n glRotatef(rotate_y, 0.0, 1.0, 0.0);\n\n // render from the vbo\n glBindBuffer(GL_ARRAY_BUFFER, vbo);\n glVertexPointer(4, GL_FLOAT, 0, 0);\n\n glEnableClientState(GL_VERTEX_ARRAY);\n glColor3f(1.0, 0.0, 0.0);\n glDrawArrays(GL_POINTS, 0, mesh_width * mesh_height);\n glDisableClientState(GL_VERTEX_ARRAY);\n\n glutSwapBuffers();\n\n g_fAnim += 0.01f;\n\n sdkStopTimer(&timer);\n computeFPS();\n}\n\nvoid timerEvent(int value)\n{\n if (glutGetWindow()) {\n glutPostRedisplay();\n glutTimerFunc(REFRESH_DELAY, timerEvent, 0);\n }\n}\n\nvoid cleanup()\n{\n sdkDeleteTimer(&timer);\n\n if (vbo) {\n deleteVBO(&vbo, cuda_vbo_resource);\n }\n}\n\n\n////////////////////////////////////////////////////////////////////////////////\n//! Keyboard events handler\n////////////////////////////////////////////////////////////////////////////////\nvoid keyboard(unsigned char key, int /*x*/, int /*y*/)\n{\n switch (key) {\n case (27):\n#if defined(__APPLE__) || defined(MACOSX)\n exit(EXIT_SUCCESS);\n#else\n glutDestroyWindow(glutGetWindow());\n return;\n#endif\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Mouse event handlers\n////////////////////////////////////////////////////////////////////////////////\nvoid mouse(int button, int state, int x, int y)\n{\n if (state == GLUT_DOWN) {\n mouse_buttons |= 1 << button;\n }\n else if (state == GLUT_UP) {\n mouse_buttons = 0;\n }\n\n mouse_old_x = x;\n mouse_old_y = y;\n}\n\nvoid motion(int x, int y)\n{\n float dx, dy;\n dx = (float)(x - mouse_old_x);\n dy = (float)(y - mouse_old_y);\n\n if (mouse_buttons & 1) {\n rotate_x += dy * 0.2f;\n rotate_y += dx * 0.2f;\n }\n else if (mouse_buttons & 4) {\n translate_z += dy * 0.01f;\n }\n\n mouse_old_x = x;\n mouse_old_y = y;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Check if the result is correct or write data to file for external\n//! regression testing\n////////////////////////////////////////////////////////////////////////////////\nvoid checkResultCuda(int argc, char **argv, const GLuint &vbo)\n{\n if (!d_vbo_buffer) {\n checkCudaErrors(cudaGraphicsUnregisterResource(cuda_vbo_resource));\n\n // map buffer object\n glBindBuffer(GL_ARRAY_BUFFER, vbo);\n float *data = (float *)glMapBuffer(GL_ARRAY_BUFFER, GL_READ_ONLY);\n\n // check result\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // write file for regression test\n sdkWriteFile(\"./data/regression.dat\", data, mesh_width * mesh_height * 3, 0.0, false);\n }\n\n // unmap GL buffer object\n if (!glUnmapBuffer(GL_ARRAY_BUFFER)) {\n fprintf(stderr, \"Unmap buffer failed.\\n\");\n fflush(stderr);\n }\n\n checkCudaErrors(cudaGraphicsGLRegisterBuffer(&cuda_vbo_resource, vbo, cudaGraphicsMapFlagsWriteDiscard));\n\n SDK_CHECK_ERROR_GL();\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n This example demonstrates how to use the Cuda OpenGL bindings to\n dynamically modify a vertex buffer using a Cuda kernel.\n\n The steps are:\n 1. Create an empty vertex buffer object (VBO)\n 2. Register the VBO with Cuda\n 3. Map the VBO for writing from Cuda\n 4. Run Cuda kernel to modify the vertex positions\n 5. Unmap the VBO\n 6. Render the results using OpenGL\n\n Host code\n*/\n\n// includes, system\n#include \n#include \n#include \n#include \n\n#ifdef _WIN32\n#define WINDOWS_LEAN_AND_MEAN\n#define NOMINMAX\n#include \n#endif\n\n// OpenGL Graphics includes\n#include \n#if defined(__APPLE__) || defined(MACOSX)\n#pragma clang diagnostic ignored \"-Wdeprecated-declarations\"\n#include \n#ifndef glutCloseFunc\n#define glutCloseFunc glutWMCloseFunc\n#endif\n#else\n#include \n#endif\n\n// includes, cuda\n#include \n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n#include \n\n#define MAX_EPSILON_ERROR 10.0f\n#define THRESHOLD 0.30f\n#define REFRESH_DELAY 10 // ms\n\n////////////////////////////////////////////////////////////////////////////////\n// constants\nconst unsigned int window_width = 512;\nconst unsigned int window_height = 512;\n\nconst unsigned int mesh_width = 256;\nconst unsigned int mesh_height = 256;\n\n// vbo variables\nGLuint vbo;\nstruct hipGraphicsResource *cuda_vbo_resource;\nvoid *d_vbo_buffer = NULL;\n\nfloat g_fAnim = 0.0;\n\n// mouse controls\nint mouse_old_x, mouse_old_y;\nint mouse_buttons = 0;\nfloat rotate_x = 0.0, rotate_y = 0.0;\nfloat translate_z = -3.0;\n\nStopWatchInterface *timer = NULL;\n\n// Auto-Verification Code\nint fpsCount = 0; // FPS count for averaging\nint fpsLimit = 1; // FPS limit for sampling\nint g_Index = 0;\nfloat avgFPS = 0.0f;\nunsigned int frameCount = 0;\nunsigned int g_TotalErrors = 0;\nbool g_bQAReadback = false;\n\nint *pArgc = NULL;\nchar **pArgv = NULL;\n\n#define MAX(a, b) ((a > b) ? a : b)\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\nbool runTest(int argc, char **argv, char *ref_file);\nvoid cleanup();\n\n// GL functionality\nbool initGL(int *argc, char **argv);\nvoid createVBO(GLuint *vbo, struct hipGraphicsResource **vbo_res, unsigned int vbo_res_flags);\nvoid deleteVBO(GLuint *vbo, struct hipGraphicsResource *vbo_res);\n\n// rendering callbacks\nvoid display();\nvoid keyboard(unsigned char key, int x, int y);\nvoid mouse(int button, int state, int x, int y);\nvoid motion(int x, int y);\nvoid timerEvent(int value);\n\n// Cuda functionality\nvoid runCuda(struct hipGraphicsResource **vbo_resource);\nvoid runAutoTest(int devID, char **argv, char *ref_file);\nvoid checkResultCuda(int argc, char **argv, const GLuint &vbo);\n\nconst char *sSDKsample = \"simpleGL (VBO)\";\n\n///////////////////////////////////////////////////////////////////////////////\n//! Simple kernel to modify vertex positions in sine wave pattern\n//! @param data data in global memory\n///////////////////////////////////////////////////////////////////////////////\n__global__ void simple_vbo_kernel(float4 *pos, unsigned int width, unsigned int height, float time)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // calculate uv coordinates\n float u = x / (float)width;\n float v = y / (float)height;\n u = u * 2.0f - 1.0f;\n v = v * 2.0f - 1.0f;\n\n // calculate simple sine wave pattern\n float freq = 4.0f;\n float w = sinf(u * freq + time) * cosf(v * freq + time) * 0.5f;\n\n // write output vertex\n pos[y * width + x] = make_float4(u, w, v, 1.0f);\n}\n\n\nvoid launch_kernel(float4 *pos, unsigned int mesh_width, unsigned int mesh_height, float time)\n{\n // execute the kernel\n dim3 block(8, 8, 1);\n dim3 grid(mesh_width / block.x, mesh_height / block.y, 1);\n simple_vbo_kernel<<>>(pos, mesh_width, mesh_height, time);\n}\n\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n char *ref_file = NULL;\n\n pArgc = &argc;\n pArgv = argv;\n\n#if defined(__linux__)\n setenv(\"DISPLAY\", \":0\", 0);\n#endif\n\n printf(\"%s starting...\\n\", sSDKsample);\n\n if (argc > 1) {\n if (checkCmdLineFlag(argc, (const char **)argv, \"file\")) {\n // In this mode, we are running non-OpenGL and doing a compare of the VBO was generated correctly\n getCmdLineArgumentString(argc, (const char **)argv, \"file\", (char **)&ref_file);\n }\n }\n\n printf(\"\\n\");\n\n runTest(argc, argv, ref_file);\n\n printf(\"%s completed, returned %s\\n\", sSDKsample, (g_TotalErrors == 0) ? \"OK\" : \"ERROR!\");\n exit(g_TotalErrors == 0 ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\nvoid computeFPS()\n{\n frameCount++;\n fpsCount++;\n\n if (fpsCount == fpsLimit) {\n avgFPS = 1.f / (sdkGetAverageTimerValue(&timer) / 1000.f);\n fpsCount = 0;\n fpsLimit = (int)MAX(avgFPS, 1.f);\n\n sdkResetTimer(&timer);\n }\n\n char fps[256];\n sprintf(fps, \"Cuda GL Interop (VBO): %3.1f fps (Max 100Hz)\", avgFPS);\n glutSetWindowTitle(fps);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Initialize GL\n////////////////////////////////////////////////////////////////////////////////\nbool initGL(int *argc, char **argv)\n{\n glutInit(argc, argv);\n glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE);\n glutInitWindowSize(window_width, window_height);\n glutCreateWindow(\"Cuda GL Interop (VBO)\");\n glutDisplayFunc(display);\n glutKeyboardFunc(keyboard);\n glutMotionFunc(motion);\n glutTimerFunc(REFRESH_DELAY, timerEvent, 0);\n\n // initialize necessary OpenGL extensions\n if (!isGLVersionSupported(2, 0)) {\n fprintf(stderr, \"ERROR: Support for necessary OpenGL extensions missing.\");\n fflush(stderr);\n return false;\n }\n\n // default initialization\n glClearColor(0.0, 0.0, 0.0, 1.0);\n glDisable(GL_DEPTH_TEST);\n\n // viewport\n glViewport(0, 0, window_width, window_height);\n\n // projection\n glMatrixMode(GL_PROJECTION);\n glLoadIdentity();\n gluPerspective(60.0, (GLfloat)window_width / (GLfloat)window_height, 0.1, 10.0);\n\n SDK_CHECK_ERROR_GL();\n\n return true;\n}\n\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nbool runTest(int argc, char **argv, char *ref_file)\n{\n // Create the CUTIL timer\n sdkCreateTimer(&timer);\n\n // use command-line specified CUDA device, otherwise use device with highest Gflops/s\n int devID = findCudaDevice(argc, (const char **)argv);\n\n // command line mode only\n if (ref_file != NULL) {\n // create VBO\n checkCudaErrors(hipMalloc((void **)&d_vbo_buffer, mesh_width * mesh_height * 4 * sizeof(float)));\n\n // run the cuda part\n runAutoTest(devID, argv, ref_file);\n\n // check result of Cuda step\n checkResultCuda(argc, argv, vbo);\n\n hipFree(d_vbo_buffer);\n d_vbo_buffer = NULL;\n }\n else {\n // First initialize OpenGL context, so we can properly set the GL for CUDA.\n // This is necessary in order to achieve optimal performance with OpenGL/CUDA interop.\n if (false == initGL(&argc, argv)) {\n return false;\n }\n\n // register callbacks\n glutDisplayFunc(display);\n glutKeyboardFunc(keyboard);\n glutMouseFunc(mouse);\n glutMotionFunc(motion);\n#if defined(__APPLE__) || defined(MACOSX)\n atexit(cleanup);\n#else\n glutCloseFunc(cleanup);\n#endif\n\n // create VBO\n createVBO(&vbo, &cuda_vbo_resource, cudaGraphicsMapFlagsWriteDiscard);\n\n // run the cuda part\n runCuda(&cuda_vbo_resource);\n\n // start rendering mainloop\n glutMainLoop();\n }\n\n return true;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run the Cuda part of the computation\n////////////////////////////////////////////////////////////////////////////////\nvoid runCuda(struct hipGraphicsResource **vbo_resource)\n{\n // map OpenGL buffer object for writing from CUDA\n float4 *dptr;\n checkCudaErrors(hipGraphicsMapResources(1, vbo_resource, 0));\n size_t num_bytes;\n checkCudaErrors(hipGraphicsResourceGetMappedPointer((void **)&dptr, &num_bytes, *vbo_resource));\n // printf(\"CUDA mapped VBO: May access %ld bytes\\n\", num_bytes);\n\n // execute the kernel\n // dim3 block(8, 8, 1);\n // dim3 grid(mesh_width / block.x, mesh_height / block.y, 1);\n // kernel<<< grid, block>>>(dptr, mesh_width, mesh_height, g_fAnim);\n\n launch_kernel(dptr, mesh_width, mesh_height, g_fAnim);\n\n // unmap buffer object\n checkCudaErrors(hipGraphicsUnmapResources(1, vbo_resource, 0));\n}\n\n#ifdef _WIN32\n#ifndef FOPEN\n#define FOPEN(fHandle, filename, mode) fopen_s(&fHandle, filename, mode)\n#endif\n#else\n#ifndef FOPEN\n#define FOPEN(fHandle, filename, mode) (fHandle = fopen(filename, mode))\n#endif\n#endif\n\nvoid sdkDumpBin2(void *data, unsigned int bytes, const char *filename)\n{\n printf(\"sdkDumpBin: <%s>\\n\", filename);\n FILE *fp;\n FOPEN(fp, filename, \"wb\");\n fwrite(data, bytes, 1, fp);\n fflush(fp);\n fclose(fp);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run the Cuda part of the computation\n////////////////////////////////////////////////////////////////////////////////\nvoid runAutoTest(int devID, char **argv, char *ref_file)\n{\n char *reference_file = NULL;\n void *imageData = malloc(mesh_width * mesh_height * sizeof(float));\n\n // execute the kernel\n launch_kernel((float4 *)d_vbo_buffer, mesh_width, mesh_height, g_fAnim);\n\n hipDeviceSynchronize();\n getLastCudaError(\"launch_kernel failed\");\n\n checkCudaErrors(\n hipMemcpy(imageData, d_vbo_buffer, mesh_width * mesh_height * sizeof(float), hipMemcpyDeviceToHost));\n\n sdkDumpBin2(imageData, mesh_width * mesh_height * sizeof(float), \"simpleGL.bin\");\n reference_file = sdkFindFilePath(ref_file, argv[0]);\n\n if (reference_file\n && !sdkCompareBin2BinFloat(\"simpleGL.bin\",\n reference_file,\n mesh_width * mesh_height * sizeof(float),\n MAX_EPSILON_ERROR,\n THRESHOLD,\n pArgv[0])) {\n g_TotalErrors++;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Create VBO\n////////////////////////////////////////////////////////////////////////////////\nvoid createVBO(GLuint *vbo, struct hipGraphicsResource **vbo_res, unsigned int vbo_res_flags)\n{\n assert(vbo);\n\n // create buffer object\n glGenBuffers(1, vbo);\n glBindBuffer(GL_ARRAY_BUFFER, *vbo);\n\n // initialize buffer object\n unsigned int size = mesh_width * mesh_height * 4 * sizeof(float);\n glBufferData(GL_ARRAY_BUFFER, size, 0, GL_DYNAMIC_DRAW);\n\n glBindBuffer(GL_ARRAY_BUFFER, 0);\n\n // register this buffer object with CUDA\n checkCudaErrors(hipGraphicsGLRegisterBuffer(vbo_res, *vbo, vbo_res_flags));\n\n SDK_CHECK_ERROR_GL();\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Delete VBO\n////////////////////////////////////////////////////////////////////////////////\nvoid deleteVBO(GLuint *vbo, struct hipGraphicsResource *vbo_res)\n{\n\n // unregister this buffer object with CUDA\n checkCudaErrors(hipGraphicsUnregisterResource(vbo_res));\n\n glBindBuffer(1, *vbo);\n glDeleteBuffers(1, vbo);\n\n *vbo = 0;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Display callback\n////////////////////////////////////////////////////////////////////////////////\nvoid display()\n{\n sdkStartTimer(&timer);\n\n // run CUDA kernel to generate vertex positions\n runCuda(&cuda_vbo_resource);\n\n glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);\n\n // set view matrix\n glMatrixMode(GL_MODELVIEW);\n glLoadIdentity();\n glTranslatef(0.0, 0.0, translate_z);\n glRotatef(rotate_x, 1.0, 0.0, 0.0);\n glRotatef(rotate_y, 0.0, 1.0, 0.0);\n\n // render from the vbo\n glBindBuffer(GL_ARRAY_BUFFER, vbo);\n glVertexPointer(4, GL_FLOAT, 0, 0);\n\n glEnableClientState(GL_VERTEX_ARRAY);\n glColor3f(1.0, 0.0, 0.0);\n glDrawArrays(GL_POINTS, 0, mesh_width * mesh_height);\n glDisableClientState(GL_VERTEX_ARRAY);\n\n glutSwapBuffers();\n\n g_fAnim += 0.01f;\n\n sdkStopTimer(&timer);\n computeFPS();\n}\n\nvoid timerEvent(int value)\n{\n if (glutGetWindow()) {\n glutPostRedisplay();\n glutTimerFunc(REFRESH_DELAY, timerEvent, 0);\n }\n}\n\nvoid cleanup()\n{\n sdkDeleteTimer(&timer);\n\n if (vbo) {\n deleteVBO(&vbo, cuda_vbo_resource);\n }\n}\n\n\n////////////////////////////////////////////////////////////////////////////////\n//! Keyboard events handler\n////////////////////////////////////////////////////////////////////////////////\nvoid keyboard(unsigned char key, int /*x*/, int /*y*/)\n{\n switch (key) {\n case (27):\n#if defined(__APPLE__) || defined(MACOSX)\n exit(EXIT_SUCCESS);\n#else\n glutDestroyWindow(glutGetWindow());\n return;\n#endif\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Mouse event handlers\n////////////////////////////////////////////////////////////////////////////////\nvoid mouse(int button, int state, int x, int y)\n{\n if (state == GLUT_DOWN) {\n mouse_buttons |= 1 << button;\n }\n else if (state == GLUT_UP) {\n mouse_buttons = 0;\n }\n\n mouse_old_x = x;\n mouse_old_y = y;\n}\n\nvoid motion(int x, int y)\n{\n float dx, dy;\n dx = (float)(x - mouse_old_x);\n dy = (float)(y - mouse_old_y);\n\n if (mouse_buttons & 1) {\n rotate_x += dy * 0.2f;\n rotate_y += dx * 0.2f;\n }\n else if (mouse_buttons & 4) {\n translate_z += dy * 0.01f;\n }\n\n mouse_old_x = x;\n mouse_old_y = y;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Check if the result is correct or write data to file for external\n//! regression testing\n////////////////////////////////////////////////////////////////////////////////\nvoid checkResultCuda(int argc, char **argv, const GLuint &vbo)\n{\n if (!d_vbo_buffer) {\n checkCudaErrors(hipGraphicsUnregisterResource(cuda_vbo_resource));\n\n // map buffer object\n glBindBuffer(GL_ARRAY_BUFFER, vbo);\n float *data = (float *)glMapBuffer(GL_ARRAY_BUFFER, GL_READ_ONLY);\n\n // check result\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // write file for regression test\n sdkWriteFile(\"./data/regression.dat\", data, mesh_width * mesh_height * 3, 0.0, false);\n }\n\n // unmap GL buffer object\n if (!glUnmapBuffer(GL_ARRAY_BUFFER)) {\n fprintf(stderr, \"Unmap buffer failed.\\n\");\n fflush(stderr);\n }\n\n checkCudaErrors(hipGraphicsGLRegisterBuffer(&cuda_vbo_resource, vbo, cudaGraphicsMapFlagsWriteDiscard));\n\n SDK_CHECK_ERROR_GL();\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/simpleGL/simpleGL.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Simple 3D volume renderer\n\n#ifndef _VOLUMERENDER_KERNEL_CU_\n#define _VOLUMERENDER_KERNEL_CU_\n\n#include \n#include \n\ntypedef unsigned int uint;\ntypedef unsigned char uchar;\n\ncudaArray *d_volumeArray = 0;\ncudaArray *d_transferFuncArray;\n\ntypedef unsigned char VolumeType;\n// typedef unsigned short VolumeType;\n\ncudaTextureObject_t texObject; // For 3D texture\ncudaTextureObject_t transferTex; // For 1D transfer function texture\n\ntypedef struct\n{\n float4 m[3];\n} float3x4;\n\n__constant__ float3x4 c_invViewMatrix; // inverse view matrix\n\nstruct Ray\n{\n float3 o; // origin\n float3 d; // direction\n};\n\n// intersect ray with a box\n// http://www.siggraph.org/education/materials/HyperGraph/raytrace/rtinter3.htm\n\n__device__ int intersectBox(Ray r, float3 boxmin, float3 boxmax, float *tnear, float *tfar)\n{\n // compute intersection of ray with all six bbox planes\n float3 invR = make_float3(1.0f) / r.d;\n float3 tbot = invR * (boxmin - r.o);\n float3 ttop = invR * (boxmax - r.o);\n\n // re-order intersections to find smallest and largest on each axis\n float3 tmin = fminf(ttop, tbot);\n float3 tmax = fmaxf(ttop, tbot);\n\n // find the largest tmin and the smallest tmax\n float largest_tmin = fmaxf(fmaxf(tmin.x, tmin.y), fmaxf(tmin.x, tmin.z));\n float smallest_tmax = fminf(fminf(tmax.x, tmax.y), fminf(tmax.x, tmax.z));\n\n *tnear = largest_tmin;\n *tfar = smallest_tmax;\n\n return smallest_tmax > largest_tmin;\n}\n\n// transform vector by matrix (no translation)\n__device__ float3 mul(const float3x4 &M, const float3 &v)\n{\n float3 r;\n r.x = dot(v, make_float3(M.m[0]));\n r.y = dot(v, make_float3(M.m[1]));\n r.z = dot(v, make_float3(M.m[2]));\n return r;\n}\n\n// transform vector by matrix with translation\n__device__ float4 mul(const float3x4 &M, const float4 &v)\n{\n float4 r;\n r.x = dot(v, M.m[0]);\n r.y = dot(v, M.m[1]);\n r.z = dot(v, M.m[2]);\n r.w = 1.0f;\n return r;\n}\n\n__device__ uint rgbaFloatToInt(float4 rgba)\n{\n rgba.x = __saturatef(rgba.x); // clamp to [0.0, 1.0]\n rgba.y = __saturatef(rgba.y);\n rgba.z = __saturatef(rgba.z);\n rgba.w = __saturatef(rgba.w);\n return (uint(rgba.w * 255) << 24) | (uint(rgba.z * 255) << 16) | (uint(rgba.y * 255) << 8) | uint(rgba.x * 255);\n}\n\n__global__ void d_render(uint *d_output,\n uint imageW,\n uint imageH,\n float density,\n float brightness,\n float transferOffset,\n float transferScale,\n cudaTextureObject_t tex,\n cudaTextureObject_t transferTex)\n{\n const int maxSteps = 500;\n const float tstep = 0.01f;\n const float opacityThreshold = 0.95f;\n const float3 boxMin = make_float3(-1.0f, -1.0f, -1.0f);\n const float3 boxMax = make_float3(1.0f, 1.0f, 1.0f);\n\n uint x = blockIdx.x * blockDim.x + threadIdx.x;\n uint y = blockIdx.y * blockDim.y + threadIdx.y;\n\n if ((x >= imageW) || (y >= imageH))\n return;\n\n float u = (x / (float)imageW) * 2.0f - 1.0f;\n float v = (y / (float)imageH) * 2.0f - 1.0f;\n\n // calculate eye ray in world space\n Ray eyeRay;\n eyeRay.o = make_float3(mul(c_invViewMatrix, make_float4(0.0f, 0.0f, 0.0f, 1.0f)));\n eyeRay.d = normalize(make_float3(u, v, -2.0f));\n eyeRay.d = mul(c_invViewMatrix, eyeRay.d);\n\n // find intersection with box\n float tnear, tfar;\n int hit = intersectBox(eyeRay, boxMin, boxMax, &tnear, &tfar);\n\n if (!hit)\n return;\n\n if (tnear < 0.0f)\n tnear = 0.0f; // clamp to near plane\n\n // march along ray from front to back, accumulating color\n float4 sum = make_float4(0.0f);\n float t = tnear;\n float3 pos = eyeRay.o + eyeRay.d * tnear;\n float3 step = eyeRay.d * tstep;\n\n for (int i = 0; i < maxSteps; i++) {\n // read from 3D texture\n // remap position to [0, 1] coordinates\n float sample = tex3D(tex, pos.x * 0.5f + 0.5f, pos.y * 0.5f + 0.5f, pos.z * 0.5f + 0.5f);\n // sample *= 64.0f; // scale for 10-bit data\n\n // lookup in transfer function texture\n float4 col = tex1D(transferTex, (sample - transferOffset) * transferScale);\n col.w *= density;\n\n // \"under\" operator for back-to-front blending\n // sum = lerp(sum, col, col.w);\n\n // pre-multiply alpha\n col.x *= col.w;\n col.y *= col.w;\n col.z *= col.w;\n // \"over\" operator for front-to-back blending\n sum = sum + col * (1.0f - sum.w);\n\n // exit early if opaque\n if (sum.w > opacityThreshold)\n break;\n\n t += tstep;\n\n if (t > tfar)\n break;\n\n pos += step;\n }\n\n sum *= brightness;\n\n // write output color\n d_output[y * imageW + x] = rgbaFloatToInt(sum);\n}\n\nextern \"C\" void setTextureFilterMode(bool bLinearFilter)\n{\n if (texObject) {\n checkCudaErrors(cudaDestroyTextureObject(texObject));\n }\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = d_volumeArray;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = bLinearFilter ? cudaFilterModeLinear : cudaFilterModePoint;\n\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.addressMode[2] = cudaAddressModeWrap;\n\n texDescr.readMode = cudaReadModeNormalizedFloat;\n\n checkCudaErrors(cudaCreateTextureObject(&texObject, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void initCuda(void *h_volume, cudaExtent volumeSize)\n{\n // create 3D array\n cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc();\n checkCudaErrors(cudaMalloc3DArray(&d_volumeArray, &channelDesc, volumeSize));\n\n // copy data to 3D array\n cudaMemcpy3DParms copyParams = {0};\n copyParams.srcPtr =\n make_cudaPitchedPtr(h_volume, volumeSize.width * sizeof(VolumeType), volumeSize.width, volumeSize.height);\n copyParams.dstArray = d_volumeArray;\n copyParams.extent = volumeSize;\n copyParams.kind = cudaMemcpyHostToDevice;\n checkCudaErrors(cudaMemcpy3D(©Params));\n\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = d_volumeArray;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true; // access with normalized texture coordinates\n texDescr.filterMode = cudaFilterModeLinear; // linear interpolation\n\n texDescr.addressMode[0] = cudaAddressModeClamp; // clamp texture coordinates\n texDescr.addressMode[1] = cudaAddressModeClamp;\n texDescr.addressMode[2] = cudaAddressModeClamp;\n\n texDescr.readMode = cudaReadModeNormalizedFloat;\n\n checkCudaErrors(cudaCreateTextureObject(&texObject, &texRes, &texDescr, NULL));\n\n // create transfer function texture\n float4 transferFunc[] = {\n {\n 0.0,\n 0.0,\n 0.0,\n 0.0,\n },\n {\n 1.0,\n 0.0,\n 0.0,\n 1.0,\n },\n {\n 1.0,\n 0.5,\n 0.0,\n 1.0,\n },\n {\n 1.0,\n 1.0,\n 0.0,\n 1.0,\n },\n {\n 0.0,\n 1.0,\n 0.0,\n 1.0,\n },\n {\n 0.0,\n 1.0,\n 1.0,\n 1.0,\n },\n {\n 0.0,\n 0.0,\n 1.0,\n 1.0,\n },\n {\n 1.0,\n 0.0,\n 1.0,\n 1.0,\n },\n {\n 0.0,\n 0.0,\n 0.0,\n 0.0,\n },\n };\n\n cudaChannelFormatDesc channelDesc2 = cudaCreateChannelDesc();\n cudaArray *d_transferFuncArray;\n checkCudaErrors(cudaMallocArray(&d_transferFuncArray, &channelDesc2, sizeof(transferFunc) / sizeof(float4), 1));\n checkCudaErrors(cudaMemcpy2DToArray(\n d_transferFuncArray, 0, 0, transferFunc, 0, sizeof(transferFunc), 1, cudaMemcpyHostToDevice));\n\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = d_transferFuncArray;\n\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true; // access with normalized texture coordinates\n texDescr.filterMode = cudaFilterModeLinear;\n\n texDescr.addressMode[0] = cudaAddressModeClamp; // wrap texture coordinates\n\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&transferTex, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void freeCudaBuffers()\n{\n checkCudaErrors(cudaDestroyTextureObject(texObject));\n checkCudaErrors(cudaDestroyTextureObject(transferTex));\n checkCudaErrors(cudaFreeArray(d_volumeArray));\n checkCudaErrors(cudaFreeArray(d_transferFuncArray));\n}\n\nextern \"C\" void render_kernel(dim3 gridSize,\n dim3 blockSize,\n uint *d_output,\n uint imageW,\n uint imageH,\n float density,\n float brightness,\n float transferOffset,\n float transferScale)\n{\n d_render<<>>(\n d_output, imageW, imageH, density, brightness, transferOffset, transferScale, texObject, transferTex);\n}\n\nextern \"C\" void copyInvViewMatrix(float *invViewMatrix, size_t sizeofMatrix)\n{\n checkCudaErrors(cudaMemcpyToSymbol(c_invViewMatrix, invViewMatrix, sizeofMatrix));\n}\n\n#endif // #ifndef _VOLUMERENDER_KERNEL_CU_\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Simple 3D volume renderer\n\n#include \"hip/hip_runtime.h\"\n#ifndef _VOLUMERENDER_KERNEL_CU_\n#define _VOLUMERENDER_KERNEL_CU_\n\n#include \n#include \n\ntypedef unsigned int uint;\ntypedef unsigned char uchar;\n\nhipArray *d_volumeArray = 0;\nhipArray *d_transferFuncArray;\n\ntypedef unsigned char VolumeType;\n// typedef unsigned short VolumeType;\n\nhipTextureObject_t texObject; // For 3D texture\nhipTextureObject_t transferTex; // For 1D transfer function texture\n\ntypedef struct\n{\n float4 m[3];\n} float3x4;\n\n__constant__ float3x4 c_invViewMatrix; // inverse view matrix\n\nstruct Ray\n{\n float3 o; // origin\n float3 d; // direction\n};\n\n// intersect ray with a box\n// http://www.siggraph.org/education/materials/HyperGraph/raytrace/rtinter3.htm\n\n__device__ int intersectBox(Ray r, float3 boxmin, float3 boxmax, float *tnear, float *tfar)\n{\n // compute intersection of ray with all six bbox planes\n float3 invR = make_float3(1.0f) / r.d;\n float3 tbot = invR * (boxmin - r.o);\n float3 ttop = invR * (boxmax - r.o);\n\n // re-order intersections to find smallest and largest on each axis\n float3 tmin = fminf(ttop, tbot);\n float3 tmax = fmaxf(ttop, tbot);\n\n // find the largest tmin and the smallest tmax\n float largest_tmin = fmaxf(fmaxf(tmin.x, tmin.y), fmaxf(tmin.x, tmin.z));\n float smallest_tmax = fminf(fminf(tmax.x, tmax.y), fminf(tmax.x, tmax.z));\n\n *tnear = largest_tmin;\n *tfar = smallest_tmax;\n\n return smallest_tmax > largest_tmin;\n}\n\n// transform vector by matrix (no translation)\n__device__ float3 mul(const float3x4 &M, const float3 &v)\n{\n float3 r;\n r.x = dot(v, make_float3(M.m[0]));\n r.y = dot(v, make_float3(M.m[1]));\n r.z = dot(v, make_float3(M.m[2]));\n return r;\n}\n\n// transform vector by matrix with translation\n__device__ float4 mul(const float3x4 &M, const float4 &v)\n{\n float4 r;\n r.x = dot(v, M.m[0]);\n r.y = dot(v, M.m[1]);\n r.z = dot(v, M.m[2]);\n r.w = 1.0f;\n return r;\n}\n\n__device__ uint rgbaFloatToInt(float4 rgba)\n{\n rgba.x = __saturatef(rgba.x); // clamp to [0.0, 1.0]\n rgba.y = __saturatef(rgba.y);\n rgba.z = __saturatef(rgba.z);\n rgba.w = __saturatef(rgba.w);\n return (uint(rgba.w * 255) << 24) | (uint(rgba.z * 255) << 16) | (uint(rgba.y * 255) << 8) | uint(rgba.x * 255);\n}\n\n__global__ void d_render(uint *d_output,\n uint imageW,\n uint imageH,\n float density,\n float brightness,\n float transferOffset,\n float transferScale,\n hipTextureObject_t tex,\n hipTextureObject_t transferTex)\n{\n const int maxSteps = 500;\n const float tstep = 0.01f;\n const float opacityThreshold = 0.95f;\n const float3 boxMin = make_float3(-1.0f, -1.0f, -1.0f);\n const float3 boxMax = make_float3(1.0f, 1.0f, 1.0f);\n\n uint x = blockIdx.x * blockDim.x + threadIdx.x;\n uint y = blockIdx.y * blockDim.y + threadIdx.y;\n\n if ((x >= imageW) || (y >= imageH))\n return;\n\n float u = (x / (float)imageW) * 2.0f - 1.0f;\n float v = (y / (float)imageH) * 2.0f - 1.0f;\n\n // calculate eye ray in world space\n Ray eyeRay;\n eyeRay.o = make_float3(mul(c_invViewMatrix, make_float4(0.0f, 0.0f, 0.0f, 1.0f)));\n eyeRay.d = normalize(make_float3(u, v, -2.0f));\n eyeRay.d = mul(c_invViewMatrix, eyeRay.d);\n\n // find intersection with box\n float tnear, tfar;\n int hit = intersectBox(eyeRay, boxMin, boxMax, &tnear, &tfar);\n\n if (!hit)\n return;\n\n if (tnear < 0.0f)\n tnear = 0.0f; // clamp to near plane\n\n // march along ray from front to back, accumulating color\n float4 sum = make_float4(0.0f);\n float t = tnear;\n float3 pos = eyeRay.o + eyeRay.d * tnear;\n float3 step = eyeRay.d * tstep;\n\n for (int i = 0; i < maxSteps; i++) {\n // read from 3D texture\n // remap position to [0, 1] coordinates\n float sample = tex3D(tex, pos.x * 0.5f + 0.5f, pos.y * 0.5f + 0.5f, pos.z * 0.5f + 0.5f);\n // sample *= 64.0f; // scale for 10-bit data\n\n // lookup in transfer function texture\n float4 col = tex1D(transferTex, (sample - transferOffset) * transferScale);\n col.w *= density;\n\n // \"under\" operator for back-to-front blending\n // sum = lerp(sum, col, col.w);\n\n // pre-multiply alpha\n col.x *= col.w;\n col.y *= col.w;\n col.z *= col.w;\n // \"over\" operator for front-to-back blending\n sum = sum + col * (1.0f - sum.w);\n\n // exit early if opaque\n if (sum.w > opacityThreshold)\n break;\n\n t += tstep;\n\n if (t > tfar)\n break;\n\n pos += step;\n }\n\n sum *= brightness;\n\n // write output color\n d_output[y * imageW + x] = rgbaFloatToInt(sum);\n}\n\nextern \"C\" void setTextureFilterMode(bool bLinearFilter)\n{\n if (texObject) {\n checkCudaErrors(hipDestroyTextureObject(texObject));\n }\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = d_volumeArray;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = bLinearFilter ? hipFilterModeLinear : hipFilterModePoint;\n\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.addressMode[2] = hipAddressModeWrap;\n\n texDescr.readMode = hipReadModeNormalizedFloat;\n\n checkCudaErrors(hipCreateTextureObject(&texObject, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void initCuda(void *h_volume, hipExtent volumeSize)\n{\n // create 3D array\n hipChannelFormatDesc channelDesc = hipCreateChannelDesc();\n checkCudaErrors(hipMalloc3DArray(&d_volumeArray, &channelDesc, volumeSize));\n\n // copy data to 3D array\n hipMemcpy3DParms copyParams = {0};\n copyParams.srcPtr =\n make_hipPitchedPtr(h_volume, volumeSize.width * sizeof(VolumeType), volumeSize.width, volumeSize.height);\n copyParams.dstArray = d_volumeArray;\n copyParams.extent = volumeSize;\n copyParams.kind = hipMemcpyHostToDevice;\n checkCudaErrors(hipMemcpy3D(©Params));\n\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = d_volumeArray;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true; // access with normalized texture coordinates\n texDescr.filterMode = hipFilterModeLinear; // linear interpolation\n\n texDescr.addressMode[0] = hipAddressModeClamp; // clamp texture coordinates\n texDescr.addressMode[1] = hipAddressModeClamp;\n texDescr.addressMode[2] = hipAddressModeClamp;\n\n texDescr.readMode = hipReadModeNormalizedFloat;\n\n checkCudaErrors(hipCreateTextureObject(&texObject, &texRes, &texDescr, NULL));\n\n // create transfer function texture\n float4 transferFunc[] = {\n {\n 0.0,\n 0.0,\n 0.0,\n 0.0,\n },\n {\n 1.0,\n 0.0,\n 0.0,\n 1.0,\n },\n {\n 1.0,\n 0.5,\n 0.0,\n 1.0,\n },\n {\n 1.0,\n 1.0,\n 0.0,\n 1.0,\n },\n {\n 0.0,\n 1.0,\n 0.0,\n 1.0,\n },\n {\n 0.0,\n 1.0,\n 1.0,\n 1.0,\n },\n {\n 0.0,\n 0.0,\n 1.0,\n 1.0,\n },\n {\n 1.0,\n 0.0,\n 1.0,\n 1.0,\n },\n {\n 0.0,\n 0.0,\n 0.0,\n 0.0,\n },\n };\n\n hipChannelFormatDesc channelDesc2 = hipCreateChannelDesc();\n hipArray *d_transferFuncArray;\n checkCudaErrors(hipMallocArray(&d_transferFuncArray, &channelDesc2, sizeof(transferFunc) / sizeof(float4), 1));\n checkCudaErrors(hipMemcpy2DToArray(\n d_transferFuncArray, 0, 0, transferFunc, 0, sizeof(transferFunc), 1, hipMemcpyHostToDevice));\n\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = d_transferFuncArray;\n\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true; // access with normalized texture coordinates\n texDescr.filterMode = hipFilterModeLinear;\n\n texDescr.addressMode[0] = hipAddressModeClamp; // wrap texture coordinates\n\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&transferTex, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void freeCudaBuffers()\n{\n checkCudaErrors(hipDestroyTextureObject(texObject));\n checkCudaErrors(hipDestroyTextureObject(transferTex));\n checkCudaErrors(hipFreeArray(d_volumeArray));\n checkCudaErrors(hipFreeArray(d_transferFuncArray));\n}\n\nextern \"C\" void render_kernel(dim3 gridSize,\n dim3 blockSize,\n uint *d_output,\n uint imageW,\n uint imageH,\n float density,\n float brightness,\n float transferOffset,\n float transferScale)\n{\n d_render<<>>(\n d_output, imageW, imageH, density, brightness, transferOffset, transferScale, texObject, transferTex);\n}\n\nextern \"C\" void copyInvViewMatrix(float *invViewMatrix, size_t sizeofMatrix)\n{\n checkCudaErrors(hipMemcpyToSymbol(HIP_SYMBOL(c_invViewMatrix), invViewMatrix, sizeofMatrix));\n}\n\n#endif // #ifndef _VOLUMERENDER_KERNEL_CU_\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/volumeRender/volumeRender_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#define GLFW_INCLUDE_VULKAN\n#ifdef _WIN64\n// Add windows.h to the include path firstly as dependency for other Windows headers\n#include \n// Add other Windows headers\n#include \n#include \n#include \n#define _USE_MATH_DEFINES\n#endif\n\n#include \n#include \n#ifdef _WIN64\n#include \n#endif\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"linmath.h\"\n\n#define WIDTH 800\n#define HEIGHT 600\n\nconst int MAX_FRAMES = 4;\n\nconst std::vector validationLayers = {\"VK_LAYER_KHRONOS_validation\"};\n\n#ifdef NDEBUG\nconst bool enableValidationLayers = false;\n#else\nconst bool enableValidationLayers = true;\n#endif\n\nstd::string execution_path;\n\nVkResult CreateDebugUtilsMessengerEXT(VkInstance instance,\n const VkDebugUtilsMessengerCreateInfoEXT *pCreateInfo,\n const VkAllocationCallbacks *pAllocator,\n VkDebugUtilsMessengerEXT *pDebugMessenger)\n{\n auto func = (PFN_vkCreateDebugUtilsMessengerEXT)vkGetInstanceProcAddr(instance, \"vkCreateDebugUtilsMessengerEXT\");\n if (func != nullptr) {\n return func(instance, pCreateInfo, pAllocator, pDebugMessenger);\n }\n else {\n return VK_ERROR_EXTENSION_NOT_PRESENT;\n }\n};\n\nconst std::vector deviceExtensions = {\n VK_KHR_SWAPCHAIN_EXTENSION_NAME,\n VK_KHR_EXTERNAL_MEMORY_EXTENSION_NAME,\n VK_KHR_EXTERNAL_SEMAPHORE_EXTENSION_NAME,\n#ifdef _WIN64\n VK_KHR_EXTERNAL_MEMORY_WIN32_EXTENSION_NAME,\n VK_KHR_EXTERNAL_SEMAPHORE_WIN32_EXTENSION_NAME,\n#else\n VK_KHR_EXTERNAL_MEMORY_FD_EXTENSION_NAME,\n VK_KHR_EXTERNAL_SEMAPHORE_FD_EXTENSION_NAME,\n#endif\n};\n\n#ifdef _WIN64\nclass WindowsSecurityAttributes\n{\nprotected:\n SECURITY_ATTRIBUTES m_winSecurityAttributes;\n PSECURITY_DESCRIPTOR m_winPSecurityDescriptor;\n\npublic:\n WindowsSecurityAttributes();\n SECURITY_ATTRIBUTES *operator&();\n ~WindowsSecurityAttributes();\n};\n\nWindowsSecurityAttributes::WindowsSecurityAttributes()\n{\n m_winPSecurityDescriptor = (PSECURITY_DESCRIPTOR)calloc(1, SECURITY_DESCRIPTOR_MIN_LENGTH + 2 * sizeof(void **));\n\n PSID *ppSID = (PSID *)((PBYTE)m_winPSecurityDescriptor + SECURITY_DESCRIPTOR_MIN_LENGTH);\n PACL *ppACL = (PACL *)((PBYTE)ppSID + sizeof(PSID *));\n\n InitializeSecurityDescriptor(m_winPSecurityDescriptor, SECURITY_DESCRIPTOR_REVISION);\n\n SID_IDENTIFIER_AUTHORITY sidIdentifierAuthority = SECURITY_WORLD_SID_AUTHORITY;\n AllocateAndInitializeSid(&sidIdentifierAuthority, 1, SECURITY_WORLD_RID, 0, 0, 0, 0, 0, 0, 0, ppSID);\n\n EXPLICIT_ACCESS explicitAccess;\n ZeroMemory(&explicitAccess, sizeof(EXPLICIT_ACCESS));\n explicitAccess.grfAccessPermissions = STANDARD_RIGHTS_ALL | SPECIFIC_RIGHTS_ALL;\n explicitAccess.grfAccessMode = SET_ACCESS;\n explicitAccess.grfInheritance = INHERIT_ONLY;\n explicitAccess.Trustee.TrusteeForm = TRUSTEE_IS_SID;\n explicitAccess.Trustee.TrusteeType = TRUSTEE_IS_WELL_KNOWN_GROUP;\n explicitAccess.Trustee.ptstrName = (LPTSTR)*ppSID;\n\n SetEntriesInAcl(1, &explicitAccess, NULL, ppACL);\n\n SetSecurityDescriptorDacl(m_winPSecurityDescriptor, TRUE, *ppACL, FALSE);\n\n m_winSecurityAttributes.nLength = sizeof(m_winSecurityAttributes);\n m_winSecurityAttributes.lpSecurityDescriptor = m_winPSecurityDescriptor;\n m_winSecurityAttributes.bInheritHandle = TRUE;\n}\n\nSECURITY_ATTRIBUTES *WindowsSecurityAttributes::operator&() { return &m_winSecurityAttributes; }\n\nWindowsSecurityAttributes::~WindowsSecurityAttributes()\n{\n PSID *ppSID = (PSID *)((PBYTE)m_winPSecurityDescriptor + SECURITY_DESCRIPTOR_MIN_LENGTH);\n PACL *ppACL = (PACL *)((PBYTE)ppSID + sizeof(PSID *));\n\n if (*ppSID) {\n FreeSid(*ppSID);\n }\n if (*ppACL) {\n LocalFree(*ppACL);\n }\n free(m_winPSecurityDescriptor);\n}\n#endif\n\nvoid DestroyDebugUtilsMessengerEXT(VkInstance instance,\n VkDebugUtilsMessengerEXT debugMessenger,\n const VkAllocationCallbacks *pAllocator)\n{\n auto func = (PFN_vkDestroyDebugUtilsMessengerEXT)vkGetInstanceProcAddr(instance, \"vkDestroyDebugUtilsMessengerEXT\");\n if (func != nullptr) {\n func(instance, debugMessenger, pAllocator);\n }\n}\n\nstruct QueueFamilyIndices\n{\n int graphicsFamily = -1;\n int presentFamily = -1;\n\n bool isComplete() { return graphicsFamily >= 0 && presentFamily >= 0; }\n};\n\nstruct SwapChainSupportDetails\n{\n VkSurfaceCapabilitiesKHR capabilities;\n std::vector formats;\n std::vector presentModes;\n};\n\ntypedef float vec2[2];\n\nstruct Vertex\n{\n vec4 pos;\n vec3 color;\n vec2 texCoord;\n\n static VkVertexInputBindingDescription getBindingDescription()\n {\n VkVertexInputBindingDescription bindingDescription = {};\n bindingDescription.binding = 0;\n bindingDescription.stride = sizeof(Vertex);\n bindingDescription.inputRate = VK_VERTEX_INPUT_RATE_VERTEX;\n\n return bindingDescription;\n }\n\n static std::array getAttributeDescriptions()\n {\n std::array attributeDescriptions = {};\n\n attributeDescriptions[0].binding = 0;\n attributeDescriptions[0].location = 0;\n attributeDescriptions[0].format = VK_FORMAT_R32G32B32A32_SFLOAT;\n attributeDescriptions[0].offset = offsetof(Vertex, pos);\n\n attributeDescriptions[1].binding = 0;\n attributeDescriptions[1].location = 1;\n attributeDescriptions[1].format = VK_FORMAT_R32G32B32_SFLOAT;\n attributeDescriptions[1].offset = offsetof(Vertex, color);\n\n attributeDescriptions[2].binding = 0;\n attributeDescriptions[2].location = 2;\n attributeDescriptions[2].format = VK_FORMAT_R32G32_SFLOAT;\n attributeDescriptions[2].offset = offsetof(Vertex, texCoord);\n\n return attributeDescriptions;\n }\n};\n\nstruct UniformBufferObject\n{\n alignas(16) mat4x4 model;\n alignas(16) mat4x4 view;\n alignas(16) mat4x4 proj;\n};\n\nconst std::vector vertices = {{{-1.0f, -1.0f, 0.0f, 1.0f}, {1.0f, 0.0f, 0.0f}, {0.0f, 0.0f}},\n {{1.0f, -1.0f, 0.0f, 1.0f}, {0.0f, 1.0f, 0.0f}, {1.0f, 0.0f}},\n {{1.0f, 1.0f, 0.0f, 1.0f}, {0.0f, 0.0f, 1.0f}, {1.0f, 1.0f}},\n {{-1.0f, 1.0f, 0.0f, 1.0f}, {1.0f, 1.0f, 1.0f}, {0.0f, 1.0f}}};\n\nconst std::vector indices = {0, 1, 2, 2, 3, 0};\n\n// convert floating point rgba color to 32-bit integer\n__device__ unsigned int rgbaFloatToInt(float4 rgba)\n{\n rgba.x = __saturatef(rgba.x); // clamp to [0.0, 1.0]\n rgba.y = __saturatef(rgba.y);\n rgba.z = __saturatef(rgba.z);\n rgba.w = __saturatef(rgba.w);\n return ((unsigned int)(rgba.w * 255.0f) << 24) | ((unsigned int)(rgba.z * 255.0f) << 16)\n | ((unsigned int)(rgba.y * 255.0f) << 8) | ((unsigned int)(rgba.x * 255.0f));\n}\n\n__device__ float4 rgbaIntToFloat(unsigned int c)\n{\n float4 rgba;\n rgba.x = (c & 0xff) * 0.003921568627f; // /255.0f;\n rgba.y = ((c >> 8) & 0xff) * 0.003921568627f; // /255.0f;\n rgba.z = ((c >> 16) & 0xff) * 0.003921568627f; // /255.0f;\n rgba.w = ((c >> 24) & 0xff) * 0.003921568627f; // /255.0f;\n return rgba;\n}\n\nint filter_radius = 14;\nint g_nFilterSign = 1;\n\n// This varies the filter radius, so we can see automatic animation\nvoid varySigma()\n{\n filter_radius += g_nFilterSign;\n\n if (filter_radius > 64) {\n filter_radius = 64; // clamp to 64 and then negate sign\n g_nFilterSign = -1;\n }\n else if (filter_radius < 0) {\n filter_radius = 0;\n g_nFilterSign = 1;\n }\n}\n\n// row pass using texture lookups\n__global__ void d_boxfilter_rgba_x(cudaSurfaceObject_t *dstSurfMipMapArray,\n cudaTextureObject_t textureMipMapInput,\n size_t baseWidth,\n size_t baseHeight,\n size_t mipLevels,\n int filter_radius)\n{\n float scale = 1.0f / (float)((filter_radius << 1) + 1);\n unsigned int y = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (y < baseHeight) {\n for (uint32_t mipLevelIdx = 0; mipLevelIdx < mipLevels; mipLevelIdx++) {\n uint32_t width = (baseWidth >> mipLevelIdx) ? (baseWidth >> mipLevelIdx) : 1;\n uint32_t height = (baseHeight >> mipLevelIdx) ? (baseHeight >> mipLevelIdx) : 1;\n if (y < height && filter_radius < width) {\n float px = 1.0 / width;\n float py = 1.0 / height;\n float4 t = make_float4(0.0f);\n for (int x = -filter_radius; x <= filter_radius; x++) {\n t += tex2DLod(textureMipMapInput, x * px, y * py, (float)mipLevelIdx);\n }\n\n unsigned int dataB = rgbaFloatToInt(t * scale);\n surf2Dwrite(dataB, dstSurfMipMapArray[mipLevelIdx], 0, y);\n\n for (int x = 1; x < width; x++) {\n t += tex2DLod(textureMipMapInput, (x + filter_radius) * px, y * py, (float)mipLevelIdx);\n t -= tex2DLod(textureMipMapInput, (x - filter_radius - 1) * px, y * py, (float)mipLevelIdx);\n unsigned int dataB = rgbaFloatToInt(t * scale);\n surf2Dwrite(dataB, dstSurfMipMapArray[mipLevelIdx], x * sizeof(uchar4), y);\n }\n }\n }\n }\n}\n\n// column pass using coalesced global memory reads\n__global__ void d_boxfilter_rgba_y(cudaSurfaceObject_t *dstSurfMipMapArray,\n cudaSurfaceObject_t *srcSurfMipMapArray,\n size_t baseWidth,\n size_t baseHeight,\n size_t mipLevels,\n int filter_radius)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n float scale = 1.0f / (float)((filter_radius << 1) + 1);\n\n for (uint32_t mipLevelIdx = 0; mipLevelIdx < mipLevels; mipLevelIdx++) {\n uint32_t width = (baseWidth >> mipLevelIdx) ? (baseWidth >> mipLevelIdx) : 1;\n uint32_t height = (baseHeight >> mipLevelIdx) ? (baseHeight >> mipLevelIdx) : 1;\n\n if (x < width && height > filter_radius) {\n float4 t;\n // do left edge\n int colInBytes = x * sizeof(uchar4);\n unsigned int pixFirst = surf2Dread(srcSurfMipMapArray[mipLevelIdx], colInBytes, 0);\n t = rgbaIntToFloat(pixFirst) * filter_radius;\n\n for (int y = 0; (y < (filter_radius + 1)) && (y < height); y++) {\n unsigned int pix = surf2Dread(srcSurfMipMapArray[mipLevelIdx], colInBytes, y);\n t += rgbaIntToFloat(pix);\n }\n\n unsigned int dataB = rgbaFloatToInt(t * scale);\n surf2Dwrite(dataB, dstSurfMipMapArray[mipLevelIdx], colInBytes, 0);\n\n for (int y = 1; (y < filter_radius + 1) && ((y + filter_radius) < height); y++) {\n unsigned int pix =\n surf2Dread(srcSurfMipMapArray[mipLevelIdx], colInBytes, y + filter_radius);\n t += rgbaIntToFloat(pix);\n t -= rgbaIntToFloat(pixFirst);\n\n dataB = rgbaFloatToInt(t * scale);\n surf2Dwrite(dataB, dstSurfMipMapArray[mipLevelIdx], colInBytes, y);\n }\n\n // main loop\n for (int y = (filter_radius + 1); y < (height - filter_radius); y++) {\n unsigned int pix =\n surf2Dread(srcSurfMipMapArray[mipLevelIdx], colInBytes, y + filter_radius);\n t += rgbaIntToFloat(pix);\n\n pix = surf2Dread(srcSurfMipMapArray[mipLevelIdx], colInBytes, y - filter_radius - 1);\n t -= rgbaIntToFloat(pix);\n\n dataB = rgbaFloatToInt(t * scale);\n surf2Dwrite(dataB, dstSurfMipMapArray[mipLevelIdx], colInBytes, y);\n }\n\n // do right edge\n unsigned int pixLast = surf2Dread(srcSurfMipMapArray[mipLevelIdx], colInBytes, height - 1);\n for (int y = height - filter_radius; (y < height) && ((y - filter_radius - 1) > 1); y++) {\n t += rgbaIntToFloat(pixLast);\n unsigned int pix =\n surf2Dread(srcSurfMipMapArray[mipLevelIdx], colInBytes, y - filter_radius - 1);\n t -= rgbaIntToFloat(pix);\n dataB = rgbaFloatToInt(t * scale);\n surf2Dwrite(dataB, dstSurfMipMapArray[mipLevelIdx], colInBytes, y);\n }\n }\n }\n}\n\nclass vulkanImageCUDA\n{\npublic:\n void loadImageData(const std::string &filename)\n {\n // load image (needed so we can get the width and height before we create\n // the window\n char *image_path = sdkFindFilePath(filename.c_str(), execution_path.c_str());\n\n if (image_path == 0) {\n printf(\"Error finding image file '%s'\\n\", filename.c_str());\n exit(EXIT_FAILURE);\n }\n\n sdkLoadPPM4(image_path, (unsigned char **)&image_data, &imageWidth, &imageHeight);\n\n if (!image_data) {\n printf(\"Error opening file '%s'\\n\", image_path);\n exit(EXIT_FAILURE);\n }\n\n printf(\"Loaded '%s', %d x %d pixels\\n\", image_path, imageWidth, imageHeight);\n }\n\n void run()\n {\n initWindow();\n initVulkan();\n initCuda();\n mainLoop();\n cleanup();\n }\n\nprivate:\n GLFWwindow *window;\n\n VkInstance instance;\n VkDebugUtilsMessengerEXT debugMessenger;\n VkSurfaceKHR surface;\n\n VkPhysicalDevice physicalDevice = VK_NULL_HANDLE;\n VkDevice device;\n uint8_t vkDeviceUUID[VK_UUID_SIZE];\n\n VkQueue graphicsQueue;\n VkQueue presentQueue;\n\n VkSwapchainKHR swapChain;\n std::vector swapChainImages;\n VkFormat swapChainImageFormat;\n VkExtent2D swapChainExtent;\n std::vector swapChainImageViews;\n std::vector swapChainFramebuffers;\n\n VkRenderPass renderPass;\n VkDescriptorSetLayout descriptorSetLayout;\n VkPipelineLayout pipelineLayout;\n VkPipeline graphicsPipeline;\n\n VkCommandPool commandPool;\n\n VkImage textureImage;\n VkDeviceMemory textureImageMemory;\n VkImageView textureImageView;\n VkSampler textureSampler;\n\n VkBuffer vertexBuffer;\n VkDeviceMemory vertexBufferMemory;\n VkBuffer indexBuffer;\n VkDeviceMemory indexBufferMemory;\n\n std::vector uniformBuffers;\n std::vector uniformBuffersMemory;\n\n VkDescriptorPool descriptorPool;\n std::vector descriptorSets;\n\n std::vector commandBuffers;\n\n std::vector imageAvailableSemaphores;\n std::vector renderFinishedSemaphores;\n VkSemaphore cudaUpdateVkSemaphore, vkUpdateCudaSemaphore;\n std::vector inFlightFences;\n\n size_t currentFrame = 0;\n\n bool framebufferResized = false;\n\n#ifdef _WIN64\n PFN_vkGetMemoryWin32HandleKHR fpGetMemoryWin32HandleKHR;\n PFN_vkGetSemaphoreWin32HandleKHR fpGetSemaphoreWin32HandleKHR;\n#else\n PFN_vkGetMemoryFdKHR fpGetMemoryFdKHR = NULL;\n PFN_vkGetSemaphoreFdKHR fpGetSemaphoreFdKHR = NULL;\n#endif\n\n PFN_vkGetPhysicalDeviceProperties2 fpGetPhysicalDeviceProperties2;\n\n unsigned int *image_data = NULL;\n unsigned int imageWidth, imageHeight;\n unsigned int mipLevels = 1;\n size_t totalImageMemSize;\n\n // CUDA objects\n cudaExternalMemory_t cudaExtMemImageBuffer;\n cudaMipmappedArray_t cudaMipmappedImageArray, cudaMipmappedImageArrayTemp, cudaMipmappedImageArrayOrig;\n std::vector surfaceObjectList, surfaceObjectListTemp;\n cudaSurfaceObject_t *d_surfaceObjectList, *d_surfaceObjectListTemp;\n cudaTextureObject_t textureObjMipMapInput;\n\n cudaExternalSemaphore_t cudaExtCudaUpdateVkSemaphore;\n cudaExternalSemaphore_t cudaExtVkUpdateCudaSemaphore;\n cudaStream_t streamToRun;\n\n void initWindow()\n {\n glfwInit();\n\n glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);\n\n window = glfwCreateWindow(WIDTH, HEIGHT, \"Vulkan Image CUDA Box Filter\", nullptr, nullptr);\n glfwSetWindowUserPointer(window, this);\n glfwSetFramebufferSizeCallback(window, framebufferResizeCallback);\n }\n\n static void framebufferResizeCallback(GLFWwindow *window, int width, int height)\n {\n auto app = reinterpret_cast(glfwGetWindowUserPointer(window));\n app->framebufferResized = true;\n }\n\n void initVulkan()\n {\n createInstance();\n setupDebugMessenger();\n createSurface();\n pickPhysicalDevice();\n createLogicalDevice();\n getKhrExtensionsFn();\n createSwapChain();\n createImageViews();\n createRenderPass();\n createDescriptorSetLayout();\n createGraphicsPipeline();\n createFramebuffers();\n createCommandPool();\n createTextureImage();\n createTextureImageView();\n createTextureSampler();\n createVertexBuffer();\n createIndexBuffer();\n createUniformBuffers();\n createDescriptorPool();\n createDescriptorSets();\n createCommandBuffers();\n createSyncObjects();\n createSyncObjectsExt();\n }\n\n void initCuda()\n {\n setCudaVkDevice();\n checkCudaErrors(cudaStreamCreate(&streamToRun));\n cudaVkImportImageMem();\n cudaVkImportSemaphore();\n }\n\n void mainLoop()\n {\n updateUniformBuffer();\n while (!glfwWindowShouldClose(window)) {\n glfwPollEvents();\n drawFrame();\n }\n\n vkDeviceWaitIdle(device);\n }\n\n void cleanupSwapChain()\n {\n for (auto framebuffer : swapChainFramebuffers) {\n vkDestroyFramebuffer(device, framebuffer, nullptr);\n }\n\n vkFreeCommandBuffers(device, commandPool, static_cast(commandBuffers.size()), commandBuffers.data());\n\n vkDestroyPipeline(device, graphicsPipeline, nullptr);\n vkDestroyPipelineLayout(device, pipelineLayout, nullptr);\n vkDestroyRenderPass(device, renderPass, nullptr);\n\n for (auto imageView : swapChainImageViews) {\n vkDestroyImageView(device, imageView, nullptr);\n }\n\n vkDestroySwapchainKHR(device, swapChain, nullptr);\n\n for (size_t i = 0; i < swapChainImages.size(); i++) {\n vkDestroyBuffer(device, uniformBuffers[i], nullptr);\n vkFreeMemory(device, uniformBuffersMemory[i], nullptr);\n }\n\n vkDestroyDescriptorPool(device, descriptorPool, nullptr);\n }\n\n void cleanup()\n {\n cleanupSwapChain();\n\n vkDestroySampler(device, textureSampler, nullptr);\n vkDestroyImageView(device, textureImageView, nullptr);\n\n for (int i = 0; i < mipLevels; i++) {\n checkCudaErrors(cudaDestroySurfaceObject(surfaceObjectList[i]));\n checkCudaErrors(cudaDestroySurfaceObject(surfaceObjectListTemp[i]));\n }\n\n checkCudaErrors(cudaFree(d_surfaceObjectList));\n checkCudaErrors(cudaFree(d_surfaceObjectListTemp));\n checkCudaErrors(cudaFreeMipmappedArray(cudaMipmappedImageArrayTemp));\n checkCudaErrors(cudaFreeMipmappedArray(cudaMipmappedImageArrayOrig));\n checkCudaErrors(cudaFreeMipmappedArray(cudaMipmappedImageArray));\n checkCudaErrors(cudaDestroyTextureObject(textureObjMipMapInput));\n checkCudaErrors(cudaDestroyExternalMemory(cudaExtMemImageBuffer));\n checkCudaErrors(cudaDestroyExternalSemaphore(cudaExtCudaUpdateVkSemaphore));\n checkCudaErrors(cudaDestroyExternalSemaphore(cudaExtVkUpdateCudaSemaphore));\n\n vkDestroyImage(device, textureImage, nullptr);\n vkFreeMemory(device, textureImageMemory, nullptr);\n\n vkDestroyDescriptorSetLayout(device, descriptorSetLayout, nullptr);\n\n vkDestroyBuffer(device, indexBuffer, nullptr);\n vkFreeMemory(device, indexBufferMemory, nullptr);\n\n vkDestroyBuffer(device, vertexBuffer, nullptr);\n vkFreeMemory(device, vertexBufferMemory, nullptr);\n\n vkDestroySemaphore(device, cudaUpdateVkSemaphore, nullptr);\n vkDestroySemaphore(device, vkUpdateCudaSemaphore, nullptr);\n\n for (size_t i = 0; i < MAX_FRAMES; i++) {\n vkDestroySemaphore(device, renderFinishedSemaphores[i], nullptr);\n vkDestroySemaphore(device, imageAvailableSemaphores[i], nullptr);\n vkDestroyFence(device, inFlightFences[i], nullptr);\n }\n\n vkDestroyCommandPool(device, commandPool, nullptr);\n\n vkDestroyDevice(device, nullptr);\n\n if (enableValidationLayers) {\n DestroyDebugUtilsMessengerEXT(instance, debugMessenger, nullptr);\n }\n\n vkDestroySurfaceKHR(instance, surface, nullptr);\n vkDestroyInstance(instance, nullptr);\n\n glfwDestroyWindow(window);\n\n glfwTerminate();\n }\n\n void recreateSwapChain()\n {\n int width = 0, height = 0;\n while (width == 0 || height == 0) {\n glfwGetFramebufferSize(window, &width, &height);\n glfwWaitEvents();\n }\n\n vkDeviceWaitIdle(device);\n\n cleanupSwapChain();\n\n createSwapChain();\n createImageViews();\n createRenderPass();\n createGraphicsPipeline();\n createFramebuffers();\n createUniformBuffers();\n createDescriptorPool();\n createDescriptorSets();\n createCommandBuffers();\n }\n\n void createInstance()\n {\n if (enableValidationLayers && !checkValidationLayerSupport()) {\n throw std::runtime_error(\"validation layers requested, but not available!\");\n }\n\n VkApplicationInfo appInfo = {};\n appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;\n appInfo.pApplicationName = \"Vulkan Image CUDA Interop\";\n appInfo.applicationVersion = VK_MAKE_VERSION(1, 0, 0);\n appInfo.pEngineName = \"No Engine\";\n appInfo.engineVersion = VK_MAKE_VERSION(1, 0, 0);\n appInfo.apiVersion = VK_API_VERSION_1_1;\n\n VkInstanceCreateInfo createInfo = {};\n createInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;\n createInfo.pApplicationInfo = &appInfo;\n\n auto extensions = getRequiredExtensions();\n createInfo.enabledExtensionCount = static_cast(extensions.size());\n createInfo.ppEnabledExtensionNames = extensions.data();\n\n VkDebugUtilsMessengerCreateInfoEXT debugCreateInfo;\n if (enableValidationLayers) {\n createInfo.enabledLayerCount = static_cast(validationLayers.size());\n createInfo.ppEnabledLayerNames = validationLayers.data();\n\n populateDebugMessengerCreateInfo(debugCreateInfo);\n createInfo.pNext = (VkDebugUtilsMessengerCreateInfoEXT *)&debugCreateInfo;\n }\n else {\n createInfo.enabledLayerCount = 0;\n\n createInfo.pNext = nullptr;\n }\n\n if (vkCreateInstance(&createInfo, nullptr, &instance) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create instance!\");\n }\n\n fpGetPhysicalDeviceProperties2 =\n (PFN_vkGetPhysicalDeviceProperties2)vkGetInstanceProcAddr(instance, \"vkGetPhysicalDeviceProperties2\");\n if (fpGetPhysicalDeviceProperties2 == NULL) {\n throw std::runtime_error(\"Vulkan: Proc address for \\\"vkGetPhysicalDeviceProperties2KHR\\\" not \"\n \"found.\\n\");\n }\n\n#ifdef _WIN64\n fpGetMemoryWin32HandleKHR =\n (PFN_vkGetMemoryWin32HandleKHR)vkGetInstanceProcAddr(instance, \"vkGetMemoryWin32HandleKHR\");\n if (fpGetMemoryWin32HandleKHR == NULL) {\n throw std::runtime_error(\"Vulkan: Proc address for \\\"vkGetMemoryWin32HandleKHR\\\" not \"\n \"found.\\n\");\n }\n#else\n fpGetMemoryFdKHR = (PFN_vkGetMemoryFdKHR)vkGetInstanceProcAddr(instance, \"vkGetMemoryFdKHR\");\n if (fpGetMemoryFdKHR == NULL) {\n throw std::runtime_error(\"Vulkan: Proc address for \\\"vkGetMemoryFdKHR\\\" not found.\\n\");\n }\n else {\n std::cout << \"Vulkan proc address for vkGetMemoryFdKHR - \" << fpGetMemoryFdKHR << std::endl;\n }\n#endif\n }\n\n void populateDebugMessengerCreateInfo(VkDebugUtilsMessengerCreateInfoEXT &createInfo)\n {\n createInfo = {};\n createInfo.sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_MESSENGER_CREATE_INFO_EXT;\n createInfo.messageSeverity = VK_DEBUG_UTILS_MESSAGE_SEVERITY_VERBOSE_BIT_EXT\n | VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT\n | VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT;\n createInfo.messageType = VK_DEBUG_UTILS_MESSAGE_TYPE_GENERAL_BIT_EXT\n | VK_DEBUG_UTILS_MESSAGE_TYPE_VALIDATION_BIT_EXT\n | VK_DEBUG_UTILS_MESSAGE_TYPE_PERFORMANCE_BIT_EXT;\n createInfo.pfnUserCallback = debugCallback;\n }\n\n void setupDebugMessenger()\n {\n if (!enableValidationLayers)\n return;\n\n VkDebugUtilsMessengerCreateInfoEXT createInfo;\n populateDebugMessengerCreateInfo(createInfo);\n\n if (CreateDebugUtilsMessengerEXT(instance, &createInfo, nullptr, &debugMessenger) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to set up debug messenger!\");\n }\n }\n\n void createSurface()\n {\n if (glfwCreateWindowSurface(instance, window, nullptr, &surface) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create window surface!\");\n }\n }\n\n void pickPhysicalDevice()\n {\n uint32_t deviceCount = 0;\n vkEnumeratePhysicalDevices(instance, &deviceCount, nullptr);\n\n if (deviceCount == 0) {\n throw std::runtime_error(\"failed to find GPUs with Vulkan support!\");\n }\n\n std::vector devices(deviceCount);\n vkEnumeratePhysicalDevices(instance, &deviceCount, devices.data());\n\n for (const auto &device : devices) {\n if (isDeviceSuitable(device)) {\n physicalDevice = device;\n break;\n }\n }\n\n if (physicalDevice == VK_NULL_HANDLE) {\n throw std::runtime_error(\"failed to find a suitable GPU!\");\n }\n\n std::cout << \"Selected physical device = \" << physicalDevice << std::endl;\n\n VkPhysicalDeviceIDProperties vkPhysicalDeviceIDProperties = {};\n vkPhysicalDeviceIDProperties.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ID_PROPERTIES;\n vkPhysicalDeviceIDProperties.pNext = NULL;\n\n VkPhysicalDeviceProperties2 vkPhysicalDeviceProperties2 = {};\n vkPhysicalDeviceProperties2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2;\n vkPhysicalDeviceProperties2.pNext = &vkPhysicalDeviceIDProperties;\n\n fpGetPhysicalDeviceProperties2(physicalDevice, &vkPhysicalDeviceProperties2);\n\n memcpy(vkDeviceUUID, vkPhysicalDeviceIDProperties.deviceUUID, sizeof(vkDeviceUUID));\n }\n\n void getKhrExtensionsFn()\n {\n#ifdef _WIN64\n\n fpGetSemaphoreWin32HandleKHR =\n (PFN_vkGetSemaphoreWin32HandleKHR)vkGetDeviceProcAddr(device, \"vkGetSemaphoreWin32HandleKHR\");\n if (fpGetSemaphoreWin32HandleKHR == NULL) {\n throw std::runtime_error(\"Vulkan: Proc address for \\\"vkGetSemaphoreWin32HandleKHR\\\" not \"\n \"found.\\n\");\n }\n#else\n fpGetSemaphoreFdKHR = (PFN_vkGetSemaphoreFdKHR)vkGetDeviceProcAddr(device, \"vkGetSemaphoreFdKHR\");\n if (fpGetSemaphoreFdKHR == NULL) {\n throw std::runtime_error(\"Vulkan: Proc address for \\\"vkGetSemaphoreFdKHR\\\" not found.\\n\");\n }\n#endif\n }\n\n int setCudaVkDevice()\n {\n int current_device = 0;\n int device_count = 0;\n int devices_prohibited = 0;\n\n cudaDeviceProp deviceProp;\n int computeMode;\n checkCudaErrors(cudaGetDeviceCount(&device_count));\n\n if (device_count == 0) {\n fprintf(stderr, \"CUDA error: no devices supporting CUDA.\\n\");\n exit(EXIT_FAILURE);\n }\n\n // Find the GPU which is selected by Vulkan\n while (current_device < device_count) {\n cudaGetDeviceProperties(&deviceProp, current_device);\n checkCudaErrors(cudaDeviceGetAttribute(&computeMode, cudaDevAttrComputeMode, current_device));\n if ((computeMode != cudaComputeModeProhibited)) {\n // Compare the cuda device UUID with vulkan UUID\n int ret = memcmp(&deviceProp.uuid, &vkDeviceUUID, VK_UUID_SIZE);\n if (ret == 0) {\n checkCudaErrors(cudaSetDevice(current_device));\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, current_device));\n printf(\"GPU Device %d: \\\"%s\\\" with compute capability %d.%d\\n\\n\",\n current_device,\n deviceProp.name,\n deviceProp.major,\n deviceProp.minor);\n\n return current_device;\n }\n }\n else {\n devices_prohibited++;\n }\n\n current_device++;\n }\n\n if (devices_prohibited == device_count) {\n fprintf(stderr,\n \"CUDA error:\"\n \" No Vulkan-CUDA Interop capable GPU found.\\n\");\n exit(EXIT_FAILURE);\n }\n\n return -1;\n }\n\n void createLogicalDevice()\n {\n QueueFamilyIndices indices = findQueueFamilies(physicalDevice);\n\n std::vector queueCreateInfos;\n std::set uniqueQueueFamilies = {indices.graphicsFamily, indices.presentFamily};\n\n float queuePriority = 1.0f;\n for (int queueFamily : uniqueQueueFamilies) {\n VkDeviceQueueCreateInfo queueCreateInfo = {};\n queueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;\n queueCreateInfo.queueFamilyIndex = queueFamily;\n queueCreateInfo.queueCount = 1;\n queueCreateInfo.pQueuePriorities = &queuePriority;\n queueCreateInfos.push_back(queueCreateInfo);\n }\n\n VkPhysicalDeviceFeatures deviceFeatures = {};\n deviceFeatures.samplerAnisotropy = VK_TRUE;\n\n VkDeviceCreateInfo createInfo = {};\n createInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;\n\n createInfo.pQueueCreateInfos = queueCreateInfos.data();\n createInfo.queueCreateInfoCount = queueCreateInfos.size();\n\n createInfo.pEnabledFeatures = &deviceFeatures;\n std::vector enabledExtensionNameList;\n\n for (int i = 0; i < deviceExtensions.size(); i++) {\n enabledExtensionNameList.push_back(deviceExtensions[i]);\n }\n if (enableValidationLayers) {\n createInfo.enabledLayerCount = static_cast(validationLayers.size());\n createInfo.ppEnabledLayerNames = validationLayers.data();\n }\n else {\n createInfo.enabledLayerCount = 0;\n }\n createInfo.enabledExtensionCount = static_cast(enabledExtensionNameList.size());\n createInfo.ppEnabledExtensionNames = enabledExtensionNameList.data();\n\n if (vkCreateDevice(physicalDevice, &createInfo, nullptr, &device) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create logical device!\");\n }\n vkGetDeviceQueue(device, indices.graphicsFamily, 0, &graphicsQueue);\n vkGetDeviceQueue(device, indices.presentFamily, 0, &presentQueue);\n }\n\n void createSwapChain()\n {\n SwapChainSupportDetails swapChainSupport = querySwapChainSupport(physicalDevice);\n\n VkSurfaceFormatKHR surfaceFormat = chooseSwapSurfaceFormat(swapChainSupport.formats);\n VkPresentModeKHR presentMode = chooseSwapPresentMode(swapChainSupport.presentModes);\n VkExtent2D extent = chooseSwapExtent(swapChainSupport.capabilities);\n\n uint32_t imageCount = swapChainSupport.capabilities.minImageCount + 1;\n if (swapChainSupport.capabilities.maxImageCount > 0\n && imageCount > swapChainSupport.capabilities.maxImageCount) {\n imageCount = swapChainSupport.capabilities.maxImageCount;\n }\n\n VkSwapchainCreateInfoKHR createInfo = {};\n createInfo.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;\n createInfo.surface = surface;\n\n createInfo.minImageCount = imageCount;\n createInfo.imageFormat = surfaceFormat.format;\n createInfo.imageColorSpace = surfaceFormat.colorSpace;\n createInfo.imageExtent = extent;\n createInfo.imageArrayLayers = 1;\n createInfo.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;\n\n QueueFamilyIndices indices = findQueueFamilies(physicalDevice);\n uint32_t queueFamilyIndices[] = {(uint32_t)indices.graphicsFamily, (uint32_t)indices.presentFamily};\n\n if (indices.graphicsFamily != indices.presentFamily) {\n createInfo.imageSharingMode = VK_SHARING_MODE_CONCURRENT;\n createInfo.queueFamilyIndexCount = 2;\n createInfo.pQueueFamilyIndices = queueFamilyIndices;\n }\n else {\n createInfo.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;\n }\n\n createInfo.preTransform = swapChainSupport.capabilities.currentTransform;\n createInfo.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;\n createInfo.presentMode = presentMode;\n createInfo.clipped = VK_TRUE;\n\n if (vkCreateSwapchainKHR(device, &createInfo, nullptr, &swapChain) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create swap chain!\");\n }\n\n vkGetSwapchainImagesKHR(device, swapChain, &imageCount, nullptr);\n swapChainImages.resize(imageCount);\n vkGetSwapchainImagesKHR(device, swapChain, &imageCount, swapChainImages.data());\n\n swapChainImageFormat = surfaceFormat.format;\n swapChainExtent = extent;\n }\n\n void createImageViews()\n {\n swapChainImageViews.resize(swapChainImages.size());\n\n for (size_t i = 0; i < swapChainImages.size(); i++) {\n swapChainImageViews[i] = createImageView(swapChainImages[i], swapChainImageFormat);\n }\n }\n\n void createRenderPass()\n {\n VkAttachmentDescription colorAttachment = {};\n colorAttachment.format = swapChainImageFormat;\n colorAttachment.samples = VK_SAMPLE_COUNT_1_BIT;\n colorAttachment.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;\n colorAttachment.storeOp = VK_ATTACHMENT_STORE_OP_STORE;\n colorAttachment.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;\n colorAttachment.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;\n colorAttachment.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;\n colorAttachment.finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;\n\n VkAttachmentReference colorAttachmentRef = {};\n colorAttachmentRef.attachment = 0;\n colorAttachmentRef.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;\n\n VkSubpassDescription subpass = {};\n subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;\n subpass.colorAttachmentCount = 1;\n subpass.pColorAttachments = &colorAttachmentRef;\n\n VkSubpassDependency dependency = {};\n dependency.srcSubpass = VK_SUBPASS_EXTERNAL;\n dependency.dstSubpass = 0;\n dependency.srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;\n dependency.srcAccessMask = 0;\n dependency.dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;\n dependency.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_READ_BIT | VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;\n\n VkRenderPassCreateInfo renderPassInfo = {};\n renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;\n renderPassInfo.attachmentCount = 1;\n renderPassInfo.pAttachments = &colorAttachment;\n renderPassInfo.subpassCount = 1;\n renderPassInfo.pSubpasses = &subpass;\n renderPassInfo.dependencyCount = 1;\n renderPassInfo.pDependencies = &dependency;\n\n if (vkCreateRenderPass(device, &renderPassInfo, nullptr, &renderPass) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create render pass!\");\n }\n }\n\n void createDescriptorSetLayout()\n {\n VkDescriptorSetLayoutBinding uboLayoutBinding = {};\n uboLayoutBinding.binding = 0;\n uboLayoutBinding.descriptorCount = 1;\n uboLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;\n uboLayoutBinding.pImmutableSamplers = nullptr;\n uboLayoutBinding.stageFlags = VK_SHADER_STAGE_VERTEX_BIT;\n\n VkDescriptorSetLayoutBinding samplerLayoutBinding = {};\n samplerLayoutBinding.binding = 1;\n samplerLayoutBinding.descriptorCount = 1;\n samplerLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;\n samplerLayoutBinding.pImmutableSamplers = nullptr;\n samplerLayoutBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;\n\n std::array bindings = {uboLayoutBinding, samplerLayoutBinding};\n VkDescriptorSetLayoutCreateInfo layoutInfo = {};\n layoutInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;\n layoutInfo.bindingCount = static_cast(bindings.size());\n layoutInfo.pBindings = bindings.data();\n\n if (vkCreateDescriptorSetLayout(device, &layoutInfo, nullptr, &descriptorSetLayout) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create descriptor set layout!\");\n }\n }\n\n void createGraphicsPipeline()\n {\n auto vertShaderCode = readFile(\"vert.spv\");\n auto fragShaderCode = readFile(\"frag.spv\");\n\n VkShaderModule vertShaderModule = createShaderModule(vertShaderCode);\n VkShaderModule fragShaderModule = createShaderModule(fragShaderCode);\n\n VkPipelineShaderStageCreateInfo vertShaderStageInfo = {};\n vertShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;\n vertShaderStageInfo.stage = VK_SHADER_STAGE_VERTEX_BIT;\n vertShaderStageInfo.module = vertShaderModule;\n vertShaderStageInfo.pName = \"main\";\n\n VkPipelineShaderStageCreateInfo fragShaderStageInfo = {};\n fragShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;\n fragShaderStageInfo.stage = VK_SHADER_STAGE_FRAGMENT_BIT;\n fragShaderStageInfo.module = fragShaderModule;\n fragShaderStageInfo.pName = \"main\";\n\n VkPipelineShaderStageCreateInfo shaderStages[] = {vertShaderStageInfo, fragShaderStageInfo};\n\n VkPipelineVertexInputStateCreateInfo vertexInputInfo = {};\n vertexInputInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;\n\n auto bindingDescription = Vertex::getBindingDescription();\n auto attributeDescriptions = Vertex::getAttributeDescriptions();\n\n vertexInputInfo.vertexBindingDescriptionCount = 1;\n vertexInputInfo.vertexAttributeDescriptionCount = static_cast(attributeDescriptions.size());\n vertexInputInfo.pVertexBindingDescriptions = &bindingDescription;\n vertexInputInfo.pVertexAttributeDescriptions = attributeDescriptions.data();\n\n VkPipelineInputAssemblyStateCreateInfo inputAssembly = {};\n inputAssembly.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;\n inputAssembly.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;\n inputAssembly.primitiveRestartEnable = VK_FALSE;\n\n VkViewport viewport = {};\n viewport.x = 0.0f;\n viewport.y = 0.0f;\n viewport.width = (float)swapChainExtent.width;\n viewport.height = (float)swapChainExtent.height;\n viewport.minDepth = 0.0f;\n viewport.maxDepth = 1.0f;\n\n VkRect2D scissor = {};\n scissor.offset = {0, 0};\n scissor.extent = swapChainExtent;\n\n VkPipelineViewportStateCreateInfo viewportState = {};\n viewportState.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;\n viewportState.viewportCount = 1;\n viewportState.pViewports = &viewport;\n viewportState.scissorCount = 1;\n viewportState.pScissors = &scissor;\n\n VkPipelineRasterizationStateCreateInfo rasterizer = {};\n rasterizer.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;\n rasterizer.depthClampEnable = VK_FALSE;\n rasterizer.rasterizerDiscardEnable = VK_FALSE;\n rasterizer.polygonMode = VK_POLYGON_MODE_FILL;\n rasterizer.lineWidth = 1.0f;\n rasterizer.cullMode = VK_CULL_MODE_BACK_BIT;\n rasterizer.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;\n rasterizer.depthBiasEnable = VK_FALSE;\n\n VkPipelineMultisampleStateCreateInfo multisampling = {};\n multisampling.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;\n multisampling.sampleShadingEnable = VK_FALSE;\n multisampling.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;\n\n VkPipelineColorBlendAttachmentState colorBlendAttachment = {};\n colorBlendAttachment.colorWriteMask =\n VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;\n colorBlendAttachment.blendEnable = VK_FALSE;\n\n VkPipelineColorBlendStateCreateInfo colorBlending = {};\n colorBlending.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;\n colorBlending.logicOpEnable = VK_FALSE;\n colorBlending.logicOp = VK_LOGIC_OP_COPY;\n colorBlending.attachmentCount = 1;\n colorBlending.pAttachments = &colorBlendAttachment;\n colorBlending.blendConstants[0] = 0.0f;\n colorBlending.blendConstants[1] = 0.0f;\n colorBlending.blendConstants[2] = 0.0f;\n colorBlending.blendConstants[3] = 0.0f;\n\n VkPipelineLayoutCreateInfo pipelineLayoutInfo = {};\n pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;\n pipelineLayoutInfo.setLayoutCount = 1;\n pipelineLayoutInfo.pSetLayouts = &descriptorSetLayout;\n\n if (vkCreatePipelineLayout(device, &pipelineLayoutInfo, nullptr, &pipelineLayout) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create pipeline layout!\");\n }\n\n VkGraphicsPipelineCreateInfo pipelineInfo = {};\n pipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;\n pipelineInfo.stageCount = 2;\n pipelineInfo.pStages = shaderStages;\n pipelineInfo.pVertexInputState = &vertexInputInfo;\n pipelineInfo.pInputAssemblyState = &inputAssembly;\n pipelineInfo.pViewportState = &viewportState;\n pipelineInfo.pRasterizationState = &rasterizer;\n pipelineInfo.pMultisampleState = &multisampling;\n pipelineInfo.pColorBlendState = &colorBlending;\n pipelineInfo.layout = pipelineLayout;\n pipelineInfo.renderPass = renderPass;\n pipelineInfo.subpass = 0;\n pipelineInfo.basePipelineHandle = VK_NULL_HANDLE;\n\n if (vkCreateGraphicsPipelines(device, VK_NULL_HANDLE, 1, &pipelineInfo, nullptr, &graphicsPipeline)\n != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create graphics pipeline!\");\n }\n\n vkDestroyShaderModule(device, fragShaderModule, nullptr);\n vkDestroyShaderModule(device, vertShaderModule, nullptr);\n }\n\n void createFramebuffers()\n {\n swapChainFramebuffers.resize(swapChainImageViews.size());\n\n for (size_t i = 0; i < swapChainImageViews.size(); i++) {\n VkImageView attachments[] = {swapChainImageViews[i]};\n\n VkFramebufferCreateInfo framebufferInfo = {};\n framebufferInfo.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;\n framebufferInfo.renderPass = renderPass;\n framebufferInfo.attachmentCount = 1;\n framebufferInfo.pAttachments = attachments;\n framebufferInfo.width = swapChainExtent.width;\n framebufferInfo.height = swapChainExtent.height;\n framebufferInfo.layers = 1;\n\n if (vkCreateFramebuffer(device, &framebufferInfo, nullptr, &swapChainFramebuffers[i]) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create framebuffer!\");\n }\n }\n }\n\n void createCommandPool()\n {\n QueueFamilyIndices queueFamilyIndices = findQueueFamilies(physicalDevice);\n\n VkCommandPoolCreateInfo poolInfo = {};\n poolInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;\n poolInfo.queueFamilyIndex = queueFamilyIndices.graphicsFamily;\n\n if (vkCreateCommandPool(device, &poolInfo, nullptr, &commandPool) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create graphics command pool!\");\n }\n }\n\n void createTextureImage()\n {\n VkDeviceSize imageSize = imageWidth * imageHeight * 4;\n mipLevels = static_cast(std::floor(std::log2(std::max(imageWidth, imageHeight)))) + 1;\n printf(\"mipLevels = %d\\n\", mipLevels);\n\n if (!image_data) {\n throw std::runtime_error(\"failed to load texture image!\");\n }\n\n VkBuffer stagingBuffer;\n VkDeviceMemory stagingBufferMemory;\n createBuffer(imageSize,\n VK_BUFFER_USAGE_TRANSFER_SRC_BIT,\n VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,\n stagingBuffer,\n stagingBufferMemory);\n\n void *data;\n vkMapMemory(device, stagingBufferMemory, 0, imageSize, 0, &data);\n memcpy(data, image_data, static_cast(imageSize));\n vkUnmapMemory(device, stagingBufferMemory);\n\n // VK_FORMAT_R8G8B8A8_UNORM changed to VK_FORMAT_R8G8B8A8_UINT\n createImage(imageWidth,\n imageHeight,\n VK_FORMAT_R8G8B8A8_UNORM,\n VK_IMAGE_TILING_OPTIMAL,\n VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT\n | VK_IMAGE_USAGE_SAMPLED_BIT,\n VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,\n textureImage,\n textureImageMemory);\n\n transitionImageLayout(\n textureImage, VK_FORMAT_R8G8B8A8_UINT, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);\n copyBufferToImage(\n stagingBuffer, textureImage, static_cast(imageWidth), static_cast(imageHeight));\n\n vkDestroyBuffer(device, stagingBuffer, nullptr);\n vkFreeMemory(device, stagingBufferMemory, nullptr);\n\n generateMipmaps(textureImage, VK_FORMAT_R8G8B8A8_UNORM);\n }\n\n void generateMipmaps(VkImage image, VkFormat imageFormat)\n {\n VkFormatProperties formatProperties;\n vkGetPhysicalDeviceFormatProperties(physicalDevice, imageFormat, &formatProperties);\n\n if (!(formatProperties.optimalTilingFeatures & VK_FORMAT_FEATURE_SAMPLED_IMAGE_FILTER_LINEAR_BIT)) {\n throw std::runtime_error(\"texture image format does not support linear blitting!\");\n }\n\n VkCommandBuffer commandBuffer = beginSingleTimeCommands();\n\n VkImageMemoryBarrier barrier = {};\n barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;\n barrier.image = image;\n barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;\n barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;\n barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;\n barrier.subresourceRange.baseArrayLayer = 0;\n barrier.subresourceRange.layerCount = 1;\n barrier.subresourceRange.levelCount = 1;\n\n int32_t mipWidth = imageWidth;\n int32_t mipHeight = imageHeight;\n\n for (uint32_t i = 1; i < mipLevels; i++) {\n barrier.subresourceRange.baseMipLevel = i - 1;\n barrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;\n barrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;\n barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;\n barrier.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT;\n\n vkCmdPipelineBarrier(commandBuffer,\n VK_PIPELINE_STAGE_TRANSFER_BIT,\n VK_PIPELINE_STAGE_TRANSFER_BIT,\n 0,\n 0,\n nullptr,\n 0,\n nullptr,\n 1,\n &barrier);\n\n VkImageBlit blit = {};\n blit.srcOffsets[0] = {0, 0, 0};\n blit.srcOffsets[1] = {mipWidth, mipHeight, 1};\n blit.srcSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;\n blit.srcSubresource.mipLevel = i - 1;\n blit.srcSubresource.baseArrayLayer = 0;\n blit.srcSubresource.layerCount = 1;\n blit.dstOffsets[0] = {0, 0, 0};\n blit.dstOffsets[1] = {mipWidth > 1 ? mipWidth / 2 : 1, mipHeight > 1 ? mipHeight / 2 : 1, 1};\n blit.dstSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;\n blit.dstSubresource.mipLevel = i;\n blit.dstSubresource.baseArrayLayer = 0;\n blit.dstSubresource.layerCount = 1;\n\n vkCmdBlitImage(commandBuffer,\n image,\n VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,\n image,\n VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,\n 1,\n &blit,\n VK_FILTER_LINEAR);\n\n barrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;\n barrier.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;\n barrier.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT;\n barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;\n\n vkCmdPipelineBarrier(commandBuffer,\n VK_PIPELINE_STAGE_TRANSFER_BIT,\n VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,\n 0,\n 0,\n nullptr,\n 0,\n nullptr,\n 1,\n &barrier);\n\n if (mipWidth > 1)\n mipWidth /= 2;\n if (mipHeight > 1)\n mipHeight /= 2;\n }\n\n barrier.subresourceRange.baseMipLevel = mipLevels - 1;\n barrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;\n barrier.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;\n barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;\n barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;\n\n vkCmdPipelineBarrier(commandBuffer,\n VK_PIPELINE_STAGE_TRANSFER_BIT,\n VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,\n 0,\n 0,\n nullptr,\n 0,\n nullptr,\n 1,\n &barrier);\n\n endSingleTimeCommands(commandBuffer);\n }\n\n#ifdef _WIN64 // For windows\n HANDLE getVkImageMemHandle(VkExternalMemoryHandleTypeFlagsKHR externalMemoryHandleType)\n {\n HANDLE handle;\n\n VkMemoryGetWin32HandleInfoKHR vkMemoryGetWin32HandleInfoKHR = {};\n vkMemoryGetWin32HandleInfoKHR.sType = VK_STRUCTURE_TYPE_MEMORY_GET_WIN32_HANDLE_INFO_KHR;\n vkMemoryGetWin32HandleInfoKHR.pNext = NULL;\n vkMemoryGetWin32HandleInfoKHR.memory = textureImageMemory;\n vkMemoryGetWin32HandleInfoKHR.handleType = (VkExternalMemoryHandleTypeFlagBitsKHR)externalMemoryHandleType;\n\n fpGetMemoryWin32HandleKHR(device, &vkMemoryGetWin32HandleInfoKHR, &handle);\n return handle;\n }\n HANDLE getVkSemaphoreHandle(VkExternalSemaphoreHandleTypeFlagBitsKHR externalSemaphoreHandleType,\n VkSemaphore &semVkCuda)\n {\n HANDLE handle;\n\n VkSemaphoreGetWin32HandleInfoKHR vulkanSemaphoreGetWin32HandleInfoKHR = {};\n vulkanSemaphoreGetWin32HandleInfoKHR.sType = VK_STRUCTURE_TYPE_SEMAPHORE_GET_WIN32_HANDLE_INFO_KHR;\n vulkanSemaphoreGetWin32HandleInfoKHR.pNext = NULL;\n vulkanSemaphoreGetWin32HandleInfoKHR.semaphore = semVkCuda;\n vulkanSemaphoreGetWin32HandleInfoKHR.handleType = externalSemaphoreHandleType;\n\n fpGetSemaphoreWin32HandleKHR(device, &vulkanSemaphoreGetWin32HandleInfoKHR, &handle);\n\n return handle;\n }\n#else\n int getVkImageMemHandle(VkExternalMemoryHandleTypeFlagsKHR externalMemoryHandleType)\n {\n if (externalMemoryHandleType == VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT_KHR) {\n int fd;\n\n VkMemoryGetFdInfoKHR vkMemoryGetFdInfoKHR = {};\n vkMemoryGetFdInfoKHR.sType = VK_STRUCTURE_TYPE_MEMORY_GET_FD_INFO_KHR;\n vkMemoryGetFdInfoKHR.pNext = NULL;\n vkMemoryGetFdInfoKHR.memory = textureImageMemory;\n vkMemoryGetFdInfoKHR.handleType = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT_KHR;\n\n fpGetMemoryFdKHR(device, &vkMemoryGetFdInfoKHR, &fd);\n\n return fd;\n }\n return -1;\n }\n\n int getVkSemaphoreHandle(VkExternalSemaphoreHandleTypeFlagBitsKHR externalSemaphoreHandleType,\n VkSemaphore &semVkCuda)\n {\n if (externalSemaphoreHandleType == VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_FD_BIT) {\n int fd;\n\n VkSemaphoreGetFdInfoKHR vulkanSemaphoreGetFdInfoKHR = {};\n vulkanSemaphoreGetFdInfoKHR.sType = VK_STRUCTURE_TYPE_SEMAPHORE_GET_FD_INFO_KHR;\n vulkanSemaphoreGetFdInfoKHR.pNext = NULL;\n vulkanSemaphoreGetFdInfoKHR.semaphore = semVkCuda;\n vulkanSemaphoreGetFdInfoKHR.handleType = VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_FD_BIT_KHR;\n\n fpGetSemaphoreFdKHR(device, &vulkanSemaphoreGetFdInfoKHR, &fd);\n\n return fd;\n }\n return -1;\n }\n#endif\n\n void createTextureImageView() { textureImageView = createImageView(textureImage, VK_FORMAT_R8G8B8A8_UNORM); }\n\n void createTextureSampler()\n {\n VkSamplerCreateInfo samplerInfo = {};\n samplerInfo.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;\n samplerInfo.magFilter = VK_FILTER_LINEAR;\n samplerInfo.minFilter = VK_FILTER_LINEAR;\n samplerInfo.addressModeU = VK_SAMPLER_ADDRESS_MODE_REPEAT;\n samplerInfo.addressModeV = VK_SAMPLER_ADDRESS_MODE_REPEAT;\n samplerInfo.addressModeW = VK_SAMPLER_ADDRESS_MODE_REPEAT;\n samplerInfo.anisotropyEnable = VK_TRUE;\n samplerInfo.maxAnisotropy = 16;\n samplerInfo.borderColor = VK_BORDER_COLOR_INT_OPAQUE_BLACK;\n samplerInfo.unnormalizedCoordinates = VK_FALSE;\n samplerInfo.compareEnable = VK_FALSE;\n samplerInfo.compareOp = VK_COMPARE_OP_ALWAYS;\n samplerInfo.mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR;\n samplerInfo.minLod = 0; // Optional\n samplerInfo.maxLod = static_cast(mipLevels);\n samplerInfo.mipLodBias = 0; // Optional\n\n if (vkCreateSampler(device, &samplerInfo, nullptr, &textureSampler) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create texture sampler!\");\n }\n }\n\n VkImageView createImageView(VkImage image, VkFormat format)\n {\n VkImageViewCreateInfo viewInfo = {};\n viewInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;\n viewInfo.image = image;\n viewInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;\n viewInfo.format = format;\n viewInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;\n viewInfo.subresourceRange.baseMipLevel = 0;\n viewInfo.subresourceRange.levelCount = mipLevels;\n viewInfo.subresourceRange.baseArrayLayer = 0;\n viewInfo.subresourceRange.layerCount = 1;\n\n VkImageView imageView;\n if (vkCreateImageView(device, &viewInfo, nullptr, &imageView) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create texture image view!\");\n }\n\n return imageView;\n }\n\n void createImage(uint32_t width,\n uint32_t height,\n VkFormat format,\n VkImageTiling tiling,\n VkImageUsageFlags usage,\n VkMemoryPropertyFlags properties,\n VkImage &image,\n VkDeviceMemory &imageMemory)\n {\n VkImageCreateInfo imageInfo = {};\n imageInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;\n imageInfo.imageType = VK_IMAGE_TYPE_2D;\n imageInfo.extent.width = width;\n imageInfo.extent.height = height;\n imageInfo.extent.depth = 1;\n imageInfo.mipLevels = mipLevels;\n imageInfo.arrayLayers = 1;\n imageInfo.format = format;\n imageInfo.tiling = tiling;\n imageInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;\n imageInfo.usage = usage;\n imageInfo.samples = VK_SAMPLE_COUNT_1_BIT;\n imageInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;\n\n VkExternalMemoryImageCreateInfo vkExternalMemImageCreateInfo = {};\n vkExternalMemImageCreateInfo.sType = VK_STRUCTURE_TYPE_EXTERNAL_MEMORY_IMAGE_CREATE_INFO;\n vkExternalMemImageCreateInfo.pNext = NULL;\n#ifdef _WIN64\n vkExternalMemImageCreateInfo.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT;\n#else\n vkExternalMemImageCreateInfo.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT_KHR;\n#endif\n\n imageInfo.pNext = &vkExternalMemImageCreateInfo;\n\n if (vkCreateImage(device, &imageInfo, nullptr, &image) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create image!\");\n }\n\n VkMemoryRequirements memRequirements;\n vkGetImageMemoryRequirements(device, image, &memRequirements);\n\n#ifdef _WIN64\n WindowsSecurityAttributes winSecurityAttributes;\n\n VkExportMemoryWin32HandleInfoKHR vulkanExportMemoryWin32HandleInfoKHR = {};\n vulkanExportMemoryWin32HandleInfoKHR.sType = VK_STRUCTURE_TYPE_EXPORT_MEMORY_WIN32_HANDLE_INFO_KHR;\n vulkanExportMemoryWin32HandleInfoKHR.pNext = NULL;\n vulkanExportMemoryWin32HandleInfoKHR.pAttributes = &winSecurityAttributes;\n vulkanExportMemoryWin32HandleInfoKHR.dwAccess = DXGI_SHARED_RESOURCE_READ | DXGI_SHARED_RESOURCE_WRITE;\n vulkanExportMemoryWin32HandleInfoKHR.name = (LPCWSTR)NULL;\n#endif\n VkExportMemoryAllocateInfoKHR vulkanExportMemoryAllocateInfoKHR = {};\n vulkanExportMemoryAllocateInfoKHR.sType = VK_STRUCTURE_TYPE_EXPORT_MEMORY_ALLOCATE_INFO_KHR;\n#ifdef _WIN64\n vulkanExportMemoryAllocateInfoKHR.pNext = IsWindows8OrGreater() ? &vulkanExportMemoryWin32HandleInfoKHR : NULL;\n vulkanExportMemoryAllocateInfoKHR.handleTypes = IsWindows8OrGreater()\n ? VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT\n : VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_KMT_BIT;\n#else\n vulkanExportMemoryAllocateInfoKHR.pNext = NULL;\n vulkanExportMemoryAllocateInfoKHR.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT_KHR;\n#endif\n\n VkMemoryAllocateInfo allocInfo = {};\n allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;\n allocInfo.allocationSize = memRequirements.size;\n allocInfo.pNext = &vulkanExportMemoryAllocateInfoKHR;\n allocInfo.memoryTypeIndex = findMemoryType(memRequirements.memoryTypeBits, properties);\n\n VkMemoryRequirements vkMemoryRequirements = {};\n vkGetImageMemoryRequirements(device, image, &vkMemoryRequirements);\n totalImageMemSize = vkMemoryRequirements.size;\n\n if (vkAllocateMemory(device, &allocInfo, nullptr, &textureImageMemory) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to allocate image memory!\");\n }\n\n vkBindImageMemory(device, image, textureImageMemory, 0);\n }\n\n void cudaVkImportSemaphore()\n {\n cudaExternalSemaphoreHandleDesc externalSemaphoreHandleDesc;\n memset(&externalSemaphoreHandleDesc, 0, sizeof(externalSemaphoreHandleDesc));\n#ifdef _WIN64\n externalSemaphoreHandleDesc.type = IsWindows8OrGreater() ? cudaExternalSemaphoreHandleTypeOpaqueWin32\n : cudaExternalSemaphoreHandleTypeOpaqueWin32Kmt;\n externalSemaphoreHandleDesc.handle.win32.handle =\n getVkSemaphoreHandle(IsWindows8OrGreater() ? VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_WIN32_BIT\n : VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_WIN32_KMT_BIT,\n cudaUpdateVkSemaphore);\n#else\n externalSemaphoreHandleDesc.type = cudaExternalSemaphoreHandleTypeOpaqueFd;\n externalSemaphoreHandleDesc.handle.fd =\n getVkSemaphoreHandle(VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_FD_BIT, cudaUpdateVkSemaphore);\n#endif\n externalSemaphoreHandleDesc.flags = 0;\n\n checkCudaErrors(cudaImportExternalSemaphore(&cudaExtCudaUpdateVkSemaphore, &externalSemaphoreHandleDesc));\n\n memset(&externalSemaphoreHandleDesc, 0, sizeof(externalSemaphoreHandleDesc));\n#ifdef _WIN64\n externalSemaphoreHandleDesc.type = IsWindows8OrGreater() ? cudaExternalSemaphoreHandleTypeOpaqueWin32\n : cudaExternalSemaphoreHandleTypeOpaqueWin32Kmt;\n ;\n externalSemaphoreHandleDesc.handle.win32.handle =\n getVkSemaphoreHandle(IsWindows8OrGreater() ? VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_WIN32_BIT\n : VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_WIN32_KMT_BIT,\n vkUpdateCudaSemaphore);\n#else\n externalSemaphoreHandleDesc.type = cudaExternalSemaphoreHandleTypeOpaqueFd;\n externalSemaphoreHandleDesc.handle.fd =\n getVkSemaphoreHandle(VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_FD_BIT, vkUpdateCudaSemaphore);\n#endif\n externalSemaphoreHandleDesc.flags = 0;\n checkCudaErrors(cudaImportExternalSemaphore(&cudaExtVkUpdateCudaSemaphore, &externalSemaphoreHandleDesc));\n printf(\"CUDA Imported Vulkan semaphore\\n\");\n }\n\n void cudaVkImportImageMem()\n {\n cudaExternalMemoryHandleDesc cudaExtMemHandleDesc;\n memset(&cudaExtMemHandleDesc, 0, sizeof(cudaExtMemHandleDesc));\n#ifdef _WIN64\n cudaExtMemHandleDesc.type = IsWindows8OrGreater() ? cudaExternalMemoryHandleTypeOpaqueWin32\n : cudaExternalMemoryHandleTypeOpaqueWin32Kmt;\n cudaExtMemHandleDesc.handle.win32.handle =\n getVkImageMemHandle(IsWindows8OrGreater() ? VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT\n : VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_KMT_BIT);\n#else\n cudaExtMemHandleDesc.type = cudaExternalMemoryHandleTypeOpaqueFd;\n\n cudaExtMemHandleDesc.handle.fd = getVkImageMemHandle(VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT_KHR);\n#endif\n cudaExtMemHandleDesc.size = totalImageMemSize;\n\n checkCudaErrors(cudaImportExternalMemory(&cudaExtMemImageBuffer, &cudaExtMemHandleDesc));\n\n cudaExternalMemoryMipmappedArrayDesc externalMemoryMipmappedArrayDesc;\n\n memset(&externalMemoryMipmappedArrayDesc, 0, sizeof(externalMemoryMipmappedArrayDesc));\n\n cudaExtent extent = make_cudaExtent(imageWidth, imageHeight, 0);\n cudaChannelFormatDesc formatDesc;\n formatDesc.x = 8;\n formatDesc.y = 8;\n formatDesc.z = 8;\n formatDesc.w = 8;\n formatDesc.f = cudaChannelFormatKindUnsigned;\n\n externalMemoryMipmappedArrayDesc.offset = 0;\n externalMemoryMipmappedArrayDesc.formatDesc = formatDesc;\n externalMemoryMipmappedArrayDesc.extent = extent;\n externalMemoryMipmappedArrayDesc.flags = 0;\n externalMemoryMipmappedArrayDesc.numLevels = mipLevels;\n\n checkCudaErrors(cudaExternalMemoryGetMappedMipmappedArray(\n &cudaMipmappedImageArray, cudaExtMemImageBuffer, &externalMemoryMipmappedArrayDesc));\n\n checkCudaErrors(cudaMallocMipmappedArray(&cudaMipmappedImageArrayTemp, &formatDesc, extent, mipLevels));\n checkCudaErrors(cudaMallocMipmappedArray(&cudaMipmappedImageArrayOrig, &formatDesc, extent, mipLevels));\n\n for (int mipLevelIdx = 0; mipLevelIdx < mipLevels; mipLevelIdx++) {\n cudaArray_t cudaMipLevelArray, cudaMipLevelArrayTemp, cudaMipLevelArrayOrig;\n cudaResourceDesc resourceDesc;\n\n checkCudaErrors(cudaGetMipmappedArrayLevel(&cudaMipLevelArray, cudaMipmappedImageArray, mipLevelIdx));\n checkCudaErrors(\n cudaGetMipmappedArrayLevel(&cudaMipLevelArrayTemp, cudaMipmappedImageArrayTemp, mipLevelIdx));\n checkCudaErrors(\n cudaGetMipmappedArrayLevel(&cudaMipLevelArrayOrig, cudaMipmappedImageArrayOrig, mipLevelIdx));\n\n uint32_t width = (imageWidth >> mipLevelIdx) ? (imageWidth >> mipLevelIdx) : 1;\n uint32_t height = (imageHeight >> mipLevelIdx) ? (imageHeight >> mipLevelIdx) : 1;\n checkCudaErrors(cudaMemcpy2DArrayToArray(cudaMipLevelArrayOrig,\n 0,\n 0,\n cudaMipLevelArray,\n 0,\n 0,\n width * sizeof(uchar4),\n height,\n cudaMemcpyDeviceToDevice));\n\n memset(&resourceDesc, 0, sizeof(resourceDesc));\n resourceDesc.resType = cudaResourceTypeArray;\n resourceDesc.res.array.array = cudaMipLevelArray;\n\n cudaSurfaceObject_t surfaceObject;\n checkCudaErrors(cudaCreateSurfaceObject(&surfaceObject, &resourceDesc));\n\n surfaceObjectList.push_back(surfaceObject);\n\n memset(&resourceDesc, 0, sizeof(resourceDesc));\n resourceDesc.resType = cudaResourceTypeArray;\n resourceDesc.res.array.array = cudaMipLevelArrayTemp;\n\n cudaSurfaceObject_t surfaceObjectTemp;\n checkCudaErrors(cudaCreateSurfaceObject(&surfaceObjectTemp, &resourceDesc));\n surfaceObjectListTemp.push_back(surfaceObjectTemp);\n }\n\n cudaResourceDesc resDescr;\n memset(&resDescr, 0, sizeof(cudaResourceDesc));\n\n resDescr.resType = cudaResourceTypeMipmappedArray;\n resDescr.res.mipmap.mipmap = cudaMipmappedImageArrayOrig;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = cudaFilterModeLinear;\n texDescr.mipmapFilterMode = cudaFilterModeLinear;\n\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n\n texDescr.maxMipmapLevelClamp = float(mipLevels - 1);\n\n texDescr.readMode = cudaReadModeNormalizedFloat;\n\n checkCudaErrors(cudaCreateTextureObject(&textureObjMipMapInput, &resDescr, &texDescr, NULL));\n\n checkCudaErrors(cudaMalloc((void **)&d_surfaceObjectList, sizeof(cudaSurfaceObject_t) * mipLevels));\n checkCudaErrors(cudaMalloc((void **)&d_surfaceObjectListTemp, sizeof(cudaSurfaceObject_t) * mipLevels));\n\n checkCudaErrors(cudaMemcpy(d_surfaceObjectList,\n surfaceObjectList.data(),\n sizeof(cudaSurfaceObject_t) * mipLevels,\n cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(d_surfaceObjectListTemp,\n surfaceObjectListTemp.data(),\n sizeof(cudaSurfaceObject_t) * mipLevels,\n cudaMemcpyHostToDevice));\n\n printf(\"CUDA Kernel Vulkan image buffer\\n\");\n }\n\n void cudaUpdateVkImage()\n {\n cudaVkSemaphoreWait(cudaExtVkUpdateCudaSemaphore);\n\n int nthreads = 128;\n\n /*Perform 2D box filter on image using CUDA */\n d_boxfilter_rgba_x<<>>(\n d_surfaceObjectListTemp, textureObjMipMapInput, imageWidth, imageHeight, mipLevels, filter_radius);\n\n d_boxfilter_rgba_y<<>>(\n d_surfaceObjectList, d_surfaceObjectListTemp, imageWidth, imageHeight, mipLevels, filter_radius);\n\n varySigma();\n\n cudaVkSemaphoreSignal(cudaExtCudaUpdateVkSemaphore);\n }\n\n void transitionImageLayout(VkImage image, VkFormat format, VkImageLayout oldLayout, VkImageLayout newLayout)\n {\n VkCommandBuffer commandBuffer = beginSingleTimeCommands();\n\n VkImageMemoryBarrier barrier = {};\n barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;\n barrier.oldLayout = oldLayout;\n barrier.newLayout = newLayout;\n barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;\n barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;\n barrier.image = image;\n barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;\n barrier.subresourceRange.baseMipLevel = 0;\n barrier.subresourceRange.levelCount = mipLevels;\n barrier.subresourceRange.baseArrayLayer = 0;\n barrier.subresourceRange.layerCount = 1;\n\n VkPipelineStageFlags sourceStage;\n VkPipelineStageFlags destinationStage;\n\n if (oldLayout == VK_IMAGE_LAYOUT_UNDEFINED && newLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL) {\n barrier.srcAccessMask = 0;\n barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;\n\n sourceStage = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;\n destinationStage = VK_PIPELINE_STAGE_TRANSFER_BIT;\n }\n else if (oldLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL\n && newLayout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL) {\n barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;\n barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;\n\n sourceStage = VK_PIPELINE_STAGE_TRANSFER_BIT;\n destinationStage = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;\n }\n else {\n throw std::invalid_argument(\"unsupported layout transition!\");\n }\n\n vkCmdPipelineBarrier(commandBuffer, sourceStage, destinationStage, 0, 0, nullptr, 0, nullptr, 1, &barrier);\n\n endSingleTimeCommands(commandBuffer);\n }\n\n void copyBufferToImage(VkBuffer buffer, VkImage image, uint32_t width, uint32_t height)\n {\n VkCommandBuffer commandBuffer = beginSingleTimeCommands();\n\n VkBufferImageCopy region = {};\n region.bufferOffset = 0;\n region.bufferRowLength = 0;\n region.bufferImageHeight = 0;\n region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;\n region.imageSubresource.mipLevel = 0;\n region.imageSubresource.baseArrayLayer = 0;\n region.imageSubresource.layerCount = 1;\n region.imageOffset = {0, 0, 0};\n region.imageExtent = {width, height, 1};\n\n vkCmdCopyBufferToImage(commandBuffer, buffer, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, ®ion);\n\n endSingleTimeCommands(commandBuffer);\n }\n\n void createVertexBuffer()\n {\n VkDeviceSize bufferSize = sizeof(vertices[0]) * vertices.size();\n\n VkBuffer stagingBuffer;\n VkDeviceMemory stagingBufferMemory;\n createBuffer(bufferSize,\n VK_BUFFER_USAGE_TRANSFER_SRC_BIT,\n VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,\n stagingBuffer,\n stagingBufferMemory);\n\n void *data;\n vkMapMemory(device, stagingBufferMemory, 0, bufferSize, 0, &data);\n memcpy(data, vertices.data(), (size_t)bufferSize);\n vkUnmapMemory(device, stagingBufferMemory);\n\n createBuffer(bufferSize,\n VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,\n VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,\n vertexBuffer,\n vertexBufferMemory);\n\n copyBuffer(stagingBuffer, vertexBuffer, bufferSize);\n\n vkDestroyBuffer(device, stagingBuffer, nullptr);\n vkFreeMemory(device, stagingBufferMemory, nullptr);\n }\n\n void createIndexBuffer()\n {\n VkDeviceSize bufferSize = sizeof(indices[0]) * indices.size();\n\n VkBuffer stagingBuffer;\n VkDeviceMemory stagingBufferMemory;\n createBuffer(bufferSize,\n VK_BUFFER_USAGE_TRANSFER_SRC_BIT,\n VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,\n stagingBuffer,\n stagingBufferMemory);\n\n void *data;\n vkMapMemory(device, stagingBufferMemory, 0, bufferSize, 0, &data);\n memcpy(data, indices.data(), (size_t)bufferSize);\n vkUnmapMemory(device, stagingBufferMemory);\n\n createBuffer(bufferSize,\n VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_INDEX_BUFFER_BIT,\n VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,\n indexBuffer,\n indexBufferMemory);\n\n copyBuffer(stagingBuffer, indexBuffer, bufferSize);\n\n vkDestroyBuffer(device, stagingBuffer, nullptr);\n vkFreeMemory(device, stagingBufferMemory, nullptr);\n }\n\n void createUniformBuffers()\n {\n VkDeviceSize bufferSize = sizeof(UniformBufferObject);\n\n uniformBuffers.resize(swapChainImages.size());\n uniformBuffersMemory.resize(swapChainImages.size());\n\n for (size_t i = 0; i < swapChainImages.size(); i++) {\n createBuffer(bufferSize,\n VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,\n VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,\n uniformBuffers[i],\n uniformBuffersMemory[i]);\n }\n }\n\n void createDescriptorPool()\n {\n std::array poolSizes = {};\n poolSizes[0].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;\n poolSizes[0].descriptorCount = static_cast(swapChainImages.size());\n poolSizes[1].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;\n poolSizes[1].descriptorCount = static_cast(swapChainImages.size());\n\n VkDescriptorPoolCreateInfo poolInfo = {};\n poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;\n poolInfo.poolSizeCount = static_cast(poolSizes.size());\n poolInfo.pPoolSizes = poolSizes.data();\n poolInfo.maxSets = static_cast(swapChainImages.size());\n\n if (vkCreateDescriptorPool(device, &poolInfo, nullptr, &descriptorPool) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create descriptor pool!\");\n }\n }\n\n void createDescriptorSets()\n {\n std::vector layouts(swapChainImages.size(), descriptorSetLayout);\n VkDescriptorSetAllocateInfo allocInfo = {};\n allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;\n allocInfo.descriptorPool = descriptorPool;\n allocInfo.descriptorSetCount = static_cast(swapChainImages.size());\n allocInfo.pSetLayouts = layouts.data();\n\n descriptorSets.resize(swapChainImages.size());\n if (vkAllocateDescriptorSets(device, &allocInfo, descriptorSets.data()) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to allocate descriptor sets!\");\n }\n\n for (size_t i = 0; i < swapChainImages.size(); i++) {\n VkDescriptorBufferInfo bufferInfo = {};\n bufferInfo.buffer = uniformBuffers[i];\n bufferInfo.offset = 0;\n bufferInfo.range = sizeof(UniformBufferObject);\n\n VkDescriptorImageInfo imageInfo = {};\n imageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;\n imageInfo.imageView = textureImageView;\n imageInfo.sampler = textureSampler;\n\n std::array descriptorWrites = {};\n\n descriptorWrites[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;\n descriptorWrites[0].dstSet = descriptorSets[i];\n descriptorWrites[0].dstBinding = 0;\n descriptorWrites[0].dstArrayElement = 0;\n descriptorWrites[0].descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;\n descriptorWrites[0].descriptorCount = 1;\n descriptorWrites[0].pBufferInfo = &bufferInfo;\n\n descriptorWrites[1].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;\n descriptorWrites[1].dstSet = descriptorSets[i];\n descriptorWrites[1].dstBinding = 1;\n descriptorWrites[1].dstArrayElement = 0;\n descriptorWrites[1].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;\n descriptorWrites[1].descriptorCount = 1;\n descriptorWrites[1].pImageInfo = &imageInfo;\n\n vkUpdateDescriptorSets(\n device, static_cast(descriptorWrites.size()), descriptorWrites.data(), 0, nullptr);\n }\n }\n\n void createBuffer(VkDeviceSize size,\n VkBufferUsageFlags usage,\n VkMemoryPropertyFlags properties,\n VkBuffer &buffer,\n VkDeviceMemory &bufferMemory)\n {\n VkBufferCreateInfo bufferInfo = {};\n bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;\n bufferInfo.size = size;\n bufferInfo.usage = usage;\n bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;\n\n if (vkCreateBuffer(device, &bufferInfo, nullptr, &buffer) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create buffer!\");\n }\n\n VkMemoryRequirements memRequirements;\n vkGetBufferMemoryRequirements(device, buffer, &memRequirements);\n\n VkMemoryAllocateInfo allocInfo = {};\n allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;\n allocInfo.allocationSize = memRequirements.size;\n allocInfo.memoryTypeIndex = findMemoryType(memRequirements.memoryTypeBits, properties);\n\n if (vkAllocateMemory(device, &allocInfo, nullptr, &bufferMemory) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to allocate buffer memory!\");\n }\n\n vkBindBufferMemory(device, buffer, bufferMemory, 0);\n }\n\n VkCommandBuffer beginSingleTimeCommands()\n {\n VkCommandBufferAllocateInfo allocInfo = {};\n allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;\n allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;\n allocInfo.commandPool = commandPool;\n allocInfo.commandBufferCount = 1;\n\n VkCommandBuffer commandBuffer;\n vkAllocateCommandBuffers(device, &allocInfo, &commandBuffer);\n\n VkCommandBufferBeginInfo beginInfo = {};\n beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;\n beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;\n\n vkBeginCommandBuffer(commandBuffer, &beginInfo);\n\n return commandBuffer;\n }\n\n void endSingleTimeCommands(VkCommandBuffer commandBuffer)\n {\n vkEndCommandBuffer(commandBuffer);\n\n VkSubmitInfo submitInfo = {};\n submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;\n submitInfo.commandBufferCount = 1;\n submitInfo.pCommandBuffers = &commandBuffer;\n\n vkQueueSubmit(graphicsQueue, 1, &submitInfo, VK_NULL_HANDLE);\n vkQueueWaitIdle(graphicsQueue);\n\n vkFreeCommandBuffers(device, commandPool, 1, &commandBuffer);\n }\n\n void copyBuffer(VkBuffer srcBuffer, VkBuffer dstBuffer, VkDeviceSize size)\n {\n VkCommandBuffer commandBuffer = beginSingleTimeCommands();\n\n VkBufferCopy copyRegion = {};\n copyRegion.size = size;\n vkCmdCopyBuffer(commandBuffer, srcBuffer, dstBuffer, 1, ©Region);\n\n endSingleTimeCommands(commandBuffer);\n }\n\n uint32_t findMemoryType(uint32_t typeFilter, VkMemoryPropertyFlags properties)\n {\n VkPhysicalDeviceMemoryProperties memProperties;\n vkGetPhysicalDeviceMemoryProperties(physicalDevice, &memProperties);\n\n for (uint32_t i = 0; i < memProperties.memoryTypeCount; i++) {\n if ((typeFilter & (1 << i)) && (memProperties.memoryTypes[i].propertyFlags & properties) == properties) {\n return i;\n }\n }\n\n throw std::runtime_error(\"failed to find suitable memory type!\");\n }\n\n void createCommandBuffers()\n {\n commandBuffers.resize(swapChainFramebuffers.size());\n\n VkCommandBufferAllocateInfo allocInfo = {};\n allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;\n allocInfo.commandPool = commandPool;\n allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;\n allocInfo.commandBufferCount = (uint32_t)commandBuffers.size();\n\n if (vkAllocateCommandBuffers(device, &allocInfo, commandBuffers.data()) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to allocate command buffers!\");\n }\n\n for (size_t i = 0; i < commandBuffers.size(); i++) {\n VkCommandBufferBeginInfo beginInfo = {};\n beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;\n beginInfo.flags = VK_COMMAND_BUFFER_USAGE_SIMULTANEOUS_USE_BIT;\n\n if (vkBeginCommandBuffer(commandBuffers[i], &beginInfo) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to begin recording command buffer!\");\n }\n\n VkRenderPassBeginInfo renderPassInfo = {};\n renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;\n renderPassInfo.renderPass = renderPass;\n renderPassInfo.framebuffer = swapChainFramebuffers[i];\n renderPassInfo.renderArea.offset = {0, 0};\n renderPassInfo.renderArea.extent = swapChainExtent;\n\n VkClearValue clearColor = {0.0f, 0.0f, 0.0f, 1.0f};\n renderPassInfo.clearValueCount = 1;\n renderPassInfo.pClearValues = &clearColor;\n\n vkCmdBeginRenderPass(commandBuffers[i], &renderPassInfo, VK_SUBPASS_CONTENTS_INLINE);\n\n vkCmdBindPipeline(commandBuffers[i], VK_PIPELINE_BIND_POINT_GRAPHICS, graphicsPipeline);\n\n VkBuffer vertexBuffers[] = {vertexBuffer};\n VkDeviceSize offsets[] = {0};\n vkCmdBindVertexBuffers(commandBuffers[i], 0, 1, vertexBuffers, offsets);\n\n vkCmdBindIndexBuffer(commandBuffers[i], indexBuffer, 0, VK_INDEX_TYPE_UINT16);\n\n vkCmdBindDescriptorSets(commandBuffers[i],\n VK_PIPELINE_BIND_POINT_GRAPHICS,\n pipelineLayout,\n 0,\n 1,\n &descriptorSets[i],\n 0,\n nullptr);\n\n vkCmdDrawIndexed(commandBuffers[i], static_cast(indices.size()), 1, 0, 0, 0);\n // vkCmdDraw(commandBuffers[i], static_cast(vertices.size()), 1,\n // 0, 0);\n\n vkCmdEndRenderPass(commandBuffers[i]);\n\n if (vkEndCommandBuffer(commandBuffers[i]) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to record command buffer!\");\n }\n }\n }\n\n void createSyncObjects()\n {\n imageAvailableSemaphores.resize(MAX_FRAMES);\n renderFinishedSemaphores.resize(MAX_FRAMES);\n inFlightFences.resize(MAX_FRAMES);\n\n VkSemaphoreCreateInfo semaphoreInfo = {};\n semaphoreInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;\n\n VkFenceCreateInfo fenceInfo = {};\n fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;\n fenceInfo.flags = VK_FENCE_CREATE_SIGNALED_BIT;\n\n for (size_t i = 0; i < MAX_FRAMES; i++) {\n if (vkCreateSemaphore(device, &semaphoreInfo, nullptr, &imageAvailableSemaphores[i]) != VK_SUCCESS\n || vkCreateSemaphore(device, &semaphoreInfo, nullptr, &renderFinishedSemaphores[i]) != VK_SUCCESS\n || vkCreateFence(device, &fenceInfo, nullptr, &inFlightFences[i]) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create synchronization objects for a frame!\");\n }\n }\n }\n\n void createSyncObjectsExt()\n {\n VkSemaphoreCreateInfo semaphoreInfo = {};\n semaphoreInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;\n\n memset(&semaphoreInfo, 0, sizeof(semaphoreInfo));\n semaphoreInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;\n\n#ifdef _WIN64\n WindowsSecurityAttributes winSecurityAttributes;\n\n VkExportSemaphoreWin32HandleInfoKHR vulkanExportSemaphoreWin32HandleInfoKHR = {};\n vulkanExportSemaphoreWin32HandleInfoKHR.sType = VK_STRUCTURE_TYPE_EXPORT_SEMAPHORE_WIN32_HANDLE_INFO_KHR;\n vulkanExportSemaphoreWin32HandleInfoKHR.pNext = NULL;\n vulkanExportSemaphoreWin32HandleInfoKHR.pAttributes = &winSecurityAttributes;\n vulkanExportSemaphoreWin32HandleInfoKHR.dwAccess = DXGI_SHARED_RESOURCE_READ | DXGI_SHARED_RESOURCE_WRITE;\n vulkanExportSemaphoreWin32HandleInfoKHR.name = (LPCWSTR)NULL;\n#endif\n VkExportSemaphoreCreateInfoKHR vulkanExportSemaphoreCreateInfo = {};\n vulkanExportSemaphoreCreateInfo.sType = VK_STRUCTURE_TYPE_EXPORT_SEMAPHORE_CREATE_INFO_KHR;\n#ifdef _WIN64\n vulkanExportSemaphoreCreateInfo.pNext = IsWindows8OrGreater() ? &vulkanExportSemaphoreWin32HandleInfoKHR : NULL;\n vulkanExportSemaphoreCreateInfo.handleTypes = IsWindows8OrGreater()\n ? VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_WIN32_BIT\n : VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_WIN32_KMT_BIT;\n#else\n vulkanExportSemaphoreCreateInfo.pNext = NULL;\n vulkanExportSemaphoreCreateInfo.handleTypes = VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_FD_BIT;\n#endif\n semaphoreInfo.pNext = &vulkanExportSemaphoreCreateInfo;\n\n if (vkCreateSemaphore(device, &semaphoreInfo, nullptr, &cudaUpdateVkSemaphore) != VK_SUCCESS\n || vkCreateSemaphore(device, &semaphoreInfo, nullptr, &vkUpdateCudaSemaphore) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create synchronization objects for a CUDA-Vulkan!\");\n }\n }\n\n void updateUniformBuffer()\n {\n UniformBufferObject ubo = {};\n\n mat4x4_identity(ubo.model);\n mat4x4 Model;\n mat4x4_dup(Model, ubo.model);\n mat4x4_rotate(ubo.model, Model, 0.0f, 0.0f, 1.0f, degreesToRadians(135.0f));\n\n vec3 eye = {2.0f, 2.0f, 2.0f};\n vec3 center = {0.0f, 0.0f, 0.0f};\n vec3 up = {0.0f, 0.0f, 1.0f};\n mat4x4_look_at(ubo.view, eye, center, up);\n\n mat4x4_perspective(\n ubo.proj, degreesToRadians(45.0f), swapChainExtent.width / (float)swapChainExtent.height, 0.1f, 10.0f);\n ubo.proj[1][1] *= -1;\n\n for (size_t i = 0; i < swapChainImages.size(); i++) {\n void *data;\n vkMapMemory(device, uniformBuffersMemory[i], 0, sizeof(ubo), 0, &data);\n memcpy(data, &ubo, sizeof(ubo));\n vkUnmapMemory(device, uniformBuffersMemory[i]);\n }\n }\n\n void drawFrame()\n {\n static int startSubmit = 0;\n\n vkWaitForFences(device, 1, &inFlightFences[currentFrame], VK_TRUE, std::numeric_limits::max());\n\n uint32_t imageIndex;\n VkResult result = vkAcquireNextImageKHR(device,\n swapChain,\n std::numeric_limits::max(),\n imageAvailableSemaphores[currentFrame],\n VK_NULL_HANDLE,\n &imageIndex);\n\n if (result == VK_ERROR_OUT_OF_DATE_KHR) {\n recreateSwapChain();\n return;\n }\n else if (result != VK_SUCCESS && result != VK_SUBOPTIMAL_KHR) {\n throw std::runtime_error(\"failed to acquire swap chain image!\");\n }\n\n vkResetFences(device, 1, &inFlightFences[currentFrame]);\n\n if (!startSubmit) {\n submitVulkan(imageIndex);\n startSubmit = 1;\n }\n else {\n submitVulkanCuda(imageIndex);\n }\n\n VkPresentInfoKHR presentInfo = {};\n presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;\n\n VkSemaphore signalSemaphores[] = {renderFinishedSemaphores[currentFrame]};\n\n presentInfo.waitSemaphoreCount = 1;\n presentInfo.pWaitSemaphores = signalSemaphores;\n\n VkSwapchainKHR swapChains[] = {swapChain};\n presentInfo.swapchainCount = 1;\n presentInfo.pSwapchains = swapChains;\n presentInfo.pImageIndices = &imageIndex;\n presentInfo.pResults = nullptr; // Optional\n\n result = vkQueuePresentKHR(presentQueue, &presentInfo);\n\n if (result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR || framebufferResized) {\n framebufferResized = false;\n recreateSwapChain();\n }\n else if (result != VK_SUCCESS) {\n throw std::runtime_error(\"failed to present swap chain image!\");\n }\n\n cudaUpdateVkImage();\n\n currentFrame = (currentFrame + 1) % MAX_FRAMES;\n // Added sleep of 10 millisecs so that CPU does not submit too much work to\n // GPU\n std::this_thread::sleep_for(std::chrono::microseconds(10000));\n char title[256];\n sprintf(title, \"Vulkan Image CUDA Box Filter (radius=%d)\", filter_radius);\n glfwSetWindowTitle(window, title);\n }\n\n void cudaVkSemaphoreSignal(cudaExternalSemaphore_t &extSemaphore)\n {\n cudaExternalSemaphoreSignalParams extSemaphoreSignalParams;\n memset(&extSemaphoreSignalParams, 0, sizeof(extSemaphoreSignalParams));\n\n extSemaphoreSignalParams.params.fence.value = 0;\n extSemaphoreSignalParams.flags = 0;\n checkCudaErrors(cudaSignalExternalSemaphoresAsync(&extSemaphore, &extSemaphoreSignalParams, 1, streamToRun));\n }\n\n void cudaVkSemaphoreWait(cudaExternalSemaphore_t &extSemaphore)\n {\n cudaExternalSemaphoreWaitParams extSemaphoreWaitParams;\n\n memset(&extSemaphoreWaitParams, 0, sizeof(extSemaphoreWaitParams));\n\n extSemaphoreWaitParams.params.fence.value = 0;\n extSemaphoreWaitParams.flags = 0;\n\n checkCudaErrors(cudaWaitExternalSemaphoresAsync(&extSemaphore, &extSemaphoreWaitParams, 1, streamToRun));\n }\n\n void submitVulkan(uint32_t imageIndex)\n {\n VkSubmitInfo submitInfo = {};\n submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;\n\n VkSemaphore waitSemaphores[] = {imageAvailableSemaphores[currentFrame]};\n VkPipelineStageFlags waitStages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};\n submitInfo.waitSemaphoreCount = 1;\n submitInfo.pWaitSemaphores = waitSemaphores;\n submitInfo.pWaitDstStageMask = waitStages;\n submitInfo.commandBufferCount = 1;\n submitInfo.pCommandBuffers = &commandBuffers[imageIndex];\n\n VkSemaphore signalSemaphores[] = {renderFinishedSemaphores[currentFrame], vkUpdateCudaSemaphore};\n\n submitInfo.signalSemaphoreCount = 2;\n submitInfo.pSignalSemaphores = signalSemaphores;\n\n if (vkQueueSubmit(graphicsQueue, 1, &submitInfo, inFlightFences[currentFrame]) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to submit draw command buffer!\");\n }\n }\n\n void submitVulkanCuda(uint32_t imageIndex)\n {\n VkSubmitInfo submitInfo = {};\n submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;\n\n VkSemaphore waitSemaphores[] = {imageAvailableSemaphores[currentFrame], cudaUpdateVkSemaphore};\n VkPipelineStageFlags waitStages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,\n VK_PIPELINE_STAGE_ALL_COMMANDS_BIT};\n submitInfo.waitSemaphoreCount = 2;\n submitInfo.pWaitSemaphores = waitSemaphores;\n submitInfo.pWaitDstStageMask = waitStages;\n submitInfo.commandBufferCount = 1;\n submitInfo.pCommandBuffers = &commandBuffers[imageIndex];\n\n VkSemaphore signalSemaphores[] = {renderFinishedSemaphores[currentFrame], vkUpdateCudaSemaphore};\n\n submitInfo.signalSemaphoreCount = 2;\n submitInfo.pSignalSemaphores = signalSemaphores;\n\n if (vkQueueSubmit(graphicsQueue, 1, &submitInfo, inFlightFences[currentFrame]) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to submit draw command buffer!\");\n }\n }\n\n VkShaderModule createShaderModule(const std::vector &code)\n {\n VkShaderModuleCreateInfo createInfo = {};\n createInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;\n createInfo.codeSize = code.size();\n createInfo.pCode = reinterpret_cast(code.data());\n\n VkShaderModule shaderModule;\n if (vkCreateShaderModule(device, &createInfo, nullptr, &shaderModule) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create shader module!\");\n }\n\n return shaderModule;\n }\n\n VkSurfaceFormatKHR chooseSwapSurfaceFormat(const std::vector &availableFormats)\n {\n if (availableFormats.size() == 1 && availableFormats[0].format == VK_FORMAT_UNDEFINED) {\n return {VK_FORMAT_B8G8R8A8_UNORM, VK_COLOR_SPACE_SRGB_NONLINEAR_KHR};\n }\n\n for (const auto &availableFormat : availableFormats) {\n if (availableFormat.format == VK_FORMAT_B8G8R8A8_UNORM\n && availableFormat.colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR) {\n return availableFormat;\n }\n }\n\n return availableFormats[0];\n }\n\n VkPresentModeKHR chooseSwapPresentMode(const std::vector &availablePresentModes)\n {\n VkPresentModeKHR bestMode = VK_PRESENT_MODE_FIFO_KHR;\n\n for (const auto &availablePresentMode : availablePresentModes) {\n if (availablePresentMode == VK_PRESENT_MODE_MAILBOX_KHR) {\n return availablePresentMode;\n }\n else if (availablePresentMode == VK_PRESENT_MODE_IMMEDIATE_KHR) {\n bestMode = availablePresentMode;\n }\n }\n\n return bestMode;\n }\n\n VkExtent2D chooseSwapExtent(const VkSurfaceCapabilitiesKHR &capabilities)\n {\n if (capabilities.currentExtent.width != std::numeric_limits::max()) {\n return capabilities.currentExtent;\n }\n else {\n int width, height;\n glfwGetFramebufferSize(window, &width, &height);\n\n VkExtent2D actualExtent = {static_cast(width), static_cast(height)};\n\n actualExtent.width = std::max(capabilities.minImageExtent.width,\n std::min(capabilities.maxImageExtent.width, actualExtent.width));\n actualExtent.height = std::max(capabilities.minImageExtent.height,\n std::min(capabilities.maxImageExtent.height, actualExtent.height));\n\n return actualExtent;\n }\n }\n\n SwapChainSupportDetails querySwapChainSupport(VkPhysicalDevice device)\n {\n SwapChainSupportDetails details;\n\n vkGetPhysicalDeviceSurfaceCapabilitiesKHR(device, surface, &details.capabilities);\n\n uint32_t formatCount;\n vkGetPhysicalDeviceSurfaceFormatsKHR(device, surface, &formatCount, nullptr);\n\n if (formatCount != 0) {\n details.formats.resize(formatCount);\n vkGetPhysicalDeviceSurfaceFormatsKHR(device, surface, &formatCount, details.formats.data());\n }\n\n uint32_t presentModeCount;\n vkGetPhysicalDeviceSurfacePresentModesKHR(device, surface, &presentModeCount, nullptr);\n\n if (presentModeCount != 0) {\n details.presentModes.resize(presentModeCount);\n vkGetPhysicalDeviceSurfacePresentModesKHR(device, surface, &presentModeCount, details.presentModes.data());\n }\n\n return details;\n }\n\n bool isDeviceSuitable(VkPhysicalDevice device)\n {\n QueueFamilyIndices indices = findQueueFamilies(device);\n\n bool extensionsSupported = checkDeviceExtensionSupport(device);\n\n bool swapChainAdequate = false;\n if (extensionsSupported) {\n SwapChainSupportDetails swapChainSupport = querySwapChainSupport(device);\n swapChainAdequate = !swapChainSupport.formats.empty() && !swapChainSupport.presentModes.empty();\n }\n\n VkPhysicalDeviceFeatures supportedFeatures;\n vkGetPhysicalDeviceFeatures(device, &supportedFeatures);\n\n return indices.isComplete() && extensionsSupported && swapChainAdequate && supportedFeatures.samplerAnisotropy;\n }\n\n bool checkDeviceExtensionSupport(VkPhysicalDevice device)\n {\n uint32_t extensionCount;\n vkEnumerateDeviceExtensionProperties(device, nullptr, &extensionCount, nullptr);\n\n std::vector availableExtensions(extensionCount);\n vkEnumerateDeviceExtensionProperties(device, nullptr, &extensionCount, availableExtensions.data());\n\n std::set requiredExtensions(deviceExtensions.begin(), deviceExtensions.end());\n\n for (const auto &extension : availableExtensions) {\n requiredExtensions.erase(extension.extensionName);\n }\n\n return requiredExtensions.empty();\n }\n\n QueueFamilyIndices findQueueFamilies(VkPhysicalDevice device)\n {\n QueueFamilyIndices indices;\n\n uint32_t queueFamilyCount = 0;\n vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, nullptr);\n\n std::vector queueFamilies(queueFamilyCount);\n vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, queueFamilies.data());\n\n int i = 0;\n for (const auto &queueFamily : queueFamilies) {\n if (queueFamily.queueCount > 0 && queueFamily.queueFlags & VK_QUEUE_GRAPHICS_BIT) {\n indices.graphicsFamily = i;\n }\n\n VkBool32 presentSupport = false;\n vkGetPhysicalDeviceSurfaceSupportKHR(device, i, surface, &presentSupport);\n\n if (queueFamily.queueCount > 0 && presentSupport) {\n indices.presentFamily = i;\n }\n\n if (indices.isComplete()) {\n break;\n }\n\n i++;\n }\n\n return indices;\n }\n\n std::vector getRequiredExtensions()\n {\n uint32_t glfwExtensionCount = 0;\n const char **glfwExtensions;\n glfwExtensions = glfwGetRequiredInstanceExtensions(&glfwExtensionCount);\n\n std::vector extensions(glfwExtensions, glfwExtensions + glfwExtensionCount);\n\n if (enableValidationLayers) {\n extensions.push_back(VK_EXT_DEBUG_UTILS_EXTENSION_NAME);\n }\n\n return extensions;\n }\n\n bool checkValidationLayerSupport()\n {\n uint32_t layerCount;\n vkEnumerateInstanceLayerProperties(&layerCount, nullptr);\n\n std::vector availableLayers(layerCount);\n vkEnumerateInstanceLayerProperties(&layerCount, availableLayers.data());\n\n for (const char *layerName : validationLayers) {\n bool layerFound = false;\n\n for (const auto &layerProperties : availableLayers) {\n if (strcmp(layerName, layerProperties.layerName) == 0) {\n layerFound = true;\n break;\n }\n }\n\n if (!layerFound) {\n return false;\n }\n }\n\n return true;\n }\n\n static std::vector readFile(const std::string &filename)\n {\n char *file_path = sdkFindFilePath(filename.c_str(), execution_path.c_str());\n std::ifstream file(file_path, std::ios::ate | std::ios::binary);\n\n if (!file.is_open()) {\n throw std::runtime_error(\"failed to open file!\");\n }\n\n size_t fileSize = (size_t)file.tellg();\n std::vector buffer(fileSize);\n\n file.seekg(0);\n file.read(buffer.data(), fileSize);\n\n file.close();\n\n return buffer;\n }\n\n static VKAPI_ATTR VkBool32 VKAPI_CALL debugCallback(VkDebugUtilsMessageSeverityFlagBitsEXT messageSeverity,\n VkDebugUtilsMessageTypeFlagsEXT messageType,\n const VkDebugUtilsMessengerCallbackDataEXT *pCallbackData,\n void *pUserData)\n {\n std::cerr << \"validation layer: \" << pCallbackData->pMessage << std::endl;\n\n return VK_FALSE;\n }\n};\n\nint main(int argc, char **argv)\n{\n execution_path = argv[0];\n std::string image_filename = \"teapot1024.ppm\";\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"file\")) {\n getCmdLineArgumentString(argc, (const char **)argv, \"file\", (char **)&image_filename);\n }\n\n vulkanImageCUDA app;\n\n try {\n // This app only works on ppm images\n app.loadImageData(image_filename);\n app.run();\n }\n catch (const std::exception &e) {\n std::cerr << e.what() << std::endl;\n return EXIT_FAILURE;\n }\n\n return EXIT_SUCCESS;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#define GLFW_INCLUDE_VULKAN\n#ifdef _WIN64\n// Add windows.h to the include path firstly as dependency for other Windows headers\n#include \n// Add other Windows headers\n#include \n#include \n#include \n#define _USE_MATH_DEFINES\n#endif\n\n#include \n#include \n#ifdef _WIN64\n#include \n#endif\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"linmath.h\"\n\n#define WIDTH 800\n#define HEIGHT 600\n\nconst int MAX_FRAMES = 4;\n\nconst std::vector validationLayers = {\"VK_LAYER_KHRONOS_validation\"};\n\n#ifdef NDEBUG\nconst bool enableValidationLayers = false;\n#else\nconst bool enableValidationLayers = true;\n#endif\n\nstd::string execution_path;\n\nVkResult CreateDebugUtilsMessengerEXT(VkInstance instance,\n const VkDebugUtilsMessengerCreateInfoEXT *pCreateInfo,\n const VkAllocationCallbacks *pAllocator,\n VkDebugUtilsMessengerEXT *pDebugMessenger)\n{\n auto func = (PFN_vkCreateDebugUtilsMessengerEXT)vkGetInstanceProcAddr(instance, \"vkCreateDebugUtilsMessengerEXT\");\n if (func != nullptr) {\n return func(instance, pCreateInfo, pAllocator, pDebugMessenger);\n }\n else {\n return VK_ERROR_EXTENSION_NOT_PRESENT;\n }\n};\n\nconst std::vector deviceExtensions = {\n VK_KHR_SWAPCHAIN_EXTENSION_NAME,\n VK_KHR_EXTERNAL_MEMORY_EXTENSION_NAME,\n VK_KHR_EXTERNAL_SEMAPHORE_EXTENSION_NAME,\n#ifdef _WIN64\n VK_KHR_EXTERNAL_MEMORY_WIN32_EXTENSION_NAME,\n VK_KHR_EXTERNAL_SEMAPHORE_WIN32_EXTENSION_NAME,\n#else\n VK_KHR_EXTERNAL_MEMORY_FD_EXTENSION_NAME,\n VK_KHR_EXTERNAL_SEMAPHORE_FD_EXTENSION_NAME,\n#endif\n};\n\n#ifdef _WIN64\nclass WindowsSecurityAttributes\n{\nprotected:\n SECURITY_ATTRIBUTES m_winSecurityAttributes;\n PSECURITY_DESCRIPTOR m_winPSecurityDescriptor;\n\npublic:\n WindowsSecurityAttributes();\n SECURITY_ATTRIBUTES *operator&();\n ~WindowsSecurityAttributes();\n};\n\nWindowsSecurityAttributes::WindowsSecurityAttributes()\n{\n m_winPSecurityDescriptor = (PSECURITY_DESCRIPTOR)calloc(1, SECURITY_DESCRIPTOR_MIN_LENGTH + 2 * sizeof(void **));\n\n PSID *ppSID = (PSID *)((PBYTE)m_winPSecurityDescriptor + SECURITY_DESCRIPTOR_MIN_LENGTH);\n PACL *ppACL = (PACL *)((PBYTE)ppSID + sizeof(PSID *));\n\n InitializeSecurityDescriptor(m_winPSecurityDescriptor, SECURITY_DESCRIPTOR_REVISION);\n\n SID_IDENTIFIER_AUTHORITY sidIdentifierAuthority = SECURITY_WORLD_SID_AUTHORITY;\n AllocateAndInitializeSid(&sidIdentifierAuthority, 1, SECURITY_WORLD_RID, 0, 0, 0, 0, 0, 0, 0, ppSID);\n\n EXPLICIT_ACCESS explicitAccess;\n ZeroMemory(&explicitAccess, sizeof(EXPLICIT_ACCESS));\n explicitAccess.grfAccessPermissions = STANDARD_RIGHTS_ALL | SPECIFIC_RIGHTS_ALL;\n explicitAccess.grfAccessMode = SET_ACCESS;\n explicitAccess.grfInheritance = INHERIT_ONLY;\n explicitAccess.Trustee.TrusteeForm = TRUSTEE_IS_SID;\n explicitAccess.Trustee.TrusteeType = TRUSTEE_IS_WELL_KNOWN_GROUP;\n explicitAccess.Trustee.ptstrName = (LPTSTR)*ppSID;\n\n SetEntriesInAcl(1, &explicitAccess, NULL, ppACL);\n\n SetSecurityDescriptorDacl(m_winPSecurityDescriptor, TRUE, *ppACL, FALSE);\n\n m_winSecurityAttributes.nLength = sizeof(m_winSecurityAttributes);\n m_winSecurityAttributes.lpSecurityDescriptor = m_winPSecurityDescriptor;\n m_winSecurityAttributes.bInheritHandle = TRUE;\n}\n\nSECURITY_ATTRIBUTES *WindowsSecurityAttributes::operator&() { return &m_winSecurityAttributes; }\n\nWindowsSecurityAttributes::~WindowsSecurityAttributes()\n{\n PSID *ppSID = (PSID *)((PBYTE)m_winPSecurityDescriptor + SECURITY_DESCRIPTOR_MIN_LENGTH);\n PACL *ppACL = (PACL *)((PBYTE)ppSID + sizeof(PSID *));\n\n if (*ppSID) {\n FreeSid(*ppSID);\n }\n if (*ppACL) {\n LocalFree(*ppACL);\n }\n free(m_winPSecurityDescriptor);\n}\n#endif\n\nvoid DestroyDebugUtilsMessengerEXT(VkInstance instance,\n VkDebugUtilsMessengerEXT debugMessenger,\n const VkAllocationCallbacks *pAllocator)\n{\n auto func = (PFN_vkDestroyDebugUtilsMessengerEXT)vkGetInstanceProcAddr(instance, \"vkDestroyDebugUtilsMessengerEXT\");\n if (func != nullptr) {\n func(instance, debugMessenger, pAllocator);\n }\n}\n\nstruct QueueFamilyIndices\n{\n int graphicsFamily = -1;\n int presentFamily = -1;\n\n bool isComplete() { return graphicsFamily >= 0 && presentFamily >= 0; }\n};\n\nstruct SwapChainSupportDetails\n{\n VkSurfaceCapabilitiesKHR capabilities;\n std::vector formats;\n std::vector presentModes;\n};\n\ntypedef float vec2[2];\n\nstruct Vertex\n{\n vec4 pos;\n vec3 color;\n vec2 texCoord;\n\n static VkVertexInputBindingDescription getBindingDescription()\n {\n VkVertexInputBindingDescription bindingDescription = {};\n bindingDescription.binding = 0;\n bindingDescription.stride = sizeof(Vertex);\n bindingDescription.inputRate = VK_VERTEX_INPUT_RATE_VERTEX;\n\n return bindingDescription;\n }\n\n static std::array getAttributeDescriptions()\n {\n std::array attributeDescriptions = {};\n\n attributeDescriptions[0].binding = 0;\n attributeDescriptions[0].location = 0;\n attributeDescriptions[0].format = VK_FORMAT_R32G32B32A32_SFLOAT;\n attributeDescriptions[0].offset = offsetof(Vertex, pos);\n\n attributeDescriptions[1].binding = 0;\n attributeDescriptions[1].location = 1;\n attributeDescriptions[1].format = VK_FORMAT_R32G32B32_SFLOAT;\n attributeDescriptions[1].offset = offsetof(Vertex, color);\n\n attributeDescriptions[2].binding = 0;\n attributeDescriptions[2].location = 2;\n attributeDescriptions[2].format = VK_FORMAT_R32G32_SFLOAT;\n attributeDescriptions[2].offset = offsetof(Vertex, texCoord);\n\n return attributeDescriptions;\n }\n};\n\nstruct UniformBufferObject\n{\n alignas(16) mat4x4 model;\n alignas(16) mat4x4 view;\n alignas(16) mat4x4 proj;\n};\n\nconst std::vector vertices = {{{-1.0f, -1.0f, 0.0f, 1.0f}, {1.0f, 0.0f, 0.0f}, {0.0f, 0.0f}},\n {{1.0f, -1.0f, 0.0f, 1.0f}, {0.0f, 1.0f, 0.0f}, {1.0f, 0.0f}},\n {{1.0f, 1.0f, 0.0f, 1.0f}, {0.0f, 0.0f, 1.0f}, {1.0f, 1.0f}},\n {{-1.0f, 1.0f, 0.0f, 1.0f}, {1.0f, 1.0f, 1.0f}, {0.0f, 1.0f}}};\n\nconst std::vector indices = {0, 1, 2, 2, 3, 0};\n\n// convert floating point rgba color to 32-bit integer\n__device__ unsigned int rgbaFloatToInt(float4 rgba)\n{\n rgba.x = __saturatef(rgba.x); // clamp to [0.0, 1.0]\n rgba.y = __saturatef(rgba.y);\n rgba.z = __saturatef(rgba.z);\n rgba.w = __saturatef(rgba.w);\n return ((unsigned int)(rgba.w * 255.0f) << 24) | ((unsigned int)(rgba.z * 255.0f) << 16)\n | ((unsigned int)(rgba.y * 255.0f) << 8) | ((unsigned int)(rgba.x * 255.0f));\n}\n\n__device__ float4 rgbaIntToFloat(unsigned int c)\n{\n float4 rgba;\n rgba.x = (c & 0xff) * 0.003921568627f; // /255.0f;\n rgba.y = ((c >> 8) & 0xff) * 0.003921568627f; // /255.0f;\n rgba.z = ((c >> 16) & 0xff) * 0.003921568627f; // /255.0f;\n rgba.w = ((c >> 24) & 0xff) * 0.003921568627f; // /255.0f;\n return rgba;\n}\n\nint filter_radius = 14;\nint g_nFilterSign = 1;\n\n// This varies the filter radius, so we can see automatic animation\nvoid varySigma()\n{\n filter_radius += g_nFilterSign;\n\n if (filter_radius > 64) {\n filter_radius = 64; // clamp to 64 and then negate sign\n g_nFilterSign = -1;\n }\n else if (filter_radius < 0) {\n filter_radius = 0;\n g_nFilterSign = 1;\n }\n}\n\n// row pass using texture lookups\n__global__ void d_boxfilter_rgba_x(hipSurfaceObject_t *dstSurfMipMapArray,\n hipTextureObject_t textureMipMapInput,\n size_t baseWidth,\n size_t baseHeight,\n size_t mipLevels,\n int filter_radius)\n{\n float scale = 1.0f / (float)((filter_radius << 1) + 1);\n unsigned int y = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (y < baseHeight) {\n for (uint32_t mipLevelIdx = 0; mipLevelIdx < mipLevels; mipLevelIdx++) {\n uint32_t width = (baseWidth >> mipLevelIdx) ? (baseWidth >> mipLevelIdx) : 1;\n uint32_t height = (baseHeight >> mipLevelIdx) ? (baseHeight >> mipLevelIdx) : 1;\n if (y < height && filter_radius < width) {\n float px = 1.0 / width;\n float py = 1.0 / height;\n float4 t = make_float4(0.0f);\n for (int x = -filter_radius; x <= filter_radius; x++) {\n t += tex2DLod(textureMipMapInput, x * px, y * py, (float)mipLevelIdx);\n }\n\n unsigned int dataB = rgbaFloatToInt(t * scale);\n surf2Dwrite(dataB, dstSurfMipMapArray[mipLevelIdx], 0, y);\n\n for (int x = 1; x < width; x++) {\n t += tex2DLod(textureMipMapInput, (x + filter_radius) * px, y * py, (float)mipLevelIdx);\n t -= tex2DLod(textureMipMapInput, (x - filter_radius - 1) * px, y * py, (float)mipLevelIdx);\n unsigned int dataB = rgbaFloatToInt(t * scale);\n surf2Dwrite(dataB, dstSurfMipMapArray[mipLevelIdx], x * sizeof(uchar4), y);\n }\n }\n }\n }\n}\n\n// column pass using coalesced global memory reads\n__global__ void d_boxfilter_rgba_y(hipSurfaceObject_t *dstSurfMipMapArray,\n hipSurfaceObject_t *srcSurfMipMapArray,\n size_t baseWidth,\n size_t baseHeight,\n size_t mipLevels,\n int filter_radius)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n float scale = 1.0f / (float)((filter_radius << 1) + 1);\n\n for (uint32_t mipLevelIdx = 0; mipLevelIdx < mipLevels; mipLevelIdx++) {\n uint32_t width = (baseWidth >> mipLevelIdx) ? (baseWidth >> mipLevelIdx) : 1;\n uint32_t height = (baseHeight >> mipLevelIdx) ? (baseHeight >> mipLevelIdx) : 1;\n\n if (x < width && height > filter_radius) {\n float4 t;\n // do left edge\n int colInBytes = x * sizeof(uchar4);\n unsigned int pixFirst = surf2Dread(srcSurfMipMapArray[mipLevelIdx], colInBytes, 0);\n t = rgbaIntToFloat(pixFirst) * filter_radius;\n\n for (int y = 0; (y < (filter_radius + 1)) && (y < height); y++) {\n unsigned int pix = surf2Dread(srcSurfMipMapArray[mipLevelIdx], colInBytes, y);\n t += rgbaIntToFloat(pix);\n }\n\n unsigned int dataB = rgbaFloatToInt(t * scale);\n surf2Dwrite(dataB, dstSurfMipMapArray[mipLevelIdx], colInBytes, 0);\n\n for (int y = 1; (y < filter_radius + 1) && ((y + filter_radius) < height); y++) {\n unsigned int pix =\n surf2Dread(srcSurfMipMapArray[mipLevelIdx], colInBytes, y + filter_radius);\n t += rgbaIntToFloat(pix);\n t -= rgbaIntToFloat(pixFirst);\n\n dataB = rgbaFloatToInt(t * scale);\n surf2Dwrite(dataB, dstSurfMipMapArray[mipLevelIdx], colInBytes, y);\n }\n\n // main loop\n for (int y = (filter_radius + 1); y < (height - filter_radius); y++) {\n unsigned int pix =\n surf2Dread(srcSurfMipMapArray[mipLevelIdx], colInBytes, y + filter_radius);\n t += rgbaIntToFloat(pix);\n\n pix = surf2Dread(srcSurfMipMapArray[mipLevelIdx], colInBytes, y - filter_radius - 1);\n t -= rgbaIntToFloat(pix);\n\n dataB = rgbaFloatToInt(t * scale);\n surf2Dwrite(dataB, dstSurfMipMapArray[mipLevelIdx], colInBytes, y);\n }\n\n // do right edge\n unsigned int pixLast = surf2Dread(srcSurfMipMapArray[mipLevelIdx], colInBytes, height - 1);\n for (int y = height - filter_radius; (y < height) && ((y - filter_radius - 1) > 1); y++) {\n t += rgbaIntToFloat(pixLast);\n unsigned int pix =\n surf2Dread(srcSurfMipMapArray[mipLevelIdx], colInBytes, y - filter_radius - 1);\n t -= rgbaIntToFloat(pix);\n dataB = rgbaFloatToInt(t * scale);\n surf2Dwrite(dataB, dstSurfMipMapArray[mipLevelIdx], colInBytes, y);\n }\n }\n }\n}\n\nclass vulkanImageCUDA\n{\npublic:\n void loadImageData(const std::string &filename)\n {\n // load image (needed so we can get the width and height before we create\n // the window\n char *image_path = sdkFindFilePath(filename.c_str(), execution_path.c_str());\n\n if (image_path == 0) {\n printf(\"Error finding image file '%s'\\n\", filename.c_str());\n exit(EXIT_FAILURE);\n }\n\n sdkLoadPPM4(image_path, (unsigned char **)&image_data, &imageWidth, &imageHeight);\n\n if (!image_data) {\n printf(\"Error opening file '%s'\\n\", image_path);\n exit(EXIT_FAILURE);\n }\n\n printf(\"Loaded '%s', %d x %d pixels\\n\", image_path, imageWidth, imageHeight);\n }\n\n void run()\n {\n initWindow();\n initVulkan();\n initCuda();\n mainLoop();\n cleanup();\n }\n\nprivate:\n GLFWwindow *window;\n\n VkInstance instance;\n VkDebugUtilsMessengerEXT debugMessenger;\n VkSurfaceKHR surface;\n\n VkPhysicalDevice physicalDevice = VK_NULL_HANDLE;\n VkDevice device;\n uint8_t vkDeviceUUID[VK_UUID_SIZE];\n\n VkQueue graphicsQueue;\n VkQueue presentQueue;\n\n VkSwapchainKHR swapChain;\n std::vector swapChainImages;\n VkFormat swapChainImageFormat;\n VkExtent2D swapChainExtent;\n std::vector swapChainImageViews;\n std::vector swapChainFramebuffers;\n\n VkRenderPass renderPass;\n VkDescriptorSetLayout descriptorSetLayout;\n VkPipelineLayout pipelineLayout;\n VkPipeline graphicsPipeline;\n\n VkCommandPool commandPool;\n\n VkImage textureImage;\n VkDeviceMemory textureImageMemory;\n VkImageView textureImageView;\n VkSampler textureSampler;\n\n VkBuffer vertexBuffer;\n VkDeviceMemory vertexBufferMemory;\n VkBuffer indexBuffer;\n VkDeviceMemory indexBufferMemory;\n\n std::vector uniformBuffers;\n std::vector uniformBuffersMemory;\n\n VkDescriptorPool descriptorPool;\n std::vector descriptorSets;\n\n std::vector commandBuffers;\n\n std::vector imageAvailableSemaphores;\n std::vector renderFinishedSemaphores;\n VkSemaphore cudaUpdateVkSemaphore, vkUpdateCudaSemaphore;\n std::vector inFlightFences;\n\n size_t currentFrame = 0;\n\n bool framebufferResized = false;\n\n#ifdef _WIN64\n PFN_vkGetMemoryWin32HandleKHR fpGetMemoryWin32HandleKHR;\n PFN_vkGetSemaphoreWin32HandleKHR fpGetSemaphoreWin32HandleKHR;\n#else\n PFN_vkGetMemoryFdKHR fpGetMemoryFdKHR = NULL;\n PFN_vkGetSemaphoreFdKHR fpGetSemaphoreFdKHR = NULL;\n#endif\n\n PFN_vkGetPhysicalDeviceProperties2 fpGetPhysicalDeviceProperties2;\n\n unsigned int *image_data = NULL;\n unsigned int imageWidth, imageHeight;\n unsigned int mipLevels = 1;\n size_t totalImageMemSize;\n\n // CUDA objects\n hipExternalMemory_t cudaExtMemImageBuffer;\n hipMipmappedArray_t cudaMipmappedImageArray, cudaMipmappedImageArrayTemp, cudaMipmappedImageArrayOrig;\n std::vector surfaceObjectList, surfaceObjectListTemp;\n hipSurfaceObject_t *d_surfaceObjectList, *d_surfaceObjectListTemp;\n hipTextureObject_t textureObjMipMapInput;\n\n hipExternalSemaphore_t cudaExtCudaUpdateVkSemaphore;\n hipExternalSemaphore_t cudaExtVkUpdateCudaSemaphore;\n hipStream_t streamToRun;\n\n void initWindow()\n {\n glfwInit();\n\n glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);\n\n window = glfwCreateWindow(WIDTH, HEIGHT, \"Vulkan Image CUDA Box Filter\", nullptr, nullptr);\n glfwSetWindowUserPointer(window, this);\n glfwSetFramebufferSizeCallback(window, framebufferResizeCallback);\n }\n\n static void framebufferResizeCallback(GLFWwindow *window, int width, int height)\n {\n auto app = reinterpret_cast(glfwGetWindowUserPointer(window));\n app->framebufferResized = true;\n }\n\n void initVulkan()\n {\n createInstance();\n setupDebugMessenger();\n createSurface();\n pickPhysicalDevice();\n createLogicalDevice();\n getKhrExtensionsFn();\n createSwapChain();\n createImageViews();\n createRenderPass();\n createDescriptorSetLayout();\n createGraphicsPipeline();\n createFramebuffers();\n createCommandPool();\n createTextureImage();\n createTextureImageView();\n createTextureSampler();\n createVertexBuffer();\n createIndexBuffer();\n createUniformBuffers();\n createDescriptorPool();\n createDescriptorSets();\n createCommandBuffers();\n createSyncObjects();\n createSyncObjectsExt();\n }\n\n void initCuda()\n {\n setCudaVkDevice();\n checkCudaErrors(hipStreamCreate(&streamToRun));\n cudaVkImportImageMem();\n cudaVkImportSemaphore();\n }\n\n void mainLoop()\n {\n updateUniformBuffer();\n while (!glfwWindowShouldClose(window)) {\n glfwPollEvents();\n drawFrame();\n }\n\n vkDeviceWaitIdle(device);\n }\n\n void cleanupSwapChain()\n {\n for (auto framebuffer : swapChainFramebuffers) {\n vkDestroyFramebuffer(device, framebuffer, nullptr);\n }\n\n vkFreeCommandBuffers(device, commandPool, static_cast(commandBuffers.size()), commandBuffers.data());\n\n vkDestroyPipeline(device, graphicsPipeline, nullptr);\n vkDestroyPipelineLayout(device, pipelineLayout, nullptr);\n vkDestroyRenderPass(device, renderPass, nullptr);\n\n for (auto imageView : swapChainImageViews) {\n vkDestroyImageView(device, imageView, nullptr);\n }\n\n vkDestroySwapchainKHR(device, swapChain, nullptr);\n\n for (size_t i = 0; i < swapChainImages.size(); i++) {\n vkDestroyBuffer(device, uniformBuffers[i], nullptr);\n vkFreeMemory(device, uniformBuffersMemory[i], nullptr);\n }\n\n vkDestroyDescriptorPool(device, descriptorPool, nullptr);\n }\n\n void cleanup()\n {\n cleanupSwapChain();\n\n vkDestroySampler(device, textureSampler, nullptr);\n vkDestroyImageView(device, textureImageView, nullptr);\n\n for (int i = 0; i < mipLevels; i++) {\n checkCudaErrors(hipDestroySurfaceObject(surfaceObjectList[i]));\n checkCudaErrors(hipDestroySurfaceObject(surfaceObjectListTemp[i]));\n }\n\n checkCudaErrors(hipFree(d_surfaceObjectList));\n checkCudaErrors(hipFree(d_surfaceObjectListTemp));\n checkCudaErrors(hipFreeMipmappedArray(cudaMipmappedImageArrayTemp));\n checkCudaErrors(hipFreeMipmappedArray(cudaMipmappedImageArrayOrig));\n checkCudaErrors(hipFreeMipmappedArray(cudaMipmappedImageArray));\n checkCudaErrors(hipDestroyTextureObject(textureObjMipMapInput));\n checkCudaErrors(hipDestroyExternalMemory(cudaExtMemImageBuffer));\n checkCudaErrors(hipDestroyExternalSemaphore(cudaExtCudaUpdateVkSemaphore));\n checkCudaErrors(hipDestroyExternalSemaphore(cudaExtVkUpdateCudaSemaphore));\n\n vkDestroyImage(device, textureImage, nullptr);\n vkFreeMemory(device, textureImageMemory, nullptr);\n\n vkDestroyDescriptorSetLayout(device, descriptorSetLayout, nullptr);\n\n vkDestroyBuffer(device, indexBuffer, nullptr);\n vkFreeMemory(device, indexBufferMemory, nullptr);\n\n vkDestroyBuffer(device, vertexBuffer, nullptr);\n vkFreeMemory(device, vertexBufferMemory, nullptr);\n\n vkDestroySemaphore(device, cudaUpdateVkSemaphore, nullptr);\n vkDestroySemaphore(device, vkUpdateCudaSemaphore, nullptr);\n\n for (size_t i = 0; i < MAX_FRAMES; i++) {\n vkDestroySemaphore(device, renderFinishedSemaphores[i], nullptr);\n vkDestroySemaphore(device, imageAvailableSemaphores[i], nullptr);\n vkDestroyFence(device, inFlightFences[i], nullptr);\n }\n\n vkDestroyCommandPool(device, commandPool, nullptr);\n\n vkDestroyDevice(device, nullptr);\n\n if (enableValidationLayers) {\n DestroyDebugUtilsMessengerEXT(instance, debugMessenger, nullptr);\n }\n\n vkDestroySurfaceKHR(instance, surface, nullptr);\n vkDestroyInstance(instance, nullptr);\n\n glfwDestroyWindow(window);\n\n glfwTerminate();\n }\n\n void recreateSwapChain()\n {\n int width = 0, height = 0;\n while (width == 0 || height == 0) {\n glfwGetFramebufferSize(window, &width, &height);\n glfwWaitEvents();\n }\n\n vkDeviceWaitIdle(device);\n\n cleanupSwapChain();\n\n createSwapChain();\n createImageViews();\n createRenderPass();\n createGraphicsPipeline();\n createFramebuffers();\n createUniformBuffers();\n createDescriptorPool();\n createDescriptorSets();\n createCommandBuffers();\n }\n\n void createInstance()\n {\n if (enableValidationLayers && !checkValidationLayerSupport()) {\n throw std::runtime_error(\"validation layers requested, but not available!\");\n }\n\n VkApplicationInfo appInfo = {};\n appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;\n appInfo.pApplicationName = \"Vulkan Image CUDA Interop\";\n appInfo.applicationVersion = VK_MAKE_VERSION(1, 0, 0);\n appInfo.pEngineName = \"No Engine\";\n appInfo.engineVersion = VK_MAKE_VERSION(1, 0, 0);\n appInfo.apiVersion = VK_API_VERSION_1_1;\n\n VkInstanceCreateInfo createInfo = {};\n createInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;\n createInfo.pApplicationInfo = &appInfo;\n\n auto extensions = getRequiredExtensions();\n createInfo.enabledExtensionCount = static_cast(extensions.size());\n createInfo.ppEnabledExtensionNames = extensions.data();\n\n VkDebugUtilsMessengerCreateInfoEXT debugCreateInfo;\n if (enableValidationLayers) {\n createInfo.enabledLayerCount = static_cast(validationLayers.size());\n createInfo.ppEnabledLayerNames = validationLayers.data();\n\n populateDebugMessengerCreateInfo(debugCreateInfo);\n createInfo.pNext = (VkDebugUtilsMessengerCreateInfoEXT *)&debugCreateInfo;\n }\n else {\n createInfo.enabledLayerCount = 0;\n\n createInfo.pNext = nullptr;\n }\n\n if (vkCreateInstance(&createInfo, nullptr, &instance) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create instance!\");\n }\n\n fpGetPhysicalDeviceProperties2 =\n (PFN_vkGetPhysicalDeviceProperties2)vkGetInstanceProcAddr(instance, \"vkGetPhysicalDeviceProperties2\");\n if (fpGetPhysicalDeviceProperties2 == NULL) {\n throw std::runtime_error(\"Vulkan: Proc address for \\\"vkGetPhysicalDeviceProperties2KHR\\\" not \"\n \"found.\\n\");\n }\n\n#ifdef _WIN64\n fpGetMemoryWin32HandleKHR =\n (PFN_vkGetMemoryWin32HandleKHR)vkGetInstanceProcAddr(instance, \"vkGetMemoryWin32HandleKHR\");\n if (fpGetMemoryWin32HandleKHR == NULL) {\n throw std::runtime_error(\"Vulkan: Proc address for \\\"vkGetMemoryWin32HandleKHR\\\" not \"\n \"found.\\n\");\n }\n#else\n fpGetMemoryFdKHR = (PFN_vkGetMemoryFdKHR)vkGetInstanceProcAddr(instance, \"vkGetMemoryFdKHR\");\n if (fpGetMemoryFdKHR == NULL) {\n throw std::runtime_error(\"Vulkan: Proc address for \\\"vkGetMemoryFdKHR\\\" not found.\\n\");\n }\n else {\n std::cout << \"Vulkan proc address for vkGetMemoryFdKHR - \" << fpGetMemoryFdKHR << std::endl;\n }\n#endif\n }\n\n void populateDebugMessengerCreateInfo(VkDebugUtilsMessengerCreateInfoEXT &createInfo)\n {\n createInfo = {};\n createInfo.sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_MESSENGER_CREATE_INFO_EXT;\n createInfo.messageSeverity = VK_DEBUG_UTILS_MESSAGE_SEVERITY_VERBOSE_BIT_EXT\n | VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT\n | VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT;\n createInfo.messageType = VK_DEBUG_UTILS_MESSAGE_TYPE_GENERAL_BIT_EXT\n | VK_DEBUG_UTILS_MESSAGE_TYPE_VALIDATION_BIT_EXT\n | VK_DEBUG_UTILS_MESSAGE_TYPE_PERFORMANCE_BIT_EXT;\n createInfo.pfnUserCallback = debugCallback;\n }\n\n void setupDebugMessenger()\n {\n if (!enableValidationLayers)\n return;\n\n VkDebugUtilsMessengerCreateInfoEXT createInfo;\n populateDebugMessengerCreateInfo(createInfo);\n\n if (CreateDebugUtilsMessengerEXT(instance, &createInfo, nullptr, &debugMessenger) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to set up debug messenger!\");\n }\n }\n\n void createSurface()\n {\n if (glfwCreateWindowSurface(instance, window, nullptr, &surface) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create window surface!\");\n }\n }\n\n void pickPhysicalDevice()\n {\n uint32_t deviceCount = 0;\n vkEnumeratePhysicalDevices(instance, &deviceCount, nullptr);\n\n if (deviceCount == 0) {\n throw std::runtime_error(\"failed to find GPUs with Vulkan support!\");\n }\n\n std::vector devices(deviceCount);\n vkEnumeratePhysicalDevices(instance, &deviceCount, devices.data());\n\n for (const auto &device : devices) {\n if (isDeviceSuitable(device)) {\n physicalDevice = device;\n break;\n }\n }\n\n if (physicalDevice == VK_NULL_HANDLE) {\n throw std::runtime_error(\"failed to find a suitable GPU!\");\n }\n\n std::cout << \"Selected physical device = \" << physicalDevice << std::endl;\n\n VkPhysicalDeviceIDProperties vkPhysicalDeviceIDProperties = {};\n vkPhysicalDeviceIDProperties.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ID_PROPERTIES;\n vkPhysicalDeviceIDProperties.pNext = NULL;\n\n VkPhysicalDeviceProperties2 vkPhysicalDeviceProperties2 = {};\n vkPhysicalDeviceProperties2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2;\n vkPhysicalDeviceProperties2.pNext = &vkPhysicalDeviceIDProperties;\n\n fpGetPhysicalDeviceProperties2(physicalDevice, &vkPhysicalDeviceProperties2);\n\n memcpy(vkDeviceUUID, vkPhysicalDeviceIDProperties.deviceUUID, sizeof(vkDeviceUUID));\n }\n\n void getKhrExtensionsFn()\n {\n#ifdef _WIN64\n\n fpGetSemaphoreWin32HandleKHR =\n (PFN_vkGetSemaphoreWin32HandleKHR)vkGetDeviceProcAddr(device, \"vkGetSemaphoreWin32HandleKHR\");\n if (fpGetSemaphoreWin32HandleKHR == NULL) {\n throw std::runtime_error(\"Vulkan: Proc address for \\\"vkGetSemaphoreWin32HandleKHR\\\" not \"\n \"found.\\n\");\n }\n#else\n fpGetSemaphoreFdKHR = (PFN_vkGetSemaphoreFdKHR)vkGetDeviceProcAddr(device, \"vkGetSemaphoreFdKHR\");\n if (fpGetSemaphoreFdKHR == NULL) {\n throw std::runtime_error(\"Vulkan: Proc address for \\\"vkGetSemaphoreFdKHR\\\" not found.\\n\");\n }\n#endif\n }\n\n int setCudaVkDevice()\n {\n int current_device = 0;\n int device_count = 0;\n int devices_prohibited = 0;\n\n hipDeviceProp_t deviceProp;\n int computeMode;\n checkCudaErrors(hipGetDeviceCount(&device_count));\n\n if (device_count == 0) {\n fprintf(stderr, \"CUDA error: no devices supporting CUDA.\\n\");\n exit(EXIT_FAILURE);\n }\n\n // Find the GPU which is selected by Vulkan\n while (current_device < device_count) {\n hipGetDeviceProperties(&deviceProp, current_device);\n checkCudaErrors(hipDeviceGetAttribute(&computeMode, hipDeviceAttributeComputeMode, current_device));\n if ((computeMode != hipComputeModeProhibited)) {\n // Compare the cuda device UUID with vulkan UUID\n int ret = memcmp(&deviceProp.uuid, &vkDeviceUUID, VK_UUID_SIZE);\n if (ret == 0) {\n checkCudaErrors(hipSetDevice(current_device));\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, current_device));\n printf(\"GPU Device %d: \\\"%s\\\" with compute capability %d.%d\\n\\n\",\n current_device,\n deviceProp.name,\n deviceProp.major,\n deviceProp.minor);\n\n return current_device;\n }\n }\n else {\n devices_prohibited++;\n }\n\n current_device++;\n }\n\n if (devices_prohibited == device_count) {\n fprintf(stderr,\n \"CUDA error:\"\n \" No Vulkan-CUDA Interop capable GPU found.\\n\");\n exit(EXIT_FAILURE);\n }\n\n return -1;\n }\n\n void createLogicalDevice()\n {\n QueueFamilyIndices indices = findQueueFamilies(physicalDevice);\n\n std::vector queueCreateInfos;\n std::set uniqueQueueFamilies = {indices.graphicsFamily, indices.presentFamily};\n\n float queuePriority = 1.0f;\n for (int queueFamily : uniqueQueueFamilies) {\n VkDeviceQueueCreateInfo queueCreateInfo = {};\n queueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;\n queueCreateInfo.queueFamilyIndex = queueFamily;\n queueCreateInfo.queueCount = 1;\n queueCreateInfo.pQueuePriorities = &queuePriority;\n queueCreateInfos.push_back(queueCreateInfo);\n }\n\n VkPhysicalDeviceFeatures deviceFeatures = {};\n deviceFeatures.samplerAnisotropy = VK_TRUE;\n\n VkDeviceCreateInfo createInfo = {};\n createInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;\n\n createInfo.pQueueCreateInfos = queueCreateInfos.data();\n createInfo.queueCreateInfoCount = queueCreateInfos.size();\n\n createInfo.pEnabledFeatures = &deviceFeatures;\n std::vector enabledExtensionNameList;\n\n for (int i = 0; i < deviceExtensions.size(); i++) {\n enabledExtensionNameList.push_back(deviceExtensions[i]);\n }\n if (enableValidationLayers) {\n createInfo.enabledLayerCount = static_cast(validationLayers.size());\n createInfo.ppEnabledLayerNames = validationLayers.data();\n }\n else {\n createInfo.enabledLayerCount = 0;\n }\n createInfo.enabledExtensionCount = static_cast(enabledExtensionNameList.size());\n createInfo.ppEnabledExtensionNames = enabledExtensionNameList.data();\n\n if (vkCreateDevice(physicalDevice, &createInfo, nullptr, &device) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create logical device!\");\n }\n vkGetDeviceQueue(device, indices.graphicsFamily, 0, &graphicsQueue);\n vkGetDeviceQueue(device, indices.presentFamily, 0, &presentQueue);\n }\n\n void createSwapChain()\n {\n SwapChainSupportDetails swapChainSupport = querySwapChainSupport(physicalDevice);\n\n VkSurfaceFormatKHR surfaceFormat = chooseSwapSurfaceFormat(swapChainSupport.formats);\n VkPresentModeKHR presentMode = chooseSwapPresentMode(swapChainSupport.presentModes);\n VkExtent2D extent = chooseSwapExtent(swapChainSupport.capabilities);\n\n uint32_t imageCount = swapChainSupport.capabilities.minImageCount + 1;\n if (swapChainSupport.capabilities.maxImageCount > 0\n && imageCount > swapChainSupport.capabilities.maxImageCount) {\n imageCount = swapChainSupport.capabilities.maxImageCount;\n }\n\n VkSwapchainCreateInfoKHR createInfo = {};\n createInfo.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;\n createInfo.surface = surface;\n\n createInfo.minImageCount = imageCount;\n createInfo.imageFormat = surfaceFormat.format;\n createInfo.imageColorSpace = surfaceFormat.colorSpace;\n createInfo.imageExtent = extent;\n createInfo.imageArrayLayers = 1;\n createInfo.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;\n\n QueueFamilyIndices indices = findQueueFamilies(physicalDevice);\n uint32_t queueFamilyIndices[] = {(uint32_t)indices.graphicsFamily, (uint32_t)indices.presentFamily};\n\n if (indices.graphicsFamily != indices.presentFamily) {\n createInfo.imageSharingMode = VK_SHARING_MODE_CONCURRENT;\n createInfo.queueFamilyIndexCount = 2;\n createInfo.pQueueFamilyIndices = queueFamilyIndices;\n }\n else {\n createInfo.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;\n }\n\n createInfo.preTransform = swapChainSupport.capabilities.currentTransform;\n createInfo.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;\n createInfo.presentMode = presentMode;\n createInfo.clipped = VK_TRUE;\n\n if (vkCreateSwapchainKHR(device, &createInfo, nullptr, &swapChain) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create swap chain!\");\n }\n\n vkGetSwapchainImagesKHR(device, swapChain, &imageCount, nullptr);\n swapChainImages.resize(imageCount);\n vkGetSwapchainImagesKHR(device, swapChain, &imageCount, swapChainImages.data());\n\n swapChainImageFormat = surfaceFormat.format;\n swapChainExtent = extent;\n }\n\n void createImageViews()\n {\n swapChainImageViews.resize(swapChainImages.size());\n\n for (size_t i = 0; i < swapChainImages.size(); i++) {\n swapChainImageViews[i] = createImageView(swapChainImages[i], swapChainImageFormat);\n }\n }\n\n void createRenderPass()\n {\n VkAttachmentDescription colorAttachment = {};\n colorAttachment.format = swapChainImageFormat;\n colorAttachment.samples = VK_SAMPLE_COUNT_1_BIT;\n colorAttachment.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;\n colorAttachment.storeOp = VK_ATTACHMENT_STORE_OP_STORE;\n colorAttachment.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;\n colorAttachment.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;\n colorAttachment.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;\n colorAttachment.finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;\n\n VkAttachmentReference colorAttachmentRef = {};\n colorAttachmentRef.attachment = 0;\n colorAttachmentRef.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;\n\n VkSubpassDescription subpass = {};\n subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;\n subpass.colorAttachmentCount = 1;\n subpass.pColorAttachments = &colorAttachmentRef;\n\n VkSubpassDependency dependency = {};\n dependency.srcSubpass = VK_SUBPASS_EXTERNAL;\n dependency.dstSubpass = 0;\n dependency.srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;\n dependency.srcAccessMask = 0;\n dependency.dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;\n dependency.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_READ_BIT | VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;\n\n VkRenderPassCreateInfo renderPassInfo = {};\n renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;\n renderPassInfo.attachmentCount = 1;\n renderPassInfo.pAttachments = &colorAttachment;\n renderPassInfo.subpassCount = 1;\n renderPassInfo.pSubpasses = &subpass;\n renderPassInfo.dependencyCount = 1;\n renderPassInfo.pDependencies = &dependency;\n\n if (vkCreateRenderPass(device, &renderPassInfo, nullptr, &renderPass) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create render pass!\");\n }\n }\n\n void createDescriptorSetLayout()\n {\n VkDescriptorSetLayoutBinding uboLayoutBinding = {};\n uboLayoutBinding.binding = 0;\n uboLayoutBinding.descriptorCount = 1;\n uboLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;\n uboLayoutBinding.pImmutableSamplers = nullptr;\n uboLayoutBinding.stageFlags = VK_SHADER_STAGE_VERTEX_BIT;\n\n VkDescriptorSetLayoutBinding samplerLayoutBinding = {};\n samplerLayoutBinding.binding = 1;\n samplerLayoutBinding.descriptorCount = 1;\n samplerLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;\n samplerLayoutBinding.pImmutableSamplers = nullptr;\n samplerLayoutBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;\n\n std::array bindings = {uboLayoutBinding, samplerLayoutBinding};\n VkDescriptorSetLayoutCreateInfo layoutInfo = {};\n layoutInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;\n layoutInfo.bindingCount = static_cast(bindings.size());\n layoutInfo.pBindings = bindings.data();\n\n if (vkCreateDescriptorSetLayout(device, &layoutInfo, nullptr, &descriptorSetLayout) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create descriptor set layout!\");\n }\n }\n\n void createGraphicsPipeline()\n {\n auto vertShaderCode = readFile(\"vert.spv\");\n auto fragShaderCode = readFile(\"frag.spv\");\n\n VkShaderModule vertShaderModule = createShaderModule(vertShaderCode);\n VkShaderModule fragShaderModule = createShaderModule(fragShaderCode);\n\n VkPipelineShaderStageCreateInfo vertShaderStageInfo = {};\n vertShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;\n vertShaderStageInfo.stage = VK_SHADER_STAGE_VERTEX_BIT;\n vertShaderStageInfo.module = vertShaderModule;\n vertShaderStageInfo.pName = \"main\";\n\n VkPipelineShaderStageCreateInfo fragShaderStageInfo = {};\n fragShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;\n fragShaderStageInfo.stage = VK_SHADER_STAGE_FRAGMENT_BIT;\n fragShaderStageInfo.module = fragShaderModule;\n fragShaderStageInfo.pName = \"main\";\n\n VkPipelineShaderStageCreateInfo shaderStages[] = {vertShaderStageInfo, fragShaderStageInfo};\n\n VkPipelineVertexInputStateCreateInfo vertexInputInfo = {};\n vertexInputInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;\n\n auto bindingDescription = Vertex::getBindingDescription();\n auto attributeDescriptions = Vertex::getAttributeDescriptions();\n\n vertexInputInfo.vertexBindingDescriptionCount = 1;\n vertexInputInfo.vertexAttributeDescriptionCount = static_cast(attributeDescriptions.size());\n vertexInputInfo.pVertexBindingDescriptions = &bindingDescription;\n vertexInputInfo.pVertexAttributeDescriptions = attributeDescriptions.data();\n\n VkPipelineInputAssemblyStateCreateInfo inputAssembly = {};\n inputAssembly.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;\n inputAssembly.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;\n inputAssembly.primitiveRestartEnable = VK_FALSE;\n\n VkViewport viewport = {};\n viewport.x = 0.0f;\n viewport.y = 0.0f;\n viewport.width = (float)swapChainExtent.width;\n viewport.height = (float)swapChainExtent.height;\n viewport.minDepth = 0.0f;\n viewport.maxDepth = 1.0f;\n\n VkRect2D scissor = {};\n scissor.offset = {0, 0};\n scissor.extent = swapChainExtent;\n\n VkPipelineViewportStateCreateInfo viewportState = {};\n viewportState.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;\n viewportState.viewportCount = 1;\n viewportState.pViewports = &viewport;\n viewportState.scissorCount = 1;\n viewportState.pScissors = &scissor;\n\n VkPipelineRasterizationStateCreateInfo rasterizer = {};\n rasterizer.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;\n rasterizer.depthClampEnable = VK_FALSE;\n rasterizer.rasterizerDiscardEnable = VK_FALSE;\n rasterizer.polygonMode = VK_POLYGON_MODE_FILL;\n rasterizer.lineWidth = 1.0f;\n rasterizer.cullMode = VK_CULL_MODE_BACK_BIT;\n rasterizer.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;\n rasterizer.depthBiasEnable = VK_FALSE;\n\n VkPipelineMultisampleStateCreateInfo multisampling = {};\n multisampling.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;\n multisampling.sampleShadingEnable = VK_FALSE;\n multisampling.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;\n\n VkPipelineColorBlendAttachmentState colorBlendAttachment = {};\n colorBlendAttachment.colorWriteMask =\n VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;\n colorBlendAttachment.blendEnable = VK_FALSE;\n\n VkPipelineColorBlendStateCreateInfo colorBlending = {};\n colorBlending.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;\n colorBlending.logicOpEnable = VK_FALSE;\n colorBlending.logicOp = VK_LOGIC_OP_COPY;\n colorBlending.attachmentCount = 1;\n colorBlending.pAttachments = &colorBlendAttachment;\n colorBlending.blendConstants[0] = 0.0f;\n colorBlending.blendConstants[1] = 0.0f;\n colorBlending.blendConstants[2] = 0.0f;\n colorBlending.blendConstants[3] = 0.0f;\n\n VkPipelineLayoutCreateInfo pipelineLayoutInfo = {};\n pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;\n pipelineLayoutInfo.setLayoutCount = 1;\n pipelineLayoutInfo.pSetLayouts = &descriptorSetLayout;\n\n if (vkCreatePipelineLayout(device, &pipelineLayoutInfo, nullptr, &pipelineLayout) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create pipeline layout!\");\n }\n\n VkGraphicsPipelineCreateInfo pipelineInfo = {};\n pipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;\n pipelineInfo.stageCount = 2;\n pipelineInfo.pStages = shaderStages;\n pipelineInfo.pVertexInputState = &vertexInputInfo;\n pipelineInfo.pInputAssemblyState = &inputAssembly;\n pipelineInfo.pViewportState = &viewportState;\n pipelineInfo.pRasterizationState = &rasterizer;\n pipelineInfo.pMultisampleState = &multisampling;\n pipelineInfo.pColorBlendState = &colorBlending;\n pipelineInfo.layout = pipelineLayout;\n pipelineInfo.renderPass = renderPass;\n pipelineInfo.subpass = 0;\n pipelineInfo.basePipelineHandle = VK_NULL_HANDLE;\n\n if (vkCreateGraphicsPipelines(device, VK_NULL_HANDLE, 1, &pipelineInfo, nullptr, &graphicsPipeline)\n != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create graphics pipeline!\");\n }\n\n vkDestroyShaderModule(device, fragShaderModule, nullptr);\n vkDestroyShaderModule(device, vertShaderModule, nullptr);\n }\n\n void createFramebuffers()\n {\n swapChainFramebuffers.resize(swapChainImageViews.size());\n\n for (size_t i = 0; i < swapChainImageViews.size(); i++) {\n VkImageView attachments[] = {swapChainImageViews[i]};\n\n VkFramebufferCreateInfo framebufferInfo = {};\n framebufferInfo.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;\n framebufferInfo.renderPass = renderPass;\n framebufferInfo.attachmentCount = 1;\n framebufferInfo.pAttachments = attachments;\n framebufferInfo.width = swapChainExtent.width;\n framebufferInfo.height = swapChainExtent.height;\n framebufferInfo.layers = 1;\n\n if (vkCreateFramebuffer(device, &framebufferInfo, nullptr, &swapChainFramebuffers[i]) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create framebuffer!\");\n }\n }\n }\n\n void createCommandPool()\n {\n QueueFamilyIndices queueFamilyIndices = findQueueFamilies(physicalDevice);\n\n VkCommandPoolCreateInfo poolInfo = {};\n poolInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;\n poolInfo.queueFamilyIndex = queueFamilyIndices.graphicsFamily;\n\n if (vkCreateCommandPool(device, &poolInfo, nullptr, &commandPool) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create graphics command pool!\");\n }\n }\n\n void createTextureImage()\n {\n VkDeviceSize imageSize = imageWidth * imageHeight * 4;\n mipLevels = static_cast(std::floor(std::log2(std::max(imageWidth, imageHeight)))) + 1;\n printf(\"mipLevels = %d\\n\", mipLevels);\n\n if (!image_data) {\n throw std::runtime_error(\"failed to load texture image!\");\n }\n\n VkBuffer stagingBuffer;\n VkDeviceMemory stagingBufferMemory;\n createBuffer(imageSize,\n VK_BUFFER_USAGE_TRANSFER_SRC_BIT,\n VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,\n stagingBuffer,\n stagingBufferMemory);\n\n void *data;\n vkMapMemory(device, stagingBufferMemory, 0, imageSize, 0, &data);\n memcpy(data, image_data, static_cast(imageSize));\n vkUnmapMemory(device, stagingBufferMemory);\n\n // VK_FORMAT_R8G8B8A8_UNORM changed to VK_FORMAT_R8G8B8A8_UINT\n createImage(imageWidth,\n imageHeight,\n VK_FORMAT_R8G8B8A8_UNORM,\n VK_IMAGE_TILING_OPTIMAL,\n VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT\n | VK_IMAGE_USAGE_SAMPLED_BIT,\n VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,\n textureImage,\n textureImageMemory);\n\n transitionImageLayout(\n textureImage, VK_FORMAT_R8G8B8A8_UINT, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);\n copyBufferToImage(\n stagingBuffer, textureImage, static_cast(imageWidth), static_cast(imageHeight));\n\n vkDestroyBuffer(device, stagingBuffer, nullptr);\n vkFreeMemory(device, stagingBufferMemory, nullptr);\n\n generateMipmaps(textureImage, VK_FORMAT_R8G8B8A8_UNORM);\n }\n\n void generateMipmaps(VkImage image, VkFormat imageFormat)\n {\n VkFormatProperties formatProperties;\n vkGetPhysicalDeviceFormatProperties(physicalDevice, imageFormat, &formatProperties);\n\n if (!(formatProperties.optimalTilingFeatures & VK_FORMAT_FEATURE_SAMPLED_IMAGE_FILTER_LINEAR_BIT)) {\n throw std::runtime_error(\"texture image format does not support linear blitting!\");\n }\n\n VkCommandBuffer commandBuffer = beginSingleTimeCommands();\n\n VkImageMemoryBarrier barrier = {};\n barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;\n barrier.image = image;\n barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;\n barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;\n barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;\n barrier.subresourceRange.baseArrayLayer = 0;\n barrier.subresourceRange.layerCount = 1;\n barrier.subresourceRange.levelCount = 1;\n\n int32_t mipWidth = imageWidth;\n int32_t mipHeight = imageHeight;\n\n for (uint32_t i = 1; i < mipLevels; i++) {\n barrier.subresourceRange.baseMipLevel = i - 1;\n barrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;\n barrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;\n barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;\n barrier.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT;\n\n vkCmdPipelineBarrier(commandBuffer,\n VK_PIPELINE_STAGE_TRANSFER_BIT,\n VK_PIPELINE_STAGE_TRANSFER_BIT,\n 0,\n 0,\n nullptr,\n 0,\n nullptr,\n 1,\n &barrier);\n\n VkImageBlit blit = {};\n blit.srcOffsets[0] = {0, 0, 0};\n blit.srcOffsets[1] = {mipWidth, mipHeight, 1};\n blit.srcSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;\n blit.srcSubresource.mipLevel = i - 1;\n blit.srcSubresource.baseArrayLayer = 0;\n blit.srcSubresource.layerCount = 1;\n blit.dstOffsets[0] = {0, 0, 0};\n blit.dstOffsets[1] = {mipWidth > 1 ? mipWidth / 2 : 1, mipHeight > 1 ? mipHeight / 2 : 1, 1};\n blit.dstSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;\n blit.dstSubresource.mipLevel = i;\n blit.dstSubresource.baseArrayLayer = 0;\n blit.dstSubresource.layerCount = 1;\n\n vkCmdBlitImage(commandBuffer,\n image,\n VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,\n image,\n VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,\n 1,\n &blit,\n VK_FILTER_LINEAR);\n\n barrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;\n barrier.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;\n barrier.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT;\n barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;\n\n vkCmdPipelineBarrier(commandBuffer,\n VK_PIPELINE_STAGE_TRANSFER_BIT,\n VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,\n 0,\n 0,\n nullptr,\n 0,\n nullptr,\n 1,\n &barrier);\n\n if (mipWidth > 1)\n mipWidth /= 2;\n if (mipHeight > 1)\n mipHeight /= 2;\n }\n\n barrier.subresourceRange.baseMipLevel = mipLevels - 1;\n barrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;\n barrier.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;\n barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;\n barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;\n\n vkCmdPipelineBarrier(commandBuffer,\n VK_PIPELINE_STAGE_TRANSFER_BIT,\n VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,\n 0,\n 0,\n nullptr,\n 0,\n nullptr,\n 1,\n &barrier);\n\n endSingleTimeCommands(commandBuffer);\n }\n\n#ifdef _WIN64 // For windows\n HANDLE getVkImageMemHandle(VkExternalMemoryHandleTypeFlagsKHR externalMemoryHandleType)\n {\n HANDLE handle;\n\n VkMemoryGetWin32HandleInfoKHR vkMemoryGetWin32HandleInfoKHR = {};\n vkMemoryGetWin32HandleInfoKHR.sType = VK_STRUCTURE_TYPE_MEMORY_GET_WIN32_HANDLE_INFO_KHR;\n vkMemoryGetWin32HandleInfoKHR.pNext = NULL;\n vkMemoryGetWin32HandleInfoKHR.memory = textureImageMemory;\n vkMemoryGetWin32HandleInfoKHR.handleType = (VkExternalMemoryHandleTypeFlagBitsKHR)externalMemoryHandleType;\n\n fpGetMemoryWin32HandleKHR(device, &vkMemoryGetWin32HandleInfoKHR, &handle);\n return handle;\n }\n HANDLE getVkSemaphoreHandle(VkExternalSemaphoreHandleTypeFlagBitsKHR externalSemaphoreHandleType,\n VkSemaphore &semVkCuda)\n {\n HANDLE handle;\n\n VkSemaphoreGetWin32HandleInfoKHR vulkanSemaphoreGetWin32HandleInfoKHR = {};\n vulkanSemaphoreGetWin32HandleInfoKHR.sType = VK_STRUCTURE_TYPE_SEMAPHORE_GET_WIN32_HANDLE_INFO_KHR;\n vulkanSemaphoreGetWin32HandleInfoKHR.pNext = NULL;\n vulkanSemaphoreGetWin32HandleInfoKHR.semaphore = semVkCuda;\n vulkanSemaphoreGetWin32HandleInfoKHR.handleType = externalSemaphoreHandleType;\n\n fpGetSemaphoreWin32HandleKHR(device, &vulkanSemaphoreGetWin32HandleInfoKHR, &handle);\n\n return handle;\n }\n#else\n int getVkImageMemHandle(VkExternalMemoryHandleTypeFlagsKHR externalMemoryHandleType)\n {\n if (externalMemoryHandleType == VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT_KHR) {\n int fd;\n\n VkMemoryGetFdInfoKHR vkMemoryGetFdInfoKHR = {};\n vkMemoryGetFdInfoKHR.sType = VK_STRUCTURE_TYPE_MEMORY_GET_FD_INFO_KHR;\n vkMemoryGetFdInfoKHR.pNext = NULL;\n vkMemoryGetFdInfoKHR.memory = textureImageMemory;\n vkMemoryGetFdInfoKHR.handleType = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT_KHR;\n\n fpGetMemoryFdKHR(device, &vkMemoryGetFdInfoKHR, &fd);\n\n return fd;\n }\n return -1;\n }\n\n int getVkSemaphoreHandle(VkExternalSemaphoreHandleTypeFlagBitsKHR externalSemaphoreHandleType,\n VkSemaphore &semVkCuda)\n {\n if (externalSemaphoreHandleType == VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_FD_BIT) {\n int fd;\n\n VkSemaphoreGetFdInfoKHR vulkanSemaphoreGetFdInfoKHR = {};\n vulkanSemaphoreGetFdInfoKHR.sType = VK_STRUCTURE_TYPE_SEMAPHORE_GET_FD_INFO_KHR;\n vulkanSemaphoreGetFdInfoKHR.pNext = NULL;\n vulkanSemaphoreGetFdInfoKHR.semaphore = semVkCuda;\n vulkanSemaphoreGetFdInfoKHR.handleType = VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_FD_BIT_KHR;\n\n fpGetSemaphoreFdKHR(device, &vulkanSemaphoreGetFdInfoKHR, &fd);\n\n return fd;\n }\n return -1;\n }\n#endif\n\n void createTextureImageView() { textureImageView = createImageView(textureImage, VK_FORMAT_R8G8B8A8_UNORM); }\n\n void createTextureSampler()\n {\n VkSamplerCreateInfo samplerInfo = {};\n samplerInfo.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;\n samplerInfo.magFilter = VK_FILTER_LINEAR;\n samplerInfo.minFilter = VK_FILTER_LINEAR;\n samplerInfo.addressModeU = VK_SAMPLER_ADDRESS_MODE_REPEAT;\n samplerInfo.addressModeV = VK_SAMPLER_ADDRESS_MODE_REPEAT;\n samplerInfo.addressModeW = VK_SAMPLER_ADDRESS_MODE_REPEAT;\n samplerInfo.anisotropyEnable = VK_TRUE;\n samplerInfo.maxAnisotropy = 16;\n samplerInfo.borderColor = VK_BORDER_COLOR_INT_OPAQUE_BLACK;\n samplerInfo.unnormalizedCoordinates = VK_FALSE;\n samplerInfo.compareEnable = VK_FALSE;\n samplerInfo.compareOp = VK_COMPARE_OP_ALWAYS;\n samplerInfo.mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR;\n samplerInfo.minLod = 0; // Optional\n samplerInfo.maxLod = static_cast(mipLevels);\n samplerInfo.mipLodBias = 0; // Optional\n\n if (vkCreateSampler(device, &samplerInfo, nullptr, &textureSampler) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create texture sampler!\");\n }\n }\n\n VkImageView createImageView(VkImage image, VkFormat format)\n {\n VkImageViewCreateInfo viewInfo = {};\n viewInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;\n viewInfo.image = image;\n viewInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;\n viewInfo.format = format;\n viewInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;\n viewInfo.subresourceRange.baseMipLevel = 0;\n viewInfo.subresourceRange.levelCount = mipLevels;\n viewInfo.subresourceRange.baseArrayLayer = 0;\n viewInfo.subresourceRange.layerCount = 1;\n\n VkImageView imageView;\n if (vkCreateImageView(device, &viewInfo, nullptr, &imageView) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create texture image view!\");\n }\n\n return imageView;\n }\n\n void createImage(uint32_t width,\n uint32_t height,\n VkFormat format,\n VkImageTiling tiling,\n VkImageUsageFlags usage,\n VkMemoryPropertyFlags properties,\n VkImage &image,\n VkDeviceMemory &imageMemory)\n {\n VkImageCreateInfo imageInfo = {};\n imageInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;\n imageInfo.imageType = VK_IMAGE_TYPE_2D;\n imageInfo.extent.width = width;\n imageInfo.extent.height = height;\n imageInfo.extent.depth = 1;\n imageInfo.mipLevels = mipLevels;\n imageInfo.arrayLayers = 1;\n imageInfo.format = format;\n imageInfo.tiling = tiling;\n imageInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;\n imageInfo.usage = usage;\n imageInfo.samples = VK_SAMPLE_COUNT_1_BIT;\n imageInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;\n\n VkExternalMemoryImageCreateInfo vkExternalMemImageCreateInfo = {};\n vkExternalMemImageCreateInfo.sType = VK_STRUCTURE_TYPE_EXTERNAL_MEMORY_IMAGE_CREATE_INFO;\n vkExternalMemImageCreateInfo.pNext = NULL;\n#ifdef _WIN64\n vkExternalMemImageCreateInfo.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT;\n#else\n vkExternalMemImageCreateInfo.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT_KHR;\n#endif\n\n imageInfo.pNext = &vkExternalMemImageCreateInfo;\n\n if (vkCreateImage(device, &imageInfo, nullptr, &image) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create image!\");\n }\n\n VkMemoryRequirements memRequirements;\n vkGetImageMemoryRequirements(device, image, &memRequirements);\n\n#ifdef _WIN64\n WindowsSecurityAttributes winSecurityAttributes;\n\n VkExportMemoryWin32HandleInfoKHR vulkanExportMemoryWin32HandleInfoKHR = {};\n vulkanExportMemoryWin32HandleInfoKHR.sType = VK_STRUCTURE_TYPE_EXPORT_MEMORY_WIN32_HANDLE_INFO_KHR;\n vulkanExportMemoryWin32HandleInfoKHR.pNext = NULL;\n vulkanExportMemoryWin32HandleInfoKHR.pAttributes = &winSecurityAttributes;\n vulkanExportMemoryWin32HandleInfoKHR.dwAccess = DXGI_SHARED_RESOURCE_READ | DXGI_SHARED_RESOURCE_WRITE;\n vulkanExportMemoryWin32HandleInfoKHR.name = (LPCWSTR)NULL;\n#endif\n VkExportMemoryAllocateInfoKHR vulkanExportMemoryAllocateInfoKHR = {};\n vulkanExportMemoryAllocateInfoKHR.sType = VK_STRUCTURE_TYPE_EXPORT_MEMORY_ALLOCATE_INFO_KHR;\n#ifdef _WIN64\n vulkanExportMemoryAllocateInfoKHR.pNext = IsWindows8OrGreater() ? &vulkanExportMemoryWin32HandleInfoKHR : NULL;\n vulkanExportMemoryAllocateInfoKHR.handleTypes = IsWindows8OrGreater()\n ? VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT\n : VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_KMT_BIT;\n#else\n vulkanExportMemoryAllocateInfoKHR.pNext = NULL;\n vulkanExportMemoryAllocateInfoKHR.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT_KHR;\n#endif\n\n VkMemoryAllocateInfo allocInfo = {};\n allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;\n allocInfo.allocationSize = memRequirements.size;\n allocInfo.pNext = &vulkanExportMemoryAllocateInfoKHR;\n allocInfo.memoryTypeIndex = findMemoryType(memRequirements.memoryTypeBits, properties);\n\n VkMemoryRequirements vkMemoryRequirements = {};\n vkGetImageMemoryRequirements(device, image, &vkMemoryRequirements);\n totalImageMemSize = vkMemoryRequirements.size;\n\n if (vkAllocateMemory(device, &allocInfo, nullptr, &textureImageMemory) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to allocate image memory!\");\n }\n\n vkBindImageMemory(device, image, textureImageMemory, 0);\n }\n\n void cudaVkImportSemaphore()\n {\n hipExternalSemaphoreHandleDesc externalSemaphoreHandleDesc;\n memset(&externalSemaphoreHandleDesc, 0, sizeof(externalSemaphoreHandleDesc));\n#ifdef _WIN64\n externalSemaphoreHandleDesc.type = IsWindows8OrGreater() ? hipExternalSemaphoreHandleTypeOpaqueWin32\n : hipExternalSemaphoreHandleTypeOpaqueWin32Kmt;\n externalSemaphoreHandleDesc.handle.win32.handle =\n getVkSemaphoreHandle(IsWindows8OrGreater() ? VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_WIN32_BIT\n : VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_WIN32_KMT_BIT,\n cudaUpdateVkSemaphore);\n#else\n externalSemaphoreHandleDesc.type = hipExternalSemaphoreHandleTypeOpaqueFd;\n externalSemaphoreHandleDesc.handle.fd =\n getVkSemaphoreHandle(VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_FD_BIT, cudaUpdateVkSemaphore);\n#endif\n externalSemaphoreHandleDesc.flags = 0;\n\n checkCudaErrors(hipImportExternalSemaphore(&cudaExtCudaUpdateVkSemaphore, &externalSemaphoreHandleDesc));\n\n memset(&externalSemaphoreHandleDesc, 0, sizeof(externalSemaphoreHandleDesc));\n#ifdef _WIN64\n externalSemaphoreHandleDesc.type = IsWindows8OrGreater() ? hipExternalSemaphoreHandleTypeOpaqueWin32\n : hipExternalSemaphoreHandleTypeOpaqueWin32Kmt;\n ;\n externalSemaphoreHandleDesc.handle.win32.handle =\n getVkSemaphoreHandle(IsWindows8OrGreater() ? VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_WIN32_BIT\n : VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_WIN32_KMT_BIT,\n vkUpdateCudaSemaphore);\n#else\n externalSemaphoreHandleDesc.type = hipExternalSemaphoreHandleTypeOpaqueFd;\n externalSemaphoreHandleDesc.handle.fd =\n getVkSemaphoreHandle(VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_FD_BIT, vkUpdateCudaSemaphore);\n#endif\n externalSemaphoreHandleDesc.flags = 0;\n checkCudaErrors(hipImportExternalSemaphore(&cudaExtVkUpdateCudaSemaphore, &externalSemaphoreHandleDesc));\n printf(\"CUDA Imported Vulkan semaphore\\n\");\n }\n\n void cudaVkImportImageMem()\n {\n hipExternalMemoryHandleDesc cudaExtMemHandleDesc;\n memset(&cudaExtMemHandleDesc, 0, sizeof(cudaExtMemHandleDesc));\n#ifdef _WIN64\n cudaExtMemHandleDesc.type = IsWindows8OrGreater() ? hipExternalMemoryHandleTypeOpaqueWin32\n : hipExternalMemoryHandleTypeOpaqueWin32Kmt;\n cudaExtMemHandleDesc.handle.win32.handle =\n getVkImageMemHandle(IsWindows8OrGreater() ? VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT\n : VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_KMT_BIT);\n#else\n cudaExtMemHandleDesc.type = hipExternalMemoryHandleTypeOpaqueFd;\n\n cudaExtMemHandleDesc.handle.fd = getVkImageMemHandle(VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT_KHR);\n#endif\n cudaExtMemHandleDesc.size = totalImageMemSize;\n\n checkCudaErrors(hipImportExternalMemory(&cudaExtMemImageBuffer, &cudaExtMemHandleDesc));\n\n cudaExternalMemoryMipmappedArrayDesc externalMemoryMipmappedArrayDesc;\n\n memset(&externalMemoryMipmappedArrayDesc, 0, sizeof(externalMemoryMipmappedArrayDesc));\n\n hipExtent extent = make_hipExtent(imageWidth, imageHeight, 0);\n hipChannelFormatDesc formatDesc;\n formatDesc.x = 8;\n formatDesc.y = 8;\n formatDesc.z = 8;\n formatDesc.w = 8;\n formatDesc.f = hipChannelFormatKindUnsigned;\n\n externalMemoryMipmappedArrayDesc.offset = 0;\n externalMemoryMipmappedArrayDesc.formatDesc = formatDesc;\n externalMemoryMipmappedArrayDesc.extent = extent;\n externalMemoryMipmappedArrayDesc.flags = 0;\n externalMemoryMipmappedArrayDesc.numLevels = mipLevels;\n\n checkCudaErrors(cudaExternalMemoryGetMappedMipmappedArray(\n &cudaMipmappedImageArray, cudaExtMemImageBuffer, &externalMemoryMipmappedArrayDesc));\n\n checkCudaErrors(hipMallocMipmappedArray(&cudaMipmappedImageArrayTemp, &formatDesc, extent, mipLevels));\n checkCudaErrors(hipMallocMipmappedArray(&cudaMipmappedImageArrayOrig, &formatDesc, extent, mipLevels));\n\n for (int mipLevelIdx = 0; mipLevelIdx < mipLevels; mipLevelIdx++) {\n hipArray_t cudaMipLevelArray, cudaMipLevelArrayTemp, cudaMipLevelArrayOrig;\n hipResourceDesc resourceDesc;\n\n checkCudaErrors(hipGetMipmappedArrayLevel(&cudaMipLevelArray, cudaMipmappedImageArray, mipLevelIdx));\n checkCudaErrors(\n hipGetMipmappedArrayLevel(&cudaMipLevelArrayTemp, cudaMipmappedImageArrayTemp, mipLevelIdx));\n checkCudaErrors(\n hipGetMipmappedArrayLevel(&cudaMipLevelArrayOrig, cudaMipmappedImageArrayOrig, mipLevelIdx));\n\n uint32_t width = (imageWidth >> mipLevelIdx) ? (imageWidth >> mipLevelIdx) : 1;\n uint32_t height = (imageHeight >> mipLevelIdx) ? (imageHeight >> mipLevelIdx) : 1;\n checkCudaErrors(hipMemcpy2DArrayToArray(cudaMipLevelArrayOrig,\n 0,\n 0,\n cudaMipLevelArray,\n 0,\n 0,\n width * sizeof(uchar4),\n height,\n hipMemcpyDeviceToDevice));\n\n memset(&resourceDesc, 0, sizeof(resourceDesc));\n resourceDesc.resType = hipResourceTypeArray;\n resourceDesc.res.array.array = cudaMipLevelArray;\n\n hipSurfaceObject_t surfaceObject;\n checkCudaErrors(hipCreateSurfaceObject(&surfaceObject, &resourceDesc));\n\n surfaceObjectList.push_back(surfaceObject);\n\n memset(&resourceDesc, 0, sizeof(resourceDesc));\n resourceDesc.resType = hipResourceTypeArray;\n resourceDesc.res.array.array = cudaMipLevelArrayTemp;\n\n hipSurfaceObject_t surfaceObjectTemp;\n checkCudaErrors(hipCreateSurfaceObject(&surfaceObjectTemp, &resourceDesc));\n surfaceObjectListTemp.push_back(surfaceObjectTemp);\n }\n\n hipResourceDesc resDescr;\n memset(&resDescr, 0, sizeof(hipResourceDesc));\n\n resDescr.resType = hipResourceTypeMipmappedArray;\n resDescr.res.mipmap.mipmap = cudaMipmappedImageArrayOrig;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = hipFilterModeLinear;\n texDescr.mipmapFilterMode = hipFilterModeLinear;\n\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n\n texDescr.maxMipmapLevelClamp = float(mipLevels - 1);\n\n texDescr.readMode = hipReadModeNormalizedFloat;\n\n checkCudaErrors(hipCreateTextureObject(&textureObjMipMapInput, &resDescr, &texDescr, NULL));\n\n checkCudaErrors(hipMalloc((void **)&d_surfaceObjectList, sizeof(hipSurfaceObject_t) * mipLevels));\n checkCudaErrors(hipMalloc((void **)&d_surfaceObjectListTemp, sizeof(hipSurfaceObject_t) * mipLevels));\n\n checkCudaErrors(hipMemcpy(d_surfaceObjectList,\n surfaceObjectList.data(),\n sizeof(hipSurfaceObject_t) * mipLevels,\n hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(d_surfaceObjectListTemp,\n surfaceObjectListTemp.data(),\n sizeof(hipSurfaceObject_t) * mipLevels,\n hipMemcpyHostToDevice));\n\n printf(\"CUDA Kernel Vulkan image buffer\\n\");\n }\n\n void cudaUpdateVkImage()\n {\n cudaVkSemaphoreWait(cudaExtVkUpdateCudaSemaphore);\n\n int nthreads = 128;\n\n /*Perform 2D box filter on image using CUDA */\n d_boxfilter_rgba_x<<>>(\n d_surfaceObjectListTemp, textureObjMipMapInput, imageWidth, imageHeight, mipLevels, filter_radius);\n\n d_boxfilter_rgba_y<<>>(\n d_surfaceObjectList, d_surfaceObjectListTemp, imageWidth, imageHeight, mipLevels, filter_radius);\n\n varySigma();\n\n cudaVkSemaphoreSignal(cudaExtCudaUpdateVkSemaphore);\n }\n\n void transitionImageLayout(VkImage image, VkFormat format, VkImageLayout oldLayout, VkImageLayout newLayout)\n {\n VkCommandBuffer commandBuffer = beginSingleTimeCommands();\n\n VkImageMemoryBarrier barrier = {};\n barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;\n barrier.oldLayout = oldLayout;\n barrier.newLayout = newLayout;\n barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;\n barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;\n barrier.image = image;\n barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;\n barrier.subresourceRange.baseMipLevel = 0;\n barrier.subresourceRange.levelCount = mipLevels;\n barrier.subresourceRange.baseArrayLayer = 0;\n barrier.subresourceRange.layerCount = 1;\n\n VkPipelineStageFlags sourceStage;\n VkPipelineStageFlags destinationStage;\n\n if (oldLayout == VK_IMAGE_LAYOUT_UNDEFINED && newLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL) {\n barrier.srcAccessMask = 0;\n barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;\n\n sourceStage = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;\n destinationStage = VK_PIPELINE_STAGE_TRANSFER_BIT;\n }\n else if (oldLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL\n && newLayout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL) {\n barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;\n barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;\n\n sourceStage = VK_PIPELINE_STAGE_TRANSFER_BIT;\n destinationStage = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;\n }\n else {\n throw std::invalid_argument(\"unsupported layout transition!\");\n }\n\n vkCmdPipelineBarrier(commandBuffer, sourceStage, destinationStage, 0, 0, nullptr, 0, nullptr, 1, &barrier);\n\n endSingleTimeCommands(commandBuffer);\n }\n\n void copyBufferToImage(VkBuffer buffer, VkImage image, uint32_t width, uint32_t height)\n {\n VkCommandBuffer commandBuffer = beginSingleTimeCommands();\n\n VkBufferImageCopy region = {};\n region.bufferOffset = 0;\n region.bufferRowLength = 0;\n region.bufferImageHeight = 0;\n region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;\n region.imageSubresource.mipLevel = 0;\n region.imageSubresource.baseArrayLayer = 0;\n region.imageSubresource.layerCount = 1;\n region.imageOffset = {0, 0, 0};\n region.imageExtent = {width, height, 1};\n\n vkCmdCopyBufferToImage(commandBuffer, buffer, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, ®ion);\n\n endSingleTimeCommands(commandBuffer);\n }\n\n void createVertexBuffer()\n {\n VkDeviceSize bufferSize = sizeof(vertices[0]) * vertices.size();\n\n VkBuffer stagingBuffer;\n VkDeviceMemory stagingBufferMemory;\n createBuffer(bufferSize,\n VK_BUFFER_USAGE_TRANSFER_SRC_BIT,\n VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,\n stagingBuffer,\n stagingBufferMemory);\n\n void *data;\n vkMapMemory(device, stagingBufferMemory, 0, bufferSize, 0, &data);\n memcpy(data, vertices.data(), (size_t)bufferSize);\n vkUnmapMemory(device, stagingBufferMemory);\n\n createBuffer(bufferSize,\n VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,\n VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,\n vertexBuffer,\n vertexBufferMemory);\n\n copyBuffer(stagingBuffer, vertexBuffer, bufferSize);\n\n vkDestroyBuffer(device, stagingBuffer, nullptr);\n vkFreeMemory(device, stagingBufferMemory, nullptr);\n }\n\n void createIndexBuffer()\n {\n VkDeviceSize bufferSize = sizeof(indices[0]) * indices.size();\n\n VkBuffer stagingBuffer;\n VkDeviceMemory stagingBufferMemory;\n createBuffer(bufferSize,\n VK_BUFFER_USAGE_TRANSFER_SRC_BIT,\n VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,\n stagingBuffer,\n stagingBufferMemory);\n\n void *data;\n vkMapMemory(device, stagingBufferMemory, 0, bufferSize, 0, &data);\n memcpy(data, indices.data(), (size_t)bufferSize);\n vkUnmapMemory(device, stagingBufferMemory);\n\n createBuffer(bufferSize,\n VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_INDEX_BUFFER_BIT,\n VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,\n indexBuffer,\n indexBufferMemory);\n\n copyBuffer(stagingBuffer, indexBuffer, bufferSize);\n\n vkDestroyBuffer(device, stagingBuffer, nullptr);\n vkFreeMemory(device, stagingBufferMemory, nullptr);\n }\n\n void createUniformBuffers()\n {\n VkDeviceSize bufferSize = sizeof(UniformBufferObject);\n\n uniformBuffers.resize(swapChainImages.size());\n uniformBuffersMemory.resize(swapChainImages.size());\n\n for (size_t i = 0; i < swapChainImages.size(); i++) {\n createBuffer(bufferSize,\n VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,\n VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,\n uniformBuffers[i],\n uniformBuffersMemory[i]);\n }\n }\n\n void createDescriptorPool()\n {\n std::array poolSizes = {};\n poolSizes[0].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;\n poolSizes[0].descriptorCount = static_cast(swapChainImages.size());\n poolSizes[1].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;\n poolSizes[1].descriptorCount = static_cast(swapChainImages.size());\n\n VkDescriptorPoolCreateInfo poolInfo = {};\n poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;\n poolInfo.poolSizeCount = static_cast(poolSizes.size());\n poolInfo.pPoolSizes = poolSizes.data();\n poolInfo.maxSets = static_cast(swapChainImages.size());\n\n if (vkCreateDescriptorPool(device, &poolInfo, nullptr, &descriptorPool) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create descriptor pool!\");\n }\n }\n\n void createDescriptorSets()\n {\n std::vector layouts(swapChainImages.size(), descriptorSetLayout);\n VkDescriptorSetAllocateInfo allocInfo = {};\n allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;\n allocInfo.descriptorPool = descriptorPool;\n allocInfo.descriptorSetCount = static_cast(swapChainImages.size());\n allocInfo.pSetLayouts = layouts.data();\n\n descriptorSets.resize(swapChainImages.size());\n if (vkAllocateDescriptorSets(device, &allocInfo, descriptorSets.data()) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to allocate descriptor sets!\");\n }\n\n for (size_t i = 0; i < swapChainImages.size(); i++) {\n VkDescriptorBufferInfo bufferInfo = {};\n bufferInfo.buffer = uniformBuffers[i];\n bufferInfo.offset = 0;\n bufferInfo.range = sizeof(UniformBufferObject);\n\n VkDescriptorImageInfo imageInfo = {};\n imageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;\n imageInfo.imageView = textureImageView;\n imageInfo.sampler = textureSampler;\n\n std::array descriptorWrites = {};\n\n descriptorWrites[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;\n descriptorWrites[0].dstSet = descriptorSets[i];\n descriptorWrites[0].dstBinding = 0;\n descriptorWrites[0].dstArrayElement = 0;\n descriptorWrites[0].descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;\n descriptorWrites[0].descriptorCount = 1;\n descriptorWrites[0].pBufferInfo = &bufferInfo;\n\n descriptorWrites[1].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;\n descriptorWrites[1].dstSet = descriptorSets[i];\n descriptorWrites[1].dstBinding = 1;\n descriptorWrites[1].dstArrayElement = 0;\n descriptorWrites[1].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;\n descriptorWrites[1].descriptorCount = 1;\n descriptorWrites[1].pImageInfo = &imageInfo;\n\n vkUpdateDescriptorSets(\n device, static_cast(descriptorWrites.size()), descriptorWrites.data(), 0, nullptr);\n }\n }\n\n void createBuffer(VkDeviceSize size,\n VkBufferUsageFlags usage,\n VkMemoryPropertyFlags properties,\n VkBuffer &buffer,\n VkDeviceMemory &bufferMemory)\n {\n VkBufferCreateInfo bufferInfo = {};\n bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;\n bufferInfo.size = size;\n bufferInfo.usage = usage;\n bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;\n\n if (vkCreateBuffer(device, &bufferInfo, nullptr, &buffer) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create buffer!\");\n }\n\n VkMemoryRequirements memRequirements;\n vkGetBufferMemoryRequirements(device, buffer, &memRequirements);\n\n VkMemoryAllocateInfo allocInfo = {};\n allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;\n allocInfo.allocationSize = memRequirements.size;\n allocInfo.memoryTypeIndex = findMemoryType(memRequirements.memoryTypeBits, properties);\n\n if (vkAllocateMemory(device, &allocInfo, nullptr, &bufferMemory) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to allocate buffer memory!\");\n }\n\n vkBindBufferMemory(device, buffer, bufferMemory, 0);\n }\n\n VkCommandBuffer beginSingleTimeCommands()\n {\n VkCommandBufferAllocateInfo allocInfo = {};\n allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;\n allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;\n allocInfo.commandPool = commandPool;\n allocInfo.commandBufferCount = 1;\n\n VkCommandBuffer commandBuffer;\n vkAllocateCommandBuffers(device, &allocInfo, &commandBuffer);\n\n VkCommandBufferBeginInfo beginInfo = {};\n beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;\n beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;\n\n vkBeginCommandBuffer(commandBuffer, &beginInfo);\n\n return commandBuffer;\n }\n\n void endSingleTimeCommands(VkCommandBuffer commandBuffer)\n {\n vkEndCommandBuffer(commandBuffer);\n\n VkSubmitInfo submitInfo = {};\n submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;\n submitInfo.commandBufferCount = 1;\n submitInfo.pCommandBuffers = &commandBuffer;\n\n vkQueueSubmit(graphicsQueue, 1, &submitInfo, VK_NULL_HANDLE);\n vkQueueWaitIdle(graphicsQueue);\n\n vkFreeCommandBuffers(device, commandPool, 1, &commandBuffer);\n }\n\n void copyBuffer(VkBuffer srcBuffer, VkBuffer dstBuffer, VkDeviceSize size)\n {\n VkCommandBuffer commandBuffer = beginSingleTimeCommands();\n\n VkBufferCopy copyRegion = {};\n copyRegion.size = size;\n vkCmdCopyBuffer(commandBuffer, srcBuffer, dstBuffer, 1, ©Region);\n\n endSingleTimeCommands(commandBuffer);\n }\n\n uint32_t findMemoryType(uint32_t typeFilter, VkMemoryPropertyFlags properties)\n {\n VkPhysicalDeviceMemoryProperties memProperties;\n vkGetPhysicalDeviceMemoryProperties(physicalDevice, &memProperties);\n\n for (uint32_t i = 0; i < memProperties.memoryTypeCount; i++) {\n if ((typeFilter & (1 << i)) && (memProperties.memoryTypes[i].propertyFlags & properties) == properties) {\n return i;\n }\n }\n\n throw std::runtime_error(\"failed to find suitable memory type!\");\n }\n\n void createCommandBuffers()\n {\n commandBuffers.resize(swapChainFramebuffers.size());\n\n VkCommandBufferAllocateInfo allocInfo = {};\n allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;\n allocInfo.commandPool = commandPool;\n allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;\n allocInfo.commandBufferCount = (uint32_t)commandBuffers.size();\n\n if (vkAllocateCommandBuffers(device, &allocInfo, commandBuffers.data()) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to allocate command buffers!\");\n }\n\n for (size_t i = 0; i < commandBuffers.size(); i++) {\n VkCommandBufferBeginInfo beginInfo = {};\n beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;\n beginInfo.flags = VK_COMMAND_BUFFER_USAGE_SIMULTANEOUS_USE_BIT;\n\n if (vkBeginCommandBuffer(commandBuffers[i], &beginInfo) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to begin recording command buffer!\");\n }\n\n VkRenderPassBeginInfo renderPassInfo = {};\n renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;\n renderPassInfo.renderPass = renderPass;\n renderPassInfo.framebuffer = swapChainFramebuffers[i];\n renderPassInfo.renderArea.offset = {0, 0};\n renderPassInfo.renderArea.extent = swapChainExtent;\n\n VkClearValue clearColor = {0.0f, 0.0f, 0.0f, 1.0f};\n renderPassInfo.clearValueCount = 1;\n renderPassInfo.pClearValues = &clearColor;\n\n vkCmdBeginRenderPass(commandBuffers[i], &renderPassInfo, VK_SUBPASS_CONTENTS_INLINE);\n\n vkCmdBindPipeline(commandBuffers[i], VK_PIPELINE_BIND_POINT_GRAPHICS, graphicsPipeline);\n\n VkBuffer vertexBuffers[] = {vertexBuffer};\n VkDeviceSize offsets[] = {0};\n vkCmdBindVertexBuffers(commandBuffers[i], 0, 1, vertexBuffers, offsets);\n\n vkCmdBindIndexBuffer(commandBuffers[i], indexBuffer, 0, VK_INDEX_TYPE_UINT16);\n\n vkCmdBindDescriptorSets(commandBuffers[i],\n VK_PIPELINE_BIND_POINT_GRAPHICS,\n pipelineLayout,\n 0,\n 1,\n &descriptorSets[i],\n 0,\n nullptr);\n\n vkCmdDrawIndexed(commandBuffers[i], static_cast(indices.size()), 1, 0, 0, 0);\n // vkCmdDraw(commandBuffers[i], static_cast(vertices.size()), 1,\n // 0, 0);\n\n vkCmdEndRenderPass(commandBuffers[i]);\n\n if (vkEndCommandBuffer(commandBuffers[i]) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to record command buffer!\");\n }\n }\n }\n\n void createSyncObjects()\n {\n imageAvailableSemaphores.resize(MAX_FRAMES);\n renderFinishedSemaphores.resize(MAX_FRAMES);\n inFlightFences.resize(MAX_FRAMES);\n\n VkSemaphoreCreateInfo semaphoreInfo = {};\n semaphoreInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;\n\n VkFenceCreateInfo fenceInfo = {};\n fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;\n fenceInfo.flags = VK_FENCE_CREATE_SIGNALED_BIT;\n\n for (size_t i = 0; i < MAX_FRAMES; i++) {\n if (vkCreateSemaphore(device, &semaphoreInfo, nullptr, &imageAvailableSemaphores[i]) != VK_SUCCESS\n || vkCreateSemaphore(device, &semaphoreInfo, nullptr, &renderFinishedSemaphores[i]) != VK_SUCCESS\n || vkCreateFence(device, &fenceInfo, nullptr, &inFlightFences[i]) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create synchronization objects for a frame!\");\n }\n }\n }\n\n void createSyncObjectsExt()\n {\n VkSemaphoreCreateInfo semaphoreInfo = {};\n semaphoreInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;\n\n memset(&semaphoreInfo, 0, sizeof(semaphoreInfo));\n semaphoreInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;\n\n#ifdef _WIN64\n WindowsSecurityAttributes winSecurityAttributes;\n\n VkExportSemaphoreWin32HandleInfoKHR vulkanExportSemaphoreWin32HandleInfoKHR = {};\n vulkanExportSemaphoreWin32HandleInfoKHR.sType = VK_STRUCTURE_TYPE_EXPORT_SEMAPHORE_WIN32_HANDLE_INFO_KHR;\n vulkanExportSemaphoreWin32HandleInfoKHR.pNext = NULL;\n vulkanExportSemaphoreWin32HandleInfoKHR.pAttributes = &winSecurityAttributes;\n vulkanExportSemaphoreWin32HandleInfoKHR.dwAccess = DXGI_SHARED_RESOURCE_READ | DXGI_SHARED_RESOURCE_WRITE;\n vulkanExportSemaphoreWin32HandleInfoKHR.name = (LPCWSTR)NULL;\n#endif\n VkExportSemaphoreCreateInfoKHR vulkanExportSemaphoreCreateInfo = {};\n vulkanExportSemaphoreCreateInfo.sType = VK_STRUCTURE_TYPE_EXPORT_SEMAPHORE_CREATE_INFO_KHR;\n#ifdef _WIN64\n vulkanExportSemaphoreCreateInfo.pNext = IsWindows8OrGreater() ? &vulkanExportSemaphoreWin32HandleInfoKHR : NULL;\n vulkanExportSemaphoreCreateInfo.handleTypes = IsWindows8OrGreater()\n ? VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_WIN32_BIT\n : VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_WIN32_KMT_BIT;\n#else\n vulkanExportSemaphoreCreateInfo.pNext = NULL;\n vulkanExportSemaphoreCreateInfo.handleTypes = VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_FD_BIT;\n#endif\n semaphoreInfo.pNext = &vulkanExportSemaphoreCreateInfo;\n\n if (vkCreateSemaphore(device, &semaphoreInfo, nullptr, &cudaUpdateVkSemaphore) != VK_SUCCESS\n || vkCreateSemaphore(device, &semaphoreInfo, nullptr, &vkUpdateCudaSemaphore) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create synchronization objects for a CUDA-Vulkan!\");\n }\n }\n\n void updateUniformBuffer()\n {\n UniformBufferObject ubo = {};\n\n mat4x4_identity(ubo.model);\n mat4x4 Model;\n mat4x4_dup(Model, ubo.model);\n mat4x4_rotate(ubo.model, Model, 0.0f, 0.0f, 1.0f, degreesToRadians(135.0f));\n\n vec3 eye = {2.0f, 2.0f, 2.0f};\n vec3 center = {0.0f, 0.0f, 0.0f};\n vec3 up = {0.0f, 0.0f, 1.0f};\n mat4x4_look_at(ubo.view, eye, center, up);\n\n mat4x4_perspective(\n ubo.proj, degreesToRadians(45.0f), swapChainExtent.width / (float)swapChainExtent.height, 0.1f, 10.0f);\n ubo.proj[1][1] *= -1;\n\n for (size_t i = 0; i < swapChainImages.size(); i++) {\n void *data;\n vkMapMemory(device, uniformBuffersMemory[i], 0, sizeof(ubo), 0, &data);\n memcpy(data, &ubo, sizeof(ubo));\n vkUnmapMemory(device, uniformBuffersMemory[i]);\n }\n }\n\n void drawFrame()\n {\n static int startSubmit = 0;\n\n vkWaitForFences(device, 1, &inFlightFences[currentFrame], VK_TRUE, std::numeric_limits::max());\n\n uint32_t imageIndex;\n VkResult result = vkAcquireNextImageKHR(device,\n swapChain,\n std::numeric_limits::max(),\n imageAvailableSemaphores[currentFrame],\n VK_NULL_HANDLE,\n &imageIndex);\n\n if (result == VK_ERROR_OUT_OF_DATE_KHR) {\n recreateSwapChain();\n return;\n }\n else if (result != VK_SUCCESS && result != VK_SUBOPTIMAL_KHR) {\n throw std::runtime_error(\"failed to acquire swap chain image!\");\n }\n\n vkResetFences(device, 1, &inFlightFences[currentFrame]);\n\n if (!startSubmit) {\n submitVulkan(imageIndex);\n startSubmit = 1;\n }\n else {\n submitVulkanCuda(imageIndex);\n }\n\n VkPresentInfoKHR presentInfo = {};\n presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;\n\n VkSemaphore signalSemaphores[] = {renderFinishedSemaphores[currentFrame]};\n\n presentInfo.waitSemaphoreCount = 1;\n presentInfo.pWaitSemaphores = signalSemaphores;\n\n VkSwapchainKHR swapChains[] = {swapChain};\n presentInfo.swapchainCount = 1;\n presentInfo.pSwapchains = swapChains;\n presentInfo.pImageIndices = &imageIndex;\n presentInfo.pResults = nullptr; // Optional\n\n result = vkQueuePresentKHR(presentQueue, &presentInfo);\n\n if (result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR || framebufferResized) {\n framebufferResized = false;\n recreateSwapChain();\n }\n else if (result != VK_SUCCESS) {\n throw std::runtime_error(\"failed to present swap chain image!\");\n }\n\n cudaUpdateVkImage();\n\n currentFrame = (currentFrame + 1) % MAX_FRAMES;\n // Added sleep of 10 millisecs so that CPU does not submit too much work to\n // GPU\n std::this_thread::sleep_for(std::chrono::microseconds(10000));\n char title[256];\n sprintf(title, \"Vulkan Image CUDA Box Filter (radius=%d)\", filter_radius);\n glfwSetWindowTitle(window, title);\n }\n\n void cudaVkSemaphoreSignal(hipExternalSemaphore_t &extSemaphore)\n {\n hipExternalSemaphoreSignalParams extSemaphoreSignalParams;\n memset(&extSemaphoreSignalParams, 0, sizeof(extSemaphoreSignalParams));\n\n extSemaphoreSignalParams.params.fence.value = 0;\n extSemaphoreSignalParams.flags = 0;\n checkCudaErrors(hipSignalExternalSemaphoresAsync(&extSemaphore, &extSemaphoreSignalParams, 1, streamToRun));\n }\n\n void cudaVkSemaphoreWait(hipExternalSemaphore_t &extSemaphore)\n {\n hipExternalSemaphoreWaitParams extSemaphoreWaitParams;\n\n memset(&extSemaphoreWaitParams, 0, sizeof(extSemaphoreWaitParams));\n\n extSemaphoreWaitParams.params.fence.value = 0;\n extSemaphoreWaitParams.flags = 0;\n\n checkCudaErrors(hipWaitExternalSemaphoresAsync(&extSemaphore, &extSemaphoreWaitParams, 1, streamToRun));\n }\n\n void submitVulkan(uint32_t imageIndex)\n {\n VkSubmitInfo submitInfo = {};\n submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;\n\n VkSemaphore waitSemaphores[] = {imageAvailableSemaphores[currentFrame]};\n VkPipelineStageFlags waitStages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};\n submitInfo.waitSemaphoreCount = 1;\n submitInfo.pWaitSemaphores = waitSemaphores;\n submitInfo.pWaitDstStageMask = waitStages;\n submitInfo.commandBufferCount = 1;\n submitInfo.pCommandBuffers = &commandBuffers[imageIndex];\n\n VkSemaphore signalSemaphores[] = {renderFinishedSemaphores[currentFrame], vkUpdateCudaSemaphore};\n\n submitInfo.signalSemaphoreCount = 2;\n submitInfo.pSignalSemaphores = signalSemaphores;\n\n if (vkQueueSubmit(graphicsQueue, 1, &submitInfo, inFlightFences[currentFrame]) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to submit draw command buffer!\");\n }\n }\n\n void submitVulkanCuda(uint32_t imageIndex)\n {\n VkSubmitInfo submitInfo = {};\n submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;\n\n VkSemaphore waitSemaphores[] = {imageAvailableSemaphores[currentFrame], cudaUpdateVkSemaphore};\n VkPipelineStageFlags waitStages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,\n VK_PIPELINE_STAGE_ALL_COMMANDS_BIT};\n submitInfo.waitSemaphoreCount = 2;\n submitInfo.pWaitSemaphores = waitSemaphores;\n submitInfo.pWaitDstStageMask = waitStages;\n submitInfo.commandBufferCount = 1;\n submitInfo.pCommandBuffers = &commandBuffers[imageIndex];\n\n VkSemaphore signalSemaphores[] = {renderFinishedSemaphores[currentFrame], vkUpdateCudaSemaphore};\n\n submitInfo.signalSemaphoreCount = 2;\n submitInfo.pSignalSemaphores = signalSemaphores;\n\n if (vkQueueSubmit(graphicsQueue, 1, &submitInfo, inFlightFences[currentFrame]) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to submit draw command buffer!\");\n }\n }\n\n VkShaderModule createShaderModule(const std::vector &code)\n {\n VkShaderModuleCreateInfo createInfo = {};\n createInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;\n createInfo.codeSize = code.size();\n createInfo.pCode = reinterpret_cast(code.data());\n\n VkShaderModule shaderModule;\n if (vkCreateShaderModule(device, &createInfo, nullptr, &shaderModule) != VK_SUCCESS) {\n throw std::runtime_error(\"failed to create shader module!\");\n }\n\n return shaderModule;\n }\n\n VkSurfaceFormatKHR chooseSwapSurfaceFormat(const std::vector &availableFormats)\n {\n if (availableFormats.size() == 1 && availableFormats[0].format == VK_FORMAT_UNDEFINED) {\n return {VK_FORMAT_B8G8R8A8_UNORM, VK_COLOR_SPACE_SRGB_NONLINEAR_KHR};\n }\n\n for (const auto &availableFormat : availableFormats) {\n if (availableFormat.format == VK_FORMAT_B8G8R8A8_UNORM\n && availableFormat.colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR) {\n return availableFormat;\n }\n }\n\n return availableFormats[0];\n }\n\n VkPresentModeKHR chooseSwapPresentMode(const std::vector &availablePresentModes)\n {\n VkPresentModeKHR bestMode = VK_PRESENT_MODE_FIFO_KHR;\n\n for (const auto &availablePresentMode : availablePresentModes) {\n if (availablePresentMode == VK_PRESENT_MODE_MAILBOX_KHR) {\n return availablePresentMode;\n }\n else if (availablePresentMode == VK_PRESENT_MODE_IMMEDIATE_KHR) {\n bestMode = availablePresentMode;\n }\n }\n\n return bestMode;\n }\n\n VkExtent2D chooseSwapExtent(const VkSurfaceCapabilitiesKHR &capabilities)\n {\n if (capabilities.currentExtent.width != std::numeric_limits::max()) {\n return capabilities.currentExtent;\n }\n else {\n int width, height;\n glfwGetFramebufferSize(window, &width, &height);\n\n VkExtent2D actualExtent = {static_cast(width), static_cast(height)};\n\n actualExtent.width = std::max(capabilities.minImageExtent.width,\n std::min(capabilities.maxImageExtent.width, actualExtent.width));\n actualExtent.height = std::max(capabilities.minImageExtent.height,\n std::min(capabilities.maxImageExtent.height, actualExtent.height));\n\n return actualExtent;\n }\n }\n\n SwapChainSupportDetails querySwapChainSupport(VkPhysicalDevice device)\n {\n SwapChainSupportDetails details;\n\n vkGetPhysicalDeviceSurfaceCapabilitiesKHR(device, surface, &details.capabilities);\n\n uint32_t formatCount;\n vkGetPhysicalDeviceSurfaceFormatsKHR(device, surface, &formatCount, nullptr);\n\n if (formatCount != 0) {\n details.formats.resize(formatCount);\n vkGetPhysicalDeviceSurfaceFormatsKHR(device, surface, &formatCount, details.formats.data());\n }\n\n uint32_t presentModeCount;\n vkGetPhysicalDeviceSurfacePresentModesKHR(device, surface, &presentModeCount, nullptr);\n\n if (presentModeCount != 0) {\n details.presentModes.resize(presentModeCount);\n vkGetPhysicalDeviceSurfacePresentModesKHR(device, surface, &presentModeCount, details.presentModes.data());\n }\n\n return details;\n }\n\n bool isDeviceSuitable(VkPhysicalDevice device)\n {\n QueueFamilyIndices indices = findQueueFamilies(device);\n\n bool extensionsSupported = checkDeviceExtensionSupport(device);\n\n bool swapChainAdequate = false;\n if (extensionsSupported) {\n SwapChainSupportDetails swapChainSupport = querySwapChainSupport(device);\n swapChainAdequate = !swapChainSupport.formats.empty() && !swapChainSupport.presentModes.empty();\n }\n\n VkPhysicalDeviceFeatures supportedFeatures;\n vkGetPhysicalDeviceFeatures(device, &supportedFeatures);\n\n return indices.isComplete() && extensionsSupported && swapChainAdequate && supportedFeatures.samplerAnisotropy;\n }\n\n bool checkDeviceExtensionSupport(VkPhysicalDevice device)\n {\n uint32_t extensionCount;\n vkEnumerateDeviceExtensionProperties(device, nullptr, &extensionCount, nullptr);\n\n std::vector availableExtensions(extensionCount);\n vkEnumerateDeviceExtensionProperties(device, nullptr, &extensionCount, availableExtensions.data());\n\n std::set requiredExtensions(deviceExtensions.begin(), deviceExtensions.end());\n\n for (const auto &extension : availableExtensions) {\n requiredExtensions.erase(extension.extensionName);\n }\n\n return requiredExtensions.empty();\n }\n\n QueueFamilyIndices findQueueFamilies(VkPhysicalDevice device)\n {\n QueueFamilyIndices indices;\n\n uint32_t queueFamilyCount = 0;\n vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, nullptr);\n\n std::vector queueFamilies(queueFamilyCount);\n vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, queueFamilies.data());\n\n int i = 0;\n for (const auto &queueFamily : queueFamilies) {\n if (queueFamily.queueCount > 0 && queueFamily.queueFlags & VK_QUEUE_GRAPHICS_BIT) {\n indices.graphicsFamily = i;\n }\n\n VkBool32 presentSupport = false;\n vkGetPhysicalDeviceSurfaceSupportKHR(device, i, surface, &presentSupport);\n\n if (queueFamily.queueCount > 0 && presentSupport) {\n indices.presentFamily = i;\n }\n\n if (indices.isComplete()) {\n break;\n }\n\n i++;\n }\n\n return indices;\n }\n\n std::vector getRequiredExtensions()\n {\n uint32_t glfwExtensionCount = 0;\n const char **glfwExtensions;\n glfwExtensions = glfwGetRequiredInstanceExtensions(&glfwExtensionCount);\n\n std::vector extensions(glfwExtensions, glfwExtensions + glfwExtensionCount);\n\n if (enableValidationLayers) {\n extensions.push_back(VK_EXT_DEBUG_UTILS_EXTENSION_NAME);\n }\n\n return extensions;\n }\n\n bool checkValidationLayerSupport()\n {\n uint32_t layerCount;\n vkEnumerateInstanceLayerProperties(&layerCount, nullptr);\n\n std::vector availableLayers(layerCount);\n vkEnumerateInstanceLayerProperties(&layerCount, availableLayers.data());\n\n for (const char *layerName : validationLayers) {\n bool layerFound = false;\n\n for (const auto &layerProperties : availableLayers) {\n if (strcmp(layerName, layerProperties.layerName) == 0) {\n layerFound = true;\n break;\n }\n }\n\n if (!layerFound) {\n return false;\n }\n }\n\n return true;\n }\n\n static std::vector readFile(const std::string &filename)\n {\n char *file_path = sdkFindFilePath(filename.c_str(), execution_path.c_str());\n std::ifstream file(file_path, std::ios::ate | std::ios::binary);\n\n if (!file.is_open()) {\n throw std::runtime_error(\"failed to open file!\");\n }\n\n size_t fileSize = (size_t)file.tellg();\n std::vector buffer(fileSize);\n\n file.seekg(0);\n file.read(buffer.data(), fileSize);\n\n file.close();\n\n return buffer;\n }\n\n static VKAPI_ATTR VkBool32 VKAPI_CALL debugCallback(VkDebugUtilsMessageSeverityFlagBitsEXT messageSeverity,\n VkDebugUtilsMessageTypeFlagsEXT messageType,\n const VkDebugUtilsMessengerCallbackDataEXT *pCallbackData,\n void *pUserData)\n {\n std::cerr << \"validation layer: \" << pCallbackData->pMessage << std::endl;\n\n return VK_FALSE;\n }\n};\n\nint main(int argc, char **argv)\n{\n execution_path = argv[0];\n std::string image_filename = \"teapot1024.ppm\";\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"file\")) {\n getCmdLineArgumentString(argc, (const char **)argv, \"file\", (char **)&image_filename);\n }\n\n vulkanImageCUDA app;\n\n try {\n // This app only works on ppm images\n app.loadImageData(image_filename);\n app.run();\n }\n catch (const std::exception &e) {\n std::cerr << e.what() << std::endl;\n return EXIT_FAILURE;\n }\n\n return EXIT_SUCCESS;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/vulkanImageCUDA/vulkanImageCUDA.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n\nusing namespace std;\n\nconst char *sSampleName = \"P2P (Peer-to-Peer) GPU Bandwidth Latency Test\";\n\ntypedef enum {\n P2P_WRITE = 0,\n P2P_READ = 1,\n} P2PDataTransfer;\n\ntypedef enum {\n CE = 0,\n SM = 1,\n} P2PEngine;\n\nP2PEngine p2p_mechanism = CE; // By default use Copy Engine\n\n// Macro for checking cuda errors following a cuda launch or api call\n#define cudaCheckError() \\\n { \\\n cudaError_t e = cudaGetLastError(); \\\n if (e != cudaSuccess) { \\\n printf(\"Cuda failure %s:%d: '%s'\\n\", __FILE__, __LINE__, cudaGetErrorString(e)); \\\n exit(EXIT_FAILURE); \\\n } \\\n }\n__global__ void delay(volatile int *flag, unsigned long long timeout_clocks = 10000000)\n{\n // Wait until the application notifies us that it has completed queuing up the\n // experiment, or timeout and exit, allowing the application to make progress\n long long int start_clock, sample_clock;\n start_clock = clock64();\n\n while (!*flag) {\n sample_clock = clock64();\n\n if (sample_clock - start_clock > timeout_clocks) {\n break;\n }\n }\n}\n\n// This kernel is for demonstration purposes only, not a performant kernel for\n// p2p transfers.\n__global__ void copyp2p(int4 *__restrict__ dest, int4 const *__restrict__ src, size_t num_elems)\n{\n size_t globalId = blockIdx.x * blockDim.x + threadIdx.x;\n size_t gridSize = blockDim.x * gridDim.x;\n\n#pragma unroll(5)\n for (size_t i = globalId; i < num_elems; i += gridSize) {\n dest[i] = src[i];\n }\n}\n\n///////////////////////////////////////////////////////////////////////////\n// Print help screen\n///////////////////////////////////////////////////////////////////////////\nvoid printHelp(void)\n{\n printf(\"Usage: p2pBandwidthLatencyTest [OPTION]...\\n\");\n printf(\"Tests bandwidth/latency of GPU pairs using P2P and without P2P\\n\");\n printf(\"\\n\");\n\n printf(\"Options:\\n\");\n printf(\"--help\\t\\tDisplay this help menu\\n\");\n printf(\"--p2p_read\\tUse P2P reads for data transfers between GPU pairs and show \"\n \"corresponding results.\\n \\t\\tDefault used is P2P write operation.\\n\");\n printf(\"--sm_copy Use SM intiated p2p transfers instead of Copy Engine\\n\");\n printf(\"--numElems= Number of integer elements to be used in p2p copy.\\n\");\n}\n\n///////////////////////////////////////////////////////////////////////////\n// Detect if cudaMemcpyPeerAsync will automatically fall back to\n// host-staged copies when P2P is disabled.\n//\n// We probe a single representative pair (device 0 -> device 1).\n// On a given system, confidential-computing (CC) and security\n// policies are uniform across GPUs, so if this pair is blocked\n// with cudaErrorNotSupported in P2P-off mode, it is reasonable\n// to assume all cross-GPU pairs behave the same.\n//\n// For a production application that must handle heterogeneous\n// environments, users may want to probe all device pairs.\n///////////////////////////////////////////////////////////////////////////\nbool detectFallback(int numGPUs)\n{\n if (numGPUs <= 1)\n return false;\n\n cudaSetDevice(0);\n int *tmp0 = nullptr, *tmp1 = nullptr;\n cudaStream_t s;\n cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking);\n\n size_t testElems = 1;\n cudaMalloc(&tmp0, testElems * sizeof(int));\n cudaSetDevice(1);\n cudaMalloc(&tmp1, testElems * sizeof(int));\n cudaCheckError();\n\n // Explicitly ensure P2P is disabled for this test\n // (Clear any pre-existing P2P access if it happens to be enabled)\n cudaSetDevice(0);\n cudaDeviceDisablePeerAccess(1);\n cudaGetLastError(); // Clear error if peer access was not enabled\n\n cudaSetDevice(1);\n cudaDeviceDisablePeerAccess(0);\n cudaGetLastError(); // Clear error if peer access was not enabled\n\n bool needsFallback = false;\n cudaError_t testErr = cudaMemcpyPeerAsync(tmp1, 1, tmp0, 0, testElems * sizeof(int), s);\n\n if (testErr == cudaErrorPeerAccessNotEnabled || testErr == cudaErrorNotSupported) {\n needsFallback = true;\n printf(\"Note: cudaMemcpyPeerAsync reported '%s' - will use host-mediated copy when P2P is disabled\\n\",\n cudaGetErrorString(testErr));\n cudaGetLastError();\n }\n\n cudaStreamSynchronize(s);\n cudaFree(tmp0);\n cudaFree(tmp1);\n cudaStreamDestroy(s);\n cudaCheckError();\n\n return needsFallback;\n}\n\nvoid checkP2Paccess(int numGPUs)\n{\n for (int i = 0; i < numGPUs; i++) {\n cudaSetDevice(i);\n cudaCheckError();\n\n for (int j = 0; j < numGPUs; j++) {\n int access;\n if (i != j) {\n cudaDeviceCanAccessPeer(&access, i, j);\n cudaCheckError();\n printf(\"Device=%d %s Access Peer Device=%d\\n\", i, access ? \"CAN\" : \"CANNOT\", j);\n }\n }\n }\n printf(\"\\n***NOTE: In case a device doesn't have P2P access to other one, it \"\n \"falls back to normal memcopy procedure.\\nSo you can see lesser \"\n \"Bandwidth (GB/s) and unstable Latency (us) in those cases.\\n\\n\");\n}\n\nvoid performP2PCopy(int *dest,\n int destDevice,\n int *src,\n int srcDevice,\n int num_elems,\n int repeat,\n bool p2paccess,\n cudaStream_t streamToRun,\n bool useFallback,\n int *hostBuffer)\n{\n int blockSize = 0;\n int numBlocks = 0;\n\n cudaOccupancyMaxPotentialBlockSize(&numBlocks, &blockSize, copyp2p);\n cudaCheckError();\n\n if (p2p_mechanism == SM && p2paccess) {\n for (int r = 0; r < repeat; r++) {\n copyp2p<<>>((int4 *)dest, (int4 *)src, num_elems / 4);\n }\n cudaCheckError();\n }\n else if (useFallback && srcDevice != destDevice) {\n // Use host-mediated copy for cross-GPU transfers when cudaMemcpyPeerAsync is not supported\n for (int r = 0; r < repeat; r++) {\n cudaMemcpyAsync(hostBuffer, src, sizeof(int) * num_elems, cudaMemcpyDeviceToHost, streamToRun);\n cudaMemcpyAsync(dest, hostBuffer, sizeof(int) * num_elems, cudaMemcpyHostToDevice, streamToRun);\n }\n cudaCheckError();\n }\n else if (useFallback && srcDevice == destDevice) {\n // Same device copy\n for (int r = 0; r < repeat; r++) {\n cudaMemcpyAsync(dest, src, sizeof(int) * num_elems, cudaMemcpyDeviceToDevice, streamToRun);\n }\n cudaCheckError();\n }\n else {\n // Use cudaMemcpyPeerAsync\n for (int r = 0; r < repeat; r++) {\n cudaMemcpyPeerAsync(dest, destDevice, src, srcDevice, sizeof(int) * num_elems, streamToRun);\n }\n cudaCheckError();\n }\n}\n\nvoid outputBandwidthMatrix(int numElems, int numGPUs, bool p2p, P2PDataTransfer p2p_method, bool needsFallback)\n{\n int repeat = 5;\n volatile int *flag = NULL;\n vector buffers(numGPUs);\n vector buffersD2D(numGPUs); // buffer for D2D, that is, intra-GPU copy\n vector start(numGPUs);\n vector stop(numGPUs);\n vector stream(numGPUs);\n\n cudaHostAlloc((void **)&flag, sizeof(*flag), cudaHostAllocPortable);\n cudaCheckError();\n\n for (int d = 0; d < numGPUs; d++) {\n cudaSetDevice(d);\n cudaStreamCreateWithFlags(&stream[d], cudaStreamNonBlocking);\n cudaMalloc(&buffers[d], numElems * sizeof(int));\n cudaCheckError();\n cudaMemset(buffers[d], 0, numElems * sizeof(int));\n cudaCheckError();\n cudaMalloc(&buffersD2D[d], numElems * sizeof(int));\n cudaCheckError();\n cudaMemset(buffersD2D[d], 0, numElems * sizeof(int));\n cudaCheckError();\n cudaEventCreate(&start[d]);\n cudaCheckError();\n cudaEventCreate(&stop[d]);\n cudaCheckError();\n }\n\n // Decide if fallback is needed based on global detection result and current test scenario\n bool useFallback = (!p2p && needsFallback);\n int *hostBuffer = nullptr;\n\n if (useFallback) {\n // Allocate hostBuffer for this function's numElems\n cudaHostAlloc((void **)&hostBuffer, sizeof(int) * numElems, cudaHostAllocDefault);\n cudaCheckError();\n }\n\n vector bandwidthMatrix(numGPUs * numGPUs);\n\n for (int i = 0; i < numGPUs; i++) {\n cudaSetDevice(i);\n\n for (int j = 0; j < numGPUs; j++) {\n int access = 0;\n if (p2p) {\n cudaDeviceCanAccessPeer(&access, i, j);\n if (access) {\n cudaDeviceEnablePeerAccess(j, 0);\n cudaCheckError();\n cudaSetDevice(j);\n cudaCheckError();\n cudaDeviceEnablePeerAccess(i, 0);\n cudaCheckError();\n cudaSetDevice(i);\n cudaCheckError();\n }\n }\n\n cudaStreamSynchronize(stream[i]);\n cudaCheckError();\n\n // Block the stream until all the work is queued up\n // DANGER! - cudaMemcpy*Async may infinitely block waiting for\n // room to push the operation, so keep the number of repeatitions\n // relatively low. Higher repeatitions will cause the delay kernel\n // to timeout and lead to unstable results.\n *flag = 0;\n delay<<<1, 1, 0, stream[i]>>>(flag);\n cudaCheckError();\n cudaEventRecord(start[i], stream[i]);\n cudaCheckError();\n\n if (i == j) {\n // Perform intra-GPU, D2D copies\n performP2PCopy(\n buffers[i], i, buffersD2D[i], i, numElems, repeat, access, stream[i], useFallback, hostBuffer);\n }\n else {\n if (p2p_method == P2P_WRITE) {\n performP2PCopy(\n buffers[j], j, buffers[i], i, numElems, repeat, access, stream[i], useFallback, hostBuffer);\n }\n else {\n performP2PCopy(\n buffers[i], i, buffers[j], j, numElems, repeat, access, stream[i], useFallback, hostBuffer);\n }\n }\n cudaCheckError();\n\n cudaEventRecord(stop[i], stream[i]);\n cudaCheckError();\n\n // Release the queued events\n *flag = 1;\n cudaStreamSynchronize(stream[i]);\n cudaCheckError();\n\n float time_ms;\n cudaEventElapsedTime(&time_ms, start[i], stop[i]);\n double time_s = time_ms / 1e3;\n\n double gb = numElems * sizeof(int) * repeat / (double)1e9;\n if (i == j) {\n gb *= 2; // must count both the read and the write here\n }\n bandwidthMatrix[i * numGPUs + j] = gb / time_s;\n if (p2p && access) {\n cudaDeviceDisablePeerAccess(j);\n cudaSetDevice(j);\n cudaDeviceDisablePeerAccess(i);\n cudaSetDevice(i);\n cudaCheckError();\n }\n }\n }\n\n printf(\" D\\\\D\");\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6d \", j);\n }\n\n printf(\"\\n\");\n\n for (int i = 0; i < numGPUs; i++) {\n printf(\"%6d \", i);\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6.02f \", bandwidthMatrix[i * numGPUs + j]);\n }\n\n printf(\"\\n\");\n }\n\n for (int d = 0; d < numGPUs; d++) {\n cudaSetDevice(d);\n cudaFree(buffers[d]);\n cudaFree(buffersD2D[d]);\n cudaCheckError();\n cudaEventDestroy(start[d]);\n cudaCheckError();\n cudaEventDestroy(stop[d]);\n cudaCheckError();\n cudaStreamDestroy(stream[d]);\n cudaCheckError();\n }\n\n if (hostBuffer) {\n cudaFreeHost(hostBuffer);\n cudaCheckError();\n }\n\n cudaFreeHost((void *)flag);\n cudaCheckError();\n}\n\nvoid outputBidirectionalBandwidthMatrix(int numElems, int numGPUs, bool p2p, bool needsFallback)\n{\n int repeat = 5;\n volatile int *flag = NULL;\n vector buffers(numGPUs);\n vector buffersD2D(numGPUs);\n vector start(numGPUs);\n vector stop(numGPUs);\n vector stream0(numGPUs);\n vector stream1(numGPUs);\n\n cudaHostAlloc((void **)&flag, sizeof(*flag), cudaHostAllocPortable);\n cudaCheckError();\n\n for (int d = 0; d < numGPUs; d++) {\n cudaSetDevice(d);\n cudaMalloc(&buffers[d], numElems * sizeof(int));\n cudaMemset(buffers[d], 0, numElems * sizeof(int));\n cudaMalloc(&buffersD2D[d], numElems * sizeof(int));\n cudaMemset(buffersD2D[d], 0, numElems * sizeof(int));\n cudaCheckError();\n cudaEventCreate(&start[d]);\n cudaCheckError();\n cudaEventCreate(&stop[d]);\n cudaCheckError();\n cudaStreamCreateWithFlags(&stream0[d], cudaStreamNonBlocking);\n cudaCheckError();\n cudaStreamCreateWithFlags(&stream1[d], cudaStreamNonBlocking);\n cudaCheckError();\n }\n\n // Decide if fallback is needed based on global detection result and current test scenario\n bool useFallback = (!p2p && needsFallback);\n int *hostBuffer = nullptr;\n\n if (useFallback) {\n // Allocate hostBuffer for this function's numElems\n cudaHostAlloc((void **)&hostBuffer, sizeof(int) * numElems, cudaHostAllocDefault);\n cudaCheckError();\n }\n\n vector bandwidthMatrix(numGPUs * numGPUs);\n\n for (int i = 0; i < numGPUs; i++) {\n cudaSetDevice(i);\n\n for (int j = 0; j < numGPUs; j++) {\n int access = 0;\n if (p2p) {\n cudaDeviceCanAccessPeer(&access, i, j);\n if (access) {\n cudaSetDevice(i);\n cudaDeviceEnablePeerAccess(j, 0);\n cudaCheckError();\n cudaSetDevice(j);\n cudaDeviceEnablePeerAccess(i, 0);\n cudaCheckError();\n }\n }\n\n cudaSetDevice(i);\n cudaStreamSynchronize(stream0[i]);\n cudaStreamSynchronize(stream1[j]);\n cudaCheckError();\n\n // Block the stream until all the work is queued up\n // DANGER! - cudaMemcpy*Async may infinitely block waiting for\n // room to push the operation, so keep the number of repeatitions\n // relatively low. Higher repeatitions will cause the delay kernel\n // to timeout and lead to unstable results.\n *flag = 0;\n cudaSetDevice(i);\n // No need to block stream1 since it'll be blocked on stream0's event\n delay<<<1, 1, 0, stream0[i]>>>(flag);\n cudaCheckError();\n\n // Force stream1 not to start until stream0 does, in order to ensure\n // the events on stream0 fully encompass the time needed for all\n // operations\n cudaEventRecord(start[i], stream0[i]);\n cudaStreamWaitEvent(stream1[j], start[i], 0);\n\n if (i == j) {\n // For intra-GPU perform 2 memcopies buffersD2D <-> buffers\n performP2PCopy(\n buffers[i], i, buffersD2D[i], i, numElems, repeat, access, stream0[i], useFallback, hostBuffer);\n performP2PCopy(\n buffersD2D[i], i, buffers[i], i, numElems, repeat, access, stream1[i], useFallback, hostBuffer);\n }\n else {\n if (access && p2p_mechanism == SM) {\n cudaSetDevice(j);\n }\n performP2PCopy(\n buffers[i], i, buffers[j], j, numElems, repeat, access, stream1[j], useFallback, hostBuffer);\n if (access && p2p_mechanism == SM) {\n cudaSetDevice(i);\n }\n performP2PCopy(\n buffers[j], j, buffers[i], i, numElems, repeat, access, stream0[i], useFallback, hostBuffer);\n }\n\n // Notify stream0 that stream1 is complete and record the time of\n // the total transaction\n cudaEventRecord(stop[j], stream1[j]);\n cudaStreamWaitEvent(stream0[i], stop[j], 0);\n cudaEventRecord(stop[i], stream0[i]);\n\n // Release the queued operations\n *flag = 1;\n cudaStreamSynchronize(stream0[i]);\n cudaStreamSynchronize(stream1[j]);\n cudaCheckError();\n\n float time_ms;\n cudaEventElapsedTime(&time_ms, start[i], stop[i]);\n double time_s = time_ms / 1e3;\n\n double gb = 2.0 * numElems * sizeof(int) * repeat / (double)1e9;\n if (i == j) {\n gb *= 2; // must count both the read and the write here\n }\n bandwidthMatrix[i * numGPUs + j] = gb / time_s;\n if (p2p && access) {\n cudaSetDevice(i);\n cudaDeviceDisablePeerAccess(j);\n cudaSetDevice(j);\n cudaDeviceDisablePeerAccess(i);\n }\n }\n }\n\n printf(\" D\\\\D\");\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6d \", j);\n }\n\n printf(\"\\n\");\n\n for (int i = 0; i < numGPUs; i++) {\n printf(\"%6d \", i);\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6.02f \", bandwidthMatrix[i * numGPUs + j]);\n }\n\n printf(\"\\n\");\n }\n\n for (int d = 0; d < numGPUs; d++) {\n cudaSetDevice(d);\n cudaFree(buffers[d]);\n cudaFree(buffersD2D[d]);\n cudaCheckError();\n cudaEventDestroy(start[d]);\n cudaCheckError();\n cudaEventDestroy(stop[d]);\n cudaCheckError();\n cudaStreamDestroy(stream0[d]);\n cudaCheckError();\n cudaStreamDestroy(stream1[d]);\n cudaCheckError();\n }\n\n if (hostBuffer) {\n cudaFreeHost(hostBuffer);\n cudaCheckError();\n }\n\n cudaFreeHost((void *)flag);\n cudaCheckError();\n}\n\nvoid outputLatencyMatrix(int numGPUs, bool p2p, P2PDataTransfer p2p_method, bool needsFallback)\n{\n int repeat = 100;\n int numElems = 4; // perform 1-int4 transfer.\n volatile int *flag = NULL;\n StopWatchInterface *stopWatch = NULL;\n vector buffers(numGPUs);\n vector buffersD2D(numGPUs); // buffer for D2D, that is, intra-GPU copy\n vector stream(numGPUs);\n vector start(numGPUs);\n vector stop(numGPUs);\n\n cudaHostAlloc((void **)&flag, sizeof(*flag), cudaHostAllocPortable);\n cudaCheckError();\n\n if (!sdkCreateTimer(&stopWatch)) {\n printf(\"Failed to create stop watch\\n\");\n exit(EXIT_FAILURE);\n }\n sdkStartTimer(&stopWatch);\n\n for (int d = 0; d < numGPUs; d++) {\n cudaSetDevice(d);\n cudaStreamCreateWithFlags(&stream[d], cudaStreamNonBlocking);\n cudaMalloc(&buffers[d], sizeof(int) * numElems);\n cudaMemset(buffers[d], 0, sizeof(int) * numElems);\n cudaMalloc(&buffersD2D[d], sizeof(int) * numElems);\n cudaMemset(buffersD2D[d], 0, sizeof(int) * numElems);\n cudaCheckError();\n cudaEventCreate(&start[d]);\n cudaCheckError();\n cudaEventCreate(&stop[d]);\n cudaCheckError();\n }\n\n // Decide if fallback is needed based on global detection result and current test scenario\n bool useFallback = (!p2p && needsFallback);\n int *hostBuffer = nullptr;\n\n if (useFallback) {\n // Allocate hostBuffer for this function's numElems\n cudaHostAlloc((void **)&hostBuffer, sizeof(int) * numElems, cudaHostAllocDefault);\n cudaCheckError();\n }\n\n vector gpuLatencyMatrix(numGPUs * numGPUs);\n vector cpuLatencyMatrix(numGPUs * numGPUs);\n\n for (int i = 0; i < numGPUs; i++) {\n cudaSetDevice(i);\n\n for (int j = 0; j < numGPUs; j++) {\n int access = 0;\n if (p2p) {\n cudaDeviceCanAccessPeer(&access, i, j);\n if (access) {\n cudaDeviceEnablePeerAccess(j, 0);\n cudaCheckError();\n cudaSetDevice(j);\n cudaDeviceEnablePeerAccess(i, 0);\n cudaSetDevice(i);\n cudaCheckError();\n }\n }\n cudaStreamSynchronize(stream[i]);\n cudaCheckError();\n\n // Block the stream until all the work is queued up\n // DANGER! - cudaMemcpy*Async may infinitely block waiting for\n // room to push the operation, so keep the number of repeatitions\n // relatively low. Higher repeatitions will cause the delay kernel\n // to timeout and lead to unstable results.\n *flag = 0;\n delay<<<1, 1, 0, stream[i]>>>(flag);\n cudaCheckError();\n cudaEventRecord(start[i], stream[i]);\n\n sdkResetTimer(&stopWatch);\n if (i == j) {\n // Perform intra-GPU, D2D copies\n performP2PCopy(\n buffers[i], i, buffersD2D[i], i, numElems, repeat, access, stream[i], useFallback, hostBuffer);\n }\n else {\n if (p2p_method == P2P_WRITE) {\n performP2PCopy(\n buffers[j], j, buffers[i], i, numElems, repeat, access, stream[i], useFallback, hostBuffer);\n }\n else {\n performP2PCopy(\n buffers[i], i, buffers[j], j, numElems, repeat, access, stream[i], useFallback, hostBuffer);\n }\n }\n float cpu_time_ms = sdkGetTimerValue(&stopWatch);\n\n cudaEventRecord(stop[i], stream[i]);\n // Now that the work has been queued up, release the stream\n *flag = 1;\n cudaStreamSynchronize(stream[i]);\n cudaCheckError();\n\n float gpu_time_ms;\n cudaEventElapsedTime(&gpu_time_ms, start[i], stop[i]);\n\n gpuLatencyMatrix[i * numGPUs + j] = gpu_time_ms * 1e3 / repeat;\n cpuLatencyMatrix[i * numGPUs + j] = cpu_time_ms * 1e3 / repeat;\n if (p2p && access) {\n cudaDeviceDisablePeerAccess(j);\n cudaSetDevice(j);\n cudaDeviceDisablePeerAccess(i);\n cudaSetDevice(i);\n cudaCheckError();\n }\n }\n }\n\n printf(\" GPU\");\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6d \", j);\n }\n\n printf(\"\\n\");\n\n for (int i = 0; i < numGPUs; i++) {\n printf(\"%6d \", i);\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6.02f \", gpuLatencyMatrix[i * numGPUs + j]);\n }\n\n printf(\"\\n\");\n }\n\n printf(\"\\n CPU\");\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6d \", j);\n }\n\n printf(\"\\n\");\n\n for (int i = 0; i < numGPUs; i++) {\n printf(\"%6d \", i);\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6.02f \", cpuLatencyMatrix[i * numGPUs + j]);\n }\n\n printf(\"\\n\");\n }\n\n for (int d = 0; d < numGPUs; d++) {\n cudaSetDevice(d);\n cudaFree(buffers[d]);\n cudaFree(buffersD2D[d]);\n cudaCheckError();\n cudaEventDestroy(start[d]);\n cudaCheckError();\n cudaEventDestroy(stop[d]);\n cudaCheckError();\n cudaStreamDestroy(stream[d]);\n cudaCheckError();\n }\n\n sdkDeleteTimer(&stopWatch);\n\n if (hostBuffer) {\n cudaFreeHost(hostBuffer);\n cudaCheckError();\n }\n\n cudaFreeHost((void *)flag);\n cudaCheckError();\n}\n\nint main(int argc, char **argv)\n{\n int numGPUs, numElems = 40000000;\n P2PDataTransfer p2p_method = P2P_WRITE;\n\n cudaGetDeviceCount(&numGPUs);\n cudaCheckError();\n\n // process command line args\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\")) {\n printHelp();\n return 0;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"p2p_read\")) {\n p2p_method = P2P_READ;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"sm_copy\")) {\n p2p_mechanism = SM;\n }\n\n // number of elements of int to be used in copy.\n if (checkCmdLineFlag(argc, (const char **)argv, \"numElems\")) {\n numElems = getCmdLineArgumentInt(argc, (const char **)argv, \"numElems\");\n }\n\n printf(\"[%s]\\n\", sSampleName);\n\n // output devices\n for (int i = 0; i < numGPUs; i++) {\n cudaDeviceProp prop;\n cudaGetDeviceProperties(&prop, i);\n cudaCheckError();\n printf(\"Device: %d, %s, pciBusID: %x, pciDeviceID: %x, pciDomainID:%x\\n\",\n i,\n prop.name,\n prop.pciBusID,\n prop.pciDeviceID,\n prop.pciDomainID);\n }\n\n checkP2Paccess(numGPUs);\n\n // Environment detection: One-time check if cudaMemcpyPeerAsync is supported when P2P is disabled\n bool needsFallback = detectFallback(numGPUs);\n\n // Check peer-to-peer connectivity\n printf(\"P2P Connectivity Matrix\\n\");\n printf(\" D\\\\D\");\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6d\", j);\n }\n printf(\"\\n\");\n\n for (int i = 0; i < numGPUs; i++) {\n printf(\"%6d\\t\", i);\n for (int j = 0; j < numGPUs; j++) {\n if (i != j) {\n int access;\n cudaDeviceCanAccessPeer(&access, i, j);\n cudaCheckError();\n printf(\"%6d\", (access) ? 1 : 0);\n }\n else {\n printf(\"%6d\", 1);\n }\n }\n printf(\"\\n\");\n }\n\n printf(\"Unidirectional P2P=Disabled Bandwidth Matrix (GB/s)\\n\");\n outputBandwidthMatrix(numElems, numGPUs, false, P2P_WRITE, needsFallback);\n printf(\"Unidirectional P2P=Enabled Bandwidth (P2P Writes) Matrix (GB/s)\\n\");\n outputBandwidthMatrix(numElems, numGPUs, true, P2P_WRITE, needsFallback);\n if (p2p_method == P2P_READ) {\n printf(\"Unidirectional P2P=Enabled Bandwidth (P2P Reads) Matrix (GB/s)\\n\");\n outputBandwidthMatrix(numElems, numGPUs, true, p2p_method, needsFallback);\n }\n printf(\"Bidirectional P2P=Disabled Bandwidth Matrix (GB/s)\\n\");\n outputBidirectionalBandwidthMatrix(numElems, numGPUs, false, needsFallback);\n printf(\"Bidirectional P2P=Enabled Bandwidth Matrix (GB/s)\\n\");\n outputBidirectionalBandwidthMatrix(numElems, numGPUs, true, needsFallback);\n\n printf(\"P2P=Disabled Latency Matrix (us)\\n\");\n outputLatencyMatrix(numGPUs, false, P2P_WRITE, needsFallback);\n printf(\"P2P=Enabled Latency (P2P Writes) Matrix (us)\\n\");\n outputLatencyMatrix(numGPUs, true, P2P_WRITE, needsFallback);\n if (p2p_method == P2P_READ) {\n printf(\"P2P=Enabled Latency (P2P Reads) Matrix (us)\\n\");\n outputLatencyMatrix(numGPUs, true, p2p_method, needsFallback);\n }\n\n printf(\"\\nNOTE: The CUDA Samples are not meant for performance measurements. \"\n \"Results may vary when GPU Boost is enabled.\\n\");\n\n exit(EXIT_SUCCESS);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\nusing namespace std;\n\nconst char *sSampleName = \"P2P (Peer-to-Peer) GPU Bandwidth Latency Test\";\n\ntypedef enum {\n P2P_WRITE = 0,\n P2P_READ = 1,\n} P2PDataTransfer;\n\ntypedef enum {\n CE = 0,\n SM = 1,\n} P2PEngine;\n\nP2PEngine p2p_mechanism = CE; // By default use Copy Engine\n\n// Macro for checking cuda errors following a cuda launch or api call\n#define cudaCheckError() \\\n { \\\n hipError_t e = hipGetLastError(); \\\n if (e != hipSuccess) { \\\n printf(\"Cuda failure %s:%d: '%s'\\n\", __FILE__, __LINE__, hipGetErrorString(e)); \\\n exit(EXIT_FAILURE); \\\n } \\\n }\n__global__ void delay(volatile int *flag, unsigned long long timeout_clocks = 10000000)\n{\n // Wait until the application notifies us that it has completed queuing up the\n // experiment, or timeout and exit, allowing the application to make progress\n long long int start_clock, sample_clock;\n start_clock = clock64();\n\n while (!*flag) {\n sample_clock = clock64();\n\n if (sample_clock - start_clock > timeout_clocks) {\n break;\n }\n }\n}\n\n// This kernel is for demonstration purposes only, not a performant kernel for\n// p2p transfers.\n__global__ void copyp2p(int4 *__restrict__ dest, int4 const *__restrict__ src, size_t num_elems)\n{\n size_t globalId = blockIdx.x * blockDim.x + threadIdx.x;\n size_t gridSize = blockDim.x * gridDim.x;\n\n#pragma unroll(5)\n for (size_t i = globalId; i < num_elems; i += gridSize) {\n dest[i] = src[i];\n }\n}\n\n///////////////////////////////////////////////////////////////////////////\n// Print help screen\n///////////////////////////////////////////////////////////////////////////\nvoid printHelp(void)\n{\n printf(\"Usage: p2pBandwidthLatencyTest [OPTION]...\\n\");\n printf(\"Tests bandwidth/latency of GPU pairs using P2P and without P2P\\n\");\n printf(\"\\n\");\n\n printf(\"Options:\\n\");\n printf(\"--help\\t\\tDisplay this help menu\\n\");\n printf(\"--p2p_read\\tUse P2P reads for data transfers between GPU pairs and show \"\n \"corresponding results.\\n \\t\\tDefault used is P2P write operation.\\n\");\n printf(\"--sm_copy Use SM intiated p2p transfers instead of Copy Engine\\n\");\n printf(\"--numElems= Number of integer elements to be used in p2p copy.\\n\");\n}\n\n///////////////////////////////////////////////////////////////////////////\n// Detect if cudaMemcpyPeerAsync will automatically fall back to\n// host-staged copies when P2P is disabled.\n//\n// We probe a single representative pair (device 0 -> device 1).\n// On a given system, confidential-computing (CC) and security\n// policies are uniform across GPUs, so if this pair is blocked\n// with cudaErrorNotSupported in P2P-off mode, it is reasonable\n// to assume all cross-GPU pairs behave the same.\n//\n// For a production application that must handle heterogeneous\n// environments, users may want to probe all device pairs.\n///////////////////////////////////////////////////////////////////////////\nbool detectFallback(int numGPUs)\n{\n if (numGPUs <= 1)\n return false;\n\n hipSetDevice(0);\n int *tmp0 = nullptr, *tmp1 = nullptr;\n hipStream_t s;\n hipStreamCreateWithFlags(&s, hipStreamNonBlocking);\n\n size_t testElems = 1;\n hipMalloc(&tmp0, testElems * sizeof(int));\n hipSetDevice(1);\n hipMalloc(&tmp1, testElems * sizeof(int));\n cudaCheckError();\n\n // Explicitly ensure P2P is disabled for this test\n // (Clear any pre-existing P2P access if it happens to be enabled)\n hipSetDevice(0);\n hipDeviceDisablePeerAccess(1);\n hipGetLastError(); // Clear error if peer access was not enabled\n\n hipSetDevice(1);\n hipDeviceDisablePeerAccess(0);\n hipGetLastError(); // Clear error if peer access was not enabled\n\n bool needsFallback = false;\n hipError_t testErr = hipMemcpyPeerAsync(tmp1, 1, tmp0, 0, testElems * sizeof(int), s);\n\n if (testErr == hipErrorPeerAccessNotEnabled || testErr == hipErrorNotSupported) {\n needsFallback = true;\n printf(\"Note: hipMemcpyPeerAsync reported '%s' - will use host-mediated copy when P2P is disabled\\n\",\n hipGetErrorString(testErr));\n hipGetLastError();\n }\n\n hipStreamSynchronize(s);\n hipFree(tmp0);\n hipFree(tmp1);\n hipStreamDestroy(s);\n cudaCheckError();\n\n return needsFallback;\n}\n\nvoid checkP2Paccess(int numGPUs)\n{\n for (int i = 0; i < numGPUs; i++) {\n hipSetDevice(i);\n cudaCheckError();\n\n for (int j = 0; j < numGPUs; j++) {\n int access;\n if (i != j) {\n hipDeviceCanAccessPeer(&access, i, j);\n cudaCheckError();\n printf(\"Device=%d %s Access Peer Device=%d\\n\", i, access ? \"CAN\" : \"CANNOT\", j);\n }\n }\n }\n printf(\"\\n***NOTE: In case a device doesn't have P2P access to other one, it \"\n \"falls back to normal memcopy procedure.\\nSo you can see lesser \"\n \"Bandwidth (GB/s) and unstable Latency (us) in those cases.\\n\\n\");\n}\n\nvoid performP2PCopy(int *dest,\n int destDevice,\n int *src,\n int srcDevice,\n int num_elems,\n int repeat,\n bool p2paccess,\n hipStream_t streamToRun,\n bool useFallback,\n int *hostBuffer)\n{\n int blockSize = 0;\n int numBlocks = 0;\n\n hipOccupancyMaxPotentialBlockSize(&numBlocks, &blockSize, copyp2p);\n cudaCheckError();\n\n if (p2p_mechanism == SM && p2paccess) {\n for (int r = 0; r < repeat; r++) {\n copyp2p<<>>((int4 *)dest, (int4 *)src, num_elems / 4);\n }\n cudaCheckError();\n }\n else if (useFallback && srcDevice != destDevice) {\n // Use host-mediated copy for cross-GPU transfers when cudaMemcpyPeerAsync is not supported\n for (int r = 0; r < repeat; r++) {\n hipMemcpyAsync(hostBuffer, src, sizeof(int) * num_elems, hipMemcpyDeviceToHost, streamToRun);\n hipMemcpyAsync(dest, hostBuffer, sizeof(int) * num_elems, hipMemcpyHostToDevice, streamToRun);\n }\n cudaCheckError();\n }\n else if (useFallback && srcDevice == destDevice) {\n // Same device copy\n for (int r = 0; r < repeat; r++) {\n hipMemcpyAsync(dest, src, sizeof(int) * num_elems, hipMemcpyDeviceToDevice, streamToRun);\n }\n cudaCheckError();\n }\n else {\n // Use cudaMemcpyPeerAsync\n for (int r = 0; r < repeat; r++) {\n hipMemcpyPeerAsync(dest, destDevice, src, srcDevice, sizeof(int) * num_elems, streamToRun);\n }\n cudaCheckError();\n }\n}\n\nvoid outputBandwidthMatrix(int numElems, int numGPUs, bool p2p, P2PDataTransfer p2p_method, bool needsFallback)\n{\n int repeat = 5;\n volatile int *flag = NULL;\n vector buffers(numGPUs);\n vector buffersD2D(numGPUs); // buffer for D2D, that is, intra-GPU copy\n vector start(numGPUs);\n vector stop(numGPUs);\n vector stream(numGPUs);\n\n hipHostAlloc((void **)&flag, sizeof(*flag), hipHostMallocPortable);\n cudaCheckError();\n\n for (int d = 0; d < numGPUs; d++) {\n hipSetDevice(d);\n hipStreamCreateWithFlags(&stream[d], hipStreamNonBlocking);\n hipMalloc(&buffers[d], numElems * sizeof(int));\n cudaCheckError();\n hipMemset(buffers[d], 0, numElems * sizeof(int));\n cudaCheckError();\n hipMalloc(&buffersD2D[d], numElems * sizeof(int));\n cudaCheckError();\n hipMemset(buffersD2D[d], 0, numElems * sizeof(int));\n cudaCheckError();\n hipEventCreate(&start[d]);\n cudaCheckError();\n hipEventCreate(&stop[d]);\n cudaCheckError();\n }\n\n // Decide if fallback is needed based on global detection result and current test scenario\n bool useFallback = (!p2p && needsFallback);\n int *hostBuffer = nullptr;\n\n if (useFallback) {\n // Allocate hostBuffer for this function's numElems\n hipHostAlloc((void **)&hostBuffer, sizeof(int) * numElems, hipHostMallocDefault);\n cudaCheckError();\n }\n\n vector bandwidthMatrix(numGPUs * numGPUs);\n\n for (int i = 0; i < numGPUs; i++) {\n hipSetDevice(i);\n\n for (int j = 0; j < numGPUs; j++) {\n int access = 0;\n if (p2p) {\n hipDeviceCanAccessPeer(&access, i, j);\n if (access) {\n hipDeviceEnablePeerAccess(j, 0);\n cudaCheckError();\n hipSetDevice(j);\n cudaCheckError();\n hipDeviceEnablePeerAccess(i, 0);\n cudaCheckError();\n hipSetDevice(i);\n cudaCheckError();\n }\n }\n\n hipStreamSynchronize(stream[i]);\n cudaCheckError();\n\n // Block the stream until all the work is queued up\n // DANGER! - cudaMemcpy*Async may infinitely block waiting for\n // room to push the operation, so keep the number of repeatitions\n // relatively low. Higher repeatitions will cause the delay kernel\n // to timeout and lead to unstable results.\n *flag = 0;\n delay<<<1, 1, 0, stream[i]>>>(flag);\n cudaCheckError();\n hipEventRecord(start[i], stream[i]);\n cudaCheckError();\n\n if (i == j) {\n // Perform intra-GPU, D2D copies\n performP2PCopy(\n buffers[i], i, buffersD2D[i], i, numElems, repeat, access, stream[i], useFallback, hostBuffer);\n }\n else {\n if (p2p_method == P2P_WRITE) {\n performP2PCopy(\n buffers[j], j, buffers[i], i, numElems, repeat, access, stream[i], useFallback, hostBuffer);\n }\n else {\n performP2PCopy(\n buffers[i], i, buffers[j], j, numElems, repeat, access, stream[i], useFallback, hostBuffer);\n }\n }\n cudaCheckError();\n\n hipEventRecord(stop[i], stream[i]);\n cudaCheckError();\n\n // Release the queued events\n *flag = 1;\n hipStreamSynchronize(stream[i]);\n cudaCheckError();\n\n float time_ms;\n hipEventElapsedTime(&time_ms, start[i], stop[i]);\n double time_s = time_ms / 1e3;\n\n double gb = numElems * sizeof(int) * repeat / (double)1e9;\n if (i == j) {\n gb *= 2; // must count both the read and the write here\n }\n bandwidthMatrix[i * numGPUs + j] = gb / time_s;\n if (p2p && access) {\n hipDeviceDisablePeerAccess(j);\n hipSetDevice(j);\n hipDeviceDisablePeerAccess(i);\n hipSetDevice(i);\n cudaCheckError();\n }\n }\n }\n\n printf(\" D\\\\D\");\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6d \", j);\n }\n\n printf(\"\\n\");\n\n for (int i = 0; i < numGPUs; i++) {\n printf(\"%6d \", i);\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6.02f \", bandwidthMatrix[i * numGPUs + j]);\n }\n\n printf(\"\\n\");\n }\n\n for (int d = 0; d < numGPUs; d++) {\n hipSetDevice(d);\n hipFree(buffers[d]);\n hipFree(buffersD2D[d]);\n cudaCheckError();\n hipEventDestroy(start[d]);\n cudaCheckError();\n hipEventDestroy(stop[d]);\n cudaCheckError();\n hipStreamDestroy(stream[d]);\n cudaCheckError();\n }\n\n if (hostBuffer) {\n hipHostFree(hostBuffer);\n cudaCheckError();\n }\n\n hipHostFree((void *)flag);\n cudaCheckError();\n}\n\nvoid outputBidirectionalBandwidthMatrix(int numElems, int numGPUs, bool p2p, bool needsFallback)\n{\n int repeat = 5;\n volatile int *flag = NULL;\n vector buffers(numGPUs);\n vector buffersD2D(numGPUs);\n vector start(numGPUs);\n vector stop(numGPUs);\n vector stream0(numGPUs);\n vector stream1(numGPUs);\n\n hipHostAlloc((void **)&flag, sizeof(*flag), hipHostMallocPortable);\n cudaCheckError();\n\n for (int d = 0; d < numGPUs; d++) {\n hipSetDevice(d);\n hipMalloc(&buffers[d], numElems * sizeof(int));\n hipMemset(buffers[d], 0, numElems * sizeof(int));\n hipMalloc(&buffersD2D[d], numElems * sizeof(int));\n hipMemset(buffersD2D[d], 0, numElems * sizeof(int));\n cudaCheckError();\n hipEventCreate(&start[d]);\n cudaCheckError();\n hipEventCreate(&stop[d]);\n cudaCheckError();\n hipStreamCreateWithFlags(&stream0[d], hipStreamNonBlocking);\n cudaCheckError();\n hipStreamCreateWithFlags(&stream1[d], hipStreamNonBlocking);\n cudaCheckError();\n }\n\n // Decide if fallback is needed based on global detection result and current test scenario\n bool useFallback = (!p2p && needsFallback);\n int *hostBuffer = nullptr;\n\n if (useFallback) {\n // Allocate hostBuffer for this function's numElems\n hipHostAlloc((void **)&hostBuffer, sizeof(int) * numElems, hipHostMallocDefault);\n cudaCheckError();\n }\n\n vector bandwidthMatrix(numGPUs * numGPUs);\n\n for (int i = 0; i < numGPUs; i++) {\n hipSetDevice(i);\n\n for (int j = 0; j < numGPUs; j++) {\n int access = 0;\n if (p2p) {\n hipDeviceCanAccessPeer(&access, i, j);\n if (access) {\n hipSetDevice(i);\n hipDeviceEnablePeerAccess(j, 0);\n cudaCheckError();\n hipSetDevice(j);\n hipDeviceEnablePeerAccess(i, 0);\n cudaCheckError();\n }\n }\n\n hipSetDevice(i);\n hipStreamSynchronize(stream0[i]);\n hipStreamSynchronize(stream1[j]);\n cudaCheckError();\n\n // Block the stream until all the work is queued up\n // DANGER! - cudaMemcpy*Async may infinitely block waiting for\n // room to push the operation, so keep the number of repeatitions\n // relatively low. Higher repeatitions will cause the delay kernel\n // to timeout and lead to unstable results.\n *flag = 0;\n hipSetDevice(i);\n // No need to block stream1 since it'll be blocked on stream0's event\n delay<<<1, 1, 0, stream0[i]>>>(flag);\n cudaCheckError();\n\n // Force stream1 not to start until stream0 does, in order to ensure\n // the events on stream0 fully encompass the time needed for all\n // operations\n hipEventRecord(start[i], stream0[i]);\n hipStreamWaitEvent(stream1[j], start[i], 0);\n\n if (i == j) {\n // For intra-GPU perform 2 memcopies buffersD2D <-> buffers\n performP2PCopy(\n buffers[i], i, buffersD2D[i], i, numElems, repeat, access, stream0[i], useFallback, hostBuffer);\n performP2PCopy(\n buffersD2D[i], i, buffers[i], i, numElems, repeat, access, stream1[i], useFallback, hostBuffer);\n }\n else {\n if (access && p2p_mechanism == SM) {\n hipSetDevice(j);\n }\n performP2PCopy(\n buffers[i], i, buffers[j], j, numElems, repeat, access, stream1[j], useFallback, hostBuffer);\n if (access && p2p_mechanism == SM) {\n hipSetDevice(i);\n }\n performP2PCopy(\n buffers[j], j, buffers[i], i, numElems, repeat, access, stream0[i], useFallback, hostBuffer);\n }\n\n // Notify stream0 that stream1 is complete and record the time of\n // the total transaction\n hipEventRecord(stop[j], stream1[j]);\n hipStreamWaitEvent(stream0[i], stop[j], 0);\n hipEventRecord(stop[i], stream0[i]);\n\n // Release the queued operations\n *flag = 1;\n hipStreamSynchronize(stream0[i]);\n hipStreamSynchronize(stream1[j]);\n cudaCheckError();\n\n float time_ms;\n hipEventElapsedTime(&time_ms, start[i], stop[i]);\n double time_s = time_ms / 1e3;\n\n double gb = 2.0 * numElems * sizeof(int) * repeat / (double)1e9;\n if (i == j) {\n gb *= 2; // must count both the read and the write here\n }\n bandwidthMatrix[i * numGPUs + j] = gb / time_s;\n if (p2p && access) {\n hipSetDevice(i);\n hipDeviceDisablePeerAccess(j);\n hipSetDevice(j);\n hipDeviceDisablePeerAccess(i);\n }\n }\n }\n\n printf(\" D\\\\D\");\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6d \", j);\n }\n\n printf(\"\\n\");\n\n for (int i = 0; i < numGPUs; i++) {\n printf(\"%6d \", i);\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6.02f \", bandwidthMatrix[i * numGPUs + j]);\n }\n\n printf(\"\\n\");\n }\n\n for (int d = 0; d < numGPUs; d++) {\n hipSetDevice(d);\n hipFree(buffers[d]);\n hipFree(buffersD2D[d]);\n cudaCheckError();\n hipEventDestroy(start[d]);\n cudaCheckError();\n hipEventDestroy(stop[d]);\n cudaCheckError();\n hipStreamDestroy(stream0[d]);\n cudaCheckError();\n hipStreamDestroy(stream1[d]);\n cudaCheckError();\n }\n\n if (hostBuffer) {\n hipHostFree(hostBuffer);\n cudaCheckError();\n }\n\n hipHostFree((void *)flag);\n cudaCheckError();\n}\n\nvoid outputLatencyMatrix(int numGPUs, bool p2p, P2PDataTransfer p2p_method, bool needsFallback)\n{\n int repeat = 100;\n int numElems = 4; // perform 1-int4 transfer.\n volatile int *flag = NULL;\n StopWatchInterface *stopWatch = NULL;\n vector buffers(numGPUs);\n vector buffersD2D(numGPUs); // buffer for D2D, that is, intra-GPU copy\n vector stream(numGPUs);\n vector start(numGPUs);\n vector stop(numGPUs);\n\n hipHostAlloc((void **)&flag, sizeof(*flag), hipHostMallocPortable);\n cudaCheckError();\n\n if (!sdkCreateTimer(&stopWatch)) {\n printf(\"Failed to create stop watch\\n\");\n exit(EXIT_FAILURE);\n }\n sdkStartTimer(&stopWatch);\n\n for (int d = 0; d < numGPUs; d++) {\n hipSetDevice(d);\n hipStreamCreateWithFlags(&stream[d], hipStreamNonBlocking);\n hipMalloc(&buffers[d], sizeof(int) * numElems);\n hipMemset(buffers[d], 0, sizeof(int) * numElems);\n hipMalloc(&buffersD2D[d], sizeof(int) * numElems);\n hipMemset(buffersD2D[d], 0, sizeof(int) * numElems);\n cudaCheckError();\n hipEventCreate(&start[d]);\n cudaCheckError();\n hipEventCreate(&stop[d]);\n cudaCheckError();\n }\n\n // Decide if fallback is needed based on global detection result and current test scenario\n bool useFallback = (!p2p && needsFallback);\n int *hostBuffer = nullptr;\n\n if (useFallback) {\n // Allocate hostBuffer for this function's numElems\n hipHostAlloc((void **)&hostBuffer, sizeof(int) * numElems, hipHostMallocDefault);\n cudaCheckError();\n }\n\n vector gpuLatencyMatrix(numGPUs * numGPUs);\n vector cpuLatencyMatrix(numGPUs * numGPUs);\n\n for (int i = 0; i < numGPUs; i++) {\n hipSetDevice(i);\n\n for (int j = 0; j < numGPUs; j++) {\n int access = 0;\n if (p2p) {\n hipDeviceCanAccessPeer(&access, i, j);\n if (access) {\n hipDeviceEnablePeerAccess(j, 0);\n cudaCheckError();\n hipSetDevice(j);\n hipDeviceEnablePeerAccess(i, 0);\n hipSetDevice(i);\n cudaCheckError();\n }\n }\n hipStreamSynchronize(stream[i]);\n cudaCheckError();\n\n // Block the stream until all the work is queued up\n // DANGER! - cudaMemcpy*Async may infinitely block waiting for\n // room to push the operation, so keep the number of repeatitions\n // relatively low. Higher repeatitions will cause the delay kernel\n // to timeout and lead to unstable results.\n *flag = 0;\n delay<<<1, 1, 0, stream[i]>>>(flag);\n cudaCheckError();\n hipEventRecord(start[i], stream[i]);\n\n sdkResetTimer(&stopWatch);\n if (i == j) {\n // Perform intra-GPU, D2D copies\n performP2PCopy(\n buffers[i], i, buffersD2D[i], i, numElems, repeat, access, stream[i], useFallback, hostBuffer);\n }\n else {\n if (p2p_method == P2P_WRITE) {\n performP2PCopy(\n buffers[j], j, buffers[i], i, numElems, repeat, access, stream[i], useFallback, hostBuffer);\n }\n else {\n performP2PCopy(\n buffers[i], i, buffers[j], j, numElems, repeat, access, stream[i], useFallback, hostBuffer);\n }\n }\n float cpu_time_ms = sdkGetTimerValue(&stopWatch);\n\n hipEventRecord(stop[i], stream[i]);\n // Now that the work has been queued up, release the stream\n *flag = 1;\n hipStreamSynchronize(stream[i]);\n cudaCheckError();\n\n float gpu_time_ms;\n hipEventElapsedTime(&gpu_time_ms, start[i], stop[i]);\n\n gpuLatencyMatrix[i * numGPUs + j] = gpu_time_ms * 1e3 / repeat;\n cpuLatencyMatrix[i * numGPUs + j] = cpu_time_ms * 1e3 / repeat;\n if (p2p && access) {\n hipDeviceDisablePeerAccess(j);\n hipSetDevice(j);\n hipDeviceDisablePeerAccess(i);\n hipSetDevice(i);\n cudaCheckError();\n }\n }\n }\n\n printf(\" GPU\");\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6d \", j);\n }\n\n printf(\"\\n\");\n\n for (int i = 0; i < numGPUs; i++) {\n printf(\"%6d \", i);\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6.02f \", gpuLatencyMatrix[i * numGPUs + j]);\n }\n\n printf(\"\\n\");\n }\n\n printf(\"\\n CPU\");\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6d \", j);\n }\n\n printf(\"\\n\");\n\n for (int i = 0; i < numGPUs; i++) {\n printf(\"%6d \", i);\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6.02f \", cpuLatencyMatrix[i * numGPUs + j]);\n }\n\n printf(\"\\n\");\n }\n\n for (int d = 0; d < numGPUs; d++) {\n hipSetDevice(d);\n hipFree(buffers[d]);\n hipFree(buffersD2D[d]);\n cudaCheckError();\n hipEventDestroy(start[d]);\n cudaCheckError();\n hipEventDestroy(stop[d]);\n cudaCheckError();\n hipStreamDestroy(stream[d]);\n cudaCheckError();\n }\n\n sdkDeleteTimer(&stopWatch);\n\n if (hostBuffer) {\n hipHostFree(hostBuffer);\n cudaCheckError();\n }\n\n hipHostFree((void *)flag);\n cudaCheckError();\n}\n\nint main(int argc, char **argv)\n{\n int numGPUs, numElems = 40000000;\n P2PDataTransfer p2p_method = P2P_WRITE;\n\n hipGetDeviceCount(&numGPUs);\n cudaCheckError();\n\n // process command line args\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\")) {\n printHelp();\n return 0;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"p2p_read\")) {\n p2p_method = P2P_READ;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"sm_copy\")) {\n p2p_mechanism = SM;\n }\n\n // number of elements of int to be used in copy.\n if (checkCmdLineFlag(argc, (const char **)argv, \"numElems\")) {\n numElems = getCmdLineArgumentInt(argc, (const char **)argv, \"numElems\");\n }\n\n printf(\"[%s]\\n\", sSampleName);\n\n // output devices\n for (int i = 0; i < numGPUs; i++) {\n hipDeviceProp_t prop;\n hipGetDeviceProperties(&prop, i);\n cudaCheckError();\n printf(\"Device: %d, %s, pciBusID: %x, pciDeviceID: %x, pciDomainID:%x\\n\",\n i,\n prop.name,\n prop.pciBusID,\n prop.pciDeviceID,\n prop.pciDomainID);\n }\n\n checkP2Paccess(numGPUs);\n\n // Environment detection: One-time check if cudaMemcpyPeerAsync is supported when P2P is disabled\n bool needsFallback = detectFallback(numGPUs);\n\n // Check peer-to-peer connectivity\n printf(\"P2P Connectivity Matrix\\n\");\n printf(\" D\\\\D\");\n\n for (int j = 0; j < numGPUs; j++) {\n printf(\"%6d\", j);\n }\n printf(\"\\n\");\n\n for (int i = 0; i < numGPUs; i++) {\n printf(\"%6d\\t\", i);\n for (int j = 0; j < numGPUs; j++) {\n if (i != j) {\n int access;\n hipDeviceCanAccessPeer(&access, i, j);\n cudaCheckError();\n printf(\"%6d\", (access) ? 1 : 0);\n }\n else {\n printf(\"%6d\", 1);\n }\n }\n printf(\"\\n\");\n }\n\n printf(\"Unidirectional P2P=Disabled Bandwidth Matrix (GB/s)\\n\");\n outputBandwidthMatrix(numElems, numGPUs, false, P2P_WRITE, needsFallback);\n printf(\"Unidirectional P2P=Enabled Bandwidth (P2P Writes) Matrix (GB/s)\\n\");\n outputBandwidthMatrix(numElems, numGPUs, true, P2P_WRITE, needsFallback);\n if (p2p_method == P2P_READ) {\n printf(\"Unidirectional P2P=Enabled Bandwidth (P2P Reads) Matrix (GB/s)\\n\");\n outputBandwidthMatrix(numElems, numGPUs, true, p2p_method, needsFallback);\n }\n printf(\"Bidirectional P2P=Disabled Bandwidth Matrix (GB/s)\\n\");\n outputBidirectionalBandwidthMatrix(numElems, numGPUs, false, needsFallback);\n printf(\"Bidirectional P2P=Enabled Bandwidth Matrix (GB/s)\\n\");\n outputBidirectionalBandwidthMatrix(numElems, numGPUs, true, needsFallback);\n\n printf(\"P2P=Disabled Latency Matrix (us)\\n\");\n outputLatencyMatrix(numGPUs, false, P2P_WRITE, needsFallback);\n printf(\"P2P=Enabled Latency (P2P Writes) Matrix (us)\\n\");\n outputLatencyMatrix(numGPUs, true, P2P_WRITE, needsFallback);\n if (p2p_method == P2P_READ) {\n printf(\"P2P=Enabled Latency (P2P Reads) Matrix (us)\\n\");\n outputLatencyMatrix(numGPUs, true, p2p_method, needsFallback);\n }\n\n printf(\"\\nNOTE: The CUDA Samples are not meant for performance measurements. \"\n \"Results may vary when GPU Boost is enabled.\\n\");\n\n exit(EXIT_SUCCESS);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/p2pBandwidthLatencyTest/p2pBandwidthLatencyTest.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n////////////////////////////////////////////////////////////////////////////////\n// Global types\n////////////////////////////////////////////////////////////////////////////////\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \n#include \n\n#include \"MonteCarlo_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Helper reduction template\n// Please see the \"reduction\" CUDA Sample for more information\n////////////////////////////////////////////////////////////////////////////////\n#include \"MonteCarlo_reduction.cuh\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Internal GPU-side data structures\n////////////////////////////////////////////////////////////////////////////////\n#define MAX_OPTIONS (1024 * 1024)\n\n// Preprocessed input option data\ntypedef struct\n{\n real S;\n real X;\n real MuByT;\n real VBySqrtT;\n} __TOptionData;\n\n////////////////////////////////////////////////////////////////////////////////\n// Overloaded shortcut payoff functions for different precision modes\n////////////////////////////////////////////////////////////////////////////////\n__device__ inline float endCallValue(float S, float X, float r, float MuByT, float VBySqrtT)\n{\n float callValue = S * __expf(MuByT + VBySqrtT * r) - X;\n return (callValue > 0.0F) ? callValue : 0.0F;\n}\n\n__device__ inline double endCallValue(double S, double X, double r, double MuByT, double VBySqrtT)\n{\n double callValue = S * exp(MuByT + VBySqrtT * r) - X;\n return (callValue > 0.0) ? callValue : 0.0;\n}\n\n#define THREAD_N 256\n\n////////////////////////////////////////////////////////////////////////////////\n// This kernel computes the integral over all paths using a single thread block\n// per option. It is fastest when the number of thread blocks times the work per\n// block is high enough to keep the GPU busy.\n////////////////////////////////////////////////////////////////////////////////\nstatic __global__ void MonteCarloOneBlockPerOption(curandState *__restrict rngStates,\n const __TOptionData *__restrict d_OptionData,\n __TOptionValue *__restrict d_CallValue,\n int pathN,\n int optionN)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n const int SUM_N = THREAD_N;\n __shared__ real s_SumCall[SUM_N];\n __shared__ real s_Sum2Call[SUM_N];\n\n // determine global thread id\n int tid = threadIdx.x + blockIdx.x * blockDim.x;\n\n // Copy random number state to local memory for efficiency\n curandState localState = rngStates[tid];\n for (int optionIndex = blockIdx.x; optionIndex < optionN; optionIndex += gridDim.x) {\n const real S = d_OptionData[optionIndex].S;\n const real X = d_OptionData[optionIndex].X;\n const real MuByT = d_OptionData[optionIndex].MuByT;\n const real VBySqrtT = d_OptionData[optionIndex].VBySqrtT;\n\n // Cycle through the entire samples array:\n // derive end stock price for each path\n // accumulate partial integrals into intermediate shared memory buffer\n for (int iSum = threadIdx.x; iSum < SUM_N; iSum += blockDim.x) {\n __TOptionValue sumCall = {0, 0};\n\n#pragma unroll 8\n for (int i = iSum; i < pathN; i += SUM_N) {\n real r = curand_normal(&localState);\n real callValue = endCallValue(S, X, r, MuByT, VBySqrtT);\n sumCall.Expected += callValue;\n sumCall.Confidence += callValue * callValue;\n }\n\n s_SumCall[iSum] = sumCall.Expected;\n s_Sum2Call[iSum] = sumCall.Confidence;\n }\n\n // Reduce shared memory accumulators\n // and write final result to global memory\n cg::sync(cta);\n sumReduce(s_SumCall, s_Sum2Call, cta, tile32, &d_CallValue[optionIndex]);\n }\n}\n\nstatic __global__ void rngSetupStates(curandState *rngState, int device_id)\n{\n // determine global thread id\n int tid = threadIdx.x + blockIdx.x * blockDim.x;\n // Each threadblock gets different seed,\n // Threads within a threadblock get different sequence numbers\n curand_init(blockIdx.x + gridDim.x * device_id, threadIdx.x, 0, &rngState[tid]);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Host-side interface to GPU Monte Carlo\n////////////////////////////////////////////////////////////////////////////////\n\nextern \"C\" void initMonteCarloGPU(TOptionPlan *plan)\n{\n checkCudaErrors(cudaMalloc(&plan->d_OptionData, sizeof(__TOptionData) * (plan->optionCount)));\n checkCudaErrors(cudaMalloc(&plan->d_CallValue, sizeof(__TOptionValue) * (plan->optionCount)));\n checkCudaErrors(cudaMallocHost(&plan->h_OptionData, sizeof(__TOptionData) * (plan->optionCount)));\n // Allocate internal device memory\n checkCudaErrors(cudaMallocHost(&plan->h_CallValue, sizeof(__TOptionValue) * (plan->optionCount)));\n // Allocate states for pseudo random number generators\n checkCudaErrors(cudaMalloc((void **)&plan->rngStates, plan->gridSize * THREAD_N * sizeof(curandState)));\n checkCudaErrors(cudaMemset(plan->rngStates, 0, plan->gridSize * THREAD_N * sizeof(curandState)));\n\n // place each device pathN random numbers apart on the random number sequence\n rngSetupStates<<gridSize, THREAD_N>>>(plan->rngStates, plan->device);\n getLastCudaError(\"rngSetupStates kernel failed.\\n\");\n}\n\n// Compute statistics and deallocate internal device memory\nextern \"C\" void closeMonteCarloGPU(TOptionPlan *plan)\n{\n for (int i = 0; i < plan->optionCount; i++) {\n const double RT = plan->optionData[i].R * plan->optionData[i].T;\n const double sum = plan->h_CallValue[i].Expected;\n const double sum2 = plan->h_CallValue[i].Confidence;\n const double pathN = plan->pathN;\n // Derive average from the total sum and discount by riskfree rate\n plan->callValue[i].Expected = (float)(exp(-RT) * sum / pathN);\n // Standard deviation\n double stdDev = sqrt((pathN * sum2 - sum * sum) / (pathN * (pathN - 1)));\n // Confidence width; in 95% of all cases theoretical value lies within these\n // borders\n plan->callValue[i].Confidence = (float)(exp(-RT) * 1.96 * stdDev / sqrt(pathN));\n }\n\n checkCudaErrors(cudaFree(plan->rngStates));\n checkCudaErrors(cudaFreeHost(plan->h_CallValue));\n checkCudaErrors(cudaFreeHost(plan->h_OptionData));\n checkCudaErrors(cudaFree(plan->d_CallValue));\n checkCudaErrors(cudaFree(plan->d_OptionData));\n}\n\n// Main computations\nextern \"C\" void MonteCarloGPU(TOptionPlan *plan, cudaStream_t stream)\n{\n __TOptionValue *h_CallValue = plan->h_CallValue;\n\n if (plan->optionCount <= 0 || plan->optionCount > MAX_OPTIONS) {\n printf(\"MonteCarloGPU(): bad option count.\\n\");\n return;\n }\n\n __TOptionData *h_OptionData = (__TOptionData *)plan->h_OptionData;\n\n for (int i = 0; i < plan->optionCount; i++) {\n const double T = plan->optionData[i].T;\n const double R = plan->optionData[i].R;\n const double V = plan->optionData[i].V;\n const double MuByT = (R - 0.5 * V * V) * T;\n const double VBySqrtT = V * sqrt(T);\n h_OptionData[i].S = (real)plan->optionData[i].S;\n h_OptionData[i].X = (real)plan->optionData[i].X;\n h_OptionData[i].MuByT = (real)MuByT;\n h_OptionData[i].VBySqrtT = (real)VBySqrtT;\n }\n\n checkCudaErrors(cudaMemcpyAsync(\n plan->d_OptionData, h_OptionData, plan->optionCount * sizeof(__TOptionData), cudaMemcpyHostToDevice, stream));\n\n MonteCarloOneBlockPerOption<<gridSize, THREAD_N, 0, stream>>>(plan->rngStates,\n (__TOptionData *)(plan->d_OptionData),\n (__TOptionValue *)(plan->d_CallValue),\n plan->pathN,\n plan->optionCount);\n getLastCudaError(\"MonteCarloOneBlockPerOption() execution failed\\n\");\n\n checkCudaErrors(cudaMemcpyAsync(\n h_CallValue, plan->d_CallValue, plan->optionCount * sizeof(__TOptionValue), cudaMemcpyDeviceToHost, stream));\n\n // cudaDeviceSynchronize();\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n////////////////////////////////////////////////////////////////////////////////\n// Global types\n////////////////////////////////////////////////////////////////////////////////\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \n#include \n\n#include \"MonteCarlo_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Helper reduction template\n// Please see the \"reduction\" CUDA Sample for more information\n////////////////////////////////////////////////////////////////////////////////\n#include \"MonteCarlo_reduction.cuh\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Internal GPU-side data structures\n////////////////////////////////////////////////////////////////////////////////\n#define MAX_OPTIONS (1024 * 1024)\n\n// Preprocessed input option data\ntypedef struct\n{\n real S;\n real X;\n real MuByT;\n real VBySqrtT;\n} __TOptionData;\n\n////////////////////////////////////////////////////////////////////////////////\n// Overloaded shortcut payoff functions for different precision modes\n////////////////////////////////////////////////////////////////////////////////\n__device__ inline float endCallValue(float S, float X, float r, float MuByT, float VBySqrtT)\n{\n float callValue = S * __expf(MuByT + VBySqrtT * r) - X;\n return (callValue > 0.0F) ? callValue : 0.0F;\n}\n\n__device__ inline double endCallValue(double S, double X, double r, double MuByT, double VBySqrtT)\n{\n double callValue = S * exp(MuByT + VBySqrtT * r) - X;\n return (callValue > 0.0) ? callValue : 0.0;\n}\n\n#define THREAD_N 256\n\n////////////////////////////////////////////////////////////////////////////////\n// This kernel computes the integral over all paths using a single thread block\n// per option. It is fastest when the number of thread blocks times the work per\n// block is high enough to keep the GPU busy.\n////////////////////////////////////////////////////////////////////////////////\nstatic __global__ void MonteCarloOneBlockPerOption(hiprandState *__restrict rngStates,\n const __TOptionData *__restrict d_OptionData,\n __TOptionValue *__restrict d_CallValue,\n int pathN,\n int optionN)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n const int SUM_N = THREAD_N;\n __shared__ real s_SumCall[SUM_N];\n __shared__ real s_Sum2Call[SUM_N];\n\n // determine global thread id\n int tid = threadIdx.x + blockIdx.x * blockDim.x;\n\n // Copy random number state to local memory for efficiency\n hiprandState localState = rngStates[tid];\n for (int optionIndex = blockIdx.x; optionIndex < optionN; optionIndex += gridDim.x) {\n const real S = d_OptionData[optionIndex].S;\n const real X = d_OptionData[optionIndex].X;\n const real MuByT = d_OptionData[optionIndex].MuByT;\n const real VBySqrtT = d_OptionData[optionIndex].VBySqrtT;\n\n // Cycle through the entire samples array:\n // derive end stock price for each path\n // accumulate partial integrals into intermediate shared memory buffer\n for (int iSum = threadIdx.x; iSum < SUM_N; iSum += blockDim.x) {\n __TOptionValue sumCall = {0, 0};\n\n#pragma unroll 8\n for (int i = iSum; i < pathN; i += SUM_N) {\n real r = hiprand_normal(&localState);\n real callValue = endCallValue(S, X, r, MuByT, VBySqrtT);\n sumCall.Expected += callValue;\n sumCall.Confidence += callValue * callValue;\n }\n\n s_SumCall[iSum] = sumCall.Expected;\n s_Sum2Call[iSum] = sumCall.Confidence;\n }\n\n // Reduce shared memory accumulators\n // and write final result to global memory\n cg::sync(cta);\n sumReduce(s_SumCall, s_Sum2Call, cta, tile32, &d_CallValue[optionIndex]);\n }\n}\n\nstatic __global__ void rngSetupStates(hiprandState *rngState, int device_id)\n{\n // determine global thread id\n int tid = threadIdx.x + blockIdx.x * blockDim.x;\n // Each threadblock gets different seed,\n // Threads within a threadblock get different sequence numbers\n hiprand_init(blockIdx.x + gridDim.x * device_id, threadIdx.x, 0, &rngState[tid]);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Host-side interface to GPU Monte Carlo\n////////////////////////////////////////////////////////////////////////////////\n\nextern \"C\" void initMonteCarloGPU(TOptionPlan *plan)\n{\n checkCudaErrors(hipMalloc(&plan->d_OptionData, sizeof(__TOptionData) * (plan->optionCount)));\n checkCudaErrors(hipMalloc(&plan->d_CallValue, sizeof(__TOptionValue) * (plan->optionCount)));\n checkCudaErrors(hipHostMalloc(&plan->h_OptionData, sizeof(__TOptionData) * (plan->optionCount)));\n // Allocate internal device memory\n checkCudaErrors(hipHostMalloc(&plan->h_CallValue, sizeof(__TOptionValue) * (plan->optionCount)));\n // Allocate states for pseudo random number generators\n checkCudaErrors(hipMalloc((void **)&plan->rngStates, plan->gridSize * THREAD_N * sizeof(hiprandState)));\n checkCudaErrors(hipMemset(plan->rngStates, 0, plan->gridSize * THREAD_N * sizeof(hiprandState)));\n\n // place each device pathN random numbers apart on the random number sequence\n rngSetupStates<<gridSize, THREAD_N>>>(plan->rngStates, plan->device);\n getLastCudaError(\"rngSetupStates kernel failed.\\n\");\n}\n\n// Compute statistics and deallocate internal device memory\nextern \"C\" void closeMonteCarloGPU(TOptionPlan *plan)\n{\n for (int i = 0; i < plan->optionCount; i++) {\n const double RT = plan->optionData[i].R * plan->optionData[i].T;\n const double sum = plan->h_CallValue[i].Expected;\n const double sum2 = plan->h_CallValue[i].Confidence;\n const double pathN = plan->pathN;\n // Derive average from the total sum and discount by riskfree rate\n plan->callValue[i].Expected = (float)(exp(-RT) * sum / pathN);\n // Standard deviation\n double stdDev = sqrt((pathN * sum2 - sum * sum) / (pathN * (pathN - 1)));\n // Confidence width; in 95% of all cases theoretical value lies within these\n // borders\n plan->callValue[i].Confidence = (float)(exp(-RT) * 1.96 * stdDev / sqrt(pathN));\n }\n\n checkCudaErrors(hipFree(plan->rngStates));\n checkCudaErrors(hipHostFree(plan->h_CallValue));\n checkCudaErrors(hipHostFree(plan->h_OptionData));\n checkCudaErrors(hipFree(plan->d_CallValue));\n checkCudaErrors(hipFree(plan->d_OptionData));\n}\n\n// Main computations\nextern \"C\" void MonteCarloGPU(TOptionPlan *plan, hipStream_t stream)\n{\n __TOptionValue *h_CallValue = plan->h_CallValue;\n\n if (plan->optionCount <= 0 || plan->optionCount > MAX_OPTIONS) {\n printf(\"MonteCarloGPU(): bad option count.\\n\");\n return;\n }\n\n __TOptionData *h_OptionData = (__TOptionData *)plan->h_OptionData;\n\n for (int i = 0; i < plan->optionCount; i++) {\n const double T = plan->optionData[i].T;\n const double R = plan->optionData[i].R;\n const double V = plan->optionData[i].V;\n const double MuByT = (R - 0.5 * V * V) * T;\n const double VBySqrtT = V * sqrt(T);\n h_OptionData[i].S = (real)plan->optionData[i].S;\n h_OptionData[i].X = (real)plan->optionData[i].X;\n h_OptionData[i].MuByT = (real)MuByT;\n h_OptionData[i].VBySqrtT = (real)VBySqrtT;\n }\n\n checkCudaErrors(hipMemcpyAsync(\n plan->d_OptionData, h_OptionData, plan->optionCount * sizeof(__TOptionData), hipMemcpyHostToDevice, stream));\n\n MonteCarloOneBlockPerOption<<gridSize, THREAD_N, 0, stream>>>(plan->rngStates,\n (__TOptionData *)(plan->d_OptionData),\n (__TOptionValue *)(plan->d_CallValue),\n plan->pathN,\n plan->optionCount);\n getLastCudaError(\"MonteCarloOneBlockPerOption() execution failed\\n\");\n\n checkCudaErrors(hipMemcpyAsync(\n h_CallValue, plan->d_CallValue, plan->optionCount * sizeof(__TOptionValue), hipMemcpyDeviceToHost, stream));\n\n // cudaDeviceSynchronize();\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/MonteCarloMultiGPU/MonteCarlo_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Utilities and system includes\n#include \n\nnamespace cg = cooperative_groups;\n\n#include // for FLT_MAX\n#include \n#include \n#include \n\n#include \"CudaMath.h\"\n#include \"dds.h\"\n#include \"permutations.h\"\n\n// Definitions\n#define INPUT_IMAGE \"teapot512_std.ppm\"\n#define REFERENCE_IMAGE \"teapot512_ref.dds\"\n\n#define ERROR_THRESHOLD 0.02f\n\n#define NUM_THREADS 64 // Number of threads per block.\n\n#define __debugsync()\n\ntemplate __device__ inline void swap(T &a, T &b)\n{\n T tmp = a;\n a = b;\n b = tmp;\n}\n\n//__constant__ float3 kColorMetric = { 0.2126f, 0.7152f, 0.0722f };\n__constant__ float3 kColorMetric = {1.0f, 1.0f, 1.0f};\n\n////////////////////////////////////////////////////////////////////////////////\n// Sort colors\n////////////////////////////////////////////////////////////////////////////////\n__device__ void sortColors(const float *values, int *ranks, cg::thread_group tile)\n{\n const int tid = threadIdx.x;\n\n int rank = 0;\n\n#pragma unroll\n\n for (int i = 0; i < 16; i++) {\n rank += (values[i] < values[tid]);\n }\n\n ranks[tid] = rank;\n\n cg::sync(tile);\n\n // Resolve elements with the same index.\n for (int i = 0; i < 15; i++) {\n if (tid > i && ranks[tid] == ranks[i]) {\n ++ranks[tid];\n }\n cg::sync(tile);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Load color block to shared mem\n////////////////////////////////////////////////////////////////////////////////\n__device__ void loadColorBlock(const uint *image,\n float3 colors[16],\n float3 sums[16],\n int xrefs[16],\n int blockOffset,\n cg::thread_block cta)\n{\n const int bid = blockIdx.x + blockOffset;\n const int idx = threadIdx.x;\n\n __shared__ float dps[16];\n\n float3 tmp;\n\n cg::thread_group tile = cg::tiled_partition(cta, 16);\n\n if (idx < 16) {\n // Read color and copy to shared mem.\n uint c = image[(bid) * 16 + idx];\n\n colors[idx].x = ((c >> 0) & 0xFF) * (1.0f / 255.0f);\n colors[idx].y = ((c >> 8) & 0xFF) * (1.0f / 255.0f);\n colors[idx].z = ((c >> 16) & 0xFF) * (1.0f / 255.0f);\n\n cg::sync(tile);\n // Sort colors along the best fit line.\n colorSums(colors, sums, tile);\n\n cg::sync(tile);\n\n float3 axis = bestFitLine(colors, sums[0], tile);\n\n cg::sync(tile);\n\n dps[idx] = dot(colors[idx], axis);\n\n cg::sync(tile);\n\n sortColors(dps, xrefs, tile);\n\n cg::sync(tile);\n\n tmp = colors[idx];\n\n cg::sync(tile);\n\n colors[xrefs[idx]] = tmp;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Round color to RGB565 and expand\n////////////////////////////////////////////////////////////////////////////////\ninline __device__ float3 roundAndExpand(float3 v, ushort *w)\n{\n v.x = rintf(__saturatef(v.x) * 31.0f);\n v.y = rintf(__saturatef(v.y) * 63.0f);\n v.z = rintf(__saturatef(v.z) * 31.0f);\n\n *w = ((ushort)v.x << 11) | ((ushort)v.y << 5) | (ushort)v.z;\n v.x *= 0.03227752766457f; // approximate integer bit expansion.\n v.y *= 0.01583151765563f;\n v.z *= 0.03227752766457f;\n return v;\n}\n\n__constant__ float alphaTable4[4] = {9.0f, 0.0f, 6.0f, 3.0f};\n__constant__ float alphaTable3[4] = {4.0f, 0.0f, 2.0f, 2.0f};\n__constant__ const int prods4[4] = {0x090000, 0x000900, 0x040102, 0x010402};\n__constant__ const int prods3[4] = {0x040000, 0x000400, 0x040101, 0x010401};\n\n#define USE_TABLES 1\n\n////////////////////////////////////////////////////////////////////////////////\n// Evaluate permutations\n////////////////////////////////////////////////////////////////////////////////\nstatic __device__ float\nevalPermutation4(const float3 *colors, uint permutation, ushort *start, ushort *end, float3 color_sum)\n{\n// Compute endpoints using least squares.\n#if USE_TABLES\n float3 alphax_sum = make_float3(0.0f, 0.0f, 0.0f);\n\n int akku = 0;\n\n // Compute alpha & beta for this permutation.\n for (int i = 0; i < 16; i++) {\n const uint bits = permutation >> (2 * i);\n\n alphax_sum += alphaTable4[bits & 3] * colors[i];\n akku += prods4[bits & 3];\n }\n\n float alpha2_sum = float(akku >> 16);\n float beta2_sum = float((akku >> 8) & 0xff);\n float alphabeta_sum = float((akku >> 0) & 0xff);\n float3 betax_sum = (9.0f * color_sum) - alphax_sum;\n#else\n float alpha2_sum = 0.0f;\n float beta2_sum = 0.0f;\n float alphabeta_sum = 0.0f;\n float3 alphax_sum = make_float3(0.0f, 0.0f, 0.0f);\n\n // Compute alpha & beta for this permutation.\n for (int i = 0; i < 16; i++) {\n const uint bits = permutation >> (2 * i);\n\n float beta = (bits & 1);\n\n if (bits & 2) {\n beta = (1 + beta) * (1.0f / 3.0f);\n }\n\n float alpha = 1.0f - beta;\n\n alpha2_sum += alpha * alpha;\n beta2_sum += beta * beta;\n alphabeta_sum += alpha * beta;\n alphax_sum += alpha * colors[i];\n }\n\n float3 betax_sum = color_sum - alphax_sum;\n#endif\n\n // alpha2, beta2, alphabeta and factor could be precomputed for each\n // permutation, but it's faster to recompute them.\n const float factor = 1.0f / (alpha2_sum * beta2_sum - alphabeta_sum * alphabeta_sum);\n\n float3 a = (alphax_sum * beta2_sum - betax_sum * alphabeta_sum) * factor;\n float3 b = (betax_sum * alpha2_sum - alphax_sum * alphabeta_sum) * factor;\n\n // Round a, b to the closest 5-6-5 color and expand...\n a = roundAndExpand(a, start);\n b = roundAndExpand(b, end);\n\n // compute the error\n float3 e = a * a * alpha2_sum + b * b * beta2_sum + 2.0f * (a * b * alphabeta_sum - a * alphax_sum - b * betax_sum);\n\n return (0.111111111111f) * dot(e, kColorMetric);\n}\n\nstatic __device__ float\nevalPermutation3(const float3 *colors, uint permutation, ushort *start, ushort *end, float3 color_sum)\n{\n// Compute endpoints using least squares.\n#if USE_TABLES\n float3 alphax_sum = make_float3(0.0f, 0.0f, 0.0f);\n\n int akku = 0;\n\n // Compute alpha & beta for this permutation.\n for (int i = 0; i < 16; i++) {\n const uint bits = permutation >> (2 * i);\n\n alphax_sum += alphaTable3[bits & 3] * colors[i];\n akku += prods3[bits & 3];\n }\n\n float alpha2_sum = float(akku >> 16);\n float beta2_sum = float((akku >> 8) & 0xff);\n float alphabeta_sum = float((akku >> 0) & 0xff);\n float3 betax_sum = (4.0f * color_sum) - alphax_sum;\n#else\n float alpha2_sum = 0.0f;\n float beta2_sum = 0.0f;\n float alphabeta_sum = 0.0f;\n float3 alphax_sum = make_float3(0.0f, 0.0f, 0.0f);\n\n // Compute alpha & beta for this permutation.\n for (int i = 0; i < 16; i++) {\n const uint bits = permutation >> (2 * i);\n\n float beta = (bits & 1);\n\n if (bits & 2) {\n beta = 0.5f;\n }\n\n float alpha = 1.0f - beta;\n\n alpha2_sum += alpha * alpha;\n beta2_sum += beta * beta;\n alphabeta_sum += alpha * beta;\n alphax_sum += alpha * colors[i];\n }\n\n float3 betax_sum = color_sum - alphax_sum;\n#endif\n\n const float factor = 1.0f / (alpha2_sum * beta2_sum - alphabeta_sum * alphabeta_sum);\n\n float3 a = (alphax_sum * beta2_sum - betax_sum * alphabeta_sum) * factor;\n float3 b = (betax_sum * alpha2_sum - alphax_sum * alphabeta_sum) * factor;\n\n // Round a, b to the closest 5-6-5 color and expand...\n a = roundAndExpand(a, start);\n b = roundAndExpand(b, end);\n\n // compute the error\n float3 e = a * a * alpha2_sum + b * b * beta2_sum + 2.0f * (a * b * alphabeta_sum - a * alphax_sum - b * betax_sum);\n\n return (0.25f) * dot(e, kColorMetric);\n}\n\n__device__ void evalAllPermutations(const float3 *colors,\n const uint *permutations,\n ushort &bestStart,\n ushort &bestEnd,\n uint &bestPermutation,\n float *errors,\n float3 color_sum,\n cg::thread_block cta)\n{\n const int idx = threadIdx.x;\n\n float bestError = FLT_MAX;\n\n __shared__ uint s_permutations[160];\n\n for (int i = 0; i < 16; i++) {\n int pidx = idx + NUM_THREADS * i;\n\n if (pidx >= 992) {\n break;\n }\n\n ushort start, end;\n uint permutation = permutations[pidx];\n\n if (pidx < 160) {\n s_permutations[pidx] = permutation;\n }\n\n float error = evalPermutation4(colors, permutation, &start, &end, color_sum);\n\n if (error < bestError) {\n bestError = error;\n bestPermutation = permutation;\n bestStart = start;\n bestEnd = end;\n }\n }\n\n if (bestStart < bestEnd) {\n swap(bestEnd, bestStart);\n bestPermutation ^= 0x55555555; // Flip indices.\n }\n\n cg::sync(cta); // Sync here to ensure s_permutations is valid going forward\n\n for (int i = 0; i < 3; i++) {\n int pidx = idx + NUM_THREADS * i;\n\n if (pidx >= 160) {\n break;\n }\n\n ushort start, end;\n uint permutation = s_permutations[pidx];\n float error = evalPermutation3(colors, permutation, &start, &end, color_sum);\n\n if (error < bestError) {\n bestError = error;\n bestPermutation = permutation;\n bestStart = start;\n bestEnd = end;\n\n if (bestStart > bestEnd) {\n swap(bestEnd, bestStart);\n bestPermutation ^= (~bestPermutation >> 1) & 0x55555555; // Flip indices.\n }\n }\n }\n\n errors[idx] = bestError;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Find index with minimum error\n////////////////////////////////////////////////////////////////////////////////\n__device__ int findMinError(float *errors, cg::thread_block cta)\n{\n const int idx = threadIdx.x;\n __shared__ int indices[NUM_THREADS];\n indices[idx] = idx;\n\n cg::sync(cta);\n\n for (int d = NUM_THREADS / 2; d > 0; d >>= 1) {\n float err0 = errors[idx];\n float err1 = (idx + d) < NUM_THREADS ? errors[idx + d] : FLT_MAX;\n int index1 = (idx + d) < NUM_THREADS ? indices[idx + d] : 0;\n\n cg::sync(cta);\n\n if (err1 < err0) {\n errors[idx] = err1;\n indices[idx] = index1;\n }\n\n cg::sync(cta);\n }\n\n return indices[0];\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Save DXT block\n////////////////////////////////////////////////////////////////////////////////\n__device__ void saveBlockDXT1(ushort start, ushort end, uint permutation, int xrefs[16], uint2 *result, int blockOffset)\n{\n const int bid = blockIdx.x + blockOffset;\n\n if (start == end) {\n permutation = 0;\n }\n\n // Reorder permutation.\n uint indices = 0;\n\n for (int i = 0; i < 16; i++) {\n int ref = xrefs[i];\n indices |= ((permutation >> (2 * ref)) & 3) << (2 * i);\n }\n\n // Write endpoints.\n result[bid].x = (end << 16) | start;\n\n // Write palette indices.\n result[bid].y = indices;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Compress color block\n////////////////////////////////////////////////////////////////////////////////\n__global__ void compress(const uint *permutations, const uint *image, uint2 *result, int blockOffset)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n\n const int idx = threadIdx.x;\n\n __shared__ float3 colors[16];\n __shared__ float3 sums[16];\n __shared__ int xrefs[16];\n\n loadColorBlock(image, colors, sums, xrefs, blockOffset, cta);\n\n cg::sync(cta);\n\n ushort bestStart, bestEnd;\n uint bestPermutation;\n\n __shared__ float errors[NUM_THREADS];\n\n evalAllPermutations(colors, permutations, bestStart, bestEnd, bestPermutation, errors, sums[0], cta);\n\n // Use a parallel reduction to find minimum error.\n const int minIdx = findMinError(errors, cta);\n\n cg::sync(cta);\n\n // Only write the result of the winner thread.\n if (idx == minIdx) {\n saveBlockDXT1(bestStart, bestEnd, bestPermutation, xrefs, result, blockOffset);\n }\n}\n\n// Helper structs and functions to validate the output of the compressor.\n// We cannot simply do a bitwise compare, because different compilers produce\n// different\n// results for different targets due to floating point arithmetic.\n\nunion Color32\n{\n struct\n {\n unsigned char b, g, r, a;\n };\n unsigned int u;\n};\n\nunion Color16\n{\n struct\n {\n unsigned short b : 5;\n unsigned short g : 6;\n unsigned short r : 5;\n };\n unsigned short u;\n};\n\nstruct BlockDXT1\n{\n Color16 col0;\n Color16 col1;\n union\n {\n unsigned char row[4];\n unsigned int indices;\n };\n\n void decompress(Color32 colors[16]) const;\n};\n\nvoid BlockDXT1::decompress(Color32 *colors) const\n{\n Color32 palette[4];\n\n // Does bit expansion before interpolation.\n palette[0].b = (col0.b << 3) | (col0.b >> 2);\n palette[0].g = (col0.g << 2) | (col0.g >> 4);\n palette[0].r = (col0.r << 3) | (col0.r >> 2);\n palette[0].a = 0xFF;\n\n palette[1].r = (col1.r << 3) | (col1.r >> 2);\n palette[1].g = (col1.g << 2) | (col1.g >> 4);\n palette[1].b = (col1.b << 3) | (col1.b >> 2);\n palette[1].a = 0xFF;\n\n if (col0.u > col1.u) {\n // Four-color block: derive the other two colors.\n palette[2].r = (2 * palette[0].r + palette[1].r) / 3;\n palette[2].g = (2 * palette[0].g + palette[1].g) / 3;\n palette[2].b = (2 * palette[0].b + palette[1].b) / 3;\n palette[2].a = 0xFF;\n\n palette[3].r = (2 * palette[1].r + palette[0].r) / 3;\n palette[3].g = (2 * palette[1].g + palette[0].g) / 3;\n palette[3].b = (2 * palette[1].b + palette[0].b) / 3;\n palette[3].a = 0xFF;\n }\n else {\n // Three-color block: derive the other color.\n palette[2].r = (palette[0].r + palette[1].r) / 2;\n palette[2].g = (palette[0].g + palette[1].g) / 2;\n palette[2].b = (palette[0].b + palette[1].b) / 2;\n palette[2].a = 0xFF;\n\n palette[3].r = 0x00;\n palette[3].g = 0x00;\n palette[3].b = 0x00;\n palette[3].a = 0x00;\n }\n\n for (int i = 0; i < 16; i++) {\n colors[i] = palette[(indices >> (2 * i)) & 0x3];\n }\n}\n\nstatic int compareColors(const Color32 *b0, const Color32 *b1)\n{\n int sum = 0;\n\n for (int i = 0; i < 16; i++) {\n int r = (b0[i].r - b1[i].r);\n int g = (b0[i].g - b1[i].g);\n int b = (b0[i].b - b1[i].b);\n sum += r * r + g * g + b * b;\n }\n\n return sum;\n}\n\nstatic int compareBlock(const BlockDXT1 *b0, const BlockDXT1 *b1)\n{\n Color32 colors0[16];\n Color32 colors1[16];\n\n if (memcmp(b0, b1, sizeof(BlockDXT1)) == 0) {\n return 0;\n }\n else {\n b0->decompress(colors0);\n b1->decompress(colors1);\n\n return compareColors(colors0, colors1);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"%s Starting...\\n\\n\", argv[0]);\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n findCudaDevice(argc, (const char **)argv);\n\n // Load input image.\n unsigned char *data = NULL;\n uint W, H;\n\n char *image_path = sdkFindFilePath(INPUT_IMAGE, argv[0]);\n\n if (image_path == 0) {\n printf(\"Error, unable to find source image <%s>\\n\", image_path);\n exit(EXIT_FAILURE);\n }\n\n if (!sdkLoadPPM4ub(image_path, &data, &W, &H)) {\n printf(\"Error, unable to open source image file <%s>\\n\", image_path);\n\n exit(EXIT_FAILURE);\n }\n\n uint w = W, h = H;\n\n printf(\"Image Loaded '%s', %d x %d pixels\\n\\n\", image_path, w, h);\n\n // Allocate input image.\n const uint memSize = w * h * 4;\n assert(0 != memSize);\n uint *block_image = (uint *)malloc(memSize);\n\n // Convert linear image to block linear.\n for (uint by = 0; by < h / 4; by++) {\n for (uint bx = 0; bx < w / 4; bx++) {\n for (int i = 0; i < 16; i++) {\n const int x = i & 3;\n const int y = i / 4;\n block_image[(by * w / 4 + bx) * 16 + i] = ((uint *)data)[(by * 4 + y) * 4 * (W / 4) + bx * 4 + x];\n }\n }\n }\n\n // copy into global mem\n uint *d_data = NULL;\n checkCudaErrors(cudaMalloc((void **)&d_data, memSize));\n\n // Result\n uint *d_result = NULL;\n const uint compressedSize = (w / 4) * (h / 4) * 8;\n checkCudaErrors(cudaMalloc((void **)&d_result, compressedSize));\n uint *h_result = (uint *)malloc(compressedSize);\n\n // Compute permutations.\n uint permutations[1024];\n computePermutations(permutations);\n\n // Copy permutations host to devie.\n uint *d_permutations = NULL;\n checkCudaErrors(cudaMalloc((void **)&d_permutations, 1024 * sizeof(uint)));\n checkCudaErrors(cudaMemcpy(d_permutations, permutations, 1024 * sizeof(uint), cudaMemcpyHostToDevice));\n\n // create a timer\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n\n // Copy image from host to device\n checkCudaErrors(cudaMemcpy(d_data, block_image, memSize, cudaMemcpyHostToDevice));\n\n // Determine launch configuration and run timed computation numIterations\n // times\n uint blocks = ((w + 3) / 4) * ((h + 3) / 4); // rounds up by 1 block in each dim if %4 != 0\n\n int devID;\n cudaDeviceProp deviceProp;\n\n // get number of SMs on this GPU\n checkCudaErrors(cudaGetDevice(&devID));\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, devID));\n\n // Restrict the numbers of blocks to launch on low end GPUs to avoid kernel\n // timeout\n int blocksPerLaunch = min(blocks, 768 * deviceProp.multiProcessorCount);\n\n printf(\"Running DXT Compression on %u x %u image...\\n\", w, h);\n printf(\"\\n%u Blocks, %u Threads per Block, %u Threads in Grid...\\n\\n\", blocks, NUM_THREADS, blocks * NUM_THREADS);\n int numIterations = 1;\n\n for (int i = -1; i < numIterations; ++i) {\n if (i == 0) {\n checkCudaErrors(cudaDeviceSynchronize());\n sdkStartTimer(&timer);\n }\n\n for (int j = 0; j < (int)blocks; j += blocksPerLaunch) {\n compress<<>>(d_permutations, d_data, (uint2 *)d_result, j);\n }\n }\n\n getLastCudaError(\"compress\");\n\n // sync to host, stop timer, record perf\n checkCudaErrors(cudaDeviceSynchronize());\n sdkStopTimer(&timer);\n double dAvgTime = 1.0e-3 * sdkGetTimerValue(&timer) / (double)numIterations;\n printf(\"dxtc, Throughput = %.4f MPixels/s, Time = %.5f s, Size = %u Pixels, \"\n \"NumDevsUsed = %i, Workgroup = %d\\n\",\n (1.0e-6 * (double)(W * H) / dAvgTime),\n dAvgTime,\n (W * H),\n 1,\n NUM_THREADS);\n\n // copy result data from device to host\n checkCudaErrors(cudaMemcpy(h_result, d_result, compressedSize, cudaMemcpyDeviceToHost));\n\n // Write out result data to DDS file\n char output_filename[1024];\n strcpy(output_filename, image_path);\n strcpy(output_filename + strlen(image_path) - 3, \"dds\");\n FILE *fp = fopen(output_filename, \"wb\");\n\n if (fp == 0) {\n printf(\"Error, unable to open output image <%s>\\n\", output_filename);\n exit(EXIT_FAILURE);\n }\n\n DDSHeader header;\n header.fourcc = FOURCC_DDS;\n header.size = 124;\n header.flags = (DDSD_WIDTH | DDSD_HEIGHT | DDSD_CAPS | DDSD_PIXELFORMAT | DDSD_LINEARSIZE);\n header.height = h;\n header.width = w;\n header.pitch = compressedSize;\n header.depth = 0;\n header.mipmapcount = 0;\n memset(header.reserved, 0, sizeof(header.reserved));\n header.pf.size = 32;\n header.pf.flags = DDPF_FOURCC;\n header.pf.fourcc = FOURCC_DXT1;\n header.pf.bitcount = 0;\n header.pf.rmask = 0;\n header.pf.gmask = 0;\n header.pf.bmask = 0;\n header.pf.amask = 0;\n header.caps.caps1 = DDSCAPS_TEXTURE;\n header.caps.caps2 = 0;\n header.caps.caps3 = 0;\n header.caps.caps4 = 0;\n header.notused = 0;\n fwrite(&header, sizeof(DDSHeader), 1, fp);\n fwrite(h_result, compressedSize, 1, fp);\n fclose(fp);\n\n // Make sure the generated image is correct.\n const char *reference_image_path = sdkFindFilePath(REFERENCE_IMAGE, argv[0]);\n\n if (reference_image_path == 0) {\n printf(\"Error, unable to find reference image\\n\");\n\n exit(EXIT_FAILURE);\n }\n\n fp = fopen(reference_image_path, \"rb\");\n\n if (fp == 0) {\n printf(\"Error, unable to open reference image\\n\");\n\n exit(EXIT_FAILURE);\n }\n\n fseek(fp, sizeof(DDSHeader), SEEK_SET);\n uint referenceSize = (W / 4) * (H / 4) * 8;\n uint *reference = (uint *)malloc(referenceSize);\n fread(reference, referenceSize, 1, fp);\n fclose(fp);\n\n printf(\"\\nChecking accuracy...\\n\");\n float rms = 0;\n\n for (uint y = 0; y < h; y += 4) {\n for (uint x = 0; x < w; x += 4) {\n uint referenceBlockIdx = ((y / 4) * (W / 4) + (x / 4));\n uint resultBlockIdx = ((y / 4) * (w / 4) + (x / 4));\n\n int cmp =\n compareBlock(((BlockDXT1 *)h_result) + resultBlockIdx, ((BlockDXT1 *)reference) + referenceBlockIdx);\n\n if (cmp != 0.0f) {\n printf(\"Deviation at (%4d,%4d):\\t%f rms\\n\", x / 4, y / 4, float(cmp) / 16 / 3);\n }\n\n rms += cmp;\n }\n }\n\n rms /= w * h * 3;\n\n // Free allocated resources and exit\n checkCudaErrors(cudaFree(d_permutations));\n checkCudaErrors(cudaFree(d_data));\n checkCudaErrors(cudaFree(d_result));\n free(image_path);\n free(data);\n free(block_image);\n free(h_result);\n free(reference);\n sdkDeleteTimer(&timer);\n\n printf(\"RMS(reference, result) = %f\\n\\n\", rms);\n printf(rms <= ERROR_THRESHOLD ? \"Test passed\\n\" : \"Test failed!\\n\");\n /* Return zero if test passed, one otherwise */\n return rms > ERROR_THRESHOLD;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Utilities and system includes\n#include \"hip/hip_runtime.h\"\n#include \n\nnamespace cg = cooperative_groups;\n\n#include // for FLT_MAX\n#include \n#include \n#include \n\n#include \"CudaMath.h\"\n#include \"dds.h\"\n#include \"permutations.h\"\n\n// Definitions\n#define INPUT_IMAGE \"teapot512_std.ppm\"\n#define REFERENCE_IMAGE \"teapot512_ref.dds\"\n\n#define ERROR_THRESHOLD 0.02f\n\n#define NUM_THREADS 64 // Number of threads per block.\n\n#define __debugsync()\n\ntemplate __device__ inline void swap(T &a, T &b)\n{\n T tmp = a;\n a = b;\n b = tmp;\n}\n\n//__constant__ float3 kColorMetric = { 0.2126f, 0.7152f, 0.0722f };\n__constant__ float3 kColorMetric = {1.0f, 1.0f, 1.0f};\n\n////////////////////////////////////////////////////////////////////////////////\n// Sort colors\n////////////////////////////////////////////////////////////////////////////////\n__device__ void sortColors(const float *values, int *ranks, cg::thread_group tile)\n{\n const int tid = threadIdx.x;\n\n int rank = 0;\n\n#pragma unroll\n\n for (int i = 0; i < 16; i++) {\n rank += (values[i] < values[tid]);\n }\n\n ranks[tid] = rank;\n\n cg::sync(tile);\n\n // Resolve elements with the same index.\n for (int i = 0; i < 15; i++) {\n if (tid > i && ranks[tid] == ranks[i]) {\n ++ranks[tid];\n }\n cg::sync(tile);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Load color block to shared mem\n////////////////////////////////////////////////////////////////////////////////\n__device__ void loadColorBlock(const uint *image,\n float3 colors[16],\n float3 sums[16],\n int xrefs[16],\n int blockOffset,\n cg::thread_block cta)\n{\n const int bid = blockIdx.x + blockOffset;\n const int idx = threadIdx.x;\n\n __shared__ float dps[16];\n\n float3 tmp;\n\n cg::thread_group tile = cg::tiled_partition(cta, 16);\n\n if (idx < 16) {\n // Read color and copy to shared mem.\n uint c = image[(bid) * 16 + idx];\n\n colors[idx].x = ((c >> 0) & 0xFF) * (1.0f / 255.0f);\n colors[idx].y = ((c >> 8) & 0xFF) * (1.0f / 255.0f);\n colors[idx].z = ((c >> 16) & 0xFF) * (1.0f / 255.0f);\n\n cg::sync(tile);\n // Sort colors along the best fit line.\n colorSums(colors, sums, tile);\n\n cg::sync(tile);\n\n float3 axis = bestFitLine(colors, sums[0], tile);\n\n cg::sync(tile);\n\n dps[idx] = dot(colors[idx], axis);\n\n cg::sync(tile);\n\n sortColors(dps, xrefs, tile);\n\n cg::sync(tile);\n\n tmp = colors[idx];\n\n cg::sync(tile);\n\n colors[xrefs[idx]] = tmp;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Round color to RGB565 and expand\n////////////////////////////////////////////////////////////////////////////////\ninline __device__ float3 roundAndExpand(float3 v, ushort *w)\n{\n v.x = rintf(__saturatef(v.x) * 31.0f);\n v.y = rintf(__saturatef(v.y) * 63.0f);\n v.z = rintf(__saturatef(v.z) * 31.0f);\n\n *w = ((ushort)v.x << 11) | ((ushort)v.y << 5) | (ushort)v.z;\n v.x *= 0.03227752766457f; // approximate integer bit expansion.\n v.y *= 0.01583151765563f;\n v.z *= 0.03227752766457f;\n return v;\n}\n\n__constant__ float alphaTable4[4] = {9.0f, 0.0f, 6.0f, 3.0f};\n__constant__ float alphaTable3[4] = {4.0f, 0.0f, 2.0f, 2.0f};\n__constant__ const int prods4[4] = {0x090000, 0x000900, 0x040102, 0x010402};\n__constant__ const int prods3[4] = {0x040000, 0x000400, 0x040101, 0x010401};\n\n#define USE_TABLES 1\n\n////////////////////////////////////////////////////////////////////////////////\n// Evaluate permutations\n////////////////////////////////////////////////////////////////////////////////\nstatic __device__ float\nevalPermutation4(const float3 *colors, uint permutation, ushort *start, ushort *end, float3 color_sum)\n{\n// Compute endpoints using least squares.\n#if USE_TABLES\n float3 alphax_sum = make_float3(0.0f, 0.0f, 0.0f);\n\n int akku = 0;\n\n // Compute alpha & beta for this permutation.\n for (int i = 0; i < 16; i++) {\n const uint bits = permutation >> (2 * i);\n\n alphax_sum += alphaTable4[bits & 3] * colors[i];\n akku += prods4[bits & 3];\n }\n\n float alpha2_sum = float(akku >> 16);\n float beta2_sum = float((akku >> 8) & 0xff);\n float alphabeta_sum = float((akku >> 0) & 0xff);\n float3 betax_sum = (9.0f * color_sum) - alphax_sum;\n#else\n float alpha2_sum = 0.0f;\n float beta2_sum = 0.0f;\n float alphabeta_sum = 0.0f;\n float3 alphax_sum = make_float3(0.0f, 0.0f, 0.0f);\n\n // Compute alpha & beta for this permutation.\n for (int i = 0; i < 16; i++) {\n const uint bits = permutation >> (2 * i);\n\n float beta = (bits & 1);\n\n if (bits & 2) {\n beta = (1 + beta) * (1.0f / 3.0f);\n }\n\n float alpha = 1.0f - beta;\n\n alpha2_sum += alpha * alpha;\n beta2_sum += beta * beta;\n alphabeta_sum += alpha * beta;\n alphax_sum += alpha * colors[i];\n }\n\n float3 betax_sum = color_sum - alphax_sum;\n#endif\n\n // alpha2, beta2, alphabeta and factor could be precomputed for each\n // permutation, but it's faster to recompute them.\n const float factor = 1.0f / (alpha2_sum * beta2_sum - alphabeta_sum * alphabeta_sum);\n\n float3 a = (alphax_sum * beta2_sum - betax_sum * alphabeta_sum) * factor;\n float3 b = (betax_sum * alpha2_sum - alphax_sum * alphabeta_sum) * factor;\n\n // Round a, b to the closest 5-6-5 color and expand...\n a = roundAndExpand(a, start);\n b = roundAndExpand(b, end);\n\n // compute the error\n float3 e = a * a * alpha2_sum + b * b * beta2_sum + 2.0f * (a * b * alphabeta_sum - a * alphax_sum - b * betax_sum);\n\n return (0.111111111111f) * dot(e, kColorMetric);\n}\n\nstatic __device__ float\nevalPermutation3(const float3 *colors, uint permutation, ushort *start, ushort *end, float3 color_sum)\n{\n// Compute endpoints using least squares.\n#if USE_TABLES\n float3 alphax_sum = make_float3(0.0f, 0.0f, 0.0f);\n\n int akku = 0;\n\n // Compute alpha & beta for this permutation.\n for (int i = 0; i < 16; i++) {\n const uint bits = permutation >> (2 * i);\n\n alphax_sum += alphaTable3[bits & 3] * colors[i];\n akku += prods3[bits & 3];\n }\n\n float alpha2_sum = float(akku >> 16);\n float beta2_sum = float((akku >> 8) & 0xff);\n float alphabeta_sum = float((akku >> 0) & 0xff);\n float3 betax_sum = (4.0f * color_sum) - alphax_sum;\n#else\n float alpha2_sum = 0.0f;\n float beta2_sum = 0.0f;\n float alphabeta_sum = 0.0f;\n float3 alphax_sum = make_float3(0.0f, 0.0f, 0.0f);\n\n // Compute alpha & beta for this permutation.\n for (int i = 0; i < 16; i++) {\n const uint bits = permutation >> (2 * i);\n\n float beta = (bits & 1);\n\n if (bits & 2) {\n beta = 0.5f;\n }\n\n float alpha = 1.0f - beta;\n\n alpha2_sum += alpha * alpha;\n beta2_sum += beta * beta;\n alphabeta_sum += alpha * beta;\n alphax_sum += alpha * colors[i];\n }\n\n float3 betax_sum = color_sum - alphax_sum;\n#endif\n\n const float factor = 1.0f / (alpha2_sum * beta2_sum - alphabeta_sum * alphabeta_sum);\n\n float3 a = (alphax_sum * beta2_sum - betax_sum * alphabeta_sum) * factor;\n float3 b = (betax_sum * alpha2_sum - alphax_sum * alphabeta_sum) * factor;\n\n // Round a, b to the closest 5-6-5 color and expand...\n a = roundAndExpand(a, start);\n b = roundAndExpand(b, end);\n\n // compute the error\n float3 e = a * a * alpha2_sum + b * b * beta2_sum + 2.0f * (a * b * alphabeta_sum - a * alphax_sum - b * betax_sum);\n\n return (0.25f) * dot(e, kColorMetric);\n}\n\n__device__ void evalAllPermutations(const float3 *colors,\n const uint *permutations,\n ushort &bestStart,\n ushort &bestEnd,\n uint &bestPermutation,\n float *errors,\n float3 color_sum,\n cg::thread_block cta)\n{\n const int idx = threadIdx.x;\n\n float bestError = FLT_MAX;\n\n __shared__ uint s_permutations[160];\n\n for (int i = 0; i < 16; i++) {\n int pidx = idx + NUM_THREADS * i;\n\n if (pidx >= 992) {\n break;\n }\n\n ushort start, end;\n uint permutation = permutations[pidx];\n\n if (pidx < 160) {\n s_permutations[pidx] = permutation;\n }\n\n float error = evalPermutation4(colors, permutation, &start, &end, color_sum);\n\n if (error < bestError) {\n bestError = error;\n bestPermutation = permutation;\n bestStart = start;\n bestEnd = end;\n }\n }\n\n if (bestStart < bestEnd) {\n swap(bestEnd, bestStart);\n bestPermutation ^= 0x55555555; // Flip indices.\n }\n\n cg::sync(cta); // Sync here to ensure s_permutations is valid going forward\n\n for (int i = 0; i < 3; i++) {\n int pidx = idx + NUM_THREADS * i;\n\n if (pidx >= 160) {\n break;\n }\n\n ushort start, end;\n uint permutation = s_permutations[pidx];\n float error = evalPermutation3(colors, permutation, &start, &end, color_sum);\n\n if (error < bestError) {\n bestError = error;\n bestPermutation = permutation;\n bestStart = start;\n bestEnd = end;\n\n if (bestStart > bestEnd) {\n swap(bestEnd, bestStart);\n bestPermutation ^= (~bestPermutation >> 1) & 0x55555555; // Flip indices.\n }\n }\n }\n\n errors[idx] = bestError;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Find index with minimum error\n////////////////////////////////////////////////////////////////////////////////\n__device__ int findMinError(float *errors, cg::thread_block cta)\n{\n const int idx = threadIdx.x;\n __shared__ int indices[NUM_THREADS];\n indices[idx] = idx;\n\n cg::sync(cta);\n\n for (int d = NUM_THREADS / 2; d > 0; d >>= 1) {\n float err0 = errors[idx];\n float err1 = (idx + d) < NUM_THREADS ? errors[idx + d] : FLT_MAX;\n int index1 = (idx + d) < NUM_THREADS ? indices[idx + d] : 0;\n\n cg::sync(cta);\n\n if (err1 < err0) {\n errors[idx] = err1;\n indices[idx] = index1;\n }\n\n cg::sync(cta);\n }\n\n return indices[0];\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Save DXT block\n////////////////////////////////////////////////////////////////////////////////\n__device__ void saveBlockDXT1(ushort start, ushort end, uint permutation, int xrefs[16], uint2 *result, int blockOffset)\n{\n const int bid = blockIdx.x + blockOffset;\n\n if (start == end) {\n permutation = 0;\n }\n\n // Reorder permutation.\n uint indices = 0;\n\n for (int i = 0; i < 16; i++) {\n int ref = xrefs[i];\n indices |= ((permutation >> (2 * ref)) & 3) << (2 * i);\n }\n\n // Write endpoints.\n result[bid].x = (end << 16) | start;\n\n // Write palette indices.\n result[bid].y = indices;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Compress color block\n////////////////////////////////////////////////////////////////////////////////\n__global__ void compress(const uint *permutations, const uint *image, uint2 *result, int blockOffset)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n\n const int idx = threadIdx.x;\n\n __shared__ float3 colors[16];\n __shared__ float3 sums[16];\n __shared__ int xrefs[16];\n\n loadColorBlock(image, colors, sums, xrefs, blockOffset, cta);\n\n cg::sync(cta);\n\n ushort bestStart, bestEnd;\n uint bestPermutation;\n\n __shared__ float errors[NUM_THREADS];\n\n evalAllPermutations(colors, permutations, bestStart, bestEnd, bestPermutation, errors, sums[0], cta);\n\n // Use a parallel reduction to find minimum error.\n const int minIdx = findMinError(errors, cta);\n\n cg::sync(cta);\n\n // Only write the result of the winner thread.\n if (idx == minIdx) {\n saveBlockDXT1(bestStart, bestEnd, bestPermutation, xrefs, result, blockOffset);\n }\n}\n\n// Helper structs and functions to validate the output of the compressor.\n// We cannot simply do a bitwise compare, because different compilers produce\n// different\n// results for different targets due to floating point arithmetic.\n\nunion Color32\n{\n struct\n {\n unsigned char b, g, r, a;\n };\n unsigned int u;\n};\n\nunion Color16\n{\n struct\n {\n unsigned short b : 5;\n unsigned short g : 6;\n unsigned short r : 5;\n };\n unsigned short u;\n};\n\nstruct BlockDXT1\n{\n Color16 col0;\n Color16 col1;\n union\n {\n unsigned char row[4];\n unsigned int indices;\n };\n\n void decompress(Color32 colors[16]) const;\n};\n\nvoid BlockDXT1::decompress(Color32 *colors) const\n{\n Color32 palette[4];\n\n // Does bit expansion before interpolation.\n palette[0].b = (col0.b << 3) | (col0.b >> 2);\n palette[0].g = (col0.g << 2) | (col0.g >> 4);\n palette[0].r = (col0.r << 3) | (col0.r >> 2);\n palette[0].a = 0xFF;\n\n palette[1].r = (col1.r << 3) | (col1.r >> 2);\n palette[1].g = (col1.g << 2) | (col1.g >> 4);\n palette[1].b = (col1.b << 3) | (col1.b >> 2);\n palette[1].a = 0xFF;\n\n if (col0.u > col1.u) {\n // Four-color block: derive the other two colors.\n palette[2].r = (2 * palette[0].r + palette[1].r) / 3;\n palette[2].g = (2 * palette[0].g + palette[1].g) / 3;\n palette[2].b = (2 * palette[0].b + palette[1].b) / 3;\n palette[2].a = 0xFF;\n\n palette[3].r = (2 * palette[1].r + palette[0].r) / 3;\n palette[3].g = (2 * palette[1].g + palette[0].g) / 3;\n palette[3].b = (2 * palette[1].b + palette[0].b) / 3;\n palette[3].a = 0xFF;\n }\n else {\n // Three-color block: derive the other color.\n palette[2].r = (palette[0].r + palette[1].r) / 2;\n palette[2].g = (palette[0].g + palette[1].g) / 2;\n palette[2].b = (palette[0].b + palette[1].b) / 2;\n palette[2].a = 0xFF;\n\n palette[3].r = 0x00;\n palette[3].g = 0x00;\n palette[3].b = 0x00;\n palette[3].a = 0x00;\n }\n\n for (int i = 0; i < 16; i++) {\n colors[i] = palette[(indices >> (2 * i)) & 0x3];\n }\n}\n\nstatic int compareColors(const Color32 *b0, const Color32 *b1)\n{\n int sum = 0;\n\n for (int i = 0; i < 16; i++) {\n int r = (b0[i].r - b1[i].r);\n int g = (b0[i].g - b1[i].g);\n int b = (b0[i].b - b1[i].b);\n sum += r * r + g * g + b * b;\n }\n\n return sum;\n}\n\nstatic int compareBlock(const BlockDXT1 *b0, const BlockDXT1 *b1)\n{\n Color32 colors0[16];\n Color32 colors1[16];\n\n if (memcmp(b0, b1, sizeof(BlockDXT1)) == 0) {\n return 0;\n }\n else {\n b0->decompress(colors0);\n b1->decompress(colors1);\n\n return compareColors(colors0, colors1);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"%s Starting...\\n\\n\", argv[0]);\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n findCudaDevice(argc, (const char **)argv);\n\n // Load input image.\n unsigned char *data = NULL;\n uint W, H;\n\n char *image_path = sdkFindFilePath(INPUT_IMAGE, argv[0]);\n\n if (image_path == 0) {\n printf(\"Error, unable to find source image <%s>\\n\", image_path);\n exit(EXIT_FAILURE);\n }\n\n if (!sdkLoadPPM4ub(image_path, &data, &W, &H)) {\n printf(\"Error, unable to open source image file <%s>\\n\", image_path);\n\n exit(EXIT_FAILURE);\n }\n\n uint w = W, h = H;\n\n printf(\"Image Loaded '%s', %d x %d pixels\\n\\n\", image_path, w, h);\n\n // Allocate input image.\n const uint memSize = w * h * 4;\n assert(0 != memSize);\n uint *block_image = (uint *)malloc(memSize);\n\n // Convert linear image to block linear.\n for (uint by = 0; by < h / 4; by++) {\n for (uint bx = 0; bx < w / 4; bx++) {\n for (int i = 0; i < 16; i++) {\n const int x = i & 3;\n const int y = i / 4;\n block_image[(by * w / 4 + bx) * 16 + i] = ((uint *)data)[(by * 4 + y) * 4 * (W / 4) + bx * 4 + x];\n }\n }\n }\n\n // copy into global mem\n uint *d_data = NULL;\n checkCudaErrors(hipMalloc((void **)&d_data, memSize));\n\n // Result\n uint *d_result = NULL;\n const uint compressedSize = (w / 4) * (h / 4) * 8;\n checkCudaErrors(hipMalloc((void **)&d_result, compressedSize));\n uint *h_result = (uint *)malloc(compressedSize);\n\n // Compute permutations.\n uint permutations[1024];\n computePermutations(permutations);\n\n // Copy permutations host to devie.\n uint *d_permutations = NULL;\n checkCudaErrors(hipMalloc((void **)&d_permutations, 1024 * sizeof(uint)));\n checkCudaErrors(hipMemcpy(d_permutations, permutations, 1024 * sizeof(uint), hipMemcpyHostToDevice));\n\n // create a timer\n StopWatchInterface *timer = NULL;\n sdkCreateTimer(&timer);\n\n // Copy image from host to device\n checkCudaErrors(hipMemcpy(d_data, block_image, memSize, hipMemcpyHostToDevice));\n\n // Determine launch configuration and run timed computation numIterations\n // times\n uint blocks = ((w + 3) / 4) * ((h + 3) / 4); // rounds up by 1 block in each dim if %4 != 0\n\n int devID;\n hipDeviceProp_t deviceProp;\n\n // get number of SMs on this GPU\n checkCudaErrors(hipGetDevice(&devID));\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, devID));\n\n // Restrict the numbers of blocks to launch on low end GPUs to avoid kernel\n // timeout\n int blocksPerLaunch = min(blocks, 768 * deviceProp.multiProcessorCount);\n\n printf(\"Running DXT Compression on %u x %u image...\\n\", w, h);\n printf(\"\\n%u Blocks, %u Threads per Block, %u Threads in Grid...\\n\\n\", blocks, NUM_THREADS, blocks * NUM_THREADS);\n int numIterations = 1;\n\n for (int i = -1; i < numIterations; ++i) {\n if (i == 0) {\n checkCudaErrors(hipDeviceSynchronize());\n sdkStartTimer(&timer);\n }\n\n for (int j = 0; j < (int)blocks; j += blocksPerLaunch) {\n compress<<>>(d_permutations, d_data, (uint2 *)d_result, j);\n }\n }\n\n getLastCudaError(\"compress\");\n\n // sync to host, stop timer, record perf\n checkCudaErrors(hipDeviceSynchronize());\n sdkStopTimer(&timer);\n double dAvgTime = 1.0e-3 * sdkGetTimerValue(&timer) / (double)numIterations;\n printf(\"dxtc, Throughput = %.4f MPixels/s, Time = %.5f s, Size = %u Pixels, \"\n \"NumDevsUsed = %i, Workgroup = %d\\n\",\n (1.0e-6 * (double)(W * H) / dAvgTime),\n dAvgTime,\n (W * H),\n 1,\n NUM_THREADS);\n\n // copy result data from device to host\n checkCudaErrors(hipMemcpy(h_result, d_result, compressedSize, hipMemcpyDeviceToHost));\n\n // Write out result data to DDS file\n char output_filename[1024];\n strcpy(output_filename, image_path);\n strcpy(output_filename + strlen(image_path) - 3, \"dds\");\n FILE *fp = fopen(output_filename, \"wb\");\n\n if (fp == 0) {\n printf(\"Error, unable to open output image <%s>\\n\", output_filename);\n exit(EXIT_FAILURE);\n }\n\n DDSHeader header;\n header.fourcc = FOURCC_DDS;\n header.size = 124;\n header.flags = (DDSD_WIDTH | DDSD_HEIGHT | DDSD_CAPS | DDSD_PIXELFORMAT | DDSD_LINEARSIZE);\n header.height = h;\n header.width = w;\n header.pitch = compressedSize;\n header.depth = 0;\n header.mipmapcount = 0;\n memset(header.reserved, 0, sizeof(header.reserved));\n header.pf.size = 32;\n header.pf.flags = DDPF_FOURCC;\n header.pf.fourcc = FOURCC_DXT1;\n header.pf.bitcount = 0;\n header.pf.rmask = 0;\n header.pf.gmask = 0;\n header.pf.bmask = 0;\n header.pf.amask = 0;\n header.caps.caps1 = DDSCAPS_TEXTURE;\n header.caps.caps2 = 0;\n header.caps.caps3 = 0;\n header.caps.caps4 = 0;\n header.notused = 0;\n fwrite(&header, sizeof(DDSHeader), 1, fp);\n fwrite(h_result, compressedSize, 1, fp);\n fclose(fp);\n\n // Make sure the generated image is correct.\n const char *reference_image_path = sdkFindFilePath(REFERENCE_IMAGE, argv[0]);\n\n if (reference_image_path == 0) {\n printf(\"Error, unable to find reference image\\n\");\n\n exit(EXIT_FAILURE);\n }\n\n fp = fopen(reference_image_path, \"rb\");\n\n if (fp == 0) {\n printf(\"Error, unable to open reference image\\n\");\n\n exit(EXIT_FAILURE);\n }\n\n fseek(fp, sizeof(DDSHeader), SEEK_SET);\n uint referenceSize = (W / 4) * (H / 4) * 8;\n uint *reference = (uint *)malloc(referenceSize);\n fread(reference, referenceSize, 1, fp);\n fclose(fp);\n\n printf(\"\\nChecking accuracy...\\n\");\n float rms = 0;\n\n for (uint y = 0; y < h; y += 4) {\n for (uint x = 0; x < w; x += 4) {\n uint referenceBlockIdx = ((y / 4) * (W / 4) + (x / 4));\n uint resultBlockIdx = ((y / 4) * (w / 4) + (x / 4));\n\n int cmp =\n compareBlock(((BlockDXT1 *)h_result) + resultBlockIdx, ((BlockDXT1 *)reference) + referenceBlockIdx);\n\n if (cmp != 0.0f) {\n printf(\"Deviation at (%4d,%4d):\\t%f rms\\n\", x / 4, y / 4, float(cmp) / 16 / 3);\n }\n\n rms += cmp;\n }\n }\n\n rms /= w * h * 3;\n\n // Free allocated resources and exit\n checkCudaErrors(hipFree(d_permutations));\n checkCudaErrors(hipFree(d_data));\n checkCudaErrors(hipFree(d_result));\n free(image_path);\n free(data);\n free(block_image);\n free(h_result);\n free(reference);\n sdkDeleteTimer(&timer);\n\n printf(\"RMS(reference, result) = %f\\n\\n\", rms);\n printf(rms <= ERROR_THRESHOLD ? \"Test passed\\n\" : \"Test failed!\\n\");\n /* Return zero if test passed, one otherwise */\n return rms > ERROR_THRESHOLD;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/dxtc/dxtc.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"binomialOptions_common.h\"\n#include \"common_gpu_header.h\"\n#include \"realtype.h\"\n\n// Preprocessed input option data\ntypedef struct\n{\n real S;\n real X;\n real vDt;\n real puByDf;\n real pdByDf;\n} __TOptionData;\nstatic __constant__ __TOptionData d_OptionData[MAX_OPTIONS];\n__device__ real d_CallValue[MAX_OPTIONS];\n\n#define THREADBLOCK_SIZE 128\n#define ELEMS_PER_THREAD (NUM_STEPS / THREADBLOCK_SIZE)\n#if NUM_STEPS % THREADBLOCK_SIZE\n#error Bad constants\n#endif\n\n////////////////////////////////////////////////////////////////////////////////\n// Overloaded shortcut functions for different precision modes\n////////////////////////////////////////////////////////////////////////////////\n\n#ifndef DOUBLE_PRECISION\n__device__ inline float expiryCallValue(float S, float X, float vDt, int i)\n{\n float d = S * __expf(vDt * (2.0f * i - NUM_STEPS)) - X;\n return (d > 0.0F) ? d : 0.0F;\n}\n\n#else\n__device__ inline double expiryCallValue(double S, double X, double vDt, int i)\n{\n double d = S * exp(vDt * (2.0 * i - NUM_STEPS)) - X;\n return (d > 0.0) ? d : 0.0;\n}\n#endif\n\n////////////////////////////////////////////////////////////////////////////////\n// GPU kernel\n////////////////////////////////////////////////////////////////////////////////\nextern \"C\" __global__ void binomialOptionsKernel()\n{\n __shared__ real call_exchange[THREADBLOCK_SIZE + 1];\n\n const int tid = threadIdx.x;\n const real S = d_OptionData[blockIdx.x].S;\n const real X = d_OptionData[blockIdx.x].X;\n const real vDt = d_OptionData[blockIdx.x].vDt;\n const real puByDf = d_OptionData[blockIdx.x].puByDf;\n const real pdByDf = d_OptionData[blockIdx.x].pdByDf;\n\n real call[ELEMS_PER_THREAD + 1];\n#pragma unroll\n for (int i = 0; i < ELEMS_PER_THREAD; ++i)\n call[i] = expiryCallValue(S, X, vDt, tid * ELEMS_PER_THREAD + i);\n\n if (tid == 0)\n call_exchange[THREADBLOCK_SIZE] = expiryCallValue(S, X, vDt, NUM_STEPS);\n\n int final_it = max(0, tid * ELEMS_PER_THREAD - 1);\n\n#pragma unroll 16\n for (int i = NUM_STEPS; i > 0; --i) {\n call_exchange[tid] = call[0];\n __syncthreads();\n call[ELEMS_PER_THREAD] = call_exchange[tid + 1];\n __syncthreads();\n\n if (i > final_it) {\n#pragma unroll\n for (int j = 0; j < ELEMS_PER_THREAD; ++j)\n call[j] = puByDf * call[j + 1] + pdByDf * call[j];\n }\n }\n\n if (tid == 0) {\n d_CallValue[blockIdx.x] = call[0];\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \"binomialOptions_common.h\"\n#include \"common_gpu_header.h\"\n#include \"realtype.h\"\n\n// Preprocessed input option data\ntypedef struct\n{\n real S;\n real X;\n real vDt;\n real puByDf;\n real pdByDf;\n} __TOptionData;\nstatic __constant__ __TOptionData d_OptionData[MAX_OPTIONS];\n__device__ real d_CallValue[MAX_OPTIONS];\n\n#define THREADBLOCK_SIZE 128\n#define ELEMS_PER_THREAD (NUM_STEPS / THREADBLOCK_SIZE)\n#if NUM_STEPS % THREADBLOCK_SIZE\n#error Bad constants\n#endif\n\n////////////////////////////////////////////////////////////////////////////////\n// Overloaded shortcut functions for different precision modes\n////////////////////////////////////////////////////////////////////////////////\n\n#ifndef DOUBLE_PRECISION\n__device__ inline float expiryCallValue(float S, float X, float vDt, int i)\n{\n float d = S * __expf(vDt * (2.0f * i - NUM_STEPS)) - X;\n return (d > 0.0F) ? d : 0.0F;\n}\n\n#else\n__device__ inline double expiryCallValue(double S, double X, double vDt, int i)\n{\n double d = S * exp(vDt * (2.0 * i - NUM_STEPS)) - X;\n return (d > 0.0) ? d : 0.0;\n}\n#endif\n\n////////////////////////////////////////////////////////////////////////////////\n// GPU kernel\n////////////////////////////////////////////////////////////////////////////////\nextern \"C\" __global__ void binomialOptionsKernel()\n{\n __shared__ real call_exchange[THREADBLOCK_SIZE + 1];\n\n const int tid = threadIdx.x;\n const real S = d_OptionData[blockIdx.x].S;\n const real X = d_OptionData[blockIdx.x].X;\n const real vDt = d_OptionData[blockIdx.x].vDt;\n const real puByDf = d_OptionData[blockIdx.x].puByDf;\n const real pdByDf = d_OptionData[blockIdx.x].pdByDf;\n\n real call[ELEMS_PER_THREAD + 1];\n#pragma unroll\n for (int i = 0; i < ELEMS_PER_THREAD; ++i)\n call[i] = expiryCallValue(S, X, vDt, tid * ELEMS_PER_THREAD + i);\n\n if (tid == 0)\n call_exchange[THREADBLOCK_SIZE] = expiryCallValue(S, X, vDt, NUM_STEPS);\n\n int final_it = max(0, tid * ELEMS_PER_THREAD - 1);\n\n#pragma unroll 16\n for (int i = NUM_STEPS; i > 0; --i) {\n call_exchange[tid] = call[0];\n __syncthreads();\n call[ELEMS_PER_THREAD] = call_exchange[tid + 1];\n __syncthreads();\n\n if (i > final_it) {\n#pragma unroll\n for (int j = 0; j < ELEMS_PER_THREAD; ++j)\n call[j] = puByDf * call[j + 1] + pdByDf * call[j];\n }\n }\n\n if (tid == 0) {\n d_CallValue[blockIdx.x] = call[0];\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/binomialOptions_nvrtc/binomialOptions_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \n\n#include \"SobelFilter_kernels.h\"\n\n// Texture object for reading image\ncudaTextureObject_t texObject;\nextern __shared__ unsigned char LocalBlock[];\nstatic cudaArray *array = NULL;\n\n#define RADIUS 1\n\n#ifdef FIXED_BLOCKWIDTH\n#define BlockWidth 80\n#define SharedPitch 384\n#endif\n\n// This will output the proper CUDA error strings in the event that a CUDA host\n// call returns an error\n#define checkCudaErrors(err) __checkCudaErrors(err, __FILE__, __LINE__)\n\ninline void __checkCudaErrors(cudaError err, const char *file, const int line)\n{\n if (cudaSuccess != err) {\n fprintf(stderr, \"%s(%i) : CUDA Runtime API error %d: %s.\\n\", file, line, (int)err, cudaGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n}\n\n__device__ unsigned char ComputeSobel(unsigned char ul, // upper left\n unsigned char um, // upper middle\n unsigned char ur, // upper right\n unsigned char ml, // middle left\n unsigned char mm, // middle (unused)\n unsigned char mr, // middle right\n unsigned char ll, // lower left\n unsigned char lm, // lower middle\n unsigned char lr, // lower right\n float fScale)\n{\n short Horz = ur + 2 * mr + lr - ul - 2 * ml - ll;\n short Vert = ul + 2 * um + ur - ll - 2 * lm - lr;\n short Sum = (short)(fScale * (abs((int)Horz) + abs((int)Vert)));\n\n if (Sum < 0) {\n return 0;\n }\n else if (Sum > 0xff) {\n return 0xff;\n }\n\n return (unsigned char)Sum;\n}\n\n__global__ void SobelShared(uchar4 *pSobelOriginal,\n unsigned short SobelPitch,\n#ifndef FIXED_BLOCKWIDTH\n short BlockWidth,\n short SharedPitch,\n#endif\n short w,\n short h,\n float fScale,\n cudaTextureObject_t tex)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n short u = 4 * blockIdx.x * BlockWidth;\n short v = blockIdx.y * blockDim.y + threadIdx.y;\n short ib;\n\n int SharedIdx = threadIdx.y * SharedPitch;\n\n for (ib = threadIdx.x; ib < BlockWidth + 2 * RADIUS; ib += blockDim.x) {\n LocalBlock[SharedIdx + 4 * ib + 0] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 0), (float)(v - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 1] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 1), (float)(v - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 2] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 2), (float)(v - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 3] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 3), (float)(v - RADIUS));\n }\n\n if (threadIdx.y < RADIUS * 2) {\n //\n // copy trailing RADIUS*2 rows of pixels into shared\n //\n SharedIdx = (blockDim.y + threadIdx.y) * SharedPitch;\n\n for (ib = threadIdx.x; ib < BlockWidth + 2 * RADIUS; ib += blockDim.x) {\n LocalBlock[SharedIdx + 4 * ib + 0] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 0), (float)(v + blockDim.y - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 1] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 1), (float)(v + blockDim.y - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 2] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 2), (float)(v + blockDim.y - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 3] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 3), (float)(v + blockDim.y - RADIUS));\n }\n }\n\n cg::sync(cta);\n\n u >>= 2; // index as uchar4 from here\n uchar4 *pSobel = (uchar4 *)(((char *)pSobelOriginal) + v * SobelPitch);\n SharedIdx = threadIdx.y * SharedPitch;\n\n for (ib = threadIdx.x; ib < BlockWidth; ib += blockDim.x) {\n unsigned char pix00 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 0];\n unsigned char pix01 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 1];\n unsigned char pix02 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 2];\n unsigned char pix10 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 0];\n unsigned char pix11 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 1];\n unsigned char pix12 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 2];\n unsigned char pix20 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 0];\n unsigned char pix21 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 1];\n unsigned char pix22 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 2];\n\n uchar4 out;\n\n out.x = ComputeSobel(pix00, pix01, pix02, pix10, pix11, pix12, pix20, pix21, pix22, fScale);\n\n pix00 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 3];\n pix10 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 3];\n pix20 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 3];\n out.y = ComputeSobel(pix01, pix02, pix00, pix11, pix12, pix10, pix21, pix22, pix20, fScale);\n\n pix01 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 4];\n pix11 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 4];\n pix21 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 4];\n out.z = ComputeSobel(pix02, pix00, pix01, pix12, pix10, pix11, pix22, pix20, pix21, fScale);\n\n pix02 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 5];\n pix12 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 5];\n pix22 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 5];\n out.w = ComputeSobel(pix00, pix01, pix02, pix10, pix11, pix12, pix20, pix21, pix22, fScale);\n\n if (u + ib < w / 4 && v < h) {\n pSobel[u + ib] = out;\n }\n }\n\n cg::sync(cta);\n}\n\n__global__ void\nSobelCopyImage(Pixel *pSobelOriginal, unsigned int Pitch, int w, int h, float fscale, cudaTextureObject_t tex)\n{\n unsigned char *pSobel = (unsigned char *)(((char *)pSobelOriginal) + blockIdx.x * Pitch);\n\n for (int i = threadIdx.x; i < w; i += blockDim.x) {\n pSobel[i] = min(max((tex2D(tex, (float)i, (float)blockIdx.x) * fscale), 0.f), 255.f);\n }\n}\n\n__global__ void SobelTex(Pixel *pSobelOriginal, unsigned int Pitch, int w, int h, float fScale, cudaTextureObject_t tex)\n{\n unsigned char *pSobel = (unsigned char *)(((char *)pSobelOriginal) + blockIdx.x * Pitch);\n\n for (int i = threadIdx.x; i < w; i += blockDim.x) {\n unsigned char pix00 = tex2D(tex, (float)i - 1, (float)blockIdx.x - 1);\n unsigned char pix01 = tex2D(tex, (float)i + 0, (float)blockIdx.x - 1);\n unsigned char pix02 = tex2D(tex, (float)i + 1, (float)blockIdx.x - 1);\n unsigned char pix10 = tex2D(tex, (float)i - 1, (float)blockIdx.x + 0);\n unsigned char pix11 = tex2D(tex, (float)i + 0, (float)blockIdx.x + 0);\n unsigned char pix12 = tex2D(tex, (float)i + 1, (float)blockIdx.x + 0);\n unsigned char pix20 = tex2D(tex, (float)i - 1, (float)blockIdx.x + 1);\n unsigned char pix21 = tex2D(tex, (float)i + 0, (float)blockIdx.x + 1);\n unsigned char pix22 = tex2D(tex, (float)i + 1, (float)blockIdx.x + 1);\n pSobel[i] = ComputeSobel(pix00, pix01, pix02, pix10, pix11, pix12, pix20, pix21, pix22, fScale);\n }\n}\n\nextern \"C\" void setupTexture(int iw, int ih, Pixel *data, int Bpp)\n{\n cudaChannelFormatDesc desc;\n\n if (Bpp == 1) {\n desc = cudaCreateChannelDesc();\n }\n else {\n desc = cudaCreateChannelDesc();\n }\n\n checkCudaErrors(cudaMallocArray(&array, &desc, iw, ih));\n checkCudaErrors(cudaMemcpy2DToArray(\n array, 0, 0, data, iw * Bpp * sizeof(Pixel), iw * Bpp * sizeof(Pixel), ih, cudaMemcpyHostToDevice));\n\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = array;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = cudaFilterModePoint;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&texObject, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void deleteTexture(void)\n{\n checkCudaErrors(cudaFreeArray(array));\n checkCudaErrors(cudaDestroyTextureObject(texObject));\n}\n\n// Wrapper for the __global__ call that sets up the texture and threads\nextern \"C\" void sobelFilter(Pixel *odata, int iw, int ih, enum SobelDisplayMode mode, float fScale)\n{\n switch (mode) {\n case SOBELDISPLAY_IMAGE:\n SobelCopyImage<<>>(odata, iw, iw, ih, fScale, texObject);\n break;\n\n case SOBELDISPLAY_SOBELTEX:\n SobelTex<<>>(odata, iw, iw, ih, fScale, texObject);\n break;\n\n case SOBELDISPLAY_SOBELSHARED: {\n dim3 threads(16, 4);\n#ifndef FIXED_BLOCKWIDTH\n int BlockWidth = 80; // must be divisible by 16 for coalescing\n#endif\n dim3 blocks =\n dim3(iw / (4 * BlockWidth) + (0 != iw % (4 * BlockWidth)), ih / threads.y + (0 != ih % threads.y));\n int SharedPitch = ~0x3f & (4 * (BlockWidth + 2 * RADIUS) + 0x3f);\n int sharedMem = SharedPitch * (threads.y + 2 * RADIUS);\n\n // for the shared kernel, width must be divisible by 4\n iw &= ~3;\n\n SobelShared<<>>((uchar4 *)odata,\n iw,\n#ifndef FIXED_BLOCKWIDTH\n BlockWidth,\n SharedPitch,\n#endif\n iw,\n ih,\n fScale,\n texObject);\n } break;\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \n\n#include \"SobelFilter_kernels.h\"\n\n// Texture object for reading image\nhipTextureObject_t texObject;\nextern __shared__ unsigned char LocalBlock[];\nstatic hipArray *array = NULL;\n\n#define RADIUS 1\n\n#ifdef FIXED_BLOCKWIDTH\n#define BlockWidth 80\n#define SharedPitch 384\n#endif\n\n// This will output the proper CUDA error strings in the event that a CUDA host\n// call returns an error\n#define checkCudaErrors(err) __checkCudaErrors(err, __FILE__, __LINE__)\n\ninline void __checkCudaErrors(hipError_t err, const char *file, const int line)\n{\n if (hipSuccess != err) {\n fprintf(stderr, \"%s(%i) : CUDA Runtime API error %d: %s.\\n\", file, line, (int)err, hipGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n}\n\n__device__ unsigned char ComputeSobel(unsigned char ul, // upper left\n unsigned char um, // upper middle\n unsigned char ur, // upper right\n unsigned char ml, // middle left\n unsigned char mm, // middle (unused)\n unsigned char mr, // middle right\n unsigned char ll, // lower left\n unsigned char lm, // lower middle\n unsigned char lr, // lower right\n float fScale)\n{\n short Horz = ur + 2 * mr + lr - ul - 2 * ml - ll;\n short Vert = ul + 2 * um + ur - ll - 2 * lm - lr;\n short Sum = (short)(fScale * (abs((int)Horz) + abs((int)Vert)));\n\n if (Sum < 0) {\n return 0;\n }\n else if (Sum > 0xff) {\n return 0xff;\n }\n\n return (unsigned char)Sum;\n}\n\n__global__ void SobelShared(uchar4 *pSobelOriginal,\n unsigned short SobelPitch,\n#ifndef FIXED_BLOCKWIDTH\n short BlockWidth,\n short SharedPitch,\n#endif\n short w,\n short h,\n float fScale,\n hipTextureObject_t tex)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n short u = 4 * blockIdx.x * BlockWidth;\n short v = blockIdx.y * blockDim.y + threadIdx.y;\n short ib;\n\n int SharedIdx = threadIdx.y * SharedPitch;\n\n for (ib = threadIdx.x; ib < BlockWidth + 2 * RADIUS; ib += blockDim.x) {\n LocalBlock[SharedIdx + 4 * ib + 0] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 0), (float)(v - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 1] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 1), (float)(v - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 2] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 2), (float)(v - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 3] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 3), (float)(v - RADIUS));\n }\n\n if (threadIdx.y < RADIUS * 2) {\n //\n // copy trailing RADIUS*2 rows of pixels into shared\n //\n SharedIdx = (blockDim.y + threadIdx.y) * SharedPitch;\n\n for (ib = threadIdx.x; ib < BlockWidth + 2 * RADIUS; ib += blockDim.x) {\n LocalBlock[SharedIdx + 4 * ib + 0] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 0), (float)(v + blockDim.y - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 1] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 1), (float)(v + blockDim.y - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 2] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 2), (float)(v + blockDim.y - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 3] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 3), (float)(v + blockDim.y - RADIUS));\n }\n }\n\n cg::sync(cta);\n\n u >>= 2; // index as uchar4 from here\n uchar4 *pSobel = (uchar4 *)(((char *)pSobelOriginal) + v * SobelPitch);\n SharedIdx = threadIdx.y * SharedPitch;\n\n for (ib = threadIdx.x; ib < BlockWidth; ib += blockDim.x) {\n unsigned char pix00 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 0];\n unsigned char pix01 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 1];\n unsigned char pix02 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 2];\n unsigned char pix10 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 0];\n unsigned char pix11 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 1];\n unsigned char pix12 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 2];\n unsigned char pix20 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 0];\n unsigned char pix21 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 1];\n unsigned char pix22 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 2];\n\n uchar4 out;\n\n out.x = ComputeSobel(pix00, pix01, pix02, pix10, pix11, pix12, pix20, pix21, pix22, fScale);\n\n pix00 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 3];\n pix10 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 3];\n pix20 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 3];\n out.y = ComputeSobel(pix01, pix02, pix00, pix11, pix12, pix10, pix21, pix22, pix20, fScale);\n\n pix01 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 4];\n pix11 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 4];\n pix21 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 4];\n out.z = ComputeSobel(pix02, pix00, pix01, pix12, pix10, pix11, pix22, pix20, pix21, fScale);\n\n pix02 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 5];\n pix12 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 5];\n pix22 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 5];\n out.w = ComputeSobel(pix00, pix01, pix02, pix10, pix11, pix12, pix20, pix21, pix22, fScale);\n\n if (u + ib < w / 4 && v < h) {\n pSobel[u + ib] = out;\n }\n }\n\n cg::sync(cta);\n}\n\n__global__ void\nSobelCopyImage(Pixel *pSobelOriginal, unsigned int Pitch, int w, int h, float fscale, hipTextureObject_t tex)\n{\n unsigned char *pSobel = (unsigned char *)(((char *)pSobelOriginal) + blockIdx.x * Pitch);\n\n for (int i = threadIdx.x; i < w; i += blockDim.x) {\n pSobel[i] = min(max((tex2D(tex, (float)i, (float)blockIdx.x) * fscale), 0.f), 255.f);\n }\n}\n\n__global__ void SobelTex(Pixel *pSobelOriginal, unsigned int Pitch, int w, int h, float fScale, hipTextureObject_t tex)\n{\n unsigned char *pSobel = (unsigned char *)(((char *)pSobelOriginal) + blockIdx.x * Pitch);\n\n for (int i = threadIdx.x; i < w; i += blockDim.x) {\n unsigned char pix00 = tex2D(tex, (float)i - 1, (float)blockIdx.x - 1);\n unsigned char pix01 = tex2D(tex, (float)i + 0, (float)blockIdx.x - 1);\n unsigned char pix02 = tex2D(tex, (float)i + 1, (float)blockIdx.x - 1);\n unsigned char pix10 = tex2D(tex, (float)i - 1, (float)blockIdx.x + 0);\n unsigned char pix11 = tex2D(tex, (float)i + 0, (float)blockIdx.x + 0);\n unsigned char pix12 = tex2D(tex, (float)i + 1, (float)blockIdx.x + 0);\n unsigned char pix20 = tex2D(tex, (float)i - 1, (float)blockIdx.x + 1);\n unsigned char pix21 = tex2D(tex, (float)i + 0, (float)blockIdx.x + 1);\n unsigned char pix22 = tex2D(tex, (float)i + 1, (float)blockIdx.x + 1);\n pSobel[i] = ComputeSobel(pix00, pix01, pix02, pix10, pix11, pix12, pix20, pix21, pix22, fScale);\n }\n}\n\nextern \"C\" void setupTexture(int iw, int ih, Pixel *data, int Bpp)\n{\n hipChannelFormatDesc desc;\n\n if (Bpp == 1) {\n desc = hipCreateChannelDesc();\n }\n else {\n desc = hipCreateChannelDesc();\n }\n\n checkCudaErrors(hipMallocArray(&array, &desc, iw, ih));\n checkCudaErrors(hipMemcpy2DToArray(\n array, 0, 0, data, iw * Bpp * sizeof(Pixel), iw * Bpp * sizeof(Pixel), ih, hipMemcpyHostToDevice));\n\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = array;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = hipFilterModePoint;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&texObject, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void deleteTexture(void)\n{\n checkCudaErrors(hipFreeArray(array));\n checkCudaErrors(hipDestroyTextureObject(texObject));\n}\n\n// Wrapper for the __global__ call that sets up the texture and threads\nextern \"C\" void sobelFilter(Pixel *odata, int iw, int ih, enum SobelDisplayMode mode, float fScale)\n{\n switch (mode) {\n case SOBELDISPLAY_IMAGE:\n SobelCopyImage<<>>(odata, iw, iw, ih, fScale, texObject);\n break;\n\n case SOBELDISPLAY_SOBELTEX:\n SobelTex<<>>(odata, iw, iw, ih, fScale, texObject);\n break;\n\n case SOBELDISPLAY_SOBELSHARED: {\n dim3 threads(16, 4);\n#ifndef FIXED_BLOCKWIDTH\n int BlockWidth = 80; // must be divisible by 16 for coalescing\n#endif\n dim3 blocks =\n dim3(iw / (4 * BlockWidth) + (0 != iw % (4 * BlockWidth)), ih / threads.y + (0 != ih % threads.y));\n int SharedPitch = ~0x3f & (4 * (BlockWidth + 2 * RADIUS) + 0x3f);\n int sharedMem = SharedPitch * (threads.y + 2 * RADIUS);\n\n // for the shared kernel, width must be divisible by 4\n iw &= ~3;\n\n SobelShared<<>>((uchar4 *)odata,\n iw,\n#ifndef FIXED_BLOCKWIDTH\n BlockWidth,\n SharedPitch,\n#endif\n iw,\n ih,\n fScale,\n texObject);\n } break;\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/SobelFilter/SobelFilter_kernels.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#ifndef _VOLUMEFILTER_KERNEL_CU_\n#define _VOLUMEFILTER_KERNEL_CU_\n\n#include \n#include \n\n#include \"volumeFilter.h\"\n\ntypedef unsigned int uint;\ntypedef unsigned char uchar;\ntypedef unsigned short ushort;\n\n__constant__ float4 c_filterData[VOLUMEFILTER_MAXWEIGHTS];\n\n__global__ void d_filter_surface3d(int filterSize,\n float filter_offset,\n cudaExtent volumeSize,\n cudaTextureObject_t volumeTexIn,\n cudaSurfaceObject_t volumeTexOut)\n{\n int x = blockIdx.x * blockDim.x + threadIdx.x;\n int y = blockIdx.y * blockDim.y + threadIdx.y;\n int z = blockIdx.z * blockDim.z + threadIdx.z;\n\n if (x >= volumeSize.width || y >= volumeSize.height || z >= volumeSize.depth) {\n return;\n }\n\n float filtered = 0;\n float4 basecoord = make_float4(x, y, z, 0);\n\n for (int i = 0; i < filterSize; i++) {\n float4 coord = basecoord + c_filterData[i];\n filtered += tex3D(volumeTexIn, coord.x, coord.y, coord.z) * c_filterData[i].w;\n }\n\n filtered += filter_offset;\n\n VolumeType output = VolumeTypeInfo::convert(filtered);\n\n // surface writes need byte offsets for x!\n surf3Dwrite(output, volumeTexOut, x * sizeof(VolumeType), y, z);\n}\n\nstatic unsigned int iDivUp(size_t a, size_t b)\n{\n size_t val = (a % b != 0) ? (a / b + 1) : (a / b);\n if (val > UINT_MAX) {\n fprintf(stderr, \"\\nUINT_MAX limit exceeded in iDivUp() exiting.....\\n\");\n exit(EXIT_FAILURE); // val exceeds limit\n }\n\n return static_cast(val);\n}\n\nextern \"C\" Volume *VolumeFilter_runFilter(Volume *input,\n Volume *output0,\n Volume *output1,\n int iterations,\n int numWeights,\n float4 *weights,\n float postWeightOffset)\n{\n Volume *swap = 0;\n cudaExtent size = input->size;\n unsigned int dim = 32 / sizeof(VolumeType);\n dim3 blockSize(dim, dim, 1);\n dim3 gridSize(iDivUp(size.width, blockSize.x), iDivUp(size.height, blockSize.y), iDivUp(size.depth, blockSize.z));\n\n // set weights\n checkCudaErrors(cudaMemcpyToSymbol(c_filterData, weights, sizeof(float4) * numWeights));\n\n for (int i = 0; i < iterations; i++) {\n d_filter_surface3d<<>>(\n numWeights, postWeightOffset, size, input->volumeTex, output0->volumeSurf);\n\n getLastCudaError(\"filter kernel failed\");\n\n swap = input;\n input = output0;\n output0 = swap;\n\n if (i == 0) {\n output0 = output1;\n }\n }\n\n return input;\n}\n#endif\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#ifndef _VOLUMEFILTER_KERNEL_CU_\n#define _VOLUMEFILTER_KERNEL_CU_\n\n#include \n#include \n\n#include \"volumeFilter.h\"\n\ntypedef unsigned int uint;\ntypedef unsigned char uchar;\ntypedef unsigned short ushort;\n\n__constant__ float4 c_filterData[VOLUMEFILTER_MAXWEIGHTS];\n\n__global__ void d_filter_surface3d(int filterSize,\n float filter_offset,\n hipExtent volumeSize,\n hipTextureObject_t volumeTexIn,\n hipSurfaceObject_t volumeTexOut)\n{\n int x = blockIdx.x * blockDim.x + threadIdx.x;\n int y = blockIdx.y * blockDim.y + threadIdx.y;\n int z = blockIdx.z * blockDim.z + threadIdx.z;\n\n if (x >= volumeSize.width || y >= volumeSize.height || z >= volumeSize.depth) {\n return;\n }\n\n float filtered = 0;\n float4 basecoord = make_float4(x, y, z, 0);\n\n for (int i = 0; i < filterSize; i++) {\n float4 coord = basecoord + c_filterData[i];\n filtered += tex3D(volumeTexIn, coord.x, coord.y, coord.z) * c_filterData[i].w;\n }\n\n filtered += filter_offset;\n\n VolumeType output = VolumeTypeInfo::convert(filtered);\n\n // surface writes need byte offsets for x!\n surf3Dwrite(output, volumeTexOut, x * sizeof(VolumeType), y, z);\n}\n\nstatic unsigned int iDivUp(size_t a, size_t b)\n{\n size_t val = (a % b != 0) ? (a / b + 1) : (a / b);\n if (val > UINT_MAX) {\n fprintf(stderr, \"\\nUINT_MAX limit exceeded in iDivUp() exiting.....\\n\");\n exit(EXIT_FAILURE); // val exceeds limit\n }\n\n return static_cast(val);\n}\n\nextern \"C\" Volume *VolumeFilter_runFilter(Volume *input,\n Volume *output0,\n Volume *output1,\n int iterations,\n int numWeights,\n float4 *weights,\n float postWeightOffset)\n{\n Volume *swap = 0;\n hipExtent size = input->size;\n unsigned int dim = 32 / sizeof(VolumeType);\n dim3 blockSize(dim, dim, 1);\n dim3 gridSize(iDivUp(size.width, blockSize.x), iDivUp(size.height, blockSize.y), iDivUp(size.depth, blockSize.z));\n\n // set weights\n checkCudaErrors(hipMemcpyToSymbol(HIP_SYMBOL(c_filterData), weights, sizeof(float4) * numWeights));\n\n for (int i = 0; i < iterations; i++) {\n d_filter_surface3d<<>>(\n numWeights, postWeightOffset, size, input->volumeTex, output0->volumeSurf);\n\n getLastCudaError(\"filter kernel failed\");\n\n swap = input;\n input = output0;\n output0 = swap;\n\n if (i == 0) {\n output0 = output1;\n }\n }\n\n return input;\n}\n#endif\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/volumeFiltering/volumeFilter_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Simple 3D volume renderer\n\n#ifndef _VOLUMERENDER_KERNEL_CU_\n#define _VOLUMERENDER_KERNEL_CU_\n\n#include \n#include \n\n#include \"volumeRender.h\"\n\n#define VOLUMERENDER_TFS 2\n#define VOLUMERENDER_TF_PREINTSIZE 1024\n#define VOLUMERENDER_TF_PREINTSTEPS 1024\n#define VOLUMERENDER_TF_PREINTRAY 4\n\nenum TFMode {\n TF_SINGLE_1D = 0, // single 1D TF for everything\n TF_LAYERED_2D_PREINT = 1, // layered 2D TF uses pre-integration\n TF_LAYERED_2D = 2, // layered 2D TF without pre-integration behavior\n};\n\ntypedef unsigned int uint;\ntypedef unsigned char uchar;\n\nstatic bool usePreInt = true;\nstatic cudaArray *d_transferIntegrate = 0;\nstatic cudaArray *d_transferFunc = 0;\nstatic cudaArray *d_transferArray = 0;\n\n// 1D transfer function texture\ncudaTextureObject_t transferTex;\n// 1D transfer integration texture\ncudaTextureObject_t transferIntegrateTex;\ncudaSurfaceObject_t transferIntegrateSurf;\n// 2D layered preintegrated transfer function texture\ncudaTextureObject_t transferLayerPreintTex;\ncudaSurfaceObject_t transferLayerPreintSurf;\n\ntypedef struct\n{\n float4 m[3];\n} float3x4;\n\n__constant__ float3x4 c_invViewMatrix; // inverse view matrix\n\nstruct Ray\n{\n float3 o; // origin\n float3 d; // direction\n};\n\n// intersect ray with a box\n// http://www.siggraph.org/education/materials/HyperGraph/raytrace/rtinter3.htm\n\n__device__ int intersectBox(Ray r, float3 boxmin, float3 boxmax, float *tnear, float *tfar)\n{\n // compute intersection of ray with all six bbox planes\n float3 invR = make_float3(1.0f) / r.d;\n float3 tbot = invR * (boxmin - r.o);\n float3 ttop = invR * (boxmax - r.o);\n\n // re-order intersections to find smallest and largest on each axis\n float3 tmin = fminf(ttop, tbot);\n float3 tmax = fmaxf(ttop, tbot);\n\n // find the largest tmin and the smallest tmax\n float largest_tmin = fmaxf(fmaxf(tmin.x, tmin.y), fmaxf(tmin.x, tmin.z));\n float smallest_tmax = fminf(fminf(tmax.x, tmax.y), fminf(tmax.x, tmax.z));\n\n *tnear = largest_tmin;\n *tfar = smallest_tmax;\n\n return smallest_tmax > largest_tmin;\n}\n\n// transform vector by matrix (no translation)\n__device__ float3 mul(const float3x4 &M, const float3 &v)\n{\n float3 r;\n r.x = dot(v, make_float3(M.m[0]));\n r.y = dot(v, make_float3(M.m[1]));\n r.z = dot(v, make_float3(M.m[2]));\n return r;\n}\n\n// transform vector by matrix with translation\n__device__ float4 mul(const float3x4 &M, const float4 &v)\n{\n float4 r;\n r.x = dot(v, M.m[0]);\n r.y = dot(v, M.m[1]);\n r.z = dot(v, M.m[2]);\n r.w = 1.0f;\n return r;\n}\n\n__device__ uint rgbaFloatToInt(float4 rgba)\n{\n rgba.x = __saturatef(rgba.x); // clamp to [0.0, 1.0]\n rgba.y = __saturatef(rgba.y);\n rgba.z = __saturatef(rgba.z);\n rgba.w = __saturatef(rgba.w);\n return (uint(rgba.w * 255) << 24) | (uint(rgba.z * 255) << 16) | (uint(rgba.y * 255) << 8) | uint(rgba.x * 255);\n}\n\ntemplate \n__device__ void d_render(uint *d_output,\n uint imageW,\n uint imageH,\n float density,\n float brightness,\n float transferOffset,\n float transferScale,\n cudaTextureObject_t volumeTex,\n cudaTextureObject_t transferTex,\n cudaTextureObject_t transferLayerPreintTex,\n float transferWeight = 0.0f)\n{\n const float rayscale = float(TFMODE != TF_SINGLE_1D ? VOLUMERENDER_TF_PREINTRAY : 1);\n const int maxSteps = 512;\n const float tstep = 0.01f * rayscale;\n const float opacityThreshold = 0.95f;\n const float3 boxMin = make_float3(-1.0f, -1.0f, -1.0f);\n const float3 boxMax = make_float3(1.0f, 1.0f, 1.0f);\n\n density *= rayscale;\n\n uint x = blockIdx.x * blockDim.x + threadIdx.x;\n uint y = blockIdx.y * blockDim.y + threadIdx.y;\n\n if ((x >= imageW) || (y >= imageH))\n return;\n\n float u = (x / (float)imageW) * 2.0f - 1.0f;\n float v = (y / (float)imageH) * 2.0f - 1.0f;\n\n // calculate eye ray in world space\n Ray eyeRay;\n eyeRay.o = make_float3(mul(c_invViewMatrix, make_float4(0.0f, 0.0f, 0.0f, 1.0f)));\n eyeRay.d = normalize(make_float3(u, v, -2.0f));\n eyeRay.d = mul(c_invViewMatrix, eyeRay.d);\n\n // find intersection with box\n float tnear, tfar;\n int hit = intersectBox(eyeRay, boxMin, boxMax, &tnear, &tfar);\n\n if (!hit)\n return;\n\n if (tnear < 0.0f)\n tnear = 0.0f; // clamp to near plane\n\n // march along ray from front to back, accumulating color\n float4 sum = make_float4(0.0f);\n float t = tnear;\n float3 pos = eyeRay.o + eyeRay.d * tnear;\n float3 step = eyeRay.d * tstep;\n\n float lastsample = 0;\n\n // lastsample = (lastsample-transferOffset)*transferScale;\n for (int i = 0; i < maxSteps; i++) {\n // read from 3D texture\n // remap position to [0, 1] coordinates\n float3 coord = make_float3(pos.x * 0.5f + 0.5f, pos.y * 0.5f + 0.5f, pos.z * 0.5f + 0.5f);\n float sample = tex3D(volumeTex, coord.x, coord.y, coord.z);\n // sample = (sample-transferOffset)*transferScale;\n // sample *= 64.0f; // scale for 10-bit data\n\n // lookup in transfer function texture\n float4 col;\n int tfid = (pos.x < 0);\n\n if (TFMODE != TF_SINGLE_1D) {\n col = tex2DLayered(\n transferLayerPreintTex, sample, TFMODE == TF_LAYERED_2D ? sample : lastsample, tfid);\n col.w *= density;\n lastsample = sample;\n }\n else {\n col = tex1D(transferTex, sample);\n col.w *= 0;\n }\n\n // \"under\" operator for back-to-front blending\n // sum = lerp(sum, col, col.w);\n\n // pre-multiply alpha\n col.x *= col.w;\n col.y *= col.w;\n col.z *= col.w;\n // \"over\" operator for front-to-back blending\n sum = sum + col * (1.0f - sum.w);\n\n // exit early if opaque\n if (sum.w > opacityThreshold)\n break;\n\n t += tstep;\n\n if (t > tfar)\n break;\n\n pos += step;\n }\n\n sum *= brightness;\n\n // write output color\n d_output[y * imageW + x] = rgbaFloatToInt(sum);\n}\n\n__global__ void d_render_regular(uint *d_output,\n uint imageW,\n uint imageH,\n float density,\n float brightness,\n float transferOffset,\n float transferScale,\n cudaTextureObject_t volumeTex,\n cudaTextureObject_t transferTex,\n cudaTextureObject_t transferLayerPreintTex,\n float transferWeight = 0.0f)\n{\n d_render(d_output,\n imageW,\n imageH,\n density,\n brightness,\n transferOffset,\n transferScale,\n volumeTex,\n transferTex,\n transferLayerPreintTex,\n transferWeight);\n}\n\n__global__ void d_render_preint(uint *d_output,\n uint imageW,\n uint imageH,\n float density,\n float brightness,\n float transferOffset,\n float transferScale,\n cudaTextureObject_t volumeTex,\n cudaTextureObject_t transferTex,\n cudaTextureObject_t transferLayerPreintTex,\n float transferWeight = 0.0f)\n{\n d_render(d_output,\n imageW,\n imageH,\n density,\n brightness,\n transferOffset,\n transferScale,\n volumeTex,\n transferTex,\n transferLayerPreintTex,\n transferWeight);\n}\n\n__global__ void d_render_preint_off(uint *d_output,\n uint imageW,\n uint imageH,\n float density,\n float brightness,\n float transferOffset,\n float transferScale,\n cudaTextureObject_t volumeTex,\n cudaTextureObject_t transferTex,\n cudaTextureObject_t transferLayerPreintTex,\n float transferWeight = 0.0f)\n{\n d_render(d_output,\n imageW,\n imageH,\n density,\n brightness,\n transferOffset,\n transferScale,\n volumeTex,\n transferTex,\n transferLayerPreintTex,\n transferWeight);\n}\n\n//////////////////////////////////////////////////////////////////////////\n\n__global__ void\nd_integrate_trapezoidal(cudaExtent extent, cudaTextureObject_t transferTex, cudaSurfaceObject_t transferIntegrateSurf)\n{\n uint x = blockIdx.x * blockDim.x + threadIdx.x;\n\n // for higher speed could use hierarchical approach for sum\n if (x >= extent.width) {\n return;\n }\n\n float stepsize = 1.0 / float(extent.width - 1);\n float to = float(x) * stepsize;\n\n float4 outclr = make_float4(0, 0, 0, 0);\n float incr = stepsize;\n\n float4 lastval = tex1D(transferTex, 0);\n\n float cur = incr;\n\n while (cur < to + incr * 0.5) {\n float4 val = tex1D(transferTex, cur);\n float4 trapezoid = (lastval + val) / 2.0f;\n lastval = val;\n\n outclr += trapezoid;\n cur += incr;\n }\n\n // surface writes need byte offsets for x!\n surf1Dwrite(outclr, transferIntegrateSurf, x * sizeof(float4));\n}\n\n__global__ void d_preintegrate(int layer,\n float steps,\n cudaExtent extent,\n cudaTextureObject_t transferTex,\n cudaTextureObject_t transferIntegrateTex,\n cudaSurfaceObject_t transferLayerPreintSurf)\n{\n uint x = blockIdx.x * blockDim.x + threadIdx.x;\n uint y = blockIdx.y * blockDim.y + threadIdx.y;\n\n if (x >= extent.width || y >= extent.height) {\n return;\n }\n\n float sx = float(x) / float(extent.width);\n float sy = float(y) / float(extent.height);\n\n float smax = max(sx, sy);\n float smin = min(sx, sy);\n\n float4 iv;\n\n if (x != y) {\n // assumes square textures!\n float fracc = smax - smin;\n fracc = 1.0 / (fracc * steps);\n\n float4 intmax = tex1D(transferIntegrateTex, smax);\n float4 intmin = tex1D(transferIntegrateTex, smin);\n iv.x = (intmax.x - intmin.x) * fracc;\n iv.y = (intmax.y - intmin.y) * fracc;\n iv.z = (intmax.z - intmin.z) * fracc;\n // iv.w = (intmax.w - intmin.w)*fracc;\n iv.w = (1.0 - exp(-(intmax.w - intmin.w) * fracc));\n }\n else {\n float4 sample = tex1D(transferTex, smin);\n iv.x = sample.x;\n iv.y = sample.y;\n iv.z = sample.z;\n // iv.w = sample.w;\n iv.w = (1.0 - exp(-sample.w));\n }\n\n iv.x = __saturatef(iv.x);\n iv.y = __saturatef(iv.y);\n iv.z = __saturatef(iv.z);\n iv.w = __saturatef(iv.w);\n\n // surface writes need byte offsets for x!\n surf2DLayeredwrite(iv, transferLayerPreintSurf, x * sizeof(float4), y, layer);\n}\n\n//////////////////////////////////////////////////////////////////////////\n\nvoid VolumeRender_setTextureFilterMode(bool bLinearFilter, Volume *vol)\n{\n if (vol->volumeTex) {\n checkCudaErrors(cudaDestroyTextureObject(vol->volumeTex));\n }\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = vol->content;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = bLinearFilter ? cudaFilterModeLinear : cudaFilterModePoint;\n\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.addressMode[2] = cudaAddressModeWrap;\n\n texDescr.readMode = VolumeTypeInfo::readMode;\n\n checkCudaErrors(cudaCreateTextureObject(&vol->volumeTex, &texRes, &texDescr, NULL));\n}\n\nstatic unsigned int iDivUp(size_t a, size_t b)\n{\n size_t val = (a % b != 0) ? (a / b + 1) : (a / b);\n if (val > UINT_MAX) {\n fprintf(stderr, \"\\nUINT_MAX limit exceeded in iDivUp() exiting.....\\n\");\n exit(EXIT_FAILURE); // val exceeds limit\n }\n\n return static_cast(val);\n}\n\nvoid VolumeRender_updateTF(int tfIdx, int numColors, float4 *colors)\n{\n if (d_transferFunc) {\n checkCudaErrors(cudaFreeArray(d_transferFunc));\n d_transferFunc = 0;\n }\n\n cudaChannelFormatDesc channelFloat4 = cudaCreateChannelDesc();\n checkCudaErrors(cudaMallocArray(&d_transferFunc, &channelFloat4, numColors, 1));\n checkCudaErrors(\n cudaMemcpy2DToArray(d_transferFunc, 0, 0, colors, 0, sizeof(float4) * numColors, 1, cudaMemcpyHostToDevice));\n\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = d_transferFunc;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = cudaFilterModeLinear;\n texDescr.addressMode[0] = cudaAddressModeClamp;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&transferTex, &texRes, &texDescr, NULL));\n\n if (tfIdx < 0 || tfIdx >= VOLUMERENDER_TFS) {\n return;\n }\n\n {\n cudaExtent extent = {VOLUMERENDER_TF_PREINTSTEPS, 0, 0};\n dim3 blockSize(32, 1, 1);\n dim3 gridSize(iDivUp(extent.width, blockSize.x), 1, 1);\n d_integrate_trapezoidal<<>>(extent, transferTex, transferIntegrateSurf);\n }\n\n {\n cudaExtent extent = {VOLUMERENDER_TF_PREINTSIZE, VOLUMERENDER_TF_PREINTSIZE, VOLUMERENDER_TFS};\n dim3 blockSize(16, 16, 1);\n dim3 gridSize(iDivUp(extent.width, blockSize.x), iDivUp(extent.height, blockSize.y), 1);\n d_preintegrate<<>>(tfIdx,\n float(VOLUMERENDER_TF_PREINTSTEPS),\n extent,\n transferTex,\n transferIntegrateTex,\n transferLayerPreintSurf);\n }\n}\n\nvoid VolumeRender_init()\n{\n cudaResourceDesc texRes;\n cudaTextureDesc texDescr;\n\n cudaChannelFormatDesc channelFloat4 = cudaCreateChannelDesc();\n cudaExtent extent = {VOLUMERENDER_TF_PREINTSIZE, VOLUMERENDER_TF_PREINTSIZE, VOLUMERENDER_TFS};\n checkCudaErrors(\n cudaMalloc3DArray(&d_transferArray, &channelFloat4, extent, cudaArrayLayered | cudaArraySurfaceLoadStore));\n\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = d_transferArray;\n\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = cudaFilterModeLinear;\n texDescr.addressMode[0] = cudaAddressModeClamp;\n texDescr.addressMode[1] = cudaAddressModeClamp;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&transferLayerPreintTex, &texRes, &texDescr, NULL));\n\n cudaResourceDesc surfRes;\n memset(&surfRes, 0, sizeof(cudaResourceDesc));\n surfRes.resType = cudaResourceTypeArray;\n surfRes.res.array.array = d_transferArray;\n\n checkCudaErrors(cudaCreateSurfaceObject(&transferLayerPreintSurf, &surfRes));\n\n checkCudaErrors(cudaMallocArray(\n &d_transferIntegrate, &channelFloat4, VOLUMERENDER_TF_PREINTSTEPS, 0, cudaArraySurfaceLoadStore));\n\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = d_transferIntegrate;\n\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = cudaFilterModeLinear;\n\n texDescr.addressMode[0] = cudaAddressModeClamp;\n texDescr.addressMode[1] = cudaAddressModeClamp;\n texDescr.addressMode[2] = cudaAddressModeClamp;\n\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&transferIntegrateTex, &texRes, &texDescr, NULL));\n\n memset(&surfRes, 0, sizeof(cudaResourceDesc));\n surfRes.resType = cudaResourceTypeArray;\n surfRes.res.array.array = d_transferIntegrate;\n\n checkCudaErrors(cudaCreateSurfaceObject(&transferIntegrateSurf, &surfRes));\n\n // create transfer function texture\n float4 transferFunc0[] = {\n {\n 0.0,\n 0.0,\n 0.0,\n 0.0,\n },\n {\n 1.0,\n 0.0,\n 0.0,\n 1.0,\n },\n {\n 1.0,\n 0.5,\n 0.0,\n 1.0,\n },\n {\n 1.0,\n 1.0,\n 0.0,\n 1.0,\n },\n {\n 0.0,\n 1.0,\n 0.0,\n 1.0,\n },\n {\n 0.0,\n 1.0,\n 1.0,\n 1.0,\n },\n {\n 0.0,\n 0.0,\n 1.0,\n 1.0,\n },\n {\n 1.0,\n 0.0,\n 1.0,\n 1.0,\n },\n {\n 0.0,\n 0.0,\n 0.0,\n 0.0,\n },\n };\n\n float4 transferFunc1[] = {\n {\n 0.0,\n 0.0,\n 0.0,\n 0.0,\n },\n {\n 0.0,\n 1.0,\n 0.0,\n 0.125,\n },\n {\n 0.0,\n 0.5,\n 1.0,\n 0.125,\n },\n {\n 0.0,\n 1.0,\n 1.0,\n 0.125,\n },\n {\n 0.0,\n 1.0,\n 0.0,\n 0.125,\n },\n {\n 0.25,\n 0.75,\n 0.0,\n 1.0,\n },\n {\n 0.75,\n 0.25,\n 0.0,\n 0.125,\n },\n {\n 1.0,\n 0.75,\n 0.0,\n 0.125,\n },\n {\n 0.0,\n 0.0,\n 0.0,\n 0.0,\n },\n };\n\n VolumeRender_updateTF(1, sizeof(transferFunc1) / sizeof(float4), transferFunc1);\n VolumeRender_updateTF(0, sizeof(transferFunc0) / sizeof(float4), transferFunc0);\n}\n\nvoid VolumeRender_deinit()\n{\n checkCudaErrors(cudaDestroyTextureObject(transferTex));\n checkCudaErrors(cudaDestroyTextureObject(transferIntegrateTex));\n checkCudaErrors(cudaDestroySurfaceObject(transferIntegrateSurf));\n checkCudaErrors(cudaDestroyTextureObject(transferLayerPreintTex));\n checkCudaErrors(cudaDestroySurfaceObject(transferLayerPreintSurf));\n checkCudaErrors(cudaFreeArray(d_transferFunc));\n checkCudaErrors(cudaFreeArray(d_transferArray));\n checkCudaErrors(cudaFreeArray(d_transferIntegrate));\n d_transferArray = 0;\n d_transferFunc = 0;\n d_transferIntegrate = 0;\n}\n\nvoid VolumeRender_setPreIntegrated(int state) { usePreInt = !!state; }\n\nvoid VolumeRender_render(dim3 gridSize,\n dim3 blockSize,\n uint *d_output,\n uint imageW,\n uint imageH,\n float density,\n float brightness,\n float transferOffset,\n float transferScale,\n cudaTextureObject_t volumeTex)\n{\n if (usePreInt) {\n d_render_preint<<>>(d_output,\n imageW,\n imageH,\n density,\n brightness,\n transferOffset,\n transferScale,\n volumeTex,\n transferTex,\n transferLayerPreintTex);\n }\n else {\n d_render_preint_off<<>>(d_output,\n imageW,\n imageH,\n density,\n brightness,\n transferOffset,\n transferScale,\n volumeTex,\n transferTex,\n transferLayerPreintTex);\n }\n}\n\nvoid VolumeRender_copyInvViewMatrix(float *invViewMatrix, size_t sizeofMatrix)\n{\n checkCudaErrors(cudaMemcpyToSymbol(c_invViewMatrix, invViewMatrix, sizeofMatrix));\n}\n\n#endif // #ifndef _VOLUMERENDER_KERNEL_CU_\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Simple 3D volume renderer\n\n#include \"hip/hip_runtime.h\"\n#ifndef _VOLUMERENDER_KERNEL_CU_\n#define _VOLUMERENDER_KERNEL_CU_\n\n#include \n#include \n\n#include \"volumeRender.h\"\n\n#define VOLUMERENDER_TFS 2\n#define VOLUMERENDER_TF_PREINTSIZE 1024\n#define VOLUMERENDER_TF_PREINTSTEPS 1024\n#define VOLUMERENDER_TF_PREINTRAY 4\n\nenum TFMode {\n TF_SINGLE_1D = 0, // single 1D TF for everything\n TF_LAYERED_2D_PREINT = 1, // layered 2D TF uses pre-integration\n TF_LAYERED_2D = 2, // layered 2D TF without pre-integration behavior\n};\n\ntypedef unsigned int uint;\ntypedef unsigned char uchar;\n\nstatic bool usePreInt = true;\nstatic hipArray *d_transferIntegrate = 0;\nstatic hipArray *d_transferFunc = 0;\nstatic hipArray *d_transferArray = 0;\n\n// 1D transfer function texture\nhipTextureObject_t transferTex;\n// 1D transfer integration texture\nhipTextureObject_t transferIntegrateTex;\nhipSurfaceObject_t transferIntegrateSurf;\n// 2D layered preintegrated transfer function texture\nhipTextureObject_t transferLayerPreintTex;\nhipSurfaceObject_t transferLayerPreintSurf;\n\ntypedef struct\n{\n float4 m[3];\n} float3x4;\n\n__constant__ float3x4 c_invViewMatrix; // inverse view matrix\n\nstruct Ray\n{\n float3 o; // origin\n float3 d; // direction\n};\n\n// intersect ray with a box\n// http://www.siggraph.org/education/materials/HyperGraph/raytrace/rtinter3.htm\n\n__device__ int intersectBox(Ray r, float3 boxmin, float3 boxmax, float *tnear, float *tfar)\n{\n // compute intersection of ray with all six bbox planes\n float3 invR = make_float3(1.0f) / r.d;\n float3 tbot = invR * (boxmin - r.o);\n float3 ttop = invR * (boxmax - r.o);\n\n // re-order intersections to find smallest and largest on each axis\n float3 tmin = fminf(ttop, tbot);\n float3 tmax = fmaxf(ttop, tbot);\n\n // find the largest tmin and the smallest tmax\n float largest_tmin = fmaxf(fmaxf(tmin.x, tmin.y), fmaxf(tmin.x, tmin.z));\n float smallest_tmax = fminf(fminf(tmax.x, tmax.y), fminf(tmax.x, tmax.z));\n\n *tnear = largest_tmin;\n *tfar = smallest_tmax;\n\n return smallest_tmax > largest_tmin;\n}\n\n// transform vector by matrix (no translation)\n__device__ float3 mul(const float3x4 &M, const float3 &v)\n{\n float3 r;\n r.x = dot(v, make_float3(M.m[0]));\n r.y = dot(v, make_float3(M.m[1]));\n r.z = dot(v, make_float3(M.m[2]));\n return r;\n}\n\n// transform vector by matrix with translation\n__device__ float4 mul(const float3x4 &M, const float4 &v)\n{\n float4 r;\n r.x = dot(v, M.m[0]);\n r.y = dot(v, M.m[1]);\n r.z = dot(v, M.m[2]);\n r.w = 1.0f;\n return r;\n}\n\n__device__ uint rgbaFloatToInt(float4 rgba)\n{\n rgba.x = __saturatef(rgba.x); // clamp to [0.0, 1.0]\n rgba.y = __saturatef(rgba.y);\n rgba.z = __saturatef(rgba.z);\n rgba.w = __saturatef(rgba.w);\n return (uint(rgba.w * 255) << 24) | (uint(rgba.z * 255) << 16) | (uint(rgba.y * 255) << 8) | uint(rgba.x * 255);\n}\n\ntemplate \n__device__ void d_render(uint *d_output,\n uint imageW,\n uint imageH,\n float density,\n float brightness,\n float transferOffset,\n float transferScale,\n hipTextureObject_t volumeTex,\n hipTextureObject_t transferTex,\n hipTextureObject_t transferLayerPreintTex,\n float transferWeight = 0.0f)\n{\n const float rayscale = float(TFMODE != TF_SINGLE_1D ? VOLUMERENDER_TF_PREINTRAY : 1);\n const int maxSteps = 512;\n const float tstep = 0.01f * rayscale;\n const float opacityThreshold = 0.95f;\n const float3 boxMin = make_float3(-1.0f, -1.0f, -1.0f);\n const float3 boxMax = make_float3(1.0f, 1.0f, 1.0f);\n\n density *= rayscale;\n\n uint x = blockIdx.x * blockDim.x + threadIdx.x;\n uint y = blockIdx.y * blockDim.y + threadIdx.y;\n\n if ((x >= imageW) || (y >= imageH))\n return;\n\n float u = (x / (float)imageW) * 2.0f - 1.0f;\n float v = (y / (float)imageH) * 2.0f - 1.0f;\n\n // calculate eye ray in world space\n Ray eyeRay;\n eyeRay.o = make_float3(mul(c_invViewMatrix, make_float4(0.0f, 0.0f, 0.0f, 1.0f)));\n eyeRay.d = normalize(make_float3(u, v, -2.0f));\n eyeRay.d = mul(c_invViewMatrix, eyeRay.d);\n\n // find intersection with box\n float tnear, tfar;\n int hit = intersectBox(eyeRay, boxMin, boxMax, &tnear, &tfar);\n\n if (!hit)\n return;\n\n if (tnear < 0.0f)\n tnear = 0.0f; // clamp to near plane\n\n // march along ray from front to back, accumulating color\n float4 sum = make_float4(0.0f);\n float t = tnear;\n float3 pos = eyeRay.o + eyeRay.d * tnear;\n float3 step = eyeRay.d * tstep;\n\n float lastsample = 0;\n\n // lastsample = (lastsample-transferOffset)*transferScale;\n for (int i = 0; i < maxSteps; i++) {\n // read from 3D texture\n // remap position to [0, 1] coordinates\n float3 coord = make_float3(pos.x * 0.5f + 0.5f, pos.y * 0.5f + 0.5f, pos.z * 0.5f + 0.5f);\n float sample = tex3D(volumeTex, coord.x, coord.y, coord.z);\n // sample = (sample-transferOffset)*transferScale;\n // sample *= 64.0f; // scale for 10-bit data\n\n // lookup in transfer function texture\n float4 col;\n int tfid = (pos.x < 0);\n\n if (TFMODE != TF_SINGLE_1D) {\n col = tex2DLayered(\n transferLayerPreintTex, sample, TFMODE == TF_LAYERED_2D ? sample : lastsample, tfid);\n col.w *= density;\n lastsample = sample;\n }\n else {\n col = tex1D(transferTex, sample);\n col.w *= 0;\n }\n\n // \"under\" operator for back-to-front blending\n // sum = lerp(sum, col, col.w);\n\n // pre-multiply alpha\n col.x *= col.w;\n col.y *= col.w;\n col.z *= col.w;\n // \"over\" operator for front-to-back blending\n sum = sum + col * (1.0f - sum.w);\n\n // exit early if opaque\n if (sum.w > opacityThreshold)\n break;\n\n t += tstep;\n\n if (t > tfar)\n break;\n\n pos += step;\n }\n\n sum *= brightness;\n\n // write output color\n d_output[y * imageW + x] = rgbaFloatToInt(sum);\n}\n\n__global__ void d_render_regular(uint *d_output,\n uint imageW,\n uint imageH,\n float density,\n float brightness,\n float transferOffset,\n float transferScale,\n hipTextureObject_t volumeTex,\n hipTextureObject_t transferTex,\n hipTextureObject_t transferLayerPreintTex,\n float transferWeight = 0.0f)\n{\n d_render(d_output,\n imageW,\n imageH,\n density,\n brightness,\n transferOffset,\n transferScale,\n volumeTex,\n transferTex,\n transferLayerPreintTex,\n transferWeight);\n}\n\n__global__ void d_render_preint(uint *d_output,\n uint imageW,\n uint imageH,\n float density,\n float brightness,\n float transferOffset,\n float transferScale,\n hipTextureObject_t volumeTex,\n hipTextureObject_t transferTex,\n hipTextureObject_t transferLayerPreintTex,\n float transferWeight = 0.0f)\n{\n d_render(d_output,\n imageW,\n imageH,\n density,\n brightness,\n transferOffset,\n transferScale,\n volumeTex,\n transferTex,\n transferLayerPreintTex,\n transferWeight);\n}\n\n__global__ void d_render_preint_off(uint *d_output,\n uint imageW,\n uint imageH,\n float density,\n float brightness,\n float transferOffset,\n float transferScale,\n hipTextureObject_t volumeTex,\n hipTextureObject_t transferTex,\n hipTextureObject_t transferLayerPreintTex,\n float transferWeight = 0.0f)\n{\n d_render(d_output,\n imageW,\n imageH,\n density,\n brightness,\n transferOffset,\n transferScale,\n volumeTex,\n transferTex,\n transferLayerPreintTex,\n transferWeight);\n}\n\n//////////////////////////////////////////////////////////////////////////\n\n__global__ void\nd_integrate_trapezoidal(hipExtent extent, hipTextureObject_t transferTex, hipSurfaceObject_t transferIntegrateSurf)\n{\n uint x = blockIdx.x * blockDim.x + threadIdx.x;\n\n // for higher speed could use hierarchical approach for sum\n if (x >= extent.width) {\n return;\n }\n\n float stepsize = 1.0 / float(extent.width - 1);\n float to = float(x) * stepsize;\n\n float4 outclr = make_float4(0, 0, 0, 0);\n float incr = stepsize;\n\n float4 lastval = tex1D(transferTex, 0);\n\n float cur = incr;\n\n while (cur < to + incr * 0.5) {\n float4 val = tex1D(transferTex, cur);\n float4 trapezoid = (lastval + val) / 2.0f;\n lastval = val;\n\n outclr += trapezoid;\n cur += incr;\n }\n\n // surface writes need byte offsets for x!\n surf1Dwrite(outclr, transferIntegrateSurf, x * sizeof(float4));\n}\n\n__global__ void d_preintegrate(int layer,\n float steps,\n hipExtent extent,\n hipTextureObject_t transferTex,\n hipTextureObject_t transferIntegrateTex,\n hipSurfaceObject_t transferLayerPreintSurf)\n{\n uint x = blockIdx.x * blockDim.x + threadIdx.x;\n uint y = blockIdx.y * blockDim.y + threadIdx.y;\n\n if (x >= extent.width || y >= extent.height) {\n return;\n }\n\n float sx = float(x) / float(extent.width);\n float sy = float(y) / float(extent.height);\n\n float smax = max(sx, sy);\n float smin = min(sx, sy);\n\n float4 iv;\n\n if (x != y) {\n // assumes square textures!\n float fracc = smax - smin;\n fracc = 1.0 / (fracc * steps);\n\n float4 intmax = tex1D(transferIntegrateTex, smax);\n float4 intmin = tex1D(transferIntegrateTex, smin);\n iv.x = (intmax.x - intmin.x) * fracc;\n iv.y = (intmax.y - intmin.y) * fracc;\n iv.z = (intmax.z - intmin.z) * fracc;\n // iv.w = (intmax.w - intmin.w)*fracc;\n iv.w = (1.0 - exp(-(intmax.w - intmin.w) * fracc));\n }\n else {\n float4 sample = tex1D(transferTex, smin);\n iv.x = sample.x;\n iv.y = sample.y;\n iv.z = sample.z;\n // iv.w = sample.w;\n iv.w = (1.0 - exp(-sample.w));\n }\n\n iv.x = __saturatef(iv.x);\n iv.y = __saturatef(iv.y);\n iv.z = __saturatef(iv.z);\n iv.w = __saturatef(iv.w);\n\n // surface writes need byte offsets for x!\n surf2DLayeredwrite(iv, transferLayerPreintSurf, x * sizeof(float4), y, layer);\n}\n\n//////////////////////////////////////////////////////////////////////////\n\nvoid VolumeRender_setTextureFilterMode(bool bLinearFilter, Volume *vol)\n{\n if (vol->volumeTex) {\n checkCudaErrors(hipDestroyTextureObject(vol->volumeTex));\n }\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = vol->content;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = bLinearFilter ? hipFilterModeLinear : hipFilterModePoint;\n\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.addressMode[2] = hipAddressModeWrap;\n\n texDescr.readMode = VolumeTypeInfo::readMode;\n\n checkCudaErrors(hipCreateTextureObject(&vol->volumeTex, &texRes, &texDescr, NULL));\n}\n\nstatic unsigned int iDivUp(size_t a, size_t b)\n{\n size_t val = (a % b != 0) ? (a / b + 1) : (a / b);\n if (val > UINT_MAX) {\n fprintf(stderr, \"\\nUINT_MAX limit exceeded in iDivUp() exiting.....\\n\");\n exit(EXIT_FAILURE); // val exceeds limit\n }\n\n return static_cast(val);\n}\n\nvoid VolumeRender_updateTF(int tfIdx, int numColors, float4 *colors)\n{\n if (d_transferFunc) {\n checkCudaErrors(hipFreeArray(d_transferFunc));\n d_transferFunc = 0;\n }\n\n hipChannelFormatDesc channelFloat4 = hipCreateChannelDesc();\n checkCudaErrors(hipMallocArray(&d_transferFunc, &channelFloat4, numColors, 1));\n checkCudaErrors(\n hipMemcpy2DToArray(d_transferFunc, 0, 0, colors, 0, sizeof(float4) * numColors, 1, hipMemcpyHostToDevice));\n\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = d_transferFunc;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = hipFilterModeLinear;\n texDescr.addressMode[0] = hipAddressModeClamp;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&transferTex, &texRes, &texDescr, NULL));\n\n if (tfIdx < 0 || tfIdx >= VOLUMERENDER_TFS) {\n return;\n }\n\n {\n hipExtent extent = {VOLUMERENDER_TF_PREINTSTEPS, 0, 0};\n dim3 blockSize(32, 1, 1);\n dim3 gridSize(iDivUp(extent.width, blockSize.x), 1, 1);\n d_integrate_trapezoidal<<>>(extent, transferTex, transferIntegrateSurf);\n }\n\n {\n hipExtent extent = {VOLUMERENDER_TF_PREINTSIZE, VOLUMERENDER_TF_PREINTSIZE, VOLUMERENDER_TFS};\n dim3 blockSize(16, 16, 1);\n dim3 gridSize(iDivUp(extent.width, blockSize.x), iDivUp(extent.height, blockSize.y), 1);\n d_preintegrate<<>>(tfIdx,\n float(VOLUMERENDER_TF_PREINTSTEPS),\n extent,\n transferTex,\n transferIntegrateTex,\n transferLayerPreintSurf);\n }\n}\n\nvoid VolumeRender_init()\n{\n hipResourceDesc texRes;\n hipTextureDesc texDescr;\n\n hipChannelFormatDesc channelFloat4 = hipCreateChannelDesc();\n hipExtent extent = {VOLUMERENDER_TF_PREINTSIZE, VOLUMERENDER_TF_PREINTSIZE, VOLUMERENDER_TFS};\n checkCudaErrors(\n hipMalloc3DArray(&d_transferArray, &channelFloat4, extent, hipArrayLayered | hipArraySurfaceLoadStore));\n\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = d_transferArray;\n\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = hipFilterModeLinear;\n texDescr.addressMode[0] = hipAddressModeClamp;\n texDescr.addressMode[1] = hipAddressModeClamp;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&transferLayerPreintTex, &texRes, &texDescr, NULL));\n\n hipResourceDesc surfRes;\n memset(&surfRes, 0, sizeof(hipResourceDesc));\n surfRes.resType = hipResourceTypeArray;\n surfRes.res.array.array = d_transferArray;\n\n checkCudaErrors(hipCreateSurfaceObject(&transferLayerPreintSurf, &surfRes));\n\n checkCudaErrors(hipMallocArray(\n &d_transferIntegrate, &channelFloat4, VOLUMERENDER_TF_PREINTSTEPS, 0, hipArraySurfaceLoadStore));\n\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = d_transferIntegrate;\n\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = hipFilterModeLinear;\n\n texDescr.addressMode[0] = hipAddressModeClamp;\n texDescr.addressMode[1] = hipAddressModeClamp;\n texDescr.addressMode[2] = hipAddressModeClamp;\n\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&transferIntegrateTex, &texRes, &texDescr, NULL));\n\n memset(&surfRes, 0, sizeof(hipResourceDesc));\n surfRes.resType = hipResourceTypeArray;\n surfRes.res.array.array = d_transferIntegrate;\n\n checkCudaErrors(hipCreateSurfaceObject(&transferIntegrateSurf, &surfRes));\n\n // create transfer function texture\n float4 transferFunc0[] = {\n {\n 0.0,\n 0.0,\n 0.0,\n 0.0,\n },\n {\n 1.0,\n 0.0,\n 0.0,\n 1.0,\n },\n {\n 1.0,\n 0.5,\n 0.0,\n 1.0,\n },\n {\n 1.0,\n 1.0,\n 0.0,\n 1.0,\n },\n {\n 0.0,\n 1.0,\n 0.0,\n 1.0,\n },\n {\n 0.0,\n 1.0,\n 1.0,\n 1.0,\n },\n {\n 0.0,\n 0.0,\n 1.0,\n 1.0,\n },\n {\n 1.0,\n 0.0,\n 1.0,\n 1.0,\n },\n {\n 0.0,\n 0.0,\n 0.0,\n 0.0,\n },\n };\n\n float4 transferFunc1[] = {\n {\n 0.0,\n 0.0,\n 0.0,\n 0.0,\n },\n {\n 0.0,\n 1.0,\n 0.0,\n 0.125,\n },\n {\n 0.0,\n 0.5,\n 1.0,\n 0.125,\n },\n {\n 0.0,\n 1.0,\n 1.0,\n 0.125,\n },\n {\n 0.0,\n 1.0,\n 0.0,\n 0.125,\n },\n {\n 0.25,\n 0.75,\n 0.0,\n 1.0,\n },\n {\n 0.75,\n 0.25,\n 0.0,\n 0.125,\n },\n {\n 1.0,\n 0.75,\n 0.0,\n 0.125,\n },\n {\n 0.0,\n 0.0,\n 0.0,\n 0.0,\n },\n };\n\n VolumeRender_updateTF(1, sizeof(transferFunc1) / sizeof(float4), transferFunc1);\n VolumeRender_updateTF(0, sizeof(transferFunc0) / sizeof(float4), transferFunc0);\n}\n\nvoid VolumeRender_deinit()\n{\n checkCudaErrors(hipDestroyTextureObject(transferTex));\n checkCudaErrors(hipDestroyTextureObject(transferIntegrateTex));\n checkCudaErrors(hipDestroySurfaceObject(transferIntegrateSurf));\n checkCudaErrors(hipDestroyTextureObject(transferLayerPreintTex));\n checkCudaErrors(hipDestroySurfaceObject(transferLayerPreintSurf));\n checkCudaErrors(hipFreeArray(d_transferFunc));\n checkCudaErrors(hipFreeArray(d_transferArray));\n checkCudaErrors(hipFreeArray(d_transferIntegrate));\n d_transferArray = 0;\n d_transferFunc = 0;\n d_transferIntegrate = 0;\n}\n\nvoid VolumeRender_setPreIntegrated(int state) { usePreInt = !!state; }\n\nvoid VolumeRender_render(dim3 gridSize,\n dim3 blockSize,\n uint *d_output,\n uint imageW,\n uint imageH,\n float density,\n float brightness,\n float transferOffset,\n float transferScale,\n hipTextureObject_t volumeTex)\n{\n if (usePreInt) {\n d_render_preint<<>>(d_output,\n imageW,\n imageH,\n density,\n brightness,\n transferOffset,\n transferScale,\n volumeTex,\n transferTex,\n transferLayerPreintTex);\n }\n else {\n d_render_preint_off<<>>(d_output,\n imageW,\n imageH,\n density,\n brightness,\n transferOffset,\n transferScale,\n volumeTex,\n transferTex,\n transferLayerPreintTex);\n }\n}\n\nvoid VolumeRender_copyInvViewMatrix(float *invViewMatrix, size_t sizeofMatrix)\n{\n checkCudaErrors(hipMemcpyToSymbol(HIP_SYMBOL(c_invViewMatrix), invViewMatrix, sizeofMatrix));\n}\n\n#endif // #ifndef _VOLUMERENDER_KERNEL_CU_\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/volumeFiltering/volumeRender_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n\n#include \"ShaderStructs.h\"\n#include \"helper_cuda.h\"\n#include \"sinewave_cuda.h\"\n\n__global__ void sinewave_gen_kernel(Vertex *vertices, unsigned int width, unsigned int height, float time)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // calculate uv coordinates\n float u = x / (float)width;\n float v = y / (float)height;\n u = u * 2.0f - 1.0f;\n v = v * 2.0f - 1.0f;\n\n // calculate simple sine wave pattern\n float freq = 4.0f;\n float w = sinf(u * freq + time) * cosf(v * freq + time) * 0.5f;\n\n if (y < height && x < width) {\n // write output vertex\n vertices[y * width + x].position.x = u;\n vertices[y * width + x].position.y = w;\n vertices[y * width + x].position.z = v;\n vertices[y * width + x].color.x = 1.0f;\n vertices[y * width + x].color.y = 0.0f;\n vertices[y * width + x].color.z = 0.0f;\n vertices[y * width + x].color.w = 0.0f;\n }\n}\n\nVertex *cudaImportVertexBuffer(void *sharedHandle, cudaExternalMemory_t &externalMemory, int meshWidth, int meshHeight)\n{\n cudaExternalMemoryHandleDesc externalMemoryHandleDesc;\n memset(&externalMemoryHandleDesc, 0, sizeof(externalMemoryHandleDesc));\n\n externalMemoryHandleDesc.type = cudaExternalMemoryHandleTypeD3D11ResourceKmt;\n externalMemoryHandleDesc.size = sizeof(Vertex) * meshHeight * meshWidth;\n externalMemoryHandleDesc.flags = cudaExternalMemoryDedicated;\n externalMemoryHandleDesc.handle.win32.handle = sharedHandle;\n\n checkCudaErrors(cudaImportExternalMemory(&externalMemory, &externalMemoryHandleDesc));\n\n cudaExternalMemoryBufferDesc externalMemoryBufferDesc;\n memset(&externalMemoryBufferDesc, 0, sizeof(externalMemoryBufferDesc));\n externalMemoryBufferDesc.offset = 0;\n externalMemoryBufferDesc.size = sizeof(Vertex) * meshHeight * meshWidth;\n externalMemoryBufferDesc.flags = 0;\n\n Vertex *cudaDevVertptr = NULL;\n checkCudaErrors(\n cudaExternalMemoryGetMappedBuffer((void **)&cudaDevVertptr, externalMemory, &externalMemoryBufferDesc));\n\n return cudaDevVertptr;\n}\n\nvoid cudaImportKeyedMutex(void *sharedHandle, cudaExternalSemaphore_t &extSemaphore)\n{\n cudaExternalSemaphoreHandleDesc extSemaDesc;\n memset(&extSemaDesc, 0, sizeof(extSemaDesc));\n extSemaDesc.type = cudaExternalSemaphoreHandleTypeKeyedMutexKmt;\n extSemaDesc.handle.win32.handle = sharedHandle;\n extSemaDesc.flags = 0;\n\n checkCudaErrors(cudaImportExternalSemaphore(&extSemaphore, &extSemaDesc));\n}\n\nvoid cudaAcquireSync(cudaExternalSemaphore_t &extSemaphore,\n uint64_t key,\n unsigned int timeoutMs,\n cudaStream_t streamToRun)\n{\n cudaExternalSemaphoreWaitParams extSemWaitParams;\n memset(&extSemWaitParams, 0, sizeof(extSemWaitParams));\n extSemWaitParams.params.keyedMutex.key = key;\n extSemWaitParams.params.keyedMutex.timeoutMs = timeoutMs;\n\n checkCudaErrors(cudaWaitExternalSemaphoresAsync(&extSemaphore, &extSemWaitParams, 1, streamToRun));\n}\n\nvoid cudaReleaseSync(cudaExternalSemaphore_t &extSemaphore, uint64_t key, cudaStream_t streamToRun)\n{\n cudaExternalSemaphoreSignalParams extSemSigParams;\n memset(&extSemSigParams, 0, sizeof(extSemSigParams));\n extSemSigParams.params.keyedMutex.key = key;\n\n checkCudaErrors(cudaSignalExternalSemaphoresAsync(&extSemaphore, &extSemSigParams, 1, streamToRun));\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run the Cuda part of the computation\n////////////////////////////////////////////////////////////////////////////////\nvoid RunSineWaveKernel(cudaExternalSemaphore_t &extSemaphore,\n uint64_t &key,\n unsigned int timeoutMs,\n size_t mesh_width,\n size_t mesh_height,\n Vertex *cudaDevVertptr,\n cudaStream_t streamToRun)\n{\n static float t = 0.0f;\n cudaAcquireSync(extSemaphore, key++, timeoutMs, streamToRun);\n\n dim3 block(16, 16, 1);\n dim3 grid(mesh_width / 16, mesh_height / 16, 1);\n sinewave_gen_kernel<<>>(cudaDevVertptr, mesh_width, mesh_height, t);\n getLastCudaError(\"sinewave_gen_kernel execution failed.\\n\");\n\n cudaReleaseSync(extSemaphore, key, streamToRun);\n t += 0.01f;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \"ShaderStructs.h\"\n#include \"hip/hip_runtime_api.h\"\n#include \"sinewave_cuda.h\"\n\n__global__ void sinewave_gen_kernel(Vertex *vertices, unsigned int width, unsigned int height, float time)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // calculate uv coordinates\n float u = x / (float)width;\n float v = y / (float)height;\n u = u * 2.0f - 1.0f;\n v = v * 2.0f - 1.0f;\n\n // calculate simple sine wave pattern\n float freq = 4.0f;\n float w = sinf(u * freq + time) * cosf(v * freq + time) * 0.5f;\n\n if (y < height && x < width) {\n // write output vertex\n vertices[y * width + x].position.x = u;\n vertices[y * width + x].position.y = w;\n vertices[y * width + x].position.z = v;\n vertices[y * width + x].color.x = 1.0f;\n vertices[y * width + x].color.y = 0.0f;\n vertices[y * width + x].color.z = 0.0f;\n vertices[y * width + x].color.w = 0.0f;\n }\n}\n\nVertex *cudaImportVertexBuffer(void *sharedHandle, hipExternalMemory_t &externalMemory, int meshWidth, int meshHeight)\n{\n hipExternalMemoryHandleDesc externalMemoryHandleDesc;\n memset(&externalMemoryHandleDesc, 0, sizeof(externalMemoryHandleDesc));\n\n externalMemoryHandleDesc.type = hipExternalMemoryHandleTypeD3D11ResourceKmt;\n externalMemoryHandleDesc.size = sizeof(Vertex) * meshHeight * meshWidth;\n externalMemoryHandleDesc.flags = hipExternalMemoryDedicated;\n externalMemoryHandleDesc.handle.win32.handle = sharedHandle;\n\n checkCudaErrors(hipImportExternalMemory(&externalMemory, &externalMemoryHandleDesc));\n\n hipExternalMemoryBufferDesc externalMemoryBufferDesc;\n memset(&externalMemoryBufferDesc, 0, sizeof(externalMemoryBufferDesc));\n externalMemoryBufferDesc.offset = 0;\n externalMemoryBufferDesc.size = sizeof(Vertex) * meshHeight * meshWidth;\n externalMemoryBufferDesc.flags = 0;\n\n Vertex *cudaDevVertptr = NULL;\n checkCudaErrors(\n hipExternalMemoryGetMappedBuffer((void **)&cudaDevVertptr, externalMemory, &externalMemoryBufferDesc));\n\n return cudaDevVertptr;\n}\n\nvoid cudaImportKeyedMutex(void *sharedHandle, hipExternalSemaphore_t &extSemaphore)\n{\n hipExternalSemaphoreHandleDesc extSemaDesc;\n memset(&extSemaDesc, 0, sizeof(extSemaDesc));\n extSemaDesc.type = cudaExternalSemaphoreHandleTypeKeyedMutexKmt;\n extSemaDesc.handle.win32.handle = sharedHandle;\n extSemaDesc.flags = 0;\n\n checkCudaErrors(hipImportExternalSemaphore(&extSemaphore, &extSemaDesc));\n}\n\nvoid cudaAcquireSync(hipExternalSemaphore_t &extSemaphore,\n uint64_t key,\n unsigned int timeoutMs,\n hipStream_t streamToRun)\n{\n hipExternalSemaphoreWaitParams extSemWaitParams;\n memset(&extSemWaitParams, 0, sizeof(extSemWaitParams));\n extSemWaitParams.params.keyedMutex.key = key;\n extSemWaitParams.params.keyedMutex.timeoutMs = timeoutMs;\n\n checkCudaErrors(hipWaitExternalSemaphoresAsync(&extSemaphore, &extSemWaitParams, 1, streamToRun));\n}\n\nvoid cudaReleaseSync(hipExternalSemaphore_t &extSemaphore, uint64_t key, hipStream_t streamToRun)\n{\n hipExternalSemaphoreSignalParams extSemSigParams;\n memset(&extSemSigParams, 0, sizeof(extSemSigParams));\n extSemSigParams.params.keyedMutex.key = key;\n\n checkCudaErrors(hipSignalExternalSemaphoresAsync(&extSemaphore, &extSemSigParams, 1, streamToRun));\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run the Cuda part of the computation\n////////////////////////////////////////////////////////////////////////////////\nvoid RunSineWaveKernel(hipExternalSemaphore_t &extSemaphore,\n uint64_t &key,\n unsigned int timeoutMs,\n size_t mesh_width,\n size_t mesh_height,\n Vertex *cudaDevVertptr,\n hipStream_t streamToRun)\n{\n static float t = 0.0f;\n cudaAcquireSync(extSemaphore, key++, timeoutMs, streamToRun);\n\n dim3 block(16, 16, 1);\n dim3 grid(mesh_width / 16, mesh_height / 16, 1);\n sinewave_gen_kernel<<>>(cudaDevVertptr, mesh_width, mesh_height, t);\n getLastCudaError(\"sinewave_gen_kernel execution failed.\\n\");\n\n cudaReleaseSync(extSemaphore, key, streamToRun);\n t += 0.01f;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/simpleD3D11/sinewave_cuda.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * See: https://www.piday.org/million/\n */\n\n#include \n\n#include \"MonteCarloPi.h\"\n#define CUDA_DRIVER_API\n#include \n#include \n\n#define ROUND_UP_TO_GRANULARITY(x, n) (((x + n - 1) / n) * n)\n\n// `ipcHandleTypeFlag` specifies the platform specific handle type this sample\n// uses for importing and exporting memory allocation. On Linux this sample\n// specifies the type as CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR meaning that\n// file descriptors will be used. On Windows this sample specifies the type as\n// CU_MEM_HANDLE_TYPE_WIN32 meaning that NT HANDLEs will be used. The\n// ipcHandleTypeFlag variable is a convenience variable and is passed by value\n// to individual requests.\n#if defined(__linux__)\nCUmemAllocationHandleType ipcHandleTypeFlag = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;\n#else\nCUmemAllocationHandleType ipcHandleTypeFlag = CU_MEM_HANDLE_TYPE_WIN32;\n#endif\n\n// Windows-specific LPSECURITYATTRIBUTES\nvoid getDefaultSecurityDescriptor(CUmemAllocationProp *prop)\n{\n#if defined(__linux__)\n return;\n#elif defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)\n static const char sddl[] = \"D:P(OA;;GARCSDWDWOCCDCLCSWLODTWPRPCRFA;;;WD)\";\n static OBJECT_ATTRIBUTES objAttributes;\n static bool objAttributesConfigured = false;\n\n if (!objAttributesConfigured) {\n PSECURITY_DESCRIPTOR secDesc;\n BOOL result = ConvertStringSecurityDescriptorToSecurityDescriptorA(sddl, SDDL_REVISION_1, &secDesc, NULL);\n if (result == 0) {\n printf(\"IPC failure: getDefaultSecurityDescriptor Failed! (%d)\\n\", GetLastError());\n }\n\n InitializeObjectAttributes(&objAttributes, NULL, 0, NULL, secDesc);\n\n objAttributesConfigured = true;\n }\n\n prop->win32HandleMetaData = &objAttributes;\n return;\n#endif\n}\n\n__global__ void monte_carlo_kernel(vec2 *xyVector,\n float *pointsInsideCircle,\n float *numPointsInCircle,\n unsigned int numPoints,\n float time)\n{\n const size_t stride = gridDim.x * blockDim.x;\n size_t tid = blockIdx.x * blockDim.x + threadIdx.x;\n float count = 0.0f;\n\n curandState rgnState;\n curand_init((unsigned long long)time, tid, 0, &rgnState);\n\n for (; tid < numPoints; tid += stride) {\n float x = curand_uniform(&rgnState);\n float y = curand_uniform(&rgnState);\n x = (2.0f * x) - 1.0f;\n y = (2.0f * y) - 1.0f;\n xyVector[tid][0] = x;\n xyVector[tid][1] = y;\n\n // Compute the distance of this point form the center(0, 0)\n float dist = sqrtf((x * x) + (y * y));\n\n // If distance is less than the radius of the unit circle, the point lies in\n // the circle.\n pointsInsideCircle[tid] = (dist <= 1.0f);\n count += (dist <= 1.0f);\n }\n atomicAdd(numPointsInCircle, count);\n}\n\nMonteCarloPiSimulation::MonteCarloPiSimulation(size_t num_points)\n : m_xyVector(nullptr)\n , m_pointsInsideCircle(nullptr)\n , m_totalPointsInsideCircle(0)\n , m_totalPointsSimulated(0)\n , m_numPoints(num_points)\n{\n}\n\nMonteCarloPiSimulation::~MonteCarloPiSimulation()\n{\n if (m_numPointsInCircle) {\n checkCudaErrors(cudaFree(m_numPointsInCircle));\n m_numPointsInCircle = nullptr;\n }\n if (m_hostNumPointsInCircle) {\n checkCudaErrors(cudaFreeHost(m_hostNumPointsInCircle));\n m_hostNumPointsInCircle = nullptr;\n }\n\n cleanupSimulationAllocations();\n}\n\nvoid MonteCarloPiSimulation::initSimulation(int cudaDevice, cudaStream_t stream)\n{\n m_cudaDevice = cudaDevice;\n getIdealExecutionConfiguration();\n\n // Allocate a position buffer that contains random location of the points in\n // XY cartesian plane.\n // Allocate a bitmap buffer which holds information of whether a point in the\n // position buffer is inside the unit circle or not.\n setupSimulationAllocations();\n\n checkCudaErrors(cudaMalloc((float **)&m_numPointsInCircle, sizeof(*m_numPointsInCircle)));\n checkCudaErrors(cudaMallocHost((float **)&m_hostNumPointsInCircle, sizeof(*m_hostNumPointsInCircle)));\n}\n\nvoid MonteCarloPiSimulation::stepSimulation(float time, cudaStream_t stream)\n{\n checkCudaErrors(cudaMemsetAsync(m_numPointsInCircle, 0, sizeof(*m_numPointsInCircle), stream));\n\n monte_carlo_kernel<<>>(\n m_xyVector, m_pointsInsideCircle, m_numPointsInCircle, m_numPoints, time);\n getLastCudaError(\"Failed to launch CUDA simulation\");\n\n checkCudaErrors(cudaMemcpyAsync(\n m_hostNumPointsInCircle, m_numPointsInCircle, sizeof(*m_numPointsInCircle), cudaMemcpyDeviceToHost, stream));\n\n // Queue up a stream callback to compute and print the PI value.\n checkCudaErrors(cudaLaunchHostFunc(stream, this->computePiCallback, (void *)this));\n}\n\nvoid MonteCarloPiSimulation::computePiCallback(void *args)\n{\n MonteCarloPiSimulation *cbData = (MonteCarloPiSimulation *)args;\n cbData->m_totalPointsInsideCircle += *(cbData->m_hostNumPointsInCircle);\n cbData->m_totalPointsSimulated += cbData->m_numPoints;\n double piValue = 4.0 * ((double)cbData->m_totalPointsInsideCircle / (double)cbData->m_totalPointsSimulated);\n printf(\"Approximate Pi value for %zd data points: %lf \\n\", cbData->m_totalPointsSimulated, piValue);\n}\n\nvoid MonteCarloPiSimulation::getIdealExecutionConfiguration()\n{\n int warpSize = 0;\n int multiProcessorCount = 0;\n\n checkCudaErrors(cudaSetDevice(m_cudaDevice));\n checkCudaErrors(cudaDeviceGetAttribute(&warpSize, cudaDevAttrWarpSize, m_cudaDevice));\n\n // We don't need large block sizes, since there's not much inter-thread\n // communication\n m_threads = warpSize;\n\n // Use the occupancy calculator and fill the gpu as best as we can\n checkCudaErrors(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&m_blocks, monte_carlo_kernel, warpSize, 0));\n\n checkCudaErrors(cudaDeviceGetAttribute(&multiProcessorCount, cudaDevAttrMultiProcessorCount, m_cudaDevice));\n m_blocks *= multiProcessorCount;\n\n // Go ahead and the clamp the blocks to the minimum needed for this\n // height/width\n m_blocks = std::min(m_blocks, (int)((m_numPoints + m_threads - 1) / m_threads));\n}\n\nvoid MonteCarloPiSimulation::setupSimulationAllocations()\n{\n CUdeviceptr d_ptr = 0U;\n size_t granularity = 0;\n CUmemGenericAllocationHandle cudaPositionHandle, cudaInCircleHandle;\n\n CUmemAllocationProp allocProp = {};\n allocProp.type = CU_MEM_ALLOCATION_TYPE_PINNED;\n allocProp.location.type = CU_MEM_LOCATION_TYPE_DEVICE;\n allocProp.location.id = m_cudaDevice;\n allocProp.win32HandleMetaData = NULL;\n allocProp.requestedHandleTypes = ipcHandleTypeFlag;\n\n // Windows-specific LPSECURITYATTRIBUTES is required when\n // CU_MEM_HANDLE_TYPE_WIN32 is used. The security attribute defines the scope\n // of which exported allocations may be tranferred to other processes. For all\n // other handle types, pass NULL.\n getDefaultSecurityDescriptor(&allocProp);\n\n // Get the recommended granularity for m_cudaDevice.\n checkCudaErrors(cuMemGetAllocationGranularity(&granularity, &allocProp, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED));\n\n size_t xyPositionVecSize = m_numPoints * sizeof(*m_xyVector);\n size_t inCircleVecSize = m_numPoints * sizeof(*m_pointsInsideCircle);\n\n size_t xyPositionSize = ROUND_UP_TO_GRANULARITY(xyPositionVecSize, granularity);\n size_t inCircleSize = ROUND_UP_TO_GRANULARITY(inCircleVecSize, granularity);\n m_totalAllocationSize = (xyPositionSize + inCircleSize);\n\n // Reserve the required contiguous VA space for the allocations\n checkCudaErrors(cuMemAddressReserve(&d_ptr, m_totalAllocationSize, granularity, 0U, 0));\n\n // Create the allocations as a pinned allocation on this device.\n // Create an allocation to store all the positions of points on the xy plane\n // and a second allocation which stores information if the corresponding\n // position is inside the unit circle or not.\n checkCudaErrors(cuMemCreate(&cudaPositionHandle, xyPositionSize, &allocProp, 0));\n checkCudaErrors(cuMemCreate(&cudaInCircleHandle, inCircleSize, &allocProp, 0));\n\n // Export the allocation to a platform-specific handle. The type of handle\n // requested here must match the requestedHandleTypes field in the prop\n // structure passed to cuMemCreate. The handle obtained here will be passed to\n // vulkan to import the allocation.\n checkCudaErrors(\n cuMemExportToShareableHandle((void *)&m_posShareableHandle, cudaPositionHandle, ipcHandleTypeFlag, 0));\n checkCudaErrors(\n cuMemExportToShareableHandle((void *)&m_inCircleShareableHandle, cudaInCircleHandle, ipcHandleTypeFlag, 0));\n\n CUdeviceptr va_position = d_ptr;\n CUdeviceptr va_InCircle = va_position + xyPositionSize;\n m_pointsInsideCircle = (float *)va_InCircle;\n m_xyVector = (vec2 *)va_position;\n\n // Assign the chunk to the appropriate VA range\n checkCudaErrors(cuMemMap(va_position, xyPositionSize, 0, cudaPositionHandle, 0));\n checkCudaErrors(cuMemMap(va_InCircle, inCircleSize, 0, cudaInCircleHandle, 0));\n\n // Release the handles for the allocation. Since the allocation is currently\n // mapped to a VA range with a previous call to cuMemMap the actual freeing of\n // memory allocation will happen on an eventual call to cuMemUnmap. Thus the\n // allocation will be kept live until it is unmapped.\n checkCudaErrors(cuMemRelease(cudaPositionHandle));\n checkCudaErrors(cuMemRelease(cudaInCircleHandle));\n\n CUmemAccessDesc accessDescriptor = {};\n accessDescriptor.location.id = m_cudaDevice;\n accessDescriptor.location.type = CU_MEM_LOCATION_TYPE_DEVICE;\n accessDescriptor.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;\n\n // Apply the access descriptor to the whole VA range. Essentially enables\n // Read-Write access to the range.\n checkCudaErrors(cuMemSetAccess(d_ptr, m_totalAllocationSize, &accessDescriptor, 1));\n}\n\nvoid MonteCarloPiSimulation::cleanupSimulationAllocations()\n{\n if (m_xyVector && m_pointsInsideCircle) {\n // Unmap the mapped virtual memory region\n // Since the handles to the mapped backing stores have already been released\n // by cuMemRelease, and these are the only/last mappings referencing them,\n // The backing stores will be freed.\n checkCudaErrors(cuMemUnmap((CUdeviceptr)m_xyVector, m_totalAllocationSize));\n\n checkIpcErrors(ipcCloseShareableHandle(m_posShareableHandle));\n checkIpcErrors(ipcCloseShareableHandle(m_inCircleShareableHandle));\n\n // Free the virtual address region.\n checkCudaErrors(cuMemAddressFree((CUdeviceptr)m_xyVector, m_totalAllocationSize));\n\n m_xyVector = nullptr;\n m_pointsInsideCircle = nullptr;\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * See: https://www.piday.org/million/\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \"MonteCarloPi.h\"\n#define CUDA_DRIVER_API\n#include \n#include \n\n#define ROUND_UP_TO_GRANULARITY(x, n) (((x + n - 1) / n) * n)\n\n// `ipcHandleTypeFlag` specifies the platform specific handle type this sample\n// uses for importing and exporting memory allocation. On Linux this sample\n// specifies the type as CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR meaning that\n// file descriptors will be used. On Windows this sample specifies the type as\n// CU_MEM_HANDLE_TYPE_WIN32 meaning that NT HANDLEs will be used. The\n// ipcHandleTypeFlag variable is a convenience variable and is passed by value\n// to individual requests.\n#if defined(__linux__)\nhipMemAllocationHandleType ipcHandleTypeFlag = hipMemHandleTypePosixFileDescriptor;\n#else\nhipMemAllocationHandleType ipcHandleTypeFlag = hipMemHandleTypeWin32;\n#endif\n\n// Windows-specific LPSECURITYATTRIBUTES\nvoid getDefaultSecurityDescriptor(hipMemAllocationProp *prop)\n{\n#if defined(__linux__)\n return;\n#elif defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)\n static const char sddl[] = \"D:P(OA;;GARCSDWDWOCCDCLCSWLODTWPRPCRFA;;;WD)\";\n static OBJECT_ATTRIBUTES objAttributes;\n static bool objAttributesConfigured = false;\n\n if (!objAttributesConfigured) {\n PSECURITY_DESCRIPTOR secDesc;\n BOOL result = ConvertStringSecurityDescriptorToSecurityDescriptorA(sddl, SDDL_REVISION_1, &secDesc, NULL);\n if (result == 0) {\n printf(\"IPC failure: getDefaultSecurityDescriptor Failed! (%d)\\n\", GetLastError());\n }\n\n InitializeObjectAttributes(&objAttributes, NULL, 0, NULL, secDesc);\n\n objAttributesConfigured = true;\n }\n\n prop->win32HandleMetaData = &objAttributes;\n return;\n#endif\n}\n\n__global__ void monte_carlo_kernel(vec2 *xyVector,\n float *pointsInsideCircle,\n float *numPointsInCircle,\n unsigned int numPoints,\n float time)\n{\n const size_t stride = gridDim.x * blockDim.x;\n size_t tid = blockIdx.x * blockDim.x + threadIdx.x;\n float count = 0.0f;\n\n hiprandState rgnState;\n hiprand_init((unsigned long long)time, tid, 0, &rgnState);\n\n for (; tid < numPoints; tid += stride) {\n float x = hiprand_uniform(&rgnState);\n float y = hiprand_uniform(&rgnState);\n x = (2.0f * x) - 1.0f;\n y = (2.0f * y) - 1.0f;\n xyVector[tid][0] = x;\n xyVector[tid][1] = y;\n\n // Compute the distance of this point form the center(0, 0)\n float dist = sqrtf((x * x) + (y * y));\n\n // If distance is less than the radius of the unit circle, the point lies in\n // the circle.\n pointsInsideCircle[tid] = (dist <= 1.0f);\n count += (dist <= 1.0f);\n }\n atomicAdd(numPointsInCircle, count);\n}\n\nMonteCarloPiSimulation::MonteCarloPiSimulation(size_t num_points)\n : m_xyVector(nullptr)\n , m_pointsInsideCircle(nullptr)\n , m_totalPointsInsideCircle(0)\n , m_totalPointsSimulated(0)\n , m_numPoints(num_points)\n{\n}\n\nMonteCarloPiSimulation::~MonteCarloPiSimulation()\n{\n if (m_numPointsInCircle) {\n checkCudaErrors(hipFree(m_numPointsInCircle));\n m_numPointsInCircle = nullptr;\n }\n if (m_hostNumPointsInCircle) {\n checkCudaErrors(hipHostFree(m_hostNumPointsInCircle));\n m_hostNumPointsInCircle = nullptr;\n }\n\n cleanupSimulationAllocations();\n}\n\nvoid MonteCarloPiSimulation::initSimulation(int cudaDevice, hipStream_t stream)\n{\n m_cudaDevice = cudaDevice;\n getIdealExecutionConfiguration();\n\n // Allocate a position buffer that contains random location of the points in\n // XY cartesian plane.\n // Allocate a bitmap buffer which holds information of whether a point in the\n // position buffer is inside the unit circle or not.\n setupSimulationAllocations();\n\n checkCudaErrors(hipMalloc((float **)&m_numPointsInCircle, sizeof(*m_numPointsInCircle)));\n checkCudaErrors(hipHostMalloc((float **)&m_hostNumPointsInCircle, sizeof(*m_hostNumPointsInCircle)));\n}\n\nvoid MonteCarloPiSimulation::stepSimulation(float time, hipStream_t stream)\n{\n checkCudaErrors(hipMemsetAsync(m_numPointsInCircle, 0, sizeof(*m_numPointsInCircle), stream));\n\n monte_carlo_kernel<<>>(\n m_xyVector, m_pointsInsideCircle, m_numPointsInCircle, m_numPoints, time);\n getLastCudaError(\"Failed to launch CUDA simulation\");\n\n checkCudaErrors(hipMemcpyAsync(\n m_hostNumPointsInCircle, m_numPointsInCircle, sizeof(*m_numPointsInCircle), hipMemcpyDeviceToHost, stream));\n\n // Queue up a stream callback to compute and print the PI value.\n checkCudaErrors(hipLaunchHostFunc(stream, this->computePiCallback, (void *)this));\n}\n\nvoid MonteCarloPiSimulation::computePiCallback(void *args)\n{\n MonteCarloPiSimulation *cbData = (MonteCarloPiSimulation *)args;\n cbData->m_totalPointsInsideCircle += *(cbData->m_hostNumPointsInCircle);\n cbData->m_totalPointsSimulated += cbData->m_numPoints;\n double piValue = 4.0 * ((double)cbData->m_totalPointsInsideCircle / (double)cbData->m_totalPointsSimulated);\n printf(\"Approximate Pi value for %zd data points: %lf \\n\", cbData->m_totalPointsSimulated, piValue);\n}\n\nvoid MonteCarloPiSimulation::getIdealExecutionConfiguration()\n{\n int warpSize = 0;\n int multiProcessorCount = 0;\n\n checkCudaErrors(hipSetDevice(m_cudaDevice));\n checkCudaErrors(hipDeviceGetAttribute(&warpSize, hipDeviceAttributeWarpSize, m_cudaDevice));\n\n // We don't need large block sizes, since there's not much inter-thread\n // communication\n m_threads = warpSize;\n\n // Use the occupancy calculator and fill the gpu as best as we can\n checkCudaErrors(hipOccupancyMaxActiveBlocksPerMultiprocessor(&m_blocks, monte_carlo_kernel, warpSize, 0));\n\n checkCudaErrors(hipDeviceGetAttribute(&multiProcessorCount, hipDeviceAttributeMultiprocessorCount, m_cudaDevice));\n m_blocks *= multiProcessorCount;\n\n // Go ahead and the clamp the blocks to the minimum needed for this\n // height/width\n m_blocks = std::min(m_blocks, (int)((m_numPoints + m_threads - 1) / m_threads));\n}\n\nvoid MonteCarloPiSimulation::setupSimulationAllocations()\n{\n hipDeviceptr_t d_ptr = 0U;\n size_t granularity = 0;\n hipMemGenericAllocationHandle_t cudaPositionHandle, cudaInCircleHandle;\n\n hipMemAllocationProp allocProp = {};\n allocProp.type = hipMemAllocationTypePinned;\n allocProp.location.type = hipMemLocationTypeDevice;\n allocProp.location.id = m_cudaDevice;\n allocProp.win32HandleMetaData = NULL;\n allocProp.requestedHandleTypes = ipcHandleTypeFlag;\n\n // Windows-specific LPSECURITYATTRIBUTES is required when\n // CU_MEM_HANDLE_TYPE_WIN32 is used. The security attribute defines the scope\n // of which exported allocations may be tranferred to other processes. For all\n // other handle types, pass NULL.\n getDefaultSecurityDescriptor(&allocProp);\n\n // Get the recommended granularity for m_cudaDevice.\n checkCudaErrors(hipMemGetAllocationGranularity(&granularity, &allocProp, hipMemAllocationGranularityRecommended));\n\n size_t xyPositionVecSize = m_numPoints * sizeof(*m_xyVector);\n size_t inCircleVecSize = m_numPoints * sizeof(*m_pointsInsideCircle);\n\n size_t xyPositionSize = ROUND_UP_TO_GRANULARITY(xyPositionVecSize, granularity);\n size_t inCircleSize = ROUND_UP_TO_GRANULARITY(inCircleVecSize, granularity);\n m_totalAllocationSize = (xyPositionSize + inCircleSize);\n\n // Reserve the required contiguous VA space for the allocations\n checkCudaErrors(hipMemAddressReserve(&d_ptr, m_totalAllocationSize, granularity, 0U, 0));\n\n // Create the allocations as a pinned allocation on this device.\n // Create an allocation to store all the positions of points on the xy plane\n // and a second allocation which stores information if the corresponding\n // position is inside the unit circle or not.\n checkCudaErrors(hipMemCreate(&cudaPositionHandle, xyPositionSize, &allocProp, 0));\n checkCudaErrors(hipMemCreate(&cudaInCircleHandle, inCircleSize, &allocProp, 0));\n\n // Export the allocation to a platform-specific handle. The type of handle\n // requested here must match the requestedHandleTypes field in the prop\n // structure passed to cuMemCreate. The handle obtained here will be passed to\n // vulkan to import the allocation.\n checkCudaErrors(\n hipMemExportToShareableHandle((void *)&m_posShareableHandle, cudaPositionHandle, ipcHandleTypeFlag, 0));\n checkCudaErrors(\n hipMemExportToShareableHandle((void *)&m_inCircleShareableHandle, cudaInCircleHandle, ipcHandleTypeFlag, 0));\n\n hipDeviceptr_t va_position = d_ptr;\n hipDeviceptr_t va_InCircle = va_position + xyPositionSize;\n m_pointsInsideCircle = (float *)va_InCircle;\n m_xyVector = (vec2 *)va_position;\n\n // Assign the chunk to the appropriate VA range\n checkCudaErrors(hipMemMap(va_position, xyPositionSize, 0, cudaPositionHandle, 0));\n checkCudaErrors(hipMemMap(va_InCircle, inCircleSize, 0, cudaInCircleHandle, 0));\n\n // Release the handles for the allocation. Since the allocation is currently\n // mapped to a VA range with a previous call to cuMemMap the actual freeing of\n // memory allocation will happen on an eventual call to cuMemUnmap. Thus the\n // allocation will be kept live until it is unmapped.\n checkCudaErrors(hipMemRelease(cudaPositionHandle));\n checkCudaErrors(hipMemRelease(cudaInCircleHandle));\n\n hipMemAccessDesc accessDescriptor = {};\n accessDescriptor.location.id = m_cudaDevice;\n accessDescriptor.location.type = hipMemLocationTypeDevice;\n accessDescriptor.flags = hipMemAccessFlagsProtReadWrite;\n\n // Apply the access descriptor to the whole VA range. Essentially enables\n // Read-Write access to the range.\n checkCudaErrors(hipMemSetAccess(d_ptr, m_totalAllocationSize, &accessDescriptor, 1));\n}\n\nvoid MonteCarloPiSimulation::cleanupSimulationAllocations()\n{\n if (m_xyVector && m_pointsInsideCircle) {\n // Unmap the mapped virtual memory region\n // Since the handles to the mapped backing stores have already been released\n // by cuMemRelease, and these are the only/last mappings referencing them,\n // The backing stores will be freed.\n checkCudaErrors(hipMemUnmap((hipDeviceptr_t)m_xyVector, m_totalAllocationSize));\n\n checkIpcErrors(ipcCloseShareableHandle(m_posShareableHandle));\n checkIpcErrors(ipcCloseShareableHandle(m_inCircleShareableHandle));\n\n // Free the virtual address region.\n checkCudaErrors(hipMemAddressFree((hipDeviceptr_t)m_xyVector, m_totalAllocationSize));\n\n m_xyVector = nullptr;\n m_pointsInsideCircle = nullptr;\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/simpleVulkanMMAP/MonteCarloPi.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include // CUDA device initialization helper functions\n#include \n#include \n\n__constant__ float cGaussian[64]; // gaussian array in device side\n\ncudaTextureObject_t rgbaTexdImage;\ncudaTextureObject_t rgbaTexdTemp;\n\nuint *dImage = NULL; // original image\nuint *dTemp = NULL; // temp array for iterations\nsize_t pitch;\n\n/*\n Perform a simple bilateral filter.\n\n Bilateral filter is a nonlinear filter that is a mixture of range\n filter and domain filter, the previous one preserves crisp edges and\n the latter one filters noise. The intensity value at each pixel in\n an image is replaced by a weighted average of intensity values from\n nearby pixels.\n\n The weight factor is calculated by the product of domain filter\n component(using the gaussian distribution as a spatial distance) as\n well as range filter component(Euclidean distance between center pixel\n and the current neighbor pixel). Because this process is nonlinear,\n the sample just uses a simple pixel by pixel step.\n\n Texture fetches automatically clamp to edge of image. 1D gaussian array\n is mapped to a 1D texture instead of using shared memory, which may\n cause severe bank conflict.\n\n Threads are y-pass(column-pass), because the output is coalesced.\n\n Parameters\n od - pointer to output data in global memory\n d_f - pointer to the 1D gaussian array\n e_d - euclidean delta\n w - image width\n h - image height\n r - filter radius\n*/\n\n// Euclidean Distance (x, y, d) = exp((|x - y| / d)^2 / 2)\n__device__ float euclideanLen(float4 a, float4 b, float d)\n{\n float mod = (b.x - a.x) * (b.x - a.x) + (b.y - a.y) * (b.y - a.y) + (b.z - a.z) * (b.z - a.z);\n\n return __expf(-mod / (2.f * d * d));\n}\n\n__device__ uint rgbaFloatToInt(float4 rgba)\n{\n rgba.x = __saturatef(fabs(rgba.x)); // clamp to [0.0, 1.0]\n rgba.y = __saturatef(fabs(rgba.y));\n rgba.z = __saturatef(fabs(rgba.z));\n rgba.w = __saturatef(fabs(rgba.w));\n return (uint(rgba.w * 255.0f) << 24) | (uint(rgba.z * 255.0f) << 16) | (uint(rgba.y * 255.0f) << 8)\n | uint(rgba.x * 255.0f);\n}\n\n__device__ float4 rgbaIntToFloat(uint c)\n{\n float4 rgba;\n rgba.x = (c & 0xff) * 0.003921568627f; // /255.0f;\n rgba.y = ((c >> 8) & 0xff) * 0.003921568627f; // /255.0f;\n rgba.z = ((c >> 16) & 0xff) * 0.003921568627f; // /255.0f;\n rgba.w = ((c >> 24) & 0xff) * 0.003921568627f; // /255.0f;\n return rgba;\n}\n\n// column pass using coalesced global memory reads\n__global__ void d_bilateral_filter(uint *od, int w, int h, float e_d, int r, cudaTextureObject_t rgbaTex)\n{\n int x = blockIdx.x * blockDim.x + threadIdx.x;\n int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n if (x >= w || y >= h) {\n return;\n }\n\n float sum = 0.0f;\n float factor;\n float4 t = {0.f, 0.f, 0.f, 0.f};\n float4 center = tex2D(rgbaTex, x, y);\n\n for (int i = -r; i <= r; i++) {\n for (int j = -r; j <= r; j++) {\n float4 curPix = tex2D(rgbaTex, x + j, y + i);\n factor = cGaussian[i + r] * cGaussian[j + r] * // domain factor\n euclideanLen(curPix, center, e_d); // range factor\n\n t += factor * curPix;\n sum += factor;\n }\n }\n\n od[y * w + x] = rgbaFloatToInt(t / sum);\n}\n\nextern \"C\" void initTexture(int width, int height, uint *hImage)\n{\n // copy image data to array\n checkCudaErrors(cudaMallocPitch(&dImage, &pitch, sizeof(uint) * width, height));\n checkCudaErrors(cudaMallocPitch(&dTemp, &pitch, sizeof(uint) * width, height));\n checkCudaErrors(cudaMemcpy2D(\n dImage, pitch, hImage, sizeof(uint) * width, sizeof(uint) * width, height, cudaMemcpyHostToDevice));\n\n // texture rgbaTex;\n cudaChannelFormatDesc desc = cudaCreateChannelDesc();\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypePitch2D;\n texRes.res.pitch2D.devPtr = dImage;\n texRes.res.pitch2D.desc = desc;\n texRes.res.pitch2D.width = width;\n texRes.res.pitch2D.height = height;\n texRes.res.pitch2D.pitchInBytes = pitch;\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = cudaFilterModePoint;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeNormalizedFloat;\n\n checkCudaErrors(cudaCreateTextureObject(&rgbaTexdImage, &texRes, &texDescr, NULL));\n\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypePitch2D;\n texRes.res.pitch2D.devPtr = dTemp;\n texRes.res.pitch2D.desc = desc;\n texRes.res.pitch2D.width = width;\n texRes.res.pitch2D.height = height;\n texRes.res.pitch2D.pitchInBytes = pitch;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = cudaFilterModeLinear;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeNormalizedFloat;\n\n checkCudaErrors(cudaCreateTextureObject(&rgbaTexdTemp, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void freeTextures()\n{\n checkCudaErrors(cudaDestroyTextureObject(rgbaTexdImage));\n checkCudaErrors(cudaDestroyTextureObject(rgbaTexdTemp));\n checkCudaErrors(cudaFree(dImage));\n checkCudaErrors(cudaFree(dTemp));\n}\n\n/*\n Because a 2D gaussian mask is symmetry in row and column,\n here only generate a 1D mask, and use the product by row\n and column index later.\n\n 1D gaussian distribution :\n g(x, d) -- C * exp(-x^2/d^2), C is a constant amplifier\n\n parameters:\n og - output gaussian array in global memory\n delta - the 2nd parameter 'd' in the above function\n radius - half of the filter size\n (total filter size = 2 * radius + 1)\n*/\nextern \"C\" void updateGaussian(float delta, int radius)\n{\n float fGaussian[64];\n\n for (int i = 0; i < 2 * radius + 1; ++i) {\n float x = (float)(i - radius);\n fGaussian[i] = expf(-(x * x) / (2 * delta * delta));\n }\n\n checkCudaErrors(cudaMemcpyToSymbol(cGaussian, fGaussian, sizeof(float) * (2 * radius + 1)));\n}\n\n/*\n Perform 2D bilateral filter on image using CUDA\n\n Parameters:\n d_dest - pointer to destination image in device memory\n width - image width\n height - image height\n e_d - euclidean delta\n radius - filter radius\n iterations - number of iterations\n*/\n\n// RGBA version\nextern \"C\" double bilateralFilterRGBA(uint *dDest,\n int width,\n int height,\n float e_d,\n int radius,\n int iterations,\n StopWatchInterface *timer)\n{\n // var for kernel computation timing\n double dKernelTime;\n\n for (int i = 0; i < iterations; i++) {\n // sync host and start kernel computation timer\n dKernelTime = 0.0;\n checkCudaErrors(cudaDeviceSynchronize());\n sdkResetTimer(&timer);\n\n dim3 gridSize((width + 16 - 1) / 16, (height + 16 - 1) / 16);\n dim3 blockSize(16, 16);\n\n if (iterations > 1) {\n d_bilateral_filter<<>>(dDest, width, height, e_d, radius, rgbaTexdTemp);\n }\n else {\n d_bilateral_filter<<>>(dDest, width, height, e_d, radius, rgbaTexdImage);\n }\n\n // sync host and stop computation timer\n checkCudaErrors(cudaDeviceSynchronize());\n dKernelTime += sdkGetTimerValue(&timer);\n\n if (iterations > 1) {\n // copy result back from global memory to array\n checkCudaErrors(cudaMemcpy2D(\n dTemp, pitch, dDest, sizeof(int) * width, sizeof(int) * width, height, cudaMemcpyDeviceToDevice));\n }\n }\n\n return ((dKernelTime / 1000.) / (double)iterations);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include // CUDA device initialization helper functions\n#include \n#include \n\n__constant__ float cGaussian[64]; // gaussian array in device side\n\nhipTextureObject_t rgbaTexdImage;\nhipTextureObject_t rgbaTexdTemp;\n\nuint *dImage = NULL; // original image\nuint *dTemp = NULL; // temp array for iterations\nsize_t pitch;\n\n/*\n Perform a simple bilateral filter.\n\n Bilateral filter is a nonlinear filter that is a mixture of range\n filter and domain filter, the previous one preserves crisp edges and\n the latter one filters noise. The intensity value at each pixel in\n an image is replaced by a weighted average of intensity values from\n nearby pixels.\n\n The weight factor is calculated by the product of domain filter\n component(using the gaussian distribution as a spatial distance) as\n well as range filter component(Euclidean distance between center pixel\n and the current neighbor pixel). Because this process is nonlinear,\n the sample just uses a simple pixel by pixel step.\n\n Texture fetches automatically clamp to edge of image. 1D gaussian array\n is mapped to a 1D texture instead of using shared memory, which may\n cause severe bank conflict.\n\n Threads are y-pass(column-pass), because the output is coalesced.\n\n Parameters\n od - pointer to output data in global memory\n d_f - pointer to the 1D gaussian array\n e_d - euclidean delta\n w - image width\n h - image height\n r - filter radius\n*/\n\n// Euclidean Distance (x, y, d) = exp((|x - y| / d)^2 / 2)\n__device__ float euclideanLen(float4 a, float4 b, float d)\n{\n float mod = (b.x - a.x) * (b.x - a.x) + (b.y - a.y) * (b.y - a.y) + (b.z - a.z) * (b.z - a.z);\n\n return __expf(-mod / (2.f * d * d));\n}\n\n__device__ uint rgbaFloatToInt(float4 rgba)\n{\n rgba.x = __saturatef(fabs(rgba.x)); // clamp to [0.0, 1.0]\n rgba.y = __saturatef(fabs(rgba.y));\n rgba.z = __saturatef(fabs(rgba.z));\n rgba.w = __saturatef(fabs(rgba.w));\n return (uint(rgba.w * 255.0f) << 24) | (uint(rgba.z * 255.0f) << 16) | (uint(rgba.y * 255.0f) << 8)\n | uint(rgba.x * 255.0f);\n}\n\n__device__ float4 rgbaIntToFloat(uint c)\n{\n float4 rgba;\n rgba.x = (c & 0xff) * 0.003921568627f; // /255.0f;\n rgba.y = ((c >> 8) & 0xff) * 0.003921568627f; // /255.0f;\n rgba.z = ((c >> 16) & 0xff) * 0.003921568627f; // /255.0f;\n rgba.w = ((c >> 24) & 0xff) * 0.003921568627f; // /255.0f;\n return rgba;\n}\n\n// column pass using coalesced global memory reads\n__global__ void d_bilateral_filter(uint *od, int w, int h, float e_d, int r, hipTextureObject_t rgbaTex)\n{\n int x = blockIdx.x * blockDim.x + threadIdx.x;\n int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n if (x >= w || y >= h) {\n return;\n }\n\n float sum = 0.0f;\n float factor;\n float4 t = {0.f, 0.f, 0.f, 0.f};\n float4 center = tex2D(rgbaTex, x, y);\n\n for (int i = -r; i <= r; i++) {\n for (int j = -r; j <= r; j++) {\n float4 curPix = tex2D(rgbaTex, x + j, y + i);\n factor = cGaussian[i + r] * cGaussian[j + r] * // domain factor\n euclideanLen(curPix, center, e_d); // range factor\n\n t += factor * curPix;\n sum += factor;\n }\n }\n\n od[y * w + x] = rgbaFloatToInt(t / sum);\n}\n\nextern \"C\" void initTexture(int width, int height, uint *hImage)\n{\n // copy image data to array\n checkCudaErrors(hipMallocPitch(&dImage, &pitch, sizeof(uint) * width, height));\n checkCudaErrors(hipMallocPitch(&dTemp, &pitch, sizeof(uint) * width, height));\n checkCudaErrors(hipMemcpy2D(\n dImage, pitch, hImage, sizeof(uint) * width, sizeof(uint) * width, height, hipMemcpyHostToDevice));\n\n // texture rgbaTex;\n hipChannelFormatDesc desc = hipCreateChannelDesc();\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypePitch2D;\n texRes.res.pitch2D.devPtr = dImage;\n texRes.res.pitch2D.desc = desc;\n texRes.res.pitch2D.width = width;\n texRes.res.pitch2D.height = height;\n texRes.res.pitch2D.pitchInBytes = pitch;\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = hipFilterModePoint;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeNormalizedFloat;\n\n checkCudaErrors(hipCreateTextureObject(&rgbaTexdImage, &texRes, &texDescr, NULL));\n\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypePitch2D;\n texRes.res.pitch2D.devPtr = dTemp;\n texRes.res.pitch2D.desc = desc;\n texRes.res.pitch2D.width = width;\n texRes.res.pitch2D.height = height;\n texRes.res.pitch2D.pitchInBytes = pitch;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = hipFilterModeLinear;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeNormalizedFloat;\n\n checkCudaErrors(hipCreateTextureObject(&rgbaTexdTemp, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void freeTextures()\n{\n checkCudaErrors(hipDestroyTextureObject(rgbaTexdImage));\n checkCudaErrors(hipDestroyTextureObject(rgbaTexdTemp));\n checkCudaErrors(hipFree(dImage));\n checkCudaErrors(hipFree(dTemp));\n}\n\n/*\n Because a 2D gaussian mask is symmetry in row and column,\n here only generate a 1D mask, and use the product by row\n and column index later.\n\n 1D gaussian distribution :\n g(x, d) -- C * exp(-x^2/d^2), C is a constant amplifier\n\n parameters:\n og - output gaussian array in global memory\n delta - the 2nd parameter 'd' in the above function\n radius - half of the filter size\n (total filter size = 2 * radius + 1)\n*/\nextern \"C\" void updateGaussian(float delta, int radius)\n{\n float fGaussian[64];\n\n for (int i = 0; i < 2 * radius + 1; ++i) {\n float x = (float)(i - radius);\n fGaussian[i] = expf(-(x * x) / (2 * delta * delta));\n }\n\n checkCudaErrors(hipMemcpyToSymbol(HIP_SYMBOL(cGaussian), fGaussian, sizeof(float) * (2 * radius + 1)));\n}\n\n/*\n Perform 2D bilateral filter on image using CUDA\n\n Parameters:\n d_dest - pointer to destination image in device memory\n width - image width\n height - image height\n e_d - euclidean delta\n radius - filter radius\n iterations - number of iterations\n*/\n\n// RGBA version\nextern \"C\" double bilateralFilterRGBA(uint *dDest,\n int width,\n int height,\n float e_d,\n int radius,\n int iterations,\n StopWatchInterface *timer)\n{\n // var for kernel computation timing\n double dKernelTime;\n\n for (int i = 0; i < iterations; i++) {\n // sync host and start kernel computation timer\n dKernelTime = 0.0;\n checkCudaErrors(hipDeviceSynchronize());\n sdkResetTimer(&timer);\n\n dim3 gridSize((width + 16 - 1) / 16, (height + 16 - 1) / 16);\n dim3 blockSize(16, 16);\n\n if (iterations > 1) {\n d_bilateral_filter<<>>(dDest, width, height, e_d, radius, rgbaTexdTemp);\n }\n else {\n d_bilateral_filter<<>>(dDest, width, height, e_d, radius, rgbaTexdImage);\n }\n\n // sync host and stop computation timer\n checkCudaErrors(hipDeviceSynchronize());\n dKernelTime += sdkGetTimerValue(&timer);\n\n if (iterations > 1) {\n // copy result back from global memory to array\n checkCudaErrors(hipMemcpy2D(\n dTemp, pitch, dDest, sizeof(int) * width, sizeof(int) * width, height, hipMemcpyDeviceToDevice));\n }\n }\n\n return ((dKernelTime / 1000.) / (double)iterations);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/bilateralFilter/bilateral_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n\n#include \"Mandelbrot_kernel.cuh\"\n#include \"Mandelbrot_kernel.h\"\n#include \"helper_cuda.h\"\n\n// The Mandelbrot CUDA GPU thread function\n\ntemplate \n__global__ void Mandelbrot0(uchar4 *dst,\n const int imageW,\n const int imageH,\n const int crunch,\n const T xOff,\n const T yOff,\n const T xJP,\n const T yJP,\n const T scale,\n const uchar4 colors,\n const int frame,\n const int animationFrame,\n const int gridWidth,\n const int numBlocks,\n const bool isJ)\n{\n // loop until all blocks completed\n for (unsigned int blockIndex = blockIdx.x; blockIndex < numBlocks; blockIndex += gridDim.x) {\n unsigned int blockX = blockIndex % gridWidth;\n unsigned int blockY = blockIndex / gridWidth;\n\n // process this block\n const int ix = blockDim.x * blockX + threadIdx.x;\n const int iy = blockDim.y * blockY + threadIdx.y;\n\n if ((ix < imageW) && (iy < imageH)) {\n // Calculate the location\n const T xPos = (T)ix * scale + xOff;\n const T yPos = (T)iy * scale + yOff;\n\n // Calculate the Mandelbrot index for the current location\n int m = CalcMandelbrot(xPos, yPos, xJP, yJP, crunch, isJ);\n // int m = blockIdx.x; // uncomment to see scheduling\n // order\n m = m > 0 ? crunch - m : 0;\n\n // Convert the Mandelbrot index into a color\n uchar4 color;\n\n if (m) {\n m += animationFrame;\n color.x = m * colors.x;\n color.y = m * colors.y;\n color.z = m * colors.z;\n }\n else {\n color.x = 0;\n color.y = 0;\n color.z = 0;\n }\n\n // Output the pixel\n int pixel = imageW * iy + ix;\n\n if (frame == 0) {\n color.w = 0;\n dst[pixel] = color;\n }\n else {\n int frame1 = frame + 1;\n int frame2 = frame1 / 2;\n dst[pixel].x = (dst[pixel].x * frame + color.x + frame2) / frame1;\n dst[pixel].y = (dst[pixel].y * frame + color.y + frame2) / frame1;\n dst[pixel].z = (dst[pixel].z * frame + color.z + frame2) / frame1;\n }\n }\n }\n\n} // Mandelbrot0\n\n// The Mandelbrot CUDA GPU thread function (double single version)\n__global__ void MandelbrotDS0(uchar4 *dst,\n const int imageW,\n const int imageH,\n const int crunch,\n const float xOff0,\n const float xOff1,\n const float yOff0,\n const float yOff1,\n const float xJP,\n const float yJP,\n const float scale,\n const uchar4 colors,\n const int frame,\n const int animationFrame,\n const int gridWidth,\n const int numBlocks,\n const bool isJ)\n{\n // loop until all blocks completed\n for (unsigned int blockIndex = blockIdx.x; blockIndex < numBlocks; blockIndex += gridDim.x) {\n unsigned int blockX = blockIndex % gridWidth;\n unsigned int blockY = blockIndex / gridWidth;\n\n // process this block\n const int ix = blockDim.x * blockX + threadIdx.x;\n const int iy = blockDim.y * blockY + threadIdx.y;\n\n if ((ix < imageW) && (iy < imageH)) {\n // Calculate the location\n float xPos0 = (float)ix * scale;\n float xPos1 = 0.0f;\n float yPos0 = (float)iy * scale;\n float yPos1 = 0.0f;\n dsadd(xPos0, xPos1, xPos0, xPos1, xOff0, xOff1);\n dsadd(yPos0, yPos1, yPos0, yPos1, yOff0, yOff1);\n\n // Calculate the Mandelbrot index for the current location\n int m = CalcMandelbrotDS(xPos0, xPos1, yPos0, yPos1, xJP, yJP, crunch, isJ);\n m = m > 0 ? crunch - m : 0;\n\n // Convert the Mandelbrot index into a color\n uchar4 color;\n\n if (m) {\n m += animationFrame;\n color.x = m * colors.x;\n color.y = m * colors.y;\n color.z = m * colors.z;\n }\n else {\n color.x = 0;\n color.y = 0;\n color.z = 0;\n }\n\n // Output the pixel\n int pixel = imageW * iy + ix;\n\n if (frame == 0) {\n color.w = 0;\n dst[pixel] = color;\n }\n else {\n int frame1 = frame + 1;\n int frame2 = frame1 / 2;\n dst[pixel].x = (dst[pixel].x * frame + color.x + frame2) / frame1;\n dst[pixel].y = (dst[pixel].y * frame + color.y + frame2) / frame1;\n dst[pixel].z = (dst[pixel].z * frame + color.z + frame2) / frame1;\n }\n }\n }\n} // MandelbrotDS0\n\n// The Mandelbrot secondary AA pass CUDA GPU thread function\ntemplate \n__global__ void Mandelbrot1(uchar4 *dst,\n const int imageW,\n const int imageH,\n const int crunch,\n const T xOff,\n const T yOff,\n const T xJP,\n const T yJP,\n const T scale,\n const uchar4 colors,\n const int frame,\n const int animationFrame,\n const int gridWidth,\n const int numBlocks,\n const bool isJ)\n{\n // loop until all blocks completed\n for (unsigned int blockIndex = blockIdx.x; blockIndex < numBlocks; blockIndex += gridDim.x) {\n unsigned int blockX = blockIndex % gridWidth;\n unsigned int blockY = blockIndex / gridWidth;\n\n // process this block\n const int ix = blockDim.x * blockX + threadIdx.x;\n const int iy = blockDim.y * blockY + threadIdx.y;\n\n if ((ix < imageW) && (iy < imageH)) {\n // Get the current pixel color\n int pixel = imageW * iy + ix;\n uchar4 pixelColor = dst[pixel];\n int count = 0;\n\n // Search for pixels out of tolerance surrounding the current pixel\n if (ix > 0) {\n count += CheckColors(pixelColor, dst[pixel - 1]);\n }\n\n if (ix + 1 < imageW) {\n count += CheckColors(pixelColor, dst[pixel + 1]);\n }\n\n if (iy > 0) {\n count += CheckColors(pixelColor, dst[pixel - imageW]);\n }\n\n if (iy + 1 < imageH) {\n count += CheckColors(pixelColor, dst[pixel + imageW]);\n }\n\n if (count) {\n // Calculate the location\n const T xPos = (T)ix * scale + xOff;\n const T yPos = (T)iy * scale + yOff;\n\n // Calculate the Mandelbrot index for the current location\n int m = CalcMandelbrot(xPos, yPos, xJP, yJP, crunch, isJ);\n m = m > 0 ? crunch - m : 0;\n\n // Convert the Mandelbrot index into a color\n uchar4 color;\n\n if (m) {\n m += animationFrame;\n color.x = m * colors.x;\n color.y = m * colors.y;\n color.z = m * colors.z;\n }\n else {\n color.x = 0;\n color.y = 0;\n color.z = 0;\n }\n\n // Output the pixel\n int frame1 = frame + 1;\n int frame2 = frame1 / 2;\n dst[pixel].x = (pixelColor.x * frame + color.x + frame2) / frame1;\n dst[pixel].y = (pixelColor.y * frame + color.y + frame2) / frame1;\n dst[pixel].z = (pixelColor.z * frame + color.z + frame2) / frame1;\n }\n }\n }\n\n} // Mandelbrot1\n\n// The Mandelbrot secondary AA pass CUDA GPU thread function (double single\n// version)\n__global__ void MandelbrotDS1(uchar4 *dst,\n const int imageW,\n const int imageH,\n const int crunch,\n const float xOff0,\n const float xOff1,\n const float yOff0,\n const float yOff1,\n const float xJP,\n const float yJP,\n const float scale,\n const uchar4 colors,\n const int frame,\n const int animationFrame,\n const int gridWidth,\n const int numBlocks,\n const bool isJ)\n{\n // loop until all blocks completed\n for (unsigned int blockIndex = blockIdx.x; blockIndex < numBlocks; blockIndex += gridDim.x) {\n unsigned int blockX = blockIndex % gridWidth;\n unsigned int blockY = blockIndex / gridWidth;\n\n // process this block\n const int ix = blockDim.x * blockX + threadIdx.x;\n const int iy = blockDim.y * blockY + threadIdx.y;\n\n if ((ix < imageW) && (iy < imageH)) {\n // Get the current pixel color\n int pixel = imageW * iy + ix;\n uchar4 pixelColor = dst[pixel];\n int count = 0;\n\n // Search for pixels out of tolerance surrounding the current pixel\n if (ix > 0) {\n count += CheckColors(pixelColor, dst[pixel - 1]);\n }\n\n if (ix + 1 < imageW) {\n count += CheckColors(pixelColor, dst[pixel + 1]);\n }\n\n if (iy > 0) {\n count += CheckColors(pixelColor, dst[pixel - imageW]);\n }\n\n if (iy + 1 < imageH) {\n count += CheckColors(pixelColor, dst[pixel + imageW]);\n }\n\n if (count) {\n // Calculate the location\n float xPos0 = (float)ix * scale;\n float xPos1 = 0.0f;\n float yPos0 = (float)iy * scale;\n float yPos1 = 0.0f;\n dsadd(xPos0, xPos1, xPos0, xPos1, xOff0, xOff1);\n dsadd(yPos0, yPos1, yPos0, yPos1, yOff0, yOff1);\n\n // Calculate the Mandelbrot index for the current location\n int m = CalcMandelbrotDS(xPos0, xPos1, yPos0, yPos1, xJP, yJP, crunch, isJ);\n m = m > 0 ? crunch - m : 0;\n\n // Convert the Mandelbrot index into a color\n uchar4 color;\n\n if (m) {\n m += animationFrame;\n color.x = m * colors.x;\n color.y = m * colors.y;\n color.z = m * colors.z;\n }\n else {\n color.x = 0;\n color.y = 0;\n color.z = 0;\n }\n\n // Output the pixel\n int frame1 = frame + 1;\n int frame2 = frame1 / 2;\n dst[pixel].x = (pixelColor.x * frame + color.x + frame2) / frame1;\n dst[pixel].y = (pixelColor.y * frame + color.y + frame2) / frame1;\n dst[pixel].z = (pixelColor.z * frame + color.z + frame2) / frame1;\n }\n }\n }\n\n} // MandelbrotDS1\n\n// The host CPU Mandelbrot thread spawner\nvoid RunMandelbrot0(uchar4 *dst,\n const int imageW,\n const int imageH,\n const int crunch,\n const double xOff,\n const double yOff,\n const double xjp,\n const double yjp,\n const double scale,\n const uchar4 colors,\n const int frame,\n const int animationFrame,\n const int mode,\n const int numSMs,\n const bool isJ,\n int version)\n{\n dim3 threads(BLOCKDIM_X, BLOCKDIM_Y);\n dim3 grid(iDivUp(imageW, BLOCKDIM_X), iDivUp(imageH, BLOCKDIM_Y));\n\n int numWorkerBlocks = numSMs;\n\n switch (mode) {\n default:\n case 0:\n Mandelbrot0<<>>(dst,\n imageW,\n imageH,\n crunch,\n (float)xOff,\n (float)yOff,\n (float)xjp,\n (float)yjp,\n (float)scale,\n colors,\n frame,\n animationFrame,\n grid.x,\n grid.x * grid.y,\n isJ);\n break;\n case 1:\n float x0, x1, y0, y1;\n dsdeq(x0, x1, xOff);\n dsdeq(y0, y1, yOff);\n MandelbrotDS0<<>>(dst,\n imageW,\n imageH,\n crunch,\n x0,\n x1,\n y0,\n y1,\n (float)xjp,\n (float)yjp,\n (float)scale,\n colors,\n frame,\n animationFrame,\n grid.x,\n grid.x * grid.y,\n isJ);\n break;\n case 2:\n Mandelbrot0<<>>(dst,\n imageW,\n imageH,\n crunch,\n xOff,\n yOff,\n xjp,\n yjp,\n scale,\n colors,\n frame,\n animationFrame,\n grid.x,\n grid.x * grid.y,\n isJ);\n break;\n }\n\n getLastCudaError(\"Mandelbrot0 kernel execution failed.\\n\");\n} // RunMandelbrot0\n\n// The host CPU Mandelbrot thread spawner\nvoid RunMandelbrot1(uchar4 *dst,\n const int imageW,\n const int imageH,\n const int crunch,\n const double xOff,\n const double yOff,\n const double xjp,\n const double yjp,\n const double scale,\n const uchar4 colors,\n const int frame,\n const int animationFrame,\n const int mode,\n const int numSMs,\n const bool isJ,\n int version)\n{\n dim3 threads(BLOCKDIM_X, BLOCKDIM_Y);\n dim3 grid(iDivUp(imageW, BLOCKDIM_X), iDivUp(imageH, BLOCKDIM_Y));\n\n int numWorkerBlocks = numSMs;\n\n switch (mode) {\n default:\n case 0:\n Mandelbrot1<<>>(dst,\n imageW,\n imageH,\n crunch,\n (float)xOff,\n (float)yOff,\n (float)xjp,\n (float)yjp,\n (float)scale,\n colors,\n frame,\n animationFrame,\n grid.x,\n grid.x * grid.y,\n isJ);\n break;\n case 1:\n float x0, x1, y0, y1;\n dsdeq(x0, x1, xOff);\n dsdeq(y0, y1, yOff);\n MandelbrotDS1<<>>(dst,\n imageW,\n imageH,\n crunch,\n x0,\n x1,\n y0,\n y1,\n (float)xjp,\n (float)yjp,\n (float)scale,\n colors,\n frame,\n animationFrame,\n grid.x,\n grid.x * grid.y,\n isJ);\n break;\n case 2:\n Mandelbrot1<<>>(dst,\n imageW,\n imageH,\n crunch,\n xOff,\n yOff,\n xjp,\n yjp,\n scale,\n colors,\n frame,\n animationFrame,\n grid.x,\n grid.x * grid.y,\n isJ);\n break;\n }\n\n getLastCudaError(\"Mandelbrot1 kernel execution failed.\\n\");\n} // RunMandelbrot1\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \"Mandelbrot_kernel.cuh\"\n#include \"Mandelbrot_kernel.h\"\n#include \"hip/hip_runtime_api.h\"\n\n// The Mandelbrot CUDA GPU thread function\n\ntemplate \n__global__ void Mandelbrot0(uchar4 *dst,\n const int imageW,\n const int imageH,\n const int crunch,\n const T xOff,\n const T yOff,\n const T xJP,\n const T yJP,\n const T scale,\n const uchar4 colors,\n const int frame,\n const int animationFrame,\n const int gridWidth,\n const int numBlocks,\n const bool isJ)\n{\n // loop until all blocks completed\n for (unsigned int blockIndex = blockIdx.x; blockIndex < numBlocks; blockIndex += gridDim.x) {\n unsigned int blockX = blockIndex % gridWidth;\n unsigned int blockY = blockIndex / gridWidth;\n\n // process this block\n const int ix = blockDim.x * blockX + threadIdx.x;\n const int iy = blockDim.y * blockY + threadIdx.y;\n\n if ((ix < imageW) && (iy < imageH)) {\n // Calculate the location\n const T xPos = (T)ix * scale + xOff;\n const T yPos = (T)iy * scale + yOff;\n\n // Calculate the Mandelbrot index for the current location\n int m = CalcMandelbrot(xPos, yPos, xJP, yJP, crunch, isJ);\n // int m = blockIdx.x; // uncomment to see scheduling\n // order\n m = m > 0 ? crunch - m : 0;\n\n // Convert the Mandelbrot index into a color\n uchar4 color;\n\n if (m) {\n m += animationFrame;\n color.x = m * colors.x;\n color.y = m * colors.y;\n color.z = m * colors.z;\n }\n else {\n color.x = 0;\n color.y = 0;\n color.z = 0;\n }\n\n // Output the pixel\n int pixel = imageW * iy + ix;\n\n if (frame == 0) {\n color.w = 0;\n dst[pixel] = color;\n }\n else {\n int frame1 = frame + 1;\n int frame2 = frame1 / 2;\n dst[pixel].x = (dst[pixel].x * frame + color.x + frame2) / frame1;\n dst[pixel].y = (dst[pixel].y * frame + color.y + frame2) / frame1;\n dst[pixel].z = (dst[pixel].z * frame + color.z + frame2) / frame1;\n }\n }\n }\n\n} // Mandelbrot0\n\n// The Mandelbrot CUDA GPU thread function (double single version)\n__global__ void MandelbrotDS0(uchar4 *dst,\n const int imageW,\n const int imageH,\n const int crunch,\n const float xOff0,\n const float xOff1,\n const float yOff0,\n const float yOff1,\n const float xJP,\n const float yJP,\n const float scale,\n const uchar4 colors,\n const int frame,\n const int animationFrame,\n const int gridWidth,\n const int numBlocks,\n const bool isJ)\n{\n // loop until all blocks completed\n for (unsigned int blockIndex = blockIdx.x; blockIndex < numBlocks; blockIndex += gridDim.x) {\n unsigned int blockX = blockIndex % gridWidth;\n unsigned int blockY = blockIndex / gridWidth;\n\n // process this block\n const int ix = blockDim.x * blockX + threadIdx.x;\n const int iy = blockDim.y * blockY + threadIdx.y;\n\n if ((ix < imageW) && (iy < imageH)) {\n // Calculate the location\n float xPos0 = (float)ix * scale;\n float xPos1 = 0.0f;\n float yPos0 = (float)iy * scale;\n float yPos1 = 0.0f;\n dsadd(xPos0, xPos1, xPos0, xPos1, xOff0, xOff1);\n dsadd(yPos0, yPos1, yPos0, yPos1, yOff0, yOff1);\n\n // Calculate the Mandelbrot index for the current location\n int m = CalcMandelbrotDS(xPos0, xPos1, yPos0, yPos1, xJP, yJP, crunch, isJ);\n m = m > 0 ? crunch - m : 0;\n\n // Convert the Mandelbrot index into a color\n uchar4 color;\n\n if (m) {\n m += animationFrame;\n color.x = m * colors.x;\n color.y = m * colors.y;\n color.z = m * colors.z;\n }\n else {\n color.x = 0;\n color.y = 0;\n color.z = 0;\n }\n\n // Output the pixel\n int pixel = imageW * iy + ix;\n\n if (frame == 0) {\n color.w = 0;\n dst[pixel] = color;\n }\n else {\n int frame1 = frame + 1;\n int frame2 = frame1 / 2;\n dst[pixel].x = (dst[pixel].x * frame + color.x + frame2) / frame1;\n dst[pixel].y = (dst[pixel].y * frame + color.y + frame2) / frame1;\n dst[pixel].z = (dst[pixel].z * frame + color.z + frame2) / frame1;\n }\n }\n }\n} // MandelbrotDS0\n\n// The Mandelbrot secondary AA pass CUDA GPU thread function\ntemplate \n__global__ void Mandelbrot1(uchar4 *dst,\n const int imageW,\n const int imageH,\n const int crunch,\n const T xOff,\n const T yOff,\n const T xJP,\n const T yJP,\n const T scale,\n const uchar4 colors,\n const int frame,\n const int animationFrame,\n const int gridWidth,\n const int numBlocks,\n const bool isJ)\n{\n // loop until all blocks completed\n for (unsigned int blockIndex = blockIdx.x; blockIndex < numBlocks; blockIndex += gridDim.x) {\n unsigned int blockX = blockIndex % gridWidth;\n unsigned int blockY = blockIndex / gridWidth;\n\n // process this block\n const int ix = blockDim.x * blockX + threadIdx.x;\n const int iy = blockDim.y * blockY + threadIdx.y;\n\n if ((ix < imageW) && (iy < imageH)) {\n // Get the current pixel color\n int pixel = imageW * iy + ix;\n uchar4 pixelColor = dst[pixel];\n int count = 0;\n\n // Search for pixels out of tolerance surrounding the current pixel\n if (ix > 0) {\n count += CheckColors(pixelColor, dst[pixel - 1]);\n }\n\n if (ix + 1 < imageW) {\n count += CheckColors(pixelColor, dst[pixel + 1]);\n }\n\n if (iy > 0) {\n count += CheckColors(pixelColor, dst[pixel - imageW]);\n }\n\n if (iy + 1 < imageH) {\n count += CheckColors(pixelColor, dst[pixel + imageW]);\n }\n\n if (count) {\n // Calculate the location\n const T xPos = (T)ix * scale + xOff;\n const T yPos = (T)iy * scale + yOff;\n\n // Calculate the Mandelbrot index for the current location\n int m = CalcMandelbrot(xPos, yPos, xJP, yJP, crunch, isJ);\n m = m > 0 ? crunch - m : 0;\n\n // Convert the Mandelbrot index into a color\n uchar4 color;\n\n if (m) {\n m += animationFrame;\n color.x = m * colors.x;\n color.y = m * colors.y;\n color.z = m * colors.z;\n }\n else {\n color.x = 0;\n color.y = 0;\n color.z = 0;\n }\n\n // Output the pixel\n int frame1 = frame + 1;\n int frame2 = frame1 / 2;\n dst[pixel].x = (pixelColor.x * frame + color.x + frame2) / frame1;\n dst[pixel].y = (pixelColor.y * frame + color.y + frame2) / frame1;\n dst[pixel].z = (pixelColor.z * frame + color.z + frame2) / frame1;\n }\n }\n }\n\n} // Mandelbrot1\n\n// The Mandelbrot secondary AA pass CUDA GPU thread function (double single\n// version)\n__global__ void MandelbrotDS1(uchar4 *dst,\n const int imageW,\n const int imageH,\n const int crunch,\n const float xOff0,\n const float xOff1,\n const float yOff0,\n const float yOff1,\n const float xJP,\n const float yJP,\n const float scale,\n const uchar4 colors,\n const int frame,\n const int animationFrame,\n const int gridWidth,\n const int numBlocks,\n const bool isJ)\n{\n // loop until all blocks completed\n for (unsigned int blockIndex = blockIdx.x; blockIndex < numBlocks; blockIndex += gridDim.x) {\n unsigned int blockX = blockIndex % gridWidth;\n unsigned int blockY = blockIndex / gridWidth;\n\n // process this block\n const int ix = blockDim.x * blockX + threadIdx.x;\n const int iy = blockDim.y * blockY + threadIdx.y;\n\n if ((ix < imageW) && (iy < imageH)) {\n // Get the current pixel color\n int pixel = imageW * iy + ix;\n uchar4 pixelColor = dst[pixel];\n int count = 0;\n\n // Search for pixels out of tolerance surrounding the current pixel\n if (ix > 0) {\n count += CheckColors(pixelColor, dst[pixel - 1]);\n }\n\n if (ix + 1 < imageW) {\n count += CheckColors(pixelColor, dst[pixel + 1]);\n }\n\n if (iy > 0) {\n count += CheckColors(pixelColor, dst[pixel - imageW]);\n }\n\n if (iy + 1 < imageH) {\n count += CheckColors(pixelColor, dst[pixel + imageW]);\n }\n\n if (count) {\n // Calculate the location\n float xPos0 = (float)ix * scale;\n float xPos1 = 0.0f;\n float yPos0 = (float)iy * scale;\n float yPos1 = 0.0f;\n dsadd(xPos0, xPos1, xPos0, xPos1, xOff0, xOff1);\n dsadd(yPos0, yPos1, yPos0, yPos1, yOff0, yOff1);\n\n // Calculate the Mandelbrot index for the current location\n int m = CalcMandelbrotDS(xPos0, xPos1, yPos0, yPos1, xJP, yJP, crunch, isJ);\n m = m > 0 ? crunch - m : 0;\n\n // Convert the Mandelbrot index into a color\n uchar4 color;\n\n if (m) {\n m += animationFrame;\n color.x = m * colors.x;\n color.y = m * colors.y;\n color.z = m * colors.z;\n }\n else {\n color.x = 0;\n color.y = 0;\n color.z = 0;\n }\n\n // Output the pixel\n int frame1 = frame + 1;\n int frame2 = frame1 / 2;\n dst[pixel].x = (pixelColor.x * frame + color.x + frame2) / frame1;\n dst[pixel].y = (pixelColor.y * frame + color.y + frame2) / frame1;\n dst[pixel].z = (pixelColor.z * frame + color.z + frame2) / frame1;\n }\n }\n }\n\n} // MandelbrotDS1\n\n// The host CPU Mandelbrot thread spawner\nvoid RunMandelbrot0(uchar4 *dst,\n const int imageW,\n const int imageH,\n const int crunch,\n const double xOff,\n const double yOff,\n const double xjp,\n const double yjp,\n const double scale,\n const uchar4 colors,\n const int frame,\n const int animationFrame,\n const int mode,\n const int numSMs,\n const bool isJ,\n int version)\n{\n dim3 threads(BLOCKDIM_X, BLOCKDIM_Y);\n dim3 grid(iDivUp(imageW, BLOCKDIM_X), iDivUp(imageH, BLOCKDIM_Y));\n\n int numWorkerBlocks = numSMs;\n\n switch (mode) {\n default:\n case 0:\n Mandelbrot0<<>>(dst,\n imageW,\n imageH,\n crunch,\n (float)xOff,\n (float)yOff,\n (float)xjp,\n (float)yjp,\n (float)scale,\n colors,\n frame,\n animationFrame,\n grid.x,\n grid.x * grid.y,\n isJ);\n break;\n case 1:\n float x0, x1, y0, y1;\n dsdeq(x0, x1, xOff);\n dsdeq(y0, y1, yOff);\n MandelbrotDS0<<>>(dst,\n imageW,\n imageH,\n crunch,\n x0,\n x1,\n y0,\n y1,\n (float)xjp,\n (float)yjp,\n (float)scale,\n colors,\n frame,\n animationFrame,\n grid.x,\n grid.x * grid.y,\n isJ);\n break;\n case 2:\n Mandelbrot0<<>>(dst,\n imageW,\n imageH,\n crunch,\n xOff,\n yOff,\n xjp,\n yjp,\n scale,\n colors,\n frame,\n animationFrame,\n grid.x,\n grid.x * grid.y,\n isJ);\n break;\n }\n\n getLastCudaError(\"Mandelbrot0 kernel execution failed.\\n\");\n} // RunMandelbrot0\n\n// The host CPU Mandelbrot thread spawner\nvoid RunMandelbrot1(uchar4 *dst,\n const int imageW,\n const int imageH,\n const int crunch,\n const double xOff,\n const double yOff,\n const double xjp,\n const double yjp,\n const double scale,\n const uchar4 colors,\n const int frame,\n const int animationFrame,\n const int mode,\n const int numSMs,\n const bool isJ,\n int version)\n{\n dim3 threads(BLOCKDIM_X, BLOCKDIM_Y);\n dim3 grid(iDivUp(imageW, BLOCKDIM_X), iDivUp(imageH, BLOCKDIM_Y));\n\n int numWorkerBlocks = numSMs;\n\n switch (mode) {\n default:\n case 0:\n Mandelbrot1<<>>(dst,\n imageW,\n imageH,\n crunch,\n (float)xOff,\n (float)yOff,\n (float)xjp,\n (float)yjp,\n (float)scale,\n colors,\n frame,\n animationFrame,\n grid.x,\n grid.x * grid.y,\n isJ);\n break;\n case 1:\n float x0, x1, y0, y1;\n dsdeq(x0, x1, xOff);\n dsdeq(y0, y1, yOff);\n MandelbrotDS1<<>>(dst,\n imageW,\n imageH,\n crunch,\n x0,\n x1,\n y0,\n y1,\n (float)xjp,\n (float)yjp,\n (float)scale,\n colors,\n frame,\n animationFrame,\n grid.x,\n grid.x * grid.y,\n isJ);\n break;\n case 2:\n Mandelbrot1<<>>(dst,\n imageW,\n imageH,\n crunch,\n xOff,\n yOff,\n xjp,\n yjp,\n scale,\n colors,\n frame,\n animationFrame,\n grid.x,\n grid.x * grid.y,\n isJ);\n break;\n }\n\n getLastCudaError(\"Mandelbrot1 kernel execution failed.\\n\");\n} // RunMandelbrot1\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/Mandelbrot/Mandelbrot_cuda.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n\n#define PI 3.1415926536f\n\n/*\n * Paint a 2D surface with a moving bulls-eye pattern. The \"face\" parameter\n * selects\n * between 6 different colors to use. We will use a different color on each\n * face of a\n * cube map.\n */\n__global__ void cuda_kernel_texture_cube(char *surface, int width, int height, size_t pitch, int face, float t)\n{\n int x = blockIdx.x * blockDim.x + threadIdx.x;\n int y = blockIdx.y * blockDim.y + threadIdx.y;\n unsigned char *pixel;\n\n // in the case where, due to quantization into grids, we have\n // more threads than pixels, skip the threads which don't\n // correspond to valid pixels\n if (x >= width || y >= height)\n return;\n\n // get a pointer to this pixel\n pixel = (unsigned char *)(surface + y * pitch) + 4 * x;\n\n // populate it\n float theta_x = (2.0f * x) / width - 1.0f;\n float theta_y = (2.0f * y) / height - 1.0f;\n float theta = 2.0f * PI * sqrt(theta_x * theta_x + theta_y * theta_y);\n unsigned char value = 255 * (0.6f + 0.4f * cos(theta + t));\n\n pixel[3] = 255; // alpha\n\n if (face % 2) {\n pixel[0] = // blue\n pixel[1] = // green\n pixel[2] = 0.5; // red\n pixel[face / 2] = value;\n }\n else {\n pixel[0] = // blue\n pixel[1] = // green\n pixel[2] = value; // red\n pixel[face / 2] = 0.5;\n }\n}\n\nextern \"C\" void cuda_texture_cube(void *surface, int width, int height, size_t pitch, int face, float t)\n{\n cudaError_t error = cudaSuccess;\n\n dim3 Db = dim3(16, 16); // block dimensions are fixed to be 256 threads\n dim3 Dg = dim3((width + Db.x - 1) / Db.x, (height + Db.y - 1) / Db.y);\n\n cuda_kernel_texture_cube<<>>((char *)surface, width, height, pitch, face, t);\n\n error = cudaGetLastError();\n\n if (error != cudaSuccess) {\n printf(\"cuda_kernel_texture_cube() failed to launch error = %d\\n\", error);\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n#define PI 3.1415926536f\n\n/*\n * Paint a 2D surface with a moving bulls-eye pattern. The \"face\" parameter\n * selects\n * between 6 different colors to use. We will use a different color on each\n * face of a\n * cube map.\n */\n__global__ void cuda_kernel_texture_cube(char *surface, int width, int height, size_t pitch, int face, float t)\n{\n int x = blockIdx.x * blockDim.x + threadIdx.x;\n int y = blockIdx.y * blockDim.y + threadIdx.y;\n unsigned char *pixel;\n\n // in the case where, due to quantization into grids, we have\n // more threads than pixels, skip the threads which don't\n // correspond to valid pixels\n if (x >= width || y >= height)\n return;\n\n // get a pointer to this pixel\n pixel = (unsigned char *)(surface + y * pitch) + 4 * x;\n\n // populate it\n float theta_x = (2.0f * x) / width - 1.0f;\n float theta_y = (2.0f * y) / height - 1.0f;\n float theta = 2.0f * PI * sqrt(theta_x * theta_x + theta_y * theta_y);\n unsigned char value = 255 * (0.6f + 0.4f * cos(theta + t));\n\n pixel[3] = 255; // alpha\n\n if (face % 2) {\n pixel[0] = // blue\n pixel[1] = // green\n pixel[2] = 0.5; // red\n pixel[face / 2] = value;\n }\n else {\n pixel[0] = // blue\n pixel[1] = // green\n pixel[2] = value; // red\n pixel[face / 2] = 0.5;\n }\n}\n\nextern \"C\" void cuda_texture_cube(void *surface, int width, int height, size_t pitch, int face, float t)\n{\n hipError_t error = hipSuccess;\n\n dim3 Db = dim3(16, 16); // block dimensions are fixed to be 256 threads\n dim3 Dg = dim3((width + Db.x - 1) / Db.x, (height + Db.y - 1) / Db.y);\n\n cuda_kernel_texture_cube<<>>((char *)surface, width, height, pitch, face, t);\n\n error = hipGetLastError();\n\n if (error != hipSuccess) {\n printf(\"cuda_kernel_texture_cube() failed to launch error = %d\\n\", error);\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/simpleD3D11Texture/texture_cube.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n\n/*\n * Paint a 3D texture with a gradient in X (blue) and Z (green), and have every\n * other Z slice have full red.\n */\n__global__ void cuda_kernel_texture_3d(unsigned char *surface,\n int width,\n int height,\n int depth,\n size_t pitch,\n size_t pitchSlice,\n float t)\n{\n int x = blockIdx.x * blockDim.x + threadIdx.x;\n int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // in the case where, due to quantization into grids, we have\n // more threads than pixels, skip the threads which don't\n // correspond to valid pixels\n if (x >= width || y >= height)\n return;\n\n // walk across the Z slices of this texture. it should be noted that\n // this is far from optimal data access.\n for (int z = 0; z < depth; ++z) {\n // get a pointer to this pixel\n unsigned char *pixel = surface + z * pitchSlice + y * pitch + 4 * x;\n pixel[0] = (unsigned char)(255.f * (0.5f + 0.5f * cos(t + (x * x + y * y + z * z) * 0.0001f * 3.14f))); // red\n pixel[1] = (unsigned char)(255.f * (0.5f + 0.5f * sin(t + (x * x + y * y + z * z) * 0.0001f * 3.14f))); // green\n pixel[2] = (unsigned char)0; // blue\n pixel[3] = 255; // alpha\n }\n}\n\nextern \"C\" void\ncuda_texture_3d(void *surface, int width, int height, int depth, size_t pitch, size_t pitchSlice, float t)\n{\n cudaError_t error = cudaSuccess;\n\n dim3 Db = dim3(16, 16); // block dimensions are fixed to be 256 threads\n dim3 Dg = dim3((width + Db.x - 1) / Db.x, (height + Db.y - 1) / Db.y);\n\n cuda_kernel_texture_3d<<>>((unsigned char *)surface, width, height, depth, pitch, pitchSlice, t);\n\n error = cudaGetLastError();\n\n if (error != cudaSuccess) {\n printf(\"cuda_kernel_texture_3d() failed to launch error = %d\\n\", error);\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n/*\n * Paint a 3D texture with a gradient in X (blue) and Z (green), and have every\n * other Z slice have full red.\n */\n__global__ void cuda_kernel_texture_3d(unsigned char *surface,\n int width,\n int height,\n int depth,\n size_t pitch,\n size_t pitchSlice,\n float t)\n{\n int x = blockIdx.x * blockDim.x + threadIdx.x;\n int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // in the case where, due to quantization into grids, we have\n // more threads than pixels, skip the threads which don't\n // correspond to valid pixels\n if (x >= width || y >= height)\n return;\n\n // walk across the Z slices of this texture. it should be noted that\n // this is far from optimal data access.\n for (int z = 0; z < depth; ++z) {\n // get a pointer to this pixel\n unsigned char *pixel = surface + z * pitchSlice + y * pitch + 4 * x;\n pixel[0] = (unsigned char)(255.f * (0.5f + 0.5f * cos(t + (x * x + y * y + z * z) * 0.0001f * 3.14f))); // red\n pixel[1] = (unsigned char)(255.f * (0.5f + 0.5f * sin(t + (x * x + y * y + z * z) * 0.0001f * 3.14f))); // green\n pixel[2] = (unsigned char)0; // blue\n pixel[3] = 255; // alpha\n }\n}\n\nextern \"C\" void\ncuda_texture_3d(void *surface, int width, int height, int depth, size_t pitch, size_t pitchSlice, float t)\n{\n hipError_t error = hipSuccess;\n\n dim3 Db = dim3(16, 16); // block dimensions are fixed to be 256 threads\n dim3 Dg = dim3((width + Db.x - 1) / Db.x, (height + Db.y - 1) / Db.y);\n\n cuda_kernel_texture_3d<<>>((unsigned char *)surface, width, height, depth, pitch, pitchSlice, t);\n\n error = hipGetLastError();\n\n if (error != hipSuccess) {\n printf(\"cuda_kernel_texture_3d() failed to launch error = %d\\n\", error);\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/simpleD3D11Texture/texture_3d.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n\n#define PI 3.1415926536f\n\n/*\n * Paint a 2D texture with a moving red/green hatch pattern on a\n * strobing blue background. Note that this kernel reads to and\n * writes from the texture, hence why this texture was not mapped\n * as WriteDiscard.\n */\n__global__ void cuda_kernel_texture_2d(unsigned char *surface, int width, int height, size_t pitch, float t)\n{\n int x = blockIdx.x * blockDim.x + threadIdx.x;\n int y = blockIdx.y * blockDim.y + threadIdx.y;\n float *pixel;\n\n // in the case where, due to quantization into grids, we have\n // more threads than pixels, skip the threads which don't\n // correspond to valid pixels\n if (x >= width || y >= height)\n return;\n\n // get a pointer to the pixel at (x,y)\n pixel = (float *)(surface + y * pitch) + 4 * x;\n\n // populate it\n float value_x = 0.5f + 0.5f * cos(t + 10.0f * ((2.0f * x) / width - 1.0f));\n float value_y = 0.5f + 0.5f * cos(t + 10.0f * ((2.0f * y) / height - 1.0f));\n pixel[0] = 0.5 * pixel[0] + 0.5 * pow(value_x, 3.0f); // red\n pixel[1] = 0.5 * pixel[1] + 0.5 * pow(value_y, 3.0f); // green\n pixel[2] = 0.5f + 0.5f * cos(t); // blue\n pixel[3] = 1; // alpha\n}\n\nextern \"C\" void cuda_texture_2d(void *surface, int width, int height, size_t pitch, float t)\n{\n cudaError_t error = cudaSuccess;\n\n dim3 Db = dim3(16, 16); // block dimensions are fixed to be 256 threads\n dim3 Dg = dim3((width + Db.x - 1) / Db.x, (height + Db.y - 1) / Db.y);\n\n cuda_kernel_texture_2d<<>>((unsigned char *)surface, width, height, pitch, t);\n\n error = cudaGetLastError();\n\n if (error != cudaSuccess) {\n printf(\"cuda_kernel_texture_2d() failed to launch error = %d\\n\", error);\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n#define PI 3.1415926536f\n\n/*\n * Paint a 2D texture with a moving red/green hatch pattern on a\n * strobing blue background. Note that this kernel reads to and\n * writes from the texture, hence why this texture was not mapped\n * as WriteDiscard.\n */\n__global__ void cuda_kernel_texture_2d(unsigned char *surface, int width, int height, size_t pitch, float t)\n{\n int x = blockIdx.x * blockDim.x + threadIdx.x;\n int y = blockIdx.y * blockDim.y + threadIdx.y;\n float *pixel;\n\n // in the case where, due to quantization into grids, we have\n // more threads than pixels, skip the threads which don't\n // correspond to valid pixels\n if (x >= width || y >= height)\n return;\n\n // get a pointer to the pixel at (x,y)\n pixel = (float *)(surface + y * pitch) + 4 * x;\n\n // populate it\n float value_x = 0.5f + 0.5f * cos(t + 10.0f * ((2.0f * x) / width - 1.0f));\n float value_y = 0.5f + 0.5f * cos(t + 10.0f * ((2.0f * y) / height - 1.0f));\n pixel[0] = 0.5 * pixel[0] + 0.5 * pow(value_x, 3.0f); // red\n pixel[1] = 0.5 * pixel[1] + 0.5 * pow(value_y, 3.0f); // green\n pixel[2] = 0.5f + 0.5f * cos(t); // blue\n pixel[3] = 1; // alpha\n}\n\nextern \"C\" void cuda_texture_2d(void *surface, int width, int height, size_t pitch, float t)\n{\n hipError_t error = hipSuccess;\n\n dim3 Db = dim3(16, 16); // block dimensions are fixed to be 256 threads\n dim3 Dg = dim3((width + Db.x - 1) / Db.x, (height + Db.y - 1) / Db.y);\n\n cuda_kernel_texture_2d<<>>((unsigned char *)surface, width, height, pitch, t);\n\n error = hipGetLastError();\n\n if (error != hipSuccess) {\n printf(\"cuda_kernel_texture_2d() failed to launch error = %d\\n\", error);\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/simpleD3D11Texture/texture_2d.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n\n#if defined(__APPLE__) || defined(MACOSX)\n#pragma clang diagnostic ignored \"-Wdeprecated-declarations\"\n#include \n#else\n#include \n#endif\n\n// CUDA standard includes\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \"bodysystem.h\"\n\n__constant__ float softeningSquared;\n__constant__ double softeningSquared_fp64;\n\ncudaError_t setSofteningSquared(float softeningSq)\n{\n return cudaMemcpyToSymbol(softeningSquared, &softeningSq, sizeof(float), 0, cudaMemcpyHostToDevice);\n}\n\ncudaError_t setSofteningSquared(double softeningSq)\n{\n return cudaMemcpyToSymbol(softeningSquared_fp64, &softeningSq, sizeof(double), 0, cudaMemcpyHostToDevice);\n}\n\ntemplate struct SharedMemory\n{\n __device__ inline operator T *()\n {\n extern __shared__ int __smem[];\n return (T *)__smem;\n }\n\n __device__ inline operator const T *() const\n {\n extern __shared__ int __smem[];\n return (T *)__smem;\n }\n};\n\ntemplate __device__ T rsqrt_T(T x) { return rsqrt(x); }\n\ntemplate <> __device__ float rsqrt_T(float x) { return rsqrtf(x); }\n\ntemplate <> __device__ double rsqrt_T(double x) { return rsqrt(x); }\n\n// Macros to simplify shared memory addressing\n#define SX(i) sharedPos[i + blockDim.x * threadIdx.y]\n// This macro is only used when multithreadBodies is true (below)\n#define SX_SUM(i, j) sharedPos[i + blockDim.x * j]\n\ntemplate __device__ T getSofteningSquared() { return softeningSquared; }\ntemplate <> __device__ double getSofteningSquared() { return softeningSquared_fp64; }\n\ntemplate struct DeviceData\n{\n T *dPos[2]; // mapped host pointers\n T *dVel;\n cudaEvent_t event;\n unsigned int offset;\n unsigned int numBodies;\n};\n\ntemplate \n__device__ typename vec3::Type\nbodyBodyInteraction(typename vec3::Type ai, typename vec4::Type bi, typename vec4::Type bj)\n{\n typename vec3::Type r;\n\n // r_ij [3 FLOPS]\n r.x = bj.x - bi.x;\n r.y = bj.y - bi.y;\n r.z = bj.z - bi.z;\n\n // distSqr = dot(r_ij, r_ij) + EPS^2 [6 FLOPS]\n T distSqr = r.x * r.x + r.y * r.y + r.z * r.z;\n distSqr += getSofteningSquared();\n\n // invDistCube =1/distSqr^(3/2) [4 FLOPS (2 mul, 1 sqrt, 1 inv)]\n T invDist = rsqrt_T(distSqr);\n T invDistCube = invDist * invDist * invDist;\n\n // s = m_j * invDistCube [1 FLOP]\n T s = bj.w * invDistCube;\n\n // a_i = a_i + s * r_ij [6 FLOPS]\n ai.x += r.x * s;\n ai.y += r.y * s;\n ai.z += r.z * s;\n\n return ai;\n}\n\ntemplate \n__device__ typename vec3::Type\ncomputeBodyAccel(typename vec4::Type bodyPos, typename vec4::Type *positions, int numTiles, cg::thread_block cta)\n{\n typename vec4::Type *sharedPos = SharedMemory::Type>();\n\n typename vec3::Type acc = {0.0f, 0.0f, 0.0f};\n\n for (int tile = 0; tile < numTiles; tile++) {\n sharedPos[threadIdx.x] = positions[tile * blockDim.x + threadIdx.x];\n\n cg::sync(cta);\n\n// This is the \"tile_calculation\" from the GPUG3 article.\n#pragma unroll 128\n\n for (unsigned int counter = 0; counter < blockDim.x; counter++) {\n acc = bodyBodyInteraction(acc, bodyPos, sharedPos[counter]);\n }\n\n cg::sync(cta);\n }\n\n return acc;\n}\n\ntemplate \n__global__ void integrateBodies(typename vec4::Type *__restrict__ newPos,\n typename vec4::Type *__restrict__ oldPos,\n typename vec4::Type *vel,\n unsigned int deviceOffset,\n unsigned int deviceNumBodies,\n float deltaTime,\n float damping,\n int numTiles)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (index >= deviceNumBodies) {\n return;\n }\n\n typename vec4::Type position = oldPos[deviceOffset + index];\n\n typename vec3::Type accel = computeBodyAccel(position, oldPos, numTiles, cta);\n\n // acceleration = force / mass;\n // new velocity = old velocity + acceleration * deltaTime\n // note we factor out the body's mass from the equation, here and in\n // bodyBodyInteraction\n // (because they cancel out). Thus here force == acceleration\n typename vec4::Type velocity = vel[deviceOffset + index];\n\n velocity.x += accel.x * deltaTime;\n velocity.y += accel.y * deltaTime;\n velocity.z += accel.z * deltaTime;\n\n velocity.x *= damping;\n velocity.y *= damping;\n velocity.z *= damping;\n\n // new position = old position + velocity * deltaTime\n position.x += velocity.x * deltaTime;\n position.y += velocity.y * deltaTime;\n position.z += velocity.z * deltaTime;\n\n // store new position and velocity\n newPos[deviceOffset + index] = position;\n vel[deviceOffset + index] = velocity;\n}\n\ntemplate \nvoid integrateNbodySystem(DeviceData *deviceData,\n cudaGraphicsResource **pgres,\n unsigned int currentRead,\n float deltaTime,\n float damping,\n unsigned int numBodies,\n unsigned int numDevices,\n int blockSize,\n bool bUsePBO)\n{\n if (bUsePBO) {\n checkCudaErrors(cudaGraphicsResourceSetMapFlags(pgres[currentRead], cudaGraphicsMapFlagsReadOnly));\n checkCudaErrors(cudaGraphicsResourceSetMapFlags(pgres[1 - currentRead], cudaGraphicsMapFlagsWriteDiscard));\n checkCudaErrors(cudaGraphicsMapResources(2, pgres, 0));\n size_t bytes;\n checkCudaErrors(cudaGraphicsResourceGetMappedPointer(\n (void **)&(deviceData[0].dPos[currentRead]), &bytes, pgres[currentRead]));\n checkCudaErrors(cudaGraphicsResourceGetMappedPointer(\n (void **)&(deviceData[0].dPos[1 - currentRead]), &bytes, pgres[1 - currentRead]));\n }\n\n for (unsigned int dev = 0; dev != numDevices; dev++) {\n if (numDevices > 1) {\n cudaSetDevice(dev);\n }\n\n int numBlocks = (deviceData[dev].numBodies + blockSize - 1) / blockSize;\n int numTiles = (numBodies + blockSize - 1) / blockSize;\n int sharedMemSize = blockSize * 4 * sizeof(T); // 4 floats for pos\n\n integrateBodies\n <<>>((typename vec4::Type *)deviceData[dev].dPos[1 - currentRead],\n (typename vec4::Type *)deviceData[dev].dPos[currentRead],\n (typename vec4::Type *)deviceData[dev].dVel,\n deviceData[dev].offset,\n deviceData[dev].numBodies,\n deltaTime,\n damping,\n numTiles);\n\n if (numDevices > 1) {\n checkCudaErrors(cudaEventRecord(deviceData[dev].event));\n // MJH: Hack on older driver versions to force kernel launches to flush!\n cudaStreamQuery(0);\n }\n\n // check if kernel invocation generated an error\n getLastCudaError(\"Kernel execution failed\");\n }\n\n if (numDevices > 1) {\n for (unsigned int dev = 0; dev < numDevices; dev++) {\n checkCudaErrors(cudaEventSynchronize(deviceData[dev].event));\n }\n }\n\n if (bUsePBO) {\n checkCudaErrors(cudaGraphicsUnmapResources(2, pgres, 0));\n }\n}\n\n// Explicit specializations needed to generate code\ntemplate void integrateNbodySystem(DeviceData *deviceData,\n cudaGraphicsResource **pgres,\n unsigned int currentRead,\n float deltaTime,\n float damping,\n unsigned int numBodies,\n unsigned int numDevices,\n int blockSize,\n bool bUsePBO);\n\ntemplate void integrateNbodySystem(DeviceData *deviceData,\n cudaGraphicsResource **pgres,\n unsigned int currentRead,\n float deltaTime,\n float damping,\n unsigned int numBodies,\n unsigned int numDevices,\n int blockSize,\n bool bUsePBO);\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n\n#if defined(__APPLE__) || defined(MACOSX)\n#pragma clang diagnostic ignored \"-Wdeprecated-declarations\"\n#include \n#else\n#include \n#endif\n\n// CUDA standard includes\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \"bodysystem.h\"\n\n__constant__ float softeningSquared;\n__constant__ double softeningSquared_fp64;\n\nhipError_t setSofteningSquared(float softeningSq)\n{\n return hipMemcpyToSymbol(HIP_SYMBOL(softeningSquared), &softeningSq, sizeof(float), 0, hipMemcpyHostToDevice);\n}\n\nhipError_t setSofteningSquared(double softeningSq)\n{\n return hipMemcpyToSymbol(HIP_SYMBOL(softeningSquared_fp64), &softeningSq, sizeof(double), 0, hipMemcpyHostToDevice);\n}\n\ntemplate struct SharedMemory\n{\n __device__ inline operator T *()\n {\n extern __shared__ int __smem[];\n return (T *)__smem;\n }\n\n __device__ inline operator const T *() const\n {\n extern __shared__ int __smem[];\n return (T *)__smem;\n }\n};\n\ntemplate __device__ T rsqrt_T(T x) { return rsqrt(x); }\n\ntemplate <> __device__ float rsqrt_T(float x) { return rsqrtf(x); }\n\ntemplate <> __device__ double rsqrt_T(double x) { return rsqrt(x); }\n\n// Macros to simplify shared memory addressing\n#define SX(i) sharedPos[i + blockDim.x * threadIdx.y]\n// This macro is only used when multithreadBodies is true (below)\n#define SX_SUM(i, j) sharedPos[i + blockDim.x * j]\n\ntemplate __device__ T getSofteningSquared() { return softeningSquared; }\ntemplate <> __device__ double getSofteningSquared() { return softeningSquared_fp64; }\n\ntemplate struct DeviceData\n{\n T *dPos[2]; // mapped host pointers\n T *dVel;\n hipEvent_t event;\n unsigned int offset;\n unsigned int numBodies;\n};\n\ntemplate \n__device__ typename vec3::Type\nbodyBodyInteraction(typename vec3::Type ai, typename vec4::Type bi, typename vec4::Type bj)\n{\n typename vec3::Type r;\n\n // r_ij [3 FLOPS]\n r.x = bj.x - bi.x;\n r.y = bj.y - bi.y;\n r.z = bj.z - bi.z;\n\n // distSqr = dot(r_ij, r_ij) + EPS^2 [6 FLOPS]\n T distSqr = r.x * r.x + r.y * r.y + r.z * r.z;\n distSqr += getSofteningSquared();\n\n // invDistCube =1/distSqr^(3/2) [4 FLOPS (2 mul, 1 sqrt, 1 inv)]\n T invDist = rsqrt_T(distSqr);\n T invDistCube = invDist * invDist * invDist;\n\n // s = m_j * invDistCube [1 FLOP]\n T s = bj.w * invDistCube;\n\n // a_i = a_i + s * r_ij [6 FLOPS]\n ai.x += r.x * s;\n ai.y += r.y * s;\n ai.z += r.z * s;\n\n return ai;\n}\n\ntemplate \n__device__ typename vec3::Type\ncomputeBodyAccel(typename vec4::Type bodyPos, typename vec4::Type *positions, int numTiles, cg::thread_block cta)\n{\n typename vec4::Type *sharedPos = SharedMemory::Type>();\n\n typename vec3::Type acc = {0.0f, 0.0f, 0.0f};\n\n for (int tile = 0; tile < numTiles; tile++) {\n sharedPos[threadIdx.x] = positions[tile * blockDim.x + threadIdx.x];\n\n cg::sync(cta);\n\n// This is the \"tile_calculation\" from the GPUG3 article.\n#pragma unroll 128\n\n for (unsigned int counter = 0; counter < blockDim.x; counter++) {\n acc = bodyBodyInteraction(acc, bodyPos, sharedPos[counter]);\n }\n\n cg::sync(cta);\n }\n\n return acc;\n}\n\ntemplate \n__global__ void integrateBodies(typename vec4::Type *__restrict__ newPos,\n typename vec4::Type *__restrict__ oldPos,\n typename vec4::Type *vel,\n unsigned int deviceOffset,\n unsigned int deviceNumBodies,\n float deltaTime,\n float damping,\n int numTiles)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (index >= deviceNumBodies) {\n return;\n }\n\n typename vec4::Type position = oldPos[deviceOffset + index];\n\n typename vec3::Type accel = computeBodyAccel(position, oldPos, numTiles, cta);\n\n // acceleration = force / mass;\n // new velocity = old velocity + acceleration * deltaTime\n // note we factor out the body's mass from the equation, here and in\n // bodyBodyInteraction\n // (because they cancel out). Thus here force == acceleration\n typename vec4::Type velocity = vel[deviceOffset + index];\n\n velocity.x += accel.x * deltaTime;\n velocity.y += accel.y * deltaTime;\n velocity.z += accel.z * deltaTime;\n\n velocity.x *= damping;\n velocity.y *= damping;\n velocity.z *= damping;\n\n // new position = old position + velocity * deltaTime\n position.x += velocity.x * deltaTime;\n position.y += velocity.y * deltaTime;\n position.z += velocity.z * deltaTime;\n\n // store new position and velocity\n newPos[deviceOffset + index] = position;\n vel[deviceOffset + index] = velocity;\n}\n\ntemplate \nvoid integrateNbodySystem(DeviceData *deviceData,\n hipGraphicsResource **pgres,\n unsigned int currentRead,\n float deltaTime,\n float damping,\n unsigned int numBodies,\n unsigned int numDevices,\n int blockSize,\n bool bUsePBO)\n{\n if (bUsePBO) {\n checkCudaErrors(cudaGraphicsResourceSetMapFlags(pgres[currentRead], cudaGraphicsMapFlagsReadOnly));\n checkCudaErrors(cudaGraphicsResourceSetMapFlags(pgres[1 - currentRead], cudaGraphicsMapFlagsWriteDiscard));\n checkCudaErrors(hipGraphicsMapResources(2, pgres, 0));\n size_t bytes;\n checkCudaErrors(hipGraphicsResourceGetMappedPointer(\n (void **)&(deviceData[0].dPos[currentRead]), &bytes, pgres[currentRead]));\n checkCudaErrors(hipGraphicsResourceGetMappedPointer(\n (void **)&(deviceData[0].dPos[1 - currentRead]), &bytes, pgres[1 - currentRead]));\n }\n\n for (unsigned int dev = 0; dev != numDevices; dev++) {\n if (numDevices > 1) {\n hipSetDevice(dev);\n }\n\n int numBlocks = (deviceData[dev].numBodies + blockSize - 1) / blockSize;\n int numTiles = (numBodies + blockSize - 1) / blockSize;\n int sharedMemSize = blockSize * 4 * sizeof(T); // 4 floats for pos\n\n integrateBodies\n <<>>((typename vec4::Type *)deviceData[dev].dPos[1 - currentRead],\n (typename vec4::Type *)deviceData[dev].dPos[currentRead],\n (typename vec4::Type *)deviceData[dev].dVel,\n deviceData[dev].offset,\n deviceData[dev].numBodies,\n deltaTime,\n damping,\n numTiles);\n\n if (numDevices > 1) {\n checkCudaErrors(hipEventRecord(deviceData[dev].event));\n // MJH: Hack on older driver versions to force kernel launches to flush!\n hipStreamQuery(0);\n }\n\n // check if kernel invocation generated an error\n getLastCudaError(\"Kernel execution failed\");\n }\n\n if (numDevices > 1) {\n for (unsigned int dev = 0; dev < numDevices; dev++) {\n checkCudaErrors(hipEventSynchronize(deviceData[dev].event));\n }\n }\n\n if (bUsePBO) {\n checkCudaErrors(hipGraphicsUnmapResources(2, pgres, 0));\n }\n}\n\n// Explicit specializations needed to generate code\ntemplate void integrateNbodySystem(DeviceData *deviceData,\n hipGraphicsResource **pgres,\n unsigned int currentRead,\n float deltaTime,\n float damping,\n unsigned int numBodies,\n unsigned int numDevices,\n int blockSize,\n bool bUsePBO);\n\ntemplate void integrateNbodySystem(DeviceData *deviceData,\n hipGraphicsResource **pgres,\n unsigned int currentRead,\n float deltaTime,\n float damping,\n unsigned int numBodies,\n unsigned int numDevices,\n int blockSize,\n bool bUsePBO);\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/nbody/bodysystemcuda.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include // CUDA FFT Libraries\n#include // Helper functions for CUDA Error handling\n#include \n#include \n\n// OpenGL Graphics includes\n#define HELPERGL_EXTERN_GL_FUNC_IMPLEMENTATION\n#include \n\n// FluidsGL CUDA kernel definitions\n#include \"fluidsGL_kernels.cuh\"\n\n// Texture object for reading velocity field\ncudaTextureObject_t texObj;\nstatic cudaArray *array = NULL;\n\n// Particle data\nextern GLuint vbo; // OpenGL vertex buffer object\nextern struct cudaGraphicsResource *cuda_vbo_resource; // handles OpenGL-CUDA exchange\n\n// Texture pitch\nextern size_t tPitch;\nextern cufftHandle planr2c;\nextern cufftHandle planc2r;\ncData *vxfield = NULL;\ncData *vyfield = NULL;\n\nvoid setupTexture(int x, int y)\n{\n cudaChannelFormatDesc desc = cudaCreateChannelDesc();\n\n cudaMallocArray(&array, &desc, y, x);\n getLastCudaError(\"cudaMalloc failed\");\n\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = array;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = cudaFilterModeLinear;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&texObj, &texRes, &texDescr, NULL));\n}\n\nvoid updateTexture(cData *data, size_t wib, size_t h, size_t pitch)\n{\n checkCudaErrors(cudaMemcpy2DToArray(array, 0, 0, data, pitch, wib, h, cudaMemcpyDeviceToDevice));\n}\n\nvoid deleteTexture(void)\n{\n checkCudaErrors(cudaDestroyTextureObject(texObj));\n checkCudaErrors(cudaFreeArray(array));\n}\n\n// Note that these kernels are designed to work with arbitrary\n// domain sizes, not just domains that are multiples of the tile\n// size. Therefore, we have extra code that checks to make sure\n// a given thread location falls within the domain boundaries in\n// both X and Y. Also, the domain is covered by looping over\n// multiple elements in the Y direction, while there is a one-to-one\n// mapping between threads in X and the tile size in X.\n// Nolan Goodnight 9/22/06\n\n// This method adds constant force vectors to the velocity field\n// stored in 'v' according to v(x,t+1) = v(x,t) + dt * f.\n__global__ void addForces_k(cData *v, int dx, int dy, int spx, int spy, float fx, float fy, int r, size_t pitch)\n{\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n cData *fj = (cData *)((char *)v + (ty + spy) * pitch) + tx + spx;\n\n cData vterm = *fj;\n tx -= r;\n ty -= r;\n float s = 1.f / (1.f + tx * tx * tx * tx + ty * ty * ty * ty);\n vterm.x += s * fx;\n vterm.y += s * fy;\n *fj = vterm;\n}\n\n// This method performs the velocity advection step, where we\n// trace velocity vectors back in time to update each grid cell.\n// That is, v(x,t+1) = v(p(x,-dt),t). Here we perform bilinear\n// interpolation in the velocity space.\n__global__ void advectVelocity_k(cData *v,\n float *vx,\n float *vy,\n int dx,\n int pdx,\n int dy,\n float dt,\n int lb,\n cudaTextureObject_t texObject)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n cData vterm, ploc;\n float vxterm, vyterm;\n\n // gtidx is the domain location in x for this thread\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fj = fi * pdx + gtidx;\n vterm = tex2D(texObject, (float)gtidx, (float)fi);\n ploc.x = (gtidx + 0.5f) - (dt * vterm.x * dx);\n ploc.y = (fi + 0.5f) - (dt * vterm.y * dy);\n vterm = tex2D(texObject, ploc.x, ploc.y);\n vxterm = vterm.x;\n vyterm = vterm.y;\n vx[fj] = vxterm;\n vy[fj] = vyterm;\n }\n }\n }\n}\n\n// This method performs velocity diffusion and forces mass conservation\n// in the frequency domain. The inputs 'vx' and 'vy' are complex-valued\n// arrays holding the Fourier coefficients of the velocity field in\n// X and Y. Diffusion in this space takes a simple form described as:\n// v(k,t) = v(k,t) / (1 + visc * dt * k^2), where visc is the viscosity,\n// and k is the wavenumber. The projection step forces the Fourier\n// velocity vectors to be orthogonal to the vectors for each\n// wavenumber: v(k,t) = v(k,t) - ((k dot v(k,t) * k) / k^2.\n__global__ void diffuseProject_k(cData *vx, cData *vy, int dx, int dy, float dt, float visc, int lb)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n cData xterm, yterm;\n\n // gtidx is the domain location in x for this thread\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fj = fi * dx + gtidx;\n xterm = vx[fj];\n yterm = vy[fj];\n\n // Compute the index of the wavenumber based on the\n // data order produced by a standard NN FFT.\n int iix = gtidx;\n int iiy = (fi > dy / 2) ? (fi - (dy)) : fi;\n\n // Velocity diffusion\n float kk = (float)(iix * iix + iiy * iiy); // k^2\n float diff = 1.f / (1.f + visc * dt * kk);\n xterm.x *= diff;\n xterm.y *= diff;\n yterm.x *= diff;\n yterm.y *= diff;\n\n // Velocity projection\n if (kk > 0.f) {\n float rkk = 1.f / kk;\n // Real portion of velocity projection\n float rkp = (iix * xterm.x + iiy * yterm.x);\n // Imaginary portion of velocity projection\n float ikp = (iix * xterm.y + iiy * yterm.y);\n xterm.x -= rkk * rkp * iix;\n xterm.y -= rkk * ikp * iix;\n yterm.x -= rkk * rkp * iiy;\n yterm.y -= rkk * ikp * iiy;\n }\n\n vx[fj] = xterm;\n vy[fj] = yterm;\n }\n }\n }\n}\n\n// This method updates the velocity field 'v' using the two complex\n// arrays from the previous step: 'vx' and 'vy'. Here we scale the\n// real components by 1/(dx*dy) to account for an unnormalized FFT.\n__global__ void updateVelocity_k(cData *v, float *vx, float *vy, int dx, int pdx, int dy, int lb, size_t pitch)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n float vxterm, vyterm;\n cData nvterm;\n\n // gtidx is the domain location in x for this thread\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fjr = fi * pdx + gtidx;\n vxterm = vx[fjr];\n vyterm = vy[fjr];\n\n // Normalize the result of the inverse FFT\n float scale = 1.f / (dx * dy);\n nvterm.x = vxterm * scale;\n nvterm.y = vyterm * scale;\n\n cData *fj = (cData *)((char *)v + fi * pitch) + gtidx;\n *fj = nvterm;\n }\n } // If this thread is inside the domain in Y\n } // If this thread is inside the domain in X\n}\n\n// This method updates the particles by moving particle positions\n// according to the velocity field and time step. That is, for each\n// particle: p(t+1) = p(t) + dt * v(p(t)).\n__global__ void advectParticles_k(cData *part, cData *v, int dx, int dy, float dt, int lb, size_t pitch)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n // gtidx is the domain location in x for this thread\n cData pterm, vterm;\n\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fj = fi * dx + gtidx;\n pterm = part[fj];\n\n int xvi = ((int)(pterm.x * dx));\n int yvi = ((int)(pterm.y * dy));\n vterm = *((cData *)((char *)v + yvi * pitch) + xvi);\n\n pterm.x += dt * vterm.x;\n pterm.x = pterm.x - (int)pterm.x;\n pterm.x += 1.f;\n pterm.x = pterm.x - (int)pterm.x;\n pterm.y += dt * vterm.y;\n pterm.y = pterm.y - (int)pterm.y;\n pterm.y += 1.f;\n pterm.y = pterm.y - (int)pterm.y;\n\n part[fj] = pterm;\n }\n } // If this thread is inside the domain in Y\n } // If this thread is inside the domain in X\n}\n\n// These are the external function calls necessary for launching fluid\n// simulation\nextern \"C\" void addForces(cData *v, int dx, int dy, int spx, int spy, float fx, float fy, int r)\n{\n dim3 tids(2 * r + 1, 2 * r + 1);\n\n addForces_k<<<1, tids>>>(v, dx, dy, spx, spy, fx, fy, r, tPitch);\n getLastCudaError(\"addForces_k failed.\");\n}\n\nextern \"C\" void advectVelocity(cData *v, float *vx, float *vy, int dx, int pdx, int dy, float dt)\n{\n dim3 grid((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1));\n\n dim3 tids(TIDSX, TIDSY);\n\n updateTexture(v, DIM * sizeof(cData), DIM, tPitch);\n advectVelocity_k<<>>(v, vx, vy, dx, pdx, dy, dt, TILEY / TIDSY, texObj);\n\n getLastCudaError(\"advectVelocity_k failed.\");\n}\n\nextern \"C\" void diffuseProject(cData *vx, cData *vy, int dx, int dy, float dt, float visc)\n{\n // Forward FFT\n checkCudaErrors(cufftExecR2C(planr2c, (cufftReal *)vx, (cufftComplex *)vx));\n checkCudaErrors(cufftExecR2C(planr2c, (cufftReal *)vy, (cufftComplex *)vy));\n\n uint3 grid = make_uint3((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1), 1);\n uint3 tids = make_uint3(TIDSX, TIDSY, 1);\n\n diffuseProject_k<<>>(vx, vy, dx, dy, dt, visc, TILEY / TIDSY);\n getLastCudaError(\"diffuseProject_k failed.\");\n\n // Inverse FFT\n checkCudaErrors(cufftExecC2R(planc2r, (cufftComplex *)vx, (cufftReal *)vx));\n checkCudaErrors(cufftExecC2R(planc2r, (cufftComplex *)vy, (cufftReal *)vy));\n}\n\nextern \"C\" void updateVelocity(cData *v, float *vx, float *vy, int dx, int pdx, int dy)\n{\n dim3 grid((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1));\n dim3 tids(TIDSX, TIDSY);\n\n updateVelocity_k<<>>(v, vx, vy, dx, pdx, dy, TILEY / TIDSY, tPitch);\n getLastCudaError(\"updateVelocity_k failed.\");\n}\n\nextern \"C\" void advectParticles(GLuint vbo, cData *v, int dx, int dy, float dt)\n{\n dim3 grid((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1));\n dim3 tids(TIDSX, TIDSY);\n\n cData *p;\n cudaGraphicsMapResources(1, &cuda_vbo_resource, 0);\n getLastCudaError(\"cudaGraphicsMapResources failed\");\n\n size_t num_bytes;\n cudaGraphicsResourceGetMappedPointer((void **)&p, &num_bytes, cuda_vbo_resource);\n getLastCudaError(\"cudaGraphicsResourceGetMappedPointer failed\");\n\n advectParticles_k<<>>(p, v, dx, dy, dt, TILEY / TIDSY, tPitch);\n getLastCudaError(\"advectParticles_k failed.\");\n\n cudaGraphicsUnmapResources(1, &cuda_vbo_resource, 0);\n getLastCudaError(\"cudaGraphicsUnmapResources failed\");\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include // CUDA FFT Libraries\n#include // Helper functions for CUDA Error handling\n#include \n#include \n\n// OpenGL Graphics includes\n#define HELPERGL_EXTERN_GL_FUNC_IMPLEMENTATION\n#include \n\n// FluidsGL CUDA kernel definitions\n#include \"fluidsGL_kernels.cuh\"\n\n// Texture object for reading velocity field\nhipTextureObject_t texObj;\nstatic hipArray *array = NULL;\n\n// Particle data\nextern GLuint vbo; // OpenGL vertex buffer object\nextern struct hipGraphicsResource *cuda_vbo_resource; // handles OpenGL-CUDA exchange\n\n// Texture pitch\nextern size_t tPitch;\nextern hipfftHandle planr2c;\nextern hipfftHandle planc2r;\ncData *vxfield = NULL;\ncData *vyfield = NULL;\n\nvoid setupTexture(int x, int y)\n{\n hipChannelFormatDesc desc = hipCreateChannelDesc();\n\n hipMallocArray(&array, &desc, y, x);\n getLastCudaError(\"hipMalloc failed\");\n\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = array;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = hipFilterModeLinear;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&texObj, &texRes, &texDescr, NULL));\n}\n\nvoid updateTexture(cData *data, size_t wib, size_t h, size_t pitch)\n{\n checkCudaErrors(hipMemcpy2DToArray(array, 0, 0, data, pitch, wib, h, hipMemcpyDeviceToDevice));\n}\n\nvoid deleteTexture(void)\n{\n checkCudaErrors(hipDestroyTextureObject(texObj));\n checkCudaErrors(hipFreeArray(array));\n}\n\n// Note that these kernels are designed to work with arbitrary\n// domain sizes, not just domains that are multiples of the tile\n// size. Therefore, we have extra code that checks to make sure\n// a given thread location falls within the domain boundaries in\n// both X and Y. Also, the domain is covered by looping over\n// multiple elements in the Y direction, while there is a one-to-one\n// mapping between threads in X and the tile size in X.\n// Nolan Goodnight 9/22/06\n\n// This method adds constant force vectors to the velocity field\n// stored in 'v' according to v(x,t+1) = v(x,t) + dt * f.\n__global__ void addForces_k(cData *v, int dx, int dy, int spx, int spy, float fx, float fy, int r, size_t pitch)\n{\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n cData *fj = (cData *)((char *)v + (ty + spy) * pitch) + tx + spx;\n\n cData vterm = *fj;\n tx -= r;\n ty -= r;\n float s = 1.f / (1.f + tx * tx * tx * tx + ty * ty * ty * ty);\n vterm.x += s * fx;\n vterm.y += s * fy;\n *fj = vterm;\n}\n\n// This method performs the velocity advection step, where we\n// trace velocity vectors back in time to update each grid cell.\n// That is, v(x,t+1) = v(p(x,-dt),t). Here we perform bilinear\n// interpolation in the velocity space.\n__global__ void advectVelocity_k(cData *v,\n float *vx,\n float *vy,\n int dx,\n int pdx,\n int dy,\n float dt,\n int lb,\n hipTextureObject_t texObject)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n cData vterm, ploc;\n float vxterm, vyterm;\n\n // gtidx is the domain location in x for this thread\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fj = fi * pdx + gtidx;\n vterm = tex2D(texObject, (float)gtidx, (float)fi);\n ploc.x = (gtidx + 0.5f) - (dt * vterm.x * dx);\n ploc.y = (fi + 0.5f) - (dt * vterm.y * dy);\n vterm = tex2D(texObject, ploc.x, ploc.y);\n vxterm = vterm.x;\n vyterm = vterm.y;\n vx[fj] = vxterm;\n vy[fj] = vyterm;\n }\n }\n }\n}\n\n// This method performs velocity diffusion and forces mass conservation\n// in the frequency domain. The inputs 'vx' and 'vy' are complex-valued\n// arrays holding the Fourier coefficients of the velocity field in\n// X and Y. Diffusion in this space takes a simple form described as:\n// v(k,t) = v(k,t) / (1 + visc * dt * k^2), where visc is the viscosity,\n// and k is the wavenumber. The projection step forces the Fourier\n// velocity vectors to be orthogonal to the vectors for each\n// wavenumber: v(k,t) = v(k,t) - ((k dot v(k,t) * k) / k^2.\n__global__ void diffuseProject_k(cData *vx, cData *vy, int dx, int dy, float dt, float visc, int lb)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n cData xterm, yterm;\n\n // gtidx is the domain location in x for this thread\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fj = fi * dx + gtidx;\n xterm = vx[fj];\n yterm = vy[fj];\n\n // Compute the index of the wavenumber based on the\n // data order produced by a standard NN FFT.\n int iix = gtidx;\n int iiy = (fi > dy / 2) ? (fi - (dy)) : fi;\n\n // Velocity diffusion\n float kk = (float)(iix * iix + iiy * iiy); // k^2\n float diff = 1.f / (1.f + visc * dt * kk);\n xterm.x *= diff;\n xterm.y *= diff;\n yterm.x *= diff;\n yterm.y *= diff;\n\n // Velocity projection\n if (kk > 0.f) {\n float rkk = 1.f / kk;\n // Real portion of velocity projection\n float rkp = (iix * xterm.x + iiy * yterm.x);\n // Imaginary portion of velocity projection\n float ikp = (iix * xterm.y + iiy * yterm.y);\n xterm.x -= rkk * rkp * iix;\n xterm.y -= rkk * ikp * iix;\n yterm.x -= rkk * rkp * iiy;\n yterm.y -= rkk * ikp * iiy;\n }\n\n vx[fj] = xterm;\n vy[fj] = yterm;\n }\n }\n }\n}\n\n// This method updates the velocity field 'v' using the two complex\n// arrays from the previous step: 'vx' and 'vy'. Here we scale the\n// real components by 1/(dx*dy) to account for an unnormalized FFT.\n__global__ void updateVelocity_k(cData *v, float *vx, float *vy, int dx, int pdx, int dy, int lb, size_t pitch)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n float vxterm, vyterm;\n cData nvterm;\n\n // gtidx is the domain location in x for this thread\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fjr = fi * pdx + gtidx;\n vxterm = vx[fjr];\n vyterm = vy[fjr];\n\n // Normalize the result of the inverse FFT\n float scale = 1.f / (dx * dy);\n nvterm.x = vxterm * scale;\n nvterm.y = vyterm * scale;\n\n cData *fj = (cData *)((char *)v + fi * pitch) + gtidx;\n *fj = nvterm;\n }\n } // If this thread is inside the domain in Y\n } // If this thread is inside the domain in X\n}\n\n// This method updates the particles by moving particle positions\n// according to the velocity field and time step. That is, for each\n// particle: p(t+1) = p(t) + dt * v(p(t)).\n__global__ void advectParticles_k(cData *part, cData *v, int dx, int dy, float dt, int lb, size_t pitch)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n // gtidx is the domain location in x for this thread\n cData pterm, vterm;\n\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fj = fi * dx + gtidx;\n pterm = part[fj];\n\n int xvi = ((int)(pterm.x * dx));\n int yvi = ((int)(pterm.y * dy));\n vterm = *((cData *)((char *)v + yvi * pitch) + xvi);\n\n pterm.x += dt * vterm.x;\n pterm.x = pterm.x - (int)pterm.x;\n pterm.x += 1.f;\n pterm.x = pterm.x - (int)pterm.x;\n pterm.y += dt * vterm.y;\n pterm.y = pterm.y - (int)pterm.y;\n pterm.y += 1.f;\n pterm.y = pterm.y - (int)pterm.y;\n\n part[fj] = pterm;\n }\n } // If this thread is inside the domain in Y\n } // If this thread is inside the domain in X\n}\n\n// These are the external function calls necessary for launching fluid\n// simulation\nextern \"C\" void addForces(cData *v, int dx, int dy, int spx, int spy, float fx, float fy, int r)\n{\n dim3 tids(2 * r + 1, 2 * r + 1);\n\n addForces_k<<<1, tids>>>(v, dx, dy, spx, spy, fx, fy, r, tPitch);\n getLastCudaError(\"addForces_k failed.\");\n}\n\nextern \"C\" void advectVelocity(cData *v, float *vx, float *vy, int dx, int pdx, int dy, float dt)\n{\n dim3 grid((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1));\n\n dim3 tids(TIDSX, TIDSY);\n\n updateTexture(v, DIM * sizeof(cData), DIM, tPitch);\n advectVelocity_k<<>>(v, vx, vy, dx, pdx, dy, dt, TILEY / TIDSY, texObj);\n\n getLastCudaError(\"advectVelocity_k failed.\");\n}\n\nextern \"C\" void diffuseProject(cData *vx, cData *vy, int dx, int dy, float dt, float visc)\n{\n // Forward FFT\n checkCudaErrors(hipfftExecR2C(planr2c, (hipfftReal *)vx, (hipfftComplex *)vx));\n checkCudaErrors(hipfftExecR2C(planr2c, (hipfftReal *)vy, (hipfftComplex *)vy));\n\n uint3 grid = make_uint3((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1), 1);\n uint3 tids = make_uint3(TIDSX, TIDSY, 1);\n\n diffuseProject_k<<>>(vx, vy, dx, dy, dt, visc, TILEY / TIDSY);\n getLastCudaError(\"diffuseProject_k failed.\");\n\n // Inverse FFT\n checkCudaErrors(hipfftExecC2R(planc2r, (hipfftComplex *)vx, (hipfftReal *)vx));\n checkCudaErrors(hipfftExecC2R(planc2r, (hipfftComplex *)vy, (hipfftReal *)vy));\n}\n\nextern \"C\" void updateVelocity(cData *v, float *vx, float *vy, int dx, int pdx, int dy)\n{\n dim3 grid((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1));\n dim3 tids(TIDSX, TIDSY);\n\n updateVelocity_k<<>>(v, vx, vy, dx, pdx, dy, TILEY / TIDSY, tPitch);\n getLastCudaError(\"updateVelocity_k failed.\");\n}\n\nextern \"C\" void advectParticles(GLuint vbo, cData *v, int dx, int dy, float dt)\n{\n dim3 grid((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1));\n dim3 tids(TIDSX, TIDSY);\n\n cData *p;\n hipGraphicsMapResources(1, &cuda_vbo_resource, 0);\n getLastCudaError(\"hipGraphicsMapResources failed\");\n\n size_t num_bytes;\n hipGraphicsResourceGetMappedPointer((void **)&p, &num_bytes, cuda_vbo_resource);\n getLastCudaError(\"hipGraphicsResourceGetMappedPointer failed\");\n\n advectParticles_k<<>>(p, v, dx, dy, dt, TILEY / TIDSY, tPitch);\n getLastCudaError(\"advectParticles_k failed.\");\n\n hipGraphicsUnmapResources(1, &cuda_vbo_resource, 0);\n getLastCudaError(\"hipGraphicsUnmapResources failed\");\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/fluidsGL/fluidsGL_kernels.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Utilities and system includes\n\n#include \n\nnamespace cg = cooperative_groups;\n\n#include \n\ncudaTextureObject_t inTexObject;\n\n// clamp x to range [a, b]\n__device__ float clamp(float x, float a, float b) { return max(a, min(b, x)); }\n\n__device__ int clamp(int x, int a, int b) { return max(a, min(b, x)); }\n\n// convert floating point rgb color to 8-bit integer\n__device__ int rgbToInt(float r, float g, float b)\n{\n r = clamp(r, 0.0f, 255.0f);\n g = clamp(g, 0.0f, 255.0f);\n b = clamp(b, 0.0f, 255.0f);\n return (int(b) << 16) | (int(g) << 8) | int(r);\n}\n\n// get pixel from 2D image, with clamping to border\n__device__ uchar4 getPixel(int x, int y, cudaTextureObject_t inTex)\n{\n#ifndef USE_TEXTURE_RGBA8UI\n float4 res = tex2D(inTex, x, y);\n uchar4 ucres = make_uchar4(res.x * 255.0f, res.y * 255.0f, res.z * 255.0f, res.w * 255.0f);\n#else\n uchar4 ucres = tex2D(inTex, x, y);\n#endif\n return ucres;\n}\n\n// macros to make indexing shared memory easier\n#define SMEM(X, Y) sdata[(Y) * tilew + (X)]\n\n/*\n 2D convolution using shared memory\n - operates on 8-bit RGB data stored in 32-bit int\n - assumes kernel radius is less than or equal to block size\n - not optimized for performance\n _____________\n | : : |\n |_ _:_____:_ _|\n | | | |\n | | | |\n |_ _|_____|_ _|\n r | : : |\n |___:_____:___|\n r bw r\n <----tilew---->\n*/\n\n__global__ void cudaProcess(unsigned int *g_odata,\n int imgw,\n int imgh,\n int tilew,\n int r,\n float threshold,\n float highlight,\n cudaTextureObject_t inTex)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n extern __shared__ uchar4 sdata[];\n\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n int bw = blockDim.x;\n int bh = blockDim.y;\n int x = blockIdx.x * bw + tx;\n int y = blockIdx.y * bh + ty;\n\n#if 0\n uchar4 c4 = getPixel(x, y);\n g_odata[y*imgw+x] = rgbToInt(c4.z, c4.y, c4.x);\n#else\n // copy tile to shared memory\n // center region\n SMEM(r + tx, r + ty) = getPixel(x, y, inTex);\n\n // borders\n if (threadIdx.x < r) {\n // left\n SMEM(tx, r + ty) = getPixel(x - r, y, inTex);\n // right\n SMEM(r + bw + tx, r + ty) = getPixel(x + bw, y, inTex);\n }\n\n if (threadIdx.y < r) {\n // top\n SMEM(r + tx, ty) = getPixel(x, y - r, inTex);\n // bottom\n SMEM(r + tx, r + bh + ty) = getPixel(x, y + bh, inTex);\n }\n\n // load corners\n if ((threadIdx.x < r) && (threadIdx.y < r)) {\n // tl\n SMEM(tx, ty) = getPixel(x - r, y - r, inTex);\n // bl\n SMEM(tx, r + bh + ty) = getPixel(x - r, y + bh, inTex);\n // tr\n SMEM(r + bw + tx, ty) = getPixel(x + bh, y - r, inTex);\n // br\n SMEM(r + bw + tx, r + bh + ty) = getPixel(x + bw, y + bh, inTex);\n }\n\n // wait for loads to complete\n cg::sync(cta);\n\n // perform convolution\n float rsum = 0.0f;\n float gsum = 0.0f;\n float bsum = 0.0f;\n float samples = 0.0f;\n\n for (int dy = -r; dy <= r; dy++) {\n for (int dx = -r; dx <= r; dx++) {\n#if 0\n // try this to see the benefit of using shared memory\n uchar4 pixel = getPixel(x+dx, y+dy);\n#else\n uchar4 pixel = SMEM(r + tx + dx, r + ty + dy);\n#endif\n\n // only sum pixels within disc-shaped kernel\n float l = dx * dx + dy * dy;\n\n if (l <= r * r) {\n float r = float(pixel.x);\n float g = float(pixel.y);\n float b = float(pixel.z);\n#if 1\n // brighten highlights\n float lum = (r + g + b) / (255 * 3);\n\n if (lum > threshold) {\n r *= highlight;\n g *= highlight;\n b *= highlight;\n }\n\n#endif\n rsum += r;\n gsum += g;\n bsum += b;\n samples += 1.0f;\n }\n }\n }\n\n rsum /= samples;\n gsum /= samples;\n bsum /= samples;\n // ABGR\n g_odata[y * imgw + x] = rgbToInt(rsum, gsum, bsum);\n// g_odata[y*imgw+x] = rgbToInt(x,y,0);\n#endif\n}\n\nextern \"C\" void launch_cudaProcess(dim3 grid,\n dim3 block,\n int sbytes,\n cudaArray *g_data_array,\n unsigned int *g_odata,\n int imgw,\n int imgh,\n int tilew,\n int radius,\n float threshold,\n float highlight)\n{\n struct cudaChannelFormatDesc desc;\n checkCudaErrors(cudaGetChannelDesc(&desc, g_data_array));\n\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = g_data_array;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = cudaFilterModePoint;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&inTexObject, &texRes, &texDescr, NULL));\n\n#if 0\n printf(\"CUDA Array channel descriptor, bits per component:\\n\");\n printf(\"X %d Y %d Z %d W %d, kind %d\\n\",\n desc.x,desc.y,desc.z,desc.w,desc.f);\n\n printf(\"Possible values for channel format kind: i %d, u%d, f%d:\\n\",\n cudaChannelFormatKindSigned, cudaChannelFormatKindUnsigned,\n cudaChannelFormatKindFloat);\n#endif\n\n// printf(\"\\n\");\n#ifdef GPU_PROFILING\n StopWatchInterface *timer = 0;\n sdkCreateTimer(&timer);\n\n int nIter = 30;\n\n for (int i = -1; i < nIter; ++i) {\n if (i == 0) {\n sdkStartTimer(&timer);\n }\n\n#endif\n\n cudaProcess<<>>(\n g_odata, imgw, imgh, block.x + (2 * radius), radius, 0.8f, 4.0f, inTexObject);\n\n#ifdef GPU_PROFILING\n }\n\n cudaDeviceSynchronize();\n sdkStopTimer(&timer);\n double dSeconds = sdkGetTimerValue(&timer) / ((double)nIter * 1000.0);\n double dNumTexels = (double)imgw * (double)imgh;\n double mtexps = 1.0e-6 * dNumTexels / dSeconds;\n\n if (radius == 4) {\n printf(\"\\n\");\n printf(\"postprocessGL, Throughput = %.4f MTexels/s, Time = %.5f s, Size = \"\n \"%.0f Texels, NumDevsUsed = %d, Workgroup = %u\\n\",\n mtexps,\n dSeconds,\n dNumTexels,\n 1,\n block.x * block.y);\n }\n\n#endif\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Utilities and system includes\n\n#include \"hip/hip_runtime.h\"\n#include \n\nnamespace cg = cooperative_groups;\n\n#include \n\nhipTextureObject_t inTexObject;\n\n// clamp x to range [a, b]\n__device__ float clamp(float x, float a, float b) { return max(a, min(b, x)); }\n\n__device__ int clamp(int x, int a, int b) { return max(a, min(b, x)); }\n\n// convert floating point rgb color to 8-bit integer\n__device__ int rgbToInt(float r, float g, float b)\n{\n r = clamp(r, 0.0f, 255.0f);\n g = clamp(g, 0.0f, 255.0f);\n b = clamp(b, 0.0f, 255.0f);\n return (int(b) << 16) | (int(g) << 8) | int(r);\n}\n\n// get pixel from 2D image, with clamping to border\n__device__ uchar4 getPixel(int x, int y, hipTextureObject_t inTex)\n{\n#ifndef USE_TEXTURE_RGBA8UI\n float4 res = tex2D(inTex, x, y);\n uchar4 ucres = make_uchar4(res.x * 255.0f, res.y * 255.0f, res.z * 255.0f, res.w * 255.0f);\n#else\n uchar4 ucres = tex2D(inTex, x, y);\n#endif\n return ucres;\n}\n\n// macros to make indexing shared memory easier\n#define SMEM(X, Y) sdata[(Y) * tilew + (X)]\n\n/*\n 2D convolution using shared memory\n - operates on 8-bit RGB data stored in 32-bit int\n - assumes kernel radius is less than or equal to block size\n - not optimized for performance\n _____________\n | : : |\n |_ _:_____:_ _|\n | | | |\n | | | |\n |_ _|_____|_ _|\n r | : : |\n |___:_____:___|\n r bw r\n <----tilew---->\n*/\n\n__global__ void cudaProcess(unsigned int *g_odata,\n int imgw,\n int imgh,\n int tilew,\n int r,\n float threshold,\n float highlight,\n hipTextureObject_t inTex)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n extern __shared__ uchar4 sdata[];\n\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n int bw = blockDim.x;\n int bh = blockDim.y;\n int x = blockIdx.x * bw + tx;\n int y = blockIdx.y * bh + ty;\n\n#if 0\n uchar4 c4 = getPixel(x, y);\n g_odata[y*imgw+x] = rgbToInt(c4.z, c4.y, c4.x);\n#else\n // copy tile to shared memory\n // center region\n SMEM(r + tx, r + ty) = getPixel(x, y, inTex);\n\n // borders\n if (threadIdx.x < r) {\n // left\n SMEM(tx, r + ty) = getPixel(x - r, y, inTex);\n // right\n SMEM(r + bw + tx, r + ty) = getPixel(x + bw, y, inTex);\n }\n\n if (threadIdx.y < r) {\n // top\n SMEM(r + tx, ty) = getPixel(x, y - r, inTex);\n // bottom\n SMEM(r + tx, r + bh + ty) = getPixel(x, y + bh, inTex);\n }\n\n // load corners\n if ((threadIdx.x < r) && (threadIdx.y < r)) {\n // tl\n SMEM(tx, ty) = getPixel(x - r, y - r, inTex);\n // bl\n SMEM(tx, r + bh + ty) = getPixel(x - r, y + bh, inTex);\n // tr\n SMEM(r + bw + tx, ty) = getPixel(x + bh, y - r, inTex);\n // br\n SMEM(r + bw + tx, r + bh + ty) = getPixel(x + bw, y + bh, inTex);\n }\n\n // wait for loads to complete\n cg::sync(cta);\n\n // perform convolution\n float rsum = 0.0f;\n float gsum = 0.0f;\n float bsum = 0.0f;\n float samples = 0.0f;\n\n for (int dy = -r; dy <= r; dy++) {\n for (int dx = -r; dx <= r; dx++) {\n#if 0\n // try this to see the benefit of using shared memory\n uchar4 pixel = getPixel(x+dx, y+dy);\n#else\n uchar4 pixel = SMEM(r + tx + dx, r + ty + dy);\n#endif\n\n // only sum pixels within disc-shaped kernel\n float l = dx * dx + dy * dy;\n\n if (l <= r * r) {\n float r = float(pixel.x);\n float g = float(pixel.y);\n float b = float(pixel.z);\n#if 1\n // brighten highlights\n float lum = (r + g + b) / (255 * 3);\n\n if (lum > threshold) {\n r *= highlight;\n g *= highlight;\n b *= highlight;\n }\n\n#endif\n rsum += r;\n gsum += g;\n bsum += b;\n samples += 1.0f;\n }\n }\n }\n\n rsum /= samples;\n gsum /= samples;\n bsum /= samples;\n // ABGR\n g_odata[y * imgw + x] = rgbToInt(rsum, gsum, bsum);\n// g_odata[y*imgw+x] = rgbToInt(x,y,0);\n#endif\n}\n\nextern \"C\" void launch_cudaProcess(dim3 grid,\n dim3 block,\n int sbytes,\n hipArray *g_data_array,\n unsigned int *g_odata,\n int imgw,\n int imgh,\n int tilew,\n int radius,\n float threshold,\n float highlight)\n{\n struct hipChannelFormatDesc desc;\n checkCudaErrors(hipGetChannelDesc(&desc, g_data_array));\n\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = g_data_array;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = hipFilterModePoint;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&inTexObject, &texRes, &texDescr, NULL));\n\n#if 0\n printf(\"CUDA Array channel descriptor, bits per component:\\n\");\n printf(\"X %d Y %d Z %d W %d, kind %d\\n\",\n desc.x,desc.y,desc.z,desc.w,desc.f);\n\n printf(\"Possible values for channel format kind: i %d, u%d, f%d:\\n\",\n hipChannelFormatKindSigned, hipChannelFormatKindUnsigned,\n hipChannelFormatKindFloat);\n#endif\n\n// printf(\"\\n\");\n#ifdef GPU_PROFILING\n StopWatchInterface *timer = 0;\n sdkCreateTimer(&timer);\n\n int nIter = 30;\n\n for (int i = -1; i < nIter; ++i) {\n if (i == 0) {\n sdkStartTimer(&timer);\n }\n\n#endif\n\n cudaProcess<<>>(\n g_odata, imgw, imgh, block.x + (2 * radius), radius, 0.8f, 4.0f, inTexObject);\n\n#ifdef GPU_PROFILING\n }\n\n hipDeviceSynchronize();\n sdkStopTimer(&timer);\n double dSeconds = sdkGetTimerValue(&timer) / ((double)nIter * 1000.0);\n double dNumTexels = (double)imgw * (double)imgh;\n double mtexps = 1.0e-6 * dNumTexels / dSeconds;\n\n if (radius == 4) {\n printf(\"\\n\");\n printf(\"postprocessGL, Throughput = %.4f MTexels/s, Time = %.5f s, Size = \"\n \"%.0f Texels, NumDevsUsed = %d, Workgroup = %u\\n\",\n mtexps,\n dSeconds,\n dNumTexels,\n 1,\n block.x * block.y);\n }\n\n#endif\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/postProcessGL/postProcessGL.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n\n#include \"SineWaveSimulation.h\"\n\n__global__ void sinewave(float *heightMap, unsigned int width, unsigned int height, float time)\n{\n const float freq = 4.0f;\n const size_t stride = gridDim.x * blockDim.x;\n\n // Iterate through the entire array in a way that is\n // independent of the grid configuration\n for (size_t tid = blockIdx.x * blockDim.x + threadIdx.x; tid < width * height; tid += stride) {\n // Calculate the x, y coordinates\n const size_t y = tid / width;\n const size_t x = tid - y * width;\n // Normalize x, y to [0,1]\n const float u = ((2.0f * x) / width) - 1.0f;\n const float v = ((2.0f * y) / height) - 1.0f;\n // Calculate the new height value\n const float w = 0.5f * sinf(u * freq + time) * cosf(v * freq + time);\n // Store this new height value\n heightMap[tid] = w;\n }\n}\n\nSineWaveSimulation::SineWaveSimulation(size_t width, size_t height)\n : m_heightMap(nullptr)\n , m_width(width)\n , m_height(height)\n{\n}\n\nvoid SineWaveSimulation::initCudaLaunchConfig(int device)\n{\n cudaDeviceProp prop = {};\n checkCudaErrors(cudaSetDevice(device));\n checkCudaErrors(cudaGetDeviceProperties(&prop, device));\n\n // We don't need large block sizes, since there's not much inter-thread\n // communication\n m_threads = prop.warpSize;\n\n // Use the occupancy calculator and fill the gpu as best as we can\n checkCudaErrors(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&m_blocks, sinewave, prop.warpSize, 0));\n m_blocks *= prop.multiProcessorCount;\n\n // Go ahead and the clamp the blocks to the minimum needed for this\n // height/width\n m_blocks = std::min(m_blocks, (int)((m_width * m_height + m_threads - 1) / m_threads));\n}\n\nint SineWaveSimulation::initCuda(uint8_t *vkDeviceUUID, size_t UUID_SIZE)\n{\n int current_device = 0;\n int device_count = 0;\n int devices_prohibited = 0;\n\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDeviceCount(&device_count));\n\n if (device_count == 0) {\n fprintf(stderr, \"CUDA error: no devices supporting CUDA.\\n\");\n exit(EXIT_FAILURE);\n }\n\n // Find the GPU which is selected by Vulkan\n while (current_device < device_count) {\n cudaGetDeviceProperties(&deviceProp, current_device);\n int computeMode;\n checkCudaErrors(cudaDeviceGetAttribute(&computeMode, cudaDevAttrComputeMode, current_device));\n\n if ((computeMode != cudaComputeModeProhibited)) {\n // Compare the cuda device UUID with vulkan UUID\n int ret = memcmp((void *)&deviceProp.uuid, vkDeviceUUID, UUID_SIZE);\n if (ret == 0) {\n checkCudaErrors(cudaSetDevice(current_device));\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, current_device));\n printf(\"GPU Device %d: \\\"%s\\\" with compute capability %d.%d\\n\\n\",\n current_device,\n deviceProp.name,\n deviceProp.major,\n deviceProp.minor);\n\n return current_device;\n }\n }\n else {\n devices_prohibited++;\n }\n\n current_device++;\n }\n\n if (devices_prohibited == device_count) {\n fprintf(stderr,\n \"CUDA error:\"\n \" No Vulkan-CUDA Interop capable GPU found.\\n\");\n exit(EXIT_FAILURE);\n }\n\n return -1;\n}\n\nSineWaveSimulation::~SineWaveSimulation() { m_heightMap = NULL; }\n\nvoid SineWaveSimulation::initSimulation(float *heights) { m_heightMap = heights; }\n\nvoid SineWaveSimulation::stepSimulation(float time, cudaStream_t stream)\n{\n sinewave<<>>(m_heightMap, m_width, m_height, time);\n getLastCudaError(\"Failed to launch CUDA simulation\");\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n#include \"SineWaveSimulation.h\"\n\n__global__ void sinewave(float *heightMap, unsigned int width, unsigned int height, float time)\n{\n const float freq = 4.0f;\n const size_t stride = gridDim.x * blockDim.x;\n\n // Iterate through the entire array in a way that is\n // independent of the grid configuration\n for (size_t tid = blockIdx.x * blockDim.x + threadIdx.x; tid < width * height; tid += stride) {\n // Calculate the x, y coordinates\n const size_t y = tid / width;\n const size_t x = tid - y * width;\n // Normalize x, y to [0,1]\n const float u = ((2.0f * x) / width) - 1.0f;\n const float v = ((2.0f * y) / height) - 1.0f;\n // Calculate the new height value\n const float w = 0.5f * sinf(u * freq + time) * cosf(v * freq + time);\n // Store this new height value\n heightMap[tid] = w;\n }\n}\n\nSineWaveSimulation::SineWaveSimulation(size_t width, size_t height)\n : m_heightMap(nullptr)\n , m_width(width)\n , m_height(height)\n{\n}\n\nvoid SineWaveSimulation::initCudaLaunchConfig(int device)\n{\n hipDeviceProp_t prop = {};\n checkCudaErrors(hipSetDevice(device));\n checkCudaErrors(hipGetDeviceProperties(&prop, device));\n\n // We don't need large block sizes, since there's not much inter-thread\n // communication\n m_threads = prop.warpSize;\n\n // Use the occupancy calculator and fill the gpu as best as we can\n checkCudaErrors(hipOccupancyMaxActiveBlocksPerMultiprocessor(&m_blocks, sinewave, prop.warpSize, 0));\n m_blocks *= prop.multiProcessorCount;\n\n // Go ahead and the clamp the blocks to the minimum needed for this\n // height/width\n m_blocks = std::min(m_blocks, (int)((m_width * m_height + m_threads - 1) / m_threads));\n}\n\nint SineWaveSimulation::initCuda(uint8_t *vkDeviceUUID, size_t UUID_SIZE)\n{\n int current_device = 0;\n int device_count = 0;\n int devices_prohibited = 0;\n\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDeviceCount(&device_count));\n\n if (device_count == 0) {\n fprintf(stderr, \"CUDA error: no devices supporting CUDA.\\n\");\n exit(EXIT_FAILURE);\n }\n\n // Find the GPU which is selected by Vulkan\n while (current_device < device_count) {\n hipGetDeviceProperties(&deviceProp, current_device);\n int computeMode;\n checkCudaErrors(hipDeviceGetAttribute(&computeMode, hipDeviceAttributeComputeMode, current_device));\n\n if ((computeMode != hipComputeModeProhibited)) {\n // Compare the cuda device UUID with vulkan UUID\n int ret = memcmp((void *)&deviceProp.uuid, vkDeviceUUID, UUID_SIZE);\n if (ret == 0) {\n checkCudaErrors(hipSetDevice(current_device));\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, current_device));\n printf(\"GPU Device %d: \\\"%s\\\" with compute capability %d.%d\\n\\n\",\n current_device,\n deviceProp.name,\n deviceProp.major,\n deviceProp.minor);\n\n return current_device;\n }\n }\n else {\n devices_prohibited++;\n }\n\n current_device++;\n }\n\n if (devices_prohibited == device_count) {\n fprintf(stderr,\n \"CUDA error:\"\n \" No Vulkan-CUDA Interop capable GPU found.\\n\");\n exit(EXIT_FAILURE);\n }\n\n return -1;\n}\n\nSineWaveSimulation::~SineWaveSimulation() { m_heightMap = NULL; }\n\nvoid SineWaveSimulation::initSimulation(float *heights) { m_heightMap = heights; }\n\nvoid SineWaveSimulation::stepSimulation(float time, hipStream_t stream)\n{\n sinewave<<>>(m_heightMap, m_width, m_height, time);\n getLastCudaError(\"Failed to launch CUDA simulation\");\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/simpleVulkan/SineWaveSimulation.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"ShaderStructs.h\"\n\n__global__ void sinewave_gen_kernel(Vertex *vertices, unsigned int width, unsigned int height, float time)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // calculate uv coordinates\n float u = x / (float)width;\n float v = y / (float)height;\n u = u * 2.0f - 1.0f;\n v = v * 2.0f - 1.0f;\n\n // calculate simple sine wave pattern\n float freq = 4.0f;\n float w = sinf(u * freq + time) * cosf(v * freq + time) * 0.5f;\n\n if (y < height && x < width) {\n // write output vertex\n vertices[y * width + x].position.x = u;\n vertices[y * width + x].position.y = w;\n vertices[y * width + x].position.z = v;\n // vertices[y*width+x].position[3] = 1.0f;\n vertices[y * width + x].color.x = 1.0f;\n vertices[y * width + x].color.y = 0.0f;\n vertices[y * width + x].color.z = 0.0f;\n vertices[y * width + x].color.w = 0.0f;\n }\n}\n\n// The host CPU Sinewave thread spawner\nvoid RunSineWaveKernel(size_t mesh_width,\n size_t mesh_height,\n Vertex *cudaDevVertptr,\n cudaStream_t streamToRun,\n float AnimTime)\n{\n dim3 block(16, 16, 1);\n dim3 grid(mesh_width / 16, mesh_height / 16, 1);\n Vertex *vertices = (Vertex *)cudaDevVertptr;\n sinewave_gen_kernel<<>>(vertices, mesh_width, mesh_height, AnimTime);\n\n getLastCudaError(\"sinewave_gen_kernel execution failed.\\n\");\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \"ShaderStructs.h\"\n\n__global__ void sinewave_gen_kernel(Vertex *vertices, unsigned int width, unsigned int height, float time)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // calculate uv coordinates\n float u = x / (float)width;\n float v = y / (float)height;\n u = u * 2.0f - 1.0f;\n v = v * 2.0f - 1.0f;\n\n // calculate simple sine wave pattern\n float freq = 4.0f;\n float w = sinf(u * freq + time) * cosf(v * freq + time) * 0.5f;\n\n if (y < height && x < width) {\n // write output vertex\n vertices[y * width + x].position.x = u;\n vertices[y * width + x].position.y = w;\n vertices[y * width + x].position.z = v;\n // vertices[y*width+x].position[3] = 1.0f;\n vertices[y * width + x].color.x = 1.0f;\n vertices[y * width + x].color.y = 0.0f;\n vertices[y * width + x].color.z = 0.0f;\n vertices[y * width + x].color.w = 0.0f;\n }\n}\n\n// The host CPU Sinewave thread spawner\nvoid RunSineWaveKernel(size_t mesh_width,\n size_t mesh_height,\n Vertex *cudaDevVertptr,\n hipStream_t streamToRun,\n float AnimTime)\n{\n dim3 block(16, 16, 1);\n dim3 grid(mesh_width / 16, mesh_height / 16, 1);\n Vertex *vertices = (Vertex *)cudaDevVertptr;\n sinewave_gen_kernel<<>>(vertices, mesh_width, mesh_height, AnimTime);\n\n getLastCudaError(\"sinewave_gen_kernel execution failed.\\n\");\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/simpleD3D12/sinewave_cuda.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * Portions Copyright (c) 2009 Mike Giles, Oxford University. All rights\n * reserved.\n * Portions Copyright (c) 2008 Frances Y. Kuo and Stephen Joe. All rights\n * reserved.\n *\n * Sobol Quasi-random Number Generator example\n *\n * Based on CUDA code submitted by Mike Giles, Oxford University, United Kingdom\n * http://people.maths.ox.ac.uk/~gilesm/\n *\n * and C code developed by Stephen Joe, University of Waikato, New Zealand\n * and Frances Kuo, University of New South Wales, Australia\n * http://web.maths.unsw.edu.au/~fkuo/sobol/\n *\n * For theoretical background see:\n *\n * P. Bratley and B.L. Fox.\n * Implementing Sobol's quasirandom sequence generator\n * http://portal.acm.org/citation.cfm?id=42288\n * ACM Trans. on Math. Software, 14(1):88-100, 1988\n *\n * S. Joe and F. Kuo.\n * Remark on algorithm 659: implementing Sobol's quasirandom sequence generator.\n * http://portal.acm.org/citation.cfm?id=641879\n * ACM Trans. on Math. Software, 29(1):49-57, 2003\n *\n */\n\n#include \n\n#include \"sobol.h\"\n#include \"sobol_gpu.h\"\n\nnamespace cg = cooperative_groups;\n#include \n\n#define k_2powneg32 2.3283064E-10F\n\n__global__ void sobolGPU_kernel(unsigned n_vectors, unsigned n_dimensions, unsigned *d_directions, float *d_output)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ unsigned int v[n_directions];\n\n // Offset into the correct dimension as specified by the\n // block y coordinate\n d_directions = d_directions + n_directions * blockIdx.y;\n d_output = d_output + n_vectors * blockIdx.y;\n\n // Copy the direction numbers for this dimension into shared\n // memory - there are only 32 direction numbers so only the\n // first 32 (n_directions) threads need participate.\n if (threadIdx.x < n_directions) {\n v[threadIdx.x] = d_directions[threadIdx.x];\n }\n\n cg::sync(cta);\n\n // Set initial index (i.e. which vector this thread is\n // computing first) and stride (i.e. step to the next vector\n // for this thread)\n int i0 = threadIdx.x + blockIdx.x * blockDim.x;\n int stride = gridDim.x * blockDim.x;\n\n // Get the gray code of the index\n // c.f. Numerical Recipes in C, chapter 20\n // http://www.nrbook.com/a/bookcpdf/c20-2.pdf\n unsigned int g = i0 ^ (i0 >> 1);\n\n // Initialisation for first point x[i0]\n // In the Bratley and Fox paper this is equation (*), where\n // we are computing the value for x[n] without knowing the\n // value of x[n-1].\n unsigned int X = 0;\n unsigned int mask;\n\n for (unsigned int k = 0; k < __ffs(stride) - 1; k++) {\n // We want X ^= g_k * v[k], where g_k is one or zero.\n // We do this by setting a mask with all bits equal to\n // g_k. In reality we keep shifting g so that g_k is the\n // LSB of g. This way we avoid multiplication.\n mask = -(g & 1);\n X ^= mask & v[k];\n g = g >> 1;\n }\n\n if (i0 < n_vectors) {\n d_output[i0] = (float)X * k_2powneg32;\n }\n\n // Now do rest of points, using the stride\n // Here we want to generate x[i] from x[i-stride] where we\n // don't have any of the x in between, therefore we have to\n // revisit the equation (**), this is easiest with an example\n // so assume stride is 16.\n // From x[n] to x[n+16] there will be:\n // 8 changes in the first bit\n // 4 changes in the second bit\n // 2 changes in the third bit\n // 1 change in the fourth\n // 1 change in one of the remaining bits\n //\n // What this means is that in the equation:\n // x[n+1] = x[n] ^ v[p]\n // x[n+2] = x[n+1] ^ v[q] = x[n] ^ v[p] ^ v[q]\n // ...\n // We will apply xor with v[1] eight times, v[2] four times,\n // v[3] twice, v[4] once and one other direction number once.\n // Since two xors cancel out, we can skip even applications\n // and just apply xor with v[4] (i.e. log2(16)) and with\n // the current applicable direction number.\n // Note that all these indices count from 1, so we need to\n // subtract 1 from them all to account for C arrays counting\n // from zero.\n unsigned int v_log2stridem1 = v[__ffs(stride) - 2];\n unsigned int v_stridemask = stride - 1;\n\n for (unsigned int i = i0 + stride; i < n_vectors; i += stride) {\n // x[i] = x[i-stride] ^ v[b] ^ v[c]\n // where b is log2(stride) minus 1 for C array indexing\n // where c is the index of the rightmost zero bit in i,\n // not including the bottom log2(stride) bits, minus 1\n // for C array indexing\n // In the Bratley and Fox paper this is equation (**)\n X ^= v_log2stridem1 ^ v[__ffs(~((i - stride) | v_stridemask)) - 1];\n d_output[i] = (float)X * k_2powneg32;\n }\n}\n\nextern \"C\" void sobolGPU(int n_vectors, int n_dimensions, unsigned int *d_directions, float *d_output)\n{\n const int threadsperblock = 64;\n\n // Set up the execution configuration\n dim3 dimGrid;\n dim3 dimBlock;\n\n int device;\n cudaDeviceProp prop;\n checkCudaErrors(cudaGetDevice(&device));\n checkCudaErrors(cudaGetDeviceProperties(&prop, device));\n\n // This implementation of the generator outputs all the draws for\n // one dimension in a contiguous region of memory, followed by the\n // next dimension and so on.\n // Therefore all threads within a block will be processing different\n // vectors from the same dimension. As a result we want the total\n // number of blocks to be a multiple of the number of dimensions.\n dimGrid.y = n_dimensions;\n\n // If the number of dimensions is large then we will set the number\n // of blocks to equal the number of dimensions (i.e. dimGrid.x = 1)\n // but if the number of dimensions is small (e.g. less than four per\n // multiprocessor) then we'll partition the vectors across blocks\n // (as well as threads).\n if (n_dimensions < (4 * prop.multiProcessorCount)) {\n dimGrid.x = 4 * prop.multiProcessorCount;\n }\n else {\n dimGrid.x = 1;\n }\n\n // Cap the dimGrid.x if the number of vectors is small\n if (dimGrid.x > (unsigned int)(n_vectors / threadsperblock)) {\n dimGrid.x = (n_vectors + threadsperblock - 1) / threadsperblock;\n }\n\n // Round up to a power of two, required for the algorithm so that\n // stride is a power of two.\n unsigned int targetDimGridX = dimGrid.x;\n\n for (dimGrid.x = 1; dimGrid.x < targetDimGridX; dimGrid.x *= 2)\n ;\n\n // Fix the number of threads\n dimBlock.x = threadsperblock;\n\n // Execute GPU kernel\n sobolGPU_kernel<<>>(n_vectors, n_dimensions, d_directions, d_output);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * Portions Copyright (c) 2009 Mike Giles, Oxford University. All rights\n * reserved.\n * Portions Copyright (c) 2008 Frances Y. Kuo and Stephen Joe. All rights\n * reserved.\n *\n * Sobol Quasi-random Number Generator example\n *\n * Based on CUDA code submitted by Mike Giles, Oxford University, United Kingdom\n * http://people.maths.ox.ac.uk/~gilesm/\n *\n * and C code developed by Stephen Joe, University of Waikato, New Zealand\n * and Frances Kuo, University of New South Wales, Australia\n * http://web.maths.unsw.edu.au/~fkuo/sobol/\n *\n * For theoretical background see:\n *\n * P. Bratley and B.L. Fox.\n * Implementing Sobol's quasirandom sequence generator\n * http://portal.acm.org/citation.cfm?id=42288\n * ACM Trans. on Math. Software, 14(1):88-100, 1988\n *\n * S. Joe and F. Kuo.\n * Remark on algorithm 659: implementing Sobol's quasirandom sequence generator.\n * http://portal.acm.org/citation.cfm?id=641879\n * ACM Trans. on Math. Software, 29(1):49-57, 2003\n *\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \"sobol.h\"\n#include \"sobol_gpu.h\"\n\nnamespace cg = cooperative_groups;\n#include \n\n#define k_2powneg32 2.3283064E-10F\n\n__global__ void sobolGPU_kernel(unsigned n_vectors, unsigned n_dimensions, unsigned *d_directions, float *d_output)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ unsigned int v[n_directions];\n\n // Offset into the correct dimension as specified by the\n // block y coordinate\n d_directions = d_directions + n_directions * blockIdx.y;\n d_output = d_output + n_vectors * blockIdx.y;\n\n // Copy the direction numbers for this dimension into shared\n // memory - there are only 32 direction numbers so only the\n // first 32 (n_directions) threads need participate.\n if (threadIdx.x < n_directions) {\n v[threadIdx.x] = d_directions[threadIdx.x];\n }\n\n cg::sync(cta);\n\n // Set initial index (i.e. which vector this thread is\n // computing first) and stride (i.e. step to the next vector\n // for this thread)\n int i0 = threadIdx.x + blockIdx.x * blockDim.x;\n int stride = gridDim.x * blockDim.x;\n\n // Get the gray code of the index\n // c.f. Numerical Recipes in C, chapter 20\n // http://www.nrbook.com/a/bookcpdf/c20-2.pdf\n unsigned int g = i0 ^ (i0 >> 1);\n\n // Initialisation for first point x[i0]\n // In the Bratley and Fox paper this is equation (*), where\n // we are computing the value for x[n] without knowing the\n // value of x[n-1].\n unsigned int X = 0;\n unsigned int mask;\n\n for (unsigned int k = 0; k < __ffs(stride) - 1; k++) {\n // We want X ^= g_k * v[k], where g_k is one or zero.\n // We do this by setting a mask with all bits equal to\n // g_k. In reality we keep shifting g so that g_k is the\n // LSB of g. This way we avoid multiplication.\n mask = -(g & 1);\n X ^= mask & v[k];\n g = g >> 1;\n }\n\n if (i0 < n_vectors) {\n d_output[i0] = (float)X * k_2powneg32;\n }\n\n // Now do rest of points, using the stride\n // Here we want to generate x[i] from x[i-stride] where we\n // don't have any of the x in between, therefore we have to\n // revisit the equation (**), this is easiest with an example\n // so assume stride is 16.\n // From x[n] to x[n+16] there will be:\n // 8 changes in the first bit\n // 4 changes in the second bit\n // 2 changes in the third bit\n // 1 change in the fourth\n // 1 change in one of the remaining bits\n //\n // What this means is that in the equation:\n // x[n+1] = x[n] ^ v[p]\n // x[n+2] = x[n+1] ^ v[q] = x[n] ^ v[p] ^ v[q]\n // ...\n // We will apply xor with v[1] eight times, v[2] four times,\n // v[3] twice, v[4] once and one other direction number once.\n // Since two xors cancel out, we can skip even applications\n // and just apply xor with v[4] (i.e. log2(16)) and with\n // the current applicable direction number.\n // Note that all these indices count from 1, so we need to\n // subtract 1 from them all to account for C arrays counting\n // from zero.\n unsigned int v_log2stridem1 = v[__ffs(stride) - 2];\n unsigned int v_stridemask = stride - 1;\n\n for (unsigned int i = i0 + stride; i < n_vectors; i += stride) {\n // x[i] = x[i-stride] ^ v[b] ^ v[c]\n // where b is log2(stride) minus 1 for C array indexing\n // where c is the index of the rightmost zero bit in i,\n // not including the bottom log2(stride) bits, minus 1\n // for C array indexing\n // In the Bratley and Fox paper this is equation (**)\n X ^= v_log2stridem1 ^ v[__ffs(~((i - stride) | v_stridemask)) - 1];\n d_output[i] = (float)X * k_2powneg32;\n }\n}\n\nextern \"C\" void sobolGPU(int n_vectors, int n_dimensions, unsigned int *d_directions, float *d_output)\n{\n const int threadsperblock = 64;\n\n // Set up the execution configuration\n dim3 dimGrid;\n dim3 dimBlock;\n\n int device;\n hipDeviceProp_t prop;\n checkCudaErrors(hipGetDevice(&device));\n checkCudaErrors(hipGetDeviceProperties(&prop, device));\n\n // This implementation of the generator outputs all the draws for\n // one dimension in a contiguous region of memory, followed by the\n // next dimension and so on.\n // Therefore all threads within a block will be processing different\n // vectors from the same dimension. As a result we want the total\n // number of blocks to be a multiple of the number of dimensions.\n dimGrid.y = n_dimensions;\n\n // If the number of dimensions is large then we will set the number\n // of blocks to equal the number of dimensions (i.e. dimGrid.x = 1)\n // but if the number of dimensions is small (e.g. less than four per\n // multiprocessor) then we'll partition the vectors across blocks\n // (as well as threads).\n if (n_dimensions < (4 * prop.multiProcessorCount)) {\n dimGrid.x = 4 * prop.multiProcessorCount;\n }\n else {\n dimGrid.x = 1;\n }\n\n // Cap the dimGrid.x if the number of vectors is small\n if (dimGrid.x > (unsigned int)(n_vectors / threadsperblock)) {\n dimGrid.x = (n_vectors + threadsperblock - 1) / threadsperblock;\n }\n\n // Round up to a power of two, required for the algorithm so that\n // stride is a power of two.\n unsigned int targetDimGridX = dimGrid.x;\n\n for (dimGrid.x = 1; dimGrid.x < targetDimGridX; dimGrid.x *= 2)\n ;\n\n // Fix the number of threads\n dimBlock.x = threadsperblock;\n\n // Execute GPU kernel\n sobolGPU_kernel<<>>(n_vectors, n_dimensions, d_directions, d_output);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/SobolQRNG/sobol_gpu.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#ifndef _MARCHING_CUBES_KERNEL_CU_\n#define _MARCHING_CUBES_KERNEL_CU_\n\n#include \n#include // includes for helper CUDA functions\n#include \n#include \n#include \n#include \n#include \n\n#include \"defines.h\"\n#include \"tables.h\"\n\n// textures containing look-up tables\ncudaTextureObject_t triTex;\ncudaTextureObject_t numVertsTex;\n\n// volume data\ncudaTextureObject_t volumeTex;\n\nextern \"C\" void allocateTextures(uint **d_edgeTable, uint **d_triTable, uint **d_numVertsTable)\n{\n checkCudaErrors(cudaMalloc((void **)d_edgeTable, 256 * sizeof(uint)));\n checkCudaErrors(cudaMemcpy((void *)*d_edgeTable, (void *)edgeTable, 256 * sizeof(uint), cudaMemcpyHostToDevice));\n cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc(32, 0, 0, 0, cudaChannelFormatKindUnsigned);\n\n checkCudaErrors(cudaMalloc((void **)d_triTable, 256 * 16 * sizeof(uint)));\n checkCudaErrors(cudaMemcpy((void *)*d_triTable, (void *)triTable, 256 * 16 * sizeof(uint), cudaMemcpyHostToDevice));\n\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeLinear;\n texRes.res.linear.devPtr = *d_triTable;\n texRes.res.linear.sizeInBytes = 256 * 16 * sizeof(uint);\n texRes.res.linear.desc = channelDesc;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = cudaFilterModePoint;\n texDescr.addressMode[0] = cudaAddressModeClamp;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&triTex, &texRes, &texDescr, NULL));\n\n checkCudaErrors(cudaMalloc((void **)d_numVertsTable, 256 * sizeof(uint)));\n checkCudaErrors(\n cudaMemcpy((void *)*d_numVertsTable, (void *)numVertsTable, 256 * sizeof(uint), cudaMemcpyHostToDevice));\n\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeLinear;\n texRes.res.linear.devPtr = *d_numVertsTable;\n texRes.res.linear.sizeInBytes = 256 * sizeof(uint);\n texRes.res.linear.desc = channelDesc;\n\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = cudaFilterModePoint;\n texDescr.addressMode[0] = cudaAddressModeClamp;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&numVertsTex, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void createVolumeTexture(uchar *d_volume, size_t buffSize)\n{\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeLinear;\n texRes.res.linear.devPtr = d_volume;\n texRes.res.linear.sizeInBytes = buffSize;\n texRes.res.linear.desc = cudaCreateChannelDesc(8, 0, 0, 0, cudaChannelFormatKindUnsigned);\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = cudaFilterModePoint;\n texDescr.addressMode[0] = cudaAddressModeClamp;\n texDescr.readMode = cudaReadModeNormalizedFloat;\n\n checkCudaErrors(cudaCreateTextureObject(&volumeTex, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void destroyAllTextureObjects()\n{\n checkCudaErrors(cudaDestroyTextureObject(triTex));\n checkCudaErrors(cudaDestroyTextureObject(numVertsTex));\n checkCudaErrors(cudaDestroyTextureObject(volumeTex));\n}\n\n// an interesting field function\n__device__ float tangle(float x, float y, float z)\n{\n x *= 3.0f;\n y *= 3.0f;\n z *= 3.0f;\n return (x * x * x * x - 5.0f * x * x + y * y * y * y - 5.0f * y * y + z * z * z * z - 5.0f * z * z + 11.8f) * 0.2f\n + 0.5f;\n}\n\n// evaluate field function at point\n__device__ float fieldFunc(float3 p) { return tangle(p.x, p.y, p.z); }\n\n// evaluate field function at a point\n// returns value and gradient in float4\n__device__ float4 fieldFunc4(float3 p)\n{\n float v = tangle(p.x, p.y, p.z);\n const float d = 0.001f;\n float dx = tangle(p.x + d, p.y, p.z) - v;\n float dy = tangle(p.x, p.y + d, p.z) - v;\n float dz = tangle(p.x, p.y, p.z + d) - v;\n return make_float4(dx, dy, dz, v);\n}\n\n// sample volume data set at a point\n__device__ float sampleVolume(cudaTextureObject_t volumeTex, uchar *data, uint3 p, uint3 gridSize)\n{\n p.x = min(p.x, gridSize.x - 1);\n p.y = min(p.y, gridSize.y - 1);\n p.z = min(p.z, gridSize.z - 1);\n uint i = (p.z * gridSize.x * gridSize.y) + (p.y * gridSize.x) + p.x;\n // return (float) data[i] / 255.0f;\n return tex1Dfetch(volumeTex, i);\n}\n\n// compute position in 3d grid from 1d index\n// only works for power of 2 sizes\n__device__ uint3 calcGridPos(uint i, uint3 gridSizeShift, uint3 gridSizeMask)\n{\n uint3 gridPos;\n gridPos.x = i & gridSizeMask.x;\n gridPos.y = (i >> gridSizeShift.y) & gridSizeMask.y;\n gridPos.z = (i >> gridSizeShift.z) & gridSizeMask.z;\n return gridPos;\n}\n\n// classify voxel based on number of vertices it will generate\n// one thread per voxel\n__global__ void classifyVoxel(uint *voxelVerts,\n uint *voxelOccupied,\n uchar *volume,\n uint3 gridSize,\n uint3 gridSizeShift,\n uint3 gridSizeMask,\n uint numVoxels,\n float3 voxelSize,\n float isoValue,\n cudaTextureObject_t numVertsTex,\n cudaTextureObject_t volumeTex)\n{\n uint blockId = __mul24(blockIdx.y, gridDim.x) + blockIdx.x;\n uint i = __mul24(blockId, blockDim.x) + threadIdx.x;\n\n uint3 gridPos = calcGridPos(i, gridSizeShift, gridSizeMask);\n\n// read field values at neighbouring grid vertices\n#if SAMPLE_VOLUME\n float field[8];\n field[0] = sampleVolume(volumeTex, volume, gridPos, gridSize);\n field[1] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 0, 0), gridSize);\n field[2] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 1, 0), gridSize);\n field[3] = sampleVolume(volumeTex, volume, gridPos + make_uint3(0, 1, 0), gridSize);\n field[4] = sampleVolume(volumeTex, volume, gridPos + make_uint3(0, 0, 1), gridSize);\n field[5] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 0, 1), gridSize);\n field[6] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 1, 1), gridSize);\n field[7] = sampleVolume(volumeTex, volume, gridPos + make_uint3(0, 1, 1), gridSize);\n#else\n float3 p;\n p.x = -1.0f + (gridPos.x * voxelSize.x);\n p.y = -1.0f + (gridPos.y * voxelSize.y);\n p.z = -1.0f + (gridPos.z * voxelSize.z);\n\n float field[8];\n field[0] = fieldFunc(p);\n field[1] = fieldFunc(p + make_float3(voxelSize.x, 0, 0));\n field[2] = fieldFunc(p + make_float3(voxelSize.x, voxelSize.y, 0));\n field[3] = fieldFunc(p + make_float3(0, voxelSize.y, 0));\n field[4] = fieldFunc(p + make_float3(0, 0, voxelSize.z));\n field[5] = fieldFunc(p + make_float3(voxelSize.x, 0, voxelSize.z));\n field[6] = fieldFunc(p + make_float3(voxelSize.x, voxelSize.y, voxelSize.z));\n field[7] = fieldFunc(p + make_float3(0, voxelSize.y, voxelSize.z));\n#endif\n\n // calculate flag indicating if each vertex is inside or outside isosurface\n uint cubeindex;\n cubeindex = uint(field[0] < isoValue);\n cubeindex += uint(field[1] < isoValue) * 2;\n cubeindex += uint(field[2] < isoValue) * 4;\n cubeindex += uint(field[3] < isoValue) * 8;\n cubeindex += uint(field[4] < isoValue) * 16;\n cubeindex += uint(field[5] < isoValue) * 32;\n cubeindex += uint(field[6] < isoValue) * 64;\n cubeindex += uint(field[7] < isoValue) * 128;\n\n // read number of vertices from texture\n uint numVerts = tex1Dfetch(numVertsTex, cubeindex);\n\n if (i < numVoxels) {\n voxelVerts[i] = numVerts;\n voxelOccupied[i] = (numVerts > 0);\n }\n}\n\nextern \"C\" void launch_classifyVoxel(dim3 grid,\n dim3 threads,\n uint *voxelVerts,\n uint *voxelOccupied,\n uchar *volume,\n uint3 gridSize,\n uint3 gridSizeShift,\n uint3 gridSizeMask,\n uint numVoxels,\n float3 voxelSize,\n float isoValue)\n{\n // calculate number of vertices need per voxel\n classifyVoxel<<>>(voxelVerts,\n voxelOccupied,\n volume,\n gridSize,\n gridSizeShift,\n gridSizeMask,\n numVoxels,\n voxelSize,\n isoValue,\n numVertsTex,\n volumeTex);\n getLastCudaError(\"classifyVoxel failed\");\n}\n\n// compact voxel array\n__global__ void compactVoxels(uint *compactedVoxelArray, uint *voxelOccupied, uint *voxelOccupiedScan, uint numVoxels)\n{\n uint blockId = __mul24(blockIdx.y, gridDim.x) + blockIdx.x;\n uint i = __mul24(blockId, blockDim.x) + threadIdx.x;\n\n if (voxelOccupied[i] && (i < numVoxels)) {\n compactedVoxelArray[voxelOccupiedScan[i]] = i;\n }\n}\n\nextern \"C\" void launch_compactVoxels(dim3 grid,\n dim3 threads,\n uint *compactedVoxelArray,\n uint *voxelOccupied,\n uint *voxelOccupiedScan,\n uint numVoxels)\n{\n compactVoxels<<>>(compactedVoxelArray, voxelOccupied, voxelOccupiedScan, numVoxels);\n getLastCudaError(\"compactVoxels failed\");\n}\n\n// compute interpolated vertex along an edge\n__device__ float3 vertexInterp(float isolevel, float3 p0, float3 p1, float f0, float f1)\n{\n float t = (isolevel - f0) / (f1 - f0);\n return lerp(p0, p1, t);\n}\n\n// compute interpolated vertex position and normal along an edge\n__device__ void vertexInterp2(float isolevel, float3 p0, float3 p1, float4 f0, float4 f1, float3 &p, float3 &n)\n{\n float t = (isolevel - f0.w) / (f1.w - f0.w);\n p = lerp(p0, p1, t);\n n.x = lerp(f0.x, f1.x, t);\n n.y = lerp(f0.y, f1.y, t);\n n.z = lerp(f0.z, f1.z, t);\n // n = normalize(n);\n}\n\n// generate triangles for each voxel using marching cubes\n// interpolates normals from field function\n__global__ void generateTriangles(float4 *pos,\n float4 *norm,\n uint *compactedVoxelArray,\n uint *numVertsScanned,\n uint3 gridSize,\n uint3 gridSizeShift,\n uint3 gridSizeMask,\n float3 voxelSize,\n float isoValue,\n uint activeVoxels,\n uint maxVerts,\n cudaTextureObject_t triTex,\n cudaTextureObject_t numVertsTex)\n{\n uint blockId = __mul24(blockIdx.y, gridDim.x) + blockIdx.x;\n uint i = __mul24(blockId, blockDim.x) + threadIdx.x;\n\n if (i > activeVoxels - 1) {\n // can't return here because of syncthreads()\n i = activeVoxels - 1;\n }\n\n#if SKIP_EMPTY_VOXELS\n uint voxel = compactedVoxelArray[i];\n#else\n uint voxel = i;\n#endif\n\n // compute position in 3d grid\n uint3 gridPos = calcGridPos(voxel, gridSizeShift, gridSizeMask);\n\n float3 p;\n p.x = -1.0f + (gridPos.x * voxelSize.x);\n p.y = -1.0f + (gridPos.y * voxelSize.y);\n p.z = -1.0f + (gridPos.z * voxelSize.z);\n\n // calculate cell vertex positions\n float3 v[8];\n v[0] = p;\n v[1] = p + make_float3(voxelSize.x, 0, 0);\n v[2] = p + make_float3(voxelSize.x, voxelSize.y, 0);\n v[3] = p + make_float3(0, voxelSize.y, 0);\n v[4] = p + make_float3(0, 0, voxelSize.z);\n v[5] = p + make_float3(voxelSize.x, 0, voxelSize.z);\n v[6] = p + make_float3(voxelSize.x, voxelSize.y, voxelSize.z);\n v[7] = p + make_float3(0, voxelSize.y, voxelSize.z);\n\n // evaluate field values\n float4 field[8];\n field[0] = fieldFunc4(v[0]);\n field[1] = fieldFunc4(v[1]);\n field[2] = fieldFunc4(v[2]);\n field[3] = fieldFunc4(v[3]);\n field[4] = fieldFunc4(v[4]);\n field[5] = fieldFunc4(v[5]);\n field[6] = fieldFunc4(v[6]);\n field[7] = fieldFunc4(v[7]);\n\n // recalculate flag\n // (this is faster than storing it in global memory)\n uint cubeindex;\n cubeindex = uint(field[0].w < isoValue);\n cubeindex += uint(field[1].w < isoValue) * 2;\n cubeindex += uint(field[2].w < isoValue) * 4;\n cubeindex += uint(field[3].w < isoValue) * 8;\n cubeindex += uint(field[4].w < isoValue) * 16;\n cubeindex += uint(field[5].w < isoValue) * 32;\n cubeindex += uint(field[6].w < isoValue) * 64;\n cubeindex += uint(field[7].w < isoValue) * 128;\n\n // find the vertices where the surface intersects the cube\n\n#if USE_SHARED\n // use partioned shared memory to avoid using local memory\n __shared__ float3 vertlist[12 * NTHREADS];\n __shared__ float3 normlist[12 * NTHREADS];\n\n vertexInterp2(isoValue, v[0], v[1], field[0], field[1], vertlist[threadIdx.x], normlist[threadIdx.x]);\n vertexInterp2(\n isoValue, v[1], v[2], field[1], field[2], vertlist[threadIdx.x + NTHREADS], normlist[threadIdx.x + NTHREADS]);\n vertexInterp2(isoValue,\n v[2],\n v[3],\n field[2],\n field[3],\n vertlist[threadIdx.x + (NTHREADS * 2)],\n normlist[threadIdx.x + (NTHREADS * 2)]);\n vertexInterp2(isoValue,\n v[3],\n v[0],\n field[3],\n field[0],\n vertlist[threadIdx.x + (NTHREADS * 3)],\n normlist[threadIdx.x + (NTHREADS * 3)]);\n vertexInterp2(isoValue,\n v[4],\n v[5],\n field[4],\n field[5],\n vertlist[threadIdx.x + (NTHREADS * 4)],\n normlist[threadIdx.x + (NTHREADS * 4)]);\n vertexInterp2(isoValue,\n v[5],\n v[6],\n field[5],\n field[6],\n vertlist[threadIdx.x + (NTHREADS * 5)],\n normlist[threadIdx.x + (NTHREADS * 5)]);\n vertexInterp2(isoValue,\n v[6],\n v[7],\n field[6],\n field[7],\n vertlist[threadIdx.x + (NTHREADS * 6)],\n normlist[threadIdx.x + (NTHREADS * 6)]);\n vertexInterp2(isoValue,\n v[7],\n v[4],\n field[7],\n field[4],\n vertlist[threadIdx.x + (NTHREADS * 7)],\n normlist[threadIdx.x + (NTHREADS * 7)]);\n vertexInterp2(isoValue,\n v[0],\n v[4],\n field[0],\n field[4],\n vertlist[threadIdx.x + (NTHREADS * 8)],\n normlist[threadIdx.x + (NTHREADS * 8)]);\n vertexInterp2(isoValue,\n v[1],\n v[5],\n field[1],\n field[5],\n vertlist[threadIdx.x + (NTHREADS * 9)],\n normlist[threadIdx.x + (NTHREADS * 9)]);\n vertexInterp2(isoValue,\n v[2],\n v[6],\n field[2],\n field[6],\n vertlist[threadIdx.x + (NTHREADS * 10)],\n normlist[threadIdx.x + (NTHREADS * 10)]);\n vertexInterp2(isoValue,\n v[3],\n v[7],\n field[3],\n field[7],\n vertlist[threadIdx.x + (NTHREADS * 11)],\n normlist[threadIdx.x + (NTHREADS * 11)]);\n __syncthreads();\n\n#else\n float3 vertlist[12];\n float3 normlist[12];\n\n vertexInterp2(isoValue, v[0], v[1], field[0], field[1], vertlist[0], normlist[0]);\n vertexInterp2(isoValue, v[1], v[2], field[1], field[2], vertlist[1], normlist[1]);\n vertexInterp2(isoValue, v[2], v[3], field[2], field[3], vertlist[2], normlist[2]);\n vertexInterp2(isoValue, v[3], v[0], field[3], field[0], vertlist[3], normlist[3]);\n\n vertexInterp2(isoValue, v[4], v[5], field[4], field[5], vertlist[4], normlist[4]);\n vertexInterp2(isoValue, v[5], v[6], field[5], field[6], vertlist[5], normlist[5]);\n vertexInterp2(isoValue, v[6], v[7], field[6], field[7], vertlist[6], normlist[6]);\n vertexInterp2(isoValue, v[7], v[4], field[7], field[4], vertlist[7], normlist[7]);\n\n vertexInterp2(isoValue, v[0], v[4], field[0], field[4], vertlist[8], normlist[8]);\n vertexInterp2(isoValue, v[1], v[5], field[1], field[5], vertlist[9], normlist[9]);\n vertexInterp2(isoValue, v[2], v[6], field[2], field[6], vertlist[10], normlist[10]);\n vertexInterp2(isoValue, v[3], v[7], field[3], field[7], vertlist[11], normlist[11]);\n#endif\n\n // output triangle vertices\n uint numVerts = tex1Dfetch(numVertsTex, cubeindex);\n\n for (int i = 0; i < numVerts; i++) {\n uint edge = tex1Dfetch(triTex, cubeindex * 16 + i);\n\n uint index = numVertsScanned[voxel] + i;\n\n if (index < maxVerts) {\n#if USE_SHARED\n pos[index] = make_float4(vertlist[(edge * NTHREADS) + threadIdx.x], 1.0f);\n norm[index] = make_float4(normlist[(edge * NTHREADS) + threadIdx.x], 0.0f);\n#else\n pos[index] = make_float4(vertlist[edge], 1.0f);\n norm[index] = make_float4(normlist[edge], 0.0f);\n#endif\n }\n }\n}\n\nextern \"C\" void launch_generateTriangles(dim3 grid,\n dim3 threads,\n float4 *pos,\n float4 *norm,\n uint *compactedVoxelArray,\n uint *numVertsScanned,\n uint3 gridSize,\n uint3 gridSizeShift,\n uint3 gridSizeMask,\n float3 voxelSize,\n float isoValue,\n uint activeVoxels,\n uint maxVerts)\n{\n generateTriangles<<>>(pos,\n norm,\n compactedVoxelArray,\n numVertsScanned,\n gridSize,\n gridSizeShift,\n gridSizeMask,\n voxelSize,\n isoValue,\n activeVoxels,\n maxVerts,\n triTex,\n numVertsTex);\n getLastCudaError(\"generateTriangles failed\");\n}\n\n// calculate triangle normal\n__device__ float3 calcNormal(float3 *v0, float3 *v1, float3 *v2)\n{\n float3 edge0 = *v1 - *v0;\n float3 edge1 = *v2 - *v0;\n // note - it's faster to perform normalization in vertex shader rather than\n // here\n return cross(edge0, edge1);\n}\n\n// version that calculates flat surface normal for each triangle\n__global__ void generateTriangles2(float4 *pos,\n float4 *norm,\n uint *compactedVoxelArray,\n uint *numVertsScanned,\n uchar *volume,\n uint3 gridSize,\n uint3 gridSizeShift,\n uint3 gridSizeMask,\n float3 voxelSize,\n float isoValue,\n uint activeVoxels,\n uint maxVerts,\n cudaTextureObject_t triTex,\n cudaTextureObject_t numVertsTex,\n cudaTextureObject_t volumeTex)\n{\n uint blockId = __mul24(blockIdx.y, gridDim.x) + blockIdx.x;\n uint i = __mul24(blockId, blockDim.x) + threadIdx.x;\n\n if (i > activeVoxels - 1) {\n i = activeVoxels - 1;\n }\n\n#if SKIP_EMPTY_VOXELS\n uint voxel = compactedVoxelArray[i];\n#else\n uint voxel = i;\n#endif\n\n // compute position in 3d grid\n uint3 gridPos = calcGridPos(voxel, gridSizeShift, gridSizeMask);\n\n float3 p;\n p.x = -1.0f + (gridPos.x * voxelSize.x);\n p.y = -1.0f + (gridPos.y * voxelSize.y);\n p.z = -1.0f + (gridPos.z * voxelSize.z);\n\n // calculate cell vertex positions\n float3 v[8];\n v[0] = p;\n v[1] = p + make_float3(voxelSize.x, 0, 0);\n v[2] = p + make_float3(voxelSize.x, voxelSize.y, 0);\n v[3] = p + make_float3(0, voxelSize.y, 0);\n v[4] = p + make_float3(0, 0, voxelSize.z);\n v[5] = p + make_float3(voxelSize.x, 0, voxelSize.z);\n v[6] = p + make_float3(voxelSize.x, voxelSize.y, voxelSize.z);\n v[7] = p + make_float3(0, voxelSize.y, voxelSize.z);\n\n#if SAMPLE_VOLUME\n float field[8];\n field[0] = sampleVolume(volumeTex, volume, gridPos, gridSize);\n field[1] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 0, 0), gridSize);\n field[2] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 1, 0), gridSize);\n field[3] = sampleVolume(volumeTex, volume, gridPos + make_uint3(0, 1, 0), gridSize);\n field[4] = sampleVolume(volumeTex, volume, gridPos + make_uint3(0, 0, 1), gridSize);\n field[5] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 0, 1), gridSize);\n field[6] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 1, 1), gridSize);\n field[7] = sampleVolume(volumeTex, volume, gridPos + make_uint3(0, 1, 1), gridSize);\n#else\n // evaluate field values\n float field[8];\n field[0] = fieldFunc(v[0]);\n field[1] = fieldFunc(v[1]);\n field[2] = fieldFunc(v[2]);\n field[3] = fieldFunc(v[3]);\n field[4] = fieldFunc(v[4]);\n field[5] = fieldFunc(v[5]);\n field[6] = fieldFunc(v[6]);\n field[7] = fieldFunc(v[7]);\n#endif\n\n // recalculate flag\n uint cubeindex;\n cubeindex = uint(field[0] < isoValue);\n cubeindex += uint(field[1] < isoValue) * 2;\n cubeindex += uint(field[2] < isoValue) * 4;\n cubeindex += uint(field[3] < isoValue) * 8;\n cubeindex += uint(field[4] < isoValue) * 16;\n cubeindex += uint(field[5] < isoValue) * 32;\n cubeindex += uint(field[6] < isoValue) * 64;\n cubeindex += uint(field[7] < isoValue) * 128;\n\n // find the vertices where the surface intersects the cube\n\n#if USE_SHARED\n // use shared memory to avoid using local\n __shared__ float3 vertlist[12 * NTHREADS];\n\n vertlist[threadIdx.x] = vertexInterp(isoValue, v[0], v[1], field[0], field[1]);\n vertlist[NTHREADS + threadIdx.x] = vertexInterp(isoValue, v[1], v[2], field[1], field[2]);\n vertlist[(NTHREADS * 2) + threadIdx.x] = vertexInterp(isoValue, v[2], v[3], field[2], field[3]);\n vertlist[(NTHREADS * 3) + threadIdx.x] = vertexInterp(isoValue, v[3], v[0], field[3], field[0]);\n vertlist[(NTHREADS * 4) + threadIdx.x] = vertexInterp(isoValue, v[4], v[5], field[4], field[5]);\n vertlist[(NTHREADS * 5) + threadIdx.x] = vertexInterp(isoValue, v[5], v[6], field[5], field[6]);\n vertlist[(NTHREADS * 6) + threadIdx.x] = vertexInterp(isoValue, v[6], v[7], field[6], field[7]);\n vertlist[(NTHREADS * 7) + threadIdx.x] = vertexInterp(isoValue, v[7], v[4], field[7], field[4]);\n vertlist[(NTHREADS * 8) + threadIdx.x] = vertexInterp(isoValue, v[0], v[4], field[0], field[4]);\n vertlist[(NTHREADS * 9) + threadIdx.x] = vertexInterp(isoValue, v[1], v[5], field[1], field[5]);\n vertlist[(NTHREADS * 10) + threadIdx.x] = vertexInterp(isoValue, v[2], v[6], field[2], field[6]);\n vertlist[(NTHREADS * 11) + threadIdx.x] = vertexInterp(isoValue, v[3], v[7], field[3], field[7]);\n __syncthreads();\n#else\n\n float3 vertlist[12];\n\n vertlist[0] = vertexInterp(isoValue, v[0], v[1], field[0], field[1]);\n vertlist[1] = vertexInterp(isoValue, v[1], v[2], field[1], field[2]);\n vertlist[2] = vertexInterp(isoValue, v[2], v[3], field[2], field[3]);\n vertlist[3] = vertexInterp(isoValue, v[3], v[0], field[3], field[0]);\n\n vertlist[4] = vertexInterp(isoValue, v[4], v[5], field[4], field[5]);\n vertlist[5] = vertexInterp(isoValue, v[5], v[6], field[5], field[6]);\n vertlist[6] = vertexInterp(isoValue, v[6], v[7], field[6], field[7]);\n vertlist[7] = vertexInterp(isoValue, v[7], v[4], field[7], field[4]);\n\n vertlist[8] = vertexInterp(isoValue, v[0], v[4], field[0], field[4]);\n vertlist[9] = vertexInterp(isoValue, v[1], v[5], field[1], field[5]);\n vertlist[10] = vertexInterp(isoValue, v[2], v[6], field[2], field[6]);\n vertlist[11] = vertexInterp(isoValue, v[3], v[7], field[3], field[7]);\n#endif\n\n // output triangle vertices\n uint numVerts = tex1Dfetch(numVertsTex, cubeindex);\n\n for (int i = 0; i < numVerts; i += 3) {\n uint index = numVertsScanned[voxel] + i;\n\n float3 *v[3];\n uint edge;\n edge = tex1Dfetch(triTex, (cubeindex * 16) + i);\n#if USE_SHARED\n v[0] = &vertlist[(edge * NTHREADS) + threadIdx.x];\n#else\n v[0] = &vertlist[edge];\n#endif\n\n edge = tex1Dfetch(triTex, (cubeindex * 16) + i + 1);\n#if USE_SHARED\n v[1] = &vertlist[(edge * NTHREADS) + threadIdx.x];\n#else\n v[1] = &vertlist[edge];\n#endif\n\n edge = tex1Dfetch(triTex, (cubeindex * 16) + i + 2);\n#if USE_SHARED\n v[2] = &vertlist[(edge * NTHREADS) + threadIdx.x];\n#else\n v[2] = &vertlist[edge];\n#endif\n\n // calculate triangle surface normal\n float3 n = calcNormal(v[0], v[1], v[2]);\n\n if (index < (maxVerts - 3)) {\n pos[index] = make_float4(*v[0], 1.0f);\n norm[index] = make_float4(n, 0.0f);\n\n pos[index + 1] = make_float4(*v[1], 1.0f);\n norm[index + 1] = make_float4(n, 0.0f);\n\n pos[index + 2] = make_float4(*v[2], 1.0f);\n norm[index + 2] = make_float4(n, 0.0f);\n }\n }\n}\n\nextern \"C\" void launch_generateTriangles2(dim3 grid,\n dim3 threads,\n float4 *pos,\n float4 *norm,\n uint *compactedVoxelArray,\n uint *numVertsScanned,\n uchar *volume,\n uint3 gridSize,\n uint3 gridSizeShift,\n uint3 gridSizeMask,\n float3 voxelSize,\n float isoValue,\n uint activeVoxels,\n uint maxVerts)\n{\n generateTriangles2<<>>(pos,\n norm,\n compactedVoxelArray,\n numVertsScanned,\n volume,\n gridSize,\n gridSizeShift,\n gridSizeMask,\n voxelSize,\n isoValue,\n activeVoxels,\n maxVerts,\n triTex,\n numVertsTex,\n volumeTex);\n getLastCudaError(\"generateTriangles2 failed\");\n}\n\nextern \"C\" void ThrustScanWrapper(unsigned int *output, unsigned int *input, unsigned int numElements)\n{\n thrust::exclusive_scan(thrust::device_ptr(input),\n thrust::device_ptr(input + numElements),\n thrust::device_ptr(output));\n}\n\n#endif\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#ifndef _MARCHING_CUBES_KERNEL_CU_\n#define _MARCHING_CUBES_KERNEL_CU_\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"defines.h\"\n#include \"tables.h\"\n\n// textures containing look-up tables\nhipTextureObject_t triTex;\nhipTextureObject_t numVertsTex;\n\n// volume data\nhipTextureObject_t volumeTex;\n\nextern \"C\" void allocateTextures(uint **d_edgeTable, uint **d_triTable, uint **d_numVertsTable)\n{\n checkCudaErrors(hipMalloc((void **)d_edgeTable, 256 * sizeof(uint)));\n checkCudaErrors(hipMemcpy((void *)*d_edgeTable, (void *)edgeTable, 256 * sizeof(uint), hipMemcpyHostToDevice));\n hipChannelFormatDesc channelDesc = hipCreateChannelDesc(32, 0, 0, 0, hipChannelFormatKindUnsigned);\n\n checkCudaErrors(hipMalloc((void **)d_triTable, 256 * 16 * sizeof(uint)));\n checkCudaErrors(hipMemcpy((void *)*d_triTable, (void *)triTable, 256 * 16 * sizeof(uint), hipMemcpyHostToDevice));\n\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeLinear;\n texRes.res.linear.devPtr = *d_triTable;\n texRes.res.linear.sizeInBytes = 256 * 16 * sizeof(uint);\n texRes.res.linear.desc = channelDesc;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = hipFilterModePoint;\n texDescr.addressMode[0] = hipAddressModeClamp;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&triTex, &texRes, &texDescr, NULL));\n\n checkCudaErrors(hipMalloc((void **)d_numVertsTable, 256 * sizeof(uint)));\n checkCudaErrors(\n hipMemcpy((void *)*d_numVertsTable, (void *)numVertsTable, 256 * sizeof(uint), hipMemcpyHostToDevice));\n\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeLinear;\n texRes.res.linear.devPtr = *d_numVertsTable;\n texRes.res.linear.sizeInBytes = 256 * sizeof(uint);\n texRes.res.linear.desc = channelDesc;\n\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = hipFilterModePoint;\n texDescr.addressMode[0] = hipAddressModeClamp;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&numVertsTex, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void createVolumeTexture(uchar *d_volume, size_t buffSize)\n{\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeLinear;\n texRes.res.linear.devPtr = d_volume;\n texRes.res.linear.sizeInBytes = buffSize;\n texRes.res.linear.desc = hipCreateChannelDesc(8, 0, 0, 0, hipChannelFormatKindUnsigned);\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = hipFilterModePoint;\n texDescr.addressMode[0] = hipAddressModeClamp;\n texDescr.readMode = hipReadModeNormalizedFloat;\n\n checkCudaErrors(hipCreateTextureObject(&volumeTex, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void destroyAllTextureObjects()\n{\n checkCudaErrors(hipDestroyTextureObject(triTex));\n checkCudaErrors(hipDestroyTextureObject(numVertsTex));\n checkCudaErrors(hipDestroyTextureObject(volumeTex));\n}\n\n// an interesting field function\n__device__ float tangle(float x, float y, float z)\n{\n x *= 3.0f;\n y *= 3.0f;\n z *= 3.0f;\n return (x * x * x * x - 5.0f * x * x + y * y * y * y - 5.0f * y * y + z * z * z * z - 5.0f * z * z + 11.8f) * 0.2f\n + 0.5f;\n}\n\n// evaluate field function at point\n__device__ float fieldFunc(float3 p) { return tangle(p.x, p.y, p.z); }\n\n// evaluate field function at a point\n// returns value and gradient in float4\n__device__ float4 fieldFunc4(float3 p)\n{\n float v = tangle(p.x, p.y, p.z);\n const float d = 0.001f;\n float dx = tangle(p.x + d, p.y, p.z) - v;\n float dy = tangle(p.x, p.y + d, p.z) - v;\n float dz = tangle(p.x, p.y, p.z + d) - v;\n return make_float4(dx, dy, dz, v);\n}\n\n// sample volume data set at a point\n__device__ float sampleVolume(hipTextureObject_t volumeTex, uchar *data, uint3 p, uint3 gridSize)\n{\n p.x = min(p.x, gridSize.x - 1);\n p.y = min(p.y, gridSize.y - 1);\n p.z = min(p.z, gridSize.z - 1);\n uint i = (p.z * gridSize.x * gridSize.y) + (p.y * gridSize.x) + p.x;\n // return (float) data[i] / 255.0f;\n return tex1Dfetch(volumeTex, i);\n}\n\n// compute position in 3d grid from 1d index\n// only works for power of 2 sizes\n__device__ uint3 calcGridPos(uint i, uint3 gridSizeShift, uint3 gridSizeMask)\n{\n uint3 gridPos;\n gridPos.x = i & gridSizeMask.x;\n gridPos.y = (i >> gridSizeShift.y) & gridSizeMask.y;\n gridPos.z = (i >> gridSizeShift.z) & gridSizeMask.z;\n return gridPos;\n}\n\n// classify voxel based on number of vertices it will generate\n// one thread per voxel\n__global__ void classifyVoxel(uint *voxelVerts,\n uint *voxelOccupied,\n uchar *volume,\n uint3 gridSize,\n uint3 gridSizeShift,\n uint3 gridSizeMask,\n uint numVoxels,\n float3 voxelSize,\n float isoValue,\n hipTextureObject_t numVertsTex,\n hipTextureObject_t volumeTex)\n{\n uint blockId = __mul24(blockIdx.y, gridDim.x) + blockIdx.x;\n uint i = __mul24(blockId, blockDim.x) + threadIdx.x;\n\n uint3 gridPos = calcGridPos(i, gridSizeShift, gridSizeMask);\n\n// read field values at neighbouring grid vertices\n#if SAMPLE_VOLUME\n float field[8];\n field[0] = sampleVolume(volumeTex, volume, gridPos, gridSize);\n field[1] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 0, 0), gridSize);\n field[2] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 1, 0), gridSize);\n field[3] = sampleVolume(volumeTex, volume, gridPos + make_uint3(0, 1, 0), gridSize);\n field[4] = sampleVolume(volumeTex, volume, gridPos + make_uint3(0, 0, 1), gridSize);\n field[5] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 0, 1), gridSize);\n field[6] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 1, 1), gridSize);\n field[7] = sampleVolume(volumeTex, volume, gridPos + make_uint3(0, 1, 1), gridSize);\n#else\n float3 p;\n p.x = -1.0f + (gridPos.x * voxelSize.x);\n p.y = -1.0f + (gridPos.y * voxelSize.y);\n p.z = -1.0f + (gridPos.z * voxelSize.z);\n\n float field[8];\n field[0] = fieldFunc(p);\n field[1] = fieldFunc(p + make_float3(voxelSize.x, 0, 0));\n field[2] = fieldFunc(p + make_float3(voxelSize.x, voxelSize.y, 0));\n field[3] = fieldFunc(p + make_float3(0, voxelSize.y, 0));\n field[4] = fieldFunc(p + make_float3(0, 0, voxelSize.z));\n field[5] = fieldFunc(p + make_float3(voxelSize.x, 0, voxelSize.z));\n field[6] = fieldFunc(p + make_float3(voxelSize.x, voxelSize.y, voxelSize.z));\n field[7] = fieldFunc(p + make_float3(0, voxelSize.y, voxelSize.z));\n#endif\n\n // calculate flag indicating if each vertex is inside or outside isosurface\n uint cubeindex;\n cubeindex = uint(field[0] < isoValue);\n cubeindex += uint(field[1] < isoValue) * 2;\n cubeindex += uint(field[2] < isoValue) * 4;\n cubeindex += uint(field[3] < isoValue) * 8;\n cubeindex += uint(field[4] < isoValue) * 16;\n cubeindex += uint(field[5] < isoValue) * 32;\n cubeindex += uint(field[6] < isoValue) * 64;\n cubeindex += uint(field[7] < isoValue) * 128;\n\n // read number of vertices from texture\n uint numVerts = tex1Dfetch(numVertsTex, cubeindex);\n\n if (i < numVoxels) {\n voxelVerts[i] = numVerts;\n voxelOccupied[i] = (numVerts > 0);\n }\n}\n\nextern \"C\" void launch_classifyVoxel(dim3 grid,\n dim3 threads,\n uint *voxelVerts,\n uint *voxelOccupied,\n uchar *volume,\n uint3 gridSize,\n uint3 gridSizeShift,\n uint3 gridSizeMask,\n uint numVoxels,\n float3 voxelSize,\n float isoValue)\n{\n // calculate number of vertices need per voxel\n classifyVoxel<<>>(voxelVerts,\n voxelOccupied,\n volume,\n gridSize,\n gridSizeShift,\n gridSizeMask,\n numVoxels,\n voxelSize,\n isoValue,\n numVertsTex,\n volumeTex);\n getLastCudaError(\"classifyVoxel failed\");\n}\n\n// compact voxel array\n__global__ void compactVoxels(uint *compactedVoxelArray, uint *voxelOccupied, uint *voxelOccupiedScan, uint numVoxels)\n{\n uint blockId = __mul24(blockIdx.y, gridDim.x) + blockIdx.x;\n uint i = __mul24(blockId, blockDim.x) + threadIdx.x;\n\n if (voxelOccupied[i] && (i < numVoxels)) {\n compactedVoxelArray[voxelOccupiedScan[i]] = i;\n }\n}\n\nextern \"C\" void launch_compactVoxels(dim3 grid,\n dim3 threads,\n uint *compactedVoxelArray,\n uint *voxelOccupied,\n uint *voxelOccupiedScan,\n uint numVoxels)\n{\n compactVoxels<<>>(compactedVoxelArray, voxelOccupied, voxelOccupiedScan, numVoxels);\n getLastCudaError(\"compactVoxels failed\");\n}\n\n// compute interpolated vertex along an edge\n__device__ float3 vertexInterp(float isolevel, float3 p0, float3 p1, float f0, float f1)\n{\n float t = (isolevel - f0) / (f1 - f0);\n return lerp(p0, p1, t);\n}\n\n// compute interpolated vertex position and normal along an edge\n__device__ void vertexInterp2(float isolevel, float3 p0, float3 p1, float4 f0, float4 f1, float3 &p, float3 &n)\n{\n float t = (isolevel - f0.w) / (f1.w - f0.w);\n p = lerp(p0, p1, t);\n n.x = lerp(f0.x, f1.x, t);\n n.y = lerp(f0.y, f1.y, t);\n n.z = lerp(f0.z, f1.z, t);\n // n = normalize(n);\n}\n\n// generate triangles for each voxel using marching cubes\n// interpolates normals from field function\n__global__ void generateTriangles(float4 *pos,\n float4 *norm,\n uint *compactedVoxelArray,\n uint *numVertsScanned,\n uint3 gridSize,\n uint3 gridSizeShift,\n uint3 gridSizeMask,\n float3 voxelSize,\n float isoValue,\n uint activeVoxels,\n uint maxVerts,\n hipTextureObject_t triTex,\n hipTextureObject_t numVertsTex)\n{\n uint blockId = __mul24(blockIdx.y, gridDim.x) + blockIdx.x;\n uint i = __mul24(blockId, blockDim.x) + threadIdx.x;\n\n if (i > activeVoxels - 1) {\n // can't return here because of syncthreads()\n i = activeVoxels - 1;\n }\n\n#if SKIP_EMPTY_VOXELS\n uint voxel = compactedVoxelArray[i];\n#else\n uint voxel = i;\n#endif\n\n // compute position in 3d grid\n uint3 gridPos = calcGridPos(voxel, gridSizeShift, gridSizeMask);\n\n float3 p;\n p.x = -1.0f + (gridPos.x * voxelSize.x);\n p.y = -1.0f + (gridPos.y * voxelSize.y);\n p.z = -1.0f + (gridPos.z * voxelSize.z);\n\n // calculate cell vertex positions\n float3 v[8];\n v[0] = p;\n v[1] = p + make_float3(voxelSize.x, 0, 0);\n v[2] = p + make_float3(voxelSize.x, voxelSize.y, 0);\n v[3] = p + make_float3(0, voxelSize.y, 0);\n v[4] = p + make_float3(0, 0, voxelSize.z);\n v[5] = p + make_float3(voxelSize.x, 0, voxelSize.z);\n v[6] = p + make_float3(voxelSize.x, voxelSize.y, voxelSize.z);\n v[7] = p + make_float3(0, voxelSize.y, voxelSize.z);\n\n // evaluate field values\n float4 field[8];\n field[0] = fieldFunc4(v[0]);\n field[1] = fieldFunc4(v[1]);\n field[2] = fieldFunc4(v[2]);\n field[3] = fieldFunc4(v[3]);\n field[4] = fieldFunc4(v[4]);\n field[5] = fieldFunc4(v[5]);\n field[6] = fieldFunc4(v[6]);\n field[7] = fieldFunc4(v[7]);\n\n // recalculate flag\n // (this is faster than storing it in global memory)\n uint cubeindex;\n cubeindex = uint(field[0].w < isoValue);\n cubeindex += uint(field[1].w < isoValue) * 2;\n cubeindex += uint(field[2].w < isoValue) * 4;\n cubeindex += uint(field[3].w < isoValue) * 8;\n cubeindex += uint(field[4].w < isoValue) * 16;\n cubeindex += uint(field[5].w < isoValue) * 32;\n cubeindex += uint(field[6].w < isoValue) * 64;\n cubeindex += uint(field[7].w < isoValue) * 128;\n\n // find the vertices where the surface intersects the cube\n\n#if USE_SHARED\n // use partioned shared memory to avoid using local memory\n __shared__ float3 vertlist[12 * NTHREADS];\n __shared__ float3 normlist[12 * NTHREADS];\n\n vertexInterp2(isoValue, v[0], v[1], field[0], field[1], vertlist[threadIdx.x], normlist[threadIdx.x]);\n vertexInterp2(\n isoValue, v[1], v[2], field[1], field[2], vertlist[threadIdx.x + NTHREADS], normlist[threadIdx.x + NTHREADS]);\n vertexInterp2(isoValue,\n v[2],\n v[3],\n field[2],\n field[3],\n vertlist[threadIdx.x + (NTHREADS * 2)],\n normlist[threadIdx.x + (NTHREADS * 2)]);\n vertexInterp2(isoValue,\n v[3],\n v[0],\n field[3],\n field[0],\n vertlist[threadIdx.x + (NTHREADS * 3)],\n normlist[threadIdx.x + (NTHREADS * 3)]);\n vertexInterp2(isoValue,\n v[4],\n v[5],\n field[4],\n field[5],\n vertlist[threadIdx.x + (NTHREADS * 4)],\n normlist[threadIdx.x + (NTHREADS * 4)]);\n vertexInterp2(isoValue,\n v[5],\n v[6],\n field[5],\n field[6],\n vertlist[threadIdx.x + (NTHREADS * 5)],\n normlist[threadIdx.x + (NTHREADS * 5)]);\n vertexInterp2(isoValue,\n v[6],\n v[7],\n field[6],\n field[7],\n vertlist[threadIdx.x + (NTHREADS * 6)],\n normlist[threadIdx.x + (NTHREADS * 6)]);\n vertexInterp2(isoValue,\n v[7],\n v[4],\n field[7],\n field[4],\n vertlist[threadIdx.x + (NTHREADS * 7)],\n normlist[threadIdx.x + (NTHREADS * 7)]);\n vertexInterp2(isoValue,\n v[0],\n v[4],\n field[0],\n field[4],\n vertlist[threadIdx.x + (NTHREADS * 8)],\n normlist[threadIdx.x + (NTHREADS * 8)]);\n vertexInterp2(isoValue,\n v[1],\n v[5],\n field[1],\n field[5],\n vertlist[threadIdx.x + (NTHREADS * 9)],\n normlist[threadIdx.x + (NTHREADS * 9)]);\n vertexInterp2(isoValue,\n v[2],\n v[6],\n field[2],\n field[6],\n vertlist[threadIdx.x + (NTHREADS * 10)],\n normlist[threadIdx.x + (NTHREADS * 10)]);\n vertexInterp2(isoValue,\n v[3],\n v[7],\n field[3],\n field[7],\n vertlist[threadIdx.x + (NTHREADS * 11)],\n normlist[threadIdx.x + (NTHREADS * 11)]);\n __syncthreads();\n\n#else\n float3 vertlist[12];\n float3 normlist[12];\n\n vertexInterp2(isoValue, v[0], v[1], field[0], field[1], vertlist[0], normlist[0]);\n vertexInterp2(isoValue, v[1], v[2], field[1], field[2], vertlist[1], normlist[1]);\n vertexInterp2(isoValue, v[2], v[3], field[2], field[3], vertlist[2], normlist[2]);\n vertexInterp2(isoValue, v[3], v[0], field[3], field[0], vertlist[3], normlist[3]);\n\n vertexInterp2(isoValue, v[4], v[5], field[4], field[5], vertlist[4], normlist[4]);\n vertexInterp2(isoValue, v[5], v[6], field[5], field[6], vertlist[5], normlist[5]);\n vertexInterp2(isoValue, v[6], v[7], field[6], field[7], vertlist[6], normlist[6]);\n vertexInterp2(isoValue, v[7], v[4], field[7], field[4], vertlist[7], normlist[7]);\n\n vertexInterp2(isoValue, v[0], v[4], field[0], field[4], vertlist[8], normlist[8]);\n vertexInterp2(isoValue, v[1], v[5], field[1], field[5], vertlist[9], normlist[9]);\n vertexInterp2(isoValue, v[2], v[6], field[2], field[6], vertlist[10], normlist[10]);\n vertexInterp2(isoValue, v[3], v[7], field[3], field[7], vertlist[11], normlist[11]);\n#endif\n\n // output triangle vertices\n uint numVerts = tex1Dfetch(numVertsTex, cubeindex);\n\n for (int i = 0; i < numVerts; i++) {\n uint edge = tex1Dfetch(triTex, cubeindex * 16 + i);\n\n uint index = numVertsScanned[voxel] + i;\n\n if (index < maxVerts) {\n#if USE_SHARED\n pos[index] = make_float4(vertlist[(edge * NTHREADS) + threadIdx.x], 1.0f);\n norm[index] = make_float4(normlist[(edge * NTHREADS) + threadIdx.x], 0.0f);\n#else\n pos[index] = make_float4(vertlist[edge], 1.0f);\n norm[index] = make_float4(normlist[edge], 0.0f);\n#endif\n }\n }\n}\n\nextern \"C\" void launch_generateTriangles(dim3 grid,\n dim3 threads,\n float4 *pos,\n float4 *norm,\n uint *compactedVoxelArray,\n uint *numVertsScanned,\n uint3 gridSize,\n uint3 gridSizeShift,\n uint3 gridSizeMask,\n float3 voxelSize,\n float isoValue,\n uint activeVoxels,\n uint maxVerts)\n{\n generateTriangles<<>>(pos,\n norm,\n compactedVoxelArray,\n numVertsScanned,\n gridSize,\n gridSizeShift,\n gridSizeMask,\n voxelSize,\n isoValue,\n activeVoxels,\n maxVerts,\n triTex,\n numVertsTex);\n getLastCudaError(\"generateTriangles failed\");\n}\n\n// calculate triangle normal\n__device__ float3 calcNormal(float3 *v0, float3 *v1, float3 *v2)\n{\n float3 edge0 = *v1 - *v0;\n float3 edge1 = *v2 - *v0;\n // note - it's faster to perform normalization in vertex shader rather than\n // here\n return cross(edge0, edge1);\n}\n\n// version that calculates flat surface normal for each triangle\n__global__ void generateTriangles2(float4 *pos,\n float4 *norm,\n uint *compactedVoxelArray,\n uint *numVertsScanned,\n uchar *volume,\n uint3 gridSize,\n uint3 gridSizeShift,\n uint3 gridSizeMask,\n float3 voxelSize,\n float isoValue,\n uint activeVoxels,\n uint maxVerts,\n hipTextureObject_t triTex,\n hipTextureObject_t numVertsTex,\n hipTextureObject_t volumeTex)\n{\n uint blockId = __mul24(blockIdx.y, gridDim.x) + blockIdx.x;\n uint i = __mul24(blockId, blockDim.x) + threadIdx.x;\n\n if (i > activeVoxels - 1) {\n i = activeVoxels - 1;\n }\n\n#if SKIP_EMPTY_VOXELS\n uint voxel = compactedVoxelArray[i];\n#else\n uint voxel = i;\n#endif\n\n // compute position in 3d grid\n uint3 gridPos = calcGridPos(voxel, gridSizeShift, gridSizeMask);\n\n float3 p;\n p.x = -1.0f + (gridPos.x * voxelSize.x);\n p.y = -1.0f + (gridPos.y * voxelSize.y);\n p.z = -1.0f + (gridPos.z * voxelSize.z);\n\n // calculate cell vertex positions\n float3 v[8];\n v[0] = p;\n v[1] = p + make_float3(voxelSize.x, 0, 0);\n v[2] = p + make_float3(voxelSize.x, voxelSize.y, 0);\n v[3] = p + make_float3(0, voxelSize.y, 0);\n v[4] = p + make_float3(0, 0, voxelSize.z);\n v[5] = p + make_float3(voxelSize.x, 0, voxelSize.z);\n v[6] = p + make_float3(voxelSize.x, voxelSize.y, voxelSize.z);\n v[7] = p + make_float3(0, voxelSize.y, voxelSize.z);\n\n#if SAMPLE_VOLUME\n float field[8];\n field[0] = sampleVolume(volumeTex, volume, gridPos, gridSize);\n field[1] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 0, 0), gridSize);\n field[2] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 1, 0), gridSize);\n field[3] = sampleVolume(volumeTex, volume, gridPos + make_uint3(0, 1, 0), gridSize);\n field[4] = sampleVolume(volumeTex, volume, gridPos + make_uint3(0, 0, 1), gridSize);\n field[5] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 0, 1), gridSize);\n field[6] = sampleVolume(volumeTex, volume, gridPos + make_uint3(1, 1, 1), gridSize);\n field[7] = sampleVolume(volumeTex, volume, gridPos + make_uint3(0, 1, 1), gridSize);\n#else\n // evaluate field values\n float field[8];\n field[0] = fieldFunc(v[0]);\n field[1] = fieldFunc(v[1]);\n field[2] = fieldFunc(v[2]);\n field[3] = fieldFunc(v[3]);\n field[4] = fieldFunc(v[4]);\n field[5] = fieldFunc(v[5]);\n field[6] = fieldFunc(v[6]);\n field[7] = fieldFunc(v[7]);\n#endif\n\n // recalculate flag\n uint cubeindex;\n cubeindex = uint(field[0] < isoValue);\n cubeindex += uint(field[1] < isoValue) * 2;\n cubeindex += uint(field[2] < isoValue) * 4;\n cubeindex += uint(field[3] < isoValue) * 8;\n cubeindex += uint(field[4] < isoValue) * 16;\n cubeindex += uint(field[5] < isoValue) * 32;\n cubeindex += uint(field[6] < isoValue) * 64;\n cubeindex += uint(field[7] < isoValue) * 128;\n\n // find the vertices where the surface intersects the cube\n\n#if USE_SHARED\n // use shared memory to avoid using local\n __shared__ float3 vertlist[12 * NTHREADS];\n\n vertlist[threadIdx.x] = vertexInterp(isoValue, v[0], v[1], field[0], field[1]);\n vertlist[NTHREADS + threadIdx.x] = vertexInterp(isoValue, v[1], v[2], field[1], field[2]);\n vertlist[(NTHREADS * 2) + threadIdx.x] = vertexInterp(isoValue, v[2], v[3], field[2], field[3]);\n vertlist[(NTHREADS * 3) + threadIdx.x] = vertexInterp(isoValue, v[3], v[0], field[3], field[0]);\n vertlist[(NTHREADS * 4) + threadIdx.x] = vertexInterp(isoValue, v[4], v[5], field[4], field[5]);\n vertlist[(NTHREADS * 5) + threadIdx.x] = vertexInterp(isoValue, v[5], v[6], field[5], field[6]);\n vertlist[(NTHREADS * 6) + threadIdx.x] = vertexInterp(isoValue, v[6], v[7], field[6], field[7]);\n vertlist[(NTHREADS * 7) + threadIdx.x] = vertexInterp(isoValue, v[7], v[4], field[7], field[4]);\n vertlist[(NTHREADS * 8) + threadIdx.x] = vertexInterp(isoValue, v[0], v[4], field[0], field[4]);\n vertlist[(NTHREADS * 9) + threadIdx.x] = vertexInterp(isoValue, v[1], v[5], field[1], field[5]);\n vertlist[(NTHREADS * 10) + threadIdx.x] = vertexInterp(isoValue, v[2], v[6], field[2], field[6]);\n vertlist[(NTHREADS * 11) + threadIdx.x] = vertexInterp(isoValue, v[3], v[7], field[3], field[7]);\n __syncthreads();\n#else\n\n float3 vertlist[12];\n\n vertlist[0] = vertexInterp(isoValue, v[0], v[1], field[0], field[1]);\n vertlist[1] = vertexInterp(isoValue, v[1], v[2], field[1], field[2]);\n vertlist[2] = vertexInterp(isoValue, v[2], v[3], field[2], field[3]);\n vertlist[3] = vertexInterp(isoValue, v[3], v[0], field[3], field[0]);\n\n vertlist[4] = vertexInterp(isoValue, v[4], v[5], field[4], field[5]);\n vertlist[5] = vertexInterp(isoValue, v[5], v[6], field[5], field[6]);\n vertlist[6] = vertexInterp(isoValue, v[6], v[7], field[6], field[7]);\n vertlist[7] = vertexInterp(isoValue, v[7], v[4], field[7], field[4]);\n\n vertlist[8] = vertexInterp(isoValue, v[0], v[4], field[0], field[4]);\n vertlist[9] = vertexInterp(isoValue, v[1], v[5], field[1], field[5]);\n vertlist[10] = vertexInterp(isoValue, v[2], v[6], field[2], field[6]);\n vertlist[11] = vertexInterp(isoValue, v[3], v[7], field[3], field[7]);\n#endif\n\n // output triangle vertices\n uint numVerts = tex1Dfetch(numVertsTex, cubeindex);\n\n for (int i = 0; i < numVerts; i += 3) {\n uint index = numVertsScanned[voxel] + i;\n\n float3 *v[3];\n uint edge;\n edge = tex1Dfetch(triTex, (cubeindex * 16) + i);\n#if USE_SHARED\n v[0] = &vertlist[(edge * NTHREADS) + threadIdx.x];\n#else\n v[0] = &vertlist[edge];\n#endif\n\n edge = tex1Dfetch(triTex, (cubeindex * 16) + i + 1);\n#if USE_SHARED\n v[1] = &vertlist[(edge * NTHREADS) + threadIdx.x];\n#else\n v[1] = &vertlist[edge];\n#endif\n\n edge = tex1Dfetch(triTex, (cubeindex * 16) + i + 2);\n#if USE_SHARED\n v[2] = &vertlist[(edge * NTHREADS) + threadIdx.x];\n#else\n v[2] = &vertlist[edge];\n#endif\n\n // calculate triangle surface normal\n float3 n = calcNormal(v[0], v[1], v[2]);\n\n if (index < (maxVerts - 3)) {\n pos[index] = make_float4(*v[0], 1.0f);\n norm[index] = make_float4(n, 0.0f);\n\n pos[index + 1] = make_float4(*v[1], 1.0f);\n norm[index + 1] = make_float4(n, 0.0f);\n\n pos[index + 2] = make_float4(*v[2], 1.0f);\n norm[index + 2] = make_float4(n, 0.0f);\n }\n }\n}\n\nextern \"C\" void launch_generateTriangles2(dim3 grid,\n dim3 threads,\n float4 *pos,\n float4 *norm,\n uint *compactedVoxelArray,\n uint *numVertsScanned,\n uchar *volume,\n uint3 gridSize,\n uint3 gridSizeShift,\n uint3 gridSizeMask,\n float3 voxelSize,\n float isoValue,\n uint activeVoxels,\n uint maxVerts)\n{\n generateTriangles2<<>>(pos,\n norm,\n compactedVoxelArray,\n numVertsScanned,\n volume,\n gridSize,\n gridSizeShift,\n gridSizeMask,\n voxelSize,\n isoValue,\n activeVoxels,\n maxVerts,\n triTex,\n numVertsTex,\n volumeTex);\n getLastCudaError(\"generateTriangles2 failed\");\n}\n\nextern \"C\" void ThrustScanWrapper(unsigned int *output, unsigned int *input, unsigned int numElements)\n{\n thrust::exclusive_scan(thrust::device_ptr(input),\n thrust::device_ptr(input + numElements),\n thrust::device_ptr(output));\n}\n\n#endif\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/marchingCubes/marchingCubes_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n\n// Implements NV12 to BGR batch conversion\n\n#include \n#include \n\n#include \"resize_convert.h\"\n\n#define CONV_THREADS_X 64\n#define CONV_THREADS_Y 10\n\n__forceinline__ __device__ static float clampF(float x, float lower, float upper)\n{\n return x < lower ? lower : (x > upper ? upper : x);\n}\n\n__global__ static void nv12ToBGRplanarBatchKernel(const uint8_t *pNv12,\n int nNv12Pitch,\n float *pBgr,\n int nRgbPitch,\n int nWidth,\n int nHeight,\n int nBatchSize)\n{\n int x = threadIdx.x + blockIdx.x * blockDim.x;\n int y = threadIdx.y + blockIdx.y * blockDim.y;\n\n if ((x << 2) + 1 > nWidth || (y << 1) + 1 > nHeight)\n return;\n\n const uint8_t *__restrict__ pSrc = pNv12;\n\n for (int i = blockIdx.z; i < nBatchSize; i += gridDim.z) {\n pSrc = pNv12 + i * ((nHeight * nNv12Pitch * 3) >> 1) + (x << 2) + (y << 1) * nNv12Pitch;\n uchar4 luma2x01, luma2x23, uv2;\n *(uint32_t *)&luma2x01 = *(uint32_t *)pSrc;\n *(uint32_t *)&luma2x23 = *(uint32_t *)(pSrc + nNv12Pitch);\n *(uint32_t *)&uv2 = *(uint32_t *)(pSrc + (nHeight - y) * nNv12Pitch);\n\n float *pDstBlock = (pBgr + i * ((nHeight * nRgbPitch * 3) >> 2) + ((blockIdx.x * blockDim.x) << 2)\n + ((blockIdx.y * blockDim.y) << 1) * (nRgbPitch >> 2));\n\n float2 add1;\n float2 add2;\n float2 add3;\n float2 add00, add01, add02, add03;\n float2 d, e;\n\n add00.x = 1.1644f * luma2x01.x;\n add01.x = 1.1644f * luma2x01.y;\n add00.y = 1.1644f * luma2x01.z;\n add01.y = 1.1644f * luma2x01.w;\n\n add02.x = 1.1644f * luma2x23.x;\n add03.x = 1.1644f * luma2x23.y;\n add02.y = 1.1644f * luma2x23.z;\n add03.y = 1.1644f * luma2x23.w;\n\n d.x = uv2.x - 128.0f;\n e.x = uv2.y - 128.0f;\n d.y = uv2.z - 128.0f;\n e.y = uv2.w - 128.0f;\n\n add1.x = 2.0172f * d.x;\n add1.y = 2.0172f * d.y;\n\n add2.x = (-0.3918f) * d.x + (-0.8130f) * e.x;\n add2.y = (-0.3918f) * d.y + (-0.8130f) * e.y;\n\n add3.x = 1.5960f * e.x;\n add3.y = 1.5960f * e.y;\n\n int rowStride = (threadIdx.y << 1) * (nRgbPitch >> 2);\n int nextRowStride = ((threadIdx.y << 1) + 1) * (nRgbPitch >> 2);\n // B\n *((float4 *)&pDstBlock[rowStride + (threadIdx.x << 2)]) = make_float4(clampF(add00.x + add1.x, 0.0f, 255.0f),\n clampF(add01.x + add1.x, 0.0f, 255.0f),\n clampF(add00.y + add1.y, 0.0f, 255.0f),\n clampF(add01.y + add1.y, 0.0f, 255.0f));\n *((float4 *)&pDstBlock[nextRowStride + (threadIdx.x << 2)]) =\n make_float4(clampF(add02.x + add1.x, 0.0f, 255.0f),\n clampF(add03.x + add1.x, 0.0f, 255.0f),\n clampF(add02.y + add1.y, 0.0f, 255.0f),\n clampF(add03.y + add1.y, 0.0f, 255.0f));\n\n int planeStride = nHeight * nRgbPitch >> 2;\n // G\n *((float4 *)&pDstBlock[planeStride + rowStride + (threadIdx.x << 2)]) =\n make_float4(clampF(add00.x + add2.x, 0.0f, 255.0f),\n clampF(add01.x + add2.x, 0.0f, 255.0f),\n clampF(add00.y + add2.y, 0.0f, 255.0f),\n clampF(add01.y + add2.y, 0.0f, 255.0f));\n *((float4 *)&pDstBlock[planeStride + nextRowStride + (threadIdx.x << 2)]) =\n make_float4(clampF(add02.x + add2.x, 0.0f, 255.0f),\n clampF(add03.x + add2.x, 0.0f, 255.0f),\n clampF(add02.y + add2.y, 0.0f, 255.0f),\n clampF(add03.y + add2.y, 0.0f, 255.0f));\n\n // R\n *((float4 *)&pDstBlock[(planeStride << 1) + rowStride + (threadIdx.x << 2)]) =\n make_float4(clampF(add00.x + add3.x, 0.0f, 255.0f),\n clampF(add01.x + add3.x, 0.0f, 255.0f),\n clampF(add00.y + add3.y, 0.0f, 255.0f),\n clampF(add01.y + add3.y, 0.0f, 255.0f));\n *((float4 *)&pDstBlock[(planeStride << 1) + nextRowStride + (threadIdx.x << 2)]) =\n make_float4(clampF(add02.x + add3.x, 0.0f, 255.0f),\n clampF(add03.x + add3.x, 0.0f, 255.0f),\n clampF(add02.y + add3.y, 0.0f, 255.0f),\n clampF(add03.y + add3.y, 0.0f, 255.0f));\n }\n}\n\nvoid nv12ToBGRplanarBatch(uint8_t *pNv12,\n int nNv12Pitch,\n float *pBgr,\n int nRgbPitch,\n int nWidth,\n int nHeight,\n int nBatchSize,\n cudaStream_t stream)\n{\n dim3 threads(CONV_THREADS_X, CONV_THREADS_Y);\n\n size_t blockDimZ = nBatchSize;\n\n // Restricting blocks in Z-dim till 32 to not launch too many blocks\n blockDimZ = (blockDimZ > 32) ? 32 : blockDimZ;\n\n dim3 blocks((nWidth / 4 - 1) / threads.x + 1, (nHeight / 2 - 1) / threads.y + 1, blockDimZ);\n nv12ToBGRplanarBatchKernel<<>>(\n pNv12, nNv12Pitch, pBgr, nRgbPitch, nWidth, nHeight, nBatchSize);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n\n// Implements NV12 to BGR batch conversion\n\n#include \n\n#include \"resize_convert.h\"\n\n#define CONV_THREADS_X 64\n#define CONV_THREADS_Y 10\n\n__forceinline__ __device__ static float clampF(float x, float lower, float upper)\n{\n return x < lower ? lower : (x > upper ? upper : x);\n}\n\n__global__ static void nv12ToBGRplanarBatchKernel(const uint8_t *pNv12,\n int nNv12Pitch,\n float *pBgr,\n int nRgbPitch,\n int nWidth,\n int nHeight,\n int nBatchSize)\n{\n int x = threadIdx.x + blockIdx.x * blockDim.x;\n int y = threadIdx.y + blockIdx.y * blockDim.y;\n\n if ((x << 2) + 1 > nWidth || (y << 1) + 1 > nHeight)\n return;\n\n const uint8_t *__restrict__ pSrc = pNv12;\n\n for (int i = blockIdx.z; i < nBatchSize; i += gridDim.z) {\n pSrc = pNv12 + i * ((nHeight * nNv12Pitch * 3) >> 1) + (x << 2) + (y << 1) * nNv12Pitch;\n uchar4 luma2x01, luma2x23, uv2;\n *(uint32_t *)&luma2x01 = *(uint32_t *)pSrc;\n *(uint32_t *)&luma2x23 = *(uint32_t *)(pSrc + nNv12Pitch);\n *(uint32_t *)&uv2 = *(uint32_t *)(pSrc + (nHeight - y) * nNv12Pitch);\n\n float *pDstBlock = (pBgr + i * ((nHeight * nRgbPitch * 3) >> 2) + ((blockIdx.x * blockDim.x) << 2)\n + ((blockIdx.y * blockDim.y) << 1) * (nRgbPitch >> 2));\n\n float2 add1;\n float2 add2;\n float2 add3;\n float2 add00, add01, add02, add03;\n float2 d, e;\n\n add00.x = 1.1644f * luma2x01.x;\n add01.x = 1.1644f * luma2x01.y;\n add00.y = 1.1644f * luma2x01.z;\n add01.y = 1.1644f * luma2x01.w;\n\n add02.x = 1.1644f * luma2x23.x;\n add03.x = 1.1644f * luma2x23.y;\n add02.y = 1.1644f * luma2x23.z;\n add03.y = 1.1644f * luma2x23.w;\n\n d.x = uv2.x - 128.0f;\n e.x = uv2.y - 128.0f;\n d.y = uv2.z - 128.0f;\n e.y = uv2.w - 128.0f;\n\n add1.x = 2.0172f * d.x;\n add1.y = 2.0172f * d.y;\n\n add2.x = (-0.3918f) * d.x + (-0.8130f) * e.x;\n add2.y = (-0.3918f) * d.y + (-0.8130f) * e.y;\n\n add3.x = 1.5960f * e.x;\n add3.y = 1.5960f * e.y;\n\n int rowStride = (threadIdx.y << 1) * (nRgbPitch >> 2);\n int nextRowStride = ((threadIdx.y << 1) + 1) * (nRgbPitch >> 2);\n // B\n *((float4 *)&pDstBlock[rowStride + (threadIdx.x << 2)]) = make_float4(clampF(add00.x + add1.x, 0.0f, 255.0f),\n clampF(add01.x + add1.x, 0.0f, 255.0f),\n clampF(add00.y + add1.y, 0.0f, 255.0f),\n clampF(add01.y + add1.y, 0.0f, 255.0f));\n *((float4 *)&pDstBlock[nextRowStride + (threadIdx.x << 2)]) =\n make_float4(clampF(add02.x + add1.x, 0.0f, 255.0f),\n clampF(add03.x + add1.x, 0.0f, 255.0f),\n clampF(add02.y + add1.y, 0.0f, 255.0f),\n clampF(add03.y + add1.y, 0.0f, 255.0f));\n\n int planeStride = nHeight * nRgbPitch >> 2;\n // G\n *((float4 *)&pDstBlock[planeStride + rowStride + (threadIdx.x << 2)]) =\n make_float4(clampF(add00.x + add2.x, 0.0f, 255.0f),\n clampF(add01.x + add2.x, 0.0f, 255.0f),\n clampF(add00.y + add2.y, 0.0f, 255.0f),\n clampF(add01.y + add2.y, 0.0f, 255.0f));\n *((float4 *)&pDstBlock[planeStride + nextRowStride + (threadIdx.x << 2)]) =\n make_float4(clampF(add02.x + add2.x, 0.0f, 255.0f),\n clampF(add03.x + add2.x, 0.0f, 255.0f),\n clampF(add02.y + add2.y, 0.0f, 255.0f),\n clampF(add03.y + add2.y, 0.0f, 255.0f));\n\n // R\n *((float4 *)&pDstBlock[(planeStride << 1) + rowStride + (threadIdx.x << 2)]) =\n make_float4(clampF(add00.x + add3.x, 0.0f, 255.0f),\n clampF(add01.x + add3.x, 0.0f, 255.0f),\n clampF(add00.y + add3.y, 0.0f, 255.0f),\n clampF(add01.y + add3.y, 0.0f, 255.0f));\n *((float4 *)&pDstBlock[(planeStride << 1) + nextRowStride + (threadIdx.x << 2)]) =\n make_float4(clampF(add02.x + add3.x, 0.0f, 255.0f),\n clampF(add03.x + add3.x, 0.0f, 255.0f),\n clampF(add02.y + add3.y, 0.0f, 255.0f),\n clampF(add03.y + add3.y, 0.0f, 255.0f));\n }\n}\n\nvoid nv12ToBGRplanarBatch(uint8_t *pNv12,\n int nNv12Pitch,\n float *pBgr,\n int nRgbPitch,\n int nWidth,\n int nHeight,\n int nBatchSize,\n hipStream_t stream)\n{\n dim3 threads(CONV_THREADS_X, CONV_THREADS_Y);\n\n size_t blockDimZ = nBatchSize;\n\n // Restricting blocks in Z-dim till 32 to not launch too many blocks\n blockDimZ = (blockDimZ > 32) ? 32 : blockDimZ;\n\n dim3 blocks((nWidth / 4 - 1) / threads.x + 1, (nHeight / 2 - 1) / threads.y + 1, blockDimZ);\n nv12ToBGRplanarBatchKernel<<>>(\n pNv12, nNv12Pitch, pBgr, nRgbPitch, nWidth, nHeight, nBatchSize);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/NV12toBGRandResize/nv12_to_bgr_planar.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"resize_convert.h\"\n#include \"utils.h\"\n\n__global__ void floatToChar(float *src, unsigned char *dst, int height, int width, int batchSize)\n{\n int x = threadIdx.x + blockIdx.x * blockDim.x;\n\n if (x >= height * width)\n return;\n\n int offset = height * width * 3;\n\n for (int j = 0; j < batchSize; j++) {\n // b\n *(dst + j * offset + x * 3 + 0) = (unsigned char)*(src + j * offset + height * width * 0 + x);\n // g\n *(dst + j * offset + x * 3 + 1) = (unsigned char)*(src + j * offset + height * width * 1 + x);\n // r\n *(dst + j * offset + x * 3 + 2) = (unsigned char)*(src + j * offset + height * width * 2 + x);\n }\n}\n\nvoid floatPlanarToChar(float *src, unsigned char *dst, int height, int width, int batchSize)\n{\n floatToChar<<<(height * width - 1) / 1024 + 1, 1024, 0, NULL>>>(src, dst, height, width, batchSize);\n}\n\nvoid dumpRawBGR(float *d_srcBGR, int pitch, int width, int height, int batchSize, char *folder, char *tag)\n{\n float *bgr, *d_bgr;\n int frameSize;\n char directory[120];\n char mkdir_cmd[256];\n#if !defined(_WIN32)\n sprintf(directory, \"output/%s\", folder);\n sprintf(mkdir_cmd, \"mkdir -p %s 2> /dev/null\", directory);\n#else\n sprintf(directory, \"output\\\\%s\", folder);\n sprintf(mkdir_cmd, \"mkdir %s 2> nul\", directory);\n#endif\n\n int ret = system(mkdir_cmd);\n\n frameSize = width * height * 3 * sizeof(float);\n bgr = (float *)malloc(frameSize);\n if (bgr == NULL) {\n std::cerr << \"Failed malloc for bgr\\n\";\n return;\n }\n\n d_bgr = d_srcBGR;\n for (int i = 0; i < batchSize; i++) {\n char filename[256];\n std::ofstream *outputFile;\n\n checkCudaErrors(cudaMemcpy((void *)bgr, (void *)d_bgr, frameSize, cudaMemcpyDeviceToHost));\n snprintf(filename, sizeof(filename), \"%s/%s_%d.raw\", directory, tag, (i + 1));\n\n outputFile = new std::ofstream(filename);\n if (outputFile) {\n outputFile->write((char *)bgr, frameSize);\n delete outputFile;\n }\n\n d_bgr += pitch * height * 3;\n }\n\n free(bgr);\n}\n\nvoid dumpBGR(float *d_srcBGR, int pitch, int width, int height, int batchSize, char *folder, char *tag)\n{\n dumpRawBGR(d_srcBGR, pitch, width, height, batchSize, folder, tag);\n}\n\nvoid dumpYUV(unsigned char *d_nv12, int size, char *folder, char *tag)\n{\n unsigned char *nv12Data;\n std::ofstream *nv12File;\n char filename[256];\n char directory[60];\n char mkdir_cmd[256];\n#if !defined(_WIN32)\n sprintf(directory, \"output/%s\", folder);\n sprintf(mkdir_cmd, \"mkdir -p %s 2> /dev/null\", directory);\n#else\n sprintf(directory, \"output\\\\%s\", folder);\n sprintf(mkdir_cmd, \"mkdir %s 2> nul\", directory);\n#endif\n\n int ret = system(mkdir_cmd);\n\n snprintf(filename, sizeof(filename), \"%s/%s.nv12\", directory, tag);\n\n nv12File = new std::ofstream(filename);\n if (nv12File == NULL) {\n std::cerr << \"Failed to new \" << filename;\n return;\n }\n\n nv12Data = (unsigned char *)malloc(size * (sizeof(char)));\n if (nv12Data == NULL) {\n std::cerr << \"Failed to allcoate memory\\n\";\n return;\n }\n\n cudaMemcpy((void *)nv12Data, (void *)d_nv12, size, cudaMemcpyDeviceToHost);\n\n nv12File->write((const char *)nv12Data, size);\n\n free(nv12Data);\n delete nv12File;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"resize_convert.h\"\n#include \"utils.h\"\n\n__global__ void floatToChar(float *src, unsigned char *dst, int height, int width, int batchSize)\n{\n int x = threadIdx.x + blockIdx.x * blockDim.x;\n\n if (x >= height * width)\n return;\n\n int offset = height * width * 3;\n\n for (int j = 0; j < batchSize; j++) {\n // b\n *(dst + j * offset + x * 3 + 0) = (unsigned char)*(src + j * offset + height * width * 0 + x);\n // g\n *(dst + j * offset + x * 3 + 1) = (unsigned char)*(src + j * offset + height * width * 1 + x);\n // r\n *(dst + j * offset + x * 3 + 2) = (unsigned char)*(src + j * offset + height * width * 2 + x);\n }\n}\n\nvoid floatPlanarToChar(float *src, unsigned char *dst, int height, int width, int batchSize)\n{\n floatToChar<<<(height * width - 1) / 1024 + 1, 1024, 0, NULL>>>(src, dst, height, width, batchSize);\n}\n\nvoid dumpRawBGR(float *d_srcBGR, int pitch, int width, int height, int batchSize, char *folder, char *tag)\n{\n float *bgr, *d_bgr;\n int frameSize;\n char directory[120];\n char mkdir_cmd[256];\n#if !defined(_WIN32)\n sprintf(directory, \"output/%s\", folder);\n sprintf(mkdir_cmd, \"mkdir -p %s 2> /dev/null\", directory);\n#else\n sprintf(directory, \"output\\\\%s\", folder);\n sprintf(mkdir_cmd, \"mkdir %s 2> nul\", directory);\n#endif\n\n int ret = system(mkdir_cmd);\n\n frameSize = width * height * 3 * sizeof(float);\n bgr = (float *)malloc(frameSize);\n if (bgr == NULL) {\n std::cerr << \"Failed malloc for bgr\\n\";\n return;\n }\n\n d_bgr = d_srcBGR;\n for (int i = 0; i < batchSize; i++) {\n char filename[256];\n std::ofstream *outputFile;\n\n checkCudaErrors(hipMemcpy((void *)bgr, (void *)d_bgr, frameSize, hipMemcpyDeviceToHost));\n snprintf(filename, sizeof(filename), \"%s/%s_%d.raw\", directory, tag, (i + 1));\n\n outputFile = new std::ofstream(filename);\n if (outputFile) {\n outputFile->write((char *)bgr, frameSize);\n delete outputFile;\n }\n\n d_bgr += pitch * height * 3;\n }\n\n free(bgr);\n}\n\nvoid dumpBGR(float *d_srcBGR, int pitch, int width, int height, int batchSize, char *folder, char *tag)\n{\n dumpRawBGR(d_srcBGR, pitch, width, height, batchSize, folder, tag);\n}\n\nvoid dumpYUV(unsigned char *d_nv12, int size, char *folder, char *tag)\n{\n unsigned char *nv12Data;\n std::ofstream *nv12File;\n char filename[256];\n char directory[60];\n char mkdir_cmd[256];\n#if !defined(_WIN32)\n sprintf(directory, \"output/%s\", folder);\n sprintf(mkdir_cmd, \"mkdir -p %s 2> /dev/null\", directory);\n#else\n sprintf(directory, \"output\\\\%s\", folder);\n sprintf(mkdir_cmd, \"mkdir %s 2> nul\", directory);\n#endif\n\n int ret = system(mkdir_cmd);\n\n snprintf(filename, sizeof(filename), \"%s/%s.nv12\", directory, tag);\n\n nv12File = new std::ofstream(filename);\n if (nv12File == NULL) {\n std::cerr << \"Failed to new \" << filename;\n return;\n }\n\n nv12Data = (unsigned char *)malloc(size * (sizeof(char)));\n if (nv12Data == NULL) {\n std::cerr << \"Failed to allcoate memory\\n\";\n return;\n }\n\n hipMemcpy((void *)nv12Data, (void *)d_nv12, size, hipMemcpyDeviceToHost);\n\n nv12File->write((const char *)nv12Data, size);\n\n free(nv12Data);\n delete nv12File;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/NV12toBGRandResize/utils.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Implements interlace NV12 frames batch resize\n\n#include \n#include \n\n#include \"resize_convert.h\"\n\n__global__ static void resizeNV12BatchKernel(cudaTextureObject_t texSrcLuma,\n cudaTextureObject_t texSrcChroma,\n uint8_t *pDstNv12,\n int nSrcWidth,\n int nSrcHeight,\n int nDstPitch,\n int nDstWidth,\n int nDstHeight,\n int nBatchSize)\n{\n int x = threadIdx.x + blockIdx.x * blockDim.x;\n int y = threadIdx.y + blockIdx.y * blockDim.y;\n\n int px = x * 2, py = y * 2;\n\n if ((px + 1) >= nDstWidth || (py + 1) >= nDstHeight)\n return;\n\n float fxScale = 1.0f * nSrcWidth / nDstWidth;\n float fyScale = 1.0f * nSrcHeight / nDstHeight;\n\n uint8_t *p = pDstNv12 + px + py * nDstPitch;\n int hh = nDstHeight * 3 / 2;\n int nByte = nDstPitch * hh;\n int px_fxScale = px * fxScale;\n int px_fxScale_1 = (px + 1) * fxScale;\n int py_fyScale = py * fyScale;\n int py_fyScale_1 = (py + 1) * fyScale;\n\n for (int i = blockIdx.z; i < nBatchSize; i += gridDim.z) {\n *(uchar2 *)p = make_uchar2(tex2D(texSrcLuma, px_fxScale, py_fyScale),\n tex2D(texSrcLuma, px_fxScale_1, py_fyScale));\n *(uchar2 *)(p + nDstPitch) = make_uchar2(tex2D(texSrcLuma, px_fxScale, py_fyScale_1),\n tex2D(texSrcLuma, px_fxScale_1, py_fyScale_1));\n *(uchar2 *)(p + (nDstHeight - y) * nDstPitch) =\n tex2D(texSrcChroma, x * fxScale, (hh * i + nDstHeight + y) * fyScale);\n p += nByte;\n py += hh;\n }\n}\n\nvoid resizeNV12Batch(uint8_t *dpSrc,\n int nSrcPitch,\n int nSrcWidth,\n int nSrcHeight,\n uint8_t *dpDst,\n int nDstPitch,\n int nDstWidth,\n int nDstHeight,\n int nBatchSize,\n cudaStream_t stream)\n{\n int hhSrc = ceilf(nSrcHeight * 3.0f / 2.0f);\n cudaResourceDesc resDesc = {};\n resDesc.resType = cudaResourceTypePitch2D;\n resDesc.res.pitch2D.devPtr = dpSrc;\n resDesc.res.pitch2D.desc = cudaCreateChannelDesc();\n resDesc.res.pitch2D.width = nSrcWidth;\n resDesc.res.pitch2D.height = hhSrc * nBatchSize;\n resDesc.res.pitch2D.pitchInBytes = nSrcPitch;\n\n cudaTextureDesc texDesc = {};\n texDesc.filterMode = cudaFilterModePoint;\n texDesc.readMode = cudaReadModeElementType;\n\n cudaTextureObject_t texLuma = 0;\n checkCudaErrors(cudaCreateTextureObject(&texLuma, &resDesc, &texDesc, NULL));\n\n resDesc.res.pitch2D.desc = cudaCreateChannelDesc();\n resDesc.res.pitch2D.width /= 2;\n\n cudaTextureObject_t texChroma = 0;\n checkCudaErrors(cudaCreateTextureObject(&texChroma, &resDesc, &texDesc, NULL));\n\n dim3 block(32, 32, 1);\n\n size_t blockDimZ = nBatchSize;\n\n // Restricting blocks in Z-dim till 32 to not launch too many blocks\n blockDimZ = (blockDimZ > 32) ? 32 : blockDimZ;\n\n dim3 grid((nDstWidth / 2 + block.x) / block.x, (nDstHeight / 2 + block.y) / block.y, blockDimZ);\n resizeNV12BatchKernel<<>>(\n texLuma, texChroma, dpDst, nSrcWidth, nSrcHeight, nDstPitch, nDstWidth, nDstHeight, nBatchSize);\n\n checkCudaErrors(cudaDestroyTextureObject(texLuma));\n checkCudaErrors(cudaDestroyTextureObject(texChroma));\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Implements interlace NV12 frames batch resize\n\n#include \n\n#include \"resize_convert.h\"\n\n__global__ static void resizeNV12BatchKernel(hipTextureObject_t texSrcLuma,\n hipTextureObject_t texSrcChroma,\n uint8_t *pDstNv12,\n int nSrcWidth,\n int nSrcHeight,\n int nDstPitch,\n int nDstWidth,\n int nDstHeight,\n int nBatchSize)\n{\n int x = threadIdx.x + blockIdx.x * blockDim.x;\n int y = threadIdx.y + blockIdx.y * blockDim.y;\n\n int px = x * 2, py = y * 2;\n\n if ((px + 1) >= nDstWidth || (py + 1) >= nDstHeight)\n return;\n\n float fxScale = 1.0f * nSrcWidth / nDstWidth;\n float fyScale = 1.0f * nSrcHeight / nDstHeight;\n\n uint8_t *p = pDstNv12 + px + py * nDstPitch;\n int hh = nDstHeight * 3 / 2;\n int nByte = nDstPitch * hh;\n int px_fxScale = px * fxScale;\n int px_fxScale_1 = (px + 1) * fxScale;\n int py_fyScale = py * fyScale;\n int py_fyScale_1 = (py + 1) * fyScale;\n\n for (int i = blockIdx.z; i < nBatchSize; i += gridDim.z) {\n *(uchar2 *)p = make_uchar2(tex2D(texSrcLuma, px_fxScale, py_fyScale),\n tex2D(texSrcLuma, px_fxScale_1, py_fyScale));\n *(uchar2 *)(p + nDstPitch) = make_uchar2(tex2D(texSrcLuma, px_fxScale, py_fyScale_1),\n tex2D(texSrcLuma, px_fxScale_1, py_fyScale_1));\n *(uchar2 *)(p + (nDstHeight - y) * nDstPitch) =\n tex2D(texSrcChroma, x * fxScale, (hh * i + nDstHeight + y) * fyScale);\n p += nByte;\n py += hh;\n }\n}\n\nvoid resizeNV12Batch(uint8_t *dpSrc,\n int nSrcPitch,\n int nSrcWidth,\n int nSrcHeight,\n uint8_t *dpDst,\n int nDstPitch,\n int nDstWidth,\n int nDstHeight,\n int nBatchSize,\n hipStream_t stream)\n{\n int hhSrc = ceilf(nSrcHeight * 3.0f / 2.0f);\n hipResourceDesc resDesc = {};\n resDesc.resType = hipResourceTypePitch2D;\n resDesc.res.pitch2D.devPtr = dpSrc;\n resDesc.res.pitch2D.desc = hipCreateChannelDesc();\n resDesc.res.pitch2D.width = nSrcWidth;\n resDesc.res.pitch2D.height = hhSrc * nBatchSize;\n resDesc.res.pitch2D.pitchInBytes = nSrcPitch;\n\n hipTextureDesc texDesc = {};\n texDesc.filterMode = hipFilterModePoint;\n texDesc.readMode = hipReadModeElementType;\n\n hipTextureObject_t texLuma = 0;\n checkCudaErrors(hipCreateTextureObject(&texLuma, &resDesc, &texDesc, NULL));\n\n resDesc.res.pitch2D.desc = hipCreateChannelDesc();\n resDesc.res.pitch2D.width /= 2;\n\n hipTextureObject_t texChroma = 0;\n checkCudaErrors(hipCreateTextureObject(&texChroma, &resDesc, &texDesc, NULL));\n\n dim3 block(32, 32, 1);\n\n size_t blockDimZ = nBatchSize;\n\n // Restricting blocks in Z-dim till 32 to not launch too many blocks\n blockDimZ = (blockDimZ > 32) ? 32 : blockDimZ;\n\n dim3 grid((nDstWidth / 2 + block.x) / block.x, (nDstHeight / 2 + block.y) / block.y, blockDimZ);\n resizeNV12BatchKernel<<>>(\n texLuma, texChroma, dpDst, nSrcWidth, nSrcHeight, nDstPitch, nDstWidth, nDstHeight, nBatchSize);\n\n checkCudaErrors(hipDestroyTextureObject(texLuma));\n checkCudaErrors(hipDestroyTextureObject(texChroma));\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/NV12toBGRandResize/nv12_resize.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n\n// Implements BGR 3 progressive planars frames batch resize\n\n#include \n#include \n\n#include \"resize_convert.h\"\n\n__global__ void resizeBGRplanarBatchKernel(cudaTextureObject_t texSrc,\n float *pDst,\n int nDstPitch,\n int nDstHeight,\n int nSrcHeight,\n int batch,\n float scaleX,\n float scaleY,\n int cropX,\n int cropY,\n int cropW,\n int cropH)\n{\n int x = threadIdx.x + blockIdx.x * blockDim.x;\n int y = threadIdx.y + blockIdx.y * blockDim.y;\n\n if (x >= (int)(cropW / scaleX) || y >= (int)(cropH / scaleY))\n return;\n\n int frameSize = nDstPitch * nDstHeight;\n float *p = NULL;\n for (int i = blockIdx.z; i < batch; i += gridDim.z) {\n#pragma unroll\n for (int channel = 0; channel < 3; channel++) {\n p = pDst + i * 3 * frameSize + y * nDstPitch + x + channel * frameSize;\n *p = tex2D(texSrc, x * scaleX + cropX, ((3 * i + channel) * nSrcHeight + y * scaleY + cropY));\n }\n }\n}\n\n\nstatic void resizeBGRplanarBatchCore(float *dpSrc,\n int nSrcPitch,\n int nSrcWidth,\n int nSrcHeight,\n float *dpDst,\n int nDstPitch,\n int nDstWidth,\n int nDstHeight,\n int nBatchSize,\n cudaStream_t stream,\n bool whSameResizeRatio,\n int cropX,\n int cropY,\n int cropW,\n int cropH)\n{\n cudaTextureObject_t texSrc[2];\n int nTiles = 1, h, iTile;\n\n h = nSrcHeight * 3 * nBatchSize;\n while ((h + nTiles - 1) / nTiles > 65536)\n nTiles++;\n\n if (nTiles > 2)\n return;\n\n int batchTile = nBatchSize / nTiles;\n int batchTileLast = nBatchSize - batchTile * (nTiles - 1);\n\n for (iTile = 0; iTile < nTiles; ++iTile) {\n int bs = (iTile == nTiles - 1) ? batchTileLast : batchTile;\n float *dpSrcNew = dpSrc + iTile * (batchTile * 3 * nSrcHeight * nSrcPitch);\n\n cudaResourceDesc resDesc = {};\n resDesc.resType = cudaResourceTypePitch2D;\n resDesc.res.pitch2D.devPtr = dpSrcNew;\n resDesc.res.pitch2D.desc = cudaCreateChannelDesc();\n resDesc.res.pitch2D.width = nSrcWidth;\n resDesc.res.pitch2D.height = bs * 3 * nSrcHeight;\n resDesc.res.pitch2D.pitchInBytes = nSrcPitch * sizeof(float);\n cudaTextureDesc texDesc = {};\n texDesc.filterMode = cudaFilterModeLinear;\n texDesc.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&texSrc[iTile], &resDesc, &texDesc, NULL));\n float *dpDstNew = dpDst + iTile * (batchTile * 3 * nDstHeight * nDstPitch);\n\n if (cropW == 0 || cropH == 0) {\n cropX = 0;\n cropY = 0;\n cropW = nSrcWidth;\n cropH = nSrcHeight;\n }\n\n float scaleX = (cropW * 1.0f / nDstWidth);\n float scaleY = (cropH * 1.0f / nDstHeight);\n\n if (whSameResizeRatio == true)\n scaleX = scaleY = scaleX > scaleY ? scaleX : scaleY;\n dim3 block(32, 32, 1);\n\n size_t blockDimZ = bs;\n // Restricting blocks in Z-dim till 32 to not launch too many blocks\n blockDimZ = (blockDimZ > 32) ? 32 : blockDimZ;\n dim3 grid((cropW * 1.0f / scaleX + block.x - 1) / block.x,\n (cropH * 1.0f / scaleY + block.y - 1) / block.y,\n blockDimZ);\n\n resizeBGRplanarBatchKernel<<>>(\n texSrc[iTile], dpDstNew, nDstPitch, nDstHeight, nSrcHeight, bs, scaleX, scaleY, cropX, cropY, cropW, cropH);\n }\n\n for (iTile = 0; iTile < nTiles; ++iTile)\n checkCudaErrors(cudaDestroyTextureObject(texSrc[iTile]));\n}\n\nvoid resizeBGRplanarBatch(float *dpSrc,\n int nSrcPitch,\n int nSrcWidth,\n int nSrcHeight,\n float *dpDst,\n int nDstPitch,\n int nDstWidth,\n int nDstHeight,\n int nBatchSize,\n cudaStream_t stream,\n int cropX,\n int cropY,\n int cropW,\n int cropH,\n bool whSameResizeRatio)\n{\n resizeBGRplanarBatchCore(dpSrc,\n nSrcPitch,\n nSrcWidth,\n nSrcHeight,\n dpDst,\n nDstPitch,\n nDstWidth,\n nDstHeight,\n nBatchSize,\n stream,\n whSameResizeRatio,\n cropX,\n cropY,\n cropW,\n cropH);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n\n// Implements BGR 3 progressive planars frames batch resize\n\n#include \n\n#include \"resize_convert.h\"\n\n__global__ void resizeBGRplanarBatchKernel(hipTextureObject_t texSrc,\n float *pDst,\n int nDstPitch,\n int nDstHeight,\n int nSrcHeight,\n int batch,\n float scaleX,\n float scaleY,\n int cropX,\n int cropY,\n int cropW,\n int cropH)\n{\n int x = threadIdx.x + blockIdx.x * blockDim.x;\n int y = threadIdx.y + blockIdx.y * blockDim.y;\n\n if (x >= (int)(cropW / scaleX) || y >= (int)(cropH / scaleY))\n return;\n\n int frameSize = nDstPitch * nDstHeight;\n float *p = NULL;\n for (int i = blockIdx.z; i < batch; i += gridDim.z) {\n#pragma unroll\n for (int channel = 0; channel < 3; channel++) {\n p = pDst + i * 3 * frameSize + y * nDstPitch + x + channel * frameSize;\n *p = tex2D(texSrc, x * scaleX + cropX, ((3 * i + channel) * nSrcHeight + y * scaleY + cropY));\n }\n }\n}\n\n\nstatic void resizeBGRplanarBatchCore(float *dpSrc,\n int nSrcPitch,\n int nSrcWidth,\n int nSrcHeight,\n float *dpDst,\n int nDstPitch,\n int nDstWidth,\n int nDstHeight,\n int nBatchSize,\n hipStream_t stream,\n bool whSameResizeRatio,\n int cropX,\n int cropY,\n int cropW,\n int cropH)\n{\n hipTextureObject_t texSrc[2];\n int nTiles = 1, h, iTile;\n\n h = nSrcHeight * 3 * nBatchSize;\n while ((h + nTiles - 1) / nTiles > 65536)\n nTiles++;\n\n if (nTiles > 2)\n return;\n\n int batchTile = nBatchSize / nTiles;\n int batchTileLast = nBatchSize - batchTile * (nTiles - 1);\n\n for (iTile = 0; iTile < nTiles; ++iTile) {\n int bs = (iTile == nTiles - 1) ? batchTileLast : batchTile;\n float *dpSrcNew = dpSrc + iTile * (batchTile * 3 * nSrcHeight * nSrcPitch);\n\n hipResourceDesc resDesc = {};\n resDesc.resType = hipResourceTypePitch2D;\n resDesc.res.pitch2D.devPtr = dpSrcNew;\n resDesc.res.pitch2D.desc = hipCreateChannelDesc();\n resDesc.res.pitch2D.width = nSrcWidth;\n resDesc.res.pitch2D.height = bs * 3 * nSrcHeight;\n resDesc.res.pitch2D.pitchInBytes = nSrcPitch * sizeof(float);\n hipTextureDesc texDesc = {};\n texDesc.filterMode = hipFilterModeLinear;\n texDesc.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&texSrc[iTile], &resDesc, &texDesc, NULL));\n float *dpDstNew = dpDst + iTile * (batchTile * 3 * nDstHeight * nDstPitch);\n\n if (cropW == 0 || cropH == 0) {\n cropX = 0;\n cropY = 0;\n cropW = nSrcWidth;\n cropH = nSrcHeight;\n }\n\n float scaleX = (cropW * 1.0f / nDstWidth);\n float scaleY = (cropH * 1.0f / nDstHeight);\n\n if (whSameResizeRatio == true)\n scaleX = scaleY = scaleX > scaleY ? scaleX : scaleY;\n dim3 block(32, 32, 1);\n\n size_t blockDimZ = bs;\n // Restricting blocks in Z-dim till 32 to not launch too many blocks\n blockDimZ = (blockDimZ > 32) ? 32 : blockDimZ;\n dim3 grid((cropW * 1.0f / scaleX + block.x - 1) / block.x,\n (cropH * 1.0f / scaleY + block.y - 1) / block.y,\n blockDimZ);\n\n resizeBGRplanarBatchKernel<<>>(\n texSrc[iTile], dpDstNew, nDstPitch, nDstHeight, nSrcHeight, bs, scaleX, scaleY, cropX, cropY, cropW, cropH);\n }\n\n for (iTile = 0; iTile < nTiles; ++iTile)\n checkCudaErrors(hipDestroyTextureObject(texSrc[iTile]));\n}\n\nvoid resizeBGRplanarBatch(float *dpSrc,\n int nSrcPitch,\n int nSrcWidth,\n int nSrcHeight,\n float *dpDst,\n int nDstPitch,\n int nDstWidth,\n int nDstHeight,\n int nBatchSize,\n hipStream_t stream,\n int cropX,\n int cropY,\n int cropW,\n int cropH,\n bool whSameResizeRatio)\n{\n resizeBGRplanarBatchCore(dpSrc,\n nSrcPitch,\n nSrcWidth,\n nSrcHeight,\n dpDst,\n nDstPitch,\n nDstWidth,\n nDstHeight,\n nBatchSize,\n stream,\n whSameResizeRatio,\n cropX,\n cropY,\n cropW,\n cropH);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/NV12toBGRandResize/bgr_resize.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#ifndef QUASIRANDOMGENERATOR_KERNEL_CUH\n#define QUASIRANDOMGENERATOR_KERNEL_CUH\n\n#include \"quasirandomGenerator_common.h\"\n\n// Fast integer multiplication\n#define MUL(a, b) __umul24(a, b)\n\n////////////////////////////////////////////////////////////////////////////////\n// Niederreiter quasirandom number generation kernel\n////////////////////////////////////////////////////////////////////////////////\n__constant__ unsigned int c_Table[QRNG_DIMENSIONS][QRNG_RESOLUTION];\n\nextern \"C\" __global__ void quasirandomGeneratorKernel(float *d_Output, unsigned int seed, unsigned int N)\n{\n unsigned int *dimBase = &c_Table[threadIdx.y][0];\n unsigned int tid = MUL(blockDim.x, blockIdx.x) + threadIdx.x;\n unsigned int threadN = MUL(blockDim.x, gridDim.x);\n\n for (unsigned int pos = tid; pos < N; pos += threadN) {\n unsigned int result = 0;\n unsigned int data = seed + pos;\n\n for (int bit = 0; bit < QRNG_RESOLUTION; bit++, data >>= 1)\n if (data & 1) {\n result ^= dimBase[bit];\n }\n\n d_Output[MUL(threadIdx.y, N) + pos] = (float)(result + 1) * INT_SCALE;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Moro's Inverse Cumulative Normal Distribution function approximation\n////////////////////////////////////////////////////////////////////////////////\n__device__ inline float MoroInvCNDgpu(unsigned int x)\n{\n const float a1 = 2.50662823884f;\n const float a2 = -18.61500062529f;\n const float a3 = 41.39119773534f;\n const float a4 = -25.44106049637f;\n const float b1 = -8.4735109309f;\n const float b2 = 23.08336743743f;\n const float b3 = -21.06224101826f;\n const float b4 = 3.13082909833f;\n const float c1 = 0.337475482272615f;\n const float c2 = 0.976169019091719f;\n const float c3 = 0.160797971491821f;\n const float c4 = 2.76438810333863E-02f;\n const float c5 = 3.8405729373609E-03f;\n const float c6 = 3.951896511919E-04f;\n const float c7 = 3.21767881768E-05f;\n const float c8 = 2.888167364E-07f;\n const float c9 = 3.960315187E-07f;\n\n float z;\n\n bool negate = false;\n\n // Ensure the conversion to floating point will give a value in the\n // range (0,0.5] by restricting the input to the bottom half of the\n // input domain. We will later reflect the result if the input was\n // originally in the top half of the input domain\n if (x >= 0x80000000UL) {\n x = 0xffffffffUL - x;\n negate = true;\n }\n\n // x is now in the range [0,0x80000000) (i.e. [0,0x7fffffff])\n // Convert to floating point in (0,0.5]\n const float x1 = 1.0f / static_cast(0xffffffffUL);\n const float x2 = x1 / 2.0f;\n float p1 = x * x1 + x2;\n // Convert to floating point in (-0.5,0]\n float p2 = p1 - 0.5f;\n\n // The input to the Moro inversion is p2 which is in the range\n // (-0.5,0]. This means that our output will be the negative side\n // of the bell curve (which we will reflect if \"negate\" is true).\n\n // Main body of the bell curve for |p| < 0.42\n if (p2 > -0.42f) {\n z = p2 * p2;\n z = p2 * (((a4 * z + a3) * z + a2) * z + a1) / ((((b4 * z + b3) * z + b2) * z + b1) * z + 1.0f);\n }\n // Special case (Chebychev) for tail\n else {\n z = __logf(-__logf(p1));\n z = -(c1 + z * (c2 + z * (c3 + z * (c4 + z * (c5 + z * (c6 + z * (c7 + z * (c8 + z * c9))))))));\n }\n\n // If the original input (x) was in the top half of the range, reflect\n // to get the positive side of the bell curve\n\n return negate ? -z : z;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Main kernel. Choose between transforming\n// input sequence and uniform ascending (0, 1) sequence\n////////////////////////////////////////////////////////////////////////////////\n\nextern \"C\" __global__ void inverseCNDKernel(float *d_Output, unsigned int pathN)\n{\n unsigned int distance = ((unsigned int)-1) / (pathN + 1);\n unsigned int tid = MUL(blockDim.x, blockIdx.x) + threadIdx.x;\n unsigned int threadN = MUL(blockDim.x, gridDim.x);\n\n // Transform input number sequence if it's supplied\n if (0) // d_Input)\n {\n /*\n for (unsigned int pos = tid; pos < pathN; pos += threadN)\n {\n unsigned int d = d_Input[pos];\n d_Output[pos] = (float)MoroInvCNDgpu(d);\n }\n */\n }\n // Else generate input uniformly placed samples on the fly\n // and write to destination\n else {\n for (unsigned int pos = tid; pos < pathN; pos += threadN) {\n unsigned int d = (pos + 1) * distance;\n d_Output[pos] = (float)MoroInvCNDgpu(d);\n }\n }\n}\n\n#endif\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#ifndef QUASIRANDOMGENERATOR_KERNEL_CUH\n#define QUASIRANDOMGENERATOR_KERNEL_CUH\n\n#include \"quasirandomGenerator_common.h\"\n\n// Fast integer multiplication\n#define MUL(a, b) __umul24(a, b)\n\n////////////////////////////////////////////////////////////////////////////////\n// Niederreiter quasirandom number generation kernel\n////////////////////////////////////////////////////////////////////////////////\n__constant__ unsigned int c_Table[QRNG_DIMENSIONS][QRNG_RESOLUTION];\n\nextern \"C\" __global__ void quasirandomGeneratorKernel(float *d_Output, unsigned int seed, unsigned int N)\n{\n unsigned int *dimBase = &c_Table[threadIdx.y][0];\n unsigned int tid = MUL(blockDim.x, blockIdx.x) + threadIdx.x;\n unsigned int threadN = MUL(blockDim.x, gridDim.x);\n\n for (unsigned int pos = tid; pos < N; pos += threadN) {\n unsigned int result = 0;\n unsigned int data = seed + pos;\n\n for (int bit = 0; bit < QRNG_RESOLUTION; bit++, data >>= 1)\n if (data & 1) {\n result ^= dimBase[bit];\n }\n\n d_Output[MUL(threadIdx.y, N) + pos] = (float)(result + 1) * INT_SCALE;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Moro's Inverse Cumulative Normal Distribution function approximation\n////////////////////////////////////////////////////////////////////////////////\n__device__ inline float MoroInvCNDgpu(unsigned int x)\n{\n const float a1 = 2.50662823884f;\n const float a2 = -18.61500062529f;\n const float a3 = 41.39119773534f;\n const float a4 = -25.44106049637f;\n const float b1 = -8.4735109309f;\n const float b2 = 23.08336743743f;\n const float b3 = -21.06224101826f;\n const float b4 = 3.13082909833f;\n const float c1 = 0.337475482272615f;\n const float c2 = 0.976169019091719f;\n const float c3 = 0.160797971491821f;\n const float c4 = 2.76438810333863E-02f;\n const float c5 = 3.8405729373609E-03f;\n const float c6 = 3.951896511919E-04f;\n const float c7 = 3.21767881768E-05f;\n const float c8 = 2.888167364E-07f;\n const float c9 = 3.960315187E-07f;\n\n float z;\n\n bool negate = false;\n\n // Ensure the conversion to floating point will give a value in the\n // range (0,0.5] by restricting the input to the bottom half of the\n // input domain. We will later reflect the result if the input was\n // originally in the top half of the input domain\n if (x >= 0x80000000UL) {\n x = 0xffffffffUL - x;\n negate = true;\n }\n\n // x is now in the range [0,0x80000000) (i.e. [0,0x7fffffff])\n // Convert to floating point in (0,0.5]\n const float x1 = 1.0f / static_cast(0xffffffffUL);\n const float x2 = x1 / 2.0f;\n float p1 = x * x1 + x2;\n // Convert to floating point in (-0.5,0]\n float p2 = p1 - 0.5f;\n\n // The input to the Moro inversion is p2 which is in the range\n // (-0.5,0]. This means that our output will be the negative side\n // of the bell curve (which we will reflect if \"negate\" is true).\n\n // Main body of the bell curve for |p| < 0.42\n if (p2 > -0.42f) {\n z = p2 * p2;\n z = p2 * (((a4 * z + a3) * z + a2) * z + a1) / ((((b4 * z + b3) * z + b2) * z + b1) * z + 1.0f);\n }\n // Special case (Chebychev) for tail\n else {\n z = __logf(-__logf(p1));\n z = -(c1 + z * (c2 + z * (c3 + z * (c4 + z * (c5 + z * (c6 + z * (c7 + z * (c8 + z * c9))))))));\n }\n\n // If the original input (x) was in the top half of the range, reflect\n // to get the positive side of the bell curve\n\n return negate ? -z : z;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Main kernel. Choose between transforming\n// input sequence and uniform ascending (0, 1) sequence\n////////////////////////////////////////////////////////////////////////////////\n\nextern \"C\" __global__ void inverseCNDKernel(float *d_Output, unsigned int pathN)\n{\n unsigned int distance = ((unsigned int)-1) / (pathN + 1);\n unsigned int tid = MUL(blockDim.x, blockIdx.x) + threadIdx.x;\n unsigned int threadN = MUL(blockDim.x, gridDim.x);\n\n // Transform input number sequence if it's supplied\n if (0) // d_Input)\n {\n /*\n for (unsigned int pos = tid; pos < pathN; pos += threadN)\n {\n unsigned int d = d_Input[pos];\n d_Output[pos] = (float)MoroInvCNDgpu(d);\n }\n */\n }\n // Else generate input uniformly placed samples on the fly\n // and write to destination\n else {\n for (unsigned int pos = tid; pos < pathN; pos += threadN) {\n unsigned int d = (pos + 1) * distance;\n d_Output[pos] = (float)MoroInvCNDgpu(d);\n }\n }\n}\n\n#endif\n", "source": "/kaggle/working/cuda-samples/Samples/5_Domain_Specific/quasirandomGenerator_nvrtc/quasirandomGenerator_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Simple interop app demonstrating EGLImage + EGLSync interop with CUDA.\n// Using EGLSync - CUDA Event interop one can achieve synchronization on GPU\n// itself for GL-EGL-CUDA operations instead of blocking CPU for\n// synchronization. This app requires GLES 3.2 or higher\n\n//---------------------------INCLUDES---------------------------------//\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"egl_common.h\"\n#include \"graphics_interface.h\"\n\n//---------------------------DEFINES---------------------------------//\n#define MAX_ITR 100\n\n#define FAILURE 0\n#define SUCCESS 1\n#define WAIVED 2\n\n#define BLOCK_SIZE 16\n\n#define GL_READ 0\n#define GL_WRITE 1\n//---------------------------MACROS---------------------------------//\n\n// Error-checking wrapper around GL calls\n#define GL_SAFE_CALL(call) \\\n { \\\n GLenum err; \\\n call; \\\n err = glGetError(); \\\n if (err != GL_NO_ERROR) { \\\n fprintf(stderr, \"%s:%d GL error: %d\\n\", __FILE__, __LINE__, err); \\\n cleanup(FAILURE); \\\n } \\\n }\n\n#define GL_SAFE_CALL_NO_CLEANUP(call, err) \\\n { \\\n GLenum status; \\\n call; \\\n status = glGetError(); \\\n if (status != GL_NO_ERROR) { \\\n fprintf(stderr, \"%s:%d GL error: %d\\n\", __FILE__, __LINE__, status); \\\n err = status; \\\n } \\\n }\n\n// Error-checking wrapper around CUDA calls (taken from cutil.h)\n#define CUDA_SAFE_CALL(call) \\\n do { \\\n cudaError err = call; \\\n if (cudaSuccess != err) { \\\n fprintf( \\\n stderr, \"Cuda error in file '%s' in line %i : %s.\\n\", __FILE__, __LINE__, cudaGetErrorString(err)); \\\n cleanup(FAILURE); \\\n } \\\n } while (0)\n\n#define CUDA_SAFE_CALL_NO_CLEANUP(call, err) \\\n do { \\\n cudaError status = call; \\\n if (cudaSuccess != status) { \\\n fprintf( \\\n stderr, \"Cuda error in file '%s' in line %i : %s.\\n\", __FILE__, __LINE__, cudaGetErrorString(status)); \\\n err = status; \\\n } \\\n } while (0)\n\n#if defined(EXTENSION_LIST)\nEXTENSION_LIST(EXTLST_DECL)\ntypedef void (*extlst_fnptr_t)(void);\nstatic struct\n{\n extlst_fnptr_t *fnptr;\n char const *name;\n} extensionList[] = {EXTENSION_LIST(EXTLST_ENTRY)};\n\nint eglSetupExtensions(void)\n{\n unsigned int i;\n\n for (i = 0; i < (sizeof(extensionList) / sizeof(*extensionList)); i++) {\n *extensionList[i].fnptr = eglGetProcAddress(extensionList[i].name);\n if (*extensionList[i].fnptr == NULL) {\n printf(\"Couldn't get address of %s()\\n\", extensionList[i].name);\n return 0;\n }\n }\n\n return 1;\n}\n#endif\n\n#if defined(EXTENSION_LIST)\nEXTENSION_LIST(EXTLST_EXTERN)\n#endif\n\n//------------------------GLOBAL VARIABLES--------------------------//\n\n// GL texture\nGLuint tex[2] = {0};\n\n// Used to catch unexpected termination from GLUT\nint cleanExit = 0;\n\n// Use CPU Sync or GPU sync; Default GPU\nint useGpu = 1;\n\n// CUDA Resource\nCUgraphicsResource writeResource = NULL;\nCUgraphicsResource readResource = NULL;\nCUarray writeArray, readArray;\nCUdevice device;\nCUcontext context;\n\n// Which device to run on\nunsigned int dev = 0;\n\n// Default width, height, and iterations value\nint width = 2048;\nint height = 2048;\nint itr = MAX_ITR;\n\n// Error check variable\n__device__ static unsigned int numErrors = 0;\n\n//-----------------------FUNCTION PROTOTYPES------------------------//\n\nvoid checkSync(int argc, char **argv);\nint parseCmdLine(int argc, char **argv);\nvoid printUsage(void);\nvoid cleanup(int status);\nvoid exitHandler(void);\nvoid printStatus(int status);\nvoid checkSyncOnCPU(void);\nvoid checkSyncOnGPU(EGLDisplay dpy);\n\n__global__ void\nverify_and_update_kernel(CUsurfObject write, CUsurfObject read, char expected, char newval, int width, int height);\nextern \"C\" cudaError_t cudaGetValueMismatch();\n\n//-----------------------FUNCTION DEFINITIONS------------------------//\n\nint main(int argc, char *argv[])\n{\n#if defined(__linux__)\n setenv(\"DISPLAY\", \":0\", 0);\n#endif\n\n parseCmdLine(argc, argv);\n atexit(exitHandler);\n\n checkSync(argc, argv);\n return 0;\n}\n\nint parseCmdLine(int argc, char **argv)\n{\n int i;\n for (i = 1; i < argc; i++) {\n if (strcmp(argv[i], \"-cpu\") == 0) {\n useGpu = 0;\n }\n\n if (strcmp(argv[i], \"-h\") == 0) {\n printUsage();\n cleanup(SUCCESS);\n }\n\n if (strcmp(argv[i], \"-width\") == 0) {\n ++i;\n if (i == argc) {\n printf(\"width option must be followed by value\\n\");\n return FAILURE;\n }\n if (sscanf(argv[i], \"%d\", &width) != 1) {\n printf(\"Error: invalid width value\\n\");\n return FAILURE;\n }\n }\n\n if (strcmp(argv[i], \"-height\") == 0) {\n ++i;\n if (i == argc) {\n printf(\"height option must be followed by value\\n\");\n return FAILURE;\n }\n if (sscanf(argv[i], \"%d\", &height) != 1) {\n printf(\"Error: invalid height value\\n\");\n return FAILURE;\n }\n }\n if (strcmp(argv[i], \"-itr\") == 0) {\n ++i;\n if (i == argc) {\n printf(\"itr option must be followed by iteration value\\n\");\n return FAILURE;\n }\n if (sscanf(argv[i], \"%d\", &itr) != 1) {\n printf(\"Error: invalid iteration value\\n\");\n return FAILURE;\n }\n }\n }\n\n return SUCCESS;\n}\n\nvoid printUsage(void)\n{\n printf(\"Usage:\\n\");\n printf(\"\\t-h\\tPrint command line options\\n\");\n printf(\"\\t-cpu\\tSync on the CPU instead of the GPU\\n\");\n printf(\"\\t-width w\\tSet the width to w\\n\");\n printf(\"\\t-height h\\tSet the height to h\\n\");\n printf(\"\\t-itr i\\tSet number of iterations to i\\n\");\n}\n\nvoid checkSync(int argc, char **argv)\n{\n int x, y;\n int bufferSize = width * height * 4;\n unsigned char *pSurf_read = NULL, *pSurf_write = NULL;\n int integrated;\n\n CUresult status = CUDA_SUCCESS;\n CUctxCreateParams ctxCreateParams = {};\n\n // Init values for variables\n x = y = 0;\n\n if (CUDA_SUCCESS != (status = cuInit(0))) {\n printf(\"Failed to initialize CUDA\\n\");\n }\n device = findCudaDeviceDRV(argc, (const char **)argv);\n\n if (CUDA_SUCCESS != (status = cuCtxCreate(&context, &ctxCreateParams, 0, device))) {\n printf(\"failed to create CUDA context\\n\");\n }\n cuCtxPushCurrent(context);\n\n status = cuDeviceGetAttribute(&integrated, CU_DEVICE_ATTRIBUTE_INTEGRATED, device);\n if (status != CUDA_SUCCESS) {\n printf(\"Failed to get device attribute CU_DEVICE_ATTRIBUTE_INTEGRATED\\n\");\n cleanup(FAILURE);\n }\n\n if (integrated != 1) {\n printf(\"EGLSync_CUDAEvent_Interop does not support dGPU. Waiving sample.\\n\");\n cleanup(WAIVED);\n }\n\n#if (defined(__arm__) || defined(__aarch64__)) && defined(__linux__)\n graphics_setup_window(0, 0, width, height, \"EGLSync_CUDA_Interop\");\n#endif\n\n pSurf_read = (unsigned char *)malloc(bufferSize);\n pSurf_write = (unsigned char *)malloc(bufferSize);\n if (pSurf_read == NULL || pSurf_write == NULL) {\n printf(\"malloc failed\\n\");\n cleanup(FAILURE);\n }\n\n for (x = 0; x < width; x++) {\n for (y = 0; y < height; y++) {\n pSurf_read[(y * width + x) * 4] = 1;\n pSurf_read[(y * width + x) * 4 + 1] = 1;\n pSurf_read[(y * width + x) * 4 + 2] = 1;\n pSurf_read[(y * width + x) * 4 + 3] = 1;\n pSurf_write[(y * width + x) * 4] = 0;\n pSurf_write[(y * width + x) * 4 + 1] = 0;\n pSurf_write[(y * width + x) * 4 + 2] = 0;\n pSurf_write[(y * width + x) * 4 + 3] = 0;\n }\n }\n\n // NOP call to error-check the above glut calls\n GL_SAFE_CALL({});\n\n // Init texture\n GL_SAFE_CALL(glGenTextures(2, tex));\n\n GL_SAFE_CALL(glBindTexture(GL_TEXTURE_2D, tex[GL_READ]));\n GL_SAFE_CALL(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST));\n GL_SAFE_CALL(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST));\n GL_SAFE_CALL(glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, width, height, 0, GL_RGBA, GL_UNSIGNED_BYTE, pSurf_read));\n GL_SAFE_CALL(glBindTexture(GL_TEXTURE_2D, 0));\n\n GL_SAFE_CALL(glBindTexture(GL_TEXTURE_2D, tex[GL_WRITE]));\n GL_SAFE_CALL(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST));\n GL_SAFE_CALL(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST));\n GL_SAFE_CALL(glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, width, height, 0, GL_RGBA, GL_UNSIGNED_BYTE, pSurf_write));\n GL_SAFE_CALL(glBindTexture(GL_TEXTURE_2D, 0));\n\n glFinish();\n\n EGLDisplay eglDisplayHandle = eglGetCurrentDisplay();\n if (eglDisplayHandle == EGL_NO_DISPLAY) {\n printf(\"eglDisplayHandle failed \\n\");\n cleanup(FAILURE);\n }\n else {\n printf(\"eglDisplay Handle created \\n\");\n }\n\n if (!eglSetupExtensions()) {\n printf(\"SetupExtentions failed \\n\");\n cleanup(FAILURE);\n }\n\n EGLContext eglCtx = eglGetCurrentContext();\n if (eglCtx == EGL_NO_CONTEXT) {\n printf(\"Context1 create failed with error %d\\n\", eglGetError());\n cleanup(FAILURE);\n }\n\n // Create the EGL_Image\n EGLint eglImgAttrs[] = {EGL_IMAGE_PRESERVED_KHR, EGL_TRUE, EGL_NONE, EGL_NONE};\n\n EGLImageKHR eglImage1 = eglCreateImageKHR(\n eglDisplayHandle, eglCtx, EGL_GL_TEXTURE_2D_KHR, (EGLClientBuffer)(intptr_t)tex[GL_READ], eglImgAttrs);\n if (eglImage1 == EGL_NO_IMAGE_KHR) {\n printf(\"EGLImage create failed for read texture with error %d\\n\", eglGetError());\n cleanup(FAILURE);\n }\n else {\n printf(\"EGLImage1 created \\n\");\n }\n\n EGLImageKHR eglImage2 = eglCreateImageKHR(\n eglDisplayHandle, eglCtx, EGL_GL_TEXTURE_2D_KHR, (EGLClientBuffer)(intptr_t)tex[GL_WRITE], eglImgAttrs);\n if (eglImage2 == EGL_NO_IMAGE_KHR) {\n printf(\"EGLImage create failed for write texture with error %d\\n\", eglGetError());\n cleanup(FAILURE);\n }\n else {\n printf(\"EGLImage2 created \\n\");\n }\n\n glFinish();\n\n status = cuGraphicsEGLRegisterImage(&writeResource, eglImage1, CU_GRAPHICS_REGISTER_FLAGS_NONE);\n if (status != CUDA_SUCCESS) {\n printf(\"cuGraphicsEGLRegisterImage failed with Texture 1\\n\");\n cleanup(FAILURE);\n }\n else {\n printf(\"cuGraphicsEGLRegisterImage Passed, writeResource created with texture \"\n \"1\\n\");\n }\n\n status = cuGraphicsSubResourceGetMappedArray(&writeArray, writeResource, 0, 0);\n if (status != CUDA_SUCCESS) {\n printf(\"cuGraphicsSubResourceGetMappedArray failed for writeResource with \"\n \"texture 1\\n\");\n cleanup(FAILURE);\n }\n\n status = cuGraphicsEGLRegisterImage(&readResource, eglImage2, CU_GRAPHICS_REGISTER_FLAGS_NONE);\n if (status != CUDA_SUCCESS) {\n printf(\"cuGraphicsEGLRegisterImage failed for readResource with Texture 2\\n\");\n cleanup(FAILURE);\n }\n else {\n printf(\"cuGraphicsEGLRegisterImage Passed, readResource created with texture \"\n \"2\\n\");\n }\n\n status = cuGraphicsSubResourceGetMappedArray(&readArray, readResource, 0, 0);\n if (status != CUDA_SUCCESS) {\n printf(\"cuGraphicsSubResourceGetMappedArray failed for texture 2\\n\");\n cleanup(FAILURE);\n }\n\n if (useGpu) {\n printf(\"Using GPU Sync path\\n\");\n checkSyncOnGPU(eglDisplayHandle);\n }\n else {\n printf(\"Using CPU Sync path\\n\");\n checkSyncOnCPU();\n }\n\n free(pSurf_read);\n free(pSurf_write);\n cleanup(SUCCESS);\n}\n\nvoid checkSyncOnCPU(void)\n{\n int z = 0;\n unsigned char expectedData, newData;\n CUresult status = CUDA_SUCCESS;\n CUDA_RESOURCE_DESC wdsc, rdsc;\n memset(&wdsc, 0, sizeof(wdsc));\n memset(&rdsc, 0, sizeof(rdsc));\n\n expectedData = 0;\n newData = 1;\n\n wdsc.resType = CU_RESOURCE_TYPE_ARRAY;\n wdsc.res.array.hArray = writeArray;\n CUsurfObject writeSurface;\n rdsc.resType = CU_RESOURCE_TYPE_ARRAY;\n rdsc.res.array.hArray = readArray;\n CUsurfObject readSurface;\n\n status = cuSurfObjectCreate(&writeSurface, &wdsc);\n if (status != CUDA_SUCCESS) {\n printf(\"Surface bounding failed with status %d\\n\", status);\n cleanup(FAILURE);\n }\n status = cuSurfObjectCreate(&readSurface, &rdsc);\n if (status != CUDA_SUCCESS) {\n printf(\"Surface bounding failed\\n\");\n cleanup(FAILURE);\n }\n\n for (z = 0; z < itr; z++) {\n // GL call to copy from read texture to write texture\n GL_SAFE_CALL(glCopyImageSubData(\n tex[GL_READ], GL_TEXTURE_2D, 0, 0, 0, 0, tex[GL_WRITE], GL_TEXTURE_2D, 0, 0, 0, 0, width, height, 1));\n\n glFinish();\n\n newData++;\n expectedData++;\n\n verify_and_update_kernel<<<(width * height) / 256, 256>>>(\n writeSurface, readSurface, expectedData, newData, width, height);\n\n status = cuCtxSynchronize();\n if (status != CUDA_SUCCESS) {\n printf(\"cuCtxSynchronize failed \\n\");\n }\n }\n\n cudaError_t err = cudaGetValueMismatch();\n if (err != cudaSuccess) {\n printf(\"Value mismatch seen when using CPU sync\\n\");\n cleanup(FAILURE);\n }\n\n // Clean up CUDA writeResource\n status = cuGraphicsUnregisterResource(writeResource);\n if (status != CUDA_SUCCESS) {\n printf(\"Failed to unregister %d\", status);\n cleanup(FAILURE);\n }\n else {\n printf(\"Unregistered writeResource. \\n\");\n }\n\n // Clean up CUDA readResource\n status = cuGraphicsUnregisterResource(readResource);\n if (status != CUDA_SUCCESS) {\n printf(\"Failed to unregister %d\", status);\n cleanup(FAILURE);\n }\n else {\n printf(\"Unregistered readResource. \\n\");\n }\n}\n\n/*\n Performs same function as checkSyncOnCPU\n Here instead of glFinish() and cuCtxSynchronize like in checkSyncOnCPU,\n we make use of EGLSync, CUDA Event and cuStreamWaitEvent, eglWaitSyncKHR to\n achieve the synchronization due to this CPU is not blocked for any\n synchronization needed between GL-EGL & CUDA operations all synchronizations\n happens on the GPU only.\n*/\nvoid checkSyncOnGPU(EGLDisplay dpy)\n{\n int z = 0;\n unsigned char expectedData, newData;\n cudaError_t err;\n CUresult status = CUDA_SUCCESS;\n CUstream stream;\n CUevent timingDisabledEvent;\n CUDA_RESOURCE_DESC wdsc, rdsc;\n memset(&wdsc, 0, sizeof(wdsc));\n memset(&rdsc, 0, sizeof(rdsc));\n\n expectedData = 0;\n newData = 1;\n\n wdsc.resType = CU_RESOURCE_TYPE_ARRAY;\n wdsc.res.array.hArray = writeArray;\n CUsurfObject writeSurface;\n rdsc.resType = CU_RESOURCE_TYPE_ARRAY;\n rdsc.res.array.hArray = readArray;\n CUsurfObject readSurface;\n\n status = cuSurfObjectCreate(&writeSurface, &wdsc);\n if (status != CUDA_SUCCESS) {\n printf(\"Surface bounding failed with status %d\\n\", status);\n cleanup(FAILURE);\n }\n status = cuSurfObjectCreate(&readSurface, &rdsc);\n if (status != CUDA_SUCCESS) {\n printf(\"Surface bounding failed\\n\");\n cleanup(FAILURE);\n }\n\n status = cuStreamCreate(&stream, CU_STREAM_DEFAULT);\n if (status != CUDA_SUCCESS) {\n printf(\"Stream creation failed\\n\");\n cleanup(FAILURE);\n }\n\n // Creates timing disabled event which uses non-blocking synchronization\n status = cuEventCreate(&timingDisabledEvent, CU_EVENT_DISABLE_TIMING);\n if (status != CUDA_SUCCESS) {\n printf(\"Default event creation failed\\n\");\n cleanup(FAILURE);\n }\n\n /*\n 1. We perform texture-to-texture copy in GLES which is async function\n 2. Followed by creating EGLSync and a CUDA Event from that EGLSync object\n 3. Using cuStreamWaitEvent() we wait in GPU for the GLES to finish texture\n copy.\n 4. CUDA kernel verfiy_and_update_kernel verifies if the copied data by\n GLES is correct, and it updates the buffer with new values.\n 5. This is followed by eglWaitSyncKHR() which waits for the cuda kernel to\n finish, so that in the next iteration GLES can perform the copying of the\n updated buffer to write texture,\n */\n for (z = 0; z < itr; z++) {\n // GL call to copy from read texture to write texture\n GL_SAFE_CALL(glCopyImageSubData(\n tex[GL_READ], GL_TEXTURE_2D, 0, 0, 0, 0, tex[GL_WRITE], GL_TEXTURE_2D, 0, 0, 0, 0, width, height, 1));\n\n EGLSyncKHR eglSyncForGL, eglSyncForCuda;\n EGLBoolean egl_status = EGL_TRUE;\n EGLAttribKHR eglattrib[] = {EGL_CUDA_EVENT_HANDLE_NV, (EGLAttrib)timingDisabledEvent, EGL_NONE};\n\n CUevent cudaEGLSyncEvent;\n\n eglSyncForGL = eglCreateSyncKHR(dpy, EGL_SYNC_FENCE_KHR, NULL);\n\n if (eglSyncForGL == EGL_NO_SYNC_KHR) {\n printf(\" EGL Sync creation failed\\n\");\n cleanup(FAILURE);\n }\n\n status = cuEventCreateFromEGLSync(&cudaEGLSyncEvent, eglSyncForGL, CU_EVENT_DEFAULT);\n if (status != CUDA_SUCCESS) {\n printf(\"CUDA event creation from EGLSync failed\\n\");\n cleanup(FAILURE);\n }\n\n // We wait from CUDA in GPU for GL-EGL operation completion\n status = cuStreamWaitEvent(stream, cudaEGLSyncEvent, 0);\n if (status != CUDA_SUCCESS) {\n printf(\"Stream wait for event created from EGLSync failed\\n\");\n cleanup(FAILURE);\n }\n\n egl_status = eglDestroySyncKHR(dpy, eglSyncForGL);\n if (egl_status != EGL_TRUE) {\n printf(\"EGL sync object destruction failed\\n\");\n cleanup(FAILURE);\n }\n\n newData++;\n expectedData++;\n\n // Verifies the values in readSurface which is copied by\n // glCopyImageSubData() And writes value of newData into writeSurface\n verify_and_update_kernel<<<(width * height) / 256, 256, 0, stream>>>(\n writeSurface, readSurface, expectedData, newData, width, height);\n\n status = cuEventDestroy(cudaEGLSyncEvent);\n if (status != CUDA_SUCCESS) {\n printf(\"Event Destroy failed\\n\");\n cleanup(FAILURE);\n }\n\n status = cuEventRecord(timingDisabledEvent, stream);\n if (status != CUDA_SUCCESS) {\n printf(\"Event Record failed\\n\");\n cleanup(FAILURE);\n }\n\n // creating an EGL sync object linked to a CUDA event object\n eglSyncForCuda = eglCreateSync64KHR(dpy, EGL_SYNC_CUDA_EVENT_NV, eglattrib);\n\n // We wait from EGL for CUDA operation completion\n egl_status = eglWaitSyncKHR(dpy, eglSyncForCuda, 0);\n if (egl_status != EGL_TRUE) {\n printf(\"eglWaitSyncKHR failed\\n\");\n cleanup(FAILURE);\n }\n egl_status = eglDestroySyncKHR(dpy, eglSyncForCuda);\n if (egl_status != EGL_TRUE) {\n printf(\"EGL sync object destruction failed\\n\");\n cleanup(FAILURE);\n }\n }\n\n err = cudaGetValueMismatch();\n if (err != cudaSuccess) {\n printf(\"Value mismatch seen when using GPU sync\\n\");\n cleanup(FAILURE);\n }\n\n // Clean up CUDA writeResource\n status = cuGraphicsUnregisterResource(writeResource);\n if (status != CUDA_SUCCESS) {\n printf(\"Failed to unregister %d\", status);\n cleanup(FAILURE);\n }\n else {\n printf(\"Unregistered writeResource. \\n\");\n }\n\n // Clean up CUDA readResource\n status = cuGraphicsUnregisterResource(readResource);\n if (status != CUDA_SUCCESS) {\n printf(\"Failed to unregister %d\", status);\n cleanup(FAILURE);\n }\n else {\n printf(\"Unregistered readResource. \\n\");\n }\n}\n\n// Verifies the values in readSurface whether they are expected ones\n// And writes value of newData into writeSurface\n__global__ void\nverify_and_update_kernel(CUsurfObject write, CUsurfObject read, char expected, char newval, int width, int height)\n{\n unsigned int x = blockDim.x * blockIdx.x + threadIdx.x;\n unsigned int y = blockDim.y * blockIdx.y + threadIdx.y;\n\n if (x < width && y < height) {\n uchar4 check;\n surf2Dread(&check, read, x * 4, y);\n if (check.x != expected || check.y != expected || check.z != expected || check.w != expected) {\n printf(\"Mismatch found in values read[0]= %u read[1]= %u read[2]= %u \"\n \"read[3]= %u expected is %u\\n\",\n check.x,\n check.y,\n check.z,\n check.w,\n expected);\n numErrors++;\n return;\n }\n uchar4 data = make_uchar4(newval, newval, newval, newval);\n surf2Dwrite(data, write, x * 4, y);\n }\n}\n\n__global__ void getNumErrors(int *numErr) { *numErr = numErrors; }\n\nextern \"C\" cudaError_t cudaGetValueMismatch()\n{\n int numErr_h;\n int *numErr_d = NULL;\n cudaError_t err = cudaSuccess;\n\n err = cudaMalloc(&numErr_d, sizeof(int));\n if (err != cudaSuccess) {\n printf(\"Cuda Main: cudaMemcpy failed with %s\\n\", cudaGetErrorString(err));\n cudaFree(numErr_d);\n return err;\n }\n\n getNumErrors<<<1, 1>>>(numErr_d);\n err = cudaDeviceSynchronize();\n if (err != cudaSuccess) {\n printf(\"Cuda Main: cudaDeviceSynchronize failed with %s\\n\", cudaGetErrorString(err));\n }\n err = cudaMemcpy(&numErr_h, numErr_d, sizeof(int), cudaMemcpyDeviceToHost);\n if (err != cudaSuccess) {\n printf(\"Cuda Main: cudaMemcpy failed with %s\\n\", cudaGetErrorString(err));\n cudaFree(numErr_d);\n return err;\n }\n err = cudaFree(numErr_d);\n if (err != cudaSuccess) {\n printf(\"Cuda Main: cudaFree failed with %s\\n\", cudaGetErrorString(err));\n return err;\n }\n if (numErr_h > 0) {\n return cudaErrorUnknown;\n }\n return cudaSuccess;\n}\n\n// Clean up state and exit. If status is SUCCESS, regression success is printed\n// to stdout. This will happen if the glut timer is triggered. If status is\n// anything else, the regression failure message is printed.\nvoid cleanup(int status)\n{\n GLenum glErr = GL_NO_ERROR;\n cudaError cudaErr = cudaSuccess;\n int exitStatus = status;\n\n // Clean up GL\n if (*tex) {\n GL_SAFE_CALL_NO_CLEANUP(glDeleteTextures(2, tex), glErr);\n }\n\n // Print test status and exit\n if (glErr != GL_NO_ERROR || cudaErr != cudaSuccess)\n exitStatus = FAILURE;\n\n printStatus(exitStatus);\n\n cleanExit = 1;\n\n graphics_close_window();\n\n if (exitStatus == FAILURE)\n exit(EXIT_FAILURE);\n\n if (exitStatus == WAIVED)\n exit(EXIT_WAIVED);\n\n exit(0);\n}\n\nvoid exitHandler(void)\n{\n if (!cleanExit) {\n printf(\"&&&& EGLSync_CUDAEvent_Interop unexpected failure \\n\");\n printStatus(FAILURE);\n }\n}\n\n// Print test success or fail for regression testing\nvoid printStatus(int status)\n{\n switch (status) {\n case SUCCESS:\n printf(\"&&&& EGLSync_CUDAEvent_Interop PASSED\\n\");\n break;\n case WAIVED:\n printf(\"&&&& EGLSync_CUDAEvent_Interop WAIVED\\n\");\n break;\n default:\n printf(\"&&&& EGLSync_CUDAEvent_Interop FAILED\\n\");\n break;\n }\n fflush(stdout);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Simple interop app demonstrating EGLImage + EGLSync interop with CUDA.\n// Using EGLSync - CUDA Event interop one can achieve synchronization on GPU\n// itself for GL-EGL-CUDA operations instead of blocking CPU for\n// synchronization. This app requires GLES 3.2 or higher\n\n//---------------------------INCLUDES---------------------------------//\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"egl_common.h\"\n#include \"graphics_interface.h\"\n\n//---------------------------DEFINES---------------------------------//\n#define MAX_ITR 100\n\n#define FAILURE 0\n#define SUCCESS 1\n#define WAIVED 2\n\n#define BLOCK_SIZE 16\n\n#define GL_READ 0\n#define GL_WRITE 1\n//---------------------------MACROS---------------------------------//\n\n// Error-checking wrapper around GL calls\n#define GL_SAFE_CALL(call) \\\n { \\\n GLenum err; \\\n call; \\\n err = glGetError(); \\\n if (err != GL_NO_ERROR) { \\\n fprintf(stderr, \"%s:%d GL error: %d\\n\", __FILE__, __LINE__, err); \\\n cleanup(FAILURE); \\\n } \\\n }\n\n#define GL_SAFE_CALL_NO_CLEANUP(call, err) \\\n { \\\n GLenum status; \\\n call; \\\n status = glGetError(); \\\n if (status != GL_NO_ERROR) { \\\n fprintf(stderr, \"%s:%d GL error: %d\\n\", __FILE__, __LINE__, status); \\\n err = status; \\\n } \\\n }\n\n// Error-checking wrapper around CUDA calls (taken from cutil.h)\n#define CUDA_SAFE_CALL(call) \\\n do { \\\n hipError_t err = call; \\\n if (hipSuccess != err) { \\\n fprintf( \\\n stderr, \"Cuda error in file '%s' in line %i : %s.\\n\", __FILE__, __LINE__, hipGetErrorString(err)); \\\n cleanup(FAILURE); \\\n } \\\n } while (0)\n\n#define CUDA_SAFE_CALL_NO_CLEANUP(call, err) \\\n do { \\\n hipError_t status = call; \\\n if (hipSuccess != status) { \\\n fprintf( \\\n stderr, \"Cuda error in file '%s' in line %i : %s.\\n\", __FILE__, __LINE__, hipGetErrorString(status)); \\\n err = status; \\\n } \\\n } while (0)\n\n#if defined(EXTENSION_LIST)\nEXTENSION_LIST(EXTLST_DECL)\ntypedef void (*extlst_fnptr_t)(void);\nstatic struct\n{\n extlst_fnptr_t *fnptr;\n char const *name;\n} extensionList[] = {EXTENSION_LIST(EXTLST_ENTRY)};\n\nint eglSetupExtensions(void)\n{\n unsigned int i;\n\n for (i = 0; i < (sizeof(extensionList) / sizeof(*extensionList)); i++) {\n *extensionList[i].fnptr = eglGetProcAddress(extensionList[i].name);\n if (*extensionList[i].fnptr == NULL) {\n printf(\"Couldn't get address of %s()\\n\", extensionList[i].name);\n return 0;\n }\n }\n\n return 1;\n}\n#endif\n\n#if defined(EXTENSION_LIST)\nEXTENSION_LIST(EXTLST_EXTERN)\n#endif\n\n//------------------------GLOBAL VARIABLES--------------------------//\n\n// GL texture\nGLuint tex[2] = {0};\n\n// Used to catch unexpected termination from GLUT\nint cleanExit = 0;\n\n// Use CPU Sync or GPU sync; Default GPU\nint useGpu = 1;\n\n// CUDA Resource\nhipGraphicsResource_t writeResource = NULL;\nhipGraphicsResource_t readResource = NULL;\nhipArray_t writeArray, readArray;\nhipDevice_t device;\nhipCtx_t context;\n\n// Which device to run on\nunsigned int dev = 0;\n\n// Default width, height, and iterations value\nint width = 2048;\nint height = 2048;\nint itr = MAX_ITR;\n\n// Error check variable\n__device__ static unsigned int numErrors = 0;\n\n//-----------------------FUNCTION PROTOTYPES------------------------//\n\nvoid checkSync(int argc, char **argv);\nint parseCmdLine(int argc, char **argv);\nvoid printUsage(void);\nvoid cleanup(int status);\nvoid exitHandler(void);\nvoid printStatus(int status);\nvoid checkSyncOnCPU(void);\nvoid checkSyncOnGPU(EGLDisplay dpy);\n\n__global__ void\nverify_and_update_kernel(hipSurfaceObject_t write, hipSurfaceObject_t read, char expected, char newval, int width, int height);\nextern \"C\" hipError_t cudaGetValueMismatch();\n\n//-----------------------FUNCTION DEFINITIONS------------------------//\n\nint main(int argc, char *argv[])\n{\n#if defined(__linux__)\n setenv(\"DISPLAY\", \":0\", 0);\n#endif\n\n parseCmdLine(argc, argv);\n atexit(exitHandler);\n\n checkSync(argc, argv);\n return 0;\n}\n\nint parseCmdLine(int argc, char **argv)\n{\n int i;\n for (i = 1; i < argc; i++) {\n if (strcmp(argv[i], \"-cpu\") == 0) {\n useGpu = 0;\n }\n\n if (strcmp(argv[i], \"-h\") == 0) {\n printUsage();\n cleanup(SUCCESS);\n }\n\n if (strcmp(argv[i], \"-width\") == 0) {\n ++i;\n if (i == argc) {\n printf(\"width option must be followed by value\\n\");\n return FAILURE;\n }\n if (sscanf(argv[i], \"%d\", &width) != 1) {\n printf(\"Error: invalid width value\\n\");\n return FAILURE;\n }\n }\n\n if (strcmp(argv[i], \"-height\") == 0) {\n ++i;\n if (i == argc) {\n printf(\"height option must be followed by value\\n\");\n return FAILURE;\n }\n if (sscanf(argv[i], \"%d\", &height) != 1) {\n printf(\"Error: invalid height value\\n\");\n return FAILURE;\n }\n }\n if (strcmp(argv[i], \"-itr\") == 0) {\n ++i;\n if (i == argc) {\n printf(\"itr option must be followed by iteration value\\n\");\n return FAILURE;\n }\n if (sscanf(argv[i], \"%d\", &itr) != 1) {\n printf(\"Error: invalid iteration value\\n\");\n return FAILURE;\n }\n }\n }\n\n return SUCCESS;\n}\n\nvoid printUsage(void)\n{\n printf(\"Usage:\\n\");\n printf(\"\\t-h\\tPrint command line options\\n\");\n printf(\"\\t-cpu\\tSync on the CPU instead of the GPU\\n\");\n printf(\"\\t-width w\\tSet the width to w\\n\");\n printf(\"\\t-height h\\tSet the height to h\\n\");\n printf(\"\\t-itr i\\tSet number of iterations to i\\n\");\n}\n\nvoid checkSync(int argc, char **argv)\n{\n int x, y;\n int bufferSize = width * height * 4;\n unsigned char *pSurf_read = NULL, *pSurf_write = NULL;\n int integrated;\n\n hipError_t status = hipSuccess;\n CUctxCreateParams ctxCreateParams = {};\n\n // Init values for variables\n x = y = 0;\n\n if (hipSuccess != (status = hipInit(0))) {\n printf(\"Failed to initialize CUDA\\n\");\n }\n device = findCudaDeviceDRV(argc, (const char **)argv);\n\n if (hipSuccess != (status = hipCtxCreate(&context, &ctxCreateParams, 0, device))) {\n printf(\"failed to create CUDA context\\n\");\n }\n hipCtxPushCurrent(context);\n\n status = hipDeviceGetAttribute(&integrated, hipDeviceAttributeIntegrated, device);\n if (status != hipSuccess) {\n printf(\"Failed to get device attribute hipDeviceAttributeIntegrated\\n\");\n cleanup(FAILURE);\n }\n\n if (integrated != 1) {\n printf(\"EGLSync_CUDAEvent_Interop does not support dGPU. Waiving sample.\\n\");\n cleanup(WAIVED);\n }\n\n#if (defined(__arm__) || defined(__aarch64__)) && defined(__linux__)\n graphics_setup_window(0, 0, width, height, \"EGLSync_CUDA_Interop\");\n#endif\n\n pSurf_read = (unsigned char *)malloc(bufferSize);\n pSurf_write = (unsigned char *)malloc(bufferSize);\n if (pSurf_read == NULL || pSurf_write == NULL) {\n printf(\"malloc failed\\n\");\n cleanup(FAILURE);\n }\n\n for (x = 0; x < width; x++) {\n for (y = 0; y < height; y++) {\n pSurf_read[(y * width + x) * 4] = 1;\n pSurf_read[(y * width + x) * 4 + 1] = 1;\n pSurf_read[(y * width + x) * 4 + 2] = 1;\n pSurf_read[(y * width + x) * 4 + 3] = 1;\n pSurf_write[(y * width + x) * 4] = 0;\n pSurf_write[(y * width + x) * 4 + 1] = 0;\n pSurf_write[(y * width + x) * 4 + 2] = 0;\n pSurf_write[(y * width + x) * 4 + 3] = 0;\n }\n }\n\n // NOP call to error-check the above glut calls\n GL_SAFE_CALL({});\n\n // Init texture\n GL_SAFE_CALL(glGenTextures(2, tex));\n\n GL_SAFE_CALL(glBindTexture(GL_TEXTURE_2D, tex[GL_READ]));\n GL_SAFE_CALL(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST));\n GL_SAFE_CALL(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST));\n GL_SAFE_CALL(glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, width, height, 0, GL_RGBA, GL_UNSIGNED_BYTE, pSurf_read));\n GL_SAFE_CALL(glBindTexture(GL_TEXTURE_2D, 0));\n\n GL_SAFE_CALL(glBindTexture(GL_TEXTURE_2D, tex[GL_WRITE]));\n GL_SAFE_CALL(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST));\n GL_SAFE_CALL(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST));\n GL_SAFE_CALL(glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, width, height, 0, GL_RGBA, GL_UNSIGNED_BYTE, pSurf_write));\n GL_SAFE_CALL(glBindTexture(GL_TEXTURE_2D, 0));\n\n glFinish();\n\n EGLDisplay eglDisplayHandle = eglGetCurrentDisplay();\n if (eglDisplayHandle == EGL_NO_DISPLAY) {\n printf(\"eglDisplayHandle failed \\n\");\n cleanup(FAILURE);\n }\n else {\n printf(\"eglDisplay Handle created \\n\");\n }\n\n if (!eglSetupExtensions()) {\n printf(\"SetupExtentions failed \\n\");\n cleanup(FAILURE);\n }\n\n EGLContext eglCtx = eglGetCurrentContext();\n if (eglCtx == EGL_NO_CONTEXT) {\n printf(\"Context1 create failed with error %d\\n\", eglGetError());\n cleanup(FAILURE);\n }\n\n // Create the EGL_Image\n EGLint eglImgAttrs[] = {EGL_IMAGE_PRESERVED_KHR, EGL_TRUE, EGL_NONE, EGL_NONE};\n\n EGLImageKHR eglImage1 = eglCreateImageKHR(\n eglDisplayHandle, eglCtx, EGL_GL_TEXTURE_2D_KHR, (EGLClientBuffer)(intptr_t)tex[GL_READ], eglImgAttrs);\n if (eglImage1 == EGL_NO_IMAGE_KHR) {\n printf(\"EGLImage create failed for read texture with error %d\\n\", eglGetError());\n cleanup(FAILURE);\n }\n else {\n printf(\"EGLImage1 created \\n\");\n }\n\n EGLImageKHR eglImage2 = eglCreateImageKHR(\n eglDisplayHandle, eglCtx, EGL_GL_TEXTURE_2D_KHR, (EGLClientBuffer)(intptr_t)tex[GL_WRITE], eglImgAttrs);\n if (eglImage2 == EGL_NO_IMAGE_KHR) {\n printf(\"EGLImage create failed for write texture with error %d\\n\", eglGetError());\n cleanup(FAILURE);\n }\n else {\n printf(\"EGLImage2 created \\n\");\n }\n\n glFinish();\n\n status = cuGraphicsEGLRegisterImage(&writeResource, eglImage1, hipGraphicsRegisterFlagsNone);\n if (status != hipSuccess) {\n printf(\"cuGraphicsEGLRegisterImage failed with Texture 1\\n\");\n cleanup(FAILURE);\n }\n else {\n printf(\"cuGraphicsEGLRegisterImage Passed, writeResource created with texture \"\n \"1\\n\");\n }\n\n status = hipGraphicsSubResourceGetMappedArray(&writeArray, writeResource, 0, 0);\n if (status != hipSuccess) {\n printf(\"hipGraphicsSubResourceGetMappedArray failed for writeResource with \"\n \"texture 1\\n\");\n cleanup(FAILURE);\n }\n\n status = cuGraphicsEGLRegisterImage(&readResource, eglImage2, hipGraphicsRegisterFlagsNone);\n if (status != hipSuccess) {\n printf(\"cuGraphicsEGLRegisterImage failed for readResource with Texture 2\\n\");\n cleanup(FAILURE);\n }\n else {\n printf(\"cuGraphicsEGLRegisterImage Passed, readResource created with texture \"\n \"2\\n\");\n }\n\n status = hipGraphicsSubResourceGetMappedArray(&readArray, readResource, 0, 0);\n if (status != hipSuccess) {\n printf(\"hipGraphicsSubResourceGetMappedArray failed for texture 2\\n\");\n cleanup(FAILURE);\n }\n\n if (useGpu) {\n printf(\"Using GPU Sync path\\n\");\n checkSyncOnGPU(eglDisplayHandle);\n }\n else {\n printf(\"Using CPU Sync path\\n\");\n checkSyncOnCPU();\n }\n\n free(pSurf_read);\n free(pSurf_write);\n cleanup(SUCCESS);\n}\n\nvoid checkSyncOnCPU(void)\n{\n int z = 0;\n unsigned char expectedData, newData;\n hipError_t status = hipSuccess;\n HIP_RESOURCE_DESC wdsc, rdsc;\n memset(&wdsc, 0, sizeof(wdsc));\n memset(&rdsc, 0, sizeof(rdsc));\n\n expectedData = 0;\n newData = 1;\n\n wdsc.resType = HIP_RESOURCE_TYPE_ARRAY;\n wdsc.res.array.hArray = writeArray;\n hipSurfaceObject_t writeSurface;\n rdsc.resType = HIP_RESOURCE_TYPE_ARRAY;\n rdsc.res.array.hArray = readArray;\n hipSurfaceObject_t readSurface;\n\n status = cuSurfObjectCreate(&writeSurface, &wdsc);\n if (status != hipSuccess) {\n printf(\"Surface bounding failed with status %d\\n\", status);\n cleanup(FAILURE);\n }\n status = cuSurfObjectCreate(&readSurface, &rdsc);\n if (status != hipSuccess) {\n printf(\"Surface bounding failed\\n\");\n cleanup(FAILURE);\n }\n\n for (z = 0; z < itr; z++) {\n // GL call to copy from read texture to write texture\n GL_SAFE_CALL(glCopyImageSubData(\n tex[GL_READ], GL_TEXTURE_2D, 0, 0, 0, 0, tex[GL_WRITE], GL_TEXTURE_2D, 0, 0, 0, 0, width, height, 1));\n\n glFinish();\n\n newData++;\n expectedData++;\n\n verify_and_update_kernel<<<(width * height) / 256, 256>>>(\n writeSurface, readSurface, expectedData, newData, width, height);\n\n status = hipCtxSynchronize();\n if (status != hipSuccess) {\n printf(\"hipCtxSynchronize failed \\n\");\n }\n }\n\n hipError_t err = cudaGetValueMismatch();\n if (err != hipSuccess) {\n printf(\"Value mismatch seen when using CPU sync\\n\");\n cleanup(FAILURE);\n }\n\n // Clean up CUDA writeResource\n status = hipGraphicsUnregisterResource(writeResource);\n if (status != hipSuccess) {\n printf(\"Failed to unregister %d\", status);\n cleanup(FAILURE);\n }\n else {\n printf(\"Unregistered writeResource. \\n\");\n }\n\n // Clean up CUDA readResource\n status = hipGraphicsUnregisterResource(readResource);\n if (status != hipSuccess) {\n printf(\"Failed to unregister %d\", status);\n cleanup(FAILURE);\n }\n else {\n printf(\"Unregistered readResource. \\n\");\n }\n}\n\n/*\n Performs same function as checkSyncOnCPU\n Here instead of glFinish() and cuCtxSynchronize like in checkSyncOnCPU,\n we make use of EGLSync, CUDA Event and cuStreamWaitEvent, eglWaitSyncKHR to\n achieve the synchronization due to this CPU is not blocked for any\n synchronization needed between GL-EGL & CUDA operations all synchronizations\n happens on the GPU only.\n*/\nvoid checkSyncOnGPU(EGLDisplay dpy)\n{\n int z = 0;\n unsigned char expectedData, newData;\n hipError_t err;\n hipError_t status = hipSuccess;\n hipStream_t stream;\n hipEvent_t timingDisabledEvent;\n HIP_RESOURCE_DESC wdsc, rdsc;\n memset(&wdsc, 0, sizeof(wdsc));\n memset(&rdsc, 0, sizeof(rdsc));\n\n expectedData = 0;\n newData = 1;\n\n wdsc.resType = HIP_RESOURCE_TYPE_ARRAY;\n wdsc.res.array.hArray = writeArray;\n hipSurfaceObject_t writeSurface;\n rdsc.resType = HIP_RESOURCE_TYPE_ARRAY;\n rdsc.res.array.hArray = readArray;\n hipSurfaceObject_t readSurface;\n\n status = cuSurfObjectCreate(&writeSurface, &wdsc);\n if (status != hipSuccess) {\n printf(\"Surface bounding failed with status %d\\n\", status);\n cleanup(FAILURE);\n }\n status = cuSurfObjectCreate(&readSurface, &rdsc);\n if (status != hipSuccess) {\n printf(\"Surface bounding failed\\n\");\n cleanup(FAILURE);\n }\n\n status = hipStreamCreateWithFlags(&stream, hipStreamDefault);\n if (status != hipSuccess) {\n printf(\"Stream creation failed\\n\");\n cleanup(FAILURE);\n }\n\n // Creates timing disabled event which uses non-blocking synchronization\n status = hipEventCreateWithFlags(&timingDisabledEvent, hipEventDisableTiming);\n if (status != hipSuccess) {\n printf(\"Default event creation failed\\n\");\n cleanup(FAILURE);\n }\n\n /*\n 1. We perform texture-to-texture copy in GLES which is async function\n 2. Followed by creating EGLSync and a CUDA Event from that EGLSync object\n 3. Using cuStreamWaitEvent() we wait in GPU for the GLES to finish texture\n copy.\n 4. CUDA kernel verfiy_and_update_kernel verifies if the copied data by\n GLES is correct, and it updates the buffer with new values.\n 5. This is followed by eglWaitSyncKHR() which waits for the cuda kernel to\n finish, so that in the next iteration GLES can perform the copying of the\n updated buffer to write texture,\n */\n for (z = 0; z < itr; z++) {\n // GL call to copy from read texture to write texture\n GL_SAFE_CALL(glCopyImageSubData(\n tex[GL_READ], GL_TEXTURE_2D, 0, 0, 0, 0, tex[GL_WRITE], GL_TEXTURE_2D, 0, 0, 0, 0, width, height, 1));\n\n EGLSyncKHR eglSyncForGL, eglSyncForCuda;\n EGLBoolean egl_status = EGL_TRUE;\n EGLAttribKHR eglattrib[] = {EGL_CUDA_EVENT_HANDLE_NV, (EGLAttrib)timingDisabledEvent, EGL_NONE};\n\n hipEvent_t cudaEGLSyncEvent;\n\n eglSyncForGL = eglCreateSyncKHR(dpy, EGL_SYNC_FENCE_KHR, NULL);\n\n if (eglSyncForGL == EGL_NO_SYNC_KHR) {\n printf(\" EGL Sync creation failed\\n\");\n cleanup(FAILURE);\n }\n\n status = cuEventCreateFromEGLSync(&cudaEGLSyncEvent, eglSyncForGL, hipEventDefault);\n if (status != hipSuccess) {\n printf(\"CUDA event creation from EGLSync failed\\n\");\n cleanup(FAILURE);\n }\n\n // We wait from CUDA in GPU for GL-EGL operation completion\n status = hipStreamWaitEvent(stream, cudaEGLSyncEvent, 0);\n if (status != hipSuccess) {\n printf(\"Stream wait for event created from EGLSync failed\\n\");\n cleanup(FAILURE);\n }\n\n egl_status = eglDestroySyncKHR(dpy, eglSyncForGL);\n if (egl_status != EGL_TRUE) {\n printf(\"EGL sync object destruction failed\\n\");\n cleanup(FAILURE);\n }\n\n newData++;\n expectedData++;\n\n // Verifies the values in readSurface which is copied by\n // glCopyImageSubData() And writes value of newData into writeSurface\n verify_and_update_kernel<<<(width * height) / 256, 256, 0, stream>>>(\n writeSurface, readSurface, expectedData, newData, width, height);\n\n status = hipEventDestroy(cudaEGLSyncEvent);\n if (status != hipSuccess) {\n printf(\"Event Destroy failed\\n\");\n cleanup(FAILURE);\n }\n\n status = hipEventRecord(timingDisabledEvent, stream);\n if (status != hipSuccess) {\n printf(\"Event Record failed\\n\");\n cleanup(FAILURE);\n }\n\n // creating an EGL sync object linked to a CUDA event object\n eglSyncForCuda = eglCreateSync64KHR(dpy, EGL_SYNC_CUDA_EVENT_NV, eglattrib);\n\n // We wait from EGL for CUDA operation completion\n egl_status = eglWaitSyncKHR(dpy, eglSyncForCuda, 0);\n if (egl_status != EGL_TRUE) {\n printf(\"eglWaitSyncKHR failed\\n\");\n cleanup(FAILURE);\n }\n egl_status = eglDestroySyncKHR(dpy, eglSyncForCuda);\n if (egl_status != EGL_TRUE) {\n printf(\"EGL sync object destruction failed\\n\");\n cleanup(FAILURE);\n }\n }\n\n err = cudaGetValueMismatch();\n if (err != hipSuccess) {\n printf(\"Value mismatch seen when using GPU sync\\n\");\n cleanup(FAILURE);\n }\n\n // Clean up CUDA writeResource\n status = hipGraphicsUnregisterResource(writeResource);\n if (status != hipSuccess) {\n printf(\"Failed to unregister %d\", status);\n cleanup(FAILURE);\n }\n else {\n printf(\"Unregistered writeResource. \\n\");\n }\n\n // Clean up CUDA readResource\n status = hipGraphicsUnregisterResource(readResource);\n if (status != hipSuccess) {\n printf(\"Failed to unregister %d\", status);\n cleanup(FAILURE);\n }\n else {\n printf(\"Unregistered readResource. \\n\");\n }\n}\n\n// Verifies the values in readSurface whether they are expected ones\n// And writes value of newData into writeSurface\n__global__ void\nverify_and_update_kernel(hipSurfaceObject_t write, hipSurfaceObject_t read, char expected, char newval, int width, int height)\n{\n unsigned int x = blockDim.x * blockIdx.x + threadIdx.x;\n unsigned int y = blockDim.y * blockIdx.y + threadIdx.y;\n\n if (x < width && y < height) {\n uchar4 check;\n surf2Dread(&check, read, x * 4, y);\n if (check.x != expected || check.y != expected || check.z != expected || check.w != expected) {\n printf(\"Mismatch found in values read[0]= %u read[1]= %u read[2]= %u \"\n \"read[3]= %u expected is %u\\n\",\n check.x,\n check.y,\n check.z,\n check.w,\n expected);\n numErrors++;\n return;\n }\n uchar4 data = make_uchar4(newval, newval, newval, newval);\n surf2Dwrite(data, write, x * 4, y);\n }\n}\n\n__global__ void getNumErrors(int *numErr) { *numErr = numErrors; }\n\nextern \"C\" hipError_t cudaGetValueMismatch()\n{\n int numErr_h;\n int *numErr_d = NULL;\n hipError_t err = hipSuccess;\n\n err = hipMalloc(&numErr_d, sizeof(int));\n if (err != hipSuccess) {\n printf(\"Cuda Main: hipMemcpy failed with %s\\n\", hipGetErrorString(err));\n hipFree(numErr_d);\n return err;\n }\n\n getNumErrors<<<1, 1>>>(numErr_d);\n err = hipDeviceSynchronize();\n if (err != hipSuccess) {\n printf(\"Cuda Main: hipDeviceSynchronize failed with %s\\n\", hipGetErrorString(err));\n }\n err = hipMemcpy(&numErr_h, numErr_d, sizeof(int), hipMemcpyDeviceToHost);\n if (err != hipSuccess) {\n printf(\"Cuda Main: hipMemcpy failed with %s\\n\", hipGetErrorString(err));\n hipFree(numErr_d);\n return err;\n }\n err = hipFree(numErr_d);\n if (err != hipSuccess) {\n printf(\"Cuda Main: hipFree failed with %s\\n\", hipGetErrorString(err));\n return err;\n }\n if (numErr_h > 0) {\n return hipErrorUnknown;\n }\n return hipSuccess;\n}\n\n// Clean up state and exit. If status is SUCCESS, regression success is printed\n// to stdout. This will happen if the glut timer is triggered. If status is\n// anything else, the regression failure message is printed.\nvoid cleanup(int status)\n{\n GLenum glErr = GL_NO_ERROR;\n hipError_t cudaErr = hipSuccess;\n int exitStatus = status;\n\n // Clean up GL\n if (*tex) {\n GL_SAFE_CALL_NO_CLEANUP(glDeleteTextures(2, tex), glErr);\n }\n\n // Print test status and exit\n if (glErr != GL_NO_ERROR || cudaErr != hipSuccess)\n exitStatus = FAILURE;\n\n printStatus(exitStatus);\n\n cleanExit = 1;\n\n graphics_close_window();\n\n if (exitStatus == FAILURE)\n exit(EXIT_FAILURE);\n\n if (exitStatus == WAIVED)\n exit(EXIT_WAIVED);\n\n exit(0);\n}\n\nvoid exitHandler(void)\n{\n if (!cleanExit) {\n printf(\"&&&& EGLSync_CUDAEvent_Interop unexpected failure \\n\");\n printStatus(FAILURE);\n }\n}\n\n// Print test success or fail for regression testing\nvoid printStatus(int status)\n{\n switch (status) {\n case SUCCESS:\n printf(\"&&&& EGLSync_CUDAEvent_Interop PASSED\\n\");\n break;\n case WAIVED:\n printf(\"&&&& EGLSync_CUDAEvent_Interop WAIVED\\n\");\n break;\n default:\n printf(\"&&&& EGLSync_CUDAEvent_Interop FAILED\\n\");\n break;\n }\n fflush(stdout);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/8_Platform_Specific/Tegra/EGLSync_CUDAEvent_Interop/EGLSync_CUDAEvent_Interop.cu"} {"input": "/* Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n#include \n#include \n\nstatic __global__ void flipSurfaceBits(cudaSurfaceObject_t surfObj, int width, int height)\n{\n char data;\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n if (x < width && y < height) {\n // Read from input surface\n surf2Dread(&data, surfObj, x, y);\n // Write to output surface\n data = ~data;\n surf2Dwrite(data, surfObj, x, y);\n }\n}\n\n// Copy cudaArray to surface memory and launch the CUDA kernel\nvoid launchFlipSurfaceBitsKernel(cudaArray_t *levelArray,\n int32_t *multiPlanarWidth,\n int32_t *multiPlanarHeight,\n int numPlanes)\n{\n\n cudaSurfaceObject_t surfObject[numPlanes] = {0};\n cudaResourceDesc resDesc;\n\n for (int i = 0; i < numPlanes; i++) {\n memset(&resDesc, 0, sizeof(resDesc));\n resDesc.resType = cudaResourceTypeArray;\n resDesc.res.array.array = levelArray[i];\n checkCudaErrors(cudaCreateSurfaceObject(&surfObject[i], &resDesc));\n dim3 threadsperBlock(16, 16);\n dim3 numBlocks((multiPlanarWidth[i] + threadsperBlock.x - 1) / threadsperBlock.x,\n (multiPlanarHeight[i] + threadsperBlock.y - 1) / threadsperBlock.y);\n flipSurfaceBits<<>>(surfObject[i], multiPlanarWidth[i], multiPlanarHeight[i]);\n }\n}\n", "output": "/* Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n#include \n#include \n\nstatic __global__ void flipSurfaceBits(hipSurfaceObject_t surfObj, int width, int height)\n{\n char data;\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n if (x < width && y < height) {\n // Read from input surface\n surf2Dread(&data, surfObj, x, y);\n // Write to output surface\n data = ~data;\n surf2Dwrite(data, surfObj, x, y);\n }\n}\n\n// Copy cudaArray to surface memory and launch the CUDA kernel\nvoid launchFlipSurfaceBitsKernel(hipArray_t *levelArray,\n int32_t *multiPlanarWidth,\n int32_t *multiPlanarHeight,\n int numPlanes)\n{\n\n hipSurfaceObject_t surfObject[numPlanes] = {0};\n hipResourceDesc resDesc;\n\n for (int i = 0; i < numPlanes; i++) {\n memset(&resDesc, 0, sizeof(resDesc));\n resDesc.resType = hipResourceTypeArray;\n resDesc.res.array.array = levelArray[i];\n checkCudaErrors(hipCreateSurfaceObject(&surfObject[i], &resDesc));\n dim3 threadsperBlock(16, 16);\n dim3 numBlocks((multiPlanarWidth[i] + threadsperBlock.x - 1) / threadsperBlock.x,\n (multiPlanarHeight[i] + threadsperBlock.y - 1) / threadsperBlock.y);\n flipSurfaceBits<<>>(surfObject[i], multiPlanarWidth[i], multiPlanarHeight[i]);\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/8_Platform_Specific/Tegra/cudaNvSciBufMultiplanar/imageKernels.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n This example demonstrates how to use the CUDA C bindings to OpenGL ES to\n dynamically modify a vertex buffer using a CUDA C kernel.\n\n The steps are:\n 1. Create an empty vertex buffer object (VBO)\n 2. Register the VBO with CUDA C\n 3. Map the VBO for writing from CUDA C\n 4. Run CUDA C kernel to modify the vertex positions\n 5. Unmap the VBO\n 6. Render the results using OpenGL ES\n\n Host code\n*/\n\n// includes, system\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nvoid error_exit(const char *format, ...)\n{\n va_list args;\n va_start(args, format);\n vfprintf(stderr, format, args);\n va_end(args);\n exit(1);\n}\n\n#include \"graphics_interface.c\"\n\n#ifdef _WIN32\n#define WINDOWS_LEAN_AND_MEAN\n#define NOMINMAX\n#include \n#endif\n\n// includes, cuda\n#include \n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n// #include // helper functions for CUDA/GL interop\n\n#include \n\n#define MAX_EPSILON_ERROR 0.0f\n#define THRESHOLD 0.0f\n#define REFRESH_DELAY 1 // ms\n\n#define GUI_IDLE 0x100\n#define GUI_ROTATE 0x101\n#define GUI_TRANSLATE 0x102\n\nint gui_mode;\n\n////////////////////////////////////////////////////////////////////////////////\n// constants\nconst unsigned int window_width = 512;\nconst unsigned int window_height = 512;\n\nconst unsigned int mesh_width = 256;\nconst unsigned int mesh_height = 256;\n\n// OpenGL ES variables and interop with CUDA C\nGLuint mesh_vao, mesh_vbo;\nstruct cudaGraphicsResource *cuda_vbo_resource;\nvoid *d_vbo_buffer = NULL;\n\nfloat g_fAnim = 0.0;\n\n// UI / mouse controls\nint mouse_old_x, mouse_old_y;\nint mouse_buttons = 0;\nfloat rotate_x = 0.0, rotate_y = 0.0;\nfloat translate_z = -3.0;\n\nStopWatchInterface *timer = NULL;\n\n// Frame statistics\nint frame;\nint fpsCount = 0; // FPS count for averaging\nint fpsLimit = 1; // FPS limit for sampling\nint g_Index = 0;\nfloat avgFPS = 0.0f;\nunsigned int frameCount = 0;\nunsigned int g_TotalErrors = 0;\n\n// Auto-Verification Code\nbool g_bQAReadback = false;\n\nint *pArgc = NULL;\nchar **pArgv = NULL;\n\n#define MAX(a, b) ((a > b) ? a : b)\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\n\n// CUDA functionality\nvoid runCuda(struct cudaGraphicsResource **vbo_resource);\nvoid runAutoTest(int devID, char **argv, char *ref_file);\nvoid checkResultCuda(int argc, char **argv, const GLuint &vbo);\n\nconst char *sSDKsample = \"simpleGLES (VBO)\";\n\nvoid computeFPS()\n{\n frameCount++;\n fpsCount++;\n\n if (fpsCount == fpsLimit) {\n avgFPS = 1.f / (sdkGetAverageTimerValue(&timer) / 1000.f);\n fpsCount = 0;\n fpsLimit = (int)MAX(avgFPS, 1.f);\n\n sdkResetTimer(&timer);\n }\n\n char fps[256];\n sprintf(fps, \"Cuda/OpenGL ES Interop (VBO): %3.1f fps (Max 1000 fps)\", avgFPS);\n graphics_set_windowtitle(fps);\n}\n\n///////////////////////////////////////////////////////////////////////////////\n//! Simple kernel to modify vertex positions in sine wave pattern\n//! @param data data in global memory\n///////////////////////////////////////////////////////////////////////////////\n__global__ void simple_vbo_kernel(float4 *pos, unsigned int width, unsigned int height, float time)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // calculate uv coordinates\n float u = x / (float)width;\n float v = y / (float)height;\n u = u * 2.0f - 1.0f;\n v = v * 2.0f - 1.0f;\n\n // calculate simple sine wave pattern\n float freq = 4.0f;\n float w = sinf(u * freq + time) * cosf(v * freq + time) * 0.5f;\n\n // write output vertex\n pos[y * width + x] = make_float4(u, w, v, 1.0f);\n}\n\nvoid launch_kernel(float4 *pos, unsigned int mesh_width, unsigned int mesh_height, float time)\n{\n // execute the kernel\n dim3 block(8, 8, 1);\n dim3 grid(mesh_width / block.x, mesh_height / block.y, 1);\n simple_vbo_kernel<<>>(pos, mesh_width, mesh_height, time);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run the Cuda part of the computation\n////////////////////////////////////////////////////////////////////////////////\nvoid runCuda(struct cudaGraphicsResource **vbo_resource)\n{\n // map OpenGL buffer object for writing from CUDA\n float4 *dptr;\n cudaGraphicsMapResources(1, vbo_resource, 0);\n size_t num_bytes;\n cudaGraphicsResourceGetMappedPointer((void **)&dptr, &num_bytes, *vbo_resource);\n // printf(\"Sample CUDA mapped VBO: May access %ld bytes\\n\", num_bytes);\n\n // execute the kernel\n // dim3 block(8, 8, 1);\n // dim3 grid(mesh_width / block.x, mesh_height / block.y, 1);\n // kernel<<< grid, block>>>(dptr, mesh_width, mesh_height, g_fAnim);\n\n launch_kernel(dptr, mesh_width, mesh_height, g_fAnim);\n\n // unmap buffer object\n cudaGraphicsUnmapResources(1, vbo_resource, 0);\n}\n\n#ifdef _WIN32\n#ifndef FOPEN\n#define FOPEN(fHandle, filename, mode) fopen_s(&fHandle, filename, mode)\n#endif\n#else\n#ifndef FOPEN\n#define FOPEN(fHandle, filename, mode) (fHandle = fopen(filename, mode))\n#endif\n#endif\n\nvoid sdkDumpBin2(void *data, unsigned int bytes, const char *filename)\n{\n printf(\"sdkDumpBin: <%s>\\n\", filename);\n FILE *fp;\n FOPEN(fp, filename, \"wb\");\n fwrite(data, bytes, 1, fp);\n fflush(fp);\n fclose(fp);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run the Cuda part of the computation\n////////////////////////////////////////////////////////////////////////////////\nvoid runAutoTest(int devID, char **argv, char *ref_file)\n{\n char *reference_file = NULL;\n void *imageData = malloc(mesh_width * mesh_height * sizeof(float));\n\n // execute the kernel\n launch_kernel((float4 *)d_vbo_buffer, mesh_width, mesh_height, g_fAnim);\n\n cudaDeviceSynchronize();\n getLastCudaError(\"launch_kernel failed\");\n\n cudaMemcpy(imageData, d_vbo_buffer, mesh_width * mesh_height * sizeof(float), cudaMemcpyDeviceToHost);\n\n sdkDumpBin2(imageData, mesh_width * mesh_height * sizeof(float), \"simpleGL.bin\");\n reference_file = sdkFindFilePath(ref_file, argv[0]);\n\n if (reference_file\n && !sdkCompareBin2BinFloat(\"simpleGL.bin\",\n reference_file,\n mesh_width * mesh_height * sizeof(float),\n MAX_EPSILON_ERROR,\n THRESHOLD,\n pArgv[0])) {\n g_TotalErrors++;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Display callback\n////////////////////////////////////////////////////////////////////////////////\nvoid display_thisframe(float time_delta)\n{\n sdkStartTimer(&timer);\n\n // run CUDA kernel to generate vertex positions\n runCuda(&cuda_vbo_resource);\n\n glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);\n // GET_GLERROR(0);\n\n // set view matrix: broken, it doesn't work in OpenGL ES! Must put into shader\n // glMatrixMode(GL_MODELVIEW);\n // glLoadIdentity();\n // glTranslatef(0.0, 0.0, translate_z);\n // glRotatef(rotate_x, 1.0, 0.0, 0.0);\n // glRotatef(rotate_y, 0.0, 1.0, 0.0);\n\n glDrawArrays(GL_POINTS, 0, mesh_width * mesh_height);\n\n // GET_GLERROR(0);\n glFinish();\n // GET_GLERROR(0);\n\n g_fAnim += time_delta;\n\n sdkStopTimer(&timer);\n computeFPS();\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Check if the result is correct or write data to file for external\n//! regression testing\n////////////////////////////////////////////////////////////////////////////////\nvoid checkResultCuda(int argc, char **argv, const GLuint &vbo)\n{\n if (!d_vbo_buffer) {\n printf(\"%s: Mapping result buffer from OpenGL ES\\n\", __FUNCTION__);\n\n cudaGraphicsUnregisterResource(cuda_vbo_resource);\n\n // map buffer object\n glBindBuffer(GL_ARRAY_BUFFER, vbo);\n float *data =\n (float *)glMapBufferRange(GL_ARRAY_BUFFER, 0, mesh_width * mesh_height * 4 * sizeof(float), GL_READ_ONLY);\n\n // check result\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // write file for regression test\n sdkWriteFile(\"./data/regression.dat\", data, mesh_width * mesh_height * 3, 0.0, false);\n }\n\n // unmap GL buffer object\n if (!glUnmapBuffer(GL_ARRAY_BUFFER)) {\n fprintf(stderr, \"Unmap buffer failed.\\n\");\n fflush(stderr);\n }\n\n checkCudaErrors(cudaGraphicsGLRegisterBuffer(&cuda_vbo_resource, vbo, cudaGraphicsMapFlagsWriteDiscard));\n\n GET_GLERROR(0);\n }\n}\n\nGLuint mesh_shader = 0;\n\nvoid readAndCompileShaderFromGLSLFile(GLuint new_shaderprogram, const char *filename, GLenum shaderType)\n{\n FILE *file = fopen(filename, \"rb\"); // open shader text file\n if (!file)\n error_exit(\"Filename %s does not exist\\n\", filename);\n\n /* get the size of the file and read it */\n fseek(file, 0, SEEK_END);\n GLint size = ftell(file);\n char *data = (char *)malloc(sizeof(char) * (size + 1));\n memset(data, 0, sizeof(char) * (size + 1));\n fseek(file, 0, SEEK_SET);\n size_t res = fread(data, 1, size, file);\n fclose(file);\n\n GLuint shader = glCreateShader(shaderType);\n glShaderSource(shader, 1, (const GLchar **)&data, &size);\n glCompileShader(shader);\n\n GET_GLERROR(0);\n GLint compile_success = 0;\n glGetShaderiv(shader, GL_COMPILE_STATUS, &compile_success);\n GET_GLERROR(0);\n\n if (compile_success == GL_FALSE) {\n printf(\"Compilation of %s failed!\\n Reason:\\n\", filename);\n\n GLint maxLength = 0;\n glGetShaderiv(shader, GL_INFO_LOG_LENGTH, &maxLength);\n\n char errorLog[maxLength];\n glGetShaderInfoLog(shader, maxLength, &maxLength, &errorLog[0]);\n\n printf(\"%s\", errorLog);\n\n glDeleteShader(shader);\n exit(1);\n }\n\n glAttachShader(new_shaderprogram, shader);\n glDeleteShader(shader); // good to do?\n\n free(data);\n}\n\nGLuint ShaderCreate(const char *vshader_filename, const char *fshader_filename)\n{\n printf(\"Loading GLSL shaders %s %s\\n\", vshader_filename, fshader_filename);\n\n GLuint new_shaderprogram = glCreateProgram();\n\n GET_GLERROR(0);\n if (vshader_filename)\n readAndCompileShaderFromGLSLFile(new_shaderprogram, vshader_filename, GL_VERTEX_SHADER);\n\n GET_GLERROR(0);\n if (fshader_filename)\n readAndCompileShaderFromGLSLFile(new_shaderprogram, fshader_filename, GL_FRAGMENT_SHADER);\n\n GET_GLERROR(0);\n\n glLinkProgram(new_shaderprogram);\n\n GET_GLERROR(0);\n GLint link_success;\n glGetProgramiv(new_shaderprogram, GL_LINK_STATUS, &link_success);\n\n if (link_success == GL_FALSE) {\n printf(\"Linking of %s with %s failed!\\n Reason:\\n\", vshader_filename, fshader_filename);\n\n GLint maxLength = 0;\n glGetShaderiv(new_shaderprogram, GL_INFO_LOG_LENGTH, &maxLength);\n\n char errorLog[maxLength];\n glGetShaderInfoLog(new_shaderprogram, maxLength, &maxLength, &errorLog[0]);\n\n printf(\"%s\", errorLog);\n\n exit(EXIT_FAILURE);\n }\n\n return new_shaderprogram;\n}\n\n//===========================================================================\n// InitGraphicsState() - initialize OpenGL\n//===========================================================================\nstatic void InitGraphicsState(void)\n{\n char *GL_version = (char *)glGetString(GL_VERSION);\n char *GL_vendor = (char *)glGetString(GL_VENDOR);\n char *GL_renderer = (char *)glGetString(GL_RENDERER);\n\n printf(\"Version: %s\\n\", GL_version);\n printf(\"Vendor: %s\\n\", GL_vendor);\n printf(\"Renderer: %s\\n\", GL_renderer);\n\n // RENDERING SETUP (OpenGL ES or OpenGL Core Profile!)\n glGenVertexArrays(1, &mesh_vao); // Features' Vertex Array Object allocation\n glBindVertexArray(mesh_vao); // bind VAO\n\n // initialize buffer object\n glGenBuffers(1, &mesh_vbo);\n glBindBuffer(GL_ARRAY_BUFFER, mesh_vbo);\n\n unsigned int size = mesh_width * mesh_height * 4 * sizeof(float);\n glBufferData(GL_ARRAY_BUFFER, size, NULL, GL_DYNAMIC_DRAW);\n glVertexAttribPointer((GLuint)0, 4, GL_FLOAT, GL_FALSE, 0, 0);\n glEnableVertexAttribArray(0);\n\n checkCudaErrors(cudaGraphicsGLRegisterBuffer(&cuda_vbo_resource, mesh_vbo, cudaGraphicsMapFlagsNone));\n // glBindVertexArray(0); // keep above Vertex Array Object bound (it's the\n // only one throughout)\n\n // GLSL stuff\n char *vertex_shader_path = sdkFindFilePath(\"mesh.vert.glsl\", pArgv[0]);\n char *fragment_shader_path = sdkFindFilePath(\"mesh.frag.glsl\", pArgv[0]);\n\n if (vertex_shader_path == NULL || fragment_shader_path == NULL) {\n printf(\"Error finding shader file\\n\");\n exit(EXIT_FAILURE);\n }\n\n mesh_shader = ShaderCreate(vertex_shader_path, fragment_shader_path);\n GET_GLERROR(0);\n\n free(vertex_shader_path);\n free(fragment_shader_path);\n\n glUseProgram(mesh_shader);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nbool runTest(int argc, char **argv, char *ref_file)\n{\n // Create the CUTIL timer\n sdkCreateTimer(&timer);\n\n int devID = 0;\n#if defined(__aarch64__) || defined(__arm__)\n // find iGPU on the system which is compute capable which will perform\n // GLES-CUDA interop\n devID = findIntegratedGPU();\n#else\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n devID = findCudaDevice(argc, (const char **)argv);\n#endif\n\n // command line mode only\n if (ref_file != NULL) {\n // create VBO\n checkCudaErrors(cudaMalloc((void **)&d_vbo_buffer, mesh_width * mesh_height * 4 * sizeof(float)));\n\n // run the cuda part\n runAutoTest(devID, argv, ref_file);\n\n // check result of Cuda step\n checkResultCuda(argc, argv, mesh_vbo);\n\n cudaFree(d_vbo_buffer);\n d_vbo_buffer = NULL;\n }\n else {\n // this would use command-line specified CUDA device, note that CUDA\n // defaults to highest Gflops/s device\n if (checkCmdLineFlag(argc, (const char **)argv, \"device\"))\n error_exit(\"Device setting not yet implemented!\\n\");\n\n // create X11 window and set up associated OpenGL ES context\n graphics_setup_window(0, 0, window_width, window_height, sSDKsample);\n\n InitGraphicsState(); // set up GLES stuff\n\n glClearColor(0, 0.5, 1, 1); // blue-ish background\n glClear(GL_COLOR_BUFFER_BIT);\n\n // printf(\"WP%d\\n\", __LINE__);\n graphics_swap_buffers();\n\n XEvent event;\n KeySym key;\n char text[255];\n\n int frame = 0;\n\n while (frame < 100000) {\n if (XPending(display)) {\n XNextEvent(display, &event);\n\n if (event.type == Expose && event.xexpose.count == 0) {\n printf(\"Redraw requested!\\n\");\n }\n if (event.type == KeyPress && XLookupString(&event.xkey, text, 255, &key, 0) == 1) {\n if (text[0] == 27)\n goto label_stop_x;\n\n printf(\"You pressed the %c key!\\n\", text[0]);\n }\n\n if (event.type == ButtonPress) {\n printf(\n \"Mouse button %d press at (%d,%d)\\n\", event.xbutton.button, event.xbutton.x, event.xbutton.y);\n\n if (event.xbutton.button == Button1)\n gui_mode = GUI_TRANSLATE;\n if (event.xbutton.button == Button3)\n gui_mode = GUI_ROTATE;\n mouse_old_x = event.xbutton.x;\n mouse_old_y = event.xbutton.y;\n }\n\n if (event.type == ButtonRelease) {\n printf(\"Mouse button %d released at (%d,%d)\\n\",\n event.xbutton.button,\n event.xbutton.x,\n event.xbutton.y);\n\n gui_mode = GUI_IDLE;\n mouse_old_x = event.xbutton.x;\n mouse_old_y = event.xbutton.y;\n }\n\n if (event.type == MotionNotify) {\n // printf(\"Mouse motion towards %d %d, GUI mode is 0x%x\\n\",\n //\t event.xmotion.x, event.xmotion.y, gui_mode);\n float dx, dy;\n dx = (float)(event.xmotion.x - mouse_old_x);\n dy = (float)(event.xmotion.y - mouse_old_y);\n\n if (gui_mode == GUI_ROTATE) {\n rotate_x += dy * 0.2f;\n rotate_y += dx * 0.2f;\n printf(\"rot x %f y %f\\n\", rotate_x, rotate_y);\n }\n if (gui_mode == GUI_TRANSLATE) {\n translate_z += dy * 0.01f;\n printf(\"translate z %f\\n\", translate_z);\n }\n\n mouse_old_x = event.xmotion.x;\n mouse_old_y = event.xmotion.y;\n }\n }\n\n display_thisframe(0.010);\n usleep(1000); // need not take full CPU and GPU\n\n graphics_swap_buffers();\n // printf(\"frame %d\\n\",frame++);\n }\n\n label_stop_x:\n // NOTE: Before destroying OpenGL ES context, must unregister all shared\n // resources from CUDA !\n cudaGraphicsUnregisterResource(cuda_vbo_resource);\n\n graphics_close_window(); // close window and destroy OpenGL ES context\n\n sdkDeleteTimer(&timer);\n }\n\n return true;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n char *ref_file = NULL;\n\n pArgc = &argc;\n pArgv = argv;\n\n#if defined(__linux__)\n setenv(\"DISPLAY\", \":0\", 0);\n#endif\n\n printf(\"%s starting...\\n\", sSDKsample);\n\n if (argc > 1) {\n if (checkCmdLineFlag(argc, (const char **)argv, \"file\")) {\n // In this mode, we run without OpenGL and see if VBO is generated\n // correctly\n getCmdLineArgumentString(argc, (const char **)argv, \"file\", (char **)&ref_file);\n }\n }\n\n printf(\"\\n\");\n\n runTest(argc, argv, ref_file);\n\n printf(\"%s completed, returned %s\\n\", sSDKsample, (g_TotalErrors == 0) ? \"OK\" : \"ERROR!\");\n\n exit(g_TotalErrors == 0 ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n This example demonstrates how to use the CUDA C bindings to OpenGL ES to\n dynamically modify a vertex buffer using a CUDA C kernel.\n\n The steps are:\n 1. Create an empty vertex buffer object (VBO)\n 2. Register the VBO with CUDA C\n 3. Map the VBO for writing from CUDA C\n 4. Run CUDA C kernel to modify the vertex positions\n 5. Unmap the VBO\n 6. Render the results using OpenGL ES\n\n Host code\n*/\n\n// includes, system\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nvoid error_exit(const char *format, ...)\n{\n va_list args;\n va_start(args, format);\n vfprintf(stderr, format, args);\n va_end(args);\n exit(1);\n}\n\n#include \"graphics_interface.c\"\n\n#ifdef _WIN32\n#define WINDOWS_LEAN_AND_MEAN\n#define NOMINMAX\n#include \n#endif\n\n// includes, cuda\n#include \n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n// #include // helper functions for CUDA/GL interop\n\n#include \n\n#define MAX_EPSILON_ERROR 0.0f\n#define THRESHOLD 0.0f\n#define REFRESH_DELAY 1 // ms\n\n#define GUI_IDLE 0x100\n#define GUI_ROTATE 0x101\n#define GUI_TRANSLATE 0x102\n\nint gui_mode;\n\n////////////////////////////////////////////////////////////////////////////////\n// constants\nconst unsigned int window_width = 512;\nconst unsigned int window_height = 512;\n\nconst unsigned int mesh_width = 256;\nconst unsigned int mesh_height = 256;\n\n// OpenGL ES variables and interop with CUDA C\nGLuint mesh_vao, mesh_vbo;\nstruct hipGraphicsResource *cuda_vbo_resource;\nvoid *d_vbo_buffer = NULL;\n\nfloat g_fAnim = 0.0;\n\n// UI / mouse controls\nint mouse_old_x, mouse_old_y;\nint mouse_buttons = 0;\nfloat rotate_x = 0.0, rotate_y = 0.0;\nfloat translate_z = -3.0;\n\nStopWatchInterface *timer = NULL;\n\n// Frame statistics\nint frame;\nint fpsCount = 0; // FPS count for averaging\nint fpsLimit = 1; // FPS limit for sampling\nint g_Index = 0;\nfloat avgFPS = 0.0f;\nunsigned int frameCount = 0;\nunsigned int g_TotalErrors = 0;\n\n// Auto-Verification Code\nbool g_bQAReadback = false;\n\nint *pArgc = NULL;\nchar **pArgv = NULL;\n\n#define MAX(a, b) ((a > b) ? a : b)\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\n\n// CUDA functionality\nvoid runCuda(struct hipGraphicsResource **vbo_resource);\nvoid runAutoTest(int devID, char **argv, char *ref_file);\nvoid checkResultCuda(int argc, char **argv, const GLuint &vbo);\n\nconst char *sSDKsample = \"simpleGLES (VBO)\";\n\nvoid computeFPS()\n{\n frameCount++;\n fpsCount++;\n\n if (fpsCount == fpsLimit) {\n avgFPS = 1.f / (sdkGetAverageTimerValue(&timer) / 1000.f);\n fpsCount = 0;\n fpsLimit = (int)MAX(avgFPS, 1.f);\n\n sdkResetTimer(&timer);\n }\n\n char fps[256];\n sprintf(fps, \"Cuda/OpenGL ES Interop (VBO): %3.1f fps (Max 1000 fps)\", avgFPS);\n graphics_set_windowtitle(fps);\n}\n\n///////////////////////////////////////////////////////////////////////////////\n//! Simple kernel to modify vertex positions in sine wave pattern\n//! @param data data in global memory\n///////////////////////////////////////////////////////////////////////////////\n__global__ void simple_vbo_kernel(float4 *pos, unsigned int width, unsigned int height, float time)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // calculate uv coordinates\n float u = x / (float)width;\n float v = y / (float)height;\n u = u * 2.0f - 1.0f;\n v = v * 2.0f - 1.0f;\n\n // calculate simple sine wave pattern\n float freq = 4.0f;\n float w = sinf(u * freq + time) * cosf(v * freq + time) * 0.5f;\n\n // write output vertex\n pos[y * width + x] = make_float4(u, w, v, 1.0f);\n}\n\nvoid launch_kernel(float4 *pos, unsigned int mesh_width, unsigned int mesh_height, float time)\n{\n // execute the kernel\n dim3 block(8, 8, 1);\n dim3 grid(mesh_width / block.x, mesh_height / block.y, 1);\n simple_vbo_kernel<<>>(pos, mesh_width, mesh_height, time);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run the Cuda part of the computation\n////////////////////////////////////////////////////////////////////////////////\nvoid runCuda(struct hipGraphicsResource **vbo_resource)\n{\n // map OpenGL buffer object for writing from CUDA\n float4 *dptr;\n hipGraphicsMapResources(1, vbo_resource, 0);\n size_t num_bytes;\n hipGraphicsResourceGetMappedPointer((void **)&dptr, &num_bytes, *vbo_resource);\n // printf(\"Sample CUDA mapped VBO: May access %ld bytes\\n\", num_bytes);\n\n // execute the kernel\n // dim3 block(8, 8, 1);\n // dim3 grid(mesh_width / block.x, mesh_height / block.y, 1);\n // kernel<<< grid, block>>>(dptr, mesh_width, mesh_height, g_fAnim);\n\n launch_kernel(dptr, mesh_width, mesh_height, g_fAnim);\n\n // unmap buffer object\n hipGraphicsUnmapResources(1, vbo_resource, 0);\n}\n\n#ifdef _WIN32\n#ifndef FOPEN\n#define FOPEN(fHandle, filename, mode) fopen_s(&fHandle, filename, mode)\n#endif\n#else\n#ifndef FOPEN\n#define FOPEN(fHandle, filename, mode) (fHandle = fopen(filename, mode))\n#endif\n#endif\n\nvoid sdkDumpBin2(void *data, unsigned int bytes, const char *filename)\n{\n printf(\"sdkDumpBin: <%s>\\n\", filename);\n FILE *fp;\n FOPEN(fp, filename, \"wb\");\n fwrite(data, bytes, 1, fp);\n fflush(fp);\n fclose(fp);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run the Cuda part of the computation\n////////////////////////////////////////////////////////////////////////////////\nvoid runAutoTest(int devID, char **argv, char *ref_file)\n{\n char *reference_file = NULL;\n void *imageData = malloc(mesh_width * mesh_height * sizeof(float));\n\n // execute the kernel\n launch_kernel((float4 *)d_vbo_buffer, mesh_width, mesh_height, g_fAnim);\n\n hipDeviceSynchronize();\n getLastCudaError(\"launch_kernel failed\");\n\n hipMemcpy(imageData, d_vbo_buffer, mesh_width * mesh_height * sizeof(float), hipMemcpyDeviceToHost);\n\n sdkDumpBin2(imageData, mesh_width * mesh_height * sizeof(float), \"simpleGL.bin\");\n reference_file = sdkFindFilePath(ref_file, argv[0]);\n\n if (reference_file\n && !sdkCompareBin2BinFloat(\"simpleGL.bin\",\n reference_file,\n mesh_width * mesh_height * sizeof(float),\n MAX_EPSILON_ERROR,\n THRESHOLD,\n pArgv[0])) {\n g_TotalErrors++;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Display callback\n////////////////////////////////////////////////////////////////////////////////\nvoid display_thisframe(float time_delta)\n{\n sdkStartTimer(&timer);\n\n // run CUDA kernel to generate vertex positions\n runCuda(&cuda_vbo_resource);\n\n glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);\n // GET_GLERROR(0);\n\n // set view matrix: broken, it doesn't work in OpenGL ES! Must put into shader\n // glMatrixMode(GL_MODELVIEW);\n // glLoadIdentity();\n // glTranslatef(0.0, 0.0, translate_z);\n // glRotatef(rotate_x, 1.0, 0.0, 0.0);\n // glRotatef(rotate_y, 0.0, 1.0, 0.0);\n\n glDrawArrays(GL_POINTS, 0, mesh_width * mesh_height);\n\n // GET_GLERROR(0);\n glFinish();\n // GET_GLERROR(0);\n\n g_fAnim += time_delta;\n\n sdkStopTimer(&timer);\n computeFPS();\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Check if the result is correct or write data to file for external\n//! regression testing\n////////////////////////////////////////////////////////////////////////////////\nvoid checkResultCuda(int argc, char **argv, const GLuint &vbo)\n{\n if (!d_vbo_buffer) {\n printf(\"%s: Mapping result buffer from OpenGL ES\\n\", __FUNCTION__);\n\n hipGraphicsUnregisterResource(cuda_vbo_resource);\n\n // map buffer object\n glBindBuffer(GL_ARRAY_BUFFER, vbo);\n float *data =\n (float *)glMapBufferRange(GL_ARRAY_BUFFER, 0, mesh_width * mesh_height * 4 * sizeof(float), GL_READ_ONLY);\n\n // check result\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // write file for regression test\n sdkWriteFile(\"./data/regression.dat\", data, mesh_width * mesh_height * 3, 0.0, false);\n }\n\n // unmap GL buffer object\n if (!glUnmapBuffer(GL_ARRAY_BUFFER)) {\n fprintf(stderr, \"Unmap buffer failed.\\n\");\n fflush(stderr);\n }\n\n checkCudaErrors(hipGraphicsGLRegisterBuffer(&cuda_vbo_resource, vbo, cudaGraphicsMapFlagsWriteDiscard));\n\n GET_GLERROR(0);\n }\n}\n\nGLuint mesh_shader = 0;\n\nvoid readAndCompileShaderFromGLSLFile(GLuint new_shaderprogram, const char *filename, GLenum shaderType)\n{\n FILE *file = fopen(filename, \"rb\"); // open shader text file\n if (!file)\n error_exit(\"Filename %s does not exist\\n\", filename);\n\n /* get the size of the file and read it */\n fseek(file, 0, SEEK_END);\n GLint size = ftell(file);\n char *data = (char *)malloc(sizeof(char) * (size + 1));\n memset(data, 0, sizeof(char) * (size + 1));\n fseek(file, 0, SEEK_SET);\n size_t res = fread(data, 1, size, file);\n fclose(file);\n\n GLuint shader = glCreateShader(shaderType);\n glShaderSource(shader, 1, (const GLchar **)&data, &size);\n glCompileShader(shader);\n\n GET_GLERROR(0);\n GLint compile_success = 0;\n glGetShaderiv(shader, GL_COMPILE_STATUS, &compile_success);\n GET_GLERROR(0);\n\n if (compile_success == GL_FALSE) {\n printf(\"Compilation of %s failed!\\n Reason:\\n\", filename);\n\n GLint maxLength = 0;\n glGetShaderiv(shader, GL_INFO_LOG_LENGTH, &maxLength);\n\n char errorLog[maxLength];\n glGetShaderInfoLog(shader, maxLength, &maxLength, &errorLog[0]);\n\n printf(\"%s\", errorLog);\n\n glDeleteShader(shader);\n exit(1);\n }\n\n glAttachShader(new_shaderprogram, shader);\n glDeleteShader(shader); // good to do?\n\n free(data);\n}\n\nGLuint ShaderCreate(const char *vshader_filename, const char *fshader_filename)\n{\n printf(\"Loading GLSL shaders %s %s\\n\", vshader_filename, fshader_filename);\n\n GLuint new_shaderprogram = glCreateProgram();\n\n GET_GLERROR(0);\n if (vshader_filename)\n readAndCompileShaderFromGLSLFile(new_shaderprogram, vshader_filename, GL_VERTEX_SHADER);\n\n GET_GLERROR(0);\n if (fshader_filename)\n readAndCompileShaderFromGLSLFile(new_shaderprogram, fshader_filename, GL_FRAGMENT_SHADER);\n\n GET_GLERROR(0);\n\n glLinkProgram(new_shaderprogram);\n\n GET_GLERROR(0);\n GLint link_success;\n glGetProgramiv(new_shaderprogram, GL_LINK_STATUS, &link_success);\n\n if (link_success == GL_FALSE) {\n printf(\"Linking of %s with %s failed!\\n Reason:\\n\", vshader_filename, fshader_filename);\n\n GLint maxLength = 0;\n glGetShaderiv(new_shaderprogram, GL_INFO_LOG_LENGTH, &maxLength);\n\n char errorLog[maxLength];\n glGetShaderInfoLog(new_shaderprogram, maxLength, &maxLength, &errorLog[0]);\n\n printf(\"%s\", errorLog);\n\n exit(EXIT_FAILURE);\n }\n\n return new_shaderprogram;\n}\n\n//===========================================================================\n// InitGraphicsState() - initialize OpenGL\n//===========================================================================\nstatic void InitGraphicsState(void)\n{\n char *GL_version = (char *)glGetString(GL_VERSION);\n char *GL_vendor = (char *)glGetString(GL_VENDOR);\n char *GL_renderer = (char *)glGetString(GL_RENDERER);\n\n printf(\"Version: %s\\n\", GL_version);\n printf(\"Vendor: %s\\n\", GL_vendor);\n printf(\"Renderer: %s\\n\", GL_renderer);\n\n // RENDERING SETUP (OpenGL ES or OpenGL Core Profile!)\n glGenVertexArrays(1, &mesh_vao); // Features' Vertex Array Object allocation\n glBindVertexArray(mesh_vao); // bind VAO\n\n // initialize buffer object\n glGenBuffers(1, &mesh_vbo);\n glBindBuffer(GL_ARRAY_BUFFER, mesh_vbo);\n\n unsigned int size = mesh_width * mesh_height * 4 * sizeof(float);\n glBufferData(GL_ARRAY_BUFFER, size, NULL, GL_DYNAMIC_DRAW);\n glVertexAttribPointer((GLuint)0, 4, GL_FLOAT, GL_FALSE, 0, 0);\n glEnableVertexAttribArray(0);\n\n checkCudaErrors(hipGraphicsGLRegisterBuffer(&cuda_vbo_resource, mesh_vbo, cudaGraphicsMapFlagsNone));\n // glBindVertexArray(0); // keep above Vertex Array Object bound (it's the\n // only one throughout)\n\n // GLSL stuff\n char *vertex_shader_path = sdkFindFilePath(\"mesh.vert.glsl\", pArgv[0]);\n char *fragment_shader_path = sdkFindFilePath(\"mesh.frag.glsl\", pArgv[0]);\n\n if (vertex_shader_path == NULL || fragment_shader_path == NULL) {\n printf(\"Error finding shader file\\n\");\n exit(EXIT_FAILURE);\n }\n\n mesh_shader = ShaderCreate(vertex_shader_path, fragment_shader_path);\n GET_GLERROR(0);\n\n free(vertex_shader_path);\n free(fragment_shader_path);\n\n glUseProgram(mesh_shader);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nbool runTest(int argc, char **argv, char *ref_file)\n{\n // Create the CUTIL timer\n sdkCreateTimer(&timer);\n\n int devID = 0;\n#if defined(__aarch64__) || defined(__arm__)\n // find iGPU on the system which is compute capable which will perform\n // GLES-CUDA interop\n devID = findIntegratedGPU();\n#else\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n devID = findCudaDevice(argc, (const char **)argv);\n#endif\n\n // command line mode only\n if (ref_file != NULL) {\n // create VBO\n checkCudaErrors(hipMalloc((void **)&d_vbo_buffer, mesh_width * mesh_height * 4 * sizeof(float)));\n\n // run the cuda part\n runAutoTest(devID, argv, ref_file);\n\n // check result of Cuda step\n checkResultCuda(argc, argv, mesh_vbo);\n\n hipFree(d_vbo_buffer);\n d_vbo_buffer = NULL;\n }\n else {\n // this would use command-line specified CUDA device, note that CUDA\n // defaults to highest Gflops/s device\n if (checkCmdLineFlag(argc, (const char **)argv, \"device\"))\n error_exit(\"Device setting not yet implemented!\\n\");\n\n // create X11 window and set up associated OpenGL ES context\n graphics_setup_window(0, 0, window_width, window_height, sSDKsample);\n\n InitGraphicsState(); // set up GLES stuff\n\n glClearColor(0, 0.5, 1, 1); // blue-ish background\n glClear(GL_COLOR_BUFFER_BIT);\n\n // printf(\"WP%d\\n\", __LINE__);\n graphics_swap_buffers();\n\n XEvent event;\n KeySym key;\n char text[255];\n\n int frame = 0;\n\n while (frame < 100000) {\n if (XPending(display)) {\n XNextEvent(display, &event);\n\n if (event.type == Expose && event.xexpose.count == 0) {\n printf(\"Redraw requested!\\n\");\n }\n if (event.type == KeyPress && XLookupString(&event.xkey, text, 255, &key, 0) == 1) {\n if (text[0] == 27)\n goto label_stop_x;\n\n printf(\"You pressed the %c key!\\n\", text[0]);\n }\n\n if (event.type == ButtonPress) {\n printf(\n \"Mouse button %d press at (%d,%d)\\n\", event.xbutton.button, event.xbutton.x, event.xbutton.y);\n\n if (event.xbutton.button == Button1)\n gui_mode = GUI_TRANSLATE;\n if (event.xbutton.button == Button3)\n gui_mode = GUI_ROTATE;\n mouse_old_x = event.xbutton.x;\n mouse_old_y = event.xbutton.y;\n }\n\n if (event.type == ButtonRelease) {\n printf(\"Mouse button %d released at (%d,%d)\\n\",\n event.xbutton.button,\n event.xbutton.x,\n event.xbutton.y);\n\n gui_mode = GUI_IDLE;\n mouse_old_x = event.xbutton.x;\n mouse_old_y = event.xbutton.y;\n }\n\n if (event.type == MotionNotify) {\n // printf(\"Mouse motion towards %d %d, GUI mode is 0x%x\\n\",\n //\t event.xmotion.x, event.xmotion.y, gui_mode);\n float dx, dy;\n dx = (float)(event.xmotion.x - mouse_old_x);\n dy = (float)(event.xmotion.y - mouse_old_y);\n\n if (gui_mode == GUI_ROTATE) {\n rotate_x += dy * 0.2f;\n rotate_y += dx * 0.2f;\n printf(\"rot x %f y %f\\n\", rotate_x, rotate_y);\n }\n if (gui_mode == GUI_TRANSLATE) {\n translate_z += dy * 0.01f;\n printf(\"translate z %f\\n\", translate_z);\n }\n\n mouse_old_x = event.xmotion.x;\n mouse_old_y = event.xmotion.y;\n }\n }\n\n display_thisframe(0.010);\n usleep(1000); // need not take full CPU and GPU\n\n graphics_swap_buffers();\n // printf(\"frame %d\\n\",frame++);\n }\n\n label_stop_x:\n // NOTE: Before destroying OpenGL ES context, must unregister all shared\n // resources from CUDA !\n hipGraphicsUnregisterResource(cuda_vbo_resource);\n\n graphics_close_window(); // close window and destroy OpenGL ES context\n\n sdkDeleteTimer(&timer);\n }\n\n return true;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n char *ref_file = NULL;\n\n pArgc = &argc;\n pArgv = argv;\n\n#if defined(__linux__)\n setenv(\"DISPLAY\", \":0\", 0);\n#endif\n\n printf(\"%s starting...\\n\", sSDKsample);\n\n if (argc > 1) {\n if (checkCmdLineFlag(argc, (const char **)argv, \"file\")) {\n // In this mode, we run without OpenGL and see if VBO is generated\n // correctly\n getCmdLineArgumentString(argc, (const char **)argv, \"file\", (char **)&ref_file);\n }\n }\n\n printf(\"\\n\");\n\n runTest(argc, argv, ref_file);\n\n printf(\"%s completed, returned %s\\n\", sSDKsample, (g_TotalErrors == 0) ? \"OK\" : \"ERROR!\");\n\n exit(g_TotalErrors == 0 ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/8_Platform_Specific/Tegra/simpleGLES/simpleGLES.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include // CUDA FFT Libraries\n#include // Helper functions for CUDA Error handling\n#include \n#include \n\n// OpenGL Graphics includes\n#include \n\n// FluidsGLES CUDA kernel definitions\n#include \"fluidsGLES_kernels.cuh\"\n\n// Texture object for reading velocity field\ncudaTextureObject_t texObj;\nstatic cudaArray *array = NULL;\n\n// Particle data\nextern GLuint vbo; // OpenGL vertex buffer object\nextern struct cudaGraphicsResource *cuda_vbo_resource; // handles OpenGL-CUDA exchange\n\n// Texture pitch\nextern size_t tPitch;\nextern cufftHandle planr2c;\nextern cufftHandle planc2r;\ncData *vxfield = NULL;\ncData *vyfield = NULL;\n\nvoid setupTexture(int x, int y)\n{\n cudaChannelFormatDesc desc = cudaCreateChannelDesc();\n\n cudaMallocArray(&array, &desc, y, x);\n getLastCudaError(\"cudaMalloc failed\");\n\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = array;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = cudaFilterModeLinear;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&texObj, &texRes, &texDescr, NULL));\n}\n\nvoid updateTexture(cData *data, size_t wib, size_t h, size_t pitch)\n{\n checkCudaErrors(cudaMemcpy2DToArray(array, 0, 0, data, pitch, wib, h, cudaMemcpyDeviceToDevice));\n}\n\nvoid deleteTexture(void)\n{\n checkCudaErrors(cudaDestroyTextureObject(texObj));\n checkCudaErrors(cudaFreeArray(array));\n}\n\n// Note that these kernels are designed to work with arbitrary\n// domain sizes, not just domains that are multiples of the tile\n// size. Therefore, we have extra code that checks to make sure\n// a given thread location falls within the domain boundaries in\n// both X and Y. Also, the domain is covered by looping over\n// multiple elements in the Y direction, while there is a one-to-one\n// mapping between threads in X and the tile size in X.\n// Nolan Goodnight 9/22/06\n\n// This method adds constant force vectors to the velocity field\n// stored in 'v' according to v(x,t+1) = v(x,t) + dt * f.\n__global__ void addForces_k(cData *v, int dx, int dy, int spx, int spy, float fx, float fy, int r, size_t pitch)\n{\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n cData *fj = (cData *)((char *)v + (ty + spy) * pitch) + tx + spx;\n\n cData vterm = *fj;\n tx -= r;\n ty -= r;\n float s = 1.f / (1.f + tx * tx * tx * tx + ty * ty * ty * ty);\n vterm.x += s * fx;\n vterm.y += s * fy;\n *fj = vterm;\n}\n\n// This method performs the velocity advection step, where we\n// trace velocity vectors back in time to update each grid cell.\n// That is, v(x,t+1) = v(p(x,-dt),t). Here we perform bilinear\n// interpolation in the velocity space.\n__global__ void advectVelocity_k(cData *v,\n float *vx,\n float *vy,\n int dx,\n int pdx,\n int dy,\n float dt,\n int lb,\n cudaTextureObject_t texObject)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n cData vterm, ploc;\n float vxterm, vyterm;\n\n // gtidx is the domain location in x for this thread\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fj = fi * pdx + gtidx;\n vterm = tex2D(texObject, (float)gtidx, (float)fi);\n ploc.x = (gtidx + 0.5f) - (dt * vterm.x * dx);\n ploc.y = (fi + 0.5f) - (dt * vterm.y * dy);\n vterm = tex2D(texObject, ploc.x, ploc.y);\n vxterm = vterm.x;\n vyterm = vterm.y;\n vx[fj] = vxterm;\n vy[fj] = vyterm;\n }\n }\n }\n}\n\n// This method performs velocity diffusion and forces mass conservation\n// in the frequency domain. The inputs 'vx' and 'vy' are complex-valued\n// arrays holding the Fourier coefficients of the velocity field in\n// X and Y. Diffusion in this space takes a simple form described as:\n// v(k,t) = v(k,t) / (1 + visc * dt * k^2), where visc is the viscosity,\n// and k is the wavenumber. The projection step forces the Fourier\n// velocity vectors to be orthogonal to the vectors for each\n// wavenumber: v(k,t) = v(k,t) - ((k dot v(k,t) * k) / k^2.\n__global__ void diffuseProject_k(cData *vx, cData *vy, int dx, int dy, float dt, float visc, int lb)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n cData xterm, yterm;\n\n // gtidx is the domain location in x for this thread\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fj = fi * dx + gtidx;\n xterm = vx[fj];\n yterm = vy[fj];\n\n // Compute the index of the wavenumber based on the\n // data order produced by a standard NN FFT.\n int iix = gtidx;\n int iiy = (fi > dy / 2) ? (fi - (dy)) : fi;\n\n // Velocity diffusion\n float kk = (float)(iix * iix + iiy * iiy); // k^2\n float diff = 1.f / (1.f + visc * dt * kk);\n xterm.x *= diff;\n xterm.y *= diff;\n yterm.x *= diff;\n yterm.y *= diff;\n\n // Velocity projection\n if (kk > 0.f) {\n float rkk = 1.f / kk;\n // Real portion of velocity projection\n float rkp = (iix * xterm.x + iiy * yterm.x);\n // Imaginary portion of velocity projection\n float ikp = (iix * xterm.y + iiy * yterm.y);\n xterm.x -= rkk * rkp * iix;\n xterm.y -= rkk * ikp * iix;\n yterm.x -= rkk * rkp * iiy;\n yterm.y -= rkk * ikp * iiy;\n }\n\n vx[fj] = xterm;\n vy[fj] = yterm;\n }\n }\n }\n}\n\n// This method updates the velocity field 'v' using the two complex\n// arrays from the previous step: 'vx' and 'vy'. Here we scale the\n// real components by 1/(dx*dy) to account for an unnormalized FFT.\n__global__ void updateVelocity_k(cData *v, float *vx, float *vy, int dx, int pdx, int dy, int lb, size_t pitch)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n float vxterm, vyterm;\n cData nvterm;\n\n // gtidx is the domain location in x for this thread\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fjr = fi * pdx + gtidx;\n vxterm = vx[fjr];\n vyterm = vy[fjr];\n\n // Normalize the result of the inverse FFT\n float scale = 1.f / (dx * dy);\n nvterm.x = vxterm * scale;\n nvterm.y = vyterm * scale;\n\n cData *fj = (cData *)((char *)v + fi * pitch) + gtidx;\n *fj = nvterm;\n }\n } // If this thread is inside the domain in Y\n } // If this thread is inside the domain in X\n}\n\n// This method updates the particles by moving particle positions\n// according to the velocity field and time step. That is, for each\n// particle: p(t+1) = p(t) + dt * v(p(t)).\n__global__ void advectParticles_k(cData *part, cData *v, int dx, int dy, float dt, int lb, size_t pitch)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n // gtidx is the domain location in x for this thread\n cData pterm, vterm;\n\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fj = fi * dx + gtidx;\n pterm = part[fj];\n\n int xvi = ((int)(pterm.x * dx));\n int yvi = ((int)(pterm.y * dy));\n vterm = *((cData *)((char *)v + yvi * pitch) + xvi);\n\n pterm.x += dt * vterm.x;\n pterm.x = pterm.x - (int)pterm.x;\n pterm.x += 1.f;\n pterm.x = pterm.x - (int)pterm.x;\n pterm.y += dt * vterm.y;\n pterm.y = pterm.y - (int)pterm.y;\n pterm.y += 1.f;\n pterm.y = pterm.y - (int)pterm.y;\n\n part[fj] = pterm;\n }\n } // If this thread is inside the domain in Y\n } // If this thread is inside the domain in X\n}\n\n// These are the external function calls necessary for launching fluid simuation\nextern \"C\" void addForces(cData *v, int dx, int dy, int spx, int spy, float fx, float fy, int r)\n{\n dim3 tids(2 * r + 1, 2 * r + 1);\n\n addForces_k<<<1, tids>>>(v, dx, dy, spx, spy, fx, fy, r, tPitch);\n getLastCudaError(\"addForces_k failed.\");\n}\n\nextern \"C\" void advectVelocity(cData *v, float *vx, float *vy, int dx, int pdx, int dy, float dt)\n{\n dim3 grid((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1));\n\n dim3 tids(TIDSX, TIDSY);\n\n updateTexture(v, DIM * sizeof(cData), DIM, tPitch);\n advectVelocity_k<<>>(v, vx, vy, dx, pdx, dy, dt, TILEY / TIDSY, texObj);\n getLastCudaError(\"advectVelocity_k failed.\");\n}\n\nextern \"C\" void diffuseProject(cData *vx, cData *vy, int dx, int dy, float dt, float visc)\n{\n // Forward FFT\n checkCudaErrors(cufftExecR2C(planr2c, (cufftReal *)vx, (cufftComplex *)vx));\n checkCudaErrors(cufftExecR2C(planr2c, (cufftReal *)vy, (cufftComplex *)vy));\n\n uint3 grid = make_uint3((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1), 1);\n uint3 tids = make_uint3(TIDSX, TIDSY, 1);\n\n diffuseProject_k<<>>(vx, vy, dx, dy, dt, visc, TILEY / TIDSY);\n getLastCudaError(\"diffuseProject_k failed.\");\n\n // Inverse FFT\n checkCudaErrors(cufftExecC2R(planc2r, (cufftComplex *)vx, (cufftReal *)vx));\n checkCudaErrors(cufftExecC2R(planc2r, (cufftComplex *)vy, (cufftReal *)vy));\n}\n\nextern \"C\" void updateVelocity(cData *v, float *vx, float *vy, int dx, int pdx, int dy)\n{\n dim3 grid((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1));\n dim3 tids(TIDSX, TIDSY);\n\n updateVelocity_k<<>>(v, vx, vy, dx, pdx, dy, TILEY / TIDSY, tPitch);\n getLastCudaError(\"updateVelocity_k failed.\");\n}\n\nextern \"C\" void advectParticles(GLuint vbo, cData *v, int dx, int dy, float dt)\n{\n dim3 grid((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1));\n dim3 tids(TIDSX, TIDSY);\n\n cData *p;\n checkCudaErrors(cudaGraphicsMapResources(1, &cuda_vbo_resource, 0));\n getLastCudaError(\"cudaGraphicsMapResources failed\");\n\n size_t num_bytes;\n checkCudaErrors(cudaGraphicsResourceGetMappedPointer((void **)&p, &num_bytes, cuda_vbo_resource));\n getLastCudaError(\"cudaGraphicsResourceGetMappedPointer failed\");\n\n advectParticles_k<<>>(p, v, dx, dy, dt, TILEY / TIDSY, tPitch);\n getLastCudaError(\"advectParticles_k failed.\");\n\n checkCudaErrors(cudaGraphicsUnmapResources(1, &cuda_vbo_resource, 0));\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include // CUDA FFT Libraries\n#include // Helper functions for CUDA Error handling\n#include \n#include \n\n// OpenGL Graphics includes\n#include \n\n// FluidsGLES CUDA kernel definitions\n#include \"fluidsGLES_kernels.cuh\"\n\n// Texture object for reading velocity field\nhipTextureObject_t texObj;\nstatic hipArray *array = NULL;\n\n// Particle data\nextern GLuint vbo; // OpenGL vertex buffer object\nextern struct hipGraphicsResource *cuda_vbo_resource; // handles OpenGL-CUDA exchange\n\n// Texture pitch\nextern size_t tPitch;\nextern hipfftHandle planr2c;\nextern hipfftHandle planc2r;\ncData *vxfield = NULL;\ncData *vyfield = NULL;\n\nvoid setupTexture(int x, int y)\n{\n hipChannelFormatDesc desc = hipCreateChannelDesc();\n\n hipMallocArray(&array, &desc, y, x);\n getLastCudaError(\"hipMalloc failed\");\n\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = array;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = hipFilterModeLinear;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&texObj, &texRes, &texDescr, NULL));\n}\n\nvoid updateTexture(cData *data, size_t wib, size_t h, size_t pitch)\n{\n checkCudaErrors(hipMemcpy2DToArray(array, 0, 0, data, pitch, wib, h, hipMemcpyDeviceToDevice));\n}\n\nvoid deleteTexture(void)\n{\n checkCudaErrors(hipDestroyTextureObject(texObj));\n checkCudaErrors(hipFreeArray(array));\n}\n\n// Note that these kernels are designed to work with arbitrary\n// domain sizes, not just domains that are multiples of the tile\n// size. Therefore, we have extra code that checks to make sure\n// a given thread location falls within the domain boundaries in\n// both X and Y. Also, the domain is covered by looping over\n// multiple elements in the Y direction, while there is a one-to-one\n// mapping between threads in X and the tile size in X.\n// Nolan Goodnight 9/22/06\n\n// This method adds constant force vectors to the velocity field\n// stored in 'v' according to v(x,t+1) = v(x,t) + dt * f.\n__global__ void addForces_k(cData *v, int dx, int dy, int spx, int spy, float fx, float fy, int r, size_t pitch)\n{\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n cData *fj = (cData *)((char *)v + (ty + spy) * pitch) + tx + spx;\n\n cData vterm = *fj;\n tx -= r;\n ty -= r;\n float s = 1.f / (1.f + tx * tx * tx * tx + ty * ty * ty * ty);\n vterm.x += s * fx;\n vterm.y += s * fy;\n *fj = vterm;\n}\n\n// This method performs the velocity advection step, where we\n// trace velocity vectors back in time to update each grid cell.\n// That is, v(x,t+1) = v(p(x,-dt),t). Here we perform bilinear\n// interpolation in the velocity space.\n__global__ void advectVelocity_k(cData *v,\n float *vx,\n float *vy,\n int dx,\n int pdx,\n int dy,\n float dt,\n int lb,\n hipTextureObject_t texObject)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n cData vterm, ploc;\n float vxterm, vyterm;\n\n // gtidx is the domain location in x for this thread\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fj = fi * pdx + gtidx;\n vterm = tex2D(texObject, (float)gtidx, (float)fi);\n ploc.x = (gtidx + 0.5f) - (dt * vterm.x * dx);\n ploc.y = (fi + 0.5f) - (dt * vterm.y * dy);\n vterm = tex2D(texObject, ploc.x, ploc.y);\n vxterm = vterm.x;\n vyterm = vterm.y;\n vx[fj] = vxterm;\n vy[fj] = vyterm;\n }\n }\n }\n}\n\n// This method performs velocity diffusion and forces mass conservation\n// in the frequency domain. The inputs 'vx' and 'vy' are complex-valued\n// arrays holding the Fourier coefficients of the velocity field in\n// X and Y. Diffusion in this space takes a simple form described as:\n// v(k,t) = v(k,t) / (1 + visc * dt * k^2), where visc is the viscosity,\n// and k is the wavenumber. The projection step forces the Fourier\n// velocity vectors to be orthogonal to the vectors for each\n// wavenumber: v(k,t) = v(k,t) - ((k dot v(k,t) * k) / k^2.\n__global__ void diffuseProject_k(cData *vx, cData *vy, int dx, int dy, float dt, float visc, int lb)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n cData xterm, yterm;\n\n // gtidx is the domain location in x for this thread\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fj = fi * dx + gtidx;\n xterm = vx[fj];\n yterm = vy[fj];\n\n // Compute the index of the wavenumber based on the\n // data order produced by a standard NN FFT.\n int iix = gtidx;\n int iiy = (fi > dy / 2) ? (fi - (dy)) : fi;\n\n // Velocity diffusion\n float kk = (float)(iix * iix + iiy * iiy); // k^2\n float diff = 1.f / (1.f + visc * dt * kk);\n xterm.x *= diff;\n xterm.y *= diff;\n yterm.x *= diff;\n yterm.y *= diff;\n\n // Velocity projection\n if (kk > 0.f) {\n float rkk = 1.f / kk;\n // Real portion of velocity projection\n float rkp = (iix * xterm.x + iiy * yterm.x);\n // Imaginary portion of velocity projection\n float ikp = (iix * xterm.y + iiy * yterm.y);\n xterm.x -= rkk * rkp * iix;\n xterm.y -= rkk * ikp * iix;\n yterm.x -= rkk * rkp * iiy;\n yterm.y -= rkk * ikp * iiy;\n }\n\n vx[fj] = xterm;\n vy[fj] = yterm;\n }\n }\n }\n}\n\n// This method updates the velocity field 'v' using the two complex\n// arrays from the previous step: 'vx' and 'vy'. Here we scale the\n// real components by 1/(dx*dy) to account for an unnormalized FFT.\n__global__ void updateVelocity_k(cData *v, float *vx, float *vy, int dx, int pdx, int dy, int lb, size_t pitch)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n float vxterm, vyterm;\n cData nvterm;\n\n // gtidx is the domain location in x for this thread\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fjr = fi * pdx + gtidx;\n vxterm = vx[fjr];\n vyterm = vy[fjr];\n\n // Normalize the result of the inverse FFT\n float scale = 1.f / (dx * dy);\n nvterm.x = vxterm * scale;\n nvterm.y = vyterm * scale;\n\n cData *fj = (cData *)((char *)v + fi * pitch) + gtidx;\n *fj = nvterm;\n }\n } // If this thread is inside the domain in Y\n } // If this thread is inside the domain in X\n}\n\n// This method updates the particles by moving particle positions\n// according to the velocity field and time step. That is, for each\n// particle: p(t+1) = p(t) + dt * v(p(t)).\n__global__ void advectParticles_k(cData *part, cData *v, int dx, int dy, float dt, int lb, size_t pitch)\n{\n int gtidx = blockIdx.x * blockDim.x + threadIdx.x;\n int gtidy = blockIdx.y * (lb * blockDim.y) + threadIdx.y * lb;\n int p;\n\n // gtidx is the domain location in x for this thread\n cData pterm, vterm;\n\n if (gtidx < dx) {\n for (p = 0; p < lb; p++) {\n // fi is the domain location in y for this thread\n int fi = gtidy + p;\n\n if (fi < dy) {\n int fj = fi * dx + gtidx;\n pterm = part[fj];\n\n int xvi = ((int)(pterm.x * dx));\n int yvi = ((int)(pterm.y * dy));\n vterm = *((cData *)((char *)v + yvi * pitch) + xvi);\n\n pterm.x += dt * vterm.x;\n pterm.x = pterm.x - (int)pterm.x;\n pterm.x += 1.f;\n pterm.x = pterm.x - (int)pterm.x;\n pterm.y += dt * vterm.y;\n pterm.y = pterm.y - (int)pterm.y;\n pterm.y += 1.f;\n pterm.y = pterm.y - (int)pterm.y;\n\n part[fj] = pterm;\n }\n } // If this thread is inside the domain in Y\n } // If this thread is inside the domain in X\n}\n\n// These are the external function calls necessary for launching fluid simuation\nextern \"C\" void addForces(cData *v, int dx, int dy, int spx, int spy, float fx, float fy, int r)\n{\n dim3 tids(2 * r + 1, 2 * r + 1);\n\n addForces_k<<<1, tids>>>(v, dx, dy, spx, spy, fx, fy, r, tPitch);\n getLastCudaError(\"addForces_k failed.\");\n}\n\nextern \"C\" void advectVelocity(cData *v, float *vx, float *vy, int dx, int pdx, int dy, float dt)\n{\n dim3 grid((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1));\n\n dim3 tids(TIDSX, TIDSY);\n\n updateTexture(v, DIM * sizeof(cData), DIM, tPitch);\n advectVelocity_k<<>>(v, vx, vy, dx, pdx, dy, dt, TILEY / TIDSY, texObj);\n getLastCudaError(\"advectVelocity_k failed.\");\n}\n\nextern \"C\" void diffuseProject(cData *vx, cData *vy, int dx, int dy, float dt, float visc)\n{\n // Forward FFT\n checkCudaErrors(hipfftExecR2C(planr2c, (hipfftReal *)vx, (hipfftComplex *)vx));\n checkCudaErrors(hipfftExecR2C(planr2c, (hipfftReal *)vy, (hipfftComplex *)vy));\n\n uint3 grid = make_uint3((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1), 1);\n uint3 tids = make_uint3(TIDSX, TIDSY, 1);\n\n diffuseProject_k<<>>(vx, vy, dx, dy, dt, visc, TILEY / TIDSY);\n getLastCudaError(\"diffuseProject_k failed.\");\n\n // Inverse FFT\n checkCudaErrors(hipfftExecC2R(planc2r, (hipfftComplex *)vx, (hipfftReal *)vx));\n checkCudaErrors(hipfftExecC2R(planc2r, (hipfftComplex *)vy, (hipfftReal *)vy));\n}\n\nextern \"C\" void updateVelocity(cData *v, float *vx, float *vy, int dx, int pdx, int dy)\n{\n dim3 grid((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1));\n dim3 tids(TIDSX, TIDSY);\n\n updateVelocity_k<<>>(v, vx, vy, dx, pdx, dy, TILEY / TIDSY, tPitch);\n getLastCudaError(\"updateVelocity_k failed.\");\n}\n\nextern \"C\" void advectParticles(GLuint vbo, cData *v, int dx, int dy, float dt)\n{\n dim3 grid((dx / TILEX) + (!(dx % TILEX) ? 0 : 1), (dy / TILEY) + (!(dy % TILEY) ? 0 : 1));\n dim3 tids(TIDSX, TIDSY);\n\n cData *p;\n checkCudaErrors(hipGraphicsMapResources(1, &cuda_vbo_resource, 0));\n getLastCudaError(\"hipGraphicsMapResources failed\");\n\n size_t num_bytes;\n checkCudaErrors(hipGraphicsResourceGetMappedPointer((void **)&p, &num_bytes, cuda_vbo_resource));\n getLastCudaError(\"hipGraphicsResourceGetMappedPointer failed\");\n\n advectParticles_k<<>>(p, v, dx, dy, dt, TILEY / TIDSY, tPitch);\n getLastCudaError(\"advectParticles_k failed.\");\n\n checkCudaErrors(hipGraphicsUnmapResources(1, &cuda_vbo_resource, 0));\n}\n", "source": "/kaggle/working/cuda-samples/Samples/8_Platform_Specific/Tegra/fluidsGLES/fluidsGLES_kernels.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n This example demonstrates how to use the CUDA C bindings to OpenGL ES to\n dynamically modify a vertex buffer using a CUDA C kernel.\n\n The steps are:\n 1. Create an empty vertex buffer object (VBO)\n 2. Register the VBO with CUDA C\n 3. Map the VBO for writing from CUDA C\n 4. Run CUDA C kernel to modify the vertex positions\n 5. Unmap the VBO\n 6. Render the results using OpenGL ES\n\n Host code\n*/\n\n// includes, system\n#include \n#include \n#include \n#include \n#include \n#include \n\nvoid error_exit(const char *format, ...)\n{\n va_list args;\n va_start(args, format);\n vfprintf(stderr, format, args);\n va_end(args);\n exit(1);\n}\n\n#if 0\n#include \"graphics_interface.c\"\n#else\n#include \"graphics_interface_egloutput_via_egl.c\"\n#endif\n\n#ifdef _WIN32\n#define WINDOWS_LEAN_AND_MEAN\n#define NOMINMAX\n#include \n#endif\n\n// includes, cuda\n#include \n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n // #include // helper functions for CUDA/GL interop\n\n#include \n\n#define MAX_EPSILON_ERROR 0.0f\n#define THRESHOLD 0.0f\n#define REFRESH_DELAY 1 // ms\n\n#define GUI_IDLE 0x100\n#define GUI_ROTATE 0x101\n#define GUI_TRANSLATE 0x102\n\nint gui_mode;\n\n////////////////////////////////////////////////////////////////////////////////\n// constants\nconst unsigned int window_width = 512;\nconst unsigned int window_height = 512;\n\nconst unsigned int mesh_width = 256;\nconst unsigned int mesh_height = 256;\n\n// OpenGL ES variables and interop with CUDA C\nGLuint mesh_vao, mesh_vbo;\nstruct cudaGraphicsResource *cuda_vbo_resource;\nvoid *d_vbo_buffer = NULL;\n\nfloat g_fAnim = 0.0;\n\n// UI / mouse controls\nint mouse_old_x, mouse_old_y;\nint mouse_buttons = 0;\nfloat rotate_x = 0.0, rotate_y = 0.0;\nfloat translate_z = -3.0;\n\nStopWatchInterface *timer = NULL;\n\n// Frame statistics\nint frame;\nint fpsCount = 0; // FPS count for averaging\nint fpsLimit = 1; // FPS limit for sampling\nint g_Index = 0;\nfloat avgFPS = 0.0f;\nunsigned int frameCount = 0;\nunsigned int g_TotalErrors = 0;\n\n// Auto-Verification Code\nbool g_bQAReadback = false;\n\nint *pArgc = NULL;\nchar **pArgv = NULL;\n\n#define MAX(a, b) ((a > b) ? a : b)\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\n\n// CUDA functionality\nvoid runCuda(struct cudaGraphicsResource **vbo_resource);\nvoid runAutoTest(int devID, char **argv, char *ref_file);\nvoid checkResultCuda(int argc, char **argv, const GLuint &vbo);\n\nconst char *sSDKsample = \"simpleGLES (VBO)\";\n\nvoid computeFPS()\n{\n frameCount++;\n fpsCount++;\n\n if (fpsCount == fpsLimit) {\n avgFPS = 1.f / (sdkGetAverageTimerValue(&timer) / 1000.f);\n fpsCount = 0;\n fpsLimit = (int)MAX(avgFPS, 1.f);\n\n sdkResetTimer(&timer);\n }\n\n char fps[256];\n sprintf(fps, \"Cuda/OpenGL ES Interop (VBO): %3.1f fps (Max 1000 fps)\", avgFPS);\n graphics_set_windowtitle(fps);\n}\n\n///////////////////////////////////////////////////////////////////////////////\n//! Simple kernel to modify vertex positions in sine wave pattern\n//! @param data data in global memory\n///////////////////////////////////////////////////////////////////////////////\n__global__ void simple_vbo_kernel(float4 *pos, unsigned int width, unsigned int height, float time)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // calculate uv coordinates\n float u = x / (float)width;\n float v = y / (float)height;\n u = u * 2.0f - 1.0f;\n v = v * 2.0f - 1.0f;\n\n // calculate simple sine wave pattern\n float freq = 4.0f;\n float w = sinf(u * freq + time) * cosf(v * freq + time) * 0.5f;\n\n // write output vertex\n pos[y * width + x] = make_float4(u, w, v, 1.0f);\n}\n\nvoid launch_kernel(float4 *pos, unsigned int mesh_width, unsigned int mesh_height, float time)\n{\n // execute the kernel\n dim3 block(8, 8, 1);\n dim3 grid(mesh_width / block.x, mesh_height / block.y, 1);\n simple_vbo_kernel<<>>(pos, mesh_width, mesh_height, time);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run the Cuda part of the computation\n////////////////////////////////////////////////////////////////////////////////\nvoid runCuda(struct cudaGraphicsResource **vbo_resource)\n{\n // map OpenGL buffer object for writing from CUDA\n float4 *dptr;\n cudaGraphicsMapResources(1, vbo_resource, 0);\n size_t num_bytes;\n cudaGraphicsResourceGetMappedPointer((void **)&dptr, &num_bytes, *vbo_resource);\n // printf(\"Sample CUDA mapped VBO: May access %ld bytes\\n\", num_bytes);\n\n // execute the kernel\n // dim3 block(8, 8, 1);\n // dim3 grid(mesh_width / block.x, mesh_height / block.y, 1);\n // kernel<<< grid, block>>>(dptr, mesh_width, mesh_height, g_fAnim);\n\n launch_kernel(dptr, mesh_width, mesh_height, g_fAnim);\n\n // unmap buffer object\n cudaGraphicsUnmapResources(1, vbo_resource, 0);\n}\n\n#ifdef _WIN32\n#ifndef FOPEN\n#define FOPEN(fHandle, filename, mode) fopen_s(&fHandle, filename, mode)\n#endif\n#else\n#ifndef FOPEN\n#define FOPEN(fHandle, filename, mode) (fHandle = fopen(filename, mode))\n#endif\n#endif\n\nvoid sdkDumpBin2(void *data, unsigned int bytes, const char *filename)\n{\n printf(\"sdkDumpBin: <%s>\\n\", filename);\n FILE *fp;\n FOPEN(fp, filename, \"wb\");\n fwrite(data, bytes, 1, fp);\n fflush(fp);\n fclose(fp);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run the Cuda part of the computation\n////////////////////////////////////////////////////////////////////////////////\nvoid runAutoTest(int devID, char **argv, char *ref_file)\n{\n char *reference_file = NULL;\n void *imageData = malloc(mesh_width * mesh_height * sizeof(float));\n\n // execute the kernel\n launch_kernel((float4 *)d_vbo_buffer, mesh_width, mesh_height, g_fAnim);\n\n cudaDeviceSynchronize();\n getLastCudaError(\"launch_kernel failed\");\n\n cudaMemcpy(imageData, d_vbo_buffer, mesh_width * mesh_height * sizeof(float), cudaMemcpyDeviceToHost);\n\n sdkDumpBin2(imageData, mesh_width * mesh_height * sizeof(float), \"simpleGL.bin\");\n reference_file = sdkFindFilePath(ref_file, argv[0]);\n\n if (reference_file\n && !sdkCompareBin2BinFloat(\"simpleGL.bin\",\n reference_file,\n mesh_width * mesh_height * sizeof(float),\n MAX_EPSILON_ERROR,\n THRESHOLD,\n pArgv[0])) {\n g_TotalErrors++;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Display callback\n////////////////////////////////////////////////////////////////////////////////\nvoid display_thisframe(float time_delta)\n{\n sdkStartTimer(&timer);\n\n // run CUDA kernel to generate vertex positions\n runCuda(&cuda_vbo_resource);\n\n glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);\n // GET_GLERROR(0);\n\n // set view matrix: broken, it doesn't work in OpenGL ES! Must put into shader\n // glMatrixMode(GL_MODELVIEW);\n // glLoadIdentity();\n // glTranslatef(0.0, 0.0, translate_z);\n // glRotatef(rotate_x, 1.0, 0.0, 0.0);\n // glRotatef(rotate_y, 0.0, 1.0, 0.0);\n\n glDrawArrays(GL_POINTS, 0, mesh_width * mesh_height);\n\n // GET_GLERROR(0);\n glFinish();\n // GET_GLERROR(0);\n\n g_fAnim += time_delta;\n\n sdkStopTimer(&timer);\n computeFPS();\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Check if the result is correct or write data to file for external\n//! regression testing\n////////////////////////////////////////////////////////////////////////////////\nvoid checkResultCuda(int argc, char **argv, const GLuint &vbo)\n{\n if (!d_vbo_buffer) {\n printf(\"%s: Mapping result buffer from OpenGL ES\\n\", __FUNCTION__);\n\n cudaGraphicsUnregisterResource(cuda_vbo_resource);\n\n // map buffer object\n glBindBuffer(GL_ARRAY_BUFFER, vbo);\n float *data =\n (float *)glMapBufferRange(GL_ARRAY_BUFFER, 0, mesh_width * mesh_height * 4 * sizeof(float), GL_READ_ONLY);\n\n // check result\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // write file for regression test\n sdkWriteFile(\"./data/regression.dat\", data, mesh_width * mesh_height * 3, 0.0, false);\n }\n\n // unmap GL buffer object\n if (!glUnmapBuffer(GL_ARRAY_BUFFER)) {\n fprintf(stderr, \"Unmap buffer failed.\\n\");\n fflush(stderr);\n }\n\n checkCudaErrors(cudaGraphicsGLRegisterBuffer(&cuda_vbo_resource, vbo, cudaGraphicsMapFlagsWriteDiscard));\n\n GET_GLERROR(0);\n }\n}\n\nGLuint mesh_shader = 0;\n\nvoid readAndCompileShaderFromGLSLFile(GLuint new_shaderprogram, const char *filename, GLenum shaderType)\n{\n FILE *file = fopen(filename, \"rb\"); // open shader text file\n if (!file)\n error_exit(\"Filename %s does not exist\\n\", filename);\n\n /* get the size of the file and read it */\n fseek(file, 0, SEEK_END);\n GLint size = ftell(file);\n char *data = (char *)malloc(sizeof(char) * (size + 1));\n memset(data, 0, sizeof(char) * (size + 1));\n fseek(file, 0, SEEK_SET);\n size_t res = fread(data, 1, size, file);\n fclose(file);\n\n GLuint shader = glCreateShader(shaderType);\n glShaderSource(shader, 1, (const GLchar **)&data, &size);\n glCompileShader(shader);\n\n GET_GLERROR(0);\n GLint compile_success = 0;\n glGetShaderiv(shader, GL_COMPILE_STATUS, &compile_success);\n GET_GLERROR(0);\n\n if (compile_success == GL_FALSE) {\n printf(\"Compilation of %s failed!\\n Reason:\\n\", filename);\n\n GLint maxLength = 0;\n glGetShaderiv(shader, GL_INFO_LOG_LENGTH, &maxLength);\n\n char errorLog[maxLength];\n glGetShaderInfoLog(shader, maxLength, &maxLength, &errorLog[0]);\n\n printf(\"%s\", errorLog);\n\n glDeleteShader(shader);\n exit(1);\n }\n\n glAttachShader(new_shaderprogram, shader);\n glDeleteShader(shader); // good to do?\n\n free(data);\n}\n\nGLuint ShaderCreate(const char *vshader_filename, const char *fshader_filename)\n{\n printf(\"Loading GLSL shaders %s %s\\n\", vshader_filename, fshader_filename);\n\n GLuint new_shaderprogram = glCreateProgram();\n\n GET_GLERROR(0);\n if (vshader_filename)\n readAndCompileShaderFromGLSLFile(new_shaderprogram, vshader_filename, GL_VERTEX_SHADER);\n\n GET_GLERROR(0);\n if (fshader_filename)\n readAndCompileShaderFromGLSLFile(new_shaderprogram, fshader_filename, GL_FRAGMENT_SHADER);\n\n GET_GLERROR(0);\n\n glLinkProgram(new_shaderprogram);\n\n GET_GLERROR(0);\n GLint link_success;\n glGetProgramiv(new_shaderprogram, GL_LINK_STATUS, &link_success);\n\n if (link_success == GL_FALSE) {\n printf(\"Linking of %s with %s failed!\\n Reason:\\n\", vshader_filename, fshader_filename);\n\n GLint maxLength = 0;\n glGetShaderiv(new_shaderprogram, GL_INFO_LOG_LENGTH, &maxLength);\n\n char errorLog[maxLength];\n glGetShaderInfoLog(new_shaderprogram, maxLength, &maxLength, &errorLog[0]);\n\n printf(\"%s\", errorLog);\n\n exit(EXIT_FAILURE);\n }\n\n return new_shaderprogram;\n}\n\n//===========================================================================\n// InitGraphicsState() - initialize OpenGL\n//===========================================================================\nstatic void InitGraphicsState(char **argv)\n{\n char *GL_version = (char *)glGetString(GL_VERSION);\n char *GL_vendor = (char *)glGetString(GL_VENDOR);\n char *GL_renderer = (char *)glGetString(GL_RENDERER);\n\n printf(\"Version: %s\\n\", GL_version);\n printf(\"Vendor: %s\\n\", GL_vendor);\n printf(\"Renderer: %s\\n\", GL_renderer);\n\n // RENDERING SETUP (OpenGL ES or OpenGL Core Profile!)\n glGenVertexArrays(1, &mesh_vao); // Features' Vertex Array Object allocation\n glBindVertexArray(mesh_vao); // bind VAO\n\n // initialize buffer object\n glGenBuffers(1, &mesh_vbo);\n glBindBuffer(GL_ARRAY_BUFFER, mesh_vbo);\n\n unsigned int size = mesh_width * mesh_height * 4 * sizeof(float);\n glBufferData(GL_ARRAY_BUFFER, size, NULL, GL_DYNAMIC_DRAW);\n glVertexAttribPointer((GLuint)0, 4, GL_FLOAT, GL_FALSE, 0, 0);\n glEnableVertexAttribArray(0);\n\n checkCudaErrors(cudaGraphicsGLRegisterBuffer(&cuda_vbo_resource, mesh_vbo, cudaGraphicsMapFlagsNone));\n\n // glBindVertexArray(0); // keep above Vertex Array Object bound (it's the\n // only one throughout)\n\n // GLSL stuff\n char *vertex_shader_path = sdkFindFilePath(\"mesh.vert.glsl\", argv[0]);\n char *fragment_shader_path = sdkFindFilePath(\"mesh.frag.glsl\", argv[0]);\n\n if (vertex_shader_path == NULL || fragment_shader_path == NULL) {\n printf(\"Error finding shader file\\n\");\n exit(EXIT_FAILURE);\n }\n\n mesh_shader = ShaderCreate(vertex_shader_path, fragment_shader_path);\n GET_GLERROR(0);\n\n free(vertex_shader_path);\n free(fragment_shader_path);\n\n glUseProgram(mesh_shader);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nbool runTest(int argc, char **argv, char *ref_file)\n{\n // Create the CUTIL timer\n sdkCreateTimer(&timer);\n\n int devID = 0;\n#if defined(__aarch64__) || defined(__arm__)\n // find iGPU on the system which is compute capable which will perform\n // GLES-CUDA interop\n devID = findIntegratedGPU();\n#else\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n devID = findCudaDevice(argc, (const char **)argv);\n#endif\n\n // command line mode only\n if (ref_file != NULL) {\n // create VBO\n checkCudaErrors(cudaMalloc((void **)&d_vbo_buffer, mesh_width * mesh_height * 4 * sizeof(float)));\n\n // run the cuda part\n runAutoTest(devID, argv, ref_file);\n\n // check result of Cuda step\n checkResultCuda(argc, argv, mesh_vbo);\n\n cudaFree(d_vbo_buffer);\n d_vbo_buffer = NULL;\n }\n else {\n // this would use command-line specified CUDA device, note that CUDA\n // defaults to highest Gflops/s device\n if (checkCmdLineFlag(argc, (const char **)argv, \"device\"))\n error_exit(\"Device setting not yet implemented!\\n\");\n\n // create X11 window and set up associated OpenGL ES context\n graphics_setup_window(0, 0, window_width, window_height, sSDKsample);\n\n InitGraphicsState(argv); // set up GLES stuff\n\n glClearColor(0, 0.5, 1, 1); // blue-ish background\n glClear(GL_COLOR_BUFFER_BIT);\n\n // printf(\"WP%d\\n\", __LINE__);\n graphics_swap_buffers();\n\n int frame = 0;\n\n while (frame < 1000) {\n display_thisframe(0.010);\n usleep(1000); // need not take full CPU and GPU\n\n graphics_swap_buffers();\n // printf(\"frame %d\\n\",frame++);\n }\n\n // NOTE: Before destroying OpenGL ES context, must unregister all shared\n // resources from CUDA !\n cudaGraphicsUnregisterResource(cuda_vbo_resource);\n\n graphics_close_window(); // close window and destroy OpenGL ES context\n\n sdkDeleteTimer(&timer);\n }\n\n return true;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n char *ref_file = NULL;\n\n pArgc = &argc;\n pArgv = argv;\n\n#if defined(__linux__)\n setenv(\"DISPLAY\", \":0\", 0);\n#endif\n\n printf(\"%s starting...\\n\", sSDKsample);\n\n if (argc > 1) {\n if (checkCmdLineFlag(argc, (const char **)argv, \"file\")) {\n // In this mode, we run without OpenGL and see if VBO is generated\n // correctly\n getCmdLineArgumentString(argc, (const char **)argv, \"file\", (char **)&ref_file);\n }\n }\n\n printf(\"\\n\");\n\n runTest(argc, argv, ref_file);\n\n printf(\"%s completed, returned %s\\n\", sSDKsample, (g_TotalErrors == 0) ? \"OK\" : \"ERROR!\");\n\n exit(g_TotalErrors == 0 ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n This example demonstrates how to use the CUDA C bindings to OpenGL ES to\n dynamically modify a vertex buffer using a CUDA C kernel.\n\n The steps are:\n 1. Create an empty vertex buffer object (VBO)\n 2. Register the VBO with CUDA C\n 3. Map the VBO for writing from CUDA C\n 4. Run CUDA C kernel to modify the vertex positions\n 5. Unmap the VBO\n 6. Render the results using OpenGL ES\n\n Host code\n*/\n\n// includes, system\n#include \n#include \n#include \n#include \n#include \n#include \n\nvoid error_exit(const char *format, ...)\n{\n va_list args;\n va_start(args, format);\n vfprintf(stderr, format, args);\n va_end(args);\n exit(1);\n}\n\n#if 0\n#include \"graphics_interface.c\"\n#else\n#include \"graphics_interface_egloutput_via_egl.c\"\n#endif\n\n#ifdef _WIN32\n#define WINDOWS_LEAN_AND_MEAN\n#define NOMINMAX\n#include \n#endif\n\n// includes, cuda\n#include \n#include \n\n// Utilities and timing functions\n#include // includes cuda.h and cuda_runtime_api.h\n\n// CUDA helper functions\n#include // helper functions for CUDA error check\n // #include // helper functions for CUDA/GL interop\n\n#include \n\n#define MAX_EPSILON_ERROR 0.0f\n#define THRESHOLD 0.0f\n#define REFRESH_DELAY 1 // ms\n\n#define GUI_IDLE 0x100\n#define GUI_ROTATE 0x101\n#define GUI_TRANSLATE 0x102\n\nint gui_mode;\n\n////////////////////////////////////////////////////////////////////////////////\n// constants\nconst unsigned int window_width = 512;\nconst unsigned int window_height = 512;\n\nconst unsigned int mesh_width = 256;\nconst unsigned int mesh_height = 256;\n\n// OpenGL ES variables and interop with CUDA C\nGLuint mesh_vao, mesh_vbo;\nstruct hipGraphicsResource *cuda_vbo_resource;\nvoid *d_vbo_buffer = NULL;\n\nfloat g_fAnim = 0.0;\n\n// UI / mouse controls\nint mouse_old_x, mouse_old_y;\nint mouse_buttons = 0;\nfloat rotate_x = 0.0, rotate_y = 0.0;\nfloat translate_z = -3.0;\n\nStopWatchInterface *timer = NULL;\n\n// Frame statistics\nint frame;\nint fpsCount = 0; // FPS count for averaging\nint fpsLimit = 1; // FPS limit for sampling\nint g_Index = 0;\nfloat avgFPS = 0.0f;\nunsigned int frameCount = 0;\nunsigned int g_TotalErrors = 0;\n\n// Auto-Verification Code\nbool g_bQAReadback = false;\n\nint *pArgc = NULL;\nchar **pArgv = NULL;\n\n#define MAX(a, b) ((a > b) ? a : b)\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\n\n// CUDA functionality\nvoid runCuda(struct hipGraphicsResource **vbo_resource);\nvoid runAutoTest(int devID, char **argv, char *ref_file);\nvoid checkResultCuda(int argc, char **argv, const GLuint &vbo);\n\nconst char *sSDKsample = \"simpleGLES (VBO)\";\n\nvoid computeFPS()\n{\n frameCount++;\n fpsCount++;\n\n if (fpsCount == fpsLimit) {\n avgFPS = 1.f / (sdkGetAverageTimerValue(&timer) / 1000.f);\n fpsCount = 0;\n fpsLimit = (int)MAX(avgFPS, 1.f);\n\n sdkResetTimer(&timer);\n }\n\n char fps[256];\n sprintf(fps, \"Cuda/OpenGL ES Interop (VBO): %3.1f fps (Max 1000 fps)\", avgFPS);\n graphics_set_windowtitle(fps);\n}\n\n///////////////////////////////////////////////////////////////////////////////\n//! Simple kernel to modify vertex positions in sine wave pattern\n//! @param data data in global memory\n///////////////////////////////////////////////////////////////////////////////\n__global__ void simple_vbo_kernel(float4 *pos, unsigned int width, unsigned int height, float time)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n // calculate uv coordinates\n float u = x / (float)width;\n float v = y / (float)height;\n u = u * 2.0f - 1.0f;\n v = v * 2.0f - 1.0f;\n\n // calculate simple sine wave pattern\n float freq = 4.0f;\n float w = sinf(u * freq + time) * cosf(v * freq + time) * 0.5f;\n\n // write output vertex\n pos[y * width + x] = make_float4(u, w, v, 1.0f);\n}\n\nvoid launch_kernel(float4 *pos, unsigned int mesh_width, unsigned int mesh_height, float time)\n{\n // execute the kernel\n dim3 block(8, 8, 1);\n dim3 grid(mesh_width / block.x, mesh_height / block.y, 1);\n simple_vbo_kernel<<>>(pos, mesh_width, mesh_height, time);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run the Cuda part of the computation\n////////////////////////////////////////////////////////////////////////////////\nvoid runCuda(struct hipGraphicsResource **vbo_resource)\n{\n // map OpenGL buffer object for writing from CUDA\n float4 *dptr;\n hipGraphicsMapResources(1, vbo_resource, 0);\n size_t num_bytes;\n hipGraphicsResourceGetMappedPointer((void **)&dptr, &num_bytes, *vbo_resource);\n // printf(\"Sample CUDA mapped VBO: May access %ld bytes\\n\", num_bytes);\n\n // execute the kernel\n // dim3 block(8, 8, 1);\n // dim3 grid(mesh_width / block.x, mesh_height / block.y, 1);\n // kernel<<< grid, block>>>(dptr, mesh_width, mesh_height, g_fAnim);\n\n launch_kernel(dptr, mesh_width, mesh_height, g_fAnim);\n\n // unmap buffer object\n hipGraphicsUnmapResources(1, vbo_resource, 0);\n}\n\n#ifdef _WIN32\n#ifndef FOPEN\n#define FOPEN(fHandle, filename, mode) fopen_s(&fHandle, filename, mode)\n#endif\n#else\n#ifndef FOPEN\n#define FOPEN(fHandle, filename, mode) (fHandle = fopen(filename, mode))\n#endif\n#endif\n\nvoid sdkDumpBin2(void *data, unsigned int bytes, const char *filename)\n{\n printf(\"sdkDumpBin: <%s>\\n\", filename);\n FILE *fp;\n FOPEN(fp, filename, \"wb\");\n fwrite(data, bytes, 1, fp);\n fflush(fp);\n fclose(fp);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run the Cuda part of the computation\n////////////////////////////////////////////////////////////////////////////////\nvoid runAutoTest(int devID, char **argv, char *ref_file)\n{\n char *reference_file = NULL;\n void *imageData = malloc(mesh_width * mesh_height * sizeof(float));\n\n // execute the kernel\n launch_kernel((float4 *)d_vbo_buffer, mesh_width, mesh_height, g_fAnim);\n\n hipDeviceSynchronize();\n getLastCudaError(\"launch_kernel failed\");\n\n hipMemcpy(imageData, d_vbo_buffer, mesh_width * mesh_height * sizeof(float), hipMemcpyDeviceToHost);\n\n sdkDumpBin2(imageData, mesh_width * mesh_height * sizeof(float), \"simpleGL.bin\");\n reference_file = sdkFindFilePath(ref_file, argv[0]);\n\n if (reference_file\n && !sdkCompareBin2BinFloat(\"simpleGL.bin\",\n reference_file,\n mesh_width * mesh_height * sizeof(float),\n MAX_EPSILON_ERROR,\n THRESHOLD,\n pArgv[0])) {\n g_TotalErrors++;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Display callback\n////////////////////////////////////////////////////////////////////////////////\nvoid display_thisframe(float time_delta)\n{\n sdkStartTimer(&timer);\n\n // run CUDA kernel to generate vertex positions\n runCuda(&cuda_vbo_resource);\n\n glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);\n // GET_GLERROR(0);\n\n // set view matrix: broken, it doesn't work in OpenGL ES! Must put into shader\n // glMatrixMode(GL_MODELVIEW);\n // glLoadIdentity();\n // glTranslatef(0.0, 0.0, translate_z);\n // glRotatef(rotate_x, 1.0, 0.0, 0.0);\n // glRotatef(rotate_y, 0.0, 1.0, 0.0);\n\n glDrawArrays(GL_POINTS, 0, mesh_width * mesh_height);\n\n // GET_GLERROR(0);\n glFinish();\n // GET_GLERROR(0);\n\n g_fAnim += time_delta;\n\n sdkStopTimer(&timer);\n computeFPS();\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Check if the result is correct or write data to file for external\n//! regression testing\n////////////////////////////////////////////////////////////////////////////////\nvoid checkResultCuda(int argc, char **argv, const GLuint &vbo)\n{\n if (!d_vbo_buffer) {\n printf(\"%s: Mapping result buffer from OpenGL ES\\n\", __FUNCTION__);\n\n hipGraphicsUnregisterResource(cuda_vbo_resource);\n\n // map buffer object\n glBindBuffer(GL_ARRAY_BUFFER, vbo);\n float *data =\n (float *)glMapBufferRange(GL_ARRAY_BUFFER, 0, mesh_width * mesh_height * 4 * sizeof(float), GL_READ_ONLY);\n\n // check result\n if (checkCmdLineFlag(argc, (const char **)argv, \"regression\")) {\n // write file for regression test\n sdkWriteFile(\"./data/regression.dat\", data, mesh_width * mesh_height * 3, 0.0, false);\n }\n\n // unmap GL buffer object\n if (!glUnmapBuffer(GL_ARRAY_BUFFER)) {\n fprintf(stderr, \"Unmap buffer failed.\\n\");\n fflush(stderr);\n }\n\n checkCudaErrors(hipGraphicsGLRegisterBuffer(&cuda_vbo_resource, vbo, cudaGraphicsMapFlagsWriteDiscard));\n\n GET_GLERROR(0);\n }\n}\n\nGLuint mesh_shader = 0;\n\nvoid readAndCompileShaderFromGLSLFile(GLuint new_shaderprogram, const char *filename, GLenum shaderType)\n{\n FILE *file = fopen(filename, \"rb\"); // open shader text file\n if (!file)\n error_exit(\"Filename %s does not exist\\n\", filename);\n\n /* get the size of the file and read it */\n fseek(file, 0, SEEK_END);\n GLint size = ftell(file);\n char *data = (char *)malloc(sizeof(char) * (size + 1));\n memset(data, 0, sizeof(char) * (size + 1));\n fseek(file, 0, SEEK_SET);\n size_t res = fread(data, 1, size, file);\n fclose(file);\n\n GLuint shader = glCreateShader(shaderType);\n glShaderSource(shader, 1, (const GLchar **)&data, &size);\n glCompileShader(shader);\n\n GET_GLERROR(0);\n GLint compile_success = 0;\n glGetShaderiv(shader, GL_COMPILE_STATUS, &compile_success);\n GET_GLERROR(0);\n\n if (compile_success == GL_FALSE) {\n printf(\"Compilation of %s failed!\\n Reason:\\n\", filename);\n\n GLint maxLength = 0;\n glGetShaderiv(shader, GL_INFO_LOG_LENGTH, &maxLength);\n\n char errorLog[maxLength];\n glGetShaderInfoLog(shader, maxLength, &maxLength, &errorLog[0]);\n\n printf(\"%s\", errorLog);\n\n glDeleteShader(shader);\n exit(1);\n }\n\n glAttachShader(new_shaderprogram, shader);\n glDeleteShader(shader); // good to do?\n\n free(data);\n}\n\nGLuint ShaderCreate(const char *vshader_filename, const char *fshader_filename)\n{\n printf(\"Loading GLSL shaders %s %s\\n\", vshader_filename, fshader_filename);\n\n GLuint new_shaderprogram = glCreateProgram();\n\n GET_GLERROR(0);\n if (vshader_filename)\n readAndCompileShaderFromGLSLFile(new_shaderprogram, vshader_filename, GL_VERTEX_SHADER);\n\n GET_GLERROR(0);\n if (fshader_filename)\n readAndCompileShaderFromGLSLFile(new_shaderprogram, fshader_filename, GL_FRAGMENT_SHADER);\n\n GET_GLERROR(0);\n\n glLinkProgram(new_shaderprogram);\n\n GET_GLERROR(0);\n GLint link_success;\n glGetProgramiv(new_shaderprogram, GL_LINK_STATUS, &link_success);\n\n if (link_success == GL_FALSE) {\n printf(\"Linking of %s with %s failed!\\n Reason:\\n\", vshader_filename, fshader_filename);\n\n GLint maxLength = 0;\n glGetShaderiv(new_shaderprogram, GL_INFO_LOG_LENGTH, &maxLength);\n\n char errorLog[maxLength];\n glGetShaderInfoLog(new_shaderprogram, maxLength, &maxLength, &errorLog[0]);\n\n printf(\"%s\", errorLog);\n\n exit(EXIT_FAILURE);\n }\n\n return new_shaderprogram;\n}\n\n//===========================================================================\n// InitGraphicsState() - initialize OpenGL\n//===========================================================================\nstatic void InitGraphicsState(char **argv)\n{\n char *GL_version = (char *)glGetString(GL_VERSION);\n char *GL_vendor = (char *)glGetString(GL_VENDOR);\n char *GL_renderer = (char *)glGetString(GL_RENDERER);\n\n printf(\"Version: %s\\n\", GL_version);\n printf(\"Vendor: %s\\n\", GL_vendor);\n printf(\"Renderer: %s\\n\", GL_renderer);\n\n // RENDERING SETUP (OpenGL ES or OpenGL Core Profile!)\n glGenVertexArrays(1, &mesh_vao); // Features' Vertex Array Object allocation\n glBindVertexArray(mesh_vao); // bind VAO\n\n // initialize buffer object\n glGenBuffers(1, &mesh_vbo);\n glBindBuffer(GL_ARRAY_BUFFER, mesh_vbo);\n\n unsigned int size = mesh_width * mesh_height * 4 * sizeof(float);\n glBufferData(GL_ARRAY_BUFFER, size, NULL, GL_DYNAMIC_DRAW);\n glVertexAttribPointer((GLuint)0, 4, GL_FLOAT, GL_FALSE, 0, 0);\n glEnableVertexAttribArray(0);\n\n checkCudaErrors(hipGraphicsGLRegisterBuffer(&cuda_vbo_resource, mesh_vbo, cudaGraphicsMapFlagsNone));\n\n // glBindVertexArray(0); // keep above Vertex Array Object bound (it's the\n // only one throughout)\n\n // GLSL stuff\n char *vertex_shader_path = sdkFindFilePath(\"mesh.vert.glsl\", argv[0]);\n char *fragment_shader_path = sdkFindFilePath(\"mesh.frag.glsl\", argv[0]);\n\n if (vertex_shader_path == NULL || fragment_shader_path == NULL) {\n printf(\"Error finding shader file\\n\");\n exit(EXIT_FAILURE);\n }\n\n mesh_shader = ShaderCreate(vertex_shader_path, fragment_shader_path);\n GET_GLERROR(0);\n\n free(vertex_shader_path);\n free(fragment_shader_path);\n\n glUseProgram(mesh_shader);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nbool runTest(int argc, char **argv, char *ref_file)\n{\n // Create the CUTIL timer\n sdkCreateTimer(&timer);\n\n int devID = 0;\n#if defined(__aarch64__) || defined(__arm__)\n // find iGPU on the system which is compute capable which will perform\n // GLES-CUDA interop\n devID = findIntegratedGPU();\n#else\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n devID = findCudaDevice(argc, (const char **)argv);\n#endif\n\n // command line mode only\n if (ref_file != NULL) {\n // create VBO\n checkCudaErrors(hipMalloc((void **)&d_vbo_buffer, mesh_width * mesh_height * 4 * sizeof(float)));\n\n // run the cuda part\n runAutoTest(devID, argv, ref_file);\n\n // check result of Cuda step\n checkResultCuda(argc, argv, mesh_vbo);\n\n hipFree(d_vbo_buffer);\n d_vbo_buffer = NULL;\n }\n else {\n // this would use command-line specified CUDA device, note that CUDA\n // defaults to highest Gflops/s device\n if (checkCmdLineFlag(argc, (const char **)argv, \"device\"))\n error_exit(\"Device setting not yet implemented!\\n\");\n\n // create X11 window and set up associated OpenGL ES context\n graphics_setup_window(0, 0, window_width, window_height, sSDKsample);\n\n InitGraphicsState(argv); // set up GLES stuff\n\n glClearColor(0, 0.5, 1, 1); // blue-ish background\n glClear(GL_COLOR_BUFFER_BIT);\n\n // printf(\"WP%d\\n\", __LINE__);\n graphics_swap_buffers();\n\n int frame = 0;\n\n while (frame < 1000) {\n display_thisframe(0.010);\n usleep(1000); // need not take full CPU and GPU\n\n graphics_swap_buffers();\n // printf(\"frame %d\\n\",frame++);\n }\n\n // NOTE: Before destroying OpenGL ES context, must unregister all shared\n // resources from CUDA !\n hipGraphicsUnregisterResource(cuda_vbo_resource);\n\n graphics_close_window(); // close window and destroy OpenGL ES context\n\n sdkDeleteTimer(&timer);\n }\n\n return true;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n char *ref_file = NULL;\n\n pArgc = &argc;\n pArgv = argv;\n\n#if defined(__linux__)\n setenv(\"DISPLAY\", \":0\", 0);\n#endif\n\n printf(\"%s starting...\\n\", sSDKsample);\n\n if (argc > 1) {\n if (checkCmdLineFlag(argc, (const char **)argv, \"file\")) {\n // In this mode, we run without OpenGL and see if VBO is generated\n // correctly\n getCmdLineArgumentString(argc, (const char **)argv, \"file\", (char **)&ref_file);\n }\n }\n\n printf(\"\\n\");\n\n runTest(argc, argv, ref_file);\n\n printf(\"%s completed, returned %s\\n\", sSDKsample, (g_TotalErrors == 0) ? \"OK\" : \"ERROR!\");\n\n exit(g_TotalErrors == 0 ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/8_Platform_Specific/Tegra/simpleGLES_EGLOutput/simpleGLES_EGLOutput.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n\n// #include \n// #include \n\n// CUDA standard includes\n#include \n// #include \n\n#include \"bodysystem.h\"\n\n__constant__ float softeningSquared;\n__constant__ double softeningSquared_fp64;\n\ncudaError_t setSofteningSquared(float softeningSq)\n{\n return cudaMemcpyToSymbol(softeningSquared, &softeningSq, sizeof(float), 0, cudaMemcpyHostToDevice);\n}\n\ncudaError_t setSofteningSquared(double softeningSq)\n{\n return cudaMemcpyToSymbol(softeningSquared_fp64, &softeningSq, sizeof(double), 0, cudaMemcpyHostToDevice);\n}\n\ntemplate struct SharedMemory\n{\n __device__ inline operator T *()\n {\n extern __shared__ int __smem[];\n return (T *)__smem;\n }\n\n __device__ inline operator const T *() const\n {\n extern __shared__ int __smem[];\n return (T *)__smem;\n }\n};\n\ntemplate __device__ T rsqrt_T(T x) { return rsqrt(x); }\n\ntemplate <> __device__ float rsqrt_T(float x) { return rsqrtf(x); }\n\ntemplate <> __device__ double rsqrt_T(double x) { return rsqrt(x); }\n\n// Macros to simplify shared memory addressing\n#define SX(i) sharedPos[i + blockDim.x * threadIdx.y]\n// This macro is only used when multithreadBodies is true (below)\n#define SX_SUM(i, j) sharedPos[i + blockDim.x * j]\n\ntemplate __device__ T getSofteningSquared() { return softeningSquared; }\ntemplate <> __device__ double getSofteningSquared() { return softeningSquared_fp64; }\n\ntemplate struct DeviceData\n{\n T *dPos[2]; // mapped host pointers\n T *dVel;\n cudaEvent_t event;\n unsigned int offset;\n unsigned int numBodies;\n};\n\ntemplate \n__device__ typename vec3::Type\nbodyBodyInteraction(typename vec3::Type ai, typename vec4::Type bi, typename vec4::Type bj)\n{\n typename vec3::Type r;\n\n // r_ij [3 FLOPS]\n r.x = bj.x - bi.x;\n r.y = bj.y - bi.y;\n r.z = bj.z - bi.z;\n\n // distSqr = dot(r_ij, r_ij) + EPS^2 [6 FLOPS]\n T distSqr = r.x * r.x + r.y * r.y + r.z * r.z;\n distSqr += getSofteningSquared();\n\n // invDistCube =1/distSqr^(3/2) [4 FLOPS (2 mul, 1 sqrt, 1 inv)]\n T invDist = rsqrt_T(distSqr);\n T invDistCube = invDist * invDist * invDist;\n\n // s = m_j * invDistCube [1 FLOP]\n T s = bj.w * invDistCube;\n\n // a_i = a_i + s * r_ij [6 FLOPS]\n ai.x += r.x * s;\n ai.y += r.y * s;\n ai.z += r.z * s;\n\n return ai;\n}\n\ntemplate \n__device__ typename vec3::Type\ncomputeBodyAccel(typename vec4::Type bodyPos, typename vec4::Type *positions, int numTiles)\n{\n typename vec4::Type *sharedPos = SharedMemory::Type>();\n\n typename vec3::Type acc = {0.0f, 0.0f, 0.0f};\n\n for (int tile = 0; tile < numTiles; tile++) {\n sharedPos[threadIdx.x] = positions[tile * blockDim.x + threadIdx.x];\n\n __syncthreads();\n\n // This is the \"tile_calculation\" from the GPUG3 article.\n#pragma unroll 128\n\n for (unsigned int counter = 0; counter < blockDim.x; counter++) {\n acc = bodyBodyInteraction(acc, bodyPos, sharedPos[counter]);\n }\n\n __syncthreads();\n }\n\n return acc;\n}\n\ntemplate \n__global__ void integrateBodies(typename vec4::Type *__restrict__ newPos,\n typename vec4::Type *__restrict__ oldPos,\n typename vec4::Type *vel,\n unsigned int deviceOffset,\n unsigned int deviceNumBodies,\n float deltaTime,\n float damping,\n int numTiles)\n{\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (index >= deviceNumBodies) {\n return;\n }\n\n typename vec4::Type position = oldPos[deviceOffset + index];\n\n typename vec3::Type accel = computeBodyAccel(position, oldPos, numTiles);\n\n // acceleration = force / mass;\n // new velocity = old velocity + acceleration * deltaTime\n // note we factor out the body's mass from the equation, here and in\n // bodyBodyInteraction (because they cancel out). Thus here force ==\n // acceleration\n typename vec4::Type velocity = vel[deviceOffset + index];\n\n velocity.x += accel.x * deltaTime;\n velocity.y += accel.y * deltaTime;\n velocity.z += accel.z * deltaTime;\n\n velocity.x *= damping;\n velocity.y *= damping;\n velocity.z *= damping;\n\n // new position = old position + velocity * deltaTime\n position.x += velocity.x * deltaTime;\n position.y += velocity.y * deltaTime;\n position.z += velocity.z * deltaTime;\n\n // store new position and velocity\n newPos[deviceOffset + index] = position;\n vel[deviceOffset + index] = velocity;\n}\n\ntemplate \nvoid integrateNbodySystem(DeviceData *deviceData,\n cudaGraphicsResource **pgres,\n unsigned int currentRead,\n float deltaTime,\n float damping,\n unsigned int numBodies,\n unsigned int numDevices,\n int blockSize,\n bool bUsePBO)\n{\n if (bUsePBO) {\n checkCudaErrors(cudaGraphicsResourceSetMapFlags(pgres[currentRead], cudaGraphicsMapFlagsReadOnly));\n checkCudaErrors(cudaGraphicsResourceSetMapFlags(pgres[1 - currentRead], cudaGraphicsMapFlagsWriteDiscard));\n checkCudaErrors(cudaGraphicsMapResources(2, pgres, 0));\n size_t bytes;\n checkCudaErrors(cudaGraphicsResourceGetMappedPointer(\n (void **)&(deviceData[0].dPos[currentRead]), &bytes, pgres[currentRead]));\n checkCudaErrors(cudaGraphicsResourceGetMappedPointer(\n (void **)&(deviceData[0].dPos[1 - currentRead]), &bytes, pgres[1 - currentRead]));\n }\n\n for (unsigned int dev = 0; dev != numDevices; dev++) {\n if (numDevices > 1) {\n cudaSetDevice(dev);\n }\n\n int numBlocks = (deviceData[dev].numBodies + blockSize - 1) / blockSize;\n int numTiles = (numBodies + blockSize - 1) / blockSize;\n int sharedMemSize = blockSize * 4 * sizeof(T); // 4 floats for pos\n\n integrateBodies\n <<>>((typename vec4::Type *)deviceData[dev].dPos[1 - currentRead],\n (typename vec4::Type *)deviceData[dev].dPos[currentRead],\n (typename vec4::Type *)deviceData[dev].dVel,\n deviceData[dev].offset,\n deviceData[dev].numBodies,\n deltaTime,\n damping,\n numTiles);\n\n if (numDevices > 1) {\n checkCudaErrors(cudaEventRecord(deviceData[dev].event));\n // MJH: Hack on older driver versions to force kernel launches to flush!\n cudaStreamQuery(0);\n }\n\n // check if kernel invocation generated an error\n getLastCudaError(\"Kernel execution failed\");\n }\n\n if (numDevices > 1) {\n for (unsigned int dev = 0; dev < numDevices; dev++) {\n checkCudaErrors(cudaEventSynchronize(deviceData[dev].event));\n }\n }\n\n if (bUsePBO) {\n checkCudaErrors(cudaGraphicsUnmapResources(2, pgres, 0));\n }\n}\n\n// Explicit specializations needed to generate code\ntemplate void integrateNbodySystem(DeviceData *deviceData,\n cudaGraphicsResource **pgres,\n unsigned int currentRead,\n float deltaTime,\n float damping,\n unsigned int numBodies,\n unsigned int numDevices,\n int blockSize,\n bool bUsePBO);\n\ntemplate void integrateNbodySystem(DeviceData *deviceData,\n cudaGraphicsResource **pgres,\n unsigned int currentRead,\n float deltaTime,\n float damping,\n unsigned int numBodies,\n unsigned int numDevices,\n int blockSize,\n bool bUsePBO);\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n\n// #include \n// #include \n\n// CUDA standard includes\n#include \n// #include \n\n#include \"bodysystem.h\"\n\n__constant__ float softeningSquared;\n__constant__ double softeningSquared_fp64;\n\nhipError_t setSofteningSquared(float softeningSq)\n{\n return hipMemcpyToSymbol(HIP_SYMBOL(softeningSquared), &softeningSq, sizeof(float), 0, hipMemcpyHostToDevice);\n}\n\nhipError_t setSofteningSquared(double softeningSq)\n{\n return hipMemcpyToSymbol(HIP_SYMBOL(softeningSquared_fp64), &softeningSq, sizeof(double), 0, hipMemcpyHostToDevice);\n}\n\ntemplate struct SharedMemory\n{\n __device__ inline operator T *()\n {\n extern __shared__ int __smem[];\n return (T *)__smem;\n }\n\n __device__ inline operator const T *() const\n {\n extern __shared__ int __smem[];\n return (T *)__smem;\n }\n};\n\ntemplate __device__ T rsqrt_T(T x) { return rsqrt(x); }\n\ntemplate <> __device__ float rsqrt_T(float x) { return rsqrtf(x); }\n\ntemplate <> __device__ double rsqrt_T(double x) { return rsqrt(x); }\n\n// Macros to simplify shared memory addressing\n#define SX(i) sharedPos[i + blockDim.x * threadIdx.y]\n// This macro is only used when multithreadBodies is true (below)\n#define SX_SUM(i, j) sharedPos[i + blockDim.x * j]\n\ntemplate __device__ T getSofteningSquared() { return softeningSquared; }\ntemplate <> __device__ double getSofteningSquared() { return softeningSquared_fp64; }\n\ntemplate struct DeviceData\n{\n T *dPos[2]; // mapped host pointers\n T *dVel;\n hipEvent_t event;\n unsigned int offset;\n unsigned int numBodies;\n};\n\ntemplate \n__device__ typename vec3::Type\nbodyBodyInteraction(typename vec3::Type ai, typename vec4::Type bi, typename vec4::Type bj)\n{\n typename vec3::Type r;\n\n // r_ij [3 FLOPS]\n r.x = bj.x - bi.x;\n r.y = bj.y - bi.y;\n r.z = bj.z - bi.z;\n\n // distSqr = dot(r_ij, r_ij) + EPS^2 [6 FLOPS]\n T distSqr = r.x * r.x + r.y * r.y + r.z * r.z;\n distSqr += getSofteningSquared();\n\n // invDistCube =1/distSqr^(3/2) [4 FLOPS (2 mul, 1 sqrt, 1 inv)]\n T invDist = rsqrt_T(distSqr);\n T invDistCube = invDist * invDist * invDist;\n\n // s = m_j * invDistCube [1 FLOP]\n T s = bj.w * invDistCube;\n\n // a_i = a_i + s * r_ij [6 FLOPS]\n ai.x += r.x * s;\n ai.y += r.y * s;\n ai.z += r.z * s;\n\n return ai;\n}\n\ntemplate \n__device__ typename vec3::Type\ncomputeBodyAccel(typename vec4::Type bodyPos, typename vec4::Type *positions, int numTiles)\n{\n typename vec4::Type *sharedPos = SharedMemory::Type>();\n\n typename vec3::Type acc = {0.0f, 0.0f, 0.0f};\n\n for (int tile = 0; tile < numTiles; tile++) {\n sharedPos[threadIdx.x] = positions[tile * blockDim.x + threadIdx.x];\n\n __syncthreads();\n\n // This is the \"tile_calculation\" from the GPUG3 article.\n#pragma unroll 128\n\n for (unsigned int counter = 0; counter < blockDim.x; counter++) {\n acc = bodyBodyInteraction(acc, bodyPos, sharedPos[counter]);\n }\n\n __syncthreads();\n }\n\n return acc;\n}\n\ntemplate \n__global__ void integrateBodies(typename vec4::Type *__restrict__ newPos,\n typename vec4::Type *__restrict__ oldPos,\n typename vec4::Type *vel,\n unsigned int deviceOffset,\n unsigned int deviceNumBodies,\n float deltaTime,\n float damping,\n int numTiles)\n{\n int index = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (index >= deviceNumBodies) {\n return;\n }\n\n typename vec4::Type position = oldPos[deviceOffset + index];\n\n typename vec3::Type accel = computeBodyAccel(position, oldPos, numTiles);\n\n // acceleration = force / mass;\n // new velocity = old velocity + acceleration * deltaTime\n // note we factor out the body's mass from the equation, here and in\n // bodyBodyInteraction (because they cancel out). Thus here force ==\n // acceleration\n typename vec4::Type velocity = vel[deviceOffset + index];\n\n velocity.x += accel.x * deltaTime;\n velocity.y += accel.y * deltaTime;\n velocity.z += accel.z * deltaTime;\n\n velocity.x *= damping;\n velocity.y *= damping;\n velocity.z *= damping;\n\n // new position = old position + velocity * deltaTime\n position.x += velocity.x * deltaTime;\n position.y += velocity.y * deltaTime;\n position.z += velocity.z * deltaTime;\n\n // store new position and velocity\n newPos[deviceOffset + index] = position;\n vel[deviceOffset + index] = velocity;\n}\n\ntemplate \nvoid integrateNbodySystem(DeviceData *deviceData,\n hipGraphicsResource **pgres,\n unsigned int currentRead,\n float deltaTime,\n float damping,\n unsigned int numBodies,\n unsigned int numDevices,\n int blockSize,\n bool bUsePBO)\n{\n if (bUsePBO) {\n checkCudaErrors(cudaGraphicsResourceSetMapFlags(pgres[currentRead], cudaGraphicsMapFlagsReadOnly));\n checkCudaErrors(cudaGraphicsResourceSetMapFlags(pgres[1 - currentRead], cudaGraphicsMapFlagsWriteDiscard));\n checkCudaErrors(hipGraphicsMapResources(2, pgres, 0));\n size_t bytes;\n checkCudaErrors(hipGraphicsResourceGetMappedPointer(\n (void **)&(deviceData[0].dPos[currentRead]), &bytes, pgres[currentRead]));\n checkCudaErrors(hipGraphicsResourceGetMappedPointer(\n (void **)&(deviceData[0].dPos[1 - currentRead]), &bytes, pgres[1 - currentRead]));\n }\n\n for (unsigned int dev = 0; dev != numDevices; dev++) {\n if (numDevices > 1) {\n hipSetDevice(dev);\n }\n\n int numBlocks = (deviceData[dev].numBodies + blockSize - 1) / blockSize;\n int numTiles = (numBodies + blockSize - 1) / blockSize;\n int sharedMemSize = blockSize * 4 * sizeof(T); // 4 floats for pos\n\n integrateBodies\n <<>>((typename vec4::Type *)deviceData[dev].dPos[1 - currentRead],\n (typename vec4::Type *)deviceData[dev].dPos[currentRead],\n (typename vec4::Type *)deviceData[dev].dVel,\n deviceData[dev].offset,\n deviceData[dev].numBodies,\n deltaTime,\n damping,\n numTiles);\n\n if (numDevices > 1) {\n checkCudaErrors(hipEventRecord(deviceData[dev].event));\n // MJH: Hack on older driver versions to force kernel launches to flush!\n hipStreamQuery(0);\n }\n\n // check if kernel invocation generated an error\n getLastCudaError(\"Kernel execution failed\");\n }\n\n if (numDevices > 1) {\n for (unsigned int dev = 0; dev < numDevices; dev++) {\n checkCudaErrors(hipEventSynchronize(deviceData[dev].event));\n }\n }\n\n if (bUsePBO) {\n checkCudaErrors(hipGraphicsUnmapResources(2, pgres, 0));\n }\n}\n\n// Explicit specializations needed to generate code\ntemplate void integrateNbodySystem(DeviceData *deviceData,\n hipGraphicsResource **pgres,\n unsigned int currentRead,\n float deltaTime,\n float damping,\n unsigned int numBodies,\n unsigned int numDevices,\n int blockSize,\n bool bUsePBO);\n\ntemplate void integrateNbodySystem(DeviceData *deviceData,\n hipGraphicsResource **pgres,\n unsigned int currentRead,\n float deltaTime,\n float damping,\n unsigned int numBodies,\n unsigned int numDevices,\n int blockSize,\n bool bUsePBO);\n", "source": "/kaggle/working/cuda-samples/Samples/8_Platform_Specific/Tegra/nbody_opengles/bodysystemcuda.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n\n#include \"cuda_consumer.h\"\n#include \"nvmedia_image_nvscibuf.h\"\n#include \"nvmedia_utils/cmdline.h\"\n\n// Enable this to 1 if require cuda processed output to ppm file.\n#define WRITE_OUTPUT_IMAGE 0\n\n#define checkNvSciErrors(call) \\\n do { \\\n NvSciError _status = call; \\\n if (NvSciError_Success != _status) { \\\n printf(\"NVSCI call in file '%s' in line %i returned\" \\\n \" %d, expected %d\\n\", \\\n __FILE__, \\\n __LINE__, \\\n _status, \\\n NvSciError_Success); \\\n fflush(stdout); \\\n exit(EXIT_FAILURE); \\\n } \\\n } while (0)\n\n__global__ static void\nyuvToGrayscale(cudaSurfaceObject_t surfaceObject, unsigned int *dstImage, int32_t imageWidth, int32_t imageHeight)\n{\n size_t x = blockIdx.x * blockDim.x + threadIdx.x;\n size_t y = blockIdx.y * blockDim.y + threadIdx.y;\n\n uchar4 *dstImageUchar4 = (uchar4 *)dstImage;\n for (; x < imageWidth && y < imageHeight; x += gridDim.x * blockDim.x, y += gridDim.y * blockDim.y) {\n int colInBytes = x * sizeof(unsigned char);\n unsigned char luma = surf2Dread(surfaceObject, colInBytes, y);\n uchar4 grayscalePix = make_uchar4(luma, luma, luma, 0);\n\n dstImageUchar4[y * imageWidth + x] = grayscalePix;\n }\n}\n\nstatic void cudaImportNvSciSync(cudaExternalSemaphore_t &extSem, NvSciSyncObj &syncObj)\n{\n cudaExternalSemaphoreHandleDesc extSemDesc;\n memset(&extSemDesc, 0, sizeof(extSemDesc));\n extSemDesc.type = cudaExternalSemaphoreHandleTypeNvSciSync;\n extSemDesc.handle.nvSciSyncObj = (void *)syncObj;\n\n checkCudaErrors(cudaImportExternalSemaphore(&extSem, &extSemDesc));\n}\n\nstatic void waitExternalSemaphore(cudaExternalSemaphore_t &waitSem, NvSciSyncFence *fence, cudaStream_t stream)\n{\n cudaExternalSemaphoreWaitParams waitParams;\n memset(&waitParams, 0, sizeof(waitParams));\n // For cross-process signaler-waiter applications need to use NvSciIpc\n // and NvSciSync[Export|Import] utilities to share the NvSciSyncFence\n // across process. This step is optional in single-process.\n waitParams.params.nvSciSync.fence = (void *)fence;\n waitParams.flags = 0;\n\n checkCudaErrors(cudaWaitExternalSemaphoresAsync(&waitSem, &waitParams, 1, stream));\n}\n\nstatic void signalExternalSemaphore(cudaExternalSemaphore_t &signalSem, NvSciSyncFence *fence, cudaStream_t stream)\n{\n cudaExternalSemaphoreSignalParams signalParams;\n memset(&signalParams, 0, sizeof(signalParams));\n // For cross-process signaler-waiter applications need to use NvSciIpc\n // and NvSciSync[Export|Import] utilities to share the NvSciSyncFence\n // across process. This step is optional in single-process.\n signalParams.params.nvSciSync.fence = (void *)fence;\n signalParams.flags = 0;\n\n checkCudaErrors(cudaSignalExternalSemaphoresAsync(&signalSem, &signalParams, 1, stream));\n}\n\nstatic void yuvToGrayscaleCudaKernel(cudaExternalResInterop &cudaExtResObj, int32_t imageWidth, int32_t imageHeight)\n{\n#if WRITE_OUTPUT_IMAGE\n unsigned int *h_dstImage;\n checkCudaErrors(cudaMallocHost(&h_dstImage, sizeof(unsigned int) * imageHeight * imageWidth));\n#endif\n dim3 block(16, 16, 1);\n dim3 grid((imageWidth / block.x) + 1, (imageHeight / block.y) + 1, 1);\n\n yuvToGrayscale<<>>(\n cudaExtResObj.cudaSurfaceNvmediaBuf[0], cudaExtResObj.d_outputImage, imageWidth, imageHeight);\n\n#if WRITE_OUTPUT_IMAGE\n checkCudaErrors(cudaMemcpyAsync(h_dstImage,\n cudaExtResObj.d_outputImage,\n sizeof(unsigned int) * imageHeight * imageWidth,\n cudaMemcpyDeviceToHost,\n cudaExtResObj.stream));\n checkCudaErrors(cudaStreamSynchronize(cudaExtResObj.stream));\n char outputFilename[1024];\n std::string image_filename = \"Grayscale\";\n strcpy(outputFilename, image_filename.c_str());\n strcpy(outputFilename + image_filename.length(), \"_nvsci_out.ppm\");\n sdkSavePPM4ub(outputFilename, (unsigned char *)h_dstImage, imageWidth, imageHeight);\n printf(\"Wrote '%s'\\n\", outputFilename);\n checkCudaErrors(cudaFreeHost(h_dstImage));\n#endif\n}\n\nstatic void cudaImportNvSciImage(cudaExternalResInterop &cudaExtResObj, NvSciBufObj &inputBufObj)\n{\n NvSciBufModule module = NULL;\n NvSciBufAttrList attrlist = NULL;\n NvSciBufAttrKeyValuePair pairArrayOut[10];\n\n checkNvSciErrors(NvSciBufModuleOpen(&module));\n checkNvSciErrors(NvSciBufAttrListCreate(module, &attrlist));\n checkNvSciErrors(NvSciBufObjGetAttrList(inputBufObj, &attrlist));\n\n memset(pairArrayOut, 0, sizeof(NvSciBufAttrKeyValuePair) * 10);\n\n int numAttrs = 0;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_Size;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_PlaneChannelCount;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_PlaneCount;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_PlaneWidth;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_PlaneHeight;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_Layout;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_PlaneBitsPerPixel;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_PlaneOffset;\n\n checkNvSciErrors(NvSciBufAttrListGetAttrs(attrlist, pairArrayOut, numAttrs));\n\n uint64_t size = *(uint64_t *)pairArrayOut[0].value;\n uint8_t channelCount = *(uint8_t *)pairArrayOut[1].value;\n cudaExtResObj.planeCount = *(int32_t *)pairArrayOut[2].value;\n cudaExtResObj.imageWidth = (int32_t *)malloc(sizeof(int32_t) * cudaExtResObj.planeCount);\n cudaExtResObj.imageHeight = (int32_t *)malloc(sizeof(int32_t) * cudaExtResObj.planeCount);\n cudaExtResObj.planeOffset = (uint64_t *)malloc(sizeof(uint64_t) * cudaExtResObj.planeCount);\n\n memcpy(cudaExtResObj.imageWidth, (int32_t *)pairArrayOut[3].value, cudaExtResObj.planeCount * sizeof(int32_t));\n memcpy(cudaExtResObj.imageHeight, (int32_t *)pairArrayOut[4].value, cudaExtResObj.planeCount * sizeof(int32_t));\n memcpy(cudaExtResObj.planeOffset, (uint64_t *)pairArrayOut[7].value, cudaExtResObj.planeCount * sizeof(uint64_t));\n\n NvSciBufAttrValImageLayoutType layout = *(NvSciBufAttrValImageLayoutType *)pairArrayOut[5].value;\n uint32_t bitsPerPixel = *(uint32_t *)pairArrayOut[6].value;\n\n if (layout != NvSciBufImage_BlockLinearType) {\n printf(\"Image layout is not block linear.. waiving execution\\n\");\n exit(EXIT_WAIVED);\n }\n\n cudaExternalMemoryHandleDesc memHandleDesc;\n memset(&memHandleDesc, 0, sizeof(memHandleDesc));\n memHandleDesc.type = cudaExternalMemoryHandleTypeNvSciBuf;\n memHandleDesc.handle.nvSciBufObject = inputBufObj;\n memHandleDesc.size = size;\n checkCudaErrors(cudaImportExternalMemory(&cudaExtResObj.extMemImageBuf, &memHandleDesc));\n\n cudaExtResObj.d_mipmapArray =\n (cudaMipmappedArray_t *)malloc(sizeof(cudaMipmappedArray_t) * cudaExtResObj.planeCount);\n\n for (int i = 0; i < cudaExtResObj.planeCount; i++) {\n cudaExtent extent = {};\n memset(&extent, 0, sizeof(extent));\n extent.width = cudaExtResObj.imageWidth[i];\n extent.height = cudaExtResObj.imageHeight[i];\n extent.depth = 0;\n cudaChannelFormatDesc desc;\n switch (channelCount) {\n case 1:\n default:\n desc = cudaCreateChannelDesc(bitsPerPixel, 0, 0, 0, cudaChannelFormatKindUnsigned);\n break;\n case 2:\n desc = cudaCreateChannelDesc(bitsPerPixel, bitsPerPixel, 0, 0, cudaChannelFormatKindUnsigned);\n break;\n case 3:\n desc = cudaCreateChannelDesc(bitsPerPixel, bitsPerPixel, bitsPerPixel, 0, cudaChannelFormatKindUnsigned);\n break;\n case 4:\n desc = cudaCreateChannelDesc(\n bitsPerPixel, bitsPerPixel, bitsPerPixel, bitsPerPixel, cudaChannelFormatKindUnsigned);\n break;\n }\n\n cudaExternalMemoryMipmappedArrayDesc mipmapDesc = {0};\n mipmapDesc.offset = cudaExtResObj.planeOffset[i];\n mipmapDesc.formatDesc = desc;\n mipmapDesc.extent = extent;\n mipmapDesc.flags = 0;\n mipmapDesc.numLevels = 1;\n checkCudaErrors(cudaExternalMemoryGetMappedMipmappedArray(\n &cudaExtResObj.d_mipmapArray[i], cudaExtResObj.extMemImageBuf, &mipmapDesc));\n }\n}\n\nstatic cudaSurfaceObject_t createCudaSurface(cudaArray_t &d_mipLevelArray)\n{\n cudaResourceDesc resourceDesc;\n memset(&resourceDesc, 0, sizeof(resourceDesc));\n resourceDesc.resType = cudaResourceTypeArray;\n resourceDesc.res.array.array = d_mipLevelArray;\n\n cudaSurfaceObject_t surfaceObject;\n checkCudaErrors(cudaCreateSurfaceObject(&surfaceObject, &resourceDesc));\n return surfaceObject;\n}\n\nstatic cudaStream_t createCudaStream(int deviceId)\n{\n checkCudaErrors(cudaSetDevice(deviceId));\n cudaStream_t stream;\n checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));\n return stream;\n}\n\n// CUDA setup buffers/synchronization objects for interop via NvSci API.\nvoid setupCuda(cudaExternalResInterop &cudaExtResObj,\n NvSciBufObj &inputBufObj,\n NvSciSyncObj &syncObj,\n NvSciSyncObj &cudaSignalerSyncObj,\n int deviceId)\n{\n checkCudaErrors(cudaSetDevice(deviceId));\n cudaImportNvSciSync(cudaExtResObj.waitSem, syncObj);\n cudaImportNvSciSync(cudaExtResObj.signalSem, cudaSignalerSyncObj);\n\n cudaImportNvSciImage(cudaExtResObj, inputBufObj);\n cudaExtResObj.d_mipLevelArray = (cudaArray_t *)malloc(sizeof(cudaArray_t) * cudaExtResObj.planeCount);\n cudaExtResObj.cudaSurfaceNvmediaBuf =\n (cudaSurfaceObject_t *)malloc(sizeof(cudaSurfaceObject_t) * cudaExtResObj.planeCount);\n\n for (int i = 0; i < cudaExtResObj.planeCount; ++i) {\n uint32_t mipLevelId = 0;\n checkCudaErrors(\n cudaGetMipmappedArrayLevel(&cudaExtResObj.d_mipLevelArray[i], cudaExtResObj.d_mipmapArray[i], mipLevelId));\n cudaExtResObj.cudaSurfaceNvmediaBuf[i] = createCudaSurface(cudaExtResObj.d_mipLevelArray[i]);\n }\n\n cudaExtResObj.stream = createCudaStream(deviceId);\n checkCudaErrors(cudaMalloc(&cudaExtResObj.d_outputImage,\n sizeof(unsigned int) * cudaExtResObj.imageWidth[0] * cudaExtResObj.imageHeight[0]));\n}\n\n// CUDA clean up buffers used **with** NvSci API.\nvoid cleanupCuda(cudaExternalResInterop &cudaExtResObj)\n{\n for (int i = 0; i < cudaExtResObj.planeCount; i++) {\n checkCudaErrors(cudaDestroySurfaceObject(cudaExtResObj.cudaSurfaceNvmediaBuf[i]));\n checkCudaErrors(cudaFreeMipmappedArray(cudaExtResObj.d_mipmapArray[i]));\n }\n free(cudaExtResObj.d_mipmapArray);\n free(cudaExtResObj.d_mipLevelArray);\n free(cudaExtResObj.cudaSurfaceNvmediaBuf);\n free(cudaExtResObj.imageWidth);\n free(cudaExtResObj.imageHeight);\n checkCudaErrors(cudaDestroyExternalSemaphore(cudaExtResObj.waitSem));\n checkCudaErrors(cudaDestroyExternalSemaphore(cudaExtResObj.signalSem));\n checkCudaErrors(cudaDestroyExternalMemory(cudaExtResObj.extMemImageBuf));\n checkCudaErrors(cudaStreamDestroy(cudaExtResObj.stream));\n checkCudaErrors(cudaFree(cudaExtResObj.d_outputImage));\n}\n\nvoid runCudaOperation(cudaExternalResInterop &cudaExtResObj,\n NvSciSyncFence *cudaWaitFence,\n NvSciSyncFence *cudaSignalFence,\n int deviceId,\n int iterations)\n{\n checkCudaErrors(cudaSetDevice(deviceId));\n static int64_t launch = 0;\n\n waitExternalSemaphore(cudaExtResObj.waitSem, cudaWaitFence, cudaExtResObj.stream);\n\n // run cuda kernel over surface object of the LUMA surface part to extract\n // grayscale.\n yuvToGrayscaleCudaKernel(cudaExtResObj, cudaExtResObj.imageWidth[0], cudaExtResObj.imageHeight[0]);\n\n // signal fence till the second last iterations for NvMedia2DBlit to wait for\n // cuda signal and for final iteration as there is no corresponding NvMedia\n // operation pending therefore we end with cudaStreamSynchronize()\n if (launch < iterations - 1) {\n signalExternalSemaphore(cudaExtResObj.signalSem, cudaSignalFence, cudaExtResObj.stream);\n }\n else {\n checkCudaErrors(cudaStreamSynchronize(cudaExtResObj.stream));\n }\n launch++;\n}\n\n// CUDA imports and operates on NvSci buffer/synchronization objects\nvoid setupCuda(Blit2DTest *ctx, cudaResources &cudaResObj, int deviceId)\n{\n checkCudaErrors(cudaSetDevice(deviceId));\n cudaResObj.d_yuvArray = (cudaArray_t *)malloc(sizeof(cudaArray_t) * ctx->numSurfaces);\n cudaResObj.cudaSurfaceNvmediaBuf = (cudaSurfaceObject_t *)malloc(sizeof(cudaSurfaceObject_t) * ctx->numSurfaces);\n cudaChannelFormatDesc channelDesc;\n switch (ctx->bytesPerPixel) {\n case 1:\n default:\n channelDesc = cudaCreateChannelDesc(8, 0, 0, 0, cudaChannelFormatKindUnsigned);\n break;\n }\n\n for (int k = 0; k < ctx->numSurfaces; k++) {\n checkCudaErrors(cudaMallocArray(&cudaResObj.d_yuvArray[k],\n &channelDesc,\n ctx->widthSurface * ctx->xScalePtr[k] * ctx->bytesPerPixel,\n ctx->heightSurface * ctx->yScalePtr[k]));\n cudaResObj.cudaSurfaceNvmediaBuf[k] = createCudaSurface(cudaResObj.d_yuvArray[k]);\n }\n checkCudaErrors(\n cudaMalloc(&cudaResObj.d_outputImage, sizeof(unsigned int) * ctx->widthSurface * ctx->heightSurface));\n\n cudaResObj.stream = createCudaStream(deviceId);\n}\n\n// CUDA clean up buffers used **without** NvSci API.\nvoid cleanupCuda(Blit2DTest *ctx, cudaResources &cudaResObj)\n{\n for (int k = 0; k < ctx->numSurfaces; k++) {\n checkCudaErrors(cudaDestroySurfaceObject(cudaResObj.cudaSurfaceNvmediaBuf[k]));\n checkCudaErrors(cudaFreeArray(cudaResObj.d_yuvArray[k]));\n }\n\n free(cudaResObj.cudaSurfaceNvmediaBuf);\n\n checkCudaErrors(cudaStreamDestroy(cudaResObj.stream));\n checkCudaErrors(cudaFree(cudaResObj.d_outputImage));\n}\n\nstatic void\nyuvToGrayscaleCudaKernelNonNvSci(cudaResources &cudaResObj, int deviceId, int32_t imageWidth, int32_t imageHeight)\n{\n#if WRITE_OUTPUT_IMAGE\n unsigned int *h_dstImage;\n checkCudaErrors(cudaMallocHost(&h_dstImage, sizeof(unsigned int) * imageHeight * imageWidth));\n#endif\n dim3 block(16, 16, 1);\n dim3 grid((imageWidth / block.x) + 1, (imageHeight / block.y) + 1, 1);\n\n yuvToGrayscale<<>>(\n cudaResObj.cudaSurfaceNvmediaBuf[0], cudaResObj.d_outputImage, imageWidth, imageHeight);\n\n#if WRITE_OUTPUT_IMAGE\n checkCudaErrors(cudaMemcpyAsync(h_dstImage,\n cudaResObj.d_outputImage,\n sizeof(unsigned int) * imageHeight * imageWidth,\n cudaMemcpyDeviceToHost,\n cudaResObj.stream));\n checkCudaErrors(cudaStreamSynchronize(cudaResObj.stream));\n char outputFilename[1024];\n std::string image_filename = \"Grayscale\";\n strcpy(outputFilename, image_filename.c_str());\n strcpy(outputFilename + image_filename.length(), \"_non-nvsci_out.ppm\");\n sdkSavePPM4ub(outputFilename, (unsigned char *)h_dstImage, imageWidth, imageHeight);\n printf(\"Wrote '%s'\\n\", outputFilename);\n checkCudaErrors(cudaFreeHost(h_dstImage));\n#else\n checkCudaErrors(cudaStreamSynchronize(cudaResObj.stream));\n#endif\n}\n\n// CUDA operates **without** NvSci APIs buffer/synchronization objects.\nvoid runCudaOperation(Blit2DTest *ctx, cudaResources &cudaResObj, int deviceId)\n{\n for (int k = 0; k < ctx->numSurfaces; k++) {\n checkCudaErrors(cudaMemcpy2DToArray(cudaResObj.d_yuvArray[k],\n 0,\n 0,\n ctx->dstBuff[k],\n ctx->widthSurface * ctx->xScalePtr[k] * ctx->bytesPerPixel,\n ctx->widthSurface * ctx->xScalePtr[k] * ctx->bytesPerPixel,\n ctx->heightSurface * ctx->yScalePtr[k],\n cudaMemcpyHostToDevice));\n }\n // run cuda kernel over surface object of the LUMA surface part to extract\n // grayscale.\n yuvToGrayscaleCudaKernelNonNvSci(cudaResObj, deviceId, ctx->widthSurface, ctx->heightSurface);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n\n#include \"cuda_consumer.h\"\n#include \"nvmedia_image_nvscibuf.h\"\n#include \"nvmedia_utils/cmdline.h\"\n\n// Enable this to 1 if require cuda processed output to ppm file.\n#define WRITE_OUTPUT_IMAGE 0\n\n#define checkNvSciErrors(call) \\\n do { \\\n NvSciError _status = call; \\\n if (NvSciError_Success != _status) { \\\n printf(\"NVSCI call in file '%s' in line %i returned\" \\\n \" %d, expected %d\\n\", \\\n __FILE__, \\\n __LINE__, \\\n _status, \\\n NvSciError_Success); \\\n fflush(stdout); \\\n exit(EXIT_FAILURE); \\\n } \\\n } while (0)\n\n__global__ static void\nyuvToGrayscale(hipSurfaceObject_t surfaceObject, unsigned int *dstImage, int32_t imageWidth, int32_t imageHeight)\n{\n size_t x = blockIdx.x * blockDim.x + threadIdx.x;\n size_t y = blockIdx.y * blockDim.y + threadIdx.y;\n\n uchar4 *dstImageUchar4 = (uchar4 *)dstImage;\n for (; x < imageWidth && y < imageHeight; x += gridDim.x * blockDim.x, y += gridDim.y * blockDim.y) {\n int colInBytes = x * sizeof(unsigned char);\n unsigned char luma = surf2Dread(surfaceObject, colInBytes, y);\n uchar4 grayscalePix = make_uchar4(luma, luma, luma, 0);\n\n dstImageUchar4[y * imageWidth + x] = grayscalePix;\n }\n}\n\nstatic void cudaImportNvSciSync(hipExternalSemaphore_t &extSem, NvSciSyncObj &syncObj)\n{\n hipExternalSemaphoreHandleDesc extSemDesc;\n memset(&extSemDesc, 0, sizeof(extSemDesc));\n extSemDesc.type = cudaExternalSemaphoreHandleTypeNvSciSync;\n extSemDesc.handle.nvSciSyncObj = (void *)syncObj;\n\n checkCudaErrors(hipImportExternalSemaphore(&extSem, &extSemDesc));\n}\n\nstatic void waitExternalSemaphore(hipExternalSemaphore_t &waitSem, NvSciSyncFence *fence, hipStream_t stream)\n{\n hipExternalSemaphoreWaitParams waitParams;\n memset(&waitParams, 0, sizeof(waitParams));\n // For cross-process signaler-waiter applications need to use NvSciIpc\n // and NvSciSync[Export|Import] utilities to share the NvSciSyncFence\n // across process. This step is optional in single-process.\n waitParams.params.nvSciSync.fence = (void *)fence;\n waitParams.flags = 0;\n\n checkCudaErrors(hipWaitExternalSemaphoresAsync(&waitSem, &waitParams, 1, stream));\n}\n\nstatic void signalExternalSemaphore(hipExternalSemaphore_t &signalSem, NvSciSyncFence *fence, hipStream_t stream)\n{\n hipExternalSemaphoreSignalParams signalParams;\n memset(&signalParams, 0, sizeof(signalParams));\n // For cross-process signaler-waiter applications need to use NvSciIpc\n // and NvSciSync[Export|Import] utilities to share the NvSciSyncFence\n // across process. This step is optional in single-process.\n signalParams.params.nvSciSync.fence = (void *)fence;\n signalParams.flags = 0;\n\n checkCudaErrors(hipSignalExternalSemaphoresAsync(&signalSem, &signalParams, 1, stream));\n}\n\nstatic void yuvToGrayscaleCudaKernel(cudaExternalResInterop &cudaExtResObj, int32_t imageWidth, int32_t imageHeight)\n{\n#if WRITE_OUTPUT_IMAGE\n unsigned int *h_dstImage;\n checkCudaErrors(hipHostMalloc(&h_dstImage, sizeof(unsigned int) * imageHeight * imageWidth));\n#endif\n dim3 block(16, 16, 1);\n dim3 grid((imageWidth / block.x) + 1, (imageHeight / block.y) + 1, 1);\n\n yuvToGrayscale<<>>(\n cudaExtResObj.cudaSurfaceNvmediaBuf[0], cudaExtResObj.d_outputImage, imageWidth, imageHeight);\n\n#if WRITE_OUTPUT_IMAGE\n checkCudaErrors(hipMemcpyAsync(h_dstImage,\n cudaExtResObj.d_outputImage,\n sizeof(unsigned int) * imageHeight * imageWidth,\n hipMemcpyDeviceToHost,\n cudaExtResObj.stream));\n checkCudaErrors(hipStreamSynchronize(cudaExtResObj.stream));\n char outputFilename[1024];\n std::string image_filename = \"Grayscale\";\n strcpy(outputFilename, image_filename.c_str());\n strcpy(outputFilename + image_filename.length(), \"_nvsci_out.ppm\");\n sdkSavePPM4ub(outputFilename, (unsigned char *)h_dstImage, imageWidth, imageHeight);\n printf(\"Wrote '%s'\\n\", outputFilename);\n checkCudaErrors(hipHostFree(h_dstImage));\n#endif\n}\n\nstatic void cudaImportNvSciImage(cudaExternalResInterop &cudaExtResObj, NvSciBufObj &inputBufObj)\n{\n NvSciBufModule module = NULL;\n NvSciBufAttrList attrlist = NULL;\n NvSciBufAttrKeyValuePair pairArrayOut[10];\n\n checkNvSciErrors(NvSciBufModuleOpen(&module));\n checkNvSciErrors(NvSciBufAttrListCreate(module, &attrlist));\n checkNvSciErrors(NvSciBufObjGetAttrList(inputBufObj, &attrlist));\n\n memset(pairArrayOut, 0, sizeof(NvSciBufAttrKeyValuePair) * 10);\n\n int numAttrs = 0;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_Size;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_PlaneChannelCount;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_PlaneCount;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_PlaneWidth;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_PlaneHeight;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_Layout;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_PlaneBitsPerPixel;\n pairArrayOut[numAttrs++].key = NvSciBufImageAttrKey_PlaneOffset;\n\n checkNvSciErrors(NvSciBufAttrListGetAttrs(attrlist, pairArrayOut, numAttrs));\n\n uint64_t size = *(uint64_t *)pairArrayOut[0].value;\n uint8_t channelCount = *(uint8_t *)pairArrayOut[1].value;\n cudaExtResObj.planeCount = *(int32_t *)pairArrayOut[2].value;\n cudaExtResObj.imageWidth = (int32_t *)malloc(sizeof(int32_t) * cudaExtResObj.planeCount);\n cudaExtResObj.imageHeight = (int32_t *)malloc(sizeof(int32_t) * cudaExtResObj.planeCount);\n cudaExtResObj.planeOffset = (uint64_t *)malloc(sizeof(uint64_t) * cudaExtResObj.planeCount);\n\n memcpy(cudaExtResObj.imageWidth, (int32_t *)pairArrayOut[3].value, cudaExtResObj.planeCount * sizeof(int32_t));\n memcpy(cudaExtResObj.imageHeight, (int32_t *)pairArrayOut[4].value, cudaExtResObj.planeCount * sizeof(int32_t));\n memcpy(cudaExtResObj.planeOffset, (uint64_t *)pairArrayOut[7].value, cudaExtResObj.planeCount * sizeof(uint64_t));\n\n NvSciBufAttrValImageLayoutType layout = *(NvSciBufAttrValImageLayoutType *)pairArrayOut[5].value;\n uint32_t bitsPerPixel = *(uint32_t *)pairArrayOut[6].value;\n\n if (layout != NvSciBufImage_BlockLinearType) {\n printf(\"Image layout is not block linear.. waiving execution\\n\");\n exit(EXIT_WAIVED);\n }\n\n hipExternalMemoryHandleDesc memHandleDesc;\n memset(&memHandleDesc, 0, sizeof(memHandleDesc));\n memHandleDesc.type = cudaExternalMemoryHandleTypeNvSciBuf;\n memHandleDesc.handle.nvSciBufObject = inputBufObj;\n memHandleDesc.size = size;\n checkCudaErrors(hipImportExternalMemory(&cudaExtResObj.extMemImageBuf, &memHandleDesc));\n\n cudaExtResObj.d_mipmapArray =\n (hipMipmappedArray_t *)malloc(sizeof(hipMipmappedArray_t) * cudaExtResObj.planeCount);\n\n for (int i = 0; i < cudaExtResObj.planeCount; i++) {\n hipExtent extent = {};\n memset(&extent, 0, sizeof(extent));\n extent.width = cudaExtResObj.imageWidth[i];\n extent.height = cudaExtResObj.imageHeight[i];\n extent.depth = 0;\n hipChannelFormatDesc desc;\n switch (channelCount) {\n case 1:\n default:\n desc = hipCreateChannelDesc(bitsPerPixel, 0, 0, 0, hipChannelFormatKindUnsigned);\n break;\n case 2:\n desc = hipCreateChannelDesc(bitsPerPixel, bitsPerPixel, 0, 0, hipChannelFormatKindUnsigned);\n break;\n case 3:\n desc = hipCreateChannelDesc(bitsPerPixel, bitsPerPixel, bitsPerPixel, 0, hipChannelFormatKindUnsigned);\n break;\n case 4:\n desc = hipCreateChannelDesc(\n bitsPerPixel, bitsPerPixel, bitsPerPixel, bitsPerPixel, hipChannelFormatKindUnsigned);\n break;\n }\n\n cudaExternalMemoryMipmappedArrayDesc mipmapDesc = {0};\n mipmapDesc.offset = cudaExtResObj.planeOffset[i];\n mipmapDesc.formatDesc = desc;\n mipmapDesc.extent = extent;\n mipmapDesc.flags = 0;\n mipmapDesc.numLevels = 1;\n checkCudaErrors(cudaExternalMemoryGetMappedMipmappedArray(\n &cudaExtResObj.d_mipmapArray[i], cudaExtResObj.extMemImageBuf, &mipmapDesc));\n }\n}\n\nstatic hipSurfaceObject_t createCudaSurface(hipArray_t &d_mipLevelArray)\n{\n hipResourceDesc resourceDesc;\n memset(&resourceDesc, 0, sizeof(resourceDesc));\n resourceDesc.resType = hipResourceTypeArray;\n resourceDesc.res.array.array = d_mipLevelArray;\n\n hipSurfaceObject_t surfaceObject;\n checkCudaErrors(hipCreateSurfaceObject(&surfaceObject, &resourceDesc));\n return surfaceObject;\n}\n\nstatic hipStream_t createCudaStream(int deviceId)\n{\n checkCudaErrors(hipSetDevice(deviceId));\n hipStream_t stream;\n checkCudaErrors(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));\n return stream;\n}\n\n// CUDA setup buffers/synchronization objects for interop via NvSci API.\nvoid setupCuda(cudaExternalResInterop &cudaExtResObj,\n NvSciBufObj &inputBufObj,\n NvSciSyncObj &syncObj,\n NvSciSyncObj &cudaSignalerSyncObj,\n int deviceId)\n{\n checkCudaErrors(hipSetDevice(deviceId));\n cudaImportNvSciSync(cudaExtResObj.waitSem, syncObj);\n cudaImportNvSciSync(cudaExtResObj.signalSem, cudaSignalerSyncObj);\n\n cudaImportNvSciImage(cudaExtResObj, inputBufObj);\n cudaExtResObj.d_mipLevelArray = (hipArray_t *)malloc(sizeof(hipArray_t) * cudaExtResObj.planeCount);\n cudaExtResObj.cudaSurfaceNvmediaBuf =\n (hipSurfaceObject_t *)malloc(sizeof(hipSurfaceObject_t) * cudaExtResObj.planeCount);\n\n for (int i = 0; i < cudaExtResObj.planeCount; ++i) {\n uint32_t mipLevelId = 0;\n checkCudaErrors(\n hipGetMipmappedArrayLevel(&cudaExtResObj.d_mipLevelArray[i], cudaExtResObj.d_mipmapArray[i], mipLevelId));\n cudaExtResObj.cudaSurfaceNvmediaBuf[i] = createCudaSurface(cudaExtResObj.d_mipLevelArray[i]);\n }\n\n cudaExtResObj.stream = createCudaStream(deviceId);\n checkCudaErrors(hipMalloc(&cudaExtResObj.d_outputImage,\n sizeof(unsigned int) * cudaExtResObj.imageWidth[0] * cudaExtResObj.imageHeight[0]));\n}\n\n// CUDA clean up buffers used **with** NvSci API.\nvoid cleanupCuda(cudaExternalResInterop &cudaExtResObj)\n{\n for (int i = 0; i < cudaExtResObj.planeCount; i++) {\n checkCudaErrors(hipDestroySurfaceObject(cudaExtResObj.cudaSurfaceNvmediaBuf[i]));\n checkCudaErrors(hipFreeMipmappedArray(cudaExtResObj.d_mipmapArray[i]));\n }\n free(cudaExtResObj.d_mipmapArray);\n free(cudaExtResObj.d_mipLevelArray);\n free(cudaExtResObj.cudaSurfaceNvmediaBuf);\n free(cudaExtResObj.imageWidth);\n free(cudaExtResObj.imageHeight);\n checkCudaErrors(hipDestroyExternalSemaphore(cudaExtResObj.waitSem));\n checkCudaErrors(hipDestroyExternalSemaphore(cudaExtResObj.signalSem));\n checkCudaErrors(hipDestroyExternalMemory(cudaExtResObj.extMemImageBuf));\n checkCudaErrors(hipStreamDestroy(cudaExtResObj.stream));\n checkCudaErrors(hipFree(cudaExtResObj.d_outputImage));\n}\n\nvoid runCudaOperation(cudaExternalResInterop &cudaExtResObj,\n NvSciSyncFence *cudaWaitFence,\n NvSciSyncFence *cudaSignalFence,\n int deviceId,\n int iterations)\n{\n checkCudaErrors(hipSetDevice(deviceId));\n static int64_t launch = 0;\n\n waitExternalSemaphore(cudaExtResObj.waitSem, cudaWaitFence, cudaExtResObj.stream);\n\n // run cuda kernel over surface object of the LUMA surface part to extract\n // grayscale.\n yuvToGrayscaleCudaKernel(cudaExtResObj, cudaExtResObj.imageWidth[0], cudaExtResObj.imageHeight[0]);\n\n // signal fence till the second last iterations for NvMedia2DBlit to wait for\n // cuda signal and for final iteration as there is no corresponding NvMedia\n // operation pending therefore we end with cudaStreamSynchronize()\n if (launch < iterations - 1) {\n signalExternalSemaphore(cudaExtResObj.signalSem, cudaSignalFence, cudaExtResObj.stream);\n }\n else {\n checkCudaErrors(hipStreamSynchronize(cudaExtResObj.stream));\n }\n launch++;\n}\n\n// CUDA imports and operates on NvSci buffer/synchronization objects\nvoid setupCuda(Blit2DTest *ctx, cudaResources &cudaResObj, int deviceId)\n{\n checkCudaErrors(hipSetDevice(deviceId));\n cudaResObj.d_yuvArray = (hipArray_t *)malloc(sizeof(hipArray_t) * ctx->numSurfaces);\n cudaResObj.cudaSurfaceNvmediaBuf = (hipSurfaceObject_t *)malloc(sizeof(hipSurfaceObject_t) * ctx->numSurfaces);\n hipChannelFormatDesc channelDesc;\n switch (ctx->bytesPerPixel) {\n case 1:\n default:\n channelDesc = hipCreateChannelDesc(8, 0, 0, 0, hipChannelFormatKindUnsigned);\n break;\n }\n\n for (int k = 0; k < ctx->numSurfaces; k++) {\n checkCudaErrors(hipMallocArray(&cudaResObj.d_yuvArray[k],\n &channelDesc,\n ctx->widthSurface * ctx->xScalePtr[k] * ctx->bytesPerPixel,\n ctx->heightSurface * ctx->yScalePtr[k]));\n cudaResObj.cudaSurfaceNvmediaBuf[k] = createCudaSurface(cudaResObj.d_yuvArray[k]);\n }\n checkCudaErrors(\n hipMalloc(&cudaResObj.d_outputImage, sizeof(unsigned int) * ctx->widthSurface * ctx->heightSurface));\n\n cudaResObj.stream = createCudaStream(deviceId);\n}\n\n// CUDA clean up buffers used **without** NvSci API.\nvoid cleanupCuda(Blit2DTest *ctx, cudaResources &cudaResObj)\n{\n for (int k = 0; k < ctx->numSurfaces; k++) {\n checkCudaErrors(hipDestroySurfaceObject(cudaResObj.cudaSurfaceNvmediaBuf[k]));\n checkCudaErrors(hipFreeArray(cudaResObj.d_yuvArray[k]));\n }\n\n free(cudaResObj.cudaSurfaceNvmediaBuf);\n\n checkCudaErrors(hipStreamDestroy(cudaResObj.stream));\n checkCudaErrors(hipFree(cudaResObj.d_outputImage));\n}\n\nstatic void\nyuvToGrayscaleCudaKernelNonNvSci(cudaResources &cudaResObj, int deviceId, int32_t imageWidth, int32_t imageHeight)\n{\n#if WRITE_OUTPUT_IMAGE\n unsigned int *h_dstImage;\n checkCudaErrors(hipHostMalloc(&h_dstImage, sizeof(unsigned int) * imageHeight * imageWidth));\n#endif\n dim3 block(16, 16, 1);\n dim3 grid((imageWidth / block.x) + 1, (imageHeight / block.y) + 1, 1);\n\n yuvToGrayscale<<>>(\n cudaResObj.cudaSurfaceNvmediaBuf[0], cudaResObj.d_outputImage, imageWidth, imageHeight);\n\n#if WRITE_OUTPUT_IMAGE\n checkCudaErrors(hipMemcpyAsync(h_dstImage,\n cudaResObj.d_outputImage,\n sizeof(unsigned int) * imageHeight * imageWidth,\n hipMemcpyDeviceToHost,\n cudaResObj.stream));\n checkCudaErrors(hipStreamSynchronize(cudaResObj.stream));\n char outputFilename[1024];\n std::string image_filename = \"Grayscale\";\n strcpy(outputFilename, image_filename.c_str());\n strcpy(outputFilename + image_filename.length(), \"_non-nvsci_out.ppm\");\n sdkSavePPM4ub(outputFilename, (unsigned char *)h_dstImage, imageWidth, imageHeight);\n printf(\"Wrote '%s'\\n\", outputFilename);\n checkCudaErrors(hipHostFree(h_dstImage));\n#else\n checkCudaErrors(hipStreamSynchronize(cudaResObj.stream));\n#endif\n}\n\n// CUDA operates **without** NvSci APIs buffer/synchronization objects.\nvoid runCudaOperation(Blit2DTest *ctx, cudaResources &cudaResObj, int deviceId)\n{\n for (int k = 0; k < ctx->numSurfaces; k++) {\n checkCudaErrors(hipMemcpy2DToArray(cudaResObj.d_yuvArray[k],\n 0,\n 0,\n ctx->dstBuff[k],\n ctx->widthSurface * ctx->xScalePtr[k] * ctx->bytesPerPixel,\n ctx->widthSurface * ctx->xScalePtr[k] * ctx->bytesPerPixel,\n ctx->heightSurface * ctx->yScalePtr[k],\n hipMemcpyHostToDevice));\n }\n // run cuda kernel over surface object of the LUMA surface part to extract\n // grayscale.\n yuvToGrayscaleCudaKernelNonNvSci(cudaResObj, deviceId, ctx->widthSurface, ctx->heightSurface);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/8_Platform_Specific/Tegra/cudaNvSciNvMedia/cuda_consumer.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// std::system includes\n#include \n\n// CUDA-C includes\n#include \n#include \n\n#define TOTAL_SIZE 256 * 1024 * 1024\n#define EACH_SIZE 128 * 1024 * 1024\n\n// # threadblocks\n#define TBLOCKS 1024\n#define THREADS 512\n\n// throw error on equality\n#define ERR_EQ(X, Y) \\\n do { \\\n if ((X) == (Y)) { \\\n fprintf(stderr, \"Error in %s at %s:%d\\n\", __func__, __FILE__, __LINE__); \\\n exit(-1); \\\n } \\\n } while (0)\n\n// throw error on difference\n#define ERR_NE(X, Y) \\\n do { \\\n if ((X) != (Y)) { \\\n fprintf(stderr, \"Error in %s at %s:%d\\n\", __func__, __FILE__, __LINE__); \\\n exit(-1); \\\n } \\\n } while (0)\n\n// copy from source -> destination arrays\n__global__ void memcpy_kernel(int *dst, int *src, size_t n)\n{\n int num = gridDim.x * blockDim.x;\n int id = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int i = id; i < n / sizeof(int); i += num) {\n dst[i] = src[i];\n }\n}\n\n// initialise memory\nvoid mem_init(int *buf, size_t n)\n{\n for (int i = 0; i < n / sizeof(int); i++) {\n buf[i] = i;\n }\n}\n\nint main(int argc, char **argv)\n{\n cudaDeviceProp device_prop;\n int dev_id;\n\n printf(\"Starting [%s]...\\n\", argv[0]);\n\n // set device\n dev_id = findCudaDevice(argc, (const char **)argv);\n checkCudaErrors(cudaGetDeviceProperties(&device_prop, dev_id));\n\n if ((device_prop.major << 4) + device_prop.minor < 0x35) {\n fprintf(stderr,\n \"%s requires Compute Capability of SM 3.5 or higher to \"\n \"run.\\nexiting...\\n\",\n argv[0]);\n exit(EXIT_WAIVED);\n }\n\n // get the range of priorities available\n // [ greatest_priority, lowest_priority ]\n int priority_low;\n int priority_hi;\n checkCudaErrors(cudaDeviceGetStreamPriorityRange(&priority_low, &priority_hi));\n\n printf(\"CUDA stream priority range: LOW: %d to HIGH: %d\\n\", priority_low, priority_hi);\n\n // create streams with highest and lowest available priorities\n cudaStream_t st_low;\n cudaStream_t st_hi;\n checkCudaErrors(cudaStreamCreateWithPriority(&st_low, cudaStreamNonBlocking, priority_low));\n checkCudaErrors(cudaStreamCreateWithPriority(&st_hi, cudaStreamNonBlocking, priority_hi));\n\n size_t size;\n size = TOTAL_SIZE;\n\n // initialise host data\n int *h_src_low;\n int *h_src_hi;\n ERR_EQ(h_src_low = (int *)malloc(size), NULL);\n ERR_EQ(h_src_hi = (int *)malloc(size), NULL);\n mem_init(h_src_low, size);\n mem_init(h_src_hi, size);\n\n // initialise device data\n int *h_dst_low;\n int *h_dst_hi;\n ERR_EQ(h_dst_low = (int *)malloc(size), NULL);\n ERR_EQ(h_dst_hi = (int *)malloc(size), NULL);\n memset(h_dst_low, 0, size);\n memset(h_dst_hi, 0, size);\n\n // copy source data -> device\n int *d_src_low;\n int *d_src_hi;\n checkCudaErrors(cudaMalloc(&d_src_low, size));\n checkCudaErrors(cudaMalloc(&d_src_hi, size));\n checkCudaErrors(cudaMemcpy(d_src_low, h_src_low, size, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(d_src_hi, h_src_hi, size, cudaMemcpyHostToDevice));\n\n // allocate memory for memcopy destination\n int *d_dst_low;\n int *d_dst_hi;\n checkCudaErrors(cudaMalloc(&d_dst_low, size));\n checkCudaErrors(cudaMalloc(&d_dst_hi, size));\n\n // create some events\n cudaEvent_t ev_start_low;\n cudaEvent_t ev_start_hi;\n cudaEvent_t ev_end_low;\n cudaEvent_t ev_end_hi;\n checkCudaErrors(cudaEventCreate(&ev_start_low));\n checkCudaErrors(cudaEventCreate(&ev_start_hi));\n checkCudaErrors(cudaEventCreate(&ev_end_low));\n checkCudaErrors(cudaEventCreate(&ev_end_hi));\n\n /* */\n\n // call pair of kernels repeatedly (with different priority streams)\n checkCudaErrors(cudaEventRecord(ev_start_low, st_low));\n checkCudaErrors(cudaEventRecord(ev_start_hi, st_hi));\n\n for (int i = 0; i < TOTAL_SIZE; i += EACH_SIZE) {\n int j = i / sizeof(int);\n memcpy_kernel<<>>(d_dst_low + j, d_src_low + j, EACH_SIZE);\n memcpy_kernel<<>>(d_dst_hi + j, d_src_hi + j, EACH_SIZE);\n }\n\n checkCudaErrors(cudaEventRecord(ev_end_low, st_low));\n checkCudaErrors(cudaEventRecord(ev_end_hi, st_hi));\n\n checkCudaErrors(cudaEventSynchronize(ev_end_low));\n checkCudaErrors(cudaEventSynchronize(ev_end_hi));\n\n /* */\n\n size = TOTAL_SIZE;\n checkCudaErrors(cudaMemcpy(h_dst_low, d_dst_low, size, cudaMemcpyDeviceToHost));\n checkCudaErrors(cudaMemcpy(h_dst_hi, d_dst_hi, size, cudaMemcpyDeviceToHost));\n\n // check results of kernels\n ERR_NE(memcmp(h_dst_low, h_src_low, size), 0);\n ERR_NE(memcmp(h_dst_hi, h_src_hi, size), 0);\n\n // check timings\n float ms_low;\n float ms_hi;\n checkCudaErrors(cudaEventElapsedTime(&ms_low, ev_start_low, ev_end_low));\n checkCudaErrors(cudaEventElapsedTime(&ms_hi, ev_start_hi, ev_end_hi));\n\n printf(\"elapsed time of kernels launched to LOW priority stream: %.3lf ms\\n\", ms_low);\n printf(\"elapsed time of kernels launched to HI priority stream: %.3lf ms\\n\", ms_hi);\n\n exit(EXIT_SUCCESS);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// std::system includes\n#include \n\n// CUDA-C includes\n#include \n#include \n\n#define TOTAL_SIZE 256 * 1024 * 1024\n#define EACH_SIZE 128 * 1024 * 1024\n\n// # threadblocks\n#define TBLOCKS 1024\n#define THREADS 512\n\n// throw error on equality\n#define ERR_EQ(X, Y) \\\n do { \\\n if ((X) == (Y)) { \\\n fprintf(stderr, \"Error in %s at %s:%d\\n\", __func__, __FILE__, __LINE__); \\\n exit(-1); \\\n } \\\n } while (0)\n\n// throw error on difference\n#define ERR_NE(X, Y) \\\n do { \\\n if ((X) != (Y)) { \\\n fprintf(stderr, \"Error in %s at %s:%d\\n\", __func__, __FILE__, __LINE__); \\\n exit(-1); \\\n } \\\n } while (0)\n\n// copy from source -> destination arrays\n__global__ void memcpy_kernel(int *dst, int *src, size_t n)\n{\n int num = gridDim.x * blockDim.x;\n int id = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int i = id; i < n / sizeof(int); i += num) {\n dst[i] = src[i];\n }\n}\n\n// initialise memory\nvoid mem_init(int *buf, size_t n)\n{\n for (int i = 0; i < n / sizeof(int); i++) {\n buf[i] = i;\n }\n}\n\nint main(int argc, char **argv)\n{\n hipDeviceProp_t device_prop;\n int dev_id;\n\n printf(\"Starting [%s]...\\n\", argv[0]);\n\n // set device\n dev_id = findCudaDevice(argc, (const char **)argv);\n checkCudaErrors(hipGetDeviceProperties(&device_prop, dev_id));\n\n if ((device_prop.major << 4) + device_prop.minor < 0x35) {\n fprintf(stderr,\n \"%s requires Compute Capability of SM 3.5 or higher to \"\n \"run.\\nexiting...\\n\",\n argv[0]);\n exit(EXIT_WAIVED);\n }\n\n // get the range of priorities available\n // [ greatest_priority, lowest_priority ]\n int priority_low;\n int priority_hi;\n checkCudaErrors(hipDeviceGetStreamPriorityRange(&priority_low, &priority_hi));\n\n printf(\"CUDA stream priority range: LOW: %d to HIGH: %d\\n\", priority_low, priority_hi);\n\n // create streams with highest and lowest available priorities\n hipStream_t st_low;\n hipStream_t st_hi;\n checkCudaErrors(hipStreamCreateWithPriority(&st_low, hipStreamNonBlocking, priority_low));\n checkCudaErrors(hipStreamCreateWithPriority(&st_hi, hipStreamNonBlocking, priority_hi));\n\n size_t size;\n size = TOTAL_SIZE;\n\n // initialise host data\n int *h_src_low;\n int *h_src_hi;\n ERR_EQ(h_src_low = (int *)malloc(size), NULL);\n ERR_EQ(h_src_hi = (int *)malloc(size), NULL);\n mem_init(h_src_low, size);\n mem_init(h_src_hi, size);\n\n // initialise device data\n int *h_dst_low;\n int *h_dst_hi;\n ERR_EQ(h_dst_low = (int *)malloc(size), NULL);\n ERR_EQ(h_dst_hi = (int *)malloc(size), NULL);\n memset(h_dst_low, 0, size);\n memset(h_dst_hi, 0, size);\n\n // copy source data -> device\n int *d_src_low;\n int *d_src_hi;\n checkCudaErrors(hipMalloc(&d_src_low, size));\n checkCudaErrors(hipMalloc(&d_src_hi, size));\n checkCudaErrors(hipMemcpy(d_src_low, h_src_low, size, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(d_src_hi, h_src_hi, size, hipMemcpyHostToDevice));\n\n // allocate memory for memcopy destination\n int *d_dst_low;\n int *d_dst_hi;\n checkCudaErrors(hipMalloc(&d_dst_low, size));\n checkCudaErrors(hipMalloc(&d_dst_hi, size));\n\n // create some events\n hipEvent_t ev_start_low;\n hipEvent_t ev_start_hi;\n hipEvent_t ev_end_low;\n hipEvent_t ev_end_hi;\n checkCudaErrors(hipEventCreate(&ev_start_low));\n checkCudaErrors(hipEventCreate(&ev_start_hi));\n checkCudaErrors(hipEventCreate(&ev_end_low));\n checkCudaErrors(hipEventCreate(&ev_end_hi));\n\n /* */\n\n // call pair of kernels repeatedly (with different priority streams)\n checkCudaErrors(hipEventRecord(ev_start_low, st_low));\n checkCudaErrors(hipEventRecord(ev_start_hi, st_hi));\n\n for (int i = 0; i < TOTAL_SIZE; i += EACH_SIZE) {\n int j = i / sizeof(int);\n memcpy_kernel<<>>(d_dst_low + j, d_src_low + j, EACH_SIZE);\n memcpy_kernel<<>>(d_dst_hi + j, d_src_hi + j, EACH_SIZE);\n }\n\n checkCudaErrors(hipEventRecord(ev_end_low, st_low));\n checkCudaErrors(hipEventRecord(ev_end_hi, st_hi));\n\n checkCudaErrors(hipEventSynchronize(ev_end_low));\n checkCudaErrors(hipEventSynchronize(ev_end_hi));\n\n /* */\n\n size = TOTAL_SIZE;\n checkCudaErrors(hipMemcpy(h_dst_low, d_dst_low, size, hipMemcpyDeviceToHost));\n checkCudaErrors(hipMemcpy(h_dst_hi, d_dst_hi, size, hipMemcpyDeviceToHost));\n\n // check results of kernels\n ERR_NE(memcmp(h_dst_low, h_src_low, size), 0);\n ERR_NE(memcmp(h_dst_hi, h_src_hi, size), 0);\n\n // check timings\n float ms_low;\n float ms_hi;\n checkCudaErrors(hipEventElapsedTime(&ms_low, ev_start_low, ev_end_low));\n checkCudaErrors(hipEventElapsedTime(&ms_hi, ev_start_hi, ev_end_hi));\n\n printf(\"elapsed time of kernels launched to LOW priority stream: %.3lf ms\\n\", ms_low);\n printf(\"elapsed time of kernels launched to HI priority stream: %.3lf ms\\n\", ms_hi);\n\n exit(EXIT_SUCCESS);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/StreamPriorities/StreamPriorities.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// CUDA sample demonstrating a integer GEMM computation using the Warp Matrix\n// Multiply and Accumulate API.\n\n// In this program, the compute_gemm kernel computes the result of a matrix\n// multiplication and addition: D = alpha * A * B + beta * C. The dimensions of\n// both C and D matrices are M_GLOBAL x N_GLOBAL. The A matrix is M_GLOBAL x\n// K_GLOBAL (row-major), the B matrix is K_GLOBAL x N_GLOBAL (column-major). In\n// that kernel, each CTA computes one 128 x 128 tile of the resulting matrix per\n// iteration. When the tile is computed, the CTA stores it to the global memory\n// and begins a new iteration, selecting a new 128 x 128 tile to compute.\n// Each CTA consists of eight warps. For the 128 x 128 tile, each warp computes\n// eight 16 x 16 subtiles, organized in a 2 x 4 two-dimensional array. Warps\n// compute the 16 x 16 subtiles using nvcuda::wmma::mma_sync operations by\n// moving through the K_GLOBAL dimension of the A and B matrices and\n// accumulating the intermediate result in the local thread state.\n\n// There are a number of simple optimizations used in the algorithm:\n// - The CTA copies the 128 x 128 tile of the C matrix from the global memory to\n// shared memory. After that is done, each warp loads the C matrix fragments\n// from shared memory, thus avoiding a random global memory access.\n// - On each internal iteration, the CTA copies a portion of the A and B\n// matrices from\n// global memory to shared memory. After that, all warps in the CTA reuse the\n// A and B data from shared memory, thus reducing the number of data copies\n// from global memory.\n// - The portions of the A and B matrices are stored in shared memory with an\n// additional\n// padding (skew) to reduce the number of shared memory access bank conflicts.\n// (See a detailed explanation near the SKEW_HALF macro definition.)\n// - When the CTA finishes computing the tiles of the resulting matrix, each\n// warp stores\n// its subtiles to shared memory. The CTA then copies the shared memory\n// contents to global memory, again avoiding redundant random global memory\n// accesses.\n// - Note that the CTA tile size is chosen to maximize the GPU register\n// utilization,\n// but carefully enough to avoid local memory use.\n\n#include \n#include \n#include \n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Externally configurable parameters.\n\n#ifndef CPU_DEBUG\n// Set this to 1 to verify the correctness of the GPU-computed matrix.\n#define CPU_DEBUG 0\n#endif\n\n#ifndef SHARED_MEMORY_LIMIT_64K\n// Set this to 0 to use more than 64 Kb of shared memory to cache data, to\n// improve the performance of the computations on GPU.\n// Note that you need a GPU that can have more than 64 Kb of shared memory\n// per multiprocessor.\n#define SHARED_MEMORY_LIMIT_64K 1\n#endif\n\n// GPU configuration.\n\n#define WARP_SIZE 32\n\n// MMA matrix tile dimensions.\n\n#define M 16\n#define N 16\n#define K 16\n\n#define WMMA_M 16\n#define WMMA_N 16\n#define WMMA_K 16\n\n// GEMM configuration.\n\n#define M_TILES 256\n#define N_TILES 256\n#define K_TILES 256\n\n#define M_GLOBAL (M * M_TILES)\n#define N_GLOBAL (N * N_TILES)\n#define K_GLOBAL (K * K_TILES)\n\n#define C_LAYOUT wmma::mem_row_major\n\n// Implementation constants.\n\n#define WARPS_PER_BLOCK 8\n#define THREADS_PER_BLOCK (WARP_SIZE * WARPS_PER_BLOCK)\n\n#if SHARED_MEMORY_LIMIT_64K\n// With only 64 Kb shared memory available, we can fit two 8-tile chunks of\n// the A and B matrix data, that are 16 * 16 * 8 * 8 * 2 = 32 Kb each\n// (i.e. two 8x8 arrays of tiles of 16x16 uint8_t-typed elements per CTA).\n// But we cannot account the 8 Kb total skew overhead, without which the\n// performance would be severely impacted. So we choose to reduce the chunk size\n// in half, i.e. the amount of A and B matrix data we cache in shared memory.\n// Accordingly, this doubles the number of outer iterations across the global K\n// dimension, which only slightly impacts the performance.\n#define CHUNK_K 8\n#else\n#define CHUNK_K 16\n#endif\n\n#define CHUNK_LINE_BYTES (CHUNK_K * K * sizeof(uint8_t))\n#define WARP_COPY_BYTES (WARP_SIZE * sizeof(int4))\n#define CHUNK_COPY_LINES_PER_WARP (WARP_COPY_BYTES / CHUNK_LINE_BYTES)\n#define CHUNK_COPY_LINE_LANES (WARP_SIZE / CHUNK_COPY_LINES_PER_WARP)\n\n#define BLOCK_ROW_WARPS 2\n#define BLOCK_COL_WARPS 4\n\n#define WARP_ROW_TILES 4\n#define WARP_COL_TILES 2\n\n#define BLOCK_ROW_TILES (WARP_ROW_TILES * BLOCK_ROW_WARPS)\n#define BLOCK_COL_TILES (WARP_COL_TILES * BLOCK_COL_WARPS)\n\n#define GLOBAL_MEM_STRIDE N_GLOBAL\n\n#define SHMEM_STRIDE (N * BLOCK_ROW_TILES)\n#define SHMEM_OFFSET (N * WARP_ROW_TILES)\n\n// The macro below is used to shift rows of the A matrix and columns of the B\n// matrix in shared memory to minimize possible bank conflicts. Before\n// performing the nvcuda::wmma::mma_sync operation, the warp must load the\n// matrix data using the nvcuda::wmma::load_matrix_sync operation. Although the\n// memory access pattern is not specified for that function, each lane in the\n// warp can read one or multiple matrix elements from different matrix rows or\n// columns. For shared memory, such access can result in bank conflicts if\n// different rows / columns of the matrix map to the same bank. By shifting each\n// row and column by a few bytes, we make sure that they map to different banks,\n// thus reducing the number of possible bank conflicts. The number of 32\n// one-byte \"uint8_t\" elements is chosen as the minimum possible shift because\n// we must keep each row and column 256-bit aligned, as required by\n// nvcuda::wmma::load_matrix_sync.\n#define SKEW_UINT8 32\n\n#define checkKernelErrors(expr) \\\n do { \\\n expr; \\\n \\\n cudaError_t __err = cudaGetLastError(); \\\n if (__err != cudaSuccess) { \\\n printf(\"Line %d: '%s' failed: %s\\n\", __LINE__, #expr, cudaGetErrorString(__err)); \\\n abort(); \\\n } \\\n } while (0)\n\nusing namespace nvcuda;\n\n__host__ void init_host_matrices(uint8_t *a, uint8_t *b, int *c)\n{\n for (int i = 0; i < M_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n a[i * K_GLOBAL + j] = (uint8_t)(rand() % 3);\n }\n }\n\n for (int i = 0; i < N_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n b[i * K_GLOBAL + j] = (uint8_t)(rand() % 3);\n }\n }\n\n for (int t = 0; t < M_GLOBAL * N_GLOBAL; t++) {\n c[t] = (rand() % 3);\n }\n}\n\n__global__ void compute_gemm_imma(const uint8_t *A, const uint8_t *B, const int *C, int *D, int alpha, int beta)\n{\n extern __shared__ uint8_t shmem[][CHUNK_K * K + SKEW_UINT8];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n const size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n int *shmem_warp_tile_ptr = (int *)&shmem[0][0] + (warpId / 2) * SHMEM_STRIDE * K * 2 + (warpId % 2) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and\n // from shared memory.\n int *shmem_warp_stream_ptr = (int *)&shmem[0][0] + warpId * SHMEM_STRIDE * K;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may\n // result in a loss of precision). Zero still needs to be specially handled\n // though.\n beta /= alpha;\n\n // Each CTA slides along the 128 x 128 tiles from the top left corner of the\n // matrix to the right and down, and selects the next tile to compute. Once\n // there's no such tile, all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared\n // memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const int *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < K; i++) {\n typedef int4 copy_t;\n\n *((copy_t *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId) =\n *((copy_t *)(src_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment\n // multiplications along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const int *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * K + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Scale the C matrix.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const uint8_t *warp_ptr = (warpId < 4) ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % 4) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % 4) * 2);\n\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy\n // the B matrix.\n size_t shmem_idx = warpId < (WARPS_PER_BLOCK / 2)\n ? (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (N * (warpId % (WARPS_PER_BLOCK / 2)) * 2 + shmem_idx_b_off);\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n int4 *lane_ptr = (int4 *)(warp_ptr + tile_k * K + (laneId / CHUNK_COPY_LINE_LANES) * K_GLOBAL)\n + (laneId % CHUNK_COPY_LINE_LANES);\n\n // Shift the second half of the warp to the next row / column in the\n // shared memory.\n shmem_idx += laneId / CHUNK_COPY_LINE_LANES;\n\n#pragma unroll\n for (int i = 0; i < ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP) * 2; i++) {\n // Copy 16 bytes at once in each lane.\n *((int4 *)&shmem[shmem_idx][0] + (laneId % CHUNK_COPY_LINE_LANES)) = *lane_ptr;\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = (int4 *)((uint8_t *)lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP);\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / 2) * M * 2 + (i * M);\n const uint8_t *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_UINT8);\n\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be\n // reused against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const uint8_t *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_UINT8);\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n // Uniform, point-wise transformations of ALL fragment elements by ALL\n // threads in the warp are well-defined even though element indices\n // within fragment storage are not defined.\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n int *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * K + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global\n // memory.\n int *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < K; i++) {\n *((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n}\n\n// Performs an MxNxK GEMM (C=alpha*A*B + beta*C) assuming:\n// 1) Matrices are packed in memory.\n// 2) M, N and K are multiples of 16.\n// 3) Neither A nor B are transposed.\n// Note: This is a less performant version of the compute_gemm_imma kernel. It\n// is designed for\n// demonstration purposes only to show the CUDA WMMA API use without\n// relying on availability of the shared memory.\n__global__ void simple_wmma_gemm_imma(const uint8_t *a,\n const uint8_t *b,\n const int *c,\n int *d,\n int m_ld,\n int n_ld,\n int k_ld,\n int alpha,\n int beta)\n{\n // Leading dimensions. Packed with no transpositions.\n int lda = m_ld;\n int ldb = k_ld;\n int ldc = n_ld;\n\n // Tile using a 2D grid\n int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;\n int warpN = (blockIdx.y * blockDim.y + threadIdx.y);\n\n // Declare the fragments\n wmma::fragment a_frag;\n wmma::fragment b_frag;\n wmma::fragment acc_frag;\n wmma::fragment c_frag;\n\n wmma::fill_fragment(acc_frag, 0.0f);\n\n // Loop over k\n for (int i = 0; i < k_ld; i += WMMA_K) {\n int aCol = i;\n int aRow = warpM * WMMA_M;\n\n int bCol = i;\n int bRow = warpN * WMMA_N;\n\n // Bounds checking\n if (aRow < m_ld && aCol < k_ld && bRow < k_ld && bCol < n_ld) {\n // Load the inputs\n wmma::load_matrix_sync(a_frag, a + aCol + aRow * lda, lda);\n wmma::load_matrix_sync(b_frag, b + bCol + bRow * ldb, ldb);\n\n // Perform the matrix multiplication\n wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);\n }\n }\n\n // Load in the current value of c, scale it by beta, and add this our result\n // scaled by alpha\n int cCol = warpN * WMMA_N;\n int cRow = warpM * WMMA_M;\n\n if (cRow < m_ld && cCol < n_ld) {\n wmma::load_matrix_sync(c_frag, c + cCol + cRow * ldc, ldc, wmma::mem_row_major);\n\n for (int i = 0; i < c_frag.num_elements; i++) {\n c_frag.x[i] = alpha * acc_frag.x[i] + beta * c_frag.x[i];\n }\n\n // Store the output\n wmma::store_matrix_sync(d + cCol + cRow * ldc, c_frag, ldc, wmma::mem_row_major);\n }\n}\n\n__host__ void matMultiplyOnHost(uint8_t *A,\n uint8_t *B,\n int *C,\n int alpha,\n int beta,\n int numARows,\n int numAColumns,\n int numBRows,\n int numBColumns,\n int numCRows,\n int numCColumns)\n{\n for (int i = 0; i < numCRows; i++) {\n for (int j = 0; j < numCColumns; j++) {\n int temp = 0;\n\n for (int k = 0; k < numAColumns; k++) {\n temp += A[i * numAColumns + k] * B[j * numBRows + k];\n }\n\n C[i * numCColumns + j] = temp * alpha + beta * C[i * numCColumns + j];\n }\n }\n}\n\nint main(int argc, char **argv)\n{\n printf(\"Initializing...\\n\");\n\n int dev = findCudaDevice(argc, (const char **)argv);\n\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, dev));\n\n // Tensor cores require a GPU of Volta (SM72) architecture or higher.\n if (deviceProp.major < 7 || (deviceProp.major <= 7 && deviceProp.minor < 2)) {\n printf(\"immaTensorCoreGemm requires SM 7.2 or higher to use Tensor Cores. \"\n \"Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n printf(\"M: %d (%d x %d)\\n\", M_GLOBAL, M, M_TILES);\n printf(\"N: %d (%d x %d)\\n\", N_GLOBAL, N, N_TILES);\n printf(\"K: %d (%d x %d)\\n\", K_GLOBAL, K, K_TILES);\n\n uint8_t *A_h = NULL;\n uint8_t *B_h = NULL;\n int *C_h = NULL;\n#if CPU_DEBUG\n int *result_hD = NULL;\n int *result_host = NULL;\n#endif\n\n A_h = (uint8_t *)malloc(sizeof(uint8_t) * M_GLOBAL * K_GLOBAL);\n B_h = (uint8_t *)malloc(sizeof(uint8_t) * K_GLOBAL * N_GLOBAL);\n C_h = (int *)malloc(sizeof(int) * M_GLOBAL * N_GLOBAL);\n#if CPU_DEBUG\n result_hD = (int *)malloc(sizeof(int) * M_GLOBAL * N_GLOBAL);\n result_host = (int *)malloc(sizeof(int) * M_GLOBAL * N_GLOBAL);\n#endif\n\n uint8_t *A = NULL;\n uint8_t *B = NULL;\n int *C = NULL;\n int *D = NULL;\n\n checkCudaErrors(cudaMalloc(reinterpret_cast(&A), sizeof(uint8_t) * M_GLOBAL * K_GLOBAL));\n checkCudaErrors(cudaMalloc(reinterpret_cast(&B), sizeof(uint8_t) * N_GLOBAL * K_GLOBAL));\n checkCudaErrors(cudaMalloc(reinterpret_cast(&C), sizeof(int) * M_GLOBAL * N_GLOBAL));\n checkCudaErrors(cudaMalloc(reinterpret_cast(&D), sizeof(int) * M_GLOBAL * N_GLOBAL));\n\n assert(((unsigned long long)A) % 128 == 0);\n assert(((unsigned long long)B) % 128 == 0);\n assert(((unsigned long long)C) % 128 == 0);\n assert(((unsigned long long)D) % 128 == 0);\n\n init_host_matrices(A_h, B_h, C_h);\n\n checkCudaErrors(cudaMemcpy(A, A_h, sizeof(uint8_t) * M_GLOBAL * K_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(B, B_h, sizeof(uint8_t) * N_GLOBAL * K_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(C, C_h, sizeof(int) * M_GLOBAL * N_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemset(D, 0, sizeof(int) * M_GLOBAL * N_GLOBAL));\n\n printf(\"Preparing data for GPU...\\n\");\n\n assert(((unsigned long long)A) % 128 == 0);\n assert(((unsigned long long)B) % 128 == 0);\n assert(((unsigned long long)C) % 128 == 0);\n assert(((unsigned long long)D) % 128 == 0);\n\n enum {\n // Compute the right amount of shared memory to request.\n // We need shared memory to hold per-CTA C and D matrix tiles, and to cache\n // per-CTA chunks\n // of the A and B matrices. Therefore, the right amount to request is the\n // maximum of those\n // two numbers.\n SHMEM_SZ = MAX(sizeof(uint8_t) * (BLOCK_COL_TILES * M) * (CHUNK_K * K + SKEW_UINT8) * 2,\n M * (BLOCK_ROW_WARPS * WARP_ROW_TILES) * N * (BLOCK_COL_WARPS * WARP_COL_TILES) * sizeof(int))\n };\n\n printf(\"Required shared memory size: %lu Kb\\n\", SHMEM_SZ / 1024UL);\n\n int alpha = 1;\n int beta = 1;\n\n cudaEvent_t start, stop;\n\n checkCudaErrors(cudaEventCreate(&start));\n checkCudaErrors(cudaEventCreate(&stop));\n checkCudaErrors(cudaEventRecord(start));\n\n // If enough shared memory available on the GPU use high performant kernel\n if (deviceProp.sharedMemPerMultiprocessor >= SHMEM_SZ) {\n printf(\"Computing... using high performance kernel compute_gemm_imma \\n\");\n\n checkCudaErrors(cudaFuncSetAttribute(compute_gemm_imma, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors((compute_gemm_imma<<>>(\n A, B, C, D, alpha, beta)));\n#if CPU_DEBUG\n checkCudaErrors(cudaMemcpy(result_hD, D, sizeof(int) * M_GLOBAL * N_GLOBAL, cudaMemcpyDeviceToHost));\n#endif\n }\n else {\n dim3 gridDim;\n dim3 blockDim;\n\n // blockDim.x must be a multiple of warpSize\n // 128x4 means we have 16 warps and a block computes a 64x64 output tile\n blockDim.x = 128;\n blockDim.y = 4;\n\n gridDim.x = (M_GLOBAL + (WMMA_M * blockDim.x / 32 - 1)) / (WMMA_M * blockDim.x / 32);\n gridDim.y = (N_GLOBAL + WMMA_N * blockDim.y - 1) / (WMMA_N * blockDim.y);\n\n printf(\"Computing... using simple_wmma_gemm_imma kernel\\n\");\n simple_wmma_gemm_imma<<>>(A, B, C, D, M_GLOBAL, N_GLOBAL, K_GLOBAL, alpha, beta);\n#if CPU_DEBUG\n checkCudaErrors(cudaMemcpy(result_hD, D, sizeof(int) * M_GLOBAL * N_GLOBAL, cudaMemcpyDeviceToHost));\n#endif\n }\n\n checkCudaErrors(cudaEventRecord(stop));\n checkCudaErrors(cudaEventSynchronize(stop));\n\n#if CPU_DEBUG\n printf(\"Verifying correctness of the computations...\\n\");\n\n memcpy(result_host, C_h, sizeof(int) * M_GLOBAL * N_GLOBAL);\n\n matMultiplyOnHost(A_h, B_h, result_host, alpha, beta, M_GLOBAL, K_GLOBAL, K_GLOBAL, N_GLOBAL, M_GLOBAL, N_GLOBAL);\n\n for (int i = 0; i < N_GLOBAL * M_GLOBAL; i++) {\n if (abs(result_hD[i] - result_host[i]) > 0) {\n printf(\"mismatch i=%d result_hD=%d result_host=%d\\n\", i, result_hD[i], result_host[i]);\n }\n }\n free(result_host);\n free(result_hD);\n#endif\n\n float milliseconds = 0;\n\n checkCudaErrors(cudaEventElapsedTime(&milliseconds, start, stop));\n\n printf(\"Time: %f ms\\n\", milliseconds);\n printf(\"TOPS: %.2f\\n\", (((double)M_GLOBAL * N_GLOBAL * K_GLOBAL * 2) / (milliseconds / 1000.)) / 1e12);\n\n free(A_h);\n free(B_h);\n free(C_h);\n checkCudaErrors(cudaFree(reinterpret_cast(A)));\n checkCudaErrors(cudaFree(reinterpret_cast(B)));\n checkCudaErrors(cudaFree(reinterpret_cast(C)));\n checkCudaErrors(cudaFree(reinterpret_cast(D)));\n\n return EXIT_SUCCESS;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// CUDA sample demonstrating a integer GEMM computation using the Warp Matrix\n// Multiply and Accumulate API.\n\n// In this program, the compute_gemm kernel computes the result of a matrix\n// multiplication and addition: D = alpha * A * B + beta * C. The dimensions of\n// both C and D matrices are M_GLOBAL x N_GLOBAL. The A matrix is M_GLOBAL x\n// K_GLOBAL (row-major), the B matrix is K_GLOBAL x N_GLOBAL (column-major). In\n// that kernel, each CTA computes one 128 x 128 tile of the resulting matrix per\n// iteration. When the tile is computed, the CTA stores it to the global memory\n// and begins a new iteration, selecting a new 128 x 128 tile to compute.\n// Each CTA consists of eight warps. For the 128 x 128 tile, each warp computes\n// eight 16 x 16 subtiles, organized in a 2 x 4 two-dimensional array. Warps\n// compute the 16 x 16 subtiles using nvcuda::wmma::mma_sync operations by\n// moving through the K_GLOBAL dimension of the A and B matrices and\n// accumulating the intermediate result in the local thread state.\n\n// There are a number of simple optimizations used in the algorithm:\n// - The CTA copies the 128 x 128 tile of the C matrix from the global memory to\n// shared memory. After that is done, each warp loads the C matrix fragments\n// from shared memory, thus avoiding a random global memory access.\n// - On each internal iteration, the CTA copies a portion of the A and B\n// matrices from\n// global memory to shared memory. After that, all warps in the CTA reuse the\n// A and B data from shared memory, thus reducing the number of data copies\n// from global memory.\n// - The portions of the A and B matrices are stored in shared memory with an\n// additional\n// padding (skew) to reduce the number of shared memory access bank conflicts.\n// (See a detailed explanation near the SKEW_HALF macro definition.)\n// - When the CTA finishes computing the tiles of the resulting matrix, each\n// warp stores\n// its subtiles to shared memory. The CTA then copies the shared memory\n// contents to global memory, again avoiding redundant random global memory\n// accesses.\n// - Note that the CTA tile size is chosen to maximize the GPU register\n// utilization,\n// but carefully enough to avoid local memory use.\n\n#include \n#include \n#include \n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Externally configurable parameters.\n\n#ifndef CPU_DEBUG\n// Set this to 1 to verify the correctness of the GPU-computed matrix.\n#define CPU_DEBUG 0\n#endif\n\n#ifndef SHARED_MEMORY_LIMIT_64K\n// Set this to 0 to use more than 64 Kb of shared memory to cache data, to\n// improve the performance of the computations on GPU.\n// Note that you need a GPU that can have more than 64 Kb of shared memory\n// per multiprocessor.\n#define SHARED_MEMORY_LIMIT_64K 1\n#endif\n\n// GPU configuration.\n\n#define WARP_SIZE 32\n\n// MMA matrix tile dimensions.\n\n#define M 16\n#define N 16\n#define K 16\n\n#define WMMA_M 16\n#define WMMA_N 16\n#define WMMA_K 16\n\n// GEMM configuration.\n\n#define M_TILES 256\n#define N_TILES 256\n#define K_TILES 256\n\n#define M_GLOBAL (M * M_TILES)\n#define N_GLOBAL (N * N_TILES)\n#define K_GLOBAL (K * K_TILES)\n\n#define C_LAYOUT wmma::mem_row_major\n\n// Implementation constants.\n\n#define WARPS_PER_BLOCK 8\n#define THREADS_PER_BLOCK (WARP_SIZE * WARPS_PER_BLOCK)\n\n#if SHARED_MEMORY_LIMIT_64K\n// With only 64 Kb shared memory available, we can fit two 8-tile chunks of\n// the A and B matrix data, that are 16 * 16 * 8 * 8 * 2 = 32 Kb each\n// (i.e. two 8x8 arrays of tiles of 16x16 uint8_t-typed elements per CTA).\n// But we cannot account the 8 Kb total skew overhead, without which the\n// performance would be severely impacted. So we choose to reduce the chunk size\n// in half, i.e. the amount of A and B matrix data we cache in shared memory.\n// Accordingly, this doubles the number of outer iterations across the global K\n// dimension, which only slightly impacts the performance.\n#define CHUNK_K 8\n#else\n#define CHUNK_K 16\n#endif\n\n#define CHUNK_LINE_BYTES (CHUNK_K * K * sizeof(uint8_t))\n#define WARP_COPY_BYTES (WARP_SIZE * sizeof(int4))\n#define CHUNK_COPY_LINES_PER_WARP (WARP_COPY_BYTES / CHUNK_LINE_BYTES)\n#define CHUNK_COPY_LINE_LANES (WARP_SIZE / CHUNK_COPY_LINES_PER_WARP)\n\n#define BLOCK_ROW_WARPS 2\n#define BLOCK_COL_WARPS 4\n\n#define WARP_ROW_TILES 4\n#define WARP_COL_TILES 2\n\n#define BLOCK_ROW_TILES (WARP_ROW_TILES * BLOCK_ROW_WARPS)\n#define BLOCK_COL_TILES (WARP_COL_TILES * BLOCK_COL_WARPS)\n\n#define GLOBAL_MEM_STRIDE N_GLOBAL\n\n#define SHMEM_STRIDE (N * BLOCK_ROW_TILES)\n#define SHMEM_OFFSET (N * WARP_ROW_TILES)\n\n// The macro below is used to shift rows of the A matrix and columns of the B\n// matrix in shared memory to minimize possible bank conflicts. Before\n// performing the nvcuda::wmma::mma_sync operation, the warp must load the\n// matrix data using the nvcuda::wmma::load_matrix_sync operation. Although the\n// memory access pattern is not specified for that function, each lane in the\n// warp can read one or multiple matrix elements from different matrix rows or\n// columns. For shared memory, such access can result in bank conflicts if\n// different rows / columns of the matrix map to the same bank. By shifting each\n// row and column by a few bytes, we make sure that they map to different banks,\n// thus reducing the number of possible bank conflicts. The number of 32\n// one-byte \"uint8_t\" elements is chosen as the minimum possible shift because\n// we must keep each row and column 256-bit aligned, as required by\n// nvcuda::wmma::load_matrix_sync.\n#define SKEW_UINT8 32\n\n#define checkKernelErrors(expr) \\\n do { \\\n expr; \\\n \\\n hipError_t __err = hipGetLastError(); \\\n if (__err != hipSuccess) { \\\n printf(\"Line %d: '%s' failed: %s\\n\", __LINE__, #expr, hipGetErrorString(__err)); \\\n abort(); \\\n } \\\n } while (0)\n\nusing namespace nvcuda;\n\n__host__ void init_host_matrices(uint8_t *a, uint8_t *b, int *c)\n{\n for (int i = 0; i < M_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n a[i * K_GLOBAL + j] = (uint8_t)(rand() % 3);\n }\n }\n\n for (int i = 0; i < N_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n b[i * K_GLOBAL + j] = (uint8_t)(rand() % 3);\n }\n }\n\n for (int t = 0; t < M_GLOBAL * N_GLOBAL; t++) {\n c[t] = (rand() % 3);\n }\n}\n\n__global__ void compute_gemm_imma(const uint8_t *A, const uint8_t *B, const int *C, int *D, int alpha, int beta)\n{\n extern __shared__ uint8_t shmem[][CHUNK_K * K + SKEW_UINT8];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n const size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n int *shmem_warp_tile_ptr = (int *)&shmem[0][0] + (warpId / 2) * SHMEM_STRIDE * K * 2 + (warpId % 2) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and\n // from shared memory.\n int *shmem_warp_stream_ptr = (int *)&shmem[0][0] + warpId * SHMEM_STRIDE * K;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may\n // result in a loss of precision). Zero still needs to be specially handled\n // though.\n beta /= alpha;\n\n // Each CTA slides along the 128 x 128 tiles from the top left corner of the\n // matrix to the right and down, and selects the next tile to compute. Once\n // there's no such tile, all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared\n // memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const int *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < K; i++) {\n typedef int4 copy_t;\n\n *((copy_t *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId) =\n *((copy_t *)(src_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment\n // multiplications along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const int *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * K + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Scale the C matrix.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const uint8_t *warp_ptr = (warpId < 4) ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % 4) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % 4) * 2);\n\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy\n // the B matrix.\n size_t shmem_idx = warpId < (WARPS_PER_BLOCK / 2)\n ? (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (N * (warpId % (WARPS_PER_BLOCK / 2)) * 2 + shmem_idx_b_off);\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n int4 *lane_ptr = (int4 *)(warp_ptr + tile_k * K + (laneId / CHUNK_COPY_LINE_LANES) * K_GLOBAL)\n + (laneId % CHUNK_COPY_LINE_LANES);\n\n // Shift the second half of the warp to the next row / column in the\n // shared memory.\n shmem_idx += laneId / CHUNK_COPY_LINE_LANES;\n\n#pragma unroll\n for (int i = 0; i < ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP) * 2; i++) {\n // Copy 16 bytes at once in each lane.\n *((int4 *)&shmem[shmem_idx][0] + (laneId % CHUNK_COPY_LINE_LANES)) = *lane_ptr;\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = (int4 *)((uint8_t *)lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP);\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / 2) * M * 2 + (i * M);\n const uint8_t *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_UINT8);\n\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be\n // reused against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const uint8_t *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_UINT8);\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n // Uniform, point-wise transformations of ALL fragment elements by ALL\n // threads in the warp are well-defined even though element indices\n // within fragment storage are not defined.\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n int *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * K + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global\n // memory.\n int *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < K; i++) {\n *((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n}\n\n// Performs an MxNxK GEMM (C=alpha*A*B + beta*C) assuming:\n// 1) Matrices are packed in memory.\n// 2) M, N and K are multiples of 16.\n// 3) Neither A nor B are transposed.\n// Note: This is a less performant version of the compute_gemm_imma kernel. It\n// is designed for\n// demonstration purposes only to show the CUDA WMMA API use without\n// relying on availability of the shared memory.\n__global__ void simple_wmma_gemm_imma(const uint8_t *a,\n const uint8_t *b,\n const int *c,\n int *d,\n int m_ld,\n int n_ld,\n int k_ld,\n int alpha,\n int beta)\n{\n // Leading dimensions. Packed with no transpositions.\n int lda = m_ld;\n int ldb = k_ld;\n int ldc = n_ld;\n\n // Tile using a 2D grid\n int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;\n int warpN = (blockIdx.y * blockDim.y + threadIdx.y);\n\n // Declare the fragments\n wmma::fragment a_frag;\n wmma::fragment b_frag;\n wmma::fragment acc_frag;\n wmma::fragment c_frag;\n\n wmma::fill_fragment(acc_frag, 0.0f);\n\n // Loop over k\n for (int i = 0; i < k_ld; i += WMMA_K) {\n int aCol = i;\n int aRow = warpM * WMMA_M;\n\n int bCol = i;\n int bRow = warpN * WMMA_N;\n\n // Bounds checking\n if (aRow < m_ld && aCol < k_ld && bRow < k_ld && bCol < n_ld) {\n // Load the inputs\n wmma::load_matrix_sync(a_frag, a + aCol + aRow * lda, lda);\n wmma::load_matrix_sync(b_frag, b + bCol + bRow * ldb, ldb);\n\n // Perform the matrix multiplication\n wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);\n }\n }\n\n // Load in the current value of c, scale it by beta, and add this our result\n // scaled by alpha\n int cCol = warpN * WMMA_N;\n int cRow = warpM * WMMA_M;\n\n if (cRow < m_ld && cCol < n_ld) {\n wmma::load_matrix_sync(c_frag, c + cCol + cRow * ldc, ldc, wmma::mem_row_major);\n\n for (int i = 0; i < c_frag.num_elements; i++) {\n c_frag.x[i] = alpha * acc_frag.x[i] + beta * c_frag.x[i];\n }\n\n // Store the output\n wmma::store_matrix_sync(d + cCol + cRow * ldc, c_frag, ldc, wmma::mem_row_major);\n }\n}\n\n__host__ void matMultiplyOnHost(uint8_t *A,\n uint8_t *B,\n int *C,\n int alpha,\n int beta,\n int numARows,\n int numAColumns,\n int numBRows,\n int numBColumns,\n int numCRows,\n int numCColumns)\n{\n for (int i = 0; i < numCRows; i++) {\n for (int j = 0; j < numCColumns; j++) {\n int temp = 0;\n\n for (int k = 0; k < numAColumns; k++) {\n temp += A[i * numAColumns + k] * B[j * numBRows + k];\n }\n\n C[i * numCColumns + j] = temp * alpha + beta * C[i * numCColumns + j];\n }\n }\n}\n\nint main(int argc, char **argv)\n{\n printf(\"Initializing...\\n\");\n\n int dev = findCudaDevice(argc, (const char **)argv);\n\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, dev));\n\n // Tensor cores require a GPU of Volta (SM72) architecture or higher.\n if (deviceProp.major < 7 || (deviceProp.major <= 7 && deviceProp.minor < 2)) {\n printf(\"immaTensorCoreGemm requires SM 7.2 or higher to use Tensor Cores. \"\n \"Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n printf(\"M: %d (%d x %d)\\n\", M_GLOBAL, M, M_TILES);\n printf(\"N: %d (%d x %d)\\n\", N_GLOBAL, N, N_TILES);\n printf(\"K: %d (%d x %d)\\n\", K_GLOBAL, K, K_TILES);\n\n uint8_t *A_h = NULL;\n uint8_t *B_h = NULL;\n int *C_h = NULL;\n#if CPU_DEBUG\n int *result_hD = NULL;\n int *result_host = NULL;\n#endif\n\n A_h = (uint8_t *)malloc(sizeof(uint8_t) * M_GLOBAL * K_GLOBAL);\n B_h = (uint8_t *)malloc(sizeof(uint8_t) * K_GLOBAL * N_GLOBAL);\n C_h = (int *)malloc(sizeof(int) * M_GLOBAL * N_GLOBAL);\n#if CPU_DEBUG\n result_hD = (int *)malloc(sizeof(int) * M_GLOBAL * N_GLOBAL);\n result_host = (int *)malloc(sizeof(int) * M_GLOBAL * N_GLOBAL);\n#endif\n\n uint8_t *A = NULL;\n uint8_t *B = NULL;\n int *C = NULL;\n int *D = NULL;\n\n checkCudaErrors(hipMalloc(reinterpret_cast(&A), sizeof(uint8_t) * M_GLOBAL * K_GLOBAL));\n checkCudaErrors(hipMalloc(reinterpret_cast(&B), sizeof(uint8_t) * N_GLOBAL * K_GLOBAL));\n checkCudaErrors(hipMalloc(reinterpret_cast(&C), sizeof(int) * M_GLOBAL * N_GLOBAL));\n checkCudaErrors(hipMalloc(reinterpret_cast(&D), sizeof(int) * M_GLOBAL * N_GLOBAL));\n\n assert(((unsigned long long)A) % 128 == 0);\n assert(((unsigned long long)B) % 128 == 0);\n assert(((unsigned long long)C) % 128 == 0);\n assert(((unsigned long long)D) % 128 == 0);\n\n init_host_matrices(A_h, B_h, C_h);\n\n checkCudaErrors(hipMemcpy(A, A_h, sizeof(uint8_t) * M_GLOBAL * K_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(B, B_h, sizeof(uint8_t) * N_GLOBAL * K_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(C, C_h, sizeof(int) * M_GLOBAL * N_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemset(D, 0, sizeof(int) * M_GLOBAL * N_GLOBAL));\n\n printf(\"Preparing data for GPU...\\n\");\n\n assert(((unsigned long long)A) % 128 == 0);\n assert(((unsigned long long)B) % 128 == 0);\n assert(((unsigned long long)C) % 128 == 0);\n assert(((unsigned long long)D) % 128 == 0);\n\n enum {\n // Compute the right amount of shared memory to request.\n // We need shared memory to hold per-CTA C and D matrix tiles, and to cache\n // per-CTA chunks\n // of the A and B matrices. Therefore, the right amount to request is the\n // maximum of those\n // two numbers.\n SHMEM_SZ = MAX(sizeof(uint8_t) * (BLOCK_COL_TILES * M) * (CHUNK_K * K + SKEW_UINT8) * 2,\n M * (BLOCK_ROW_WARPS * WARP_ROW_TILES) * N * (BLOCK_COL_WARPS * WARP_COL_TILES) * sizeof(int))\n };\n\n printf(\"Required shared memory size: %lu Kb\\n\", SHMEM_SZ / 1024UL);\n\n int alpha = 1;\n int beta = 1;\n\n hipEvent_t start, stop;\n\n checkCudaErrors(hipEventCreate(&start));\n checkCudaErrors(hipEventCreate(&stop));\n checkCudaErrors(hipEventRecord(start));\n\n // If enough shared memory available on the GPU use high performant kernel\n if (deviceProp.sharedMemPerMultiprocessor >= SHMEM_SZ) {\n printf(\"Computing... using high performance kernel compute_gemm_imma \\n\");\n\n checkCudaErrors(hipFuncSetAttribute(reinterpret_cast(compute_gemm_imma), hipFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors((compute_gemm_imma<<>>(\n A, B, C, D, alpha, beta)));\n#if CPU_DEBUG\n checkCudaErrors(hipMemcpy(result_hD, D, sizeof(int) * M_GLOBAL * N_GLOBAL, hipMemcpyDeviceToHost));\n#endif\n }\n else {\n dim3 gridDim;\n dim3 blockDim;\n\n // blockDim.x must be a multiple of warpSize\n // 128x4 means we have 16 warps and a block computes a 64x64 output tile\n blockDim.x = 128;\n blockDim.y = 4;\n\n gridDim.x = (M_GLOBAL + (WMMA_M * blockDim.x / 32 - 1)) / (WMMA_M * blockDim.x / 32);\n gridDim.y = (N_GLOBAL + WMMA_N * blockDim.y - 1) / (WMMA_N * blockDim.y);\n\n printf(\"Computing... using simple_wmma_gemm_imma kernel\\n\");\n simple_wmma_gemm_imma<<>>(A, B, C, D, M_GLOBAL, N_GLOBAL, K_GLOBAL, alpha, beta);\n#if CPU_DEBUG\n checkCudaErrors(hipMemcpy(result_hD, D, sizeof(int) * M_GLOBAL * N_GLOBAL, hipMemcpyDeviceToHost));\n#endif\n }\n\n checkCudaErrors(hipEventRecord(stop));\n checkCudaErrors(hipEventSynchronize(stop));\n\n#if CPU_DEBUG\n printf(\"Verifying correctness of the computations...\\n\");\n\n memcpy(result_host, C_h, sizeof(int) * M_GLOBAL * N_GLOBAL);\n\n matMultiplyOnHost(A_h, B_h, result_host, alpha, beta, M_GLOBAL, K_GLOBAL, K_GLOBAL, N_GLOBAL, M_GLOBAL, N_GLOBAL);\n\n for (int i = 0; i < N_GLOBAL * M_GLOBAL; i++) {\n if (abs(result_hD[i] - result_host[i]) > 0) {\n printf(\"mismatch i=%d result_hD=%d result_host=%d\\n\", i, result_hD[i], result_host[i]);\n }\n }\n free(result_host);\n free(result_hD);\n#endif\n\n float milliseconds = 0;\n\n checkCudaErrors(hipEventElapsedTime(&milliseconds, start, stop));\n\n printf(\"Time: %f ms\\n\", milliseconds);\n printf(\"TOPS: %.2f\\n\", (((double)M_GLOBAL * N_GLOBAL * K_GLOBAL * 2) / (milliseconds / 1000.)) / 1e12);\n\n free(A_h);\n free(B_h);\n free(C_h);\n checkCudaErrors(hipFree(reinterpret_cast(A)));\n checkCudaErrors(hipFree(reinterpret_cast(B)));\n checkCudaErrors(hipFree(reinterpret_cast(C)));\n checkCudaErrors(hipFree(reinterpret_cast(D)));\n\n return EXIT_SUCCESS;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/immaTensorCoreGemm/immaTensorCoreGemm.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// This sample demonstrates dynamic global memory allocation through device C++\n// new and delete operators and virtual function declarations available with\n// CUDA 4.0.\n\n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \n#include \n#include \n#include \n\nconst char *sSDKsample = \"newdelete\";\n\n#include \"container.hpp\"\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// Kernels to allocate and instantiate Container objects on the device heap\n//\n////////////////////////////////////////////////////////////////////////////////\n\n__global__ void vectorCreate(Container **g_container, int max_size)\n{\n // The Vector object and the data storage are allocated in device heap memory.\n // This makes it persistent for the lifetime of the CUDA context.\n // The grid has only one thread as only a single object instance is needed.\n\n *g_container = new Vector(max_size);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// Kernels to fill and consume shared Container objects.\n//\n////////////////////////////////////////////////////////////////////////////////\n\n__global__ void containerFill(Container **g_container)\n{\n // All threads of the grid cooperatively populate the shared Container object\n // with data.\n if (threadIdx.x == 0) {\n (*g_container)->push(blockIdx.x);\n }\n}\n\n__global__ void containerConsume(Container **g_container, int *d_result)\n{\n // All threads of the grid cooperatively consume the data from the shared\n // Container object.\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n int v;\n\n if ((*g_container)->pop(v)) {\n d_result[idx] = v;\n }\n else {\n d_result[idx] = -1;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// Kernel to delete shared Container objects.\n//\n////////////////////////////////////////////////////////////////////////////////\n\n__global__ void containerDelete(Container **g_container) { delete *g_container; }\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// Kernels to using of placement new to put shared Vector objects and data in\n// shared memory\n//\n////////////////////////////////////////////////////////////////////////////////\n\n__global__ void placementNew(int *d_result)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ unsigned char __align__(8) s_buffer[sizeof(Vector)];\n __shared__ int __align__(8) s_data[1024];\n __shared__ Vector *s_vector;\n\n // The first thread of the block initializes the shared Vector object.\n // The placement new operator enables the Vector object and the data array top\n // be placed in shared memory.\n if (threadIdx.x == 0) {\n s_vector = new (s_buffer) Vector(1024, s_data);\n }\n\n cg::sync(cta);\n\n if ((threadIdx.x & 1) == 0) {\n s_vector->push(threadIdx.x >> 1);\n }\n\n // Need to sync as the vector implementation does not support concurrent\n // push/pop operations.\n cg::sync(cta);\n\n int v;\n\n if (s_vector->pop(v)) {\n d_result[threadIdx.x] = v;\n }\n else {\n d_result[threadIdx.x] = -1;\n }\n\n // Note: deleting objects placed in shared memory is not necessary (lifetime\n // of shared memory is that of the block)\n}\n\nstruct ComplexType_t\n{\n int a;\n int b;\n float c;\n float d;\n};\n\n__global__ void complexVector(int *d_result)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ unsigned char __align__(8) s_buffer[sizeof(Vector)];\n __shared__ ComplexType_t __align__(8) s_data[1024];\n __shared__ Vector *s_vector;\n\n // The first thread of the block initializes the shared Vector object.\n // The placement new operator enables the Vector object and the data array top\n // be placed in shared memory.\n if (threadIdx.x == 0) {\n s_vector = new (s_buffer) Vector(1024, s_data);\n }\n\n cg::sync(cta);\n\n if ((threadIdx.x & 1) == 0) {\n ComplexType_t data;\n data.a = threadIdx.x >> 1;\n data.b = blockIdx.x;\n data.c = threadIdx.x / (float)(blockDim.x);\n data.d = blockIdx.x / (float)(gridDim.x);\n\n s_vector->push(data);\n }\n\n cg::sync(cta);\n\n ComplexType_t v;\n\n if (s_vector->pop(v)) {\n d_result[threadIdx.x] = v.a;\n }\n else {\n d_result[threadIdx.x] = -1;\n }\n\n // Note: deleting objects placed in shared memory is not necessary (lifetime\n // of shared memory is that of the block)\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// Host code\n//\n////////////////////////////////////////////////////////////////////////////////\n\nbool checkResult(int *d_result, int N)\n{\n std::vector h_result;\n h_result.resize(N);\n\n checkCudaErrors(cudaMemcpy(&h_result[0], d_result, N * sizeof(int), cudaMemcpyDeviceToHost));\n std::sort(h_result.begin(), h_result.end());\n\n bool success = true;\n bool test = false;\n\n int value = 0;\n\n for (int i = 0; i < N; ++i) {\n if (h_result[i] != -1) {\n test = true;\n }\n\n if (test && (value++) != h_result[i]) {\n success = false;\n }\n }\n\n return success;\n}\n\nbool testContainer(Container **d_container, int blocks, int threads)\n{\n int *d_result;\n cudaMalloc(&d_result, blocks * threads * sizeof(int));\n\n containerFill<<>>(d_container);\n containerConsume<<>>(d_container, d_result);\n containerDelete<<<1, 1>>>(d_container);\n checkCudaErrors(cudaDeviceSynchronize());\n\n bool success = checkResult(d_result, blocks * threads);\n\n cudaFree(d_result);\n\n return success;\n}\n\nbool testPlacementNew(int threads)\n{\n int *d_result;\n cudaMalloc(&d_result, threads * sizeof(int));\n\n placementNew<<<1, threads>>>(d_result);\n checkCudaErrors(cudaDeviceSynchronize());\n\n bool success = checkResult(d_result, threads);\n\n cudaFree(d_result);\n\n return success;\n}\n\nbool testComplexType(int threads)\n{\n int *d_result;\n cudaMalloc(&d_result, threads * sizeof(int));\n\n complexVector<<<1, threads>>>(d_result);\n checkCudaErrors(cudaDeviceSynchronize());\n\n bool success = checkResult(d_result, threads);\n\n cudaFree(d_result);\n\n return success;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// MAIN\n//\n////////////////////////////////////////////////////////////////////////////////\n\nint main(int argc, char **argv)\n{\n printf(\"%s Starting...\\n\\n\", sSDKsample);\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n findCudaDevice(argc, (const char **)argv);\n\n // set the heap size for device size new/delete to 128 MB\n checkCudaErrors(cudaDeviceSetLimit(cudaLimitMallocHeapSize, 128 * (1 << 20)));\n\n Container **d_container;\n checkCudaErrors(cudaMalloc(&d_container, sizeof(Container **)));\n\n bool bTest = false;\n int test_passed = 0;\n\n printf(\" > Container = Vector test \");\n vectorCreate<<<1, 1>>>(d_container, 128 * 128);\n bTest = testContainer(d_container, 128, 128);\n printf(bTest ? \"OK\\n\\n\" : \"NOT OK\\n\\n\");\n test_passed += (bTest ? 1 : 0);\n\n checkCudaErrors(cudaFree(d_container));\n\n printf(\" > Container = Vector, using placement new on SMEM buffer test \");\n bTest = testPlacementNew(1024);\n printf(bTest ? \"OK\\n\\n\" : \"NOT OK\\n\\n\");\n test_passed += (bTest ? 1 : 0);\n\n printf(\" > Container = Vector, with user defined datatype test \");\n bTest = testComplexType(1024);\n printf(bTest ? \"OK\\n\\n\" : \"NOT OK\\n\\n\");\n test_passed += (bTest ? 1 : 0);\n\n printf(\"Test Summary: %d/3 succesfully run\\n\", test_passed);\n\n exit(test_passed == 3 ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// This sample demonstrates dynamic global memory allocation through device C++\n// new and delete operators and virtual function declarations available with\n// CUDA 4.0.\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \n#include \n#include \n#include \n\nconst char *sSDKsample = \"newdelete\";\n\n#include \"container.hpp\"\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// Kernels to allocate and instantiate Container objects on the device heap\n//\n////////////////////////////////////////////////////////////////////////////////\n\n__global__ void vectorCreate(Container **g_container, int max_size)\n{\n // The Vector object and the data storage are allocated in device heap memory.\n // This makes it persistent for the lifetime of the CUDA context.\n // The grid has only one thread as only a single object instance is needed.\n\n *g_container = new Vector(max_size);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// Kernels to fill and consume shared Container objects.\n//\n////////////////////////////////////////////////////////////////////////////////\n\n__global__ void containerFill(Container **g_container)\n{\n // All threads of the grid cooperatively populate the shared Container object\n // with data.\n if (threadIdx.x == 0) {\n (*g_container)->push(blockIdx.x);\n }\n}\n\n__global__ void containerConsume(Container **g_container, int *d_result)\n{\n // All threads of the grid cooperatively consume the data from the shared\n // Container object.\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n int v;\n\n if ((*g_container)->pop(v)) {\n d_result[idx] = v;\n }\n else {\n d_result[idx] = -1;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// Kernel to delete shared Container objects.\n//\n////////////////////////////////////////////////////////////////////////////////\n\n__global__ void containerDelete(Container **g_container) { delete *g_container; }\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// Kernels to using of placement new to put shared Vector objects and data in\n// shared memory\n//\n////////////////////////////////////////////////////////////////////////////////\n\n__global__ void placementNew(int *d_result)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ unsigned char __align__(8) s_buffer[sizeof(Vector)];\n __shared__ int __align__(8) s_data[1024];\n __shared__ Vector *s_vector;\n\n // The first thread of the block initializes the shared Vector object.\n // The placement new operator enables the Vector object and the data array top\n // be placed in shared memory.\n if (threadIdx.x == 0) {\n s_vector = new (s_buffer) Vector(1024, s_data);\n }\n\n cg::sync(cta);\n\n if ((threadIdx.x & 1) == 0) {\n s_vector->push(threadIdx.x >> 1);\n }\n\n // Need to sync as the vector implementation does not support concurrent\n // push/pop operations.\n cg::sync(cta);\n\n int v;\n\n if (s_vector->pop(v)) {\n d_result[threadIdx.x] = v;\n }\n else {\n d_result[threadIdx.x] = -1;\n }\n\n // Note: deleting objects placed in shared memory is not necessary (lifetime\n // of shared memory is that of the block)\n}\n\nstruct ComplexType_t\n{\n int a;\n int b;\n float c;\n float d;\n};\n\n__global__ void complexVector(int *d_result)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ unsigned char __align__(8) s_buffer[sizeof(Vector)];\n __shared__ ComplexType_t __align__(8) s_data[1024];\n __shared__ Vector *s_vector;\n\n // The first thread of the block initializes the shared Vector object.\n // The placement new operator enables the Vector object and the data array top\n // be placed in shared memory.\n if (threadIdx.x == 0) {\n s_vector = new (s_buffer) Vector(1024, s_data);\n }\n\n cg::sync(cta);\n\n if ((threadIdx.x & 1) == 0) {\n ComplexType_t data;\n data.a = threadIdx.x >> 1;\n data.b = blockIdx.x;\n data.c = threadIdx.x / (float)(blockDim.x);\n data.d = blockIdx.x / (float)(gridDim.x);\n\n s_vector->push(data);\n }\n\n cg::sync(cta);\n\n ComplexType_t v;\n\n if (s_vector->pop(v)) {\n d_result[threadIdx.x] = v.a;\n }\n else {\n d_result[threadIdx.x] = -1;\n }\n\n // Note: deleting objects placed in shared memory is not necessary (lifetime\n // of shared memory is that of the block)\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// Host code\n//\n////////////////////////////////////////////////////////////////////////////////\n\nbool checkResult(int *d_result, int N)\n{\n std::vector h_result;\n h_result.resize(N);\n\n checkCudaErrors(hipMemcpy(&h_result[0], d_result, N * sizeof(int), hipMemcpyDeviceToHost));\n std::sort(h_result.begin(), h_result.end());\n\n bool success = true;\n bool test = false;\n\n int value = 0;\n\n for (int i = 0; i < N; ++i) {\n if (h_result[i] != -1) {\n test = true;\n }\n\n if (test && (value++) != h_result[i]) {\n success = false;\n }\n }\n\n return success;\n}\n\nbool testContainer(Container **d_container, int blocks, int threads)\n{\n int *d_result;\n hipMalloc(&d_result, blocks * threads * sizeof(int));\n\n containerFill<<>>(d_container);\n containerConsume<<>>(d_container, d_result);\n containerDelete<<<1, 1>>>(d_container);\n checkCudaErrors(hipDeviceSynchronize());\n\n bool success = checkResult(d_result, blocks * threads);\n\n hipFree(d_result);\n\n return success;\n}\n\nbool testPlacementNew(int threads)\n{\n int *d_result;\n hipMalloc(&d_result, threads * sizeof(int));\n\n placementNew<<<1, threads>>>(d_result);\n checkCudaErrors(hipDeviceSynchronize());\n\n bool success = checkResult(d_result, threads);\n\n hipFree(d_result);\n\n return success;\n}\n\nbool testComplexType(int threads)\n{\n int *d_result;\n hipMalloc(&d_result, threads * sizeof(int));\n\n complexVector<<<1, threads>>>(d_result);\n checkCudaErrors(hipDeviceSynchronize());\n\n bool success = checkResult(d_result, threads);\n\n hipFree(d_result);\n\n return success;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// MAIN\n//\n////////////////////////////////////////////////////////////////////////////////\n\nint main(int argc, char **argv)\n{\n printf(\"%s Starting...\\n\\n\", sSDKsample);\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n findCudaDevice(argc, (const char **)argv);\n\n // set the heap size for device size new/delete to 128 MB\n checkCudaErrors(hipDeviceSetLimit(hipLimitMallocHeapSize, 128 * (1 << 20)));\n\n Container **d_container;\n checkCudaErrors(hipMalloc(&d_container, sizeof(Container **)));\n\n bool bTest = false;\n int test_passed = 0;\n\n printf(\" > Container = Vector test \");\n vectorCreate<<<1, 1>>>(d_container, 128 * 128);\n bTest = testContainer(d_container, 128, 128);\n printf(bTest ? \"OK\\n\\n\" : \"NOT OK\\n\\n\");\n test_passed += (bTest ? 1 : 0);\n\n checkCudaErrors(hipFree(d_container));\n\n printf(\" > Container = Vector, using placement new on SMEM buffer test \");\n bTest = testPlacementNew(1024);\n printf(bTest ? \"OK\\n\\n\" : \"NOT OK\\n\\n\");\n test_passed += (bTest ? 1 : 0);\n\n printf(\" > Container = Vector, with user defined datatype test \");\n bTest = testComplexType(1024);\n printf(bTest ? \"OK\\n\\n\" : \"NOT OK\\n\\n\");\n test_passed += (bTest ? 1 : 0);\n\n printf(\"Test Summary: %d/3 succesfully run\\n\", test_passed);\n\n exit(test_passed == 3 ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/newdelete/newdelete.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n\n// Device code\nextern \"C\" __global__ void memMapIpc_kernel(char *ptr, int sz, char val)\n{\n // Dummy kernel\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n for (; idx < sz; idx += (gridDim.x * blockDim.x)) {\n ptr[idx] = val;\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n\n// Device code\n#include \"hip/hip_runtime.h\"\nextern \"C\" __global__ void memMapIpc_kernel(char *ptr, int sz, char val)\n{\n // Dummy kernel\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n for (; idx < sz; idx += (gridDim.x * blockDim.x)) {\n ptr[idx] = val;\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/memMapIPCDrv/memMapIpc_kernel.cu"} {"input": "/* Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This file demonstrates the usage of conditional graph nodes with\n * a series of *simple* example graphs.\n *\n * For more information on conditional nodes, see the programming guide:\n *\n * https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#conditional-graph-nodes\n *\n */\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n/*\n * Create a graph containing two nodes.\n * The first node, A, is a kernel and the second node, B, is a conditional IF node.\n * The kernel sets the condition variable to true if a device memory location\n * contains an odd number. Otherwise the condition variable is set to false.\n * There is a single kernel, C, within the conditional body which prints a message.\n *\n * A -> B [ C ]\n *\n */\n\n__global__ void ifGraphKernelA(char *dPtr, cudaGraphConditionalHandle handle)\n{\n // In this example, condition is set if *dPtr is odd\n unsigned int value = *dPtr & 0x01;\n cudaGraphSetConditional(handle, value);\n printf(\"GPU: Handle set to %d\\n\", value);\n}\n\n// This kernel will only be executed if the condition is true\n__global__ void ifGraphKernelC(void) { printf(\"GPU: Hello from the GPU! The condition was true.\\n\"); }\n\n// Setup and launch the graph\nvoid simpleIfGraph(void)\n{\n cudaGraph_t graph;\n cudaGraphExec_t graphExec;\n cudaGraphNode_t kernelNode;\n cudaGraphNode_t conditionalNode;\n\n void *kernelArgs[2];\n\n // Allocate a byte of device memory to use as input\n char *dPtr;\n checkCudaErrors(cudaMalloc((void **)&dPtr, 1));\n\n printf(\"simpleIfGraph: Building graph...\\n\");\n cudaGraphCreate(&graph, 0);\n\n // Create conditional handle.\n cudaGraphConditionalHandle handle;\n cudaGraphConditionalHandleCreate(&handle, graph);\n\n // Use a kernel upstream of the conditional to set the handle value\n cudaGraphNodeParams params = {cudaGraphNodeTypeKernel};\n params.kernel.func = (void *)ifGraphKernelA;\n params.kernel.blockDim.x = params.kernel.blockDim.y = params.kernel.blockDim.z = 1;\n params.kernel.gridDim.x = params.kernel.gridDim.y = params.kernel.gridDim.z = 1;\n params.kernel.kernelParams = kernelArgs;\n kernelArgs[0] = &dPtr;\n kernelArgs[1] = &handle;\n checkCudaErrors(cudaGraphAddNode(&kernelNode, graph, NULL, NULL, 0, ¶ms));\n\n cudaGraphNodeParams cParams = {cudaGraphNodeTypeConditional};\n cParams.conditional.handle = handle;\n cParams.conditional.type = cudaGraphCondTypeIf;\n cParams.conditional.size = 1;\n checkCudaErrors(cudaGraphAddNode(&conditionalNode, graph, &kernelNode, NULL, 0, &cParams));\n\n cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];\n\n // Populate the body of the conditional node\n cudaGraphNode_t bodyNode;\n params.kernel.func = (void *)ifGraphKernelC;\n params.kernel.kernelParams = nullptr;\n checkCudaErrors(cudaGraphAddNode(&bodyNode, bodyGraph, NULL, NULL, 0, ¶ms));\n\n checkCudaErrors(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n // Initialize device memory and launch the graph\n checkCudaErrors(cudaMemset(dPtr, 0, 1)); // Set dPtr to 0\n printf(\"Host: Launching graph with device memory set to 0\\n\");\n checkCudaErrors(cudaGraphLaunch(graphExec, 0));\n checkCudaErrors(cudaDeviceSynchronize());\n\n // Initialize device memory and launch the graph\n checkCudaErrors(cudaMemset(dPtr, 1, 1)); // Set dPtr to 1\n printf(\"Host: Launching graph with device memory set to 1\\n\");\n checkCudaErrors(cudaGraphLaunch(graphExec, 0));\n checkCudaErrors(cudaDeviceSynchronize());\n\n // Cleanup\n checkCudaErrors(cudaGraphExecDestroy(graphExec));\n checkCudaErrors(cudaGraphDestroy(graph));\n checkCudaErrors(cudaFree(dPtr));\n\n printf(\"simpleIfGraph: Complete\\n\\n\");\n}\n\n/*\n * Create a graph containing a single conditional while node.\n * The default value of the conditional variable is set to true, so this\n * effectively becomes a do-while loop as the conditional body will always\n * execute at least once. The body of the conditional contains 3 kernel nodes:\n * A [ B -> C -> D ]\n * Nodes B and C are just dummy nodes for demonstrative purposes. Node D\n * will decrement a device memory location and set the condition value to false\n * when the value reaches zero, terminating the loop.\n * In this example, stream capture is used to populate the conditional body.\n */\n\n// This kernel will only be executed if the condition is true\n__global__ void doWhileEmptyKernel(void)\n{\n printf(\"GPU: doWhileEmptyKernel()\\n\");\n return;\n}\n\n__global__ void doWhileLoopKernel(char *dPtr, cudaGraphConditionalHandle handle)\n{\n if (--(*dPtr) == 0) {\n cudaGraphSetConditional(handle, 0);\n }\n printf(\"GPU: counter = %d\\n\", *dPtr);\n}\n\nvoid simpleDoWhileGraph(void)\n{\n cudaGraph_t graph;\n cudaGraphExec_t graphExec;\n cudaGraphNode_t conditionalNode;\n\n // Allocate a byte of device memory to use as input\n char *dPtr;\n checkCudaErrors(cudaMalloc((void **)&dPtr, 1));\n\n printf(\"simpleDoWhileGraph: Building graph...\\n\");\n checkCudaErrors(cudaGraphCreate(&graph, 0));\n\n cudaGraphConditionalHandle handle;\n checkCudaErrors(cudaGraphConditionalHandleCreate(&handle, graph, 1, cudaGraphCondAssignDefault));\n\n cudaGraphNodeParams cParams = {cudaGraphNodeTypeConditional};\n cParams.conditional.handle = handle;\n cParams.conditional.type = cudaGraphCondTypeWhile;\n cParams.conditional.size = 1;\n checkCudaErrors(cudaGraphAddNode(&conditionalNode, graph, NULL, NULL, 0, &cParams));\n\n cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];\n\n cudaStream_t captureStream;\n checkCudaErrors(cudaStreamCreate(&captureStream));\n\n checkCudaErrors(\n cudaStreamBeginCaptureToGraph(captureStream, bodyGraph, nullptr, nullptr, 0, cudaStreamCaptureModeGlobal));\n doWhileEmptyKernel<<<1, 1, 0, captureStream>>>();\n doWhileEmptyKernel<<<1, 1, 0, captureStream>>>();\n doWhileLoopKernel<<<1, 1, 0, captureStream>>>(dPtr, handle);\n checkCudaErrors(cudaStreamEndCapture(captureStream, nullptr));\n checkCudaErrors(cudaStreamDestroy(captureStream));\n\n checkCudaErrors(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n // Initialize device memory and launch the graph\n checkCudaErrors(cudaMemset(dPtr, 10, 1)); // Set dPtr to 10\n printf(\"Host: Launching graph with loop counter set to 10\\n\");\n checkCudaErrors(cudaGraphLaunch(graphExec, 0));\n checkCudaErrors(cudaDeviceSynchronize());\n\n // Cleanup\n checkCudaErrors(cudaGraphExecDestroy(graphExec));\n checkCudaErrors(cudaGraphDestroy(graph));\n checkCudaErrors(cudaFree(dPtr));\n\n printf(\"simpleDoWhileGraph: Complete\\n\\n\");\n}\n\n/*\n * Create a graph containing a conditional while loop using stream capture.\n * This demonstrates how to insert a conditional node into a stream which is\n * being captured. The graph consists of a kernel node, A, followed by a\n * conditional while node, B, followed by a kernel node, D. The conditional\n * body is populated by a single kernel node, C:\n *\n * A -> B [ C ] -> D\n *\n * The same kernel will be used for both nodes A and C. This kernel will test\n * a device memory location and set the condition when the location is non-zero.\n * We must run the kernel before the loop as well as inside the loop in order\n * to behave like a while loop as opposed to a do-while loop. We need to evaluate\n * the device memory location before the conditional node is evaluated in order\n * to set the condition variable properly. Because we're using a kernel upstream\n * of the conditional node, there is no need to use the handle default value to\n * initialize the conditional value.\n */\n\n__global__ void capturedWhileKernel(char *dPtr, cudaGraphConditionalHandle handle)\n{\n printf(\"GPU: counter = %d\\n\", *dPtr);\n if (*dPtr) {\n (*dPtr)--;\n }\n cudaGraphSetConditional(handle, *dPtr);\n}\n\n__global__ void capturedWhileEmptyKernel(void)\n{\n printf(\"GPU: capturedWhileEmptyKernel()\\n\");\n return;\n}\n\nvoid capturedWhileGraph(void)\n{\n cudaGraph_t graph;\n cudaGraphExec_t graphExec;\n\n cudaStreamCaptureStatus status;\n const cudaGraphNode_t *dependencies;\n size_t numDependencies;\n\n // Allocate a byte of device memory to use as input\n char *dPtr;\n checkCudaErrors(cudaMalloc((void **)&dPtr, 1));\n\n printf(\"capturedWhileGraph: Building graph...\\n\");\n cudaStream_t captureStream;\n checkCudaErrors(cudaStreamCreate(&captureStream));\n\n checkCudaErrors(cudaStreamBeginCapture(captureStream, cudaStreamCaptureModeGlobal));\n\n // Obtain the handle of the graph\n checkCudaErrors(\n cudaStreamGetCaptureInfo(captureStream, &status, NULL, &graph, &dependencies, NULL, &numDependencies));\n\n // Create the conditional handle\n cudaGraphConditionalHandle handle;\n checkCudaErrors(cudaGraphConditionalHandleCreate(&handle, graph));\n\n // Insert kernel node A\n capturedWhileKernel<<<1, 1, 0, captureStream>>>(dPtr, handle);\n\n // Obtain the handle for node A\n checkCudaErrors(\n cudaStreamGetCaptureInfo(captureStream, &status, NULL, &graph, &dependencies, NULL, &numDependencies));\n\n // Insert conditional node B\n cudaGraphNode_t conditionalNode;\n cudaGraphNodeParams cParams = {cudaGraphNodeTypeConditional};\n cParams.conditional.handle = handle;\n cParams.conditional.type = cudaGraphCondTypeWhile;\n cParams.conditional.size = 1;\n checkCudaErrors(cudaGraphAddNode(&conditionalNode, graph, dependencies, NULL, numDependencies, &cParams));\n\n cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];\n\n // Update stream capture dependencies to account for the node we manually added\n checkCudaErrors(cudaStreamUpdateCaptureDependencies(\n captureStream, &conditionalNode, NULL, 1, cudaStreamSetCaptureDependencies));\n\n // Insert kernel node D\n capturedWhileEmptyKernel<<<1, 1, 0, captureStream>>>();\n\n checkCudaErrors(cudaStreamEndCapture(captureStream, &graph));\n checkCudaErrors(cudaStreamDestroy(captureStream));\n\n // Populate conditional body graph using stream capture\n cudaStream_t bodyStream;\n checkCudaErrors(cudaStreamCreate(&bodyStream));\n\n checkCudaErrors(\n cudaStreamBeginCaptureToGraph(bodyStream, bodyGraph, nullptr, nullptr, 0, cudaStreamCaptureModeGlobal));\n\n // Insert kernel node C\n capturedWhileKernel<<<1, 1, 0, bodyStream>>>(dPtr, handle);\n checkCudaErrors(cudaStreamEndCapture(bodyStream, nullptr));\n checkCudaErrors(cudaStreamDestroy(bodyStream));\n\n checkCudaErrors(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n // Initialize device memory and launch the graph\n // Device memory is zero, so the conditional node will not execute\n checkCudaErrors(cudaMemset(dPtr, 0, 1)); // Set dPtr to 0\n printf(\"Host: Launching graph with loop counter set to 0\\n\");\n checkCudaErrors(cudaGraphLaunch(graphExec, 0));\n checkCudaErrors(cudaDeviceSynchronize());\n\n // Initialize device memory and launch the graph\n checkCudaErrors(cudaMemset(dPtr, 10, 1)); // Set dPtr to 10\n printf(\"Host: Launching graph with loop counter set to 10\\n\");\n checkCudaErrors(cudaGraphLaunch(graphExec, 0));\n checkCudaErrors(cudaDeviceSynchronize());\n\n // Cleanup\n checkCudaErrors(cudaGraphExecDestroy(graphExec));\n checkCudaErrors(cudaGraphDestroy(graph));\n checkCudaErrors(cudaFree(dPtr));\n\n printf(\"capturedWhileGraph: Complete\\n\\n\");\n}\n\n/*\n * Create a graph containing two nodes.\n * The first node, A, is a kernel and the second node, B, is a conditional IF node containing two graphs.\n * The first graph within the conditional will be executed when the condition is true, while the second\n * graph will be executed when the conditional is false.\n * The kernel sets the condition variable to true if a device memory location\n * contains an odd number. Otherwise the condition variable is set to false.\n * There is a single kernel(C & D) within each conditional body which prints a message.\n *\n * A -> B [ C | D ]\n *\n * This example requires CUDA >= 12.8.\n */\n\n// This kernel will only be executed if the condition is false\n__global__ void ifGraphKernelD(void) { printf(\"GPU: Hello from the GPU! The condition was false.\\n\"); }\n\n// Setup and launch the graph\nvoid simpleIfElseGraph(void)\n{\n cudaGraph_t graph;\n cudaGraphExec_t graphExec;\n cudaGraphNode_t kernelNode;\n cudaGraphNode_t conditionalNode;\n\n void *kernelArgs[2];\n\n // Allocate a byte of device memory to use as input\n char *dPtr;\n checkCudaErrors(cudaMalloc((void **)&dPtr, 1));\n\n printf(\"simpleIfElseGraph: Building graph...\\n\");\n cudaGraphCreate(&graph, 0);\n\n // Create conditional handle.\n cudaGraphConditionalHandle handle;\n cudaGraphConditionalHandleCreate(&handle, graph);\n\n // Use a kernel upstream of the conditional to set the handle value\n cudaGraphNodeParams params = {cudaGraphNodeTypeKernel};\n params.kernel.func = (void *)ifGraphKernelA;\n params.kernel.blockDim.x = params.kernel.blockDim.y = params.kernel.blockDim.z = 1;\n params.kernel.gridDim.x = params.kernel.gridDim.y = params.kernel.gridDim.z = 1;\n params.kernel.kernelParams = kernelArgs;\n kernelArgs[0] = &dPtr;\n kernelArgs[1] = &handle;\n checkCudaErrors(cudaGraphAddNode(&kernelNode, graph, NULL, NULL, 0, ¶ms));\n\n cudaGraphNodeParams cParams = {cudaGraphNodeTypeConditional};\n cParams.conditional.handle = handle;\n cParams.conditional.type = cudaGraphCondTypeIf;\n cParams.conditional.size = 2; // Set size to 2 to indicate an ELSE graph will be used\n checkCudaErrors(cudaGraphAddNode(&conditionalNode, graph, &kernelNode, NULL, 0, &cParams));\n\n cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];\n\n // Populate the body of the first graph in the conditional node, executed if the condition is true\n cudaGraphNode_t trueBodyNode;\n params.kernel.func = (void *)ifGraphKernelC;\n params.kernel.kernelParams = nullptr;\n checkCudaErrors(cudaGraphAddNode(&trueBodyNode, bodyGraph, NULL, NULL, 0, ¶ms));\n\n // Populate the body of the second graph in the conditional node, executed if the condition is false\n bodyGraph = cParams.conditional.phGraph_out[1];\n\n cudaGraphNode_t falseBodyNode;\n params.kernel.func = (void *)ifGraphKernelD;\n params.kernel.kernelParams = nullptr;\n checkCudaErrors(cudaGraphAddNode(&falseBodyNode, bodyGraph, NULL, NULL, 0, ¶ms));\n\n checkCudaErrors(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n // Initialize device memory and launch the graph\n checkCudaErrors(cudaMemset(dPtr, 0, 1)); // Set dPtr to 0\n printf(\"Host: Launching graph with device memory set to 0\\n\");\n checkCudaErrors(cudaGraphLaunch(graphExec, 0));\n checkCudaErrors(cudaDeviceSynchronize());\n\n // Initialize device memory and launch the graph\n checkCudaErrors(cudaMemset(dPtr, 1, 1)); // Set dPtr to 1\n printf(\"Host: Launching graph with device memory set to 1\\n\");\n checkCudaErrors(cudaGraphLaunch(graphExec, 0));\n checkCudaErrors(cudaDeviceSynchronize());\n\n // Cleanup\n checkCudaErrors(cudaGraphExecDestroy(graphExec));\n checkCudaErrors(cudaGraphDestroy(graph));\n checkCudaErrors(cudaFree(dPtr));\n\n printf(\"simpleIfElseGraph: Complete\\n\\n\");\n}\n\n/*\n * Create a graph containing two nodes.\n * The first node, A, is a kernel and the second node, B, is a conditional SWITCH node containing four graphs.\n * The nth graph within the conditional will be executed when the condition is n. If conditional >= n,\n * no graph will be executed.\n * Kernel A sets the condition variable to the value stored in a device memory location.\n * This device location is updated from the host with each launch to demonstrate the behavior.\n * There is a single kernel(nodes C, D, E and F) within each conditional body which prints a message.\n *\n * A -> B [ C | D | E | F ]\n *\n * This example requires CUDA >= 12.8.\n */\n\n__global__ void switchGraphKernelA(char *dPtr, cudaGraphConditionalHandle handle)\n{\n unsigned int value = *dPtr;\n cudaGraphSetConditional(handle, value);\n printf(\"GPU: Handle set to %d\\n\", value);\n}\n\n__global__ void switchGraphKernelC(void) { printf(\"GPU: Hello from switchGraphKernelC(), running on the GPU!\\n\"); }\n\n__global__ void switchGraphKernelD(void) { printf(\"GPU: Hello from switchGraphKernelD(), running on the GPU!\\n\"); }\n\n__global__ void switchGraphKernelE(void) { printf(\"GPU: Hello from switchGraphKernelE(), running on the GPU!\\n\"); }\n\n__global__ void switchGraphKernelF(void) { printf(\"GPU: Hello from switchGraphKernelF(), running on the GPU!\\n\"); }\n\n// Setup and launch the graph\nvoid simpleSwitchGraph(void)\n{\n cudaGraph_t graph;\n cudaGraphExec_t graphExec;\n cudaGraphNode_t kernelNode;\n cudaGraphNode_t conditionalNode;\n\n void *kernelArgs[2];\n\n // Allocate a byte of device memory to use as input\n char *dPtr;\n checkCudaErrors(cudaMalloc((void **)&dPtr, 1));\n\n printf(\"simpleSwitchGraph: Building graph...\\n\");\n cudaGraphCreate(&graph, 0);\n\n // Create conditional handle.\n cudaGraphConditionalHandle handle;\n cudaGraphConditionalHandleCreate(&handle, graph);\n\n // Use a kernel upstream of the conditional to set the handle value\n cudaGraphNodeParams params = {cudaGraphNodeTypeKernel};\n params.kernel.func = (void *)switchGraphKernelA;\n params.kernel.blockDim.x = params.kernel.blockDim.y = params.kernel.blockDim.z = 1;\n params.kernel.gridDim.x = params.kernel.gridDim.y = params.kernel.gridDim.z = 1;\n params.kernel.kernelParams = kernelArgs;\n kernelArgs[0] = &dPtr;\n kernelArgs[1] = &handle;\n checkCudaErrors(cudaGraphAddNode(&kernelNode, graph, NULL, NULL, 0, ¶ms));\n\n cudaGraphNodeParams cParams = {cudaGraphNodeTypeConditional};\n cParams.conditional.handle = handle;\n cParams.conditional.type = cudaGraphCondTypeSwitch;\n cParams.conditional.size = 4;\n checkCudaErrors(cudaGraphAddNode(&conditionalNode, graph, &kernelNode, NULL, 0, &cParams));\n\n // Populate the four graph bodies within the SWITCH conditional graph\n cudaGraphNode_t bodyNode;\n params.kernel.kernelParams = nullptr;\n params.kernel.func = (void *)switchGraphKernelC;\n checkCudaErrors(cudaGraphAddNode(&bodyNode, cParams.conditional.phGraph_out[0], NULL, NULL, 0, ¶ms));\n params.kernel.func = (void *)switchGraphKernelD;\n checkCudaErrors(cudaGraphAddNode(&bodyNode, cParams.conditional.phGraph_out[1], NULL, NULL, 0, ¶ms));\n params.kernel.func = (void *)switchGraphKernelE;\n checkCudaErrors(cudaGraphAddNode(&bodyNode, cParams.conditional.phGraph_out[2], NULL, NULL, 0, ¶ms));\n params.kernel.func = (void *)switchGraphKernelF;\n checkCudaErrors(cudaGraphAddNode(&bodyNode, cParams.conditional.phGraph_out[3], NULL, NULL, 0, ¶ms));\n\n checkCudaErrors(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n for (char i = 0; i < 5; i++) {\n // Initialize device memory and launch the graph\n checkCudaErrors(cudaMemset(dPtr, i, 1));\n printf(\"Host: Launching graph with device memory set to %d\\n\", i);\n checkCudaErrors(cudaGraphLaunch(graphExec, 0));\n checkCudaErrors(cudaDeviceSynchronize());\n }\n\n // Cleanup\n checkCudaErrors(cudaGraphExecDestroy(graphExec));\n checkCudaErrors(cudaGraphDestroy(graph));\n checkCudaErrors(cudaFree(dPtr));\n\n printf(\"simpleSwitchGraph: Complete\\n\\n\");\n}\n\nint main(int argc, char **argv)\n{\n int device = findCudaDevice(argc, (const char **)argv);\n\n int driverVersion = 0;\n\n cudaDriverGetVersion(&driverVersion);\n printf(\"Driver version is: %d.%d\\n\", driverVersion / 1000, (driverVersion % 100) / 10);\n\n if (driverVersion < 12030) {\n printf(\"Skipping execution as driver does not support Graph Conditional Nodes\\n\");\n return 0;\n }\n\n simpleIfGraph();\n simpleDoWhileGraph();\n capturedWhileGraph();\n\n if (driverVersion < 12080) {\n printf(\"Skipping execution as driver does not support if/else and switch type Graph Conditional Nodes\\n\");\n return 0;\n }\n\n simpleIfElseGraph();\n simpleSwitchGraph();\n\n return 0;\n}\n", "output": "/* Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This file demonstrates the usage of conditional graph nodes with\n * a series of *simple* example graphs.\n *\n * For more information on conditional nodes, see the programming guide:\n *\n * https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#conditional-graph-nodes\n *\n */\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n/*\n * Create a graph containing two nodes.\n * The first node, A, is a kernel and the second node, B, is a conditional IF node.\n * The kernel sets the condition variable to true if a device memory location\n * contains an odd number. Otherwise the condition variable is set to false.\n * There is a single kernel, C, within the conditional body which prints a message.\n *\n * A -> B [ C ]\n *\n */\n\n__global__ void ifGraphKernelA(char *dPtr, cudaGraphConditionalHandle handle)\n{\n // In this example, condition is set if *dPtr is odd\n unsigned int value = *dPtr & 0x01;\n cudaGraphSetConditional(handle, value);\n printf(\"GPU: Handle set to %d\\n\", value);\n}\n\n// This kernel will only be executed if the condition is true\n__global__ void ifGraphKernelC(void) { printf(\"GPU: Hello from the GPU! The condition was true.\\n\"); }\n\n// Setup and launch the graph\nvoid simpleIfGraph(void)\n{\n hipGraph_t graph;\n hipGraphExec_t graphExec;\n hipGraphNode_t kernelNode;\n hipGraphNode_t conditionalNode;\n\n void *kernelArgs[2];\n\n // Allocate a byte of device memory to use as input\n char *dPtr;\n checkCudaErrors(hipMalloc((void **)&dPtr, 1));\n\n printf(\"simpleIfGraph: Building graph...\\n\");\n hipGraphCreate(&graph, 0);\n\n // Create conditional handle.\n cudaGraphConditionalHandle handle;\n cudaGraphConditionalHandleCreate(&handle, graph);\n\n // Use a kernel upstream of the conditional to set the handle value\n hipGraphNodeParams params = {hipGraphNodeTypeKernel};\n params.kernel.func = (void *)ifGraphKernelA;\n params.kernel.blockDim.x = params.kernel.blockDim.y = params.kernel.blockDim.z = 1;\n params.kernel.gridDim.x = params.kernel.gridDim.y = params.kernel.gridDim.z = 1;\n params.kernel.kernelParams = kernelArgs;\n kernelArgs[0] = &dPtr;\n kernelArgs[1] = &handle;\n checkCudaErrors(hipGraphAddNode(&kernelNode, graph, NULL, NULL, 0, ¶ms));\n\n hipGraphNodeParams cParams = {hipGraphNodeTypeConditional};\n cParams.conditional.handle = handle;\n cParams.conditional.type = cudaGraphCondTypeIf;\n cParams.conditional.size = 1;\n checkCudaErrors(hipGraphAddNode(&conditionalNode, graph, &kernelNode, NULL, 0, &cParams));\n\n hipGraph_t bodyGraph = cParams.conditional.phGraph_out[0];\n\n // Populate the body of the conditional node\n hipGraphNode_t bodyNode;\n params.kernel.func = (void *)ifGraphKernelC;\n params.kernel.kernelParams = nullptr;\n checkCudaErrors(hipGraphAddNode(&bodyNode, bodyGraph, NULL, NULL, 0, ¶ms));\n\n checkCudaErrors(hipGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n // Initialize device memory and launch the graph\n checkCudaErrors(hipMemset(dPtr, 0, 1)); // Set dPtr to 0\n printf(\"Host: Launching graph with device memory set to 0\\n\");\n checkCudaErrors(hipGraphLaunch(graphExec, 0));\n checkCudaErrors(hipDeviceSynchronize());\n\n // Initialize device memory and launch the graph\n checkCudaErrors(hipMemset(dPtr, 1, 1)); // Set dPtr to 1\n printf(\"Host: Launching graph with device memory set to 1\\n\");\n checkCudaErrors(hipGraphLaunch(graphExec, 0));\n checkCudaErrors(hipDeviceSynchronize());\n\n // Cleanup\n checkCudaErrors(hipGraphExecDestroy(graphExec));\n checkCudaErrors(hipGraphDestroy(graph));\n checkCudaErrors(hipFree(dPtr));\n\n printf(\"simpleIfGraph: Complete\\n\\n\");\n}\n\n/*\n * Create a graph containing a single conditional while node.\n * The default value of the conditional variable is set to true, so this\n * effectively becomes a do-while loop as the conditional body will always\n * execute at least once. The body of the conditional contains 3 kernel nodes:\n * A [ B -> C -> D ]\n * Nodes B and C are just dummy nodes for demonstrative purposes. Node D\n * will decrement a device memory location and set the condition value to false\n * when the value reaches zero, terminating the loop.\n * In this example, stream capture is used to populate the conditional body.\n */\n\n// This kernel will only be executed if the condition is true\n__global__ void doWhileEmptyKernel(void)\n{\n printf(\"GPU: doWhileEmptyKernel()\\n\");\n return;\n}\n\n__global__ void doWhileLoopKernel(char *dPtr, cudaGraphConditionalHandle handle)\n{\n if (--(*dPtr) == 0) {\n cudaGraphSetConditional(handle, 0);\n }\n printf(\"GPU: counter = %d\\n\", *dPtr);\n}\n\nvoid simpleDoWhileGraph(void)\n{\n hipGraph_t graph;\n hipGraphExec_t graphExec;\n hipGraphNode_t conditionalNode;\n\n // Allocate a byte of device memory to use as input\n char *dPtr;\n checkCudaErrors(hipMalloc((void **)&dPtr, 1));\n\n printf(\"simpleDoWhileGraph: Building graph...\\n\");\n checkCudaErrors(hipGraphCreate(&graph, 0));\n\n cudaGraphConditionalHandle handle;\n checkCudaErrors(cudaGraphConditionalHandleCreate(&handle, graph, 1, cudaGraphCondAssignDefault));\n\n hipGraphNodeParams cParams = {hipGraphNodeTypeConditional};\n cParams.conditional.handle = handle;\n cParams.conditional.type = cudaGraphCondTypeWhile;\n cParams.conditional.size = 1;\n checkCudaErrors(hipGraphAddNode(&conditionalNode, graph, NULL, NULL, 0, &cParams));\n\n hipGraph_t bodyGraph = cParams.conditional.phGraph_out[0];\n\n hipStream_t captureStream;\n checkCudaErrors(hipStreamCreate(&captureStream));\n\n checkCudaErrors(\n hipStreamBeginCaptureToGraph(captureStream, bodyGraph, nullptr, nullptr, 0, hipStreamCaptureModeGlobal));\n doWhileEmptyKernel<<<1, 1, 0, captureStream>>>();\n doWhileEmptyKernel<<<1, 1, 0, captureStream>>>();\n doWhileLoopKernel<<<1, 1, 0, captureStream>>>(dPtr, handle);\n checkCudaErrors(hipStreamEndCapture(captureStream, nullptr));\n checkCudaErrors(hipStreamDestroy(captureStream));\n\n checkCudaErrors(hipGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n // Initialize device memory and launch the graph\n checkCudaErrors(hipMemset(dPtr, 10, 1)); // Set dPtr to 10\n printf(\"Host: Launching graph with loop counter set to 10\\n\");\n checkCudaErrors(hipGraphLaunch(graphExec, 0));\n checkCudaErrors(hipDeviceSynchronize());\n\n // Cleanup\n checkCudaErrors(hipGraphExecDestroy(graphExec));\n checkCudaErrors(hipGraphDestroy(graph));\n checkCudaErrors(hipFree(dPtr));\n\n printf(\"simpleDoWhileGraph: Complete\\n\\n\");\n}\n\n/*\n * Create a graph containing a conditional while loop using stream capture.\n * This demonstrates how to insert a conditional node into a stream which is\n * being captured. The graph consists of a kernel node, A, followed by a\n * conditional while node, B, followed by a kernel node, D. The conditional\n * body is populated by a single kernel node, C:\n *\n * A -> B [ C ] -> D\n *\n * The same kernel will be used for both nodes A and C. This kernel will test\n * a device memory location and set the condition when the location is non-zero.\n * We must run the kernel before the loop as well as inside the loop in order\n * to behave like a while loop as opposed to a do-while loop. We need to evaluate\n * the device memory location before the conditional node is evaluated in order\n * to set the condition variable properly. Because we're using a kernel upstream\n * of the conditional node, there is no need to use the handle default value to\n * initialize the conditional value.\n */\n\n__global__ void capturedWhileKernel(char *dPtr, cudaGraphConditionalHandle handle)\n{\n printf(\"GPU: counter = %d\\n\", *dPtr);\n if (*dPtr) {\n (*dPtr)--;\n }\n cudaGraphSetConditional(handle, *dPtr);\n}\n\n__global__ void capturedWhileEmptyKernel(void)\n{\n printf(\"GPU: capturedWhileEmptyKernel()\\n\");\n return;\n}\n\nvoid capturedWhileGraph(void)\n{\n hipGraph_t graph;\n hipGraphExec_t graphExec;\n\n hipStreamCaptureStatus status;\n const hipGraphNode_t *dependencies;\n size_t numDependencies;\n\n // Allocate a byte of device memory to use as input\n char *dPtr;\n checkCudaErrors(hipMalloc((void **)&dPtr, 1));\n\n printf(\"capturedWhileGraph: Building graph...\\n\");\n hipStream_t captureStream;\n checkCudaErrors(hipStreamCreate(&captureStream));\n\n checkCudaErrors(hipStreamBeginCapture(captureStream, hipStreamCaptureModeGlobal));\n\n // Obtain the handle of the graph\n checkCudaErrors(\n hipStreamGetCaptureInfo(captureStream, &status, NULL, &graph, &dependencies, NULL, &numDependencies));\n\n // Create the conditional handle\n cudaGraphConditionalHandle handle;\n checkCudaErrors(cudaGraphConditionalHandleCreate(&handle, graph));\n\n // Insert kernel node A\n capturedWhileKernel<<<1, 1, 0, captureStream>>>(dPtr, handle);\n\n // Obtain the handle for node A\n checkCudaErrors(\n hipStreamGetCaptureInfo(captureStream, &status, NULL, &graph, &dependencies, NULL, &numDependencies));\n\n // Insert conditional node B\n hipGraphNode_t conditionalNode;\n hipGraphNodeParams cParams = {hipGraphNodeTypeConditional};\n cParams.conditional.handle = handle;\n cParams.conditional.type = cudaGraphCondTypeWhile;\n cParams.conditional.size = 1;\n checkCudaErrors(hipGraphAddNode(&conditionalNode, graph, dependencies, NULL, numDependencies, &cParams));\n\n hipGraph_t bodyGraph = cParams.conditional.phGraph_out[0];\n\n // Update stream capture dependencies to account for the node we manually added\n checkCudaErrors(hipStreamUpdateCaptureDependencies(\n captureStream, &conditionalNode, NULL, 1, hipStreamSetCaptureDependencies));\n\n // Insert kernel node D\n capturedWhileEmptyKernel<<<1, 1, 0, captureStream>>>();\n\n checkCudaErrors(hipStreamEndCapture(captureStream, &graph));\n checkCudaErrors(hipStreamDestroy(captureStream));\n\n // Populate conditional body graph using stream capture\n hipStream_t bodyStream;\n checkCudaErrors(hipStreamCreate(&bodyStream));\n\n checkCudaErrors(\n hipStreamBeginCaptureToGraph(bodyStream, bodyGraph, nullptr, nullptr, 0, hipStreamCaptureModeGlobal));\n\n // Insert kernel node C\n capturedWhileKernel<<<1, 1, 0, bodyStream>>>(dPtr, handle);\n checkCudaErrors(hipStreamEndCapture(bodyStream, nullptr));\n checkCudaErrors(hipStreamDestroy(bodyStream));\n\n checkCudaErrors(hipGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n // Initialize device memory and launch the graph\n // Device memory is zero, so the conditional node will not execute\n checkCudaErrors(hipMemset(dPtr, 0, 1)); // Set dPtr to 0\n printf(\"Host: Launching graph with loop counter set to 0\\n\");\n checkCudaErrors(hipGraphLaunch(graphExec, 0));\n checkCudaErrors(hipDeviceSynchronize());\n\n // Initialize device memory and launch the graph\n checkCudaErrors(hipMemset(dPtr, 10, 1)); // Set dPtr to 10\n printf(\"Host: Launching graph with loop counter set to 10\\n\");\n checkCudaErrors(hipGraphLaunch(graphExec, 0));\n checkCudaErrors(hipDeviceSynchronize());\n\n // Cleanup\n checkCudaErrors(hipGraphExecDestroy(graphExec));\n checkCudaErrors(hipGraphDestroy(graph));\n checkCudaErrors(hipFree(dPtr));\n\n printf(\"capturedWhileGraph: Complete\\n\\n\");\n}\n\n/*\n * Create a graph containing two nodes.\n * The first node, A, is a kernel and the second node, B, is a conditional IF node containing two graphs.\n * The first graph within the conditional will be executed when the condition is true, while the second\n * graph will be executed when the conditional is false.\n * The kernel sets the condition variable to true if a device memory location\n * contains an odd number. Otherwise the condition variable is set to false.\n * There is a single kernel(C & D) within each conditional body which prints a message.\n *\n * A -> B [ C | D ]\n *\n * This example requires CUDA >= 12.8.\n */\n\n// This kernel will only be executed if the condition is false\n__global__ void ifGraphKernelD(void) { printf(\"GPU: Hello from the GPU! The condition was false.\\n\"); }\n\n// Setup and launch the graph\nvoid simpleIfElseGraph(void)\n{\n hipGraph_t graph;\n hipGraphExec_t graphExec;\n hipGraphNode_t kernelNode;\n hipGraphNode_t conditionalNode;\n\n void *kernelArgs[2];\n\n // Allocate a byte of device memory to use as input\n char *dPtr;\n checkCudaErrors(hipMalloc((void **)&dPtr, 1));\n\n printf(\"simpleIfElseGraph: Building graph...\\n\");\n hipGraphCreate(&graph, 0);\n\n // Create conditional handle.\n cudaGraphConditionalHandle handle;\n cudaGraphConditionalHandleCreate(&handle, graph);\n\n // Use a kernel upstream of the conditional to set the handle value\n hipGraphNodeParams params = {hipGraphNodeTypeKernel};\n params.kernel.func = (void *)ifGraphKernelA;\n params.kernel.blockDim.x = params.kernel.blockDim.y = params.kernel.blockDim.z = 1;\n params.kernel.gridDim.x = params.kernel.gridDim.y = params.kernel.gridDim.z = 1;\n params.kernel.kernelParams = kernelArgs;\n kernelArgs[0] = &dPtr;\n kernelArgs[1] = &handle;\n checkCudaErrors(hipGraphAddNode(&kernelNode, graph, NULL, NULL, 0, ¶ms));\n\n hipGraphNodeParams cParams = {hipGraphNodeTypeConditional};\n cParams.conditional.handle = handle;\n cParams.conditional.type = cudaGraphCondTypeIf;\n cParams.conditional.size = 2; // Set size to 2 to indicate an ELSE graph will be used\n checkCudaErrors(hipGraphAddNode(&conditionalNode, graph, &kernelNode, NULL, 0, &cParams));\n\n hipGraph_t bodyGraph = cParams.conditional.phGraph_out[0];\n\n // Populate the body of the first graph in the conditional node, executed if the condition is true\n hipGraphNode_t trueBodyNode;\n params.kernel.func = (void *)ifGraphKernelC;\n params.kernel.kernelParams = nullptr;\n checkCudaErrors(hipGraphAddNode(&trueBodyNode, bodyGraph, NULL, NULL, 0, ¶ms));\n\n // Populate the body of the second graph in the conditional node, executed if the condition is false\n bodyGraph = cParams.conditional.phGraph_out[1];\n\n hipGraphNode_t falseBodyNode;\n params.kernel.func = (void *)ifGraphKernelD;\n params.kernel.kernelParams = nullptr;\n checkCudaErrors(hipGraphAddNode(&falseBodyNode, bodyGraph, NULL, NULL, 0, ¶ms));\n\n checkCudaErrors(hipGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n // Initialize device memory and launch the graph\n checkCudaErrors(hipMemset(dPtr, 0, 1)); // Set dPtr to 0\n printf(\"Host: Launching graph with device memory set to 0\\n\");\n checkCudaErrors(hipGraphLaunch(graphExec, 0));\n checkCudaErrors(hipDeviceSynchronize());\n\n // Initialize device memory and launch the graph\n checkCudaErrors(hipMemset(dPtr, 1, 1)); // Set dPtr to 1\n printf(\"Host: Launching graph with device memory set to 1\\n\");\n checkCudaErrors(hipGraphLaunch(graphExec, 0));\n checkCudaErrors(hipDeviceSynchronize());\n\n // Cleanup\n checkCudaErrors(hipGraphExecDestroy(graphExec));\n checkCudaErrors(hipGraphDestroy(graph));\n checkCudaErrors(hipFree(dPtr));\n\n printf(\"simpleIfElseGraph: Complete\\n\\n\");\n}\n\n/*\n * Create a graph containing two nodes.\n * The first node, A, is a kernel and the second node, B, is a conditional SWITCH node containing four graphs.\n * The nth graph within the conditional will be executed when the condition is n. If conditional >= n,\n * no graph will be executed.\n * Kernel A sets the condition variable to the value stored in a device memory location.\n * This device location is updated from the host with each launch to demonstrate the behavior.\n * There is a single kernel(nodes C, D, E and F) within each conditional body which prints a message.\n *\n * A -> B [ C | D | E | F ]\n *\n * This example requires CUDA >= 12.8.\n */\n\n__global__ void switchGraphKernelA(char *dPtr, cudaGraphConditionalHandle handle)\n{\n unsigned int value = *dPtr;\n cudaGraphSetConditional(handle, value);\n printf(\"GPU: Handle set to %d\\n\", value);\n}\n\n__global__ void switchGraphKernelC(void) { printf(\"GPU: Hello from switchGraphKernelC(), running on the GPU!\\n\"); }\n\n__global__ void switchGraphKernelD(void) { printf(\"GPU: Hello from switchGraphKernelD(), running on the GPU!\\n\"); }\n\n__global__ void switchGraphKernelE(void) { printf(\"GPU: Hello from switchGraphKernelE(), running on the GPU!\\n\"); }\n\n__global__ void switchGraphKernelF(void) { printf(\"GPU: Hello from switchGraphKernelF(), running on the GPU!\\n\"); }\n\n// Setup and launch the graph\nvoid simpleSwitchGraph(void)\n{\n hipGraph_t graph;\n hipGraphExec_t graphExec;\n hipGraphNode_t kernelNode;\n hipGraphNode_t conditionalNode;\n\n void *kernelArgs[2];\n\n // Allocate a byte of device memory to use as input\n char *dPtr;\n checkCudaErrors(hipMalloc((void **)&dPtr, 1));\n\n printf(\"simpleSwitchGraph: Building graph...\\n\");\n hipGraphCreate(&graph, 0);\n\n // Create conditional handle.\n cudaGraphConditionalHandle handle;\n cudaGraphConditionalHandleCreate(&handle, graph);\n\n // Use a kernel upstream of the conditional to set the handle value\n hipGraphNodeParams params = {hipGraphNodeTypeKernel};\n params.kernel.func = (void *)switchGraphKernelA;\n params.kernel.blockDim.x = params.kernel.blockDim.y = params.kernel.blockDim.z = 1;\n params.kernel.gridDim.x = params.kernel.gridDim.y = params.kernel.gridDim.z = 1;\n params.kernel.kernelParams = kernelArgs;\n kernelArgs[0] = &dPtr;\n kernelArgs[1] = &handle;\n checkCudaErrors(hipGraphAddNode(&kernelNode, graph, NULL, NULL, 0, ¶ms));\n\n hipGraphNodeParams cParams = {hipGraphNodeTypeConditional};\n cParams.conditional.handle = handle;\n cParams.conditional.type = cudaGraphCondTypeSwitch;\n cParams.conditional.size = 4;\n checkCudaErrors(hipGraphAddNode(&conditionalNode, graph, &kernelNode, NULL, 0, &cParams));\n\n // Populate the four graph bodies within the SWITCH conditional graph\n hipGraphNode_t bodyNode;\n params.kernel.kernelParams = nullptr;\n params.kernel.func = (void *)switchGraphKernelC;\n checkCudaErrors(hipGraphAddNode(&bodyNode, cParams.conditional.phGraph_out[0], NULL, NULL, 0, ¶ms));\n params.kernel.func = (void *)switchGraphKernelD;\n checkCudaErrors(hipGraphAddNode(&bodyNode, cParams.conditional.phGraph_out[1], NULL, NULL, 0, ¶ms));\n params.kernel.func = (void *)switchGraphKernelE;\n checkCudaErrors(hipGraphAddNode(&bodyNode, cParams.conditional.phGraph_out[2], NULL, NULL, 0, ¶ms));\n params.kernel.func = (void *)switchGraphKernelF;\n checkCudaErrors(hipGraphAddNode(&bodyNode, cParams.conditional.phGraph_out[3], NULL, NULL, 0, ¶ms));\n\n checkCudaErrors(hipGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n for (char i = 0; i < 5; i++) {\n // Initialize device memory and launch the graph\n checkCudaErrors(hipMemset(dPtr, i, 1));\n printf(\"Host: Launching graph with device memory set to %d\\n\", i);\n checkCudaErrors(hipGraphLaunch(graphExec, 0));\n checkCudaErrors(hipDeviceSynchronize());\n }\n\n // Cleanup\n checkCudaErrors(hipGraphExecDestroy(graphExec));\n checkCudaErrors(hipGraphDestroy(graph));\n checkCudaErrors(hipFree(dPtr));\n\n printf(\"simpleSwitchGraph: Complete\\n\\n\");\n}\n\nint main(int argc, char **argv)\n{\n int device = findCudaDevice(argc, (const char **)argv);\n\n int driverVersion = 0;\n\n hipDriverGetVersion(&driverVersion);\n printf(\"Driver version is: %d.%d\\n\", driverVersion / 1000, (driverVersion % 100) / 10);\n\n if (driverVersion < 12030) {\n printf(\"Skipping execution as driver does not support Graph Conditional Nodes\\n\");\n return 0;\n }\n\n simpleIfGraph();\n simpleDoWhileGraph();\n capturedWhileGraph();\n\n if (driverVersion < 12080) {\n printf(\"Skipping execution as driver does not support if/else and switch type Graph Conditional Nodes\\n\");\n return 0;\n }\n\n simpleIfElseGraph();\n simpleSwitchGraph();\n\n return 0;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/graphConditionalNodes/graphConditionalNodes.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// CUDA sample demonstrating a __nv_bfloat16 (E8M7) GEMM computation using the Warp Matrix Multiply\n// and Accumulate API introduced in CUDA 11.0.\n\n// In this program, the compute_gemm kernel computes the result of a matrix multiplication\n// and addition: D = alpha * A * B + beta * C. The dimensions of both C and D matrices\n// are M_GLOBAL x N_GLOBAL. The A matrix is M_GLOBAL x K_GLOBAL (row-major), the B matrix\n// is K_GLOBAL x N_GLOBAL (column-major).\n// In that kernel, each CTA computes one 128 x 128 tile of the resulting matrix\n// per iteration. When the tile is computed, the CTA stores it to the global memory\n// and begins a new iteration, selecting a new 128 x 128 tile to compute.\n// Each CTA consists of eight warps. For the 128 x 128 tile, each warp computes eight\n// 16 x 16 subtiles, organized in a 2 x 4 two-dimensional array.\n// Warps compute the 16 x 16 subtiles using nvcuda::wmma::mma_sync operations by\n// moving through the K_GLOBAL dimension of the A and B matrices and accumulating\n// the intermediate result in the local thread state.\n\n// There are a number of simple optimizations used in the algorithm:\n// - The CTA copies the 128 x 128 tile of the C matrix from the global memory to\n// shared memory. After that is done, each warp loads the C matrix fragments from\n// shared memory, thus avoiding a random global memory access.\n// - On each internal iteration, the CTA copies a portion of the A and B matrices from\n// global memory to shared memory. After that, all warps in the CTA reuse the A and B\n// data from shared memory, thus reducing the number of data copies from global memory.\n// - The portions of the A and B matrices are stored in shared memory with an additional\n// padding (skew) to reduce the number of shared memory access bank conflicts.\n// (See a detailed explanation near the SKEW_BF16 macro definition.)\n// - When the CTA finishes computing the tiles of the resulting matrix, each warp stores\n// its subtiles to shared memory. The CTA then copies the shared memory contents to\n// global memory, again avoiding redundant random global memory accesses.\n// - Note that the CTA tile size is chosen to maximize the GPU register utilization,\n// but carefully enough to avoid local memory use.\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Externally configurable parameters.\n\n#ifndef CPU_DEBUG\n// Set this to 1 to verify the correctness of the GPU-computed matrix.\n#define CPU_DEBUG 0\n#endif\n\n#ifndef SHARED_MEMORY_LIMIT_64K\n// Set this to 0 to use more than 64 Kb of shared memory to cache data, to\n// improve the performance of the computations on GPU.\n// Note that you need a GPU that can have more than 64 Kb of shared memory\n// per multiprocessor.\n#define SHARED_MEMORY_LIMIT_64K 0\n#endif\n\n// GPU configuration.\n\n#define WARP_SIZE 32\n\n// MMA matrix tile dimensions.\n\n#define M 16\n#define N 16\n#define K 16\n\n// GEMM configuration.\n\n#define M_TILES 512\n#define N_TILES 512\n#define K_TILES 512\n\n#define M_GLOBAL (M * M_TILES)\n#define N_GLOBAL (N * N_TILES)\n#define K_GLOBAL (K * K_TILES)\n\n#define C_LAYOUT wmma::mem_row_major\n\n// Implementation constants.\n\n#define WARPS_PER_BLOCK 8\n#define THREADS_PER_BLOCK (WARP_SIZE * WARPS_PER_BLOCK)\n\n#if SHARED_MEMORY_LIMIT_64K\n// With only 64 Kb shared memory available, we can fit two 8-tile chunks of\n// the A and B matrix data, that is (M = 16) * (K = 16) * 8 * (CHUNK_K = 8)\n// * sizeof(__nv_bfloat16) = 32 Kb each.\n// (i.e. two 8x8 arrays of tiles of 16x16 __nv_bfloat16-typed elements per CTA).\n// But we cannot account the 8 Kb total skew overhead, without which the performance\n// would be severely impacted. So we choose to reduce the chunk size in half,\n// i.e. the amount of A and B matrix data we cache in shared memory.\n// Accordingly, this doubles the number of outer iterations across the global K\n// dimension, which only slightly impacts the performance.\n#define CHUNK_K 4\n#else\n#define CHUNK_K 8\n#endif\n\n#define CHUNK_LINE_BYTES (CHUNK_K * K * sizeof(__nv_bfloat16))\n#define WARP_COPY_BYTES (WARP_SIZE * sizeof(int4))\n#define CHUNK_COPY_LINES_PER_WARP (WARP_COPY_BYTES / CHUNK_LINE_BYTES)\n#define CHUNK_COPY_LINE_LANES (WARP_SIZE / CHUNK_COPY_LINES_PER_WARP)\n\n#define BLOCK_ROW_WARPS 2\n#define BLOCK_COL_WARPS 4\n\n#define WARP_ROW_TILES 4\n#define WARP_COL_TILES 2\n\n#define BLOCK_ROW_TILES (WARP_ROW_TILES * BLOCK_ROW_WARPS)\n#define BLOCK_COL_TILES (WARP_COL_TILES * BLOCK_COL_WARPS)\n\n#define GLOBAL_MEM_STRIDE N_GLOBAL\n\n#define SHMEM_STRIDE (N * BLOCK_ROW_TILES)\n#define SHMEM_OFFSET (N * WARP_ROW_TILES)\n\n// The macro below is used to shift rows of the A matrix and columns of the B matrix\n// in shared memory to minimize possible bank conflicts.\n// Before performing the nvcuda::wmma::mma_sync operation, the warp must load the matrix\n// data using the nvcuda::wmma::load_matrix_sync operation. Although the memory access pattern\n// is not specified for that function, each lane in the warp can read one or multiple matrix\n// elements from different matrix rows or columns.\n// For shared memory, such access can result in bank conflicts if different rows / columns\n// of the matrix map to the same bank. By shifting each row and column by a few bytes, we\n// make sure that they map to different banks, thus reducing the number of possible bank\n// conflicts.\n// The number of 16 two-byte \"__nv_bfloat16\" elements is chosen as the minimum possible shift because\n// we must keep each row and column 256-bit aligned, as required by nvcuda::wmma::load_matrix_sync.\n#define SKEW_BF16 16\n\n#define checkKernelErrors(expr) \\\n do { \\\n expr; \\\n \\\n cudaError_t __err = cudaGetLastError(); \\\n if (__err != cudaSuccess) { \\\n printf(\"Line %d: '%s' failed: %s\\n\", __LINE__, #expr, cudaGetErrorString(__err)); \\\n abort(); \\\n } \\\n } while (0)\n\nenum kernels {\n bf16mma_shmem_gemm_async_copy = 0, // __nv_bfloat16 MMA shmem using kernel with async_copy\n bf16mma_shmem_gemm = 1, // __nv_bfloat16 MMA shmem using kernel normal copy (without async_copy).\n simple_bf16mma_gemm = 2 // __nv_bfloat16 MMA non-shmem using simple kernel.\n};\n\nconst char *kernelNames[] = {\"compute_bf16gemm_async_copy\", \"compute_bf16gemm\", \"simple_wmma_bf16gemm\"};\n\nusing namespace nvcuda;\n\n__host__ void init_host_matrices(__nv_bfloat16 *a, __nv_bfloat16 *b, float *c)\n{\n for (int i = 0; i < M_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n a[i * K_GLOBAL + j] = (__nv_bfloat16)(float)(rand() % 3);\n }\n }\n\n for (int i = 0; i < N_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n b[i * K_GLOBAL + j] = (__nv_bfloat16)(float)(rand() % 3);\n }\n }\n\n for (int t = 0; t < M_GLOBAL * N_GLOBAL; t++) {\n c[t] = (float)(rand() % 3);\n }\n}\n\n__global__ void\ncompute_bf16gemm(const __nv_bfloat16 *A, const __nv_bfloat16 *B, const float *C, float *D, float alpha, float beta)\n{\n#if __CUDA_ARCH__ >= 800\n extern __shared__ __nv_bfloat16 shmem[][CHUNK_K * K + SKEW_BF16];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n const size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n float *shmem_warp_tile_ptr = (float *)&shmem[0][0] + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS\n + (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.\n float *shmem_warp_stream_ptr = (float *)&shmem[0][0] + warpId * SHMEM_STRIDE * N;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may result\n // in a loss of precision). Zero still needs to be specially handled though.\n beta /= alpha;\n\n // Each CTA slides along the 128 x 128 tiles from the top left corner of the matrix to the\n // right and down, and selects the next tile to compute. Once there's no such tile,\n // all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const float *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId) =\n *((int4 *)(src_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment multiplications\n // along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Scale the C matrix.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const __nv_bfloat16 *warp_ptr =\n (warpId < (WARPS_PER_BLOCK / 2))\n ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);\n\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.\n size_t shmem_idx = warpId < (WARPS_PER_BLOCK / 2)\n ? (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (N * (warpId % (WARPS_PER_BLOCK / 2)) * 2 + shmem_idx_b_off);\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n const __nv_bfloat16 *lane_ptr = (warp_ptr + tile_k * K + (laneId / CHUNK_COPY_LINE_LANES) * K_GLOBAL);\n\n // Shift the second half of the warp to the next row / column in the shared memory.\n shmem_idx += laneId / CHUNK_COPY_LINE_LANES;\n\n#pragma unroll\n for (int i = 0; i < ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP) * 2; i++) {\n // Copy 16 bytes at once in each lane.\n *((int4 *)&shmem[shmem_idx][0] + (laneId % CHUNK_COPY_LINE_LANES)) =\n *((int4 *)lane_ptr + (laneId % CHUNK_COPY_LINE_LANES));\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / BLOCK_ROW_WARPS) * M * BLOCK_ROW_WARPS + (i * M);\n const __nv_bfloat16 *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_BF16);\n\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be reused\n // against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const __nv_bfloat16 *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_BF16);\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n // Uniform, point-wise transformations of ALL fragment elements by ALL threads in the\n // warp are well-defined even though element indices within fragment storage are not defined.\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * K + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global memory.\n float *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((float4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((float4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n#endif\n}\n\n__global__ void compute_bf16gemm_async_copy(const __nv_bfloat16 *A,\n const __nv_bfloat16 *B,\n const float *C,\n float *D,\n float alpha,\n float beta)\n{\n#if __CUDA_ARCH__ >= 800\n extern __shared__ __nv_bfloat16 shmem[][CHUNK_K * K + SKEW_BF16];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n constexpr size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n float *shmem_warp_tile_ptr = (float *)&shmem[0][0] + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS\n + (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.\n float *shmem_warp_stream_ptr = (float *)&shmem[0][0] + warpId * SHMEM_STRIDE * N;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may result\n // in a loss of precision). Zero still needs to be specially handled though.\n beta /= alpha;\n\n cuda::pipeline pipe = cuda::make_pipeline();\n const auto shape4 = cuda::aligned_size_t(sizeof(float4));\n constexpr int loadStride = 2; // load 4 floats, left-shift by 2.\n\n // Each CTA slides along the 128 x 128 tiles from the top left corner of the matrix to the\n // right and down, and selects the next tile to compute. Once there's no such tile,\n // all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const float *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < N; i++) {\n pipe.producer_acquire();\n\n cuda::memcpy_async(&shmem_warp_stream_ptr[(SHMEM_STRIDE * i) + (laneId << loadStride)],\n &src_gmem_warp_stream_ptr[(GLOBAL_MEM_STRIDE * i) + (laneId << loadStride)],\n shape4,\n pipe);\n\n pipe.producer_commit();\n }\n\n // Now wait for all the above issued 8 batches to complete.\n cuda::pipeline_consumer_wait_prior<0>(pipe);\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment multiplications\n // along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n // Scale the C matrix.\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n pipe.consumer_release();\n\n // sync here so that shared memory can then be used for loading A & B matrices.\n __syncthreads();\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const __nv_bfloat16 *warp_ptr =\n (warpId < (WARPS_PER_BLOCK / 2))\n ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);\n\n constexpr int chunksPerLane = ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP) * 2;\n constexpr int loadStrideBfloat8 = 3; // load 8 bfloats, left-shift by 3.\n const int laneLoadElem = (laneId % CHUNK_COPY_LINE_LANES) << loadStrideBfloat8;\n const int stridePerLaneCopy = (laneId / CHUNK_COPY_LINE_LANES);\n\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.\n // As for bf16 MMA M == N we use M for warp 4-7 + shmem_idx_b_off.\n size_t shmem_idx =\n (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2) + ((warpId / (WARPS_PER_BLOCK / 2)) * shmem_idx_b_off);\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n const __nv_bfloat16 *lane_ptr = warp_ptr + tile_k * K + stridePerLaneCopy * K_GLOBAL + laneLoadElem;\n\n // Shift the second half of the warp to the next row / column in the shared memory.\n shmem_idx += stridePerLaneCopy;\n\n#pragma unroll\n for (int i = 0; i < chunksPerLane; i++) {\n // Copy 16 bytes at once in each lane.\n pipe.producer_acquire();\n cuda::memcpy_async(&shmem[shmem_idx][laneLoadElem], lane_ptr, shape4, pipe);\n pipe.producer_commit();\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n cuda::pipeline_consumer_wait_prior<0>(pipe);\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / BLOCK_ROW_WARPS) * M * BLOCK_ROW_WARPS + (i * M);\n const __nv_bfloat16 *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_BF16);\n\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be reused\n // against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const __nv_bfloat16 *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_BF16);\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n\n pipe.consumer_release();\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n // Uniform, point-wise transformations of ALL fragment elements by ALL threads in the\n // warp are well-defined even though element indices within fragment storage are not defined.\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * K + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global memory.\n float *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n#endif\n}\n\n// Performs an MxNxK bf16 GEMM (C=alpha*A*B + beta*C) assuming:\n// 1) Matrices are packed in memory.\n// 2) M, N and K are multiples of 16, 16 and 16 respectively.\n// 3) A is row major, B is column major matrix.\n// Note: This is a less performant version of the compute_bf16gemm kernel. It is designed for\n// demonstration purposes only to show the CUDA WMMA API use without relying on\n// availability of the shared memory.\n__global__ void simple_wmma_bf16gemm(__nv_bfloat16 *a,\n __nv_bfloat16 *b,\n float *c,\n float *d,\n int m_ld,\n int n_ld,\n int k_ld,\n float alpha,\n float beta)\n{\n#if __CUDA_ARCH__ >= 800\n // Leading dimensions. Packed with no transpositions.\n int lda = k_ld;\n int ldb = k_ld;\n int ldc = n_ld;\n\n // Tile using a 2D grid\n int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;\n int warpN = (blockIdx.y * blockDim.y + threadIdx.y);\n\n // Declare the fragments\n wmma::fragment a_frag;\n wmma::fragment b_frag;\n wmma::fragment acc_frag;\n wmma::fragment c_frag;\n\n wmma::fill_fragment(acc_frag, 0.0f);\n\n // Loop over k\n for (int i = 0; i < k_ld; i += K) {\n int aCol = i;\n int aRow = warpM * M;\n\n int bCol = i;\n int bRow = warpN * N;\n\n // Bounds checking\n if (aRow < m_ld && aCol < k_ld && bRow < k_ld && bCol < n_ld) {\n // Load the inputs\n wmma::load_matrix_sync(a_frag, a + aCol + aRow * lda, lda);\n wmma::load_matrix_sync(b_frag, b + bRow + bCol * ldb, ldb);\n\n // Perform the matrix multiplication\n wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);\n }\n }\n\n // Load in the current value of c, scale it by beta, and add this our result scaled by alpha\n int cCol = warpN * N;\n int cRow = warpM * M;\n\n if (cRow < m_ld && cCol < n_ld) {\n wmma::load_matrix_sync(c_frag, c + cCol + cRow * ldc, ldc, wmma::mem_row_major);\n\n for (int i = 0; i < c_frag.num_elements; i++) {\n c_frag.x[i] = alpha * acc_frag.x[i] + beta * c_frag.x[i];\n }\n\n // Store the output\n wmma::store_matrix_sync(d + cCol + cRow * ldc, c_frag, ldc, wmma::mem_row_major);\n }\n#endif\n}\n\n__host__ void matMultiplyOnHost(__nv_bfloat16 *A,\n __nv_bfloat16 *B,\n float *C,\n float alpha,\n float beta,\n int numARows,\n int numAColumns,\n int numBRows,\n int numBColumns,\n int numCRows,\n int numCColumns)\n{\n for (int i = 0; i < numCRows; i++) {\n for (int j = 0; j < numCColumns; j++) {\n float temp = 0.0;\n\n for (int k = 0; k < numAColumns; k++) {\n temp += (float)A[i * numAColumns + k] * (float)B[j * numBRows + k];\n }\n\n C[i * numCColumns + j] = temp * alpha + beta * C[i * numCColumns + j];\n }\n }\n}\n\nint main(int argc, char **argv)\n{\n printf(\"Initializing...\\n\");\n\n int dev = findCudaDevice(argc, (const char **)argv);\n\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, dev));\n\n // Tensor cores require a GPU of Volta (SM8X) architecture or higher.\n if (deviceProp.major < 8) {\n printf(\"bf16TensorCoreGemm requires requires SM 8.0 or higher to use Tensor Cores. Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n printf(\"M: %d (%d x %d)\\n\", M_GLOBAL, M, M_TILES);\n printf(\"N: %d (%d x %d)\\n\", N_GLOBAL, N, N_TILES);\n printf(\"K: %d (%d x %d)\\n\", K_GLOBAL, K, K_TILES);\n\n __nv_bfloat16 *A_h = NULL;\n __nv_bfloat16 *B_h = NULL;\n float *C_h = NULL;\n#if CPU_DEBUG\n float *result_hD = NULL;\n float *result_host = NULL;\n#endif\n\n A_h = (__nv_bfloat16 *)malloc(sizeof(__nv_bfloat16) * M_GLOBAL * K_GLOBAL);\n B_h = (__nv_bfloat16 *)malloc(sizeof(__nv_bfloat16) * K_GLOBAL * N_GLOBAL);\n C_h = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n#if CPU_DEBUG\n result_hD = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n result_host = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n#endif\n\n __nv_bfloat16 *A = NULL;\n __nv_bfloat16 *B = NULL;\n float *C = NULL;\n float *D = NULL;\n\n checkCudaErrors(cudaMalloc((void **)&A, sizeof(__nv_bfloat16) * M_GLOBAL * K_GLOBAL));\n checkCudaErrors(cudaMalloc((void **)&B, sizeof(__nv_bfloat16) * N_GLOBAL * K_GLOBAL));\n checkCudaErrors(cudaMalloc((void **)&C, sizeof(float) * M_GLOBAL * N_GLOBAL));\n checkCudaErrors(cudaMalloc((void **)&D, sizeof(float) * M_GLOBAL * N_GLOBAL));\n\n assert(((unsigned long long)A) % 128 == 0);\n assert(((unsigned long long)B) % 128 == 0);\n assert(((unsigned long long)C) % 128 == 0);\n assert(((unsigned long long)D) % 128 == 0);\n\n init_host_matrices(A_h, B_h, C_h);\n\n printf(\"Preparing data for GPU...\\n\");\n\n checkCudaErrors(cudaMemcpy(A, A_h, sizeof(__nv_bfloat16) * M_GLOBAL * K_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(B, B_h, sizeof(__nv_bfloat16) * N_GLOBAL * K_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(C, C_h, sizeof(float) * M_GLOBAL * N_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemset(D, 0, sizeof(float) * M_GLOBAL * N_GLOBAL));\n\n enum {\n // Compute the right amount of shared memory to request.\n // We need shared memory to hold per-CTA C and D matrix tiles, and to cache per-CTA chunks\n // of the A and B matrices. Therefore, the right amount to request is the maximum of those\n // two numbers.\n SHMEM_SZ = MAX(sizeof(__nv_bfloat16) * (BLOCK_COL_TILES * M) * (CHUNK_K * K + SKEW_BF16) * 2,\n M * (BLOCK_ROW_WARPS * WARP_ROW_TILES) * N * (BLOCK_COL_WARPS * WARP_COL_TILES) * sizeof(float))\n };\n\n printf(\"Required shared memory size: %lu Kb\\n\", SHMEM_SZ / 1024UL);\n\n const float alpha = 1.1f;\n const float beta = 1.2f;\n\n cudaEvent_t start, stop;\n\n checkCudaErrors(cudaEventCreate(&start));\n checkCudaErrors(cudaEventCreate(&stop));\n checkCudaErrors(cudaEventRecord(start));\n\n // kernel to run - default (b16mma_shmem_gemm_async_copy == 0)\n kernels selected_kernel = bf16mma_shmem_gemm_async_copy;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"kernel\")) {\n int kernel_number = getCmdLineArgumentInt(argc, (const char **)argv, \"kernel\");\n if (kernel_number < 3) {\n selected_kernel = (kernels)kernel_number;\n }\n else {\n printf(\"Error: kernel number should be between 0 to 2, you have entered %d\\n\", kernel_number);\n exit(EXIT_FAILURE);\n }\n }\n\n // If enough shared memory available on the GPU use high performant kernel\n if ((deviceProp.sharedMemPerMultiprocessor >= SHMEM_SZ) && (selected_kernel != simple_bf16mma_gemm)) {\n printf(\"Computing using high performance kernel = %d - %s\\n\", selected_kernel, kernelNames[selected_kernel]);\n\n switch (selected_kernel) {\n case bf16mma_shmem_gemm_async_copy:\n default:\n checkCudaErrors(cudaFuncSetAttribute(\n compute_bf16gemm_async_copy, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors(\n (compute_bf16gemm_async_copy<<>>(\n A, B, C, D, alpha, beta)));\n break;\n case bf16mma_shmem_gemm:\n checkCudaErrors(\n cudaFuncSetAttribute(compute_bf16gemm, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors((compute_bf16gemm<<>>(\n A, B, C, D, alpha, beta)));\n break;\n }\n#if CPU_DEBUG\n checkCudaErrors(cudaMemcpy(result_hD, D, sizeof(float) * M_GLOBAL * N_GLOBAL, cudaMemcpyDeviceToHost));\n#endif\n }\n else {\n dim3 gridDim;\n dim3 blockDim;\n\n // blockDim.x must be a multple of warpSize\n // 128x4 means we have 16 warps and a block computes a 64x64 output tile\n blockDim.x = 128;\n blockDim.y = 4;\n\n gridDim.x = (M_GLOBAL + (M * blockDim.x / 32 - 1)) / (M * blockDim.x / 32);\n gridDim.y = (N_GLOBAL + N * blockDim.y - 1) / (N * blockDim.y);\n\n printf(\"Computing... using simple_wmma_gemm kernel\\n\");\n simple_wmma_bf16gemm<<>>(A, B, C, D, M_GLOBAL, N_GLOBAL, K_GLOBAL, alpha, beta);\n#if CPU_DEBUG\n checkCudaErrors(cudaMemcpy(result_hD, D, sizeof(float) * M_GLOBAL * N_GLOBAL, cudaMemcpyDeviceToHost));\n#endif\n }\n\n checkCudaErrors(cudaEventRecord(stop));\n checkCudaErrors(cudaEventSynchronize(stop));\n\n#if CPU_DEBUG\n printf(\"Verifying correctness of the computations...\\n\");\n\n memcpy(result_host, C_h, sizeof(float) * M_GLOBAL * N_GLOBAL);\n\n matMultiplyOnHost(A_h, B_h, result_host, alpha, beta, M_GLOBAL, K_GLOBAL, K_GLOBAL, N_GLOBAL, M_GLOBAL, N_GLOBAL);\n\n for (int i = 0; i < N_GLOBAL * M_GLOBAL; i++) {\n if (fabs(result_hD[i] - result_host[i]) > 0.1f) {\n printf(\"mismatch i=%d result_hD=%f result_host=%f\\n\", i, result_hD[i], result_host[i]);\n }\n }\n free(result_hD);\n free(result_host);\n#endif\n\n float milliseconds = 0;\n\n checkCudaErrors(cudaEventElapsedTime(&milliseconds, start, stop));\n\n printf(\"Time: %f ms\\n\", milliseconds);\n printf(\"TFLOPS: %.2f\\n\", (((double)M_GLOBAL * N_GLOBAL * K_GLOBAL * 2) / (milliseconds / 1000.)) / 1e12);\n\n free(A_h);\n free(B_h);\n free(C_h);\n checkCudaErrors(cudaFree((void *)A));\n checkCudaErrors(cudaFree((void *)B));\n checkCudaErrors(cudaFree((void *)C));\n checkCudaErrors(cudaFree((void *)D));\n\n return 0;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// CUDA sample demonstrating a __nv_bfloat16 (E8M7) GEMM computation using the Warp Matrix Multiply\n// and Accumulate API introduced in CUDA 11.0.\n\n// In this program, the compute_gemm kernel computes the result of a matrix multiplication\n// and addition: D = alpha * A * B + beta * C. The dimensions of both C and D matrices\n// are M_GLOBAL x N_GLOBAL. The A matrix is M_GLOBAL x K_GLOBAL (row-major), the B matrix\n// is K_GLOBAL x N_GLOBAL (column-major).\n// In that kernel, each CTA computes one 128 x 128 tile of the resulting matrix\n// per iteration. When the tile is computed, the CTA stores it to the global memory\n// and begins a new iteration, selecting a new 128 x 128 tile to compute.\n// Each CTA consists of eight warps. For the 128 x 128 tile, each warp computes eight\n// 16 x 16 subtiles, organized in a 2 x 4 two-dimensional array.\n// Warps compute the 16 x 16 subtiles using nvcuda::wmma::mma_sync operations by\n// moving through the K_GLOBAL dimension of the A and B matrices and accumulating\n// the intermediate result in the local thread state.\n\n// There are a number of simple optimizations used in the algorithm:\n// - The CTA copies the 128 x 128 tile of the C matrix from the global memory to\n// shared memory. After that is done, each warp loads the C matrix fragments from\n// shared memory, thus avoiding a random global memory access.\n// - On each internal iteration, the CTA copies a portion of the A and B matrices from\n// global memory to shared memory. After that, all warps in the CTA reuse the A and B\n// data from shared memory, thus reducing the number of data copies from global memory.\n// - The portions of the A and B matrices are stored in shared memory with an additional\n// padding (skew) to reduce the number of shared memory access bank conflicts.\n// (See a detailed explanation near the SKEW_BF16 macro definition.)\n// - When the CTA finishes computing the tiles of the resulting matrix, each warp stores\n// its subtiles to shared memory. The CTA then copies the shared memory contents to\n// global memory, again avoiding redundant random global memory accesses.\n// - Note that the CTA tile size is chosen to maximize the GPU register utilization,\n// but carefully enough to avoid local memory use.\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Externally configurable parameters.\n\n#ifndef CPU_DEBUG\n// Set this to 1 to verify the correctness of the GPU-computed matrix.\n#define CPU_DEBUG 0\n#endif\n\n#ifndef SHARED_MEMORY_LIMIT_64K\n// Set this to 0 to use more than 64 Kb of shared memory to cache data, to\n// improve the performance of the computations on GPU.\n// Note that you need a GPU that can have more than 64 Kb of shared memory\n// per multiprocessor.\n#define SHARED_MEMORY_LIMIT_64K 0\n#endif\n\n// GPU configuration.\n\n#define WARP_SIZE 32\n\n// MMA matrix tile dimensions.\n\n#define M 16\n#define N 16\n#define K 16\n\n// GEMM configuration.\n\n#define M_TILES 512\n#define N_TILES 512\n#define K_TILES 512\n\n#define M_GLOBAL (M * M_TILES)\n#define N_GLOBAL (N * N_TILES)\n#define K_GLOBAL (K * K_TILES)\n\n#define C_LAYOUT wmma::mem_row_major\n\n// Implementation constants.\n\n#define WARPS_PER_BLOCK 8\n#define THREADS_PER_BLOCK (WARP_SIZE * WARPS_PER_BLOCK)\n\n#if SHARED_MEMORY_LIMIT_64K\n// With only 64 Kb shared memory available, we can fit two 8-tile chunks of\n// the A and B matrix data, that is (M = 16) * (K = 16) * 8 * (CHUNK_K = 8)\n// * sizeof(__nv_bfloat16) = 32 Kb each.\n// (i.e. two 8x8 arrays of tiles of 16x16 __nv_bfloat16-typed elements per CTA).\n// But we cannot account the 8 Kb total skew overhead, without which the performance\n// would be severely impacted. So we choose to reduce the chunk size in half,\n// i.e. the amount of A and B matrix data we cache in shared memory.\n// Accordingly, this doubles the number of outer iterations across the global K\n// dimension, which only slightly impacts the performance.\n#define CHUNK_K 4\n#else\n#define CHUNK_K 8\n#endif\n\n#define CHUNK_LINE_BYTES (CHUNK_K * K * sizeof(__hip_bfloat16))\n#define WARP_COPY_BYTES (WARP_SIZE * sizeof(int4))\n#define CHUNK_COPY_LINES_PER_WARP (WARP_COPY_BYTES / CHUNK_LINE_BYTES)\n#define CHUNK_COPY_LINE_LANES (WARP_SIZE / CHUNK_COPY_LINES_PER_WARP)\n\n#define BLOCK_ROW_WARPS 2\n#define BLOCK_COL_WARPS 4\n\n#define WARP_ROW_TILES 4\n#define WARP_COL_TILES 2\n\n#define BLOCK_ROW_TILES (WARP_ROW_TILES * BLOCK_ROW_WARPS)\n#define BLOCK_COL_TILES (WARP_COL_TILES * BLOCK_COL_WARPS)\n\n#define GLOBAL_MEM_STRIDE N_GLOBAL\n\n#define SHMEM_STRIDE (N * BLOCK_ROW_TILES)\n#define SHMEM_OFFSET (N * WARP_ROW_TILES)\n\n// The macro below is used to shift rows of the A matrix and columns of the B matrix\n// in shared memory to minimize possible bank conflicts.\n// Before performing the nvcuda::wmma::mma_sync operation, the warp must load the matrix\n// data using the nvcuda::wmma::load_matrix_sync operation. Although the memory access pattern\n// is not specified for that function, each lane in the warp can read one or multiple matrix\n// elements from different matrix rows or columns.\n// For shared memory, such access can result in bank conflicts if different rows / columns\n// of the matrix map to the same bank. By shifting each row and column by a few bytes, we\n// make sure that they map to different banks, thus reducing the number of possible bank\n// conflicts.\n// The number of 16 two-byte \"__nv_bfloat16\" elements is chosen as the minimum possible shift because\n// we must keep each row and column 256-bit aligned, as required by nvcuda::wmma::load_matrix_sync.\n#define SKEW_BF16 16\n\n#define checkKernelErrors(expr) \\\n do { \\\n expr; \\\n \\\n hipError_t __err = hipGetLastError(); \\\n if (__err != hipSuccess) { \\\n printf(\"Line %d: '%s' failed: %s\\n\", __LINE__, #expr, hipGetErrorString(__err)); \\\n abort(); \\\n } \\\n } while (0)\n\nenum kernels {\n bf16mma_shmem_gemm_async_copy = 0, // __nv_bfloat16 MMA shmem using kernel with async_copy\n bf16mma_shmem_gemm = 1, // __nv_bfloat16 MMA shmem using kernel normal copy (without async_copy).\n simple_bf16mma_gemm = 2 // __nv_bfloat16 MMA non-shmem using simple kernel.\n};\n\nconst char *kernelNames[] = {\"compute_bf16gemm_async_copy\", \"compute_bf16gemm\", \"simple_wmma_bf16gemm\"};\n\nusing namespace nvcuda;\n\n__host__ void init_host_matrices(__hip_bfloat16 *a, __hip_bfloat16 *b, float *c)\n{\n for (int i = 0; i < M_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n a[i * K_GLOBAL + j] = (__hip_bfloat16)(float)(rand() % 3);\n }\n }\n\n for (int i = 0; i < N_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n b[i * K_GLOBAL + j] = (__hip_bfloat16)(float)(rand() % 3);\n }\n }\n\n for (int t = 0; t < M_GLOBAL * N_GLOBAL; t++) {\n c[t] = (float)(rand() % 3);\n }\n}\n\n__global__ void\ncompute_bf16gemm(const __hip_bfloat16 *A, const __hip_bfloat16 *B, const float *C, float *D, float alpha, float beta)\n{\n#if __CUDA_ARCH__ >= 800\n extern __shared__ __hip_bfloat16 shmem[][CHUNK_K * K + SKEW_BF16];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n const size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n float *shmem_warp_tile_ptr = (float *)&shmem[0][0] + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS\n + (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.\n float *shmem_warp_stream_ptr = (float *)&shmem[0][0] + warpId * SHMEM_STRIDE * N;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may result\n // in a loss of precision). Zero still needs to be specially handled though.\n beta /= alpha;\n\n // Each CTA slides along the 128 x 128 tiles from the top left corner of the matrix to the\n // right and down, and selects the next tile to compute. Once there's no such tile,\n // all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const float *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId) =\n *((int4 *)(src_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment multiplications\n // along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Scale the C matrix.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const __hip_bfloat16 *warp_ptr =\n (warpId < (WARPS_PER_BLOCK / 2))\n ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);\n\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.\n size_t shmem_idx = warpId < (WARPS_PER_BLOCK / 2)\n ? (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (N * (warpId % (WARPS_PER_BLOCK / 2)) * 2 + shmem_idx_b_off);\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n const __hip_bfloat16 *lane_ptr = (warp_ptr + tile_k * K + (laneId / CHUNK_COPY_LINE_LANES) * K_GLOBAL);\n\n // Shift the second half of the warp to the next row / column in the shared memory.\n shmem_idx += laneId / CHUNK_COPY_LINE_LANES;\n\n#pragma unroll\n for (int i = 0; i < ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP) * 2; i++) {\n // Copy 16 bytes at once in each lane.\n *((int4 *)&shmem[shmem_idx][0] + (laneId % CHUNK_COPY_LINE_LANES)) =\n *((int4 *)lane_ptr + (laneId % CHUNK_COPY_LINE_LANES));\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / BLOCK_ROW_WARPS) * M * BLOCK_ROW_WARPS + (i * M);\n const __hip_bfloat16 *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_BF16);\n\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be reused\n // against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const __hip_bfloat16 *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_BF16);\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n // Uniform, point-wise transformations of ALL fragment elements by ALL threads in the\n // warp are well-defined even though element indices within fragment storage are not defined.\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * K + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global memory.\n float *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((float4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((float4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n#endif\n}\n\n__global__ void compute_bf16gemm_async_copy(const __hip_bfloat16 *A,\n const __hip_bfloat16 *B,\n const float *C,\n float *D,\n float alpha,\n float beta)\n{\n#if __CUDA_ARCH__ >= 800\n extern __shared__ __hip_bfloat16 shmem[][CHUNK_K * K + SKEW_BF16];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n constexpr size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n float *shmem_warp_tile_ptr = (float *)&shmem[0][0] + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS\n + (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.\n float *shmem_warp_stream_ptr = (float *)&shmem[0][0] + warpId * SHMEM_STRIDE * N;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may result\n // in a loss of precision). Zero still needs to be specially handled though.\n beta /= alpha;\n\n cuda::pipeline pipe = cuda::make_pipeline();\n const auto shape4 = cuda::aligned_size_t(sizeof(float4));\n constexpr int loadStride = 2; // load 4 floats, left-shift by 2.\n\n // Each CTA slides along the 128 x 128 tiles from the top left corner of the matrix to the\n // right and down, and selects the next tile to compute. Once there's no such tile,\n // all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const float *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < N; i++) {\n pipe.producer_acquire();\n\n cuda::memcpy_async(&shmem_warp_stream_ptr[(SHMEM_STRIDE * i) + (laneId << loadStride)],\n &src_gmem_warp_stream_ptr[(GLOBAL_MEM_STRIDE * i) + (laneId << loadStride)],\n shape4,\n pipe);\n\n pipe.producer_commit();\n }\n\n // Now wait for all the above issued 8 batches to complete.\n cuda::pipeline_consumer_wait_prior<0>(pipe);\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment multiplications\n // along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n // Scale the C matrix.\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n pipe.consumer_release();\n\n // sync here so that shared memory can then be used for loading A & B matrices.\n __syncthreads();\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const __hip_bfloat16 *warp_ptr =\n (warpId < (WARPS_PER_BLOCK / 2))\n ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);\n\n constexpr int chunksPerLane = ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP) * 2;\n constexpr int loadStrideBfloat8 = 3; // load 8 bfloats, left-shift by 3.\n const int laneLoadElem = (laneId % CHUNK_COPY_LINE_LANES) << loadStrideBfloat8;\n const int stridePerLaneCopy = (laneId / CHUNK_COPY_LINE_LANES);\n\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.\n // As for bf16 MMA M == N we use M for warp 4-7 + shmem_idx_b_off.\n size_t shmem_idx =\n (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2) + ((warpId / (WARPS_PER_BLOCK / 2)) * shmem_idx_b_off);\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n const __hip_bfloat16 *lane_ptr = warp_ptr + tile_k * K + stridePerLaneCopy * K_GLOBAL + laneLoadElem;\n\n // Shift the second half of the warp to the next row / column in the shared memory.\n shmem_idx += stridePerLaneCopy;\n\n#pragma unroll\n for (int i = 0; i < chunksPerLane; i++) {\n // Copy 16 bytes at once in each lane.\n pipe.producer_acquire();\n cuda::memcpy_async(&shmem[shmem_idx][laneLoadElem], lane_ptr, shape4, pipe);\n pipe.producer_commit();\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n cuda::pipeline_consumer_wait_prior<0>(pipe);\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / BLOCK_ROW_WARPS) * M * BLOCK_ROW_WARPS + (i * M);\n const __hip_bfloat16 *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_BF16);\n\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be reused\n // against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const __hip_bfloat16 *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_BF16);\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n\n pipe.consumer_release();\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n // Uniform, point-wise transformations of ALL fragment elements by ALL threads in the\n // warp are well-defined even though element indices within fragment storage are not defined.\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * K + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global memory.\n float *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n#endif\n}\n\n// Performs an MxNxK bf16 GEMM (C=alpha*A*B + beta*C) assuming:\n// 1) Matrices are packed in memory.\n// 2) M, N and K are multiples of 16, 16 and 16 respectively.\n// 3) A is row major, B is column major matrix.\n// Note: This is a less performant version of the compute_bf16gemm kernel. It is designed for\n// demonstration purposes only to show the CUDA WMMA API use without relying on\n// availability of the shared memory.\n__global__ void simple_wmma_bf16gemm(__hip_bfloat16 *a,\n __hip_bfloat16 *b,\n float *c,\n float *d,\n int m_ld,\n int n_ld,\n int k_ld,\n float alpha,\n float beta)\n{\n#if __CUDA_ARCH__ >= 800\n // Leading dimensions. Packed with no transpositions.\n int lda = k_ld;\n int ldb = k_ld;\n int ldc = n_ld;\n\n // Tile using a 2D grid\n int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;\n int warpN = (blockIdx.y * blockDim.y + threadIdx.y);\n\n // Declare the fragments\n wmma::fragment a_frag;\n wmma::fragment b_frag;\n wmma::fragment acc_frag;\n wmma::fragment c_frag;\n\n wmma::fill_fragment(acc_frag, 0.0f);\n\n // Loop over k\n for (int i = 0; i < k_ld; i += K) {\n int aCol = i;\n int aRow = warpM * M;\n\n int bCol = i;\n int bRow = warpN * N;\n\n // Bounds checking\n if (aRow < m_ld && aCol < k_ld && bRow < k_ld && bCol < n_ld) {\n // Load the inputs\n wmma::load_matrix_sync(a_frag, a + aCol + aRow * lda, lda);\n wmma::load_matrix_sync(b_frag, b + bRow + bCol * ldb, ldb);\n\n // Perform the matrix multiplication\n wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);\n }\n }\n\n // Load in the current value of c, scale it by beta, and add this our result scaled by alpha\n int cCol = warpN * N;\n int cRow = warpM * M;\n\n if (cRow < m_ld && cCol < n_ld) {\n wmma::load_matrix_sync(c_frag, c + cCol + cRow * ldc, ldc, wmma::mem_row_major);\n\n for (int i = 0; i < c_frag.num_elements; i++) {\n c_frag.x[i] = alpha * acc_frag.x[i] + beta * c_frag.x[i];\n }\n\n // Store the output\n wmma::store_matrix_sync(d + cCol + cRow * ldc, c_frag, ldc, wmma::mem_row_major);\n }\n#endif\n}\n\n__host__ void matMultiplyOnHost(__hip_bfloat16 *A,\n __hip_bfloat16 *B,\n float *C,\n float alpha,\n float beta,\n int numARows,\n int numAColumns,\n int numBRows,\n int numBColumns,\n int numCRows,\n int numCColumns)\n{\n for (int i = 0; i < numCRows; i++) {\n for (int j = 0; j < numCColumns; j++) {\n float temp = 0.0;\n\n for (int k = 0; k < numAColumns; k++) {\n temp += (float)A[i * numAColumns + k] * (float)B[j * numBRows + k];\n }\n\n C[i * numCColumns + j] = temp * alpha + beta * C[i * numCColumns + j];\n }\n }\n}\n\nint main(int argc, char **argv)\n{\n printf(\"Initializing...\\n\");\n\n int dev = findCudaDevice(argc, (const char **)argv);\n\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, dev));\n\n // Tensor cores require a GPU of Volta (SM8X) architecture or higher.\n if (deviceProp.major < 8) {\n printf(\"bf16TensorCoreGemm requires requires SM 8.0 or higher to use Tensor Cores. Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n printf(\"M: %d (%d x %d)\\n\", M_GLOBAL, M, M_TILES);\n printf(\"N: %d (%d x %d)\\n\", N_GLOBAL, N, N_TILES);\n printf(\"K: %d (%d x %d)\\n\", K_GLOBAL, K, K_TILES);\n\n __hip_bfloat16 *A_h = NULL;\n __hip_bfloat16 *B_h = NULL;\n float *C_h = NULL;\n#if CPU_DEBUG\n float *result_hD = NULL;\n float *result_host = NULL;\n#endif\n\n A_h = (__hip_bfloat16 *)malloc(sizeof(__hip_bfloat16) * M_GLOBAL * K_GLOBAL);\n B_h = (__hip_bfloat16 *)malloc(sizeof(__hip_bfloat16) * K_GLOBAL * N_GLOBAL);\n C_h = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n#if CPU_DEBUG\n result_hD = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n result_host = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n#endif\n\n __hip_bfloat16 *A = NULL;\n __hip_bfloat16 *B = NULL;\n float *C = NULL;\n float *D = NULL;\n\n checkCudaErrors(hipMalloc((void **)&A, sizeof(__hip_bfloat16) * M_GLOBAL * K_GLOBAL));\n checkCudaErrors(hipMalloc((void **)&B, sizeof(__hip_bfloat16) * N_GLOBAL * K_GLOBAL));\n checkCudaErrors(hipMalloc((void **)&C, sizeof(float) * M_GLOBAL * N_GLOBAL));\n checkCudaErrors(hipMalloc((void **)&D, sizeof(float) * M_GLOBAL * N_GLOBAL));\n\n assert(((unsigned long long)A) % 128 == 0);\n assert(((unsigned long long)B) % 128 == 0);\n assert(((unsigned long long)C) % 128 == 0);\n assert(((unsigned long long)D) % 128 == 0);\n\n init_host_matrices(A_h, B_h, C_h);\n\n printf(\"Preparing data for GPU...\\n\");\n\n checkCudaErrors(hipMemcpy(A, A_h, sizeof(__hip_bfloat16) * M_GLOBAL * K_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(B, B_h, sizeof(__hip_bfloat16) * N_GLOBAL * K_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(C, C_h, sizeof(float) * M_GLOBAL * N_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemset(D, 0, sizeof(float) * M_GLOBAL * N_GLOBAL));\n\n enum {\n // Compute the right amount of shared memory to request.\n // We need shared memory to hold per-CTA C and D matrix tiles, and to cache per-CTA chunks\n // of the A and B matrices. Therefore, the right amount to request is the maximum of those\n // two numbers.\n SHMEM_SZ = MAX(sizeof(__hip_bfloat16) * (BLOCK_COL_TILES * M) * (CHUNK_K * K + SKEW_BF16) * 2,\n M * (BLOCK_ROW_WARPS * WARP_ROW_TILES) * N * (BLOCK_COL_WARPS * WARP_COL_TILES) * sizeof(float))\n };\n\n printf(\"Required shared memory size: %lu Kb\\n\", SHMEM_SZ / 1024UL);\n\n const float alpha = 1.1f;\n const float beta = 1.2f;\n\n hipEvent_t start, stop;\n\n checkCudaErrors(hipEventCreate(&start));\n checkCudaErrors(hipEventCreate(&stop));\n checkCudaErrors(hipEventRecord(start));\n\n // kernel to run - default (b16mma_shmem_gemm_async_copy == 0)\n kernels selected_kernel = bf16mma_shmem_gemm_async_copy;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"kernel\")) {\n int kernel_number = getCmdLineArgumentInt(argc, (const char **)argv, \"kernel\");\n if (kernel_number < 3) {\n selected_kernel = (kernels)kernel_number;\n }\n else {\n printf(\"Error: kernel number should be between 0 to 2, you have entered %d\\n\", kernel_number);\n exit(EXIT_FAILURE);\n }\n }\n\n // If enough shared memory available on the GPU use high performant kernel\n if ((deviceProp.sharedMemPerMultiprocessor >= SHMEM_SZ) && (selected_kernel != simple_bf16mma_gemm)) {\n printf(\"Computing using high performance kernel = %d - %s\\n\", selected_kernel, kernelNames[selected_kernel]);\n\n switch (selected_kernel) {\n case bf16mma_shmem_gemm_async_copy:\n default:\n checkCudaErrors(hipFuncSetAttribute(reinterpret_cast(compute_bf16gemm_async_copy), hipFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors(\n (compute_bf16gemm_async_copy<<>>(\n A, B, C, D, alpha, beta)));\n break;\n case bf16mma_shmem_gemm:\n checkCudaErrors(\n hipFuncSetAttribute(reinterpret_cast(compute_bf16gemm), hipFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors((compute_bf16gemm<<>>(\n A, B, C, D, alpha, beta)));\n break;\n }\n#if CPU_DEBUG\n checkCudaErrors(hipMemcpy(result_hD, D, sizeof(float) * M_GLOBAL * N_GLOBAL, hipMemcpyDeviceToHost));\n#endif\n }\n else {\n dim3 gridDim;\n dim3 blockDim;\n\n // blockDim.x must be a multple of warpSize\n // 128x4 means we have 16 warps and a block computes a 64x64 output tile\n blockDim.x = 128;\n blockDim.y = 4;\n\n gridDim.x = (M_GLOBAL + (M * blockDim.x / 32 - 1)) / (M * blockDim.x / 32);\n gridDim.y = (N_GLOBAL + N * blockDim.y - 1) / (N * blockDim.y);\n\n printf(\"Computing... using simple_wmma_gemm kernel\\n\");\n simple_wmma_bf16gemm<<>>(A, B, C, D, M_GLOBAL, N_GLOBAL, K_GLOBAL, alpha, beta);\n#if CPU_DEBUG\n checkCudaErrors(hipMemcpy(result_hD, D, sizeof(float) * M_GLOBAL * N_GLOBAL, hipMemcpyDeviceToHost));\n#endif\n }\n\n checkCudaErrors(hipEventRecord(stop));\n checkCudaErrors(hipEventSynchronize(stop));\n\n#if CPU_DEBUG\n printf(\"Verifying correctness of the computations...\\n\");\n\n memcpy(result_host, C_h, sizeof(float) * M_GLOBAL * N_GLOBAL);\n\n matMultiplyOnHost(A_h, B_h, result_host, alpha, beta, M_GLOBAL, K_GLOBAL, K_GLOBAL, N_GLOBAL, M_GLOBAL, N_GLOBAL);\n\n for (int i = 0; i < N_GLOBAL * M_GLOBAL; i++) {\n if (fabs(result_hD[i] - result_host[i]) > 0.1f) {\n printf(\"mismatch i=%d result_hD=%f result_host=%f\\n\", i, result_hD[i], result_host[i]);\n }\n }\n free(result_hD);\n free(result_host);\n#endif\n\n float milliseconds = 0;\n\n checkCudaErrors(hipEventElapsedTime(&milliseconds, start, stop));\n\n printf(\"Time: %f ms\\n\", milliseconds);\n printf(\"TFLOPS: %.2f\\n\", (((double)M_GLOBAL * N_GLOBAL * K_GLOBAL * 2) / (milliseconds / 1000.)) / 1e12);\n\n free(A_h);\n free(B_h);\n free(C_h);\n checkCudaErrors(hipFree((void *)A));\n checkCudaErrors(hipFree((void *)B));\n checkCudaErrors(hipFree((void *)C));\n checkCudaErrors(hipFree((void *)D));\n\n return 0;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/bf16TensorCoreGemm/bf16TensorCoreGemm.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n\n#include \"jacobi.h\"\n\nnamespace cg = cooperative_groups;\n\n// 8 Rows of square-matrix A processed by each CTA.\n// This can be max 32 and only power of 2 (i.e., 2/4/8/16/32).\n#define ROWS_PER_CTA 8\n\n#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 600\n#else\n__device__ double atomicAdd(double *address, double val)\n{\n unsigned long long int *address_as_ull = (unsigned long long int *)address;\n unsigned long long int old = *address_as_ull, assumed;\n\n do {\n assumed = old;\n old = atomicCAS(address_as_ull, assumed, __double_as_longlong(val + __longlong_as_double(assumed)));\n\n // Note: uses integer comparison to avoid hang in case of NaN (since NaN !=\n // NaN)\n } while (assumed != old);\n\n return __longlong_as_double(old);\n}\n#endif\n\nstatic __global__ void\nJacobiMethod(const float *A, const double *b, const float conv_threshold, double *x, double *x_new, double *sum)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ double x_shared[N_ROWS]; // N_ROWS == n\n __shared__ double b_shared[ROWS_PER_CTA + 1];\n\n for (int i = threadIdx.x; i < N_ROWS; i += blockDim.x) {\n x_shared[i] = x[i];\n }\n\n if (threadIdx.x < ROWS_PER_CTA) {\n int k = threadIdx.x;\n for (int i = k + (blockIdx.x * ROWS_PER_CTA); (k < ROWS_PER_CTA) && (i < N_ROWS);\n k += ROWS_PER_CTA, i += ROWS_PER_CTA) {\n b_shared[i % (ROWS_PER_CTA + 1)] = b[i];\n }\n }\n\n cg::sync(cta);\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n for (int k = 0, i = blockIdx.x * ROWS_PER_CTA; (k < ROWS_PER_CTA) && (i < N_ROWS); k++, i++) {\n double rowThreadSum = 0.0;\n for (int j = threadIdx.x; j < N_ROWS; j += blockDim.x) {\n rowThreadSum += (A[i * N_ROWS + j] * x_shared[j]);\n }\n\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n rowThreadSum += tile32.shfl_down(rowThreadSum, offset);\n }\n\n if (tile32.thread_rank() == 0) {\n atomicAdd(&b_shared[i % (ROWS_PER_CTA + 1)], -rowThreadSum);\n }\n }\n\n cg::sync(cta);\n\n if (threadIdx.x < ROWS_PER_CTA) {\n cg::thread_block_tile tile8 = cg::tiled_partition(cta);\n double temp_sum = 0.0;\n\n int k = threadIdx.x;\n\n for (int i = k + (blockIdx.x * ROWS_PER_CTA); (k < ROWS_PER_CTA) && (i < N_ROWS);\n k += ROWS_PER_CTA, i += ROWS_PER_CTA) {\n double dx = b_shared[i % (ROWS_PER_CTA + 1)];\n dx /= A[i * N_ROWS + i];\n\n x_new[i] = (x_shared[i] + dx);\n temp_sum += fabs(dx);\n }\n\n for (int offset = tile8.size() / 2; offset > 0; offset /= 2) {\n temp_sum += tile8.shfl_down(temp_sum, offset);\n }\n\n if (tile8.thread_rank() == 0) {\n atomicAdd(sum, temp_sum);\n }\n }\n}\n\n// Thread block size for finalError kernel should be multiple of 32\nstatic __global__ void finalError(double *x, double *g_sum)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n extern __shared__ double warpSum[];\n double sum = 0.0;\n\n int globalThreadId = blockIdx.x * blockDim.x + threadIdx.x;\n\n for (int i = globalThreadId; i < N_ROWS; i += blockDim.x * gridDim.x) {\n double d = x[i] - 1.0;\n sum += fabs(d);\n }\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n sum += tile32.shfl_down(sum, offset);\n }\n\n if (tile32.thread_rank() == 0) {\n warpSum[threadIdx.x / warpSize] = sum;\n }\n\n cg::sync(cta);\n\n double blockSum = 0.0;\n if (threadIdx.x < (blockDim.x / warpSize)) {\n blockSum = warpSum[threadIdx.x];\n }\n\n if (threadIdx.x < 32) {\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n blockSum += tile32.shfl_down(blockSum, offset);\n }\n if (tile32.thread_rank() == 0) {\n atomicAdd(g_sum, blockSum);\n }\n }\n}\n\ndouble JacobiMethodGpuCudaGraphExecKernelSetParams(const float *A,\n const double *b,\n const float conv_threshold,\n const int max_iter,\n double *x,\n double *x_new,\n cudaStream_t stream)\n{\n // CTA size\n dim3 nthreads(256, 1, 1);\n // grid size\n dim3 nblocks((N_ROWS / ROWS_PER_CTA) + 2, 1, 1);\n cudaGraph_t graph;\n cudaGraphExec_t graphExec = NULL;\n\n double sum = 0.0;\n double *d_sum = NULL;\n checkCudaErrors(cudaMalloc(&d_sum, sizeof(double)));\n\n std::vector nodeDependencies;\n cudaGraphNode_t memcpyNode, jacobiKernelNode, memsetNode;\n cudaMemcpy3DParms memcpyParams = {0};\n cudaMemsetParams memsetParams = {0};\n\n memsetParams.dst = (void *)d_sum;\n memsetParams.value = 0;\n memsetParams.pitch = 0;\n // elementSize can be max 4 bytes, so we take sizeof(float) and width=2\n memsetParams.elementSize = sizeof(float);\n memsetParams.width = 2;\n memsetParams.height = 1;\n\n checkCudaErrors(cudaGraphCreate(&graph, 0));\n checkCudaErrors(cudaGraphAddMemsetNode(&memsetNode, graph, NULL, 0, &memsetParams));\n nodeDependencies.push_back(memsetNode);\n\n cudaKernelNodeParams NodeParams0, NodeParams1;\n NodeParams0.func = (void *)JacobiMethod;\n NodeParams0.gridDim = nblocks;\n NodeParams0.blockDim = nthreads;\n NodeParams0.sharedMemBytes = 0;\n void *kernelArgs0[6] = {\n (void *)&A, (void *)&b, (void *)&conv_threshold, (void *)&x, (void *)&x_new, (void *)&d_sum};\n NodeParams0.kernelParams = kernelArgs0;\n NodeParams0.extra = NULL;\n\n checkCudaErrors(cudaGraphAddKernelNode(\n &jacobiKernelNode, graph, nodeDependencies.data(), nodeDependencies.size(), &NodeParams0));\n\n nodeDependencies.clear();\n nodeDependencies.push_back(jacobiKernelNode);\n\n memcpyParams.srcArray = NULL;\n memcpyParams.srcPos = make_cudaPos(0, 0, 0);\n memcpyParams.srcPtr = make_cudaPitchedPtr(d_sum, sizeof(double), 1, 1);\n memcpyParams.dstArray = NULL;\n memcpyParams.dstPos = make_cudaPos(0, 0, 0);\n memcpyParams.dstPtr = make_cudaPitchedPtr(&sum, sizeof(double), 1, 1);\n memcpyParams.extent = make_cudaExtent(sizeof(double), 1, 1);\n memcpyParams.kind = cudaMemcpyDeviceToHost;\n\n checkCudaErrors(\n cudaGraphAddMemcpyNode(&memcpyNode, graph, nodeDependencies.data(), nodeDependencies.size(), &memcpyParams));\n\n checkCudaErrors(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n NodeParams1.func = (void *)JacobiMethod;\n NodeParams1.gridDim = nblocks;\n NodeParams1.blockDim = nthreads;\n NodeParams1.sharedMemBytes = 0;\n void *kernelArgs1[6] = {\n (void *)&A, (void *)&b, (void *)&conv_threshold, (void *)&x_new, (void *)&x, (void *)&d_sum};\n NodeParams1.kernelParams = kernelArgs1;\n NodeParams1.extra = NULL;\n\n int k = 0;\n for (k = 0; k < max_iter; k++) {\n checkCudaErrors(cudaGraphExecKernelNodeSetParams(\n graphExec, jacobiKernelNode, ((k & 1) == 0) ? &NodeParams0 : &NodeParams1));\n checkCudaErrors(cudaGraphLaunch(graphExec, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n if (sum <= conv_threshold) {\n checkCudaErrors(cudaMemsetAsync(d_sum, 0, sizeof(double), stream));\n nblocks.x = (N_ROWS / nthreads.x) + 1;\n size_t sharedMemSize = ((nthreads.x / 32) + 1) * sizeof(double);\n if ((k & 1) == 0) {\n finalError<<>>(x_new, d_sum);\n }\n else {\n finalError<<>>(x, d_sum);\n }\n\n checkCudaErrors(cudaMemcpyAsync(&sum, d_sum, sizeof(double), cudaMemcpyDeviceToHost, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n printf(\"GPU iterations : %d\\n\", k + 1);\n printf(\"GPU error : %.3e\\n\", sum);\n break;\n }\n }\n\n checkCudaErrors(cudaFree(d_sum));\n return sum;\n}\n\ndouble JacobiMethodGpuCudaGraphExecUpdate(const float *A,\n const double *b,\n const float conv_threshold,\n const int max_iter,\n double *x,\n double *x_new,\n cudaStream_t stream)\n{\n // CTA size\n dim3 nthreads(256, 1, 1);\n // grid size\n dim3 nblocks((N_ROWS / ROWS_PER_CTA) + 2, 1, 1);\n cudaGraph_t graph;\n cudaGraphExec_t graphExec = NULL;\n\n double sum = 0.0;\n double *d_sum;\n checkCudaErrors(cudaMalloc(&d_sum, sizeof(double)));\n\n int k = 0;\n for (k = 0; k < max_iter; k++) {\n checkCudaErrors(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal));\n checkCudaErrors(cudaMemsetAsync(d_sum, 0, sizeof(double), stream));\n if ((k & 1) == 0) {\n JacobiMethod<<>>(A, b, conv_threshold, x, x_new, d_sum);\n }\n else {\n JacobiMethod<<>>(A, b, conv_threshold, x_new, x, d_sum);\n }\n checkCudaErrors(cudaMemcpyAsync(&sum, d_sum, sizeof(double), cudaMemcpyDeviceToHost, stream));\n checkCudaErrors(cudaStreamEndCapture(stream, &graph));\n\n if (graphExec == NULL) {\n checkCudaErrors(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n }\n else {\n cudaGraphExecUpdateResult updateResult_out;\n checkCudaErrors(cudaGraphExecUpdate(graphExec, graph, NULL, &updateResult_out));\n if (updateResult_out != cudaGraphExecUpdateSuccess) {\n if (graphExec != NULL) {\n checkCudaErrors(cudaGraphExecDestroy(graphExec));\n }\n printf(\"k = %d graph update failed with error - %d\\n\", k, updateResult_out);\n checkCudaErrors(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n }\n }\n checkCudaErrors(cudaGraphLaunch(graphExec, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n if (sum <= conv_threshold) {\n checkCudaErrors(cudaMemsetAsync(d_sum, 0, sizeof(double), stream));\n nblocks.x = (N_ROWS / nthreads.x) + 1;\n size_t sharedMemSize = ((nthreads.x / 32) + 1) * sizeof(double);\n if ((k & 1) == 0) {\n finalError<<>>(x_new, d_sum);\n }\n else {\n finalError<<>>(x, d_sum);\n }\n\n checkCudaErrors(cudaMemcpyAsync(&sum, d_sum, sizeof(double), cudaMemcpyDeviceToHost, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n printf(\"GPU iterations : %d\\n\", k + 1);\n printf(\"GPU error : %.3e\\n\", sum);\n break;\n }\n }\n\n checkCudaErrors(cudaFree(d_sum));\n return sum;\n}\n\ndouble JacobiMethodGpu(const float *A,\n const double *b,\n const float conv_threshold,\n const int max_iter,\n double *x,\n double *x_new,\n cudaStream_t stream)\n{\n // CTA size\n dim3 nthreads(256, 1, 1);\n // grid size\n dim3 nblocks((N_ROWS / ROWS_PER_CTA) + 2, 1, 1);\n\n double sum = 0.0;\n double *d_sum;\n checkCudaErrors(cudaMalloc(&d_sum, sizeof(double)));\n int k = 0;\n\n for (k = 0; k < max_iter; k++) {\n checkCudaErrors(cudaMemsetAsync(d_sum, 0, sizeof(double), stream));\n if ((k & 1) == 0) {\n JacobiMethod<<>>(A, b, conv_threshold, x, x_new, d_sum);\n }\n else {\n JacobiMethod<<>>(A, b, conv_threshold, x_new, x, d_sum);\n }\n checkCudaErrors(cudaMemcpyAsync(&sum, d_sum, sizeof(double), cudaMemcpyDeviceToHost, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n if (sum <= conv_threshold) {\n checkCudaErrors(cudaMemsetAsync(d_sum, 0, sizeof(double), stream));\n nblocks.x = (N_ROWS / nthreads.x) + 1;\n size_t sharedMemSize = ((nthreads.x / 32) + 1) * sizeof(double);\n if ((k & 1) == 0) {\n finalError<<>>(x_new, d_sum);\n }\n else {\n finalError<<>>(x, d_sum);\n }\n\n checkCudaErrors(cudaMemcpyAsync(&sum, d_sum, sizeof(double), cudaMemcpyDeviceToHost, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n printf(\"GPU iterations : %d\\n\", k + 1);\n printf(\"GPU error : %.3e\\n\", sum);\n break;\n }\n }\n\n checkCudaErrors(cudaFree(d_sum));\n return sum;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n\n#include \"jacobi.h\"\n\nnamespace cg = cooperative_groups;\n\n// 8 Rows of square-matrix A processed by each CTA.\n// This can be max 32 and only power of 2 (i.e., 2/4/8/16/32).\n#define ROWS_PER_CTA 8\n\n#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 600\n#else\n__device__ double atomicAdd(double *address, double val)\n{\n unsigned long long int *address_as_ull = (unsigned long long int *)address;\n unsigned long long int old = *address_as_ull, assumed;\n\n do {\n assumed = old;\n old = atomicCAS(address_as_ull, assumed, __double_as_longlong(val + __longlong_as_double(assumed)));\n\n // Note: uses integer comparison to avoid hang in case of NaN (since NaN !=\n // NaN)\n } while (assumed != old);\n\n return __longlong_as_double(old);\n}\n#endif\n\nstatic __global__ void\nJacobiMethod(const float *A, const double *b, const float conv_threshold, double *x, double *x_new, double *sum)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ double x_shared[N_ROWS]; // N_ROWS == n\n __shared__ double b_shared[ROWS_PER_CTA + 1];\n\n for (int i = threadIdx.x; i < N_ROWS; i += blockDim.x) {\n x_shared[i] = x[i];\n }\n\n if (threadIdx.x < ROWS_PER_CTA) {\n int k = threadIdx.x;\n for (int i = k + (blockIdx.x * ROWS_PER_CTA); (k < ROWS_PER_CTA) && (i < N_ROWS);\n k += ROWS_PER_CTA, i += ROWS_PER_CTA) {\n b_shared[i % (ROWS_PER_CTA + 1)] = b[i];\n }\n }\n\n cg::sync(cta);\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n for (int k = 0, i = blockIdx.x * ROWS_PER_CTA; (k < ROWS_PER_CTA) && (i < N_ROWS); k++, i++) {\n double rowThreadSum = 0.0;\n for (int j = threadIdx.x; j < N_ROWS; j += blockDim.x) {\n rowThreadSum += (A[i * N_ROWS + j] * x_shared[j]);\n }\n\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n rowThreadSum += tile32.shfl_down(rowThreadSum, offset);\n }\n\n if (tile32.thread_rank() == 0) {\n atomicAdd(&b_shared[i % (ROWS_PER_CTA + 1)], -rowThreadSum);\n }\n }\n\n cg::sync(cta);\n\n if (threadIdx.x < ROWS_PER_CTA) {\n cg::thread_block_tile tile8 = cg::tiled_partition(cta);\n double temp_sum = 0.0;\n\n int k = threadIdx.x;\n\n for (int i = k + (blockIdx.x * ROWS_PER_CTA); (k < ROWS_PER_CTA) && (i < N_ROWS);\n k += ROWS_PER_CTA, i += ROWS_PER_CTA) {\n double dx = b_shared[i % (ROWS_PER_CTA + 1)];\n dx /= A[i * N_ROWS + i];\n\n x_new[i] = (x_shared[i] + dx);\n temp_sum += fabs(dx);\n }\n\n for (int offset = tile8.size() / 2; offset > 0; offset /= 2) {\n temp_sum += tile8.shfl_down(temp_sum, offset);\n }\n\n if (tile8.thread_rank() == 0) {\n atomicAdd(sum, temp_sum);\n }\n }\n}\n\n// Thread block size for finalError kernel should be multiple of 32\nstatic __global__ void finalError(double *x, double *g_sum)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n extern __shared__ double warpSum[];\n double sum = 0.0;\n\n int globalThreadId = blockIdx.x * blockDim.x + threadIdx.x;\n\n for (int i = globalThreadId; i < N_ROWS; i += blockDim.x * gridDim.x) {\n double d = x[i] - 1.0;\n sum += fabs(d);\n }\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n sum += tile32.shfl_down(sum, offset);\n }\n\n if (tile32.thread_rank() == 0) {\n warpSum[threadIdx.x / warpSize] = sum;\n }\n\n cg::sync(cta);\n\n double blockSum = 0.0;\n if (threadIdx.x < (blockDim.x / warpSize)) {\n blockSum = warpSum[threadIdx.x];\n }\n\n if (threadIdx.x < 32) {\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n blockSum += tile32.shfl_down(blockSum, offset);\n }\n if (tile32.thread_rank() == 0) {\n atomicAdd(g_sum, blockSum);\n }\n }\n}\n\ndouble JacobiMethodGpuCudaGraphExecKernelSetParams(const float *A,\n const double *b,\n const float conv_threshold,\n const int max_iter,\n double *x,\n double *x_new,\n hipStream_t stream)\n{\n // CTA size\n dim3 nthreads(256, 1, 1);\n // grid size\n dim3 nblocks((N_ROWS / ROWS_PER_CTA) + 2, 1, 1);\n hipGraph_t graph;\n hipGraphExec_t graphExec = NULL;\n\n double sum = 0.0;\n double *d_sum = NULL;\n checkCudaErrors(hipMalloc(&d_sum, sizeof(double)));\n\n std::vector nodeDependencies;\n hipGraphNode_t memcpyNode, jacobiKernelNode, memsetNode;\n hipMemcpy3DParms memcpyParams = {0};\n hipMemsetParams memsetParams = {0};\n\n memsetParams.dst = (void *)d_sum;\n memsetParams.value = 0;\n memsetParams.pitch = 0;\n // elementSize can be max 4 bytes, so we take sizeof(float) and width=2\n memsetParams.elementSize = sizeof(float);\n memsetParams.width = 2;\n memsetParams.height = 1;\n\n checkCudaErrors(hipGraphCreate(&graph, 0));\n checkCudaErrors(hipGraphAddMemsetNode(&memsetNode, graph, NULL, 0, &memsetParams));\n nodeDependencies.push_back(memsetNode);\n\n hipKernelNodeParams NodeParams0, NodeParams1;\n NodeParams0.func = (void *)JacobiMethod;\n NodeParams0.gridDim = nblocks;\n NodeParams0.blockDim = nthreads;\n NodeParams0.sharedMemBytes = 0;\n void *kernelArgs0[6] = {\n (void *)&A, (void *)&b, (void *)&conv_threshold, (void *)&x, (void *)&x_new, (void *)&d_sum};\n NodeParams0.kernelParams = kernelArgs0;\n NodeParams0.extra = NULL;\n\n checkCudaErrors(hipGraphAddKernelNode(\n &jacobiKernelNode, graph, nodeDependencies.data(), nodeDependencies.size(), &NodeParams0));\n\n nodeDependencies.clear();\n nodeDependencies.push_back(jacobiKernelNode);\n\n memcpyParams.srcArray = NULL;\n memcpyParams.srcPos = make_hipPos(0, 0, 0);\n memcpyParams.srcPtr = make_hipPitchedPtr(d_sum, sizeof(double), 1, 1);\n memcpyParams.dstArray = NULL;\n memcpyParams.dstPos = make_hipPos(0, 0, 0);\n memcpyParams.dstPtr = make_hipPitchedPtr(&sum, sizeof(double), 1, 1);\n memcpyParams.extent = make_hipExtent(sizeof(double), 1, 1);\n memcpyParams.kind = hipMemcpyDeviceToHost;\n\n checkCudaErrors(\n hipGraphAddMemcpyNode(&memcpyNode, graph, nodeDependencies.data(), nodeDependencies.size(), &memcpyParams));\n\n checkCudaErrors(hipGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n NodeParams1.func = (void *)JacobiMethod;\n NodeParams1.gridDim = nblocks;\n NodeParams1.blockDim = nthreads;\n NodeParams1.sharedMemBytes = 0;\n void *kernelArgs1[6] = {\n (void *)&A, (void *)&b, (void *)&conv_threshold, (void *)&x_new, (void *)&x, (void *)&d_sum};\n NodeParams1.kernelParams = kernelArgs1;\n NodeParams1.extra = NULL;\n\n int k = 0;\n for (k = 0; k < max_iter; k++) {\n checkCudaErrors(hipGraphExecKernelNodeSetParams(\n graphExec, jacobiKernelNode, ((k & 1) == 0) ? &NodeParams0 : &NodeParams1));\n checkCudaErrors(hipGraphLaunch(graphExec, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n\n if (sum <= conv_threshold) {\n checkCudaErrors(hipMemsetAsync(d_sum, 0, sizeof(double), stream));\n nblocks.x = (N_ROWS / nthreads.x) + 1;\n size_t sharedMemSize = ((nthreads.x / 32) + 1) * sizeof(double);\n if ((k & 1) == 0) {\n finalError<<>>(x_new, d_sum);\n }\n else {\n finalError<<>>(x, d_sum);\n }\n\n checkCudaErrors(hipMemcpyAsync(&sum, d_sum, sizeof(double), hipMemcpyDeviceToHost, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n printf(\"GPU iterations : %d\\n\", k + 1);\n printf(\"GPU error : %.3e\\n\", sum);\n break;\n }\n }\n\n checkCudaErrors(hipFree(d_sum));\n return sum;\n}\n\ndouble JacobiMethodGpuCudaGraphExecUpdate(const float *A,\n const double *b,\n const float conv_threshold,\n const int max_iter,\n double *x,\n double *x_new,\n hipStream_t stream)\n{\n // CTA size\n dim3 nthreads(256, 1, 1);\n // grid size\n dim3 nblocks((N_ROWS / ROWS_PER_CTA) + 2, 1, 1);\n hipGraph_t graph;\n hipGraphExec_t graphExec = NULL;\n\n double sum = 0.0;\n double *d_sum;\n checkCudaErrors(hipMalloc(&d_sum, sizeof(double)));\n\n int k = 0;\n for (k = 0; k < max_iter; k++) {\n checkCudaErrors(hipStreamBeginCapture(stream, hipStreamCaptureModeGlobal));\n checkCudaErrors(hipMemsetAsync(d_sum, 0, sizeof(double), stream));\n if ((k & 1) == 0) {\n JacobiMethod<<>>(A, b, conv_threshold, x, x_new, d_sum);\n }\n else {\n JacobiMethod<<>>(A, b, conv_threshold, x_new, x, d_sum);\n }\n checkCudaErrors(hipMemcpyAsync(&sum, d_sum, sizeof(double), hipMemcpyDeviceToHost, stream));\n checkCudaErrors(hipStreamEndCapture(stream, &graph));\n\n if (graphExec == NULL) {\n checkCudaErrors(hipGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n }\n else {\n hipGraphExecUpdateResult updateResult_out;\n checkCudaErrors(hipGraphExecUpdate(graphExec, graph, NULL, &updateResult_out));\n if (updateResult_out != hipGraphExecUpdateSuccess) {\n if (graphExec != NULL) {\n checkCudaErrors(hipGraphExecDestroy(graphExec));\n }\n printf(\"k = %d graph update failed with error - %d\\n\", k, updateResult_out);\n checkCudaErrors(hipGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n }\n }\n checkCudaErrors(hipGraphLaunch(graphExec, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n\n if (sum <= conv_threshold) {\n checkCudaErrors(hipMemsetAsync(d_sum, 0, sizeof(double), stream));\n nblocks.x = (N_ROWS / nthreads.x) + 1;\n size_t sharedMemSize = ((nthreads.x / 32) + 1) * sizeof(double);\n if ((k & 1) == 0) {\n finalError<<>>(x_new, d_sum);\n }\n else {\n finalError<<>>(x, d_sum);\n }\n\n checkCudaErrors(hipMemcpyAsync(&sum, d_sum, sizeof(double), hipMemcpyDeviceToHost, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n printf(\"GPU iterations : %d\\n\", k + 1);\n printf(\"GPU error : %.3e\\n\", sum);\n break;\n }\n }\n\n checkCudaErrors(hipFree(d_sum));\n return sum;\n}\n\ndouble JacobiMethodGpu(const float *A,\n const double *b,\n const float conv_threshold,\n const int max_iter,\n double *x,\n double *x_new,\n hipStream_t stream)\n{\n // CTA size\n dim3 nthreads(256, 1, 1);\n // grid size\n dim3 nblocks((N_ROWS / ROWS_PER_CTA) + 2, 1, 1);\n\n double sum = 0.0;\n double *d_sum;\n checkCudaErrors(hipMalloc(&d_sum, sizeof(double)));\n int k = 0;\n\n for (k = 0; k < max_iter; k++) {\n checkCudaErrors(hipMemsetAsync(d_sum, 0, sizeof(double), stream));\n if ((k & 1) == 0) {\n JacobiMethod<<>>(A, b, conv_threshold, x, x_new, d_sum);\n }\n else {\n JacobiMethod<<>>(A, b, conv_threshold, x_new, x, d_sum);\n }\n checkCudaErrors(hipMemcpyAsync(&sum, d_sum, sizeof(double), hipMemcpyDeviceToHost, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n\n if (sum <= conv_threshold) {\n checkCudaErrors(hipMemsetAsync(d_sum, 0, sizeof(double), stream));\n nblocks.x = (N_ROWS / nthreads.x) + 1;\n size_t sharedMemSize = ((nthreads.x / 32) + 1) * sizeof(double);\n if ((k & 1) == 0) {\n finalError<<>>(x_new, d_sum);\n }\n else {\n finalError<<>>(x, d_sum);\n }\n\n checkCudaErrors(hipMemcpyAsync(&sum, d_sum, sizeof(double), hipMemcpyDeviceToHost, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n printf(\"GPU iterations : %d\\n\", k + 1);\n printf(\"GPU error : %.3e\\n\", sum);\n break;\n }\n }\n\n checkCudaErrors(hipFree(d_sum));\n return sum;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/jacobiCudaGraphs/jacobi.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n#include \n\n////////////////////////////////////////////////////////////////////////////////\n// Variable on the GPU used to generate unique identifiers of blocks.\n////////////////////////////////////////////////////////////////////////////////\n__device__ int g_uids = 0;\n\n////////////////////////////////////////////////////////////////////////////////\n// Print a simple message to signal the block which is currently executing.\n////////////////////////////////////////////////////////////////////////////////\n__device__ void print_info(int depth, int thread, int uid, int parent_uid)\n{\n if (threadIdx.x == 0) {\n if (depth == 0)\n printf(\"BLOCK %d launched by the host\\n\", uid);\n else {\n char buffer[32];\n\n for (int i = 0; i < depth; ++i) {\n buffer[3 * i + 0] = '|';\n buffer[3 * i + 1] = ' ';\n buffer[3 * i + 2] = ' ';\n }\n\n buffer[3 * depth] = '\\0';\n printf(\"%sBLOCK %d launched by thread %d of block %d\\n\", buffer, uid, thread, parent_uid);\n }\n }\n\n __syncthreads();\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// The kernel using CUDA dynamic parallelism.\n//\n// It generates a unique identifier for each block. Prints the information\n// about that block. Finally, if the 'max_depth' has not been reached, the\n// block launches new blocks directly from the GPU.\n////////////////////////////////////////////////////////////////////////////////\n__global__ void cdp_kernel(int max_depth, int depth, int thread, int parent_uid)\n{\n // We create a unique ID per block. Thread 0 does that and shares the value\n // with the other threads.\n __shared__ int s_uid;\n\n if (threadIdx.x == 0) {\n s_uid = atomicAdd(&g_uids, 1);\n }\n\n __syncthreads();\n\n // We print the ID of the block and information about its parent.\n print_info(depth, thread, s_uid, parent_uid);\n\n // We launch new blocks if we haven't reached the max_depth yet.\n if (++depth >= max_depth) {\n return;\n }\n\n cdp_kernel<<>>(max_depth, depth, threadIdx.x, s_uid);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Main entry point.\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"starting Simple Print (CUDA Dynamic Parallelism)\\n\");\n\n // Parse a few command-line arguments.\n int max_depth = 2;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\") || checkCmdLineFlag(argc, (const char **)argv, \"h\")) {\n printf(\"Usage: %s depth=\\t(where max_depth is a value between 1 \"\n \"and 8).\\n\",\n argv[0]);\n exit(EXIT_SUCCESS);\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"depth\")) {\n max_depth = getCmdLineArgumentInt(argc, (const char **)argv, \"depth\");\n\n if (max_depth < 1 || max_depth > 8) {\n printf(\"depth parameter has to be between 1 and 8\\n\");\n exit(EXIT_FAILURE);\n }\n }\n\n // Find/set the device.\n int device = -1;\n cudaDeviceProp deviceProp;\n device = findCudaDevice(argc, (const char **)argv);\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, device));\n\n if (!(deviceProp.major > 3 || (deviceProp.major == 3 && deviceProp.minor >= 5))) {\n printf(\"GPU %d - %s does not support CUDA Dynamic Parallelism\\n Exiting.\", device, deviceProp.name);\n exit(EXIT_WAIVED);\n }\n\n // Print a message describing what the sample does.\n printf(\"*********************************************************************\"\n \"******\\n\");\n printf(\"The CPU launches 2 blocks of 2 threads each. On the device each thread \"\n \"will\\n\");\n printf(\"launch 2 blocks of 2 threads each. The GPU we will do that \"\n \"recursively\\n\");\n printf(\"until it reaches max_depth=%d\\n\\n\", max_depth);\n printf(\"In total 2\");\n int num_blocks = 2, sum = 2;\n\n for (int i = 1; i < max_depth; ++i) {\n num_blocks *= 4;\n printf(\"+%d\", num_blocks);\n sum += num_blocks;\n }\n\n printf(\"=%d blocks are launched!!! (%d from the GPU)\\n\", sum, sum - 2);\n printf(\"************************************************************************\"\n \"***\\n\\n\");\n\n // Launch the kernel from the CPU.\n printf(\"Launching cdp_kernel() with CUDA Dynamic Parallelism:\\n\\n\");\n cdp_kernel<<<2, 2>>>(max_depth, 0, 0, -1);\n checkCudaErrors(cudaGetLastError());\n\n // Finalize.\n checkCudaErrors(cudaDeviceSynchronize());\n\n exit(EXIT_SUCCESS);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n\n////////////////////////////////////////////////////////////////////////////////\n// Variable on the GPU used to generate unique identifiers of blocks.\n////////////////////////////////////////////////////////////////////////////////\n__device__ int g_uids = 0;\n\n////////////////////////////////////////////////////////////////////////////////\n// Print a simple message to signal the block which is currently executing.\n////////////////////////////////////////////////////////////////////////////////\n__device__ void print_info(int depth, int thread, int uid, int parent_uid)\n{\n if (threadIdx.x == 0) {\n if (depth == 0)\n printf(\"BLOCK %d launched by the host\\n\", uid);\n else {\n char buffer[32];\n\n for (int i = 0; i < depth; ++i) {\n buffer[3 * i + 0] = '|';\n buffer[3 * i + 1] = ' ';\n buffer[3 * i + 2] = ' ';\n }\n\n buffer[3 * depth] = '\\0';\n printf(\"%sBLOCK %d launched by thread %d of block %d\\n\", buffer, uid, thread, parent_uid);\n }\n }\n\n __syncthreads();\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// The kernel using CUDA dynamic parallelism.\n//\n// It generates a unique identifier for each block. Prints the information\n// about that block. Finally, if the 'max_depth' has not been reached, the\n// block launches new blocks directly from the GPU.\n////////////////////////////////////////////////////////////////////////////////\n__global__ void cdp_kernel(int max_depth, int depth, int thread, int parent_uid)\n{\n // We create a unique ID per block. Thread 0 does that and shares the value\n // with the other threads.\n __shared__ int s_uid;\n\n if (threadIdx.x == 0) {\n s_uid = atomicAdd(&g_uids, 1);\n }\n\n __syncthreads();\n\n // We print the ID of the block and information about its parent.\n print_info(depth, thread, s_uid, parent_uid);\n\n // We launch new blocks if we haven't reached the max_depth yet.\n if (++depth >= max_depth) {\n return;\n }\n\n cdp_kernel<<>>(max_depth, depth, threadIdx.x, s_uid);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Main entry point.\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"starting Simple Print (CUDA Dynamic Parallelism)\\n\");\n\n // Parse a few command-line arguments.\n int max_depth = 2;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\") || checkCmdLineFlag(argc, (const char **)argv, \"h\")) {\n printf(\"Usage: %s depth=\\t(where max_depth is a value between 1 \"\n \"and 8).\\n\",\n argv[0]);\n exit(EXIT_SUCCESS);\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"depth\")) {\n max_depth = getCmdLineArgumentInt(argc, (const char **)argv, \"depth\");\n\n if (max_depth < 1 || max_depth > 8) {\n printf(\"depth parameter has to be between 1 and 8\\n\");\n exit(EXIT_FAILURE);\n }\n }\n\n // Find/set the device.\n int device = -1;\n hipDeviceProp_t deviceProp;\n device = findCudaDevice(argc, (const char **)argv);\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, device));\n\n if (!(deviceProp.major > 3 || (deviceProp.major == 3 && deviceProp.minor >= 5))) {\n printf(\"GPU %d - %s does not support CUDA Dynamic Parallelism\\n Exiting.\", device, deviceProp.name);\n exit(EXIT_WAIVED);\n }\n\n // Print a message describing what the sample does.\n printf(\"*********************************************************************\"\n \"******\\n\");\n printf(\"The CPU launches 2 blocks of 2 threads each. On the device each thread \"\n \"will\\n\");\n printf(\"launch 2 blocks of 2 threads each. The GPU we will do that \"\n \"recursively\\n\");\n printf(\"until it reaches max_depth=%d\\n\\n\", max_depth);\n printf(\"In total 2\");\n int num_blocks = 2, sum = 2;\n\n for (int i = 1; i < max_depth; ++i) {\n num_blocks *= 4;\n printf(\"+%d\", num_blocks);\n sum += num_blocks;\n }\n\n printf(\"=%d blocks are launched!!! (%d from the GPU)\\n\", sum, sum - 2);\n printf(\"************************************************************************\"\n \"***\\n\\n\");\n\n // Launch the kernel from the CPU.\n printf(\"Launching cdp_kernel() with CUDA Dynamic Parallelism:\\n\\n\");\n cdp_kernel<<<2, 2>>>(max_depth, 0, 0, -1);\n checkCudaErrors(hipGetLastError());\n\n // Finalize.\n checkCudaErrors(hipDeviceSynchronize());\n\n exit(EXIT_SUCCESS);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/cdpSimplePrint/cdpSimplePrint.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// System includes\n#include \n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#define NUM_GRAPHS 8\n#define THREADS_PER_BLOCK 512\n\nvoid printMemoryFootprint(int device)\n{\n size_t footprint;\n checkCudaErrors(cudaDeviceGetGraphMemAttribute(device, (cudaGraphMemAttributeType)0, &footprint));\n printf(\" FOOTPRINT: %lu bytes\\n\", footprint);\n}\n\nvoid prepareAllocParams(cudaMemAllocNodeParams *allocParams, size_t bytes, int device)\n{\n memset(allocParams, 0, sizeof(*allocParams));\n\n allocParams->bytesize = bytes;\n allocParams->poolProps.allocType = cudaMemAllocationTypePinned;\n allocParams->poolProps.location.id = device;\n allocParams->poolProps.location.type = cudaMemLocationTypeDevice;\n}\n\nvoid createVirtAddrReuseGraph(cudaGraphExec_t *graphExec, size_t bytes, int device)\n{\n cudaGraph_t graph;\n cudaGraphNode_t allocNodeA, allocNodeB, freeNodeA, freeNodeB;\n cudaMemAllocNodeParams allocParams;\n float *d_a, *d_b;\n\n checkCudaErrors(cudaGraphCreate(&graph, 0));\n prepareAllocParams(&allocParams, bytes, device);\n\n checkCudaErrors(cudaGraphAddMemAllocNode(&allocNodeA, graph, NULL, 0, &allocParams));\n d_a = (float *)allocParams.dptr;\n checkCudaErrors(cudaGraphAddMemFreeNode(&freeNodeA, graph, &allocNodeA, 1, (void *)d_a));\n\n // The dependency between the allocation of d_b and the free of d_a allows d_b\n // to reuse the same VA.\n checkCudaErrors(cudaGraphAddMemAllocNode(&allocNodeB, graph, &freeNodeA, 1, &allocParams));\n d_b = (float *)allocParams.dptr;\n\n if (d_a == d_b) {\n printf(\"Check confirms that d_a and d_b share a virtual address.\\n\");\n }\n else {\n printf(\"Check shows that d_a and d_b DO NOT share a virtual address.\\n\");\n }\n\n checkCudaErrors(cudaGraphAddMemFreeNode(&freeNodeB, graph, &allocNodeB, 1, (void *)d_b));\n\n checkCudaErrors(cudaGraphInstantiate(graphExec, graph, NULL, NULL, 0));\n checkCudaErrors(cudaGraphDestroy(graph));\n}\n\nvoid virtualAddressReuseSingleGraph(size_t bytes, int device)\n{\n cudaStream_t stream;\n cudaGraphExec_t graphExec;\n\n printf(\"================================\\n\");\n printf(\"Running virtual address reuse example.\\n\");\n printf(\"Sequential allocations & frees within a single graph enable CUDA to \"\n \"reuse virtual addresses.\\n\\n\");\n\n createVirtAddrReuseGraph(&graphExec, bytes, device);\n checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));\n\n checkCudaErrors(cudaGraphLaunch(graphExec, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n printMemoryFootprint(device);\n\n checkCudaErrors(cudaGraphExecDestroy(graphExec));\n checkCudaErrors(cudaStreamDestroy(stream));\n}\n\n// This is a kernel that does no real work but runs at least for a specified\n// number of clocks\n__global__ void clockBlock(clock_t clock_count)\n{\n unsigned int start_clock = (unsigned int)clock();\n\n clock_t clock_offset = 0;\n\n while (clock_offset < clock_count) {\n unsigned int end_clock = (unsigned int)clock();\n\n // The code below should work like\n // this (thanks to modular arithmetics):\n //\n // clock_offset = (clock_t) (end_clock > start_clock ?\n // end_clock - start_clock :\n // end_clock + (0xffffffffu - start_clock));\n //\n // Indeed, let m = 2^32 then\n // end - start = end + m - start (mod m).\n\n clock_offset = (clock_t)(end_clock - start_clock);\n }\n}\n\n// A pointer to the allocated device buffer is returned in dPtr so the caller\n// can compare virtual addresses. The kernel node is added to increase the\n// graph's runtime.\nvoid createSimpleAllocFreeGraph(cudaGraphExec_t *graphExec, float **dPtr, size_t bytes, int device)\n{\n cudaGraph_t graph;\n cudaGraphNode_t allocNodeA, freeNodeA, blockDeviceNode;\n cudaMemAllocNodeParams allocParams;\n cudaKernelNodeParams blockDeviceNodeParams = {0};\n int numElements = bytes / sizeof(float);\n float kernelTime = 5; // time for each thread to run in microseconds\n\n checkCudaErrors(cudaGraphCreate(&graph, 0));\n prepareAllocParams(&allocParams, bytes, device);\n\n checkCudaErrors(cudaGraphAddMemAllocNode(&allocNodeA, graph, NULL, 0, &allocParams));\n *dPtr = (float *)allocParams.dptr;\n\n int clockRate;\n checkCudaErrors(cudaDeviceGetAttribute(&clockRate, cudaDevAttrClockRate, device));\n clock_t time_clocks = (clock_t)((kernelTime / 1000.0) * clockRate);\n\n void *blockDeviceArgs[1] = {(void *)&time_clocks};\n\n size_t numBlocks = numElements / (size_t)THREADS_PER_BLOCK;\n blockDeviceNodeParams.gridDim = dim3(numBlocks, 1, 1);\n blockDeviceNodeParams.blockDim = dim3(THREADS_PER_BLOCK, 1, 1);\n blockDeviceNodeParams.sharedMemBytes = 0;\n blockDeviceNodeParams.extra = NULL;\n blockDeviceNodeParams.func = (void *)clockBlock;\n blockDeviceNodeParams.kernelParams = (void **)blockDeviceArgs;\n checkCudaErrors(cudaGraphAddKernelNode(&blockDeviceNode, graph, &allocNodeA, 1, &blockDeviceNodeParams));\n\n checkCudaErrors(cudaGraphAddMemFreeNode(&freeNodeA, graph, &blockDeviceNode, 1, (void *)*dPtr));\n\n checkCudaErrors(cudaGraphInstantiate(graphExec, graph, NULL, NULL, 0));\n checkCudaErrors(cudaGraphDestroy(graph));\n}\n\nvoid physicalMemoryReuseSingleStream(size_t bytes, int device)\n{\n cudaStream_t stream;\n cudaGraphExec_t graphExecs[NUM_GRAPHS];\n float *dPtrs[NUM_GRAPHS];\n bool virtualAddrDiffer = true;\n\n printf(\"================================\\n\");\n printf(\"Running physical memory reuse example.\\n\");\n printf(\"CUDA reuses the same physical memory for allocations from separate \"\n \"graphs when the allocation lifetimes don't overlap.\\n\\n\");\n\n for (int i = 0; i < NUM_GRAPHS; i++) {\n createSimpleAllocFreeGraph(&graphExecs[i], &dPtrs[i], bytes, device);\n }\n\n printf(\"Creating the graph execs does not reserve any physical memory.\\n\");\n printMemoryFootprint(device);\n\n checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));\n\n checkCudaErrors(cudaGraphLaunch(graphExecs[0], stream));\n printf(\"\\nThe first graph launched reserves the memory it needs.\\n\");\n printMemoryFootprint(device);\n\n checkCudaErrors(cudaGraphLaunch(graphExecs[0], stream));\n printf(\"A subsequent launch of the same graph in the same stream reuses the \"\n \"same physical memory. \");\n printf(\"Thus the memory footprint does not grow here.\\n\");\n printMemoryFootprint(device);\n\n printf(\"\\nSubsequent launches of other graphs in the same stream also reuse the \"\n \"physical memory. \");\n printf(\"Thus the memory footprint does not grow here.\\n\");\n for (int i = 1; i < NUM_GRAPHS; i++) {\n checkCudaErrors(cudaGraphLaunch(graphExecs[i], stream));\n printf(\"%02d: \", i);\n printMemoryFootprint(device);\n }\n\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n for (int i = 0; i < NUM_GRAPHS; i++) {\n for (int j = i + 1; j < NUM_GRAPHS; j++) {\n if (dPtrs[i] == dPtrs[j]) {\n virtualAddrDiffer = false;\n printf(\"Error: Graph exec %d and %d have the same virtual address!\\n\", i - 1, i);\n }\n }\n checkCudaErrors(cudaGraphExecDestroy(graphExecs[i]));\n }\n if (virtualAddrDiffer) {\n printf(\"\\nCheck confirms all graphs use a different virtual address.\\n\");\n }\n else {\n printf(\"\\nAll graphs do NOT use different virtual addresses. Exiting test.\\n\");\n exit(EXIT_FAILURE);\n }\n\n checkCudaErrors(cudaStreamDestroy(stream));\n}\n\nvoid simultaneousStreams(size_t bytes, int device)\n{\n cudaStream_t streams[NUM_GRAPHS];\n cudaGraphExec_t graphExecs[NUM_GRAPHS];\n float *dPtrs[NUM_GRAPHS];\n\n printf(\"================================\\n\");\n printf(\"Running simultaneous streams example.\\n\");\n printf(\"Graphs that can run concurrently need separate physical memory. \");\n printf(\"In this example, each graph launched in a separate stream increases the \"\n \"total memory footprint.\\n\\n\");\n\n printf(\"When launching a new graph, CUDA may reuse physical memory from a graph \"\n \"whose execution has already \");\n printf(\"finished -- even if the new graph is being launched in a different \"\n \"stream from the completed graph. \");\n printf(\"Therefore, a kernel node is added to the graphs to increase \"\n \"runtime.\\n\\n\");\n\n for (int i = 0; i < NUM_GRAPHS; i++) {\n createSimpleAllocFreeGraph(&graphExecs[i], &dPtrs[i], bytes, device);\n checkCudaErrors(cudaStreamCreateWithFlags(&streams[i], cudaStreamNonBlocking));\n }\n\n printf(\"Initial footprint:\\n\");\n printMemoryFootprint(device);\n\n printf(\"\\nEach graph launch in a seperate stream grows the memory footprint:\\n\");\n for (int i = 1; i < NUM_GRAPHS; i++) {\n checkCudaErrors(cudaGraphLaunch(graphExecs[i], streams[i]));\n printf(\"%02d: \", i);\n printMemoryFootprint(device);\n }\n\n for (int i = 0; i < NUM_GRAPHS; i++) {\n checkCudaErrors(cudaStreamSynchronize(streams[i]));\n checkCudaErrors(cudaGraphExecDestroy(graphExecs[i]));\n checkCudaErrors(cudaStreamDestroy(streams[i]));\n }\n}\n\nvoid createSimpleAllocNoFreeGraph(cudaGraphExec_t *graphExec, float **dPtr, size_t bytes, int device)\n{\n cudaGraph_t graph;\n cudaGraphNode_t allocNodeA;\n cudaMemAllocNodeParams allocParams;\n\n checkCudaErrors(cudaGraphCreate(&graph, 0));\n prepareAllocParams(&allocParams, bytes, device);\n\n checkCudaErrors(cudaGraphAddMemAllocNode(&allocNodeA, graph, NULL, 0, &allocParams));\n *dPtr = (float *)allocParams.dptr;\n\n checkCudaErrors(cudaGraphInstantiate(graphExec, graph, NULL, NULL, 0));\n checkCudaErrors(cudaGraphDestroy(graph));\n}\n\nvoid unfreedAllocations(size_t bytes, int device)\n{\n cudaStream_t stream;\n cudaGraphExec_t graphExecs[NUM_GRAPHS];\n float *dPtrs[NUM_GRAPHS];\n\n printf(\"================================\\n\");\n printf(\"Running unfreed streams example.\\n\");\n printf(\"CUDA cannot reuse phyiscal memory from graphs which do not free their \"\n \"allocations.\\n\\n\");\n\n for (int i = 0; i < NUM_GRAPHS; i++) {\n createSimpleAllocNoFreeGraph(&graphExecs[i], &dPtrs[i], bytes, device);\n }\n\n checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));\n\n printf(\"Despite being launched in the same stream, each graph launch grows the \"\n \"memory footprint. \");\n printf(\"Since the allocation is not freed, CUDA keeps the memory valid for \"\n \"use.\\n\");\n for (int i = 0; i < NUM_GRAPHS; i++) {\n checkCudaErrors(cudaGraphLaunch(graphExecs[i], stream));\n printf(\"%02d: \", i);\n printMemoryFootprint(device);\n }\n\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n checkCudaErrors(cudaDeviceGraphMemTrim(device));\n printf(\"\\nTrimming does not impact the memory footprint since the un-freed \"\n \"allocations are still holding onto the memory.\\n\");\n printMemoryFootprint(device);\n\n for (int i = 0; i < NUM_GRAPHS; i++) {\n checkCudaErrors(cudaFree(dPtrs[i]));\n }\n printf(\"\\nFreeing the allocations does not shrink the footprint.\\n\");\n printMemoryFootprint(device);\n\n checkCudaErrors(cudaDeviceGraphMemTrim(device));\n printf(\"\\nSince the allocations are now freed, trimming does reduce the \"\n \"footprint even when the graph execs are not yet destroyed.\\n\");\n printMemoryFootprint(device);\n\n for (int i = 0; i < NUM_GRAPHS; i++) {\n checkCudaErrors(cudaGraphExecDestroy(graphExecs[i]));\n }\n checkCudaErrors(cudaStreamDestroy(stream));\n}\n\nvoid cleanupMemory(int device)\n{\n checkCudaErrors(cudaDeviceGraphMemTrim(device));\n printf(\"\\nCleaning up example by trimming device memory.\\n\");\n printMemoryFootprint(device);\n printf(\"\\n\");\n}\n\nint main(int argc, char **argv)\n{\n size_t bytes = 64 * 1024 * 1024;\n int device = findCudaDevice(argc, (const char **)argv);\n\n int driverVersion = 0;\n int deviceSupportsMemoryPools = 0;\n\n cudaDriverGetVersion(&driverVersion);\n printf(\"Driver version is: %d.%d\\n\", driverVersion / 1000, (driverVersion % 100) / 10);\n\n if (driverVersion < 11040) {\n printf(\"Waiving execution as driver does not support Graph Memory Nodes\\n\");\n exit(EXIT_WAIVED);\n }\n\n cudaDeviceGetAttribute(&deviceSupportsMemoryPools, cudaDevAttrMemoryPoolsSupported, device);\n if (!deviceSupportsMemoryPools) {\n printf(\"Waiving execution as device does not support Memory Pools\\n\");\n exit(EXIT_WAIVED);\n }\n else {\n printf(\"Running sample.\\n\");\n }\n\n virtualAddressReuseSingleGraph(bytes, device);\n cleanupMemory(device);\n\n physicalMemoryReuseSingleStream(bytes, device);\n cleanupMemory(device);\n\n simultaneousStreams(bytes, device);\n cleanupMemory(device);\n\n unfreedAllocations(bytes, device);\n cleanupMemory(device);\n\n printf(\"================================\\n\");\n printf(\"Sample complete.\\n\");\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// System includes\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#define NUM_GRAPHS 8\n#define THREADS_PER_BLOCK 512\n\nvoid printMemoryFootprint(int device)\n{\n size_t footprint;\n checkCudaErrors(hipDeviceGetGraphMemAttribute(device, (hipGraphMemAttributeType)0, &footprint));\n printf(\" FOOTPRINT: %lu bytes\\n\", footprint);\n}\n\nvoid prepareAllocParams(hipMemAllocNodeParams *allocParams, size_t bytes, int device)\n{\n memset(allocParams, 0, sizeof(*allocParams));\n\n allocParams->bytesize = bytes;\n allocParams->poolProps.allocType = hipMemAllocationTypePinned;\n allocParams->poolProps.location.id = device;\n allocParams->poolProps.location.type = hipMemLocationTypeDevice;\n}\n\nvoid createVirtAddrReuseGraph(hipGraphExec_t *graphExec, size_t bytes, int device)\n{\n hipGraph_t graph;\n hipGraphNode_t allocNodeA, allocNodeB, freeNodeA, freeNodeB;\n hipMemAllocNodeParams allocParams;\n float *d_a, *d_b;\n\n checkCudaErrors(hipGraphCreate(&graph, 0));\n prepareAllocParams(&allocParams, bytes, device);\n\n checkCudaErrors(hipGraphAddMemAllocNode(&allocNodeA, graph, NULL, 0, &allocParams));\n d_a = (float *)allocParams.dptr;\n checkCudaErrors(hipGraphAddMemFreeNode(&freeNodeA, graph, &allocNodeA, 1, (void *)d_a));\n\n // The dependency between the allocation of d_b and the free of d_a allows d_b\n // to reuse the same VA.\n checkCudaErrors(hipGraphAddMemAllocNode(&allocNodeB, graph, &freeNodeA, 1, &allocParams));\n d_b = (float *)allocParams.dptr;\n\n if (d_a == d_b) {\n printf(\"Check confirms that d_a and d_b share a virtual address.\\n\");\n }\n else {\n printf(\"Check shows that d_a and d_b DO NOT share a virtual address.\\n\");\n }\n\n checkCudaErrors(hipGraphAddMemFreeNode(&freeNodeB, graph, &allocNodeB, 1, (void *)d_b));\n\n checkCudaErrors(hipGraphInstantiate(graphExec, graph, NULL, NULL, 0));\n checkCudaErrors(hipGraphDestroy(graph));\n}\n\nvoid virtualAddressReuseSingleGraph(size_t bytes, int device)\n{\n hipStream_t stream;\n hipGraphExec_t graphExec;\n\n printf(\"================================\\n\");\n printf(\"Running virtual address reuse example.\\n\");\n printf(\"Sequential allocations & frees within a single graph enable CUDA to \"\n \"reuse virtual addresses.\\n\\n\");\n\n createVirtAddrReuseGraph(&graphExec, bytes, device);\n checkCudaErrors(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));\n\n checkCudaErrors(hipGraphLaunch(graphExec, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n printMemoryFootprint(device);\n\n checkCudaErrors(hipGraphExecDestroy(graphExec));\n checkCudaErrors(hipStreamDestroy(stream));\n}\n\n// This is a kernel that does no real work but runs at least for a specified\n// number of clocks\n__global__ void clockBlock(clock_t clock_count)\n{\n unsigned int start_clock = (unsigned int)clock();\n\n clock_t clock_offset = 0;\n\n while (clock_offset < clock_count) {\n unsigned int end_clock = (unsigned int)clock();\n\n // The code below should work like\n // this (thanks to modular arithmetics):\n //\n // clock_offset = (clock_t) (end_clock > start_clock ?\n // end_clock - start_clock :\n // end_clock + (0xffffffffu - start_clock));\n //\n // Indeed, let m = 2^32 then\n // end - start = end + m - start (mod m).\n\n clock_offset = (clock_t)(end_clock - start_clock);\n }\n}\n\n// A pointer to the allocated device buffer is returned in dPtr so the caller\n// can compare virtual addresses. The kernel node is added to increase the\n// graph's runtime.\nvoid createSimpleAllocFreeGraph(hipGraphExec_t *graphExec, float **dPtr, size_t bytes, int device)\n{\n hipGraph_t graph;\n hipGraphNode_t allocNodeA, freeNodeA, blockDeviceNode;\n hipMemAllocNodeParams allocParams;\n hipKernelNodeParams blockDeviceNodeParams = {0};\n int numElements = bytes / sizeof(float);\n float kernelTime = 5; // time for each thread to run in microseconds\n\n checkCudaErrors(hipGraphCreate(&graph, 0));\n prepareAllocParams(&allocParams, bytes, device);\n\n checkCudaErrors(hipGraphAddMemAllocNode(&allocNodeA, graph, NULL, 0, &allocParams));\n *dPtr = (float *)allocParams.dptr;\n\n int clockRate;\n checkCudaErrors(hipDeviceGetAttribute(&clockRate, hipDeviceAttributeClockRate, device));\n clock_t time_clocks = (clock_t)((kernelTime / 1000.0) * clockRate);\n\n void *blockDeviceArgs[1] = {(void *)&time_clocks};\n\n size_t numBlocks = numElements / (size_t)THREADS_PER_BLOCK;\n blockDeviceNodeParams.gridDim = dim3(numBlocks, 1, 1);\n blockDeviceNodeParams.blockDim = dim3(THREADS_PER_BLOCK, 1, 1);\n blockDeviceNodeParams.sharedMemBytes = 0;\n blockDeviceNodeParams.extra = NULL;\n blockDeviceNodeParams.func = (void *)clockBlock;\n blockDeviceNodeParams.kernelParams = (void **)blockDeviceArgs;\n checkCudaErrors(hipGraphAddKernelNode(&blockDeviceNode, graph, &allocNodeA, 1, &blockDeviceNodeParams));\n\n checkCudaErrors(hipGraphAddMemFreeNode(&freeNodeA, graph, &blockDeviceNode, 1, (void *)*dPtr));\n\n checkCudaErrors(hipGraphInstantiate(graphExec, graph, NULL, NULL, 0));\n checkCudaErrors(hipGraphDestroy(graph));\n}\n\nvoid physicalMemoryReuseSingleStream(size_t bytes, int device)\n{\n hipStream_t stream;\n hipGraphExec_t graphExecs[NUM_GRAPHS];\n float *dPtrs[NUM_GRAPHS];\n bool virtualAddrDiffer = true;\n\n printf(\"================================\\n\");\n printf(\"Running physical memory reuse example.\\n\");\n printf(\"CUDA reuses the same physical memory for allocations from separate \"\n \"graphs when the allocation lifetimes don't overlap.\\n\\n\");\n\n for (int i = 0; i < NUM_GRAPHS; i++) {\n createSimpleAllocFreeGraph(&graphExecs[i], &dPtrs[i], bytes, device);\n }\n\n printf(\"Creating the graph execs does not reserve any physical memory.\\n\");\n printMemoryFootprint(device);\n\n checkCudaErrors(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));\n\n checkCudaErrors(hipGraphLaunch(graphExecs[0], stream));\n printf(\"\\nThe first graph launched reserves the memory it needs.\\n\");\n printMemoryFootprint(device);\n\n checkCudaErrors(hipGraphLaunch(graphExecs[0], stream));\n printf(\"A subsequent launch of the same graph in the same stream reuses the \"\n \"same physical memory. \");\n printf(\"Thus the memory footprint does not grow here.\\n\");\n printMemoryFootprint(device);\n\n printf(\"\\nSubsequent launches of other graphs in the same stream also reuse the \"\n \"physical memory. \");\n printf(\"Thus the memory footprint does not grow here.\\n\");\n for (int i = 1; i < NUM_GRAPHS; i++) {\n checkCudaErrors(hipGraphLaunch(graphExecs[i], stream));\n printf(\"%02d: \", i);\n printMemoryFootprint(device);\n }\n\n checkCudaErrors(hipStreamSynchronize(stream));\n\n for (int i = 0; i < NUM_GRAPHS; i++) {\n for (int j = i + 1; j < NUM_GRAPHS; j++) {\n if (dPtrs[i] == dPtrs[j]) {\n virtualAddrDiffer = false;\n printf(\"Error: Graph exec %d and %d have the same virtual address!\\n\", i - 1, i);\n }\n }\n checkCudaErrors(hipGraphExecDestroy(graphExecs[i]));\n }\n if (virtualAddrDiffer) {\n printf(\"\\nCheck confirms all graphs use a different virtual address.\\n\");\n }\n else {\n printf(\"\\nAll graphs do NOT use different virtual addresses. Exiting test.\\n\");\n exit(EXIT_FAILURE);\n }\n\n checkCudaErrors(hipStreamDestroy(stream));\n}\n\nvoid simultaneousStreams(size_t bytes, int device)\n{\n hipStream_t streams[NUM_GRAPHS];\n hipGraphExec_t graphExecs[NUM_GRAPHS];\n float *dPtrs[NUM_GRAPHS];\n\n printf(\"================================\\n\");\n printf(\"Running simultaneous streams example.\\n\");\n printf(\"Graphs that can run concurrently need separate physical memory. \");\n printf(\"In this example, each graph launched in a separate stream increases the \"\n \"total memory footprint.\\n\\n\");\n\n printf(\"When launching a new graph, CUDA may reuse physical memory from a graph \"\n \"whose execution has already \");\n printf(\"finished -- even if the new graph is being launched in a different \"\n \"stream from the completed graph. \");\n printf(\"Therefore, a kernel node is added to the graphs to increase \"\n \"runtime.\\n\\n\");\n\n for (int i = 0; i < NUM_GRAPHS; i++) {\n createSimpleAllocFreeGraph(&graphExecs[i], &dPtrs[i], bytes, device);\n checkCudaErrors(hipStreamCreateWithFlags(&streams[i], hipStreamNonBlocking));\n }\n\n printf(\"Initial footprint:\\n\");\n printMemoryFootprint(device);\n\n printf(\"\\nEach graph launch in a seperate stream grows the memory footprint:\\n\");\n for (int i = 1; i < NUM_GRAPHS; i++) {\n checkCudaErrors(hipGraphLaunch(graphExecs[i], streams[i]));\n printf(\"%02d: \", i);\n printMemoryFootprint(device);\n }\n\n for (int i = 0; i < NUM_GRAPHS; i++) {\n checkCudaErrors(hipStreamSynchronize(streams[i]));\n checkCudaErrors(hipGraphExecDestroy(graphExecs[i]));\n checkCudaErrors(hipStreamDestroy(streams[i]));\n }\n}\n\nvoid createSimpleAllocNoFreeGraph(hipGraphExec_t *graphExec, float **dPtr, size_t bytes, int device)\n{\n hipGraph_t graph;\n hipGraphNode_t allocNodeA;\n hipMemAllocNodeParams allocParams;\n\n checkCudaErrors(hipGraphCreate(&graph, 0));\n prepareAllocParams(&allocParams, bytes, device);\n\n checkCudaErrors(hipGraphAddMemAllocNode(&allocNodeA, graph, NULL, 0, &allocParams));\n *dPtr = (float *)allocParams.dptr;\n\n checkCudaErrors(hipGraphInstantiate(graphExec, graph, NULL, NULL, 0));\n checkCudaErrors(hipGraphDestroy(graph));\n}\n\nvoid unfreedAllocations(size_t bytes, int device)\n{\n hipStream_t stream;\n hipGraphExec_t graphExecs[NUM_GRAPHS];\n float *dPtrs[NUM_GRAPHS];\n\n printf(\"================================\\n\");\n printf(\"Running unfreed streams example.\\n\");\n printf(\"CUDA cannot reuse phyiscal memory from graphs which do not free their \"\n \"allocations.\\n\\n\");\n\n for (int i = 0; i < NUM_GRAPHS; i++) {\n createSimpleAllocNoFreeGraph(&graphExecs[i], &dPtrs[i], bytes, device);\n }\n\n checkCudaErrors(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));\n\n printf(\"Despite being launched in the same stream, each graph launch grows the \"\n \"memory footprint. \");\n printf(\"Since the allocation is not freed, CUDA keeps the memory valid for \"\n \"use.\\n\");\n for (int i = 0; i < NUM_GRAPHS; i++) {\n checkCudaErrors(hipGraphLaunch(graphExecs[i], stream));\n printf(\"%02d: \", i);\n printMemoryFootprint(device);\n }\n\n checkCudaErrors(hipStreamSynchronize(stream));\n\n checkCudaErrors(hipDeviceGraphMemTrim(device));\n printf(\"\\nTrimming does not impact the memory footprint since the un-freed \"\n \"allocations are still holding onto the memory.\\n\");\n printMemoryFootprint(device);\n\n for (int i = 0; i < NUM_GRAPHS; i++) {\n checkCudaErrors(hipFree(dPtrs[i]));\n }\n printf(\"\\nFreeing the allocations does not shrink the footprint.\\n\");\n printMemoryFootprint(device);\n\n checkCudaErrors(hipDeviceGraphMemTrim(device));\n printf(\"\\nSince the allocations are now freed, trimming does reduce the \"\n \"footprint even when the graph execs are not yet destroyed.\\n\");\n printMemoryFootprint(device);\n\n for (int i = 0; i < NUM_GRAPHS; i++) {\n checkCudaErrors(hipGraphExecDestroy(graphExecs[i]));\n }\n checkCudaErrors(hipStreamDestroy(stream));\n}\n\nvoid cleanupMemory(int device)\n{\n checkCudaErrors(hipDeviceGraphMemTrim(device));\n printf(\"\\nCleaning up example by trimming device memory.\\n\");\n printMemoryFootprint(device);\n printf(\"\\n\");\n}\n\nint main(int argc, char **argv)\n{\n size_t bytes = 64 * 1024 * 1024;\n int device = findCudaDevice(argc, (const char **)argv);\n\n int driverVersion = 0;\n int deviceSupportsMemoryPools = 0;\n\n hipDriverGetVersion(&driverVersion);\n printf(\"Driver version is: %d.%d\\n\", driverVersion / 1000, (driverVersion % 100) / 10);\n\n if (driverVersion < 11040) {\n printf(\"Waiving execution as driver does not support Graph Memory Nodes\\n\");\n exit(EXIT_WAIVED);\n }\n\n hipDeviceGetAttribute(&deviceSupportsMemoryPools, hipDeviceAttributeMemoryPoolsSupported, device);\n if (!deviceSupportsMemoryPools) {\n printf(\"Waiving execution as device does not support Memory Pools\\n\");\n exit(EXIT_WAIVED);\n }\n else {\n printf(\"Running sample.\\n\");\n }\n\n virtualAddressReuseSingleGraph(bytes, device);\n cleanupMemory(device);\n\n physicalMemoryReuseSingleStream(bytes, device);\n cleanupMemory(device);\n\n simultaneousStreams(bytes, device);\n cleanupMemory(device);\n\n unfreedAllocations(bytes, device);\n cleanupMemory(device);\n\n printf(\"================================\\n\");\n printf(\"Sample complete.\\n\");\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/graphMemoryFootprint/graphMemoryFootprint.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// CUDA sample demonstrating a tf32 (E8M10) GEMM computation using the Warp Matrix Multiply\n// and Accumulate API introduced in CUDA 11.0.\n\n// In this program, the compute_gemm kernel computes the result of a matrix multiplication\n// and addition: D = alpha * A * B + beta * C. The dimensions of both C and D matrices\n// are M_GLOBAL x N_GLOBAL. The A matrix is M_GLOBAL x K_GLOBAL (row-major), the B matrix\n// is K_GLOBAL x N_GLOBAL (column-major).\n// In that kernel, each CTA computes one 128 x 128 tile of the resulting matrix\n// per iteration. When the tile is computed, the CTA stores it to the global memory\n// and begins a new iteration, selecting a new 128 x 128 tile to compute.\n// Each CTA consists of eight warps. For the 128 x 128 tile, each warp computes eight\n// 16 x 16 subtiles, organized in a 2 x 4 two-dimensional array.\n// Warps compute the 16 x 16 subtiles using nvcuda::wmma::mma_sync operations by\n// moving through the K_GLOBAL dimension of the A and B matrices and accumulating\n// the intermediate result in the local thread state.\n\n// There are a number of simple optimizations used in the algorithm:\n// - The CTA copies the 128 x 128 tile of the C matrix from the global memory to\n// shared memory. After that is done, each warp loads the C matrix fragments from\n// shared memory, thus avoiding a random global memory access.\n// - On each internal iteration, the CTA copies a portion of the A and B matrices from\n// global memory to shared memory. After that, all warps in the CTA reuse the A and B\n// data from shared memory, thus reducing the number of data copies from global memory.\n// - The portions of the A and B matrices are stored in shared memory with an additional\n// padding (skew) to reduce the number of shared memory access bank conflicts.\n// (See a detailed explanation near the SKEW_FLOAT macro definition.)\n// - When the CTA finishes computing the tiles of the resulting matrix, each warp stores\n// its subtiles to shared memory. The CTA then copies the shared memory contents to\n// global memory, again avoiding redundant random global memory accesses.\n// - Note that the CTA tile size is chosen to maximize the GPU register utilization,\n// but carefully enough to avoid local memory use.\n\n#include \n#include \n#include \n#include \n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Externally configurable parameters.\n\n#ifndef CPU_DEBUG\n// Set this to 1 to verify the correctness of the GPU-computed matrix.\n#define CPU_DEBUG 0\n#endif\n\n#ifndef SHARED_MEMORY_LIMIT_64K\n// Set this to 0 to use more than 64 Kb of shared memory to cache data, to\n// improve the performance of the computations on GPU.\n// Note that you need a GPU that can have more than 64 Kb of shared memory\n// per multiprocessor.\n#define SHARED_MEMORY_LIMIT_64K 0\n#endif\n\n// GPU configuration.\n\n#define WARP_SIZE 32\n\n// MMA matrix tile dimensions.\n\n#define M 16\n#define N 16\n#define K 8\n\n// GEMM configuration.\n\n#define M_TILES 512\n#define N_TILES 512\n#define K_TILES 512\n\n#define M_GLOBAL (M * M_TILES)\n#define N_GLOBAL (N * N_TILES)\n#define K_GLOBAL (K * K_TILES)\n\n#define C_LAYOUT wmma::mem_row_major\n\n// Implementation constants.\n\n#define WARPS_PER_BLOCK 8\n#define THREADS_PER_BLOCK (WARP_SIZE * WARPS_PER_BLOCK)\n\n#if SHARED_MEMORY_LIMIT_64K\n// With only 64 Kb shared memory available, we can fit two 8-tile chunks of\n// the A and B matrix data, that is (M = 16) * (K = 8) * 8 * (CHUNK_K = 8)\n// * sizeof(float) = 32 Kb each.\n// (i.e. two 8x8 arrays of tiles of 16x8 float-typed elements per CTA).\n// But we cannot account the 8 Kb total skew overhead, without which the performance\n// would be severely impacted. So we choose to reduce the chunk size in half,\n// i.e. the amount of A and B matrix data we cache in shared memory.\n// Accordingly, this doubles the number of outer iterations across the global K\n// dimension, which only slightly impacts the performance.\n#define CHUNK_K 4\n#else\n#define CHUNK_K 8\n#endif\n\n#define CHUNK_LINE_BYTES (CHUNK_K * K * sizeof(float))\n#define WARP_COPY_BYTES (WARP_SIZE * sizeof(int4))\n#define CHUNK_COPY_LINES_PER_WARP (WARP_COPY_BYTES / CHUNK_LINE_BYTES)\n#define CHUNK_COPY_LINE_LANES (WARP_SIZE / CHUNK_COPY_LINES_PER_WARP)\n\n#define BLOCK_ROW_WARPS 2\n#define BLOCK_COL_WARPS 4\n\n#define WARP_ROW_TILES 4\n#define WARP_COL_TILES 2\n\n#define BLOCK_ROW_TILES (WARP_ROW_TILES * BLOCK_ROW_WARPS)\n#define BLOCK_COL_TILES (WARP_COL_TILES * BLOCK_COL_WARPS)\n\n#define GLOBAL_MEM_STRIDE N_GLOBAL\n\n#define SHMEM_STRIDE (N * BLOCK_ROW_TILES)\n#define SHMEM_OFFSET (N * WARP_ROW_TILES)\n\n// The macro below is used to shift rows of the A matrix and columns of the B matrix\n// in shared memory to minimize possible bank conflicts.\n// Before performing the nvcuda::wmma::mma_sync operation, the warp must load the matrix\n// data using the nvcuda::wmma::load_matrix_sync operation. Although the memory access pattern\n// is not specified for that function, each lane in the warp can read one or multiple matrix\n// elements from different matrix rows or columns.\n// For shared memory, such access can result in bank conflicts if different rows / columns\n// of the matrix map to the same bank. By shifting each row and column by a few bytes, we\n// make sure that they map to different banks, thus reducing the number of possible bank\n// conflicts.\n// The number of 8 four-byte \"float\" elements is chosen as the minimum possible shift because\n// we must keep each row and column 256-bit aligned, as required by nvcuda::wmma::load_matrix_sync.\n#define SKEW_FLOAT 8\n\n#define checkKernelErrors(expr) \\\n do { \\\n expr; \\\n \\\n cudaError_t __err = cudaGetLastError(); \\\n if (__err != cudaSuccess) { \\\n printf(\"Line %d: '%s' failed: %s\\n\", __LINE__, #expr, cudaGetErrorString(__err)); \\\n abort(); \\\n } \\\n } while (0)\n\nenum kernels {\n tf32mma_shmem_gemm_async_copy = 0, // tf32 MMA shmem using kernel with async_copy\n tf32mma_shmem_gemm = 1, // tf32 MMA shmem using kernel normal copy (without async_copy).\n simple_tf32mma_gemm = 2 // tf32 MMA non-shmem using simple kernel.\n};\n\nconst char *kernelNames[] = {\"compute_tf32gemm_async_copy\", \"compute_tf32gemm\", \"simple_wmma_tf32gemm\"};\n\nusing namespace nvcuda;\n\n__host__ void init_host_matrices(float *a, float *b, float *c)\n{\n for (int i = 0; i < M_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n a[i * K_GLOBAL + j] = (float)(rand() % 3);\n }\n }\n\n for (int i = 0; i < N_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n b[i * K_GLOBAL + j] = (float)(rand() % 3);\n }\n }\n\n for (int t = 0; t < M_GLOBAL * N_GLOBAL; t++) {\n c[t] = (float)(rand() % 3);\n }\n}\n\n__global__ void compute_tf32gemm(const float *A, const float *B, const float *C, float *D, float alpha, float beta)\n{\n#if __CUDA_ARCH__ >= 800\n extern __shared__ float shmem[][CHUNK_K * K + SKEW_FLOAT];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n const size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n float *shmem_warp_tile_ptr = (float *)&shmem[0][0] + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS\n + (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.\n float *shmem_warp_stream_ptr = (float *)&shmem[0][0] + warpId * SHMEM_STRIDE * N;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may result\n // in a loss of precision). Zero still needs to be specially handled though.\n beta /= alpha;\n\n // Each CTA slides along the 128 x 128 tiles from the top left corner of the matrix to the\n // right and down, and selects the next tile to compute. Once there's no such tile,\n // all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const float *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId) =\n *((int4 *)(src_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment multiplications\n // along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Scale the C matrix.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const float *warp_ptr =\n (warpId < (WARPS_PER_BLOCK / 2))\n ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);\n\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.\n size_t shmem_idx = warpId < (WARPS_PER_BLOCK / 2)\n ? (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (N * (warpId % (WARPS_PER_BLOCK / 2)) * 2 + shmem_idx_b_off);\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n const float *lane_ptr = (warp_ptr + tile_k * K + (laneId / CHUNK_COPY_LINE_LANES) * K_GLOBAL);\n\n // Shift the second half of the warp to the next row / column in the shared memory.\n shmem_idx += laneId / CHUNK_COPY_LINE_LANES;\n\n#pragma unroll\n for (int i = 0; i < ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP) * 2; i++) {\n // Copy 16 bytes at once in each lane.\n *((int4 *)&shmem[shmem_idx][0] + (laneId % CHUNK_COPY_LINE_LANES)) =\n *((int4 *)lane_ptr + (laneId % CHUNK_COPY_LINE_LANES));\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / BLOCK_ROW_WARPS) * M * BLOCK_ROW_WARPS + (i * M);\n const float *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_FLOAT);\n#pragma unroll\n for (int t = 0; t < a[i].num_elements; t++) {\n a[i].x[t] = wmma::__float_to_tf32(a[i].x[t]);\n }\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be reused\n // against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const float *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_FLOAT);\n#pragma unroll\n for (int t = 0; t < b[j].num_elements; t++) {\n b[j].x[t] = wmma::__float_to_tf32(b[j].x[t]);\n }\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n // Uniform, point-wise transformations of ALL fragment elements by ALL threads in the\n // warp are well-defined even though element indices within fragment storage are not defined.\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global memory.\n float *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n#endif\n}\n\n__global__ void\ncompute_tf32gemm_async_copy(const float *A, const float *B, const float *C, float *D, const float alpha, float beta)\n{\n#if __CUDA_ARCH__ >= 800\n extern __shared__ float shmem[][CHUNK_K * K + SKEW_FLOAT];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n float *shmem_warp_tile_ptr = (float *)&shmem[0][0] + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS\n + (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.\n float *shmem_warp_stream_ptr = (float *)&shmem[0][0] + warpId * SHMEM_STRIDE * N;\n\n // Offset in shared memory from which the B matrix is stored.\n constexpr size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may result\n // in a loss of precision). Zero still needs to be specially handled though.\n beta /= alpha;\n\n cuda::pipeline pipe = cuda::make_pipeline();\n const auto shape4 = cuda::aligned_size_t(sizeof(float4));\n constexpr int loadStride = 2; // load 4 floats, so left-shift by 2.\n\n // Each CTA slides along the 128 x 128 tiles from the top left corner of the matrix to the\n // right and down, and selects the next tile to compute. Once there's no such tile,\n // all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const float *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < N; i++) {\n pipe.producer_acquire();\n cuda::memcpy_async(&shmem_warp_stream_ptr[(SHMEM_STRIDE * i) + (laneId << loadStride)],\n &src_gmem_warp_stream_ptr[(GLOBAL_MEM_STRIDE * i) + (laneId << loadStride)],\n shape4,\n pipe);\n pipe.producer_commit();\n }\n // Now wait for all the above issued 8 batches to complete.\n cuda::pipeline_consumer_wait_prior<0>(pipe);\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment multiplications\n // along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n // Scale the C matrix.\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n pipe.consumer_release();\n\n // sync here so that shared memory can then be used for loading A & B matrices.\n __syncthreads();\n\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const float *warp_ptr =\n (warpId < (WARPS_PER_BLOCK / 2))\n ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);\n\n constexpr int chunksPerLane = ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP) * 2;\n const int laneLoadElem = (laneId % CHUNK_COPY_LINE_LANES) << loadStride;\n const int stridePerLaneCopy = (laneId / CHUNK_COPY_LINE_LANES);\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.\n // As for tf32 MMA M == N we use M for warp 4-7 + shmem_idx_b_off.\n size_t shmem_idx =\n (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2) + ((warpId / (WARPS_PER_BLOCK / 2)) * shmem_idx_b_off);\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n const float *lane_ptr = (warp_ptr + tile_k * K + stridePerLaneCopy * K_GLOBAL + laneLoadElem);\n\n // Shift the second half of the warp to the next row / column in the shared memory.\n shmem_idx += stridePerLaneCopy;\n\n#pragma unroll\n for (int i = 0; i < chunksPerLane; i++) {\n // Copy 16 bytes at once in each lane.\n pipe.producer_acquire();\n cuda::memcpy_async(&shmem[shmem_idx][laneLoadElem], lane_ptr, shape4, pipe);\n pipe.producer_commit();\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n\n cuda::pipeline_consumer_wait_prior<0>(pipe);\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / BLOCK_ROW_WARPS) * M * BLOCK_ROW_WARPS + (i * M);\n const float *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_FLOAT);\n\n#pragma unroll\n for (int t = 0; t < a[i].num_elements; t++) {\n a[i].x[t] = wmma::__float_to_tf32(a[i].x[t]);\n }\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be reused\n // against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const float *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_FLOAT);\n#pragma unroll\n for (int t = 0; t < b[j].num_elements; t++) {\n b[j].x[t] = wmma::__float_to_tf32(b[j].x[t]);\n }\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n pipe.consumer_release();\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n // Uniform, point-wise transformations of ALL fragment elements by ALL threads in the\n // warp are well-defined even though element indices within fragment storage are not defined.\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global memory.\n float *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((float4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((float4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n#endif\n}\n\n// Performs an MxNxK tf32 GEMM (C=alpha*A*B + beta*C) assuming:\n// 1) Matrices are packed in memory.\n// 2) M, N and K are multiples of 16, 16 and 8 respectively.\n// 3) A is row major, B is column major matrix.\n// Note: This is a less performant version of the compute_tf32gemm kernel. It is designed for\n// demonstration purposes only to show the CUDA WMMA API use without relying on\n// availability of the shared memory.\n__global__ void\nsimple_wmma_tf32gemm(float *a, float *b, float *c, float *d, int m_ld, int n_ld, int k_ld, float alpha, float beta)\n{\n#if __CUDA_ARCH__ >= 800\n // Leading dimensions. Packed with no transpositions.\n int lda = k_ld;\n int ldb = k_ld;\n int ldc = n_ld;\n\n // Tile using a 2D grid\n int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;\n int warpN = (blockIdx.y * blockDim.y + threadIdx.y);\n\n // Declare the fragments\n wmma::fragment a_frag;\n wmma::fragment b_frag;\n wmma::fragment acc_frag;\n wmma::fragment c_frag;\n\n wmma::fill_fragment(acc_frag, 0.0f);\n\n // Loop over k\n for (int i = 0; i < k_ld; i += K) {\n int aCol = i;\n int aRow = warpM * M;\n\n // int bCol = i;\n // int bRow = warpN * N;\n int bCol = warpN * N;\n int bRow = i;\n\n // Bounds checking\n if (aRow < m_ld && aCol < k_ld && bRow < k_ld && bCol < n_ld) {\n // Load the inputs\n wmma::load_matrix_sync(a_frag, a + aCol + aRow * lda, lda);\n wmma::load_matrix_sync(b_frag, b + bRow + bCol * ldb, ldb);\n\n#pragma unroll\n for (int t = 0; t < a_frag.num_elements; t++) {\n a_frag.x[t] = wmma::__float_to_tf32(a_frag.x[t]);\n }\n\n#pragma unroll\n for (int t = 0; t < b_frag.num_elements; t++) {\n b_frag.x[t] = wmma::__float_to_tf32(b_frag.x[t]);\n }\n // Perform the matrix multiplication\n wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);\n }\n }\n\n // Load in the current value of c, scale it by beta, and add this our result scaled by alpha\n int cCol = warpN * N;\n int cRow = warpM * M;\n\n if (cRow < m_ld && cCol < n_ld) {\n wmma::load_matrix_sync(c_frag, c + cCol + cRow * ldc, ldc, wmma::mem_row_major);\n\n for (int i = 0; i < c_frag.num_elements; i++) {\n c_frag.x[i] = alpha * acc_frag.x[i] + beta * c_frag.x[i];\n }\n\n // Store the output\n wmma::store_matrix_sync(d + cCol + cRow * ldc, c_frag, ldc, wmma::mem_row_major);\n }\n#endif\n}\n\n__host__ void matMultiplyOnHost(float *A,\n float *B,\n float *C,\n float alpha,\n float beta,\n int numARows,\n int numAColumns,\n int numBRows,\n int numBColumns,\n int numCRows,\n int numCColumns)\n{\n for (int i = 0; i < numCRows; i++) {\n for (int j = 0; j < numCColumns; j++) {\n float temp = 0.0;\n\n for (int k = 0; k < numAColumns; k++) {\n temp += A[i * numAColumns + k] * B[j * numBRows + k];\n }\n\n C[i * numCColumns + j] = temp * alpha + beta * C[i * numCColumns + j];\n }\n }\n}\n\nint main(int argc, char **argv)\n{\n printf(\"Initializing...\\n\");\n\n int dev = findCudaDevice(argc, (const char **)argv);\n\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, dev));\n\n // Tensor cores require a GPU of Volta (SM8X) architecture or higher.\n if (deviceProp.major < 8) {\n printf(\"tf32TensorCoreGemm requires requires SM 8.0 or higher to use Tensor Cores. Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n printf(\"M: %d (%d x %d)\\n\", M_GLOBAL, M, M_TILES);\n printf(\"N: %d (%d x %d)\\n\", N_GLOBAL, N, N_TILES);\n printf(\"K: %d (%d x %d)\\n\", K_GLOBAL, K, K_TILES);\n\n float *A_h = NULL;\n float *B_h = NULL;\n float *C_h = NULL;\n#if CPU_DEBUG\n float *result_hD = NULL;\n float *result_host = NULL;\n#endif\n\n A_h = (float *)malloc(sizeof(float) * M_GLOBAL * K_GLOBAL);\n B_h = (float *)malloc(sizeof(float) * K_GLOBAL * N_GLOBAL);\n C_h = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n#if CPU_DEBUG\n result_hD = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n result_host = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n#endif\n\n float *A = NULL;\n float *B = NULL;\n float *C = NULL;\n float *D = NULL;\n\n checkCudaErrors(cudaMalloc((void **)&A, sizeof(float) * M_GLOBAL * K_GLOBAL));\n checkCudaErrors(cudaMalloc((void **)&B, sizeof(float) * N_GLOBAL * K_GLOBAL));\n checkCudaErrors(cudaMalloc((void **)&C, sizeof(float) * M_GLOBAL * N_GLOBAL));\n checkCudaErrors(cudaMalloc((void **)&D, sizeof(float) * M_GLOBAL * N_GLOBAL));\n\n assert(((unsigned long long)A) % 128 == 0);\n assert(((unsigned long long)B) % 128 == 0);\n assert(((unsigned long long)C) % 128 == 0);\n assert(((unsigned long long)D) % 128 == 0);\n\n init_host_matrices(A_h, B_h, C_h);\n\n printf(\"Preparing data for GPU...\\n\");\n\n checkCudaErrors(cudaMemcpy(A, A_h, sizeof(float) * M_GLOBAL * K_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(B, B_h, sizeof(float) * N_GLOBAL * K_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(C, C_h, sizeof(float) * M_GLOBAL * N_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemset(D, 0, sizeof(float) * M_GLOBAL * N_GLOBAL));\n\n enum {\n // Compute the right amount of shared memory to request.\n // We need shared memory to hold per-CTA C and D matrix tiles, and to cache per-CTA chunks\n // of the A and B matrices. Therefore, the right amount to request is the maximum of those\n // two numbers.\n SHMEM_SZ = MAX(sizeof(float) * (BLOCK_COL_TILES * M) * (CHUNK_K * K + SKEW_FLOAT) * 2,\n M * (BLOCK_ROW_WARPS * WARP_ROW_TILES) * N * (BLOCK_COL_WARPS * WARP_COL_TILES) * sizeof(float))\n };\n\n printf(\"Required shared memory size: %lu Kb\\n\", SHMEM_SZ / 1024UL);\n\n const float alpha = 1.1f;\n const float beta = 1.2f;\n\n cudaEvent_t start, stop;\n\n checkCudaErrors(cudaEventCreate(&start));\n checkCudaErrors(cudaEventCreate(&stop));\n checkCudaErrors(cudaEventRecord(start));\n\n // kernel to run - default (tf32mma_shmem_gemm_async_copy == 0)\n kernels selected_kernel = tf32mma_shmem_gemm_async_copy;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"kernel\")) {\n int kernel_number = getCmdLineArgumentInt(argc, (const char **)argv, \"kernel\");\n if (kernel_number < 3) {\n selected_kernel = (kernels)kernel_number;\n }\n else {\n printf(\"Error: kernel number should be between 0 to 2, you have entered %d\\n\", kernel_number);\n exit(EXIT_FAILURE);\n }\n }\n\n // If enough shared memory available on the GPU use high performant kernel\n if ((deviceProp.sharedMemPerMultiprocessor >= SHMEM_SZ) && (selected_kernel != simple_tf32mma_gemm)) {\n printf(\"Computing using high performance kernel = %d - %s\\n\", selected_kernel, kernelNames[selected_kernel]);\n\n switch (selected_kernel) {\n case tf32mma_shmem_gemm_async_copy:\n default:\n checkCudaErrors(cudaFuncSetAttribute(\n compute_tf32gemm_async_copy, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors(\n (compute_tf32gemm_async_copy<<>>(\n A, B, C, D, alpha, beta)));\n break;\n case tf32mma_shmem_gemm:\n checkCudaErrors(\n cudaFuncSetAttribute(compute_tf32gemm, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors((compute_tf32gemm<<>>(\n A, B, C, D, alpha, beta)));\n break;\n }\n#if CPU_DEBUG\n checkCudaErrors(cudaMemcpy(result_hD, D, sizeof(float) * M_GLOBAL * N_GLOBAL, cudaMemcpyDeviceToHost));\n#endif\n }\n else {\n dim3 gridDim;\n dim3 blockDim;\n\n // blockDim.x must be a multple of warpSize\n // 128x4 means we have 16 warps and a block computes a 64x64 output tile\n blockDim.x = 128;\n blockDim.y = 4;\n\n gridDim.x = (M_GLOBAL + (M * blockDim.x / 32 - 1)) / (M * blockDim.x / 32);\n gridDim.y = (N_GLOBAL + N * blockDim.y - 1) / (N * blockDim.y);\n\n printf(\"Computing... using simple_wmma_gemm kernel\\n\");\n simple_wmma_tf32gemm<<>>(A, B, C, D, M_GLOBAL, N_GLOBAL, K_GLOBAL, alpha, beta);\n#if CPU_DEBUG\n checkCudaErrors(cudaMemcpy(result_hD, D, sizeof(float) * M_GLOBAL * N_GLOBAL, cudaMemcpyDeviceToHost));\n#endif\n }\n\n checkCudaErrors(cudaEventRecord(stop));\n checkCudaErrors(cudaEventSynchronize(stop));\n\n#if CPU_DEBUG\n printf(\"Verifying correctness of the computations...\\n\");\n\n memcpy(result_host, C_h, sizeof(float) * M_GLOBAL * N_GLOBAL);\n\n matMultiplyOnHost(A_h, B_h, result_host, alpha, beta, M_GLOBAL, K_GLOBAL, K_GLOBAL, N_GLOBAL, M_GLOBAL, N_GLOBAL);\n\n for (int i = 0; i < N_GLOBAL * M_GLOBAL; i++) {\n if (fabs(result_hD[i] - result_host[i]) > 0.1f) {\n printf(\"mismatch i=%d result_hD=%f result_host=%f\\n\", i, result_hD[i], result_host[i]);\n }\n }\n free(result_hD);\n free(result_host);\n#endif\n\n float milliseconds = 0;\n\n checkCudaErrors(cudaEventElapsedTime(&milliseconds, start, stop));\n\n printf(\"Time: %f ms\\n\", milliseconds);\n printf(\"TFLOPS: %.2f\\n\", (((double)M_GLOBAL * N_GLOBAL * K_GLOBAL * 2) / (milliseconds / 1000.)) / 1e12);\n\n free(A_h);\n free(B_h);\n free(C_h);\n checkCudaErrors(cudaFree((void *)A));\n checkCudaErrors(cudaFree((void *)B));\n checkCudaErrors(cudaFree((void *)C));\n checkCudaErrors(cudaFree((void *)D));\n\n return 0;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// CUDA sample demonstrating a tf32 (E8M10) GEMM computation using the Warp Matrix Multiply\n// and Accumulate API introduced in CUDA 11.0.\n\n// In this program, the compute_gemm kernel computes the result of a matrix multiplication\n// and addition: D = alpha * A * B + beta * C. The dimensions of both C and D matrices\n// are M_GLOBAL x N_GLOBAL. The A matrix is M_GLOBAL x K_GLOBAL (row-major), the B matrix\n// is K_GLOBAL x N_GLOBAL (column-major).\n// In that kernel, each CTA computes one 128 x 128 tile of the resulting matrix\n// per iteration. When the tile is computed, the CTA stores it to the global memory\n// and begins a new iteration, selecting a new 128 x 128 tile to compute.\n// Each CTA consists of eight warps. For the 128 x 128 tile, each warp computes eight\n// 16 x 16 subtiles, organized in a 2 x 4 two-dimensional array.\n// Warps compute the 16 x 16 subtiles using nvcuda::wmma::mma_sync operations by\n// moving through the K_GLOBAL dimension of the A and B matrices and accumulating\n// the intermediate result in the local thread state.\n\n// There are a number of simple optimizations used in the algorithm:\n// - The CTA copies the 128 x 128 tile of the C matrix from the global memory to\n// shared memory. After that is done, each warp loads the C matrix fragments from\n// shared memory, thus avoiding a random global memory access.\n// - On each internal iteration, the CTA copies a portion of the A and B matrices from\n// global memory to shared memory. After that, all warps in the CTA reuse the A and B\n// data from shared memory, thus reducing the number of data copies from global memory.\n// - The portions of the A and B matrices are stored in shared memory with an additional\n// padding (skew) to reduce the number of shared memory access bank conflicts.\n// (See a detailed explanation near the SKEW_FLOAT macro definition.)\n// - When the CTA finishes computing the tiles of the resulting matrix, each warp stores\n// its subtiles to shared memory. The CTA then copies the shared memory contents to\n// global memory, again avoiding redundant random global memory accesses.\n// - Note that the CTA tile size is chosen to maximize the GPU register utilization,\n// but carefully enough to avoid local memory use.\n\n#include \n#include \n#include \n#include \n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Externally configurable parameters.\n\n#ifndef CPU_DEBUG\n// Set this to 1 to verify the correctness of the GPU-computed matrix.\n#define CPU_DEBUG 0\n#endif\n\n#ifndef SHARED_MEMORY_LIMIT_64K\n// Set this to 0 to use more than 64 Kb of shared memory to cache data, to\n// improve the performance of the computations on GPU.\n// Note that you need a GPU that can have more than 64 Kb of shared memory\n// per multiprocessor.\n#define SHARED_MEMORY_LIMIT_64K 0\n#endif\n\n// GPU configuration.\n\n#define WARP_SIZE 32\n\n// MMA matrix tile dimensions.\n\n#define M 16\n#define N 16\n#define K 8\n\n// GEMM configuration.\n\n#define M_TILES 512\n#define N_TILES 512\n#define K_TILES 512\n\n#define M_GLOBAL (M * M_TILES)\n#define N_GLOBAL (N * N_TILES)\n#define K_GLOBAL (K * K_TILES)\n\n#define C_LAYOUT wmma::mem_row_major\n\n// Implementation constants.\n\n#define WARPS_PER_BLOCK 8\n#define THREADS_PER_BLOCK (WARP_SIZE * WARPS_PER_BLOCK)\n\n#if SHARED_MEMORY_LIMIT_64K\n// With only 64 Kb shared memory available, we can fit two 8-tile chunks of\n// the A and B matrix data, that is (M = 16) * (K = 8) * 8 * (CHUNK_K = 8)\n// * sizeof(float) = 32 Kb each.\n// (i.e. two 8x8 arrays of tiles of 16x8 float-typed elements per CTA).\n// But we cannot account the 8 Kb total skew overhead, without which the performance\n// would be severely impacted. So we choose to reduce the chunk size in half,\n// i.e. the amount of A and B matrix data we cache in shared memory.\n// Accordingly, this doubles the number of outer iterations across the global K\n// dimension, which only slightly impacts the performance.\n#define CHUNK_K 4\n#else\n#define CHUNK_K 8\n#endif\n\n#define CHUNK_LINE_BYTES (CHUNK_K * K * sizeof(float))\n#define WARP_COPY_BYTES (WARP_SIZE * sizeof(int4))\n#define CHUNK_COPY_LINES_PER_WARP (WARP_COPY_BYTES / CHUNK_LINE_BYTES)\n#define CHUNK_COPY_LINE_LANES (WARP_SIZE / CHUNK_COPY_LINES_PER_WARP)\n\n#define BLOCK_ROW_WARPS 2\n#define BLOCK_COL_WARPS 4\n\n#define WARP_ROW_TILES 4\n#define WARP_COL_TILES 2\n\n#define BLOCK_ROW_TILES (WARP_ROW_TILES * BLOCK_ROW_WARPS)\n#define BLOCK_COL_TILES (WARP_COL_TILES * BLOCK_COL_WARPS)\n\n#define GLOBAL_MEM_STRIDE N_GLOBAL\n\n#define SHMEM_STRIDE (N * BLOCK_ROW_TILES)\n#define SHMEM_OFFSET (N * WARP_ROW_TILES)\n\n// The macro below is used to shift rows of the A matrix and columns of the B matrix\n// in shared memory to minimize possible bank conflicts.\n// Before performing the nvcuda::wmma::mma_sync operation, the warp must load the matrix\n// data using the nvcuda::wmma::load_matrix_sync operation. Although the memory access pattern\n// is not specified for that function, each lane in the warp can read one or multiple matrix\n// elements from different matrix rows or columns.\n// For shared memory, such access can result in bank conflicts if different rows / columns\n// of the matrix map to the same bank. By shifting each row and column by a few bytes, we\n// make sure that they map to different banks, thus reducing the number of possible bank\n// conflicts.\n// The number of 8 four-byte \"float\" elements is chosen as the minimum possible shift because\n// we must keep each row and column 256-bit aligned, as required by nvcuda::wmma::load_matrix_sync.\n#define SKEW_FLOAT 8\n\n#define checkKernelErrors(expr) \\\n do { \\\n expr; \\\n \\\n hipError_t __err = hipGetLastError(); \\\n if (__err != hipSuccess) { \\\n printf(\"Line %d: '%s' failed: %s\\n\", __LINE__, #expr, hipGetErrorString(__err)); \\\n abort(); \\\n } \\\n } while (0)\n\nenum kernels {\n tf32mma_shmem_gemm_async_copy = 0, // tf32 MMA shmem using kernel with async_copy\n tf32mma_shmem_gemm = 1, // tf32 MMA shmem using kernel normal copy (without async_copy).\n simple_tf32mma_gemm = 2 // tf32 MMA non-shmem using simple kernel.\n};\n\nconst char *kernelNames[] = {\"compute_tf32gemm_async_copy\", \"compute_tf32gemm\", \"simple_wmma_tf32gemm\"};\n\nusing namespace nvcuda;\n\n__host__ void init_host_matrices(float *a, float *b, float *c)\n{\n for (int i = 0; i < M_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n a[i * K_GLOBAL + j] = (float)(rand() % 3);\n }\n }\n\n for (int i = 0; i < N_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n b[i * K_GLOBAL + j] = (float)(rand() % 3);\n }\n }\n\n for (int t = 0; t < M_GLOBAL * N_GLOBAL; t++) {\n c[t] = (float)(rand() % 3);\n }\n}\n\n__global__ void compute_tf32gemm(const float *A, const float *B, const float *C, float *D, float alpha, float beta)\n{\n#if __CUDA_ARCH__ >= 800\n extern __shared__ float shmem[][CHUNK_K * K + SKEW_FLOAT];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n const size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n float *shmem_warp_tile_ptr = (float *)&shmem[0][0] + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS\n + (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.\n float *shmem_warp_stream_ptr = (float *)&shmem[0][0] + warpId * SHMEM_STRIDE * N;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may result\n // in a loss of precision). Zero still needs to be specially handled though.\n beta /= alpha;\n\n // Each CTA slides along the 128 x 128 tiles from the top left corner of the matrix to the\n // right and down, and selects the next tile to compute. Once there's no such tile,\n // all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const float *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId) =\n *((int4 *)(src_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment multiplications\n // along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Scale the C matrix.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const float *warp_ptr =\n (warpId < (WARPS_PER_BLOCK / 2))\n ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);\n\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.\n size_t shmem_idx = warpId < (WARPS_PER_BLOCK / 2)\n ? (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (N * (warpId % (WARPS_PER_BLOCK / 2)) * 2 + shmem_idx_b_off);\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n const float *lane_ptr = (warp_ptr + tile_k * K + (laneId / CHUNK_COPY_LINE_LANES) * K_GLOBAL);\n\n // Shift the second half of the warp to the next row / column in the shared memory.\n shmem_idx += laneId / CHUNK_COPY_LINE_LANES;\n\n#pragma unroll\n for (int i = 0; i < ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP) * 2; i++) {\n // Copy 16 bytes at once in each lane.\n *((int4 *)&shmem[shmem_idx][0] + (laneId % CHUNK_COPY_LINE_LANES)) =\n *((int4 *)lane_ptr + (laneId % CHUNK_COPY_LINE_LANES));\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / BLOCK_ROW_WARPS) * M * BLOCK_ROW_WARPS + (i * M);\n const float *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_FLOAT);\n#pragma unroll\n for (int t = 0; t < a[i].num_elements; t++) {\n a[i].x[t] = wmma::__float_to_tf32(a[i].x[t]);\n }\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be reused\n // against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const float *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_FLOAT);\n#pragma unroll\n for (int t = 0; t < b[j].num_elements; t++) {\n b[j].x[t] = wmma::__float_to_tf32(b[j].x[t]);\n }\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n // Uniform, point-wise transformations of ALL fragment elements by ALL threads in the\n // warp are well-defined even though element indices within fragment storage are not defined.\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global memory.\n float *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n#endif\n}\n\n__global__ void\ncompute_tf32gemm_async_copy(const float *A, const float *B, const float *C, float *D, const float alpha, float beta)\n{\n#if __CUDA_ARCH__ >= 800\n extern __shared__ float shmem[][CHUNK_K * K + SKEW_FLOAT];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n float *shmem_warp_tile_ptr = (float *)&shmem[0][0] + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS\n + (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.\n float *shmem_warp_stream_ptr = (float *)&shmem[0][0] + warpId * SHMEM_STRIDE * N;\n\n // Offset in shared memory from which the B matrix is stored.\n constexpr size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may result\n // in a loss of precision). Zero still needs to be specially handled though.\n beta /= alpha;\n\n cuda::pipeline pipe = cuda::make_pipeline();\n const auto shape4 = cuda::aligned_size_t(sizeof(float4));\n constexpr int loadStride = 2; // load 4 floats, so left-shift by 2.\n\n // Each CTA slides along the 128 x 128 tiles from the top left corner of the matrix to the\n // right and down, and selects the next tile to compute. Once there's no such tile,\n // all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const float *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < N; i++) {\n pipe.producer_acquire();\n cuda::memcpy_async(&shmem_warp_stream_ptr[(SHMEM_STRIDE * i) + (laneId << loadStride)],\n &src_gmem_warp_stream_ptr[(GLOBAL_MEM_STRIDE * i) + (laneId << loadStride)],\n shape4,\n pipe);\n pipe.producer_commit();\n }\n // Now wait for all the above issued 8 batches to complete.\n cuda::pipeline_consumer_wait_prior<0>(pipe);\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment multiplications\n // along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n // Scale the C matrix.\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n pipe.consumer_release();\n\n // sync here so that shared memory can then be used for loading A & B matrices.\n __syncthreads();\n\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const float *warp_ptr =\n (warpId < (WARPS_PER_BLOCK / 2))\n ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);\n\n constexpr int chunksPerLane = ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP) * 2;\n const int laneLoadElem = (laneId % CHUNK_COPY_LINE_LANES) << loadStride;\n const int stridePerLaneCopy = (laneId / CHUNK_COPY_LINE_LANES);\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.\n // As for tf32 MMA M == N we use M for warp 4-7 + shmem_idx_b_off.\n size_t shmem_idx =\n (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2) + ((warpId / (WARPS_PER_BLOCK / 2)) * shmem_idx_b_off);\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n const float *lane_ptr = (warp_ptr + tile_k * K + stridePerLaneCopy * K_GLOBAL + laneLoadElem);\n\n // Shift the second half of the warp to the next row / column in the shared memory.\n shmem_idx += stridePerLaneCopy;\n\n#pragma unroll\n for (int i = 0; i < chunksPerLane; i++) {\n // Copy 16 bytes at once in each lane.\n pipe.producer_acquire();\n cuda::memcpy_async(&shmem[shmem_idx][laneLoadElem], lane_ptr, shape4, pipe);\n pipe.producer_commit();\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n\n cuda::pipeline_consumer_wait_prior<0>(pipe);\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / BLOCK_ROW_WARPS) * M * BLOCK_ROW_WARPS + (i * M);\n const float *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_FLOAT);\n\n#pragma unroll\n for (int t = 0; t < a[i].num_elements; t++) {\n a[i].x[t] = wmma::__float_to_tf32(a[i].x[t]);\n }\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be reused\n // against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const float *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_FLOAT);\n#pragma unroll\n for (int t = 0; t < b[j].num_elements; t++) {\n b[j].x[t] = wmma::__float_to_tf32(b[j].x[t]);\n }\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n pipe.consumer_release();\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n // Uniform, point-wise transformations of ALL fragment elements by ALL threads in the\n // warp are well-defined even though element indices within fragment storage are not defined.\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global memory.\n float *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((float4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((float4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n#endif\n}\n\n// Performs an MxNxK tf32 GEMM (C=alpha*A*B + beta*C) assuming:\n// 1) Matrices are packed in memory.\n// 2) M, N and K are multiples of 16, 16 and 8 respectively.\n// 3) A is row major, B is column major matrix.\n// Note: This is a less performant version of the compute_tf32gemm kernel. It is designed for\n// demonstration purposes only to show the CUDA WMMA API use without relying on\n// availability of the shared memory.\n__global__ void\nsimple_wmma_tf32gemm(float *a, float *b, float *c, float *d, int m_ld, int n_ld, int k_ld, float alpha, float beta)\n{\n#if __CUDA_ARCH__ >= 800\n // Leading dimensions. Packed with no transpositions.\n int lda = k_ld;\n int ldb = k_ld;\n int ldc = n_ld;\n\n // Tile using a 2D grid\n int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;\n int warpN = (blockIdx.y * blockDim.y + threadIdx.y);\n\n // Declare the fragments\n wmma::fragment a_frag;\n wmma::fragment b_frag;\n wmma::fragment acc_frag;\n wmma::fragment c_frag;\n\n wmma::fill_fragment(acc_frag, 0.0f);\n\n // Loop over k\n for (int i = 0; i < k_ld; i += K) {\n int aCol = i;\n int aRow = warpM * M;\n\n // int bCol = i;\n // int bRow = warpN * N;\n int bCol = warpN * N;\n int bRow = i;\n\n // Bounds checking\n if (aRow < m_ld && aCol < k_ld && bRow < k_ld && bCol < n_ld) {\n // Load the inputs\n wmma::load_matrix_sync(a_frag, a + aCol + aRow * lda, lda);\n wmma::load_matrix_sync(b_frag, b + bRow + bCol * ldb, ldb);\n\n#pragma unroll\n for (int t = 0; t < a_frag.num_elements; t++) {\n a_frag.x[t] = wmma::__float_to_tf32(a_frag.x[t]);\n }\n\n#pragma unroll\n for (int t = 0; t < b_frag.num_elements; t++) {\n b_frag.x[t] = wmma::__float_to_tf32(b_frag.x[t]);\n }\n // Perform the matrix multiplication\n wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);\n }\n }\n\n // Load in the current value of c, scale it by beta, and add this our result scaled by alpha\n int cCol = warpN * N;\n int cRow = warpM * M;\n\n if (cRow < m_ld && cCol < n_ld) {\n wmma::load_matrix_sync(c_frag, c + cCol + cRow * ldc, ldc, wmma::mem_row_major);\n\n for (int i = 0; i < c_frag.num_elements; i++) {\n c_frag.x[i] = alpha * acc_frag.x[i] + beta * c_frag.x[i];\n }\n\n // Store the output\n wmma::store_matrix_sync(d + cCol + cRow * ldc, c_frag, ldc, wmma::mem_row_major);\n }\n#endif\n}\n\n__host__ void matMultiplyOnHost(float *A,\n float *B,\n float *C,\n float alpha,\n float beta,\n int numARows,\n int numAColumns,\n int numBRows,\n int numBColumns,\n int numCRows,\n int numCColumns)\n{\n for (int i = 0; i < numCRows; i++) {\n for (int j = 0; j < numCColumns; j++) {\n float temp = 0.0;\n\n for (int k = 0; k < numAColumns; k++) {\n temp += A[i * numAColumns + k] * B[j * numBRows + k];\n }\n\n C[i * numCColumns + j] = temp * alpha + beta * C[i * numCColumns + j];\n }\n }\n}\n\nint main(int argc, char **argv)\n{\n printf(\"Initializing...\\n\");\n\n int dev = findCudaDevice(argc, (const char **)argv);\n\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, dev));\n\n // Tensor cores require a GPU of Volta (SM8X) architecture or higher.\n if (deviceProp.major < 8) {\n printf(\"tf32TensorCoreGemm requires requires SM 8.0 or higher to use Tensor Cores. Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n printf(\"M: %d (%d x %d)\\n\", M_GLOBAL, M, M_TILES);\n printf(\"N: %d (%d x %d)\\n\", N_GLOBAL, N, N_TILES);\n printf(\"K: %d (%d x %d)\\n\", K_GLOBAL, K, K_TILES);\n\n float *A_h = NULL;\n float *B_h = NULL;\n float *C_h = NULL;\n#if CPU_DEBUG\n float *result_hD = NULL;\n float *result_host = NULL;\n#endif\n\n A_h = (float *)malloc(sizeof(float) * M_GLOBAL * K_GLOBAL);\n B_h = (float *)malloc(sizeof(float) * K_GLOBAL * N_GLOBAL);\n C_h = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n#if CPU_DEBUG\n result_hD = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n result_host = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n#endif\n\n float *A = NULL;\n float *B = NULL;\n float *C = NULL;\n float *D = NULL;\n\n checkCudaErrors(hipMalloc((void **)&A, sizeof(float) * M_GLOBAL * K_GLOBAL));\n checkCudaErrors(hipMalloc((void **)&B, sizeof(float) * N_GLOBAL * K_GLOBAL));\n checkCudaErrors(hipMalloc((void **)&C, sizeof(float) * M_GLOBAL * N_GLOBAL));\n checkCudaErrors(hipMalloc((void **)&D, sizeof(float) * M_GLOBAL * N_GLOBAL));\n\n assert(((unsigned long long)A) % 128 == 0);\n assert(((unsigned long long)B) % 128 == 0);\n assert(((unsigned long long)C) % 128 == 0);\n assert(((unsigned long long)D) % 128 == 0);\n\n init_host_matrices(A_h, B_h, C_h);\n\n printf(\"Preparing data for GPU...\\n\");\n\n checkCudaErrors(hipMemcpy(A, A_h, sizeof(float) * M_GLOBAL * K_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(B, B_h, sizeof(float) * N_GLOBAL * K_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(C, C_h, sizeof(float) * M_GLOBAL * N_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemset(D, 0, sizeof(float) * M_GLOBAL * N_GLOBAL));\n\n enum {\n // Compute the right amount of shared memory to request.\n // We need shared memory to hold per-CTA C and D matrix tiles, and to cache per-CTA chunks\n // of the A and B matrices. Therefore, the right amount to request is the maximum of those\n // two numbers.\n SHMEM_SZ = MAX(sizeof(float) * (BLOCK_COL_TILES * M) * (CHUNK_K * K + SKEW_FLOAT) * 2,\n M * (BLOCK_ROW_WARPS * WARP_ROW_TILES) * N * (BLOCK_COL_WARPS * WARP_COL_TILES) * sizeof(float))\n };\n\n printf(\"Required shared memory size: %lu Kb\\n\", SHMEM_SZ / 1024UL);\n\n const float alpha = 1.1f;\n const float beta = 1.2f;\n\n hipEvent_t start, stop;\n\n checkCudaErrors(hipEventCreate(&start));\n checkCudaErrors(hipEventCreate(&stop));\n checkCudaErrors(hipEventRecord(start));\n\n // kernel to run - default (tf32mma_shmem_gemm_async_copy == 0)\n kernels selected_kernel = tf32mma_shmem_gemm_async_copy;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"kernel\")) {\n int kernel_number = getCmdLineArgumentInt(argc, (const char **)argv, \"kernel\");\n if (kernel_number < 3) {\n selected_kernel = (kernels)kernel_number;\n }\n else {\n printf(\"Error: kernel number should be between 0 to 2, you have entered %d\\n\", kernel_number);\n exit(EXIT_FAILURE);\n }\n }\n\n // If enough shared memory available on the GPU use high performant kernel\n if ((deviceProp.sharedMemPerMultiprocessor >= SHMEM_SZ) && (selected_kernel != simple_tf32mma_gemm)) {\n printf(\"Computing using high performance kernel = %d - %s\\n\", selected_kernel, kernelNames[selected_kernel]);\n\n switch (selected_kernel) {\n case tf32mma_shmem_gemm_async_copy:\n default:\n checkCudaErrors(hipFuncSetAttribute(reinterpret_cast(compute_tf32gemm_async_copy), hipFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors(\n (compute_tf32gemm_async_copy<<>>(\n A, B, C, D, alpha, beta)));\n break;\n case tf32mma_shmem_gemm:\n checkCudaErrors(\n hipFuncSetAttribute(reinterpret_cast(compute_tf32gemm), hipFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors((compute_tf32gemm<<>>(\n A, B, C, D, alpha, beta)));\n break;\n }\n#if CPU_DEBUG\n checkCudaErrors(hipMemcpy(result_hD, D, sizeof(float) * M_GLOBAL * N_GLOBAL, hipMemcpyDeviceToHost));\n#endif\n }\n else {\n dim3 gridDim;\n dim3 blockDim;\n\n // blockDim.x must be a multple of warpSize\n // 128x4 means we have 16 warps and a block computes a 64x64 output tile\n blockDim.x = 128;\n blockDim.y = 4;\n\n gridDim.x = (M_GLOBAL + (M * blockDim.x / 32 - 1)) / (M * blockDim.x / 32);\n gridDim.y = (N_GLOBAL + N * blockDim.y - 1) / (N * blockDim.y);\n\n printf(\"Computing... using simple_wmma_gemm kernel\\n\");\n simple_wmma_tf32gemm<<>>(A, B, C, D, M_GLOBAL, N_GLOBAL, K_GLOBAL, alpha, beta);\n#if CPU_DEBUG\n checkCudaErrors(hipMemcpy(result_hD, D, sizeof(float) * M_GLOBAL * N_GLOBAL, hipMemcpyDeviceToHost));\n#endif\n }\n\n checkCudaErrors(hipEventRecord(stop));\n checkCudaErrors(hipEventSynchronize(stop));\n\n#if CPU_DEBUG\n printf(\"Verifying correctness of the computations...\\n\");\n\n memcpy(result_host, C_h, sizeof(float) * M_GLOBAL * N_GLOBAL);\n\n matMultiplyOnHost(A_h, B_h, result_host, alpha, beta, M_GLOBAL, K_GLOBAL, K_GLOBAL, N_GLOBAL, M_GLOBAL, N_GLOBAL);\n\n for (int i = 0; i < N_GLOBAL * M_GLOBAL; i++) {\n if (fabs(result_hD[i] - result_host[i]) > 0.1f) {\n printf(\"mismatch i=%d result_hD=%f result_host=%f\\n\", i, result_hD[i], result_host[i]);\n }\n }\n free(result_hD);\n free(result_host);\n#endif\n\n float milliseconds = 0;\n\n checkCudaErrors(hipEventElapsedTime(&milliseconds, start, stop));\n\n printf(\"Time: %f ms\\n\", milliseconds);\n printf(\"TFLOPS: %.2f\\n\", (((double)M_GLOBAL * N_GLOBAL * K_GLOBAL * 2) / (milliseconds / 1000.)) / 1e12);\n\n free(A_h);\n free(B_h);\n free(C_h);\n checkCudaErrors(hipFree((void *)A));\n checkCudaErrors(hipFree((void *)B));\n checkCudaErrors(hipFree((void *)C));\n checkCudaErrors(hipFree((void *)D));\n\n return 0;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/tf32TensorCoreGemm/tf32TensorCoreGemm.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * Simple kernel for ptxjit demonstration.\n *\n */\nextern \"C\" __global__ void myKernel(int *data)\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n data[tid] = tid;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * Simple kernel for ptxjit demonstration.\n *\n */\n#include \"hip/hip_runtime.h\"\nextern \"C\" __global__ void myKernel(int *data)\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n data[tid] = tid;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/ptxjit/ptxjit_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n//\n// This sample uses the compressible memory allocation if device supports it\n// and performs saxpy on it.\n// Compressible memory may give better performance if the data is amenable to\n// compression.\n\n#include \n#include \n#define CUDA_DRIVER_API\n#include \"compMalloc.h\"\n#include \"helper_cuda.h\"\n\n__global__ void saxpy(const float a, const float4 *x, const float4 *y, float4 *z, const size_t n)\n{\n for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {\n const float4 x4 = x[i];\n const float4 y4 = y[i];\n z[i] = make_float4(a * x4.x + y4.x, a * x4.y + y4.y, a * x4.z + y4.z, a * x4.w + y4.w);\n }\n}\n\n__global__ void init(float4 *x, float4 *y, const float val, const size_t n)\n{\n const float4 val4 = make_float4(val, val, val, val);\n for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {\n x[i] = y[i] = val4;\n }\n}\n\nvoid launchSaxpy(const float a,\n float4 *x,\n float4 *y,\n float4 *z,\n const size_t n,\n const float init_val,\n const bool compressibleZbuf)\n{\n cudaEvent_t start, stop;\n float ms;\n int blockSize;\n int minGridSize;\n dim3 threads, blocks;\n\n if (!compressibleZbuf) {\n // We are on config where compressible buffer can only be initialized through cudaMemcpy\n // hence, x & y buffers are allocated as compressible and initialized via cudaMemcpy\n // whereas z buffer is allocated as non-compressible.\n float4 *h_x = (float4 *)malloc(sizeof(float4) * n);\n float4 *h_y = (float4 *)malloc(sizeof(float4) * n);\n for (int i = 0; i < n; i++) {\n h_x[i].x = h_x[i].y = h_x[i].z = h_x[i].w = init_val;\n h_y[i].x = h_y[i].y = h_y[i].z = h_y[i].w = init_val;\n }\n checkCudaErrors(cudaMemcpy(x, h_x, sizeof(float4) * n, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(y, h_y, sizeof(float4) * n, cudaMemcpyHostToDevice));\n free(h_x);\n free(h_y);\n }\n else {\n checkCudaErrors(cudaOccupancyMaxPotentialBlockSize(&minGridSize, &blockSize, (void *)init));\n threads = dim3(blockSize, 1, 1);\n blocks = dim3(minGridSize, 1, 1);\n init<<>>(x, y, init_val, n);\n }\n\n checkCudaErrors(cudaOccupancyMaxPotentialBlockSize(&minGridSize, &blockSize, (void *)saxpy));\n threads = dim3(blockSize, 1, 1);\n blocks = dim3(minGridSize, 1, 1);\n\n checkCudaErrors(cudaEventCreate(&start));\n checkCudaErrors(cudaEventCreate(&stop));\n checkCudaErrors(cudaEventRecord(start));\n saxpy<<>>(a, x, y, z, n);\n checkCudaErrors(cudaEventRecord(stop));\n checkCudaErrors(cudaEventSynchronize(stop));\n checkCudaErrors(cudaEventElapsedTime(&ms, start, stop));\n\n const size_t size = n * sizeof(float4);\n printf(\"Running saxpy with %d blocks x %d threads = %.3f ms %.3f TB/s\\n\",\n blocks.x,\n threads.x,\n ms,\n (size * 3) / ms / 1e9);\n}\n\nint main(int argc, char **argv)\n{\n const size_t n = 10485760;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\") || checkCmdLineFlag(argc, (const char **)argv, \"?\")) {\n printf(\"Usage -device=n (n >= 0 for deviceID)\\n\");\n exit(EXIT_SUCCESS);\n }\n\n findCudaDevice(argc, (const char **)argv);\n CUdevice currentDevice;\n checkCudaErrors(cuCtxGetDevice(¤tDevice));\n\n // Check that the selected device supports virtual memory management\n int vmm_supported = -1;\n checkCudaErrors(\n cuDeviceGetAttribute(&vmm_supported, CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED, currentDevice));\n if (vmm_supported == 0) {\n printf(\"Device %d doesn't support Virtual Memory Management, waiving the execution.\\n\", currentDevice);\n exit(EXIT_WAIVED);\n }\n\n int isCompressionAvailable;\n checkCudaErrors(cuDeviceGetAttribute(\n &isCompressionAvailable, CU_DEVICE_ATTRIBUTE_GENERIC_COMPRESSION_SUPPORTED, currentDevice));\n if (isCompressionAvailable == 0) {\n printf(\"Device %d doesn't support Generic memory compression, waiving the execution.\\n\", currentDevice);\n exit(EXIT_WAIVED);\n }\n\n printf(\"Generic memory compression support is available\\n\");\n\n int major, minor;\n checkCudaErrors(cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, currentDevice));\n checkCudaErrors(cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, currentDevice));\n float4 *x, *y, *z;\n const size_t size = n * sizeof(float4);\n\n // Allocating compressible memory\n checkCudaErrors(allocateCompressible((void **)&x, size, true));\n checkCudaErrors(allocateCompressible((void **)&y, size, true));\n bool compressibleZbuf = 0;\n if ((major == 8 && minor == 0) || (major == 8 && minor == 6)) {\n // On SM 8.0 and 8.6 GPUs compressible buffer can only be initialized\n // through cudaMemcpy.\n printf(\"allocating non-compressible Z buffer\\n\");\n checkCudaErrors(allocateCompressible((void **)&z, size, false));\n compressibleZbuf = 0;\n }\n else {\n checkCudaErrors(allocateCompressible((void **)&z, size, true));\n compressibleZbuf = 1;\n }\n\n printf(\"Running saxpy on %zu bytes of Compressible memory\\n\", size);\n\n const float a = 1.0f;\n const float init_val = 1.0f;\n launchSaxpy(a, x, y, z, n, init_val, compressibleZbuf);\n\n checkCudaErrors(freeCompressible(x, size, true));\n checkCudaErrors(freeCompressible(y, size, true));\n checkCudaErrors(freeCompressible(z, size, true));\n\n printf(\"Running saxpy on %zu bytes of Non-Compressible memory\\n\", size);\n // Allocating non-compressible memory\n checkCudaErrors(allocateCompressible((void **)&x, size, false));\n checkCudaErrors(allocateCompressible((void **)&y, size, false));\n checkCudaErrors(allocateCompressible((void **)&z, size, false));\n\n launchSaxpy(a, x, y, z, n, init_val, compressibleZbuf);\n\n checkCudaErrors(freeCompressible(x, size, false));\n checkCudaErrors(freeCompressible(y, size, false));\n checkCudaErrors(freeCompressible(z, size, false));\n\n printf(\"\\nNOTE: The CUDA Samples are not meant for performance measurements. \"\n \"Results may vary when GPU Boost is enabled.\\n\");\n return EXIT_SUCCESS;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n//\n// This sample uses the compressible memory allocation if device supports it\n// and performs saxpy on it.\n// Compressible memory may give better performance if the data is amenable to\n// compression.\n\n#include \n#include \n#define CUDA_DRIVER_API\n#include \"compMalloc.h\"\n#include \"hip/hip_runtime_api.h\"\n\n__global__ void saxpy(const float a, const float4 *x, const float4 *y, float4 *z, const size_t n)\n{\n for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {\n const float4 x4 = x[i];\n const float4 y4 = y[i];\n z[i] = make_float4(a * x4.x + y4.x, a * x4.y + y4.y, a * x4.z + y4.z, a * x4.w + y4.w);\n }\n}\n\n__global__ void init(float4 *x, float4 *y, const float val, const size_t n)\n{\n const float4 val4 = make_float4(val, val, val, val);\n for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {\n x[i] = y[i] = val4;\n }\n}\n\nvoid launchSaxpy(const float a,\n float4 *x,\n float4 *y,\n float4 *z,\n const size_t n,\n const float init_val,\n const bool compressibleZbuf)\n{\n hipEvent_t start, stop;\n float ms;\n int blockSize;\n int minGridSize;\n dim3 threads, blocks;\n\n if (!compressibleZbuf) {\n // We are on config where compressible buffer can only be initialized through cudaMemcpy\n // hence, x & y buffers are allocated as compressible and initialized via cudaMemcpy\n // whereas z buffer is allocated as non-compressible.\n float4 *h_x = (float4 *)malloc(sizeof(float4) * n);\n float4 *h_y = (float4 *)malloc(sizeof(float4) * n);\n for (int i = 0; i < n; i++) {\n h_x[i].x = h_x[i].y = h_x[i].z = h_x[i].w = init_val;\n h_y[i].x = h_y[i].y = h_y[i].z = h_y[i].w = init_val;\n }\n checkCudaErrors(hipMemcpy(x, h_x, sizeof(float4) * n, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(y, h_y, sizeof(float4) * n, hipMemcpyHostToDevice));\n free(h_x);\n free(h_y);\n }\n else {\n checkCudaErrors(hipOccupancyMaxPotentialBlockSize(&minGridSize, &blockSize, (void *)init));\n threads = dim3(blockSize, 1, 1);\n blocks = dim3(minGridSize, 1, 1);\n init<<>>(x, y, init_val, n);\n }\n\n checkCudaErrors(hipOccupancyMaxPotentialBlockSize(&minGridSize, &blockSize, (void *)saxpy));\n threads = dim3(blockSize, 1, 1);\n blocks = dim3(minGridSize, 1, 1);\n\n checkCudaErrors(hipEventCreate(&start));\n checkCudaErrors(hipEventCreate(&stop));\n checkCudaErrors(hipEventRecord(start));\n saxpy<<>>(a, x, y, z, n);\n checkCudaErrors(hipEventRecord(stop));\n checkCudaErrors(hipEventSynchronize(stop));\n checkCudaErrors(hipEventElapsedTime(&ms, start, stop));\n\n const size_t size = n * sizeof(float4);\n printf(\"Running saxpy with %d blocks x %d threads = %.3f ms %.3f TB/s\\n\",\n blocks.x,\n threads.x,\n ms,\n (size * 3) / ms / 1e9);\n}\n\nint main(int argc, char **argv)\n{\n const size_t n = 10485760;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\") || checkCmdLineFlag(argc, (const char **)argv, \"?\")) {\n printf(\"Usage -device=n (n >= 0 for deviceID)\\n\");\n exit(EXIT_SUCCESS);\n }\n\n findCudaDevice(argc, (const char **)argv);\n hipDevice_t currentDevice;\n checkCudaErrors(hipCtxGetDevice(¤tDevice));\n\n // Check that the selected device supports virtual memory management\n int vmm_supported = -1;\n checkCudaErrors(\n hipDeviceGetAttribute(&vmm_supported, CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED, currentDevice));\n if (vmm_supported == 0) {\n printf(\"Device %d doesn't support Virtual Memory Management, waiving the execution.\\n\", currentDevice);\n exit(EXIT_WAIVED);\n }\n\n int isCompressionAvailable;\n checkCudaErrors(hipDeviceGetAttribute(\n &isCompressionAvailable, CU_DEVICE_ATTRIBUTE_GENERIC_COMPRESSION_SUPPORTED, currentDevice));\n if (isCompressionAvailable == 0) {\n printf(\"Device %d doesn't support Generic memory compression, waiving the execution.\\n\", currentDevice);\n exit(EXIT_WAIVED);\n }\n\n printf(\"Generic memory compression support is available\\n\");\n\n int major, minor;\n checkCudaErrors(hipDeviceGetAttribute(&major, hipDeviceAttributeComputeCapabilityMajor, currentDevice));\n checkCudaErrors(hipDeviceGetAttribute(&minor, hipDeviceAttributeComputeCapabilityMinor, currentDevice));\n float4 *x, *y, *z;\n const size_t size = n * sizeof(float4);\n\n // Allocating compressible memory\n checkCudaErrors(allocateCompressible((void **)&x, size, true));\n checkCudaErrors(allocateCompressible((void **)&y, size, true));\n bool compressibleZbuf = 0;\n if ((major == 8 && minor == 0) || (major == 8 && minor == 6)) {\n // On SM 8.0 and 8.6 GPUs compressible buffer can only be initialized\n // through cudaMemcpy.\n printf(\"allocating non-compressible Z buffer\\n\");\n checkCudaErrors(allocateCompressible((void **)&z, size, false));\n compressibleZbuf = 0;\n }\n else {\n checkCudaErrors(allocateCompressible((void **)&z, size, true));\n compressibleZbuf = 1;\n }\n\n printf(\"Running saxpy on %zu bytes of Compressible memory\\n\", size);\n\n const float a = 1.0f;\n const float init_val = 1.0f;\n launchSaxpy(a, x, y, z, n, init_val, compressibleZbuf);\n\n checkCudaErrors(freeCompressible(x, size, true));\n checkCudaErrors(freeCompressible(y, size, true));\n checkCudaErrors(freeCompressible(z, size, true));\n\n printf(\"Running saxpy on %zu bytes of Non-Compressible memory\\n\", size);\n // Allocating non-compressible memory\n checkCudaErrors(allocateCompressible((void **)&x, size, false));\n checkCudaErrors(allocateCompressible((void **)&y, size, false));\n checkCudaErrors(allocateCompressible((void **)&z, size, false));\n\n launchSaxpy(a, x, y, z, n, init_val, compressibleZbuf);\n\n checkCudaErrors(freeCompressible(x, size, false));\n checkCudaErrors(freeCompressible(y, size, false));\n checkCudaErrors(freeCompressible(z, size, false));\n\n printf(\"\\nNOTE: The CUDA Samples are not meant for performance measurements. \"\n \"Results may vary when GPU Boost is enabled.\\n\");\n return EXIT_SUCCESS;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/cudaCompressibleMemory/saxpy.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// CUDA sample demonstrating a GEMM computation using the Warp Matrix Multiply\n// and Accumulate API introduced in CUDA 9.\n\n// In this program, the compute_gemm kernel computes the result of a matrix\n// multiplication and addition: D = alpha * A * B + beta * C. The dimensions of\n// both C and D matrices are M_GLOBAL x N_GLOBAL. The A matrix is M_GLOBAL x\n// K_GLOBAL (row-major), the B matrix is K_GLOBAL x N_GLOBAL (column-major). In\n// that kernel, each CTA computes one 128 x 128 tile of the resulting matrix per\n// iteration. When the tile is computed, the CTA stores it to the global memory\n// and begins a new iteration, selecting a new 128 x 128 tile to compute.\n// Each CTA consists of eight warps. For the 128 x 128 tile, each warp computes\n// eight 16 x 16 subtiles, organized in a 2 x 4 two-dimensional array. Warps\n// compute the 16 x 16 subtiles using nvcuda::wmma::mma_sync operations by\n// moving through the K_GLOBAL dimension of the A and B matrices and\n// accumulating the intermediate result in the local thread state.\n\n// There are a number of simple optimizations used in the algorithm:\n// - The CTA copies the 128 x 128 tile of the C matrix from the global memory to\n// shared memory. After that is done, each warp loads the C matrix fragments\n// from shared memory, thus avoiding a random global memory access.\n// - On each internal iteration, the CTA copies a portion of the A and B\n// matrices from global memory to shared memory. After that, all warps in the\n// CTA reuse the A and B data from shared memory, thus reducing the number of\n// data copies from global memory.\n// - The portions of the A and B matrices are stored in shared memory with an\n// additional padding (skew) to reduce the number of shared memory access bank\n// conflicts.\n// (See a detailed explanation near the SKEW_HALF macro definition.)\n// - When the CTA finishes computing the tiles of the resulting matrix, each\n// warp stores its subtiles to shared memory. The CTA then copies the shared\n// memory contents to global memory, again avoiding redundant random global\n// memory accesses.\n// - Note that the CTA tile size is chosen to maximize the GPU register\n// utilization, but carefully enough to avoid local memory use.\n\n#include \n#include \n#include \n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Externally configurable parameters.\n\n#ifndef CPU_DEBUG\n// Set this to 1 to verify the correctness of the GPU-computed matrix.\n#define CPU_DEBUG 0\n#endif\n\n#ifndef SHARED_MEMORY_LIMIT_64K\n// Set this to 0 to use more than 64 Kb of shared memory to cache data, to\n// improve the performance of the computations on GPU.\n// Note that you need a GPU that can have more than 64 Kb of shared memory\n// per multiprocessor.\n#define SHARED_MEMORY_LIMIT_64K 1\n#endif\n\n// GPU configuration.\n\n#define WARP_SIZE 32\n\n// MMA matrix tile dimensions.\n\n#define M 16\n#define N 16\n#define K 16\n\n#define WMMA_M 16\n#define WMMA_N 16\n#define WMMA_K 16\n\n// GEMM configuration.\n\n#define M_TILES 256\n#define N_TILES 256\n#define K_TILES 256\n\n#define M_GLOBAL (M * M_TILES)\n#define N_GLOBAL (N * N_TILES)\n#define K_GLOBAL (K * K_TILES)\n\n#define C_LAYOUT wmma::mem_row_major\n\n// Implementation constants.\n\n#define WARPS_PER_BLOCK 8\n#define THREADS_PER_BLOCK (WARP_SIZE * WARPS_PER_BLOCK)\n\n#if SHARED_MEMORY_LIMIT_64K\n// With only 64 Kb shared memory available, we can fit two 8-tile chunks of\n// the A and B matrix data, that are 16 * 16 * 8 * 8 * 2 = 32 Kb each\n// (i.e. two 8x8 arrays of tiles of 16x16 half-typed elements per CTA).\n// But we cannot account the 8 Kb total skew overhead, without which the\n// performance would be severely impacted. So we choose to reduce the chunk size\n// in half, i.e. the amount of A and B matrix data we cache in shared memory.\n// Accordingly, this doubles the number of outer iterations across the global K\n// dimension, which only slightly impacts the performance.\n#define CHUNK_K 4\n#else\n#define CHUNK_K 8\n#endif\n\n#define CHUNK_LINE_BYTES (CHUNK_K * K * sizeof(half))\n#define WARP_COPY_BYTES (WARP_SIZE * sizeof(int4))\n#define CHUNK_COPY_LINES_PER_WARP (WARP_COPY_BYTES / CHUNK_LINE_BYTES)\n#define CHUNK_COPY_LINE_LANES (WARP_SIZE / CHUNK_COPY_LINES_PER_WARP)\n\n#define BLOCK_ROW_WARPS 2\n#define BLOCK_COL_WARPS 4\n\n#define WARP_ROW_TILES 4\n#define WARP_COL_TILES 2\n\n#define BLOCK_ROW_TILES (WARP_ROW_TILES * BLOCK_ROW_WARPS)\n#define BLOCK_COL_TILES (WARP_COL_TILES * BLOCK_COL_WARPS)\n\n#define GLOBAL_MEM_STRIDE N_GLOBAL\n\n#define SHMEM_STRIDE (N * BLOCK_ROW_TILES)\n#define SHMEM_OFFSET (N * WARP_ROW_TILES)\n\n// The macro below is used to shift rows of the A matrix and columns of the B matrix\n// in shared memory to minimize possible bank conflicts.\n// Before performing the nvcuda::wmma::mma_sync operation, the warp must load the matrix\n// data using the nvcuda::wmma::load_matrix_sync operation. Although the memory access pattern\n// is not specified for that function, each lane in the warp can read one or multiple matrix\n// elements from different matrix rows or columns.\n// For shared memory, such access can result in bank conflicts if different rows / columns\n// of the matrix map to the same bank. By shifting each row and column by a few bytes, we\n// make sure that they map to different banks, thus reducing the number of possible bank\n// conflicts.\n// The number of 16 two-byte \"half\" elements is chosen as the minimum possible shift because\n// we must keep each row and column 256-bit aligned, as required by nvcuda::wmma::load_matrix_sync.\n#define SKEW_HALF 16\n\n#define checkKernelErrors(expr) \\\n do { \\\n expr; \\\n \\\n cudaError_t __err = cudaGetLastError(); \\\n if (__err != cudaSuccess) { \\\n printf(\"Line %d: '%s' failed: %s\\n\", __LINE__, #expr, cudaGetErrorString(__err)); \\\n abort(); \\\n } \\\n } while (0)\n\nusing namespace nvcuda;\n\n__host__ void init_host_matrices(half *a, half *b, float *c)\n{\n for (int i = 0; i < M_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n a[i * K_GLOBAL + j] = (half)(rand() % 3);\n }\n }\n\n for (int i = 0; i < N_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n b[i * K_GLOBAL + j] = (half)(rand() % 3);\n }\n }\n\n for (int t = 0; t < M_GLOBAL * N_GLOBAL; t++) {\n c[t] = static_cast(rand() % 3);\n }\n}\n\n__global__ void compute_gemm(const half *A, const half *B, const float *C, float *D, float alpha, float beta)\n{\n extern __shared__ half shmem[][CHUNK_K * K + SKEW_HALF];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n const size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n float *shmem_warp_tile_ptr =\n (float *)&shmem[0][0] + (warpId / 2) * SHMEM_STRIDE * K * 2 + (warpId % 2) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and\n // from shared memory.\n float *shmem_warp_stream_ptr = (float *)&shmem[0][0] + warpId * SHMEM_STRIDE * K;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may\n // result in a loss of precision). Zero still needs to be specially handled\n // though.\n beta /= alpha;\n\n // Each CTA slides along the 128 x 128 tiles from the top left corner of the\n // matrix to the right and down, and selects the next tile to compute. Once\n // there's no such tile, all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_ROW_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared\n // memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const float *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < K; i++) {\n typedef int4 copy_t;\n\n *((copy_t *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId) =\n *((copy_t *)(src_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment\n // multiplications along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * K + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Scale the C matrix.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const half *warp_ptr = (warpId < 4) ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % 4) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % 4) * 2);\n\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy\n // the B matrix.\n size_t shmem_idx = warpId < (WARPS_PER_BLOCK / 2)\n ? (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (N * (warpId % (WARPS_PER_BLOCK / 2)) * 2 + shmem_idx_b_off);\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n int4 *lane_ptr = (int4 *)(warp_ptr + tile_k * K + (laneId / CHUNK_COPY_LINE_LANES) * K_GLOBAL)\n + (laneId % CHUNK_COPY_LINE_LANES);\n\n // Shift the second half of the warp to the next row / column in the\n // shared memory.\n shmem_idx += laneId / CHUNK_COPY_LINE_LANES;\n\n#pragma unroll\n for (int i = 0; i < ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP) * 2; i++) {\n // Copy 16 bytes at once in each lane.\n *((int4 *)&shmem[shmem_idx][0] + (laneId % CHUNK_COPY_LINE_LANES)) = *lane_ptr;\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = (int4 *)((half *)lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP);\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / 2) * M * 2 + (i * M);\n const half *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_HALF);\n\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be\n // reused against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const half *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_HALF);\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n // Uniform, point-wise transformations of ALL fragment elements by ALL\n // threads in the warp are well-defined even though element indices\n // within fragment storage are not defined.\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * K + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global\n // memory.\n float *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < K; i++) {\n *((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n}\n\n// Performs an MxNxK GEMM (C=alpha*A*B + beta*C) assuming:\n// 1) Matrices are packed in memory.\n// 2) M, N and K are multiples of 16.\n// 3) Neither A nor B are transposed.\n// Note: This is a less performant version of the compute_gemm kernel. It is\n// designed for\n// demonstration purposes only to show the CUDA WMMA API use without\n// relying on availability of the shared memory.\n__global__ void\nsimple_wmma_gemm(half *a, half *b, float *c, float *d, int m_ld, int n_ld, int k_ld, float alpha, float beta)\n{\n // Leading dimensions. Packed with no transpositions.\n int lda = k_ld;\n int ldb = k_ld;\n int ldc = n_ld;\n\n // Tile using a 2D grid\n int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;\n int warpN = (blockIdx.y * blockDim.y + threadIdx.y);\n\n // Declare the fragments\n wmma::fragment a_frag;\n wmma::fragment b_frag;\n wmma::fragment acc_frag;\n wmma::fragment c_frag;\n\n wmma::fill_fragment(acc_frag, 0.0f);\n\n // Loop over k\n for (int i = 0; i < k_ld; i += WMMA_K) {\n int aCol = i;\n int aRow = warpM * WMMA_M;\n int bCol = warpN * N;\n int bRow = i;\n\n // Bounds checking\n if (aRow < m_ld && aCol < k_ld && bRow < k_ld && bCol < n_ld) {\n // Load the inputs\n wmma::load_matrix_sync(a_frag, a + aCol + aRow * lda, lda);\n wmma::load_matrix_sync(b_frag, b + bRow + bCol * ldb, ldb);\n\n // Perform the matrix multiplication\n wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);\n }\n }\n\n // Load in the current value of c, scale it by beta, and add this our result\n // scaled by alpha\n int cCol = warpN * WMMA_N;\n int cRow = warpM * WMMA_M;\n\n if (cRow < m_ld && cCol < n_ld) {\n wmma::load_matrix_sync(c_frag, c + cCol + cRow * ldc, ldc, wmma::mem_row_major);\n\n for (int i = 0; i < c_frag.num_elements; i++) {\n c_frag.x[i] = alpha * acc_frag.x[i] + beta * c_frag.x[i];\n }\n\n // Store the output\n wmma::store_matrix_sync(d + cCol + cRow * ldc, c_frag, ldc, wmma::mem_row_major);\n }\n}\n\n__host__ void matMultiplyOnHost(half *A,\n half *B,\n float *C,\n float alpha,\n float beta,\n int numARows,\n int numAColumns,\n int numBRows,\n int numBColumns,\n int numCRows,\n int numCColumns)\n{\n for (int i = 0; i < numCRows; i++) {\n for (int j = 0; j < numCColumns; j++) {\n float temp = 0.0;\n\n for (int k = 0; k < numAColumns; k++) {\n temp += (float)A[i * numAColumns + k] * (float)B[j * numBRows + k];\n }\n\n C[i * numCColumns + j] = temp * alpha + beta * C[i * numCColumns + j];\n }\n }\n}\n\nint main(int argc, char **argv)\n{\n printf(\"Initializing...\\n\");\n\n int dev = findCudaDevice(argc, (const char **)argv);\n\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, dev));\n\n // Tensor cores require a GPU of Volta (SM7X) architecture or higher.\n if (deviceProp.major < 7) {\n printf(\"cudaTensorCoreGemm requires SM 7.0 or higher to use Tensor \"\n \"Cores. Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n printf(\"M: %d (%d x %d)\\n\", M_GLOBAL, M, M_TILES);\n printf(\"N: %d (%d x %d)\\n\", N_GLOBAL, N, N_TILES);\n printf(\"K: %d (%d x %d)\\n\", K_GLOBAL, K, K_TILES);\n\n half *A_h = NULL;\n half *B_h = NULL;\n float *C_h = NULL;\n#if CPU_DEBUG\n float *result_hD = NULL;\n float *result_host = NULL;\n#endif\n\n A_h = (half *)malloc(sizeof(half) * M_GLOBAL * K_GLOBAL);\n B_h = (half *)malloc(sizeof(half) * K_GLOBAL * N_GLOBAL);\n C_h = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n#if CPU_DEBUG\n result_hD = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n result_host = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n#endif\n\n half *A = NULL;\n half *B = NULL;\n float *C = NULL;\n float *D = NULL;\n\n checkCudaErrors(cudaMalloc(reinterpret_cast(&A), sizeof(half) * M_GLOBAL * K_GLOBAL));\n checkCudaErrors(cudaMalloc(reinterpret_cast(&B), sizeof(half) * N_GLOBAL * K_GLOBAL));\n checkCudaErrors(cudaMalloc(reinterpret_cast(&C), sizeof(float) * M_GLOBAL * N_GLOBAL));\n checkCudaErrors(cudaMalloc(reinterpret_cast(&D), sizeof(float) * M_GLOBAL * N_GLOBAL));\n\n assert(((unsigned long long)A) % 128 == 0);\n assert(((unsigned long long)B) % 128 == 0);\n assert(((unsigned long long)C) % 128 == 0);\n assert(((unsigned long long)D) % 128 == 0);\n\n init_host_matrices(A_h, B_h, C_h);\n\n printf(\"Preparing data for GPU...\\n\");\n\n checkCudaErrors(cudaMemcpy(A, A_h, sizeof(half) * M_GLOBAL * K_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(B, B_h, sizeof(half) * N_GLOBAL * K_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(C, C_h, sizeof(float) * M_GLOBAL * N_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemset(D, 0, sizeof(float) * M_GLOBAL * N_GLOBAL));\n\n enum {\n // Compute the right amount of shared memory to request.\n // We need shared memory to hold per-CTA C and D matrix tiles, and to cache\n // per-CTA chunks\n // of the A and B matrices. Therefore, the right amount to request is the\n // maximum of those\n // two numbers.\n SHMEM_SZ = MAX(sizeof(half) * (BLOCK_COL_TILES * M) * (CHUNK_K * K + SKEW_HALF) * 2,\n M * (BLOCK_ROW_WARPS * WARP_ROW_TILES) * N * (BLOCK_COL_WARPS * WARP_COL_TILES) * sizeof(float))\n };\n\n printf(\"Required shared memory size: %lu Kb\\n\", SHMEM_SZ / 1024UL);\n\n const float alpha = 1.1f;\n const float beta = 1.2f;\n\n cudaEvent_t start, stop;\n\n checkCudaErrors(cudaEventCreate(&start));\n checkCudaErrors(cudaEventCreate(&stop));\n checkCudaErrors(cudaEventRecord(start));\n\n // If enough shared memory available on the GPU use high performant kernel\n if (deviceProp.sharedMemPerMultiprocessor >= SHMEM_SZ) {\n printf(\"Computing... using high performance kernel compute_gemm \\n\");\n\n checkCudaErrors(cudaFuncSetAttribute(compute_gemm, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors(\n (compute_gemm<<>>(A, B, C, D, alpha, beta)));\n#if CPU_DEBUG\n checkCudaErrors(cudaMemcpy(result_hD, D, sizeof(float) * M_GLOBAL * N_GLOBAL, cudaMemcpyDeviceToHost));\n#endif\n }\n else {\n dim3 gridDim;\n dim3 blockDim;\n\n // blockDim.x must be a multple of warpSize\n // 128x4 means we have 16 warps and a block computes a 64x64 output tile\n blockDim.x = 128;\n blockDim.y = 4;\n\n gridDim.x = (M_GLOBAL + (WMMA_M * blockDim.x / 32 - 1)) / (WMMA_M * blockDim.x / 32);\n gridDim.y = (N_GLOBAL + WMMA_N * blockDim.y - 1) / (WMMA_N * blockDim.y);\n\n printf(\"Computing... using simple_wmma_gemm kernel\\n\");\n simple_wmma_gemm<<>>(A, B, C, D, M_GLOBAL, N_GLOBAL, K_GLOBAL, alpha, beta);\n#if CPU_DEBUG\n checkCudaErrors(cudaMemcpy(result_hD, D, sizeof(float) * M_GLOBAL * N_GLOBAL, cudaMemcpyDeviceToHost));\n#endif\n }\n\n checkCudaErrors(cudaEventRecord(stop));\n checkCudaErrors(cudaEventSynchronize(stop));\n\n#if CPU_DEBUG\n printf(\"Verifying correctness of the computations...\\n\");\n\n memcpy(result_host, C_h, sizeof(float) * M_GLOBAL * N_GLOBAL);\n\n matMultiplyOnHost(A_h, B_h, result_host, alpha, beta, M_GLOBAL, K_GLOBAL, K_GLOBAL, N_GLOBAL, M_GLOBAL, N_GLOBAL);\n\n for (int i = 0; i < N_GLOBAL * M_GLOBAL; i++) {\n if (fabs(result_hD[i] - result_host[i]) > 0.1f)\n printf(\"mismatch i=%d result_hD=%f result_host=%f\\n\", i, result_hD[i], result_host[i]);\n }\n free(result_hD);\n free(result_host);\n#endif\n\n float milliseconds = 0;\n\n checkCudaErrors(cudaEventElapsedTime(&milliseconds, start, stop));\n\n printf(\"Time: %f ms\\n\", milliseconds);\n printf(\"TFLOPS: %.2f\\n\",\n static_cast((static_cast(M_GLOBAL) * N_GLOBAL * K_GLOBAL * 2) / (milliseconds / 1000.))\n / 1e12);\n\n free(A_h);\n free(B_h);\n free(C_h);\n checkCudaErrors(cudaFree(reinterpret_cast(A)));\n checkCudaErrors(cudaFree(reinterpret_cast(B)));\n checkCudaErrors(cudaFree(reinterpret_cast(C)));\n checkCudaErrors(cudaFree(reinterpret_cast(D)));\n\n return 0;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// CUDA sample demonstrating a GEMM computation using the Warp Matrix Multiply\n// and Accumulate API introduced in CUDA 9.\n\n// In this program, the compute_gemm kernel computes the result of a matrix\n// multiplication and addition: D = alpha * A * B + beta * C. The dimensions of\n// both C and D matrices are M_GLOBAL x N_GLOBAL. The A matrix is M_GLOBAL x\n// K_GLOBAL (row-major), the B matrix is K_GLOBAL x N_GLOBAL (column-major). In\n// that kernel, each CTA computes one 128 x 128 tile of the resulting matrix per\n// iteration. When the tile is computed, the CTA stores it to the global memory\n// and begins a new iteration, selecting a new 128 x 128 tile to compute.\n// Each CTA consists of eight warps. For the 128 x 128 tile, each warp computes\n// eight 16 x 16 subtiles, organized in a 2 x 4 two-dimensional array. Warps\n// compute the 16 x 16 subtiles using nvcuda::wmma::mma_sync operations by\n// moving through the K_GLOBAL dimension of the A and B matrices and\n// accumulating the intermediate result in the local thread state.\n\n// There are a number of simple optimizations used in the algorithm:\n// - The CTA copies the 128 x 128 tile of the C matrix from the global memory to\n// shared memory. After that is done, each warp loads the C matrix fragments\n// from shared memory, thus avoiding a random global memory access.\n// - On each internal iteration, the CTA copies a portion of the A and B\n// matrices from global memory to shared memory. After that, all warps in the\n// CTA reuse the A and B data from shared memory, thus reducing the number of\n// data copies from global memory.\n// - The portions of the A and B matrices are stored in shared memory with an\n// additional padding (skew) to reduce the number of shared memory access bank\n// conflicts.\n// (See a detailed explanation near the SKEW_HALF macro definition.)\n// - When the CTA finishes computing the tiles of the resulting matrix, each\n// warp stores its subtiles to shared memory. The CTA then copies the shared\n// memory contents to global memory, again avoiding redundant random global\n// memory accesses.\n// - Note that the CTA tile size is chosen to maximize the GPU register\n// utilization, but carefully enough to avoid local memory use.\n\n#include \n#include \n#include \n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Externally configurable parameters.\n\n#ifndef CPU_DEBUG\n// Set this to 1 to verify the correctness of the GPU-computed matrix.\n#define CPU_DEBUG 0\n#endif\n\n#ifndef SHARED_MEMORY_LIMIT_64K\n// Set this to 0 to use more than 64 Kb of shared memory to cache data, to\n// improve the performance of the computations on GPU.\n// Note that you need a GPU that can have more than 64 Kb of shared memory\n// per multiprocessor.\n#define SHARED_MEMORY_LIMIT_64K 1\n#endif\n\n// GPU configuration.\n\n#define WARP_SIZE 32\n\n// MMA matrix tile dimensions.\n\n#define M 16\n#define N 16\n#define K 16\n\n#define WMMA_M 16\n#define WMMA_N 16\n#define WMMA_K 16\n\n// GEMM configuration.\n\n#define M_TILES 256\n#define N_TILES 256\n#define K_TILES 256\n\n#define M_GLOBAL (M * M_TILES)\n#define N_GLOBAL (N * N_TILES)\n#define K_GLOBAL (K * K_TILES)\n\n#define C_LAYOUT wmma::mem_row_major\n\n// Implementation constants.\n\n#define WARPS_PER_BLOCK 8\n#define THREADS_PER_BLOCK (WARP_SIZE * WARPS_PER_BLOCK)\n\n#if SHARED_MEMORY_LIMIT_64K\n// With only 64 Kb shared memory available, we can fit two 8-tile chunks of\n// the A and B matrix data, that are 16 * 16 * 8 * 8 * 2 = 32 Kb each\n// (i.e. two 8x8 arrays of tiles of 16x16 half-typed elements per CTA).\n// But we cannot account the 8 Kb total skew overhead, without which the\n// performance would be severely impacted. So we choose to reduce the chunk size\n// in half, i.e. the amount of A and B matrix data we cache in shared memory.\n// Accordingly, this doubles the number of outer iterations across the global K\n// dimension, which only slightly impacts the performance.\n#define CHUNK_K 4\n#else\n#define CHUNK_K 8\n#endif\n\n#define CHUNK_LINE_BYTES (CHUNK_K * K * sizeof(half))\n#define WARP_COPY_BYTES (WARP_SIZE * sizeof(int4))\n#define CHUNK_COPY_LINES_PER_WARP (WARP_COPY_BYTES / CHUNK_LINE_BYTES)\n#define CHUNK_COPY_LINE_LANES (WARP_SIZE / CHUNK_COPY_LINES_PER_WARP)\n\n#define BLOCK_ROW_WARPS 2\n#define BLOCK_COL_WARPS 4\n\n#define WARP_ROW_TILES 4\n#define WARP_COL_TILES 2\n\n#define BLOCK_ROW_TILES (WARP_ROW_TILES * BLOCK_ROW_WARPS)\n#define BLOCK_COL_TILES (WARP_COL_TILES * BLOCK_COL_WARPS)\n\n#define GLOBAL_MEM_STRIDE N_GLOBAL\n\n#define SHMEM_STRIDE (N * BLOCK_ROW_TILES)\n#define SHMEM_OFFSET (N * WARP_ROW_TILES)\n\n// The macro below is used to shift rows of the A matrix and columns of the B matrix\n// in shared memory to minimize possible bank conflicts.\n// Before performing the nvcuda::wmma::mma_sync operation, the warp must load the matrix\n// data using the nvcuda::wmma::load_matrix_sync operation. Although the memory access pattern\n// is not specified for that function, each lane in the warp can read one or multiple matrix\n// elements from different matrix rows or columns.\n// For shared memory, such access can result in bank conflicts if different rows / columns\n// of the matrix map to the same bank. By shifting each row and column by a few bytes, we\n// make sure that they map to different banks, thus reducing the number of possible bank\n// conflicts.\n// The number of 16 two-byte \"half\" elements is chosen as the minimum possible shift because\n// we must keep each row and column 256-bit aligned, as required by nvcuda::wmma::load_matrix_sync.\n#define SKEW_HALF 16\n\n#define checkKernelErrors(expr) \\\n do { \\\n expr; \\\n \\\n hipError_t __err = hipGetLastError(); \\\n if (__err != hipSuccess) { \\\n printf(\"Line %d: '%s' failed: %s\\n\", __LINE__, #expr, hipGetErrorString(__err)); \\\n abort(); \\\n } \\\n } while (0)\n\nusing namespace nvcuda;\n\n__host__ void init_host_matrices(half *a, half *b, float *c)\n{\n for (int i = 0; i < M_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n a[i * K_GLOBAL + j] = (half)(rand() % 3);\n }\n }\n\n for (int i = 0; i < N_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n b[i * K_GLOBAL + j] = (half)(rand() % 3);\n }\n }\n\n for (int t = 0; t < M_GLOBAL * N_GLOBAL; t++) {\n c[t] = static_cast(rand() % 3);\n }\n}\n\n__global__ void compute_gemm(const half *A, const half *B, const float *C, float *D, float alpha, float beta)\n{\n extern __shared__ half shmem[][CHUNK_K * K + SKEW_HALF];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n const size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n float *shmem_warp_tile_ptr =\n (float *)&shmem[0][0] + (warpId / 2) * SHMEM_STRIDE * K * 2 + (warpId % 2) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and\n // from shared memory.\n float *shmem_warp_stream_ptr = (float *)&shmem[0][0] + warpId * SHMEM_STRIDE * K;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may\n // result in a loss of precision). Zero still needs to be specially handled\n // though.\n beta /= alpha;\n\n // Each CTA slides along the 128 x 128 tiles from the top left corner of the\n // matrix to the right and down, and selects the next tile to compute. Once\n // there's no such tile, all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_ROW_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared\n // memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const float *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < K; i++) {\n typedef int4 copy_t;\n\n *((copy_t *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId) =\n *((copy_t *)(src_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment\n // multiplications along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * K + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Scale the C matrix.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const half *warp_ptr = (warpId < 4) ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % 4) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % 4) * 2);\n\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy\n // the B matrix.\n size_t shmem_idx = warpId < (WARPS_PER_BLOCK / 2)\n ? (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (N * (warpId % (WARPS_PER_BLOCK / 2)) * 2 + shmem_idx_b_off);\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n int4 *lane_ptr = (int4 *)(warp_ptr + tile_k * K + (laneId / CHUNK_COPY_LINE_LANES) * K_GLOBAL)\n + (laneId % CHUNK_COPY_LINE_LANES);\n\n // Shift the second half of the warp to the next row / column in the\n // shared memory.\n shmem_idx += laneId / CHUNK_COPY_LINE_LANES;\n\n#pragma unroll\n for (int i = 0; i < ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP) * 2; i++) {\n // Copy 16 bytes at once in each lane.\n *((int4 *)&shmem[shmem_idx][0] + (laneId % CHUNK_COPY_LINE_LANES)) = *lane_ptr;\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = (int4 *)((half *)lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP);\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / 2) * M * 2 + (i * M);\n const half *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_HALF);\n\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be\n // reused against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const half *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_HALF);\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n // Uniform, point-wise transformations of ALL fragment elements by ALL\n // threads in the warp are well-defined even though element indices\n // within fragment storage are not defined.\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n float *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * K + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global\n // memory.\n float *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < K; i++) {\n *((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n}\n\n// Performs an MxNxK GEMM (C=alpha*A*B + beta*C) assuming:\n// 1) Matrices are packed in memory.\n// 2) M, N and K are multiples of 16.\n// 3) Neither A nor B are transposed.\n// Note: This is a less performant version of the compute_gemm kernel. It is\n// designed for\n// demonstration purposes only to show the CUDA WMMA API use without\n// relying on availability of the shared memory.\n__global__ void\nsimple_wmma_gemm(half *a, half *b, float *c, float *d, int m_ld, int n_ld, int k_ld, float alpha, float beta)\n{\n // Leading dimensions. Packed with no transpositions.\n int lda = k_ld;\n int ldb = k_ld;\n int ldc = n_ld;\n\n // Tile using a 2D grid\n int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;\n int warpN = (blockIdx.y * blockDim.y + threadIdx.y);\n\n // Declare the fragments\n wmma::fragment a_frag;\n wmma::fragment b_frag;\n wmma::fragment acc_frag;\n wmma::fragment c_frag;\n\n wmma::fill_fragment(acc_frag, 0.0f);\n\n // Loop over k\n for (int i = 0; i < k_ld; i += WMMA_K) {\n int aCol = i;\n int aRow = warpM * WMMA_M;\n int bCol = warpN * N;\n int bRow = i;\n\n // Bounds checking\n if (aRow < m_ld && aCol < k_ld && bRow < k_ld && bCol < n_ld) {\n // Load the inputs\n wmma::load_matrix_sync(a_frag, a + aCol + aRow * lda, lda);\n wmma::load_matrix_sync(b_frag, b + bRow + bCol * ldb, ldb);\n\n // Perform the matrix multiplication\n wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);\n }\n }\n\n // Load in the current value of c, scale it by beta, and add this our result\n // scaled by alpha\n int cCol = warpN * WMMA_N;\n int cRow = warpM * WMMA_M;\n\n if (cRow < m_ld && cCol < n_ld) {\n wmma::load_matrix_sync(c_frag, c + cCol + cRow * ldc, ldc, wmma::mem_row_major);\n\n for (int i = 0; i < c_frag.num_elements; i++) {\n c_frag.x[i] = alpha * acc_frag.x[i] + beta * c_frag.x[i];\n }\n\n // Store the output\n wmma::store_matrix_sync(d + cCol + cRow * ldc, c_frag, ldc, wmma::mem_row_major);\n }\n}\n\n__host__ void matMultiplyOnHost(half *A,\n half *B,\n float *C,\n float alpha,\n float beta,\n int numARows,\n int numAColumns,\n int numBRows,\n int numBColumns,\n int numCRows,\n int numCColumns)\n{\n for (int i = 0; i < numCRows; i++) {\n for (int j = 0; j < numCColumns; j++) {\n float temp = 0.0;\n\n for (int k = 0; k < numAColumns; k++) {\n temp += (float)A[i * numAColumns + k] * (float)B[j * numBRows + k];\n }\n\n C[i * numCColumns + j] = temp * alpha + beta * C[i * numCColumns + j];\n }\n }\n}\n\nint main(int argc, char **argv)\n{\n printf(\"Initializing...\\n\");\n\n int dev = findCudaDevice(argc, (const char **)argv);\n\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, dev));\n\n // Tensor cores require a GPU of Volta (SM7X) architecture or higher.\n if (deviceProp.major < 7) {\n printf(\"cudaTensorCoreGemm requires SM 7.0 or higher to use Tensor \"\n \"Cores. Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n printf(\"M: %d (%d x %d)\\n\", M_GLOBAL, M, M_TILES);\n printf(\"N: %d (%d x %d)\\n\", N_GLOBAL, N, N_TILES);\n printf(\"K: %d (%d x %d)\\n\", K_GLOBAL, K, K_TILES);\n\n half *A_h = NULL;\n half *B_h = NULL;\n float *C_h = NULL;\n#if CPU_DEBUG\n float *result_hD = NULL;\n float *result_host = NULL;\n#endif\n\n A_h = (half *)malloc(sizeof(half) * M_GLOBAL * K_GLOBAL);\n B_h = (half *)malloc(sizeof(half) * K_GLOBAL * N_GLOBAL);\n C_h = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n#if CPU_DEBUG\n result_hD = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n result_host = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);\n#endif\n\n half *A = NULL;\n half *B = NULL;\n float *C = NULL;\n float *D = NULL;\n\n checkCudaErrors(hipMalloc(reinterpret_cast(&A), sizeof(half) * M_GLOBAL * K_GLOBAL));\n checkCudaErrors(hipMalloc(reinterpret_cast(&B), sizeof(half) * N_GLOBAL * K_GLOBAL));\n checkCudaErrors(hipMalloc(reinterpret_cast(&C), sizeof(float) * M_GLOBAL * N_GLOBAL));\n checkCudaErrors(hipMalloc(reinterpret_cast(&D), sizeof(float) * M_GLOBAL * N_GLOBAL));\n\n assert(((unsigned long long)A) % 128 == 0);\n assert(((unsigned long long)B) % 128 == 0);\n assert(((unsigned long long)C) % 128 == 0);\n assert(((unsigned long long)D) % 128 == 0);\n\n init_host_matrices(A_h, B_h, C_h);\n\n printf(\"Preparing data for GPU...\\n\");\n\n checkCudaErrors(hipMemcpy(A, A_h, sizeof(half) * M_GLOBAL * K_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(B, B_h, sizeof(half) * N_GLOBAL * K_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(C, C_h, sizeof(float) * M_GLOBAL * N_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemset(D, 0, sizeof(float) * M_GLOBAL * N_GLOBAL));\n\n enum {\n // Compute the right amount of shared memory to request.\n // We need shared memory to hold per-CTA C and D matrix tiles, and to cache\n // per-CTA chunks\n // of the A and B matrices. Therefore, the right amount to request is the\n // maximum of those\n // two numbers.\n SHMEM_SZ = MAX(sizeof(half) * (BLOCK_COL_TILES * M) * (CHUNK_K * K + SKEW_HALF) * 2,\n M * (BLOCK_ROW_WARPS * WARP_ROW_TILES) * N * (BLOCK_COL_WARPS * WARP_COL_TILES) * sizeof(float))\n };\n\n printf(\"Required shared memory size: %lu Kb\\n\", SHMEM_SZ / 1024UL);\n\n const float alpha = 1.1f;\n const float beta = 1.2f;\n\n hipEvent_t start, stop;\n\n checkCudaErrors(hipEventCreate(&start));\n checkCudaErrors(hipEventCreate(&stop));\n checkCudaErrors(hipEventRecord(start));\n\n // If enough shared memory available on the GPU use high performant kernel\n if (deviceProp.sharedMemPerMultiprocessor >= SHMEM_SZ) {\n printf(\"Computing... using high performance kernel compute_gemm \\n\");\n\n checkCudaErrors(hipFuncSetAttribute(reinterpret_cast(compute_gemm), hipFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors(\n (compute_gemm<<>>(A, B, C, D, alpha, beta)));\n#if CPU_DEBUG\n checkCudaErrors(hipMemcpy(result_hD, D, sizeof(float) * M_GLOBAL * N_GLOBAL, hipMemcpyDeviceToHost));\n#endif\n }\n else {\n dim3 gridDim;\n dim3 blockDim;\n\n // blockDim.x must be a multple of warpSize\n // 128x4 means we have 16 warps and a block computes a 64x64 output tile\n blockDim.x = 128;\n blockDim.y = 4;\n\n gridDim.x = (M_GLOBAL + (WMMA_M * blockDim.x / 32 - 1)) / (WMMA_M * blockDim.x / 32);\n gridDim.y = (N_GLOBAL + WMMA_N * blockDim.y - 1) / (WMMA_N * blockDim.y);\n\n printf(\"Computing... using simple_wmma_gemm kernel\\n\");\n simple_wmma_gemm<<>>(A, B, C, D, M_GLOBAL, N_GLOBAL, K_GLOBAL, alpha, beta);\n#if CPU_DEBUG\n checkCudaErrors(hipMemcpy(result_hD, D, sizeof(float) * M_GLOBAL * N_GLOBAL, hipMemcpyDeviceToHost));\n#endif\n }\n\n checkCudaErrors(hipEventRecord(stop));\n checkCudaErrors(hipEventSynchronize(stop));\n\n#if CPU_DEBUG\n printf(\"Verifying correctness of the computations...\\n\");\n\n memcpy(result_host, C_h, sizeof(float) * M_GLOBAL * N_GLOBAL);\n\n matMultiplyOnHost(A_h, B_h, result_host, alpha, beta, M_GLOBAL, K_GLOBAL, K_GLOBAL, N_GLOBAL, M_GLOBAL, N_GLOBAL);\n\n for (int i = 0; i < N_GLOBAL * M_GLOBAL; i++) {\n if (fabs(result_hD[i] - result_host[i]) > 0.1f)\n printf(\"mismatch i=%d result_hD=%f result_host=%f\\n\", i, result_hD[i], result_host[i]);\n }\n free(result_hD);\n free(result_host);\n#endif\n\n float milliseconds = 0;\n\n checkCudaErrors(hipEventElapsedTime(&milliseconds, start, stop));\n\n printf(\"Time: %f ms\\n\", milliseconds);\n printf(\"TFLOPS: %.2f\\n\",\n static_cast((static_cast(M_GLOBAL) * N_GLOBAL * K_GLOBAL * 2) / (milliseconds / 1000.))\n / 1e12);\n\n free(A_h);\n free(B_h);\n free(C_h);\n checkCudaErrors(hipFree(reinterpret_cast(A)));\n checkCudaErrors(hipFree(reinterpret_cast(B)));\n checkCudaErrors(hipFree(reinterpret_cast(C)));\n checkCudaErrors(hipFree(reinterpret_cast(D)));\n\n return 0;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/cudaTensorCoreGemm/cudaTensorCoreGemm.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n\n// CUDA sample demonstrating a Double precision GEMM computation using the Warp\n// Matrix Multiply and Accumulate API introduced in CUDA 11.0.\n\n// In this program, the compute_dgemm kernel computes the result of a matrix multiplication\n// and addition: D = alpha * A * B + beta * C. The dimensions of both C and D matrices\n// are M_GLOBAL x N_GLOBAL. The A matrix is M_GLOBAL x K_GLOBAL (row-major), the B matrix\n// is K_GLOBAL x N_GLOBAL (column-major).\n// In that kernel, each CTA computes one 64 x 64 tile of the resulting matrix\n// per iteration. When the tile is computed, the CTA stores it to the global memory\n// and begins a new iteration, selecting a new 64 x 64 tile to compute.\n// Each CTA consists of eight warps. For the 64 x 64 tile, each warp computes eight\n// 8 x 8 subtiles, organized in a 2 x 4 two-dimensional array.\n// Warps compute the 8 x 8 subtiles using nvcuda::wmma::mma_sync operations by\n// moving through the K_GLOBAL dimension of the A and B matrices and accumulating\n// the intermediate result in the local thread state.\n\n// There are a number of simple optimizations used in the algorithm:\n// - The CTA copies the 64 x 64 tile of the C matrix from the global memory to\n// shared memory. After that is done, each warp loads the C matrix fragments from\n// shared memory, thus avoiding a random global memory access.\n// - On each internal iteration, the CTA copies a portion of the A and B matrices from\n// global memory to shared memory. After that, all warps in the CTA reuse the A and B\n// data from shared memory, thus reducing the number of data copies from global memory.\n// - The portions of the A and B matrices are stored in shared memory with an additional\n// padding (skew) to reduce the number of shared memory access bank conflicts.\n// (See a detailed explanation near the SKEW_DOUBLE macro definition.)\n// - When the CTA finishes computing the tiles of the resulting matrix, each warp stores\n// its subtiles to shared memory. The CTA then copies the shared memory contents to\n// global memory, again avoiding redundant random global memory accesses.\n// - Note that the CTA tile size is chosen to maximize the GPU register utilization,\n// but carefully enough to avoid local memory use.\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Externally configurable parameters.\n\n#ifndef CPU_DEBUG\n// Set this to 1 to verify the correctness of the GPU-computed matrix.\n#define CPU_DEBUG 0\n#endif\n\n#ifndef SHARED_MEMORY_LIMIT_64K\n// Set this to 0 to use more than 64 Kb of shared memory to cache data, to\n// improve the performance of the computations on GPU.\n// Note that you need a GPU that can have more than 64 Kb of shared memory\n// per multiprocessor.\n#define SHARED_MEMORY_LIMIT_64K 0\n#endif\n\n// GPU configuration.\n\n#define WARP_SIZE 32\n\n// MMA matrix tile dimensions.\n\n#define M 8\n#define N 8\n#define K 4\n\n// GEMM configuration.\n\n#define M_TILES 1024\n#define N_TILES 1024\n#define K_TILES 1024\n\n#define M_GLOBAL (M * M_TILES)\n#define N_GLOBAL (N * N_TILES)\n#define K_GLOBAL (K * K_TILES)\n\n#define C_LAYOUT wmma::mem_row_major\n\n// Implementation constants.\n\n#define WARPS_PER_BLOCK 8\n#define THREADS_PER_BLOCK (WARP_SIZE * WARPS_PER_BLOCK)\n\n#if SHARED_MEMORY_LIMIT_64K\n// With only 64 Kb shared memory available, we can fit 8x16-tile chunks of each\n// the A and B matrix data, that are (M = 8) * (K = 4) * 8 * (CHUNK_K = 16) * sizeof(double) = 32 Kb each\n// But we cannot account the 4 Kb total skew overhead, without which the performance\n// would be severely impacted. So we choose to reduce the chunk size in half,\n// i.e. the amount of A and B matrix data we cache in shared memory.\n// Accordingly, this doubles the number of outer iterations across the global K\n// dimension, which only slightly impacts the performance.\n#define CHUNK_K 8\n#else\n#define CHUNK_K 16\n#endif\n\n#define CHUNK_LINE_BYTES (CHUNK_K * K * sizeof(double))\n#define WARP_COPY_BYTES (WARP_SIZE * sizeof(int4))\n#define CHUNK_COPY_LINES_PER_WARP (WARP_COPY_BYTES / CHUNK_LINE_BYTES)\n#define CHUNK_COPY_LINE_LANES (WARP_SIZE / CHUNK_COPY_LINES_PER_WARP)\n\n#define BLOCK_ROW_WARPS 2\n#define BLOCK_COL_WARPS 4\n\n#define WARP_ROW_TILES 4\n#define WARP_COL_TILES 2\n\n#define BLOCK_ROW_TILES (WARP_ROW_TILES * BLOCK_ROW_WARPS)\n#define BLOCK_COL_TILES (WARP_COL_TILES * BLOCK_COL_WARPS)\n\n#define GLOBAL_MEM_STRIDE N_GLOBAL\n\n#define SHMEM_STRIDE (N * BLOCK_ROW_TILES)\n#define SHMEM_OFFSET (N * WARP_ROW_TILES)\n\n// The macro below is used to shift rows of the A matrix and columns of the B matrix\n// in shared memory to minimize possible bank conflicts.\n// Before performing the nvcuda::wmma::mma_sync operation, the warp must load the matrix\n// data using the nvcuda::wmma::load_matrix_sync operation. Although the memory access pattern\n// is not specified for that function, each lane in the warp can read one or multiple matrix\n// elements from different matrix rows or columns.\n// For shared memory, such access can result in bank conflicts if different rows / columns\n// of the matrix map to the same bank. By shifting each row and column by a few bytes, we\n// make sure that they map to different banks, thus reducing the number of possible bank\n// conflicts.\n// The number of 4 eight-byte \"double\" elements is chosen as the minimum possible shift because\n// we must keep each row and column 256-bit aligned, as required by nvcuda::wmma::load_matrix_sync.\n#define SKEW_DOUBLE 4\n\n#define checkKernelErrors(expr) \\\n do { \\\n expr; \\\n \\\n cudaError_t __err = cudaGetLastError(); \\\n if (__err != cudaSuccess) { \\\n printf(\"Line %d: '%s' failed: %s\\n\", __LINE__, #expr, cudaGetErrorString(__err)); \\\n abort(); \\\n } \\\n } while (0)\n\nenum kernels {\n dmma_shmem_gemm_async_copy = 0, // DMMA shmem using kernel with async_copy\n dmma_shmem_gemm_cg_async_copy = 1, // DMMA shmem using kernel with cooperative groups async_copy\n dmma_shmem_gemm = 2, // DMMA shmem using kernel normal copy (without async_copy).\n simple_dmma_gemm = 3 // DMMA non-shmem using simple kernel.\n};\n\nconst char *kernelNames[] = {\"compute_dgemm_async_copy\",\n \"compute_dgemm_cg_async_copy\",\n \"compute_dgemm\",\n \"simple_wmma_gemm\"};\n\nusing namespace nvcuda;\nnamespace cg = cooperative_groups;\n\n__host__ void init_host_matrices(double *a, double *b, double *c)\n{\n for (int i = 0; i < M_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n a[i * K_GLOBAL + j] = (double)(rand() % 3);\n }\n }\n\n for (int i = 0; i < N_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n b[i * K_GLOBAL + j] = (double)(rand() % 3);\n }\n }\n\n for (int t = 0; t < M_GLOBAL * N_GLOBAL; t++) {\n c[t] = (double)(rand() % 3);\n }\n}\n\n__global__ void compute_dgemm(const double *A, const double *B, const double *C, double *D, double alpha, double beta)\n{\n#if __CUDA_ARCH__ >= 800\n extern __shared__ double shmem[][CHUNK_K * K + SKEW_DOUBLE];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n const size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n double *shmem_warp_tile_ptr = (double *)&shmem[0][0]\n + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS\n + (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.\n double *shmem_warp_stream_ptr = (double *)&shmem[0][0] + warpId * SHMEM_STRIDE * N;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may result\n // in a loss of precision). Zero still needs to be specially handled though.\n beta /= alpha;\n\n // Each CTA slides along the 64 x 64 tiles from the top left corner of the matrix to the\n // right and down, and selects the next tile to compute. Once there's no such tile,\n // all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const double *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId) =\n *((int4 *)(src_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment multiplications\n // along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const double *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Scale the C matrix.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const double *warp_ptr =\n (warpId < (WARPS_PER_BLOCK / 2))\n ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);\n\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.\n size_t shmem_idx = warpId < (WARPS_PER_BLOCK / 2)\n ? (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (N * (warpId % (WARPS_PER_BLOCK / 2)) * 2 + shmem_idx_b_off);\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n const double *lane_ptr = warp_ptr + tile_k * K + (laneId / CHUNK_COPY_LINE_LANES) * K_GLOBAL;\n\n // Shift the second half of the warp to the next row / column in the shared memory.\n shmem_idx += laneId / CHUNK_COPY_LINE_LANES;\n\n#pragma unroll\n for (int i = 0; i < ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP); i++) {\n // Copy 16 bytes at once in each lane.\n *((int4 *)&shmem[shmem_idx][0] + (laneId % CHUNK_COPY_LINE_LANES)) =\n *((int4 *)lane_ptr + (laneId % CHUNK_COPY_LINE_LANES));\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / 2) * M * 2 + (i * M);\n const double *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_DOUBLE);\n\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be reused\n // against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const double *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_DOUBLE);\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n // Uniform, point-wise transformations of ALL fragment elements by ALL threads in the\n // warp are well-defined even though element indices within fragment storage are not defined.\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n double *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global memory.\n double *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n#endif\n}\n\n__global__ void\ncompute_dgemm_async_copy(const double *A, const double *B, const double *C, double *D, double alpha, double beta)\n{\n#if __CUDA_ARCH__ >= 800\n extern __shared__ double shmem[][CHUNK_K * K + SKEW_DOUBLE];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n constexpr size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n double *shmem_warp_tile_ptr = &shmem[0][0] + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS\n + (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.\n double *shmem_warp_stream_ptr = &shmem[0][0] + warpId * SHMEM_STRIDE * N;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may result\n // in a loss of precision). Zero still needs to be specially handled though.\n beta /= alpha;\n\n cuda::pipeline pipe = cuda::make_pipeline();\n\n const auto shape2 = cuda::aligned_size_t(sizeof(double2));\n constexpr int loadStride = 1; // load 2 double, left-shift by 1.\n\n // Each CTA slides along the 64 x 64 tiles from the top left corner of the matrix to the\n // right and down, and selects the next tile to compute. Once there's no such tile,\n // all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const double *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < N; i++) {\n pipe.producer_acquire();\n cuda::memcpy_async(&shmem_warp_stream_ptr[(SHMEM_STRIDE * i) + (laneId << loadStride)],\n &src_gmem_warp_stream_ptr[(GLOBAL_MEM_STRIDE * i) + (laneId << loadStride)],\n shape2,\n pipe);\n\n pipe.producer_commit();\n }\n // Now wait for all the above issued 8 batches to complete.\n cuda::pipeline_consumer_wait_prior<0>(pipe);\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment multiplications\n // along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const double *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n // Scale the C matrix.\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n pipe.consumer_release();\n // sync here so that shared memory can then be used for loading A & B matrices.\n __syncthreads();\n\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const double *warp_ptr =\n (warpId < (WARPS_PER_BLOCK / 2))\n ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);\n\n const int stridePerLaneCopy = (laneId / CHUNK_COPY_LINE_LANES);\n constexpr int chunksPerLane = ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP);\n const int laneLoadElem = (laneId % CHUNK_COPY_LINE_LANES) << loadStride;\n\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.\n // As for DMMA M == N we use M for warp 4-7 + shmem_idx_b_off.\n size_t shmem_idx =\n (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2) + (shmem_idx_b_off * (warpId / (WARPS_PER_BLOCK / 2)));\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n const double *lane_ptr = warp_ptr + tile_k * K + stridePerLaneCopy * K_GLOBAL + laneLoadElem;\n\n // Shift the second half of the warp to the next row / column in the shared memory.\n shmem_idx += stridePerLaneCopy;\n#pragma unroll\n for (int i = 0; i < chunksPerLane; i++) {\n // Copy 16 bytes at once in each lane.\n pipe.producer_acquire();\n\n cuda::memcpy_async(&shmem[shmem_idx][laneLoadElem], lane_ptr, shape2, pipe);\n\n pipe.producer_commit();\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n\n cuda::pipeline_consumer_wait_prior<0>(pipe);\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / 2) * M * 2 + (i * M);\n const double *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_DOUBLE);\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be reused\n // against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const double *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_DOUBLE);\n }\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n pipe.consumer_release();\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n // Uniform, point-wise transformations of ALL fragment elements by ALL threads in the\n // warp are well-defined even though element indices within fragment storage are not defined.\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n double *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global memory.\n double *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n#endif\n}\n\n__global__ void\ncompute_dgemm_cg_async_copy(const double *A, const double *B, const double *C, double *D, double alpha, double beta)\n{\n#if __CUDA_ARCH__ >= 800\n extern __shared__ double shmem[][CHUNK_K * K + SKEW_DOUBLE];\n auto cta = cg::this_thread_block();\n auto tile32 = cg::tiled_partition<32>(cta);\n\n constexpr int tileChunkCopySize = WARP_SIZE / CHUNK_COPY_LINES_PER_WARP;\n auto tileChunkCopy = cg::tiled_partition(cta);\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n constexpr size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n double *shmem_warp_tile_ptr =\n (double *)&shmem[0][0] + (warpId / 2) * SHMEM_STRIDE * N * 2 + (warpId % 2) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.\n double *shmem_warp_stream_ptr = (double *)&shmem[0][0] + warpId * SHMEM_STRIDE * N;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may result\n // in a loss of precision). Zero still needs to be specially handled though.\n beta /= alpha;\n\n // Each CTA slides along the 64 x 64 tiles from the top left corner of the matrix to the\n // right and down, and selects the next tile to compute. Once there's no such tile,\n // all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const double *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < N; i++) {\n auto dst_ptr = &shmem_warp_stream_ptr[(SHMEM_STRIDE * i)];\n auto src_ptr = &src_gmem_warp_stream_ptr[(GLOBAL_MEM_STRIDE * i)];\n cg::memcpy_async(\n tile32, dst_ptr, src_ptr, cuda::aligned_size_t{tile32.size() * sizeof(double2)});\n }\n\n cg::wait(cta);\n\n // These fragments will accumulate the result of A and B matrix fragment multiplications\n // along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const double *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n // Scale the C matrix.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n // sync here so that shared memory can then be used for loading A & B matrices.\n cg::wait(cta);\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const double *warp_ptr =\n (warpId < 4) ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);\n\n const int stridePerLaneCopy = (laneId / CHUNK_COPY_LINE_LANES);\n constexpr int chunksPerLane = ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP);\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.\n // As for DMMA M == N we use M for warp 4-7 + shmem_idx_b_off.\n size_t shmem_idx =\n (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2) + (shmem_idx_b_off * (warpId / (WARPS_PER_BLOCK / 2)));\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n auto lane_ptr = warp_ptr + tile_k * K + stridePerLaneCopy * K_GLOBAL;\n\n // Shift the second half of the warp to the next row / column in the shared memory.\n shmem_idx += stridePerLaneCopy;\n\n#pragma unroll\n for (int i = 0; i < chunksPerLane; i++) {\n // Copy 16 bytes at once in each lane.\n auto dst_ptr = &shmem[shmem_idx][0];\n auto src_ptr = lane_ptr;\n\n cg::memcpy_async(tileChunkCopy,\n dst_ptr,\n src_ptr,\n cuda::aligned_size_t{tileChunkCopySize * sizeof(double2)});\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n cg::wait(cta);\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / 2) * M * 2 + (i * M);\n const double *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_DOUBLE);\n\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be reused\n // against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const double *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_DOUBLE);\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n cg::sync(cta);\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n // Uniform, point-wise transformations of ALL fragment elements by ALL threads in the\n // warp are well-defined even though element indices within fragment storage are not defined.\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n double *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n cg::sync(cta);\n\n // Now that shared memory contains all the D tiles, stream them to global memory.\n double *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n cg::sync(cta);\n }\n#endif\n}\n\n// Performs an MxNxK DGEMM (C=alpha*A*B + beta*C) assuming:\n// 1) Matrices are packed in memory.\n// 2) M, N and K are multiples of 8, 8 and 4 respectively.\n// 3) A is row major, B is column major matrix.\n// Note: This is a less performant version of the compute_dgemm kernel. It is designed for\n// demonstration purposes only to show the CUDA WMMA API use without relying on\n// availability of the shared memory.\n__global__ void\nsimple_wmma_gemm(double *a, double *b, double *c, double *d, int m_ld, int n_ld, int k_ld, double alpha, double beta)\n{\n#if __CUDA_ARCH__ >= 800\n // Leading dimensions. Packed with no transpositions.\n int lda = k_ld;\n int ldb = k_ld;\n int ldc = n_ld;\n\n // Tile using a 2D grid\n int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;\n int warpN = (blockIdx.y * blockDim.y + threadIdx.y);\n\n // Declare the fragments\n wmma::fragment a_frag;\n wmma::fragment b_frag;\n wmma::fragment acc_frag;\n wmma::fragment c_frag;\n\n wmma::fill_fragment(acc_frag, 0.0f);\n\n // Loop over k\n for (int i = 0; i < k_ld; i += K) {\n int aCol = i;\n int aRow = warpM * M;\n\n int bCol = warpN * N;\n int bRow = i;\n\n // Bounds checking\n if (aRow < m_ld && aCol < k_ld && bRow < k_ld && bCol < n_ld) {\n // Load the inputs\n wmma::load_matrix_sync(a_frag, a + aCol + aRow * lda, lda);\n wmma::load_matrix_sync(b_frag, b + bRow + bCol * ldb, ldb);\n\n // Perform the matrix multiplication\n wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);\n }\n }\n\n // Load in the current value of c, scale it by beta, and add this our result scaled by alpha\n int cCol = warpN * N;\n int cRow = warpM * M;\n\n if (cRow < m_ld && cCol < n_ld) {\n wmma::load_matrix_sync(c_frag, c + cCol + cRow * ldc, ldc, wmma::mem_row_major);\n\n for (int i = 0; i < c_frag.num_elements; i++) {\n c_frag.x[i] = alpha * acc_frag.x[i] + beta * c_frag.x[i];\n }\n\n // Store the output\n wmma::store_matrix_sync(d + cCol + cRow * ldc, c_frag, ldc, wmma::mem_row_major);\n }\n#endif\n}\n\n__host__ void matMultiplyOnHost(double *A,\n double *B,\n double *C,\n float alpha,\n float beta,\n int numARows,\n int numAColumns,\n int numBRows,\n int numBColumns,\n int numCRows,\n int numCColumns)\n{\n for (int i = 0; i < numCRows; i++) {\n for (int j = 0; j < numCColumns; j++) {\n double temp = 0.0;\n\n for (int k = 0; k < numAColumns; k++) {\n // B matrix is column major. A matrix is row major.\n temp += A[i * numAColumns + k] * B[j * numBRows + k];\n }\n\n C[i * numCColumns + j] = temp * alpha + beta * C[i * numCColumns + j];\n }\n }\n}\n\nint main(int argc, char **argv)\n{\n printf(\"Initializing...\\n\");\n\n int dev = findCudaDevice(argc, (const char **)argv);\n\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, dev));\n\n // Double precision Tensor cores require a GPU of Ampere (SM8X) architecture or higher.\n if (deviceProp.major < 8) {\n printf(\"dmmaTensorCoreGemm requires SM 8.0 or higher. Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n printf(\"M: %d (%d x %d)\\n\", M_GLOBAL, M, M_TILES);\n printf(\"N: %d (%d x %d)\\n\", N_GLOBAL, N, N_TILES);\n printf(\"K: %d (%d x %d)\\n\", K_GLOBAL, K, K_TILES);\n\n double *A_h = NULL;\n double *B_h = NULL;\n double *C_h = NULL;\n#if CPU_DEBUG\n double *result_hD = NULL;\n double *result_host = NULL;\n#endif\n\n A_h = (double *)malloc(sizeof(double) * M_GLOBAL * K_GLOBAL);\n B_h = (double *)malloc(sizeof(double) * K_GLOBAL * N_GLOBAL);\n C_h = (double *)malloc(sizeof(double) * M_GLOBAL * N_GLOBAL);\n#if CPU_DEBUG\n result_hD = (double *)malloc(sizeof(double) * M_GLOBAL * N_GLOBAL);\n result_host = (double *)malloc(sizeof(double) * M_GLOBAL * N_GLOBAL);\n#endif\n\n double *A = NULL;\n double *B = NULL;\n double *C = NULL;\n double *D = NULL;\n\n checkCudaErrors(cudaMalloc((void **)&A, sizeof(double) * M_GLOBAL * K_GLOBAL));\n checkCudaErrors(cudaMalloc((void **)&B, sizeof(double) * N_GLOBAL * K_GLOBAL));\n checkCudaErrors(cudaMalloc((void **)&C, sizeof(double) * M_GLOBAL * N_GLOBAL));\n checkCudaErrors(cudaMalloc((void **)&D, sizeof(double) * M_GLOBAL * N_GLOBAL));\n\n assert(((unsigned long long)A) % 128 == 0);\n assert(((unsigned long long)B) % 128 == 0);\n assert(((unsigned long long)C) % 128 == 0);\n assert(((unsigned long long)D) % 128 == 0);\n\n init_host_matrices(A_h, B_h, C_h);\n\n printf(\"Preparing data for GPU...\\n\");\n\n checkCudaErrors(cudaMemcpy(A, A_h, sizeof(double) * M_GLOBAL * K_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(B, B_h, sizeof(double) * N_GLOBAL * K_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(C, C_h, sizeof(double) * M_GLOBAL * N_GLOBAL, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemset(D, 0, sizeof(double) * M_GLOBAL * N_GLOBAL));\n\n enum {\n // Compute the right amount of shared memory to request.\n // We need shared memory to hold per-CTA C and D matrix tiles, and to cache per-CTA chunks\n // of the A and B matrices. Therefore, the right amount to request is the maximum of those\n // two numbers.\n SHMEM_SZ = MAX(sizeof(double) * (BLOCK_COL_TILES * M) * (CHUNK_K * K + SKEW_DOUBLE) * 2,\n M * (BLOCK_ROW_WARPS * WARP_ROW_TILES) * N * (BLOCK_COL_WARPS * WARP_COL_TILES) * sizeof(double))\n };\n\n printf(\"Required shared memory size: %lu Kb\\n\", SHMEM_SZ / 1024UL);\n\n const double alpha = 1.1f;\n const double beta = 1.2f;\n\n cudaEvent_t start, stop;\n\n checkCudaErrors(cudaEventCreate(&start));\n checkCudaErrors(cudaEventCreate(&stop));\n checkCudaErrors(cudaEventRecord(start));\n\n kernels selected_kernel = dmma_shmem_gemm_async_copy;\n\n // kernel to run - default (dmma_shmem_gemm_async_copy == 0)\n if (checkCmdLineFlag(argc, (const char **)argv, \"kernel\")) {\n int kernel_number = getCmdLineArgumentInt(argc, (const char **)argv, \"kernel\");\n if (kernel_number < 4) {\n selected_kernel = (kernels)kernel_number;\n }\n else {\n printf(\"Error: kernel number should be between 0 to 3, you have entered %d\\n\", kernel_number);\n exit(EXIT_FAILURE);\n }\n }\n\n // If enough shared memory available on the GPU use high performant kernel\n if ((deviceProp.sharedMemPerMultiprocessor >= SHMEM_SZ) && (selected_kernel != simple_dmma_gemm)) {\n printf(\"Computing using high performance kernel = %d - %s\\n\", selected_kernel, kernelNames[selected_kernel]);\n\n switch (selected_kernel) {\n case dmma_shmem_gemm_async_copy:\n default:\n checkCudaErrors(\n cudaFuncSetAttribute(compute_dgemm_async_copy, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors(\n (compute_dgemm_async_copy<<>>(\n A, B, C, D, alpha, beta)));\n break;\n case dmma_shmem_gemm_cg_async_copy:\n checkCudaErrors(cudaFuncSetAttribute(\n compute_dgemm_cg_async_copy, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors(\n (compute_dgemm_cg_async_copy<<>>(\n A, B, C, D, alpha, beta)));\n break;\n case dmma_shmem_gemm:\n checkCudaErrors(cudaFuncSetAttribute(compute_dgemm, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors((compute_dgemm<<>>(\n A, B, C, D, alpha, beta)));\n break;\n }\n\n#if CPU_DEBUG\n checkCudaErrors(cudaMemcpy(result_hD, D, sizeof(double) * M_GLOBAL * N_GLOBAL, cudaMemcpyDeviceToHost));\n#endif\n }\n else {\n dim3 gridDim;\n dim3 blockDim;\n\n // blockDim.x must be a multple of warpSize\n // 128x4 means we have 16 warps and a block computes a 64x64 output tile\n blockDim.x = 128;\n blockDim.y = 4;\n\n gridDim.x = (M_GLOBAL + (M * blockDim.x / 32 - 1)) / (M * blockDim.x / 32);\n gridDim.y = (N_GLOBAL + N * blockDim.y - 1) / (N * blockDim.y);\n\n printf(\"Computing... using simple_wmma_gemm kernel\\n\");\n simple_wmma_gemm<<>>(A, B, C, D, M_GLOBAL, N_GLOBAL, K_GLOBAL, alpha, beta);\n#if CPU_DEBUG\n checkCudaErrors(cudaMemcpy(result_hD, D, sizeof(double) * M_GLOBAL * N_GLOBAL, cudaMemcpyDeviceToHost));\n#endif\n }\n\n checkCudaErrors(cudaEventRecord(stop));\n checkCudaErrors(cudaEventSynchronize(stop));\n\n#if CPU_DEBUG\n printf(\"Verifying correctness of the computations...\\n\");\n\n memcpy(result_host, C_h, sizeof(double) * M_GLOBAL * N_GLOBAL);\n\n matMultiplyOnHost(A_h, B_h, result_host, alpha, beta, M_GLOBAL, K_GLOBAL, K_GLOBAL, N_GLOBAL, M_GLOBAL, N_GLOBAL);\n\n size_t number_of_matches = 0;\n for (int i = 0; i < N_GLOBAL * M_GLOBAL; i++) {\n if (fabs(result_hD[i] - result_host[i]) > 0.1f) {\n printf(\"mismatch i=%d result_hD=%f result_host=%f\\n\", i, result_hD[i], result_host[i]);\n break;\n }\n else {\n number_of_matches++;\n }\n }\n printf(\"number_of_matches = %zu out of = %d \\n\", number_of_matches, N_GLOBAL * M_GLOBAL);\n free(result_hD);\n free(result_host);\n#endif\n\n float milliseconds = 0;\n\n checkCudaErrors(cudaEventElapsedTime(&milliseconds, start, stop));\n\n printf(\"Time: %f ms\\n\", milliseconds);\n printf(\"FP64 TFLOPS: %.2f\\n\", (((double)M_GLOBAL * N_GLOBAL * K_GLOBAL * 2) / (milliseconds / 1000.)) / 1e12);\n\n free(A_h);\n free(B_h);\n free(C_h);\n checkCudaErrors(cudaFree((void *)A));\n checkCudaErrors(cudaFree((void *)B));\n checkCudaErrors(cudaFree((void *)C));\n checkCudaErrors(cudaFree((void *)D));\n\n return 0;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n\n// CUDA sample demonstrating a Double precision GEMM computation using the Warp\n// Matrix Multiply and Accumulate API introduced in CUDA 11.0.\n\n// In this program, the compute_dgemm kernel computes the result of a matrix multiplication\n// and addition: D = alpha * A * B + beta * C. The dimensions of both C and D matrices\n// are M_GLOBAL x N_GLOBAL. The A matrix is M_GLOBAL x K_GLOBAL (row-major), the B matrix\n// is K_GLOBAL x N_GLOBAL (column-major).\n// In that kernel, each CTA computes one 64 x 64 tile of the resulting matrix\n// per iteration. When the tile is computed, the CTA stores it to the global memory\n// and begins a new iteration, selecting a new 64 x 64 tile to compute.\n// Each CTA consists of eight warps. For the 64 x 64 tile, each warp computes eight\n// 8 x 8 subtiles, organized in a 2 x 4 two-dimensional array.\n// Warps compute the 8 x 8 subtiles using nvcuda::wmma::mma_sync operations by\n// moving through the K_GLOBAL dimension of the A and B matrices and accumulating\n// the intermediate result in the local thread state.\n\n// There are a number of simple optimizations used in the algorithm:\n// - The CTA copies the 64 x 64 tile of the C matrix from the global memory to\n// shared memory. After that is done, each warp loads the C matrix fragments from\n// shared memory, thus avoiding a random global memory access.\n// - On each internal iteration, the CTA copies a portion of the A and B matrices from\n// global memory to shared memory. After that, all warps in the CTA reuse the A and B\n// data from shared memory, thus reducing the number of data copies from global memory.\n// - The portions of the A and B matrices are stored in shared memory with an additional\n// padding (skew) to reduce the number of shared memory access bank conflicts.\n// (See a detailed explanation near the SKEW_DOUBLE macro definition.)\n// - When the CTA finishes computing the tiles of the resulting matrix, each warp stores\n// its subtiles to shared memory. The CTA then copies the shared memory contents to\n// global memory, again avoiding redundant random global memory accesses.\n// - Note that the CTA tile size is chosen to maximize the GPU register utilization,\n// but carefully enough to avoid local memory use.\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Externally configurable parameters.\n\n#ifndef CPU_DEBUG\n// Set this to 1 to verify the correctness of the GPU-computed matrix.\n#define CPU_DEBUG 0\n#endif\n\n#ifndef SHARED_MEMORY_LIMIT_64K\n// Set this to 0 to use more than 64 Kb of shared memory to cache data, to\n// improve the performance of the computations on GPU.\n// Note that you need a GPU that can have more than 64 Kb of shared memory\n// per multiprocessor.\n#define SHARED_MEMORY_LIMIT_64K 0\n#endif\n\n// GPU configuration.\n\n#define WARP_SIZE 32\n\n// MMA matrix tile dimensions.\n\n#define M 8\n#define N 8\n#define K 4\n\n// GEMM configuration.\n\n#define M_TILES 1024\n#define N_TILES 1024\n#define K_TILES 1024\n\n#define M_GLOBAL (M * M_TILES)\n#define N_GLOBAL (N * N_TILES)\n#define K_GLOBAL (K * K_TILES)\n\n#define C_LAYOUT wmma::mem_row_major\n\n// Implementation constants.\n\n#define WARPS_PER_BLOCK 8\n#define THREADS_PER_BLOCK (WARP_SIZE * WARPS_PER_BLOCK)\n\n#if SHARED_MEMORY_LIMIT_64K\n// With only 64 Kb shared memory available, we can fit 8x16-tile chunks of each\n// the A and B matrix data, that are (M = 8) * (K = 4) * 8 * (CHUNK_K = 16) * sizeof(double) = 32 Kb each\n// But we cannot account the 4 Kb total skew overhead, without which the performance\n// would be severely impacted. So we choose to reduce the chunk size in half,\n// i.e. the amount of A and B matrix data we cache in shared memory.\n// Accordingly, this doubles the number of outer iterations across the global K\n// dimension, which only slightly impacts the performance.\n#define CHUNK_K 8\n#else\n#define CHUNK_K 16\n#endif\n\n#define CHUNK_LINE_BYTES (CHUNK_K * K * sizeof(double))\n#define WARP_COPY_BYTES (WARP_SIZE * sizeof(int4))\n#define CHUNK_COPY_LINES_PER_WARP (WARP_COPY_BYTES / CHUNK_LINE_BYTES)\n#define CHUNK_COPY_LINE_LANES (WARP_SIZE / CHUNK_COPY_LINES_PER_WARP)\n\n#define BLOCK_ROW_WARPS 2\n#define BLOCK_COL_WARPS 4\n\n#define WARP_ROW_TILES 4\n#define WARP_COL_TILES 2\n\n#define BLOCK_ROW_TILES (WARP_ROW_TILES * BLOCK_ROW_WARPS)\n#define BLOCK_COL_TILES (WARP_COL_TILES * BLOCK_COL_WARPS)\n\n#define GLOBAL_MEM_STRIDE N_GLOBAL\n\n#define SHMEM_STRIDE (N * BLOCK_ROW_TILES)\n#define SHMEM_OFFSET (N * WARP_ROW_TILES)\n\n// The macro below is used to shift rows of the A matrix and columns of the B matrix\n// in shared memory to minimize possible bank conflicts.\n// Before performing the nvcuda::wmma::mma_sync operation, the warp must load the matrix\n// data using the nvcuda::wmma::load_matrix_sync operation. Although the memory access pattern\n// is not specified for that function, each lane in the warp can read one or multiple matrix\n// elements from different matrix rows or columns.\n// For shared memory, such access can result in bank conflicts if different rows / columns\n// of the matrix map to the same bank. By shifting each row and column by a few bytes, we\n// make sure that they map to different banks, thus reducing the number of possible bank\n// conflicts.\n// The number of 4 eight-byte \"double\" elements is chosen as the minimum possible shift because\n// we must keep each row and column 256-bit aligned, as required by nvcuda::wmma::load_matrix_sync.\n#define SKEW_DOUBLE 4\n\n#define checkKernelErrors(expr) \\\n do { \\\n expr; \\\n \\\n hipError_t __err = hipGetLastError(); \\\n if (__err != hipSuccess) { \\\n printf(\"Line %d: '%s' failed: %s\\n\", __LINE__, #expr, hipGetErrorString(__err)); \\\n abort(); \\\n } \\\n } while (0)\n\nenum kernels {\n dmma_shmem_gemm_async_copy = 0, // DMMA shmem using kernel with async_copy\n dmma_shmem_gemm_cg_async_copy = 1, // DMMA shmem using kernel with cooperative groups async_copy\n dmma_shmem_gemm = 2, // DMMA shmem using kernel normal copy (without async_copy).\n simple_dmma_gemm = 3 // DMMA non-shmem using simple kernel.\n};\n\nconst char *kernelNames[] = {\"compute_dgemm_async_copy\",\n \"compute_dgemm_cg_async_copy\",\n \"compute_dgemm\",\n \"simple_wmma_gemm\"};\n\nusing namespace nvcuda;\nnamespace cg = cooperative_groups;\n\n__host__ void init_host_matrices(double *a, double *b, double *c)\n{\n for (int i = 0; i < M_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n a[i * K_GLOBAL + j] = (double)(rand() % 3);\n }\n }\n\n for (int i = 0; i < N_GLOBAL; i++) {\n for (int j = 0; j < K_GLOBAL; j++) {\n b[i * K_GLOBAL + j] = (double)(rand() % 3);\n }\n }\n\n for (int t = 0; t < M_GLOBAL * N_GLOBAL; t++) {\n c[t] = (double)(rand() % 3);\n }\n}\n\n__global__ void compute_dgemm(const double *A, const double *B, const double *C, double *D, double alpha, double beta)\n{\n#if __CUDA_ARCH__ >= 800\n extern __shared__ double shmem[][CHUNK_K * K + SKEW_DOUBLE];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n const size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n double *shmem_warp_tile_ptr = (double *)&shmem[0][0]\n + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS\n + (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.\n double *shmem_warp_stream_ptr = (double *)&shmem[0][0] + warpId * SHMEM_STRIDE * N;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may result\n // in a loss of precision). Zero still needs to be specially handled though.\n beta /= alpha;\n\n // Each CTA slides along the 64 x 64 tiles from the top left corner of the matrix to the\n // right and down, and selects the next tile to compute. Once there's no such tile,\n // all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const double *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId) =\n *((int4 *)(src_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment multiplications\n // along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const double *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Scale the C matrix.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const double *warp_ptr =\n (warpId < (WARPS_PER_BLOCK / 2))\n ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);\n\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.\n size_t shmem_idx = warpId < (WARPS_PER_BLOCK / 2)\n ? (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (N * (warpId % (WARPS_PER_BLOCK / 2)) * 2 + shmem_idx_b_off);\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n const double *lane_ptr = warp_ptr + tile_k * K + (laneId / CHUNK_COPY_LINE_LANES) * K_GLOBAL;\n\n // Shift the second half of the warp to the next row / column in the shared memory.\n shmem_idx += laneId / CHUNK_COPY_LINE_LANES;\n\n#pragma unroll\n for (int i = 0; i < ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP); i++) {\n // Copy 16 bytes at once in each lane.\n *((int4 *)&shmem[shmem_idx][0] + (laneId % CHUNK_COPY_LINE_LANES)) =\n *((int4 *)lane_ptr + (laneId % CHUNK_COPY_LINE_LANES));\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / 2) * M * 2 + (i * M);\n const double *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_DOUBLE);\n\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be reused\n // against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const double *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_DOUBLE);\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n // Uniform, point-wise transformations of ALL fragment elements by ALL threads in the\n // warp are well-defined even though element indices within fragment storage are not defined.\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n double *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global memory.\n double *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n#endif\n}\n\n__global__ void\ncompute_dgemm_async_copy(const double *A, const double *B, const double *C, double *D, double alpha, double beta)\n{\n#if __CUDA_ARCH__ >= 800\n extern __shared__ double shmem[][CHUNK_K * K + SKEW_DOUBLE];\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n constexpr size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n double *shmem_warp_tile_ptr = &shmem[0][0] + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS\n + (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.\n double *shmem_warp_stream_ptr = &shmem[0][0] + warpId * SHMEM_STRIDE * N;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may result\n // in a loss of precision). Zero still needs to be specially handled though.\n beta /= alpha;\n\n cuda::pipeline pipe = cuda::make_pipeline();\n\n const auto shape2 = cuda::aligned_size_t(sizeof(double2));\n constexpr int loadStride = 1; // load 2 double, left-shift by 1.\n\n // Each CTA slides along the 64 x 64 tiles from the top left corner of the matrix to the\n // right and down, and selects the next tile to compute. Once there's no such tile,\n // all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const double *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < N; i++) {\n pipe.producer_acquire();\n cuda::memcpy_async(&shmem_warp_stream_ptr[(SHMEM_STRIDE * i) + (laneId << loadStride)],\n &src_gmem_warp_stream_ptr[(GLOBAL_MEM_STRIDE * i) + (laneId << loadStride)],\n shape2,\n pipe);\n\n pipe.producer_commit();\n }\n // Now wait for all the above issued 8 batches to complete.\n cuda::pipeline_consumer_wait_prior<0>(pipe);\n __syncthreads();\n\n // These fragments will accumulate the result of A and B matrix fragment multiplications\n // along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const double *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n // Scale the C matrix.\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n pipe.consumer_release();\n // sync here so that shared memory can then be used for loading A & B matrices.\n __syncthreads();\n\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const double *warp_ptr =\n (warpId < (WARPS_PER_BLOCK / 2))\n ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);\n\n const int stridePerLaneCopy = (laneId / CHUNK_COPY_LINE_LANES);\n constexpr int chunksPerLane = ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP);\n const int laneLoadElem = (laneId % CHUNK_COPY_LINE_LANES) << loadStride;\n\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.\n // As for DMMA M == N we use M for warp 4-7 + shmem_idx_b_off.\n size_t shmem_idx =\n (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2) + (shmem_idx_b_off * (warpId / (WARPS_PER_BLOCK / 2)));\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n const double *lane_ptr = warp_ptr + tile_k * K + stridePerLaneCopy * K_GLOBAL + laneLoadElem;\n\n // Shift the second half of the warp to the next row / column in the shared memory.\n shmem_idx += stridePerLaneCopy;\n#pragma unroll\n for (int i = 0; i < chunksPerLane; i++) {\n // Copy 16 bytes at once in each lane.\n pipe.producer_acquire();\n\n cuda::memcpy_async(&shmem[shmem_idx][laneLoadElem], lane_ptr, shape2, pipe);\n\n pipe.producer_commit();\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n\n cuda::pipeline_consumer_wait_prior<0>(pipe);\n __syncthreads();\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / 2) * M * 2 + (i * M);\n const double *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_DOUBLE);\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be reused\n // against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const double *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_DOUBLE);\n }\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n pipe.consumer_release();\n __syncthreads();\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n // Uniform, point-wise transformations of ALL fragment elements by ALL threads in the\n // warp are well-defined even though element indices within fragment storage are not defined.\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n double *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n __syncthreads();\n\n // Now that shared memory contains all the D tiles, stream them to global memory.\n double *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n\n __syncthreads();\n }\n#endif\n}\n\n__global__ void\ncompute_dgemm_cg_async_copy(const double *A, const double *B, const double *C, double *D, double alpha, double beta)\n{\n#if __CUDA_ARCH__ >= 800\n extern __shared__ double shmem[][CHUNK_K * K + SKEW_DOUBLE];\n auto cta = cg::this_thread_block();\n auto tile32 = cg::tiled_partition<32>(cta);\n\n constexpr int tileChunkCopySize = WARP_SIZE / CHUNK_COPY_LINES_PER_WARP;\n auto tileChunkCopy = cg::tiled_partition(cta);\n\n // Warp and lane identification.\n const unsigned int warpId = threadIdx.x / WARP_SIZE;\n const unsigned int laneId = threadIdx.x % WARP_SIZE;\n\n // Offset in shared memory from which the B matrix is stored.\n constexpr size_t shmem_idx_b_off = BLOCK_COL_TILES * M;\n\n // This pointer is used to access the C and D matrix tiles this warp computes.\n double *shmem_warp_tile_ptr =\n (double *)&shmem[0][0] + (warpId / 2) * SHMEM_STRIDE * N * 2 + (warpId % 2) * SHMEM_OFFSET;\n\n // This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.\n double *shmem_warp_stream_ptr = (double *)&shmem[0][0] + warpId * SHMEM_STRIDE * N;\n\n // Adjust the beta scaler, as it'll be multiplied by alpha at the end of\n // each tile computation. Technically this is not generally correct (may result\n // in a loss of precision). Zero still needs to be specially handled though.\n beta /= alpha;\n\n // Each CTA slides along the 64 x 64 tiles from the top left corner of the matrix to the\n // right and down, and selects the next tile to compute. Once there's no such tile,\n // all warps in this CTA exit.\n for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {\n const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);\n const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;\n\n // Stop when there are no more D matrix tiles to compute in this CTA.\n if (block_tile_i >= M_TILES) {\n break;\n }\n\n // This warp's pointer to the C matrix data to copy memory from to shared memory.\n const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;\n const double *src_gmem_warp_stream_ptr = &C[gmem_idx];\n\n // Stream multiple C tiles to shared memory.\n#pragma unroll\n for (int i = 0; i < N; i++) {\n auto dst_ptr = &shmem_warp_stream_ptr[(SHMEM_STRIDE * i)];\n auto src_ptr = &src_gmem_warp_stream_ptr[(GLOBAL_MEM_STRIDE * i)];\n cg::memcpy_async(\n tile32, dst_ptr, src_ptr, cuda::aligned_size_t{tile32.size() * sizeof(double2)});\n }\n\n cg::wait(cta);\n\n // These fragments will accumulate the result of A and B matrix fragment multiplications\n // along the K_GLOBAL dimension.\n wmma::fragment c[WARP_COL_TILES][WARP_ROW_TILES];\n\n // Load the C matrix tiles into fragments from shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n const double *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n wmma::load_matrix_sync(c[i][j], tile_ptr, SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n // Scale the C matrix.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++) {\n c[i][j].x[t] *= beta;\n }\n }\n }\n\n // sync here so that shared memory can then be used for loading A & B matrices.\n cg::wait(cta);\n // Select what warp copies what matrix to shared memory.\n // Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.\n const double *warp_ptr =\n (warpId < 4) ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)\n : (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);\n\n const int stridePerLaneCopy = (laneId / CHUNK_COPY_LINE_LANES);\n constexpr int chunksPerLane = ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP);\n // Go through the global K dimension by a fixed step at a time.\n#pragma unroll\n for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {\n // Copy slices of the A and B matrices to shared memory.\n // The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.\n // As for DMMA M == N we use M for warp 4-7 + shmem_idx_b_off.\n size_t shmem_idx =\n (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2) + (shmem_idx_b_off * (warpId / (WARPS_PER_BLOCK / 2)));\n\n // First half of the warp copies the first row / column of the matrix,\n // the second half of the warp copies the next.\n auto lane_ptr = warp_ptr + tile_k * K + stridePerLaneCopy * K_GLOBAL;\n\n // Shift the second half of the warp to the next row / column in the shared memory.\n shmem_idx += stridePerLaneCopy;\n\n#pragma unroll\n for (int i = 0; i < chunksPerLane; i++) {\n // Copy 16 bytes at once in each lane.\n auto dst_ptr = &shmem[shmem_idx][0];\n auto src_ptr = lane_ptr;\n\n cg::memcpy_async(tileChunkCopy,\n dst_ptr,\n src_ptr,\n cuda::aligned_size_t{tileChunkCopySize * sizeof(double2)});\n\n // Advance the global memory pointer and the shared memory index.\n lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;\n shmem_idx += CHUNK_COPY_LINES_PER_WARP;\n }\n cg::wait(cta);\n\n // Compute a grid of C matrix tiles in each warp.\n#pragma unroll\n for (int k_step = 0; k_step < CHUNK_K; k_step++) {\n wmma::fragment a[WARP_COL_TILES];\n wmma::fragment b[WARP_ROW_TILES];\n\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n size_t shmem_idx_a = (warpId / 2) * M * 2 + (i * M);\n const double *tile_ptr = &shmem[shmem_idx_a][k_step * K];\n\n wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_DOUBLE);\n\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n if (i == 0) {\n // Load the B matrix fragment once, because it is going to be reused\n // against the other A matrix fragments.\n size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);\n const double *tile_ptr = &shmem[shmem_idx_b][k_step * K];\n\n wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_DOUBLE);\n }\n\n wmma::mma_sync(c[i][j], a[i], b[j], c[i][j]);\n }\n }\n }\n cg::sync(cta);\n }\n\n // Store the D fragments to shared memory.\n#pragma unroll\n for (int i = 0; i < WARP_COL_TILES; i++) {\n#pragma unroll\n for (int j = 0; j < WARP_ROW_TILES; j++) {\n // Uniform, point-wise transformations of ALL fragment elements by ALL threads in the\n // warp are well-defined even though element indices within fragment storage are not defined.\n#pragma unroll\n for (int t = 0; t < c[i][j].num_elements; t++)\n c[i][j].x[t] *= alpha;\n\n double *tile_ptr = shmem_warp_tile_ptr + i * SHMEM_STRIDE * N + j * N;\n\n wmma::store_matrix_sync(tile_ptr, c[i][j], SHMEM_STRIDE, C_LAYOUT);\n }\n }\n\n cg::sync(cta);\n\n // Now that shared memory contains all the D tiles, stream them to global memory.\n double *dst_gmem_warp_stream_ptr = &D[gmem_idx];\n\n#pragma unroll\n for (int i = 0; i < N; i++) {\n *((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =\n *((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);\n }\n cg::sync(cta);\n }\n#endif\n}\n\n// Performs an MxNxK DGEMM (C=alpha*A*B + beta*C) assuming:\n// 1) Matrices are packed in memory.\n// 2) M, N and K are multiples of 8, 8 and 4 respectively.\n// 3) A is row major, B is column major matrix.\n// Note: This is a less performant version of the compute_dgemm kernel. It is designed for\n// demonstration purposes only to show the CUDA WMMA API use without relying on\n// availability of the shared memory.\n__global__ void\nsimple_wmma_gemm(double *a, double *b, double *c, double *d, int m_ld, int n_ld, int k_ld, double alpha, double beta)\n{\n#if __CUDA_ARCH__ >= 800\n // Leading dimensions. Packed with no transpositions.\n int lda = k_ld;\n int ldb = k_ld;\n int ldc = n_ld;\n\n // Tile using a 2D grid\n int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;\n int warpN = (blockIdx.y * blockDim.y + threadIdx.y);\n\n // Declare the fragments\n wmma::fragment a_frag;\n wmma::fragment b_frag;\n wmma::fragment acc_frag;\n wmma::fragment c_frag;\n\n wmma::fill_fragment(acc_frag, 0.0f);\n\n // Loop over k\n for (int i = 0; i < k_ld; i += K) {\n int aCol = i;\n int aRow = warpM * M;\n\n int bCol = warpN * N;\n int bRow = i;\n\n // Bounds checking\n if (aRow < m_ld && aCol < k_ld && bRow < k_ld && bCol < n_ld) {\n // Load the inputs\n wmma::load_matrix_sync(a_frag, a + aCol + aRow * lda, lda);\n wmma::load_matrix_sync(b_frag, b + bRow + bCol * ldb, ldb);\n\n // Perform the matrix multiplication\n wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);\n }\n }\n\n // Load in the current value of c, scale it by beta, and add this our result scaled by alpha\n int cCol = warpN * N;\n int cRow = warpM * M;\n\n if (cRow < m_ld && cCol < n_ld) {\n wmma::load_matrix_sync(c_frag, c + cCol + cRow * ldc, ldc, wmma::mem_row_major);\n\n for (int i = 0; i < c_frag.num_elements; i++) {\n c_frag.x[i] = alpha * acc_frag.x[i] + beta * c_frag.x[i];\n }\n\n // Store the output\n wmma::store_matrix_sync(d + cCol + cRow * ldc, c_frag, ldc, wmma::mem_row_major);\n }\n#endif\n}\n\n__host__ void matMultiplyOnHost(double *A,\n double *B,\n double *C,\n float alpha,\n float beta,\n int numARows,\n int numAColumns,\n int numBRows,\n int numBColumns,\n int numCRows,\n int numCColumns)\n{\n for (int i = 0; i < numCRows; i++) {\n for (int j = 0; j < numCColumns; j++) {\n double temp = 0.0;\n\n for (int k = 0; k < numAColumns; k++) {\n // B matrix is column major. A matrix is row major.\n temp += A[i * numAColumns + k] * B[j * numBRows + k];\n }\n\n C[i * numCColumns + j] = temp * alpha + beta * C[i * numCColumns + j];\n }\n }\n}\n\nint main(int argc, char **argv)\n{\n printf(\"Initializing...\\n\");\n\n int dev = findCudaDevice(argc, (const char **)argv);\n\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, dev));\n\n // Double precision Tensor cores require a GPU of Ampere (SM8X) architecture or higher.\n if (deviceProp.major < 8) {\n printf(\"dmmaTensorCoreGemm requires SM 8.0 or higher. Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n printf(\"M: %d (%d x %d)\\n\", M_GLOBAL, M, M_TILES);\n printf(\"N: %d (%d x %d)\\n\", N_GLOBAL, N, N_TILES);\n printf(\"K: %d (%d x %d)\\n\", K_GLOBAL, K, K_TILES);\n\n double *A_h = NULL;\n double *B_h = NULL;\n double *C_h = NULL;\n#if CPU_DEBUG\n double *result_hD = NULL;\n double *result_host = NULL;\n#endif\n\n A_h = (double *)malloc(sizeof(double) * M_GLOBAL * K_GLOBAL);\n B_h = (double *)malloc(sizeof(double) * K_GLOBAL * N_GLOBAL);\n C_h = (double *)malloc(sizeof(double) * M_GLOBAL * N_GLOBAL);\n#if CPU_DEBUG\n result_hD = (double *)malloc(sizeof(double) * M_GLOBAL * N_GLOBAL);\n result_host = (double *)malloc(sizeof(double) * M_GLOBAL * N_GLOBAL);\n#endif\n\n double *A = NULL;\n double *B = NULL;\n double *C = NULL;\n double *D = NULL;\n\n checkCudaErrors(hipMalloc((void **)&A, sizeof(double) * M_GLOBAL * K_GLOBAL));\n checkCudaErrors(hipMalloc((void **)&B, sizeof(double) * N_GLOBAL * K_GLOBAL));\n checkCudaErrors(hipMalloc((void **)&C, sizeof(double) * M_GLOBAL * N_GLOBAL));\n checkCudaErrors(hipMalloc((void **)&D, sizeof(double) * M_GLOBAL * N_GLOBAL));\n\n assert(((unsigned long long)A) % 128 == 0);\n assert(((unsigned long long)B) % 128 == 0);\n assert(((unsigned long long)C) % 128 == 0);\n assert(((unsigned long long)D) % 128 == 0);\n\n init_host_matrices(A_h, B_h, C_h);\n\n printf(\"Preparing data for GPU...\\n\");\n\n checkCudaErrors(hipMemcpy(A, A_h, sizeof(double) * M_GLOBAL * K_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(B, B_h, sizeof(double) * N_GLOBAL * K_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(C, C_h, sizeof(double) * M_GLOBAL * N_GLOBAL, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemset(D, 0, sizeof(double) * M_GLOBAL * N_GLOBAL));\n\n enum {\n // Compute the right amount of shared memory to request.\n // We need shared memory to hold per-CTA C and D matrix tiles, and to cache per-CTA chunks\n // of the A and B matrices. Therefore, the right amount to request is the maximum of those\n // two numbers.\n SHMEM_SZ = MAX(sizeof(double) * (BLOCK_COL_TILES * M) * (CHUNK_K * K + SKEW_DOUBLE) * 2,\n M * (BLOCK_ROW_WARPS * WARP_ROW_TILES) * N * (BLOCK_COL_WARPS * WARP_COL_TILES) * sizeof(double))\n };\n\n printf(\"Required shared memory size: %lu Kb\\n\", SHMEM_SZ / 1024UL);\n\n const double alpha = 1.1f;\n const double beta = 1.2f;\n\n hipEvent_t start, stop;\n\n checkCudaErrors(hipEventCreate(&start));\n checkCudaErrors(hipEventCreate(&stop));\n checkCudaErrors(hipEventRecord(start));\n\n kernels selected_kernel = dmma_shmem_gemm_async_copy;\n\n // kernel to run - default (dmma_shmem_gemm_async_copy == 0)\n if (checkCmdLineFlag(argc, (const char **)argv, \"kernel\")) {\n int kernel_number = getCmdLineArgumentInt(argc, (const char **)argv, \"kernel\");\n if (kernel_number < 4) {\n selected_kernel = (kernels)kernel_number;\n }\n else {\n printf(\"Error: kernel number should be between 0 to 3, you have entered %d\\n\", kernel_number);\n exit(EXIT_FAILURE);\n }\n }\n\n // If enough shared memory available on the GPU use high performant kernel\n if ((deviceProp.sharedMemPerMultiprocessor >= SHMEM_SZ) && (selected_kernel != simple_dmma_gemm)) {\n printf(\"Computing using high performance kernel = %d - %s\\n\", selected_kernel, kernelNames[selected_kernel]);\n\n switch (selected_kernel) {\n case dmma_shmem_gemm_async_copy:\n default:\n checkCudaErrors(\n hipFuncSetAttribute(reinterpret_cast(compute_dgemm_async_copy), hipFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors(\n (compute_dgemm_async_copy<<>>(\n A, B, C, D, alpha, beta)));\n break;\n case dmma_shmem_gemm_cg_async_copy:\n checkCudaErrors(hipFuncSetAttribute(reinterpret_cast(compute_dgemm_cg_async_copy), hipFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors(\n (compute_dgemm_cg_async_copy<<>>(\n A, B, C, D, alpha, beta)));\n break;\n case dmma_shmem_gemm:\n checkCudaErrors(hipFuncSetAttribute(reinterpret_cast(compute_dgemm), hipFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));\n checkKernelErrors((compute_dgemm<<>>(\n A, B, C, D, alpha, beta)));\n break;\n }\n\n#if CPU_DEBUG\n checkCudaErrors(hipMemcpy(result_hD, D, sizeof(double) * M_GLOBAL * N_GLOBAL, hipMemcpyDeviceToHost));\n#endif\n }\n else {\n dim3 gridDim;\n dim3 blockDim;\n\n // blockDim.x must be a multple of warpSize\n // 128x4 means we have 16 warps and a block computes a 64x64 output tile\n blockDim.x = 128;\n blockDim.y = 4;\n\n gridDim.x = (M_GLOBAL + (M * blockDim.x / 32 - 1)) / (M * blockDim.x / 32);\n gridDim.y = (N_GLOBAL + N * blockDim.y - 1) / (N * blockDim.y);\n\n printf(\"Computing... using simple_wmma_gemm kernel\\n\");\n simple_wmma_gemm<<>>(A, B, C, D, M_GLOBAL, N_GLOBAL, K_GLOBAL, alpha, beta);\n#if CPU_DEBUG\n checkCudaErrors(hipMemcpy(result_hD, D, sizeof(double) * M_GLOBAL * N_GLOBAL, hipMemcpyDeviceToHost));\n#endif\n }\n\n checkCudaErrors(hipEventRecord(stop));\n checkCudaErrors(hipEventSynchronize(stop));\n\n#if CPU_DEBUG\n printf(\"Verifying correctness of the computations...\\n\");\n\n memcpy(result_host, C_h, sizeof(double) * M_GLOBAL * N_GLOBAL);\n\n matMultiplyOnHost(A_h, B_h, result_host, alpha, beta, M_GLOBAL, K_GLOBAL, K_GLOBAL, N_GLOBAL, M_GLOBAL, N_GLOBAL);\n\n size_t number_of_matches = 0;\n for (int i = 0; i < N_GLOBAL * M_GLOBAL; i++) {\n if (fabs(result_hD[i] - result_host[i]) > 0.1f) {\n printf(\"mismatch i=%d result_hD=%f result_host=%f\\n\", i, result_hD[i], result_host[i]);\n break;\n }\n else {\n number_of_matches++;\n }\n }\n printf(\"number_of_matches = %zu out of = %d \\n\", number_of_matches, N_GLOBAL * M_GLOBAL);\n free(result_hD);\n free(result_host);\n#endif\n\n float milliseconds = 0;\n\n checkCudaErrors(hipEventElapsedTime(&milliseconds, start, stop));\n\n printf(\"Time: %f ms\\n\", milliseconds);\n printf(\"FP64 TFLOPS: %.2f\\n\", (((double)M_GLOBAL * N_GLOBAL * K_GLOBAL * 2) / (milliseconds / 1000.)) / 1e12);\n\n free(A_h);\n free(B_h);\n free(C_h);\n checkCudaErrors(hipFree((void *)A));\n checkCudaErrors(hipFree((void *)B));\n checkCudaErrors(hipFree((void *)C));\n checkCudaErrors(hipFree((void *)D));\n\n return 0;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/dmmaTensorCoreGemm/dmmaTensorCoreGemm.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n\n#define MAX_DEPTH 16\n#define INSERTION_SORT 32\n\n////////////////////////////////////////////////////////////////////////////////\n// Selection sort used when depth gets too big or the number of elements drops\n// below a threshold.\n////////////////////////////////////////////////////////////////////////////////\n__device__ void selection_sort(unsigned int *data, int left, int right)\n{\n for (int i = left; i <= right; ++i) {\n unsigned min_val = data[i];\n int min_idx = i;\n\n // Find the smallest value in the range [left, right].\n for (int j = i + 1; j <= right; ++j) {\n unsigned val_j = data[j];\n\n if (val_j < min_val) {\n min_idx = j;\n min_val = val_j;\n }\n }\n\n // Swap the values.\n if (i != min_idx) {\n data[min_idx] = data[i];\n data[i] = min_val;\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Very basic quicksort algorithm, recursively launching the next level.\n////////////////////////////////////////////////////////////////////////////////\n__global__ void cdp_simple_quicksort(unsigned int *data, int left, int right, int depth)\n{\n // If we're too deep or there are few elements left, we use an insertion\n // sort...\n if (depth >= MAX_DEPTH || right - left <= INSERTION_SORT) {\n selection_sort(data, left, right);\n return;\n }\n\n unsigned int *lptr = data + left;\n unsigned int *rptr = data + right;\n unsigned int pivot = data[(left + right) / 2];\n\n // Do the partitioning.\n while (lptr <= rptr) {\n // Find the next left- and right-hand values to swap\n unsigned int lval = *lptr;\n unsigned int rval = *rptr;\n\n // Move the left pointer as long as the pointed element is smaller than the\n // pivot.\n while (lval < pivot) {\n lptr++;\n lval = *lptr;\n }\n\n // Move the right pointer as long as the pointed element is larger than the\n // pivot.\n while (rval > pivot) {\n rptr--;\n rval = *rptr;\n }\n\n // If the swap points are valid, do the swap!\n if (lptr <= rptr) {\n *lptr++ = rval;\n *rptr-- = lval;\n }\n }\n\n // Now the recursive part\n int nright = rptr - data;\n int nleft = lptr - data;\n\n // Launch a new block to sort the left part.\n if (left < (rptr - data)) {\n cudaStream_t s;\n cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking);\n cdp_simple_quicksort<<<1, 1, 0, s>>>(data, left, nright, depth + 1);\n cudaStreamDestroy(s);\n }\n\n // Launch a new block to sort the right part.\n if ((lptr - data) < right) {\n cudaStream_t s1;\n cudaStreamCreateWithFlags(&s1, cudaStreamNonBlocking);\n cdp_simple_quicksort<<<1, 1, 0, s1>>>(data, nleft, right, depth + 1);\n cudaStreamDestroy(s1);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Call the quicksort kernel from the host.\n////////////////////////////////////////////////////////////////////////////////\nvoid run_qsort(unsigned int *data, unsigned int nitems)\n{\n // Prepare CDP for the max depth 'MAX_DEPTH'.\n\n // Launch on device\n int left = 0;\n int right = nitems - 1;\n std::cout << \"Launching kernel on the GPU\" << std::endl;\n cdp_simple_quicksort<<<1, 1>>>(data, left, right, 0);\n checkCudaErrors(cudaDeviceSynchronize());\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Initialize data on the host.\n////////////////////////////////////////////////////////////////////////////////\nvoid initialize_data(unsigned int *dst, unsigned int nitems)\n{\n // Fixed seed for illustration\n srand(2047);\n\n // Fill dst with random values\n for (unsigned i = 0; i < nitems; i++)\n dst[i] = rand() % nitems;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Verify the results.\n////////////////////////////////////////////////////////////////////////////////\nvoid check_results(int n, unsigned int *results_d)\n{\n unsigned int *results_h = new unsigned[n];\n checkCudaErrors(cudaMemcpy(results_h, results_d, n * sizeof(unsigned), cudaMemcpyDeviceToHost));\n\n for (int i = 1; i < n; ++i)\n if (results_h[i - 1] > results_h[i]) {\n std::cout << \"Invalid item[\" << i - 1 << \"]: \" << results_h[i - 1] << \" greater than \" << results_h[i]\n << std::endl;\n exit(EXIT_FAILURE);\n }\n\n std::cout << \"OK\" << std::endl;\n delete[] results_h;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Main entry point.\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n int num_items = 128;\n bool verbose = false;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\") || checkCmdLineFlag(argc, (const char **)argv, \"h\")) {\n std::cerr << \"Usage: \" << argv[0]\n << \" num_items=\\twhere num_items is the number of \"\n \"items to sort\"\n << std::endl;\n exit(EXIT_SUCCESS);\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"v\")) {\n verbose = true;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"num_items\")) {\n num_items = getCmdLineArgumentInt(argc, (const char **)argv, \"num_items\");\n\n if (num_items < 1) {\n std::cerr << \"ERROR: num_items has to be greater than 1\" << std::endl;\n exit(EXIT_FAILURE);\n }\n }\n\n // Find/set device and get device properties\n int device = -1;\n cudaDeviceProp deviceProp;\n device = findCudaDevice(argc, (const char **)argv);\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, device));\n\n if (!(deviceProp.major > 3 || (deviceProp.major == 3 && deviceProp.minor >= 5))) {\n printf(\"GPU %d - %s does not support CUDA Dynamic Parallelism\\n Exiting.\", device, deviceProp.name);\n exit(EXIT_WAIVED);\n }\n\n // Create input data\n unsigned int *h_data = 0;\n unsigned int *d_data = 0;\n\n // Allocate CPU memory and initialize data.\n std::cout << \"Initializing data:\" << std::endl;\n h_data = (unsigned int *)malloc(num_items * sizeof(unsigned int));\n initialize_data(h_data, num_items);\n\n if (verbose) {\n for (int i = 0; i < num_items; i++)\n std::cout << \"Data [\" << i << \"]: \" << h_data[i] << std::endl;\n }\n\n // Allocate GPU memory.\n checkCudaErrors(cudaMalloc((void **)&d_data, num_items * sizeof(unsigned int)));\n checkCudaErrors(cudaMemcpy(d_data, h_data, num_items * sizeof(unsigned int), cudaMemcpyHostToDevice));\n\n // Execute\n std::cout << \"Running quicksort on \" << num_items << \" elements\" << std::endl;\n run_qsort(d_data, num_items);\n\n // Check result\n std::cout << \"Validating results: \";\n check_results(num_items, d_data);\n\n free(h_data);\n checkCudaErrors(cudaFree(d_data));\n\n exit(EXIT_SUCCESS);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n#define MAX_DEPTH 16\n#define INSERTION_SORT 32\n\n////////////////////////////////////////////////////////////////////////////////\n// Selection sort used when depth gets too big or the number of elements drops\n// below a threshold.\n////////////////////////////////////////////////////////////////////////////////\n__device__ void selection_sort(unsigned int *data, int left, int right)\n{\n for (int i = left; i <= right; ++i) {\n unsigned min_val = data[i];\n int min_idx = i;\n\n // Find the smallest value in the range [left, right].\n for (int j = i + 1; j <= right; ++j) {\n unsigned val_j = data[j];\n\n if (val_j < min_val) {\n min_idx = j;\n min_val = val_j;\n }\n }\n\n // Swap the values.\n if (i != min_idx) {\n data[min_idx] = data[i];\n data[i] = min_val;\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Very basic quicksort algorithm, recursively launching the next level.\n////////////////////////////////////////////////////////////////////////////////\n__global__ void cdp_simple_quicksort(unsigned int *data, int left, int right, int depth)\n{\n // If we're too deep or there are few elements left, we use an insertion\n // sort...\n if (depth >= MAX_DEPTH || right - left <= INSERTION_SORT) {\n selection_sort(data, left, right);\n return;\n }\n\n unsigned int *lptr = data + left;\n unsigned int *rptr = data + right;\n unsigned int pivot = data[(left + right) / 2];\n\n // Do the partitioning.\n while (lptr <= rptr) {\n // Find the next left- and right-hand values to swap\n unsigned int lval = *lptr;\n unsigned int rval = *rptr;\n\n // Move the left pointer as long as the pointed element is smaller than the\n // pivot.\n while (lval < pivot) {\n lptr++;\n lval = *lptr;\n }\n\n // Move the right pointer as long as the pointed element is larger than the\n // pivot.\n while (rval > pivot) {\n rptr--;\n rval = *rptr;\n }\n\n // If the swap points are valid, do the swap!\n if (lptr <= rptr) {\n *lptr++ = rval;\n *rptr-- = lval;\n }\n }\n\n // Now the recursive part\n int nright = rptr - data;\n int nleft = lptr - data;\n\n // Launch a new block to sort the left part.\n if (left < (rptr - data)) {\n hipStream_t s;\n hipStreamCreateWithFlags(&s, hipStreamNonBlocking);\n cdp_simple_quicksort<<<1, 1, 0, s>>>(data, left, nright, depth + 1);\n hipStreamDestroy(s);\n }\n\n // Launch a new block to sort the right part.\n if ((lptr - data) < right) {\n hipStream_t s1;\n hipStreamCreateWithFlags(&s1, hipStreamNonBlocking);\n cdp_simple_quicksort<<<1, 1, 0, s1>>>(data, nleft, right, depth + 1);\n hipStreamDestroy(s1);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Call the quicksort kernel from the host.\n////////////////////////////////////////////////////////////////////////////////\nvoid run_qsort(unsigned int *data, unsigned int nitems)\n{\n // Prepare CDP for the max depth 'MAX_DEPTH'.\n\n // Launch on device\n int left = 0;\n int right = nitems - 1;\n std::cout << \"Launching kernel on the GPU\" << std::endl;\n cdp_simple_quicksort<<<1, 1>>>(data, left, right, 0);\n checkCudaErrors(hipDeviceSynchronize());\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Initialize data on the host.\n////////////////////////////////////////////////////////////////////////////////\nvoid initialize_data(unsigned int *dst, unsigned int nitems)\n{\n // Fixed seed for illustration\n srand(2047);\n\n // Fill dst with random values\n for (unsigned i = 0; i < nitems; i++)\n dst[i] = rand() % nitems;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Verify the results.\n////////////////////////////////////////////////////////////////////////////////\nvoid check_results(int n, unsigned int *results_d)\n{\n unsigned int *results_h = new unsigned[n];\n checkCudaErrors(hipMemcpy(results_h, results_d, n * sizeof(unsigned), hipMemcpyDeviceToHost));\n\n for (int i = 1; i < n; ++i)\n if (results_h[i - 1] > results_h[i]) {\n std::cout << \"Invalid item[\" << i - 1 << \"]: \" << results_h[i - 1] << \" greater than \" << results_h[i]\n << std::endl;\n exit(EXIT_FAILURE);\n }\n\n std::cout << \"OK\" << std::endl;\n delete[] results_h;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Main entry point.\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n int num_items = 128;\n bool verbose = false;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\") || checkCmdLineFlag(argc, (const char **)argv, \"h\")) {\n std::cerr << \"Usage: \" << argv[0]\n << \" num_items=\\twhere num_items is the number of \"\n \"items to sort\"\n << std::endl;\n exit(EXIT_SUCCESS);\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"v\")) {\n verbose = true;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"num_items\")) {\n num_items = getCmdLineArgumentInt(argc, (const char **)argv, \"num_items\");\n\n if (num_items < 1) {\n std::cerr << \"ERROR: num_items has to be greater than 1\" << std::endl;\n exit(EXIT_FAILURE);\n }\n }\n\n // Find/set device and get device properties\n int device = -1;\n hipDeviceProp_t deviceProp;\n device = findCudaDevice(argc, (const char **)argv);\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, device));\n\n if (!(deviceProp.major > 3 || (deviceProp.major == 3 && deviceProp.minor >= 5))) {\n printf(\"GPU %d - %s does not support CUDA Dynamic Parallelism\\n Exiting.\", device, deviceProp.name);\n exit(EXIT_WAIVED);\n }\n\n // Create input data\n unsigned int *h_data = 0;\n unsigned int *d_data = 0;\n\n // Allocate CPU memory and initialize data.\n std::cout << \"Initializing data:\" << std::endl;\n h_data = (unsigned int *)malloc(num_items * sizeof(unsigned int));\n initialize_data(h_data, num_items);\n\n if (verbose) {\n for (int i = 0; i < num_items; i++)\n std::cout << \"Data [\" << i << \"]: \" << h_data[i] << std::endl;\n }\n\n // Allocate GPU memory.\n checkCudaErrors(hipMalloc((void **)&d_data, num_items * sizeof(unsigned int)));\n checkCudaErrors(hipMemcpy(d_data, h_data, num_items * sizeof(unsigned int), hipMemcpyHostToDevice));\n\n // Execute\n std::cout << \"Running quicksort on \" << num_items << \" elements\" << std::endl;\n run_qsort(d_data, num_items);\n\n // Check result\n std::cout << \"Validating results: \";\n check_results(num_items, d_data);\n\n free(h_data);\n checkCudaErrors(hipFree(d_data));\n\n exit(EXIT_SUCCESS);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/cdpSimpleQuicksort/cdpSimpleQuicksort.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// This is a basic, recursive bitonic sort taken from\n// http://www.iti.fh-flensburg.de/lang/algorithmen/sortieren/bitonic/oddn.htm\n//\n// The parallel code is based on:\n// http://www.tools-of-computing.com/tc/CS/Sorts/bitonic_sort.htm\n//\n// The multithread code is from me.\n\n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \"cdpQuicksort.h\"\n\n// Inline PTX call to return index of highest non-zero bit in a word\nstatic __device__ __forceinline__ unsigned int __btflo(unsigned int word)\n{\n unsigned int ret;\n asm volatile(\"bfind.u32 %0, %1;\" : \"=r\"(ret) : \"r\"(word));\n return ret;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// qcompare\n//\n// Comparison function. Note difference from libc standard in\n// that we take by reference, not by pointer. I can't figure\n// out how to get a template-to-pointer specialisation working.\n// Perhaps it requires a class?\n//\n////////////////////////////////////////////////////////////////////////////////\n__device__ __forceinline__ int qcompare(unsigned &val1, unsigned &val2)\n{\n return (val1 > val2) ? 1 : (val1 == val2) ? 0 : -1;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// Basic any-N bitonic sort. We sort \"len\" elements of \"indata\", starting\n// from the \"offset\" elements into the input data array. Note that \"outdata\"\n// can safely be the same as \"indata\" for an in-place sort (we stage through\n// shared memory).\n//\n// We handle non-power-of-2 sizes by padding out to the next largest power of\n// 2.\n// This is the fully-generic version, for sorting arbitrary data which does not\n// have a clear \"maximum\" value. We track \"invalid\" entries in a separate array\n// to make sure that they always sorts as \"max value\" elements. A template\n// parameter \"OOR\" allows specialisation to optimise away the invalid tracking.\n//\n// We can do a more specialised version for int/longlong/flat/double, in which\n// we pad out the array with max-value-of-type elements. That's another\n// function.\n//\n// The last step copies from indata -> outdata... the rest is done in-place.\n// We use shared memory as temporary storage, which puts an upper limit on\n// how much data we can sort per block.\n//\n////////////////////////////////////////////////////////////////////////////////\nstatic __device__ __forceinline__ void\nbitonicsort_kernel(unsigned *indata, unsigned *outdata, unsigned int offset, unsigned int len, cg::thread_block cta)\n{\n __shared__ unsigned sortbuf[1024]; // Max of 1024 elements - TODO: make this dynamic\n\n // First copy data into shared memory.\n unsigned int inside = (threadIdx.x < len);\n sortbuf[threadIdx.x] = inside ? indata[threadIdx.x + offset] : 0xffffffffu;\n cg::sync(cta);\n\n // Now the sort loops\n // Here, \"k\" is the sort level (remember bitonic does a multi-level butterfly\n // style sort)\n // and \"j\" is the partner element in the butterfly.\n // Two threads each work on one butterfly, because the read/write needs to\n // happen\n // simultaneously\n for (unsigned int k = 2; k <= blockDim.x; k *= 2) // Butterfly stride increments in powers of 2\n {\n for (unsigned int j = k >> 1; j > 0; j >>= 1) // Strides also in powers of to, up to swap_idx, then ascending means mine threadIdx.x);\n bool swap = false;\n\n if ((threadIdx.x & k) == ascend) {\n if (my_elem > swap_elem)\n swap = true;\n }\n\n if ((threadIdx.x & k) == descend) {\n if (my_elem < swap_elem)\n swap = true;\n }\n\n // If we had to swap, then write my data to the other element's position.\n // Don't forget to track out-of-range status too!\n if (swap) {\n sortbuf[swap_idx] = my_elem;\n }\n\n cg::sync(cta);\n }\n }\n\n // Copy the sorted data from shared memory back to the output buffer\n if (threadIdx.x < len)\n outdata[threadIdx.x + offset] = sortbuf[threadIdx.x];\n}\n\n//////////////////////////////////////////////////////////////////////////////////\n// This is an emergency-CTA sort, which sorts an arbitrary sized chunk\n// using a single block. Useful for if qsort runs out of nesting depth.\n//\n// Note that bitonic sort needs enough storage to pad up to the nearest\n// power of 2. This means that the double-buffer is always large enough\n// (when combined with the main buffer), but we do not get enough space\n// to keep OOR information.\n//\n// This in turn means that this sort does not work with a generic data\n// type. It must be a directly-comparable (i.e. with max value) type.\n//\n////////////////////////////////////////////////////////////////////////////////\nstatic __device__ __forceinline__ void big_bitonicsort_kernel(unsigned *indata,\n unsigned *outdata,\n unsigned *backbuf,\n unsigned int offset,\n unsigned int len,\n cg::thread_block cta)\n{\n unsigned int len2 = 1 << (__btflo(len - 1U) + 1); // Round up len to nearest power-of-2\n\n if (threadIdx.x >= len2)\n return; // Early out for case where more threads launched than there is\n // data\n\n // First, set up our unused values to be the max data type.\n for (unsigned int i = len; i < len2; i += blockDim.x) {\n unsigned int index = i + threadIdx.x;\n\n if (index < len2) {\n // Must split our index between two buffers\n if (index < len)\n indata[index + offset] = 0xffffffffu;\n else\n backbuf[index + offset - len] = 0xffffffffu;\n }\n }\n\n cg::sync(cta);\n\n // Now the sort loops\n // Here, \"k\" is the sort level (remember bitonic does a multi-level butterfly\n // style sort)\n // and \"j\" is the partner element in the butterfly.\n // Two threads each work on one butterfly, because the read/write needs to\n // happen\n // simultaneously\n for (unsigned int k = 2; k <= len2; k *= 2) // Butterfly stride increments in powers of 2\n {\n for (unsigned int j = k >> 1; j > 0; j >>= 1) // Strides also in powers of to, up to index) {\n unsigned my_elem, swap_elem;\n\n if (index < len)\n my_elem = indata[index + offset];\n else\n my_elem = backbuf[index + offset - len];\n\n if (swap_idx < len)\n swap_elem = indata[swap_idx + offset];\n else\n swap_elem = backbuf[swap_idx + offset - len];\n\n // The k'th bit of my index (and hence my sort item ID)\n // determines if we sort ascending or descending.\n // Also, if either my_elem or swap_elem is out of range, then it\n // ALWAYS acts like it's the largest number.\n bool swap = false;\n\n if ((index & k) == 0) {\n if (my_elem > swap_elem)\n swap = true;\n }\n\n if ((index & k) == k) {\n if (my_elem < swap_elem)\n swap = true;\n }\n\n // If we had to swap, then write my data to the other element's\n // position.\n if (swap) {\n if (swap_idx < len)\n indata[swap_idx + offset] = my_elem;\n else\n backbuf[swap_idx + offset - len] = my_elem;\n\n if (index < len)\n indata[index + offset] = swap_elem;\n else\n backbuf[index + offset - len] = swap_elem;\n }\n }\n }\n\n cg::sync(cta); // Only need to sync for each \"j\" pass\n }\n }\n\n // Copy the sorted data from the input to the output buffer, because we sort\n // in-place\n if (outdata != indata) {\n for (unsigned int i = 0; i < len; i += blockDim.x) {\n unsigned int index = i + threadIdx.x;\n\n if (index < len)\n outdata[index + offset] = indata[index + offset];\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// KERNELS\n////////////////////////////////////////////////////////////////////////////////\n\n__global__ void bitonicsort(unsigned *indata, unsigned *outdata, unsigned int offset, unsigned int len)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n bitonicsort_kernel(indata, outdata, offset, len, cta);\n}\n\n__global__ void\nbig_bitonicsort(unsigned *indata, unsigned *outdata, unsigned *backbuf, unsigned int offset, unsigned int len)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n big_bitonicsort_kernel(indata, outdata, backbuf, offset, len, cta);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// This is a basic, recursive bitonic sort taken from\n// http://www.iti.fh-flensburg.de/lang/algorithmen/sortieren/bitonic/oddn.htm\n//\n// The parallel code is based on:\n// http://www.tools-of-computing.com/tc/CS/Sorts/bitonic_sort.htm\n//\n// The multithread code is from me.\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \"cdpQuicksort.h\"\n\n// Inline PTX call to return index of highest non-zero bit in a word\nstatic __device__ __forceinline__ unsigned int __btflo(unsigned int word)\n{\n unsigned int ret;\n asm volatile(\"bfind.u32 %0, %1;\" : \"=r\"(ret) : \"r\"(word));\n return ret;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// qcompare\n//\n// Comparison function. Note difference from libc standard in\n// that we take by reference, not by pointer. I can't figure\n// out how to get a template-to-pointer specialisation working.\n// Perhaps it requires a class?\n//\n////////////////////////////////////////////////////////////////////////////////\n__device__ __forceinline__ int qcompare(unsigned &val1, unsigned &val2)\n{\n return (val1 > val2) ? 1 : (val1 == val2) ? 0 : -1;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// Basic any-N bitonic sort. We sort \"len\" elements of \"indata\", starting\n// from the \"offset\" elements into the input data array. Note that \"outdata\"\n// can safely be the same as \"indata\" for an in-place sort (we stage through\n// shared memory).\n//\n// We handle non-power-of-2 sizes by padding out to the next largest power of\n// 2.\n// This is the fully-generic version, for sorting arbitrary data which does not\n// have a clear \"maximum\" value. We track \"invalid\" entries in a separate array\n// to make sure that they always sorts as \"max value\" elements. A template\n// parameter \"OOR\" allows specialisation to optimise away the invalid tracking.\n//\n// We can do a more specialised version for int/longlong/flat/double, in which\n// we pad out the array with max-value-of-type elements. That's another\n// function.\n//\n// The last step copies from indata -> outdata... the rest is done in-place.\n// We use shared memory as temporary storage, which puts an upper limit on\n// how much data we can sort per block.\n//\n////////////////////////////////////////////////////////////////////////////////\nstatic __device__ __forceinline__ void\nbitonicsort_kernel(unsigned *indata, unsigned *outdata, unsigned int offset, unsigned int len, cg::thread_block cta)\n{\n __shared__ unsigned sortbuf[1024]; // Max of 1024 elements - TODO: make this dynamic\n\n // First copy data into shared memory.\n unsigned int inside = (threadIdx.x < len);\n sortbuf[threadIdx.x] = inside ? indata[threadIdx.x + offset] : 0xffffffffu;\n cg::sync(cta);\n\n // Now the sort loops\n // Here, \"k\" is the sort level (remember bitonic does a multi-level butterfly\n // style sort)\n // and \"j\" is the partner element in the butterfly.\n // Two threads each work on one butterfly, because the read/write needs to\n // happen\n // simultaneously\n for (unsigned int k = 2; k <= blockDim.x; k *= 2) // Butterfly stride increments in powers of 2\n {\n for (unsigned int j = k >> 1; j > 0; j >>= 1) // Strides also in powers of to, up to swap_idx, then ascending means mine threadIdx.x);\n bool swap = false;\n\n if ((threadIdx.x & k) == ascend) {\n if (my_elem > swap_elem)\n swap = true;\n }\n\n if ((threadIdx.x & k) == descend) {\n if (my_elem < swap_elem)\n swap = true;\n }\n\n // If we had to swap, then write my data to the other element's position.\n // Don't forget to track out-of-range status too!\n if (swap) {\n sortbuf[swap_idx] = my_elem;\n }\n\n cg::sync(cta);\n }\n }\n\n // Copy the sorted data from shared memory back to the output buffer\n if (threadIdx.x < len)\n outdata[threadIdx.x + offset] = sortbuf[threadIdx.x];\n}\n\n//////////////////////////////////////////////////////////////////////////////////\n// This is an emergency-CTA sort, which sorts an arbitrary sized chunk\n// using a single block. Useful for if qsort runs out of nesting depth.\n//\n// Note that bitonic sort needs enough storage to pad up to the nearest\n// power of 2. This means that the double-buffer is always large enough\n// (when combined with the main buffer), but we do not get enough space\n// to keep OOR information.\n//\n// This in turn means that this sort does not work with a generic data\n// type. It must be a directly-comparable (i.e. with max value) type.\n//\n////////////////////////////////////////////////////////////////////////////////\nstatic __device__ __forceinline__ void big_bitonicsort_kernel(unsigned *indata,\n unsigned *outdata,\n unsigned *backbuf,\n unsigned int offset,\n unsigned int len,\n cg::thread_block cta)\n{\n unsigned int len2 = 1 << (__btflo(len - 1U) + 1); // Round up len to nearest power-of-2\n\n if (threadIdx.x >= len2)\n return; // Early out for case where more threads launched than there is\n // data\n\n // First, set up our unused values to be the max data type.\n for (unsigned int i = len; i < len2; i += blockDim.x) {\n unsigned int index = i + threadIdx.x;\n\n if (index < len2) {\n // Must split our index between two buffers\n if (index < len)\n indata[index + offset] = 0xffffffffu;\n else\n backbuf[index + offset - len] = 0xffffffffu;\n }\n }\n\n cg::sync(cta);\n\n // Now the sort loops\n // Here, \"k\" is the sort level (remember bitonic does a multi-level butterfly\n // style sort)\n // and \"j\" is the partner element in the butterfly.\n // Two threads each work on one butterfly, because the read/write needs to\n // happen\n // simultaneously\n for (unsigned int k = 2; k <= len2; k *= 2) // Butterfly stride increments in powers of 2\n {\n for (unsigned int j = k >> 1; j > 0; j >>= 1) // Strides also in powers of to, up to index) {\n unsigned my_elem, swap_elem;\n\n if (index < len)\n my_elem = indata[index + offset];\n else\n my_elem = backbuf[index + offset - len];\n\n if (swap_idx < len)\n swap_elem = indata[swap_idx + offset];\n else\n swap_elem = backbuf[swap_idx + offset - len];\n\n // The k'th bit of my index (and hence my sort item ID)\n // determines if we sort ascending or descending.\n // Also, if either my_elem or swap_elem is out of range, then it\n // ALWAYS acts like it's the largest number.\n bool swap = false;\n\n if ((index & k) == 0) {\n if (my_elem > swap_elem)\n swap = true;\n }\n\n if ((index & k) == k) {\n if (my_elem < swap_elem)\n swap = true;\n }\n\n // If we had to swap, then write my data to the other element's\n // position.\n if (swap) {\n if (swap_idx < len)\n indata[swap_idx + offset] = my_elem;\n else\n backbuf[swap_idx + offset - len] = my_elem;\n\n if (index < len)\n indata[index + offset] = swap_elem;\n else\n backbuf[index + offset - len] = swap_elem;\n }\n }\n }\n\n cg::sync(cta); // Only need to sync for each \"j\" pass\n }\n }\n\n // Copy the sorted data from the input to the output buffer, because we sort\n // in-place\n if (outdata != indata) {\n for (unsigned int i = 0; i < len; i += blockDim.x) {\n unsigned int index = i + threadIdx.x;\n\n if (index < len)\n outdata[index + offset] = indata[index + offset];\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// KERNELS\n////////////////////////////////////////////////////////////////////////////////\n\n__global__ void bitonicsort(unsigned *indata, unsigned *outdata, unsigned int offset, unsigned int len)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n bitonicsort_kernel(indata, outdata, offset, len, cta);\n}\n\n__global__ void\nbig_bitonicsort(unsigned *indata, unsigned *outdata, unsigned *backbuf, unsigned int offset, unsigned int len)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n big_bitonicsort_kernel(indata, outdata, backbuf, offset, len, cta);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/cdpAdvancedQuicksort/cdpBitonicSort.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// cdpAdvancedQuicksort.cu\n//\n// Implementation of a parallel quicksort in CUDA. It comes in\n// several parts:\n//\n// 1. A small-set insertion sort. We do this on any set with <=32 elements\n// 2. A partitioning kernel, which - given a pivot - separates an input\n// array into elements <=pivot, and >pivot. Two quicksorts will then\n// be launched to resolve each of these.\n// 3. A quicksort co-ordinator, which figures out what kernels to launch\n// and when.\n//\n////////////////////////////////////////////////////////////////////////////////\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \n#include \n\n#include \"cdpQuicksort.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Inline PTX call to return index of highest non-zero bit in a word\n////////////////////////////////////////////////////////////////////////////////\nstatic __device__ __forceinline__ unsigned int __qsflo(unsigned int word)\n{\n unsigned int ret;\n asm volatile(\"bfind.u32 %0, %1;\" : \"=r\"(ret) : \"r\"(word));\n return ret;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// ringbufAlloc\n//\n// Allocates from a ringbuffer. Allows for not failing when we run out\n// of stack for tracking the offset counts for each sort subsection.\n//\n// We use the atomicMax trick to allow out-of-order retirement. If we\n// hit the size limit on the ringbuffer, then we spin-wait for people\n// to complete.\n//\n////////////////////////////////////////////////////////////////////////////////\ntemplate static __device__ T *ringbufAlloc(qsortRingbuf *ringbuf)\n{\n // Wait for there to be space in the ring buffer. We'll retry only a fixed\n // number of times and then fail, to avoid an out-of-memory deadlock.\n unsigned int loop = 10000;\n\n while (((ringbuf->head - ringbuf->tail) >= ringbuf->stacksize) && (loop-- > 0))\n ;\n\n if (loop == 0)\n return NULL;\n\n // Note that the element includes a little index book-keeping, for freeing\n // later.\n unsigned int index = atomicAdd((unsigned int *)&ringbuf->head, 1);\n T *ret = (T *)(ringbuf->stackbase) + (index & (ringbuf->stacksize - 1));\n ret->index = index;\n\n return ret;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// ringBufFree\n//\n// Releases an element from the ring buffer. If every element is released\n// up to and including this one, we can advance the tail to indicate that\n// space is now available.\n//\n////////////////////////////////////////////////////////////////////////////////\ntemplate static __device__ void ringbufFree(qsortRingbuf *ringbuf, T *data)\n{\n unsigned int index = data->index; // Non-wrapped index to free\n unsigned int count = atomicAdd((unsigned int *)&(ringbuf->count), 1) + 1;\n unsigned int max = atomicMax((unsigned int *)&(ringbuf->max), index + 1);\n\n // Update the tail if need be. Note we update \"max\" to be the new value in\n // ringbuf->max\n if (max < (index + 1))\n max = index + 1;\n\n if (max == count)\n atomicMax((unsigned int *)&(ringbuf->tail), count);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// qsort_warp\n//\n// Simplest possible implementation, does a per-warp quicksort with no\n// inter-warp\n// communication. This has a high atomic issue rate, but the rest should\n// actually\n// be fairly quick because of low work per thread.\n//\n// A warp finds its section of the data, then writes all data pivot to the other. Atomics are used to get a unique\n// section of the buffer.\n//\n// Obvious optimisation: do multiple chunks per warp, to increase in-flight\n// loads\n// and cover the instruction overhead.\n//\n////////////////////////////////////////////////////////////////////////////////\n__global__ void qsort_warp(unsigned *indata,\n unsigned *outdata,\n unsigned int offset,\n unsigned int len,\n qsortAtomicData *atomicData,\n qsortRingbuf *atomicDataStack,\n unsigned int source_is_indata,\n unsigned int depth)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Find my data offset, based on warp ID\n unsigned int thread_id = threadIdx.x + (blockIdx.x << QSORT_BLOCKSIZE_SHIFT);\n // unsigned int warp_id = threadIdx.x >> 5; // Used for debug only\n unsigned int lane_id = threadIdx.x & (warpSize - 1);\n\n // Exit if I'm outside the range of sort to be done\n if (thread_id >= len)\n return;\n\n //\n // First part of the algorithm. Each warp counts the number of elements that\n // are\n // greater/less than the pivot.\n //\n // When a warp knows its count, it updates an atomic counter.\n //\n\n // Read in the data and the pivot. Arbitrary pivot selection for now.\n unsigned pivot = indata[offset + len / 2];\n unsigned data = indata[offset + thread_id];\n\n // Count how many are <= and how many are > pivot.\n // If all are <= pivot then we adjust the comparison\n // because otherwise the sort will move nothing and\n // we'll iterate forever.\n cg::coalesced_group active = cg::coalesced_threads();\n unsigned int greater = (data > pivot);\n unsigned int gt_mask = active.ballot(greater);\n\n if (gt_mask == 0) {\n greater = (data >= pivot);\n gt_mask = active.ballot(greater); // Must re-ballot for adjusted comparator\n }\n\n unsigned int lt_mask = active.ballot(!greater);\n unsigned int gt_count = __popc(gt_mask);\n unsigned int lt_count = __popc(lt_mask);\n\n // Atomically adjust the lt_ and gt_offsets by this amount. Only one thread\n // need do this. Share the result using shfl\n unsigned int lt_offset, gt_offset;\n\n if (lane_id == 0) {\n if (lt_count > 0)\n lt_offset = atomicAdd((unsigned int *)&atomicData->lt_offset, lt_count);\n\n if (gt_count > 0)\n gt_offset = len - (atomicAdd((unsigned int *)&atomicData->gt_offset, gt_count) + gt_count);\n }\n\n lt_offset = active.shfl((int)lt_offset, 0); // Everyone pulls the offsets from lane 0\n gt_offset = active.shfl((int)gt_offset, 0);\n\n // Now compute my own personal offset within this. I need to know how many\n // threads with a lane ID less than mine are going to write to the same buffer\n // as me. We can use popc to implement a single-operation warp scan in this\n // case.\n unsigned lane_mask_lt;\n asm(\"mov.u32 %0, %%lanemask_lt;\" : \"=r\"(lane_mask_lt));\n unsigned int my_mask = greater ? gt_mask : lt_mask;\n unsigned int my_offset = __popc(my_mask & lane_mask_lt);\n\n // Move data.\n my_offset += greater ? gt_offset : lt_offset;\n outdata[offset + my_offset] = data;\n\n // Count up if we're the last warp in. If so, then Kepler will launch the next\n // set of sorts directly from here.\n if (lane_id == 0) {\n // Count \"elements written\". If I wrote the last one, then trigger the next\n // qsorts\n unsigned int mycount = lt_count + gt_count;\n\n if (atomicAdd((unsigned int *)&atomicData->sorted_count, mycount) + mycount == len) {\n // We're the last warp to do any sorting. Therefore it's up to us to\n // launch the next stage.\n unsigned int lt_len = atomicData->lt_offset;\n unsigned int gt_len = atomicData->gt_offset;\n\n cudaStream_t lstream, rstream;\n cudaStreamCreateWithFlags(&lstream, cudaStreamNonBlocking);\n cudaStreamCreateWithFlags(&rstream, cudaStreamNonBlocking);\n\n // Begin by freeing our atomicData storage. It's better for the ringbuffer\n // algorithm\n // if we free when we're done, rather than re-using (makes for less\n // fragmentation).\n ringbufFree(atomicDataStack, atomicData);\n\n // Exceptional case: if \"lt_len\" is zero, then all values in the batch\n // are equal. We are then done (may need to copy into correct buffer,\n // though)\n if (lt_len == 0) {\n if (source_is_indata)\n cudaMemcpyAsync(indata + offset,\n outdata + offset,\n gt_len * sizeof(unsigned),\n cudaMemcpyDeviceToDevice,\n lstream);\n\n return;\n }\n\n // Start with lower half first\n if (lt_len > BITONICSORT_LEN) {\n // If we've exceeded maximum depth, fall through to backup\n // big_bitonicsort\n if (depth >= QSORT_MAXDEPTH) {\n // The final bitonic stage sorts in-place in \"outdata\". We therefore\n // re-use \"indata\" as the out-of-range tracking buffer. For (2^n)+1\n // elements we need (2^(n+1)) bytes of oor buffer. The backup qsort\n // buffer is at least this large when sizeof(QTYPE) >= 2.\n big_bitonicsort<<<1, BITONICSORT_LEN, 0, lstream>>>(\n outdata, source_is_indata ? indata : outdata, indata, offset, lt_len);\n }\n else {\n // Launch another quicksort. We need to allocate more storage for the\n // atomic data.\n if ((atomicData = ringbufAlloc(atomicDataStack)) == NULL)\n printf(\"Stack-allocation error. Failing left child launch.\\n\");\n else {\n atomicData->lt_offset = atomicData->gt_offset = atomicData->sorted_count = 0;\n unsigned int numblocks =\n (unsigned int)(lt_len + (QSORT_BLOCKSIZE - 1)) >> QSORT_BLOCKSIZE_SHIFT;\n qsort_warp<<>>(\n outdata, indata, offset, lt_len, atomicData, atomicDataStack, !source_is_indata, depth + 1);\n }\n }\n }\n else if (lt_len > 1) {\n // Final stage uses a bitonic sort instead. It's important to\n // make sure the final stage ends up in the correct (original) buffer.\n // We launch the smallest power-of-2 number of threads that we can.\n unsigned int bitonic_len = 1 << (__qsflo(lt_len - 1U) + 1);\n bitonicsort<<<1, bitonic_len, 0, lstream>>>(\n outdata, source_is_indata ? indata : outdata, offset, lt_len);\n }\n // Finally, if we sorted just one single element, we must still make\n // sure that it winds up in the correct place.\n else if (source_is_indata && (lt_len == 1))\n indata[offset] = outdata[offset];\n\n if (cudaPeekAtLastError() != cudaSuccess)\n printf(\"Left-side launch fail: %s\\n\", cudaGetErrorString(cudaGetLastError()));\n\n // Now the upper half.\n if (gt_len > BITONICSORT_LEN) {\n // If we've exceeded maximum depth, fall through to backup\n // big_bitonicsort\n if (depth >= QSORT_MAXDEPTH)\n big_bitonicsort<<<1, BITONICSORT_LEN, 0, rstream>>>(\n outdata, source_is_indata ? indata : outdata, indata, offset + lt_len, gt_len);\n else {\n // Allocate new atomic storage for this launch\n if ((atomicData = ringbufAlloc(atomicDataStack)) == NULL)\n printf(\"Stack allocation error! Failing right-side launch.\\n\");\n else {\n atomicData->lt_offset = atomicData->gt_offset = atomicData->sorted_count = 0;\n unsigned int numblocks =\n (unsigned int)(gt_len + (QSORT_BLOCKSIZE - 1)) >> QSORT_BLOCKSIZE_SHIFT;\n qsort_warp<<>>(outdata,\n indata,\n offset + lt_len,\n gt_len,\n atomicData,\n atomicDataStack,\n !source_is_indata,\n depth + 1);\n }\n }\n }\n else if (gt_len > 1) {\n unsigned int bitonic_len = 1 << (__qsflo(gt_len - 1U) + 1);\n bitonicsort<<<1, bitonic_len, 0, rstream>>>(\n outdata, source_is_indata ? indata : outdata, offset + lt_len, gt_len);\n }\n else if (source_is_indata && (gt_len == 1))\n indata[offset + lt_len] = outdata[offset + lt_len];\n\n if (cudaPeekAtLastError() != cudaSuccess)\n printf(\"Right-side launch fail: %s\\n\", cudaGetErrorString(cudaGetLastError()));\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// run_quicksort\n//\n// Host-side code to run the Kepler version of quicksort. It's pretty\n// simple, because all launch control is handled on the device via CDP.\n//\n// All parallel quicksorts require an equal-sized scratch buffer. This\n// must be passed in ahead of time.\n//\n// Returns the time elapsed for the sort.\n//\n////////////////////////////////////////////////////////////////////////////////\nfloat run_quicksort_cdp(unsigned *gpudata, unsigned *scratchdata, unsigned int count, cudaStream_t stream)\n{\n unsigned int stacksize = QSORT_STACK_ELEMS;\n\n // This is the stack, for atomic tracking of each sort's status\n qsortAtomicData *gpustack;\n checkCudaErrors(cudaMalloc((void **)&gpustack, stacksize * sizeof(qsortAtomicData)));\n checkCudaErrors(cudaMemset(gpustack, 0, sizeof(qsortAtomicData))); // Only need set first entry to 0\n\n // Create the memory ringbuffer used for handling the stack.\n // Initialise everything to where it needs to be.\n qsortRingbuf buf;\n qsortRingbuf *ringbuf;\n checkCudaErrors(cudaMalloc((void **)&ringbuf, sizeof(qsortRingbuf)));\n buf.head = 1; // We start with one allocation\n buf.tail = 0;\n buf.count = 0;\n buf.max = 0;\n buf.stacksize = stacksize;\n buf.stackbase = gpustack;\n checkCudaErrors(cudaMemcpy(ringbuf, &buf, sizeof(buf), cudaMemcpyHostToDevice));\n\n // Timing events...\n cudaEvent_t ev1, ev2;\n checkCudaErrors(cudaEventCreate(&ev1));\n checkCudaErrors(cudaEventCreate(&ev2));\n checkCudaErrors(cudaEventRecord(ev1));\n\n // Now we trivially launch the qsort kernel\n if (count > BITONICSORT_LEN) {\n unsigned int numblocks = (unsigned int)(count + (QSORT_BLOCKSIZE - 1)) >> QSORT_BLOCKSIZE_SHIFT;\n qsort_warp<<>>(\n gpudata, scratchdata, 0U, count, gpustack, ringbuf, true, 0);\n }\n else {\n bitonicsort<<<1, BITONICSORT_LEN>>>(gpudata, gpudata, 0, count);\n }\n\n checkCudaErrors(cudaGetLastError());\n checkCudaErrors(cudaEventRecord(ev2));\n checkCudaErrors(cudaDeviceSynchronize());\n\n float elapse = 0.0f;\n\n if (cudaPeekAtLastError() != cudaSuccess)\n printf(\"Launch failure: %s\\n\", cudaGetErrorString(cudaGetLastError()));\n else\n checkCudaErrors(cudaEventElapsedTime(&elapse, ev1, ev2));\n\n // Sanity check that the stack allocator is doing the right thing\n checkCudaErrors(cudaMemcpy(&buf, ringbuf, sizeof(*ringbuf), cudaMemcpyDeviceToHost));\n\n if (count > BITONICSORT_LEN && buf.head != buf.tail) {\n printf(\"Stack allocation error!\\nRingbuf:\\n\");\n printf(\"\\t head = %u\\n\", buf.head);\n printf(\"\\t tail = %u\\n\", buf.tail);\n printf(\"\\tcount = %u\\n\", buf.count);\n printf(\"\\t max = %u\\n\", buf.max);\n }\n\n // Release our stack data once we're done\n checkCudaErrors(cudaFree(ringbuf));\n checkCudaErrors(cudaFree(gpustack));\n\n return elapse;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n////////////////////////////////////////////////////////////////////////////////\nint run_qsort(unsigned int size, int seed, int debug, int loop, int verbose)\n{\n if (seed > 0)\n srand(seed);\n\n // Create and set up our test\n unsigned *gpudata, *scratchdata;\n checkCudaErrors(cudaMalloc((void **)&gpudata, size * sizeof(unsigned)));\n checkCudaErrors(cudaMalloc((void **)&scratchdata, size * sizeof(unsigned)));\n\n // Create CPU data.\n unsigned *data = new unsigned[size];\n unsigned int min = loop ? loop : size;\n unsigned int max = size;\n loop = (loop == 0) ? 1 : loop;\n\n for (size = min; size <= max; size += loop) {\n if (verbose)\n printf(\" Input: \");\n\n for (unsigned int i = 0; i < size; i++) {\n // Build data 8 bits at a time\n data[i] = 0;\n char *ptr = (char *)&(data[i]);\n\n for (unsigned j = 0; j < sizeof(unsigned); j++) {\n // Easy-to-read data in debug mode\n if (debug) {\n *ptr++ = (char)(rand() % 10);\n break;\n }\n\n *ptr++ = (char)(rand() & 255);\n }\n\n if (verbose) {\n if (i && !(i % 32))\n printf(\"\\n \");\n\n printf(\"%u \", data[i]);\n }\n }\n\n if (verbose)\n printf(\"\\n\");\n\n checkCudaErrors(cudaMemcpy(gpudata, data, size * sizeof(unsigned), cudaMemcpyHostToDevice));\n\n // So we're now populated and ready to go! We size our launch as\n // blocks of up to BLOCKSIZE threads, and appropriate grid size.\n // One thread is launched per element.\n float elapse;\n elapse = run_quicksort_cdp(gpudata, scratchdata, size, NULL);\n\n // run_bitonicsort(gpudata, scratchdata, size, verbose);\n checkCudaErrors(cudaDeviceSynchronize());\n\n // Copy back the data and verify correct sort\n checkCudaErrors(cudaMemcpy(data, gpudata, size * sizeof(unsigned), cudaMemcpyDeviceToHost));\n\n if (verbose) {\n printf(\"Output: \");\n\n for (unsigned int i = 0; i < size; i++) {\n if (i && !(i % 32))\n printf(\"\\n \");\n\n printf(\"%u \", data[i]);\n }\n\n printf(\"\\n\");\n }\n\n unsigned int check;\n\n for (check = 1; check < size; check++) {\n if (data[check] < data[check - 1]) {\n printf(\"FAILED at element: %d\\n\", check);\n break;\n }\n }\n\n if (check != size) {\n printf(\" cdpAdvancedQuicksort FAILED\\n\");\n exit(EXIT_FAILURE);\n }\n else\n printf(\" cdpAdvancedQuicksort PASSED\\n\");\n\n // Display the time between event recordings\n printf(\"Sorted %u elems in %.3f ms (%.3f Melems/sec)\\n\", size, elapse, (float)size / (elapse * 1000.0f));\n fflush(stdout);\n }\n\n // Release everything and we're done\n checkCudaErrors(cudaFree(scratchdata));\n checkCudaErrors(cudaFree(gpudata));\n delete (data);\n return 0;\n}\n\nstatic void usage()\n{\n printf(\"Syntax: cdpAdvancedQuicksort [-size=] [-seed=] [-debug] \"\n \"[-loop-step=] [-verbose]\\n\");\n printf(\"If loop_step is non-zero, will run from 1->array_len in steps of \"\n \"loop_step\\n\");\n}\n\n// Host side entry\nint main(int argc, char *argv[])\n{\n int size = 1000000;\n unsigned int seed = 0;\n int debug = 0;\n int loop = 0;\n int verbose = 0;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\") || checkCmdLineFlag(argc, (const char **)argv, \"h\")) {\n usage();\n printf(\"&&&& cdpAdvancedQuicksort WAIVED\\n\");\n exit(EXIT_WAIVED);\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"size\")) {\n size = getCmdLineArgumentInt(argc, (const char **)argv, \"size\");\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"seed\")) {\n seed = getCmdLineArgumentInt(argc, (const char **)argv, \"seed\");\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"loop-step\")) {\n loop = getCmdLineArgumentInt(argc, (const char **)argv, \"loop-step\");\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"debug\")) {\n debug = 1;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"verbose\")) {\n verbose = 1;\n }\n\n // Get device properties\n int cuda_device = findCudaDevice(argc, (const char **)argv);\n cudaDeviceProp properties;\n checkCudaErrors(cudaGetDeviceProperties(&properties, cuda_device));\n int cdpCapable = (properties.major == 3 && properties.minor >= 5) || properties.major >= 4;\n\n printf(\"GPU device %s has compute capabilities (SM %d.%d)\\n\", properties.name, properties.major, properties.minor);\n\n if (!cdpCapable) {\n printf(\"cdpAdvancedQuicksort requires SM 3.5 or higher to use CUDA Dynamic \"\n \"Parallelism. Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n printf(\"Running qsort on %d elements with seed %d, on %s\\n\", size, seed, properties.name);\n\n // Increase the pending kernel launch limit\n cudaDeviceSetLimit(cudaLimitDevRuntimePendingLaunchCount, 4096);\n\n run_qsort(size, seed, debug, loop, verbose);\n\n exit(EXIT_SUCCESS);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// cdpAdvancedQuicksort.cu\n//\n// Implementation of a parallel quicksort in CUDA. It comes in\n// several parts:\n//\n// 1. A small-set insertion sort. We do this on any set with <=32 elements\n// 2. A partitioning kernel, which - given a pivot - separates an input\n// array into elements <=pivot, and >pivot. Two quicksorts will then\n// be launched to resolve each of these.\n// 3. A quicksort co-ordinator, which figures out what kernels to launch\n// and when.\n//\n////////////////////////////////////////////////////////////////////////////////\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \n#include \n\n#include \"cdpQuicksort.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Inline PTX call to return index of highest non-zero bit in a word\n////////////////////////////////////////////////////////////////////////////////\nstatic __device__ __forceinline__ unsigned int __qsflo(unsigned int word)\n{\n unsigned int ret;\n asm volatile(\"bfind.u32 %0, %1;\" : \"=r\"(ret) : \"r\"(word));\n return ret;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// ringbufAlloc\n//\n// Allocates from a ringbuffer. Allows for not failing when we run out\n// of stack for tracking the offset counts for each sort subsection.\n//\n// We use the atomicMax trick to allow out-of-order retirement. If we\n// hit the size limit on the ringbuffer, then we spin-wait for people\n// to complete.\n//\n////////////////////////////////////////////////////////////////////////////////\ntemplate static __device__ T *ringbufAlloc(qsortRingbuf *ringbuf)\n{\n // Wait for there to be space in the ring buffer. We'll retry only a fixed\n // number of times and then fail, to avoid an out-of-memory deadlock.\n unsigned int loop = 10000;\n\n while (((ringbuf->head - ringbuf->tail) >= ringbuf->stacksize) && (loop-- > 0))\n ;\n\n if (loop == 0)\n return NULL;\n\n // Note that the element includes a little index book-keeping, for freeing\n // later.\n unsigned int index = atomicAdd((unsigned int *)&ringbuf->head, 1);\n T *ret = (T *)(ringbuf->stackbase) + (index & (ringbuf->stacksize - 1));\n ret->index = index;\n\n return ret;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// ringBufFree\n//\n// Releases an element from the ring buffer. If every element is released\n// up to and including this one, we can advance the tail to indicate that\n// space is now available.\n//\n////////////////////////////////////////////////////////////////////////////////\ntemplate static __device__ void ringbufFree(qsortRingbuf *ringbuf, T *data)\n{\n unsigned int index = data->index; // Non-wrapped index to free\n unsigned int count = atomicAdd((unsigned int *)&(ringbuf->count), 1) + 1;\n unsigned int max = atomicMax((unsigned int *)&(ringbuf->max), index + 1);\n\n // Update the tail if need be. Note we update \"max\" to be the new value in\n // ringbuf->max\n if (max < (index + 1))\n max = index + 1;\n\n if (max == count)\n atomicMax((unsigned int *)&(ringbuf->tail), count);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// qsort_warp\n//\n// Simplest possible implementation, does a per-warp quicksort with no\n// inter-warp\n// communication. This has a high atomic issue rate, but the rest should\n// actually\n// be fairly quick because of low work per thread.\n//\n// A warp finds its section of the data, then writes all data pivot to the other. Atomics are used to get a unique\n// section of the buffer.\n//\n// Obvious optimisation: do multiple chunks per warp, to increase in-flight\n// loads\n// and cover the instruction overhead.\n//\n////////////////////////////////////////////////////////////////////////////////\n__global__ void qsort_warp(unsigned *indata,\n unsigned *outdata,\n unsigned int offset,\n unsigned int len,\n qsortAtomicData *atomicData,\n qsortRingbuf *atomicDataStack,\n unsigned int source_is_indata,\n unsigned int depth)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Find my data offset, based on warp ID\n unsigned int thread_id = threadIdx.x + (blockIdx.x << QSORT_BLOCKSIZE_SHIFT);\n // unsigned int warp_id = threadIdx.x >> 5; // Used for debug only\n unsigned int lane_id = threadIdx.x & (warpSize - 1);\n\n // Exit if I'm outside the range of sort to be done\n if (thread_id >= len)\n return;\n\n //\n // First part of the algorithm. Each warp counts the number of elements that\n // are\n // greater/less than the pivot.\n //\n // When a warp knows its count, it updates an atomic counter.\n //\n\n // Read in the data and the pivot. Arbitrary pivot selection for now.\n unsigned pivot = indata[offset + len / 2];\n unsigned data = indata[offset + thread_id];\n\n // Count how many are <= and how many are > pivot.\n // If all are <= pivot then we adjust the comparison\n // because otherwise the sort will move nothing and\n // we'll iterate forever.\n cg::coalesced_group active = cg::coalesced_threads();\n unsigned int greater = (data > pivot);\n unsigned int gt_mask = active.ballot(greater);\n\n if (gt_mask == 0) {\n greater = (data >= pivot);\n gt_mask = active.ballot(greater); // Must re-ballot for adjusted comparator\n }\n\n unsigned int lt_mask = active.ballot(!greater);\n unsigned int gt_count = __popc(gt_mask);\n unsigned int lt_count = __popc(lt_mask);\n\n // Atomically adjust the lt_ and gt_offsets by this amount. Only one thread\n // need do this. Share the result using shfl\n unsigned int lt_offset, gt_offset;\n\n if (lane_id == 0) {\n if (lt_count > 0)\n lt_offset = atomicAdd((unsigned int *)&atomicData->lt_offset, lt_count);\n\n if (gt_count > 0)\n gt_offset = len - (atomicAdd((unsigned int *)&atomicData->gt_offset, gt_count) + gt_count);\n }\n\n lt_offset = active.shfl((int)lt_offset, 0); // Everyone pulls the offsets from lane 0\n gt_offset = active.shfl((int)gt_offset, 0);\n\n // Now compute my own personal offset within this. I need to know how many\n // threads with a lane ID less than mine are going to write to the same buffer\n // as me. We can use popc to implement a single-operation warp scan in this\n // case.\n unsigned lane_mask_lt;\n asm(\"mov.u32 %0, %%lanemask_lt;\" : \"=r\"(lane_mask_lt));\n unsigned int my_mask = greater ? gt_mask : lt_mask;\n unsigned int my_offset = __popc(my_mask & lane_mask_lt);\n\n // Move data.\n my_offset += greater ? gt_offset : lt_offset;\n outdata[offset + my_offset] = data;\n\n // Count up if we're the last warp in. If so, then Kepler will launch the next\n // set of sorts directly from here.\n if (lane_id == 0) {\n // Count \"elements written\". If I wrote the last one, then trigger the next\n // qsorts\n unsigned int mycount = lt_count + gt_count;\n\n if (atomicAdd((unsigned int *)&atomicData->sorted_count, mycount) + mycount == len) {\n // We're the last warp to do any sorting. Therefore it's up to us to\n // launch the next stage.\n unsigned int lt_len = atomicData->lt_offset;\n unsigned int gt_len = atomicData->gt_offset;\n\n hipStream_t lstream, rstream;\n hipStreamCreateWithFlags(&lstream, hipStreamNonBlocking);\n hipStreamCreateWithFlags(&rstream, hipStreamNonBlocking);\n\n // Begin by freeing our atomicData storage. It's better for the ringbuffer\n // algorithm\n // if we free when we're done, rather than re-using (makes for less\n // fragmentation).\n ringbufFree(atomicDataStack, atomicData);\n\n // Exceptional case: if \"lt_len\" is zero, then all values in the batch\n // are equal. We are then done (may need to copy into correct buffer,\n // though)\n if (lt_len == 0) {\n if (source_is_indata)\n hipMemcpyAsync(indata + offset,\n outdata + offset,\n gt_len * sizeof(unsigned),\n hipMemcpyDeviceToDevice,\n lstream);\n\n return;\n }\n\n // Start with lower half first\n if (lt_len > BITONICSORT_LEN) {\n // If we've exceeded maximum depth, fall through to backup\n // big_bitonicsort\n if (depth >= QSORT_MAXDEPTH) {\n // The final bitonic stage sorts in-place in \"outdata\". We therefore\n // re-use \"indata\" as the out-of-range tracking buffer. For (2^n)+1\n // elements we need (2^(n+1)) bytes of oor buffer. The backup qsort\n // buffer is at least this large when sizeof(QTYPE) >= 2.\n big_bitonicsort<<<1, BITONICSORT_LEN, 0, lstream>>>(\n outdata, source_is_indata ? indata : outdata, indata, offset, lt_len);\n }\n else {\n // Launch another quicksort. We need to allocate more storage for the\n // atomic data.\n if ((atomicData = ringbufAlloc(atomicDataStack)) == NULL)\n printf(\"Stack-allocation error. Failing left child launch.\\n\");\n else {\n atomicData->lt_offset = atomicData->gt_offset = atomicData->sorted_count = 0;\n unsigned int numblocks =\n (unsigned int)(lt_len + (QSORT_BLOCKSIZE - 1)) >> QSORT_BLOCKSIZE_SHIFT;\n qsort_warp<<>>(\n outdata, indata, offset, lt_len, atomicData, atomicDataStack, !source_is_indata, depth + 1);\n }\n }\n }\n else if (lt_len > 1) {\n // Final stage uses a bitonic sort instead. It's important to\n // make sure the final stage ends up in the correct (original) buffer.\n // We launch the smallest power-of-2 number of threads that we can.\n unsigned int bitonic_len = 1 << (__qsflo(lt_len - 1U) + 1);\n bitonicsort<<<1, bitonic_len, 0, lstream>>>(\n outdata, source_is_indata ? indata : outdata, offset, lt_len);\n }\n // Finally, if we sorted just one single element, we must still make\n // sure that it winds up in the correct place.\n else if (source_is_indata && (lt_len == 1))\n indata[offset] = outdata[offset];\n\n if (hipPeekAtLastError() != hipSuccess)\n printf(\"Left-side launch fail: %s\\n\", hipGetErrorString(hipGetLastError()));\n\n // Now the upper half.\n if (gt_len > BITONICSORT_LEN) {\n // If we've exceeded maximum depth, fall through to backup\n // big_bitonicsort\n if (depth >= QSORT_MAXDEPTH)\n big_bitonicsort<<<1, BITONICSORT_LEN, 0, rstream>>>(\n outdata, source_is_indata ? indata : outdata, indata, offset + lt_len, gt_len);\n else {\n // Allocate new atomic storage for this launch\n if ((atomicData = ringbufAlloc(atomicDataStack)) == NULL)\n printf(\"Stack allocation error! Failing right-side launch.\\n\");\n else {\n atomicData->lt_offset = atomicData->gt_offset = atomicData->sorted_count = 0;\n unsigned int numblocks =\n (unsigned int)(gt_len + (QSORT_BLOCKSIZE - 1)) >> QSORT_BLOCKSIZE_SHIFT;\n qsort_warp<<>>(outdata,\n indata,\n offset + lt_len,\n gt_len,\n atomicData,\n atomicDataStack,\n !source_is_indata,\n depth + 1);\n }\n }\n }\n else if (gt_len > 1) {\n unsigned int bitonic_len = 1 << (__qsflo(gt_len - 1U) + 1);\n bitonicsort<<<1, bitonic_len, 0, rstream>>>(\n outdata, source_is_indata ? indata : outdata, offset + lt_len, gt_len);\n }\n else if (source_is_indata && (gt_len == 1))\n indata[offset + lt_len] = outdata[offset + lt_len];\n\n if (hipPeekAtLastError() != hipSuccess)\n printf(\"Right-side launch fail: %s\\n\", hipGetErrorString(hipGetLastError()));\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// run_quicksort\n//\n// Host-side code to run the Kepler version of quicksort. It's pretty\n// simple, because all launch control is handled on the device via CDP.\n//\n// All parallel quicksorts require an equal-sized scratch buffer. This\n// must be passed in ahead of time.\n//\n// Returns the time elapsed for the sort.\n//\n////////////////////////////////////////////////////////////////////////////////\nfloat run_quicksort_cdp(unsigned *gpudata, unsigned *scratchdata, unsigned int count, hipStream_t stream)\n{\n unsigned int stacksize = QSORT_STACK_ELEMS;\n\n // This is the stack, for atomic tracking of each sort's status\n qsortAtomicData *gpustack;\n checkCudaErrors(hipMalloc((void **)&gpustack, stacksize * sizeof(qsortAtomicData)));\n checkCudaErrors(hipMemset(gpustack, 0, sizeof(qsortAtomicData))); // Only need set first entry to 0\n\n // Create the memory ringbuffer used for handling the stack.\n // Initialise everything to where it needs to be.\n qsortRingbuf buf;\n qsortRingbuf *ringbuf;\n checkCudaErrors(hipMalloc((void **)&ringbuf, sizeof(qsortRingbuf)));\n buf.head = 1; // We start with one allocation\n buf.tail = 0;\n buf.count = 0;\n buf.max = 0;\n buf.stacksize = stacksize;\n buf.stackbase = gpustack;\n checkCudaErrors(hipMemcpy(ringbuf, &buf, sizeof(buf), hipMemcpyHostToDevice));\n\n // Timing events...\n hipEvent_t ev1, ev2;\n checkCudaErrors(hipEventCreate(&ev1));\n checkCudaErrors(hipEventCreate(&ev2));\n checkCudaErrors(hipEventRecord(ev1));\n\n // Now we trivially launch the qsort kernel\n if (count > BITONICSORT_LEN) {\n unsigned int numblocks = (unsigned int)(count + (QSORT_BLOCKSIZE - 1)) >> QSORT_BLOCKSIZE_SHIFT;\n qsort_warp<<>>(\n gpudata, scratchdata, 0U, count, gpustack, ringbuf, true, 0);\n }\n else {\n bitonicsort<<<1, BITONICSORT_LEN>>>(gpudata, gpudata, 0, count);\n }\n\n checkCudaErrors(hipGetLastError());\n checkCudaErrors(hipEventRecord(ev2));\n checkCudaErrors(hipDeviceSynchronize());\n\n float elapse = 0.0f;\n\n if (hipPeekAtLastError() != hipSuccess)\n printf(\"Launch failure: %s\\n\", hipGetErrorString(hipGetLastError()));\n else\n checkCudaErrors(hipEventElapsedTime(&elapse, ev1, ev2));\n\n // Sanity check that the stack allocator is doing the right thing\n checkCudaErrors(hipMemcpy(&buf, ringbuf, sizeof(*ringbuf), hipMemcpyDeviceToHost));\n\n if (count > BITONICSORT_LEN && buf.head != buf.tail) {\n printf(\"Stack allocation error!\\nRingbuf:\\n\");\n printf(\"\\t head = %u\\n\", buf.head);\n printf(\"\\t tail = %u\\n\", buf.tail);\n printf(\"\\tcount = %u\\n\", buf.count);\n printf(\"\\t max = %u\\n\", buf.max);\n }\n\n // Release our stack data once we're done\n checkCudaErrors(hipFree(ringbuf));\n checkCudaErrors(hipFree(gpustack));\n\n return elapse;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n////////////////////////////////////////////////////////////////////////////////\nint run_qsort(unsigned int size, int seed, int debug, int loop, int verbose)\n{\n if (seed > 0)\n srand(seed);\n\n // Create and set up our test\n unsigned *gpudata, *scratchdata;\n checkCudaErrors(hipMalloc((void **)&gpudata, size * sizeof(unsigned)));\n checkCudaErrors(hipMalloc((void **)&scratchdata, size * sizeof(unsigned)));\n\n // Create CPU data.\n unsigned *data = new unsigned[size];\n unsigned int min = loop ? loop : size;\n unsigned int max = size;\n loop = (loop == 0) ? 1 : loop;\n\n for (size = min; size <= max; size += loop) {\n if (verbose)\n printf(\" Input: \");\n\n for (unsigned int i = 0; i < size; i++) {\n // Build data 8 bits at a time\n data[i] = 0;\n char *ptr = (char *)&(data[i]);\n\n for (unsigned j = 0; j < sizeof(unsigned); j++) {\n // Easy-to-read data in debug mode\n if (debug) {\n *ptr++ = (char)(rand() % 10);\n break;\n }\n\n *ptr++ = (char)(rand() & 255);\n }\n\n if (verbose) {\n if (i && !(i % 32))\n printf(\"\\n \");\n\n printf(\"%u \", data[i]);\n }\n }\n\n if (verbose)\n printf(\"\\n\");\n\n checkCudaErrors(hipMemcpy(gpudata, data, size * sizeof(unsigned), hipMemcpyHostToDevice));\n\n // So we're now populated and ready to go! We size our launch as\n // blocks of up to BLOCKSIZE threads, and appropriate grid size.\n // One thread is launched per element.\n float elapse;\n elapse = run_quicksort_cdp(gpudata, scratchdata, size, NULL);\n\n // run_bitonicsort(gpudata, scratchdata, size, verbose);\n checkCudaErrors(hipDeviceSynchronize());\n\n // Copy back the data and verify correct sort\n checkCudaErrors(hipMemcpy(data, gpudata, size * sizeof(unsigned), hipMemcpyDeviceToHost));\n\n if (verbose) {\n printf(\"Output: \");\n\n for (unsigned int i = 0; i < size; i++) {\n if (i && !(i % 32))\n printf(\"\\n \");\n\n printf(\"%u \", data[i]);\n }\n\n printf(\"\\n\");\n }\n\n unsigned int check;\n\n for (check = 1; check < size; check++) {\n if (data[check] < data[check - 1]) {\n printf(\"FAILED at element: %d\\n\", check);\n break;\n }\n }\n\n if (check != size) {\n printf(\" cdpAdvancedQuicksort FAILED\\n\");\n exit(EXIT_FAILURE);\n }\n else\n printf(\" cdpAdvancedQuicksort PASSED\\n\");\n\n // Display the time between event recordings\n printf(\"Sorted %u elems in %.3f ms (%.3f Melems/sec)\\n\", size, elapse, (float)size / (elapse * 1000.0f));\n fflush(stdout);\n }\n\n // Release everything and we're done\n checkCudaErrors(hipFree(scratchdata));\n checkCudaErrors(hipFree(gpudata));\n delete (data);\n return 0;\n}\n\nstatic void usage()\n{\n printf(\"Syntax: cdpAdvancedQuicksort [-size=] [-seed=] [-debug] \"\n \"[-loop-step=] [-verbose]\\n\");\n printf(\"If loop_step is non-zero, will run from 1->array_len in steps of \"\n \"loop_step\\n\");\n}\n\n// Host side entry\nint main(int argc, char *argv[])\n{\n int size = 1000000;\n unsigned int seed = 0;\n int debug = 0;\n int loop = 0;\n int verbose = 0;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\") || checkCmdLineFlag(argc, (const char **)argv, \"h\")) {\n usage();\n printf(\"&&&& cdpAdvancedQuicksort WAIVED\\n\");\n exit(EXIT_WAIVED);\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"size\")) {\n size = getCmdLineArgumentInt(argc, (const char **)argv, \"size\");\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"seed\")) {\n seed = getCmdLineArgumentInt(argc, (const char **)argv, \"seed\");\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"loop-step\")) {\n loop = getCmdLineArgumentInt(argc, (const char **)argv, \"loop-step\");\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"debug\")) {\n debug = 1;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"verbose\")) {\n verbose = 1;\n }\n\n // Get device properties\n int cuda_device = findCudaDevice(argc, (const char **)argv);\n hipDeviceProp_t properties;\n checkCudaErrors(hipGetDeviceProperties(&properties, cuda_device));\n int cdpCapable = (properties.major == 3 && properties.minor >= 5) || properties.major >= 4;\n\n printf(\"GPU device %s has compute capabilities (SM %d.%d)\\n\", properties.name, properties.major, properties.minor);\n\n if (!cdpCapable) {\n printf(\"cdpAdvancedQuicksort requires SM 3.5 or higher to use CUDA Dynamic \"\n \"Parallelism. Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n printf(\"Running qsort on %d elements with seed %d, on %s\\n\", size, seed, properties.name);\n\n // Increase the pending kernel launch limit\n hipDeviceSetLimit(cudaLimitDevRuntimePendingLaunchCount, 4096);\n\n run_qsort(size, seed, debug, loop, verbose);\n\n exit(EXIT_SUCCESS);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/cdpAdvancedQuicksort/cdpAdvancedQuicksort.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/**\n * Matrix multiplication: C = A * B.\n *\n * This sample demonstrates implements matrix multiplication which makes use of\n * shared memory to ensure data reuse, the matrix multiplication is done using\n * tiling approach.\n * With compute capability 8.0 or higher the CUDA kernels involved uses\n * asynchronously copy data from global to shared memory; a.k.a., async-copy.\n * This sample has been written for clarity of exposition to illustrate various\n * CUDA programming principles, not with the goal of providing the most\n * performant generic kernel for matrix multiplication.\n */\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n#include \n\n#if __CUDA_ARCH__ >= 700\n#include \n#endif\n#include \n\nnamespace cg = cooperative_groups;\n\n// Helper functions and utilities to work with CUDA\n#include \n#include \n\nenum kernels {\n AsyncCopyMultiStageLargeChunk = 0,\n AsyncCopyLargeChunk = 1,\n AsyncCopyLargeChunkAWBarrier = 2,\n AsyncCopyMultiStageSharedState = 3,\n AsyncCopyMultiStage = 4,\n AsyncCopySingleStage = 5,\n Naive = 6,\n NaiveLargeChunk = 7\n};\n\nconst char *kernelNames[] = {\"AsyncCopyMultiStageLargeChunk\",\n \"AsyncCopyLargeChunk\",\n \"AsyncCopyLargeChunkAWBarrier\",\n \"AsyncCopyMultiStageSharedState\",\n \"AsyncCopyMultiStage\",\n \"AsyncCopySingleStage\",\n \"Naive\",\n \"NaiveLargeChunk\"};\n\nconstexpr int blockSize = 16;\n\n// Multi Stage memcpy_async pipeline with large chunk copy\ntemplate \n__global__ void MatrixMulAsyncCopyMultiStageLargeChunk(float *__restrict__ C,\n const float *__restrict__ A,\n const float *__restrict__ B,\n int wA,\n int wB)\n{\n // Requires BLOCK_SIZE % 4 == 0\n\n // Multi-stage pipeline version\n constexpr size_t maxPipelineStages = 4;\n\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A for each stage\n __shared__ alignas(alignof(float4)) float As[maxPipelineStages][BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B for each stage\n __shared__ alignas(alignof(float4)) float Bs[maxPipelineStages][BLOCK_SIZE][BLOCK_SIZE];\n\n float Csub = 0.0;\n\n // Index of the first sub-matrix of A processed by the block\n const int aBegin = wA * (BLOCK_SIZE)*blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n const int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n const int bBegin = BLOCK_SIZE * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n const int t4x = threadIdx.x * 4;\n const auto shape4 = cuda::aligned_size_t(sizeof(float4));\n\n cuda::pipeline pipe = cuda::make_pipeline();\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin, i = 0, aStage = aBegin, bStage = bBegin, iStage = 0; a <= aEnd;\n a += aStep, b += bStep, ++i) {\n // Load the matrices from device memory to shared memory; each thread loads\n // one element of each matrix\n for (; aStage <= a + aStep * maxPipelineStages; aStage += aStep, bStage += bStep, ++iStage) {\n pipe.producer_acquire();\n if (aStage <= aEnd && t4x < BLOCK_SIZE) {\n // Rotating buffer\n const int j = iStage % maxPipelineStages;\n cuda::memcpy_async(&As[j][threadIdx.y][t4x], &A[aStage + wA * threadIdx.y + t4x], shape4, pipe);\n cuda::memcpy_async(&Bs[j][threadIdx.y][t4x], &B[aStage + wA * threadIdx.y + t4x], shape4, pipe);\n }\n pipe.producer_commit();\n }\n\n pipe.consumer_wait();\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n // Rotating buffer\n const int j = i % maxPipelineStages;\n\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[j][threadIdx.y][k] * Bs[j][k][threadIdx.x];\n }\n pipe.consumer_release();\n\n // Don't have to synchronize because maxPipelineStages is greater than one\n // therefore next iteration is loading to a different buffer.\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes four element\n int c = wB * BLOCK_SIZE * blockIdx.y + BLOCK_SIZE * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n}\n\n// Single Stage memcpy_async pipeline with Large copy chunk (float4)\ntemplate \n__global__ void MatrixMulAsyncCopyLargeChunk(float *__restrict__ C,\n const float *__restrict__ A,\n const float *__restrict__ B,\n int wA,\n int wB)\n{\n // Requires BLOCK_SIZE % 4 == 0\n\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ alignas(alignof(float4)) float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ alignas(alignof(float4)) float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = wA * BLOCK_SIZE * blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n // Single-stage pipeline version\n float Csub = 0.0;\n\n const int t4x = threadIdx.x * 4;\n const auto shape4 = cuda::aligned_size_t(sizeof(float4));\n cuda::pipeline pipe = cuda::make_pipeline();\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Load the matrices from device memory to shared memory;\n // a subset of threads loads a contiguous chunk of elements.\n\n // Previously, per-thread:\n // As[ty][tx] = A[a + wA * ty + tx];\n // Bs[ty][tx] = B[b + wB * ty + tx];\n\n // Now, one fourth of the threads load four elements of each matrix\n if (t4x < BLOCK_SIZE) {\n pipe.producer_acquire();\n\n cuda::memcpy_async(&As[threadIdx.y][t4x], &A[a + wA * threadIdx.y + t4x], shape4, pipe);\n cuda::memcpy_async(&Bs[threadIdx.y][t4x], &B[a + wA * threadIdx.y + t4x], shape4, pipe);\n\n pipe.producer_commit();\n pipe.consumer_wait();\n }\n\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[threadIdx.y][k] * Bs[k][threadIdx.x];\n }\n\n pipe.consumer_release();\n\n // Synchronize to make sure that the preceding\n // computation is done before overwriting the\n // shared memory sub-matrix buffers As and Bs in the next iteration.\n __syncthreads();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes four element\n int c = wB * BLOCK_SIZE * blockIdx.y + BLOCK_SIZE * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n}\n\n// Single Stage memcpy_async pipeline with Large copy chunk (float4) using\n// arrive-wait barrier\ntemplate \n__global__ void MatrixMulAsyncCopyLargeChunkAWBarrier(float *__restrict__ C,\n const float *__restrict__ A,\n const float *__restrict__ B,\n int wA,\n int wB)\n{\n#if __CUDA_ARCH__ >= 700\n#pragma diag_suppress static_var_with_dynamic_init\n // Requires BLOCK_SIZE % 4 == 0\n\n __shared__ cuda::barrier bar;\n\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ alignas(alignof(float4)) float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ alignas(alignof(float4)) float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n if (threadIdx.x == 0) {\n init(&bar, blockDim.x * blockDim.y);\n }\n __syncthreads();\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = wA * BLOCK_SIZE * blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n float Csub = 0.0;\n\n const int t4x = threadIdx.x * 4;\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Load the matrices from device memory to shared memory;\n // a subset of threads loads a contiguous chunk of elements.\n\n // Now, one fourth of the threads load four elements of each matrix\n if (t4x < BLOCK_SIZE) {\n float4 *const A4s = reinterpret_cast(&As[threadIdx.y][t4x]);\n float4 *const B4s = reinterpret_cast(&Bs[threadIdx.y][t4x]);\n const float4 *const A4 = reinterpret_cast(&A[a + wA * threadIdx.y + t4x]);\n const float4 *const B4 = reinterpret_cast(&B[a + wA * threadIdx.y + t4x]);\n\n cuda::memcpy_async(A4s, A4, sizeof(float4), bar);\n cuda::memcpy_async(B4s, B4, sizeof(float4), bar);\n }\n\n // Synchronize to make sure the matrices are loaded\n bar.arrive_and_wait();\n\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[threadIdx.y][k] * Bs[k][threadIdx.x];\n }\n\n // Synchronize to make sure that the preceding\n // computation is done before overwriting the\n // shared memory sub-matrix buffers As and Bs in the next iteration.\n bar.arrive_and_wait();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes four element\n int c = wB * BLOCK_SIZE * blockIdx.y + BLOCK_SIZE * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n#endif\n}\n\n// Single Stage memcpy_async pipeline with float copy\ntemplate \n__global__ void MatrixMulAsyncCopySingleStage(float *C, const float *A, const float *B, int wA, int wB)\n{\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = wA * BLOCK_SIZE * blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n // Single-stage pipeline version\n float Csub = 0.0;\n\n cuda::pipeline pipe = cuda::make_pipeline();\n const auto shape1 = cuda::aligned_size_t(sizeof(float));\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Load the matrices from device memory to shared memory; each thread loads\n // one element of each matrix\n {\n pipe.producer_acquire();\n\n cuda::memcpy_async(&As[threadIdx.y][threadIdx.x], &A[a + wA * threadIdx.y + threadIdx.x], shape1, pipe);\n cuda::memcpy_async(&Bs[threadIdx.y][threadIdx.x], &B[b + wB * threadIdx.y + threadIdx.x], shape1, pipe);\n\n pipe.producer_commit();\n }\n\n pipe.consumer_wait();\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[threadIdx.y][k] * Bs[k][threadIdx.x];\n }\n\n // Synchronize to make sure that the preceding\n // computation is done before overwriting the\n // shared memory sub-matrix buffers As and Bs in the next iteration.\n __syncthreads();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes four element\n int c = wB * BLOCK_SIZE * blockIdx.y + BLOCK_SIZE * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n}\n\n// Multi Stage memcpy_async thread_scope_thread pipeline with single-element\n// async-copy\ntemplate \n__global__ void MatrixMulAsyncCopyMultiStage(float *__restrict__ C,\n const float *__restrict__ A,\n const float *__restrict__ B,\n int wA,\n int wB)\n{\n // Multi-stage pipeline version\n constexpr size_t maxPipelineStages = 4;\n\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A for each stage\n __shared__ float As[maxPipelineStages][BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B for each stage\n __shared__ float Bs[maxPipelineStages][BLOCK_SIZE][BLOCK_SIZE];\n\n float Csub = 0.0;\n\n // Index of the first sub-matrix of A processed by the block\n const int aBegin = wA * BLOCK_SIZE * blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n const int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n const int bBegin = BLOCK_SIZE * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n cuda::pipeline pipe = cuda::make_pipeline();\n const auto shape1 = cuda::aligned_size_t(sizeof(float));\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin, i = 0, aStage = aBegin, bStage = bBegin, iStage = 0; a <= aEnd;\n a += aStep, b += bStep, ++i) {\n // Load the matrices from device memory to shared memory; each thread loads\n // one element of each matrix\n\n for (; aStage <= a + aStep * maxPipelineStages; aStage += aStep, bStage += bStep, ++iStage) {\n if (aStage <= aEnd) {\n // Rotating buffer\n const int j = iStage % maxPipelineStages;\n\n pipe.producer_acquire();\n\n cuda::memcpy_async(\n &As[j][threadIdx.y][threadIdx.x], &A[aStage + wA * threadIdx.y + threadIdx.x], shape1, pipe);\n cuda::memcpy_async(\n &Bs[j][threadIdx.y][threadIdx.x], &B[bStage + wB * threadIdx.y + threadIdx.x], shape1, pipe);\n\n pipe.producer_commit();\n }\n }\n pipe.consumer_wait();\n\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n const int j = i % maxPipelineStages;\n\n // Multiply the two matrices together;\n // each thread computes one element\n // of the block sub-matrix\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[j][threadIdx.y][k] * Bs[j][k][threadIdx.x];\n }\n\n pipe.consumer_release();\n // Don't have to synchronize because maxPipelineStages is greater than one\n // therefore next iteration is loading to a different buffer.\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes four element\n int c = wB * BLOCK_SIZE * blockIdx.y + BLOCK_SIZE * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n}\n\n// Multi Stage shared state memcpy_async pipeline thread_scope_block\n// with parititioned producer & consumer, here we've 1 warp as producer\n// group which issues memcpy_async operations and rest all warps are part of\n// consumer group which perform gemm computation on the loaded matrices by\n// producer.\ntemplate \n__global__ void MatrixMulAsyncCopyMultiStageSharedState(float *__restrict__ C,\n const float *__restrict__ A,\n const float *__restrict__ B,\n int wA,\n int wB)\n{\n // Multi-stage pipeline version\n constexpr size_t maxPipelineStages = 4;\n\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A for each stage\n __shared__ float As[maxPipelineStages][BLOCK_SIZE_X][BLOCK_SIZE_X];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B for each stage\n __shared__ float Bs[maxPipelineStages][BLOCK_SIZE_X][BLOCK_SIZE_X];\n\n float Csub = 0.0;\n\n // Index of the first sub-matrix of A processed by the block\n const int aBegin = wA * BLOCK_SIZE_X * blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n const int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n constexpr int aStep = BLOCK_SIZE_X;\n\n // Index of the first sub-matrix of B processed by the block\n const int bBegin = BLOCK_SIZE_X * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE_X * wB;\n\n auto cta = cg::this_thread_block();\n\n const auto shape1 = cuda::aligned_size_t(sizeof(float));\n __shared__ cuda::pipeline_shared_state shared_state;\n constexpr int consumer_row_count = BLOCK_SIZE_X;\n\n const auto thread_role =\n (cta.thread_index().y < consumer_row_count) ? cuda::pipeline_role::consumer : cuda::pipeline_role::producer;\n auto pipe = cuda::make_pipeline(cta, &shared_state, thread_role);\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin, i = 0, aStage = aBegin, bStage = bBegin, iStage = 0; a <= aEnd;\n a += aStep, b += bStep, ++i) {\n if (threadIdx.y >= consumer_row_count) {\n // this is a whole producer warp because threadIdx.y >= 16 where 16 ==\n // consumer_row_count,\n // which loads the matrices from device memory to shared memory;\n for (; aStage <= a + aStep * maxPipelineStages; aStage += aStep, bStage += bStep, ++iStage) {\n if (aStage <= aEnd) {\n // Rotating buffer\n const int j = iStage % maxPipelineStages;\n const int strideRows = (blockDim.y - consumer_row_count);\n pipe.producer_acquire();\n for (int rowId = threadIdx.y - consumer_row_count; rowId < BLOCK_SIZE_X; rowId += strideRows) {\n cuda::memcpy_async(\n &As[j][rowId][threadIdx.x], &A[aStage + wA * rowId + threadIdx.x], shape1, pipe);\n cuda::memcpy_async(\n &Bs[j][rowId][threadIdx.x], &B[bStage + wB * rowId + threadIdx.x], shape1, pipe);\n }\n pipe.producer_commit();\n }\n }\n }\n else {\n // this is a whole set of consumer group because threadIdx.y <\n // consumer_row_count where consumer_row_count == 16,\n // which computes gemm operation on matrices loaded in shared memory by\n // producer warp.\n const int j = i % maxPipelineStages;\n // Synchronize consumer group to make sure the matrices are loaded by\n // producer group.\n pipe.consumer_wait();\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE_X; ++k) {\n Csub += As[j][threadIdx.y][k] * Bs[j][k][threadIdx.x];\n }\n pipe.consumer_release();\n }\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes four element\n if (threadIdx.y < consumer_row_count) {\n const int c = wB * BLOCK_SIZE_X * blockIdx.y + BLOCK_SIZE_X * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n }\n}\n\n/**\n * Matrix multiplication (CUDA Kernel) on the device: C = A * B\n * wA is A's width and wB is B's width\n */\ntemplate __global__ void MatrixMulNaive(float *C, float *A, float *B, int wA, int wB)\n{\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = wA * BLOCK_SIZE * blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n // Csub is used to store the element of the block sub-matrix\n // that is computed by the thread\n float Csub = 0;\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Load the matrices from device memory\n // to shared memory; each thread loads\n // one element of each matrix\n As[threadIdx.y][threadIdx.x] = A[a + wA * threadIdx.y + threadIdx.x];\n Bs[threadIdx.y][threadIdx.x] = B[b + wB * threadIdx.y + threadIdx.x];\n\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[threadIdx.y][k] * Bs[k][threadIdx.x];\n }\n\n // Synchronize to make sure that the preceding\n // computation is done before loading two new\n // sub-matrices of A and B in the next iteration\n __syncthreads();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes one element\n int c = wB * BLOCK_SIZE * blockIdx.y + BLOCK_SIZE * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n}\n\ntemplate __global__ void MatrixMulNaiveLargeChunk(float *C, float *A, float *B, int wA, int wB)\n{\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ alignas(alignof(float4)) float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ alignas(alignof(float4)) float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n int t4x = threadIdx.x * 4;\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = wA * BLOCK_SIZE * blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n // Csub is used to store the element of the block sub-matrix\n // that is computed by the thread\n float Csub = 0;\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Load the matrices from device memory\n // to shared memory;\n\n // One fourth of the threads load four elements of each matrix\n if (t4x < BLOCK_SIZE) {\n float4 *const A4s = reinterpret_cast(&As[threadIdx.y][t4x]);\n float4 *const B4s = reinterpret_cast(&Bs[threadIdx.y][t4x]);\n const float4 *const A4 = reinterpret_cast(&A[a + wA * threadIdx.y + t4x]);\n const float4 *const B4 = reinterpret_cast(&B[a + wA * threadIdx.y + t4x]);\n *A4s = *A4;\n *B4s = *B4;\n }\n\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[threadIdx.y][k] * Bs[k][threadIdx.x];\n }\n\n // Synchronize to make sure that the preceding\n // computation is done before loading two new\n // sub-matrices of A and B in the next iteration\n __syncthreads();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes one element\n int c = wB * BLOCK_SIZE * blockIdx.y + BLOCK_SIZE * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n}\n\nvoid ConstantInit(float *data, int size, float val)\n{\n for (int i = 0; i < size; ++i) {\n data[i] = val;\n }\n}\n\n/**\n * Run matrix multiplication using CUDA\n */\nint MatrixMultiply(int argc, char **argv, const dim3 &dimsA, const dim3 &dimsB, kernels kernel_number)\n{\n // Allocate host memory for matrices A and B\n unsigned int size_A = dimsA.x * dimsA.y;\n unsigned int mem_size_A = sizeof(float) * size_A;\n float *h_A;\n checkCudaErrors(cudaMallocHost(&h_A, mem_size_A));\n unsigned int size_B = dimsB.x * dimsB.y;\n unsigned int mem_size_B = sizeof(float) * size_B;\n float *h_B;\n checkCudaErrors(cudaMallocHost(&h_B, mem_size_B));\n cudaStream_t stream;\n\n // Initialize host memory\n const float valB = 2.10f;\n ConstantInit(h_A, size_A, 1.0f);\n ConstantInit(h_B, size_B, valB);\n\n // Allocate device memory\n float *d_A, *d_B, *d_C;\n\n // Allocate host matrix C\n dim3 dimsC(dimsB.x, dimsA.y, 1);\n unsigned int mem_size_C = dimsC.x * dimsC.y * sizeof(float);\n float *h_C;\n checkCudaErrors(cudaMallocHost(&h_C, mem_size_C));\n\n if (h_C == NULL) {\n fprintf(stderr, \"Failed to allocate host matrix C!\\n\");\n exit(EXIT_FAILURE);\n }\n\n checkCudaErrors(cudaMalloc(reinterpret_cast(&d_A), mem_size_A));\n checkCudaErrors(cudaMalloc(reinterpret_cast(&d_B), mem_size_B));\n checkCudaErrors(cudaMalloc(reinterpret_cast(&d_C), mem_size_C));\n // Allocate CUDA events that we'll use for timing\n cudaEvent_t start, stop;\n checkCudaErrors(cudaEventCreate(&start));\n checkCudaErrors(cudaEventCreate(&stop));\n\n checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));\n\n // copy host memory to device\n checkCudaErrors(cudaMemcpyAsync(d_A, h_A, mem_size_A, cudaMemcpyHostToDevice, stream));\n checkCudaErrors(cudaMemcpyAsync(d_B, h_B, mem_size_B, cudaMemcpyHostToDevice, stream));\n checkCudaErrors(cudaMemsetAsync(d_C, 0, mem_size_C, stream));\n\n // Setup execution parameters\n dim3 threads(blockSize, blockSize);\n dim3 grid(dimsB.x / threads.x, dimsA.y / threads.y);\n\n // Here the block size is 16x18, where first 16 rows are consumer thread group\n // and last 2 rows (1 warp) is producer thread group\n dim3 threadsSharedStateKernel(blockSize, blockSize + 2, 1);\n dim3 gridSharedStateKernel(dimsB.x / threadsSharedStateKernel.x, dimsA.y / threadsSharedStateKernel.x);\n\n printf(\"Running kernel = %d - %s\\n\", kernel_number, kernelNames[kernel_number]);\n // Create and start timer\n printf(\"Computing result using CUDA Kernel...\\n\");\n\n // Performs warmup operation using matrixMul CUDA kernel\n switch (kernel_number) {\n case AsyncCopyMultiStageLargeChunk:\n default:\n MatrixMulAsyncCopyMultiStageLargeChunk\n <<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyLargeChunk:\n MatrixMulAsyncCopyLargeChunk<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyLargeChunkAWBarrier:\n MatrixMulAsyncCopyLargeChunkAWBarrier<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyMultiStageSharedState:\n MatrixMulAsyncCopyMultiStageSharedState\n <<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyMultiStage:\n MatrixMulAsyncCopyMultiStage<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopySingleStage:\n MatrixMulAsyncCopySingleStage<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case Naive:\n MatrixMulNaive<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case NaiveLargeChunk:\n MatrixMulNaiveLargeChunk<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n }\n\n printf(\"done\\n\");\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n // Execute the kernel\n int nIter = 100;\n\n // Record the start event\n checkCudaErrors(cudaEventRecord(start, stream));\n\n for (int j = 0; j < nIter; j++) {\n switch (kernel_number) {\n case AsyncCopyMultiStageLargeChunk:\n default:\n MatrixMulAsyncCopyMultiStageLargeChunk\n <<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyLargeChunk:\n MatrixMulAsyncCopyLargeChunk<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyLargeChunkAWBarrier:\n MatrixMulAsyncCopyLargeChunkAWBarrier\n <<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyMultiStageSharedState:\n MatrixMulAsyncCopyMultiStageSharedState\n <<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyMultiStage:\n MatrixMulAsyncCopyMultiStage<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopySingleStage:\n MatrixMulAsyncCopySingleStage<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case Naive:\n MatrixMulNaive<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case NaiveLargeChunk:\n MatrixMulNaiveLargeChunk<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n }\n }\n\n // Record the stop event\n checkCudaErrors(cudaEventRecord(stop, stream));\n\n // Wait for the stop event to complete\n checkCudaErrors(cudaEventSynchronize(stop));\n\n float msecTotal = 0.0f;\n checkCudaErrors(cudaEventElapsedTime(&msecTotal, start, stop));\n\n // Compute and print the performance\n float msecPerMatrixMul = msecTotal / nIter;\n double flopsPerMatrixMul =\n 2.0 * static_cast(dimsA.x) * static_cast(dimsA.y) * static_cast(dimsB.x);\n double gigaFlops = (flopsPerMatrixMul * 1.0e-9f) / (msecPerMatrixMul / 1000.0f);\n printf(\"Performance= %.2f GFlop/s, Time= %.3f msec, Size= %.0f Ops,\"\n \" WorkgroupSize= %u threads/block\\n\",\n gigaFlops,\n msecPerMatrixMul,\n flopsPerMatrixMul,\n threads.x * threads.y);\n\n // Copy result from device to host\n checkCudaErrors(cudaMemcpyAsync(h_C, d_C, mem_size_C, cudaMemcpyDeviceToHost, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n printf(\"Checking computed result for correctness: \");\n bool correct = true;\n\n // test relative error by the formula\n // |_cpu - _gpu|/<|x|, |y|> < eps\n double eps = 1.e-6; // machine zero\n\n for (int i = 0; i < static_cast(dimsC.x * dimsC.y); i++) {\n double abs_err = fabs(h_C[i] - (dimsA.x * valB));\n double dot_length = dimsA.x;\n double abs_val = fabs(h_C[i]);\n double rel_err = abs_err / abs_val / dot_length;\n\n if (rel_err > eps) {\n printf(\"Error! Matrix[%05d]=%.8f, ref=%.8f error term is > %E\\n\", i, h_C[i], dimsA.x * valB, eps);\n correct = false;\n }\n }\n\n printf(\"%s\\n\", correct ? \"Result = PASS\" : \"Result = FAIL\");\n\n // Clean up memory\n checkCudaErrors(cudaFreeHost(h_A));\n checkCudaErrors(cudaFreeHost(h_B));\n checkCudaErrors(cudaFreeHost(h_C));\n checkCudaErrors(cudaFree(d_A));\n checkCudaErrors(cudaFree(d_B));\n checkCudaErrors(cudaFree(d_C));\n checkCudaErrors(cudaEventDestroy(start));\n checkCudaErrors(cudaEventDestroy(stop));\n printf(\"\\nNOTE: The CUDA Samples are not meant for performance \"\n \"measurements. Results may vary when GPU Boost is enabled.\\n\");\n\n if (correct) {\n return EXIT_SUCCESS;\n }\n else {\n return EXIT_FAILURE;\n }\n}\n\nint main(int argc, char **argv)\n{\n printf(\"[globalToShmemAsyncCopy] - Starting...\\n\");\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\") || checkCmdLineFlag(argc, (const char **)argv, \"?\")) {\n printf(\"Usage -device=n (n >= 0 for deviceID)\\n\");\n printf(\" -wA=WidthA -hA=HeightA (Width x Height of Matrix A)\\n\");\n printf(\" -wB=WidthB -hB=HeightB (Width x Height of Matrix B)\\n\");\n printf(\" -kernel=kernel_number (0 - AsyncCopyMultiStageLargeChunk; 1 - \"\n \"AsyncCopyLargeChunk)\\n\");\n printf(\" (2 - AsyncCopyLargeChunkAWBarrier; 3 - \"\n \"AsyncCopyMultiStageSharedState)\\n\");\n printf(\" (4 - AsyncCopyMultiStage; 5 - \"\n \"AsyncCopySingleStage; 6 - Naive without memcpy_async)\\n\");\n printf(\" (7 - NaiveLargeChunk without \"\n \"memcpy_async)\\n\");\n printf(\" Note: Outer matrix dimensions of A & B matrices must be equal.\\n\");\n\n exit(EXIT_SUCCESS);\n }\n\n // This will pick the best possible CUDA capable device, otherwise\n // override the device ID based on input provided at the command line\n int dev = findCudaDevice(argc, (const char **)argv);\n\n int matrixBlock = 32;\n dim3 dimsA(10 * 4 * matrixBlock, 10 * 4 * matrixBlock, 1);\n dim3 dimsB(10 * 4 * matrixBlock, 10 * 4 * matrixBlock, 1);\n\n // width of Matrix A\n if (checkCmdLineFlag(argc, (const char **)argv, \"wA\")) {\n dimsA.x = getCmdLineArgumentInt(argc, (const char **)argv, \"wA\");\n }\n\n // height of Matrix A\n if (checkCmdLineFlag(argc, (const char **)argv, \"hA\")) {\n dimsA.y = getCmdLineArgumentInt(argc, (const char **)argv, \"hA\");\n }\n\n // width of Matrix B\n if (checkCmdLineFlag(argc, (const char **)argv, \"wB\")) {\n dimsB.x = getCmdLineArgumentInt(argc, (const char **)argv, \"wB\");\n }\n\n // height of Matrix B\n if (checkCmdLineFlag(argc, (const char **)argv, \"hB\")) {\n dimsB.y = getCmdLineArgumentInt(argc, (const char **)argv, \"hB\");\n }\n\n if (dimsA.x != dimsB.y) {\n printf(\"Error: outer matrix dimensions must be equal. (%d != %d)\\n\", dimsA.x, dimsB.y);\n exit(EXIT_FAILURE);\n }\n\n kernels selected_kernel = AsyncCopyMultiStageLargeChunk;\n\n // kernel to run - default (AsyncCopyMultiStageLargeChunk == 0)\n if (checkCmdLineFlag(argc, (const char **)argv, \"kernel\")) {\n int kernel_number = getCmdLineArgumentInt(argc, (const char **)argv, \"kernel\");\n if (kernel_number < 8) {\n selected_kernel = (kernels)kernel_number;\n }\n else {\n printf(\"Error: kernel number should be between 0 to 6, you have entered \"\n \"%d\\n\",\n kernel_number);\n exit(EXIT_FAILURE);\n }\n }\n\n int major = 0;\n checkCudaErrors(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev));\n if (major < 7) {\n printf(\"globalToShmemAsyncCopy requires SM 7.0 or higher. Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n printf(\"MatrixA(%d,%d), MatrixB(%d,%d)\\n\", dimsA.x, dimsA.y, dimsB.x, dimsB.y);\n\n int matrix_result = MatrixMultiply(argc, argv, dimsA, dimsB, selected_kernel);\n\n exit(matrix_result);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/**\n * Matrix multiplication: C = A * B.\n *\n * This sample demonstrates implements matrix multiplication which makes use of\n * shared memory to ensure data reuse, the matrix multiplication is done using\n * tiling approach.\n * With compute capability 8.0 or higher the CUDA kernels involved uses\n * asynchronously copy data from global to shared memory; a.k.a., async-copy.\n * This sample has been written for clarity of exposition to illustrate various\n * CUDA programming principles, not with the goal of providing the most\n * performant generic kernel for matrix multiplication.\n */\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n#include \n\n#if __CUDA_ARCH__ >= 700\n#include \n#endif\n#include \n\nnamespace cg = cooperative_groups;\n\n// Helper functions and utilities to work with CUDA\n#include \n#include \n\nenum kernels {\n AsyncCopyMultiStageLargeChunk = 0,\n AsyncCopyLargeChunk = 1,\n AsyncCopyLargeChunkAWBarrier = 2,\n AsyncCopyMultiStageSharedState = 3,\n AsyncCopyMultiStage = 4,\n AsyncCopySingleStage = 5,\n Naive = 6,\n NaiveLargeChunk = 7\n};\n\nconst char *kernelNames[] = {\"AsyncCopyMultiStageLargeChunk\",\n \"AsyncCopyLargeChunk\",\n \"AsyncCopyLargeChunkAWBarrier\",\n \"AsyncCopyMultiStageSharedState\",\n \"AsyncCopyMultiStage\",\n \"AsyncCopySingleStage\",\n \"Naive\",\n \"NaiveLargeChunk\"};\n\nconstexpr int blockSize = 16;\n\n// Multi Stage memcpy_async pipeline with large chunk copy\ntemplate \n__global__ void MatrixMulAsyncCopyMultiStageLargeChunk(float *__restrict__ C,\n const float *__restrict__ A,\n const float *__restrict__ B,\n int wA,\n int wB)\n{\n // Requires BLOCK_SIZE % 4 == 0\n\n // Multi-stage pipeline version\n constexpr size_t maxPipelineStages = 4;\n\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A for each stage\n __shared__ alignas(alignof(float4)) float As[maxPipelineStages][BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B for each stage\n __shared__ alignas(alignof(float4)) float Bs[maxPipelineStages][BLOCK_SIZE][BLOCK_SIZE];\n\n float Csub = 0.0;\n\n // Index of the first sub-matrix of A processed by the block\n const int aBegin = wA * (BLOCK_SIZE)*blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n const int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n const int bBegin = BLOCK_SIZE * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n const int t4x = threadIdx.x * 4;\n const auto shape4 = cuda::aligned_size_t(sizeof(float4));\n\n cuda::pipeline pipe = cuda::make_pipeline();\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin, i = 0, aStage = aBegin, bStage = bBegin, iStage = 0; a <= aEnd;\n a += aStep, b += bStep, ++i) {\n // Load the matrices from device memory to shared memory; each thread loads\n // one element of each matrix\n for (; aStage <= a + aStep * maxPipelineStages; aStage += aStep, bStage += bStep, ++iStage) {\n pipe.producer_acquire();\n if (aStage <= aEnd && t4x < BLOCK_SIZE) {\n // Rotating buffer\n const int j = iStage % maxPipelineStages;\n cuda::memcpy_async(&As[j][threadIdx.y][t4x], &A[aStage + wA * threadIdx.y + t4x], shape4, pipe);\n cuda::memcpy_async(&Bs[j][threadIdx.y][t4x], &B[aStage + wA * threadIdx.y + t4x], shape4, pipe);\n }\n pipe.producer_commit();\n }\n\n pipe.consumer_wait();\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n // Rotating buffer\n const int j = i % maxPipelineStages;\n\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[j][threadIdx.y][k] * Bs[j][k][threadIdx.x];\n }\n pipe.consumer_release();\n\n // Don't have to synchronize because maxPipelineStages is greater than one\n // therefore next iteration is loading to a different buffer.\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes four element\n int c = wB * BLOCK_SIZE * blockIdx.y + BLOCK_SIZE * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n}\n\n// Single Stage memcpy_async pipeline with Large copy chunk (float4)\ntemplate \n__global__ void MatrixMulAsyncCopyLargeChunk(float *__restrict__ C,\n const float *__restrict__ A,\n const float *__restrict__ B,\n int wA,\n int wB)\n{\n // Requires BLOCK_SIZE % 4 == 0\n\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ alignas(alignof(float4)) float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ alignas(alignof(float4)) float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = wA * BLOCK_SIZE * blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n // Single-stage pipeline version\n float Csub = 0.0;\n\n const int t4x = threadIdx.x * 4;\n const auto shape4 = cuda::aligned_size_t(sizeof(float4));\n cuda::pipeline pipe = cuda::make_pipeline();\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Load the matrices from device memory to shared memory;\n // a subset of threads loads a contiguous chunk of elements.\n\n // Previously, per-thread:\n // As[ty][tx] = A[a + wA * ty + tx];\n // Bs[ty][tx] = B[b + wB * ty + tx];\n\n // Now, one fourth of the threads load four elements of each matrix\n if (t4x < BLOCK_SIZE) {\n pipe.producer_acquire();\n\n cuda::memcpy_async(&As[threadIdx.y][t4x], &A[a + wA * threadIdx.y + t4x], shape4, pipe);\n cuda::memcpy_async(&Bs[threadIdx.y][t4x], &B[a + wA * threadIdx.y + t4x], shape4, pipe);\n\n pipe.producer_commit();\n pipe.consumer_wait();\n }\n\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[threadIdx.y][k] * Bs[k][threadIdx.x];\n }\n\n pipe.consumer_release();\n\n // Synchronize to make sure that the preceding\n // computation is done before overwriting the\n // shared memory sub-matrix buffers As and Bs in the next iteration.\n __syncthreads();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes four element\n int c = wB * BLOCK_SIZE * blockIdx.y + BLOCK_SIZE * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n}\n\n// Single Stage memcpy_async pipeline with Large copy chunk (float4) using\n// arrive-wait barrier\ntemplate \n__global__ void MatrixMulAsyncCopyLargeChunkAWBarrier(float *__restrict__ C,\n const float *__restrict__ A,\n const float *__restrict__ B,\n int wA,\n int wB)\n{\n#if __CUDA_ARCH__ >= 700\n#pragma diag_suppress static_var_with_dynamic_init\n // Requires BLOCK_SIZE % 4 == 0\n\n __shared__ cuda::barrier bar;\n\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ alignas(alignof(float4)) float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ alignas(alignof(float4)) float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n if (threadIdx.x == 0) {\n init(&bar, blockDim.x * blockDim.y);\n }\n __syncthreads();\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = wA * BLOCK_SIZE * blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n float Csub = 0.0;\n\n const int t4x = threadIdx.x * 4;\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Load the matrices from device memory to shared memory;\n // a subset of threads loads a contiguous chunk of elements.\n\n // Now, one fourth of the threads load four elements of each matrix\n if (t4x < BLOCK_SIZE) {\n float4 *const A4s = reinterpret_cast(&As[threadIdx.y][t4x]);\n float4 *const B4s = reinterpret_cast(&Bs[threadIdx.y][t4x]);\n const float4 *const A4 = reinterpret_cast(&A[a + wA * threadIdx.y + t4x]);\n const float4 *const B4 = reinterpret_cast(&B[a + wA * threadIdx.y + t4x]);\n\n cuda::memcpy_async(A4s, A4, sizeof(float4), bar);\n cuda::memcpy_async(B4s, B4, sizeof(float4), bar);\n }\n\n // Synchronize to make sure the matrices are loaded\n bar.arrive_and_wait();\n\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[threadIdx.y][k] * Bs[k][threadIdx.x];\n }\n\n // Synchronize to make sure that the preceding\n // computation is done before overwriting the\n // shared memory sub-matrix buffers As and Bs in the next iteration.\n bar.arrive_and_wait();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes four element\n int c = wB * BLOCK_SIZE * blockIdx.y + BLOCK_SIZE * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n#endif\n}\n\n// Single Stage memcpy_async pipeline with float copy\ntemplate \n__global__ void MatrixMulAsyncCopySingleStage(float *C, const float *A, const float *B, int wA, int wB)\n{\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = wA * BLOCK_SIZE * blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n // Single-stage pipeline version\n float Csub = 0.0;\n\n cuda::pipeline pipe = cuda::make_pipeline();\n const auto shape1 = cuda::aligned_size_t(sizeof(float));\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Load the matrices from device memory to shared memory; each thread loads\n // one element of each matrix\n {\n pipe.producer_acquire();\n\n cuda::memcpy_async(&As[threadIdx.y][threadIdx.x], &A[a + wA * threadIdx.y + threadIdx.x], shape1, pipe);\n cuda::memcpy_async(&Bs[threadIdx.y][threadIdx.x], &B[b + wB * threadIdx.y + threadIdx.x], shape1, pipe);\n\n pipe.producer_commit();\n }\n\n pipe.consumer_wait();\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[threadIdx.y][k] * Bs[k][threadIdx.x];\n }\n\n // Synchronize to make sure that the preceding\n // computation is done before overwriting the\n // shared memory sub-matrix buffers As and Bs in the next iteration.\n __syncthreads();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes four element\n int c = wB * BLOCK_SIZE * blockIdx.y + BLOCK_SIZE * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n}\n\n// Multi Stage memcpy_async thread_scope_thread pipeline with single-element\n// async-copy\ntemplate \n__global__ void MatrixMulAsyncCopyMultiStage(float *__restrict__ C,\n const float *__restrict__ A,\n const float *__restrict__ B,\n int wA,\n int wB)\n{\n // Multi-stage pipeline version\n constexpr size_t maxPipelineStages = 4;\n\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A for each stage\n __shared__ float As[maxPipelineStages][BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B for each stage\n __shared__ float Bs[maxPipelineStages][BLOCK_SIZE][BLOCK_SIZE];\n\n float Csub = 0.0;\n\n // Index of the first sub-matrix of A processed by the block\n const int aBegin = wA * BLOCK_SIZE * blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n const int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n const int bBegin = BLOCK_SIZE * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n cuda::pipeline pipe = cuda::make_pipeline();\n const auto shape1 = cuda::aligned_size_t(sizeof(float));\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin, i = 0, aStage = aBegin, bStage = bBegin, iStage = 0; a <= aEnd;\n a += aStep, b += bStep, ++i) {\n // Load the matrices from device memory to shared memory; each thread loads\n // one element of each matrix\n\n for (; aStage <= a + aStep * maxPipelineStages; aStage += aStep, bStage += bStep, ++iStage) {\n if (aStage <= aEnd) {\n // Rotating buffer\n const int j = iStage % maxPipelineStages;\n\n pipe.producer_acquire();\n\n cuda::memcpy_async(\n &As[j][threadIdx.y][threadIdx.x], &A[aStage + wA * threadIdx.y + threadIdx.x], shape1, pipe);\n cuda::memcpy_async(\n &Bs[j][threadIdx.y][threadIdx.x], &B[bStage + wB * threadIdx.y + threadIdx.x], shape1, pipe);\n\n pipe.producer_commit();\n }\n }\n pipe.consumer_wait();\n\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n const int j = i % maxPipelineStages;\n\n // Multiply the two matrices together;\n // each thread computes one element\n // of the block sub-matrix\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[j][threadIdx.y][k] * Bs[j][k][threadIdx.x];\n }\n\n pipe.consumer_release();\n // Don't have to synchronize because maxPipelineStages is greater than one\n // therefore next iteration is loading to a different buffer.\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes four element\n int c = wB * BLOCK_SIZE * blockIdx.y + BLOCK_SIZE * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n}\n\n// Multi Stage shared state memcpy_async pipeline thread_scope_block\n// with parititioned producer & consumer, here we've 1 warp as producer\n// group which issues memcpy_async operations and rest all warps are part of\n// consumer group which perform gemm computation on the loaded matrices by\n// producer.\ntemplate \n__global__ void MatrixMulAsyncCopyMultiStageSharedState(float *__restrict__ C,\n const float *__restrict__ A,\n const float *__restrict__ B,\n int wA,\n int wB)\n{\n // Multi-stage pipeline version\n constexpr size_t maxPipelineStages = 4;\n\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A for each stage\n __shared__ float As[maxPipelineStages][BLOCK_SIZE_X][BLOCK_SIZE_X];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B for each stage\n __shared__ float Bs[maxPipelineStages][BLOCK_SIZE_X][BLOCK_SIZE_X];\n\n float Csub = 0.0;\n\n // Index of the first sub-matrix of A processed by the block\n const int aBegin = wA * BLOCK_SIZE_X * blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n const int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n constexpr int aStep = BLOCK_SIZE_X;\n\n // Index of the first sub-matrix of B processed by the block\n const int bBegin = BLOCK_SIZE_X * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE_X * wB;\n\n auto cta = cg::this_thread_block();\n\n const auto shape1 = cuda::aligned_size_t(sizeof(float));\n __shared__ cuda::pipeline_shared_state shared_state;\n constexpr int consumer_row_count = BLOCK_SIZE_X;\n\n const auto thread_role =\n (cta.thread_index().y < consumer_row_count) ? cuda::pipeline_role::consumer : cuda::pipeline_role::producer;\n auto pipe = cuda::make_pipeline(cta, &shared_state, thread_role);\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin, i = 0, aStage = aBegin, bStage = bBegin, iStage = 0; a <= aEnd;\n a += aStep, b += bStep, ++i) {\n if (threadIdx.y >= consumer_row_count) {\n // this is a whole producer warp because threadIdx.y >= 16 where 16 ==\n // consumer_row_count,\n // which loads the matrices from device memory to shared memory;\n for (; aStage <= a + aStep * maxPipelineStages; aStage += aStep, bStage += bStep, ++iStage) {\n if (aStage <= aEnd) {\n // Rotating buffer\n const int j = iStage % maxPipelineStages;\n const int strideRows = (blockDim.y - consumer_row_count);\n pipe.producer_acquire();\n for (int rowId = threadIdx.y - consumer_row_count; rowId < BLOCK_SIZE_X; rowId += strideRows) {\n cuda::memcpy_async(\n &As[j][rowId][threadIdx.x], &A[aStage + wA * rowId + threadIdx.x], shape1, pipe);\n cuda::memcpy_async(\n &Bs[j][rowId][threadIdx.x], &B[bStage + wB * rowId + threadIdx.x], shape1, pipe);\n }\n pipe.producer_commit();\n }\n }\n }\n else {\n // this is a whole set of consumer group because threadIdx.y <\n // consumer_row_count where consumer_row_count == 16,\n // which computes gemm operation on matrices loaded in shared memory by\n // producer warp.\n const int j = i % maxPipelineStages;\n // Synchronize consumer group to make sure the matrices are loaded by\n // producer group.\n pipe.consumer_wait();\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE_X; ++k) {\n Csub += As[j][threadIdx.y][k] * Bs[j][k][threadIdx.x];\n }\n pipe.consumer_release();\n }\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes four element\n if (threadIdx.y < consumer_row_count) {\n const int c = wB * BLOCK_SIZE_X * blockIdx.y + BLOCK_SIZE_X * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n }\n}\n\n/**\n * Matrix multiplication (CUDA Kernel) on the device: C = A * B\n * wA is A's width and wB is B's width\n */\ntemplate __global__ void MatrixMulNaive(float *C, float *A, float *B, int wA, int wB)\n{\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = wA * BLOCK_SIZE * blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n // Csub is used to store the element of the block sub-matrix\n // that is computed by the thread\n float Csub = 0;\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Load the matrices from device memory\n // to shared memory; each thread loads\n // one element of each matrix\n As[threadIdx.y][threadIdx.x] = A[a + wA * threadIdx.y + threadIdx.x];\n Bs[threadIdx.y][threadIdx.x] = B[b + wB * threadIdx.y + threadIdx.x];\n\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[threadIdx.y][k] * Bs[k][threadIdx.x];\n }\n\n // Synchronize to make sure that the preceding\n // computation is done before loading two new\n // sub-matrices of A and B in the next iteration\n __syncthreads();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes one element\n int c = wB * BLOCK_SIZE * blockIdx.y + BLOCK_SIZE * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n}\n\ntemplate __global__ void MatrixMulNaiveLargeChunk(float *C, float *A, float *B, int wA, int wB)\n{\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ alignas(alignof(float4)) float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ alignas(alignof(float4)) float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n int t4x = threadIdx.x * 4;\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = wA * BLOCK_SIZE * blockIdx.y;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * blockIdx.x;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n // Csub is used to store the element of the block sub-matrix\n // that is computed by the thread\n float Csub = 0;\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Load the matrices from device memory\n // to shared memory;\n\n // One fourth of the threads load four elements of each matrix\n if (t4x < BLOCK_SIZE) {\n float4 *const A4s = reinterpret_cast(&As[threadIdx.y][t4x]);\n float4 *const B4s = reinterpret_cast(&Bs[threadIdx.y][t4x]);\n const float4 *const A4 = reinterpret_cast(&A[a + wA * threadIdx.y + t4x]);\n const float4 *const B4 = reinterpret_cast(&B[a + wA * threadIdx.y + t4x]);\n *A4s = *A4;\n *B4s = *B4;\n }\n\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n// Multiply the two matrices together;\n// each thread computes one element\n// of the block sub-matrix\n#pragma unroll\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[threadIdx.y][k] * Bs[k][threadIdx.x];\n }\n\n // Synchronize to make sure that the preceding\n // computation is done before loading two new\n // sub-matrices of A and B in the next iteration\n __syncthreads();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes one element\n int c = wB * BLOCK_SIZE * blockIdx.y + BLOCK_SIZE * blockIdx.x;\n C[c + wB * threadIdx.y + threadIdx.x] = Csub;\n}\n\nvoid ConstantInit(float *data, int size, float val)\n{\n for (int i = 0; i < size; ++i) {\n data[i] = val;\n }\n}\n\n/**\n * Run matrix multiplication using CUDA\n */\nint MatrixMultiply(int argc, char **argv, const dim3 &dimsA, const dim3 &dimsB, kernels kernel_number)\n{\n // Allocate host memory for matrices A and B\n unsigned int size_A = dimsA.x * dimsA.y;\n unsigned int mem_size_A = sizeof(float) * size_A;\n float *h_A;\n checkCudaErrors(hipHostMalloc(&h_A, mem_size_A));\n unsigned int size_B = dimsB.x * dimsB.y;\n unsigned int mem_size_B = sizeof(float) * size_B;\n float *h_B;\n checkCudaErrors(hipHostMalloc(&h_B, mem_size_B));\n hipStream_t stream;\n\n // Initialize host memory\n const float valB = 2.10f;\n ConstantInit(h_A, size_A, 1.0f);\n ConstantInit(h_B, size_B, valB);\n\n // Allocate device memory\n float *d_A, *d_B, *d_C;\n\n // Allocate host matrix C\n dim3 dimsC(dimsB.x, dimsA.y, 1);\n unsigned int mem_size_C = dimsC.x * dimsC.y * sizeof(float);\n float *h_C;\n checkCudaErrors(hipHostMalloc(&h_C, mem_size_C));\n\n if (h_C == NULL) {\n fprintf(stderr, \"Failed to allocate host matrix C!\\n\");\n exit(EXIT_FAILURE);\n }\n\n checkCudaErrors(hipMalloc(reinterpret_cast(&d_A), mem_size_A));\n checkCudaErrors(hipMalloc(reinterpret_cast(&d_B), mem_size_B));\n checkCudaErrors(hipMalloc(reinterpret_cast(&d_C), mem_size_C));\n // Allocate CUDA events that we'll use for timing\n hipEvent_t start, stop;\n checkCudaErrors(hipEventCreate(&start));\n checkCudaErrors(hipEventCreate(&stop));\n\n checkCudaErrors(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));\n\n // copy host memory to device\n checkCudaErrors(hipMemcpyAsync(d_A, h_A, mem_size_A, hipMemcpyHostToDevice, stream));\n checkCudaErrors(hipMemcpyAsync(d_B, h_B, mem_size_B, hipMemcpyHostToDevice, stream));\n checkCudaErrors(hipMemsetAsync(d_C, 0, mem_size_C, stream));\n\n // Setup execution parameters\n dim3 threads(blockSize, blockSize);\n dim3 grid(dimsB.x / threads.x, dimsA.y / threads.y);\n\n // Here the block size is 16x18, where first 16 rows are consumer thread group\n // and last 2 rows (1 warp) is producer thread group\n dim3 threadsSharedStateKernel(blockSize, blockSize + 2, 1);\n dim3 gridSharedStateKernel(dimsB.x / threadsSharedStateKernel.x, dimsA.y / threadsSharedStateKernel.x);\n\n printf(\"Running kernel = %d - %s\\n\", kernel_number, kernelNames[kernel_number]);\n // Create and start timer\n printf(\"Computing result using CUDA Kernel...\\n\");\n\n // Performs warmup operation using matrixMul CUDA kernel\n switch (kernel_number) {\n case AsyncCopyMultiStageLargeChunk:\n default:\n MatrixMulAsyncCopyMultiStageLargeChunk\n <<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyLargeChunk:\n MatrixMulAsyncCopyLargeChunk<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyLargeChunkAWBarrier:\n MatrixMulAsyncCopyLargeChunkAWBarrier<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyMultiStageSharedState:\n MatrixMulAsyncCopyMultiStageSharedState\n <<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyMultiStage:\n MatrixMulAsyncCopyMultiStage<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopySingleStage:\n MatrixMulAsyncCopySingleStage<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case Naive:\n MatrixMulNaive<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case NaiveLargeChunk:\n MatrixMulNaiveLargeChunk<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n }\n\n printf(\"done\\n\");\n checkCudaErrors(hipStreamSynchronize(stream));\n\n // Execute the kernel\n int nIter = 100;\n\n // Record the start event\n checkCudaErrors(hipEventRecord(start, stream));\n\n for (int j = 0; j < nIter; j++) {\n switch (kernel_number) {\n case AsyncCopyMultiStageLargeChunk:\n default:\n MatrixMulAsyncCopyMultiStageLargeChunk\n <<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyLargeChunk:\n MatrixMulAsyncCopyLargeChunk<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyLargeChunkAWBarrier:\n MatrixMulAsyncCopyLargeChunkAWBarrier\n <<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyMultiStageSharedState:\n MatrixMulAsyncCopyMultiStageSharedState\n <<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopyMultiStage:\n MatrixMulAsyncCopyMultiStage<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case AsyncCopySingleStage:\n MatrixMulAsyncCopySingleStage<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case Naive:\n MatrixMulNaive<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n case NaiveLargeChunk:\n MatrixMulNaiveLargeChunk<<>>(d_C, d_A, d_B, dimsA.x, dimsB.x);\n break;\n }\n }\n\n // Record the stop event\n checkCudaErrors(hipEventRecord(stop, stream));\n\n // Wait for the stop event to complete\n checkCudaErrors(hipEventSynchronize(stop));\n\n float msecTotal = 0.0f;\n checkCudaErrors(hipEventElapsedTime(&msecTotal, start, stop));\n\n // Compute and print the performance\n float msecPerMatrixMul = msecTotal / nIter;\n double flopsPerMatrixMul =\n 2.0 * static_cast(dimsA.x) * static_cast(dimsA.y) * static_cast(dimsB.x);\n double gigaFlops = (flopsPerMatrixMul * 1.0e-9f) / (msecPerMatrixMul / 1000.0f);\n printf(\"Performance= %.2f GFlop/s, Time= %.3f msec, Size= %.0f Ops,\"\n \" WorkgroupSize= %u threads/block\\n\",\n gigaFlops,\n msecPerMatrixMul,\n flopsPerMatrixMul,\n threads.x * threads.y);\n\n // Copy result from device to host\n checkCudaErrors(hipMemcpyAsync(h_C, d_C, mem_size_C, hipMemcpyDeviceToHost, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n\n printf(\"Checking computed result for correctness: \");\n bool correct = true;\n\n // test relative error by the formula\n // |_cpu - _gpu|/<|x|, |y|> < eps\n double eps = 1.e-6; // machine zero\n\n for (int i = 0; i < static_cast(dimsC.x * dimsC.y); i++) {\n double abs_err = fabs(h_C[i] - (dimsA.x * valB));\n double dot_length = dimsA.x;\n double abs_val = fabs(h_C[i]);\n double rel_err = abs_err / abs_val / dot_length;\n\n if (rel_err > eps) {\n printf(\"Error! Matrix[%05d]=%.8f, ref=%.8f error term is > %E\\n\", i, h_C[i], dimsA.x * valB, eps);\n correct = false;\n }\n }\n\n printf(\"%s\\n\", correct ? \"Result = PASS\" : \"Result = FAIL\");\n\n // Clean up memory\n checkCudaErrors(hipHostFree(h_A));\n checkCudaErrors(hipHostFree(h_B));\n checkCudaErrors(hipHostFree(h_C));\n checkCudaErrors(hipFree(d_A));\n checkCudaErrors(hipFree(d_B));\n checkCudaErrors(hipFree(d_C));\n checkCudaErrors(hipEventDestroy(start));\n checkCudaErrors(hipEventDestroy(stop));\n printf(\"\\nNOTE: The CUDA Samples are not meant for performance \"\n \"measurements. Results may vary when GPU Boost is enabled.\\n\");\n\n if (correct) {\n return EXIT_SUCCESS;\n }\n else {\n return EXIT_FAILURE;\n }\n}\n\nint main(int argc, char **argv)\n{\n printf(\"[globalToShmemAsyncCopy] - Starting...\\n\");\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\") || checkCmdLineFlag(argc, (const char **)argv, \"?\")) {\n printf(\"Usage -device=n (n >= 0 for deviceID)\\n\");\n printf(\" -wA=WidthA -hA=HeightA (Width x Height of Matrix A)\\n\");\n printf(\" -wB=WidthB -hB=HeightB (Width x Height of Matrix B)\\n\");\n printf(\" -kernel=kernel_number (0 - AsyncCopyMultiStageLargeChunk; 1 - \"\n \"AsyncCopyLargeChunk)\\n\");\n printf(\" (2 - AsyncCopyLargeChunkAWBarrier; 3 - \"\n \"AsyncCopyMultiStageSharedState)\\n\");\n printf(\" (4 - AsyncCopyMultiStage; 5 - \"\n \"AsyncCopySingleStage; 6 - Naive without memcpy_async)\\n\");\n printf(\" (7 - NaiveLargeChunk without \"\n \"memcpy_async)\\n\");\n printf(\" Note: Outer matrix dimensions of A & B matrices must be equal.\\n\");\n\n exit(EXIT_SUCCESS);\n }\n\n // This will pick the best possible CUDA capable device, otherwise\n // override the device ID based on input provided at the command line\n int dev = findCudaDevice(argc, (const char **)argv);\n\n int matrixBlock = 32;\n dim3 dimsA(10 * 4 * matrixBlock, 10 * 4 * matrixBlock, 1);\n dim3 dimsB(10 * 4 * matrixBlock, 10 * 4 * matrixBlock, 1);\n\n // width of Matrix A\n if (checkCmdLineFlag(argc, (const char **)argv, \"wA\")) {\n dimsA.x = getCmdLineArgumentInt(argc, (const char **)argv, \"wA\");\n }\n\n // height of Matrix A\n if (checkCmdLineFlag(argc, (const char **)argv, \"hA\")) {\n dimsA.y = getCmdLineArgumentInt(argc, (const char **)argv, \"hA\");\n }\n\n // width of Matrix B\n if (checkCmdLineFlag(argc, (const char **)argv, \"wB\")) {\n dimsB.x = getCmdLineArgumentInt(argc, (const char **)argv, \"wB\");\n }\n\n // height of Matrix B\n if (checkCmdLineFlag(argc, (const char **)argv, \"hB\")) {\n dimsB.y = getCmdLineArgumentInt(argc, (const char **)argv, \"hB\");\n }\n\n if (dimsA.x != dimsB.y) {\n printf(\"Error: outer matrix dimensions must be equal. (%d != %d)\\n\", dimsA.x, dimsB.y);\n exit(EXIT_FAILURE);\n }\n\n kernels selected_kernel = AsyncCopyMultiStageLargeChunk;\n\n // kernel to run - default (AsyncCopyMultiStageLargeChunk == 0)\n if (checkCmdLineFlag(argc, (const char **)argv, \"kernel\")) {\n int kernel_number = getCmdLineArgumentInt(argc, (const char **)argv, \"kernel\");\n if (kernel_number < 8) {\n selected_kernel = (kernels)kernel_number;\n }\n else {\n printf(\"Error: kernel number should be between 0 to 6, you have entered \"\n \"%d\\n\",\n kernel_number);\n exit(EXIT_FAILURE);\n }\n }\n\n int major = 0;\n checkCudaErrors(hipDeviceGetAttribute(&major, hipDeviceAttributeComputeCapabilityMajor, dev));\n if (major < 7) {\n printf(\"globalToShmemAsyncCopy requires SM 7.0 or higher. Exiting...\\n\");\n exit(EXIT_WAIVED);\n }\n\n printf(\"MatrixA(%d,%d), MatrixB(%d,%d)\\n\", dimsA.x, dimsA.y, dimsB.x, dimsB.y);\n\n int matrix_result = MatrixMultiply(argc, argv, dimsA, dimsB, selected_kernel);\n\n exit(matrix_result);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/globalToShmemAsyncCopy/globalToShmemAsyncCopy.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n// includes, project\n#include \n#include \n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#define NUM_ELEMS 10000000\n#define NUM_THREADS_PER_BLOCK 512\n\n// warp-aggregated atomic increment\n__device__ int atomicAggInc(int *counter)\n{\n cg::coalesced_group active = cg::coalesced_threads();\n\n // leader does the update\n int res = 0;\n if (active.thread_rank() == 0) {\n res = atomicAdd(counter, active.size());\n }\n\n // broadcast result\n res = active.shfl(res, 0);\n\n // each thread computes its own value\n return res + active.thread_rank();\n}\n\n__global__ void filter_arr(int *dst, int *nres, const int *src, int n)\n{\n int id = threadIdx.x + blockIdx.x * blockDim.x;\n\n for (int i = id; i < n; i += gridDim.x * blockDim.x) {\n if (src[i] > 0)\n dst[atomicAggInc(nres)] = src[i];\n }\n}\n\n// warp-aggregated atomic multi bucket increment\n#if __CUDA_ARCH__ >= 700\n__device__ int atomicAggIncMulti(const int bucket, int *counter)\n{\n cg::coalesced_group active = cg::coalesced_threads();\n // group all threads with same bucket value.\n auto labeledGroup = cg::labeled_partition(active, bucket);\n\n int res = 0;\n if (labeledGroup.thread_rank() == 0) {\n res = atomicAdd(&counter[bucket], labeledGroup.size());\n }\n\n // broadcast result\n res = labeledGroup.shfl(res, 0);\n\n // each thread computes its own value\n return res + labeledGroup.thread_rank();\n}\n#endif\n\n// Places individual value indices into its corresponding buckets.\n__global__ void\nmapToBuckets(const int *srcArr, int *indicesBuckets, int *bucketCounters, const int srcSize, const int numOfBuckets)\n{\n#if __CUDA_ARCH__ >= 700\n cg::grid_group grid = cg::this_grid();\n\n for (int i = grid.thread_rank(); i < srcSize; i += grid.size()) {\n const int bucket = srcArr[i];\n if (bucket < numOfBuckets) {\n indicesBuckets[atomicAggIncMulti(bucket, bucketCounters)] = i;\n }\n }\n#endif\n}\n\nint mapIndicesToBuckets(int *h_srcArr, int *d_srcArr, int numOfBuckets)\n{\n int *d_indicesBuckets, *d_bucketCounters;\n int *cpuBucketCounters = new int[numOfBuckets];\n int *h_bucketCounters = new int[numOfBuckets];\n\n memset(cpuBucketCounters, 0, sizeof(int) * numOfBuckets);\n // Initialize each bucket counters.\n for (int i = 0; i < numOfBuckets; i++) {\n h_bucketCounters[i] = i * NUM_ELEMS;\n }\n\n checkCudaErrors(cudaMalloc(&d_indicesBuckets, sizeof(int) * NUM_ELEMS * numOfBuckets));\n checkCudaErrors(cudaMalloc(&d_bucketCounters, sizeof(int) * numOfBuckets));\n\n checkCudaErrors(cudaMemcpy(d_bucketCounters, h_bucketCounters, sizeof(int) * numOfBuckets, cudaMemcpyHostToDevice));\n\n dim3 dimBlock(NUM_THREADS_PER_BLOCK, 1, 1);\n dim3 dimGrid((NUM_ELEMS / NUM_THREADS_PER_BLOCK), 1, 1);\n\n mapToBuckets<<>>(d_srcArr, d_indicesBuckets, d_bucketCounters, NUM_ELEMS, numOfBuckets);\n\n checkCudaErrors(cudaMemcpy(h_bucketCounters, d_bucketCounters, sizeof(int) * numOfBuckets, cudaMemcpyDeviceToHost));\n\n for (int i = 0; i < NUM_ELEMS; i++) {\n cpuBucketCounters[h_srcArr[i]]++;\n }\n\n bool allMatch = true;\n int finalElems = 0;\n for (int i = 0; i < numOfBuckets; i++) {\n finalElems += (h_bucketCounters[i] - i * NUM_ELEMS);\n if (cpuBucketCounters[i] != (h_bucketCounters[i] - i * NUM_ELEMS)) {\n allMatch = false;\n break;\n }\n }\n\n if (!allMatch && finalElems != NUM_ELEMS) {\n return EXIT_FAILURE;\n }\n return EXIT_SUCCESS;\n}\n\n// Warp-aggregated atomic Max in multi bucket\n#if __CUDA_ARCH__ >= 700\n__device__ void atomicAggMaxMulti(const int bucket, int *counter, const int valueForMax)\n{\n cg::coalesced_group active = cg::coalesced_threads();\n // group all threads with same bucket value.\n auto labeledGroup = cg::labeled_partition(active, bucket);\n\n const int maxValueInGroup = cg::reduce(labeledGroup, valueForMax, cg::greater());\n\n if (labeledGroup.thread_rank() == 0) {\n atomicMax(&counter[bucket], maxValueInGroup);\n }\n}\n#endif\n\n// Performs max calculation in each buckets.\n__global__ void calculateMaxInEachBuckets(const int *srcArr,\n const int *valueInBuckets,\n int *bucketsMax,\n const int srcSize,\n const int numOfBuckets)\n{\n#if __CUDA_ARCH__ >= 700\n cg::grid_group grid = cg::this_grid();\n\n for (int i = grid.thread_rank(); i < srcSize; i += grid.size()) {\n const int bucket = srcArr[i];\n if (bucket < numOfBuckets) {\n atomicAggMaxMulti(bucket, bucketsMax, valueInBuckets[i]);\n }\n }\n#endif\n}\n\nint calculateMaxInBuckets(int *h_srcArr, int *d_srcArr, int numOfBuckets)\n{\n int *d_valueInBuckets, *d_bucketsMax;\n int *h_valueInBuckets = new int[NUM_ELEMS];\n int *cpuBucketsMax = new int[numOfBuckets];\n int *h_bucketsMax = new int[numOfBuckets];\n\n memset(cpuBucketsMax, 0, sizeof(int) * numOfBuckets);\n\n // Here we create values which is assumed to correspond to each\n // buckets of srcArr at same array index.\n for (int i = 0; i < NUM_ELEMS; i++) {\n h_valueInBuckets[i] = rand();\n }\n\n checkCudaErrors(cudaMalloc(&d_valueInBuckets, sizeof(int) * NUM_ELEMS));\n checkCudaErrors(cudaMalloc(&d_bucketsMax, sizeof(int) * numOfBuckets));\n\n checkCudaErrors(cudaMemset(d_bucketsMax, 0, sizeof(int) * numOfBuckets));\n checkCudaErrors(cudaMemcpy(d_valueInBuckets, h_valueInBuckets, sizeof(int) * NUM_ELEMS, cudaMemcpyHostToDevice));\n\n dim3 dimBlock(NUM_THREADS_PER_BLOCK, 1, 1);\n dim3 dimGrid((NUM_ELEMS / NUM_THREADS_PER_BLOCK), 1, 1);\n\n calculateMaxInEachBuckets<<>>(d_srcArr, d_valueInBuckets, d_bucketsMax, NUM_ELEMS, numOfBuckets);\n\n checkCudaErrors(cudaMemcpy(h_bucketsMax, d_bucketsMax, sizeof(int) * numOfBuckets, cudaMemcpyDeviceToHost));\n\n for (int i = 0; i < NUM_ELEMS; i++) {\n if (cpuBucketsMax[h_srcArr[i]] < h_valueInBuckets[i]) {\n cpuBucketsMax[h_srcArr[i]] = h_valueInBuckets[i];\n }\n }\n\n bool allMatch = true;\n int finalElems = 0;\n for (int i = 0; i < numOfBuckets; i++) {\n if (cpuBucketsMax[i] != h_bucketsMax[i]) {\n allMatch = false;\n printf(\"CPU i=%d max = %d mismatches GPU max = %d\\n\", i, cpuBucketsMax[i], h_bucketsMax[i]);\n break;\n }\n }\n if (allMatch) {\n printf(\"CPU max matches GPU max\\n\");\n }\n\n delete[] h_valueInBuckets;\n delete[] cpuBucketsMax;\n delete[] h_bucketsMax;\n checkCudaErrors(cudaFree(d_valueInBuckets));\n checkCudaErrors(cudaFree(d_bucketsMax));\n\n if (!allMatch && finalElems != NUM_ELEMS) {\n return EXIT_FAILURE;\n }\n\n return EXIT_SUCCESS;\n}\n\nint main(int argc, char **argv)\n{\n int *data_to_filter, *filtered_data, nres = 0;\n int *d_data_to_filter, *d_filtered_data, *d_nres;\n\n int numOfBuckets = 5;\n\n data_to_filter = reinterpret_cast(malloc(sizeof(int) * NUM_ELEMS));\n\n // Generate input data.\n for (int i = 0; i < NUM_ELEMS; i++) {\n data_to_filter[i] = rand() % numOfBuckets;\n }\n\n int devId = findCudaDevice(argc, (const char **)argv);\n\n checkCudaErrors(cudaMalloc(&d_data_to_filter, sizeof(int) * NUM_ELEMS));\n checkCudaErrors(cudaMalloc(&d_filtered_data, sizeof(int) * NUM_ELEMS));\n checkCudaErrors(cudaMalloc(&d_nres, sizeof(int)));\n\n checkCudaErrors(cudaMemcpy(d_data_to_filter, data_to_filter, sizeof(int) * NUM_ELEMS, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemset(d_nres, 0, sizeof(int)));\n\n dim3 dimBlock(NUM_THREADS_PER_BLOCK, 1, 1);\n dim3 dimGrid((NUM_ELEMS / NUM_THREADS_PER_BLOCK) + 1, 1, 1);\n\n filter_arr<<>>(d_filtered_data, d_nres, d_data_to_filter, NUM_ELEMS);\n\n checkCudaErrors(cudaMemcpy(&nres, d_nres, sizeof(int), cudaMemcpyDeviceToHost));\n\n filtered_data = reinterpret_cast(malloc(sizeof(int) * nres));\n\n checkCudaErrors(cudaMemcpy(filtered_data, d_filtered_data, sizeof(int) * nres, cudaMemcpyDeviceToHost));\n\n int *host_filtered_data = reinterpret_cast(malloc(sizeof(int) * NUM_ELEMS));\n\n // Generate host output with host filtering code.\n int host_flt_count = 0;\n for (int i = 0; i < NUM_ELEMS; i++) {\n if (data_to_filter[i] > 0) {\n host_filtered_data[host_flt_count++] = data_to_filter[i];\n }\n }\n\n int major = 0;\n checkCudaErrors(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, devId));\n\n int mapIndicesToBucketsStatus = EXIT_SUCCESS;\n int calculateMaxInBucketsStatus = EXIT_SUCCESS;\n // atomicAggIncMulti & atomicAggMaxMulti require a GPU of Volta (SM7X) architecture or higher,\n // so that it can take advantage of the new MATCH capability of Volta hardware\n if (major >= 7) {\n mapIndicesToBucketsStatus = mapIndicesToBuckets(data_to_filter, d_data_to_filter, numOfBuckets);\n calculateMaxInBucketsStatus = calculateMaxInBuckets(data_to_filter, d_data_to_filter, numOfBuckets);\n }\n\n printf(\"\\nWarp Aggregated Atomics %s \\n\",\n (host_flt_count == nres) && (mapIndicesToBucketsStatus == EXIT_SUCCESS)\n && (calculateMaxInBucketsStatus == EXIT_SUCCESS)\n ? \"PASSED\"\n : \"FAILED\");\n\n checkCudaErrors(cudaFree(d_data_to_filter));\n checkCudaErrors(cudaFree(d_filtered_data));\n checkCudaErrors(cudaFree(d_nres));\n free(data_to_filter);\n free(filtered_data);\n free(host_filtered_data);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n// includes, project\n#include \n#include \n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#define NUM_ELEMS 10000000\n#define NUM_THREADS_PER_BLOCK 512\n\n// warp-aggregated atomic increment\n__device__ int atomicAggInc(int *counter)\n{\n cg::coalesced_group active = cg::coalesced_threads();\n\n // leader does the update\n int res = 0;\n if (active.thread_rank() == 0) {\n res = atomicAdd(counter, active.size());\n }\n\n // broadcast result\n res = active.shfl(res, 0);\n\n // each thread computes its own value\n return res + active.thread_rank();\n}\n\n__global__ void filter_arr(int *dst, int *nres, const int *src, int n)\n{\n int id = threadIdx.x + blockIdx.x * blockDim.x;\n\n for (int i = id; i < n; i += gridDim.x * blockDim.x) {\n if (src[i] > 0)\n dst[atomicAggInc(nres)] = src[i];\n }\n}\n\n// warp-aggregated atomic multi bucket increment\n#if __CUDA_ARCH__ >= 700\n__device__ int atomicAggIncMulti(const int bucket, int *counter)\n{\n cg::coalesced_group active = cg::coalesced_threads();\n // group all threads with same bucket value.\n auto labeledGroup = cg::labeled_partition(active, bucket);\n\n int res = 0;\n if (labeledGroup.thread_rank() == 0) {\n res = atomicAdd(&counter[bucket], labeledGroup.size());\n }\n\n // broadcast result\n res = labeledGroup.shfl(res, 0);\n\n // each thread computes its own value\n return res + labeledGroup.thread_rank();\n}\n#endif\n\n// Places individual value indices into its corresponding buckets.\n__global__ void\nmapToBuckets(const int *srcArr, int *indicesBuckets, int *bucketCounters, const int srcSize, const int numOfBuckets)\n{\n#if __CUDA_ARCH__ >= 700\n cg::grid_group grid = cg::this_grid();\n\n for (int i = grid.thread_rank(); i < srcSize; i += grid.size()) {\n const int bucket = srcArr[i];\n if (bucket < numOfBuckets) {\n indicesBuckets[atomicAggIncMulti(bucket, bucketCounters)] = i;\n }\n }\n#endif\n}\n\nint mapIndicesToBuckets(int *h_srcArr, int *d_srcArr, int numOfBuckets)\n{\n int *d_indicesBuckets, *d_bucketCounters;\n int *cpuBucketCounters = new int[numOfBuckets];\n int *h_bucketCounters = new int[numOfBuckets];\n\n memset(cpuBucketCounters, 0, sizeof(int) * numOfBuckets);\n // Initialize each bucket counters.\n for (int i = 0; i < numOfBuckets; i++) {\n h_bucketCounters[i] = i * NUM_ELEMS;\n }\n\n checkCudaErrors(hipMalloc(&d_indicesBuckets, sizeof(int) * NUM_ELEMS * numOfBuckets));\n checkCudaErrors(hipMalloc(&d_bucketCounters, sizeof(int) * numOfBuckets));\n\n checkCudaErrors(hipMemcpy(d_bucketCounters, h_bucketCounters, sizeof(int) * numOfBuckets, hipMemcpyHostToDevice));\n\n dim3 dimBlock(NUM_THREADS_PER_BLOCK, 1, 1);\n dim3 dimGrid((NUM_ELEMS / NUM_THREADS_PER_BLOCK), 1, 1);\n\n mapToBuckets<<>>(d_srcArr, d_indicesBuckets, d_bucketCounters, NUM_ELEMS, numOfBuckets);\n\n checkCudaErrors(hipMemcpy(h_bucketCounters, d_bucketCounters, sizeof(int) * numOfBuckets, hipMemcpyDeviceToHost));\n\n for (int i = 0; i < NUM_ELEMS; i++) {\n cpuBucketCounters[h_srcArr[i]]++;\n }\n\n bool allMatch = true;\n int finalElems = 0;\n for (int i = 0; i < numOfBuckets; i++) {\n finalElems += (h_bucketCounters[i] - i * NUM_ELEMS);\n if (cpuBucketCounters[i] != (h_bucketCounters[i] - i * NUM_ELEMS)) {\n allMatch = false;\n break;\n }\n }\n\n if (!allMatch && finalElems != NUM_ELEMS) {\n return EXIT_FAILURE;\n }\n return EXIT_SUCCESS;\n}\n\n// Warp-aggregated atomic Max in multi bucket\n#if __CUDA_ARCH__ >= 700\n__device__ void atomicAggMaxMulti(const int bucket, int *counter, const int valueForMax)\n{\n cg::coalesced_group active = cg::coalesced_threads();\n // group all threads with same bucket value.\n auto labeledGroup = cg::labeled_partition(active, bucket);\n\n const int maxValueInGroup = cg::reduce(labeledGroup, valueForMax, cg::greater());\n\n if (labeledGroup.thread_rank() == 0) {\n atomicMax(&counter[bucket], maxValueInGroup);\n }\n}\n#endif\n\n// Performs max calculation in each buckets.\n__global__ void calculateMaxInEachBuckets(const int *srcArr,\n const int *valueInBuckets,\n int *bucketsMax,\n const int srcSize,\n const int numOfBuckets)\n{\n#if __CUDA_ARCH__ >= 700\n cg::grid_group grid = cg::this_grid();\n\n for (int i = grid.thread_rank(); i < srcSize; i += grid.size()) {\n const int bucket = srcArr[i];\n if (bucket < numOfBuckets) {\n atomicAggMaxMulti(bucket, bucketsMax, valueInBuckets[i]);\n }\n }\n#endif\n}\n\nint calculateMaxInBuckets(int *h_srcArr, int *d_srcArr, int numOfBuckets)\n{\n int *d_valueInBuckets, *d_bucketsMax;\n int *h_valueInBuckets = new int[NUM_ELEMS];\n int *cpuBucketsMax = new int[numOfBuckets];\n int *h_bucketsMax = new int[numOfBuckets];\n\n memset(cpuBucketsMax, 0, sizeof(int) * numOfBuckets);\n\n // Here we create values which is assumed to correspond to each\n // buckets of srcArr at same array index.\n for (int i = 0; i < NUM_ELEMS; i++) {\n h_valueInBuckets[i] = rand();\n }\n\n checkCudaErrors(hipMalloc(&d_valueInBuckets, sizeof(int) * NUM_ELEMS));\n checkCudaErrors(hipMalloc(&d_bucketsMax, sizeof(int) * numOfBuckets));\n\n checkCudaErrors(hipMemset(d_bucketsMax, 0, sizeof(int) * numOfBuckets));\n checkCudaErrors(hipMemcpy(d_valueInBuckets, h_valueInBuckets, sizeof(int) * NUM_ELEMS, hipMemcpyHostToDevice));\n\n dim3 dimBlock(NUM_THREADS_PER_BLOCK, 1, 1);\n dim3 dimGrid((NUM_ELEMS / NUM_THREADS_PER_BLOCK), 1, 1);\n\n calculateMaxInEachBuckets<<>>(d_srcArr, d_valueInBuckets, d_bucketsMax, NUM_ELEMS, numOfBuckets);\n\n checkCudaErrors(hipMemcpy(h_bucketsMax, d_bucketsMax, sizeof(int) * numOfBuckets, hipMemcpyDeviceToHost));\n\n for (int i = 0; i < NUM_ELEMS; i++) {\n if (cpuBucketsMax[h_srcArr[i]] < h_valueInBuckets[i]) {\n cpuBucketsMax[h_srcArr[i]] = h_valueInBuckets[i];\n }\n }\n\n bool allMatch = true;\n int finalElems = 0;\n for (int i = 0; i < numOfBuckets; i++) {\n if (cpuBucketsMax[i] != h_bucketsMax[i]) {\n allMatch = false;\n printf(\"CPU i=%d max = %d mismatches GPU max = %d\\n\", i, cpuBucketsMax[i], h_bucketsMax[i]);\n break;\n }\n }\n if (allMatch) {\n printf(\"CPU max matches GPU max\\n\");\n }\n\n delete[] h_valueInBuckets;\n delete[] cpuBucketsMax;\n delete[] h_bucketsMax;\n checkCudaErrors(hipFree(d_valueInBuckets));\n checkCudaErrors(hipFree(d_bucketsMax));\n\n if (!allMatch && finalElems != NUM_ELEMS) {\n return EXIT_FAILURE;\n }\n\n return EXIT_SUCCESS;\n}\n\nint main(int argc, char **argv)\n{\n int *data_to_filter, *filtered_data, nres = 0;\n int *d_data_to_filter, *d_filtered_data, *d_nres;\n\n int numOfBuckets = 5;\n\n data_to_filter = reinterpret_cast(malloc(sizeof(int) * NUM_ELEMS));\n\n // Generate input data.\n for (int i = 0; i < NUM_ELEMS; i++) {\n data_to_filter[i] = rand() % numOfBuckets;\n }\n\n int devId = findCudaDevice(argc, (const char **)argv);\n\n checkCudaErrors(hipMalloc(&d_data_to_filter, sizeof(int) * NUM_ELEMS));\n checkCudaErrors(hipMalloc(&d_filtered_data, sizeof(int) * NUM_ELEMS));\n checkCudaErrors(hipMalloc(&d_nres, sizeof(int)));\n\n checkCudaErrors(hipMemcpy(d_data_to_filter, data_to_filter, sizeof(int) * NUM_ELEMS, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemset(d_nres, 0, sizeof(int)));\n\n dim3 dimBlock(NUM_THREADS_PER_BLOCK, 1, 1);\n dim3 dimGrid((NUM_ELEMS / NUM_THREADS_PER_BLOCK) + 1, 1, 1);\n\n filter_arr<<>>(d_filtered_data, d_nres, d_data_to_filter, NUM_ELEMS);\n\n checkCudaErrors(hipMemcpy(&nres, d_nres, sizeof(int), hipMemcpyDeviceToHost));\n\n filtered_data = reinterpret_cast(malloc(sizeof(int) * nres));\n\n checkCudaErrors(hipMemcpy(filtered_data, d_filtered_data, sizeof(int) * nres, hipMemcpyDeviceToHost));\n\n int *host_filtered_data = reinterpret_cast(malloc(sizeof(int) * NUM_ELEMS));\n\n // Generate host output with host filtering code.\n int host_flt_count = 0;\n for (int i = 0; i < NUM_ELEMS; i++) {\n if (data_to_filter[i] > 0) {\n host_filtered_data[host_flt_count++] = data_to_filter[i];\n }\n }\n\n int major = 0;\n checkCudaErrors(hipDeviceGetAttribute(&major, hipDeviceAttributeComputeCapabilityMajor, devId));\n\n int mapIndicesToBucketsStatus = EXIT_SUCCESS;\n int calculateMaxInBucketsStatus = EXIT_SUCCESS;\n // atomicAggIncMulti & atomicAggMaxMulti require a GPU of Volta (SM7X) architecture or higher,\n // so that it can take advantage of the new MATCH capability of Volta hardware\n if (major >= 7) {\n mapIndicesToBucketsStatus = mapIndicesToBuckets(data_to_filter, d_data_to_filter, numOfBuckets);\n calculateMaxInBucketsStatus = calculateMaxInBuckets(data_to_filter, d_data_to_filter, numOfBuckets);\n }\n\n printf(\"\\nWarp Aggregated Atomics %s \\n\",\n (host_flt_count == nres) && (mapIndicesToBucketsStatus == EXIT_SUCCESS)\n && (calculateMaxInBucketsStatus == EXIT_SUCCESS)\n ? \"PASSED\"\n : \"FAILED\");\n\n checkCudaErrors(hipFree(d_data_to_filter));\n checkCudaErrors(hipFree(d_filtered_data));\n checkCudaErrors(hipFree(d_nres));\n free(data_to_filter);\n free(filtered_data);\n free(host_filtered_data);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/warpAggregatedAtomicsCG/warpAggregatedAtomicsCG.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample illustrates basic usage of binary partition cooperative groups\n * within the thread block tile when divergent path exists.\n * 1.) Each thread loads a value from random array.\n * 2.) then checks if it is odd or even.\n * 3.) create binary partition group based on the above predicate\n * 4.) we count the number of odd/even in the group based on size of the binary\n groups\n * 5.) write it global counter of odd.\n * 6.) sum the values loaded by individual threads(using reduce) and write it to\n global even & odd elements sum.\n *\n * **NOTE** :\n * binary_partition results in splitting warp into divergent thread groups\n * this is not good from performance perspective, but in cases where warp\n * divergence is inevitable one can use binary_partition group.\n*/\n\n#include \n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\nvoid initOddEvenArr(int *inputArr, unsigned int size)\n{\n for (int i = 0; i < size; i++) {\n inputArr[i] = rand() % 50;\n }\n}\n\n/**\n * CUDA kernel device code\n *\n * Creates cooperative groups and performs odd/even counting & summation.\n */\n__global__ void oddEvenCountAndSumCG(int *inputArr, int *numOfOdds, int *sumOfOddAndEvens, unsigned int size)\n{\n cg::thread_block cta = cg::this_thread_block();\n cg::grid_group grid = cg::this_grid();\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n for (int i = grid.thread_rank(); i < size; i += grid.size()) {\n int elem = inputArr[i];\n auto subTile = cg::binary_partition(tile32, elem & 1);\n if (elem & 1) // Odd numbers group\n {\n int oddGroupSum = cg::reduce(subTile, elem, cg::plus());\n\n if (subTile.thread_rank() == 0) {\n // Add number of odds present in this group of Odds.\n atomicAdd(numOfOdds, subTile.size());\n\n // Add local reduction of odds present in this group of Odds.\n atomicAdd(&sumOfOddAndEvens[0], oddGroupSum);\n }\n }\n else // Even numbers group\n {\n int evenGroupSum = cg::reduce(subTile, elem, cg::plus());\n\n if (subTile.thread_rank() == 0) {\n // Add local reduction of even present in this group of evens.\n atomicAdd(&sumOfOddAndEvens[1], evenGroupSum);\n }\n }\n // reconverge warp so for next loop iteration we ensure convergence of\n // above diverged threads to perform coalesced loads of inputArr.\n cg::sync(tile32);\n }\n}\n\n/**\n * Host main routine\n */\nint main(int argc, const char **argv)\n{\n int deviceId = findCudaDevice(argc, argv);\n int *h_inputArr, *d_inputArr;\n int *h_numOfOdds, *d_numOfOdds;\n int *h_sumOfOddEvenElems, *d_sumOfOddEvenElems;\n unsigned int arrSize = 1024 * 100;\n\n checkCudaErrors(cudaMallocHost(&h_inputArr, sizeof(int) * arrSize));\n checkCudaErrors(cudaMallocHost(&h_numOfOdds, sizeof(int)));\n checkCudaErrors(cudaMallocHost(&h_sumOfOddEvenElems, sizeof(int) * 2));\n initOddEvenArr(h_inputArr, arrSize);\n\n cudaStream_t stream;\n checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));\n checkCudaErrors(cudaMalloc(&d_inputArr, sizeof(int) * arrSize));\n checkCudaErrors(cudaMalloc(&d_numOfOdds, sizeof(int)));\n checkCudaErrors(cudaMalloc(&d_sumOfOddEvenElems, sizeof(int) * 2));\n\n checkCudaErrors(cudaMemcpyAsync(d_inputArr, h_inputArr, sizeof(int) * arrSize, cudaMemcpyHostToDevice, stream));\n checkCudaErrors(cudaMemsetAsync(d_numOfOdds, 0, sizeof(int), stream));\n checkCudaErrors(cudaMemsetAsync(d_sumOfOddEvenElems, 0, 2 * sizeof(int), stream));\n\n // Launch the kernel\n int threadsPerBlock = 0;\n int blocksPerGrid = 0;\n checkCudaErrors(cudaOccupancyMaxPotentialBlockSize(&blocksPerGrid, &threadsPerBlock, oddEvenCountAndSumCG, 0, 0));\n\n printf(\"\\nLaunching %d blocks with %d threads...\\n\\n\", blocksPerGrid, threadsPerBlock);\n\n oddEvenCountAndSumCG<<>>(\n d_inputArr, d_numOfOdds, d_sumOfOddEvenElems, arrSize);\n\n checkCudaErrors(cudaMemcpyAsync(h_numOfOdds, d_numOfOdds, sizeof(int), cudaMemcpyDeviceToHost, stream));\n checkCudaErrors(\n cudaMemcpyAsync(h_sumOfOddEvenElems, d_sumOfOddEvenElems, 2 * sizeof(int), cudaMemcpyDeviceToHost, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n printf(\"Array size = %d Num of Odds = %d Sum of Odds = %d Sum of Evens %d\\n\",\n arrSize,\n h_numOfOdds[0],\n h_sumOfOddEvenElems[0],\n h_sumOfOddEvenElems[1]);\n printf(\"\\n...Done.\\n\\n\");\n\n checkCudaErrors(cudaFreeHost(h_inputArr));\n checkCudaErrors(cudaFreeHost(h_numOfOdds));\n checkCudaErrors(cudaFreeHost(h_sumOfOddEvenElems));\n\n checkCudaErrors(cudaFree(d_inputArr));\n checkCudaErrors(cudaFree(d_numOfOdds));\n checkCudaErrors(cudaFree(d_sumOfOddEvenElems));\n\n return EXIT_SUCCESS;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample illustrates basic usage of binary partition cooperative groups\n * within the thread block tile when divergent path exists.\n * 1.) Each thread loads a value from random array.\n * 2.) then checks if it is odd or even.\n * 3.) create binary partition group based on the above predicate\n * 4.) we count the number of odd/even in the group based on size of the binary\n groups\n * 5.) write it global counter of odd.\n * 6.) sum the values loaded by individual threads(using reduce) and write it to\n global even & odd elements sum.\n *\n * **NOTE** :\n * binary_partition results in splitting warp into divergent thread groups\n * this is not good from performance perspective, but in cases where warp\n * divergence is inevitable one can use binary_partition group.\n*/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\nvoid initOddEvenArr(int *inputArr, unsigned int size)\n{\n for (int i = 0; i < size; i++) {\n inputArr[i] = rand() % 50;\n }\n}\n\n/**\n * CUDA kernel device code\n *\n * Creates cooperative groups and performs odd/even counting & summation.\n */\n__global__ void oddEvenCountAndSumCG(int *inputArr, int *numOfOdds, int *sumOfOddAndEvens, unsigned int size)\n{\n cg::thread_block cta = cg::this_thread_block();\n cg::grid_group grid = cg::this_grid();\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n for (int i = grid.thread_rank(); i < size; i += grid.size()) {\n int elem = inputArr[i];\n auto subTile = cg::binary_partition(tile32, elem & 1);\n if (elem & 1) // Odd numbers group\n {\n int oddGroupSum = cg::reduce(subTile, elem, cg::plus());\n\n if (subTile.thread_rank() == 0) {\n // Add number of odds present in this group of Odds.\n atomicAdd(numOfOdds, subTile.size());\n\n // Add local reduction of odds present in this group of Odds.\n atomicAdd(&sumOfOddAndEvens[0], oddGroupSum);\n }\n }\n else // Even numbers group\n {\n int evenGroupSum = cg::reduce(subTile, elem, cg::plus());\n\n if (subTile.thread_rank() == 0) {\n // Add local reduction of even present in this group of evens.\n atomicAdd(&sumOfOddAndEvens[1], evenGroupSum);\n }\n }\n // reconverge warp so for next loop iteration we ensure convergence of\n // above diverged threads to perform coalesced loads of inputArr.\n cg::sync(tile32);\n }\n}\n\n/**\n * Host main routine\n */\nint main(int argc, const char **argv)\n{\n int deviceId = findCudaDevice(argc, argv);\n int *h_inputArr, *d_inputArr;\n int *h_numOfOdds, *d_numOfOdds;\n int *h_sumOfOddEvenElems, *d_sumOfOddEvenElems;\n unsigned int arrSize = 1024 * 100;\n\n checkCudaErrors(hipHostMalloc(&h_inputArr, sizeof(int) * arrSize));\n checkCudaErrors(hipHostMalloc(&h_numOfOdds, sizeof(int)));\n checkCudaErrors(hipHostMalloc(&h_sumOfOddEvenElems, sizeof(int) * 2));\n initOddEvenArr(h_inputArr, arrSize);\n\n hipStream_t stream;\n checkCudaErrors(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));\n checkCudaErrors(hipMalloc(&d_inputArr, sizeof(int) * arrSize));\n checkCudaErrors(hipMalloc(&d_numOfOdds, sizeof(int)));\n checkCudaErrors(hipMalloc(&d_sumOfOddEvenElems, sizeof(int) * 2));\n\n checkCudaErrors(hipMemcpyAsync(d_inputArr, h_inputArr, sizeof(int) * arrSize, hipMemcpyHostToDevice, stream));\n checkCudaErrors(hipMemsetAsync(d_numOfOdds, 0, sizeof(int), stream));\n checkCudaErrors(hipMemsetAsync(d_sumOfOddEvenElems, 0, 2 * sizeof(int), stream));\n\n // Launch the kernel\n int threadsPerBlock = 0;\n int blocksPerGrid = 0;\n checkCudaErrors(hipOccupancyMaxPotentialBlockSize(&blocksPerGrid, &threadsPerBlock, oddEvenCountAndSumCG, 0, 0));\n\n printf(\"\\nLaunching %d blocks with %d threads...\\n\\n\", blocksPerGrid, threadsPerBlock);\n\n oddEvenCountAndSumCG<<>>(\n d_inputArr, d_numOfOdds, d_sumOfOddEvenElems, arrSize);\n\n checkCudaErrors(hipMemcpyAsync(h_numOfOdds, d_numOfOdds, sizeof(int), hipMemcpyDeviceToHost, stream));\n checkCudaErrors(\n hipMemcpyAsync(h_sumOfOddEvenElems, d_sumOfOddEvenElems, 2 * sizeof(int), hipMemcpyDeviceToHost, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n\n printf(\"Array size = %d Num of Odds = %d Sum of Odds = %d Sum of Evens %d\\n\",\n arrSize,\n h_numOfOdds[0],\n h_sumOfOddEvenElems[0],\n h_sumOfOddEvenElems[1]);\n printf(\"\\n...Done.\\n\\n\");\n\n checkCudaErrors(hipHostFree(h_inputArr));\n checkCudaErrors(hipHostFree(h_numOfOdds));\n checkCudaErrors(hipHostFree(h_sumOfOddEvenElems));\n\n checkCudaErrors(hipFree(d_inputArr));\n checkCudaErrors(hipFree(d_numOfOdds));\n checkCudaErrors(hipFree(d_sumOfOddEvenElems));\n\n return EXIT_SUCCESS;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/binaryPartitionCG/binaryPartitionCG.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n\n// for cuda::std::tuple\n#include \n\nnamespace cg = cooperative_groups;\n#include \n\n////////////////////////////////////////////////////////////////////////////////\n// A structure of 2D points (structure of arrays).\n////////////////////////////////////////////////////////////////////////////////\nclass Points\n{\n float *m_x;\n float *m_y;\n\npublic:\n // Constructor.\n __host__ __device__ Points()\n : m_x(NULL)\n , m_y(NULL)\n {\n }\n\n // Constructor.\n __host__ __device__ Points(float *x, float *y)\n : m_x(x)\n , m_y(y)\n {\n }\n\n // Get a point.\n __host__ __device__ __forceinline__ float2 get_point(int idx) const { return make_float2(m_x[idx], m_y[idx]); }\n\n // Set a point.\n __host__ __device__ __forceinline__ void set_point(int idx, const float2 &p)\n {\n m_x[idx] = p.x;\n m_y[idx] = p.y;\n }\n\n // Set the pointers.\n __host__ __device__ __forceinline__ void set(float *x, float *y)\n {\n m_x = x;\n m_y = y;\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n// A 2D bounding box\n////////////////////////////////////////////////////////////////////////////////\nclass Bounding_box\n{\n // Extreme points of the bounding box.\n float2 m_p_min;\n float2 m_p_max;\n\npublic:\n // Constructor. Create a unit box.\n __host__ __device__ Bounding_box()\n {\n m_p_min = make_float2(0.0f, 0.0f);\n m_p_max = make_float2(1.0f, 1.0f);\n }\n\n // Compute the center of the bounding-box.\n __host__ __device__ void compute_center(float2 ¢er) const\n {\n center.x = 0.5f * (m_p_min.x + m_p_max.x);\n center.y = 0.5f * (m_p_min.y + m_p_max.y);\n }\n\n // The points of the box.\n __host__ __device__ __forceinline__ const float2 &get_max() const { return m_p_max; }\n\n __host__ __device__ __forceinline__ const float2 &get_min() const { return m_p_min; }\n\n // Does a box contain a point.\n __host__ __device__ bool contains(const float2 &p) const\n {\n return p.x >= m_p_min.x && p.x < m_p_max.x && p.y >= m_p_min.y && p.y < m_p_max.y;\n }\n\n // Define the bounding box.\n __host__ __device__ void set(float min_x, float min_y, float max_x, float max_y)\n {\n m_p_min.x = min_x;\n m_p_min.y = min_y;\n m_p_max.x = max_x;\n m_p_max.y = max_y;\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n// A node of a quadree.\n////////////////////////////////////////////////////////////////////////////////\nclass Quadtree_node\n{\n // The identifier of the node.\n int m_id;\n // The bounding box of the tree.\n Bounding_box m_bounding_box;\n // The range of points.\n int m_begin, m_end;\n\npublic:\n // Constructor.\n __host__ __device__ Quadtree_node()\n : m_id(0)\n , m_begin(0)\n , m_end(0)\n {\n }\n\n // The ID of a node at its level.\n __host__ __device__ int id() const { return m_id; }\n\n // The ID of a node at its level.\n __host__ __device__ void set_id(int new_id) { m_id = new_id; }\n\n // The bounding box.\n __host__ __device__ __forceinline__ const Bounding_box &bounding_box() const { return m_bounding_box; }\n\n // Set the bounding box.\n __host__ __device__ __forceinline__ void set_bounding_box(float min_x, float min_y, float max_x, float max_y)\n {\n m_bounding_box.set(min_x, min_y, max_x, max_y);\n }\n\n // The number of points in the tree.\n __host__ __device__ __forceinline__ int num_points() const { return m_end - m_begin; }\n\n // The range of points in the tree.\n __host__ __device__ __forceinline__ int points_begin() const { return m_begin; }\n\n __host__ __device__ __forceinline__ int points_end() const { return m_end; }\n\n // Define the range for that node.\n __host__ __device__ __forceinline__ void set_range(int begin, int end)\n {\n m_begin = begin;\n m_end = end;\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n// Algorithm parameters.\n////////////////////////////////////////////////////////////////////////////////\nstruct Parameters\n{\n // Choose the right set of points to use as in/out.\n int point_selector;\n // The number of nodes at a given level (2^k for level k).\n int num_nodes_at_this_level;\n // The recursion depth.\n int depth;\n // The max value for depth.\n const int max_depth;\n // The minimum number of points in a node to stop recursion.\n const int min_points_per_node;\n\n // Constructor set to default values.\n __host__ __device__ Parameters(int max_depth, int min_points_per_node)\n : point_selector(0)\n , num_nodes_at_this_level(1)\n , depth(0)\n , max_depth(max_depth)\n , min_points_per_node(min_points_per_node)\n {\n }\n\n // Copy constructor. Changes the values for next iteration.\n __host__ __device__ Parameters(const Parameters ¶ms, bool)\n : point_selector((params.point_selector + 1) % 2)\n , num_nodes_at_this_level(4 * params.num_nodes_at_this_level)\n , depth(params.depth + 1)\n , max_depth(params.max_depth)\n , min_points_per_node(params.min_points_per_node)\n {\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n// Build a quadtree on the GPU. Use CUDA Dynamic Parallelism.\n//\n// The algorithm works as follows. The host (CPU) launches one block of\n// NUM_THREADS_PER_BLOCK threads. That block will do the following steps:\n//\n// 1- Check the number of points and its depth.\n//\n// We impose a maximum depth to the tree and a minimum number of points per\n// node. If the maximum depth is exceeded or the minimum number of points is\n// reached. The threads in the block exit.\n//\n// Before exiting, they perform a buffer swap if it is needed. Indeed, the\n// algorithm uses two buffers to permute the points and make sure they are\n// properly distributed in the quadtree. By design we want all points to be\n// in the first buffer of points at the end of the algorithm. It is the reason\n// why we may have to swap the buffer before leavin (if the points are in the\n// 2nd buffer).\n//\n// 2- Count the number of points in each child.\n//\n// If the depth is not too high and the number of points is sufficient, the\n// block has to dispatch the points into four geometrical buckets: Its\n// children. For that purpose, we compute the center of the bounding box and\n// count the number of points in each quadrant.\n//\n// The set of points is divided into sections. Each section is given to a\n// warp of threads (32 threads). Warps use __ballot and __popc intrinsics\n// to count the points. See the Programming Guide for more information about\n// those functions.\n//\n// 3- Scan the warps' results to know the \"global\" numbers.\n//\n// Warps work independently from each other. At the end, each warp knows the\n// number of points in its section. To know the numbers for the block, the\n// block has to run a scan/reduce at the block level. It's a traditional\n// approach. The implementation in that sample is not as optimized as what\n// could be found in fast radix sorts, for example, but it relies on the same\n// idea.\n//\n// 4- Move points.\n//\n// Now that the block knows how many points go in each of its 4 children, it\n// remains to dispatch the points. It is straightforward.\n//\n// 5- Launch new blocks.\n//\n// The block launches four new blocks: One per children. Each of the four blocks\n// will apply the same algorithm.\n////////////////////////////////////////////////////////////////////////////////\ntemplate \n__global__ void build_quadtree_kernel(Quadtree_node *nodes, Points *points, Parameters params)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // The number of warps in a block.\n const int NUM_WARPS_PER_BLOCK = NUM_THREADS_PER_BLOCK / warpSize;\n\n // Shared memory to store the number of points.\n extern __shared__ int smem[];\n\n // s_num_pts[4][NUM_WARPS_PER_BLOCK];\n // Addresses of shared memory.\n volatile int *s_num_pts[4];\n\n for (int i = 0; i < 4; ++i)\n s_num_pts[i] = (volatile int *)&smem[i * NUM_WARPS_PER_BLOCK];\n\n // Compute the coordinates of the threads in the block.\n const int warp_id = threadIdx.x / warpSize;\n const int lane_id = threadIdx.x % warpSize;\n\n // Mask for compaction.\n // Same as: asm( \"mov.u32 %0, %%lanemask_lt;\" : \"=r\"(lane_mask_lt) );\n int lane_mask_lt = (1 << lane_id) - 1;\n\n // The current node.\n Quadtree_node &node = nodes[blockIdx.x];\n\n // The number of points in the node.\n int num_points = node.num_points();\n\n float2 center;\n int range_begin, range_end;\n int warp_cnts[4] = {0, 0, 0, 0};\n //\n // 1- Check the number of points and its depth.\n //\n\n // Stop the recursion here. Make sure points[0] contains all the points.\n if (params.depth >= params.max_depth || num_points <= params.min_points_per_node) {\n if (params.point_selector == 1) {\n int it = node.points_begin(), end = node.points_end();\n\n for (it += threadIdx.x; it < end; it += NUM_THREADS_PER_BLOCK)\n if (it < end)\n points[0].set_point(it, points[1].get_point(it));\n }\n\n return;\n }\n\n // Compute the center of the bounding box of the points.\n const Bounding_box &bbox = node.bounding_box();\n\n bbox.compute_center(center);\n\n // Find how many points to give to each warp.\n int num_points_per_warp = max(warpSize, (num_points + NUM_WARPS_PER_BLOCK - 1) / NUM_WARPS_PER_BLOCK);\n\n // Each warp of threads will compute the number of points to move to each\n // quadrant.\n range_begin = node.points_begin() + warp_id * num_points_per_warp;\n range_end = min(range_begin + num_points_per_warp, node.points_end());\n\n //\n // 2- Count the number of points in each child.\n //\n\n // Input points.\n const Points &in_points = points[params.point_selector];\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n // Compute the number of points.\n for (int range_it = range_begin + tile32.thread_rank(); tile32.any(range_it < range_end); range_it += warpSize) {\n // Is it still an active thread?\n bool is_active = range_it < range_end;\n\n // Load the coordinates of the point.\n float2 p = is_active ? in_points.get_point(range_it) : make_float2(0.0f, 0.0f);\n\n // Count top-left points.\n int num_pts = __popc(tile32.ballot(is_active && p.x < center.x && p.y >= center.y));\n warp_cnts[0] += tile32.shfl(num_pts, 0);\n\n // Count top-right points.\n num_pts = __popc(tile32.ballot(is_active && p.x >= center.x && p.y >= center.y));\n warp_cnts[1] += tile32.shfl(num_pts, 0);\n\n // Count bottom-left points.\n num_pts = __popc(tile32.ballot(is_active && p.x < center.x && p.y < center.y));\n warp_cnts[2] += tile32.shfl(num_pts, 0);\n\n // Count bottom-right points.\n num_pts = __popc(tile32.ballot(is_active && p.x >= center.x && p.y < center.y));\n warp_cnts[3] += tile32.shfl(num_pts, 0);\n }\n\n if (tile32.thread_rank() == 0) {\n s_num_pts[0][warp_id] = warp_cnts[0];\n s_num_pts[1][warp_id] = warp_cnts[1];\n s_num_pts[2][warp_id] = warp_cnts[2];\n s_num_pts[3][warp_id] = warp_cnts[3];\n }\n\n // Make sure warps have finished counting.\n cg::sync(cta);\n\n //\n // 3- Scan the warps' results to know the \"global\" numbers.\n //\n\n // First 4 warps scan the numbers of points per child (inclusive scan).\n if (warp_id < 4) {\n int num_pts = tile32.thread_rank() < NUM_WARPS_PER_BLOCK ? s_num_pts[warp_id][tile32.thread_rank()] : 0;\n#pragma unroll\n\n for (int offset = 1; offset < NUM_WARPS_PER_BLOCK; offset *= 2) {\n int n = tile32.shfl_up(num_pts, offset);\n\n if (tile32.thread_rank() >= offset)\n num_pts += n;\n }\n\n if (tile32.thread_rank() < NUM_WARPS_PER_BLOCK)\n s_num_pts[warp_id][tile32.thread_rank()] = num_pts;\n }\n\n cg::sync(cta);\n\n // Compute global offsets.\n if (warp_id == 0) {\n int sum = s_num_pts[0][NUM_WARPS_PER_BLOCK - 1];\n\n for (int row = 1; row < 4; ++row) {\n int tmp = s_num_pts[row][NUM_WARPS_PER_BLOCK - 1];\n cg::sync(tile32);\n\n if (tile32.thread_rank() < NUM_WARPS_PER_BLOCK)\n s_num_pts[row][tile32.thread_rank()] += sum;\n\n cg::sync(tile32);\n sum += tmp;\n }\n }\n\n cg::sync(cta);\n\n // Make the scan exclusive.\n int val = 0;\n if (threadIdx.x < 4 * NUM_WARPS_PER_BLOCK) {\n val = threadIdx.x == 0 ? 0 : smem[threadIdx.x - 1];\n val += node.points_begin();\n }\n\n cg::sync(cta);\n\n if (threadIdx.x < 4 * NUM_WARPS_PER_BLOCK) {\n smem[threadIdx.x] = val;\n }\n\n cg::sync(cta);\n\n //\n // 4- Move points.\n //\n if (!(params.depth >= params.max_depth || num_points <= params.min_points_per_node)) {\n // Output points.\n Points &out_points = points[(params.point_selector + 1) % 2];\n\n warp_cnts[0] = s_num_pts[0][warp_id];\n warp_cnts[1] = s_num_pts[1][warp_id];\n warp_cnts[2] = s_num_pts[2][warp_id];\n warp_cnts[3] = s_num_pts[3][warp_id];\n\n const Points &in_points = points[params.point_selector];\n // Reorder points.\n for (int range_it = range_begin + tile32.thread_rank(); tile32.any(range_it < range_end);\n range_it += warpSize) {\n // Is it still an active thread?\n bool is_active = range_it < range_end;\n\n // Load the coordinates of the point.\n float2 p = is_active ? in_points.get_point(range_it) : make_float2(0.0f, 0.0f);\n\n // Count top-left points.\n bool pred = is_active && p.x < center.x && p.y >= center.y;\n int vote = tile32.ballot(pred);\n int dest = warp_cnts[0] + __popc(vote & lane_mask_lt);\n\n if (pred)\n out_points.set_point(dest, p);\n\n warp_cnts[0] += tile32.shfl(__popc(vote), 0);\n\n // Count top-right points.\n pred = is_active && p.x >= center.x && p.y >= center.y;\n vote = tile32.ballot(pred);\n dest = warp_cnts[1] + __popc(vote & lane_mask_lt);\n\n if (pred)\n out_points.set_point(dest, p);\n\n warp_cnts[1] += tile32.shfl(__popc(vote), 0);\n\n // Count bottom-left points.\n pred = is_active && p.x < center.x && p.y < center.y;\n vote = tile32.ballot(pred);\n dest = warp_cnts[2] + __popc(vote & lane_mask_lt);\n\n if (pred)\n out_points.set_point(dest, p);\n\n warp_cnts[2] += tile32.shfl(__popc(vote), 0);\n\n // Count bottom-right points.\n pred = is_active && p.x >= center.x && p.y < center.y;\n vote = tile32.ballot(pred);\n dest = warp_cnts[3] + __popc(vote & lane_mask_lt);\n\n if (pred)\n out_points.set_point(dest, p);\n\n warp_cnts[3] += tile32.shfl(__popc(vote), 0);\n }\n }\n\n cg::sync(cta);\n\n if (tile32.thread_rank() == 0) {\n s_num_pts[0][warp_id] = warp_cnts[0];\n s_num_pts[1][warp_id] = warp_cnts[1];\n s_num_pts[2][warp_id] = warp_cnts[2];\n s_num_pts[3][warp_id] = warp_cnts[3];\n }\n\n cg::sync(cta);\n\n //\n // 5- Launch new blocks.\n //\n if (!(params.depth >= params.max_depth || num_points <= params.min_points_per_node)) {\n // The last thread launches new blocks.\n if (threadIdx.x == NUM_THREADS_PER_BLOCK - 1) {\n // The children.\n Quadtree_node *children = &nodes[params.num_nodes_at_this_level - (node.id() & ~3)];\n\n // The offsets of the children at their level.\n int child_offset = 4 * node.id();\n\n // Set IDs.\n children[child_offset + 0].set_id(4 * node.id() + 0);\n children[child_offset + 1].set_id(4 * node.id() + 1);\n children[child_offset + 2].set_id(4 * node.id() + 2);\n children[child_offset + 3].set_id(4 * node.id() + 3);\n\n const Bounding_box &bbox = node.bounding_box();\n // Points of the bounding-box.\n const float2 &p_min = bbox.get_min();\n const float2 &p_max = bbox.get_max();\n\n // Set the bounding boxes of the children.\n children[child_offset + 0].set_bounding_box(p_min.x, center.y, center.x,\n p_max.y); // Top-left.\n children[child_offset + 1].set_bounding_box(center.x, center.y, p_max.x,\n p_max.y); // Top-right.\n children[child_offset + 2].set_bounding_box(p_min.x, p_min.y, center.x,\n center.y); // Bottom-left.\n children[child_offset + 3].set_bounding_box(center.x, p_min.y, p_max.x,\n center.y); // Bottom-right.\n\n // Set the ranges of the children.\n\n children[child_offset + 0].set_range(node.points_begin(), s_num_pts[0][warp_id]);\n children[child_offset + 1].set_range(s_num_pts[0][warp_id], s_num_pts[1][warp_id]);\n children[child_offset + 2].set_range(s_num_pts[1][warp_id], s_num_pts[2][warp_id]);\n children[child_offset + 3].set_range(s_num_pts[2][warp_id], s_num_pts[3][warp_id]);\n\n // Launch 4 children.\n build_quadtree_kernel\n <<<4, NUM_THREADS_PER_BLOCK, 4 * NUM_WARPS_PER_BLOCK * sizeof(int)>>>(\n &children[child_offset], points, Parameters(params, true));\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Make sure a Quadtree is properly defined.\n////////////////////////////////////////////////////////////////////////////////\nbool check_quadtree(const Quadtree_node *nodes, int idx, int num_pts, Points *pts, Parameters params)\n{\n const Quadtree_node &node = nodes[idx];\n int num_points = node.num_points();\n\n if (!(params.depth == params.max_depth || num_points <= params.min_points_per_node)) {\n int num_points_in_children = 0;\n\n num_points_in_children += nodes[params.num_nodes_at_this_level + 4 * idx + 0].num_points();\n num_points_in_children += nodes[params.num_nodes_at_this_level + 4 * idx + 1].num_points();\n num_points_in_children += nodes[params.num_nodes_at_this_level + 4 * idx + 2].num_points();\n num_points_in_children += nodes[params.num_nodes_at_this_level + 4 * idx + 3].num_points();\n\n if (num_points_in_children != node.num_points())\n return false;\n\n return check_quadtree(\n &nodes[params.num_nodes_at_this_level], 4 * idx + 0, num_pts, pts, Parameters(params, true))\n && check_quadtree(\n &nodes[params.num_nodes_at_this_level], 4 * idx + 1, num_pts, pts, Parameters(params, true))\n && check_quadtree(\n &nodes[params.num_nodes_at_this_level], 4 * idx + 2, num_pts, pts, Parameters(params, true))\n && check_quadtree(\n &nodes[params.num_nodes_at_this_level], 4 * idx + 3, num_pts, pts, Parameters(params, true));\n }\n\n const Bounding_box &bbox = node.bounding_box();\n\n for (int it = node.points_begin(); it < node.points_end(); ++it) {\n if (it >= num_pts)\n return false;\n\n float2 p = pts->get_point(it);\n\n if (!bbox.contains(p))\n return false;\n }\n\n return true;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Parallel random number generator.\n////////////////////////////////////////////////////////////////////////////////\nstruct Random_generator\n{\n int count;\n\n __host__ __device__ Random_generator()\n : count(0)\n {\n }\n __host__ __device__ unsigned int hash(unsigned int a)\n {\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) ^ (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a + 0xd3a2646c) ^ (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) ^ (a >> 16);\n return a;\n }\n\n __host__ __device__ __forceinline__ cuda::std::tuple operator()()\n {\n#ifdef __CUDA_ARCH__\n unsigned seed = hash(blockIdx.x * blockDim.x + threadIdx.x + count);\n // thrust::generate may call operator() more than once per thread.\n // Hence, increment count by grid size to ensure uniqueness of seed\n count += blockDim.x * gridDim.x;\n#else\n unsigned seed = hash(0);\n#endif\n thrust::default_random_engine rng(seed);\n thrust::random::uniform_real_distribution distrib;\n return cuda::std::make_tuple(distrib(rng), distrib(rng));\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n// Allocate GPU structs, launch kernel and clean up\n////////////////////////////////////////////////////////////////////////////////\nbool cdpQuadtree(int warp_size)\n{\n // Constants to control the algorithm.\n const int num_points = 1024;\n const int max_depth = 8;\n const int min_points_per_node = 16;\n\n // Allocate memory for points.\n thrust::device_vector x_d0(num_points);\n thrust::device_vector x_d1(num_points);\n thrust::device_vector y_d0(num_points);\n thrust::device_vector y_d1(num_points);\n\n // Generate random points.\n Random_generator rnd;\n thrust::generate(\n thrust::make_zip_iterator(x_d0.begin(), y_d0.begin()), thrust::make_zip_iterator(x_d0.end(), y_d0.end()), rnd);\n\n // Host structures to analyze the device ones.\n Points points_init[2];\n points_init[0].set(thrust::raw_pointer_cast(&x_d0[0]), thrust::raw_pointer_cast(&y_d0[0]));\n points_init[1].set(thrust::raw_pointer_cast(&x_d1[0]), thrust::raw_pointer_cast(&y_d1[0]));\n\n // Allocate memory to store points.\n Points *points;\n checkCudaErrors(cudaMalloc((void **)&points, 2 * sizeof(Points)));\n checkCudaErrors(cudaMemcpy(points, points_init, 2 * sizeof(Points), cudaMemcpyHostToDevice));\n\n // We could use a close form...\n int max_nodes = 0;\n\n for (int i = 0, num_nodes_at_level = 1; i < max_depth; ++i, num_nodes_at_level *= 4)\n max_nodes += num_nodes_at_level;\n\n // Allocate memory to store the tree.\n Quadtree_node root;\n root.set_range(0, num_points);\n Quadtree_node *nodes;\n checkCudaErrors(cudaMalloc((void **)&nodes, max_nodes * sizeof(Quadtree_node)));\n checkCudaErrors(cudaMemcpy(nodes, &root, sizeof(Quadtree_node), cudaMemcpyHostToDevice));\n\n // Build the quadtree.\n Parameters params(max_depth, min_points_per_node);\n std::cout << \"Launching CDP kernel to build the quadtree\" << std::endl;\n const int NUM_THREADS_PER_BLOCK = 128; // Do not use less than 128 threads.\n const int NUM_WARPS_PER_BLOCK = NUM_THREADS_PER_BLOCK / warp_size;\n const size_t smem_size = 4 * NUM_WARPS_PER_BLOCK * sizeof(int);\n build_quadtree_kernel<<<1, NUM_THREADS_PER_BLOCK, smem_size>>>(nodes, points, params);\n checkCudaErrors(cudaGetLastError());\n\n // Copy points to CPU.\n thrust::host_vector x_h(x_d0);\n thrust::host_vector y_h(y_d0);\n Points host_points;\n host_points.set(thrust::raw_pointer_cast(&x_h[0]), thrust::raw_pointer_cast(&y_h[0]));\n\n // Copy nodes to CPU.\n Quadtree_node *host_nodes = new Quadtree_node[max_nodes];\n checkCudaErrors(cudaMemcpy(host_nodes, nodes, max_nodes * sizeof(Quadtree_node), cudaMemcpyDeviceToHost));\n\n // Validate the results.\n bool ok = check_quadtree(host_nodes, 0, num_points, &host_points, params);\n std::cout << \"Results: \" << (ok ? \"OK\" : \"FAILED\") << std::endl;\n\n // Free CPU memory.\n delete[] host_nodes;\n\n // Free memory.\n checkCudaErrors(cudaFree(nodes));\n checkCudaErrors(cudaFree(points));\n\n return ok;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Main entry point.\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n // Find/set the device.\n // The test requires an architecture SM35 or greater (CDP capable).\n int cuda_device = findCudaDevice(argc, (const char **)argv);\n cudaDeviceProp deviceProps;\n checkCudaErrors(cudaGetDeviceProperties(&deviceProps, cuda_device));\n int cdpCapable = (deviceProps.major == 3 && deviceProps.minor >= 5) || deviceProps.major >= 4;\n\n printf(\n \"GPU device %s has compute capabilities (SM %d.%d)\\n\", deviceProps.name, deviceProps.major, deviceProps.minor);\n\n if (!cdpCapable) {\n std::cerr << \"cdpQuadTree requires SM 3.5 or higher to use CUDA Dynamic \"\n \"Parallelism. Exiting...\\n\"\n << std::endl;\n exit(EXIT_WAIVED);\n }\n\n bool ok = cdpQuadtree(deviceProps.warpSize);\n\n return (ok ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n// for cuda::std::tuple\n#include \n\nnamespace cg = cooperative_groups;\n#include \n\n////////////////////////////////////////////////////////////////////////////////\n// A structure of 2D points (structure of arrays).\n////////////////////////////////////////////////////////////////////////////////\nclass Points\n{\n float *m_x;\n float *m_y;\n\npublic:\n // Constructor.\n __host__ __device__ Points()\n : m_x(NULL)\n , m_y(NULL)\n {\n }\n\n // Constructor.\n __host__ __device__ Points(float *x, float *y)\n : m_x(x)\n , m_y(y)\n {\n }\n\n // Get a point.\n __host__ __device__ __forceinline__ float2 get_point(int idx) const { return make_float2(m_x[idx], m_y[idx]); }\n\n // Set a point.\n __host__ __device__ __forceinline__ void set_point(int idx, const float2 &p)\n {\n m_x[idx] = p.x;\n m_y[idx] = p.y;\n }\n\n // Set the pointers.\n __host__ __device__ __forceinline__ void set(float *x, float *y)\n {\n m_x = x;\n m_y = y;\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n// A 2D bounding box\n////////////////////////////////////////////////////////////////////////////////\nclass Bounding_box\n{\n // Extreme points of the bounding box.\n float2 m_p_min;\n float2 m_p_max;\n\npublic:\n // Constructor. Create a unit box.\n __host__ __device__ Bounding_box()\n {\n m_p_min = make_float2(0.0f, 0.0f);\n m_p_max = make_float2(1.0f, 1.0f);\n }\n\n // Compute the center of the bounding-box.\n __host__ __device__ void compute_center(float2 ¢er) const\n {\n center.x = 0.5f * (m_p_min.x + m_p_max.x);\n center.y = 0.5f * (m_p_min.y + m_p_max.y);\n }\n\n // The points of the box.\n __host__ __device__ __forceinline__ const float2 &get_max() const { return m_p_max; }\n\n __host__ __device__ __forceinline__ const float2 &get_min() const { return m_p_min; }\n\n // Does a box contain a point.\n __host__ __device__ bool contains(const float2 &p) const\n {\n return p.x >= m_p_min.x && p.x < m_p_max.x && p.y >= m_p_min.y && p.y < m_p_max.y;\n }\n\n // Define the bounding box.\n __host__ __device__ void set(float min_x, float min_y, float max_x, float max_y)\n {\n m_p_min.x = min_x;\n m_p_min.y = min_y;\n m_p_max.x = max_x;\n m_p_max.y = max_y;\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n// A node of a quadree.\n////////////////////////////////////////////////////////////////////////////////\nclass Quadtree_node\n{\n // The identifier of the node.\n int m_id;\n // The bounding box of the tree.\n Bounding_box m_bounding_box;\n // The range of points.\n int m_begin, m_end;\n\npublic:\n // Constructor.\n __host__ __device__ Quadtree_node()\n : m_id(0)\n , m_begin(0)\n , m_end(0)\n {\n }\n\n // The ID of a node at its level.\n __host__ __device__ int id() const { return m_id; }\n\n // The ID of a node at its level.\n __host__ __device__ void set_id(int new_id) { m_id = new_id; }\n\n // The bounding box.\n __host__ __device__ __forceinline__ const Bounding_box &bounding_box() const { return m_bounding_box; }\n\n // Set the bounding box.\n __host__ __device__ __forceinline__ void set_bounding_box(float min_x, float min_y, float max_x, float max_y)\n {\n m_bounding_box.set(min_x, min_y, max_x, max_y);\n }\n\n // The number of points in the tree.\n __host__ __device__ __forceinline__ int num_points() const { return m_end - m_begin; }\n\n // The range of points in the tree.\n __host__ __device__ __forceinline__ int points_begin() const { return m_begin; }\n\n __host__ __device__ __forceinline__ int points_end() const { return m_end; }\n\n // Define the range for that node.\n __host__ __device__ __forceinline__ void set_range(int begin, int end)\n {\n m_begin = begin;\n m_end = end;\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n// Algorithm parameters.\n////////////////////////////////////////////////////////////////////////////////\nstruct Parameters\n{\n // Choose the right set of points to use as in/out.\n int point_selector;\n // The number of nodes at a given level (2^k for level k).\n int num_nodes_at_this_level;\n // The recursion depth.\n int depth;\n // The max value for depth.\n const int max_depth;\n // The minimum number of points in a node to stop recursion.\n const int min_points_per_node;\n\n // Constructor set to default values.\n __host__ __device__ Parameters(int max_depth, int min_points_per_node)\n : point_selector(0)\n , num_nodes_at_this_level(1)\n , depth(0)\n , max_depth(max_depth)\n , min_points_per_node(min_points_per_node)\n {\n }\n\n // Copy constructor. Changes the values for next iteration.\n __host__ __device__ Parameters(const Parameters ¶ms, bool)\n : point_selector((params.point_selector + 1) % 2)\n , num_nodes_at_this_level(4 * params.num_nodes_at_this_level)\n , depth(params.depth + 1)\n , max_depth(params.max_depth)\n , min_points_per_node(params.min_points_per_node)\n {\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n// Build a quadtree on the GPU. Use CUDA Dynamic Parallelism.\n//\n// The algorithm works as follows. The host (CPU) launches one block of\n// NUM_THREADS_PER_BLOCK threads. That block will do the following steps:\n//\n// 1- Check the number of points and its depth.\n//\n// We impose a maximum depth to the tree and a minimum number of points per\n// node. If the maximum depth is exceeded or the minimum number of points is\n// reached. The threads in the block exit.\n//\n// Before exiting, they perform a buffer swap if it is needed. Indeed, the\n// algorithm uses two buffers to permute the points and make sure they are\n// properly distributed in the quadtree. By design we want all points to be\n// in the first buffer of points at the end of the algorithm. It is the reason\n// why we may have to swap the buffer before leavin (if the points are in the\n// 2nd buffer).\n//\n// 2- Count the number of points in each child.\n//\n// If the depth is not too high and the number of points is sufficient, the\n// block has to dispatch the points into four geometrical buckets: Its\n// children. For that purpose, we compute the center of the bounding box and\n// count the number of points in each quadrant.\n//\n// The set of points is divided into sections. Each section is given to a\n// warp of threads (32 threads). Warps use __ballot and __popc intrinsics\n// to count the points. See the Programming Guide for more information about\n// those functions.\n//\n// 3- Scan the warps' results to know the \"global\" numbers.\n//\n// Warps work independently from each other. At the end, each warp knows the\n// number of points in its section. To know the numbers for the block, the\n// block has to run a scan/reduce at the block level. It's a traditional\n// approach. The implementation in that sample is not as optimized as what\n// could be found in fast radix sorts, for example, but it relies on the same\n// idea.\n//\n// 4- Move points.\n//\n// Now that the block knows how many points go in each of its 4 children, it\n// remains to dispatch the points. It is straightforward.\n//\n// 5- Launch new blocks.\n//\n// The block launches four new blocks: One per children. Each of the four blocks\n// will apply the same algorithm.\n////////////////////////////////////////////////////////////////////////////////\ntemplate \n__global__ void build_quadtree_kernel(Quadtree_node *nodes, Points *points, Parameters params)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // The number of warps in a block.\n const int NUM_WARPS_PER_BLOCK = NUM_THREADS_PER_BLOCK / warpSize;\n\n // Shared memory to store the number of points.\n extern __shared__ int smem[];\n\n // s_num_pts[4][NUM_WARPS_PER_BLOCK];\n // Addresses of shared memory.\n volatile int *s_num_pts[4];\n\n for (int i = 0; i < 4; ++i)\n s_num_pts[i] = (volatile int *)&smem[i * NUM_WARPS_PER_BLOCK];\n\n // Compute the coordinates of the threads in the block.\n const int warp_id = threadIdx.x / warpSize;\n const int lane_id = threadIdx.x % warpSize;\n\n // Mask for compaction.\n // Same as: asm( \"mov.u32 %0, %%lanemask_lt;\" : \"=r\"(lane_mask_lt) );\n int lane_mask_lt = (1 << lane_id) - 1;\n\n // The current node.\n Quadtree_node &node = nodes[blockIdx.x];\n\n // The number of points in the node.\n int num_points = node.num_points();\n\n float2 center;\n int range_begin, range_end;\n int warp_cnts[4] = {0, 0, 0, 0};\n //\n // 1- Check the number of points and its depth.\n //\n\n // Stop the recursion here. Make sure points[0] contains all the points.\n if (params.depth >= params.max_depth || num_points <= params.min_points_per_node) {\n if (params.point_selector == 1) {\n int it = node.points_begin(), end = node.points_end();\n\n for (it += threadIdx.x; it < end; it += NUM_THREADS_PER_BLOCK)\n if (it < end)\n points[0].set_point(it, points[1].get_point(it));\n }\n\n return;\n }\n\n // Compute the center of the bounding box of the points.\n const Bounding_box &bbox = node.bounding_box();\n\n bbox.compute_center(center);\n\n // Find how many points to give to each warp.\n int num_points_per_warp = max(warpSize, (num_points + NUM_WARPS_PER_BLOCK - 1) / NUM_WARPS_PER_BLOCK);\n\n // Each warp of threads will compute the number of points to move to each\n // quadrant.\n range_begin = node.points_begin() + warp_id * num_points_per_warp;\n range_end = min(range_begin + num_points_per_warp, node.points_end());\n\n //\n // 2- Count the number of points in each child.\n //\n\n // Input points.\n const Points &in_points = points[params.point_selector];\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n // Compute the number of points.\n for (int range_it = range_begin + tile32.thread_rank(); tile32.any(range_it < range_end); range_it += warpSize) {\n // Is it still an active thread?\n bool is_active = range_it < range_end;\n\n // Load the coordinates of the point.\n float2 p = is_active ? in_points.get_point(range_it) : make_float2(0.0f, 0.0f);\n\n // Count top-left points.\n int num_pts = __popc(tile32.ballot(is_active && p.x < center.x && p.y >= center.y));\n warp_cnts[0] += tile32.shfl(num_pts, 0);\n\n // Count top-right points.\n num_pts = __popc(tile32.ballot(is_active && p.x >= center.x && p.y >= center.y));\n warp_cnts[1] += tile32.shfl(num_pts, 0);\n\n // Count bottom-left points.\n num_pts = __popc(tile32.ballot(is_active && p.x < center.x && p.y < center.y));\n warp_cnts[2] += tile32.shfl(num_pts, 0);\n\n // Count bottom-right points.\n num_pts = __popc(tile32.ballot(is_active && p.x >= center.x && p.y < center.y));\n warp_cnts[3] += tile32.shfl(num_pts, 0);\n }\n\n if (tile32.thread_rank() == 0) {\n s_num_pts[0][warp_id] = warp_cnts[0];\n s_num_pts[1][warp_id] = warp_cnts[1];\n s_num_pts[2][warp_id] = warp_cnts[2];\n s_num_pts[3][warp_id] = warp_cnts[3];\n }\n\n // Make sure warps have finished counting.\n cg::sync(cta);\n\n //\n // 3- Scan the warps' results to know the \"global\" numbers.\n //\n\n // First 4 warps scan the numbers of points per child (inclusive scan).\n if (warp_id < 4) {\n int num_pts = tile32.thread_rank() < NUM_WARPS_PER_BLOCK ? s_num_pts[warp_id][tile32.thread_rank()] : 0;\n#pragma unroll\n\n for (int offset = 1; offset < NUM_WARPS_PER_BLOCK; offset *= 2) {\n int n = tile32.shfl_up(num_pts, offset);\n\n if (tile32.thread_rank() >= offset)\n num_pts += n;\n }\n\n if (tile32.thread_rank() < NUM_WARPS_PER_BLOCK)\n s_num_pts[warp_id][tile32.thread_rank()] = num_pts;\n }\n\n cg::sync(cta);\n\n // Compute global offsets.\n if (warp_id == 0) {\n int sum = s_num_pts[0][NUM_WARPS_PER_BLOCK - 1];\n\n for (int row = 1; row < 4; ++row) {\n int tmp = s_num_pts[row][NUM_WARPS_PER_BLOCK - 1];\n cg::sync(tile32);\n\n if (tile32.thread_rank() < NUM_WARPS_PER_BLOCK)\n s_num_pts[row][tile32.thread_rank()] += sum;\n\n cg::sync(tile32);\n sum += tmp;\n }\n }\n\n cg::sync(cta);\n\n // Make the scan exclusive.\n int val = 0;\n if (threadIdx.x < 4 * NUM_WARPS_PER_BLOCK) {\n val = threadIdx.x == 0 ? 0 : smem[threadIdx.x - 1];\n val += node.points_begin();\n }\n\n cg::sync(cta);\n\n if (threadIdx.x < 4 * NUM_WARPS_PER_BLOCK) {\n smem[threadIdx.x] = val;\n }\n\n cg::sync(cta);\n\n //\n // 4- Move points.\n //\n if (!(params.depth >= params.max_depth || num_points <= params.min_points_per_node)) {\n // Output points.\n Points &out_points = points[(params.point_selector + 1) % 2];\n\n warp_cnts[0] = s_num_pts[0][warp_id];\n warp_cnts[1] = s_num_pts[1][warp_id];\n warp_cnts[2] = s_num_pts[2][warp_id];\n warp_cnts[3] = s_num_pts[3][warp_id];\n\n const Points &in_points = points[params.point_selector];\n // Reorder points.\n for (int range_it = range_begin + tile32.thread_rank(); tile32.any(range_it < range_end);\n range_it += warpSize) {\n // Is it still an active thread?\n bool is_active = range_it < range_end;\n\n // Load the coordinates of the point.\n float2 p = is_active ? in_points.get_point(range_it) : make_float2(0.0f, 0.0f);\n\n // Count top-left points.\n bool pred = is_active && p.x < center.x && p.y >= center.y;\n int vote = tile32.ballot(pred);\n int dest = warp_cnts[0] + __popc(vote & lane_mask_lt);\n\n if (pred)\n out_points.set_point(dest, p);\n\n warp_cnts[0] += tile32.shfl(__popc(vote), 0);\n\n // Count top-right points.\n pred = is_active && p.x >= center.x && p.y >= center.y;\n vote = tile32.ballot(pred);\n dest = warp_cnts[1] + __popc(vote & lane_mask_lt);\n\n if (pred)\n out_points.set_point(dest, p);\n\n warp_cnts[1] += tile32.shfl(__popc(vote), 0);\n\n // Count bottom-left points.\n pred = is_active && p.x < center.x && p.y < center.y;\n vote = tile32.ballot(pred);\n dest = warp_cnts[2] + __popc(vote & lane_mask_lt);\n\n if (pred)\n out_points.set_point(dest, p);\n\n warp_cnts[2] += tile32.shfl(__popc(vote), 0);\n\n // Count bottom-right points.\n pred = is_active && p.x >= center.x && p.y < center.y;\n vote = tile32.ballot(pred);\n dest = warp_cnts[3] + __popc(vote & lane_mask_lt);\n\n if (pred)\n out_points.set_point(dest, p);\n\n warp_cnts[3] += tile32.shfl(__popc(vote), 0);\n }\n }\n\n cg::sync(cta);\n\n if (tile32.thread_rank() == 0) {\n s_num_pts[0][warp_id] = warp_cnts[0];\n s_num_pts[1][warp_id] = warp_cnts[1];\n s_num_pts[2][warp_id] = warp_cnts[2];\n s_num_pts[3][warp_id] = warp_cnts[3];\n }\n\n cg::sync(cta);\n\n //\n // 5- Launch new blocks.\n //\n if (!(params.depth >= params.max_depth || num_points <= params.min_points_per_node)) {\n // The last thread launches new blocks.\n if (threadIdx.x == NUM_THREADS_PER_BLOCK - 1) {\n // The children.\n Quadtree_node *children = &nodes[params.num_nodes_at_this_level - (node.id() & ~3)];\n\n // The offsets of the children at their level.\n int child_offset = 4 * node.id();\n\n // Set IDs.\n children[child_offset + 0].set_id(4 * node.id() + 0);\n children[child_offset + 1].set_id(4 * node.id() + 1);\n children[child_offset + 2].set_id(4 * node.id() + 2);\n children[child_offset + 3].set_id(4 * node.id() + 3);\n\n const Bounding_box &bbox = node.bounding_box();\n // Points of the bounding-box.\n const float2 &p_min = bbox.get_min();\n const float2 &p_max = bbox.get_max();\n\n // Set the bounding boxes of the children.\n children[child_offset + 0].set_bounding_box(p_min.x, center.y, center.x,\n p_max.y); // Top-left.\n children[child_offset + 1].set_bounding_box(center.x, center.y, p_max.x,\n p_max.y); // Top-right.\n children[child_offset + 2].set_bounding_box(p_min.x, p_min.y, center.x,\n center.y); // Bottom-left.\n children[child_offset + 3].set_bounding_box(center.x, p_min.y, p_max.x,\n center.y); // Bottom-right.\n\n // Set the ranges of the children.\n\n children[child_offset + 0].set_range(node.points_begin(), s_num_pts[0][warp_id]);\n children[child_offset + 1].set_range(s_num_pts[0][warp_id], s_num_pts[1][warp_id]);\n children[child_offset + 2].set_range(s_num_pts[1][warp_id], s_num_pts[2][warp_id]);\n children[child_offset + 3].set_range(s_num_pts[2][warp_id], s_num_pts[3][warp_id]);\n\n // Launch 4 children.\n build_quadtree_kernel\n <<<4, NUM_THREADS_PER_BLOCK, 4 * NUM_WARPS_PER_BLOCK * sizeof(int)>>>(\n &children[child_offset], points, Parameters(params, true));\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Make sure a Quadtree is properly defined.\n////////////////////////////////////////////////////////////////////////////////\nbool check_quadtree(const Quadtree_node *nodes, int idx, int num_pts, Points *pts, Parameters params)\n{\n const Quadtree_node &node = nodes[idx];\n int num_points = node.num_points();\n\n if (!(params.depth == params.max_depth || num_points <= params.min_points_per_node)) {\n int num_points_in_children = 0;\n\n num_points_in_children += nodes[params.num_nodes_at_this_level + 4 * idx + 0].num_points();\n num_points_in_children += nodes[params.num_nodes_at_this_level + 4 * idx + 1].num_points();\n num_points_in_children += nodes[params.num_nodes_at_this_level + 4 * idx + 2].num_points();\n num_points_in_children += nodes[params.num_nodes_at_this_level + 4 * idx + 3].num_points();\n\n if (num_points_in_children != node.num_points())\n return false;\n\n return check_quadtree(\n &nodes[params.num_nodes_at_this_level], 4 * idx + 0, num_pts, pts, Parameters(params, true))\n && check_quadtree(\n &nodes[params.num_nodes_at_this_level], 4 * idx + 1, num_pts, pts, Parameters(params, true))\n && check_quadtree(\n &nodes[params.num_nodes_at_this_level], 4 * idx + 2, num_pts, pts, Parameters(params, true))\n && check_quadtree(\n &nodes[params.num_nodes_at_this_level], 4 * idx + 3, num_pts, pts, Parameters(params, true));\n }\n\n const Bounding_box &bbox = node.bounding_box();\n\n for (int it = node.points_begin(); it < node.points_end(); ++it) {\n if (it >= num_pts)\n return false;\n\n float2 p = pts->get_point(it);\n\n if (!bbox.contains(p))\n return false;\n }\n\n return true;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Parallel random number generator.\n////////////////////////////////////////////////////////////////////////////////\nstruct Random_generator\n{\n int count;\n\n __host__ __device__ Random_generator()\n : count(0)\n {\n }\n __host__ __device__ unsigned int hash(unsigned int a)\n {\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) ^ (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a + 0xd3a2646c) ^ (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) ^ (a >> 16);\n return a;\n }\n\n __host__ __device__ __forceinline__ cuda::std::tuple operator()()\n {\n#ifdef __CUDA_ARCH__\n unsigned seed = hash(blockIdx.x * blockDim.x + threadIdx.x + count);\n // thrust::generate may call operator() more than once per thread.\n // Hence, increment count by grid size to ensure uniqueness of seed\n count += blockDim.x * gridDim.x;\n#else\n unsigned seed = hash(0);\n#endif\n thrust::default_random_engine rng(seed);\n thrust::random::uniform_real_distribution distrib;\n return cuda::std::make_tuple(distrib(rng), distrib(rng));\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n// Allocate GPU structs, launch kernel and clean up\n////////////////////////////////////////////////////////////////////////////////\nbool cdpQuadtree(int warp_size)\n{\n // Constants to control the algorithm.\n const int num_points = 1024;\n const int max_depth = 8;\n const int min_points_per_node = 16;\n\n // Allocate memory for points.\n thrust::device_vector x_d0(num_points);\n thrust::device_vector x_d1(num_points);\n thrust::device_vector y_d0(num_points);\n thrust::device_vector y_d1(num_points);\n\n // Generate random points.\n Random_generator rnd;\n thrust::generate(\n thrust::make_zip_iterator(x_d0.begin(), y_d0.begin()), thrust::make_zip_iterator(x_d0.end(), y_d0.end()), rnd);\n\n // Host structures to analyze the device ones.\n Points points_init[2];\n points_init[0].set(thrust::raw_pointer_cast(&x_d0[0]), thrust::raw_pointer_cast(&y_d0[0]));\n points_init[1].set(thrust::raw_pointer_cast(&x_d1[0]), thrust::raw_pointer_cast(&y_d1[0]));\n\n // Allocate memory to store points.\n Points *points;\n checkCudaErrors(hipMalloc((void **)&points, 2 * sizeof(Points)));\n checkCudaErrors(hipMemcpy(points, points_init, 2 * sizeof(Points), hipMemcpyHostToDevice));\n\n // We could use a close form...\n int max_nodes = 0;\n\n for (int i = 0, num_nodes_at_level = 1; i < max_depth; ++i, num_nodes_at_level *= 4)\n max_nodes += num_nodes_at_level;\n\n // Allocate memory to store the tree.\n Quadtree_node root;\n root.set_range(0, num_points);\n Quadtree_node *nodes;\n checkCudaErrors(hipMalloc((void **)&nodes, max_nodes * sizeof(Quadtree_node)));\n checkCudaErrors(hipMemcpy(nodes, &root, sizeof(Quadtree_node), hipMemcpyHostToDevice));\n\n // Build the quadtree.\n Parameters params(max_depth, min_points_per_node);\n std::cout << \"Launching CDP kernel to build the quadtree\" << std::endl;\n const int NUM_THREADS_PER_BLOCK = 128; // Do not use less than 128 threads.\n const int NUM_WARPS_PER_BLOCK = NUM_THREADS_PER_BLOCK / warp_size;\n const size_t smem_size = 4 * NUM_WARPS_PER_BLOCK * sizeof(int);\n build_quadtree_kernel<<<1, NUM_THREADS_PER_BLOCK, smem_size>>>(nodes, points, params);\n checkCudaErrors(hipGetLastError());\n\n // Copy points to CPU.\n thrust::host_vector x_h(x_d0);\n thrust::host_vector y_h(y_d0);\n Points host_points;\n host_points.set(thrust::raw_pointer_cast(&x_h[0]), thrust::raw_pointer_cast(&y_h[0]));\n\n // Copy nodes to CPU.\n Quadtree_node *host_nodes = new Quadtree_node[max_nodes];\n checkCudaErrors(hipMemcpy(host_nodes, nodes, max_nodes * sizeof(Quadtree_node), hipMemcpyDeviceToHost));\n\n // Validate the results.\n bool ok = check_quadtree(host_nodes, 0, num_points, &host_points, params);\n std::cout << \"Results: \" << (ok ? \"OK\" : \"FAILED\") << std::endl;\n\n // Free CPU memory.\n delete[] host_nodes;\n\n // Free memory.\n checkCudaErrors(hipFree(nodes));\n checkCudaErrors(hipFree(points));\n\n return ok;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Main entry point.\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n // Find/set the device.\n // The test requires an architecture SM35 or greater (CDP capable).\n int cuda_device = findCudaDevice(argc, (const char **)argv);\n hipDeviceProp_t deviceProps;\n checkCudaErrors(hipGetDeviceProperties(&deviceProps, cuda_device));\n int cdpCapable = (deviceProps.major == 3 && deviceProps.minor >= 5) || deviceProps.major >= 4;\n\n printf(\n \"GPU device %s has compute capabilities (SM %d.%d)\\n\", deviceProps.name, deviceProps.major, deviceProps.minor);\n\n if (!cdpCapable) {\n std::cerr << \"cdpQuadTree requires SM 3.5 or higher to use CUDA Dynamic \"\n \"Parallelism. Exiting...\\n\"\n << std::endl;\n exit(EXIT_WAIVED);\n }\n\n bool ok = cdpQuadtree(deviceProps.warpSize);\n\n return (ok ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/cdpQuadtree/cdpQuadtree.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// System includes\n#include \n#include \n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#define THREADS_PER_BLOCK 512\n#define ALLOWABLE_VARIANCE 1.e-6f\n#define NUM_ELEMENTS 8000000\n\n// Stores the square of each input element in output array\n__global__ void squareArray(const float *input, float *output, int numElements)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx < numElements) {\n output[idx] = input[idx] * input[idx];\n }\n}\n\n// Stores the negative of each input element in output array\n__global__ void negateArray(const float *input, float *output, int numElements)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx < numElements) {\n output[idx] = input[idx] * -1;\n }\n}\n\nstruct negSquareArrays\n{\n float *input;\n float *square;\n float *negSquare;\n int numElements;\n size_t bytes;\n size_t numBlocks;\n};\n\nvoid fillRandomly(float *array, int numElements)\n{\n for (int n = 0; n < numElements; n++) {\n array[n] = rand() / (float)RAND_MAX;\n }\n}\n\nvoid resetOutputArrays(negSquareArrays *hostArrays)\n{\n fillRandomly(hostArrays->square, hostArrays->numElements);\n fillRandomly(hostArrays->negSquare, hostArrays->numElements);\n}\n\nvoid prepareHostArrays(negSquareArrays *hostArrays)\n{\n hostArrays->numElements = NUM_ELEMENTS;\n size_t bytes = hostArrays->numElements * sizeof(float);\n\n size_t numBlocks = hostArrays->numElements / (size_t)THREADS_PER_BLOCK;\n if ((numBlocks % (size_t)THREADS_PER_BLOCK) != 0) {\n numBlocks++;\n }\n\n hostArrays->input = (float *)malloc(bytes);\n hostArrays->square = (float *)malloc(bytes);\n hostArrays->negSquare = (float *)malloc(bytes);\n hostArrays->bytes = bytes;\n hostArrays->numBlocks = numBlocks;\n\n fillRandomly(hostArrays->input, hostArrays->numElements);\n fillRandomly(hostArrays->square, hostArrays->numElements);\n fillRandomly(hostArrays->negSquare, hostArrays->numElements);\n}\n\nvoid createFreeGraph(cudaGraphExec_t *graphExec, float *dPtr)\n{\n cudaGraph_t graph;\n cudaGraphNode_t freeNode;\n\n checkCudaErrors(cudaGraphCreate(&graph, 0));\n\n checkCudaErrors(cudaGraphAddMemFreeNode(&freeNode, graph, NULL, 0, (void *)dPtr));\n\n checkCudaErrors(cudaGraphInstantiate(graphExec, graph, NULL, NULL, 0));\n checkCudaErrors(cudaGraphDestroy(graph));\n}\n\n/**\n * Demonstrates explicitly creating a CUDA graph including memory nodes.\n * createNegateSquaresGraphWithStreamCapture constructs an equivalent graph\n * using stream capture.\n *\n * If d_negSquare_out is non null, then:\n * 1) d_negSquare will not be freed;\n * 2) the value of d_negSquare_out will be set to d_negSquare.\n *\n * Diagram of the graph constructed by createNegateSquaresGraphExplicitly:\n *\n * alloc d_input\n * |\n * alloc d_square\n * |\n * Memcpy a to device\n * |\n * launch kernel squareArray ------->---- Memcpy d_square to host\n * | |\n * free d_input |\n * | |\n * allocate d_negSquare |\n * | |\n * launch kernel negateArray -------->--- free d_square\n * |\n * Memcpy d_negSquare to host\n * |\n * free d_negSquare\n */\nvoid createNegateSquaresGraphExplicitly(cudaGraphExec_t *graphExec,\n int device,\n negSquareArrays *hostArrays,\n float **d_negSquare_out = NULL)\n{\n // Array buffers on device\n float *d_input, *d_square, *d_negSquare;\n\n // Memory allocation parameters\n cudaMemAllocNodeParams allocParams;\n memset(&allocParams, 0, sizeof(allocParams));\n allocParams.bytesize = hostArrays->bytes;\n allocParams.poolProps.allocType = cudaMemAllocationTypePinned;\n allocParams.poolProps.location.id = device;\n allocParams.poolProps.location.type = cudaMemLocationTypeDevice;\n\n // Kernel launch parameters\n cudaKernelNodeParams kernelNodeParams = {0};\n kernelNodeParams.gridDim = dim3(hostArrays->numBlocks, 1, 1);\n kernelNodeParams.blockDim = dim3(THREADS_PER_BLOCK, 1, 1);\n kernelNodeParams.sharedMemBytes = 0;\n kernelNodeParams.extra = NULL;\n\n cudaGraph_t graph;\n cudaGraphNode_t allocNodeInput, allocNodeSquare, allocNodeNegSquare;\n cudaGraphNode_t copyNodeInput, copyNodeSquare, copyNodeNegSquare;\n cudaGraphNode_t squareKernelNode, negateKernelNode;\n cudaGraphNode_t freeNodeInput, freeNodeSquare;\n\n // Buffer for storing graph node dependencies\n std::vector nodeDependencies;\n\n checkCudaErrors(cudaGraphCreate(&graph, 0));\n\n checkCudaErrors(cudaGraphAddMemAllocNode(&allocNodeInput, graph, NULL, 0, &allocParams));\n d_input = (float *)allocParams.dptr;\n\n // To keep the graph structure simple (fewer branching dependencies),\n // allocNodeSquare should depend on allocNodeInput\n checkCudaErrors(cudaGraphAddMemAllocNode(&allocNodeSquare, graph, &allocNodeInput, 1, &allocParams));\n d_square = (float *)allocParams.dptr;\n\n // copyNodeInput needs to depend on allocNodeInput because copyNodeInput\n // writes to d_input. It does so here indirectly through allocNodeSquare.\n checkCudaErrors(cudaGraphAddMemcpyNode1D(©NodeInput,\n graph,\n &allocNodeSquare,\n 1,\n d_input,\n hostArrays->input,\n hostArrays->bytes,\n cudaMemcpyHostToDevice));\n\n void *squareKernelArgs[3] = {(void *)&d_input, (void *)&d_square, (void *)&(hostArrays->numElements)};\n kernelNodeParams.func = (void *)squareArray;\n kernelNodeParams.kernelParams = (void **)squareKernelArgs;\n\n // Square kernel depends on copyNodeInput to ensure all data is on the device\n // before kernel launch.\n checkCudaErrors(cudaGraphAddKernelNode(&squareKernelNode, graph, ©NodeInput, 1, &kernelNodeParams));\n\n checkCudaErrors(cudaGraphAddMemcpyNode1D(©NodeSquare,\n graph,\n &squareKernelNode,\n 1,\n hostArrays->square,\n d_square,\n hostArrays->bytes,\n cudaMemcpyDeviceToHost));\n\n // Free of d_input depends on the square kernel to ensure that d_input is not\n // freed while being read by the kernel. It also depends on the alloc of\n // d_input via squareKernelNode > copyNodeInput > allocNodeSquare >\n // allocNodeInput.\n checkCudaErrors(cudaGraphAddMemFreeNode(&freeNodeInput, graph, &squareKernelNode, 1, d_input));\n\n // Allocation of C depends on free of A so CUDA can reuse the virtual address.\n checkCudaErrors(cudaGraphAddMemAllocNode(&allocNodeNegSquare, graph, &freeNodeInput, 1, &allocParams));\n d_negSquare = (float *)allocParams.dptr;\n\n if (d_negSquare == d_input) {\n printf(\"Check verified that d_negSquare and d_input share a virtual \"\n \"address.\\n\");\n }\n\n void *negateKernelArgs[3] = {(void *)&d_square, (void *)&d_negSquare, (void *)&(hostArrays->numElements)};\n kernelNodeParams.func = (void *)negateArray;\n kernelNodeParams.kernelParams = (void **)negateKernelArgs;\n\n checkCudaErrors(cudaGraphAddKernelNode(&negateKernelNode, graph, &allocNodeNegSquare, 1, &kernelNodeParams));\n\n nodeDependencies.push_back(copyNodeSquare);\n nodeDependencies.push_back(negateKernelNode);\n checkCudaErrors(\n cudaGraphAddMemFreeNode(&freeNodeSquare, graph, nodeDependencies.data(), nodeDependencies.size(), d_square));\n nodeDependencies.clear();\n\n checkCudaErrors(cudaGraphAddMemcpyNode1D(©NodeNegSquare,\n graph,\n &negateKernelNode,\n 1,\n hostArrays->negSquare,\n d_negSquare,\n hostArrays->bytes,\n cudaMemcpyDeviceToHost));\n\n if (d_negSquare_out == NULL) {\n cudaGraphNode_t freeNodeNegSquare;\n checkCudaErrors(cudaGraphAddMemFreeNode(&freeNodeNegSquare, graph, ©NodeNegSquare, 1, d_negSquare));\n }\n else {\n *d_negSquare_out = d_negSquare;\n }\n\n checkCudaErrors(cudaGraphInstantiate(graphExec, graph, NULL, NULL, 0));\n checkCudaErrors(cudaGraphDestroy(graph));\n}\n\n/**\n * Adds work to a CUDA stream which negates the square of values in the input\n * array.\n *\n * If d_negSquare_out is non null, then:\n * 1) d_negSquare will not be freed;\n * 2) the value of d_negSquare_out will be set to d_negSquare.\n *\n * Diagram of the stream operations in doNegateSquaresInStream\n * ---------------------------------------------------------------------\n * | STREAM | STREAM2 |\n * ---------------------------------------------------------------------\n *\n * alloc d_input\n * |\n * alloc d_square\n * |\n * Memcpy a to device\n * |\n * launch kernel squareArray\n * |\n * record squareKernelCompleteEvent -->-- wait squareKernelCompleteEvent\n * | |\n * free d_input |\n * | |\n * allocate d_negSquare Memcpy d_square to host\n * | |\n * launch kernel negateArray |\n * | |\n * record negateKernelCompleteEvent -->-- wait negateKernelCompleteEvent\n * | |\n * Memcpy d_negSquare to host |\n * | free d_square\n * free d_negSquare |\n * | |\n * wait squareFreeEvent --------------<---- record squareFreeEvent\n */\nvoid doNegateSquaresInStream(cudaStream_t stream1, negSquareArrays *hostArrays, float **d_negSquare_out = NULL)\n{\n float *d_input, *d_square, *d_negSquare;\n cudaStream_t stream2;\n cudaEvent_t squareKernelCompleteEvent, negateKernelCompleteEvent, squareFreeEvent;\n\n checkCudaErrors(cudaStreamCreateWithFlags(&stream2, cudaStreamNonBlocking));\n\n checkCudaErrors(cudaEventCreate(&squareKernelCompleteEvent));\n checkCudaErrors(cudaEventCreate(&negateKernelCompleteEvent));\n checkCudaErrors(cudaEventCreate(&squareFreeEvent));\n\n // Virtual addresses are assigned synchronously when cudaMallocAsync is\n // called, thus there is no performace benefit gained by separating the\n // allocations into two streams.\n checkCudaErrors(cudaMallocAsync(&d_input, hostArrays->bytes, stream1));\n checkCudaErrors(cudaMallocAsync(&d_square, hostArrays->bytes, stream1));\n\n checkCudaErrors(cudaMemcpyAsync(d_input, hostArrays->input, hostArrays->bytes, cudaMemcpyHostToDevice, stream1));\n squareArray<<numBlocks, THREADS_PER_BLOCK, 0, stream1>>>(d_input, d_square, hostArrays->numElements);\n checkCudaErrors(cudaEventRecord(squareKernelCompleteEvent, stream1));\n\n checkCudaErrors(cudaStreamWaitEvent(stream2, squareKernelCompleteEvent, 0));\n checkCudaErrors(cudaMemcpyAsync(hostArrays->square, d_square, hostArrays->bytes, cudaMemcpyDeviceToHost, stream2));\n\n checkCudaErrors(cudaFreeAsync(d_input, stream1));\n checkCudaErrors(cudaMallocAsync(&d_negSquare, hostArrays->bytes, stream1));\n negateArray<<numBlocks, THREADS_PER_BLOCK, 0, stream1>>>(\n d_square, d_negSquare, hostArrays->numElements);\n checkCudaErrors(cudaEventRecord(negateKernelCompleteEvent, stream1));\n checkCudaErrors(\n cudaMemcpyAsync(hostArrays->negSquare, d_negSquare, hostArrays->bytes, cudaMemcpyDeviceToHost, stream1));\n if (d_negSquare_out == NULL) {\n checkCudaErrors(cudaFreeAsync(d_negSquare, stream1));\n }\n else {\n *d_negSquare_out = d_negSquare;\n }\n\n checkCudaErrors(cudaStreamWaitEvent(stream2, negateKernelCompleteEvent, 0));\n checkCudaErrors(cudaFreeAsync(d_square, stream2));\n checkCudaErrors(cudaEventRecord(squareFreeEvent, stream2));\n\n checkCudaErrors(cudaStreamWaitEvent(stream1, squareFreeEvent, 0));\n\n checkCudaErrors(cudaStreamDestroy(stream2));\n checkCudaErrors(cudaEventDestroy(squareKernelCompleteEvent));\n checkCudaErrors(cudaEventDestroy(negateKernelCompleteEvent));\n checkCudaErrors(cudaEventDestroy(squareFreeEvent));\n}\n\n/**\n * Demonstrates creating a CUDA graph including memory nodes using stream\n * capture. createNegateSquaresGraphExplicitly constructs an equivalent graph\n * without stream capture.\n */\nvoid createNegateSquaresGraphWithStreamCapture(cudaGraphExec_t *graphExec,\n negSquareArrays *hostArrays,\n float **d_negSquare_out = NULL)\n{\n cudaGraph_t graph;\n cudaStream_t stream;\n\n checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));\n\n checkCudaErrors(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal));\n doNegateSquaresInStream(stream, hostArrays, d_negSquare_out);\n checkCudaErrors(cudaStreamEndCapture(stream, &graph));\n\n checkCudaErrors(cudaGraphInstantiate(graphExec, graph, NULL, NULL, 0));\n checkCudaErrors(cudaStreamDestroy(stream));\n checkCudaErrors(cudaGraphDestroy(graph));\n}\n\nvoid prepareRefArrays(negSquareArrays *hostArrays, negSquareArrays *deviceRefArrays, bool **foundValidationFailure)\n{\n deviceRefArrays->bytes = hostArrays->bytes;\n deviceRefArrays->numElements = hostArrays->numElements;\n\n for (int i = 0; i < hostArrays->numElements; i++) {\n hostArrays->square[i] = hostArrays->input[i] * hostArrays->input[i];\n hostArrays->negSquare[i] = hostArrays->square[i] * -1;\n }\n\n checkCudaErrors(cudaMalloc((void **)&deviceRefArrays->negSquare, deviceRefArrays->bytes));\n checkCudaErrors(\n cudaMemcpy(deviceRefArrays->negSquare, hostArrays->negSquare, hostArrays->bytes, cudaMemcpyHostToDevice));\n\n checkCudaErrors(cudaMallocManaged((void **)foundValidationFailure, sizeof(bool)));\n}\n\nint checkValidationFailure(bool *foundValidationFailure)\n{\n if (*foundValidationFailure) {\n printf(\"Validation FAILURE!\\n\\n\");\n *foundValidationFailure = false;\n return EXIT_FAILURE;\n }\n else {\n printf(\"Validation PASSED!\\n\\n\");\n return EXIT_SUCCESS;\n }\n}\n\n__global__ void validateGPU(float *d_negSquare, negSquareArrays devRefArrays, bool *foundValidationFailure)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n float ref, diff;\n\n if (idx < devRefArrays.numElements) {\n ref = devRefArrays.negSquare[idx];\n diff = d_negSquare[idx] - ref;\n diff *= diff;\n ref *= ref;\n if (diff / ref > ALLOWABLE_VARIANCE) {\n *foundValidationFailure = true;\n }\n }\n}\n\nvoid validateHost(negSquareArrays *hostArrays, bool *foundValidationFailure)\n{\n float ref, diff;\n\n for (int i = 0; i < hostArrays->numElements; i++) {\n ref = hostArrays->input[i] * hostArrays->input[i] * -1;\n diff = hostArrays->negSquare[i] - ref;\n diff *= diff;\n ref *= ref;\n if (diff / ref > ALLOWABLE_VARIANCE) {\n *foundValidationFailure = true;\n }\n }\n}\n\nint main(int argc, char **argv)\n{\n negSquareArrays hostArrays, deviceRefArrays;\n cudaStream_t stream;\n cudaGraphExec_t graphExec, graphExecFreeC;\n\n // Declare pointers for GPU buffers\n float *d_negSquare = NULL;\n bool *foundValidationFailure = NULL;\n\n srand(time(0));\n int device = findCudaDevice(argc, (const char **)argv);\n\n int driverVersion = 0;\n int deviceSupportsMemoryPools = 0;\n\n cudaDriverGetVersion(&driverVersion);\n printf(\"Driver version is: %d.%d\\n\", driverVersion / 1000, (driverVersion % 100) / 10);\n\n if (driverVersion < 11040) {\n printf(\"Waiving execution as driver does not support Graph Memory Nodes\\n\");\n exit(EXIT_WAIVED);\n }\n\n cudaDeviceGetAttribute(&deviceSupportsMemoryPools, cudaDevAttrMemoryPoolsSupported, device);\n if (!deviceSupportsMemoryPools) {\n printf(\"Waiving execution as device does not support Memory Pools\\n\");\n exit(EXIT_WAIVED);\n }\n else {\n printf(\"Setting up sample.\\n\");\n }\n\n prepareHostArrays(&hostArrays);\n prepareRefArrays(&hostArrays, &deviceRefArrays, &foundValidationFailure);\n checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));\n printf(\"Setup complete.\\n\\n\");\n\n printf(\"Running negateSquares in a stream.\\n\");\n doNegateSquaresInStream(stream, &hostArrays);\n checkCudaErrors(cudaStreamSynchronize(stream));\n printf(\"Validating negateSquares in a stream...\\n\");\n validateHost(&hostArrays, foundValidationFailure);\n checkValidationFailure(foundValidationFailure);\n resetOutputArrays(&hostArrays);\n\n printf(\"Running negateSquares in a stream-captured graph.\\n\");\n createNegateSquaresGraphWithStreamCapture(&graphExec, &hostArrays);\n checkCudaErrors(cudaGraphLaunch(graphExec, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n printf(\"Validating negateSquares in a stream-captured graph...\\n\");\n validateHost(&hostArrays, foundValidationFailure);\n checkValidationFailure(foundValidationFailure);\n resetOutputArrays(&hostArrays);\n\n printf(\"Running negateSquares in an explicitly constructed graph.\\n\");\n createNegateSquaresGraphExplicitly(&graphExec, device, &hostArrays);\n checkCudaErrors(cudaGraphLaunch(graphExec, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n printf(\"Validating negateSquares in an explicitly constructed graph...\\n\");\n validateHost(&hostArrays, foundValidationFailure);\n checkValidationFailure(foundValidationFailure);\n resetOutputArrays(&hostArrays);\n\n // Each of the three examples below free d_negSquare outside the graph. As\n // demonstrated by validateGPU, d_negSquare can be accessed by outside the\n // graph before d_negSquare is freed.\n\n printf(\"Running negateSquares with d_negSquare freed outside the stream.\\n\");\n createNegateSquaresGraphExplicitly(&graphExec, device, &hostArrays, &d_negSquare);\n checkCudaErrors(cudaGraphLaunch(graphExec, stream));\n validateGPU<<>>(\n d_negSquare, deviceRefArrays, foundValidationFailure);\n // Since cudaFree is synchronous, the stream must synchronize before freeing\n // d_negSquare to ensure d_negSquare no longer being accessed.\n checkCudaErrors(cudaStreamSynchronize(stream));\n checkCudaErrors(cudaFree(d_negSquare));\n printf(\"Validating negateSquares with d_negSquare freed outside the \"\n \"stream...\\n\");\n validateHost(&hostArrays, foundValidationFailure);\n checkValidationFailure(foundValidationFailure);\n resetOutputArrays(&hostArrays);\n\n printf(\"Running negateSquares with d_negSquare freed outside the graph.\\n\");\n checkCudaErrors(cudaGraphLaunch(graphExec, stream));\n validateGPU<<>>(\n d_negSquare, deviceRefArrays, foundValidationFailure);\n checkCudaErrors(cudaFreeAsync(d_negSquare, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n printf(\"Validating negateSquares with d_negSquare freed outside the graph...\\n\");\n checkValidationFailure(foundValidationFailure);\n resetOutputArrays(&hostArrays);\n\n printf(\"Running negateSquares with d_negSquare freed in a different graph.\\n\");\n createFreeGraph(&graphExecFreeC, d_negSquare);\n checkCudaErrors(cudaGraphLaunch(graphExec, stream));\n validateGPU<<>>(\n d_negSquare, deviceRefArrays, foundValidationFailure);\n checkCudaErrors(cudaGraphLaunch(graphExecFreeC, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n printf(\"Validating negateSquares with d_negSquare freed in a different \"\n \"graph...\\n\");\n checkValidationFailure(foundValidationFailure);\n\n printf(\"Cleaning up sample.\\n\");\n checkCudaErrors(cudaGraphExecDestroy(graphExec));\n checkCudaErrors(cudaGraphExecDestroy(graphExecFreeC));\n checkCudaErrors(cudaStreamDestroy(stream));\n checkCudaErrors(cudaFree(foundValidationFailure));\n checkCudaErrors(cudaFree(deviceRefArrays.negSquare));\n free(hostArrays.input);\n free(hostArrays.square);\n free(hostArrays.negSquare);\n printf(\"Cleanup complete. Exiting sample.\\n\");\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// System includes\n#include \n#include \n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#define THREADS_PER_BLOCK 512\n#define ALLOWABLE_VARIANCE 1.e-6f\n#define NUM_ELEMENTS 8000000\n\n// Stores the square of each input element in output array\n__global__ void squareArray(const float *input, float *output, int numElements)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx < numElements) {\n output[idx] = input[idx] * input[idx];\n }\n}\n\n// Stores the negative of each input element in output array\n__global__ void negateArray(const float *input, float *output, int numElements)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx < numElements) {\n output[idx] = input[idx] * -1;\n }\n}\n\nstruct negSquareArrays\n{\n float *input;\n float *square;\n float *negSquare;\n int numElements;\n size_t bytes;\n size_t numBlocks;\n};\n\nvoid fillRandomly(float *array, int numElements)\n{\n for (int n = 0; n < numElements; n++) {\n array[n] = rand() / (float)RAND_MAX;\n }\n}\n\nvoid resetOutputArrays(negSquareArrays *hostArrays)\n{\n fillRandomly(hostArrays->square, hostArrays->numElements);\n fillRandomly(hostArrays->negSquare, hostArrays->numElements);\n}\n\nvoid prepareHostArrays(negSquareArrays *hostArrays)\n{\n hostArrays->numElements = NUM_ELEMENTS;\n size_t bytes = hostArrays->numElements * sizeof(float);\n\n size_t numBlocks = hostArrays->numElements / (size_t)THREADS_PER_BLOCK;\n if ((numBlocks % (size_t)THREADS_PER_BLOCK) != 0) {\n numBlocks++;\n }\n\n hostArrays->input = (float *)malloc(bytes);\n hostArrays->square = (float *)malloc(bytes);\n hostArrays->negSquare = (float *)malloc(bytes);\n hostArrays->bytes = bytes;\n hostArrays->numBlocks = numBlocks;\n\n fillRandomly(hostArrays->input, hostArrays->numElements);\n fillRandomly(hostArrays->square, hostArrays->numElements);\n fillRandomly(hostArrays->negSquare, hostArrays->numElements);\n}\n\nvoid createFreeGraph(hipGraphExec_t *graphExec, float *dPtr)\n{\n hipGraph_t graph;\n hipGraphNode_t freeNode;\n\n checkCudaErrors(hipGraphCreate(&graph, 0));\n\n checkCudaErrors(hipGraphAddMemFreeNode(&freeNode, graph, NULL, 0, (void *)dPtr));\n\n checkCudaErrors(hipGraphInstantiate(graphExec, graph, NULL, NULL, 0));\n checkCudaErrors(hipGraphDestroy(graph));\n}\n\n/**\n * Demonstrates explicitly creating a CUDA graph including memory nodes.\n * createNegateSquaresGraphWithStreamCapture constructs an equivalent graph\n * using stream capture.\n *\n * If d_negSquare_out is non null, then:\n * 1) d_negSquare will not be freed;\n * 2) the value of d_negSquare_out will be set to d_negSquare.\n *\n * Diagram of the graph constructed by createNegateSquaresGraphExplicitly:\n *\n * alloc d_input\n * |\n * alloc d_square\n * |\n * Memcpy a to device\n * |\n * launch kernel squareArray ------->---- Memcpy d_square to host\n * | |\n * free d_input |\n * | |\n * allocate d_negSquare |\n * | |\n * launch kernel negateArray -------->--- free d_square\n * |\n * Memcpy d_negSquare to host\n * |\n * free d_negSquare\n */\nvoid createNegateSquaresGraphExplicitly(hipGraphExec_t *graphExec,\n int device,\n negSquareArrays *hostArrays,\n float **d_negSquare_out = NULL)\n{\n // Array buffers on device\n float *d_input, *d_square, *d_negSquare;\n\n // Memory allocation parameters\n hipMemAllocNodeParams allocParams;\n memset(&allocParams, 0, sizeof(allocParams));\n allocParams.bytesize = hostArrays->bytes;\n allocParams.poolProps.allocType = hipMemAllocationTypePinned;\n allocParams.poolProps.location.id = device;\n allocParams.poolProps.location.type = hipMemLocationTypeDevice;\n\n // Kernel launch parameters\n hipKernelNodeParams kernelNodeParams = {0};\n kernelNodeParams.gridDim = dim3(hostArrays->numBlocks, 1, 1);\n kernelNodeParams.blockDim = dim3(THREADS_PER_BLOCK, 1, 1);\n kernelNodeParams.sharedMemBytes = 0;\n kernelNodeParams.extra = NULL;\n\n hipGraph_t graph;\n hipGraphNode_t allocNodeInput, allocNodeSquare, allocNodeNegSquare;\n hipGraphNode_t copyNodeInput, copyNodeSquare, copyNodeNegSquare;\n hipGraphNode_t squareKernelNode, negateKernelNode;\n hipGraphNode_t freeNodeInput, freeNodeSquare;\n\n // Buffer for storing graph node dependencies\n std::vector nodeDependencies;\n\n checkCudaErrors(hipGraphCreate(&graph, 0));\n\n checkCudaErrors(hipGraphAddMemAllocNode(&allocNodeInput, graph, NULL, 0, &allocParams));\n d_input = (float *)allocParams.dptr;\n\n // To keep the graph structure simple (fewer branching dependencies),\n // allocNodeSquare should depend on allocNodeInput\n checkCudaErrors(hipGraphAddMemAllocNode(&allocNodeSquare, graph, &allocNodeInput, 1, &allocParams));\n d_square = (float *)allocParams.dptr;\n\n // copyNodeInput needs to depend on allocNodeInput because copyNodeInput\n // writes to d_input. It does so here indirectly through allocNodeSquare.\n checkCudaErrors(hipGraphAddMemcpyNode1D(©NodeInput,\n graph,\n &allocNodeSquare,\n 1,\n d_input,\n hostArrays->input,\n hostArrays->bytes,\n hipMemcpyHostToDevice));\n\n void *squareKernelArgs[3] = {(void *)&d_input, (void *)&d_square, (void *)&(hostArrays->numElements)};\n kernelNodeParams.func = (void *)squareArray;\n kernelNodeParams.kernelParams = (void **)squareKernelArgs;\n\n // Square kernel depends on copyNodeInput to ensure all data is on the device\n // before kernel launch.\n checkCudaErrors(hipGraphAddKernelNode(&squareKernelNode, graph, ©NodeInput, 1, &kernelNodeParams));\n\n checkCudaErrors(hipGraphAddMemcpyNode1D(©NodeSquare,\n graph,\n &squareKernelNode,\n 1,\n hostArrays->square,\n d_square,\n hostArrays->bytes,\n hipMemcpyDeviceToHost));\n\n // Free of d_input depends on the square kernel to ensure that d_input is not\n // freed while being read by the kernel. It also depends on the alloc of\n // d_input via squareKernelNode > copyNodeInput > allocNodeSquare >\n // allocNodeInput.\n checkCudaErrors(hipGraphAddMemFreeNode(&freeNodeInput, graph, &squareKernelNode, 1, d_input));\n\n // Allocation of C depends on free of A so CUDA can reuse the virtual address.\n checkCudaErrors(hipGraphAddMemAllocNode(&allocNodeNegSquare, graph, &freeNodeInput, 1, &allocParams));\n d_negSquare = (float *)allocParams.dptr;\n\n if (d_negSquare == d_input) {\n printf(\"Check verified that d_negSquare and d_input share a virtual \"\n \"address.\\n\");\n }\n\n void *negateKernelArgs[3] = {(void *)&d_square, (void *)&d_negSquare, (void *)&(hostArrays->numElements)};\n kernelNodeParams.func = (void *)negateArray;\n kernelNodeParams.kernelParams = (void **)negateKernelArgs;\n\n checkCudaErrors(hipGraphAddKernelNode(&negateKernelNode, graph, &allocNodeNegSquare, 1, &kernelNodeParams));\n\n nodeDependencies.push_back(copyNodeSquare);\n nodeDependencies.push_back(negateKernelNode);\n checkCudaErrors(\n hipGraphAddMemFreeNode(&freeNodeSquare, graph, nodeDependencies.data(), nodeDependencies.size(), d_square));\n nodeDependencies.clear();\n\n checkCudaErrors(hipGraphAddMemcpyNode1D(©NodeNegSquare,\n graph,\n &negateKernelNode,\n 1,\n hostArrays->negSquare,\n d_negSquare,\n hostArrays->bytes,\n hipMemcpyDeviceToHost));\n\n if (d_negSquare_out == NULL) {\n hipGraphNode_t freeNodeNegSquare;\n checkCudaErrors(hipGraphAddMemFreeNode(&freeNodeNegSquare, graph, ©NodeNegSquare, 1, d_negSquare));\n }\n else {\n *d_negSquare_out = d_negSquare;\n }\n\n checkCudaErrors(hipGraphInstantiate(graphExec, graph, NULL, NULL, 0));\n checkCudaErrors(hipGraphDestroy(graph));\n}\n\n/**\n * Adds work to a CUDA stream which negates the square of values in the input\n * array.\n *\n * If d_negSquare_out is non null, then:\n * 1) d_negSquare will not be freed;\n * 2) the value of d_negSquare_out will be set to d_negSquare.\n *\n * Diagram of the stream operations in doNegateSquaresInStream\n * ---------------------------------------------------------------------\n * | STREAM | STREAM2 |\n * ---------------------------------------------------------------------\n *\n * alloc d_input\n * |\n * alloc d_square\n * |\n * Memcpy a to device\n * |\n * launch kernel squareArray\n * |\n * record squareKernelCompleteEvent -->-- wait squareKernelCompleteEvent\n * | |\n * free d_input |\n * | |\n * allocate d_negSquare Memcpy d_square to host\n * | |\n * launch kernel negateArray |\n * | |\n * record negateKernelCompleteEvent -->-- wait negateKernelCompleteEvent\n * | |\n * Memcpy d_negSquare to host |\n * | free d_square\n * free d_negSquare |\n * | |\n * wait squareFreeEvent --------------<---- record squareFreeEvent\n */\nvoid doNegateSquaresInStream(hipStream_t stream1, negSquareArrays *hostArrays, float **d_negSquare_out = NULL)\n{\n float *d_input, *d_square, *d_negSquare;\n hipStream_t stream2;\n hipEvent_t squareKernelCompleteEvent, negateKernelCompleteEvent, squareFreeEvent;\n\n checkCudaErrors(hipStreamCreateWithFlags(&stream2, hipStreamNonBlocking));\n\n checkCudaErrors(hipEventCreate(&squareKernelCompleteEvent));\n checkCudaErrors(hipEventCreate(&negateKernelCompleteEvent));\n checkCudaErrors(hipEventCreate(&squareFreeEvent));\n\n // Virtual addresses are assigned synchronously when cudaMallocAsync is\n // called, thus there is no performace benefit gained by separating the\n // allocations into two streams.\n checkCudaErrors(hipMallocAsync(&d_input, hostArrays->bytes, stream1));\n checkCudaErrors(hipMallocAsync(&d_square, hostArrays->bytes, stream1));\n\n checkCudaErrors(hipMemcpyAsync(d_input, hostArrays->input, hostArrays->bytes, hipMemcpyHostToDevice, stream1));\n squareArray<<numBlocks, THREADS_PER_BLOCK, 0, stream1>>>(d_input, d_square, hostArrays->numElements);\n checkCudaErrors(hipEventRecord(squareKernelCompleteEvent, stream1));\n\n checkCudaErrors(hipStreamWaitEvent(stream2, squareKernelCompleteEvent, 0));\n checkCudaErrors(hipMemcpyAsync(hostArrays->square, d_square, hostArrays->bytes, hipMemcpyDeviceToHost, stream2));\n\n checkCudaErrors(hipFreeAsync(d_input, stream1));\n checkCudaErrors(hipMallocAsync(&d_negSquare, hostArrays->bytes, stream1));\n negateArray<<numBlocks, THREADS_PER_BLOCK, 0, stream1>>>(\n d_square, d_negSquare, hostArrays->numElements);\n checkCudaErrors(hipEventRecord(negateKernelCompleteEvent, stream1));\n checkCudaErrors(\n hipMemcpyAsync(hostArrays->negSquare, d_negSquare, hostArrays->bytes, hipMemcpyDeviceToHost, stream1));\n if (d_negSquare_out == NULL) {\n checkCudaErrors(hipFreeAsync(d_negSquare, stream1));\n }\n else {\n *d_negSquare_out = d_negSquare;\n }\n\n checkCudaErrors(hipStreamWaitEvent(stream2, negateKernelCompleteEvent, 0));\n checkCudaErrors(hipFreeAsync(d_square, stream2));\n checkCudaErrors(hipEventRecord(squareFreeEvent, stream2));\n\n checkCudaErrors(hipStreamWaitEvent(stream1, squareFreeEvent, 0));\n\n checkCudaErrors(hipStreamDestroy(stream2));\n checkCudaErrors(hipEventDestroy(squareKernelCompleteEvent));\n checkCudaErrors(hipEventDestroy(negateKernelCompleteEvent));\n checkCudaErrors(hipEventDestroy(squareFreeEvent));\n}\n\n/**\n * Demonstrates creating a CUDA graph including memory nodes using stream\n * capture. createNegateSquaresGraphExplicitly constructs an equivalent graph\n * without stream capture.\n */\nvoid createNegateSquaresGraphWithStreamCapture(hipGraphExec_t *graphExec,\n negSquareArrays *hostArrays,\n float **d_negSquare_out = NULL)\n{\n hipGraph_t graph;\n hipStream_t stream;\n\n checkCudaErrors(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));\n\n checkCudaErrors(hipStreamBeginCapture(stream, hipStreamCaptureModeGlobal));\n doNegateSquaresInStream(stream, hostArrays, d_negSquare_out);\n checkCudaErrors(hipStreamEndCapture(stream, &graph));\n\n checkCudaErrors(hipGraphInstantiate(graphExec, graph, NULL, NULL, 0));\n checkCudaErrors(hipStreamDestroy(stream));\n checkCudaErrors(hipGraphDestroy(graph));\n}\n\nvoid prepareRefArrays(negSquareArrays *hostArrays, negSquareArrays *deviceRefArrays, bool **foundValidationFailure)\n{\n deviceRefArrays->bytes = hostArrays->bytes;\n deviceRefArrays->numElements = hostArrays->numElements;\n\n for (int i = 0; i < hostArrays->numElements; i++) {\n hostArrays->square[i] = hostArrays->input[i] * hostArrays->input[i];\n hostArrays->negSquare[i] = hostArrays->square[i] * -1;\n }\n\n checkCudaErrors(hipMalloc((void **)&deviceRefArrays->negSquare, deviceRefArrays->bytes));\n checkCudaErrors(\n hipMemcpy(deviceRefArrays->negSquare, hostArrays->negSquare, hostArrays->bytes, hipMemcpyHostToDevice));\n\n checkCudaErrors(hipMallocManaged((void **)foundValidationFailure, sizeof(bool)));\n}\n\nint checkValidationFailure(bool *foundValidationFailure)\n{\n if (*foundValidationFailure) {\n printf(\"Validation FAILURE!\\n\\n\");\n *foundValidationFailure = false;\n return EXIT_FAILURE;\n }\n else {\n printf(\"Validation PASSED!\\n\\n\");\n return EXIT_SUCCESS;\n }\n}\n\n__global__ void validateGPU(float *d_negSquare, negSquareArrays devRefArrays, bool *foundValidationFailure)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n float ref, diff;\n\n if (idx < devRefArrays.numElements) {\n ref = devRefArrays.negSquare[idx];\n diff = d_negSquare[idx] - ref;\n diff *= diff;\n ref *= ref;\n if (diff / ref > ALLOWABLE_VARIANCE) {\n *foundValidationFailure = true;\n }\n }\n}\n\nvoid validateHost(negSquareArrays *hostArrays, bool *foundValidationFailure)\n{\n float ref, diff;\n\n for (int i = 0; i < hostArrays->numElements; i++) {\n ref = hostArrays->input[i] * hostArrays->input[i] * -1;\n diff = hostArrays->negSquare[i] - ref;\n diff *= diff;\n ref *= ref;\n if (diff / ref > ALLOWABLE_VARIANCE) {\n *foundValidationFailure = true;\n }\n }\n}\n\nint main(int argc, char **argv)\n{\n negSquareArrays hostArrays, deviceRefArrays;\n hipStream_t stream;\n hipGraphExec_t graphExec, graphExecFreeC;\n\n // Declare pointers for GPU buffers\n float *d_negSquare = NULL;\n bool *foundValidationFailure = NULL;\n\n srand(time(0));\n int device = findCudaDevice(argc, (const char **)argv);\n\n int driverVersion = 0;\n int deviceSupportsMemoryPools = 0;\n\n hipDriverGetVersion(&driverVersion);\n printf(\"Driver version is: %d.%d\\n\", driverVersion / 1000, (driverVersion % 100) / 10);\n\n if (driverVersion < 11040) {\n printf(\"Waiving execution as driver does not support Graph Memory Nodes\\n\");\n exit(EXIT_WAIVED);\n }\n\n hipDeviceGetAttribute(&deviceSupportsMemoryPools, hipDeviceAttributeMemoryPoolsSupported, device);\n if (!deviceSupportsMemoryPools) {\n printf(\"Waiving execution as device does not support Memory Pools\\n\");\n exit(EXIT_WAIVED);\n }\n else {\n printf(\"Setting up sample.\\n\");\n }\n\n prepareHostArrays(&hostArrays);\n prepareRefArrays(&hostArrays, &deviceRefArrays, &foundValidationFailure);\n checkCudaErrors(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));\n printf(\"Setup complete.\\n\\n\");\n\n printf(\"Running negateSquares in a stream.\\n\");\n doNegateSquaresInStream(stream, &hostArrays);\n checkCudaErrors(hipStreamSynchronize(stream));\n printf(\"Validating negateSquares in a stream...\\n\");\n validateHost(&hostArrays, foundValidationFailure);\n checkValidationFailure(foundValidationFailure);\n resetOutputArrays(&hostArrays);\n\n printf(\"Running negateSquares in a stream-captured graph.\\n\");\n createNegateSquaresGraphWithStreamCapture(&graphExec, &hostArrays);\n checkCudaErrors(hipGraphLaunch(graphExec, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n printf(\"Validating negateSquares in a stream-captured graph...\\n\");\n validateHost(&hostArrays, foundValidationFailure);\n checkValidationFailure(foundValidationFailure);\n resetOutputArrays(&hostArrays);\n\n printf(\"Running negateSquares in an explicitly constructed graph.\\n\");\n createNegateSquaresGraphExplicitly(&graphExec, device, &hostArrays);\n checkCudaErrors(hipGraphLaunch(graphExec, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n printf(\"Validating negateSquares in an explicitly constructed graph...\\n\");\n validateHost(&hostArrays, foundValidationFailure);\n checkValidationFailure(foundValidationFailure);\n resetOutputArrays(&hostArrays);\n\n // Each of the three examples below free d_negSquare outside the graph. As\n // demonstrated by validateGPU, d_negSquare can be accessed by outside the\n // graph before d_negSquare is freed.\n\n printf(\"Running negateSquares with d_negSquare freed outside the stream.\\n\");\n createNegateSquaresGraphExplicitly(&graphExec, device, &hostArrays, &d_negSquare);\n checkCudaErrors(hipGraphLaunch(graphExec, stream));\n validateGPU<<>>(\n d_negSquare, deviceRefArrays, foundValidationFailure);\n // Since cudaFree is synchronous, the stream must synchronize before freeing\n // d_negSquare to ensure d_negSquare no longer being accessed.\n checkCudaErrors(hipStreamSynchronize(stream));\n checkCudaErrors(hipFree(d_negSquare));\n printf(\"Validating negateSquares with d_negSquare freed outside the \"\n \"stream...\\n\");\n validateHost(&hostArrays, foundValidationFailure);\n checkValidationFailure(foundValidationFailure);\n resetOutputArrays(&hostArrays);\n\n printf(\"Running negateSquares with d_negSquare freed outside the graph.\\n\");\n checkCudaErrors(hipGraphLaunch(graphExec, stream));\n validateGPU<<>>(\n d_negSquare, deviceRefArrays, foundValidationFailure);\n checkCudaErrors(hipFreeAsync(d_negSquare, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n printf(\"Validating negateSquares with d_negSquare freed outside the graph...\\n\");\n checkValidationFailure(foundValidationFailure);\n resetOutputArrays(&hostArrays);\n\n printf(\"Running negateSquares with d_negSquare freed in a different graph.\\n\");\n createFreeGraph(&graphExecFreeC, d_negSquare);\n checkCudaErrors(hipGraphLaunch(graphExec, stream));\n validateGPU<<>>(\n d_negSquare, deviceRefArrays, foundValidationFailure);\n checkCudaErrors(hipGraphLaunch(graphExecFreeC, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n printf(\"Validating negateSquares with d_negSquare freed in a different \"\n \"graph...\\n\");\n checkValidationFailure(foundValidationFailure);\n\n printf(\"Cleaning up sample.\\n\");\n checkCudaErrors(hipGraphExecDestroy(graphExec));\n checkCudaErrors(hipGraphExecDestroy(graphExecFreeC));\n checkCudaErrors(hipStreamDestroy(stream));\n checkCudaErrors(hipFree(foundValidationFailure));\n checkCudaErrors(hipFree(deviceRefArrays.negSquare));\n free(hostArrays.input);\n free(hostArrays.square);\n free(hostArrays.negSquare);\n printf(\"Cleanup complete. Exiting sample.\\n\");\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/graphMemoryNodes/graphMemoryNodes.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n This sample has two kernels, one doing the rendering every frame, and\n another one used to generate the mip map levels at startup.\n\n For rendering we use a \"virtual\" texturing approach, where one 2d texture\n stores pointers to the actual textures used. This can be achieved by the\n new cudaTextureObject introduced in CUDA 5.0 and requiring sm3+ hardware.\n\n The mipmap generation kernel uses cudaSurfaceObject and cudaTextureObject\n passed as kernel arguments to compute the higher mip map level based on\n the lower.\n*/\n\n#ifndef _BINDLESSTEXTURE_KERNEL_CU_\n#define _BINDLESSTEXTURE_KERNEL_CU_\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"bindlessTexture.h\"\n\n// set this to just see the mipmap chain of first image\n// #define SHOW_MIPMAPS\n\n// local references to resources\n\nImage atlasImage;\nstd::vector contentImages;\nfloat highestLod = 1.0f;\n\n#ifndef MAX\n#define MAX(a, b) ((a > b) ? a : b)\n#endif\n\n//////////////////////////////////////////////////////////////////////////\n\n__host__ __device__ __inline__ uint2 encodeTextureObject(cudaTextureObject_t obj)\n{\n return make_uint2((uint)(obj & 0xFFFFFFFF), (uint)(obj >> 32));\n}\n\n__host__ __device__ __inline__ cudaTextureObject_t decodeTextureObject(uint2 obj)\n{\n return (((cudaTextureObject_t)obj.x) | ((cudaTextureObject_t)obj.y) << 32);\n}\n\n__device__ __inline__ float4 to_float4(uchar4 vec) { return make_float4(vec.x, vec.y, vec.z, vec.w); }\n\n__device__ __inline__ uchar4 to_uchar4(float4 vec)\n{\n return make_uchar4((uchar)vec.x, (uchar)vec.y, (uchar)vec.z, (uchar)vec.w);\n}\n\n//////////////////////////////////////////////////////////////////////////\n// Rendering\n\n// the atlas texture stores the 64 bit cudaTextureObjects\n// we use it for \"virtual\" texturing\n\n__global__ void d_render(uchar4 *d_output, uint imageW, uint imageH, float lod, cudaTextureObject_t atlasTexture)\n{\n uint x = blockIdx.x * blockDim.x + threadIdx.x;\n uint y = blockIdx.y * blockDim.y + threadIdx.y;\n\n float u = x / (float)imageW;\n float v = y / (float)imageH;\n\n if ((x < imageW) && (y < imageH)) {\n // read from 2D atlas texture and decode texture object\n uint2 texCoded = tex2D(atlasTexture, u, v);\n cudaTextureObject_t tex = decodeTextureObject(texCoded);\n\n // read from cuda texture object, use template to specify what data will be\n // returned. tex2DLod allows us to pass the lod (mip map level) directly.\n // There is other functions with CUDA 5, e.g. tex2DGrad, that allow you\n // to pass derivatives to perform automatic mipmap/anisotropic filtering.\n float4 color = tex2DLod(tex, u, 1 - v, lod);\n // In our sample tex is always valid, but for something like your own\n // sparse texturing you would need to make sure to handle the zero case.\n\n // write output color\n uint i = y * imageW + x;\n d_output[i] = to_uchar4(color * 255.0);\n }\n}\n\nextern \"C\" void renderAtlasImage(dim3 gridSize, dim3 blockSize, uchar4 *d_output, uint imageW, uint imageH, float lod)\n{\n // psuedo animate lod\n lod = fmodf(lod, highestLod * 2);\n lod = highestLod - fabs(lod - highestLod);\n\n#ifdef SHOW_MIPMAPS\n lod = 0.0f;\n#endif\n\n d_render<<>>(d_output, imageW, imageH, lod, atlasImage.textureObject);\n\n checkCudaErrors(cudaGetLastError());\n}\n\n//////////////////////////////////////////////////////////////////////////\n// MipMap Generation\n\n// A key benefit of using the new surface objects is that we don't need any\n// global binding points anymore. We can directly pass them as function\n// arguments.\n\n__global__ void d_mipmap(cudaSurfaceObject_t mipOutput, cudaTextureObject_t mipInput, uint imageW, uint imageH)\n{\n uint x = blockIdx.x * blockDim.x + threadIdx.x;\n uint y = blockIdx.y * blockDim.y + threadIdx.y;\n\n float px = 1.0 / float(imageW);\n float py = 1.0 / float(imageH);\n\n if ((x < imageW) && (y < imageH)) {\n // take the average of 4 samples\n\n // we are using the normalized access to make sure non-power-of-two textures\n // behave well when downsized.\n float4 color = (tex2D(mipInput, (x + 0) * px, (y + 0) * py))\n + (tex2D(mipInput, (x + 1) * px, (y + 0) * py))\n + (tex2D(mipInput, (x + 1) * px, (y + 1) * py))\n + (tex2D(mipInput, (x + 0) * px, (y + 1) * py));\n\n color /= 4.0;\n color *= 255.0;\n color = fminf(color, make_float4(255.0));\n\n surf2Dwrite(to_uchar4(color), mipOutput, x * sizeof(uchar4), y);\n }\n}\n\nvoid generateMipMaps(cudaMipmappedArray_t mipmapArray, cudaExtent size)\n{\n size_t width = size.width;\n size_t height = size.height;\n\n#ifdef SHOW_MIPMAPS\n cudaArray_t levelFirst;\n checkCudaErrors(cudaGetMipmappedArrayLevel(&levelFirst, mipmapArray, 0));\n#endif\n\n uint level = 0;\n\n while (width != 1 || height != 1) {\n width /= 2;\n width = MAX((size_t)1, width);\n height /= 2;\n height = MAX((size_t)1, height);\n\n cudaArray_t levelFrom;\n checkCudaErrors(cudaGetMipmappedArrayLevel(&levelFrom, mipmapArray, level));\n cudaArray_t levelTo;\n checkCudaErrors(cudaGetMipmappedArrayLevel(&levelTo, mipmapArray, level + 1));\n\n cudaExtent levelToSize;\n checkCudaErrors(cudaArrayGetInfo(NULL, &levelToSize, NULL, levelTo));\n checkHost(levelToSize.width == width);\n checkHost(levelToSize.height == height);\n checkHost(levelToSize.depth == 0);\n\n // generate texture object for reading\n cudaTextureObject_t texInput;\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = levelFrom;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = 1;\n texDescr.filterMode = cudaFilterModeLinear;\n\n texDescr.addressMode[0] = cudaAddressModeClamp;\n texDescr.addressMode[1] = cudaAddressModeClamp;\n texDescr.addressMode[2] = cudaAddressModeClamp;\n\n texDescr.readMode = cudaReadModeNormalizedFloat;\n\n checkCudaErrors(cudaCreateTextureObject(&texInput, &texRes, &texDescr, NULL));\n\n // generate surface object for writing\n\n cudaSurfaceObject_t surfOutput;\n cudaResourceDesc surfRes;\n memset(&surfRes, 0, sizeof(cudaResourceDesc));\n surfRes.resType = cudaResourceTypeArray;\n surfRes.res.array.array = levelTo;\n\n checkCudaErrors(cudaCreateSurfaceObject(&surfOutput, &surfRes));\n\n // run mipmap kernel\n dim3 blockSize(16, 16, 1);\n dim3 gridSize(((uint)width + blockSize.x - 1) / blockSize.x, ((uint)height + blockSize.y - 1) / blockSize.y, 1);\n\n d_mipmap<<>>(surfOutput, texInput, (uint)width, (uint)height);\n\n checkCudaErrors(cudaDeviceSynchronize());\n checkCudaErrors(cudaGetLastError());\n\n checkCudaErrors(cudaDestroySurfaceObject(surfOutput));\n\n checkCudaErrors(cudaDestroyTextureObject(texInput));\n\n#ifdef SHOW_MIPMAPS\n // we blit the current mipmap back into first level\n cudaMemcpy3DParms copyParams = {0};\n copyParams.dstArray = levelFirst;\n copyParams.srcArray = levelTo;\n copyParams.extent = make_cudaExtent(width, height, 1);\n copyParams.kind = cudaMemcpyDeviceToDevice;\n checkCudaErrors(cudaMemcpy3D(©Params));\n#endif\n\n level++;\n }\n}\n\nuint getMipMapLevels(cudaExtent size)\n{\n size_t sz = MAX(MAX(size.width, size.height), size.depth);\n\n uint levels = 0;\n\n while (sz) {\n sz /= 2;\n levels++;\n }\n\n return levels;\n}\n\n//////////////////////////////////////////////////////////////////////////\n// Initalization\n\nextern \"C\" void randomizeAtlas()\n{\n uint2 *h_data = (uint2 *)atlasImage.h_data;\n\n // assign random texture object handles to our atlas image tiles\n for (size_t i = 0; i < atlasImage.size.width * atlasImage.size.height; i++) {\n#ifdef SHOW_MIPMAPS\n h_data[i] = encodeTextureObject(contentImages[0].textureObject);\n#else\n h_data[i] = encodeTextureObject(contentImages[rand() % contentImages.size()].textureObject);\n#endif\n }\n\n // copy data to atlas array\n cudaMemcpy3DParms copyParams = {0};\n copyParams.srcPtr = make_cudaPitchedPtr(\n atlasImage.h_data, atlasImage.size.width * sizeof(uint2), atlasImage.size.width, atlasImage.size.height);\n copyParams.dstArray = atlasImage.dataArray;\n copyParams.extent = atlasImage.size;\n copyParams.extent.depth = 1;\n copyParams.kind = cudaMemcpyHostToDevice;\n checkCudaErrors(cudaMemcpy3D(©Params));\n};\n\nextern \"C\" void deinitAtlasAndImages()\n{\n for (size_t i = 0; i < contentImages.size(); i++) {\n Image &image = contentImages[i];\n\n if (image.h_data) {\n free(image.h_data);\n }\n\n if (image.textureObject) {\n checkCudaErrors(cudaDestroyTextureObject(image.textureObject));\n }\n\n if (image.mipmapArray) {\n checkCudaErrors(cudaFreeMipmappedArray(image.mipmapArray));\n }\n }\n\n if (atlasImage.h_data) {\n free(atlasImage.h_data);\n }\n\n if (atlasImage.textureObject) {\n checkCudaErrors(cudaDestroyTextureObject(atlasImage.textureObject));\n }\n\n if (atlasImage.dataArray) {\n checkCudaErrors(cudaFreeArray(atlasImage.dataArray));\n }\n}\n\nextern \"C\" void initAtlasAndImages(const Image *images, size_t numImages, cudaExtent atlasSize)\n{\n // create individual textures\n contentImages.resize(numImages);\n\n for (size_t i = 0; i < numImages; i++) {\n Image &image = contentImages[i];\n image.size = images[i].size;\n image.size.depth = 0;\n image.type = cudaResourceTypeMipmappedArray;\n\n // how many mipmaps we need\n uint levels = getMipMapLevels(image.size);\n highestLod = MAX(highestLod, (float)levels - 1);\n\n cudaChannelFormatDesc desc = cudaCreateChannelDesc();\n checkCudaErrors(cudaMallocMipmappedArray(&image.mipmapArray, &desc, image.size, levels));\n\n // upload level 0\n cudaArray_t level0;\n checkCudaErrors(cudaGetMipmappedArrayLevel(&level0, image.mipmapArray, 0));\n\n cudaMemcpy3DParms copyParams = {0};\n copyParams.srcPtr = make_cudaPitchedPtr(\n images[i].h_data, image.size.width * sizeof(uchar4), image.size.width, image.size.height);\n copyParams.dstArray = level0;\n copyParams.extent = image.size;\n copyParams.extent.depth = 1;\n copyParams.kind = cudaMemcpyHostToDevice;\n checkCudaErrors(cudaMemcpy3D(©Params));\n\n // compute rest of mipmaps based on level 0\n generateMipMaps(image.mipmapArray, image.size);\n\n // generate bindless texture object\n\n cudaResourceDesc resDescr;\n memset(&resDescr, 0, sizeof(cudaResourceDesc));\n\n resDescr.resType = cudaResourceTypeMipmappedArray;\n resDescr.res.mipmap.mipmap = image.mipmapArray;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = 1;\n texDescr.filterMode = cudaFilterModeLinear;\n texDescr.mipmapFilterMode = cudaFilterModeLinear;\n\n texDescr.addressMode[0] = cudaAddressModeClamp;\n texDescr.addressMode[1] = cudaAddressModeClamp;\n texDescr.addressMode[2] = cudaAddressModeClamp;\n\n texDescr.maxMipmapLevelClamp = float(levels - 1);\n\n texDescr.readMode = cudaReadModeNormalizedFloat;\n\n checkCudaErrors(cudaCreateTextureObject(&image.textureObject, &resDescr, &texDescr, NULL));\n }\n\n // create atlas array\n cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc();\n checkCudaErrors(cudaMallocArray(&atlasImage.dataArray, &channelDesc, atlasSize.width, atlasSize.height));\n atlasImage.h_data = malloc(atlasSize.width * atlasSize.height * sizeof(uint2));\n atlasImage.type = cudaResourceTypeArray;\n atlasImage.size = atlasSize;\n\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = atlasImage.dataArray;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = cudaFilterModePoint;\n texDescr.addressMode[0] = cudaAddressModeClamp;\n texDescr.addressMode[1] = cudaAddressModeClamp;\n texDescr.addressMode[1] = cudaAddressModeClamp;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&atlasImage.textureObject, &texRes, &texDescr, NULL));\n\n randomizeAtlas();\n}\n\n#endif // #ifndef _SIMPLETEXTURE3D_KERNEL_CU_\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n This sample has two kernels, one doing the rendering every frame, and\n another one used to generate the mip map levels at startup.\n\n For rendering we use a \"virtual\" texturing approach, where one 2d texture\n stores pointers to the actual textures used. This can be achieved by the\n new cudaTextureObject introduced in CUDA 5.0 and requiring sm3+ hardware.\n\n The mipmap generation kernel uses cudaSurfaceObject and cudaTextureObject\n passed as kernel arguments to compute the higher mip map level based on\n the lower.\n*/\n\n#include \"hip/hip_runtime.h\"\n#ifndef _BINDLESSTEXTURE_KERNEL_CU_\n#define _BINDLESSTEXTURE_KERNEL_CU_\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"bindlessTexture.h\"\n\n// set this to just see the mipmap chain of first image\n// #define SHOW_MIPMAPS\n\n// local references to resources\n\nImage atlasImage;\nstd::vector contentImages;\nfloat highestLod = 1.0f;\n\n#ifndef MAX\n#define MAX(a, b) ((a > b) ? a : b)\n#endif\n\n//////////////////////////////////////////////////////////////////////////\n\n__host__ __device__ __inline__ uint2 encodeTextureObject(hipTextureObject_t obj)\n{\n return make_uint2((uint)(obj & 0xFFFFFFFF), (uint)(obj >> 32));\n}\n\n__host__ __device__ __inline__ hipTextureObject_t decodeTextureObject(uint2 obj)\n{\n return (((hipTextureObject_t)obj.x) | ((hipTextureObject_t)obj.y) << 32);\n}\n\n__device__ __inline__ float4 to_float4(uchar4 vec) { return make_float4(vec.x, vec.y, vec.z, vec.w); }\n\n__device__ __inline__ uchar4 to_uchar4(float4 vec)\n{\n return make_uchar4((uchar)vec.x, (uchar)vec.y, (uchar)vec.z, (uchar)vec.w);\n}\n\n//////////////////////////////////////////////////////////////////////////\n// Rendering\n\n// the atlas texture stores the 64 bit cudaTextureObjects\n// we use it for \"virtual\" texturing\n\n__global__ void d_render(uchar4 *d_output, uint imageW, uint imageH, float lod, hipTextureObject_t atlasTexture)\n{\n uint x = blockIdx.x * blockDim.x + threadIdx.x;\n uint y = blockIdx.y * blockDim.y + threadIdx.y;\n\n float u = x / (float)imageW;\n float v = y / (float)imageH;\n\n if ((x < imageW) && (y < imageH)) {\n // read from 2D atlas texture and decode texture object\n uint2 texCoded = tex2D(atlasTexture, u, v);\n hipTextureObject_t tex = decodeTextureObject(texCoded);\n\n // read from cuda texture object, use template to specify what data will be\n // returned. tex2DLod allows us to pass the lod (mip map level) directly.\n // There is other functions with CUDA 5, e.g. tex2DGrad, that allow you\n // to pass derivatives to perform automatic mipmap/anisotropic filtering.\n float4 color = tex2DLod(tex, u, 1 - v, lod);\n // In our sample tex is always valid, but for something like your own\n // sparse texturing you would need to make sure to handle the zero case.\n\n // write output color\n uint i = y * imageW + x;\n d_output[i] = to_uchar4(color * 255.0);\n }\n}\n\nextern \"C\" void renderAtlasImage(dim3 gridSize, dim3 blockSize, uchar4 *d_output, uint imageW, uint imageH, float lod)\n{\n // psuedo animate lod\n lod = fmodf(lod, highestLod * 2);\n lod = highestLod - fabs(lod - highestLod);\n\n#ifdef SHOW_MIPMAPS\n lod = 0.0f;\n#endif\n\n d_render<<>>(d_output, imageW, imageH, lod, atlasImage.textureObject);\n\n checkCudaErrors(hipGetLastError());\n}\n\n//////////////////////////////////////////////////////////////////////////\n// MipMap Generation\n\n// A key benefit of using the new surface objects is that we don't need any\n// global binding points anymore. We can directly pass them as function\n// arguments.\n\n__global__ void d_mipmap(hipSurfaceObject_t mipOutput, hipTextureObject_t mipInput, uint imageW, uint imageH)\n{\n uint x = blockIdx.x * blockDim.x + threadIdx.x;\n uint y = blockIdx.y * blockDim.y + threadIdx.y;\n\n float px = 1.0 / float(imageW);\n float py = 1.0 / float(imageH);\n\n if ((x < imageW) && (y < imageH)) {\n // take the average of 4 samples\n\n // we are using the normalized access to make sure non-power-of-two textures\n // behave well when downsized.\n float4 color = (tex2D(mipInput, (x + 0) * px, (y + 0) * py))\n + (tex2D(mipInput, (x + 1) * px, (y + 0) * py))\n + (tex2D(mipInput, (x + 1) * px, (y + 1) * py))\n + (tex2D(mipInput, (x + 0) * px, (y + 1) * py));\n\n color /= 4.0;\n color *= 255.0;\n color = fminf(color, make_float4(255.0));\n\n surf2Dwrite(to_uchar4(color), mipOutput, x * sizeof(uchar4), y);\n }\n}\n\nvoid generateMipMaps(hipMipmappedArray_t mipmapArray, hipExtent size)\n{\n size_t width = size.width;\n size_t height = size.height;\n\n#ifdef SHOW_MIPMAPS\n hipArray_t levelFirst;\n checkCudaErrors(hipGetMipmappedArrayLevel(&levelFirst, mipmapArray, 0));\n#endif\n\n uint level = 0;\n\n while (width != 1 || height != 1) {\n width /= 2;\n width = MAX((size_t)1, width);\n height /= 2;\n height = MAX((size_t)1, height);\n\n hipArray_t levelFrom;\n checkCudaErrors(hipGetMipmappedArrayLevel(&levelFrom, mipmapArray, level));\n hipArray_t levelTo;\n checkCudaErrors(hipGetMipmappedArrayLevel(&levelTo, mipmapArray, level + 1));\n\n hipExtent levelToSize;\n checkCudaErrors(hipArrayGetInfo(NULL, &levelToSize, NULL, levelTo));\n checkHost(levelToSize.width == width);\n checkHost(levelToSize.height == height);\n checkHost(levelToSize.depth == 0);\n\n // generate texture object for reading\n hipTextureObject_t texInput;\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = levelFrom;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = 1;\n texDescr.filterMode = hipFilterModeLinear;\n\n texDescr.addressMode[0] = hipAddressModeClamp;\n texDescr.addressMode[1] = hipAddressModeClamp;\n texDescr.addressMode[2] = hipAddressModeClamp;\n\n texDescr.readMode = hipReadModeNormalizedFloat;\n\n checkCudaErrors(hipCreateTextureObject(&texInput, &texRes, &texDescr, NULL));\n\n // generate surface object for writing\n\n hipSurfaceObject_t surfOutput;\n hipResourceDesc surfRes;\n memset(&surfRes, 0, sizeof(hipResourceDesc));\n surfRes.resType = hipResourceTypeArray;\n surfRes.res.array.array = levelTo;\n\n checkCudaErrors(hipCreateSurfaceObject(&surfOutput, &surfRes));\n\n // run mipmap kernel\n dim3 blockSize(16, 16, 1);\n dim3 gridSize(((uint)width + blockSize.x - 1) / blockSize.x, ((uint)height + blockSize.y - 1) / blockSize.y, 1);\n\n d_mipmap<<>>(surfOutput, texInput, (uint)width, (uint)height);\n\n checkCudaErrors(hipDeviceSynchronize());\n checkCudaErrors(hipGetLastError());\n\n checkCudaErrors(hipDestroySurfaceObject(surfOutput));\n\n checkCudaErrors(hipDestroyTextureObject(texInput));\n\n#ifdef SHOW_MIPMAPS\n // we blit the current mipmap back into first level\n hipMemcpy3DParms copyParams = {0};\n copyParams.dstArray = levelFirst;\n copyParams.srcArray = levelTo;\n copyParams.extent = make_hipExtent(width, height, 1);\n copyParams.kind = hipMemcpyDeviceToDevice;\n checkCudaErrors(hipMemcpy3D(©Params));\n#endif\n\n level++;\n }\n}\n\nuint getMipMapLevels(hipExtent size)\n{\n size_t sz = MAX(MAX(size.width, size.height), size.depth);\n\n uint levels = 0;\n\n while (sz) {\n sz /= 2;\n levels++;\n }\n\n return levels;\n}\n\n//////////////////////////////////////////////////////////////////////////\n// Initalization\n\nextern \"C\" void randomizeAtlas()\n{\n uint2 *h_data = (uint2 *)atlasImage.h_data;\n\n // assign random texture object handles to our atlas image tiles\n for (size_t i = 0; i < atlasImage.size.width * atlasImage.size.height; i++) {\n#ifdef SHOW_MIPMAPS\n h_data[i] = encodeTextureObject(contentImages[0].textureObject);\n#else\n h_data[i] = encodeTextureObject(contentImages[rand() % contentImages.size()].textureObject);\n#endif\n }\n\n // copy data to atlas array\n hipMemcpy3DParms copyParams = {0};\n copyParams.srcPtr = make_hipPitchedPtr(\n atlasImage.h_data, atlasImage.size.width * sizeof(uint2), atlasImage.size.width, atlasImage.size.height);\n copyParams.dstArray = atlasImage.dataArray;\n copyParams.extent = atlasImage.size;\n copyParams.extent.depth = 1;\n copyParams.kind = hipMemcpyHostToDevice;\n checkCudaErrors(hipMemcpy3D(©Params));\n};\n\nextern \"C\" void deinitAtlasAndImages()\n{\n for (size_t i = 0; i < contentImages.size(); i++) {\n Image &image = contentImages[i];\n\n if (image.h_data) {\n free(image.h_data);\n }\n\n if (image.textureObject) {\n checkCudaErrors(hipDestroyTextureObject(image.textureObject));\n }\n\n if (image.mipmapArray) {\n checkCudaErrors(hipFreeMipmappedArray(image.mipmapArray));\n }\n }\n\n if (atlasImage.h_data) {\n free(atlasImage.h_data);\n }\n\n if (atlasImage.textureObject) {\n checkCudaErrors(hipDestroyTextureObject(atlasImage.textureObject));\n }\n\n if (atlasImage.dataArray) {\n checkCudaErrors(hipFreeArray(atlasImage.dataArray));\n }\n}\n\nextern \"C\" void initAtlasAndImages(const Image *images, size_t numImages, hipExtent atlasSize)\n{\n // create individual textures\n contentImages.resize(numImages);\n\n for (size_t i = 0; i < numImages; i++) {\n Image &image = contentImages[i];\n image.size = images[i].size;\n image.size.depth = 0;\n image.type = hipResourceTypeMipmappedArray;\n\n // how many mipmaps we need\n uint levels = getMipMapLevels(image.size);\n highestLod = MAX(highestLod, (float)levels - 1);\n\n hipChannelFormatDesc desc = hipCreateChannelDesc();\n checkCudaErrors(hipMallocMipmappedArray(&image.mipmapArray, &desc, image.size, levels));\n\n // upload level 0\n hipArray_t level0;\n checkCudaErrors(hipGetMipmappedArrayLevel(&level0, image.mipmapArray, 0));\n\n hipMemcpy3DParms copyParams = {0};\n copyParams.srcPtr = make_hipPitchedPtr(\n images[i].h_data, image.size.width * sizeof(uchar4), image.size.width, image.size.height);\n copyParams.dstArray = level0;\n copyParams.extent = image.size;\n copyParams.extent.depth = 1;\n copyParams.kind = hipMemcpyHostToDevice;\n checkCudaErrors(hipMemcpy3D(©Params));\n\n // compute rest of mipmaps based on level 0\n generateMipMaps(image.mipmapArray, image.size);\n\n // generate bindless texture object\n\n hipResourceDesc resDescr;\n memset(&resDescr, 0, sizeof(hipResourceDesc));\n\n resDescr.resType = hipResourceTypeMipmappedArray;\n resDescr.res.mipmap.mipmap = image.mipmapArray;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = 1;\n texDescr.filterMode = hipFilterModeLinear;\n texDescr.mipmapFilterMode = hipFilterModeLinear;\n\n texDescr.addressMode[0] = hipAddressModeClamp;\n texDescr.addressMode[1] = hipAddressModeClamp;\n texDescr.addressMode[2] = hipAddressModeClamp;\n\n texDescr.maxMipmapLevelClamp = float(levels - 1);\n\n texDescr.readMode = hipReadModeNormalizedFloat;\n\n checkCudaErrors(hipCreateTextureObject(&image.textureObject, &resDescr, &texDescr, NULL));\n }\n\n // create atlas array\n hipChannelFormatDesc channelDesc = hipCreateChannelDesc();\n checkCudaErrors(hipMallocArray(&atlasImage.dataArray, &channelDesc, atlasSize.width, atlasSize.height));\n atlasImage.h_data = malloc(atlasSize.width * atlasSize.height * sizeof(uint2));\n atlasImage.type = hipResourceTypeArray;\n atlasImage.size = atlasSize;\n\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = atlasImage.dataArray;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = hipFilterModePoint;\n texDescr.addressMode[0] = hipAddressModeClamp;\n texDescr.addressMode[1] = hipAddressModeClamp;\n texDescr.addressMode[1] = hipAddressModeClamp;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&atlasImage.textureObject, &texRes, &texDescr, NULL));\n\n randomizeAtlas();\n}\n\n#endif // #ifndef _SIMPLETEXTURE3D_KERNEL_CU_\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/bindlessTexture/bindlessTexture_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n\n__forceinline__ __device__ float2 operator+(float2 a, float2 b)\n{\n float2 c;\n c.x = a.x + b.x;\n c.y = a.y + b.y;\n return c;\n}\n\n__forceinline__ __device__ float2 operator-(float2 a, float2 b)\n{\n float2 c;\n c.x = a.x - b.x;\n c.y = a.y - b.y;\n return c;\n}\n\n__forceinline__ __device__ float2 operator*(float a, float2 b)\n{\n float2 c;\n c.x = a * b.x;\n c.y = a * b.y;\n return c;\n}\n\n__forceinline__ __device__ float length(float2 a) { return sqrtf(a.x * a.x + a.y * a.y); }\n\n#define MAX_TESSELLATION 32\nstruct BezierLine\n{\n float2 CP[3];\n float2 *vertexPos;\n int nVertices;\n};\n\n__global__ void computeBezierLinePositions(int lidx, BezierLine *bLines, int nTessPoints)\n{\n int idx = threadIdx.x + blockDim.x * blockIdx.x;\n\n if (idx < nTessPoints) {\n float u = (float)idx / (float)(nTessPoints - 1);\n float omu = 1.0f - u;\n\n float B3u[3];\n\n B3u[0] = omu * omu;\n B3u[1] = 2.0f * u * omu;\n B3u[2] = u * u;\n\n float2 position = {0, 0};\n\n for (int i = 0; i < 3; i++) {\n position = position + B3u[i] * bLines[lidx].CP[i];\n }\n\n bLines[lidx].vertexPos[idx] = position;\n }\n}\n\n__global__ void computeBezierLinesCDP(BezierLine *bLines, int nLines)\n{\n int lidx = threadIdx.x + blockDim.x * blockIdx.x;\n\n if (lidx < nLines) {\n float curvature = length(bLines[lidx].CP[1] - 0.5f * (bLines[lidx].CP[0] + bLines[lidx].CP[2]))\n / length(bLines[lidx].CP[2] - bLines[lidx].CP[0]);\n int nTessPoints = min(max((int)(curvature * 16.0f), 4), MAX_TESSELLATION);\n\n if (bLines[lidx].vertexPos == NULL) {\n bLines[lidx].nVertices = nTessPoints;\n cudaMalloc((void **)&bLines[lidx].vertexPos, nTessPoints * sizeof(float2));\n }\n\n computeBezierLinePositions<<>>(\n lidx, bLines, bLines[lidx].nVertices);\n }\n}\n\n__global__ void freeVertexMem(BezierLine *bLines, int nLines)\n{\n int lidx = threadIdx.x + blockDim.x * blockIdx.x;\n\n if (lidx < nLines)\n cudaFree(bLines[lidx].vertexPos);\n}\n\nunsigned int checkCapableSM35Device(int argc, char **argv)\n{\n // Get device properties\n cudaDeviceProp properties;\n int device_count = 0, device = -1;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"device\")) {\n device = getCmdLineArgumentInt(argc, (const char **)argv, \"device\");\n\n cudaDeviceProp properties;\n checkCudaErrors(cudaGetDeviceProperties(&properties, device));\n\n if (properties.major > 3 || (properties.major == 3 && properties.minor >= 5)) {\n printf(\"Running on GPU %d (%s)\\n\", device, properties.name);\n }\n else {\n printf(\"cdpBezierTessellation requires GPU devices with compute SM 3.5 or \"\n \"higher.\");\n printf(\"Current GPU device has compute SM %d.%d. Exiting...\\n\", properties.major, properties.minor);\n return EXIT_FAILURE;\n }\n }\n else {\n checkCudaErrors(cudaGetDeviceCount(&device_count));\n\n for (int i = 0; i < device_count; ++i) {\n checkCudaErrors(cudaGetDeviceProperties(&properties, i));\n\n if (properties.major > 3 || (properties.major == 3 && properties.minor >= 5)) {\n device = i;\n printf(\"Running on GPU %d (%s)\\n\", i, properties.name);\n break;\n }\n\n printf(\"GPU %d %s does not support CUDA Dynamic Parallelism\\n\", i, properties.name);\n }\n }\n if (device == -1) {\n fprintf(stderr,\n \"cdpBezierTessellation requires GPU devices with compute SM 3.5 or \"\n \"higher. Exiting...\\n\");\n return EXIT_WAIVED;\n }\n\n return EXIT_SUCCESS;\n}\n\n#define N_LINES 256\n#define BLOCK_DIM 64\nint main(int argc, char **argv)\n{\n BezierLine *bLines_h = new BezierLine[N_LINES];\n\n float2 last = {0, 0};\n\n for (int i = 0; i < N_LINES; i++) {\n bLines_h[i].CP[0] = last;\n\n for (int j = 1; j < 3; j++) {\n bLines_h[i].CP[j].x = (float)rand() / (float)RAND_MAX;\n bLines_h[i].CP[j].y = (float)rand() / (float)RAND_MAX;\n }\n\n last = bLines_h[i].CP[2];\n bLines_h[i].vertexPos = NULL;\n bLines_h[i].nVertices = 0;\n }\n\n unsigned int sm35Ret = checkCapableSM35Device(argc, argv);\n if (sm35Ret != EXIT_SUCCESS) {\n exit(sm35Ret);\n }\n\n BezierLine *bLines_d;\n checkCudaErrors(cudaMalloc((void **)&bLines_d, N_LINES * sizeof(BezierLine)));\n checkCudaErrors(cudaMemcpy(bLines_d, bLines_h, N_LINES * sizeof(BezierLine), cudaMemcpyHostToDevice));\n printf(\"Computing Bezier Lines (CUDA Dynamic Parallelism Version) ... \");\n computeBezierLinesCDP<<<(unsigned int)ceil((float)N_LINES / (float)BLOCK_DIM), BLOCK_DIM>>>(bLines_d, N_LINES);\n printf(\"Done!\\n\");\n\n // Do something to draw the lines here\n\n freeVertexMem<<<(unsigned int)ceil((float)N_LINES / (float)BLOCK_DIM), BLOCK_DIM>>>(bLines_d, N_LINES);\n checkCudaErrors(cudaFree(bLines_d));\n delete[] bLines_h;\n\n exit(EXIT_SUCCESS);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n__forceinline__ __device__ float2 operator+(float2 a, float2 b)\n{\n float2 c;\n c.x = a.x + b.x;\n c.y = a.y + b.y;\n return c;\n}\n\n__forceinline__ __device__ float2 operator-(float2 a, float2 b)\n{\n float2 c;\n c.x = a.x - b.x;\n c.y = a.y - b.y;\n return c;\n}\n\n__forceinline__ __device__ float2 operator*(float a, float2 b)\n{\n float2 c;\n c.x = a * b.x;\n c.y = a * b.y;\n return c;\n}\n\n__forceinline__ __device__ float length(float2 a) { return sqrtf(a.x * a.x + a.y * a.y); }\n\n#define MAX_TESSELLATION 32\nstruct BezierLine\n{\n float2 CP[3];\n float2 *vertexPos;\n int nVertices;\n};\n\n__global__ void computeBezierLinePositions(int lidx, BezierLine *bLines, int nTessPoints)\n{\n int idx = threadIdx.x + blockDim.x * blockIdx.x;\n\n if (idx < nTessPoints) {\n float u = (float)idx / (float)(nTessPoints - 1);\n float omu = 1.0f - u;\n\n float B3u[3];\n\n B3u[0] = omu * omu;\n B3u[1] = 2.0f * u * omu;\n B3u[2] = u * u;\n\n float2 position = {0, 0};\n\n for (int i = 0; i < 3; i++) {\n position = position + B3u[i] * bLines[lidx].CP[i];\n }\n\n bLines[lidx].vertexPos[idx] = position;\n }\n}\n\n__global__ void computeBezierLinesCDP(BezierLine *bLines, int nLines)\n{\n int lidx = threadIdx.x + blockDim.x * blockIdx.x;\n\n if (lidx < nLines) {\n float curvature = length(bLines[lidx].CP[1] - 0.5f * (bLines[lidx].CP[0] + bLines[lidx].CP[2]))\n / length(bLines[lidx].CP[2] - bLines[lidx].CP[0]);\n int nTessPoints = min(max((int)(curvature * 16.0f), 4), MAX_TESSELLATION);\n\n if (bLines[lidx].vertexPos == NULL) {\n bLines[lidx].nVertices = nTessPoints;\n hipMalloc((void **)&bLines[lidx].vertexPos, nTessPoints * sizeof(float2));\n }\n\n computeBezierLinePositions<<>>(\n lidx, bLines, bLines[lidx].nVertices);\n }\n}\n\n__global__ void freeVertexMem(BezierLine *bLines, int nLines)\n{\n int lidx = threadIdx.x + blockDim.x * blockIdx.x;\n\n if (lidx < nLines)\n hipFree(bLines[lidx].vertexPos);\n}\n\nunsigned int checkCapableSM35Device(int argc, char **argv)\n{\n // Get device properties\n hipDeviceProp_t properties;\n int device_count = 0, device = -1;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"device\")) {\n device = getCmdLineArgumentInt(argc, (const char **)argv, \"device\");\n\n hipDeviceProp_t properties;\n checkCudaErrors(hipGetDeviceProperties(&properties, device));\n\n if (properties.major > 3 || (properties.major == 3 && properties.minor >= 5)) {\n printf(\"Running on GPU %d (%s)\\n\", device, properties.name);\n }\n else {\n printf(\"cdpBezierTessellation requires GPU devices with compute SM 3.5 or \"\n \"higher.\");\n printf(\"Current GPU device has compute SM %d.%d. Exiting...\\n\", properties.major, properties.minor);\n return EXIT_FAILURE;\n }\n }\n else {\n checkCudaErrors(hipGetDeviceCount(&device_count));\n\n for (int i = 0; i < device_count; ++i) {\n checkCudaErrors(hipGetDeviceProperties(&properties, i));\n\n if (properties.major > 3 || (properties.major == 3 && properties.minor >= 5)) {\n device = i;\n printf(\"Running on GPU %d (%s)\\n\", i, properties.name);\n break;\n }\n\n printf(\"GPU %d %s does not support CUDA Dynamic Parallelism\\n\", i, properties.name);\n }\n }\n if (device == -1) {\n fprintf(stderr,\n \"cdpBezierTessellation requires GPU devices with compute SM 3.5 or \"\n \"higher. Exiting...\\n\");\n return EXIT_WAIVED;\n }\n\n return EXIT_SUCCESS;\n}\n\n#define N_LINES 256\n#define BLOCK_DIM 64\nint main(int argc, char **argv)\n{\n BezierLine *bLines_h = new BezierLine[N_LINES];\n\n float2 last = {0, 0};\n\n for (int i = 0; i < N_LINES; i++) {\n bLines_h[i].CP[0] = last;\n\n for (int j = 1; j < 3; j++) {\n bLines_h[i].CP[j].x = (float)rand() / (float)RAND_MAX;\n bLines_h[i].CP[j].y = (float)rand() / (float)RAND_MAX;\n }\n\n last = bLines_h[i].CP[2];\n bLines_h[i].vertexPos = NULL;\n bLines_h[i].nVertices = 0;\n }\n\n unsigned int sm35Ret = checkCapableSM35Device(argc, argv);\n if (sm35Ret != EXIT_SUCCESS) {\n exit(sm35Ret);\n }\n\n BezierLine *bLines_d;\n checkCudaErrors(hipMalloc((void **)&bLines_d, N_LINES * sizeof(BezierLine)));\n checkCudaErrors(hipMemcpy(bLines_d, bLines_h, N_LINES * sizeof(BezierLine), hipMemcpyHostToDevice));\n printf(\"Computing Bezier Lines (CUDA Dynamic Parallelism Version) ... \");\n computeBezierLinesCDP<<<(unsigned int)ceil((float)N_LINES / (float)BLOCK_DIM), BLOCK_DIM>>>(bLines_d, N_LINES);\n printf(\"Done!\\n\");\n\n // Do something to draw the lines here\n\n freeVertexMem<<<(unsigned int)ceil((float)N_LINES / (float)BLOCK_DIM), BLOCK_DIM>>>(bLines_d, N_LINES);\n checkCudaErrors(hipFree(bLines_d));\n delete[] bLines_h;\n\n exit(EXIT_SUCCESS);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/cdpBezierTessellation/BezierLineCDP.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#define THREADS_PER_BLOCK 512\n#define GRAPH_LAUNCH_ITERATIONS 3\n\ntypedef struct callBackData\n{\n const char *fn_name;\n double *data;\n} callBackData_t;\n\n__global__ void reduce(float *inputVec, double *outputVec, size_t inputSize, size_t outputSize)\n{\n __shared__ double tmp[THREADS_PER_BLOCK];\n\n cg::thread_block cta = cg::this_thread_block();\n size_t globaltid = blockIdx.x * blockDim.x + threadIdx.x;\n\n double temp_sum = 0.0;\n for (int i = globaltid; i < inputSize; i += gridDim.x * blockDim.x) {\n temp_sum += (double)inputVec[i];\n }\n tmp[cta.thread_rank()] = temp_sum;\n\n cg::sync(cta);\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n double beta = temp_sum;\n double temp;\n\n for (int i = tile32.size() / 2; i > 0; i >>= 1) {\n if (tile32.thread_rank() < i) {\n temp = tmp[cta.thread_rank() + i];\n beta += temp;\n tmp[cta.thread_rank()] = beta;\n }\n cg::sync(tile32);\n }\n cg::sync(cta);\n\n if (cta.thread_rank() == 0 && blockIdx.x < outputSize) {\n beta = 0.0;\n for (int i = 0; i < cta.size(); i += tile32.size()) {\n beta += tmp[i];\n }\n outputVec[blockIdx.x] = beta;\n }\n}\n\n__global__ void reduceFinal(double *inputVec, double *result, size_t inputSize)\n{\n __shared__ double tmp[THREADS_PER_BLOCK];\n\n cg::thread_block cta = cg::this_thread_block();\n size_t globaltid = blockIdx.x * blockDim.x + threadIdx.x;\n\n double temp_sum = 0.0;\n for (int i = globaltid; i < inputSize; i += gridDim.x * blockDim.x) {\n temp_sum += (double)inputVec[i];\n }\n tmp[cta.thread_rank()] = temp_sum;\n\n cg::sync(cta);\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n // do reduction in shared mem\n if ((blockDim.x >= 512) && (cta.thread_rank() < 256)) {\n tmp[cta.thread_rank()] = temp_sum = temp_sum + tmp[cta.thread_rank() + 256];\n }\n\n cg::sync(cta);\n\n if ((blockDim.x >= 256) && (cta.thread_rank() < 128)) {\n tmp[cta.thread_rank()] = temp_sum = temp_sum + tmp[cta.thread_rank() + 128];\n }\n\n cg::sync(cta);\n\n if ((blockDim.x >= 128) && (cta.thread_rank() < 64)) {\n tmp[cta.thread_rank()] = temp_sum = temp_sum + tmp[cta.thread_rank() + 64];\n }\n\n cg::sync(cta);\n\n if (cta.thread_rank() < 32) {\n // Fetch final intermediate sum from 2nd warp\n if (blockDim.x >= 64)\n temp_sum += tmp[cta.thread_rank() + 32];\n // Reduce final warp using shuffle\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n temp_sum += tile32.shfl_down(temp_sum, offset);\n }\n }\n // write result for this block to global mem\n if (cta.thread_rank() == 0)\n result[0] = temp_sum;\n}\n\nvoid init_input(float *a, size_t size)\n{\n for (size_t i = 0; i < size; i++)\n a[i] = (rand() & 0xFF) / (float)RAND_MAX;\n}\n\nvoid CUDART_CB myHostNodeCallback(void *data)\n{\n // Check status of GPU after stream operations are done\n callBackData_t *tmp = (callBackData_t *)(data);\n // checkCudaErrors(tmp->status);\n\n double *result = (double *)(tmp->data);\n char *function = (char *)(tmp->fn_name);\n printf(\"[%s] Host callback final reduced sum = %lf\\n\", function, *result);\n *result = 0.0; // reset the result\n}\n\nvoid cudaGraphsManual(float *inputVec_h,\n float *inputVec_d,\n double *outputVec_d,\n double *result_d,\n size_t inputSize,\n size_t numOfBlocks)\n{\n cudaStream_t streamForGraph;\n cudaGraph_t graph;\n std::vector nodeDependencies;\n cudaGraphNode_t memcpyNode, kernelNode, memsetNode;\n double result_h = 0.0;\n\n checkCudaErrors(cudaStreamCreate(&streamForGraph));\n\n cudaKernelNodeParams kernelNodeParams = {0};\n cudaMemcpy3DParms memcpyParams = {0};\n cudaMemsetParams memsetParams = {0};\n\n memcpyParams.srcArray = NULL;\n memcpyParams.srcPos = make_cudaPos(0, 0, 0);\n memcpyParams.srcPtr = make_cudaPitchedPtr(inputVec_h, sizeof(float) * inputSize, inputSize, 1);\n memcpyParams.dstArray = NULL;\n memcpyParams.dstPos = make_cudaPos(0, 0, 0);\n memcpyParams.dstPtr = make_cudaPitchedPtr(inputVec_d, sizeof(float) * inputSize, inputSize, 1);\n memcpyParams.extent = make_cudaExtent(sizeof(float) * inputSize, 1, 1);\n memcpyParams.kind = cudaMemcpyHostToDevice;\n\n memsetParams.dst = (void *)outputVec_d;\n memsetParams.value = 0;\n memsetParams.pitch = 0;\n memsetParams.elementSize = sizeof(float); // elementSize can be max 4 bytes\n memsetParams.width = numOfBlocks * 2;\n memsetParams.height = 1;\n\n checkCudaErrors(cudaGraphCreate(&graph, 0));\n checkCudaErrors(cudaGraphAddMemcpyNode(&memcpyNode, graph, NULL, 0, &memcpyParams));\n checkCudaErrors(cudaGraphAddMemsetNode(&memsetNode, graph, NULL, 0, &memsetParams));\n\n nodeDependencies.push_back(memsetNode);\n nodeDependencies.push_back(memcpyNode);\n\n void *kernelArgs[4] = {(void *)&inputVec_d, (void *)&outputVec_d, &inputSize, &numOfBlocks};\n\n kernelNodeParams.func = (void *)reduce;\n kernelNodeParams.gridDim = dim3(numOfBlocks, 1, 1);\n kernelNodeParams.blockDim = dim3(THREADS_PER_BLOCK, 1, 1);\n kernelNodeParams.sharedMemBytes = 0;\n kernelNodeParams.kernelParams = (void **)kernelArgs;\n kernelNodeParams.extra = NULL;\n\n checkCudaErrors(cudaGraphAddKernelNode(\n &kernelNode, graph, nodeDependencies.data(), nodeDependencies.size(), &kernelNodeParams));\n\n nodeDependencies.clear();\n nodeDependencies.push_back(kernelNode);\n\n memset(&memsetParams, 0, sizeof(memsetParams));\n memsetParams.dst = result_d;\n memsetParams.value = 0;\n memsetParams.elementSize = sizeof(float);\n memsetParams.width = 2;\n memsetParams.height = 1;\n checkCudaErrors(cudaGraphAddMemsetNode(&memsetNode, graph, NULL, 0, &memsetParams));\n\n nodeDependencies.push_back(memsetNode);\n\n memset(&kernelNodeParams, 0, sizeof(kernelNodeParams));\n kernelNodeParams.func = (void *)reduceFinal;\n kernelNodeParams.gridDim = dim3(1, 1, 1);\n kernelNodeParams.blockDim = dim3(THREADS_PER_BLOCK, 1, 1);\n kernelNodeParams.sharedMemBytes = 0;\n void *kernelArgs2[3] = {(void *)&outputVec_d, (void *)&result_d, &numOfBlocks};\n kernelNodeParams.kernelParams = kernelArgs2;\n kernelNodeParams.extra = NULL;\n\n checkCudaErrors(cudaGraphAddKernelNode(\n &kernelNode, graph, nodeDependencies.data(), nodeDependencies.size(), &kernelNodeParams));\n nodeDependencies.clear();\n nodeDependencies.push_back(kernelNode);\n\n memset(&memcpyParams, 0, sizeof(memcpyParams));\n\n memcpyParams.srcArray = NULL;\n memcpyParams.srcPos = make_cudaPos(0, 0, 0);\n memcpyParams.srcPtr = make_cudaPitchedPtr(result_d, sizeof(double), 1, 1);\n memcpyParams.dstArray = NULL;\n memcpyParams.dstPos = make_cudaPos(0, 0, 0);\n memcpyParams.dstPtr = make_cudaPitchedPtr(&result_h, sizeof(double), 1, 1);\n memcpyParams.extent = make_cudaExtent(sizeof(double), 1, 1);\n memcpyParams.kind = cudaMemcpyDeviceToHost;\n checkCudaErrors(\n cudaGraphAddMemcpyNode(&memcpyNode, graph, nodeDependencies.data(), nodeDependencies.size(), &memcpyParams));\n nodeDependencies.clear();\n nodeDependencies.push_back(memcpyNode);\n\n cudaGraphNode_t hostNode;\n cudaHostNodeParams hostParams = {0};\n hostParams.fn = myHostNodeCallback;\n callBackData_t hostFnData;\n hostFnData.data = &result_h;\n hostFnData.fn_name = \"cudaGraphsManual\";\n hostParams.userData = &hostFnData;\n\n checkCudaErrors(\n cudaGraphAddHostNode(&hostNode, graph, nodeDependencies.data(), nodeDependencies.size(), &hostParams));\n\n cudaGraphNode_t *nodes = NULL;\n size_t numNodes = 0;\n checkCudaErrors(cudaGraphGetNodes(graph, nodes, &numNodes));\n printf(\"\\nNum of nodes in the graph created manually = %zu\\n\", numNodes);\n\n cudaGraphExec_t graphExec;\n checkCudaErrors(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n cudaGraph_t clonedGraph;\n cudaGraphExec_t clonedGraphExec;\n checkCudaErrors(cudaGraphClone(&clonedGraph, graph));\n checkCudaErrors(cudaGraphInstantiate(&clonedGraphExec, clonedGraph, NULL, NULL, 0));\n\n for (int i = 0; i < GRAPH_LAUNCH_ITERATIONS; i++) {\n checkCudaErrors(cudaGraphLaunch(graphExec, streamForGraph));\n }\n\n checkCudaErrors(cudaStreamSynchronize(streamForGraph));\n\n printf(\"Cloned Graph Output.. \\n\");\n for (int i = 0; i < GRAPH_LAUNCH_ITERATIONS; i++) {\n checkCudaErrors(cudaGraphLaunch(clonedGraphExec, streamForGraph));\n }\n checkCudaErrors(cudaStreamSynchronize(streamForGraph));\n\n checkCudaErrors(cudaGraphExecDestroy(graphExec));\n checkCudaErrors(cudaGraphExecDestroy(clonedGraphExec));\n checkCudaErrors(cudaGraphDestroy(graph));\n checkCudaErrors(cudaGraphDestroy(clonedGraph));\n checkCudaErrors(cudaStreamDestroy(streamForGraph));\n}\n\nvoid cudaGraphsUsingStreamCapture(float *inputVec_h,\n float *inputVec_d,\n double *outputVec_d,\n double *result_d,\n size_t inputSize,\n size_t numOfBlocks)\n{\n cudaStream_t stream1, stream2, stream3, streamForGraph;\n cudaEvent_t forkStreamEvent, memsetEvent1, memsetEvent2;\n cudaGraph_t graph;\n double result_h = 0.0;\n\n checkCudaErrors(cudaStreamCreate(&stream1));\n checkCudaErrors(cudaStreamCreate(&stream2));\n checkCudaErrors(cudaStreamCreate(&stream3));\n checkCudaErrors(cudaStreamCreate(&streamForGraph));\n\n checkCudaErrors(cudaEventCreate(&forkStreamEvent));\n checkCudaErrors(cudaEventCreate(&memsetEvent1));\n checkCudaErrors(cudaEventCreate(&memsetEvent2));\n\n checkCudaErrors(cudaStreamBeginCapture(stream1, cudaStreamCaptureModeGlobal));\n\n checkCudaErrors(cudaEventRecord(forkStreamEvent, stream1));\n checkCudaErrors(cudaStreamWaitEvent(stream2, forkStreamEvent, 0));\n checkCudaErrors(cudaStreamWaitEvent(stream3, forkStreamEvent, 0));\n\n checkCudaErrors(cudaMemcpyAsync(inputVec_d, inputVec_h, sizeof(float) * inputSize, cudaMemcpyDefault, stream1));\n\n checkCudaErrors(cudaMemsetAsync(outputVec_d, 0, sizeof(double) * numOfBlocks, stream2));\n\n checkCudaErrors(cudaEventRecord(memsetEvent1, stream2));\n\n checkCudaErrors(cudaMemsetAsync(result_d, 0, sizeof(double), stream3));\n checkCudaErrors(cudaEventRecord(memsetEvent2, stream3));\n\n checkCudaErrors(cudaStreamWaitEvent(stream1, memsetEvent1, 0));\n\n reduce<<>>(inputVec_d, outputVec_d, inputSize, numOfBlocks);\n\n checkCudaErrors(cudaStreamWaitEvent(stream1, memsetEvent2, 0));\n\n reduceFinal<<<1, THREADS_PER_BLOCK, 0, stream1>>>(outputVec_d, result_d, numOfBlocks);\n checkCudaErrors(cudaMemcpyAsync(&result_h, result_d, sizeof(double), cudaMemcpyDefault, stream1));\n\n callBackData_t hostFnData = {0};\n hostFnData.data = &result_h;\n hostFnData.fn_name = \"cudaGraphsUsingStreamCapture\";\n cudaHostFn_t fn = myHostNodeCallback;\n checkCudaErrors(cudaLaunchHostFunc(stream1, fn, &hostFnData));\n checkCudaErrors(cudaStreamEndCapture(stream1, &graph));\n\n cudaGraphNode_t *nodes = NULL;\n size_t numNodes = 0;\n checkCudaErrors(cudaGraphGetNodes(graph, nodes, &numNodes));\n printf(\"\\nNum of nodes in the graph created using stream capture API = %zu\\n\", numNodes);\n\n cudaGraphExec_t graphExec;\n checkCudaErrors(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n cudaGraph_t clonedGraph;\n cudaGraphExec_t clonedGraphExec;\n checkCudaErrors(cudaGraphClone(&clonedGraph, graph));\n checkCudaErrors(cudaGraphInstantiate(&clonedGraphExec, clonedGraph, NULL, NULL, 0));\n\n for (int i = 0; i < GRAPH_LAUNCH_ITERATIONS; i++) {\n checkCudaErrors(cudaGraphLaunch(graphExec, streamForGraph));\n }\n\n checkCudaErrors(cudaStreamSynchronize(streamForGraph));\n\n printf(\"Cloned Graph Output.. \\n\");\n for (int i = 0; i < GRAPH_LAUNCH_ITERATIONS; i++) {\n checkCudaErrors(cudaGraphLaunch(clonedGraphExec, streamForGraph));\n }\n\n checkCudaErrors(cudaStreamSynchronize(streamForGraph));\n\n checkCudaErrors(cudaGraphExecDestroy(graphExec));\n checkCudaErrors(cudaGraphExecDestroy(clonedGraphExec));\n checkCudaErrors(cudaGraphDestroy(graph));\n checkCudaErrors(cudaGraphDestroy(clonedGraph));\n checkCudaErrors(cudaStreamDestroy(stream1));\n checkCudaErrors(cudaStreamDestroy(stream2));\n checkCudaErrors(cudaStreamDestroy(streamForGraph));\n}\n\nint main(int argc, char **argv)\n{\n size_t size = 1 << 24; // number of elements to reduce\n size_t maxBlocks = 512;\n\n // This will pick the best possible CUDA capable device\n int devID = findCudaDevice(argc, (const char **)argv);\n\n printf(\"%zu elements\\n\", size);\n printf(\"threads per block = %d\\n\", THREADS_PER_BLOCK);\n printf(\"Graph Launch iterations = %d\\n\", GRAPH_LAUNCH_ITERATIONS);\n\n float *inputVec_d = NULL, *inputVec_h = NULL;\n double *outputVec_d = NULL, *result_d;\n\n checkCudaErrors(cudaMallocHost(&inputVec_h, sizeof(float) * size));\n checkCudaErrors(cudaMalloc(&inputVec_d, sizeof(float) * size));\n checkCudaErrors(cudaMalloc(&outputVec_d, sizeof(double) * maxBlocks));\n checkCudaErrors(cudaMalloc(&result_d, sizeof(double)));\n\n init_input(inputVec_h, size);\n\n cudaGraphsManual(inputVec_h, inputVec_d, outputVec_d, result_d, size, maxBlocks);\n cudaGraphsUsingStreamCapture(inputVec_h, inputVec_d, outputVec_d, result_d, size, maxBlocks);\n\n checkCudaErrors(cudaFree(inputVec_d));\n checkCudaErrors(cudaFree(outputVec_d));\n checkCudaErrors(cudaFree(result_d));\n checkCudaErrors(cudaFreeHost(inputVec_h));\n return EXIT_SUCCESS;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#define THREADS_PER_BLOCK 512\n#define GRAPH_LAUNCH_ITERATIONS 3\n\ntypedef struct callBackData\n{\n const char *fn_name;\n double *data;\n} callBackData_t;\n\n__global__ void reduce(float *inputVec, double *outputVec, size_t inputSize, size_t outputSize)\n{\n __shared__ double tmp[THREADS_PER_BLOCK];\n\n cg::thread_block cta = cg::this_thread_block();\n size_t globaltid = blockIdx.x * blockDim.x + threadIdx.x;\n\n double temp_sum = 0.0;\n for (int i = globaltid; i < inputSize; i += gridDim.x * blockDim.x) {\n temp_sum += (double)inputVec[i];\n }\n tmp[cta.thread_rank()] = temp_sum;\n\n cg::sync(cta);\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n double beta = temp_sum;\n double temp;\n\n for (int i = tile32.size() / 2; i > 0; i >>= 1) {\n if (tile32.thread_rank() < i) {\n temp = tmp[cta.thread_rank() + i];\n beta += temp;\n tmp[cta.thread_rank()] = beta;\n }\n cg::sync(tile32);\n }\n cg::sync(cta);\n\n if (cta.thread_rank() == 0 && blockIdx.x < outputSize) {\n beta = 0.0;\n for (int i = 0; i < cta.size(); i += tile32.size()) {\n beta += tmp[i];\n }\n outputVec[blockIdx.x] = beta;\n }\n}\n\n__global__ void reduceFinal(double *inputVec, double *result, size_t inputSize)\n{\n __shared__ double tmp[THREADS_PER_BLOCK];\n\n cg::thread_block cta = cg::this_thread_block();\n size_t globaltid = blockIdx.x * blockDim.x + threadIdx.x;\n\n double temp_sum = 0.0;\n for (int i = globaltid; i < inputSize; i += gridDim.x * blockDim.x) {\n temp_sum += (double)inputVec[i];\n }\n tmp[cta.thread_rank()] = temp_sum;\n\n cg::sync(cta);\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n // do reduction in shared mem\n if ((blockDim.x >= 512) && (cta.thread_rank() < 256)) {\n tmp[cta.thread_rank()] = temp_sum = temp_sum + tmp[cta.thread_rank() + 256];\n }\n\n cg::sync(cta);\n\n if ((blockDim.x >= 256) && (cta.thread_rank() < 128)) {\n tmp[cta.thread_rank()] = temp_sum = temp_sum + tmp[cta.thread_rank() + 128];\n }\n\n cg::sync(cta);\n\n if ((blockDim.x >= 128) && (cta.thread_rank() < 64)) {\n tmp[cta.thread_rank()] = temp_sum = temp_sum + tmp[cta.thread_rank() + 64];\n }\n\n cg::sync(cta);\n\n if (cta.thread_rank() < 32) {\n // Fetch final intermediate sum from 2nd warp\n if (blockDim.x >= 64)\n temp_sum += tmp[cta.thread_rank() + 32];\n // Reduce final warp using shuffle\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n temp_sum += tile32.shfl_down(temp_sum, offset);\n }\n }\n // write result for this block to global mem\n if (cta.thread_rank() == 0)\n result[0] = temp_sum;\n}\n\nvoid init_input(float *a, size_t size)\n{\n for (size_t i = 0; i < size; i++)\n a[i] = (rand() & 0xFF) / (float)RAND_MAX;\n}\n\nvoid CUDART_CB myHostNodeCallback(void *data)\n{\n // Check status of GPU after stream operations are done\n callBackData_t *tmp = (callBackData_t *)(data);\n // checkCudaErrors(tmp->status);\n\n double *result = (double *)(tmp->data);\n char *function = (char *)(tmp->fn_name);\n printf(\"[%s] Host callback final reduced sum = %lf\\n\", function, *result);\n *result = 0.0; // reset the result\n}\n\nvoid cudaGraphsManual(float *inputVec_h,\n float *inputVec_d,\n double *outputVec_d,\n double *result_d,\n size_t inputSize,\n size_t numOfBlocks)\n{\n hipStream_t streamForGraph;\n hipGraph_t graph;\n std::vector nodeDependencies;\n hipGraphNode_t memcpyNode, kernelNode, memsetNode;\n double result_h = 0.0;\n\n checkCudaErrors(hipStreamCreate(&streamForGraph));\n\n hipKernelNodeParams kernelNodeParams = {0};\n hipMemcpy3DParms memcpyParams = {0};\n hipMemsetParams memsetParams = {0};\n\n memcpyParams.srcArray = NULL;\n memcpyParams.srcPos = make_hipPos(0, 0, 0);\n memcpyParams.srcPtr = make_hipPitchedPtr(inputVec_h, sizeof(float) * inputSize, inputSize, 1);\n memcpyParams.dstArray = NULL;\n memcpyParams.dstPos = make_hipPos(0, 0, 0);\n memcpyParams.dstPtr = make_hipPitchedPtr(inputVec_d, sizeof(float) * inputSize, inputSize, 1);\n memcpyParams.extent = make_hipExtent(sizeof(float) * inputSize, 1, 1);\n memcpyParams.kind = hipMemcpyHostToDevice;\n\n memsetParams.dst = (void *)outputVec_d;\n memsetParams.value = 0;\n memsetParams.pitch = 0;\n memsetParams.elementSize = sizeof(float); // elementSize can be max 4 bytes\n memsetParams.width = numOfBlocks * 2;\n memsetParams.height = 1;\n\n checkCudaErrors(hipGraphCreate(&graph, 0));\n checkCudaErrors(hipGraphAddMemcpyNode(&memcpyNode, graph, NULL, 0, &memcpyParams));\n checkCudaErrors(hipGraphAddMemsetNode(&memsetNode, graph, NULL, 0, &memsetParams));\n\n nodeDependencies.push_back(memsetNode);\n nodeDependencies.push_back(memcpyNode);\n\n void *kernelArgs[4] = {(void *)&inputVec_d, (void *)&outputVec_d, &inputSize, &numOfBlocks};\n\n kernelNodeParams.func = (void *)reduce;\n kernelNodeParams.gridDim = dim3(numOfBlocks, 1, 1);\n kernelNodeParams.blockDim = dim3(THREADS_PER_BLOCK, 1, 1);\n kernelNodeParams.sharedMemBytes = 0;\n kernelNodeParams.kernelParams = (void **)kernelArgs;\n kernelNodeParams.extra = NULL;\n\n checkCudaErrors(hipGraphAddKernelNode(\n &kernelNode, graph, nodeDependencies.data(), nodeDependencies.size(), &kernelNodeParams));\n\n nodeDependencies.clear();\n nodeDependencies.push_back(kernelNode);\n\n memset(&memsetParams, 0, sizeof(memsetParams));\n memsetParams.dst = result_d;\n memsetParams.value = 0;\n memsetParams.elementSize = sizeof(float);\n memsetParams.width = 2;\n memsetParams.height = 1;\n checkCudaErrors(hipGraphAddMemsetNode(&memsetNode, graph, NULL, 0, &memsetParams));\n\n nodeDependencies.push_back(memsetNode);\n\n memset(&kernelNodeParams, 0, sizeof(kernelNodeParams));\n kernelNodeParams.func = (void *)reduceFinal;\n kernelNodeParams.gridDim = dim3(1, 1, 1);\n kernelNodeParams.blockDim = dim3(THREADS_PER_BLOCK, 1, 1);\n kernelNodeParams.sharedMemBytes = 0;\n void *kernelArgs2[3] = {(void *)&outputVec_d, (void *)&result_d, &numOfBlocks};\n kernelNodeParams.kernelParams = kernelArgs2;\n kernelNodeParams.extra = NULL;\n\n checkCudaErrors(hipGraphAddKernelNode(\n &kernelNode, graph, nodeDependencies.data(), nodeDependencies.size(), &kernelNodeParams));\n nodeDependencies.clear();\n nodeDependencies.push_back(kernelNode);\n\n memset(&memcpyParams, 0, sizeof(memcpyParams));\n\n memcpyParams.srcArray = NULL;\n memcpyParams.srcPos = make_hipPos(0, 0, 0);\n memcpyParams.srcPtr = make_hipPitchedPtr(result_d, sizeof(double), 1, 1);\n memcpyParams.dstArray = NULL;\n memcpyParams.dstPos = make_hipPos(0, 0, 0);\n memcpyParams.dstPtr = make_hipPitchedPtr(&result_h, sizeof(double), 1, 1);\n memcpyParams.extent = make_hipExtent(sizeof(double), 1, 1);\n memcpyParams.kind = hipMemcpyDeviceToHost;\n checkCudaErrors(\n hipGraphAddMemcpyNode(&memcpyNode, graph, nodeDependencies.data(), nodeDependencies.size(), &memcpyParams));\n nodeDependencies.clear();\n nodeDependencies.push_back(memcpyNode);\n\n hipGraphNode_t hostNode;\n hipHostNodeParams hostParams = {0};\n hostParams.fn = myHostNodeCallback;\n callBackData_t hostFnData;\n hostFnData.data = &result_h;\n hostFnData.fn_name = \"cudaGraphsManual\";\n hostParams.userData = &hostFnData;\n\n checkCudaErrors(\n hipGraphAddHostNode(&hostNode, graph, nodeDependencies.data(), nodeDependencies.size(), &hostParams));\n\n hipGraphNode_t *nodes = NULL;\n size_t numNodes = 0;\n checkCudaErrors(hipGraphGetNodes(graph, nodes, &numNodes));\n printf(\"\\nNum of nodes in the graph created manually = %zu\\n\", numNodes);\n\n hipGraphExec_t graphExec;\n checkCudaErrors(hipGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n hipGraph_t clonedGraph;\n hipGraphExec_t clonedGraphExec;\n checkCudaErrors(hipGraphClone(&clonedGraph, graph));\n checkCudaErrors(hipGraphInstantiate(&clonedGraphExec, clonedGraph, NULL, NULL, 0));\n\n for (int i = 0; i < GRAPH_LAUNCH_ITERATIONS; i++) {\n checkCudaErrors(hipGraphLaunch(graphExec, streamForGraph));\n }\n\n checkCudaErrors(hipStreamSynchronize(streamForGraph));\n\n printf(\"Cloned Graph Output.. \\n\");\n for (int i = 0; i < GRAPH_LAUNCH_ITERATIONS; i++) {\n checkCudaErrors(hipGraphLaunch(clonedGraphExec, streamForGraph));\n }\n checkCudaErrors(hipStreamSynchronize(streamForGraph));\n\n checkCudaErrors(hipGraphExecDestroy(graphExec));\n checkCudaErrors(hipGraphExecDestroy(clonedGraphExec));\n checkCudaErrors(hipGraphDestroy(graph));\n checkCudaErrors(hipGraphDestroy(clonedGraph));\n checkCudaErrors(hipStreamDestroy(streamForGraph));\n}\n\nvoid cudaGraphsUsingStreamCapture(float *inputVec_h,\n float *inputVec_d,\n double *outputVec_d,\n double *result_d,\n size_t inputSize,\n size_t numOfBlocks)\n{\n hipStream_t stream1, stream2, stream3, streamForGraph;\n hipEvent_t forkStreamEvent, memsetEvent1, memsetEvent2;\n hipGraph_t graph;\n double result_h = 0.0;\n\n checkCudaErrors(hipStreamCreate(&stream1));\n checkCudaErrors(hipStreamCreate(&stream2));\n checkCudaErrors(hipStreamCreate(&stream3));\n checkCudaErrors(hipStreamCreate(&streamForGraph));\n\n checkCudaErrors(hipEventCreate(&forkStreamEvent));\n checkCudaErrors(hipEventCreate(&memsetEvent1));\n checkCudaErrors(hipEventCreate(&memsetEvent2));\n\n checkCudaErrors(hipStreamBeginCapture(stream1, hipStreamCaptureModeGlobal));\n\n checkCudaErrors(hipEventRecord(forkStreamEvent, stream1));\n checkCudaErrors(hipStreamWaitEvent(stream2, forkStreamEvent, 0));\n checkCudaErrors(hipStreamWaitEvent(stream3, forkStreamEvent, 0));\n\n checkCudaErrors(hipMemcpyAsync(inputVec_d, inputVec_h, sizeof(float) * inputSize, hipMemcpyDefault, stream1));\n\n checkCudaErrors(hipMemsetAsync(outputVec_d, 0, sizeof(double) * numOfBlocks, stream2));\n\n checkCudaErrors(hipEventRecord(memsetEvent1, stream2));\n\n checkCudaErrors(hipMemsetAsync(result_d, 0, sizeof(double), stream3));\n checkCudaErrors(hipEventRecord(memsetEvent2, stream3));\n\n checkCudaErrors(hipStreamWaitEvent(stream1, memsetEvent1, 0));\n\n reduce<<>>(inputVec_d, outputVec_d, inputSize, numOfBlocks);\n\n checkCudaErrors(hipStreamWaitEvent(stream1, memsetEvent2, 0));\n\n reduceFinal<<<1, THREADS_PER_BLOCK, 0, stream1>>>(outputVec_d, result_d, numOfBlocks);\n checkCudaErrors(hipMemcpyAsync(&result_h, result_d, sizeof(double), hipMemcpyDefault, stream1));\n\n callBackData_t hostFnData = {0};\n hostFnData.data = &result_h;\n hostFnData.fn_name = \"cudaGraphsUsingStreamCapture\";\n hipHostFn_t fn = myHostNodeCallback;\n checkCudaErrors(hipLaunchHostFunc(stream1, fn, &hostFnData));\n checkCudaErrors(hipStreamEndCapture(stream1, &graph));\n\n hipGraphNode_t *nodes = NULL;\n size_t numNodes = 0;\n checkCudaErrors(hipGraphGetNodes(graph, nodes, &numNodes));\n printf(\"\\nNum of nodes in the graph created using stream capture API = %zu\\n\", numNodes);\n\n hipGraphExec_t graphExec;\n checkCudaErrors(hipGraphInstantiate(&graphExec, graph, NULL, NULL, 0));\n\n hipGraph_t clonedGraph;\n hipGraphExec_t clonedGraphExec;\n checkCudaErrors(hipGraphClone(&clonedGraph, graph));\n checkCudaErrors(hipGraphInstantiate(&clonedGraphExec, clonedGraph, NULL, NULL, 0));\n\n for (int i = 0; i < GRAPH_LAUNCH_ITERATIONS; i++) {\n checkCudaErrors(hipGraphLaunch(graphExec, streamForGraph));\n }\n\n checkCudaErrors(hipStreamSynchronize(streamForGraph));\n\n printf(\"Cloned Graph Output.. \\n\");\n for (int i = 0; i < GRAPH_LAUNCH_ITERATIONS; i++) {\n checkCudaErrors(hipGraphLaunch(clonedGraphExec, streamForGraph));\n }\n\n checkCudaErrors(hipStreamSynchronize(streamForGraph));\n\n checkCudaErrors(hipGraphExecDestroy(graphExec));\n checkCudaErrors(hipGraphExecDestroy(clonedGraphExec));\n checkCudaErrors(hipGraphDestroy(graph));\n checkCudaErrors(hipGraphDestroy(clonedGraph));\n checkCudaErrors(hipStreamDestroy(stream1));\n checkCudaErrors(hipStreamDestroy(stream2));\n checkCudaErrors(hipStreamDestroy(streamForGraph));\n}\n\nint main(int argc, char **argv)\n{\n size_t size = 1 << 24; // number of elements to reduce\n size_t maxBlocks = 512;\n\n // This will pick the best possible CUDA capable device\n int devID = findCudaDevice(argc, (const char **)argv);\n\n printf(\"%zu elements\\n\", size);\n printf(\"threads per block = %d\\n\", THREADS_PER_BLOCK);\n printf(\"Graph Launch iterations = %d\\n\", GRAPH_LAUNCH_ITERATIONS);\n\n float *inputVec_d = NULL, *inputVec_h = NULL;\n double *outputVec_d = NULL, *result_d;\n\n checkCudaErrors(hipHostMalloc(&inputVec_h, sizeof(float) * size));\n checkCudaErrors(hipMalloc(&inputVec_d, sizeof(float) * size));\n checkCudaErrors(hipMalloc(&outputVec_d, sizeof(double) * maxBlocks));\n checkCudaErrors(hipMalloc(&result_d, sizeof(double)));\n\n init_input(inputVec_h, size);\n\n cudaGraphsManual(inputVec_h, inputVec_d, outputVec_d, result_d, size, maxBlocks);\n cudaGraphsUsingStreamCapture(inputVec_h, inputVec_d, outputVec_d, result_d, size, maxBlocks);\n\n checkCudaErrors(hipFree(inputVec_d));\n checkCudaErrors(hipFree(outputVec_d));\n checkCudaErrors(hipFree(result_d));\n checkCudaErrors(hipHostFree(inputVec_h));\n return EXIT_SUCCESS;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/3_CUDA_Features/simpleCudaGraphs/simpleCudaGraphs.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* Example showing the use of CUFFT for fast 1D-convolution using FFT. */\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// includes, project\n#include \n#include \n#include \n#include \n#include \n\n// Complex data type\ntypedef float2 Complex;\nstatic __device__ __host__ inline Complex ComplexAdd(Complex, Complex);\nstatic __device__ __host__ inline Complex ComplexScale(Complex, float);\nstatic __device__ __host__ inline Complex ComplexMul(Complex, Complex);\nstatic __global__ void ComplexPointwiseMulAndScale(Complex *, const Complex *, int, float);\n\n// Filtering functions\nvoid Convolve(const Complex *, int, const Complex *, int, Complex *);\n\n// Padding functions\nint PadData(const Complex *, Complex **, int, const Complex *, Complex **, int);\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\nvoid runTest(int argc, char **argv);\n\n// The filter size is assumed to be a number smaller than the signal size\n#define SIGNAL_SIZE 50\n#define FILTER_KERNEL_SIZE 11\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv) { runTest(argc, argv); }\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nvoid runTest(int argc, char **argv)\n{\n printf(\"[simpleCUFFT] is starting...\\n\");\n\n findCudaDevice(argc, (const char **)argv);\n\n // Allocate host memory for the signal\n Complex *h_signal = reinterpret_cast(malloc(sizeof(Complex) * SIGNAL_SIZE));\n\n // Initialize the memory for the signal\n for (unsigned int i = 0; i < SIGNAL_SIZE; ++i) {\n h_signal[i].x = rand() / static_cast(RAND_MAX);\n h_signal[i].y = 0;\n }\n\n // Allocate host memory for the filter\n Complex *h_filter_kernel = reinterpret_cast(malloc(sizeof(Complex) * FILTER_KERNEL_SIZE));\n\n // Initialize the memory for the filter\n for (unsigned int i = 0; i < FILTER_KERNEL_SIZE; ++i) {\n h_filter_kernel[i].x = rand() / static_cast(RAND_MAX);\n h_filter_kernel[i].y = 0;\n }\n\n // Pad signal and filter kernel\n Complex *h_padded_signal;\n Complex *h_padded_filter_kernel;\n int new_size =\n PadData(h_signal, &h_padded_signal, SIGNAL_SIZE, h_filter_kernel, &h_padded_filter_kernel, FILTER_KERNEL_SIZE);\n int mem_size = sizeof(Complex) * new_size;\n\n // Allocate device memory for signal\n Complex *d_signal;\n checkCudaErrors(cudaMalloc(reinterpret_cast(&d_signal), mem_size));\n // Copy host memory to device\n checkCudaErrors(cudaMemcpy(d_signal, h_padded_signal, mem_size, cudaMemcpyHostToDevice));\n\n // Allocate device memory for filter kernel\n Complex *d_filter_kernel;\n checkCudaErrors(cudaMalloc(reinterpret_cast(&d_filter_kernel), mem_size));\n\n // Copy host memory to device\n checkCudaErrors(cudaMemcpy(d_filter_kernel, h_padded_filter_kernel, mem_size, cudaMemcpyHostToDevice));\n\n // CUFFT plan simple API\n cufftHandle plan;\n checkCudaErrors(cufftPlan1d(&plan, new_size, CUFFT_C2C, 1));\n\n // CUFFT plan advanced API\n cufftHandle plan_adv;\n size_t workSize;\n long long int new_size_long = new_size;\n\n checkCudaErrors(cufftCreate(&plan_adv));\n checkCudaErrors(cufftXtMakePlanMany(\n plan_adv, 1, &new_size_long, NULL, 1, 1, CUDA_C_32F, NULL, 1, 1, CUDA_C_32F, 1, &workSize, CUDA_C_32F));\n printf(\"Temporary buffer size %li bytes\\n\", workSize);\n\n // Transform signal and kernel\n printf(\"Transforming signal cufftExecC2C\\n\");\n checkCudaErrors(cufftExecC2C(\n plan, reinterpret_cast(d_signal), reinterpret_cast(d_signal), CUFFT_FORWARD));\n checkCudaErrors(cufftExecC2C(plan_adv,\n reinterpret_cast(d_filter_kernel),\n reinterpret_cast(d_filter_kernel),\n CUFFT_FORWARD));\n\n // Multiply the coefficients together and normalize the result\n printf(\"Launching ComplexPointwiseMulAndScale<<< >>>\\n\");\n ComplexPointwiseMulAndScale<<<32, 256>>>(d_signal, d_filter_kernel, new_size, 1.0f / new_size);\n\n // Check if kernel execution generated and error\n getLastCudaError(\"Kernel execution failed [ ComplexPointwiseMulAndScale ]\");\n\n // Transform signal back\n printf(\"Transforming signal back cufftExecC2C\\n\");\n checkCudaErrors(cufftExecC2C(\n plan, reinterpret_cast(d_signal), reinterpret_cast(d_signal), CUFFT_INVERSE));\n\n // Copy device memory to host\n Complex *h_convolved_signal = h_padded_signal;\n checkCudaErrors(cudaMemcpy(h_convolved_signal, d_signal, mem_size, cudaMemcpyDeviceToHost));\n\n // Allocate host memory for the convolution result\n Complex *h_convolved_signal_ref = reinterpret_cast(malloc(sizeof(Complex) * SIGNAL_SIZE));\n\n // Convolve on the host\n Convolve(h_signal, SIGNAL_SIZE, h_filter_kernel, FILTER_KERNEL_SIZE, h_convolved_signal_ref);\n\n // check result\n bool bTestResult = sdkCompareL2fe(reinterpret_cast(h_convolved_signal_ref),\n reinterpret_cast(h_convolved_signal),\n 2 * SIGNAL_SIZE,\n 1e-5f);\n\n // Destroy CUFFT context\n checkCudaErrors(cufftDestroy(plan));\n checkCudaErrors(cufftDestroy(plan_adv));\n\n // cleanup memory\n free(h_signal);\n free(h_filter_kernel);\n free(h_padded_signal);\n free(h_padded_filter_kernel);\n free(h_convolved_signal_ref);\n checkCudaErrors(cudaFree(d_signal));\n checkCudaErrors(cudaFree(d_filter_kernel));\n\n exit(bTestResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\n// Pad data\nint PadData(const Complex *signal,\n Complex **padded_signal,\n int signal_size,\n const Complex *filter_kernel,\n Complex **padded_filter_kernel,\n int filter_kernel_size)\n{\n int minRadius = filter_kernel_size / 2;\n int maxRadius = filter_kernel_size - minRadius;\n int new_size = signal_size + maxRadius;\n\n // Pad signal\n Complex *new_data = reinterpret_cast(malloc(sizeof(Complex) * new_size));\n memcpy(new_data + 0, signal, signal_size * sizeof(Complex));\n memset(new_data + signal_size, 0, (new_size - signal_size) * sizeof(Complex));\n *padded_signal = new_data;\n\n // Pad filter\n new_data = reinterpret_cast(malloc(sizeof(Complex) * new_size));\n memcpy(new_data + 0, filter_kernel + minRadius, maxRadius * sizeof(Complex));\n memset(new_data + maxRadius, 0, (new_size - filter_kernel_size) * sizeof(Complex));\n memcpy(new_data + new_size - minRadius, filter_kernel, minRadius * sizeof(Complex));\n *padded_filter_kernel = new_data;\n\n return new_size;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Filtering operations\n////////////////////////////////////////////////////////////////////////////////\n\n// Computes convolution on the host\nvoid Convolve(const Complex *signal,\n int signal_size,\n const Complex *filter_kernel,\n int filter_kernel_size,\n Complex *filtered_signal)\n{\n int minRadius = filter_kernel_size / 2;\n int maxRadius = filter_kernel_size - minRadius;\n\n // Loop over output element indices\n for (int i = 0; i < signal_size; ++i) {\n filtered_signal[i].x = filtered_signal[i].y = 0;\n\n // Loop over convolution indices\n for (int j = -maxRadius + 1; j <= minRadius; ++j) {\n int k = i + j;\n\n if (k >= 0 && k < signal_size) {\n filtered_signal[i] =\n ComplexAdd(filtered_signal[i], ComplexMul(signal[k], filter_kernel[minRadius - j]));\n }\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Complex operations\n////////////////////////////////////////////////////////////////////////////////\n\n// Complex addition\nstatic __device__ __host__ inline Complex ComplexAdd(Complex a, Complex b)\n{\n Complex c;\n c.x = a.x + b.x;\n c.y = a.y + b.y;\n return c;\n}\n\n// Complex scale\nstatic __device__ __host__ inline Complex ComplexScale(Complex a, float s)\n{\n Complex c;\n c.x = s * a.x;\n c.y = s * a.y;\n return c;\n}\n\n// Complex multiplication\nstatic __device__ __host__ inline Complex ComplexMul(Complex a, Complex b)\n{\n Complex c;\n c.x = a.x * b.x - a.y * b.y;\n c.y = a.x * b.y + a.y * b.x;\n return c;\n}\n\n// Complex pointwise multiplication\nstatic __global__ void ComplexPointwiseMulAndScale(Complex *a, const Complex *b, int size, float scale)\n{\n const int numThreads = blockDim.x * gridDim.x;\n const int threadID = blockIdx.x * blockDim.x + threadIdx.x;\n\n for (int i = threadID; i < size; i += numThreads) {\n a[i] = ComplexScale(ComplexMul(a[i], b[i]), scale);\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* Example showing the use of CUFFT for fast 1D-convolution using FFT. */\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// includes, project\n#include \n#include \n#include \n#include \n#include \n\n// Complex data type\ntypedef float2 Complex;\nstatic __device__ __host__ inline Complex ComplexAdd(Complex, Complex);\nstatic __device__ __host__ inline Complex ComplexScale(Complex, float);\nstatic __device__ __host__ inline Complex ComplexMul(Complex, Complex);\nstatic __global__ void ComplexPointwiseMulAndScale(Complex *, const Complex *, int, float);\n\n// Filtering functions\nvoid Convolve(const Complex *, int, const Complex *, int, Complex *);\n\n// Padding functions\nint PadData(const Complex *, Complex **, int, const Complex *, Complex **, int);\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\nvoid runTest(int argc, char **argv);\n\n// The filter size is assumed to be a number smaller than the signal size\n#define SIGNAL_SIZE 50\n#define FILTER_KERNEL_SIZE 11\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv) { runTest(argc, argv); }\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a simple test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nvoid runTest(int argc, char **argv)\n{\n printf(\"[simpleCUFFT] is starting...\\n\");\n\n findCudaDevice(argc, (const char **)argv);\n\n // Allocate host memory for the signal\n Complex *h_signal = reinterpret_cast(malloc(sizeof(Complex) * SIGNAL_SIZE));\n\n // Initialize the memory for the signal\n for (unsigned int i = 0; i < SIGNAL_SIZE; ++i) {\n h_signal[i].x = rand() / static_cast(RAND_MAX);\n h_signal[i].y = 0;\n }\n\n // Allocate host memory for the filter\n Complex *h_filter_kernel = reinterpret_cast(malloc(sizeof(Complex) * FILTER_KERNEL_SIZE));\n\n // Initialize the memory for the filter\n for (unsigned int i = 0; i < FILTER_KERNEL_SIZE; ++i) {\n h_filter_kernel[i].x = rand() / static_cast(RAND_MAX);\n h_filter_kernel[i].y = 0;\n }\n\n // Pad signal and filter kernel\n Complex *h_padded_signal;\n Complex *h_padded_filter_kernel;\n int new_size =\n PadData(h_signal, &h_padded_signal, SIGNAL_SIZE, h_filter_kernel, &h_padded_filter_kernel, FILTER_KERNEL_SIZE);\n int mem_size = sizeof(Complex) * new_size;\n\n // Allocate device memory for signal\n Complex *d_signal;\n checkCudaErrors(hipMalloc(reinterpret_cast(&d_signal), mem_size));\n // Copy host memory to device\n checkCudaErrors(hipMemcpy(d_signal, h_padded_signal, mem_size, hipMemcpyHostToDevice));\n\n // Allocate device memory for filter kernel\n Complex *d_filter_kernel;\n checkCudaErrors(hipMalloc(reinterpret_cast(&d_filter_kernel), mem_size));\n\n // Copy host memory to device\n checkCudaErrors(hipMemcpy(d_filter_kernel, h_padded_filter_kernel, mem_size, hipMemcpyHostToDevice));\n\n // CUFFT plan simple API\n hipfftHandle plan;\n checkCudaErrors(hipfftPlan1d(&plan, new_size, HIPFFT_C2C, 1));\n\n // CUFFT plan advanced API\n hipfftHandle plan_adv;\n size_t workSize;\n long long int new_size_long = new_size;\n\n checkCudaErrors(hipfftCreate(&plan_adv));\n checkCudaErrors(hipfftXtMakePlanMany(\n plan_adv, 1, &new_size_long, NULL, 1, 1, HIP_C_32F, NULL, 1, 1, HIP_C_32F, 1, &workSize, HIP_C_32F));\n printf(\"Temporary buffer size %li bytes\\n\", workSize);\n\n // Transform signal and kernel\n printf(\"Transforming signal hipfftExecC2C\\n\");\n checkCudaErrors(hipfftExecC2C(\n plan, reinterpret_cast(d_signal), reinterpret_cast(d_signal), HIPFFT_FORWARD));\n checkCudaErrors(hipfftExecC2C(plan_adv,\n reinterpret_cast(d_filter_kernel),\n reinterpret_cast(d_filter_kernel),\n HIPFFT_FORWARD));\n\n // Multiply the coefficients together and normalize the result\n printf(\"Launching ComplexPointwiseMulAndScale<<< >>>\\n\");\n ComplexPointwiseMulAndScale<<<32, 256>>>(d_signal, d_filter_kernel, new_size, 1.0f / new_size);\n\n // Check if kernel execution generated and error\n getLastCudaError(\"Kernel execution failed [ ComplexPointwiseMulAndScale ]\");\n\n // Transform signal back\n printf(\"Transforming signal back hipfftExecC2C\\n\");\n checkCudaErrors(hipfftExecC2C(\n plan, reinterpret_cast(d_signal), reinterpret_cast(d_signal), HIPFFT_BACKWARD));\n\n // Copy device memory to host\n Complex *h_convolved_signal = h_padded_signal;\n checkCudaErrors(hipMemcpy(h_convolved_signal, d_signal, mem_size, hipMemcpyDeviceToHost));\n\n // Allocate host memory for the convolution result\n Complex *h_convolved_signal_ref = reinterpret_cast(malloc(sizeof(Complex) * SIGNAL_SIZE));\n\n // Convolve on the host\n Convolve(h_signal, SIGNAL_SIZE, h_filter_kernel, FILTER_KERNEL_SIZE, h_convolved_signal_ref);\n\n // check result\n bool bTestResult = sdkCompareL2fe(reinterpret_cast(h_convolved_signal_ref),\n reinterpret_cast(h_convolved_signal),\n 2 * SIGNAL_SIZE,\n 1e-5f);\n\n // Destroy CUFFT context\n checkCudaErrors(hipfftDestroy(plan));\n checkCudaErrors(hipfftDestroy(plan_adv));\n\n // cleanup memory\n free(h_signal);\n free(h_filter_kernel);\n free(h_padded_signal);\n free(h_padded_filter_kernel);\n free(h_convolved_signal_ref);\n checkCudaErrors(hipFree(d_signal));\n checkCudaErrors(hipFree(d_filter_kernel));\n\n exit(bTestResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\n// Pad data\nint PadData(const Complex *signal,\n Complex **padded_signal,\n int signal_size,\n const Complex *filter_kernel,\n Complex **padded_filter_kernel,\n int filter_kernel_size)\n{\n int minRadius = filter_kernel_size / 2;\n int maxRadius = filter_kernel_size - minRadius;\n int new_size = signal_size + maxRadius;\n\n // Pad signal\n Complex *new_data = reinterpret_cast(malloc(sizeof(Complex) * new_size));\n memcpy(new_data + 0, signal, signal_size * sizeof(Complex));\n memset(new_data + signal_size, 0, (new_size - signal_size) * sizeof(Complex));\n *padded_signal = new_data;\n\n // Pad filter\n new_data = reinterpret_cast(malloc(sizeof(Complex) * new_size));\n memcpy(new_data + 0, filter_kernel + minRadius, maxRadius * sizeof(Complex));\n memset(new_data + maxRadius, 0, (new_size - filter_kernel_size) * sizeof(Complex));\n memcpy(new_data + new_size - minRadius, filter_kernel, minRadius * sizeof(Complex));\n *padded_filter_kernel = new_data;\n\n return new_size;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Filtering operations\n////////////////////////////////////////////////////////////////////////////////\n\n// Computes convolution on the host\nvoid Convolve(const Complex *signal,\n int signal_size,\n const Complex *filter_kernel,\n int filter_kernel_size,\n Complex *filtered_signal)\n{\n int minRadius = filter_kernel_size / 2;\n int maxRadius = filter_kernel_size - minRadius;\n\n // Loop over output element indices\n for (int i = 0; i < signal_size; ++i) {\n filtered_signal[i].x = filtered_signal[i].y = 0;\n\n // Loop over convolution indices\n for (int j = -maxRadius + 1; j <= minRadius; ++j) {\n int k = i + j;\n\n if (k >= 0 && k < signal_size) {\n filtered_signal[i] =\n ComplexAdd(filtered_signal[i], ComplexMul(signal[k], filter_kernel[minRadius - j]));\n }\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Complex operations\n////////////////////////////////////////////////////////////////////////////////\n\n// Complex addition\nstatic __device__ __host__ inline Complex ComplexAdd(Complex a, Complex b)\n{\n Complex c;\n c.x = a.x + b.x;\n c.y = a.y + b.y;\n return c;\n}\n\n// Complex scale\nstatic __device__ __host__ inline Complex ComplexScale(Complex a, float s)\n{\n Complex c;\n c.x = s * a.x;\n c.y = s * a.y;\n return c;\n}\n\n// Complex multiplication\nstatic __device__ __host__ inline Complex ComplexMul(Complex a, Complex b)\n{\n Complex c;\n c.x = a.x * b.x - a.y * b.y;\n c.y = a.x * b.y + a.y * b.x;\n return c;\n}\n\n// Complex pointwise multiplication\nstatic __global__ void ComplexPointwiseMulAndScale(Complex *a, const Complex *b, int size, float scale)\n{\n const int numThreads = blockDim.x * gridDim.x;\n const int threadID = blockIdx.x * blockDim.x + threadIdx.x;\n\n for (int i = threadID; i < size; i += numThreads) {\n a[i] = ComplexScale(ComplexMul(a[i], b[i]), scale);\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/4_CUDA_Libraries/simpleCUFFT/simpleCUFFT.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// This sample is an implementation of a simple line-of-sight algorithm:\n// Given a height map and a ray originating at some observation point,\n// it computes all the points along the ray that are visible from the\n// observation point.\n// It is based on the description made in \"Guy E. Blelloch. Vector models\n// for data-parallel computing. MIT Press, 1990\" and uses open source CUDA\n// Thrust Library\n\n#ifdef _WIN32\n#define NOMINMAX\n#endif\n\n// includes, system\n#include \n#include \n#include \n#include \n#include \n\n// includes, project\n#include \n#include \n#include \n\n// includes, library\n#include \n#include \n#include \n#include \n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, types\n\n// Boolean\ntypedef unsigned char Bool;\nenum { False = 0, True = 1 };\n\n// 2D height field\nstruct HeightField\n{\n int width;\n float *height;\n};\n\n// Ray\nstruct Ray\n{\n float3 origin;\n float2 dir;\n int length;\n float oneOverLength;\n};\n\n////////////////////////////////////////////////////////////////////////////////\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\nint runTest(int argc, char **argv);\n__global__ void computeAngles_kernel(const Ray, float *, cudaTextureObject_t);\n__global__ void computeVisibilities_kernel(const float *, const float *, int, Bool *);\nvoid lineOfSight_gold(const HeightField, const Ray, Bool *);\n__device__ __host__ float2 getLocation(const Ray, int);\n__device__ __host__ float getAngle(const Ray, float2, float);\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n int res = runTest(argc, argv);\n\n if (res != 1) {\n printf(\"Test failed!\\n\");\n exit(EXIT_FAILURE);\n }\n\n printf(\"Test passed\\n\");\n exit(EXIT_SUCCESS);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a line-of-sight test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nint runTest(int argc, char **argv)\n{\n ////////////////////////////////////////////////////////////////////////////\n // Device initialization\n\n printf(\"[%s] - Starting...\\n\", argv[0]);\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n findCudaDevice(argc, (const char **)argv);\n\n ////////////////////////////////////////////////////////////////////////////\n // Timer\n\n // Create\n StopWatchInterface *timer;\n sdkCreateTimer(&timer);\n\n // Number of iterations to get accurate timing\n uint numIterations = 100;\n\n ////////////////////////////////////////////////////////////////////////////\n // Height field\n\n HeightField heightField;\n\n // Allocate in host memory\n int2 dim = make_int2(10000, 100);\n heightField.width = dim.x;\n thrust::host_vector height(dim.x * dim.y);\n heightField.height = (float *)&height[0];\n\n //\n // Fill in with an arbitrary sine surface\n for (int x = 0; x < dim.x; ++x)\n for (int y = 0; y < dim.y; ++y) {\n float amp = 0.1f * (x + y);\n float period = 2.0f + amp;\n *(heightField.height + dim.x * y + x) =\n amp * (sinf(sqrtf((float)(x * x + y * y)) * 2.0f * 3.1416f / period) + 1.0f);\n }\n\n // Allocate CUDA array in device memory\n cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc(32, 0, 0, 0, cudaChannelFormatKindFloat);\n cudaArray *heightFieldArray;\n checkCudaErrors(cudaMallocArray(&heightFieldArray, &channelDesc, dim.x, dim.y));\n\n // Initialize device memory\n checkCudaErrors(cudaMemcpy2DToArray(heightFieldArray,\n 0,\n 0,\n heightField.height,\n dim.x * sizeof(float),\n dim.x * sizeof(float),\n dim.y,\n cudaMemcpyHostToDevice));\n\n cudaTextureObject_t heightFieldTex;\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = heightFieldArray;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n texDescr.normalizedCoords = false;\n texDescr.filterMode = cudaFilterModePoint;\n texDescr.addressMode[0] = cudaAddressModeClamp;\n texDescr.addressMode[1] = cudaAddressModeClamp;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&heightFieldTex, &texRes, &texDescr, NULL));\n\n //////////////////////////////////////////////////////////////////////////////\n // Ray (starts at origin and traverses the height field diagonally)\n\n Ray ray;\n ray.origin = make_float3(0, 0, 2.0f);\n int2 dir = make_int2(dim.x - 1, dim.y - 1);\n ray.dir = make_float2((float)dir.x, (float)dir.y);\n ray.length = max(abs(dir.x), abs(dir.y));\n ray.oneOverLength = 1.0f / ray.length;\n\n //////////////////////////////////////////////////////////////////////////////\n // View angles\n\n // Allocate view angles for each point along the ray\n thrust::device_vector d_angles(ray.length);\n\n // Allocate result of max-scan operation on the array of view angles\n thrust::device_vector d_scannedAngles(ray.length);\n\n //////////////////////////////////////////////////////////////////////////////\n // Visibility results\n\n // Allocate visibility results for each point along the ray\n thrust::device_vector d_visibilities(ray.length);\n thrust::host_vector h_visibilities(ray.length);\n thrust::host_vector h_visibilitiesRef(ray.length);\n\n //////////////////////////////////////////////////////////////////////////////\n // Reference solution\n lineOfSight_gold(heightField, ray, (Bool *)&h_visibilitiesRef[0]);\n\n //////////////////////////////////////////////////////////////////////////////\n // Device solution\n\n // Execution configuration\n dim3 block(256);\n dim3 grid((uint)ceil(ray.length / (double)block.x));\n\n // Compute device solution\n printf(\"Line of sight\\n\");\n sdkStartTimer(&timer);\n\n for (uint i = 0; i < numIterations; ++i) {\n // Compute view angle for each point along the ray\n computeAngles_kernel<<>>(ray, thrust::raw_pointer_cast(&d_angles[0]), heightFieldTex);\n getLastCudaError(\"Kernel execution failed\");\n\n // Perform a max-scan operation on the array of view angles\n thrust::inclusive_scan(d_angles.begin(), d_angles.end(), d_scannedAngles.begin(), thrust::maximum());\n getLastCudaError(\"Kernel execution failed\");\n\n // Compute visibility results based on the array of view angles\n // and its scanned version\n computeVisibilities_kernel<<>>(thrust::raw_pointer_cast(&d_angles[0]),\n thrust::raw_pointer_cast(&d_scannedAngles[0]),\n ray.length,\n thrust::raw_pointer_cast(&d_visibilities[0]));\n getLastCudaError(\"Kernel execution failed\");\n }\n\n cudaDeviceSynchronize();\n sdkStopTimer(&timer);\n getLastCudaError(\"Kernel execution failed\");\n\n // Copy visibility results back to the host\n thrust::copy(d_visibilities.begin(), d_visibilities.end(), h_visibilities.begin());\n\n // Compare device visibility results against reference results\n bool res = compareData(thrust::raw_pointer_cast(&h_visibilitiesRef[0]),\n thrust::raw_pointer_cast(&h_visibilities[0]),\n ray.length,\n 0.0f,\n 0.0f);\n printf(\"Average time: %f ms\\n\\n\", sdkGetTimerValue(&timer) / numIterations);\n sdkResetTimer(&timer);\n\n // Cleanup memory\n checkCudaErrors(cudaFreeArray(heightFieldArray));\n return res;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Compute view angles for each point along the ray\n//! @param ray ray\n//! @param angles view angles\n////////////////////////////////////////////////////////////////////////////////\n__global__ void computeAngles_kernel(const Ray ray, float *angles, cudaTextureObject_t HeightFieldTex)\n{\n uint i = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (i < ray.length) {\n float2 location = getLocation(ray, i + 1);\n float height = tex2D(HeightFieldTex, location.x, location.y);\n float angle = getAngle(ray, location, height);\n angles[i] = angle;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Compute visibility for each point along the ray\n//! @param angles view angles\n//! @param scannedAngles max-scanned view angles\n//! @param numAngles number of view angles\n//! @param visibilities boolean array indicating the visibility of each point\n//! along the ray\n////////////////////////////////////////////////////////////////////////////////\n__global__ void\ncomputeVisibilities_kernel(const float *angles, const float *scannedAngles, int numAngles, Bool *visibilities)\n{\n uint i = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (i < numAngles) {\n visibilities[i] = scannedAngles[i] <= angles[i];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Compute reference data set\n//! @param heightField height field\n//! @param ray ray\n//! @param visibilities boolean array indicating the visibility of each point\n//! along the ray\n////////////////////////////////////////////////////////////////////////////////\nvoid lineOfSight_gold(const HeightField heightField, const Ray ray, Bool *visibilities)\n{\n float angleMax = asinf(-1.0f);\n\n for (int i = 0; i < ray.length; ++i) {\n float2 location = getLocation(ray, i + 1);\n float height = *(heightField.height + heightField.width * (int)floorf(location.y) + (int)floorf(location.x));\n float angle = getAngle(ray, location, height);\n\n if (angle > angleMax) {\n angleMax = angle;\n visibilities[i] = True;\n }\n else {\n visibilities[i] = False;\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Compute the 2D coordinates of the point located at i steps from the origin\n//! of the ray\n//! @param ray ray\n//! @param i integer offset along the ray\n////////////////////////////////////////////////////////////////////////////////\n__device__ __host__ float2 getLocation(const Ray ray, int i)\n{\n float step = i * ray.oneOverLength;\n return make_float2(ray.origin.x, ray.origin.y) + ray.dir * step;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Compute the angle of view between a 3D point and the origin of the ray\n//! @param ray ray\n//! @param location 2D coordinates of the input point\n//! @param height height of the input point\n////////////////////////////////////////////////////////////////////////////////\n__device__ __host__ float getAngle(const Ray ray, float2 location, float height)\n{\n float2 dir = location - make_float2(ray.origin.x, ray.origin.y);\n return atanf((height - ray.origin.z) / length(dir));\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// This sample is an implementation of a simple line-of-sight algorithm:\n// Given a height map and a ray originating at some observation point,\n// it computes all the points along the ray that are visible from the\n// observation point.\n// It is based on the description made in \"Guy E. Blelloch. Vector models\n// for data-parallel computing. MIT Press, 1990\" and uses open source CUDA\n// Thrust Library\n\n#include \"hip/hip_runtime.h\"\n#ifdef _WIN32\n#define NOMINMAX\n#endif\n\n// includes, system\n#include \n#include \n#include \n#include \n#include \n\n// includes, project\n#include \n#include \n#include \n\n// includes, library\n#include \n#include \n#include \n#include \n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, types\n\n// Boolean\ntypedef unsigned char Bool;\nenum { False = 0, True = 1 };\n\n// 2D height field\nstruct HeightField\n{\n int width;\n float *height;\n};\n\n// Ray\nstruct Ray\n{\n float3 origin;\n float2 dir;\n int length;\n float oneOverLength;\n};\n\n////////////////////////////////////////////////////////////////////////////////\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\nint runTest(int argc, char **argv);\n__global__ void computeAngles_kernel(const Ray, float *, hipTextureObject_t);\n__global__ void computeVisibilities_kernel(const float *, const float *, int, Bool *);\nvoid lineOfSight_gold(const HeightField, const Ray, Bool *);\n__device__ __host__ float2 getLocation(const Ray, int);\n__device__ __host__ float getAngle(const Ray, float2, float);\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n int res = runTest(argc, argv);\n\n if (res != 1) {\n printf(\"Test failed!\\n\");\n exit(EXIT_FAILURE);\n }\n\n printf(\"Test passed\\n\");\n exit(EXIT_SUCCESS);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Run a line-of-sight test for CUDA\n////////////////////////////////////////////////////////////////////////////////\nint runTest(int argc, char **argv)\n{\n ////////////////////////////////////////////////////////////////////////////\n // Device initialization\n\n printf(\"[%s] - Starting...\\n\", argv[0]);\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n findCudaDevice(argc, (const char **)argv);\n\n ////////////////////////////////////////////////////////////////////////////\n // Timer\n\n // Create\n StopWatchInterface *timer;\n sdkCreateTimer(&timer);\n\n // Number of iterations to get accurate timing\n uint numIterations = 100;\n\n ////////////////////////////////////////////////////////////////////////////\n // Height field\n\n HeightField heightField;\n\n // Allocate in host memory\n int2 dim = make_int2(10000, 100);\n heightField.width = dim.x;\n thrust::host_vector height(dim.x * dim.y);\n heightField.height = (float *)&height[0];\n\n //\n // Fill in with an arbitrary sine surface\n for (int x = 0; x < dim.x; ++x)\n for (int y = 0; y < dim.y; ++y) {\n float amp = 0.1f * (x + y);\n float period = 2.0f + amp;\n *(heightField.height + dim.x * y + x) =\n amp * (sinf(sqrtf((float)(x * x + y * y)) * 2.0f * 3.1416f / period) + 1.0f);\n }\n\n // Allocate CUDA array in device memory\n hipChannelFormatDesc channelDesc = hipCreateChannelDesc(32, 0, 0, 0, hipChannelFormatKindFloat);\n hipArray *heightFieldArray;\n checkCudaErrors(hipMallocArray(&heightFieldArray, &channelDesc, dim.x, dim.y));\n\n // Initialize device memory\n checkCudaErrors(hipMemcpy2DToArray(heightFieldArray,\n 0,\n 0,\n heightField.height,\n dim.x * sizeof(float),\n dim.x * sizeof(float),\n dim.y,\n hipMemcpyHostToDevice));\n\n hipTextureObject_t heightFieldTex;\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = heightFieldArray;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n texDescr.normalizedCoords = false;\n texDescr.filterMode = hipFilterModePoint;\n texDescr.addressMode[0] = hipAddressModeClamp;\n texDescr.addressMode[1] = hipAddressModeClamp;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&heightFieldTex, &texRes, &texDescr, NULL));\n\n //////////////////////////////////////////////////////////////////////////////\n // Ray (starts at origin and traverses the height field diagonally)\n\n Ray ray;\n ray.origin = make_float3(0, 0, 2.0f);\n int2 dir = make_int2(dim.x - 1, dim.y - 1);\n ray.dir = make_float2((float)dir.x, (float)dir.y);\n ray.length = max(abs(dir.x), abs(dir.y));\n ray.oneOverLength = 1.0f / ray.length;\n\n //////////////////////////////////////////////////////////////////////////////\n // View angles\n\n // Allocate view angles for each point along the ray\n thrust::device_vector d_angles(ray.length);\n\n // Allocate result of max-scan operation on the array of view angles\n thrust::device_vector d_scannedAngles(ray.length);\n\n //////////////////////////////////////////////////////////////////////////////\n // Visibility results\n\n // Allocate visibility results for each point along the ray\n thrust::device_vector d_visibilities(ray.length);\n thrust::host_vector h_visibilities(ray.length);\n thrust::host_vector h_visibilitiesRef(ray.length);\n\n //////////////////////////////////////////////////////////////////////////////\n // Reference solution\n lineOfSight_gold(heightField, ray, (Bool *)&h_visibilitiesRef[0]);\n\n //////////////////////////////////////////////////////////////////////////////\n // Device solution\n\n // Execution configuration\n dim3 block(256);\n dim3 grid((uint)ceil(ray.length / (double)block.x));\n\n // Compute device solution\n printf(\"Line of sight\\n\");\n sdkStartTimer(&timer);\n\n for (uint i = 0; i < numIterations; ++i) {\n // Compute view angle for each point along the ray\n computeAngles_kernel<<>>(ray, thrust::raw_pointer_cast(&d_angles[0]), heightFieldTex);\n getLastCudaError(\"Kernel execution failed\");\n\n // Perform a max-scan operation on the array of view angles\n thrust::inclusive_scan(d_angles.begin(), d_angles.end(), d_scannedAngles.begin(), thrust::maximum());\n getLastCudaError(\"Kernel execution failed\");\n\n // Compute visibility results based on the array of view angles\n // and its scanned version\n computeVisibilities_kernel<<>>(thrust::raw_pointer_cast(&d_angles[0]),\n thrust::raw_pointer_cast(&d_scannedAngles[0]),\n ray.length,\n thrust::raw_pointer_cast(&d_visibilities[0]));\n getLastCudaError(\"Kernel execution failed\");\n }\n\n hipDeviceSynchronize();\n sdkStopTimer(&timer);\n getLastCudaError(\"Kernel execution failed\");\n\n // Copy visibility results back to the host\n thrust::copy(d_visibilities.begin(), d_visibilities.end(), h_visibilities.begin());\n\n // Compare device visibility results against reference results\n bool res = compareData(thrust::raw_pointer_cast(&h_visibilitiesRef[0]),\n thrust::raw_pointer_cast(&h_visibilities[0]),\n ray.length,\n 0.0f,\n 0.0f);\n printf(\"Average time: %f ms\\n\\n\", sdkGetTimerValue(&timer) / numIterations);\n sdkResetTimer(&timer);\n\n // Cleanup memory\n checkCudaErrors(hipFreeArray(heightFieldArray));\n return res;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Compute view angles for each point along the ray\n//! @param ray ray\n//! @param angles view angles\n////////////////////////////////////////////////////////////////////////////////\n__global__ void computeAngles_kernel(const Ray ray, float *angles, hipTextureObject_t HeightFieldTex)\n{\n uint i = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (i < ray.length) {\n float2 location = getLocation(ray, i + 1);\n float height = tex2D(HeightFieldTex, location.x, location.y);\n float angle = getAngle(ray, location, height);\n angles[i] = angle;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Compute visibility for each point along the ray\n//! @param angles view angles\n//! @param scannedAngles max-scanned view angles\n//! @param numAngles number of view angles\n//! @param visibilities boolean array indicating the visibility of each point\n//! along the ray\n////////////////////////////////////////////////////////////////////////////////\n__global__ void\ncomputeVisibilities_kernel(const float *angles, const float *scannedAngles, int numAngles, Bool *visibilities)\n{\n uint i = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (i < numAngles) {\n visibilities[i] = scannedAngles[i] <= angles[i];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Compute reference data set\n//! @param heightField height field\n//! @param ray ray\n//! @param visibilities boolean array indicating the visibility of each point\n//! along the ray\n////////////////////////////////////////////////////////////////////////////////\nvoid lineOfSight_gold(const HeightField heightField, const Ray ray, Bool *visibilities)\n{\n float angleMax = asinf(-1.0f);\n\n for (int i = 0; i < ray.length; ++i) {\n float2 location = getLocation(ray, i + 1);\n float height = *(heightField.height + heightField.width * (int)floorf(location.y) + (int)floorf(location.x));\n float angle = getAngle(ray, location, height);\n\n if (angle > angleMax) {\n angleMax = angle;\n visibilities[i] = True;\n }\n else {\n visibilities[i] = False;\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Compute the 2D coordinates of the point located at i steps from the origin\n//! of the ray\n//! @param ray ray\n//! @param i integer offset along the ray\n////////////////////////////////////////////////////////////////////////////////\n__device__ __host__ float2 getLocation(const Ray ray, int i)\n{\n float step = i * ray.oneOverLength;\n return make_float2(ray.origin.x, ray.origin.y) + ray.dir * step;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Compute the angle of view between a 3D point and the origin of the ray\n//! @param ray ray\n//! @param location 2D coordinates of the input point\n//! @param height height of the input point\n////////////////////////////////////////////////////////////////////////////////\n__device__ __host__ float getAngle(const Ray ray, float2 location, float height)\n{\n float2 dir = location - make_float2(ray.origin.x, ray.origin.y);\n return atanf((height - ray.origin.z) / length(dir));\n}\n", "source": "/kaggle/working/cuda-samples/Samples/4_CUDA_Libraries/lineOfSight/lineOfSight.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n\n// convert floating point rgba color to 32-bit integer\n__device__ unsigned int rgbaFloatToInt(float4 rgba)\n{\n rgba.x = __saturatef(rgba.x); // clamp to [0.0, 1.0]\n rgba.y = __saturatef(rgba.y);\n rgba.z = __saturatef(rgba.z);\n rgba.w = __saturatef(rgba.w);\n return ((unsigned int)(rgba.w * 255.0f) << 24) | ((unsigned int)(rgba.z * 255.0f) << 16)\n | ((unsigned int)(rgba.y * 255.0f) << 8) | ((unsigned int)(rgba.x * 255.0f));\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Rotate an image using texture lookups\n//! @param outputData output data in global memory\n////////////////////////////////////////////////////////////////////////////////\nstatic __global__ void\ntransformKernel(unsigned int *outputData, int width, int height, float theta, cudaTextureObject_t tex)\n{\n // calculate normalized texture coordinates\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n float u = (float)x - (float)width / 2;\n float v = (float)y - (float)height / 2;\n float tu = u * cosf(theta) - v * sinf(theta);\n float tv = v * cosf(theta) + u * sinf(theta);\n\n tu /= (float)width;\n tv /= (float)height;\n\n // read from texture and write to global memory\n float4 pix = tex2D(tex, tu + 0.5f, tv + 0.5f);\n unsigned int pixelInt = rgbaFloatToInt(pix);\n outputData[y * width + x] = pixelInt;\n}\n\nstatic __global__ void rgbToGrayscaleKernel(unsigned int *rgbaImage, size_t imageWidth, size_t imageHeight)\n{\n size_t gidX = blockDim.x * blockIdx.x + threadIdx.x;\n\n uchar4 *pixArray = (uchar4 *)rgbaImage;\n\n for (int pixId = gidX; pixId < imageWidth * imageHeight; pixId += gridDim.x * blockDim.x) {\n uchar4 dataA = pixArray[pixId];\n unsigned char grayscale = (unsigned char)(dataA.x * 0.3 + dataA.y * 0.59 + dataA.z * 0.11);\n uchar4 dataB = make_uchar4(grayscale, grayscale, grayscale, 0);\n pixArray[pixId] = dataB;\n }\n}\n\nvoid launchGrayScaleKernel(unsigned int *d_rgbaImage,\n std::string image_filename,\n size_t imageWidth,\n size_t imageHeight,\n cudaStream_t stream)\n{\n int numThreadsPerBlock = 1024;\n int numOfBlocks = (imageWidth * imageHeight) / numThreadsPerBlock;\n\n rgbToGrayscaleKernel<<>>(d_rgbaImage, imageWidth, imageHeight);\n\n unsigned int *outputData;\n checkCudaErrors(cudaMallocHost((void **)&outputData, sizeof(unsigned int) * imageWidth * imageHeight));\n checkCudaErrors(cudaMemcpyAsync(\n outputData, d_rgbaImage, sizeof(unsigned int) * imageWidth * imageHeight, cudaMemcpyDeviceToHost, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n char outputFilename[1024];\n strcpy(outputFilename, image_filename.c_str());\n strcpy(outputFilename + image_filename.length() - 4, \"_out.ppm\");\n sdkSavePPM4ub(outputFilename, (unsigned char *)outputData, imageWidth, imageHeight);\n printf(\"Wrote '%s'\\n\", outputFilename);\n\n checkCudaErrors(cudaFreeHost(outputData));\n}\n\nvoid rotateKernel(cudaTextureObject_t &texObj,\n const float angle,\n unsigned int *d_outputData,\n const int imageWidth,\n const int imageHeight,\n cudaStream_t stream)\n{\n dim3 dimBlock(8, 8, 1);\n dim3 dimGrid(imageWidth / dimBlock.x, imageHeight / dimBlock.y, 1);\n\n transformKernel<<>>(d_outputData, imageWidth, imageHeight, angle, texObj);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n\n// convert floating point rgba color to 32-bit integer\n__device__ unsigned int rgbaFloatToInt(float4 rgba)\n{\n rgba.x = __saturatef(rgba.x); // clamp to [0.0, 1.0]\n rgba.y = __saturatef(rgba.y);\n rgba.z = __saturatef(rgba.z);\n rgba.w = __saturatef(rgba.w);\n return ((unsigned int)(rgba.w * 255.0f) << 24) | ((unsigned int)(rgba.z * 255.0f) << 16)\n | ((unsigned int)(rgba.y * 255.0f) << 8) | ((unsigned int)(rgba.x * 255.0f));\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Rotate an image using texture lookups\n//! @param outputData output data in global memory\n////////////////////////////////////////////////////////////////////////////////\nstatic __global__ void\ntransformKernel(unsigned int *outputData, int width, int height, float theta, hipTextureObject_t tex)\n{\n // calculate normalized texture coordinates\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n\n float u = (float)x - (float)width / 2;\n float v = (float)y - (float)height / 2;\n float tu = u * cosf(theta) - v * sinf(theta);\n float tv = v * cosf(theta) + u * sinf(theta);\n\n tu /= (float)width;\n tv /= (float)height;\n\n // read from texture and write to global memory\n float4 pix = tex2D(tex, tu + 0.5f, tv + 0.5f);\n unsigned int pixelInt = rgbaFloatToInt(pix);\n outputData[y * width + x] = pixelInt;\n}\n\nstatic __global__ void rgbToGrayscaleKernel(unsigned int *rgbaImage, size_t imageWidth, size_t imageHeight)\n{\n size_t gidX = blockDim.x * blockIdx.x + threadIdx.x;\n\n uchar4 *pixArray = (uchar4 *)rgbaImage;\n\n for (int pixId = gidX; pixId < imageWidth * imageHeight; pixId += gridDim.x * blockDim.x) {\n uchar4 dataA = pixArray[pixId];\n unsigned char grayscale = (unsigned char)(dataA.x * 0.3 + dataA.y * 0.59 + dataA.z * 0.11);\n uchar4 dataB = make_uchar4(grayscale, grayscale, grayscale, 0);\n pixArray[pixId] = dataB;\n }\n}\n\nvoid launchGrayScaleKernel(unsigned int *d_rgbaImage,\n std::string image_filename,\n size_t imageWidth,\n size_t imageHeight,\n hipStream_t stream)\n{\n int numThreadsPerBlock = 1024;\n int numOfBlocks = (imageWidth * imageHeight) / numThreadsPerBlock;\n\n rgbToGrayscaleKernel<<>>(d_rgbaImage, imageWidth, imageHeight);\n\n unsigned int *outputData;\n checkCudaErrors(hipHostMalloc((void **)&outputData, sizeof(unsigned int) * imageWidth * imageHeight));\n checkCudaErrors(hipMemcpyAsync(\n outputData, d_rgbaImage, sizeof(unsigned int) * imageWidth * imageHeight, hipMemcpyDeviceToHost, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n\n char outputFilename[1024];\n strcpy(outputFilename, image_filename.c_str());\n strcpy(outputFilename + image_filename.length() - 4, \"_out.ppm\");\n sdkSavePPM4ub(outputFilename, (unsigned char *)outputData, imageWidth, imageHeight);\n printf(\"Wrote '%s'\\n\", outputFilename);\n\n checkCudaErrors(hipHostFree(outputData));\n}\n\nvoid rotateKernel(hipTextureObject_t &texObj,\n const float angle,\n unsigned int *d_outputData,\n const int imageWidth,\n const int imageHeight,\n hipStream_t stream)\n{\n dim3 dimBlock(8, 8, 1);\n dim3 dimGrid(imageWidth / dimBlock.x, imageHeight / dimBlock.y, 1);\n\n transformKernel<<>>(d_outputData, imageWidth, imageHeight, angle, texObj);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/4_CUDA_Libraries/cudaNvSci/imageKernels.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample implements a conjugate gradient solver on multiple GPU using\n * Unified Memory optimized prefetching and usage hints.\n *\n */\n\n// includes, system\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n// Utilities and system includes\n#include \n#include \n#include // helper function CUDA error checking and initialization\n#include // helper for shared functions common to CUDA Samples\n\nnamespace cg = cooperative_groups;\n\nconst char *sSDKname = \"conjugateGradientMultiDeviceCG\";\n\n#define ENABLE_CPU_DEBUG_CODE 0\n#define THREADS_PER_BLOCK 512\n\n__device__ double grid_dot_result = 0.0;\n\n/* genTridiag: generate a random tridiagonal symmetric matrix */\nvoid genTridiag(int *I, int *J, float *val, int N, int nz)\n{\n I[0] = 0, J[0] = 0, J[1] = 1;\n val[0] = (float)rand() / RAND_MAX + 10.0f;\n val[1] = (float)rand() / RAND_MAX;\n int start;\n\n for (int i = 1; i < N; i++) {\n if (i > 1) {\n I[i] = I[i - 1] + 3;\n }\n else {\n I[1] = 2;\n }\n\n start = (i - 1) * 3 + 2;\n J[start] = i - 1;\n J[start + 1] = i;\n\n if (i < N - 1) {\n J[start + 2] = i + 1;\n }\n\n val[start] = val[start - 1];\n val[start + 1] = (float)rand() / RAND_MAX + 10.0f;\n\n if (i < N - 1) {\n val[start + 2] = (float)rand() / RAND_MAX;\n }\n }\n\n I[N] = nz;\n}\n\n// I - contains location of the given non-zero element in the row of the matrix\n// J - contains location of the given non-zero element in the column of the\n// matrix val - contains values of the given non-zero elements of the matrix\n// inputVecX - input vector to be multiplied\n// outputVecY - resultant vector\nvoid cpuSpMV(int *I, int *J, float *val, int nnz, int num_rows, float alpha, float *inputVecX, float *outputVecY)\n{\n for (int i = 0; i < num_rows; i++) {\n int num_elems_this_row = I[i + 1] - I[i];\n\n float output = 0.0;\n for (int j = 0; j < num_elems_this_row; j++) {\n output += alpha * val[I[i] + j] * inputVecX[J[I[i] + j]];\n }\n outputVecY[i] = output;\n }\n\n return;\n}\n\nfloat dotProduct(float *vecA, float *vecB, int size)\n{\n float result = 0.0;\n\n for (int i = 0; i < size; i++) {\n result = result + (vecA[i] * vecB[i]);\n }\n\n return result;\n}\n\nvoid scaleVector(float *vec, float alpha, int size)\n{\n for (int i = 0; i < size; i++) {\n vec[i] = alpha * vec[i];\n }\n}\n\nvoid saxpy(float *x, float *y, float a, int size)\n{\n for (int i = 0; i < size; i++) {\n y[i] = a * x[i] + y[i];\n }\n}\n\nvoid cpuConjugateGrad(int *I, int *J, float *val, float *x, float *Ax, float *p, float *r, int nnz, int N, float tol)\n{\n int max_iter = 10000;\n\n float alpha = 1.0;\n float alpham1 = -1.0;\n float r0 = 0.0, b, a, na;\n\n cpuSpMV(I, J, val, nnz, N, alpha, x, Ax);\n saxpy(Ax, r, alpham1, N);\n\n float r1 = dotProduct(r, r, N);\n\n int k = 1;\n\n while (r1 > tol * tol && k <= max_iter) {\n if (k > 1) {\n b = r1 / r0;\n scaleVector(p, b, N);\n\n saxpy(r, p, alpha, N);\n }\n else {\n for (int i = 0; i < N; i++)\n p[i] = r[i];\n }\n\n cpuSpMV(I, J, val, nnz, N, alpha, p, Ax);\n\n float dot = dotProduct(p, Ax, N);\n a = r1 / dot;\n\n saxpy(p, x, a, N);\n na = -a;\n saxpy(Ax, r, na, N);\n\n r0 = r1;\n r1 = dotProduct(r, r, N);\n\n printf(\"\\nCPU code iteration = %3d, residual = %e\\n\", k, sqrt(r1));\n k++;\n }\n}\n\n// Data filled on CPU needed for MultiGPU operations.\nstruct MultiDeviceData\n{\n unsigned char *hostMemoryArrivedList;\n unsigned int numDevices;\n unsigned int deviceRank;\n};\n\n// Class used for coordination of multiple devices.\nclass PeerGroup\n{\n const MultiDeviceData &data;\n const cg::grid_group &grid;\n\n __device__ unsigned char load_arrived(unsigned char *arrived) const\n {\n#if __CUDA_ARCH__ < 700\n return *(volatile unsigned char *)arrived;\n#else\n unsigned int result;\n asm volatile(\"ld.acquire.sys.global.u8 %0, [%1];\" : \"=r\"(result) : \"l\"(arrived) : \"memory\");\n return result;\n#endif\n }\n\n __device__ void store_arrived(unsigned char *arrived, unsigned char val) const\n {\n#if __CUDA_ARCH__ < 700\n *(volatile unsigned char *)arrived = val;\n#else\n unsigned int reg_val = val;\n asm volatile(\"st.release.sys.global.u8 [%1], %0;\" ::\"r\"(reg_val) \"l\"(arrived) : \"memory\");\n\n // Avoids compiler warnings from unused variable val.\n (void)(reg_val = reg_val);\n#endif\n }\n\npublic:\n __device__ PeerGroup(const MultiDeviceData &data, const cg::grid_group &grid)\n : data(data)\n , grid(grid) {};\n\n __device__ unsigned int size() const { return data.numDevices * grid.size(); }\n\n __device__ unsigned int thread_rank() const { return data.deviceRank * grid.size() + grid.thread_rank(); }\n\n __device__ void sync() const\n {\n grid.sync();\n\n // One thread from each grid participates in the sync.\n if (grid.thread_rank() == 0) {\n if (data.deviceRank == 0) {\n // Leader grid waits for others to join and then releases them.\n // Other GPUs can arrive in any order, so the leader have to wait for\n // all others.\n for (int i = 0; i < data.numDevices - 1; i++) {\n while (load_arrived(&data.hostMemoryArrivedList[i]) == 0)\n ;\n }\n for (int i = 0; i < data.numDevices - 1; i++) {\n store_arrived(&data.hostMemoryArrivedList[i], 0);\n }\n __threadfence_system();\n }\n else {\n // Other grids note their arrival and wait to be released.\n store_arrived(&data.hostMemoryArrivedList[data.deviceRank - 1], 1);\n while (load_arrived(&data.hostMemoryArrivedList[data.deviceRank - 1]) == 1)\n ;\n }\n }\n\n grid.sync();\n }\n};\n\n__device__ void gpuSpMV(int *I,\n int *J,\n float *val,\n int nnz,\n int num_rows,\n float alpha,\n float *inputVecX,\n float *outputVecY,\n const PeerGroup &peer_group)\n{\n for (int i = peer_group.thread_rank(); i < num_rows; i += peer_group.size()) {\n int row_elem = I[i];\n int next_row_elem = I[i + 1];\n int num_elems_this_row = next_row_elem - row_elem;\n\n float output = 0.0;\n for (int j = 0; j < num_elems_this_row; j++) {\n output += alpha * val[row_elem + j] * inputVecX[J[row_elem + j]];\n }\n\n outputVecY[i] = output;\n }\n}\n\n__device__ void gpuSaxpy(float *x, float *y, float a, int size, const PeerGroup &peer_group)\n{\n for (int i = peer_group.thread_rank(); i < size; i += peer_group.size()) {\n y[i] = a * x[i] + y[i];\n }\n}\n\n__device__ void\ngpuDotProduct(float *vecA, float *vecB, int size, const cg::thread_block &cta, const PeerGroup &peer_group)\n{\n extern __shared__ double tmp[];\n\n double temp_sum = 0.0;\n\n for (int i = peer_group.thread_rank(); i < size; i += peer_group.size()) {\n temp_sum += (double)(vecA[i] * vecB[i]);\n }\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n temp_sum = cg::reduce(tile32, temp_sum, cg::plus());\n\n if (tile32.thread_rank() == 0) {\n tmp[tile32.meta_group_rank()] = temp_sum;\n }\n\n cg::sync(cta);\n\n if (tile32.meta_group_rank() == 0) {\n temp_sum = tile32.thread_rank() < tile32.meta_group_size() ? tmp[tile32.thread_rank()] : 0.0;\n temp_sum = cg::reduce(tile32, temp_sum, cg::plus());\n\n if (tile32.thread_rank() == 0) {\n atomicAdd(&grid_dot_result, temp_sum);\n }\n }\n}\n\n__device__ void gpuCopyVector(float *srcA, float *destB, int size, const PeerGroup &peer_group)\n{\n for (int i = peer_group.thread_rank(); i < size; i += peer_group.size()) {\n destB[i] = srcA[i];\n }\n}\n\n__device__ void gpuScaleVectorAndSaxpy(float *x, float *y, float a, float scale, int size, const PeerGroup &peer_group)\n{\n for (int i = peer_group.thread_rank(); i < size; i += peer_group.size()) {\n y[i] = a * x[i] + scale * y[i];\n }\n}\n\nextern \"C\" __global__ void multiGpuConjugateGradient(int *I,\n int *J,\n float *val,\n float *x,\n float *Ax,\n float *p,\n float *r,\n double *dot_result,\n int nnz,\n int N,\n float tol,\n MultiDeviceData multi_device_data)\n{\n cg::thread_block cta = cg::this_thread_block();\n cg::grid_group grid = cg::this_grid();\n PeerGroup peer_group(multi_device_data, grid);\n\n const int max_iter = 10000;\n\n float alpha = 1.0;\n float alpham1 = -1.0;\n float r0 = 0.0, r1, b, a, na;\n\n for (int i = peer_group.thread_rank(); i < N; i += peer_group.size()) {\n r[i] = 1.0;\n x[i] = 0.0;\n }\n\n cg::sync(grid);\n\n gpuSpMV(I, J, val, nnz, N, alpha, x, Ax, peer_group);\n\n cg::sync(grid);\n\n gpuSaxpy(Ax, r, alpham1, N, peer_group);\n\n cg::sync(grid);\n\n gpuDotProduct(r, r, N, cta, peer_group);\n\n cg::sync(grid);\n\n if (grid.thread_rank() == 0) {\n atomicAdd_system(dot_result, grid_dot_result);\n grid_dot_result = 0.0;\n }\n peer_group.sync();\n\n r1 = *dot_result;\n\n int k = 1;\n while (r1 > tol * tol && k <= max_iter) {\n if (k > 1) {\n b = r1 / r0;\n gpuScaleVectorAndSaxpy(r, p, alpha, b, N, peer_group);\n }\n else {\n gpuCopyVector(r, p, N, peer_group);\n }\n\n peer_group.sync();\n\n gpuSpMV(I, J, val, nnz, N, alpha, p, Ax, peer_group);\n\n if (peer_group.thread_rank() == 0) {\n *dot_result = 0.0;\n }\n peer_group.sync();\n\n gpuDotProduct(p, Ax, N, cta, peer_group);\n\n cg::sync(grid);\n\n if (grid.thread_rank() == 0) {\n atomicAdd_system(dot_result, grid_dot_result);\n grid_dot_result = 0.0;\n }\n peer_group.sync();\n\n a = r1 / *dot_result;\n\n gpuSaxpy(p, x, a, N, peer_group);\n\n na = -a;\n\n gpuSaxpy(Ax, r, na, N, peer_group);\n\n r0 = r1;\n\n peer_group.sync();\n\n if (peer_group.thread_rank() == 0) {\n *dot_result = 0.0;\n }\n\n peer_group.sync();\n\n gpuDotProduct(r, r, N, cta, peer_group);\n\n cg::sync(grid);\n\n if (grid.thread_rank() == 0) {\n atomicAdd_system(dot_result, grid_dot_result);\n grid_dot_result = 0.0;\n }\n peer_group.sync();\n\n r1 = *dot_result;\n k++;\n }\n}\n\n// Map of device version to device number\nstd::multimap, int> getIdenticalGPUs()\n{\n int numGpus = 0;\n checkCudaErrors(cudaGetDeviceCount(&numGpus));\n\n std::multimap, int> identicalGpus;\n\n for (int i = 0; i < numGpus; i++) {\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, i));\n\n // Filter unsupported devices\n if (deviceProp.cooperativeLaunch && deviceProp.concurrentManagedAccess) {\n identicalGpus.emplace(std::make_pair(deviceProp.major, deviceProp.minor), i);\n }\n printf(\"GPU Device %d: \\\"%s\\\" with compute capability %d.%d\\n\",\n i,\n deviceProp.name,\n deviceProp.major,\n deviceProp.minor);\n }\n\n return identicalGpus;\n}\n\nint main(int argc, char **argv)\n{\n constexpr size_t kNumGpusRequired = 2;\n int N = 0, nz = 0, *I = NULL, *J = NULL;\n float *val = NULL;\n const float tol = 1e-5f;\n float *x;\n float rhs = 1.0;\n float r1;\n float *r, *p, *Ax;\n\n printf(\"Starting [%s]...\\n\", sSDKname);\n auto gpusByArch = getIdenticalGPUs();\n\n auto it = gpusByArch.begin();\n auto end = gpusByArch.end();\n\n auto bestFit = std::make_pair(it, it);\n // use std::distance to find the largest number of GPUs amongst architectures\n auto distance = [](decltype(bestFit) p) { return std::distance(p.first, p.second); };\n\n // Read each unique key/pair element in order\n for (; it != end; it = gpusByArch.upper_bound(it->first)) {\n // first and second are iterators bounded within the architecture group\n auto testFit = gpusByArch.equal_range(it->first);\n // Always use devices with highest architecture version or whichever has the\n // most devices available\n if (distance(bestFit) <= distance(testFit))\n bestFit = testFit;\n }\n\n if (distance(bestFit) < kNumGpusRequired) {\n printf(\"No two or more GPUs with same architecture capable of \"\n \"concurrentManagedAccess found. \"\n \"\\nWaiving the sample\\n\");\n exit(EXIT_WAIVED);\n }\n\n std::set bestFitDeviceIds;\n\n // Check & select peer-to-peer access capable GPU devices as enabling p2p\n // access between participating GPUs gives better performance.\n for (auto itr = bestFit.first; itr != bestFit.second; itr++) {\n int deviceId = itr->second;\n checkCudaErrors(cudaSetDevice(deviceId));\n\n std::for_each(itr, bestFit.second, [&deviceId, &bestFitDeviceIds, &kNumGpusRequired](decltype(*itr) mapPair) {\n if (deviceId != mapPair.second) {\n int access = 0;\n checkCudaErrors(cudaDeviceCanAccessPeer(&access, deviceId, mapPair.second));\n printf(\"Device=%d %s Access Peer Device=%d\\n\", deviceId, access ? \"CAN\" : \"CANNOT\", mapPair.second);\n if (access && bestFitDeviceIds.size() < kNumGpusRequired) {\n bestFitDeviceIds.emplace(deviceId);\n bestFitDeviceIds.emplace(mapPair.second);\n }\n else {\n printf(\"Ignoring device %i (max devices exceeded)\\n\", mapPair.second);\n }\n }\n });\n\n if (bestFitDeviceIds.size() >= kNumGpusRequired) {\n printf(\"Selected p2p capable devices - \");\n for (auto devicesItr = bestFitDeviceIds.begin(); devicesItr != bestFitDeviceIds.end(); devicesItr++) {\n printf(\"deviceId = %d \", *devicesItr);\n }\n printf(\"\\n\");\n break;\n }\n }\n\n // if bestFitDeviceIds.size() == 0 it means the GPUs in system are not p2p\n // capable, hence we add it without p2p capability check.\n if (!bestFitDeviceIds.size()) {\n printf(\"Devices involved are not p2p capable.. selecting %zu of them\\n\", kNumGpusRequired);\n std::for_each(\n bestFit.first, bestFit.second, [&bestFitDeviceIds, &kNumGpusRequired](decltype(*bestFit.first) mapPair) {\n if (bestFitDeviceIds.size() < kNumGpusRequired) {\n bestFitDeviceIds.emplace(mapPair.second);\n }\n else {\n printf(\"Ignoring device %i (max devices exceeded)\\n\", mapPair.second);\n }\n // Insert the sequence into the deviceIds set\n });\n }\n else {\n // perform cudaDeviceEnablePeerAccess in both directions for all\n // participating devices.\n for (auto p1_itr = bestFitDeviceIds.begin(); p1_itr != bestFitDeviceIds.end(); p1_itr++) {\n checkCudaErrors(cudaSetDevice(*p1_itr));\n for (auto p2_itr = bestFitDeviceIds.begin(); p2_itr != bestFitDeviceIds.end(); p2_itr++) {\n if (*p1_itr != *p2_itr) {\n checkCudaErrors(cudaDeviceEnablePeerAccess(*p2_itr, 0));\n checkCudaErrors(cudaSetDevice(*p1_itr));\n }\n }\n }\n }\n\n /* Generate a random tridiagonal symmetric matrix in CSR format */\n N = 10485760 * 2;\n nz = (N - 2) * 3 + 4;\n\n checkCudaErrors(cudaMallocManaged((void **)&I, sizeof(int) * (N + 1)));\n checkCudaErrors(cudaMallocManaged((void **)&J, sizeof(int) * nz));\n checkCudaErrors(cudaMallocManaged((void **)&val, sizeof(float) * nz));\n\n float *val_cpu = (float *)malloc(sizeof(float) * nz);\n\n genTridiag(I, J, val_cpu, N, nz);\n\n memcpy(val, val_cpu, sizeof(float) * nz);\n cudaMemLocation deviceLoc;\n deviceLoc.type = cudaMemLocationTypeDevice;\n deviceLoc.id = 0; // Device location with initial device 0\n checkCudaErrors(cudaMemAdvise(I, sizeof(int) * (N + 1), cudaMemAdviseSetReadMostly, deviceLoc));\n checkCudaErrors(cudaMemAdvise(J, sizeof(int) * nz, cudaMemAdviseSetReadMostly, deviceLoc));\n checkCudaErrors(cudaMemAdvise(val, sizeof(float) * nz, cudaMemAdviseSetReadMostly, deviceLoc));\n\n checkCudaErrors(cudaMallocManaged((void **)&x, sizeof(float) * N));\n\n double *dot_result;\n checkCudaErrors(cudaMallocManaged((void **)&dot_result, sizeof(double)));\n\n checkCudaErrors(cudaMemset(dot_result, 0, sizeof(double)));\n\n // temp memory for ConjugateGradient\n checkCudaErrors(cudaMallocManaged((void **)&r, N * sizeof(float)));\n checkCudaErrors(cudaMallocManaged((void **)&p, N * sizeof(float)));\n checkCudaErrors(cudaMallocManaged((void **)&Ax, N * sizeof(float)));\n\n std::cout << \"\\nRunning on GPUs = \" << kNumGpusRequired << std::endl;\n cudaStream_t nStreams[kNumGpusRequired];\n\n int sMemSize = sizeof(double) * ((THREADS_PER_BLOCK / 32) + 1);\n int numBlocksPerSm = INT_MAX;\n int numThreads = THREADS_PER_BLOCK;\n int numSms = INT_MAX;\n auto deviceId = bestFitDeviceIds.begin();\n\n // set numSms & numBlocksPerSm to be lowest of 2 devices\n while (deviceId != bestFitDeviceIds.end()) {\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaSetDevice(*deviceId));\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, *deviceId));\n\n int numBlocksPerSm_current = 0;\n checkCudaErrors(cudaOccupancyMaxActiveBlocksPerMultiprocessor(\n &numBlocksPerSm_current, multiGpuConjugateGradient, numThreads, sMemSize));\n\n if (numBlocksPerSm > numBlocksPerSm_current) {\n numBlocksPerSm = numBlocksPerSm_current;\n }\n if (numSms > deviceProp.multiProcessorCount) {\n numSms = deviceProp.multiProcessorCount;\n }\n deviceId++;\n }\n\n if (!numBlocksPerSm) {\n printf(\"Max active blocks per SM is returned as 0.\\n Hence, Waiving the \"\n \"sample\\n\");\n exit(EXIT_WAIVED);\n }\n\n int device_count = 0;\n int totalThreadsPerGPU = numSms * numBlocksPerSm * THREADS_PER_BLOCK;\n deviceId = bestFitDeviceIds.begin();\n while (deviceId != bestFitDeviceIds.end()) {\n checkCudaErrors(cudaSetDevice(*deviceId));\n checkCudaErrors(cudaStreamCreate(&nStreams[device_count]));\n\n int perGPUIter = N / (totalThreadsPerGPU * kNumGpusRequired);\n int offset_Ax = device_count * totalThreadsPerGPU;\n int offset_r = device_count * totalThreadsPerGPU;\n int offset_p = device_count * totalThreadsPerGPU;\n int offset_x = device_count * totalThreadsPerGPU;\n\n // Create device location with specific device ID\n cudaMemLocation deviceLoc;\n deviceLoc.type = cudaMemLocationTypeDevice;\n deviceLoc.id = *deviceId;\n checkCudaErrors(cudaMemPrefetchAsync(I, sizeof(int) * N, deviceLoc, 0, nStreams[device_count]));\n checkCudaErrors(cudaMemPrefetchAsync(val, sizeof(float) * nz, deviceLoc, 0, nStreams[device_count]));\n checkCudaErrors(cudaMemPrefetchAsync(J, sizeof(float) * nz, deviceLoc, 0, nStreams[device_count]));\n\n if (offset_Ax <= N) {\n for (int i = 0; i < perGPUIter; i++) {\n cudaMemAdvise(\n Ax + offset_Ax, sizeof(float) * totalThreadsPerGPU, cudaMemAdviseSetPreferredLocation, deviceLoc);\n cudaMemAdvise(\n r + offset_r, sizeof(float) * totalThreadsPerGPU, cudaMemAdviseSetPreferredLocation, deviceLoc);\n cudaMemAdvise(\n x + offset_x, sizeof(float) * totalThreadsPerGPU, cudaMemAdviseSetPreferredLocation, deviceLoc);\n cudaMemAdvise(\n p + offset_p, sizeof(float) * totalThreadsPerGPU, cudaMemAdviseSetPreferredLocation, deviceLoc);\n\n cudaMemAdvise(\n Ax + offset_Ax, sizeof(float) * totalThreadsPerGPU, cudaMemAdviseSetAccessedBy, deviceLoc);\n cudaMemAdvise(r + offset_r, sizeof(float) * totalThreadsPerGPU, cudaMemAdviseSetAccessedBy, deviceLoc);\n cudaMemAdvise(p + offset_p, sizeof(float) * totalThreadsPerGPU, cudaMemAdviseSetAccessedBy, deviceLoc);\n cudaMemAdvise(x + offset_x, sizeof(float) * totalThreadsPerGPU, cudaMemAdviseSetAccessedBy, deviceLoc);\n\n offset_Ax += totalThreadsPerGPU * kNumGpusRequired;\n offset_r += totalThreadsPerGPU * kNumGpusRequired;\n offset_p += totalThreadsPerGPU * kNumGpusRequired;\n offset_x += totalThreadsPerGPU * kNumGpusRequired;\n\n if (offset_Ax >= N) {\n break;\n }\n }\n }\n\n device_count++;\n deviceId++;\n }\n\n#if ENABLE_CPU_DEBUG_CODE\n float *Ax_cpu = (float *)malloc(sizeof(float) * N);\n float *r_cpu = (float *)malloc(sizeof(float) * N);\n float *p_cpu = (float *)malloc(sizeof(float) * N);\n float *x_cpu = (float *)malloc(sizeof(float) * N);\n\n for (int i = 0; i < N; i++) {\n r_cpu[i] = 1.0;\n Ax_cpu[i] = x_cpu[i] = 0.0;\n }\n#endif\n\n printf(\"Total threads per GPU = %d numBlocksPerSm = %d\\n\",\n numSms * numBlocksPerSm * THREADS_PER_BLOCK,\n numBlocksPerSm);\n dim3 dimGrid(numSms * numBlocksPerSm, 1, 1), dimBlock(THREADS_PER_BLOCK, 1, 1);\n\n // Structure used for cross-grid synchronization.\n MultiDeviceData multi_device_data;\n checkCudaErrors(cudaHostAlloc(&multi_device_data.hostMemoryArrivedList,\n (kNumGpusRequired - 1) * sizeof(*multi_device_data.hostMemoryArrivedList),\n cudaHostAllocPortable));\n memset(multi_device_data.hostMemoryArrivedList,\n 0,\n (kNumGpusRequired - 1) * sizeof(*multi_device_data.hostMemoryArrivedList));\n multi_device_data.numDevices = kNumGpusRequired;\n multi_device_data.deviceRank = 0;\n\n void *kernelArgs[] = {\n (void *)&I,\n (void *)&J,\n (void *)&val,\n (void *)&x,\n (void *)&Ax,\n (void *)&p,\n (void *)&r,\n (void *)&dot_result,\n (void *)&nz,\n (void *)&N,\n (void *)&tol,\n (void *)&multi_device_data,\n };\n\n printf(\"Launching kernel\\n\");\n\n deviceId = bestFitDeviceIds.begin();\n device_count = 0;\n while (deviceId != bestFitDeviceIds.end()) {\n checkCudaErrors(cudaSetDevice(*deviceId));\n checkCudaErrors(cudaLaunchCooperativeKernel(\n (void *)multiGpuConjugateGradient, dimGrid, dimBlock, kernelArgs, sMemSize, nStreams[device_count++]));\n multi_device_data.deviceRank++;\n deviceId++;\n }\n\n // Use cudaMemLocationTypeHost for optimal host memory location\n cudaMemLocation hostLoc;\n hostLoc.type = cudaMemLocationTypeHost;\n checkCudaErrors(cudaMemPrefetchAsync(x, sizeof(float) * N, hostLoc, 0));\n checkCudaErrors(cudaMemPrefetchAsync(dot_result, sizeof(double), hostLoc, 0));\n\n deviceId = bestFitDeviceIds.begin();\n device_count = 0;\n while (deviceId != bestFitDeviceIds.end()) {\n checkCudaErrors(cudaSetDevice(*deviceId));\n checkCudaErrors(cudaStreamSynchronize(nStreams[device_count++]));\n deviceId++;\n }\n\n r1 = (float)*dot_result;\n\n printf(\"GPU Final, residual = %e \\n \", sqrt(r1));\n\n#if ENABLE_CPU_DEBUG_CODE\n cpuConjugateGrad(I, J, val, x_cpu, Ax_cpu, p_cpu, r_cpu, nz, N, tol);\n#endif\n\n float rsum, diff, err = 0.0;\n\n for (int i = 0; i < N; i++) {\n rsum = 0.0;\n\n for (int j = I[i]; j < I[i + 1]; j++) {\n rsum += val_cpu[j] * x[J[j]];\n }\n\n diff = fabs(rsum - rhs);\n\n if (diff > err) {\n err = diff;\n }\n }\n\n checkCudaErrors(cudaFreeHost(multi_device_data.hostMemoryArrivedList));\n checkCudaErrors(cudaFree(I));\n checkCudaErrors(cudaFree(J));\n checkCudaErrors(cudaFree(val));\n checkCudaErrors(cudaFree(x));\n checkCudaErrors(cudaFree(r));\n checkCudaErrors(cudaFree(p));\n checkCudaErrors(cudaFree(Ax));\n checkCudaErrors(cudaFree(dot_result));\n free(val_cpu);\n\n#if ENABLE_CPU_DEBUG_CODE\n free(Ax_cpu);\n free(r_cpu);\n free(p_cpu);\n free(x_cpu);\n#endif\n\n printf(\"Test Summary: Error amount = %f \\n\", err);\n fprintf(stdout, \"&&&& conjugateGradientMultiDeviceCG %s\\n\", (sqrt(r1) < tol) ? \"PASSED\" : \"FAILED\");\n exit((sqrt(r1) < tol) ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample implements a conjugate gradient solver on multiple GPU using\n * Unified Memory optimized prefetching and usage hints.\n *\n */\n\n// includes, system\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n// Utilities and system includes\n#include \n#include \n#include // helper function CUDA error checking and initialization\n#include // helper for shared functions common to CUDA Samples\n\nnamespace cg = cooperative_groups;\n\nconst char *sSDKname = \"conjugateGradientMultiDeviceCG\";\n\n#define ENABLE_CPU_DEBUG_CODE 0\n#define THREADS_PER_BLOCK 512\n\n__device__ double grid_dot_result = 0.0;\n\n/* genTridiag: generate a random tridiagonal symmetric matrix */\nvoid genTridiag(int *I, int *J, float *val, int N, int nz)\n{\n I[0] = 0, J[0] = 0, J[1] = 1;\n val[0] = (float)rand() / RAND_MAX + 10.0f;\n val[1] = (float)rand() / RAND_MAX;\n int start;\n\n for (int i = 1; i < N; i++) {\n if (i > 1) {\n I[i] = I[i - 1] + 3;\n }\n else {\n I[1] = 2;\n }\n\n start = (i - 1) * 3 + 2;\n J[start] = i - 1;\n J[start + 1] = i;\n\n if (i < N - 1) {\n J[start + 2] = i + 1;\n }\n\n val[start] = val[start - 1];\n val[start + 1] = (float)rand() / RAND_MAX + 10.0f;\n\n if (i < N - 1) {\n val[start + 2] = (float)rand() / RAND_MAX;\n }\n }\n\n I[N] = nz;\n}\n\n// I - contains location of the given non-zero element in the row of the matrix\n// J - contains location of the given non-zero element in the column of the\n// matrix val - contains values of the given non-zero elements of the matrix\n// inputVecX - input vector to be multiplied\n// outputVecY - resultant vector\nvoid cpuSpMV(int *I, int *J, float *val, int nnz, int num_rows, float alpha, float *inputVecX, float *outputVecY)\n{\n for (int i = 0; i < num_rows; i++) {\n int num_elems_this_row = I[i + 1] - I[i];\n\n float output = 0.0;\n for (int j = 0; j < num_elems_this_row; j++) {\n output += alpha * val[I[i] + j] * inputVecX[J[I[i] + j]];\n }\n outputVecY[i] = output;\n }\n\n return;\n}\n\nfloat dotProduct(float *vecA, float *vecB, int size)\n{\n float result = 0.0;\n\n for (int i = 0; i < size; i++) {\n result = result + (vecA[i] * vecB[i]);\n }\n\n return result;\n}\n\nvoid scaleVector(float *vec, float alpha, int size)\n{\n for (int i = 0; i < size; i++) {\n vec[i] = alpha * vec[i];\n }\n}\n\nvoid saxpy(float *x, float *y, float a, int size)\n{\n for (int i = 0; i < size; i++) {\n y[i] = a * x[i] + y[i];\n }\n}\n\nvoid cpuConjugateGrad(int *I, int *J, float *val, float *x, float *Ax, float *p, float *r, int nnz, int N, float tol)\n{\n int max_iter = 10000;\n\n float alpha = 1.0;\n float alpham1 = -1.0;\n float r0 = 0.0, b, a, na;\n\n cpuSpMV(I, J, val, nnz, N, alpha, x, Ax);\n saxpy(Ax, r, alpham1, N);\n\n float r1 = dotProduct(r, r, N);\n\n int k = 1;\n\n while (r1 > tol * tol && k <= max_iter) {\n if (k > 1) {\n b = r1 / r0;\n scaleVector(p, b, N);\n\n saxpy(r, p, alpha, N);\n }\n else {\n for (int i = 0; i < N; i++)\n p[i] = r[i];\n }\n\n cpuSpMV(I, J, val, nnz, N, alpha, p, Ax);\n\n float dot = dotProduct(p, Ax, N);\n a = r1 / dot;\n\n saxpy(p, x, a, N);\n na = -a;\n saxpy(Ax, r, na, N);\n\n r0 = r1;\n r1 = dotProduct(r, r, N);\n\n printf(\"\\nCPU code iteration = %3d, residual = %e\\n\", k, sqrt(r1));\n k++;\n }\n}\n\n// Data filled on CPU needed for MultiGPU operations.\nstruct MultiDeviceData\n{\n unsigned char *hostMemoryArrivedList;\n unsigned int numDevices;\n unsigned int deviceRank;\n};\n\n// Class used for coordination of multiple devices.\nclass PeerGroup\n{\n const MultiDeviceData &data;\n const cg::grid_group &grid;\n\n __device__ unsigned char load_arrived(unsigned char *arrived) const\n {\n#if __CUDA_ARCH__ < 700\n return *(volatile unsigned char *)arrived;\n#else\n unsigned int result;\n asm volatile(\"ld.acquire.sys.global.u8 %0, [%1];\" : \"=r\"(result) : \"l\"(arrived) : \"memory\");\n return result;\n#endif\n }\n\n __device__ void store_arrived(unsigned char *arrived, unsigned char val) const\n {\n#if __CUDA_ARCH__ < 700\n *(volatile unsigned char *)arrived = val;\n#else\n unsigned int reg_val = val;\n asm volatile(\"st.release.sys.global.u8 [%1], %0;\" ::\"r\"(reg_val) \"l\"(arrived) : \"memory\");\n\n // Avoids compiler warnings from unused variable val.\n (void)(reg_val = reg_val);\n#endif\n }\n\npublic:\n __device__ PeerGroup(const MultiDeviceData &data, const cg::grid_group &grid)\n : data(data)\n , grid(grid) {};\n\n __device__ unsigned int size() const { return data.numDevices * grid.size(); }\n\n __device__ unsigned int thread_rank() const { return data.deviceRank * grid.size() + grid.thread_rank(); }\n\n __device__ void sync() const\n {\n grid.sync();\n\n // One thread from each grid participates in the sync.\n if (grid.thread_rank() == 0) {\n if (data.deviceRank == 0) {\n // Leader grid waits for others to join and then releases them.\n // Other GPUs can arrive in any order, so the leader have to wait for\n // all others.\n for (int i = 0; i < data.numDevices - 1; i++) {\n while (load_arrived(&data.hostMemoryArrivedList[i]) == 0)\n ;\n }\n for (int i = 0; i < data.numDevices - 1; i++) {\n store_arrived(&data.hostMemoryArrivedList[i], 0);\n }\n __threadfence_system();\n }\n else {\n // Other grids note their arrival and wait to be released.\n store_arrived(&data.hostMemoryArrivedList[data.deviceRank - 1], 1);\n while (load_arrived(&data.hostMemoryArrivedList[data.deviceRank - 1]) == 1)\n ;\n }\n }\n\n grid.sync();\n }\n};\n\n__device__ void gpuSpMV(int *I,\n int *J,\n float *val,\n int nnz,\n int num_rows,\n float alpha,\n float *inputVecX,\n float *outputVecY,\n const PeerGroup &peer_group)\n{\n for (int i = peer_group.thread_rank(); i < num_rows; i += peer_group.size()) {\n int row_elem = I[i];\n int next_row_elem = I[i + 1];\n int num_elems_this_row = next_row_elem - row_elem;\n\n float output = 0.0;\n for (int j = 0; j < num_elems_this_row; j++) {\n output += alpha * val[row_elem + j] * inputVecX[J[row_elem + j]];\n }\n\n outputVecY[i] = output;\n }\n}\n\n__device__ void gpuSaxpy(float *x, float *y, float a, int size, const PeerGroup &peer_group)\n{\n for (int i = peer_group.thread_rank(); i < size; i += peer_group.size()) {\n y[i] = a * x[i] + y[i];\n }\n}\n\n__device__ void\ngpuDotProduct(float *vecA, float *vecB, int size, const cg::thread_block &cta, const PeerGroup &peer_group)\n{\n extern __shared__ double tmp[];\n\n double temp_sum = 0.0;\n\n for (int i = peer_group.thread_rank(); i < size; i += peer_group.size()) {\n temp_sum += (double)(vecA[i] * vecB[i]);\n }\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n temp_sum = cg::reduce(tile32, temp_sum, cg::plus());\n\n if (tile32.thread_rank() == 0) {\n tmp[tile32.meta_group_rank()] = temp_sum;\n }\n\n cg::sync(cta);\n\n if (tile32.meta_group_rank() == 0) {\n temp_sum = tile32.thread_rank() < tile32.meta_group_size() ? tmp[tile32.thread_rank()] : 0.0;\n temp_sum = cg::reduce(tile32, temp_sum, cg::plus());\n\n if (tile32.thread_rank() == 0) {\n atomicAdd(&grid_dot_result, temp_sum);\n }\n }\n}\n\n__device__ void gpuCopyVector(float *srcA, float *destB, int size, const PeerGroup &peer_group)\n{\n for (int i = peer_group.thread_rank(); i < size; i += peer_group.size()) {\n destB[i] = srcA[i];\n }\n}\n\n__device__ void gpuScaleVectorAndSaxpy(float *x, float *y, float a, float scale, int size, const PeerGroup &peer_group)\n{\n for (int i = peer_group.thread_rank(); i < size; i += peer_group.size()) {\n y[i] = a * x[i] + scale * y[i];\n }\n}\n\nextern \"C\" __global__ void multiGpuConjugateGradient(int *I,\n int *J,\n float *val,\n float *x,\n float *Ax,\n float *p,\n float *r,\n double *dot_result,\n int nnz,\n int N,\n float tol,\n MultiDeviceData multi_device_data)\n{\n cg::thread_block cta = cg::this_thread_block();\n cg::grid_group grid = cg::this_grid();\n PeerGroup peer_group(multi_device_data, grid);\n\n const int max_iter = 10000;\n\n float alpha = 1.0;\n float alpham1 = -1.0;\n float r0 = 0.0, r1, b, a, na;\n\n for (int i = peer_group.thread_rank(); i < N; i += peer_group.size()) {\n r[i] = 1.0;\n x[i] = 0.0;\n }\n\n cg::sync(grid);\n\n gpuSpMV(I, J, val, nnz, N, alpha, x, Ax, peer_group);\n\n cg::sync(grid);\n\n gpuSaxpy(Ax, r, alpham1, N, peer_group);\n\n cg::sync(grid);\n\n gpuDotProduct(r, r, N, cta, peer_group);\n\n cg::sync(grid);\n\n if (grid.thread_rank() == 0) {\n atomicAdd_system(dot_result, grid_dot_result);\n grid_dot_result = 0.0;\n }\n peer_group.sync();\n\n r1 = *dot_result;\n\n int k = 1;\n while (r1 > tol * tol && k <= max_iter) {\n if (k > 1) {\n b = r1 / r0;\n gpuScaleVectorAndSaxpy(r, p, alpha, b, N, peer_group);\n }\n else {\n gpuCopyVector(r, p, N, peer_group);\n }\n\n peer_group.sync();\n\n gpuSpMV(I, J, val, nnz, N, alpha, p, Ax, peer_group);\n\n if (peer_group.thread_rank() == 0) {\n *dot_result = 0.0;\n }\n peer_group.sync();\n\n gpuDotProduct(p, Ax, N, cta, peer_group);\n\n cg::sync(grid);\n\n if (grid.thread_rank() == 0) {\n atomicAdd_system(dot_result, grid_dot_result);\n grid_dot_result = 0.0;\n }\n peer_group.sync();\n\n a = r1 / *dot_result;\n\n gpuSaxpy(p, x, a, N, peer_group);\n\n na = -a;\n\n gpuSaxpy(Ax, r, na, N, peer_group);\n\n r0 = r1;\n\n peer_group.sync();\n\n if (peer_group.thread_rank() == 0) {\n *dot_result = 0.0;\n }\n\n peer_group.sync();\n\n gpuDotProduct(r, r, N, cta, peer_group);\n\n cg::sync(grid);\n\n if (grid.thread_rank() == 0) {\n atomicAdd_system(dot_result, grid_dot_result);\n grid_dot_result = 0.0;\n }\n peer_group.sync();\n\n r1 = *dot_result;\n k++;\n }\n}\n\n// Map of device version to device number\nstd::multimap, int> getIdenticalGPUs()\n{\n int numGpus = 0;\n checkCudaErrors(hipGetDeviceCount(&numGpus));\n\n std::multimap, int> identicalGpus;\n\n for (int i = 0; i < numGpus; i++) {\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, i));\n\n // Filter unsupported devices\n if (deviceProp.cooperativeLaunch && deviceProp.concurrentManagedAccess) {\n identicalGpus.emplace(std::make_pair(deviceProp.major, deviceProp.minor), i);\n }\n printf(\"GPU Device %d: \\\"%s\\\" with compute capability %d.%d\\n\",\n i,\n deviceProp.name,\n deviceProp.major,\n deviceProp.minor);\n }\n\n return identicalGpus;\n}\n\nint main(int argc, char **argv)\n{\n constexpr size_t kNumGpusRequired = 2;\n int N = 0, nz = 0, *I = NULL, *J = NULL;\n float *val = NULL;\n const float tol = 1e-5f;\n float *x;\n float rhs = 1.0;\n float r1;\n float *r, *p, *Ax;\n\n printf(\"Starting [%s]...\\n\", sSDKname);\n auto gpusByArch = getIdenticalGPUs();\n\n auto it = gpusByArch.begin();\n auto end = gpusByArch.end();\n\n auto bestFit = std::make_pair(it, it);\n // use std::distance to find the largest number of GPUs amongst architectures\n auto distance = [](decltype(bestFit) p) { return std::distance(p.first, p.second); };\n\n // Read each unique key/pair element in order\n for (; it != end; it = gpusByArch.upper_bound(it->first)) {\n // first and second are iterators bounded within the architecture group\n auto testFit = gpusByArch.equal_range(it->first);\n // Always use devices with highest architecture version or whichever has the\n // most devices available\n if (distance(bestFit) <= distance(testFit))\n bestFit = testFit;\n }\n\n if (distance(bestFit) < kNumGpusRequired) {\n printf(\"No two or more GPUs with same architecture capable of \"\n \"concurrentManagedAccess found. \"\n \"\\nWaiving the sample\\n\");\n exit(EXIT_WAIVED);\n }\n\n std::set bestFitDeviceIds;\n\n // Check & select peer-to-peer access capable GPU devices as enabling p2p\n // access between participating GPUs gives better performance.\n for (auto itr = bestFit.first; itr != bestFit.second; itr++) {\n int deviceId = itr->second;\n checkCudaErrors(hipSetDevice(deviceId));\n\n std::for_each(itr, bestFit.second, [&deviceId, &bestFitDeviceIds, &kNumGpusRequired](decltype(*itr) mapPair) {\n if (deviceId != mapPair.second) {\n int access = 0;\n checkCudaErrors(hipDeviceCanAccessPeer(&access, deviceId, mapPair.second));\n printf(\"Device=%d %s Access Peer Device=%d\\n\", deviceId, access ? \"CAN\" : \"CANNOT\", mapPair.second);\n if (access && bestFitDeviceIds.size() < kNumGpusRequired) {\n bestFitDeviceIds.emplace(deviceId);\n bestFitDeviceIds.emplace(mapPair.second);\n }\n else {\n printf(\"Ignoring device %i (max devices exceeded)\\n\", mapPair.second);\n }\n }\n });\n\n if (bestFitDeviceIds.size() >= kNumGpusRequired) {\n printf(\"Selected p2p capable devices - \");\n for (auto devicesItr = bestFitDeviceIds.begin(); devicesItr != bestFitDeviceIds.end(); devicesItr++) {\n printf(\"deviceId = %d \", *devicesItr);\n }\n printf(\"\\n\");\n break;\n }\n }\n\n // if bestFitDeviceIds.size() == 0 it means the GPUs in system are not p2p\n // capable, hence we add it without p2p capability check.\n if (!bestFitDeviceIds.size()) {\n printf(\"Devices involved are not p2p capable.. selecting %zu of them\\n\", kNumGpusRequired);\n std::for_each(\n bestFit.first, bestFit.second, [&bestFitDeviceIds, &kNumGpusRequired](decltype(*bestFit.first) mapPair) {\n if (bestFitDeviceIds.size() < kNumGpusRequired) {\n bestFitDeviceIds.emplace(mapPair.second);\n }\n else {\n printf(\"Ignoring device %i (max devices exceeded)\\n\", mapPair.second);\n }\n // Insert the sequence into the deviceIds set\n });\n }\n else {\n // perform cudaDeviceEnablePeerAccess in both directions for all\n // participating devices.\n for (auto p1_itr = bestFitDeviceIds.begin(); p1_itr != bestFitDeviceIds.end(); p1_itr++) {\n checkCudaErrors(hipSetDevice(*p1_itr));\n for (auto p2_itr = bestFitDeviceIds.begin(); p2_itr != bestFitDeviceIds.end(); p2_itr++) {\n if (*p1_itr != *p2_itr) {\n checkCudaErrors(hipDeviceEnablePeerAccess(*p2_itr, 0));\n checkCudaErrors(hipSetDevice(*p1_itr));\n }\n }\n }\n }\n\n /* Generate a random tridiagonal symmetric matrix in CSR format */\n N = 10485760 * 2;\n nz = (N - 2) * 3 + 4;\n\n checkCudaErrors(hipMallocManaged((void **)&I, sizeof(int) * (N + 1)));\n checkCudaErrors(hipMallocManaged((void **)&J, sizeof(int) * nz));\n checkCudaErrors(hipMallocManaged((void **)&val, sizeof(float) * nz));\n\n float *val_cpu = (float *)malloc(sizeof(float) * nz);\n\n genTridiag(I, J, val_cpu, N, nz);\n\n memcpy(val, val_cpu, sizeof(float) * nz);\n hipMemLocation deviceLoc;\n deviceLoc.type = hipMemLocationTypeDevice;\n deviceLoc.id = 0; // Device location with initial device 0\n checkCudaErrors(hipMemAdvise(I, sizeof(int) * (N + 1), hipMemAdviseSetReadMostly, deviceLoc));\n checkCudaErrors(hipMemAdvise(J, sizeof(int) * nz, hipMemAdviseSetReadMostly, deviceLoc));\n checkCudaErrors(hipMemAdvise(val, sizeof(float) * nz, hipMemAdviseSetReadMostly, deviceLoc));\n\n checkCudaErrors(hipMallocManaged((void **)&x, sizeof(float) * N));\n\n double *dot_result;\n checkCudaErrors(hipMallocManaged((void **)&dot_result, sizeof(double)));\n\n checkCudaErrors(hipMemset(dot_result, 0, sizeof(double)));\n\n // temp memory for ConjugateGradient\n checkCudaErrors(hipMallocManaged((void **)&r, N * sizeof(float)));\n checkCudaErrors(hipMallocManaged((void **)&p, N * sizeof(float)));\n checkCudaErrors(hipMallocManaged((void **)&Ax, N * sizeof(float)));\n\n std::cout << \"\\nRunning on GPUs = \" << kNumGpusRequired << std::endl;\n hipStream_t nStreams[kNumGpusRequired];\n\n int sMemSize = sizeof(double) * ((THREADS_PER_BLOCK / 32) + 1);\n int numBlocksPerSm = INT_MAX;\n int numThreads = THREADS_PER_BLOCK;\n int numSms = INT_MAX;\n auto deviceId = bestFitDeviceIds.begin();\n\n // set numSms & numBlocksPerSm to be lowest of 2 devices\n while (deviceId != bestFitDeviceIds.end()) {\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipSetDevice(*deviceId));\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, *deviceId));\n\n int numBlocksPerSm_current = 0;\n checkCudaErrors(hipOccupancyMaxActiveBlocksPerMultiprocessor(\n &numBlocksPerSm_current, multiGpuConjugateGradient, numThreads, sMemSize));\n\n if (numBlocksPerSm > numBlocksPerSm_current) {\n numBlocksPerSm = numBlocksPerSm_current;\n }\n if (numSms > deviceProp.multiProcessorCount) {\n numSms = deviceProp.multiProcessorCount;\n }\n deviceId++;\n }\n\n if (!numBlocksPerSm) {\n printf(\"Max active blocks per SM is returned as 0.\\n Hence, Waiving the \"\n \"sample\\n\");\n exit(EXIT_WAIVED);\n }\n\n int device_count = 0;\n int totalThreadsPerGPU = numSms * numBlocksPerSm * THREADS_PER_BLOCK;\n deviceId = bestFitDeviceIds.begin();\n while (deviceId != bestFitDeviceIds.end()) {\n checkCudaErrors(hipSetDevice(*deviceId));\n checkCudaErrors(hipStreamCreate(&nStreams[device_count]));\n\n int perGPUIter = N / (totalThreadsPerGPU * kNumGpusRequired);\n int offset_Ax = device_count * totalThreadsPerGPU;\n int offset_r = device_count * totalThreadsPerGPU;\n int offset_p = device_count * totalThreadsPerGPU;\n int offset_x = device_count * totalThreadsPerGPU;\n\n // Create device location with specific device ID\n hipMemLocation deviceLoc;\n deviceLoc.type = hipMemLocationTypeDevice;\n deviceLoc.id = *deviceId;\n checkCudaErrors(hipMemPrefetchAsync(I, sizeof(int) * N, deviceLoc, 0, nStreams[device_count]));\n checkCudaErrors(hipMemPrefetchAsync(val, sizeof(float) * nz, deviceLoc, 0, nStreams[device_count]));\n checkCudaErrors(hipMemPrefetchAsync(J, sizeof(float) * nz, deviceLoc, 0, nStreams[device_count]));\n\n if (offset_Ax <= N) {\n for (int i = 0; i < perGPUIter; i++) {\n hipMemAdvise(\n Ax + offset_Ax, sizeof(float) * totalThreadsPerGPU, hipMemAdviseSetPreferredLocation, deviceLoc);\n hipMemAdvise(\n r + offset_r, sizeof(float) * totalThreadsPerGPU, hipMemAdviseSetPreferredLocation, deviceLoc);\n hipMemAdvise(\n x + offset_x, sizeof(float) * totalThreadsPerGPU, hipMemAdviseSetPreferredLocation, deviceLoc);\n hipMemAdvise(\n p + offset_p, sizeof(float) * totalThreadsPerGPU, hipMemAdviseSetPreferredLocation, deviceLoc);\n\n hipMemAdvise(\n Ax + offset_Ax, sizeof(float) * totalThreadsPerGPU, hipMemAdviseSetAccessedBy, deviceLoc);\n hipMemAdvise(r + offset_r, sizeof(float) * totalThreadsPerGPU, hipMemAdviseSetAccessedBy, deviceLoc);\n hipMemAdvise(p + offset_p, sizeof(float) * totalThreadsPerGPU, hipMemAdviseSetAccessedBy, deviceLoc);\n hipMemAdvise(x + offset_x, sizeof(float) * totalThreadsPerGPU, hipMemAdviseSetAccessedBy, deviceLoc);\n\n offset_Ax += totalThreadsPerGPU * kNumGpusRequired;\n offset_r += totalThreadsPerGPU * kNumGpusRequired;\n offset_p += totalThreadsPerGPU * kNumGpusRequired;\n offset_x += totalThreadsPerGPU * kNumGpusRequired;\n\n if (offset_Ax >= N) {\n break;\n }\n }\n }\n\n device_count++;\n deviceId++;\n }\n\n#if ENABLE_CPU_DEBUG_CODE\n float *Ax_cpu = (float *)malloc(sizeof(float) * N);\n float *r_cpu = (float *)malloc(sizeof(float) * N);\n float *p_cpu = (float *)malloc(sizeof(float) * N);\n float *x_cpu = (float *)malloc(sizeof(float) * N);\n\n for (int i = 0; i < N; i++) {\n r_cpu[i] = 1.0;\n Ax_cpu[i] = x_cpu[i] = 0.0;\n }\n#endif\n\n printf(\"Total threads per GPU = %d numBlocksPerSm = %d\\n\",\n numSms * numBlocksPerSm * THREADS_PER_BLOCK,\n numBlocksPerSm);\n dim3 dimGrid(numSms * numBlocksPerSm, 1, 1), dimBlock(THREADS_PER_BLOCK, 1, 1);\n\n // Structure used for cross-grid synchronization.\n MultiDeviceData multi_device_data;\n checkCudaErrors(hipHostAlloc(&multi_device_data.hostMemoryArrivedList,\n (kNumGpusRequired - 1) * sizeof(*multi_device_data.hostMemoryArrivedList),\n hipHostMallocPortable));\n memset(multi_device_data.hostMemoryArrivedList,\n 0,\n (kNumGpusRequired - 1) * sizeof(*multi_device_data.hostMemoryArrivedList));\n multi_device_data.numDevices = kNumGpusRequired;\n multi_device_data.deviceRank = 0;\n\n void *kernelArgs[] = {\n (void *)&I,\n (void *)&J,\n (void *)&val,\n (void *)&x,\n (void *)&Ax,\n (void *)&p,\n (void *)&r,\n (void *)&dot_result,\n (void *)&nz,\n (void *)&N,\n (void *)&tol,\n (void *)&multi_device_data,\n };\n\n printf(\"Launching kernel\\n\");\n\n deviceId = bestFitDeviceIds.begin();\n device_count = 0;\n while (deviceId != bestFitDeviceIds.end()) {\n checkCudaErrors(hipSetDevice(*deviceId));\n checkCudaErrors(hipLaunchCooperativeKernel(reinterpret_cast((void *)multiGpuConjugateGradient), dimGrid, dimBlock, kernelArgs, sMemSize, nStreams[device_count++]));\n multi_device_data.deviceRank++;\n deviceId++;\n }\n\n // Use cudaMemLocationTypeHost for optimal host memory location\n hipMemLocation hostLoc;\n hostLoc.type = hipMemLocationTypeHost;\n checkCudaErrors(hipMemPrefetchAsync(x, sizeof(float) * N, hostLoc, 0));\n checkCudaErrors(hipMemPrefetchAsync(dot_result, sizeof(double), hostLoc, 0));\n\n deviceId = bestFitDeviceIds.begin();\n device_count = 0;\n while (deviceId != bestFitDeviceIds.end()) {\n checkCudaErrors(hipSetDevice(*deviceId));\n checkCudaErrors(hipStreamSynchronize(nStreams[device_count++]));\n deviceId++;\n }\n\n r1 = (float)*dot_result;\n\n printf(\"GPU Final, residual = %e \\n \", sqrt(r1));\n\n#if ENABLE_CPU_DEBUG_CODE\n cpuConjugateGrad(I, J, val, x_cpu, Ax_cpu, p_cpu, r_cpu, nz, N, tol);\n#endif\n\n float rsum, diff, err = 0.0;\n\n for (int i = 0; i < N; i++) {\n rsum = 0.0;\n\n for (int j = I[i]; j < I[i + 1]; j++) {\n rsum += val_cpu[j] * x[J[j]];\n }\n\n diff = fabs(rsum - rhs);\n\n if (diff > err) {\n err = diff;\n }\n }\n\n checkCudaErrors(hipHostFree(multi_device_data.hostMemoryArrivedList));\n checkCudaErrors(hipFree(I));\n checkCudaErrors(hipFree(J));\n checkCudaErrors(hipFree(val));\n checkCudaErrors(hipFree(x));\n checkCudaErrors(hipFree(r));\n checkCudaErrors(hipFree(p));\n checkCudaErrors(hipFree(Ax));\n checkCudaErrors(hipFree(dot_result));\n free(val_cpu);\n\n#if ENABLE_CPU_DEBUG_CODE\n free(Ax_cpu);\n free(r_cpu);\n free(p_cpu);\n free(x_cpu);\n#endif\n\n printf(\"Test Summary: Error amount = %f \\n\", err);\n fprintf(stdout, \"&&&& conjugateGradientMultiDeviceCG %s\\n\", (sqrt(r1) < tol) ? \"PASSED\" : \"FAILED\");\n exit((sqrt(r1) < tol) ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/4_CUDA_Libraries/conjugateGradientMultiDeviceCG/conjugateGradientMultiDeviceCG.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n///////////////////////////////////////////////////////////////////////////////\n#include \n#include \n\n// Round a / b to nearest higher integer value\nint cuda_iDivUp(int a, int b) { return (a + (b - 1)) / b; }\n\n// complex math functions\n__device__ float2 conjugate(float2 arg) { return make_float2(arg.x, -arg.y); }\n\n__device__ float2 complex_exp(float arg) { return make_float2(cosf(arg), sinf(arg)); }\n\n__device__ float2 complex_add(float2 a, float2 b) { return make_float2(a.x + b.x, a.y + b.y); }\n\n__device__ float2 complex_mult(float2 ab, float2 cd)\n{\n return make_float2(ab.x * cd.x - ab.y * cd.y, ab.x * cd.y + ab.y * cd.x);\n}\n\n// generate wave heightfield at time t based on initial heightfield and\n// dispersion relationship\n__global__ void generateSpectrumKernel(float2 *h0,\n float2 *ht,\n unsigned int in_width,\n unsigned int out_width,\n unsigned int out_height,\n float t,\n float patchSize)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n unsigned int in_index = y * in_width + x;\n unsigned int in_mindex = (out_height - y) * in_width + (out_width - x); // mirrored\n unsigned int out_index = y * out_width + x;\n\n // calculate wave vector\n float2 k;\n k.x = (-(int)out_width / 2.0f + x) * (2.0f * CUDART_PI_F / patchSize);\n k.y = (-(int)out_width / 2.0f + y) * (2.0f * CUDART_PI_F / patchSize);\n\n // calculate dispersion w(k)\n float k_len = sqrtf(k.x * k.x + k.y * k.y);\n float w = sqrtf(9.81f * k_len);\n\n if ((x < out_width) && (y < out_height)) {\n float2 h0_k = h0[in_index];\n float2 h0_mk = h0[in_mindex];\n\n // output frequency-space complex values\n ht[out_index] =\n complex_add(complex_mult(h0_k, complex_exp(w * t)), complex_mult(conjugate(h0_mk), complex_exp(-w * t)));\n // ht[out_index] = h0_k;\n }\n}\n\n// update height map values based on output of FFT\n__global__ void updateHeightmapKernel(float *heightMap, float2 *ht, unsigned int width)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n unsigned int i = y * width + x;\n\n // cos(pi * (m1 + m2))\n float sign_correction = ((x + y) & 0x01) ? -1.0f : 1.0f;\n\n heightMap[i] = ht[i].x * sign_correction;\n}\n\n// update height map values based on output of FFT\n__global__ void updateHeightmapKernel_y(float *heightMap, float2 *ht, unsigned int width)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n unsigned int i = y * width + x;\n\n // cos(pi * (m1 + m2))\n float sign_correction = ((x + y) & 0x01) ? -1.0f : 1.0f;\n\n heightMap[i] = ht[i].y * sign_correction;\n}\n\n// generate slope by partial differences in spatial domain\n__global__ void calculateSlopeKernel(float *h, float2 *slopeOut, unsigned int width, unsigned int height)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n unsigned int i = y * width + x;\n\n float2 slope = make_float2(0.0f, 0.0f);\n\n if ((x > 0) && (y > 0) && (x < width - 1) && (y < height - 1)) {\n slope.x = h[i + 1] - h[i - 1];\n slope.y = h[i + width] - h[i - width];\n }\n\n slopeOut[i] = slope;\n}\n\n// wrapper functions\nextern \"C\" void cudaGenerateSpectrumKernel(float2 *d_h0,\n float2 *d_ht,\n unsigned int in_width,\n unsigned int out_width,\n unsigned int out_height,\n float animTime,\n float patchSize)\n{\n dim3 block(8, 8, 1);\n dim3 grid(cuda_iDivUp(out_width, block.x), cuda_iDivUp(out_height, block.y), 1);\n generateSpectrumKernel<<>>(d_h0, d_ht, in_width, out_width, out_height, animTime, patchSize);\n}\n\nextern \"C\" void\ncudaUpdateHeightmapKernel(float *d_heightMap, float2 *d_ht, unsigned int width, unsigned int height, bool autoTest)\n{\n dim3 block(8, 8, 1);\n dim3 grid(cuda_iDivUp(width, block.x), cuda_iDivUp(height, block.y), 1);\n if (autoTest) {\n updateHeightmapKernel_y<<>>(d_heightMap, d_ht, width);\n }\n else {\n updateHeightmapKernel<<>>(d_heightMap, d_ht, width);\n }\n}\n\nextern \"C\" void cudaCalculateSlopeKernel(float *hptr, float2 *slopeOut, unsigned int width, unsigned int height)\n{\n dim3 block(8, 8, 1);\n dim3 grid2(cuda_iDivUp(width, block.x), cuda_iDivUp(height, block.y), 1);\n calculateSlopeKernel<<>>(hptr, slopeOut, width, height);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n///////////////////////////////////////////////////////////////////////////////\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n// Round a / b to nearest higher integer value\nint cuda_iDivUp(int a, int b) { return (a + (b - 1)) / b; }\n\n// complex math functions\n__device__ float2 conjugate(float2 arg) { return make_float2(arg.x, -arg.y); }\n\n__device__ float2 complex_exp(float arg) { return make_float2(cosf(arg), sinf(arg)); }\n\n__device__ float2 complex_add(float2 a, float2 b) { return make_float2(a.x + b.x, a.y + b.y); }\n\n__device__ float2 complex_mult(float2 ab, float2 cd)\n{\n return make_float2(ab.x * cd.x - ab.y * cd.y, ab.x * cd.y + ab.y * cd.x);\n}\n\n// generate wave heightfield at time t based on initial heightfield and\n// dispersion relationship\n__global__ void generateSpectrumKernel(float2 *h0,\n float2 *ht,\n unsigned int in_width,\n unsigned int out_width,\n unsigned int out_height,\n float t,\n float patchSize)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n unsigned int in_index = y * in_width + x;\n unsigned int in_mindex = (out_height - y) * in_width + (out_width - x); // mirrored\n unsigned int out_index = y * out_width + x;\n\n // calculate wave vector\n float2 k;\n k.x = (-(int)out_width / 2.0f + x) * (2.0f * HIP_PI_F / patchSize);\n k.y = (-(int)out_width / 2.0f + y) * (2.0f * HIP_PI_F / patchSize);\n\n // calculate dispersion w(k)\n float k_len = sqrtf(k.x * k.x + k.y * k.y);\n float w = sqrtf(9.81f * k_len);\n\n if ((x < out_width) && (y < out_height)) {\n float2 h0_k = h0[in_index];\n float2 h0_mk = h0[in_mindex];\n\n // output frequency-space complex values\n ht[out_index] =\n complex_add(complex_mult(h0_k, complex_exp(w * t)), complex_mult(conjugate(h0_mk), complex_exp(-w * t)));\n // ht[out_index] = h0_k;\n }\n}\n\n// update height map values based on output of FFT\n__global__ void updateHeightmapKernel(float *heightMap, float2 *ht, unsigned int width)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n unsigned int i = y * width + x;\n\n // cos(pi * (m1 + m2))\n float sign_correction = ((x + y) & 0x01) ? -1.0f : 1.0f;\n\n heightMap[i] = ht[i].x * sign_correction;\n}\n\n// update height map values based on output of FFT\n__global__ void updateHeightmapKernel_y(float *heightMap, float2 *ht, unsigned int width)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n unsigned int i = y * width + x;\n\n // cos(pi * (m1 + m2))\n float sign_correction = ((x + y) & 0x01) ? -1.0f : 1.0f;\n\n heightMap[i] = ht[i].y * sign_correction;\n}\n\n// generate slope by partial differences in spatial domain\n__global__ void calculateSlopeKernel(float *h, float2 *slopeOut, unsigned int width, unsigned int height)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;\n unsigned int i = y * width + x;\n\n float2 slope = make_float2(0.0f, 0.0f);\n\n if ((x > 0) && (y > 0) && (x < width - 1) && (y < height - 1)) {\n slope.x = h[i + 1] - h[i - 1];\n slope.y = h[i + width] - h[i - width];\n }\n\n slopeOut[i] = slope;\n}\n\n// wrapper functions\nextern \"C\" void cudaGenerateSpectrumKernel(float2 *d_h0,\n float2 *d_ht,\n unsigned int in_width,\n unsigned int out_width,\n unsigned int out_height,\n float animTime,\n float patchSize)\n{\n dim3 block(8, 8, 1);\n dim3 grid(cuda_iDivUp(out_width, block.x), cuda_iDivUp(out_height, block.y), 1);\n generateSpectrumKernel<<>>(d_h0, d_ht, in_width, out_width, out_height, animTime, patchSize);\n}\n\nextern \"C\" void\ncudaUpdateHeightmapKernel(float *d_heightMap, float2 *d_ht, unsigned int width, unsigned int height, bool autoTest)\n{\n dim3 block(8, 8, 1);\n dim3 grid(cuda_iDivUp(width, block.x), cuda_iDivUp(height, block.y), 1);\n if (autoTest) {\n updateHeightmapKernel_y<<>>(d_heightMap, d_ht, width);\n }\n else {\n updateHeightmapKernel<<>>(d_heightMap, d_ht, width);\n }\n}\n\nextern \"C\" void cudaCalculateSlopeKernel(float *hptr, float2 *slopeOut, unsigned int width, unsigned int height)\n{\n dim3 block(8, 8, 1);\n dim3 grid2(cuda_iDivUp(width, block.x), cuda_iDivUp(height, block.y), 1);\n calculateSlopeKernel<<>>(hptr, slopeOut, width, height);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/4_CUDA_Libraries/oceanFFT/oceanFFT_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample implements a conjugate gradient solver on GPU using\n * Multi Block Cooperative Groups, also uses Unified Memory.\n *\n */\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// Utilities and system includes\n#include \n#include \n#include // helper function CUDA error checking and initialization\n#include // helper for shared functions common to CUDA Samples\n\nnamespace cg = cooperative_groups;\n\nconst char *sSDKname = \"conjugateGradientMultiBlockCG\";\n\n#define ENABLE_CPU_DEBUG_CODE 0\n#define THREADS_PER_BLOCK 512\n\n/* genTridiag: generate a random tridiagonal symmetric matrix */\nvoid genTridiag(int *I, int *J, float *val, int N, int nz)\n{\n I[0] = 0, J[0] = 0, J[1] = 1;\n val[0] = static_cast(rand()) / RAND_MAX + 10.0f;\n val[1] = static_cast(rand()) / RAND_MAX;\n int start;\n\n for (int i = 1; i < N; i++) {\n if (i > 1) {\n I[i] = I[i - 1] + 3;\n }\n else {\n I[1] = 2;\n }\n\n start = (i - 1) * 3 + 2;\n J[start] = i - 1;\n J[start + 1] = i;\n\n if (i < N - 1) {\n J[start + 2] = i + 1;\n }\n\n val[start] = val[start - 1];\n val[start + 1] = static_cast(rand()) / RAND_MAX + 10.0f;\n\n if (i < N - 1) {\n val[start + 2] = static_cast(rand()) / RAND_MAX;\n }\n }\n\n I[N] = nz;\n}\n\n// I - contains location of the given non-zero element in the row of the matrix\n// J - contains location of the given non-zero element in the column of the\n// matrix val - contains values of the given non-zero elements of the matrix\n// inputVecX - input vector to be multiplied\n// outputVecY - resultant vector\nvoid cpuSpMV(int *I, int *J, float *val, int nnz, int num_rows, float alpha, float *inputVecX, float *outputVecY)\n{\n for (int i = 0; i < num_rows; i++) {\n int num_elems_this_row = I[i + 1] - I[i];\n\n float output = 0.0;\n for (int j = 0; j < num_elems_this_row; j++) {\n output += alpha * val[I[i] + j] * inputVecX[J[I[i] + j]];\n }\n outputVecY[i] = output;\n }\n\n return;\n}\n\ndouble dotProduct(float *vecA, float *vecB, int size)\n{\n double result = 0.0;\n\n for (int i = 0; i < size; i++) {\n result = result + (vecA[i] * vecB[i]);\n }\n\n return result;\n}\n\nvoid scaleVector(float *vec, float alpha, int size)\n{\n for (int i = 0; i < size; i++) {\n vec[i] = alpha * vec[i];\n }\n}\n\nvoid saxpy(float *x, float *y, float a, int size)\n{\n for (int i = 0; i < size; i++) {\n y[i] = a * x[i] + y[i];\n }\n}\n\nvoid cpuConjugateGrad(int *I, int *J, float *val, float *x, float *Ax, float *p, float *r, int nnz, int N, float tol)\n{\n int max_iter = 10000;\n\n float alpha = 1.0;\n float alpham1 = -1.0;\n float r0 = 0.0, b, a, na;\n\n cpuSpMV(I, J, val, nnz, N, alpha, x, Ax);\n saxpy(Ax, r, alpham1, N);\n\n float r1 = dotProduct(r, r, N);\n\n int k = 1;\n\n while (r1 > tol * tol && k <= max_iter) {\n if (k > 1) {\n b = r1 / r0;\n scaleVector(p, b, N);\n\n saxpy(r, p, alpha, N);\n }\n else {\n for (int i = 0; i < N; i++)\n p[i] = r[i];\n }\n\n cpuSpMV(I, J, val, nnz, N, alpha, p, Ax);\n\n float dot = dotProduct(p, Ax, N);\n a = r1 / dot;\n\n saxpy(p, x, a, N);\n na = -a;\n saxpy(Ax, r, na, N);\n\n r0 = r1;\n r1 = dotProduct(r, r, N);\n\n printf(\"\\nCPU code iteration = %3d, residual = %e\\n\", k, sqrt(r1));\n k++;\n }\n}\n\n__device__ void gpuSpMV(int *I,\n int *J,\n float *val,\n int nnz,\n int num_rows,\n float alpha,\n float *inputVecX,\n float *outputVecY,\n cg::thread_block &cta,\n const cg::grid_group &grid)\n{\n for (int i = grid.thread_rank(); i < num_rows; i += grid.size()) {\n int row_elem = I[i];\n int next_row_elem = I[i + 1];\n int num_elems_this_row = next_row_elem - row_elem;\n\n float output = 0.0;\n for (int j = 0; j < num_elems_this_row; j++) {\n // I or J or val arrays - can be put in shared memory\n // as the access is random and reused in next calls of gpuSpMV function.\n output += alpha * val[row_elem + j] * inputVecX[J[row_elem + j]];\n }\n\n outputVecY[i] = output;\n }\n}\n\n__device__ void gpuSaxpy(float *x, float *y, float a, int size, const cg::grid_group &grid)\n{\n for (int i = grid.thread_rank(); i < size; i += grid.size()) {\n y[i] = a * x[i] + y[i];\n }\n}\n\n__device__ void gpuDotProduct(float *vecA,\n float *vecB,\n double *result,\n int size,\n const cg::thread_block &cta,\n const cg::grid_group &grid)\n{\n extern __shared__ double tmp[];\n\n double temp_sum = 0.0;\n for (int i = grid.thread_rank(); i < size; i += grid.size()) {\n temp_sum += static_cast(vecA[i] * vecB[i]);\n }\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n temp_sum = cg::reduce(tile32, temp_sum, cg::plus());\n\n if (tile32.thread_rank() == 0) {\n tmp[tile32.meta_group_rank()] = temp_sum;\n }\n\n cg::sync(cta);\n\n if (tile32.meta_group_rank() == 0) {\n temp_sum = tile32.thread_rank() < tile32.meta_group_size() ? tmp[tile32.thread_rank()] : 0.0;\n temp_sum = cg::reduce(tile32, temp_sum, cg::plus());\n\n if (tile32.thread_rank() == 0) {\n atomicAdd(result, temp_sum);\n }\n }\n}\n\n__device__ void gpuCopyVector(float *srcA, float *destB, int size, const cg::grid_group &grid)\n{\n for (int i = grid.thread_rank(); i < size; i += grid.size()) {\n destB[i] = srcA[i];\n }\n}\n\n__device__ void\ngpuScaleVectorAndSaxpy(const float *x, float *y, float a, float scale, int size, const cg::grid_group &grid)\n{\n for (int i = grid.thread_rank(); i < size; i += grid.size()) {\n y[i] = a * x[i] + scale * y[i];\n }\n}\n\nextern \"C\" __global__ void gpuConjugateGradient(int *I,\n int *J,\n float *val,\n float *x,\n float *Ax,\n float *p,\n float *r,\n double *dot_result,\n int nnz,\n int N,\n float tol)\n{\n cg::thread_block cta = cg::this_thread_block();\n cg::grid_group grid = cg::this_grid();\n\n int max_iter = 10000;\n\n float alpha = 1.0;\n float alpham1 = -1.0;\n float r0 = 0.0, r1, b, a, na;\n\n gpuSpMV(I, J, val, nnz, N, alpha, x, Ax, cta, grid);\n\n cg::sync(grid);\n\n gpuSaxpy(Ax, r, alpham1, N, grid);\n\n cg::sync(grid);\n\n gpuDotProduct(r, r, dot_result, N, cta, grid);\n\n cg::sync(grid);\n\n r1 = *dot_result;\n\n int k = 1;\n while (r1 > tol * tol && k <= max_iter) {\n if (k > 1) {\n b = r1 / r0;\n gpuScaleVectorAndSaxpy(r, p, alpha, b, N, grid);\n }\n else {\n gpuCopyVector(r, p, N, grid);\n }\n\n cg::sync(grid);\n\n gpuSpMV(I, J, val, nnz, N, alpha, p, Ax, cta, grid);\n\n if (threadIdx.x == 0 && blockIdx.x == 0)\n *dot_result = 0.0;\n\n cg::sync(grid);\n\n gpuDotProduct(p, Ax, dot_result, N, cta, grid);\n\n cg::sync(grid);\n\n a = r1 / *dot_result;\n\n gpuSaxpy(p, x, a, N, grid);\n na = -a;\n gpuSaxpy(Ax, r, na, N, grid);\n\n r0 = r1;\n\n cg::sync(grid);\n if (threadIdx.x == 0 && blockIdx.x == 0)\n *dot_result = 0.0;\n\n cg::sync(grid);\n\n gpuDotProduct(r, r, dot_result, N, cta, grid);\n\n cg::sync(grid);\n\n r1 = *dot_result;\n k++;\n }\n}\n\nbool areAlmostEqual(float a, float b, float maxRelDiff)\n{\n float diff = fabsf(a - b);\n float abs_a = fabsf(a);\n float abs_b = fabsf(b);\n float largest = abs_a > abs_b ? abs_a : abs_b;\n\n if (diff <= largest * maxRelDiff) {\n return true;\n }\n else {\n printf(\"maxRelDiff = %.8e\\n\", maxRelDiff);\n printf(\"diff %.8e > largest * maxRelDiff %.8e therefore %.8e and %.8e are not \"\n \"same\\n\",\n diff,\n largest * maxRelDiff,\n a,\n b);\n return false;\n }\n}\n\nint main(int argc, char **argv)\n{\n int N = 0, nz = 0, *I = NULL, *J = NULL;\n float *val = NULL;\n const float tol = 1e-5f;\n float *x;\n float *rhs;\n float r1;\n float *r, *p, *Ax;\n cudaEvent_t start, stop;\n\n printf(\"Starting [%s]...\\n\", sSDKname);\n\n // This will pick the best possible CUDA capable device\n cudaDeviceProp deviceProp;\n int devID = findCudaDevice(argc, (const char **)argv);\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, devID));\n\n if (!deviceProp.managedMemory) {\n // This sample requires being run on a device that supports Unified Memory\n fprintf(stderr, \"Unified Memory not supported on this device\\n\");\n exit(EXIT_WAIVED);\n }\n\n // This sample requires being run on a device that supports Cooperative Kernel\n // Launch\n if (!deviceProp.cooperativeLaunch) {\n printf(\"\\nSelected GPU (%d) does not support Cooperative Kernel Launch, \"\n \"Waiving the run\\n\",\n devID);\n exit(EXIT_WAIVED);\n }\n\n // Statistics about the GPU device\n printf(\"> GPU device has %d Multi-Processors, SM %d.%d compute capabilities\\n\\n\",\n deviceProp.multiProcessorCount,\n deviceProp.major,\n deviceProp.minor);\n\n /* Generate a random tridiagonal symmetric matrix in CSR format */\n N = 1048576;\n nz = (N - 2) * 3 + 4;\n\n cudaMallocManaged(reinterpret_cast(&I), sizeof(int) * (N + 1));\n cudaMallocManaged(reinterpret_cast(&J), sizeof(int) * nz);\n cudaMallocManaged(reinterpret_cast(&val), sizeof(float) * nz);\n\n genTridiag(I, J, val, N, nz);\n\n cudaMallocManaged(reinterpret_cast(&x), sizeof(float) * N);\n cudaMallocManaged(reinterpret_cast(&rhs), sizeof(float) * N);\n\n double *dot_result;\n\n cudaMallocManaged(reinterpret_cast(&dot_result), sizeof(double));\n\n *dot_result = 0.0;\n\n // temp memory for CG\n checkCudaErrors(cudaMallocManaged(reinterpret_cast(&r), N * sizeof(float)));\n checkCudaErrors(cudaMallocManaged(reinterpret_cast(&p), N * sizeof(float)));\n checkCudaErrors(cudaMallocManaged(reinterpret_cast(&Ax), N * sizeof(float)));\n\n cudaDeviceSynchronize();\n\n checkCudaErrors(cudaEventCreate(&start));\n checkCudaErrors(cudaEventCreate(&stop));\n\n#if ENABLE_CPU_DEBUG_CODE\n float *Ax_cpu = reinterpret_cast(malloc(sizeof(float) * N));\n float *r_cpu = reinterpret_cast(malloc(sizeof(float) * N));\n float *p_cpu = reinterpret_cast(malloc(sizeof(float) * N));\n float *x_cpu = reinterpret_cast(malloc(sizeof(float) * N));\n\n for (int i = 0; i < N; i++) {\n r_cpu[i] = 1.0;\n Ax_cpu[i] = x_cpu[i] = 0.0;\n }\n\n#endif\n\n for (int i = 0; i < N; i++) {\n r[i] = rhs[i] = 1.0;\n x[i] = 0.0;\n }\n\n void *kernelArgs[] = {\n (void *)&I,\n (void *)&J,\n (void *)&val,\n (void *)&x,\n (void *)&Ax,\n (void *)&p,\n (void *)&r,\n (void *)&dot_result,\n (void *)&nz,\n (void *)&N,\n (void *)&tol,\n };\n\n int sMemSize = sizeof(double) * ((THREADS_PER_BLOCK / 32) + 1);\n int numBlocksPerSm = 0;\n int numThreads = THREADS_PER_BLOCK;\n\n checkCudaErrors(\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&numBlocksPerSm, gpuConjugateGradient, numThreads, sMemSize));\n\n int numSms = deviceProp.multiProcessorCount;\n dim3 dimGrid(numSms * numBlocksPerSm, 1, 1), dimBlock(THREADS_PER_BLOCK, 1, 1);\n checkCudaErrors(cudaEventRecord(start, 0));\n checkCudaErrors(\n cudaLaunchCooperativeKernel((void *)gpuConjugateGradient, dimGrid, dimBlock, kernelArgs, sMemSize, NULL));\n checkCudaErrors(cudaEventRecord(stop, 0));\n checkCudaErrors(cudaDeviceSynchronize());\n\n float time;\n checkCudaErrors(cudaEventElapsedTime(&time, start, stop));\n\n r1 = *dot_result;\n\n printf(\"GPU Final, residual = %e, kernel execution time = %f ms\\n\", sqrt(r1), time);\n\n#if ENABLE_CPU_DEBUG_CODE\n cpuConjugateGrad(I, J, val, x_cpu, Ax_cpu, p_cpu, r_cpu, nz, N, tol);\n#endif\n\n float rsum, diff, err = 0.0;\n\n for (int i = 0; i < N; i++) {\n rsum = 0.0;\n\n for (int j = I[i]; j < I[i + 1]; j++) {\n rsum += val[j] * x[J[j]];\n }\n\n diff = fabs(rsum - rhs[i]);\n\n if (diff > err) {\n err = diff;\n }\n }\n\n checkCudaErrors(cudaFree(I));\n checkCudaErrors(cudaFree(J));\n checkCudaErrors(cudaFree(val));\n checkCudaErrors(cudaFree(x));\n checkCudaErrors(cudaFree(rhs));\n checkCudaErrors(cudaFree(r));\n checkCudaErrors(cudaFree(p));\n checkCudaErrors(cudaFree(Ax));\n checkCudaErrors(cudaFree(dot_result));\n checkCudaErrors(cudaEventDestroy(start));\n checkCudaErrors(cudaEventDestroy(stop));\n\n#if ENABLE_CPU_DEBUG_CODE\n free(Ax_cpu);\n free(r_cpu);\n free(p_cpu);\n free(x_cpu);\n#endif\n\n printf(\"Test Summary: Error amount = %f \\n\", err);\n fprintf(stdout, \"&&&& conjugateGradientMultiBlockCG %s\\n\", (sqrt(r1) < tol) ? \"PASSED\" : \"FAILED\");\n exit((sqrt(r1) < tol) ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample implements a conjugate gradient solver on GPU using\n * Multi Block Cooperative Groups, also uses Unified Memory.\n *\n */\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// Utilities and system includes\n#include \n#include \n#include // helper function CUDA error checking and initialization\n#include // helper for shared functions common to CUDA Samples\n\nnamespace cg = cooperative_groups;\n\nconst char *sSDKname = \"conjugateGradientMultiBlockCG\";\n\n#define ENABLE_CPU_DEBUG_CODE 0\n#define THREADS_PER_BLOCK 512\n\n/* genTridiag: generate a random tridiagonal symmetric matrix */\nvoid genTridiag(int *I, int *J, float *val, int N, int nz)\n{\n I[0] = 0, J[0] = 0, J[1] = 1;\n val[0] = static_cast(rand()) / RAND_MAX + 10.0f;\n val[1] = static_cast(rand()) / RAND_MAX;\n int start;\n\n for (int i = 1; i < N; i++) {\n if (i > 1) {\n I[i] = I[i - 1] + 3;\n }\n else {\n I[1] = 2;\n }\n\n start = (i - 1) * 3 + 2;\n J[start] = i - 1;\n J[start + 1] = i;\n\n if (i < N - 1) {\n J[start + 2] = i + 1;\n }\n\n val[start] = val[start - 1];\n val[start + 1] = static_cast(rand()) / RAND_MAX + 10.0f;\n\n if (i < N - 1) {\n val[start + 2] = static_cast(rand()) / RAND_MAX;\n }\n }\n\n I[N] = nz;\n}\n\n// I - contains location of the given non-zero element in the row of the matrix\n// J - contains location of the given non-zero element in the column of the\n// matrix val - contains values of the given non-zero elements of the matrix\n// inputVecX - input vector to be multiplied\n// outputVecY - resultant vector\nvoid cpuSpMV(int *I, int *J, float *val, int nnz, int num_rows, float alpha, float *inputVecX, float *outputVecY)\n{\n for (int i = 0; i < num_rows; i++) {\n int num_elems_this_row = I[i + 1] - I[i];\n\n float output = 0.0;\n for (int j = 0; j < num_elems_this_row; j++) {\n output += alpha * val[I[i] + j] * inputVecX[J[I[i] + j]];\n }\n outputVecY[i] = output;\n }\n\n return;\n}\n\ndouble dotProduct(float *vecA, float *vecB, int size)\n{\n double result = 0.0;\n\n for (int i = 0; i < size; i++) {\n result = result + (vecA[i] * vecB[i]);\n }\n\n return result;\n}\n\nvoid scaleVector(float *vec, float alpha, int size)\n{\n for (int i = 0; i < size; i++) {\n vec[i] = alpha * vec[i];\n }\n}\n\nvoid saxpy(float *x, float *y, float a, int size)\n{\n for (int i = 0; i < size; i++) {\n y[i] = a * x[i] + y[i];\n }\n}\n\nvoid cpuConjugateGrad(int *I, int *J, float *val, float *x, float *Ax, float *p, float *r, int nnz, int N, float tol)\n{\n int max_iter = 10000;\n\n float alpha = 1.0;\n float alpham1 = -1.0;\n float r0 = 0.0, b, a, na;\n\n cpuSpMV(I, J, val, nnz, N, alpha, x, Ax);\n saxpy(Ax, r, alpham1, N);\n\n float r1 = dotProduct(r, r, N);\n\n int k = 1;\n\n while (r1 > tol * tol && k <= max_iter) {\n if (k > 1) {\n b = r1 / r0;\n scaleVector(p, b, N);\n\n saxpy(r, p, alpha, N);\n }\n else {\n for (int i = 0; i < N; i++)\n p[i] = r[i];\n }\n\n cpuSpMV(I, J, val, nnz, N, alpha, p, Ax);\n\n float dot = dotProduct(p, Ax, N);\n a = r1 / dot;\n\n saxpy(p, x, a, N);\n na = -a;\n saxpy(Ax, r, na, N);\n\n r0 = r1;\n r1 = dotProduct(r, r, N);\n\n printf(\"\\nCPU code iteration = %3d, residual = %e\\n\", k, sqrt(r1));\n k++;\n }\n}\n\n__device__ void gpuSpMV(int *I,\n int *J,\n float *val,\n int nnz,\n int num_rows,\n float alpha,\n float *inputVecX,\n float *outputVecY,\n cg::thread_block &cta,\n const cg::grid_group &grid)\n{\n for (int i = grid.thread_rank(); i < num_rows; i += grid.size()) {\n int row_elem = I[i];\n int next_row_elem = I[i + 1];\n int num_elems_this_row = next_row_elem - row_elem;\n\n float output = 0.0;\n for (int j = 0; j < num_elems_this_row; j++) {\n // I or J or val arrays - can be put in shared memory\n // as the access is random and reused in next calls of gpuSpMV function.\n output += alpha * val[row_elem + j] * inputVecX[J[row_elem + j]];\n }\n\n outputVecY[i] = output;\n }\n}\n\n__device__ void gpuSaxpy(float *x, float *y, float a, int size, const cg::grid_group &grid)\n{\n for (int i = grid.thread_rank(); i < size; i += grid.size()) {\n y[i] = a * x[i] + y[i];\n }\n}\n\n__device__ void gpuDotProduct(float *vecA,\n float *vecB,\n double *result,\n int size,\n const cg::thread_block &cta,\n const cg::grid_group &grid)\n{\n extern __shared__ double tmp[];\n\n double temp_sum = 0.0;\n for (int i = grid.thread_rank(); i < size; i += grid.size()) {\n temp_sum += static_cast(vecA[i] * vecB[i]);\n }\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n temp_sum = cg::reduce(tile32, temp_sum, cg::plus());\n\n if (tile32.thread_rank() == 0) {\n tmp[tile32.meta_group_rank()] = temp_sum;\n }\n\n cg::sync(cta);\n\n if (tile32.meta_group_rank() == 0) {\n temp_sum = tile32.thread_rank() < tile32.meta_group_size() ? tmp[tile32.thread_rank()] : 0.0;\n temp_sum = cg::reduce(tile32, temp_sum, cg::plus());\n\n if (tile32.thread_rank() == 0) {\n atomicAdd(result, temp_sum);\n }\n }\n}\n\n__device__ void gpuCopyVector(float *srcA, float *destB, int size, const cg::grid_group &grid)\n{\n for (int i = grid.thread_rank(); i < size; i += grid.size()) {\n destB[i] = srcA[i];\n }\n}\n\n__device__ void\ngpuScaleVectorAndSaxpy(const float *x, float *y, float a, float scale, int size, const cg::grid_group &grid)\n{\n for (int i = grid.thread_rank(); i < size; i += grid.size()) {\n y[i] = a * x[i] + scale * y[i];\n }\n}\n\nextern \"C\" __global__ void gpuConjugateGradient(int *I,\n int *J,\n float *val,\n float *x,\n float *Ax,\n float *p,\n float *r,\n double *dot_result,\n int nnz,\n int N,\n float tol)\n{\n cg::thread_block cta = cg::this_thread_block();\n cg::grid_group grid = cg::this_grid();\n\n int max_iter = 10000;\n\n float alpha = 1.0;\n float alpham1 = -1.0;\n float r0 = 0.0, r1, b, a, na;\n\n gpuSpMV(I, J, val, nnz, N, alpha, x, Ax, cta, grid);\n\n cg::sync(grid);\n\n gpuSaxpy(Ax, r, alpham1, N, grid);\n\n cg::sync(grid);\n\n gpuDotProduct(r, r, dot_result, N, cta, grid);\n\n cg::sync(grid);\n\n r1 = *dot_result;\n\n int k = 1;\n while (r1 > tol * tol && k <= max_iter) {\n if (k > 1) {\n b = r1 / r0;\n gpuScaleVectorAndSaxpy(r, p, alpha, b, N, grid);\n }\n else {\n gpuCopyVector(r, p, N, grid);\n }\n\n cg::sync(grid);\n\n gpuSpMV(I, J, val, nnz, N, alpha, p, Ax, cta, grid);\n\n if (threadIdx.x == 0 && blockIdx.x == 0)\n *dot_result = 0.0;\n\n cg::sync(grid);\n\n gpuDotProduct(p, Ax, dot_result, N, cta, grid);\n\n cg::sync(grid);\n\n a = r1 / *dot_result;\n\n gpuSaxpy(p, x, a, N, grid);\n na = -a;\n gpuSaxpy(Ax, r, na, N, grid);\n\n r0 = r1;\n\n cg::sync(grid);\n if (threadIdx.x == 0 && blockIdx.x == 0)\n *dot_result = 0.0;\n\n cg::sync(grid);\n\n gpuDotProduct(r, r, dot_result, N, cta, grid);\n\n cg::sync(grid);\n\n r1 = *dot_result;\n k++;\n }\n}\n\nbool areAlmostEqual(float a, float b, float maxRelDiff)\n{\n float diff = fabsf(a - b);\n float abs_a = fabsf(a);\n float abs_b = fabsf(b);\n float largest = abs_a > abs_b ? abs_a : abs_b;\n\n if (diff <= largest * maxRelDiff) {\n return true;\n }\n else {\n printf(\"maxRelDiff = %.8e\\n\", maxRelDiff);\n printf(\"diff %.8e > largest * maxRelDiff %.8e therefore %.8e and %.8e are not \"\n \"same\\n\",\n diff,\n largest * maxRelDiff,\n a,\n b);\n return false;\n }\n}\n\nint main(int argc, char **argv)\n{\n int N = 0, nz = 0, *I = NULL, *J = NULL;\n float *val = NULL;\n const float tol = 1e-5f;\n float *x;\n float *rhs;\n float r1;\n float *r, *p, *Ax;\n hipEvent_t start, stop;\n\n printf(\"Starting [%s]...\\n\", sSDKname);\n\n // This will pick the best possible CUDA capable device\n hipDeviceProp_t deviceProp;\n int devID = findCudaDevice(argc, (const char **)argv);\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, devID));\n\n if (!deviceProp.managedMemory) {\n // This sample requires being run on a device that supports Unified Memory\n fprintf(stderr, \"Unified Memory not supported on this device\\n\");\n exit(EXIT_WAIVED);\n }\n\n // This sample requires being run on a device that supports Cooperative Kernel\n // Launch\n if (!deviceProp.cooperativeLaunch) {\n printf(\"\\nSelected GPU (%d) does not support Cooperative Kernel Launch, \"\n \"Waiving the run\\n\",\n devID);\n exit(EXIT_WAIVED);\n }\n\n // Statistics about the GPU device\n printf(\"> GPU device has %d Multi-Processors, SM %d.%d compute capabilities\\n\\n\",\n deviceProp.multiProcessorCount,\n deviceProp.major,\n deviceProp.minor);\n\n /* Generate a random tridiagonal symmetric matrix in CSR format */\n N = 1048576;\n nz = (N - 2) * 3 + 4;\n\n hipMallocManaged(reinterpret_cast(&I), sizeof(int) * (N + 1));\n hipMallocManaged(reinterpret_cast(&J), sizeof(int) * nz);\n hipMallocManaged(reinterpret_cast(&val), sizeof(float) * nz);\n\n genTridiag(I, J, val, N, nz);\n\n hipMallocManaged(reinterpret_cast(&x), sizeof(float) * N);\n hipMallocManaged(reinterpret_cast(&rhs), sizeof(float) * N);\n\n double *dot_result;\n\n hipMallocManaged(reinterpret_cast(&dot_result), sizeof(double));\n\n *dot_result = 0.0;\n\n // temp memory for CG\n checkCudaErrors(hipMallocManaged(reinterpret_cast(&r), N * sizeof(float)));\n checkCudaErrors(hipMallocManaged(reinterpret_cast(&p), N * sizeof(float)));\n checkCudaErrors(hipMallocManaged(reinterpret_cast(&Ax), N * sizeof(float)));\n\n hipDeviceSynchronize();\n\n checkCudaErrors(hipEventCreate(&start));\n checkCudaErrors(hipEventCreate(&stop));\n\n#if ENABLE_CPU_DEBUG_CODE\n float *Ax_cpu = reinterpret_cast(malloc(sizeof(float) * N));\n float *r_cpu = reinterpret_cast(malloc(sizeof(float) * N));\n float *p_cpu = reinterpret_cast(malloc(sizeof(float) * N));\n float *x_cpu = reinterpret_cast(malloc(sizeof(float) * N));\n\n for (int i = 0; i < N; i++) {\n r_cpu[i] = 1.0;\n Ax_cpu[i] = x_cpu[i] = 0.0;\n }\n\n#endif\n\n for (int i = 0; i < N; i++) {\n r[i] = rhs[i] = 1.0;\n x[i] = 0.0;\n }\n\n void *kernelArgs[] = {\n (void *)&I,\n (void *)&J,\n (void *)&val,\n (void *)&x,\n (void *)&Ax,\n (void *)&p,\n (void *)&r,\n (void *)&dot_result,\n (void *)&nz,\n (void *)&N,\n (void *)&tol,\n };\n\n int sMemSize = sizeof(double) * ((THREADS_PER_BLOCK / 32) + 1);\n int numBlocksPerSm = 0;\n int numThreads = THREADS_PER_BLOCK;\n\n checkCudaErrors(\n hipOccupancyMaxActiveBlocksPerMultiprocessor(&numBlocksPerSm, gpuConjugateGradient, numThreads, sMemSize));\n\n int numSms = deviceProp.multiProcessorCount;\n dim3 dimGrid(numSms * numBlocksPerSm, 1, 1), dimBlock(THREADS_PER_BLOCK, 1, 1);\n checkCudaErrors(hipEventRecord(start, 0));\n checkCudaErrors(\n hipLaunchCooperativeKernel(reinterpret_cast((void *)gpuConjugateGradient), dimGrid, dimBlock, kernelArgs, sMemSize, NULL));\n checkCudaErrors(hipEventRecord(stop, 0));\n checkCudaErrors(hipDeviceSynchronize());\n\n float time;\n checkCudaErrors(hipEventElapsedTime(&time, start, stop));\n\n r1 = *dot_result;\n\n printf(\"GPU Final, residual = %e, kernel execution time = %f ms\\n\", sqrt(r1), time);\n\n#if ENABLE_CPU_DEBUG_CODE\n cpuConjugateGrad(I, J, val, x_cpu, Ax_cpu, p_cpu, r_cpu, nz, N, tol);\n#endif\n\n float rsum, diff, err = 0.0;\n\n for (int i = 0; i < N; i++) {\n rsum = 0.0;\n\n for (int j = I[i]; j < I[i + 1]; j++) {\n rsum += val[j] * x[J[j]];\n }\n\n diff = fabs(rsum - rhs[i]);\n\n if (diff > err) {\n err = diff;\n }\n }\n\n checkCudaErrors(hipFree(I));\n checkCudaErrors(hipFree(J));\n checkCudaErrors(hipFree(val));\n checkCudaErrors(hipFree(x));\n checkCudaErrors(hipFree(rhs));\n checkCudaErrors(hipFree(r));\n checkCudaErrors(hipFree(p));\n checkCudaErrors(hipFree(Ax));\n checkCudaErrors(hipFree(dot_result));\n checkCudaErrors(hipEventDestroy(start));\n checkCudaErrors(hipEventDestroy(stop));\n\n#if ENABLE_CPU_DEBUG_CODE\n free(Ax_cpu);\n free(r_cpu);\n free(p_cpu);\n free(x_cpu);\n#endif\n\n printf(\"Test Summary: Error amount = %f \\n\", err);\n fprintf(stdout, \"&&&& conjugateGradientMultiBlockCG %s\\n\", (sqrt(r1) < tol) ? \"PASSED\" : \"FAILED\");\n exit((sqrt(r1) < tol) ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/4_CUDA_Libraries/conjugateGradientMultiBlockCG/conjugateGradientMultiBlockCG.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample implements a conjugate gradient solver on GPU\n * using CUBLAS and CUSPARSE with CUDA Graphs\n *\n */\n\n// includes, system\n#include \n#include \n#include \n\n/* Using updated (v2) interfaces to cublas */\n#include \n#include \n#include \n\n// Utilities and system includes\n#include // helper function CUDA error checking and initialization\n#include // helper for shared functions common to CUDA Samples\n\nconst char *sSDKname = \"conjugateGradientCudaGraphs\";\n\n#ifndef WITH_GRAPH\n#define WITH_GRAPH 1\n#endif\n\n/* genTridiag: generate a random tridiagonal symmetric matrix */\nvoid genTridiag(int *I, int *J, float *val, int N, int nz)\n{\n I[0] = 0, J[0] = 0, J[1] = 1;\n val[0] = (float)rand() / RAND_MAX + 10.0f;\n val[1] = (float)rand() / RAND_MAX;\n int start;\n\n for (int i = 1; i < N; i++) {\n if (i > 1) {\n I[i] = I[i - 1] + 3;\n }\n else {\n I[1] = 2;\n }\n\n start = (i - 1) * 3 + 2;\n J[start] = i - 1;\n J[start + 1] = i;\n\n if (i < N - 1) {\n J[start + 2] = i + 1;\n }\n\n val[start] = val[start - 1];\n val[start + 1] = (float)rand() / RAND_MAX + 10.0f;\n\n if (i < N - 1) {\n val[start + 2] = (float)rand() / RAND_MAX;\n }\n }\n\n I[N] = nz;\n}\n\n__global__ void initVectors(float *rhs, float *x, int N)\n{\n size_t gid = blockIdx.x * blockDim.x + threadIdx.x;\n\n for (size_t i = gid; i < N; i += gridDim.x * blockDim.x) {\n rhs[i] = 1.0;\n x[i] = 0.0;\n }\n}\n\n__global__ void r1_div_x(float *r1, float *r0, float *b)\n{\n int gid = blockIdx.x * blockDim.x + threadIdx.x;\n if (gid == 0) {\n b[0] = r1[0] / r0[0];\n }\n}\n\n__global__ void a_minus(float *a, float *na)\n{\n int gid = blockIdx.x * blockDim.x + threadIdx.x;\n if (gid == 0) {\n na[0] = -(a[0]);\n }\n}\n\nint main(int argc, char **argv)\n{\n int N = 0, nz = 0, *I = NULL, *J = NULL;\n float *val = NULL;\n const float tol = 1e-5f;\n const int max_iter = 10000;\n float *x;\n float *rhs;\n float r1;\n\n int *d_col, *d_row;\n float *d_val, *d_x;\n float *d_r, *d_p, *d_Ax;\n int k;\n float alpha, beta, alpham1;\n\n cudaStream_t stream1, streamForGraph;\n\n // This will pick the best possible CUDA capable device\n cudaDeviceProp deviceProp;\n int devID = findCudaDevice(argc, (const char **)argv);\n\n if (devID < 0) {\n printf(\"exiting...\\n\");\n exit(EXIT_SUCCESS);\n }\n\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, devID));\n\n // Statistics about the GPU device\n printf(\"> GPU device has %d Multi-Processors, SM %d.%d compute capabilities\\n\\n\",\n deviceProp.multiProcessorCount,\n deviceProp.major,\n deviceProp.minor);\n\n /* Generate a random tridiagonal symmetric matrix in CSR format */\n N = 1048576;\n nz = (N - 2) * 3 + 4;\n checkCudaErrors(cudaMallocHost(&I, sizeof(int) * (N + 1)));\n checkCudaErrors(cudaMallocHost(&J, sizeof(int) * nz));\n checkCudaErrors(cudaMallocHost(&val, sizeof(float) * nz));\n genTridiag(I, J, val, N, nz);\n\n checkCudaErrors(cudaMallocHost(&x, sizeof(float) * N));\n rhs = (float *)malloc(sizeof(float) * N);\n\n for (int i = 0; i < N; i++) {\n rhs[i] = 1.0;\n x[i] = 0.0;\n }\n\n /* Get handle to the CUBLAS context */\n cublasHandle_t cublasHandle = 0;\n cublasStatus_t cublasStatus;\n cublasStatus = cublasCreate(&cublasHandle);\n\n checkCudaErrors(cublasStatus);\n\n /* Get handle to the CUSPARSE context */\n cusparseHandle_t cusparseHandle = 0;\n cusparseStatus_t cusparseStatus;\n cusparseStatus = cusparseCreate(&cusparseHandle);\n\n checkCudaErrors(cusparseStatus);\n\n checkCudaErrors(cudaStreamCreate(&stream1));\n\n checkCudaErrors(cudaMalloc((void **)&d_col, nz * sizeof(int)));\n checkCudaErrors(cudaMalloc((void **)&d_row, (N + 1) * sizeof(int)));\n checkCudaErrors(cudaMalloc((void **)&d_val, nz * sizeof(float)));\n checkCudaErrors(cudaMalloc((void **)&d_x, N * sizeof(float)));\n checkCudaErrors(cudaMalloc((void **)&d_r, N * sizeof(float)));\n checkCudaErrors(cudaMalloc((void **)&d_p, N * sizeof(float)));\n checkCudaErrors(cudaMalloc((void **)&d_Ax, N * sizeof(float)));\n\n float *d_r1, *d_r0, *d_dot, *d_a, *d_na, *d_b;\n checkCudaErrors(cudaMalloc((void **)&d_r1, sizeof(float)));\n checkCudaErrors(cudaMalloc((void **)&d_r0, sizeof(float)));\n checkCudaErrors(cudaMalloc((void **)&d_dot, sizeof(float)));\n checkCudaErrors(cudaMalloc((void **)&d_a, sizeof(float)));\n checkCudaErrors(cudaMalloc((void **)&d_na, sizeof(float)));\n checkCudaErrors(cudaMalloc((void **)&d_b, sizeof(float)));\n\n /* Wrap raw data into cuSPARSE generic API objects */\n cusparseSpMatDescr_t matA = NULL;\n checkCudaErrors(cusparseCreateCsr(&matA,\n N,\n N,\n nz,\n d_row,\n d_col,\n d_val,\n CUSPARSE_INDEX_32I,\n CUSPARSE_INDEX_32I,\n CUSPARSE_INDEX_BASE_ZERO,\n CUDA_R_32F));\n cusparseDnVecDescr_t vecx = NULL;\n checkCudaErrors(cusparseCreateDnVec(&vecx, N, d_x, CUDA_R_32F));\n cusparseDnVecDescr_t vecp = NULL;\n checkCudaErrors(cusparseCreateDnVec(&vecp, N, d_p, CUDA_R_32F));\n cusparseDnVecDescr_t vecAx = NULL;\n checkCudaErrors(cusparseCreateDnVec(&vecAx, N, d_Ax, CUDA_R_32F));\n\n /* Allocate workspace for cuSPARSE */\n size_t bufferSize = 0;\n checkCudaErrors(cusparseSpMV_bufferSize(cusparseHandle,\n CUSPARSE_OPERATION_NON_TRANSPOSE,\n &alpha,\n matA,\n vecx,\n &beta,\n vecAx,\n CUDA_R_32F,\n CUSPARSE_SPMV_ALG_DEFAULT,\n &bufferSize));\n void *buffer = NULL;\n checkCudaErrors(cudaMalloc(&buffer, bufferSize));\n\n cusparseMatDescr_t descr = 0;\n checkCudaErrors(cusparseCreateMatDescr(&descr));\n\n checkCudaErrors(cusparseSetMatType(descr, CUSPARSE_MATRIX_TYPE_GENERAL));\n checkCudaErrors(cusparseSetMatIndexBase(descr, CUSPARSE_INDEX_BASE_ZERO));\n\n int numBlocks = 0, blockSize = 0;\n checkCudaErrors(cudaOccupancyMaxPotentialBlockSize(&numBlocks, &blockSize, initVectors));\n\n checkCudaErrors(cudaMemcpyAsync(d_col, J, nz * sizeof(int), cudaMemcpyHostToDevice, stream1));\n checkCudaErrors(cudaMemcpyAsync(d_row, I, (N + 1) * sizeof(int), cudaMemcpyHostToDevice, stream1));\n checkCudaErrors(cudaMemcpyAsync(d_val, val, nz * sizeof(float), cudaMemcpyHostToDevice, stream1));\n\n initVectors<<>>(d_r, d_x, N);\n\n alpha = 1.0;\n alpham1 = -1.0;\n beta = 0.0;\n\n checkCudaErrors(cusparseSetStream(cusparseHandle, stream1));\n checkCudaErrors(cusparseSpMV(cusparseHandle,\n CUSPARSE_OPERATION_NON_TRANSPOSE,\n &alpha,\n matA,\n vecx,\n &beta,\n vecAx,\n CUDA_R_32F,\n CUSPARSE_SPMV_ALG_DEFAULT,\n buffer));\n\n checkCudaErrors(cublasSetStream(cublasHandle, stream1));\n checkCudaErrors(cublasSaxpy(cublasHandle, N, &alpham1, d_Ax, 1, d_r, 1));\n\n checkCudaErrors(cublasSetPointerMode(cublasHandle, CUBLAS_POINTER_MODE_DEVICE));\n checkCudaErrors(cublasSdot(cublasHandle, N, d_r, 1, d_r, 1, d_r1));\n\n k = 1;\n // First Iteration when k=1 starts\n checkCudaErrors(cublasScopy(cublasHandle, N, d_r, 1, d_p, 1));\n checkCudaErrors(cusparseSpMV(cusparseHandle,\n CUSPARSE_OPERATION_NON_TRANSPOSE,\n &alpha,\n matA,\n vecp,\n &beta,\n vecAx,\n CUDA_R_32F,\n CUSPARSE_SPMV_ALG_DEFAULT,\n buffer));\n\n checkCudaErrors(cublasSdot(cublasHandle, N, d_p, 1, d_Ax, 1, d_dot));\n\n r1_div_x<<<1, 1, 0, stream1>>>(d_r1, d_dot, d_a);\n\n checkCudaErrors(cublasSaxpy(cublasHandle, N, d_a, d_p, 1, d_x, 1));\n\n a_minus<<<1, 1, 0, stream1>>>(d_a, d_na);\n\n checkCudaErrors(cublasSaxpy(cublasHandle, N, d_na, d_Ax, 1, d_r, 1));\n\n checkCudaErrors(cudaMemcpyAsync(d_r0, d_r1, sizeof(float), cudaMemcpyDeviceToDevice, stream1));\n\n checkCudaErrors(cublasSdot(cublasHandle, N, d_r, 1, d_r, 1, d_r1));\n\n checkCudaErrors(cudaMemcpyAsync(&r1, d_r1, sizeof(float), cudaMemcpyDeviceToHost, stream1));\n checkCudaErrors(cudaStreamSynchronize(stream1));\n printf(\"iteration = %3d, residual = %e\\n\", k, sqrt(r1));\n // First Iteration when k=1 ends\n k++;\n\n#if WITH_GRAPH\n cudaGraph_t initGraph;\n checkCudaErrors(cudaStreamCreate(&streamForGraph));\n checkCudaErrors(cublasSetStream(cublasHandle, stream1));\n checkCudaErrors(cusparseSetStream(cusparseHandle, stream1));\n checkCudaErrors(cudaStreamBeginCapture(stream1, cudaStreamCaptureModeGlobal));\n\n r1_div_x<<<1, 1, 0, stream1>>>(d_r1, d_r0, d_b);\n cublasSetPointerMode(cublasHandle, CUBLAS_POINTER_MODE_DEVICE);\n checkCudaErrors(cublasSscal(cublasHandle, N, d_b, d_p, 1));\n cublasSetPointerMode(cublasHandle, CUBLAS_POINTER_MODE_HOST);\n checkCudaErrors(cublasSaxpy(cublasHandle, N, &alpha, d_r, 1, d_p, 1));\n cublasSetPointerMode(cublasHandle, CUBLAS_POINTER_MODE_DEVICE);\n\n checkCudaErrors(cusparseSetPointerMode(cusparseHandle, CUSPARSE_POINTER_MODE_HOST));\n checkCudaErrors(cusparseSpMV(cusparseHandle,\n CUSPARSE_OPERATION_NON_TRANSPOSE,\n &alpha,\n matA,\n vecp,\n &beta,\n vecAx,\n CUDA_R_32F,\n CUSPARSE_SPMV_ALG_DEFAULT,\n buffer));\n\n checkCudaErrors(cudaMemsetAsync(d_dot, 0, sizeof(float), stream1));\n checkCudaErrors(cublasSdot(cublasHandle, N, d_p, 1, d_Ax, 1, d_dot));\n\n r1_div_x<<<1, 1, 0, stream1>>>(d_r1, d_dot, d_a);\n\n checkCudaErrors(cublasSaxpy(cublasHandle, N, d_a, d_p, 1, d_x, 1));\n\n a_minus<<<1, 1, 0, stream1>>>(d_a, d_na);\n\n checkCudaErrors(cublasSaxpy(cublasHandle, N, d_na, d_Ax, 1, d_r, 1));\n\n checkCudaErrors(cudaMemcpyAsync(d_r0, d_r1, sizeof(float), cudaMemcpyDeviceToDevice, stream1));\n checkCudaErrors(cudaMemsetAsync(d_r1, 0, sizeof(float), stream1));\n\n checkCudaErrors(cublasSdot(cublasHandle, N, d_r, 1, d_r, 1, d_r1));\n\n checkCudaErrors(cudaMemcpyAsync((float *)&r1, d_r1, sizeof(float), cudaMemcpyDeviceToHost, stream1));\n\n checkCudaErrors(cudaStreamEndCapture(stream1, &initGraph));\n cudaGraphExec_t graphExec;\n checkCudaErrors(cudaGraphInstantiate(&graphExec, initGraph, NULL, NULL, 0));\n#endif\n\n checkCudaErrors(cublasSetStream(cublasHandle, stream1));\n checkCudaErrors(cusparseSetStream(cusparseHandle, stream1));\n\n while (r1 > tol * tol && k <= max_iter) {\n#if WITH_GRAPH\n checkCudaErrors(cudaGraphLaunch(graphExec, streamForGraph));\n checkCudaErrors(cudaStreamSynchronize(streamForGraph));\n#else\n r1_div_x<<<1, 1, 0, stream1>>>(d_r1, d_r0, d_b);\n cublasSetPointerMode(cublasHandle, CUBLAS_POINTER_MODE_DEVICE);\n checkCudaErrors(cublasSscal(cublasHandle, N, d_b, d_p, 1));\n\n cublasSetPointerMode(cublasHandle, CUBLAS_POINTER_MODE_HOST);\n checkCudaErrors(cublasSaxpy(cublasHandle, N, &alpha, d_r, 1, d_p, 1));\n\n checkCudaErrors(cusparseSpMV(cusparseHandle,\n CUSPARSE_OPERATION_NON_TRANSPOSE,\n &alpha,\n matA,\n vecp,\n &beta,\n vecAx,\n CUDA_R_32F,\n CUSPARSE_SPMV_ALG_DEFAULT,\n buffer));\n\n cublasSetPointerMode(cublasHandle, CUBLAS_POINTER_MODE_DEVICE);\n checkCudaErrors(cublasSdot(cublasHandle, N, d_p, 1, d_Ax, 1, d_dot));\n\n r1_div_x<<<1, 1, 0, stream1>>>(d_r1, d_dot, d_a);\n\n checkCudaErrors(cublasSaxpy(cublasHandle, N, d_a, d_p, 1, d_x, 1));\n\n a_minus<<<1, 1, 0, stream1>>>(d_a, d_na);\n checkCudaErrors(cublasSaxpy(cublasHandle, N, d_na, d_Ax, 1, d_r, 1));\n\n checkCudaErrors(cudaMemcpyAsync(d_r0, d_r1, sizeof(float), cudaMemcpyDeviceToDevice, stream1));\n\n checkCudaErrors(cublasSdot(cublasHandle, N, d_r, 1, d_r, 1, d_r1));\n checkCudaErrors(cudaMemcpyAsync((float *)&r1, d_r1, sizeof(float), cudaMemcpyDeviceToHost, stream1));\n checkCudaErrors(cudaStreamSynchronize(stream1));\n#endif\n printf(\"iteration = %3d, residual = %e\\n\", k, sqrt(r1));\n k++;\n }\n\n#if WITH_GRAPH\n checkCudaErrors(cudaMemcpyAsync(x, d_x, N * sizeof(float), cudaMemcpyDeviceToHost, streamForGraph));\n checkCudaErrors(cudaStreamSynchronize(streamForGraph));\n#else\n checkCudaErrors(cudaMemcpyAsync(x, d_x, N * sizeof(float), cudaMemcpyDeviceToHost, stream1));\n checkCudaErrors(cudaStreamSynchronize(stream1));\n#endif\n\n float rsum, diff, err = 0.0;\n\n for (int i = 0; i < N; i++) {\n rsum = 0.0;\n\n for (int j = I[i]; j < I[i + 1]; j++) {\n rsum += val[j] * x[J[j]];\n }\n\n diff = fabs(rsum - rhs[i]);\n\n if (diff > err) {\n err = diff;\n }\n }\n\n#if WITH_GRAPH\n checkCudaErrors(cudaGraphExecDestroy(graphExec));\n checkCudaErrors(cudaGraphDestroy(initGraph));\n checkCudaErrors(cudaStreamDestroy(streamForGraph));\n#endif\n checkCudaErrors(cudaStreamDestroy(stream1));\n cusparseDestroy(cusparseHandle);\n cublasDestroy(cublasHandle);\n\n if (matA) {\n checkCudaErrors(cusparseDestroySpMat(matA));\n }\n if (vecx) {\n checkCudaErrors(cusparseDestroyDnVec(vecx));\n }\n if (vecAx) {\n checkCudaErrors(cusparseDestroyDnVec(vecAx));\n }\n if (vecp) {\n checkCudaErrors(cusparseDestroyDnVec(vecp));\n }\n\n checkCudaErrors(cudaFreeHost(I));\n checkCudaErrors(cudaFreeHost(J));\n checkCudaErrors(cudaFreeHost(val));\n checkCudaErrors(cudaFreeHost(x));\n free(rhs);\n checkCudaErrors(cudaFree(d_col));\n checkCudaErrors(cudaFree(d_row));\n checkCudaErrors(cudaFree(d_val));\n checkCudaErrors(cudaFree(d_x));\n checkCudaErrors(cudaFree(d_r));\n checkCudaErrors(cudaFree(d_p));\n checkCudaErrors(cudaFree(d_Ax));\n\n printf(\"Test Summary: Error amount = %f\\n\", err);\n exit((k <= max_iter) ? 0 : 1);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample implements a conjugate gradient solver on GPU\n * using CUBLAS and CUSPARSE with CUDA Graphs\n *\n */\n\n// includes, system\n#include \n#include \n#include \n\n/* Using updated (v2) interfaces to cublas */\n#include \n#include \n#include \n\n// Utilities and system includes\n#include // helper function CUDA error checking and initialization\n#include // helper for shared functions common to CUDA Samples\n\nconst char *sSDKname = \"conjugateGradientCudaGraphs\";\n\n#ifndef WITH_GRAPH\n#define WITH_GRAPH 1\n#endif\n\n/* genTridiag: generate a random tridiagonal symmetric matrix */\nvoid genTridiag(int *I, int *J, float *val, int N, int nz)\n{\n I[0] = 0, J[0] = 0, J[1] = 1;\n val[0] = (float)rand() / RAND_MAX + 10.0f;\n val[1] = (float)rand() / RAND_MAX;\n int start;\n\n for (int i = 1; i < N; i++) {\n if (i > 1) {\n I[i] = I[i - 1] + 3;\n }\n else {\n I[1] = 2;\n }\n\n start = (i - 1) * 3 + 2;\n J[start] = i - 1;\n J[start + 1] = i;\n\n if (i < N - 1) {\n J[start + 2] = i + 1;\n }\n\n val[start] = val[start - 1];\n val[start + 1] = (float)rand() / RAND_MAX + 10.0f;\n\n if (i < N - 1) {\n val[start + 2] = (float)rand() / RAND_MAX;\n }\n }\n\n I[N] = nz;\n}\n\n__global__ void initVectors(float *rhs, float *x, int N)\n{\n size_t gid = blockIdx.x * blockDim.x + threadIdx.x;\n\n for (size_t i = gid; i < N; i += gridDim.x * blockDim.x) {\n rhs[i] = 1.0;\n x[i] = 0.0;\n }\n}\n\n__global__ void r1_div_x(float *r1, float *r0, float *b)\n{\n int gid = blockIdx.x * blockDim.x + threadIdx.x;\n if (gid == 0) {\n b[0] = r1[0] / r0[0];\n }\n}\n\n__global__ void a_minus(float *a, float *na)\n{\n int gid = blockIdx.x * blockDim.x + threadIdx.x;\n if (gid == 0) {\n na[0] = -(a[0]);\n }\n}\n\nint main(int argc, char **argv)\n{\n int N = 0, nz = 0, *I = NULL, *J = NULL;\n float *val = NULL;\n const float tol = 1e-5f;\n const int max_iter = 10000;\n float *x;\n float *rhs;\n float r1;\n\n int *d_col, *d_row;\n float *d_val, *d_x;\n float *d_r, *d_p, *d_Ax;\n int k;\n float alpha, beta, alpham1;\n\n hipStream_t stream1, streamForGraph;\n\n // This will pick the best possible CUDA capable device\n hipDeviceProp_t deviceProp;\n int devID = findCudaDevice(argc, (const char **)argv);\n\n if (devID < 0) {\n printf(\"exiting...\\n\");\n exit(EXIT_SUCCESS);\n }\n\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, devID));\n\n // Statistics about the GPU device\n printf(\"> GPU device has %d Multi-Processors, SM %d.%d compute capabilities\\n\\n\",\n deviceProp.multiProcessorCount,\n deviceProp.major,\n deviceProp.minor);\n\n /* Generate a random tridiagonal symmetric matrix in CSR format */\n N = 1048576;\n nz = (N - 2) * 3 + 4;\n checkCudaErrors(hipHostMalloc(&I, sizeof(int) * (N + 1)));\n checkCudaErrors(hipHostMalloc(&J, sizeof(int) * nz));\n checkCudaErrors(hipHostMalloc(&val, sizeof(float) * nz));\n genTridiag(I, J, val, N, nz);\n\n checkCudaErrors(hipHostMalloc(&x, sizeof(float) * N));\n rhs = (float *)malloc(sizeof(float) * N);\n\n for (int i = 0; i < N; i++) {\n rhs[i] = 1.0;\n x[i] = 0.0;\n }\n\n /* Get handle to the CUBLAS context */\n hipblasHandle_t cublasHandle = 0;\n hipblasStatus_t hipblasStatus_t;\n hipblasStatus_t = hipblasCreate(&cublasHandle);\n\n checkCudaErrors(hipblasStatus_t);\n\n /* Get handle to the CUSPARSE context */\n hipsparseHandle_t cusparseHandle = 0;\n hipsparseStatus_t cusparseStatus;\n cusparseStatus = hipsparseCreate(&cusparseHandle);\n\n checkCudaErrors(cusparseStatus);\n\n checkCudaErrors(hipStreamCreate(&stream1));\n\n checkCudaErrors(hipMalloc((void **)&d_col, nz * sizeof(int)));\n checkCudaErrors(hipMalloc((void **)&d_row, (N + 1) * sizeof(int)));\n checkCudaErrors(hipMalloc((void **)&d_val, nz * sizeof(float)));\n checkCudaErrors(hipMalloc((void **)&d_x, N * sizeof(float)));\n checkCudaErrors(hipMalloc((void **)&d_r, N * sizeof(float)));\n checkCudaErrors(hipMalloc((void **)&d_p, N * sizeof(float)));\n checkCudaErrors(hipMalloc((void **)&d_Ax, N * sizeof(float)));\n\n float *d_r1, *d_r0, *d_dot, *d_a, *d_na, *d_b;\n checkCudaErrors(hipMalloc((void **)&d_r1, sizeof(float)));\n checkCudaErrors(hipMalloc((void **)&d_r0, sizeof(float)));\n checkCudaErrors(hipMalloc((void **)&d_dot, sizeof(float)));\n checkCudaErrors(hipMalloc((void **)&d_a, sizeof(float)));\n checkCudaErrors(hipMalloc((void **)&d_na, sizeof(float)));\n checkCudaErrors(hipMalloc((void **)&d_b, sizeof(float)));\n\n /* Wrap raw data into cuSPARSE generic API objects */\n hipsparseSpMatDescr_t matA = NULL;\n checkCudaErrors(hipsparseCreateCsr(&matA,\n N,\n N,\n nz,\n d_row,\n d_col,\n d_val,\n HIPSPARSE_INDEX_32I,\n HIPSPARSE_INDEX_32I,\n HIPSPARSE_INDEX_BASE_ZERO,\n HIP_R_32F));\n hipsparseDnVecDescr_t vecx = NULL;\n checkCudaErrors(hipsparseCreateDnVec(&vecx, N, d_x, HIP_R_32F));\n hipsparseDnVecDescr_t vecp = NULL;\n checkCudaErrors(hipsparseCreateDnVec(&vecp, N, d_p, HIP_R_32F));\n hipsparseDnVecDescr_t vecAx = NULL;\n checkCudaErrors(hipsparseCreateDnVec(&vecAx, N, d_Ax, HIP_R_32F));\n\n /* Allocate workspace for cuSPARSE */\n size_t bufferSize = 0;\n checkCudaErrors(hipsparseSpMV_bufferSize(cusparseHandle,\n HIPSPARSE_OPERATION_NON_TRANSPOSE,\n &alpha,\n matA,\n vecx,\n &beta,\n vecAx,\n HIP_R_32F,\n HIPSPARSE_SPMV_ALG_DEFAULT,\n &bufferSize));\n void *buffer = NULL;\n checkCudaErrors(hipMalloc(&buffer, bufferSize));\n\n hipsparseMatDescr_t descr = 0;\n checkCudaErrors(hipsparseCreateMatDescr(&descr));\n\n checkCudaErrors(hipsparseSetMatType(descr, HIPSPARSE_MATRIX_TYPE_GENERAL));\n checkCudaErrors(hipsparseSetMatIndexBase(descr, HIPSPARSE_INDEX_BASE_ZERO));\n\n int numBlocks = 0, blockSize = 0;\n checkCudaErrors(hipOccupancyMaxPotentialBlockSize(&numBlocks, &blockSize, initVectors));\n\n checkCudaErrors(hipMemcpyAsync(d_col, J, nz * sizeof(int), hipMemcpyHostToDevice, stream1));\n checkCudaErrors(hipMemcpyAsync(d_row, I, (N + 1) * sizeof(int), hipMemcpyHostToDevice, stream1));\n checkCudaErrors(hipMemcpyAsync(d_val, val, nz * sizeof(float), hipMemcpyHostToDevice, stream1));\n\n initVectors<<>>(d_r, d_x, N);\n\n alpha = 1.0;\n alpham1 = -1.0;\n beta = 0.0;\n\n checkCudaErrors(hipsparseSetStream(cusparseHandle, stream1));\n checkCudaErrors(hipsparseSpMV(cusparseHandle,\n HIPSPARSE_OPERATION_NON_TRANSPOSE,\n &alpha,\n matA,\n vecx,\n &beta,\n vecAx,\n HIP_R_32F,\n HIPSPARSE_SPMV_ALG_DEFAULT,\n buffer));\n\n checkCudaErrors(hipblasSetStream(cublasHandle, stream1));\n checkCudaErrors(hipblasSaxpy(cublasHandle, N, &alpham1, d_Ax, 1, d_r, 1));\n\n checkCudaErrors(hipblasSetPointerMode(cublasHandle, HIPBLAS_POINTER_MODE_DEVICE));\n checkCudaErrors(hipblasSdot(cublasHandle, N, d_r, 1, d_r, 1, d_r1));\n\n k = 1;\n // First Iteration when k=1 starts\n checkCudaErrors(hipblasScopy(cublasHandle, N, d_r, 1, d_p, 1));\n checkCudaErrors(hipsparseSpMV(cusparseHandle,\n HIPSPARSE_OPERATION_NON_TRANSPOSE,\n &alpha,\n matA,\n vecp,\n &beta,\n vecAx,\n HIP_R_32F,\n HIPSPARSE_SPMV_ALG_DEFAULT,\n buffer));\n\n checkCudaErrors(hipblasSdot(cublasHandle, N, d_p, 1, d_Ax, 1, d_dot));\n\n r1_div_x<<<1, 1, 0, stream1>>>(d_r1, d_dot, d_a);\n\n checkCudaErrors(hipblasSaxpy(cublasHandle, N, d_a, d_p, 1, d_x, 1));\n\n a_minus<<<1, 1, 0, stream1>>>(d_a, d_na);\n\n checkCudaErrors(hipblasSaxpy(cublasHandle, N, d_na, d_Ax, 1, d_r, 1));\n\n checkCudaErrors(hipMemcpyAsync(d_r0, d_r1, sizeof(float), hipMemcpyDeviceToDevice, stream1));\n\n checkCudaErrors(hipblasSdot(cublasHandle, N, d_r, 1, d_r, 1, d_r1));\n\n checkCudaErrors(hipMemcpyAsync(&r1, d_r1, sizeof(float), hipMemcpyDeviceToHost, stream1));\n checkCudaErrors(hipStreamSynchronize(stream1));\n printf(\"iteration = %3d, residual = %e\\n\", k, sqrt(r1));\n // First Iteration when k=1 ends\n k++;\n\n#if WITH_GRAPH\n hipGraph_t initGraph;\n checkCudaErrors(hipStreamCreate(&streamForGraph));\n checkCudaErrors(hipblasSetStream(cublasHandle, stream1));\n checkCudaErrors(hipsparseSetStream(cusparseHandle, stream1));\n checkCudaErrors(hipStreamBeginCapture(stream1, hipStreamCaptureModeGlobal));\n\n r1_div_x<<<1, 1, 0, stream1>>>(d_r1, d_r0, d_b);\n hipblasSetPointerMode(cublasHandle, HIPBLAS_POINTER_MODE_DEVICE);\n checkCudaErrors(hipblasSscal(cublasHandle, N, d_b, d_p, 1));\n hipblasSetPointerMode(cublasHandle, HIPBLAS_POINTER_MODE_HOST);\n checkCudaErrors(hipblasSaxpy(cublasHandle, N, &alpha, d_r, 1, d_p, 1));\n hipblasSetPointerMode(cublasHandle, HIPBLAS_POINTER_MODE_DEVICE);\n\n checkCudaErrors(hipsparseSetPointerMode(cusparseHandle, HIPSPARSE_POINTER_MODE_HOST));\n checkCudaErrors(hipsparseSpMV(cusparseHandle,\n HIPSPARSE_OPERATION_NON_TRANSPOSE,\n &alpha,\n matA,\n vecp,\n &beta,\n vecAx,\n HIP_R_32F,\n HIPSPARSE_SPMV_ALG_DEFAULT,\n buffer));\n\n checkCudaErrors(hipMemsetAsync(d_dot, 0, sizeof(float), stream1));\n checkCudaErrors(hipblasSdot(cublasHandle, N, d_p, 1, d_Ax, 1, d_dot));\n\n r1_div_x<<<1, 1, 0, stream1>>>(d_r1, d_dot, d_a);\n\n checkCudaErrors(hipblasSaxpy(cublasHandle, N, d_a, d_p, 1, d_x, 1));\n\n a_minus<<<1, 1, 0, stream1>>>(d_a, d_na);\n\n checkCudaErrors(hipblasSaxpy(cublasHandle, N, d_na, d_Ax, 1, d_r, 1));\n\n checkCudaErrors(hipMemcpyAsync(d_r0, d_r1, sizeof(float), hipMemcpyDeviceToDevice, stream1));\n checkCudaErrors(hipMemsetAsync(d_r1, 0, sizeof(float), stream1));\n\n checkCudaErrors(hipblasSdot(cublasHandle, N, d_r, 1, d_r, 1, d_r1));\n\n checkCudaErrors(hipMemcpyAsync((float *)&r1, d_r1, sizeof(float), hipMemcpyDeviceToHost, stream1));\n\n checkCudaErrors(hipStreamEndCapture(stream1, &initGraph));\n hipGraphExec_t graphExec;\n checkCudaErrors(hipGraphInstantiate(&graphExec, initGraph, NULL, NULL, 0));\n#endif\n\n checkCudaErrors(hipblasSetStream(cublasHandle, stream1));\n checkCudaErrors(hipsparseSetStream(cusparseHandle, stream1));\n\n while (r1 > tol * tol && k <= max_iter) {\n#if WITH_GRAPH\n checkCudaErrors(hipGraphLaunch(graphExec, streamForGraph));\n checkCudaErrors(hipStreamSynchronize(streamForGraph));\n#else\n r1_div_x<<<1, 1, 0, stream1>>>(d_r1, d_r0, d_b);\n hipblasSetPointerMode(cublasHandle, HIPBLAS_POINTER_MODE_DEVICE);\n checkCudaErrors(hipblasSscal(cublasHandle, N, d_b, d_p, 1));\n\n hipblasSetPointerMode(cublasHandle, HIPBLAS_POINTER_MODE_HOST);\n checkCudaErrors(hipblasSaxpy(cublasHandle, N, &alpha, d_r, 1, d_p, 1));\n\n checkCudaErrors(hipsparseSpMV(cusparseHandle,\n HIPSPARSE_OPERATION_NON_TRANSPOSE,\n &alpha,\n matA,\n vecp,\n &beta,\n vecAx,\n HIP_R_32F,\n HIPSPARSE_SPMV_ALG_DEFAULT,\n buffer));\n\n hipblasSetPointerMode(cublasHandle, HIPBLAS_POINTER_MODE_DEVICE);\n checkCudaErrors(hipblasSdot(cublasHandle, N, d_p, 1, d_Ax, 1, d_dot));\n\n r1_div_x<<<1, 1, 0, stream1>>>(d_r1, d_dot, d_a);\n\n checkCudaErrors(hipblasSaxpy(cublasHandle, N, d_a, d_p, 1, d_x, 1));\n\n a_minus<<<1, 1, 0, stream1>>>(d_a, d_na);\n checkCudaErrors(hipblasSaxpy(cublasHandle, N, d_na, d_Ax, 1, d_r, 1));\n\n checkCudaErrors(hipMemcpyAsync(d_r0, d_r1, sizeof(float), hipMemcpyDeviceToDevice, stream1));\n\n checkCudaErrors(hipblasSdot(cublasHandle, N, d_r, 1, d_r, 1, d_r1));\n checkCudaErrors(hipMemcpyAsync((float *)&r1, d_r1, sizeof(float), hipMemcpyDeviceToHost, stream1));\n checkCudaErrors(hipStreamSynchronize(stream1));\n#endif\n printf(\"iteration = %3d, residual = %e\\n\", k, sqrt(r1));\n k++;\n }\n\n#if WITH_GRAPH\n checkCudaErrors(hipMemcpyAsync(x, d_x, N * sizeof(float), hipMemcpyDeviceToHost, streamForGraph));\n checkCudaErrors(hipStreamSynchronize(streamForGraph));\n#else\n checkCudaErrors(hipMemcpyAsync(x, d_x, N * sizeof(float), hipMemcpyDeviceToHost, stream1));\n checkCudaErrors(hipStreamSynchronize(stream1));\n#endif\n\n float rsum, diff, err = 0.0;\n\n for (int i = 0; i < N; i++) {\n rsum = 0.0;\n\n for (int j = I[i]; j < I[i + 1]; j++) {\n rsum += val[j] * x[J[j]];\n }\n\n diff = fabs(rsum - rhs[i]);\n\n if (diff > err) {\n err = diff;\n }\n }\n\n#if WITH_GRAPH\n checkCudaErrors(hipGraphExecDestroy(graphExec));\n checkCudaErrors(hipGraphDestroy(initGraph));\n checkCudaErrors(hipStreamDestroy(streamForGraph));\n#endif\n checkCudaErrors(hipStreamDestroy(stream1));\n hipsparseDestroy(cusparseHandle);\n hipblasDestroy(cublasHandle);\n\n if (matA) {\n checkCudaErrors(hipsparseDestroySpMat(matA));\n }\n if (vecx) {\n checkCudaErrors(hipsparseDestroyDnVec(vecx));\n }\n if (vecAx) {\n checkCudaErrors(hipsparseDestroyDnVec(vecAx));\n }\n if (vecp) {\n checkCudaErrors(hipsparseDestroyDnVec(vecp));\n }\n\n checkCudaErrors(hipHostFree(I));\n checkCudaErrors(hipHostFree(J));\n checkCudaErrors(hipHostFree(val));\n checkCudaErrors(hipHostFree(x));\n free(rhs);\n checkCudaErrors(hipFree(d_col));\n checkCudaErrors(hipFree(d_row));\n checkCudaErrors(hipFree(d_val));\n checkCudaErrors(hipFree(d_x));\n checkCudaErrors(hipFree(d_r));\n checkCudaErrors(hipFree(d_p));\n checkCudaErrors(hipFree(d_Ax));\n\n printf(\"Test Summary: Error amount = %f\\n\", err);\n exit((k <= max_iter) ? 0 : 1);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/4_CUDA_Libraries/conjugateGradientCudaGraphs/conjugateGradientCudaGraphs.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* Example showing the use of CUFFT for fast 1D-convolution using FFT. */\n\n// System includes\n#include \n#include \n#include \n#include \n\n// CUDA runtime\n#include \n\n// CUFFT Header file\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Complex data type\ntypedef float2 Complex;\n\nstatic __device__ __host__ inline Complex ComplexAdd(Complex, Complex);\nstatic __device__ __host__ inline Complex ComplexScale(Complex, float);\nstatic __device__ __host__ inline Complex ComplexMul(Complex, Complex);\nstatic __global__ void ComplexPointwiseMulAndScale(cufftComplex *, cufftComplex *, int, float);\n\n// Kernel for GPU\nvoid multiplyCoefficient(cudaLibXtDesc *, cudaLibXtDesc *, int, float, int);\n\n// Filtering functions\nvoid Convolve(const Complex *, int, const Complex *, int, Complex *);\n\n// Padding functions\nint PadData(const Complex *, Complex **, int, const Complex *, Complex **, int);\n\n////////////////////////////////////////////////////////////////////////////////\n// Data configuration\n// The filter size is assumed to be a number smaller than the signal size\n///////////////////////////////////////////////////////////////////////////////\nconst int SIGNAL_SIZE = 1018;\nconst int FILTER_KERNEL_SIZE = 11;\nconst int GPU_COUNT = 2;\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"\\n[simpleCUFFT_MGPU] is starting...\\n\\n\");\n\n int GPU_N;\n checkCudaErrors(cudaGetDeviceCount(&GPU_N));\n\n if (GPU_N < GPU_COUNT) {\n printf(\"No. of GPU on node %d\\n\", GPU_N);\n printf(\"Two GPUs are required to run simpleCUFFT_MGPU sample code\\n\");\n exit(EXIT_WAIVED);\n }\n\n int *major_minor = (int *)malloc(sizeof(int) * GPU_N * 2);\n int found2IdenticalGPUs = 0;\n int nGPUs = 2;\n int *whichGPUs;\n whichGPUs = (int *)malloc(sizeof(int) * nGPUs);\n\n for (int i = 0; i < GPU_N; i++) {\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, i));\n major_minor[i * 2] = deviceProp.major;\n major_minor[i * 2 + 1] = deviceProp.minor;\n printf(\"GPU Device %d: \\\"%s\\\" with compute capability %d.%d\\n\",\n i,\n deviceProp.name,\n deviceProp.major,\n deviceProp.minor);\n }\n\n for (int i = 0; i < GPU_N; i++) {\n for (int j = i + 1; j < GPU_N; j++) {\n if ((major_minor[i * 2] == major_minor[j * 2]) && (major_minor[i * 2 + 1] == major_minor[j * 2 + 1])) {\n whichGPUs[0] = i;\n whichGPUs[1] = j;\n found2IdenticalGPUs = 1;\n break;\n }\n }\n if (found2IdenticalGPUs) {\n break;\n }\n }\n\n free(major_minor);\n if (!found2IdenticalGPUs) {\n printf(\"No Two GPUs with same architecture found\\nWaiving simpleCUFFT_2d_MGPU \"\n \"sample\\n\");\n exit(EXIT_WAIVED);\n }\n\n // Allocate host memory for the signal\n Complex *h_signal = (Complex *)malloc(sizeof(Complex) * SIGNAL_SIZE);\n\n // Initialize the memory for the signal\n for (int i = 0; i < SIGNAL_SIZE; ++i) {\n h_signal[i].x = rand() / (float)RAND_MAX;\n h_signal[i].y = 0;\n }\n\n // Allocate host memory for the filter\n Complex *h_filter_kernel = (Complex *)malloc(sizeof(Complex) * FILTER_KERNEL_SIZE);\n\n // Initialize the memory for the filter\n for (int i = 0; i < FILTER_KERNEL_SIZE; ++i) {\n h_filter_kernel[i].x = rand() / (float)RAND_MAX;\n h_filter_kernel[i].y = 0;\n }\n\n // Pad signal and filter kernel\n Complex *h_padded_signal;\n Complex *h_padded_filter_kernel;\n int new_size =\n PadData(h_signal, &h_padded_signal, SIGNAL_SIZE, h_filter_kernel, &h_padded_filter_kernel, FILTER_KERNEL_SIZE);\n\n // cufftCreate() - Create an empty plan\n cufftResult result;\n cufftHandle plan_input;\n checkCudaErrors(cufftCreate(&plan_input));\n\n // cufftXtSetGPUs() - Define which GPUs to use\n result = cufftXtSetGPUs(plan_input, nGPUs, whichGPUs);\n\n if (result == CUFFT_INVALID_DEVICE) {\n printf(\"This sample requires two GPUs on the same board.\\n\");\n printf(\"No such board was found. Waiving sample.\\n\");\n exit(EXIT_WAIVED);\n }\n else if (result != CUFFT_SUCCESS) {\n printf(\"cufftXtSetGPUs failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // Print the device information to run the code\n printf(\"\\nRunning on GPUs\\n\");\n for (int i = 0; i < nGPUs; i++) {\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, whichGPUs[i]));\n printf(\"GPU Device %d: \\\"%s\\\" with compute capability %d.%d\\n\",\n whichGPUs[i],\n deviceProp.name,\n deviceProp.major,\n deviceProp.minor);\n }\n\n size_t *worksize;\n worksize = (size_t *)malloc(sizeof(size_t) * nGPUs);\n\n // cufftMakePlan1d() - Create the plan\n checkCudaErrors(cufftMakePlan1d(plan_input, new_size, CUFFT_C2C, 1, worksize));\n\n // cufftXtMalloc() - Malloc data on multiple GPUs\n cudaLibXtDesc *d_signal;\n checkCudaErrors(cufftXtMalloc(plan_input, (cudaLibXtDesc **)&d_signal, CUFFT_XT_FORMAT_INPLACE));\n cudaLibXtDesc *d_out_signal;\n checkCudaErrors(cufftXtMalloc(plan_input, (cudaLibXtDesc **)&d_out_signal, CUFFT_XT_FORMAT_INPLACE));\n cudaLibXtDesc *d_filter_kernel;\n checkCudaErrors(cufftXtMalloc(plan_input, (cudaLibXtDesc **)&d_filter_kernel, CUFFT_XT_FORMAT_INPLACE));\n cudaLibXtDesc *d_out_filter_kernel;\n checkCudaErrors(cufftXtMalloc(plan_input, (cudaLibXtDesc **)&d_out_filter_kernel, CUFFT_XT_FORMAT_INPLACE));\n\n // cufftXtMemcpy() - Copy data from host to multiple GPUs\n checkCudaErrors(cufftXtMemcpy(plan_input, d_signal, h_padded_signal, CUFFT_COPY_HOST_TO_DEVICE));\n checkCudaErrors(cufftXtMemcpy(plan_input, d_filter_kernel, h_padded_filter_kernel, CUFFT_COPY_HOST_TO_DEVICE));\n\n // cufftXtExecDescriptorC2C() - Execute FFT on data on multiple GPUs\n checkCudaErrors(cufftXtExecDescriptorC2C(plan_input, d_signal, d_signal, CUFFT_FORWARD));\n checkCudaErrors(cufftXtExecDescriptorC2C(plan_input, d_filter_kernel, d_filter_kernel, CUFFT_FORWARD));\n\n // cufftXtMemcpy() - Copy the data to natural order on GPUs\n checkCudaErrors(cufftXtMemcpy(plan_input, d_out_signal, d_signal, CUFFT_COPY_DEVICE_TO_DEVICE));\n checkCudaErrors(cufftXtMemcpy(plan_input, d_out_filter_kernel, d_filter_kernel, CUFFT_COPY_DEVICE_TO_DEVICE));\n\n printf(\"\\n\\nValue of Library Descriptor\\n\");\n printf(\"Number of GPUs %d\\n\", d_out_signal->descriptor->nGPUs);\n printf(\"Device id %d %d\\n\", d_out_signal->descriptor->GPUs[0], d_out_signal->descriptor->GPUs[1]);\n printf(\"Data size on GPU %ld %ld\\n\",\n (long)(d_out_signal->descriptor->size[0] / sizeof(cufftComplex)),\n (long)(d_out_signal->descriptor->size[1] / sizeof(cufftComplex)));\n\n // Multiply the coefficients together and normalize the result\n printf(\"Launching ComplexPointwiseMulAndScale<<< >>>\\n\");\n multiplyCoefficient(d_out_signal, d_out_filter_kernel, new_size, 1.0f / new_size, nGPUs);\n\n // cufftXtExecDescriptorC2C() - Execute inverse FFT on data on multiple GPUs\n printf(\"Transforming signal back cufftExecC2C\\n\");\n checkCudaErrors(cufftXtExecDescriptorC2C(plan_input, d_out_signal, d_out_signal, CUFFT_INVERSE));\n\n // Create host pointer pointing to padded signal\n Complex *h_convolved_signal = h_padded_signal;\n\n // Allocate host memory for the convolution result\n Complex *h_convolved_signal_ref = (Complex *)malloc(sizeof(Complex) * SIGNAL_SIZE);\n\n // cufftXtMemcpy() - Copy data from multiple GPUs to host\n checkCudaErrors(cufftXtMemcpy(plan_input, h_convolved_signal, d_out_signal, CUFFT_COPY_DEVICE_TO_HOST));\n\n // Convolve on the host\n Convolve(h_signal, SIGNAL_SIZE, h_filter_kernel, FILTER_KERNEL_SIZE, h_convolved_signal_ref);\n\n // Compare CPU and GPU result\n bool bTestResult =\n sdkCompareL2fe((float *)h_convolved_signal_ref, (float *)h_convolved_signal, 2 * SIGNAL_SIZE, 1e-5f);\n printf(\"\\nvalue of TestResult %d\\n\", bTestResult);\n\n // Cleanup memory\n free(whichGPUs);\n free(worksize);\n free(h_signal);\n free(h_filter_kernel);\n free(h_padded_signal);\n free(h_padded_filter_kernel);\n free(h_convolved_signal_ref);\n\n // cudaXtFree() - Free GPU memory\n checkCudaErrors(cufftXtFree(d_signal));\n checkCudaErrors(cufftXtFree(d_filter_kernel));\n checkCudaErrors(cufftXtFree(d_out_signal));\n checkCudaErrors(cufftXtFree(d_out_filter_kernel));\n\n // cufftDestroy() - Destroy FFT plan\n checkCudaErrors(cufftDestroy(plan_input));\n\n exit(bTestResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\n///////////////////////////////////////////////////////////////////////////////////\n// Function for padding original data\n//////////////////////////////////////////////////////////////////////////////////\nint PadData(const Complex *signal,\n Complex **padded_signal,\n int signal_size,\n const Complex *filter_kernel,\n Complex **padded_filter_kernel,\n int filter_kernel_size)\n{\n int minRadius = filter_kernel_size / 2;\n int maxRadius = filter_kernel_size - minRadius;\n int new_size = signal_size + maxRadius;\n\n // Pad signal\n Complex *new_data = (Complex *)malloc(sizeof(Complex) * new_size);\n memcpy(new_data + 0, signal, signal_size * sizeof(Complex));\n memset(new_data + signal_size, 0, (new_size - signal_size) * sizeof(Complex));\n *padded_signal = new_data;\n\n // Pad filter\n new_data = (Complex *)malloc(sizeof(Complex) * new_size);\n memcpy(new_data + 0, filter_kernel + minRadius, maxRadius * sizeof(Complex));\n memset(new_data + maxRadius, 0, (new_size - filter_kernel_size) * sizeof(Complex));\n memcpy(new_data + new_size - minRadius, filter_kernel, minRadius * sizeof(Complex));\n *padded_filter_kernel = new_data;\n\n return new_size;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Filtering operations - Computing Convolution on the host\n////////////////////////////////////////////////////////////////////////////////\nvoid Convolve(const Complex *signal,\n int signal_size,\n const Complex *filter_kernel,\n int filter_kernel_size,\n Complex *filtered_signal)\n{\n int minRadius = filter_kernel_size / 2;\n int maxRadius = filter_kernel_size - minRadius;\n\n // Loop over output element indices\n for (int i = 0; i < signal_size; ++i) {\n filtered_signal[i].x = filtered_signal[i].y = 0;\n\n // Loop over convolution indices\n for (int j = -maxRadius + 1; j <= minRadius; ++j) {\n int k = i + j;\n\n if (k >= 0 && k < signal_size) {\n filtered_signal[i] =\n ComplexAdd(filtered_signal[i], ComplexMul(signal[k], filter_kernel[minRadius - j]));\n }\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Launch Kernel on multiple GPU\n////////////////////////////////////////////////////////////////////////////////\nvoid multiplyCoefficient(cudaLibXtDesc *d_signal, cudaLibXtDesc *d_filter_kernel, int new_size, float val, int nGPUs)\n{\n int device;\n // Launch the ComplexPointwiseMulAndScale<<< >>> kernel on multiple GPU\n for (int i = 0; i < nGPUs; i++) {\n device = d_signal->descriptor->GPUs[i];\n\n // Set device\n checkCudaErrors(cudaSetDevice(device));\n\n // Perform GPU computations\n ComplexPointwiseMulAndScale<<<32, 256>>>((cufftComplex *)d_signal->descriptor->data[i],\n (cufftComplex *)d_filter_kernel->descriptor->data[i],\n int(d_signal->descriptor->size[i] / sizeof(cufftComplex)),\n val);\n }\n\n // Wait for device to finish all operation\n for (int i = 0; i < nGPUs; i++) {\n device = d_signal->descriptor->GPUs[i];\n checkCudaErrors(cudaSetDevice(device));\n cudaDeviceSynchronize();\n // Check if kernel execution generated and error\n getLastCudaError(\"Kernel execution failed [ ComplexPointwiseMulAndScale ]\");\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Complex operations\n////////////////////////////////////////////////////////////////////////////////\n\n// Complex addition\nstatic __device__ __host__ inline Complex ComplexAdd(Complex a, Complex b)\n{\n Complex c;\n c.x = a.x + b.x;\n c.y = a.y + b.y;\n return c;\n}\n\n// Complex scale\nstatic __device__ __host__ inline Complex ComplexScale(Complex a, float s)\n{\n Complex c;\n c.x = s * a.x;\n c.y = s * a.y;\n return c;\n}\n\n// Complex multiplication\nstatic __device__ __host__ inline Complex ComplexMul(Complex a, Complex b)\n{\n Complex c;\n c.x = a.x * b.x - a.y * b.y;\n c.y = a.x * b.y + a.y * b.x;\n return c;\n}\n// Complex pointwise multiplication\nstatic __global__ void ComplexPointwiseMulAndScale(cufftComplex *a, cufftComplex *b, int size, float scale)\n{\n const int numThreads = blockDim.x * gridDim.x;\n const int threadID = blockIdx.x * blockDim.x + threadIdx.x;\n for (int i = threadID; i < size; i += numThreads) {\n a[i] = ComplexScale(ComplexMul(a[i], b[i]), scale);\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/* Example showing the use of CUFFT for fast 1D-convolution using FFT. */\n\n// System includes\n#include \n#include \n#include \n#include \n\n// CUDA runtime\n#include \n\n// CUFFT Header file\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Complex data type\ntypedef float2 Complex;\n\nstatic __device__ __host__ inline Complex ComplexAdd(Complex, Complex);\nstatic __device__ __host__ inline Complex ComplexScale(Complex, float);\nstatic __device__ __host__ inline Complex ComplexMul(Complex, Complex);\nstatic __global__ void ComplexPointwiseMulAndScale(hipfftComplex *, hipfftComplex *, int, float);\n\n// Kernel for GPU\nvoid multiplyCoefficient(hipLibXtDesc *, hipLibXtDesc *, int, float, int);\n\n// Filtering functions\nvoid Convolve(const Complex *, int, const Complex *, int, Complex *);\n\n// Padding functions\nint PadData(const Complex *, Complex **, int, const Complex *, Complex **, int);\n\n////////////////////////////////////////////////////////////////////////////////\n// Data configuration\n// The filter size is assumed to be a number smaller than the signal size\n///////////////////////////////////////////////////////////////////////////////\nconst int SIGNAL_SIZE = 1018;\nconst int FILTER_KERNEL_SIZE = 11;\nconst int GPU_COUNT = 2;\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"\\n[simpleCUFFT_MGPU] is starting...\\n\\n\");\n\n int GPU_N;\n checkCudaErrors(hipGetDeviceCount(&GPU_N));\n\n if (GPU_N < GPU_COUNT) {\n printf(\"No. of GPU on node %d\\n\", GPU_N);\n printf(\"Two GPUs are required to run simpleCUFFT_MGPU sample code\\n\");\n exit(EXIT_WAIVED);\n }\n\n int *major_minor = (int *)malloc(sizeof(int) * GPU_N * 2);\n int found2IdenticalGPUs = 0;\n int nGPUs = 2;\n int *whichGPUs;\n whichGPUs = (int *)malloc(sizeof(int) * nGPUs);\n\n for (int i = 0; i < GPU_N; i++) {\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, i));\n major_minor[i * 2] = deviceProp.major;\n major_minor[i * 2 + 1] = deviceProp.minor;\n printf(\"GPU Device %d: \\\"%s\\\" with compute capability %d.%d\\n\",\n i,\n deviceProp.name,\n deviceProp.major,\n deviceProp.minor);\n }\n\n for (int i = 0; i < GPU_N; i++) {\n for (int j = i + 1; j < GPU_N; j++) {\n if ((major_minor[i * 2] == major_minor[j * 2]) && (major_minor[i * 2 + 1] == major_minor[j * 2 + 1])) {\n whichGPUs[0] = i;\n whichGPUs[1] = j;\n found2IdenticalGPUs = 1;\n break;\n }\n }\n if (found2IdenticalGPUs) {\n break;\n }\n }\n\n free(major_minor);\n if (!found2IdenticalGPUs) {\n printf(\"No Two GPUs with same architecture found\\nWaiving simpleCUFFT_2d_MGPU \"\n \"sample\\n\");\n exit(EXIT_WAIVED);\n }\n\n // Allocate host memory for the signal\n Complex *h_signal = (Complex *)malloc(sizeof(Complex) * SIGNAL_SIZE);\n\n // Initialize the memory for the signal\n for (int i = 0; i < SIGNAL_SIZE; ++i) {\n h_signal[i].x = rand() / (float)RAND_MAX;\n h_signal[i].y = 0;\n }\n\n // Allocate host memory for the filter\n Complex *h_filter_kernel = (Complex *)malloc(sizeof(Complex) * FILTER_KERNEL_SIZE);\n\n // Initialize the memory for the filter\n for (int i = 0; i < FILTER_KERNEL_SIZE; ++i) {\n h_filter_kernel[i].x = rand() / (float)RAND_MAX;\n h_filter_kernel[i].y = 0;\n }\n\n // Pad signal and filter kernel\n Complex *h_padded_signal;\n Complex *h_padded_filter_kernel;\n int new_size =\n PadData(h_signal, &h_padded_signal, SIGNAL_SIZE, h_filter_kernel, &h_padded_filter_kernel, FILTER_KERNEL_SIZE);\n\n // cufftCreate() - Create an empty plan\n hipfftResult result;\n hipfftHandle plan_input;\n checkCudaErrors(hipfftCreate(&plan_input));\n\n // cufftXtSetGPUs() - Define which GPUs to use\n result = hipfftXtSetGPUs(plan_input, nGPUs, whichGPUs);\n\n if (result == HIPFFT_INVALID_DEVICE) {\n printf(\"This sample requires two GPUs on the same board.\\n\");\n printf(\"No such board was found. Waiving sample.\\n\");\n exit(EXIT_WAIVED);\n }\n else if (result != HIPFFT_SUCCESS) {\n printf(\"hipfftXtSetGPUs failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // Print the device information to run the code\n printf(\"\\nRunning on GPUs\\n\");\n for (int i = 0; i < nGPUs; i++) {\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, whichGPUs[i]));\n printf(\"GPU Device %d: \\\"%s\\\" with compute capability %d.%d\\n\",\n whichGPUs[i],\n deviceProp.name,\n deviceProp.major,\n deviceProp.minor);\n }\n\n size_t *worksize;\n worksize = (size_t *)malloc(sizeof(size_t) * nGPUs);\n\n // cufftMakePlan1d() - Create the plan\n checkCudaErrors(hipfftMakePlan1d(plan_input, new_size, HIPFFT_C2C, 1, worksize));\n\n // cufftXtMalloc() - Malloc data on multiple GPUs\n hipLibXtDesc *d_signal;\n checkCudaErrors(hipfftXtMalloc(plan_input, (hipLibXtDesc **)&d_signal, HIPFFT_XT_FORMAT_INPLACE));\n hipLibXtDesc *d_out_signal;\n checkCudaErrors(hipfftXtMalloc(plan_input, (hipLibXtDesc **)&d_out_signal, HIPFFT_XT_FORMAT_INPLACE));\n hipLibXtDesc *d_filter_kernel;\n checkCudaErrors(hipfftXtMalloc(plan_input, (hipLibXtDesc **)&d_filter_kernel, HIPFFT_XT_FORMAT_INPLACE));\n hipLibXtDesc *d_out_filter_kernel;\n checkCudaErrors(hipfftXtMalloc(plan_input, (hipLibXtDesc **)&d_out_filter_kernel, HIPFFT_XT_FORMAT_INPLACE));\n\n // cufftXtMemcpy() - Copy data from host to multiple GPUs\n checkCudaErrors(hipfftXtMemcpy(plan_input, d_signal, h_padded_signal, HIPFFT_COPY_HOST_TO_DEVICE));\n checkCudaErrors(hipfftXtMemcpy(plan_input, d_filter_kernel, h_padded_filter_kernel, HIPFFT_COPY_HOST_TO_DEVICE));\n\n // cufftXtExecDescriptorC2C() - Execute FFT on data on multiple GPUs\n checkCudaErrors(hipfftXtExecDescriptorC2C(plan_input, d_signal, d_signal, HIPFFT_FORWARD));\n checkCudaErrors(hipfftXtExecDescriptorC2C(plan_input, d_filter_kernel, d_filter_kernel, HIPFFT_FORWARD));\n\n // cufftXtMemcpy() - Copy the data to natural order on GPUs\n checkCudaErrors(hipfftXtMemcpy(plan_input, d_out_signal, d_signal, HIPFFT_COPY_DEVICE_TO_DEVICE));\n checkCudaErrors(hipfftXtMemcpy(plan_input, d_out_filter_kernel, d_filter_kernel, HIPFFT_COPY_DEVICE_TO_DEVICE));\n\n printf(\"\\n\\nValue of Library Descriptor\\n\");\n printf(\"Number of GPUs %d\\n\", d_out_signal->descriptor->nGPUs);\n printf(\"Device id %d %d\\n\", d_out_signal->descriptor->GPUs[0], d_out_signal->descriptor->GPUs[1]);\n printf(\"Data size on GPU %ld %ld\\n\",\n (long)(d_out_signal->descriptor->size[0] / sizeof(hipfftComplex)),\n (long)(d_out_signal->descriptor->size[1] / sizeof(hipfftComplex)));\n\n // Multiply the coefficients together and normalize the result\n printf(\"Launching ComplexPointwiseMulAndScale<<< >>>\\n\");\n multiplyCoefficient(d_out_signal, d_out_filter_kernel, new_size, 1.0f / new_size, nGPUs);\n\n // cufftXtExecDescriptorC2C() - Execute inverse FFT on data on multiple GPUs\n printf(\"Transforming signal back hipfftExecC2C\\n\");\n checkCudaErrors(hipfftXtExecDescriptorC2C(plan_input, d_out_signal, d_out_signal, HIPFFT_BACKWARD));\n\n // Create host pointer pointing to padded signal\n Complex *h_convolved_signal = h_padded_signal;\n\n // Allocate host memory for the convolution result\n Complex *h_convolved_signal_ref = (Complex *)malloc(sizeof(Complex) * SIGNAL_SIZE);\n\n // cufftXtMemcpy() - Copy data from multiple GPUs to host\n checkCudaErrors(hipfftXtMemcpy(plan_input, h_convolved_signal, d_out_signal, HIPFFT_COPY_DEVICE_TO_HOST));\n\n // Convolve on the host\n Convolve(h_signal, SIGNAL_SIZE, h_filter_kernel, FILTER_KERNEL_SIZE, h_convolved_signal_ref);\n\n // Compare CPU and GPU result\n bool bTestResult =\n sdkCompareL2fe((float *)h_convolved_signal_ref, (float *)h_convolved_signal, 2 * SIGNAL_SIZE, 1e-5f);\n printf(\"\\nvalue of TestResult %d\\n\", bTestResult);\n\n // Cleanup memory\n free(whichGPUs);\n free(worksize);\n free(h_signal);\n free(h_filter_kernel);\n free(h_padded_signal);\n free(h_padded_filter_kernel);\n free(h_convolved_signal_ref);\n\n // cudaXtFree() - Free GPU memory\n checkCudaErrors(hipfftXtFree(d_signal));\n checkCudaErrors(hipfftXtFree(d_filter_kernel));\n checkCudaErrors(hipfftXtFree(d_out_signal));\n checkCudaErrors(hipfftXtFree(d_out_filter_kernel));\n\n // cufftDestroy() - Destroy FFT plan\n checkCudaErrors(hipfftDestroy(plan_input));\n\n exit(bTestResult ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\n///////////////////////////////////////////////////////////////////////////////////\n// Function for padding original data\n//////////////////////////////////////////////////////////////////////////////////\nint PadData(const Complex *signal,\n Complex **padded_signal,\n int signal_size,\n const Complex *filter_kernel,\n Complex **padded_filter_kernel,\n int filter_kernel_size)\n{\n int minRadius = filter_kernel_size / 2;\n int maxRadius = filter_kernel_size - minRadius;\n int new_size = signal_size + maxRadius;\n\n // Pad signal\n Complex *new_data = (Complex *)malloc(sizeof(Complex) * new_size);\n memcpy(new_data + 0, signal, signal_size * sizeof(Complex));\n memset(new_data + signal_size, 0, (new_size - signal_size) * sizeof(Complex));\n *padded_signal = new_data;\n\n // Pad filter\n new_data = (Complex *)malloc(sizeof(Complex) * new_size);\n memcpy(new_data + 0, filter_kernel + minRadius, maxRadius * sizeof(Complex));\n memset(new_data + maxRadius, 0, (new_size - filter_kernel_size) * sizeof(Complex));\n memcpy(new_data + new_size - minRadius, filter_kernel, minRadius * sizeof(Complex));\n *padded_filter_kernel = new_data;\n\n return new_size;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Filtering operations - Computing Convolution on the host\n////////////////////////////////////////////////////////////////////////////////\nvoid Convolve(const Complex *signal,\n int signal_size,\n const Complex *filter_kernel,\n int filter_kernel_size,\n Complex *filtered_signal)\n{\n int minRadius = filter_kernel_size / 2;\n int maxRadius = filter_kernel_size - minRadius;\n\n // Loop over output element indices\n for (int i = 0; i < signal_size; ++i) {\n filtered_signal[i].x = filtered_signal[i].y = 0;\n\n // Loop over convolution indices\n for (int j = -maxRadius + 1; j <= minRadius; ++j) {\n int k = i + j;\n\n if (k >= 0 && k < signal_size) {\n filtered_signal[i] =\n ComplexAdd(filtered_signal[i], ComplexMul(signal[k], filter_kernel[minRadius - j]));\n }\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Launch Kernel on multiple GPU\n////////////////////////////////////////////////////////////////////////////////\nvoid multiplyCoefficient(hipLibXtDesc *d_signal, hipLibXtDesc *d_filter_kernel, int new_size, float val, int nGPUs)\n{\n int device;\n // Launch the ComplexPointwiseMulAndScale<<< >>> kernel on multiple GPU\n for (int i = 0; i < nGPUs; i++) {\n device = d_signal->descriptor->GPUs[i];\n\n // Set device\n checkCudaErrors(hipSetDevice(device));\n\n // Perform GPU computations\n ComplexPointwiseMulAndScale<<<32, 256>>>((hipfftComplex *)d_signal->descriptor->data[i],\n (hipfftComplex *)d_filter_kernel->descriptor->data[i],\n int(d_signal->descriptor->size[i] / sizeof(hipfftComplex)),\n val);\n }\n\n // Wait for device to finish all operation\n for (int i = 0; i < nGPUs; i++) {\n device = d_signal->descriptor->GPUs[i];\n checkCudaErrors(hipSetDevice(device));\n hipDeviceSynchronize();\n // Check if kernel execution generated and error\n getLastCudaError(\"Kernel execution failed [ ComplexPointwiseMulAndScale ]\");\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Complex operations\n////////////////////////////////////////////////////////////////////////////////\n\n// Complex addition\nstatic __device__ __host__ inline Complex ComplexAdd(Complex a, Complex b)\n{\n Complex c;\n c.x = a.x + b.x;\n c.y = a.y + b.y;\n return c;\n}\n\n// Complex scale\nstatic __device__ __host__ inline Complex ComplexScale(Complex a, float s)\n{\n Complex c;\n c.x = s * a.x;\n c.y = s * a.y;\n return c;\n}\n\n// Complex multiplication\nstatic __device__ __host__ inline Complex ComplexMul(Complex a, Complex b)\n{\n Complex c;\n c.x = a.x * b.x - a.y * b.y;\n c.y = a.x * b.y + a.y * b.x;\n return c;\n}\n// Complex pointwise multiplication\nstatic __global__ void ComplexPointwiseMulAndScale(hipfftComplex *a, hipfftComplex *b, int size, float scale)\n{\n const int numThreads = blockDim.x * gridDim.x;\n const int threadID = blockIdx.x * blockDim.x + threadIdx.x;\n for (int i = threadID; i < size; i += numThreads) {\n a[i] = ComplexScale(ComplexMul(a[i], b[i]), scale);\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/4_CUDA_Libraries/simpleCUFFT_MGPU/simpleCUFFT_MGPU.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// simpleCUFFT_2d_MGPU.cu\n//\n// This sample code demonstrate the use of CUFFT library for 2D data on multiple GPU.\n// Example showing the use of CUFFT for solving 2D-POISSON equation using FFT on multiple GPU.\n// For reference we have used the equation given in http://www.bu.edu/pasi/files/2011/07/\n// Lecture83.pdf\n//\n////////////////////////////////////////////////////////////////////////////////\n\n\n// System includes\n#include \n#include \n#include \n#include \n\n// CUDA runtime\n#include \n\n// CUFFT Header file\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Complex data type\ntypedef float2 Complex;\n\n// Data configuration\nconst int GPU_COUNT = 2;\nconst int BSZ_Y = 4;\nconst int BSZ_X = 4;\n\n// Forward Declaration\nvoid solvePoissonEquation(cudaLibXtDesc *, cudaLibXtDesc *, float **, int, int);\n\n__global__ void solvePoisson(cufftComplex *, cufftComplex *, float *, int, int, int n_gpu);\n\n///////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"\\nPoisson equation using CUFFT library on Multiple GPUs is \"\n \"starting...\\n\\n\");\n\n int GPU_N;\n checkCudaErrors(cudaGetDeviceCount(&GPU_N));\n\n if (GPU_N < GPU_COUNT) {\n printf(\"No. of GPU on node %d\\n\", GPU_N);\n printf(\"Two GPUs are required to run simpleCUFFT_2d_MGPU sample code\\n\");\n exit(EXIT_WAIVED);\n }\n\n int *major_minor = (int *)malloc(sizeof(int) * GPU_N * 2);\n int found2IdenticalGPUs = 0;\n int nGPUs = 2;\n int *whichGPUs;\n whichGPUs = (int *)malloc(sizeof(int) * nGPUs);\n\n for (int i = 0; i < GPU_N; i++) {\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, i));\n major_minor[i * 2] = deviceProp.major;\n major_minor[i * 2 + 1] = deviceProp.minor;\n printf(\"GPU Device %d: \\\"%s\\\" with compute capability %d.%d\\n\",\n i,\n deviceProp.name,\n deviceProp.major,\n deviceProp.minor);\n }\n\n for (int i = 0; i < GPU_N; i++) {\n for (int j = i + 1; j < GPU_N; j++) {\n if ((major_minor[i * 2] == major_minor[j * 2]) && (major_minor[i * 2 + 1] == major_minor[j * 2 + 1])) {\n whichGPUs[0] = i;\n whichGPUs[1] = j;\n found2IdenticalGPUs = 1;\n break;\n }\n }\n if (found2IdenticalGPUs) {\n break;\n }\n }\n\n free(major_minor);\n if (!found2IdenticalGPUs) {\n printf(\"No Two GPUs with same architecture found\\nWaiving simpleCUFFT_2d_MGPU \"\n \"sample\\n\");\n exit(EXIT_WAIVED);\n }\n\n int N = 64;\n float xMAX = 1.0f, xMIN = 0.0f, yMIN = 0.0f, h = (xMAX - xMIN) / ((float)N), s = 0.1f, s2 = s * s;\n float *x, *y, *f, *u_a, r2;\n\n x = (float *)malloc(sizeof(float) * N * N);\n y = (float *)malloc(sizeof(float) * N * N);\n f = (float *)malloc(sizeof(float) * N * N);\n u_a = (float *)malloc(sizeof(float) * N * N);\n\n for (int j = 0; j < N; j++)\n for (int i = 0; i < N; i++) {\n x[N * j + i] = xMIN + i * h;\n y[N * j + i] = yMIN + j * h;\n r2 = (x[N * j + i] - 0.5f) * (x[N * j + i] - 0.5f) + (y[N * j + i] - 0.5f) * (y[N * j + i] - 0.5f);\n f[N * j + i] = (r2 - 2 * s2) / (s2 * s2) * exp(-r2 / (2 * s2));\n u_a[N * j + i] = exp(-r2 / (2 * s2)); // analytical solution\n }\n\n float *k, *d_k[GPU_COUNT];\n k = (float *)malloc(sizeof(float) * N);\n for (int i = 0; i <= N / 2; i++) {\n k[i] = i * 2 * (float)M_PI;\n }\n for (int i = N / 2 + 1; i < N; i++) {\n k[i] = (i - N) * 2 * (float)M_PI;\n }\n\n // Create a complex variable on host\n Complex *h_f = (Complex *)malloc(sizeof(Complex) * N * N);\n\n // Initialize the memory for the signal\n for (int i = 0; i < (N * N); i++) {\n h_f[i].x = f[i];\n h_f[i].y = 0.0f;\n }\n\n // cufftCreate() - Create an empty plan\n cufftResult result;\n cufftHandle planComplex;\n result = cufftCreate(&planComplex);\n if (result != CUFFT_SUCCESS) {\n printf(\"cufftCreate failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // cufftXtSetGPUs() - Define which GPUs to use\n result = cufftXtSetGPUs(planComplex, nGPUs, whichGPUs);\n\n if (result == CUFFT_INVALID_DEVICE) {\n printf(\"This sample requires two GPUs on the same board.\\n\");\n printf(\"No such board was found. Waiving sample.\\n\");\n exit(EXIT_WAIVED);\n }\n else if (result != CUFFT_SUCCESS) {\n printf(\"cufftXtSetGPUs failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // Print the device information to run the code\n printf(\"\\nRunning on GPUs\\n\");\n for (int i = 0; i < 2; i++) {\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, whichGPUs[i]));\n printf(\"GPU Device %d: \\\"%s\\\" with compute capability %d.%d\\n\",\n whichGPUs[i],\n deviceProp.name,\n deviceProp.major,\n deviceProp.minor);\n }\n\n size_t *worksize;\n worksize = (size_t *)malloc(sizeof(size_t) * nGPUs);\n\n // cufftMakePlan2d() - Create the plan\n result = cufftMakePlan2d(planComplex, N, N, CUFFT_C2C, worksize);\n if (result != CUFFT_SUCCESS) {\n printf(\"*MakePlan* failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n for (int i = 0; i < nGPUs; i++) {\n cudaSetDevice(whichGPUs[i]);\n cudaMalloc((void **)&d_k[i], sizeof(float) * N);\n cudaMemcpy(d_k[i], k, sizeof(float) * N, cudaMemcpyHostToDevice);\n }\n\n // Create a variable on device\n // d_f - variable on device to store the input data\n // d_d_f - variable that store the natural order of d_f data\n // d_out - device output\n cudaLibXtDesc *d_f, *d_d_f, *d_out;\n\n // cufftXtMalloc() - Malloc data on multiple GPUs\n\n result = cufftXtMalloc(planComplex, (cudaLibXtDesc **)&d_f, CUFFT_XT_FORMAT_INPLACE);\n if (result != CUFFT_SUCCESS) {\n printf(\"*XtMalloc failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n result = cufftXtMalloc(planComplex, (cudaLibXtDesc **)&d_d_f, CUFFT_XT_FORMAT_INPLACE);\n if (result != CUFFT_SUCCESS) {\n printf(\"*XtMalloc failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n result = cufftXtMalloc(planComplex, (cudaLibXtDesc **)&d_out, CUFFT_XT_FORMAT_INPLACE);\n if (result != CUFFT_SUCCESS) {\n printf(\"*XtMalloc failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // cufftXtMemcpy() - Copy the data from host to device\n result = cufftXtMemcpy(planComplex, d_f, h_f, CUFFT_COPY_HOST_TO_DEVICE);\n if (result != CUFFT_SUCCESS) {\n printf(\"*XtMemcpy failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // cufftXtExecDescriptorC2C() - Execute FFT on data on multiple GPUs\n printf(\"Forward 2d FFT on multiple GPUs\\n\");\n result = cufftXtExecDescriptorC2C(planComplex, d_f, d_f, CUFFT_FORWARD);\n if (result != CUFFT_SUCCESS) {\n printf(\"*XtExecC2C failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // cufftXtMemcpy() - Copy the data to natural order on GPUs\n result = cufftXtMemcpy(planComplex, d_d_f, d_f, CUFFT_COPY_DEVICE_TO_DEVICE);\n if (result != CUFFT_SUCCESS) {\n printf(\"*XtMemcpy failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n printf(\"Solve Poisson Equation\\n\");\n solvePoissonEquation(d_d_f, d_out, d_k, N, nGPUs);\n\n printf(\"Inverse 2d FFT on multiple GPUs\\n\");\n // cufftXtExecDescriptorC2C() - Execute inverse FFT on data on multiple GPUs\n result = cufftXtExecDescriptorC2C(planComplex, d_out, d_out, CUFFT_INVERSE);\n if (result != CUFFT_SUCCESS) {\n printf(\"*XtExecC2C failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // Create a variable on host to copy the data from device\n // h_d_out - variable store the output of device\n Complex *h_d_out = (Complex *)malloc(sizeof(Complex) * N * N);\n\n // cufftXtMemcpy() - Copy data from multiple GPUs to host\n result = cufftXtMemcpy(planComplex, h_d_out, d_out, CUFFT_COPY_DEVICE_TO_HOST);\n if (result != CUFFT_SUCCESS) {\n printf(\"*XtMemcpy failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n float *out = (float *)malloc(sizeof(float) * N * N);\n float constant = h_d_out[0].x / N * N;\n for (int i = 0; i < N * N; i++) {\n // subtract u[0] to force the arbitrary constant to be 0\n out[i] = (h_d_out[i].x / (N * N)) - constant;\n }\n\n // cleanup memory\n\n free(h_f);\n free(k);\n free(out);\n free(h_d_out);\n free(x);\n free(whichGPUs);\n free(y);\n free(f);\n free(u_a);\n free(worksize);\n\n // cudaXtFree() - Free GPU memory\n for (int i = 0; i < GPU_COUNT; i++) {\n cudaFree(d_k[i]);\n }\n result = cufftXtFree(d_out);\n if (result != CUFFT_SUCCESS) {\n printf(\"*XtFree failed\\n\");\n exit(EXIT_FAILURE);\n }\n result = cufftXtFree(d_f);\n if (result != CUFFT_SUCCESS) {\n printf(\"*XtFree failed\\n\");\n exit(EXIT_FAILURE);\n }\n result = cufftXtFree(d_d_f);\n if (result != CUFFT_SUCCESS) {\n printf(\"*XtFree failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // cufftDestroy() - Destroy FFT plan\n result = cufftDestroy(planComplex);\n if (result != CUFFT_SUCCESS) {\n printf(\"cufftDestroy failed: code %d\\n\", (int)result);\n exit(EXIT_FAILURE);\n }\n\n exit(EXIT_SUCCESS);\n}\n\n////////////////////////////////////////////////////////////////////////////////////\n// Launch kernel on multiple GPU\n///////////////////////////////////////////////////////////////////////////////////\nvoid solvePoissonEquation(cudaLibXtDesc *d_ft, cudaLibXtDesc *d_ft_k, float **k, int N, int nGPUs)\n{\n int device;\n dim3 dimGrid(int(N / BSZ_X), int((N / 2) / BSZ_Y));\n dim3 dimBlock(BSZ_X, BSZ_Y);\n\n for (int i = 0; i < nGPUs; i++) {\n device = d_ft_k->descriptor->GPUs[i];\n cudaSetDevice(device);\n solvePoisson<<>>(\n (cufftComplex *)d_ft->descriptor->data[i], (cufftComplex *)d_ft_k->descriptor->data[i], k[i], N, i, nGPUs);\n }\n\n // Wait for device to finish all operation\n for (int i = 0; i < nGPUs; i++) {\n device = d_ft_k->descriptor->GPUs[i];\n cudaSetDevice(device);\n cudaDeviceSynchronize();\n\n // Check if kernel execution generated and error\n getLastCudaError(\"Kernel execution failed [ solvePoisson ]\");\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Kernel for Solving Poisson equation on GPU\n////////////////////////////////////////////////////////////////////////////////\n__global__ void solvePoisson(cufftComplex *ft, cufftComplex *ft_k, float *k, int N, int gpu_id, int n_gpu)\n{\n int i = threadIdx.x + blockIdx.x * blockDim.x;\n int j = threadIdx.y + blockIdx.y * blockDim.y;\n int index = j * N + i;\n if (i < N && j < N / n_gpu) {\n float k2 = k[i] * k[i] + k[j + gpu_id * N / n_gpu] * k[j + gpu_id * N / n_gpu];\n if (i == 0 && j == 0 && gpu_id == 0) {\n k2 = 1.0f;\n }\n\n ft_k[index].x = -ft[index].x * 1 / k2;\n ft_k[index].y = -ft[index].y * 1 / k2;\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n////////////////////////////////////////////////////////////////////////////////\n//\n// simpleCUFFT_2d_MGPU.cu\n//\n// This sample code demonstrate the use of CUFFT library for 2D data on multiple GPU.\n// Example showing the use of CUFFT for solving 2D-POISSON equation using FFT on multiple GPU.\n// For reference we have used the equation given in http://www.bu.edu/pasi/files/2011/07/\n// Lecture83.pdf\n//\n////////////////////////////////////////////////////////////////////////////////\n\n\n// System includes\n#include \n#include \n#include \n#include \n\n// CUDA runtime\n#include \n\n// CUFFT Header file\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Complex data type\ntypedef float2 Complex;\n\n// Data configuration\nconst int GPU_COUNT = 2;\nconst int BSZ_Y = 4;\nconst int BSZ_X = 4;\n\n// Forward Declaration\nvoid solvePoissonEquation(hipLibXtDesc *, hipLibXtDesc *, float **, int, int);\n\n__global__ void solvePoisson(hipfftComplex *, hipfftComplex *, float *, int, int, int n_gpu);\n\n///////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n printf(\"\\nPoisson equation using CUFFT library on Multiple GPUs is \"\n \"starting...\\n\\n\");\n\n int GPU_N;\n checkCudaErrors(hipGetDeviceCount(&GPU_N));\n\n if (GPU_N < GPU_COUNT) {\n printf(\"No. of GPU on node %d\\n\", GPU_N);\n printf(\"Two GPUs are required to run simpleCUFFT_2d_MGPU sample code\\n\");\n exit(EXIT_WAIVED);\n }\n\n int *major_minor = (int *)malloc(sizeof(int) * GPU_N * 2);\n int found2IdenticalGPUs = 0;\n int nGPUs = 2;\n int *whichGPUs;\n whichGPUs = (int *)malloc(sizeof(int) * nGPUs);\n\n for (int i = 0; i < GPU_N; i++) {\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, i));\n major_minor[i * 2] = deviceProp.major;\n major_minor[i * 2 + 1] = deviceProp.minor;\n printf(\"GPU Device %d: \\\"%s\\\" with compute capability %d.%d\\n\",\n i,\n deviceProp.name,\n deviceProp.major,\n deviceProp.minor);\n }\n\n for (int i = 0; i < GPU_N; i++) {\n for (int j = i + 1; j < GPU_N; j++) {\n if ((major_minor[i * 2] == major_minor[j * 2]) && (major_minor[i * 2 + 1] == major_minor[j * 2 + 1])) {\n whichGPUs[0] = i;\n whichGPUs[1] = j;\n found2IdenticalGPUs = 1;\n break;\n }\n }\n if (found2IdenticalGPUs) {\n break;\n }\n }\n\n free(major_minor);\n if (!found2IdenticalGPUs) {\n printf(\"No Two GPUs with same architecture found\\nWaiving simpleCUFFT_2d_MGPU \"\n \"sample\\n\");\n exit(EXIT_WAIVED);\n }\n\n int N = 64;\n float xMAX = 1.0f, xMIN = 0.0f, yMIN = 0.0f, h = (xMAX - xMIN) / ((float)N), s = 0.1f, s2 = s * s;\n float *x, *y, *f, *u_a, r2;\n\n x = (float *)malloc(sizeof(float) * N * N);\n y = (float *)malloc(sizeof(float) * N * N);\n f = (float *)malloc(sizeof(float) * N * N);\n u_a = (float *)malloc(sizeof(float) * N * N);\n\n for (int j = 0; j < N; j++)\n for (int i = 0; i < N; i++) {\n x[N * j + i] = xMIN + i * h;\n y[N * j + i] = yMIN + j * h;\n r2 = (x[N * j + i] - 0.5f) * (x[N * j + i] - 0.5f) + (y[N * j + i] - 0.5f) * (y[N * j + i] - 0.5f);\n f[N * j + i] = (r2 - 2 * s2) / (s2 * s2) * exp(-r2 / (2 * s2));\n u_a[N * j + i] = exp(-r2 / (2 * s2)); // analytical solution\n }\n\n float *k, *d_k[GPU_COUNT];\n k = (float *)malloc(sizeof(float) * N);\n for (int i = 0; i <= N / 2; i++) {\n k[i] = i * 2 * (float)M_PI;\n }\n for (int i = N / 2 + 1; i < N; i++) {\n k[i] = (i - N) * 2 * (float)M_PI;\n }\n\n // Create a complex variable on host\n Complex *h_f = (Complex *)malloc(sizeof(Complex) * N * N);\n\n // Initialize the memory for the signal\n for (int i = 0; i < (N * N); i++) {\n h_f[i].x = f[i];\n h_f[i].y = 0.0f;\n }\n\n // cufftCreate() - Create an empty plan\n hipfftResult result;\n hipfftHandle planComplex;\n result = hipfftCreate(&planComplex);\n if (result != HIPFFT_SUCCESS) {\n printf(\"hipfftCreate failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // cufftXtSetGPUs() - Define which GPUs to use\n result = hipfftXtSetGPUs(planComplex, nGPUs, whichGPUs);\n\n if (result == HIPFFT_INVALID_DEVICE) {\n printf(\"This sample requires two GPUs on the same board.\\n\");\n printf(\"No such board was found. Waiving sample.\\n\");\n exit(EXIT_WAIVED);\n }\n else if (result != HIPFFT_SUCCESS) {\n printf(\"hipfftXtSetGPUs failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // Print the device information to run the code\n printf(\"\\nRunning on GPUs\\n\");\n for (int i = 0; i < 2; i++) {\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, whichGPUs[i]));\n printf(\"GPU Device %d: \\\"%s\\\" with compute capability %d.%d\\n\",\n whichGPUs[i],\n deviceProp.name,\n deviceProp.major,\n deviceProp.minor);\n }\n\n size_t *worksize;\n worksize = (size_t *)malloc(sizeof(size_t) * nGPUs);\n\n // cufftMakePlan2d() - Create the plan\n result = hipfftMakePlan2d(planComplex, N, N, HIPFFT_C2C, worksize);\n if (result != HIPFFT_SUCCESS) {\n printf(\"*MakePlan* failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n for (int i = 0; i < nGPUs; i++) {\n hipSetDevice(whichGPUs[i]);\n hipMalloc((void **)&d_k[i], sizeof(float) * N);\n hipMemcpy(d_k[i], k, sizeof(float) * N, hipMemcpyHostToDevice);\n }\n\n // Create a variable on device\n // d_f - variable on device to store the input data\n // d_d_f - variable that store the natural order of d_f data\n // d_out - device output\n hipLibXtDesc *d_f, *d_d_f, *d_out;\n\n // cufftXtMalloc() - Malloc data on multiple GPUs\n\n result = hipfftXtMalloc(planComplex, (hipLibXtDesc **)&d_f, HIPFFT_XT_FORMAT_INPLACE);\n if (result != HIPFFT_SUCCESS) {\n printf(\"*XtMalloc failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n result = hipfftXtMalloc(planComplex, (hipLibXtDesc **)&d_d_f, HIPFFT_XT_FORMAT_INPLACE);\n if (result != HIPFFT_SUCCESS) {\n printf(\"*XtMalloc failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n result = hipfftXtMalloc(planComplex, (hipLibXtDesc **)&d_out, HIPFFT_XT_FORMAT_INPLACE);\n if (result != HIPFFT_SUCCESS) {\n printf(\"*XtMalloc failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // cufftXtMemcpy() - Copy the data from host to device\n result = hipfftXtMemcpy(planComplex, d_f, h_f, HIPFFT_COPY_HOST_TO_DEVICE);\n if (result != HIPFFT_SUCCESS) {\n printf(\"*XtMemcpy failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // cufftXtExecDescriptorC2C() - Execute FFT on data on multiple GPUs\n printf(\"Forward 2d FFT on multiple GPUs\\n\");\n result = hipfftXtExecDescriptorC2C(planComplex, d_f, d_f, HIPFFT_FORWARD);\n if (result != HIPFFT_SUCCESS) {\n printf(\"*XtExecC2C failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // cufftXtMemcpy() - Copy the data to natural order on GPUs\n result = hipfftXtMemcpy(planComplex, d_d_f, d_f, HIPFFT_COPY_DEVICE_TO_DEVICE);\n if (result != HIPFFT_SUCCESS) {\n printf(\"*XtMemcpy failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n printf(\"Solve Poisson Equation\\n\");\n solvePoissonEquation(d_d_f, d_out, d_k, N, nGPUs);\n\n printf(\"Inverse 2d FFT on multiple GPUs\\n\");\n // cufftXtExecDescriptorC2C() - Execute inverse FFT on data on multiple GPUs\n result = hipfftXtExecDescriptorC2C(planComplex, d_out, d_out, HIPFFT_BACKWARD);\n if (result != HIPFFT_SUCCESS) {\n printf(\"*XtExecC2C failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // Create a variable on host to copy the data from device\n // h_d_out - variable store the output of device\n Complex *h_d_out = (Complex *)malloc(sizeof(Complex) * N * N);\n\n // cufftXtMemcpy() - Copy data from multiple GPUs to host\n result = hipfftXtMemcpy(planComplex, h_d_out, d_out, HIPFFT_COPY_DEVICE_TO_HOST);\n if (result != HIPFFT_SUCCESS) {\n printf(\"*XtMemcpy failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n float *out = (float *)malloc(sizeof(float) * N * N);\n float constant = h_d_out[0].x / N * N;\n for (int i = 0; i < N * N; i++) {\n // subtract u[0] to force the arbitrary constant to be 0\n out[i] = (h_d_out[i].x / (N * N)) - constant;\n }\n\n // cleanup memory\n\n free(h_f);\n free(k);\n free(out);\n free(h_d_out);\n free(x);\n free(whichGPUs);\n free(y);\n free(f);\n free(u_a);\n free(worksize);\n\n // cudaXtFree() - Free GPU memory\n for (int i = 0; i < GPU_COUNT; i++) {\n hipFree(d_k[i]);\n }\n result = hipfftXtFree(d_out);\n if (result != HIPFFT_SUCCESS) {\n printf(\"*XtFree failed\\n\");\n exit(EXIT_FAILURE);\n }\n result = hipfftXtFree(d_f);\n if (result != HIPFFT_SUCCESS) {\n printf(\"*XtFree failed\\n\");\n exit(EXIT_FAILURE);\n }\n result = hipfftXtFree(d_d_f);\n if (result != HIPFFT_SUCCESS) {\n printf(\"*XtFree failed\\n\");\n exit(EXIT_FAILURE);\n }\n\n // cufftDestroy() - Destroy FFT plan\n result = hipfftDestroy(planComplex);\n if (result != HIPFFT_SUCCESS) {\n printf(\"hipfftDestroy failed: code %d\\n\", (int)result);\n exit(EXIT_FAILURE);\n }\n\n exit(EXIT_SUCCESS);\n}\n\n////////////////////////////////////////////////////////////////////////////////////\n// Launch kernel on multiple GPU\n///////////////////////////////////////////////////////////////////////////////////\nvoid solvePoissonEquation(hipLibXtDesc *d_ft, hipLibXtDesc *d_ft_k, float **k, int N, int nGPUs)\n{\n int device;\n dim3 dimGrid(int(N / BSZ_X), int((N / 2) / BSZ_Y));\n dim3 dimBlock(BSZ_X, BSZ_Y);\n\n for (int i = 0; i < nGPUs; i++) {\n device = d_ft_k->descriptor->GPUs[i];\n hipSetDevice(device);\n solvePoisson<<>>(\n (hipfftComplex *)d_ft->descriptor->data[i], (hipfftComplex *)d_ft_k->descriptor->data[i], k[i], N, i, nGPUs);\n }\n\n // Wait for device to finish all operation\n for (int i = 0; i < nGPUs; i++) {\n device = d_ft_k->descriptor->GPUs[i];\n hipSetDevice(device);\n hipDeviceSynchronize();\n\n // Check if kernel execution generated and error\n getLastCudaError(\"Kernel execution failed [ solvePoisson ]\");\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Kernel for Solving Poisson equation on GPU\n////////////////////////////////////////////////////////////////////////////////\n__global__ void solvePoisson(hipfftComplex *ft, hipfftComplex *ft_k, float *k, int N, int gpu_id, int n_gpu)\n{\n int i = threadIdx.x + blockIdx.x * blockDim.x;\n int j = threadIdx.y + blockIdx.y * blockDim.y;\n int index = j * N + i;\n if (i < N && j < N / n_gpu) {\n float k2 = k[i] * k[i] + k[j + gpu_id * N / n_gpu] * k[j + gpu_id * N / n_gpu];\n if (i == 0 && j == 0 && gpu_id == 0) {\n k2 = 1.0f;\n }\n\n ft_k[index].x = -ft[index].x * 1 / k2;\n ft_k[index].y = -ft[index].y * 1 / k2;\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/4_CUDA_Libraries/simpleCUFFT_2d_MGPU/simpleCUFFT_2d_MGPU.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"commonKernels.hpp\"\n\n__global__ void spinWhileLessThanOne(volatile unsigned int *latch)\n{\n while (latch[0] < 1)\n ;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \"commonKernels.hpp\"\n\n__global__ void spinWhileLessThanOne(volatile unsigned int *latch)\n{\n while (latch[0] < 1)\n ;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/6_Performance/UnifiedMemoryPerf/commonKernels.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n\n#include \"commonDefs.hpp\"\n#include \"commonKernels.hpp\"\n\n#define VERIFY_GPU_CORRECTNESS 0\n\nsize_t maxSampleSizeInMb = 64;\nint numKernelRuns = 20;\nint verboseResults = 0;\n\nconst char *memAllocTypeStr[MEMALLOC_TYPE_COUNT] = {\"Managed_Memory_With_Hints\",\n \"Managed_Memory_With_Hints_FullyAsync\",\n \"Managed_Memory_NoHints\",\n \"Zero_Copy\",\n \"Memcpy_HostMalloc_DeviceCudaMalloc\",\n \"MemcpyAsync_HostMalloc_DeviceCudaMalloc\",\n \"Memcpy_HostCudaHostAlloc_DeviceCudaMalloc\",\n \"MemcpyAsync_HostCudaHostAlloc_DeviceCudaMalloc\"};\n\nconst char *memAllocTypeShortStr[MEMALLOC_TYPE_COUNT] = {\n \"UMhint\", // Managed Memory With Hints\n \"UMhntAs\", // Managed Memory With_Hints Async\n \"UMeasy\", // Managed_Memory with No Hints\n \"0Copy\", // Zero Copy\n \"MemCopy\", // USE HOST PAGEABLE AND DEVICE_MEMORY\n \"CpAsync\", // USE HOST PAGEABLE AND DEVICE_MEMORY ASYNC\n \"CpHpglk\", // USE HOST PAGELOCKED AND DEVICE MEMORY\n \"CpPglAs\" // USE HOST PAGELOCKED AND DEVICE MEMORY ASYNC\n};\n\nstatic float RandFloat(float low, float high)\n{\n float t = (float)rand() / (float)RAND_MAX;\n return (1.0f - t) * low + t * high;\n}\n\nvoid fillMatrixWithRandomValues(float *matrix, unsigned int matrixDim)\n{\n unsigned int i, j;\n for (i = 0; i < matrixDim; ++i) {\n for (j = 0; j < matrixDim; ++j) {\n matrix[j + i * matrixDim] = RandFloat(0.0f, 10.0f);\n }\n }\n}\n\n#if VERIFY_GPU_CORRECTNESS\nvoid verifyMatrixMultiplyCorrectness(float *C, float *A, float *B, unsigned int matrixDim)\n{\n unsigned int i, j, k, numErrors = 0;\n for (i = 0; i < matrixDim; ++i) {\n for (j = 0; j < matrixDim; ++j) {\n float result = 0.0f;\n for (k = 0; k < matrixDim; ++k) {\n result += A[k + i * matrixDim] * B[j + k * matrixDim];\n }\n if (fabs(C[j + i * matrixDim] - result) > 0.001 * matrixDim) {\n printf(\"At [%u, %u]: Expected %f, Found %f\\n\", i, j, result, C[j + i * matrixDim]);\n ++numErrors;\n }\n }\n }\n if (numErrors != 0) {\n printf(\"%d value mismatches occured\\n\", numErrors);\n fflush(stdout);\n exit(EXIT_FAILURE); // exit since value mismatches occured\n }\n}\n#endif\n\nvoid copyMatrix(float *dstMatrix, float *srcMatrix, unsigned int matrixDim)\n{\n size_t size = matrixDim * matrixDim * sizeof(float);\n memcpy(dstMatrix, srcMatrix, size);\n}\n\nvoid verifyMatrixData(float *expectedData, float *observedData, unsigned int matrixDim)\n{\n unsigned int i, j, numErrors = 0;\n for (i = 0; i < matrixDim; ++i) {\n for (j = 0; j < matrixDim; ++j) {\n if (expectedData[j + i * matrixDim] != observedData[j + i * matrixDim]) {\n ++numErrors;\n if (verboseResults) {\n printf(\"At [%u, %u]: Expected %f, Found %f\\n\",\n i,\n j,\n expectedData[j + i * matrixDim],\n observedData[j + i * matrixDim]);\n }\n }\n }\n }\n if (numErrors != 0) {\n printf(\"%d value mismatches occured\\n\", numErrors);\n fflush(stdout);\n exit(EXIT_FAILURE); // exit since value mismatches occured\n }\n}\n\n#define BLOCK_SIZE 32\n__global__ void matrixMultiplyKernel(float *C, float *A, float *B, unsigned int matrixDim)\n{\n // Block index\n int bx = blockIdx.x;\n int by = blockIdx.y;\n\n // Thread index\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n\n unsigned int wA = matrixDim;\n unsigned int wB = matrixDim;\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = matrixDim * BLOCK_SIZE * by;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * bx;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n // Csub is used to store the element of the block sub-matrix\n // that is computed by the thread\n float Csub = 0;\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n // Load the matrices from device memory\n // to shared memory; each thread loads\n // one element of each matrix\n As[ty][tx] = A[a + wA * ty + tx];\n Bs[ty][tx] = B[b + wB * ty + tx];\n\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n // Multiply the two matrices together;\n // each thread computes one element\n // of the block sub-matrix\n#pragma unroll\n\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[ty][k] * Bs[k][tx];\n }\n\n // Synchronize to make sure that the preceding\n // computation is done before loading two new\n // sub-matrices of A and B in the next iteration\n __syncthreads();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes one element\n int c = wB * BLOCK_SIZE * by + BLOCK_SIZE * bx;\n C[c + wB * ty + tx] = Csub;\n}\n\nvoid runMatrixMultiplyKernel(unsigned int matrixDim,\n int allocType,\n unsigned int numLoops,\n double *gpuLaunchCallsTimes,\n double *gpuTransferToCallsTimes,\n double *gpuTransferFromCallsTimes,\n double *gpuLaunchAndTransferCallsTimes,\n double *gpuLaunchTransferSyncTimes,\n double *cpuAccessTimes,\n double *overallTimes,\n int device_id)\n{\n float *dptrA = NULL, *hptrA = NULL;\n float *dptrB = NULL, *hptrB = NULL;\n float *dptrC = NULL, *hptrC = NULL;\n float *randValuesX = NULL, *randValuesY = NULL;\n float *randValuesVerifyXmulY = NULL, *randValuesVerifyYmulX = NULL;\n bool copyRequired = false, hintsRequired = false;\n bool someTransferOpRequired;\n bool isAsync = false;\n cudaStream_t streamToRunOn;\n unsigned int *latch;\n size_t size = matrixDim * matrixDim * sizeof(float);\n dim3 threads(32, 32);\n dim3 grid(matrixDim / threads.x, matrixDim / threads.y);\n StopWatchInterface *gpuLaunchCallsTimer = 0, *gpuTransferCallsTimer = 0;\n StopWatchInterface *gpuSyncTimer = 0, *cpuAccessTimer = 0;\n sdkCreateTimer(&gpuLaunchCallsTimer);\n sdkCreateTimer(&gpuTransferCallsTimer);\n sdkCreateTimer(&gpuSyncTimer);\n sdkCreateTimer(&cpuAccessTimer);\n unsigned int i;\n\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, device_id));\n checkCudaErrors(cudaStreamCreate(&streamToRunOn));\n\n randValuesX = (float *)malloc(size);\n if (!randValuesX) {\n exit(EXIT_FAILURE); // exit since memory allocation error\n }\n randValuesY = (float *)malloc(size);\n if (!randValuesY) {\n exit(EXIT_FAILURE); // exit since memory allocation error\n }\n randValuesVerifyXmulY = (float *)malloc(size);\n if (!randValuesVerifyXmulY) {\n exit(EXIT_FAILURE); // exit since memory allocation error\n }\n randValuesVerifyYmulX = (float *)malloc(size);\n if (!randValuesVerifyYmulX) {\n exit(EXIT_FAILURE); // exit since memory allocation error\n }\n checkCudaErrors(cudaMalloc(&dptrA, size));\n checkCudaErrors(cudaMalloc(&dptrB, size));\n checkCudaErrors(cudaMalloc(&dptrC, size));\n\n fillMatrixWithRandomValues(randValuesX, matrixDim);\n fillMatrixWithRandomValues(randValuesY, matrixDim);\n\n checkCudaErrors(cudaMemcpyAsync(dptrA, randValuesX, size, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpyAsync(dptrB, randValuesY, size, cudaMemcpyHostToDevice));\n matrixMultiplyKernel<<>>(dptrC, dptrA, dptrB, matrixDim);\n checkCudaErrors(cudaMemcpyAsync(randValuesVerifyXmulY, dptrC, size, cudaMemcpyDeviceToHost));\n checkCudaErrors(cudaStreamSynchronize(NULL));\n matrixMultiplyKernel<<>>(dptrC, dptrB, dptrA, matrixDim);\n checkCudaErrors(cudaMemcpyAsync(randValuesVerifyYmulX, dptrC, size, cudaMemcpyDeviceToHost));\n checkCudaErrors(cudaStreamSynchronize(NULL));\n#if VERIFY_GPU_CORRECTNESS\n verifyMatrixMultiplyCorrectness(randValuesVerifyXmulY, randValuesX, randValuesY, matrixDim);\n verifyMatrixMultiplyCorrectness(randValuesVerifyYmulX, randValuesY, randValuesX, matrixDim);\n#endif\n checkCudaErrors(cudaFree(dptrA));\n checkCudaErrors(cudaFree(dptrB));\n checkCudaErrors(cudaFree(dptrC));\n\n checkCudaErrors(cudaMallocHost(&latch, sizeof(unsigned int)));\n\n switch (allocType) {\n case USE_HOST_PAGEABLE_AND_DEVICE_MEMORY:\n case USE_HOST_PAGEABLE_AND_DEVICE_MEMORY_ASYNC:\n hptrA = (float *)malloc(size);\n if (!hptrA) {\n exit(EXIT_FAILURE); // exit since memory allocation error\n }\n hptrB = (float *)malloc(size);\n if (!hptrB) {\n exit(EXIT_FAILURE); // exit since memory allocation error\n }\n hptrC = (float *)malloc(size);\n if (!hptrC) {\n exit(EXIT_FAILURE); // exit since memory allocation error\n }\n checkCudaErrors(cudaMalloc(&dptrA, size));\n checkCudaErrors(cudaMalloc(&dptrB, size));\n checkCudaErrors(cudaMalloc(&dptrC, size));\n copyRequired = true;\n break;\n\n case USE_HOST_PAGELOCKED_AND_DEVICE_MEMORY:\n case USE_HOST_PAGELOCKED_AND_DEVICE_MEMORY_ASYNC:\n checkCudaErrors(cudaMallocHost(&hptrA, size));\n checkCudaErrors(cudaMallocHost(&hptrB, size));\n checkCudaErrors(cudaMallocHost(&hptrC, size));\n checkCudaErrors(cudaMalloc(&dptrA, size));\n checkCudaErrors(cudaMalloc(&dptrB, size));\n checkCudaErrors(cudaMalloc(&dptrC, size));\n copyRequired = true;\n break;\n\n case USE_ZERO_COPY:\n checkCudaErrors(cudaMallocHost(&hptrA, size));\n checkCudaErrors(cudaMallocHost(&hptrB, size));\n checkCudaErrors(cudaMallocHost(&hptrC, size));\n checkCudaErrors(cudaHostGetDevicePointer(&dptrA, hptrA, 0));\n checkCudaErrors(cudaHostGetDevicePointer(&dptrB, hptrB, 0));\n checkCudaErrors(cudaHostGetDevicePointer(&dptrC, hptrC, 0));\n break;\n\n case USE_MANAGED_MEMORY:\n checkCudaErrors(cudaMallocManaged(&dptrA, size));\n checkCudaErrors(cudaMallocManaged(&dptrB, size));\n checkCudaErrors(cudaMallocManaged(&dptrC, size));\n hptrA = dptrA;\n hptrB = dptrB;\n hptrC = dptrC;\n break;\n\n case USE_MANAGED_MEMORY_WITH_HINTS:\n case USE_MANAGED_MEMORY_WITH_HINTS_ASYNC:\n if (deviceProp.concurrentManagedAccess) {\n checkCudaErrors(cudaMallocManaged(&dptrA, size));\n checkCudaErrors(cudaMallocManaged(&dptrB, size));\n checkCudaErrors(cudaMallocManaged(&dptrC, size));\n cudaMemLocation hostLoc;\n hostLoc.type = cudaMemLocationTypeHost;\n checkCudaErrors(cudaMemPrefetchAsync(dptrA, size, hostLoc, 0));\n checkCudaErrors(cudaMemPrefetchAsync(dptrB, size, hostLoc, 0));\n checkCudaErrors(cudaMemPrefetchAsync(dptrC, size, hostLoc, 0));\n }\n else {\n checkCudaErrors(cudaMallocManaged(&dptrA, size, cudaMemAttachHost));\n checkCudaErrors(cudaMallocManaged(&dptrB, size, cudaMemAttachHost));\n checkCudaErrors(cudaMallocManaged(&dptrC, size, cudaMemAttachHost));\n }\n hptrA = dptrA;\n hptrB = dptrB;\n hptrC = dptrC;\n hintsRequired = true;\n break;\n\n default:\n exit(EXIT_FAILURE); // exit with error\n }\n\n if (allocType == USE_HOST_PAGEABLE_AND_DEVICE_MEMORY_ASYNC\n || allocType == USE_HOST_PAGELOCKED_AND_DEVICE_MEMORY_ASYNC\n || allocType == USE_MANAGED_MEMORY_WITH_HINTS_ASYNC) {\n isAsync = true;\n }\n\n someTransferOpRequired = copyRequired || hintsRequired;\n\n // fill buffers with 0 to avoid any first access page-fault overheads.\n memset(hptrA, 0, size);\n memset(hptrB, 0, size);\n memset(hptrC, 0, size);\n\n for (i = 0; i < numLoops; i++) {\n cpuAccessTimes[i] = 0.0;\n gpuLaunchCallsTimes[i] = 0.0;\n gpuTransferToCallsTimes[i] = 0.0;\n gpuTransferFromCallsTimes[i] = 0.0;\n\n sdkStartTimer(&cpuAccessTimer);\n {\n copyMatrix(hptrA, (i & 0x1 == 0) ? randValuesX : randValuesY, matrixDim);\n copyMatrix(hptrB, (i & 0x1 == 0) ? randValuesY : randValuesX, matrixDim);\n }\n sdkStopTimer(&cpuAccessTimer);\n cpuAccessTimes[i] += sdkGetAverageTimerValue(&cpuAccessTimer);\n sdkResetTimer(&cpuAccessTimer);\n\n if (isAsync && hintsRequired) {\n *latch = 0;\n // Prevent any work on stream from starting until all work is pushed\n spinWhileLessThanOne<<<1, 1, 0, streamToRunOn>>>(latch);\n }\n\n if (someTransferOpRequired) {\n sdkStartTimer(&gpuTransferCallsTimer);\n if (copyRequired) {\n if (isAsync) {\n checkCudaErrors(cudaMemcpyAsync(dptrA, hptrA, size, cudaMemcpyHostToDevice, streamToRunOn));\n checkCudaErrors(cudaMemcpyAsync(dptrB, hptrB, size, cudaMemcpyHostToDevice, streamToRunOn));\n }\n else {\n checkCudaErrors(cudaMemcpy(dptrA, hptrA, size, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(dptrB, hptrB, size, cudaMemcpyHostToDevice));\n }\n }\n if (hintsRequired) {\n if (deviceProp.concurrentManagedAccess) {\n cudaMemLocation deviceLoc;\n deviceLoc.type = cudaMemLocationTypeDevice;\n deviceLoc.id = device_id;\n checkCudaErrors(cudaMemPrefetchAsync(dptrA, size, deviceLoc, 0, streamToRunOn));\n checkCudaErrors(cudaMemPrefetchAsync(dptrB, size, deviceLoc, 0, streamToRunOn));\n checkCudaErrors(cudaMemPrefetchAsync(dptrC, size, deviceLoc, 0, streamToRunOn));\n }\n else {\n checkCudaErrors(cudaStreamAttachMemAsync(streamToRunOn, dptrA, 0, cudaMemAttachGlobal));\n checkCudaErrors(cudaStreamAttachMemAsync(streamToRunOn, dptrB, 0, cudaMemAttachGlobal));\n checkCudaErrors(cudaStreamAttachMemAsync(streamToRunOn, dptrC, 0, cudaMemAttachGlobal));\n }\n if (!isAsync) {\n checkCudaErrors(cudaStreamSynchronize(streamToRunOn));\n }\n }\n\n sdkStopTimer(&gpuTransferCallsTimer);\n gpuTransferToCallsTimes[i] += sdkGetAverageTimerValue(&gpuTransferCallsTimer);\n sdkResetTimer(&gpuTransferCallsTimer);\n }\n\n sdkStartTimer(&gpuLaunchCallsTimer);\n {\n matrixMultiplyKernel<<>>(dptrC, dptrA, dptrB, matrixDim);\n if (!isAsync) {\n checkCudaErrors(cudaStreamSynchronize(streamToRunOn));\n }\n }\n sdkStopTimer(&gpuLaunchCallsTimer);\n\n gpuLaunchCallsTimes[i] += sdkGetAverageTimerValue(&gpuLaunchCallsTimer);\n sdkResetTimer(&gpuLaunchCallsTimer);\n\n if (someTransferOpRequired) {\n sdkStartTimer(&gpuTransferCallsTimer);\n if (hintsRequired) {\n if (deviceProp.concurrentManagedAccess) {\n cudaMemLocation hostLoc;\n hostLoc.type = cudaMemLocationTypeHost;\n checkCudaErrors(cudaMemPrefetchAsync(dptrA, size, hostLoc, 0));\n checkCudaErrors(cudaMemPrefetchAsync(dptrB, size, hostLoc, 0));\n checkCudaErrors(cudaMemPrefetchAsync(dptrC, size, hostLoc, 0));\n }\n else {\n checkCudaErrors(cudaStreamAttachMemAsync(streamToRunOn, dptrA, 0, cudaMemAttachHost));\n checkCudaErrors(cudaStreamAttachMemAsync(streamToRunOn, dptrB, 0, cudaMemAttachHost));\n checkCudaErrors(cudaStreamAttachMemAsync(streamToRunOn, dptrC, 0, cudaMemAttachHost));\n }\n if (!isAsync) {\n checkCudaErrors(cudaStreamSynchronize(streamToRunOn));\n }\n }\n if (copyRequired) {\n if (isAsync) {\n checkCudaErrors(cudaMemcpyAsync(hptrC, dptrC, size, cudaMemcpyDeviceToHost, streamToRunOn));\n }\n else {\n checkCudaErrors(cudaMemcpy(hptrC, dptrC, size, cudaMemcpyDeviceToHost));\n }\n }\n sdkStopTimer(&gpuTransferCallsTimer);\n gpuTransferFromCallsTimes[i] += sdkGetAverageTimerValue(&gpuTransferCallsTimer);\n sdkResetTimer(&gpuTransferCallsTimer);\n }\n gpuLaunchAndTransferCallsTimes[i] =\n gpuLaunchCallsTimes[i] + gpuTransferToCallsTimes[i] + gpuTransferFromCallsTimes[i];\n gpuLaunchTransferSyncTimes[i] = gpuLaunchAndTransferCallsTimes[i];\n if (isAsync) {\n sdkStartTimer(&gpuSyncTimer);\n {\n if (hintsRequired) {\n *latch = 1;\n }\n checkCudaErrors(cudaStreamSynchronize(streamToRunOn));\n }\n sdkStopTimer(&gpuSyncTimer);\n gpuLaunchTransferSyncTimes[i] += sdkGetAverageTimerValue(&gpuSyncTimer);\n sdkResetTimer(&gpuSyncTimer);\n }\n\n sdkStartTimer(&cpuAccessTimer);\n {\n verifyMatrixData((i & 0x1 == 0) ? randValuesVerifyXmulY : randValuesVerifyYmulX, hptrC, matrixDim);\n }\n sdkStopTimer(&cpuAccessTimer);\n cpuAccessTimes[i] += sdkGetAverageTimerValue(&cpuAccessTimer);\n sdkResetTimer(&cpuAccessTimer);\n overallTimes[i] = cpuAccessTimes[i] + gpuLaunchTransferSyncTimes[i];\n }\n\n switch (allocType) {\n case USE_HOST_PAGEABLE_AND_DEVICE_MEMORY:\n case USE_HOST_PAGEABLE_AND_DEVICE_MEMORY_ASYNC:\n free(hptrA);\n free(hptrB);\n free(hptrC);\n checkCudaErrors(cudaFree(dptrA));\n checkCudaErrors(cudaFree(dptrB));\n checkCudaErrors(cudaFree(dptrC));\n break;\n\n case USE_HOST_PAGELOCKED_AND_DEVICE_MEMORY:\n case USE_HOST_PAGELOCKED_AND_DEVICE_MEMORY_ASYNC:\n checkCudaErrors(cudaFreeHost(hptrA));\n checkCudaErrors(cudaFreeHost(hptrB));\n checkCudaErrors(cudaFreeHost(hptrC));\n checkCudaErrors(cudaFree(dptrA));\n checkCudaErrors(cudaFree(dptrB));\n checkCudaErrors(cudaFree(dptrC));\n break;\n\n case USE_ZERO_COPY:\n checkCudaErrors(cudaFreeHost(hptrA));\n checkCudaErrors(cudaFreeHost(hptrB));\n checkCudaErrors(cudaFreeHost(hptrC));\n break;\n\n case USE_MANAGED_MEMORY:\n case USE_MANAGED_MEMORY_WITH_HINTS:\n case USE_MANAGED_MEMORY_WITH_HINTS_ASYNC:\n checkCudaErrors(cudaFree(dptrA));\n checkCudaErrors(cudaFree(dptrB));\n checkCudaErrors(cudaFree(dptrC));\n break;\n\n default:\n exit(EXIT_FAILURE); // exit due to error\n }\n\n checkCudaErrors(cudaStreamDestroy(streamToRunOn));\n checkCudaErrors(cudaFreeHost(latch));\n free(randValuesX);\n free(randValuesY);\n free(randValuesVerifyXmulY);\n free(randValuesVerifyYmulX);\n sdkDeleteTimer(&gpuLaunchCallsTimer);\n sdkDeleteTimer(&gpuTransferCallsTimer);\n sdkDeleteTimer(&gpuSyncTimer);\n sdkDeleteTimer(&cpuAccessTimer);\n}\n\nvoid matrixMultiplyPerfRunner(bool reportAsBandwidth,\n bool print_launch_transfer_results,\n bool print_std_deviation,\n int device_id)\n{\n int i;\n unsigned int minMatrixDim = 32;\n unsigned int multiplierDim = 2;\n unsigned int matrixDim;\n unsigned int minSize = minMatrixDim * minMatrixDim * sizeof(float);\n unsigned int maxSize = (maxSampleSizeInMb * ONE_MB) / 4; // 3 buffers are used, but dividing by 4 (power of 2)\n unsigned int multiplier = multiplierDim * multiplierDim;\n unsigned int numSizesToTest;\n\n struct testResults *results;\n struct resultsData *gpuLaunchCallsTimes;\n struct resultsData *gpuTransferToCallsTimes;\n struct resultsData *gpuTransferFromCallsTimes;\n struct resultsData *gpuLaunchAndTransferCallsTimes;\n struct resultsData *gpuLaunchTransferSyncTimes;\n struct resultsData *cpuAccessTimes;\n struct resultsData *overallTimes;\n unsigned long *sizesToTest;\n unsigned int j;\n\n numSizesToTest = findNumSizesToTest(minSize, maxSize, multiplier);\n\n createAndInitTestResults(&results, \"matrixMultiplyPerf\", numKernelRuns, numSizesToTest);\n\n sizesToTest = getPtrSizesToTest(results);\n\n createResultDataAndAddToTestResults(\n &gpuLaunchCallsTimes, results, \"GPU Kernel Launch Call Time\", false, reportAsBandwidth);\n createResultDataAndAddToTestResults(\n &gpuTransferToCallsTimes, results, \"CPU to GPU Transfer Calls Time\", false, reportAsBandwidth);\n createResultDataAndAddToTestResults(\n &gpuTransferFromCallsTimes, results, \"GPU to CPU Transfer Calls Time\", false, reportAsBandwidth);\n createResultDataAndAddToTestResults(\n &gpuLaunchAndTransferCallsTimes, results, \"GPU Launch and Transfer Calls Time\", false, reportAsBandwidth);\n createResultDataAndAddToTestResults(\n &gpuLaunchTransferSyncTimes, results, \"GPU Launch Transfer and Sync Time\", false, reportAsBandwidth);\n createResultDataAndAddToTestResults(&cpuAccessTimes, results, \"CPU Access Time\", false, reportAsBandwidth);\n createResultDataAndAddToTestResults(&overallTimes, results, \"Overall Time\", false, reportAsBandwidth);\n\n printf(\"Running \");\n for (matrixDim = minMatrixDim, j = 0; matrixDim * matrixDim <= maxSize / sizeof(float);\n matrixDim *= multiplierDim, ++j) {\n sizesToTest[j] = matrixDim * matrixDim * sizeof(float);\n for (i = MEMALLOC_TYPE_START; i <= MEMALLOC_TYPE_END; i++) {\n printf(\".\");\n fflush(stdout);\n runMatrixMultiplyKernel(matrixDim,\n i,\n numKernelRuns,\n getPtrRunTimesInMs(gpuLaunchCallsTimes, i, j),\n getPtrRunTimesInMs(gpuTransferToCallsTimes, i, j),\n getPtrRunTimesInMs(gpuTransferFromCallsTimes, i, j),\n getPtrRunTimesInMs(gpuLaunchAndTransferCallsTimes, i, j),\n getPtrRunTimesInMs(gpuLaunchTransferSyncTimes, i, j),\n getPtrRunTimesInMs(cpuAccessTimes, i, j),\n getPtrRunTimesInMs(overallTimes, i, j),\n device_id);\n }\n }\n printf(\"\\n\");\n printResults(results, print_launch_transfer_results, print_std_deviation);\n freeTestResultsAndAllResultsData(results);\n}\n\nstatic void usage()\n{\n printf(\"./cudaMemoryTypesPerf [-device=] [-reportAsBandwidth] \"\n \"[-print-launch-transfer-results] [-print-std-deviation] [-verbose]\\n\");\n printf(\"Options:\\n\");\n printf(\"-reportAsBandwidth: By default time taken is printed, this \"\n \"option allows to instead print bandwidth.\\n\");\n printf(\"-print-launch-transfer-results: By default overall results are printed, \"\n \"this option allows to print data transfers and kernel time as well.\\n\");\n printf(\"-print-std-deviation: Prints std deviation of the results.\\n\");\n printf(\"-kernel-iterations=: Number of times the kernel tests should \"\n \"be run[default is 100 iterations].\\n\");\n printf(\"-device=: Allows to pass GPU Device ID on which \"\n \"the tests will be run.\\n\");\n printf(\"-verbose: Prints highly verbose output.\\n\");\n}\n\nint main(int argc, char **argv)\n{\n bool reportAsBandwidth = false;\n bool print_launch_transfer_results = false;\n bool print_std_deviation = false;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\") || checkCmdLineFlag(argc, (const char **)argv, \"h\")) {\n usage();\n printf(\"&&&& %s WAIVED\\n\", argv[0]);\n exit(EXIT_WAIVED);\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"reportAsBandwidth\")) {\n reportAsBandwidth = true;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"print-launch-transfer-results\")) {\n print_launch_transfer_results = true;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"print-std-deviation\")) {\n print_std_deviation = true;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"kernel-iterations\")) {\n numKernelRuns = getCmdLineArgumentInt(argc, (const char **)argv, \"kernel-iterations\");\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"verbose\")) {\n verboseResults = 1;\n }\n\n // set device\n cudaDeviceProp device_prop;\n int device_id = findCudaDevice(argc, (const char **)argv);\n checkCudaErrors(cudaGetDeviceProperties(&device_prop, device_id));\n\n if (!device_prop.managedMemory) {\n // This samples requires being run on a device that supports Unified Memory\n fprintf(stderr, \"Unified Memory not supported on this device\\n\");\n\n exit(EXIT_WAIVED);\n }\n\n matrixMultiplyPerfRunner(reportAsBandwidth, print_launch_transfer_results, print_std_deviation, device_id);\n\n printf(\"\\nNOTE: The CUDA Samples are not meant for performance measurements. \"\n \"Results may vary when GPU Boost is enabled.\\n\");\n exit(EXIT_SUCCESS);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n#include \"commonDefs.hpp\"\n#include \"commonKernels.hpp\"\n\n#define VERIFY_GPU_CORRECTNESS 0\n\nsize_t maxSampleSizeInMb = 64;\nint numKernelRuns = 20;\nint verboseResults = 0;\n\nconst char *memAllocTypeStr[MEMALLOC_TYPE_COUNT] = {\"Managed_Memory_With_Hints\",\n \"Managed_Memory_With_Hints_FullyAsync\",\n \"Managed_Memory_NoHints\",\n \"Zero_Copy\",\n \"Memcpy_HostMalloc_DeviceCudaMalloc\",\n \"MemcpyAsync_HostMalloc_DeviceCudaMalloc\",\n \"Memcpy_HostCudaHostAlloc_DeviceCudaMalloc\",\n \"MemcpyAsync_HostCudaHostAlloc_DeviceCudaMalloc\"};\n\nconst char *memAllocTypeShortStr[MEMALLOC_TYPE_COUNT] = {\n \"UMhint\", // Managed Memory With Hints\n \"UMhntAs\", // Managed Memory With_Hints Async\n \"UMeasy\", // Managed_Memory with No Hints\n \"0Copy\", // Zero Copy\n \"MemCopy\", // USE HOST PAGEABLE AND DEVICE_MEMORY\n \"CpAsync\", // USE HOST PAGEABLE AND DEVICE_MEMORY ASYNC\n \"CpHpglk\", // USE HOST PAGELOCKED AND DEVICE MEMORY\n \"CpPglAs\" // USE HOST PAGELOCKED AND DEVICE MEMORY ASYNC\n};\n\nstatic float RandFloat(float low, float high)\n{\n float t = (float)rand() / (float)RAND_MAX;\n return (1.0f - t) * low + t * high;\n}\n\nvoid fillMatrixWithRandomValues(float *matrix, unsigned int matrixDim)\n{\n unsigned int i, j;\n for (i = 0; i < matrixDim; ++i) {\n for (j = 0; j < matrixDim; ++j) {\n matrix[j + i * matrixDim] = RandFloat(0.0f, 10.0f);\n }\n }\n}\n\n#if VERIFY_GPU_CORRECTNESS\nvoid verifyMatrixMultiplyCorrectness(float *C, float *A, float *B, unsigned int matrixDim)\n{\n unsigned int i, j, k, numErrors = 0;\n for (i = 0; i < matrixDim; ++i) {\n for (j = 0; j < matrixDim; ++j) {\n float result = 0.0f;\n for (k = 0; k < matrixDim; ++k) {\n result += A[k + i * matrixDim] * B[j + k * matrixDim];\n }\n if (fabs(C[j + i * matrixDim] - result) > 0.001 * matrixDim) {\n printf(\"At [%u, %u]: Expected %f, Found %f\\n\", i, j, result, C[j + i * matrixDim]);\n ++numErrors;\n }\n }\n }\n if (numErrors != 0) {\n printf(\"%d value mismatches occured\\n\", numErrors);\n fflush(stdout);\n exit(EXIT_FAILURE); // exit since value mismatches occured\n }\n}\n#endif\n\nvoid copyMatrix(float *dstMatrix, float *srcMatrix, unsigned int matrixDim)\n{\n size_t size = matrixDim * matrixDim * sizeof(float);\n memcpy(dstMatrix, srcMatrix, size);\n}\n\nvoid verifyMatrixData(float *expectedData, float *observedData, unsigned int matrixDim)\n{\n unsigned int i, j, numErrors = 0;\n for (i = 0; i < matrixDim; ++i) {\n for (j = 0; j < matrixDim; ++j) {\n if (expectedData[j + i * matrixDim] != observedData[j + i * matrixDim]) {\n ++numErrors;\n if (verboseResults) {\n printf(\"At [%u, %u]: Expected %f, Found %f\\n\",\n i,\n j,\n expectedData[j + i * matrixDim],\n observedData[j + i * matrixDim]);\n }\n }\n }\n }\n if (numErrors != 0) {\n printf(\"%d value mismatches occured\\n\", numErrors);\n fflush(stdout);\n exit(EXIT_FAILURE); // exit since value mismatches occured\n }\n}\n\n#define BLOCK_SIZE 32\n__global__ void matrixMultiplyKernel(float *C, float *A, float *B, unsigned int matrixDim)\n{\n // Block index\n int bx = blockIdx.x;\n int by = blockIdx.y;\n\n // Thread index\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n\n unsigned int wA = matrixDim;\n unsigned int wB = matrixDim;\n\n // Index of the first sub-matrix of A processed by the block\n int aBegin = matrixDim * BLOCK_SIZE * by;\n\n // Index of the last sub-matrix of A processed by the block\n int aEnd = aBegin + wA - 1;\n\n // Step size used to iterate through the sub-matrices of A\n int aStep = BLOCK_SIZE;\n\n // Index of the first sub-matrix of B processed by the block\n int bBegin = BLOCK_SIZE * bx;\n\n // Step size used to iterate through the sub-matrices of B\n int bStep = BLOCK_SIZE * wB;\n\n // Csub is used to store the element of the block sub-matrix\n // that is computed by the thread\n float Csub = 0;\n\n // Loop over all the sub-matrices of A and B\n // required to compute the block sub-matrix\n for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep) {\n // Declaration of the shared memory array As used to\n // store the sub-matrix of A\n __shared__ float As[BLOCK_SIZE][BLOCK_SIZE];\n\n // Declaration of the shared memory array Bs used to\n // store the sub-matrix of B\n __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];\n\n // Load the matrices from device memory\n // to shared memory; each thread loads\n // one element of each matrix\n As[ty][tx] = A[a + wA * ty + tx];\n Bs[ty][tx] = B[b + wB * ty + tx];\n\n // Synchronize to make sure the matrices are loaded\n __syncthreads();\n\n // Multiply the two matrices together;\n // each thread computes one element\n // of the block sub-matrix\n#pragma unroll\n\n for (int k = 0; k < BLOCK_SIZE; ++k) {\n Csub += As[ty][k] * Bs[k][tx];\n }\n\n // Synchronize to make sure that the preceding\n // computation is done before loading two new\n // sub-matrices of A and B in the next iteration\n __syncthreads();\n }\n\n // Write the block sub-matrix to device memory;\n // each thread writes one element\n int c = wB * BLOCK_SIZE * by + BLOCK_SIZE * bx;\n C[c + wB * ty + tx] = Csub;\n}\n\nvoid runMatrixMultiplyKernel(unsigned int matrixDim,\n int allocType,\n unsigned int numLoops,\n double *gpuLaunchCallsTimes,\n double *gpuTransferToCallsTimes,\n double *gpuTransferFromCallsTimes,\n double *gpuLaunchAndTransferCallsTimes,\n double *gpuLaunchTransferSyncTimes,\n double *cpuAccessTimes,\n double *overallTimes,\n int device_id)\n{\n float *dptrA = NULL, *hptrA = NULL;\n float *dptrB = NULL, *hptrB = NULL;\n float *dptrC = NULL, *hptrC = NULL;\n float *randValuesX = NULL, *randValuesY = NULL;\n float *randValuesVerifyXmulY = NULL, *randValuesVerifyYmulX = NULL;\n bool copyRequired = false, hintsRequired = false;\n bool someTransferOpRequired;\n bool isAsync = false;\n hipStream_t streamToRunOn;\n unsigned int *latch;\n size_t size = matrixDim * matrixDim * sizeof(float);\n dim3 threads(32, 32);\n dim3 grid(matrixDim / threads.x, matrixDim / threads.y);\n StopWatchInterface *gpuLaunchCallsTimer = 0, *gpuTransferCallsTimer = 0;\n StopWatchInterface *gpuSyncTimer = 0, *cpuAccessTimer = 0;\n sdkCreateTimer(&gpuLaunchCallsTimer);\n sdkCreateTimer(&gpuTransferCallsTimer);\n sdkCreateTimer(&gpuSyncTimer);\n sdkCreateTimer(&cpuAccessTimer);\n unsigned int i;\n\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, device_id));\n checkCudaErrors(hipStreamCreate(&streamToRunOn));\n\n randValuesX = (float *)malloc(size);\n if (!randValuesX) {\n exit(EXIT_FAILURE); // exit since memory allocation error\n }\n randValuesY = (float *)malloc(size);\n if (!randValuesY) {\n exit(EXIT_FAILURE); // exit since memory allocation error\n }\n randValuesVerifyXmulY = (float *)malloc(size);\n if (!randValuesVerifyXmulY) {\n exit(EXIT_FAILURE); // exit since memory allocation error\n }\n randValuesVerifyYmulX = (float *)malloc(size);\n if (!randValuesVerifyYmulX) {\n exit(EXIT_FAILURE); // exit since memory allocation error\n }\n checkCudaErrors(hipMalloc(&dptrA, size));\n checkCudaErrors(hipMalloc(&dptrB, size));\n checkCudaErrors(hipMalloc(&dptrC, size));\n\n fillMatrixWithRandomValues(randValuesX, matrixDim);\n fillMatrixWithRandomValues(randValuesY, matrixDim);\n\n checkCudaErrors(hipMemcpyAsync(dptrA, randValuesX, size, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpyAsync(dptrB, randValuesY, size, hipMemcpyHostToDevice));\n matrixMultiplyKernel<<>>(dptrC, dptrA, dptrB, matrixDim);\n checkCudaErrors(hipMemcpyAsync(randValuesVerifyXmulY, dptrC, size, hipMemcpyDeviceToHost));\n checkCudaErrors(hipStreamSynchronize(NULL));\n matrixMultiplyKernel<<>>(dptrC, dptrB, dptrA, matrixDim);\n checkCudaErrors(hipMemcpyAsync(randValuesVerifyYmulX, dptrC, size, hipMemcpyDeviceToHost));\n checkCudaErrors(hipStreamSynchronize(NULL));\n#if VERIFY_GPU_CORRECTNESS\n verifyMatrixMultiplyCorrectness(randValuesVerifyXmulY, randValuesX, randValuesY, matrixDim);\n verifyMatrixMultiplyCorrectness(randValuesVerifyYmulX, randValuesY, randValuesX, matrixDim);\n#endif\n checkCudaErrors(hipFree(dptrA));\n checkCudaErrors(hipFree(dptrB));\n checkCudaErrors(hipFree(dptrC));\n\n checkCudaErrors(hipHostMalloc(&latch, sizeof(unsigned int)));\n\n switch (allocType) {\n case USE_HOST_PAGEABLE_AND_DEVICE_MEMORY:\n case USE_HOST_PAGEABLE_AND_DEVICE_MEMORY_ASYNC:\n hptrA = (float *)malloc(size);\n if (!hptrA) {\n exit(EXIT_FAILURE); // exit since memory allocation error\n }\n hptrB = (float *)malloc(size);\n if (!hptrB) {\n exit(EXIT_FAILURE); // exit since memory allocation error\n }\n hptrC = (float *)malloc(size);\n if (!hptrC) {\n exit(EXIT_FAILURE); // exit since memory allocation error\n }\n checkCudaErrors(hipMalloc(&dptrA, size));\n checkCudaErrors(hipMalloc(&dptrB, size));\n checkCudaErrors(hipMalloc(&dptrC, size));\n copyRequired = true;\n break;\n\n case USE_HOST_PAGELOCKED_AND_DEVICE_MEMORY:\n case USE_HOST_PAGELOCKED_AND_DEVICE_MEMORY_ASYNC:\n checkCudaErrors(hipHostMalloc(&hptrA, size));\n checkCudaErrors(hipHostMalloc(&hptrB, size));\n checkCudaErrors(hipHostMalloc(&hptrC, size));\n checkCudaErrors(hipMalloc(&dptrA, size));\n checkCudaErrors(hipMalloc(&dptrB, size));\n checkCudaErrors(hipMalloc(&dptrC, size));\n copyRequired = true;\n break;\n\n case USE_ZERO_COPY:\n checkCudaErrors(hipHostMalloc(&hptrA, size));\n checkCudaErrors(hipHostMalloc(&hptrB, size));\n checkCudaErrors(hipHostMalloc(&hptrC, size));\n checkCudaErrors(hipHostGetDevicePointer(&dptrA, hptrA, 0));\n checkCudaErrors(hipHostGetDevicePointer(&dptrB, hptrB, 0));\n checkCudaErrors(hipHostGetDevicePointer(&dptrC, hptrC, 0));\n break;\n\n case USE_MANAGED_MEMORY:\n checkCudaErrors(hipMallocManaged(&dptrA, size));\n checkCudaErrors(hipMallocManaged(&dptrB, size));\n checkCudaErrors(hipMallocManaged(&dptrC, size));\n hptrA = dptrA;\n hptrB = dptrB;\n hptrC = dptrC;\n break;\n\n case USE_MANAGED_MEMORY_WITH_HINTS:\n case USE_MANAGED_MEMORY_WITH_HINTS_ASYNC:\n if (deviceProp.concurrentManagedAccess) {\n checkCudaErrors(hipMallocManaged(&dptrA, size));\n checkCudaErrors(hipMallocManaged(&dptrB, size));\n checkCudaErrors(hipMallocManaged(&dptrC, size));\n hipMemLocation hostLoc;\n hostLoc.type = hipMemLocationTypeHost;\n checkCudaErrors(hipMemPrefetchAsync(dptrA, size, hostLoc, 0));\n checkCudaErrors(hipMemPrefetchAsync(dptrB, size, hostLoc, 0));\n checkCudaErrors(hipMemPrefetchAsync(dptrC, size, hostLoc, 0));\n }\n else {\n checkCudaErrors(hipMallocManaged(&dptrA, size, hipMemAttachHost));\n checkCudaErrors(hipMallocManaged(&dptrB, size, hipMemAttachHost));\n checkCudaErrors(hipMallocManaged(&dptrC, size, hipMemAttachHost));\n }\n hptrA = dptrA;\n hptrB = dptrB;\n hptrC = dptrC;\n hintsRequired = true;\n break;\n\n default:\n exit(EXIT_FAILURE); // exit with error\n }\n\n if (allocType == USE_HOST_PAGEABLE_AND_DEVICE_MEMORY_ASYNC\n || allocType == USE_HOST_PAGELOCKED_AND_DEVICE_MEMORY_ASYNC\n || allocType == USE_MANAGED_MEMORY_WITH_HINTS_ASYNC) {\n isAsync = true;\n }\n\n someTransferOpRequired = copyRequired || hintsRequired;\n\n // fill buffers with 0 to avoid any first access page-fault overheads.\n memset(hptrA, 0, size);\n memset(hptrB, 0, size);\n memset(hptrC, 0, size);\n\n for (i = 0; i < numLoops; i++) {\n cpuAccessTimes[i] = 0.0;\n gpuLaunchCallsTimes[i] = 0.0;\n gpuTransferToCallsTimes[i] = 0.0;\n gpuTransferFromCallsTimes[i] = 0.0;\n\n sdkStartTimer(&cpuAccessTimer);\n {\n copyMatrix(hptrA, (i & 0x1 == 0) ? randValuesX : randValuesY, matrixDim);\n copyMatrix(hptrB, (i & 0x1 == 0) ? randValuesY : randValuesX, matrixDim);\n }\n sdkStopTimer(&cpuAccessTimer);\n cpuAccessTimes[i] += sdkGetAverageTimerValue(&cpuAccessTimer);\n sdkResetTimer(&cpuAccessTimer);\n\n if (isAsync && hintsRequired) {\n *latch = 0;\n // Prevent any work on stream from starting until all work is pushed\n spinWhileLessThanOne<<<1, 1, 0, streamToRunOn>>>(latch);\n }\n\n if (someTransferOpRequired) {\n sdkStartTimer(&gpuTransferCallsTimer);\n if (copyRequired) {\n if (isAsync) {\n checkCudaErrors(hipMemcpyAsync(dptrA, hptrA, size, hipMemcpyHostToDevice, streamToRunOn));\n checkCudaErrors(hipMemcpyAsync(dptrB, hptrB, size, hipMemcpyHostToDevice, streamToRunOn));\n }\n else {\n checkCudaErrors(hipMemcpy(dptrA, hptrA, size, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(dptrB, hptrB, size, hipMemcpyHostToDevice));\n }\n }\n if (hintsRequired) {\n if (deviceProp.concurrentManagedAccess) {\n hipMemLocation deviceLoc;\n deviceLoc.type = hipMemLocationTypeDevice;\n deviceLoc.id = device_id;\n checkCudaErrors(hipMemPrefetchAsync(dptrA, size, deviceLoc, 0, streamToRunOn));\n checkCudaErrors(hipMemPrefetchAsync(dptrB, size, deviceLoc, 0, streamToRunOn));\n checkCudaErrors(hipMemPrefetchAsync(dptrC, size, deviceLoc, 0, streamToRunOn));\n }\n else {\n checkCudaErrors(hipStreamAttachMemAsync(streamToRunOn, dptrA, 0, hipMemAttachGlobal));\n checkCudaErrors(hipStreamAttachMemAsync(streamToRunOn, dptrB, 0, hipMemAttachGlobal));\n checkCudaErrors(hipStreamAttachMemAsync(streamToRunOn, dptrC, 0, hipMemAttachGlobal));\n }\n if (!isAsync) {\n checkCudaErrors(hipStreamSynchronize(streamToRunOn));\n }\n }\n\n sdkStopTimer(&gpuTransferCallsTimer);\n gpuTransferToCallsTimes[i] += sdkGetAverageTimerValue(&gpuTransferCallsTimer);\n sdkResetTimer(&gpuTransferCallsTimer);\n }\n\n sdkStartTimer(&gpuLaunchCallsTimer);\n {\n matrixMultiplyKernel<<>>(dptrC, dptrA, dptrB, matrixDim);\n if (!isAsync) {\n checkCudaErrors(hipStreamSynchronize(streamToRunOn));\n }\n }\n sdkStopTimer(&gpuLaunchCallsTimer);\n\n gpuLaunchCallsTimes[i] += sdkGetAverageTimerValue(&gpuLaunchCallsTimer);\n sdkResetTimer(&gpuLaunchCallsTimer);\n\n if (someTransferOpRequired) {\n sdkStartTimer(&gpuTransferCallsTimer);\n if (hintsRequired) {\n if (deviceProp.concurrentManagedAccess) {\n hipMemLocation hostLoc;\n hostLoc.type = hipMemLocationTypeHost;\n checkCudaErrors(hipMemPrefetchAsync(dptrA, size, hostLoc, 0));\n checkCudaErrors(hipMemPrefetchAsync(dptrB, size, hostLoc, 0));\n checkCudaErrors(hipMemPrefetchAsync(dptrC, size, hostLoc, 0));\n }\n else {\n checkCudaErrors(hipStreamAttachMemAsync(streamToRunOn, dptrA, 0, hipMemAttachHost));\n checkCudaErrors(hipStreamAttachMemAsync(streamToRunOn, dptrB, 0, hipMemAttachHost));\n checkCudaErrors(hipStreamAttachMemAsync(streamToRunOn, dptrC, 0, hipMemAttachHost));\n }\n if (!isAsync) {\n checkCudaErrors(hipStreamSynchronize(streamToRunOn));\n }\n }\n if (copyRequired) {\n if (isAsync) {\n checkCudaErrors(hipMemcpyAsync(hptrC, dptrC, size, hipMemcpyDeviceToHost, streamToRunOn));\n }\n else {\n checkCudaErrors(hipMemcpy(hptrC, dptrC, size, hipMemcpyDeviceToHost));\n }\n }\n sdkStopTimer(&gpuTransferCallsTimer);\n gpuTransferFromCallsTimes[i] += sdkGetAverageTimerValue(&gpuTransferCallsTimer);\n sdkResetTimer(&gpuTransferCallsTimer);\n }\n gpuLaunchAndTransferCallsTimes[i] =\n gpuLaunchCallsTimes[i] + gpuTransferToCallsTimes[i] + gpuTransferFromCallsTimes[i];\n gpuLaunchTransferSyncTimes[i] = gpuLaunchAndTransferCallsTimes[i];\n if (isAsync) {\n sdkStartTimer(&gpuSyncTimer);\n {\n if (hintsRequired) {\n *latch = 1;\n }\n checkCudaErrors(hipStreamSynchronize(streamToRunOn));\n }\n sdkStopTimer(&gpuSyncTimer);\n gpuLaunchTransferSyncTimes[i] += sdkGetAverageTimerValue(&gpuSyncTimer);\n sdkResetTimer(&gpuSyncTimer);\n }\n\n sdkStartTimer(&cpuAccessTimer);\n {\n verifyMatrixData((i & 0x1 == 0) ? randValuesVerifyXmulY : randValuesVerifyYmulX, hptrC, matrixDim);\n }\n sdkStopTimer(&cpuAccessTimer);\n cpuAccessTimes[i] += sdkGetAverageTimerValue(&cpuAccessTimer);\n sdkResetTimer(&cpuAccessTimer);\n overallTimes[i] = cpuAccessTimes[i] + gpuLaunchTransferSyncTimes[i];\n }\n\n switch (allocType) {\n case USE_HOST_PAGEABLE_AND_DEVICE_MEMORY:\n case USE_HOST_PAGEABLE_AND_DEVICE_MEMORY_ASYNC:\n free(hptrA);\n free(hptrB);\n free(hptrC);\n checkCudaErrors(hipFree(dptrA));\n checkCudaErrors(hipFree(dptrB));\n checkCudaErrors(hipFree(dptrC));\n break;\n\n case USE_HOST_PAGELOCKED_AND_DEVICE_MEMORY:\n case USE_HOST_PAGELOCKED_AND_DEVICE_MEMORY_ASYNC:\n checkCudaErrors(hipHostFree(hptrA));\n checkCudaErrors(hipHostFree(hptrB));\n checkCudaErrors(hipHostFree(hptrC));\n checkCudaErrors(hipFree(dptrA));\n checkCudaErrors(hipFree(dptrB));\n checkCudaErrors(hipFree(dptrC));\n break;\n\n case USE_ZERO_COPY:\n checkCudaErrors(hipHostFree(hptrA));\n checkCudaErrors(hipHostFree(hptrB));\n checkCudaErrors(hipHostFree(hptrC));\n break;\n\n case USE_MANAGED_MEMORY:\n case USE_MANAGED_MEMORY_WITH_HINTS:\n case USE_MANAGED_MEMORY_WITH_HINTS_ASYNC:\n checkCudaErrors(hipFree(dptrA));\n checkCudaErrors(hipFree(dptrB));\n checkCudaErrors(hipFree(dptrC));\n break;\n\n default:\n exit(EXIT_FAILURE); // exit due to error\n }\n\n checkCudaErrors(hipStreamDestroy(streamToRunOn));\n checkCudaErrors(hipHostFree(latch));\n free(randValuesX);\n free(randValuesY);\n free(randValuesVerifyXmulY);\n free(randValuesVerifyYmulX);\n sdkDeleteTimer(&gpuLaunchCallsTimer);\n sdkDeleteTimer(&gpuTransferCallsTimer);\n sdkDeleteTimer(&gpuSyncTimer);\n sdkDeleteTimer(&cpuAccessTimer);\n}\n\nvoid matrixMultiplyPerfRunner(bool reportAsBandwidth,\n bool print_launch_transfer_results,\n bool print_std_deviation,\n int device_id)\n{\n int i;\n unsigned int minMatrixDim = 32;\n unsigned int multiplierDim = 2;\n unsigned int matrixDim;\n unsigned int minSize = minMatrixDim * minMatrixDim * sizeof(float);\n unsigned int maxSize = (maxSampleSizeInMb * ONE_MB) / 4; // 3 buffers are used, but dividing by 4 (power of 2)\n unsigned int multiplier = multiplierDim * multiplierDim;\n unsigned int numSizesToTest;\n\n struct testResults *results;\n struct resultsData *gpuLaunchCallsTimes;\n struct resultsData *gpuTransferToCallsTimes;\n struct resultsData *gpuTransferFromCallsTimes;\n struct resultsData *gpuLaunchAndTransferCallsTimes;\n struct resultsData *gpuLaunchTransferSyncTimes;\n struct resultsData *cpuAccessTimes;\n struct resultsData *overallTimes;\n unsigned long *sizesToTest;\n unsigned int j;\n\n numSizesToTest = findNumSizesToTest(minSize, maxSize, multiplier);\n\n createAndInitTestResults(&results, \"matrixMultiplyPerf\", numKernelRuns, numSizesToTest);\n\n sizesToTest = getPtrSizesToTest(results);\n\n createResultDataAndAddToTestResults(\n &gpuLaunchCallsTimes, results, \"GPU Kernel Launch Call Time\", false, reportAsBandwidth);\n createResultDataAndAddToTestResults(\n &gpuTransferToCallsTimes, results, \"CPU to GPU Transfer Calls Time\", false, reportAsBandwidth);\n createResultDataAndAddToTestResults(\n &gpuTransferFromCallsTimes, results, \"GPU to CPU Transfer Calls Time\", false, reportAsBandwidth);\n createResultDataAndAddToTestResults(\n &gpuLaunchAndTransferCallsTimes, results, \"GPU Launch and Transfer Calls Time\", false, reportAsBandwidth);\n createResultDataAndAddToTestResults(\n &gpuLaunchTransferSyncTimes, results, \"GPU Launch Transfer and Sync Time\", false, reportAsBandwidth);\n createResultDataAndAddToTestResults(&cpuAccessTimes, results, \"CPU Access Time\", false, reportAsBandwidth);\n createResultDataAndAddToTestResults(&overallTimes, results, \"Overall Time\", false, reportAsBandwidth);\n\n printf(\"Running \");\n for (matrixDim = minMatrixDim, j = 0; matrixDim * matrixDim <= maxSize / sizeof(float);\n matrixDim *= multiplierDim, ++j) {\n sizesToTest[j] = matrixDim * matrixDim * sizeof(float);\n for (i = MEMALLOC_TYPE_START; i <= MEMALLOC_TYPE_END; i++) {\n printf(\".\");\n fflush(stdout);\n runMatrixMultiplyKernel(matrixDim,\n i,\n numKernelRuns,\n getPtrRunTimesInMs(gpuLaunchCallsTimes, i, j),\n getPtrRunTimesInMs(gpuTransferToCallsTimes, i, j),\n getPtrRunTimesInMs(gpuTransferFromCallsTimes, i, j),\n getPtrRunTimesInMs(gpuLaunchAndTransferCallsTimes, i, j),\n getPtrRunTimesInMs(gpuLaunchTransferSyncTimes, i, j),\n getPtrRunTimesInMs(cpuAccessTimes, i, j),\n getPtrRunTimesInMs(overallTimes, i, j),\n device_id);\n }\n }\n printf(\"\\n\");\n printResults(results, print_launch_transfer_results, print_std_deviation);\n freeTestResultsAndAllResultsData(results);\n}\n\nstatic void usage()\n{\n printf(\"./cudaMemoryTypesPerf [-device=] [-reportAsBandwidth] \"\n \"[-print-launch-transfer-results] [-print-std-deviation] [-verbose]\\n\");\n printf(\"Options:\\n\");\n printf(\"-reportAsBandwidth: By default time taken is printed, this \"\n \"option allows to instead print bandwidth.\\n\");\n printf(\"-print-launch-transfer-results: By default overall results are printed, \"\n \"this option allows to print data transfers and kernel time as well.\\n\");\n printf(\"-print-std-deviation: Prints std deviation of the results.\\n\");\n printf(\"-kernel-iterations=: Number of times the kernel tests should \"\n \"be run[default is 100 iterations].\\n\");\n printf(\"-device=: Allows to pass GPU Device ID on which \"\n \"the tests will be run.\\n\");\n printf(\"-verbose: Prints highly verbose output.\\n\");\n}\n\nint main(int argc, char **argv)\n{\n bool reportAsBandwidth = false;\n bool print_launch_transfer_results = false;\n bool print_std_deviation = false;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\") || checkCmdLineFlag(argc, (const char **)argv, \"h\")) {\n usage();\n printf(\"&&&& %s WAIVED\\n\", argv[0]);\n exit(EXIT_WAIVED);\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"reportAsBandwidth\")) {\n reportAsBandwidth = true;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"print-launch-transfer-results\")) {\n print_launch_transfer_results = true;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"print-std-deviation\")) {\n print_std_deviation = true;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"kernel-iterations\")) {\n numKernelRuns = getCmdLineArgumentInt(argc, (const char **)argv, \"kernel-iterations\");\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"verbose\")) {\n verboseResults = 1;\n }\n\n // set device\n hipDeviceProp_t device_prop;\n int device_id = findCudaDevice(argc, (const char **)argv);\n checkCudaErrors(hipGetDeviceProperties(&device_prop, device_id));\n\n if (!device_prop.managedMemory) {\n // This samples requires being run on a device that supports Unified Memory\n fprintf(stderr, \"Unified Memory not supported on this device\\n\");\n\n exit(EXIT_WAIVED);\n }\n\n matrixMultiplyPerfRunner(reportAsBandwidth, print_launch_transfer_results, print_std_deviation, device_id);\n\n printf(\"\\nNOTE: The CUDA Samples are not meant for performance measurements. \"\n \"Results may vary when GPU Boost is enabled.\\n\");\n exit(EXIT_SUCCESS);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/6_Performance/UnifiedMemoryPerf/matrixMultiplyPerf.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// -----------------------------------------------------------------------------\n// Transpose\n//\n// This file contains both device and host code for transposing a floating-point\n// matrix. It performs several transpose kernels, which incrementally improve\n// performance through coalescing, removing shared memory bank conflicts, and\n// eliminating partition camping. Several of the kernels perform a copy, used\n// to represent the best case performance that a transpose can achieve.\n//\n// Please see the whitepaper in the docs folder of the transpose project for a\n// detailed description of this performance study.\n// -----------------------------------------------------------------------------\n\n#include \n\nnamespace cg = cooperative_groups;\n// Utilities and system includes\n#include // helper for cuda error checking functions\n#include // helper for image and data comparison\n#include // helper for string parsing\n\nconst char *sSDKsample = \"Transpose\";\n\n// Each block transposes/copies a tile of TILE_DIM x TILE_DIM elements\n// using TILE_DIM x BLOCK_ROWS threads, so that each thread transposes\n// TILE_DIM/BLOCK_ROWS elements. TILE_DIM must be an integral multiple of\n// BLOCK_ROWS\n\n#define TILE_DIM 32\n#define BLOCK_ROWS 16\n\n// This sample assumes that MATRIX_SIZE_X = MATRIX_SIZE_Y\nint MATRIX_SIZE_X = 1024;\nint MATRIX_SIZE_Y = 1024;\nint MUL_FACTOR = TILE_DIM;\n\n#define FLOOR(a, b) (a - (a % b))\n\n// Compute the tile size necessary to illustrate performance cases for SM20+\n// hardware\nint MAX_TILES = (FLOOR(MATRIX_SIZE_X, 512) * FLOOR(MATRIX_SIZE_Y, 512)) / (TILE_DIM * TILE_DIM);\n\n// Number of repetitions used for timing. Two sets of repetitions are\n// performed: 1) over kernel launches and 2) inside the kernel over just the\n// loads and stores\n\n#define NUM_REPS 100\n\n// -------------------------------------------------------\n// Copies\n// width and height must be integral multiples of TILE_DIM\n// -------------------------------------------------------\n\n__global__ void copy(float *odata, float *idata, int width, int height)\n{\n int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;\n\n int index = xIndex + width * yIndex;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n odata[index + i * width] = idata[index + i * width];\n }\n}\n\n__global__ void copySharedMem(float *odata, float *idata, int width, int height)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float tile[TILE_DIM][TILE_DIM];\n\n int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;\n\n int index = xIndex + width * yIndex;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n if (xIndex < width && yIndex < height) {\n tile[threadIdx.y + i][threadIdx.x] = idata[index + i * width];\n tile[threadIdx.y + i][threadIdx.x] = idata[index + i * width];\n }\n }\n\n cg::sync(cta);\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n if (xIndex < height && yIndex < width) {\n odata[index + i * width] = tile[threadIdx.y + i][threadIdx.x];\n odata[index + i * width] = tile[threadIdx.y + i][threadIdx.x];\n }\n }\n}\n\n// -------------------------------------------------------\n// Transposes\n// width and height must be integral multiples of TILE_DIM\n// -------------------------------------------------------\n\n__global__ void transposeNaive(float *odata, float *idata, int width, int height)\n{\n int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;\n\n int index_in = xIndex + width * yIndex;\n int index_out = yIndex + height * xIndex;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n odata[index_out + i] = idata[index_in + i * width];\n }\n}\n\n// coalesced transpose (with bank conflicts)\n\n__global__ void transposeCoalesced(float *odata, float *idata, int width, int height)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float tile[TILE_DIM][TILE_DIM];\n\n int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;\n int index_in = xIndex + (yIndex)*width;\n\n xIndex = blockIdx.y * TILE_DIM + threadIdx.x;\n yIndex = blockIdx.x * TILE_DIM + threadIdx.y;\n int index_out = xIndex + (yIndex)*height;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n tile[threadIdx.y + i][threadIdx.x] = idata[index_in + i * width];\n }\n\n cg::sync(cta);\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n odata[index_out + i * height] = tile[threadIdx.x][threadIdx.y + i];\n }\n}\n\n// Coalesced transpose with no bank conflicts\n\n__global__ void transposeNoBankConflicts(float *odata, float *idata, int width, int height)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float tile[TILE_DIM][TILE_DIM + 1];\n\n int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;\n int index_in = xIndex + (yIndex)*width;\n\n xIndex = blockIdx.y * TILE_DIM + threadIdx.x;\n yIndex = blockIdx.x * TILE_DIM + threadIdx.y;\n int index_out = xIndex + (yIndex)*height;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n tile[threadIdx.y + i][threadIdx.x] = idata[index_in + i * width];\n }\n\n cg::sync(cta);\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n odata[index_out + i * height] = tile[threadIdx.x][threadIdx.y + i];\n }\n}\n\n// Transpose that effectively reorders execution of thread blocks along\n// diagonals of the matrix (also coalesced and has no bank conflicts)\n//\n// Here blockIdx.x is interpreted as the distance along a diagonal and\n// blockIdx.y as corresponding to different diagonals\n//\n// blockIdx_x and blockIdx_y expressions map the diagonal coordinates to the\n// more commonly used cartesian coordinates so that the only changes to the code\n// from the coalesced version are the calculation of the blockIdx_x and\n// blockIdx_y and replacement of blockIdx.x and bloclIdx.y with the subscripted\n// versions in the remaining code\n\n__global__ void transposeDiagonal(float *odata, float *idata, int width, int height)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float tile[TILE_DIM][TILE_DIM + 1];\n\n int blockIdx_x, blockIdx_y;\n\n // do diagonal reordering\n if (width == height) {\n blockIdx_y = blockIdx.x;\n blockIdx_x = (blockIdx.x + blockIdx.y) % gridDim.x;\n }\n else {\n int bid = blockIdx.x + gridDim.x * blockIdx.y;\n blockIdx_y = bid % gridDim.y;\n blockIdx_x = ((bid / gridDim.y) + blockIdx_y) % gridDim.x;\n }\n\n // from here on the code is same as previous kernel except blockIdx_x replaces\n // blockIdx.x and similarly for y\n\n int xIndex = blockIdx_x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx_y * TILE_DIM + threadIdx.y;\n int index_in = xIndex + (yIndex)*width;\n\n xIndex = blockIdx_y * TILE_DIM + threadIdx.x;\n yIndex = blockIdx_x * TILE_DIM + threadIdx.y;\n int index_out = xIndex + (yIndex)*height;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n tile[threadIdx.y + i][threadIdx.x] = idata[index_in + i * width];\n }\n\n cg::sync(cta);\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n odata[index_out + i * height] = tile[threadIdx.x][threadIdx.y + i];\n }\n}\n\n// --------------------------------------------------------------------\n// Partial transposes\n// NB: the coarse- and fine-grained routines only perform part of a\n// transpose and will fail the test against the reference solution\n//\n// They are used to assess performance characteristics of different\n// components of a full transpose\n// --------------------------------------------------------------------\n\n__global__ void transposeFineGrained(float *odata, float *idata, int width, int height)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float block[TILE_DIM][TILE_DIM + 1];\n\n int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;\n int index = xIndex + (yIndex)*width;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n block[threadIdx.y + i][threadIdx.x] = idata[index + i * width];\n }\n\n cg::sync(cta);\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n odata[index + i * height] = block[threadIdx.x][threadIdx.y + i];\n }\n}\n\n__global__ void transposeCoarseGrained(float *odata, float *idata, int width, int height)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float block[TILE_DIM][TILE_DIM + 1];\n\n int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;\n int index_in = xIndex + (yIndex)*width;\n\n xIndex = blockIdx.y * TILE_DIM + threadIdx.x;\n yIndex = blockIdx.x * TILE_DIM + threadIdx.y;\n int index_out = xIndex + (yIndex)*height;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n block[threadIdx.y + i][threadIdx.x] = idata[index_in + i * width];\n }\n\n cg::sync(cta);\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n odata[index_out + i * height] = block[threadIdx.y + i][threadIdx.x];\n }\n}\n\n// ---------------------\n// host utility routines\n// ---------------------\n\nvoid computeTransposeGold(float *gold, float *idata, const int size_x, const int size_y)\n{\n for (int y = 0; y < size_y; ++y) {\n for (int x = 0; x < size_x; ++x) {\n gold[(x * size_y) + y] = idata[(y * size_x) + x];\n }\n }\n}\n\nvoid getParams(int argc, char **argv, cudaDeviceProp &deviceProp, int &size_x, int &size_y, int max_tile_dim)\n{\n // set matrix size (if (x,y) dim of matrix is not square, then this will have\n // to be modified\n if (checkCmdLineFlag(argc, (const char **)argv, \"dimX\")) {\n size_x = getCmdLineArgumentInt(argc, (const char **)argv, \"dimX\");\n\n if (size_x > max_tile_dim) {\n printf(\"> MatrixSize X = %d is greater than the recommended size = %d\\n\", size_x, max_tile_dim);\n }\n else {\n printf(\"> MatrixSize X = %d\\n\", size_x);\n }\n }\n else {\n size_x = max_tile_dim;\n size_x = FLOOR(size_x, 512);\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"dimY\")) {\n size_y = getCmdLineArgumentInt(argc, (const char **)argv, \"dimY\");\n\n if (size_y > max_tile_dim) {\n printf(\"> MatrixSize Y = %d is greater than the recommended size = %d\\n\", size_y, max_tile_dim);\n }\n else {\n printf(\"> MatrixSize Y = %d\\n\", size_y);\n }\n }\n else {\n size_y = max_tile_dim;\n size_y = FLOOR(size_y, 512);\n }\n}\n\nvoid showHelp()\n{\n printf(\"\\n%s : Command line options\\n\", sSDKsample);\n printf(\"\\t-device=n (where n=0,1,2.... for the GPU device)\\n\\n\");\n printf(\"> The default matrix size can be overridden with these parameters\\n\");\n printf(\"\\t-dimX=row_dim_size (matrix row dimensions)\\n\");\n printf(\"\\t-dimY=col_dim_size (matrix column dimensions)\\n\");\n}\n\n// ----\n// main\n// ----\n\nint main(int argc, char **argv)\n{\n // Start logs\n printf(\"%s Starting...\\n\\n\", sSDKsample);\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\")) {\n showHelp();\n return 0;\n }\n\n int devID = findCudaDevice(argc, (const char **)argv);\n cudaDeviceProp deviceProp;\n\n // get number of SMs on this GPU\n checkCudaErrors(cudaGetDevice(&devID));\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, devID));\n\n // compute the scaling factor (for GPUs with fewer MPs)\n float scale_factor, total_tiles;\n scale_factor = max(\n (192.0f / (_ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * (float)deviceProp.multiProcessorCount)),\n 1.0f);\n\n printf(\"> Device %d: \\\"%s\\\"\\n\", devID, deviceProp.name);\n printf(\"> SM Capability %d.%d detected:\\n\", deviceProp.major, deviceProp.minor);\n\n // Calculate number of tiles we will run for the Matrix Transpose performance\n // tests\n int size_x, size_y, max_matrix_dim, matrix_size_test;\n\n matrix_size_test = 512; // we round down max_matrix_dim for this perf test\n total_tiles = (float)MAX_TILES / scale_factor;\n\n max_matrix_dim = FLOOR((int)(floor(sqrt(total_tiles)) * TILE_DIM), matrix_size_test);\n\n // This is the minimum size allowed\n if (max_matrix_dim == 0) {\n max_matrix_dim = matrix_size_test;\n }\n\n printf(\"> [%s] has %d MP(s) x %d (Cores/MP) = %d (Cores)\\n\",\n deviceProp.name,\n deviceProp.multiProcessorCount,\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor),\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * deviceProp.multiProcessorCount);\n\n printf(\"> Compute performance scaling factor = %4.2f\\n\", scale_factor);\n\n // Extract parameters if there are any, command line -dimx and -dimy can\n // override any of these settings\n getParams(argc, argv, deviceProp, size_x, size_y, max_matrix_dim);\n\n if (size_x != size_y) {\n printf(\"\\n[%s] does not support non-square matrices (row_dim_size(%d) != \"\n \"col_dim_size(%d))\\nExiting...\\n\\n\",\n sSDKsample,\n size_x,\n size_y);\n exit(EXIT_FAILURE);\n }\n\n if (size_x % TILE_DIM != 0 || size_y % TILE_DIM != 0) {\n printf(\"[%s] Matrix size must be integral multiple of tile \"\n \"size\\nExiting...\\n\\n\",\n sSDKsample);\n exit(EXIT_FAILURE);\n }\n\n // kernel pointer and descriptor\n void (*kernel)(float *, float *, int, int);\n const char *kernelName;\n\n // execution configuration parameters\n dim3 grid(size_x / TILE_DIM, size_y / TILE_DIM), threads(TILE_DIM, BLOCK_ROWS);\n\n if (grid.x < 1 || grid.y < 1) {\n printf(\"[%s] grid size computation incorrect in test \\nExiting...\\n\\n\", sSDKsample);\n exit(EXIT_FAILURE);\n }\n\n // CUDA events\n cudaEvent_t start, stop;\n\n // size of memory required to store the matrix\n size_t mem_size = static_cast(sizeof(float) * size_x * size_y);\n\n if (2 * mem_size > deviceProp.totalGlobalMem) {\n printf(\"Input matrix size is larger than the available device memory!\\n\");\n printf(\"Please choose a smaller size matrix\\n\");\n exit(EXIT_FAILURE);\n }\n\n // allocate host memory\n float *h_idata = (float *)malloc(mem_size);\n float *h_odata = (float *)malloc(mem_size);\n float *transposeGold = (float *)malloc(mem_size);\n float *gold;\n\n // allocate device memory\n float *d_idata, *d_odata;\n checkCudaErrors(cudaMalloc((void **)&d_idata, mem_size));\n checkCudaErrors(cudaMalloc((void **)&d_odata, mem_size));\n\n // initialize host data\n for (int i = 0; i < (size_x * size_y); ++i) {\n h_idata[i] = (float)i;\n }\n\n // copy host data to device\n checkCudaErrors(cudaMemcpy(d_idata, h_idata, mem_size, cudaMemcpyHostToDevice));\n\n // Compute reference transpose solution\n computeTransposeGold(transposeGold, h_idata, size_x, size_y);\n\n // print out common data for all kernels\n printf(\"\\nMatrix size: %dx%d (%dx%d tiles), tile size: %dx%d, block size: \"\n \"%dx%d\\n\\n\",\n size_x,\n size_y,\n size_x / TILE_DIM,\n size_y / TILE_DIM,\n TILE_DIM,\n TILE_DIM,\n TILE_DIM,\n BLOCK_ROWS);\n\n // initialize events\n checkCudaErrors(cudaEventCreate(&start));\n checkCudaErrors(cudaEventCreate(&stop));\n\n //\n // loop over different kernels\n //\n\n bool success = true;\n\n for (int k = 0; k < 8; k++) {\n // set kernel pointer\n switch (k) {\n case 0:\n kernel = ©\n kernelName = \"simple copy \";\n break;\n\n case 1:\n kernel = ©SharedMem;\n kernelName = \"shared memory copy\";\n break;\n\n case 2:\n kernel = &transposeNaive;\n kernelName = \"naive \";\n break;\n\n case 3:\n kernel = &transposeCoalesced;\n kernelName = \"coalesced \";\n break;\n\n case 4:\n kernel = &transposeNoBankConflicts;\n kernelName = \"optimized \";\n break;\n\n case 5:\n kernel = &transposeCoarseGrained;\n kernelName = \"coarse-grained \";\n break;\n\n case 6:\n kernel = &transposeFineGrained;\n kernelName = \"fine-grained \";\n break;\n\n case 7:\n kernel = &transposeDiagonal;\n kernelName = \"diagonal \";\n break;\n }\n\n // set reference solution\n if (kernel == © || kernel == ©SharedMem) {\n gold = h_idata;\n }\n else if (kernel == &transposeCoarseGrained || kernel == &transposeFineGrained) {\n gold = h_odata; // fine- and coarse-grained kernels are not full\n // transposes, so bypass check\n }\n else {\n gold = transposeGold;\n }\n\n // Clear error status\n checkCudaErrors(cudaGetLastError());\n\n // warmup to avoid timing startup\n kernel<<>>(d_odata, d_idata, size_x, size_y);\n\n // take measurements for loop over kernel launches\n checkCudaErrors(cudaEventRecord(start, 0));\n\n for (int i = 0; i < NUM_REPS; i++) {\n kernel<<>>(d_odata, d_idata, size_x, size_y);\n // Ensure no launch failure\n checkCudaErrors(cudaGetLastError());\n }\n\n checkCudaErrors(cudaEventRecord(stop, 0));\n checkCudaErrors(cudaEventSynchronize(stop));\n float kernelTime;\n checkCudaErrors(cudaEventElapsedTime(&kernelTime, start, stop));\n\n checkCudaErrors(cudaMemcpy(h_odata, d_odata, mem_size, cudaMemcpyDeviceToHost));\n bool res = compareData(gold, h_odata, size_x * size_y, 0.01f, 0.0f);\n\n if (res == false) {\n printf(\"*** %s kernel FAILED ***\\n\", kernelName);\n success = false;\n }\n\n // take measurements for loop inside kernel\n checkCudaErrors(cudaMemcpy(h_odata, d_odata, mem_size, cudaMemcpyDeviceToHost));\n res = compareData(gold, h_odata, size_x * size_y, 0.01f, 0.0f);\n\n if (res == false) {\n printf(\"*** %s kernel FAILED ***\\n\", kernelName);\n success = false;\n }\n\n // report effective bandwidths\n float kernelBandwidth = 2.0f * 1000.0f * mem_size / (1024 * 1024 * 1024) / (kernelTime / NUM_REPS);\n printf(\"transpose %s, Throughput = %.4f GB/s, Time = %.5f ms, Size = %u fp32 \"\n \"elements, NumDevsUsed = %u, Workgroup = %u\\n\",\n kernelName,\n kernelBandwidth,\n kernelTime / NUM_REPS,\n (size_x * size_y),\n 1,\n TILE_DIM * BLOCK_ROWS);\n\n // Reset d_odata to zero before starting the next loop iteration to avoid\n // carrying over results from previous kernels. Without this reset, residual\n // data from a prior kernel (e.g., 'copy') could make a subsequent\n // kernel (e.g., 'copySharedMem') appear correct even if it performs no work,\n // leading to false positives in compareData.\n for (int i = 0; i < (size_x * size_y); ++i) {\n h_odata[i] = 0;\n }\n // copy host data to device\n checkCudaErrors(cudaMemcpy(d_odata, h_odata, mem_size, cudaMemcpyHostToDevice));\n\n // Reset d_odata to zero before starting the next loop iteration to avoid\n // carrying over results from previous kernels. Without this reset, residual\n // data from a prior kernel (e.g., 'copy') could make a subsequent\n // kernel (e.g., 'copySharedMem') appear correct even if it performs no work,\n // leading to false positives in compareData.\n for (int i = 0; i < (size_x * size_y); ++i) {\n h_odata[i] = 0;\n }\n // copy host data to device\n checkCudaErrors(cudaMemcpy(d_odata, h_odata, mem_size, cudaMemcpyHostToDevice));\n }\n\n // cleanup\n free(h_idata);\n free(h_odata);\n free(transposeGold);\n cudaFree(d_idata);\n cudaFree(d_odata);\n\n checkCudaErrors(cudaEventDestroy(start));\n checkCudaErrors(cudaEventDestroy(stop));\n\n if (!success) {\n printf(\"Test failed!\\n\");\n exit(EXIT_FAILURE);\n }\n\n printf(\"Test passed\\n\");\n exit(EXIT_SUCCESS);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// -----------------------------------------------------------------------------\n// Transpose\n//\n// This file contains both device and host code for transposing a floating-point\n// matrix. It performs several transpose kernels, which incrementally improve\n// performance through coalescing, removing shared memory bank conflicts, and\n// eliminating partition camping. Several of the kernels perform a copy, used\n// to represent the best case performance that a transpose can achieve.\n//\n// Please see the whitepaper in the docs folder of the transpose project for a\n// detailed description of this performance study.\n// -----------------------------------------------------------------------------\n\n#include \"hip/hip_runtime.h\"\n#include \n\nnamespace cg = cooperative_groups;\n// Utilities and system includes\n#include // helper for cuda error checking functions\n#include // helper for image and data comparison\n#include // helper for string parsing\n\nconst char *sSDKsample = \"Transpose\";\n\n// Each block transposes/copies a tile of TILE_DIM x TILE_DIM elements\n// using TILE_DIM x BLOCK_ROWS threads, so that each thread transposes\n// TILE_DIM/BLOCK_ROWS elements. TILE_DIM must be an integral multiple of\n// BLOCK_ROWS\n\n#define TILE_DIM 32\n#define BLOCK_ROWS 16\n\n// This sample assumes that MATRIX_SIZE_X = MATRIX_SIZE_Y\nint MATRIX_SIZE_X = 1024;\nint MATRIX_SIZE_Y = 1024;\nint MUL_FACTOR = TILE_DIM;\n\n#define FLOOR(a, b) (a - (a % b))\n\n// Compute the tile size necessary to illustrate performance cases for SM20+\n// hardware\nint MAX_TILES = (FLOOR(MATRIX_SIZE_X, 512) * FLOOR(MATRIX_SIZE_Y, 512)) / (TILE_DIM * TILE_DIM);\n\n// Number of repetitions used for timing. Two sets of repetitions are\n// performed: 1) over kernel launches and 2) inside the kernel over just the\n// loads and stores\n\n#define NUM_REPS 100\n\n// -------------------------------------------------------\n// Copies\n// width and height must be integral multiples of TILE_DIM\n// -------------------------------------------------------\n\n__global__ void copy(float *odata, float *idata, int width, int height)\n{\n int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;\n\n int index = xIndex + width * yIndex;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n odata[index + i * width] = idata[index + i * width];\n }\n}\n\n__global__ void copySharedMem(float *odata, float *idata, int width, int height)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float tile[TILE_DIM][TILE_DIM];\n\n int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;\n\n int index = xIndex + width * yIndex;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n if (xIndex < width && yIndex < height) {\n tile[threadIdx.y + i][threadIdx.x] = idata[index + i * width];\n tile[threadIdx.y + i][threadIdx.x] = idata[index + i * width];\n }\n }\n\n cg::sync(cta);\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n if (xIndex < height && yIndex < width) {\n odata[index + i * width] = tile[threadIdx.y + i][threadIdx.x];\n odata[index + i * width] = tile[threadIdx.y + i][threadIdx.x];\n }\n }\n}\n\n// -------------------------------------------------------\n// Transposes\n// width and height must be integral multiples of TILE_DIM\n// -------------------------------------------------------\n\n__global__ void transposeNaive(float *odata, float *idata, int width, int height)\n{\n int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;\n\n int index_in = xIndex + width * yIndex;\n int index_out = yIndex + height * xIndex;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n odata[index_out + i] = idata[index_in + i * width];\n }\n}\n\n// coalesced transpose (with bank conflicts)\n\n__global__ void transposeCoalesced(float *odata, float *idata, int width, int height)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float tile[TILE_DIM][TILE_DIM];\n\n int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;\n int index_in = xIndex + (yIndex)*width;\n\n xIndex = blockIdx.y * TILE_DIM + threadIdx.x;\n yIndex = blockIdx.x * TILE_DIM + threadIdx.y;\n int index_out = xIndex + (yIndex)*height;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n tile[threadIdx.y + i][threadIdx.x] = idata[index_in + i * width];\n }\n\n cg::sync(cta);\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n odata[index_out + i * height] = tile[threadIdx.x][threadIdx.y + i];\n }\n}\n\n// Coalesced transpose with no bank conflicts\n\n__global__ void transposeNoBankConflicts(float *odata, float *idata, int width, int height)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float tile[TILE_DIM][TILE_DIM + 1];\n\n int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;\n int index_in = xIndex + (yIndex)*width;\n\n xIndex = blockIdx.y * TILE_DIM + threadIdx.x;\n yIndex = blockIdx.x * TILE_DIM + threadIdx.y;\n int index_out = xIndex + (yIndex)*height;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n tile[threadIdx.y + i][threadIdx.x] = idata[index_in + i * width];\n }\n\n cg::sync(cta);\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n odata[index_out + i * height] = tile[threadIdx.x][threadIdx.y + i];\n }\n}\n\n// Transpose that effectively reorders execution of thread blocks along\n// diagonals of the matrix (also coalesced and has no bank conflicts)\n//\n// Here blockIdx.x is interpreted as the distance along a diagonal and\n// blockIdx.y as corresponding to different diagonals\n//\n// blockIdx_x and blockIdx_y expressions map the diagonal coordinates to the\n// more commonly used cartesian coordinates so that the only changes to the code\n// from the coalesced version are the calculation of the blockIdx_x and\n// blockIdx_y and replacement of blockIdx.x and bloclIdx.y with the subscripted\n// versions in the remaining code\n\n__global__ void transposeDiagonal(float *odata, float *idata, int width, int height)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float tile[TILE_DIM][TILE_DIM + 1];\n\n int blockIdx_x, blockIdx_y;\n\n // do diagonal reordering\n if (width == height) {\n blockIdx_y = blockIdx.x;\n blockIdx_x = (blockIdx.x + blockIdx.y) % gridDim.x;\n }\n else {\n int bid = blockIdx.x + gridDim.x * blockIdx.y;\n blockIdx_y = bid % gridDim.y;\n blockIdx_x = ((bid / gridDim.y) + blockIdx_y) % gridDim.x;\n }\n\n // from here on the code is same as previous kernel except blockIdx_x replaces\n // blockIdx.x and similarly for y\n\n int xIndex = blockIdx_x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx_y * TILE_DIM + threadIdx.y;\n int index_in = xIndex + (yIndex)*width;\n\n xIndex = blockIdx_y * TILE_DIM + threadIdx.x;\n yIndex = blockIdx_x * TILE_DIM + threadIdx.y;\n int index_out = xIndex + (yIndex)*height;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n tile[threadIdx.y + i][threadIdx.x] = idata[index_in + i * width];\n }\n\n cg::sync(cta);\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n odata[index_out + i * height] = tile[threadIdx.x][threadIdx.y + i];\n }\n}\n\n// --------------------------------------------------------------------\n// Partial transposes\n// NB: the coarse- and fine-grained routines only perform part of a\n// transpose and will fail the test against the reference solution\n//\n// They are used to assess performance characteristics of different\n// components of a full transpose\n// --------------------------------------------------------------------\n\n__global__ void transposeFineGrained(float *odata, float *idata, int width, int height)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float block[TILE_DIM][TILE_DIM + 1];\n\n int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;\n int index = xIndex + (yIndex)*width;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n block[threadIdx.y + i][threadIdx.x] = idata[index + i * width];\n }\n\n cg::sync(cta);\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n odata[index + i * height] = block[threadIdx.x][threadIdx.y + i];\n }\n}\n\n__global__ void transposeCoarseGrained(float *odata, float *idata, int width, int height)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float block[TILE_DIM][TILE_DIM + 1];\n\n int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;\n int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;\n int index_in = xIndex + (yIndex)*width;\n\n xIndex = blockIdx.y * TILE_DIM + threadIdx.x;\n yIndex = blockIdx.x * TILE_DIM + threadIdx.y;\n int index_out = xIndex + (yIndex)*height;\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n block[threadIdx.y + i][threadIdx.x] = idata[index_in + i * width];\n }\n\n cg::sync(cta);\n\n for (int i = 0; i < TILE_DIM; i += BLOCK_ROWS) {\n odata[index_out + i * height] = block[threadIdx.y + i][threadIdx.x];\n }\n}\n\n// ---------------------\n// host utility routines\n// ---------------------\n\nvoid computeTransposeGold(float *gold, float *idata, const int size_x, const int size_y)\n{\n for (int y = 0; y < size_y; ++y) {\n for (int x = 0; x < size_x; ++x) {\n gold[(x * size_y) + y] = idata[(y * size_x) + x];\n }\n }\n}\n\nvoid getParams(int argc, char **argv, hipDeviceProp_t &deviceProp, int &size_x, int &size_y, int max_tile_dim)\n{\n // set matrix size (if (x,y) dim of matrix is not square, then this will have\n // to be modified\n if (checkCmdLineFlag(argc, (const char **)argv, \"dimX\")) {\n size_x = getCmdLineArgumentInt(argc, (const char **)argv, \"dimX\");\n\n if (size_x > max_tile_dim) {\n printf(\"> MatrixSize X = %d is greater than the recommended size = %d\\n\", size_x, max_tile_dim);\n }\n else {\n printf(\"> MatrixSize X = %d\\n\", size_x);\n }\n }\n else {\n size_x = max_tile_dim;\n size_x = FLOOR(size_x, 512);\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"dimY\")) {\n size_y = getCmdLineArgumentInt(argc, (const char **)argv, \"dimY\");\n\n if (size_y > max_tile_dim) {\n printf(\"> MatrixSize Y = %d is greater than the recommended size = %d\\n\", size_y, max_tile_dim);\n }\n else {\n printf(\"> MatrixSize Y = %d\\n\", size_y);\n }\n }\n else {\n size_y = max_tile_dim;\n size_y = FLOOR(size_y, 512);\n }\n}\n\nvoid showHelp()\n{\n printf(\"\\n%s : Command line options\\n\", sSDKsample);\n printf(\"\\t-device=n (where n=0,1,2.... for the GPU device)\\n\\n\");\n printf(\"> The default matrix size can be overridden with these parameters\\n\");\n printf(\"\\t-dimX=row_dim_size (matrix row dimensions)\\n\");\n printf(\"\\t-dimY=col_dim_size (matrix column dimensions)\\n\");\n}\n\n// ----\n// main\n// ----\n\nint main(int argc, char **argv)\n{\n // Start logs\n printf(\"%s Starting...\\n\\n\", sSDKsample);\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\")) {\n showHelp();\n return 0;\n }\n\n int devID = findCudaDevice(argc, (const char **)argv);\n hipDeviceProp_t deviceProp;\n\n // get number of SMs on this GPU\n checkCudaErrors(hipGetDevice(&devID));\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, devID));\n\n // compute the scaling factor (for GPUs with fewer MPs)\n float scale_factor, total_tiles;\n scale_factor = max(\n (192.0f / (_ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * (float)deviceProp.multiProcessorCount)),\n 1.0f);\n\n printf(\"> Device %d: \\\"%s\\\"\\n\", devID, deviceProp.name);\n printf(\"> SM Capability %d.%d detected:\\n\", deviceProp.major, deviceProp.minor);\n\n // Calculate number of tiles we will run for the Matrix Transpose performance\n // tests\n int size_x, size_y, max_matrix_dim, matrix_size_test;\n\n matrix_size_test = 512; // we round down max_matrix_dim for this perf test\n total_tiles = (float)MAX_TILES / scale_factor;\n\n max_matrix_dim = FLOOR((int)(floor(sqrt(total_tiles)) * TILE_DIM), matrix_size_test);\n\n // This is the minimum size allowed\n if (max_matrix_dim == 0) {\n max_matrix_dim = matrix_size_test;\n }\n\n printf(\"> [%s] has %d MP(s) x %d (Cores/MP) = %d (Cores)\\n\",\n deviceProp.name,\n deviceProp.multiProcessorCount,\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor),\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * deviceProp.multiProcessorCount);\n\n printf(\"> Compute performance scaling factor = %4.2f\\n\", scale_factor);\n\n // Extract parameters if there are any, command line -dimx and -dimy can\n // override any of these settings\n getParams(argc, argv, deviceProp, size_x, size_y, max_matrix_dim);\n\n if (size_x != size_y) {\n printf(\"\\n[%s] does not support non-square matrices (row_dim_size(%d) != \"\n \"col_dim_size(%d))\\nExiting...\\n\\n\",\n sSDKsample,\n size_x,\n size_y);\n exit(EXIT_FAILURE);\n }\n\n if (size_x % TILE_DIM != 0 || size_y % TILE_DIM != 0) {\n printf(\"[%s] Matrix size must be integral multiple of tile \"\n \"size\\nExiting...\\n\\n\",\n sSDKsample);\n exit(EXIT_FAILURE);\n }\n\n // kernel pointer and descriptor\n void (*kernel)(float *, float *, int, int);\n const char *kernelName;\n\n // execution configuration parameters\n dim3 grid(size_x / TILE_DIM, size_y / TILE_DIM), threads(TILE_DIM, BLOCK_ROWS);\n\n if (grid.x < 1 || grid.y < 1) {\n printf(\"[%s] grid size computation incorrect in test \\nExiting...\\n\\n\", sSDKsample);\n exit(EXIT_FAILURE);\n }\n\n // CUDA events\n hipEvent_t start, stop;\n\n // size of memory required to store the matrix\n size_t mem_size = static_cast(sizeof(float) * size_x * size_y);\n\n if (2 * mem_size > deviceProp.totalGlobalMem) {\n printf(\"Input matrix size is larger than the available device memory!\\n\");\n printf(\"Please choose a smaller size matrix\\n\");\n exit(EXIT_FAILURE);\n }\n\n // allocate host memory\n float *h_idata = (float *)malloc(mem_size);\n float *h_odata = (float *)malloc(mem_size);\n float *transposeGold = (float *)malloc(mem_size);\n float *gold;\n\n // allocate device memory\n float *d_idata, *d_odata;\n checkCudaErrors(hipMalloc((void **)&d_idata, mem_size));\n checkCudaErrors(hipMalloc((void **)&d_odata, mem_size));\n\n // initialize host data\n for (int i = 0; i < (size_x * size_y); ++i) {\n h_idata[i] = (float)i;\n }\n\n // copy host data to device\n checkCudaErrors(hipMemcpy(d_idata, h_idata, mem_size, hipMemcpyHostToDevice));\n\n // Compute reference transpose solution\n computeTransposeGold(transposeGold, h_idata, size_x, size_y);\n\n // print out common data for all kernels\n printf(\"\\nMatrix size: %dx%d (%dx%d tiles), tile size: %dx%d, block size: \"\n \"%dx%d\\n\\n\",\n size_x,\n size_y,\n size_x / TILE_DIM,\n size_y / TILE_DIM,\n TILE_DIM,\n TILE_DIM,\n TILE_DIM,\n BLOCK_ROWS);\n\n // initialize events\n checkCudaErrors(hipEventCreate(&start));\n checkCudaErrors(hipEventCreate(&stop));\n\n //\n // loop over different kernels\n //\n\n bool success = true;\n\n for (int k = 0; k < 8; k++) {\n // set kernel pointer\n switch (k) {\n case 0:\n kernel = ©\n kernelName = \"simple copy \";\n break;\n\n case 1:\n kernel = ©SharedMem;\n kernelName = \"shared memory copy\";\n break;\n\n case 2:\n kernel = &transposeNaive;\n kernelName = \"naive \";\n break;\n\n case 3:\n kernel = &transposeCoalesced;\n kernelName = \"coalesced \";\n break;\n\n case 4:\n kernel = &transposeNoBankConflicts;\n kernelName = \"optimized \";\n break;\n\n case 5:\n kernel = &transposeCoarseGrained;\n kernelName = \"coarse-grained \";\n break;\n\n case 6:\n kernel = &transposeFineGrained;\n kernelName = \"fine-grained \";\n break;\n\n case 7:\n kernel = &transposeDiagonal;\n kernelName = \"diagonal \";\n break;\n }\n\n // set reference solution\n if (kernel == © || kernel == ©SharedMem) {\n gold = h_idata;\n }\n else if (kernel == &transposeCoarseGrained || kernel == &transposeFineGrained) {\n gold = h_odata; // fine- and coarse-grained kernels are not full\n // transposes, so bypass check\n }\n else {\n gold = transposeGold;\n }\n\n // Clear error status\n checkCudaErrors(hipGetLastError());\n\n // warmup to avoid timing startup\n kernel<<>>(d_odata, d_idata, size_x, size_y);\n\n // take measurements for loop over kernel launches\n checkCudaErrors(hipEventRecord(start, 0));\n\n for (int i = 0; i < NUM_REPS; i++) {\n kernel<<>>(d_odata, d_idata, size_x, size_y);\n // Ensure no launch failure\n checkCudaErrors(hipGetLastError());\n }\n\n checkCudaErrors(hipEventRecord(stop, 0));\n checkCudaErrors(hipEventSynchronize(stop));\n float kernelTime;\n checkCudaErrors(hipEventElapsedTime(&kernelTime, start, stop));\n\n checkCudaErrors(hipMemcpy(h_odata, d_odata, mem_size, hipMemcpyDeviceToHost));\n bool res = compareData(gold, h_odata, size_x * size_y, 0.01f, 0.0f);\n\n if (res == false) {\n printf(\"*** %s kernel FAILED ***\\n\", kernelName);\n success = false;\n }\n\n // take measurements for loop inside kernel\n checkCudaErrors(hipMemcpy(h_odata, d_odata, mem_size, hipMemcpyDeviceToHost));\n res = compareData(gold, h_odata, size_x * size_y, 0.01f, 0.0f);\n\n if (res == false) {\n printf(\"*** %s kernel FAILED ***\\n\", kernelName);\n success = false;\n }\n\n // report effective bandwidths\n float kernelBandwidth = 2.0f * 1000.0f * mem_size / (1024 * 1024 * 1024) / (kernelTime / NUM_REPS);\n printf(\"transpose %s, Throughput = %.4f GB/s, Time = %.5f ms, Size = %u fp32 \"\n \"elements, NumDevsUsed = %u, Workgroup = %u\\n\",\n kernelName,\n kernelBandwidth,\n kernelTime / NUM_REPS,\n (size_x * size_y),\n 1,\n TILE_DIM * BLOCK_ROWS);\n\n // Reset d_odata to zero before starting the next loop iteration to avoid\n // carrying over results from previous kernels. Without this reset, residual\n // data from a prior kernel (e.g., 'copy') could make a subsequent\n // kernel (e.g., 'copySharedMem') appear correct even if it performs no work,\n // leading to false positives in compareData.\n for (int i = 0; i < (size_x * size_y); ++i) {\n h_odata[i] = 0;\n }\n // copy host data to device\n checkCudaErrors(hipMemcpy(d_odata, h_odata, mem_size, hipMemcpyHostToDevice));\n\n // Reset d_odata to zero before starting the next loop iteration to avoid\n // carrying over results from previous kernels. Without this reset, residual\n // data from a prior kernel (e.g., 'copy') could make a subsequent\n // kernel (e.g., 'copySharedMem') appear correct even if it performs no work,\n // leading to false positives in compareData.\n for (int i = 0; i < (size_x * size_y); ++i) {\n h_odata[i] = 0;\n }\n // copy host data to device\n checkCudaErrors(hipMemcpy(d_odata, h_odata, mem_size, hipMemcpyHostToDevice));\n }\n\n // cleanup\n free(h_idata);\n free(h_odata);\n free(transposeGold);\n hipFree(d_idata);\n hipFree(d_odata);\n\n checkCudaErrors(hipEventDestroy(start));\n checkCudaErrors(hipEventDestroy(stop));\n\n if (!success) {\n printf(\"Test failed!\\n\");\n exit(EXIT_FAILURE);\n }\n\n printf(\"Test passed\\n\");\n exit(EXIT_SUCCESS);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/6_Performance/transpose/transpose.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This is a simple test showing huge access speed gap\n * between aligned and misaligned structures\n * (those having/missing __align__ keyword).\n * It measures per-element copy throughput for\n * aligned and misaligned structures on\n * big chunks of data.\n */\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// includes, project\n#include // helper functions for CUDA error checking and initialization\n#include // helper utility functions\n\n////////////////////////////////////////////////////////////////////////////////\n// Misaligned types\n////////////////////////////////////////////////////////////////////////////////\ntypedef unsigned char uint8;\n\ntypedef unsigned short int uint16;\n\ntypedef struct\n{\n unsigned char r, g, b, a;\n} RGBA8_misaligned;\n\ntypedef struct\n{\n unsigned int l, a;\n} LA32_misaligned;\n\ntypedef struct\n{\n unsigned int r, g, b;\n} RGB32_misaligned;\n\ntypedef struct\n{\n unsigned int r, g, b, a;\n} RGBA32_misaligned;\n\n////////////////////////////////////////////////////////////////////////////////\n// Aligned types\n////////////////////////////////////////////////////////////////////////////////\ntypedef struct __align__(4)\n{\n unsigned char r, g, b, a;\n} RGBA8;\n\ntypedef unsigned int I32;\n\ntypedef struct __align__(8)\n{\n unsigned int l, a;\n} LA32;\n\ntypedef struct __align__(16)\n{\n unsigned int r, g, b;\n} RGB32;\n\ntypedef struct __align__(16)\n{\n unsigned int r, g, b, a;\n} RGBA32;\n\n////////////////////////////////////////////////////////////////////////////////\n// Because G80 class hardware natively supports global memory operations\n// only with data elements of 4, 8 and 16 bytes, if structure size\n// exceeds 16 bytes, it can't be efficiently read or written,\n// since more than one global memory non-coalescable load/store instructions\n// will be generated, even if __align__ option is supplied.\n// \"Structure of arrays\" storage strategy offers best performance\n// in general case. See section 5.1.2 of the Programming Guide.\n////////////////////////////////////////////////////////////////////////////////\ntypedef struct __align__(16)\n{\n RGBA32 c1, c2;\n} RGBA32_2;\n\n////////////////////////////////////////////////////////////////////////////////\n// Common host and device functions\n////////////////////////////////////////////////////////////////////////////////\n// Round a / b to nearest higher integer value\nint iDivUp(int a, int b) { return (a % b != 0) ? (a / b + 1) : (a / b); }\n\n// Round a / b to nearest lower integer value\nint iDivDown(int a, int b) { return a / b; }\n\n// Align a to nearest higher multiple of b\nint iAlignUp(int a, int b) { return (a % b != 0) ? (a - a % b + b) : a; }\n\n// Align a to nearest lower multiple of b\nint iAlignDown(int a, int b) { return a - a % b; }\n\n////////////////////////////////////////////////////////////////////////////////\n// Simple CUDA kernel.\n// Copy is carried out on per-element basis,\n// so it's not per-byte in case of padded structures.\n////////////////////////////////////////////////////////////////////////////////\ntemplate __global__ void testKernel(TData *d_odata, TData *d_idata, int numElements)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n const int numThreads = blockDim.x * gridDim.x;\n\n for (int pos = tid; pos < numElements; pos += numThreads) {\n d_odata[pos] = d_idata[pos];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Validation routine for simple copy kernel.\n// We must know \"packed\" size of TData (number_of_fields * sizeof(simple_type))\n// and compare only these \"packed\" parts of the structure,\n// containing actual user data. The compiler behavior with padding bytes\n// is undefined, since padding is merely a placeholder\n// and doesn't contain any user data.\n////////////////////////////////////////////////////////////////////////////////\ntemplate int testCPU(TData *h_odata, TData *h_idata, int numElements, int packedElementSize)\n{\n for (int pos = 0; pos < numElements; pos++) {\n TData src = h_idata[pos];\n TData dst = h_odata[pos];\n\n for (int i = 0; i < packedElementSize; i++)\n if (((char *)&src)[i] != ((char *)&dst)[i]) {\n return 0;\n }\n }\n\n return 1;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Data configuration\n////////////////////////////////////////////////////////////////////////////////\n// Memory chunk size in bytes. Reused for test\nconst int MEM_SIZE = 50000000;\nconst int NUM_ITERATIONS = 32;\n\n// GPU input and output data\nunsigned char *d_idata, *d_odata;\n// CPU input data and instance of GPU output data\nunsigned char *h_idataCPU, *h_odataGPU;\nStopWatchInterface *hTimer = NULL;\n\ntemplate int runTest(int packedElementSize, int memory_size)\n{\n const int totalMemSizeAligned = iAlignDown(memory_size, sizeof(TData));\n const int numElements = iDivDown(memory_size, sizeof(TData));\n\n // Clean output buffer before current test\n checkCudaErrors(cudaMemset(d_odata, 0, memory_size));\n // Run test\n checkCudaErrors(cudaDeviceSynchronize());\n sdkResetTimer(&hTimer);\n sdkStartTimer(&hTimer);\n\n for (int i = 0; i < NUM_ITERATIONS; i++) {\n testKernel<<<64, 256>>>((TData *)d_odata, (TData *)d_idata, numElements);\n getLastCudaError(\"testKernel() execution failed\\n\");\n }\n\n checkCudaErrors(cudaDeviceSynchronize());\n sdkStopTimer(&hTimer);\n double gpuTime = sdkGetTimerValue(&hTimer) / NUM_ITERATIONS;\n printf(\"Avg. time: %f ms / Copy throughput: %f GB/s.\\n\",\n gpuTime,\n (double)totalMemSizeAligned / (gpuTime * 0.001 * 1073741824.0));\n\n // Read back GPU results and run validation\n checkCudaErrors(cudaMemcpy(h_odataGPU, d_odata, memory_size, cudaMemcpyDeviceToHost));\n int flag = testCPU((TData *)h_odataGPU, (TData *)h_idataCPU, numElements, packedElementSize);\n\n printf(flag ? \"\\tTEST OK\\n\" : \"\\tTEST FAILURE\\n\");\n\n return !flag;\n}\n\nint main(int argc, char **argv)\n{\n int i, nTotalFailures = 0;\n\n int devID;\n cudaDeviceProp deviceProp;\n printf(\"[%s] - Starting...\\n\", argv[0]);\n\n // find first CUDA device\n devID = findCudaDevice(argc, (const char **)argv);\n\n // get number of SMs on this GPU\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, devID));\n printf(\"[%s] has %d MP(s) x %d (Cores/MP) = %d (Cores)\\n\",\n deviceProp.name,\n deviceProp.multiProcessorCount,\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor),\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * deviceProp.multiProcessorCount);\n\n // Anything that is less than 192 Cores will have a scaled down workload\n float scale_factor = max(\n (192.0f / (_ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * (float)deviceProp.multiProcessorCount)),\n 1.0f);\n\n int MemorySize = (int)(MEM_SIZE / scale_factor) & 0xffffff00; // force multiple of 256 bytes\n\n printf(\"> Compute scaling value = %4.2f\\n\", scale_factor);\n printf(\"> Memory Size = %d\\n\", MemorySize);\n\n sdkCreateTimer(&hTimer);\n\n printf(\"Allocating memory...\\n\");\n h_idataCPU = (unsigned char *)malloc(MemorySize);\n h_odataGPU = (unsigned char *)malloc(MemorySize);\n checkCudaErrors(cudaMalloc((void **)&d_idata, MemorySize));\n checkCudaErrors(cudaMalloc((void **)&d_odata, MemorySize));\n\n printf(\"Generating host input data array...\\n\");\n\n for (i = 0; i < MemorySize; i++) {\n h_idataCPU[i] = (i & 0xFF) + 1;\n }\n\n printf(\"Uploading input data to GPU memory...\\n\");\n checkCudaErrors(cudaMemcpy(d_idata, h_idataCPU, MemorySize, cudaMemcpyHostToDevice));\n\n printf(\"Testing misaligned types...\\n\");\n printf(\"uint8...\\n\");\n nTotalFailures += runTest(1, MemorySize);\n\n printf(\"uint16...\\n\");\n nTotalFailures += runTest(2, MemorySize);\n\n printf(\"RGBA8_misaligned...\\n\");\n nTotalFailures += runTest(4, MemorySize);\n\n printf(\"LA32_misaligned...\\n\");\n nTotalFailures += runTest(8, MemorySize);\n\n printf(\"RGB32_misaligned...\\n\");\n nTotalFailures += runTest(12, MemorySize);\n\n printf(\"RGBA32_misaligned...\\n\");\n nTotalFailures += runTest(16, MemorySize);\n\n printf(\"Testing aligned types...\\n\");\n printf(\"RGBA8...\\n\");\n nTotalFailures += runTest(4, MemorySize);\n\n printf(\"I32...\\n\");\n nTotalFailures += runTest(4, MemorySize);\n\n printf(\"LA32...\\n\");\n nTotalFailures += runTest(8, MemorySize);\n\n printf(\"RGB32...\\n\");\n nTotalFailures += runTest(12, MemorySize);\n\n printf(\"RGBA32...\\n\");\n nTotalFailures += runTest(16, MemorySize);\n\n printf(\"RGBA32_2...\\n\");\n nTotalFailures += runTest(32, MemorySize);\n\n printf(\"\\n[alignedTypes] -> Test Results: %d Failures\\n\", nTotalFailures);\n\n printf(\"Shutting down...\\n\");\n checkCudaErrors(cudaFree(d_idata));\n checkCudaErrors(cudaFree(d_odata));\n free(h_odataGPU);\n free(h_idataCPU);\n\n sdkDeleteTimer(&hTimer);\n\n if (nTotalFailures != 0) {\n printf(\"Test failed!\\n\");\n exit(EXIT_FAILURE);\n }\n\n printf(\"Test passed\\n\");\n exit(EXIT_SUCCESS);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This is a simple test showing huge access speed gap\n * between aligned and misaligned structures\n * (those having/missing __align__ keyword).\n * It measures per-element copy throughput for\n * aligned and misaligned structures on\n * big chunks of data.\n */\n\n// includes, system\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n// includes, project\n#include // helper functions for CUDA error checking and initialization\n#include // helper utility functions\n\n////////////////////////////////////////////////////////////////////////////////\n// Misaligned types\n////////////////////////////////////////////////////////////////////////////////\ntypedef unsigned char uint8;\n\ntypedef unsigned short int uint16;\n\ntypedef struct\n{\n unsigned char r, g, b, a;\n} RGBA8_misaligned;\n\ntypedef struct\n{\n unsigned int l, a;\n} LA32_misaligned;\n\ntypedef struct\n{\n unsigned int r, g, b;\n} RGB32_misaligned;\n\ntypedef struct\n{\n unsigned int r, g, b, a;\n} RGBA32_misaligned;\n\n////////////////////////////////////////////////////////////////////////////////\n// Aligned types\n////////////////////////////////////////////////////////////////////////////////\ntypedef struct __align__(4)\n{\n unsigned char r, g, b, a;\n} RGBA8;\n\ntypedef unsigned int I32;\n\ntypedef struct __align__(8)\n{\n unsigned int l, a;\n} LA32;\n\ntypedef struct __align__(16)\n{\n unsigned int r, g, b;\n} RGB32;\n\ntypedef struct __align__(16)\n{\n unsigned int r, g, b, a;\n} RGBA32;\n\n////////////////////////////////////////////////////////////////////////////////\n// Because G80 class hardware natively supports global memory operations\n// only with data elements of 4, 8 and 16 bytes, if structure size\n// exceeds 16 bytes, it can't be efficiently read or written,\n// since more than one global memory non-coalescable load/store instructions\n// will be generated, even if __align__ option is supplied.\n// \"Structure of arrays\" storage strategy offers best performance\n// in general case. See section 5.1.2 of the Programming Guide.\n////////////////////////////////////////////////////////////////////////////////\ntypedef struct __align__(16)\n{\n RGBA32 c1, c2;\n} RGBA32_2;\n\n////////////////////////////////////////////////////////////////////////////////\n// Common host and device functions\n////////////////////////////////////////////////////////////////////////////////\n// Round a / b to nearest higher integer value\nint iDivUp(int a, int b) { return (a % b != 0) ? (a / b + 1) : (a / b); }\n\n// Round a / b to nearest lower integer value\nint iDivDown(int a, int b) { return a / b; }\n\n// Align a to nearest higher multiple of b\nint iAlignUp(int a, int b) { return (a % b != 0) ? (a - a % b + b) : a; }\n\n// Align a to nearest lower multiple of b\nint iAlignDown(int a, int b) { return a - a % b; }\n\n////////////////////////////////////////////////////////////////////////////////\n// Simple CUDA kernel.\n// Copy is carried out on per-element basis,\n// so it's not per-byte in case of padded structures.\n////////////////////////////////////////////////////////////////////////////////\ntemplate __global__ void testKernel(TData *d_odata, TData *d_idata, int numElements)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n const int numThreads = blockDim.x * gridDim.x;\n\n for (int pos = tid; pos < numElements; pos += numThreads) {\n d_odata[pos] = d_idata[pos];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Validation routine for simple copy kernel.\n// We must know \"packed\" size of TData (number_of_fields * sizeof(simple_type))\n// and compare only these \"packed\" parts of the structure,\n// containing actual user data. The compiler behavior with padding bytes\n// is undefined, since padding is merely a placeholder\n// and doesn't contain any user data.\n////////////////////////////////////////////////////////////////////////////////\ntemplate int testCPU(TData *h_odata, TData *h_idata, int numElements, int packedElementSize)\n{\n for (int pos = 0; pos < numElements; pos++) {\n TData src = h_idata[pos];\n TData dst = h_odata[pos];\n\n for (int i = 0; i < packedElementSize; i++)\n if (((char *)&src)[i] != ((char *)&dst)[i]) {\n return 0;\n }\n }\n\n return 1;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Data configuration\n////////////////////////////////////////////////////////////////////////////////\n// Memory chunk size in bytes. Reused for test\nconst int MEM_SIZE = 50000000;\nconst int NUM_ITERATIONS = 32;\n\n// GPU input and output data\nunsigned char *d_idata, *d_odata;\n// CPU input data and instance of GPU output data\nunsigned char *h_idataCPU, *h_odataGPU;\nStopWatchInterface *hTimer = NULL;\n\ntemplate int runTest(int packedElementSize, int memory_size)\n{\n const int totalMemSizeAligned = iAlignDown(memory_size, sizeof(TData));\n const int numElements = iDivDown(memory_size, sizeof(TData));\n\n // Clean output buffer before current test\n checkCudaErrors(hipMemset(d_odata, 0, memory_size));\n // Run test\n checkCudaErrors(hipDeviceSynchronize());\n sdkResetTimer(&hTimer);\n sdkStartTimer(&hTimer);\n\n for (int i = 0; i < NUM_ITERATIONS; i++) {\n testKernel<<<64, 256>>>((TData *)d_odata, (TData *)d_idata, numElements);\n getLastCudaError(\"testKernel() execution failed\\n\");\n }\n\n checkCudaErrors(hipDeviceSynchronize());\n sdkStopTimer(&hTimer);\n double gpuTime = sdkGetTimerValue(&hTimer) / NUM_ITERATIONS;\n printf(\"Avg. time: %f ms / Copy throughput: %f GB/s.\\n\",\n gpuTime,\n (double)totalMemSizeAligned / (gpuTime * 0.001 * 1073741824.0));\n\n // Read back GPU results and run validation\n checkCudaErrors(hipMemcpy(h_odataGPU, d_odata, memory_size, hipMemcpyDeviceToHost));\n int flag = testCPU((TData *)h_odataGPU, (TData *)h_idataCPU, numElements, packedElementSize);\n\n printf(flag ? \"\\tTEST OK\\n\" : \"\\tTEST FAILURE\\n\");\n\n return !flag;\n}\n\nint main(int argc, char **argv)\n{\n int i, nTotalFailures = 0;\n\n int devID;\n hipDeviceProp_t deviceProp;\n printf(\"[%s] - Starting...\\n\", argv[0]);\n\n // find first CUDA device\n devID = findCudaDevice(argc, (const char **)argv);\n\n // get number of SMs on this GPU\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, devID));\n printf(\"[%s] has %d MP(s) x %d (Cores/MP) = %d (Cores)\\n\",\n deviceProp.name,\n deviceProp.multiProcessorCount,\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor),\n _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * deviceProp.multiProcessorCount);\n\n // Anything that is less than 192 Cores will have a scaled down workload\n float scale_factor = max(\n (192.0f / (_ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * (float)deviceProp.multiProcessorCount)),\n 1.0f);\n\n int MemorySize = (int)(MEM_SIZE / scale_factor) & 0xffffff00; // force multiple of 256 bytes\n\n printf(\"> Compute scaling value = %4.2f\\n\", scale_factor);\n printf(\"> Memory Size = %d\\n\", MemorySize);\n\n sdkCreateTimer(&hTimer);\n\n printf(\"Allocating memory...\\n\");\n h_idataCPU = (unsigned char *)malloc(MemorySize);\n h_odataGPU = (unsigned char *)malloc(MemorySize);\n checkCudaErrors(hipMalloc((void **)&d_idata, MemorySize));\n checkCudaErrors(hipMalloc((void **)&d_odata, MemorySize));\n\n printf(\"Generating host input data array...\\n\");\n\n for (i = 0; i < MemorySize; i++) {\n h_idataCPU[i] = (i & 0xFF) + 1;\n }\n\n printf(\"Uploading input data to GPU memory...\\n\");\n checkCudaErrors(hipMemcpy(d_idata, h_idataCPU, MemorySize, hipMemcpyHostToDevice));\n\n printf(\"Testing misaligned types...\\n\");\n printf(\"uint8...\\n\");\n nTotalFailures += runTest(1, MemorySize);\n\n printf(\"uint16...\\n\");\n nTotalFailures += runTest(2, MemorySize);\n\n printf(\"RGBA8_misaligned...\\n\");\n nTotalFailures += runTest(4, MemorySize);\n\n printf(\"LA32_misaligned...\\n\");\n nTotalFailures += runTest(8, MemorySize);\n\n printf(\"RGB32_misaligned...\\n\");\n nTotalFailures += runTest(12, MemorySize);\n\n printf(\"RGBA32_misaligned...\\n\");\n nTotalFailures += runTest(16, MemorySize);\n\n printf(\"Testing aligned types...\\n\");\n printf(\"RGBA8...\\n\");\n nTotalFailures += runTest(4, MemorySize);\n\n printf(\"I32...\\n\");\n nTotalFailures += runTest(4, MemorySize);\n\n printf(\"LA32...\\n\");\n nTotalFailures += runTest(8, MemorySize);\n\n printf(\"RGB32...\\n\");\n nTotalFailures += runTest(12, MemorySize);\n\n printf(\"RGBA32...\\n\");\n nTotalFailures += runTest(16, MemorySize);\n\n printf(\"RGBA32_2...\\n\");\n nTotalFailures += runTest(32, MemorySize);\n\n printf(\"\\n[alignedTypes] -> Test Results: %d Failures\\n\", nTotalFailures);\n\n printf(\"Shutting down...\\n\");\n checkCudaErrors(hipFree(d_idata));\n checkCudaErrors(hipFree(d_odata));\n free(h_odataGPU);\n free(h_idataCPU);\n\n sdkDeleteTimer(&hTimer);\n\n if (nTotalFailures != 0) {\n printf(\"Test failed!\\n\");\n exit(EXIT_FAILURE);\n }\n\n printf(\"Test passed\\n\");\n exit(EXIT_SUCCESS);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/6_Performance/alignedTypes/alignedTypes.cu"} {"input": "/* Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This is a simple application showing the performance characteristics of cudaGraphs.\n */\n\n#define USE_NVTX\n\n#include \n#include \n#include \n#include \n\ntypedef volatile int LatchType;\n\nstd::chrono::time_point getCpuTime()\n{\n return std::chrono::high_resolution_clock::now();\n}\n\ntemplate float getMicroSecondDuration(T start, T end)\n{\n return std::chrono::duration_cast(end - start).count() * .001f;\n}\n\nfloat getAsyncMicroSecondDuration(cudaEvent_t start, cudaEvent_t end)\n{\n float ms;\n cudaEventElapsedTime(&ms, start, end);\n return ms * 1000;\n}\n\n#ifdef USE_NVTX\n#include \n\nclass Tracer\n{\npublic:\n Tracer(const char *name) { nvtxRangePushA(name); }\n ~Tracer() { nvtxRangePop(); }\n};\n#define RANGE(name) Tracer uniq_name_using_macros(name);\n#define RANGE_PUSH(name) nvtxRangePushA(name)\n#define RANGE_POP() nvtxRangePop();\n#else\n#define RANGE(name)\n#endif\n\nstd::vector stream;\ncudaEvent_t event[1];\ncudaEvent_t timingEvent[2];\n\nstruct hostData\n{\n long long timeElapsed;\n bool timeoutDetected;\n long long timeElapsed2;\n bool timeoutDetected2;\n LatchType latch;\n LatchType latch2;\n};\n\nstruct hostData *hostData;\n\n__global__ void empty() {}\n\n// Function to read the GPU nanosecond timer in a kernel\n__device__ __forceinline__ unsigned long long __globaltimer()\n{\n unsigned long long globaltimer;\n asm volatile(\"mov.u64 %0, %globaltimer;\" : \"=l\"(globaltimer));\n return globaltimer;\n}\n\n__global__ void delay(long long ticks)\n{\n long long endTime = clock64() + ticks;\n while (clock64() < endTime)\n ;\n}\n\n__global__ void waitWithTimeout(long long nanoseconds, bool *timeoutDetected, long long *timeElapsed, LatchType *latch)\n{\n long long startTime = __globaltimer();\n long long endTime = startTime + nanoseconds;\n long long time = 0;\n do {\n time = __globaltimer();\n } while (time < endTime && (latch == NULL || *latch == 0));\n if (timeElapsed != NULL) {\n *timeElapsed = time - startTime;\n }\n if (timeoutDetected) {\n // report timeout if latch not detected\n *timeoutDetected = (latch == NULL || *latch == 0);\n }\n}\n\n__global__ void preUploadAnnotation() {}\n\n__global__ void postUploadAnnotation() {}\n\ncudaGraph_t createParallelChain(int length, int width, bool singleEntry = false)\n{\n RANGE_PUSH(__func__);\n RANGE(\"capture\");\n cudaGraph_t graph;\n cudaStreamBeginCapture(stream[0], cudaStreamCaptureModeGlobal);\n int streamIdx = 0;\n if (singleEntry) {\n empty<<<1, 1, 0, stream[streamIdx]>>>();\n }\n\n cudaEventRecord(event[0], stream[0]);\n for (int i = 1; i < width; i++) {\n cudaStreamWaitEvent(stream[i], event[0]);\n }\n\n for (int i = 0; i < width; i++) {\n streamIdx = i;\n for (int j = 0; j < length; j++) {\n empty<<<1, 1, 0, stream[streamIdx]>>>();\n }\n }\n\n for (int i = 1; i < width; i++) {\n cudaEventRecord(event[0], stream[i]);\n cudaStreamWaitEvent(stream[0], event[0]);\n }\n\n cudaStreamEndCapture(stream[0], &graph);\n return graph;\n}\n\nstd::vector metricName;\nstd::vector metricValue;\n\nint counter2 = 0;\nvoid runDemo(cudaGraph_t graph, int length, int width)\n{\n cudaGraphExec_t graphExec;\n {\n auto start = getCpuTime();\n cudaGraphInstantiateWithFlags(&graphExec, graph, 0);\n auto end = getCpuTime();\n metricName.push_back(\"instantiation\");\n metricValue.push_back(getMicroSecondDuration(start, end));\n }\n {\n RANGE(\"launch including upload\");\n auto start = getCpuTime();\n cudaGraphLaunch(graphExec, stream[0]);\n auto apiReturn = getCpuTime();\n cudaStreamSynchronize(stream[0]);\n auto streamSync = getCpuTime();\n metricName.push_back(\"first_launch_api\");\n metricValue.push_back(getMicroSecondDuration(start, apiReturn));\n metricName.push_back(\"first_launch_total\");\n metricValue.push_back(getMicroSecondDuration(start, streamSync));\n }\n {\n RANGE(\"repeat lauch in empty stream\");\n auto start = getCpuTime();\n cudaGraphLaunch(graphExec, stream[0]);\n auto apiReturn = getCpuTime();\n cudaStreamSynchronize(stream[0]);\n auto streamSync = getCpuTime();\n metricName.push_back(\"repeat_launch_api\");\n metricValue.push_back(getMicroSecondDuration(start, apiReturn));\n metricName.push_back(\"repeat_launch_total\");\n metricValue.push_back(getMicroSecondDuration(start, streamSync));\n }\n {\n // re-instantiating the exec to simulate first launch into a busy stream.\n cudaGraphExecDestroy(graphExec);\n cudaGraphInstantiateWithFlags(&graphExec, graph, 0);\n\n long long maxTimeoutNanoSeconds = 4000 + 500 * length * width;\n waitWithTimeout<<<1, 1, 0, stream[0]>>>(\n maxTimeoutNanoSeconds, &hostData->timeoutDetected, &hostData->timeElapsed, &hostData->latch);\n\n RANGE(\"launch including upload in busy stream\");\n cudaEventRecord(timingEvent[0], stream[0]);\n cudaGraphLaunch(graphExec, stream[0]);\n cudaEventRecord(timingEvent[1], stream[0]);\n\n hostData->latch = 1;\n cudaStreamSynchronize(stream[0]);\n\n metricName.push_back(\"first_launch_device\");\n metricValue.push_back(getAsyncMicroSecondDuration(timingEvent[0], timingEvent[1]));\n metricName.push_back(\"blockingKernelTimeoutDetected\");\n metricValue.push_back(hostData->timeoutDetected);\n hostData->latch = 0;\n hostData->timeoutDetected = 0;\n }\n {\n RANGE(\"repeat lauch in busy stream\");\n long long maxTimeoutNanoSeconds = 4000 + 500 * length * width;\n waitWithTimeout<<<1, 1, 0, stream[0]>>>(\n maxTimeoutNanoSeconds, &hostData->timeoutDetected, &hostData->timeElapsed, &hostData->latch);\n cudaEventRecord(timingEvent[0], stream[0]);\n cudaGraphLaunch(graphExec, stream[0]);\n cudaEventRecord(timingEvent[1], stream[0]);\n\n hostData->latch = 1;\n cudaStreamSynchronize(stream[0]);\n\n metricName.push_back(\"repeat_launch_device\");\n metricValue.push_back(getAsyncMicroSecondDuration(timingEvent[0], timingEvent[1]));\n metricName.push_back(\"blockingKernelTimeoutDetected\");\n metricValue.push_back(hostData->timeoutDetected);\n hostData->latch = 0;\n hostData->timeoutDetected = 0;\n }\n {\n // re-instantiating the exec to provide upload with work to do.\n cudaGraphExecDestroy(graphExec);\n cudaGraphInstantiateWithFlags(&graphExec, graph, 0);\n long long maxTimeoutNanoSeconds = 4000 + 1000 * length * width;\n waitWithTimeout<<<1, 1, 0, stream[0]>>>(\n maxTimeoutNanoSeconds, &hostData->timeoutDetected2, &hostData->timeElapsed2, &hostData->latch2);\n maxTimeoutNanoSeconds = 2000 + 500 * length * width;\n waitWithTimeout<<<1, 1, 0, stream[1]>>>(\n maxTimeoutNanoSeconds, &hostData->timeoutDetected, &hostData->timeElapsed, &hostData->latch);\n\n RANGE(\"uploading a graph off of the critical path\");\n preUploadAnnotation<<<1, 1, 0, stream[1]>>>();\n cudaEventRecord(timingEvent[0], stream[0]);\n auto start = getCpuTime();\n cudaGraphUpload(graphExec, stream[1]);\n auto apiReturn = getCpuTime();\n cudaEventRecord(event[0], stream[1]);\n cudaEventRecord(timingEvent[1], stream[0]);\n postUploadAnnotation<<<1, 1, 0, stream[1]>>>();\n\n hostData->latch = 1; // release the blocking kernel for the upload\n cudaStreamWaitEvent(stream[0], event[0]);\n cudaGraphLaunch(graphExec, stream[0]);\n cudaEventSynchronize(event[0]); // upload done, similuate critical path being ready for the graph to run by the\n // release of the second latch\n\n hostData->latch2 = 1; // release the work\n cudaStreamSynchronize(stream[0]);\n\n metricName.push_back(\"upload_api_time\");\n metricValue.push_back(getMicroSecondDuration(start, apiReturn));\n metricName.push_back(\"updoad_device_time\");\n metricValue.push_back(getAsyncMicroSecondDuration(timingEvent[0], timingEvent[1]));\n metricName.push_back(\"blockingKernelTimeoutDetected\");\n metricValue.push_back(hostData->timeoutDetected);\n\n hostData->latch = 0;\n hostData->latch2 = 0;\n hostData->timeoutDetected = 0;\n hostData->timeoutDetected2 = 0;\n }\n cudaGraphExecDestroy(graphExec);\n cudaGraphDestroy(graph);\n RANGE_POP();\n}\n\nvoid usage()\n{\n printf(\"programName [outputFmt] [numTrials] [length] [width] [pattern] [stride] [maxLength] \\n\");\n printf(\"\\toutputFmt - program output, default=3 (see below)\\n\");\n printf(\"\\tnumTrials (per length)\\n\");\n printf(\"\\tstarting length of the topology\\n\");\n printf(\"\\twidth - width of the graph topology\\n\");\n printf(\"\\tpattern - Structure of graph, default=0 (see below)\\n\");\n printf(\"\\tstride - how to grow the length between each set of trials \\n\");\n printf(\"\\tmaxLength - maximum lenght to try \\n\");\n printf(\"\\n\");\n printf(\"outputFmt can be:\\n\");\n printf(\"\\t0: this help message\\n\");\n printf(\"\\t1: csv data headers\\n\");\n printf(\"\\t2: per trial csv data\\n\");\n printf(\"\\t3: csv data & headers\\n\");\n printf(\"\\t4: csv data is printed and trials are averaged for each length\\n\");\n printf(\"\\t5: csv data is printed and trials are averaged for each length and headers are printed\\n\");\n printf(\"\\n\");\n printf(\"Pattern can be:\\n\");\n printf(\"\\t0: No interconnect between branches\\n\");\n printf(\"\\t1: Adds an extra root node before the initial fork\\n\");\n}\n\nint main(int argc, char **argv)\n{\n if (argc < 1) {\n usage();\n return 0;\n }\n\n int numTrials = 1, length = 20, width = 1, outputFmt = 3, pattern = 0, stride = 1;\n if (argc > 1)\n outputFmt = atoi(argv[1]);\n if (argc > 2)\n numTrials = atoi(argv[2]);\n if (argc > 3)\n length = atoi(argv[3]);\n if (argc > 4)\n width = atoi(argv[4]);\n if (argc > 5)\n pattern = atoi(argv[5]);\n if (argc > 6)\n stride = atoi(argv[6]);\n int maxLength = length;\n if (argc > 7)\n maxLength = atoi(argv[7]);\n if (maxLength < length) {\n maxLength = length;\n }\n\n if ((outputFmt & 4) && (outputFmt & 2)) {\n printf(\"printing average and all samples doesn't make sense\\n\");\n }\n\n if (length == 0 || width == 0 || outputFmt == 0 || outputFmt > 5 || pattern > 1) {\n usage();\n return 0;\n }\n\n bool singleEntry = (pattern == 1);\n\n cudaGraph_t graph;\n\n cudaFree(0);\n cudaMallocHost(&hostData, sizeof(*hostData));\n int numStreams = width;\n if (numStreams == 1)\n numStreams = 2; // demo needs two streams even if capture only needs 1.\n stream.resize(numStreams);\n for (int i = 0; i < numStreams; i++) {\n cudaStreamCreate(&stream[i]);\n }\n\n cudaEventCreate(&event[0], cudaEventDisableTiming);\n cudaEventCreate(&timingEvent[0], 0);\n cudaEventCreate(&timingEvent[1], 0);\n\n {\n RANGE(\"warmup\");\n for (int i = 0; i < width; i++) {\n empty<<<1, 1, 0, stream[i]>>>();\n }\n cudaStreamSynchronize(stream[0]);\n\n auto start = getCpuTime();\n graph = createParallelChain(length, width, singleEntry);\n auto end = getCpuTime();\n metricValue.push_back(getMicroSecondDuration(start, end));\n metricName.push_back(\"capture\");\n runDemo(graph, length, width);\n }\n\n if (outputFmt & 1) {\n printf(\"length, width, pattern, \");\n for (int i = 0; i < metricName.size(); i++) {\n printf(\"%s, \", metricName[i]);\n }\n printf(\"\\r\\n\");\n }\n\n if (!(outputFmt & 6)) {\n printf(\"skipping trials since no output is expected\\n\");\n return 1;\n }\n\n std::vector metricTotal;\n metricTotal.resize(metricValue.size());\n\n while (length <= maxLength) {\n for (int i = 0; i < numTrials; i++) {\n metricName.clear();\n metricValue.clear();\n auto start = getCpuTime();\n graph = createParallelChain(length, width, singleEntry);\n auto end = getCpuTime();\n metricValue.push_back(getMicroSecondDuration(start, end));\n\n runDemo(graph, length, width);\n\n if (outputFmt & 2) {\n printf(\"%d, %d, %d, \", length, width, pattern);\n for (int i = 0; i < metricValue.size(); i++) {\n printf(\"%0.3f, \", metricValue[i]);\n }\n printf(\"\\r\\n\");\n }\n if (outputFmt & 4) {\n for (int i = 0; i < metricTotal.size(); i++) {\n metricTotal[i] += metricValue[i];\n }\n }\n }\n\n if (outputFmt & 4) {\n printf(\"%d, %d, %d, \", length, width, pattern);\n for (int i = 0; i < metricTotal.size(); i++) {\n printf(\"%0.3f, \", metricTotal[i] / numTrials);\n metricTotal[i] = 0;\n }\n printf(\"\\r\\n\");\n }\n\n length += stride;\n }\n\n cudaFreeHost(hostData);\n\n printf(\"\\n\");\n printf(\"Test passed\\n\");\n return 0;\n}\n", "output": "/* Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This is a simple application showing the performance characteristics of cudaGraphs.\n */\n\n#define USE_NVTX\n\n#include \n#include \n#include \n#include \n\ntypedef volatile int LatchType;\n\nstd::chrono::time_point getCpuTime()\n{\n return std::chrono::high_resolution_clock::now();\n}\n\ntemplate float getMicroSecondDuration(T start, T end)\n{\n return std::chrono::duration_cast(end - start).count() * .001f;\n}\n\nfloat getAsyncMicroSecondDuration(hipEvent_t start, hipEvent_t end)\n{\n float ms;\n hipEventElapsedTime(&ms, start, end);\n return ms * 1000;\n}\n\n#ifdef USE_NVTX\n#include \n\nclass Tracer\n{\npublic:\n Tracer(const char *name) { nvtxRangePushA(name); }\n ~Tracer() { nvtxRangePop(); }\n};\n#define RANGE(name) Tracer uniq_name_using_macros(name);\n#define RANGE_PUSH(name) nvtxRangePushA(name)\n#define RANGE_POP() nvtxRangePop();\n#else\n#define RANGE(name)\n#endif\n\nstd::vector stream;\nhipEvent_t event[1];\nhipEvent_t timingEvent[2];\n\nstruct hostData\n{\n long long timeElapsed;\n bool timeoutDetected;\n long long timeElapsed2;\n bool timeoutDetected2;\n LatchType latch;\n LatchType latch2;\n};\n\nstruct hostData *hostData;\n\n__global__ void empty() {}\n\n// Function to read the GPU nanosecond timer in a kernel\n__device__ __forceinline__ unsigned long long __globaltimer()\n{\n unsigned long long globaltimer;\n asm volatile(\"mov.u64 %0, %globaltimer;\" : \"=l\"(globaltimer));\n return globaltimer;\n}\n\n__global__ void delay(long long ticks)\n{\n long long endTime = clock64() + ticks;\n while (clock64() < endTime)\n ;\n}\n\n__global__ void waitWithTimeout(long long nanoseconds, bool *timeoutDetected, long long *timeElapsed, LatchType *latch)\n{\n long long startTime = __globaltimer();\n long long endTime = startTime + nanoseconds;\n long long time = 0;\n do {\n time = __globaltimer();\n } while (time < endTime && (latch == NULL || *latch == 0));\n if (timeElapsed != NULL) {\n *timeElapsed = time - startTime;\n }\n if (timeoutDetected) {\n // report timeout if latch not detected\n *timeoutDetected = (latch == NULL || *latch == 0);\n }\n}\n\n__global__ void preUploadAnnotation() {}\n\n__global__ void postUploadAnnotation() {}\n\nhipGraph_t createParallelChain(int length, int width, bool singleEntry = false)\n{\n RANGE_PUSH(__func__);\n RANGE(\"capture\");\n hipGraph_t graph;\n hipStreamBeginCapture(stream[0], hipStreamCaptureModeGlobal);\n int streamIdx = 0;\n if (singleEntry) {\n empty<<<1, 1, 0, stream[streamIdx]>>>();\n }\n\n hipEventRecord(event[0], stream[0]);\n for (int i = 1; i < width; i++) {\n hipStreamWaitEvent(stream[i], event[0]);\n }\n\n for (int i = 0; i < width; i++) {\n streamIdx = i;\n for (int j = 0; j < length; j++) {\n empty<<<1, 1, 0, stream[streamIdx]>>>();\n }\n }\n\n for (int i = 1; i < width; i++) {\n hipEventRecord(event[0], stream[i]);\n hipStreamWaitEvent(stream[0], event[0]);\n }\n\n hipStreamEndCapture(stream[0], &graph);\n return graph;\n}\n\nstd::vector metricName;\nstd::vector metricValue;\n\nint counter2 = 0;\nvoid runDemo(hipGraph_t graph, int length, int width)\n{\n hipGraphExec_t graphExec;\n {\n auto start = getCpuTime();\n hipGraphInstantiateWithFlags(&graphExec, graph, 0);\n auto end = getCpuTime();\n metricName.push_back(\"instantiation\");\n metricValue.push_back(getMicroSecondDuration(start, end));\n }\n {\n RANGE(\"launch including upload\");\n auto start = getCpuTime();\n hipGraphLaunch(graphExec, stream[0]);\n auto apiReturn = getCpuTime();\n hipStreamSynchronize(stream[0]);\n auto streamSync = getCpuTime();\n metricName.push_back(\"first_launch_api\");\n metricValue.push_back(getMicroSecondDuration(start, apiReturn));\n metricName.push_back(\"first_launch_total\");\n metricValue.push_back(getMicroSecondDuration(start, streamSync));\n }\n {\n RANGE(\"repeat lauch in empty stream\");\n auto start = getCpuTime();\n hipGraphLaunch(graphExec, stream[0]);\n auto apiReturn = getCpuTime();\n hipStreamSynchronize(stream[0]);\n auto streamSync = getCpuTime();\n metricName.push_back(\"repeat_launch_api\");\n metricValue.push_back(getMicroSecondDuration(start, apiReturn));\n metricName.push_back(\"repeat_launch_total\");\n metricValue.push_back(getMicroSecondDuration(start, streamSync));\n }\n {\n // re-instantiating the exec to simulate first launch into a busy stream.\n hipGraphExecDestroy(graphExec);\n hipGraphInstantiateWithFlags(&graphExec, graph, 0);\n\n long long maxTimeoutNanoSeconds = 4000 + 500 * length * width;\n waitWithTimeout<<<1, 1, 0, stream[0]>>>(\n maxTimeoutNanoSeconds, &hostData->timeoutDetected, &hostData->timeElapsed, &hostData->latch);\n\n RANGE(\"launch including upload in busy stream\");\n hipEventRecord(timingEvent[0], stream[0]);\n hipGraphLaunch(graphExec, stream[0]);\n hipEventRecord(timingEvent[1], stream[0]);\n\n hostData->latch = 1;\n hipStreamSynchronize(stream[0]);\n\n metricName.push_back(\"first_launch_device\");\n metricValue.push_back(getAsyncMicroSecondDuration(timingEvent[0], timingEvent[1]));\n metricName.push_back(\"blockingKernelTimeoutDetected\");\n metricValue.push_back(hostData->timeoutDetected);\n hostData->latch = 0;\n hostData->timeoutDetected = 0;\n }\n {\n RANGE(\"repeat lauch in busy stream\");\n long long maxTimeoutNanoSeconds = 4000 + 500 * length * width;\n waitWithTimeout<<<1, 1, 0, stream[0]>>>(\n maxTimeoutNanoSeconds, &hostData->timeoutDetected, &hostData->timeElapsed, &hostData->latch);\n hipEventRecord(timingEvent[0], stream[0]);\n hipGraphLaunch(graphExec, stream[0]);\n hipEventRecord(timingEvent[1], stream[0]);\n\n hostData->latch = 1;\n hipStreamSynchronize(stream[0]);\n\n metricName.push_back(\"repeat_launch_device\");\n metricValue.push_back(getAsyncMicroSecondDuration(timingEvent[0], timingEvent[1]));\n metricName.push_back(\"blockingKernelTimeoutDetected\");\n metricValue.push_back(hostData->timeoutDetected);\n hostData->latch = 0;\n hostData->timeoutDetected = 0;\n }\n {\n // re-instantiating the exec to provide upload with work to do.\n hipGraphExecDestroy(graphExec);\n hipGraphInstantiateWithFlags(&graphExec, graph, 0);\n long long maxTimeoutNanoSeconds = 4000 + 1000 * length * width;\n waitWithTimeout<<<1, 1, 0, stream[0]>>>(\n maxTimeoutNanoSeconds, &hostData->timeoutDetected2, &hostData->timeElapsed2, &hostData->latch2);\n maxTimeoutNanoSeconds = 2000 + 500 * length * width;\n waitWithTimeout<<<1, 1, 0, stream[1]>>>(\n maxTimeoutNanoSeconds, &hostData->timeoutDetected, &hostData->timeElapsed, &hostData->latch);\n\n RANGE(\"uploading a graph off of the critical path\");\n preUploadAnnotation<<<1, 1, 0, stream[1]>>>();\n hipEventRecord(timingEvent[0], stream[0]);\n auto start = getCpuTime();\n hipGraphUpload(graphExec, stream[1]);\n auto apiReturn = getCpuTime();\n hipEventRecord(event[0], stream[1]);\n hipEventRecord(timingEvent[1], stream[0]);\n postUploadAnnotation<<<1, 1, 0, stream[1]>>>();\n\n hostData->latch = 1; // release the blocking kernel for the upload\n hipStreamWaitEvent(stream[0], event[0]);\n hipGraphLaunch(graphExec, stream[0]);\n hipEventSynchronize(event[0]); // upload done, similuate critical path being ready for the graph to run by the\n // release of the second latch\n\n hostData->latch2 = 1; // release the work\n hipStreamSynchronize(stream[0]);\n\n metricName.push_back(\"upload_api_time\");\n metricValue.push_back(getMicroSecondDuration(start, apiReturn));\n metricName.push_back(\"updoad_device_time\");\n metricValue.push_back(getAsyncMicroSecondDuration(timingEvent[0], timingEvent[1]));\n metricName.push_back(\"blockingKernelTimeoutDetected\");\n metricValue.push_back(hostData->timeoutDetected);\n\n hostData->latch = 0;\n hostData->latch2 = 0;\n hostData->timeoutDetected = 0;\n hostData->timeoutDetected2 = 0;\n }\n hipGraphExecDestroy(graphExec);\n hipGraphDestroy(graph);\n RANGE_POP();\n}\n\nvoid usage()\n{\n printf(\"programName [outputFmt] [numTrials] [length] [width] [pattern] [stride] [maxLength] \\n\");\n printf(\"\\toutputFmt - program output, default=3 (see below)\\n\");\n printf(\"\\tnumTrials (per length)\\n\");\n printf(\"\\tstarting length of the topology\\n\");\n printf(\"\\twidth - width of the graph topology\\n\");\n printf(\"\\tpattern - Structure of graph, default=0 (see below)\\n\");\n printf(\"\\tstride - how to grow the length between each set of trials \\n\");\n printf(\"\\tmaxLength - maximum lenght to try \\n\");\n printf(\"\\n\");\n printf(\"outputFmt can be:\\n\");\n printf(\"\\t0: this help message\\n\");\n printf(\"\\t1: csv data headers\\n\");\n printf(\"\\t2: per trial csv data\\n\");\n printf(\"\\t3: csv data & headers\\n\");\n printf(\"\\t4: csv data is printed and trials are averaged for each length\\n\");\n printf(\"\\t5: csv data is printed and trials are averaged for each length and headers are printed\\n\");\n printf(\"\\n\");\n printf(\"Pattern can be:\\n\");\n printf(\"\\t0: No interconnect between branches\\n\");\n printf(\"\\t1: Adds an extra root node before the initial fork\\n\");\n}\n\nint main(int argc, char **argv)\n{\n if (argc < 1) {\n usage();\n return 0;\n }\n\n int numTrials = 1, length = 20, width = 1, outputFmt = 3, pattern = 0, stride = 1;\n if (argc > 1)\n outputFmt = atoi(argv[1]);\n if (argc > 2)\n numTrials = atoi(argv[2]);\n if (argc > 3)\n length = atoi(argv[3]);\n if (argc > 4)\n width = atoi(argv[4]);\n if (argc > 5)\n pattern = atoi(argv[5]);\n if (argc > 6)\n stride = atoi(argv[6]);\n int maxLength = length;\n if (argc > 7)\n maxLength = atoi(argv[7]);\n if (maxLength < length) {\n maxLength = length;\n }\n\n if ((outputFmt & 4) && (outputFmt & 2)) {\n printf(\"printing average and all samples doesn't make sense\\n\");\n }\n\n if (length == 0 || width == 0 || outputFmt == 0 || outputFmt > 5 || pattern > 1) {\n usage();\n return 0;\n }\n\n bool singleEntry = (pattern == 1);\n\n hipGraph_t graph;\n\n hipFree(0);\n hipHostMalloc(&hostData, sizeof(*hostData));\n int numStreams = width;\n if (numStreams == 1)\n numStreams = 2; // demo needs two streams even if capture only needs 1.\n stream.resize(numStreams);\n for (int i = 0; i < numStreams; i++) {\n hipStreamCreate(&stream[i]);\n }\n\n hipEventCreate(&event[0], hipEventDisableTiming);\n hipEventCreate(&timingEvent[0], 0);\n hipEventCreate(&timingEvent[1], 0);\n\n {\n RANGE(\"warmup\");\n for (int i = 0; i < width; i++) {\n empty<<<1, 1, 0, stream[i]>>>();\n }\n hipStreamSynchronize(stream[0]);\n\n auto start = getCpuTime();\n graph = createParallelChain(length, width, singleEntry);\n auto end = getCpuTime();\n metricValue.push_back(getMicroSecondDuration(start, end));\n metricName.push_back(\"capture\");\n runDemo(graph, length, width);\n }\n\n if (outputFmt & 1) {\n printf(\"length, width, pattern, \");\n for (int i = 0; i < metricName.size(); i++) {\n printf(\"%s, \", metricName[i]);\n }\n printf(\"\\r\\n\");\n }\n\n if (!(outputFmt & 6)) {\n printf(\"skipping trials since no output is expected\\n\");\n return 1;\n }\n\n std::vector metricTotal;\n metricTotal.resize(metricValue.size());\n\n while (length <= maxLength) {\n for (int i = 0; i < numTrials; i++) {\n metricName.clear();\n metricValue.clear();\n auto start = getCpuTime();\n graph = createParallelChain(length, width, singleEntry);\n auto end = getCpuTime();\n metricValue.push_back(getMicroSecondDuration(start, end));\n\n runDemo(graph, length, width);\n\n if (outputFmt & 2) {\n printf(\"%d, %d, %d, \", length, width, pattern);\n for (int i = 0; i < metricValue.size(); i++) {\n printf(\"%0.3f, \", metricValue[i]);\n }\n printf(\"\\r\\n\");\n }\n if (outputFmt & 4) {\n for (int i = 0; i < metricTotal.size(); i++) {\n metricTotal[i] += metricValue[i];\n }\n }\n }\n\n if (outputFmt & 4) {\n printf(\"%d, %d, %d, \", length, width, pattern);\n for (int i = 0; i < metricTotal.size(); i++) {\n printf(\"%0.3f, \", metricTotal[i] / numTrials);\n metricTotal[i] = 0;\n }\n printf(\"\\r\\n\");\n }\n\n length += stride;\n }\n\n hipHostFree(hostData);\n\n printf(\"\\n\");\n printf(\"Test passed\\n\");\n return 0;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/6_Performance/cudaGraphsPerfScaling/cudaGraphPerfScaling.cu"} {"input": "/* Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This is a simple test showing performance and usability\n * improvements with large kernel parameters introduced in CUDA 12.1\n */\n#include \n#include \n#include \n\n// Utility includes\n#include \n\nusing namespace std;\nusing namespace std::chrono;\n\n#define TEST_ITERATIONS (1000)\n#define TOTAL_PARAMS (8000) // ints\n#define KERNEL_PARAM_LIMIT (1024) // ints\n#define CONST_COPIED_PARAMS (TOTAL_PARAMS - KERNEL_PARAM_LIMIT)\n\n__constant__ int excess_params[CONST_COPIED_PARAMS];\n\ntypedef struct\n{\n int param[KERNEL_PARAM_LIMIT];\n} param_t;\n\ntypedef struct\n{\n int param[TOTAL_PARAMS];\n} param_large_t;\n\n// Kernel with 4KB kernel parameter limit\n__global__ void kernelDefault(__grid_constant__ const param_t p, int *result)\n{\n int tmp = 0;\n\n // accumulate kernel parameters\n for (int i = 0; i < KERNEL_PARAM_LIMIT; ++i) {\n tmp += p.param[i];\n }\n\n // accumulate excess values passed via const memory\n for (int i = 0; i < CONST_COPIED_PARAMS; ++i) {\n tmp += excess_params[i];\n }\n\n *result = tmp;\n}\n\n// Kernel with 32,764 byte kernel parameter limit\n__global__ void kernelLargeParam(__grid_constant__ const param_large_t p, int *result)\n{\n int tmp = 0;\n\n // accumulate kernel parameters\n for (int i = 0; i < TOTAL_PARAMS; ++i) {\n tmp += p.param[i];\n }\n\n *result = tmp;\n}\n\nstatic void report_time(std::chrono::time_point start,\n std::chrono::time_point end,\n int iters)\n{\n auto usecs = duration_cast>(end - start);\n cout << usecs.count() / iters << endl;\n}\n\nint main()\n{\n int rc;\n cudaFree(0);\n\n param_t p;\n param_large_t p_large;\n\n // pageable host memory that holds excess constants passed via constant memory\n int *copied_params = (int *)malloc(CONST_COPIED_PARAMS * sizeof(int));\n assert(copied_params);\n\n // storage for computed result\n int *d_result;\n int h_result;\n checkCudaErrors(cudaMalloc(&d_result, sizeof(int)));\n\n int expected_result = 0;\n\n // fill in data for validation\n for (int i = 0; i < KERNEL_PARAM_LIMIT; ++i) {\n p.param[i] = (i & 0xFF);\n }\n for (int i = KERNEL_PARAM_LIMIT; i < TOTAL_PARAMS; ++i) {\n copied_params[i - KERNEL_PARAM_LIMIT] = (i & 0xFF);\n }\n for (int i = 0; i < TOTAL_PARAMS; ++i) {\n p_large.param[i] = (i & 0xFF);\n expected_result += (i & 0xFF);\n }\n\n // warmup, verify correctness\n checkCudaErrors(\n cudaMemcpyToSymbol(excess_params, copied_params, CONST_COPIED_PARAMS * sizeof(int), 0, cudaMemcpyHostToDevice));\n kernelDefault<<<1, 1>>>(p, d_result);\n checkCudaErrors(cudaMemcpy(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost));\n checkCudaErrors(cudaDeviceSynchronize());\n if (h_result != expected_result) {\n std::cout << \"Test failed\" << std::endl;\n rc = -1;\n goto Exit;\n }\n\n kernelLargeParam<<<1, 1>>>(p_large, d_result);\n checkCudaErrors(cudaMemcpy(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost));\n checkCudaErrors(cudaDeviceSynchronize());\n if (h_result != expected_result) {\n std::cout << \"Test failed\" << std::endl;\n rc = -1;\n goto Exit;\n }\n\n // benchmark default kernel parameter limit\n {\n auto start = steady_clock::now();\n for (int i = 0; i < TEST_ITERATIONS; ++i) {\n checkCudaErrors(cudaMemcpyToSymbol(\n excess_params, copied_params, CONST_COPIED_PARAMS * sizeof(int), 0, cudaMemcpyHostToDevice));\n kernelDefault<<<1, 1>>>(p, d_result);\n }\n checkCudaErrors(cudaDeviceSynchronize());\n auto end = steady_clock::now();\n std::cout << \"Kernel 4KB parameter limit - time (us):\";\n report_time(start, end, TEST_ITERATIONS);\n\n // benchmark large kernel parameter limit\n start = steady_clock::now();\n for (int i = 0; i < TEST_ITERATIONS; ++i) {\n kernelLargeParam<<<1, 1>>>(p_large, d_result);\n }\n checkCudaErrors(cudaDeviceSynchronize());\n end = steady_clock::now();\n std::cout << \"Kernel 32,764 byte parameter limit - time (us):\";\n report_time(start, end, TEST_ITERATIONS);\n }\n std::cout << \"Test passed!\" << std::endl;\n rc = 0;\nExit:\n // cleanup\n cudaFree(d_result);\n free(copied_params);\n return rc;\n}\n", "output": "/* Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This is a simple test showing performance and usability\n * improvements with large kernel parameters introduced in CUDA 12.1\n */\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n// Utility includes\n#include \n\nusing namespace std;\nusing namespace std::chrono;\n\n#define TEST_ITERATIONS (1000)\n#define TOTAL_PARAMS (8000) // ints\n#define KERNEL_PARAM_LIMIT (1024) // ints\n#define CONST_COPIED_PARAMS (TOTAL_PARAMS - KERNEL_PARAM_LIMIT)\n\n__constant__ int excess_params[CONST_COPIED_PARAMS];\n\ntypedef struct\n{\n int param[KERNEL_PARAM_LIMIT];\n} param_t;\n\ntypedef struct\n{\n int param[TOTAL_PARAMS];\n} param_large_t;\n\n// Kernel with 4KB kernel parameter limit\n__global__ void kernelDefault(__grid_constant__ const param_t p, int *result)\n{\n int tmp = 0;\n\n // accumulate kernel parameters\n for (int i = 0; i < KERNEL_PARAM_LIMIT; ++i) {\n tmp += p.param[i];\n }\n\n // accumulate excess values passed via const memory\n for (int i = 0; i < CONST_COPIED_PARAMS; ++i) {\n tmp += excess_params[i];\n }\n\n *result = tmp;\n}\n\n// Kernel with 32,764 byte kernel parameter limit\n__global__ void kernelLargeParam(__grid_constant__ const param_large_t p, int *result)\n{\n int tmp = 0;\n\n // accumulate kernel parameters\n for (int i = 0; i < TOTAL_PARAMS; ++i) {\n tmp += p.param[i];\n }\n\n *result = tmp;\n}\n\nstatic void report_time(std::chrono::time_point start,\n std::chrono::time_point end,\n int iters)\n{\n auto usecs = duration_cast>(end - start);\n cout << usecs.count() / iters << endl;\n}\n\nint main()\n{\n int rc;\n hipFree(0);\n\n param_t p;\n param_large_t p_large;\n\n // pageable host memory that holds excess constants passed via constant memory\n int *copied_params = (int *)malloc(CONST_COPIED_PARAMS * sizeof(int));\n assert(copied_params);\n\n // storage for computed result\n int *d_result;\n int h_result;\n checkCudaErrors(hipMalloc(&d_result, sizeof(int)));\n\n int expected_result = 0;\n\n // fill in data for validation\n for (int i = 0; i < KERNEL_PARAM_LIMIT; ++i) {\n p.param[i] = (i & 0xFF);\n }\n for (int i = KERNEL_PARAM_LIMIT; i < TOTAL_PARAMS; ++i) {\n copied_params[i - KERNEL_PARAM_LIMIT] = (i & 0xFF);\n }\n for (int i = 0; i < TOTAL_PARAMS; ++i) {\n p_large.param[i] = (i & 0xFF);\n expected_result += (i & 0xFF);\n }\n\n // warmup, verify correctness\n checkCudaErrors(\n hipMemcpyToSymbol(HIP_SYMBOL(excess_params), copied_params, CONST_COPIED_PARAMS * sizeof(int), 0, hipMemcpyHostToDevice));\n kernelDefault<<<1, 1>>>(p, d_result);\n checkCudaErrors(hipMemcpy(&h_result, d_result, sizeof(int), hipMemcpyDeviceToHost));\n checkCudaErrors(hipDeviceSynchronize());\n if (h_result != expected_result) {\n std::cout << \"Test failed\" << std::endl;\n rc = -1;\n goto Exit;\n }\n\n kernelLargeParam<<<1, 1>>>(p_large, d_result);\n checkCudaErrors(hipMemcpy(&h_result, d_result, sizeof(int), hipMemcpyDeviceToHost));\n checkCudaErrors(hipDeviceSynchronize());\n if (h_result != expected_result) {\n std::cout << \"Test failed\" << std::endl;\n rc = -1;\n goto Exit;\n }\n\n // benchmark default kernel parameter limit\n {\n auto start = steady_clock::now();\n for (int i = 0; i < TEST_ITERATIONS; ++i) {\n checkCudaErrors(hipMemcpyToSymbol(HIP_SYMBOL(excess_params), copied_params, CONST_COPIED_PARAMS * sizeof(int), 0, hipMemcpyHostToDevice));\n kernelDefault<<<1, 1>>>(p, d_result);\n }\n checkCudaErrors(hipDeviceSynchronize());\n auto end = steady_clock::now();\n std::cout << \"Kernel 4KB parameter limit - time (us):\";\n report_time(start, end, TEST_ITERATIONS);\n\n // benchmark large kernel parameter limit\n start = steady_clock::now();\n for (int i = 0; i < TEST_ITERATIONS; ++i) {\n kernelLargeParam<<<1, 1>>>(p_large, d_result);\n }\n checkCudaErrors(hipDeviceSynchronize());\n end = steady_clock::now();\n std::cout << \"Kernel 32,764 byte parameter limit - time (us):\";\n report_time(start, end, TEST_ITERATIONS);\n }\n std::cout << \"Test passed!\" << std::endl;\n rc = 0;\nExit:\n // cleanup\n hipFree(d_result);\n free(copied_params);\n return rc;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/6_Performance/LargeKernelParameter/LargeKernelParameter.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"../inc/pricingengine.h\"\n\nnamespace cg = cooperative_groups;\n#include \n\n#include \"../inc/asianoption.h\"\n#include \"../inc/cudasharedmem.h\"\n\nusing std::string;\nusing std::vector;\n\n// RNG init kernel\n__global__ void initRNG(curandState *const rngStates, const unsigned int seed)\n{\n // Determine thread ID\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Initialise the RNG\n curand_init(seed, tid, 0, &rngStates[tid]);\n}\n\n__device__ inline float getPathStep(float &drift, float &diffusion, curandState &state)\n{\n return expf(drift + diffusion * curand_normal(&state));\n}\n__device__ inline double getPathStep(double &drift, double &diffusion, curandState &state)\n{\n return exp(drift + diffusion * curand_normal_double(&state));\n}\n\n// Path generation kernel\ntemplate \n__global__ void generatePaths(Real *const paths,\n curandState *const rngStates,\n const AsianOption *const option,\n const unsigned int numSims,\n const unsigned int numTimesteps)\n{\n // Determine thread ID\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int step = gridDim.x * blockDim.x;\n\n // Compute parameters\n Real drift = (option->r - static_cast(0.5) * option->sigma * option->sigma) * option->dt;\n Real diffusion = option->sigma * sqrt(option->dt);\n\n // Initialise the RNG\n curandState localState = rngStates[tid];\n\n for (unsigned int i = tid; i < numSims; i += step) {\n // Shift the output pointer\n Real *output = paths + i;\n\n // Simulate the path\n Real s = static_cast(1);\n\n for (unsigned int t = 0; t < numTimesteps; t++, output += numSims) {\n s *= getPathStep(drift, diffusion, localState);\n *output = s;\n }\n }\n}\n\ntemplate __device__ Real reduce_sum(Real in, cg::thread_block cta)\n{\n SharedMemory sdata;\n\n // Perform first level of reduction:\n // - Write to shared memory\n unsigned int ltid = threadIdx.x;\n\n sdata[ltid] = in;\n cg::sync(cta);\n\n // Do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (ltid < s) {\n sdata[ltid] += sdata[ltid + s];\n }\n\n cg::sync(cta);\n }\n\n return sdata[0];\n}\n\n// Valuation kernel\ntemplate \n__global__ void computeValue(Real *const values,\n const Real *const paths,\n const AsianOption *const option,\n const unsigned int numSims,\n const unsigned int numTimesteps)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Determine thread ID\n unsigned int bid = blockIdx.x;\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int step = gridDim.x * blockDim.x;\n\n Real sumPayoffs = static_cast(0);\n\n for (unsigned int i = tid; i < numSims; i += step) {\n // Shift the input pointer\n const Real *path = paths + i;\n // Compute the arithmetic average\n Real avg = static_cast(0);\n\n for (unsigned int t = 0; t < numTimesteps; t++, path += numSims) {\n avg += *path;\n }\n\n avg = avg * option->spot / numTimesteps;\n // Compute the payoff\n Real payoff = avg - option->strike;\n\n if (option->type == AsianOption::Put) {\n payoff = -payoff;\n }\n\n payoff = max(static_cast(0), payoff);\n // Accumulate payoff locally\n sumPayoffs += payoff;\n }\n\n // Reduce within the block\n sumPayoffs = reduce_sum(sumPayoffs, cta);\n\n // Store the result\n if (threadIdx.x == 0) {\n values[bid] = sumPayoffs;\n }\n}\n\ntemplate \nPricingEngine::PricingEngine(unsigned int numSims,\n unsigned int device,\n unsigned int threadBlockSize,\n unsigned int seed)\n : m_numSims(numSims)\n , m_device(device)\n , m_threadBlockSize(threadBlockSize)\n , m_seed(seed)\n{\n}\n\ntemplate void PricingEngine::operator()(AsianOption &option)\n{\n cudaError_t cudaResult = cudaSuccess;\n struct cudaDeviceProp deviceProperties;\n struct cudaFuncAttributes funcAttributes;\n\n // Get device properties\n cudaResult = cudaGetDeviceProperties(&deviceProperties, m_device);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not get device properties: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Check precision is valid\n unsigned int deviceVersion = deviceProperties.major * 10 + deviceProperties.minor;\n\n if (typeid(Real) == typeid(double) && deviceVersion < 13) {\n throw std::runtime_error(\"Device does not have double precision support\");\n }\n\n // Attach to GPU\n cudaResult = cudaSetDevice(m_device);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not set CUDA device: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Determine how to divide the work between cores\n dim3 block;\n dim3 grid;\n block.x = m_threadBlockSize;\n grid.x = (m_numSims + m_threadBlockSize - 1) / m_threadBlockSize;\n\n // Aim to launch around ten or more times as many blocks as there\n // are multiprocessors on the target device.\n unsigned int blocksPerSM = 10;\n unsigned int numSMs = deviceProperties.multiProcessorCount;\n\n while (grid.x > 2 * blocksPerSM * numSMs) {\n grid.x >>= 1;\n }\n\n // Get initRNG function properties and check the maximum block size\n cudaResult = cudaFuncGetAttributes(&funcAttributes, initRNG);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for initRNG kernel\");\n }\n\n // Get generatePaths function properties and check the maximum block size\n cudaResult = cudaFuncGetAttributes(&funcAttributes, generatePaths);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for generatePaths kernel\");\n }\n\n // Get computeValue function properties and check the maximum block size\n cudaResult = cudaFuncGetAttributes(&funcAttributes, computeValue);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for computeValue kernel\");\n }\n\n // Setup problem on GPU\n AsianOption *d_option = 0;\n cudaResult = cudaMalloc((void **)&d_option, sizeof(AsianOption));\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not allocate memory on device for option data: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n cudaResult = cudaMemcpy(d_option, &option, sizeof(AsianOption), cudaMemcpyHostToDevice);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not copy data to device: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Allocate memory for paths\n Real *d_paths = 0;\n int numTimesteps = static_cast(option.tenor / option.dt);\n cudaResult = cudaMalloc((void **)&d_paths, m_numSims * numTimesteps * sizeof(Real));\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not allocate memory on device for paths: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Allocate memory for RNG states\n curandState *d_rngStates = 0;\n cudaResult = cudaMalloc((void **)&d_rngStates, grid.x * block.x * sizeof(curandState));\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not allocate memory on device for RNG state: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Allocate memory for result\n Real *d_values = 0;\n cudaResult = cudaMalloc((void **)&d_values, grid.x * sizeof(Real));\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not allocate memory on device for partial results: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Initialise RNG\n initRNG<<>>(d_rngStates, m_seed);\n\n // Generate paths\n generatePaths<<>>(d_paths, d_rngStates, d_option, m_numSims, numTimesteps);\n\n // Compute value\n computeValue<<>>(d_values, d_paths, d_option, m_numSims, numTimesteps);\n\n // Copy partial results back\n vector values(grid.x);\n cudaResult = cudaMemcpy(&values[0], d_values, grid.x * sizeof(Real), cudaMemcpyDeviceToHost);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not copy partial results to host: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Complete sum-reduction on host\n option.value = std::accumulate(values.begin(), values.end(), static_cast(0));\n\n // Compute the mean\n option.value /= m_numSims;\n\n // Discount to present value\n option.value *= exp(-option.r * option.tenor);\n\n // Cleanup\n if (d_option) {\n cudaFree(d_option);\n d_option = 0;\n }\n\n if (d_paths) {\n cudaFree(d_paths);\n d_paths = 0;\n }\n\n if (d_rngStates) {\n cudaFree(d_rngStates);\n d_rngStates = 0;\n }\n\n if (d_values) {\n cudaFree(d_values);\n d_values = 0;\n }\n}\n\n// Explicit template instantiation\ntemplate class PricingEngine;\ntemplate class PricingEngine;\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"../inc/pricingengine.h\"\n\nnamespace cg = cooperative_groups;\n#include \n\n#include \"../inc/asianoption.h\"\n#include \"../inc/cudasharedmem.h\"\n\nusing std::string;\nusing std::vector;\n\n// RNG init kernel\n__global__ void initRNG(hiprandState *const rngStates, const unsigned int seed)\n{\n // Determine thread ID\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Initialise the RNG\n hiprand_init(seed, tid, 0, &rngStates[tid]);\n}\n\n__device__ inline float getPathStep(float &drift, float &diffusion, hiprandState &state)\n{\n return expf(drift + diffusion * hiprand_normal(&state));\n}\n__device__ inline double getPathStep(double &drift, double &diffusion, hiprandState &state)\n{\n return exp(drift + diffusion * hiprand_normal_double(&state));\n}\n\n// Path generation kernel\ntemplate \n__global__ void generatePaths(Real *const paths,\n hiprandState *const rngStates,\n const AsianOption *const option,\n const unsigned int numSims,\n const unsigned int numTimesteps)\n{\n // Determine thread ID\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int step = gridDim.x * blockDim.x;\n\n // Compute parameters\n Real drift = (option->r - static_cast(0.5) * option->sigma * option->sigma) * option->dt;\n Real diffusion = option->sigma * sqrt(option->dt);\n\n // Initialise the RNG\n hiprandState localState = rngStates[tid];\n\n for (unsigned int i = tid; i < numSims; i += step) {\n // Shift the output pointer\n Real *output = paths + i;\n\n // Simulate the path\n Real s = static_cast(1);\n\n for (unsigned int t = 0; t < numTimesteps; t++, output += numSims) {\n s *= getPathStep(drift, diffusion, localState);\n *output = s;\n }\n }\n}\n\ntemplate __device__ Real reduce_sum(Real in, cg::thread_block cta)\n{\n SharedMemory sdata;\n\n // Perform first level of reduction:\n // - Write to shared memory\n unsigned int ltid = threadIdx.x;\n\n sdata[ltid] = in;\n cg::sync(cta);\n\n // Do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (ltid < s) {\n sdata[ltid] += sdata[ltid + s];\n }\n\n cg::sync(cta);\n }\n\n return sdata[0];\n}\n\n// Valuation kernel\ntemplate \n__global__ void computeValue(Real *const values,\n const Real *const paths,\n const AsianOption *const option,\n const unsigned int numSims,\n const unsigned int numTimesteps)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Determine thread ID\n unsigned int bid = blockIdx.x;\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int step = gridDim.x * blockDim.x;\n\n Real sumPayoffs = static_cast(0);\n\n for (unsigned int i = tid; i < numSims; i += step) {\n // Shift the input pointer\n const Real *path = paths + i;\n // Compute the arithmetic average\n Real avg = static_cast(0);\n\n for (unsigned int t = 0; t < numTimesteps; t++, path += numSims) {\n avg += *path;\n }\n\n avg = avg * option->spot / numTimesteps;\n // Compute the payoff\n Real payoff = avg - option->strike;\n\n if (option->type == AsianOption::Put) {\n payoff = -payoff;\n }\n\n payoff = max(static_cast(0), payoff);\n // Accumulate payoff locally\n sumPayoffs += payoff;\n }\n\n // Reduce within the block\n sumPayoffs = reduce_sum(sumPayoffs, cta);\n\n // Store the result\n if (threadIdx.x == 0) {\n values[bid] = sumPayoffs;\n }\n}\n\ntemplate \nPricingEngine::PricingEngine(unsigned int numSims,\n unsigned int device,\n unsigned int threadBlockSize,\n unsigned int seed)\n : m_numSims(numSims)\n , m_device(device)\n , m_threadBlockSize(threadBlockSize)\n , m_seed(seed)\n{\n}\n\ntemplate void PricingEngine::operator()(AsianOption &option)\n{\n hipError_t cudaResult = hipSuccess;\n struct hipDeviceProp_t deviceProperties;\n struct hipFuncAttributes funcAttributes;\n\n // Get device properties\n cudaResult = hipGetDeviceProperties(&deviceProperties, m_device);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not get device properties: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Check precision is valid\n unsigned int deviceVersion = deviceProperties.major * 10 + deviceProperties.minor;\n\n if (typeid(Real) == typeid(double) && deviceVersion < 13) {\n throw std::runtime_error(\"Device does not have double precision support\");\n }\n\n // Attach to GPU\n cudaResult = hipSetDevice(m_device);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not set CUDA device: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Determine how to divide the work between cores\n dim3 block;\n dim3 grid;\n block.x = m_threadBlockSize;\n grid.x = (m_numSims + m_threadBlockSize - 1) / m_threadBlockSize;\n\n // Aim to launch around ten or more times as many blocks as there\n // are multiprocessors on the target device.\n unsigned int blocksPerSM = 10;\n unsigned int numSMs = deviceProperties.multiProcessorCount;\n\n while (grid.x > 2 * blocksPerSM * numSMs) {\n grid.x >>= 1;\n }\n\n // Get initRNG function properties and check the maximum block size\n cudaResult = hipFuncGetAttributes(&funcAttributes,reinterpret_cast(initRNG));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for initRNG kernel\");\n }\n\n // Get generatePaths function properties and check the maximum block size\n cudaResult = hipFuncGetAttributes(&funcAttributes,reinterpret_cast(generatePaths));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for generatePaths kernel\");\n }\n\n // Get computeValue function properties and check the maximum block size\n cudaResult = hipFuncGetAttributes(&funcAttributes,reinterpret_cast(computeValue));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for computeValue kernel\");\n }\n\n // Setup problem on GPU\n AsianOption *d_option = 0;\n cudaResult = hipMalloc((void **)&d_option, sizeof(AsianOption));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not allocate memory on device for option data: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n cudaResult = hipMemcpy(d_option, &option, sizeof(AsianOption), hipMemcpyHostToDevice);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not copy data to device: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Allocate memory for paths\n Real *d_paths = 0;\n int numTimesteps = static_cast(option.tenor / option.dt);\n cudaResult = hipMalloc((void **)&d_paths, m_numSims * numTimesteps * sizeof(Real));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not allocate memory on device for paths: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Allocate memory for RNG states\n hiprandState *d_rngStates = 0;\n cudaResult = hipMalloc((void **)&d_rngStates, grid.x * block.x * sizeof(hiprandState));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not allocate memory on device for RNG state: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Allocate memory for result\n Real *d_values = 0;\n cudaResult = hipMalloc((void **)&d_values, grid.x * sizeof(Real));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not allocate memory on device for partial results: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Initialise RNG\n initRNG<<>>(d_rngStates, m_seed);\n\n // Generate paths\n generatePaths<<>>(d_paths, d_rngStates, d_option, m_numSims, numTimesteps);\n\n // Compute value\n computeValue<<>>(d_values, d_paths, d_option, m_numSims, numTimesteps);\n\n // Copy partial results back\n vector values(grid.x);\n cudaResult = hipMemcpy(&values[0], d_values, grid.x * sizeof(Real), hipMemcpyDeviceToHost);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not copy partial results to host: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Complete sum-reduction on host\n option.value = std::accumulate(values.begin(), values.end(), static_cast(0));\n\n // Compute the mean\n option.value /= m_numSims;\n\n // Discount to present value\n option.value *= exp(-option.r * option.tenor);\n\n // Cleanup\n if (d_option) {\n hipFree(d_option);\n d_option = 0;\n }\n\n if (d_paths) {\n hipFree(d_paths);\n d_paths = 0;\n }\n\n if (d_rngStates) {\n hipFree(d_rngStates);\n d_rngStates = 0;\n }\n\n if (d_values) {\n hipFree(d_values);\n d_values = 0;\n }\n}\n\n// Explicit template instantiation\ntemplate class PricingEngine;\ntemplate class PricingEngine;\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/MC_SingleAsianOptionP/src/pricingengine.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates Inter Process Communication\n * using one process per GPU for computation.\n */\n\n#include \n#include \n#include \n#include \n#define CUDA_DRIVER_API 1\n#include \"helper_cuda.h\"\n#include \"helper_cuda_drvapi.h\"\n#include \"helper_multiprocess.h\"\n\nstatic const char shmName[] = \"streamOrderedAllocationIPCshm\";\nstatic const char ipcName[] = \"streamOrderedAllocationIPC_pipe\";\n// For direct NVLINK and PCI-E peers, at max 8 simultaneous peers are allowed\n// For NVSWITCH connected peers like DGX-2, simultaneous peers are not limited\n// in the same way.\n#define MAX_DEVICES (32)\n#define DATA_SIZE (64ULL << 20ULL) // 64MB\n\n#if defined(__linux__)\n#define cpu_atomic_add32(a, x) __sync_add_and_fetch(a, x)\n#elif defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)\n#define cpu_atomic_add32(a, x) InterlockedAdd((volatile LONG *)a, x)\n#else\n#error Unsupported system\n#endif\n\ntypedef struct shmStruct_st\n{\n size_t nprocesses;\n int barrier;\n int sense;\n cudaMemAllocationHandleType handleType;\n int devices[MAX_DEVICES];\n cudaMemPoolPtrExportData exportPtrData[MAX_DEVICES];\n} shmStruct;\n\n__global__ void simpleKernel(char *ptr, int sz, char val)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n for (; idx < sz; idx += (gridDim.x * blockDim.x)) {\n ptr[idx] = val;\n }\n}\n\nstatic void barrierWait(volatile int *barrier, volatile int *sense, unsigned int n)\n{\n int count;\n\n // Check-in\n count = cpu_atomic_add32(barrier, 1);\n if (count == n) // Last one in\n *sense = 1;\n while (!*sense)\n ;\n\n // Check-out\n count = cpu_atomic_add32(barrier, -1);\n if (count == 0) // Last one out\n *sense = 0;\n while (*sense)\n ;\n}\n\nstatic void childProcess(int id)\n{\n volatile shmStruct *shm = NULL;\n cudaStream_t stream;\n sharedMemoryInfo info;\n size_t procCount, i;\n int blocks = 0;\n int threads = 128;\n cudaDeviceProp prop;\n std::vector ptrs;\n pid_t pid;\n char pidString[20] = {0};\n char lshmName[40] = {0};\n\n std::vector verification_buffer(DATA_SIZE);\n\n pid = getppid();\n snprintf(pidString, sizeof(pidString), \"%d\", pid);\n strcat(lshmName, shmName);\n strcat(lshmName, pidString);\n\n printf(\"CP: lshmName = %s\\n\", lshmName);\n\n ipcHandle *ipcChildHandle = NULL;\n checkIpcErrors(ipcOpenSocket(ipcChildHandle));\n\n if (sharedMemoryOpen(lshmName, sizeof(shmStruct), &info) != 0) {\n printf(\"Failed to create shared memory slab\\n\");\n exit(EXIT_FAILURE);\n }\n shm = (volatile shmStruct *)info.addr;\n procCount = shm->nprocesses;\n\n barrierWait(&shm->barrier, &shm->sense, (unsigned int)(procCount + 1));\n\n // Receive all allocation handles shared by Parent.\n std::vector shHandle(shm->nprocesses);\n checkIpcErrors(ipcRecvShareableHandles(ipcChildHandle, shHandle));\n\n checkCudaErrors(cudaSetDevice(shm->devices[id]));\n checkCudaErrors(cudaGetDeviceProperties(&prop, shm->devices[id]));\n checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));\n checkCudaErrors(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, simpleKernel, threads, 0));\n blocks *= prop.multiProcessorCount;\n\n std::vector pools(shm->nprocesses);\n\n cudaMemAllocationHandleType handleType = shm->handleType;\n\n // Import mem pools from all the devices created in the master\n // process using shareable handles received via socket\n // and import the pointer to the allocated buffer using\n // exportData filled in shared memory by the master process.\n for (i = 0; i < procCount; i++) {\n checkCudaErrors(cudaMemPoolImportFromShareableHandle(&pools[i], (void *)shHandle[i], handleType, 0));\n\n cudaMemAccessFlags accessFlags;\n cudaMemLocation location;\n location.type = cudaMemLocationTypeDevice;\n location.id = shm->devices[id];\n checkCudaErrors(cudaMemPoolGetAccess(&accessFlags, pools[i], &location));\n if (accessFlags != cudaMemAccessFlagsProtReadWrite) {\n cudaMemAccessDesc desc;\n memset(&desc, 0, sizeof(cudaMemAccessDesc));\n desc.location.type = cudaMemLocationTypeDevice;\n desc.location.id = shm->devices[id];\n desc.flags = cudaMemAccessFlagsProtReadWrite;\n checkCudaErrors(cudaMemPoolSetAccess(pools[i], &desc, 1));\n }\n\n // Import the allocation from each memory pool by iterating over exportData\n // until import is success.\n for (int j = 0; j < procCount; j++) {\n void *ptr = NULL;\n // Import the allocation using the opaque export data retrieved through\n // the shared memory\".\n cudaError_t ret =\n cudaMemPoolImportPointer(&ptr, pools[i], (cudaMemPoolPtrExportData *)&shm->exportPtrData[j]);\n\n if (ret == cudaSuccess) {\n // Pointer import is successful hence add it to the ptrs bag.\n ptrs.push_back(ptr);\n break;\n }\n else {\n // Reset failure error received from cudaMemPoolImportPointer\n // for further try.\n cudaGetLastError();\n }\n }\n // Since we have imported allocations shared by the parent with us, we can\n // close this ShareableHandle.\n checkIpcErrors(ipcCloseShareableHandle(shHandle[i]));\n }\n\n // Since we have imported allocations shared by the parent with us, we can\n // close the socket.\n checkIpcErrors(ipcCloseSocket(ipcChildHandle));\n\n // At each iteration of the loop, each sibling process will push work on\n // their respective devices accessing the next peer mapped buffer allocated\n // by the master process (these can come from other sibling processes as\n // well). To coordinate each process' access, we force the stream to wait for\n // the work already accessing this buffer.\n for (i = 0; i < procCount; i++) {\n size_t bufferId = (i + id) % procCount;\n\n // Push a simple kernel on it\n simpleKernel<<>>((char *)ptrs[bufferId], DATA_SIZE, id);\n checkCudaErrors(cudaGetLastError());\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n // Wait for all my sibling processes to push this stage of their work\n // before proceeding to the next. This prevents siblings from racing\n // ahead and clobbering the recorded event or waiting on the wrong\n // recorded event.\n barrierWait(&shm->barrier, &shm->sense, (unsigned int)procCount);\n if (id == 0) {\n printf(\"Step %lld done\\n\", (unsigned long long)i);\n }\n }\n\n // Now wait for my buffer to be ready so I can copy it locally and verify it\n checkCudaErrors(cudaMemcpyAsync(&verification_buffer[0], ptrs[id], DATA_SIZE, cudaMemcpyDeviceToHost, stream));\n\n // And wait for all the queued up work to complete\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n printf(\"Process %d: verifying...\\n\", id);\n\n // The contents should have the id of the sibling just after me\n char compareId = (char)((id + 1) % procCount);\n for (unsigned long long j = 0; j < DATA_SIZE; j++) {\n if (verification_buffer[j] != compareId) {\n printf(\"Process %d: Verification mismatch at %lld: %d != %d\\n\",\n id,\n j,\n (int)verification_buffer[j],\n (int)compareId);\n }\n }\n\n // Clean up!\n for (i = 0; i < procCount; i++) {\n // Free the memory before the exporter process frees it\n checkCudaErrors(cudaFreeAsync(ptrs[i], stream));\n }\n\n // And wait for all the queued up work to complete\n checkCudaErrors(cudaStreamSynchronize(stream));\n checkCudaErrors(cudaStreamDestroy(stream));\n\n printf(\"Process %d complete!\\n\", id);\n}\n\nstatic void parentProcess(char *app)\n{\n sharedMemoryInfo info;\n int devCount, i;\n volatile shmStruct *shm = NULL;\n std::vector ptrs;\n std::vector processes;\n cudaMemAllocationHandleType handleType = cudaMemHandleTypeNone;\n pid_t pid;\n char pidString[20] = {0};\n char lshmName[40] = {0};\n\n pid = getpid();\n snprintf(pidString, sizeof(pidString), \"%d\", pid);\n strcat(lshmName, shmName);\n strcat(lshmName, pidString);\n\n printf(\"PP: lshmName = %s\\n\", lshmName);\n\n checkCudaErrors(cudaGetDeviceCount(&devCount));\n std::vector devices(devCount);\n for (i = 0; i < devCount; i++) {\n cuDeviceGet(&devices[i], i);\n }\n\n if (sharedMemoryCreate(lshmName, sizeof(*shm), &info) != 0) {\n printf(\"Failed to create shared memory slab\\n\");\n exit(EXIT_FAILURE);\n }\n shm = (volatile shmStruct *)info.addr;\n memset((void *)shm, 0, sizeof(*shm));\n\n // Pick all the devices that can access each other's memory for this test\n // Keep in mind that CUDA has minimal support for fork() without a\n // corresponding exec() in the child process, but in this case our\n // spawnProcess will always exec, so no need to worry.\n for (i = 0; i < devCount; i++) {\n bool allPeers = true;\n cudaDeviceProp prop;\n checkCudaErrors(cudaGetDeviceProperties(&prop, i));\n\n int isMemPoolSupported = 0;\n checkCudaErrors(cudaDeviceGetAttribute(&isMemPoolSupported, cudaDevAttrMemoryPoolsSupported, i));\n // CUDA IPC is only supported on devices with unified addressing\n if (!isMemPoolSupported) {\n printf(\"Device %d does not support cuda memory pools, skipping...\\n\", i);\n continue;\n }\n int supportedHandleTypes = 0;\n checkCudaErrors(cudaDeviceGetAttribute(&supportedHandleTypes, cudaDevAttrMemoryPoolSupportedHandleTypes, i));\n if (supportedHandleTypes == 0) {\n printf(\"Device %d does not support Memory pool based IPC, skipping...\\n\", i);\n continue;\n }\n\n if (handleType == cudaMemHandleTypeNone) {\n if (supportedHandleTypes & cudaMemHandleTypePosixFileDescriptor) {\n handleType = cudaMemHandleTypePosixFileDescriptor;\n }\n else if (supportedHandleTypes & cudaMemHandleTypeWin32) {\n handleType = cudaMemHandleTypeWin32;\n }\n else {\n printf(\"Device %d does not support any supported handle types, skipping...\\n\", i);\n continue;\n }\n }\n else {\n if ((supportedHandleTypes & handleType) != handleType) {\n printf(\"Mixed handle types are not supported, waiving test\\n\");\n exit(EXIT_WAIVED);\n }\n }\n // This sample requires two processes accessing each device, so we need\n // to ensure exclusive or prohibited mode is not set\n int computeMode;\n checkCudaErrors(cudaDeviceGetAttribute(&computeMode, cudaDevAttrComputeMode, i));\n if (computeMode != cudaComputeModeDefault) {\n printf(\"Device %d is in an unsupported compute mode for this sample\\n\", i);\n continue;\n }\n#if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)\n // CUDA IPC on Windows is only supported on TCC\n if (!prop.tccDriver) {\n printf(\"Device %d is not in TCC mode\\n\", i);\n continue;\n }\n#endif\n\n for (int j = 0; j < shm->nprocesses; j++) {\n int canAccessPeerIJ, canAccessPeerJI;\n checkCudaErrors(cudaDeviceCanAccessPeer(&canAccessPeerJI, shm->devices[j], i));\n checkCudaErrors(cudaDeviceCanAccessPeer(&canAccessPeerIJ, i, shm->devices[j]));\n if (!canAccessPeerIJ || !canAccessPeerJI) {\n allPeers = false;\n break;\n }\n }\n if (allPeers) {\n // Enable peers here. This isn't necessary for IPC, but it will\n // setup the peers for the device. For systems that only allow 8\n // peers per GPU at a time, this acts to remove devices from CanAccessPeer\n for (int j = 0; j < shm->nprocesses; j++) {\n checkCudaErrors(cudaSetDevice(i));\n checkCudaErrors(cudaDeviceEnablePeerAccess(shm->devices[j], 0));\n checkCudaErrors(cudaSetDevice(shm->devices[j]));\n checkCudaErrors(cudaDeviceEnablePeerAccess(i, 0));\n }\n shm->devices[shm->nprocesses++] = i;\n if (shm->nprocesses >= MAX_DEVICES)\n break;\n }\n else {\n printf(\"Device %d is not peer capable with some other selected peers, \"\n \"skipping\\n\",\n i);\n }\n }\n\n if (shm->nprocesses == 0) {\n printf(\"No CUDA devices support IPC\\n\");\n exit(EXIT_WAIVED);\n }\n\n if (handleType == cudaMemHandleTypeNone) {\n printf(\"No supported handle types found, waiving test\\n\");\n exit(EXIT_WAIVED);\n }\n\n std::vector shareableHandles(shm->nprocesses);\n std::vector streams(shm->nprocesses);\n std::vector pools(shm->nprocesses);\n\n // Now allocate memory for each process and fill the shared\n // memory buffer with the export data and get memPool handles to communicate\n for (i = 0; i < shm->nprocesses; i++) {\n void *ptr = NULL;\n checkCudaErrors(cudaSetDevice(shm->devices[i]));\n checkCudaErrors(cudaStreamCreateWithFlags(&streams[i], cudaStreamNonBlocking));\n // Allocate an explicit pool with IPC capabilities\n cudaMemPoolProps poolProps;\n memset(&poolProps, 0, sizeof(cudaMemPoolProps));\n poolProps.allocType = cudaMemAllocationTypePinned;\n poolProps.handleTypes = handleType;\n\n poolProps.location.type = cudaMemLocationTypeDevice;\n poolProps.location.id = shm->devices[i];\n\n checkCudaErrors(cudaMemPoolCreate(&pools[i], &poolProps));\n\n // Query the shareable handle for the pool\n // Allocate memory in a stream from the pool just created\n checkCudaErrors(cudaMallocAsync(&ptr, DATA_SIZE, pools[i], streams[i]));\n\n checkCudaErrors(cudaMemPoolExportToShareableHandle(&shareableHandles[i], pools[i], handleType, 0));\n\n // Memset handle to 0 to make sure call to `cudaMemPoolImportPointer` in\n // childProcess will fail if the following call to\n // `cudaMemPoolExportPointer` fails.\n memset((void *)&shm->exportPtrData[i], 0, sizeof(cudaMemPoolPtrExportData));\n // Get the opaque \u2018bag-of-bits\u2019 representing the allocation\n checkCudaErrors(cudaMemPoolExportPointer((cudaMemPoolPtrExportData *)&shm->exportPtrData[i], ptr));\n ptrs.push_back(ptr);\n }\n\n shm->handleType = handleType;\n\n // Launch the child processes!\n for (i = 0; i < shm->nprocesses; i++) {\n char devIdx[10];\n char *const args[] = {app, devIdx, NULL};\n Process process;\n\n SPRINTF(devIdx, \"%d\", i);\n\n if (spawnProcess(&process, app, args)) {\n printf(\"Failed to create process\\n\");\n exit(EXIT_FAILURE);\n }\n\n processes.push_back(process);\n }\n\n barrierWait(&shm->barrier, &shm->sense, (unsigned int)(shm->nprocesses + 1));\n\n ipcHandle *ipcParentHandle = NULL;\n checkIpcErrors(ipcCreateSocket(ipcParentHandle, ipcName, processes));\n checkIpcErrors(ipcSendShareableHandles(ipcParentHandle, shareableHandles, processes));\n\n // Close the shareable handles as they are not needed anymore.\n for (int i = 0; i < shm->nprocesses; i++) {\n checkIpcErrors(ipcCloseShareableHandle(shareableHandles[i]));\n }\n checkIpcErrors(ipcCloseSocket(ipcParentHandle));\n\n // And wait for them to finish\n for (i = 0; i < processes.size(); i++) {\n if (waitProcess(&processes[i]) != EXIT_SUCCESS) {\n printf(\"Process %d failed!\\n\", i);\n exit(EXIT_FAILURE);\n }\n }\n\n // Clean up!\n for (i = 0; i < shm->nprocesses; i++) {\n checkCudaErrors(cudaSetDevice(shm->devices[i]));\n checkCudaErrors(cudaFreeAsync(ptrs[i], streams[i]));\n checkCudaErrors(cudaStreamSynchronize(streams[i]));\n checkCudaErrors(cudaMemPoolDestroy(pools[i]));\n }\n\n sharedMemoryClose(&info);\n}\n\n// Host code\nint main(int argc, char **argv)\n{\n#if defined(__arm__) || defined(__aarch64__) || defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)\n printf(\"Not supported on ARM or Windows\\n\");\n return EXIT_WAIVED;\n#else\n if (argc == 1) {\n parentProcess(argv[0]);\n }\n else {\n childProcess(atoi(argv[1]));\n }\n return EXIT_SUCCESS;\n#endif\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates Inter Process Communication\n * using one process per GPU for computation.\n */\n\n#include \n#include \n#include \n#include \n#define CUDA_DRIVER_API 1\n#include \"hip/hip_runtime_api.h\"\n#include \"helper_cuda_drvapi.h\"\n#include \"helper_multiprocess.h\"\n\nstatic const char shmName[] = \"streamOrderedAllocationIPCshm\";\nstatic const char ipcName[] = \"streamOrderedAllocationIPC_pipe\";\n// For direct NVLINK and PCI-E peers, at max 8 simultaneous peers are allowed\n// For NVSWITCH connected peers like DGX-2, simultaneous peers are not limited\n// in the same way.\n#define MAX_DEVICES (32)\n#define DATA_SIZE (64ULL << 20ULL) // 64MB\n\n#if defined(__linux__)\n#define cpu_atomic_add32(a, x) __sync_add_and_fetch(a, x)\n#elif defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)\n#define cpu_atomic_add32(a, x) InterlockedAdd((volatile LONG *)a, x)\n#else\n#error Unsupported system\n#endif\n\ntypedef struct shmStruct_st\n{\n size_t nprocesses;\n int barrier;\n int sense;\n hipMemAllocationHandleType handleType;\n int devices[MAX_DEVICES];\n hipMemPoolPtrExportData exportPtrData[MAX_DEVICES];\n} shmStruct;\n\n__global__ void simpleKernel(char *ptr, int sz, char val)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n for (; idx < sz; idx += (gridDim.x * blockDim.x)) {\n ptr[idx] = val;\n }\n}\n\nstatic void barrierWait(volatile int *barrier, volatile int *sense, unsigned int n)\n{\n int count;\n\n // Check-in\n count = cpu_atomic_add32(barrier, 1);\n if (count == n) // Last one in\n *sense = 1;\n while (!*sense)\n ;\n\n // Check-out\n count = cpu_atomic_add32(barrier, -1);\n if (count == 0) // Last one out\n *sense = 0;\n while (*sense)\n ;\n}\n\nstatic void childProcess(int id)\n{\n volatile shmStruct *shm = NULL;\n hipStream_t stream;\n sharedMemoryInfo info;\n size_t procCount, i;\n int blocks = 0;\n int threads = 128;\n hipDeviceProp_t prop;\n std::vector ptrs;\n pid_t pid;\n char pidString[20] = {0};\n char lshmName[40] = {0};\n\n std::vector verification_buffer(DATA_SIZE);\n\n pid = getppid();\n snprintf(pidString, sizeof(pidString), \"%d\", pid);\n strcat(lshmName, shmName);\n strcat(lshmName, pidString);\n\n printf(\"CP: lshmName = %s\\n\", lshmName);\n\n ipcHandle *ipcChildHandle = NULL;\n checkIpcErrors(ipcOpenSocket(ipcChildHandle));\n\n if (sharedMemoryOpen(lshmName, sizeof(shmStruct), &info) != 0) {\n printf(\"Failed to create shared memory slab\\n\");\n exit(EXIT_FAILURE);\n }\n shm = (volatile shmStruct *)info.addr;\n procCount = shm->nprocesses;\n\n barrierWait(&shm->barrier, &shm->sense, (unsigned int)(procCount + 1));\n\n // Receive all allocation handles shared by Parent.\n std::vector shHandle(shm->nprocesses);\n checkIpcErrors(ipcRecvShareableHandles(ipcChildHandle, shHandle));\n\n checkCudaErrors(hipSetDevice(shm->devices[id]));\n checkCudaErrors(hipGetDeviceProperties(&prop, shm->devices[id]));\n checkCudaErrors(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));\n checkCudaErrors(hipOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, simpleKernel, threads, 0));\n blocks *= prop.multiProcessorCount;\n\n std::vector pools(shm->nprocesses);\n\n hipMemAllocationHandleType handleType = shm->handleType;\n\n // Import mem pools from all the devices created in the master\n // process using shareable handles received via socket\n // and import the pointer to the allocated buffer using\n // exportData filled in shared memory by the master process.\n for (i = 0; i < procCount; i++) {\n checkCudaErrors(hipMemPoolImportFromShareableHandle(&pools[i], (void *)shHandle[i], handleType, 0));\n\n hipMemAccessFlags accessFlags;\n hipMemLocation location;\n location.type = hipMemLocationTypeDevice;\n location.id = shm->devices[id];\n checkCudaErrors(hipMemPoolGetAccess(&accessFlags, pools[i], &location));\n if (accessFlags != hipMemAccessFlagsProtReadWrite) {\n hipMemAccessDesc desc;\n memset(&desc, 0, sizeof(hipMemAccessDesc));\n desc.location.type = hipMemLocationTypeDevice;\n desc.location.id = shm->devices[id];\n desc.flags = hipMemAccessFlagsProtReadWrite;\n checkCudaErrors(hipMemPoolSetAccess(pools[i], &desc, 1));\n }\n\n // Import the allocation from each memory pool by iterating over exportData\n // until import is success.\n for (int j = 0; j < procCount; j++) {\n void *ptr = NULL;\n // Import the allocation using the opaque export data retrieved through\n // the shared memory\".\n hipError_t ret =\n hipMemPoolImportPointer(&ptr, pools[i], (hipMemPoolPtrExportData *)&shm->exportPtrData[j]);\n\n if (ret == hipSuccess) {\n // Pointer import is successful hence add it to the ptrs bag.\n ptrs.push_back(ptr);\n break;\n }\n else {\n // Reset failure error received from cudaMemPoolImportPointer\n // for further try.\n hipGetLastError();\n }\n }\n // Since we have imported allocations shared by the parent with us, we can\n // close this ShareableHandle.\n checkIpcErrors(ipcCloseShareableHandle(shHandle[i]));\n }\n\n // Since we have imported allocations shared by the parent with us, we can\n // close the socket.\n checkIpcErrors(ipcCloseSocket(ipcChildHandle));\n\n // At each iteration of the loop, each sibling process will push work on\n // their respective devices accessing the next peer mapped buffer allocated\n // by the master process (these can come from other sibling processes as\n // well). To coordinate each process' access, we force the stream to wait for\n // the work already accessing this buffer.\n for (i = 0; i < procCount; i++) {\n size_t bufferId = (i + id) % procCount;\n\n // Push a simple kernel on it\n simpleKernel<<>>((char *)ptrs[bufferId], DATA_SIZE, id);\n checkCudaErrors(hipGetLastError());\n checkCudaErrors(hipStreamSynchronize(stream));\n\n // Wait for all my sibling processes to push this stage of their work\n // before proceeding to the next. This prevents siblings from racing\n // ahead and clobbering the recorded event or waiting on the wrong\n // recorded event.\n barrierWait(&shm->barrier, &shm->sense, (unsigned int)procCount);\n if (id == 0) {\n printf(\"Step %lld done\\n\", (unsigned long long)i);\n }\n }\n\n // Now wait for my buffer to be ready so I can copy it locally and verify it\n checkCudaErrors(hipMemcpyAsync(&verification_buffer[0], ptrs[id], DATA_SIZE, hipMemcpyDeviceToHost, stream));\n\n // And wait for all the queued up work to complete\n checkCudaErrors(hipStreamSynchronize(stream));\n\n printf(\"Process %d: verifying...\\n\", id);\n\n // The contents should have the id of the sibling just after me\n char compareId = (char)((id + 1) % procCount);\n for (unsigned long long j = 0; j < DATA_SIZE; j++) {\n if (verification_buffer[j] != compareId) {\n printf(\"Process %d: Verification mismatch at %lld: %d != %d\\n\",\n id,\n j,\n (int)verification_buffer[j],\n (int)compareId);\n }\n }\n\n // Clean up!\n for (i = 0; i < procCount; i++) {\n // Free the memory before the exporter process frees it\n checkCudaErrors(hipFreeAsync(ptrs[i], stream));\n }\n\n // And wait for all the queued up work to complete\n checkCudaErrors(hipStreamSynchronize(stream));\n checkCudaErrors(hipStreamDestroy(stream));\n\n printf(\"Process %d complete!\\n\", id);\n}\n\nstatic void parentProcess(char *app)\n{\n sharedMemoryInfo info;\n int devCount, i;\n volatile shmStruct *shm = NULL;\n std::vector ptrs;\n std::vector processes;\n hipMemAllocationHandleType handleType = hipMemHandleTypeNone;\n pid_t pid;\n char pidString[20] = {0};\n char lshmName[40] = {0};\n\n pid = getpid();\n snprintf(pidString, sizeof(pidString), \"%d\", pid);\n strcat(lshmName, shmName);\n strcat(lshmName, pidString);\n\n printf(\"PP: lshmName = %s\\n\", lshmName);\n\n checkCudaErrors(hipGetDeviceCount(&devCount));\n std::vector devices(devCount);\n for (i = 0; i < devCount; i++) {\n hipDeviceGet(&devices[i], i);\n }\n\n if (sharedMemoryCreate(lshmName, sizeof(*shm), &info) != 0) {\n printf(\"Failed to create shared memory slab\\n\");\n exit(EXIT_FAILURE);\n }\n shm = (volatile shmStruct *)info.addr;\n memset((void *)shm, 0, sizeof(*shm));\n\n // Pick all the devices that can access each other's memory for this test\n // Keep in mind that CUDA has minimal support for fork() without a\n // corresponding exec() in the child process, but in this case our\n // spawnProcess will always exec, so no need to worry.\n for (i = 0; i < devCount; i++) {\n bool allPeers = true;\n hipDeviceProp_t prop;\n checkCudaErrors(hipGetDeviceProperties(&prop, i));\n\n int isMemPoolSupported = 0;\n checkCudaErrors(hipDeviceGetAttribute(&isMemPoolSupported, hipDeviceAttributeMemoryPoolsSupported, i));\n // CUDA IPC is only supported on devices with unified addressing\n if (!isMemPoolSupported) {\n printf(\"Device %d does not support cuda memory pools, skipping...\\n\", i);\n continue;\n }\n int supportedHandleTypes = 0;\n checkCudaErrors(hipDeviceGetAttribute(&supportedHandleTypes, cudaDevAttrMemoryPoolSupportedHandleTypes, i));\n if (supportedHandleTypes == 0) {\n printf(\"Device %d does not support Memory pool based IPC, skipping...\\n\", i);\n continue;\n }\n\n if (handleType == hipMemHandleTypeNone) {\n if (supportedHandleTypes & hipMemHandleTypePosixFileDescriptor) {\n handleType = hipMemHandleTypePosixFileDescriptor;\n }\n else if (supportedHandleTypes & hipMemHandleTypeWin32) {\n handleType = hipMemHandleTypeWin32;\n }\n else {\n printf(\"Device %d does not support any supported handle types, skipping...\\n\", i);\n continue;\n }\n }\n else {\n if ((supportedHandleTypes & handleType) != handleType) {\n printf(\"Mixed handle types are not supported, waiving test\\n\");\n exit(EXIT_WAIVED);\n }\n }\n // This sample requires two processes accessing each device, so we need\n // to ensure exclusive or prohibited mode is not set\n int computeMode;\n checkCudaErrors(hipDeviceGetAttribute(&computeMode, hipDeviceAttributeComputeMode, i));\n if (computeMode != hipComputeModeDefault) {\n printf(\"Device %d is in an unsupported compute mode for this sample\\n\", i);\n continue;\n }\n#if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)\n // CUDA IPC on Windows is only supported on TCC\n if (!prop.tccDriver) {\n printf(\"Device %d is not in TCC mode\\n\", i);\n continue;\n }\n#endif\n\n for (int j = 0; j < shm->nprocesses; j++) {\n int canAccessPeerIJ, canAccessPeerJI;\n checkCudaErrors(hipDeviceCanAccessPeer(&canAccessPeerJI, shm->devices[j], i));\n checkCudaErrors(hipDeviceCanAccessPeer(&canAccessPeerIJ, i, shm->devices[j]));\n if (!canAccessPeerIJ || !canAccessPeerJI) {\n allPeers = false;\n break;\n }\n }\n if (allPeers) {\n // Enable peers here. This isn't necessary for IPC, but it will\n // setup the peers for the device. For systems that only allow 8\n // peers per GPU at a time, this acts to remove devices from CanAccessPeer\n for (int j = 0; j < shm->nprocesses; j++) {\n checkCudaErrors(hipSetDevice(i));\n checkCudaErrors(hipDeviceEnablePeerAccess(shm->devices[j], 0));\n checkCudaErrors(hipSetDevice(shm->devices[j]));\n checkCudaErrors(hipDeviceEnablePeerAccess(i, 0));\n }\n shm->devices[shm->nprocesses++] = i;\n if (shm->nprocesses >= MAX_DEVICES)\n break;\n }\n else {\n printf(\"Device %d is not peer capable with some other selected peers, \"\n \"skipping\\n\",\n i);\n }\n }\n\n if (shm->nprocesses == 0) {\n printf(\"No CUDA devices support IPC\\n\");\n exit(EXIT_WAIVED);\n }\n\n if (handleType == hipMemHandleTypeNone) {\n printf(\"No supported handle types found, waiving test\\n\");\n exit(EXIT_WAIVED);\n }\n\n std::vector shareableHandles(shm->nprocesses);\n std::vector streams(shm->nprocesses);\n std::vector pools(shm->nprocesses);\n\n // Now allocate memory for each process and fill the shared\n // memory buffer with the export data and get memPool handles to communicate\n for (i = 0; i < shm->nprocesses; i++) {\n void *ptr = NULL;\n checkCudaErrors(hipSetDevice(shm->devices[i]));\n checkCudaErrors(hipStreamCreateWithFlags(&streams[i], hipStreamNonBlocking));\n // Allocate an explicit pool with IPC capabilities\n hipMemPoolProps poolProps;\n memset(&poolProps, 0, sizeof(hipMemPoolProps));\n poolProps.allocType = hipMemAllocationTypePinned;\n poolProps.handleTypes = handleType;\n\n poolProps.location.type = hipMemLocationTypeDevice;\n poolProps.location.id = shm->devices[i];\n\n checkCudaErrors(hipMemPoolCreate(&pools[i], &poolProps));\n\n // Query the shareable handle for the pool\n // Allocate memory in a stream from the pool just created\n checkCudaErrors(hipMallocAsync(&ptr, DATA_SIZE, pools[i], streams[i]));\n\n checkCudaErrors(hipMemPoolExportToShareableHandle(&shareableHandles[i], pools[i], handleType, 0));\n\n // Memset handle to 0 to make sure call to `cudaMemPoolImportPointer` in\n // childProcess will fail if the following call to\n // `cudaMemPoolExportPointer` fails.\n memset((void *)&shm->exportPtrData[i], 0, sizeof(hipMemPoolPtrExportData));\n // Get the opaque \u2018bag-of-bits\u2019 representing the allocation\n checkCudaErrors(hipMemPoolExportPointer((hipMemPoolPtrExportData *)&shm->exportPtrData[i], ptr));\n ptrs.push_back(ptr);\n }\n\n shm->handleType = handleType;\n\n // Launch the child processes!\n for (i = 0; i < shm->nprocesses; i++) {\n char devIdx[10];\n char *const args[] = {app, devIdx, NULL};\n Process process;\n\n SPRINTF(devIdx, \"%d\", i);\n\n if (spawnProcess(&process, app, args)) {\n printf(\"Failed to create process\\n\");\n exit(EXIT_FAILURE);\n }\n\n processes.push_back(process);\n }\n\n barrierWait(&shm->barrier, &shm->sense, (unsigned int)(shm->nprocesses + 1));\n\n ipcHandle *ipcParentHandle = NULL;\n checkIpcErrors(ipcCreateSocket(ipcParentHandle, ipcName, processes));\n checkIpcErrors(ipcSendShareableHandles(ipcParentHandle, shareableHandles, processes));\n\n // Close the shareable handles as they are not needed anymore.\n for (int i = 0; i < shm->nprocesses; i++) {\n checkIpcErrors(ipcCloseShareableHandle(shareableHandles[i]));\n }\n checkIpcErrors(ipcCloseSocket(ipcParentHandle));\n\n // And wait for them to finish\n for (i = 0; i < processes.size(); i++) {\n if (waitProcess(&processes[i]) != EXIT_SUCCESS) {\n printf(\"Process %d failed!\\n\", i);\n exit(EXIT_FAILURE);\n }\n }\n\n // Clean up!\n for (i = 0; i < shm->nprocesses; i++) {\n checkCudaErrors(hipSetDevice(shm->devices[i]));\n checkCudaErrors(hipFreeAsync(ptrs[i], streams[i]));\n checkCudaErrors(hipStreamSynchronize(streams[i]));\n checkCudaErrors(hipMemPoolDestroy(pools[i]));\n }\n\n sharedMemoryClose(&info);\n}\n\n// Host code\nint main(int argc, char **argv)\n{\n#if defined(__arm__) || defined(__aarch64__) || defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)\n printf(\"Not supported on ARM or Windows\\n\");\n return EXIT_WAIVED;\n#else\n if (argc == 1) {\n parentProcess(argv[0]);\n }\n else {\n childProcess(atoi(argv[1]));\n }\n return EXIT_SUCCESS;\n#endif\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/streamOrderedAllocationIPC/streamOrderedAllocationIPC.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n Parallel reduction\n\n This sample shows how to perform a reduction operation on an array of values\n to produce a single value in a single kernel (as opposed to two or more\n kernel calls as shown in the \"reduction\" CUDA Sample). Single-pass\n reduction requires Cooperative Groups.\n\n Reductions are a very common computation in parallel algorithms. Any time\n an array of values needs to be reduced to a single value using a binary\n associative operator, a reduction can be used. Example applications include\n statistics computations such as mean and standard deviation, and image\n processing applications such as finding the total luminance of an\n image.\n\n This code performs sum reductions, but any associative operator such as\n min() or max() could also be used.\n\n It assumes the input size is a power of 2.\n\n COMMAND LINE ARGUMENTS\n\n \"--n=\" :Specify the number of elements to reduce (default 33554432)\n \"--threads=\" :Specify the number of threads per block (default 128)\n \"--maxblocks=\" :Specify the maximum number of thread blocks to launch\n (kernel 6 only, default 64)\n*/\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// includes, project\n#include \n#include \n#include \n\nconst char *sSDKsample = \"reductionMultiBlockCG\";\n\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n/*\n Parallel sum reduction using shared memory\n - takes log(n) steps for n input elements\n - uses n/2 threads\n - only works for power-of-2 arrays\n\n This version adds multiple elements per thread sequentially. This reduces the\n overall cost of the algorithm while keeping the work complexity O(n) and the\n step complexity O(log n).\n (Brent's Theorem optimization)\n\n See the CUDA SDK \"reduction\" sample for more information.\n*/\n\n__device__ void reduceBlock(double *sdata, const cg::thread_block &cta)\n{\n const unsigned int tid = cta.thread_rank();\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n sdata[tid] = cg::reduce(tile32, sdata[tid], cg::plus());\n cg::sync(cta);\n\n double beta = 0.0;\n if (cta.thread_rank() == 0) {\n beta = 0;\n for (int i = 0; i < blockDim.x; i += tile32.size()) {\n beta += sdata[i];\n }\n sdata[0] = beta;\n }\n cg::sync(cta);\n}\n\n// This reduction kernel reduces an arbitrary size array in a single kernel\n// invocation\n//\n// For more details on the reduction algorithm (notably the multi-pass\n// approach), see the \"reduction\" sample in the CUDA SDK.\nextern \"C\" __global__ void reduceSinglePassMultiBlockCG(const float *g_idata, float *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block block = cg::this_thread_block();\n cg::grid_group grid = cg::this_grid();\n\n extern double __shared__ sdata[];\n\n // Stride over grid and add the values to a shared memory buffer\n sdata[block.thread_rank()] = 0;\n\n for (int i = grid.thread_rank(); i < n; i += grid.size()) {\n sdata[block.thread_rank()] += g_idata[i];\n }\n\n cg::sync(block);\n\n // Reduce each block (called once per block)\n reduceBlock(sdata, block);\n // Write out the result to global memory\n if (block.thread_rank() == 0) {\n g_odata[blockIdx.x] = sdata[0];\n }\n cg::sync(grid);\n\n if (grid.thread_rank() == 0) {\n for (int block = 1; block < gridDim.x; block++) {\n g_odata[0] += g_odata[block];\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Wrapper function for kernel launch\n////////////////////////////////////////////////////////////////////////////////\nvoid call_reduceSinglePassMultiBlockCG(int size, int threads, int numBlocks, float *d_idata, float *d_odata)\n{\n int smemSize = threads * sizeof(double);\n void *kernelArgs[] = {\n (void *)&d_idata,\n (void *)&d_odata,\n (void *)&size,\n };\n\n dim3 dimBlock(threads, 1, 1);\n dim3 dimGrid(numBlocks, 1, 1);\n\n cudaLaunchCooperativeKernel((void *)reduceSinglePassMultiBlockCG, dimGrid, dimBlock, kernelArgs, smemSize, NULL);\n // check if kernel execution generated an error\n getLastCudaError(\"Kernel execution failed\");\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\nbool runTest(int argc, char **argv, int device);\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n cudaDeviceProp deviceProp = {0};\n int dev;\n\n printf(\"%s Starting...\\n\\n\", sSDKsample);\n\n dev = findCudaDevice(argc, (const char **)argv);\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, dev));\n if (!deviceProp.cooperativeLaunch) {\n printf(\"\\nSelected GPU (%d) does not support Cooperative Kernel Launch, \"\n \"Waiving the run\\n\",\n dev);\n exit(EXIT_WAIVED);\n }\n\n bool bTestPassed = false;\n bTestPassed = runTest(argc, argv, dev);\n\n exit(bTestPassed ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Compute sum reduction on CPU\n//! We use Kahan summation for an accurate sum of large arrays.\n//! http://en.wikipedia.org/wiki/Kahan_summation_algorithm\n//!\n//! @param data pointer to input data\n//! @param size number of input data elements\n////////////////////////////////////////////////////////////////////////////////\ntemplate T reduceCPU(T *data, int size)\n{\n T sum = data[0];\n T c = (T)0.0;\n\n for (int i = 1; i < size; i++) {\n T y = data[i] - c;\n T t = sum + y;\n c = (t - sum) - y;\n sum = t;\n }\n\n return sum;\n}\n\nunsigned int nextPow2(unsigned int x)\n{\n --x;\n x |= x >> 1;\n x |= x >> 2;\n x |= x >> 4;\n x |= x >> 8;\n x |= x >> 16;\n return ++x;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Compute the number of threads and blocks to use for the reduction\n// We set threads / block to the minimum of maxThreads and n/2.\n////////////////////////////////////////////////////////////////////////////////\nvoid getNumBlocksAndThreads(int n, int maxBlocks, int maxThreads, int &blocks, int &threads)\n{\n if (n == 1) {\n threads = 1;\n blocks = 1;\n }\n else {\n checkCudaErrors(cudaOccupancyMaxPotentialBlockSize(&blocks, &threads, reduceSinglePassMultiBlockCG));\n }\n\n blocks = min(maxBlocks, blocks);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// This function performs a reduction of the input data multiple times and\n// measures the average reduction time.\n////////////////////////////////////////////////////////////////////////////////\nfloat benchmarkReduce(int n,\n int numThreads,\n int numBlocks,\n int maxThreads,\n int maxBlocks,\n int testIterations,\n StopWatchInterface *timer,\n float *h_odata,\n float *d_idata,\n float *d_odata)\n{\n float gpu_result = 0;\n cudaError_t error;\n\n printf(\"\\nLaunching %s kernel\\n\", \"SinglePass Multi Block Cooperative Groups\");\n for (int i = 0; i < testIterations; ++i) {\n gpu_result = 0;\n sdkStartTimer(&timer);\n call_reduceSinglePassMultiBlockCG(n, numThreads, numBlocks, d_idata, d_odata);\n cudaDeviceSynchronize();\n sdkStopTimer(&timer);\n }\n\n // copy final sum from device to host\n error = cudaMemcpy(&gpu_result, d_odata, sizeof(float), cudaMemcpyDeviceToHost);\n checkCudaErrors(error);\n\n return gpu_result;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// The main function which runs the reduction test.\n////////////////////////////////////////////////////////////////////////////////\nbool runTest(int argc, char **argv, int device)\n{\n int size = 1 << 25; // number of elements to reduce\n bool bTestPassed = false;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"n\")) {\n size = getCmdLineArgumentInt(argc, (const char **)argv, \"n\");\n }\n\n printf(\"%d elements\\n\", size);\n\n // Set the device to be used\n cudaDeviceProp prop = {0};\n checkCudaErrors(cudaSetDevice(device));\n checkCudaErrors(cudaGetDeviceProperties(&prop, device));\n\n // create random input data on CPU\n unsigned int bytes = size * sizeof(float);\n\n float *h_idata = (float *)malloc(bytes);\n\n for (int i = 0; i < size; i++) {\n // Keep the numbers small so we don't get truncation error in the sum\n h_idata[i] = (rand() & 0xFF) / (float)RAND_MAX;\n }\n\n // Determine the launch configuration (threads, blocks)\n int maxThreads = 0;\n int maxBlocks = 0;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"threads\")) {\n maxThreads = getCmdLineArgumentInt(argc, (const char **)argv, \"threads\");\n }\n else {\n maxThreads = prop.maxThreadsPerBlock;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"maxblocks\")) {\n maxBlocks = getCmdLineArgumentInt(argc, (const char **)argv, \"maxblocks\");\n }\n else {\n maxBlocks = prop.multiProcessorCount * (prop.maxThreadsPerMultiProcessor / prop.maxThreadsPerBlock);\n }\n\n int numBlocks = 0;\n int numThreads = 0;\n getNumBlocksAndThreads(size, maxBlocks, maxThreads, numBlocks, numThreads);\n\n // We calculate the occupancy to know how many block can actually fit on the\n // GPU\n int numBlocksPerSm = 0;\n checkCudaErrors(cudaOccupancyMaxActiveBlocksPerMultiprocessor(\n &numBlocksPerSm, reduceSinglePassMultiBlockCG, numThreads, numThreads * sizeof(double)));\n\n int numSms = prop.multiProcessorCount;\n if (numBlocks > numBlocksPerSm * numSms) {\n numBlocks = numBlocksPerSm * numSms;\n }\n printf(\"numThreads: %d\\n\", numThreads);\n printf(\"numBlocks: %d\\n\", numBlocks);\n\n // allocate mem for the result on host side\n float *h_odata = (float *)malloc(numBlocks * sizeof(float));\n\n // allocate device memory and data\n float *d_idata = NULL;\n float *d_odata = NULL;\n\n checkCudaErrors(cudaMalloc((void **)&d_idata, bytes));\n checkCudaErrors(cudaMalloc((void **)&d_odata, numBlocks * sizeof(float)));\n\n // copy data directly to device memory\n checkCudaErrors(cudaMemcpy(d_idata, h_idata, bytes, cudaMemcpyHostToDevice));\n checkCudaErrors(cudaMemcpy(d_odata, h_idata, numBlocks * sizeof(float), cudaMemcpyHostToDevice));\n\n int testIterations = 100;\n\n StopWatchInterface *timer = 0;\n sdkCreateTimer(&timer);\n\n float gpu_result = 0;\n\n gpu_result = benchmarkReduce(\n size, numThreads, numBlocks, maxThreads, maxBlocks, testIterations, timer, h_odata, d_idata, d_odata);\n\n float reduceTime = sdkGetAverageTimerValue(&timer);\n printf(\"Average time: %f ms\\n\", reduceTime);\n printf(\"Bandwidth: %f GB/s\\n\\n\", (size * sizeof(int)) / (reduceTime * 1.0e6));\n\n // compute reference solution\n float cpu_result = reduceCPU(h_idata, size);\n printf(\"GPU result = %0.12f\\n\", gpu_result);\n printf(\"CPU result = %0.12f\\n\", cpu_result);\n\n double threshold = 1e-8 * size;\n double diff = abs((double)gpu_result - (double)cpu_result);\n bTestPassed = (diff < threshold);\n\n // cleanup\n sdkDeleteTimer(&timer);\n\n free(h_idata);\n free(h_odata);\n cudaFree(d_idata);\n cudaFree(d_odata);\n\n return bTestPassed;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n Parallel reduction\n\n This sample shows how to perform a reduction operation on an array of values\n to produce a single value in a single kernel (as opposed to two or more\n kernel calls as shown in the \"reduction\" CUDA Sample). Single-pass\n reduction requires Cooperative Groups.\n\n Reductions are a very common computation in parallel algorithms. Any time\n an array of values needs to be reduced to a single value using a binary\n associative operator, a reduction can be used. Example applications include\n statistics computations such as mean and standard deviation, and image\n processing applications such as finding the total luminance of an\n image.\n\n This code performs sum reductions, but any associative operator such as\n min() or max() could also be used.\n\n It assumes the input size is a power of 2.\n\n COMMAND LINE ARGUMENTS\n\n \"--n=\" :Specify the number of elements to reduce (default 33554432)\n \"--threads=\" :Specify the number of threads per block (default 128)\n \"--maxblocks=\" :Specify the maximum number of thread blocks to launch\n (kernel 6 only, default 64)\n*/\n\n// includes, system\n#include \n#include \n#include \n#include \n\n// includes, project\n#include \n#include \n#include \n\nconst char *sSDKsample = \"reductionMultiBlockCG\";\n\n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n/*\n Parallel sum reduction using shared memory\n - takes log(n) steps for n input elements\n - uses n/2 threads\n - only works for power-of-2 arrays\n\n This version adds multiple elements per thread sequentially. This reduces the\n overall cost of the algorithm while keeping the work complexity O(n) and the\n step complexity O(log n).\n (Brent's Theorem optimization)\n\n See the CUDA SDK \"reduction\" sample for more information.\n*/\n\n__device__ void reduceBlock(double *sdata, const cg::thread_block &cta)\n{\n const unsigned int tid = cta.thread_rank();\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n sdata[tid] = cg::reduce(tile32, sdata[tid], cg::plus());\n cg::sync(cta);\n\n double beta = 0.0;\n if (cta.thread_rank() == 0) {\n beta = 0;\n for (int i = 0; i < blockDim.x; i += tile32.size()) {\n beta += sdata[i];\n }\n sdata[0] = beta;\n }\n cg::sync(cta);\n}\n\n// This reduction kernel reduces an arbitrary size array in a single kernel\n// invocation\n//\n// For more details on the reduction algorithm (notably the multi-pass\n// approach), see the \"reduction\" sample in the CUDA SDK.\nextern \"C\" __global__ void reduceSinglePassMultiBlockCG(const float *g_idata, float *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block block = cg::this_thread_block();\n cg::grid_group grid = cg::this_grid();\n\n extern double __shared__ sdata[];\n\n // Stride over grid and add the values to a shared memory buffer\n sdata[block.thread_rank()] = 0;\n\n for (int i = grid.thread_rank(); i < n; i += grid.size()) {\n sdata[block.thread_rank()] += g_idata[i];\n }\n\n cg::sync(block);\n\n // Reduce each block (called once per block)\n reduceBlock(sdata, block);\n // Write out the result to global memory\n if (block.thread_rank() == 0) {\n g_odata[blockIdx.x] = sdata[0];\n }\n cg::sync(grid);\n\n if (grid.thread_rank() == 0) {\n for (int block = 1; block < gridDim.x; block++) {\n g_odata[0] += g_odata[block];\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Wrapper function for kernel launch\n////////////////////////////////////////////////////////////////////////////////\nvoid call_reduceSinglePassMultiBlockCG(int size, int threads, int numBlocks, float *d_idata, float *d_odata)\n{\n int smemSize = threads * sizeof(double);\n void *kernelArgs[] = {\n (void *)&d_idata,\n (void *)&d_odata,\n (void *)&size,\n };\n\n dim3 dimBlock(threads, 1, 1);\n dim3 dimGrid(numBlocks, 1, 1);\n\n hipLaunchCooperativeKernel(reinterpret_cast((void *)reduceSinglePassMultiBlockCG), dimGrid, dimBlock, kernelArgs, smemSize, NULL);\n // check if kernel execution generated an error\n getLastCudaError(\"Kernel execution failed\");\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// declaration, forward\nbool runTest(int argc, char **argv, int device);\n\n////////////////////////////////////////////////////////////////////////////////\n// Program main\n////////////////////////////////////////////////////////////////////////////////\nint main(int argc, char **argv)\n{\n hipDeviceProp_t deviceProp = {0};\n int dev;\n\n printf(\"%s Starting...\\n\\n\", sSDKsample);\n\n dev = findCudaDevice(argc, (const char **)argv);\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, dev));\n if (!deviceProp.cooperativeLaunch) {\n printf(\"\\nSelected GPU (%d) does not support Cooperative Kernel Launch, \"\n \"Waiving the run\\n\",\n dev);\n exit(EXIT_WAIVED);\n }\n\n bool bTestPassed = false;\n bTestPassed = runTest(argc, argv, dev);\n\n exit(bTestPassed ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n//! Compute sum reduction on CPU\n//! We use Kahan summation for an accurate sum of large arrays.\n//! http://en.wikipedia.org/wiki/Kahan_summation_algorithm\n//!\n//! @param data pointer to input data\n//! @param size number of input data elements\n////////////////////////////////////////////////////////////////////////////////\ntemplate T reduceCPU(T *data, int size)\n{\n T sum = data[0];\n T c = (T)0.0;\n\n for (int i = 1; i < size; i++) {\n T y = data[i] - c;\n T t = sum + y;\n c = (t - sum) - y;\n sum = t;\n }\n\n return sum;\n}\n\nunsigned int nextPow2(unsigned int x)\n{\n --x;\n x |= x >> 1;\n x |= x >> 2;\n x |= x >> 4;\n x |= x >> 8;\n x |= x >> 16;\n return ++x;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Compute the number of threads and blocks to use for the reduction\n// We set threads / block to the minimum of maxThreads and n/2.\n////////////////////////////////////////////////////////////////////////////////\nvoid getNumBlocksAndThreads(int n, int maxBlocks, int maxThreads, int &blocks, int &threads)\n{\n if (n == 1) {\n threads = 1;\n blocks = 1;\n }\n else {\n checkCudaErrors(hipOccupancyMaxPotentialBlockSize(&blocks, &threads, reduceSinglePassMultiBlockCG));\n }\n\n blocks = min(maxBlocks, blocks);\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// This function performs a reduction of the input data multiple times and\n// measures the average reduction time.\n////////////////////////////////////////////////////////////////////////////////\nfloat benchmarkReduce(int n,\n int numThreads,\n int numBlocks,\n int maxThreads,\n int maxBlocks,\n int testIterations,\n StopWatchInterface *timer,\n float *h_odata,\n float *d_idata,\n float *d_odata)\n{\n float gpu_result = 0;\n hipError_t error;\n\n printf(\"\\nLaunching %s kernel\\n\", \"SinglePass Multi Block Cooperative Groups\");\n for (int i = 0; i < testIterations; ++i) {\n gpu_result = 0;\n sdkStartTimer(&timer);\n call_reduceSinglePassMultiBlockCG(n, numThreads, numBlocks, d_idata, d_odata);\n hipDeviceSynchronize();\n sdkStopTimer(&timer);\n }\n\n // copy final sum from device to host\n error = hipMemcpy(&gpu_result, d_odata, sizeof(float), hipMemcpyDeviceToHost);\n checkCudaErrors(error);\n\n return gpu_result;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// The main function which runs the reduction test.\n////////////////////////////////////////////////////////////////////////////////\nbool runTest(int argc, char **argv, int device)\n{\n int size = 1 << 25; // number of elements to reduce\n bool bTestPassed = false;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"n\")) {\n size = getCmdLineArgumentInt(argc, (const char **)argv, \"n\");\n }\n\n printf(\"%d elements\\n\", size);\n\n // Set the device to be used\n hipDeviceProp_t prop = {0};\n checkCudaErrors(hipSetDevice(device));\n checkCudaErrors(hipGetDeviceProperties(&prop, device));\n\n // create random input data on CPU\n unsigned int bytes = size * sizeof(float);\n\n float *h_idata = (float *)malloc(bytes);\n\n for (int i = 0; i < size; i++) {\n // Keep the numbers small so we don't get truncation error in the sum\n h_idata[i] = (rand() & 0xFF) / (float)RAND_MAX;\n }\n\n // Determine the launch configuration (threads, blocks)\n int maxThreads = 0;\n int maxBlocks = 0;\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"threads\")) {\n maxThreads = getCmdLineArgumentInt(argc, (const char **)argv, \"threads\");\n }\n else {\n maxThreads = prop.maxThreadsPerBlock;\n }\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"maxblocks\")) {\n maxBlocks = getCmdLineArgumentInt(argc, (const char **)argv, \"maxblocks\");\n }\n else {\n maxBlocks = prop.multiProcessorCount * (prop.maxThreadsPerMultiProcessor / prop.maxThreadsPerBlock);\n }\n\n int numBlocks = 0;\n int numThreads = 0;\n getNumBlocksAndThreads(size, maxBlocks, maxThreads, numBlocks, numThreads);\n\n // We calculate the occupancy to know how many block can actually fit on the\n // GPU\n int numBlocksPerSm = 0;\n checkCudaErrors(hipOccupancyMaxActiveBlocksPerMultiprocessor(\n &numBlocksPerSm, reduceSinglePassMultiBlockCG, numThreads, numThreads * sizeof(double)));\n\n int numSms = prop.multiProcessorCount;\n if (numBlocks > numBlocksPerSm * numSms) {\n numBlocks = numBlocksPerSm * numSms;\n }\n printf(\"numThreads: %d\\n\", numThreads);\n printf(\"numBlocks: %d\\n\", numBlocks);\n\n // allocate mem for the result on host side\n float *h_odata = (float *)malloc(numBlocks * sizeof(float));\n\n // allocate device memory and data\n float *d_idata = NULL;\n float *d_odata = NULL;\n\n checkCudaErrors(hipMalloc((void **)&d_idata, bytes));\n checkCudaErrors(hipMalloc((void **)&d_odata, numBlocks * sizeof(float)));\n\n // copy data directly to device memory\n checkCudaErrors(hipMemcpy(d_idata, h_idata, bytes, hipMemcpyHostToDevice));\n checkCudaErrors(hipMemcpy(d_odata, h_idata, numBlocks * sizeof(float), hipMemcpyHostToDevice));\n\n int testIterations = 100;\n\n StopWatchInterface *timer = 0;\n sdkCreateTimer(&timer);\n\n float gpu_result = 0;\n\n gpu_result = benchmarkReduce(\n size, numThreads, numBlocks, maxThreads, maxBlocks, testIterations, timer, h_odata, d_idata, d_odata);\n\n float reduceTime = sdkGetAverageTimerValue(&timer);\n printf(\"Average time: %f ms\\n\", reduceTime);\n printf(\"Bandwidth: %f GB/s\\n\\n\", (size * sizeof(int)) / (reduceTime * 1.0e6));\n\n // compute reference solution\n float cpu_result = reduceCPU(h_idata, size);\n printf(\"GPU result = %0.12f\\n\", gpu_result);\n printf(\"CPU result = %0.12f\\n\", cpu_result);\n\n double threshold = 1e-8 * size;\n double diff = abs((double)gpu_result - (double)cpu_result);\n bTestPassed = (diff < threshold);\n\n // cleanup\n sdkDeleteTimer(&timer);\n\n free(h_idata);\n free(h_odata);\n hipFree(d_idata);\n hipFree(d_odata);\n\n return bTestPassed;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/reductionMultiBlockCG/reductionMultiBlockCG.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n Parallel reduction kernels\n*/\n\n#ifndef _REDUCE_KERNEL_H_\n#define _REDUCE_KERNEL_H_\n\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n// Utility class used to avoid linker errors with extern\n// unsized shared memory arrays with templated type\ntemplate struct SharedMemory\n{\n __device__ inline operator T *()\n {\n extern __shared__ int __smem[];\n return (T *)__smem;\n }\n\n __device__ inline operator const T *() const\n {\n extern __shared__ int __smem[];\n return (T *)__smem;\n }\n};\n\n// specialize for double to avoid unaligned memory\n// access compile errors\ntemplate <> struct SharedMemory\n{\n __device__ inline operator double *()\n {\n extern __shared__ double __smem_d[];\n return (double *)__smem_d;\n }\n\n __device__ inline operator const double *() const\n {\n extern __shared__ double __smem_d[];\n return (double *)__smem_d;\n }\n};\n\ntemplate __device__ __forceinline__ T warpReduceSum(unsigned int mask, T mySum)\n{\n for (int offset = warpSize / 2; offset > 0; offset /= 2) {\n mySum += __shfl_down_sync(mask, mySum, offset);\n }\n return mySum;\n}\n\n#if __CUDA_ARCH__ >= 800\n// Specialize warpReduceFunc for int inputs to use __reduce_add_sync intrinsic\n// when on SM 8.0 or higher\ntemplate <> __device__ __forceinline__ int warpReduceSum(unsigned int mask, int mySum)\n{\n mySum = __reduce_add_sync(mask, mySum);\n return mySum;\n}\n#endif\n\n/*\n Parallel sum reduction using shared memory\n - takes log(n) steps for n input elements\n - uses n threads\n - only works for power-of-2 arrays\n*/\n\n/* This reduction interleaves which threads are active by using the modulo\n operator. This operator is very expensive on GPUs, and the interleaved\n inactivity means that no whole warps are active, which is also very\n inefficient */\ntemplate __global__ void reduce0(T *g_idata, T *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n T *sdata = SharedMemory();\n\n // load shared mem\n unsigned int tid = threadIdx.x;\n unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n sdata[tid] = (i < n) ? g_idata[i] : 0;\n\n cg::sync(cta);\n\n // do reduction in shared mem\n for (unsigned int s = 1; s < blockDim.x; s *= 2) {\n // modulo arithmetic is slow!\n if ((tid % (2 * s)) == 0) {\n sdata[tid] += sdata[tid + s];\n }\n\n cg::sync(cta);\n }\n\n // write result for this block to global mem\n if (tid == 0)\n g_odata[blockIdx.x] = sdata[0];\n}\n\n/* This version uses contiguous threads, but its interleaved\n addressing results in many shared memory bank conflicts.\n*/\ntemplate __global__ void reduce1(T *g_idata, T *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n T *sdata = SharedMemory();\n\n // load shared mem\n unsigned int tid = threadIdx.x;\n unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n sdata[tid] = (i < n) ? g_idata[i] : 0;\n\n cg::sync(cta);\n\n // do reduction in shared mem\n for (unsigned int s = 1; s < blockDim.x; s *= 2) {\n int index = 2 * s * tid;\n\n if (index < blockDim.x) {\n sdata[index] += sdata[index + s];\n }\n\n cg::sync(cta);\n }\n\n // write result for this block to global mem\n if (tid == 0)\n g_odata[blockIdx.x] = sdata[0];\n}\n\n/*\n This version uses sequential addressing -- no divergence or bank conflicts.\n*/\ntemplate __global__ void reduce2(T *g_idata, T *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n T *sdata = SharedMemory();\n\n // load shared mem\n unsigned int tid = threadIdx.x;\n unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n sdata[tid] = (i < n) ? g_idata[i] : 0;\n\n cg::sync(cta);\n\n // do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (tid < s) {\n sdata[tid] += sdata[tid + s];\n }\n\n cg::sync(cta);\n }\n\n // write result for this block to global mem\n if (tid == 0)\n g_odata[blockIdx.x] = sdata[0];\n}\n\n/*\n This version uses n/2 threads --\n it performs the first level of reduction when reading from global memory.\n*/\ntemplate __global__ void reduce3(T *g_idata, T *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n T *sdata = SharedMemory();\n\n // perform first level of reduction,\n // reading from global memory, writing to shared memory\n unsigned int tid = threadIdx.x;\n unsigned int i = blockIdx.x * (blockDim.x * 2) + threadIdx.x;\n\n T mySum = (i < n) ? g_idata[i] : 0;\n\n if (i + blockDim.x < n)\n mySum += g_idata[i + blockDim.x];\n\n sdata[tid] = mySum;\n cg::sync(cta);\n\n // do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (tid < s) {\n sdata[tid] = mySum = mySum + sdata[tid + s];\n }\n\n cg::sync(cta);\n }\n\n // write result for this block to global mem\n if (tid == 0)\n g_odata[blockIdx.x] = mySum;\n}\n\n/*\n This version uses the warp shuffle operation if available to reduce\n warp synchronization. When shuffle is not available the final warp's\n worth of work is unrolled to reduce looping overhead.\n\n See\n http://devblogs.nvidia.com/parallelforall/faster-parallel-reductions-kepler/\n for additional information about using shuffle to perform a reduction\n within a warp.\n\n Note, this kernel needs a minimum of 64*sizeof(T) bytes of shared memory.\n In other words if blockSize <= 32, allocate 64*sizeof(T) bytes.\n If blockSize > 32, allocate blockSize*sizeof(T) bytes.\n*/\ntemplate __global__ void reduce4(T *g_idata, T *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n T *sdata = SharedMemory();\n\n // perform first level of reduction,\n // reading from global memory, writing to shared memory\n unsigned int tid = threadIdx.x;\n unsigned int i = blockIdx.x * (blockDim.x * 2) + threadIdx.x;\n\n T mySum = (i < n) ? g_idata[i] : 0;\n\n if (i + blockSize < n)\n mySum += g_idata[i + blockSize];\n\n sdata[tid] = mySum;\n cg::sync(cta);\n\n // do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 32; s >>= 1) {\n if (tid < s) {\n sdata[tid] = mySum = mySum + sdata[tid + s];\n }\n\n cg::sync(cta);\n }\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n if (cta.thread_rank() < 32) {\n // Fetch final intermediate sum from 2nd warp\n if (blockSize >= 64)\n mySum += sdata[tid + 32];\n // Reduce final warp using shuffle\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n mySum += tile32.shfl_down(mySum, offset);\n }\n }\n\n // write result for this block to global mem\n if (cta.thread_rank() == 0)\n g_odata[blockIdx.x] = mySum;\n}\n\n/*\n This version is completely unrolled, unless warp shuffle is available, then\n shuffle is used within a loop. It uses a template parameter to achieve\n optimal code for any (power of 2) number of threads. This requires a switch\n statement in the host code to handle all the different thread block sizes at\n compile time. When shuffle is available, it is used to reduce warp\n synchronization.\n\n Note, this kernel needs a minimum of 64*sizeof(T) bytes of shared memory.\n In other words if blockSize <= 32, allocate 64*sizeof(T) bytes.\n If blockSize > 32, allocate blockSize*sizeof(T) bytes.\n*/\ntemplate __global__ void reduce5(T *g_idata, T *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n T *sdata = SharedMemory();\n\n // perform first level of reduction,\n // reading from global memory, writing to shared memory\n unsigned int tid = threadIdx.x;\n unsigned int i = blockIdx.x * (blockSize * 2) + threadIdx.x;\n\n T mySum = (i < n) ? g_idata[i] : 0;\n\n if (i + blockSize < n)\n mySum += g_idata[i + blockSize];\n\n sdata[tid] = mySum;\n cg::sync(cta);\n\n // do reduction in shared mem\n if ((blockSize >= 512) && (tid < 256)) {\n sdata[tid] = mySum = mySum + sdata[tid + 256];\n }\n\n cg::sync(cta);\n\n if ((blockSize >= 256) && (tid < 128)) {\n sdata[tid] = mySum = mySum + sdata[tid + 128];\n }\n\n cg::sync(cta);\n\n if ((blockSize >= 128) && (tid < 64)) {\n sdata[tid] = mySum = mySum + sdata[tid + 64];\n }\n\n cg::sync(cta);\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n if (cta.thread_rank() < 32) {\n // Fetch final intermediate sum from 2nd warp\n if (blockSize >= 64)\n mySum += sdata[tid + 32];\n // Reduce final warp using shuffle\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n mySum += tile32.shfl_down(mySum, offset);\n }\n }\n\n // write result for this block to global mem\n if (cta.thread_rank() == 0)\n g_odata[blockIdx.x] = mySum;\n}\n\n/*\n This version adds multiple elements per thread sequentially. This reduces\n the overall cost of the algorithm while keeping the work complexity O(n) and\n the step complexity O(log n). (Brent's Theorem optimization)\n\n Note, this kernel needs a minimum of 64*sizeof(T) bytes of shared memory.\n In other words if blockSize <= 32, allocate 64*sizeof(T) bytes.\n If blockSize > 32, allocate blockSize*sizeof(T) bytes.\n*/\ntemplate __global__ void reduce6(T *g_idata, T *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n T *sdata = SharedMemory();\n\n // perform first level of reduction,\n // reading from global memory, writing to shared memory\n unsigned int tid = threadIdx.x;\n unsigned int gridSize = blockSize * gridDim.x;\n\n T mySum = 0;\n\n // we reduce multiple elements per thread. The number is determined by the\n // number of active thread blocks (via gridDim). More blocks will result\n // in a larger gridSize and therefore fewer elements per thread\n if (nIsPow2) {\n unsigned int i = blockIdx.x * blockSize * 2 + threadIdx.x;\n gridSize = gridSize << 1;\n\n while (i < n) {\n mySum += g_idata[i];\n // ensure we don't read out of bounds -- this is optimized away for\n // powerOf2 sized arrays\n if ((i + blockSize) < n) {\n mySum += g_idata[i + blockSize];\n }\n i += gridSize;\n }\n }\n else {\n unsigned int i = blockIdx.x * blockSize + threadIdx.x;\n while (i < n) {\n mySum += g_idata[i];\n i += gridSize;\n }\n }\n\n // each thread puts its local sum into shared memory\n sdata[tid] = mySum;\n cg::sync(cta);\n\n // do reduction in shared mem\n if ((blockSize >= 512) && (tid < 256)) {\n sdata[tid] = mySum = mySum + sdata[tid + 256];\n }\n\n cg::sync(cta);\n\n if ((blockSize >= 256) && (tid < 128)) {\n sdata[tid] = mySum = mySum + sdata[tid + 128];\n }\n\n cg::sync(cta);\n\n if ((blockSize >= 128) && (tid < 64)) {\n sdata[tid] = mySum = mySum + sdata[tid + 64];\n }\n\n cg::sync(cta);\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n if (cta.thread_rank() < 32) {\n // Fetch final intermediate sum from 2nd warp\n if (blockSize >= 64)\n mySum += sdata[tid + 32];\n // Reduce final warp using shuffle\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n mySum += tile32.shfl_down(mySum, offset);\n }\n }\n\n // write result for this block to global mem\n if (cta.thread_rank() == 0)\n g_odata[blockIdx.x] = mySum;\n}\n\ntemplate \n__global__ void reduce7(const T *__restrict__ g_idata, T *__restrict__ g_odata, unsigned int n)\n{\n T *sdata = SharedMemory();\n\n // perform first level of reduction,\n // reading from global memory, writing to shared memory\n unsigned int tid = threadIdx.x;\n unsigned int gridSize = blockSize * gridDim.x;\n unsigned int maskLength = (blockSize & 31); // 31 = warpSize-1\n maskLength = (maskLength > 0) ? (32 - maskLength) : maskLength;\n const unsigned int mask = (0xffffffff) >> maskLength;\n\n T mySum = 0;\n\n // we reduce multiple elements per thread. The number is determined by the\n // number of active thread blocks (via gridDim). More blocks will result\n // in a larger gridSize and therefore fewer elements per thread\n if (nIsPow2) {\n unsigned int i = blockIdx.x * blockSize * 2 + threadIdx.x;\n gridSize = gridSize << 1;\n\n while (i < n) {\n mySum += g_idata[i];\n // ensure we don't read out of bounds -- this is optimized away for\n // powerOf2 sized arrays\n if ((i + blockSize) < n) {\n mySum += g_idata[i + blockSize];\n }\n i += gridSize;\n }\n }\n else {\n unsigned int i = blockIdx.x * blockSize + threadIdx.x;\n while (i < n) {\n mySum += g_idata[i];\n i += gridSize;\n }\n }\n\n // Reduce within warp using shuffle or reduce_add if T==int & CUDA_ARCH ==\n // SM 8.0\n mySum = warpReduceSum(mask, mySum);\n\n // each thread puts its local sum into shared memory\n if ((tid % warpSize) == 0) {\n sdata[tid / warpSize] = mySum;\n }\n\n __syncthreads();\n\n const unsigned int shmem_extent = (blockSize / warpSize) > 0 ? (blockSize / warpSize) : 1;\n const unsigned int ballot_result = __ballot_sync(mask, tid < shmem_extent);\n if (tid < shmem_extent) {\n mySum = sdata[tid];\n // Reduce final warp using shuffle or reduce_add if T==int & CUDA_ARCH ==\n // SM 8.0\n mySum = warpReduceSum(ballot_result, mySum);\n }\n\n // write result for this block to global mem\n if (tid == 0) {\n g_odata[blockIdx.x] = mySum;\n }\n}\n\n// Performs a reduction step and updates numTotal with how many are remaining\ntemplate __device__ T cg_reduce_n(T in, Group &threads)\n{\n return cg::reduce(threads, in, cg::plus());\n}\n\ntemplate __global__ void cg_reduce(T *g_idata, T *g_odata, unsigned int n)\n{\n // Shared memory for intermediate steps\n T *sdata = SharedMemory();\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Handle to tile in thread block\n cg::thread_block_tile<32> tile = cg::tiled_partition<32>(cta);\n\n unsigned int ctaSize = cta.size();\n unsigned int numCtas = gridDim.x;\n unsigned int threadRank = cta.thread_rank();\n unsigned int threadIndex = (blockIdx.x * ctaSize) + threadRank;\n\n T threadVal = 0;\n {\n unsigned int i = threadIndex;\n unsigned int indexStride = (numCtas * ctaSize);\n while (i < n) {\n threadVal += g_idata[i];\n i += indexStride;\n }\n sdata[threadRank] = threadVal;\n }\n\n // Wait for all tiles to finish and reduce within CTA\n {\n unsigned int ctaSteps = tile.meta_group_size();\n unsigned int ctaIndex = ctaSize >> 1;\n while (ctaIndex >= 32) {\n cta.sync();\n if (threadRank < ctaIndex) {\n threadVal += sdata[threadRank + ctaIndex];\n sdata[threadRank] = threadVal;\n }\n ctaSteps >>= 1;\n ctaIndex >>= 1;\n }\n }\n\n // Shuffle redux instead of smem redux\n {\n cta.sync();\n if (tile.meta_group_rank() == 0) {\n threadVal = cg_reduce_n(threadVal, tile);\n }\n }\n\n if (threadRank == 0)\n g_odata[blockIdx.x] = threadVal;\n}\n\ntemplate \n__global__ void multi_warp_cg_reduce(T *g_idata, T *g_odata, unsigned int n)\n{\n // Shared memory for intermediate steps\n T *sdata = SharedMemory();\n __shared__ cg::block_tile_memory scratch;\n\n // Handle to thread block group\n auto cta = cg::this_thread_block(scratch);\n // Handle to multiWarpTile in thread block\n auto multiWarpTile = cg::tiled_partition(cta);\n\n unsigned int gridSize = BlockSize * gridDim.x;\n T threadVal = 0;\n\n // we reduce multiple elements per thread. The number is determined by the\n // number of active thread blocks (via gridDim). More blocks will result\n // in a larger gridSize and therefore fewer elements per thread\n int nIsPow2 = !(n & n - 1);\n if (nIsPow2) {\n unsigned int i = blockIdx.x * BlockSize * 2 + threadIdx.x;\n gridSize = gridSize << 1;\n\n while (i < n) {\n threadVal += g_idata[i];\n // ensure we don't read out of bounds -- this is optimized away for\n // powerOf2 sized arrays\n if ((i + BlockSize) < n) {\n threadVal += g_idata[i + blockDim.x];\n }\n i += gridSize;\n }\n }\n else {\n unsigned int i = blockIdx.x * BlockSize + threadIdx.x;\n while (i < n) {\n threadVal += g_idata[i];\n i += gridSize;\n }\n }\n\n threadVal = cg_reduce_n(threadVal, multiWarpTile);\n\n if (multiWarpTile.thread_rank() == 0) {\n sdata[multiWarpTile.meta_group_rank()] = threadVal;\n }\n cg::sync(cta);\n\n if (threadIdx.x == 0) {\n threadVal = 0;\n for (int i = 0; i < multiWarpTile.meta_group_size(); i++) {\n threadVal += sdata[i];\n }\n g_odata[blockIdx.x] = threadVal;\n }\n}\n\nextern \"C\" bool isPow2(unsigned int x);\n\n////////////////////////////////////////////////////////////////////////////////\n// Wrapper function for kernel launch\n////////////////////////////////////////////////////////////////////////////////\ntemplate void reduce(int size, int threads, int blocks, int whichKernel, T *d_idata, T *d_odata)\n{\n dim3 dimBlock(threads, 1, 1);\n dim3 dimGrid(blocks, 1, 1);\n\n // when there is only one warp per block, we need to allocate two warps\n // worth of shared memory so that we don't index shared memory out of bounds\n int smemSize = (threads <= 32) ? 2 * threads * sizeof(T) : threads * sizeof(T);\n\n // as kernel 9 - multi_warp_cg_reduce cannot work for more than 64 threads\n // we choose to set kernel 7 for this purpose.\n if (threads < 64 && whichKernel == 9) {\n whichKernel = 7;\n }\n\n // choose which of the optimized versions of reduction to launch\n switch (whichKernel) {\n case 0:\n reduce0<<>>(d_idata, d_odata, size);\n break;\n\n case 1:\n reduce1<<>>(d_idata, d_odata, size);\n break;\n\n case 2:\n reduce2<<>>(d_idata, d_odata, size);\n break;\n\n case 3:\n reduce3<<>>(d_idata, d_odata, size);\n break;\n\n case 4:\n switch (threads) {\n case 512:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 256:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 128:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 64:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 32:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 16:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 8:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 4:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 2:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 1:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n }\n\n break;\n\n case 5:\n switch (threads) {\n case 512:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 256:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 128:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 64:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 32:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 16:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 8:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 4:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 2:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 1:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n }\n\n break;\n\n case 6:\n if (isPow2(size)) {\n switch (threads) {\n case 512:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 256:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 128:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 64:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 32:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 16:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 8:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 4:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 2:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 1:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n }\n }\n else {\n switch (threads) {\n case 512:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 256:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 128:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 64:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 32:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 16:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 8:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 4:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 2:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 1:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n }\n }\n\n break;\n\n case 7:\n // For reduce7 kernel we require only blockSize/warpSize\n // number of elements in shared memory\n smemSize = ((threads / 32) + 1) * sizeof(T);\n if (isPow2(size)) {\n switch (threads) {\n case 1024:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n case 512:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 256:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 128:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 64:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 32:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 16:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 8:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 4:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 2:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 1:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n }\n }\n else {\n switch (threads) {\n case 1024:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n case 512:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 256:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 128:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 64:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 32:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 16:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 8:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 4:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 2:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 1:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n }\n }\n\n break;\n case 8:\n cg_reduce<<>>(d_idata, d_odata, size);\n break;\n case 9:\n constexpr int numOfMultiWarpGroups = 2;\n smemSize = numOfMultiWarpGroups * sizeof(T);\n switch (threads) {\n case 1024:\n multi_warp_cg_reduce\n <<>>(d_idata, d_odata, size);\n break;\n\n case 512:\n multi_warp_cg_reduce\n <<>>(d_idata, d_odata, size);\n break;\n\n case 256:\n multi_warp_cg_reduce\n <<>>(d_idata, d_odata, size);\n break;\n\n case 128:\n multi_warp_cg_reduce\n <<>>(d_idata, d_odata, size);\n break;\n\n case 64:\n multi_warp_cg_reduce\n <<>>(d_idata, d_odata, size);\n break;\n\n default:\n printf(\"thread block size of < 64 is not supported for this kernel\\n\");\n break;\n }\n break;\n }\n}\n\n// Instantiate the reduction function for 3 types\ntemplate void reduce(int size, int threads, int blocks, int whichKernel, int *d_idata, int *d_odata);\n\ntemplate void reduce(int size, int threads, int blocks, int whichKernel, float *d_idata, float *d_odata);\n\ntemplate void reduce(int size, int threads, int blocks, int whichKernel, double *d_idata, double *d_odata);\n\n#endif // #ifndef _REDUCE_KERNEL_H_\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n Parallel reduction kernels\n*/\n\n#include \"hip/hip_runtime.h\"\n#ifndef _REDUCE_KERNEL_H_\n#define _REDUCE_KERNEL_H_\n\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n// Utility class used to avoid linker errors with extern\n// unsized shared memory arrays with templated type\ntemplate struct SharedMemory\n{\n __device__ inline operator T *()\n {\n extern __shared__ int __smem[];\n return (T *)__smem;\n }\n\n __device__ inline operator const T *() const\n {\n extern __shared__ int __smem[];\n return (T *)__smem;\n }\n};\n\n// specialize for double to avoid unaligned memory\n// access compile errors\ntemplate <> struct SharedMemory\n{\n __device__ inline operator double *()\n {\n extern __shared__ double __smem_d[];\n return (double *)__smem_d;\n }\n\n __device__ inline operator const double *() const\n {\n extern __shared__ double __smem_d[];\n return (double *)__smem_d;\n }\n};\n\ntemplate __device__ __forceinline__ T warpReduceSum(unsigned int mask, T mySum)\n{\n for (int offset = warpSize / 2; offset > 0; offset /= 2) {\n mySum += __shfl_down_sync(mask, mySum, offset);\n }\n return mySum;\n}\n\n#if __CUDA_ARCH__ >= 800\n// Specialize warpReduceFunc for int inputs to use __reduce_add_sync intrinsic\n// when on SM 8.0 or higher\ntemplate <> __device__ __forceinline__ int warpReduceSum(unsigned int mask, int mySum)\n{\n mySum = __reduce_add_sync(mask, mySum);\n return mySum;\n}\n#endif\n\n/*\n Parallel sum reduction using shared memory\n - takes log(n) steps for n input elements\n - uses n threads\n - only works for power-of-2 arrays\n*/\n\n/* This reduction interleaves which threads are active by using the modulo\n operator. This operator is very expensive on GPUs, and the interleaved\n inactivity means that no whole warps are active, which is also very\n inefficient */\ntemplate __global__ void reduce0(T *g_idata, T *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n T *sdata = SharedMemory();\n\n // load shared mem\n unsigned int tid = threadIdx.x;\n unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n sdata[tid] = (i < n) ? g_idata[i] : 0;\n\n cg::sync(cta);\n\n // do reduction in shared mem\n for (unsigned int s = 1; s < blockDim.x; s *= 2) {\n // modulo arithmetic is slow!\n if ((tid % (2 * s)) == 0) {\n sdata[tid] += sdata[tid + s];\n }\n\n cg::sync(cta);\n }\n\n // write result for this block to global mem\n if (tid == 0)\n g_odata[blockIdx.x] = sdata[0];\n}\n\n/* This version uses contiguous threads, but its interleaved\n addressing results in many shared memory bank conflicts.\n*/\ntemplate __global__ void reduce1(T *g_idata, T *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n T *sdata = SharedMemory();\n\n // load shared mem\n unsigned int tid = threadIdx.x;\n unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n sdata[tid] = (i < n) ? g_idata[i] : 0;\n\n cg::sync(cta);\n\n // do reduction in shared mem\n for (unsigned int s = 1; s < blockDim.x; s *= 2) {\n int index = 2 * s * tid;\n\n if (index < blockDim.x) {\n sdata[index] += sdata[index + s];\n }\n\n cg::sync(cta);\n }\n\n // write result for this block to global mem\n if (tid == 0)\n g_odata[blockIdx.x] = sdata[0];\n}\n\n/*\n This version uses sequential addressing -- no divergence or bank conflicts.\n*/\ntemplate __global__ void reduce2(T *g_idata, T *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n T *sdata = SharedMemory();\n\n // load shared mem\n unsigned int tid = threadIdx.x;\n unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n sdata[tid] = (i < n) ? g_idata[i] : 0;\n\n cg::sync(cta);\n\n // do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (tid < s) {\n sdata[tid] += sdata[tid + s];\n }\n\n cg::sync(cta);\n }\n\n // write result for this block to global mem\n if (tid == 0)\n g_odata[blockIdx.x] = sdata[0];\n}\n\n/*\n This version uses n/2 threads --\n it performs the first level of reduction when reading from global memory.\n*/\ntemplate __global__ void reduce3(T *g_idata, T *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n T *sdata = SharedMemory();\n\n // perform first level of reduction,\n // reading from global memory, writing to shared memory\n unsigned int tid = threadIdx.x;\n unsigned int i = blockIdx.x * (blockDim.x * 2) + threadIdx.x;\n\n T mySum = (i < n) ? g_idata[i] : 0;\n\n if (i + blockDim.x < n)\n mySum += g_idata[i + blockDim.x];\n\n sdata[tid] = mySum;\n cg::sync(cta);\n\n // do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (tid < s) {\n sdata[tid] = mySum = mySum + sdata[tid + s];\n }\n\n cg::sync(cta);\n }\n\n // write result for this block to global mem\n if (tid == 0)\n g_odata[blockIdx.x] = mySum;\n}\n\n/*\n This version uses the warp shuffle operation if available to reduce\n warp synchronization. When shuffle is not available the final warp's\n worth of work is unrolled to reduce looping overhead.\n\n See\n http://devblogs.nvidia.com/parallelforall/faster-parallel-reductions-kepler/\n for additional information about using shuffle to perform a reduction\n within a warp.\n\n Note, this kernel needs a minimum of 64*sizeof(T) bytes of shared memory.\n In other words if blockSize <= 32, allocate 64*sizeof(T) bytes.\n If blockSize > 32, allocate blockSize*sizeof(T) bytes.\n*/\ntemplate __global__ void reduce4(T *g_idata, T *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n T *sdata = SharedMemory();\n\n // perform first level of reduction,\n // reading from global memory, writing to shared memory\n unsigned int tid = threadIdx.x;\n unsigned int i = blockIdx.x * (blockDim.x * 2) + threadIdx.x;\n\n T mySum = (i < n) ? g_idata[i] : 0;\n\n if (i + blockSize < n)\n mySum += g_idata[i + blockSize];\n\n sdata[tid] = mySum;\n cg::sync(cta);\n\n // do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 32; s >>= 1) {\n if (tid < s) {\n sdata[tid] = mySum = mySum + sdata[tid + s];\n }\n\n cg::sync(cta);\n }\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n if (cta.thread_rank() < 32) {\n // Fetch final intermediate sum from 2nd warp\n if (blockSize >= 64)\n mySum += sdata[tid + 32];\n // Reduce final warp using shuffle\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n mySum += tile32.shfl_down(mySum, offset);\n }\n }\n\n // write result for this block to global mem\n if (cta.thread_rank() == 0)\n g_odata[blockIdx.x] = mySum;\n}\n\n/*\n This version is completely unrolled, unless warp shuffle is available, then\n shuffle is used within a loop. It uses a template parameter to achieve\n optimal code for any (power of 2) number of threads. This requires a switch\n statement in the host code to handle all the different thread block sizes at\n compile time. When shuffle is available, it is used to reduce warp\n synchronization.\n\n Note, this kernel needs a minimum of 64*sizeof(T) bytes of shared memory.\n In other words if blockSize <= 32, allocate 64*sizeof(T) bytes.\n If blockSize > 32, allocate blockSize*sizeof(T) bytes.\n*/\ntemplate __global__ void reduce5(T *g_idata, T *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n T *sdata = SharedMemory();\n\n // perform first level of reduction,\n // reading from global memory, writing to shared memory\n unsigned int tid = threadIdx.x;\n unsigned int i = blockIdx.x * (blockSize * 2) + threadIdx.x;\n\n T mySum = (i < n) ? g_idata[i] : 0;\n\n if (i + blockSize < n)\n mySum += g_idata[i + blockSize];\n\n sdata[tid] = mySum;\n cg::sync(cta);\n\n // do reduction in shared mem\n if ((blockSize >= 512) && (tid < 256)) {\n sdata[tid] = mySum = mySum + sdata[tid + 256];\n }\n\n cg::sync(cta);\n\n if ((blockSize >= 256) && (tid < 128)) {\n sdata[tid] = mySum = mySum + sdata[tid + 128];\n }\n\n cg::sync(cta);\n\n if ((blockSize >= 128) && (tid < 64)) {\n sdata[tid] = mySum = mySum + sdata[tid + 64];\n }\n\n cg::sync(cta);\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n if (cta.thread_rank() < 32) {\n // Fetch final intermediate sum from 2nd warp\n if (blockSize >= 64)\n mySum += sdata[tid + 32];\n // Reduce final warp using shuffle\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n mySum += tile32.shfl_down(mySum, offset);\n }\n }\n\n // write result for this block to global mem\n if (cta.thread_rank() == 0)\n g_odata[blockIdx.x] = mySum;\n}\n\n/*\n This version adds multiple elements per thread sequentially. This reduces\n the overall cost of the algorithm while keeping the work complexity O(n) and\n the step complexity O(log n). (Brent's Theorem optimization)\n\n Note, this kernel needs a minimum of 64*sizeof(T) bytes of shared memory.\n In other words if blockSize <= 32, allocate 64*sizeof(T) bytes.\n If blockSize > 32, allocate blockSize*sizeof(T) bytes.\n*/\ntemplate __global__ void reduce6(T *g_idata, T *g_odata, unsigned int n)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n T *sdata = SharedMemory();\n\n // perform first level of reduction,\n // reading from global memory, writing to shared memory\n unsigned int tid = threadIdx.x;\n unsigned int gridSize = blockSize * gridDim.x;\n\n T mySum = 0;\n\n // we reduce multiple elements per thread. The number is determined by the\n // number of active thread blocks (via gridDim). More blocks will result\n // in a larger gridSize and therefore fewer elements per thread\n if (nIsPow2) {\n unsigned int i = blockIdx.x * blockSize * 2 + threadIdx.x;\n gridSize = gridSize << 1;\n\n while (i < n) {\n mySum += g_idata[i];\n // ensure we don't read out of bounds -- this is optimized away for\n // powerOf2 sized arrays\n if ((i + blockSize) < n) {\n mySum += g_idata[i + blockSize];\n }\n i += gridSize;\n }\n }\n else {\n unsigned int i = blockIdx.x * blockSize + threadIdx.x;\n while (i < n) {\n mySum += g_idata[i];\n i += gridSize;\n }\n }\n\n // each thread puts its local sum into shared memory\n sdata[tid] = mySum;\n cg::sync(cta);\n\n // do reduction in shared mem\n if ((blockSize >= 512) && (tid < 256)) {\n sdata[tid] = mySum = mySum + sdata[tid + 256];\n }\n\n cg::sync(cta);\n\n if ((blockSize >= 256) && (tid < 128)) {\n sdata[tid] = mySum = mySum + sdata[tid + 128];\n }\n\n cg::sync(cta);\n\n if ((blockSize >= 128) && (tid < 64)) {\n sdata[tid] = mySum = mySum + sdata[tid + 64];\n }\n\n cg::sync(cta);\n\n cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);\n\n if (cta.thread_rank() < 32) {\n // Fetch final intermediate sum from 2nd warp\n if (blockSize >= 64)\n mySum += sdata[tid + 32];\n // Reduce final warp using shuffle\n for (int offset = tile32.size() / 2; offset > 0; offset /= 2) {\n mySum += tile32.shfl_down(mySum, offset);\n }\n }\n\n // write result for this block to global mem\n if (cta.thread_rank() == 0)\n g_odata[blockIdx.x] = mySum;\n}\n\ntemplate \n__global__ void reduce7(const T *__restrict__ g_idata, T *__restrict__ g_odata, unsigned int n)\n{\n T *sdata = SharedMemory();\n\n // perform first level of reduction,\n // reading from global memory, writing to shared memory\n unsigned int tid = threadIdx.x;\n unsigned int gridSize = blockSize * gridDim.x;\n unsigned int maskLength = (blockSize & 31); // 31 = warpSize-1\n maskLength = (maskLength > 0) ? (32 - maskLength) : maskLength;\n const unsigned int mask = (0xffffffff) >> maskLength;\n\n T mySum = 0;\n\n // we reduce multiple elements per thread. The number is determined by the\n // number of active thread blocks (via gridDim). More blocks will result\n // in a larger gridSize and therefore fewer elements per thread\n if (nIsPow2) {\n unsigned int i = blockIdx.x * blockSize * 2 + threadIdx.x;\n gridSize = gridSize << 1;\n\n while (i < n) {\n mySum += g_idata[i];\n // ensure we don't read out of bounds -- this is optimized away for\n // powerOf2 sized arrays\n if ((i + blockSize) < n) {\n mySum += g_idata[i + blockSize];\n }\n i += gridSize;\n }\n }\n else {\n unsigned int i = blockIdx.x * blockSize + threadIdx.x;\n while (i < n) {\n mySum += g_idata[i];\n i += gridSize;\n }\n }\n\n // Reduce within warp using shuffle or reduce_add if T==int & CUDA_ARCH ==\n // SM 8.0\n mySum = warpReduceSum(mask, mySum);\n\n // each thread puts its local sum into shared memory\n if ((tid % warpSize) == 0) {\n sdata[tid / warpSize] = mySum;\n }\n\n __syncthreads();\n\n const unsigned int shmem_extent = (blockSize / warpSize) > 0 ? (blockSize / warpSize) : 1;\n const unsigned int ballot_result = __ballot_sync(mask, tid < shmem_extent);\n if (tid < shmem_extent) {\n mySum = sdata[tid];\n // Reduce final warp using shuffle or reduce_add if T==int & CUDA_ARCH ==\n // SM 8.0\n mySum = warpReduceSum(ballot_result, mySum);\n }\n\n // write result for this block to global mem\n if (tid == 0) {\n g_odata[blockIdx.x] = mySum;\n }\n}\n\n// Performs a reduction step and updates numTotal with how many are remaining\ntemplate __device__ T cg_reduce_n(T in, Group &threads)\n{\n return cg::reduce(threads, in, cg::plus());\n}\n\ntemplate __global__ void cg_reduce(T *g_idata, T *g_odata, unsigned int n)\n{\n // Shared memory for intermediate steps\n T *sdata = SharedMemory();\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Handle to tile in thread block\n cg::thread_block_tile<32> tile = cg::tiled_partition<32>(cta);\n\n unsigned int ctaSize = cta.size();\n unsigned int numCtas = gridDim.x;\n unsigned int threadRank = cta.thread_rank();\n unsigned int threadIndex = (blockIdx.x * ctaSize) + threadRank;\n\n T threadVal = 0;\n {\n unsigned int i = threadIndex;\n unsigned int indexStride = (numCtas * ctaSize);\n while (i < n) {\n threadVal += g_idata[i];\n i += indexStride;\n }\n sdata[threadRank] = threadVal;\n }\n\n // Wait for all tiles to finish and reduce within CTA\n {\n unsigned int ctaSteps = tile.meta_group_size();\n unsigned int ctaIndex = ctaSize >> 1;\n while (ctaIndex >= 32) {\n cta.sync();\n if (threadRank < ctaIndex) {\n threadVal += sdata[threadRank + ctaIndex];\n sdata[threadRank] = threadVal;\n }\n ctaSteps >>= 1;\n ctaIndex >>= 1;\n }\n }\n\n // Shuffle redux instead of smem redux\n {\n cta.sync();\n if (tile.meta_group_rank() == 0) {\n threadVal = cg_reduce_n(threadVal, tile);\n }\n }\n\n if (threadRank == 0)\n g_odata[blockIdx.x] = threadVal;\n}\n\ntemplate \n__global__ void multi_warp_cg_reduce(T *g_idata, T *g_odata, unsigned int n)\n{\n // Shared memory for intermediate steps\n T *sdata = SharedMemory();\n __shared__ cg::block_tile_memory scratch;\n\n // Handle to thread block group\n auto cta = cg::this_thread_block(scratch);\n // Handle to multiWarpTile in thread block\n auto multiWarpTile = cg::tiled_partition(cta);\n\n unsigned int gridSize = BlockSize * gridDim.x;\n T threadVal = 0;\n\n // we reduce multiple elements per thread. The number is determined by the\n // number of active thread blocks (via gridDim). More blocks will result\n // in a larger gridSize and therefore fewer elements per thread\n int nIsPow2 = !(n & n - 1);\n if (nIsPow2) {\n unsigned int i = blockIdx.x * BlockSize * 2 + threadIdx.x;\n gridSize = gridSize << 1;\n\n while (i < n) {\n threadVal += g_idata[i];\n // ensure we don't read out of bounds -- this is optimized away for\n // powerOf2 sized arrays\n if ((i + BlockSize) < n) {\n threadVal += g_idata[i + blockDim.x];\n }\n i += gridSize;\n }\n }\n else {\n unsigned int i = blockIdx.x * BlockSize + threadIdx.x;\n while (i < n) {\n threadVal += g_idata[i];\n i += gridSize;\n }\n }\n\n threadVal = cg_reduce_n(threadVal, multiWarpTile);\n\n if (multiWarpTile.thread_rank() == 0) {\n sdata[multiWarpTile.meta_group_rank()] = threadVal;\n }\n cg::sync(cta);\n\n if (threadIdx.x == 0) {\n threadVal = 0;\n for (int i = 0; i < multiWarpTile.meta_group_size(); i++) {\n threadVal += sdata[i];\n }\n g_odata[blockIdx.x] = threadVal;\n }\n}\n\nextern \"C\" bool isPow2(unsigned int x);\n\n////////////////////////////////////////////////////////////////////////////////\n// Wrapper function for kernel launch\n////////////////////////////////////////////////////////////////////////////////\ntemplate void reduce(int size, int threads, int blocks, int whichKernel, T *d_idata, T *d_odata)\n{\n dim3 dimBlock(threads, 1, 1);\n dim3 dimGrid(blocks, 1, 1);\n\n // when there is only one warp per block, we need to allocate two warps\n // worth of shared memory so that we don't index shared memory out of bounds\n int smemSize = (threads <= 32) ? 2 * threads * sizeof(T) : threads * sizeof(T);\n\n // as kernel 9 - multi_warp_cg_reduce cannot work for more than 64 threads\n // we choose to set kernel 7 for this purpose.\n if (threads < 64 && whichKernel == 9) {\n whichKernel = 7;\n }\n\n // choose which of the optimized versions of reduction to launch\n switch (whichKernel) {\n case 0:\n reduce0<<>>(d_idata, d_odata, size);\n break;\n\n case 1:\n reduce1<<>>(d_idata, d_odata, size);\n break;\n\n case 2:\n reduce2<<>>(d_idata, d_odata, size);\n break;\n\n case 3:\n reduce3<<>>(d_idata, d_odata, size);\n break;\n\n case 4:\n switch (threads) {\n case 512:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 256:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 128:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 64:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 32:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 16:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 8:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 4:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 2:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n\n case 1:\n reduce4<<>>(d_idata, d_odata, size);\n break;\n }\n\n break;\n\n case 5:\n switch (threads) {\n case 512:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 256:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 128:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 64:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 32:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 16:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 8:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 4:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 2:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n\n case 1:\n reduce5<<>>(d_idata, d_odata, size);\n break;\n }\n\n break;\n\n case 6:\n if (isPow2(size)) {\n switch (threads) {\n case 512:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 256:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 128:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 64:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 32:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 16:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 8:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 4:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 2:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 1:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n }\n }\n else {\n switch (threads) {\n case 512:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 256:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 128:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 64:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 32:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 16:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 8:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 4:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 2:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n\n case 1:\n reduce6<<>>(d_idata, d_odata, size);\n break;\n }\n }\n\n break;\n\n case 7:\n // For reduce7 kernel we require only blockSize/warpSize\n // number of elements in shared memory\n smemSize = ((threads / 32) + 1) * sizeof(T);\n if (isPow2(size)) {\n switch (threads) {\n case 1024:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n case 512:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 256:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 128:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 64:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 32:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 16:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 8:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 4:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 2:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 1:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n }\n }\n else {\n switch (threads) {\n case 1024:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n case 512:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 256:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 128:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 64:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 32:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 16:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 8:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 4:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 2:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n\n case 1:\n reduce7<<>>(d_idata, d_odata, size);\n break;\n }\n }\n\n break;\n case 8:\n cg_reduce<<>>(d_idata, d_odata, size);\n break;\n case 9:\n constexpr int numOfMultiWarpGroups = 2;\n smemSize = numOfMultiWarpGroups * sizeof(T);\n switch (threads) {\n case 1024:\n multi_warp_cg_reduce\n <<>>(d_idata, d_odata, size);\n break;\n\n case 512:\n multi_warp_cg_reduce\n <<>>(d_idata, d_odata, size);\n break;\n\n case 256:\n multi_warp_cg_reduce\n <<>>(d_idata, d_odata, size);\n break;\n\n case 128:\n multi_warp_cg_reduce\n <<>>(d_idata, d_odata, size);\n break;\n\n case 64:\n multi_warp_cg_reduce\n <<>>(d_idata, d_odata, size);\n break;\n\n default:\n printf(\"thread block size of < 64 is not supported for this kernel\\n\");\n break;\n }\n break;\n }\n}\n\n// Instantiate the reduction function for 3 types\ntemplate void reduce(int size, int threads, int blocks, int whichKernel, int *d_idata, int *d_odata);\n\ntemplate void reduce(int size, int threads, int blocks, int whichKernel, float *d_idata, float *d_odata);\n\ntemplate void reduce(int size, int threads, int blocks, int whichKernel, double *d_idata, double *d_odata);\n\n#endif // #ifndef _REDUCE_KERNEL_H_\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/reduction/reduction_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"../inc/piestimator.h\"\n\nnamespace cg = cooperative_groups;\n#include \n#include \n\n#include \"../inc/cudasharedmem.h\"\n\nusing std::string;\nusing std::vector;\n\n// Helper templates to support float and double in same code\ntemplate struct TYPE_IS\n{\n static const bool test = false;\n};\ntemplate struct TYPE_IS\n{\n static const bool test = true;\n};\ntemplate struct IF\n{\n typedef R type;\n};\ntemplate struct IF\n{\n typedef L type;\n};\n\n// RNG init kernel\ntemplate \n__global__ void\ninitRNG(rngState_t *const rngStates, rngDirectionVectors_t *const rngDirections, unsigned int numDrawsPerDirection)\n{\n // Determine thread ID\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int step = gridDim.x * blockDim.x;\n\n // Determine offset to avoid overlapping sub-sequences\n unsigned int offset = tid * ((numDrawsPerDirection + step - 1) / step);\n\n // Initialise the RNG\n curand_init(rngDirections[0], offset, &rngStates[tid]);\n curand_init(rngDirections[1], offset, &rngStates[tid + step]);\n}\n\n__device__ unsigned int reduce_sum(unsigned int in, cg::thread_block cta)\n{\n extern __shared__ unsigned int sdata[];\n\n // Perform first level of reduction:\n // - Write to shared memory\n unsigned int ltid = threadIdx.x;\n\n sdata[ltid] = in;\n cg::sync(cta);\n\n // Do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (ltid < s) {\n sdata[ltid] += sdata[ltid + s];\n }\n\n cg::sync(cta);\n }\n\n return sdata[0];\n}\n\n__device__ inline void getPoint(float &x, float &y, curandStateSobol32 &state1, curandStateSobol32 &state2)\n{\n x = curand_uniform(&state1);\n y = curand_uniform(&state2);\n}\n__device__ inline void getPoint(double &x, double &y, curandStateSobol64 &state1, curandStateSobol64 &state2)\n{\n x = curand_uniform_double(&state1);\n y = curand_uniform_double(&state2);\n}\n\n// Estimator kernel\ntemplate \n__global__ void computeValue(unsigned int *const results, rngState_t *const rngStates, const unsigned int numSims)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Determine thread ID\n unsigned int bid = blockIdx.x;\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int step = gridDim.x * blockDim.x;\n\n // Initialise the RNG\n rngState_t localState1 = rngStates[tid];\n rngState_t localState2 = rngStates[tid + step];\n\n // Count the number of points which lie inside the unit quarter-circle\n unsigned int pointsInside = 0;\n\n for (unsigned int i = tid; i < numSims; i += step) {\n Real x;\n Real y;\n getPoint(x, y, localState1, localState2);\n Real l2norm2 = x * x + y * y;\n\n if (l2norm2 < static_cast(1)) {\n pointsInside++;\n }\n }\n\n // Reduce within the block\n pointsInside = reduce_sum(pointsInside, cta);\n\n // Store the result\n if (threadIdx.x == 0) {\n results[bid] = pointsInside;\n }\n}\n\ntemplate \nPiEstimator::PiEstimator(unsigned int numSims, unsigned int device, unsigned int threadBlockSize)\n : m_numSims(numSims)\n , m_device(device)\n , m_threadBlockSize(threadBlockSize)\n{\n}\n\ntemplate Real PiEstimator::operator()()\n{\n cudaError_t cudaResult = cudaSuccess;\n struct cudaDeviceProp deviceProperties;\n struct cudaFuncAttributes funcAttributes;\n\n // Determine type of generator to use (32- or 64-bit)\n typedef\n typename IF::test, curandStateSobol64_t, curandStateSobol32_t>::type curandStateSobol_sz;\n typedef typename IF::test, curandDirectionVectors64_t, curandDirectionVectors32_t>::type\n curandDirectionVectors_sz;\n\n // Get device properties\n cudaResult = cudaGetDeviceProperties(&deviceProperties, m_device);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not get device properties: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Check precision is valid\n if (typeid(Real) == typeid(double)\n && (deviceProperties.major < 1 || (deviceProperties.major == 1 && deviceProperties.minor < 3))) {\n throw std::runtime_error(\"Device does not have double precision support\");\n }\n\n // Attach to GPU\n cudaResult = cudaSetDevice(m_device);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not set CUDA device: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Determine how to divide the work between cores\n dim3 block;\n dim3 grid;\n block.x = m_threadBlockSize;\n grid.x = (m_numSims + m_threadBlockSize - 1) / m_threadBlockSize;\n\n // Aim to launch around ten or more times as many blocks as there\n // are multiprocessors on the target device.\n unsigned int blocksPerSM = 10;\n unsigned int numSMs = deviceProperties.multiProcessorCount;\n\n while (grid.x > 2 * blocksPerSM * numSMs) {\n grid.x >>= 1;\n }\n\n // Get initRNG function properties and check the maximum block size\n cudaResult = cudaFuncGetAttributes(&funcAttributes, initRNG);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for initRNG kernel\");\n }\n\n // Get computeValue function properties and check the maximum block size\n cudaResult = cudaFuncGetAttributes(&funcAttributes, computeValue);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for computeValue kernel\");\n }\n\n // Check the dimensions are valid\n if (block.x > (unsigned int)deviceProperties.maxThreadsDim[0]) {\n throw std::runtime_error(\"Block X dimension is too large for device\");\n }\n\n if (grid.x > (unsigned int)deviceProperties.maxGridSize[0]) {\n throw std::runtime_error(\"Grid X dimension is too large for device\");\n }\n\n // Allocate memory for RNG states and direction vectors\n curandStateSobol_sz *d_rngStates = 0;\n curandDirectionVectors_sz *d_rngDirections = 0;\n cudaResult = cudaMalloc((void **)&d_rngStates, 2 * grid.x * block.x * sizeof(curandStateSobol_sz));\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not allocate memory on device for RNG states: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n cudaResult = cudaMalloc((void **)&d_rngDirections, 2 * sizeof(curandDirectionVectors_sz));\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not allocate memory on device for RNG direction vectors: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Allocate memory for result\n // Each thread block will produce one result\n unsigned int *d_results = 0;\n cudaResult = cudaMalloc((void **)&d_results, grid.x * sizeof(unsigned int));\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not allocate memory on device for partial results: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Generate direction vectors on the host and copy to the device\n if (typeid(Real) == typeid(float)) {\n curandDirectionVectors32_t *rngDirections;\n curandStatus_t curandResult = curandGetDirectionVectors32(&rngDirections, CURAND_DIRECTION_VECTORS_32_JOEKUO6);\n\n if (curandResult != CURAND_STATUS_SUCCESS) {\n string msg(\"Could not get direction vectors for quasi-random number \"\n \"generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n cudaResult =\n cudaMemcpy(d_rngDirections, rngDirections, 2 * sizeof(curandDirectionVectors32_t), cudaMemcpyHostToDevice);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not copy direction vectors to device: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n }\n else if (typeid(Real) == typeid(double)) {\n curandDirectionVectors64_t *rngDirections;\n curandStatus_t curandResult = curandGetDirectionVectors64(&rngDirections, CURAND_DIRECTION_VECTORS_64_JOEKUO6);\n\n if (curandResult != CURAND_STATUS_SUCCESS) {\n string msg(\"Could not get direction vectors for quasi-random number \"\n \"generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n cudaResult =\n cudaMemcpy(d_rngDirections, rngDirections, 2 * sizeof(curandDirectionVectors64_t), cudaMemcpyHostToDevice);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not copy direction vectors to device: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n }\n else {\n string msg(\"Could not get direction vectors for random number generator of \"\n \"specified type\");\n throw std::runtime_error(msg);\n }\n\n // Initialise RNG\n initRNG<<>>(d_rngStates, d_rngDirections, m_numSims);\n\n // Count the points inside unit quarter-circle\n computeValue<<>>(d_results, d_rngStates, m_numSims);\n\n // Copy partial results back\n vector results(grid.x);\n cudaResult = cudaMemcpy(&results[0], d_results, grid.x * sizeof(unsigned int), cudaMemcpyDeviceToHost);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not copy partial results to host: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Complete sum-reduction on host\n Real value = static_cast(std::accumulate(results.begin(), results.end(), 0));\n\n // Determine the proportion of points inside the quarter-circle,\n // i.e. the area of the unit quarter-circle\n value /= m_numSims;\n\n // Value is currently an estimate of the area of a unit quarter-circle, so we\n // can scale to a full circle by multiplying by four. Now since the area of a\n // circle is pi * r^2, and r is one, the value will be an estimate for the\n // value of pi.\n value *= 4;\n\n // Cleanup\n if (d_rngStates) {\n cudaFree(d_rngStates);\n d_rngStates = 0;\n }\n\n if (d_rngDirections) {\n cudaFree(d_rngDirections);\n d_rngDirections = 0;\n }\n\n if (d_results) {\n cudaFree(d_results);\n d_results = 0;\n }\n\n return value;\n}\n\n// Explicit template instantiation\ntemplate class PiEstimator;\ntemplate class PiEstimator;\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"../inc/piestimator.h\"\n\nnamespace cg = cooperative_groups;\n#include \n#include \n\n#include \"../inc/cudasharedmem.h\"\n\nusing std::string;\nusing std::vector;\n\n// Helper templates to support float and double in same code\ntemplate struct TYPE_IS\n{\n static const bool test = false;\n};\ntemplate struct TYPE_IS\n{\n static const bool test = true;\n};\ntemplate struct IF\n{\n typedef R type;\n};\ntemplate struct IF\n{\n typedef L type;\n};\n\n// RNG init kernel\ntemplate \n__global__ void\ninitRNG(rngState_t *const rngStates, rngDirectionVectors_t *const rngDirections, unsigned int numDrawsPerDirection)\n{\n // Determine thread ID\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int step = gridDim.x * blockDim.x;\n\n // Determine offset to avoid overlapping sub-sequences\n unsigned int offset = tid * ((numDrawsPerDirection + step - 1) / step);\n\n // Initialise the RNG\n hiprand_init(rngDirections[0], offset, &rngStates[tid]);\n hiprand_init(rngDirections[1], offset, &rngStates[tid + step]);\n}\n\n__device__ unsigned int reduce_sum(unsigned int in, cg::thread_block cta)\n{\n extern __shared__ unsigned int sdata[];\n\n // Perform first level of reduction:\n // - Write to shared memory\n unsigned int ltid = threadIdx.x;\n\n sdata[ltid] = in;\n cg::sync(cta);\n\n // Do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (ltid < s) {\n sdata[ltid] += sdata[ltid + s];\n }\n\n cg::sync(cta);\n }\n\n return sdata[0];\n}\n\n__device__ inline void getPoint(float &x, float &y, hiprandStateSobol32 &state1, hiprandStateSobol32 &state2)\n{\n x = hiprand_uniform(&state1);\n y = hiprand_uniform(&state2);\n}\n__device__ inline void getPoint(double &x, double &y, hiprandStateSobol64 &state1, hiprandStateSobol64 &state2)\n{\n x = hiprand_uniform_double(&state1);\n y = hiprand_uniform_double(&state2);\n}\n\n// Estimator kernel\ntemplate \n__global__ void computeValue(unsigned int *const results, rngState_t *const rngStates, const unsigned int numSims)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Determine thread ID\n unsigned int bid = blockIdx.x;\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int step = gridDim.x * blockDim.x;\n\n // Initialise the RNG\n rngState_t localState1 = rngStates[tid];\n rngState_t localState2 = rngStates[tid + step];\n\n // Count the number of points which lie inside the unit quarter-circle\n unsigned int pointsInside = 0;\n\n for (unsigned int i = tid; i < numSims; i += step) {\n Real x;\n Real y;\n getPoint(x, y, localState1, localState2);\n Real l2norm2 = x * x + y * y;\n\n if (l2norm2 < static_cast(1)) {\n pointsInside++;\n }\n }\n\n // Reduce within the block\n pointsInside = reduce_sum(pointsInside, cta);\n\n // Store the result\n if (threadIdx.x == 0) {\n results[bid] = pointsInside;\n }\n}\n\ntemplate \nPiEstimator::PiEstimator(unsigned int numSims, unsigned int device, unsigned int threadBlockSize)\n : m_numSims(numSims)\n , m_device(device)\n , m_threadBlockSize(threadBlockSize)\n{\n}\n\ntemplate Real PiEstimator::operator()()\n{\n hipError_t cudaResult = hipSuccess;\n struct hipDeviceProp_t deviceProperties;\n struct hipFuncAttributes funcAttributes;\n\n // Determine type of generator to use (32- or 64-bit)\n typedef\n typename IF::test, hiprandStateSobol64_t, hiprandStateSobol32_t>::type curandStateSobol_sz;\n typedef typename IF::test, hiprandDirectionVectors64_t, hiprandDirectionVectors32_t>::type\n curandDirectionVectors_sz;\n\n // Get device properties\n cudaResult = hipGetDeviceProperties(&deviceProperties, m_device);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not get device properties: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Check precision is valid\n if (typeid(Real) == typeid(double)\n && (deviceProperties.major < 1 || (deviceProperties.major == 1 && deviceProperties.minor < 3))) {\n throw std::runtime_error(\"Device does not have double precision support\");\n }\n\n // Attach to GPU\n cudaResult = hipSetDevice(m_device);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not set CUDA device: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Determine how to divide the work between cores\n dim3 block;\n dim3 grid;\n block.x = m_threadBlockSize;\n grid.x = (m_numSims + m_threadBlockSize - 1) / m_threadBlockSize;\n\n // Aim to launch around ten or more times as many blocks as there\n // are multiprocessors on the target device.\n unsigned int blocksPerSM = 10;\n unsigned int numSMs = deviceProperties.multiProcessorCount;\n\n while (grid.x > 2 * blocksPerSM * numSMs) {\n grid.x >>= 1;\n }\n\n // Get initRNG function properties and check the maximum block size\n cudaResult = hipFuncGetAttributes(&funcAttributes,reinterpret_cast(initRNG));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for initRNG kernel\");\n }\n\n // Get computeValue function properties and check the maximum block size\n cudaResult = hipFuncGetAttributes(&funcAttributes,reinterpret_cast(computeValue));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for computeValue kernel\");\n }\n\n // Check the dimensions are valid\n if (block.x > (unsigned int)deviceProperties.maxThreadsDim[0]) {\n throw std::runtime_error(\"Block X dimension is too large for device\");\n }\n\n if (grid.x > (unsigned int)deviceProperties.maxGridSize[0]) {\n throw std::runtime_error(\"Grid X dimension is too large for device\");\n }\n\n // Allocate memory for RNG states and direction vectors\n curandStateSobol_sz *d_rngStates = 0;\n curandDirectionVectors_sz *d_rngDirections = 0;\n cudaResult = hipMalloc((void **)&d_rngStates, 2 * grid.x * block.x * sizeof(curandStateSobol_sz));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not allocate memory on device for RNG states: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n cudaResult = hipMalloc((void **)&d_rngDirections, 2 * sizeof(curandDirectionVectors_sz));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not allocate memory on device for RNG direction vectors: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Allocate memory for result\n // Each thread block will produce one result\n unsigned int *d_results = 0;\n cudaResult = hipMalloc((void **)&d_results, grid.x * sizeof(unsigned int));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not allocate memory on device for partial results: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Generate direction vectors on the host and copy to the device\n if (typeid(Real) == typeid(float)) {\n hiprandDirectionVectors32_t *rngDirections;\n hiprandStatus_t curandResult = hiprandGetDirectionVectors32(&rngDirections, HIPRAND_DIRECTION_VECTORS_32_JOEKUO6);\n\n if (curandResult != HIPRAND_STATUS_SUCCESS) {\n string msg(\"Could not get direction vectors for quasi-random number \"\n \"generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n cudaResult =\n hipMemcpy(d_rngDirections, rngDirections, 2 * sizeof(hiprandDirectionVectors32_t), hipMemcpyHostToDevice);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not copy direction vectors to device: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n }\n else if (typeid(Real) == typeid(double)) {\n hiprandDirectionVectors64_t *rngDirections;\n hiprandStatus_t curandResult = hiprandGetDirectionVectors64(&rngDirections, HIPRAND_DIRECTION_VECTORS_64_JOEKUO6);\n\n if (curandResult != HIPRAND_STATUS_SUCCESS) {\n string msg(\"Could not get direction vectors for quasi-random number \"\n \"generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n cudaResult =\n hipMemcpy(d_rngDirections, rngDirections, 2 * sizeof(hiprandDirectionVectors64_t), hipMemcpyHostToDevice);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not copy direction vectors to device: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n }\n else {\n string msg(\"Could not get direction vectors for random number generator of \"\n \"specified type\");\n throw std::runtime_error(msg);\n }\n\n // Initialise RNG\n initRNG<<>>(d_rngStates, d_rngDirections, m_numSims);\n\n // Count the points inside unit quarter-circle\n computeValue<<>>(d_results, d_rngStates, m_numSims);\n\n // Copy partial results back\n vector results(grid.x);\n cudaResult = hipMemcpy(&results[0], d_results, grid.x * sizeof(unsigned int), hipMemcpyDeviceToHost);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not copy partial results to host: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Complete sum-reduction on host\n Real value = static_cast(std::accumulate(results.begin(), results.end(), 0));\n\n // Determine the proportion of points inside the quarter-circle,\n // i.e. the area of the unit quarter-circle\n value /= m_numSims;\n\n // Value is currently an estimate of the area of a unit quarter-circle, so we\n // can scale to a full circle by multiplying by four. Now since the area of a\n // circle is pi * r^2, and r is one, the value will be an estimate for the\n // value of pi.\n value *= 4;\n\n // Cleanup\n if (d_rngStates) {\n hipFree(d_rngStates);\n d_rngStates = 0;\n }\n\n if (d_rngDirections) {\n hipFree(d_rngDirections);\n d_rngDirections = 0;\n }\n\n if (d_results) {\n hipFree(d_results);\n d_results = 0;\n }\n\n return value;\n}\n\n// Explicit template instantiation\ntemplate class PiEstimator;\ntemplate class PiEstimator;\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/MC_EstimatePiInlineQ/src/piestimator.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \n\n#include \"histogram_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// GPU-specific common definitions\n////////////////////////////////////////////////////////////////////////////////\n// Data type used for input data fetches\ntypedef uint4 data_t;\n\n// May change on future hardware, so better parametrize the code\n#define SHARED_MEMORY_BANKS 16\n\n////////////////////////////////////////////////////////////////////////////////\n// Main computation pass: compute gridDim.x partial histograms\n////////////////////////////////////////////////////////////////////////////////\n// Count a byte into shared-memory storage\ninline __device__ void addByte(uchar *s_ThreadBase, uint data)\n{\n s_ThreadBase[UMUL(data, HISTOGRAM64_THREADBLOCK_SIZE)]++;\n}\n\n// Count four bytes of a word\ninline __device__ void addWord(uchar *s_ThreadBase, uint data)\n{\n // Only higher 6 bits of each byte matter, as this is a 64-bin histogram\n addByte(s_ThreadBase, (data >> 2) & 0x3FU);\n addByte(s_ThreadBase, (data >> 10) & 0x3FU);\n addByte(s_ThreadBase, (data >> 18) & 0x3FU);\n addByte(s_ThreadBase, (data >> 26) & 0x3FU);\n}\n\n__global__ void histogram64Kernel(uint *d_PartialHistograms, data_t *d_Data, uint dataCount)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Encode thread index in order to avoid bank conflicts in s_Hist[] access:\n // each group of SHARED_MEMORY_BANKS threads accesses consecutive shared\n // memory banks\n // and the same bytes [0..3] within the banks\n // Because of this permutation block size should be a multiple of 4 *\n // SHARED_MEMORY_BANKS\n const uint threadPos = ((threadIdx.x & ~(SHARED_MEMORY_BANKS * 4 - 1)) << 0)\n | ((threadIdx.x & (SHARED_MEMORY_BANKS - 1)) << 2)\n | ((threadIdx.x & (SHARED_MEMORY_BANKS * 3)) >> 4);\n\n // Per-thread histogram storage\n __shared__ uchar s_Hist[HISTOGRAM64_THREADBLOCK_SIZE * HISTOGRAM64_BIN_COUNT];\n uchar *s_ThreadBase = s_Hist + threadPos;\n\n// Initialize shared memory (writing 32-bit words)\n#pragma unroll\n\n for (uint i = 0; i < (HISTOGRAM64_BIN_COUNT / 4); i++) {\n ((uint *)s_Hist)[threadIdx.x + i * HISTOGRAM64_THREADBLOCK_SIZE] = 0;\n }\n\n // Read data from global memory and submit to the shared-memory histogram\n // Since histogram counters are byte-sized, every single thread can't do more\n // than 255 submission\n cg::sync(cta);\n\n for (uint pos = UMAD(blockIdx.x, blockDim.x, threadIdx.x); pos < dataCount; pos += UMUL(blockDim.x, gridDim.x)) {\n data_t data = d_Data[pos];\n addWord(s_ThreadBase, data.x);\n addWord(s_ThreadBase, data.y);\n addWord(s_ThreadBase, data.z);\n addWord(s_ThreadBase, data.w);\n }\n\n // Accumulate per-thread histograms into per-block and write to global memory\n cg::sync(cta);\n\n if (threadIdx.x < HISTOGRAM64_BIN_COUNT) {\n uchar *s_HistBase = s_Hist + UMUL(threadIdx.x, HISTOGRAM64_THREADBLOCK_SIZE);\n\n uint sum = 0;\n uint pos = 4 * (threadIdx.x & (SHARED_MEMORY_BANKS - 1));\n\n#pragma unroll\n\n for (uint i = 0; i < (HISTOGRAM64_THREADBLOCK_SIZE / 4); i++) {\n sum += s_HistBase[pos + 0] + s_HistBase[pos + 1] + s_HistBase[pos + 2] + s_HistBase[pos + 3];\n pos = (pos + 4) & (HISTOGRAM64_THREADBLOCK_SIZE - 1);\n }\n\n d_PartialHistograms[blockIdx.x * HISTOGRAM64_BIN_COUNT + threadIdx.x] = sum;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Merge histogram64() output\n// Run one threadblock per bin; each threadbock adds up the same bin counter\n// from every partial histogram. Reads are uncoalesced, but mergeHistogram64\n// takes only a fraction of total processing time\n////////////////////////////////////////////////////////////////////////////////\n#define MERGE_THREADBLOCK_SIZE 256\n\n__global__ void mergeHistogram64Kernel(uint *d_Histogram, uint *d_PartialHistograms, uint histogramCount)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ uint data[MERGE_THREADBLOCK_SIZE];\n\n uint sum = 0;\n\n for (uint i = threadIdx.x; i < histogramCount; i += MERGE_THREADBLOCK_SIZE) {\n sum += d_PartialHistograms[blockIdx.x + i * HISTOGRAM64_BIN_COUNT];\n }\n\n data[threadIdx.x] = sum;\n\n for (uint stride = MERGE_THREADBLOCK_SIZE / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n\n if (threadIdx.x < stride) {\n data[threadIdx.x] += data[threadIdx.x + stride];\n }\n }\n\n if (threadIdx.x == 0) {\n d_Histogram[blockIdx.x] = data[0];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// CPU interface to GPU histogram calculator\n////////////////////////////////////////////////////////////////////////////////\n// histogram64kernel() intermediate results buffer\n// MAX_PARTIAL_HISTOGRAM64_COUNT == 32768 and HISTOGRAM64_THREADBLOCK_SIZE == 64\n// amounts to max. 480MB of input data\nstatic const uint MAX_PARTIAL_HISTOGRAM64_COUNT = 32768;\nstatic uint *d_PartialHistograms;\n\n// Internal memory allocation\nextern \"C\" void initHistogram64(void)\n{\n assert(HISTOGRAM64_THREADBLOCK_SIZE % (4 * SHARED_MEMORY_BANKS) == 0);\n checkCudaErrors(cudaMalloc((void **)&d_PartialHistograms,\n MAX_PARTIAL_HISTOGRAM64_COUNT * HISTOGRAM64_BIN_COUNT * sizeof(uint)));\n}\n\n// Internal memory deallocation\nextern \"C\" void closeHistogram64(void) { checkCudaErrors(cudaFree(d_PartialHistograms)); }\n\n// Round a / b to nearest higher integer value\ninline uint iDivUp(uint a, uint b) { return (a % b != 0) ? (a / b + 1) : (a / b); }\n\n// Snap a to nearest lower multiple of b\ninline uint iSnapDown(uint a, uint b) { return a - a % b; }\n\nextern \"C\" void histogram64(uint *d_Histogram, void *d_Data, uint byteCount)\n{\n const uint histogramCount = iDivUp(byteCount, HISTOGRAM64_THREADBLOCK_SIZE * iSnapDown(255, sizeof(data_t)));\n\n assert(byteCount % sizeof(data_t) == 0);\n assert(histogramCount <= MAX_PARTIAL_HISTOGRAM64_COUNT);\n\n histogram64Kernel<<>>(\n d_PartialHistograms, (data_t *)d_Data, byteCount / sizeof(data_t));\n getLastCudaError(\"histogram64Kernel() execution failed\\n\");\n\n mergeHistogram64Kernel<<>>(\n d_Histogram, d_PartialHistograms, histogramCount);\n getLastCudaError(\"mergeHistogram64() execution failed\\n\");\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \n\n#include \"histogram_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// GPU-specific common definitions\n////////////////////////////////////////////////////////////////////////////////\n// Data type used for input data fetches\ntypedef uint4 data_t;\n\n// May change on future hardware, so better parametrize the code\n#define SHARED_MEMORY_BANKS 16\n\n////////////////////////////////////////////////////////////////////////////////\n// Main computation pass: compute gridDim.x partial histograms\n////////////////////////////////////////////////////////////////////////////////\n// Count a byte into shared-memory storage\ninline __device__ void addByte(uchar *s_ThreadBase, uint data)\n{\n s_ThreadBase[UMUL(data, HISTOGRAM64_THREADBLOCK_SIZE)]++;\n}\n\n// Count four bytes of a word\ninline __device__ void addWord(uchar *s_ThreadBase, uint data)\n{\n // Only higher 6 bits of each byte matter, as this is a 64-bin histogram\n addByte(s_ThreadBase, (data >> 2) & 0x3FU);\n addByte(s_ThreadBase, (data >> 10) & 0x3FU);\n addByte(s_ThreadBase, (data >> 18) & 0x3FU);\n addByte(s_ThreadBase, (data >> 26) & 0x3FU);\n}\n\n__global__ void histogram64Kernel(uint *d_PartialHistograms, data_t *d_Data, uint dataCount)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Encode thread index in order to avoid bank conflicts in s_Hist[] access:\n // each group of SHARED_MEMORY_BANKS threads accesses consecutive shared\n // memory banks\n // and the same bytes [0..3] within the banks\n // Because of this permutation block size should be a multiple of 4 *\n // SHARED_MEMORY_BANKS\n const uint threadPos = ((threadIdx.x & ~(SHARED_MEMORY_BANKS * 4 - 1)) << 0)\n | ((threadIdx.x & (SHARED_MEMORY_BANKS - 1)) << 2)\n | ((threadIdx.x & (SHARED_MEMORY_BANKS * 3)) >> 4);\n\n // Per-thread histogram storage\n __shared__ uchar s_Hist[HISTOGRAM64_THREADBLOCK_SIZE * HISTOGRAM64_BIN_COUNT];\n uchar *s_ThreadBase = s_Hist + threadPos;\n\n// Initialize shared memory (writing 32-bit words)\n#pragma unroll\n\n for (uint i = 0; i < (HISTOGRAM64_BIN_COUNT / 4); i++) {\n ((uint *)s_Hist)[threadIdx.x + i * HISTOGRAM64_THREADBLOCK_SIZE] = 0;\n }\n\n // Read data from global memory and submit to the shared-memory histogram\n // Since histogram counters are byte-sized, every single thread can't do more\n // than 255 submission\n cg::sync(cta);\n\n for (uint pos = UMAD(blockIdx.x, blockDim.x, threadIdx.x); pos < dataCount; pos += UMUL(blockDim.x, gridDim.x)) {\n data_t data = d_Data[pos];\n addWord(s_ThreadBase, data.x);\n addWord(s_ThreadBase, data.y);\n addWord(s_ThreadBase, data.z);\n addWord(s_ThreadBase, data.w);\n }\n\n // Accumulate per-thread histograms into per-block and write to global memory\n cg::sync(cta);\n\n if (threadIdx.x < HISTOGRAM64_BIN_COUNT) {\n uchar *s_HistBase = s_Hist + UMUL(threadIdx.x, HISTOGRAM64_THREADBLOCK_SIZE);\n\n uint sum = 0;\n uint pos = 4 * (threadIdx.x & (SHARED_MEMORY_BANKS - 1));\n\n#pragma unroll\n\n for (uint i = 0; i < (HISTOGRAM64_THREADBLOCK_SIZE / 4); i++) {\n sum += s_HistBase[pos + 0] + s_HistBase[pos + 1] + s_HistBase[pos + 2] + s_HistBase[pos + 3];\n pos = (pos + 4) & (HISTOGRAM64_THREADBLOCK_SIZE - 1);\n }\n\n d_PartialHistograms[blockIdx.x * HISTOGRAM64_BIN_COUNT + threadIdx.x] = sum;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Merge histogram64() output\n// Run one threadblock per bin; each threadbock adds up the same bin counter\n// from every partial histogram. Reads are uncoalesced, but mergeHistogram64\n// takes only a fraction of total processing time\n////////////////////////////////////////////////////////////////////////////////\n#define MERGE_THREADBLOCK_SIZE 256\n\n__global__ void mergeHistogram64Kernel(uint *d_Histogram, uint *d_PartialHistograms, uint histogramCount)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ uint data[MERGE_THREADBLOCK_SIZE];\n\n uint sum = 0;\n\n for (uint i = threadIdx.x; i < histogramCount; i += MERGE_THREADBLOCK_SIZE) {\n sum += d_PartialHistograms[blockIdx.x + i * HISTOGRAM64_BIN_COUNT];\n }\n\n data[threadIdx.x] = sum;\n\n for (uint stride = MERGE_THREADBLOCK_SIZE / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n\n if (threadIdx.x < stride) {\n data[threadIdx.x] += data[threadIdx.x + stride];\n }\n }\n\n if (threadIdx.x == 0) {\n d_Histogram[blockIdx.x] = data[0];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// CPU interface to GPU histogram calculator\n////////////////////////////////////////////////////////////////////////////////\n// histogram64kernel() intermediate results buffer\n// MAX_PARTIAL_HISTOGRAM64_COUNT == 32768 and HISTOGRAM64_THREADBLOCK_SIZE == 64\n// amounts to max. 480MB of input data\nstatic const uint MAX_PARTIAL_HISTOGRAM64_COUNT = 32768;\nstatic uint *d_PartialHistograms;\n\n// Internal memory allocation\nextern \"C\" void initHistogram64(void)\n{\n assert(HISTOGRAM64_THREADBLOCK_SIZE % (4 * SHARED_MEMORY_BANKS) == 0);\n checkCudaErrors(hipMalloc((void **)&d_PartialHistograms,\n MAX_PARTIAL_HISTOGRAM64_COUNT * HISTOGRAM64_BIN_COUNT * sizeof(uint)));\n}\n\n// Internal memory deallocation\nextern \"C\" void closeHistogram64(void) { checkCudaErrors(hipFree(d_PartialHistograms)); }\n\n// Round a / b to nearest higher integer value\ninline uint iDivUp(uint a, uint b) { return (a % b != 0) ? (a / b + 1) : (a / b); }\n\n// Snap a to nearest lower multiple of b\ninline uint iSnapDown(uint a, uint b) { return a - a % b; }\n\nextern \"C\" void histogram64(uint *d_Histogram, void *d_Data, uint byteCount)\n{\n const uint histogramCount = iDivUp(byteCount, HISTOGRAM64_THREADBLOCK_SIZE * iSnapDown(255, sizeof(data_t)));\n\n assert(byteCount % sizeof(data_t) == 0);\n assert(histogramCount <= MAX_PARTIAL_HISTOGRAM64_COUNT);\n\n histogram64Kernel<<>>(\n d_PartialHistograms, (data_t *)d_Data, byteCount / sizeof(data_t));\n getLastCudaError(\"histogram64Kernel() execution failed\\n\");\n\n mergeHistogram64Kernel<<>>(\n d_Histogram, d_PartialHistograms, histogramCount);\n getLastCudaError(\"mergeHistogram64() execution failed\\n\");\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/histogram/histogram64.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \n\n#include \"histogram_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Shortcut shared memory atomic addition functions\n////////////////////////////////////////////////////////////////////////////////\n\n#define TAG_MASK 0xFFFFFFFFU\ninline __device__ void addByte(uint *s_WarpHist, uint data, uint threadTag) { atomicAdd(s_WarpHist + data, 1); }\n\ninline __device__ void addWord(uint *s_WarpHist, uint data, uint tag)\n{\n addByte(s_WarpHist, (data >> 0) & 0xFFU, tag);\n addByte(s_WarpHist, (data >> 8) & 0xFFU, tag);\n addByte(s_WarpHist, (data >> 16) & 0xFFU, tag);\n addByte(s_WarpHist, (data >> 24) & 0xFFU, tag);\n}\n\n__global__ void histogram256Kernel(uint *d_PartialHistograms, uint *d_Data, uint dataCount)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Per-warp subhistogram storage\n __shared__ uint s_Hist[HISTOGRAM256_THREADBLOCK_MEMORY];\n uint *s_WarpHist = s_Hist + (threadIdx.x >> LOG2_WARP_SIZE) * HISTOGRAM256_BIN_COUNT;\n\n// Clear shared memory storage for current threadblock before processing\n#pragma unroll\n\n for (uint i = 0; i < (HISTOGRAM256_THREADBLOCK_MEMORY / HISTOGRAM256_THREADBLOCK_SIZE); i++) {\n s_Hist[threadIdx.x + i * HISTOGRAM256_THREADBLOCK_SIZE] = 0;\n }\n\n // Cycle through the entire data set, update subhistograms for each warp\n const uint tag = threadIdx.x << (UINT_BITS - LOG2_WARP_SIZE);\n\n cg::sync(cta);\n\n for (uint pos = UMAD(blockIdx.x, blockDim.x, threadIdx.x); pos < dataCount; pos += UMUL(blockDim.x, gridDim.x)) {\n uint data = d_Data[pos];\n addWord(s_WarpHist, data, tag);\n }\n\n // Merge per-warp histograms into per-block and write to global memory\n cg::sync(cta);\n\n for (uint bin = threadIdx.x; bin < HISTOGRAM256_BIN_COUNT; bin += HISTOGRAM256_THREADBLOCK_SIZE) {\n uint sum = 0;\n\n for (uint i = 0; i < WARP_COUNT; i++) {\n sum += s_Hist[bin + i * HISTOGRAM256_BIN_COUNT] & TAG_MASK;\n }\n\n d_PartialHistograms[blockIdx.x * HISTOGRAM256_BIN_COUNT + bin] = sum;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Merge histogram256() output\n// Run one threadblock per bin; each threadblock adds up the same bin counter\n// from every partial histogram. Reads are uncoalesced, but mergeHistogram256\n// takes only a fraction of total processing time\n////////////////////////////////////////////////////////////////////////////////\n#define MERGE_THREADBLOCK_SIZE 256\n\n__global__ void mergeHistogram256Kernel(uint *d_Histogram, uint *d_PartialHistograms, uint histogramCount)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n\n uint sum = 0;\n\n for (uint i = threadIdx.x; i < histogramCount; i += MERGE_THREADBLOCK_SIZE) {\n sum += d_PartialHistograms[blockIdx.x + i * HISTOGRAM256_BIN_COUNT];\n }\n\n __shared__ uint data[MERGE_THREADBLOCK_SIZE];\n data[threadIdx.x] = sum;\n\n for (uint stride = MERGE_THREADBLOCK_SIZE / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n\n if (threadIdx.x < stride) {\n data[threadIdx.x] += data[threadIdx.x + stride];\n }\n }\n\n if (threadIdx.x == 0) {\n d_Histogram[blockIdx.x] = data[0];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Host interface to GPU histogram\n////////////////////////////////////////////////////////////////////////////////\n// histogram256kernel() intermediate results buffer\nstatic const uint PARTIAL_HISTOGRAM256_COUNT = 240;\nstatic uint *d_PartialHistograms;\n\n// Internal memory allocation\nextern \"C\" void initHistogram256(void)\n{\n checkCudaErrors(\n cudaMalloc((void **)&d_PartialHistograms, PARTIAL_HISTOGRAM256_COUNT * HISTOGRAM256_BIN_COUNT * sizeof(uint)));\n}\n\n// Internal memory deallocation\nextern \"C\" void closeHistogram256(void) { checkCudaErrors(cudaFree(d_PartialHistograms)); }\n\nextern \"C\" void histogram256(uint *d_Histogram, void *d_Data, uint byteCount)\n{\n assert(byteCount % sizeof(uint) == 0);\n histogram256Kernel<<>>(\n d_PartialHistograms, (uint *)d_Data, byteCount / sizeof(uint));\n getLastCudaError(\"histogram256Kernel() execution failed\\n\");\n\n mergeHistogram256Kernel<<>>(\n d_Histogram, d_PartialHistograms, PARTIAL_HISTOGRAM256_COUNT);\n getLastCudaError(\"mergeHistogram256Kernel() execution failed\\n\");\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \n\n#include \"histogram_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Shortcut shared memory atomic addition functions\n////////////////////////////////////////////////////////////////////////////////\n\n#define TAG_MASK 0xFFFFFFFFU\ninline __device__ void addByte(uint *s_WarpHist, uint data, uint threadTag) { atomicAdd(s_WarpHist + data, 1); }\n\ninline __device__ void addWord(uint *s_WarpHist, uint data, uint tag)\n{\n addByte(s_WarpHist, (data >> 0) & 0xFFU, tag);\n addByte(s_WarpHist, (data >> 8) & 0xFFU, tag);\n addByte(s_WarpHist, (data >> 16) & 0xFFU, tag);\n addByte(s_WarpHist, (data >> 24) & 0xFFU, tag);\n}\n\n__global__ void histogram256Kernel(uint *d_PartialHistograms, uint *d_Data, uint dataCount)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Per-warp subhistogram storage\n __shared__ uint s_Hist[HISTOGRAM256_THREADBLOCK_MEMORY];\n uint *s_WarpHist = s_Hist + (threadIdx.x >> LOG2_WARP_SIZE) * HISTOGRAM256_BIN_COUNT;\n\n// Clear shared memory storage for current threadblock before processing\n#pragma unroll\n\n for (uint i = 0; i < (HISTOGRAM256_THREADBLOCK_MEMORY / HISTOGRAM256_THREADBLOCK_SIZE); i++) {\n s_Hist[threadIdx.x + i * HISTOGRAM256_THREADBLOCK_SIZE] = 0;\n }\n\n // Cycle through the entire data set, update subhistograms for each warp\n const uint tag = threadIdx.x << (UINT_BITS - LOG2_WARP_SIZE);\n\n cg::sync(cta);\n\n for (uint pos = UMAD(blockIdx.x, blockDim.x, threadIdx.x); pos < dataCount; pos += UMUL(blockDim.x, gridDim.x)) {\n uint data = d_Data[pos];\n addWord(s_WarpHist, data, tag);\n }\n\n // Merge per-warp histograms into per-block and write to global memory\n cg::sync(cta);\n\n for (uint bin = threadIdx.x; bin < HISTOGRAM256_BIN_COUNT; bin += HISTOGRAM256_THREADBLOCK_SIZE) {\n uint sum = 0;\n\n for (uint i = 0; i < WARP_COUNT; i++) {\n sum += s_Hist[bin + i * HISTOGRAM256_BIN_COUNT] & TAG_MASK;\n }\n\n d_PartialHistograms[blockIdx.x * HISTOGRAM256_BIN_COUNT + bin] = sum;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Merge histogram256() output\n// Run one threadblock per bin; each threadblock adds up the same bin counter\n// from every partial histogram. Reads are uncoalesced, but mergeHistogram256\n// takes only a fraction of total processing time\n////////////////////////////////////////////////////////////////////////////////\n#define MERGE_THREADBLOCK_SIZE 256\n\n__global__ void mergeHistogram256Kernel(uint *d_Histogram, uint *d_PartialHistograms, uint histogramCount)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n\n uint sum = 0;\n\n for (uint i = threadIdx.x; i < histogramCount; i += MERGE_THREADBLOCK_SIZE) {\n sum += d_PartialHistograms[blockIdx.x + i * HISTOGRAM256_BIN_COUNT];\n }\n\n __shared__ uint data[MERGE_THREADBLOCK_SIZE];\n data[threadIdx.x] = sum;\n\n for (uint stride = MERGE_THREADBLOCK_SIZE / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n\n if (threadIdx.x < stride) {\n data[threadIdx.x] += data[threadIdx.x + stride];\n }\n }\n\n if (threadIdx.x == 0) {\n d_Histogram[blockIdx.x] = data[0];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Host interface to GPU histogram\n////////////////////////////////////////////////////////////////////////////////\n// histogram256kernel() intermediate results buffer\nstatic const uint PARTIAL_HISTOGRAM256_COUNT = 240;\nstatic uint *d_PartialHistograms;\n\n// Internal memory allocation\nextern \"C\" void initHistogram256(void)\n{\n checkCudaErrors(\n hipMalloc((void **)&d_PartialHistograms, PARTIAL_HISTOGRAM256_COUNT * HISTOGRAM256_BIN_COUNT * sizeof(uint)));\n}\n\n// Internal memory deallocation\nextern \"C\" void closeHistogram256(void) { checkCudaErrors(hipFree(d_PartialHistograms)); }\n\nextern \"C\" void histogram256(uint *d_Histogram, void *d_Data, uint byteCount)\n{\n assert(byteCount % sizeof(uint) == 0);\n histogram256Kernel<<>>(\n d_PartialHistograms, (uint *)d_Data, byteCount / sizeof(uint));\n getLastCudaError(\"histogram256Kernel() execution failed\\n\");\n\n mergeHistogram256Kernel<<>>(\n d_Histogram, d_PartialHistograms, PARTIAL_HISTOGRAM256_COUNT);\n getLastCudaError(\"mergeHistogram256Kernel() execution failed\\n\");\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/histogram/histogram256.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"../inc/piestimator.h\"\n\nnamespace cg = cooperative_groups;\n#include \n\nusing std::string;\nusing std::vector;\n\n__device__ unsigned int reduce_sum(unsigned int in, cg::thread_block cta)\n{\n extern __shared__ unsigned int sdata[];\n\n // Perform first level of reduction:\n // - Write to shared memory\n unsigned int ltid = threadIdx.x;\n\n sdata[ltid] = in;\n cg::sync(cta);\n\n // Do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (ltid < s) {\n sdata[ltid] += sdata[ltid + s];\n }\n\n cg::sync(cta);\n }\n\n return sdata[0];\n}\n\n// Estimator kernel\ntemplate \n__global__ void computeValue(unsigned int *const results, const Real *const points, const unsigned int numSims)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Determine thread ID\n unsigned int bid = blockIdx.x;\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int step = gridDim.x * blockDim.x;\n\n // Shift the input/output pointers\n const Real *pointx = points + tid;\n const Real *pointy = pointx + numSims;\n\n // Count the number of points which lie inside the unit quarter-circle\n unsigned int pointsInside = 0;\n\n for (unsigned int i = tid; i < numSims; i += step, pointx += step, pointy += step) {\n Real x = *pointx;\n Real y = *pointy;\n Real l2norm2 = x * x + y * y;\n\n if (l2norm2 < static_cast(1)) {\n pointsInside++;\n }\n }\n\n // Reduce within the block\n pointsInside = reduce_sum(pointsInside, cta);\n\n // Store the result\n if (threadIdx.x == 0) {\n results[bid] = pointsInside;\n }\n}\n\ntemplate \nPiEstimator::PiEstimator(unsigned int numSims, unsigned int device, unsigned int threadBlockSize)\n : m_numSims(numSims)\n , m_device(device)\n , m_threadBlockSize(threadBlockSize)\n{\n}\n\ntemplate Real PiEstimator::operator()()\n{\n cudaError_t cudaResult = cudaSuccess;\n struct cudaDeviceProp deviceProperties;\n struct cudaFuncAttributes funcAttributes;\n\n // Get device properties\n cudaResult = cudaGetDeviceProperties(&deviceProperties, m_device);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not get device properties: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Check precision is valid\n if (typeid(Real) == typeid(double)\n && (deviceProperties.major < 1 || (deviceProperties.major == 1 && deviceProperties.minor < 3))) {\n throw std::runtime_error(\"Device does not have double precision support\");\n }\n\n // Attach to GPU\n cudaResult = cudaSetDevice(m_device);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not set CUDA device: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Determine how to divide the work between cores\n dim3 block;\n dim3 grid;\n block.x = m_threadBlockSize;\n grid.x = (m_numSims + m_threadBlockSize - 1) / m_threadBlockSize;\n\n // Aim to launch around ten or more times as many blocks as there\n // are multiprocessors on the target device.\n unsigned int blocksPerSM = 10;\n unsigned int numSMs = deviceProperties.multiProcessorCount;\n\n while (grid.x > 2 * blocksPerSM * numSMs) {\n grid.x >>= 1;\n }\n\n // Get computeValue function properties and check the maximum block size\n cudaResult = cudaFuncGetAttributes(&funcAttributes, computeValue);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for computeValue kernel\");\n }\n\n // Check the dimensions are valid\n if (block.x > (unsigned int)deviceProperties.maxThreadsDim[0]) {\n throw std::runtime_error(\"Block X dimension is too large for device\");\n }\n\n if (grid.x > (unsigned int)deviceProperties.maxGridSize[0]) {\n throw std::runtime_error(\"Grid X dimension is too large for device\");\n }\n\n // Allocate memory for points\n // Each simulation has two random numbers to give X and Y coordinate\n Real *d_points = 0;\n cudaResult = cudaMalloc((void **)&d_points, 2 * m_numSims * sizeof(Real));\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not allocate memory on device for random numbers: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Allocate memory for result\n // Each thread block will produce one result\n unsigned int *d_results = 0;\n cudaResult = cudaMalloc((void **)&d_results, grid.x * sizeof(unsigned int));\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not allocate memory on device for partial results: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Generate random points in unit square\n curandStatus_t curandResult;\n curandGenerator_t qrng;\n\n if (typeid(Real) == typeid(float)) {\n curandResult = curandCreateGenerator(&qrng, CURAND_RNG_QUASI_SOBOL32);\n }\n else if (typeid(Real) == typeid(double)) {\n curandResult = curandCreateGenerator(&qrng, CURAND_RNG_QUASI_SOBOL64);\n }\n else {\n string msg(\"Could not create random number generator of specified type\");\n throw std::runtime_error(msg);\n }\n\n if (curandResult != CURAND_STATUS_SUCCESS) {\n string msg(\"Could not create quasi-random number generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n curandResult = curandSetQuasiRandomGeneratorDimensions(qrng, 2);\n\n if (curandResult != CURAND_STATUS_SUCCESS) {\n string msg(\"Could not set number of dimensions for quasi-random number \"\n \"generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n curandResult = curandSetGeneratorOrdering(qrng, CURAND_ORDERING_QUASI_DEFAULT);\n\n if (curandResult != CURAND_STATUS_SUCCESS) {\n string msg(\"Could not set order for quasi-random number generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n if (typeid(Real) == typeid(float)) {\n curandResult = curandGenerateUniform(qrng, (float *)d_points, 2 * m_numSims);\n }\n else if (typeid(Real) == typeid(double)) {\n curandResult = curandGenerateUniformDouble(qrng, (double *)d_points, 2 * m_numSims);\n }\n else {\n string msg(\"Could not generate random numbers of specified type\");\n throw std::runtime_error(msg);\n }\n\n if (curandResult != CURAND_STATUS_SUCCESS) {\n string msg(\"Could not generate quasi-random numbers: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n curandResult = curandDestroyGenerator(qrng);\n\n if (curandResult != CURAND_STATUS_SUCCESS) {\n string msg(\"Could not destroy quasi-random number generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n // Count the points inside unit quarter-circle\n computeValue<<>>(d_results, d_points, m_numSims);\n\n // Copy partial results back\n vector results(grid.x);\n cudaResult = cudaMemcpy(&results[0], d_results, grid.x * sizeof(unsigned int), cudaMemcpyDeviceToHost);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not copy partial results to host: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Complete sum-reduction on host\n Real value = static_cast(std::accumulate(results.begin(), results.end(), 0));\n\n // Determine the proportion of points inside the quarter-circle,\n // i.e. the area of the unit quarter-circle\n value /= m_numSims;\n\n // Value is currently an estimate of the area of a unit quarter-circle, so we\n // can scale to a full circle by multiplying by four. Now since the area of a\n // circle is pi * r^2, and r is one, the value will be an estimate for the\n // value of pi.\n value *= 4;\n\n // Cleanup\n if (d_points) {\n cudaFree(d_points);\n d_points = 0;\n }\n\n if (d_results) {\n cudaFree(d_results);\n d_results = 0;\n }\n\n return value;\n}\n\n// Explicit template instantiation\ntemplate class PiEstimator;\ntemplate class PiEstimator;\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"../inc/piestimator.h\"\n\nnamespace cg = cooperative_groups;\n#include \n\nusing std::string;\nusing std::vector;\n\n__device__ unsigned int reduce_sum(unsigned int in, cg::thread_block cta)\n{\n extern __shared__ unsigned int sdata[];\n\n // Perform first level of reduction:\n // - Write to shared memory\n unsigned int ltid = threadIdx.x;\n\n sdata[ltid] = in;\n cg::sync(cta);\n\n // Do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (ltid < s) {\n sdata[ltid] += sdata[ltid + s];\n }\n\n cg::sync(cta);\n }\n\n return sdata[0];\n}\n\n// Estimator kernel\ntemplate \n__global__ void computeValue(unsigned int *const results, const Real *const points, const unsigned int numSims)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Determine thread ID\n unsigned int bid = blockIdx.x;\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int step = gridDim.x * blockDim.x;\n\n // Shift the input/output pointers\n const Real *pointx = points + tid;\n const Real *pointy = pointx + numSims;\n\n // Count the number of points which lie inside the unit quarter-circle\n unsigned int pointsInside = 0;\n\n for (unsigned int i = tid; i < numSims; i += step, pointx += step, pointy += step) {\n Real x = *pointx;\n Real y = *pointy;\n Real l2norm2 = x * x + y * y;\n\n if (l2norm2 < static_cast(1)) {\n pointsInside++;\n }\n }\n\n // Reduce within the block\n pointsInside = reduce_sum(pointsInside, cta);\n\n // Store the result\n if (threadIdx.x == 0) {\n results[bid] = pointsInside;\n }\n}\n\ntemplate \nPiEstimator::PiEstimator(unsigned int numSims, unsigned int device, unsigned int threadBlockSize)\n : m_numSims(numSims)\n , m_device(device)\n , m_threadBlockSize(threadBlockSize)\n{\n}\n\ntemplate Real PiEstimator::operator()()\n{\n hipError_t cudaResult = hipSuccess;\n struct hipDeviceProp_t deviceProperties;\n struct hipFuncAttributes funcAttributes;\n\n // Get device properties\n cudaResult = hipGetDeviceProperties(&deviceProperties, m_device);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not get device properties: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Check precision is valid\n if (typeid(Real) == typeid(double)\n && (deviceProperties.major < 1 || (deviceProperties.major == 1 && deviceProperties.minor < 3))) {\n throw std::runtime_error(\"Device does not have double precision support\");\n }\n\n // Attach to GPU\n cudaResult = hipSetDevice(m_device);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not set CUDA device: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Determine how to divide the work between cores\n dim3 block;\n dim3 grid;\n block.x = m_threadBlockSize;\n grid.x = (m_numSims + m_threadBlockSize - 1) / m_threadBlockSize;\n\n // Aim to launch around ten or more times as many blocks as there\n // are multiprocessors on the target device.\n unsigned int blocksPerSM = 10;\n unsigned int numSMs = deviceProperties.multiProcessorCount;\n\n while (grid.x > 2 * blocksPerSM * numSMs) {\n grid.x >>= 1;\n }\n\n // Get computeValue function properties and check the maximum block size\n cudaResult = hipFuncGetAttributes(&funcAttributes,reinterpret_cast(computeValue));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for computeValue kernel\");\n }\n\n // Check the dimensions are valid\n if (block.x > (unsigned int)deviceProperties.maxThreadsDim[0]) {\n throw std::runtime_error(\"Block X dimension is too large for device\");\n }\n\n if (grid.x > (unsigned int)deviceProperties.maxGridSize[0]) {\n throw std::runtime_error(\"Grid X dimension is too large for device\");\n }\n\n // Allocate memory for points\n // Each simulation has two random numbers to give X and Y coordinate\n Real *d_points = 0;\n cudaResult = hipMalloc((void **)&d_points, 2 * m_numSims * sizeof(Real));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not allocate memory on device for random numbers: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Allocate memory for result\n // Each thread block will produce one result\n unsigned int *d_results = 0;\n cudaResult = hipMalloc((void **)&d_results, grid.x * sizeof(unsigned int));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not allocate memory on device for partial results: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Generate random points in unit square\n hiprandStatus_t curandResult;\n hiprandGenerator_t qrng;\n\n if (typeid(Real) == typeid(float)) {\n curandResult = hiprandCreateGenerator(&qrng, HIPRAND_RNG_QUASI_SOBOL32);\n }\n else if (typeid(Real) == typeid(double)) {\n curandResult = hiprandCreateGenerator(&qrng, HIPRAND_RNG_QUASI_SOBOL64);\n }\n else {\n string msg(\"Could not create random number generator of specified type\");\n throw std::runtime_error(msg);\n }\n\n if (curandResult != HIPRAND_STATUS_SUCCESS) {\n string msg(\"Could not create quasi-random number generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n curandResult = hiprandSetQuasiRandomGeneratorDimensions(qrng, 2);\n\n if (curandResult != HIPRAND_STATUS_SUCCESS) {\n string msg(\"Could not set number of dimensions for quasi-random number \"\n \"generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n curandResult = hiprandSetGeneratorOrdering(qrng, HIPRAND_ORDERING_QUASI_DEFAULT);\n\n if (curandResult != HIPRAND_STATUS_SUCCESS) {\n string msg(\"Could not set order for quasi-random number generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n if (typeid(Real) == typeid(float)) {\n curandResult = hiprandGenerateUniform(qrng, (float *)d_points, 2 * m_numSims);\n }\n else if (typeid(Real) == typeid(double)) {\n curandResult = hiprandGenerateUniformDouble(qrng, (double *)d_points, 2 * m_numSims);\n }\n else {\n string msg(\"Could not generate random numbers of specified type\");\n throw std::runtime_error(msg);\n }\n\n if (curandResult != HIPRAND_STATUS_SUCCESS) {\n string msg(\"Could not generate quasi-random numbers: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n curandResult = hiprandDestroyGenerator(qrng);\n\n if (curandResult != HIPRAND_STATUS_SUCCESS) {\n string msg(\"Could not destroy quasi-random number generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n // Count the points inside unit quarter-circle\n computeValue<<>>(d_results, d_points, m_numSims);\n\n // Copy partial results back\n vector results(grid.x);\n cudaResult = hipMemcpy(&results[0], d_results, grid.x * sizeof(unsigned int), hipMemcpyDeviceToHost);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not copy partial results to host: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Complete sum-reduction on host\n Real value = static_cast(std::accumulate(results.begin(), results.end(), 0));\n\n // Determine the proportion of points inside the quarter-circle,\n // i.e. the area of the unit quarter-circle\n value /= m_numSims;\n\n // Value is currently an estimate of the area of a unit quarter-circle, so we\n // can scale to a full circle by multiplying by four. Now since the area of a\n // circle is pi * r^2, and r is one, the value will be an estimate for the\n // value of pi.\n value *= 4;\n\n // Cleanup\n if (d_points) {\n hipFree(d_points);\n d_points = 0;\n }\n\n if (d_results) {\n hipFree(d_results);\n d_results = 0;\n }\n\n return value;\n}\n\n// Explicit template instantiation\ntemplate class PiEstimator;\ntemplate class PiEstimator;\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/MC_EstimatePiQ/src/piestimator.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates peer-to-peer access of stream ordered memory\n * allocated with cudaMallocAsync and cudaMemPool family of APIs through simple\n * kernel which does peer-to-peer to access & scales vector elements.\n */\n\n// System includes\n#include \n#include \n#include \n#include \n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Simple kernel to demonstrate copying cudaMallocAsync memory via P2P to peer\n// device\n__global__ void copyP2PAndScale(const int *src, int *dst, int N)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx < N) {\n // scale & store src vector.\n dst[idx] = 2 * src[idx];\n }\n}\n\n// Map of device version to device number\nstd::multimap, int> getIdenticalGPUs()\n{\n int numGpus = 0;\n checkCudaErrors(cudaGetDeviceCount(&numGpus));\n\n std::multimap, int> identicalGpus;\n\n for (int i = 0; i < numGpus; i++) {\n int isMemPoolSupported = 0;\n checkCudaErrors(cudaDeviceGetAttribute(&isMemPoolSupported, cudaDevAttrMemoryPoolsSupported, i));\n\n // Filter unsupported devices\n if (isMemPoolSupported) {\n int major = 0, minor = 0;\n checkCudaErrors(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, i));\n checkCudaErrors(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, i));\n identicalGpus.emplace(std::make_pair(major, minor), i);\n }\n }\n\n return identicalGpus;\n}\n\nstd::pair getP2PCapableGpuPair()\n{\n constexpr size_t kNumGpusRequired = 2;\n\n auto gpusByArch = getIdenticalGPUs();\n\n auto it = gpusByArch.begin();\n auto end = gpusByArch.end();\n\n auto bestFit = std::make_pair(it, it);\n // use std::distance to find the largest number of GPUs amongst architectures\n auto distance = [](decltype(bestFit) p) { return std::distance(p.first, p.second); };\n\n // Read each unique key/pair element in order\n for (; it != end; it = gpusByArch.upper_bound(it->first)) {\n // first and second are iterators bounded within the architecture group\n auto testFit = gpusByArch.equal_range(it->first);\n // Always use devices with highest architecture version or whichever has the\n // most devices available\n if (distance(bestFit) <= distance(testFit))\n bestFit = testFit;\n }\n\n if (distance(bestFit) < kNumGpusRequired) {\n printf(\"No Two or more GPUs with same architecture capable of cuda Memory \"\n \"Pools found.\"\n \"\\nWaiving the sample\\n\");\n exit(EXIT_WAIVED);\n }\n\n std::set bestFitDeviceIds;\n\n // check & select peer-to-peer access capable GPU devices.\n int devIds[2];\n for (auto itr = bestFit.first; itr != bestFit.second; itr++) {\n int deviceId = itr->second;\n checkCudaErrors(cudaSetDevice(deviceId));\n\n std::for_each(itr, bestFit.second, [&deviceId, &bestFitDeviceIds, &kNumGpusRequired](decltype(*itr) mapPair) {\n if (deviceId != mapPair.second) {\n int access = 0;\n checkCudaErrors(cudaDeviceCanAccessPeer(&access, deviceId, mapPair.second));\n printf(\"Device=%d %s Access Peer Device=%d\\n\", deviceId, access ? \"CAN\" : \"CANNOT\", mapPair.second);\n if (access && bestFitDeviceIds.size() < kNumGpusRequired) {\n bestFitDeviceIds.emplace(deviceId);\n bestFitDeviceIds.emplace(mapPair.second);\n }\n else {\n printf(\"Ignoring device %i (max devices exceeded)\\n\", mapPair.second);\n }\n }\n });\n\n if (bestFitDeviceIds.size() >= kNumGpusRequired) {\n printf(\"Selected p2p capable devices - \");\n int i = 0;\n for (auto devicesItr = bestFitDeviceIds.begin(); devicesItr != bestFitDeviceIds.end(); devicesItr++) {\n devIds[i++] = *devicesItr;\n printf(\"deviceId = %d \", *devicesItr);\n }\n printf(\"\\n\");\n break;\n }\n }\n\n // if bestFitDeviceIds.size() == 0 it means the GPUs in system are not p2p\n // capable, hence we add it without p2p capability check.\n if (!bestFitDeviceIds.size()) {\n printf(\"No Two or more Devices p2p capable found.. exiting..\\n\");\n exit(EXIT_WAIVED);\n }\n\n auto p2pGpuPair = std::make_pair(devIds[0], devIds[1]);\n\n return p2pGpuPair;\n}\n\nint memPoolP2PCopy()\n{\n int *dev0_srcVec, *dev1_dstVec; // Device buffers\n cudaStream_t stream1, stream2;\n cudaMemPool_t memPool;\n cudaEvent_t waitOnStream1;\n\n // Allocate CPU memory.\n size_t nelem = 1048576;\n size_t bytes = nelem * sizeof(int);\n\n int *a = (int *)malloc(bytes);\n int *output = (int *)malloc(bytes);\n\n /* Initialize the vectors. */\n for (int n = 0; n < nelem; n++) {\n a[n] = rand() / (int)RAND_MAX;\n }\n\n auto p2pDevices = getP2PCapableGpuPair();\n printf(\"selected devices = %d & %d\\n\", p2pDevices.first, p2pDevices.second);\n checkCudaErrors(cudaSetDevice(p2pDevices.first));\n checkCudaErrors(cudaEventCreate(&waitOnStream1));\n\n checkCudaErrors(cudaStreamCreateWithFlags(&stream1, cudaStreamNonBlocking));\n\n // Get the default mempool for device p2pDevices.first from the pair\n checkCudaErrors(cudaDeviceGetDefaultMemPool(&memPool, p2pDevices.first));\n\n // Allocate memory in a stream from the pool set above.\n checkCudaErrors(cudaMallocAsync(&dev0_srcVec, bytes, stream1));\n\n checkCudaErrors(cudaMemcpyAsync(dev0_srcVec, a, bytes, cudaMemcpyHostToDevice, stream1));\n checkCudaErrors(cudaEventRecord(waitOnStream1, stream1));\n\n checkCudaErrors(cudaSetDevice(p2pDevices.second));\n checkCudaErrors(cudaStreamCreateWithFlags(&stream2, cudaStreamNonBlocking));\n\n // Allocate memory in p2pDevices.second device\n checkCudaErrors(cudaMallocAsync(&dev1_dstVec, bytes, stream2));\n\n // Setup peer mappings for p2pDevices.second device\n cudaMemAccessDesc desc;\n memset(&desc, 0, sizeof(cudaMemAccessDesc));\n desc.location.type = cudaMemLocationTypeDevice;\n desc.location.id = p2pDevices.second;\n desc.flags = cudaMemAccessFlagsProtReadWrite;\n checkCudaErrors(cudaMemPoolSetAccess(memPool, &desc, 1));\n\n printf(\"> copyP2PAndScale kernel running ...\\n\");\n dim3 block(256);\n dim3 grid((unsigned int)ceil(nelem / (int)block.x));\n checkCudaErrors(cudaStreamWaitEvent(stream2, waitOnStream1));\n copyP2PAndScale<<>>(dev0_srcVec, dev1_dstVec, nelem);\n\n checkCudaErrors(cudaMemcpyAsync(output, dev1_dstVec, bytes, cudaMemcpyDeviceToHost, stream2));\n checkCudaErrors(cudaFreeAsync(dev0_srcVec, stream2));\n checkCudaErrors(cudaFreeAsync(dev1_dstVec, stream2));\n checkCudaErrors(cudaStreamSynchronize(stream2));\n\n /* Compare the results */\n printf(\"> Checking the results from copyP2PAndScale() ...\\n\");\n\n for (int n = 0; n < nelem; n++) {\n if ((2 * a[n]) != output[n]) {\n printf(\"mismatch i = %d expected = %d val = %d\\n\", n, 2 * a[n], output[n]);\n return EXIT_FAILURE;\n }\n }\n\n free(a);\n free(output);\n checkCudaErrors(cudaStreamDestroy(stream1));\n checkCudaErrors(cudaStreamDestroy(stream2));\n printf(\"PASSED\\n\");\n\n return EXIT_SUCCESS;\n}\n\nint main(int argc, char **argv)\n{\n int ret = memPoolP2PCopy();\n return ret;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates peer-to-peer access of stream ordered memory\n * allocated with cudaMallocAsync and cudaMemPool family of APIs through simple\n * kernel which does peer-to-peer to access & scales vector elements.\n */\n\n// System includes\n#include \n#include \n#include \n#include \n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n// Simple kernel to demonstrate copying cudaMallocAsync memory via P2P to peer\n// device\n__global__ void copyP2PAndScale(const int *src, int *dst, int N)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx < N) {\n // scale & store src vector.\n dst[idx] = 2 * src[idx];\n }\n}\n\n// Map of device version to device number\nstd::multimap, int> getIdenticalGPUs()\n{\n int numGpus = 0;\n checkCudaErrors(hipGetDeviceCount(&numGpus));\n\n std::multimap, int> identicalGpus;\n\n for (int i = 0; i < numGpus; i++) {\n int isMemPoolSupported = 0;\n checkCudaErrors(hipDeviceGetAttribute(&isMemPoolSupported, hipDeviceAttributeMemoryPoolsSupported, i));\n\n // Filter unsupported devices\n if (isMemPoolSupported) {\n int major = 0, minor = 0;\n checkCudaErrors(hipDeviceGetAttribute(&major, hipDeviceAttributeComputeCapabilityMajor, i));\n checkCudaErrors(hipDeviceGetAttribute(&minor, hipDeviceAttributeComputeCapabilityMinor, i));\n identicalGpus.emplace(std::make_pair(major, minor), i);\n }\n }\n\n return identicalGpus;\n}\n\nstd::pair getP2PCapableGpuPair()\n{\n constexpr size_t kNumGpusRequired = 2;\n\n auto gpusByArch = getIdenticalGPUs();\n\n auto it = gpusByArch.begin();\n auto end = gpusByArch.end();\n\n auto bestFit = std::make_pair(it, it);\n // use std::distance to find the largest number of GPUs amongst architectures\n auto distance = [](decltype(bestFit) p) { return std::distance(p.first, p.second); };\n\n // Read each unique key/pair element in order\n for (; it != end; it = gpusByArch.upper_bound(it->first)) {\n // first and second are iterators bounded within the architecture group\n auto testFit = gpusByArch.equal_range(it->first);\n // Always use devices with highest architecture version or whichever has the\n // most devices available\n if (distance(bestFit) <= distance(testFit))\n bestFit = testFit;\n }\n\n if (distance(bestFit) < kNumGpusRequired) {\n printf(\"No Two or more GPUs with same architecture capable of cuda Memory \"\n \"Pools found.\"\n \"\\nWaiving the sample\\n\");\n exit(EXIT_WAIVED);\n }\n\n std::set bestFitDeviceIds;\n\n // check & select peer-to-peer access capable GPU devices.\n int devIds[2];\n for (auto itr = bestFit.first; itr != bestFit.second; itr++) {\n int deviceId = itr->second;\n checkCudaErrors(hipSetDevice(deviceId));\n\n std::for_each(itr, bestFit.second, [&deviceId, &bestFitDeviceIds, &kNumGpusRequired](decltype(*itr) mapPair) {\n if (deviceId != mapPair.second) {\n int access = 0;\n checkCudaErrors(hipDeviceCanAccessPeer(&access, deviceId, mapPair.second));\n printf(\"Device=%d %s Access Peer Device=%d\\n\", deviceId, access ? \"CAN\" : \"CANNOT\", mapPair.second);\n if (access && bestFitDeviceIds.size() < kNumGpusRequired) {\n bestFitDeviceIds.emplace(deviceId);\n bestFitDeviceIds.emplace(mapPair.second);\n }\n else {\n printf(\"Ignoring device %i (max devices exceeded)\\n\", mapPair.second);\n }\n }\n });\n\n if (bestFitDeviceIds.size() >= kNumGpusRequired) {\n printf(\"Selected p2p capable devices - \");\n int i = 0;\n for (auto devicesItr = bestFitDeviceIds.begin(); devicesItr != bestFitDeviceIds.end(); devicesItr++) {\n devIds[i++] = *devicesItr;\n printf(\"deviceId = %d \", *devicesItr);\n }\n printf(\"\\n\");\n break;\n }\n }\n\n // if bestFitDeviceIds.size() == 0 it means the GPUs in system are not p2p\n // capable, hence we add it without p2p capability check.\n if (!bestFitDeviceIds.size()) {\n printf(\"No Two or more Devices p2p capable found.. exiting..\\n\");\n exit(EXIT_WAIVED);\n }\n\n auto p2pGpuPair = std::make_pair(devIds[0], devIds[1]);\n\n return p2pGpuPair;\n}\n\nint memPoolP2PCopy()\n{\n int *dev0_srcVec, *dev1_dstVec; // Device buffers\n hipStream_t stream1, stream2;\n hipMemPool_t memPool;\n hipEvent_t waitOnStream1;\n\n // Allocate CPU memory.\n size_t nelem = 1048576;\n size_t bytes = nelem * sizeof(int);\n\n int *a = (int *)malloc(bytes);\n int *output = (int *)malloc(bytes);\n\n /* Initialize the vectors. */\n for (int n = 0; n < nelem; n++) {\n a[n] = rand() / (int)RAND_MAX;\n }\n\n auto p2pDevices = getP2PCapableGpuPair();\n printf(\"selected devices = %d & %d\\n\", p2pDevices.first, p2pDevices.second);\n checkCudaErrors(hipSetDevice(p2pDevices.first));\n checkCudaErrors(hipEventCreate(&waitOnStream1));\n\n checkCudaErrors(hipStreamCreateWithFlags(&stream1, hipStreamNonBlocking));\n\n // Get the default mempool for device p2pDevices.first from the pair\n checkCudaErrors(hipDeviceGetDefaultMemPool(&memPool, p2pDevices.first));\n\n // Allocate memory in a stream from the pool set above.\n checkCudaErrors(hipMallocAsync(&dev0_srcVec, bytes, stream1));\n\n checkCudaErrors(hipMemcpyAsync(dev0_srcVec, a, bytes, hipMemcpyHostToDevice, stream1));\n checkCudaErrors(hipEventRecord(waitOnStream1, stream1));\n\n checkCudaErrors(hipSetDevice(p2pDevices.second));\n checkCudaErrors(hipStreamCreateWithFlags(&stream2, hipStreamNonBlocking));\n\n // Allocate memory in p2pDevices.second device\n checkCudaErrors(hipMallocAsync(&dev1_dstVec, bytes, stream2));\n\n // Setup peer mappings for p2pDevices.second device\n hipMemAccessDesc desc;\n memset(&desc, 0, sizeof(hipMemAccessDesc));\n desc.location.type = hipMemLocationTypeDevice;\n desc.location.id = p2pDevices.second;\n desc.flags = hipMemAccessFlagsProtReadWrite;\n checkCudaErrors(hipMemPoolSetAccess(memPool, &desc, 1));\n\n printf(\"> copyP2PAndScale kernel running ...\\n\");\n dim3 block(256);\n dim3 grid((unsigned int)ceil(nelem / (int)block.x));\n checkCudaErrors(hipStreamWaitEvent(stream2, waitOnStream1));\n copyP2PAndScale<<>>(dev0_srcVec, dev1_dstVec, nelem);\n\n checkCudaErrors(hipMemcpyAsync(output, dev1_dstVec, bytes, hipMemcpyDeviceToHost, stream2));\n checkCudaErrors(hipFreeAsync(dev0_srcVec, stream2));\n checkCudaErrors(hipFreeAsync(dev1_dstVec, stream2));\n checkCudaErrors(hipStreamSynchronize(stream2));\n\n /* Compare the results */\n printf(\"> Checking the results from copyP2PAndScale() ...\\n\");\n\n for (int n = 0; n < nelem; n++) {\n if ((2 * a[n]) != output[n]) {\n printf(\"mismatch i = %d expected = %d val = %d\\n\", n, 2 * a[n], output[n]);\n return EXIT_FAILURE;\n }\n }\n\n free(a);\n free(output);\n checkCudaErrors(hipStreamDestroy(stream1));\n checkCudaErrors(hipStreamDestroy(stream2));\n printf(\"PASSED\\n\");\n\n return EXIT_SUCCESS;\n}\n\nint main(int argc, char **argv)\n{\n int ret = memPoolP2PCopy();\n return ret;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/streamOrderedAllocationP2P/streamOrderedAllocationP2P.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * Demonstration of inline PTX (assembly language) usage in CUDA kernels\n */\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n__global__ void sequence_gpu(int *d_ptr, int length)\n{\n int elemID = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (elemID < length) {\n unsigned int laneid;\n // This command gets the lane ID within the current warp\n asm(\"mov.u32 %0, %%laneid;\" : \"=r\"(laneid));\n d_ptr[elemID] = laneid;\n }\n}\n\n\nvoid sequence_cpu(int *h_ptr, int length)\n{\n for (int elemID = 0; elemID < length; elemID++) {\n h_ptr[elemID] = elemID % 32;\n }\n}\n\nint main(int argc, char **argv)\n{\n printf(\"CUDA inline PTX assembler sample\\n\");\n\n const int N = 1000;\n\n int dev = findCudaDevice(argc, (const char **)argv);\n\n if (dev == -1) {\n return EXIT_FAILURE;\n }\n\n int *d_ptr;\n checkCudaErrors(cudaMalloc(&d_ptr, N * sizeof(int)));\n\n int *h_ptr;\n checkCudaErrors(cudaMallocHost(&h_ptr, N * sizeof(int)));\n\n dim3 cudaBlockSize(256, 1, 1);\n dim3 cudaGridSize((N + cudaBlockSize.x - 1) / cudaBlockSize.x, 1, 1);\n sequence_gpu<<>>(d_ptr, N);\n checkCudaErrors(cudaGetLastError());\n checkCudaErrors(cudaDeviceSynchronize());\n\n sequence_cpu(h_ptr, N);\n\n int *h_d_ptr;\n checkCudaErrors(cudaMallocHost(&h_d_ptr, N * sizeof(int)));\n checkCudaErrors(cudaMemcpy(h_d_ptr, d_ptr, N * sizeof(int), cudaMemcpyDeviceToHost));\n\n bool bValid = true;\n\n for (int i = 0; i < N && bValid; i++) {\n if (h_ptr[i] != h_d_ptr[i]) {\n bValid = false;\n }\n }\n\n printf(\"Test %s.\\n\", bValid ? \"Successful\" : \"Failed\");\n\n checkCudaErrors(cudaFree(d_ptr));\n checkCudaErrors(cudaFreeHost(h_ptr));\n checkCudaErrors(cudaFreeHost(h_d_ptr));\n\n return bValid ? EXIT_SUCCESS : EXIT_FAILURE;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * Demonstration of inline PTX (assembly language) usage in CUDA kernels\n */\n\n// System includes\n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n__global__ void sequence_gpu(int *d_ptr, int length)\n{\n int elemID = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (elemID < length) {\n unsigned int laneid;\n // This command gets the lane ID within the current warp\n asm(\"mov.u32 %0, %%laneid;\" : \"=r\"(laneid));\n d_ptr[elemID] = laneid;\n }\n}\n\n\nvoid sequence_cpu(int *h_ptr, int length)\n{\n for (int elemID = 0; elemID < length; elemID++) {\n h_ptr[elemID] = elemID % 32;\n }\n}\n\nint main(int argc, char **argv)\n{\n printf(\"CUDA inline PTX assembler sample\\n\");\n\n const int N = 1000;\n\n int dev = findCudaDevice(argc, (const char **)argv);\n\n if (dev == -1) {\n return EXIT_FAILURE;\n }\n\n int *d_ptr;\n checkCudaErrors(hipMalloc(&d_ptr, N * sizeof(int)));\n\n int *h_ptr;\n checkCudaErrors(hipHostMalloc(&h_ptr, N * sizeof(int)));\n\n dim3 cudaBlockSize(256, 1, 1);\n dim3 cudaGridSize((N + cudaBlockSize.x - 1) / cudaBlockSize.x, 1, 1);\n sequence_gpu<<>>(d_ptr, N);\n checkCudaErrors(hipGetLastError());\n checkCudaErrors(hipDeviceSynchronize());\n\n sequence_cpu(h_ptr, N);\n\n int *h_d_ptr;\n checkCudaErrors(hipHostMalloc(&h_d_ptr, N * sizeof(int)));\n checkCudaErrors(hipMemcpy(h_d_ptr, d_ptr, N * sizeof(int), hipMemcpyDeviceToHost));\n\n bool bValid = true;\n\n for (int i = 0; i < N && bValid; i++) {\n if (h_ptr[i] != h_d_ptr[i]) {\n bValid = false;\n }\n }\n\n printf(\"Test %s.\\n\", bValid ? \"Successful\" : \"Failed\");\n\n checkCudaErrors(hipFree(d_ptr));\n checkCudaErrors(hipHostFree(h_ptr));\n checkCudaErrors(hipHostFree(h_d_ptr));\n\n return bValid ? EXIT_SUCCESS : EXIT_FAILURE;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/inlinePTX/inlinePTX.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#ifndef _BOXFILTER_KERNEL_CH_\n#define _BOXFILTER_KERNEL_CH_\n\n#include \n#include \n\ncudaTextureObject_t tex;\ncudaTextureObject_t texTempArray;\ncudaTextureObject_t rgbaTex;\ncudaTextureObject_t rgbaTexTempArray;\ncudaArray *d_array, *d_tempArray;\n\n////////////////////////////////////////////////////////////////////////////////\n// These are CUDA Helper functions\n\n// This will output the proper CUDA error strings in the event that a CUDA host\n// call returns an error\n#define checkCudaErrors(err) __checkCudaErrors(err, __FILE__, __LINE__)\n\ninline void __checkCudaErrors(cudaError err, const char *file, const int line)\n{\n if (cudaSuccess != err) {\n fprintf(stderr, \"%s(%i) : CUDA Runtime API error %d: %s.\\n\", file, line, (int)err, cudaGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n}\n\n/*\n Perform a fast box filter using the sliding window method.\n\n As the kernel moves from left to right, we add in the contribution of the\n new sample on the right, and subtract the value of the exiting sample on the\n left. This only requires 2 adds and a mul per output value, independent of the\n filter radius. The box filter is separable, so to perform a 2D box filter we\n perform the filter in the x direction, followed by the same filter in the y\n direction. Applying multiple iterations of the box filter converges towards a\n Gaussian blur. Using CUDA, rows or columns of the image are processed in\n parallel. This version duplicates edge pixels.\n\n Note that the x (row) pass suffers from uncoalesced global memory reads,\n since each thread is reading from a different row. For this reason it is\n better to use texture lookups for the x pass.\n The y (column) pass is perfectly coalesced.\n\n Parameters\n id - pointer to input data in global memory\n od - pointer to output data in global memory\n w - image width\n h - image height\n r - filter radius\n\n e.g. for r = 2, w = 8:\n\n 0 1 2 3 4 5 6 7\n x - -\n - x - -\n - - x - -\n - - x - -\n - - x - -\n - - x - -\n - - x -\n - - x\n*/\n\n// process row\n__device__ void d_boxfilter_x(float *id, float *od, int w, int h, int r)\n{\n float scale = 1.0f / (float)((r << 1) + 1);\n\n float t;\n // do left edge\n t = id[0] * r;\n\n for (int x = 0; x < (r + 1); x++) {\n t += id[x];\n }\n\n od[0] = t * scale;\n\n for (int x = 1; x < (r + 1); x++) {\n t += id[x + r];\n t -= id[0];\n od[x] = t * scale;\n }\n\n // main loop\n for (int x = (r + 1); x < w - r; x++) {\n t += id[x + r];\n t -= id[x - r - 1];\n od[x] = t * scale;\n }\n\n // do right edge\n for (int x = w - r; x < w; x++) {\n t += id[w - 1];\n t -= id[x - r - 1];\n od[x] = t * scale;\n }\n}\n\n// process column\n__device__ void d_boxfilter_y(float *id, float *od, int w, int h, int r)\n{\n float scale = 1.0f / (float)((r << 1) + 1);\n\n float t;\n // do left edge\n t = id[0] * r;\n\n for (int y = 0; y < (r + 1); y++) {\n t += id[y * w];\n }\n\n od[0] = t * scale;\n\n for (int y = 1; y < (r + 1); y++) {\n t += id[(y + r) * w];\n t -= id[0];\n od[y * w] = t * scale;\n }\n\n // main loop\n for (int y = (r + 1); y < (h - r); y++) {\n t += id[(y + r) * w];\n t -= id[((y - r) * w) - w];\n od[y * w] = t * scale;\n }\n\n // do right edge\n for (int y = h - r; y < h; y++) {\n t += id[(h - 1) * w];\n t -= id[((y - r) * w) - w];\n od[y * w] = t * scale;\n }\n}\n\n__global__ void d_boxfilter_x_global(float *id, float *od, int w, int h, int r)\n{\n unsigned int y = blockIdx.x * blockDim.x + threadIdx.x;\n d_boxfilter_x(&id[y * w], &od[y * w], w, h, r);\n}\n\n__global__ void d_boxfilter_y_global(float *id, float *od, int w, int h, int r)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n d_boxfilter_y(&id[x], &od[x], w, h, r);\n}\n\n// texture version\n// texture fetches automatically clamp to edge of image\n__global__ void d_boxfilter_x_tex(float *od, int w, int h, int r, cudaTextureObject_t tex)\n{\n float scale = 1.0f / (float)((r << 1) + 1);\n unsigned int y = blockIdx.x * blockDim.x + threadIdx.x;\n\n float t = 0.0f;\n\n for (int x = -r; x <= r; x++) {\n t += tex2D(tex, x, y);\n }\n\n od[y * w] = t * scale;\n\n for (int x = 1; x < w; x++) {\n t += tex2D(tex, x + r, y);\n t -= tex2D(tex, x - r - 1, y);\n od[y * w + x] = t * scale;\n }\n}\n\n__global__ void d_boxfilter_y_tex(float *od, int w, int h, int r, cudaTextureObject_t tex)\n{\n float scale = 1.0f / (float)((r << 1) + 1);\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n\n float t = 0.0f;\n\n for (int y = -r; y <= r; y++) {\n t += tex2D(tex, x, y);\n }\n\n od[x] = t * scale;\n\n for (int y = 1; y < h; y++) {\n t += tex2D(tex, x, y + r);\n t -= tex2D(tex, x, y - r - 1);\n od[y * w + x] = t * scale;\n }\n}\n\n// RGBA version\n// reads from 32-bit unsigned int array holding 8-bit RGBA\n\n// convert floating point rgba color to 32-bit integer\n__device__ unsigned int rgbaFloatToInt(float4 rgba)\n{\n rgba.x = __saturatef(rgba.x); // clamp to [0.0, 1.0]\n rgba.y = __saturatef(rgba.y);\n rgba.z = __saturatef(rgba.z);\n rgba.w = __saturatef(rgba.w);\n return ((unsigned int)(rgba.w * 255.0f) << 24) | ((unsigned int)(rgba.z * 255.0f) << 16)\n | ((unsigned int)(rgba.y * 255.0f) << 8) | ((unsigned int)(rgba.x * 255.0f));\n}\n\n__device__ float4 rgbaIntToFloat(unsigned int c)\n{\n float4 rgba;\n rgba.x = (c & 0xff) * 0.003921568627f; // /255.0f;\n rgba.y = ((c >> 8) & 0xff) * 0.003921568627f; // /255.0f;\n rgba.z = ((c >> 16) & 0xff) * 0.003921568627f; // /255.0f;\n rgba.w = ((c >> 24) & 0xff) * 0.003921568627f; // /255.0f;\n return rgba;\n}\n\n// row pass using texture lookups\n__global__ void d_boxfilter_rgba_x(unsigned int *od, int w, int h, int r, cudaTextureObject_t rgbaTex)\n{\n float scale = 1.0f / (float)((r << 1) + 1);\n unsigned int y = blockIdx.x * blockDim.x + threadIdx.x;\n\n // as long as address is always less than height, we do work\n if (y < h) {\n float4 t = make_float4(0.0f);\n\n for (int x = -r; x <= r; x++) {\n t += tex2D(rgbaTex, x, y);\n }\n\n od[y * w] = rgbaFloatToInt(t * scale);\n\n for (int x = 1; x < w; x++) {\n t += tex2D(rgbaTex, x + r, y);\n t -= tex2D(rgbaTex, x - r - 1, y);\n od[y * w + x] = rgbaFloatToInt(t * scale);\n }\n }\n}\n\n// column pass using coalesced global memory reads\n__global__ void d_boxfilter_rgba_y(unsigned int *id, unsigned int *od, int w, int h, int r)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n id = &id[x];\n od = &od[x];\n\n float scale = 1.0f / (float)((r << 1) + 1);\n\n float4 t;\n // do left edge\n t = rgbaIntToFloat(id[0]) * r;\n\n for (int y = 0; y < (r + 1); y++) {\n t += rgbaIntToFloat(id[y * w]);\n }\n\n od[0] = rgbaFloatToInt(t * scale);\n\n for (int y = 1; y < (r + 1); y++) {\n t += rgbaIntToFloat(id[(y + r) * w]);\n t -= rgbaIntToFloat(id[0]);\n od[y * w] = rgbaFloatToInt(t * scale);\n }\n\n // main loop\n for (int y = (r + 1); y < (h - r); y++) {\n t += rgbaIntToFloat(id[(y + r) * w]);\n t -= rgbaIntToFloat(id[((y - r) * w) - w]);\n od[y * w] = rgbaFloatToInt(t * scale);\n }\n\n // do right edge\n for (int y = h - r; y < h; y++) {\n t += rgbaIntToFloat(id[(h - 1) * w]);\n t -= rgbaIntToFloat(id[((y - r) * w) - w]);\n od[y * w] = rgbaFloatToInt(t * scale);\n }\n}\n\nextern \"C\" void initTexture(int width, int height, void *pImage, bool useRGBA)\n{\n // copy image data to array\n cudaChannelFormatDesc channelDesc;\n if (useRGBA) {\n channelDesc = cudaCreateChannelDesc(8, 8, 8, 8, cudaChannelFormatKindUnsigned);\n }\n else {\n channelDesc = cudaCreateChannelDesc(32, 0, 0, 0, cudaChannelFormatKindFloat);\n }\n checkCudaErrors(cudaMallocArray(&d_array, &channelDesc, width, height));\n\n size_t bytesPerElem = (useRGBA ? sizeof(uchar4) : sizeof(float));\n checkCudaErrors(cudaMemcpy2DToArray(\n d_array, 0, 0, pImage, width * bytesPerElem, width * bytesPerElem, height, cudaMemcpyHostToDevice));\n\n checkCudaErrors(cudaMallocArray(&d_tempArray, &channelDesc, width, height));\n\n // set texture parameters\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = d_array;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = cudaFilterModeLinear;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeNormalizedFloat;\n\n checkCudaErrors(cudaCreateTextureObject(&rgbaTex, &texRes, &texDescr, NULL));\n\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = d_tempArray;\n\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = cudaFilterModeLinear;\n texDescr.addressMode[0] = cudaAddressModeClamp;\n texDescr.addressMode[1] = cudaAddressModeClamp;\n texDescr.readMode = cudaReadModeNormalizedFloat;\n\n checkCudaErrors(cudaCreateTextureObject(&rgbaTexTempArray, &texRes, &texDescr, NULL));\n\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = d_array;\n\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = cudaFilterModePoint;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = d_tempArray;\n\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = cudaFilterModePoint;\n texDescr.addressMode[0] = cudaAddressModeWrap;\n texDescr.addressMode[1] = cudaAddressModeWrap;\n texDescr.readMode = cudaReadModeElementType;\n\n checkCudaErrors(cudaCreateTextureObject(&texTempArray, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void freeTextures()\n{\n checkCudaErrors(cudaDestroyTextureObject(tex));\n checkCudaErrors(cudaDestroyTextureObject(texTempArray));\n checkCudaErrors(cudaDestroyTextureObject(rgbaTex));\n checkCudaErrors(cudaDestroyTextureObject(rgbaTexTempArray));\n checkCudaErrors(cudaFreeArray(d_array));\n checkCudaErrors(cudaFreeArray(d_tempArray));\n}\n\n/*\n Perform 2D box filter on image using CUDA\n\n Parameters:\n d_temp - pointer to temporary storage in device memory\n d_dest - pointer to destination image in device memory\n width - image width\n height - image height\n radius - filter radius\n iterations - number of iterations\n\n*/\nextern \"C\" double boxFilter(float *d_temp,\n float *d_dest,\n int width,\n int height,\n int radius,\n int iterations,\n int nthreads,\n StopWatchInterface *timer)\n{\n // var for kernel timing\n double dKernelTime = 0.0;\n\n // sync host and start computation timer_kernel\n checkCudaErrors(cudaDeviceSynchronize());\n\n for (int i = 0; i < iterations; i++) {\n sdkResetTimer(&timer);\n // use texture for horizontal pass\n if (iterations > 1) {\n d_boxfilter_x_tex<<>>(d_temp, width, height, radius, texTempArray);\n }\n else {\n d_boxfilter_x_tex<<>>(d_temp, width, height, radius, tex);\n }\n\n d_boxfilter_y_global<<>>(d_temp, d_dest, width, height, radius);\n\n // sync host and stop computation timer_kernel\n checkCudaErrors(cudaDeviceSynchronize());\n dKernelTime += sdkGetTimerValue(&timer);\n\n if (iterations > 1) {\n // copy result back from global memory to array\n checkCudaErrors(cudaMemcpy2DToArray(d_tempArray,\n 0,\n 0,\n d_dest,\n width * sizeof(float),\n width * sizeof(float),\n height,\n cudaMemcpyDeviceToDevice));\n }\n }\n\n return ((dKernelTime / 1000.) / (double)iterations);\n}\n\n// RGBA version\nextern \"C\" double boxFilterRGBA(unsigned int *d_temp,\n unsigned int *d_dest,\n int width,\n int height,\n int radius,\n int iterations,\n int nthreads,\n StopWatchInterface *timer)\n{\n // var for kernel computation timing\n double dKernelTime;\n\n for (int i = 0; i < iterations; i++) {\n // sync host and start kernel computation timer_kernel\n dKernelTime = 0.0;\n checkCudaErrors(cudaDeviceSynchronize());\n sdkResetTimer(&timer);\n\n // use texture for horizontal pass\n if (iterations > 1) {\n d_boxfilter_rgba_x<<>>(d_temp, width, height, radius, rgbaTexTempArray);\n }\n else {\n d_boxfilter_rgba_x<<>>(d_temp, width, height, radius, rgbaTex);\n }\n\n d_boxfilter_rgba_y<<>>(d_temp, d_dest, width, height, radius);\n\n // sync host and stop computation timer_kernel\n checkCudaErrors(cudaDeviceSynchronize());\n dKernelTime += sdkGetTimerValue(&timer);\n\n if (iterations > 1) {\n // copy result back from global memory to array\n checkCudaErrors(cudaMemcpy2DToArray(d_tempArray,\n 0,\n 0,\n d_dest,\n width * sizeof(unsigned int),\n width * sizeof(unsigned int),\n height,\n cudaMemcpyDeviceToDevice));\n }\n }\n\n return ((dKernelTime / 1000.) / (double)iterations);\n}\n\n#endif // #ifndef _BOXFILTER_KERNEL_H_\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#ifndef _BOXFILTER_KERNEL_CH_\n#define _BOXFILTER_KERNEL_CH_\n\n#include \n#include \n\nhipTextureObject_t tex;\nhipTextureObject_t texTempArray;\nhipTextureObject_t rgbaTex;\nhipTextureObject_t rgbaTexTempArray;\nhipArray *d_array, *d_tempArray;\n\n////////////////////////////////////////////////////////////////////////////////\n// These are CUDA Helper functions\n\n// This will output the proper CUDA error strings in the event that a CUDA host\n// call returns an error\n#define checkCudaErrors(err) __checkCudaErrors(err, __FILE__, __LINE__)\n\ninline void __checkCudaErrors(hipError_t err, const char *file, const int line)\n{\n if (hipSuccess != err) {\n fprintf(stderr, \"%s(%i) : CUDA Runtime API error %d: %s.\\n\", file, line, (int)err, hipGetErrorString(err));\n exit(EXIT_FAILURE);\n }\n}\n\n/*\n Perform a fast box filter using the sliding window method.\n\n As the kernel moves from left to right, we add in the contribution of the\n new sample on the right, and subtract the value of the exiting sample on the\n left. This only requires 2 adds and a mul per output value, independent of the\n filter radius. The box filter is separable, so to perform a 2D box filter we\n perform the filter in the x direction, followed by the same filter in the y\n direction. Applying multiple iterations of the box filter converges towards a\n Gaussian blur. Using CUDA, rows or columns of the image are processed in\n parallel. This version duplicates edge pixels.\n\n Note that the x (row) pass suffers from uncoalesced global memory reads,\n since each thread is reading from a different row. For this reason it is\n better to use texture lookups for the x pass.\n The y (column) pass is perfectly coalesced.\n\n Parameters\n id - pointer to input data in global memory\n od - pointer to output data in global memory\n w - image width\n h - image height\n r - filter radius\n\n e.g. for r = 2, w = 8:\n\n 0 1 2 3 4 5 6 7\n x - -\n - x - -\n - - x - -\n - - x - -\n - - x - -\n - - x - -\n - - x -\n - - x\n*/\n\n// process row\n__device__ void d_boxfilter_x(float *id, float *od, int w, int h, int r)\n{\n float scale = 1.0f / (float)((r << 1) + 1);\n\n float t;\n // do left edge\n t = id[0] * r;\n\n for (int x = 0; x < (r + 1); x++) {\n t += id[x];\n }\n\n od[0] = t * scale;\n\n for (int x = 1; x < (r + 1); x++) {\n t += id[x + r];\n t -= id[0];\n od[x] = t * scale;\n }\n\n // main loop\n for (int x = (r + 1); x < w - r; x++) {\n t += id[x + r];\n t -= id[x - r - 1];\n od[x] = t * scale;\n }\n\n // do right edge\n for (int x = w - r; x < w; x++) {\n t += id[w - 1];\n t -= id[x - r - 1];\n od[x] = t * scale;\n }\n}\n\n// process column\n__device__ void d_boxfilter_y(float *id, float *od, int w, int h, int r)\n{\n float scale = 1.0f / (float)((r << 1) + 1);\n\n float t;\n // do left edge\n t = id[0] * r;\n\n for (int y = 0; y < (r + 1); y++) {\n t += id[y * w];\n }\n\n od[0] = t * scale;\n\n for (int y = 1; y < (r + 1); y++) {\n t += id[(y + r) * w];\n t -= id[0];\n od[y * w] = t * scale;\n }\n\n // main loop\n for (int y = (r + 1); y < (h - r); y++) {\n t += id[(y + r) * w];\n t -= id[((y - r) * w) - w];\n od[y * w] = t * scale;\n }\n\n // do right edge\n for (int y = h - r; y < h; y++) {\n t += id[(h - 1) * w];\n t -= id[((y - r) * w) - w];\n od[y * w] = t * scale;\n }\n}\n\n__global__ void d_boxfilter_x_global(float *id, float *od, int w, int h, int r)\n{\n unsigned int y = blockIdx.x * blockDim.x + threadIdx.x;\n d_boxfilter_x(&id[y * w], &od[y * w], w, h, r);\n}\n\n__global__ void d_boxfilter_y_global(float *id, float *od, int w, int h, int r)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n d_boxfilter_y(&id[x], &od[x], w, h, r);\n}\n\n// texture version\n// texture fetches automatically clamp to edge of image\n__global__ void d_boxfilter_x_tex(float *od, int w, int h, int r, hipTextureObject_t tex)\n{\n float scale = 1.0f / (float)((r << 1) + 1);\n unsigned int y = blockIdx.x * blockDim.x + threadIdx.x;\n\n float t = 0.0f;\n\n for (int x = -r; x <= r; x++) {\n t += tex2D(tex, x, y);\n }\n\n od[y * w] = t * scale;\n\n for (int x = 1; x < w; x++) {\n t += tex2D(tex, x + r, y);\n t -= tex2D(tex, x - r - 1, y);\n od[y * w + x] = t * scale;\n }\n}\n\n__global__ void d_boxfilter_y_tex(float *od, int w, int h, int r, hipTextureObject_t tex)\n{\n float scale = 1.0f / (float)((r << 1) + 1);\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n\n float t = 0.0f;\n\n for (int y = -r; y <= r; y++) {\n t += tex2D(tex, x, y);\n }\n\n od[x] = t * scale;\n\n for (int y = 1; y < h; y++) {\n t += tex2D(tex, x, y + r);\n t -= tex2D(tex, x, y - r - 1);\n od[y * w + x] = t * scale;\n }\n}\n\n// RGBA version\n// reads from 32-bit unsigned int array holding 8-bit RGBA\n\n// convert floating point rgba color to 32-bit integer\n__device__ unsigned int rgbaFloatToInt(float4 rgba)\n{\n rgba.x = __saturatef(rgba.x); // clamp to [0.0, 1.0]\n rgba.y = __saturatef(rgba.y);\n rgba.z = __saturatef(rgba.z);\n rgba.w = __saturatef(rgba.w);\n return ((unsigned int)(rgba.w * 255.0f) << 24) | ((unsigned int)(rgba.z * 255.0f) << 16)\n | ((unsigned int)(rgba.y * 255.0f) << 8) | ((unsigned int)(rgba.x * 255.0f));\n}\n\n__device__ float4 rgbaIntToFloat(unsigned int c)\n{\n float4 rgba;\n rgba.x = (c & 0xff) * 0.003921568627f; // /255.0f;\n rgba.y = ((c >> 8) & 0xff) * 0.003921568627f; // /255.0f;\n rgba.z = ((c >> 16) & 0xff) * 0.003921568627f; // /255.0f;\n rgba.w = ((c >> 24) & 0xff) * 0.003921568627f; // /255.0f;\n return rgba;\n}\n\n// row pass using texture lookups\n__global__ void d_boxfilter_rgba_x(unsigned int *od, int w, int h, int r, hipTextureObject_t rgbaTex)\n{\n float scale = 1.0f / (float)((r << 1) + 1);\n unsigned int y = blockIdx.x * blockDim.x + threadIdx.x;\n\n // as long as address is always less than height, we do work\n if (y < h) {\n float4 t = make_float4(0.0f);\n\n for (int x = -r; x <= r; x++) {\n t += tex2D(rgbaTex, x, y);\n }\n\n od[y * w] = rgbaFloatToInt(t * scale);\n\n for (int x = 1; x < w; x++) {\n t += tex2D(rgbaTex, x + r, y);\n t -= tex2D(rgbaTex, x - r - 1, y);\n od[y * w + x] = rgbaFloatToInt(t * scale);\n }\n }\n}\n\n// column pass using coalesced global memory reads\n__global__ void d_boxfilter_rgba_y(unsigned int *id, unsigned int *od, int w, int h, int r)\n{\n unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;\n id = &id[x];\n od = &od[x];\n\n float scale = 1.0f / (float)((r << 1) + 1);\n\n float4 t;\n // do left edge\n t = rgbaIntToFloat(id[0]) * r;\n\n for (int y = 0; y < (r + 1); y++) {\n t += rgbaIntToFloat(id[y * w]);\n }\n\n od[0] = rgbaFloatToInt(t * scale);\n\n for (int y = 1; y < (r + 1); y++) {\n t += rgbaIntToFloat(id[(y + r) * w]);\n t -= rgbaIntToFloat(id[0]);\n od[y * w] = rgbaFloatToInt(t * scale);\n }\n\n // main loop\n for (int y = (r + 1); y < (h - r); y++) {\n t += rgbaIntToFloat(id[(y + r) * w]);\n t -= rgbaIntToFloat(id[((y - r) * w) - w]);\n od[y * w] = rgbaFloatToInt(t * scale);\n }\n\n // do right edge\n for (int y = h - r; y < h; y++) {\n t += rgbaIntToFloat(id[(h - 1) * w]);\n t -= rgbaIntToFloat(id[((y - r) * w) - w]);\n od[y * w] = rgbaFloatToInt(t * scale);\n }\n}\n\nextern \"C\" void initTexture(int width, int height, void *pImage, bool useRGBA)\n{\n // copy image data to array\n hipChannelFormatDesc channelDesc;\n if (useRGBA) {\n channelDesc = hipCreateChannelDesc(8, 8, 8, 8, hipChannelFormatKindUnsigned);\n }\n else {\n channelDesc = hipCreateChannelDesc(32, 0, 0, 0, hipChannelFormatKindFloat);\n }\n checkCudaErrors(hipMallocArray(&d_array, &channelDesc, width, height));\n\n size_t bytesPerElem = (useRGBA ? sizeof(uchar4) : sizeof(float));\n checkCudaErrors(hipMemcpy2DToArray(\n d_array, 0, 0, pImage, width * bytesPerElem, width * bytesPerElem, height, hipMemcpyHostToDevice));\n\n checkCudaErrors(hipMallocArray(&d_tempArray, &channelDesc, width, height));\n\n // set texture parameters\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = d_array;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = hipFilterModeLinear;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeNormalizedFloat;\n\n checkCudaErrors(hipCreateTextureObject(&rgbaTex, &texRes, &texDescr, NULL));\n\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = d_tempArray;\n\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = false;\n texDescr.filterMode = hipFilterModeLinear;\n texDescr.addressMode[0] = hipAddressModeClamp;\n texDescr.addressMode[1] = hipAddressModeClamp;\n texDescr.readMode = hipReadModeNormalizedFloat;\n\n checkCudaErrors(hipCreateTextureObject(&rgbaTexTempArray, &texRes, &texDescr, NULL));\n\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = d_array;\n\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = hipFilterModePoint;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = d_tempArray;\n\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n texDescr.normalizedCoords = true;\n texDescr.filterMode = hipFilterModePoint;\n texDescr.addressMode[0] = hipAddressModeWrap;\n texDescr.addressMode[1] = hipAddressModeWrap;\n texDescr.readMode = hipReadModeElementType;\n\n checkCudaErrors(hipCreateTextureObject(&texTempArray, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void freeTextures()\n{\n checkCudaErrors(hipDestroyTextureObject(tex));\n checkCudaErrors(hipDestroyTextureObject(texTempArray));\n checkCudaErrors(hipDestroyTextureObject(rgbaTex));\n checkCudaErrors(hipDestroyTextureObject(rgbaTexTempArray));\n checkCudaErrors(hipFreeArray(d_array));\n checkCudaErrors(hipFreeArray(d_tempArray));\n}\n\n/*\n Perform 2D box filter on image using CUDA\n\n Parameters:\n d_temp - pointer to temporary storage in device memory\n d_dest - pointer to destination image in device memory\n width - image width\n height - image height\n radius - filter radius\n iterations - number of iterations\n\n*/\nextern \"C\" double boxFilter(float *d_temp,\n float *d_dest,\n int width,\n int height,\n int radius,\n int iterations,\n int nthreads,\n StopWatchInterface *timer)\n{\n // var for kernel timing\n double dKernelTime = 0.0;\n\n // sync host and start computation timer_kernel\n checkCudaErrors(hipDeviceSynchronize());\n\n for (int i = 0; i < iterations; i++) {\n sdkResetTimer(&timer);\n // use texture for horizontal pass\n if (iterations > 1) {\n d_boxfilter_x_tex<<>>(d_temp, width, height, radius, texTempArray);\n }\n else {\n d_boxfilter_x_tex<<>>(d_temp, width, height, radius, tex);\n }\n\n d_boxfilter_y_global<<>>(d_temp, d_dest, width, height, radius);\n\n // sync host and stop computation timer_kernel\n checkCudaErrors(hipDeviceSynchronize());\n dKernelTime += sdkGetTimerValue(&timer);\n\n if (iterations > 1) {\n // copy result back from global memory to array\n checkCudaErrors(hipMemcpy2DToArray(d_tempArray,\n 0,\n 0,\n d_dest,\n width * sizeof(float),\n width * sizeof(float),\n height,\n hipMemcpyDeviceToDevice));\n }\n }\n\n return ((dKernelTime / 1000.) / (double)iterations);\n}\n\n// RGBA version\nextern \"C\" double boxFilterRGBA(unsigned int *d_temp,\n unsigned int *d_dest,\n int width,\n int height,\n int radius,\n int iterations,\n int nthreads,\n StopWatchInterface *timer)\n{\n // var for kernel computation timing\n double dKernelTime;\n\n for (int i = 0; i < iterations; i++) {\n // sync host and start kernel computation timer_kernel\n dKernelTime = 0.0;\n checkCudaErrors(hipDeviceSynchronize());\n sdkResetTimer(&timer);\n\n // use texture for horizontal pass\n if (iterations > 1) {\n d_boxfilter_rgba_x<<>>(d_temp, width, height, radius, rgbaTexTempArray);\n }\n else {\n d_boxfilter_rgba_x<<>>(d_temp, width, height, radius, rgbaTex);\n }\n\n d_boxfilter_rgba_y<<>>(d_temp, d_dest, width, height, radius);\n\n // sync host and stop computation timer_kernel\n checkCudaErrors(hipDeviceSynchronize());\n dKernelTime += sdkGetTimerValue(&timer);\n\n if (iterations > 1) {\n // copy result back from global memory to array\n checkCudaErrors(hipMemcpy2DToArray(d_tempArray,\n 0,\n 0,\n d_dest,\n width * sizeof(unsigned int),\n width * sizeof(unsigned int),\n height,\n hipMemcpyDeviceToDevice));\n }\n }\n\n return ((dKernelTime / 1000.) / (double)iterations);\n}\n\n#endif // #ifndef _BOXFILTER_KERNEL_H_\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/boxFilter/boxFilter_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\nextern \"C\" __global__ void sequence_gpu(int *d_ptr, int length)\n{\n int elemID = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (elemID < length) {\n unsigned int laneid;\n\n // This command gets the lane ID within the current warp\n asm(\"mov.u32 %0, %%laneid;\" : \"=r\"(laneid));\n\n d_ptr[elemID] = laneid;\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\nextern \"C\" __global__ void sequence_gpu(int *d_ptr, int length)\n{\n int elemID = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (elemID < length) {\n unsigned int laneid;\n\n // This command gets the lane ID within the current warp\n asm(\"mov.u32 %0, %%laneid;\" : \"=r\"(laneid));\n\n d_ptr[elemID] = laneid;\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/inlinePTX_nvrtc/inlinePTX_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Based on http://www.iti.fh-flensburg.de/lang/algorithmen/sortieren/bitonic/bitonicen.htm\n\n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \n\n#include \"sortingNetworks_common.cuh\"\n#include \"sortingNetworks_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Monolithic bitonic sort kernel for short arrays fitting into shared memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void\nbitonicSortShared(uint *d_DstKey, uint *d_DstVal, uint *d_SrcKey, uint *d_SrcVal, uint arrayLength, uint dir)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Shared memory storage for one or more short vectors\n __shared__ uint s_key[SHARED_SIZE_LIMIT];\n __shared__ uint s_val[SHARED_SIZE_LIMIT];\n\n // Offset to the beginning of subbatch and load data\n d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n s_key[threadIdx.x + 0] = d_SrcKey[0];\n s_val[threadIdx.x + 0] = d_SrcVal[0];\n s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];\n s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];\n\n for (uint size = 2; size < arrayLength; size <<= 1) {\n // Bitonic merge\n uint ddd = dir ^ ((threadIdx.x & (size / 2)) != 0);\n\n for (uint stride = size / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], ddd);\n }\n }\n\n // ddd == dir for the last bitonic merge step\n {\n for (uint stride = arrayLength / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], dir);\n }\n }\n\n cg::sync(cta);\n d_DstKey[0] = s_key[threadIdx.x + 0];\n d_DstVal[0] = s_val[threadIdx.x + 0];\n d_DstKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n d_DstVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Bitonic sort kernel for large arrays (not fitting into shared memory)\n////////////////////////////////////////////////////////////////////////////////\n// Bottom-level bitonic sort\n// Almost the same as bitonicSortShared with the exception of\n// even / odd subarrays being sorted in opposite directions\n// Bitonic merge accepts both\n// Ascending | descending or descending | ascending sorted pairs\n__global__ void bitonicSortShared1(uint *d_DstKey, uint *d_DstVal, uint *d_SrcKey, uint *d_SrcVal)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Shared memory storage for current subarray\n __shared__ uint s_key[SHARED_SIZE_LIMIT];\n __shared__ uint s_val[SHARED_SIZE_LIMIT];\n\n // Offset to the beginning of subarray and load data\n d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n s_key[threadIdx.x + 0] = d_SrcKey[0];\n s_val[threadIdx.x + 0] = d_SrcVal[0];\n s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];\n s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];\n\n for (uint size = 2; size < SHARED_SIZE_LIMIT; size <<= 1) {\n // Bitonic merge\n uint ddd = (threadIdx.x & (size / 2)) != 0;\n\n for (uint stride = size / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], ddd);\n }\n }\n\n // Odd / even arrays of SHARED_SIZE_LIMIT elements\n // sorted in opposite directions\n uint ddd = blockIdx.x & 1;\n {\n for (uint stride = SHARED_SIZE_LIMIT / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], ddd);\n }\n }\n\n cg::sync(cta);\n d_DstKey[0] = s_key[threadIdx.x + 0];\n d_DstVal[0] = s_val[threadIdx.x + 0];\n d_DstKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n d_DstVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n}\n\n// Bitonic merge iteration for stride >= SHARED_SIZE_LIMIT\n__global__ void bitonicMergeGlobal(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint arrayLength,\n uint size,\n uint stride,\n uint dir)\n{\n uint global_comparatorI = blockIdx.x * blockDim.x + threadIdx.x;\n uint comparatorI = global_comparatorI & (arrayLength / 2 - 1);\n\n // Bitonic merge\n uint ddd = dir ^ ((comparatorI & (size / 2)) != 0);\n uint pos = 2 * global_comparatorI - (global_comparatorI & (stride - 1));\n\n uint keyA = d_SrcKey[pos + 0];\n uint valA = d_SrcVal[pos + 0];\n uint keyB = d_SrcKey[pos + stride];\n uint valB = d_SrcVal[pos + stride];\n\n Comparator(keyA, valA, keyB, valB, ddd);\n\n d_DstKey[pos + 0] = keyA;\n d_DstVal[pos + 0] = valA;\n d_DstKey[pos + stride] = keyB;\n d_DstVal[pos + stride] = valB;\n}\n\n// Combined bitonic merge steps for\n// size > SHARED_SIZE_LIMIT and stride = [1 .. SHARED_SIZE_LIMIT / 2]\n__global__ void bitonicMergeShared(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint arrayLength,\n uint size,\n uint dir)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Shared memory storage for current subarray\n __shared__ uint s_key[SHARED_SIZE_LIMIT];\n __shared__ uint s_val[SHARED_SIZE_LIMIT];\n\n d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n s_key[threadIdx.x + 0] = d_SrcKey[0];\n s_val[threadIdx.x + 0] = d_SrcVal[0];\n s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];\n s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];\n\n // Bitonic merge\n uint comparatorI = UMAD(blockIdx.x, blockDim.x, threadIdx.x) & ((arrayLength / 2) - 1);\n uint ddd = dir ^ ((comparatorI & (size / 2)) != 0);\n\n for (uint stride = SHARED_SIZE_LIMIT / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], ddd);\n }\n\n cg::sync(cta);\n d_DstKey[0] = s_key[threadIdx.x + 0];\n d_DstVal[0] = s_val[threadIdx.x + 0];\n d_DstKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n d_DstVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Interface function\n////////////////////////////////////////////////////////////////////////////////\n// Helper function (also used by odd-even merge sort)\nextern \"C\" uint factorRadix2(uint *log2L, uint L)\n{\n if (!L) {\n *log2L = 0;\n return 0;\n }\n else {\n for (*log2L = 0; (L & 1) == 0; L >>= 1, *log2L++)\n ;\n\n return L;\n }\n}\n\nextern \"C\" uint\nbitonicSort(uint *d_DstKey, uint *d_DstVal, uint *d_SrcKey, uint *d_SrcVal, uint batchSize, uint arrayLength, uint dir)\n{\n // Nothing to sort\n if (arrayLength < 2)\n return 0;\n\n // Only power-of-two array lengths are supported by this implementation\n uint log2L;\n uint factorizationRemainder = factorRadix2(&log2L, arrayLength);\n assert(factorizationRemainder == 1);\n\n dir = (dir != 0);\n\n uint blockCount = batchSize * arrayLength / SHARED_SIZE_LIMIT;\n uint threadCount = SHARED_SIZE_LIMIT / 2;\n\n if (arrayLength <= SHARED_SIZE_LIMIT) {\n assert((batchSize * arrayLength) % SHARED_SIZE_LIMIT == 0);\n bitonicSortShared<<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, arrayLength, dir);\n }\n else {\n bitonicSortShared1<<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal);\n\n for (uint size = 2 * SHARED_SIZE_LIMIT; size <= arrayLength; size <<= 1)\n for (unsigned stride = size / 2; stride > 0; stride >>= 1)\n if (stride >= SHARED_SIZE_LIMIT) {\n bitonicMergeGlobal<<<(batchSize * arrayLength) / 512, 256>>>(\n d_DstKey, d_DstVal, d_DstKey, d_DstVal, arrayLength, size, stride, dir);\n }\n else {\n bitonicMergeShared<<>>(\n d_DstKey, d_DstVal, d_DstKey, d_DstVal, arrayLength, size, dir);\n break;\n }\n }\n\n return threadCount;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Based on http://www.iti.fh-flensburg.de/lang/algorithmen/sortieren/bitonic/bitonicen.htm\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \n\n#include \"sortingNetworks_common.cuh\"\n#include \"sortingNetworks_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Monolithic bitonic sort kernel for short arrays fitting into shared memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void\nbitonicSortShared(uint *d_DstKey, uint *d_DstVal, uint *d_SrcKey, uint *d_SrcVal, uint arrayLength, uint dir)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Shared memory storage for one or more short vectors\n __shared__ uint s_key[SHARED_SIZE_LIMIT];\n __shared__ uint s_val[SHARED_SIZE_LIMIT];\n\n // Offset to the beginning of subbatch and load data\n d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n s_key[threadIdx.x + 0] = d_SrcKey[0];\n s_val[threadIdx.x + 0] = d_SrcVal[0];\n s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];\n s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];\n\n for (uint size = 2; size < arrayLength; size <<= 1) {\n // Bitonic merge\n uint ddd = dir ^ ((threadIdx.x & (size / 2)) != 0);\n\n for (uint stride = size / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], ddd);\n }\n }\n\n // ddd == dir for the last bitonic merge step\n {\n for (uint stride = arrayLength / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], dir);\n }\n }\n\n cg::sync(cta);\n d_DstKey[0] = s_key[threadIdx.x + 0];\n d_DstVal[0] = s_val[threadIdx.x + 0];\n d_DstKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n d_DstVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Bitonic sort kernel for large arrays (not fitting into shared memory)\n////////////////////////////////////////////////////////////////////////////////\n// Bottom-level bitonic sort\n// Almost the same as bitonicSortShared with the exception of\n// even / odd subarrays being sorted in opposite directions\n// Bitonic merge accepts both\n// Ascending | descending or descending | ascending sorted pairs\n__global__ void bitonicSortShared1(uint *d_DstKey, uint *d_DstVal, uint *d_SrcKey, uint *d_SrcVal)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Shared memory storage for current subarray\n __shared__ uint s_key[SHARED_SIZE_LIMIT];\n __shared__ uint s_val[SHARED_SIZE_LIMIT];\n\n // Offset to the beginning of subarray and load data\n d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n s_key[threadIdx.x + 0] = d_SrcKey[0];\n s_val[threadIdx.x + 0] = d_SrcVal[0];\n s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];\n s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];\n\n for (uint size = 2; size < SHARED_SIZE_LIMIT; size <<= 1) {\n // Bitonic merge\n uint ddd = (threadIdx.x & (size / 2)) != 0;\n\n for (uint stride = size / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], ddd);\n }\n }\n\n // Odd / even arrays of SHARED_SIZE_LIMIT elements\n // sorted in opposite directions\n uint ddd = blockIdx.x & 1;\n {\n for (uint stride = SHARED_SIZE_LIMIT / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], ddd);\n }\n }\n\n cg::sync(cta);\n d_DstKey[0] = s_key[threadIdx.x + 0];\n d_DstVal[0] = s_val[threadIdx.x + 0];\n d_DstKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n d_DstVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n}\n\n// Bitonic merge iteration for stride >= SHARED_SIZE_LIMIT\n__global__ void bitonicMergeGlobal(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint arrayLength,\n uint size,\n uint stride,\n uint dir)\n{\n uint global_comparatorI = blockIdx.x * blockDim.x + threadIdx.x;\n uint comparatorI = global_comparatorI & (arrayLength / 2 - 1);\n\n // Bitonic merge\n uint ddd = dir ^ ((comparatorI & (size / 2)) != 0);\n uint pos = 2 * global_comparatorI - (global_comparatorI & (stride - 1));\n\n uint keyA = d_SrcKey[pos + 0];\n uint valA = d_SrcVal[pos + 0];\n uint keyB = d_SrcKey[pos + stride];\n uint valB = d_SrcVal[pos + stride];\n\n Comparator(keyA, valA, keyB, valB, ddd);\n\n d_DstKey[pos + 0] = keyA;\n d_DstVal[pos + 0] = valA;\n d_DstKey[pos + stride] = keyB;\n d_DstVal[pos + stride] = valB;\n}\n\n// Combined bitonic merge steps for\n// size > SHARED_SIZE_LIMIT and stride = [1 .. SHARED_SIZE_LIMIT / 2]\n__global__ void bitonicMergeShared(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint arrayLength,\n uint size,\n uint dir)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Shared memory storage for current subarray\n __shared__ uint s_key[SHARED_SIZE_LIMIT];\n __shared__ uint s_val[SHARED_SIZE_LIMIT];\n\n d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n s_key[threadIdx.x + 0] = d_SrcKey[0];\n s_val[threadIdx.x + 0] = d_SrcVal[0];\n s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];\n s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];\n\n // Bitonic merge\n uint comparatorI = UMAD(blockIdx.x, blockDim.x, threadIdx.x) & ((arrayLength / 2) - 1);\n uint ddd = dir ^ ((comparatorI & (size / 2)) != 0);\n\n for (uint stride = SHARED_SIZE_LIMIT / 2; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], ddd);\n }\n\n cg::sync(cta);\n d_DstKey[0] = s_key[threadIdx.x + 0];\n d_DstVal[0] = s_val[threadIdx.x + 0];\n d_DstKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n d_DstVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Interface function\n////////////////////////////////////////////////////////////////////////////////\n// Helper function (also used by odd-even merge sort)\nextern \"C\" uint factorRadix2(uint *log2L, uint L)\n{\n if (!L) {\n *log2L = 0;\n return 0;\n }\n else {\n for (*log2L = 0; (L & 1) == 0; L >>= 1, *log2L++)\n ;\n\n return L;\n }\n}\n\nextern \"C\" uint\nbitonicSort(uint *d_DstKey, uint *d_DstVal, uint *d_SrcKey, uint *d_SrcVal, uint batchSize, uint arrayLength, uint dir)\n{\n // Nothing to sort\n if (arrayLength < 2)\n return 0;\n\n // Only power-of-two array lengths are supported by this implementation\n uint log2L;\n uint factorizationRemainder = factorRadix2(&log2L, arrayLength);\n assert(factorizationRemainder == 1);\n\n dir = (dir != 0);\n\n uint blockCount = batchSize * arrayLength / SHARED_SIZE_LIMIT;\n uint threadCount = SHARED_SIZE_LIMIT / 2;\n\n if (arrayLength <= SHARED_SIZE_LIMIT) {\n assert((batchSize * arrayLength) % SHARED_SIZE_LIMIT == 0);\n bitonicSortShared<<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, arrayLength, dir);\n }\n else {\n bitonicSortShared1<<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal);\n\n for (uint size = 2 * SHARED_SIZE_LIMIT; size <= arrayLength; size <<= 1)\n for (unsigned stride = size / 2; stride > 0; stride >>= 1)\n if (stride >= SHARED_SIZE_LIMIT) {\n bitonicMergeGlobal<<<(batchSize * arrayLength) / 512, 256>>>(\n d_DstKey, d_DstVal, d_DstKey, d_DstVal, arrayLength, size, stride, dir);\n }\n else {\n bitonicMergeShared<<>>(\n d_DstKey, d_DstVal, d_DstKey, d_DstVal, arrayLength, size, dir);\n break;\n }\n }\n\n return threadCount;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/sortingNetworks/bitonicSort.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Based on http://www.iti.fh-flensburg.de/lang/algorithmen/sortieren/networks/oemen.htm\n\n\n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \n\n#include \"sortingNetworks_common.cuh\"\n#include \"sortingNetworks_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Monolithic Bacther's sort kernel for short arrays fitting into shared memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void\noddEvenMergeSortShared(uint *d_DstKey, uint *d_DstVal, uint *d_SrcKey, uint *d_SrcVal, uint arrayLength, uint dir)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Shared memory storage for one or more small vectors\n __shared__ uint s_key[SHARED_SIZE_LIMIT];\n __shared__ uint s_val[SHARED_SIZE_LIMIT];\n\n // Offset to the beginning of subbatch and load data\n d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n s_key[threadIdx.x + 0] = d_SrcKey[0];\n s_val[threadIdx.x + 0] = d_SrcVal[0];\n s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];\n s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];\n\n for (uint size = 2; size <= arrayLength; size <<= 1) {\n uint stride = size / 2;\n uint offset = threadIdx.x & (stride - 1);\n\n {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], dir);\n stride >>= 1;\n }\n\n for (; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n\n if (offset >= stride)\n Comparator(s_key[pos - stride], s_val[pos - stride], s_key[pos + 0], s_val[pos + 0], dir);\n }\n }\n\n cg::sync(cta);\n d_DstKey[0] = s_key[threadIdx.x + 0];\n d_DstVal[0] = s_val[threadIdx.x + 0];\n d_DstKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n d_DstVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Odd-even merge sort iteration kernel\n// for large arrays (not fitting into shared memory)\n////////////////////////////////////////////////////////////////////////////////\n__global__ void oddEvenMergeGlobal(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint arrayLength,\n uint size,\n uint stride,\n uint dir)\n{\n uint global_comparatorI = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Odd-even merge\n uint pos = 2 * global_comparatorI - (global_comparatorI & (stride - 1));\n\n if (stride < size / 2) {\n uint offset = global_comparatorI & ((size / 2) - 1);\n\n if (offset >= stride) {\n uint keyA = d_SrcKey[pos - stride];\n uint valA = d_SrcVal[pos - stride];\n uint keyB = d_SrcKey[pos + 0];\n uint valB = d_SrcVal[pos + 0];\n\n Comparator(keyA, valA, keyB, valB, dir);\n\n d_DstKey[pos - stride] = keyA;\n d_DstVal[pos - stride] = valA;\n d_DstKey[pos + 0] = keyB;\n d_DstVal[pos + 0] = valB;\n }\n }\n else {\n uint keyA = d_SrcKey[pos + 0];\n uint valA = d_SrcVal[pos + 0];\n uint keyB = d_SrcKey[pos + stride];\n uint valB = d_SrcVal[pos + stride];\n\n Comparator(keyA, valA, keyB, valB, dir);\n\n d_DstKey[pos + 0] = keyA;\n d_DstVal[pos + 0] = valA;\n d_DstKey[pos + stride] = keyB;\n d_DstVal[pos + stride] = valB;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Interface function\n////////////////////////////////////////////////////////////////////////////////\n// Helper function\nextern \"C\" uint factorRadix2(uint *log2L, uint L);\n\nextern \"C\" void oddEvenMergeSort(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint batchSize,\n uint arrayLength,\n uint dir)\n{\n // Nothing to sort\n if (arrayLength < 2)\n return;\n\n // Only power-of-two array lengths are supported by this implementation\n uint log2L;\n uint factorizationRemainder = factorRadix2(&log2L, arrayLength);\n assert(factorizationRemainder == 1);\n\n dir = (dir != 0);\n\n uint blockCount = (batchSize * arrayLength) / SHARED_SIZE_LIMIT;\n uint threadCount = SHARED_SIZE_LIMIT / 2;\n\n if (arrayLength <= SHARED_SIZE_LIMIT) {\n assert(SHARED_SIZE_LIMIT % arrayLength == 0);\n oddEvenMergeSortShared<<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, arrayLength, dir);\n }\n else {\n oddEvenMergeSortShared<<>>(\n d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, SHARED_SIZE_LIMIT, dir);\n\n for (uint size = 2 * SHARED_SIZE_LIMIT; size <= arrayLength; size <<= 1)\n for (unsigned stride = size / 2; stride > 0; stride >>= 1) {\n // Unlike with bitonic sort, combining bitonic merge steps with\n // stride = [SHARED_SIZE_LIMIT / 2 .. 1] seems to be impossible as there\n // are dependencies between data elements crossing the SHARED_SIZE_LIMIT\n // borders\n oddEvenMergeGlobal<<<(batchSize * arrayLength) / 512, 256>>>(\n d_DstKey, d_DstVal, d_DstKey, d_DstVal, arrayLength, size, stride, dir);\n }\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Based on http://www.iti.fh-flensburg.de/lang/algorithmen/sortieren/networks/oemen.htm\n\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \n\n#include \"sortingNetworks_common.cuh\"\n#include \"sortingNetworks_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Monolithic Bacther's sort kernel for short arrays fitting into shared memory\n////////////////////////////////////////////////////////////////////////////////\n__global__ void\noddEvenMergeSortShared(uint *d_DstKey, uint *d_DstVal, uint *d_SrcKey, uint *d_SrcVal, uint arrayLength, uint dir)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Shared memory storage for one or more small vectors\n __shared__ uint s_key[SHARED_SIZE_LIMIT];\n __shared__ uint s_val[SHARED_SIZE_LIMIT];\n\n // Offset to the beginning of subbatch and load data\n d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n d_DstVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;\n s_key[threadIdx.x + 0] = d_SrcKey[0];\n s_val[threadIdx.x + 0] = d_SrcVal[0];\n s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];\n s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];\n\n for (uint size = 2; size <= arrayLength; size <<= 1) {\n uint stride = size / 2;\n uint offset = threadIdx.x & (stride - 1);\n\n {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n Comparator(s_key[pos + 0], s_val[pos + 0], s_key[pos + stride], s_val[pos + stride], dir);\n stride >>= 1;\n }\n\n for (; stride > 0; stride >>= 1) {\n cg::sync(cta);\n uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));\n\n if (offset >= stride)\n Comparator(s_key[pos - stride], s_val[pos - stride], s_key[pos + 0], s_val[pos + 0], dir);\n }\n }\n\n cg::sync(cta);\n d_DstKey[0] = s_key[threadIdx.x + 0];\n d_DstVal[0] = s_val[threadIdx.x + 0];\n d_DstKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n d_DstVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Odd-even merge sort iteration kernel\n// for large arrays (not fitting into shared memory)\n////////////////////////////////////////////////////////////////////////////////\n__global__ void oddEvenMergeGlobal(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint arrayLength,\n uint size,\n uint stride,\n uint dir)\n{\n uint global_comparatorI = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Odd-even merge\n uint pos = 2 * global_comparatorI - (global_comparatorI & (stride - 1));\n\n if (stride < size / 2) {\n uint offset = global_comparatorI & ((size / 2) - 1);\n\n if (offset >= stride) {\n uint keyA = d_SrcKey[pos - stride];\n uint valA = d_SrcVal[pos - stride];\n uint keyB = d_SrcKey[pos + 0];\n uint valB = d_SrcVal[pos + 0];\n\n Comparator(keyA, valA, keyB, valB, dir);\n\n d_DstKey[pos - stride] = keyA;\n d_DstVal[pos - stride] = valA;\n d_DstKey[pos + 0] = keyB;\n d_DstVal[pos + 0] = valB;\n }\n }\n else {\n uint keyA = d_SrcKey[pos + 0];\n uint valA = d_SrcVal[pos + 0];\n uint keyB = d_SrcKey[pos + stride];\n uint valB = d_SrcVal[pos + stride];\n\n Comparator(keyA, valA, keyB, valB, dir);\n\n d_DstKey[pos + 0] = keyA;\n d_DstVal[pos + 0] = valA;\n d_DstKey[pos + stride] = keyB;\n d_DstVal[pos + stride] = valB;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Interface function\n////////////////////////////////////////////////////////////////////////////////\n// Helper function\nextern \"C\" uint factorRadix2(uint *log2L, uint L);\n\nextern \"C\" void oddEvenMergeSort(uint *d_DstKey,\n uint *d_DstVal,\n uint *d_SrcKey,\n uint *d_SrcVal,\n uint batchSize,\n uint arrayLength,\n uint dir)\n{\n // Nothing to sort\n if (arrayLength < 2)\n return;\n\n // Only power-of-two array lengths are supported by this implementation\n uint log2L;\n uint factorizationRemainder = factorRadix2(&log2L, arrayLength);\n assert(factorizationRemainder == 1);\n\n dir = (dir != 0);\n\n uint blockCount = (batchSize * arrayLength) / SHARED_SIZE_LIMIT;\n uint threadCount = SHARED_SIZE_LIMIT / 2;\n\n if (arrayLength <= SHARED_SIZE_LIMIT) {\n assert(SHARED_SIZE_LIMIT % arrayLength == 0);\n oddEvenMergeSortShared<<>>(d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, arrayLength, dir);\n }\n else {\n oddEvenMergeSortShared<<>>(\n d_DstKey, d_DstVal, d_SrcKey, d_SrcVal, SHARED_SIZE_LIMIT, dir);\n\n for (uint size = 2 * SHARED_SIZE_LIMIT; size <= arrayLength; size <<= 1)\n for (unsigned stride = size / 2; stride > 0; stride >>= 1) {\n // Unlike with bitonic sort, combining bitonic merge steps with\n // stride = [SHARED_SIZE_LIMIT / 2 .. 1] seems to be impossible as there\n // are dependencies between data elements crossing the SHARED_SIZE_LIMIT\n // borders\n oddEvenMergeGlobal<<<(batchSize * arrayLength) / 512, 256>>>(\n d_DstKey, d_DstVal, d_DstKey, d_DstVal, arrayLength, size, stride, dir);\n }\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/sortingNetworks/oddEvenMergeSort.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \n\n#include \"scan_common.h\"\n\n// All three kernels run 512 threads per workgroup\n// Must be a power of two\n#define THREADBLOCK_SIZE 256\n\n////////////////////////////////////////////////////////////////////////////////\n// Basic scan codelets\n////////////////////////////////////////////////////////////////////////////////\n// Naive inclusive scan: O(N * log2(N)) operations\n// Allocate 2 * 'size' local memory, initialize the first half\n// with 'size' zeros avoiding if(pos >= offset) condition evaluation\n// and saving instructions\ninline __device__ uint scan1Inclusive(uint idata, volatile uint *s_Data, uint size, cg::thread_block cta)\n{\n uint pos = 2 * threadIdx.x - (threadIdx.x & (size - 1));\n s_Data[pos] = 0;\n pos += size;\n s_Data[pos] = idata;\n\n for (uint offset = 1; offset < size; offset <<= 1) {\n cg::sync(cta);\n uint t = s_Data[pos] + s_Data[pos - offset];\n cg::sync(cta);\n s_Data[pos] = t;\n }\n\n return s_Data[pos];\n}\n\ninline __device__ uint scan1Exclusive(uint idata, volatile uint *s_Data, uint size, cg::thread_block cta)\n{\n return scan1Inclusive(idata, s_Data, size, cta) - idata;\n}\n\ninline __device__ uint4 scan4Inclusive(uint4 idata4, volatile uint *s_Data, uint size, cg::thread_block cta)\n{\n // Level-0 inclusive scan\n idata4.y += idata4.x;\n idata4.z += idata4.y;\n idata4.w += idata4.z;\n\n // Level-1 exclusive scan\n uint oval = scan1Exclusive(idata4.w, s_Data, size / 4, cta);\n\n idata4.x += oval;\n idata4.y += oval;\n idata4.z += oval;\n idata4.w += oval;\n\n return idata4;\n}\n\n// Exclusive vector scan: the array to be scanned is stored\n// in local thread memory scope as uint4\ninline __device__ uint4 scan4Exclusive(uint4 idata4, volatile uint *s_Data, uint size, cg::thread_block cta)\n{\n uint4 odata4 = scan4Inclusive(idata4, s_Data, size, cta);\n odata4.x -= idata4.x;\n odata4.y -= idata4.y;\n odata4.z -= idata4.z;\n odata4.w -= idata4.w;\n return odata4;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Scan kernels\n////////////////////////////////////////////////////////////////////////////////\n__global__ void scanExclusiveShared(uint4 *d_Dst, uint4 *d_Src, uint size)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ uint s_Data[2 * THREADBLOCK_SIZE];\n\n uint pos = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Load data\n uint4 idata4 = d_Src[pos];\n\n // Calculate exclusive scan\n uint4 odata4 = scan4Exclusive(idata4, s_Data, size, cta);\n\n // Write back\n d_Dst[pos] = odata4;\n}\n\n// Exclusive scan of top elements of bottom-level scans (4 * THREADBLOCK_SIZE)\n__global__ void scanExclusiveShared2(uint *d_Buf, uint *d_Dst, uint *d_Src, uint N, uint arrayLength)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ uint s_Data[2 * THREADBLOCK_SIZE];\n\n // Skip loads and stores for inactive threads of last threadblock (pos >= N)\n uint pos = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Load top elements\n // Convert results of bottom-level scan back to inclusive\n uint idata = 0;\n\n if (pos < N)\n idata = d_Dst[(4 * THREADBLOCK_SIZE) - 1 + (4 * THREADBLOCK_SIZE) * pos]\n + d_Src[(4 * THREADBLOCK_SIZE) - 1 + (4 * THREADBLOCK_SIZE) * pos];\n\n // Compute\n uint odata = scan1Exclusive(idata, s_Data, arrayLength, cta);\n\n // Avoid out-of-bound access\n if (pos < N) {\n d_Buf[pos] = odata;\n }\n}\n\n// Final step of large-array scan: combine basic inclusive scan with exclusive\n// scan of top elements of input arrays\n__global__ void uniformUpdate(uint4 *d_Data, uint *d_Buffer)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ uint buf;\n uint pos = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (threadIdx.x == 0) {\n buf = d_Buffer[blockIdx.x];\n }\n\n cg::sync(cta);\n\n uint4 data4 = d_Data[pos];\n data4.x += buf;\n data4.y += buf;\n data4.z += buf;\n data4.w += buf;\n d_Data[pos] = data4;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Interface function\n////////////////////////////////////////////////////////////////////////////////\n// Derived as 32768 (max power-of-two gridDim.x) * 4 * THREADBLOCK_SIZE\n// Due to scanExclusiveShared<<<>>>() 1D block addressing\nextern \"C\" const uint MAX_BATCH_ELEMENTS = 64 * 1048576;\nextern \"C\" const uint MIN_SHORT_ARRAY_SIZE = 4;\nextern \"C\" const uint MAX_SHORT_ARRAY_SIZE = 4 * THREADBLOCK_SIZE;\nextern \"C\" const uint MIN_LARGE_ARRAY_SIZE = 8 * THREADBLOCK_SIZE;\nextern \"C\" const uint MAX_LARGE_ARRAY_SIZE = 4 * THREADBLOCK_SIZE * THREADBLOCK_SIZE;\n\n// Internal exclusive scan buffer\nstatic uint *d_Buf;\n\nextern \"C\" void initScan(void)\n{\n checkCudaErrors(cudaMalloc((void **)&d_Buf, (MAX_BATCH_ELEMENTS / (4 * THREADBLOCK_SIZE)) * sizeof(uint)));\n}\n\nextern \"C\" void closeScan(void) { checkCudaErrors(cudaFree(d_Buf)); }\n\nstatic uint factorRadix2(uint &log2L, uint L)\n{\n if (!L) {\n log2L = 0;\n return 0;\n }\n else {\n for (log2L = 0; (L & 1) == 0; L >>= 1, log2L++)\n ;\n\n return L;\n }\n}\n\nstatic uint iDivUp(uint dividend, uint divisor)\n{\n return ((dividend % divisor) == 0) ? (dividend / divisor) : (dividend / divisor + 1);\n}\n\nextern \"C\" size_t scanExclusiveShort(uint *d_Dst, uint *d_Src, uint batchSize, uint arrayLength)\n{\n // Check power-of-two factorization\n uint log2L;\n uint factorizationRemainder = factorRadix2(log2L, arrayLength);\n assert(factorizationRemainder == 1);\n\n // Check supported size range\n assert((arrayLength >= MIN_SHORT_ARRAY_SIZE) && (arrayLength <= MAX_SHORT_ARRAY_SIZE));\n\n // Check total batch size limit\n assert((batchSize * arrayLength) <= MAX_BATCH_ELEMENTS);\n\n // Check all threadblocks to be fully packed with data\n assert((batchSize * arrayLength) % (4 * THREADBLOCK_SIZE) == 0);\n\n scanExclusiveShared<<<(batchSize * arrayLength) / (4 * THREADBLOCK_SIZE), THREADBLOCK_SIZE>>>(\n (uint4 *)d_Dst, (uint4 *)d_Src, arrayLength);\n getLastCudaError(\"scanExclusiveShared() execution FAILED\\n\");\n\n return THREADBLOCK_SIZE;\n}\n\nextern \"C\" size_t scanExclusiveLarge(uint *d_Dst, uint *d_Src, uint batchSize, uint arrayLength)\n{\n // Check power-of-two factorization\n uint log2L;\n uint factorizationRemainder = factorRadix2(log2L, arrayLength);\n assert(factorizationRemainder == 1);\n\n // Check supported size range\n assert((arrayLength >= MIN_LARGE_ARRAY_SIZE) && (arrayLength <= MAX_LARGE_ARRAY_SIZE));\n\n // Check total batch size limit\n assert((batchSize * arrayLength) <= MAX_BATCH_ELEMENTS);\n\n scanExclusiveShared<<<(batchSize * arrayLength) / (4 * THREADBLOCK_SIZE), THREADBLOCK_SIZE>>>(\n (uint4 *)d_Dst, (uint4 *)d_Src, 4 * THREADBLOCK_SIZE);\n getLastCudaError(\"scanExclusiveShared() execution FAILED\\n\");\n\n // Not all threadblocks need to be packed with input data:\n // inactive threads of highest threadblock just don't do global reads and\n // writes\n const uint blockCount2 = iDivUp((batchSize * arrayLength) / (4 * THREADBLOCK_SIZE), THREADBLOCK_SIZE);\n scanExclusiveShared2<<>>((uint *)d_Buf,\n (uint *)d_Dst,\n (uint *)d_Src,\n (batchSize * arrayLength) / (4 * THREADBLOCK_SIZE),\n arrayLength / (4 * THREADBLOCK_SIZE));\n getLastCudaError(\"scanExclusiveShared2() execution FAILED\\n\");\n\n uniformUpdate<<<(batchSize * arrayLength) / (4 * THREADBLOCK_SIZE), THREADBLOCK_SIZE>>>((uint4 *)d_Dst,\n (uint *)d_Buf);\n getLastCudaError(\"uniformUpdate() execution FAILED\\n\");\n\n return THREADBLOCK_SIZE;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \n\n#include \"scan_common.h\"\n\n// All three kernels run 512 threads per workgroup\n// Must be a power of two\n#define THREADBLOCK_SIZE 256\n\n////////////////////////////////////////////////////////////////////////////////\n// Basic scan codelets\n////////////////////////////////////////////////////////////////////////////////\n// Naive inclusive scan: O(N * log2(N)) operations\n// Allocate 2 * 'size' local memory, initialize the first half\n// with 'size' zeros avoiding if(pos >= offset) condition evaluation\n// and saving instructions\ninline __device__ uint scan1Inclusive(uint idata, volatile uint *s_Data, uint size, cg::thread_block cta)\n{\n uint pos = 2 * threadIdx.x - (threadIdx.x & (size - 1));\n s_Data[pos] = 0;\n pos += size;\n s_Data[pos] = idata;\n\n for (uint offset = 1; offset < size; offset <<= 1) {\n cg::sync(cta);\n uint t = s_Data[pos] + s_Data[pos - offset];\n cg::sync(cta);\n s_Data[pos] = t;\n }\n\n return s_Data[pos];\n}\n\ninline __device__ uint scan1Exclusive(uint idata, volatile uint *s_Data, uint size, cg::thread_block cta)\n{\n return scan1Inclusive(idata, s_Data, size, cta) - idata;\n}\n\ninline __device__ uint4 scan4Inclusive(uint4 idata4, volatile uint *s_Data, uint size, cg::thread_block cta)\n{\n // Level-0 inclusive scan\n idata4.y += idata4.x;\n idata4.z += idata4.y;\n idata4.w += idata4.z;\n\n // Level-1 exclusive scan\n uint oval = scan1Exclusive(idata4.w, s_Data, size / 4, cta);\n\n idata4.x += oval;\n idata4.y += oval;\n idata4.z += oval;\n idata4.w += oval;\n\n return idata4;\n}\n\n// Exclusive vector scan: the array to be scanned is stored\n// in local thread memory scope as uint4\ninline __device__ uint4 scan4Exclusive(uint4 idata4, volatile uint *s_Data, uint size, cg::thread_block cta)\n{\n uint4 odata4 = scan4Inclusive(idata4, s_Data, size, cta);\n odata4.x -= idata4.x;\n odata4.y -= idata4.y;\n odata4.z -= idata4.z;\n odata4.w -= idata4.w;\n return odata4;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Scan kernels\n////////////////////////////////////////////////////////////////////////////////\n__global__ void scanExclusiveShared(uint4 *d_Dst, uint4 *d_Src, uint size)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ uint s_Data[2 * THREADBLOCK_SIZE];\n\n uint pos = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Load data\n uint4 idata4 = d_Src[pos];\n\n // Calculate exclusive scan\n uint4 odata4 = scan4Exclusive(idata4, s_Data, size, cta);\n\n // Write back\n d_Dst[pos] = odata4;\n}\n\n// Exclusive scan of top elements of bottom-level scans (4 * THREADBLOCK_SIZE)\n__global__ void scanExclusiveShared2(uint *d_Buf, uint *d_Dst, uint *d_Src, uint N, uint arrayLength)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ uint s_Data[2 * THREADBLOCK_SIZE];\n\n // Skip loads and stores for inactive threads of last threadblock (pos >= N)\n uint pos = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Load top elements\n // Convert results of bottom-level scan back to inclusive\n uint idata = 0;\n\n if (pos < N)\n idata = d_Dst[(4 * THREADBLOCK_SIZE) - 1 + (4 * THREADBLOCK_SIZE) * pos]\n + d_Src[(4 * THREADBLOCK_SIZE) - 1 + (4 * THREADBLOCK_SIZE) * pos];\n\n // Compute\n uint odata = scan1Exclusive(idata, s_Data, arrayLength, cta);\n\n // Avoid out-of-bound access\n if (pos < N) {\n d_Buf[pos] = odata;\n }\n}\n\n// Final step of large-array scan: combine basic inclusive scan with exclusive\n// scan of top elements of input arrays\n__global__ void uniformUpdate(uint4 *d_Data, uint *d_Buffer)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ uint buf;\n uint pos = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (threadIdx.x == 0) {\n buf = d_Buffer[blockIdx.x];\n }\n\n cg::sync(cta);\n\n uint4 data4 = d_Data[pos];\n data4.x += buf;\n data4.y += buf;\n data4.z += buf;\n data4.w += buf;\n d_Data[pos] = data4;\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Interface function\n////////////////////////////////////////////////////////////////////////////////\n// Derived as 32768 (max power-of-two gridDim.x) * 4 * THREADBLOCK_SIZE\n// Due to scanExclusiveShared<<<>>>() 1D block addressing\nextern \"C\" const uint MAX_BATCH_ELEMENTS = 64 * 1048576;\nextern \"C\" const uint MIN_SHORT_ARRAY_SIZE = 4;\nextern \"C\" const uint MAX_SHORT_ARRAY_SIZE = 4 * THREADBLOCK_SIZE;\nextern \"C\" const uint MIN_LARGE_ARRAY_SIZE = 8 * THREADBLOCK_SIZE;\nextern \"C\" const uint MAX_LARGE_ARRAY_SIZE = 4 * THREADBLOCK_SIZE * THREADBLOCK_SIZE;\n\n// Internal exclusive scan buffer\nstatic uint *d_Buf;\n\nextern \"C\" void initScan(void)\n{\n checkCudaErrors(hipMalloc((void **)&d_Buf, (MAX_BATCH_ELEMENTS / (4 * THREADBLOCK_SIZE)) * sizeof(uint)));\n}\n\nextern \"C\" void closeScan(void) { checkCudaErrors(hipFree(d_Buf)); }\n\nstatic uint factorRadix2(uint &log2L, uint L)\n{\n if (!L) {\n log2L = 0;\n return 0;\n }\n else {\n for (log2L = 0; (L & 1) == 0; L >>= 1, log2L++)\n ;\n\n return L;\n }\n}\n\nstatic uint iDivUp(uint dividend, uint divisor)\n{\n return ((dividend % divisor) == 0) ? (dividend / divisor) : (dividend / divisor + 1);\n}\n\nextern \"C\" size_t scanExclusiveShort(uint *d_Dst, uint *d_Src, uint batchSize, uint arrayLength)\n{\n // Check power-of-two factorization\n uint log2L;\n uint factorizationRemainder = factorRadix2(log2L, arrayLength);\n assert(factorizationRemainder == 1);\n\n // Check supported size range\n assert((arrayLength >= MIN_SHORT_ARRAY_SIZE) && (arrayLength <= MAX_SHORT_ARRAY_SIZE));\n\n // Check total batch size limit\n assert((batchSize * arrayLength) <= MAX_BATCH_ELEMENTS);\n\n // Check all threadblocks to be fully packed with data\n assert((batchSize * arrayLength) % (4 * THREADBLOCK_SIZE) == 0);\n\n scanExclusiveShared<<<(batchSize * arrayLength) / (4 * THREADBLOCK_SIZE), THREADBLOCK_SIZE>>>(\n (uint4 *)d_Dst, (uint4 *)d_Src, arrayLength);\n getLastCudaError(\"scanExclusiveShared() execution FAILED\\n\");\n\n return THREADBLOCK_SIZE;\n}\n\nextern \"C\" size_t scanExclusiveLarge(uint *d_Dst, uint *d_Src, uint batchSize, uint arrayLength)\n{\n // Check power-of-two factorization\n uint log2L;\n uint factorizationRemainder = factorRadix2(log2L, arrayLength);\n assert(factorizationRemainder == 1);\n\n // Check supported size range\n assert((arrayLength >= MIN_LARGE_ARRAY_SIZE) && (arrayLength <= MAX_LARGE_ARRAY_SIZE));\n\n // Check total batch size limit\n assert((batchSize * arrayLength) <= MAX_BATCH_ELEMENTS);\n\n scanExclusiveShared<<<(batchSize * arrayLength) / (4 * THREADBLOCK_SIZE), THREADBLOCK_SIZE>>>(\n (uint4 *)d_Dst, (uint4 *)d_Src, 4 * THREADBLOCK_SIZE);\n getLastCudaError(\"scanExclusiveShared() execution FAILED\\n\");\n\n // Not all threadblocks need to be packed with input data:\n // inactive threads of highest threadblock just don't do global reads and\n // writes\n const uint blockCount2 = iDivUp((batchSize * arrayLength) / (4 * THREADBLOCK_SIZE), THREADBLOCK_SIZE);\n scanExclusiveShared2<<>>((uint *)d_Buf,\n (uint *)d_Dst,\n (uint *)d_Src,\n (batchSize * arrayLength) / (4 * THREADBLOCK_SIZE),\n arrayLength / (4 * THREADBLOCK_SIZE));\n getLastCudaError(\"scanExclusiveShared2() execution FAILED\\n\");\n\n uniformUpdate<<<(batchSize * arrayLength) / (4 * THREADBLOCK_SIZE), THREADBLOCK_SIZE>>>((uint4 *)d_Dst,\n (uint *)d_Buf);\n getLastCudaError(\"uniformUpdate() execution FAILED\\n\");\n\n return THREADBLOCK_SIZE;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/scan/scan.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \n\n#include \"FunctionPointers_kernels.h\"\n\n// Texture object for reading image\ncudaTextureObject_t tex;\nextern __shared__ unsigned char LocalBlock[];\nstatic cudaArray *array = NULL;\n\n#define RADIUS 1\n\n// pixel value used for thresholding function,\n// works well with sample image 'teapot512'\n#define THRESHOLD 150.0f\n\n#ifdef FIXED_BLOCKWIDTH\n#define BlockWidth 80\n#define SharedPitch 384\n#endif\n\n// A function pointer can be declared explicitly like this line:\n//__device__ unsigned char (*pointFunction)(unsigned char, float ) = NULL;\n// or by using typedef's like below:\n\ntypedef unsigned char (*blockFunction_t)(unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n float);\n\ntypedef unsigned char (*pointFunction_t)(unsigned char, float);\n\n__device__ blockFunction_t blockFunction;\n\n__device__ unsigned char ComputeSobel(unsigned char ul, // upper left\n unsigned char um, // upper middle\n unsigned char ur, // upper right\n unsigned char ml, // middle left\n unsigned char mm, // middle (unused)\n unsigned char mr, // middle right\n unsigned char ll, // lower left\n unsigned char lm, // lower middle\n unsigned char lr, // lower right\n float fScale)\n{\n short Horz = ur + 2 * mr + lr - ul - 2 * ml - ll;\n short Vert = ul + 2 * um + ur - ll - 2 * lm - lr;\n short Sum = (short)(fScale * (abs((int)Horz) + abs((int)Vert)));\n return (unsigned char)((Sum < 0) ? 0 : ((Sum > 255) ? 255 : Sum));\n}\n\n// define a function pointer and initialize to NULL\n__device__ unsigned char (*varFunction)(unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n float x) = NULL;\n\n__device__ unsigned char ComputeBox(unsigned char ul, // upper left\n unsigned char um, // upper middle\n unsigned char ur, // upper right\n unsigned char ml, // middle left\n unsigned char mm, // middle...middle\n unsigned char mr, // middle right\n unsigned char ll, // lower left\n unsigned char lm, // lower middle\n unsigned char lr, // lower right\n float fscale)\n{\n short Sum = (short)(ul + um + ur + ml + mm + mr + ll + lm + lr) / 9;\n Sum *= fscale;\n return (unsigned char)((Sum < 0) ? 0 : ((Sum > 255) ? 255 : Sum));\n}\n__device__ unsigned char Threshold(unsigned char in, float thresh)\n{\n if (in > thresh) {\n return 0xFF;\n }\n else {\n return 0;\n }\n}\n\n// Declare function tables, one for the point function chosen, one for the\n// block function chosen. The number of entries is determined by the\n// enum in FunctionPointers_kernels.h\n__device__ blockFunction_t blockFunction_table[LAST_BLOCK_FILTER];\n__device__ pointFunction_t pointFunction_table[LAST_POINT_FILTER];\n\n// Declare device side function pointers. We retrieve them later with\n// cudaMemcpyFromSymbol to set our function tables above in some\n// particular order specified at runtime.\n__device__ blockFunction_t pComputeSobel = ComputeSobel;\n__device__ blockFunction_t pComputeBox = ComputeBox;\n__device__ pointFunction_t pComputeThreshold = Threshold;\n\n// Allocate host side tables to mirror the device side, and later, we\n// fill these tables with the function pointers. This lets us send\n// the pointers to the kernel on invocation, as a method of choosing\n// which function to run.\nblockFunction_t h_blockFunction_table[2];\npointFunction_t h_pointFunction_table[2];\n\n// Perform a filter operation on the data, using shared memory\n// The actual operation performed is\n// determined by the function pointer \"blockFunction\" and selected\n// by the integer argument \"blockOperation\" and has access\n// to an apron around the current pixel being processed.\n// Following the block operation, a per-pixel operation,\n// pointed to by pPointFunction is performed before the final\n// pixel is produced.\n__global__ void SobelShared(uchar4 *pSobelOriginal,\n unsigned short SobelPitch,\n#ifndef FIXED_BLOCKWIDTH\n short BlockWidth,\n short SharedPitch,\n#endif\n short w,\n short h,\n float fScale,\n int blockOperation,\n pointFunction_t pPointFunction,\n cudaTextureObject_t tex)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n short u = 4 * blockIdx.x * BlockWidth;\n short v = blockIdx.y * blockDim.y + threadIdx.y;\n short ib;\n\n int SharedIdx = threadIdx.y * SharedPitch;\n\n for (ib = threadIdx.x; ib < BlockWidth + 2 * RADIUS; ib += blockDim.x) {\n LocalBlock[SharedIdx + 4 * ib + 0] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 0), (float)(v - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 1] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 1), (float)(v - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 2] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 2), (float)(v - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 3] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 3), (float)(v - RADIUS));\n }\n\n if (threadIdx.y < RADIUS * 2) {\n //\n // copy trailing RADIUS*2 rows of pixels into shared\n //\n SharedIdx = (blockDim.y + threadIdx.y) * SharedPitch;\n\n for (ib = threadIdx.x; ib < BlockWidth + 2 * RADIUS; ib += blockDim.x) {\n LocalBlock[SharedIdx + 4 * ib + 0] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 0), (float)(v + blockDim.y - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 1] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 1), (float)(v + blockDim.y - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 2] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 2), (float)(v + blockDim.y - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 3] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 3), (float)(v + blockDim.y - RADIUS));\n }\n }\n\n cg::sync(cta);\n\n u >>= 2; // index as uchar4 from here\n uchar4 *pSobel = (uchar4 *)(((char *)pSobelOriginal) + v * SobelPitch);\n SharedIdx = threadIdx.y * SharedPitch;\n\n blockFunction = blockFunction_table[blockOperation];\n\n for (ib = threadIdx.x; ib < BlockWidth; ib += blockDim.x) {\n uchar4 out;\n\n unsigned char pix00 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 0];\n unsigned char pix01 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 1];\n unsigned char pix02 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 2];\n unsigned char pix10 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 0];\n unsigned char pix11 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 1];\n unsigned char pix12 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 2];\n unsigned char pix20 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 0];\n unsigned char pix21 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 1];\n unsigned char pix22 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 2];\n\n out.x = (*blockFunction)(pix00, pix01, pix02, pix10, pix11, pix12, pix20, pix21, pix22, fScale);\n\n pix00 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 3];\n pix10 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 3];\n pix20 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 3];\n out.y = (*blockFunction)(pix01, pix02, pix00, pix11, pix12, pix10, pix21, pix22, pix20, fScale);\n\n pix01 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 4];\n pix11 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 4];\n pix21 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 4];\n out.z = (*blockFunction)(pix02, pix00, pix01, pix12, pix10, pix11, pix22, pix20, pix21, fScale);\n\n pix02 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 5];\n pix12 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 5];\n pix22 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 5];\n out.w = (*blockFunction)(pix00, pix01, pix02, pix10, pix11, pix12, pix20, pix21, pix22, fScale);\n\n if (pPointFunction != NULL) {\n out.x = (*pPointFunction)(out.x, THRESHOLD);\n out.y = (*pPointFunction)(out.y, THRESHOLD);\n out.z = (*pPointFunction)(out.z, THRESHOLD);\n out.w = (*pPointFunction)(out.w, THRESHOLD);\n }\n\n if (u + ib < w / 4 && v < h) {\n pSobel[u + ib] = out;\n }\n }\n\n cg::sync(cta);\n}\n\n__global__ void\nSobelCopyImage(Pixel *pSobelOriginal, unsigned int Pitch, int w, int h, float fscale, cudaTextureObject_t tex)\n{\n unsigned char *pSobel = (unsigned char *)(((char *)pSobelOriginal) + blockIdx.x * Pitch);\n\n for (int i = threadIdx.x; i < w; i += blockDim.x) {\n pSobel[i] = min(max((tex2D(tex, (float)i, (float)blockIdx.x) * fscale), 0.f), 255.f);\n }\n}\n\n// Perform block and pointer filtering using texture lookups.\n// The block and point operations are determined by the\n// input argument (see comment above for \"SobelShared\" function)\n__global__ void SobelTex(Pixel *pSobelOriginal,\n unsigned int Pitch,\n int w,\n int h,\n float fScale,\n int blockOperation,\n pointFunction_t pPointOperation,\n cudaTextureObject_t tex)\n{\n unsigned char *pSobel = (unsigned char *)(((char *)pSobelOriginal) + blockIdx.x * Pitch);\n unsigned char tmp = 0;\n\n for (int i = threadIdx.x; i < w; i += blockDim.x) {\n unsigned char pix00 = tex2D(tex, (float)i - 1, (float)blockIdx.x - 1);\n unsigned char pix01 = tex2D(tex, (float)i + 0, (float)blockIdx.x - 1);\n unsigned char pix02 = tex2D(tex, (float)i + 1, (float)blockIdx.x - 1);\n unsigned char pix10 = tex2D(tex, (float)i - 1, (float)blockIdx.x + 0);\n unsigned char pix11 = tex2D(tex, (float)i + 0, (float)blockIdx.x + 0);\n unsigned char pix12 = tex2D(tex, (float)i + 1, (float)blockIdx.x + 0);\n unsigned char pix20 = tex2D(tex, (float)i - 1, (float)blockIdx.x + 1);\n unsigned char pix21 = tex2D(tex, (float)i + 0, (float)blockIdx.x + 1);\n unsigned char pix22 = tex2D(tex, (float)i + 1, (float)blockIdx.x + 1);\n tmp = (*(blockFunction_table[blockOperation]))(\n pix00, pix01, pix02, pix10, pix11, pix12, pix20, pix21, pix22, fScale);\n\n if (pPointOperation != NULL) {\n tmp = (*pPointOperation)(tmp, 150.0);\n }\n\n pSobel[i] = tmp;\n }\n}\n\nextern \"C\" void setupTexture(int iw, int ih, Pixel *data, int Bpp)\n{\n cudaChannelFormatDesc desc;\n\n if (Bpp == 1) {\n desc = cudaCreateChannelDesc();\n }\n else {\n desc = cudaCreateChannelDesc();\n }\n\n checkCudaErrors(cudaMallocArray(&array, &desc, iw, ih));\n checkCudaErrors(cudaMemcpy2DToArray(\n array, 0, 0, data, iw * Bpp * sizeof(Pixel), iw * Bpp * sizeof(Pixel), ih, cudaMemcpyHostToDevice));\n\n cudaResourceDesc texRes;\n memset(&texRes, 0, sizeof(cudaResourceDesc));\n\n texRes.resType = cudaResourceTypeArray;\n texRes.res.array.array = array;\n\n cudaTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(cudaTextureDesc));\n\n checkCudaErrors(cudaCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void deleteTexture(void)\n{\n checkCudaErrors(cudaFreeArray(array));\n checkCudaErrors(cudaDestroyTextureObject(tex));\n}\n\n// Copy the pointers from the function tables to the host side\nvoid setupFunctionTables()\n{\n // Dynamically assign the function table.\n // Copy the function pointers to their appropriate locations according to the\n // enum\n checkCudaErrors(cudaMemcpyFromSymbol(&h_blockFunction_table[SOBEL_FILTER], pComputeSobel, sizeof(blockFunction_t)));\n checkCudaErrors(cudaMemcpyFromSymbol(&h_blockFunction_table[BOX_FILTER], pComputeBox, sizeof(blockFunction_t)));\n\n // do the same for the point function, where the 2nd function is NULL (\"no-op\"\n // filter, skipped in kernel code)\n checkCudaErrors(\n cudaMemcpyFromSymbol(&h_pointFunction_table[THRESHOLD_FILTER], pComputeThreshold, sizeof(pointFunction_t)));\n h_pointFunction_table[NULL_FILTER] = NULL;\n\n // now copy the function tables back to the device, so if we wish we can use\n // an index into the table to choose them\n // We have now set the order in the function table according to our enum.\n checkCudaErrors(\n cudaMemcpyToSymbol(blockFunction_table, h_blockFunction_table, sizeof(blockFunction_t) * LAST_BLOCK_FILTER));\n checkCudaErrors(\n cudaMemcpyToSymbol(pointFunction_table, h_pointFunction_table, sizeof(pointFunction_t) * LAST_POINT_FILTER));\n}\n\n// Wrapper for the __global__ call that sets up the texture and threads\n// Below two methods for selecting the image processing function to run are\n// shown.\n// BlockOperation is an integer kernel argument used as an index into the\n// blockFunction_table on the device side\n// pPointOp is itself a function pointer passed as a kernel argument, retrieved\n// from a host side copy of the function table\nextern \"C\" void sobelFilter(Pixel *odata,\n int iw,\n int ih,\n enum SobelDisplayMode mode,\n float fScale,\n int blockOperation,\n int pointOperation)\n{\n pointFunction_t pPointOp = h_pointFunction_table[pointOperation];\n\n switch (mode) {\n case SOBELDISPLAY_IMAGE:\n SobelCopyImage<<>>(odata, iw, iw, ih, fScale, tex);\n break;\n\n case SOBELDISPLAY_SOBELTEX:\n SobelTex<<>>(odata, iw, iw, ih, fScale, blockOperation, pPointOp, tex);\n break;\n\n case SOBELDISPLAY_SOBELSHARED: {\n dim3 threads(16, 4);\n#ifndef FIXED_BLOCKWIDTH\n int BlockWidth = 80; // must be divisible by 16 for coalescing\n#endif\n dim3 blocks =\n dim3(iw / (4 * BlockWidth) + (0 != iw % (4 * BlockWidth)), ih / threads.y + (0 != ih % threads.y));\n int SharedPitch = ~0x3f & (4 * (BlockWidth + 2 * RADIUS) + 0x3f);\n int sharedMem = SharedPitch * (threads.y + 2 * RADIUS);\n\n // for the shared kernel, width must be divisible by 4\n iw &= ~3;\n\n SobelShared<<>>((uchar4 *)odata,\n iw,\n#ifndef FIXED_BLOCKWIDTH\n BlockWidth,\n SharedPitch,\n#endif\n iw,\n ih,\n fScale,\n blockOperation,\n pPointOp,\n tex);\n } break;\n }\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n\n#include \n\n#include \"FunctionPointers_kernels.h\"\n\n// Texture object for reading image\nhipTextureObject_t tex;\nextern __shared__ unsigned char LocalBlock[];\nstatic hipArray *array = NULL;\n\n#define RADIUS 1\n\n// pixel value used for thresholding function,\n// works well with sample image 'teapot512'\n#define THRESHOLD 150.0f\n\n#ifdef FIXED_BLOCKWIDTH\n#define BlockWidth 80\n#define SharedPitch 384\n#endif\n\n// A function pointer can be declared explicitly like this line:\n//__device__ unsigned char (*pointFunction)(unsigned char, float ) = NULL;\n// or by using typedef's like below:\n\ntypedef unsigned char (*blockFunction_t)(unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n float);\n\ntypedef unsigned char (*pointFunction_t)(unsigned char, float);\n\n__device__ blockFunction_t blockFunction;\n\n__device__ unsigned char ComputeSobel(unsigned char ul, // upper left\n unsigned char um, // upper middle\n unsigned char ur, // upper right\n unsigned char ml, // middle left\n unsigned char mm, // middle (unused)\n unsigned char mr, // middle right\n unsigned char ll, // lower left\n unsigned char lm, // lower middle\n unsigned char lr, // lower right\n float fScale)\n{\n short Horz = ur + 2 * mr + lr - ul - 2 * ml - ll;\n short Vert = ul + 2 * um + ur - ll - 2 * lm - lr;\n short Sum = (short)(fScale * (abs((int)Horz) + abs((int)Vert)));\n return (unsigned char)((Sum < 0) ? 0 : ((Sum > 255) ? 255 : Sum));\n}\n\n// define a function pointer and initialize to NULL\n__device__ unsigned char (*varFunction)(unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n unsigned char,\n float x) = NULL;\n\n__device__ unsigned char ComputeBox(unsigned char ul, // upper left\n unsigned char um, // upper middle\n unsigned char ur, // upper right\n unsigned char ml, // middle left\n unsigned char mm, // middle...middle\n unsigned char mr, // middle right\n unsigned char ll, // lower left\n unsigned char lm, // lower middle\n unsigned char lr, // lower right\n float fscale)\n{\n short Sum = (short)(ul + um + ur + ml + mm + mr + ll + lm + lr) / 9;\n Sum *= fscale;\n return (unsigned char)((Sum < 0) ? 0 : ((Sum > 255) ? 255 : Sum));\n}\n__device__ unsigned char Threshold(unsigned char in, float thresh)\n{\n if (in > thresh) {\n return 0xFF;\n }\n else {\n return 0;\n }\n}\n\n// Declare function tables, one for the point function chosen, one for the\n// block function chosen. The number of entries is determined by the\n// enum in FunctionPointers_kernels.h\n__device__ blockFunction_t blockFunction_table[LAST_BLOCK_FILTER];\n__device__ pointFunction_t pointFunction_table[LAST_POINT_FILTER];\n\n// Declare device side function pointers. We retrieve them later with\n// cudaMemcpyFromSymbol to set our function tables above in some\n// particular order specified at runtime.\n__device__ blockFunction_t pComputeSobel = ComputeSobel;\n__device__ blockFunction_t pComputeBox = ComputeBox;\n__device__ pointFunction_t pComputeThreshold = Threshold;\n\n// Allocate host side tables to mirror the device side, and later, we\n// fill these tables with the function pointers. This lets us send\n// the pointers to the kernel on invocation, as a method of choosing\n// which function to run.\nblockFunction_t h_blockFunction_table[2];\npointFunction_t h_pointFunction_table[2];\n\n// Perform a filter operation on the data, using shared memory\n// The actual operation performed is\n// determined by the function pointer \"blockFunction\" and selected\n// by the integer argument \"blockOperation\" and has access\n// to an apron around the current pixel being processed.\n// Following the block operation, a per-pixel operation,\n// pointed to by pPointFunction is performed before the final\n// pixel is produced.\n__global__ void SobelShared(uchar4 *pSobelOriginal,\n unsigned short SobelPitch,\n#ifndef FIXED_BLOCKWIDTH\n short BlockWidth,\n short SharedPitch,\n#endif\n short w,\n short h,\n float fScale,\n int blockOperation,\n pointFunction_t pPointFunction,\n hipTextureObject_t tex)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n short u = 4 * blockIdx.x * BlockWidth;\n short v = blockIdx.y * blockDim.y + threadIdx.y;\n short ib;\n\n int SharedIdx = threadIdx.y * SharedPitch;\n\n for (ib = threadIdx.x; ib < BlockWidth + 2 * RADIUS; ib += blockDim.x) {\n LocalBlock[SharedIdx + 4 * ib + 0] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 0), (float)(v - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 1] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 1), (float)(v - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 2] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 2), (float)(v - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 3] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 3), (float)(v - RADIUS));\n }\n\n if (threadIdx.y < RADIUS * 2) {\n //\n // copy trailing RADIUS*2 rows of pixels into shared\n //\n SharedIdx = (blockDim.y + threadIdx.y) * SharedPitch;\n\n for (ib = threadIdx.x; ib < BlockWidth + 2 * RADIUS; ib += blockDim.x) {\n LocalBlock[SharedIdx + 4 * ib + 0] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 0), (float)(v + blockDim.y - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 1] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 1), (float)(v + blockDim.y - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 2] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 2), (float)(v + blockDim.y - RADIUS));\n LocalBlock[SharedIdx + 4 * ib + 3] =\n tex2D(tex, (float)(u + 4 * ib - RADIUS + 3), (float)(v + blockDim.y - RADIUS));\n }\n }\n\n cg::sync(cta);\n\n u >>= 2; // index as uchar4 from here\n uchar4 *pSobel = (uchar4 *)(((char *)pSobelOriginal) + v * SobelPitch);\n SharedIdx = threadIdx.y * SharedPitch;\n\n blockFunction = blockFunction_table[blockOperation];\n\n for (ib = threadIdx.x; ib < BlockWidth; ib += blockDim.x) {\n uchar4 out;\n\n unsigned char pix00 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 0];\n unsigned char pix01 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 1];\n unsigned char pix02 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 2];\n unsigned char pix10 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 0];\n unsigned char pix11 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 1];\n unsigned char pix12 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 2];\n unsigned char pix20 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 0];\n unsigned char pix21 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 1];\n unsigned char pix22 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 2];\n\n out.x = (*blockFunction)(pix00, pix01, pix02, pix10, pix11, pix12, pix20, pix21, pix22, fScale);\n\n pix00 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 3];\n pix10 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 3];\n pix20 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 3];\n out.y = (*blockFunction)(pix01, pix02, pix00, pix11, pix12, pix10, pix21, pix22, pix20, fScale);\n\n pix01 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 4];\n pix11 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 4];\n pix21 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 4];\n out.z = (*blockFunction)(pix02, pix00, pix01, pix12, pix10, pix11, pix22, pix20, pix21, fScale);\n\n pix02 = LocalBlock[SharedIdx + 4 * ib + 0 * SharedPitch + 5];\n pix12 = LocalBlock[SharedIdx + 4 * ib + 1 * SharedPitch + 5];\n pix22 = LocalBlock[SharedIdx + 4 * ib + 2 * SharedPitch + 5];\n out.w = (*blockFunction)(pix00, pix01, pix02, pix10, pix11, pix12, pix20, pix21, pix22, fScale);\n\n if (pPointFunction != NULL) {\n out.x = (*pPointFunction)(out.x, THRESHOLD);\n out.y = (*pPointFunction)(out.y, THRESHOLD);\n out.z = (*pPointFunction)(out.z, THRESHOLD);\n out.w = (*pPointFunction)(out.w, THRESHOLD);\n }\n\n if (u + ib < w / 4 && v < h) {\n pSobel[u + ib] = out;\n }\n }\n\n cg::sync(cta);\n}\n\n__global__ void\nSobelCopyImage(Pixel *pSobelOriginal, unsigned int Pitch, int w, int h, float fscale, hipTextureObject_t tex)\n{\n unsigned char *pSobel = (unsigned char *)(((char *)pSobelOriginal) + blockIdx.x * Pitch);\n\n for (int i = threadIdx.x; i < w; i += blockDim.x) {\n pSobel[i] = min(max((tex2D(tex, (float)i, (float)blockIdx.x) * fscale), 0.f), 255.f);\n }\n}\n\n// Perform block and pointer filtering using texture lookups.\n// The block and point operations are determined by the\n// input argument (see comment above for \"SobelShared\" function)\n__global__ void SobelTex(Pixel *pSobelOriginal,\n unsigned int Pitch,\n int w,\n int h,\n float fScale,\n int blockOperation,\n pointFunction_t pPointOperation,\n hipTextureObject_t tex)\n{\n unsigned char *pSobel = (unsigned char *)(((char *)pSobelOriginal) + blockIdx.x * Pitch);\n unsigned char tmp = 0;\n\n for (int i = threadIdx.x; i < w; i += blockDim.x) {\n unsigned char pix00 = tex2D(tex, (float)i - 1, (float)blockIdx.x - 1);\n unsigned char pix01 = tex2D(tex, (float)i + 0, (float)blockIdx.x - 1);\n unsigned char pix02 = tex2D(tex, (float)i + 1, (float)blockIdx.x - 1);\n unsigned char pix10 = tex2D(tex, (float)i - 1, (float)blockIdx.x + 0);\n unsigned char pix11 = tex2D(tex, (float)i + 0, (float)blockIdx.x + 0);\n unsigned char pix12 = tex2D(tex, (float)i + 1, (float)blockIdx.x + 0);\n unsigned char pix20 = tex2D(tex, (float)i - 1, (float)blockIdx.x + 1);\n unsigned char pix21 = tex2D(tex, (float)i + 0, (float)blockIdx.x + 1);\n unsigned char pix22 = tex2D(tex, (float)i + 1, (float)blockIdx.x + 1);\n tmp = (*(blockFunction_table[blockOperation]))(\n pix00, pix01, pix02, pix10, pix11, pix12, pix20, pix21, pix22, fScale);\n\n if (pPointOperation != NULL) {\n tmp = (*pPointOperation)(tmp, 150.0);\n }\n\n pSobel[i] = tmp;\n }\n}\n\nextern \"C\" void setupTexture(int iw, int ih, Pixel *data, int Bpp)\n{\n hipChannelFormatDesc desc;\n\n if (Bpp == 1) {\n desc = hipCreateChannelDesc();\n }\n else {\n desc = hipCreateChannelDesc();\n }\n\n checkCudaErrors(hipMallocArray(&array, &desc, iw, ih));\n checkCudaErrors(hipMemcpy2DToArray(\n array, 0, 0, data, iw * Bpp * sizeof(Pixel), iw * Bpp * sizeof(Pixel), ih, hipMemcpyHostToDevice));\n\n hipResourceDesc texRes;\n memset(&texRes, 0, sizeof(hipResourceDesc));\n\n texRes.resType = hipResourceTypeArray;\n texRes.res.array.array = array;\n\n hipTextureDesc texDescr;\n memset(&texDescr, 0, sizeof(hipTextureDesc));\n\n checkCudaErrors(hipCreateTextureObject(&tex, &texRes, &texDescr, NULL));\n}\n\nextern \"C\" void deleteTexture(void)\n{\n checkCudaErrors(hipFreeArray(array));\n checkCudaErrors(hipDestroyTextureObject(tex));\n}\n\n// Copy the pointers from the function tables to the host side\nvoid setupFunctionTables()\n{\n // Dynamically assign the function table.\n // Copy the function pointers to their appropriate locations according to the\n // enum\n checkCudaErrors(hipMemcpyFromSymbol(&h_blockFunction_table[SOBEL_FILTER],HIP_SYMBOL(pComputeSobel), sizeof(blockFunction_t)));\n checkCudaErrors(hipMemcpyFromSymbol(&h_blockFunction_table[BOX_FILTER],HIP_SYMBOL(pComputeBox), sizeof(blockFunction_t)));\n\n // do the same for the point function, where the 2nd function is NULL (\"no-op\"\n // filter, skipped in kernel code)\n checkCudaErrors(\n hipMemcpyFromSymbol(&h_pointFunction_table[THRESHOLD_FILTER],HIP_SYMBOL(pComputeThreshold), sizeof(pointFunction_t)));\n h_pointFunction_table[NULL_FILTER] = NULL;\n\n // now copy the function tables back to the device, so if we wish we can use\n // an index into the table to choose them\n // We have now set the order in the function table according to our enum.\n checkCudaErrors(\n hipMemcpyToSymbol(HIP_SYMBOL(blockFunction_table), h_blockFunction_table, sizeof(blockFunction_t) * LAST_BLOCK_FILTER));\n checkCudaErrors(\n hipMemcpyToSymbol(HIP_SYMBOL(pointFunction_table), h_pointFunction_table, sizeof(pointFunction_t) * LAST_POINT_FILTER));\n}\n\n// Wrapper for the __global__ call that sets up the texture and threads\n// Below two methods for selecting the image processing function to run are\n// shown.\n// BlockOperation is an integer kernel argument used as an index into the\n// blockFunction_table on the device side\n// pPointOp is itself a function pointer passed as a kernel argument, retrieved\n// from a host side copy of the function table\nextern \"C\" void sobelFilter(Pixel *odata,\n int iw,\n int ih,\n enum SobelDisplayMode mode,\n float fScale,\n int blockOperation,\n int pointOperation)\n{\n pointFunction_t pPointOp = h_pointFunction_table[pointOperation];\n\n switch (mode) {\n case SOBELDISPLAY_IMAGE:\n SobelCopyImage<<>>(odata, iw, iw, ih, fScale, tex);\n break;\n\n case SOBELDISPLAY_SOBELTEX:\n SobelTex<<>>(odata, iw, iw, ih, fScale, blockOperation, pPointOp, tex);\n break;\n\n case SOBELDISPLAY_SOBELSHARED: {\n dim3 threads(16, 4);\n#ifndef FIXED_BLOCKWIDTH\n int BlockWidth = 80; // must be divisible by 16 for coalescing\n#endif\n dim3 blocks =\n dim3(iw / (4 * BlockWidth) + (0 != iw % (4 * BlockWidth)), ih / threads.y + (0 != ih % threads.y));\n int SharedPitch = ~0x3f & (4 * (BlockWidth + 2 * RADIUS) + 0x3f);\n int sharedMem = SharedPitch * (threads.y + 2 * RADIUS);\n\n // for the shared kernel, width must be divisible by 4\n iw &= ~3;\n\n SobelShared<<>>((uchar4 *)odata,\n iw,\n#ifndef FIXED_BLOCKWIDTH\n BlockWidth,\n SharedPitch,\n#endif\n iw,\n ih,\n fScale,\n blockOperation,\n pPointOp,\n tex);\n } break;\n }\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/FunctionPointers/FunctionPointers_kernels.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n//\n// DESCRIPTION: Simple CUDA consumer rendering sample app\n//\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"eglstrm_common.h\"\n\nextern bool isCrossDevice;\n\n__device__ static unsigned int numErrors = 0, errorFound = 0;\n__device__ void checkProducerDataGPU(char *data, int size, char expectedVal, int frameNumber)\n{\n if ((data[blockDim.x * blockIdx.x + threadIdx.x] != expectedVal) && (!errorFound)) {\n printf(\"Producer FOUND:%d expected: %d at %d for trial %d %d\\n\",\n data[blockDim.x * blockIdx.x + threadIdx.x],\n expectedVal,\n (blockDim.x * blockIdx.x + threadIdx.x),\n frameNumber,\n numErrors);\n numErrors++;\n errorFound = 1;\n return;\n }\n}\n\n__device__ void checkConsumerDataGPU(char *data, int size, char expectedVal, int frameNumber)\n{\n if ((data[blockDim.x * blockIdx.x + threadIdx.x] != expectedVal) && (!errorFound)) {\n printf(\"Consumer FOUND:%d expected: %d at %d for trial %d %d\\n\",\n data[blockDim.x * blockIdx.x + threadIdx.x],\n expectedVal,\n (blockDim.x * blockIdx.x + threadIdx.x),\n frameNumber,\n numErrors);\n numErrors++;\n errorFound = 1;\n return;\n }\n}\n\n__global__ void writeDataToBuffer(char *pSrc, char newVal) { pSrc[blockDim.x * blockIdx.x + threadIdx.x] = newVal; }\n\n__global__ void testKernelConsumer(char *pSrc, char size, char expectedVal, char newVal, int frameNumber)\n{\n checkConsumerDataGPU(pSrc, size, expectedVal, frameNumber);\n}\n\n__global__ void testKernelProducer(char *pSrc, char size, char expectedVal, char newVal, int frameNumber)\n{\n checkProducerDataGPU(pSrc, size, expectedVal, frameNumber);\n}\n__global__ void getNumErrors(int *numErr) { *numErr = numErrors; }\n\ncudaError_t cudaProducer_filter(cudaStream_t pStream,\n char *pSrc,\n int width,\n int height,\n char expectedVal,\n char newVal,\n int frameNumber)\n{\n // in case where consumer is on dgpu and producer is on igpu when return is\n // called the frame is not copied back to igpu. So the consumer changes is not\n // visible to producer\n if (isCrossDevice == 0) {\n testKernelProducer<<<(width * height) / 1024, 1024, 1, pStream>>>(\n pSrc, width * height, expectedVal, newVal, frameNumber);\n }\n writeDataToBuffer<<<(width * height) / 1024, 1024, 1, pStream>>>(pSrc, newVal);\n return cudaSuccess;\n};\n\ncudaError_t cudaConsumer_filter(cudaStream_t cStream,\n char *pSrc,\n int width,\n int height,\n char expectedVal,\n char newVal,\n int frameNumber)\n{\n testKernelConsumer<<<(width * height) / 1024, 1024, 1, cStream>>>(\n pSrc, width * height, expectedVal, newVal, frameNumber);\n writeDataToBuffer<<<(width * height) / 1024, 1024, 1, cStream>>>(pSrc, newVal);\n return cudaSuccess;\n};\n\ncudaError_t cudaGetValueMismatch()\n{\n int numErr_h;\n int *numErr_d = NULL;\n cudaError_t err = cudaSuccess;\n err = cudaMalloc(&numErr_d, sizeof(int));\n if (err != cudaSuccess) {\n printf(\"Cuda Main: cudaMalloc failed with %s\\n\", cudaGetErrorString(err));\n return err;\n }\n getNumErrors<<<1, 1>>>(numErr_d);\n err = cudaDeviceSynchronize();\n if (err != cudaSuccess) {\n printf(\"Cuda Main: cudaDeviceSynchronize failed with %s\\n\", cudaGetErrorString(err));\n }\n err = cudaMemcpy(&numErr_h, numErr_d, sizeof(int), cudaMemcpyDeviceToHost);\n if (err != cudaSuccess) {\n printf(\"Cuda Main: cudaMemcpy failed with %s\\n\", cudaGetErrorString(err));\n cudaFree(numErr_d);\n return err;\n }\n err = cudaFree(numErr_d);\n if (err != cudaSuccess) {\n printf(\"Cuda Main: cudaFree failed with %s\\n\", cudaGetErrorString(err));\n return err;\n }\n if (numErr_h > 0) {\n return cudaErrorUnknown;\n }\n return cudaSuccess;\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n//\n// DESCRIPTION: Simple CUDA consumer rendering sample app\n//\n\n#include \n#include \n#include \n#include \n#include \n\n#include \"eglstrm_common.h\"\n\nextern bool isCrossDevice;\n\n__device__ static unsigned int numErrors = 0, errorFound = 0;\n__device__ void checkProducerDataGPU(char *data, int size, char expectedVal, int frameNumber)\n{\n if ((data[blockDim.x * blockIdx.x + threadIdx.x] != expectedVal) && (!errorFound)) {\n printf(\"Producer FOUND:%d expected: %d at %d for trial %d %d\\n\",\n data[blockDim.x * blockIdx.x + threadIdx.x],\n expectedVal,\n (blockDim.x * blockIdx.x + threadIdx.x),\n frameNumber,\n numErrors);\n numErrors++;\n errorFound = 1;\n return;\n }\n}\n\n__device__ void checkConsumerDataGPU(char *data, int size, char expectedVal, int frameNumber)\n{\n if ((data[blockDim.x * blockIdx.x + threadIdx.x] != expectedVal) && (!errorFound)) {\n printf(\"Consumer FOUND:%d expected: %d at %d for trial %d %d\\n\",\n data[blockDim.x * blockIdx.x + threadIdx.x],\n expectedVal,\n (blockDim.x * blockIdx.x + threadIdx.x),\n frameNumber,\n numErrors);\n numErrors++;\n errorFound = 1;\n return;\n }\n}\n\n__global__ void writeDataToBuffer(char *pSrc, char newVal) { pSrc[blockDim.x * blockIdx.x + threadIdx.x] = newVal; }\n\n__global__ void testKernelConsumer(char *pSrc, char size, char expectedVal, char newVal, int frameNumber)\n{\n checkConsumerDataGPU(pSrc, size, expectedVal, frameNumber);\n}\n\n__global__ void testKernelProducer(char *pSrc, char size, char expectedVal, char newVal, int frameNumber)\n{\n checkProducerDataGPU(pSrc, size, expectedVal, frameNumber);\n}\n__global__ void getNumErrors(int *numErr) { *numErr = numErrors; }\n\nhipError_t cudaProducer_filter(hipStream_t pStream,\n char *pSrc,\n int width,\n int height,\n char expectedVal,\n char newVal,\n int frameNumber)\n{\n // in case where consumer is on dgpu and producer is on igpu when return is\n // called the frame is not copied back to igpu. So the consumer changes is not\n // visible to producer\n if (isCrossDevice == 0) {\n testKernelProducer<<<(width * height) / 1024, 1024, 1, pStream>>>(\n pSrc, width * height, expectedVal, newVal, frameNumber);\n }\n writeDataToBuffer<<<(width * height) / 1024, 1024, 1, pStream>>>(pSrc, newVal);\n return hipSuccess;\n};\n\nhipError_t cudaConsumer_filter(hipStream_t cStream,\n char *pSrc,\n int width,\n int height,\n char expectedVal,\n char newVal,\n int frameNumber)\n{\n testKernelConsumer<<<(width * height) / 1024, 1024, 1, cStream>>>(\n pSrc, width * height, expectedVal, newVal, frameNumber);\n writeDataToBuffer<<<(width * height) / 1024, 1024, 1, cStream>>>(pSrc, newVal);\n return hipSuccess;\n};\n\nhipError_t cudaGetValueMismatch()\n{\n int numErr_h;\n int *numErr_d = NULL;\n hipError_t err = hipSuccess;\n err = hipMalloc(&numErr_d, sizeof(int));\n if (err != hipSuccess) {\n printf(\"Cuda Main: hipMalloc failed with %s\\n\", hipGetErrorString(err));\n return err;\n }\n getNumErrors<<<1, 1>>>(numErr_d);\n err = hipDeviceSynchronize();\n if (err != hipSuccess) {\n printf(\"Cuda Main: hipDeviceSynchronize failed with %s\\n\", hipGetErrorString(err));\n }\n err = hipMemcpy(&numErr_h, numErr_d, sizeof(int), hipMemcpyDeviceToHost);\n if (err != hipSuccess) {\n printf(\"Cuda Main: hipMemcpy failed with %s\\n\", hipGetErrorString(err));\n hipFree(numErr_d);\n return err;\n }\n err = hipFree(numErr_d);\n if (err != hipSuccess) {\n printf(\"Cuda Main: hipFree failed with %s\\n\", hipGetErrorString(err));\n return err;\n }\n if (numErr_h > 0) {\n return hipErrorUnknown;\n }\n return hipSuccess;\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/EGLStream_CUDA_CrossGPU/kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\nextern \"C\" __global__ void kernelFunction(int *input) { input[threadIdx.x] = 32 - threadIdx.x; }\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\nextern \"C\" __global__ void kernelFunction(int *input) { input[threadIdx.x] = 32 - threadIdx.x; }\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/threadMigration/threadMigration_kernel.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Shuffle intrinsics CUDA Sample\n// This sample demonstrates the use of the shuffle intrinsic\n// First, a simple example of a prefix sum using the shuffle to\n// perform a scan operation is provided.\n// Secondly, a more involved example of computing an integral image\n// using the shuffle intrinsic is provided, where the shuffle\n// scan operation and shuffle xor operations are used\n\n#include \n#include \n#include \n#include \n\n#include \"shfl_integral_image.cuh\"\n\n// Scan using shfl - takes log2(n) steps\n// This function demonstrates basic use of the shuffle intrinsic, __shfl_up,\n// to perform a scan operation across a block.\n// First, it performs a scan (prefix sum in this case) inside a warp\n// Then to continue the scan operation across the block,\n// each warp's sum is placed into shared memory. A single warp\n// then performs a shuffle scan on that shared memory. The results\n// are then uniformly added to each warp's threads.\n// This pyramid type approach is continued by placing each block's\n// final sum in global memory and prefix summing that via another kernel call,\n// then uniformly adding across the input data via the uniform_add<<<>>> kernel.\n\n__global__ void shfl_scan_test(int *data, int width, int *partial_sums = NULL)\n{\n extern __shared__ int sums[];\n int id = ((blockIdx.x * blockDim.x) + threadIdx.x);\n int lane_id = id % warpSize;\n // determine a warp_id within a block\n int warp_id = threadIdx.x / warpSize;\n\n // Below is the basic structure of using a shfl instruction\n // for a scan.\n // Record \"value\" as a variable - we accumulate it along the way\n int value = data[id];\n\n // Now accumulate in log steps up the chain\n // compute sums, with another thread's value who is\n // distance delta away (i). Note\n // those threads where the thread 'i' away would have\n // been out of bounds of the warp are unaffected. This\n // creates the scan sum.\n\n#pragma unroll\n for (int i = 1; i <= width; i *= 2) {\n unsigned int mask = 0xffffffff;\n int n = __shfl_up_sync(mask, value, i, width);\n\n if (lane_id >= i)\n value += n;\n }\n\n // value now holds the scan value for the individual thread\n // next sum the largest values for each warp\n\n // write the sum of the warp to smem\n if (threadIdx.x % warpSize == warpSize - 1) {\n sums[warp_id] = value;\n }\n\n __syncthreads();\n\n //\n // scan sum the warp sums\n // the same shfl scan operation, but performed on warp sums\n //\n if (warp_id == 0 && lane_id < (blockDim.x / warpSize)) {\n int warp_sum = sums[lane_id];\n\n int mask = (1 << (blockDim.x / warpSize)) - 1;\n for (int i = 1; i <= (blockDim.x / warpSize); i *= 2) {\n int n = __shfl_up_sync(mask, warp_sum, i, (blockDim.x / warpSize));\n\n if (lane_id >= i)\n warp_sum += n;\n }\n\n sums[lane_id] = warp_sum;\n }\n\n __syncthreads();\n\n // perform a uniform add across warps in the block\n // read neighbouring warp's sum and add it to threads value\n int blockSum = 0;\n\n if (warp_id > 0) {\n blockSum = sums[warp_id - 1];\n }\n\n value += blockSum;\n\n // Now write out our result\n data[id] = value;\n\n // last thread has sum, write write out the block's sum\n if (partial_sums != NULL && threadIdx.x == blockDim.x - 1) {\n partial_sums[blockIdx.x] = value;\n }\n}\n\n// Uniform add: add partial sums array\n__global__ void uniform_add(int *data, int *partial_sums, int len)\n{\n __shared__ int buf;\n int id = ((blockIdx.x * blockDim.x) + threadIdx.x);\n\n if (id > len)\n return;\n\n if (threadIdx.x == 0) {\n buf = partial_sums[blockIdx.x];\n }\n\n __syncthreads();\n data[id] += buf;\n}\n\nstatic unsigned int iDivUp(unsigned int dividend, unsigned int divisor)\n{\n return ((dividend % divisor) == 0) ? (dividend / divisor) : (dividend / divisor + 1);\n}\n\n// This function verifies the shuffle scan result, for the simple\n// prefix sum case.\nbool CPUverify(int *h_data, int *h_result, int n_elements)\n{\n // cpu verify\n for (int i = 0; i < n_elements - 1; i++) {\n h_data[i + 1] = h_data[i] + h_data[i + 1];\n }\n\n int diff = 0;\n\n for (int i = 0; i < n_elements; i++) {\n diff += h_data[i] - h_result[i];\n }\n\n printf(\"CPU verify result diff (GPUvsCPU) = %d\\n\", diff);\n bool bTestResult = false;\n\n if (diff == 0)\n bTestResult = true;\n\n StopWatchInterface *hTimer = NULL;\n sdkCreateTimer(&hTimer);\n sdkResetTimer(&hTimer);\n sdkStartTimer(&hTimer);\n\n for (int j = 0; j < 100; j++)\n for (int i = 0; i < n_elements - 1; i++) {\n h_data[i + 1] = h_data[i] + h_data[i + 1];\n }\n\n sdkStopTimer(&hTimer);\n double cput = sdkGetTimerValue(&hTimer);\n printf(\"CPU sum (naive) took %f ms\\n\", cput / 100);\n return bTestResult;\n}\n\n// this verifies the row scan result for synthetic data of all 1's\nunsigned int verifyDataRowSums(unsigned int *h_image, int w, int h)\n{\n unsigned int diff = 0;\n\n for (int j = 0; j < h; j++) {\n for (int i = 0; i < w; i++) {\n int gold = i + 1;\n diff += abs(static_cast(gold) - static_cast(h_image[j * w + i]));\n }\n }\n\n return diff;\n}\n\nbool shuffle_simple_test(int argc, char **argv)\n{\n int *h_data, *h_partial_sums, *h_result;\n int *d_data, *d_partial_sums;\n const int n_elements = 65536;\n int sz = sizeof(int) * n_elements;\n int cuda_device = 0;\n\n printf(\"Starting shfl_scan\\n\");\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n cuda_device = findCudaDevice(argc, (const char **)argv);\n\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDevice(&cuda_device));\n\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, cuda_device));\n\n printf(\"> Detected Compute SM %d.%d hardware with %d multi-processors\\n\",\n deviceProp.major,\n deviceProp.minor,\n deviceProp.multiProcessorCount);\n\n // __shfl intrinsic needs SM 3.0 or higher\n if (deviceProp.major < 3) {\n printf(\"> __shfl() intrinsic requires device SM 3.0+\\n\");\n printf(\"> Waiving test.\\n\");\n exit(EXIT_WAIVED);\n }\n\n checkCudaErrors(cudaMallocHost(reinterpret_cast(&h_data), sizeof(int) * n_elements));\n checkCudaErrors(cudaMallocHost(reinterpret_cast(&h_result), sizeof(int) * n_elements));\n\n // initialize data:\n printf(\"Computing Simple Sum test\\n\");\n printf(\"---------------------------------------------------\\n\");\n\n printf(\"Initialize test data [1, 1, 1...]\\n\");\n\n for (int i = 0; i < n_elements; i++) {\n h_data[i] = 1;\n }\n\n int blockSize = 256;\n int gridSize = n_elements / blockSize;\n int nWarps = blockSize / 32;\n int shmem_sz = nWarps * sizeof(int);\n int n_partialSums = n_elements / blockSize;\n int partial_sz = n_partialSums * sizeof(int);\n\n printf(\"Scan summation for %d elements, %d partial sums\\n\", n_elements, n_elements / blockSize);\n\n int p_blockSize = min(n_partialSums, blockSize);\n int p_gridSize = iDivUp(n_partialSums, p_blockSize);\n printf(\"Partial summing %d elements with %d blocks of size %d\\n\", n_partialSums, p_gridSize, p_blockSize);\n\n // initialize a timer\n cudaEvent_t start, stop;\n checkCudaErrors(cudaEventCreate(&start));\n checkCudaErrors(cudaEventCreate(&stop));\n float et = 0;\n float inc = 0;\n\n checkCudaErrors(cudaMalloc(reinterpret_cast(&d_data), sz));\n checkCudaErrors(cudaMalloc(reinterpret_cast(&d_partial_sums), partial_sz));\n checkCudaErrors(cudaMemset(d_partial_sums, 0, partial_sz));\n\n checkCudaErrors(cudaMallocHost(reinterpret_cast(&h_partial_sums), partial_sz));\n checkCudaErrors(cudaMemcpy(d_data, h_data, sz, cudaMemcpyHostToDevice));\n\n checkCudaErrors(cudaEventRecord(start, 0));\n shfl_scan_test<<>>(d_data, 32, d_partial_sums);\n shfl_scan_test<<>>(d_partial_sums, 32);\n uniform_add<<>>(d_data + blockSize, d_partial_sums, n_elements);\n checkCudaErrors(cudaEventRecord(stop, 0));\n checkCudaErrors(cudaEventSynchronize(stop));\n checkCudaErrors(cudaEventElapsedTime(&inc, start, stop));\n et += inc;\n\n checkCudaErrors(cudaMemcpy(h_result, d_data, sz, cudaMemcpyDeviceToHost));\n checkCudaErrors(cudaMemcpy(h_partial_sums, d_partial_sums, partial_sz, cudaMemcpyDeviceToHost));\n\n printf(\"Test Sum: %d\\n\", h_partial_sums[n_partialSums - 1]);\n printf(\"Time (ms): %f\\n\", et);\n printf(\"%d elements scanned in %f ms -> %f MegaElements/s\\n\",\n n_elements,\n et,\n n_elements / (et / 1000.0f) / 1000000.0f);\n\n bool bTestResult = CPUverify(h_data, h_result, n_elements);\n\n checkCudaErrors(cudaFreeHost(h_data));\n checkCudaErrors(cudaFreeHost(h_result));\n checkCudaErrors(cudaFreeHost(h_partial_sums));\n checkCudaErrors(cudaFree(d_data));\n checkCudaErrors(cudaFree(d_partial_sums));\n\n return bTestResult;\n}\n\n// This function tests creation of an integral image using\n// synthetic data, of size 1920x1080 pixels greyscale.\nbool shuffle_integral_image_test()\n{\n char *d_data;\n unsigned int *h_image;\n unsigned int *d_integral_image;\n int w = 1920;\n int h = 1080;\n int n_elements = w * h;\n int sz = sizeof(unsigned int) * n_elements;\n\n printf(\"\\nComputing Integral Image Test on size %d x %d synthetic data\\n\", w, h);\n printf(\"---------------------------------------------------\\n\");\n checkCudaErrors(cudaMallocHost(reinterpret_cast(&h_image), sz));\n // fill test \"image\" with synthetic 1's data\n memset(h_image, 0, sz);\n\n // each thread handles 16 values, use 1 block/row\n int blockSize = iDivUp(w, 16);\n // launch 1 block / row\n int gridSize = h;\n\n // Create a synthetic image for testing\n checkCudaErrors(cudaMalloc(reinterpret_cast(&d_data), sz));\n checkCudaErrors(cudaMalloc(reinterpret_cast(&d_integral_image), n_elements * sizeof(int) * 4));\n checkCudaErrors(cudaMemset(d_data, 1, sz));\n checkCudaErrors(cudaMemset(d_integral_image, 0, sz));\n\n cudaEvent_t start, stop;\n cudaEventCreate(&start);\n cudaEventCreate(&stop);\n float et = 0;\n unsigned int err;\n\n // Execute scan line prefix sum kernel, and time it\n cudaEventRecord(start);\n shfl_intimage_rows<<>>(reinterpret_cast(d_data),\n reinterpret_cast(d_integral_image));\n cudaEventRecord(stop);\n checkCudaErrors(cudaEventSynchronize(stop));\n checkCudaErrors(cudaEventElapsedTime(&et, start, stop));\n printf(\"Method: Fast Time (GPU Timer): %f ms \", et);\n\n // verify the scan line results\n checkCudaErrors(cudaMemcpy(h_image, d_integral_image, sz, cudaMemcpyDeviceToHost));\n err = verifyDataRowSums(h_image, w, h);\n printf(\"Diff = %d\\n\", err);\n\n // Execute column prefix sum kernel and time it\n dim3 blockSz(32, 8);\n dim3 testGrid(w / blockSz.x, 1);\n\n cudaEventRecord(start);\n shfl_vertical_shfl<<>>((unsigned int *)d_integral_image, w, h);\n cudaEventRecord(stop);\n checkCudaErrors(cudaEventSynchronize(stop));\n checkCudaErrors(cudaEventElapsedTime(&et, start, stop));\n printf(\"Method: Vertical Scan Time (GPU Timer): %f ms \", et);\n\n // Verify the column results\n checkCudaErrors(cudaMemcpy(h_image, d_integral_image, sz, cudaMemcpyDeviceToHost));\n printf(\"\\n\");\n\n int finalSum = h_image[w * h - 1];\n printf(\"CheckSum: %d, (expect %dx%d=%d)\\n\", finalSum, w, h, w * h);\n\n checkCudaErrors(cudaFree(d_data));\n checkCudaErrors(cudaFree(d_integral_image));\n checkCudaErrors(cudaFreeHost(h_image));\n // verify final sum: if the final value in the corner is the same as the size\n // of the buffer (all 1's) then the integral image was generated successfully\n return (finalSum == w * h) ? true : false;\n}\n\nint main(int argc, char *argv[])\n{\n // Initialization. The shuffle intrinsic is not available on SM < 3.0\n // so waive the test if the hardware is not present.\n int cuda_device = 0;\n\n printf(\"Starting shfl_scan\\n\");\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n cuda_device = findCudaDevice(argc, (const char **)argv);\n\n cudaDeviceProp deviceProp;\n checkCudaErrors(cudaGetDevice(&cuda_device));\n\n checkCudaErrors(cudaGetDeviceProperties(&deviceProp, cuda_device));\n\n printf(\"> Detected Compute SM %d.%d hardware with %d multi-processors\\n\",\n deviceProp.major,\n deviceProp.minor,\n deviceProp.multiProcessorCount);\n\n // __shfl intrinsic needs SM 3.0 or higher\n if (deviceProp.major < 3) {\n printf(\"> __shfl() intrinsic requires device SM 3.0+\\n\");\n printf(\"> Waiving test.\\n\");\n exit(EXIT_WAIVED);\n }\n\n bool bTestResult = true;\n bool simpleTest = shuffle_simple_test(argc, argv);\n bool intTest = shuffle_integral_image_test();\n\n bTestResult = simpleTest & intTest;\n\n exit((bTestResult) ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n// Shuffle intrinsics CUDA Sample\n// This sample demonstrates the use of the shuffle intrinsic\n// First, a simple example of a prefix sum using the shuffle to\n// perform a scan operation is provided.\n// Secondly, a more involved example of computing an integral image\n// using the shuffle intrinsic is provided, where the shuffle\n// scan operation and shuffle xor operations are used\n\n#include \n#include \n#include \n#include \n\n#include \"shfl_integral_image.cuh\"\n\n// Scan using shfl - takes log2(n) steps\n// This function demonstrates basic use of the shuffle intrinsic, __shfl_up,\n// to perform a scan operation across a block.\n// First, it performs a scan (prefix sum in this case) inside a warp\n// Then to continue the scan operation across the block,\n// each warp's sum is placed into shared memory. A single warp\n// then performs a shuffle scan on that shared memory. The results\n// are then uniformly added to each warp's threads.\n// This pyramid type approach is continued by placing each block's\n// final sum in global memory and prefix summing that via another kernel call,\n// then uniformly adding across the input data via the uniform_add<<<>>> kernel.\n\n__global__ void shfl_scan_test(int *data, int width, int *partial_sums = NULL)\n{\n extern __shared__ int sums[];\n int id = ((blockIdx.x * blockDim.x) + threadIdx.x);\n int lane_id = id % warpSize;\n // determine a warp_id within a block\n int warp_id = threadIdx.x / warpSize;\n\n // Below is the basic structure of using a shfl instruction\n // for a scan.\n // Record \"value\" as a variable - we accumulate it along the way\n int value = data[id];\n\n // Now accumulate in log steps up the chain\n // compute sums, with another thread's value who is\n // distance delta away (i). Note\n // those threads where the thread 'i' away would have\n // been out of bounds of the warp are unaffected. This\n // creates the scan sum.\n\n#pragma unroll\n for (int i = 1; i <= width; i *= 2) {\n unsigned int mask = 0xffffffff;\n int n = __shfl_up_sync(mask, value, i, width);\n\n if (lane_id >= i)\n value += n;\n }\n\n // value now holds the scan value for the individual thread\n // next sum the largest values for each warp\n\n // write the sum of the warp to smem\n if (threadIdx.x % warpSize == warpSize - 1) {\n sums[warp_id] = value;\n }\n\n __syncthreads();\n\n //\n // scan sum the warp sums\n // the same shfl scan operation, but performed on warp sums\n //\n if (warp_id == 0 && lane_id < (blockDim.x / warpSize)) {\n int warp_sum = sums[lane_id];\n\n int mask = (1 << (blockDim.x / warpSize)) - 1;\n for (int i = 1; i <= (blockDim.x / warpSize); i *= 2) {\n int n = __shfl_up_sync(mask, warp_sum, i, (blockDim.x / warpSize));\n\n if (lane_id >= i)\n warp_sum += n;\n }\n\n sums[lane_id] = warp_sum;\n }\n\n __syncthreads();\n\n // perform a uniform add across warps in the block\n // read neighbouring warp's sum and add it to threads value\n int blockSum = 0;\n\n if (warp_id > 0) {\n blockSum = sums[warp_id - 1];\n }\n\n value += blockSum;\n\n // Now write out our result\n data[id] = value;\n\n // last thread has sum, write write out the block's sum\n if (partial_sums != NULL && threadIdx.x == blockDim.x - 1) {\n partial_sums[blockIdx.x] = value;\n }\n}\n\n// Uniform add: add partial sums array\n__global__ void uniform_add(int *data, int *partial_sums, int len)\n{\n __shared__ int buf;\n int id = ((blockIdx.x * blockDim.x) + threadIdx.x);\n\n if (id > len)\n return;\n\n if (threadIdx.x == 0) {\n buf = partial_sums[blockIdx.x];\n }\n\n __syncthreads();\n data[id] += buf;\n}\n\nstatic unsigned int iDivUp(unsigned int dividend, unsigned int divisor)\n{\n return ((dividend % divisor) == 0) ? (dividend / divisor) : (dividend / divisor + 1);\n}\n\n// This function verifies the shuffle scan result, for the simple\n// prefix sum case.\nbool CPUverify(int *h_data, int *h_result, int n_elements)\n{\n // cpu verify\n for (int i = 0; i < n_elements - 1; i++) {\n h_data[i + 1] = h_data[i] + h_data[i + 1];\n }\n\n int diff = 0;\n\n for (int i = 0; i < n_elements; i++) {\n diff += h_data[i] - h_result[i];\n }\n\n printf(\"CPU verify result diff (GPUvsCPU) = %d\\n\", diff);\n bool bTestResult = false;\n\n if (diff == 0)\n bTestResult = true;\n\n StopWatchInterface *hTimer = NULL;\n sdkCreateTimer(&hTimer);\n sdkResetTimer(&hTimer);\n sdkStartTimer(&hTimer);\n\n for (int j = 0; j < 100; j++)\n for (int i = 0; i < n_elements - 1; i++) {\n h_data[i + 1] = h_data[i] + h_data[i + 1];\n }\n\n sdkStopTimer(&hTimer);\n double cput = sdkGetTimerValue(&hTimer);\n printf(\"CPU sum (naive) took %f ms\\n\", cput / 100);\n return bTestResult;\n}\n\n// this verifies the row scan result for synthetic data of all 1's\nunsigned int verifyDataRowSums(unsigned int *h_image, int w, int h)\n{\n unsigned int diff = 0;\n\n for (int j = 0; j < h; j++) {\n for (int i = 0; i < w; i++) {\n int gold = i + 1;\n diff += abs(static_cast(gold) - static_cast(h_image[j * w + i]));\n }\n }\n\n return diff;\n}\n\nbool shuffle_simple_test(int argc, char **argv)\n{\n int *h_data, *h_partial_sums, *h_result;\n int *d_data, *d_partial_sums;\n const int n_elements = 65536;\n int sz = sizeof(int) * n_elements;\n int cuda_device = 0;\n\n printf(\"Starting shfl_scan\\n\");\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n cuda_device = findCudaDevice(argc, (const char **)argv);\n\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDevice(&cuda_device));\n\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, cuda_device));\n\n printf(\"> Detected Compute SM %d.%d hardware with %d multi-processors\\n\",\n deviceProp.major,\n deviceProp.minor,\n deviceProp.multiProcessorCount);\n\n // __shfl intrinsic needs SM 3.0 or higher\n if (deviceProp.major < 3) {\n printf(\"> __shfl() intrinsic requires device SM 3.0+\\n\");\n printf(\"> Waiving test.\\n\");\n exit(EXIT_WAIVED);\n }\n\n checkCudaErrors(hipHostMalloc(reinterpret_cast(&h_data), sizeof(int) * n_elements));\n checkCudaErrors(hipHostMalloc(reinterpret_cast(&h_result), sizeof(int) * n_elements));\n\n // initialize data:\n printf(\"Computing Simple Sum test\\n\");\n printf(\"---------------------------------------------------\\n\");\n\n printf(\"Initialize test data [1, 1, 1...]\\n\");\n\n for (int i = 0; i < n_elements; i++) {\n h_data[i] = 1;\n }\n\n int blockSize = 256;\n int gridSize = n_elements / blockSize;\n int nWarps = blockSize / 32;\n int shmem_sz = nWarps * sizeof(int);\n int n_partialSums = n_elements / blockSize;\n int partial_sz = n_partialSums * sizeof(int);\n\n printf(\"Scan summation for %d elements, %d partial sums\\n\", n_elements, n_elements / blockSize);\n\n int p_blockSize = min(n_partialSums, blockSize);\n int p_gridSize = iDivUp(n_partialSums, p_blockSize);\n printf(\"Partial summing %d elements with %d blocks of size %d\\n\", n_partialSums, p_gridSize, p_blockSize);\n\n // initialize a timer\n hipEvent_t start, stop;\n checkCudaErrors(hipEventCreate(&start));\n checkCudaErrors(hipEventCreate(&stop));\n float et = 0;\n float inc = 0;\n\n checkCudaErrors(hipMalloc(reinterpret_cast(&d_data), sz));\n checkCudaErrors(hipMalloc(reinterpret_cast(&d_partial_sums), partial_sz));\n checkCudaErrors(hipMemset(d_partial_sums, 0, partial_sz));\n\n checkCudaErrors(hipHostMalloc(reinterpret_cast(&h_partial_sums), partial_sz));\n checkCudaErrors(hipMemcpy(d_data, h_data, sz, hipMemcpyHostToDevice));\n\n checkCudaErrors(hipEventRecord(start, 0));\n shfl_scan_test<<>>(d_data, 32, d_partial_sums);\n shfl_scan_test<<>>(d_partial_sums, 32);\n uniform_add<<>>(d_data + blockSize, d_partial_sums, n_elements);\n checkCudaErrors(hipEventRecord(stop, 0));\n checkCudaErrors(hipEventSynchronize(stop));\n checkCudaErrors(hipEventElapsedTime(&inc, start, stop));\n et += inc;\n\n checkCudaErrors(hipMemcpy(h_result, d_data, sz, hipMemcpyDeviceToHost));\n checkCudaErrors(hipMemcpy(h_partial_sums, d_partial_sums, partial_sz, hipMemcpyDeviceToHost));\n\n printf(\"Test Sum: %d\\n\", h_partial_sums[n_partialSums - 1]);\n printf(\"Time (ms): %f\\n\", et);\n printf(\"%d elements scanned in %f ms -> %f MegaElements/s\\n\",\n n_elements,\n et,\n n_elements / (et / 1000.0f) / 1000000.0f);\n\n bool bTestResult = CPUverify(h_data, h_result, n_elements);\n\n checkCudaErrors(hipHostFree(h_data));\n checkCudaErrors(hipHostFree(h_result));\n checkCudaErrors(hipHostFree(h_partial_sums));\n checkCudaErrors(hipFree(d_data));\n checkCudaErrors(hipFree(d_partial_sums));\n\n return bTestResult;\n}\n\n// This function tests creation of an integral image using\n// synthetic data, of size 1920x1080 pixels greyscale.\nbool shuffle_integral_image_test()\n{\n char *d_data;\n unsigned int *h_image;\n unsigned int *d_integral_image;\n int w = 1920;\n int h = 1080;\n int n_elements = w * h;\n int sz = sizeof(unsigned int) * n_elements;\n\n printf(\"\\nComputing Integral Image Test on size %d x %d synthetic data\\n\", w, h);\n printf(\"---------------------------------------------------\\n\");\n checkCudaErrors(hipHostMalloc(reinterpret_cast(&h_image), sz));\n // fill test \"image\" with synthetic 1's data\n memset(h_image, 0, sz);\n\n // each thread handles 16 values, use 1 block/row\n int blockSize = iDivUp(w, 16);\n // launch 1 block / row\n int gridSize = h;\n\n // Create a synthetic image for testing\n checkCudaErrors(hipMalloc(reinterpret_cast(&d_data), sz));\n checkCudaErrors(hipMalloc(reinterpret_cast(&d_integral_image), n_elements * sizeof(int) * 4));\n checkCudaErrors(hipMemset(d_data, 1, sz));\n checkCudaErrors(hipMemset(d_integral_image, 0, sz));\n\n hipEvent_t start, stop;\n hipEventCreate(&start);\n hipEventCreate(&stop);\n float et = 0;\n unsigned int err;\n\n // Execute scan line prefix sum kernel, and time it\n hipEventRecord(start);\n shfl_intimage_rows<<>>(reinterpret_cast(d_data),\n reinterpret_cast(d_integral_image));\n hipEventRecord(stop);\n checkCudaErrors(hipEventSynchronize(stop));\n checkCudaErrors(hipEventElapsedTime(&et, start, stop));\n printf(\"Method: Fast Time (GPU Timer): %f ms \", et);\n\n // verify the scan line results\n checkCudaErrors(hipMemcpy(h_image, d_integral_image, sz, hipMemcpyDeviceToHost));\n err = verifyDataRowSums(h_image, w, h);\n printf(\"Diff = %d\\n\", err);\n\n // Execute column prefix sum kernel and time it\n dim3 blockSz(32, 8);\n dim3 testGrid(w / blockSz.x, 1);\n\n hipEventRecord(start);\n shfl_vertical_shfl<<>>((unsigned int *)d_integral_image, w, h);\n hipEventRecord(stop);\n checkCudaErrors(hipEventSynchronize(stop));\n checkCudaErrors(hipEventElapsedTime(&et, start, stop));\n printf(\"Method: Vertical Scan Time (GPU Timer): %f ms \", et);\n\n // Verify the column results\n checkCudaErrors(hipMemcpy(h_image, d_integral_image, sz, hipMemcpyDeviceToHost));\n printf(\"\\n\");\n\n int finalSum = h_image[w * h - 1];\n printf(\"CheckSum: %d, (expect %dx%d=%d)\\n\", finalSum, w, h, w * h);\n\n checkCudaErrors(hipFree(d_data));\n checkCudaErrors(hipFree(d_integral_image));\n checkCudaErrors(hipHostFree(h_image));\n // verify final sum: if the final value in the corner is the same as the size\n // of the buffer (all 1's) then the integral image was generated successfully\n return (finalSum == w * h) ? true : false;\n}\n\nint main(int argc, char *argv[])\n{\n // Initialization. The shuffle intrinsic is not available on SM < 3.0\n // so waive the test if the hardware is not present.\n int cuda_device = 0;\n\n printf(\"Starting shfl_scan\\n\");\n\n // use command-line specified CUDA device, otherwise use device with highest\n // Gflops/s\n cuda_device = findCudaDevice(argc, (const char **)argv);\n\n hipDeviceProp_t deviceProp;\n checkCudaErrors(hipGetDevice(&cuda_device));\n\n checkCudaErrors(hipGetDeviceProperties(&deviceProp, cuda_device));\n\n printf(\"> Detected Compute SM %d.%d hardware with %d multi-processors\\n\",\n deviceProp.major,\n deviceProp.minor,\n deviceProp.multiProcessorCount);\n\n // __shfl intrinsic needs SM 3.0 or higher\n if (deviceProp.major < 3) {\n printf(\"> __shfl() intrinsic requires device SM 3.0+\\n\");\n printf(\"> Waiving test.\\n\");\n exit(EXIT_WAIVED);\n }\n\n bool bTestResult = true;\n bool simpleTest = shuffle_simple_test(argc, argv);\n bool intTest = shuffle_integral_image_test();\n\n bTestResult = simpleTest & intTest;\n\n exit((bTestResult) ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/shfl_scan/shfl_scan.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"../inc/piestimator.h\"\n\nnamespace cg = cooperative_groups;\n#include \n\nusing std::string;\nusing std::vector;\n\n// RNG init kernel\n__global__ void initRNG(curandState *const rngStates, const unsigned int seed)\n{\n // Determine thread ID\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Initialise the RNG\n curand_init(seed, tid, 0, &rngStates[tid]);\n}\n\n__device__ unsigned int reduce_sum(unsigned int in, cg::thread_block cta)\n{\n extern __shared__ unsigned int sdata[];\n\n // Perform first level of reduction:\n // - Write to shared memory\n unsigned int ltid = threadIdx.x;\n\n sdata[ltid] = in;\n cg::sync(cta);\n\n // Do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (ltid < s) {\n sdata[ltid] += sdata[ltid + s];\n }\n\n cg::sync(cta);\n }\n\n return sdata[0];\n}\n\n__device__ inline void getPoint(float &x, float &y, curandState &state)\n{\n x = curand_uniform(&state);\n y = curand_uniform(&state);\n}\n__device__ inline void getPoint(double &x, double &y, curandState &state)\n{\n x = curand_uniform_double(&state);\n y = curand_uniform_double(&state);\n}\n\n// Estimator kernel\ntemplate \n__global__ void computeValue(unsigned int *const results, curandState *const rngStates, const unsigned int numSims)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Determine thread ID\n unsigned int bid = blockIdx.x;\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int step = gridDim.x * blockDim.x;\n\n // Initialise the RNG\n curandState localState = rngStates[tid];\n\n // Count the number of points which lie inside the unit quarter-circle\n unsigned int pointsInside = 0;\n\n for (unsigned int i = tid; i < numSims; i += step) {\n Real x;\n Real y;\n getPoint(x, y, localState);\n Real l2norm2 = x * x + y * y;\n\n if (l2norm2 < static_cast(1)) {\n pointsInside++;\n }\n }\n\n // Reduce within the block\n pointsInside = reduce_sum(pointsInside, cta);\n\n // Store the result\n if (threadIdx.x == 0) {\n results[bid] = pointsInside;\n }\n}\n\ntemplate \nPiEstimator::PiEstimator(unsigned int numSims,\n unsigned int device,\n unsigned int threadBlockSize,\n unsigned int seed)\n : m_numSims(numSims)\n , m_device(device)\n , m_threadBlockSize(threadBlockSize)\n , m_seed(seed)\n{\n}\n\ntemplate Real PiEstimator::operator()()\n{\n cudaError_t cudaResult = cudaSuccess;\n struct cudaDeviceProp deviceProperties;\n struct cudaFuncAttributes funcAttributes;\n\n // Get device properties\n cudaResult = cudaGetDeviceProperties(&deviceProperties, m_device);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not get device properties: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Check precision is valid\n if (typeid(Real) == typeid(double)\n && (deviceProperties.major < 1 || (deviceProperties.major == 1 && deviceProperties.minor < 3))) {\n throw std::runtime_error(\"Device does not have double precision support\");\n }\n\n // Attach to GPU\n cudaResult = cudaSetDevice(m_device);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not set CUDA device: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Determine how to divide the work between cores\n dim3 block;\n dim3 grid;\n block.x = m_threadBlockSize;\n grid.x = (m_numSims + m_threadBlockSize - 1) / m_threadBlockSize;\n\n // Aim to launch around ten or more times as many blocks as there\n // are multiprocessors on the target device.\n unsigned int blocksPerSM = 10;\n unsigned int numSMs = deviceProperties.multiProcessorCount;\n\n while (grid.x > 2 * blocksPerSM * numSMs) {\n grid.x >>= 1;\n }\n\n // Get initRNG function properties and check the maximum block size\n cudaResult = cudaFuncGetAttributes(&funcAttributes, initRNG);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for initRNG kernel\");\n }\n\n // Get computeValue function properties and check the maximum block size\n cudaResult = cudaFuncGetAttributes(&funcAttributes, computeValue);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for computeValue kernel\");\n }\n\n // Check the dimensions are valid\n if (block.x > (unsigned int)deviceProperties.maxThreadsDim[0]) {\n throw std::runtime_error(\"Block X dimension is too large for device\");\n }\n\n if (grid.x > (unsigned int)deviceProperties.maxGridSize[0]) {\n throw std::runtime_error(\"Grid X dimension is too large for device\");\n }\n\n // Allocate memory for RNG states\n curandState *d_rngStates = 0;\n cudaResult = cudaMalloc((void **)&d_rngStates, grid.x * block.x * sizeof(curandState));\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not allocate memory on device for RNG states: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Allocate memory for result\n // Each thread block will produce one result\n unsigned int *d_results = 0;\n cudaResult = cudaMalloc((void **)&d_results, grid.x * sizeof(unsigned int));\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not allocate memory on device for partial results: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Initialise RNG\n initRNG<<>>(d_rngStates, m_seed);\n\n // Count the points inside unit quarter-circle\n computeValue<<>>(d_results, d_rngStates, m_numSims);\n\n // Copy partial results back\n vector results(grid.x);\n cudaResult = cudaMemcpy(&results[0], d_results, grid.x * sizeof(unsigned int), cudaMemcpyDeviceToHost);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not copy partial results to host: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Complete sum-reduction on host\n Real value = static_cast(std::accumulate(results.begin(), results.end(), 0));\n\n // Determine the proportion of points inside the quarter-circle,\n // i.e. the area of the unit quarter-circle\n value /= m_numSims;\n\n // Value is currently an estimate of the area of a unit quarter-circle, so we\n // can scale to a full circle by multiplying by four. Now since the area of a\n // circle is pi * r^2, and r is one, the value will be an estimate for the\n // value of pi.\n value *= 4;\n\n // Cleanup\n if (d_rngStates) {\n cudaFree(d_rngStates);\n d_rngStates = 0;\n }\n\n if (d_results) {\n cudaFree(d_results);\n d_results = 0;\n }\n\n return value;\n}\n\n// Explicit template instantiation\ntemplate class PiEstimator;\ntemplate class PiEstimator;\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"../inc/piestimator.h\"\n\nnamespace cg = cooperative_groups;\n#include \n\nusing std::string;\nusing std::vector;\n\n// RNG init kernel\n__global__ void initRNG(hiprandState *const rngStates, const unsigned int seed)\n{\n // Determine thread ID\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Initialise the RNG\n hiprand_init(seed, tid, 0, &rngStates[tid]);\n}\n\n__device__ unsigned int reduce_sum(unsigned int in, cg::thread_block cta)\n{\n extern __shared__ unsigned int sdata[];\n\n // Perform first level of reduction:\n // - Write to shared memory\n unsigned int ltid = threadIdx.x;\n\n sdata[ltid] = in;\n cg::sync(cta);\n\n // Do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (ltid < s) {\n sdata[ltid] += sdata[ltid + s];\n }\n\n cg::sync(cta);\n }\n\n return sdata[0];\n}\n\n__device__ inline void getPoint(float &x, float &y, hiprandState &state)\n{\n x = hiprand_uniform(&state);\n y = hiprand_uniform(&state);\n}\n__device__ inline void getPoint(double &x, double &y, hiprandState &state)\n{\n x = hiprand_uniform_double(&state);\n y = hiprand_uniform_double(&state);\n}\n\n// Estimator kernel\ntemplate \n__global__ void computeValue(unsigned int *const results, hiprandState *const rngStates, const unsigned int numSims)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Determine thread ID\n unsigned int bid = blockIdx.x;\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int step = gridDim.x * blockDim.x;\n\n // Initialise the RNG\n hiprandState localState = rngStates[tid];\n\n // Count the number of points which lie inside the unit quarter-circle\n unsigned int pointsInside = 0;\n\n for (unsigned int i = tid; i < numSims; i += step) {\n Real x;\n Real y;\n getPoint(x, y, localState);\n Real l2norm2 = x * x + y * y;\n\n if (l2norm2 < static_cast(1)) {\n pointsInside++;\n }\n }\n\n // Reduce within the block\n pointsInside = reduce_sum(pointsInside, cta);\n\n // Store the result\n if (threadIdx.x == 0) {\n results[bid] = pointsInside;\n }\n}\n\ntemplate \nPiEstimator::PiEstimator(unsigned int numSims,\n unsigned int device,\n unsigned int threadBlockSize,\n unsigned int seed)\n : m_numSims(numSims)\n , m_device(device)\n , m_threadBlockSize(threadBlockSize)\n , m_seed(seed)\n{\n}\n\ntemplate Real PiEstimator::operator()()\n{\n hipError_t cudaResult = hipSuccess;\n struct hipDeviceProp_t deviceProperties;\n struct hipFuncAttributes funcAttributes;\n\n // Get device properties\n cudaResult = hipGetDeviceProperties(&deviceProperties, m_device);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not get device properties: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Check precision is valid\n if (typeid(Real) == typeid(double)\n && (deviceProperties.major < 1 || (deviceProperties.major == 1 && deviceProperties.minor < 3))) {\n throw std::runtime_error(\"Device does not have double precision support\");\n }\n\n // Attach to GPU\n cudaResult = hipSetDevice(m_device);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not set CUDA device: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Determine how to divide the work between cores\n dim3 block;\n dim3 grid;\n block.x = m_threadBlockSize;\n grid.x = (m_numSims + m_threadBlockSize - 1) / m_threadBlockSize;\n\n // Aim to launch around ten or more times as many blocks as there\n // are multiprocessors on the target device.\n unsigned int blocksPerSM = 10;\n unsigned int numSMs = deviceProperties.multiProcessorCount;\n\n while (grid.x > 2 * blocksPerSM * numSMs) {\n grid.x >>= 1;\n }\n\n // Get initRNG function properties and check the maximum block size\n cudaResult = hipFuncGetAttributes(&funcAttributes,reinterpret_cast(initRNG));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for initRNG kernel\");\n }\n\n // Get computeValue function properties and check the maximum block size\n cudaResult = hipFuncGetAttributes(&funcAttributes,reinterpret_cast(computeValue));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for computeValue kernel\");\n }\n\n // Check the dimensions are valid\n if (block.x > (unsigned int)deviceProperties.maxThreadsDim[0]) {\n throw std::runtime_error(\"Block X dimension is too large for device\");\n }\n\n if (grid.x > (unsigned int)deviceProperties.maxGridSize[0]) {\n throw std::runtime_error(\"Grid X dimension is too large for device\");\n }\n\n // Allocate memory for RNG states\n hiprandState *d_rngStates = 0;\n cudaResult = hipMalloc((void **)&d_rngStates, grid.x * block.x * sizeof(hiprandState));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not allocate memory on device for RNG states: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Allocate memory for result\n // Each thread block will produce one result\n unsigned int *d_results = 0;\n cudaResult = hipMalloc((void **)&d_results, grid.x * sizeof(unsigned int));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not allocate memory on device for partial results: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Initialise RNG\n initRNG<<>>(d_rngStates, m_seed);\n\n // Count the points inside unit quarter-circle\n computeValue<<>>(d_results, d_rngStates, m_numSims);\n\n // Copy partial results back\n vector results(grid.x);\n cudaResult = hipMemcpy(&results[0], d_results, grid.x * sizeof(unsigned int), hipMemcpyDeviceToHost);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not copy partial results to host: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Complete sum-reduction on host\n Real value = static_cast(std::accumulate(results.begin(), results.end(), 0));\n\n // Determine the proportion of points inside the quarter-circle,\n // i.e. the area of the unit quarter-circle\n value /= m_numSims;\n\n // Value is currently an estimate of the area of a unit quarter-circle, so we\n // can scale to a full circle by multiplying by four. Now since the area of a\n // circle is pi * r^2, and r is one, the value will be an estimate for the\n // value of pi.\n value *= 4;\n\n // Cleanup\n if (d_rngStates) {\n hipFree(d_rngStates);\n d_rngStates = 0;\n }\n\n if (d_results) {\n hipFree(d_results);\n d_results = 0;\n }\n\n return value;\n}\n\n// Explicit template instantiation\ntemplate class PiEstimator;\ntemplate class PiEstimator;\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/MC_EstimatePiInlineP/src/piestimator.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n\n#include \"convolutionTexture_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// GPU-specific defines\n////////////////////////////////////////////////////////////////////////////////\n// Maps to a single instruction on G8x / G9x / G10x\n#define IMAD(a, b, c) (__mul24((a), (b)) + (c))\n\n// Use unrolled innermost convolution loop\n#define UNROLL_INNER 1\n\n// Round a / b to nearest higher integer value\ninline int iDivUp(int a, int b) { return (a % b != 0) ? (a / b + 1) : (a / b); }\n\n// Align a to nearest higher multiple of b\ninline int iAlignUp(int a, int b) { return (a % b != 0) ? (a - a % b + b) : a; }\n\n////////////////////////////////////////////////////////////////////////////////\n// Convolution kernel and input array storage\n////////////////////////////////////////////////////////////////////////////////\n__constant__ float c_Kernel[KERNEL_LENGTH];\n\nextern \"C\" void setConvolutionKernel(float *h_Kernel)\n{\n cudaMemcpyToSymbol(c_Kernel, h_Kernel, KERNEL_LENGTH * sizeof(float));\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Loop unrolling templates, needed for best performance\n////////////////////////////////////////////////////////////////////////////////\ntemplate __device__ float convolutionRow(float x, float y, cudaTextureObject_t texSrc)\n{\n return tex2D(texSrc, x + (float)(KERNEL_RADIUS - i), y) * c_Kernel[i] + convolutionRow(x, y, texSrc);\n}\n\ntemplate <> __device__ float convolutionRow<-1>(float x, float y, cudaTextureObject_t texSrc) { return 0; }\n\ntemplate __device__ float convolutionColumn(float x, float y, cudaTextureObject_t texSrc)\n{\n return tex2D(texSrc, x, y + (float)(KERNEL_RADIUS - i)) * c_Kernel[i]\n + convolutionColumn(x, y, texSrc);\n}\n\ntemplate <> __device__ float convolutionColumn<-1>(float x, float y, cudaTextureObject_t texSrc) { return 0; }\n\n////////////////////////////////////////////////////////////////////////////////\n// Row convolution filter\n////////////////////////////////////////////////////////////////////////////////\n__global__ void convolutionRowsKernel(float *d_Dst, int imageW, int imageH, cudaTextureObject_t texSrc)\n{\n const int ix = IMAD(blockDim.x, blockIdx.x, threadIdx.x);\n const int iy = IMAD(blockDim.y, blockIdx.y, threadIdx.y);\n const float x = (float)ix + 0.5f;\n const float y = (float)iy + 0.5f;\n\n if (ix >= imageW || iy >= imageH) {\n return;\n }\n\n float sum = 0;\n\n#if (UNROLL_INNER)\n sum = convolutionRow<2 * KERNEL_RADIUS>(x, y, texSrc);\n#else\n\n for (int k = -KERNEL_RADIUS; k <= KERNEL_RADIUS; k++) {\n sum += tex2D(texSrc, x + (float)k, y) * c_Kernel[KERNEL_RADIUS - k];\n }\n\n#endif\n\n d_Dst[IMAD(iy, imageW, ix)] = sum;\n}\n\nextern \"C\" void convolutionRowsGPU(float *d_Dst, cudaArray *a_Src, int imageW, int imageH, cudaTextureObject_t texSrc)\n{\n dim3 threads(16, 12);\n dim3 blocks(iDivUp(imageW, threads.x), iDivUp(imageH, threads.y));\n\n convolutionRowsKernel<<>>(d_Dst, imageW, imageH, texSrc);\n getLastCudaError(\"convolutionRowsKernel() execution failed\\n\");\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Column convolution filter\n////////////////////////////////////////////////////////////////////////////////\n__global__ void convolutionColumnsKernel(float *d_Dst, int imageW, int imageH, cudaTextureObject_t texSrc)\n{\n const int ix = IMAD(blockDim.x, blockIdx.x, threadIdx.x);\n const int iy = IMAD(blockDim.y, blockIdx.y, threadIdx.y);\n const float x = (float)ix + 0.5f;\n const float y = (float)iy + 0.5f;\n\n if (ix >= imageW || iy >= imageH) {\n return;\n }\n\n float sum = 0;\n\n#if (UNROLL_INNER)\n sum = convolutionColumn<2 * KERNEL_RADIUS>(x, y, texSrc);\n#else\n\n for (int k = -KERNEL_RADIUS; k <= KERNEL_RADIUS; k++) {\n sum += tex2D(texSrc, x, y + (float)k) * c_Kernel[KERNEL_RADIUS - k];\n }\n\n#endif\n\n d_Dst[IMAD(iy, imageW, ix)] = sum;\n}\n\nextern \"C\" void\nconvolutionColumnsGPU(float *d_Dst, cudaArray *a_Src, int imageW, int imageH, cudaTextureObject_t texSrc)\n{\n dim3 threads(16, 12);\n dim3 blocks(iDivUp(imageW, threads.x), iDivUp(imageH, threads.y));\n\n convolutionColumnsKernel<<>>(d_Dst, imageW, imageH, texSrc);\n getLastCudaError(\"convolutionColumnsKernel() execution failed\\n\");\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n#include \"convolutionTexture_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// GPU-specific defines\n////////////////////////////////////////////////////////////////////////////////\n// Maps to a single instruction on G8x / G9x / G10x\n#define IMAD(a, b, c) (__mul24((a), (b)) + (c))\n\n// Use unrolled innermost convolution loop\n#define UNROLL_INNER 1\n\n// Round a / b to nearest higher integer value\ninline int iDivUp(int a, int b) { return (a % b != 0) ? (a / b + 1) : (a / b); }\n\n// Align a to nearest higher multiple of b\ninline int iAlignUp(int a, int b) { return (a % b != 0) ? (a - a % b + b) : a; }\n\n////////////////////////////////////////////////////////////////////////////////\n// Convolution kernel and input array storage\n////////////////////////////////////////////////////////////////////////////////\n__constant__ float c_Kernel[KERNEL_LENGTH];\n\nextern \"C\" void setConvolutionKernel(float *h_Kernel)\n{\n hipMemcpyToSymbol(HIP_SYMBOL(c_Kernel), h_Kernel, KERNEL_LENGTH * sizeof(float));\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Loop unrolling templates, needed for best performance\n////////////////////////////////////////////////////////////////////////////////\ntemplate __device__ float convolutionRow(float x, float y, hipTextureObject_t texSrc)\n{\n return tex2D(texSrc, x + (float)(KERNEL_RADIUS - i), y) * c_Kernel[i] + convolutionRow(x, y, texSrc);\n}\n\ntemplate <> __device__ float convolutionRow<-1>(float x, float y, hipTextureObject_t texSrc) { return 0; }\n\ntemplate __device__ float convolutionColumn(float x, float y, hipTextureObject_t texSrc)\n{\n return tex2D(texSrc, x, y + (float)(KERNEL_RADIUS - i)) * c_Kernel[i]\n + convolutionColumn(x, y, texSrc);\n}\n\ntemplate <> __device__ float convolutionColumn<-1>(float x, float y, hipTextureObject_t texSrc) { return 0; }\n\n////////////////////////////////////////////////////////////////////////////////\n// Row convolution filter\n////////////////////////////////////////////////////////////////////////////////\n__global__ void convolutionRowsKernel(float *d_Dst, int imageW, int imageH, hipTextureObject_t texSrc)\n{\n const int ix = IMAD(blockDim.x, blockIdx.x, threadIdx.x);\n const int iy = IMAD(blockDim.y, blockIdx.y, threadIdx.y);\n const float x = (float)ix + 0.5f;\n const float y = (float)iy + 0.5f;\n\n if (ix >= imageW || iy >= imageH) {\n return;\n }\n\n float sum = 0;\n\n#if (UNROLL_INNER)\n sum = convolutionRow<2 * KERNEL_RADIUS>(x, y, texSrc);\n#else\n\n for (int k = -KERNEL_RADIUS; k <= KERNEL_RADIUS; k++) {\n sum += tex2D(texSrc, x + (float)k, y) * c_Kernel[KERNEL_RADIUS - k];\n }\n\n#endif\n\n d_Dst[IMAD(iy, imageW, ix)] = sum;\n}\n\nextern \"C\" void convolutionRowsGPU(float *d_Dst, hipArray *a_Src, int imageW, int imageH, hipTextureObject_t texSrc)\n{\n dim3 threads(16, 12);\n dim3 blocks(iDivUp(imageW, threads.x), iDivUp(imageH, threads.y));\n\n convolutionRowsKernel<<>>(d_Dst, imageW, imageH, texSrc);\n getLastCudaError(\"convolutionRowsKernel() execution failed\\n\");\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Column convolution filter\n////////////////////////////////////////////////////////////////////////////////\n__global__ void convolutionColumnsKernel(float *d_Dst, int imageW, int imageH, hipTextureObject_t texSrc)\n{\n const int ix = IMAD(blockDim.x, blockIdx.x, threadIdx.x);\n const int iy = IMAD(blockDim.y, blockIdx.y, threadIdx.y);\n const float x = (float)ix + 0.5f;\n const float y = (float)iy + 0.5f;\n\n if (ix >= imageW || iy >= imageH) {\n return;\n }\n\n float sum = 0;\n\n#if (UNROLL_INNER)\n sum = convolutionColumn<2 * KERNEL_RADIUS>(x, y, texSrc);\n#else\n\n for (int k = -KERNEL_RADIUS; k <= KERNEL_RADIUS; k++) {\n sum += tex2D(texSrc, x, y + (float)k) * c_Kernel[KERNEL_RADIUS - k];\n }\n\n#endif\n\n d_Dst[IMAD(iy, imageW, ix)] = sum;\n}\n\nextern \"C\" void\nconvolutionColumnsGPU(float *d_Dst, hipArray *a_Src, int imageW, int imageH, hipTextureObject_t texSrc)\n{\n dim3 threads(16, 12);\n dim3 blocks(iDivUp(imageW, threads.x), iDivUp(imageH, threads.y));\n\n convolutionColumnsKernel<<>>(d_Dst, imageW, imageH, texSrc);\n getLastCudaError(\"convolutionColumnsKernel() execution failed\\n\");\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/convolutionTexture/convolutionTexture.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates stream ordered memory allocation on a GPU using\n * cudaMallocAsync and cudaMemPool family of APIs.\n *\n * basicStreamOrderedAllocation(): demonstrates stream ordered allocation using\n * cudaMallocAsync/cudaFreeAsync APIs with default settings.\n *\n * streamOrderedAllocationPostSync(): demonstrates if there's a synchronization\n * in between allocations, then setting the release threshold on the pool will\n * make sure the synchronize will not free memory back to the OS.\n */\n\n// System includes\n#include \n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#define MAX_ITER 20\n\n/* Add two vectors on the GPU */\n__global__ void vectorAddGPU(const float *a, const float *b, float *c, int N)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx < N) {\n c[idx] = a[idx] + b[idx];\n }\n}\n\nint basicStreamOrderedAllocation(const int dev, const int nelem, const float *a, const float *b, float *c)\n{\n float *d_a, *d_b, *d_c; // Device buffers\n float errorNorm, refNorm, ref, diff;\n size_t bytes = nelem * sizeof(float);\n\n cudaStream_t stream;\n printf(\"Starting basicStreamOrderedAllocation()\\n\");\n checkCudaErrors(cudaSetDevice(dev));\n checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));\n\n checkCudaErrors(cudaMallocAsync(&d_a, bytes, stream));\n checkCudaErrors(cudaMallocAsync(&d_b, bytes, stream));\n checkCudaErrors(cudaMallocAsync(&d_c, bytes, stream));\n checkCudaErrors(cudaMemcpyAsync(d_a, a, bytes, cudaMemcpyHostToDevice, stream));\n checkCudaErrors(cudaMemcpyAsync(d_b, b, bytes, cudaMemcpyHostToDevice, stream));\n\n dim3 block(256);\n dim3 grid((unsigned int)ceil(nelem / (float)block.x));\n vectorAddGPU<<>>(d_a, d_b, d_c, nelem);\n\n checkCudaErrors(cudaFreeAsync(d_a, stream));\n checkCudaErrors(cudaFreeAsync(d_b, stream));\n checkCudaErrors(cudaMemcpyAsync(c, d_c, bytes, cudaMemcpyDeviceToHost, stream));\n checkCudaErrors(cudaFreeAsync(d_c, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n\n /* Compare the results */\n printf(\"> Checking the results from vectorAddGPU() ...\\n\");\n errorNorm = 0.f;\n refNorm = 0.f;\n\n for (int n = 0; n < nelem; n++) {\n ref = a[n] + b[n];\n diff = c[n] - ref;\n errorNorm += diff * diff;\n refNorm += ref * ref;\n }\n\n errorNorm = (float)sqrt((double)errorNorm);\n refNorm = (float)sqrt((double)refNorm);\n if (errorNorm / refNorm < 1.e-6f)\n printf(\"basicStreamOrderedAllocation PASSED\\n\");\n\n checkCudaErrors(cudaStreamDestroy(stream));\n\n return errorNorm / refNorm < 1.e-6f ? EXIT_SUCCESS : EXIT_FAILURE;\n}\n\n// streamOrderedAllocationPostSync(): demonstrates If the application wants the\n// memory to persist in the pool beyond synchronization, then it sets the\n// release threshold on the pool. This way, when the application reaches the\n// \"steady state\", it is no longer allocating/freeing memory from the OS.\nint streamOrderedAllocationPostSync(const int dev, const int nelem, const float *a, const float *b, float *c)\n{\n float *d_a, *d_b, *d_c; // Device buffers\n float errorNorm, refNorm, ref, diff;\n size_t bytes = nelem * sizeof(float);\n\n cudaStream_t stream;\n cudaMemPool_t memPool;\n cudaEvent_t start, end;\n printf(\"Starting streamOrderedAllocationPostSync()\\n\");\n checkCudaErrors(cudaSetDevice(dev));\n checkCudaErrors(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));\n checkCudaErrors(cudaEventCreate(&start));\n checkCudaErrors(cudaEventCreate(&end));\n\n checkCudaErrors(cudaDeviceGetDefaultMemPool(&memPool, dev));\n uint64_t thresholdVal = ULONG_MAX;\n // set high release threshold on the default pool so that cudaFreeAsync will\n // not actually release memory to the system. By default, the release\n // threshold for a memory pool is set to zero. This implies that the CUDA\n // driver is allowed to release a memory chunk back to the system as long as\n // it does not contain any active suballocations.\n checkCudaErrors(cudaMemPoolSetAttribute(memPool, cudaMemPoolAttrReleaseThreshold, (void *)&thresholdVal));\n\n // Record the start event\n checkCudaErrors(cudaEventRecord(start, stream));\n for (int i = 0; i < MAX_ITER; i++) {\n checkCudaErrors(cudaMallocAsync(&d_a, bytes, stream));\n checkCudaErrors(cudaMallocAsync(&d_b, bytes, stream));\n checkCudaErrors(cudaMallocAsync(&d_c, bytes, stream));\n checkCudaErrors(cudaMemcpyAsync(d_a, a, bytes, cudaMemcpyHostToDevice, stream));\n checkCudaErrors(cudaMemcpyAsync(d_b, b, bytes, cudaMemcpyHostToDevice, stream));\n\n dim3 block(256);\n dim3 grid((unsigned int)ceil(nelem / (float)block.x));\n vectorAddGPU<<>>(d_a, d_b, d_c, nelem);\n\n checkCudaErrors(cudaFreeAsync(d_a, stream));\n checkCudaErrors(cudaFreeAsync(d_b, stream));\n checkCudaErrors(cudaMemcpyAsync(c, d_c, bytes, cudaMemcpyDeviceToHost, stream));\n checkCudaErrors(cudaFreeAsync(d_c, stream));\n checkCudaErrors(cudaStreamSynchronize(stream));\n }\n checkCudaErrors(cudaEventRecord(end, stream));\n // Wait for the end event to complete\n checkCudaErrors(cudaEventSynchronize(end));\n\n float msecTotal = 0.0f;\n checkCudaErrors(cudaEventElapsedTime(&msecTotal, start, end));\n printf(\"Total elapsed time = %f ms over %d iterations\\n\", msecTotal, MAX_ITER);\n\n /* Compare the results */\n printf(\"> Checking the results from vectorAddGPU() ...\\n\");\n errorNorm = 0.f;\n refNorm = 0.f;\n\n for (int n = 0; n < nelem; n++) {\n ref = a[n] + b[n];\n diff = c[n] - ref;\n errorNorm += diff * diff;\n refNorm += ref * ref;\n }\n\n errorNorm = (float)sqrt((double)errorNorm);\n refNorm = (float)sqrt((double)refNorm);\n if (errorNorm / refNorm < 1.e-6f)\n printf(\"streamOrderedAllocationPostSync PASSED\\n\");\n\n checkCudaErrors(cudaStreamDestroy(stream));\n\n return errorNorm / refNorm < 1.e-6f ? EXIT_SUCCESS : EXIT_FAILURE;\n}\n\nint main(int argc, char **argv)\n{\n int nelem;\n int dev = 0; // use default device 0\n size_t bytes;\n float *a, *b, *c; // Host\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\")) {\n printf(\"Usage: streamOrderedAllocation [OPTION]\\n\\n\");\n printf(\"Options:\\n\");\n printf(\" --device=[device #] Specify the device to be used\\n\");\n return EXIT_SUCCESS;\n }\n\n dev = findCudaDevice(argc, (const char **)argv);\n\n int isMemPoolSupported = 0;\n checkCudaErrors(cudaDeviceGetAttribute(&isMemPoolSupported, cudaDevAttrMemoryPoolsSupported, dev));\n if (!isMemPoolSupported) {\n printf(\"Waiving execution as device does not support Memory Pools\\n\");\n exit(EXIT_WAIVED);\n }\n\n // Allocate CPU memory.\n nelem = 1048576;\n bytes = nelem * sizeof(float);\n\n a = (float *)malloc(bytes);\n b = (float *)malloc(bytes);\n c = (float *)malloc(bytes);\n /* Initialize the vectors. */\n for (int n = 0; n < nelem; n++) {\n a[n] = rand() / (float)RAND_MAX;\n b[n] = rand() / (float)RAND_MAX;\n }\n\n int ret1 = basicStreamOrderedAllocation(dev, nelem, a, b, c);\n int ret2 = streamOrderedAllocationPostSync(dev, nelem, a, b, c);\n\n /* Memory clean up */\n free(a);\n free(b);\n free(c);\n\n return ((ret1 == EXIT_SUCCESS && ret2 == EXIT_SUCCESS) ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n/*\n * This sample demonstrates stream ordered memory allocation on a GPU using\n * cudaMallocAsync and cudaMemPool family of APIs.\n *\n * basicStreamOrderedAllocation(): demonstrates stream ordered allocation using\n * cudaMallocAsync/cudaFreeAsync APIs with default settings.\n *\n * streamOrderedAllocationPostSync(): demonstrates if there's a synchronization\n * in between allocations, then setting the release threshold on the pool will\n * make sure the synchronize will not free memory back to the OS.\n */\n\n// System includes\n#include \n#include \n#include \n\n// CUDA runtime\n#include \n\n// helper functions and utilities to work with CUDA\n#include \n#include \n\n#define MAX_ITER 20\n\n/* Add two vectors on the GPU */\n__global__ void vectorAddGPU(const float *a, const float *b, float *c, int N)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx < N) {\n c[idx] = a[idx] + b[idx];\n }\n}\n\nint basicStreamOrderedAllocation(const int dev, const int nelem, const float *a, const float *b, float *c)\n{\n float *d_a, *d_b, *d_c; // Device buffers\n float errorNorm, refNorm, ref, diff;\n size_t bytes = nelem * sizeof(float);\n\n hipStream_t stream;\n printf(\"Starting basicStreamOrderedAllocation()\\n\");\n checkCudaErrors(hipSetDevice(dev));\n checkCudaErrors(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));\n\n checkCudaErrors(hipMallocAsync(&d_a, bytes, stream));\n checkCudaErrors(hipMallocAsync(&d_b, bytes, stream));\n checkCudaErrors(hipMallocAsync(&d_c, bytes, stream));\n checkCudaErrors(hipMemcpyAsync(d_a, a, bytes, hipMemcpyHostToDevice, stream));\n checkCudaErrors(hipMemcpyAsync(d_b, b, bytes, hipMemcpyHostToDevice, stream));\n\n dim3 block(256);\n dim3 grid((unsigned int)ceil(nelem / (float)block.x));\n vectorAddGPU<<>>(d_a, d_b, d_c, nelem);\n\n checkCudaErrors(hipFreeAsync(d_a, stream));\n checkCudaErrors(hipFreeAsync(d_b, stream));\n checkCudaErrors(hipMemcpyAsync(c, d_c, bytes, hipMemcpyDeviceToHost, stream));\n checkCudaErrors(hipFreeAsync(d_c, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n\n /* Compare the results */\n printf(\"> Checking the results from vectorAddGPU() ...\\n\");\n errorNorm = 0.f;\n refNorm = 0.f;\n\n for (int n = 0; n < nelem; n++) {\n ref = a[n] + b[n];\n diff = c[n] - ref;\n errorNorm += diff * diff;\n refNorm += ref * ref;\n }\n\n errorNorm = (float)sqrt((double)errorNorm);\n refNorm = (float)sqrt((double)refNorm);\n if (errorNorm / refNorm < 1.e-6f)\n printf(\"basicStreamOrderedAllocation PASSED\\n\");\n\n checkCudaErrors(hipStreamDestroy(stream));\n\n return errorNorm / refNorm < 1.e-6f ? EXIT_SUCCESS : EXIT_FAILURE;\n}\n\n// streamOrderedAllocationPostSync(): demonstrates If the application wants the\n// memory to persist in the pool beyond synchronization, then it sets the\n// release threshold on the pool. This way, when the application reaches the\n// \"steady state\", it is no longer allocating/freeing memory from the OS.\nint streamOrderedAllocationPostSync(const int dev, const int nelem, const float *a, const float *b, float *c)\n{\n float *d_a, *d_b, *d_c; // Device buffers\n float errorNorm, refNorm, ref, diff;\n size_t bytes = nelem * sizeof(float);\n\n hipStream_t stream;\n hipMemPool_t memPool;\n hipEvent_t start, end;\n printf(\"Starting streamOrderedAllocationPostSync()\\n\");\n checkCudaErrors(hipSetDevice(dev));\n checkCudaErrors(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));\n checkCudaErrors(hipEventCreate(&start));\n checkCudaErrors(hipEventCreate(&end));\n\n checkCudaErrors(hipDeviceGetDefaultMemPool(&memPool, dev));\n uint64_t thresholdVal = ULONG_MAX;\n // set high release threshold on the default pool so that cudaFreeAsync will\n // not actually release memory to the system. By default, the release\n // threshold for a memory pool is set to zero. This implies that the CUDA\n // driver is allowed to release a memory chunk back to the system as long as\n // it does not contain any active suballocations.\n checkCudaErrors(hipMemPoolSetAttribute(memPool, hipMemPoolAttrReleaseThreshold, (void *)&thresholdVal));\n\n // Record the start event\n checkCudaErrors(hipEventRecord(start, stream));\n for (int i = 0; i < MAX_ITER; i++) {\n checkCudaErrors(hipMallocAsync(&d_a, bytes, stream));\n checkCudaErrors(hipMallocAsync(&d_b, bytes, stream));\n checkCudaErrors(hipMallocAsync(&d_c, bytes, stream));\n checkCudaErrors(hipMemcpyAsync(d_a, a, bytes, hipMemcpyHostToDevice, stream));\n checkCudaErrors(hipMemcpyAsync(d_b, b, bytes, hipMemcpyHostToDevice, stream));\n\n dim3 block(256);\n dim3 grid((unsigned int)ceil(nelem / (float)block.x));\n vectorAddGPU<<>>(d_a, d_b, d_c, nelem);\n\n checkCudaErrors(hipFreeAsync(d_a, stream));\n checkCudaErrors(hipFreeAsync(d_b, stream));\n checkCudaErrors(hipMemcpyAsync(c, d_c, bytes, hipMemcpyDeviceToHost, stream));\n checkCudaErrors(hipFreeAsync(d_c, stream));\n checkCudaErrors(hipStreamSynchronize(stream));\n }\n checkCudaErrors(hipEventRecord(end, stream));\n // Wait for the end event to complete\n checkCudaErrors(hipEventSynchronize(end));\n\n float msecTotal = 0.0f;\n checkCudaErrors(hipEventElapsedTime(&msecTotal, start, end));\n printf(\"Total elapsed time = %f ms over %d iterations\\n\", msecTotal, MAX_ITER);\n\n /* Compare the results */\n printf(\"> Checking the results from vectorAddGPU() ...\\n\");\n errorNorm = 0.f;\n refNorm = 0.f;\n\n for (int n = 0; n < nelem; n++) {\n ref = a[n] + b[n];\n diff = c[n] - ref;\n errorNorm += diff * diff;\n refNorm += ref * ref;\n }\n\n errorNorm = (float)sqrt((double)errorNorm);\n refNorm = (float)sqrt((double)refNorm);\n if (errorNorm / refNorm < 1.e-6f)\n printf(\"streamOrderedAllocationPostSync PASSED\\n\");\n\n checkCudaErrors(hipStreamDestroy(stream));\n\n return errorNorm / refNorm < 1.e-6f ? EXIT_SUCCESS : EXIT_FAILURE;\n}\n\nint main(int argc, char **argv)\n{\n int nelem;\n int dev = 0; // use default device 0\n size_t bytes;\n float *a, *b, *c; // Host\n\n if (checkCmdLineFlag(argc, (const char **)argv, \"help\")) {\n printf(\"Usage: streamOrderedAllocation [OPTION]\\n\\n\");\n printf(\"Options:\\n\");\n printf(\" --device=[device #] Specify the device to be used\\n\");\n return EXIT_SUCCESS;\n }\n\n dev = findCudaDevice(argc, (const char **)argv);\n\n int isMemPoolSupported = 0;\n checkCudaErrors(hipDeviceGetAttribute(&isMemPoolSupported, hipDeviceAttributeMemoryPoolsSupported, dev));\n if (!isMemPoolSupported) {\n printf(\"Waiving execution as device does not support Memory Pools\\n\");\n exit(EXIT_WAIVED);\n }\n\n // Allocate CPU memory.\n nelem = 1048576;\n bytes = nelem * sizeof(float);\n\n a = (float *)malloc(bytes);\n b = (float *)malloc(bytes);\n c = (float *)malloc(bytes);\n /* Initialize the vectors. */\n for (int n = 0; n < nelem; n++) {\n a[n] = rand() / (float)RAND_MAX;\n b[n] = rand() / (float)RAND_MAX;\n }\n\n int ret1 = basicStreamOrderedAllocation(dev, nelem, a, b, c);\n int ret2 = streamOrderedAllocationPostSync(dev, nelem, a, b, c);\n\n /* Memory clean up */\n free(a);\n free(b);\n free(c);\n\n return ((ret1 == EXIT_SUCCESS && ret2 == EXIT_SUCCESS) ? EXIT_SUCCESS : EXIT_FAILURE);\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/streamOrderedAllocation/streamOrderedAllocation.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \"convolutionSeparable_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Convolution kernel storage\n////////////////////////////////////////////////////////////////////////////////\n__constant__ float c_Kernel[KERNEL_LENGTH];\n\nextern \"C\" void setConvolutionKernel(float *h_Kernel)\n{\n cudaMemcpyToSymbol(c_Kernel, h_Kernel, KERNEL_LENGTH * sizeof(float));\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Row convolution filter\n////////////////////////////////////////////////////////////////////////////////\n#define ROWS_BLOCKDIM_X 16\n#define ROWS_BLOCKDIM_Y 4\n#define ROWS_RESULT_STEPS 8\n#define ROWS_HALO_STEPS 1\n\n__global__ void convolutionRowsKernel(float *d_Dst, float *d_Src, int imageW, int imageH, int pitch)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float s_Data[ROWS_BLOCKDIM_Y][(ROWS_RESULT_STEPS + 2 * ROWS_HALO_STEPS) * ROWS_BLOCKDIM_X];\n\n // Offset to the left halo edge\n const int baseX = (blockIdx.x * ROWS_RESULT_STEPS - ROWS_HALO_STEPS) * ROWS_BLOCKDIM_X + threadIdx.x;\n const int baseY = blockIdx.y * ROWS_BLOCKDIM_Y + threadIdx.y;\n\n d_Src += baseY * pitch + baseX;\n d_Dst += baseY * pitch + baseX;\n\n// Load main data\n#pragma unroll\n\n for (int i = ROWS_HALO_STEPS; i < ROWS_HALO_STEPS + ROWS_RESULT_STEPS; i++) {\n s_Data[threadIdx.y][threadIdx.x + i * ROWS_BLOCKDIM_X] = d_Src[i * ROWS_BLOCKDIM_X];\n }\n\n// Load left halo\n#pragma unroll\n\n for (int i = 0; i < ROWS_HALO_STEPS; i++) {\n s_Data[threadIdx.y][threadIdx.x + i * ROWS_BLOCKDIM_X] =\n (baseX >= -i * ROWS_BLOCKDIM_X) ? d_Src[i * ROWS_BLOCKDIM_X] : 0;\n }\n\n// Load right halo\n#pragma unroll\n\n for (int i = ROWS_HALO_STEPS + ROWS_RESULT_STEPS; i < ROWS_HALO_STEPS + ROWS_RESULT_STEPS + ROWS_HALO_STEPS; i++) {\n s_Data[threadIdx.y][threadIdx.x + i * ROWS_BLOCKDIM_X] =\n (imageW - baseX > i * ROWS_BLOCKDIM_X) ? d_Src[i * ROWS_BLOCKDIM_X] : 0;\n }\n\n // Compute and store results\n cg::sync(cta);\n#pragma unroll\n\n for (int i = ROWS_HALO_STEPS; i < ROWS_HALO_STEPS + ROWS_RESULT_STEPS; i++) {\n float sum = 0;\n\n#pragma unroll\n\n for (int j = -KERNEL_RADIUS; j <= KERNEL_RADIUS; j++) {\n sum += c_Kernel[KERNEL_RADIUS - j] * s_Data[threadIdx.y][threadIdx.x + i * ROWS_BLOCKDIM_X + j];\n }\n\n d_Dst[i * ROWS_BLOCKDIM_X] = sum;\n }\n}\n\nextern \"C\" void convolutionRowsGPU(float *d_Dst, float *d_Src, int imageW, int imageH)\n{\n assert(ROWS_BLOCKDIM_X * ROWS_HALO_STEPS >= KERNEL_RADIUS);\n assert(imageW % (ROWS_RESULT_STEPS * ROWS_BLOCKDIM_X) == 0);\n assert(imageH % ROWS_BLOCKDIM_Y == 0);\n\n dim3 blocks(imageW / (ROWS_RESULT_STEPS * ROWS_BLOCKDIM_X), imageH / ROWS_BLOCKDIM_Y);\n dim3 threads(ROWS_BLOCKDIM_X, ROWS_BLOCKDIM_Y);\n\n convolutionRowsKernel<<>>(d_Dst, d_Src, imageW, imageH, imageW);\n getLastCudaError(\"convolutionRowsKernel() execution failed\\n\");\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Column convolution filter\n////////////////////////////////////////////////////////////////////////////////\n#define COLUMNS_BLOCKDIM_X 16\n#define COLUMNS_BLOCKDIM_Y 8\n#define COLUMNS_RESULT_STEPS 8\n#define COLUMNS_HALO_STEPS 1\n\n__global__ void convolutionColumnsKernel(float *d_Dst, float *d_Src, int imageW, int imageH, int pitch)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float s_Data[COLUMNS_BLOCKDIM_X]\n [(COLUMNS_RESULT_STEPS + 2 * COLUMNS_HALO_STEPS) * COLUMNS_BLOCKDIM_Y + 1];\n\n // Offset to the upper halo edge\n const int baseX = blockIdx.x * COLUMNS_BLOCKDIM_X + threadIdx.x;\n const int baseY = (blockIdx.y * COLUMNS_RESULT_STEPS - COLUMNS_HALO_STEPS) * COLUMNS_BLOCKDIM_Y + threadIdx.y;\n d_Src += baseY * pitch + baseX;\n d_Dst += baseY * pitch + baseX;\n\n// Main data\n#pragma unroll\n\n for (int i = COLUMNS_HALO_STEPS; i < COLUMNS_HALO_STEPS + COLUMNS_RESULT_STEPS; i++) {\n s_Data[threadIdx.x][threadIdx.y + i * COLUMNS_BLOCKDIM_Y] = d_Src[i * COLUMNS_BLOCKDIM_Y * pitch];\n }\n\n// Upper halo\n#pragma unroll\n\n for (int i = 0; i < COLUMNS_HALO_STEPS; i++) {\n s_Data[threadIdx.x][threadIdx.y + i * COLUMNS_BLOCKDIM_Y] =\n (baseY >= -i * COLUMNS_BLOCKDIM_Y) ? d_Src[i * COLUMNS_BLOCKDIM_Y * pitch] : 0;\n }\n\n// Lower halo\n#pragma unroll\n\n for (int i = COLUMNS_HALO_STEPS + COLUMNS_RESULT_STEPS;\n i < COLUMNS_HALO_STEPS + COLUMNS_RESULT_STEPS + COLUMNS_HALO_STEPS;\n i++) {\n s_Data[threadIdx.x][threadIdx.y + i * COLUMNS_BLOCKDIM_Y] =\n (imageH - baseY > i * COLUMNS_BLOCKDIM_Y) ? d_Src[i * COLUMNS_BLOCKDIM_Y * pitch] : 0;\n }\n\n // Compute and store results\n cg::sync(cta);\n#pragma unroll\n\n for (int i = COLUMNS_HALO_STEPS; i < COLUMNS_HALO_STEPS + COLUMNS_RESULT_STEPS; i++) {\n float sum = 0;\n#pragma unroll\n\n for (int j = -KERNEL_RADIUS; j <= KERNEL_RADIUS; j++) {\n sum += c_Kernel[KERNEL_RADIUS - j] * s_Data[threadIdx.x][threadIdx.y + i * COLUMNS_BLOCKDIM_Y + j];\n }\n\n d_Dst[i * COLUMNS_BLOCKDIM_Y * pitch] = sum;\n }\n}\n\nextern \"C\" void convolutionColumnsGPU(float *d_Dst, float *d_Src, int imageW, int imageH)\n{\n assert(COLUMNS_BLOCKDIM_Y * COLUMNS_HALO_STEPS >= KERNEL_RADIUS);\n assert(imageW % COLUMNS_BLOCKDIM_X == 0);\n assert(imageH % (COLUMNS_RESULT_STEPS * COLUMNS_BLOCKDIM_Y) == 0);\n\n dim3 blocks(imageW / COLUMNS_BLOCKDIM_X, imageH / (COLUMNS_RESULT_STEPS * COLUMNS_BLOCKDIM_Y));\n dim3 threads(COLUMNS_BLOCKDIM_X, COLUMNS_BLOCKDIM_Y);\n\n convolutionColumnsKernel<<>>(d_Dst, d_Src, imageW, imageH, imageW);\n getLastCudaError(\"convolutionColumnsKernel() execution failed\\n\");\n}\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\nnamespace cg = cooperative_groups;\n#include \"convolutionSeparable_common.h\"\n\n////////////////////////////////////////////////////////////////////////////////\n// Convolution kernel storage\n////////////////////////////////////////////////////////////////////////////////\n__constant__ float c_Kernel[KERNEL_LENGTH];\n\nextern \"C\" void setConvolutionKernel(float *h_Kernel)\n{\n hipMemcpyToSymbol(HIP_SYMBOL(c_Kernel), h_Kernel, KERNEL_LENGTH * sizeof(float));\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Row convolution filter\n////////////////////////////////////////////////////////////////////////////////\n#define ROWS_BLOCKDIM_X 16\n#define ROWS_BLOCKDIM_Y 4\n#define ROWS_RESULT_STEPS 8\n#define ROWS_HALO_STEPS 1\n\n__global__ void convolutionRowsKernel(float *d_Dst, float *d_Src, int imageW, int imageH, int pitch)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float s_Data[ROWS_BLOCKDIM_Y][(ROWS_RESULT_STEPS + 2 * ROWS_HALO_STEPS) * ROWS_BLOCKDIM_X];\n\n // Offset to the left halo edge\n const int baseX = (blockIdx.x * ROWS_RESULT_STEPS - ROWS_HALO_STEPS) * ROWS_BLOCKDIM_X + threadIdx.x;\n const int baseY = blockIdx.y * ROWS_BLOCKDIM_Y + threadIdx.y;\n\n d_Src += baseY * pitch + baseX;\n d_Dst += baseY * pitch + baseX;\n\n// Load main data\n#pragma unroll\n\n for (int i = ROWS_HALO_STEPS; i < ROWS_HALO_STEPS + ROWS_RESULT_STEPS; i++) {\n s_Data[threadIdx.y][threadIdx.x + i * ROWS_BLOCKDIM_X] = d_Src[i * ROWS_BLOCKDIM_X];\n }\n\n// Load left halo\n#pragma unroll\n\n for (int i = 0; i < ROWS_HALO_STEPS; i++) {\n s_Data[threadIdx.y][threadIdx.x + i * ROWS_BLOCKDIM_X] =\n (baseX >= -i * ROWS_BLOCKDIM_X) ? d_Src[i * ROWS_BLOCKDIM_X] : 0;\n }\n\n// Load right halo\n#pragma unroll\n\n for (int i = ROWS_HALO_STEPS + ROWS_RESULT_STEPS; i < ROWS_HALO_STEPS + ROWS_RESULT_STEPS + ROWS_HALO_STEPS; i++) {\n s_Data[threadIdx.y][threadIdx.x + i * ROWS_BLOCKDIM_X] =\n (imageW - baseX > i * ROWS_BLOCKDIM_X) ? d_Src[i * ROWS_BLOCKDIM_X] : 0;\n }\n\n // Compute and store results\n cg::sync(cta);\n#pragma unroll\n\n for (int i = ROWS_HALO_STEPS; i < ROWS_HALO_STEPS + ROWS_RESULT_STEPS; i++) {\n float sum = 0;\n\n#pragma unroll\n\n for (int j = -KERNEL_RADIUS; j <= KERNEL_RADIUS; j++) {\n sum += c_Kernel[KERNEL_RADIUS - j] * s_Data[threadIdx.y][threadIdx.x + i * ROWS_BLOCKDIM_X + j];\n }\n\n d_Dst[i * ROWS_BLOCKDIM_X] = sum;\n }\n}\n\nextern \"C\" void convolutionRowsGPU(float *d_Dst, float *d_Src, int imageW, int imageH)\n{\n assert(ROWS_BLOCKDIM_X * ROWS_HALO_STEPS >= KERNEL_RADIUS);\n assert(imageW % (ROWS_RESULT_STEPS * ROWS_BLOCKDIM_X) == 0);\n assert(imageH % ROWS_BLOCKDIM_Y == 0);\n\n dim3 blocks(imageW / (ROWS_RESULT_STEPS * ROWS_BLOCKDIM_X), imageH / ROWS_BLOCKDIM_Y);\n dim3 threads(ROWS_BLOCKDIM_X, ROWS_BLOCKDIM_Y);\n\n convolutionRowsKernel<<>>(d_Dst, d_Src, imageW, imageH, imageW);\n getLastCudaError(\"convolutionRowsKernel() execution failed\\n\");\n}\n\n////////////////////////////////////////////////////////////////////////////////\n// Column convolution filter\n////////////////////////////////////////////////////////////////////////////////\n#define COLUMNS_BLOCKDIM_X 16\n#define COLUMNS_BLOCKDIM_Y 8\n#define COLUMNS_RESULT_STEPS 8\n#define COLUMNS_HALO_STEPS 1\n\n__global__ void convolutionColumnsKernel(float *d_Dst, float *d_Src, int imageW, int imageH, int pitch)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n __shared__ float s_Data[COLUMNS_BLOCKDIM_X]\n [(COLUMNS_RESULT_STEPS + 2 * COLUMNS_HALO_STEPS) * COLUMNS_BLOCKDIM_Y + 1];\n\n // Offset to the upper halo edge\n const int baseX = blockIdx.x * COLUMNS_BLOCKDIM_X + threadIdx.x;\n const int baseY = (blockIdx.y * COLUMNS_RESULT_STEPS - COLUMNS_HALO_STEPS) * COLUMNS_BLOCKDIM_Y + threadIdx.y;\n d_Src += baseY * pitch + baseX;\n d_Dst += baseY * pitch + baseX;\n\n// Main data\n#pragma unroll\n\n for (int i = COLUMNS_HALO_STEPS; i < COLUMNS_HALO_STEPS + COLUMNS_RESULT_STEPS; i++) {\n s_Data[threadIdx.x][threadIdx.y + i * COLUMNS_BLOCKDIM_Y] = d_Src[i * COLUMNS_BLOCKDIM_Y * pitch];\n }\n\n// Upper halo\n#pragma unroll\n\n for (int i = 0; i < COLUMNS_HALO_STEPS; i++) {\n s_Data[threadIdx.x][threadIdx.y + i * COLUMNS_BLOCKDIM_Y] =\n (baseY >= -i * COLUMNS_BLOCKDIM_Y) ? d_Src[i * COLUMNS_BLOCKDIM_Y * pitch] : 0;\n }\n\n// Lower halo\n#pragma unroll\n\n for (int i = COLUMNS_HALO_STEPS + COLUMNS_RESULT_STEPS;\n i < COLUMNS_HALO_STEPS + COLUMNS_RESULT_STEPS + COLUMNS_HALO_STEPS;\n i++) {\n s_Data[threadIdx.x][threadIdx.y + i * COLUMNS_BLOCKDIM_Y] =\n (imageH - baseY > i * COLUMNS_BLOCKDIM_Y) ? d_Src[i * COLUMNS_BLOCKDIM_Y * pitch] : 0;\n }\n\n // Compute and store results\n cg::sync(cta);\n#pragma unroll\n\n for (int i = COLUMNS_HALO_STEPS; i < COLUMNS_HALO_STEPS + COLUMNS_RESULT_STEPS; i++) {\n float sum = 0;\n#pragma unroll\n\n for (int j = -KERNEL_RADIUS; j <= KERNEL_RADIUS; j++) {\n sum += c_Kernel[KERNEL_RADIUS - j] * s_Data[threadIdx.x][threadIdx.y + i * COLUMNS_BLOCKDIM_Y + j];\n }\n\n d_Dst[i * COLUMNS_BLOCKDIM_Y * pitch] = sum;\n }\n}\n\nextern \"C\" void convolutionColumnsGPU(float *d_Dst, float *d_Src, int imageW, int imageH)\n{\n assert(COLUMNS_BLOCKDIM_Y * COLUMNS_HALO_STEPS >= KERNEL_RADIUS);\n assert(imageW % COLUMNS_BLOCKDIM_X == 0);\n assert(imageH % (COLUMNS_RESULT_STEPS * COLUMNS_BLOCKDIM_Y) == 0);\n\n dim3 blocks(imageW / COLUMNS_BLOCKDIM_X, imageH / (COLUMNS_RESULT_STEPS * COLUMNS_BLOCKDIM_Y));\n dim3 threads(COLUMNS_BLOCKDIM_X, COLUMNS_BLOCKDIM_Y);\n\n convolutionColumnsKernel<<>>(d_Dst, d_Src, imageW, imageH, imageW);\n getLastCudaError(\"convolutionColumnsKernel() execution failed\\n\");\n}\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/convolutionSeparable/convolutionSeparable.cu"} {"input": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"../inc/piestimator.h\"\n\nnamespace cg = cooperative_groups;\n#include \n\nusing std::string;\nusing std::vector;\n\n__device__ unsigned int reduce_sum(unsigned int in, cg::thread_block cta)\n{\n extern __shared__ unsigned int sdata[];\n\n // Perform first level of reduction:\n // - Write to shared memory\n unsigned int ltid = threadIdx.x;\n\n sdata[ltid] = in;\n cg::sync(cta);\n\n // Do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (ltid < s) {\n sdata[ltid] += sdata[ltid + s];\n }\n\n cg::sync(cta);\n }\n\n return sdata[0];\n}\n\n// Estimator kernel\ntemplate \n__global__ void computeValue(unsigned int *const results, const Real *const points, const unsigned int numSims)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Determine thread ID\n unsigned int bid = blockIdx.x;\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int step = gridDim.x * blockDim.x;\n\n // Shift the input/output pointers\n const Real *pointx = points + tid;\n const Real *pointy = pointx + numSims;\n\n // Count the number of points which lie inside the unit quarter-circle\n unsigned int pointsInside = 0;\n\n for (unsigned int i = tid; i < numSims; i += step, pointx += step, pointy += step) {\n Real x = *pointx;\n Real y = *pointy;\n Real l2norm2 = x * x + y * y;\n\n if (l2norm2 < static_cast(1)) {\n pointsInside++;\n }\n }\n\n // Reduce within the block\n pointsInside = reduce_sum(pointsInside, cta);\n\n // Store the result\n if (threadIdx.x == 0) {\n results[bid] = pointsInside;\n }\n}\n\ntemplate \nPiEstimator::PiEstimator(unsigned int numSims,\n unsigned int device,\n unsigned int threadBlockSize,\n unsigned int seed)\n : m_numSims(numSims)\n , m_device(device)\n , m_threadBlockSize(threadBlockSize)\n , m_seed(seed)\n{\n}\n\ntemplate Real PiEstimator::operator()()\n{\n cudaError_t cudaResult = cudaSuccess;\n struct cudaDeviceProp deviceProperties;\n struct cudaFuncAttributes funcAttributes;\n\n // Get device properties\n cudaResult = cudaGetDeviceProperties(&deviceProperties, m_device);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not get device properties: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Check precision is valid\n if (typeid(Real) == typeid(double)\n && (deviceProperties.major < 1 || (deviceProperties.major == 1 && deviceProperties.minor < 3))) {\n throw std::runtime_error(\"Device does not have double precision support\");\n }\n\n // Attach to GPU\n cudaResult = cudaSetDevice(m_device);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not set CUDA device: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Determine how to divide the work between cores\n dim3 block;\n dim3 grid;\n block.x = m_threadBlockSize;\n grid.x = (m_numSims + m_threadBlockSize - 1) / m_threadBlockSize;\n\n // Aim to launch around ten or more times as many blocks as there\n // are multiprocessors on the target device.\n unsigned int blocksPerSM = 10;\n unsigned int numSMs = deviceProperties.multiProcessorCount;\n\n while (grid.x > 2 * blocksPerSM * numSMs) {\n grid.x >>= 1;\n }\n\n // Get computeValue function properties and check the maximum block size\n cudaResult = cudaFuncGetAttributes(&funcAttributes, computeValue);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for computeValue kernel\");\n }\n\n // Check the dimensions are valid\n if (block.x > (unsigned int)deviceProperties.maxThreadsDim[0]) {\n throw std::runtime_error(\"Block X dimension is too large for device\");\n }\n\n if (grid.x > (unsigned int)deviceProperties.maxGridSize[0]) {\n throw std::runtime_error(\"Grid X dimension is too large for device\");\n }\n\n // Allocate memory for points\n // Each simulation has two random numbers to give X and Y coordinate\n Real *d_points = 0;\n cudaResult = cudaMalloc((void **)&d_points, 2 * m_numSims * sizeof(Real));\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not allocate memory on device for random numbers: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Allocate memory for result\n // Each thread block will produce one result\n unsigned int *d_results = 0;\n cudaResult = cudaMalloc((void **)&d_results, grid.x * sizeof(unsigned int));\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not allocate memory on device for partial results: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Generate random points in unit square\n curandStatus_t curandResult;\n curandGenerator_t prng;\n curandResult = curandCreateGenerator(&prng, CURAND_RNG_PSEUDO_DEFAULT);\n\n if (curandResult != CURAND_STATUS_SUCCESS) {\n string msg(\"Could not create pseudo-random number generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n curandResult = curandSetPseudoRandomGeneratorSeed(prng, m_seed);\n\n if (curandResult != CURAND_STATUS_SUCCESS) {\n string msg(\"Could not set seed for pseudo-random number generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n if (typeid(Real) == typeid(float)) {\n curandResult = curandGenerateUniform(prng, (float *)d_points, 2 * m_numSims);\n }\n else if (typeid(Real) == typeid(double)) {\n curandResult = curandGenerateUniformDouble(prng, (double *)d_points, 2 * m_numSims);\n }\n else {\n string msg(\"Could not generate random numbers of specified type\");\n throw std::runtime_error(msg);\n }\n\n if (curandResult != CURAND_STATUS_SUCCESS) {\n string msg(\"Could not generate pseudo-random numbers: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n curandResult = curandDestroyGenerator(prng);\n\n if (curandResult != CURAND_STATUS_SUCCESS) {\n string msg(\"Could not destroy pseudo-random number generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n // Count the points inside unit quarter-circle\n computeValue<<>>(d_results, d_points, m_numSims);\n\n // Copy partial results back\n vector results(grid.x);\n cudaResult = cudaMemcpy(&results[0], d_results, grid.x * sizeof(unsigned int), cudaMemcpyDeviceToHost);\n\n if (cudaResult != cudaSuccess) {\n string msg(\"Could not copy partial results to host: \");\n msg += cudaGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Complete sum-reduction on host\n Real value = static_cast(std::accumulate(results.begin(), results.end(), 0));\n\n // Determine the proportion of points inside the quarter-circle,\n // i.e. the area of the unit quarter-circle\n value /= m_numSims;\n\n // Value is currently an estimate of the area of a unit quarter-circle, so we\n // can scale to a full circle by multiplying by four. Now since the area of a\n // circle is pi * r^2, and r is one, the value will be an estimate for the\n // value of pi.\n value *= 4;\n\n // Cleanup\n if (d_points) {\n cudaFree(d_points);\n d_points = 0;\n }\n\n if (d_results) {\n cudaFree(d_results);\n d_results = 0;\n }\n\n return value;\n}\n\n// Explicit template instantiation\ntemplate class PiEstimator;\ntemplate class PiEstimator;\n", "output": "/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of NVIDIA CORPORATION nor the names of its\n * contributors may be used to endorse or promote products derived\n * from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"../inc/piestimator.h\"\n\nnamespace cg = cooperative_groups;\n#include \n\nusing std::string;\nusing std::vector;\n\n__device__ unsigned int reduce_sum(unsigned int in, cg::thread_block cta)\n{\n extern __shared__ unsigned int sdata[];\n\n // Perform first level of reduction:\n // - Write to shared memory\n unsigned int ltid = threadIdx.x;\n\n sdata[ltid] = in;\n cg::sync(cta);\n\n // Do reduction in shared mem\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (ltid < s) {\n sdata[ltid] += sdata[ltid + s];\n }\n\n cg::sync(cta);\n }\n\n return sdata[0];\n}\n\n// Estimator kernel\ntemplate \n__global__ void computeValue(unsigned int *const results, const Real *const points, const unsigned int numSims)\n{\n // Handle to thread block group\n cg::thread_block cta = cg::this_thread_block();\n // Determine thread ID\n unsigned int bid = blockIdx.x;\n unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;\n unsigned int step = gridDim.x * blockDim.x;\n\n // Shift the input/output pointers\n const Real *pointx = points + tid;\n const Real *pointy = pointx + numSims;\n\n // Count the number of points which lie inside the unit quarter-circle\n unsigned int pointsInside = 0;\n\n for (unsigned int i = tid; i < numSims; i += step, pointx += step, pointy += step) {\n Real x = *pointx;\n Real y = *pointy;\n Real l2norm2 = x * x + y * y;\n\n if (l2norm2 < static_cast(1)) {\n pointsInside++;\n }\n }\n\n // Reduce within the block\n pointsInside = reduce_sum(pointsInside, cta);\n\n // Store the result\n if (threadIdx.x == 0) {\n results[bid] = pointsInside;\n }\n}\n\ntemplate \nPiEstimator::PiEstimator(unsigned int numSims,\n unsigned int device,\n unsigned int threadBlockSize,\n unsigned int seed)\n : m_numSims(numSims)\n , m_device(device)\n , m_threadBlockSize(threadBlockSize)\n , m_seed(seed)\n{\n}\n\ntemplate Real PiEstimator::operator()()\n{\n hipError_t cudaResult = hipSuccess;\n struct hipDeviceProp_t deviceProperties;\n struct hipFuncAttributes funcAttributes;\n\n // Get device properties\n cudaResult = hipGetDeviceProperties(&deviceProperties, m_device);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not get device properties: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Check precision is valid\n if (typeid(Real) == typeid(double)\n && (deviceProperties.major < 1 || (deviceProperties.major == 1 && deviceProperties.minor < 3))) {\n throw std::runtime_error(\"Device does not have double precision support\");\n }\n\n // Attach to GPU\n cudaResult = hipSetDevice(m_device);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not set CUDA device: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Determine how to divide the work between cores\n dim3 block;\n dim3 grid;\n block.x = m_threadBlockSize;\n grid.x = (m_numSims + m_threadBlockSize - 1) / m_threadBlockSize;\n\n // Aim to launch around ten or more times as many blocks as there\n // are multiprocessors on the target device.\n unsigned int blocksPerSM = 10;\n unsigned int numSMs = deviceProperties.multiProcessorCount;\n\n while (grid.x > 2 * blocksPerSM * numSMs) {\n grid.x >>= 1;\n }\n\n // Get computeValue function properties and check the maximum block size\n cudaResult = hipFuncGetAttributes(&funcAttributes,reinterpret_cast(computeValue));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not get function attributes: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n if (block.x > (unsigned int)funcAttributes.maxThreadsPerBlock) {\n throw std::runtime_error(\"Block X dimension is too large for computeValue kernel\");\n }\n\n // Check the dimensions are valid\n if (block.x > (unsigned int)deviceProperties.maxThreadsDim[0]) {\n throw std::runtime_error(\"Block X dimension is too large for device\");\n }\n\n if (grid.x > (unsigned int)deviceProperties.maxGridSize[0]) {\n throw std::runtime_error(\"Grid X dimension is too large for device\");\n }\n\n // Allocate memory for points\n // Each simulation has two random numbers to give X and Y coordinate\n Real *d_points = 0;\n cudaResult = hipMalloc((void **)&d_points, 2 * m_numSims * sizeof(Real));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not allocate memory on device for random numbers: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Allocate memory for result\n // Each thread block will produce one result\n unsigned int *d_results = 0;\n cudaResult = hipMalloc((void **)&d_results, grid.x * sizeof(unsigned int));\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not allocate memory on device for partial results: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Generate random points in unit square\n hiprandStatus_t curandResult;\n hiprandGenerator_t prng;\n curandResult = hiprandCreateGenerator(&prng, HIPRAND_RNG_PSEUDO_DEFAULT);\n\n if (curandResult != HIPRAND_STATUS_SUCCESS) {\n string msg(\"Could not create pseudo-random number generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n curandResult = hiprandSetPseudoRandomGeneratorSeed(prng, m_seed);\n\n if (curandResult != HIPRAND_STATUS_SUCCESS) {\n string msg(\"Could not set seed for pseudo-random number generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n if (typeid(Real) == typeid(float)) {\n curandResult = hiprandGenerateUniform(prng, (float *)d_points, 2 * m_numSims);\n }\n else if (typeid(Real) == typeid(double)) {\n curandResult = hiprandGenerateUniformDouble(prng, (double *)d_points, 2 * m_numSims);\n }\n else {\n string msg(\"Could not generate random numbers of specified type\");\n throw std::runtime_error(msg);\n }\n\n if (curandResult != HIPRAND_STATUS_SUCCESS) {\n string msg(\"Could not generate pseudo-random numbers: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n curandResult = hiprandDestroyGenerator(prng);\n\n if (curandResult != HIPRAND_STATUS_SUCCESS) {\n string msg(\"Could not destroy pseudo-random number generator: \");\n msg += curandResult;\n throw std::runtime_error(msg);\n }\n\n // Count the points inside unit quarter-circle\n computeValue<<>>(d_results, d_points, m_numSims);\n\n // Copy partial results back\n vector results(grid.x);\n cudaResult = hipMemcpy(&results[0], d_results, grid.x * sizeof(unsigned int), hipMemcpyDeviceToHost);\n\n if (cudaResult != hipSuccess) {\n string msg(\"Could not copy partial results to host: \");\n msg += hipGetErrorString(cudaResult);\n throw std::runtime_error(msg);\n }\n\n // Complete sum-reduction on host\n Real value = static_cast(std::accumulate(results.begin(), results.end(), 0));\n\n // Determine the proportion of points inside the quarter-circle,\n // i.e. the area of the unit quarter-circle\n value /= m_numSims;\n\n // Value is currently an estimate of the area of a unit quarter-circle, so we\n // can scale to a full circle by multiplying by four. Now since the area of a\n // circle is pi * r^2, and r is one, the value will be an estimate for the\n // value of pi.\n value *= 4;\n\n // Cleanup\n if (d_points) {\n hipFree(d_points);\n d_points = 0;\n }\n\n if (d_results) {\n hipFree(d_results);\n d_results = 0;\n }\n\n return value;\n}\n\n// Explicit template instantiation\ntemplate class PiEstimator;\ntemplate class PiEstimator;\n", "source": "/kaggle/working/cuda-samples/Samples/2_Concepts_and_Techniques/MC_EstimatePiP/src/piestimator.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n\n#ifdef _WIN32\n#pragma warning (push)\n#pragma warning (disable : 4244 4267 4521)\n#endif\n#include \n#include \n#include \n#include \n#ifdef _WIN32\n#pragma warning (pop)\n#endif\n\n#include \n\n#include \n#include \n\n\nnamespace amgx\n{\n\n#define USE_EXPERIMENTAL_4x4\n\ntemplate \nclass Multiply_1x1;\ntemplate \nclass Multiply_1x1_masked;\ntemplate \nclass Multiply_1x1_with_mask;\ntemplate \nclass Multiply_1x1_with_mask_restriction;\ntemplate \nclass Multiply_3x3;\ntemplate \nclass Multiply_4x4;\ntemplate \nclass Multiply_bxb;\n\ntemplate \nvoid multiply_block_size(Matrix &A, Vector &B, Vector &C, ViewType view)\n{\n typedef Matrix TMatrix;\n typedef Vector TVector;\n\n if (A.get_block_size() == 1)\n {\n Multiply_1x1::multiply_1x1(A, B, C, view);\n }\n else if (A.get_block_dimy() == 3 && A.get_block_dimx() == 3)\n {\n Multiply_3x3::multiply_3x3(A, B, C, view);\n }\n else if (A.get_block_dimy() == 4 && A.get_block_dimx() == 4)\n {\n Multiply_4x4::multiply_4x4(A, B, C, view);\n }\n else\n {\n Multiply_bxb::multiply_bxb(A, B, C, view);\n }\n}\n\ntemplate \nvoid multiply(Matrix &A, Vector &B, Vector &C, ViewType view)\n{\n typedef Matrix TMatrix;\n typedef Vector TVector;\n\n if (!A.is_initialized())\n {\n FatalError(\"Trying to multiply uninitialized matrix\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (A.get_block_dimx() != B.get_block_dimy())\n {\n std::stringstream ss;\n ss << \"Matrix and vector dimensions don't match: A.dimx == \" << A.get_block_dimx() << \", B.dimy == \" << B.get_block_dimy();\n FatalError(ss.str().c_str(), AMGX_ERR_BAD_PARAMETERS);\n }\n\n\n typedef Matrix TMatrix;\n typedef Vector TVector;\n\n bool latencyHiding = (view == A.getViewExterior() && A.getViewInterior() != A.getViewExterior() && !A.is_matrix_singleGPU() && B.dirtybit != 0);\n\n if (latencyHiding)\n {\n A.manager->exchange_halo_split_gather(B, B.tag);\n\n // Multiply interior rows\n multiply_block_size(A, B, C, A.getViewInterior());\n\n // Finish halo exchange\n A.manager->exchange_halo_split_finish(B, B.tag);\n\n // Multiply rows with halo dependencies\n ViewType bnd_view = (ViewType)(~(A.getViewInterior()) & A.getViewExterior());\n multiply_block_size(A, B, C, bnd_view);\n }\n else\n {\n if (view != INTERIOR && !A.is_matrix_singleGPU() && B.dirtybit != 0)\n {\n A.manager->exchange_halo_v2(B, B.tag);\n }\n\n multiply_block_size(A, B, C, view);\n }\n\n C.dirtybit = 1;\n C.set_block_dimy(A.get_block_dimx());\n}\n\ntemplate \nvoid multiply_masked(Matrix &A, Vector &B, Vector &C, typename Matrix::IVector &mask, ViewType view)\n{\n typedef Matrix TMatrix;\n typedef Vector TVector;\n typedef typename Matrix::IVector TIVector;\n\n if (!A.is_initialized())\n {\n FatalError(\"Trying to multiply uninitialized matrix\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if(A.get_block_size() != 1)\n {\n FatalError(\"Unsupported blocksize for multiply_masked()\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (A.get_block_dimx() != B.get_block_dimy())\n {\n std::stringstream ss;\n ss << \"Matrix and vector dimensions don't match: A.dimx == \" << A.get_block_dimx() << \", B.dimy == \" << B.get_block_dimy();\n FatalError(ss.str().c_str(), AMGX_ERR_BAD_PARAMETERS);\n }\n\n Multiply_1x1_masked::multiply_1x1_masked(A, B, C, mask, view);\n C.set_block_dimy(A.get_block_dimx());\n}\n\ntemplate\nvoid multiply_with_mask(Matrix &A, Vector &B, Vector &C)\n{\n if (!A.is_initialized())\n {\n FatalError(\"Trying to multiply uninitialized matrix\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (A.get_block_dimx() != B.get_block_dimy())\n {\n std::stringstream ss;\n ss << \"Matrix and vector dimensions don't match: A.dimx == \" << A.get_block_dimx() << \", B.dimy == \" << B.get_block_dimy();\n FatalError(ss.str().c_str(), AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (A.get_block_size() == 1)\n {\n Multiply_1x1_with_mask::multiply_1x1(A, B, C);\n }\n else\n {\n FatalError(\"multiply with mask not supported for bsize != 1\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n C.set_block_dimy(A.get_block_dimx());\n C.dirtybit = 1;\n //if (!A.is_matrix_singleGPU() && C.size() == B.size() && C.delayed_send==0)\n // A.manager->exchange_halo_async(C, C.tag);\n}\n\n\ntemplate\nvoid multiply_with_mask_restriction(Matrix &A, Vector &B, Vector &C, Matrix &P)\n{\n if (!A.is_initialized())\n {\n FatalError(\"Trying to multiply uninitialized matrix\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (A.get_block_dimx() != B.get_block_dimy())\n {\n std::stringstream ss;\n ss << \"Matrix and vector dimensions don't match: A.dimx == \" << A.get_block_dimx() << \", B.dimy == \" << B.get_block_dimy();\n FatalError(ss.str().c_str(), AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (A.get_block_size() == 1)\n {\n Multiply_1x1_with_mask_restriction::multiply_1x1(A, B, C, P);\n }\n else\n {\n FatalError(\"multiply with mask not supported for bsize != 1\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n C.set_block_dimy(A.get_block_dimx());\n C.dirtybit = 1;\n}\n\n\n\n\ntemplate\nvoid multiplyMM(const Matrix &A, const Matrix &B, Matrix &C)\n{\n if (!A.is_initialized())\n {\n FatalError(\"Trying to multiply uninitialized matrix\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (A.get_block_dimx() != B.get_block_dimx() || A.get_block_dimy() != B.get_block_dimy())\n {\n FatalError(\"Matrices dimensions do not match\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (TConfig::memSpace == AMGX_device)\n {\n FatalError(\"Error, multiplyMM not implemented on device\", AMGX_ERR_BAD_PARAMETERS);\n }\n else\n {\n if (A.get_block_size() != 1)\n {\n FatalError(\"multiplyMM only works for block_size ==1\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n typedef typename TConfig::IndPrec IndexType;\n typedef typename TConfig::MatPrec ValueType;\n typedef typename Matrix::IVector IVector;\n typedef typename Matrix::MVector MVector;\n C.set_initialized(0);\n IndexType num_nonzeros = 0;\n IVector mask(B.get_num_cols(), IndexType (-1));\n\n // Compute nnz in C (including explicit zeros)\n for (size_t i = 0; i < A.get_num_rows(); i++)\n {\n for (IndexType jj = A.row_offsets[i]; jj < A.row_offsets[i + 1]; jj++)\n {\n IndexType j = A.col_indices[jj];\n\n for (IndexType kk = B.row_offsets[j]; kk < B.row_offsets[j + 1]; kk++)\n {\n IndexType k = B.col_indices[kk];\n\n if (mask[k] != i)\n {\n mask[k] = i;\n num_nonzeros++;\n }\n }\n }\n }\n\n // Resize output\n C.resize(A.get_num_rows(), B.get_num_cols(), num_nonzeros);\n const IndexType unseen = static_cast(-1);\n const IndexType init = static_cast(-2);\n // Compute entries of C\n IVector next(B.get_num_cols(), unseen);\n MVector sums(B.get_num_cols(), types::util::get_zero());\n num_nonzeros = 0;\n C.row_offsets[0] = 0;\n\n for (size_t i = 0; i < A.get_num_rows(); i++)\n {\n IndexType head = init;\n IndexType length = 0;\n IndexType jj_start = A.row_offsets[i];\n IndexType jj_end = A.row_offsets[i + 1];\n\n for (IndexType jj = jj_start; jj < jj_end; jj++)\n {\n IndexType j = A.col_indices[jj];\n ValueType v = A.values[jj];\n IndexType kk_start = B.row_offsets[j];\n IndexType kk_end = B.row_offsets[j + 1];\n\n for (IndexType kk = kk_start; kk < kk_end; kk++)\n {\n IndexType k = B.col_indices[kk];\n sums[k] = sums[k] + v * B.values[kk];\n\n if (next[k] == unseen)\n {\n next[k] = head;\n head = k;\n length++;\n }\n }\n }\n\n for (IndexType jj = 0; jj < length; jj++)\n {\n //if(sums[head] != ValueType(0))\n //{\n C.col_indices[num_nonzeros] = head;\n C.values[num_nonzeros] = sums[head];\n num_nonzeros++;\n //}\n IndexType temp = head;\n head = next[head];\n // clear arrays\n next[temp] = unseen;\n sums[temp] = types::util::get_zero();\n }\n\n C.row_offsets[i + 1] = num_nonzeros;\n }\n\n // Resize output again since pass2 omits explict zeros\n //C.resize(A.num_rows, B.num_cols, num_nonzeros);\n C.set_initialized(1);\n }\n }\n}\n\n// --------------------------------\n// KERNELS\n// --------------------------------\n\ntemplate\n__global__\nvoid blockDiaCsrMultiplyKernel(const IndexType *row_offsets,\n const IndexType *column_indices,\n const ValueTypeA *nonzero_values,\n const ValueTypeB *B,\n ValueTypeB *C,\n const IndexType num_block_rows,\n const IndexType row_offset)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int eighthwarp_id = row_offset + (tid >> log_bsize);\n const int block_eighthwarp_id = threadIdx.x >> log_bsize;\n const int vec_entry_index = threadIdx.x & (bsize - 1);\n volatile __shared__ ValueTypeB s_xtemp[ bsize * eighthwarps_per_block ];\n ValueTypeB C_temp;\n int offset, s_offset;\n ValueTypeB temp[bsize];\n\n while (eighthwarp_id < num_block_rows)\n {\n //i = eighthwarp_id;\n C_temp = types::util::get_zero();\n // Contribution from each nonzero column\n int jmax = row_offsets[eighthwarp_id + 1];\n\n for (int jind = row_offsets[eighthwarp_id]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n offset = jcol * bsize + vec_entry_index;\n types::util::volcast(__cachingLoad(&B[offset]), s_xtemp + threadIdx.x);\n // Load nonzero_values\n s_offset = block_eighthwarp_id * bsize;\n\n if (ROW_MAJOR)\n {\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n }\n else\n {\n offset = jind * bsize * bsize + vec_entry_index;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n types::util::to_uptype(nonzero_values[offset + bsize * m], temp[m]);\n }\n }\n\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n C_temp = C_temp + temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n }\n\n C[eighthwarp_id * bsize + vec_entry_index] = C_temp;\n eighthwarp_id += gridDim.x * blockDim.x >> log_bsize;\n }\n}\n\n#ifdef USE_EXPERIMENTAL_4x4\n\ntemplate< typename IndexType, typename ValueTypeA, typename ValueTypeB, int CTA_SIZE, bool ROW_MAJOR >\n__global__\n__launch_bounds__( CTA_SIZE, 16 )\nvoid blockDiaCsrMultiplyKernelDiaProps_4x4( const IndexType *row_offsets,\n const IndexType *column_indices,\n const IndexType *dia_ptr,\n const ValueTypeA *nonzero_values,\n const ValueTypeB *B,\n ValueTypeB *C,\n const int num_block_rows,\n const int row_offset )\n{\n const int nHalfWarps = CTA_SIZE / 16; // Number of half warps per CTA.\n const int laneId = threadIdx.x % warpSize;\n const int halfWarpId = threadIdx.x / 16;\n const int halfLaneId = threadIdx.x % 16;\n const int halfLaneId_div_4 = halfLaneId / 4;\n const int halfLaneId_mod_4 = halfLaneId % 4;\n const int laneId_div_16 = laneId / 16;\n const int upperHalf = 16 * laneId_div_16;\n // Shared memory needed to exchange X and delta.\n __shared__ volatile ValueTypeB s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile ValueTypeB *my_s_mem = &s_mem[16 * halfWarpId];\n\n // Iterate over the rows of the matrix. One warp per two rows.\n for ( int aRowId = blockIdx.x * nHalfWarps + halfWarpId ; aRowId < num_block_rows ; aRowId += gridDim.x * nHalfWarps )\n {\n unsigned int active_mask = utils::activemask();\n // Load one block of B.\n ValueTypeB my_Ax = types::util::get_zero();\n\n // The diagonal.\n if ( halfLaneId_div_4 == 0 )\n {\n types::util::volcast( B[4 * aRowId + halfLaneId_mod_4], my_s_mem + halfLaneId);\n }\n\n // Load the diagonal.\n int diagId = dia_ptr[aRowId];\n // Update my values.\n ValueTypeA my_val = nonzero_values[16 * diagId + halfLaneId];\n\n if ( ROW_MAJOR )\n {\n my_Ax = my_Ax + my_val * types::util::volcast(my_s_mem[halfLaneId_mod_4]);\n }\n else\n {\n my_Ax = my_Ax + my_val * types::util::volcast(my_s_mem[halfLaneId_div_4]);\n }\n\n // The range of the rows.\n int aColBegin = row_offsets[aRowId + 0];\n int aColEnd = row_offsets[aRowId + 1];\n\n // Each warp load column indices of 16 nonzero blocks\n for ( ; utils::any( aColBegin < aColEnd, active_mask ) ; aColBegin += 16 )\n {\n int aColIt = aColBegin + halfLaneId;\n // Get the ID of the column.\n int aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColId = column_indices[aColIt];\n }\n\n // Count the number of active columns.\n int vote = utils::ballot(aColId != -1, active_mask);\n // The number of iterations.\n int nCols = max( __popc( vote & 0x0000ffff ), __popc( vote & 0xffff0000 ) );\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < nCols ; k += 4 )\n {\n int my_k = k + halfLaneId_div_4;\n // Exchange column indices.\n int waColId = utils::shfl( aColId, upperHalf + my_k, warpSize, active_mask );\n // Load 8 blocks of X if needed.\n ValueTypeB my_x = types::util::get_zero();\n\n if ( waColId != -1 )\n {\n my_x = B[4 * waColId + halfLaneId_mod_4];\n }\n\n types::util::volcast( my_x, my_s_mem + halfLaneId);\n // Load 8 blocks of A.\n#pragma unroll\n\n for ( int i = 0 ; i < 4 ; ++i )\n {\n int w_aColTmp = aColBegin + k + i, w_aColIt = -1;\n\n if ( w_aColTmp < aColEnd )\n {\n w_aColIt = w_aColTmp;\n }\n\n ValueTypeA my_val = types::util::get_zero();\n\n if ( w_aColIt != -1 )\n {\n my_val = nonzero_values[16 * w_aColIt + halfLaneId];\n }\n\n if ( ROW_MAJOR )\n {\n my_Ax = my_Ax + my_val * types::util::volcast(my_s_mem[4 * i + halfLaneId_mod_4]);\n }\n else\n {\n my_Ax = my_Ax + my_val * types::util::volcast(my_s_mem[4 * i + halfLaneId_div_4]);\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Reduce bmAx terms.\n if ( ROW_MAJOR )\n {\n my_Ax = my_Ax + utils::shfl_xor( my_Ax, 1, warpSize, active_mask );\n my_Ax = my_Ax + utils::shfl_xor( my_Ax, 2, warpSize, active_mask );\n }\n else\n {\n my_Ax = my_Ax + utils::shfl_xor( my_Ax, 4, warpSize, active_mask );\n my_Ax = my_Ax + utils::shfl_xor( my_Ax, 8, warpSize, active_mask );\n }\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n if ( halfLaneId_mod_4 == 0 )\n {\n C[4 * aRowId + halfLaneId_div_4] = my_Ax;\n }\n }\n else\n {\n if ( halfLaneId_div_4 == 0 )\n {\n C[4 * aRowId + halfLaneId_mod_4] = my_Ax;\n }\n }\n }\n}\n\n#else\n\ntemplate\n__global__\nvoid blockDiaCsrMultiplyKernelDiaProps_4x4(const IndexType *row_offsets,\n const IndexType *column_indices,\n const IndexType *dia_ptr,\n const ValueTypeA *nonzero_values,\n const ValueTypeB *B,\n ValueTypeB *C,\n const int num_block_rows,\n const int row_offset)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int eighthwarp_id = row_offset + (tid >> log_bsize);\n const int block_eighthwarp_id = threadIdx.x >> log_bsize;\n const int vec_entry_index = threadIdx.x & (bsize - 1);\n volatile __shared__ ValueTypeB s_xtemp[ bsize * eighthwarps_per_block ];\n ValueTypeB C_temp;\n int offset, s_offset;\n\n while (eighthwarp_id < num_block_rows)\n {\n //i = eighthwarp_id;\n C_temp = types::util::get_zero();\n // Contribution from diagonal\n offset = eighthwarp_id * bsize + vec_entry_index;\n types::util::volcast(__cachingLoad(&B[offset]), s_xtemp + threadIdx.x);\n // Load dia_values and do matrix multiply\n s_offset = block_eighthwarp_id * bsize;\n ValueTypeA temp[bsize];\n\n if (ROW_MAJOR)\n {\n loadAsVector(nonzero_values + bsize_sq * dia_ptr[eighthwarp_id] + vec_entry_index * bsize, temp);\n }\n else\n {\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n temp[m] = nonzero_values[bsize_sq * dia_ptr[eighthwarp_id] + vec_entry_index + bsize * m];\n }\n }\n\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n C_temp = C_temp + temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n\n // Contribution from each nonzero column\n int jmax = row_offsets[eighthwarp_id + 1];\n\n for (int jind = row_offsets[eighthwarp_id]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n offset = jcol * bsize + vec_entry_index;\n types::util::volcast(__cachingLoad(&B[offset]), s_xtemp + threadIdx.x);\n // Load nonzero_values\n s_offset = block_eighthwarp_id * bsize;\n\n if (ROW_MAJOR)\n {\n offset = jind * bsize_sq + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n }\n else\n {\n offset = jind * bsize_sq + vec_entry_index;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n temp[m] = nonzero_values[offset + bsize * m];\n }\n }\n\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n C_temp = C_temp + temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n }\n\n C[eighthwarp_id * bsize + vec_entry_index] = C_temp;\n eighthwarp_id += gridDim.x * blockDim.x >> log_bsize;\n }\n}\n\n#endif\n\n// implementation for arbitrary block size\ntemplate\n__global__\nvoid blockDiaCsrMultiplyKernelDiaProps(const IndexType *row_offsets,\n const IndexType *column_indices,\n const IndexType *dia_ptr,\n const ValueTypeA *nonzero_values,\n const ValueTypeB *B,\n ValueTypeB *C,\n const int num_block_rows,\n const int row_offset)\n{\n int warp_id = threadIdx.x / 32;\n int warp_thread_id = threadIdx.x & 31;\n\n // padding row blocks to fit in a single warp\n if ( warp_thread_id >= blockrows_per_warp * bsize ) { return; }\n\n // new thread id with padding\n int tid = warp_id * blockrows_per_warp * bsize + warp_thread_id;\n int eighthwarp_id = row_offset + blockIdx.x * blockrows_per_cta + tid / bsize;\n const int block_eighthwarp_id = tid / bsize;\n const int vec_entry_index = tid % bsize;\n const int bsize_sq = bsize * bsize;\n volatile __shared__ ValueTypeB s_xtemp[ bsize * blockrows_per_cta ];\n ValueTypeB C_temp;\n int offset, s_offset;\n\n while (eighthwarp_id < num_block_rows)\n {\n C_temp = types::util::get_zero();\n ValueTypeB temp[bsize];\n\n if ( diag )\n {\n // Contribution from diagonal\n offset = eighthwarp_id * bsize + vec_entry_index;\n types::util::volcast( __cachingLoad(&B[offset]), s_xtemp + tid);\n // Load dia_values and do matrix multiply\n s_offset = block_eighthwarp_id * bsize;\n\n if (ROW_MAJOR)\n {\n loadAsVector(nonzero_values + bsize_sq * dia_ptr[eighthwarp_id] + vec_entry_index * bsize, temp);\n }\n else\n {\n offset = dia_ptr[eighthwarp_id] * bsize_sq + vec_entry_index;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n types::util::to_uptype(nonzero_values[offset + bsize * m], temp[m]);\n }\n }\n\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n C_temp = C_temp + temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n }\n\n // Contribution from each nonzero column\n int jmax = row_offsets[eighthwarp_id + 1];\n\n for (int jind = row_offsets[eighthwarp_id]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n offset = jcol * bsize + vec_entry_index;\n types::util::volcast( __cachingLoad(&B[offset]), s_xtemp + tid);\n // Load nonzero_values\n s_offset = block_eighthwarp_id * bsize;\n\n if (ROW_MAJOR)\n {\n offset = jind * bsize_sq + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n }\n else\n {\n offset = jind * bsize_sq + vec_entry_index;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n types::util::to_uptype(nonzero_values[offset + bsize * m], temp[m]);\n }\n }\n\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n C_temp = C_temp + temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n }\n\n C[eighthwarp_id * bsize + vec_entry_index] = C_temp;\n eighthwarp_id += gridDim.x * blockrows_per_cta;\n }\n}\n\n// --------------------------------------\n// Methods\n// -------------------------------------\n\n// Method to perform BSPmV on host using block_dia_csr_matrix format\ntemplate \nvoid multiply_common_sqblock_host_diag(const Matrix &A, const Vector &B, Vector &C)\n{\n typedef typename Matrix::TConfig TConfig;\n\n if (TConfig::memSpace == AMGX_device)\n {\n FatalError(\"Executrion path error: device matrix in host path\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n //TODO:: This implementation is very inneficient, Use BLAS\n typedef typename TConfig::IndPrec IndexType;\n typedef typename TConfig::MatPrec ValueType;\n typedef typename Vector::value_type ValueTypeB;\n IndexType bsize = A.get_block_dimy();\n ValueTypeB temp;\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n // Initialize RHS to 0\n for (int m = 0; m < bsize; m++)\n {\n C[i * bsize + m] = types::util::get_zero();\n }\n\n // Contribution from diagonal blocks\n for (int n = 0; n < bsize; n++)\n {\n temp = B[i * bsize + n];\n\n for (int m = 0; m < bsize; m++)\n {\n C[i * bsize + m] = C[i * bsize + m] + A.values[A.diag[i] * bsize * bsize + m * bsize + n] * temp;\n }\n }\n\n // Contribution from nonzero off-diagonal blocks\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n IndexType jcol = A.col_indices[j];\n\n for (int n = 0; n < bsize; n++)\n {\n temp = B[jcol * bsize + n];\n\n for (int m = 0; m < bsize; m++)\n {\n C[i * bsize + m] = C[i * bsize + m] + A.values[j * bsize * bsize + m * bsize + n] * temp;\n }\n }\n }\n }\n }\n}\n\ntemplate \nvoid multiply_common_sqblock_host_nodiag(const Matrix &A, const Vector &B, Vector &C)\n{\n typedef typename Matrix::TConfig TConfig;\n\n if (TConfig::memSpace == AMGX_device)\n {\n FatalError(\"Executrion path error: device matrix in host path\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n //TODO:: This implementation is very inneficient, Use BLAS\n typedef typename TConfig::IndPrec IndexType;\n typedef typename TConfig::MatPrec ValueType;\n typedef typename Vector::value_type ValueTypeB;\n IndexType bsize = A.get_block_dimy();\n ValueTypeB temp;\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n // Initialize RHS to 0\n for (int m = 0; m < bsize; m++)\n {\n C[i * bsize + m] = types::util::get_zero();\n }\n\n // Contribution from nonzero blocks\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n IndexType jcol = A.col_indices[j];\n\n for (int n = 0; n < bsize; n++)\n {\n temp = B[jcol * bsize + n];\n\n for (int m = 0; m < bsize; m++)\n {\n C[i * bsize + m] = C[i * bsize + m] + A.values[j * bsize * bsize + m * bsize + n] * temp;\n }\n }\n }\n }\n }\n}\n\n\ntemplate \nclass Multiply_1x1\n{\n public:\n typedef typename Matrix::TConfig TConfig;\n static void multiply_1x1(Matrix &A, Vector &B, Vector &C, ViewType view)\n {\n if (TConfig::memSpace == AMGX_host)\n {\n if (A.hasProps(DIAG))\n {\n multiply_common_sqblock_host_diag(A, B, C);\n }\n else\n {\n multiply_common_sqblock_host_nodiag(A, B, C);\n }\n }\n else\n {\n typedef typename TConfig::VecPrec ValueTypeB;\n Cusparse::bsrmv(types::util::get_one(), A, B, types::util::get_zero(), C, view);\n cudaCheckError();\n // FatalError(\"Mixed precision is not supported for scalar matrix type\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n};\n\ntemplate \nclass Multiply_1x1_masked\n{\n public:\n typedef typename Matrix::TConfig TConfig;\n static void multiply_1x1_masked(Matrix &A, Vector &B, Vector &C, IVector mask, ViewType view)\n {\n if (TConfig::memSpace == AMGX_host)\n {\n FatalError(\"Masked multiply is not supported on host\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n typedef typename TConfig::VecPrec ValueTypeB;\n Cusparse::bsrxmv(types::util::get_one(), A, B, types::util::get_zero(), C, mask, view);\n cudaCheckError();\n }\n }\n};\n\ntemplate \nclass Multiply_1x1_with_mask\n{\n public:\n typedef typename Matrix::TConfig TConfig;\n static void multiply_1x1(Matrix &A, Vector &B, Vector &C)\n {\n if (TConfig::memSpace == AMGX_host)\n {\n FatalError(\"multiply with mask not supported on host\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n typedef typename TConfig::VecPrec ValueTypeB;\n Cusparse::bsrmv_with_mask(types::util::get_one(), A, B, types::util::get_zero(), C );\n cudaCheckError();\n // FatalError(\"Mixed precision is not supported for scalar matrix type\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n};\n\ntemplate \nclass Multiply_1x1_with_mask_restriction\n{\n public:\n typedef typename Matrix::TConfig TConfig;\n static void multiply_1x1(Matrix &A, Vector &B, Vector &C, Matrix &P)\n {\n if (TConfig::memSpace == AMGX_host)\n {\n FatalError(\"multiply with mask not supported on host\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n typedef typename TConfig::VecPrec ValueTypeB;\n Cusparse::bsrmv_with_mask_restriction(types::util::get_one(), A, B, types::util::get_zero(), C, P);\n cudaCheckError();\n // FatalError(\"Mixed precision is not supported for scalar matrix type\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n};\n\n\n// Method to perform BSPmV on device using block_dia_csr_matrix format\ntemplate \nclass Multiply_4x4\n{\n public:\n typedef typename Matrix::TConfig TConfig;\n static void multiply_4x4(const Matrix &A, const Vector &B, Vector &C, ViewType view)\n {\n if (TConfig::memSpace == AMGX_host)\n {\n if (A.hasProps(DIAG))\n {\n multiply_common_sqblock_host_diag(A, B, C);\n }\n else\n {\n multiply_common_sqblock_host_nodiag(A, B, C);\n }\n }\n else\n {\n //TODO: compare with cublas\n typedef typename TConfig::IndPrec IndexType;\n typedef typename TConfig::MatPrec ValueTypeA;\n typedef typename TConfig::VecPrec ValueTypeB;\n int num_rows, offset;\n A.getOffsetAndSizeForView(view, &offset, &num_rows);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_ind_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *B_ptr = B.raw();\n ValueTypeB *C_ptr = C.raw();\n cudaCheckError();\n const unsigned int threads_per_block = 128;\n const int eightwarps_per_block = threads_per_block / 4;\n const int num_warps_per_cta = threads_per_block / 32;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (num_rows + num_warps_per_cta - 1) / num_warps_per_cta); // (int) (A.get_num_rows()-1)/eightwarps_per_block + 1;\n\n if (!A.hasProps(DIAG))\n {\n if (A.getBlockFormat() == ROW_MAJOR)\n {\n cudaFuncSetCacheConfig(blockDiaCsrMultiplyKernel, cudaFuncCachePreferL1);\n blockDiaCsrMultiplyKernel <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n else\n {\n cudaFuncSetCacheConfig(blockDiaCsrMultiplyKernel, cudaFuncCachePreferL1);\n blockDiaCsrMultiplyKernel <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n\n }\n else\n {\n\n#ifdef USE_EXPERIMENTAL_4x4\n if ( A.getBlockFormat() == ROW_MAJOR )\n {\n blockDiaCsrMultiplyKernelDiaProps_4x4 <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n else\n {\n blockDiaCsrMultiplyKernelDiaProps_4x4 <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n#else\n\n if (A.getBlockFormat() == ROW_MAJOR)\n {\n cudaFuncSetCacheConfig(blockDiaCsrMultiplyKernelDiaProps_4x4, cudaFuncCachePreferL1);\n blockDiaCsrMultiplyKernelDiaProps_4x4 <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n else\n {\n cudaFuncSetCacheConfig(blockDiaCsrMultiplyKernelDiaProps_4x4, cudaFuncCachePreferL1);\n blockDiaCsrMultiplyKernelDiaProps_4x4 <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n\n#endif\n }\n cudaCheckError();\n }\n }\n};\n\n// Method to perform BSPmV on device using block_dia_csr_matrix format\ntemplate \nclass Multiply_bxb\n{\n public:\n typedef typename Matrix::TConfig TConfig;\n static void multiply_bxb(Matrix &A, Vector &B, Vector &C, ViewType view)\n {\n if (TConfig::memSpace == AMGX_host)\n {\n if (A.hasProps(DIAG))\n {\n multiply_common_sqblock_host_diag(A, B, C);\n }\n else\n {\n multiply_common_sqblock_host_nodiag(A, B, C);\n }\n }\n else\n {\n typedef typename TConfig::VecPrec ValueTypeB;\n Cusparse::bsrmv(types::util::get_one(), A, B, types::util::get_zero(), C, view);\n cudaCheckError();\n }\n }\n};\n\n\n// Method to perform BSPmV on device using block_dia_csr_matrix format\ntemplate \nclass Multiply_3x3\n{\n public:\n typedef typename Matrix::TConfig TConfig;\n static void multiply_3x3(const Matrix &A, const Vector &B, Vector &C, ViewType view)\n {\n if (TConfig::memSpace == AMGX_host)\n {\n if (A.hasProps(DIAG))\n {\n multiply_common_sqblock_host_diag(A, B, C);\n }\n else\n {\n multiply_common_sqblock_host_nodiag(A, B, C);\n }\n }\n else\n {\n //TODO: compare with cublas\n typedef typename TConfig::IndPrec IndexType;\n typedef typename TConfig::MatPrec ValueTypeA;\n typedef typename TConfig::VecPrec ValueTypeB;\n int num_rows, offset;\n A.getOffsetAndSizeForView(view, &offset, &num_rows);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_ind_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *B_ptr = B.raw();\n ValueTypeB *C_ptr = C.raw();\n cudaCheckError();\n const int threads_per_block = 64 * 3;\n const int blockrows_per_warp = 32 / 3;\n const int blockrows_per_cta = (threads_per_block / 32) * blockrows_per_warp;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (num_rows - 1) / blockrows_per_cta + 1); // (int) (A.get_num_rows()-1)/eightwarps_per_block + 1;\n\n if (!A.hasProps(DIAG))\n {\n if (A.getBlockFormat() == ROW_MAJOR)\n {\n cudaFuncSetCacheConfig(blockDiaCsrMultiplyKernelDiaProps, cudaFuncCachePreferL1);\n blockDiaCsrMultiplyKernelDiaProps <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n else\n {\n cudaFuncSetCacheConfig(blockDiaCsrMultiplyKernelDiaProps, cudaFuncCachePreferL1);\n blockDiaCsrMultiplyKernelDiaProps <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n }\n else\n {\n if (A.getBlockFormat() == ROW_MAJOR)\n {\n cudaFuncSetCacheConfig(blockDiaCsrMultiplyKernelDiaProps, cudaFuncCachePreferL1);\n blockDiaCsrMultiplyKernelDiaProps <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n else\n {\n cudaFuncSetCacheConfig(blockDiaCsrMultiplyKernelDiaProps, cudaFuncCachePreferL1);\n blockDiaCsrMultiplyKernelDiaProps <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n }\n }\n};\n\n// -------------------------------\n// Explict instantiations\n// -------------------------------\n#define AMGX_CASE_LINE(CASE) template void multiplyMM(const Matrix::Type>&, const Matrix::Type>&, Matrix::Type>&);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template void multiply(Matrix::Type> &, Vector::Type>&, Vector::Type> &, ViewType);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template void multiply_masked(Matrix::Type> &, Vector::Type> &, Vector::Type> &, typename Matrix::Type>::IVector &, ViewType);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template void multiply_with_mask(Matrix::Type> &, Vector::Type>&, Vector::Type> &);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template void multiply_with_mask_restriction(Matrix::Type> &, Vector::Type>&, Vector::Type> &, Matrix::Type> & );\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n\n#ifdef _WIN32\n#pragma warning (push)\n#pragma warning (disable : 4244 4267 4521)\n#endif\n#include \n#include \n#include \n#include \n#ifdef _WIN32\n#pragma warning (pop)\n#endif\n\n#include \n\n#include \n#include \n\n\nnamespace amgx\n{\n\n#define USE_EXPERIMENTAL_4x4\n\ntemplate \nclass Multiply_1x1;\ntemplate \nclass Multiply_1x1_masked;\ntemplate \nclass Multiply_1x1_with_mask;\ntemplate \nclass Multiply_1x1_with_mask_restriction;\ntemplate \nclass Multiply_3x3;\ntemplate \nclass Multiply_4x4;\ntemplate \nclass Multiply_bxb;\n\ntemplate \nvoid multiply_block_size(Matrix &A, Vector &B, Vector &C, ViewType view)\n{\n typedef Matrix TMatrix;\n typedef Vector TVector;\n\n if (A.get_block_size() == 1)\n {\n Multiply_1x1::multiply_1x1(A, B, C, view);\n }\n else if (A.get_block_dimy() == 3 && A.get_block_dimx() == 3)\n {\n Multiply_3x3::multiply_3x3(A, B, C, view);\n }\n else if (A.get_block_dimy() == 4 && A.get_block_dimx() == 4)\n {\n Multiply_4x4::multiply_4x4(A, B, C, view);\n }\n else\n {\n Multiply_bxb::multiply_bxb(A, B, C, view);\n }\n}\n\ntemplate \nvoid multiply(Matrix &A, Vector &B, Vector &C, ViewType view)\n{\n typedef Matrix TMatrix;\n typedef Vector TVector;\n\n if (!A.is_initialized())\n {\n FatalError(\"Trying to multiply uninitialized matrix\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (A.get_block_dimx() != B.get_block_dimy())\n {\n std::stringstream ss;\n ss << \"Matrix and vector dimensions don't match: A.dimx == \" << A.get_block_dimx() << \", B.dimy == \" << B.get_block_dimy();\n FatalError(ss.str().c_str(), AMGX_ERR_BAD_PARAMETERS);\n }\n\n\n typedef Matrix TMatrix;\n typedef Vector TVector;\n\n bool latencyHiding = (view == A.getViewExterior() && A.getViewInterior() != A.getViewExterior() && !A.is_matrix_singleGPU() && B.dirtybit != 0);\n\n if (latencyHiding)\n {\n A.manager->exchange_halo_split_gather(B, B.tag);\n\n // Multiply interior rows\n multiply_block_size(A, B, C, A.getViewInterior());\n\n // Finish halo exchange\n A.manager->exchange_halo_split_finish(B, B.tag);\n\n // Multiply rows with halo dependencies\n ViewType bnd_view = (ViewType)(~(A.getViewInterior()) & A.getViewExterior());\n multiply_block_size(A, B, C, bnd_view);\n }\n else\n {\n if (view != INTERIOR && !A.is_matrix_singleGPU() && B.dirtybit != 0)\n {\n A.manager->exchange_halo_v2(B, B.tag);\n }\n\n multiply_block_size(A, B, C, view);\n }\n\n C.dirtybit = 1;\n C.set_block_dimy(A.get_block_dimx());\n}\n\ntemplate \nvoid multiply_masked(Matrix &A, Vector &B, Vector &C, typename Matrix::IVector &mask, ViewType view)\n{\n typedef Matrix TMatrix;\n typedef Vector TVector;\n typedef typename Matrix::IVector TIVector;\n\n if (!A.is_initialized())\n {\n FatalError(\"Trying to multiply uninitialized matrix\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if(A.get_block_size() != 1)\n {\n FatalError(\"Unsupported blocksize for multiply_masked()\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (A.get_block_dimx() != B.get_block_dimy())\n {\n std::stringstream ss;\n ss << \"Matrix and vector dimensions don't match: A.dimx == \" << A.get_block_dimx() << \", B.dimy == \" << B.get_block_dimy();\n FatalError(ss.str().c_str(), AMGX_ERR_BAD_PARAMETERS);\n }\n\n Multiply_1x1_masked::multiply_1x1_masked(A, B, C, mask, view);\n C.set_block_dimy(A.get_block_dimx());\n}\n\ntemplate\nvoid multiply_with_mask(Matrix &A, Vector &B, Vector &C)\n{\n if (!A.is_initialized())\n {\n FatalError(\"Trying to multiply uninitialized matrix\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (A.get_block_dimx() != B.get_block_dimy())\n {\n std::stringstream ss;\n ss << \"Matrix and vector dimensions don't match: A.dimx == \" << A.get_block_dimx() << \", B.dimy == \" << B.get_block_dimy();\n FatalError(ss.str().c_str(), AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (A.get_block_size() == 1)\n {\n Multiply_1x1_with_mask::multiply_1x1(A, B, C);\n }\n else\n {\n FatalError(\"multiply with mask not supported for bsize != 1\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n C.set_block_dimy(A.get_block_dimx());\n C.dirtybit = 1;\n //if (!A.is_matrix_singleGPU() && C.size() == B.size() && C.delayed_send==0)\n // A.manager->exchange_halo_async(C, C.tag);\n}\n\n\ntemplate\nvoid multiply_with_mask_restriction(Matrix &A, Vector &B, Vector &C, Matrix &P)\n{\n if (!A.is_initialized())\n {\n FatalError(\"Trying to multiply uninitialized matrix\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (A.get_block_dimx() != B.get_block_dimy())\n {\n std::stringstream ss;\n ss << \"Matrix and vector dimensions don't match: A.dimx == \" << A.get_block_dimx() << \", B.dimy == \" << B.get_block_dimy();\n FatalError(ss.str().c_str(), AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (A.get_block_size() == 1)\n {\n Multiply_1x1_with_mask_restriction::multiply_1x1(A, B, C, P);\n }\n else\n {\n FatalError(\"multiply with mask not supported for bsize != 1\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n C.set_block_dimy(A.get_block_dimx());\n C.dirtybit = 1;\n}\n\n\n\n\ntemplate\nvoid multiplyMM(const Matrix &A, const Matrix &B, Matrix &C)\n{\n if (!A.is_initialized())\n {\n FatalError(\"Trying to multiply uninitialized matrix\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (A.get_block_dimx() != B.get_block_dimx() || A.get_block_dimy() != B.get_block_dimy())\n {\n FatalError(\"Matrices dimensions do not match\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (TConfig::memSpace == AMGX_device)\n {\n FatalError(\"Error, multiplyMM not implemented on device\", AMGX_ERR_BAD_PARAMETERS);\n }\n else\n {\n if (A.get_block_size() != 1)\n {\n FatalError(\"multiplyMM only works for block_size ==1\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n typedef typename TConfig::IndPrec IndexType;\n typedef typename TConfig::MatPrec ValueType;\n typedef typename Matrix::IVector IVector;\n typedef typename Matrix::MVector MVector;\n C.set_initialized(0);\n IndexType num_nonzeros = 0;\n IVector mask(B.get_num_cols(), IndexType (-1));\n\n // Compute nnz in C (including explicit zeros)\n for (size_t i = 0; i < A.get_num_rows(); i++)\n {\n for (IndexType jj = A.row_offsets[i]; jj < A.row_offsets[i + 1]; jj++)\n {\n IndexType j = A.col_indices[jj];\n\n for (IndexType kk = B.row_offsets[j]; kk < B.row_offsets[j + 1]; kk++)\n {\n IndexType k = B.col_indices[kk];\n\n if (mask[k] != i)\n {\n mask[k] = i;\n num_nonzeros++;\n }\n }\n }\n }\n\n // Resize output\n C.resize(A.get_num_rows(), B.get_num_cols(), num_nonzeros);\n const IndexType unseen = static_cast(-1);\n const IndexType init = static_cast(-2);\n // Compute entries of C\n IVector next(B.get_num_cols(), unseen);\n MVector sums(B.get_num_cols(), types::util::get_zero());\n num_nonzeros = 0;\n C.row_offsets[0] = 0;\n\n for (size_t i = 0; i < A.get_num_rows(); i++)\n {\n IndexType head = init;\n IndexType length = 0;\n IndexType jj_start = A.row_offsets[i];\n IndexType jj_end = A.row_offsets[i + 1];\n\n for (IndexType jj = jj_start; jj < jj_end; jj++)\n {\n IndexType j = A.col_indices[jj];\n ValueType v = A.values[jj];\n IndexType kk_start = B.row_offsets[j];\n IndexType kk_end = B.row_offsets[j + 1];\n\n for (IndexType kk = kk_start; kk < kk_end; kk++)\n {\n IndexType k = B.col_indices[kk];\n sums[k] = sums[k] + v * B.values[kk];\n\n if (next[k] == unseen)\n {\n next[k] = head;\n head = k;\n length++;\n }\n }\n }\n\n for (IndexType jj = 0; jj < length; jj++)\n {\n //if(sums[head] != ValueType(0))\n //{\n C.col_indices[num_nonzeros] = head;\n C.values[num_nonzeros] = sums[head];\n num_nonzeros++;\n //}\n IndexType temp = head;\n head = next[head];\n // clear arrays\n next[temp] = unseen;\n sums[temp] = types::util::get_zero();\n }\n\n C.row_offsets[i + 1] = num_nonzeros;\n }\n\n // Resize output again since pass2 omits explict zeros\n //C.resize(A.num_rows, B.num_cols, num_nonzeros);\n C.set_initialized(1);\n }\n }\n}\n\n// --------------------------------\n// KERNELS\n// --------------------------------\n\ntemplate\n__global__\nvoid blockDiaCsrMultiplyKernel(const IndexType *row_offsets,\n const IndexType *column_indices,\n const ValueTypeA *nonzero_values,\n const ValueTypeB *B,\n ValueTypeB *C,\n const IndexType num_block_rows,\n const IndexType row_offset)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int eighthwarp_id = row_offset + (tid >> log_bsize);\n const int block_eighthwarp_id = threadIdx.x >> log_bsize;\n const int vec_entry_index = threadIdx.x & (bsize - 1);\n volatile __shared__ ValueTypeB s_xtemp[ bsize * eighthwarps_per_block ];\n ValueTypeB C_temp;\n int offset, s_offset;\n ValueTypeB temp[bsize];\n\n while (eighthwarp_id < num_block_rows)\n {\n //i = eighthwarp_id;\n C_temp = types::util::get_zero();\n // Contribution from each nonzero column\n int jmax = row_offsets[eighthwarp_id + 1];\n\n for (int jind = row_offsets[eighthwarp_id]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n offset = jcol * bsize + vec_entry_index;\n types::util::volcast(__cachingLoad(&B[offset]), s_xtemp + threadIdx.x);\n // Load nonzero_values\n s_offset = block_eighthwarp_id * bsize;\n\n if (ROW_MAJOR)\n {\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n }\n else\n {\n offset = jind * bsize * bsize + vec_entry_index;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n types::util::to_uptype(nonzero_values[offset + bsize * m], temp[m]);\n }\n }\n\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n C_temp = C_temp + temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n }\n\n C[eighthwarp_id * bsize + vec_entry_index] = C_temp;\n eighthwarp_id += gridDim.x * blockDim.x >> log_bsize;\n }\n}\n\n#ifdef USE_EXPERIMENTAL_4x4\n\ntemplate< typename IndexType, typename ValueTypeA, typename ValueTypeB, int CTA_SIZE, bool ROW_MAJOR >\n__global__\n__launch_bounds__( CTA_SIZE, 16 )\nvoid blockDiaCsrMultiplyKernelDiaProps_4x4( const IndexType *row_offsets,\n const IndexType *column_indices,\n const IndexType *dia_ptr,\n const ValueTypeA *nonzero_values,\n const ValueTypeB *B,\n ValueTypeB *C,\n const int num_block_rows,\n const int row_offset )\n{\n const int nHalfWarps = CTA_SIZE / 16; // Number of half warps per CTA.\n const int laneId = threadIdx.x % warpSize;\n const int halfWarpId = threadIdx.x / 16;\n const int halfLaneId = threadIdx.x % 16;\n const int halfLaneId_div_4 = halfLaneId / 4;\n const int halfLaneId_mod_4 = halfLaneId % 4;\n const int laneId_div_16 = laneId / 16;\n const int upperHalf = 16 * laneId_div_16;\n // Shared memory needed to exchange X and delta.\n __shared__ volatile ValueTypeB s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile ValueTypeB *my_s_mem = &s_mem[16 * halfWarpId];\n\n // Iterate over the rows of the matrix. One warp per two rows.\n for ( int aRowId = blockIdx.x * nHalfWarps + halfWarpId ; aRowId < num_block_rows ; aRowId += gridDim.x * nHalfWarps )\n {\n unsigned int active_mask = utils::activemask();\n // Load one block of B.\n ValueTypeB my_Ax = types::util::get_zero();\n\n // The diagonal.\n if ( halfLaneId_div_4 == 0 )\n {\n types::util::volcast( B[4 * aRowId + halfLaneId_mod_4], my_s_mem + halfLaneId);\n }\n\n // Load the diagonal.\n int diagId = dia_ptr[aRowId];\n // Update my values.\n ValueTypeA my_val = nonzero_values[16 * diagId + halfLaneId];\n\n if ( ROW_MAJOR )\n {\n my_Ax = my_Ax + my_val * types::util::volcast(my_s_mem[halfLaneId_mod_4]);\n }\n else\n {\n my_Ax = my_Ax + my_val * types::util::volcast(my_s_mem[halfLaneId_div_4]);\n }\n\n // The range of the rows.\n int aColBegin = row_offsets[aRowId + 0];\n int aColEnd = row_offsets[aRowId + 1];\n\n // Each warp load column indices of 16 nonzero blocks\n for ( ; utils::any( aColBegin < aColEnd, active_mask ) ; aColBegin += 16 )\n {\n int aColIt = aColBegin + halfLaneId;\n // Get the ID of the column.\n int aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColId = column_indices[aColIt];\n }\n\n // Count the number of active columns.\n int vote = utils::ballot(aColId != -1, active_mask);\n // The number of iterations.\n int nCols = max( __popc( vote & 0x0000ffff ), __popc( vote & 0xffff0000 ) );\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < nCols ; k += 4 )\n {\n int my_k = k + halfLaneId_div_4;\n // Exchange column indices.\n int waColId = utils::shfl( aColId, upperHalf + my_k, warpSize, active_mask );\n // Load 8 blocks of X if needed.\n ValueTypeB my_x = types::util::get_zero();\n\n if ( waColId != -1 )\n {\n my_x = B[4 * waColId + halfLaneId_mod_4];\n }\n\n types::util::volcast( my_x, my_s_mem + halfLaneId);\n // Load 8 blocks of A.\n#pragma unroll\n\n for ( int i = 0 ; i < 4 ; ++i )\n {\n int w_aColTmp = aColBegin + k + i, w_aColIt = -1;\n\n if ( w_aColTmp < aColEnd )\n {\n w_aColIt = w_aColTmp;\n }\n\n ValueTypeA my_val = types::util::get_zero();\n\n if ( w_aColIt != -1 )\n {\n my_val = nonzero_values[16 * w_aColIt + halfLaneId];\n }\n\n if ( ROW_MAJOR )\n {\n my_Ax = my_Ax + my_val * types::util::volcast(my_s_mem[4 * i + halfLaneId_mod_4]);\n }\n else\n {\n my_Ax = my_Ax + my_val * types::util::volcast(my_s_mem[4 * i + halfLaneId_div_4]);\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Reduce bmAx terms.\n if ( ROW_MAJOR )\n {\n my_Ax = my_Ax + utils::shfl_xor( my_Ax, 1, warpSize, active_mask );\n my_Ax = my_Ax + utils::shfl_xor( my_Ax, 2, warpSize, active_mask );\n }\n else\n {\n my_Ax = my_Ax + utils::shfl_xor( my_Ax, 4, warpSize, active_mask );\n my_Ax = my_Ax + utils::shfl_xor( my_Ax, 8, warpSize, active_mask );\n }\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n if ( halfLaneId_mod_4 == 0 )\n {\n C[4 * aRowId + halfLaneId_div_4] = my_Ax;\n }\n }\n else\n {\n if ( halfLaneId_div_4 == 0 )\n {\n C[4 * aRowId + halfLaneId_mod_4] = my_Ax;\n }\n }\n }\n}\n\n#else\n\ntemplate\n__global__\nvoid blockDiaCsrMultiplyKernelDiaProps_4x4(const IndexType *row_offsets,\n const IndexType *column_indices,\n const IndexType *dia_ptr,\n const ValueTypeA *nonzero_values,\n const ValueTypeB *B,\n ValueTypeB *C,\n const int num_block_rows,\n const int row_offset)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int eighthwarp_id = row_offset + (tid >> log_bsize);\n const int block_eighthwarp_id = threadIdx.x >> log_bsize;\n const int vec_entry_index = threadIdx.x & (bsize - 1);\n volatile __shared__ ValueTypeB s_xtemp[ bsize * eighthwarps_per_block ];\n ValueTypeB C_temp;\n int offset, s_offset;\n\n while (eighthwarp_id < num_block_rows)\n {\n //i = eighthwarp_id;\n C_temp = types::util::get_zero();\n // Contribution from diagonal\n offset = eighthwarp_id * bsize + vec_entry_index;\n types::util::volcast(__cachingLoad(&B[offset]), s_xtemp + threadIdx.x);\n // Load dia_values and do matrix multiply\n s_offset = block_eighthwarp_id * bsize;\n ValueTypeA temp[bsize];\n\n if (ROW_MAJOR)\n {\n loadAsVector(nonzero_values + bsize_sq * dia_ptr[eighthwarp_id] + vec_entry_index * bsize, temp);\n }\n else\n {\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n temp[m] = nonzero_values[bsize_sq * dia_ptr[eighthwarp_id] + vec_entry_index + bsize * m];\n }\n }\n\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n C_temp = C_temp + temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n\n // Contribution from each nonzero column\n int jmax = row_offsets[eighthwarp_id + 1];\n\n for (int jind = row_offsets[eighthwarp_id]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n offset = jcol * bsize + vec_entry_index;\n types::util::volcast(__cachingLoad(&B[offset]), s_xtemp + threadIdx.x);\n // Load nonzero_values\n s_offset = block_eighthwarp_id * bsize;\n\n if (ROW_MAJOR)\n {\n offset = jind * bsize_sq + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n }\n else\n {\n offset = jind * bsize_sq + vec_entry_index;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n temp[m] = nonzero_values[offset + bsize * m];\n }\n }\n\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n C_temp = C_temp + temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n }\n\n C[eighthwarp_id * bsize + vec_entry_index] = C_temp;\n eighthwarp_id += gridDim.x * blockDim.x >> log_bsize;\n }\n}\n\n#endif\n\n// implementation for arbitrary block size\ntemplate\n__global__\nvoid blockDiaCsrMultiplyKernelDiaProps(const IndexType *row_offsets,\n const IndexType *column_indices,\n const IndexType *dia_ptr,\n const ValueTypeA *nonzero_values,\n const ValueTypeB *B,\n ValueTypeB *C,\n const int num_block_rows,\n const int row_offset)\n{\n int warp_id = threadIdx.x / 32;\n int warp_thread_id = threadIdx.x & 31;\n\n // padding row blocks to fit in a single warp\n if ( warp_thread_id >= blockrows_per_warp * bsize ) { return; }\n\n // new thread id with padding\n int tid = warp_id * blockrows_per_warp * bsize + warp_thread_id;\n int eighthwarp_id = row_offset + blockIdx.x * blockrows_per_cta + tid / bsize;\n const int block_eighthwarp_id = tid / bsize;\n const int vec_entry_index = tid % bsize;\n const int bsize_sq = bsize * bsize;\n volatile __shared__ ValueTypeB s_xtemp[ bsize * blockrows_per_cta ];\n ValueTypeB C_temp;\n int offset, s_offset;\n\n while (eighthwarp_id < num_block_rows)\n {\n C_temp = types::util::get_zero();\n ValueTypeB temp[bsize];\n\n if ( diag )\n {\n // Contribution from diagonal\n offset = eighthwarp_id * bsize + vec_entry_index;\n types::util::volcast( __cachingLoad(&B[offset]), s_xtemp + tid);\n // Load dia_values and do matrix multiply\n s_offset = block_eighthwarp_id * bsize;\n\n if (ROW_MAJOR)\n {\n loadAsVector(nonzero_values + bsize_sq * dia_ptr[eighthwarp_id] + vec_entry_index * bsize, temp);\n }\n else\n {\n offset = dia_ptr[eighthwarp_id] * bsize_sq + vec_entry_index;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n types::util::to_uptype(nonzero_values[offset + bsize * m], temp[m]);\n }\n }\n\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n C_temp = C_temp + temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n }\n\n // Contribution from each nonzero column\n int jmax = row_offsets[eighthwarp_id + 1];\n\n for (int jind = row_offsets[eighthwarp_id]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n offset = jcol * bsize + vec_entry_index;\n types::util::volcast( __cachingLoad(&B[offset]), s_xtemp + tid);\n // Load nonzero_values\n s_offset = block_eighthwarp_id * bsize;\n\n if (ROW_MAJOR)\n {\n offset = jind * bsize_sq + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n }\n else\n {\n offset = jind * bsize_sq + vec_entry_index;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n types::util::to_uptype(nonzero_values[offset + bsize * m], temp[m]);\n }\n }\n\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n C_temp = C_temp + temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n }\n\n C[eighthwarp_id * bsize + vec_entry_index] = C_temp;\n eighthwarp_id += gridDim.x * blockrows_per_cta;\n }\n}\n\n// --------------------------------------\n// Methods\n// -------------------------------------\n\n// Method to perform BSPmV on host using block_dia_csr_matrix format\ntemplate \nvoid multiply_common_sqblock_host_diag(const Matrix &A, const Vector &B, Vector &C)\n{\n typedef typename Matrix::TConfig TConfig;\n\n if (TConfig::memSpace == AMGX_device)\n {\n FatalError(\"Executrion path error: device matrix in host path\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n //TODO:: This implementation is very inneficient, Use BLAS\n typedef typename TConfig::IndPrec IndexType;\n typedef typename TConfig::MatPrec ValueType;\n typedef typename Vector::value_type ValueTypeB;\n IndexType bsize = A.get_block_dimy();\n ValueTypeB temp;\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n // Initialize RHS to 0\n for (int m = 0; m < bsize; m++)\n {\n C[i * bsize + m] = types::util::get_zero();\n }\n\n // Contribution from diagonal blocks\n for (int n = 0; n < bsize; n++)\n {\n temp = B[i * bsize + n];\n\n for (int m = 0; m < bsize; m++)\n {\n C[i * bsize + m] = C[i * bsize + m] + A.values[A.diag[i] * bsize * bsize + m * bsize + n] * temp;\n }\n }\n\n // Contribution from nonzero off-diagonal blocks\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n IndexType jcol = A.col_indices[j];\n\n for (int n = 0; n < bsize; n++)\n {\n temp = B[jcol * bsize + n];\n\n for (int m = 0; m < bsize; m++)\n {\n C[i * bsize + m] = C[i * bsize + m] + A.values[j * bsize * bsize + m * bsize + n] * temp;\n }\n }\n }\n }\n }\n}\n\ntemplate \nvoid multiply_common_sqblock_host_nodiag(const Matrix &A, const Vector &B, Vector &C)\n{\n typedef typename Matrix::TConfig TConfig;\n\n if (TConfig::memSpace == AMGX_device)\n {\n FatalError(\"Executrion path error: device matrix in host path\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n //TODO:: This implementation is very inneficient, Use BLAS\n typedef typename TConfig::IndPrec IndexType;\n typedef typename TConfig::MatPrec ValueType;\n typedef typename Vector::value_type ValueTypeB;\n IndexType bsize = A.get_block_dimy();\n ValueTypeB temp;\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n // Initialize RHS to 0\n for (int m = 0; m < bsize; m++)\n {\n C[i * bsize + m] = types::util::get_zero();\n }\n\n // Contribution from nonzero blocks\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n IndexType jcol = A.col_indices[j];\n\n for (int n = 0; n < bsize; n++)\n {\n temp = B[jcol * bsize + n];\n\n for (int m = 0; m < bsize; m++)\n {\n C[i * bsize + m] = C[i * bsize + m] + A.values[j * bsize * bsize + m * bsize + n] * temp;\n }\n }\n }\n }\n }\n}\n\n\ntemplate \nclass Multiply_1x1\n{\n public:\n typedef typename Matrix::TConfig TConfig;\n static void multiply_1x1(Matrix &A, Vector &B, Vector &C, ViewType view)\n {\n if (TConfig::memSpace == AMGX_host)\n {\n if (A.hasProps(DIAG))\n {\n multiply_common_sqblock_host_diag(A, B, C);\n }\n else\n {\n multiply_common_sqblock_host_nodiag(A, B, C);\n }\n }\n else\n {\n typedef typename TConfig::VecPrec ValueTypeB;\n Cusparse::bsrmv(types::util::get_one(), A, B, types::util::get_zero(), C, view);\n cudaCheckError();\n // FatalError(\"Mixed precision is not supported for scalar matrix type\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n};\n\ntemplate \nclass Multiply_1x1_masked\n{\n public:\n typedef typename Matrix::TConfig TConfig;\n static void multiply_1x1_masked(Matrix &A, Vector &B, Vector &C, IVector mask, ViewType view)\n {\n if (TConfig::memSpace == AMGX_host)\n {\n FatalError(\"Masked multiply is not supported on host\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n typedef typename TConfig::VecPrec ValueTypeB;\n Cusparse::bsrxmv(types::util::get_one(), A, B, types::util::get_zero(), C, mask, view);\n cudaCheckError();\n }\n }\n};\n\ntemplate \nclass Multiply_1x1_with_mask\n{\n public:\n typedef typename Matrix::TConfig TConfig;\n static void multiply_1x1(Matrix &A, Vector &B, Vector &C)\n {\n if (TConfig::memSpace == AMGX_host)\n {\n FatalError(\"multiply with mask not supported on host\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n typedef typename TConfig::VecPrec ValueTypeB;\n Cusparse::bsrmv_with_mask(types::util::get_one(), A, B, types::util::get_zero(), C );\n cudaCheckError();\n // FatalError(\"Mixed precision is not supported for scalar matrix type\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n};\n\ntemplate \nclass Multiply_1x1_with_mask_restriction\n{\n public:\n typedef typename Matrix::TConfig TConfig;\n static void multiply_1x1(Matrix &A, Vector &B, Vector &C, Matrix &P)\n {\n if (TConfig::memSpace == AMGX_host)\n {\n FatalError(\"multiply with mask not supported on host\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n typedef typename TConfig::VecPrec ValueTypeB;\n Cusparse::bsrmv_with_mask_restriction(types::util::get_one(), A, B, types::util::get_zero(), C, P);\n cudaCheckError();\n // FatalError(\"Mixed precision is not supported for scalar matrix type\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n};\n\n\n// Method to perform BSPmV on device using block_dia_csr_matrix format\ntemplate \nclass Multiply_4x4\n{\n public:\n typedef typename Matrix::TConfig TConfig;\n static void multiply_4x4(const Matrix &A, const Vector &B, Vector &C, ViewType view)\n {\n if (TConfig::memSpace == AMGX_host)\n {\n if (A.hasProps(DIAG))\n {\n multiply_common_sqblock_host_diag(A, B, C);\n }\n else\n {\n multiply_common_sqblock_host_nodiag(A, B, C);\n }\n }\n else\n {\n //TODO: compare with cublas\n typedef typename TConfig::IndPrec IndexType;\n typedef typename TConfig::MatPrec ValueTypeA;\n typedef typename TConfig::VecPrec ValueTypeB;\n int num_rows, offset;\n A.getOffsetAndSizeForView(view, &offset, &num_rows);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_ind_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *B_ptr = B.raw();\n ValueTypeB *C_ptr = C.raw();\n cudaCheckError();\n const unsigned int threads_per_block = 128;\n const int eightwarps_per_block = threads_per_block / 4;\n const int num_warps_per_cta = threads_per_block / 32;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (num_rows + num_warps_per_cta - 1) / num_warps_per_cta); // (int) (A.get_num_rows()-1)/eightwarps_per_block + 1;\n\n if (!A.hasProps(DIAG))\n {\n if (A.getBlockFormat() == ROW_MAJOR)\n {\n hipFuncSetCacheConfig(reinterpret_cast(blockDiaCsrMultiplyKernel), hipFuncCachePreferL1);\n blockDiaCsrMultiplyKernel <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n else\n {\n hipFuncSetCacheConfig(reinterpret_cast(blockDiaCsrMultiplyKernel), hipFuncCachePreferL1);\n blockDiaCsrMultiplyKernel <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n\n }\n else\n {\n\n#ifdef USE_EXPERIMENTAL_4x4\n if ( A.getBlockFormat() == ROW_MAJOR )\n {\n blockDiaCsrMultiplyKernelDiaProps_4x4 <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n else\n {\n blockDiaCsrMultiplyKernelDiaProps_4x4 <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n#else\n\n if (A.getBlockFormat() == ROW_MAJOR)\n {\n hipFuncSetCacheConfig(reinterpret_cast(blockDiaCsrMultiplyKernelDiaProps_4x4), hipFuncCachePreferL1);\n blockDiaCsrMultiplyKernelDiaProps_4x4 <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n else\n {\n hipFuncSetCacheConfig(reinterpret_cast(blockDiaCsrMultiplyKernelDiaProps_4x4), hipFuncCachePreferL1);\n blockDiaCsrMultiplyKernelDiaProps_4x4 <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n\n#endif\n }\n cudaCheckError();\n }\n }\n};\n\n// Method to perform BSPmV on device using block_dia_csr_matrix format\ntemplate \nclass Multiply_bxb\n{\n public:\n typedef typename Matrix::TConfig TConfig;\n static void multiply_bxb(Matrix &A, Vector &B, Vector &C, ViewType view)\n {\n if (TConfig::memSpace == AMGX_host)\n {\n if (A.hasProps(DIAG))\n {\n multiply_common_sqblock_host_diag(A, B, C);\n }\n else\n {\n multiply_common_sqblock_host_nodiag(A, B, C);\n }\n }\n else\n {\n typedef typename TConfig::VecPrec ValueTypeB;\n Cusparse::bsrmv(types::util::get_one(), A, B, types::util::get_zero(), C, view);\n cudaCheckError();\n }\n }\n};\n\n\n// Method to perform BSPmV on device using block_dia_csr_matrix format\ntemplate \nclass Multiply_3x3\n{\n public:\n typedef typename Matrix::TConfig TConfig;\n static void multiply_3x3(const Matrix &A, const Vector &B, Vector &C, ViewType view)\n {\n if (TConfig::memSpace == AMGX_host)\n {\n if (A.hasProps(DIAG))\n {\n multiply_common_sqblock_host_diag(A, B, C);\n }\n else\n {\n multiply_common_sqblock_host_nodiag(A, B, C);\n }\n }\n else\n {\n //TODO: compare with cublas\n typedef typename TConfig::IndPrec IndexType;\n typedef typename TConfig::MatPrec ValueTypeA;\n typedef typename TConfig::VecPrec ValueTypeB;\n int num_rows, offset;\n A.getOffsetAndSizeForView(view, &offset, &num_rows);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_ind_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *B_ptr = B.raw();\n ValueTypeB *C_ptr = C.raw();\n cudaCheckError();\n const int threads_per_block = 64 * 3;\n const int blockrows_per_warp = 32 / 3;\n const int blockrows_per_cta = (threads_per_block / 32) * blockrows_per_warp;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (num_rows - 1) / blockrows_per_cta + 1); // (int) (A.get_num_rows()-1)/eightwarps_per_block + 1;\n\n if (!A.hasProps(DIAG))\n {\n if (A.getBlockFormat() == ROW_MAJOR)\n {\n hipFuncSetCacheConfig(reinterpret_cast(blockDiaCsrMultiplyKernelDiaProps), hipFuncCachePreferL1);\n blockDiaCsrMultiplyKernelDiaProps <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n else\n {\n hipFuncSetCacheConfig(reinterpret_cast(blockDiaCsrMultiplyKernelDiaProps), hipFuncCachePreferL1);\n blockDiaCsrMultiplyKernelDiaProps <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n }\n else\n {\n if (A.getBlockFormat() == ROW_MAJOR)\n {\n hipFuncSetCacheConfig(reinterpret_cast(blockDiaCsrMultiplyKernelDiaProps), hipFuncCachePreferL1);\n blockDiaCsrMultiplyKernelDiaProps <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n else\n {\n hipFuncSetCacheConfig(reinterpret_cast(blockDiaCsrMultiplyKernelDiaProps), hipFuncCachePreferL1);\n blockDiaCsrMultiplyKernelDiaProps <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ind_ptr, A_nonzero_values_ptr, B_ptr, C_ptr, offset + num_rows, offset);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n }\n }\n};\n\n// -------------------------------\n// Explict instantiations\n// -------------------------------\n#define AMGX_CASE_LINE(CASE) template void multiplyMM(const Matrix::Type>&, const Matrix::Type>&, Matrix::Type>&);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template void multiply(Matrix::Type> &, Vector::Type>&, Vector::Type> &, ViewType);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template void multiply_masked(Matrix::Type> &, Vector::Type> &, Vector::Type> &, typename Matrix::Type>::IVector &, ViewType);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template void multiply_with_mask(Matrix::Type> &, Vector::Type>&, Vector::Type> &);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template void multiply_with_mask_restriction(Matrix::Type> &, Vector::Type>&, Vector::Type> &, Matrix::Type> & );\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/multiply.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n\n#ifdef _WIN32\n#pragma warning (push)\n#pragma warning (disable : 4244 4267 4521)\n#endif\n#ifdef _WIN32\n#pragma warning (pop)\n#endif\n\nnamespace amgx\n{\n\n// Kernel to invert 4by4 matrix on device using one thread\n// TODO: possible optimizations using multiple threads, possible generalisation for NxN blocks\ntemplate \n__global__\nvoid gaussianElimination4by4Kernel(const IndexType *dia_indices, const ValueTypeA *values, ValueTypeB *x, const ValueTypeB *b)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n ValueTypeB Atemp[BSIZE][BSIZE];\n ValueTypeB xtemp[BSIZE];\n ValueTypeB btemp[BSIZE];\n int offset;\n\n if (tid == 0)\n {\n#pragma unroll\n\n for (int m = 0; m < BSIZE; m++)\n {\n btemp[m] = b[m];\n }\n\n offset = 0;\n#pragma unroll\n\n for (int m = 0; m < BSIZE; m++)\n#pragma unroll\n for (int n = 0; n < BSIZE; n++)\n {\n types::util::to_uptype(values[offset++], Atemp[m][n]); // only one nonzero block - diagonal\n }\n\n gaussianElimination4by4(Atemp, xtemp, btemp);\n // Store x\n#pragma unroll\n\n for (int m = 0; m < BSIZE; m++)\n {\n x[m] = xtemp[m];\n }\n }\n}\n\n// Method to perform gaussian elimination for a matrix stored in row major format\ntemplate\nvoid gaussianEliminationRowMajor(ValueTypeA **e, ValueTypeB *x, ValueTypeB *b, const IndexType bsize)\n{\n if (bsize > 4) { FatalError(\"Warning, Gaussian Elimination code doesn't pivot and bsize > 4\", AMGX_ERR_BAD_PARAMETERS); }\n\n ValueTypeB pivot, ratio, tmp;\n\n for (int j = 0; j < bsize; j++)\n {\n types::util::to_uptype(e[j][j], pivot);\n\n for (int k = j + 1; k < bsize; k++)\n {\n ValueTypeB temp;\n types::util::to_uptype(e[k][j], temp);\n ratio = temp / pivot;\n\n for (int m = j + 1; m < bsize; m++)\n {\n e[k][m] = e[k][m] - e[j][m] * ratio;\n }\n\n b[k] = b[k] - b[j] * ratio;\n }\n }\n\n // back substitution\n for (int j = bsize - 1; j >= 0; j--)\n {\n tmp = types::util::get_zero();\n\n for (int k = j + 1; k < bsize; k++)\n {\n tmp = tmp + e[j][k] * x[k];\n }\n\n x[j] = (b[j] - tmp) / e[j][j];\n }\n}\n\n// Method to perform gaussian elimination on matrix stored in block_dia_csr_matrix_h\ntemplate \nvoid GaussianElimination >::gaussianElimination_4x4_host(const Matrix_h &A, Vector_h &x, const Vector_h &b)\n{\n typedef typename Matrix_h::value_type ValueTypeA;\n typedef typename Vector_h::value_type ValueTypeB;\n\n if (A.get_num_rows() != 1)\n {\n FatalError(\"Haven't implemented gaussian elimination solver for num_blocks != 1\", AMGX_ERR_BAD_PARAMETERS);\n }\n else\n {\n int bsize = A.get_block_dimy();\n // Allocate space for block_matrix\n ValueTypeA **Atemp = new ValueTypeA* [bsize];\n\n for ( int i = 0; i < bsize; i++)\n {\n Atemp[i] = new ValueTypeA[bsize];\n }\n\n ValueTypeB *btemp = new ValueTypeB[bsize];\n ValueTypeB *xtemp = new ValueTypeB[bsize];\n\n // Copy matrix and rhs\n for (int m = 0; m < bsize; m++)\n {\n for (int n = 0; n < bsize; n++)\n {\n Atemp[m][n] = A.values[bsize * m + n]; //diag[bsize*m + n]; since we have 1 nonzero only\n }\n\n btemp[m] = b[m];\n }\n\n ValueTypeB *x_ptr = x.raw();\n gaussianEliminationRowMajor(Atemp, x_ptr, btemp, bsize);\n }\n}\n\n// Method to perform gaussian elimination on matrix stored in block_dia_csr_matrix_d\ntemplate \nvoid GaussianElimination >::gaussianElimination_4x4_device(const Matrix_d &A, Vector_d &x, const Vector_d &b)\n{\n typedef typename Matrix_d::value_type ValueTypeA;\n typedef typename Vector_d::value_type ValueTypeB;\n typedef typename Matrix_d::index_type IndexType;\n\n if (A.get_num_rows() != 1)\n {\n FatalError(\"Haven't implemented gaussian elimination solver for num_blocks != 1 or block_size != 4\", AMGX_ERR_BAD_PARAMETERS);\n }\n else\n {\n const IndexType *dia_ptr = A.diag.raw();\n const ValueTypeA *values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n ValueTypeB *x_ptr = x.raw();\n gaussianElimination4by4Kernel <<< 1, 1>>>(dia_ptr, values_ptr, x_ptr, b_ptr);\n cudaCheckError();\n }\n}\n\ntemplate \nvoid GaussianElimination::gaussianElimination(const Matrix &A, Vector &x, const Vector &b)\n{\n FatalError(\"Matrix format not supported in gaussian_elimination\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid GaussianElimination >::gaussianElimination(const Matrix_h &A, Vector_h &x, const Vector_h &b)\n{\n if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n gaussianElimination_4x4_host(A, x, b);\n }\n else\n {\n FatalError(\"gaussElimination: Blocksize is unsupported\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\ntemplate \nvoid GaussianElimination >::gaussianElimination(const Matrix_d &A, Vector_d &x, const Vector_d &b)\n{\n if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n gaussianElimination_4x4_device(A, x, b);\n }\n else\n {\n FatalError(\"gaussElimination: Blocksize is unsupported\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n// -------------------------------\n// Explict instantiations\n// -------------------------------\n#define AMGX_CASE_LINE(CASE) template class GaussianElimination::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n#ifdef _WIN32\n#pragma warning (push)\n#pragma warning (disable : 4244 4267 4521)\n#endif\n#ifdef _WIN32\n#pragma warning (pop)\n#endif\n\nnamespace amgx\n{\n\n// Kernel to invert 4by4 matrix on device using one thread\n// TODO: possible optimizations using multiple threads, possible generalisation for NxN blocks\ntemplate \n__global__\nvoid gaussianElimination4by4Kernel(const IndexType *dia_indices, const ValueTypeA *values, ValueTypeB *x, const ValueTypeB *b)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n ValueTypeB Atemp[BSIZE][BSIZE];\n ValueTypeB xtemp[BSIZE];\n ValueTypeB btemp[BSIZE];\n int offset;\n\n if (tid == 0)\n {\n#pragma unroll\n\n for (int m = 0; m < BSIZE; m++)\n {\n btemp[m] = b[m];\n }\n\n offset = 0;\n#pragma unroll\n\n for (int m = 0; m < BSIZE; m++)\n#pragma unroll\n for (int n = 0; n < BSIZE; n++)\n {\n types::util::to_uptype(values[offset++], Atemp[m][n]); // only one nonzero block - diagonal\n }\n\n gaussianElimination4by4(Atemp, xtemp, btemp);\n // Store x\n#pragma unroll\n\n for (int m = 0; m < BSIZE; m++)\n {\n x[m] = xtemp[m];\n }\n }\n}\n\n// Method to perform gaussian elimination for a matrix stored in row major format\ntemplate\nvoid gaussianEliminationRowMajor(ValueTypeA **e, ValueTypeB *x, ValueTypeB *b, const IndexType bsize)\n{\n if (bsize > 4) { FatalError(\"Warning, Gaussian Elimination code doesn't pivot and bsize > 4\", AMGX_ERR_BAD_PARAMETERS); }\n\n ValueTypeB pivot, ratio, tmp;\n\n for (int j = 0; j < bsize; j++)\n {\n types::util::to_uptype(e[j][j], pivot);\n\n for (int k = j + 1; k < bsize; k++)\n {\n ValueTypeB temp;\n types::util::to_uptype(e[k][j], temp);\n ratio = temp / pivot;\n\n for (int m = j + 1; m < bsize; m++)\n {\n e[k][m] = e[k][m] - e[j][m] * ratio;\n }\n\n b[k] = b[k] - b[j] * ratio;\n }\n }\n\n // back substitution\n for (int j = bsize - 1; j >= 0; j--)\n {\n tmp = types::util::get_zero();\n\n for (int k = j + 1; k < bsize; k++)\n {\n tmp = tmp + e[j][k] * x[k];\n }\n\n x[j] = (b[j] - tmp) / e[j][j];\n }\n}\n\n// Method to perform gaussian elimination on matrix stored in block_dia_csr_matrix_h\ntemplate \nvoid GaussianElimination >::gaussianElimination_4x4_host(const Matrix_h &A, Vector_h &x, const Vector_h &b)\n{\n typedef typename Matrix_h::value_type ValueTypeA;\n typedef typename Vector_h::value_type ValueTypeB;\n\n if (A.get_num_rows() != 1)\n {\n FatalError(\"Haven't implemented gaussian elimination solver for num_blocks != 1\", AMGX_ERR_BAD_PARAMETERS);\n }\n else\n {\n int bsize = A.get_block_dimy();\n // Allocate space for block_matrix\n ValueTypeA **Atemp = new ValueTypeA* [bsize];\n\n for ( int i = 0; i < bsize; i++)\n {\n Atemp[i] = new ValueTypeA[bsize];\n }\n\n ValueTypeB *btemp = new ValueTypeB[bsize];\n ValueTypeB *xtemp = new ValueTypeB[bsize];\n\n // Copy matrix and rhs\n for (int m = 0; m < bsize; m++)\n {\n for (int n = 0; n < bsize; n++)\n {\n Atemp[m][n] = A.values[bsize * m + n]; //diag[bsize*m + n]; since we have 1 nonzero only\n }\n\n btemp[m] = b[m];\n }\n\n ValueTypeB *x_ptr = x.raw();\n gaussianEliminationRowMajor(Atemp, x_ptr, btemp, bsize);\n }\n}\n\n// Method to perform gaussian elimination on matrix stored in block_dia_csr_matrix_d\ntemplate \nvoid GaussianElimination >::gaussianElimination_4x4_device(const Matrix_d &A, Vector_d &x, const Vector_d &b)\n{\n typedef typename Matrix_d::value_type ValueTypeA;\n typedef typename Vector_d::value_type ValueTypeB;\n typedef typename Matrix_d::index_type IndexType;\n\n if (A.get_num_rows() != 1)\n {\n FatalError(\"Haven't implemented gaussian elimination solver for num_blocks != 1 or block_size != 4\", AMGX_ERR_BAD_PARAMETERS);\n }\n else\n {\n const IndexType *dia_ptr = A.diag.raw();\n const ValueTypeA *values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n ValueTypeB *x_ptr = x.raw();\n gaussianElimination4by4Kernel <<< 1, 1>>>(dia_ptr, values_ptr, x_ptr, b_ptr);\n cudaCheckError();\n }\n}\n\ntemplate \nvoid GaussianElimination::gaussianElimination(const Matrix &A, Vector &x, const Vector &b)\n{\n FatalError(\"Matrix format not supported in gaussian_elimination\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid GaussianElimination >::gaussianElimination(const Matrix_h &A, Vector_h &x, const Vector_h &b)\n{\n if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n gaussianElimination_4x4_host(A, x, b);\n }\n else\n {\n FatalError(\"gaussElimination: Blocksize is unsupported\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\ntemplate \nvoid GaussianElimination >::gaussianElimination(const Matrix_d &A, Vector_d &x, const Vector_d &b)\n{\n if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n gaussianElimination_4x4_device(A, x, b);\n }\n else\n {\n FatalError(\"gaussElimination: Blocksize is unsupported\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n// -------------------------------\n// Explict instantiations\n// -------------------------------\n#define AMGX_CASE_LINE(CASE) template class GaussianElimination::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/gaussian_elimination.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace amgx\n{\nnamespace testing_tools\n{\n\n\nstruct hash_path_determinism_checker_private\n{\n std::map checkpoint_counts;\n};\n\nhash_path_determinism_checker *hash_path_determinism_checker::singleton()\n{\n static hash_path_determinism_checker *m = 0;\n\n if (!m) { m = new hash_path_determinism_checker; }\n\n return m;\n}\n\nhash_path_determinism_checker::hash_path_determinism_checker()\n{\n priv = new hash_path_determinism_checker_private();\n}\nhash_path_determinism_checker::~hash_path_determinism_checker()\n{\n if (priv) { delete priv; }\n\n priv = 0;\n}\nstatic __host__ __device__ unsigned int hash_function(unsigned int a, unsigned int seed)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n\n__global__ void fast_hash_kernel(unsigned int *din, unsigned int *dout, long long int N, unsigned int seed)\n{\n long long int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n for (; tid < N; tid += blockDim.x * gridDim.x)\n {\n dout[tid] = hash_function(din[tid], seed);\n }\n}\n\n__global__ void fast_hash_kernel_no_permute(unsigned int *din, unsigned int *dout, long long int N, unsigned int seed)\n{\n long long int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n for (; tid < N; tid += blockDim.x * gridDim.x)\n {\n dout[tid] = hash_function(din[tid], seed) * hash_function((unsigned int)tid + (unsigned int)(tid >> 32), seed);\n }\n}\n\n\nvoid hash_path_determinism_checker::checkpoint(const std::string &name, void *data, long long int size_in_bytes, bool no_permute)\n{\n amgx::thrust::device_vector hash_buff(size_in_bytes / 4);\n\n if ( no_permute )\n {\n fast_hash_kernel_no_permute <<< 26, 256>>>((unsigned int *)data, amgx::thrust::raw_pointer_cast(hash_buff.data()), size_in_bytes / 4, 1987);\n cudaCheckError();\n }\n else\n {\n fast_hash_kernel <<< 26, 256>>>((unsigned int *)data, amgx::thrust::raw_pointer_cast(hash_buff.data()), size_in_bytes / 4, 1987);\n }\n\n cudaCheckError();\n unsigned long long int checksum = amgx::thrust::reduce(hash_buff.begin(), hash_buff.end(), 0ull, amgx::thrust::plus());\n cudaCheckError();\n int &checkpoint_count = priv->checkpoint_counts[name];\n printf(\"checksum %s\\t\\t%d\\t%llx\\n\", name.c_str(), checkpoint_count++, checksum);\n}\n\nunsigned long long int hash_path_determinism_checker::checksum( void *data, long long int size_in_bytes, bool no_permute )\n{\n amgx::thrust::device_vector hash_buff(size_in_bytes / 4);\n\n if ( no_permute )\n {\n fast_hash_kernel_no_permute <<< 26, 256>>>((unsigned int *)data, amgx::thrust::raw_pointer_cast(hash_buff.data()), size_in_bytes / 4, 1987);\n cudaCheckError();\n }\n else\n {\n fast_hash_kernel <<< 26, 256>>>((unsigned int *)data, amgx::thrust::raw_pointer_cast(hash_buff.data()), size_in_bytes / 4, 1987);\n }\n\n cudaCheckError();\n unsigned long long int checksum = amgx::thrust::reduce(hash_buff.begin(), hash_buff.end(), 0ull, amgx::thrust::plus());\n cudaCheckError();\n return checksum;\n}\n\n\n\n}\n}\n\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace amgx\n{\nnamespace testing_tools\n{\n\n\nstruct hash_path_determinism_checker_private\n{\n std::map checkpoint_counts;\n};\n\nhash_path_determinism_checker *hash_path_determinism_checker::singleton()\n{\n static hash_path_determinism_checker *m = 0;\n\n if (!m) { m = new hash_path_determinism_checker; }\n\n return m;\n}\n\nhash_path_determinism_checker::hash_path_determinism_checker()\n{\n priv = new hash_path_determinism_checker_private();\n}\nhash_path_determinism_checker::~hash_path_determinism_checker()\n{\n if (priv) { delete priv; }\n\n priv = 0;\n}\nstatic __host__ __device__ unsigned int hash_function(unsigned int a, unsigned int seed)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n\n__global__ void fast_hash_kernel(unsigned int *din, unsigned int *dout, long long int N, unsigned int seed)\n{\n long long int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n for (; tid < N; tid += blockDim.x * gridDim.x)\n {\n dout[tid] = hash_function(din[tid], seed);\n }\n}\n\n__global__ void fast_hash_kernel_no_permute(unsigned int *din, unsigned int *dout, long long int N, unsigned int seed)\n{\n long long int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n for (; tid < N; tid += blockDim.x * gridDim.x)\n {\n dout[tid] = hash_function(din[tid], seed) * hash_function((unsigned int)tid + (unsigned int)(tid >> 32), seed);\n }\n}\n\n\nvoid hash_path_determinism_checker::checkpoint(const std::string &name, void *data, long long int size_in_bytes, bool no_permute)\n{\n amgx::thrust::device_vector hash_buff(size_in_bytes / 4);\n\n if ( no_permute )\n {\n fast_hash_kernel_no_permute <<< 26, 256>>>((unsigned int *)data, amgx::thrust::raw_pointer_cast(hash_buff.data()), size_in_bytes / 4, 1987);\n cudaCheckError();\n }\n else\n {\n fast_hash_kernel <<< 26, 256>>>((unsigned int *)data, amgx::thrust::raw_pointer_cast(hash_buff.data()), size_in_bytes / 4, 1987);\n }\n\n cudaCheckError();\n unsigned long long int checksum = amgx::thrust::reduce(hash_buff.begin(), hash_buff.end(), 0ull, amgx::thrust::plus());\n cudaCheckError();\n int &checkpoint_count = priv->checkpoint_counts[name];\n printf(\"checksum %s\\t\\t%d\\t%llx\\n\", name.c_str(), checkpoint_count++, checksum);\n}\n\nunsigned long long int hash_path_determinism_checker::checksum( void *data, long long int size_in_bytes, bool no_permute )\n{\n amgx::thrust::device_vector hash_buff(size_in_bytes / 4);\n\n if ( no_permute )\n {\n fast_hash_kernel_no_permute <<< 26, 256>>>((unsigned int *)data, amgx::thrust::raw_pointer_cast(hash_buff.data()), size_in_bytes / 4, 1987);\n cudaCheckError();\n }\n else\n {\n fast_hash_kernel <<< 26, 256>>>((unsigned int *)data, amgx::thrust::raw_pointer_cast(hash_buff.data()), size_in_bytes / 4, 1987);\n }\n\n cudaCheckError();\n unsigned long long int checksum = amgx::thrust::reduce(hash_buff.begin(), hash_buff.end(), 0ull, amgx::thrust::plus());\n cudaCheckError();\n return checksum;\n}\n\n\n\n}\n}\n\n", "source": "/kaggle/working/amgx/src/determinism_checker.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\ntemplate\nstruct row_length\n{\n __host__ __device__ T operator()(const T &x) const\n {\n const T *next_ptr = &x;\n next_ptr++;\n return (*next_ptr) - x;\n }\n};\n\nnamespace amgx\n{\n\n__global__ void computeRowOffsetsKernel(INDEX_TYPE num_rows, INDEX_TYPE num_nz, const INDEX_TYPE *row_indices, INDEX_TYPE *row_offsets )\n{\n //one thread per non-zero\n int nz = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (nz == 0)\n {\n row_offsets[0] = 0;\n row_offsets[num_rows] = num_nz;\n }\n\n while (nz < num_nz - 1)\n {\n int row = row_indices[nz];\n int next_row = row_indices[nz + 1];\n\n while (row < next_row) //this loop should probably always execute once, but I'm making it a loop just in case...\n {\n row_offsets[++row] = nz + 1;\n }\n\n nz += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void computeRowIndicesKernel(INDEX_TYPE num_rows, const INDEX_TYPE *row_offsets, INDEX_TYPE *row_indices )\n{\n //one warp per row //possible optimziation: multiple warps per row\n int row = (blockIdx.x * blockDim.x + threadIdx.x) / AMGX_WARP_SIZE;\n int warp_id = threadIdx.x % AMGX_WARP_SIZE;\n\n while (row < num_rows)\n {\n int start = row_offsets[row] + warp_id;\n int end = row_offsets[row + 1];\n\n for (int nz = start; nz < end; nz += AMGX_WARP_SIZE)\n {\n row_indices[nz] = row;\n }\n\n row += blockDim.x * gridDim.x / AMGX_WARP_SIZE;\n }\n}\n\n__global__ void computeDiagonalKernelCOO(INDEX_TYPE num_nz, INDEX_TYPE *row_indices, INDEX_TYPE *col_indices, INDEX_TYPE *diag, INDEX_TYPE *diag_end_offsets)\n{\n //BLOCKY*BLOCKX threads per nz\n INDEX_TYPE nz = (blockIdx.x * blockDim.x + threadIdx.x);\n\n while (nz < num_nz)\n {\n INDEX_TYPE row = row_indices[nz];\n INDEX_TYPE col = col_indices[nz];\n\n if (row == col)\n {\n //copy block to diag\n diag[row] = nz;\n diag_end_offsets[row] = nz + 1;\n }\n\n nz += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void computeDiagonalKernelCSR(INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, INDEX_TYPE *diag, INDEX_TYPE null_index, INDEX_TYPE *diag_end_offsets)\n{\n INDEX_TYPE row = (blockIdx.x * blockDim.x + threadIdx.x);\n\n while (row < num_rows)\n {\n int nz = row_offsets[row];\n int last_nz = row_offsets[row + 1];\n\n while (nz < last_nz)\n {\n int col = col_indices[nz];\n\n if (row == col)\n {\n diag[row] = nz;\n diag_end_offsets[row] = nz + 1;\n break;\n }\n\n nz++;\n }\n\n row += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void computeColorOffsetsKernelCSR(INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, const INDEX_TYPE *row_colors, INDEX_TYPE *smaller_color_offsets, INDEX_TYPE *larger_color_offsets, INDEX_TYPE *diag)\n{\n INDEX_TYPE row = (blockIdx.x * blockDim.x + threadIdx.x);\n\n while (row < num_rows)\n {\n int my_color = row_colors[row];\n int nz = row_offsets[row];\n int last_nz = row_offsets[row + 1];\n int location_small = -1;\n int location_large = -1;\n\n while (nz < last_nz)\n {\n int col = col_indices[nz];\n\n if (row_colors[col] >= my_color && location_small == -1)\n {\n location_small = nz;\n }\n\n if (row_colors[col] > my_color && location_large == -1)\n {\n location_large = nz;\n break;\n }\n\n nz++;\n }\n\n if (location_large == -1) { location_large = last_nz + 1; }\n\n if (location_small == -1) { location_small = last_nz + 1; }\n\n larger_color_offsets[row] = location_large;\n smaller_color_offsets[row] = location_small;\n row += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void computeDiagonalKernelDiagProp (INDEX_TYPE num_rows, INDEX_TYPE num_nz, INDEX_TYPE *diag, INDEX_TYPE *diag_end_offsets)\n{\n INDEX_TYPE r = (blockIdx.x * blockDim.x + threadIdx.x);\n\n while (r < num_rows)\n {\n diag[r] = num_nz + r;\n diag_end_offsets[r] = num_nz + r + 1;\n r += blockDim.x * gridDim.x;\n }\n}\n\ntemplate \n__global__ void reorderElements (INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *permutation, T *data, T *temp, INDEX_TYPE max_row, INDEX_TYPE blockSize)\n{\n INDEX_TYPE rowId = blockIdx.x;\n\n while (rowId < num_rows)\n {\n INDEX_TYPE rowStart = row_offsets[rowId];\n INDEX_TYPE rowLen = row_offsets[rowId + 1] - rowStart;\n\n int i = threadIdx.x;\n\n //copy and reorder into temp storage\n while (i < rowLen * blockSize)\n {\n temp[max_row * blockIdx.x + i] = data[(permutation[rowStart + i / blockSize]) * blockSize + i % blockSize];\n i += blockDim.x;\n }\n\n __syncthreads();\n //copy back\n i = threadIdx.x;\n\n while (i < rowLen * blockSize)\n {\n data[rowStart * blockSize + i] = temp[max_row * blockIdx.x + i];\n i += blockDim.x;\n }\n\n rowId += gridDim.x;\n }\n}\n\ntemplate \nvoid\nMatrix< TemplateConfig >::apply(const Vector &v, Vector &res, ViewType view)\n{\n Vector &v_ = const_cast&>(v);\n multiply(*this, v_, res, view);\n}\n\ntemplate \nvoid\nMatrix< TemplateConfig >::apply(const Vector &v, Vector &res, ViewType view)\n{\n Vector &v_ = const_cast&>(v);\n multiply(*this, v_, res, view);\n}\n\n\n\ntemplate \nvoid\nMatrix< TemplateConfig >::print(char *f, char *s, int srows, int erows, int trank)\n{\n typedef typename TemplateConfig::MatPrec ValueType;\n int rank = 0;\n int level = 0;\n char filename[1024];\n FILE *fid = NULL;\n int printRowsStart, printRowsEnd;\n int i, j, ii, xdim, ydim, tnnz;\n ValueType a;\n#ifdef AMGX_WITH_MPI\n MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n#endif\n\n //check target rank\n if (rank == trank)\n {\n //check whether to ouput to stdout or file\n if (f == NULL)\n {\n fid = stdout;\n }\n else\n {\n level = this->amg_level_index;\n#ifdef _WIN32\n _snprintf_s(filename, 1024, 1024, \"%s_l%d_r%d.mtx\", f, level, rank);\n#else\n snprintf(filename, 1024, \"%s_l%d_r%d.mtx\", f, level, rank);\n#endif\n fid = fopen(filename, \"w\");\n }\n\n cudaDeviceSynchronize();\n cudaCheckError();\n printRowsStart = (srows < 0) ? 0 : srows;\n printRowsEnd = (erows < 0) ? this->get_num_rows() : erows;\n tnnz = this->get_num_nz() * this->get_block_size();\n\n //matrix might have separate diagonal (need to accoutn for it in nnz)\n if (this->hasProps(DIAG, this->props))\n {\n //matrix might be non-square so take std::min of # of rows and cols\n tnnz += std::min(this->get_num_rows(), this->get_num_cols()) * this->get_block_size();\n }\n\n auto trafI = [&](auto const &I, auto const &i) { return I * this->get_block_dimy() + i + 1; };\n auto trafJ = [&](auto const &J, auto const &j) { return J * this->get_block_dimx() + j + 1; };\n\n fprintf(fid, \"%%%%MatrixMarket matrix coordinate real general\\n\");\n fprintf(fid, \"%% %s\\n\", s);\n fprintf(fid, \"%d %d %d\\n\", this->get_num_rows() * this->get_block_dimx(), this->get_num_cols() * this->get_block_dimy(), tnnz);\n\n IVector_h this_row_offsets(this->row_offsets.size()); this_row_offsets = this->row_offsets;\n IVector_h this_col_indices(this->col_indices.size()); this_col_indices = this->col_indices;\n MVector_h this_values(this->values.size()); this_values = this->values;\n\n int const bsSquared = this->get_block_dimx() * this->get_block_dimy();\n\n for (i = printRowsStart; i < printRowsEnd; i++)\n {\n for (ydim = 0; ydim < this->get_block_dimy(); ydim++)\n {\n if (this->hasProps(DIAG, this->props))\n {\n if (i < std::min(this->get_num_rows(), this->get_num_cols()))\n {\n for (xdim = 0; xdim < this->get_block_dimx(); xdim++)\n {\n a = this_values[this->diag[i] * bsSquared + this->get_block_dimy() * ydim + xdim];\n fprintf(fid, \"%d %d \", trafI(i, ydim), trafI(i, xdim));\n types::util::fprintf(fid, \"%20.16f\", a);\n fprintf(fid, \"\\n\");\n }\n }\n }\n\n for (ii = this_row_offsets[i]; ii < this_row_offsets[i + 1]; ii++)\n {\n j = this_col_indices[ii];\n\n for (xdim = 0; xdim < this->get_block_dimx(); xdim++)\n {\n a = this_values[ii * bsSquared + this->get_block_dimy() * ydim + xdim];\n fprintf(fid, \"%d %d \", trafI(i, ydim), trafJ(j, xdim));\n types::util::fprintf(fid, \"%20.16f\", a);\n fprintf(fid, \"\\n\");\n }\n }\n }\n }\n\n cudaDeviceSynchronize();\n cudaCheckError();\n cudaGetLastError();\n\n if (fid != stdout)\n {\n fclose(fid);\n }\n }\n}\n\n\ntemplate \nvoid\nMatrix< TemplateConfig >::printToFile(char *f, char *s, int srows, int erows)\n{\n int rank = 0;\n#ifdef AMGX_WITH_MPI\n MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n#endif\n //notice that print will be called with different (target) rank on different on different ranks/processes\n this->print(f, s, srows, erows, rank);\n}\n\ntemplate \nvoid\nMatrix< TemplateConfig >::print(char *f, char *s, int srows, int erows, int trank)\n{\n typedef typename TemplateConfig::MatPrec ValueType;\n int rank = 0;\n int level = 0;\n char filename[1024];\n FILE *fid = NULL;\n int printRowsStart, printRowsEnd;\n int i, j, ii, xdim, ydim, tnnz;\n ValueType a;\n#ifdef AMGX_WITH_MPI\n MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n#endif\n\n //check target rank\n if (rank == trank)\n {\n //check whether to ouput to stdout or file\n if (f == NULL)\n {\n fid = stdout;\n }\n else\n {\n level = this->amg_level_index;\n#ifdef _WIN32\n _snprintf_s(filename, 1024, 1024, \"%s_r%d_l%d.mtx\", f, rank, level);\n#else\n snprintf(filename, 1024, \"%s_r%d_l%d.mtx\", f, rank, level);\n#endif\n fid = fopen(filename, \"w\");\n }\n\n cudaDeviceSynchronize();\n cudaCheckError();\n printRowsStart = (srows < 0) ? 0 : srows;\n printRowsEnd = (erows < 0) ? this->get_num_rows() : erows;\n tnnz = this->get_num_nz() * this->get_block_size();\n\n //matrix might have separate diagonal (need to accoutn for it in nnz)\n if (this->hasProps(DIAG, this->props))\n {\n //matrix might be non-square so take min of # of rows and cols\n tnnz += std::min(this->get_num_rows(), this->get_num_cols()) * this->get_block_size();\n }\n\n fprintf(fid, \"%%%%MatrixMarket matrix coordinate real general\\n\");\n fprintf(fid, \"%% %s\\n\", s);\n fprintf(fid, \"%d %d %d\\n\", this->get_num_rows() * this->get_block_dimx(), this->get_num_cols() * this->get_block_dimy(), tnnz);\n\n auto trafI = [&](auto const &I, auto const &i) { return I * this->get_block_dimy() + i + 1; };\n\n for (i = printRowsStart; i < printRowsEnd; i++)\n {\n for (ydim = 0; ydim < this->get_block_dimy(); ydim++)\n {\n if (this->hasProps(DIAG, this->props))\n {\n if (i < std::min(this->get_num_rows(), this->get_num_cols()))\n {\n for (xdim = 0; xdim < this->get_block_dimx(); xdim++)\n {\n a = this->values[this->diag[i] * this->get_block_dimx() * this->get_block_dimy() + this->get_block_dimy() * ydim + xdim];\n fprintf(fid, \"%d %d \", trafI(i, ydim), trafI(i, xdim));\n types::util::fprintf(fid, \"%20.16f\", a);\n fprintf(fid, \"\\n\");\n }\n }\n }\n\n for (ii = this->row_offsets[i]; ii < this->row_offsets[i + 1]; ii++)\n {\n j = this->col_indices[ii];\n\n for (xdim = 0; xdim < this->get_block_dimx(); xdim++)\n {\n a = this->values[ii * this->get_block_dimx() * this->get_block_dimy() + this->get_block_dimy() * ydim + xdim];\n fprintf(fid, \"%d %d \", i + 1, j + 1);\n types::util::fprintf(fid, \"%20.16f\", a);\n fprintf(fid, \"\\n\");\n }\n }\n }\n }\n\n cudaDeviceSynchronize();\n cudaCheckError();\n cudaGetLastError();\n\n if (fid != stdout)\n {\n fclose(fid);\n }\n }\n}\n\n\ntemplate \nvoid\nMatrix< TemplateConfig >::printToFile(char *f, char *s, int srows, int erows)\n{\n int rank = 0;\n#ifdef AMGX_WITH_MPI\n MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n#endif\n //notice that print will be called with different (target) rank on different on different ranks/processes\n this->print(f, s, srows, erows, rank);\n}\n\ntemplate \nvoid\nMatrix< TemplateConfig > ::convert( const Matrix &mat, unsigned int new_props, int block_dimy, int block_dimx )\n{\n if ( !mat.is_initialized() )\n {\n FatalError(\"Trying to convert from the uninitialized matrix\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n this->set_initialized(0);\n this->addProps(new_props);\n index_type block_size = block_dimx * block_dimy;\n index_type new_num_rows = ( mat.get_num_rows() * mat.block_dimy + block_dimy - 1 ) / block_dimy;\n index_type new_num_cols = ( mat.get_num_cols() * mat.block_dimx + block_dimx - 1 ) / block_dimx;\n MVector new_values;\n IVector new_col_indices;\n index_type new_num_nnz = 0;\n IVector new_row_indices;\n IVector new_row_offsets;\n MVector new_dia_values;\n new_dia_values.resize( new_num_rows * block_size, types::util::get_zero() );\n\n if ( mat.hasProps(COO) )\n {\n std::map< std::pair, index_type> ind;\n\n for ( index_type i = 0; i < mat.get_num_nz(); i++ )\n for ( index_type loc_row = 0; loc_row < mat.block_dimy; loc_row++ )\n for ( index_type loc_col = 0; loc_col < mat.block_dimx; loc_col++ )\n {\n index_type in_row = mat.row_indices[i] * mat.block_dimy + loc_row;\n index_type in_col = mat.col_indices[i] * mat.block_dimx + loc_col;\n value_type in_val = mat.values[i * mat.block_size + loc_row * mat.block_dimx + loc_col];\n\n if ( types::util::is_zero(in_val) ) { continue; } // skip zero entries\n\n index_type out_br = in_row / block_dimy;\n index_type out_bc = in_col / block_dimx;\n index_type out_lr = in_row % block_dimy;\n index_type out_lc = in_col % block_dimx;\n\n if ( ind.find( std::pair( out_br, out_bc ) ) == ind.end() )\n {\n // allocate a new block\n ind.insert( std::pair< std::pair, index_type>( std::pair( out_br, out_bc ), new_num_nnz ) );\n\n if (out_br != out_bc || !this->hasProps(DIAG))\n {\n new_row_indices.push_back( out_br );\n new_col_indices.push_back( out_bc );\n new_num_nnz++;\n\n for ( int b = 0; b < block_size; b++ )\n {\n new_values.push_back( types::util::get_zero() );\n }\n }\n }\n\n if ( out_br != out_bc || !this->hasProps(DIAG) )\n {\n new_values[ ind[std::pair( out_br, out_bc )] * block_size + out_lr * block_dimx + out_lc ] = in_val;\n }\n else\n {\n new_dia_values[ out_br * block_size + out_lr * block_dimx + out_lc ] = in_val;\n }\n }\n } // ( mat.hasProps(COO) )\n else if ( mat.hasProps(CSR))\n {\n new_num_nnz = 0;\n //MVector new_dia_values;\n //new_dia_values.resize( new_num_rows*block_size, 0.0 );\n new_row_offsets.resize( new_num_rows + 1 );\n\n // process each output block row\n for ( index_type i = 0; i < new_num_rows; i++ )\n {\n new_row_offsets[i] = new_num_nnz;\n // count non zero column blocks\n IVector non_zero_blocks( new_num_cols );\n\n for ( index_type j = i * block_dimy; j < (i + 1) * block_dimy && j < mat.get_num_rows() * mat.block_dimy; j++ )\n {\n // input row block / local position\n index_type in_br = j / mat.block_dimy;\n index_type in_lr = j % mat.block_dimy;\n\n // loop through block columns\n for ( index_type r = mat.row_offsets[in_br]; r < mat.row_offsets[in_br + 1] + ( mat.hasProps(DIAG) ); r++ )\n {\n index_type in_bc = ( r == mat.row_offsets[in_br + 1] ) ? in_br : mat.col_indices[r];\n\n // loop through local columns\n for ( index_type in_lc = 0; in_lc < mat.block_dimx; in_lc++ )\n {\n index_type in_col = in_bc * mat.block_dimx + in_lc;\n index_type out_bc = in_col / block_dimx;\n // fetch input entry value\n value_type val = ( r == mat.row_offsets[in_br + 1] ) ?\n mat.values[mat.diag[in_br] * mat.block_size + in_lr * mat.block_dimx + in_lc] :\n mat.values[r * mat.block_size + in_lr * mat.block_dimx + in_lc];\n\n if ( types::util::is_zero(val) ) { continue; } // skip zero entries\n\n // mark non_zero column\n non_zero_blocks[out_bc] = 1;\n }\n }\n }\n\n // populate non zero column blocks\n for ( int bc = 0; bc < new_num_cols; bc++ )\n if ( non_zero_blocks[bc] != 0 )\n {\n if ( i != bc || !this->hasProps(DIAG) ) // off-diagonal for DIAG\n {\n non_zero_blocks[bc] = new_num_nnz++;\n new_col_indices.push_back( bc );\n\n for ( int b = 0; b < block_size; b++ )\n {\n new_values.push_back( types::util::get_zero() );\n }\n }\n }\n\n // fill non zero values\n for ( index_type j = i * block_dimy; j < (i + 1) * block_dimy && j < mat.get_num_rows() * mat.block_dimy; j++ )\n {\n // output row block/local position\n index_type out_br = j / block_dimy;\n index_type out_lr = j % block_dimy;\n // input row block/local position\n index_type in_br = j / mat.block_dimy;\n index_type in_lr = j % mat.block_dimy;\n\n // loop through block columns\n for ( index_type r = mat.row_offsets[in_br]; r < mat.row_offsets[in_br + 1] + ( mat.hasProps(DIAG) ); r++ )\n {\n index_type in_bc = ( r == mat.row_offsets[in_br + 1] ) ? in_br : mat.col_indices[r];\n\n // loop through local columns\n for ( index_type in_lc = 0; in_lc < mat.block_dimx; in_lc++ )\n {\n index_type in_col = in_bc * mat.block_dimx + in_lc;\n index_type out_bc = in_col / block_dimx;\n index_type out_lc = in_col % block_dimx;\n // fetch input entry value\n value_type val = ( r == mat.row_offsets[in_br + 1] ) ?\n mat.values[mat.diag[in_br] * mat.block_size + in_lr * mat.block_dimx + in_lc] :\n mat.values[r * mat.block_size + in_lr * mat.block_dimx + in_lc];\n\n if ( types::util::is_zero(val) ) { continue; } // skip zero entries\n\n // write to new matrix\n if ( out_br != out_bc || !this->hasProps(DIAG) )\n {\n new_values[ non_zero_blocks[out_bc] * block_size + out_lr * block_dimx + out_lc] = val;\n }\n else\n {\n new_dia_values[ out_br * block_size + out_lr * block_dimx + out_lc ] = val;\n }\n }\n }\n }\n } // for( i < new_num_rows )\n\n // fill extra diagonal for the last block\n int extra_start = ( mat.get_num_rows() * mat.block_dimy ) % block_dimy;\n\n if ( extra_start > 0 )\n {\n for ( int r = extra_start; r < block_dimy; r++ )\n if ( this->hasProps(DIAG) )\n {\n new_dia_values[ (new_num_rows - 1) * block_size + r * block_dimx + r] = types::util::get_one();\n }\n else\n {\n new_values[ (new_num_nnz - 1) * block_size + r * block_dimx + r] = types::util::get_one();\n }\n }\n\n new_row_offsets[new_num_rows] = new_num_nnz;\n } // ( mat.hasProps(CSR) )\n else\n {\n FatalError(\"Input matrix for conversion doesn't have COO or CSR format\", AMGX_ERR_CONFIGURATION);\n }\n\n if ( this->hasProps(DIAG) )\n {\n new_values.insert(new_values.end(), new_dia_values.begin(), new_dia_values.end());\n }\n else\n for ( int b = 0; b < block_size; b++ )\n {\n new_values.push_back( types::util::get_zero() );\n }\n\n this->resize( new_num_rows, new_num_cols, new_num_nnz, block_dimy, block_dimx );\n this->values.copy( new_values );\n this->values.set_block_dimx(block_dimx);\n this->values.set_block_dimy(block_dimy);\n this->col_indices.copy( new_col_indices );\n\n if ( mat.hasProps(COO) )\n {\n this->row_indices.copy( new_row_indices );\n this->props |= COO;\n }\n\n if ( mat.hasProps(CSR) )\n {\n this->row_offsets.copy( new_row_offsets );\n this->props |= CSR;\n }\n\n if (this->hasProps(COO) && this->row_indices.size() == 0)\n {\n this->row_indices.resize(new_num_nnz);\n computeRowIndices();\n }\n\n if (this->hasProps(CSR) && this->row_offsets.size() == 0)\n {\n this->row_offsets.resize(new_num_rows + 1);\n computeRowOffsets();\n }\n\n computeDiagonal();\n this->set_initialized(1);\n}\n\ntemplate class MatrixBase;\n\ntemplate\nAMGX_ERROR\nMatrixBase::resize(index_type num_rows, index_type num_cols, index_type num_nz, int skipDiaCompute)\n{\n if (this->is_initialized())\n {\n FatalError(\"Debug throw: resizing already initialized matrix\\n\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n {\n this->num_rows = num_rows;\n this->num_cols = num_cols;\n this->num_nz = num_nz;\n\n if ( hasProps(DIAG) )\n {\n values.resize(num_nz * block_size + num_rows * block_size);\n }\n else\n {\n values.resize((num_nz + 1)*block_size);\n //thrust_wrapper::fill(values.begin() + num_nz*block_size, values.end(), static_cast(0.0));\n }\n\n diag.resize(num_rows);\n m_diag_end_offsets.resize(num_rows);\n col_indices.resize(num_nz);\n\n //if( props == NONE ) {props = CSR; props |= DIAG;}\n if ( hasProps(COO) ) { row_indices.resize(num_nz); }\n\n if ( hasProps(CSR) ) { row_offsets.resize(num_rows + 1); }\n\n m_seq_offsets.resize(num_rows + 1);\n thrust_wrapper::sequence(m_seq_offsets.begin(), m_seq_offsets.end());\n cudaCheckError();\n\n if (!skipDiaCompute )\n {\n computeDiagonal();\n }\n }\n\n return AMGX_OK;\n}\n\n\n\n\n\ntemplate\nvoid\nMatrixBase::setupMatrix(Solver *outer_solver, AMG_Config &cfg, bool reuse_matrix_structure)\n{\n // Check that matrix is initialized\n if (!this->is_initialized())\n {\n FatalError(\"Trying to setup from the uninitialized matrix\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n this->set_initialized(0);\n m_separation_interior = cfg.getParameter(\"separation_interior\", \"default\");\n m_separation_exterior = cfg.getParameter(\"separation_exterior\", \"default\");\n\n if (m_separation_interior > m_separation_exterior) { FatalError(\"Interior separation cannot be wider than the exterior separation\", AMGX_ERR_CONFIGURATION); }\n\n // If latency hiding is disabled, the interior is overwritten\n if(!isLatencyHidingEnabled(cfg))\n {\n m_separation_interior = m_separation_exterior;\n }\n\n bool is_coloring_needed = outer_solver->isColoringNeeded();\n\n if (!reuse_matrix_structure)\n {\n // Color the matrix since the structure has changed\n if ( is_coloring_needed )\n {\n // Get the scope of the solver that needs coloring\n std::string cfg_scope_for_coloring;\n outer_solver->getColoringScope(cfg_scope_for_coloring);\n this->colorMatrix(cfg, cfg_scope_for_coloring);\n }\n }\n\n // Set the matrix block format\n BlockFormat block_format = cfg.getParameter( \"block_format\", \"default\" );\n\n if ( this->getBlockFormat() != block_format )\n {\n this->setBlockFormat( block_format );\n }\n\n // Permute the values and the column indices (if necessary)\n bool reorder_cols_by_color = outer_solver->getReorderColsByColorDesired();\n bool insert_diagonal = outer_solver->getInsertDiagonalDesired();\n\n if (reorder_cols_by_color)\n {\n if (reuse_matrix_structure) // Only permute the values\n {\n this->permuteValues();\n }\n else // Permute the values and the columns\n {\n this->reorderColumnsByColor(insert_diagonal);\n this->permuteValues();\n }\n }\n\n this->set_initialized(1);\n m_is_matrix_setup = true;\n}\n\ntemplate\nbool MatrixBase::isLatencyHidingEnabled(AMG_Config &cfg)\n{\n const int min_rows_latency_hiding =\n cfg.getParameter(\"min_rows_latency_hiding\", \"default\");\n\n // Test all partitions to check if they all fall below the threshold\n if (!is_matrix_singleGPU() && min_rows_latency_hiding >= 0)\n {\n const auto& nrows_per_part = manager->getNumRowsPerPart();\n\n // Look at all partitions to check whether the number of rows falls\n // below the user defined minimum\n for(auto& nrpp : nrows_per_part)\n {\n // If any partitions still have a large enough set of rows, \n // continue latency hiding\n if(nrpp >= min_rows_latency_hiding)\n {\n return true;\n }\n }\n }\n\n return false;\n}\n\ntemplate\nvoid\nMatrixBase::reorderColumnsByColor(bool insert_diagonal)\n{\n // If columns already reordered, don't reorder again\n if (this->m_cols_reordered_by_color)\n {\n return;\n }\n\n // Check if matrix is colored\n if (!this->hasProps(COLORING))\n {\n FatalError(\"Matrix must be colored in order to be reordered by colors. Try setting coloring_level=1 in the configuration file\", AMGX_ERR_CONFIGURATION);\n }\n\n set_initialized(0);\n // Compute the row indices\n addProps(COO);\n set_allow_recompute_diag(false);\n index_type num_non_zeros = num_nz;\n\n if (hasProps(DIAG) && insert_diagonal) // Diagonal stored separetely\n {\n num_non_zeros += num_rows;\n }\n\n // Append the diagonal if stored separately\n // The new matrix will have inside diagonal\n if (hasProps(DIAG) && insert_diagonal)\n {\n amgx::thrust::counting_iterator first(0);\n amgx::thrust::counting_iterator last = first + num_rows;\n // Create new row_indices with appended diagonal\n IVector new_row_indices(num_non_zeros);\n IVector new_col_indices(num_non_zeros);\n amgx::thrust::copy(row_indices.begin(), row_indices.end(), new_row_indices.begin());\n amgx::thrust::copy(first, last, new_row_indices.begin() + num_nz);\n // Create new col_indices with appended diagonal\n amgx::thrust::copy(col_indices.begin(), col_indices.end(), new_col_indices.begin());\n amgx::thrust::copy(first, last, new_col_indices.begin() + num_nz);\n row_indices.swap(new_row_indices);\n col_indices.swap(new_col_indices);\n new_row_indices.clear();\n new_row_indices.shrink_to_fit();\n new_col_indices.clear();\n new_col_indices.shrink_to_fit();\n }\n\n // Compute the color of every column\n IVector element_colors(num_non_zeros);\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(this->getMatrixColoring().getRowColors().begin(), col_indices.begin()),\n amgx::thrust::make_permutation_iterator(this->getMatrixColoring().getRowColors().begin(), col_indices.end()),\n element_colors.begin());\n // Compute the permutation vector by sorting by rows and columns\n m_values_permutation_vector.resize(num_non_zeros);\n thrust_wrapper::sequence(m_values_permutation_vector.begin(), m_values_permutation_vector.end());\n cusp::detail::sort_by_row_and_column(row_indices, element_colors, m_values_permutation_vector);\n cudaCheckError();\n element_colors.clear();\n element_colors.shrink_to_fit();\n // Compute the new column indices sorted by color\n IVector new_column_indices(num_non_zeros);\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(col_indices.begin(), m_values_permutation_vector.begin()),\n amgx::thrust::make_permutation_iterator(col_indices.begin(), m_values_permutation_vector.end()),\n new_column_indices.begin());\n col_indices.swap(new_column_indices);\n new_column_indices.clear();\n new_column_indices.shrink_to_fit();\n\n if (hasProps(DIAG) && insert_diagonal)\n {\n // Change the number of nonzeros\n set_num_nz(num_non_zeros);\n values.resize( (num_non_zeros + 1)*this->get_block_size());\n this->m_is_permutation_inplace = false;\n }\n else\n {\n this->m_is_permutation_inplace = true;\n }\n\n if (hasProps(DIAG) && insert_diagonal)\n {\n delProps(DIAG);\n // Force recomputation of row offsets\n delProps(CSR);\n }\n\n // Compute row offsets if input matrix only had COO format or if diagonal was inserted\n addProps(CSR);\n // Recompute the diagonal\n set_allow_recompute_diag(true);\n computeDiagonal();\n // Compute the color offsets\n m_smaller_color_offsets.resize(this->get_num_rows());\n m_larger_color_offsets.resize(this->get_num_rows());\n computeColorOffsets();\n this->m_cols_reordered_by_color = true;\n set_initialized(1);\n}\n\ntemplate\nvoid\nMatrixBase::sortByRowAndColumn()\n{\n this->set_initialized(0);\n // Add row_indices array\n this->addProps(COO);\n this->set_allow_recompute_diag(false);\n\n if (this->get_block_dimx() != 1 || this->get_block_dimy() != 1)\n {\n FatalError(\"sortByRowAndColumn only works for scalar matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n size_t N = this->row_indices.size();\n IVector permutation(N);\n thrust_wrapper::sequence(permutation.begin(), permutation.end());\n cudaCheckError();\n // compute permutation and sort by (I,J)\n {\n IVector temp(this->col_indices);\n amgx::thrust::stable_sort_by_key(temp.begin(), temp.end(), permutation.begin());\n temp = this->row_indices;\n amgx::thrust::gather(permutation.begin(), permutation.end(), temp.begin(), this->row_indices.begin());\n amgx::thrust::stable_sort_by_key(this->row_indices.begin(), this->row_indices.end(), permutation.begin());\n temp = this->col_indices;\n amgx::thrust::gather(permutation.begin(), permutation.end(), temp.begin(), this->col_indices.begin());\n }\n cudaCheckError();\n // use permutation to reorder the values\n {\n MVector temp(this->values);\n amgx::thrust::gather(permutation.begin(), permutation.end(), temp.begin(), this->values.begin());\n }\n cudaCheckError();\n this->set_allow_recompute_diag(true);\n this->addProps(CSR);\n // remove row indices array\n this->delProps(COO);\n this->computeDiagonal();\n this->set_initialized(1);\n}\n\ntemplate \nvoid\nMatrix >::computeDiagonal()\n{\n ViewType oldView = this->currentView();\n\n if (this->m_initialized) { this->setView(ALL); }\n\n if (this->allow_recompute_diag)\n {\n index_type num_rows = this->get_num_rows();\n index_type num_nz = this->get_num_nz();\n IVector &row_offsets = this->row_offsets;\n IVector &row_indices = this->row_indices;\n IVector &col_indices = this->col_indices;\n IVector &diag = this->diag;\n\n if (this->diag.size() != this->get_num_rows()) { this->diag.resize(this->get_num_rows()); }\n\n if (this->m_diag_end_offsets.size() != this->get_num_rows()) { this->m_diag_end_offsets.resize(this->get_num_rows()); }\n\n if ( this->hasProps(DIAG) )\n {\n int first = num_nz;\n\n for (int r = 0; r < num_rows; r++)\n {\n diag[r] = first++;\n }\n }\n else\n {\n index_type null_index = this->get_num_nz();\n\n if ( this->hasProps(CSR) )\n {\n for (int r = 0; r < num_rows; r++)\n {\n int start = row_offsets[r];\n int end = row_offsets[r + 1];\n diag[r] = null_index;\n\n for (int j = start; j < end; j++)\n {\n if (col_indices[j] == r)\n {\n diag[r] = j;\n break;\n }\n }\n }\n }\n else if (this->hasProps(COO) )\n {\n for (int i = 0; i < num_rows; i++)\n {\n diag[i] = null_index;\n }\n\n for (int j = 0; j < num_nz; j++)\n {\n int r = row_indices[j];\n\n if (r == col_indices[j])\n {\n diag[r] = j;\n }\n }\n }\n }\n\n for (int r = 0; r < num_rows; r++)\n {\n this->m_diag_end_offsets[r] = diag[r] + 1;\n }\n }\n\n this->setView(oldView);\n}\n\ntemplate \nvoid Matrix >::computeDiagonal()\n{\n ViewType oldView;\n index_type null_index = 0; //this->get_num_nz();\n\n if (!this->allow_recompute_diag || !(this->get_num_rows() > 0))\n {\n return;\n }\n\n if (this->m_initialized)\n {\n oldView = this->currentView();\n this->setView(ALL);\n }\n\n if (this->diag.size() < this->get_num_rows())\n {\n this->diag.resize(this->get_num_rows());\n }\n\n if (this->m_diag_end_offsets.size() < this->get_num_rows())\n {\n this->m_diag_end_offsets.resize(this->get_num_rows());\n }\n\n if (this->hasProps(DIAG))\n {\n int num_blocks = std::min(4096, (this->get_num_rows() + 511) / 512);\n computeDiagonalKernelDiagProp <<< num_blocks, 512, 0, amgx::thrust::global_thread_handle::get_stream()>>>(this->get_num_rows(), this->get_num_nz(), this->diag.raw(), this->m_diag_end_offsets.raw());\n cudaCheckError();\n }\n else if (this->hasProps(COO))\n {\n int num_blocks = std::min(4096, (this->get_num_nz() + 511) / 512);\n computeDiagonalKernelCOO <<< num_blocks, 512>>>(this->get_num_nz(), this->row_indices.raw(), this->col_indices.raw(), this->diag.raw(), this->m_diag_end_offsets.raw());\n cudaCheckError();\n }\n else\n {\n int num_blocks = std::min(4096, (this->get_num_rows() + 511) / 512);\n computeDiagonalKernelCSR <<< num_blocks, 512>>>(this->get_num_rows(), this->row_offsets.raw(), this->col_indices.raw(), this->diag.raw(), null_index, this->m_diag_end_offsets.raw());\n }\n\n cudaCheckError();\n\n if (this->m_initialized)\n {\n this->setView(oldView);\n }\n\n#ifdef DEBUG\n\n if (this->diag_copy.size() == 0)\n {\n this->diag_copy = this->diag;\n }\n else\n {\n if ((this->diag.size() != this->diag_copy.size()) || (this->diag.size() == 1)) { return; }\n\n IVector_h h_diag = this->diag;\n IVector_h h_diag_copy = this->diag_copy;\n bool equal = true;\n\n for (unsigned int i = 0; i < this->diag.size(); ++i)\n {\n if (h_diag[i] != h_diag_copy[i])\n {\n equal = false;\n break;\n }\n }\n\n if (equal)\n {\n FatalError(\"ComputeDiagonal was called, but diagonal hasn't changed\", AMGX_ERR_UNKNOWN);\n }\n }\n\n#endif\n}\n\ntemplate \nvoid\nMatrix >::permuteValues()\n{\n if (this->m_cols_reordered_by_color && this->m_is_permutation_inplace)\n {\n reorderValuesInPlace();\n }\n else if (this->m_cols_reordered_by_color && !this->m_is_permutation_inplace )\n {\n // might use a lot of memory\n MVector temp_values;\n temp_values.resize(this->values.size());\n temp_values.set_block_dimx(this->values.get_block_dimx());\n temp_values.set_block_dimy(this->values.get_block_dimy());\n amgx::unpermuteVector(this->values, temp_values, this->m_values_permutation_vector, (this->get_num_nz()) * (this->get_block_size()));\n this->values.swap(temp_values);\n temp_values.clear();\n temp_values.shrink_to_fit();\n }\n else\n {\n FatalError(\"Invalid reordering level in permuteValues\", AMGX_ERR_CONFIGURATION);\n }\n}\n\n\nvoid computeRowOffsetsDevice(int num_blocks, INDEX_TYPE num_rows, INDEX_TYPE num_nz, const INDEX_TYPE *row_indices, INDEX_TYPE *row_offsets, INDEX_TYPE block_size )\n{\n computeRowOffsetsKernel <<< num_blocks, 512>>>(num_rows, num_nz, row_indices, row_offsets);\n cudaCheckError();\n}\n\nextern void computeRowIndicesDevice(int num_blocks, INDEX_TYPE num_rows, const INDEX_TYPE *row_offsets, INDEX_TYPE *row_indices, INDEX_TYPE block_size )\n{\n computeRowIndicesKernel <<< num_blocks, 512>>>(num_rows, row_offsets, row_indices);\n cudaCheckError();\n}\n\nvoid computeColorOffsetsDeviceCSR(int num_blocks, INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, const INDEX_TYPE *row_colors, INDEX_TYPE *smaller_color_offsets, INDEX_TYPE *larger_color_offsets, INDEX_TYPE block_size, INDEX_TYPE *diag )\n{\n computeColorOffsetsKernelCSR <<< num_blocks, 512>>>(num_rows, row_offsets, col_indices, row_colors, smaller_color_offsets, larger_color_offsets, diag);\n cudaCheckError();\n}\n\ntemplate \nvoid reorderElementsDeviceCSR(INDEX_TYPE num_rows,\n INDEX_TYPE *row_offsets,\n INDEX_TYPE *permutation,\n INDEX_TYPE *col_indices,\n T *values,\n INDEX_TYPE block_size)\n{\n INDEX_TYPE max_row_length = std::max(1, thrust_wrapper::transform_reduce(row_offsets, row_offsets + num_rows, row_length(), 0, amgx::thrust::maximum()));\n //TODO: optimise this in terms of storage\n INDEX_TYPE storage_space = 100 * 1024 * 1024 * sizeof(T) / sizeof(cuDoubleComplex); // because we allocate as for cuComplex\n INDEX_TYPE blocks = 1500 < storage_space / (max_row_length * block_size * sizeof(T)) ? 1500 : storage_space / (max_row_length * block_size * sizeof(T));\n blocks = blocks < num_rows ? max(1, blocks) : num_rows;\n INDEX_TYPE aligned_space = ((max_row_length * block_size * sizeof(T) / 128 + 1) * 128) / sizeof(T); //pad to 128 bytes\n Vector > tempstorage(blocks * aligned_space);\n reorderElements <<< blocks, 256>>>(num_rows, row_offsets, permutation, values, (T *)tempstorage.raw(), aligned_space, block_size);\n cudaCheckError();\n}\n\n// explicitly instantiate reorderElementsDeviceCSR, since we call it from header and it's not a part of some class\ntemplate void reorderElementsDeviceCSR(INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *permutation, INDEX_TYPE *col_indices, float *values, INDEX_TYPE block_size);\ntemplate void reorderElementsDeviceCSR(INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *permutation, INDEX_TYPE *col_indices, double *values, INDEX_TYPE block_size);\ntemplate void reorderElementsDeviceCSR(INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *permutation, INDEX_TYPE *col_indices, cuComplex *values, INDEX_TYPE block_size);\ntemplate void reorderElementsDeviceCSR(INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *permutation, INDEX_TYPE *col_indices, cuDoubleComplex *values, INDEX_TYPE block_size);\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class MatrixBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Matrix::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n\n\n\n\n}//end namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\ntemplate\nstruct row_length\n{\n __host__ __device__ T operator()(const T &x) const\n {\n const T *next_ptr = &x;\n next_ptr++;\n return (*next_ptr) - x;\n }\n};\n\nnamespace amgx\n{\n\n__global__ void computeRowOffsetsKernel(INDEX_TYPE num_rows, INDEX_TYPE num_nz, const INDEX_TYPE *row_indices, INDEX_TYPE *row_offsets )\n{\n //one thread per non-zero\n int nz = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (nz == 0)\n {\n row_offsets[0] = 0;\n row_offsets[num_rows] = num_nz;\n }\n\n while (nz < num_nz - 1)\n {\n int row = row_indices[nz];\n int next_row = row_indices[nz + 1];\n\n while (row < next_row) //this loop should probably always execute once, but I'm making it a loop just in case...\n {\n row_offsets[++row] = nz + 1;\n }\n\n nz += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void computeRowIndicesKernel(INDEX_TYPE num_rows, const INDEX_TYPE *row_offsets, INDEX_TYPE *row_indices )\n{\n //one warp per row //possible optimziation: multiple warps per row\n int row = (blockIdx.x * blockDim.x + threadIdx.x) / AMGX_WARP_SIZE;\n int warp_id = threadIdx.x % AMGX_WARP_SIZE;\n\n while (row < num_rows)\n {\n int start = row_offsets[row] + warp_id;\n int end = row_offsets[row + 1];\n\n for (int nz = start; nz < end; nz += AMGX_WARP_SIZE)\n {\n row_indices[nz] = row;\n }\n\n row += blockDim.x * gridDim.x / AMGX_WARP_SIZE;\n }\n}\n\n__global__ void computeDiagonalKernelCOO(INDEX_TYPE num_nz, INDEX_TYPE *row_indices, INDEX_TYPE *col_indices, INDEX_TYPE *diag, INDEX_TYPE *diag_end_offsets)\n{\n //BLOCKY*BLOCKX threads per nz\n INDEX_TYPE nz = (blockIdx.x * blockDim.x + threadIdx.x);\n\n while (nz < num_nz)\n {\n INDEX_TYPE row = row_indices[nz];\n INDEX_TYPE col = col_indices[nz];\n\n if (row == col)\n {\n //copy block to diag\n diag[row] = nz;\n diag_end_offsets[row] = nz + 1;\n }\n\n nz += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void computeDiagonalKernelCSR(INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, INDEX_TYPE *diag, INDEX_TYPE null_index, INDEX_TYPE *diag_end_offsets)\n{\n INDEX_TYPE row = (blockIdx.x * blockDim.x + threadIdx.x);\n\n while (row < num_rows)\n {\n int nz = row_offsets[row];\n int last_nz = row_offsets[row + 1];\n\n while (nz < last_nz)\n {\n int col = col_indices[nz];\n\n if (row == col)\n {\n diag[row] = nz;\n diag_end_offsets[row] = nz + 1;\n break;\n }\n\n nz++;\n }\n\n row += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void computeColorOffsetsKernelCSR(INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, const INDEX_TYPE *row_colors, INDEX_TYPE *smaller_color_offsets, INDEX_TYPE *larger_color_offsets, INDEX_TYPE *diag)\n{\n INDEX_TYPE row = (blockIdx.x * blockDim.x + threadIdx.x);\n\n while (row < num_rows)\n {\n int my_color = row_colors[row];\n int nz = row_offsets[row];\n int last_nz = row_offsets[row + 1];\n int location_small = -1;\n int location_large = -1;\n\n while (nz < last_nz)\n {\n int col = col_indices[nz];\n\n if (row_colors[col] >= my_color && location_small == -1)\n {\n location_small = nz;\n }\n\n if (row_colors[col] > my_color && location_large == -1)\n {\n location_large = nz;\n break;\n }\n\n nz++;\n }\n\n if (location_large == -1) { location_large = last_nz + 1; }\n\n if (location_small == -1) { location_small = last_nz + 1; }\n\n larger_color_offsets[row] = location_large;\n smaller_color_offsets[row] = location_small;\n row += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void computeDiagonalKernelDiagProp (INDEX_TYPE num_rows, INDEX_TYPE num_nz, INDEX_TYPE *diag, INDEX_TYPE *diag_end_offsets)\n{\n INDEX_TYPE r = (blockIdx.x * blockDim.x + threadIdx.x);\n\n while (r < num_rows)\n {\n diag[r] = num_nz + r;\n diag_end_offsets[r] = num_nz + r + 1;\n r += blockDim.x * gridDim.x;\n }\n}\n\ntemplate \n__global__ void reorderElements (INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *permutation, T *data, T *temp, INDEX_TYPE max_row, INDEX_TYPE blockSize)\n{\n INDEX_TYPE rowId = blockIdx.x;\n\n while (rowId < num_rows)\n {\n INDEX_TYPE rowStart = row_offsets[rowId];\n INDEX_TYPE rowLen = row_offsets[rowId + 1] - rowStart;\n\n int i = threadIdx.x;\n\n //copy and reorder into temp storage\n while (i < rowLen * blockSize)\n {\n temp[max_row * blockIdx.x + i] = data[(permutation[rowStart + i / blockSize]) * blockSize + i % blockSize];\n i += blockDim.x;\n }\n\n __syncthreads();\n //copy back\n i = threadIdx.x;\n\n while (i < rowLen * blockSize)\n {\n data[rowStart * blockSize + i] = temp[max_row * blockIdx.x + i];\n i += blockDim.x;\n }\n\n rowId += gridDim.x;\n }\n}\n\ntemplate \nvoid\nMatrix< TemplateConfig >::apply(const Vector &v, Vector &res, ViewType view)\n{\n Vector &v_ = const_cast&>(v);\n multiply(*this, v_, res, view);\n}\n\ntemplate \nvoid\nMatrix< TemplateConfig >::apply(const Vector &v, Vector &res, ViewType view)\n{\n Vector &v_ = const_cast&>(v);\n multiply(*this, v_, res, view);\n}\n\n\n\ntemplate \nvoid\nMatrix< TemplateConfig >::print(char *f, char *s, int srows, int erows, int trank)\n{\n typedef typename TemplateConfig::MatPrec ValueType;\n int rank = 0;\n int level = 0;\n char filename[1024];\n FILE *fid = NULL;\n int printRowsStart, printRowsEnd;\n int i, j, ii, xdim, ydim, tnnz;\n ValueType a;\n#ifdef AMGX_WITH_MPI\n MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n#endif\n\n //check target rank\n if (rank == trank)\n {\n //check whether to ouput to stdout or file\n if (f == NULL)\n {\n fid = stdout;\n }\n else\n {\n level = this->amg_level_index;\n#ifdef _WIN32\n _snprintf_s(filename, 1024, 1024, \"%s_l%d_r%d.mtx\", f, level, rank);\n#else\n snprintf(filename, 1024, \"%s_l%d_r%d.mtx\", f, level, rank);\n#endif\n fid = fopen(filename, \"w\");\n }\n\n hipDeviceSynchronize();\n cudaCheckError();\n printRowsStart = (srows < 0) ? 0 : srows;\n printRowsEnd = (erows < 0) ? this->get_num_rows() : erows;\n tnnz = this->get_num_nz() * this->get_block_size();\n\n //matrix might have separate diagonal (need to accoutn for it in nnz)\n if (this->hasProps(DIAG, this->props))\n {\n //matrix might be non-square so take std::min of # of rows and cols\n tnnz += std::min(this->get_num_rows(), this->get_num_cols()) * this->get_block_size();\n }\n\n auto trafI = [&](auto const &I, auto const &i) { return I * this->get_block_dimy() + i + 1; };\n auto trafJ = [&](auto const &J, auto const &j) { return J * this->get_block_dimx() + j + 1; };\n\n fprintf(fid, \"%%%%MatrixMarket matrix coordinate real general\\n\");\n fprintf(fid, \"%% %s\\n\", s);\n fprintf(fid, \"%d %d %d\\n\", this->get_num_rows() * this->get_block_dimx(), this->get_num_cols() * this->get_block_dimy(), tnnz);\n\n IVector_h this_row_offsets(this->row_offsets.size()); this_row_offsets = this->row_offsets;\n IVector_h this_col_indices(this->col_indices.size()); this_col_indices = this->col_indices;\n MVector_h this_values(this->values.size()); this_values = this->values;\n\n int const bsSquared = this->get_block_dimx() * this->get_block_dimy();\n\n for (i = printRowsStart; i < printRowsEnd; i++)\n {\n for (ydim = 0; ydim < this->get_block_dimy(); ydim++)\n {\n if (this->hasProps(DIAG, this->props))\n {\n if (i < std::min(this->get_num_rows(), this->get_num_cols()))\n {\n for (xdim = 0; xdim < this->get_block_dimx(); xdim++)\n {\n a = this_values[this->diag[i] * bsSquared + this->get_block_dimy() * ydim + xdim];\n fprintf(fid, \"%d %d \", trafI(i, ydim), trafI(i, xdim));\n types::util::fprintf(fid, \"%20.16f\", a);\n fprintf(fid, \"\\n\");\n }\n }\n }\n\n for (ii = this_row_offsets[i]; ii < this_row_offsets[i + 1]; ii++)\n {\n j = this_col_indices[ii];\n\n for (xdim = 0; xdim < this->get_block_dimx(); xdim++)\n {\n a = this_values[ii * bsSquared + this->get_block_dimy() * ydim + xdim];\n fprintf(fid, \"%d %d \", trafI(i, ydim), trafJ(j, xdim));\n types::util::fprintf(fid, \"%20.16f\", a);\n fprintf(fid, \"\\n\");\n }\n }\n }\n }\n\n hipDeviceSynchronize();\n cudaCheckError();\n hipGetLastError();\n\n if (fid != stdout)\n {\n fclose(fid);\n }\n }\n}\n\n\ntemplate \nvoid\nMatrix< TemplateConfig >::printToFile(char *f, char *s, int srows, int erows)\n{\n int rank = 0;\n#ifdef AMGX_WITH_MPI\n MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n#endif\n //notice that print will be called with different (target) rank on different on different ranks/processes\n this->print(f, s, srows, erows, rank);\n}\n\ntemplate \nvoid\nMatrix< TemplateConfig >::print(char *f, char *s, int srows, int erows, int trank)\n{\n typedef typename TemplateConfig::MatPrec ValueType;\n int rank = 0;\n int level = 0;\n char filename[1024];\n FILE *fid = NULL;\n int printRowsStart, printRowsEnd;\n int i, j, ii, xdim, ydim, tnnz;\n ValueType a;\n#ifdef AMGX_WITH_MPI\n MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n#endif\n\n //check target rank\n if (rank == trank)\n {\n //check whether to ouput to stdout or file\n if (f == NULL)\n {\n fid = stdout;\n }\n else\n {\n level = this->amg_level_index;\n#ifdef _WIN32\n _snprintf_s(filename, 1024, 1024, \"%s_r%d_l%d.mtx\", f, rank, level);\n#else\n snprintf(filename, 1024, \"%s_r%d_l%d.mtx\", f, rank, level);\n#endif\n fid = fopen(filename, \"w\");\n }\n\n hipDeviceSynchronize();\n cudaCheckError();\n printRowsStart = (srows < 0) ? 0 : srows;\n printRowsEnd = (erows < 0) ? this->get_num_rows() : erows;\n tnnz = this->get_num_nz() * this->get_block_size();\n\n //matrix might have separate diagonal (need to accoutn for it in nnz)\n if (this->hasProps(DIAG, this->props))\n {\n //matrix might be non-square so take min of # of rows and cols\n tnnz += std::min(this->get_num_rows(), this->get_num_cols()) * this->get_block_size();\n }\n\n fprintf(fid, \"%%%%MatrixMarket matrix coordinate real general\\n\");\n fprintf(fid, \"%% %s\\n\", s);\n fprintf(fid, \"%d %d %d\\n\", this->get_num_rows() * this->get_block_dimx(), this->get_num_cols() * this->get_block_dimy(), tnnz);\n\n auto trafI = [&](auto const &I, auto const &i) { return I * this->get_block_dimy() + i + 1; };\n\n for (i = printRowsStart; i < printRowsEnd; i++)\n {\n for (ydim = 0; ydim < this->get_block_dimy(); ydim++)\n {\n if (this->hasProps(DIAG, this->props))\n {\n if (i < std::min(this->get_num_rows(), this->get_num_cols()))\n {\n for (xdim = 0; xdim < this->get_block_dimx(); xdim++)\n {\n a = this->values[this->diag[i] * this->get_block_dimx() * this->get_block_dimy() + this->get_block_dimy() * ydim + xdim];\n fprintf(fid, \"%d %d \", trafI(i, ydim), trafI(i, xdim));\n types::util::fprintf(fid, \"%20.16f\", a);\n fprintf(fid, \"\\n\");\n }\n }\n }\n\n for (ii = this->row_offsets[i]; ii < this->row_offsets[i + 1]; ii++)\n {\n j = this->col_indices[ii];\n\n for (xdim = 0; xdim < this->get_block_dimx(); xdim++)\n {\n a = this->values[ii * this->get_block_dimx() * this->get_block_dimy() + this->get_block_dimy() * ydim + xdim];\n fprintf(fid, \"%d %d \", i + 1, j + 1);\n types::util::fprintf(fid, \"%20.16f\", a);\n fprintf(fid, \"\\n\");\n }\n }\n }\n }\n\n hipDeviceSynchronize();\n cudaCheckError();\n hipGetLastError();\n\n if (fid != stdout)\n {\n fclose(fid);\n }\n }\n}\n\n\ntemplate \nvoid\nMatrix< TemplateConfig >::printToFile(char *f, char *s, int srows, int erows)\n{\n int rank = 0;\n#ifdef AMGX_WITH_MPI\n MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n#endif\n //notice that print will be called with different (target) rank on different on different ranks/processes\n this->print(f, s, srows, erows, rank);\n}\n\ntemplate \nvoid\nMatrix< TemplateConfig > ::convert( const Matrix &mat, unsigned int new_props, int block_dimy, int block_dimx )\n{\n if ( !mat.is_initialized() )\n {\n FatalError(\"Trying to convert from the uninitialized matrix\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n this->set_initialized(0);\n this->addProps(new_props);\n index_type block_size = block_dimx * block_dimy;\n index_type new_num_rows = ( mat.get_num_rows() * mat.block_dimy + block_dimy - 1 ) / block_dimy;\n index_type new_num_cols = ( mat.get_num_cols() * mat.block_dimx + block_dimx - 1 ) / block_dimx;\n MVector new_values;\n IVector new_col_indices;\n index_type new_num_nnz = 0;\n IVector new_row_indices;\n IVector new_row_offsets;\n MVector new_dia_values;\n new_dia_values.resize( new_num_rows * block_size, types::util::get_zero() );\n\n if ( mat.hasProps(COO) )\n {\n std::map< std::pair, index_type> ind;\n\n for ( index_type i = 0; i < mat.get_num_nz(); i++ )\n for ( index_type loc_row = 0; loc_row < mat.block_dimy; loc_row++ )\n for ( index_type loc_col = 0; loc_col < mat.block_dimx; loc_col++ )\n {\n index_type in_row = mat.row_indices[i] * mat.block_dimy + loc_row;\n index_type in_col = mat.col_indices[i] * mat.block_dimx + loc_col;\n value_type in_val = mat.values[i * mat.block_size + loc_row * mat.block_dimx + loc_col];\n\n if ( types::util::is_zero(in_val) ) { continue; } // skip zero entries\n\n index_type out_br = in_row / block_dimy;\n index_type out_bc = in_col / block_dimx;\n index_type out_lr = in_row % block_dimy;\n index_type out_lc = in_col % block_dimx;\n\n if ( ind.find( std::pair( out_br, out_bc ) ) == ind.end() )\n {\n // allocate a new block\n ind.insert( std::pair< std::pair, index_type>( std::pair( out_br, out_bc ), new_num_nnz ) );\n\n if (out_br != out_bc || !this->hasProps(DIAG))\n {\n new_row_indices.push_back( out_br );\n new_col_indices.push_back( out_bc );\n new_num_nnz++;\n\n for ( int b = 0; b < block_size; b++ )\n {\n new_values.push_back( types::util::get_zero() );\n }\n }\n }\n\n if ( out_br != out_bc || !this->hasProps(DIAG) )\n {\n new_values[ ind[std::pair( out_br, out_bc )] * block_size + out_lr * block_dimx + out_lc ] = in_val;\n }\n else\n {\n new_dia_values[ out_br * block_size + out_lr * block_dimx + out_lc ] = in_val;\n }\n }\n } // ( mat.hasProps(COO) )\n else if ( mat.hasProps(CSR))\n {\n new_num_nnz = 0;\n //MVector new_dia_values;\n //new_dia_values.resize( new_num_rows*block_size, 0.0 );\n new_row_offsets.resize( new_num_rows + 1 );\n\n // process each output block row\n for ( index_type i = 0; i < new_num_rows; i++ )\n {\n new_row_offsets[i] = new_num_nnz;\n // count non zero column blocks\n IVector non_zero_blocks( new_num_cols );\n\n for ( index_type j = i * block_dimy; j < (i + 1) * block_dimy && j < mat.get_num_rows() * mat.block_dimy; j++ )\n {\n // input row block / local position\n index_type in_br = j / mat.block_dimy;\n index_type in_lr = j % mat.block_dimy;\n\n // loop through block columns\n for ( index_type r = mat.row_offsets[in_br]; r < mat.row_offsets[in_br + 1] + ( mat.hasProps(DIAG) ); r++ )\n {\n index_type in_bc = ( r == mat.row_offsets[in_br + 1] ) ? in_br : mat.col_indices[r];\n\n // loop through local columns\n for ( index_type in_lc = 0; in_lc < mat.block_dimx; in_lc++ )\n {\n index_type in_col = in_bc * mat.block_dimx + in_lc;\n index_type out_bc = in_col / block_dimx;\n // fetch input entry value\n value_type val = ( r == mat.row_offsets[in_br + 1] ) ?\n mat.values[mat.diag[in_br] * mat.block_size + in_lr * mat.block_dimx + in_lc] :\n mat.values[r * mat.block_size + in_lr * mat.block_dimx + in_lc];\n\n if ( types::util::is_zero(val) ) { continue; } // skip zero entries\n\n // mark non_zero column\n non_zero_blocks[out_bc] = 1;\n }\n }\n }\n\n // populate non zero column blocks\n for ( int bc = 0; bc < new_num_cols; bc++ )\n if ( non_zero_blocks[bc] != 0 )\n {\n if ( i != bc || !this->hasProps(DIAG) ) // off-diagonal for DIAG\n {\n non_zero_blocks[bc] = new_num_nnz++;\n new_col_indices.push_back( bc );\n\n for ( int b = 0; b < block_size; b++ )\n {\n new_values.push_back( types::util::get_zero() );\n }\n }\n }\n\n // fill non zero values\n for ( index_type j = i * block_dimy; j < (i + 1) * block_dimy && j < mat.get_num_rows() * mat.block_dimy; j++ )\n {\n // output row block/local position\n index_type out_br = j / block_dimy;\n index_type out_lr = j % block_dimy;\n // input row block/local position\n index_type in_br = j / mat.block_dimy;\n index_type in_lr = j % mat.block_dimy;\n\n // loop through block columns\n for ( index_type r = mat.row_offsets[in_br]; r < mat.row_offsets[in_br + 1] + ( mat.hasProps(DIAG) ); r++ )\n {\n index_type in_bc = ( r == mat.row_offsets[in_br + 1] ) ? in_br : mat.col_indices[r];\n\n // loop through local columns\n for ( index_type in_lc = 0; in_lc < mat.block_dimx; in_lc++ )\n {\n index_type in_col = in_bc * mat.block_dimx + in_lc;\n index_type out_bc = in_col / block_dimx;\n index_type out_lc = in_col % block_dimx;\n // fetch input entry value\n value_type val = ( r == mat.row_offsets[in_br + 1] ) ?\n mat.values[mat.diag[in_br] * mat.block_size + in_lr * mat.block_dimx + in_lc] :\n mat.values[r * mat.block_size + in_lr * mat.block_dimx + in_lc];\n\n if ( types::util::is_zero(val) ) { continue; } // skip zero entries\n\n // write to new matrix\n if ( out_br != out_bc || !this->hasProps(DIAG) )\n {\n new_values[ non_zero_blocks[out_bc] * block_size + out_lr * block_dimx + out_lc] = val;\n }\n else\n {\n new_dia_values[ out_br * block_size + out_lr * block_dimx + out_lc ] = val;\n }\n }\n }\n }\n } // for( i < new_num_rows )\n\n // fill extra diagonal for the last block\n int extra_start = ( mat.get_num_rows() * mat.block_dimy ) % block_dimy;\n\n if ( extra_start > 0 )\n {\n for ( int r = extra_start; r < block_dimy; r++ )\n if ( this->hasProps(DIAG) )\n {\n new_dia_values[ (new_num_rows - 1) * block_size + r * block_dimx + r] = types::util::get_one();\n }\n else\n {\n new_values[ (new_num_nnz - 1) * block_size + r * block_dimx + r] = types::util::get_one();\n }\n }\n\n new_row_offsets[new_num_rows] = new_num_nnz;\n } // ( mat.hasProps(CSR) )\n else\n {\n FatalError(\"Input matrix for conversion doesn't have COO or CSR format\", AMGX_ERR_CONFIGURATION);\n }\n\n if ( this->hasProps(DIAG) )\n {\n new_values.insert(new_values.end(), new_dia_values.begin(), new_dia_values.end());\n }\n else\n for ( int b = 0; b < block_size; b++ )\n {\n new_values.push_back( types::util::get_zero() );\n }\n\n this->resize( new_num_rows, new_num_cols, new_num_nnz, block_dimy, block_dimx );\n this->values.copy( new_values );\n this->values.set_block_dimx(block_dimx);\n this->values.set_block_dimy(block_dimy);\n this->col_indices.copy( new_col_indices );\n\n if ( mat.hasProps(COO) )\n {\n this->row_indices.copy( new_row_indices );\n this->props |= COO;\n }\n\n if ( mat.hasProps(CSR) )\n {\n this->row_offsets.copy( new_row_offsets );\n this->props |= CSR;\n }\n\n if (this->hasProps(COO) && this->row_indices.size() == 0)\n {\n this->row_indices.resize(new_num_nnz);\n computeRowIndices();\n }\n\n if (this->hasProps(CSR) && this->row_offsets.size() == 0)\n {\n this->row_offsets.resize(new_num_rows + 1);\n computeRowOffsets();\n }\n\n computeDiagonal();\n this->set_initialized(1);\n}\n\ntemplate class MatrixBase;\n\ntemplate\nAMGX_ERROR\nMatrixBase::resize(index_type num_rows, index_type num_cols, index_type num_nz, int skipDiaCompute)\n{\n if (this->is_initialized())\n {\n FatalError(\"Debug throw: resizing already initialized matrix\\n\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n {\n this->num_rows = num_rows;\n this->num_cols = num_cols;\n this->num_nz = num_nz;\n\n if ( hasProps(DIAG) )\n {\n values.resize(num_nz * block_size + num_rows * block_size);\n }\n else\n {\n values.resize((num_nz + 1)*block_size);\n //thrust_wrapper::fill(values.begin() + num_nz*block_size, values.end(), static_cast(0.0));\n }\n\n diag.resize(num_rows);\n m_diag_end_offsets.resize(num_rows);\n col_indices.resize(num_nz);\n\n //if( props == NONE ) {props = CSR; props |= DIAG;}\n if ( hasProps(COO) ) { row_indices.resize(num_nz); }\n\n if ( hasProps(CSR) ) { row_offsets.resize(num_rows + 1); }\n\n m_seq_offsets.resize(num_rows + 1);\n thrust_wrapper::sequence(m_seq_offsets.begin(), m_seq_offsets.end());\n cudaCheckError();\n\n if (!skipDiaCompute )\n {\n computeDiagonal();\n }\n }\n\n return AMGX_OK;\n}\n\n\n\n\n\ntemplate\nvoid\nMatrixBase::setupMatrix(Solver *outer_solver, AMG_Config &cfg, bool reuse_matrix_structure)\n{\n // Check that matrix is initialized\n if (!this->is_initialized())\n {\n FatalError(\"Trying to setup from the uninitialized matrix\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n this->set_initialized(0);\n m_separation_interior = cfg.getParameter(\"separation_interior\", \"default\");\n m_separation_exterior = cfg.getParameter(\"separation_exterior\", \"default\");\n\n if (m_separation_interior > m_separation_exterior) { FatalError(\"Interior separation cannot be wider than the exterior separation\", AMGX_ERR_CONFIGURATION); }\n\n // If latency hiding is disabled, the interior is overwritten\n if(!isLatencyHidingEnabled(cfg))\n {\n m_separation_interior = m_separation_exterior;\n }\n\n bool is_coloring_needed = outer_solver->isColoringNeeded();\n\n if (!reuse_matrix_structure)\n {\n // Color the matrix since the structure has changed\n if ( is_coloring_needed )\n {\n // Get the scope of the solver that needs coloring\n std::string cfg_scope_for_coloring;\n outer_solver->getColoringScope(cfg_scope_for_coloring);\n this->colorMatrix(cfg, cfg_scope_for_coloring);\n }\n }\n\n // Set the matrix block format\n BlockFormat block_format = cfg.getParameter( \"block_format\", \"default\" );\n\n if ( this->getBlockFormat() != block_format )\n {\n this->setBlockFormat( block_format );\n }\n\n // Permute the values and the column indices (if necessary)\n bool reorder_cols_by_color = outer_solver->getReorderColsByColorDesired();\n bool insert_diagonal = outer_solver->getInsertDiagonalDesired();\n\n if (reorder_cols_by_color)\n {\n if (reuse_matrix_structure) // Only permute the values\n {\n this->permuteValues();\n }\n else // Permute the values and the columns\n {\n this->reorderColumnsByColor(insert_diagonal);\n this->permuteValues();\n }\n }\n\n this->set_initialized(1);\n m_is_matrix_setup = true;\n}\n\ntemplate\nbool MatrixBase::isLatencyHidingEnabled(AMG_Config &cfg)\n{\n const int min_rows_latency_hiding =\n cfg.getParameter(\"min_rows_latency_hiding\", \"default\");\n\n // Test all partitions to check if they all fall below the threshold\n if (!is_matrix_singleGPU() && min_rows_latency_hiding >= 0)\n {\n const auto& nrows_per_part = manager->getNumRowsPerPart();\n\n // Look at all partitions to check whether the number of rows falls\n // below the user defined minimum\n for(auto& nrpp : nrows_per_part)\n {\n // If any partitions still have a large enough set of rows, \n // continue latency hiding\n if(nrpp >= min_rows_latency_hiding)\n {\n return true;\n }\n }\n }\n\n return false;\n}\n\ntemplate\nvoid\nMatrixBase::reorderColumnsByColor(bool insert_diagonal)\n{\n // If columns already reordered, don't reorder again\n if (this->m_cols_reordered_by_color)\n {\n return;\n }\n\n // Check if matrix is colored\n if (!this->hasProps(COLORING))\n {\n FatalError(\"Matrix must be colored in order to be reordered by colors. Try setting coloring_level=1 in the configuration file\", AMGX_ERR_CONFIGURATION);\n }\n\n set_initialized(0);\n // Compute the row indices\n addProps(COO);\n set_allow_recompute_diag(false);\n index_type num_non_zeros = num_nz;\n\n if (hasProps(DIAG) && insert_diagonal) // Diagonal stored separetely\n {\n num_non_zeros += num_rows;\n }\n\n // Append the diagonal if stored separately\n // The new matrix will have inside diagonal\n if (hasProps(DIAG) && insert_diagonal)\n {\n amgx::thrust::counting_iterator first(0);\n amgx::thrust::counting_iterator last = first + num_rows;\n // Create new row_indices with appended diagonal\n IVector new_row_indices(num_non_zeros);\n IVector new_col_indices(num_non_zeros);\n amgx::thrust::copy(row_indices.begin(), row_indices.end(), new_row_indices.begin());\n amgx::thrust::copy(first, last, new_row_indices.begin() + num_nz);\n // Create new col_indices with appended diagonal\n amgx::thrust::copy(col_indices.begin(), col_indices.end(), new_col_indices.begin());\n amgx::thrust::copy(first, last, new_col_indices.begin() + num_nz);\n row_indices.swap(new_row_indices);\n col_indices.swap(new_col_indices);\n new_row_indices.clear();\n new_row_indices.shrink_to_fit();\n new_col_indices.clear();\n new_col_indices.shrink_to_fit();\n }\n\n // Compute the color of every column\n IVector element_colors(num_non_zeros);\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(this->getMatrixColoring().getRowColors().begin(), col_indices.begin()),\n amgx::thrust::make_permutation_iterator(this->getMatrixColoring().getRowColors().begin(), col_indices.end()),\n element_colors.begin());\n // Compute the permutation vector by sorting by rows and columns\n m_values_permutation_vector.resize(num_non_zeros);\n thrust_wrapper::sequence(m_values_permutation_vector.begin(), m_values_permutation_vector.end());\n cusp::detail::sort_by_row_and_column(row_indices, element_colors, m_values_permutation_vector);\n cudaCheckError();\n element_colors.clear();\n element_colors.shrink_to_fit();\n // Compute the new column indices sorted by color\n IVector new_column_indices(num_non_zeros);\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(col_indices.begin(), m_values_permutation_vector.begin()),\n amgx::thrust::make_permutation_iterator(col_indices.begin(), m_values_permutation_vector.end()),\n new_column_indices.begin());\n col_indices.swap(new_column_indices);\n new_column_indices.clear();\n new_column_indices.shrink_to_fit();\n\n if (hasProps(DIAG) && insert_diagonal)\n {\n // Change the number of nonzeros\n set_num_nz(num_non_zeros);\n values.resize( (num_non_zeros + 1)*this->get_block_size());\n this->m_is_permutation_inplace = false;\n }\n else\n {\n this->m_is_permutation_inplace = true;\n }\n\n if (hasProps(DIAG) && insert_diagonal)\n {\n delProps(DIAG);\n // Force recomputation of row offsets\n delProps(CSR);\n }\n\n // Compute row offsets if input matrix only had COO format or if diagonal was inserted\n addProps(CSR);\n // Recompute the diagonal\n set_allow_recompute_diag(true);\n computeDiagonal();\n // Compute the color offsets\n m_smaller_color_offsets.resize(this->get_num_rows());\n m_larger_color_offsets.resize(this->get_num_rows());\n computeColorOffsets();\n this->m_cols_reordered_by_color = true;\n set_initialized(1);\n}\n\ntemplate\nvoid\nMatrixBase::sortByRowAndColumn()\n{\n this->set_initialized(0);\n // Add row_indices array\n this->addProps(COO);\n this->set_allow_recompute_diag(false);\n\n if (this->get_block_dimx() != 1 || this->get_block_dimy() != 1)\n {\n FatalError(\"sortByRowAndColumn only works for scalar matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n size_t N = this->row_indices.size();\n IVector permutation(N);\n thrust_wrapper::sequence(permutation.begin(), permutation.end());\n cudaCheckError();\n // compute permutation and sort by (I,J)\n {\n IVector temp(this->col_indices);\n amgx::thrust::stable_sort_by_key(temp.begin(), temp.end(), permutation.begin());\n temp = this->row_indices;\n amgx::thrust::gather(permutation.begin(), permutation.end(), temp.begin(), this->row_indices.begin());\n amgx::thrust::stable_sort_by_key(this->row_indices.begin(), this->row_indices.end(), permutation.begin());\n temp = this->col_indices;\n amgx::thrust::gather(permutation.begin(), permutation.end(), temp.begin(), this->col_indices.begin());\n }\n cudaCheckError();\n // use permutation to reorder the values\n {\n MVector temp(this->values);\n amgx::thrust::gather(permutation.begin(), permutation.end(), temp.begin(), this->values.begin());\n }\n cudaCheckError();\n this->set_allow_recompute_diag(true);\n this->addProps(CSR);\n // remove row indices array\n this->delProps(COO);\n this->computeDiagonal();\n this->set_initialized(1);\n}\n\ntemplate \nvoid\nMatrix >::computeDiagonal()\n{\n ViewType oldView = this->currentView();\n\n if (this->m_initialized) { this->setView(ALL); }\n\n if (this->allow_recompute_diag)\n {\n index_type num_rows = this->get_num_rows();\n index_type num_nz = this->get_num_nz();\n IVector &row_offsets = this->row_offsets;\n IVector &row_indices = this->row_indices;\n IVector &col_indices = this->col_indices;\n IVector &diag = this->diag;\n\n if (this->diag.size() != this->get_num_rows()) { this->diag.resize(this->get_num_rows()); }\n\n if (this->m_diag_end_offsets.size() != this->get_num_rows()) { this->m_diag_end_offsets.resize(this->get_num_rows()); }\n\n if ( this->hasProps(DIAG) )\n {\n int first = num_nz;\n\n for (int r = 0; r < num_rows; r++)\n {\n diag[r] = first++;\n }\n }\n else\n {\n index_type null_index = this->get_num_nz();\n\n if ( this->hasProps(CSR) )\n {\n for (int r = 0; r < num_rows; r++)\n {\n int start = row_offsets[r];\n int end = row_offsets[r + 1];\n diag[r] = null_index;\n\n for (int j = start; j < end; j++)\n {\n if (col_indices[j] == r)\n {\n diag[r] = j;\n break;\n }\n }\n }\n }\n else if (this->hasProps(COO) )\n {\n for (int i = 0; i < num_rows; i++)\n {\n diag[i] = null_index;\n }\n\n for (int j = 0; j < num_nz; j++)\n {\n int r = row_indices[j];\n\n if (r == col_indices[j])\n {\n diag[r] = j;\n }\n }\n }\n }\n\n for (int r = 0; r < num_rows; r++)\n {\n this->m_diag_end_offsets[r] = diag[r] + 1;\n }\n }\n\n this->setView(oldView);\n}\n\ntemplate \nvoid Matrix >::computeDiagonal()\n{\n ViewType oldView;\n index_type null_index = 0; //this->get_num_nz();\n\n if (!this->allow_recompute_diag || !(this->get_num_rows() > 0))\n {\n return;\n }\n\n if (this->m_initialized)\n {\n oldView = this->currentView();\n this->setView(ALL);\n }\n\n if (this->diag.size() < this->get_num_rows())\n {\n this->diag.resize(this->get_num_rows());\n }\n\n if (this->m_diag_end_offsets.size() < this->get_num_rows())\n {\n this->m_diag_end_offsets.resize(this->get_num_rows());\n }\n\n if (this->hasProps(DIAG))\n {\n int num_blocks = std::min(4096, (this->get_num_rows() + 511) / 512);\n computeDiagonalKernelDiagProp <<< num_blocks, 512, 0, amgx::thrust::global_thread_handle::get_stream()>>>(this->get_num_rows(), this->get_num_nz(), this->diag.raw(), this->m_diag_end_offsets.raw());\n cudaCheckError();\n }\n else if (this->hasProps(COO))\n {\n int num_blocks = std::min(4096, (this->get_num_nz() + 511) / 512);\n computeDiagonalKernelCOO <<< num_blocks, 512>>>(this->get_num_nz(), this->row_indices.raw(), this->col_indices.raw(), this->diag.raw(), this->m_diag_end_offsets.raw());\n cudaCheckError();\n }\n else\n {\n int num_blocks = std::min(4096, (this->get_num_rows() + 511) / 512);\n computeDiagonalKernelCSR <<< num_blocks, 512>>>(this->get_num_rows(), this->row_offsets.raw(), this->col_indices.raw(), this->diag.raw(), null_index, this->m_diag_end_offsets.raw());\n }\n\n cudaCheckError();\n\n if (this->m_initialized)\n {\n this->setView(oldView);\n }\n\n#ifdef DEBUG\n\n if (this->diag_copy.size() == 0)\n {\n this->diag_copy = this->diag;\n }\n else\n {\n if ((this->diag.size() != this->diag_copy.size()) || (this->diag.size() == 1)) { return; }\n\n IVector_h h_diag = this->diag;\n IVector_h h_diag_copy = this->diag_copy;\n bool equal = true;\n\n for (unsigned int i = 0; i < this->diag.size(); ++i)\n {\n if (h_diag[i] != h_diag_copy[i])\n {\n equal = false;\n break;\n }\n }\n\n if (equal)\n {\n FatalError(\"ComputeDiagonal was called, but diagonal hasn't changed\", AMGX_ERR_UNKNOWN);\n }\n }\n\n#endif\n}\n\ntemplate \nvoid\nMatrix >::permuteValues()\n{\n if (this->m_cols_reordered_by_color && this->m_is_permutation_inplace)\n {\n reorderValuesInPlace();\n }\n else if (this->m_cols_reordered_by_color && !this->m_is_permutation_inplace )\n {\n // might use a lot of memory\n MVector temp_values;\n temp_values.resize(this->values.size());\n temp_values.set_block_dimx(this->values.get_block_dimx());\n temp_values.set_block_dimy(this->values.get_block_dimy());\n amgx::unpermuteVector(this->values, temp_values, this->m_values_permutation_vector, (this->get_num_nz()) * (this->get_block_size()));\n this->values.swap(temp_values);\n temp_values.clear();\n temp_values.shrink_to_fit();\n }\n else\n {\n FatalError(\"Invalid reordering level in permuteValues\", AMGX_ERR_CONFIGURATION);\n }\n}\n\n\nvoid computeRowOffsetsDevice(int num_blocks, INDEX_TYPE num_rows, INDEX_TYPE num_nz, const INDEX_TYPE *row_indices, INDEX_TYPE *row_offsets, INDEX_TYPE block_size )\n{\n computeRowOffsetsKernel <<< num_blocks, 512>>>(num_rows, num_nz, row_indices, row_offsets);\n cudaCheckError();\n}\n\nextern void computeRowIndicesDevice(int num_blocks, INDEX_TYPE num_rows, const INDEX_TYPE *row_offsets, INDEX_TYPE *row_indices, INDEX_TYPE block_size )\n{\n computeRowIndicesKernel <<< num_blocks, 512>>>(num_rows, row_offsets, row_indices);\n cudaCheckError();\n}\n\nvoid computeColorOffsetsDeviceCSR(int num_blocks, INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, const INDEX_TYPE *row_colors, INDEX_TYPE *smaller_color_offsets, INDEX_TYPE *larger_color_offsets, INDEX_TYPE block_size, INDEX_TYPE *diag )\n{\n computeColorOffsetsKernelCSR <<< num_blocks, 512>>>(num_rows, row_offsets, col_indices, row_colors, smaller_color_offsets, larger_color_offsets, diag);\n cudaCheckError();\n}\n\ntemplate \nvoid reorderElementsDeviceCSR(INDEX_TYPE num_rows,\n INDEX_TYPE *row_offsets,\n INDEX_TYPE *permutation,\n INDEX_TYPE *col_indices,\n T *values,\n INDEX_TYPE block_size)\n{\n INDEX_TYPE max_row_length = std::max(1, thrust_wrapper::transform_reduce(row_offsets, row_offsets + num_rows, row_length(), 0, amgx::thrust::maximum()));\n //TODO: optimise this in terms of storage\n INDEX_TYPE storage_space = 100 * 1024 * 1024 * sizeof(T) / sizeof(hipDoubleComplex); // because we allocate as for cuComplex\n INDEX_TYPE blocks = 1500 < storage_space / (max_row_length * block_size * sizeof(T)) ? 1500 : storage_space / (max_row_length * block_size * sizeof(T));\n blocks = blocks < num_rows ? max(1, blocks) : num_rows;\n INDEX_TYPE aligned_space = ((max_row_length * block_size * sizeof(T) / 128 + 1) * 128) / sizeof(T); //pad to 128 bytes\n Vector > tempstorage(blocks * aligned_space);\n reorderElements <<< blocks, 256>>>(num_rows, row_offsets, permutation, values, (T *)tempstorage.raw(), aligned_space, block_size);\n cudaCheckError();\n}\n\n// explicitly instantiate reorderElementsDeviceCSR, since we call it from header and it's not a part of some class\ntemplate void reorderElementsDeviceCSR(INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *permutation, INDEX_TYPE *col_indices, float *values, INDEX_TYPE block_size);\ntemplate void reorderElementsDeviceCSR(INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *permutation, INDEX_TYPE *col_indices, double *values, INDEX_TYPE block_size);\ntemplate void reorderElementsDeviceCSR(INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *permutation, INDEX_TYPE *col_indices, hipComplex *values, INDEX_TYPE block_size);\ntemplate void reorderElementsDeviceCSR(INDEX_TYPE num_rows, INDEX_TYPE *row_offsets, INDEX_TYPE *permutation, INDEX_TYPE *col_indices, hipDoubleComplex *values, INDEX_TYPE block_size);\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class MatrixBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Matrix::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n\n\n\n\n}//end namespace amgx\n", "source": "/kaggle/working/amgx/src/matrix.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nnamespace csr_multiply_detail\n{\n\n#include \n\n#include \n#include // Included inside the namespace to solve name colisions.\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n__global__\nvoid flag_halo_rows(int *row_ids, int size, int *flagArray, int neighbor, int global_id)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n int row_id = row_ids[tidx];\n flagArray[row_id] = tidx;\n }\n}\n\n__device__ __forceinline__ int get_work( int *queue, int warp_id )\n{\n int offset = -1;\n\n if ( utils::lane_id() == 0 )\n {\n offset = atomicAdd( queue, 1 );\n }\n\n return utils::shfl( offset, 0 );\n}\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool COUNT_ONLY >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncount_non_zeroes_kernel( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const int *B_rows,\n const int *B_cols,\n int *C_rows,\n int *C_cols,\n int *Aq1,\n int *Bq1,\n int *Aq2,\n int *Bq2,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n int c_row_id = a_row_id;\n\n if (Aq1 != NULL)\n {\n a_row_id = Aq1[a_row_id];\n }\n\n // Make sure we have to proceed.\n if ( COUNT_ONLY )\n {\n volatile int *status = reinterpret_cast( wk_status );\n\n if ( set.has_failed() || *status != 0 )\n {\n return;\n }\n }\n\n // Clear the set.\n set.clear();\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = utils::Ld::load( &A_rows[a_row_id + lane_id] );\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int b_row_id = -1;\n\n if ( is_active )\n {\n b_row_id = utils::Ld::load( &A_cols[a_col_it] );\n\n //b_row_id is actually column of A\n if (Aq2 != NULL)\n {\n b_row_id = Aq2[b_row_id];\n\t \t\t\n }\n\n if (Bq1 != NULL)\n {\n b_row_id = Bq1[b_row_id];\n }\n }\n\n // The number of valid rows.\n const int num_rows = __popc( utils::ballot(is_active) );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n const int uniform_b_row_id = utils::shfl( b_row_id, k );\n // Load the range of the row of B.\n int b_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n b_col_tmp = utils::Ld::load( &B_rows[uniform_b_row_id + lane_id] );\n }\n\n int b_col_it = utils::shfl( b_col_tmp, 0 );\n int b_col_end = utils::shfl( b_col_tmp, 1 );\n\n // Iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any(b_col_it < b_col_end) ; b_col_it += WARP_SIZE )\n {\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = utils::Ld::load( &B_cols[b_col_it] );\n\n // b_col_id is actually column of B\n if (Bq2 != NULL)\n {\n b_col_id = Bq2[b_col_id];\n }\n }\n\n set.insert( b_col_id, COUNT_ONLY ? wk_status : NULL );\n }\n }\n }\n\n // Store the results.\n if ( COUNT_ONLY )\n {\n int count = set.compute_size();\n\n if ( lane_id == 0 )\n {\n C_rows[c_row_id] = count;\n }\n }\n else\n {\n int c_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n c_col_tmp = utils::Ld::load( &C_rows[c_row_id + lane_id] );\n }\n\n int c_col_it = utils::shfl( c_col_tmp, 0 );\n int c_col_end = utils::shfl( c_col_tmp, 1 );\n // Store the results.\n int count = c_col_end - c_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n set.store( count, &C_cols[c_col_it] );\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool COUNT_ONLY >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncount_non_zeroes_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict B_rows,\n const int *__restrict B_cols,\n int *__restrict C_rows,\n int *__restrict C_cols,\n int *Aq1,\n int *Bq1,\n int *Aq2,\n int *Bq2,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // The hash keys stored in shared memory.\n __shared__ /*volatile*/ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n int c_row_id = a_row_id;\n\n if (Aq1 != NULL)\n {\n a_row_id = Aq1[a_row_id];\n }\n\n // Make sure we have to proceed.\n if ( COUNT_ONLY )\n {\n volatile int *status = reinterpret_cast( wk_status );\n\n if ( set.has_failed() || *status != 0 )\n {\n return;\n }\n }\n\n // Clear the set.\n set.clear();\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = utils::Ld::load( &A_rows[a_row_id + lane_id] );\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int b_row_id = -1;\n\n if ( is_active )\n {\n b_row_id = utils::Ld::load( &A_cols[a_col_it] );\n\n //b_row_id is actually column of A\n if (Aq2 != NULL)\n {\n b_row_id = Aq2[b_row_id];\n }\n\n if (Bq1 != NULL)\n {\n b_row_id = Bq1[b_row_id];\n }\n }\n\n const int num_rows = __popc( utils::ballot(is_active) );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n const int uniform_b_row_id = utils::shfl( b_row_id, local_k );\n // Load the range of the row of B.\n int b_col_tmp = -1;\n\n if ( is_active_k && lane_id_mod_num_threads < 2 )\n {\n b_col_tmp = utils::Ld::load( &B_rows[uniform_b_row_id + lane_id_mod_num_threads] );\n }\n\n int b_col_it = utils::shfl( b_col_tmp, lane_id_div_num_threads * NUM_THREADS_PER_ROW + 0 );\n int b_col_end = utils::shfl( b_col_tmp, lane_id_div_num_threads * NUM_THREADS_PER_ROW + 1 );\n\n // Iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = utils::Ld::load( &B_cols[b_col_it] );\n\n // b_col_id is actually column of B\n if (Bq2 != NULL)\n {\n b_col_id = Bq2[b_col_id];\n }\n }\n\n set.insert( b_col_id, COUNT_ONLY ? wk_status : NULL );\n }\n }\n }\n\n // Store the results.\n if ( COUNT_ONLY )\n {\n int count = set.compute_size();\n\n if ( lane_id == 0 )\n {\n C_rows[c_row_id] = count;\n }\n }\n else\n {\n int c_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n c_col_tmp = utils::Ld::load( &C_rows[c_row_id + lane_id] );\n }\n\n int c_col_it = utils::shfl( c_col_tmp, 0 );\n int c_col_end = utils::shfl( c_col_tmp, 1 );\n // Store the results.\n int count = c_col_end - c_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n set.store( count, &C_cols[c_col_it] );\n }\n }\n}\n\ntemplate \n__device__ __forceinline__ void sparse_add_process_row(int row_id, const int *__restrict__ row_offsets, const int *__restrict__ col_indices, int lane_id, Hash_set &set, int *wk_status, int global_id, bool print_flag)\n{\n // Load the range of the row of RAP_int\n int col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n col_tmp = utils::Ld::load( &row_offsets[row_id + lane_id] );\n }\n\n int col_it = utils::shfl( col_tmp, 0 );\n int col_end = utils::shfl( col_tmp, 1 );\n\n // Iterate over the columns of RAP_int\n for ( col_it += lane_id ; utils::any(col_it < col_end) ; col_it += WARP_SIZE )\n {\n int col_id = -1;\n\n if ( col_it < col_end )\n {\n col_id = utils::Ld::load( &col_indices[col_it] );\n }\n\n set.insert( col_id, COUNT_ONLY ? wk_status : NULL );\n }\n}\n\n\ntemplate \n__device__ __forceinline__ void sparse_add_process_row_values(int row_id, const int *__restrict__ row_offsets, const int *__restrict__ col_indices, const Value_type *__restrict vals, int lane_id, Hash_map &map, int *wk_status)\n{\n // Load the range of the row.\n int col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n col_tmp = utils::Ld::load( &row_offsets[row_id + lane_id] );\n }\n\n int col_it = utils::shfl( col_tmp, 0 );\n int col_end = utils::shfl( col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( col_it += lane_id ; utils::any(col_it < col_end) ; col_it += WARP_SIZE )\n {\n const bool is_active = col_it < col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int col_id(-1);\n Value_type value = amgx::types::util::get_zero();\n\n if ( is_active )\n {\n col_id = utils::Ld::load( &col_indices[col_it] );\n value = utils::Ld::load( &vals[col_it] );\n }\n\n map.insert( col_id, value, wk_status );\n }\n}\n\n\n\n\ntemplate\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncount_non_zeroes_RAP_ext_kernel( const int RAP_int_num_rows,\n const int *__restrict RAP_int_rows,\n const int *__restrict RAP_int_cols,\n int **RAP_ext_row_ptrs,\n int **RAP_ext_col_ptrs,\n int *__restrict RAP_rows,\n int *__restrict RAP_cols,\n int **flagArray_ptrs,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status,\n int num_neighbors,\n int global_id )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // The hash keys stored in shared memory.\n __shared__ /*volatile*/ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int rap_int_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[rap_int_row_id * gmem_size], gmem_size );\n\n // Loop over rows of RAP_ext\n for ( ; rap_int_row_id < RAP_int_num_rows; rap_int_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Make sure we have to proceed.\n if ( COUNT_ONLY )\n {\n volatile int *status = reinterpret_cast( wk_status );\n\n if ( set.has_failed() || *status != 0 )\n {\n return;\n }\n }\n\n // Clear the set.\n set.clear();\n // ---------------------------------\n // First process RAP_int\n // ---------------------------------\n bool print_flag = false;\n\n if (rap_int_row_id == 3 && global_id == 0)\n {\n print_flag = true;\n }\n\n sparse_add_process_row(rap_int_row_id, RAP_int_rows, RAP_int_cols, lane_id, set, wk_status, global_id, print_flag);\n\n // ---------------------------------\n // Then process RAP_ext\n // ---------------------------------\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int flag = flagArray_ptrs[i][rap_int_row_id];\n\n if (flag != -1)\n {\n int *RAP_ext_rows = RAP_ext_row_ptrs[i];\n int *RAP_ext_cols = RAP_ext_col_ptrs[i];\n int pos_in_row_ext = flag;\n sparse_add_process_row(pos_in_row_ext, RAP_ext_rows, RAP_ext_cols, lane_id, set, wk_status, global_id, print_flag);\n }\n }\n\n // Store the results.\n if ( COUNT_ONLY )\n {\n int count = set.compute_size();\n\n if ( lane_id == 0 )\n {\n RAP_rows[rap_int_row_id] = count;\n }\n }\n else\n {\n int rap_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n rap_col_tmp = utils::Ld::load( &RAP_rows[rap_int_row_id + lane_id] );\n }\n\n int rap_col_it = utils::shfl( rap_col_tmp, 0 );\n int rap_col_end = utils::shfl( rap_col_tmp, 1 );\n // Store the results.\n int count = rap_col_end - rap_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n set.store( count, &RAP_cols[rap_col_it] );\n }\n }\n}\n\n\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Without_external_diag\n{\n static __device__ __forceinline__ bool is_active ( int a_col_it, int a_col_end ) { return a_col_it < a_col_end; }\n static __device__ __forceinline__ bool is_boundary( int a_col_it, int a_col_end ) { return false; }\n};\n\n// ====================================================================================================================\n\nstruct With_external_diag\n{\n static __device__ __forceinline__ bool is_active ( int a_col_it, int a_col_end ) { return a_col_it <= a_col_end; }\n static __device__ __forceinline__ bool is_boundary( int a_col_it, int a_col_end ) { return a_col_it == a_col_end; }\n};\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool COUNT_ONLY, typename Diag_traits >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncount_non_zeroes_ilu1_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict A_coloring,\n int *__restrict C_rows,\n int *__restrict C_cols,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Tables to broadcast values.\n __shared__ volatile int s_b_rows[CTA_SIZE], s_b_colors[CTA_SIZE];\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Make sure we have to proceed.\n if ( COUNT_ONLY )\n {\n volatile int *status = reinterpret_cast( wk_status );\n\n if ( set.has_failed() || *status != 0 )\n {\n return;\n }\n }\n\n // Clear the set.\n set.clear();\n // The color of the row.\n int a_row_color = A_coloring[a_row_id];\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = utils::Ld::load( &A_rows[a_row_id + lane_id] );\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any(Diag_traits::is_active(a_col_it, a_col_end)) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int b_row_id = -1;\n\n if ( is_active )\n {\n b_row_id = utils::Ld::load( &A_cols[a_col_it] );\n }\n\n if ( Diag_traits::is_boundary(a_col_it, a_col_end) )\n {\n b_row_id = a_row_id;\n }\n\n // Push the columns in the set.\n set.insert( b_row_id, COUNT_ONLY ? wk_status : NULL );\n\n // Skip computation if the color of the row is 0.\n if ( a_row_color != 0 )\n {\n // Gather the colors of the columns.\n int b_row_color = -1;\n\n if ( is_active )\n {\n b_row_color = A_coloring[b_row_id];\n }\n\n // The number of valid rows.\n int pred = is_active && b_row_color < a_row_color;\n int vote = utils::ballot( pred );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( pred )\n {\n s_b_rows [warp_id * WARP_SIZE + dest] = b_row_id;\n s_b_colors[warp_id * WARP_SIZE + dest] = b_row_color; // TODO: store an int2 rather than 2 ints of SM35 with 64bit banks.\n }\n\n const int num_rows = __popc( vote );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n const int uniform_b_row_id = s_b_rows [warp_id * WARP_SIZE + k];\n const int uniform_b_color = s_b_colors[warp_id * WARP_SIZE + k];\n // Load the range of the row of B.\n int b_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n b_col_tmp = utils::Ld::load( &A_rows[uniform_b_row_id + lane_id] );\n }\n\n int b_col_it = utils::shfl( b_col_tmp, 0 );\n int b_col_end = utils::shfl( b_col_tmp, 1 );\n\n // Iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any(b_col_it < b_col_end) ; b_col_it += WARP_SIZE )\n {\n int b_col_id = -1, b_col_color = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = utils::Ld::load( &A_cols[b_col_it] );\n b_col_color = utils::Ld::load( &A_coloring[b_col_id] );\n }\n\n int item = -1;\n\n if ( b_col_color >= uniform_b_color && b_col_color != a_row_color )\n {\n item = b_col_id;\n }\n\n set.insert( item, COUNT_ONLY ? wk_status : NULL );\n }\n }\n }\n }\n\n // Store the results.\n if ( COUNT_ONLY )\n {\n int count = set.compute_size();\n\n if ( lane_id == 0 )\n {\n C_rows[a_row_id] = count;\n }\n }\n else\n {\n int c_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n c_col_tmp = utils::Ld::load( &C_rows[a_row_id + lane_id] );\n }\n\n int c_col_it = utils::shfl( c_col_tmp, 0 );\n int c_col_end = utils::shfl( c_col_tmp, 1 );\n // Store the results.\n int count = c_col_end - c_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n set.store( count, &C_cols[c_col_it] );\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool COUNT_ONLY, typename Diag_traits >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncount_non_zeroes_ilu1_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict A_coloring,\n int *__restrict C_rows,\n int *__restrict C_cols,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // Tables to broadcast values.\n __shared__ volatile int s_b_rows[CTA_SIZE], s_b_colors[CTA_SIZE];\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Make sure we have to proceed.\n if ( COUNT_ONLY )\n {\n volatile int *status = reinterpret_cast( wk_status );\n\n if ( set.has_failed() || *status != 0 )\n {\n return;\n }\n }\n\n // Clear the set.\n set.clear();\n // The color of the row.\n int a_row_color = A_coloring[a_row_id];\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = utils::Ld::load( &A_rows[a_row_id + lane_id] );\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any(Diag_traits::is_active(a_col_it, a_col_end)) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int b_row_id = -1;\n\n if ( is_active )\n {\n b_row_id = utils::Ld::load( &A_cols[a_col_it] );\n }\n\n if ( Diag_traits::is_boundary(a_col_it, a_col_end) )\n {\n b_row_id = a_row_id;\n }\n\n // Push the columns in the set.\n set.insert( b_row_id, COUNT_ONLY ? wk_status : NULL );\n\n // Skip computation if the color of the row is 0.\n if ( a_row_color != 0 )\n {\n // Gather the colors of the columns.\n int b_row_color = -1;\n\n if ( is_active )\n {\n b_row_color = A_coloring[b_row_id];\n }\n\n // The number of valid rows.\n int pred = is_active && b_row_color < a_row_color;\n int vote = utils::ballot( pred );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( pred )\n {\n s_b_rows [warp_id * WARP_SIZE + dest] = b_row_id;\n s_b_colors[warp_id * WARP_SIZE + dest] = b_row_color;\n }\n\n const int num_rows = __popc( vote );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int uniform_b_row_id = -1, uniform_b_color = -1;\n\n if ( is_active_k )\n {\n uniform_b_row_id = s_b_rows [warp_id * WARP_SIZE + local_k];\n uniform_b_color = s_b_colors[warp_id * WARP_SIZE + local_k];\n }\n\n // Load the range of the row of B.\n int b_col_tmp = -1;\n\n if ( is_active_k && lane_id_mod_num_threads < 2 )\n {\n b_col_tmp = utils::Ld::load( &A_rows[uniform_b_row_id + lane_id_mod_num_threads] );\n }\n\n int b_col_it = utils::shfl( b_col_tmp, lane_id_div_num_threads * NUM_THREADS_PER_ROW + 0 );\n int b_col_end = utils::shfl( b_col_tmp, lane_id_div_num_threads * NUM_THREADS_PER_ROW + 1 );\n\n // Iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n int b_col_id = -1, b_col_color = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = utils::Ld::load( &A_cols[b_col_it] );\n b_col_color = utils::Ld::load( &A_coloring[b_col_id] );\n }\n\n int item = -1;\n\n if ( b_col_color >= uniform_b_color && b_col_color != a_row_color )\n {\n item = b_col_id;\n }\n\n set.insert( item, COUNT_ONLY ? wk_status : NULL );\n }\n }\n }\n }\n\n // Store the results.\n if ( COUNT_ONLY )\n {\n int count = set.compute_size();\n\n if ( lane_id == 0 )\n {\n C_rows[a_row_id] = count;\n }\n }\n else\n {\n int c_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n c_col_tmp = utils::Ld::load( &C_rows[a_row_id + lane_id] );\n }\n\n int c_col_it = utils::shfl( c_col_tmp, 0 );\n int c_col_end = utils::shfl( c_col_tmp, 1 );\n // Store the results.\n int count = c_col_end - c_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n set.store( count, &C_cols[c_col_it] );\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE, 6 )\nvoid\ncompute_values_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const int *__restrict B_rows,\n const int *__restrict B_cols,\n const Value_type *__restrict B_vals,\n const int *__restrict C_rows,\n int *__restrict C_cols,\n Value_type *__restrict C_vals,\n int *Aq1,\n int *Bq1,\n int *Aq2,\n int *Bq2,\n const int gmem_size,\n int *g_keys,\n Value_type *g_vals,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // The hash keys stored in shared memory.\n __shared__ /*volatile*/ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The hash values stored in shared memory.\n __shared__ Value_type s_vals[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_map map(&s_keys[warp_id * SMEM_SIZE],\n &g_keys[a_row_id * gmem_size],\n &s_vals[warp_id * SMEM_SIZE],\n &g_vals[a_row_id * gmem_size],\n gmem_size );\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n int c_row_id = a_row_id;\n\n if (Aq1 != NULL)\n {\n a_row_id = Aq1[a_row_id];\n }\n\n // Clear the map.\n map.clear();\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = utils::Ld::load( &A_rows[a_row_id + lane_id] );\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int b_row_id = -1;\n Value_type a_value = amgx::types::util::get_zero();\n\n if ( is_active )\n {\n b_row_id = utils::Ld::load( &A_cols[a_col_it] );\n a_value = utils::Ld::load( &A_vals[a_col_it] );\n\n //b_row_id is actually column of A\n if (Aq2 != NULL)\n {\n b_row_id = Aq2[b_row_id];\n }\n\n if (Bq1 != NULL)\n {\n b_row_id = Bq1[b_row_id];\n }\n }\n\n const int num_rows = __popc( utils::ballot(is_active) );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n const int uniform_b_row_id = utils::shfl( b_row_id, k );\n // The value of A.\n const Value_type uniform_a_value = utils::shfl( a_value, k );\n // Load the range of the row of B.\n int b_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n b_col_tmp = utils::Ld::load( &B_rows[uniform_b_row_id + lane_id] );\n }\n\n int b_col_it = utils::shfl( b_col_tmp, 0 );\n int b_col_end = utils::shfl( b_col_tmp, 1 );\n\n // Iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any(b_col_it < b_col_end) ; b_col_it += WARP_SIZE )\n {\n int b_col_id = -1;\n Value_type b_value = amgx::types::util::get_zero();\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = utils::Ld::load( &B_cols[b_col_it] );\n b_value = utils::Ld::load( &B_vals[b_col_it] );\n\n if (Bq2 != NULL)\n {\n b_col_id = Bq2[b_col_id];\n }\n }\n\n map.insert( b_col_id, uniform_a_value * b_value, wk_status );\n }\n }\n }\n\n // Store the results.\n int c_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n c_col_tmp = utils::Ld::load( &C_rows[c_row_id + lane_id] );\n }\n\n int c_col_it = utils::shfl( c_col_tmp, 0 );\n int c_col_end = utils::shfl( c_col_tmp, 1 );\n // Store the results.\n int count = c_col_end - c_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n map.store( count, &C_cols[c_col_it], &C_vals[c_col_it] );\n }\n\n}\n\ntemplate \n__global__ \nvoid count_non_zeroes_kernel_opt( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict B_rows,\n const int *__restrict B_cols,\n int *__restrict C_row_counts)\n{\n // Defines for hash container\n constexpr int SLOT_VACANT = -1;\n\n // Group indices\n constexpr int ngroups = CTA_SIZE / GROUP_SIZE;\n const int group_id = threadIdx.x / GROUP_SIZE;\n const int lane_id = threadIdx.x % GROUP_SIZE;\n\n // One row of A per group\n const int a_row_id = blockIdx.x * ngroups + group_id;\n\n // Block-level hash container storage\n __shared__ int key_s[ngroups*HASH_SIZE];\n __shared__ int counts[ngroups];\n\n // Initialise the keys and values.\n#pragma unroll\n for(int i = threadIdx.x; i < ngroups*HASH_SIZE; i += CTA_SIZE)\n {\n key_s[i] = SLOT_VACANT; // Inserted keys will be in range [0,N]\n }\n\n for(int i = 0; i < ngroups; ++i)\n {\n counts[i] = 0;\n }\n\n //__syncwarp();\n __syncthreads();\n\n int* key_group_s = &key_s[group_id*HASH_SIZE];\n\n if(a_row_id < A_num_rows)\n {\n // Distributed columns of row of A over threads in group.\n for (int a_col_it = A_rows[a_row_id] + lane_id; a_col_it < A_rows[a_row_id + 1]; a_col_it += GROUP_SIZE)\n {\n int a_col_id = A_cols[a_col_it];\n\n // Outer product of element of A and row of B.\n for (int b_col_it = B_rows[a_col_id]; b_col_it < B_rows[a_col_id + 1] ; ++b_col_it)\n {\n int b_col_id = B_cols[b_col_it];\n\n int hash = b_col_id % HASH_SIZE;\n\n // By construction this algorithm should guarantee \n // all keys can be inserted\n while(true)\n {\n // If the slot is vacant, then attempt acquire\n int key = key_group_s[hash];\n\n // Check if the key was already b_col_id, or was already set \n // Insert the product\n if(key == b_col_id)\n {\n break;\n }\n\n if(key == SLOT_VACANT)\n {\n int new_key = atomicCAS(&key_group_s[hash], SLOT_VACANT, b_col_id);\n\n if(new_key == SLOT_VACANT || new_key == b_col_id)\n {\n break;\n }\n }\n\n // We did not secure a slot, so linear probe to next slot\n hash = (hash + 1) % HASH_SIZE;\n }\n }\n }\n }\n\n //__syncwarp();\n __syncthreads();\n\n // Store the results.\n if(a_row_id < A_num_rows)\n {\n#pragma unroll\n for(int i = lane_id; i < HASH_SIZE; i += GROUP_SIZE)\n {\n if(key_group_s[i] != SLOT_VACANT)\n {\n atomicAdd(&counts[group_id], 1);\n }\n }\n }\n\n __syncthreads();\n\n if(a_row_id < A_num_rows && lane_id == 0)\n {\n C_row_counts[a_row_id] = counts[group_id];\n }\n}\n\ntemplate \n__global__ \nvoid compute_values_kernel_opt( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const ValueType *__restrict A_vals,\n const int *__restrict B_rows,\n const int *__restrict B_cols,\n const ValueType *__restrict B_vals,\n const int *__restrict C_rows,\n int *__restrict C_cols,\n ValueType *__restrict C_vals)\n{\n // Defines for hash container\n constexpr int SLOT_VACANT = -1;\n\n // Group indices\n constexpr int ngroups = CTA_SIZE / GROUP_SIZE;\n const int group_id = threadIdx.x / GROUP_SIZE;\n const int lane_id = threadIdx.x % GROUP_SIZE;\n\n // One row of A per group\n const int a_row_id = blockIdx.x * ngroups + group_id;\n\n // Dynamic sized shared memory\n extern __shared__ int s[];\n\n // Block-level hash container storage\n int* key_s = s;\n ValueType* data_s = (ValueType*)&key_s[ngroups*HASH_SIZE];\n int* col_ind_s = (int*)&data_s[ngroups*HASH_SIZE];\n\n // Group-level hash containers\n int* key_group_s = &key_s[group_id*HASH_SIZE];\n ValueType* data_group_s = &data_s[group_id*HASH_SIZE];\n\n // Initialise the keys and values.\n#pragma unroll\n for(int i = threadIdx.x; i < ngroups*HASH_SIZE; i += CTA_SIZE)\n {\n key_s[i] = SLOT_VACANT; // Inserted keys will be in range [0,N]\n data_s[i] = amgx::types::util::get_zero(); // We will sum into values\n }\n\n if(lane_id == 0)\n {\n col_ind_s[group_id] = 0;\n }\n\n //__syncwarp();\n __syncthreads();\n\n if(a_row_id < A_num_rows)\n {\n // Distributed columns of row of A over threads in group.\n for (int a_col_it = A_rows[a_row_id] + lane_id; a_col_it < A_rows[a_row_id + 1]; a_col_it += GROUP_SIZE)\n {\n int a_col_id = A_cols[a_col_it];\n\n // Outer product of element of A and row of B.\n for (int b_col_it = B_rows[a_col_id]; b_col_it < B_rows[a_col_id + 1] ; ++b_col_it)\n {\n ValueType val = A_vals[a_col_it]*B_vals[b_col_it];\n\n // Short circuit if zero\n if(amgx::types::util::is_zero(val)) continue;\n\n int b_col_id = B_cols[b_col_it];\n\n // XXX UPDATE HASHING APPROACH\n int hash = b_col_id % HASH_SIZE;\n\n // By construction this algorithm should guarantee \n // all keys can be inserted\n while(true)\n {\n // If the slot is vacant, then attempt acquire\n int key = key_group_s[hash];\n if(key == SLOT_VACANT)\n {\n int new_key = atomicCAS(&key_group_s[hash], SLOT_VACANT, b_col_id);\n if(new_key == SLOT_VACANT || new_key == b_col_id)\n {\n key = b_col_id;\n }\n }\n\n // Check if the key was already b_col_id, or was already set \n // Insert the product\n if(key == b_col_id)\n {\n utils::atomic_add(&data_group_s[hash], val);\n break;\n }\n\n // We did not secure a slot, so linear probe to next slot\n hash = (hash + 1) % HASH_SIZE;\n }\n }\n }\n }\n\n //__syncwarp();\n __syncthreads();\n\n // Store the results.\n int c_row_id = a_row_id;\n int c_col_it = C_rows[c_row_id];\n\n if(a_row_id < A_num_rows)\n {\n#pragma unroll\n for(int i = lane_id; i < HASH_SIZE; i += GROUP_SIZE)\n {\n if(key_group_s[i] != SLOT_VACANT)\n {\n // Short circuit if zero\n if(amgx::types::util::is_zero(data_group_s[i])) continue;\n\n int ind = atomicAdd(&col_ind_s[group_id], 1);\n\n C_cols[c_col_it + ind] = key_group_s[i];\n C_vals[c_col_it + ind] = data_group_s[i];\n }\n }\n }\n}\n\ntemplate< int NUM_THREADS_PER_ROW, typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE, 6 )\nvoid\ncompute_values_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const int *__restrict B_rows,\n const int *__restrict B_cols,\n const Value_type *__restrict B_vals,\n const int *__restrict C_rows,\n int *__restrict C_cols,\n Value_type *__restrict C_vals,\n int *Aq1,\n int *Bq1,\n int *Aq2,\n int *Bq2,\n const int gmem_size,\n int *g_keys,\n Value_type *g_vals,\n int *wk_work_queue,\n int *wk_status )\n{\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // The hash keys stored in shared memory.\n __shared__ /*volatile*/ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The hash values stored in shared memory.\n __shared__ Value_type s_vals[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_map map(&s_keys[warp_id * SMEM_SIZE],\n &g_keys[a_row_id * gmem_size],\n &s_vals[warp_id * SMEM_SIZE],\n &g_vals[a_row_id * gmem_size],\n gmem_size );\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n int c_row_id = a_row_id;\n\n if (Aq1 != NULL)\n {\n a_row_id = Aq1[a_row_id];\n }\n\n // Clear the map.\n map.clear();\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = utils::Ld::load( &A_rows[a_row_id + lane_id] );\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int b_row_id(-1);\n Value_type a_value = amgx::types::util::get_zero();\n\n if ( is_active )\n {\n b_row_id = utils::Ld::load( &A_cols[a_col_it] );\n a_value = utils::Ld::load( &A_vals[a_col_it] );\n\n //b_row_id is actually column of A\n if (Aq2 != NULL)\n {\n b_row_id = Aq2[b_row_id];\n }\n\n if (Bq1 != NULL)\n {\n b_row_id = Bq1[b_row_id];\n }\n }\n\n const int num_rows = __popc( utils::ballot(is_active) );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n const int uniform_b_row_id = utils::shfl( b_row_id, k + lane_id_div_num_threads );\n // The value of A.\n const Value_type uniform_a_value = utils::shfl( a_value, k + lane_id_div_num_threads );\n // Load the range of the row of B.\n int b_col_tmp = -1;\n\n if ( is_active_k && lane_id_mod_num_threads < 2 )\n {\n b_col_tmp = utils::Ld::load( &B_rows[uniform_b_row_id + lane_id_mod_num_threads] );\n }\n\n int b_col_it = utils::shfl( b_col_tmp, lane_id_div_num_threads * NUM_THREADS_PER_ROW + 0 );\n int b_col_end = utils::shfl( b_col_tmp, lane_id_div_num_threads * NUM_THREADS_PER_ROW + 1 );\n\n // Iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n int b_col_id(-1);\n Value_type b_value = amgx::types::util::get_zero();\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = utils::Ld::load( &B_cols[b_col_it] );\n b_value = utils::Ld::load( &B_vals[b_col_it] );\n\n //b_col_id is actually column of B\n if (Bq2 != NULL)\n {\n b_col_id = Bq2[b_col_id];\n }\n }\n\n map.insert( b_col_id, uniform_a_value * b_value, wk_status );\n }\n }\n }\n\n // Store the results.\n int c_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n c_col_tmp = utils::Ld::load( &C_rows[c_row_id + lane_id] );\n }\n\n int c_col_it = utils::shfl( c_col_tmp, 0 );\n int c_col_end = utils::shfl( c_col_tmp, 1 );\n // Store the results.\n int count = c_col_end - c_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n map.store( count, &C_cols[c_col_it], &C_vals[c_col_it] );\n }\n\n}\n\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE, 6 )\nvoid\ncompute_values_RAP_ext_kernel( const int RAP_int_num_rows,\n const int *__restrict RAP_int_rows,\n const int *__restrict RAP_int_cols,\n const Value_type *__restrict RAP_int_vals,\n int **RAP_ext_row_ptrs,\n int **RAP_ext_col_ptrs,\n Value_type **RAP_ext_val_ptrs,\n int *__restrict RAP_rows,\n int *__restrict RAP_cols,\n Value_type *__restrict RAP_vals,\n int **flagArray_ptrs,\n const int gmem_size,\n int *g_keys,\n Value_type *g_vals,\n int *wk_work_queue,\n int num_neighbors,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // The hash keys stored in shared memory.\n __shared__ /*volatile*/ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The hash values stored in shared memory.\n __shared__ Value_type s_vals[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int rap_int_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_map map( &s_keys[warp_id * SMEM_SIZE],\n &g_keys[rap_int_row_id * gmem_size],\n &s_vals[warp_id * SMEM_SIZE],\n &g_vals[rap_int_row_id * gmem_size],\n gmem_size );\n\n // Loop over rows of RAP_ext\n for ( ; rap_int_row_id < RAP_int_num_rows ; rap_int_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Clear the map.\n map.clear();\n // ---------------------------------\n // First process RAP_int\n // ---------------------------------\n sparse_add_process_row_values(rap_int_row_id, RAP_int_rows, RAP_int_cols, RAP_int_vals, lane_id, map, wk_status);\n\n // ---------------------------------\n // Then process RAP_ext\n // ---------------------------------\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int flag = flagArray_ptrs[i][rap_int_row_id];\n\n if (flag != -1)\n {\n int *RAP_ext_rows = RAP_ext_row_ptrs[i];\n int *RAP_ext_cols = RAP_ext_col_ptrs[i];\n Value_type *RAP_ext_vals = RAP_ext_val_ptrs[i];\n int pos_in_row_ext = flag;\n sparse_add_process_row_values(pos_in_row_ext, RAP_ext_rows, RAP_ext_cols, RAP_ext_vals, lane_id, map, wk_status);\n }\n }\n\n // Store the results.\n int rap_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n rap_col_tmp = utils::Ld::load( &RAP_rows[rap_int_row_id + lane_id] );\n }\n\n int rap_col_it = utils::shfl( rap_col_tmp, 0 );\n int rap_col_end = utils::shfl( rap_col_tmp, 1 );\n // Store the results.\n int count = rap_col_end - rap_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n map.store( count, &RAP_cols[rap_col_it], &RAP_vals[rap_col_it] );\n }\n}\n\n\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace csr_multiply_detail\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nnamespace amgx\n{\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nenum { WARP_SIZE = 32, SMEM_SIZE = 128 };\n\n// ====================================================================================================================\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nCSR_Multiply_Detail >::CSR_Multiply_Detail( bool allocate_values, int grid_size, int max_warp_count, int gmem_size ) :\n Base(allocate_values, grid_size, max_warp_count, gmem_size)\n{}\n\n// ====================================================================================================================\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::count_non_zeroes( const Matrix_d &A, const Matrix_d &B, Matrix_d &C, IVector *Aq1, IVector *Bq1, IVector *Aq2, IVector *Bq2 )\n{\n const int GRID_SIZE = 1024;\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Reset work queue.\n int work_offset = GRID_SIZE * NUM_WARPS;\n CUDA_SAFE_CALL( cudaMemcpy( this->m_work_queue, &work_offset, sizeof(int), cudaMemcpyHostToDevice ) );\n \n // Compute non-zero elements.\n switch ( this->m_num_threads_per_row_count )\n {\n case 2:\n csr_multiply_detail::count_non_zeroes_kernel< 2, CTA_SIZE, SMEM_SIZE, WARP_SIZE, true> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n NULL,\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status );\n cudaCheckError();\n break;\n\n case 4:\n csr_multiply_detail::count_non_zeroes_kernel< 4, CTA_SIZE, SMEM_SIZE, WARP_SIZE, true> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n NULL,\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status );\n cudaCheckError();\n break;\n\n case 8:\n csr_multiply_detail::count_non_zeroes_kernel< 8, CTA_SIZE, SMEM_SIZE, WARP_SIZE, true> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n NULL,\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status );\n cudaCheckError();\n break;\n\n case 16:\n csr_multiply_detail::count_non_zeroes_kernel<16, CTA_SIZE, SMEM_SIZE, WARP_SIZE, true> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n NULL,\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status );\n cudaCheckError();\n break;\n\n default:\n csr_multiply_detail::count_non_zeroes_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n NULL,\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status );\n }\n\n cudaCheckError();\n //CUDA_SAFE_CALL( cudaGetLastError() );\n\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::count_non_zeroes_RAP_sparse_add( Matrix_d &RAP, const Matrix_d &RAP_int, std::vector &RAP_ext_row_offsets, std::vector &RAP_ext_col_indices, std::vector &RAP_ext_values, std::vector &RAP_ext_row_ids)\n\n{\n const int GRID_SIZE = 1024;\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Reset work queue.\n int work_offset = GRID_SIZE * NUM_WARPS;\n CUDA_SAFE_CALL( cudaMemcpy( this->m_work_queue, &work_offset, sizeof(int), cudaMemcpyHostToDevice ) );\n // This is num_owned_coarse_rows\n int RAP_size = RAP.get_num_rows();\n int RAP_int_size = RAP_int.row_offsets.size() - 1;\n\n if (RAP_int_size < RAP_size)\n {\n FatalError(\"RAP_int has less rows than RAP, need to modify sparse RAP add to handle that case\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // Create a device vector of the raw pointers to the array\n // RAP_ext_row_ids_ptrs\n // RAP_ext_row_offsets_ptrs\n // RAP_ext_col_indices_ptrs\n // RAP_ext_values_ptrs\n int num_neighbors = RAP_ext_row_offsets.size();\n std::vector flagArray(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n flagArray[i].resize(RAP_size);\n thrust_wrapper::fill(flagArray[i].begin(), flagArray[i].end(), -1);\n }\n\n cudaCheckError();\n std::vector flagArray_ptrs_h(num_neighbors);\n std::vector RAP_ext_row_offsets_ptrs_h(num_neighbors);\n std::vector RAP_ext_col_indices_ptrs_h(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n flagArray_ptrs_h[i] = amgx::thrust::raw_pointer_cast(&flagArray[i][0]);\n RAP_ext_row_offsets_ptrs_h[i] = amgx::thrust::raw_pointer_cast(&RAP_ext_row_offsets[i][0]);\n RAP_ext_col_indices_ptrs_h[i] = amgx::thrust::raw_pointer_cast(&RAP_ext_col_indices[i][0]);\n }\n\n device_vector_alloc flagArray_ptrs = flagArray_ptrs_h;\n device_vector_alloc RAP_ext_row_offsets_ptrs = RAP_ext_row_offsets_ptrs_h;\n device_vector_alloc RAP_ext_col_indices_ptrs = RAP_ext_col_indices_ptrs_h;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = RAP_ext_row_ids[i].size();\n\n if (size != 0)\n {\n int num_blocks = std::min(4096, (size + 127) / 128);\n //write the position in RAP_ext_row_ids\n csr_multiply_detail::flag_halo_rows <<< num_blocks, 128>>>(\n RAP_ext_row_ids[i].raw(),\n size,\n flagArray[i].raw(),\n i,\n RAP.manager->global_id());\n cudaCheckError();\n }\n }\n\n csr_multiply_detail::count_non_zeroes_RAP_ext_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n RAP_size,\n RAP_int.row_offsets.raw(),\n RAP_int.col_indices.raw(),\n amgx::thrust::raw_pointer_cast(&RAP_ext_row_offsets_ptrs[0]),\n amgx::thrust::raw_pointer_cast(&RAP_ext_col_indices_ptrs[0]),\n RAP.row_offsets.raw(),\n (int *) NULL,\n amgx::thrust::raw_pointer_cast(&flagArray_ptrs[0]),\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status,\n num_neighbors,\n RAP.manager->global_id() );\n cudaCheckError();\n //CUDA_SAFE_CALL( cudaGetLastError() );\n}\n\n// ====================================================================================================================\n\ntemplate< int CTA_SIZE, bool COUNT_ONLY, typename Diag_traits, typename Matrix >\nstatic void\ncount_non_zeroes_ilu1_dispatch( const Matrix &A, Matrix &B, int num_threads_per_row_count, int gmem_size, int *keys, int *work_queue, int *status )\n{\n const int GRID_SIZE = 1024;\n\n switch ( num_threads_per_row_count )\n {\n case 2:\n csr_multiply_detail::count_non_zeroes_ilu1_kernel< 2, CTA_SIZE, SMEM_SIZE, WARP_SIZE, COUNT_ONLY, Diag_traits> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &A.getMatrixColoring().getRowColors()[0] ),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n gmem_size,\n keys,\n work_queue,\n status );\n cudaCheckError();\n break;\n\n case 4:\n csr_multiply_detail::count_non_zeroes_ilu1_kernel< 4, CTA_SIZE, SMEM_SIZE, WARP_SIZE, COUNT_ONLY, Diag_traits> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &A.getMatrixColoring().getRowColors()[0] ),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n gmem_size,\n keys,\n work_queue,\n status );\n cudaCheckError();\n break;\n\n case 8:\n csr_multiply_detail::count_non_zeroes_ilu1_kernel< 8, CTA_SIZE, SMEM_SIZE, WARP_SIZE, COUNT_ONLY, Diag_traits> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &A.getMatrixColoring().getRowColors()[0] ),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n gmem_size,\n keys,\n work_queue,\n status );\n cudaCheckError();\n break;\n\n case 16:\n csr_multiply_detail::count_non_zeroes_ilu1_kernel<16, CTA_SIZE, SMEM_SIZE, WARP_SIZE, COUNT_ONLY, Diag_traits> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &A.getMatrixColoring().getRowColors()[0] ),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n gmem_size,\n keys,\n work_queue,\n status );\n cudaCheckError();\n break;\n\n default:\n csr_multiply_detail::count_non_zeroes_ilu1_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &A.getMatrixColoring().getRowColors()[0] ),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n gmem_size,\n keys,\n work_queue,\n status );\n }\n\n cudaCheckError();\n}\n\n// ====================================================================================================================\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::count_non_zeroes_ilu1( const Matrix_d &A, Matrix_d &B )\n{\n const int GRID_SIZE = 1024;\n\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Reset work queue.\n int work_offset = GRID_SIZE * NUM_WARPS;\n CUDA_SAFE_CALL( cudaMemcpy( this->m_work_queue, &work_offset, sizeof(int), cudaMemcpyHostToDevice ) );\n\n // Count the number of non zeroes.\n if ( A.hasProps(DIAG) )\n count_non_zeroes_ilu1_dispatch(\n A,\n B,\n this->m_num_threads_per_row_count,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status );\n else\n count_non_zeroes_ilu1_dispatch(\n A,\n B,\n this->m_num_threads_per_row_count,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status );\n\n // Compute non-zero elements.\n CUDA_SAFE_CALL( cudaGetLastError() );\n}\n\n// ====================================================================================================================\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::compute_offsets( Matrix_d &C )\n{\n amgx::thrust::device_ptr offsets_begin(C.row_offsets.raw());\n amgx::thrust::device_ptr offsets_end (C.row_offsets.raw() + C.get_num_rows() + 1);\n thrust_wrapper::exclusive_scan( offsets_begin, offsets_end, offsets_begin );\n cudaCheckError();\n}\n\n// ====================================================================================================================\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::compute_sparsity( const Matrix_d &A, const Matrix_d &B, Matrix_d &C )\n{\n const int GRID_SIZE = 1024;\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Reset the work queue.\n int work_offset = GRID_SIZE * NUM_WARPS;\n CUDA_SAFE_CALL( cudaMemcpy( this->m_work_queue, &work_offset, sizeof(int), cudaMemcpyHostToDevice ) );\n\n // Compute the values.\n switch ( this->m_num_threads_per_row_count )\n {\n case 2:\n csr_multiply_detail::count_non_zeroes_kernel< 2, CTA_SIZE, SMEM_SIZE, WARP_SIZE, false> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n NULL,\n NULL,\n NULL,\n NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n NULL );\n cudaCheckError();\n break;\n\n case 4:\n csr_multiply_detail::count_non_zeroes_kernel< 4, CTA_SIZE, SMEM_SIZE, WARP_SIZE, false> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n NULL,\n NULL,\n NULL,\n NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n NULL );\n cudaCheckError();\n break;\n\n case 8:\n csr_multiply_detail::count_non_zeroes_kernel< 8, CTA_SIZE, SMEM_SIZE, WARP_SIZE, false> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n NULL,\n NULL,\n NULL,\n NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n NULL );\n cudaCheckError();\n break;\n\n case 16:\n csr_multiply_detail::count_non_zeroes_kernel<16, CTA_SIZE, SMEM_SIZE, WARP_SIZE, false> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n NULL,\n NULL,\n NULL,\n NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n NULL );\n cudaCheckError();\n break;\n\n default:\n csr_multiply_detail::count_non_zeroes_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n NULL,\n NULL,\n NULL,\n NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n NULL );\n }\n\n cudaCheckError();\n //CUDA_SAFE_CALL( cudaGetLastError() );\n}\n\n// ====================================================================================================================\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::compute_sparsity_ilu1( const Matrix_d &A, Matrix_d &B )\n{\n const int GRID_SIZE = 1024;\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Reset work queue.\n int work_offset = GRID_SIZE * NUM_WARPS;\n CUDA_SAFE_CALL( cudaMemcpy( this->m_work_queue, &work_offset, sizeof(int), cudaMemcpyHostToDevice ) );\n\n // Count the number of non zeroes.\n if ( A.hasProps(DIAG) )\n count_non_zeroes_ilu1_dispatch(\n A,\n B,\n this->m_num_threads_per_row_count,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n NULL );\n else\n count_non_zeroes_ilu1_dispatch(\n A,\n B,\n this->m_num_threads_per_row_count,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n NULL );\n\n // Make sure it worked properly.\n CUDA_SAFE_CALL( cudaGetLastError() );\n}\n\n// ====================================================================================================================\n\ntemplate \n__global__ \nvoid calc_max_nnz_per_row_of_C( \n const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict B_rows,\n int *__restrict C_row_max)\n{\n int a_row_id = blockIdx.x*blockDim.x + threadIdx.x;\n\n int expected_max_row_nnz = 0;\n\n if(a_row_id < A_num_rows)\n {\n for (int a_col_it = A_rows[a_row_id]; a_col_it < A_rows[a_row_id + 1]; ++a_col_it)\n {\n int a_col_id = A_cols[a_col_it];\n expected_max_row_nnz += B_rows[a_col_id+1]-B_rows[a_col_id];\n }\n }\n\n using BR = amgx::cub::BlockReduce;\n\n __shared__ typename BR::TempStorage max_s;\n int max_nnz_block = BR(max_s).Reduce(expected_max_row_nnz, [] __device__ (int a, int b) { return a > b ? a : b; });\n\n if(threadIdx.x == 0)\n {\n C_row_max[blockIdx.x] = max_nnz_block;\n }\n}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nbool CSR_Multiply_Detail >::count_non_zeroes_opt(\n const Matrix_d &A, const Matrix_d &B, Matrix_d &C, int num_threads)\n{\n constexpr int cta_size = 128;\n\n // At least for large matrices it may be optimal to determine the maximum\n // expected hash size by investigation of the matrices\n\n int grid_size = A.get_num_rows() / cta_size + 1;\n\n IVector C_row_max_block(grid_size);\n calc_max_nnz_per_row_of_C<<>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n C_row_max_block.raw());\n cudaCheckError();\n\n int max_nnz = thrust_wrapper::reduce(\n C_row_max_block.raw(), \n C_row_max_block.raw() + C_row_max_block.size(), \n 0, amgx::thrust::maximum());\n\n#define CNZ_OPT(group_size, hash_size) \\\n csr_multiply_detail::count_non_zeroes_kernel_opt \\\n <<>>( \\\n A.get_num_rows(), \\\n A.row_offsets.raw(), \\\n A.col_indices.raw(), \\\n B.row_offsets.raw(), \\\n B.col_indices.raw(), \\\n C.row_offsets.raw());\n cudaCheckError();\n\n // Operation is group per row, where group size is determined by num_threads\n switch ( num_threads )\n {\n // 16 threads per group\n case 16:\n {\n int ngroups = cta_size / 16;\n int grid_size = A.get_num_rows() / ngroups + 1;\n\n if(max_nnz < 256)\n { \n CNZ_OPT(16, 256); \n } \n else if(max_nnz < 512) \n { \n CNZ_OPT(16, 512); \n } \n else if(max_nnz < 1024)\n { \n CNZ_OPT(16, 1024); \n }\n else \n {\n return false;\n }\n }\n break;\n\n // 32 threads per group\n case 32:\n {\n int ngroups = cta_size / 32;\n int grid_size = A.get_num_rows() / ngroups + 1;\n\n if(max_nnz < 256)\n { \n CNZ_OPT(32, 256); \n } \n else if(max_nnz < 512) \n { \n CNZ_OPT(32, 512); \n } \n else if(max_nnz < 1024)\n { \n CNZ_OPT(32, 1024); \n }\n else\n {\n return false;\n }\n }\n break;\n\n default:\n FatalError(\"count_non_zeros_opt only implemented for group size = 8, 16, 32\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n cudaCheckError();\n\n return true;\n}\n\n\ntemplate \ntemplate \nvoid CSR_Multiply_Detail >::cvk_opt(const Matrix_d &A, const Matrix_d &B, Matrix_d &C)\n{\n typedef typename IndPrecisionMap::Type Index_type;\n\n constexpr int cta_size = 128;\n constexpr int ngroups = cta_size / group_size;\n\n const int grid_size = A.get_num_rows() / ngroups + 1;\n\n cudaDeviceProp deviceProps = getDeviceProperties();\n size_t max_shmem_size = deviceProps.sharedMemPerMultiprocessor;\n\n constexpr int shmem_size =\n (sizeof(Value_type)+sizeof(Index_type))*ngroups*hash_size + group_size; \n\n if(shmem_size > max_shmem_size) \n { \n FatalError(\"In compute_values_opt the requested hash size is larger than max.\\n\", \n AMGX_ERR_NOT_IMPLEMENTED); \n } \n\n cudaFuncSetAttribute(csr_multiply_detail::compute_values_kernel_opt \n , \n cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size); \n\n csr_multiply_detail::compute_values_kernel_opt \n \n <<>>( \n A.get_num_rows(), \n A.row_offsets.raw(), \n A.col_indices.raw(), \n A.values.raw(), \n B.row_offsets.raw(), \n B.col_indices.raw(), \n B.values.raw(), \n C.row_offsets.raw(), \n C.col_indices.raw(), \n C.values.raw()); \n cudaCheckError();\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::compute_values_opt( const Matrix_d &A, const Matrix_d &B, Matrix_d &C, int num_threads, int max_nnz )\n{\n int C_nrows = C.get_num_rows();\n int C_nnz = C.get_num_nz();\n int C_max_nnz_per_row = max_nnz;\n\n // The aim is to minimise the hash size while reducing the impact of the linear\n // probing. It might actually be more optimal to just use as large tables as \n // possible, reducing the linear probing cost and maximising the C write cost?\n float C_max_nnz_log2 = log2(static_cast(C_max_nnz_per_row));\n float C_max_nnz_log2_ceil = ceil(C_max_nnz_log2);\n int C_rounded_max = static_cast(2.0*pow(2.0, C_max_nnz_log2_ceil));\n\n // Operation is group per row, where group size is determined by num_threads\n switch ( num_threads )\n {\n case 16: // 16 threads per group\n {\n switch(C_rounded_max)\n {\n case 2: \n case 4:\n case 8: \n case 16:\n case 32: cvk_opt<32, 8>(A, B, C); break;\n case 64: cvk_opt<64, 8>(A, B, C); break;\n case 128: cvk_opt<128, 8>(A, B, C); break;\n case 256: cvk_opt<256, 8>(A, B, C); break;\n case 512: cvk_opt<512, 16>(A, B, C); break;\n default: \n FatalError(\"In compute_values_opt the requested hash size is too large.\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n break;\n case 32: // Warp per group\n {\n switch(C_rounded_max)\n {\n case 2: \n case 4:\n case 8: \n case 16:\n case 32: cvk_opt<32, 32>(A, B, C); break;\n case 64: cvk_opt<64, 32>(A, B, C); break;\n case 128: cvk_opt<128, 32>(A, B, C); break;\n case 256: cvk_opt<256, 32>(A, B, C); break;\n case 512: cvk_opt<512, 32>(A, B, C); break;\n case 1024: cvk_opt<1024, 32>(A, B, C); break;\n default: \n FatalError(\"In compute_values_opt the requested hash size is too large.\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n break;\n default:\n FatalError(\"compute_values_opt only implemented for group size = 16, 32\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n cudaDeviceSynchronize();\n\n cudaCheckError();\n}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::compute_values( const Matrix_d &A, const Matrix_d &B, Matrix_d &C, int num_threads, IVector *Aq1, IVector *Bq1, IVector *Aq2, IVector *Bq2 )\n{\n const int GRID_SIZE = 1024;\n const int CTA_SIZE = 128;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Reset the work queue.\n int work_offset = GRID_SIZE * NUM_WARPS;\n CUDA_SAFE_CALL( cudaMemcpy( this->m_work_queue, &work_offset, sizeof(int), cudaMemcpyHostToDevice ) );\n // Compute the values.\n int *status = NULL;\n\n if ( num_threads != this->m_num_threads_per_row_compute )\n {\n status = this->m_status;\n }\n\n \t\t\n\n switch ( num_threads )\n {\n case 2:\n csr_multiply_detail::compute_values_kernel< 2, Value_type, CTA_SIZE, SMEM_SIZE, WARP_SIZE> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n B.values.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n C.values.raw(),\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_vals,\n this->m_work_queue,\n status );\n break;\n\n case 4:\n csr_multiply_detail::compute_values_kernel< 4, Value_type, CTA_SIZE, SMEM_SIZE, WARP_SIZE> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n B.values.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n C.values.raw(),\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_vals,\n this->m_work_queue,\n status );\n break;\n\n case 8:\n csr_multiply_detail::compute_values_kernel< 8, Value_type, CTA_SIZE, SMEM_SIZE, WARP_SIZE> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n B.values.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n C.values.raw(),\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_vals,\n this->m_work_queue,\n status );\n break;\n\n case 16:\n csr_multiply_detail::compute_values_kernel<16, Value_type, CTA_SIZE, SMEM_SIZE, WARP_SIZE> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n B.values.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n C.values.raw(),\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_vals,\n this->m_work_queue,\n status );\n break;\n\n default:\n csr_multiply_detail::compute_values_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n B.values.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n C.values.raw(),\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_vals,\n this->m_work_queue,\n status );\n }\n\n cudaCheckError();\n //CUDA_SAFE_CALL( cudaGetLastError() );\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::compute_values_RAP_sparse_add( Matrix_d &RAP, const Matrix_d &RAP_int, std::vector &RAP_ext_row_offsets, std::vector &RAP_ext_col_indices, std::vector &RAP_ext_values, std::vector &RAP_ext_row_ids, int num_threads)\n{\n const int GRID_SIZE = 1024;\n const int CTA_SIZE = 128;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Reset the work queue.\n int work_offset = GRID_SIZE * NUM_WARPS;\n CUDA_SAFE_CALL( cudaMemcpy( this->m_work_queue, &work_offset, sizeof(int), cudaMemcpyHostToDevice ) );\n // Compute the values.\n int *status = NULL;\n\n if ( num_threads != this->m_num_threads_per_row_compute )\n {\n status = this->m_status;\n }\n\n // This is num_owned_coarse_rows\n int RAP_size = RAP.get_num_rows();\n int RAP_int_size = RAP_int.row_offsets.size() - 1;\n\n if (RAP_int_size < RAP_size)\n {\n FatalError(\"RAP_int has less rows than RAP, need to modify sparse RAP add to handle that case\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n //TODO: Optimize: reuse arrays from count nonzeros\n int num_neighbors = RAP_ext_row_offsets.size();\n std::vector flagArray(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n flagArray[i].resize(RAP_size);\n thrust_wrapper::fill(flagArray[i].begin(), flagArray[i].end(), -1);\n }\n\n cudaCheckError();\n std::vector flagArray_ptrs_h(num_neighbors);\n std::vector RAP_ext_row_offsets_ptrs_h(num_neighbors);\n std::vector RAP_ext_col_indices_ptrs_h(num_neighbors);\n std::vector RAP_ext_values_ptrs_h(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n flagArray_ptrs_h[i] = amgx::thrust::raw_pointer_cast(&flagArray[i][0]);\n RAP_ext_row_offsets_ptrs_h[i] = amgx::thrust::raw_pointer_cast(&RAP_ext_row_offsets[i][0]);\n RAP_ext_col_indices_ptrs_h[i] = amgx::thrust::raw_pointer_cast(&RAP_ext_col_indices[i][0]);\n RAP_ext_values_ptrs_h[i] = amgx::thrust::raw_pointer_cast(&RAP_ext_values[i][0]);\n }\n\n device_vector_alloc flagArray_ptrs = flagArray_ptrs_h;\n device_vector_alloc RAP_ext_row_offsets_ptrs = RAP_ext_row_offsets_ptrs_h;\n device_vector_alloc RAP_ext_col_indices_ptrs = RAP_ext_col_indices_ptrs_h;\n device_vector_alloc RAP_ext_values_ptrs = RAP_ext_values_ptrs_h;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = RAP_ext_row_ids[i].size();\n\n if (size != 0)\n {\n int num_blocks = std::min(4096, (size + 127) / 128);\n //write the position in RAP_ext_row_ids\n csr_multiply_detail::flag_halo_rows <<< num_blocks, 128>>>(\n RAP_ext_row_ids[i].raw(),\n size,\n flagArray[i].raw(),\n i,\n RAP.manager->global_id());\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n //CUDA_SAFE_CALL( cudaGetLastError() );\n csr_multiply_detail::compute_values_RAP_ext_kernel< Value_type, CTA_SIZE, SMEM_SIZE, WARP_SIZE> <<< GRID_SIZE, CTA_SIZE>>>(\n RAP_size,\n RAP_int.row_offsets.raw(),\n RAP_int.col_indices.raw(),\n RAP_int.values.raw(),\n amgx::thrust::raw_pointer_cast(&RAP_ext_row_offsets_ptrs[0]),\n amgx::thrust::raw_pointer_cast(&RAP_ext_col_indices_ptrs[0]),\n amgx::thrust::raw_pointer_cast(&RAP_ext_values_ptrs[0]),\n RAP.row_offsets.raw(),\n RAP.col_indices.raw(),\n RAP.values.raw(),\n amgx::thrust::raw_pointer_cast(&flagArray_ptrs[0]),\n this->m_gmem_size,\n this->m_keys,\n this->m_vals,\n this->m_work_queue,\n num_neighbors,\n status );\n cudaCheckError();\n //CUDA_SAFE_CALL( cudaGetLastError() );\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define AMGX_CASE_LINE(CASE) template class CSR_Multiply_Detail::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nnamespace csr_multiply_detail\n{\n\n#include \n\n#include \n#include // Included inside the namespace to solve name colisions.\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n__global__\nvoid flag_halo_rows(int *row_ids, int size, int *flagArray, int neighbor, int global_id)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n int row_id = row_ids[tidx];\n flagArray[row_id] = tidx;\n }\n}\n\n__device__ __forceinline__ int get_work( int *queue, int warp_id )\n{\n int offset = -1;\n\n if ( utils::lane_id() == 0 )\n {\n offset = atomicAdd( queue, 1 );\n }\n\n return utils::shfl( offset, 0 );\n}\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool COUNT_ONLY >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncount_non_zeroes_kernel( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const int *B_rows,\n const int *B_cols,\n int *C_rows,\n int *C_cols,\n int *Aq1,\n int *Bq1,\n int *Aq2,\n int *Bq2,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n int c_row_id = a_row_id;\n\n if (Aq1 != NULL)\n {\n a_row_id = Aq1[a_row_id];\n }\n\n // Make sure we have to proceed.\n if ( COUNT_ONLY )\n {\n volatile int *status = reinterpret_cast( wk_status );\n\n if ( set.has_failed() || *status != 0 )\n {\n return;\n }\n }\n\n // Clear the set.\n set.clear();\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = utils::Ld::load( &A_rows[a_row_id + lane_id] );\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int b_row_id = -1;\n\n if ( is_active )\n {\n b_row_id = utils::Ld::load( &A_cols[a_col_it] );\n\n //b_row_id is actually column of A\n if (Aq2 != NULL)\n {\n b_row_id = Aq2[b_row_id];\n\t \t\t\n }\n\n if (Bq1 != NULL)\n {\n b_row_id = Bq1[b_row_id];\n }\n }\n\n // The number of valid rows.\n const int num_rows = __popc( utils::ballot(is_active) );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n const int uniform_b_row_id = utils::shfl( b_row_id, k );\n // Load the range of the row of B.\n int b_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n b_col_tmp = utils::Ld::load( &B_rows[uniform_b_row_id + lane_id] );\n }\n\n int b_col_it = utils::shfl( b_col_tmp, 0 );\n int b_col_end = utils::shfl( b_col_tmp, 1 );\n\n // Iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any(b_col_it < b_col_end) ; b_col_it += WARP_SIZE )\n {\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = utils::Ld::load( &B_cols[b_col_it] );\n\n // b_col_id is actually column of B\n if (Bq2 != NULL)\n {\n b_col_id = Bq2[b_col_id];\n }\n }\n\n set.insert( b_col_id, COUNT_ONLY ? wk_status : NULL );\n }\n }\n }\n\n // Store the results.\n if ( COUNT_ONLY )\n {\n int count = set.compute_size();\n\n if ( lane_id == 0 )\n {\n C_rows[c_row_id] = count;\n }\n }\n else\n {\n int c_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n c_col_tmp = utils::Ld::load( &C_rows[c_row_id + lane_id] );\n }\n\n int c_col_it = utils::shfl( c_col_tmp, 0 );\n int c_col_end = utils::shfl( c_col_tmp, 1 );\n // Store the results.\n int count = c_col_end - c_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n set.store( count, &C_cols[c_col_it] );\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool COUNT_ONLY >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncount_non_zeroes_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict B_rows,\n const int *__restrict B_cols,\n int *__restrict C_rows,\n int *__restrict C_cols,\n int *Aq1,\n int *Bq1,\n int *Aq2,\n int *Bq2,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // The hash keys stored in shared memory.\n __shared__ /*volatile*/ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n int c_row_id = a_row_id;\n\n if (Aq1 != NULL)\n {\n a_row_id = Aq1[a_row_id];\n }\n\n // Make sure we have to proceed.\n if ( COUNT_ONLY )\n {\n volatile int *status = reinterpret_cast( wk_status );\n\n if ( set.has_failed() || *status != 0 )\n {\n return;\n }\n }\n\n // Clear the set.\n set.clear();\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = utils::Ld::load( &A_rows[a_row_id + lane_id] );\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int b_row_id = -1;\n\n if ( is_active )\n {\n b_row_id = utils::Ld::load( &A_cols[a_col_it] );\n\n //b_row_id is actually column of A\n if (Aq2 != NULL)\n {\n b_row_id = Aq2[b_row_id];\n }\n\n if (Bq1 != NULL)\n {\n b_row_id = Bq1[b_row_id];\n }\n }\n\n const int num_rows = __popc( utils::ballot(is_active) );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n const int uniform_b_row_id = utils::shfl( b_row_id, local_k );\n // Load the range of the row of B.\n int b_col_tmp = -1;\n\n if ( is_active_k && lane_id_mod_num_threads < 2 )\n {\n b_col_tmp = utils::Ld::load( &B_rows[uniform_b_row_id + lane_id_mod_num_threads] );\n }\n\n int b_col_it = utils::shfl( b_col_tmp, lane_id_div_num_threads * NUM_THREADS_PER_ROW + 0 );\n int b_col_end = utils::shfl( b_col_tmp, lane_id_div_num_threads * NUM_THREADS_PER_ROW + 1 );\n\n // Iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = utils::Ld::load( &B_cols[b_col_it] );\n\n // b_col_id is actually column of B\n if (Bq2 != NULL)\n {\n b_col_id = Bq2[b_col_id];\n }\n }\n\n set.insert( b_col_id, COUNT_ONLY ? wk_status : NULL );\n }\n }\n }\n\n // Store the results.\n if ( COUNT_ONLY )\n {\n int count = set.compute_size();\n\n if ( lane_id == 0 )\n {\n C_rows[c_row_id] = count;\n }\n }\n else\n {\n int c_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n c_col_tmp = utils::Ld::load( &C_rows[c_row_id + lane_id] );\n }\n\n int c_col_it = utils::shfl( c_col_tmp, 0 );\n int c_col_end = utils::shfl( c_col_tmp, 1 );\n // Store the results.\n int count = c_col_end - c_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n set.store( count, &C_cols[c_col_it] );\n }\n }\n}\n\ntemplate \n__device__ __forceinline__ void sparse_add_process_row(int row_id, const int *__restrict__ row_offsets, const int *__restrict__ col_indices, int lane_id, Hash_set &set, int *wk_status, int global_id, bool print_flag)\n{\n // Load the range of the row of RAP_int\n int col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n col_tmp = utils::Ld::load( &row_offsets[row_id + lane_id] );\n }\n\n int col_it = utils::shfl( col_tmp, 0 );\n int col_end = utils::shfl( col_tmp, 1 );\n\n // Iterate over the columns of RAP_int\n for ( col_it += lane_id ; utils::any(col_it < col_end) ; col_it += WARP_SIZE )\n {\n int col_id = -1;\n\n if ( col_it < col_end )\n {\n col_id = utils::Ld::load( &col_indices[col_it] );\n }\n\n set.insert( col_id, COUNT_ONLY ? wk_status : NULL );\n }\n}\n\n\ntemplate \n__device__ __forceinline__ void sparse_add_process_row_values(int row_id, const int *__restrict__ row_offsets, const int *__restrict__ col_indices, const Value_type *__restrict vals, int lane_id, Hash_map &map, int *wk_status)\n{\n // Load the range of the row.\n int col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n col_tmp = utils::Ld::load( &row_offsets[row_id + lane_id] );\n }\n\n int col_it = utils::shfl( col_tmp, 0 );\n int col_end = utils::shfl( col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( col_it += lane_id ; utils::any(col_it < col_end) ; col_it += WARP_SIZE )\n {\n const bool is_active = col_it < col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int col_id(-1);\n Value_type value = amgx::types::util::get_zero();\n\n if ( is_active )\n {\n col_id = utils::Ld::load( &col_indices[col_it] );\n value = utils::Ld::load( &vals[col_it] );\n }\n\n map.insert( col_id, value, wk_status );\n }\n}\n\n\n\n\ntemplate\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncount_non_zeroes_RAP_ext_kernel( const int RAP_int_num_rows,\n const int *__restrict RAP_int_rows,\n const int *__restrict RAP_int_cols,\n int **RAP_ext_row_ptrs,\n int **RAP_ext_col_ptrs,\n int *__restrict RAP_rows,\n int *__restrict RAP_cols,\n int **flagArray_ptrs,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status,\n int num_neighbors,\n int global_id )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // The hash keys stored in shared memory.\n __shared__ /*volatile*/ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int rap_int_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[rap_int_row_id * gmem_size], gmem_size );\n\n // Loop over rows of RAP_ext\n for ( ; rap_int_row_id < RAP_int_num_rows; rap_int_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Make sure we have to proceed.\n if ( COUNT_ONLY )\n {\n volatile int *status = reinterpret_cast( wk_status );\n\n if ( set.has_failed() || *status != 0 )\n {\n return;\n }\n }\n\n // Clear the set.\n set.clear();\n // ---------------------------------\n // First process RAP_int\n // ---------------------------------\n bool print_flag = false;\n\n if (rap_int_row_id == 3 && global_id == 0)\n {\n print_flag = true;\n }\n\n sparse_add_process_row(rap_int_row_id, RAP_int_rows, RAP_int_cols, lane_id, set, wk_status, global_id, print_flag);\n\n // ---------------------------------\n // Then process RAP_ext\n // ---------------------------------\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int flag = flagArray_ptrs[i][rap_int_row_id];\n\n if (flag != -1)\n {\n int *RAP_ext_rows = RAP_ext_row_ptrs[i];\n int *RAP_ext_cols = RAP_ext_col_ptrs[i];\n int pos_in_row_ext = flag;\n sparse_add_process_row(pos_in_row_ext, RAP_ext_rows, RAP_ext_cols, lane_id, set, wk_status, global_id, print_flag);\n }\n }\n\n // Store the results.\n if ( COUNT_ONLY )\n {\n int count = set.compute_size();\n\n if ( lane_id == 0 )\n {\n RAP_rows[rap_int_row_id] = count;\n }\n }\n else\n {\n int rap_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n rap_col_tmp = utils::Ld::load( &RAP_rows[rap_int_row_id + lane_id] );\n }\n\n int rap_col_it = utils::shfl( rap_col_tmp, 0 );\n int rap_col_end = utils::shfl( rap_col_tmp, 1 );\n // Store the results.\n int count = rap_col_end - rap_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n set.store( count, &RAP_cols[rap_col_it] );\n }\n }\n}\n\n\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Without_external_diag\n{\n static __device__ __forceinline__ bool is_active ( int a_col_it, int a_col_end ) { return a_col_it < a_col_end; }\n static __device__ __forceinline__ bool is_boundary( int a_col_it, int a_col_end ) { return false; }\n};\n\n// ====================================================================================================================\n\nstruct With_external_diag\n{\n static __device__ __forceinline__ bool is_active ( int a_col_it, int a_col_end ) { return a_col_it <= a_col_end; }\n static __device__ __forceinline__ bool is_boundary( int a_col_it, int a_col_end ) { return a_col_it == a_col_end; }\n};\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool COUNT_ONLY, typename Diag_traits >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncount_non_zeroes_ilu1_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict A_coloring,\n int *__restrict C_rows,\n int *__restrict C_cols,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Tables to broadcast values.\n __shared__ volatile int s_b_rows[CTA_SIZE], s_b_colors[CTA_SIZE];\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Make sure we have to proceed.\n if ( COUNT_ONLY )\n {\n volatile int *status = reinterpret_cast( wk_status );\n\n if ( set.has_failed() || *status != 0 )\n {\n return;\n }\n }\n\n // Clear the set.\n set.clear();\n // The color of the row.\n int a_row_color = A_coloring[a_row_id];\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = utils::Ld::load( &A_rows[a_row_id + lane_id] );\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any(Diag_traits::is_active(a_col_it, a_col_end)) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int b_row_id = -1;\n\n if ( is_active )\n {\n b_row_id = utils::Ld::load( &A_cols[a_col_it] );\n }\n\n if ( Diag_traits::is_boundary(a_col_it, a_col_end) )\n {\n b_row_id = a_row_id;\n }\n\n // Push the columns in the set.\n set.insert( b_row_id, COUNT_ONLY ? wk_status : NULL );\n\n // Skip computation if the color of the row is 0.\n if ( a_row_color != 0 )\n {\n // Gather the colors of the columns.\n int b_row_color = -1;\n\n if ( is_active )\n {\n b_row_color = A_coloring[b_row_id];\n }\n\n // The number of valid rows.\n int pred = is_active && b_row_color < a_row_color;\n int vote = utils::ballot( pred );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( pred )\n {\n s_b_rows [warp_id * WARP_SIZE + dest] = b_row_id;\n s_b_colors[warp_id * WARP_SIZE + dest] = b_row_color; // TODO: store an int2 rather than 2 ints of SM35 with 64bit banks.\n }\n\n const int num_rows = __popc( vote );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n const int uniform_b_row_id = s_b_rows [warp_id * WARP_SIZE + k];\n const int uniform_b_color = s_b_colors[warp_id * WARP_SIZE + k];\n // Load the range of the row of B.\n int b_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n b_col_tmp = utils::Ld::load( &A_rows[uniform_b_row_id + lane_id] );\n }\n\n int b_col_it = utils::shfl( b_col_tmp, 0 );\n int b_col_end = utils::shfl( b_col_tmp, 1 );\n\n // Iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any(b_col_it < b_col_end) ; b_col_it += WARP_SIZE )\n {\n int b_col_id = -1, b_col_color = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = utils::Ld::load( &A_cols[b_col_it] );\n b_col_color = utils::Ld::load( &A_coloring[b_col_id] );\n }\n\n int item = -1;\n\n if ( b_col_color >= uniform_b_color && b_col_color != a_row_color )\n {\n item = b_col_id;\n }\n\n set.insert( item, COUNT_ONLY ? wk_status : NULL );\n }\n }\n }\n }\n\n // Store the results.\n if ( COUNT_ONLY )\n {\n int count = set.compute_size();\n\n if ( lane_id == 0 )\n {\n C_rows[a_row_id] = count;\n }\n }\n else\n {\n int c_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n c_col_tmp = utils::Ld::load( &C_rows[a_row_id + lane_id] );\n }\n\n int c_col_it = utils::shfl( c_col_tmp, 0 );\n int c_col_end = utils::shfl( c_col_tmp, 1 );\n // Store the results.\n int count = c_col_end - c_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n set.store( count, &C_cols[c_col_it] );\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool COUNT_ONLY, typename Diag_traits >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncount_non_zeroes_ilu1_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict A_coloring,\n int *__restrict C_rows,\n int *__restrict C_cols,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // Tables to broadcast values.\n __shared__ volatile int s_b_rows[CTA_SIZE], s_b_colors[CTA_SIZE];\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Make sure we have to proceed.\n if ( COUNT_ONLY )\n {\n volatile int *status = reinterpret_cast( wk_status );\n\n if ( set.has_failed() || *status != 0 )\n {\n return;\n }\n }\n\n // Clear the set.\n set.clear();\n // The color of the row.\n int a_row_color = A_coloring[a_row_id];\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = utils::Ld::load( &A_rows[a_row_id + lane_id] );\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any(Diag_traits::is_active(a_col_it, a_col_end)) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int b_row_id = -1;\n\n if ( is_active )\n {\n b_row_id = utils::Ld::load( &A_cols[a_col_it] );\n }\n\n if ( Diag_traits::is_boundary(a_col_it, a_col_end) )\n {\n b_row_id = a_row_id;\n }\n\n // Push the columns in the set.\n set.insert( b_row_id, COUNT_ONLY ? wk_status : NULL );\n\n // Skip computation if the color of the row is 0.\n if ( a_row_color != 0 )\n {\n // Gather the colors of the columns.\n int b_row_color = -1;\n\n if ( is_active )\n {\n b_row_color = A_coloring[b_row_id];\n }\n\n // The number of valid rows.\n int pred = is_active && b_row_color < a_row_color;\n int vote = utils::ballot( pred );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( pred )\n {\n s_b_rows [warp_id * WARP_SIZE + dest] = b_row_id;\n s_b_colors[warp_id * WARP_SIZE + dest] = b_row_color;\n }\n\n const int num_rows = __popc( vote );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int uniform_b_row_id = -1, uniform_b_color = -1;\n\n if ( is_active_k )\n {\n uniform_b_row_id = s_b_rows [warp_id * WARP_SIZE + local_k];\n uniform_b_color = s_b_colors[warp_id * WARP_SIZE + local_k];\n }\n\n // Load the range of the row of B.\n int b_col_tmp = -1;\n\n if ( is_active_k && lane_id_mod_num_threads < 2 )\n {\n b_col_tmp = utils::Ld::load( &A_rows[uniform_b_row_id + lane_id_mod_num_threads] );\n }\n\n int b_col_it = utils::shfl( b_col_tmp, lane_id_div_num_threads * NUM_THREADS_PER_ROW + 0 );\n int b_col_end = utils::shfl( b_col_tmp, lane_id_div_num_threads * NUM_THREADS_PER_ROW + 1 );\n\n // Iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n int b_col_id = -1, b_col_color = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = utils::Ld::load( &A_cols[b_col_it] );\n b_col_color = utils::Ld::load( &A_coloring[b_col_id] );\n }\n\n int item = -1;\n\n if ( b_col_color >= uniform_b_color && b_col_color != a_row_color )\n {\n item = b_col_id;\n }\n\n set.insert( item, COUNT_ONLY ? wk_status : NULL );\n }\n }\n }\n }\n\n // Store the results.\n if ( COUNT_ONLY )\n {\n int count = set.compute_size();\n\n if ( lane_id == 0 )\n {\n C_rows[a_row_id] = count;\n }\n }\n else\n {\n int c_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n c_col_tmp = utils::Ld::load( &C_rows[a_row_id + lane_id] );\n }\n\n int c_col_it = utils::shfl( c_col_tmp, 0 );\n int c_col_end = utils::shfl( c_col_tmp, 1 );\n // Store the results.\n int count = c_col_end - c_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n set.store( count, &C_cols[c_col_it] );\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE, 6 )\nvoid\ncompute_values_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const int *__restrict B_rows,\n const int *__restrict B_cols,\n const Value_type *__restrict B_vals,\n const int *__restrict C_rows,\n int *__restrict C_cols,\n Value_type *__restrict C_vals,\n int *Aq1,\n int *Bq1,\n int *Aq2,\n int *Bq2,\n const int gmem_size,\n int *g_keys,\n Value_type *g_vals,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // The hash keys stored in shared memory.\n __shared__ /*volatile*/ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The hash values stored in shared memory.\n __shared__ Value_type s_vals[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_map map(&s_keys[warp_id * SMEM_SIZE],\n &g_keys[a_row_id * gmem_size],\n &s_vals[warp_id * SMEM_SIZE],\n &g_vals[a_row_id * gmem_size],\n gmem_size );\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n int c_row_id = a_row_id;\n\n if (Aq1 != NULL)\n {\n a_row_id = Aq1[a_row_id];\n }\n\n // Clear the map.\n map.clear();\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = utils::Ld::load( &A_rows[a_row_id + lane_id] );\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int b_row_id = -1;\n Value_type a_value = amgx::types::util::get_zero();\n\n if ( is_active )\n {\n b_row_id = utils::Ld::load( &A_cols[a_col_it] );\n a_value = utils::Ld::load( &A_vals[a_col_it] );\n\n //b_row_id is actually column of A\n if (Aq2 != NULL)\n {\n b_row_id = Aq2[b_row_id];\n }\n\n if (Bq1 != NULL)\n {\n b_row_id = Bq1[b_row_id];\n }\n }\n\n const int num_rows = __popc( utils::ballot(is_active) );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n const int uniform_b_row_id = utils::shfl( b_row_id, k );\n // The value of A.\n const Value_type uniform_a_value = utils::shfl( a_value, k );\n // Load the range of the row of B.\n int b_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n b_col_tmp = utils::Ld::load( &B_rows[uniform_b_row_id + lane_id] );\n }\n\n int b_col_it = utils::shfl( b_col_tmp, 0 );\n int b_col_end = utils::shfl( b_col_tmp, 1 );\n\n // Iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any(b_col_it < b_col_end) ; b_col_it += WARP_SIZE )\n {\n int b_col_id = -1;\n Value_type b_value = amgx::types::util::get_zero();\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = utils::Ld::load( &B_cols[b_col_it] );\n b_value = utils::Ld::load( &B_vals[b_col_it] );\n\n if (Bq2 != NULL)\n {\n b_col_id = Bq2[b_col_id];\n }\n }\n\n map.insert( b_col_id, uniform_a_value * b_value, wk_status );\n }\n }\n }\n\n // Store the results.\n int c_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n c_col_tmp = utils::Ld::load( &C_rows[c_row_id + lane_id] );\n }\n\n int c_col_it = utils::shfl( c_col_tmp, 0 );\n int c_col_end = utils::shfl( c_col_tmp, 1 );\n // Store the results.\n int count = c_col_end - c_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n map.store( count, &C_cols[c_col_it], &C_vals[c_col_it] );\n }\n\n}\n\ntemplate \n__global__ \nvoid count_non_zeroes_kernel_opt( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict B_rows,\n const int *__restrict B_cols,\n int *__restrict C_row_counts)\n{\n // Defines for hash container\n constexpr int SLOT_VACANT = -1;\n\n // Group indices\n constexpr int ngroups = CTA_SIZE / GROUP_SIZE;\n const int group_id = threadIdx.x / GROUP_SIZE;\n const int lane_id = threadIdx.x % GROUP_SIZE;\n\n // One row of A per group\n const int a_row_id = blockIdx.x * ngroups + group_id;\n\n // Block-level hash container storage\n __shared__ int key_s[ngroups*HASH_SIZE];\n __shared__ int counts[ngroups];\n\n // Initialise the keys and values.\n#pragma unroll\n for(int i = threadIdx.x; i < ngroups*HASH_SIZE; i += CTA_SIZE)\n {\n key_s[i] = SLOT_VACANT; // Inserted keys will be in range [0,N]\n }\n\n for(int i = 0; i < ngroups; ++i)\n {\n counts[i] = 0;\n }\n\n //__syncwarp();\n __syncthreads();\n\n int* key_group_s = &key_s[group_id*HASH_SIZE];\n\n if(a_row_id < A_num_rows)\n {\n // Distributed columns of row of A over threads in group.\n for (int a_col_it = A_rows[a_row_id] + lane_id; a_col_it < A_rows[a_row_id + 1]; a_col_it += GROUP_SIZE)\n {\n int a_col_id = A_cols[a_col_it];\n\n // Outer product of element of A and row of B.\n for (int b_col_it = B_rows[a_col_id]; b_col_it < B_rows[a_col_id + 1] ; ++b_col_it)\n {\n int b_col_id = B_cols[b_col_it];\n\n int hash = b_col_id % HASH_SIZE;\n\n // By construction this algorithm should guarantee \n // all keys can be inserted\n while(true)\n {\n // If the slot is vacant, then attempt acquire\n int key = key_group_s[hash];\n\n // Check if the key was already b_col_id, or was already set \n // Insert the product\n if(key == b_col_id)\n {\n break;\n }\n\n if(key == SLOT_VACANT)\n {\n int new_key = atomicCAS(&key_group_s[hash], SLOT_VACANT, b_col_id);\n\n if(new_key == SLOT_VACANT || new_key == b_col_id)\n {\n break;\n }\n }\n\n // We did not secure a slot, so linear probe to next slot\n hash = (hash + 1) % HASH_SIZE;\n }\n }\n }\n }\n\n //__syncwarp();\n __syncthreads();\n\n // Store the results.\n if(a_row_id < A_num_rows)\n {\n#pragma unroll\n for(int i = lane_id; i < HASH_SIZE; i += GROUP_SIZE)\n {\n if(key_group_s[i] != SLOT_VACANT)\n {\n atomicAdd(&counts[group_id], 1);\n }\n }\n }\n\n __syncthreads();\n\n if(a_row_id < A_num_rows && lane_id == 0)\n {\n C_row_counts[a_row_id] = counts[group_id];\n }\n}\n\ntemplate \n__global__ \nvoid compute_values_kernel_opt( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const ValueType *__restrict A_vals,\n const int *__restrict B_rows,\n const int *__restrict B_cols,\n const ValueType *__restrict B_vals,\n const int *__restrict C_rows,\n int *__restrict C_cols,\n ValueType *__restrict C_vals)\n{\n // Defines for hash container\n constexpr int SLOT_VACANT = -1;\n\n // Group indices\n constexpr int ngroups = CTA_SIZE / GROUP_SIZE;\n const int group_id = threadIdx.x / GROUP_SIZE;\n const int lane_id = threadIdx.x % GROUP_SIZE;\n\n // One row of A per group\n const int a_row_id = blockIdx.x * ngroups + group_id;\n\n // Dynamic sized shared memory\n extern __shared__ int s[];\n\n // Block-level hash container storage\n int* key_s = s;\n ValueType* data_s = (ValueType*)&key_s[ngroups*HASH_SIZE];\n int* col_ind_s = (int*)&data_s[ngroups*HASH_SIZE];\n\n // Group-level hash containers\n int* key_group_s = &key_s[group_id*HASH_SIZE];\n ValueType* data_group_s = &data_s[group_id*HASH_SIZE];\n\n // Initialise the keys and values.\n#pragma unroll\n for(int i = threadIdx.x; i < ngroups*HASH_SIZE; i += CTA_SIZE)\n {\n key_s[i] = SLOT_VACANT; // Inserted keys will be in range [0,N]\n data_s[i] = amgx::types::util::get_zero(); // We will sum into values\n }\n\n if(lane_id == 0)\n {\n col_ind_s[group_id] = 0;\n }\n\n //__syncwarp();\n __syncthreads();\n\n if(a_row_id < A_num_rows)\n {\n // Distributed columns of row of A over threads in group.\n for (int a_col_it = A_rows[a_row_id] + lane_id; a_col_it < A_rows[a_row_id + 1]; a_col_it += GROUP_SIZE)\n {\n int a_col_id = A_cols[a_col_it];\n\n // Outer product of element of A and row of B.\n for (int b_col_it = B_rows[a_col_id]; b_col_it < B_rows[a_col_id + 1] ; ++b_col_it)\n {\n ValueType val = A_vals[a_col_it]*B_vals[b_col_it];\n\n // Short circuit if zero\n if(amgx::types::util::is_zero(val)) continue;\n\n int b_col_id = B_cols[b_col_it];\n\n // XXX UPDATE HASHING APPROACH\n int hash = b_col_id % HASH_SIZE;\n\n // By construction this algorithm should guarantee \n // all keys can be inserted\n while(true)\n {\n // If the slot is vacant, then attempt acquire\n int key = key_group_s[hash];\n if(key == SLOT_VACANT)\n {\n int new_key = atomicCAS(&key_group_s[hash], SLOT_VACANT, b_col_id);\n if(new_key == SLOT_VACANT || new_key == b_col_id)\n {\n key = b_col_id;\n }\n }\n\n // Check if the key was already b_col_id, or was already set \n // Insert the product\n if(key == b_col_id)\n {\n utils::atomic_add(&data_group_s[hash], val);\n break;\n }\n\n // We did not secure a slot, so linear probe to next slot\n hash = (hash + 1) % HASH_SIZE;\n }\n }\n }\n }\n\n //__syncwarp();\n __syncthreads();\n\n // Store the results.\n int c_row_id = a_row_id;\n int c_col_it = C_rows[c_row_id];\n\n if(a_row_id < A_num_rows)\n {\n#pragma unroll\n for(int i = lane_id; i < HASH_SIZE; i += GROUP_SIZE)\n {\n if(key_group_s[i] != SLOT_VACANT)\n {\n // Short circuit if zero\n if(amgx::types::util::is_zero(data_group_s[i])) continue;\n\n int ind = atomicAdd(&col_ind_s[group_id], 1);\n\n C_cols[c_col_it + ind] = key_group_s[i];\n C_vals[c_col_it + ind] = data_group_s[i];\n }\n }\n }\n}\n\ntemplate< int NUM_THREADS_PER_ROW, typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE, 6 )\nvoid\ncompute_values_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const int *__restrict B_rows,\n const int *__restrict B_cols,\n const Value_type *__restrict B_vals,\n const int *__restrict C_rows,\n int *__restrict C_cols,\n Value_type *__restrict C_vals,\n int *Aq1,\n int *Bq1,\n int *Aq2,\n int *Bq2,\n const int gmem_size,\n int *g_keys,\n Value_type *g_vals,\n int *wk_work_queue,\n int *wk_status )\n{\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // The hash keys stored in shared memory.\n __shared__ /*volatile*/ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The hash values stored in shared memory.\n __shared__ Value_type s_vals[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_map map(&s_keys[warp_id * SMEM_SIZE],\n &g_keys[a_row_id * gmem_size],\n &s_vals[warp_id * SMEM_SIZE],\n &g_vals[a_row_id * gmem_size],\n gmem_size );\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n int c_row_id = a_row_id;\n\n if (Aq1 != NULL)\n {\n a_row_id = Aq1[a_row_id];\n }\n\n // Clear the map.\n map.clear();\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = utils::Ld::load( &A_rows[a_row_id + lane_id] );\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int b_row_id(-1);\n Value_type a_value = amgx::types::util::get_zero();\n\n if ( is_active )\n {\n b_row_id = utils::Ld::load( &A_cols[a_col_it] );\n a_value = utils::Ld::load( &A_vals[a_col_it] );\n\n //b_row_id is actually column of A\n if (Aq2 != NULL)\n {\n b_row_id = Aq2[b_row_id];\n }\n\n if (Bq1 != NULL)\n {\n b_row_id = Bq1[b_row_id];\n }\n }\n\n const int num_rows = __popc( utils::ballot(is_active) );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n const int uniform_b_row_id = utils::shfl( b_row_id, k + lane_id_div_num_threads );\n // The value of A.\n const Value_type uniform_a_value = utils::shfl( a_value, k + lane_id_div_num_threads );\n // Load the range of the row of B.\n int b_col_tmp = -1;\n\n if ( is_active_k && lane_id_mod_num_threads < 2 )\n {\n b_col_tmp = utils::Ld::load( &B_rows[uniform_b_row_id + lane_id_mod_num_threads] );\n }\n\n int b_col_it = utils::shfl( b_col_tmp, lane_id_div_num_threads * NUM_THREADS_PER_ROW + 0 );\n int b_col_end = utils::shfl( b_col_tmp, lane_id_div_num_threads * NUM_THREADS_PER_ROW + 1 );\n\n // Iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n int b_col_id(-1);\n Value_type b_value = amgx::types::util::get_zero();\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = utils::Ld::load( &B_cols[b_col_it] );\n b_value = utils::Ld::load( &B_vals[b_col_it] );\n\n //b_col_id is actually column of B\n if (Bq2 != NULL)\n {\n b_col_id = Bq2[b_col_id];\n }\n }\n\n map.insert( b_col_id, uniform_a_value * b_value, wk_status );\n }\n }\n }\n\n // Store the results.\n int c_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n c_col_tmp = utils::Ld::load( &C_rows[c_row_id + lane_id] );\n }\n\n int c_col_it = utils::shfl( c_col_tmp, 0 );\n int c_col_end = utils::shfl( c_col_tmp, 1 );\n // Store the results.\n int count = c_col_end - c_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n map.store( count, &C_cols[c_col_it], &C_vals[c_col_it] );\n }\n\n}\n\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE, 6 )\nvoid\ncompute_values_RAP_ext_kernel( const int RAP_int_num_rows,\n const int *__restrict RAP_int_rows,\n const int *__restrict RAP_int_cols,\n const Value_type *__restrict RAP_int_vals,\n int **RAP_ext_row_ptrs,\n int **RAP_ext_col_ptrs,\n Value_type **RAP_ext_val_ptrs,\n int *__restrict RAP_rows,\n int *__restrict RAP_cols,\n Value_type *__restrict RAP_vals,\n int **flagArray_ptrs,\n const int gmem_size,\n int *g_keys,\n Value_type *g_vals,\n int *wk_work_queue,\n int num_neighbors,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // The hash keys stored in shared memory.\n __shared__ /*volatile*/ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The hash values stored in shared memory.\n __shared__ Value_type s_vals[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int rap_int_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_map map( &s_keys[warp_id * SMEM_SIZE],\n &g_keys[rap_int_row_id * gmem_size],\n &s_vals[warp_id * SMEM_SIZE],\n &g_vals[rap_int_row_id * gmem_size],\n gmem_size );\n\n // Loop over rows of RAP_ext\n for ( ; rap_int_row_id < RAP_int_num_rows ; rap_int_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Clear the map.\n map.clear();\n // ---------------------------------\n // First process RAP_int\n // ---------------------------------\n sparse_add_process_row_values(rap_int_row_id, RAP_int_rows, RAP_int_cols, RAP_int_vals, lane_id, map, wk_status);\n\n // ---------------------------------\n // Then process RAP_ext\n // ---------------------------------\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int flag = flagArray_ptrs[i][rap_int_row_id];\n\n if (flag != -1)\n {\n int *RAP_ext_rows = RAP_ext_row_ptrs[i];\n int *RAP_ext_cols = RAP_ext_col_ptrs[i];\n Value_type *RAP_ext_vals = RAP_ext_val_ptrs[i];\n int pos_in_row_ext = flag;\n sparse_add_process_row_values(pos_in_row_ext, RAP_ext_rows, RAP_ext_cols, RAP_ext_vals, lane_id, map, wk_status);\n }\n }\n\n // Store the results.\n int rap_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n rap_col_tmp = utils::Ld::load( &RAP_rows[rap_int_row_id + lane_id] );\n }\n\n int rap_col_it = utils::shfl( rap_col_tmp, 0 );\n int rap_col_end = utils::shfl( rap_col_tmp, 1 );\n // Store the results.\n int count = rap_col_end - rap_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n map.store( count, &RAP_cols[rap_col_it], &RAP_vals[rap_col_it] );\n }\n}\n\n\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace csr_multiply_detail\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nnamespace amgx\n{\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nenum { WARP_SIZE = 32, SMEM_SIZE = 128 };\n\n// ====================================================================================================================\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nCSR_Multiply_Detail >::CSR_Multiply_Detail( bool allocate_values, int grid_size, int max_warp_count, int gmem_size ) :\n Base(allocate_values, grid_size, max_warp_count, gmem_size)\n{}\n\n// ====================================================================================================================\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::count_non_zeroes( const Matrix_d &A, const Matrix_d &B, Matrix_d &C, IVector *Aq1, IVector *Bq1, IVector *Aq2, IVector *Bq2 )\n{\n const int GRID_SIZE = 1024;\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Reset work queue.\n int work_offset = GRID_SIZE * NUM_WARPS;\n CUDA_SAFE_CALL( hipMemcpy( this->m_work_queue, &work_offset, sizeof(int), hipMemcpyHostToDevice ) );\n \n // Compute non-zero elements.\n switch ( this->m_num_threads_per_row_count )\n {\n case 2:\n csr_multiply_detail::count_non_zeroes_kernel< 2, CTA_SIZE, SMEM_SIZE, WARP_SIZE, true> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n NULL,\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status );\n cudaCheckError();\n break;\n\n case 4:\n csr_multiply_detail::count_non_zeroes_kernel< 4, CTA_SIZE, SMEM_SIZE, WARP_SIZE, true> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n NULL,\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status );\n cudaCheckError();\n break;\n\n case 8:\n csr_multiply_detail::count_non_zeroes_kernel< 8, CTA_SIZE, SMEM_SIZE, WARP_SIZE, true> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n NULL,\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status );\n cudaCheckError();\n break;\n\n case 16:\n csr_multiply_detail::count_non_zeroes_kernel<16, CTA_SIZE, SMEM_SIZE, WARP_SIZE, true> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n NULL,\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status );\n cudaCheckError();\n break;\n\n default:\n csr_multiply_detail::count_non_zeroes_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n NULL,\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status );\n }\n\n cudaCheckError();\n //CUDA_SAFE_CALL( cudaGetLastError() );\n\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::count_non_zeroes_RAP_sparse_add( Matrix_d &RAP, const Matrix_d &RAP_int, std::vector &RAP_ext_row_offsets, std::vector &RAP_ext_col_indices, std::vector &RAP_ext_values, std::vector &RAP_ext_row_ids)\n\n{\n const int GRID_SIZE = 1024;\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Reset work queue.\n int work_offset = GRID_SIZE * NUM_WARPS;\n CUDA_SAFE_CALL( hipMemcpy( this->m_work_queue, &work_offset, sizeof(int), hipMemcpyHostToDevice ) );\n // This is num_owned_coarse_rows\n int RAP_size = RAP.get_num_rows();\n int RAP_int_size = RAP_int.row_offsets.size() - 1;\n\n if (RAP_int_size < RAP_size)\n {\n FatalError(\"RAP_int has less rows than RAP, need to modify sparse RAP add to handle that case\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // Create a device vector of the raw pointers to the array\n // RAP_ext_row_ids_ptrs\n // RAP_ext_row_offsets_ptrs\n // RAP_ext_col_indices_ptrs\n // RAP_ext_values_ptrs\n int num_neighbors = RAP_ext_row_offsets.size();\n std::vector flagArray(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n flagArray[i].resize(RAP_size);\n thrust_wrapper::fill(flagArray[i].begin(), flagArray[i].end(), -1);\n }\n\n cudaCheckError();\n std::vector flagArray_ptrs_h(num_neighbors);\n std::vector RAP_ext_row_offsets_ptrs_h(num_neighbors);\n std::vector RAP_ext_col_indices_ptrs_h(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n flagArray_ptrs_h[i] = amgx::thrust::raw_pointer_cast(&flagArray[i][0]);\n RAP_ext_row_offsets_ptrs_h[i] = amgx::thrust::raw_pointer_cast(&RAP_ext_row_offsets[i][0]);\n RAP_ext_col_indices_ptrs_h[i] = amgx::thrust::raw_pointer_cast(&RAP_ext_col_indices[i][0]);\n }\n\n device_vector_alloc flagArray_ptrs = flagArray_ptrs_h;\n device_vector_alloc RAP_ext_row_offsets_ptrs = RAP_ext_row_offsets_ptrs_h;\n device_vector_alloc RAP_ext_col_indices_ptrs = RAP_ext_col_indices_ptrs_h;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = RAP_ext_row_ids[i].size();\n\n if (size != 0)\n {\n int num_blocks = std::min(4096, (size + 127) / 128);\n //write the position in RAP_ext_row_ids\n csr_multiply_detail::flag_halo_rows <<< num_blocks, 128>>>(\n RAP_ext_row_ids[i].raw(),\n size,\n flagArray[i].raw(),\n i,\n RAP.manager->global_id());\n cudaCheckError();\n }\n }\n\n csr_multiply_detail::count_non_zeroes_RAP_ext_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n RAP_size,\n RAP_int.row_offsets.raw(),\n RAP_int.col_indices.raw(),\n amgx::thrust::raw_pointer_cast(&RAP_ext_row_offsets_ptrs[0]),\n amgx::thrust::raw_pointer_cast(&RAP_ext_col_indices_ptrs[0]),\n RAP.row_offsets.raw(),\n (int *) NULL,\n amgx::thrust::raw_pointer_cast(&flagArray_ptrs[0]),\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status,\n num_neighbors,\n RAP.manager->global_id() );\n cudaCheckError();\n //CUDA_SAFE_CALL( cudaGetLastError() );\n}\n\n// ====================================================================================================================\n\ntemplate< int CTA_SIZE, bool COUNT_ONLY, typename Diag_traits, typename Matrix >\nstatic void\ncount_non_zeroes_ilu1_dispatch( const Matrix &A, Matrix &B, int num_threads_per_row_count, int gmem_size, int *keys, int *work_queue, int *status )\n{\n const int GRID_SIZE = 1024;\n\n switch ( num_threads_per_row_count )\n {\n case 2:\n csr_multiply_detail::count_non_zeroes_ilu1_kernel< 2, CTA_SIZE, SMEM_SIZE, WARP_SIZE, COUNT_ONLY, Diag_traits> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &A.getMatrixColoring().getRowColors()[0] ),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n gmem_size,\n keys,\n work_queue,\n status );\n cudaCheckError();\n break;\n\n case 4:\n csr_multiply_detail::count_non_zeroes_ilu1_kernel< 4, CTA_SIZE, SMEM_SIZE, WARP_SIZE, COUNT_ONLY, Diag_traits> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &A.getMatrixColoring().getRowColors()[0] ),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n gmem_size,\n keys,\n work_queue,\n status );\n cudaCheckError();\n break;\n\n case 8:\n csr_multiply_detail::count_non_zeroes_ilu1_kernel< 8, CTA_SIZE, SMEM_SIZE, WARP_SIZE, COUNT_ONLY, Diag_traits> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &A.getMatrixColoring().getRowColors()[0] ),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n gmem_size,\n keys,\n work_queue,\n status );\n cudaCheckError();\n break;\n\n case 16:\n csr_multiply_detail::count_non_zeroes_ilu1_kernel<16, CTA_SIZE, SMEM_SIZE, WARP_SIZE, COUNT_ONLY, Diag_traits> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &A.getMatrixColoring().getRowColors()[0] ),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n gmem_size,\n keys,\n work_queue,\n status );\n cudaCheckError();\n break;\n\n default:\n csr_multiply_detail::count_non_zeroes_ilu1_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &A.getMatrixColoring().getRowColors()[0] ),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n gmem_size,\n keys,\n work_queue,\n status );\n }\n\n cudaCheckError();\n}\n\n// ====================================================================================================================\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::count_non_zeroes_ilu1( const Matrix_d &A, Matrix_d &B )\n{\n const int GRID_SIZE = 1024;\n\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Reset work queue.\n int work_offset = GRID_SIZE * NUM_WARPS;\n CUDA_SAFE_CALL( hipMemcpy( this->m_work_queue, &work_offset, sizeof(int), hipMemcpyHostToDevice ) );\n\n // Count the number of non zeroes.\n if ( A.hasProps(DIAG) )\n count_non_zeroes_ilu1_dispatch(\n A,\n B,\n this->m_num_threads_per_row_count,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status );\n else\n count_non_zeroes_ilu1_dispatch(\n A,\n B,\n this->m_num_threads_per_row_count,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n this->m_status );\n\n // Compute non-zero elements.\n CUDA_SAFE_CALL( hipGetLastError() );\n}\n\n// ====================================================================================================================\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::compute_offsets( Matrix_d &C )\n{\n amgx::thrust::device_ptr offsets_begin(C.row_offsets.raw());\n amgx::thrust::device_ptr offsets_end (C.row_offsets.raw() + C.get_num_rows() + 1);\n thrust_wrapper::exclusive_scan( offsets_begin, offsets_end, offsets_begin );\n cudaCheckError();\n}\n\n// ====================================================================================================================\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::compute_sparsity( const Matrix_d &A, const Matrix_d &B, Matrix_d &C )\n{\n const int GRID_SIZE = 1024;\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Reset the work queue.\n int work_offset = GRID_SIZE * NUM_WARPS;\n CUDA_SAFE_CALL( hipMemcpy( this->m_work_queue, &work_offset, sizeof(int), hipMemcpyHostToDevice ) );\n\n // Compute the values.\n switch ( this->m_num_threads_per_row_count )\n {\n case 2:\n csr_multiply_detail::count_non_zeroes_kernel< 2, CTA_SIZE, SMEM_SIZE, WARP_SIZE, false> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n NULL,\n NULL,\n NULL,\n NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n NULL );\n cudaCheckError();\n break;\n\n case 4:\n csr_multiply_detail::count_non_zeroes_kernel< 4, CTA_SIZE, SMEM_SIZE, WARP_SIZE, false> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n NULL,\n NULL,\n NULL,\n NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n NULL );\n cudaCheckError();\n break;\n\n case 8:\n csr_multiply_detail::count_non_zeroes_kernel< 8, CTA_SIZE, SMEM_SIZE, WARP_SIZE, false> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n NULL,\n NULL,\n NULL,\n NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n NULL );\n cudaCheckError();\n break;\n\n case 16:\n csr_multiply_detail::count_non_zeroes_kernel<16, CTA_SIZE, SMEM_SIZE, WARP_SIZE, false> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n NULL,\n NULL,\n NULL,\n NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n NULL );\n cudaCheckError();\n break;\n\n default:\n csr_multiply_detail::count_non_zeroes_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n NULL,\n NULL,\n NULL,\n NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n NULL );\n }\n\n cudaCheckError();\n //CUDA_SAFE_CALL( cudaGetLastError() );\n}\n\n// ====================================================================================================================\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::compute_sparsity_ilu1( const Matrix_d &A, Matrix_d &B )\n{\n const int GRID_SIZE = 1024;\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Reset work queue.\n int work_offset = GRID_SIZE * NUM_WARPS;\n CUDA_SAFE_CALL( hipMemcpy( this->m_work_queue, &work_offset, sizeof(int), hipMemcpyHostToDevice ) );\n\n // Count the number of non zeroes.\n if ( A.hasProps(DIAG) )\n count_non_zeroes_ilu1_dispatch(\n A,\n B,\n this->m_num_threads_per_row_count,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n NULL );\n else\n count_non_zeroes_ilu1_dispatch(\n A,\n B,\n this->m_num_threads_per_row_count,\n this->m_gmem_size,\n this->m_keys,\n this->m_work_queue,\n NULL );\n\n // Make sure it worked properly.\n CUDA_SAFE_CALL( hipGetLastError() );\n}\n\n// ====================================================================================================================\n\ntemplate \n__global__ \nvoid calc_max_nnz_per_row_of_C( \n const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict B_rows,\n int *__restrict C_row_max)\n{\n int a_row_id = blockIdx.x*blockDim.x + threadIdx.x;\n\n int expected_max_row_nnz = 0;\n\n if(a_row_id < A_num_rows)\n {\n for (int a_col_it = A_rows[a_row_id]; a_col_it < A_rows[a_row_id + 1]; ++a_col_it)\n {\n int a_col_id = A_cols[a_col_it];\n expected_max_row_nnz += B_rows[a_col_id+1]-B_rows[a_col_id];\n }\n }\n\n using BR = amgx::hipcub::BlockReduce;\n\n __shared__ typename BR::TempStorage max_s;\n int max_nnz_block = BR(max_s).Reduce(expected_max_row_nnz, [] __device__ (int a, int b) { return a > b ? a : b; });\n\n if(threadIdx.x == 0)\n {\n C_row_max[blockIdx.x] = max_nnz_block;\n }\n}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nbool CSR_Multiply_Detail >::count_non_zeroes_opt(\n const Matrix_d &A, const Matrix_d &B, Matrix_d &C, int num_threads)\n{\n constexpr int cta_size = 128;\n\n // At least for large matrices it may be optimal to determine the maximum\n // expected hash size by investigation of the matrices\n\n int grid_size = A.get_num_rows() / cta_size + 1;\n\n IVector C_row_max_block(grid_size);\n calc_max_nnz_per_row_of_C<<>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n B.row_offsets.raw(),\n C_row_max_block.raw());\n cudaCheckError();\n\n int max_nnz = thrust_wrapper::reduce(\n C_row_max_block.raw(), \n C_row_max_block.raw() + C_row_max_block.size(), \n 0, amgx::thrust::maximum());\n\n#define CNZ_OPT(group_size, hash_size) \\\n csr_multiply_detail::count_non_zeroes_kernel_opt \\\n <<>>( \\\n A.get_num_rows(), \\\n A.row_offsets.raw(), \\\n A.col_indices.raw(), \\\n B.row_offsets.raw(), \\\n B.col_indices.raw(), \\\n C.row_offsets.raw());\n cudaCheckError();\n\n // Operation is group per row, where group size is determined by num_threads\n switch ( num_threads )\n {\n // 16 threads per group\n case 16:\n {\n int ngroups = cta_size / 16;\n int grid_size = A.get_num_rows() / ngroups + 1;\n\n if(max_nnz < 256)\n { \n CNZ_OPT(16, 256); \n } \n else if(max_nnz < 512) \n { \n CNZ_OPT(16, 512); \n } \n else if(max_nnz < 1024)\n { \n CNZ_OPT(16, 1024); \n }\n else \n {\n return false;\n }\n }\n break;\n\n // 32 threads per group\n case 32:\n {\n int ngroups = cta_size / 32;\n int grid_size = A.get_num_rows() / ngroups + 1;\n\n if(max_nnz < 256)\n { \n CNZ_OPT(32, 256); \n } \n else if(max_nnz < 512) \n { \n CNZ_OPT(32, 512); \n } \n else if(max_nnz < 1024)\n { \n CNZ_OPT(32, 1024); \n }\n else\n {\n return false;\n }\n }\n break;\n\n default:\n FatalError(\"count_non_zeros_opt only implemented for group size = 8, 16, 32\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n cudaCheckError();\n\n return true;\n}\n\n\ntemplate \ntemplate \nvoid CSR_Multiply_Detail >::cvk_opt(const Matrix_d &A, const Matrix_d &B, Matrix_d &C)\n{\n typedef typename IndPrecisionMap::Type Index_type;\n\n constexpr int cta_size = 128;\n constexpr int ngroups = cta_size / group_size;\n\n const int grid_size = A.get_num_rows() / ngroups + 1;\n\n hipDeviceProp_t deviceProps = getDeviceProperties();\n size_t max_shmem_size = deviceProps.sharedMemPerMultiprocessor;\n\n constexpr int shmem_size =\n (sizeof(Value_type)+sizeof(Index_type))*ngroups*hash_size + group_size; \n\n if(shmem_size > max_shmem_size) \n { \n FatalError(\"In compute_values_opt the requested hash size is larger than max.\\n\", \n AMGX_ERR_NOT_IMPLEMENTED); \n } \n\n hipFuncSetAttribute(reinterpret_cast(csr_multiply_detail::compute_values_kernel_opt \n ), \n hipFuncAttributeMaxDynamicSharedMemorySize, shmem_size); \n\n csr_multiply_detail::compute_values_kernel_opt \n \n <<>>( \n A.get_num_rows(), \n A.row_offsets.raw(), \n A.col_indices.raw(), \n A.values.raw(), \n B.row_offsets.raw(), \n B.col_indices.raw(), \n B.values.raw(), \n C.row_offsets.raw(), \n C.col_indices.raw(), \n C.values.raw()); \n cudaCheckError();\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::compute_values_opt( const Matrix_d &A, const Matrix_d &B, Matrix_d &C, int num_threads, int max_nnz )\n{\n int C_nrows = C.get_num_rows();\n int C_nnz = C.get_num_nz();\n int C_max_nnz_per_row = max_nnz;\n\n // The aim is to minimise the hash size while reducing the impact of the linear\n // probing. It might actually be more optimal to just use as large tables as \n // possible, reducing the linear probing cost and maximising the C write cost?\n float C_max_nnz_log2 = log2(static_cast(C_max_nnz_per_row));\n float C_max_nnz_log2_ceil = ceil(C_max_nnz_log2);\n int C_rounded_max = static_cast(2.0*pow(2.0, C_max_nnz_log2_ceil));\n\n // Operation is group per row, where group size is determined by num_threads\n switch ( num_threads )\n {\n case 16: // 16 threads per group\n {\n switch(C_rounded_max)\n {\n case 2: \n case 4:\n case 8: \n case 16:\n case 32: cvk_opt<32, 8>(A, B, C); break;\n case 64: cvk_opt<64, 8>(A, B, C); break;\n case 128: cvk_opt<128, 8>(A, B, C); break;\n case 256: cvk_opt<256, 8>(A, B, C); break;\n case 512: cvk_opt<512, 16>(A, B, C); break;\n default: \n FatalError(\"In compute_values_opt the requested hash size is too large.\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n break;\n case 32: // Warp per group\n {\n switch(C_rounded_max)\n {\n case 2: \n case 4:\n case 8: \n case 16:\n case 32: cvk_opt<32, 32>(A, B, C); break;\n case 64: cvk_opt<64, 32>(A, B, C); break;\n case 128: cvk_opt<128, 32>(A, B, C); break;\n case 256: cvk_opt<256, 32>(A, B, C); break;\n case 512: cvk_opt<512, 32>(A, B, C); break;\n case 1024: cvk_opt<1024, 32>(A, B, C); break;\n default: \n FatalError(\"In compute_values_opt the requested hash size is too large.\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n break;\n default:\n FatalError(\"compute_values_opt only implemented for group size = 16, 32\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n hipDeviceSynchronize();\n\n cudaCheckError();\n}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::compute_values( const Matrix_d &A, const Matrix_d &B, Matrix_d &C, int num_threads, IVector *Aq1, IVector *Bq1, IVector *Aq2, IVector *Bq2 )\n{\n const int GRID_SIZE = 1024;\n const int CTA_SIZE = 128;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Reset the work queue.\n int work_offset = GRID_SIZE * NUM_WARPS;\n CUDA_SAFE_CALL( hipMemcpy( this->m_work_queue, &work_offset, sizeof(int), hipMemcpyHostToDevice ) );\n // Compute the values.\n int *status = NULL;\n\n if ( num_threads != this->m_num_threads_per_row_compute )\n {\n status = this->m_status;\n }\n\n \t\t\n\n switch ( num_threads )\n {\n case 2:\n csr_multiply_detail::compute_values_kernel< 2, Value_type, CTA_SIZE, SMEM_SIZE, WARP_SIZE> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n B.values.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n C.values.raw(),\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_vals,\n this->m_work_queue,\n status );\n break;\n\n case 4:\n csr_multiply_detail::compute_values_kernel< 4, Value_type, CTA_SIZE, SMEM_SIZE, WARP_SIZE> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n B.values.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n C.values.raw(),\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_vals,\n this->m_work_queue,\n status );\n break;\n\n case 8:\n csr_multiply_detail::compute_values_kernel< 8, Value_type, CTA_SIZE, SMEM_SIZE, WARP_SIZE> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n B.values.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n C.values.raw(),\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_vals,\n this->m_work_queue,\n status );\n break;\n\n case 16:\n csr_multiply_detail::compute_values_kernel<16, Value_type, CTA_SIZE, SMEM_SIZE, WARP_SIZE> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n B.values.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n C.values.raw(),\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_vals,\n this->m_work_queue,\n status );\n break;\n\n default:\n csr_multiply_detail::compute_values_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n B.values.raw(),\n C.row_offsets.raw(),\n C.col_indices.raw(),\n C.values.raw(),\n (Aq1 != NULL) ? Aq1->raw() : NULL,\n (Bq1 != NULL) ? Bq1->raw() : NULL,\n (Aq2 != NULL) ? Aq2->raw() : NULL,\n (Bq2 != NULL) ? Bq2->raw() : NULL,\n this->m_gmem_size,\n this->m_keys,\n this->m_vals,\n this->m_work_queue,\n status );\n }\n\n cudaCheckError();\n //CUDA_SAFE_CALL( cudaGetLastError() );\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid CSR_Multiply_Detail >::compute_values_RAP_sparse_add( Matrix_d &RAP, const Matrix_d &RAP_int, std::vector &RAP_ext_row_offsets, std::vector &RAP_ext_col_indices, std::vector &RAP_ext_values, std::vector &RAP_ext_row_ids, int num_threads)\n{\n const int GRID_SIZE = 1024;\n const int CTA_SIZE = 128;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Reset the work queue.\n int work_offset = GRID_SIZE * NUM_WARPS;\n CUDA_SAFE_CALL( hipMemcpy( this->m_work_queue, &work_offset, sizeof(int), hipMemcpyHostToDevice ) );\n // Compute the values.\n int *status = NULL;\n\n if ( num_threads != this->m_num_threads_per_row_compute )\n {\n status = this->m_status;\n }\n\n // This is num_owned_coarse_rows\n int RAP_size = RAP.get_num_rows();\n int RAP_int_size = RAP_int.row_offsets.size() - 1;\n\n if (RAP_int_size < RAP_size)\n {\n FatalError(\"RAP_int has less rows than RAP, need to modify sparse RAP add to handle that case\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n //TODO: Optimize: reuse arrays from count nonzeros\n int num_neighbors = RAP_ext_row_offsets.size();\n std::vector flagArray(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n flagArray[i].resize(RAP_size);\n thrust_wrapper::fill(flagArray[i].begin(), flagArray[i].end(), -1);\n }\n\n cudaCheckError();\n std::vector flagArray_ptrs_h(num_neighbors);\n std::vector RAP_ext_row_offsets_ptrs_h(num_neighbors);\n std::vector RAP_ext_col_indices_ptrs_h(num_neighbors);\n std::vector RAP_ext_values_ptrs_h(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n flagArray_ptrs_h[i] = amgx::thrust::raw_pointer_cast(&flagArray[i][0]);\n RAP_ext_row_offsets_ptrs_h[i] = amgx::thrust::raw_pointer_cast(&RAP_ext_row_offsets[i][0]);\n RAP_ext_col_indices_ptrs_h[i] = amgx::thrust::raw_pointer_cast(&RAP_ext_col_indices[i][0]);\n RAP_ext_values_ptrs_h[i] = amgx::thrust::raw_pointer_cast(&RAP_ext_values[i][0]);\n }\n\n device_vector_alloc flagArray_ptrs = flagArray_ptrs_h;\n device_vector_alloc RAP_ext_row_offsets_ptrs = RAP_ext_row_offsets_ptrs_h;\n device_vector_alloc RAP_ext_col_indices_ptrs = RAP_ext_col_indices_ptrs_h;\n device_vector_alloc RAP_ext_values_ptrs = RAP_ext_values_ptrs_h;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = RAP_ext_row_ids[i].size();\n\n if (size != 0)\n {\n int num_blocks = std::min(4096, (size + 127) / 128);\n //write the position in RAP_ext_row_ids\n csr_multiply_detail::flag_halo_rows <<< num_blocks, 128>>>(\n RAP_ext_row_ids[i].raw(),\n size,\n flagArray[i].raw(),\n i,\n RAP.manager->global_id());\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n //CUDA_SAFE_CALL( cudaGetLastError() );\n csr_multiply_detail::compute_values_RAP_ext_kernel< Value_type, CTA_SIZE, SMEM_SIZE, WARP_SIZE> <<< GRID_SIZE, CTA_SIZE>>>(\n RAP_size,\n RAP_int.row_offsets.raw(),\n RAP_int.col_indices.raw(),\n RAP_int.values.raw(),\n amgx::thrust::raw_pointer_cast(&RAP_ext_row_offsets_ptrs[0]),\n amgx::thrust::raw_pointer_cast(&RAP_ext_col_indices_ptrs[0]),\n amgx::thrust::raw_pointer_cast(&RAP_ext_values_ptrs[0]),\n RAP.row_offsets.raw(),\n RAP.col_indices.raw(),\n RAP.values.raw(),\n amgx::thrust::raw_pointer_cast(&flagArray_ptrs[0]),\n this->m_gmem_size,\n this->m_keys,\n this->m_vals,\n this->m_work_queue,\n num_neighbors,\n status );\n cudaCheckError();\n //CUDA_SAFE_CALL( cudaGetLastError() );\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define AMGX_CASE_LINE(CASE) template class CSR_Multiply_Detail::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace amgx\n\n", "source": "/kaggle/working/amgx/src/csr_multiply_detail.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n\n#ifdef AMGX_USE_MAGMA\n#define ADD_ 1\n#define HAVE_CUBLAS 1\n#include \n#endif\n\n#include \n\nnamespace amgx\n{\n\n#define lapackCheckError(status) \\\n { \\\n if (status < 0) \\\n { \\\n std::stringstream ss; \\\n ss << \"Lapack error: argument number \" \\\n << -status << \" had an illegal value.\"; \\\n FatalError(ss.str(), AMGX_ERR_INTERNAL); \\\n } \\\n else if (status > 0) \\\n FatalError(\"Lapack error: internal error.\", \\\n AMGX_ERR_INTERNAL); \\\n } \\\n\n#define magmaCheckError(status) \\\n { \\\n if (status < 0) \\\n { \\\n std::stringstream ss; \\\n ss << \"Magma error: argument number \" \\\n << -status << \" had an illegal value.\"; \\\n FatalError(ss.str(), AMGX_ERR_INTERNAL); \\\n } \\\n else if (status > 0) \\\n FatalError(\"Magma error: internal error.\", \\\n AMGX_ERR_INTERNAL); \\\n } \\\n\ntemplate \nvoid Lapack< TemplateConfig >::check_lapack_enabled()\n{\n#ifndef AMGX_USE_LAPACK\n FatalError(\"Error: LAPACK not enabled.\", AMGX_ERR_CONFIGURATION);\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::check_magma_enabled()\n{\n#ifndef AMGX_USE_MAGMA\n FatalError(\"Error: MAGMA not enabled.\", AMGX_ERR_CONFIGURATION);\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::check_magma_enabled()\n{\n#ifndef AMGX_USE_MAGMA\n FatalError(\"Error: MAGMA not enabled.\", AMGX_ERR_CONFIGURATION);\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::not_implemented()\n{\n FatalError(\"Error: LAPACK operation not implemented on host.\", AMGX_ERR_CONFIGURATION);\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::not_implemented()\n{\n FatalError(\"Error: LAPACK operation not implemented on device.\", AMGX_ERR_CONFIGURATION);\n}\n\nnamespace\n{\n\n#ifdef AMGX_USE_LAPACK\nstruct _fcomplex { float re, im; };\ntypedef struct _fcomplex fcomplex;\nstruct _dcomplex { double re, im; };\ntypedef struct _dcomplex dcomplex;\n\nextern \"C\"\nint dgeev_(char *jobvl, char *jobvr, int *n, double *a,\n int *lda, double *wr, double *wi, double *vl,\n int *ldvl, double *vr, int *ldvr, double *work,\n int *lwork, int *info);\n\nextern \"C\"\nint sgeev_(char *jobvl, char *jobvr, int *n, float *a,\n int *lda, float *wr, float *wi, float *vl,\n int *ldvl, float *vr, int *ldvr, float *work,\n int *lwork, int *info);\n\nextern \"C\"\nint cgeev_(char *jobvl, char *jobvr, int *n, fcomplex *a,\n int *lda, fcomplex *wr, fcomplex *wi, fcomplex *vl,\n int *ldvl, fcomplex *vr, int *ldvr, fcomplex *work,\n int *lwork, int *info);\n\nextern \"C\"\nint zgeev_(char *jobvl, char *jobvr, int *n, dcomplex *a,\n int *lda, dcomplex *wr, dcomplex *wi, dcomplex *vl,\n int *ldvl, dcomplex *vr, int *ldvr, dcomplex *work,\n int *lwork, int *info);\n\nint lapack_geev_dispatch(char *jobvl, char *jobvr, int *n, double *a,\n int *lda, double *wr, double *wi, double *vl,\n int *ldvl, double *vr, int *ldvr, double *work,\n int *lwork, int *info)\n{\n return dgeev_(jobvl, jobvr, n, a, lda, wr, wi, vl, ldvl, vr, ldvr, work, lwork, info);\n}\n\nint lapack_geev_dispatch(char *jobvl, char *jobvr, int *n, float *a,\n int *lda, float *wr, float *wi, float *vl,\n int *ldvl, float *vr, int *ldvr, float *work,\n int *lwork, int *info)\n{\n return sgeev_(jobvl, jobvr, n, a, lda, wr, wi, vl, ldvl, vr, ldvr, work, lwork, info);\n}\n\nint lapack_geev_dispatch(char *jobvl, char *jobvr, int *n, cuComplex *a,\n int *lda, cuComplex *wr, cuComplex *wi, cuComplex *vl,\n int *ldvl, cuComplex *vr, int *ldvr, cuComplex *work,\n int *lwork, int *info)\n{\n return cgeev_(jobvl, jobvr, n, \n reinterpret_cast(a),\n lda, \n reinterpret_cast(wr), \n reinterpret_cast(wi), \n reinterpret_cast(vl),\n ldvl, \n reinterpret_cast(vr), \n ldvr, \n reinterpret_cast(work), \n lwork,\n info);\n}\n\nint lapack_geev_dispatch(char *jobvl, char *jobvr, int *n, cuDoubleComplex *a,\n int *lda, cuDoubleComplex *wr, cuDoubleComplex *wi, cuDoubleComplex *vl,\n int *ldvl, cuDoubleComplex *vr, int *ldvr, cuDoubleComplex *work,\n int *lwork, int *info)\n{\n return zgeev_(jobvl, jobvr, n, \n reinterpret_cast(a),\n lda, \n reinterpret_cast(wr), \n reinterpret_cast(wi), \n reinterpret_cast(vl),\n ldvl, \n reinterpret_cast(vr), \n ldvr, \n reinterpret_cast(work), \n lwork,\n info);\n}\n\ntemplate \nvoid lapack_geev(T *A, T *eigenvalues, int dim, int lda)\n{\n char job = 'N';\n T *WI = new T[dim];\n int ldv = 1;\n T *vl = 0;\n int work_size = 6 * dim;\n T *work = new T[work_size];\n int info;\n lapack_geev_dispatch(&job, &job, &dim, A, &lda, eigenvalues, WI, vl, &ldv,\n vl, &ldv, work, &work_size, &info);\n lapackCheckError(info);\n delete [] WI;\n delete [] work;\n}\n\ntemplate \nvoid lapack_geev(T *A, T *eigenvalues, T *eigenvectors, int dim, int lda, int ldvr)\n{\n char jobvl = 'N';\n char jobvr = 'V';\n T *WI = new T[dim * dim];\n int work_size = 6 * dim;\n T *vl = 0;\n int ldvl = 1;\n T *work = new T[work_size];\n int info;\n lapack_geev_dispatch(&jobvl, &jobvr, &dim, A, &lda, eigenvalues, WI, vl, &ldvl,\n eigenvectors, &ldvr, work, &work_size, &info);\n lapackCheckError(info);\n delete [] WI;\n delete [] work;\n}\n\n#endif\n\n} // end anonymous namespace\n\ntemplate \nvoid Lapack< TemplateConfig >::geev(const Vector &A, Vector &eigenvalues)\n{\n check_lapack_enabled();\n typedef typename Vector::value_type value_type;\n // It is possible the matrix has an extra row (e.g. Arnoldi).\n int dim = std::min(A.get_num_rows(), A.get_num_cols());\n int lda = A.get_lda();\n value_type *A_ptr = const_cast(A.raw());\n#ifdef AMGX_USE_LAPACK\n lapack_geev(A_ptr, eigenvalues.raw(), dim, lda);\n#else\n FatalError(\"Lapack is not supported in this build\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::geev(const Vector &A,\n Vector &eigenvalues,\n Vector &eigenvector)\n{\n check_lapack_enabled();\n typedef typename Vector::value_type value_type;\n // It is possible the matrix has an extra row (e.g. Arnoldi).\n int dim = std::min(A.get_num_rows(), A.get_num_cols());\n int lda = A.get_lda();\n value_type *A_ptr = const_cast(A.raw());\n#ifdef AMGX_USE_LAPACK\n lapack_geev(A_ptr, eigenvalues.raw(), eigenvector.raw(), dim, lda, eigenvector.get_lda());\n#else\n FatalError(\"Lapack is not supported in this build\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\n\ntemplate \nvoid Lapack< TemplateConfig >::geev(const Vector &A, Vector &eigenvalues)\n{\n not_implemented();\n}\n\nnamespace\n{\n#ifdef AMGX_USE_LAPACK\nextern \"C\"\nint dtrtri_(char *uplo, char *diag, int *n, double *\n a, int *lda, int *info);\n\nextern \"C\"\nint strtri_(char *uplo, char *diag, int *n, float *\n a, int *lda, int *info);\n\nextern \"C\"\nint ctrtri_(char *uplo, char *diag, int *n, fcomplex *\n a, int *lda, int *info);\n\nextern \"C\"\nint ztrtri_(char *uplo, char *diag, int *n, dcomplex *\n a, int *lda, int *info);\n\nint lapack_trtri_dispatch(char *uplo, char *diag, int *n, float *a,\n int *lda, int *info)\n{\n return strtri_(uplo, diag, n, a, lda, info);\n}\n\nint lapack_trtri_dispatch(char *uplo, char *diag, int *n, double *a,\n int *lda, int *info)\n{\n return dtrtri_(uplo, diag, n, a, lda, info);\n}\n\nint lapack_trtri_dispatch(char *uplo, char *diag, int *n, fcomplex *a,\n int *lda, int *info)\n{\n return ctrtri_(uplo, diag, n, a, lda, info);\n}\n\nint lapack_trtri_dispatch(char *uplo, char *diag, int *n, dcomplex *a,\n int *lda, int *info)\n{\n return ztrtri_(uplo, diag, n, a, lda, info);\n}\n\ntemplate \nvoid lapack_trtri(T *A, int dim, int lda)\n{\n char uplo = 'U';\n char diag = 'N';\n int info;\n lapack_trtri_dispatch(&uplo, &diag, &dim, A, &lda, &info);\n lapackCheckError(info);\n}\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::trtri(Vector &A)\n{\n check_lapack_enabled();\n typedef typename Vector::value_type value_type;\n int dim = std::min(A.get_num_rows(), A.get_num_cols());\n int lda = A.get_lda();\n#ifdef AMGX_USE_LAPACK\n lapack_trtri(A.raw(), dim, lda);\n#else\n FatalError(\"Lapack is not supported in this build\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\nnamespace\n{\n#ifdef AMGX_USE_MAGMA\nint magma_trtri_dispatch(magma_uplo_t uplo, magma_diag_t diag, int n, float *a,\n int lda, int *info)\n{\n return magma_strtri_gpu(uplo, diag, n, a, lda, info);\n}\n\nint magma_trtri_dispatch(magma_uplo_t uplo, magma_diag_t diag, int n, double *a,\n int lda, int *info)\n{\n return magma_dtrtri_gpu(uplo, diag, n, a, lda, info);\n}\n\nint magma_trtri_dispatch(magma_uplo_t uplo, magma_diag_t diag, int n, cuComplex *a,\n int lda, int *info)\n{\n return magma_ctrtri_gpu(uplo, diag, n, a, lda, info);\n}\n\nint magma_trtri_dispatch(magma_uplo_t uplo, magma_diag_t diag, int n, cuDoubleComplex *a,\n int lda, int *info)\n{\n return magma_ztrtri_gpu(uplo, diag, n, a, lda, info);\n}\n\ntemplate \nvoid magma_trtri(T *A, int dim, int lda)\n{\n magma_uplo_t uplo = MagmaUpper;\n magma_diag_t diag = MagmaNonUnit;\n int info;\n magma_trtri_dispatch(uplo, diag, dim, A, lda, &info);\n magmaCheckError(info);\n}\n\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::trtri(Vector &A)\n{\n check_magma_enabled();\n typedef typename Vector::value_type value_type;\n int dim = std::min(A.get_num_rows(), A.get_num_cols());\n int lda = A.get_lda();\n#ifdef AMGX_USE_MAGMA\n magma_trtri(A.raw(), dim, lda);;\n#else\n FatalError(\"Lapack is not supported in this build\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\nnamespace\n{\n#ifdef AMGX_USE_LAPACK\nextern \"C\"\nint dsygv_(int *itype, char *jobz, char *uplo, int *n,\n double *a, int *lda, double *b, int *ldb,\n double *w, double *work, int *lwork, int *info);\n\nextern \"C\"\nint ssygv_(int *itype, char *jobz, char *uplo, int *n,\n float *a, int *lda, float *b, int *ldb,\n float *w, float *work, int *lwork, int *info);\n\nextern \"C\"\nint chegv_(int *itype, char *jobz, char *uplo, int *n,\n fcomplex *a, int *lda, fcomplex *b, int *ldb,\n fcomplex *w, fcomplex *work, int *lwork, int *info);\n\nextern \"C\"\nint zhegv_(int *itype, char *jobz, char *uplo, int *n,\n dcomplex *a, int *lda, dcomplex *b, int *ldb,\n dcomplex *w, dcomplex *work, int *lwork, int *info);\n\nint lapack_sygv_dispatch(int *itype, char *jobz, char *uplo, int *n,\n double *a, int *lda, double *b, int *ldb,\n double *w, double *work, int *lwork, int *info)\n{\n return dsygv_(itype, jobz, uplo, n, a, lda, b, ldb, w, work, lwork, info);\n}\n\nint lapack_sygv_dispatch(int *itype, char *jobz, char *uplo, int *n,\n float *a, int *lda, float *b, int *ldb,\n float *w, float *work, int *lwork, int *info)\n{\n return ssygv_(itype, jobz, uplo, n, a, lda, b, ldb, w, work, lwork, info);\n}\n\nint lapack_sygv_dispatch(int *itype, char *jobz, char *uplo, int *n,\n cuComplex *a, int *lda, cuComplex *b, int *ldb,\n cuComplex *w, cuComplex *work, int *lwork, int *info)\n{\n return chegv_(itype, jobz, uplo, n,\n reinterpret_cast(a), lda, reinterpret_cast(b), ldb,\n reinterpret_cast(w), reinterpret_cast(work), lwork, info);\n}\n\nint lapack_sygv_dispatch(int *itype, char *jobz, char *uplo, int *n,\n cuComplex *a, int *lda, cuDoubleComplex *b, int *ldb,\n cuDoubleComplex *w, cuDoubleComplex *work, int *lwork, int *info)\n{\n return zhegv_(itype, jobz, uplo, n,\n reinterpret_cast(a), lda, reinterpret_cast(b), ldb,\n reinterpret_cast(w), reinterpret_cast(work), lwork, info);\n}\n\ntemplate \nvoid lapack_sygv(T *gramA, T *gramB, T *eigenvector, int dim, int lda, T *work)\n{\n int itype = 1;\n char jobz = 'V';\n char uplo = 'U';\n int ldb = lda;\n int lwork = 1024;\n int info = 0;\n lapack_sygv_dispatch(&itype, &jobz, &uplo, &dim, gramA, &lda, gramB, &ldb, eigenvector, work, &lwork, &info);\n lapackCheckError(info);\n}\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::sygv(Vector &A, Vector &B,\n Vector &eigenvalues, Vector &work)\n{\n check_lapack_enabled();\n typedef typename Vector::value_type value_type;\n int dim = std::min(A.get_num_rows(), A.get_num_cols());\n int lda = A.get_lda();\n#ifdef AMGX_USE_LAPACK\n lapack_sygv(A.raw(), B.raw(), eigenvalues.raw(), dim, lda, work.raw());\n#else\n FatalError(\"Lapack is not supported in this build\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\nnamespace\n{\n#ifdef AMGX_USE_MAGMA\nvoid magma_trsm_dispatch(magma_side_t side, magma_uplo_t uplo, magma_trans_t trans,\n magma_diag_t diag, magma_int_t m, magma_int_t n,\n float alpha, float const *dA, magma_int_t lda,\n float *dB, magma_int_t ldb)\n{\n return magma_strsm(side, uplo, trans, diag, m, n, alpha, dA, lda, dB, ldb);\n}\n\nvoid magma_trsm_dispatch(magma_side_t side, magma_uplo_t uplo, magma_trans_t trans,\n magma_diag_t diag, magma_int_t m, magma_int_t n,\n double alpha, double const *dA, magma_int_t lda,\n double *dB, magma_int_t ldb)\n{\n return magma_dtrsm(side, uplo, trans, diag, m, n, alpha, dA, lda, dB, ldb);\n}\n\nvoid magma_trmm_dispatch(magma_side_t side, magma_uplo_t uplo, magma_trans_t trans,\n magma_diag_t diag, magma_int_t m, magma_int_t n,\n float alpha, float const *dA, magma_int_t lda,\n float *dB, magma_int_t ldb)\n{\n return magma_strmm(side, uplo, trans, diag, m, n, alpha, dA, lda, dB, ldb);\n}\n\nvoid magma_trmm_dispatch(magma_side_t side, magma_uplo_t uplo, magma_trans_t trans,\n magma_diag_t diag, magma_int_t m, magma_int_t n,\n double alpha, double const *dA, magma_int_t lda,\n double *dB, magma_int_t ldb)\n{\n return magma_dtrmm(side, uplo, trans, diag, m, n, alpha, dA, lda, dB, ldb);\n}\n\nint magma_potrf_gpu_dispatch(magma_uplo_t uplo, int n, float *A, int lda, int *info)\n{\n return magma_spotrf_gpu(uplo, n, A, lda, info);\n}\n\nint magma_potrf_gpu_dispatch(magma_uplo_t uplo, int n, double *A, int lda, int *info)\n{\n return magma_dpotrf_gpu(uplo, n, A, lda, info);\n}\n\nint magma_sygst_gpu_dispatch(int itype, magma_uplo_t uplo, magma_int_t n, float *da,\n int ldda, float *B, int lddb, int *info)\n{\n return magma_ssygst_gpu(itype, uplo, n, da, ldda, B, lddb, info);\n}\n\nint magma_sygst_gpu_dispatch(int itype, magma_uplo_t uplo, magma_int_t n, double *da,\n int ldda, double *B, int lddb, int *info)\n{\n return magma_dsygst_gpu(itype, uplo, n, da, ldda, B, lddb, info);\n}\n\nint magma_syevd_gpu_dispatch(magma_vec_t jobz, magma_uplo_t uplo, int n, double *da, int ldda,\n double *w, double *wa, int ldwa, double *work,\n int lwork, int *iwork, int liwork, int *info)\n{\n return magma_dsyevd_gpu(jobz, uplo, n, da, ldda, w, wa, ldwa, work, lwork, iwork, liwork, info);\n}\n\nint magma_syevd_gpu_dispatch(magma_vec_t jobz, magma_uplo_t uplo, int n, float *da, int ldda,\n float *w, float *wa, int ldwa, float *work,\n int lwork, int *iwork, int liwork, int *info)\n{\n return magma_ssyevd_gpu(jobz, uplo, n, da, ldda, w, wa, ldwa, work, lwork, iwork, liwork, info);\n}\n\n// This is a simple modification of the magma_?sygvd() source code\n// from magma where the matrices are already on the device.\ntemplate \nmagma_int_t magma_sygvd_gpu_impl(magma_int_t itype, magma_vec_t jobz, magma_uplo_t uplo, magma_int_t n,\n T *da, magma_int_t lda, T *db, magma_int_t ldb, T *w,\n T *work, magma_int_t lwork, magma_int_t *iwork, magma_int_t liwork,\n T *wa, magma_int_t *info)\n{\n magma_uplo_t uplo_[2] = {uplo, MagmaLower}; // {uplo, 0}\n magma_vec_t jobz_[2] = {jobz, MagmaVec};//{jobz, 0};\n T d_one = MAGMA_D_ONE;\n magma_int_t ldda = n;\n magma_int_t lddb = n;\n static magma_int_t lower;\n static char trans[1];\n static magma_int_t wantz, lquery;\n static magma_int_t lopt, lwmin, liopt, liwmin;\n static cudaStream_t stream;\n magma_queue_create( &stream );\n wantz = jobz_[0] == MagmaVec;\n lower = uplo_[0] == MagmaLower;\n lquery = lwork == -1 || liwork == -1;\n *info = 0;\n\n if (itype < 1 || itype > 3)\n {\n *info = -1;\n }\n else if (! (wantz || jobz_[0] == MagmaNoVec))\n {\n *info = -2;\n }\n else if (! (lower || uplo_[0] == MagmaUpper))\n {\n *info = -3;\n }\n else if (n < 0)\n {\n *info = -4;\n }\n else if (lda < max(1, n))\n {\n *info = -6;\n }\n else if (ldb < max(1, n))\n {\n *info = -8;\n }\n\n magma_int_t nb = magma_get_dsytrd_nb(n);\n\n if (n < 1)\n {\n liwmin = 1;\n lwmin = 1;\n }\n else if (wantz)\n {\n lwmin = 1 + 6 * n * nb + 2 * n * n;\n liwmin = 5 * n + 3;\n }\n else\n {\n lwmin = 2 * n * nb + 1;\n liwmin = 1;\n }\n\n lopt = lwmin;\n liopt = liwmin;\n work[ 0] = lopt;\n iwork[0] = liopt;\n\n if (lwork < lwmin && ! lquery)\n {\n *info = -11;\n }\n else if (liwork < liwmin && ! lquery)\n {\n *info = -13;\n }\n\n if (*info != 0)\n {\n magma_xerbla( __func__, -(*info) );\n return MAGMA_ERR_ILLEGAL_VALUE;\n }\n else if (lquery)\n {\n return MAGMA_SUCCESS;\n }\n\n /* Quick return if possible */\n if (n == 0)\n {\n return 0;\n }\n\n magma_potrf_gpu_dispatch(uplo_[0], n, db, lddb, info);\n\n if (*info != 0)\n {\n *info = n + *info;\n return 0;\n }\n\n /* Transform problem to standard eigenvalue problem and solve. */\n magma_sygst_gpu_dispatch(itype, uplo_[0], n, da, ldda, db, lddb, info);\n magma_syevd_gpu_dispatch(jobz_[0], uplo_[0], n, da, ldda, w, wa, lda,\n work, lwork, iwork, liwork, info);\n lopt = max( lopt, (magma_int_t) work[0]);\n liopt = max(liopt, iwork[0]);\n\n if (wantz && *info == 0)\n {\n /* Backtransform eigenvectors to the original problem. */\n if (itype == 1 || itype == 2)\n {\n /* For A*x=(lambda)*B*x and A*B*x=(lambda)*x;\n backtransform eigenvectors: x = inv(L)'*y or inv(U)*y */\n if (lower)\n {\n *(unsigned char *)trans = MagmaTrans;\n }\n else\n {\n *(unsigned char *)trans = MagmaNoTrans;\n }\n\n magma_trsm_dispatch(MagmaLeft, uplo_[0], *trans, MagmaNonUnit,\n n, n, d_one, db, lddb, da, ldda);\n }\n else if (itype == 3)\n {\n /* For B*A*x=(lambda)*x;\n backtransform eigenvectors: x = L*y or U'*y */\n if (lower)\n {\n *(unsigned char *)trans = MagmaNoTrans;\n }\n else\n {\n *(unsigned char *)trans = MagmaTrans;\n }\n\n magma_trmm_dispatch(MagmaLeft, uplo_[0], *trans, MagmaNonUnit,\n n, n, d_one, db, lddb, da, ldda);\n }\n }\n\n magma_queue_sync( stream );\n magma_queue_destroy( stream );\n work[0] = (T) lopt;\n iwork[0] = liopt;\n return MAGMA_SUCCESS;\n}\n\ncublasStatus_t cublas_trsm(cublasHandle_t handle,\n cublasSideMode_t side, cublasFillMode_t uplo,\n cublasOperation_t trans, cublasDiagType_t diag,\n int m, int n,\n const float *alpha,\n const float *A, int lda,\n float *B, int ldb)\n{\n return cublasStrsm(handle, side, uplo, trans, diag, m, n, alpha, A, lda, B, ldb);\n}\n\ncublasStatus_t cublas_trsm(cublasHandle_t handle,\n cublasSideMode_t side, cublasFillMode_t uplo,\n cublasOperation_t trans, cublasDiagType_t diag,\n int m, int n,\n const double *alpha,\n const double *A, int lda,\n double *B, int ldb)\n{\n return cublasDtrsm(handle, side, uplo, trans, diag, m, n, alpha, A, lda, B, ldb);\n}\n\n\ntemplate \nvoid magma_sygvd_gpu(T *A, T *B, T *eigenvalues, int dim, int lda)\n{\n int itype = 1;\n magma_vec_t jobz = MagmaVec;\n magma_uplo_t uplo = MagmaUpper;\n int N = dim;\n int ldb = lda;\n int nb = 32;\n int lwork = 1 + 6 * N * nb + 2 * N * N;\n static std::vector s_work;\n s_work.resize(lwork);\n T *work = &s_work[0];\n int liwork = 3 + 5 * N;\n static std::vector s_iwork;\n s_iwork.resize(liwork);\n int *iwork = &s_iwork[0];\n static std::vector s_wa;\n s_wa.resize(lda * N);\n T *wa = &s_wa[0];\n int ldwa = N;\n int info;\n /*\n magma_sygvd_gpu_impl(itype, jobz, uplo, N, A, lda, B, ldb, eigenvalues, work, lwork, iwork, liwork, wa, &info);\n */\n magma_potrf_gpu_dispatch(uplo, N, B, lda, &info);\n magmaCheckError(info);\n magma_sygst_gpu_dispatch(itype, uplo, N, A, lda, B, ldb, &info);\n magmaCheckError(info);\n magma_syevd_gpu_dispatch(jobz, uplo, N, A, lda, eigenvalues, wa, ldwa, work, lwork, iwork, liwork, &info);\n magmaCheckError(info);\n T one = 1;\n cublasHandle_t handle = Cublas::get_handle();\n cublas_trsm(handle, CUBLAS_SIDE_LEFT, CUBLAS_FILL_MODE_UPPER, CUBLAS_OP_N, CUBLAS_DIAG_NON_UNIT, N, N, &one, B, ldb, A, lda);\n}\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::sygv(Vector &A, Vector &B,\n Vector &eigenvalues, Vector &work)\n{\n typedef typename Vector::value_type value_type;\n int dim = std::min(A.get_num_rows(), A.get_num_cols());\n int lda = A.get_lda();\n#ifdef AMGX_USE_MAGMA\n magma_sygvd_gpu(A.raw(), B.raw(), eigenvalues.raw(), dim, lda);\n#endif\n}\n\nnamespace\n{\n#ifdef AMGX_USE_MAGMA\ntemplate \nvoid magma_syevd_gpu(T *A, T *eigenvalues, int dim, int lda)\n{\n magma_vec_t jobz = MagmaVec;\n magma_uplo_t uplo = MagmaUpper;\n int N = dim;\n int nb = 32;\n int lwork = 1 + 6 * N * nb + 2 * N * N;\n static std::vector s_work;\n s_work.resize(lwork);\n T *work = &s_work[0];\n int liwork = 3 + 5 * N;\n static std::vector s_iwork;\n s_iwork.resize(liwork);\n int *iwork = &s_iwork[0];\n static std::vector s_wa;\n s_wa.resize(lda * N);\n T *wa = &s_wa[0];\n int ldwa = N;\n int info;\n magma_syevd_gpu_dispatch(jobz, uplo, N, A, lda, eigenvalues, wa, ldwa, work, lwork, iwork, liwork, &info);\n magmaCheckError(info);\n}\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::syevd(Vector &A,\n Vector &eigenvalues)\n{\n check_magma_enabled();\n typedef typename Vector::value_type value_type;\n int dim = std::min(A.get_num_rows(), A.get_num_cols());\n int lda = A.get_lda();\n#ifdef AMGX_USE_MAGMA\n magma_syevd_gpu(A.raw(), eigenvalues.raw(), dim, lda);\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::syevd(Vector &A,\n Vector &eigenvalues)\n{\n not_implemented();\n}\n\nnamespace\n{\n#ifdef AMGX_USE_MAGMA\nint magma_stedx_dispatch(magma_range_t range, int n,\n double vl, double vu,\n int il, int iu,\n double *d, double *e, double *z, int ldz,\n double *work, int lwork, int *iwork, int liwork,\n double *dwork, int *info)\n{\n return magma_dstedx(range, n, vl, vu, il, iu, d, e, z, ldz, work, lwork, iwork, liwork, dwork, info);\n}\n\nint magma_stedx_dispatch(magma_range_t range, int n,\n float vl, float vu,\n int il, int iu,\n float *d, float *e, float *z, int ldz,\n float *work, int lwork, int *iwork, int liwork,\n float *dwork, int *info)\n{\n return magma_sstedx(range, n, vl, vu, il, iu, d, e, z, ldz, work, lwork, iwork, liwork, dwork, info);\n}\n\ntemplate \nvoid magma_stedx(T *diagonal, T *subdiagonal, T *eigenvectors,\n int lower, int upper, int dim, int ldz, T *dwork, int dwork_size)\n{\n magma_range_t range = MagmaRangeI;\n int N = dim;\n T vl = 0;\n T vu = 0;\n int il = lower;\n int iu = upper;\n int lwork = 1 + 4 * N + 2 * N * N;\n static std::vector s_work;\n s_work.resize(lwork);\n int liwork = 3 + 6 * N;\n static std::vector s_iwork;\n s_iwork.resize(liwork);\n int info;\n magma_stedx_dispatch(range, N, vl, vu, il, iu, diagonal, subdiagonal, eigenvectors, ldz,\n &s_work[0], lwork, &s_iwork[0], liwork, dwork, &info);\n}\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::stedx(Vector &diagonal,\n Vector &subdiagonal,\n Vector &eigenvectors,\n int dim,\n Vector &dwork)\n{\n check_magma_enabled();\n#ifdef AMGX_USE_MAGMA\n magma_stedx(diagonal.raw(), subdiagonal.raw(), eigenvectors.raw(),\n dim, dim, dim, eigenvectors.get_lda(),\n dwork.raw(), dwork.size());\n#endif\n}\n\n\nnamespace\n{\n\ntemplate \nvoid larf(int m, int n, T *v,\n int incv, T *tau, T *c, int ldc,\n T *work)\n{\n /* Table of constant values */\n static T c_b4 = 1.;\n static T c_b5 = 0.;\n static int c1 = 1;\n /* Form H * C */\n /* w := C' * v */\n Cublas::gemv(true, m, n, &c_b4, c, ldc,\n v, incv, &c_b5, work, c1);\n /* C := C - v * w' */\n Cublas::ger(m, n, tau, v, incv, work, c1, c, ldc);\n}\n\ntemplate \n__global__\nvoid set1(T *a)\n{\n *a = 1.;\n}\n\ntemplate \n__global__\nvoid add_tau(T *a, T tau)\n{\n *a = 1 + tau;\n}\n\ntemplate \nvoid gpu_orgqr(int m, int n, int k,\n T *a, int lda, T *tau, T *work, int lwork)\n{\n int i1, i2;\n\n for (int i = k - 1; i >= 0; --i)\n {\n /* Apply H(i) to A(i:m,i:n) from the left */\n if (i < n - 1)\n {\n set1 <<< 1, 1>>>(&a[i + i * lda]);\n cudaCheckError();\n i1 = m - i;\n i2 = n - i - 1;\n larf(i1, i2, &a[i + i * lda], 1, &tau[i],\n &a[i + (i + 1) * lda], lda, work);\n }\n\n if (i < m - 1)\n {\n i1 = m - i - 1;\n Cublas::scal(i1, &tau[i], &a[i + 1 + i * lda], 1);\n }\n\n add_tau <<< 1, 1>>>(&a[i + i * lda], tau[i]);\n cudaCheckError();\n /* Set A(1:i-1,i) to zero */\n cudaError_t cuda_rc = cudaMemset(&a[i * lda], 0, sizeof(T) * i);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemset failed in geqr2\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \n__device__ __host__\nT lapy2_(T *a, T *b)\n{\n T va = *a;\n T vb = *b;\n return sqrt(va * va + vb * vb);\n}\n\ntemplate \n__device__ __host__\nT d_sign(T a, T b)\n{\n T x;\n x = (a >= 0 ? a : - a);\n return (b >= 0 ? x : -x);\n}\n\ntemplate \nvoid compute_tau_host(T *alpha, T *norm,\n T *tau, T *d1)\n{\n *d1 = lapy2_(alpha, norm);\n T beta = -d_sign(*d1, *alpha);\n // LAPACK: skipped part about scaling.\n // Negated compared to LAPACK code, avoid negating value on device later.\n *tau = -(beta - *alpha) / beta;\n *d1 = 1. / (*alpha - beta);\n *alpha = beta;\n}\n\ntemplate \nvoid larfg(int n, T *alpha, T *x,\n int incx, T *tau)\n{\n if (n <= 1)\n {\n *tau = 0.;\n return;\n }\n\n int i1 = n - 1;\n T xnorm;\n Cublas::nrm2(i1, x, incx, &xnorm);\n T h_alpha;\n cudaError_t cuda_rc = cudaMemcpy(&h_alpha, alpha, sizeof(T), cudaMemcpyDeviceToHost);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy D2H failed in larfg\", AMGX_ERR_CUDA_FAILURE);\n }\n T d1;\n compute_tau_host(&h_alpha, &xnorm, tau, &d1);\n Cublas::scal(i1, d1, x, incx);\n // Update the diagonal value on the device.\n cuda_rc = cudaMemcpy(alpha, &h_alpha, sizeof(T), cudaMemcpyHostToDevice);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy H2D failed in larfg\", AMGX_ERR_CUDA_FAILURE);\n }\n}\n\ntemplate \nvoid gpu_geqrf(int m, int n, T *a, int lda,\n T *tau, T *work)\n{\n int k = std::min(m, n);\n T *aii;\n cudaError_t cuda_rc = amgx::memory::cudaMallocAsync((void**)&aii, sizeof(T));\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMallocAsync failed in gpu_geqrf\", AMGX_ERR_CUDA_FAILURE);\n }\n\n for (int i = 0; i < k; ++i)\n {\n /* Generate elementary reflector H(i) to annihilate A(i+1:m,i) */\n int i2 = m - i;\n /* Computing MIN */\n int i3 = i + 1;\n larfg(i2, &a[i + i * lda],\n &a[std::min(i3, m - 1) + i * lda],\n 1, &tau[i]);\n\n if (i < n - 1)\n {\n /* Apply H(i) to A(i:m,i+1:n) from the left */\n cuda_rc = cudaMemcpy(aii, &a[i + i * lda], sizeof(T), cudaMemcpyDeviceToDevice);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy D2D failed in gpu_geqrf\", AMGX_ERR_CUDA_FAILURE);\n }\n set1 <<< 1, 1>>>(&a[i + i * lda]);\n cudaCheckError();\n i2 = m - i;\n i3 = n - i - 1;\n larf(i2, i3, &a[i + i * lda], 1,\n &tau[i], &a[i + (i + 1) * lda], lda, work);\n cuda_rc = cudaMemcpy(&a[i + i * lda], aii, sizeof(T), cudaMemcpyDeviceToDevice);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy D2D failed in gpu_geqrf\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n }\n\n cuda_rc = amgx::memory::cudaFreeAsync(aii);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaFreeAsync failed in gpu_geqrf\", AMGX_ERR_CUDA_FAILURE);\n }\n}\n} // end anonymous namespace\n\nnamespace\n{\n#ifdef AMGX_USE_MAGMA\nint magma_geqrf_dispatch(int m, int n, float *A, int lda,\n float *tau, float *work, int *info)\n{\n return magma_sgeqrf_gpu(m, n, A, lda, tau, work, info);\n}\n\nint magma_geqrf_dispatch(int m, int n, double *A, int lda,\n double *tau, double *work, int *info)\n{\n return magma_dgeqrf_gpu(m, n, A, lda, tau, work, info);\n}\n\ntemplate \nvoid magma_geqrf(int m, int n, T *A, int lda,\n T *tau, T *work)\n{\n int info;\n magma_geqrf_dispatch(m, n, A, lda, tau, work, &info);\n magmaCheckError(info);\n}\n\nint magma_orgqr_dispatch(int m, int n, int k, float *A, int lda,\n float *tau, float *work, int lwork, int *info)\n{\n return magma_sorgqr_gpu(m, n, k, A, lda, tau, work, lwork, info);\n}\n\nint magma_orgqr_dispatch(int m, int n, int k, double *A, int lda,\n double *tau, double *work, int lwork, int *info)\n{\n return magma_dorgqr_gpu(m, n, k, A, lda, tau, work, lwork, info);\n}\n\ntemplate \nvoid magma_orgqr(int m, int n, int k, T *A, int lda,\n T *tau, T *work, int lwork)\n{\n int info;\n magma_orgqr_dispatch(m, n, k, A, lda, tau, work, lwork, &info);\n magmaCheckError(info);\n}\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::geqrf(Vector &A,\n Vector &tau,\n Vector &work)\n{\n not_implemented();\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::geqrf(Vector &A,\n Vector &tau,\n Vector &work)\n{\n int rows = A.get_num_rows();\n int cols = A.get_num_cols();\n int lda = A.get_lda();\n#ifdef AMGX_USE_MAGMA\n magma_geqrf(rows, cols, A.raw(), lda, tau.raw(), work.raw());\n#else\n gpu_geqrf(rows, cols, A.raw(), lda, tau.raw(), work.raw());\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::orgqr(Vector &A,\n Vector &tau,\n Vector &work)\n{\n not_implemented();\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::orgqr(Vector &A,\n Vector &tau,\n Vector &work)\n{\n int rows = A.get_num_rows();\n int cols = A.get_num_cols();\n int lda = A.get_lda();\n#ifdef AMGX_USE_MAGMA\n magma_orgqr(rows, cols, cols, A.raw(), lda, tau.raw(), work.raw(), 1);\n#else\n gpu_orgqr(rows, cols, cols, A.raw(), lda, tau.raw(), work.raw(), 1);\n#endif\n}\n\n\n#define AMGX_CASE_LINE(CASE) \\\n template class Lapack::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n#ifdef AMGX_USE_MAGMA\n#define ADD_ 1\n#define HAVE_CUBLAS 1\n#include \n#endif\n\n#include \n\nnamespace amgx\n{\n\n#define lapackCheckError(status) \\\n { \\\n if (status < 0) \\\n { \\\n std::stringstream ss; \\\n ss << \"Lapack error: argument number \" \\\n << -status << \" had an illegal value.\"; \\\n FatalError(ss.str(), AMGX_ERR_INTERNAL); \\\n } \\\n else if (status > 0) \\\n FatalError(\"Lapack error: internal error.\", \\\n AMGX_ERR_INTERNAL); \\\n } \\\n\n#define magmaCheckError(status) \\\n { \\\n if (status < 0) \\\n { \\\n std::stringstream ss; \\\n ss << \"Magma error: argument number \" \\\n << -status << \" had an illegal value.\"; \\\n FatalError(ss.str(), AMGX_ERR_INTERNAL); \\\n } \\\n else if (status > 0) \\\n FatalError(\"Magma error: internal error.\", \\\n AMGX_ERR_INTERNAL); \\\n } \\\n\ntemplate \nvoid Lapack< TemplateConfig >::check_lapack_enabled()\n{\n#ifndef AMGX_USE_LAPACK\n FatalError(\"Error: LAPACK not enabled.\", AMGX_ERR_CONFIGURATION);\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::check_magma_enabled()\n{\n#ifndef AMGX_USE_MAGMA\n FatalError(\"Error: MAGMA not enabled.\", AMGX_ERR_CONFIGURATION);\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::check_magma_enabled()\n{\n#ifndef AMGX_USE_MAGMA\n FatalError(\"Error: MAGMA not enabled.\", AMGX_ERR_CONFIGURATION);\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::not_implemented()\n{\n FatalError(\"Error: LAPACK operation not implemented on host.\", AMGX_ERR_CONFIGURATION);\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::not_implemented()\n{\n FatalError(\"Error: LAPACK operation not implemented on device.\", AMGX_ERR_CONFIGURATION);\n}\n\nnamespace\n{\n\n#ifdef AMGX_USE_LAPACK\nstruct _fcomplex { float re, im; };\ntypedef struct _fcomplex fcomplex;\nstruct _dcomplex { double re, im; };\ntypedef struct _dcomplex dcomplex;\n\nextern \"C\"\nint dgeev_(char *jobvl, char *jobvr, int *n, double *a,\n int *lda, double *wr, double *wi, double *vl,\n int *ldvl, double *vr, int *ldvr, double *work,\n int *lwork, int *info);\n\nextern \"C\"\nint sgeev_(char *jobvl, char *jobvr, int *n, float *a,\n int *lda, float *wr, float *wi, float *vl,\n int *ldvl, float *vr, int *ldvr, float *work,\n int *lwork, int *info);\n\nextern \"C\"\nint cgeev_(char *jobvl, char *jobvr, int *n, fcomplex *a,\n int *lda, fcomplex *wr, fcomplex *wi, fcomplex *vl,\n int *ldvl, fcomplex *vr, int *ldvr, fcomplex *work,\n int *lwork, int *info);\n\nextern \"C\"\nint zgeev_(char *jobvl, char *jobvr, int *n, dcomplex *a,\n int *lda, dcomplex *wr, dcomplex *wi, dcomplex *vl,\n int *ldvl, dcomplex *vr, int *ldvr, dcomplex *work,\n int *lwork, int *info);\n\nint lapack_geev_dispatch(char *jobvl, char *jobvr, int *n, double *a,\n int *lda, double *wr, double *wi, double *vl,\n int *ldvl, double *vr, int *ldvr, double *work,\n int *lwork, int *info)\n{\n return dgeev_(jobvl, jobvr, n, a, lda, wr, wi, vl, ldvl, vr, ldvr, work, lwork, info);\n}\n\nint lapack_geev_dispatch(char *jobvl, char *jobvr, int *n, float *a,\n int *lda, float *wr, float *wi, float *vl,\n int *ldvl, float *vr, int *ldvr, float *work,\n int *lwork, int *info)\n{\n return sgeev_(jobvl, jobvr, n, a, lda, wr, wi, vl, ldvl, vr, ldvr, work, lwork, info);\n}\n\nint lapack_geev_dispatch(char *jobvl, char *jobvr, int *n, hipComplex *a,\n int *lda, hipComplex *wr, hipComplex *wi, hipComplex *vl,\n int *ldvl, hipComplex *vr, int *ldvr, hipComplex *work,\n int *lwork, int *info)\n{\n return cgeev_(jobvl, jobvr, n, \n reinterpret_cast(a),\n lda, \n reinterpret_cast(wr), \n reinterpret_cast(wi), \n reinterpret_cast(vl),\n ldvl, \n reinterpret_cast(vr), \n ldvr, \n reinterpret_cast(work), \n lwork,\n info);\n}\n\nint lapack_geev_dispatch(char *jobvl, char *jobvr, int *n, hipDoubleComplex *a,\n int *lda, hipDoubleComplex *wr, hipDoubleComplex *wi, hipDoubleComplex *vl,\n int *ldvl, hipDoubleComplex *vr, int *ldvr, hipDoubleComplex *work,\n int *lwork, int *info)\n{\n return zgeev_(jobvl, jobvr, n, \n reinterpret_cast(a),\n lda, \n reinterpret_cast(wr), \n reinterpret_cast(wi), \n reinterpret_cast(vl),\n ldvl, \n reinterpret_cast(vr), \n ldvr, \n reinterpret_cast(work), \n lwork,\n info);\n}\n\ntemplate \nvoid lapack_geev(T *A, T *eigenvalues, int dim, int lda)\n{\n char job = 'N';\n T *WI = new T[dim];\n int ldv = 1;\n T *vl = 0;\n int work_size = 6 * dim;\n T *work = new T[work_size];\n int info;\n lapack_geev_dispatch(&job, &job, &dim, A, &lda, eigenvalues, WI, vl, &ldv,\n vl, &ldv, work, &work_size, &info);\n lapackCheckError(info);\n delete [] WI;\n delete [] work;\n}\n\ntemplate \nvoid lapack_geev(T *A, T *eigenvalues, T *eigenvectors, int dim, int lda, int ldvr)\n{\n char jobvl = 'N';\n char jobvr = 'V';\n T *WI = new T[dim * dim];\n int work_size = 6 * dim;\n T *vl = 0;\n int ldvl = 1;\n T *work = new T[work_size];\n int info;\n lapack_geev_dispatch(&jobvl, &jobvr, &dim, A, &lda, eigenvalues, WI, vl, &ldvl,\n eigenvectors, &ldvr, work, &work_size, &info);\n lapackCheckError(info);\n delete [] WI;\n delete [] work;\n}\n\n#endif\n\n} // end anonymous namespace\n\ntemplate \nvoid Lapack< TemplateConfig >::geev(const Vector &A, Vector &eigenvalues)\n{\n check_lapack_enabled();\n typedef typename Vector::value_type value_type;\n // It is possible the matrix has an extra row (e.g. Arnoldi).\n int dim = std::min(A.get_num_rows(), A.get_num_cols());\n int lda = A.get_lda();\n value_type *A_ptr = const_cast(A.raw());\n#ifdef AMGX_USE_LAPACK\n lapack_geev(A_ptr, eigenvalues.raw(), dim, lda);\n#else\n FatalError(\"Lapack is not supported in this build\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::geev(const Vector &A,\n Vector &eigenvalues,\n Vector &eigenvector)\n{\n check_lapack_enabled();\n typedef typename Vector::value_type value_type;\n // It is possible the matrix has an extra row (e.g. Arnoldi).\n int dim = std::min(A.get_num_rows(), A.get_num_cols());\n int lda = A.get_lda();\n value_type *A_ptr = const_cast(A.raw());\n#ifdef AMGX_USE_LAPACK\n lapack_geev(A_ptr, eigenvalues.raw(), eigenvector.raw(), dim, lda, eigenvector.get_lda());\n#else\n FatalError(\"Lapack is not supported in this build\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\n\ntemplate \nvoid Lapack< TemplateConfig >::geev(const Vector &A, Vector &eigenvalues)\n{\n not_implemented();\n}\n\nnamespace\n{\n#ifdef AMGX_USE_LAPACK\nextern \"C\"\nint dtrtri_(char *uplo, char *diag, int *n, double *\n a, int *lda, int *info);\n\nextern \"C\"\nint strtri_(char *uplo, char *diag, int *n, float *\n a, int *lda, int *info);\n\nextern \"C\"\nint ctrtri_(char *uplo, char *diag, int *n, fcomplex *\n a, int *lda, int *info);\n\nextern \"C\"\nint ztrtri_(char *uplo, char *diag, int *n, dcomplex *\n a, int *lda, int *info);\n\nint lapack_trtri_dispatch(char *uplo, char *diag, int *n, float *a,\n int *lda, int *info)\n{\n return strtri_(uplo, diag, n, a, lda, info);\n}\n\nint lapack_trtri_dispatch(char *uplo, char *diag, int *n, double *a,\n int *lda, int *info)\n{\n return dtrtri_(uplo, diag, n, a, lda, info);\n}\n\nint lapack_trtri_dispatch(char *uplo, char *diag, int *n, fcomplex *a,\n int *lda, int *info)\n{\n return ctrtri_(uplo, diag, n, a, lda, info);\n}\n\nint lapack_trtri_dispatch(char *uplo, char *diag, int *n, dcomplex *a,\n int *lda, int *info)\n{\n return ztrtri_(uplo, diag, n, a, lda, info);\n}\n\ntemplate \nvoid lapack_trtri(T *A, int dim, int lda)\n{\n char uplo = 'U';\n char diag = 'N';\n int info;\n lapack_trtri_dispatch(&uplo, &diag, &dim, A, &lda, &info);\n lapackCheckError(info);\n}\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::trtri(Vector &A)\n{\n check_lapack_enabled();\n typedef typename Vector::value_type value_type;\n int dim = std::min(A.get_num_rows(), A.get_num_cols());\n int lda = A.get_lda();\n#ifdef AMGX_USE_LAPACK\n lapack_trtri(A.raw(), dim, lda);\n#else\n FatalError(\"Lapack is not supported in this build\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\nnamespace\n{\n#ifdef AMGX_USE_MAGMA\nint magma_trtri_dispatch(magma_uplo_t uplo, magma_diag_t diag, int n, float *a,\n int lda, int *info)\n{\n return magma_strtri_gpu(uplo, diag, n, a, lda, info);\n}\n\nint magma_trtri_dispatch(magma_uplo_t uplo, magma_diag_t diag, int n, double *a,\n int lda, int *info)\n{\n return magma_dtrtri_gpu(uplo, diag, n, a, lda, info);\n}\n\nint magma_trtri_dispatch(magma_uplo_t uplo, magma_diag_t diag, int n, hipComplex *a,\n int lda, int *info)\n{\n return magma_ctrtri_gpu(uplo, diag, n, a, lda, info);\n}\n\nint magma_trtri_dispatch(magma_uplo_t uplo, magma_diag_t diag, int n, hipDoubleComplex *a,\n int lda, int *info)\n{\n return magma_ztrtri_gpu(uplo, diag, n, a, lda, info);\n}\n\ntemplate \nvoid magma_trtri(T *A, int dim, int lda)\n{\n magma_uplo_t uplo = MagmaUpper;\n magma_diag_t diag = MagmaNonUnit;\n int info;\n magma_trtri_dispatch(uplo, diag, dim, A, lda, &info);\n magmaCheckError(info);\n}\n\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::trtri(Vector &A)\n{\n check_magma_enabled();\n typedef typename Vector::value_type value_type;\n int dim = std::min(A.get_num_rows(), A.get_num_cols());\n int lda = A.get_lda();\n#ifdef AMGX_USE_MAGMA\n magma_trtri(A.raw(), dim, lda);;\n#else\n FatalError(\"Lapack is not supported in this build\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\nnamespace\n{\n#ifdef AMGX_USE_LAPACK\nextern \"C\"\nint dsygv_(int *itype, char *jobz, char *uplo, int *n,\n double *a, int *lda, double *b, int *ldb,\n double *w, double *work, int *lwork, int *info);\n\nextern \"C\"\nint ssygv_(int *itype, char *jobz, char *uplo, int *n,\n float *a, int *lda, float *b, int *ldb,\n float *w, float *work, int *lwork, int *info);\n\nextern \"C\"\nint chegv_(int *itype, char *jobz, char *uplo, int *n,\n fcomplex *a, int *lda, fcomplex *b, int *ldb,\n fcomplex *w, fcomplex *work, int *lwork, int *info);\n\nextern \"C\"\nint zhegv_(int *itype, char *jobz, char *uplo, int *n,\n dcomplex *a, int *lda, dcomplex *b, int *ldb,\n dcomplex *w, dcomplex *work, int *lwork, int *info);\n\nint lapack_sygv_dispatch(int *itype, char *jobz, char *uplo, int *n,\n double *a, int *lda, double *b, int *ldb,\n double *w, double *work, int *lwork, int *info)\n{\n return dsygv_(itype, jobz, uplo, n, a, lda, b, ldb, w, work, lwork, info);\n}\n\nint lapack_sygv_dispatch(int *itype, char *jobz, char *uplo, int *n,\n float *a, int *lda, float *b, int *ldb,\n float *w, float *work, int *lwork, int *info)\n{\n return ssygv_(itype, jobz, uplo, n, a, lda, b, ldb, w, work, lwork, info);\n}\n\nint lapack_sygv_dispatch(int *itype, char *jobz, char *uplo, int *n,\n hipComplex *a, int *lda, hipComplex *b, int *ldb,\n hipComplex *w, hipComplex *work, int *lwork, int *info)\n{\n return chegv_(itype, jobz, uplo, n,\n reinterpret_cast(a), lda, reinterpret_cast(b), ldb,\n reinterpret_cast(w), reinterpret_cast(work), lwork, info);\n}\n\nint lapack_sygv_dispatch(int *itype, char *jobz, char *uplo, int *n,\n hipComplex *a, int *lda, hipDoubleComplex *b, int *ldb,\n hipDoubleComplex *w, hipDoubleComplex *work, int *lwork, int *info)\n{\n return zhegv_(itype, jobz, uplo, n,\n reinterpret_cast(a), lda, reinterpret_cast(b), ldb,\n reinterpret_cast(w), reinterpret_cast(work), lwork, info);\n}\n\ntemplate \nvoid lapack_sygv(T *gramA, T *gramB, T *eigenvector, int dim, int lda, T *work)\n{\n int itype = 1;\n char jobz = 'V';\n char uplo = 'U';\n int ldb = lda;\n int lwork = 1024;\n int info = 0;\n lapack_sygv_dispatch(&itype, &jobz, &uplo, &dim, gramA, &lda, gramB, &ldb, eigenvector, work, &lwork, &info);\n lapackCheckError(info);\n}\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::sygv(Vector &A, Vector &B,\n Vector &eigenvalues, Vector &work)\n{\n check_lapack_enabled();\n typedef typename Vector::value_type value_type;\n int dim = std::min(A.get_num_rows(), A.get_num_cols());\n int lda = A.get_lda();\n#ifdef AMGX_USE_LAPACK\n lapack_sygv(A.raw(), B.raw(), eigenvalues.raw(), dim, lda, work.raw());\n#else\n FatalError(\"Lapack is not supported in this build\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\nnamespace\n{\n#ifdef AMGX_USE_MAGMA\nvoid magma_trsm_dispatch(magma_side_t side, magma_uplo_t uplo, magma_trans_t trans,\n magma_diag_t diag, magma_int_t m, magma_int_t n,\n float alpha, float const *dA, magma_int_t lda,\n float *dB, magma_int_t ldb)\n{\n return magma_strsm(side, uplo, trans, diag, m, n, alpha, dA, lda, dB, ldb);\n}\n\nvoid magma_trsm_dispatch(magma_side_t side, magma_uplo_t uplo, magma_trans_t trans,\n magma_diag_t diag, magma_int_t m, magma_int_t n,\n double alpha, double const *dA, magma_int_t lda,\n double *dB, magma_int_t ldb)\n{\n return magma_dtrsm(side, uplo, trans, diag, m, n, alpha, dA, lda, dB, ldb);\n}\n\nvoid magma_trmm_dispatch(magma_side_t side, magma_uplo_t uplo, magma_trans_t trans,\n magma_diag_t diag, magma_int_t m, magma_int_t n,\n float alpha, float const *dA, magma_int_t lda,\n float *dB, magma_int_t ldb)\n{\n return magma_strmm(side, uplo, trans, diag, m, n, alpha, dA, lda, dB, ldb);\n}\n\nvoid magma_trmm_dispatch(magma_side_t side, magma_uplo_t uplo, magma_trans_t trans,\n magma_diag_t diag, magma_int_t m, magma_int_t n,\n double alpha, double const *dA, magma_int_t lda,\n double *dB, magma_int_t ldb)\n{\n return magma_dtrmm(side, uplo, trans, diag, m, n, alpha, dA, lda, dB, ldb);\n}\n\nint magma_potrf_gpu_dispatch(magma_uplo_t uplo, int n, float *A, int lda, int *info)\n{\n return magma_spotrf_gpu(uplo, n, A, lda, info);\n}\n\nint magma_potrf_gpu_dispatch(magma_uplo_t uplo, int n, double *A, int lda, int *info)\n{\n return magma_dpotrf_gpu(uplo, n, A, lda, info);\n}\n\nint magma_sygst_gpu_dispatch(int itype, magma_uplo_t uplo, magma_int_t n, float *da,\n int ldda, float *B, int lddb, int *info)\n{\n return magma_ssygst_gpu(itype, uplo, n, da, ldda, B, lddb, info);\n}\n\nint magma_sygst_gpu_dispatch(int itype, magma_uplo_t uplo, magma_int_t n, double *da,\n int ldda, double *B, int lddb, int *info)\n{\n return magma_dsygst_gpu(itype, uplo, n, da, ldda, B, lddb, info);\n}\n\nint magma_syevd_gpu_dispatch(magma_vec_t jobz, magma_uplo_t uplo, int n, double *da, int ldda,\n double *w, double *wa, int ldwa, double *work,\n int lwork, int *iwork, int liwork, int *info)\n{\n return magma_dsyevd_gpu(jobz, uplo, n, da, ldda, w, wa, ldwa, work, lwork, iwork, liwork, info);\n}\n\nint magma_syevd_gpu_dispatch(magma_vec_t jobz, magma_uplo_t uplo, int n, float *da, int ldda,\n float *w, float *wa, int ldwa, float *work,\n int lwork, int *iwork, int liwork, int *info)\n{\n return magma_ssyevd_gpu(jobz, uplo, n, da, ldda, w, wa, ldwa, work, lwork, iwork, liwork, info);\n}\n\n// This is a simple modification of the magma_?sygvd() source code\n// from magma where the matrices are already on the device.\ntemplate \nmagma_int_t magma_sygvd_gpu_impl(magma_int_t itype, magma_vec_t jobz, magma_uplo_t uplo, magma_int_t n,\n T *da, magma_int_t lda, T *db, magma_int_t ldb, T *w,\n T *work, magma_int_t lwork, magma_int_t *iwork, magma_int_t liwork,\n T *wa, magma_int_t *info)\n{\n magma_uplo_t uplo_[2] = {uplo, MagmaLower}; // {uplo, 0}\n magma_vec_t jobz_[2] = {jobz, MagmaVec};//{jobz, 0};\n T d_one = MAGMA_D_ONE;\n magma_int_t ldda = n;\n magma_int_t lddb = n;\n static magma_int_t lower;\n static char trans[1];\n static magma_int_t wantz, lquery;\n static magma_int_t lopt, lwmin, liopt, liwmin;\n static hipStream_t stream;\n magma_queue_create( &stream );\n wantz = jobz_[0] == MagmaVec;\n lower = uplo_[0] == MagmaLower;\n lquery = lwork == -1 || liwork == -1;\n *info = 0;\n\n if (itype < 1 || itype > 3)\n {\n *info = -1;\n }\n else if (! (wantz || jobz_[0] == MagmaNoVec))\n {\n *info = -2;\n }\n else if (! (lower || uplo_[0] == MagmaUpper))\n {\n *info = -3;\n }\n else if (n < 0)\n {\n *info = -4;\n }\n else if (lda < max(1, n))\n {\n *info = -6;\n }\n else if (ldb < max(1, n))\n {\n *info = -8;\n }\n\n magma_int_t nb = magma_get_dsytrd_nb(n);\n\n if (n < 1)\n {\n liwmin = 1;\n lwmin = 1;\n }\n else if (wantz)\n {\n lwmin = 1 + 6 * n * nb + 2 * n * n;\n liwmin = 5 * n + 3;\n }\n else\n {\n lwmin = 2 * n * nb + 1;\n liwmin = 1;\n }\n\n lopt = lwmin;\n liopt = liwmin;\n work[ 0] = lopt;\n iwork[0] = liopt;\n\n if (lwork < lwmin && ! lquery)\n {\n *info = -11;\n }\n else if (liwork < liwmin && ! lquery)\n {\n *info = -13;\n }\n\n if (*info != 0)\n {\n magma_xerbla( __func__, -(*info) );\n return MAGMA_ERR_ILLEGAL_VALUE;\n }\n else if (lquery)\n {\n return MAGMA_SUCCESS;\n }\n\n /* Quick return if possible */\n if (n == 0)\n {\n return 0;\n }\n\n magma_potrf_gpu_dispatch(uplo_[0], n, db, lddb, info);\n\n if (*info != 0)\n {\n *info = n + *info;\n return 0;\n }\n\n /* Transform problem to standard eigenvalue problem and solve. */\n magma_sygst_gpu_dispatch(itype, uplo_[0], n, da, ldda, db, lddb, info);\n magma_syevd_gpu_dispatch(jobz_[0], uplo_[0], n, da, ldda, w, wa, lda,\n work, lwork, iwork, liwork, info);\n lopt = max( lopt, (magma_int_t) work[0]);\n liopt = max(liopt, iwork[0]);\n\n if (wantz && *info == 0)\n {\n /* Backtransform eigenvectors to the original problem. */\n if (itype == 1 || itype == 2)\n {\n /* For A*x=(lambda)*B*x and A*B*x=(lambda)*x;\n backtransform eigenvectors: x = inv(L)'*y or inv(U)*y */\n if (lower)\n {\n *(unsigned char *)trans = MagmaTrans;\n }\n else\n {\n *(unsigned char *)trans = MagmaNoTrans;\n }\n\n magma_trsm_dispatch(MagmaLeft, uplo_[0], *trans, MagmaNonUnit,\n n, n, d_one, db, lddb, da, ldda);\n }\n else if (itype == 3)\n {\n /* For B*A*x=(lambda)*x;\n backtransform eigenvectors: x = L*y or U'*y */\n if (lower)\n {\n *(unsigned char *)trans = MagmaNoTrans;\n }\n else\n {\n *(unsigned char *)trans = MagmaTrans;\n }\n\n magma_trmm_dispatch(MagmaLeft, uplo_[0], *trans, MagmaNonUnit,\n n, n, d_one, db, lddb, da, ldda);\n }\n }\n\n magma_queue_sync( stream );\n magma_queue_destroy( stream );\n work[0] = (T) lopt;\n iwork[0] = liopt;\n return MAGMA_SUCCESS;\n}\n\nhipblasStatus_t cublas_trsm(hipblasHandle_t handle,\n hipblasSideMode_t side, hipblasFillMode_t uplo,\n hipblasOperation_t trans, hipblasDiagType_t diag,\n int m, int n,\n const float *alpha,\n const float *A, int lda,\n float *B, int ldb)\n{\n return hipblasStrsm(handle, side, uplo, trans, diag, m, n, alpha, A, lda, B, ldb);\n}\n\nhipblasStatus_t cublas_trsm(hipblasHandle_t handle,\n hipblasSideMode_t side, hipblasFillMode_t uplo,\n hipblasOperation_t trans, hipblasDiagType_t diag,\n int m, int n,\n const double *alpha,\n const double *A, int lda,\n double *B, int ldb)\n{\n return hipblasDtrsm(handle, side, uplo, trans, diag, m, n, alpha, A, lda, B, ldb);\n}\n\n\ntemplate \nvoid magma_sygvd_gpu(T *A, T *B, T *eigenvalues, int dim, int lda)\n{\n int itype = 1;\n magma_vec_t jobz = MagmaVec;\n magma_uplo_t uplo = MagmaUpper;\n int N = dim;\n int ldb = lda;\n int nb = 32;\n int lwork = 1 + 6 * N * nb + 2 * N * N;\n static std::vector s_work;\n s_work.resize(lwork);\n T *work = &s_work[0];\n int liwork = 3 + 5 * N;\n static std::vector s_iwork;\n s_iwork.resize(liwork);\n int *iwork = &s_iwork[0];\n static std::vector s_wa;\n s_wa.resize(lda * N);\n T *wa = &s_wa[0];\n int ldwa = N;\n int info;\n /*\n magma_sygvd_gpu_impl(itype, jobz, uplo, N, A, lda, B, ldb, eigenvalues, work, lwork, iwork, liwork, wa, &info);\n */\n magma_potrf_gpu_dispatch(uplo, N, B, lda, &info);\n magmaCheckError(info);\n magma_sygst_gpu_dispatch(itype, uplo, N, A, lda, B, ldb, &info);\n magmaCheckError(info);\n magma_syevd_gpu_dispatch(jobz, uplo, N, A, lda, eigenvalues, wa, ldwa, work, lwork, iwork, liwork, &info);\n magmaCheckError(info);\n T one = 1;\n hipblasHandle_t handle = Cublas::get_handle();\n cublas_trsm(handle, HIPBLAS_SIDE_LEFT, HIPBLAS_FILL_MODE_UPPER, HIPBLAS_OP_N, HIPBLAS_DIAG_NON_UNIT, N, N, &one, B, ldb, A, lda);\n}\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::sygv(Vector &A, Vector &B,\n Vector &eigenvalues, Vector &work)\n{\n typedef typename Vector::value_type value_type;\n int dim = std::min(A.get_num_rows(), A.get_num_cols());\n int lda = A.get_lda();\n#ifdef AMGX_USE_MAGMA\n magma_sygvd_gpu(A.raw(), B.raw(), eigenvalues.raw(), dim, lda);\n#endif\n}\n\nnamespace\n{\n#ifdef AMGX_USE_MAGMA\ntemplate \nvoid magma_syevd_gpu(T *A, T *eigenvalues, int dim, int lda)\n{\n magma_vec_t jobz = MagmaVec;\n magma_uplo_t uplo = MagmaUpper;\n int N = dim;\n int nb = 32;\n int lwork = 1 + 6 * N * nb + 2 * N * N;\n static std::vector s_work;\n s_work.resize(lwork);\n T *work = &s_work[0];\n int liwork = 3 + 5 * N;\n static std::vector s_iwork;\n s_iwork.resize(liwork);\n int *iwork = &s_iwork[0];\n static std::vector s_wa;\n s_wa.resize(lda * N);\n T *wa = &s_wa[0];\n int ldwa = N;\n int info;\n magma_syevd_gpu_dispatch(jobz, uplo, N, A, lda, eigenvalues, wa, ldwa, work, lwork, iwork, liwork, &info);\n magmaCheckError(info);\n}\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::syevd(Vector &A,\n Vector &eigenvalues)\n{\n check_magma_enabled();\n typedef typename Vector::value_type value_type;\n int dim = std::min(A.get_num_rows(), A.get_num_cols());\n int lda = A.get_lda();\n#ifdef AMGX_USE_MAGMA\n magma_syevd_gpu(A.raw(), eigenvalues.raw(), dim, lda);\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::syevd(Vector &A,\n Vector &eigenvalues)\n{\n not_implemented();\n}\n\nnamespace\n{\n#ifdef AMGX_USE_MAGMA\nint magma_stedx_dispatch(magma_range_t range, int n,\n double vl, double vu,\n int il, int iu,\n double *d, double *e, double *z, int ldz,\n double *work, int lwork, int *iwork, int liwork,\n double *dwork, int *info)\n{\n return magma_dstedx(range, n, vl, vu, il, iu, d, e, z, ldz, work, lwork, iwork, liwork, dwork, info);\n}\n\nint magma_stedx_dispatch(magma_range_t range, int n,\n float vl, float vu,\n int il, int iu,\n float *d, float *e, float *z, int ldz,\n float *work, int lwork, int *iwork, int liwork,\n float *dwork, int *info)\n{\n return magma_sstedx(range, n, vl, vu, il, iu, d, e, z, ldz, work, lwork, iwork, liwork, dwork, info);\n}\n\ntemplate \nvoid magma_stedx(T *diagonal, T *subdiagonal, T *eigenvectors,\n int lower, int upper, int dim, int ldz, T *dwork, int dwork_size)\n{\n magma_range_t range = MagmaRangeI;\n int N = dim;\n T vl = 0;\n T vu = 0;\n int il = lower;\n int iu = upper;\n int lwork = 1 + 4 * N + 2 * N * N;\n static std::vector s_work;\n s_work.resize(lwork);\n int liwork = 3 + 6 * N;\n static std::vector s_iwork;\n s_iwork.resize(liwork);\n int info;\n magma_stedx_dispatch(range, N, vl, vu, il, iu, diagonal, subdiagonal, eigenvectors, ldz,\n &s_work[0], lwork, &s_iwork[0], liwork, dwork, &info);\n}\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::stedx(Vector &diagonal,\n Vector &subdiagonal,\n Vector &eigenvectors,\n int dim,\n Vector &dwork)\n{\n check_magma_enabled();\n#ifdef AMGX_USE_MAGMA\n magma_stedx(diagonal.raw(), subdiagonal.raw(), eigenvectors.raw(),\n dim, dim, dim, eigenvectors.get_lda(),\n dwork.raw(), dwork.size());\n#endif\n}\n\n\nnamespace\n{\n\ntemplate \nvoid larf(int m, int n, T *v,\n int incv, T *tau, T *c, int ldc,\n T *work)\n{\n /* Table of constant values */\n static T c_b4 = 1.;\n static T c_b5 = 0.;\n static int c1 = 1;\n /* Form H * C */\n /* w := C' * v */\n Cublas::gemv(true, m, n, &c_b4, c, ldc,\n v, incv, &c_b5, work, c1);\n /* C := C - v * w' */\n Cublas::ger(m, n, tau, v, incv, work, c1, c, ldc);\n}\n\ntemplate \n__global__\nvoid set1(T *a)\n{\n *a = 1.;\n}\n\ntemplate \n__global__\nvoid add_tau(T *a, T tau)\n{\n *a = 1 + tau;\n}\n\ntemplate \nvoid gpu_orgqr(int m, int n, int k,\n T *a, int lda, T *tau, T *work, int lwork)\n{\n int i1, i2;\n\n for (int i = k - 1; i >= 0; --i)\n {\n /* Apply H(i) to A(i:m,i:n) from the left */\n if (i < n - 1)\n {\n set1 <<< 1, 1>>>(&a[i + i * lda]);\n cudaCheckError();\n i1 = m - i;\n i2 = n - i - 1;\n larf(i1, i2, &a[i + i * lda], 1, &tau[i],\n &a[i + (i + 1) * lda], lda, work);\n }\n\n if (i < m - 1)\n {\n i1 = m - i - 1;\n Cublas::scal(i1, &tau[i], &a[i + 1 + i * lda], 1);\n }\n\n add_tau <<< 1, 1>>>(&a[i + i * lda], tau[i]);\n cudaCheckError();\n /* Set A(1:i-1,i) to zero */\n hipError_t cuda_rc = hipMemset(&a[i * lda], 0, sizeof(T) * i);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemset failed in geqr2\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \n__device__ __host__\nT lapy2_(T *a, T *b)\n{\n T va = *a;\n T vb = *b;\n return sqrt(va * va + vb * vb);\n}\n\ntemplate \n__device__ __host__\nT d_sign(T a, T b)\n{\n T x;\n x = (a >= 0 ? a : - a);\n return (b >= 0 ? x : -x);\n}\n\ntemplate \nvoid compute_tau_host(T *alpha, T *norm,\n T *tau, T *d1)\n{\n *d1 = lapy2_(alpha, norm);\n T beta = -d_sign(*d1, *alpha);\n // LAPACK: skipped part about scaling.\n // Negated compared to LAPACK code, avoid negating value on device later.\n *tau = -(beta - *alpha) / beta;\n *d1 = 1. / (*alpha - beta);\n *alpha = beta;\n}\n\ntemplate \nvoid larfg(int n, T *alpha, T *x,\n int incx, T *tau)\n{\n if (n <= 1)\n {\n *tau = 0.;\n return;\n }\n\n int i1 = n - 1;\n T xnorm;\n Cublas::nrm2(i1, x, incx, &xnorm);\n T h_alpha;\n hipError_t cuda_rc = hipMemcpy(&h_alpha, alpha, sizeof(T), hipMemcpyDeviceToHost);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy D2H failed in larfg\", AMGX_ERR_CUDA_FAILURE);\n }\n T d1;\n compute_tau_host(&h_alpha, &xnorm, tau, &d1);\n Cublas::scal(i1, d1, x, incx);\n // Update the diagonal value on the device.\n cuda_rc = hipMemcpy(alpha, &h_alpha, sizeof(T), hipMemcpyHostToDevice);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy H2D failed in larfg\", AMGX_ERR_CUDA_FAILURE);\n }\n}\n\ntemplate \nvoid gpu_geqrf(int m, int n, T *a, int lda,\n T *tau, T *work)\n{\n int k = std::min(m, n);\n T *aii;\n hipError_t cuda_rc = amgx::memory::hipMallocAsync((void**)&aii, sizeof(T));\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMallocAsync failed in gpu_geqrf\", AMGX_ERR_CUDA_FAILURE);\n }\n\n for (int i = 0; i < k; ++i)\n {\n /* Generate elementary reflector H(i) to annihilate A(i+1:m,i) */\n int i2 = m - i;\n /* Computing MIN */\n int i3 = i + 1;\n larfg(i2, &a[i + i * lda],\n &a[std::min(i3, m - 1) + i * lda],\n 1, &tau[i]);\n\n if (i < n - 1)\n {\n /* Apply H(i) to A(i:m,i+1:n) from the left */\n cuda_rc = hipMemcpy(aii, &a[i + i * lda], sizeof(T), hipMemcpyDeviceToDevice);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy D2D failed in gpu_geqrf\", AMGX_ERR_CUDA_FAILURE);\n }\n set1 <<< 1, 1>>>(&a[i + i * lda]);\n cudaCheckError();\n i2 = m - i;\n i3 = n - i - 1;\n larf(i2, i3, &a[i + i * lda], 1,\n &tau[i], &a[i + (i + 1) * lda], lda, work);\n cuda_rc = hipMemcpy(&a[i + i * lda], aii, sizeof(T), hipMemcpyDeviceToDevice);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy D2D failed in gpu_geqrf\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n }\n\n cuda_rc = amgx::memory::hipFreeAsync(aii);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipFreeAsync failed in gpu_geqrf\", AMGX_ERR_CUDA_FAILURE);\n }\n}\n} // end anonymous namespace\n\nnamespace\n{\n#ifdef AMGX_USE_MAGMA\nint magma_geqrf_dispatch(int m, int n, float *A, int lda,\n float *tau, float *work, int *info)\n{\n return magma_sgeqrf_gpu(m, n, A, lda, tau, work, info);\n}\n\nint magma_geqrf_dispatch(int m, int n, double *A, int lda,\n double *tau, double *work, int *info)\n{\n return magma_dgeqrf_gpu(m, n, A, lda, tau, work, info);\n}\n\ntemplate \nvoid magma_geqrf(int m, int n, T *A, int lda,\n T *tau, T *work)\n{\n int info;\n magma_geqrf_dispatch(m, n, A, lda, tau, work, &info);\n magmaCheckError(info);\n}\n\nint magma_orgqr_dispatch(int m, int n, int k, float *A, int lda,\n float *tau, float *work, int lwork, int *info)\n{\n return magma_sorgqr_gpu(m, n, k, A, lda, tau, work, lwork, info);\n}\n\nint magma_orgqr_dispatch(int m, int n, int k, double *A, int lda,\n double *tau, double *work, int lwork, int *info)\n{\n return magma_dorgqr_gpu(m, n, k, A, lda, tau, work, lwork, info);\n}\n\ntemplate \nvoid magma_orgqr(int m, int n, int k, T *A, int lda,\n T *tau, T *work, int lwork)\n{\n int info;\n magma_orgqr_dispatch(m, n, k, A, lda, tau, work, lwork, &info);\n magmaCheckError(info);\n}\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::geqrf(Vector &A,\n Vector &tau,\n Vector &work)\n{\n not_implemented();\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::geqrf(Vector &A,\n Vector &tau,\n Vector &work)\n{\n int rows = A.get_num_rows();\n int cols = A.get_num_cols();\n int lda = A.get_lda();\n#ifdef AMGX_USE_MAGMA\n magma_geqrf(rows, cols, A.raw(), lda, tau.raw(), work.raw());\n#else\n gpu_geqrf(rows, cols, A.raw(), lda, tau.raw(), work.raw());\n#endif\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::orgqr(Vector &A,\n Vector &tau,\n Vector &work)\n{\n not_implemented();\n}\n\ntemplate \nvoid Lapack< TemplateConfig >::orgqr(Vector &A,\n Vector &tau,\n Vector &work)\n{\n int rows = A.get_num_rows();\n int cols = A.get_num_cols();\n int lda = A.get_lda();\n#ifdef AMGX_USE_MAGMA\n magma_orgqr(rows, cols, cols, A.raw(), lda, tau.raw(), work.raw(), 1);\n#else\n gpu_orgqr(rows, cols, cols, A.raw(), lda, tau.raw(), work.raw(), 1);\n#endif\n}\n\n\n#define AMGX_CASE_LINE(CASE) \\\n template class Lapack::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n", "source": "/kaggle/working/amgx/src/amgx_lapack.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"amgx_types/util.h\"\n\n#define USE_EXPERIMENTAL_FACTOR_COPY\n// #define USE_EXPERIMENTAL_NLARGEST\n\nnamespace amgx\n{\n\n__host__ __device__ __forceinline__\ndouble ourHash2(const int i)\n{\n unsigned int a = i;\n double result;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n result = (a ^ 0x4a51e590) / (double) UINT_MAX;\n return result;\n}\n\ntemplate \nstruct copy_pred\n{\n double _trunc_factor;\n metricType *_metric_arr;\n\n copy_pred() : _trunc_factor(0), _metric_arr(NULL) {};\n copy_pred(const double trunc_factor, metricType *metric_arr) : _trunc_factor(trunc_factor),\n _metric_arr(metric_arr) {};\n\n __host__ __device__\n bool operator()(const amgx::thrust::tuple &a)\n {\n metricType metric = _metric_arr[amgx::thrust::get<0>(a)];\n\n if ( types::util::abs(amgx::thrust::get<2>(a)) >= _trunc_factor * metric) { return true; }\n\n return false;\n }\n\n copy_pred &operator=(const copy_pred &a)\n {\n this->_trunc_factor = a._trunc_factor;\n this->_metric_arr = a._metric_arr;\n return *this;\n }\n};\n\ntemplate \nstruct scale_op\n{\n const VectorType *scale_vec;\n\n scale_op(const VectorType *s) : scale_vec(s) {};\n\n __host__ __device__\n amgx::thrust::tuple operator()(const amgx::thrust::tuple &a)\n {\n const IndexType row = amgx::thrust::get<0>(a);\n return amgx::thrust::tuple(row, amgx::thrust::get<1>(a) * scale_vec[row]);\n }\n};\n\ntemplate \n__global__\nvoid scale_kernel(const index_type *row_indices, value_type *values, const value_type *new_row_sums, const value_type *old_row_sums, const int size)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n int row_index = row_indices[tidx];\n value_type multiplier = ( types::util::abs(new_row_sums[row_index]) == 0. ) ? types::util::get_one() : (old_row_sums[row_index] / new_row_sums[row_index]);\n values[tidx] = values[tidx] * multiplier;\n\n if (values[tidx] != values[tidx]) { printf(\"%d (%d): NaN after scaling\\n\", tidx, row_index); }\n }\n}\n\ntemplate \n__global__\nvoid get_row_lengths_kernel(const IndexType *offsets, const IndexType num_rows, IndexType *row_lengths,\n const int max_elmts)\n{\n for (int row = threadIdx.x + blockIdx.x * blockDim.x; row < num_rows; row += blockDim.x * gridDim.x)\n {\n row_lengths[row] = min(offsets[row + 1] - offsets[row], max_elmts);\n }\n}\n\n// sort array by abs(values), largest element in [0]\ntemplate \n__device__\nvoid sortByFabs(IndexType *indices, ValueType *values, int elements)\n{\n int n = elements;\n\n do\n {\n int newn = 0;\n\n for (int i = 1; i < n; i++)\n {\n if (types::util::abs(values[i - 1]) < types::util::abs(values[i]))\n {\n ValueType temp_v = values[i - 1];\n IndexType temp_i = indices[i - 1];\n values[i - 1] = values[i];\n values[i] = temp_v;\n indices[i - 1] = indices[i];\n indices[i] = temp_i;\n newn = i;\n }\n }\n\n n = newn;\n }\n while (n > 0);\n}\n\n#ifdef USE_EXPERIMENTAL_NLARGEST\n\ntemplate \n__device__\nvoid insert_new_max(volatile ValueType *Mv, volatile IndexType *Mi, volatile int *T, const ValueType S,\n const IndexType I, const int thread, const int N)\n{\n if (thread < N)\n {\n T[thread] = 1;\n\n // new value larger - shift my value down\n if (Mv[thread] < S)\n {\n T[thread] = 0;\n Mv[thread + 1] = Mv[thread];\n Mi[thread + 1] = Mi[thread];\n T[thread + 1] = 1;\n }\n\n // insert new largest value here\n if (!T[thread])\n {\n Mv[thread] = S;\n Mi[thread] = I;\n // inserted to me -- note to zero out\n }\n }\n}\n\ntemplate \n__global__\nvoid truncateNandScale_kernel(const IndexType *A_offsets, const IndexType *A_indices, const ValueType *A_values,\n const IndexType num_rows,\n const IndexType *At_offsets, IndexType *At_indices, ValueType *At_values,\n const int max_elmts, const VecValueType *old_row_sums)\n{\n const int WARP_SIZE = 32;\n const int VECTORS_PER_WARP = WARP_SIZE / THREADS_PER_VECTOR;\n const int WARPS_PER_BLOCK = BLOCKSIZE / WARP_SIZE;\n const int N_LARGEST_SIZE = (THREADS_PER_VECTOR + 1) * VECTORS_PER_WARP * WARPS_PER_BLOCK;\n __shared__ volatile int T[N_LARGEST_SIZE]; // has M array been modified\n __shared__ volatile ValueType Mv[N_LARGEST_SIZE]; // maximum values\n __shared__ volatile int Mi[N_LARGEST_SIZE]; // maximum indices\n __shared__ volatile ValueType S[WARP_SIZE * WARPS_PER_BLOCK]; // temporary holding for currently read values\n __shared__ volatile IndexType I[WARP_SIZE * WARPS_PER_BLOCK]; // -- \" -- indices\n ValueType S_copy;\n IndexType I_copy;\n __shared__ volatile int Ap[WARPS_PER_BLOCK * VECTORS_PER_WARP][2]; // store row offset pointers\n __shared__ volatile ValueType old_row_sum[WARPS_PER_BLOCK * VECTORS_PER_WARP];\n const int num_vectors = blockDim.x / THREADS_PER_VECTOR * gridDim.x;\n const int vector_id = (threadIdx.x + blockIdx.x * blockDim.x) / THREADS_PER_VECTOR;\n const int thread_lane = threadIdx.x & (THREADS_PER_VECTOR - 1);\n const int vector_lane = threadIdx.x / THREADS_PER_VECTOR;\n\n // initialise relevant shared memory arrays\n for (int i = threadIdx.x; i < N_LARGEST_SIZE; i += blockDim.x)\n {\n Mv[i] = 0.;\n Mi[i] = -1;\n T[i] = 0;\n }\n\n __syncthreads();\n\n // outer loop over rows\n for (int row = vector_id; row < num_rows; row += num_vectors)\n {\n // load row offsets into shared memory\n if (thread_lane < 2)\n {\n Ap[vector_lane][thread_lane] = A_offsets[row + thread_lane];\n }\n\n const int row_start = Ap[vector_lane][0];\n const int row_end = Ap[vector_lane][1];\n const int row_len = row_end - row_start;\n const int max_elmts_row = min(At_offsets[row + 1] - At_offsets[row], max_elmts);\n\n // degenerate case -- simply copy over entire row\n if (row_len <= max_elmts && row_len > 0 && thread_lane < row_len)\n {\n At_indices[At_offsets[row] + thread_lane] = A_indices[row_start + thread_lane];\n At_values[At_offsets[row] + thread_lane] = A_values[row_start + thread_lane];\n continue;\n }\n\n // loop over chunks of the row\n const int nloops = (row_len < THREADS_PER_VECTOR) ? 1 : (int)ceil((float)row_len / THREADS_PER_VECTOR);\n\n for (int i = 0; i < nloops; i++)\n {\n // zero out temp shared memory\n S[threadIdx.x] = 0;\n I[threadIdx.x] = -2;\n\n // read chunk into shared memory\n if (i * THREADS_PER_VECTOR + thread_lane < row_len)\n {\n S[threadIdx.x] = fabs(A_values[row_start + i * THREADS_PER_VECTOR + thread_lane]);\n I[threadIdx.x] = A_indices[row_start + i * THREADS_PER_VECTOR + thread_lane];\n }\n\n // loop over n' max elements\n for (int j = 0; j < max_elmts_row; j++)\n {\n // first take a copy of this section of the row so we don't re-read it\n S_copy = S[threadIdx.x];\n I_copy = I[threadIdx.x];\n // local max & associated index (save volatile loads)\n ValueType max = S_copy; // S[threadIdx.x];\n IndexType ind = I_copy; // I[threadIdx.x];\n\n // reduce to get max\n if (THREADS_PER_VECTOR > 16)\n {\n I[threadIdx.x] = ind = max > S[threadIdx.x + 16] ? ind : I[threadIdx.x + 16];\n S[threadIdx.x] = max = max > S[threadIdx.x + 16] ? max : S[threadIdx.x + 16];\n }\n\n if (THREADS_PER_VECTOR > 8)\n {\n I[threadIdx.x] = ind = max > S[threadIdx.x + 8] ? ind : I[threadIdx.x + 8];\n S[threadIdx.x] = max = max > S[threadIdx.x + 8] ? max : S[threadIdx.x + 8];\n }\n\n if (THREADS_PER_VECTOR > 4)\n {\n I[threadIdx.x] = ind = max > S[threadIdx.x + 4] ? ind : I[threadIdx.x + 4];\n S[threadIdx.x] = max = max > S[threadIdx.x + 4] ? max : S[threadIdx.x + 4];\n }\n\n if (THREADS_PER_VECTOR > 2)\n {\n I[threadIdx.x] = ind = max > S[threadIdx.x + 2] ? ind : I[threadIdx.x + 2];\n S[threadIdx.x] = max = max > S[threadIdx.x + 2] ? max : S[threadIdx.x + 2];\n }\n\n if (THREADS_PER_VECTOR > 1)\n {\n I[threadIdx.x] = ind = max > S[threadIdx.x + 1] ? ind : I[threadIdx.x + 1];\n S[threadIdx.x] = max = max > S[threadIdx.x + 1] ? max : S[threadIdx.x + 1];\n }\n\n // maximum now in S_copy[vector_lane*32] -- insert into maximum array\n const int sm_offset = vector_lane * (THREADS_PER_VECTOR + 1);\n insert_new_max(&Mv[sm_offset], &Mi[sm_offset], &T[sm_offset],\n S[vector_lane * THREADS_PER_VECTOR], I[vector_lane * THREADS_PER_VECTOR], thread_lane,\n max_elmts);\n S[threadIdx.x] = (I[vector_lane * THREADS_PER_VECTOR] == I_copy) ? 0 : S_copy;\n I[threadIdx.x] = I_copy;\n }\n }\n\n // get scaling factor for new row\n // copy Mv -> S\n S[threadIdx.x] = 0; // intialise unneeded S array to 0\n\n if (thread_lane < max_elmts_row)\n {\n S[threadIdx.x] = Mv[vector_lane * (THREADS_PER_VECTOR + 1) + thread_lane];\n }\n\n // reduce on S\n ValueType sum = fabs(S[threadIdx.x]);\n\n if (THREADS_PER_VECTOR > 16) { S[threadIdx.x] = sum = sum + S[threadIdx.x + 16]; }\n\n if (THREADS_PER_VECTOR > 8) { S[threadIdx.x] = sum = sum + S[threadIdx.x + 8]; }\n\n if (THREADS_PER_VECTOR > 4) { S[threadIdx.x] = sum = sum + S[threadIdx.x + 4]; }\n\n if (THREADS_PER_VECTOR > 2) { S[threadIdx.x] = sum = sum + S[threadIdx.x + 2]; }\n\n if (THREADS_PER_VECTOR > 1) { S[threadIdx.x] = sum = sum + S[threadIdx.x + 1]; }\n\n const ValueType scale = old_row_sums[row] / S[vector_lane * THREADS_PER_VECTOR];\n\n // copy final scaled results into output buffer\n if (thread_lane < max_elmts_row)\n {\n At_values[At_offsets[row] + thread_lane] = Mv[vector_lane * (THREADS_PER_VECTOR + 1) + thread_lane] * scale;\n At_indices[At_offsets[row] + thread_lane] = Mi[vector_lane * (THREADS_PER_VECTOR + 1) + thread_lane];\n }\n }\n}\n\n#else\n// sort array by column index\ntemplate \n__device__\nvoid sortByIndex(IndexType *indices, ValueType *values, int elements)\n{\n int n = elements;\n\n do\n {\n int newn = 0;\n\n for (int i = 1; i < n; i++)\n {\n if (indices[i - 1] > indices[i])\n {\n ValueType temp_v = values[i - 1];\n IndexType temp_i = indices[i - 1];\n values[i - 1] = values[i];\n values[i] = temp_v;\n indices[i - 1] = indices[i];\n indices[i] = temp_i;\n newn = i;\n }\n }\n\n n = newn;\n }\n while (n > 0);\n}\n\ntemplate \n__global__\nvoid row_sum_kernel(const IndexType *A_offsets, const ValueType *A_values,\n ValueType *row_sums, const IndexType num_rows, IndexType *row_lengths, const int max_elmts, bool compute_row_length)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_rows; tidx += blockDim.x * gridDim.x)\n {\n IndexType A_start = A_offsets[tidx], A_end = A_offsets[tidx + 1];\n\n if (compute_row_length)\n {\n row_lengths[tidx] = min(A_end - A_start, max_elmts);\n }\n\n ValueType row_sum = 0.;\n bool found_nans = false;\n\n for (int i = A_start; i < A_end; i++)\n {\n row_sum += A_values[i];\n\n if (A_values[i] != A_values[i]) { found_nans = true; }\n }\n\n if (found_nans)\n {\n printf(\"row %d has NaNs in input\\n\", tidx);\n }\n\n row_sums[tidx] = row_sum;\n }\n}\n\ntemplate \n__global__\nvoid __launch_bounds__(256)\ntruncate_kernel(const IndexType *A_offsets, const IndexType *A_col_indices, const ValueType *A_values,\n const IndexType num_rows,\n IndexType *trunc_offsets, IndexType *trunc_col_indices, ValueType *trunc_values,\n const int max_elmts)\n{\n // const int tidx = threadIdx.x + blockIdx.x*blockDim.x;\n __shared__ IndexType s_ind[1024]; // assume <= 256 threads per block\n __shared__ ValueType s_val[1024];\n const int smem_start = threadIdx.x * max_elmts;\n\n // if (tidx < num_rows)\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_rows; tidx += blockDim.x * gridDim.x)\n {\n IndexType A_start = A_offsets[tidx], A_end = A_offsets[tidx + 1];\n IndexType trunc_start = trunc_offsets[tidx];\n IndexType A_len = A_end - A_start;\n\n // copy my entire row if possible\n if (A_len <= max_elmts)\n {\n for (int j = 0; j < A_len; j++)\n {\n trunc_col_indices[trunc_start + j] = A_col_indices[A_start + j];\n trunc_values[trunc_start + j] = A_values[A_start + j];\n }\n }\n else // normal case\n {\n // read first max_elmts elements into shared\n for (int j = 0; j < max_elmts; j++)\n {\n s_ind[smem_start + j] = A_col_indices[A_start + j];\n s_val[smem_start + j] = A_values[A_start + j];\n }\n\n // sort the initially inserted elements (start = largest)\n sortByFabs(&s_ind[smem_start], &s_val[smem_start], max_elmts);\n\n // loop over all other elements on the row\n for (int j = A_start + max_elmts; j < A_end; j++)\n {\n // check against current maximums\n for (int i = 0; i < max_elmts; i++)\n {\n if (types::util::abs(A_values[j]) > types::util::abs(s_val[smem_start + i]))\n {\n // shift remaining values down\n for (int k = max_elmts - 1; k > i; k--)\n {\n s_val[smem_start + k] = s_val[smem_start + k - 1];\n s_ind[smem_start + k] = s_ind[smem_start + k - 1];\n }\n\n // tentatively replace element\n s_val[smem_start + i] = A_values[j];\n s_ind[smem_start + i] = A_col_indices[j];\n break;\n }\n }\n }\n\n //sortByIndex(&s_ind[smem_start],&s_val[smem_start],max_elmts);\n // now we have the 4 largest elements (hopefully) -- output\n for (int i = 0; i < max_elmts; i++)\n {\n trunc_col_indices[trunc_start + i] = s_ind[smem_start + i];\n trunc_values[trunc_start + i] = s_val[smem_start + i];\n }\n }\n }\n}\n#endif\n\ntemplate \n__global__\nvoid perturb_kernel(const IndexType *A_offsets, const IndexType *A_col_indices, ValueType *A_values,\n const IndexType num_rows,\n const int max_elmts)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_rows; tidx += blockDim.x * gridDim.x)\n {\n IndexType A_start = A_offsets[tidx], A_end = A_offsets[tidx + 1];\n\n // perturb A slightly\n for (int j = A_start; j < A_end; j++)\n {\n A_values[j] += ourHash2(A_col_indices[j]) * A_values[j];\n }\n }\n}\n\n\n// warp-based inclusive scan\n// from: Sengupta, Harris, Garland 2008\n// idx is index of thread\ntemplate \n__device__\nint scan_warp(volatile IndexType *ptr, const unsigned int idx)\n{\n const unsigned int lane = idx & 31; // index of thread in warp\n IndexType p = ptr[idx];\n\n if (lane >= 1) { ptr[idx] = p = ptr[idx - 1] + p; }\n\n if (lane >= 2) { ptr[idx] = p = ptr[idx - 2] + p; }\n\n if (lane >= 4) { ptr[idx] = p = ptr[idx - 4] + p; }\n\n if (lane >= 8) { ptr[idx] = p = ptr[idx - 8] + p; }\n\n if (lane >= 16) { ptr[idx] = p = ptr[idx - 16] + p; }\n\n return ptr[idx];\n}\n\ntemplate \n__global__\nvoid truncateAndScale_kernel(const IndexType *A_offsets, const IndexType *A_indices,\n const matValueType *A_values,\n const IndexType num_rows,\n IndexType *At_offsets, IndexType *At_indices, matValueType *At_values,\n const vecValueType *old_row_sums, const vecValueType *new_row_sums,\n const double truncate_factor, const vecValueType *metric)\n{\n const int WARP_SIZE = 32;\n const int WARPS_PER_BLOCK = THREADS_PER_BLOCK / WARP_SIZE;\n const int VECTORS_PER_WARP = WARP_SIZE / THREADS_PER_VECTOR;\n // allocate shared memory for warp-based scans (assume up to 512 threads / block)\n __shared__ volatile int s_scan[WARP_SIZE * WARPS_PER_BLOCK];\n __shared__ volatile int s_A_ptr[WARPS_PER_BLOCK * VECTORS_PER_WARP][2]; // row offsets for A\n __shared__ volatile int s_At_ptr[WARPS_PER_BLOCK * VECTORS_PER_WARP][2]; // row offsets for At\n __shared__ double s_scale[WARPS_PER_BLOCK * VECTORS_PER_WARP]; // scaling factor for each warp\n const int vectors_per_block = THREADS_PER_BLOCK / WARP_SIZE * VECTORS_PER_WARP;\n const int num_vectors = THREADS_PER_BLOCK / THREADS_PER_VECTOR * gridDim.x;\n const int vector_id = (threadIdx.x + blockIdx.x * blockDim.x) / THREADS_PER_VECTOR;\n const int thread_lane = threadIdx.x & (THREADS_PER_VECTOR - 1);\n const int vector_lane = threadIdx.x / THREADS_PER_VECTOR;\n int scan_offset = 0;\n matValueType val = 0., fval = 0.;\n IndexType ind = 0;\n\n // outer loop over rows\n for (int row = vector_id; row < num_rows; row += num_vectors)\n {\n scan_offset = 0;\n val = 0.;\n fval = 0.; // fabs(val)\n\n if (threadIdx.x < vectors_per_block && threadIdx.x < num_rows - row) // first threads in first warp\n {\n s_scale[threadIdx.x] = old_row_sums[row + threadIdx.x] / new_row_sums[row + threadIdx.x];\n }\n\n __syncthreads();\n\n if (thread_lane < 2)\n {\n s_A_ptr[vector_lane][thread_lane] = A_offsets[row + thread_lane];\n s_At_ptr[vector_lane][thread_lane] = At_offsets[row + thread_lane];\n }\n\n const int row_start = s_A_ptr[vector_lane][0];\n const int row_end = s_A_ptr[vector_lane][1];\n const int row_len = row_end - row_start;\n const int At_row_start = s_At_ptr[vector_lane][0];\n // mark & scan over consecutive blocks of 32 threads\n int nloops = (int)ceilf((float)row_len / THREADS_PER_VECTOR);\n#pragma unroll 2\n\n for (int i = 0; i < nloops; i++)\n {\n // initialise shared memory to 0\n s_scan[threadIdx.x] = 0;\n\n // mark appropriate elements of the row\n if (i * THREADS_PER_VECTOR + thread_lane < row_len)\n {\n val = A_values[row_start + i * THREADS_PER_VECTOR + thread_lane];\n ind = A_indices[row_start + i * THREADS_PER_VECTOR + thread_lane];\n fval = fabs(val);\n s_scan[threadIdx.x] = fval >= truncate_factor * metric[row] ? 1 : 0;\n }\n\n // now prefix scan over this\n scan_warp(&s_scan[vector_lane * THREADS_PER_VECTOR], thread_lane); // threadIdx.x%32);\n\n // copy relevant parts of A to At\n if (i * THREADS_PER_VECTOR + thread_lane < row_len)\n {\n if (fval >= truncate_factor * metric[row])\n {\n // get relevant index in At\n const int At_index = At_row_start + scan_offset + s_scan[threadIdx.x] - 1;\n\n if (s_scale[vector_lane] == 0.)\n {\n At_values[At_index] = val;\n }\n else\n {\n At_values[At_index] = val / s_scale[vector_lane];\n }\n\n At_indices[At_index] = ind;\n }\n }\n\n // note new scan_offset\n scan_offset = s_scan[vector_lane * THREADS_PER_VECTOR + THREADS_PER_VECTOR - 1];\n }\n }\n}\n\ntemplate \nvoid truncateAndScale(const Matrix &A, Matrix &A_trunc, const Vector &row_sum, const Vector &new_row_sum,\n const double trunc_factor, const Vector &metric)\n{\n const int num_blocks = std::min(4096, A.get_num_rows());\n truncateAndScale_kernel\n <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(),\n A.get_num_rows(),\n A_trunc.row_offsets.raw(), A_trunc.col_indices.raw(),\n A_trunc.values.raw(),\n row_sum.raw(), new_row_sum.raw(),\n trunc_factor, metric.raw());\n cudaCheckError();\n}\n\n// host code\ntemplate \nvoid Truncate >::truncateByFactor(Matrix_h &A,\n const double trunc_factor, const AMGX_TruncateType truncType)\n{\n // throw error\n FatalError(\"Truncation on host not implemented yet\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid Truncate >::truncateByMaxElements(Matrix_h &A,\n const int max_elmts)\n{\n // throw error\n FatalError(\"Truncation on host not implemented yet\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n// device code\n#ifdef _WIN32\ntemplate \nvoid Truncate >::truncateByFactor(Matrix_d &A, const double trunc_factor, const AMGX_TruncateType truncType)\n#else\ntemplate \nvoid Truncate >::truncateByFactor(Matrix_d &A, const double trunc_factor, const AMGX_TruncateType truncType) __attribute__((noinline))\n#endif\n{\n typedef typename TConfig_d::IndPrec index_type;\n typedef typename TConfig_d::VecPrec vec_value_type;\n typedef typename TConfig_d::MatPrec value_type;\n // store the weighting values per row\n Matrix A_trunc(0, 0, 0, CSR);\n typedef Vector VVector;\n typedef typename Matrix::MVector MVector;\n typedef Vector::Type> IVector;\n VVector row_sum(A.get_num_rows());\n VVector new_row_sum(A.get_num_rows());\n VVector max_coef;\n // point to either row_sum or max_coef depending on type of truncation desired\n // vec_value_type *metric = 0;\n VVector *metric;\n\n // get either max row coefficient or row sum\n if (truncType == AMGX_TruncateByRowSum)\n {\n absRowSum(A, row_sum);\n metric = &row_sum;\n }\n else if (truncType == AMGX_TruncateByMaxCoefficient)\n {\n max_coef.resize(A.get_num_rows());\n maxCoefAndSum(A, max_coef, row_sum);\n metric = &max_coef;\n }\n else\n {\n FatalError(\"Truncation type not implemented\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n IVector row_counts(A.get_num_rows());\n\n if (truncType == AMGX_TruncateByRowSum)\n {\n countTruncElements(A, trunc_factor, row_sum, row_counts, new_row_sum);\n }\n else if (truncType == AMGX_TruncateByMaxCoefficient)\n {\n countTruncElements(A, trunc_factor, max_coef, row_counts, new_row_sum);\n }\n\n cudaCheckError();\n // initial resize (so we can scan into the row_offsets array)\n A_trunc.resize(A.get_num_rows(), A.get_num_cols(), 0);\n cudaCheckError();\n\n thrust_wrapper::exclusive_scan(row_counts.begin(), row_counts.end(), A_trunc.row_offsets.begin());\n cudaCheckError();\n \n const int nnz = A_trunc.row_offsets[A.get_num_rows() - 1] + row_counts[A.get_num_rows() - 1];\n A_trunc.row_offsets[A.get_num_rows()] = nnz;\n\n if (nnz == A.get_num_nz()) // early return -- nothing truncated\n {\n return;\n }\n\n // final resize & get row indices for both matrices\n A_trunc.resize(A.get_num_rows(), A.get_num_cols(), nnz);\n#ifndef USE_EXPERIMENTAL_FACTOR_COPY\n A_trunc.addProps(COO);\n A.addProps(COO);\n // copy relevant values from A -> A_trunc\n copy_pred pred;\n\n if (truncType == AMGX_TruncateByRowSum)\n {\n pred = copy_pred(trunc_factor, row_sum.raw());\n }\n else if (truncType == AMGX_TruncateByMaxCoefficient)\n {\n pred = copy_pred(trunc_factor, max_coef.raw());\n }\n\n using amgx::thrust::make_zip_iterator;\n using amgx::thrust::make_tuple;\n amgx::thrust::copy_if(\n make_zip_iterator(make_tuple(A.row_indices.begin(), A.col_indices.begin(), A.values.begin())),\n make_zip_iterator(make_tuple(A.row_indices.end(), A.col_indices.end(), A.values.end())),\n make_zip_iterator(make_tuple(A_trunc.row_indices.begin(), A_trunc.col_indices.begin(),\n A_trunc.values.begin())),\n pred\n );\n // scale each row of the truncated matrix - per-row scale in new_row_sum\n amgx::thrust::transform(row_sum.begin(), row_sum.end(), new_row_sum.begin(), new_row_sum.begin(),\n amgx::thrust::divides());\n scale_op op(new_row_sum.raw());\n amgx::thrust::for_each(\n make_zip_iterator(make_tuple(A_trunc.row_indices.begin(), A_trunc.values.begin())),\n make_zip_iterator(make_tuple(A_trunc.row_indices.end(), A_trunc.values.end())),\n op\n );\n#else\n const int avg_row_len = (int)ceil(sqrt((float)A.get_num_nz() / A.get_num_rows()));\n\n if (avg_row_len > 16)\n {\n truncateAndScale, Vector, 128, 32>(A, A_trunc, row_sum, new_row_sum,\n trunc_factor, *metric);\n }\n else if (avg_row_len > 8)\n {\n truncateAndScale, Vector, 128, 16>(A, A_trunc, row_sum, new_row_sum,\n trunc_factor, *metric);\n }\n else if (avg_row_len > 4)\n {\n truncateAndScale, Vector, 128, 8>(A, A_trunc, row_sum, new_row_sum,\n trunc_factor, *metric);\n }\n else\n {\n truncateAndScale, Vector, 128, 4>(A, A_trunc, row_sum, new_row_sum,\n trunc_factor, *metric);\n }\n\n#endif\n cudaCheckError();\n // copy truncated matrix to this\n A.copy(A_trunc);\n}\n\n#ifdef USE_EXPERIMENTAL_NLARGEST\ntemplate \nvoid truncateNandScale(const Matrix &A, Matrix &A_trunc, const int max_elmts, const Vector &orig_row_sums)\n{\n const int nthreads = BLOCKSIZE;\n const int nblocks = std::min(4096, A.get_num_rows() / (nthreads / THREADS_PER_VECTOR) + 1);\n truncateNandScale_kernel\n <<< nblocks, nthreads>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n A.get_num_rows(),\n A_trunc.row_offsets.raw(),\n A_trunc.col_indices.raw(),\n A_trunc.values.raw(),\n max_elmts,\n orig_row_sums.raw()\n );\n cudaCheckError();\n}\n#endif\n\ntemplate \nvoid Truncate >::truncateByMaxElements(Matrix_d &A,\n const int max_elmts)\n{\n if (max_elmts > 32) // not supported\n {\n FatalError(\"Matrix truncation to > 32 elements not supported\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n typedef typename TConfig_d::IndPrec index_type;\n typedef typename TConfig_d::VecPrec vec_value_type;\n typedef typename TConfig_d::MatPrec value_type;\n // store the weighting values per row\n typedef Vector VVector;\n typedef typename Matrix::MVector MVector;\n typedef Vector::Type> IVector;\n IVector row_lengths(A.get_num_rows());\n MVector orig_row_sums(A.get_num_rows(), 0);\n MVector new_row_sums(A.get_num_rows(), 0);\n const int blocksize = 128;\n const int num_blocks = A.get_num_rows() / blocksize + 1;\n\n // get the original row sums of A\n row_sum_kernel <<< 4096, blocksize>>>(A.row_offsets.raw(), A.values.raw(), orig_row_sums.raw(), A.get_num_rows(), row_lengths.raw(), max_elmts, true);\n cudaCheckError();\n // initial definition\n Matrix_d A_trunc(A.get_num_rows(), A.get_num_cols(), 0, CSR); // add CSR prop\n // exclusive scan to get new row structure\n thrust_wrapper::exclusive_scan(row_lengths.begin(), row_lengths.end(), A_trunc.row_offsets.begin());\n int nnz = A_trunc.row_offsets[A.get_num_rows() - 1] + row_lengths[A.get_num_rows() - 1];\n A_trunc.row_offsets[A.get_num_rows()] = nnz;\n // set final size of truncated matrix\n A_trunc.resize(A.get_num_rows(), A.get_num_cols(), nnz);\n#ifdef USE_EXPERIMENTAL_NLARGEST\n\n // run kernel with appropriate vector size\n if (max_elmts > 16)\n {\n truncateNandScale, Vector, 128, 32>(A, A_trunc, max_elmts, orig_row_sums);\n }\n else if (max_elmts > 8)\n {\n truncateNandScale, Vector, 128, 16>(A, A_trunc, max_elmts, orig_row_sums);\n }\n else if (max_elmts > 4)\n {\n truncateNandScale, Vector, 192, 8>(A, A_trunc, max_elmts, orig_row_sums);\n }\n else\n {\n truncateNandScale, Vector, 128, 4>(A, A_trunc, max_elmts, orig_row_sums);\n }\n\n#else\n cudaCheckError();\n\n truncate_kernel <<< 4096, blocksize>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n A.get_num_rows(),\n A_trunc.row_offsets.raw(),\n A_trunc.col_indices.raw(),\n A_trunc.values.raw(),\n max_elmts\n );\n cudaCheckError();\n // get the new row sums of A_trunc\n row_sum_kernel <<< 4096, blocksize>>>(A_trunc.row_offsets.raw(), A_trunc.values.raw(), new_row_sums.raw(), A_trunc.get_num_rows(), (index_type *) NULL, 0, false);\n cudaCheckError();\n A_trunc.addProps(COO);\n scale_kernel <<< 4096, blocksize>>>(A_trunc.row_indices.raw(), A_trunc.values.raw(), new_row_sums.raw(), orig_row_sums.raw(), nnz);\n cudaCheckError();\n#endif\n cudaCheckError();\n A_trunc.delProps(COO);\n A.set_initialized(0);\n A.copy(A_trunc);\n A.computeDiagonal();\n A.set_initialized(1);\n}\n\n// -------------------------------\n// Explict instantiations\n// -------------------------------\n\n#define AMGX_CASE_LINE(CASE) template class Truncate::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n} // namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"amgx_types/util.h\"\n\n#define USE_EXPERIMENTAL_FACTOR_COPY\n// #define USE_EXPERIMENTAL_NLARGEST\n\nnamespace amgx\n{\n\n__host__ __device__ __forceinline__\ndouble ourHash2(const int i)\n{\n unsigned int a = i;\n double result;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n result = (a ^ 0x4a51e590) / (double) UINT_MAX;\n return result;\n}\n\ntemplate \nstruct copy_pred\n{\n double _trunc_factor;\n metricType *_metric_arr;\n\n copy_pred() : _trunc_factor(0), _metric_arr(NULL) {};\n copy_pred(const double trunc_factor, metricType *metric_arr) : _trunc_factor(trunc_factor),\n _metric_arr(metric_arr) {};\n\n __host__ __device__\n bool operator()(const amgx::thrust::tuple &a)\n {\n metricType metric = _metric_arr[amgx::thrust::get<0>(a)];\n\n if ( types::util::abs(amgx::thrust::get<2>(a)) >= _trunc_factor * metric) { return true; }\n\n return false;\n }\n\n copy_pred &operator=(const copy_pred &a)\n {\n this->_trunc_factor = a._trunc_factor;\n this->_metric_arr = a._metric_arr;\n return *this;\n }\n};\n\ntemplate \nstruct scale_op\n{\n const VectorType *scale_vec;\n\n scale_op(const VectorType *s) : scale_vec(s) {};\n\n __host__ __device__\n amgx::thrust::tuple operator()(const amgx::thrust::tuple &a)\n {\n const IndexType row = amgx::thrust::get<0>(a);\n return amgx::thrust::tuple(row, amgx::thrust::get<1>(a) * scale_vec[row]);\n }\n};\n\ntemplate \n__global__\nvoid scale_kernel(const index_type *row_indices, value_type *values, const value_type *new_row_sums, const value_type *old_row_sums, const int size)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n int row_index = row_indices[tidx];\n value_type multiplier = ( types::util::abs(new_row_sums[row_index]) == 0. ) ? types::util::get_one() : (old_row_sums[row_index] / new_row_sums[row_index]);\n values[tidx] = values[tidx] * multiplier;\n\n if (values[tidx] != values[tidx]) { printf(\"%d (%d): NaN after scaling\\n\", tidx, row_index); }\n }\n}\n\ntemplate \n__global__\nvoid get_row_lengths_kernel(const IndexType *offsets, const IndexType num_rows, IndexType *row_lengths,\n const int max_elmts)\n{\n for (int row = threadIdx.x + blockIdx.x * blockDim.x; row < num_rows; row += blockDim.x * gridDim.x)\n {\n row_lengths[row] = min(offsets[row + 1] - offsets[row], max_elmts);\n }\n}\n\n// sort array by abs(values), largest element in [0]\ntemplate \n__device__\nvoid sortByFabs(IndexType *indices, ValueType *values, int elements)\n{\n int n = elements;\n\n do\n {\n int newn = 0;\n\n for (int i = 1; i < n; i++)\n {\n if (types::util::abs(values[i - 1]) < types::util::abs(values[i]))\n {\n ValueType temp_v = values[i - 1];\n IndexType temp_i = indices[i - 1];\n values[i - 1] = values[i];\n values[i] = temp_v;\n indices[i - 1] = indices[i];\n indices[i] = temp_i;\n newn = i;\n }\n }\n\n n = newn;\n }\n while (n > 0);\n}\n\n#ifdef USE_EXPERIMENTAL_NLARGEST\n\ntemplate \n__device__\nvoid insert_new_max(volatile ValueType *Mv, volatile IndexType *Mi, volatile int *T, const ValueType S,\n const IndexType I, const int thread, const int N)\n{\n if (thread < N)\n {\n T[thread] = 1;\n\n // new value larger - shift my value down\n if (Mv[thread] < S)\n {\n T[thread] = 0;\n Mv[thread + 1] = Mv[thread];\n Mi[thread + 1] = Mi[thread];\n T[thread + 1] = 1;\n }\n\n // insert new largest value here\n if (!T[thread])\n {\n Mv[thread] = S;\n Mi[thread] = I;\n // inserted to me -- note to zero out\n }\n }\n}\n\ntemplate \n__global__\nvoid truncateNandScale_kernel(const IndexType *A_offsets, const IndexType *A_indices, const ValueType *A_values,\n const IndexType num_rows,\n const IndexType *At_offsets, IndexType *At_indices, ValueType *At_values,\n const int max_elmts, const VecValueType *old_row_sums)\n{\n const int WARP_SIZE = 32;\n const int VECTORS_PER_WARP = WARP_SIZE / THREADS_PER_VECTOR;\n const int WARPS_PER_BLOCK = BLOCKSIZE / WARP_SIZE;\n const int N_LARGEST_SIZE = (THREADS_PER_VECTOR + 1) * VECTORS_PER_WARP * WARPS_PER_BLOCK;\n __shared__ volatile int T[N_LARGEST_SIZE]; // has M array been modified\n __shared__ volatile ValueType Mv[N_LARGEST_SIZE]; // maximum values\n __shared__ volatile int Mi[N_LARGEST_SIZE]; // maximum indices\n __shared__ volatile ValueType S[WARP_SIZE * WARPS_PER_BLOCK]; // temporary holding for currently read values\n __shared__ volatile IndexType I[WARP_SIZE * WARPS_PER_BLOCK]; // -- \" -- indices\n ValueType S_copy;\n IndexType I_copy;\n __shared__ volatile int Ap[WARPS_PER_BLOCK * VECTORS_PER_WARP][2]; // store row offset pointers\n __shared__ volatile ValueType old_row_sum[WARPS_PER_BLOCK * VECTORS_PER_WARP];\n const int num_vectors = blockDim.x / THREADS_PER_VECTOR * gridDim.x;\n const int vector_id = (threadIdx.x + blockIdx.x * blockDim.x) / THREADS_PER_VECTOR;\n const int thread_lane = threadIdx.x & (THREADS_PER_VECTOR - 1);\n const int vector_lane = threadIdx.x / THREADS_PER_VECTOR;\n\n // initialise relevant shared memory arrays\n for (int i = threadIdx.x; i < N_LARGEST_SIZE; i += blockDim.x)\n {\n Mv[i] = 0.;\n Mi[i] = -1;\n T[i] = 0;\n }\n\n __syncthreads();\n\n // outer loop over rows\n for (int row = vector_id; row < num_rows; row += num_vectors)\n {\n // load row offsets into shared memory\n if (thread_lane < 2)\n {\n Ap[vector_lane][thread_lane] = A_offsets[row + thread_lane];\n }\n\n const int row_start = Ap[vector_lane][0];\n const int row_end = Ap[vector_lane][1];\n const int row_len = row_end - row_start;\n const int max_elmts_row = min(At_offsets[row + 1] - At_offsets[row], max_elmts);\n\n // degenerate case -- simply copy over entire row\n if (row_len <= max_elmts && row_len > 0 && thread_lane < row_len)\n {\n At_indices[At_offsets[row] + thread_lane] = A_indices[row_start + thread_lane];\n At_values[At_offsets[row] + thread_lane] = A_values[row_start + thread_lane];\n continue;\n }\n\n // loop over chunks of the row\n const int nloops = (row_len < THREADS_PER_VECTOR) ? 1 : (int)ceil((float)row_len / THREADS_PER_VECTOR);\n\n for (int i = 0; i < nloops; i++)\n {\n // zero out temp shared memory\n S[threadIdx.x] = 0;\n I[threadIdx.x] = -2;\n\n // read chunk into shared memory\n if (i * THREADS_PER_VECTOR + thread_lane < row_len)\n {\n S[threadIdx.x] = fabs(A_values[row_start + i * THREADS_PER_VECTOR + thread_lane]);\n I[threadIdx.x] = A_indices[row_start + i * THREADS_PER_VECTOR + thread_lane];\n }\n\n // loop over n' max elements\n for (int j = 0; j < max_elmts_row; j++)\n {\n // first take a copy of this section of the row so we don't re-read it\n S_copy = S[threadIdx.x];\n I_copy = I[threadIdx.x];\n // local max & associated index (save volatile loads)\n ValueType max = S_copy; // S[threadIdx.x];\n IndexType ind = I_copy; // I[threadIdx.x];\n\n // reduce to get max\n if (THREADS_PER_VECTOR > 16)\n {\n I[threadIdx.x] = ind = max > S[threadIdx.x + 16] ? ind : I[threadIdx.x + 16];\n S[threadIdx.x] = max = max > S[threadIdx.x + 16] ? max : S[threadIdx.x + 16];\n }\n\n if (THREADS_PER_VECTOR > 8)\n {\n I[threadIdx.x] = ind = max > S[threadIdx.x + 8] ? ind : I[threadIdx.x + 8];\n S[threadIdx.x] = max = max > S[threadIdx.x + 8] ? max : S[threadIdx.x + 8];\n }\n\n if (THREADS_PER_VECTOR > 4)\n {\n I[threadIdx.x] = ind = max > S[threadIdx.x + 4] ? ind : I[threadIdx.x + 4];\n S[threadIdx.x] = max = max > S[threadIdx.x + 4] ? max : S[threadIdx.x + 4];\n }\n\n if (THREADS_PER_VECTOR > 2)\n {\n I[threadIdx.x] = ind = max > S[threadIdx.x + 2] ? ind : I[threadIdx.x + 2];\n S[threadIdx.x] = max = max > S[threadIdx.x + 2] ? max : S[threadIdx.x + 2];\n }\n\n if (THREADS_PER_VECTOR > 1)\n {\n I[threadIdx.x] = ind = max > S[threadIdx.x + 1] ? ind : I[threadIdx.x + 1];\n S[threadIdx.x] = max = max > S[threadIdx.x + 1] ? max : S[threadIdx.x + 1];\n }\n\n // maximum now in S_copy[vector_lane*32] -- insert into maximum array\n const int sm_offset = vector_lane * (THREADS_PER_VECTOR + 1);\n insert_new_max(&Mv[sm_offset], &Mi[sm_offset], &T[sm_offset],\n S[vector_lane * THREADS_PER_VECTOR], I[vector_lane * THREADS_PER_VECTOR], thread_lane,\n max_elmts);\n S[threadIdx.x] = (I[vector_lane * THREADS_PER_VECTOR] == I_copy) ? 0 : S_copy;\n I[threadIdx.x] = I_copy;\n }\n }\n\n // get scaling factor for new row\n // copy Mv -> S\n S[threadIdx.x] = 0; // intialise unneeded S array to 0\n\n if (thread_lane < max_elmts_row)\n {\n S[threadIdx.x] = Mv[vector_lane * (THREADS_PER_VECTOR + 1) + thread_lane];\n }\n\n // reduce on S\n ValueType sum = fabs(S[threadIdx.x]);\n\n if (THREADS_PER_VECTOR > 16) { S[threadIdx.x] = sum = sum + S[threadIdx.x + 16]; }\n\n if (THREADS_PER_VECTOR > 8) { S[threadIdx.x] = sum = sum + S[threadIdx.x + 8]; }\n\n if (THREADS_PER_VECTOR > 4) { S[threadIdx.x] = sum = sum + S[threadIdx.x + 4]; }\n\n if (THREADS_PER_VECTOR > 2) { S[threadIdx.x] = sum = sum + S[threadIdx.x + 2]; }\n\n if (THREADS_PER_VECTOR > 1) { S[threadIdx.x] = sum = sum + S[threadIdx.x + 1]; }\n\n const ValueType scale = old_row_sums[row] / S[vector_lane * THREADS_PER_VECTOR];\n\n // copy final scaled results into output buffer\n if (thread_lane < max_elmts_row)\n {\n At_values[At_offsets[row] + thread_lane] = Mv[vector_lane * (THREADS_PER_VECTOR + 1) + thread_lane] * scale;\n At_indices[At_offsets[row] + thread_lane] = Mi[vector_lane * (THREADS_PER_VECTOR + 1) + thread_lane];\n }\n }\n}\n\n#else\n// sort array by column index\ntemplate \n__device__\nvoid sortByIndex(IndexType *indices, ValueType *values, int elements)\n{\n int n = elements;\n\n do\n {\n int newn = 0;\n\n for (int i = 1; i < n; i++)\n {\n if (indices[i - 1] > indices[i])\n {\n ValueType temp_v = values[i - 1];\n IndexType temp_i = indices[i - 1];\n values[i - 1] = values[i];\n values[i] = temp_v;\n indices[i - 1] = indices[i];\n indices[i] = temp_i;\n newn = i;\n }\n }\n\n n = newn;\n }\n while (n > 0);\n}\n\ntemplate \n__global__\nvoid row_sum_kernel(const IndexType *A_offsets, const ValueType *A_values,\n ValueType *row_sums, const IndexType num_rows, IndexType *row_lengths, const int max_elmts, bool compute_row_length)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_rows; tidx += blockDim.x * gridDim.x)\n {\n IndexType A_start = A_offsets[tidx], A_end = A_offsets[tidx + 1];\n\n if (compute_row_length)\n {\n row_lengths[tidx] = min(A_end - A_start, max_elmts);\n }\n\n ValueType row_sum = 0.;\n bool found_nans = false;\n\n for (int i = A_start; i < A_end; i++)\n {\n row_sum += A_values[i];\n\n if (A_values[i] != A_values[i]) { found_nans = true; }\n }\n\n if (found_nans)\n {\n printf(\"row %d has NaNs in input\\n\", tidx);\n }\n\n row_sums[tidx] = row_sum;\n }\n}\n\ntemplate \n__global__\nvoid __launch_bounds__(256)\ntruncate_kernel(const IndexType *A_offsets, const IndexType *A_col_indices, const ValueType *A_values,\n const IndexType num_rows,\n IndexType *trunc_offsets, IndexType *trunc_col_indices, ValueType *trunc_values,\n const int max_elmts)\n{\n // const int tidx = threadIdx.x + blockIdx.x*blockDim.x;\n __shared__ IndexType s_ind[1024]; // assume <= 256 threads per block\n __shared__ ValueType s_val[1024];\n const int smem_start = threadIdx.x * max_elmts;\n\n // if (tidx < num_rows)\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_rows; tidx += blockDim.x * gridDim.x)\n {\n IndexType A_start = A_offsets[tidx], A_end = A_offsets[tidx + 1];\n IndexType trunc_start = trunc_offsets[tidx];\n IndexType A_len = A_end - A_start;\n\n // copy my entire row if possible\n if (A_len <= max_elmts)\n {\n for (int j = 0; j < A_len; j++)\n {\n trunc_col_indices[trunc_start + j] = A_col_indices[A_start + j];\n trunc_values[trunc_start + j] = A_values[A_start + j];\n }\n }\n else // normal case\n {\n // read first max_elmts elements into shared\n for (int j = 0; j < max_elmts; j++)\n {\n s_ind[smem_start + j] = A_col_indices[A_start + j];\n s_val[smem_start + j] = A_values[A_start + j];\n }\n\n // sort the initially inserted elements (start = largest)\n sortByFabs(&s_ind[smem_start], &s_val[smem_start], max_elmts);\n\n // loop over all other elements on the row\n for (int j = A_start + max_elmts; j < A_end; j++)\n {\n // check against current maximums\n for (int i = 0; i < max_elmts; i++)\n {\n if (types::util::abs(A_values[j]) > types::util::abs(s_val[smem_start + i]))\n {\n // shift remaining values down\n for (int k = max_elmts - 1; k > i; k--)\n {\n s_val[smem_start + k] = s_val[smem_start + k - 1];\n s_ind[smem_start + k] = s_ind[smem_start + k - 1];\n }\n\n // tentatively replace element\n s_val[smem_start + i] = A_values[j];\n s_ind[smem_start + i] = A_col_indices[j];\n break;\n }\n }\n }\n\n //sortByIndex(&s_ind[smem_start],&s_val[smem_start],max_elmts);\n // now we have the 4 largest elements (hopefully) -- output\n for (int i = 0; i < max_elmts; i++)\n {\n trunc_col_indices[trunc_start + i] = s_ind[smem_start + i];\n trunc_values[trunc_start + i] = s_val[smem_start + i];\n }\n }\n }\n}\n#endif\n\ntemplate \n__global__\nvoid perturb_kernel(const IndexType *A_offsets, const IndexType *A_col_indices, ValueType *A_values,\n const IndexType num_rows,\n const int max_elmts)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_rows; tidx += blockDim.x * gridDim.x)\n {\n IndexType A_start = A_offsets[tidx], A_end = A_offsets[tidx + 1];\n\n // perturb A slightly\n for (int j = A_start; j < A_end; j++)\n {\n A_values[j] += ourHash2(A_col_indices[j]) * A_values[j];\n }\n }\n}\n\n\n// warp-based inclusive scan\n// from: Sengupta, Harris, Garland 2008\n// idx is index of thread\ntemplate \n__device__\nint scan_warp(volatile IndexType *ptr, const unsigned int idx)\n{\n const unsigned int lane = idx & 31; // index of thread in warp\n IndexType p = ptr[idx];\n\n if (lane >= 1) { ptr[idx] = p = ptr[idx - 1] + p; }\n\n if (lane >= 2) { ptr[idx] = p = ptr[idx - 2] + p; }\n\n if (lane >= 4) { ptr[idx] = p = ptr[idx - 4] + p; }\n\n if (lane >= 8) { ptr[idx] = p = ptr[idx - 8] + p; }\n\n if (lane >= 16) { ptr[idx] = p = ptr[idx - 16] + p; }\n\n return ptr[idx];\n}\n\ntemplate \n__global__\nvoid truncateAndScale_kernel(const IndexType *A_offsets, const IndexType *A_indices,\n const matValueType *A_values,\n const IndexType num_rows,\n IndexType *At_offsets, IndexType *At_indices, matValueType *At_values,\n const vecValueType *old_row_sums, const vecValueType *new_row_sums,\n const double truncate_factor, const vecValueType *metric)\n{\n const int WARP_SIZE = 32;\n const int WARPS_PER_BLOCK = THREADS_PER_BLOCK / WARP_SIZE;\n const int VECTORS_PER_WARP = WARP_SIZE / THREADS_PER_VECTOR;\n // allocate shared memory for warp-based scans (assume up to 512 threads / block)\n __shared__ volatile int s_scan[WARP_SIZE * WARPS_PER_BLOCK];\n __shared__ volatile int s_A_ptr[WARPS_PER_BLOCK * VECTORS_PER_WARP][2]; // row offsets for A\n __shared__ volatile int s_At_ptr[WARPS_PER_BLOCK * VECTORS_PER_WARP][2]; // row offsets for At\n __shared__ double s_scale[WARPS_PER_BLOCK * VECTORS_PER_WARP]; // scaling factor for each warp\n const int vectors_per_block = THREADS_PER_BLOCK / WARP_SIZE * VECTORS_PER_WARP;\n const int num_vectors = THREADS_PER_BLOCK / THREADS_PER_VECTOR * gridDim.x;\n const int vector_id = (threadIdx.x + blockIdx.x * blockDim.x) / THREADS_PER_VECTOR;\n const int thread_lane = threadIdx.x & (THREADS_PER_VECTOR - 1);\n const int vector_lane = threadIdx.x / THREADS_PER_VECTOR;\n int scan_offset = 0;\n matValueType val = 0., fval = 0.;\n IndexType ind = 0;\n\n // outer loop over rows\n for (int row = vector_id; row < num_rows; row += num_vectors)\n {\n scan_offset = 0;\n val = 0.;\n fval = 0.; // fabs(val)\n\n if (threadIdx.x < vectors_per_block && threadIdx.x < num_rows - row) // first threads in first warp\n {\n s_scale[threadIdx.x] = old_row_sums[row + threadIdx.x] / new_row_sums[row + threadIdx.x];\n }\n\n __syncthreads();\n\n if (thread_lane < 2)\n {\n s_A_ptr[vector_lane][thread_lane] = A_offsets[row + thread_lane];\n s_At_ptr[vector_lane][thread_lane] = At_offsets[row + thread_lane];\n }\n\n const int row_start = s_A_ptr[vector_lane][0];\n const int row_end = s_A_ptr[vector_lane][1];\n const int row_len = row_end - row_start;\n const int At_row_start = s_At_ptr[vector_lane][0];\n // mark & scan over consecutive blocks of 32 threads\n int nloops = (int)ceilf((float)row_len / THREADS_PER_VECTOR);\n#pragma unroll 2\n\n for (int i = 0; i < nloops; i++)\n {\n // initialise shared memory to 0\n s_scan[threadIdx.x] = 0;\n\n // mark appropriate elements of the row\n if (i * THREADS_PER_VECTOR + thread_lane < row_len)\n {\n val = A_values[row_start + i * THREADS_PER_VECTOR + thread_lane];\n ind = A_indices[row_start + i * THREADS_PER_VECTOR + thread_lane];\n fval = fabs(val);\n s_scan[threadIdx.x] = fval >= truncate_factor * metric[row] ? 1 : 0;\n }\n\n // now prefix scan over this\n scan_warp(&s_scan[vector_lane * THREADS_PER_VECTOR], thread_lane); // threadIdx.x%32);\n\n // copy relevant parts of A to At\n if (i * THREADS_PER_VECTOR + thread_lane < row_len)\n {\n if (fval >= truncate_factor * metric[row])\n {\n // get relevant index in At\n const int At_index = At_row_start + scan_offset + s_scan[threadIdx.x] - 1;\n\n if (s_scale[vector_lane] == 0.)\n {\n At_values[At_index] = val;\n }\n else\n {\n At_values[At_index] = val / s_scale[vector_lane];\n }\n\n At_indices[At_index] = ind;\n }\n }\n\n // note new scan_offset\n scan_offset = s_scan[vector_lane * THREADS_PER_VECTOR + THREADS_PER_VECTOR - 1];\n }\n }\n}\n\ntemplate \nvoid truncateAndScale(const Matrix &A, Matrix &A_trunc, const Vector &row_sum, const Vector &new_row_sum,\n const double trunc_factor, const Vector &metric)\n{\n const int num_blocks = std::min(4096, A.get_num_rows());\n truncateAndScale_kernel\n <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(),\n A.get_num_rows(),\n A_trunc.row_offsets.raw(), A_trunc.col_indices.raw(),\n A_trunc.values.raw(),\n row_sum.raw(), new_row_sum.raw(),\n trunc_factor, metric.raw());\n cudaCheckError();\n}\n\n// host code\ntemplate \nvoid Truncate >::truncateByFactor(Matrix_h &A,\n const double trunc_factor, const AMGX_TruncateType truncType)\n{\n // throw error\n FatalError(\"Truncation on host not implemented yet\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid Truncate >::truncateByMaxElements(Matrix_h &A,\n const int max_elmts)\n{\n // throw error\n FatalError(\"Truncation on host not implemented yet\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n// device code\n#ifdef _WIN32\ntemplate \nvoid Truncate >::truncateByFactor(Matrix_d &A, const double trunc_factor, const AMGX_TruncateType truncType)\n#else\ntemplate \nvoid Truncate >::truncateByFactor(Matrix_d &A, const double trunc_factor, const AMGX_TruncateType truncType) __attribute__((noinline))\n#endif\n{\n typedef typename TConfig_d::IndPrec index_type;\n typedef typename TConfig_d::VecPrec vec_value_type;\n typedef typename TConfig_d::MatPrec value_type;\n // store the weighting values per row\n Matrix A_trunc(0, 0, 0, CSR);\n typedef Vector VVector;\n typedef typename Matrix::MVector MVector;\n typedef Vector::Type> IVector;\n VVector row_sum(A.get_num_rows());\n VVector new_row_sum(A.get_num_rows());\n VVector max_coef;\n // point to either row_sum or max_coef depending on type of truncation desired\n // vec_value_type *metric = 0;\n VVector *metric;\n\n // get either max row coefficient or row sum\n if (truncType == AMGX_TruncateByRowSum)\n {\n absRowSum(A, row_sum);\n metric = &row_sum;\n }\n else if (truncType == AMGX_TruncateByMaxCoefficient)\n {\n max_coef.resize(A.get_num_rows());\n maxCoefAndSum(A, max_coef, row_sum);\n metric = &max_coef;\n }\n else\n {\n FatalError(\"Truncation type not implemented\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n IVector row_counts(A.get_num_rows());\n\n if (truncType == AMGX_TruncateByRowSum)\n {\n countTruncElements(A, trunc_factor, row_sum, row_counts, new_row_sum);\n }\n else if (truncType == AMGX_TruncateByMaxCoefficient)\n {\n countTruncElements(A, trunc_factor, max_coef, row_counts, new_row_sum);\n }\n\n cudaCheckError();\n // initial resize (so we can scan into the row_offsets array)\n A_trunc.resize(A.get_num_rows(), A.get_num_cols(), 0);\n cudaCheckError();\n\n thrust_wrapper::exclusive_scan(row_counts.begin(), row_counts.end(), A_trunc.row_offsets.begin());\n cudaCheckError();\n \n const int nnz = A_trunc.row_offsets[A.get_num_rows() - 1] + row_counts[A.get_num_rows() - 1];\n A_trunc.row_offsets[A.get_num_rows()] = nnz;\n\n if (nnz == A.get_num_nz()) // early return -- nothing truncated\n {\n return;\n }\n\n // final resize & get row indices for both matrices\n A_trunc.resize(A.get_num_rows(), A.get_num_cols(), nnz);\n#ifndef USE_EXPERIMENTAL_FACTOR_COPY\n A_trunc.addProps(COO);\n A.addProps(COO);\n // copy relevant values from A -> A_trunc\n copy_pred pred;\n\n if (truncType == AMGX_TruncateByRowSum)\n {\n pred = copy_pred(trunc_factor, row_sum.raw());\n }\n else if (truncType == AMGX_TruncateByMaxCoefficient)\n {\n pred = copy_pred(trunc_factor, max_coef.raw());\n }\n\n using amgx::thrust::make_zip_iterator;\n using amgx::thrust::make_tuple;\n amgx::thrust::copy_if(\n make_zip_iterator(make_tuple(A.row_indices.begin(), A.col_indices.begin(), A.values.begin())),\n make_zip_iterator(make_tuple(A.row_indices.end(), A.col_indices.end(), A.values.end())),\n make_zip_iterator(make_tuple(A_trunc.row_indices.begin(), A_trunc.col_indices.begin(),\n A_trunc.values.begin())),\n pred\n );\n // scale each row of the truncated matrix - per-row scale in new_row_sum\n amgx::thrust::transform(row_sum.begin(), row_sum.end(), new_row_sum.begin(), new_row_sum.begin(),\n amgx::thrust::divides());\n scale_op op(new_row_sum.raw());\n amgx::thrust::for_each(\n make_zip_iterator(make_tuple(A_trunc.row_indices.begin(), A_trunc.values.begin())),\n make_zip_iterator(make_tuple(A_trunc.row_indices.end(), A_trunc.values.end())),\n op\n );\n#else\n const int avg_row_len = (int)ceil(sqrt((float)A.get_num_nz() / A.get_num_rows()));\n\n if (avg_row_len > 16)\n {\n truncateAndScale, Vector, 128, 32>(A, A_trunc, row_sum, new_row_sum,\n trunc_factor, *metric);\n }\n else if (avg_row_len > 8)\n {\n truncateAndScale, Vector, 128, 16>(A, A_trunc, row_sum, new_row_sum,\n trunc_factor, *metric);\n }\n else if (avg_row_len > 4)\n {\n truncateAndScale, Vector, 128, 8>(A, A_trunc, row_sum, new_row_sum,\n trunc_factor, *metric);\n }\n else\n {\n truncateAndScale, Vector, 128, 4>(A, A_trunc, row_sum, new_row_sum,\n trunc_factor, *metric);\n }\n\n#endif\n cudaCheckError();\n // copy truncated matrix to this\n A.copy(A_trunc);\n}\n\n#ifdef USE_EXPERIMENTAL_NLARGEST\ntemplate \nvoid truncateNandScale(const Matrix &A, Matrix &A_trunc, const int max_elmts, const Vector &orig_row_sums)\n{\n const int nthreads = BLOCKSIZE;\n const int nblocks = std::min(4096, A.get_num_rows() / (nthreads / THREADS_PER_VECTOR) + 1);\n truncateNandScale_kernel\n <<< nblocks, nthreads>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n A.get_num_rows(),\n A_trunc.row_offsets.raw(),\n A_trunc.col_indices.raw(),\n A_trunc.values.raw(),\n max_elmts,\n orig_row_sums.raw()\n );\n cudaCheckError();\n}\n#endif\n\ntemplate \nvoid Truncate >::truncateByMaxElements(Matrix_d &A,\n const int max_elmts)\n{\n if (max_elmts > 32) // not supported\n {\n FatalError(\"Matrix truncation to > 32 elements not supported\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n typedef typename TConfig_d::IndPrec index_type;\n typedef typename TConfig_d::VecPrec vec_value_type;\n typedef typename TConfig_d::MatPrec value_type;\n // store the weighting values per row\n typedef Vector VVector;\n typedef typename Matrix::MVector MVector;\n typedef Vector::Type> IVector;\n IVector row_lengths(A.get_num_rows());\n MVector orig_row_sums(A.get_num_rows(), 0);\n MVector new_row_sums(A.get_num_rows(), 0);\n const int blocksize = 128;\n const int num_blocks = A.get_num_rows() / blocksize + 1;\n\n // get the original row sums of A\n row_sum_kernel <<< 4096, blocksize>>>(A.row_offsets.raw(), A.values.raw(), orig_row_sums.raw(), A.get_num_rows(), row_lengths.raw(), max_elmts, true);\n cudaCheckError();\n // initial definition\n Matrix_d A_trunc(A.get_num_rows(), A.get_num_cols(), 0, CSR); // add CSR prop\n // exclusive scan to get new row structure\n thrust_wrapper::exclusive_scan(row_lengths.begin(), row_lengths.end(), A_trunc.row_offsets.begin());\n int nnz = A_trunc.row_offsets[A.get_num_rows() - 1] + row_lengths[A.get_num_rows() - 1];\n A_trunc.row_offsets[A.get_num_rows()] = nnz;\n // set final size of truncated matrix\n A_trunc.resize(A.get_num_rows(), A.get_num_cols(), nnz);\n#ifdef USE_EXPERIMENTAL_NLARGEST\n\n // run kernel with appropriate vector size\n if (max_elmts > 16)\n {\n truncateNandScale, Vector, 128, 32>(A, A_trunc, max_elmts, orig_row_sums);\n }\n else if (max_elmts > 8)\n {\n truncateNandScale, Vector, 128, 16>(A, A_trunc, max_elmts, orig_row_sums);\n }\n else if (max_elmts > 4)\n {\n truncateNandScale, Vector, 192, 8>(A, A_trunc, max_elmts, orig_row_sums);\n }\n else\n {\n truncateNandScale, Vector, 128, 4>(A, A_trunc, max_elmts, orig_row_sums);\n }\n\n#else\n cudaCheckError();\n\n truncate_kernel <<< 4096, blocksize>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n A.get_num_rows(),\n A_trunc.row_offsets.raw(),\n A_trunc.col_indices.raw(),\n A_trunc.values.raw(),\n max_elmts\n );\n cudaCheckError();\n // get the new row sums of A_trunc\n row_sum_kernel <<< 4096, blocksize>>>(A_trunc.row_offsets.raw(), A_trunc.values.raw(), new_row_sums.raw(), A_trunc.get_num_rows(), (index_type *) NULL, 0, false);\n cudaCheckError();\n A_trunc.addProps(COO);\n scale_kernel <<< 4096, blocksize>>>(A_trunc.row_indices.raw(), A_trunc.values.raw(), new_row_sums.raw(), orig_row_sums.raw(), nnz);\n cudaCheckError();\n#endif\n cudaCheckError();\n A_trunc.delProps(COO);\n A.set_initialized(0);\n A.copy(A_trunc);\n A.computeDiagonal();\n A.set_initialized(1);\n}\n\n// -------------------------------\n// Explict instantiations\n// -------------------------------\n\n#define AMGX_CASE_LINE(CASE) template class Truncate::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n} // namespace amgx\n\n", "source": "/kaggle/working/amgx/src/truncate.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n\n__global__ void profileLevelUp_kernel() {}\n__global__ void profileLevelDown_kernel() {}\n__global__ void profileLevelZero_kernel() {}\n__global__ void profilePhaseSetup_kernel() {}\n__global__ void profilePhaseSolve_kernel() {}\n__global__ void profilePhaseNone_kernel() {}\n__global__ void profileSubphaseMatrixColoring_kernel() {}\n__global__ void profileSubphaseSmootherSetup_kernel() {}\n__global__ void profileSubphaseFindAggregates_kernel() {}\n__global__ void profileSubphaseComputeRestriction_kernel() {}\n__global__ void profileSubphaseComputeCoarseA_kernel() {}\n__global__ void profileSubphaseNone_kernel() {}\n__global__ void profileSubphaseTruncateP_kernel() {}\n\nnamespace amgx\n{\n\nvoid profileLevelUp()\n{\n profileLevelUp_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\nvoid profileLevelDown()\n{\n profileLevelDown_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\nvoid profileLevelZero()\n{\n profileLevelZero_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\nvoid profilePhaseSetup()\n{\n profilePhaseSetup_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\nvoid profilePhaseSolve()\n{\n profilePhaseSolve_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\nvoid profilePhaseNone()\n{\n profilePhaseNone_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\nvoid profileSubphaseMatrixColoring()\n{\n profileSubphaseMatrixColoring_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\nvoid profileSubphaseSmootherSetup()\n{\n profileSubphaseSmootherSetup_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\nvoid profileSubphaseFindAggregates()\n{\n profileSubphaseFindAggregates_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\nvoid profileSubphaseComputeRestriction()\n{\n profileSubphaseComputeRestriction_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\nvoid profileSubphaseComputeCoarseA()\n{\n profileSubphaseComputeCoarseA_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\nvoid profileSubphaseNone()\n{\n profileSubphaseNone_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\nvoid profileSubphaseTruncateP()\n{\n profileSubphaseTruncateP_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\n} // namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n__global__ void profileLevelUp_kernel() {}\n__global__ void profileLevelDown_kernel() {}\n__global__ void profileLevelZero_kernel() {}\n__global__ void profilePhaseSetup_kernel() {}\n__global__ void profilePhaseSolve_kernel() {}\n__global__ void profilePhaseNone_kernel() {}\n__global__ void profileSubphaseMatrixColoring_kernel() {}\n__global__ void profileSubphaseSmootherSetup_kernel() {}\n__global__ void profileSubphaseFindAggregates_kernel() {}\n__global__ void profileSubphaseComputeRestriction_kernel() {}\n__global__ void profileSubphaseComputeCoarseA_kernel() {}\n__global__ void profileSubphaseNone_kernel() {}\n__global__ void profileSubphaseTruncateP_kernel() {}\n\nnamespace amgx\n{\n\nvoid profileLevelUp()\n{\n profileLevelUp_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\nvoid profileLevelDown()\n{\n profileLevelDown_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\nvoid profileLevelZero()\n{\n profileLevelZero_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\nvoid profilePhaseSetup()\n{\n profilePhaseSetup_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\nvoid profilePhaseSolve()\n{\n profilePhaseSolve_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\nvoid profilePhaseNone()\n{\n profilePhaseNone_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\nvoid profileSubphaseMatrixColoring()\n{\n profileSubphaseMatrixColoring_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\nvoid profileSubphaseSmootherSetup()\n{\n profileSubphaseSmootherSetup_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\nvoid profileSubphaseFindAggregates()\n{\n profileSubphaseFindAggregates_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\nvoid profileSubphaseComputeRestriction()\n{\n profileSubphaseComputeRestriction_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\nvoid profileSubphaseComputeCoarseA()\n{\n profileSubphaseComputeCoarseA_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\nvoid profileSubphaseNone()\n{\n profileSubphaseNone_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\nvoid profileSubphaseTruncateP()\n{\n profileSubphaseTruncateP_kernel <<< 1, 1, 0, thrust::global_thread_handle::get_stream()>>>();\n cudaCheckError();\n}\n\n} // namespace amgx\n\n", "source": "/kaggle/working/amgx/src/profile.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \n\n#if CUDART_VERSION < 11000\n#define CUSPARSE_SPMM_ALG_DEFAULT CUSPARSE_MM_ALG_DEFAULT\n#endif\n\n#if CUDART_VERSION >= 12000\n#define CUSPARSE_CSRMV_ALG1 CUSPARSE_SPMV_CSR_ALG1\n#define CUSPARSE_CSRMV_ALG2 CUSPARSE_SPMV_CSR_ALG2\n#endif\n\nnamespace amgx\n{\n\nCusparse::Cusparse() : m_handle(0), m_determinism_flag(false)\n{\n cusparseCheckError( cusparseCreate(&m_handle) );\n}\n\nCusparse::~Cusparse()\n{\n destroy_handle();\n}\n\nCusparse &Cusparse::get_instance()\n{\n static Cusparse s_instance;\n s_instance.create_handle();\n return s_instance;\n}\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv(\n const typename TConfig::VecPrec alphaConst,\n Matrix &A,\n Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view )\n{\n cudaStream_t null_stream = 0;\n\n // If only COO, add CSR since bsrmv doesn't support COO\n if (A.hasProps(COO) && !A.hasProps(CSR))\n {\n A.set_initialized(0);\n A.addProps(CSR);\n A.computeDiagonal();\n A.set_initialized(1);\n }\n\n // Handle cases where the view is set by the calling routine\n if(view != A.getViewExterior())\n {\n bsrmv_internal(alphaConst, A, x, betaConst, y, view, null_stream);\n return;\n }\n\n bool latencyHiding = (A.getViewInterior() != A.getViewExterior() && !A.is_matrix_singleGPU() && x.dirtybit != 0);\n\n if (latencyHiding)\n {\n A.manager->exchange_halo_split_gather(x, x.tag);\n\n // Multiply interior rows\n bsrmv_internal(alphaConst, A, x, betaConst, y, A.getViewInterior(), null_stream);\n\n // Finish halo exchange\n A.manager->exchange_halo_split_finish(x, x.tag);\n\n // Multiply rows with halo dependencies\n ViewType bnd_view = (ViewType)(~(A.getViewInterior()) & A.getViewExterior());\n bsrmv_internal(alphaConst, A, x, betaConst, y, bnd_view, null_stream);\n }\n else\n {\n if (!A.is_matrix_singleGPU() && x.dirtybit != 0)\n {\n A.manager->exchange_halo_v2(x, x.tag);\n }\n\n bsrmv_internal(alphaConst, A, x, betaConst, y, A.getViewExterior(), null_stream);\n }\n\n y.dirtybit = 1;\n}\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_with_mask(\n const typename TConfig::VecPrec alphaConst,\n Matrix &A,\n Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y)\n{\n cudaStream_t null_stream = 0;\n\n // If only COO, add CSR since bsrmv doesn't support COO\n if (A.hasProps(COO) && !A.hasProps(CSR))\n {\n A.set_initialized(0);\n A.addProps(CSR);\n A.computeDiagonal();\n A.set_initialized(1);\n }\n\n bool latencyHiding = (A.getViewInterior() != A.getViewExterior() && !A.is_matrix_singleGPU() && x.dirtybit != 0);\n\n if (latencyHiding)\n {\n A.manager->exchange_halo_split_gather(x, x.tag);\n\n // Multiply interior\n bsrmv_internal_with_mask(alphaConst, A, x, betaConst, y, INTERIOR, null_stream);\n\n A.manager->exchange_halo_split_finish(x, x.tag);\n\n // Multiply exterior\n bsrmv_internal_with_mask(alphaConst, A, x, betaConst, y, BOUNDARY, null_stream);\n }\n else\n {\n if (!A.is_matrix_singleGPU() && x.dirtybit != 0)\n {\n A.manager->exchange_halo_v2(x, x.tag);\n }\n\n bsrmv_internal(alphaConst, A, x, betaConst, y, OWNED, null_stream);\n }\n\n y.dirtybit = 1;\n}\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_with_mask_restriction(\n const typename TConfig::VecPrec alphaConst,\n Matrix &R,\n Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n Matrix &P)\n{\n // If only COO, add CSR since bsrmv doesn't support COO\n //if (A.hasProps(COO) && !A.hasProps(CSR))\n //{\n // A.set_initialized(0);\n // A.addProps(CSR);\n // A.computeDiagonal();\n // A.set_initialized(1);\n //}\n\n bool latencyHiding = (R.getViewInterior() != R.getViewExterior() && !P.is_matrix_singleGPU() && x.dirtybit != 0);\n\n if (latencyHiding)\n {\n cudaStream_t null_stream = 0;\n bsrmv_internal_with_mask_restriction(alphaConst, R, x, betaConst, y, HALO1, null_stream, P);\n P.manager->add_from_halo_split_gather(y, y.tag);\n cudaError_t cuda_rc = cudaEventRecord(P.manager->get_comm_event());\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaEventRecord failed\", AMGX_ERR_CUDA_FAILURE);\n }\n bsrmv_internal_with_mask_restriction(alphaConst, R, x, betaConst, y, OWNED, null_stream, P);\n\n if (P.manager->neighbors.size() != 0)\n {\n cuda_rc = cudaEventSynchronize(P.manager->get_comm_event());\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaEventSynchronize failed\", AMGX_ERR_CUDA_FAILURE);\n }\n P.manager->add_from_halo_split_finish(y, y.tag, P.manager->get_bdy_stream());\n cuda_rc = cudaStreamSynchronize(P.manager->get_bdy_stream());\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaStreamSynchronize failed\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n }\n else\n {\n bsrmv_internal(alphaConst, R, x, betaConst, y, OWNED, 0);\n\n // Add contribution from neighbors\n P.manager->add_from_halo_v2(y, y.tag);\n }\n\n y.dirtybit = 1;\n}\n\ntemplate< class TConfig >\nvoid Cusparse::bsrxmv(\n const typename TConfig::VecPrec alphaConst,\n Matrix &A,\n Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n typename Matrix::IVector &mask,\n ViewType view )\n{\n // If only COO, add CSR since bsrmv doesn't support COO\n if (A.hasProps(COO) && !A.hasProps(CSR))\n {\n A.set_initialized(0);\n A.addProps(CSR);\n A.computeDiagonal();\n A.set_initialized(1);\n }\n\n const int *start_offsets, *end_offsets;\n start_offsets = A.row_offsets.raw();\n end_offsets = A.row_offsets.raw() + 1;\n cusparseDirection_t direction = CUSPARSE_DIRECTION_COLUMN;\n\n if ( A.getBlockFormat() == ROW_MAJOR )\n {\n direction = CUSPARSE_DIRECTION_ROW;\n }\n\n if (view == OWNED) //This is already a view, thus do not even attempt to do latency hiding\n {\n bsrxmv_internal(Cusparse::get_instance().m_handle, direction, CUSPARSE_OPERATION_NON_TRANSPOSE,\n mask.size(),\n A.get_num_rows(), A.get_num_cols(), A.get_num_nz(), &alphaConst,\n A.cuMatDescr,\n A.values.raw(),\n mask.raw(),\n start_offsets, end_offsets,\n A.col_indices.raw(),\n A.get_block_dimx(),\n x.raw(),\n &betaConst,\n y.raw());\n }\n else //Try and do latency hiding\n {\n // latency hiding?\n FatalError(\"Trying to do latency hiding in the bsrxmv\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n// E is a vector that represents a diagonal matrix\n// operate on all rows and columns\n// y= alpha*E.x + beta*y\ntemplate< class TConfig >\nvoid Cusparse::bsrmv( const typename TConfig::VecPrec alphaConst,\n Matrix &A,\n const typename Matrix::MVector &E,\n Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view )\n{\n cudaStream_t null_stream = 0;\n\n // If only COO, add CSR since bsrmv doesn't support COO\n if (A.hasProps(COO) && !A.hasProps(CSR))\n {\n A.addProps(CSR);\n }\n\n if (view != A.getViewExterior()) //This is already a view, thus do not even attempt to do latency hiding\n {\n bsrmv_internal(alphaConst, A, E, x, betaConst, y, view, null_stream);\n }\n else //Try and do latency hiding\n {\n ViewType oldView = A.currentView();\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_async(x, x.tag);\n }\n\n if (A.getViewExterior() == A.getViewInterior())\n {\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_wait(x, x.tag);\n }\n }\n\n ViewType flags;\n bool latencyHiding = true;\n\n //if (A.manager->num_neighbors() == 0 || (x.dirtybit == 0)) {\n if (A.is_matrix_singleGPU() || (x.dirtybit == 0))\n {\n latencyHiding = false;\n A.setViewExterior();\n flags = (ViewType)(A.getViewExterior());\n }\n else\n {\n flags = (ViewType)(A.getViewInterior());\n A.setViewInterior();\n }\n\n if (latencyHiding)\n {\n bsrmv_internal(alphaConst, A, E, x, betaConst, y, flags, null_stream);\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_wait(x, x.tag);\n }\n\n A.setViewExterior();\n flags = (ViewType)(~(A.getViewInterior()) & A.getViewExterior());\n\n if (flags != 0)\n {\n bsrmv_internal(alphaConst, A, E, x, betaConst, y, flags, null_stream);\n }\n }\n else\n {\n bsrmv_internal(alphaConst, A, E, x, betaConst, y, flags, null_stream);\n }\n\n y.dirtybit = 1;\n //if (!A.is_matrix_singleGPU())\n // if (y.size() == x.size() && y.delayed_send==0) A.manager->exchange_halo_async(y, y.tag);\n A.setView(oldView);\n }\n}\n\n// operate only on columns specified by columnColorSelector, see enum ColumnColorSelector above\n// operate only on rows of specified color, given by A.offsets_rows_per_color, A.sorted_rows_by_color\n// y= alpha*A.x + beta*y\ntemplate< class TConfig >\nvoid Cusparse::bsrmv( ColumnColorSelector columnColorSelector,\n const int color,\n const typename TConfig::VecPrec alphaConst,\n Matrix &A,\n Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view )\n{\n cudaStream_t null_stream = 0;\n\n // If only COO, add CSR since bsrmv doesn't support COO\n if (A.hasProps(COO) && !A.hasProps(CSR))\n {\n A.addProps(CSR);\n }\n\n if (view != A.getViewExterior()) //This is already a view, thus do not even attempt to do latency hiding\n {\n // std::cerr << \"exterior view with selector\" << std::endl;\n bsrmv_internal(columnColorSelector, color, alphaConst, A, x, betaConst, y, view, null_stream);\n }\n else\n {\n //Try and do latency hiding\n ViewType oldView = A.currentView();\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_async(x, x.tag);\n }\n\n if (A.getViewExterior() == A.getViewInterior())\n {\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_wait(x, x.tag);\n }\n }\n\n ViewType flags;\n bool latencyHiding = true;\n\n if (A.is_matrix_singleGPU() || (x.dirtybit == 0))\n {\n latencyHiding = false;\n A.setViewExterior();\n flags = (ViewType)(A.getViewExterior());\n }\n else\n {\n flags = (ViewType)(A.getViewInterior());\n A.setViewInterior();\n }\n\n if (latencyHiding)\n {\n bsrmv_internal(columnColorSelector, color, alphaConst, A, x, betaConst, y, flags, null_stream);\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_wait(x, x.tag);\n }\n\n A.setViewExterior();\n flags = (ViewType)(~(A.getViewInterior()) & A.getViewExterior());\n\n if (flags != 0)\n {\n bsrmv_internal(columnColorSelector, color, alphaConst, A, x, betaConst, y, flags, null_stream);\n }\n }\n else\n {\n bsrmv_internal(columnColorSelector, color, alphaConst, A, x, betaConst, y, flags, null_stream);\n }\n\n y.dirtybit = 1;\n //if (!A.is_matrix_singleGPU() && y.size() == x.size() && y.delayed_send==0)\n // A.manager->exchange_halo_async(y, y.tag);\n A.setView(oldView);\n }\n}\n\n// E is a vector that represents a diagonal matrix\n// operate only on rows of specified color, given by A.offsets_rows_per_color, A.sorted_rows_by_color\n// y= alpha*E.x + beta*y\ntemplate< class TConfig >\nvoid Cusparse::bsrmv( const int color,\n const typename TConfig::VecPrec alphaConst,\n Matrix &A,\n const typename Matrix::MVector &E,\n Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view)\n{\n cudaStream_t null_stream = 0;\n\n // If only COO, add CSR since bsrmv doesn't support COO\n if (A.hasProps(COO) && !A.hasProps(CSR))\n {\n A.addProps(CSR);\n }\n\n if (view != A.getViewExterior()) //This is already a view, thus do not even attempt to do latency hiding\n {\n bsrmv_internal(color, alphaConst, A, E, x, betaConst, y, view, null_stream);\n }\n else //Try and do latency hiding\n {\n //std::ccooor << \"de\" << std::endl;\n //std::cerr << \"not an exterior view\" << std::endl;\n ViewType oldView = A.currentView();\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_async(x, x.tag);\n }\n\n if (A.getViewExterior() == A.getViewInterior())\n {\n //std::cerr << \"exchange_halo_wait\" << std::endl;\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_wait(x, x.tag);\n }\n }\n\n //std::cerr << \"xxxeded\" << std::endl;\n ViewType flags;\n bool latencyHiding = true;\n\n if (A.is_matrix_singleGPU() || (x.dirtybit == 0))\n {\n latencyHiding = false;\n A.setViewExterior();\n flags = (ViewType)(A.getViewExterior());\n }\n else\n {\n flags = (ViewType)(A.getViewInterior());\n A.setViewInterior();\n }\n\n if (latencyHiding)\n {\n bsrmv_internal(color, alphaConst, A, E, x, betaConst, y, flags, null_stream);\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_wait(x, x.tag);\n }\n\n A.setViewExterior();\n flags = (ViewType)(~(A.getViewInterior()) & A.getViewExterior());\n\n if (flags != 0)\n {\n bsrmv_internal(color, alphaConst, A, E, x, betaConst, y, flags, null_stream);\n }\n }\n else\n {\n bsrmv_internal(color, alphaConst, A, E, x, betaConst, y, flags, null_stream);\n }\n\n y.dirtybit = 1;\n //if (!A.is_matrix_singleGPU() && y.size() == x.size() && y.delayed_send==0)\n // A.manager->exchange_halo_async(y, y.tag);\n A.setView(oldView);\n }\n}\n\n__global__ void offset_by_col_off(int nrows, int* rows, const int* bsrRowPtr)\n{\n int i = threadIdx.x + blockIdx.x * blockDim.x;\n if(i >= nrows+1)\n {\n return;\n }\n\n rows[i] = bsrRowPtr[i] - bsrRowPtr[0];\n}\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_internal( const typename TConfig::VecPrec alphaConst,\n const Matrix &A,\n const Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view,\n const cudaStream_t &stream)\n{\n typedef typename TConfig::VecPrec ValueTypeB;\n int rowOff, nrows, nnz;\n A.getOffsetAndSizeForView(view, &rowOff, &nrows);\n A.getNnzForView(view, &nnz);\n\n cusparseDirection_t direction = CUSPARSE_DIRECTION_COLUMN;\n\n if ( A.getBlockFormat() == ROW_MAJOR )\n {\n direction = CUSPARSE_DIRECTION_ROW;\n }\n\n bool has_offdiag = nnz != 0;\n\n if (has_offdiag )\n {\n bsrmv( Cusparse::get_instance().m_handle, direction, CUSPARSE_OPERATION_NON_TRANSPOSE,\n nrows, A.get_num_cols(), nnz, &alphaConst,\n A.cuMatDescr,\n A.values.raw(),\n A.m_seq_offsets.raw() + rowOff,\n A.row_offsets.raw() + rowOff,\n A.col_indices.raw(),\n rowOff,\n A.get_block_dimx(),\n x.raw(),\n &betaConst,\n y.raw() + rowOff * A.get_block_dimx(),\n stream);\n }\n\n if (A.hasProps(DIAG))\n {\n ValueTypeB beta;\n\n if (!has_offdiag)\n {\n beta = betaConst;\n }\n else\n {\n beta = types::util::get_one();\n }\n\n bsrmv( Cusparse::get_instance().m_handle, direction, CUSPARSE_OPERATION_NON_TRANSPOSE,\n nrows, A.get_num_cols(), A.get_num_rows(), &alphaConst,\n A.cuMatDescr,\n A.values.raw() + A.diagOffset()*A.get_block_size(),\n A.m_seq_offsets.raw(),\n A.m_seq_offsets.raw() + rowOff,\n A.m_seq_offsets.raw(),\n rowOff,\n A.get_block_dimx(),\n x.raw(), &beta,\n y.raw() + rowOff * A.get_block_dimx(),\n stream);\n }\n}\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_internal_with_mask( const typename TConfig::VecPrec alphaConst,\n const Matrix &A,\n const Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view,\n const cudaStream_t &stream)\n{\n if (A.is_matrix_singleGPU())\n {\n FatalError(\"Should not be here in bsrmv_internal_with_mask\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if(view != INTERIOR && view != BOUNDARY)\n {\n FatalError(\"Only INTERIOR and BOUNDARY views supported for bsrmv_internal_with_mask\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename TConfig::VecPrec ValueType;\n cusparseDirection_t direction = CUSPARSE_DIRECTION_COLUMN;\n\n if ( A.getBlockFormat() == ROW_MAJOR )\n {\n direction = CUSPARSE_DIRECTION_ROW;\n }\n\n const int *start_offsets, *end_offsets;\n start_offsets = A.row_offsets.raw();\n end_offsets = A.row_offsets.raw() + 1;\n typedef typename Matrix::index_type index_type;\n\n int offset, nrows, nnz;\n A.getOffsetAndSizeForView(view, &offset, &nrows);\n A.getNnzForView(view, &nnz);\n\n if (nrows <= 0)\n {\n return; // nothing to do, early exit\n }\n\n bool has_offdiag = nnz != 0;\n\n if (has_offdiag)\n {\n cusparseCheckError(cusparseSetStream(Cusparse::get_instance().m_handle, stream));\n bsrxmv_internal( Cusparse::get_instance().m_handle, direction, CUSPARSE_OPERATION_NON_TRANSPOSE, nrows,\n nrows, A.get_num_cols(), nnz, &alphaConst,\n A.cuMatDescr,\n A.values.raw(),\n A.manager->getRowsListForView(view).raw(),\n start_offsets, end_offsets, A.col_indices.raw(),\n A.get_block_dimx(),\n x.raw(), &betaConst,\n y.raw() );\n\n // Reset to default stream\n cusparseCheckError(cusparseSetStream(Cusparse::get_instance().m_handle, 0));\n }\n\n if (A.hasProps(DIAG))\n {\n FatalError(\"Diag not supported in multiply with mask\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_internal_with_mask_restriction( const typename TConfig::VecPrec alphaConst,\n const Matrix &R,\n const Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view,\n const cudaStream_t &stream,\n Matrix &P)\n{\n if (P.is_matrix_singleGPU())\n {\n FatalError(\"Should not be here in bsrmv_internal_with_mask_with_restriction\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if(view != OWNED && view != HALO1)\n {\n FatalError(\"View not supported in restriction operation\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename TConfig::VecPrec ValueType;\n cusparseDirection_t direction = CUSPARSE_DIRECTION_COLUMN;\n\n if ( R.getBlockFormat() == ROW_MAJOR )\n {\n direction = CUSPARSE_DIRECTION_ROW;\n }\n\n int rowOff, nrows, nnz;\n R.getOffsetAndSizeForView(view, &rowOff, &nrows);\n R.getNnzForView(view, &nnz);\n\n bool has_offdiag = nnz != 0;\n typedef typename Matrix::index_type index_type;\n\n if (nrows <= 0)\n {\n return; // nothing to do, early exit\n }\n\n if (has_offdiag)\n {\n bsrmv( Cusparse::get_instance().m_handle, direction, CUSPARSE_OPERATION_NON_TRANSPOSE,\n nrows, R.get_num_cols(), nnz, &alphaConst,\n R.cuMatDescr,\n R.values.raw(),\n R.m_seq_offsets.raw() + rowOff,\n R.row_offsets.raw() + rowOff,\n R.col_indices.raw(),\n rowOff,\n R.get_block_dimx(),\n x.raw(),\n &betaConst,\n y.raw() + rowOff * R.get_block_dimx(),\n stream);\n }\n\n if (R.hasProps(DIAG))\n {\n FatalError(\"Diag not supported in multiply with mask\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_internal( const typename TConfig::VecPrec alphaConst,\n const Matrix &A,\n const typename Matrix::MVector &E,\n const Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view,\n const cudaStream_t &stream)\n{\n typedef typename TConfig::VecPrec ValueType;\n int rowOff, nrows, nnz;\n A.getOffsetAndSizeForView(view, &rowOff, &nrows);\n A.getNnzForView(view, &nnz);\n\n cusparseDirection_t direction = A.getBlockFormat() == ROW_MAJOR ? CUSPARSE_DIRECTION_ROW : CUSPARSE_DIRECTION_COLUMN;\n\n bsrmv( Cusparse::get_instance().m_handle, direction, CUSPARSE_OPERATION_NON_TRANSPOSE,\n nrows, A.get_num_cols(), nnz, &alphaConst,\n A.cuMatDescr,\n E.raw(),\n A.m_seq_offsets.raw(),\n A.m_seq_offsets.raw() + rowOff,\n A.m_seq_offsets.raw(),\n rowOff,\n A.get_block_dimx(),\n x.raw(), &betaConst,\n y.raw() + rowOff * A.get_block_dimx(),\n stream);\n}\n\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_internal( ColumnColorSelector columnColorSelector,\n const int color,\n const typename TConfig::VecPrec alphaConst,\n const Matrix &A,\n const Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view,\n const cudaStream_t &stream)\n{\n typedef typename TConfig::VecPrec ValueType;\n\n if (!A.hasProps(COLORING))\n {\n FatalError(\"Matrix is not colored, exiting\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (color < 0 || color >= A.getMatrixColoring().getNumColors())\n {\n FatalError(\"Unknown color\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n typedef typename Matrix::index_type index_type;\n index_type colorStart = 0;\n\n if ( !(view & INTERIOR) )\n {\n colorStart = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[color];\n }\n else\n {\n colorStart = A.getMatrixColoring().getOffsetsRowsPerColor()[color];\n }\n\n index_type colorNum = 0;\n\n if ( view == A.getViewInterior() )\n {\n colorNum = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[color];\n }\n else\n {\n colorNum = A.getMatrixColoring().getOffsetsRowsPerColor()[color + 1];\n }\n\n colorNum -= colorStart;\n\n if (colorNum <= 0)\n {\n return; // nothing to do, early exit\n }\n\n if (columnColorSelector == DIAG_COL)\n {\n FatalError(\"DIAG_COL has been disabled to avoid allocating diag_offsets\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n const int *start_offsets, *end_offsets;\n\n switch (columnColorSelector)\n {\n case ALL_COLORS:\n start_offsets = A.row_offsets.raw();\n end_offsets = A.row_offsets.raw() + 1;\n break;\n\n case SMALLER_COLORS:\n start_offsets = A.row_offsets.raw();\n end_offsets = A.m_smaller_color_offsets.raw();\n break;\n\n case LARGER_COLORS:\n start_offsets = A.m_larger_color_offsets.raw();\n end_offsets = A.row_offsets.raw() + 1;\n break;\n\n case DIAG_COL:\n start_offsets = A.diag.raw();\n end_offsets = A.m_diag_end_offsets.raw();\n break;\n\n default:\n FatalError(\"Unknown color selector\", AMGX_ERR_CORE);\n }\n\n cusparseDirection_t direction = CUSPARSE_DIRECTION_COLUMN;\n\n if ( A.getBlockFormat() == ROW_MAJOR )\n {\n direction = CUSPARSE_DIRECTION_ROW;\n }\n\n bool has_offdiag = A.get_num_nz() != 0;\n cusparseCheckError(cusparseSetStream(Cusparse::get_instance().m_handle, stream));\n\n if (has_offdiag)\n {\n bsrxmv_internal( Cusparse::get_instance().m_handle, direction, CUSPARSE_OPERATION_NON_TRANSPOSE, colorNum,\n A.get_num_rows(), A.get_num_cols(), A.get_num_nz(), &alphaConst,\n A.cuMatDescr,\n A.values.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + colorStart,\n start_offsets, end_offsets, A.col_indices.raw(),\n A.get_block_dimx(),\n x.raw(), &betaConst,\n y.raw() );\n }\n\n if (A.hasProps(DIAG) && columnColorSelector == ALL_COLORS)\n {\n typename TConfig::VecPrec beta = betaConst;\n\n if (has_offdiag)\n {\n beta = types::util::get_one();\n }\n\n bsrxmv_internal( Cusparse::get_instance().m_handle, direction, CUSPARSE_OPERATION_NON_TRANSPOSE, colorNum,\n A.get_num_rows(), A.get_num_cols(), A.get_num_rows(), &alphaConst,\n A.cuMatDescr,\n A.values.raw() + A.diagOffset()*A.get_block_size(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + colorStart,\n A.m_seq_offsets.raw(),\n NULL,\n A.m_seq_offsets.raw(),\n A.get_block_dimx(),\n x.raw(), &beta,\n y.raw() );\n }\n\n // Reset to default stream\n cusparseCheckError(cusparseSetStream(Cusparse::get_instance().m_handle, 0));\n}\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_internal( const int color,\n const typename TConfig::VecPrec alphaConst,\n const Matrix &A,\n const typename Matrix::MVector &E,\n const Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view,\n const cudaStream_t &stream)\n{\n typedef typename TConfig::VecPrec ValueType;\n\n if ( !A.hasProps(COLORING) )\n {\n FatalError(\"Matrix is not colored, exiting\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if ( color < 0 || color >= A.getMatrixColoring().getNumColors() )\n {\n FatalError(\"Unknown color\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n typedef typename Matrix::index_type index_type;\n //const index_type colorStart = ((view & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[color] : A.getMatrixColoring().getOffsetsRowsPerColor()[color];\n //const index_type colorNum = ((view == A.getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[color] : A.getMatrixColoring().getOffsetsRowsPerColor()[color+1]) - colorStart;\n //const index_type colorStart= A.getMatrixColoring().getOffsetsRowsPerColor()[color];\n //const index_type colorNum= A.getMatrixColoring().getOffsetsRowsPerColor()[color+1] - colorStart;\n index_type colorStart = 0;\n\n if ( !(view & INTERIOR) )\n {\n colorStart = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[color];\n }\n else\n {\n colorStart = A.getMatrixColoring().getOffsetsRowsPerColor()[color];\n }\n\n index_type colorNum = 0;\n\n if ( view == A.getViewInterior() )\n {\n colorNum = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[color];\n }\n else\n {\n colorNum = A.getMatrixColoring().getOffsetsRowsPerColor()[color + 1];\n }\n\n colorNum -= colorStart;\n\n if (colorNum <= 0)\n {\n return; // nothing to do, early exit\n }\n\n cusparseDirection_t direction = CUSPARSE_DIRECTION_COLUMN;\n\n if ( A.getBlockFormat() == ROW_MAJOR )\n {\n direction = CUSPARSE_DIRECTION_ROW;\n }\n\n cusparseCheckError(cusparseSetStream(Cusparse::get_instance().m_handle, stream));\n\n bsrxmv_internal( Cusparse::get_instance().m_handle, direction, CUSPARSE_OPERATION_NON_TRANSPOSE, colorNum,\n A.get_num_rows(), A.get_num_cols(), A.get_num_nz(), &alphaConst,\n A.cuMatDescr,\n E.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + colorStart,\n A.m_seq_offsets.raw(),\n NULL,\n A.m_seq_offsets.raw(),\n A.get_block_dimx(),\n x.raw(), &betaConst,\n y.raw() );\n // Reset to default stream\n cusparseCheckError(cusparseSetStream(Cusparse::get_instance().m_handle, 0));\n}\n\n// Simple custom implementation of matrix-vector product that has only 1 kernel.\ntemplate\n__global__ void csrmv(\n int nrows,\n const T alpha,\n const T* __restrict__ csrVal,\n const int* __restrict__ csrRow,\n const int* __restrict__ csrCol,\n const T* __restrict__ x,\n const T beta,\n T* __restrict__ y)\n{\n for(int i = threadIdx.x + blockIdx.x*blockDim.x; i < nrows; i += blockDim.x*gridDim.x)\n {\n T y_tmp = amgx::types::util::get_zero();\n\n int row_b = csrRow[i];\n int row_e = csrRow[i+1];\n\n // Unrolling is important for performance here.\n // Possible to squeeze more performance out of the key kernels if we\n // measure the sparsity and use it to inform unrolling.\n for (int col = row_b; col < row_e; col += UNROLL)\n {\n#pragma unroll UNROLL\n for(int off = 0; off < UNROLL; ++off)\n {\n int c = col + off;\n if(c < row_e) y_tmp = alpha * csrVal[c] * x[csrCol[c]] + y_tmp;\n }\n }\n\n // Don't read y unnecessarily\n if(amgx::types::util::is_zero(beta))\n {\n y[i] = y_tmp;\n }\n else\n {\n y[i] = beta*y[i] + y_tmp;\n }\n }\n}\n\ntemplate\ninline void generic_SpMV(cusparseHandle_t handle, cusparseOperation_t trans,\n int mb, int nb, int nnzb, int rowOff,\n const MatType *alpha,\n const MatType *val,\n const IndType *rowPtr,\n const IndType *colInd,\n const VecType *x,\n const VecType *beta,\n VecType *y,\n cudaDataType matType,\n cudaDataType vecType,\n const cudaStream_t& stream)\n{\n constexpr int cta_size = 128;\n const int sm_count = getSMCount();\n\n // Assuming that csrmv will be more efficient than cuSPARSE for row counts \n // that are lower than the 3 times the total number of threads\n // cuSPARSE does not like the offsetting required when latency hiding\n // it's possible to reverse the offsets, but requires extra kernel invocation \n // and usually the dependent part of the call is smaller\n if(rowOff > 0 || mb < cta_size * sm_count * 3)\n {\n // Custom single-kernel SpMV, we could actually determine unroll factor\n // more accurately here by checking non-zeros per row\n constexpr int unroll_factor = 16;\n int nblocks = mb / cta_size + 1;\n csrmv<<>>(mb, *alpha, val, rowPtr, colInd, x, *beta, y);\n cudaCheckError();\n }\n else\n {\n IndType* rows = const_cast(rowPtr);\n IndType* cols = const_cast(colInd);\n MatType* vals = const_cast(val);\n\n cusparseSpMatDescr_t matA_descr;\n cusparseDnVecDescr_t vecX_descr;\n cusparseDnVecDescr_t vecY_descr;\n cusparseCheckError(cusparseCreateDnVec(&vecX_descr, nb, const_cast(x), vecType));\n cusparseCheckError(cusparseCreateDnVec(&vecY_descr, mb, const_cast(y), vecType));\n cusparseCheckError(\n cusparseCreateCsr(&matA_descr, mb, nb, nnzb, const_cast(rows), const_cast(cols),\n const_cast(vals), CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, matType));\n\n // Use ALG2 for deterministic results when determinism_flag is set\n cusparseSpMVAlg_t spmv_alg = Cusparse::get_instance().is_deterministic() ? CUSPARSE_CSRMV_ALG2 : CUSPARSE_CSRMV_ALG1;\n\n size_t bufferSize = 0;\n cusparseCheckError(cusparseSpMV_bufferSize(handle, trans, alpha, matA_descr, vecX_descr, beta, vecY_descr, matType, spmv_alg, &bufferSize));\n\n void* dBuffer = NULL;\n if(bufferSize > 0)\n {\n cudaError_t cuda_rc = amgx::memory::cudaMallocAsync(&dBuffer, bufferSize, stream);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMallocAsync failed in generic_SpMV\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n\n cusparseCheckError(cusparseSpMV(handle, trans, alpha, matA_descr, vecX_descr, beta, vecY_descr, matType, spmv_alg, dBuffer) );\n\n cusparseCheckError(cusparseDestroySpMat(matA_descr));\n cusparseCheckError(cusparseDestroyDnVec(vecX_descr));\n cusparseCheckError(cusparseDestroyDnVec(vecY_descr));\n\n if(bufferSize > 0)\n {\n cudaError_t cuda_rc = amgx::memory::cudaFreeAsync(dBuffer);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaFreeAsync failed in generic_SpMV\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n }\n}\n\ninline void Cusparse::bsrmv( cusparseHandle_t handle, cusparseDirection_t dir, cusparseOperation_t trans,\n int mb, int nb, int nnzb,\n const float *alpha,\n const cusparseMatDescr_t descr,\n const float *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrColInd,\n int rowOff,\n int blockDim,\n const float *x,\n const float *beta,\n float *y,\n const cudaStream_t& stream)\n{\n // Run cuSparse on selected stream\n cusparseCheckError(cusparseSetStream(handle, stream));\n\n if (blockDim == 1)\n {\n generic_SpMV(handle, trans, mb, nb, nnzb, rowOff, alpha, bsrVal, bsrRowPtr, bsrColInd, x, beta, y, CUDA_R_32F, CUDA_R_32F, stream);\n }\n else\n {\n cusparseCheckError(cusparseSbsrmv(handle, dir, trans, mb, nb, nnzb, alpha, descr, bsrVal, bsrRowPtr, bsrColInd, blockDim, x, beta, y));\n }\n\n // Reset cuSparse to default stream\n cusparseCheckError(cusparseSetStream(handle, 0));\n}\n\ninline void Cusparse::bsrmv( cusparseHandle_t handle, cusparseDirection_t dir, cusparseOperation_t trans,\n int mb, int nb, int nnzb,\n const double *alpha,\n const cusparseMatDescr_t descr,\n const double *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrColInd,\n int rowOff,\n int blockDim,\n const double *x,\n const double *beta,\n double *y,\n const cudaStream_t& stream)\n{\n // Run cuSparse on selected stream\n cusparseCheckError(cusparseSetStream(handle, stream));\n\n if (blockDim == 1)\n {\n generic_SpMV(handle, trans, mb, nb, nnzb, rowOff, alpha, bsrVal, bsrRowPtr, bsrColInd, x, beta, y, CUDA_R_64F, CUDA_R_64F, stream);\n }\n else\n {\n cusparseCheckError(cusparseDbsrmv(handle, dir, trans, mb, nb, nnzb, alpha, descr, bsrVal, bsrRowPtr, bsrColInd, blockDim, x, beta, y));\n }\n\n // Reset cuSparse to default stream\n cusparseCheckError(cusparseSetStream(handle, 0));\n}\n\ninline void Cusparse::bsrmv( cusparseHandle_t handle, cusparseDirection_t dir, cusparseOperation_t trans,\n int mb, int nb, int nnzb,\n const double *alpha,\n const cusparseMatDescr_t descr,\n const float *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrColInd,\n int rowOff,\n int blockDim,\n const double *x,\n const double *beta,\n double *y,\n const cudaStream_t& stream)\n{\n #if 0\n // Run cuSparse on selected stream\n cusparseCheckError(cusparseSetStream(handle, stream));\n\n const double *d_bsrVal = reinterpret_cast(const_cast(bsrVal)); // this works due to private API call in the matrix initialization which sets cusparse matrix description in the half precision mode\n cusparseCheckError(cusparseDbsrxmv(handle, dir, trans, mb, mb, nb, nnzb, alpha, descr, d_bsrVal, bsrMaskPtr, bsrRowPtr, bsrRowPtr + 1, bsrColInd, blockDim, x, beta, y));\n\n // Reset cuSparse to default stream\n cusparseCheckError(cusparseSetStream(handle, 0));\n #else\n FatalError(\"Mixed precision modes not currently supported for CUDA 10.1 or later.\", AMGX_ERR_NOT_IMPLEMENTED);\n #endif\n}\n\n// Custom implementation of matrix-vector product to replace the original bsrxmv,\n// but with block size of 1.\ntemplate\n__global__ void csrxmv(\n int sizeOfMask,\n const T alpha,\n const T* __restrict__ csrVal,\n const int* __restrict__ csrMask,\n const int* __restrict__ csrRow,\n const int* __restrict__ csrCol,\n const T* __restrict__ x,\n const T beta,\n T* __restrict__ y)\n{\n for(int i = threadIdx.x + blockIdx.x*blockDim.x; i < sizeOfMask; i += blockDim.x*gridDim.x)\n {\n int row = csrMask[i];\n T y_tmp = amgx::types::util::get_zero();\n\n int row_b = csrRow[row];\n int row_e = csrRow[row+1];\n\n // Unrolling is important for performance here.\n // Possible to squeeze more performance out of the key kernels if we\n // measure the sparsity and use it to inform unrolling.\n for (int col = row_b; col < row_e; col += UNROLL)\n {\n#pragma unroll UNROLL\n for(int off = 0; off < UNROLL; ++off)\n {\n int c = col + off;\n if(c < row_e) y_tmp = alpha * csrVal[c] * x[csrCol[c]] + y_tmp;\n }\n }\n\n // Don't read y unnecessarily\n if(amgx::types::util::is_zero(beta))\n {\n y[row] = y_tmp;\n }\n else\n {\n y[row] = beta*y[row] + y_tmp;\n }\n }\n}\n\n// Replaces the functionality of cusparse?bsrxmv for blockDim == 1\ntemplate\ninline void Xcsrxmv( cusparseHandle_t handle, cusparseDirection_t dir, cusparseOperation_t trans, int sizeOfMask,\n int mb, int nb, int nnzb,\n const T *alpha,\n const cusparseMatDescr_t descr,\n const T *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrEndPtr,\n const int *bsrColInd,\n int blockDim,\n const T *x,\n const T *beta,\n T *y)\n{\n if(blockDim != 1)\n {\n FatalError(\"Xcsrxmv only to be called with scalar matrices.\", AMGX_ERR_INTERNAL);\n }\n if (trans != CUSPARSE_OPERATION_NON_TRANSPOSE)\n {\n FatalError(\"Cannot currently latency hide a transposed matrix.\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n if (dir != CUSPARSE_DIRECTION_ROW)\n {\n FatalError(\"Cannot currently latency hide if matrix is not row major.\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n constexpr int cta_size = 128;\n constexpr int unroll_factor = 16;\n int nblocks = sizeOfMask / cta_size + 1;\n csrxmv<<>>(sizeOfMask, *alpha, bsrVal, bsrMaskPtr, bsrRowPtr, bsrColInd, x, *beta, y);\n cudaCheckError();\n}\n\n// overloaded C++ wrappers for cusparse?bsrxmv\n// bsrxmv\n// matrix - float\n// vector - float\ninline void Cusparse::bsrxmv_internal(cusparseHandle_t handle, cusparseDirection_t dir, cusparseOperation_t trans, int sizeOfMask,\n int mb, int nb, int nnzb,\n const float *alpha,\n const cusparseMatDescr_t descr,\n const float *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrEndPtr,\n const int *bsrColInd,\n int blockDim,\n const float *x,\n const float *beta,\n float *y)\n{\n if (bsrEndPtr == NULL && bsrMaskPtr == NULL)\n {\n cusparseCheckError(CUSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n }\n else\n {\n if (bsrEndPtr == NULL)\n {\n bsrEndPtr = bsrRowPtr + 1;\n }\n\n if (blockDim == 1)\n {\n Xcsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y);\n }\n else\n {\n cusparseCheckError(cusparseSbsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y));\n }\n }\n}\n\n// bsrxmv\n// matrix - float\n// vector - double\ninline void Cusparse::bsrxmv_internal( cusparseHandle_t handle, cusparseDirection_t dir, cusparseOperation_t trans, int sizeOfMask,\n int mb, int nb, int nnzb,\n const double *alpha,\n const cusparseMatDescr_t descr,\n const float *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrEndPtr,\n const int *bsrColInd,\n int blockDim,\n const double *x,\n const double *beta,\n double *y)\n{\n cusparseCheckError(CUSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n}\n\n// bsrxmv\n// matrix - double\n// vector - double\ninline void Cusparse::bsrxmv_internal( cusparseHandle_t handle, cusparseDirection_t dir, cusparseOperation_t trans, int sizeOfMask,\n int mb, int nb, int nnzb,\n const double *alpha,\n const cusparseMatDescr_t descr,\n const double *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrEndPtr,\n const int *bsrColInd,\n int blockDim,\n const double *x,\n const double *beta,\n double *y)\n{\n if (bsrEndPtr == NULL && bsrMaskPtr == NULL)\n {\n cusparseCheckError(CUSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n }\n else\n {\n if (bsrEndPtr == NULL)\n {\n bsrEndPtr = bsrRowPtr + 1;\n }\n\n if (blockDim == 1)\n {\n Xcsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y);\n }\n else\n {\n cusparseCheckError(cusparseDbsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y));\n }\n }\n}\n\n\n\ninline void Cusparse::bsrmv( cusparseHandle_t handle, cusparseDirection_t dir, cusparseOperation_t trans,\n int mb, int nb, int nnzb,\n const cuComplex *alpha,\n const cusparseMatDescr_t descr,\n const cuComplex *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrColInd,\n int rowOff,\n int blockDim,\n const cuComplex *x,\n const cuComplex *beta,\n cuComplex *y,\n const cudaStream_t& stream)\n{\n // Run cuSparse on selected stream\n cusparseCheckError(cusparseSetStream(handle, stream));\n\n if (blockDim == 1)\n {\n generic_SpMV(handle, trans, mb, nb, nnzb, rowOff, alpha, bsrVal, bsrRowPtr, bsrColInd, x, beta, y, CUDA_C_32F, CUDA_C_32F, stream);\n }\n else\n {\n cusparseCheckError(cusparseCbsrmv(handle, dir, trans, mb, nb, nnzb, alpha, descr, bsrVal, bsrRowPtr, bsrColInd, blockDim, x, beta, y));\n }\n\n // Reset cuSparse to default stream\n cusparseCheckError(cusparseSetStream(handle, 0));\n}\n\ninline void Cusparse::bsrmv( cusparseHandle_t handle, cusparseDirection_t dir, cusparseOperation_t trans,\n int mb, int nb, int nnzb,\n const cuDoubleComplex *alpha,\n const cusparseMatDescr_t descr,\n const cuDoubleComplex *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrColInd,\n int rowOff,\n int blockDim,\n const cuDoubleComplex *x,\n const cuDoubleComplex *beta,\n cuDoubleComplex *y,\n const cudaStream_t& stream)\n{\n // Run cuSparse on selected stream\n cusparseCheckError(cusparseSetStream(handle, stream));\n\n if (blockDim == 1)\n {\n generic_SpMV(handle, trans, mb, nb, nnzb, rowOff, alpha, bsrVal, bsrRowPtr, bsrColInd, x, beta, y, CUDA_C_64F, CUDA_C_64F, stream);\n }\n else\n {\n cusparseCheckError(cusparseZbsrmv(handle, dir, trans, mb, nb, nnzb, alpha, descr, bsrVal, bsrRowPtr, bsrColInd, blockDim, x, beta, y));\n }\n\n // Reset cuSparse to default stream\n cusparseCheckError(cusparseSetStream(handle, 0));\n}\n\ninline void Cusparse::bsrmv( cusparseHandle_t handle, cusparseDirection_t dir, cusparseOperation_t trans,\n int mb, int nb, int nnzb,\n const cuDoubleComplex *alpha,\n const cusparseMatDescr_t descr,\n const cuComplex *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrColInd,\n int rowOff,\n int blockDim,\n const cuDoubleComplex *x,\n const cuDoubleComplex *beta,\n cuDoubleComplex *y,\n const cudaStream_t& stream)\n{\n #if 0\n // Run cuSparse on selected stream\n cusparseCheckError(cusparseSetStream(handle, stream));\n\n const cuDoubleComplex *d_bsrVal = reinterpret_cast(const_cast(bsrVal));\n cusparseCheckError(cusparseZbsrxmv(handle, dir, trans, mb, mb, nb, nnzb, alpha, descr, d_bsrVal, bsrMaskPtr, bsrRowPtr, bsrRowPtr + 1, bsrColInd, blockDim, x, beta, y));\n\n // Reset cuSparse to default stream\n cusparseCheckError(cusparseSetStream(handle, 0));\n #else\n FatalError(\"Mixed precision modes not currently supported for CUDA 10.1 or later.\", AMGX_ERR_NOT_IMPLEMENTED);\n #endif\n}\n\n\n// overloaded C++ wrappers for cusparse?bsrxmv\n// bsrxmv\n// matrix - cuComplex\n// vector - cuComplex\ninline void Cusparse::bsrxmv_internal( cusparseHandle_t handle, cusparseDirection_t dir, cusparseOperation_t trans, int sizeOfMask,\n int mb, int nb, int nnzb,\n const cuComplex *alpha,\n const cusparseMatDescr_t descr,\n const cuComplex *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrEndPtr,\n const int *bsrColInd,\n int blockDim,\n const cuComplex *x,\n const cuComplex *beta,\n cuComplex *y)\n{\n if (bsrEndPtr == NULL && bsrMaskPtr == NULL)\n {\n cusparseCheckError(CUSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n }\n else\n {\n if (bsrEndPtr == NULL)\n {\n bsrEndPtr = bsrRowPtr + 1;\n }\n\n if (blockDim == 1)\n {\n Xcsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y);\n }\n else\n {\n cusparseCheckError(cusparseCbsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y));\n }\n }\n}\n\n// bsrxmv\n// matrix - cuComplex\n// vector - cuDoubleComplex\ninline void Cusparse::bsrxmv_internal( cusparseHandle_t handle, cusparseDirection_t dir, cusparseOperation_t trans, int sizeOfMask,\n int mb, int nb, int nnzb,\n const cuDoubleComplex *alpha,\n const cusparseMatDescr_t descr,\n const cuComplex *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrEndPtr,\n const int *bsrColInd,\n int blockDim,\n const cuDoubleComplex *x,\n const cuDoubleComplex *beta,\n cuDoubleComplex *y)\n{\n cusparseCheckError(CUSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n}\n\n// bsrxmv\n// matrix - cuDoubleComplex\n// vector - cuDoubleComplex\ninline void Cusparse::bsrxmv_internal( cusparseHandle_t handle, cusparseDirection_t dir, cusparseOperation_t trans, int sizeOfMask,\n int mb, int nb, int nnzb,\n const cuDoubleComplex *alpha,\n const cusparseMatDescr_t descr,\n const cuDoubleComplex *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrEndPtr,\n const int *bsrColInd,\n int blockDim,\n const cuDoubleComplex *x,\n const cuDoubleComplex *beta,\n cuDoubleComplex *y)\n{\n if (bsrEndPtr == NULL && bsrMaskPtr == NULL)\n {\n cusparseCheckError(CUSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n }\n else\n {\n if (bsrEndPtr == NULL)\n {\n bsrEndPtr = bsrRowPtr + 1;\n }\n\n if (blockDim == 1)\n {\n Xcsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y);\n }\n else\n {\n cusparseCheckError(cusparseZbsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y));\n }\n }\n}\n\n\nnamespace\n{\ntemplate\ninline void\ngeneric_SpMM(cusparseHandle_t handle, cusparseOperation_t transA,\n int m, int n, int k, int nnz,\n int ldb, int ldc,\n const MatType *alpha,\n const MatType *Avals,\n const MatType *Bvals,\n MatType *Cvals,\n const IndType *rowPtr,\n const IndType *colInd,\n const MatType *beta,\n cudaDataType matType)\n{\n // Create the matrix descriptors\n cusparseSpMatDescr_t matA_descr;\n cusparseDnMatDescr_t matB_descr;\n cusparseDnMatDescr_t matC_descr;\n cusparseCheckError(\n cusparseCreateCsr(&matA_descr, m, k, nnz, const_cast(rowPtr), const_cast(colInd),\n const_cast(Avals), CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, matType));\n cusparseCheckError(\n cusparseCreateDnMat(&matB_descr, k, n, ldb, const_cast(Bvals), matType, CUSPARSE_ORDER_COL));\n cusparseCheckError(\n cusparseCreateDnMat(&matC_descr, m, n, ldc, const_cast(Cvals), matType, CUSPARSE_ORDER_COL));\n\n // Check if a buffer is required, and if so allocate it using caching allocator\n size_t bufferSize = 0;\n cusparseCheckError(\n cusparseSpMM_bufferSize(handle, transA, CUSPARSE_OPERATION_NON_TRANSPOSE, alpha, matA_descr, matB_descr,\n beta, matC_descr, matType, CUSPARSE_SPMM_ALG_DEFAULT, &bufferSize));\n\n void* dBuffer = NULL;\n if(bufferSize > 0)\n {\n cudaError_t cuda_rc = amgx::memory::cudaMallocAsync(&dBuffer, bufferSize);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMallocAsync failed in generic_SpMM\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n\n // Compute the sparse matrix - dense matrix product\n cusparseCheckError(\n cusparseSpMM(handle, transA, CUSPARSE_OPERATION_NON_TRANSPOSE, alpha, matA_descr, matB_descr, beta,\n matC_descr, matType, CUSPARSE_SPMM_ALG_DEFAULT, dBuffer));\n\n // Clean up\n cusparseCheckError(cusparseDestroySpMat(matA_descr));\n cusparseCheckError(cusparseDestroyDnMat(matB_descr));\n cusparseCheckError(cusparseDestroyDnMat(matC_descr));\n\n if(bufferSize > 0)\n {\n cudaError_t cuda_rc = amgx::memory::cudaFreeAsync(dBuffer);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaFreeAsync failed in generic_SpMM\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n}\n\nvoid\ncusparse_csrmm(cusparseHandle_t handle, cusparseOperation_t transA,\n int m, int n, int k, int nnz,\n const float *alpha,\n const cusparseMatDescr_t descrA,\n const float *csrValA,\n const int *csrRowPtrA, const int *csrColIndA,\n const float *B, int ldb,\n const float *beta, float *C, int ldc)\n{\n generic_SpMM(handle, transA, m, n, k, nnz, ldb, ldc, alpha, csrValA, B, C, csrRowPtrA, csrColIndA, beta, CUDA_R_32F);\n}\n\nvoid\ncusparse_csrmm(cusparseHandle_t handle, cusparseOperation_t transA,\n int m, int n, int k, int nnz,\n const double *alpha,\n const cusparseMatDescr_t descrA,\n const float *csrValA,\n const int *csrRowPtrA, const int *csrColIndA,\n const double *B, int ldb,\n const double *beta, double *C, int ldc)\n{\n cusparseCheckError(CUSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n}\n\nvoid\ncusparse_csrmm(cusparseHandle_t handle, cusparseOperation_t transA,\n int m, int n, int k, int nnz,\n const double *alpha,\n const cusparseMatDescr_t descrA,\n const double *csrValA,\n const int *csrRowPtrA, const int *csrColIndA,\n const double *B, int ldb,\n const double *beta, double *C, int ldc)\n{\n generic_SpMM(handle, transA, m, n, k, nnz, ldb, ldc, alpha, csrValA, B, C, csrRowPtrA, csrColIndA, beta, CUDA_R_64F);\n}\n\nvoid\ncusparse_csrmm(cusparseHandle_t handle, cusparseOperation_t transA,\n int m, int n, int k, int nnz,\n const cuComplex *alpha,\n const cusparseMatDescr_t descrA,\n const cuComplex *csrValA,\n const int *csrRowPtrA, const int *csrColIndA,\n const cuComplex *B, int ldb,\n const cuComplex *beta, cuComplex *C, int ldc)\n{\n generic_SpMM(handle, transA, m, n, k, nnz, ldb, ldc, alpha, csrValA, B, C, csrRowPtrA, csrColIndA, beta, CUDA_C_32F);\n}\n\nvoid\ncusparse_csrmm(cusparseHandle_t handle, cusparseOperation_t transA,\n int m, int n, int k, int nnz,\n const cuDoubleComplex *alpha,\n const cusparseMatDescr_t descrA,\n const cuComplex *csrValA,\n const int *csrRowPtrA, const int *csrColIndA,\n const cuDoubleComplex *B, int ldb,\n const cuDoubleComplex *beta, cuDoubleComplex *C, int ldc)\n{\n cusparseCheckError(CUSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n}\n\nvoid\ncusparse_csrmm(cusparseHandle_t handle, cusparseOperation_t transA,\n int m, int n, int k, int nnz,\n const cuDoubleComplex *alpha,\n const cusparseMatDescr_t descrA,\n const cuDoubleComplex *csrValA,\n const int *csrRowPtrA, const int *csrColIndA,\n const cuDoubleComplex *B, int ldb,\n const cuDoubleComplex *beta, cuDoubleComplex *C, int ldc)\n{\n generic_SpMM(handle, transA, m, n, k, nnz, ldb, ldc, alpha, csrValA, B, C, csrRowPtrA, csrColIndA, beta, CUDA_C_64F);\n}\n}\n\n\ntemplate \nvoid Cusparse::csrmm(typename TConfig::VecPrec alpha,\n Matrix &A,\n Vector &V,\n typename TConfig::VecPrec beta,\n Vector &Res)\n{\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo(V, V.tag);\n }\n\n if (Res.get_num_rows() != A.get_num_rows() || Res.get_num_cols() != V.get_num_cols())\n {\n FatalError(\"Cusparse::csrmm error, dimensions of result matrix do not match input matrices.\", AMGX_ERR_INTERNAL);\n }\n\n cusparseHandle_t handle = Cusparse::get_instance().m_handle;\n cusparse_csrmm(handle, CUSPARSE_OPERATION_NON_TRANSPOSE,\n A.get_num_rows(), V.get_num_cols(), A.get_num_cols(),\n A.values.size(), &alpha, A.cuMatDescr,\n A.values.raw(), A.row_offsets.raw(), A.col_indices.raw(),\n V.raw(), V.get_lda(),\n &beta, Res.raw(), Res.get_lda());\n Res.dirtybit = 1;\n}\n\ntemplate \nvoid transpose_internal(cusparseHandle_t handle, int nRows, int nCols, int nNz, const T* Avals, const int* Arows, const int* Acols, T* Bvals, int* Brows, int* Bcols, cudaDataType valType)\n{\n size_t bufferSize;\n cusparseCheckError(cusparseCsr2cscEx2_bufferSize(\n handle, nRows, nCols, nNz, Avals, Arows, Acols, Bvals, Brows, Bcols, valType,\n CUSPARSE_ACTION_NUMERIC, CUSPARSE_INDEX_BASE_ZERO, CUSPARSE_CSR2CSC_ALG1, &bufferSize));\n\n void *buffer = nullptr;\n if (bufferSize > 0)\n {\n cudaError_t cuda_rc = amgx::memory::cudaMallocAsync(&buffer, bufferSize);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMallocAsync failed in transpose_internal\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n\n cusparseCheckError(cusparseCsr2cscEx2(\n handle, nRows, nCols, nNz, Avals, Arows, Acols, Bvals, Brows, Bcols, valType,\n CUSPARSE_ACTION_NUMERIC, CUSPARSE_INDEX_BASE_ZERO, CUSPARSE_CSR2CSC_ALG1, buffer));\n\n if(bufferSize > 0)\n {\n cudaError_t cuda_rc = amgx::memory::cudaFreeAsync(buffer);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaFreeAsync failed in transpose_internal\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n}\n\nvoid transpose_internal(cusparseHandle_t handle, int nRows, int nCols, int nNz, const float* Avals, const int* Arows, const int* Acols, float* Bvals, int* Brows, int* Bcols)\n{\n transpose_internal(handle, nRows, nCols, nNz, Avals, Arows, Acols, Bvals, Brows, Bcols, CUDA_R_32F);\n}\nvoid transpose_internal(cusparseHandle_t handle, int nRows, int nCols, int nNz, const double* Avals, const int* Arows, const int* Acols, double* Bvals, int* Brows, int* Bcols)\n{\n transpose_internal(handle, nRows, nCols, nNz, Avals, Arows, Acols, Bvals, Brows, Bcols, CUDA_R_64F);\n}\nvoid transpose_internal(cusparseHandle_t handle, int nRows, int nCols, int nNz, const cuComplex* Avals, const int* Arows, const int* Acols, cuComplex* Bvals, int* Brows, int* Bcols)\n{\n transpose_internal(handle, nRows, nCols, nNz, Avals, Arows, Acols, Bvals, Brows, Bcols, CUDA_C_32F);\n}\nvoid transpose_internal(cusparseHandle_t handle, int nRows, int nCols, int nNz, const cuDoubleComplex* Avals, const int* Arows, const int* Acols, cuDoubleComplex* Bvals, int* Brows, int* Bcols)\n{\n transpose_internal(handle, nRows, nCols, nNz, Avals, Arows, Acols, Bvals, Brows, Bcols, CUDA_C_64F);\n}\n\ntemplate \nvoid Cusparse::transpose(const Matrix& A, Matrix& B, const int nRows, const int nNz)\n{\n cusparseHandle_t handle = Cusparse::get_instance().m_handle;\n transpose_internal(handle, nRows, A.get_num_cols(), nNz,\n A.values.raw(), A.row_offsets.raw(), A.col_indices.raw(),\n B.values.raw(), B.row_offsets.raw(), B.col_indices.raw());\n}\n\ntemplate \nvoid Cusparse::transpose(const Matrix& A, Matrix& B)\n{\n cusparseHandle_t handle = Cusparse::get_instance().m_handle;\n transpose_internal(handle, A.get_num_rows(), A.get_num_cols(), A.get_num_nz(),\n A.values.raw(), A.row_offsets.raw(), A.col_indices.raw(),\n B.values.raw(), B.row_offsets.raw(), B.col_indices.raw());\n}\n\n//#define AMGX_CASE_LINE(CASE) template class Cusparse::Type>;\n// AMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n//#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::bsrmv(const typename TemplateMode::Type::VecPrec , Matrix::Type>&, Vector::Type>&, const typename TemplateMode::Type::VecPrec, Vector::Type> &, ViewType);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::bsrxmv(const typename TemplateMode::Type::VecPrec , Matrix::Type>&, Vector::Type>&, const typename TemplateMode::Type::VecPrec, Vector::Type>&, typename Matrix::Type>::IVector&, ViewType);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::bsrmv_with_mask(const typename TemplateMode::Type::VecPrec , Matrix::Type>&, Vector::Type>&, const typename TemplateMode::Type::VecPrec, Vector::Type> &);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::bsrmv_with_mask_restriction(const typename TemplateMode::Type::VecPrec , Matrix::Type>&, Vector::Type>&, const typename TemplateMode::Type::VecPrec, Vector::Type> &, Matrix::Type>& );\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::bsrmv(const typename TemplateMode::Type::VecPrec , Matrix::Type>&, const typename Matrix::Type>::MVector&, Vector::Type>&, const typename TemplateMode::Type::VecPrec, Vector::Type> &, ViewType);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::bsrmv(ColumnColorSelector, const int, const typename TemplateMode::Type::VecPrec , Matrix::Type>&, Vector::Type>&, const typename TemplateMode::Type::VecPrec, Vector::Type> &, ViewType);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::bsrmv(const int, const typename TemplateMode::Type::VecPrec , Matrix::Type>&, const typename Matrix::Type>::MVector &, Vector::Type>&, const typename TemplateMode::Type::VecPrec, Vector::Type> &, ViewType);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::csrmm(typename TemplateMode::Type::VecPrec, Matrix::Type>&, Vector::Type>&, typename TemplateMode::Type::VecPrec, Vector::Type>&);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::transpose(const Matrix::Type>& A, Matrix::Type>& B); \\\n template void Cusparse::transpose(const Matrix::Type>& A, Matrix::Type>& B, const int nRows, const int nNz);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \n\n#if CUDART_VERSION < 11000\n#define HIPSPARSE_SPMM_ALG_DEFAULT HIPSPARSE_MM_ALG_DEFAULT\n#endif\n\n#if CUDART_VERSION >= 12000\n#define HIPSPARSE_CSRMV_ALG1 HIPSPARSE_SPMV_CSR_ALG1\n#define HIPSPARSE_CSRMV_ALG2 HIPSPARSE_SPMV_CSR_ALG2\n#endif\n\nnamespace amgx\n{\n\nCusparse::Cusparse() : m_handle(0), m_determinism_flag(false)\n{\n cusparseCheckError( hipsparseCreate(&m_handle) );\n}\n\nCusparse::~Cusparse()\n{\n destroy_handle();\n}\n\nCusparse &Cusparse::get_instance()\n{\n static Cusparse s_instance;\n s_instance.create_handle();\n return s_instance;\n}\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv(\n const typename TConfig::VecPrec alphaConst,\n Matrix &A,\n Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view )\n{\n hipStream_t null_stream = 0;\n\n // If only COO, add CSR since bsrmv doesn't support COO\n if (A.hasProps(COO) && !A.hasProps(CSR))\n {\n A.set_initialized(0);\n A.addProps(CSR);\n A.computeDiagonal();\n A.set_initialized(1);\n }\n\n // Handle cases where the view is set by the calling routine\n if(view != A.getViewExterior())\n {\n bsrmv_internal(alphaConst, A, x, betaConst, y, view, null_stream);\n return;\n }\n\n bool latencyHiding = (A.getViewInterior() != A.getViewExterior() && !A.is_matrix_singleGPU() && x.dirtybit != 0);\n\n if (latencyHiding)\n {\n A.manager->exchange_halo_split_gather(x, x.tag);\n\n // Multiply interior rows\n bsrmv_internal(alphaConst, A, x, betaConst, y, A.getViewInterior(), null_stream);\n\n // Finish halo exchange\n A.manager->exchange_halo_split_finish(x, x.tag);\n\n // Multiply rows with halo dependencies\n ViewType bnd_view = (ViewType)(~(A.getViewInterior()) & A.getViewExterior());\n bsrmv_internal(alphaConst, A, x, betaConst, y, bnd_view, null_stream);\n }\n else\n {\n if (!A.is_matrix_singleGPU() && x.dirtybit != 0)\n {\n A.manager->exchange_halo_v2(x, x.tag);\n }\n\n bsrmv_internal(alphaConst, A, x, betaConst, y, A.getViewExterior(), null_stream);\n }\n\n y.dirtybit = 1;\n}\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_with_mask(\n const typename TConfig::VecPrec alphaConst,\n Matrix &A,\n Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y)\n{\n hipStream_t null_stream = 0;\n\n // If only COO, add CSR since bsrmv doesn't support COO\n if (A.hasProps(COO) && !A.hasProps(CSR))\n {\n A.set_initialized(0);\n A.addProps(CSR);\n A.computeDiagonal();\n A.set_initialized(1);\n }\n\n bool latencyHiding = (A.getViewInterior() != A.getViewExterior() && !A.is_matrix_singleGPU() && x.dirtybit != 0);\n\n if (latencyHiding)\n {\n A.manager->exchange_halo_split_gather(x, x.tag);\n\n // Multiply interior\n bsrmv_internal_with_mask(alphaConst, A, x, betaConst, y, INTERIOR, null_stream);\n\n A.manager->exchange_halo_split_finish(x, x.tag);\n\n // Multiply exterior\n bsrmv_internal_with_mask(alphaConst, A, x, betaConst, y, BOUNDARY, null_stream);\n }\n else\n {\n if (!A.is_matrix_singleGPU() && x.dirtybit != 0)\n {\n A.manager->exchange_halo_v2(x, x.tag);\n }\n\n bsrmv_internal(alphaConst, A, x, betaConst, y, OWNED, null_stream);\n }\n\n y.dirtybit = 1;\n}\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_with_mask_restriction(\n const typename TConfig::VecPrec alphaConst,\n Matrix &R,\n Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n Matrix &P)\n{\n // If only COO, add CSR since bsrmv doesn't support COO\n //if (A.hasProps(COO) && !A.hasProps(CSR))\n //{\n // A.set_initialized(0);\n // A.addProps(CSR);\n // A.computeDiagonal();\n // A.set_initialized(1);\n //}\n\n bool latencyHiding = (R.getViewInterior() != R.getViewExterior() && !P.is_matrix_singleGPU() && x.dirtybit != 0);\n\n if (latencyHiding)\n {\n hipStream_t null_stream = 0;\n bsrmv_internal_with_mask_restriction(alphaConst, R, x, betaConst, y, HALO1, null_stream, P);\n P.manager->add_from_halo_split_gather(y, y.tag);\n hipError_t cuda_rc = hipEventRecord(P.manager->get_comm_event());\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipEventRecord failed\", AMGX_ERR_CUDA_FAILURE);\n }\n bsrmv_internal_with_mask_restriction(alphaConst, R, x, betaConst, y, OWNED, null_stream, P);\n\n if (P.manager->neighbors.size() != 0)\n {\n cuda_rc = hipEventSynchronize(P.manager->get_comm_event());\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipEventSynchronize failed\", AMGX_ERR_CUDA_FAILURE);\n }\n P.manager->add_from_halo_split_finish(y, y.tag, P.manager->get_bdy_stream());\n cuda_rc = hipStreamSynchronize(P.manager->get_bdy_stream());\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipStreamSynchronize failed\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n }\n else\n {\n bsrmv_internal(alphaConst, R, x, betaConst, y, OWNED, 0);\n\n // Add contribution from neighbors\n P.manager->add_from_halo_v2(y, y.tag);\n }\n\n y.dirtybit = 1;\n}\n\ntemplate< class TConfig >\nvoid Cusparse::bsrxmv(\n const typename TConfig::VecPrec alphaConst,\n Matrix &A,\n Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n typename Matrix::IVector &mask,\n ViewType view )\n{\n // If only COO, add CSR since bsrmv doesn't support COO\n if (A.hasProps(COO) && !A.hasProps(CSR))\n {\n A.set_initialized(0);\n A.addProps(CSR);\n A.computeDiagonal();\n A.set_initialized(1);\n }\n\n const int *start_offsets, *end_offsets;\n start_offsets = A.row_offsets.raw();\n end_offsets = A.row_offsets.raw() + 1;\n hipsparseDirection_t direction = HIPSPARSE_DIRECTION_COLUMN;\n\n if ( A.getBlockFormat() == ROW_MAJOR )\n {\n direction = HIPSPARSE_DIRECTION_ROW;\n }\n\n if (view == OWNED) //This is already a view, thus do not even attempt to do latency hiding\n {\n bsrxmv_internal(Cusparse::get_instance().m_handle, direction, HIPSPARSE_OPERATION_NON_TRANSPOSE,\n mask.size(),\n A.get_num_rows(), A.get_num_cols(), A.get_num_nz(), &alphaConst,\n A.cuMatDescr,\n A.values.raw(),\n mask.raw(),\n start_offsets, end_offsets,\n A.col_indices.raw(),\n A.get_block_dimx(),\n x.raw(),\n &betaConst,\n y.raw());\n }\n else //Try and do latency hiding\n {\n // latency hiding?\n FatalError(\"Trying to do latency hiding in the bsrxmv\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n// E is a vector that represents a diagonal matrix\n// operate on all rows and columns\n// y= alpha*E.x + beta*y\ntemplate< class TConfig >\nvoid Cusparse::bsrmv( const typename TConfig::VecPrec alphaConst,\n Matrix &A,\n const typename Matrix::MVector &E,\n Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view )\n{\n hipStream_t null_stream = 0;\n\n // If only COO, add CSR since bsrmv doesn't support COO\n if (A.hasProps(COO) && !A.hasProps(CSR))\n {\n A.addProps(CSR);\n }\n\n if (view != A.getViewExterior()) //This is already a view, thus do not even attempt to do latency hiding\n {\n bsrmv_internal(alphaConst, A, E, x, betaConst, y, view, null_stream);\n }\n else //Try and do latency hiding\n {\n ViewType oldView = A.currentView();\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_async(x, x.tag);\n }\n\n if (A.getViewExterior() == A.getViewInterior())\n {\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_wait(x, x.tag);\n }\n }\n\n ViewType flags;\n bool latencyHiding = true;\n\n //if (A.manager->num_neighbors() == 0 || (x.dirtybit == 0)) {\n if (A.is_matrix_singleGPU() || (x.dirtybit == 0))\n {\n latencyHiding = false;\n A.setViewExterior();\n flags = (ViewType)(A.getViewExterior());\n }\n else\n {\n flags = (ViewType)(A.getViewInterior());\n A.setViewInterior();\n }\n\n if (latencyHiding)\n {\n bsrmv_internal(alphaConst, A, E, x, betaConst, y, flags, null_stream);\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_wait(x, x.tag);\n }\n\n A.setViewExterior();\n flags = (ViewType)(~(A.getViewInterior()) & A.getViewExterior());\n\n if (flags != 0)\n {\n bsrmv_internal(alphaConst, A, E, x, betaConst, y, flags, null_stream);\n }\n }\n else\n {\n bsrmv_internal(alphaConst, A, E, x, betaConst, y, flags, null_stream);\n }\n\n y.dirtybit = 1;\n //if (!A.is_matrix_singleGPU())\n // if (y.size() == x.size() && y.delayed_send==0) A.manager->exchange_halo_async(y, y.tag);\n A.setView(oldView);\n }\n}\n\n// operate only on columns specified by columnColorSelector, see enum ColumnColorSelector above\n// operate only on rows of specified color, given by A.offsets_rows_per_color, A.sorted_rows_by_color\n// y= alpha*A.x + beta*y\ntemplate< class TConfig >\nvoid Cusparse::bsrmv( ColumnColorSelector columnColorSelector,\n const int color,\n const typename TConfig::VecPrec alphaConst,\n Matrix &A,\n Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view )\n{\n hipStream_t null_stream = 0;\n\n // If only COO, add CSR since bsrmv doesn't support COO\n if (A.hasProps(COO) && !A.hasProps(CSR))\n {\n A.addProps(CSR);\n }\n\n if (view != A.getViewExterior()) //This is already a view, thus do not even attempt to do latency hiding\n {\n // std::cerr << \"exterior view with selector\" << std::endl;\n bsrmv_internal(columnColorSelector, color, alphaConst, A, x, betaConst, y, view, null_stream);\n }\n else\n {\n //Try and do latency hiding\n ViewType oldView = A.currentView();\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_async(x, x.tag);\n }\n\n if (A.getViewExterior() == A.getViewInterior())\n {\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_wait(x, x.tag);\n }\n }\n\n ViewType flags;\n bool latencyHiding = true;\n\n if (A.is_matrix_singleGPU() || (x.dirtybit == 0))\n {\n latencyHiding = false;\n A.setViewExterior();\n flags = (ViewType)(A.getViewExterior());\n }\n else\n {\n flags = (ViewType)(A.getViewInterior());\n A.setViewInterior();\n }\n\n if (latencyHiding)\n {\n bsrmv_internal(columnColorSelector, color, alphaConst, A, x, betaConst, y, flags, null_stream);\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_wait(x, x.tag);\n }\n\n A.setViewExterior();\n flags = (ViewType)(~(A.getViewInterior()) & A.getViewExterior());\n\n if (flags != 0)\n {\n bsrmv_internal(columnColorSelector, color, alphaConst, A, x, betaConst, y, flags, null_stream);\n }\n }\n else\n {\n bsrmv_internal(columnColorSelector, color, alphaConst, A, x, betaConst, y, flags, null_stream);\n }\n\n y.dirtybit = 1;\n //if (!A.is_matrix_singleGPU() && y.size() == x.size() && y.delayed_send==0)\n // A.manager->exchange_halo_async(y, y.tag);\n A.setView(oldView);\n }\n}\n\n// E is a vector that represents a diagonal matrix\n// operate only on rows of specified color, given by A.offsets_rows_per_color, A.sorted_rows_by_color\n// y= alpha*E.x + beta*y\ntemplate< class TConfig >\nvoid Cusparse::bsrmv( const int color,\n const typename TConfig::VecPrec alphaConst,\n Matrix &A,\n const typename Matrix::MVector &E,\n Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view)\n{\n hipStream_t null_stream = 0;\n\n // If only COO, add CSR since bsrmv doesn't support COO\n if (A.hasProps(COO) && !A.hasProps(CSR))\n {\n A.addProps(CSR);\n }\n\n if (view != A.getViewExterior()) //This is already a view, thus do not even attempt to do latency hiding\n {\n bsrmv_internal(color, alphaConst, A, E, x, betaConst, y, view, null_stream);\n }\n else //Try and do latency hiding\n {\n //std::ccooor << \"de\" << std::endl;\n //std::cerr << \"not an exterior view\" << std::endl;\n ViewType oldView = A.currentView();\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_async(x, x.tag);\n }\n\n if (A.getViewExterior() == A.getViewInterior())\n {\n //std::cerr << \"exchange_halo_wait\" << std::endl;\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_wait(x, x.tag);\n }\n }\n\n //std::cerr << \"xxxeded\" << std::endl;\n ViewType flags;\n bool latencyHiding = true;\n\n if (A.is_matrix_singleGPU() || (x.dirtybit == 0))\n {\n latencyHiding = false;\n A.setViewExterior();\n flags = (ViewType)(A.getViewExterior());\n }\n else\n {\n flags = (ViewType)(A.getViewInterior());\n A.setViewInterior();\n }\n\n if (latencyHiding)\n {\n bsrmv_internal(color, alphaConst, A, E, x, betaConst, y, flags, null_stream);\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_wait(x, x.tag);\n }\n\n A.setViewExterior();\n flags = (ViewType)(~(A.getViewInterior()) & A.getViewExterior());\n\n if (flags != 0)\n {\n bsrmv_internal(color, alphaConst, A, E, x, betaConst, y, flags, null_stream);\n }\n }\n else\n {\n bsrmv_internal(color, alphaConst, A, E, x, betaConst, y, flags, null_stream);\n }\n\n y.dirtybit = 1;\n //if (!A.is_matrix_singleGPU() && y.size() == x.size() && y.delayed_send==0)\n // A.manager->exchange_halo_async(y, y.tag);\n A.setView(oldView);\n }\n}\n\n__global__ void offset_by_col_off(int nrows, int* rows, const int* bsrRowPtr)\n{\n int i = threadIdx.x + blockIdx.x * blockDim.x;\n if(i >= nrows+1)\n {\n return;\n }\n\n rows[i] = bsrRowPtr[i] - bsrRowPtr[0];\n}\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_internal( const typename TConfig::VecPrec alphaConst,\n const Matrix &A,\n const Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view,\n const hipStream_t &stream)\n{\n typedef typename TConfig::VecPrec ValueTypeB;\n int rowOff, nrows, nnz;\n A.getOffsetAndSizeForView(view, &rowOff, &nrows);\n A.getNnzForView(view, &nnz);\n\n hipsparseDirection_t direction = HIPSPARSE_DIRECTION_COLUMN;\n\n if ( A.getBlockFormat() == ROW_MAJOR )\n {\n direction = HIPSPARSE_DIRECTION_ROW;\n }\n\n bool has_offdiag = nnz != 0;\n\n if (has_offdiag )\n {\n bsrmv( Cusparse::get_instance().m_handle, direction, HIPSPARSE_OPERATION_NON_TRANSPOSE,\n nrows, A.get_num_cols(), nnz, &alphaConst,\n A.cuMatDescr,\n A.values.raw(),\n A.m_seq_offsets.raw() + rowOff,\n A.row_offsets.raw() + rowOff,\n A.col_indices.raw(),\n rowOff,\n A.get_block_dimx(),\n x.raw(),\n &betaConst,\n y.raw() + rowOff * A.get_block_dimx(),\n stream);\n }\n\n if (A.hasProps(DIAG))\n {\n ValueTypeB beta;\n\n if (!has_offdiag)\n {\n beta = betaConst;\n }\n else\n {\n beta = types::util::get_one();\n }\n\n bsrmv( Cusparse::get_instance().m_handle, direction, HIPSPARSE_OPERATION_NON_TRANSPOSE,\n nrows, A.get_num_cols(), A.get_num_rows(), &alphaConst,\n A.cuMatDescr,\n A.values.raw() + A.diagOffset()*A.get_block_size(),\n A.m_seq_offsets.raw(),\n A.m_seq_offsets.raw() + rowOff,\n A.m_seq_offsets.raw(),\n rowOff,\n A.get_block_dimx(),\n x.raw(), &beta,\n y.raw() + rowOff * A.get_block_dimx(),\n stream);\n }\n}\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_internal_with_mask( const typename TConfig::VecPrec alphaConst,\n const Matrix &A,\n const Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view,\n const hipStream_t &stream)\n{\n if (A.is_matrix_singleGPU())\n {\n FatalError(\"Should not be here in bsrmv_internal_with_mask\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if(view != INTERIOR && view != BOUNDARY)\n {\n FatalError(\"Only INTERIOR and BOUNDARY views supported for bsrmv_internal_with_mask\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename TConfig::VecPrec ValueType;\n hipsparseDirection_t direction = HIPSPARSE_DIRECTION_COLUMN;\n\n if ( A.getBlockFormat() == ROW_MAJOR )\n {\n direction = HIPSPARSE_DIRECTION_ROW;\n }\n\n const int *start_offsets, *end_offsets;\n start_offsets = A.row_offsets.raw();\n end_offsets = A.row_offsets.raw() + 1;\n typedef typename Matrix::index_type index_type;\n\n int offset, nrows, nnz;\n A.getOffsetAndSizeForView(view, &offset, &nrows);\n A.getNnzForView(view, &nnz);\n\n if (nrows <= 0)\n {\n return; // nothing to do, early exit\n }\n\n bool has_offdiag = nnz != 0;\n\n if (has_offdiag)\n {\n cusparseCheckError(hipsparseSetStream(Cusparse::get_instance().m_handle, stream));\n bsrxmv_internal( Cusparse::get_instance().m_handle, direction, HIPSPARSE_OPERATION_NON_TRANSPOSE, nrows,\n nrows, A.get_num_cols(), nnz, &alphaConst,\n A.cuMatDescr,\n A.values.raw(),\n A.manager->getRowsListForView(view).raw(),\n start_offsets, end_offsets, A.col_indices.raw(),\n A.get_block_dimx(),\n x.raw(), &betaConst,\n y.raw() );\n\n // Reset to default stream\n cusparseCheckError(hipsparseSetStream(Cusparse::get_instance().m_handle, 0));\n }\n\n if (A.hasProps(DIAG))\n {\n FatalError(\"Diag not supported in multiply with mask\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_internal_with_mask_restriction( const typename TConfig::VecPrec alphaConst,\n const Matrix &R,\n const Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view,\n const hipStream_t &stream,\n Matrix &P)\n{\n if (P.is_matrix_singleGPU())\n {\n FatalError(\"Should not be here in bsrmv_internal_with_mask_with_restriction\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if(view != OWNED && view != HALO1)\n {\n FatalError(\"View not supported in restriction operation\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename TConfig::VecPrec ValueType;\n hipsparseDirection_t direction = HIPSPARSE_DIRECTION_COLUMN;\n\n if ( R.getBlockFormat() == ROW_MAJOR )\n {\n direction = HIPSPARSE_DIRECTION_ROW;\n }\n\n int rowOff, nrows, nnz;\n R.getOffsetAndSizeForView(view, &rowOff, &nrows);\n R.getNnzForView(view, &nnz);\n\n bool has_offdiag = nnz != 0;\n typedef typename Matrix::index_type index_type;\n\n if (nrows <= 0)\n {\n return; // nothing to do, early exit\n }\n\n if (has_offdiag)\n {\n bsrmv( Cusparse::get_instance().m_handle, direction, HIPSPARSE_OPERATION_NON_TRANSPOSE,\n nrows, R.get_num_cols(), nnz, &alphaConst,\n R.cuMatDescr,\n R.values.raw(),\n R.m_seq_offsets.raw() + rowOff,\n R.row_offsets.raw() + rowOff,\n R.col_indices.raw(),\n rowOff,\n R.get_block_dimx(),\n x.raw(),\n &betaConst,\n y.raw() + rowOff * R.get_block_dimx(),\n stream);\n }\n\n if (R.hasProps(DIAG))\n {\n FatalError(\"Diag not supported in multiply with mask\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_internal( const typename TConfig::VecPrec alphaConst,\n const Matrix &A,\n const typename Matrix::MVector &E,\n const Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view,\n const hipStream_t &stream)\n{\n typedef typename TConfig::VecPrec ValueType;\n int rowOff, nrows, nnz;\n A.getOffsetAndSizeForView(view, &rowOff, &nrows);\n A.getNnzForView(view, &nnz);\n\n hipsparseDirection_t direction = A.getBlockFormat() == ROW_MAJOR ? HIPSPARSE_DIRECTION_ROW : HIPSPARSE_DIRECTION_COLUMN;\n\n bsrmv( Cusparse::get_instance().m_handle, direction, HIPSPARSE_OPERATION_NON_TRANSPOSE,\n nrows, A.get_num_cols(), nnz, &alphaConst,\n A.cuMatDescr,\n E.raw(),\n A.m_seq_offsets.raw(),\n A.m_seq_offsets.raw() + rowOff,\n A.m_seq_offsets.raw(),\n rowOff,\n A.get_block_dimx(),\n x.raw(), &betaConst,\n y.raw() + rowOff * A.get_block_dimx(),\n stream);\n}\n\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_internal( ColumnColorSelector columnColorSelector,\n const int color,\n const typename TConfig::VecPrec alphaConst,\n const Matrix &A,\n const Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view,\n const hipStream_t &stream)\n{\n typedef typename TConfig::VecPrec ValueType;\n\n if (!A.hasProps(COLORING))\n {\n FatalError(\"Matrix is not colored, exiting\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (color < 0 || color >= A.getMatrixColoring().getNumColors())\n {\n FatalError(\"Unknown color\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n typedef typename Matrix::index_type index_type;\n index_type colorStart = 0;\n\n if ( !(view & INTERIOR) )\n {\n colorStart = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[color];\n }\n else\n {\n colorStart = A.getMatrixColoring().getOffsetsRowsPerColor()[color];\n }\n\n index_type colorNum = 0;\n\n if ( view == A.getViewInterior() )\n {\n colorNum = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[color];\n }\n else\n {\n colorNum = A.getMatrixColoring().getOffsetsRowsPerColor()[color + 1];\n }\n\n colorNum -= colorStart;\n\n if (colorNum <= 0)\n {\n return; // nothing to do, early exit\n }\n\n if (columnColorSelector == DIAG_COL)\n {\n FatalError(\"DIAG_COL has been disabled to avoid allocating diag_offsets\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n const int *start_offsets, *end_offsets;\n\n switch (columnColorSelector)\n {\n case ALL_COLORS:\n start_offsets = A.row_offsets.raw();\n end_offsets = A.row_offsets.raw() + 1;\n break;\n\n case SMALLER_COLORS:\n start_offsets = A.row_offsets.raw();\n end_offsets = A.m_smaller_color_offsets.raw();\n break;\n\n case LARGER_COLORS:\n start_offsets = A.m_larger_color_offsets.raw();\n end_offsets = A.row_offsets.raw() + 1;\n break;\n\n case DIAG_COL:\n start_offsets = A.diag.raw();\n end_offsets = A.m_diag_end_offsets.raw();\n break;\n\n default:\n FatalError(\"Unknown color selector\", AMGX_ERR_CORE);\n }\n\n hipsparseDirection_t direction = HIPSPARSE_DIRECTION_COLUMN;\n\n if ( A.getBlockFormat() == ROW_MAJOR )\n {\n direction = HIPSPARSE_DIRECTION_ROW;\n }\n\n bool has_offdiag = A.get_num_nz() != 0;\n cusparseCheckError(hipsparseSetStream(Cusparse::get_instance().m_handle, stream));\n\n if (has_offdiag)\n {\n bsrxmv_internal( Cusparse::get_instance().m_handle, direction, HIPSPARSE_OPERATION_NON_TRANSPOSE, colorNum,\n A.get_num_rows(), A.get_num_cols(), A.get_num_nz(), &alphaConst,\n A.cuMatDescr,\n A.values.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + colorStart,\n start_offsets, end_offsets, A.col_indices.raw(),\n A.get_block_dimx(),\n x.raw(), &betaConst,\n y.raw() );\n }\n\n if (A.hasProps(DIAG) && columnColorSelector == ALL_COLORS)\n {\n typename TConfig::VecPrec beta = betaConst;\n\n if (has_offdiag)\n {\n beta = types::util::get_one();\n }\n\n bsrxmv_internal( Cusparse::get_instance().m_handle, direction, HIPSPARSE_OPERATION_NON_TRANSPOSE, colorNum,\n A.get_num_rows(), A.get_num_cols(), A.get_num_rows(), &alphaConst,\n A.cuMatDescr,\n A.values.raw() + A.diagOffset()*A.get_block_size(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + colorStart,\n A.m_seq_offsets.raw(),\n NULL,\n A.m_seq_offsets.raw(),\n A.get_block_dimx(),\n x.raw(), &beta,\n y.raw() );\n }\n\n // Reset to default stream\n cusparseCheckError(hipsparseSetStream(Cusparse::get_instance().m_handle, 0));\n}\n\ntemplate< class TConfig >\nvoid Cusparse::bsrmv_internal( const int color,\n const typename TConfig::VecPrec alphaConst,\n const Matrix &A,\n const typename Matrix::MVector &E,\n const Vector &x,\n const typename TConfig::VecPrec betaConst,\n Vector &y,\n ViewType view,\n const hipStream_t &stream)\n{\n typedef typename TConfig::VecPrec ValueType;\n\n if ( !A.hasProps(COLORING) )\n {\n FatalError(\"Matrix is not colored, exiting\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if ( color < 0 || color >= A.getMatrixColoring().getNumColors() )\n {\n FatalError(\"Unknown color\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n typedef typename Matrix::index_type index_type;\n //const index_type colorStart = ((view & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[color] : A.getMatrixColoring().getOffsetsRowsPerColor()[color];\n //const index_type colorNum = ((view == A.getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[color] : A.getMatrixColoring().getOffsetsRowsPerColor()[color+1]) - colorStart;\n //const index_type colorStart= A.getMatrixColoring().getOffsetsRowsPerColor()[color];\n //const index_type colorNum= A.getMatrixColoring().getOffsetsRowsPerColor()[color+1] - colorStart;\n index_type colorStart = 0;\n\n if ( !(view & INTERIOR) )\n {\n colorStart = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[color];\n }\n else\n {\n colorStart = A.getMatrixColoring().getOffsetsRowsPerColor()[color];\n }\n\n index_type colorNum = 0;\n\n if ( view == A.getViewInterior() )\n {\n colorNum = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[color];\n }\n else\n {\n colorNum = A.getMatrixColoring().getOffsetsRowsPerColor()[color + 1];\n }\n\n colorNum -= colorStart;\n\n if (colorNum <= 0)\n {\n return; // nothing to do, early exit\n }\n\n hipsparseDirection_t direction = HIPSPARSE_DIRECTION_COLUMN;\n\n if ( A.getBlockFormat() == ROW_MAJOR )\n {\n direction = HIPSPARSE_DIRECTION_ROW;\n }\n\n cusparseCheckError(hipsparseSetStream(Cusparse::get_instance().m_handle, stream));\n\n bsrxmv_internal( Cusparse::get_instance().m_handle, direction, HIPSPARSE_OPERATION_NON_TRANSPOSE, colorNum,\n A.get_num_rows(), A.get_num_cols(), A.get_num_nz(), &alphaConst,\n A.cuMatDescr,\n E.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + colorStart,\n A.m_seq_offsets.raw(),\n NULL,\n A.m_seq_offsets.raw(),\n A.get_block_dimx(),\n x.raw(), &betaConst,\n y.raw() );\n // Reset to default stream\n cusparseCheckError(hipsparseSetStream(Cusparse::get_instance().m_handle, 0));\n}\n\n// Simple custom implementation of matrix-vector product that has only 1 kernel.\ntemplate\n__global__ void csrmv(\n int nrows,\n const T alpha,\n const T* __restrict__ csrVal,\n const int* __restrict__ csrRow,\n const int* __restrict__ csrCol,\n const T* __restrict__ x,\n const T beta,\n T* __restrict__ y)\n{\n for(int i = threadIdx.x + blockIdx.x*blockDim.x; i < nrows; i += blockDim.x*gridDim.x)\n {\n T y_tmp = amgx::types::util::get_zero();\n\n int row_b = csrRow[i];\n int row_e = csrRow[i+1];\n\n // Unrolling is important for performance here.\n // Possible to squeeze more performance out of the key kernels if we\n // measure the sparsity and use it to inform unrolling.\n for (int col = row_b; col < row_e; col += UNROLL)\n {\n#pragma unroll UNROLL\n for(int off = 0; off < UNROLL; ++off)\n {\n int c = col + off;\n if(c < row_e) y_tmp = alpha * csrVal[c] * x[csrCol[c]] + y_tmp;\n }\n }\n\n // Don't read y unnecessarily\n if(amgx::types::util::is_zero(beta))\n {\n y[i] = y_tmp;\n }\n else\n {\n y[i] = beta*y[i] + y_tmp;\n }\n }\n}\n\ntemplate\ninline void generic_SpMV(hipsparseHandle_t handle, hipsparseOperation_t trans,\n int mb, int nb, int nnzb, int rowOff,\n const MatType *alpha,\n const MatType *val,\n const IndType *rowPtr,\n const IndType *colInd,\n const VecType *x,\n const VecType *beta,\n VecType *y,\n hipDataType matType,\n hipDataType vecType,\n const hipStream_t& stream)\n{\n constexpr int cta_size = 128;\n const int sm_count = getSMCount();\n\n // Assuming that csrmv will be more efficient than cuSPARSE for row counts \n // that are lower than the 3 times the total number of threads\n // cuSPARSE does not like the offsetting required when latency hiding\n // it's possible to reverse the offsets, but requires extra kernel invocation \n // and usually the dependent part of the call is smaller\n if(rowOff > 0 || mb < cta_size * sm_count * 3)\n {\n // Custom single-kernel SpMV, we could actually determine unroll factor\n // more accurately here by checking non-zeros per row\n constexpr int unroll_factor = 16;\n int nblocks = mb / cta_size + 1;\n csrmv<<>>(mb, *alpha, val, rowPtr, colInd, x, *beta, y);\n cudaCheckError();\n }\n else\n {\n IndType* rows = const_cast(rowPtr);\n IndType* cols = const_cast(colInd);\n MatType* vals = const_cast(val);\n\n hipsparseSpMatDescr_t matA_descr;\n hipsparseDnVecDescr_t vecX_descr;\n hipsparseDnVecDescr_t vecY_descr;\n cusparseCheckError(hipsparseCreateDnVec(&vecX_descr, nb, const_cast(x), vecType));\n cusparseCheckError(hipsparseCreateDnVec(&vecY_descr, mb, const_cast(y), vecType));\n cusparseCheckError(\n hipsparseCreateCsr(&matA_descr, mb, nb, nnzb, const_cast(rows), const_cast(cols),\n const_cast(vals), HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_BASE_ZERO, matType));\n\n // Use ALG2 for deterministic results when determinism_flag is set\n hipsparseSpMVAlg_t spmv_alg = Cusparse::get_instance().is_deterministic() ? HIPSPARSE_CSRMV_ALG2 : HIPSPARSE_CSRMV_ALG1;\n\n size_t bufferSize = 0;\n cusparseCheckError(hipsparseSpMV_bufferSize(handle, trans, alpha, matA_descr, vecX_descr, beta, vecY_descr, matType, spmv_alg, &bufferSize));\n\n void* dBuffer = NULL;\n if(bufferSize > 0)\n {\n hipError_t cuda_rc = amgx::memory::hipMallocAsync(&dBuffer, bufferSize, stream);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMallocAsync failed in generic_SpMV\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n\n cusparseCheckError(hipsparseSpMV(handle, trans, alpha, matA_descr, vecX_descr, beta, vecY_descr, matType, spmv_alg, dBuffer) );\n\n cusparseCheckError(hipsparseDestroySpMat(matA_descr));\n cusparseCheckError(hipsparseDestroyDnVec(vecX_descr));\n cusparseCheckError(hipsparseDestroyDnVec(vecY_descr));\n\n if(bufferSize > 0)\n {\n hipError_t cuda_rc = amgx::memory::hipFreeAsync(dBuffer);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipFreeAsync failed in generic_SpMV\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n }\n}\n\ninline void Cusparse::bsrmv( hipsparseHandle_t handle, hipsparseDirection_t dir, hipsparseOperation_t trans,\n int mb, int nb, int nnzb,\n const float *alpha,\n const hipsparseMatDescr_t descr,\n const float *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrColInd,\n int rowOff,\n int blockDim,\n const float *x,\n const float *beta,\n float *y,\n const hipStream_t& stream)\n{\n // Run cuSparse on selected stream\n cusparseCheckError(hipsparseSetStream(handle, stream));\n\n if (blockDim == 1)\n {\n generic_SpMV(handle, trans, mb, nb, nnzb, rowOff, alpha, bsrVal, bsrRowPtr, bsrColInd, x, beta, y, HIP_R_32F, HIP_R_32F, stream);\n }\n else\n {\n cusparseCheckError(hipsparseSbsrmv(handle, dir, trans, mb, nb, nnzb, alpha, descr, bsrVal, bsrRowPtr, bsrColInd, blockDim, x, beta, y));\n }\n\n // Reset cuSparse to default stream\n cusparseCheckError(hipsparseSetStream(handle, 0));\n}\n\ninline void Cusparse::bsrmv( hipsparseHandle_t handle, hipsparseDirection_t dir, hipsparseOperation_t trans,\n int mb, int nb, int nnzb,\n const double *alpha,\n const hipsparseMatDescr_t descr,\n const double *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrColInd,\n int rowOff,\n int blockDim,\n const double *x,\n const double *beta,\n double *y,\n const hipStream_t& stream)\n{\n // Run cuSparse on selected stream\n cusparseCheckError(hipsparseSetStream(handle, stream));\n\n if (blockDim == 1)\n {\n generic_SpMV(handle, trans, mb, nb, nnzb, rowOff, alpha, bsrVal, bsrRowPtr, bsrColInd, x, beta, y, HIP_R_64F, HIP_R_64F, stream);\n }\n else\n {\n cusparseCheckError(hipsparseDbsrmv(handle, dir, trans, mb, nb, nnzb, alpha, descr, bsrVal, bsrRowPtr, bsrColInd, blockDim, x, beta, y));\n }\n\n // Reset cuSparse to default stream\n cusparseCheckError(hipsparseSetStream(handle, 0));\n}\n\ninline void Cusparse::bsrmv( hipsparseHandle_t handle, hipsparseDirection_t dir, hipsparseOperation_t trans,\n int mb, int nb, int nnzb,\n const double *alpha,\n const hipsparseMatDescr_t descr,\n const float *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrColInd,\n int rowOff,\n int blockDim,\n const double *x,\n const double *beta,\n double *y,\n const hipStream_t& stream)\n{\n #if 0\n // Run cuSparse on selected stream\n cusparseCheckError(hipsparseSetStream(handle, stream));\n\n const double *d_bsrVal = reinterpret_cast(const_cast(bsrVal)); // this works due to private API call in the matrix initialization which sets cusparse matrix description in the half precision mode\n cusparseCheckError(hipsparseDbsrxmv(handle, dir, trans, mb, mb, nb, nnzb, alpha, descr, d_bsrVal, bsrMaskPtr, bsrRowPtr, bsrRowPtr + 1, bsrColInd, blockDim, x, beta, y));\n\n // Reset cuSparse to default stream\n cusparseCheckError(hipsparseSetStream(handle, 0));\n #else\n FatalError(\"Mixed precision modes not currently supported for CUDA 10.1 or later.\", AMGX_ERR_NOT_IMPLEMENTED);\n #endif\n}\n\n// Custom implementation of matrix-vector product to replace the original bsrxmv,\n// but with block size of 1.\ntemplate\n__global__ void csrxmv(\n int sizeOfMask,\n const T alpha,\n const T* __restrict__ csrVal,\n const int* __restrict__ csrMask,\n const int* __restrict__ csrRow,\n const int* __restrict__ csrCol,\n const T* __restrict__ x,\n const T beta,\n T* __restrict__ y)\n{\n for(int i = threadIdx.x + blockIdx.x*blockDim.x; i < sizeOfMask; i += blockDim.x*gridDim.x)\n {\n int row = csrMask[i];\n T y_tmp = amgx::types::util::get_zero();\n\n int row_b = csrRow[row];\n int row_e = csrRow[row+1];\n\n // Unrolling is important for performance here.\n // Possible to squeeze more performance out of the key kernels if we\n // measure the sparsity and use it to inform unrolling.\n for (int col = row_b; col < row_e; col += UNROLL)\n {\n#pragma unroll UNROLL\n for(int off = 0; off < UNROLL; ++off)\n {\n int c = col + off;\n if(c < row_e) y_tmp = alpha * csrVal[c] * x[csrCol[c]] + y_tmp;\n }\n }\n\n // Don't read y unnecessarily\n if(amgx::types::util::is_zero(beta))\n {\n y[row] = y_tmp;\n }\n else\n {\n y[row] = beta*y[row] + y_tmp;\n }\n }\n}\n\n// Replaces the functionality of cusparse?bsrxmv for blockDim == 1\ntemplate\ninline void Xcsrxmv( hipsparseHandle_t handle, hipsparseDirection_t dir, hipsparseOperation_t trans, int sizeOfMask,\n int mb, int nb, int nnzb,\n const T *alpha,\n const hipsparseMatDescr_t descr,\n const T *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrEndPtr,\n const int *bsrColInd,\n int blockDim,\n const T *x,\n const T *beta,\n T *y)\n{\n if(blockDim != 1)\n {\n FatalError(\"Xcsrxmv only to be called with scalar matrices.\", AMGX_ERR_INTERNAL);\n }\n if (trans != HIPSPARSE_OPERATION_NON_TRANSPOSE)\n {\n FatalError(\"Cannot currently latency hide a transposed matrix.\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n if (dir != HIPSPARSE_DIRECTION_ROW)\n {\n FatalError(\"Cannot currently latency hide if matrix is not row major.\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n constexpr int cta_size = 128;\n constexpr int unroll_factor = 16;\n int nblocks = sizeOfMask / cta_size + 1;\n csrxmv<<>>(sizeOfMask, *alpha, bsrVal, bsrMaskPtr, bsrRowPtr, bsrColInd, x, *beta, y);\n cudaCheckError();\n}\n\n// overloaded C++ wrappers for cusparse?bsrxmv\n// bsrxmv\n// matrix - float\n// vector - float\ninline void Cusparse::bsrxmv_internal(hipsparseHandle_t handle, hipsparseDirection_t dir, hipsparseOperation_t trans, int sizeOfMask,\n int mb, int nb, int nnzb,\n const float *alpha,\n const hipsparseMatDescr_t descr,\n const float *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrEndPtr,\n const int *bsrColInd,\n int blockDim,\n const float *x,\n const float *beta,\n float *y)\n{\n if (bsrEndPtr == NULL && bsrMaskPtr == NULL)\n {\n cusparseCheckError(HIPSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n }\n else\n {\n if (bsrEndPtr == NULL)\n {\n bsrEndPtr = bsrRowPtr + 1;\n }\n\n if (blockDim == 1)\n {\n Xcsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y);\n }\n else\n {\n cusparseCheckError(hipsparseSbsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y));\n }\n }\n}\n\n// bsrxmv\n// matrix - float\n// vector - double\ninline void Cusparse::bsrxmv_internal( hipsparseHandle_t handle, hipsparseDirection_t dir, hipsparseOperation_t trans, int sizeOfMask,\n int mb, int nb, int nnzb,\n const double *alpha,\n const hipsparseMatDescr_t descr,\n const float *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrEndPtr,\n const int *bsrColInd,\n int blockDim,\n const double *x,\n const double *beta,\n double *y)\n{\n cusparseCheckError(HIPSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n}\n\n// bsrxmv\n// matrix - double\n// vector - double\ninline void Cusparse::bsrxmv_internal( hipsparseHandle_t handle, hipsparseDirection_t dir, hipsparseOperation_t trans, int sizeOfMask,\n int mb, int nb, int nnzb,\n const double *alpha,\n const hipsparseMatDescr_t descr,\n const double *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrEndPtr,\n const int *bsrColInd,\n int blockDim,\n const double *x,\n const double *beta,\n double *y)\n{\n if (bsrEndPtr == NULL && bsrMaskPtr == NULL)\n {\n cusparseCheckError(HIPSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n }\n else\n {\n if (bsrEndPtr == NULL)\n {\n bsrEndPtr = bsrRowPtr + 1;\n }\n\n if (blockDim == 1)\n {\n Xcsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y);\n }\n else\n {\n cusparseCheckError(hipsparseDbsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y));\n }\n }\n}\n\n\n\ninline void Cusparse::bsrmv( hipsparseHandle_t handle, hipsparseDirection_t dir, hipsparseOperation_t trans,\n int mb, int nb, int nnzb,\n const hipComplex *alpha,\n const hipsparseMatDescr_t descr,\n const hipComplex *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrColInd,\n int rowOff,\n int blockDim,\n const hipComplex *x,\n const hipComplex *beta,\n hipComplex *y,\n const hipStream_t& stream)\n{\n // Run cuSparse on selected stream\n cusparseCheckError(hipsparseSetStream(handle, stream));\n\n if (blockDim == 1)\n {\n generic_SpMV(handle, trans, mb, nb, nnzb, rowOff, alpha, bsrVal, bsrRowPtr, bsrColInd, x, beta, y, HIP_C_32F, HIP_C_32F, stream);\n }\n else\n {\n cusparseCheckError(hipsparseCbsrmv(handle, dir, trans, mb, nb, nnzb, alpha, descr, bsrVal, bsrRowPtr, bsrColInd, blockDim, x, beta, y));\n }\n\n // Reset cuSparse to default stream\n cusparseCheckError(hipsparseSetStream(handle, 0));\n}\n\ninline void Cusparse::bsrmv( hipsparseHandle_t handle, hipsparseDirection_t dir, hipsparseOperation_t trans,\n int mb, int nb, int nnzb,\n const hipDoubleComplex *alpha,\n const hipsparseMatDescr_t descr,\n const hipDoubleComplex *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrColInd,\n int rowOff,\n int blockDim,\n const hipDoubleComplex *x,\n const hipDoubleComplex *beta,\n hipDoubleComplex *y,\n const hipStream_t& stream)\n{\n // Run cuSparse on selected stream\n cusparseCheckError(hipsparseSetStream(handle, stream));\n\n if (blockDim == 1)\n {\n generic_SpMV(handle, trans, mb, nb, nnzb, rowOff, alpha, bsrVal, bsrRowPtr, bsrColInd, x, beta, y, HIP_C_64F, HIP_C_64F, stream);\n }\n else\n {\n cusparseCheckError(hipsparseZbsrmv(handle, dir, trans, mb, nb, nnzb, alpha, descr, bsrVal, bsrRowPtr, bsrColInd, blockDim, x, beta, y));\n }\n\n // Reset cuSparse to default stream\n cusparseCheckError(hipsparseSetStream(handle, 0));\n}\n\ninline void Cusparse::bsrmv( hipsparseHandle_t handle, hipsparseDirection_t dir, hipsparseOperation_t trans,\n int mb, int nb, int nnzb,\n const hipDoubleComplex *alpha,\n const hipsparseMatDescr_t descr,\n const hipComplex *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrColInd,\n int rowOff,\n int blockDim,\n const hipDoubleComplex *x,\n const hipDoubleComplex *beta,\n hipDoubleComplex *y,\n const hipStream_t& stream)\n{\n #if 0\n // Run cuSparse on selected stream\n cusparseCheckError(hipsparseSetStream(handle, stream));\n\n const hipDoubleComplex *d_bsrVal = reinterpret_cast(const_cast(bsrVal));\n cusparseCheckError(hipsparseZbsrxmv(handle, dir, trans, mb, mb, nb, nnzb, alpha, descr, d_bsrVal, bsrMaskPtr, bsrRowPtr, bsrRowPtr + 1, bsrColInd, blockDim, x, beta, y));\n\n // Reset cuSparse to default stream\n cusparseCheckError(hipsparseSetStream(handle, 0));\n #else\n FatalError(\"Mixed precision modes not currently supported for CUDA 10.1 or later.\", AMGX_ERR_NOT_IMPLEMENTED);\n #endif\n}\n\n\n// overloaded C++ wrappers for cusparse?bsrxmv\n// bsrxmv\n// matrix - cuComplex\n// vector - cuComplex\ninline void Cusparse::bsrxmv_internal( hipsparseHandle_t handle, hipsparseDirection_t dir, hipsparseOperation_t trans, int sizeOfMask,\n int mb, int nb, int nnzb,\n const hipComplex *alpha,\n const hipsparseMatDescr_t descr,\n const hipComplex *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrEndPtr,\n const int *bsrColInd,\n int blockDim,\n const hipComplex *x,\n const hipComplex *beta,\n hipComplex *y)\n{\n if (bsrEndPtr == NULL && bsrMaskPtr == NULL)\n {\n cusparseCheckError(HIPSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n }\n else\n {\n if (bsrEndPtr == NULL)\n {\n bsrEndPtr = bsrRowPtr + 1;\n }\n\n if (blockDim == 1)\n {\n Xcsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y);\n }\n else\n {\n cusparseCheckError(hipsparseCbsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y));\n }\n }\n}\n\n// bsrxmv\n// matrix - cuComplex\n// vector - cuDoubleComplex\ninline void Cusparse::bsrxmv_internal( hipsparseHandle_t handle, hipsparseDirection_t dir, hipsparseOperation_t trans, int sizeOfMask,\n int mb, int nb, int nnzb,\n const hipDoubleComplex *alpha,\n const hipsparseMatDescr_t descr,\n const hipComplex *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrEndPtr,\n const int *bsrColInd,\n int blockDim,\n const hipDoubleComplex *x,\n const hipDoubleComplex *beta,\n hipDoubleComplex *y)\n{\n cusparseCheckError(HIPSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n}\n\n// bsrxmv\n// matrix - cuDoubleComplex\n// vector - cuDoubleComplex\ninline void Cusparse::bsrxmv_internal( hipsparseHandle_t handle, hipsparseDirection_t dir, hipsparseOperation_t trans, int sizeOfMask,\n int mb, int nb, int nnzb,\n const hipDoubleComplex *alpha,\n const hipsparseMatDescr_t descr,\n const hipDoubleComplex *bsrVal,\n const int *bsrMaskPtr,\n const int *bsrRowPtr,\n const int *bsrEndPtr,\n const int *bsrColInd,\n int blockDim,\n const hipDoubleComplex *x,\n const hipDoubleComplex *beta,\n hipDoubleComplex *y)\n{\n if (bsrEndPtr == NULL && bsrMaskPtr == NULL)\n {\n cusparseCheckError(HIPSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n }\n else\n {\n if (bsrEndPtr == NULL)\n {\n bsrEndPtr = bsrRowPtr + 1;\n }\n\n if (blockDim == 1)\n {\n Xcsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y);\n }\n else\n {\n cusparseCheckError(hipsparseZbsrxmv(handle, dir, trans, sizeOfMask, mb, nb, nnzb, alpha, descr, bsrVal, bsrMaskPtr, bsrRowPtr, bsrEndPtr, bsrColInd, blockDim, x, beta, y));\n }\n }\n}\n\n\nnamespace\n{\ntemplate\ninline void\ngeneric_SpMM(hipsparseHandle_t handle, hipsparseOperation_t transA,\n int m, int n, int k, int nnz,\n int ldb, int ldc,\n const MatType *alpha,\n const MatType *Avals,\n const MatType *Bvals,\n MatType *Cvals,\n const IndType *rowPtr,\n const IndType *colInd,\n const MatType *beta,\n hipDataType matType)\n{\n // Create the matrix descriptors\n hipsparseSpMatDescr_t matA_descr;\n hipsparseDnMatDescr_t matB_descr;\n hipsparseDnMatDescr_t matC_descr;\n cusparseCheckError(\n hipsparseCreateCsr(&matA_descr, m, k, nnz, const_cast(rowPtr), const_cast(colInd),\n const_cast(Avals), HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_BASE_ZERO, matType));\n cusparseCheckError(\n hipsparseCreateDnMat(&matB_descr, k, n, ldb, const_cast(Bvals), matType, HIPSPARSE_ORDER_COL));\n cusparseCheckError(\n hipsparseCreateDnMat(&matC_descr, m, n, ldc, const_cast(Cvals), matType, HIPSPARSE_ORDER_COL));\n\n // Check if a buffer is required, and if so allocate it using caching allocator\n size_t bufferSize = 0;\n cusparseCheckError(\n hipsparseSpMM_bufferSize(handle, transA, HIPSPARSE_OPERATION_NON_TRANSPOSE, alpha, matA_descr, matB_descr,\n beta, matC_descr, matType, HIPSPARSE_SPMM_ALG_DEFAULT, &bufferSize));\n\n void* dBuffer = NULL;\n if(bufferSize > 0)\n {\n hipError_t cuda_rc = amgx::memory::hipMallocAsync(&dBuffer, bufferSize);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMallocAsync failed in generic_SpMM\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n\n // Compute the sparse matrix - dense matrix product\n cusparseCheckError(\n hipsparseSpMM(handle, transA, HIPSPARSE_OPERATION_NON_TRANSPOSE, alpha, matA_descr, matB_descr, beta,\n matC_descr, matType, HIPSPARSE_SPMM_ALG_DEFAULT, dBuffer));\n\n // Clean up\n cusparseCheckError(hipsparseDestroySpMat(matA_descr));\n cusparseCheckError(hipsparseDestroyDnMat(matB_descr));\n cusparseCheckError(hipsparseDestroyDnMat(matC_descr));\n\n if(bufferSize > 0)\n {\n hipError_t cuda_rc = amgx::memory::hipFreeAsync(dBuffer);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipFreeAsync failed in generic_SpMM\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n}\n\nvoid\ncusparse_csrmm(hipsparseHandle_t handle, hipsparseOperation_t transA,\n int m, int n, int k, int nnz,\n const float *alpha,\n const hipsparseMatDescr_t descrA,\n const float *csrValA,\n const int *csrRowPtrA, const int *csrColIndA,\n const float *B, int ldb,\n const float *beta, float *C, int ldc)\n{\n generic_SpMM(handle, transA, m, n, k, nnz, ldb, ldc, alpha, csrValA, B, C, csrRowPtrA, csrColIndA, beta, HIP_R_32F);\n}\n\nvoid\ncusparse_csrmm(hipsparseHandle_t handle, hipsparseOperation_t transA,\n int m, int n, int k, int nnz,\n const double *alpha,\n const hipsparseMatDescr_t descrA,\n const float *csrValA,\n const int *csrRowPtrA, const int *csrColIndA,\n const double *B, int ldb,\n const double *beta, double *C, int ldc)\n{\n cusparseCheckError(HIPSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n}\n\nvoid\ncusparse_csrmm(hipsparseHandle_t handle, hipsparseOperation_t transA,\n int m, int n, int k, int nnz,\n const double *alpha,\n const hipsparseMatDescr_t descrA,\n const double *csrValA,\n const int *csrRowPtrA, const int *csrColIndA,\n const double *B, int ldb,\n const double *beta, double *C, int ldc)\n{\n generic_SpMM(handle, transA, m, n, k, nnz, ldb, ldc, alpha, csrValA, B, C, csrRowPtrA, csrColIndA, beta, HIP_R_64F);\n}\n\nvoid\ncusparse_csrmm(hipsparseHandle_t handle, hipsparseOperation_t transA,\n int m, int n, int k, int nnz,\n const hipComplex *alpha,\n const hipsparseMatDescr_t descrA,\n const hipComplex *csrValA,\n const int *csrRowPtrA, const int *csrColIndA,\n const hipComplex *B, int ldb,\n const hipComplex *beta, hipComplex *C, int ldc)\n{\n generic_SpMM(handle, transA, m, n, k, nnz, ldb, ldc, alpha, csrValA, B, C, csrRowPtrA, csrColIndA, beta, HIP_C_32F);\n}\n\nvoid\ncusparse_csrmm(hipsparseHandle_t handle, hipsparseOperation_t transA,\n int m, int n, int k, int nnz,\n const hipDoubleComplex *alpha,\n const hipsparseMatDescr_t descrA,\n const hipComplex *csrValA,\n const int *csrRowPtrA, const int *csrColIndA,\n const hipDoubleComplex *B, int ldb,\n const hipDoubleComplex *beta, hipDoubleComplex *C, int ldc)\n{\n cusparseCheckError(HIPSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED);\n}\n\nvoid\ncusparse_csrmm(hipsparseHandle_t handle, hipsparseOperation_t transA,\n int m, int n, int k, int nnz,\n const hipDoubleComplex *alpha,\n const hipsparseMatDescr_t descrA,\n const hipDoubleComplex *csrValA,\n const int *csrRowPtrA, const int *csrColIndA,\n const hipDoubleComplex *B, int ldb,\n const hipDoubleComplex *beta, hipDoubleComplex *C, int ldc)\n{\n generic_SpMM(handle, transA, m, n, k, nnz, ldb, ldc, alpha, csrValA, B, C, csrRowPtrA, csrColIndA, beta, HIP_C_64F);\n}\n}\n\n\ntemplate \nvoid Cusparse::csrmm(typename TConfig::VecPrec alpha,\n Matrix &A,\n Vector &V,\n typename TConfig::VecPrec beta,\n Vector &Res)\n{\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo(V, V.tag);\n }\n\n if (Res.get_num_rows() != A.get_num_rows() || Res.get_num_cols() != V.get_num_cols())\n {\n FatalError(\"Cusparse::csrmm error, dimensions of result matrix do not match input matrices.\", AMGX_ERR_INTERNAL);\n }\n\n hipsparseHandle_t handle = Cusparse::get_instance().m_handle;\n cusparse_csrmm(handle, HIPSPARSE_OPERATION_NON_TRANSPOSE,\n A.get_num_rows(), V.get_num_cols(), A.get_num_cols(),\n A.values.size(), &alpha, A.cuMatDescr,\n A.values.raw(), A.row_offsets.raw(), A.col_indices.raw(),\n V.raw(), V.get_lda(),\n &beta, Res.raw(), Res.get_lda());\n Res.dirtybit = 1;\n}\n\ntemplate \nvoid transpose_internal(hipsparseHandle_t handle, int nRows, int nCols, int nNz, const T* Avals, const int* Arows, const int* Acols, T* Bvals, int* Brows, int* Bcols, hipDataType valType)\n{\n size_t bufferSize;\n cusparseCheckError(hipsparseCsr2cscEx2_bufferSize(\n handle, nRows, nCols, nNz, Avals, Arows, Acols, Bvals, Brows, Bcols, valType,\n HIPSPARSE_ACTION_NUMERIC, HIPSPARSE_INDEX_BASE_ZERO, HIPSPARSE_CSR2CSC_ALG1, &bufferSize));\n\n void *buffer = nullptr;\n if (bufferSize > 0)\n {\n hipError_t cuda_rc = amgx::memory::hipMallocAsync(&buffer, bufferSize);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMallocAsync failed in transpose_internal\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n\n cusparseCheckError(hipsparseCsr2cscEx2(\n handle, nRows, nCols, nNz, Avals, Arows, Acols, Bvals, Brows, Bcols, valType,\n HIPSPARSE_ACTION_NUMERIC, HIPSPARSE_INDEX_BASE_ZERO, HIPSPARSE_CSR2CSC_ALG1, buffer));\n\n if(bufferSize > 0)\n {\n hipError_t cuda_rc = amgx::memory::hipFreeAsync(buffer);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipFreeAsync failed in transpose_internal\", AMGX_ERR_CUDA_FAILURE);\n }\n }\n}\n\nvoid transpose_internal(hipsparseHandle_t handle, int nRows, int nCols, int nNz, const float* Avals, const int* Arows, const int* Acols, float* Bvals, int* Brows, int* Bcols)\n{\n transpose_internal(handle, nRows, nCols, nNz, Avals, Arows, Acols, Bvals, Brows, Bcols, HIP_R_32F);\n}\nvoid transpose_internal(hipsparseHandle_t handle, int nRows, int nCols, int nNz, const double* Avals, const int* Arows, const int* Acols, double* Bvals, int* Brows, int* Bcols)\n{\n transpose_internal(handle, nRows, nCols, nNz, Avals, Arows, Acols, Bvals, Brows, Bcols, HIP_R_64F);\n}\nvoid transpose_internal(hipsparseHandle_t handle, int nRows, int nCols, int nNz, const hipComplex* Avals, const int* Arows, const int* Acols, hipComplex* Bvals, int* Brows, int* Bcols)\n{\n transpose_internal(handle, nRows, nCols, nNz, Avals, Arows, Acols, Bvals, Brows, Bcols, HIP_C_32F);\n}\nvoid transpose_internal(hipsparseHandle_t handle, int nRows, int nCols, int nNz, const hipDoubleComplex* Avals, const int* Arows, const int* Acols, hipDoubleComplex* Bvals, int* Brows, int* Bcols)\n{\n transpose_internal(handle, nRows, nCols, nNz, Avals, Arows, Acols, Bvals, Brows, Bcols, HIP_C_64F);\n}\n\ntemplate \nvoid Cusparse::transpose(const Matrix& A, Matrix& B, const int nRows, const int nNz)\n{\n hipsparseHandle_t handle = Cusparse::get_instance().m_handle;\n transpose_internal(handle, nRows, A.get_num_cols(), nNz,\n A.values.raw(), A.row_offsets.raw(), A.col_indices.raw(),\n B.values.raw(), B.row_offsets.raw(), B.col_indices.raw());\n}\n\ntemplate \nvoid Cusparse::transpose(const Matrix& A, Matrix& B)\n{\n hipsparseHandle_t handle = Cusparse::get_instance().m_handle;\n transpose_internal(handle, A.get_num_rows(), A.get_num_cols(), A.get_num_nz(),\n A.values.raw(), A.row_offsets.raw(), A.col_indices.raw(),\n B.values.raw(), B.row_offsets.raw(), B.col_indices.raw());\n}\n\n//#define AMGX_CASE_LINE(CASE) template class Cusparse::Type>;\n// AMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n//#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::bsrmv(const typename TemplateMode::Type::VecPrec , Matrix::Type>&, Vector::Type>&, const typename TemplateMode::Type::VecPrec, Vector::Type> &, ViewType);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::bsrxmv(const typename TemplateMode::Type::VecPrec , Matrix::Type>&, Vector::Type>&, const typename TemplateMode::Type::VecPrec, Vector::Type>&, typename Matrix::Type>::IVector&, ViewType);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::bsrmv_with_mask(const typename TemplateMode::Type::VecPrec , Matrix::Type>&, Vector::Type>&, const typename TemplateMode::Type::VecPrec, Vector::Type> &);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::bsrmv_with_mask_restriction(const typename TemplateMode::Type::VecPrec , Matrix::Type>&, Vector::Type>&, const typename TemplateMode::Type::VecPrec, Vector::Type> &, Matrix::Type>& );\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::bsrmv(const typename TemplateMode::Type::VecPrec , Matrix::Type>&, const typename Matrix::Type>::MVector&, Vector::Type>&, const typename TemplateMode::Type::VecPrec, Vector::Type> &, ViewType);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::bsrmv(ColumnColorSelector, const int, const typename TemplateMode::Type::VecPrec , Matrix::Type>&, Vector::Type>&, const typename TemplateMode::Type::VecPrec, Vector::Type> &, ViewType);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::bsrmv(const int, const typename TemplateMode::Type::VecPrec , Matrix::Type>&, const typename Matrix::Type>::MVector &, Vector::Type>&, const typename TemplateMode::Type::VecPrec, Vector::Type> &, ViewType);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::csrmm(typename TemplateMode::Type::VecPrec, Matrix::Type>&, Vector::Type>&, typename TemplateMode::Type::VecPrec, Vector::Type>&);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n template void Cusparse::transpose(const Matrix::Type>& A, Matrix::Type>& B); \\\n template void Cusparse::transpose(const Matrix::Type>& A, Matrix::Type>& B, const int nRows, const int nNz);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/amgx_cusparse.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#ifdef _WIN32\n#pragma warning (push)\n#pragma warning (disable : 4244 4267 4521)\n#endif\n#ifdef _WIN32\n#pragma warning (pop)\n#endif\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"strided_reduction.h\"\n\n#include \"amgx_timer.h\"\n#include \"amgx_types/util.h\"\n#include \"thrust_wrapper.h\"\n\nnamespace amgx\n{\n\n/**********************************************************\n * Returns the norm of a vector\n *********************************************************/\n\ntemplate\ntypename types::PODTypes::type get_norm(const MatrixType &A, const VectorType &r, const NormType norm_type, typename types::PODTypes::type norm_factor)\n{\n typedef typename types::PODTypes::type value_type;\n value_type nrm;\n int offset, size;\n A.getOffsetAndSizeForView(OWNED, &offset, &size);\n\n if (norm_type == L1 || norm_type == L1_SCALED)\n {\n nrm = nrm1(r, offset, size);\n\n if (A.is_matrix_distributed())\n {\n A.getManager()->global_reduce_sum(&nrm);\n }\n\n return (norm_type == L1_SCALED) ? nrm / norm_factor : nrm;\n }\n else if (norm_type == L2)\n {\n nrm = nrm2(r, offset, size);\n\n if (A.is_matrix_distributed())\n {\n nrm = nrm * nrm;\n A.getManager()->global_reduce_sum(&nrm);\n nrm = sqrt(nrm);\n }\n \n return nrm;\n }\n else if (norm_type == LMAX)\n {\n nrm = nrmmax(r, offset, size);\n\n if (A.is_matrix_distributed())\n {\n typedef TemplateConfig::vec_prec, MatrixType::TConfig::matPrec, MatrixType::TConfig::indPrec> hvector_type;\n typedef Vector HVector;\n //collect values from all neighbors, and do the \"reduction\" part\n std::vector values(0);\n HVector my_nrm(1);\n my_nrm[0] = nrm;\n A.getManager()->getComms()->global_reduce(values, my_nrm, A, 3);\n\n for (int j = 0; j < values.size(); j++)\n {\n nrm = (nrm > values[j][0] ? nrm : values[j][0]);\n }\n }\n\n return nrm;\n }\n\n return -1;\n}\n\ntemplate \nclass Norm_1x1;\n\ntemplate \nvoid get_1x1_norm(const MatrixType &A, const VectorType &r, const int block_size, const NormType norm_type, PlainVectorType &block_nrm, typename types::PODTypes::type norm_factor)\n{\n Norm_1x1::get_1x1_norm(A, r, block_size, norm_type, block_nrm, norm_factor);\n}\n\ntemplate \nclass Norm_1x1< Vector >, MatrixType, Vector>>\n{\n public:\n typedef Vector > Vector_h;\n typedef Vector> PODHostVec; // return type\n typedef typename Vector_h::value_type ValueTypeB; // vector's valuetype\n typedef TemplateConfig::vec_prec, MatrixType::TConfig::matPrec, MatrixType::TConfig::indPrec> hvector_type; // TConfig host with pod-values for ValueTypeB\n typedef Vector HVector; //vectors for saving norms from allgather\n\n static void get_1x1_norm(const MatrixType &A, const Vector_h &r, const int block_size, const NormType norm_type, PODHostVec &block_nrm, typename types::PODTypes::type norm_factor)\n {\n //collect values from all neighbors, and do the \"reduction\" part\n std::vector values(0);\n block_nrm.resize(1);\n int offset, size;\n double sum = 0.l;\n A.getOffsetAndSizeForView(OWNED, &offset, &size);\n\n if (norm_type == L1 || norm_type == L1_SCALED)\n {\n block_nrm[0] = nrm1(r, offset, size);\n\n if (A.is_matrix_distributed())\n {\n A.getManager()->getComms()->global_reduce(values, block_nrm, A, 4);\n block_nrm[0] = 0;\n\n for (int j = 0; j < values.size(); j++)\n {\n sum += values[j][0];\n }\n\n block_nrm[0] = sum;\n }\n\n if (norm_type == L1_SCALED)\n {\n block_nrm[0] /= norm_factor;\n }\n }\n else if (norm_type == L2)\n {\n block_nrm[0] = nrm2(r, offset, size);\n\n if (A.is_matrix_distributed())\n {\n block_nrm[0] *= block_nrm[0];\n A.getManager()->getComms()->global_reduce(values, block_nrm, A, 5);\n block_nrm[0] = 0;\n\n for (int j = 0; j < values.size(); j++)\n {\n sum += values[j][0];\n }\n\n block_nrm[0] = sqrt(sum);\n }\n }\n else if (norm_type == LMAX)\n {\n block_nrm[0] = nrmmax(r, offset, size);\n\n if (A.is_matrix_distributed())\n {\n A.getManager()->getComms()->global_reduce(values, block_nrm, A, 6);\n\n for (int j = 0; j < values.size(); j++)\n {\n block_nrm[0] = (block_nrm[0] > values[j][0] ? block_nrm[0] : values[j][0]);\n }\n }\n }\n else\n {\n FatalError(\"Normtype is not supported in get_1x1_norm\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n };\n};\n\ntemplate \nclass Norm_Square;\n\ntemplate \nvoid get_sq_norm(const MatrixType &A, const VectorType &r, const int block_size, const NormType norm_type, PlainVectorType &block_nrm, typename types::PODTypes::type norm_factor)\n{\n Norm_Square::get_sq_norm(A, r, block_size, norm_type, block_nrm, norm_factor);\n}\n\ntemplate \nclass Norm_Square >, MatrixType, Vector>>\n{\n public:\n typedef Vector > Vector_h; // host vector of current TConfig(r)\n typedef typename Vector_h::value_type ValueTypeB; // vector's valuetype\n typedef TemplateConfig::vec_prec, MatrixType::TConfig::matPrec, MatrixType::TConfig::indPrec> hvector_type; // TConfig host with pod-values for ValueTypeB\n typedef Vector HVector; //vectors for saving norms from allgather\n\n static void get_sq_norm(const MatrixType &A, const Vector_h &r, const int block_size, const NormType norm_type, HVector &block_nrm, typename types::PODTypes::type norm_factor)\n {\n int bsize = block_nrm.size();\n int offset, size;\n A.getOffsetAndSizeForView(OWNED, &offset, &size);\n std::vector norm(block_size, 0.l);\n\n if (norm_type == L1 || norm_type == L1_SCALED)\n {\n if ( (size * r.get_block_size()) % bsize != 0)\n {\n FatalError(\"r.size should be multiple of block size\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n block_nrm.resize(bsize, 0.);\n // for (int j=0;j::abs(r[(offset + i) * bsize + j]);\n }\n }\n\n for (int j = 0; j < bsize; j++)\n {\n block_nrm[j] = norm[j];\n }\n\n if (A.is_matrix_distributed())\n {\n //collect values from all neighbors, and do the \"reduction\" part\n std::vector values(0);\n A.getManager()->getComms()->global_reduce(values, block_nrm, A, 7);\n\n for (int i = 0; i < bsize; i++)\n {\n block_nrm[i] = 0;\n\n for (int j = 0; j < values.size(); j++) { norm[i] += values[j][i]; }\n }\n }\n\n for (int j = 0; j < bsize; j++)\n {\n block_nrm[j] = (norm_type == L1_SCALED) ? norm[j] / norm_factor : norm[j];\n }\n }\n else if (norm_type == L2)\n {\n if ( (size * r.get_block_size()) % bsize != 0)\n {\n FatalError(\"r.size should be multiple of block size\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n block_nrm.resize(bsize, 0.);\n int num_cells = (size * r.get_block_size()) / bsize;\n\n for (int i = 0; i < num_cells; i++)\n {\n for (int j = 0; j < bsize; j++)\n {\n norm[j] = norm[j] + types::util::abs(r[(offset + i) * bsize + j] * types::util::conjugate(r[(offset + i) * bsize + j]));\n }\n }\n\n for (int j = 0; j < bsize; j++)\n {\n block_nrm[j] = norm[j];\n }\n\n if (A.is_matrix_distributed())\n {\n //collect values from all neighbors, and do the \"reduction\" part\n std::vector values(0);\n A.getManager()->getComms()->global_reduce(values, block_nrm, A, 8);\n\n for (int i = 0; i < bsize; i++)\n {\n norm[i] = 0;\n\n for (int j = 0; j < values.size(); j++) { norm[i] += values[j][i]; }\n }\n }\n\n for (int j = 0; j < bsize; j++)\n {\n block_nrm[j] = sqrt(norm[j]);\n }\n }\n else\n {\n FatalError(\"Normtype not supported in get_norm\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n};\n\n__global__\nvoid createReductionMapKernel(int *map, const int bsize, const int map_size)\n{\n for (int tid = threadIdx.x + blockIdx.x * blockDim.x; tid < map_size; tid += blockDim.x * gridDim.x)\n {\n map[tid] = tid % bsize;\n }\n}\n\ntemplate \nclass Norm_Square >, MatrixType, Vector> >\n{\n public:\n typedef Vector > Vector_d;\n typedef Vector > Vector_h; // host vector of current TConfig(r)\n typedef typename Vector_h::value_type ValueTypeB; // vector's valuetype\n typedef TemplateConfig::vec_prec, MatrixType::TConfig::matPrec, MatrixType::TConfig::indPrec> hvector_type; // TConfig host with pod-values for ValueTypeB\n typedef Vector HVector; //vectors for saving norms from allgather\n static void get_sq_norm(const MatrixType &A, const Vector_d &r, const int block_size, const NormType norm_type, HVector &block_nrm, typename types::PODTypes::type norm_factor)\n {\n int bsize = block_nrm.size();\n int offset, size;\n A.getOffsetAndSizeForView(OWNED, &offset, &size);\n const int ncells = (size * r.get_block_size()) / bsize;\n\n if ( (size * r.get_block_size()) % bsize != 0)\n {\n FatalError(\"Size of vector r must be multiple of block size\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (norm_type == L1 || norm_type == L1_SCALED)\n {\n amgx::strided_reduction::reduction_generic_dispatch(\n bsize,\n &block_nrm[0], r.raw() + offset * r.get_block_size(),\n size * r.get_block_size(), amgx::strided_reduction::fabs_transform());\n\n if (A.is_matrix_distributed())\n {\n //collect values from all neighbors, and do the \"reduction\" part\n std::vector values(0);\n A.getManager()->getComms()->global_reduce(values, block_nrm, A, 9);\n\n for (int i = 0; i < bsize; i++)\n {\n block_nrm[i] = 0;\n\n for (int j = 0; j < values.size(); j++)\n {\n block_nrm[i] += values[j][i];\n }\n }\n }\n\n if (norm_type == L1_SCALED)\n {\n for (int i = 0; i < block_nrm.size(); ++i)\n {\n block_nrm[i] /= norm_factor;\n }\n }\n }\n else if (norm_type == L2)\n {\n amgx::strided_reduction::reduction_generic_dispatch(\n bsize, &block_nrm[0],\n r.raw() + offset * r.get_block_size(), size * r.get_block_size(),\n amgx::strided_reduction::square_transform());\n\n for (int i = 0; i < bsize; i++)\n {\n block_nrm[i] = sqrt(block_nrm[i]);\n }\n\n if (A.is_matrix_distributed())\n {\n for (int i = 0; i < bsize; i++)\n {\n block_nrm[i] *= block_nrm[i];\n }\n\n //collect values from all neighbors, and do the \"reduction\" part\n std::vector values(0);\n A.getManager()->getComms()->global_reduce(values, block_nrm, A, 10);\n\n for (int i = 0; i < bsize; i++)\n {\n block_nrm[i] = 0;\n\n for (int j = 0; j < values.size(); j++)\n {\n block_nrm[i] += values[j][i];\n }\n }\n\n for (int i = 0; i < bsize; i++)\n {\n block_nrm[i] = sqrt(block_nrm[i]);\n }\n }\n }\n else\n {\n FatalError(\"Normtype not supported in get_norm\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n};\n\ntemplate\nvoid get_norm(const MatrixType &A, const VectorType &r, const int block_size, const NormType norm_type, PlainVectorType &block_nrm, typename types::PODTypes::type norm_factor)\n{\n if (block_size == 1)\n {\n get_1x1_norm(A, r, block_size, norm_type, block_nrm, norm_factor);\n }\n else\n {\n get_sq_norm(A, r, block_size, norm_type, block_nrm, norm_factor);\n }\n}\n\ntemplate \nclass Norm_Factor;\n\ntemplate \nvoid compute_norm_factor(MatrixType &A, VectorType &b, VectorType &x, const NormType normType, typename types::PODTypes::type &normFactor)\n{\n if(normType == L1_SCALED)\n {\n Norm_Factor::compute_norm_factor(A, b, x, normFactor);\n }\n}\n\ntemplate \nclass Norm_Factor >, MatrixType>\n{\n public:\n typedef Vector > Vector_h;\n typedef typename Vector_h::value_type ValueTypeVec;\n typedef typename types::PODTypes::type ValueTypeNorm;\n\n static void compute_norm_factor(MatrixType &A, Vector_h &b, Vector_h &x, ValueTypeNorm& normFactor)\n {\n FatalError(\"L1 scaled norm not supported with host execution.\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n};\n\ntemplate\n__global__ void scaled_norm_factor_calc(\n int nRows, ValueTypeMat *Avals, IndexTypeVec *Arows, ValueTypeVec *Ax, ValueTypeVec *b, \n ValueTypeVec xAvg, ValueTypeNorm *localNormFactor, ValueTypeMat *diaVals = nullptr)\n{\n int r = threadIdx.x + blockIdx.x * blockDim.x;\n\n __shared__ ValueTypeNorm normFactor_s;\n if(threadIdx.x == 0)\n {\n normFactor_s = amgx::types::util::get_zero();\n }\n\n ValueTypeNorm normFactor = 0.0;\n\n if (r < nRows)\n {\n ValueTypeMat Arow_sum = amgx::types::util::get_zero();\n\n // Read in the row\n#pragma unroll\n for (int i = Arows[r]; i < Arows[r + 1]; ++i)\n {\n Arow_sum = Arow_sum + Avals[i];\n }\n if(diaVals) Arow_sum = Arow_sum + diaVals[r];\n\n normFactor =\n types::util::abs(Ax[r] - Arow_sum * xAvg) +\n types::util::abs(b[r] - Arow_sum * xAvg);\n }\n\n // Ensure normFactor_s is initialised\n __syncthreads();\n\n // Warp-local reduction to lane 0\n for(int i = warpSize/2; i > 0; i /= 2)\n {\n normFactor += utils::shfl_down(normFactor, i);\n }\n\n // Fast shared atomic add by lane 0 of each warp\n int laneId = threadIdx.x % warpSize;\n if(laneId == 0)\n {\n utils::atomic_add(&normFactor_s, normFactor);\n }\n\n // Ensure normFactor_s is final\n __syncthreads();\n\n // Final output of normFactor by first thread of each block\n if(threadIdx.x == 0)\n {\n utils::atomic_add(localNormFactor, normFactor_s);\n }\n}\n\ntemplate \nclass Norm_Factor >, MatrixType>\n{\n public:\n typedef Vector > Vector_d;\n typedef typename Vector_d::value_type ValueTypeVec;\n typedef typename Vector_d::index_type IndexTypeVec;\n typedef typename types::PODTypes::type ValueTypeNorm;\n typedef TemplateConfig::vec_prec, MatrixType::TConfig::matPrec, MatrixType::TConfig::indPrec> NormVectorType;\n typedef Vector NVector_d;\n\n static void compute_norm_factor(MatrixType &A, Vector_d &b, Vector_d &x, ValueTypeNorm &normFactor)\n {\n if (A.get_block_dimx() != 1 || A.get_block_dimy() != 1)\n {\n FatalError(\"L1 scaled norm only supported with scalar matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // Calculate Ax\n int offset, nRows;\n A.getOffsetAndSizeForView(OWNED, &offset, &nRows);\n\n Vector_d Ax(nRows);\n A.apply(x, Ax);\n\n // Calculate global average x\n ValueTypeVec xAvg = amgx::thrust::reduce(x.begin(), x.begin() + nRows, amgx::types::util::get_zero());\n\n int64_t nr = nRows;\n if (A.is_matrix_distributed())\n {\n A.manager->global_reduce_sum(&xAvg);\n nr = A.manager->num_rows_global;\n }\n amgx::types::util::divide_by_integer(xAvg, nr);\n\n // Make a copy of b\n Vector_d bTmp(b);\n\n // Calculate row sums then the local norm factors\n constexpr int nThreads = 128;\n constexpr int warpSize = 32;\n const int nBlocks = nRows/nThreads + 1;\n NVector_d localNormFactor(1, amgx::types::util::get_zero());\n if(A.hasProps(DIAG))\n {\n scaled_norm_factor_calc<<>>(\n nRows,\n A.values.raw(),\n A.row_offsets.raw(),\n Ax.raw(),\n bTmp.raw(),\n xAvg,\n localNormFactor.raw(),\n A.values.raw() + A.diagOffset()*A.get_block_size());\n cudaCheckError();\n }\n else\n {\n scaled_norm_factor_calc<<>>(\n nRows,\n A.values.raw(),\n A.row_offsets.raw(),\n Ax.raw(),\n bTmp.raw(),\n xAvg,\n localNormFactor.raw());\n cudaCheckError();\n }\n\n // Fetch the normFactor result and reduce across all ranks\n normFactor = localNormFactor[0];\n\n if (A.is_matrix_distributed())\n {\n A.manager->global_reduce_sum(&normFactor);\n }\n\n // Print the norm factor\n std::stringstream info;\n info.precision(12);\n info << \"\\tAmgX Scaled Norm Factor: \" << std::scientific << normFactor << \"\\n\";\n amgx_output(info.str().c_str(), info.str().length());\n }\n};\n\n#define AMGX_CASE_LINE(CASE) template typename types::PODTypes< typename Vector::Type>::value_type>::type get_norm(const Matrix::Type>& A, const Vector::Type>& r, const NormType norm_type, typename types::PODTypes::Type>::value_type>::type norm_factor);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template typename types::PODTypes< typename Vector::Type>::value_type>::type get_norm(const Operator::Type>& A, const Vector::Type>& r, const NormType norm_type, typename types::PODTypes::Type>::value_type>::type norm_factor);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n typedef typename Vector< TemplateMode::Type >::value_type ValueTypeMB##CASE ;\\\n typedef TemplateMode::Type::template setMemSpace::Type::template setVecPrec< types::PODTypes< ValueTypeMB##CASE >::vec_prec >::Type CurTConfigMB_h##CASE ;\\\n template void get_norm(const Matrix::Type>& A, const Vector::Type>& r, const int block_size, const NormType norm_type, Vector< CurTConfigMB_h##CASE >& block_nrm, typename types::PODTypes::Type>::value_type>::type norm_factor); \\\n template void compute_norm_factor(Matrix::Type> &A, Vector::Type> &b, Vector::Type> &x, const NormType normType, typename types::PODTypes::Type>::value_type>::type &normFactor);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n typedef typename Vector< TemplateMode::Type >::value_type ValueTypeOB##CASE ;\\\n typedef TemplateMode::Type::template setMemSpace::Type::template setVecPrec< types::PODTypes< ValueTypeOB##CASE >::vec_prec >::Type CurTConfigOB_h##CASE ;\\\n template void get_norm(const Operator::Type>& A, const Vector::Type>& r, const int block_size, const NormType norm_type, Vector< CurTConfigOB_h##CASE >& block_nrm, typename types::PODTypes::Type>::value_type>::type norm_factor);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#ifdef _WIN32\n#pragma warning (push)\n#pragma warning (disable : 4244 4267 4521)\n#endif\n#ifdef _WIN32\n#pragma warning (pop)\n#endif\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"strided_reduction.h\"\n\n#include \"amgx_timer.h\"\n#include \"amgx_types/util.h\"\n#include \"thrust_wrapper.h\"\n\nnamespace amgx\n{\n\n/**********************************************************\n * Returns the norm of a vector\n *********************************************************/\n\ntemplate\ntypename types::PODTypes::type get_norm(const MatrixType &A, const VectorType &r, const NormType norm_type, typename types::PODTypes::type norm_factor)\n{\n typedef typename types::PODTypes::type value_type;\n value_type nrm;\n int offset, size;\n A.getOffsetAndSizeForView(OWNED, &offset, &size);\n\n if (norm_type == L1 || norm_type == L1_SCALED)\n {\n nrm = nrm1(r, offset, size);\n\n if (A.is_matrix_distributed())\n {\n A.getManager()->global_reduce_sum(&nrm);\n }\n\n return (norm_type == L1_SCALED) ? nrm / norm_factor : nrm;\n }\n else if (norm_type == L2)\n {\n nrm = nrm2(r, offset, size);\n\n if (A.is_matrix_distributed())\n {\n nrm = nrm * nrm;\n A.getManager()->global_reduce_sum(&nrm);\n nrm = sqrt(nrm);\n }\n \n return nrm;\n }\n else if (norm_type == LMAX)\n {\n nrm = nrmmax(r, offset, size);\n\n if (A.is_matrix_distributed())\n {\n typedef TemplateConfig::vec_prec, MatrixType::TConfig::matPrec, MatrixType::TConfig::indPrec> hvector_type;\n typedef Vector HVector;\n //collect values from all neighbors, and do the \"reduction\" part\n std::vector values(0);\n HVector my_nrm(1);\n my_nrm[0] = nrm;\n A.getManager()->getComms()->global_reduce(values, my_nrm, A, 3);\n\n for (int j = 0; j < values.size(); j++)\n {\n nrm = (nrm > values[j][0] ? nrm : values[j][0]);\n }\n }\n\n return nrm;\n }\n\n return -1;\n}\n\ntemplate \nclass Norm_1x1;\n\ntemplate \nvoid get_1x1_norm(const MatrixType &A, const VectorType &r, const int block_size, const NormType norm_type, PlainVectorType &block_nrm, typename types::PODTypes::type norm_factor)\n{\n Norm_1x1::get_1x1_norm(A, r, block_size, norm_type, block_nrm, norm_factor);\n}\n\ntemplate \nclass Norm_1x1< Vector >, MatrixType, Vector>>\n{\n public:\n typedef Vector > Vector_h;\n typedef Vector> PODHostVec; // return type\n typedef typename Vector_h::value_type ValueTypeB; // vector's valuetype\n typedef TemplateConfig::vec_prec, MatrixType::TConfig::matPrec, MatrixType::TConfig::indPrec> hvector_type; // TConfig host with pod-values for ValueTypeB\n typedef Vector HVector; //vectors for saving norms from allgather\n\n static void get_1x1_norm(const MatrixType &A, const Vector_h &r, const int block_size, const NormType norm_type, PODHostVec &block_nrm, typename types::PODTypes::type norm_factor)\n {\n //collect values from all neighbors, and do the \"reduction\" part\n std::vector values(0);\n block_nrm.resize(1);\n int offset, size;\n double sum = 0.l;\n A.getOffsetAndSizeForView(OWNED, &offset, &size);\n\n if (norm_type == L1 || norm_type == L1_SCALED)\n {\n block_nrm[0] = nrm1(r, offset, size);\n\n if (A.is_matrix_distributed())\n {\n A.getManager()->getComms()->global_reduce(values, block_nrm, A, 4);\n block_nrm[0] = 0;\n\n for (int j = 0; j < values.size(); j++)\n {\n sum += values[j][0];\n }\n\n block_nrm[0] = sum;\n }\n\n if (norm_type == L1_SCALED)\n {\n block_nrm[0] /= norm_factor;\n }\n }\n else if (norm_type == L2)\n {\n block_nrm[0] = nrm2(r, offset, size);\n\n if (A.is_matrix_distributed())\n {\n block_nrm[0] *= block_nrm[0];\n A.getManager()->getComms()->global_reduce(values, block_nrm, A, 5);\n block_nrm[0] = 0;\n\n for (int j = 0; j < values.size(); j++)\n {\n sum += values[j][0];\n }\n\n block_nrm[0] = sqrt(sum);\n }\n }\n else if (norm_type == LMAX)\n {\n block_nrm[0] = nrmmax(r, offset, size);\n\n if (A.is_matrix_distributed())\n {\n A.getManager()->getComms()->global_reduce(values, block_nrm, A, 6);\n\n for (int j = 0; j < values.size(); j++)\n {\n block_nrm[0] = (block_nrm[0] > values[j][0] ? block_nrm[0] : values[j][0]);\n }\n }\n }\n else\n {\n FatalError(\"Normtype is not supported in get_1x1_norm\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n };\n};\n\ntemplate \nclass Norm_Square;\n\ntemplate \nvoid get_sq_norm(const MatrixType &A, const VectorType &r, const int block_size, const NormType norm_type, PlainVectorType &block_nrm, typename types::PODTypes::type norm_factor)\n{\n Norm_Square::get_sq_norm(A, r, block_size, norm_type, block_nrm, norm_factor);\n}\n\ntemplate \nclass Norm_Square >, MatrixType, Vector>>\n{\n public:\n typedef Vector > Vector_h; // host vector of current TConfig(r)\n typedef typename Vector_h::value_type ValueTypeB; // vector's valuetype\n typedef TemplateConfig::vec_prec, MatrixType::TConfig::matPrec, MatrixType::TConfig::indPrec> hvector_type; // TConfig host with pod-values for ValueTypeB\n typedef Vector HVector; //vectors for saving norms from allgather\n\n static void get_sq_norm(const MatrixType &A, const Vector_h &r, const int block_size, const NormType norm_type, HVector &block_nrm, typename types::PODTypes::type norm_factor)\n {\n int bsize = block_nrm.size();\n int offset, size;\n A.getOffsetAndSizeForView(OWNED, &offset, &size);\n std::vector norm(block_size, 0.l);\n\n if (norm_type == L1 || norm_type == L1_SCALED)\n {\n if ( (size * r.get_block_size()) % bsize != 0)\n {\n FatalError(\"r.size should be multiple of block size\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n block_nrm.resize(bsize, 0.);\n // for (int j=0;j::abs(r[(offset + i) * bsize + j]);\n }\n }\n\n for (int j = 0; j < bsize; j++)\n {\n block_nrm[j] = norm[j];\n }\n\n if (A.is_matrix_distributed())\n {\n //collect values from all neighbors, and do the \"reduction\" part\n std::vector values(0);\n A.getManager()->getComms()->global_reduce(values, block_nrm, A, 7);\n\n for (int i = 0; i < bsize; i++)\n {\n block_nrm[i] = 0;\n\n for (int j = 0; j < values.size(); j++) { norm[i] += values[j][i]; }\n }\n }\n\n for (int j = 0; j < bsize; j++)\n {\n block_nrm[j] = (norm_type == L1_SCALED) ? norm[j] / norm_factor : norm[j];\n }\n }\n else if (norm_type == L2)\n {\n if ( (size * r.get_block_size()) % bsize != 0)\n {\n FatalError(\"r.size should be multiple of block size\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n block_nrm.resize(bsize, 0.);\n int num_cells = (size * r.get_block_size()) / bsize;\n\n for (int i = 0; i < num_cells; i++)\n {\n for (int j = 0; j < bsize; j++)\n {\n norm[j] = norm[j] + types::util::abs(r[(offset + i) * bsize + j] * types::util::conjugate(r[(offset + i) * bsize + j]));\n }\n }\n\n for (int j = 0; j < bsize; j++)\n {\n block_nrm[j] = norm[j];\n }\n\n if (A.is_matrix_distributed())\n {\n //collect values from all neighbors, and do the \"reduction\" part\n std::vector values(0);\n A.getManager()->getComms()->global_reduce(values, block_nrm, A, 8);\n\n for (int i = 0; i < bsize; i++)\n {\n norm[i] = 0;\n\n for (int j = 0; j < values.size(); j++) { norm[i] += values[j][i]; }\n }\n }\n\n for (int j = 0; j < bsize; j++)\n {\n block_nrm[j] = sqrt(norm[j]);\n }\n }\n else\n {\n FatalError(\"Normtype not supported in get_norm\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n};\n\n__global__\nvoid createReductionMapKernel(int *map, const int bsize, const int map_size)\n{\n for (int tid = threadIdx.x + blockIdx.x * blockDim.x; tid < map_size; tid += blockDim.x * gridDim.x)\n {\n map[tid] = tid % bsize;\n }\n}\n\ntemplate \nclass Norm_Square >, MatrixType, Vector> >\n{\n public:\n typedef Vector > Vector_d;\n typedef Vector > Vector_h; // host vector of current TConfig(r)\n typedef typename Vector_h::value_type ValueTypeB; // vector's valuetype\n typedef TemplateConfig::vec_prec, MatrixType::TConfig::matPrec, MatrixType::TConfig::indPrec> hvector_type; // TConfig host with pod-values for ValueTypeB\n typedef Vector HVector; //vectors for saving norms from allgather\n static void get_sq_norm(const MatrixType &A, const Vector_d &r, const int block_size, const NormType norm_type, HVector &block_nrm, typename types::PODTypes::type norm_factor)\n {\n int bsize = block_nrm.size();\n int offset, size;\n A.getOffsetAndSizeForView(OWNED, &offset, &size);\n const int ncells = (size * r.get_block_size()) / bsize;\n\n if ( (size * r.get_block_size()) % bsize != 0)\n {\n FatalError(\"Size of vector r must be multiple of block size\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n if (norm_type == L1 || norm_type == L1_SCALED)\n {\n amgx::strided_reduction::reduction_generic_dispatch(\n bsize,\n &block_nrm[0], r.raw() + offset * r.get_block_size(),\n size * r.get_block_size(), amgx::strided_reduction::fabs_transform());\n\n if (A.is_matrix_distributed())\n {\n //collect values from all neighbors, and do the \"reduction\" part\n std::vector values(0);\n A.getManager()->getComms()->global_reduce(values, block_nrm, A, 9);\n\n for (int i = 0; i < bsize; i++)\n {\n block_nrm[i] = 0;\n\n for (int j = 0; j < values.size(); j++)\n {\n block_nrm[i] += values[j][i];\n }\n }\n }\n\n if (norm_type == L1_SCALED)\n {\n for (int i = 0; i < block_nrm.size(); ++i)\n {\n block_nrm[i] /= norm_factor;\n }\n }\n }\n else if (norm_type == L2)\n {\n amgx::strided_reduction::reduction_generic_dispatch(\n bsize, &block_nrm[0],\n r.raw() + offset * r.get_block_size(), size * r.get_block_size(),\n amgx::strided_reduction::square_transform());\n\n for (int i = 0; i < bsize; i++)\n {\n block_nrm[i] = sqrt(block_nrm[i]);\n }\n\n if (A.is_matrix_distributed())\n {\n for (int i = 0; i < bsize; i++)\n {\n block_nrm[i] *= block_nrm[i];\n }\n\n //collect values from all neighbors, and do the \"reduction\" part\n std::vector values(0);\n A.getManager()->getComms()->global_reduce(values, block_nrm, A, 10);\n\n for (int i = 0; i < bsize; i++)\n {\n block_nrm[i] = 0;\n\n for (int j = 0; j < values.size(); j++)\n {\n block_nrm[i] += values[j][i];\n }\n }\n\n for (int i = 0; i < bsize; i++)\n {\n block_nrm[i] = sqrt(block_nrm[i]);\n }\n }\n }\n else\n {\n FatalError(\"Normtype not supported in get_norm\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n};\n\ntemplate\nvoid get_norm(const MatrixType &A, const VectorType &r, const int block_size, const NormType norm_type, PlainVectorType &block_nrm, typename types::PODTypes::type norm_factor)\n{\n if (block_size == 1)\n {\n get_1x1_norm(A, r, block_size, norm_type, block_nrm, norm_factor);\n }\n else\n {\n get_sq_norm(A, r, block_size, norm_type, block_nrm, norm_factor);\n }\n}\n\ntemplate \nclass Norm_Factor;\n\ntemplate \nvoid compute_norm_factor(MatrixType &A, VectorType &b, VectorType &x, const NormType normType, typename types::PODTypes::type &normFactor)\n{\n if(normType == L1_SCALED)\n {\n Norm_Factor::compute_norm_factor(A, b, x, normFactor);\n }\n}\n\ntemplate \nclass Norm_Factor >, MatrixType>\n{\n public:\n typedef Vector > Vector_h;\n typedef typename Vector_h::value_type ValueTypeVec;\n typedef typename types::PODTypes::type ValueTypeNorm;\n\n static void compute_norm_factor(MatrixType &A, Vector_h &b, Vector_h &x, ValueTypeNorm& normFactor)\n {\n FatalError(\"L1 scaled norm not supported with host execution.\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n};\n\ntemplate\n__global__ void scaled_norm_factor_calc(\n int nRows, ValueTypeMat *Avals, IndexTypeVec *Arows, ValueTypeVec *Ax, ValueTypeVec *b, \n ValueTypeVec xAvg, ValueTypeNorm *localNormFactor, ValueTypeMat *diaVals = nullptr)\n{\n int r = threadIdx.x + blockIdx.x * blockDim.x;\n\n __shared__ ValueTypeNorm normFactor_s;\n if(threadIdx.x == 0)\n {\n normFactor_s = amgx::types::util::get_zero();\n }\n\n ValueTypeNorm normFactor = 0.0;\n\n if (r < nRows)\n {\n ValueTypeMat Arow_sum = amgx::types::util::get_zero();\n\n // Read in the row\n#pragma unroll\n for (int i = Arows[r]; i < Arows[r + 1]; ++i)\n {\n Arow_sum = Arow_sum + Avals[i];\n }\n if(diaVals) Arow_sum = Arow_sum + diaVals[r];\n\n normFactor =\n types::util::abs(Ax[r] - Arow_sum * xAvg) +\n types::util::abs(b[r] - Arow_sum * xAvg);\n }\n\n // Ensure normFactor_s is initialised\n __syncthreads();\n\n // Warp-local reduction to lane 0\n for(int i = warpSize/2; i > 0; i /= 2)\n {\n normFactor += utils::shfl_down(normFactor, i);\n }\n\n // Fast shared atomic add by lane 0 of each warp\n int laneId = threadIdx.x % warpSize;\n if(laneId == 0)\n {\n utils::atomic_add(&normFactor_s, normFactor);\n }\n\n // Ensure normFactor_s is final\n __syncthreads();\n\n // Final output of normFactor by first thread of each block\n if(threadIdx.x == 0)\n {\n utils::atomic_add(localNormFactor, normFactor_s);\n }\n}\n\ntemplate \nclass Norm_Factor >, MatrixType>\n{\n public:\n typedef Vector > Vector_d;\n typedef typename Vector_d::value_type ValueTypeVec;\n typedef typename Vector_d::index_type IndexTypeVec;\n typedef typename types::PODTypes::type ValueTypeNorm;\n typedef TemplateConfig::vec_prec, MatrixType::TConfig::matPrec, MatrixType::TConfig::indPrec> NormVectorType;\n typedef Vector NVector_d;\n\n static void compute_norm_factor(MatrixType &A, Vector_d &b, Vector_d &x, ValueTypeNorm &normFactor)\n {\n if (A.get_block_dimx() != 1 || A.get_block_dimy() != 1)\n {\n FatalError(\"L1 scaled norm only supported with scalar matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // Calculate Ax\n int offset, nRows;\n A.getOffsetAndSizeForView(OWNED, &offset, &nRows);\n\n Vector_d Ax(nRows);\n A.apply(x, Ax);\n\n // Calculate global average x\n ValueTypeVec xAvg = amgx::thrust::reduce(x.begin(), x.begin() + nRows, amgx::types::util::get_zero());\n\n int64_t nr = nRows;\n if (A.is_matrix_distributed())\n {\n A.manager->global_reduce_sum(&xAvg);\n nr = A.manager->num_rows_global;\n }\n amgx::types::util::divide_by_integer(xAvg, nr);\n\n // Make a copy of b\n Vector_d bTmp(b);\n\n // Calculate row sums then the local norm factors\n constexpr int nThreads = 128;\n constexpr int warpSize = 32;\n const int nBlocks = nRows/nThreads + 1;\n NVector_d localNormFactor(1, amgx::types::util::get_zero());\n if(A.hasProps(DIAG))\n {\n scaled_norm_factor_calc<<>>(\n nRows,\n A.values.raw(),\n A.row_offsets.raw(),\n Ax.raw(),\n bTmp.raw(),\n xAvg,\n localNormFactor.raw(),\n A.values.raw() + A.diagOffset()*A.get_block_size());\n cudaCheckError();\n }\n else\n {\n scaled_norm_factor_calc<<>>(\n nRows,\n A.values.raw(),\n A.row_offsets.raw(),\n Ax.raw(),\n bTmp.raw(),\n xAvg,\n localNormFactor.raw());\n cudaCheckError();\n }\n\n // Fetch the normFactor result and reduce across all ranks\n normFactor = localNormFactor[0];\n\n if (A.is_matrix_distributed())\n {\n A.manager->global_reduce_sum(&normFactor);\n }\n\n // Print the norm factor\n std::stringstream info;\n info.precision(12);\n info << \"\\tAmgX Scaled Norm Factor: \" << std::scientific << normFactor << \"\\n\";\n amgx_output(info.str().c_str(), info.str().length());\n }\n};\n\n#define AMGX_CASE_LINE(CASE) template typename types::PODTypes< typename Vector::Type>::value_type>::type get_norm(const Matrix::Type>& A, const Vector::Type>& r, const NormType norm_type, typename types::PODTypes::Type>::value_type>::type norm_factor);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template typename types::PODTypes< typename Vector::Type>::value_type>::type get_norm(const Operator::Type>& A, const Vector::Type>& r, const NormType norm_type, typename types::PODTypes::Type>::value_type>::type norm_factor);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n typedef typename Vector< TemplateMode::Type >::value_type ValueTypeMB##CASE ;\\\n typedef TemplateMode::Type::template setMemSpace::Type::template setVecPrec< types::PODTypes< ValueTypeMB##CASE >::vec_prec >::Type CurTConfigMB_h##CASE ;\\\n template void get_norm(const Matrix::Type>& A, const Vector::Type>& r, const int block_size, const NormType norm_type, Vector< CurTConfigMB_h##CASE >& block_nrm, typename types::PODTypes::Type>::value_type>::type norm_factor); \\\n template void compute_norm_factor(Matrix::Type> &A, Vector::Type> &b, Vector::Type> &x, const NormType normType, typename types::PODTypes::Type>::value_type>::type &normFactor);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) \\\n typedef typename Vector< TemplateMode::Type >::value_type ValueTypeOB##CASE ;\\\n typedef TemplateMode::Type::template setMemSpace::Type::template setVecPrec< types::PODTypes< ValueTypeOB##CASE >::vec_prec >::Type CurTConfigOB_h##CASE ;\\\n template void get_norm(const Operator::Type>& A, const Vector::Type>& r, const int block_size, const NormType norm_type, Vector< CurTConfigOB_h##CASE >& block_nrm, typename types::PODTypes::Type>::value_type>::type norm_factor);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/norm.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include // std::signbit\n#include \n\n#include \"util.h\"\n\n#include \"amgx_cublas.h\"\n\nnamespace amgx\n{\nnamespace energymin\n{\n\n\n/*************************************************************************\n* create the interpolation matrix P by computing p and using energy minimization\n************************************************************************/\n\n\n/*************************************************************************\n* create the interpolation matrix P on GPU\n************************************************************************/\n\ntemplate \n__global__\nvoid count_PnnzPerCol_kernel( const IndexType *ArowOffsets, const IndexType *AcolInd, const ValueType *Avalues,\n const ValueType *Adiag, const IndexType numCoarse,\n const int *cf_map, const int *fullToCoarseMap, const int *coarse_idx_ptr,\n int *PnnzPerCol, ValueType *thresh_per_row )\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numCoarse; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on one \"coarse\" row of A at a time\n const int row_idx = coarse_idx_ptr[tIdx];\n const int ArowBegin = ArowOffsets[row_idx];\n const int ArowEnd = ArowOffsets[row_idx + 1];\n // Matlab pseudo-code:\n // p = sparse(n, nc);\n // p(coarse,:) = speye(nc);\n // fAj = AD(j,fine)/AD(j,j) < -.1; // fine neighbors of j with oppositive sign from AD_jj\n // become the non-zeros in initial P.\n ValueType minVal(0), maxVal(0), diag_value(0);\n ValueType fine_avg_thresh = 0.;\n const ValueType thresh = (ValueType)0.25;\n\n for (int j = ArowBegin; j < ArowEnd; j++)\n {\n ValueType aValue = Avalues[j];\n\n if ( j != tIdx )\n {\n minVal = min( minVal, aValue );\n maxVal = max( maxVal, aValue );\n }\n else\n {\n diag_value = aValue;\n }\n }\n\n if ( diag_value < 0.0)\n {\n fine_avg_thresh = thresh * maxVal;\n }\n else\n {\n fine_avg_thresh = thresh * minVal;\n }\n\n thresh_per_row[tIdx] = fine_avg_thresh;\n int p_col_nnz = 0;\n\n for (int j = ArowBegin; j < ArowEnd; j++)\n {\n int AjColIndex = AcolInd[j];\n\n if (AjColIndex == tIdx)\n {\n p_col_nnz++;\n continue;\n }\n\n if (cf_map[AjColIndex] == FINE)\n {\n if (diag_value * Avalues[j] < diag_value * fine_avg_thresh)\n {\n p_col_nnz++;\n }\n }\n }\n\n PnnzPerCol[tIdx] = p_col_nnz;\n }\n}\n\n\ntemplate \n__global__\nvoid init_ProwInd_greedy_aggregation_kernel(const IndexType *ArowOffsets, const IndexType *AcolInd, const ValueType *Avalues,\n const ValueType *Adiag, const IndexType AnumRows,\n const int *cf_map, const int *fullToCoarseMap,\n const IndexType *PcolOffsets, IndexType *ProwInd,\n bool *Ma_nzDiagRows)\n{\n // for each coarse point form aggregates\n /* Matlab-like pseudo-code\n p = sparse(n, nc);\n p(coarse,:) = speye(nc);\n for (int j : coarse) // coarse point with FULL index j\n fAj = AD(j,fine)/AD(j,j) < -.1; // fine neighbors of j with oppositive sign from AD_jj\n P(fine(fAj),cj)=1; // become the non-zeros in initial P.\n end\n */\n\n // size of P and its column offsets were precomputed.\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AnumRows; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on one \"coarse\" row of A at a time (builds a column of P)\n const int ArowBegin = ArowOffsets[tIdx];\n const int ArowEnd = ArowOffsets[tIdx + 1];\n\n if (cf_map[tIdx] == COARSE) // \"coarse\" row of A\n {\n // P is CSC format since each thread works on each coarse point\n // and is therefore building a single column of P.\n // There are numCoarse columns of P\n // Columns of P correspond to rows of A (A is assumed symmetric)\n bool PcolCoarseDiagUnset = 1;\n const int PcolIndex = fullToCoarseMap[tIdx];\n int PentryNum = PcolOffsets[PcolIndex]; // starting location of the current column of P\n\n for (int j = ArowBegin; j < ArowEnd; j++)\n {\n // go thru each \"fine\" column of A\n int AjColIndex = AcolInd[j];\n\n if (cf_map[AjColIndex] == FINE && Avalues[j] / Adiag[tIdx] < -.1)\n {\n //p(AjColIndex, fullToCoarseMap(j))=1;\n if (AjColIndex > tIdx && PcolCoarseDiagUnset) // COARSE diagonal of P is identity\n {\n // set COARSE diagonal of P\n ProwInd[PentryNum] = tIdx; // set COARSE diagonal entry of P to 1.\n PentryNum++;\n PcolCoarseDiagUnset = 0;\n Ma_nzDiagRows[tIdx] = 1; // record the newly added row index of P\n }\n\n ProwInd[PentryNum] = AjColIndex; // columns of A become rows in P (since A is assumed symmetric)\n PentryNum++;\n Ma_nzDiagRows[AjColIndex] = 1; // record the newly added row index of P\n }\n }\n\n if (PcolCoarseDiagUnset) // COARSE diagonal of P is identity\n {\n // set COARSE diagonal of P\n ProwInd[PentryNum] = tIdx; // columns of A become rows in P (since A is assumed symmetric)\n PentryNum++;\n PcolCoarseDiagUnset = 0;\n Ma_nzDiagRows[tIdx] = 1; // record the newly added row index of P\n }\n }\n }\n}\n\n\ntemplate \n__global__\nvoid extract_dense_Aijs_col_major_kernel(const IndexType *ArowOffsets, const IndexType *AcolInd, const ValueType *Avalues,\n const IndexType *PcolOffsets, const IndexType *ProwInd, const IndexType PnumCols,\n const IndexType *AijOffsets, ValueType *dense_Aijs)\n{\n /* Matlab-like pseudo-code\n for (int j=1; j<=nc; j++)\n {\n // Find the nz pattern in the j-th column of P\n ij = find(p(:,j)); // col vec of subindices\n Aij{j} = A(ij,ij);\n }\n */\n\n // each thread builds one Aij\n for (unsigned int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < PnumCols; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works with one column of P\n const unsigned int PcolBegin = PcolOffsets[tIdx];\n const unsigned int PcolEnd = PcolOffsets[tIdx + 1];\n const unsigned int AijNumRows = PcolEnd - PcolBegin;\n // pointer to the beginning of the specific Aij submatrix\n ValueType *dense_Aij = dense_Aijs + AijOffsets[tIdx];\n\n for (unsigned int i = PcolBegin; i < PcolEnd; i++)\n {\n // go thru each non-zero row of the given column of P\n // and scan the corresponding rows of A\n unsigned int PiRowIndex = ProwInd[i];\n const unsigned int ArowBegin = ArowOffsets[PiRowIndex];\n const unsigned int ArowEnd = ArowOffsets[PiRowIndex + 1];\n\n for (unsigned int j = ArowBegin; j < ArowEnd; j++)\n {\n // scan the current row of A and pick the elts based on the sparsity pattern of P\n unsigned int AjColIndex = AcolInd[j];\n\n // if AjColIndex is an index of a non-zero entry in the current column of P, take it.\n for (unsigned int ii = PcolBegin; ii < PcolEnd; ii++)\n {\n if (ProwInd[ii] == AjColIndex)\n {\n // add elt to Aij in column-major format\n unsigned int AijRowIndex = i - PcolBegin;\n unsigned int AijColIndex = ii - PcolBegin;\n dense_Aij[AijColIndex * AijNumRows + AijRowIndex] = Avalues[j]; //wrong\n break;\n }\n }\n }\n }\n }\n}\n\n\n// kernel to initialize invAij's to dense identity matrices\ntemplate \n__global__\nvoid init_dense_invAijs_col_major_kernel(const IndexType *PcolOffsets, const IndexType PnumCols,\n const IndexType *AijOffsets, ValueType *dense_invAijs)\n{\n for (unsigned int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < PnumCols; tIdx += blockDim.x * gridDim.x)\n {\n // each thread initializes one invAij\n const unsigned int PcolBegin = PcolOffsets[tIdx];\n const unsigned int PcolEnd = PcolOffsets[tIdx + 1];\n const unsigned int AijNumRows = PcolEnd - PcolBegin;\n // pointer to the beginning of the specific invAij submatrix\n ValueType *dense_invAij = dense_invAijs + AijOffsets[tIdx];\n\n for (unsigned int i = 0; i < AijNumRows; i++)\n {\n dense_invAij[i * AijNumRows + i] = 1;\n }\n }\n}\n\n\n\n\n__global__\nvoid count_Ma_nnz_per_row_kernel(const int *PcolOffsets, const int *ProwInd, const int PnumCols,\n int *Ma_nnzPerRow, const int AnumRows, const bool *Ma_nzDiagRows)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AnumRows; tIdx += blockDim.x * gridDim.x)\n {\n // every thread works on a single row of Ma (figures out nnz in that row)\n bool rowNotFilled = 1; // marks whether the Ma row nnz count has been initialized\n int nnzPerRow = 0;\n\n for (int j = 0; j < PnumCols; j++)\n {\n // each thread goes thru all columns of P\n const int PjColBegin = PcolOffsets[j];\n const int PjColEnd = PcolOffsets[j + 1];\n //const int PjCol_nnz = PjColEnd - PjColBegin;\n\n // check that the column is relevant (i.e. has a non-zero in the tIdx row)\n if (tIdx < ProwInd[PjColBegin] || ProwInd[PjColEnd - 1] < tIdx) { continue; }\n\n // Make sure the column is actually relevant (i.e. has a non-zero in the tIdx row)\n bool relevantCol = 0;\n\n for (int i = PjColBegin; i < PjColEnd; i++)\n {\n if (ProwInd[i] == tIdx) { relevantCol = 1; break; }\n }\n\n if (!relevantCol) { continue; }\n\n if (rowNotFilled)\n {\n nnzPerRow = PjColEnd - PjColBegin; // initial nnz per row of Ma\n rowNotFilled = 0;\n continue;\n }\n\n // if the row nnz count has been initialized, we need to figure out which elts haven't already appeared\n for (int i = PjColBegin; i < PjColEnd; i++)\n {\n // grab a row in the current column j of P\n const int PiRowIndex = ProwInd[i];\n bool rowPrevAppeared = 0;\n\n for (int jj = j - 1; jj >= 0; jj--) // scan columns of P backwards since it's more likely that the same row entry\n {\n // appeared in the direct neighbors of the current column\n const int PjjColBegin = PcolOffsets[jj];\n const int PjjColEnd = PcolOffsets[jj + 1];\n\n // check that the column is relevant (i.e. has a non-zero in the tIdx row)\n if (tIdx < ProwInd[PjjColBegin] || ProwInd[PjjColEnd - 1] < tIdx) { continue; }\n\n if (PiRowIndex < ProwInd[PjjColBegin] || ProwInd[PjjColEnd - 1] < PiRowIndex) { continue; }\n\n // make sure the column is actually relevant (i.e. has a non-zero in the tIdx row)\n bool relevantPcol = 0;\n\n for (int ii = PjjColBegin; ii < PjjColEnd; ii++)\n {\n if (ProwInd[ii] == tIdx) { relevantPcol = 1; break; }\n }\n\n if (!relevantPcol) { continue; }\n\n // see if PiRowIndex has already appeared previously\n for (int ii = PjjColBegin; ii < PjjColEnd; ii++)\n {\n if (ProwInd[ii] == PiRowIndex) { rowPrevAppeared = 1; break; }\n }\n\n if (rowPrevAppeared) { break; } // stop scanning columns of P, since we already accounted for PiRowIndex previously\n }\n\n if (!rowPrevAppeared) // if we haven't come across the P column with the nz in the PiRowIndex,\n {\n // then we can increment the row nnz count\n nnzPerRow++;\n }\n }\n }\n\n if (nnzPerRow == 0) { Ma_nnzPerRow[tIdx] = 1; }\n else if (Ma_nzDiagRows[tIdx]) { Ma_nnzPerRow[tIdx] = nnzPerRow; }\n else { Ma_nnzPerRow[tIdx] = nnzPerRow + 1; }\n }\n}\n\n\n/*\ntemplate \n__global__\nvoid find_Ma_nz_diag_rows_kernel(const IndexType *PcolOffsets, const IndexType *ProwInd,\n const IndexType PnumCols, bool *Ma_nzDiagRows)\n{\n for (unsigned int tIdx = threadIdx.x + blockDim.x*blockIdx.x; tIdx < PnumCols; tIdx += blockDim.x*gridDim.x)\n { // each thread goes thru a single column of P\n\n const unsigned int PcolBegin = PcolOffsets[tIdx];\n const unsigned int PcolEnd = PcolOffsets[tIdx+1];\n\n for (unsigned int i=PcolBegin; i\n__global__\nvoid add_invAijs_to_Ma_col_major_kernel( const IndexType *PcolOffsets, const IndexType *ProwInd, const IndexType PnumCols,\n const IndexType *AijOffsets, const ValueType *dense_invAijs,\n const IndexType *MaRowOffsets, const IndexType *MaColInd,\n ValueType *MaValues )\n{\n // go thru all invAij matrices\n for (unsigned int j = 0; j < PnumCols; j++)\n {\n // every thread runs through every invAij matrix and works on a single row of each invAij\n const int AijOffset = AijOffsets[j];\n const ValueType *invAij = dense_invAijs + AijOffset;\n const unsigned int PcolBegin = PcolOffsets[j];\n // dimension of the current Aij within dense_Aijs\n const unsigned int AijNumRows = PcolOffsets[j + 1] - PcolOffsets[j];\n const unsigned int lda = AijNumRows; // leading dim of the current Aij (col-major)\n\n for (unsigned int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AijNumRows; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on a single row of the current invAij\n const unsigned int MaRowIndex = ProwInd[PcolBegin + tIdx];\n // get the bounds of each thread's Ma row\n const unsigned int MaRowBegin = MaRowOffsets[MaRowIndex];\n const unsigned int MaRowEnd = MaRowOffsets[MaRowIndex + 1];\n unsigned int MaEntryIndex = MaRowBegin;\n\n for (unsigned int i = 0; i < AijNumRows; i++)\n {\n // go through each element in the given row of invAij and put this element\n // into its corresponding (\"spattered\") location in Ma\n const unsigned int MaColIndex = ProwInd[PcolBegin + i];\n\n // there are potentially more nz entries in a given row of Ma\n // than the dimension of the current invAij, hence need this loop\n // (note that we do not restart count since the col indices are stored in order)\n for ( ; MaEntryIndex < MaRowEnd; MaEntryIndex++)\n {\n if (MaColInd[MaEntryIndex] == MaColIndex)\n {\n MaValues[MaEntryIndex] += invAij[tIdx + i * lda];\n break;\n }\n }\n }\n }\n\n // to avoid race conditions when accessing rows of Ma, all threads need to be in sync.\n __syncthreads();\n }\n}\n\n\n// This kernel perturbs the diagonal of Ma (i.e. adds a specified value to each diag entry)\ntemplate \n__global__\nvoid perturb_Ma_diag_kernel( const IndexType *MaRowOffsets, const IndexType *MaColInd,\n ValueType *MaValues, const IndexType AnumRows,\n const ValueType perturb_mag, const bool *Ma_nzDiagRows )\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AnumRows; tIdx += blockDim.x * gridDim.x)\n {\n // each thread perturbs one diagonal entry of Ma\n const int MaRowBegin = MaRowOffsets[tIdx];\n const int MaRowEnd = MaRowOffsets[tIdx + 1];\n bool diagIsZero = 1;\n\n for (int j = MaRowBegin; j < MaRowEnd; j++)\n {\n // find the diagonal entry in the current row\n if (tIdx == MaColInd[j])\n {\n diagIsZero = 0;\n\n if (!Ma_nzDiagRows[tIdx]) { MaValues[j] = perturb_mag; break;}\n\n MaValues[j] += perturb_mag;\n break;\n }\n }\n\n if (diagIsZero)\n {\n printf(\"\\nWARNING: found zero diagonal entry of Ma in row %d. Ma_nzDiagRows[%d] = %d \\n\",\n tIdx, tIdx, Ma_nzDiagRows[tIdx]);\n break;\n }\n }\n}\n\n// This kernel initializes each column of P to v_x.\ntemplate \n__global__\nvoid init_Pvalues_kernel( const IndexType *PcolOffsets, const IndexType *ProwInd,\n const IndexType PnumCols, const ValueTypeB *v_x,\n ValueType *Pvalues )\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < PnumCols; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on one column of P in CSC format (P^T in CSR format)\n const unsigned int PcolBegin = PcolOffsets[tIdx];\n const unsigned int PcolEnd = PcolOffsets[tIdx + 1];\n\n for (int j = PcolBegin; j < PcolEnd; j++)\n {\n // go thru each entry in the given column of P\n Pvalues[j] = v_x[ProwInd[j]];\n }\n }\n}\n\n\n\n////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nEM_Interpolator >::\nEM_Interpolator() : EM_InterpolatorBase(),\n m_cuds_handle(0),\n m_dense_Aijs(0),\n m_dense_invAijs(0),\n m_ipiv(0),\n m_cuds_info(0),\n m_cuds_wspace(0)\n{\n // Allocate a handle for cudense\n cusolverStatus_t status = cusolverDnCreate(&m_cuds_handle);\n\n if ( status != CUSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not create the CUDENSE handle\", AMGX_ERR_CUDA_FAILURE );\n }\n\n // Define the cudense stream.\n status = cusolverDnSetStream(m_cuds_handle, amgx::thrust::global_thread_handle::get_stream());\n\n if ( status != CUSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not set the stream for CUDENSE\", AMGX_ERR_CUDA_FAILURE );\n }\n\n allocMem(m_cuds_info, sizeof(int), false);\n}\n\n\ntemplate \nEM_Interpolator >::~EM_Interpolator()\n{\n if (m_cuds_handle)\n {\n cusolverDnDestroy(m_cuds_handle);\n }\n\n if (m_dense_Aijs)\n {\n amgx::memory::cudaFreeAsync(m_dense_Aijs);\n m_dense_Aijs = 0;\n }\n\n if (m_dense_invAijs)\n {\n amgx::memory::cudaFreeAsync(m_dense_invAijs);\n m_dense_invAijs = 0;\n }\n\n if (m_ipiv)\n {\n amgx::memory::cudaFreeAsync(m_ipiv);\n m_ipiv = 0;\n }\n\n if (m_cuds_wspace)\n {\n amgx::memory::cudaFreeAsync(m_cuds_wspace);\n m_cuds_wspace = 0;\n }\n\n if (m_cuds_info)\n {\n amgx::memory::cudaFreeAsync(m_cuds_info);\n m_cuds_info = 0;\n }\n}\n\n\n\n// Private functions for the implementation on device\n\n// Find the sparse dimensions of P (i.e. number of non-zeros in each column of P)\ntemplate \nvoid EM_Interpolator >\n::computePsparsity( const Matrix_d &A, const IVector &cf_map, const IVector &coarse_idx,\n Matrix_d &P, IVector &PnnzPerCol,\n BVector &Ma_nzDiagRows )\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // Assemble all the raw pointers needed\n // Matrix A\n const IndexType *ArowOffsets_ptr = A.row_offsets.raw();\n const IndexType *AcolInd_ptr = A.col_indices.raw();\n const ValueType *Avalues_ptr = A.values.raw();\n const IndexType AnumRows = (IndexType) A.get_num_rows();\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (AnumRows / blocksize + 1));\n // temporary vectors - diagonal and non-zero offsets\n VVector Adiag(A.get_num_rows(), 0.0);\n ValueType *Adiag_ptr = Adiag.raw();\n // fullToCoarseMap array gives the COARSE index corresponding to the given FULL index\n // (meaningless if the full index corresponds to a FINE pt)\n IntVector fullToCoarseMap(AnumRows, 0);\n int *fullToCoarseMap_ptr = fullToCoarseMap.raw();\n amgx::thrust::transform_exclusive_scan( cf_map.begin(), cf_map.end(), fullToCoarseMap.begin(),\n is_coarse(), 0, amgx::thrust::plus() );\n cudaCheckError();\n // Extract the diagonal of A\n find_diag_kernel_indexed_dia <<< numBlocks, blocksize>>>(AnumRows, A.diag.raw(), Avalues_ptr, Adiag_ptr);\n cudaCheckError();\n // Coarse/Fine map\n const int *cf_map_ptr = cf_map.raw();\n // Find the sparse dimensions of P (i.e. number of non-zeros in each column of P)\n // Count the number of non-zeros in the interpolation if\n const int numCoarse = coarse_idx.size();\n const int *coarse_idx_ptr = coarse_idx.raw();\n VVector thresh_per_row(AnumRows);\n count_PnnzPerCol_kernel\n <<< numBlocks, blocksize>>>(ArowOffsets_ptr, AcolInd_ptr, Avalues_ptr,\n Adiag_ptr, numCoarse,\n cf_map_ptr, fullToCoarseMap_ptr, coarse_idx_ptr,\n PnnzPerCol.raw(), thresh_per_row.raw());\n cudaCheckError();\n printMatrixStats(A);\n cudaCheckError();\n // Column offsets of P (CSC format)\n IntVector PnzOffsets(numCoarse + 1);\n PnzOffsets[0] = 0;\n // get the offsets in P with an inclusive scan\n thrust_wrapper::inclusive_scan(PnnzPerCol.begin(), PnnzPerCol.end(), PnzOffsets.begin() + 1);\n cudaCheckError();\n // get total num of non-zeros in P\n const int Pnnz = PnzOffsets[numCoarse];\n // resize P\n P.resize(0, 0, 0, 1);\n P.addProps(CSR); // IMPORTANT: P actually stores P^T in CSR format\n // because we build P by columns (and CSC format is not supported)\n P.resize(numCoarse, AnumRows, Pnnz, 1);\n // set P offsets (P column offsets or P^T row offsets)\n amgx::thrust::copy(PnzOffsets.begin(), PnzOffsets.end(), P.row_offsets.begin());\n cudaCheckError();\n // treat P as if it were in the CSC format.\n const IndexType *PcolOffsets_ptr = P.row_offsets.raw();\n IndexType *ProwInd_ptr = P.col_indices.raw();\n bool *Ma_nzDiagRows_ptr = Ma_nzDiagRows.raw();\n // Since CSC format is not supported in AMGX matrix, we instead\n // compute P^T (which is stored in P) in the CSR format.\n // Initialize P by performing greedy aggregation based on the coarse set (cf_map)\n // The values of P remain uninitialized, since at this point we only get the sparsity pattern of P.\n init_ProwInd_greedy_aggregation_kernel\n <<< numBlocks, blocksize>>>(ArowOffsets_ptr, AcolInd_ptr, Avalues_ptr,\n Adiag_ptr, AnumRows,\n cf_map_ptr, fullToCoarseMap_ptr,\n PcolOffsets_ptr, ProwInd_ptr,\n Ma_nzDiagRows_ptr);\n cudaCheckError();\n}\n\n\n// Compute the submatrices Aij's of A and their inverses invAij's\ntemplate \nvoid EM_Interpolator >\n::computeAijSubmatrices(const Matrix_d &A, const int numCoarse, const Matrix_d &P,\n ValueType *dense_Aijs, ValueType *dense_invAijs,\n const IntVector &AijOffsets, int *ipiv,\n cusolverDnHandle_t &cuds_handle, int *cuds_info)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() / blocksize + 1));\n // Assemble all the raw pointers needed\n // Matrix A\n const IndexType *ArowOffsets_ptr = A.row_offsets.raw();\n const IndexType *AcolInd_ptr = A.col_indices.raw();\n const ValueType *Avalues_ptr = A.values.raw();\n // Treat P as if it were in the CSC format.\n const IndexType *PcolOffsets_ptr = P.row_offsets.raw();\n const IndexType *ProwInd_ptr = P.col_indices.raw();\n const IndexType *AijOffsets_ptr = AijOffsets.raw();\n // Extract Aij's from A based on the sparsity pattern of each column of P.\n // Column major format is used for storage of Aij's.\n extract_dense_Aijs_col_major_kernel\n <<< numBlocks, blocksize>>>(ArowOffsets_ptr, AcolInd_ptr, Avalues_ptr,\n PcolOffsets_ptr, ProwInd_ptr, numCoarse,\n AijOffsets_ptr, dense_Aijs);\n cudaCheckError();\n // Initialize invAij's to (dense) identity matrices.\n init_dense_invAijs_col_major_kernel\n <<< numBlocks, blocksize>>>(PcolOffsets_ptr, numCoarse,\n AijOffsets_ptr, dense_invAijs);\n cudaCheckError();\n // Sequentially perform dense LU factorization on all Aij's\n // LU factors stored in place of the Aij's\n // Then, perform inversion of Aij submatrices by solving with RHS == identity matrix\n cusolverStatus_t cudsStat;\n\n if (!cuds_handle)\n {\n cusolverStatus_t status = cusolverDnCreate(&m_cuds_handle);\n\n if ( status != CUSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not create the CUDENSE handle\", AMGX_ERR_CUDA_FAILURE );\n }\n\n // Define the cudense stream.\n status = cusolverDnSetStream(m_cuds_handle, amgx::thrust::global_thread_handle::get_stream());\n\n if ( status != CUSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not set the stream for CUDENSE\", AMGX_ERR_CUDA_FAILURE );\n }\n }\n\n // Pointer to info from cudense\n if (!cuds_info) { cuds_info = (int *) malloc(sizeof(int)); }\n\n // Offset of cudense pivoting sequence for each Aij.\n int ipivOffset = 0;\n IVector_h sizes(numCoarse);\n int max_wsize = 0;\n\n for (int j = 0; j < numCoarse; j++)\n {\n int cur_wsize = 0;\n int AijOffset = AijOffsets[j];\n int AijNumRows = P.row_offsets[j + 1] - P.row_offsets[j];\n int lda = AijNumRows;\n cusolverStatus_t status = cusolverDnXgetrf_bufferSize(cuds_handle, AijNumRows, AijNumRows, dense_Aijs + AijOffset, lda, &cur_wsize);\n\n if (status != CUSOLVER_STATUS_SUCCESS)\n {\n FatalError(\"failed to perform LU factorization\", AMGX_ERR_INTERNAL);\n }\n\n max_wsize = std::max(cur_wsize, max_wsize);\n }\n\n allocMem(m_cuds_wspace, max_wsize, false);\n\n for (int j = 0; j < numCoarse; j++)\n {\n int AijOffset = AijOffsets[j]; // offset of the current Aij\n int AijNumRows = P.row_offsets[j + 1] - P.row_offsets[j]; // dimension of the next Aij within dense_Aijsc\n int lda = AijNumRows; // leading dim for the next Aij (col-major)\n sizes[j] = AijNumRows;\n // [L,U,P] = lu(Aij) (LU factors stored in place of the current Aij)\n cudsStat = cusolverDnXgetrf(cuds_handle, AijNumRows, AijNumRows,\n dense_Aijs + AijOffset, lda, m_cuds_wspace,\n ipiv + ipivOffset, cuds_info);\n\n if (cudsStat != CUSOLVER_STATUS_SUCCESS)\n {\n FatalError(\"failed to perform LU factorization\", AMGX_ERR_INTERNAL);\n }\n\n // invAij = U\\(L\\P) (i.e. find inverse of Aij).\n cudsStat = cusolverDnXgetrs(cuds_handle, CUBLAS_OP_N, AijNumRows, AijNumRows,\n dense_Aijs + AijOffset, lda,\n ipiv + ipivOffset, dense_invAijs + AijOffset, lda, cuds_info );\n\n if (cudsStat != CUSOLVER_STATUS_SUCCESS)\n {\n FatalError(\"failed to perform triangular solve\", AMGX_ERR_INTERNAL);\n }\n\n ipivOffset += AijNumRows; // starting offset of the next pivoting sequence\n }\n}\n\n\n// Compute Ma matrix\ntemplate \nvoid EM_Interpolator >\n::computeMa(Matrix_d &Ma, const int AnumRows, const int numCoarse, const Matrix_d &P,\n const ValueType *dense_invAijs, const IntVector &AijOffsets,\n const BVector &Ma_nzDiagRows,\n const bool perturb_Ma_diag, const ValueType perturb_mag)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (AnumRows / blocksize + 1));\n // Treat P as if it were in CSC format.\n const IndexType *PcolOffsets_ptr = P.row_offsets.raw();\n const IndexType *ProwInd_ptr = P.col_indices.raw();\n // Number of non-zeros in each row of Ma\n IntVector Ma_nnzPerRow(AnumRows, 0);\n // 1. Compute Ma_nnzPerRow (kernel).\n count_Ma_nnz_per_row_kernel <<< numBlocks, blocksize>>>( PcolOffsets_ptr, ProwInd_ptr, numCoarse,\n Ma_nnzPerRow.raw(), AnumRows, Ma_nzDiagRows.raw() );\n cudaCheckError();\n // 2. Resize Ma and assign MaRowOffsets by performing a scan on Ma_nnzPerRow.\n // Get the offsets in Ma with an inclusive scan (in-place)\n thrust_wrapper::inclusive_scan(Ma_nnzPerRow.begin(), Ma_nnzPerRow.end(), Ma_nnzPerRow.begin());\n cudaCheckError();\n // Get total num of non-zeros in Ma\n const int Ma_nnz = Ma_nnzPerRow[AnumRows - 1];\n // Resize Ma\n Ma.resize(0, 0, 0, 1);\n Ma.addProps(CSR);\n Ma.resize(AnumRows, AnumRows, Ma_nnz, 1);\n Ma.row_offsets[0] = 0;\n // set Ma row offsets\n amgx::thrust::copy(Ma_nnzPerRow.begin(), Ma_nnzPerRow.end(), Ma.row_offsets.begin() + 1);\n cudaCheckError();\n const IndexType *MaRowOffsets_ptr = Ma.row_offsets.raw();\n IndexType *MaColInd_ptr = Ma.col_indices.raw();\n ValueType *MaValues_ptr = Ma.values.raw();\n // 3. Get the unsorted column indices of Ma.\n get_unsorted_MaColInd_kernel <<< numBlocks, blocksize>>>(PcolOffsets_ptr, ProwInd_ptr, numCoarse,\n MaRowOffsets_ptr, MaColInd_ptr, AnumRows,\n Ma_nzDiagRows.raw());\n cudaCheckError();\n\n // 4. Sort column indices to get the final MaColInd.\n for (unsigned int MaRow = 0; MaRow < AnumRows; MaRow++)\n {\n int MaRowBegin = Ma.row_offsets[MaRow];\n int MaRowEnd = Ma.row_offsets[MaRow + 1];\n\n if (MaRowEnd - MaRowBegin > 1)\n {\n thrust_wrapper::sort( Ma.col_indices.begin() + Ma.row_offsets[MaRow],\n Ma.col_indices.begin() + Ma.row_offsets[MaRow + 1] );\n }\n }\n\n // 5. Compute MaValues by summing up the \"spattered\" inverses of Aij submatrices.\n add_invAijs_to_Ma_col_major_kernel\n <<< numBlocks, blocksize>>>(PcolOffsets_ptr, ProwInd_ptr, numCoarse,\n AijOffsets.raw(), dense_invAijs,\n MaRowOffsets_ptr, MaColInd_ptr,\n MaValues_ptr);\n cudaCheckError();\n\n // Perturb the diagonal of Ma here (prevents Ma from being singular)\n // Ma = 1.0e-8*speye(n,n);\n if (perturb_Ma_diag)\n {\n // Default: perturb_mag = 1.0e-8\n perturb_Ma_diag_kernel\n <<< numBlocks, blocksize>>>(MaRowOffsets_ptr, MaColInd_ptr,\n MaValues_ptr, AnumRows,\n perturb_mag, Ma_nzDiagRows.raw());\n cudaCheckError();\n }\n\n Ma.set_initialized(1);\n}\n\n\n// Solve for x: Ma*x=e (solved with pcg)\ntemplate \nvoid EM_Interpolator >\n::solveMa_e(Matrix_d &Ma, const int AnumRows, Vector_d &v_x)\n{\n typedef typename Matrix_d::value_type ValueType;\n // 1. Initialize x (vector of 0's) and e (vector of 1's)\n v_x.resize(AnumRows, 0.0);\n Vector_d v_e(AnumRows, 1.0);\n // ------------------------------ define pcg solver ---------------------------------\n // temporary config and pointer to main config (json format)\n AMG_Config cfg;\n std::string cfg_string = \"{ \\\"config_version\\\": 2\";\n cfg_string += \", \\\"solver\\\": { \\\"scope\\\": \\\"pcg_scope\\\", \\\"solver\\\": \\\"PCG\\\"\"; // Set the solver to PCG\n cfg_string += \" ,\\\"preconditioner\\\": \\\"NOSOLVER\\\"\"; // IMPORTANT to have this in order to allocate PCG\n cfg_string += \" } }\";\n cfg.parseParameterString(cfg_string.c_str());\n // define and allocate solver\n Solver *Ma_solver = 0;\n Ma_solver = SolverFactory::allocate(cfg, \"default\", \"solver\");\n // ------------------------------ end define pcg ------------------------------------\n const int Ma_solveIts = 20;\n const ValueType Ma_solveTol = (ValueType) 1.e-4;\n Ma_solver->setup(Ma, false);\n Ma_solver->set_max_iters(Ma_solveIts);\n Ma_solver->setTolerance(Ma_solveTol);\n // Apply pcg to solve x = Ma \\ e\n Ma_solver->solve( v_e, v_x, true );\n Vector_d v_res(AnumRows, 0.);\n Ma_solver->compute_residual( v_e, v_x, v_res );\n\n if (Ma_solver) { delete Ma_solver; Ma_solver = 0; }\n}\n\n\n// Compute the values of P (using LU factorization of Aij submatrices)\ntemplate \nvoid EM_Interpolator >\n::computePvalues( const int AnumRows, const int numCoarse, Matrix_d &P, const Vector_d &v_x,\n const ValueType *dense_Aijs, const IntVector &AijOffsets, const int *ipiv,\n cusolverDnHandle_t &cuds_handle, int *cuds_info )\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n typedef typename Vector_d::value_type ValueTypeB;\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (AnumRows / blocksize + 1));\n // Treat P as if it were in CSC format.\n const IndexType *PcolOffsets_ptr = P.row_offsets.raw();\n const IndexType *ProwInd_ptr = P.col_indices.raw();\n ValueType *Pvalues_ptr = P.values.raw();\n init_Pvalues_kernel\n <<< numBlocks, blocksize>>>(PcolOffsets_ptr, ProwInd_ptr, numCoarse,\n v_x.raw(), Pvalues_ptr);\n cudaCheckError();\n cusolverStatus_t cudsStat;\n\n if (!cuds_handle)\n {\n cusolverStatus_t status = cusolverDnCreate(&m_cuds_handle);\n\n if ( status != CUSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not create the CUDENSE handle\", AMGX_ERR_CUDA_FAILURE );\n }\n\n // Define the cudense stream.\n status = cusolverDnSetStream(m_cuds_handle, amgx::thrust::global_thread_handle::get_stream());\n\n if ( status != CUSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not set the stream for CUDENSE\", AMGX_ERR_CUDA_FAILURE );\n }\n }\n\n // Pointer to info from cudense\n if (!cuds_info) { cuds_info = (int *) malloc(sizeof(int)); }\n\n // Offset of cudense pivoting sequence for each Aij.\n int ipivOffset = 0;\n\n for (int j = 0; j < numCoarse; j++)\n {\n int AijOffset = AijOffsets[j];\n int PcolOffset = P.row_offsets[j];\n // dimension of the next Aij within dense_Aijs\n int AijNumRows = P.row_offsets[j + 1] - P.row_offsets[j];\n int lda = AijNumRows; // leading dim for the next Aij (col-major)\n\n if (lda == 0) { continue; }\n\n // invAij = U\\(L\\P) (i.e. find inverse of Aij).\n cudsStat = cusolverDnXgetrs(cuds_handle, CUBLAS_OP_N, AijNumRows, 1,\n dense_Aijs + AijOffset, lda,\n ipiv + ipivOffset, Pvalues_ptr + PcolOffset, lda, cuds_info);\n\n if (cudsStat != CUSOLVER_STATUS_SUCCESS)\n { FatalError(\"failed to perform triangular solve\", AMGX_ERR_INTERNAL); }\n\n ipivOffset += AijNumRows; // starting offset of the next pivoting sequence\n }\n\n if (cuds_info) { free(cuds_info); cuds_info = 0; }\n}\n\n\n\n// Compute the values of P (using inverses of Aij submatrices)\ntemplate \nvoid EM_Interpolator >\n::computePvalues( const int AnumRows, const int numCoarse, Matrix_d &P, const Vector_d &v_x,\n const ValueType *dense_invAijs, const IntVector &AijOffsets )\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n typedef typename Vector_d::value_type ValueTypeB;\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (AnumRows / blocksize + 1));\n // Treat P as if it were in CSC format.\n const IndexType *PcolOffsets_ptr = P.row_offsets.raw();\n const IndexType *ProwInd_ptr = P.col_indices.raw();\n ValueType *Pvalues_ptr = P.values.raw();\n init_Pvalues_kernel\n <<< numBlocks, blocksize>>>(PcolOffsets_ptr, ProwInd_ptr, numCoarse,\n v_x.raw(), Pvalues_ptr);\n cudaCheckError();\n ValueType *place_holder_ptr = 0;\n allocMem(place_holder_ptr, sizeof(ValueType), false);\n ValueType one = (ValueType) 1.;\n ValueType zero = (ValueType) 0.;\n\n for (int j = 0; j < numCoarse; j++)\n {\n int AijOffset = AijOffsets[j];\n int PcolOffset = P.row_offsets[j];\n // dimension of the next Aij within dense_Aijs\n int AijNumRows = P.row_offsets[j + 1] - P.row_offsets[j];\n int lda = AijNumRows; // leading dim for the next Aij (col-major)\n\n if (lda == 0) { continue; }\n\n Cublas::gemv(false, AijNumRows, AijNumRows, &one,\n dense_invAijs + AijOffset, lda,\n Pvalues_ptr + PcolOffset, 1,\n &zero, place_holder_ptr, 1);\n cudaDeviceSynchronize();\n cudaCheckError();\n }\n\n if (place_holder_ptr) { amgx::memory::cudaFreeAsync(place_holder_ptr); place_holder_ptr = 0; }\n}\n\n\n\n// ----------------------------\n// specialization for device\n// ----------------------------\ntemplate \nvoid EM_Interpolator >\n::generateInterpolationMatrix_1x1(const Matrix_d &A, const IntVector &cf_map,\n Matrix_d &P, void *amg)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // Number of rows in Matrix A\n const IndexType AnumRows = (IndexType) A.get_num_rows();\n // Count the number of coarse points\n const int numCoarse = (int) amgx::thrust::count(cf_map.begin(), cf_map.end(), (int)COARSE);\n cudaCheckError();\n IVector coarse_idx(numCoarse, -1);\n amgx::thrust::counting_iterator zero(0);\n amgx::thrust::counting_iterator zero_plus_nc = zero + numCoarse;\n amgx::thrust::copy_if(zero, zero_plus_nc, cf_map.begin(), coarse_idx.begin(), is_coarse());\n // ..................................................................................\n // ------------------------ Begin Compute P code ------------------------------------\n // Store the number of non-zeros in each column of the interpolation matrix P\n IntVector PnnzPerCol(numCoarse);\n BVector Ma_nzDiagRows(AnumRows, 0);\n // Find the sparse dimensions of P (i.e. number of non-zeros in each column of P)\n computePsparsity(A, cf_map, coarse_idx, P, PnnzPerCol, Ma_nzDiagRows);\n // Get total num of non-zeros in P\n int Pnnz = P.row_offsets[numCoarse];\n // ..................................................................................\n // Now follows Jacobi smoothing (skip for now since it might be unnecessary)\n // ..................................................................................\n // -------------------------- End Compute P code ------------------------------------\n // ..................................................................................\n // ..................................................................................\n // ----------------------- Begin Energy Minimization Code ---------------------------\n // Given:\n // A: the problem matrix\n // P: sparsity pattern for P (n x nc)\n // Main loop (Ma here is the inverse of Ma in the paper)\n // I. Build Ma.\n // 0. Allocate space for storing LU factors for each Aij submatrix of A (use dense format)\n // There are numCoarse of LU factors.\n // Since the number of non-zeros in each column of P determines the size of each Aij,\n // we know the dimensions of all Aij based on the row offsets of P^T.\n // Calculate offset of each Aij matrix in the dense_Aijs array\n IntVector AijOffsets(numCoarse + 1);\n AijOffsets[0] = 0;\n // Get the offsets with an inclusive scan (nnz of each column of P is the dimension of each Aij)\n amgx::thrust::transform_inclusive_scan( PnnzPerCol.begin(), PnnzPerCol.end(), AijOffsets.begin() + 1,\n square(), amgx::thrust::plus() );\n cudaCheckError();\n PnnzPerCol.clear();\n // Cumulative num of entries in all (dense) Aijs (last entry in the offsets)\n const int Aijs_size = AijOffsets[numCoarse];\n // Allocate memory for the dense Aij submatrices and initialize to zero.\n allocMem(m_dense_Aijs, Aijs_size, true);\n // Allocate memory for the dense invAij's and initialize to zero.\n allocMem(m_dense_invAijs, Aijs_size, true);\n // Allocate memory for cudense pivoting sequence.\n allocMem(m_ipiv, Pnnz, false);\n // Compute the submatrices Aij's of A and the inverses invAij's\n computeAijSubmatrices(A, numCoarse, P, m_dense_Aijs, m_dense_invAijs,\n AijOffsets, m_ipiv, m_cuds_handle, m_cuds_info);\n // Ma matrix build sequence.\n // 1. Compute Ma_nnzPerRow (kernel).\n // 2. Resize Ma and assign MaRowOffsets by performing a scan on Ma_nnzPerRow (thrust).\n // 3. Get the unsorted column indices of Ma (kernel).\n // 4. Sort column indices to get the final MaColInd (thrust).\n // 5. Compute MaValues by summing up the \"spattered\" inverses of Aij submatrices (kernel).\n // Declare Ma\n Matrix_d Ma;\n // Compute Ma matrix\n computeMa(Ma, AnumRows, numCoarse, P, m_dense_invAijs, AijOffsets, Ma_nzDiagRows);\n\n if (m_dense_invAijs) { amgx::memory::cudaFreeAsync(m_dense_invAijs); m_dense_invAijs = 0; }\n\n /* At this point, we are done with the following Matlab code analogue.\n for (int j=1; j<=nc; j++)\n {\n // Find the dofs in the j-th block\n ij = find(p(:,j)); // col vec of subindices\n ni = length(ij);\n It = sparse(n,ni);\n It(ij,:) = speye(ni);\n Aij{j} = inv(A(ij,ij));\n // Form Ma\n // A(ii,ii): submatrix of A corresponding to the dofs in the i-th block\n Ma = Ma + It*(Aij{j}*It_T);\n }\n */\n Ma_nzDiagRows.clear();\n // II. Solve for x: Ma*x=e\n //m_Mae_norm = (ValueType) 0.;\n Vector_d v_x;\n solveMa_e(Ma, AnumRows, v_x);\n cudaDeviceSynchronize();\n cudaCheckError();\n Ma.set_initialized(0);\n Ma.resize(0, 0, 0, 1);\n /* Now, we are done with the following Matlab code analogue.\n e = ones(n,1);\n Mae = zeros(size(e)); // Ma\\e;\n [Mae,flag,relres,iter,resvecCG] = pcg(Ma,e,1e-4,100);\n */\n // III. Construct the columns of P in parallel.\n // P is already initialized (and at this point stores P^T in CSR format)\n // We will be performing the following operation to build each column of P\n // P(:,j) = It{j}*(invAij{j}*(It^T*x));\n // This is done with the following steps.\n // First, apply It{j}^T mask to x to get x{j} (project x onto the dimension of Aij{j}).\n // Then solve Aij\\x{j}\n // Compute the values of P\n computePvalues( AnumRows, numCoarse, P, v_x, m_dense_Aijs, AijOffsets, m_ipiv, m_cuds_handle, m_cuds_info );\n\n if (m_dense_Aijs) { amgx::memory::cudaFreeAsync(m_dense_Aijs); m_dense_Aijs = 0; }\n\n if (m_ipiv) { amgx::memory::cudaFreeAsync(m_ipiv); m_ipiv = 0; }\n\n if (m_cuds_info) { amgx::memory::cudaFreeAsync(m_cuds_info); m_cuds_info = 0; }\n\n //if (m_dense_invAijs) { amgx::memory::cudaFreeAsync(m_dense_invAijs); m_dense_invAijs = 0; }\n /* Now, we are done building columns of P (stored as rows of P^T in CSR).\n * Here's the equivalent matlab code.\n P = zeros(n,nc);\n for (int j=1; j<=nc; j++)\n {\n // Find the dofs in the j-th block\n ij=find(p(:,j));\n ni = length(ij);\n It = zeros(n,ni);\n It(ij,:) = eye(ni);\n\n // Form j-th column of P\n P(:,j) = It*(Aij{j}*(It_T*Mae));\n }\n */\n //TODO: truncate rows/cols of P to keep the sparsity low on all levels.\n // ------------------------- End Energy Minimization Code ---------------------------\n // ..................................................................................\n}\n\n\n\ntemplate< class T_Config>\nvoid EM_InterpolatorBase\n::generateInterpolationMatrix(const Matrix &A, const IntVector &cf_map,\n Matrix &P, void *amg)\n{\n P.set_initialized(0);\n\n if (A.get_block_size() == 1)\n {\n generateInterpolationMatrix_1x1(A, cf_map, P, amg);\n }\n else\n {\n FatalError(\"Unsupported block size for distance1 interpolator\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n P.computeDiagonal();\n P.set_initialized(1);\n}\n\n#define AMGX_CASE_LINE(CASE) template class EM_InterpolatorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class EM_Interpolator::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace energymin\n} // namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include // std::signbit\n#include \n\n#include \"util.h\"\n\n#include \"amgx_cublas.h\"\n\nnamespace amgx\n{\nnamespace energymin\n{\n\n\n/*************************************************************************\n* create the interpolation matrix P by computing p and using energy minimization\n************************************************************************/\n\n\n/*************************************************************************\n* create the interpolation matrix P on GPU\n************************************************************************/\n\ntemplate \n__global__\nvoid count_PnnzPerCol_kernel( const IndexType *ArowOffsets, const IndexType *AcolInd, const ValueType *Avalues,\n const ValueType *Adiag, const IndexType numCoarse,\n const int *cf_map, const int *fullToCoarseMap, const int *coarse_idx_ptr,\n int *PnnzPerCol, ValueType *thresh_per_row )\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numCoarse; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on one \"coarse\" row of A at a time\n const int row_idx = coarse_idx_ptr[tIdx];\n const int ArowBegin = ArowOffsets[row_idx];\n const int ArowEnd = ArowOffsets[row_idx + 1];\n // Matlab pseudo-code:\n // p = sparse(n, nc);\n // p(coarse,:) = speye(nc);\n // fAj = AD(j,fine)/AD(j,j) < -.1; // fine neighbors of j with oppositive sign from AD_jj\n // become the non-zeros in initial P.\n ValueType minVal(0), maxVal(0), diag_value(0);\n ValueType fine_avg_thresh = 0.;\n const ValueType thresh = (ValueType)0.25;\n\n for (int j = ArowBegin; j < ArowEnd; j++)\n {\n ValueType aValue = Avalues[j];\n\n if ( j != tIdx )\n {\n minVal = min( minVal, aValue );\n maxVal = max( maxVal, aValue );\n }\n else\n {\n diag_value = aValue;\n }\n }\n\n if ( diag_value < 0.0)\n {\n fine_avg_thresh = thresh * maxVal;\n }\n else\n {\n fine_avg_thresh = thresh * minVal;\n }\n\n thresh_per_row[tIdx] = fine_avg_thresh;\n int p_col_nnz = 0;\n\n for (int j = ArowBegin; j < ArowEnd; j++)\n {\n int AjColIndex = AcolInd[j];\n\n if (AjColIndex == tIdx)\n {\n p_col_nnz++;\n continue;\n }\n\n if (cf_map[AjColIndex] == FINE)\n {\n if (diag_value * Avalues[j] < diag_value * fine_avg_thresh)\n {\n p_col_nnz++;\n }\n }\n }\n\n PnnzPerCol[tIdx] = p_col_nnz;\n }\n}\n\n\ntemplate \n__global__\nvoid init_ProwInd_greedy_aggregation_kernel(const IndexType *ArowOffsets, const IndexType *AcolInd, const ValueType *Avalues,\n const ValueType *Adiag, const IndexType AnumRows,\n const int *cf_map, const int *fullToCoarseMap,\n const IndexType *PcolOffsets, IndexType *ProwInd,\n bool *Ma_nzDiagRows)\n{\n // for each coarse point form aggregates\n /* Matlab-like pseudo-code\n p = sparse(n, nc);\n p(coarse,:) = speye(nc);\n for (int j : coarse) // coarse point with FULL index j\n fAj = AD(j,fine)/AD(j,j) < -.1; // fine neighbors of j with oppositive sign from AD_jj\n P(fine(fAj),cj)=1; // become the non-zeros in initial P.\n end\n */\n\n // size of P and its column offsets were precomputed.\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AnumRows; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on one \"coarse\" row of A at a time (builds a column of P)\n const int ArowBegin = ArowOffsets[tIdx];\n const int ArowEnd = ArowOffsets[tIdx + 1];\n\n if (cf_map[tIdx] == COARSE) // \"coarse\" row of A\n {\n // P is CSC format since each thread works on each coarse point\n // and is therefore building a single column of P.\n // There are numCoarse columns of P\n // Columns of P correspond to rows of A (A is assumed symmetric)\n bool PcolCoarseDiagUnset = 1;\n const int PcolIndex = fullToCoarseMap[tIdx];\n int PentryNum = PcolOffsets[PcolIndex]; // starting location of the current column of P\n\n for (int j = ArowBegin; j < ArowEnd; j++)\n {\n // go thru each \"fine\" column of A\n int AjColIndex = AcolInd[j];\n\n if (cf_map[AjColIndex] == FINE && Avalues[j] / Adiag[tIdx] < -.1)\n {\n //p(AjColIndex, fullToCoarseMap(j))=1;\n if (AjColIndex > tIdx && PcolCoarseDiagUnset) // COARSE diagonal of P is identity\n {\n // set COARSE diagonal of P\n ProwInd[PentryNum] = tIdx; // set COARSE diagonal entry of P to 1.\n PentryNum++;\n PcolCoarseDiagUnset = 0;\n Ma_nzDiagRows[tIdx] = 1; // record the newly added row index of P\n }\n\n ProwInd[PentryNum] = AjColIndex; // columns of A become rows in P (since A is assumed symmetric)\n PentryNum++;\n Ma_nzDiagRows[AjColIndex] = 1; // record the newly added row index of P\n }\n }\n\n if (PcolCoarseDiagUnset) // COARSE diagonal of P is identity\n {\n // set COARSE diagonal of P\n ProwInd[PentryNum] = tIdx; // columns of A become rows in P (since A is assumed symmetric)\n PentryNum++;\n PcolCoarseDiagUnset = 0;\n Ma_nzDiagRows[tIdx] = 1; // record the newly added row index of P\n }\n }\n }\n}\n\n\ntemplate \n__global__\nvoid extract_dense_Aijs_col_major_kernel(const IndexType *ArowOffsets, const IndexType *AcolInd, const ValueType *Avalues,\n const IndexType *PcolOffsets, const IndexType *ProwInd, const IndexType PnumCols,\n const IndexType *AijOffsets, ValueType *dense_Aijs)\n{\n /* Matlab-like pseudo-code\n for (int j=1; j<=nc; j++)\n {\n // Find the nz pattern in the j-th column of P\n ij = find(p(:,j)); // col vec of subindices\n Aij{j} = A(ij,ij);\n }\n */\n\n // each thread builds one Aij\n for (unsigned int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < PnumCols; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works with one column of P\n const unsigned int PcolBegin = PcolOffsets[tIdx];\n const unsigned int PcolEnd = PcolOffsets[tIdx + 1];\n const unsigned int AijNumRows = PcolEnd - PcolBegin;\n // pointer to the beginning of the specific Aij submatrix\n ValueType *dense_Aij = dense_Aijs + AijOffsets[tIdx];\n\n for (unsigned int i = PcolBegin; i < PcolEnd; i++)\n {\n // go thru each non-zero row of the given column of P\n // and scan the corresponding rows of A\n unsigned int PiRowIndex = ProwInd[i];\n const unsigned int ArowBegin = ArowOffsets[PiRowIndex];\n const unsigned int ArowEnd = ArowOffsets[PiRowIndex + 1];\n\n for (unsigned int j = ArowBegin; j < ArowEnd; j++)\n {\n // scan the current row of A and pick the elts based on the sparsity pattern of P\n unsigned int AjColIndex = AcolInd[j];\n\n // if AjColIndex is an index of a non-zero entry in the current column of P, take it.\n for (unsigned int ii = PcolBegin; ii < PcolEnd; ii++)\n {\n if (ProwInd[ii] == AjColIndex)\n {\n // add elt to Aij in column-major format\n unsigned int AijRowIndex = i - PcolBegin;\n unsigned int AijColIndex = ii - PcolBegin;\n dense_Aij[AijColIndex * AijNumRows + AijRowIndex] = Avalues[j]; //wrong\n break;\n }\n }\n }\n }\n }\n}\n\n\n// kernel to initialize invAij's to dense identity matrices\ntemplate \n__global__\nvoid init_dense_invAijs_col_major_kernel(const IndexType *PcolOffsets, const IndexType PnumCols,\n const IndexType *AijOffsets, ValueType *dense_invAijs)\n{\n for (unsigned int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < PnumCols; tIdx += blockDim.x * gridDim.x)\n {\n // each thread initializes one invAij\n const unsigned int PcolBegin = PcolOffsets[tIdx];\n const unsigned int PcolEnd = PcolOffsets[tIdx + 1];\n const unsigned int AijNumRows = PcolEnd - PcolBegin;\n // pointer to the beginning of the specific invAij submatrix\n ValueType *dense_invAij = dense_invAijs + AijOffsets[tIdx];\n\n for (unsigned int i = 0; i < AijNumRows; i++)\n {\n dense_invAij[i * AijNumRows + i] = 1;\n }\n }\n}\n\n\n\n\n__global__\nvoid count_Ma_nnz_per_row_kernel(const int *PcolOffsets, const int *ProwInd, const int PnumCols,\n int *Ma_nnzPerRow, const int AnumRows, const bool *Ma_nzDiagRows)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AnumRows; tIdx += blockDim.x * gridDim.x)\n {\n // every thread works on a single row of Ma (figures out nnz in that row)\n bool rowNotFilled = 1; // marks whether the Ma row nnz count has been initialized\n int nnzPerRow = 0;\n\n for (int j = 0; j < PnumCols; j++)\n {\n // each thread goes thru all columns of P\n const int PjColBegin = PcolOffsets[j];\n const int PjColEnd = PcolOffsets[j + 1];\n //const int PjCol_nnz = PjColEnd - PjColBegin;\n\n // check that the column is relevant (i.e. has a non-zero in the tIdx row)\n if (tIdx < ProwInd[PjColBegin] || ProwInd[PjColEnd - 1] < tIdx) { continue; }\n\n // Make sure the column is actually relevant (i.e. has a non-zero in the tIdx row)\n bool relevantCol = 0;\n\n for (int i = PjColBegin; i < PjColEnd; i++)\n {\n if (ProwInd[i] == tIdx) { relevantCol = 1; break; }\n }\n\n if (!relevantCol) { continue; }\n\n if (rowNotFilled)\n {\n nnzPerRow = PjColEnd - PjColBegin; // initial nnz per row of Ma\n rowNotFilled = 0;\n continue;\n }\n\n // if the row nnz count has been initialized, we need to figure out which elts haven't already appeared\n for (int i = PjColBegin; i < PjColEnd; i++)\n {\n // grab a row in the current column j of P\n const int PiRowIndex = ProwInd[i];\n bool rowPrevAppeared = 0;\n\n for (int jj = j - 1; jj >= 0; jj--) // scan columns of P backwards since it's more likely that the same row entry\n {\n // appeared in the direct neighbors of the current column\n const int PjjColBegin = PcolOffsets[jj];\n const int PjjColEnd = PcolOffsets[jj + 1];\n\n // check that the column is relevant (i.e. has a non-zero in the tIdx row)\n if (tIdx < ProwInd[PjjColBegin] || ProwInd[PjjColEnd - 1] < tIdx) { continue; }\n\n if (PiRowIndex < ProwInd[PjjColBegin] || ProwInd[PjjColEnd - 1] < PiRowIndex) { continue; }\n\n // make sure the column is actually relevant (i.e. has a non-zero in the tIdx row)\n bool relevantPcol = 0;\n\n for (int ii = PjjColBegin; ii < PjjColEnd; ii++)\n {\n if (ProwInd[ii] == tIdx) { relevantPcol = 1; break; }\n }\n\n if (!relevantPcol) { continue; }\n\n // see if PiRowIndex has already appeared previously\n for (int ii = PjjColBegin; ii < PjjColEnd; ii++)\n {\n if (ProwInd[ii] == PiRowIndex) { rowPrevAppeared = 1; break; }\n }\n\n if (rowPrevAppeared) { break; } // stop scanning columns of P, since we already accounted for PiRowIndex previously\n }\n\n if (!rowPrevAppeared) // if we haven't come across the P column with the nz in the PiRowIndex,\n {\n // then we can increment the row nnz count\n nnzPerRow++;\n }\n }\n }\n\n if (nnzPerRow == 0) { Ma_nnzPerRow[tIdx] = 1; }\n else if (Ma_nzDiagRows[tIdx]) { Ma_nnzPerRow[tIdx] = nnzPerRow; }\n else { Ma_nnzPerRow[tIdx] = nnzPerRow + 1; }\n }\n}\n\n\n/*\ntemplate \n__global__\nvoid find_Ma_nz_diag_rows_kernel(const IndexType *PcolOffsets, const IndexType *ProwInd,\n const IndexType PnumCols, bool *Ma_nzDiagRows)\n{\n for (unsigned int tIdx = threadIdx.x + blockDim.x*blockIdx.x; tIdx < PnumCols; tIdx += blockDim.x*gridDim.x)\n { // each thread goes thru a single column of P\n\n const unsigned int PcolBegin = PcolOffsets[tIdx];\n const unsigned int PcolEnd = PcolOffsets[tIdx+1];\n\n for (unsigned int i=PcolBegin; i\n__global__\nvoid add_invAijs_to_Ma_col_major_kernel( const IndexType *PcolOffsets, const IndexType *ProwInd, const IndexType PnumCols,\n const IndexType *AijOffsets, const ValueType *dense_invAijs,\n const IndexType *MaRowOffsets, const IndexType *MaColInd,\n ValueType *MaValues )\n{\n // go thru all invAij matrices\n for (unsigned int j = 0; j < PnumCols; j++)\n {\n // every thread runs through every invAij matrix and works on a single row of each invAij\n const int AijOffset = AijOffsets[j];\n const ValueType *invAij = dense_invAijs + AijOffset;\n const unsigned int PcolBegin = PcolOffsets[j];\n // dimension of the current Aij within dense_Aijs\n const unsigned int AijNumRows = PcolOffsets[j + 1] - PcolOffsets[j];\n const unsigned int lda = AijNumRows; // leading dim of the current Aij (col-major)\n\n for (unsigned int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AijNumRows; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on a single row of the current invAij\n const unsigned int MaRowIndex = ProwInd[PcolBegin + tIdx];\n // get the bounds of each thread's Ma row\n const unsigned int MaRowBegin = MaRowOffsets[MaRowIndex];\n const unsigned int MaRowEnd = MaRowOffsets[MaRowIndex + 1];\n unsigned int MaEntryIndex = MaRowBegin;\n\n for (unsigned int i = 0; i < AijNumRows; i++)\n {\n // go through each element in the given row of invAij and put this element\n // into its corresponding (\"spattered\") location in Ma\n const unsigned int MaColIndex = ProwInd[PcolBegin + i];\n\n // there are potentially more nz entries in a given row of Ma\n // than the dimension of the current invAij, hence need this loop\n // (note that we do not restart count since the col indices are stored in order)\n for ( ; MaEntryIndex < MaRowEnd; MaEntryIndex++)\n {\n if (MaColInd[MaEntryIndex] == MaColIndex)\n {\n MaValues[MaEntryIndex] += invAij[tIdx + i * lda];\n break;\n }\n }\n }\n }\n\n // to avoid race conditions when accessing rows of Ma, all threads need to be in sync.\n __syncthreads();\n }\n}\n\n\n// This kernel perturbs the diagonal of Ma (i.e. adds a specified value to each diag entry)\ntemplate \n__global__\nvoid perturb_Ma_diag_kernel( const IndexType *MaRowOffsets, const IndexType *MaColInd,\n ValueType *MaValues, const IndexType AnumRows,\n const ValueType perturb_mag, const bool *Ma_nzDiagRows )\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AnumRows; tIdx += blockDim.x * gridDim.x)\n {\n // each thread perturbs one diagonal entry of Ma\n const int MaRowBegin = MaRowOffsets[tIdx];\n const int MaRowEnd = MaRowOffsets[tIdx + 1];\n bool diagIsZero = 1;\n\n for (int j = MaRowBegin; j < MaRowEnd; j++)\n {\n // find the diagonal entry in the current row\n if (tIdx == MaColInd[j])\n {\n diagIsZero = 0;\n\n if (!Ma_nzDiagRows[tIdx]) { MaValues[j] = perturb_mag; break;}\n\n MaValues[j] += perturb_mag;\n break;\n }\n }\n\n if (diagIsZero)\n {\n printf(\"\\nWARNING: found zero diagonal entry of Ma in row %d. Ma_nzDiagRows[%d] = %d \\n\",\n tIdx, tIdx, Ma_nzDiagRows[tIdx]);\n break;\n }\n }\n}\n\n// This kernel initializes each column of P to v_x.\ntemplate \n__global__\nvoid init_Pvalues_kernel( const IndexType *PcolOffsets, const IndexType *ProwInd,\n const IndexType PnumCols, const ValueTypeB *v_x,\n ValueType *Pvalues )\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < PnumCols; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on one column of P in CSC format (P^T in CSR format)\n const unsigned int PcolBegin = PcolOffsets[tIdx];\n const unsigned int PcolEnd = PcolOffsets[tIdx + 1];\n\n for (int j = PcolBegin; j < PcolEnd; j++)\n {\n // go thru each entry in the given column of P\n Pvalues[j] = v_x[ProwInd[j]];\n }\n }\n}\n\n\n\n////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nEM_Interpolator >::\nEM_Interpolator() : EM_InterpolatorBase(),\n m_cuds_handle(0),\n m_dense_Aijs(0),\n m_dense_invAijs(0),\n m_ipiv(0),\n m_cuds_info(0),\n m_cuds_wspace(0)\n{\n // Allocate a handle for cudense\n hipsolverStatus_t status = hipsolverDnCreate(&m_cuds_handle);\n\n if ( status != HIPSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not create the CUDENSE handle\", AMGX_ERR_CUDA_FAILURE );\n }\n\n // Define the cudense stream.\n status = hipsolverSetStream(m_cuds_handle, amgx::thrust::global_thread_handle::get_stream());\n\n if ( status != HIPSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not set the stream for CUDENSE\", AMGX_ERR_CUDA_FAILURE );\n }\n\n allocMem(m_cuds_info, sizeof(int), false);\n}\n\n\ntemplate \nEM_Interpolator >::~EM_Interpolator()\n{\n if (m_cuds_handle)\n {\n hipsolverDnDestroy(m_cuds_handle);\n }\n\n if (m_dense_Aijs)\n {\n amgx::memory::hipFreeAsync(m_dense_Aijs);\n m_dense_Aijs = 0;\n }\n\n if (m_dense_invAijs)\n {\n amgx::memory::hipFreeAsync(m_dense_invAijs);\n m_dense_invAijs = 0;\n }\n\n if (m_ipiv)\n {\n amgx::memory::hipFreeAsync(m_ipiv);\n m_ipiv = 0;\n }\n\n if (m_cuds_wspace)\n {\n amgx::memory::hipFreeAsync(m_cuds_wspace);\n m_cuds_wspace = 0;\n }\n\n if (m_cuds_info)\n {\n amgx::memory::hipFreeAsync(m_cuds_info);\n m_cuds_info = 0;\n }\n}\n\n\n\n// Private functions for the implementation on device\n\n// Find the sparse dimensions of P (i.e. number of non-zeros in each column of P)\ntemplate \nvoid EM_Interpolator >\n::computePsparsity( const Matrix_d &A, const IVector &cf_map, const IVector &coarse_idx,\n Matrix_d &P, IVector &PnnzPerCol,\n BVector &Ma_nzDiagRows )\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // Assemble all the raw pointers needed\n // Matrix A\n const IndexType *ArowOffsets_ptr = A.row_offsets.raw();\n const IndexType *AcolInd_ptr = A.col_indices.raw();\n const ValueType *Avalues_ptr = A.values.raw();\n const IndexType AnumRows = (IndexType) A.get_num_rows();\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (AnumRows / blocksize + 1));\n // temporary vectors - diagonal and non-zero offsets\n VVector Adiag(A.get_num_rows(), 0.0);\n ValueType *Adiag_ptr = Adiag.raw();\n // fullToCoarseMap array gives the COARSE index corresponding to the given FULL index\n // (meaningless if the full index corresponds to a FINE pt)\n IntVector fullToCoarseMap(AnumRows, 0);\n int *fullToCoarseMap_ptr = fullToCoarseMap.raw();\n amgx::thrust::transform_exclusive_scan( cf_map.begin(), cf_map.end(), fullToCoarseMap.begin(),\n is_coarse(), 0, amgx::thrust::plus() );\n cudaCheckError();\n // Extract the diagonal of A\n find_diag_kernel_indexed_dia <<< numBlocks, blocksize>>>(AnumRows, A.diag.raw(), Avalues_ptr, Adiag_ptr);\n cudaCheckError();\n // Coarse/Fine map\n const int *cf_map_ptr = cf_map.raw();\n // Find the sparse dimensions of P (i.e. number of non-zeros in each column of P)\n // Count the number of non-zeros in the interpolation if\n const int numCoarse = coarse_idx.size();\n const int *coarse_idx_ptr = coarse_idx.raw();\n VVector thresh_per_row(AnumRows);\n count_PnnzPerCol_kernel\n <<< numBlocks, blocksize>>>(ArowOffsets_ptr, AcolInd_ptr, Avalues_ptr,\n Adiag_ptr, numCoarse,\n cf_map_ptr, fullToCoarseMap_ptr, coarse_idx_ptr,\n PnnzPerCol.raw(), thresh_per_row.raw());\n cudaCheckError();\n printMatrixStats(A);\n cudaCheckError();\n // Column offsets of P (CSC format)\n IntVector PnzOffsets(numCoarse + 1);\n PnzOffsets[0] = 0;\n // get the offsets in P with an inclusive scan\n thrust_wrapper::inclusive_scan(PnnzPerCol.begin(), PnnzPerCol.end(), PnzOffsets.begin() + 1);\n cudaCheckError();\n // get total num of non-zeros in P\n const int Pnnz = PnzOffsets[numCoarse];\n // resize P\n P.resize(0, 0, 0, 1);\n P.addProps(CSR); // IMPORTANT: P actually stores P^T in CSR format\n // because we build P by columns (and CSC format is not supported)\n P.resize(numCoarse, AnumRows, Pnnz, 1);\n // set P offsets (P column offsets or P^T row offsets)\n amgx::thrust::copy(PnzOffsets.begin(), PnzOffsets.end(), P.row_offsets.begin());\n cudaCheckError();\n // treat P as if it were in the CSC format.\n const IndexType *PcolOffsets_ptr = P.row_offsets.raw();\n IndexType *ProwInd_ptr = P.col_indices.raw();\n bool *Ma_nzDiagRows_ptr = Ma_nzDiagRows.raw();\n // Since CSC format is not supported in AMGX matrix, we instead\n // compute P^T (which is stored in P) in the CSR format.\n // Initialize P by performing greedy aggregation based on the coarse set (cf_map)\n // The values of P remain uninitialized, since at this point we only get the sparsity pattern of P.\n init_ProwInd_greedy_aggregation_kernel\n <<< numBlocks, blocksize>>>(ArowOffsets_ptr, AcolInd_ptr, Avalues_ptr,\n Adiag_ptr, AnumRows,\n cf_map_ptr, fullToCoarseMap_ptr,\n PcolOffsets_ptr, ProwInd_ptr,\n Ma_nzDiagRows_ptr);\n cudaCheckError();\n}\n\n\n// Compute the submatrices Aij's of A and their inverses invAij's\ntemplate \nvoid EM_Interpolator >\n::computeAijSubmatrices(const Matrix_d &A, const int numCoarse, const Matrix_d &P,\n ValueType *dense_Aijs, ValueType *dense_invAijs,\n const IntVector &AijOffsets, int *ipiv,\n hipsolverHandle_t &cuds_handle, int *cuds_info)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() / blocksize + 1));\n // Assemble all the raw pointers needed\n // Matrix A\n const IndexType *ArowOffsets_ptr = A.row_offsets.raw();\n const IndexType *AcolInd_ptr = A.col_indices.raw();\n const ValueType *Avalues_ptr = A.values.raw();\n // Treat P as if it were in the CSC format.\n const IndexType *PcolOffsets_ptr = P.row_offsets.raw();\n const IndexType *ProwInd_ptr = P.col_indices.raw();\n const IndexType *AijOffsets_ptr = AijOffsets.raw();\n // Extract Aij's from A based on the sparsity pattern of each column of P.\n // Column major format is used for storage of Aij's.\n extract_dense_Aijs_col_major_kernel\n <<< numBlocks, blocksize>>>(ArowOffsets_ptr, AcolInd_ptr, Avalues_ptr,\n PcolOffsets_ptr, ProwInd_ptr, numCoarse,\n AijOffsets_ptr, dense_Aijs);\n cudaCheckError();\n // Initialize invAij's to (dense) identity matrices.\n init_dense_invAijs_col_major_kernel\n <<< numBlocks, blocksize>>>(PcolOffsets_ptr, numCoarse,\n AijOffsets_ptr, dense_invAijs);\n cudaCheckError();\n // Sequentially perform dense LU factorization on all Aij's\n // LU factors stored in place of the Aij's\n // Then, perform inversion of Aij submatrices by solving with RHS == identity matrix\n hipsolverStatus_t cudsStat;\n\n if (!cuds_handle)\n {\n hipsolverStatus_t status = hipsolverDnCreate(&m_cuds_handle);\n\n if ( status != HIPSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not create the CUDENSE handle\", AMGX_ERR_CUDA_FAILURE );\n }\n\n // Define the cudense stream.\n status = hipsolverSetStream(m_cuds_handle, amgx::thrust::global_thread_handle::get_stream());\n\n if ( status != HIPSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not set the stream for CUDENSE\", AMGX_ERR_CUDA_FAILURE );\n }\n }\n\n // Pointer to info from cudense\n if (!cuds_info) { cuds_info = (int *) malloc(sizeof(int)); }\n\n // Offset of cudense pivoting sequence for each Aij.\n int ipivOffset = 0;\n IVector_h sizes(numCoarse);\n int max_wsize = 0;\n\n for (int j = 0; j < numCoarse; j++)\n {\n int cur_wsize = 0;\n int AijOffset = AijOffsets[j];\n int AijNumRows = P.row_offsets[j + 1] - P.row_offsets[j];\n int lda = AijNumRows;\n hipsolverStatus_t status = hipsolverDnXgetrf_bufferSize(cuds_handle, AijNumRows, AijNumRows, dense_Aijs + AijOffset, lda, &cur_wsize);\n\n if (status != HIPSOLVER_STATUS_SUCCESS)\n {\n FatalError(\"failed to perform LU factorization\", AMGX_ERR_INTERNAL);\n }\n\n max_wsize = std::max(cur_wsize, max_wsize);\n }\n\n allocMem(m_cuds_wspace, max_wsize, false);\n\n for (int j = 0; j < numCoarse; j++)\n {\n int AijOffset = AijOffsets[j]; // offset of the current Aij\n int AijNumRows = P.row_offsets[j + 1] - P.row_offsets[j]; // dimension of the next Aij within dense_Aijsc\n int lda = AijNumRows; // leading dim for the next Aij (col-major)\n sizes[j] = AijNumRows;\n // [L,U,P] = lu(Aij) (LU factors stored in place of the current Aij)\n cudsStat = hipsolverDnXgetrf(cuds_handle, AijNumRows, AijNumRows,\n dense_Aijs + AijOffset, lda, m_cuds_wspace,\n ipiv + ipivOffset, cuds_info);\n\n if (cudsStat != HIPSOLVER_STATUS_SUCCESS)\n {\n FatalError(\"failed to perform LU factorization\", AMGX_ERR_INTERNAL);\n }\n\n // invAij = U\\(L\\P) (i.e. find inverse of Aij).\n cudsStat = hipsolverDnXgetrs(cuds_handle, HIPBLAS_OP_N, AijNumRows, AijNumRows,\n dense_Aijs + AijOffset, lda,\n ipiv + ipivOffset, dense_invAijs + AijOffset, lda, cuds_info );\n\n if (cudsStat != HIPSOLVER_STATUS_SUCCESS)\n {\n FatalError(\"failed to perform triangular solve\", AMGX_ERR_INTERNAL);\n }\n\n ipivOffset += AijNumRows; // starting offset of the next pivoting sequence\n }\n}\n\n\n// Compute Ma matrix\ntemplate \nvoid EM_Interpolator >\n::computeMa(Matrix_d &Ma, const int AnumRows, const int numCoarse, const Matrix_d &P,\n const ValueType *dense_invAijs, const IntVector &AijOffsets,\n const BVector &Ma_nzDiagRows,\n const bool perturb_Ma_diag, const ValueType perturb_mag)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (AnumRows / blocksize + 1));\n // Treat P as if it were in CSC format.\n const IndexType *PcolOffsets_ptr = P.row_offsets.raw();\n const IndexType *ProwInd_ptr = P.col_indices.raw();\n // Number of non-zeros in each row of Ma\n IntVector Ma_nnzPerRow(AnumRows, 0);\n // 1. Compute Ma_nnzPerRow (kernel).\n count_Ma_nnz_per_row_kernel <<< numBlocks, blocksize>>>( PcolOffsets_ptr, ProwInd_ptr, numCoarse,\n Ma_nnzPerRow.raw(), AnumRows, Ma_nzDiagRows.raw() );\n cudaCheckError();\n // 2. Resize Ma and assign MaRowOffsets by performing a scan on Ma_nnzPerRow.\n // Get the offsets in Ma with an inclusive scan (in-place)\n thrust_wrapper::inclusive_scan(Ma_nnzPerRow.begin(), Ma_nnzPerRow.end(), Ma_nnzPerRow.begin());\n cudaCheckError();\n // Get total num of non-zeros in Ma\n const int Ma_nnz = Ma_nnzPerRow[AnumRows - 1];\n // Resize Ma\n Ma.resize(0, 0, 0, 1);\n Ma.addProps(CSR);\n Ma.resize(AnumRows, AnumRows, Ma_nnz, 1);\n Ma.row_offsets[0] = 0;\n // set Ma row offsets\n amgx::thrust::copy(Ma_nnzPerRow.begin(), Ma_nnzPerRow.end(), Ma.row_offsets.begin() + 1);\n cudaCheckError();\n const IndexType *MaRowOffsets_ptr = Ma.row_offsets.raw();\n IndexType *MaColInd_ptr = Ma.col_indices.raw();\n ValueType *MaValues_ptr = Ma.values.raw();\n // 3. Get the unsorted column indices of Ma.\n get_unsorted_MaColInd_kernel <<< numBlocks, blocksize>>>(PcolOffsets_ptr, ProwInd_ptr, numCoarse,\n MaRowOffsets_ptr, MaColInd_ptr, AnumRows,\n Ma_nzDiagRows.raw());\n cudaCheckError();\n\n // 4. Sort column indices to get the final MaColInd.\n for (unsigned int MaRow = 0; MaRow < AnumRows; MaRow++)\n {\n int MaRowBegin = Ma.row_offsets[MaRow];\n int MaRowEnd = Ma.row_offsets[MaRow + 1];\n\n if (MaRowEnd - MaRowBegin > 1)\n {\n thrust_wrapper::sort( Ma.col_indices.begin() + Ma.row_offsets[MaRow],\n Ma.col_indices.begin() + Ma.row_offsets[MaRow + 1] );\n }\n }\n\n // 5. Compute MaValues by summing up the \"spattered\" inverses of Aij submatrices.\n add_invAijs_to_Ma_col_major_kernel\n <<< numBlocks, blocksize>>>(PcolOffsets_ptr, ProwInd_ptr, numCoarse,\n AijOffsets.raw(), dense_invAijs,\n MaRowOffsets_ptr, MaColInd_ptr,\n MaValues_ptr);\n cudaCheckError();\n\n // Perturb the diagonal of Ma here (prevents Ma from being singular)\n // Ma = 1.0e-8*speye(n,n);\n if (perturb_Ma_diag)\n {\n // Default: perturb_mag = 1.0e-8\n perturb_Ma_diag_kernel\n <<< numBlocks, blocksize>>>(MaRowOffsets_ptr, MaColInd_ptr,\n MaValues_ptr, AnumRows,\n perturb_mag, Ma_nzDiagRows.raw());\n cudaCheckError();\n }\n\n Ma.set_initialized(1);\n}\n\n\n// Solve for x: Ma*x=e (solved with pcg)\ntemplate \nvoid EM_Interpolator >\n::solveMa_e(Matrix_d &Ma, const int AnumRows, Vector_d &v_x)\n{\n typedef typename Matrix_d::value_type ValueType;\n // 1. Initialize x (vector of 0's) and e (vector of 1's)\n v_x.resize(AnumRows, 0.0);\n Vector_d v_e(AnumRows, 1.0);\n // ------------------------------ define pcg solver ---------------------------------\n // temporary config and pointer to main config (json format)\n AMG_Config cfg;\n std::string cfg_string = \"{ \\\"config_version\\\": 2\";\n cfg_string += \", \\\"solver\\\": { \\\"scope\\\": \\\"pcg_scope\\\", \\\"solver\\\": \\\"PCG\\\"\"; // Set the solver to PCG\n cfg_string += \" ,\\\"preconditioner\\\": \\\"NOSOLVER\\\"\"; // IMPORTANT to have this in order to allocate PCG\n cfg_string += \" } }\";\n cfg.parseParameterString(cfg_string.c_str());\n // define and allocate solver\n Solver *Ma_solver = 0;\n Ma_solver = SolverFactory::allocate(cfg, \"default\", \"solver\");\n // ------------------------------ end define pcg ------------------------------------\n const int Ma_solveIts = 20;\n const ValueType Ma_solveTol = (ValueType) 1.e-4;\n Ma_solver->setup(Ma, false);\n Ma_solver->set_max_iters(Ma_solveIts);\n Ma_solver->setTolerance(Ma_solveTol);\n // Apply pcg to solve x = Ma \\ e\n Ma_solver->solve( v_e, v_x, true );\n Vector_d v_res(AnumRows, 0.);\n Ma_solver->compute_residual( v_e, v_x, v_res );\n\n if (Ma_solver) { delete Ma_solver; Ma_solver = 0; }\n}\n\n\n// Compute the values of P (using LU factorization of Aij submatrices)\ntemplate \nvoid EM_Interpolator >\n::computePvalues( const int AnumRows, const int numCoarse, Matrix_d &P, const Vector_d &v_x,\n const ValueType *dense_Aijs, const IntVector &AijOffsets, const int *ipiv,\n hipsolverHandle_t &cuds_handle, int *cuds_info )\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n typedef typename Vector_d::value_type ValueTypeB;\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (AnumRows / blocksize + 1));\n // Treat P as if it were in CSC format.\n const IndexType *PcolOffsets_ptr = P.row_offsets.raw();\n const IndexType *ProwInd_ptr = P.col_indices.raw();\n ValueType *Pvalues_ptr = P.values.raw();\n init_Pvalues_kernel\n <<< numBlocks, blocksize>>>(PcolOffsets_ptr, ProwInd_ptr, numCoarse,\n v_x.raw(), Pvalues_ptr);\n cudaCheckError();\n hipsolverStatus_t cudsStat;\n\n if (!cuds_handle)\n {\n hipsolverStatus_t status = hipsolverDnCreate(&m_cuds_handle);\n\n if ( status != HIPSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not create the CUDENSE handle\", AMGX_ERR_CUDA_FAILURE );\n }\n\n // Define the cudense stream.\n status = hipsolverSetStream(m_cuds_handle, amgx::thrust::global_thread_handle::get_stream());\n\n if ( status != HIPSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not set the stream for CUDENSE\", AMGX_ERR_CUDA_FAILURE );\n }\n }\n\n // Pointer to info from cudense\n if (!cuds_info) { cuds_info = (int *) malloc(sizeof(int)); }\n\n // Offset of cudense pivoting sequence for each Aij.\n int ipivOffset = 0;\n\n for (int j = 0; j < numCoarse; j++)\n {\n int AijOffset = AijOffsets[j];\n int PcolOffset = P.row_offsets[j];\n // dimension of the next Aij within dense_Aijs\n int AijNumRows = P.row_offsets[j + 1] - P.row_offsets[j];\n int lda = AijNumRows; // leading dim for the next Aij (col-major)\n\n if (lda == 0) { continue; }\n\n // invAij = U\\(L\\P) (i.e. find inverse of Aij).\n cudsStat = hipsolverDnXgetrs(cuds_handle, HIPBLAS_OP_N, AijNumRows, 1,\n dense_Aijs + AijOffset, lda,\n ipiv + ipivOffset, Pvalues_ptr + PcolOffset, lda, cuds_info);\n\n if (cudsStat != HIPSOLVER_STATUS_SUCCESS)\n { FatalError(\"failed to perform triangular solve\", AMGX_ERR_INTERNAL); }\n\n ipivOffset += AijNumRows; // starting offset of the next pivoting sequence\n }\n\n if (cuds_info) { free(cuds_info); cuds_info = 0; }\n}\n\n\n\n// Compute the values of P (using inverses of Aij submatrices)\ntemplate \nvoid EM_Interpolator >\n::computePvalues( const int AnumRows, const int numCoarse, Matrix_d &P, const Vector_d &v_x,\n const ValueType *dense_invAijs, const IntVector &AijOffsets )\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n typedef typename Vector_d::value_type ValueTypeB;\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (AnumRows / blocksize + 1));\n // Treat P as if it were in CSC format.\n const IndexType *PcolOffsets_ptr = P.row_offsets.raw();\n const IndexType *ProwInd_ptr = P.col_indices.raw();\n ValueType *Pvalues_ptr = P.values.raw();\n init_Pvalues_kernel\n <<< numBlocks, blocksize>>>(PcolOffsets_ptr, ProwInd_ptr, numCoarse,\n v_x.raw(), Pvalues_ptr);\n cudaCheckError();\n ValueType *place_holder_ptr = 0;\n allocMem(place_holder_ptr, sizeof(ValueType), false);\n ValueType one = (ValueType) 1.;\n ValueType zero = (ValueType) 0.;\n\n for (int j = 0; j < numCoarse; j++)\n {\n int AijOffset = AijOffsets[j];\n int PcolOffset = P.row_offsets[j];\n // dimension of the next Aij within dense_Aijs\n int AijNumRows = P.row_offsets[j + 1] - P.row_offsets[j];\n int lda = AijNumRows; // leading dim for the next Aij (col-major)\n\n if (lda == 0) { continue; }\n\n Cublas::gemv(false, AijNumRows, AijNumRows, &one,\n dense_invAijs + AijOffset, lda,\n Pvalues_ptr + PcolOffset, 1,\n &zero, place_holder_ptr, 1);\n hipDeviceSynchronize();\n cudaCheckError();\n }\n\n if (place_holder_ptr) { amgx::memory::hipFreeAsync(place_holder_ptr); place_holder_ptr = 0; }\n}\n\n\n\n// ----------------------------\n// specialization for device\n// ----------------------------\ntemplate \nvoid EM_Interpolator >\n::generateInterpolationMatrix_1x1(const Matrix_d &A, const IntVector &cf_map,\n Matrix_d &P, void *amg)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // Number of rows in Matrix A\n const IndexType AnumRows = (IndexType) A.get_num_rows();\n // Count the number of coarse points\n const int numCoarse = (int) amgx::thrust::count(cf_map.begin(), cf_map.end(), (int)COARSE);\n cudaCheckError();\n IVector coarse_idx(numCoarse, -1);\n amgx::thrust::counting_iterator zero(0);\n amgx::thrust::counting_iterator zero_plus_nc = zero + numCoarse;\n amgx::thrust::copy_if(zero, zero_plus_nc, cf_map.begin(), coarse_idx.begin(), is_coarse());\n // ..................................................................................\n // ------------------------ Begin Compute P code ------------------------------------\n // Store the number of non-zeros in each column of the interpolation matrix P\n IntVector PnnzPerCol(numCoarse);\n BVector Ma_nzDiagRows(AnumRows, 0);\n // Find the sparse dimensions of P (i.e. number of non-zeros in each column of P)\n computePsparsity(A, cf_map, coarse_idx, P, PnnzPerCol, Ma_nzDiagRows);\n // Get total num of non-zeros in P\n int Pnnz = P.row_offsets[numCoarse];\n // ..................................................................................\n // Now follows Jacobi smoothing (skip for now since it might be unnecessary)\n // ..................................................................................\n // -------------------------- End Compute P code ------------------------------------\n // ..................................................................................\n // ..................................................................................\n // ----------------------- Begin Energy Minimization Code ---------------------------\n // Given:\n // A: the problem matrix\n // P: sparsity pattern for P (n x nc)\n // Main loop (Ma here is the inverse of Ma in the paper)\n // I. Build Ma.\n // 0. Allocate space for storing LU factors for each Aij submatrix of A (use dense format)\n // There are numCoarse of LU factors.\n // Since the number of non-zeros in each column of P determines the size of each Aij,\n // we know the dimensions of all Aij based on the row offsets of P^T.\n // Calculate offset of each Aij matrix in the dense_Aijs array\n IntVector AijOffsets(numCoarse + 1);\n AijOffsets[0] = 0;\n // Get the offsets with an inclusive scan (nnz of each column of P is the dimension of each Aij)\n amgx::thrust::transform_inclusive_scan( PnnzPerCol.begin(), PnnzPerCol.end(), AijOffsets.begin() + 1,\n square(), amgx::thrust::plus() );\n cudaCheckError();\n PnnzPerCol.clear();\n // Cumulative num of entries in all (dense) Aijs (last entry in the offsets)\n const int Aijs_size = AijOffsets[numCoarse];\n // Allocate memory for the dense Aij submatrices and initialize to zero.\n allocMem(m_dense_Aijs, Aijs_size, true);\n // Allocate memory for the dense invAij's and initialize to zero.\n allocMem(m_dense_invAijs, Aijs_size, true);\n // Allocate memory for cudense pivoting sequence.\n allocMem(m_ipiv, Pnnz, false);\n // Compute the submatrices Aij's of A and the inverses invAij's\n computeAijSubmatrices(A, numCoarse, P, m_dense_Aijs, m_dense_invAijs,\n AijOffsets, m_ipiv, m_cuds_handle, m_cuds_info);\n // Ma matrix build sequence.\n // 1. Compute Ma_nnzPerRow (kernel).\n // 2. Resize Ma and assign MaRowOffsets by performing a scan on Ma_nnzPerRow (thrust).\n // 3. Get the unsorted column indices of Ma (kernel).\n // 4. Sort column indices to get the final MaColInd (thrust).\n // 5. Compute MaValues by summing up the \"spattered\" inverses of Aij submatrices (kernel).\n // Declare Ma\n Matrix_d Ma;\n // Compute Ma matrix\n computeMa(Ma, AnumRows, numCoarse, P, m_dense_invAijs, AijOffsets, Ma_nzDiagRows);\n\n if (m_dense_invAijs) { amgx::memory::hipFreeAsync(m_dense_invAijs); m_dense_invAijs = 0; }\n\n /* At this point, we are done with the following Matlab code analogue.\n for (int j=1; j<=nc; j++)\n {\n // Find the dofs in the j-th block\n ij = find(p(:,j)); // col vec of subindices\n ni = length(ij);\n It = sparse(n,ni);\n It(ij,:) = speye(ni);\n Aij{j} = inv(A(ij,ij));\n // Form Ma\n // A(ii,ii): submatrix of A corresponding to the dofs in the i-th block\n Ma = Ma + It*(Aij{j}*It_T);\n }\n */\n Ma_nzDiagRows.clear();\n // II. Solve for x: Ma*x=e\n //m_Mae_norm = (ValueType) 0.;\n Vector_d v_x;\n solveMa_e(Ma, AnumRows, v_x);\n hipDeviceSynchronize();\n cudaCheckError();\n Ma.set_initialized(0);\n Ma.resize(0, 0, 0, 1);\n /* Now, we are done with the following Matlab code analogue.\n e = ones(n,1);\n Mae = zeros(size(e)); // Ma\\e;\n [Mae,flag,relres,iter,resvecCG] = pcg(Ma,e,1e-4,100);\n */\n // III. Construct the columns of P in parallel.\n // P is already initialized (and at this point stores P^T in CSR format)\n // We will be performing the following operation to build each column of P\n // P(:,j) = It{j}*(invAij{j}*(It^T*x));\n // This is done with the following steps.\n // First, apply It{j}^T mask to x to get x{j} (project x onto the dimension of Aij{j}).\n // Then solve Aij\\x{j}\n // Compute the values of P\n computePvalues( AnumRows, numCoarse, P, v_x, m_dense_Aijs, AijOffsets, m_ipiv, m_cuds_handle, m_cuds_info );\n\n if (m_dense_Aijs) { amgx::memory::hipFreeAsync(m_dense_Aijs); m_dense_Aijs = 0; }\n\n if (m_ipiv) { amgx::memory::hipFreeAsync(m_ipiv); m_ipiv = 0; }\n\n if (m_cuds_info) { amgx::memory::hipFreeAsync(m_cuds_info); m_cuds_info = 0; }\n\n //if (m_dense_invAijs) { amgx::memory::cudaFreeAsync(m_dense_invAijs); m_dense_invAijs = 0; }\n /* Now, we are done building columns of P (stored as rows of P^T in CSR).\n * Here's the equivalent matlab code.\n P = zeros(n,nc);\n for (int j=1; j<=nc; j++)\n {\n // Find the dofs in the j-th block\n ij=find(p(:,j));\n ni = length(ij);\n It = zeros(n,ni);\n It(ij,:) = eye(ni);\n\n // Form j-th column of P\n P(:,j) = It*(Aij{j}*(It_T*Mae));\n }\n */\n //TODO: truncate rows/cols of P to keep the sparsity low on all levels.\n // ------------------------- End Energy Minimization Code ---------------------------\n // ..................................................................................\n}\n\n\n\ntemplate< class T_Config>\nvoid EM_InterpolatorBase\n::generateInterpolationMatrix(const Matrix &A, const IntVector &cf_map,\n Matrix &P, void *amg)\n{\n P.set_initialized(0);\n\n if (A.get_block_size() == 1)\n {\n generateInterpolationMatrix_1x1(A, cf_map, P, amg);\n }\n else\n {\n FatalError(\"Unsupported block size for distance1 interpolator\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n P.computeDiagonal();\n P.set_initialized(1);\n}\n\n#define AMGX_CASE_LINE(CASE) template class EM_InterpolatorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class EM_Interpolator::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace energymin\n} // namespace amgx\n\n", "source": "/kaggle/working/amgx/src/energymin/interpolators/em.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace amgx\n{\n\n/***************************************\n * Source Definitions\n ***************************************/\n\nstruct is_less_than_zero\n{\n __host__ __device__\n bool operator()(int x) const\n {\n return x < 0;\n }\n};\n\n\nstruct is_non_neg\n{\n __host__ __device__\n bool operator()(const int &x)\n {\n return x >= 0;\n }\n};\n\ntemplate\n__global__\nvoid flag_boundary_nodes_P(const IndexType *A_rows, const IndexType *A_cols, int *boundary_flags, int num_owned_coarse_pts, int A_num_rows, int* nnz_interior)\n{\n const int nWarps = cta_size / 32;\n // The coordinates of the thread inside the CTA/warp.\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n int nnz_int = 0;\n\n // Loop over rows of A.\n for ( int aRowId = blockIdx.x * nWarps + warpId ; aRowId < A_num_rows ; aRowId += gridDim.x * nWarps )\n {\n // Load A row IDs.\n int aColIt = A_rows[aRowId ];\n int aColEnd = A_rows[aRowId + 1];\n int nnz = aColEnd - aColIt;\n bool halo_found = 0;\n\n // Iterate over the columns of A.\n for ( aColIt += laneId ; utils::any( aColIt < aColEnd ) && !halo_found; aColIt += 32 )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int aColId = aColIt < aColEnd ? A_cols[aColIt] : -1;\n bool is_halo = (aColIt < aColEnd) && (aColId >= num_owned_coarse_pts);\n\n // Push coarse and strongly connected nodes in the set.\n if (utils::any(is_halo))\n {\n halo_found = true;\n break;\n }\n }\n\n if (laneId == 0)\n {\n if(halo_found)\n {\n boundary_flags[aRowId] = 1;\n }\n else\n {\n nnz_int += nnz;\n }\n }\n }\n\n if(laneId == 0)\n {\n atomicAdd(nnz_interior, nnz_int);\n }\n}\n\n\n__global__\nvoid flag_boundary_nodes(int *b2l_maps, int *boundary_flags, int size )\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n boundary_flags[b2l_maps[tidx]] = 1;\n }\n}\n\n__global__\nvoid flag_interior_nodes(int *boundary_flags, int *interior_flags, int size )\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n interior_flags[tidx] = !boundary_flags[tidx];\n }\n}\n\n__global__\nvoid fill_rows_lists(int *__restrict__ interior_rows_lists, int *__restrict__ boundary_rows_lists, const int *__restrict__ interior_flags, const int *__restrict__ boundary_flags, int size )\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n int bdy_flag = boundary_flags[tidx];\n\n if (bdy_flag != boundary_flags[tidx + 1]) { boundary_rows_lists[bdy_flag] = tidx; }\n\n int int_flag = interior_flags[tidx];\n\n if (int_flag != interior_flags[tidx + 1]) { interior_rows_lists[int_flag] = tidx; }\n }\n}\n\n\ntemplate\n__global__\nvoid fill_B_B2L_map(const IndexType *maps_offsets, const IndexType *A_b2l_map, const IndexType *row_offsets, IndexType *C_hat_b2l_map, int size, int rank)\n{\n const int nWarps = cta_size / 32;\n // The coordinates of the thread inside the CTA/warp.\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n // Loop over rows of A.\n for ( int aRowId = blockIdx.x * nWarps + warpId ; aRowId < size ; aRowId += gridDim.x * nWarps )\n {\n // Load A row IDs.\n int aColIt = maps_offsets[aRowId ];\n int size = maps_offsets[aRowId + 1] - aColIt;\n int start = row_offsets[A_b2l_map[aRowId]];\n int end = start + size;\n\n // Iterate over the columns of A.\n for ( start += laneId, aColIt += laneId ; utils::any( start < end ) ; start += 32, aColIt += 32 )\n {\n if (start < end)\n {\n C_hat_b2l_map[aColIt] = start;\n }\n }\n }\n}\n\ntemplate \n__global__ void fill_halo_rows_row_ids(int64_t *halo_rows_row_ids, INDEX_TYPE *halo_rows_local_row_ids, int64_t base_index, INDEX_TYPE num_owned_coarse_pts, INDEX_TYPE size)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n int64_t gid = halo_rows_row_ids[tidx];\n halo_rows_local_row_ids[tidx] = (INDEX_TYPE) (gid - base_index);\n\n if (gid < base_index || gid > base_index + num_owned_coarse_pts)\n {\n printf(\"Should not be here in fill_halo_rows_row_ids\\n\");\n }\n }\n}\n\n__global__\nvoid fill_local_to_global_map(int *bdy_list, int64_t *local_to_global, int64_t offset, int size )\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n local_to_global[tidx] = bdy_list[tidx] + offset;\n }\n}\n\ntemplate \n__global__ void remove_fine_halo(int64_t *cf_map_global, INDEX_TYPE *flag, INDEX_TYPE size)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n flag[tidx] = (cf_map_global[tidx] == -1) ? 0 : 1;\n }\n}\n\ntemplate \n__global__ void fill_P_local_to_global(int64_t *P_local_to_global, int64_t *cf_map_global, INDEX_TYPE *flag, INDEX_TYPE size)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n if (cf_map_global[tidx] != -1)\n {\n P_local_to_global[flag[tidx]] = cf_map_global[tidx];\n }\n }\n}\n\ntemplate \n__global__ void calc_num_neighbors(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *part_offsets, INDEX_TYPE *exists, INDEX_TYPE num_part, INDEX_TYPE my_id, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n if (col < part_offsets[my_id] || col >= part_offsets[my_id + 1])\n {\n int part = 0;\n\n while (part < num_part && (col < part_offsets[part] || col >= part_offsets[part + 1])) { part++; }\n\n if (part < num_part && (col >= part_offsets[part] && col < part_offsets[part + 1]))\n {\n exists[part] = 1;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n// this is assuming local indices (used in matrix A)\ntemplate \n__global__ void calc_num_neighbors_v2_local(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *part_offsets, INDEX_TYPE *exists, int64_t *local_to_global_map, INDEX_TYPE num_part, INDEX_TYPE my_id, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n // Check if column point to halo node\n if (col >= num_rows)\n {\n int64_t gcol = local_to_global_map[col - num_rows];\n int part = 0;\n\n while (part < num_part && (gcol < part_offsets[part] || gcol >= part_offsets[part + 1]))\n {\n part++;\n }\n\n if (part < num_part && (gcol >= part_offsets[part] && gcol < part_offsets[part + 1]))\n {\n exists[part] = 1;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n// this is assuming global indices (used in halo rows)\ntemplate \n__global__ void calc_num_neighbors_v2_global(INDEX_TYPE *row_offsets, int64_t *col_indices, int64_t *part_offsets, INDEX_TYPE *exists, INDEX_TYPE num_part, INDEX_TYPE my_id, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n int64_t start = part_offsets[my_id];\n int64_t end = part_offsets[my_id + 1];\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n int64_t col = col_indices[i];\n\n // Check if column point to halo node\n if (col < start || col >= end)\n {\n int part = 0;\n\n while (part < num_part && (col < part_offsets[part] || col >= part_offsets[part + 1]))\n {\n part++;\n }\n\n if (part < num_part && (col >= part_offsets[part] && col < part_offsets[part + 1]))\n {\n exists[part] = 1;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void flag_halo_nodes(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *part_ranges, INDEX_TYPE *flags, INDEX_TYPE *flag_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_part, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n if (col < base || col >= base + range)\n {\n int part = 0;\n\n while (part < num_part && (col < part_ranges[2 * part] || col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_part && (col >= part_ranges[2 * part] && col < part_ranges[2 * part + 1]))\n {\n flags[flag_offsets[part] + col - part_ranges[2 * part]] = 1;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void flag_halo_nodes_local(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *part_ranges, INDEX_TYPE *flags, INDEX_TYPE *flag_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_part, INDEX_TYPE num_rows, int64_t *local_to_global)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n if (col >= num_rows)\n {\n int64_t global_col = local_to_global[col - num_rows];\n int part = 0;\n\n while (part < num_part && (global_col < part_ranges[2 * part] || global_col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_part && (global_col >= part_ranges[2 * part] && global_col < part_ranges[2 * part + 1]))\n {\n flags[flag_offsets[part] + global_col - part_ranges[2 * part]] = set_val > 0 ? set_val : -col - 1; // if -1, then set -(local index)-1 to restore later\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n__global__ void flag_halo_nodes_from_local_to_global_map(int64_t *part_ranges, INDEX_TYPE *flags, INDEX_TYPE *flag_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_part, INDEX_TYPE size, int64_t *local_to_global)\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (tid < size)\n {\n int64_t global_col = local_to_global[tid];\n int part = 0;\n\n while (part < num_part && (global_col < part_ranges[2 * part] || global_col >= part_ranges[2 * part + 1]))\n {\n part++;\n }\n\n if (part < num_part && (global_col >= part_ranges[2 * part] && global_col < part_ranges[2 * part + 1]))\n {\n int col_id = range + tid;\n flags[flag_offsets[part] + global_col - part_ranges[2 * part]] = -1 - col_id;\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void flag_halo_nodes_local_v2(int64_t *part_ranges, INDEX_TYPE *flags, INDEX_TYPE *flag_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_part, INDEX_TYPE local_to_global_size, int64_t *local_to_global)\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (tid < local_to_global_size)\n {\n int64_t global_col = local_to_global[tid];\n int col = range + tid;\n // Identify the partition that owns that node\n int part = 0;\n\n while (part < num_part && (global_col < part_ranges[2 * part] || global_col >= part_ranges[2 * part + 1])) { part++; }\n\n // Flag the corresponding node in flags array\n if (part < num_part && (global_col >= part_ranges[2 * part] && global_col < part_ranges[2 * part + 1]))\n {\n flags[flag_offsets[part] + global_col - part_ranges[2 * part]] = -col - 1;\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate \n__global__ void flag_halo_nodes_global_v2(INDEX_TYPE *row_offsets, int64_t *global_col_indices, int64_t *part_ranges, INDEX_TYPE *flags, INDEX_TYPE *flag_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_part, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n int64_t global_col = global_col_indices[i];\n\n if (global_col < base || global_col >= base + (int64_t) range)\n {\n int part = 0;\n\n while (part < num_part && (global_col < part_ranges[2 * part] || global_col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_part && (global_col >= part_ranges[2 * part] && global_col < part_ranges[2 * part + 1]))\n {\n flags[flag_offsets[part] + (int) (global_col - part_ranges[2 * part])] = 1;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n\n\n\ntemplate \n__global__ void flag_halo_nodes_global(INDEX_TYPE *row_offsets, INDEX_TYPE num_rows, int64_t *global_col_indices, int64_t *part_ranges, INDEX_TYPE *flags, INDEX_TYPE *flag_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_neighbors, INDEX_TYPE *local_col_indices)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n int64_t col = global_col_indices[i];\n\n if (col < base || col >= base + range)\n {\n // exterior node\n int part = 0;\n\n while (part < num_neighbors && (col < part_ranges[2 * part] || col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_neighbors && (col >= part_ranges[2 * part] && col < part_ranges[2 * part + 1]))\n {\n // check if the flag is already set (i.e. 1-ring node) then skip it\n int flag = flags[flag_offsets[part] + col - part_ranges[2 * part]];\n\n if (flag == 0)\n {\n flags[flag_offsets[part] + col - part_ranges[2 * part]] = set_val;\n }\n else if (flag < 0) // it is possible that flag is already set to 1 in this kernel, so if it's negative then it's the 1st ring\n {\n // update local index, note that local to global mapping is unchanged for 1st ring\n local_col_indices[i] = -flags[flag_offsets[part] + col - part_ranges[2 * part]] - 1;\n }\n }\n }\n else\n {\n // interior node\n local_col_indices[i] = col - base;\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void renumber_col_indices_global(INDEX_TYPE *row_offsets, int *col_indices, int64_t *global_col_indices, int64_t *part_ranges, INDEX_TYPE *flags, INDEX_TYPE *flag_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_neighbors, INDEX_TYPE num_rows, int rank)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n int64_t col = global_col_indices[i];\n\n if (col < base || col >= base + range)\n {\n // Need to get his local index from the flag array\n int part = 0;\n\n while ( part < num_neighbors && (col < part_ranges[2 * part] || col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_neighbors && (col >= part_ranges[2 * part] && col < part_ranges[2 * part + 1]))\n {\n int pos_in_flag = flag_offsets[part] + col - part_ranges[2 * part];\n //int local_col = flags[pos_in_flag];\n col_indices[i] = range + flags[pos_in_flag];\n //bool flag = (local_col != flags[pos_in_flag+1]);\n\n if (flags[pos_in_flag] == flags[pos_in_flag + 1])\n {\n printf(\"should not be here in renumber_col_indices_global\\n\");\n }\n\n //if (flag)\n }\n }\n else\n {\n // interior node\n col_indices[i] = col - base;\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n\n\ntemplate \n__global__ void calc_new_halo_mapping(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *part_ranges, INDEX_TYPE num_part, INDEX_TYPE num_rows, int64_t *local_to_global, int64_t *new_local_to_global, INDEX_TYPE *halo_indices, int halo_base, int rank)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n if (col >= num_rows)\n {\n int64_t glob_col = local_to_global[col - num_rows];\n int part = 0;\n\n while (part < num_part && (glob_col < part_ranges[2 * part] || glob_col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_part && (glob_col >= part_ranges[2 * part] && glob_col < part_ranges[2 * part + 1]))\n {\n int pos = glob_col - part_ranges[2 * part];\n int actual_col = halo_base + halo_indices[pos];\n col_indices[i] = -1 - actual_col;\n new_local_to_global[actual_col - num_rows] = glob_col;\n // when we renumber we also store mapping to use later in residual restriction since matrix R has other order of halo indices\n // store both direct and inverse renumbering for P & R matrices\n //halo_P_renumbering[col - num_rows] = actual_col;\n //halo_R_renumbering[actual_col - num_rows] = col;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n\ntemplate \n__global__ void calc_new_halo_mapping_v3(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *part_ranges, INDEX_TYPE num_part, INDEX_TYPE num_rows, int64_t *local_to_global, int64_t *new_local_to_global, INDEX_TYPE *halo_indices, int halo_base, int rank)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n if (col >= num_rows)\n {\n int64_t glob_col = local_to_global[col - num_rows];\n int part = 0;\n\n while (part < num_part && (glob_col < part_ranges[2 * part] || glob_col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_part && (glob_col >= part_ranges[2 * part] && glob_col < part_ranges[2 * part + 1]))\n {\n int pos = glob_col - part_ranges[2 * part];\n int actual_col = halo_base + halo_indices[pos];\n col_indices[i] = -1 - actual_col;\n new_local_to_global[actual_col - num_rows] = glob_col;\n // when we renumber we also store mapping to use later in residual restriction since matrix R has other order of halo indices\n // store both direct and inverse renumbering for P & R matrices\n //halo_P_renumbering[col - num_rows] = actual_col;\n //halo_R_renumbering[actual_col - num_rows] = col;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void flip_sign_col(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n if (col < 0)\n {\n col_indices[i] = -col_indices[i] - 1;\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void calc_new_halo_mapping_ring2(INDEX_TYPE *row_offsets, INDEX_TYPE num_rows, int64_t *global_col_indices, int64_t *part_ranges, int64_t base, INDEX_TYPE range, INDEX_TYPE num_neighbors, INDEX_TYPE *halo_offsets, INDEX_TYPE *neighbor_offsets, INDEX_TYPE *neighbor_nodes, int64_t *local_to_global, INDEX_TYPE *local_col_indices)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n int64_t col = global_col_indices[i];\n\n if ((col < base || col >= base + range) && local_col_indices[i] == -1)\n {\n // update only 2nd ring halo indices here, 1st ring is updated already\n int part = 0;\n\n while (part < num_neighbors && (col < part_ranges[2 * part] || col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_neighbors && (col >= part_ranges[2 * part] && col < part_ranges[2 * part + 1]))\n {\n int pos = col - part_ranges[2 * part];\n local_col_indices[i] = halo_offsets[part] + neighbor_nodes[neighbor_offsets[part] + pos];\n local_to_global[local_col_indices[i] - range] = col;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void renumber_col_indices(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *part_ranges, INDEX_TYPE *halo_renum, INDEX_TYPE *halo_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_part, INDEX_TYPE num_rows, int64_t *local_to_global_map)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n if (col < base || col >= base + range)\n {\n int part = 0;\n\n while (part < num_part && (col < part_ranges[2 * part] || col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_part && (col >= part_ranges[2 * part] && col < part_ranges[2 * part + 1]))\n {\n int local_index = halo_renum[halo_offsets[part] + col - part_ranges[2 * part]];\n local_to_global_map[local_index - range] = col_indices[i];\n col_indices[i] = local_index;\n }\n }\n else\n {\n col_indices[i] = col - base;\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n__global__ void populate_B2L(INDEX_TYPE *indexing, INDEX_TYPE *output, INDEX_TYPE last, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n if (row == size - 1)\n {\n if (last) { output[indexing[row]] = row; }\n }\n else if (indexing[row] != indexing[row + 1])\n {\n output[indexing[row]] = row;\n }\n\n row += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void populate_bdy_list_and_l2h_map(INDEX_TYPE *l2h_map, INDEX_TYPE *halo_nodes, INDEX_TYPE *halo_indices, INDEX_TYPE *bdy_list, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n int scanned_id = halo_indices[row];\n\n if (scanned_id != halo_indices[row + 1])\n {\n bdy_list[scanned_id] = row;\n l2h_map[scanned_id] = -halo_nodes[row] - 1;\n }\n\n row += gridDim.x * blockDim.x;\n }\n}\n\n\n__global__ void populate_B2L_v2(INDEX_TYPE *indexing, INDEX_TYPE *output, INDEX_TYPE offset, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n if (indexing[row] != indexing[row + 1])\n {\n output[indexing[row] - offset] = row;\n }\n\n row += gridDim.x * blockDim.x;\n }\n}\n\n\n\ntemplate \n__global__ void export_matrix_elements(INDEX_TYPE *row_offsets, T *values, INDEX_TYPE bsize, INDEX_TYPE *maps, INDEX_TYPE *pointers, T *output, INDEX_TYPE *col_indices, INDEX_TYPE *output2, INDEX_TYPE size)\n{\n int idx = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (idx < size)\n {\n int row = maps[idx];\n INDEX_TYPE src_base = row_offsets[row];\n INDEX_TYPE dst_base = pointers[idx];\n\n for (int m = coopIdx; m < row_offsets[row + 1]*bsize - src_base * bsize; m += coop)\n {\n output[dst_base * bsize + m] = values[src_base * bsize + m];\n }\n\n for (int m = coopIdx; m < row_offsets[row + 1] - src_base; m += coop)\n {\n output2[dst_base + m] = col_indices[src_base + m];\n }\n\n idx += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void export_matrix_elements_global(INDEX_TYPE *row_offsets, T *values, INDEX_TYPE bsize, INDEX_TYPE *maps, INDEX_TYPE *pointers, T *output, INDEX_TYPE *col_indices, int64_t *output2, INDEX_TYPE size, int64_t *local_to_global, INDEX_TYPE *q, INDEX_TYPE num_owned_pts, int64_t base_index)\n{\n int idx = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (idx < size)\n {\n int row = maps[idx];\n\n if (q != NULL)\n {\n row = q[row];\n }\n\n INDEX_TYPE src_base = row_offsets[row];\n INDEX_TYPE dst_base = pointers[idx];\n\n for (int m = coopIdx; m < row_offsets[row + 1]*bsize - src_base * bsize; m += coop)\n {\n output[dst_base * bsize + m] = values[src_base * bsize + m];\n }\n\n for (int m = coopIdx; m < row_offsets[row + 1] - src_base; m += coop)\n {\n int col = col_indices[src_base + m];\n\n if (col < num_owned_pts)\n {\n output2[dst_base + m] = (int64_t) col_indices[src_base + m] + base_index;\n }\n else\n {\n output2[dst_base + m] = local_to_global[col_indices[src_base + m] - num_owned_pts];\n }\n }\n\n idx += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void export_matrix_diagonal(T *values, INDEX_TYPE bsize, INDEX_TYPE *maps, T *output, INDEX_TYPE size)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (idx < size)\n {\n int row = maps[idx];\n INDEX_TYPE src_base = row;\n INDEX_TYPE dst_base = idx;\n\n for (int m = 0; m < bsize; m++)\n {\n output[dst_base * bsize + m] = values[src_base * bsize + m];\n }\n\n idx += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void write_matrix_rowsize(INDEX_TYPE *maps, INDEX_TYPE *row_offsets, INDEX_TYPE size, INDEX_TYPE *output)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (idx < size)\n {\n int row = maps[idx];\n output[idx] = row_offsets[row + 1] - row_offsets[row];\n idx += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void find_next_ring(INDEX_TYPE *maps, INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, INDEX_TYPE size, INDEX_TYPE base, INDEX_TYPE range, INDEX_TYPE *halo_rings)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (idx < size)\n {\n int row = maps[idx];\n\n for (int col = row_offsets[row]; col < row_offsets[row + 1]; col++)\n {\n int colIdx = col_indices[col];\n\n if (colIdx >= base && colIdx < base + range)\n {\n halo_rings[colIdx - base] = 1;\n }\n }\n\n idx += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void remove_previous_rings(INDEX_TYPE *map, INDEX_TYPE *halo_rings, INDEX_TYPE size)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (idx < size)\n {\n int row = map[idx];\n halo_rings[row] = 0;\n idx += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void renumber_halo_lists(INDEX_TYPE *halo_list, INDEX_TYPE *halo_renum, int64_t base, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n halo_list[row] = halo_renum[halo_list[row] - base];\n row += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void write_ring_rows(INDEX_TYPE *indices, INDEX_TYPE *maps, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n int idx = indices[row];\n\n if (idx == 1 && row == 0)\n {\n maps[idx - 1] = row;\n }\n else if (row > 0 && idx != indices[row - 1])\n {\n maps[idx - 1] = row;\n }\n\n row += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \nvoid DistributedArrangerBase::set_part_offsets(const INDEX_TYPE num_part, const VecInt_t *part_offsets_h)\n{\n part_offsets.resize(num_part + 1);\n amgx::thrust::copy(&part_offsets_h[0], &part_offsets_h[num_part + 1], part_offsets.begin());\n cudaCheckError();\n this->num_part = num_part;\n}\n\n// TODO Make distributed reader use 64bit ints for offsets\ntemplate \nvoid DistributedArrangerBase::set_part_offsets(const INDEX_TYPE num_part, const int64_t *part_offsets_h)\n{\n part_offsets.resize(num_part + 1);\n amgx::thrust::copy(&part_offsets_h[0], &part_offsets_h[num_part + 1], part_offsets.begin());\n cudaCheckError();\n this->num_part = num_part;\n}\n\ntemplate \nvoid DistributedArranger >::create_B2L(Matrix_d &A,\n INDEX_TYPE my_id, INDEX_TYPE rings,\n int64_t &base_index, INDEX_TYPE &index_range,\n Vector &neighbors,\n I64Vector_d &halo_ranges,\n std::vector &B2L_maps,\n std::vector &L2H_maps,\n std::vector > &B2L_rings,\n DistributedComms **comms,\n std::vector **halo_rows,\n std::vector **halo_btl)\n{\n I64Vector_h halo_ranges_h;\n // initialize base_index, index_range, neighbors, halo_ranges_h\n this->create_neighbors(A, my_id, base_index, index_range, neighbors, halo_ranges_h);\n this->create_B2L_from_neighbors(A, my_id, rings, base_index, index_range, neighbors,\n halo_ranges_h, halo_ranges, B2L_maps, L2H_maps, B2L_rings, comms, halo_rows, halo_btl);\n}\n\ntemplate \nvoid DistributedArranger >::create_B2L_from_maps(Matrix_d &A,\n INDEX_TYPE my_id, INDEX_TYPE rings,\n int64_t base_index, INDEX_TYPE index_range,\n Vector &neighbors,\n I64Vector_d &halo_ranges,\n std::vector &B2L_maps,\n std::vector > &B2L_rings,\n DistributedComms **comms,\n std::vector **halo_rows,\n std::vector **halo_btl)\n{\n int num_neighbors = neighbors.size();\n this->halo_coloring = (*comms)->halo_coloring;\n std::vector halo_lists;\n halo_lists.resize(num_neighbors);\n // initialize halo_btl\n this->create_halo_btl(A, rings, num_neighbors, neighbors, base_index, index_range, my_id, B2L_maps, B2L_rings, halo_btl, halo_lists);\n // initialize halo_rows\n this->create_halo_rows(A, rings, num_neighbors, B2L_maps, B2L_rings, halo_rows, halo_lists);\n\n //this->create_halo_matrices(A, 1, base_index, index_range, my_id, neighbors, B2L_maps, B2L_rings, &(*halo_rows), &(*halo_btl));\n if (this->halo_coloring != LAST)\n {\n Matrix_d tmpA(0, 0, 0, A.getProps());\n tmpA.set_block_dimx(A.get_block_dimx());\n tmpA.set_block_dimy(A.get_block_dimy());\n tmpA.setResources(A.getResources());\n DistributedManager tmp_manager(my_id, base_index, index_range, tmpA, neighbors, halo_ranges, B2L_maps, B2L_rings, comms, NULL, NULL);\n tmpA.setManager(tmp_manager);\n tmpA.manager->getComms()->exchange_matrix_halo(**halo_rows, **halo_btl, tmpA);\n }\n}\n\ntemplate \nvoid DistributedArranger >::create_B2L_from_maps(Matrix_d &A,\n INDEX_TYPE my_id, INDEX_TYPE rings,\n Vector &neighbors,\n std::vector &B2L_maps,\n std::vector &L2H_maps,\n std::vector > &B2L_rings,\n DistributedComms **comms,\n std::vector **halo_rows,\n std::vector **halo_btl)\n{\n this->halo_coloring = (*comms)->halo_coloring;\n int num_neighbors = neighbors.size();\n int64_t base_index = 0;\n int index_range = A.get_num_rows();\n std::vector halo_lists;\n halo_lists.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n halo_lists[i] = &L2H_maps[i];\n }\n\n /*\n EXAMPLE\n Say we want 2 rings of halo. The output of these function for the example partition 1 that at this point looks like:\n A.num_rows = 4; A.num_nz = 20\n A.row_offsets = [0 4 11 15 20]\n A.col_indices = [4 0 1 2\n 4 5 0 1 2 3 7\n 0 1 2 3\n 1 2 3 6 7]\n num_neighbors=2; neighbors = [0 2]\n B2L_rings[[0 2][0 2]] B2L_maps [[0 1][1 3]] (Converted from the following input: btl_sizes = [0 2 4] btl_maps = [0 1| 1 3])\n L2H_maps (and halo_lists) [[4 5][6 7]]\n halo_rows = NULL, halo_btl=NULL\n (global indices mapping to local indices: 4-0 5-1 6-2 7-3 2-4 3-5 9-6 10-7)\n */\n // initialize ring>1 B2L_maps and B2L_rings\n this->create_rings(A, rings, num_neighbors, B2L_maps, B2L_rings);\n /*\n EXAMPLE\n having mapped out two rings of halo (already had one), B2L_maps and B2L_rings look like:\n B2L_rings[[0 2 4][0 2 4]] B2L_maps[[0 1| 2 3][1 3| 0 2]]\n */\n // initialize halo_btl, creating DistributedManager objects for each halo row matrix to be sent to each neighbor, storing B2L and L2H maps\n this->create_halo_btl(A, rings, num_neighbors, neighbors, base_index, index_range, my_id, B2L_maps, B2L_rings, halo_btl, halo_lists);\n /*\n EXAMPLE\n We created to two DistributedManager objects for the two neighbors. if halo_coloring != LAST, then the relevant fields are:\n halo_btl[0] (for neighbor ID 0)\n global_id = 1; base_index=0; index_range=6 B2L_rings[0] = [0 2 4] B2L_maps[0] = [0 1| 2 3] L2H_maps = [4 5]\n halo_btl[1] (for neighbor ID 2)\n global_id = 1; base_index=0; index_range=8 B2L_rings[0] = [0 2 4] B2L_maps[0] = [1 3| 0 2] L2H_maps = [6 7]\n */\n // Extract halo rows, ring by ring\n this->create_halo_rows(A, rings, num_neighbors, B2L_maps, B2L_rings, halo_rows, halo_lists);\n\n /*\n EXAMPLE\n Extracted two rings of rows for the two neighbors, the two halo row matrices look like:\n halo_rows[0] (for neighbor ID 0)\n halo_rows[0].row_offsets = [0 4 11 15 20]\n halo_rows[0].col_indices = [4 0 1 2\n 4 5 0 1 2 3 7\n 0 1 2 3\n 1 2 3 6 7]\n\n halo_rows[1] (for neighbor ID 2)\n halo_rows[1].row_offsets = [0 7 12 16 20]\n halo_rows[1].col_indices = [4 5 0 1 2 3 7\n 1 2 3 6 7\n 4 0 1 2\n 0 1 2 3]\n Note, that the ordering of rows corresponds to the list of indices in B2L_maps[i] (and halo_btl[i].B2L_maps[0])\n */\n //this->create_halo_matrices(A, 1, base_index, index_range, my_id, neighbors, B2L_maps, B2L_rings, &(*halo_rows), &(*halo_btl));\n if (this->halo_coloring != LAST) //For multiple rings, we still need 0-num_rings-1 rows, just not the last one.\n {\n //Set up dummy matrix and dummy manager, and exchange the halo rows\n Matrix_d tmpA(0, 0, 0, A.getProps());\n tmpA.set_block_dimx(A.get_block_dimx());\n tmpA.set_block_dimy(A.get_block_dimy());\n tmpA.setResources(A.getResources());\n I64Vector_d tmp_halo_ranges;\n DistributedManager tmp_manager(my_id, base_index, index_range, tmpA, neighbors, tmp_halo_ranges, B2L_maps, B2L_rings, comms, NULL, NULL);\n tmpA.setManager(tmp_manager);\n tmpA.manager->getComms()->exchange_matrix_halo(**halo_rows, **halo_btl, tmpA);\n }\n\n /*\n EXAMPLE\n The partitions exchange their halo_rows and halo_btl fields, partition 1 ends up with:\n halo_btl[0] (received from neighbor ID 0)\n global_id = 0; base_index=0; index_range=6; B2L_rings[0] = [0 2 4]; B2L_maps[0] = [2 3| 0 1] L2H_maps = [4 5]\n halo_rows[0].row_offsets = [0 5 13 17 21]\n halo_rows[0].col_indices = [1 2 3 4 5\n 0 1 2 3 4 5 6 7\n 0 1 3 6\n 0 1 2 3]\n\n halo_btl[1] (received from neighbor ID 2)\n global_id = 2; base_index=0; index_range=8; B2L_rings[0] = [0 2 4]; B2L_maps[0] = [1 2| 0 3] L2H_maps = [6 7]\n halo_rows[1].row_offsets = [0 4 11 16 20]\n halo_rows[1].col_indices = [7 1 2 3\n 5 6 7 0 1 2 3\n 4 5 0 2 3\n 0 1 2 3]\n */\n}\n\ntemplate \nvoid DistributedArranger >::renumber_to_local(Matrix_d &A,\n std::vector &boundary_lists,\n std::vector &halo_lists,\n INDEX_TYPE my_id,\n int64_t base_index, INDEX_TYPE index_range,\n Vector &neighbors,\n I64Vector_h &halo_ranges_h,\n I64Vector_d &halo_ranges,\n IVector &halo_nodes)\n{\n //IVector &halo_nodes = *halo_nodes_p;\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n int num_neighbors = halo_ranges_h.size() / 2;\n int total_rows_of_neighbors = 0;\n std::vector halo_offsets(num_neighbors + 1, 0);\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n total_rows_of_neighbors += halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n halo_offsets[i + 1] = halo_offsets[i] + halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n }\n\n //IVector halo_nodes(total_rows_of_neighbors);\n IVector halo_offsets_d(num_neighbors + 1);\n amgx::thrust::copy(halo_offsets.begin(), halo_offsets.end(), halo_offsets_d.begin());\n int last_halo = halo_nodes[halo_nodes.size() - 1];\n //create renumbering of halo indices (contiguously neighbor - by - neighbor, indexing starting with the number of owned rows\n thrust_wrapper::exclusive_scan(halo_nodes.begin(), halo_nodes.end(), halo_nodes.begin(), size);\n cudaCheckError();\n A.manager->local_to_global_map.resize(halo_nodes[halo_nodes.size() - 1] - size + last_halo);\n renumber_col_indices<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), halo_ranges.raw(), halo_nodes.raw(), halo_offsets_d.raw(), base_index, index_range, num_neighbors, size, A.manager->local_to_global_map.raw());\n cudaCheckError();\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n int size = halo_lists[i]->size();\n int num_blocks = std::min(4096, (size + 127) / 128);\n renumber_halo_lists <<< num_blocks, 128>>>(halo_lists[i]->raw(), halo_nodes.raw() + halo_offsets[i], halo_ranges_h[2 * i], size);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::renumber_to_local(Matrix_d &A,\n std::vector &boundary_lists, std::vector &halo_lists,\n INDEX_TYPE my_id, I64Vector_h &halo_ranges_h, IVector &halo_nodes)\n{\n //IVector &halo_nodes = *halo_nodes_p;\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n int num_neighbors = halo_ranges_h.size() / 2;\n int total_rows_of_neighbors = 0;\n std::vector halo_offsets(num_neighbors + 1, 0);\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n total_rows_of_neighbors += halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n halo_offsets[i + 1] = halo_offsets[i] + halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n }\n\n //IVector halo_nodes(total_rows_of_neighbors+1);\n IVector halo_offsets_d(num_neighbors + 1);\n amgx::thrust::copy(halo_offsets.begin(), halo_offsets.end(), halo_offsets_d.begin());\n int last_halo = (halo_nodes.size() > 0) ? halo_nodes[halo_nodes.size() - 1] : 0;\n //create renumbering of halo indices (contiguously neighbor - by - neighbor, indexing starting with the number of owned rows\n thrust_wrapper::exclusive_scan(halo_nodes.begin(), halo_nodes.end(), halo_nodes.begin(), size);\n cudaCheckError();\n\n if (halo_nodes.size() > 0)\n {\n A.manager->local_to_global_map.resize(halo_nodes[halo_nodes.size() - 1] - size + last_halo);\n }\n\n renumber_col_indices<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), A.manager->halo_ranges.raw(), halo_nodes.raw(), halo_offsets_d.raw(), A.manager->base_index(), A.manager->index_range(), num_neighbors, size, A.manager->local_to_global_map.raw());\n cudaCheckError();\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n int size = halo_lists[i]->size();\n int num_blocks = std::min(4096, (size + 127) / 128);\n renumber_halo_lists <<< num_blocks, 128>>>(halo_lists[i]->raw(), halo_nodes.raw() + halo_offsets[i], halo_ranges_h[2 * i], size);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::prepare_local(Matrix_d &A,\n INDEX_TYPE my_id, INDEX_TYPE rings,\n int64_t base_index, INDEX_TYPE index_range,\n Vector &neighbors,\n I64Vector_h &halo_ranges_h,\n I64Vector_d &halo_ranges,\n std::vector &L2H_maps,\n std::vector &boundary_lists,\n std::vector &halo_lists,\n DistributedComms **comms)\n{\n int num_neighbors = neighbors.size();\n (*comms)->set_neighbors(num_neighbors);\n IVector halo_nodes;\n // copy halo_ranges_h to device and populate boundary_lists\n this->create_boundary_lists(A, base_index, index_range, halo_ranges_h, halo_ranges, boundary_lists, halo_nodes);\n L2H_maps.resize(num_neighbors);\n halo_lists.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n L2H_maps[i].resize(boundary_lists[i].size());\n thrust_wrapper::transform(boundary_lists[i].begin(), boundary_lists[i].end(), amgx::thrust::constant_iterator (halo_ranges_h[2 * i]), L2H_maps[i].begin(), amgx::thrust::plus ());\n halo_lists[i] = &L2H_maps[i];\n }\n\n cudaCheckError();\n // setting up tmp manager in order to use Comm API\n Matrix_d tmpA(0, 0, 0, A.getProps());\n tmpA.set_block_dimx(A.get_block_dimx());\n tmpA.set_block_dimy(A.get_block_dimy());\n tmpA.setResources(A.getResources());\n std::vector B2L_maps;\n std::vector > B2L_rings;\n DistributedManager tmp_manager(my_id, base_index, index_range, tmpA, neighbors, halo_ranges, B2L_maps, B2L_rings, comms, NULL, NULL);\n tmpA.setManager(tmp_manager);\n tmpA.manager->getComms()->exchange_vectors(boundary_lists, tmpA, 0);\n this->renumber_to_local(A, boundary_lists, halo_lists, my_id, base_index, index_range, neighbors, halo_ranges_h, halo_ranges, halo_nodes);\n};\n\ntemplate \nvoid DistributedArranger >::create_B2L_from_neighbors(Matrix_d &A,\n INDEX_TYPE my_id, INDEX_TYPE rings,\n int64_t base_index, INDEX_TYPE index_range,\n Vector &neighbors,\n I64Vector_h &halo_ranges_h,\n I64Vector_d &halo_ranges,\n std::vector &B2L_maps,\n std::vector &L2H_maps,\n std::vector > &B2L_rings,\n DistributedComms **comms,\n std::vector **halo_rows,\n std::vector **halo_btl)\n{\n this->halo_coloring = (*comms)->halo_coloring;\n std::vector boundary_lists;\n std::vector halo_lists;\n this->prepare_local(A, my_id, rings, base_index, index_range, neighbors, halo_ranges_h, halo_ranges, L2H_maps, boundary_lists, halo_lists, comms);\n int num_neighbors = neighbors.size();\n // initialize BTL_maps and BTl_rings\n this->create_maps(A, rings, num_neighbors, B2L_maps, B2L_rings, boundary_lists);\n // initialize halo_btl\n this->create_halo_btl(A, rings, num_neighbors, neighbors, base_index, index_range, my_id, B2L_maps, B2L_rings, halo_btl, halo_lists);\n // initialize halo_rows\n this->create_halo_rows(A, rings, num_neighbors, B2L_maps, B2L_rings, halo_rows, halo_lists);\n\n if (this->halo_coloring != LAST)\n {\n // setting up tmp manager in order to use Comm API\n Matrix_d tmpA(0, 0, 0, A.getProps());\n tmpA.set_block_dimx(A.get_block_dimx());\n tmpA.set_block_dimy(A.get_block_dimy());\n tmpA.setResources(A.getResources());\n DistributedManager tmp_manager(my_id, base_index, index_range, tmpA, neighbors, halo_ranges, B2L_maps, B2L_rings, comms, NULL, NULL);\n tmpA.setManager(tmp_manager);\n tmpA.manager->getComms()->exchange_matrix_halo(**halo_rows, **halo_btl, tmpA);\n }\n}\n\ntemplate \nvoid DistributedArranger >::create_neighbors(Matrix_d &A, INDEX_TYPE my_id,\n int64_t &base_index, INDEX_TYPE &index_range,\n Vector &neighbors,\n I64Vector_h &halo_ranges_h)\n{\n //if (this->num_part>32) FatalError(\"Matrix slice Level 1 module supports < 32 partitions\", AMGX_ERR_NOT_IMPLEMENTED);\n if (this->num_part == 0 || this->part_offsets.size() == 0)\n {\n FatalError(\"Partitioning scheme is not set\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n IVector part_sizes(this->num_part, 0);\n IVector inv_neighbors(this->num_part, 0);\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n base_index = this->part_offsets[my_id];\n index_range = this->part_offsets[my_id + 1] - this->part_offsets[my_id];\n //\n // Step 1: flag neighbors that I have an edge to\n //\n //go through all the column indices and compare it to index ranges of nodes\n calc_num_neighbors<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), this->part_offsets.raw(), part_sizes.raw(), this->num_part, my_id, size);\n cudaCheckError();\n std::vector neighbor_flags(this->num_part);\n amgx::thrust::copy(part_sizes.begin(), part_sizes.end(), neighbor_flags.begin());\n //total number of neighbors\n int num_neighbors = amgx::thrust::reduce(neighbor_flags.begin(), neighbor_flags.end());\n cudaCheckError();\n neighbors.resize(num_neighbors);\n halo_ranges_h.resize(num_neighbors * 2);\n //initialize manager\n int active_part = 0;\n\n for (int i = 0; i < this->num_part; i++)\n {\n if (neighbor_flags[i] > 0)\n {\n neighbors[active_part] = i;\n halo_ranges_h[2 * active_part] = this->part_offsets[i];\n halo_ranges_h[2 * active_part + 1] = this->part_offsets[i + 1];\n active_part++;\n }\n }\n}\n\n\n// need for CUDA 4.2\ntemplate \nvoid DistributedArranger >::create_neighbors(Matrix_d &A, INDEX_TYPE my_id,\n I64Vector_h &halo_ranges_h)\n{\n if (this->num_part == 0 || this->part_offsets.size() == 0)\n {\n FatalError(\"Partitioning scheme is not set\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n IVector part_sizes(this->num_part, 0);\n IVector inv_neighbors(this->num_part, 0);\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n A.manager->set_base_index(this->part_offsets[my_id]);\n A.manager->set_index_range(this->part_offsets[my_id + 1] - this->part_offsets[my_id]);\n //\n // Step 1: flag neighbors that I have an edge to\n //\n //go through all the column indices and compare it to index ranges of nodes\n calc_num_neighbors<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), this->part_offsets.raw(), part_sizes.raw(), this->num_part, my_id, size);\n cudaCheckError();\n std::vector neighbor_flags(this->num_part);\n amgx::thrust::copy(part_sizes.begin(), part_sizes.end(), neighbor_flags.begin());\n //total number of neighbors\n int num_neighbors = amgx::thrust::reduce(neighbor_flags.begin(), neighbor_flags.end());\n cudaCheckError();\n A.manager->neighbors.resize(num_neighbors);\n halo_ranges_h.resize(num_neighbors * 2);\n //initialize manager\n int active_part = 0;\n\n for (int i = 0; i < this->num_part; i++)\n {\n if (neighbor_flags[i] > 0)\n {\n A.manager->neighbors[active_part] = i;\n halo_ranges_h[2 * active_part] = this->part_offsets[i];\n halo_ranges_h[2 * active_part + 1] = this->part_offsets[i + 1];\n active_part++;\n }\n }\n}\n\ntemplate \nvoid DistributedArranger >::update_neighbors_list(\n Matrix_d &A, IVector_h &neighbors, I64Vector_h &halo_ranges_h, I64Vector &halo_ranges, I64Vector_h &part_offsets_h, I64Vector &part_offsets, std::vector &halo_rows_row_offsets, std::vector &halo_rows_col_indices)\n{\n int num_partitions = A.manager->get_num_partitions();\n int my_id = A.manager->global_id();\n int num_neighbors = neighbors.size(); // current number of neighbors\n int total_halo_rows = 0;\n int total_halo_nnz = 0;\n IVector neighbor_flags(num_partitions, 0);\n\n for (int i = 0; i < halo_rows_row_offsets.size(); i++)\n {\n int num_halo_rows = halo_rows_row_offsets[i].size() - 1;\n\n if (num_halo_rows > 0)\n {\n total_halo_rows += num_halo_rows;\n total_halo_nnz += halo_rows_row_offsets[i][num_halo_rows];\n int num_blocks = std::min(4096, (num_halo_rows + 127) / 128);\n calc_num_neighbors_v2_global<16> <<< num_blocks, 128>>>(halo_rows_row_offsets[i].raw(), halo_rows_col_indices[i].raw(),\n part_offsets.raw(), neighbor_flags.raw(), num_partitions, my_id, num_halo_rows);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n IVector_h neighbor_flags_h = neighbor_flags;\n\n // unset 1-ring neighbors & myself\n for (int i = 0; i < num_neighbors; i++)\n {\n neighbor_flags_h[neighbors[i]] = 0;\n }\n\n neighbor_flags_h[my_id] = 0;\n // this will update neighbor list and halo ranges, note that we don't change 1-ring neighbors order\n append_neighbors(A, neighbors, halo_ranges_h, halo_ranges, neighbor_flags_h, part_offsets_h);\n}\n\ntemplate \nvoid DistributedArranger >::append_neighbors(Matrix_d &A, IVector_h &neighbors, I64Vector_h &halo_ranges_h, I64Vector &halo_ranges, IVector_h &neighbor_flags, I64Vector_h &part_offsets_h)\n{\n // This function creates arrays neighbors, halo_ranges_h and halo_ranges\n // base on neighbor_flags\n // Here do an MPI_allgather to figure out which partitions need data from me\n // This is required for non-symmetric matrices\n int num_part = A.manager->get_num_partitions();\n int my_id = A.manager->global_id();\n // pack 0/1 array into array of integers (size/32)\n int packed_size = (num_part + 32 - 1) / 32;\n IVector_h packed_nf(packed_size, 0);\n\n for (int i = 0; i < num_part; i++)\n {\n int packed_pos = i / 32;\n int bit_pos = i % 32;\n packed_nf[packed_pos] += (neighbor_flags[i] << bit_pos);\n }\n\n // exchange packed neighbor flags\n IVector_h gathered_packed_nf;\n A.manager->getComms()->all_gather_v(packed_nf, gathered_packed_nf, num_part);\n // assign neighbors that have edges to me\n int my_id_pos = my_id / 32;\n int my_id_bit = my_id % 32;\n\n for (int i = 0; i < num_part; i++)\n if (gathered_packed_nf[i * packed_size + my_id_pos] & (1 << my_id_bit)) // check my bit\n {\n neighbor_flags[i] = 1;\n }\n\n // compute total number of new neighbors\n int new_neighbors = amgx::thrust::reduce(neighbor_flags.begin(), neighbor_flags.end());\n cudaCheckError();\n // save old size\n int old_neighbors = neighbors.size();\n neighbors.resize(old_neighbors + new_neighbors);\n halo_ranges_h.resize(old_neighbors * 2 + new_neighbors * 2);\n // initialize manager->neighbors and manager->halo_ranges_h for the new nodes\n int active_part = old_neighbors;\n this->num_part = num_part;\n\n for (int i = 0; i < num_part; i++)\n {\n if (neighbor_flags[i] > 0)\n {\n neighbors[active_part] = i;\n halo_ranges_h[2 * active_part] = part_offsets_h[i];\n halo_ranges_h[2 * active_part + 1] = part_offsets_h[i + 1];\n active_part++;\n }\n }\n\n halo_ranges.resize(old_neighbors * 2 + new_neighbors * 2);\n amgx::thrust::copy(halo_ranges_h.begin() + old_neighbors * 2, halo_ranges_h.end(), halo_ranges.begin() + old_neighbors * 2);\n cudaCheckError();\n}\n\n\ntemplate \nvoid DistributedArranger >::create_neighbors_v2(Matrix_d &A)\n{\n // This function counts the number of neighbors, identifies the id of each neighbor,\n // And creates array halo_ranges, which is the range of rows for each neighbor\n // Outputs: A.manager->neighbors\n // A.manager->halo_ranges_h\n // A.manager->halo_ranges\n int num_part = A.manager->get_num_partitions();\n int my_id = A.manager->global_id();\n A.manager->set_base_index(A.manager->part_offsets_h[my_id]);\n A.manager->set_index_range(A.manager->part_offsets_h[my_id + 1] - A.manager->part_offsets_h[my_id]);\n IVector inv_neighbors(num_part, 0);\n IVector part_sizes(num_part, 0);\n int size = A.get_num_rows();\n\n if (size == 0)\n {\n A.manager->neighbors.resize(0);\n }\n\n // TODO: Are these the optimal block_sizes?\n int num_blocks = std::min(4096, (size + 127) / 128);\n\n // -----------------------------------------------\n // Step 1: flag neighbors that I have an edge to\n // -----------------------------------------------\n //go through all the column indices and check if they are halo vertices\n if (num_blocks != 0)\n {\n calc_num_neighbors_v2_local<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), A.manager->part_offsets.raw(), part_sizes.raw(), A.manager->local_to_global_map.raw(), num_part, my_id, size);\n cudaCheckError();\n }\n\n // D2H copy\n IVector_h neighbor_flags(num_part);\n amgx::thrust::copy(part_sizes.begin(), part_sizes.end(), neighbor_flags.begin());\n cudaCheckError();\n append_neighbors(A, A.manager->neighbors, A.manager->halo_ranges_h, A.manager->halo_ranges, neighbor_flags, A.manager->part_offsets_h);\n}\n\ntemplate \nvoid DistributedArranger >::create_neighbors_v2_global(Matrix_d &A, I64Vector_d &A_col_indices_global)\n{\n // This function counts the number of neighbors, identifies the id of each neighbor,\n // And creates array halo_ranges, which is the range of rows for each neighbor\n // Outputs: A.manager->neighbors\n // A.manager->halo_ranges_h\n // A.manager->halo_ranges\n int num_part = A.manager->get_num_partitions();\n int my_id = A.manager->global_id();\n IVector part_sizes(num_part, 0);\n int size = A.get_num_rows();\n\n if (size == 0)\n {\n A.manager->neighbors.resize(0);\n }\n\n // TODO: Are these the optimal block_sizes?\n int num_blocks = std::min(4096, (size + 127) / 128);\n\n // -----------------------------------------------\n // Step 1: flag neighbors that I have an edge to\n // -----------------------------------------------\n //go through all the column indices and check if they are halo vertices\n if (num_blocks != 0)\n {\n calc_num_neighbors_v2_global<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A_col_indices_global.raw(), A.manager->part_offsets.raw(), part_sizes.raw(), num_part, my_id, size);\n cudaCheckError();\n }\n\n // D2H copy\n IVector_h neighbor_flags(num_part);\n amgx::thrust::copy(part_sizes.begin(), part_sizes.end(), neighbor_flags.begin());\n cudaCheckError();\n append_neighbors(A, A.manager->neighbors, A.manager->halo_ranges_h, A.manager->halo_ranges, neighbor_flags, A.manager->part_offsets_h);\n}\n\ntemplate \nvoid DistributedArranger >::create_B2L_one_ring(Matrix_d &A)\n{\n std::vector boundary_lists;\n // create my boundary lists, renumber halos, create L2H maps\n create_boundary_lists_v2(A, boundary_lists);\n // compute halo offsets\n int neighbors = A.manager->num_neighbors();\n A.manager->halo_offsets.resize(neighbors + 1, 0);\n\n for (int i = 0; i < neighbors; i++)\n {\n A.manager->halo_offsets[i] = boundary_lists[i].size();\n }\n\n thrust_wrapper::exclusive_scan(A.manager->halo_offsets.begin(), A.manager->halo_offsets.end(), A.manager->halo_offsets.begin(), A.get_num_rows());\n cudaCheckError();\n // receive neighbors boundary lists for b2l maps\n A.manager->getComms()->exchange_vectors(boundary_lists, A, 0);\n // copy boundary lists to B2L maps\n // release internal boundary lists array\n A.manager->B2L_maps.resize(neighbors);\n\n for (int i = 0; i < neighbors; i++)\n {\n A.manager->B2L_maps[i] = boundary_lists[i];\n }\n\n A.manager->B2L_rings.resize(neighbors);\n\n for (int i = 0; i < neighbors; i++)\n {\n A.manager->B2L_rings[i].resize(2);\n A.manager->B2L_rings[i][0] = 0;\n A.manager->B2L_rings[i][1] = A.manager->B2L_maps[i].size();\n }\n\n this->initialize_B2L_maps_offsets(A, 1);\n A.manager->A = &A;\n}\n\ntemplate \nvoid DistributedArranger >::createRowsLists(Matrix_d &A, bool is_matrix_P)\n{\n int A_num_rows, offset;\n A.getOffsetAndSizeForView(OWNED, &offset, &A_num_rows);\n IVector boundary_flags(A_num_rows + 1, 0);\n IVector interior_flags(A_num_rows + 1, 0);\n\n int num_nz_interior;\n if (!is_matrix_P)\n {\n for (int i = 0; i < A.manager->neighbors.size(); i++)\n {\n int size = A.manager->B2L_rings[i][1];\n int num_blocks = std::min(4096, (size + 127) / 128);\n\n if ( size > 0)\n {\n flag_boundary_nodes <<< num_blocks, 128>>>(A.manager->B2L_maps[i].raw(), boundary_flags.raw(), size);\n cudaCheckError();\n }\n }\n }\n else\n {\n const int cta_size = 256;\n const int nWarps = cta_size / 32;\n int grid_size = std::min( 4096, (int) (A_num_rows + nWarps - 1) / nWarps);\n\n if (A_num_rows > 0)\n {\n IVector nnz_interior(1, 0);\n flag_boundary_nodes_P <<< grid_size, cta_size>>>(A.row_offsets.raw(), A.col_indices.raw(), boundary_flags.raw(), A.manager->halo_offsets[0], A_num_rows, nnz_interior.raw());\n cudaCheckError();\n cudaMemcpy(&num_nz_interior, nnz_interior.raw(), sizeof(int), cudaMemcpyDefault);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n int size = A_num_rows;\n int num_blocks = std::min(4096, (size + 127) / 128);\n\n if (size > 0)\n {\n flag_interior_nodes <<< num_blocks, 128>>>(boundary_flags.raw(), interior_flags.raw(), size);\n cudaCheckError();\n }\n\n if (boundary_flags.size() > 1)\n {\n thrust_wrapper::exclusive_scan(boundary_flags.begin(), boundary_flags.end(), boundary_flags.begin());\n cudaCheckError();\n }\n\n thrust_wrapper::exclusive_scan(interior_flags.begin(), interior_flags.end(), interior_flags.begin());\n cudaCheckError();\n int num_interior_rows = interior_flags[size];\n int num_boundary_rows = boundary_flags[size];\n A.manager->interior_rows_list.resize(num_interior_rows);\n A.manager->boundary_rows_list.resize(num_boundary_rows);\n\n // Create the list of interior and boundary nodes, using the interior_flag and boundary_flag\n if (size > 0)\n {\n fill_rows_lists <<< num_blocks, 128>>>(A.manager->interior_rows_list.raw(), A.manager->boundary_rows_list.raw(), interior_flags.raw(), boundary_flags.raw(), size);\n cudaCheckError();\n }\n\n // Manually set the view sizes for P\n if(is_matrix_P)\n {\n int num_rows_owned = num_interior_rows + num_boundary_rows;\n int num_nz_owned = A.row_offsets[num_rows_owned];\n A.manager->setViewSizes(num_interior_rows, num_nz_interior, num_rows_owned, num_nz_owned, num_rows_owned, num_nz_owned, A.get_num_rows(), A.get_num_nz());\n }\n}\n\ntemplate \nvoid DistributedArranger >::create_boundary_lists_v3(Matrix_d &A)\n{\n // -----------------------------------------------\n // Flag all nodes in neighbors that I will need (that is I have an edge pointing to it)\n // -----------------------------------------------\n int A_num_rows = A.get_num_rows();\n\n if (A_num_rows == 0)\n {\n return;\n }\n\n int num_neighbors = A.manager->neighbors.size();\n // compute halo offsets\n int total_rows_of_neighbors = 0;\n std::vector halo_offsets(num_neighbors + 1, 0);\n int max_neighbor_size = 0;\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n int num_rows_neighbor = A.manager->halo_ranges_h[2 * i + 1] - A.manager->halo_ranges_h[2 * i];\n total_rows_of_neighbors += num_rows_neighbor;\n halo_offsets[i + 1] = halo_offsets[i] + num_rows_neighbor;\n max_neighbor_size = max_neighbor_size > num_rows_neighbor ? max_neighbor_size : num_rows_neighbor;;\n }\n\n // copy offsets to device\n IVector halo_offsets_d(num_neighbors + 1);\n amgx::thrust::copy(halo_offsets.begin(), halo_offsets.end(), halo_offsets_d.begin());\n // flag neighbors nodes that I need\n IVector halo_nodes(total_rows_of_neighbors + 1);\n thrust_wrapper::fill(halo_nodes.begin(), halo_nodes.end(), 0);\n cudaCheckError();\n int size = A.manager->local_to_global_map.size();\n int num_blocks = std::min(4096, (size + 127) / 128);\n\n if (size > 0)\n {\n flag_halo_nodes_from_local_to_global_map <<< num_blocks, 128>>>(A.manager->halo_ranges.raw(), halo_nodes.raw(), halo_offsets_d.raw(), A.manager->base_index(), A.manager->index_range(), num_neighbors, size, A.manager->local_to_global_map.raw());\n cudaCheckError();\n }\n\n IVector halo_indices(max_neighbor_size + 1);\n A.manager->L2H_maps.resize(num_neighbors);\n A.manager->B2L_maps.resize(num_neighbors);\n is_less_than_zero pred;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n // compute halo indices map [] = 0-based halo number\n int size = halo_offsets[i + 1] - halo_offsets[i];\n amgx::thrust::replace_copy_if(halo_nodes.begin() + halo_offsets[i], halo_nodes.begin() + halo_offsets[i + 1], halo_indices.begin(), pred, 1);\n thrust_wrapper::exclusive_scan(halo_indices.begin(), halo_indices.begin() + size + 1, halo_indices.begin());\n cudaCheckError();\n int num_halo = halo_indices[size];\n A.manager->L2H_maps[i].resize(num_halo);\n A.manager->B2L_maps[i].resize(num_halo);\n // create my boundary lists = list of neighbor inner nodes corresponding to halo numbers 0..M\n // For each node in B2L_map, need to find corresponding local id and store in L2H_map\n int num_blocks = std::min(4096, (size + 127) / 128);\n populate_bdy_list_and_l2h_map <<< num_blocks, 128>>>(A.manager->L2H_maps[i].raw(), halo_nodes.raw() + halo_offsets[i], halo_indices.raw(), A.manager->B2L_maps[i].raw(), size);\n cudaCheckError();\n }\n}\n\n\ntemplate \nvoid DistributedArranger >::initialize_B2L_maps_offsets(Matrix_d &A, int num_rings)\n{\n // Compute the total number of bdy rows in each ring\n std::vector B2L_maps_offsets_h(num_rings);\n A.manager->B2L_rings_sizes.resize(num_rings);\n int num_neighbors = A.manager->neighbors.size();\n\n for (int k = 0; k < num_rings; k++)\n {\n B2L_maps_offsets_h[k].resize(num_neighbors + 1);;\n B2L_maps_offsets_h[k][0] = 0;\n\n for (int j = 0; j < num_neighbors; j++)\n {\n B2L_maps_offsets_h[k][j + 1] = B2L_maps_offsets_h[k][j] + A.manager->B2L_rings[j][k + 1];\n }\n\n A.manager->B2L_rings_sizes[k] = B2L_maps_offsets_h[k][num_neighbors];\n }\n\n // Copy maps_offsets to device\n A.manager->B2L_maps_offsets.resize(num_rings);\n\n for (int i = 0; i < num_rings; i++)\n {\n A.manager->B2L_maps_offsets[i] = B2L_maps_offsets_h[i];\n }\n\n // Store the B2L_maps ptrs on the device\n std::vector B2L_maps_ptrs_h(num_neighbors);\n\n for (int j = 0; j < num_neighbors; j++)\n {\n B2L_maps_ptrs_h[j] = A.manager->B2L_maps[j].raw();\n }\n\n A.manager->B2L_maps_ptrs = B2L_maps_ptrs_h;\n}\n\n\ntemplate \nvoid DistributedArranger >::create_one_ring_halo_rows(Matrix_d &A)\n{\n // what it does:\n // appends halo rows to matrix A\n // creates new B2L_rings, B2L_maps and L2H_maps\n // input:\n // matrix A and 1-ring B2L_maps\n int num_partitions = A.manager->get_num_partitions();\n int my_id = A.manager->global_id();\n int num_neighbors = A.manager->B2L_maps.size();\n int num_ring1_neighbors = num_neighbors;\n int num_ring1_indices = A.manager->local_to_global_map.size();\n std::vector halo_row_offsets(num_neighbors);\n std::vector halo_global_indices(num_neighbors);\n std::vector halo_local_indices(num_neighbors);\n std::vector halo_values(num_neighbors);\n // step 1: setup halo rows with global indices\n create_halo_rows_global_indices(A, halo_row_offsets, halo_global_indices, halo_values);\n // step 2: exchange halo rows\n std::vector dummy_halo_ids(0);\n A.manager->getComms()->exchange_matrix_halo(halo_row_offsets, halo_global_indices, halo_values, dummy_halo_ids, A.manager->neighbors, A.manager->global_id());\n // step 3: append the list f neighbors with the new ring-2 neighbors\n update_neighbors_list(A, A.manager->neighbors, A.manager->halo_ranges_h, A.manager->halo_ranges, A.manager->part_offsets_h, A.manager->part_offsets, halo_row_offsets, halo_global_indices);\n num_neighbors = A.manager->neighbors.size();\n // step 4: mark neighbor nodes, create local halo indices for the 2nd ring\n // this function also creates boundary lists for the 2nd ring if boundary lists array size > 0\n std::vector boundary_lists(num_neighbors);\n compute_local_halo_indices(A.row_offsets, A.col_indices, halo_row_offsets, halo_global_indices, halo_local_indices, A.manager->local_to_global_map, boundary_lists, A.manager->neighbors, A.manager->halo_ranges_h, A.manager->halo_ranges, A.manager->halo_offsets, A.manager->base_index(), A.manager->index_range(), A.get_num_rows(), 1);\n // update renumbering arrays (set identity for 2nd ring)\n //A.manager->halo_P_renumbering.resize(A.manager->local_to_global_map.size());\n //A.manager->halo_R_renumbering.resize(A.manager->local_to_global_map.size());\n //thrust_wrapper::sequence(A.manager->halo_P_renumbering.begin() + num_ring1_indices, A.manager->halo_P_renumbering.end(), num_ring1_indices + A.get_num_rows());\n //thrust_wrapper::sequence(A.manager->halo_R_renumbering.begin() + num_ring1_indices, A.manager->halo_R_renumbering.end(), num_ring1_indices + A.get_num_rows());\n // step 5: update L2H maps = identity\n A.manager->getComms()->set_neighbors(A.manager->neighbors.size());\n // step 6: receive neighbors boundary lists for 2-ring b2l maps\n A.manager->getComms()->exchange_vectors(boundary_lists, A, 0);\n // step 7: update B2L rings & B2L maps\n int rings = 2;\n A.manager->B2L_rings.resize(num_neighbors);\n A.manager->B2L_maps.resize(num_neighbors);\n // modify 1-ring where necessary\n int ring = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n // set size to 0 for 2-ring neighbors only\n if (i >= num_ring1_neighbors)\n {\n A.manager->B2L_maps[i].resize(0);\n }\n\n A.manager->B2L_rings[i].resize(rings + 1);\n A.manager->B2L_rings[i][0] = 0;\n A.manager->B2L_rings[i][ring + 1] = A.manager->B2L_maps[i].size();\n }\n\n // fill up 2-ring maps\n ring = 1;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n // append 2nd ring\n int ring1_size = A.manager->B2L_maps[i].size();\n A.manager->B2L_maps[i].resize(ring1_size + boundary_lists[i].size());\n amgx::thrust::copy(boundary_lists[i].begin(), boundary_lists[i].end(), A.manager->B2L_maps[i].begin() + ring1_size);\n A.manager->B2L_rings[i][ring + 1] = A.manager->B2L_maps[i].size();\n }\n\n cudaCheckError();\n this->initialize_B2L_maps_offsets(A, 2);\n // step 8: append halo_rows to matrix A\n // compute new # of rows & nnz\n this->append_halo_rows(A, halo_row_offsets, halo_local_indices, halo_values);\n // initialize the manager\n A.manager->set_initialized(A.row_offsets);\n // Compute the diagonal\n // TODO: Should only compute diagonal of 1-ring halo rows\n ViewType oldView = A.currentView();\n A.setView(FULL);\n A.set_allow_recompute_diag(true);\n A.computeDiagonal();\n A.setView(oldView);\n // the following steps are necessary only for latency hiding/renumbering, i.e. to use reorder_matrix()\n#if 0\n // step 6: for 2 rings: setup halo_btl\n create_halo_btl_multiple_rings(A, 2);\n#endif\n}\n\n// need for CUDA 4.2\ntemplate \nvoid DistributedArranger >::create_boundary_lists(Matrix_d &A,\n int64_t base_index, INDEX_TYPE index_range,\n I64Vector_h &halo_ranges_h, I64Vector_d &halo_ranges,\n std::vector &boundary_lists, IVector &halo_nodes)\n{\n //\n // Step 2: flag all nodes in neighbors that I will need (that is I have an edge pointing to it)\n //\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n int num_neighbors = halo_ranges_h.size() / 2;\n halo_ranges.resize(num_neighbors * 2);\n amgx::thrust::copy(halo_ranges_h.begin(), halo_ranges_h.end(), halo_ranges.begin());\n cudaCheckError();\n boundary_lists.resize(num_neighbors);\n int total_rows_of_neighbors = 0;\n std::vector halo_offsets(num_neighbors + 1, 0);\n int max_neighbor_size = 0;\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n total_rows_of_neighbors += halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n halo_offsets[i + 1] = halo_offsets[i] + halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n max_neighbor_size = max_neighbor_size > (halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i]) ? max_neighbor_size : (halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i]);;\n }\n\n IVector halo_offsets_d(num_neighbors + 1);\n amgx::thrust::copy(halo_offsets.begin(), halo_offsets.end(), halo_offsets_d.begin());\n halo_nodes.resize(total_rows_of_neighbors);\n thrust_wrapper::fill(halo_nodes.begin(), halo_nodes.end(), 0);\n cudaCheckError();\n flag_halo_nodes<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), halo_ranges.raw(), halo_nodes.raw(), halo_offsets_d.raw(), base_index, index_range, num_neighbors, size);\n cudaCheckError();\n //\n // Step 3: Write the 1-ring halo indices for each neighbor\n //\n IVector halo_indices(max_neighbor_size);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int last = halo_nodes[halo_offsets[i + 1] - 1];\n thrust_wrapper::exclusive_scan(halo_nodes.begin() + halo_offsets[i], halo_nodes.begin() + halo_offsets[i + 1], halo_indices.begin());\n int num_halo = halo_indices[halo_offsets[i + 1] - halo_offsets[i] - 1] + last;\n //A.manager.B2L_maps[i].resize(num_halo);\n boundary_lists[i].resize(num_halo);\n int size = halo_offsets[i + 1] - halo_offsets[i];\n //int num_blocks=min(4096,(size+127)/128);\n populate_B2L <<< num_blocks, 128>>>(halo_indices.raw(), boundary_lists[i].raw(),/*A.manager.B2L_maps[i].raw(),*/ last, size);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::create_boundary_lists(Matrix_d &A,\n I64Vector_h &halo_ranges_h,\n std::vector &boundary_lists, IVector &halo_nodes)\n{\n //\n // Step 2: flag all nodes in neighbors that I will need (that is I have an edge pointing to it)\n //\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n int num_neighbors = halo_ranges_h.size() / 2;\n A.manager->halo_ranges.resize(num_neighbors * 2);\n amgx::thrust::copy(halo_ranges_h.begin(), halo_ranges_h.end(), A.manager->halo_ranges.begin());\n cudaCheckError();\n boundary_lists.resize(num_neighbors);\n int total_rows_of_neighbors = 0;\n std::vector halo_offsets(num_neighbors + 1, 0);\n int max_neighbor_size = 0;\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n total_rows_of_neighbors += halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n halo_offsets[i + 1] = halo_offsets[i] + halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n max_neighbor_size = max_neighbor_size > (halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i]) ? max_neighbor_size : (halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i]);;\n }\n\n IVector halo_offsets_d(num_neighbors + 1);\n amgx::thrust::copy(halo_offsets.begin(), halo_offsets.end(), halo_offsets_d.begin());\n halo_nodes.resize(total_rows_of_neighbors);;\n thrust_wrapper::fill(halo_nodes.begin(), halo_nodes.end(), 0);\n cudaCheckError();\n flag_halo_nodes<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), A.manager->halo_ranges.raw(), halo_nodes.raw(), halo_offsets_d.raw(), A.manager->base_index(), A.manager->index_range(), num_neighbors, size);\n cudaCheckError();\n //\n // Step 3: Write the 1-ring halo indices for each neighbor\n //\n IVector halo_indices(max_neighbor_size);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int last = halo_nodes[halo_offsets[i + 1] - 1];\n thrust_wrapper::exclusive_scan(halo_nodes.begin() + halo_offsets[i], halo_nodes.begin() + halo_offsets[i + 1], halo_indices.begin());\n int num_halo = halo_indices[halo_offsets[i + 1] - halo_offsets[i] - 1] + last;\n //A.manager.B2L_maps[i].resize(num_halo);\n boundary_lists[i].resize(num_halo);\n int size = halo_offsets[i + 1] - halo_offsets[i];\n //int num_blocks=min(4096,(size+127)/128);\n populate_B2L <<< num_blocks, 128>>>(halo_indices.raw(), boundary_lists[i].raw(),/*A.manager.B2L_maps[i].raw(),*/ last, size);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::create_boundary_lists_v2(Matrix_d &A, std::vector &boundary_lists)\n{\n // -----------------------------------------------\n // Flag all nodes in neighbors that I will need (that is I have an edge pointing to it)\n // -----------------------------------------------\n int size = A.get_num_rows();\n\n if (size == 0)\n {\n return;\n }\n\n //TODO: Are these the optimal block_sizes?\n int num_blocks = std::min(4096, (size + 127) / 128);\n int num_neighbors = A.manager->neighbors.size();\n boundary_lists.resize(num_neighbors);\n // compute halo offsets\n int total_rows_of_neighbors = 0;\n std::vector halo_offsets(num_neighbors + 1, 0);\n int max_neighbor_size = 0;\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n int num_rows_neighbor = A.manager->halo_ranges_h[2 * i + 1] - A.manager->halo_ranges_h[2 * i];\n total_rows_of_neighbors += num_rows_neighbor;\n halo_offsets[i + 1] = halo_offsets[i] + num_rows_neighbor;\n max_neighbor_size = max_neighbor_size > num_rows_neighbor ? max_neighbor_size : num_rows_neighbor;;\n }\n\n // copy offsets to device\n IVector halo_offsets_d(num_neighbors + 1);\n amgx::thrust::copy(halo_offsets.begin(), halo_offsets.end(), halo_offsets_d.begin());\n // flag neighbors nodes that I need\n IVector halo_nodes(total_rows_of_neighbors);\n amgx::thrust::fill(halo_nodes.begin(), halo_nodes.end(), 0);\n cudaCheckError();\n flag_halo_nodes_local<16, 1> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), A.manager->halo_ranges.raw(), halo_nodes.raw(), halo_offsets_d.raw(), A.manager->base_index(), A.manager->index_range(), num_neighbors, size, A.manager->local_to_global_map.raw());\n cudaCheckError();\n // store renumbering so later can use on P\n //A.manager->halo_P_renumbering.resize(A.manager->local_to_global_map.size());\n //A.manager->halo_R_renumbering.resize(A.manager->local_to_global_map.size());\n // we need to update local toglobal maps as well because we will be renumbering halo nodes\n I64Vector new_local_to_global_map(A.manager->local_to_global_map.size(), -1);\n // compute new local halo indices for each neighbor node\n // renumbering: compute new column indices sorted by global index\n IVector halo_indices(max_neighbor_size);\n std::vector num_halo(num_neighbors);\n std::vector last(num_neighbors);\n std::vector halo_lists;\n halo_lists.resize(num_neighbors);\n A.manager->L2H_maps.resize(num_neighbors);\n int halo_base = size;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n // compute halo indices map [] = 0-based halo number\n last[i] = halo_nodes[halo_offsets[i + 1] - 1];\n thrust_wrapper::exclusive_scan(halo_nodes.begin() + halo_offsets[i], halo_nodes.begin() + halo_offsets[i + 1], halo_indices.begin());\n num_halo[i] = halo_indices[halo_offsets[i + 1] - halo_offsets[i] - 1] + last[i];\n // renumber my columns for this particular neighbor\n calc_new_halo_mapping<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), A.manager->halo_ranges.raw() + (int64_t) (2 * i), 1, size, A.manager->local_to_global_map.raw(), new_local_to_global_map.raw(), halo_indices.raw(), halo_base, A.manager->global_id());\n cudaCheckError();\n //calc_new_halo_mapping_v2<16><<>>( A.manager->halo_ranges.raw() + (int64_t) (2*i), 1, size, A.manager->local_to_global_map.raw(), new_local_to_global_map.raw(), halo_indices.raw(), halo_base, A.manager->local_to_global_map.size(),A.manager->global_id());\n cudaCheckError();\n // create my boundary lists = list of neighbor inner nodes corresponding to halo numbers 0..M\n boundary_lists[i].resize(num_halo[i]);\n int size = halo_offsets[i + 1] - halo_offsets[i];\n populate_B2L <<< num_blocks, 128>>>(halo_indices.raw(), boundary_lists[i].raw(), last[i], size);\n cudaCheckError();\n // identity L2H\n A.manager->L2H_maps[i].resize(num_halo[i]);\n thrust_wrapper::sequence(A.manager->L2H_maps[i].begin(), A.manager->L2H_maps[i].end(), halo_base);\n halo_lists[i] = &A.manager->L2H_maps[i];\n // continous numbering for halo indices\n halo_base += num_halo[i];\n }\n\n cudaCheckError();\n // Now change the negative col to positive\n flip_sign_col<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), size);\n cudaCheckError();\n // copy new local to global mapping\n amgx::thrust::copy(new_local_to_global_map.begin(), new_local_to_global_map.end(), A.manager->local_to_global_map.begin());\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::compute_local_halo_indices( IVector &A_row_offsets, IVector &A_col_indices, std::vector &halo_row_offsets, std::vector &halo_global_indices,\n std::vector &halo_local_indices, I64Vector &local_to_global, std::vector &boundary_lists, IVector_h &neighbors, I64Vector_h &halo_ranges_h, I64Vector &halo_ranges, IVector_h &halo_offsets, int64_t base_index, int64_t index_range, int A_num_rows, int current_num_rings)\n{\n // This function checks the halo_col_indices received from the neighbors, and identifies\n // new halo_indices and updates halo_offsets, local_to_global_map accordingly\n // input: halo row offsets,\n // halo global column indices\n // current local to global map\n // neighbors (new discovered neighbors should already be included)\n // halo_ranges for the neighbors\n //\n // output: halo offsets,\n // halo local column indices,\n // updated local to global map (in place)\n int size = A_num_rows;\n //TODO: Are these the optimal block_sizes?\n int num_blocks = std::min(4096, (size + 127) / 128);\n int num_neighbors = neighbors.size();\n // compute neighbor offsets & total number of neighbor rows\n int total_rows_of_neighbors = 0;\n std::vector neighbor_offsets_h(num_neighbors + 1, 0);\n int max_neighbor_size = 0;\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n total_rows_of_neighbors += halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n neighbor_offsets_h[i + 1] = neighbor_offsets_h[i] + halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n max_neighbor_size = max_neighbor_size > (halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i]) ? max_neighbor_size : (halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i]);\n }\n\n // copy offsets to device\n IVector neighbor_offsets(num_neighbors + 1);\n amgx::thrust::copy(neighbor_offsets_h.begin(), neighbor_offsets_h.end(), neighbor_offsets.begin());\n // store flags for all neighbor nodes\n IVector neighbor_nodes(total_rows_of_neighbors);\n // flag neighbor nodes that are in the existing rings as -(local_index)-1\n thrust_wrapper::fill(neighbor_nodes.begin(), neighbor_nodes.end(), 0);\n cudaCheckError();\n int local_to_global_size = local_to_global.size();\n int num_blocks2 = std::min(4096, (local_to_global_size + 127) / 128);\n\n if (local_to_global_size != 0)\n {\n flag_halo_nodes_local_v2 <<< num_blocks2, 128>>>(halo_ranges.raw(), neighbor_nodes.raw(), neighbor_offsets.raw(), base_index, index_range, num_neighbors, local_to_global_size, local_to_global.raw());\n cudaCheckError();\n }\n\n // 1) flag NEW neighbors nodes that I need as 1, they will be in the 2nd ring\n // 2) fill out local indices for 1st ring & internal indices\n int num_halos = halo_row_offsets.size();\n\n for (int i = 0; i < num_halos; i++)\n {\n int size = halo_global_indices[i].size();\n\n if (size > 0)\n {\n halo_local_indices[i].resize(size);\n thrust_wrapper::fill(halo_local_indices[i].begin(), halo_local_indices[i].end(), -1); // fill with -1\n // TODO: launch only on halo rows\n flag_halo_nodes_global<16, 1> <<< num_blocks, 128>>>(halo_row_offsets[i].raw(), halo_row_offsets[i].size() - 1, halo_global_indices[i].raw(), halo_ranges.raw(), neighbor_nodes.raw(), neighbor_offsets.raw(), base_index, index_range, num_neighbors, halo_local_indices[i].raw());\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n // replace all negative values with 0 in neighbor flags\n is_less_than_zero pred;\n amgx::thrust::replace_if(neighbor_nodes.begin(), neighbor_nodes.end(), pred, 0);\n cudaCheckError();\n // fill halo offsets for the current number of ring for new neighbors (it will be of size 0)\n int current_num_neighbors = (halo_offsets.size() - 1) / current_num_rings;;\n int num_rings = current_num_rings + 1;\n IVector_h new_halo_offsets(num_rings * num_neighbors + 1);\n\n for (int j = 0; j < current_num_rings; j++)\n {\n for (int i = 0 ; i <= current_num_neighbors; i++)\n {\n new_halo_offsets[num_neighbors * j + i] = halo_offsets[current_num_neighbors * j + i];\n }\n\n for (int i = current_num_neighbors; i < num_neighbors; i++)\n {\n new_halo_offsets[num_neighbors * j + i + 1] = new_halo_offsets[num_neighbors * j + i];\n }\n }\n\n halo_offsets = new_halo_offsets;\n int ring = current_num_rings;\n int current_num_halo_indices = local_to_global.size();\n //int halo_base = index_range + current_num_halo_indices;\n cudaCheckError();\n\n // compute neighbors nodes indices (in-place) for each neighbor\n for (int i = 0; i < num_neighbors; i++)\n {\n int last_node, num_halo;\n\n if (neighbor_offsets_h[i + 1] != neighbor_offsets_h[i])\n {\n last_node = neighbor_nodes[neighbor_offsets_h[i + 1] - 1];\n thrust_wrapper::exclusive_scan(neighbor_nodes.begin() + neighbor_offsets_h[i], neighbor_nodes.begin() + neighbor_offsets_h[i + 1], neighbor_nodes.begin() + neighbor_offsets_h[i]);\n num_halo = neighbor_nodes[neighbor_offsets_h[i + 1] - 1] + last_node;\n }\n else\n {\n num_halo = 0;\n last_node = 0;\n }\n\n // update halo offsets (L2H)\n halo_offsets[ring * num_neighbors + i + 1] = halo_offsets[ring * num_neighbors + i] + num_halo;\n\n // if size = 0 then we don't need to compute it\n if (boundary_lists.size() > 0)\n {\n // create my boundary lists = list of neighbor inner nodes corresponding to halo numbers 0..M\n // basically this will be our 2-ring B2L_maps\n boundary_lists[i].resize(num_halo);\n int size = neighbor_offsets_h[i + 1] - neighbor_offsets_h[i];\n num_blocks = std::min(4096, (size + 127) / 128);\n\n if (size > 0)\n {\n populate_B2L <<< num_blocks, 128>>>(neighbor_nodes.raw() + neighbor_offsets_h[i], boundary_lists[i].raw(), last_node, size);\n cudaCheckError();\n }\n }\n }\n\n cudaCheckError();\n // compute local indices and new local to global mapping\n int new_num_halo_indices = halo_offsets[num_rings * num_neighbors] - halo_offsets[current_num_rings * num_neighbors];\n local_to_global.resize(current_num_halo_indices + new_num_halo_indices);\n IVector halo_offsets_d(halo_offsets.size());\n amgx::thrust::copy(halo_offsets.begin(), halo_offsets.end(), halo_offsets_d.begin());\n cudaCheckError();\n\n // do this for all ring-1 neighbors\n for (int i = 0; i < num_halos; i++)\n {\n int num_neighbor_rows = halo_row_offsets[i].size() - 1;\n num_blocks = std::min(4096, (num_neighbor_rows + 127) / 128);\n\n if (num_blocks > 0)\n {\n calc_new_halo_mapping_ring2<16> <<< num_blocks, 128>>>(\n halo_row_offsets[i].raw(), num_neighbor_rows, halo_global_indices[i].raw(), // halo rows to process\n halo_ranges.raw(), base_index, index_range, num_neighbors, // ranges and # of neighbors\n halo_offsets_d.raw() + current_num_rings * num_neighbors, neighbor_offsets.raw(), neighbor_nodes.raw(), // halo offsets, neighbor offsets and indices\n local_to_global.raw(), halo_local_indices[i].raw()); // output\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n}\n\n\ntemplate \nvoid DistributedArranger >::create_rings(Matrix_d &A, INDEX_TYPE rings,\n INDEX_TYPE neighbors,\n std::vector &B2L_maps,\n std::vector > &B2L_rings)\n{\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n B2L_rings.resize(neighbors);\n\n for (int i = 0; i < neighbors; i ++)\n {\n B2L_rings[i].resize(rings + 1);\n B2L_rings[i][0] = 0;\n }\n\n /*\n EXAMPLE\n For example matrix. partition 1, the matrix look like:\n A.row_offsets = [0 4 11 15 20]\n A.col_indices = [4 0 1 2\n 4 5 0 1 2 3 7\n 0 1 2 3\n 1 2 3 6 7]\n num_neighbors=2; neighbors = [0 2]\n B2L_rings[[0 2][0 2]] B2L_maps [[0 1][1 3]] (Converted from the following input: btl_sizes = [0 2 4] btl_maps = [0 1| 1 3])\n L2H_maps (and halo_lists) [[4 5][6 7]]\n and we want to map further halo rings\n */\n //Scratch space computation\n for (int i = 0; i < neighbors; i++)\n {\n B2L_rings[i][1] = B2L_maps[i].size();\n }\n\n //flag array\n IVector halo_rings(A.get_num_rows());\n\n for (int i = 0; i < neighbors; i++)\n {\n //\n // For each ring: flag nodes that are neighbors of the previous ring, then remove the nodes of the previous rings from the marking\n // count the number of marked nodes, and append their row indices to the B2L_maps\n //\n for (int ring = 1; ring < rings; ring++)\n {\n size = B2L_rings[i][ring] - B2L_rings[i][ring - 1];\n thrust_wrapper::fill(halo_rings.begin(), halo_rings.begin() + size + 1, 0);\n num_blocks = std::min(4096, (size + 127) / 128);\n //Set 1 in flag array for any column index appreaing in the rows that belong to the previous ring\n find_next_ring <<< num_blocks, 128>>>(B2L_maps[i].raw() + B2L_rings[i][ring - 1], A.row_offsets.raw(), A.col_indices.raw(), size, 0/*base_index*/, /*index_range*/size, halo_rings.raw());\n cudaCheckError();\n /*\n EXAMPLE\n for neighbor ID 0, halo_rings array initially [0 0 0 0], find_next_ring look at rows in the first ring [0 1] and sets halo_rings array positions to 1 for each column index in those rows:\n halo_rings: [1 1 1 1]\n */\n //Remove all flags that belong to rows in any previous rings\n remove_previous_rings <<< num_blocks, 128>>>(B2L_maps[i].raw(), halo_rings.raw(), B2L_rings[i][ring]);\n cudaCheckError();\n /*\n EXAMPLE\n for neighbor ID 0, halo_rings array [1 1 1 1], remove_previous_rings sets halo_rings positions to 0 for each index in previous B2L rings ([0 1] in this case):\n halo_rings: [0 0 1 1]\n */\n //Count, index and append them to B2L_maps and update B2L_rings\n thrust_wrapper::inclusive_scan(halo_rings.begin(), halo_rings.begin() + size + 1, halo_rings.begin());\n B2L_rings[i][ring + 1] = B2L_rings[i][ring] + halo_rings[size];\n B2L_maps[i].resize(B2L_rings[i][ring + 1]);\n\n /*\n EXAMPLE\n after inclusive scan: [0 0 1 2], B2L_rings[0][2] = 2+2, and B2L_maps[0] gets resized to 4\n */\n if (B2L_rings[i][ring] == B2L_rings[i][ring + 1])\n {\n for (int fill = ring; fill <= rings; fill++) { B2L_rings[i][fill] = B2L_rings[i][ring]; }\n\n break;\n }\n\n num_blocks = std::min(4096, (A.get_num_rows() + 127) / 128);\n write_ring_rows <<< num_blocks, 128>>>(halo_rings.raw(), B2L_maps[i].raw() + B2L_rings[i][ring], A.get_num_rows());\n cudaCheckError();\n /*\n EXAMPLE\n appends positions i to B2L_maps, where halo_rings[i-1] != halo_rings[i]\n B2L_maps[0] = [0 1| 2 3] with B2L_rings[0] = [0 2 4]\n */\n }\n\n cudaCheckError();\n }\n}\n\n// CUDA 4.2\ntemplate \nvoid DistributedArranger >::create_rings(Matrix_d &A, INDEX_TYPE rings)\n{\n int neighbors = A.manager->num_neighbors();\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n A.manager->B2L_rings.resize(neighbors * rings + 1);\n\n for (int i = 0; i < neighbors; i ++)\n {\n A.manager->B2L_rings[i].resize(rings + 1);\n A.manager->B2L_rings[i][0] = 0;\n }\n\n //Scratch space computation\n for (int i = 0; i < neighbors; i++)\n {\n A.manager->B2L_rings[i][1] = A.manager->B2L_maps[i].size();\n }\n\n IVector halo_rings(A.get_num_rows());\n\n for (int i = 0; i < neighbors; i++)\n {\n //\n // For each ring: flag nodes that are neighbors of the previous ring, then remove the nodes of the previous ring from the marking\n // count the number of marked nodes, and append their row indices to the B2L_maps\n //\n for (int ring = 1; ring < rings; ring++)\n {\n size = A.manager->B2L_rings[i][ring] - A.manager->B2L_rings[i][ring - 1];\n thrust_wrapper::fill(halo_rings.begin(), halo_rings.begin() + size + 1, 0);\n num_blocks = std::min(4096, (size + 127) / 128);\n find_next_ring <<< num_blocks, 128>>>(A.manager->B2L_maps[i].raw() + A.manager->B2L_rings[i][ring - 1], A.row_offsets.raw(), A.col_indices.raw(), size, 0/*base_index*/, /*index_range*/size, halo_rings.raw());\n cudaCheckError();\n remove_previous_rings <<< num_blocks, 128>>>(A.manager->B2L_maps[i].raw(), halo_rings.raw(), A.manager->B2L_rings[i][ring]);\n cudaCheckError();\n thrust_wrapper::inclusive_scan(halo_rings.begin(), halo_rings.begin() + size + 1, halo_rings.begin());\n A.manager->B2L_rings[i][ring + 1] = A.manager->B2L_rings[i][ring] + halo_rings[size];\n A.manager->B2L_maps[i].resize(A.manager->B2L_rings[i][ring + 1]);\n\n if (A.manager->B2L_rings[i][ring] == A.manager->B2L_rings[i][ring + 1])\n {\n for (int fill = ring; fill <= rings; fill++) { A.manager->B2L_rings[i][fill] = A.manager->B2L_rings[i][ring]; }\n\n break;\n }\n\n num_blocks = std::min(4096, (A.get_num_rows() + 127) / 128);\n write_ring_rows <<< num_blocks, 128>>>(halo_rings.raw(), A.manager->B2L_maps[i].raw() + A.manager->B2L_rings[i][ring], A.get_num_rows());\n }\n\n cudaCheckError();\n }\n}\n\ntemplate \nvoid DistributedArranger >::create_maps(Matrix_d &A, INDEX_TYPE rings,\n INDEX_TYPE neighbors,\n std::vector &B2L_maps,\n std::vector > &B2L_rings, std::vector &boundary_lists)\n{\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n B2L_maps.resize(neighbors);\n\n //\n // Step 4: export n-ring halo matrices\n //\n for (int i = 0; i < neighbors; i++)\n {\n B2L_maps[i] = boundary_lists[i];\n }\n\n create_rings(A, rings, neighbors, B2L_maps, B2L_rings);\n //B2L_rings.resize(neighbors*rings+1);\n}\n\n// need for CUDA 4.2\ntemplate \nvoid DistributedArranger >::create_maps(Matrix_d &A, INDEX_TYPE rings, std::vector &boundary_lists)\n{\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n int neighbors = A.manager->num_neighbors();\n A.manager->B2L_maps.resize(neighbors);\n\n //\n // Step 4: export n-ring halo matrices\n //\n for (int i = 0; i < neighbors; i++)\n {\n A.manager->B2L_maps[i] = boundary_lists[i];\n }\n\n create_rings(A, rings);\n // A.manager->B2L_rings.resize(neighbors*rings+1);\n}\n\ntemplate \nvoid DistributedArranger >::create_halo_rows(Matrix_d &A, INDEX_TYPE rings, INDEX_TYPE neighbors, std::vector &B2L_maps,\n std::vector > &B2L_rings, std::vector **halo_rows_p, std::vector &halo_lists)\n{\n typedef typename MatPrecisionMap::Type ValueTypeA;\n *halo_rows_p = new std::vector >(neighbors);\n std::vector > &halo_rows = **halo_rows_p;\n\n if (this->halo_coloring == LAST) //Create dummy matrices\n {\n for (int i = 0; i < neighbors; i++)\n {\n halo_rows[i].addProps(CSR);\n\n if (A.hasProps(DIAG)) { halo_rows[i].addProps(DIAG); }\n\n halo_rows[i].resize(halo_lists[i]->size(), halo_lists[i]->size(), (A.hasProps(DIAG) ? 0 : halo_lists[i]->size()), A.get_block_dimy(), A.get_block_dimx(), 1);\n\n if (A.hasProps(DIAG))\n {\n thrust_wrapper::fill(halo_rows[i].row_offsets.begin(), halo_rows[i].row_offsets.end(), 0);\n }\n else\n {\n thrust_wrapper::sequence(halo_rows[i].row_offsets.begin(), halo_rows[i].row_offsets.end());\n amgx::thrust::copy(halo_lists[i]->begin(), halo_lists[i]->end(), halo_rows[i].col_indices.begin());\n }\n\n thrust_wrapper::fill(halo_rows[i].values.begin(), halo_rows[i].values.end(), types::util< ValueTypeA >::get_one());\n cudaCheckError();\n }\n\n return;\n }\n\n //Scratch space computation\n int max_size = 0;\n\n for (int i = 0; i < neighbors; i++)\n {\n max_size = max_size > B2L_maps[i].size() ? max_size : B2L_maps[i].size();\n //B2L_rings[i][1]=B2L_maps[i].size(); //I don't know what this was doing here, definitely stupid for rings > 1\n }\n\n IVector matrix_halo_sizes((max_size + 1)*rings);\n\n for (int i = 0; i < neighbors; i++)\n {\n //\n // Write the length of the rows in the order of B2L_maps, then calculate row_offsets\n //\n int size = B2L_rings[i][rings];\n matrix_halo_sizes.resize(size + 1);\n int num_blocks = std::min(4096, (size + 127) / 128);\n write_matrix_rowsize <<< num_blocks, 128>>>(B2L_maps[i].raw(), A.row_offsets.raw(), size, matrix_halo_sizes.raw());\n cudaCheckError();\n thrust_wrapper::exclusive_scan(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, matrix_halo_sizes.begin());\n int nnz_count = matrix_halo_sizes[size];\n //\n // Resize export halo matrix, and copy over the rows\n //\n halo_rows[i].addProps(CSR);\n\n if (A.hasProps(DIAG)) { halo_rows[i].addProps(DIAG); }\n\n halo_rows[i].resize(B2L_maps[i].size(), B2L_maps[i].size(), nnz_count, A.get_block_dimy(), A.get_block_dimx(), 1);\n export_matrix_elements<32> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.values.raw(), A.get_block_size(), B2L_maps[i].raw(), matrix_halo_sizes.raw(), halo_rows[i].values.raw(), A.col_indices.raw(), halo_rows[i].col_indices.raw(), size);\n cudaCheckError();\n amgx::thrust::copy(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, halo_rows[i].row_offsets.begin());\n\n if (A.hasProps(DIAG))\n {\n export_matrix_diagonal <<< num_blocks, 128>>>(A.values.raw() + A.row_offsets[A.get_num_rows()]*A.get_block_size(), A.get_block_size(), B2L_maps[i].raw(), halo_rows[i].values.raw() + halo_rows[i].row_offsets[halo_rows[i].get_num_rows()]*A.get_block_size(), size);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::initialize_manager(Matrix_d &A, Matrix_d &P, int num_owned_coarse_pts)\n{\n // This function initializes the manager of matrix P based on value of num_owned_coarse_pts\n // The following member of the manager are initialized:\n // global_id, num_partitions, comms (copied from manager of matrix A)\n // part_offsets_h, part_offsets, base_index, index_range, num_rows_global, halo_offsets[0] (those are initialized based on value of num_owned_coarse_pts of all partitions in communicator)\n if (P.manager == NULL)\n {\n P.manager = new DistributedManager();\n }\n\n P.setView(OWNED);\n // Copy some information from manager of matrix A\n P.manager->A = &P;\n P.manager->setComms(A.manager->getComms());\n index_type num_parts = A.manager->get_num_partitions();\n index_type my_rank = A.manager->global_id();\n P.manager->set_global_id(my_rank);\n P.manager->set_num_partitions(num_parts);\n // Do a global reduce to obtain part_offsets for RAP\n IVector_h gathered_count;\n A.manager->getComms()->all_gather(num_owned_coarse_pts, gathered_count, num_parts);\n // Compute P.manager->part_offsets\n P.manager->part_offsets_h.resize(num_parts + 1);\n P.manager->part_offsets_h[0] = (int64_t) 0;\n\n for (int i = 1; i < num_parts + 1; i++)\n {\n P.manager->part_offsets_h[i] = P.manager->part_offsets_h[i - 1] + (int64_t) gathered_count[i - 1];\n }\n\n P.manager->part_offsets = P.manager->part_offsets_h;\n // Compute P.manager->halo_ranges\n P.manager->set_base_index(P.manager->part_offsets_h[my_rank]);\n P.manager->set_index_range(num_owned_coarse_pts);\n P.manager->num_rows_global = P.manager->part_offsets_h[num_parts];\n // Set halo_offsets[0]\n P.manager->halo_offsets.resize(1);\n P.manager->halo_offsets[0] = num_owned_coarse_pts;\n}\n\ntemplate \nvoid DistributedArranger >::initialize_manager_from_global_col_indices(Matrix_d &P, I64Vector_d &P_col_indices_global)\n{\n // This function finds the number of neighbors of P\n // fills halo_offsets, local_to_global_map, B2L_maps based on global col indices of matrix P\n int num_owned_coarse_pts = P.manager->halo_offsets[0];\n // By checking the global col indices, identify neighbors\n P.manager->neighbors.resize(0);\n this->create_neighbors_v2_global(P, P_col_indices_global);\n int num_neighbors = P.manager->neighbors.size();\n P.manager->getComms()->set_neighbors(num_neighbors);\n P.manager->halo_offsets.resize(num_neighbors + 1, 0);\n int total_rows_of_neighbors = 0;\n std::vector flag_offsets_h(num_neighbors + 1, 0);\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n int num_rows = (int) (P.manager->halo_ranges_h[2 * i + 1] - P.manager->halo_ranges_h[2 * i]);\n total_rows_of_neighbors += num_rows;\n flag_offsets_h[i + 1] = flag_offsets_h[i] + num_rows;\n }\n\n // Copy to device\n IVector flag_offsets_d(num_neighbors + 1);\n amgx::thrust::copy(flag_offsets_h.begin(), flag_offsets_h.end(), flag_offsets_d.begin());\n // Allocate halo_nodes array that will be used to flag halo nodes\n // Flag_offsets array can be used to determine where each neighbor begins\n IVector halo_nodes(total_rows_of_neighbors + 1);\n thrust_wrapper::fill(halo_nodes.begin(), halo_nodes.end(), 0);\n cudaCheckError();\n int num_blocks = std::min(4096, (P.get_num_rows() + 127) / 128);\n\n if (P.get_num_rows() > 0)\n {\n flag_halo_nodes_global_v2<16> <<< num_blocks, 128>>>(P.row_offsets.raw(), P_col_indices_global.raw(), P.manager->halo_ranges.raw(), halo_nodes.raw(), flag_offsets_d.raw(), P.manager->base_index(), P.manager->index_range(), num_neighbors, P.get_num_rows());\n cudaCheckError();\n }\n\n // Step 3: Write the 1-ring global nodes for each neighbor\n thrust_wrapper::exclusive_scan(halo_nodes.begin(), halo_nodes.end(), halo_nodes.begin());\n cudaCheckError();\n std::vector bdy_lists;\n bdy_lists.resize(num_neighbors);\n int count = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int num_halo = halo_nodes[flag_offsets_h[i + 1]] - count;\n P.manager->halo_offsets[i + 1] = P.manager->halo_offsets[i] + num_halo;\n bdy_lists[i].resize(num_halo);\n\n if (num_halo > 0)\n {\n int size = flag_offsets_h[i + 1] - flag_offsets_h[i];\n int num_blocks = std::min(4096, (size + 127) / 128);\n populate_B2L_v2 <<< num_blocks, 128>>>(halo_nodes.raw() + flag_offsets_h[i], bdy_lists[i].raw(), count, size);\n cudaCheckError();\n }\n\n count += num_halo;\n }\n\n cudaCheckError();\n\n // Fill the col_indices\n if (P.get_num_rows() > 0)\n {\n renumber_col_indices_global<16> <<< num_blocks, 128>>>(P.row_offsets.raw(), P.col_indices.raw(), P_col_indices_global.raw(), P.manager->halo_ranges.raw(), halo_nodes.raw(), flag_offsets_d.raw(), P.manager->base_index(), P.manager->index_range(), num_neighbors, P.get_num_rows(), P.manager->global_id());\n cudaCheckError();\n }\n\n // Create P.local_to_global_map\n P.manager->local_to_global_map.resize(count);\n count = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int num_halo = bdy_lists[i].size();\n\n if (num_halo > 0)\n {\n int num_blocks = std::min(4096, (num_halo + 127) / 128);\n fill_local_to_global_map <<< num_blocks, 128>>>(bdy_lists[i].raw(), P.manager->local_to_global_map.raw() + count, P.manager->halo_ranges_h[2 * i], num_halo);\n cudaCheckError();\n }\n\n count += num_halo;\n }\n\n cudaCheckError();\n P.manager->getComms()->exchange_vectors(bdy_lists, P, 0);\n // copy boundary lists to B2L maps\n // release internal boundary lists array\n P.manager->B2L_maps.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n P.manager->B2L_maps[i] = bdy_lists[i];\n }\n\n P.manager->B2L_rings.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n P.manager->B2L_rings[i].resize(2);\n P.manager->B2L_rings[i][0] = 0;\n P.manager->B2L_rings[i][1] = P.manager->B2L_maps[i].size();\n }\n\n this->initialize_B2L_maps_offsets(P, 1);\n P.manager->_num_interior_nodes = P.get_num_rows();\n P.manager->_num_boundary_nodes = 0;\n P.manager->set_initialized(P.row_offsets);\n int size_one_ring, offset;\n P.manager->getOffsetAndSizeForView(FULL, &offset, &size_one_ring);\n P.setView(OWNED);\n}\n\n\ntemplate \nvoid DistributedArranger >::createTempManager(Matrix_d &B, Matrix_d &A, IVector &row_offsets)\n{\n // Initialize P.manager\n if (B.manager == NULL)\n {\n B.manager = new DistributedManager();\n }\n\n index_type num_parts = A.manager->get_num_partitions();\n index_type my_rank = A.manager->global_id();\n B.manager->A = &B;\n B.manager->setComms(A.manager->getComms());\n B.manager->set_global_id(my_rank);\n B.manager->set_num_partitions(num_parts);\n B.manager->neighbors = A.manager->neighbors;\n int num_neighbors = A.manager->neighbors.size();\n B.manager->getComms()->set_neighbors(num_neighbors);\n // --------------------------------------\n // Create B.manager->B2L_maps\n // --------------------------------------\n B.manager->B2L_maps.resize(num_neighbors);\n int max_B2L_maps_size = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int A_b2l_map_size = A.manager->B2L_rings[i][1];\n\n if (A_b2l_map_size > max_B2L_maps_size)\n {\n max_B2L_maps_size = A_b2l_map_size;\n }\n }\n\n IVector row_sizes(row_offsets.size() - 1);\n thrust_wrapper::adjacent_difference(row_offsets.begin() + 1, row_offsets.end(), row_sizes.begin());\n cudaCheckError();\n IVector B2L_maps_offsets(max_B2L_maps_size + 1);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int A_b2l_map_size = A.manager->B2L_rings[i][1];\n\n if (A_b2l_map_size > 0)\n {\n // For each node in the B2L_map for that neighbor, write the row size in B2L_maps_offsets\n thrust_wrapper::gather(A.manager->B2L_maps[i].begin(), A.manager->B2L_maps[i].begin() + A_b2l_map_size, row_sizes.begin(), B2L_maps_offsets.begin());\n // Do exclusive scan, to go from row length to row offsets\n thrust_wrapper::exclusive_scan(B2L_maps_offsets.begin(), B2L_maps_offsets.begin() + A_b2l_map_size + 1, B2L_maps_offsets.begin());\n B.manager->B2L_maps[i].resize(B2L_maps_offsets[A_b2l_map_size]);\n const int cta_size = 128;\n const int nWarps = cta_size / 32;\n int grid_size = std::min( 4096, (int) (A_b2l_map_size + nWarps - 1) / nWarps);\n // Populate the B2L_map array\n fill_B_B2L_map <<< grid_size, cta_size >>>(\n B2L_maps_offsets.raw(),\n A.manager->B2L_maps[i].raw(),\n row_offsets.raw(),\n B.manager->B2L_maps[i].raw(),\n A_b2l_map_size,\n A.manager->global_id());\n cudaCheckError();\n }\n else\n {\n B.manager->B2L_maps[i].resize(0);\n }\n\n cudaCheckError();\n }\n\n B.manager->B2L_rings.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n B.manager->B2L_rings[i].resize(2);\n B.manager->B2L_rings[i][0] = 0;\n B.manager->B2L_rings[i][1] = B.manager->B2L_maps[i].size();\n }\n\n // Initialize the B2L_maps_offsets (required by gather_v2 and scatter_v2) from the B2L_rings values\n this->initialize_B2L_maps_offsets(B, 1);\n B.manager->halo_offsets.resize(num_neighbors + 1);\n\n // Fill halo_offsets\n for (int i = 0; i < num_neighbors + 1; i++)\n {\n B.manager->halo_offsets[i] = row_offsets[A.manager->halo_offsets[i]];\n }\n}\n\n\ntemplate \nvoid DistributedArranger >::exchange_halo_rows_P(Matrix_d &A, Matrix_d &P, I64Vector &RAP_local_to_global_map, IVector_h &P_neighbors, I64Vector_h &P_halo_ranges_h, I64Vector &P_halo_ranges, IVector_h &P_halo_offsets, I64Vector_h &RAP_part_offsets_h, I64Vector &RAP_part_offsets, index_type num_owned_coarse_pts, int64_t coarse_base_index)\n{\n if (P.hasProps(DIAG) || P.get_block_size() != 1)\n {\n FatalError(\"P with external diagonal or block_size != 1 not supported\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n int num_neighbors = A.manager->num_neighbors() ;\n std::vector halo_rows_P_row_offsets;\n std::vector halo_rows_P_col_indices;\n std::vector halo_rows_P_local_col_indices;\n std::vector halo_rows_P_values;\n // In this function, store in halo_rows_P_row_offsets, halo_rows_P_col_indices and halo_rows_P_values, the rows of P that need to be sent to each neighbors\n // halo_rows_P_col_indices stores global indices\n this->pack_halo_rows_P(A, P, halo_rows_P_row_offsets, halo_rows_P_col_indices, halo_rows_P_values, RAP_local_to_global_map, num_owned_coarse_pts, coarse_base_index);\n // Do the exchange with the neighbors\n // On return, halo_rows_P_rows_offsets, halo_rows_P_col_indices and halo_rows_P_values stores the rows of P received from each neighbor (rows needed to perform A*P)\n std::vector dummy_halo_ids(0);\n A.manager->getComms()->exchange_matrix_halo(halo_rows_P_row_offsets, halo_rows_P_col_indices, halo_rows_P_values, dummy_halo_ids, A.manager->neighbors, A.manager->global_id());\n halo_rows_P_local_col_indices.resize(halo_rows_P_col_indices.size());\n\n for (int i = 0; i < halo_rows_P_col_indices.size(); i++)\n {\n halo_rows_P_local_col_indices[i].resize(halo_rows_P_col_indices[i].size());\n }\n\n // Update the list of neighbors \"P_neighbors\" and the corresponding ranges, offsets\n update_neighbors_list(A, P_neighbors, P_halo_ranges_h, P_halo_ranges, RAP_part_offsets_h, RAP_part_offsets, halo_rows_P_row_offsets, halo_rows_P_col_indices);\n std::vector dummy_boundary_list(0);\n int current_num_rings = 1;\n halo_rows_P_local_col_indices.resize(halo_rows_P_col_indices.size());\n // Convert the global indices in the rows just received to local indices\n // This function updates halo offsets, halo_rows_P_local_col_indices, local_to_global_map\n // This should not modify the existing P manager\n this->compute_local_halo_indices( P.row_offsets, P.col_indices, halo_rows_P_row_offsets, halo_rows_P_col_indices, halo_rows_P_local_col_indices, RAP_local_to_global_map, dummy_boundary_list, P_neighbors, P_halo_ranges_h, P_halo_ranges, P_halo_offsets, coarse_base_index, num_owned_coarse_pts, P.get_num_rows(), current_num_rings);\n // Append the new rows to the matrix P\n P.set_initialized(0);\n this->append_halo_rows(P, halo_rows_P_row_offsets, halo_rows_P_local_col_indices, halo_rows_P_values);\n P.set_initialized(1);\n}\n\n\ntemplate \nvoid DistributedArranger >::append_halo_rows(Matrix_d &A, std::vector &halo_row_offsets, std::vector &halo_local_indices, std::vector &halo_values )\n{\n int new_num_rows = A.get_num_rows();\n int new_num_nnz = A.row_offsets[A.get_num_rows()];\n int cur_row = A.get_num_rows();\n int cur_offset = new_num_nnz;\n int num_neighbors = halo_row_offsets.size();\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = halo_row_offsets[i].size();\n\n if (size != 0)\n {\n new_num_rows += halo_row_offsets[i].size() - 1;\n new_num_nnz += halo_local_indices[i].size();\n }\n }\n\n A.resize(new_num_rows, new_num_rows, new_num_nnz, A.get_block_dimx(), A.get_block_dimy(), 1);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int num_halo_rows = halo_row_offsets[i].size() - 1;\n\n if (num_halo_rows > 0)\n {\n // update halo row offsets in-place\n thrust_wrapper::transform(halo_row_offsets[i].begin(), halo_row_offsets[i].end(), amgx::thrust::constant_iterator(cur_offset), halo_row_offsets[i].begin(), amgx::thrust::plus());\n // insert halo rows\n amgx::thrust::copy(halo_row_offsets[i].begin(), halo_row_offsets[i].end() - 1, A.row_offsets.begin() + cur_row);\n amgx::thrust::copy(halo_local_indices[i].begin(), halo_local_indices[i].end(), A.col_indices.begin() + cur_offset);\n amgx::thrust::copy(halo_values[i].begin(), halo_values[i].end(), A.values.begin() + cur_offset);\n // update counters\n cur_offset = halo_row_offsets[i][num_halo_rows];\n cur_row += num_halo_rows;\n }\n }\n\n cudaCheckError();\n A.row_offsets[A.get_num_rows()] = cur_offset;\n int num_cols = -1;\n num_cols = thrust_wrapper::reduce(A.col_indices.begin(), A.col_indices.end(), num_cols, amgx::thrust::maximum()) + 1;\n cudaCheckError();\n A.set_num_cols(num_cols);\n}\n\n\ntemplate \nvoid DistributedArranger >::exchange_RAP_ext(Matrix_d &RAP, Matrix_d &RAP_full, Matrix_d &A, Matrix_d &P, IVector_h &P_halo_offsets, I64Vector &RAP_local_to_global_map, IVector_h &P_neighbors, I64Vector_h &P_halo_ranges_h, I64Vector &P_halo_ranges, I64Vector_h &RAP_part_offsets_h, I64Vector &RAP_part_offsets, index_type num_owned_coarse_pts, int64_t coarse_base_index, void *wk)\n{\n if (RAP_full.hasProps(DIAG) || RAP_full.get_block_size() != 1)\n {\n FatalError(\"RAP_full with external diagonal or block_size != 1 not supported\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // --------------------------------------------------------------------------------------\n // Pack the rows of RAP_full that need to be sent to neighbors, using the manager of P\n // --------------------------------------------------------------------------------------\n std::vector halo_rows_RAP_row_offsets;\n std::vector halo_rows_RAP_row_ids;\n std::vector halo_rows_RAP_col_indices;\n std::vector halo_rows_RAP_values;\n this->pack_halo_rows_RAP(RAP_full, P, halo_rows_RAP_row_offsets, halo_rows_RAP_col_indices, halo_rows_RAP_values, halo_rows_RAP_row_ids, RAP_local_to_global_map);\n // Send the rows to the neighbors, receive RAP rows from neighbors\n // Col indices received from neighbors are global indices\n // We also receive from neighbors the global row ids\n P.manager->getComms()->exchange_matrix_halo(halo_rows_RAP_row_offsets, halo_rows_RAP_col_indices, halo_rows_RAP_values, halo_rows_RAP_row_ids, P.manager->neighbors, P.manager->global_id());\n std::vector halo_rows_RAP_local_col_indices;\n halo_rows_RAP_local_col_indices.resize(halo_rows_RAP_col_indices.size());\n\n for (int i = 0; i < halo_rows_RAP_col_indices.size(); i++)\n {\n halo_rows_RAP_local_col_indices[i].resize(halo_rows_RAP_col_indices[i].size());\n }\n\n // This function finds the new number of neighbors, renumber halo_cols with local_ids, modifies halo_offsets\n update_neighbors_list(A, P_neighbors, P_halo_ranges_h, P_halo_ranges, RAP_part_offsets_h, RAP_part_offsets, halo_rows_RAP_row_offsets, halo_rows_RAP_col_indices);\n // Now renumber find local indices for global col indices, update local_to_global\n std::vector dummy_boundary_list(0);\n int current_num_rings = 2;\n this->compute_local_halo_indices(RAP_full.row_offsets, RAP_full.col_indices, halo_rows_RAP_row_offsets, halo_rows_RAP_col_indices, halo_rows_RAP_local_col_indices, RAP_local_to_global_map, dummy_boundary_list, P_neighbors, P_halo_ranges_h, P_halo_ranges, P_halo_offsets, coarse_base_index, num_owned_coarse_pts, num_owned_coarse_pts, current_num_rings );\n // Now find local row ids from global ids of rows received\n std::vector halo_rows_RAP_local_row_ids;\n halo_rows_RAP_local_row_ids.resize(halo_rows_RAP_row_ids.size());\n\n for (int i = 0; i < halo_rows_RAP_row_ids.size(); i++)\n {\n int size = halo_rows_RAP_row_ids[i].size();\n halo_rows_RAP_local_row_ids[i].resize(size);\n\n if (size != 0)\n {\n int cta_size = 128;\n int grid_size = std::min( 4096, (size + cta_size - 1) / cta_size);\n fill_halo_rows_row_ids <<< grid_size, cta_size >>>( halo_rows_RAP_row_ids[i].raw(), halo_rows_RAP_local_row_ids[i].raw(), coarse_base_index, num_owned_coarse_pts, size);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n // Add the rows received from neighbors to owned rows in RAP_full and create RAP matrix\n RAP.set_initialized( 0 );\n RAP.set_num_rows(num_owned_coarse_pts);\n RAP.set_num_cols(num_owned_coarse_pts + RAP.manager->local_to_global_map.size());\n RAP.row_offsets.resize(num_owned_coarse_pts + 1);\n RAP_full.setView(OWNED);\n CSR_Multiply::csr_RAP_sparse_add( RAP, RAP_full, halo_rows_RAP_row_offsets, halo_rows_RAP_local_col_indices, halo_rows_RAP_values, halo_rows_RAP_local_row_ids, wk );\n RAP.set_initialized( 1 );\n}\n\n\ntemplate \nvoid DistributedArranger >::pack_halo_rows_P(Matrix_d &A, Matrix_d &P, std::vector &halo_rows_P_row_offsets, std::vector &halo_rows_P_col_indices, std::vector &halo_rows_P_values, I64Vector &RAP_local_to_global_map, index_type num_owned_coarse_pts, int64_t coarse_base_index)\n{\n int num_neighbors = A.manager->num_neighbors();\n halo_rows_P_row_offsets.resize(num_neighbors);\n halo_rows_P_col_indices.resize(num_neighbors);\n halo_rows_P_values.resize(num_neighbors);\n int num_rings_to_send = 1;\n //Scratch space computation\n int max_size = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n max_size = max_size > A.manager->B2L_rings[i][num_rings_to_send] ? max_size : A.manager->B2L_rings[i][num_rings_to_send];\n }\n\n IVector matrix_halo_sizes(max_size + 1);\n\n // Here only using the 1-ring of matrix A\n for (int i = 0; i < num_neighbors; i++)\n {\n // Write the length of the rows in the order of B2L_maps, then calculate row_offsets\n int size = A.manager->B2L_rings[i][num_rings_to_send];\n\n if (size != 0)\n {\n //matrix_halo_sizes.resize(size+1);\n int num_blocks = std::min(4096, (size + 127) / 128);\n write_matrix_rowsize <<< num_blocks, 128>>>(A.manager->B2L_maps[i].raw(), P.row_offsets.raw(), size, matrix_halo_sizes.raw());\n cudaCheckError();\n thrust_wrapper::exclusive_scan(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, matrix_halo_sizes.begin());\n int nnz_count = matrix_halo_sizes[size];\n // Resize export halo matrix, and copy over the rows\n halo_rows_P_row_offsets[i].resize(size + 1);\n halo_rows_P_col_indices[i].resize(nnz_count);\n halo_rows_P_values[i].resize(nnz_count);\n /* WARNING: Since A is reordered (into interior and boundary nodes), while R and P are not reordered,\n you must unreorder A when performing R*A*P product in ordre to obtain the correct result. */\n export_matrix_elements_global<32> <<< num_blocks, 128>>>(P.row_offsets.raw(), P.values.raw(), P.get_block_size(), A.manager->B2L_maps[i].raw(), matrix_halo_sizes.raw(), halo_rows_P_values[i].raw(), P.col_indices.raw(), halo_rows_P_col_indices[i].raw(), size, RAP_local_to_global_map.raw(), NULL /*A.manager->inverse_renumbering.raw()*/, num_owned_coarse_pts, coarse_base_index);\n cudaCheckError();\n amgx::thrust::copy(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, halo_rows_P_row_offsets[i].begin());\n }\n else\n {\n halo_rows_P_row_offsets[i].resize(0);\n halo_rows_P_col_indices[i].resize(0);\n halo_rows_P_values[i].resize(0);\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::pack_halo_rows_RAP(Matrix_d &RAP, Matrix_d &P, std::vector &halo_rows_RAP_row_offsets, std::vector &halo_rows_RAP_col_indices, std::vector &halo_rows_RAP_values, std::vector &halo_rows_RAP_row_ids, I64Vector &RAP_local_to_global_map)\n{\n int num_neighbors = P.manager->neighbors.size();\n int num_owned_coarse_pts = P.manager->halo_offsets[0];\n int64_t P_col_base_index = P.manager->base_index();\n halo_rows_RAP_row_offsets.resize(num_neighbors);\n halo_rows_RAP_row_ids.resize(num_neighbors);\n halo_rows_RAP_col_indices.resize(num_neighbors);\n halo_rows_RAP_values.resize(num_neighbors);\n // Create tmp_B2L_maps that contains the list of rows to send\n std::vector tmp_B2L_maps(num_neighbors);\n\n // Fill B2L_maps with list of halo indices for all rings to send\n for (int i = 0; i < num_neighbors; i++)\n {\n tmp_B2L_maps[i].resize(0);\n halo_rows_RAP_row_ids[i].resize(0);\n int begin = P.manager->halo_offsets[i];\n int end = P.manager->halo_offsets[i + 1];\n\n if (begin > RAP.row_offsets.size() - 1)\n {\n break;\n }\n\n int local_to_global_begin = begin - P.manager->halo_offsets[0];\n int size = end - begin;\n\n if (end > RAP.row_offsets.size() - 1)\n {\n size = RAP.row_offsets.size() - 1 - begin;\n }\n\n if (size > 0)\n {\n tmp_B2L_maps[i].resize(size);\n halo_rows_RAP_row_ids[i].resize(size);\n thrust_wrapper::sequence(tmp_B2L_maps[i].begin(), tmp_B2L_maps[i].begin() + size, begin);\n amgx::thrust::copy(P.manager->local_to_global_map.begin() + (int64_t) local_to_global_begin, P.manager->local_to_global_map.begin() + (int64_t) local_to_global_begin + (int64_t) size, halo_rows_RAP_row_ids[i].begin());\n }\n }\n\n cudaCheckError();\n //Scratch space computation\n int max_size = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = tmp_B2L_maps[i].size();\n max_size = max_size > size ? max_size : size;\n }\n\n IVector matrix_halo_sizes(max_size + 1);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n // Write the length of the rows in the order of B2L_maps, then calculate row_offsets\n int size = tmp_B2L_maps[i].size();\n\n if (size > 0)\n {\n //matrix_halo_sizes.resize(size+1);\n int num_blocks = std::min(4096, (size + 127) / 128);\n write_matrix_rowsize <<< num_blocks, 128>>>(tmp_B2L_maps[i].raw(), RAP.row_offsets.raw(), size, matrix_halo_sizes.raw());\n cudaCheckError();\n thrust_wrapper::exclusive_scan(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, matrix_halo_sizes.begin());\n int nnz_count = matrix_halo_sizes[size];\n // Resize export halo matrix, and copy over the rows\n halo_rows_RAP_row_offsets[i].resize(tmp_B2L_maps[i].size() + 1);\n halo_rows_RAP_col_indices[i].resize(nnz_count);\n halo_rows_RAP_values[i].resize(nnz_count);\n export_matrix_elements_global<32> <<< num_blocks, 128>>>(RAP.row_offsets.raw(), RAP.values.raw(), RAP.get_block_size(), tmp_B2L_maps[i].raw(), matrix_halo_sizes.raw(), halo_rows_RAP_values[i].raw(), RAP.col_indices.raw(), halo_rows_RAP_col_indices[i].raw(), size, RAP_local_to_global_map.raw(), NULL, num_owned_coarse_pts, P_col_base_index);\n cudaCheckError();\n amgx::thrust::copy(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, halo_rows_RAP_row_offsets[i].begin());\n }\n else\n {\n halo_rows_RAP_row_offsets[i].resize(0);\n halo_rows_RAP_col_indices[i].resize(0);\n halo_rows_RAP_values[i].resize(0);\n }\n }\n\n cudaCheckError();\n}\n\n\ntemplate \nvoid DistributedArranger >::sparse_matrix_add(Matrix_d &RAP, std::vector &halo_rows_RAP_row_offsets, std::vector &halo_rows_RAP_col_indices, std::vector &halo_rows_RAP_values, std::vector halo_rows_RAP_row_ids, I64Vector_d &P_local_to_global_map, index_type P_col_num_owned_pts, int64_t P_col_base_index)\n{\n // Look at csr_galerkin_product implementation, use hash\n}\n\n\n// CUDA 4.2\ntemplate \nvoid DistributedArranger >::create_halo_rows(Matrix_d &A,\n std::vector &halo_lists\n )\n{\n int neighbors = A.manager->num_neighbors();\n\n if (neighbors == 0)\n {\n return;\n }\n\n int rings = (A.manager->B2L_rings.size() - 1) / neighbors;\n A.manager->halo_rows = new std::vector >(neighbors);\n std::vector > &halo_rows = *A.manager->halo_rows;\n\n if (this->halo_coloring == LAST) //Create dummy matrices, only containing diagonal elements\n {\n for (int i = 0; i < neighbors; i++)\n {\n halo_rows[i].addProps(CSR);\n\n if (A.hasProps(DIAG)) { halo_rows[i].addProps(DIAG); }\n\n halo_rows[i].resize(halo_lists[i]->size(), halo_lists[i]->size(), (A.hasProps(DIAG) ? 0 : halo_lists[i]->size()), A.get_block_dimy(), A.get_block_dimx(), 1);\n\n if (A.hasProps(DIAG))\n {\n thrust_wrapper::fill(halo_rows[i].row_offsets.begin(), halo_rows[i].row_offsets.end(), 0);\n }\n else\n {\n thrust_wrapper::sequence(halo_rows[i].row_offsets.begin(), halo_rows[i].row_offsets.end());\n amgx::thrust::copy(halo_lists[i]->begin(), halo_lists[i]->end(), halo_rows[i].col_indices.begin());\n }\n\n cudaCheckError();\n }\n\n return;\n }\n\n //Scratch space computation\n int max_size = 0;\n\n for (int i = 0; i < neighbors; i++)\n {\n max_size = max_size > A.manager->B2L_maps[i].size() ? max_size : A.manager->B2L_maps[i].size();\n //B2L_rings[i][1]=B2L_maps[i].size(); //I don't know what this was doing here, definitely stupid for rings > 1\n }\n\n IVector matrix_halo_sizes((max_size + 1)*rings);\n\n for (int i = 0; i < neighbors; i++)\n {\n //\n // Write the length of the rows in the order of B2L_maps, then calculate row_offsets\n //\n int size = A.manager->B2L_rings[i][rings];\n matrix_halo_sizes.resize(size + 1);\n int num_blocks = std::min(4096, (size + 127) / 128);\n write_matrix_rowsize <<< num_blocks, 128>>>(A.manager->B2L_maps[i].raw(), A.row_offsets.raw(), size, matrix_halo_sizes.raw());\n cudaCheckError();\n thrust_wrapper::exclusive_scan(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, matrix_halo_sizes.begin());\n cudaCheckError();\n int nnz_count = matrix_halo_sizes[size];\n //\n // Resize export halo matrix, and copy over the rows\n //\n halo_rows[i].addProps(CSR);\n\n if (A.hasProps(DIAG)) { halo_rows[i].addProps(DIAG); }\n\n halo_rows[i].resize(A.manager->B2L_maps[i].size(), A.manager->B2L_maps[i].size(), nnz_count, A.get_block_dimy(), A.get_block_dimx(), 1);\n export_matrix_elements<32> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.values.raw(), A.get_block_size(), A.manager->B2L_maps[i].raw(), matrix_halo_sizes.raw(), halo_rows[i].values.raw(), A.col_indices.raw(), halo_rows[i].col_indices.raw(), size);\n cudaCheckError();\n amgx::thrust::copy(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, halo_rows[i].row_offsets.begin());\n cudaCheckError();\n\n if (A.hasProps(DIAG))\n {\n export_matrix_diagonal <<< num_blocks, 128>>>(A.values.raw() + A.row_offsets[A.get_num_rows()]*A.get_block_size(), A.get_block_size(), A.manager->B2L_maps[i].raw(), halo_rows[i].values.raw() + halo_rows[i].row_offsets[halo_rows[i].get_num_rows()]*A.get_block_size(), size);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::create_halo_rows_global_indices(Matrix_d &A,\n std::vector &halo_rows_row_offsets,\n std::vector &halo_rows_col_indices,\n std::vector &halo_rows_values)\n{\n // in this function we assume that we only have 1-ring so far\n int neighbors = A.manager->num_neighbors();\n\n for (int i = 0; i < neighbors; i++)\n {\n // compute row offsets and nnz count\n int size = A.manager->B2L_maps[i].size();\n halo_rows_row_offsets[i].resize(size + 1);\n\n if (size > 0)\n {\n int num_blocks = std::min(4096, (size + 127) / 128);\n write_matrix_rowsize <<< num_blocks, 128>>>(A.manager->B2L_maps[i].raw(), A.row_offsets.raw(), size, halo_rows_row_offsets[i].raw());\n cudaCheckError();\n thrust_wrapper::exclusive_scan(halo_rows_row_offsets[i].begin(), halo_rows_row_offsets[i].begin() + size + 1, halo_rows_row_offsets[i].begin());\n // compute global indices\n int nnz_count = halo_rows_row_offsets[i][size];\n halo_rows_col_indices[i].resize(nnz_count);\n halo_rows_values[i].resize(nnz_count);\n export_matrix_elements_global<32> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.values.raw(), A.get_block_size(), A.manager->B2L_maps[i].raw(), halo_rows_row_offsets[i].raw(), halo_rows_values[i].raw(), A.col_indices.raw(), halo_rows_col_indices[i].raw(), size, A.manager->local_to_global_map.raw(), NULL, A.get_num_rows(), A.manager->base_index());\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::create_halo_btl(Matrix_d &A, INDEX_TYPE rings, INDEX_TYPE neighbors,\n Vector &neighbors_list,\n int64_t base_index, INDEX_TYPE index_range,\n INDEX_TYPE global_id, std::vector &B2L_maps,\n std::vector > &B2L_rings,\n std::vector **halo_btl_p,\n\n std::vector &halo_lists\n )\n{\n *halo_btl_p = new std::vector >(neighbors);\n std::vector > &halo_btl = **halo_btl_p;\n /*\n //calling constructors on the allocated memory\n for(int i=0; i();\n */\n\n if (this->halo_coloring != LAST)\n {\n //If we are sending matrices over\n for (int i = 0; i < neighbors; i++)\n {\n //\n // Transfer indexing data too (global indices of exported rows, and sizes of the halo rings)\n //\n /*\n EXAMPLE\n For example partition 1, with 2 rings, and the following B2L and L2H (called halo_lists here) info\n num_neighbors=2; neighbors = [0 2]\n B2L_rings[[0 2 4][0 2 4]] B2L_maps[[0 1| 2 3][1 3| 0 2]]\n L2H_maps (halo_lists) [[4 5][6 7]]\n */\n int size = B2L_rings[i][rings];\n halo_btl[i].set_global_id(global_id);\n halo_btl[i].set_base_index(0);\n halo_btl[i].resize(1, rings);\n halo_btl[i].B2L_maps[0].resize(size);\n halo_btl[i].L2H_maps.resize(1);\n halo_btl[i].L2H_maps[0].resize(halo_lists[i]->size());\n //send over B2L and L2H maps corresponding to current neighbor, so that it will be able to make sense of indices (local to this partition)\n amgx::thrust::copy(B2L_maps[i].begin(), B2L_maps[i].end(), halo_btl[i].B2L_maps[0].begin());\n amgx::thrust::copy(halo_lists[i]->begin(), halo_lists[i]->end(), halo_btl[i].L2H_maps[0].begin());\n int max_index = thrust_wrapper::reduce(halo_lists[i]->begin(), halo_lists[i]->end(), (int)0, amgx::thrust::maximum()) + 1;\n max_index = max_index > A.get_num_rows() ? max_index : A.get_num_rows();\n halo_btl[i].set_index_range(max_index);\n amgx::thrust::copy(B2L_rings[i].begin(), B2L_rings[i].end(), halo_btl[i].B2L_rings[0].begin());\n /*\n EXAMPLE\n These fields are going to be:\n halo_btl[0] (for neighbor ID 0)\n global_id = 1; base_index=0; index_range=6 B2L_rings[0] = [0 2 4] B2L_maps[0] = [0 1| 2 3] L2H_maps = [4 5]\n halo_btl[1] (for neighbor ID 2)\n global_id = 1; base_index=0; index_range=8 B2L_rings[0] = [0 2 4] B2L_maps[0] = [1 3| 0 2] L2H_maps = [6 7]\n */\n }\n }\n else\n {\n if (rings > 1) { FatalError(\"halo_coloring=LAST dummy halo rows creation not implemented\", AMGX_ERR_NOT_IMPLEMENTED); }\n\n //Create dummy halo matrices\n for (int i = 0; i < neighbors; i++)\n {\n int size = halo_lists[i]->size();\n halo_btl[i].set_global_id(neighbors_list[i]);\n halo_btl[i].set_base_index(0);\n halo_btl[i].resize(1, rings);\n halo_btl[i].B2L_maps[0].resize(size);\n halo_btl[i].L2H_maps.resize(1);\n halo_btl[i].L2H_maps[0].resize(B2L_rings[i][rings]);\n //Even though in this case matrices are not being sent over, we create the same data structures, so they can be handled the same way\n amgx::thrust::copy(B2L_maps[i].begin(), B2L_maps[i].end(), halo_btl[i].L2H_maps[0].begin());\n amgx::thrust::copy(halo_lists[i]->begin(), halo_lists[i]->end(), halo_btl[i].B2L_maps[0].begin());\n int max_index = thrust_wrapper::reduce(halo_lists[i]->begin(), halo_lists[i]->end(), (int)0, amgx::thrust::maximum()) + 1;\n max_index = max_index > A.get_num_rows() ? max_index : A.get_num_rows();\n halo_btl[i].set_index_range(max_index);\n halo_btl[i].B2L_rings[0][0] = 0;\n halo_btl[i].B2L_rings[0][1] = size;\n }\n }\n\n cudaCheckError();\n}\n\n// CUDA 4.2\ntemplate \nvoid DistributedArranger >::create_halo_btl(Matrix_d &A,\n std::vector &halo_lists)\n{\n int neighbors = A.manager->num_neighbors();\n\n if (neighbors == 0)\n {\n return;\n }\n\n int rings = (A.manager->B2L_rings.size() - 1) / neighbors;\n A.manager->halo_btl = new std::vector >(neighbors);\n std::vector > &halo_btl = *A.manager->halo_btl;\n\n if (this->halo_coloring != LAST)\n {\n for (int i = 0; i < neighbors; i++)\n {\n //\n // Transfer indexing data too (global indices of exported rows, and sizes of the halo rings)\n //\n int size = A.manager->B2L_rings[i][rings];\n halo_btl[i].set_global_id(A.manager->global_id());\n halo_btl[i].set_base_index(0);\n //Adding one more ring to store halo numbering\n halo_btl[i].resize(1, rings);\n halo_btl[i].B2L_maps[0].resize(size);\n halo_btl[i].L2H_maps.resize(1);\n halo_btl[i].L2H_maps[0].resize(halo_lists[i]->size());\n amgx::thrust::copy(A.manager->B2L_maps[i].begin(), A.manager->B2L_maps[i].end(), halo_btl[i].B2L_maps[0].begin());\n amgx::thrust::copy(halo_lists[i]->begin(), halo_lists[i]->end(), halo_btl[i].L2H_maps[0].begin());\n int max_index = thrust_wrapper::reduce(halo_lists[i]->begin(), halo_lists[i]->end(), (int)0, amgx::thrust::maximum()) + 1;\n max_index = max_index > A.get_num_rows() ? max_index : A.get_num_rows();\n halo_btl[i].set_index_range(max_index);\n amgx::thrust::copy(A.manager->B2L_rings[i].begin(), A.manager->B2L_rings[i].end(), halo_btl[i].B2L_rings[0].begin());\n }\n }\n else\n {\n if (rings > 1) { FatalError(\"halo_coloring=LAST dummy halo rows creation not implemented\", AMGX_ERR_NOT_IMPLEMENTED); }\n\n //Create dummy halo matrices\n for (int i = 0; i < neighbors; i++)\n {\n int size = halo_lists[i]->size();\n halo_btl[i].set_global_id(A.manager->neighbors[i]);\n halo_btl[i].set_base_index(0);\n //Adding one more ring to store halo numbering\n halo_btl[i].resize(1, rings);\n halo_btl[i].B2L_maps[0].resize(size);\n halo_btl[i].L2H_maps.resize(1);\n halo_btl[i].L2H_maps[0].resize(A.manager->B2L_rings[i][rings]);\n amgx::thrust::copy(A.manager->B2L_maps[i].begin(), A.manager->B2L_maps[i].end(), halo_btl[i].L2H_maps[0].begin());\n amgx::thrust::copy(halo_lists[i]->begin(), halo_lists[i]->end(), halo_btl[i].B2L_maps[0].begin());\n int max_index = thrust_wrapper::reduce(halo_lists[i]->begin(), halo_lists[i]->end(), (int)0, amgx::thrust::maximum()) + 1;\n max_index = max_index > A.get_num_rows() ? max_index : A.get_num_rows();\n halo_btl[i].set_index_range(max_index);\n halo_btl[i].B2L_rings[0][0] = 0;\n halo_btl[i].B2L_rings[0][1] = size;\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::create_halo_btl_multiple_rings(Matrix_d &A, int rings)\n{\n // creates halo btl managers\n // required input from A.manager: num_neighbors, B2L_rings, B2L_maps, L2H_maps\n int neighbors = A.manager->num_neighbors();\n A.manager->halo_btl = new std::vector >(neighbors);\n std::vector > &halo_btl = *A.manager->halo_btl;\n\n for (int i = 0; i < neighbors; i++)\n {\n // Transfer indexing data too (global indices of exported rows, and sizes of the halo rings)\n int size = A.manager->B2L_rings[i][rings];\n halo_btl[i].set_global_id(A.manager->global_id());\n halo_btl[i].set_base_index(0);\n // Adding one more ring to store halo numbering\n halo_btl[i].resize(1, rings);\n halo_btl[i].B2L_maps[0].resize(size);\n halo_btl[i].L2H_maps.resize(1);\n halo_btl[i].L2H_maps[0].resize(A.manager->L2H_maps[i].size());\n amgx::thrust::copy(A.manager->B2L_maps[i].begin(), A.manager->B2L_maps[i].end(), halo_btl[i].B2L_maps[0].begin());\n amgx::thrust::copy(A.manager->L2H_maps[i].begin(), A.manager->L2H_maps[i].end(), halo_btl[i].L2H_maps[0].begin());\n amgx::thrust::copy(A.manager->B2L_rings[i].begin(), A.manager->B2L_rings[i].end(), halo_btl[i].B2L_rings[0].begin());\n // compute index range\n int max_index = thrust_wrapper::reduce(A.manager->L2H_maps[i].begin(), A.manager->L2H_maps[i].end(), (int)0, amgx::thrust::maximum()) + 1;\n max_index = max_index > A.get_num_rows() ? max_index : A.get_num_rows();\n halo_btl[i].set_index_range(max_index);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::create_halo_matrices(Matrix_d &A, INDEX_TYPE rings,\n int64_t base_index, INDEX_TYPE index_range,\n INDEX_TYPE global_id, Vector &neighbors_,\n std::vector &B2L_maps,\n std::vector > &B2L_rings,\n std::vector **halo_rows_p,\n std::vector **halo_btl_p, std::vector &boundary_lists, std::vector &halo_lists)\n{\n int neighbors = neighbors_.size();\n this->create_maps(A, rings, neighbors, B2L_maps, B2L_rings, boundary_lists);\n this->create_halo_btl(A, rings, neighbors, neighbors_, base_index, index_range, global_id, B2L_maps, B2L_rings, halo_btl_p, halo_lists);\n this->create_halo_rows(A, rings, neighbors, B2L_maps, B2L_rings, halo_rows_p, halo_lists);\n}\n\ntemplate \nvoid DistributedArranger >::create_B2L(Matrix_d &A, INDEX_TYPE my_id, INDEX_TYPE rings)\n{\n this->halo_coloring = A.manager->getComms()->halo_coloring;\n I64Vector_h halo_ranges_h;\n this->create_neighbors(A, my_id, halo_ranges_h);\n A.manager->set_global_id(my_id);\n int num_neighbors = A.manager->neighbors.size();\n A.manager->resize(num_neighbors, rings);\n A.manager->getComms()->set_neighbors(num_neighbors);\n std::vector boundary_lists;\n IVector halo_nodes;\n this->create_boundary_lists(A, halo_ranges_h, boundary_lists, halo_nodes);\n A.manager->L2H_maps.resize(num_neighbors);\n std::vector halo_lists;\n halo_lists.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n A.manager->L2H_maps[i].resize(boundary_lists[i].size());\n thrust_wrapper::transform(boundary_lists[i].begin(), boundary_lists[i].end(), amgx::thrust::constant_iterator (halo_ranges_h[2 * i]), A.manager->L2H_maps[i].begin(), amgx::thrust::plus ());\n halo_lists[i] = &A.manager->L2H_maps[i];\n }\n\n cudaCheckError();\n A.manager->getComms()->exchange_vectors(boundary_lists, A, 0);\n this->renumber_to_local(A, boundary_lists, halo_lists, my_id, halo_ranges_h, halo_nodes);\n this->create_halo_matrices(A, 1, boundary_lists, halo_lists);\n\n if (this->halo_coloring != LAST)\n {\n A.manager->getComms()->exchange_matrix_halo(*A.manager->halo_rows, *A.manager->halo_btl, A);\n }\n\n A.manager->A = &A;\n}\n\ntemplate \nvoid DistributedArranger >::create_halo_matrices(Matrix_d &A, INDEX_TYPE rings, std::vector &boundary_lists, std::vector &halo_lists)\n{\n this->create_maps(A, rings, boundary_lists);\n this->create_halo_btl(A, halo_lists);\n this->create_halo_rows(A, halo_lists);\n}\n\ntemplate \nvoid DistributedArranger >::create_B2L(Matrix_h &A, INDEX_TYPE my_id, INDEX_TYPE rings)\n{\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class DistributedArrangerBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class DistributedArranger::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n\n\n\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace amgx\n{\n\n/***************************************\n * Source Definitions\n ***************************************/\n\nstruct is_less_than_zero\n{\n __host__ __device__\n bool operator()(int x) const\n {\n return x < 0;\n }\n};\n\n\nstruct is_non_neg\n{\n __host__ __device__\n bool operator()(const int &x)\n {\n return x >= 0;\n }\n};\n\ntemplate\n__global__\nvoid flag_boundary_nodes_P(const IndexType *A_rows, const IndexType *A_cols, int *boundary_flags, int num_owned_coarse_pts, int A_num_rows, int* nnz_interior)\n{\n const int nWarps = cta_size / 32;\n // The coordinates of the thread inside the CTA/warp.\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n int nnz_int = 0;\n\n // Loop over rows of A.\n for ( int aRowId = blockIdx.x * nWarps + warpId ; aRowId < A_num_rows ; aRowId += gridDim.x * nWarps )\n {\n // Load A row IDs.\n int aColIt = A_rows[aRowId ];\n int aColEnd = A_rows[aRowId + 1];\n int nnz = aColEnd - aColIt;\n bool halo_found = 0;\n\n // Iterate over the columns of A.\n for ( aColIt += laneId ; utils::any( aColIt < aColEnd ) && !halo_found; aColIt += 32 )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int aColId = aColIt < aColEnd ? A_cols[aColIt] : -1;\n bool is_halo = (aColIt < aColEnd) && (aColId >= num_owned_coarse_pts);\n\n // Push coarse and strongly connected nodes in the set.\n if (utils::any(is_halo))\n {\n halo_found = true;\n break;\n }\n }\n\n if (laneId == 0)\n {\n if(halo_found)\n {\n boundary_flags[aRowId] = 1;\n }\n else\n {\n nnz_int += nnz;\n }\n }\n }\n\n if(laneId == 0)\n {\n atomicAdd(nnz_interior, nnz_int);\n }\n}\n\n\n__global__\nvoid flag_boundary_nodes(int *b2l_maps, int *boundary_flags, int size )\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n boundary_flags[b2l_maps[tidx]] = 1;\n }\n}\n\n__global__\nvoid flag_interior_nodes(int *boundary_flags, int *interior_flags, int size )\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n interior_flags[tidx] = !boundary_flags[tidx];\n }\n}\n\n__global__\nvoid fill_rows_lists(int *__restrict__ interior_rows_lists, int *__restrict__ boundary_rows_lists, const int *__restrict__ interior_flags, const int *__restrict__ boundary_flags, int size )\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n int bdy_flag = boundary_flags[tidx];\n\n if (bdy_flag != boundary_flags[tidx + 1]) { boundary_rows_lists[bdy_flag] = tidx; }\n\n int int_flag = interior_flags[tidx];\n\n if (int_flag != interior_flags[tidx + 1]) { interior_rows_lists[int_flag] = tidx; }\n }\n}\n\n\ntemplate\n__global__\nvoid fill_B_B2L_map(const IndexType *maps_offsets, const IndexType *A_b2l_map, const IndexType *row_offsets, IndexType *C_hat_b2l_map, int size, int rank)\n{\n const int nWarps = cta_size / 32;\n // The coordinates of the thread inside the CTA/warp.\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n // Loop over rows of A.\n for ( int aRowId = blockIdx.x * nWarps + warpId ; aRowId < size ; aRowId += gridDim.x * nWarps )\n {\n // Load A row IDs.\n int aColIt = maps_offsets[aRowId ];\n int size = maps_offsets[aRowId + 1] - aColIt;\n int start = row_offsets[A_b2l_map[aRowId]];\n int end = start + size;\n\n // Iterate over the columns of A.\n for ( start += laneId, aColIt += laneId ; utils::any( start < end ) ; start += 32, aColIt += 32 )\n {\n if (start < end)\n {\n C_hat_b2l_map[aColIt] = start;\n }\n }\n }\n}\n\ntemplate \n__global__ void fill_halo_rows_row_ids(int64_t *halo_rows_row_ids, INDEX_TYPE *halo_rows_local_row_ids, int64_t base_index, INDEX_TYPE num_owned_coarse_pts, INDEX_TYPE size)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n int64_t gid = halo_rows_row_ids[tidx];\n halo_rows_local_row_ids[tidx] = (INDEX_TYPE) (gid - base_index);\n\n if (gid < base_index || gid > base_index + num_owned_coarse_pts)\n {\n printf(\"Should not be here in fill_halo_rows_row_ids\\n\");\n }\n }\n}\n\n__global__\nvoid fill_local_to_global_map(int *bdy_list, int64_t *local_to_global, int64_t offset, int size )\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n local_to_global[tidx] = bdy_list[tidx] + offset;\n }\n}\n\ntemplate \n__global__ void remove_fine_halo(int64_t *cf_map_global, INDEX_TYPE *flag, INDEX_TYPE size)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n flag[tidx] = (cf_map_global[tidx] == -1) ? 0 : 1;\n }\n}\n\ntemplate \n__global__ void fill_P_local_to_global(int64_t *P_local_to_global, int64_t *cf_map_global, INDEX_TYPE *flag, INDEX_TYPE size)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n if (cf_map_global[tidx] != -1)\n {\n P_local_to_global[flag[tidx]] = cf_map_global[tidx];\n }\n }\n}\n\ntemplate \n__global__ void calc_num_neighbors(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *part_offsets, INDEX_TYPE *exists, INDEX_TYPE num_part, INDEX_TYPE my_id, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n if (col < part_offsets[my_id] || col >= part_offsets[my_id + 1])\n {\n int part = 0;\n\n while (part < num_part && (col < part_offsets[part] || col >= part_offsets[part + 1])) { part++; }\n\n if (part < num_part && (col >= part_offsets[part] && col < part_offsets[part + 1]))\n {\n exists[part] = 1;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n// this is assuming local indices (used in matrix A)\ntemplate \n__global__ void calc_num_neighbors_v2_local(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *part_offsets, INDEX_TYPE *exists, int64_t *local_to_global_map, INDEX_TYPE num_part, INDEX_TYPE my_id, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n // Check if column point to halo node\n if (col >= num_rows)\n {\n int64_t gcol = local_to_global_map[col - num_rows];\n int part = 0;\n\n while (part < num_part && (gcol < part_offsets[part] || gcol >= part_offsets[part + 1]))\n {\n part++;\n }\n\n if (part < num_part && (gcol >= part_offsets[part] && gcol < part_offsets[part + 1]))\n {\n exists[part] = 1;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n// this is assuming global indices (used in halo rows)\ntemplate \n__global__ void calc_num_neighbors_v2_global(INDEX_TYPE *row_offsets, int64_t *col_indices, int64_t *part_offsets, INDEX_TYPE *exists, INDEX_TYPE num_part, INDEX_TYPE my_id, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n int64_t start = part_offsets[my_id];\n int64_t end = part_offsets[my_id + 1];\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n int64_t col = col_indices[i];\n\n // Check if column point to halo node\n if (col < start || col >= end)\n {\n int part = 0;\n\n while (part < num_part && (col < part_offsets[part] || col >= part_offsets[part + 1]))\n {\n part++;\n }\n\n if (part < num_part && (col >= part_offsets[part] && col < part_offsets[part + 1]))\n {\n exists[part] = 1;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void flag_halo_nodes(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *part_ranges, INDEX_TYPE *flags, INDEX_TYPE *flag_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_part, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n if (col < base || col >= base + range)\n {\n int part = 0;\n\n while (part < num_part && (col < part_ranges[2 * part] || col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_part && (col >= part_ranges[2 * part] && col < part_ranges[2 * part + 1]))\n {\n flags[flag_offsets[part] + col - part_ranges[2 * part]] = 1;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void flag_halo_nodes_local(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *part_ranges, INDEX_TYPE *flags, INDEX_TYPE *flag_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_part, INDEX_TYPE num_rows, int64_t *local_to_global)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n if (col >= num_rows)\n {\n int64_t global_col = local_to_global[col - num_rows];\n int part = 0;\n\n while (part < num_part && (global_col < part_ranges[2 * part] || global_col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_part && (global_col >= part_ranges[2 * part] && global_col < part_ranges[2 * part + 1]))\n {\n flags[flag_offsets[part] + global_col - part_ranges[2 * part]] = set_val > 0 ? set_val : -col - 1; // if -1, then set -(local index)-1 to restore later\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n__global__ void flag_halo_nodes_from_local_to_global_map(int64_t *part_ranges, INDEX_TYPE *flags, INDEX_TYPE *flag_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_part, INDEX_TYPE size, int64_t *local_to_global)\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (tid < size)\n {\n int64_t global_col = local_to_global[tid];\n int part = 0;\n\n while (part < num_part && (global_col < part_ranges[2 * part] || global_col >= part_ranges[2 * part + 1]))\n {\n part++;\n }\n\n if (part < num_part && (global_col >= part_ranges[2 * part] && global_col < part_ranges[2 * part + 1]))\n {\n int col_id = range + tid;\n flags[flag_offsets[part] + global_col - part_ranges[2 * part]] = -1 - col_id;\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void flag_halo_nodes_local_v2(int64_t *part_ranges, INDEX_TYPE *flags, INDEX_TYPE *flag_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_part, INDEX_TYPE local_to_global_size, int64_t *local_to_global)\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (tid < local_to_global_size)\n {\n int64_t global_col = local_to_global[tid];\n int col = range + tid;\n // Identify the partition that owns that node\n int part = 0;\n\n while (part < num_part && (global_col < part_ranges[2 * part] || global_col >= part_ranges[2 * part + 1])) { part++; }\n\n // Flag the corresponding node in flags array\n if (part < num_part && (global_col >= part_ranges[2 * part] && global_col < part_ranges[2 * part + 1]))\n {\n flags[flag_offsets[part] + global_col - part_ranges[2 * part]] = -col - 1;\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate \n__global__ void flag_halo_nodes_global_v2(INDEX_TYPE *row_offsets, int64_t *global_col_indices, int64_t *part_ranges, INDEX_TYPE *flags, INDEX_TYPE *flag_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_part, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n int64_t global_col = global_col_indices[i];\n\n if (global_col < base || global_col >= base + (int64_t) range)\n {\n int part = 0;\n\n while (part < num_part && (global_col < part_ranges[2 * part] || global_col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_part && (global_col >= part_ranges[2 * part] && global_col < part_ranges[2 * part + 1]))\n {\n flags[flag_offsets[part] + (int) (global_col - part_ranges[2 * part])] = 1;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n\n\n\ntemplate \n__global__ void flag_halo_nodes_global(INDEX_TYPE *row_offsets, INDEX_TYPE num_rows, int64_t *global_col_indices, int64_t *part_ranges, INDEX_TYPE *flags, INDEX_TYPE *flag_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_neighbors, INDEX_TYPE *local_col_indices)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n int64_t col = global_col_indices[i];\n\n if (col < base || col >= base + range)\n {\n // exterior node\n int part = 0;\n\n while (part < num_neighbors && (col < part_ranges[2 * part] || col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_neighbors && (col >= part_ranges[2 * part] && col < part_ranges[2 * part + 1]))\n {\n // check if the flag is already set (i.e. 1-ring node) then skip it\n int flag = flags[flag_offsets[part] + col - part_ranges[2 * part]];\n\n if (flag == 0)\n {\n flags[flag_offsets[part] + col - part_ranges[2 * part]] = set_val;\n }\n else if (flag < 0) // it is possible that flag is already set to 1 in this kernel, so if it's negative then it's the 1st ring\n {\n // update local index, note that local to global mapping is unchanged for 1st ring\n local_col_indices[i] = -flags[flag_offsets[part] + col - part_ranges[2 * part]] - 1;\n }\n }\n }\n else\n {\n // interior node\n local_col_indices[i] = col - base;\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void renumber_col_indices_global(INDEX_TYPE *row_offsets, int *col_indices, int64_t *global_col_indices, int64_t *part_ranges, INDEX_TYPE *flags, INDEX_TYPE *flag_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_neighbors, INDEX_TYPE num_rows, int rank)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n int64_t col = global_col_indices[i];\n\n if (col < base || col >= base + range)\n {\n // Need to get his local index from the flag array\n int part = 0;\n\n while ( part < num_neighbors && (col < part_ranges[2 * part] || col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_neighbors && (col >= part_ranges[2 * part] && col < part_ranges[2 * part + 1]))\n {\n int pos_in_flag = flag_offsets[part] + col - part_ranges[2 * part];\n //int local_col = flags[pos_in_flag];\n col_indices[i] = range + flags[pos_in_flag];\n //bool flag = (local_col != flags[pos_in_flag+1]);\n\n if (flags[pos_in_flag] == flags[pos_in_flag + 1])\n {\n printf(\"should not be here in renumber_col_indices_global\\n\");\n }\n\n //if (flag)\n }\n }\n else\n {\n // interior node\n col_indices[i] = col - base;\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n\n\ntemplate \n__global__ void calc_new_halo_mapping(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *part_ranges, INDEX_TYPE num_part, INDEX_TYPE num_rows, int64_t *local_to_global, int64_t *new_local_to_global, INDEX_TYPE *halo_indices, int halo_base, int rank)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n if (col >= num_rows)\n {\n int64_t glob_col = local_to_global[col - num_rows];\n int part = 0;\n\n while (part < num_part && (glob_col < part_ranges[2 * part] || glob_col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_part && (glob_col >= part_ranges[2 * part] && glob_col < part_ranges[2 * part + 1]))\n {\n int pos = glob_col - part_ranges[2 * part];\n int actual_col = halo_base + halo_indices[pos];\n col_indices[i] = -1 - actual_col;\n new_local_to_global[actual_col - num_rows] = glob_col;\n // when we renumber we also store mapping to use later in residual restriction since matrix R has other order of halo indices\n // store both direct and inverse renumbering for P & R matrices\n //halo_P_renumbering[col - num_rows] = actual_col;\n //halo_R_renumbering[actual_col - num_rows] = col;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n\ntemplate \n__global__ void calc_new_halo_mapping_v3(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *part_ranges, INDEX_TYPE num_part, INDEX_TYPE num_rows, int64_t *local_to_global, int64_t *new_local_to_global, INDEX_TYPE *halo_indices, int halo_base, int rank)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n if (col >= num_rows)\n {\n int64_t glob_col = local_to_global[col - num_rows];\n int part = 0;\n\n while (part < num_part && (glob_col < part_ranges[2 * part] || glob_col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_part && (glob_col >= part_ranges[2 * part] && glob_col < part_ranges[2 * part + 1]))\n {\n int pos = glob_col - part_ranges[2 * part];\n int actual_col = halo_base + halo_indices[pos];\n col_indices[i] = -1 - actual_col;\n new_local_to_global[actual_col - num_rows] = glob_col;\n // when we renumber we also store mapping to use later in residual restriction since matrix R has other order of halo indices\n // store both direct and inverse renumbering for P & R matrices\n //halo_P_renumbering[col - num_rows] = actual_col;\n //halo_R_renumbering[actual_col - num_rows] = col;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void flip_sign_col(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n if (col < 0)\n {\n col_indices[i] = -col_indices[i] - 1;\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void calc_new_halo_mapping_ring2(INDEX_TYPE *row_offsets, INDEX_TYPE num_rows, int64_t *global_col_indices, int64_t *part_ranges, int64_t base, INDEX_TYPE range, INDEX_TYPE num_neighbors, INDEX_TYPE *halo_offsets, INDEX_TYPE *neighbor_offsets, INDEX_TYPE *neighbor_nodes, int64_t *local_to_global, INDEX_TYPE *local_col_indices)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n int64_t col = global_col_indices[i];\n\n if ((col < base || col >= base + range) && local_col_indices[i] == -1)\n {\n // update only 2nd ring halo indices here, 1st ring is updated already\n int part = 0;\n\n while (part < num_neighbors && (col < part_ranges[2 * part] || col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_neighbors && (col >= part_ranges[2 * part] && col < part_ranges[2 * part + 1]))\n {\n int pos = col - part_ranges[2 * part];\n local_col_indices[i] = halo_offsets[part] + neighbor_nodes[neighbor_offsets[part] + pos];\n local_to_global[local_col_indices[i] - range] = col;\n }\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void renumber_col_indices(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *part_ranges, INDEX_TYPE *halo_renum, INDEX_TYPE *halo_offsets, int64_t base, INDEX_TYPE range, INDEX_TYPE num_part, INDEX_TYPE num_rows, int64_t *local_to_global_map)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n for (int i = row_offsets[row] + coopIdx; i < row_offsets[row + 1]; i += coop)\n {\n INDEX_TYPE col = col_indices[i];\n\n if (col < base || col >= base + range)\n {\n int part = 0;\n\n while (part < num_part && (col < part_ranges[2 * part] || col >= part_ranges[2 * part + 1])) { part++; }\n\n if (part < num_part && (col >= part_ranges[2 * part] && col < part_ranges[2 * part + 1]))\n {\n int local_index = halo_renum[halo_offsets[part] + col - part_ranges[2 * part]];\n local_to_global_map[local_index - range] = col_indices[i];\n col_indices[i] = local_index;\n }\n }\n else\n {\n col_indices[i] = col - base;\n }\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n__global__ void populate_B2L(INDEX_TYPE *indexing, INDEX_TYPE *output, INDEX_TYPE last, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n if (row == size - 1)\n {\n if (last) { output[indexing[row]] = row; }\n }\n else if (indexing[row] != indexing[row + 1])\n {\n output[indexing[row]] = row;\n }\n\n row += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void populate_bdy_list_and_l2h_map(INDEX_TYPE *l2h_map, INDEX_TYPE *halo_nodes, INDEX_TYPE *halo_indices, INDEX_TYPE *bdy_list, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n int scanned_id = halo_indices[row];\n\n if (scanned_id != halo_indices[row + 1])\n {\n bdy_list[scanned_id] = row;\n l2h_map[scanned_id] = -halo_nodes[row] - 1;\n }\n\n row += gridDim.x * blockDim.x;\n }\n}\n\n\n__global__ void populate_B2L_v2(INDEX_TYPE *indexing, INDEX_TYPE *output, INDEX_TYPE offset, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n if (indexing[row] != indexing[row + 1])\n {\n output[indexing[row] - offset] = row;\n }\n\n row += gridDim.x * blockDim.x;\n }\n}\n\n\n\ntemplate \n__global__ void export_matrix_elements(INDEX_TYPE *row_offsets, T *values, INDEX_TYPE bsize, INDEX_TYPE *maps, INDEX_TYPE *pointers, T *output, INDEX_TYPE *col_indices, INDEX_TYPE *output2, INDEX_TYPE size)\n{\n int idx = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (idx < size)\n {\n int row = maps[idx];\n INDEX_TYPE src_base = row_offsets[row];\n INDEX_TYPE dst_base = pointers[idx];\n\n for (int m = coopIdx; m < row_offsets[row + 1]*bsize - src_base * bsize; m += coop)\n {\n output[dst_base * bsize + m] = values[src_base * bsize + m];\n }\n\n for (int m = coopIdx; m < row_offsets[row + 1] - src_base; m += coop)\n {\n output2[dst_base + m] = col_indices[src_base + m];\n }\n\n idx += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void export_matrix_elements_global(INDEX_TYPE *row_offsets, T *values, INDEX_TYPE bsize, INDEX_TYPE *maps, INDEX_TYPE *pointers, T *output, INDEX_TYPE *col_indices, int64_t *output2, INDEX_TYPE size, int64_t *local_to_global, INDEX_TYPE *q, INDEX_TYPE num_owned_pts, int64_t base_index)\n{\n int idx = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (idx < size)\n {\n int row = maps[idx];\n\n if (q != NULL)\n {\n row = q[row];\n }\n\n INDEX_TYPE src_base = row_offsets[row];\n INDEX_TYPE dst_base = pointers[idx];\n\n for (int m = coopIdx; m < row_offsets[row + 1]*bsize - src_base * bsize; m += coop)\n {\n output[dst_base * bsize + m] = values[src_base * bsize + m];\n }\n\n for (int m = coopIdx; m < row_offsets[row + 1] - src_base; m += coop)\n {\n int col = col_indices[src_base + m];\n\n if (col < num_owned_pts)\n {\n output2[dst_base + m] = (int64_t) col_indices[src_base + m] + base_index;\n }\n else\n {\n output2[dst_base + m] = local_to_global[col_indices[src_base + m] - num_owned_pts];\n }\n }\n\n idx += gridDim.x * blockDim.x / coop;\n }\n}\n\ntemplate \n__global__ void export_matrix_diagonal(T *values, INDEX_TYPE bsize, INDEX_TYPE *maps, T *output, INDEX_TYPE size)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (idx < size)\n {\n int row = maps[idx];\n INDEX_TYPE src_base = row;\n INDEX_TYPE dst_base = idx;\n\n for (int m = 0; m < bsize; m++)\n {\n output[dst_base * bsize + m] = values[src_base * bsize + m];\n }\n\n idx += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void write_matrix_rowsize(INDEX_TYPE *maps, INDEX_TYPE *row_offsets, INDEX_TYPE size, INDEX_TYPE *output)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (idx < size)\n {\n int row = maps[idx];\n output[idx] = row_offsets[row + 1] - row_offsets[row];\n idx += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void find_next_ring(INDEX_TYPE *maps, INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, INDEX_TYPE size, INDEX_TYPE base, INDEX_TYPE range, INDEX_TYPE *halo_rings)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (idx < size)\n {\n int row = maps[idx];\n\n for (int col = row_offsets[row]; col < row_offsets[row + 1]; col++)\n {\n int colIdx = col_indices[col];\n\n if (colIdx >= base && colIdx < base + range)\n {\n halo_rings[colIdx - base] = 1;\n }\n }\n\n idx += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void remove_previous_rings(INDEX_TYPE *map, INDEX_TYPE *halo_rings, INDEX_TYPE size)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (idx < size)\n {\n int row = map[idx];\n halo_rings[row] = 0;\n idx += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void renumber_halo_lists(INDEX_TYPE *halo_list, INDEX_TYPE *halo_renum, int64_t base, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n halo_list[row] = halo_renum[halo_list[row] - base];\n row += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void write_ring_rows(INDEX_TYPE *indices, INDEX_TYPE *maps, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n int idx = indices[row];\n\n if (idx == 1 && row == 0)\n {\n maps[idx - 1] = row;\n }\n else if (row > 0 && idx != indices[row - 1])\n {\n maps[idx - 1] = row;\n }\n\n row += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \nvoid DistributedArrangerBase::set_part_offsets(const INDEX_TYPE num_part, const VecInt_t *part_offsets_h)\n{\n part_offsets.resize(num_part + 1);\n amgx::thrust::copy(&part_offsets_h[0], &part_offsets_h[num_part + 1], part_offsets.begin());\n cudaCheckError();\n this->num_part = num_part;\n}\n\n// TODO Make distributed reader use 64bit ints for offsets\ntemplate \nvoid DistributedArrangerBase::set_part_offsets(const INDEX_TYPE num_part, const int64_t *part_offsets_h)\n{\n part_offsets.resize(num_part + 1);\n amgx::thrust::copy(&part_offsets_h[0], &part_offsets_h[num_part + 1], part_offsets.begin());\n cudaCheckError();\n this->num_part = num_part;\n}\n\ntemplate \nvoid DistributedArranger >::create_B2L(Matrix_d &A,\n INDEX_TYPE my_id, INDEX_TYPE rings,\n int64_t &base_index, INDEX_TYPE &index_range,\n Vector &neighbors,\n I64Vector_d &halo_ranges,\n std::vector &B2L_maps,\n std::vector &L2H_maps,\n std::vector > &B2L_rings,\n DistributedComms **comms,\n std::vector **halo_rows,\n std::vector **halo_btl)\n{\n I64Vector_h halo_ranges_h;\n // initialize base_index, index_range, neighbors, halo_ranges_h\n this->create_neighbors(A, my_id, base_index, index_range, neighbors, halo_ranges_h);\n this->create_B2L_from_neighbors(A, my_id, rings, base_index, index_range, neighbors,\n halo_ranges_h, halo_ranges, B2L_maps, L2H_maps, B2L_rings, comms, halo_rows, halo_btl);\n}\n\ntemplate \nvoid DistributedArranger >::create_B2L_from_maps(Matrix_d &A,\n INDEX_TYPE my_id, INDEX_TYPE rings,\n int64_t base_index, INDEX_TYPE index_range,\n Vector &neighbors,\n I64Vector_d &halo_ranges,\n std::vector &B2L_maps,\n std::vector > &B2L_rings,\n DistributedComms **comms,\n std::vector **halo_rows,\n std::vector **halo_btl)\n{\n int num_neighbors = neighbors.size();\n this->halo_coloring = (*comms)->halo_coloring;\n std::vector halo_lists;\n halo_lists.resize(num_neighbors);\n // initialize halo_btl\n this->create_halo_btl(A, rings, num_neighbors, neighbors, base_index, index_range, my_id, B2L_maps, B2L_rings, halo_btl, halo_lists);\n // initialize halo_rows\n this->create_halo_rows(A, rings, num_neighbors, B2L_maps, B2L_rings, halo_rows, halo_lists);\n\n //this->create_halo_matrices(A, 1, base_index, index_range, my_id, neighbors, B2L_maps, B2L_rings, &(*halo_rows), &(*halo_btl));\n if (this->halo_coloring != LAST)\n {\n Matrix_d tmpA(0, 0, 0, A.getProps());\n tmpA.set_block_dimx(A.get_block_dimx());\n tmpA.set_block_dimy(A.get_block_dimy());\n tmpA.setResources(A.getResources());\n DistributedManager tmp_manager(my_id, base_index, index_range, tmpA, neighbors, halo_ranges, B2L_maps, B2L_rings, comms, NULL, NULL);\n tmpA.setManager(tmp_manager);\n tmpA.manager->getComms()->exchange_matrix_halo(**halo_rows, **halo_btl, tmpA);\n }\n}\n\ntemplate \nvoid DistributedArranger >::create_B2L_from_maps(Matrix_d &A,\n INDEX_TYPE my_id, INDEX_TYPE rings,\n Vector &neighbors,\n std::vector &B2L_maps,\n std::vector &L2H_maps,\n std::vector > &B2L_rings,\n DistributedComms **comms,\n std::vector **halo_rows,\n std::vector **halo_btl)\n{\n this->halo_coloring = (*comms)->halo_coloring;\n int num_neighbors = neighbors.size();\n int64_t base_index = 0;\n int index_range = A.get_num_rows();\n std::vector halo_lists;\n halo_lists.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n halo_lists[i] = &L2H_maps[i];\n }\n\n /*\n EXAMPLE\n Say we want 2 rings of halo. The output of these function for the example partition 1 that at this point looks like:\n A.num_rows = 4; A.num_nz = 20\n A.row_offsets = [0 4 11 15 20]\n A.col_indices = [4 0 1 2\n 4 5 0 1 2 3 7\n 0 1 2 3\n 1 2 3 6 7]\n num_neighbors=2; neighbors = [0 2]\n B2L_rings[[0 2][0 2]] B2L_maps [[0 1][1 3]] (Converted from the following input: btl_sizes = [0 2 4] btl_maps = [0 1| 1 3])\n L2H_maps (and halo_lists) [[4 5][6 7]]\n halo_rows = NULL, halo_btl=NULL\n (global indices mapping to local indices: 4-0 5-1 6-2 7-3 2-4 3-5 9-6 10-7)\n */\n // initialize ring>1 B2L_maps and B2L_rings\n this->create_rings(A, rings, num_neighbors, B2L_maps, B2L_rings);\n /*\n EXAMPLE\n having mapped out two rings of halo (already had one), B2L_maps and B2L_rings look like:\n B2L_rings[[0 2 4][0 2 4]] B2L_maps[[0 1| 2 3][1 3| 0 2]]\n */\n // initialize halo_btl, creating DistributedManager objects for each halo row matrix to be sent to each neighbor, storing B2L and L2H maps\n this->create_halo_btl(A, rings, num_neighbors, neighbors, base_index, index_range, my_id, B2L_maps, B2L_rings, halo_btl, halo_lists);\n /*\n EXAMPLE\n We created to two DistributedManager objects for the two neighbors. if halo_coloring != LAST, then the relevant fields are:\n halo_btl[0] (for neighbor ID 0)\n global_id = 1; base_index=0; index_range=6 B2L_rings[0] = [0 2 4] B2L_maps[0] = [0 1| 2 3] L2H_maps = [4 5]\n halo_btl[1] (for neighbor ID 2)\n global_id = 1; base_index=0; index_range=8 B2L_rings[0] = [0 2 4] B2L_maps[0] = [1 3| 0 2] L2H_maps = [6 7]\n */\n // Extract halo rows, ring by ring\n this->create_halo_rows(A, rings, num_neighbors, B2L_maps, B2L_rings, halo_rows, halo_lists);\n\n /*\n EXAMPLE\n Extracted two rings of rows for the two neighbors, the two halo row matrices look like:\n halo_rows[0] (for neighbor ID 0)\n halo_rows[0].row_offsets = [0 4 11 15 20]\n halo_rows[0].col_indices = [4 0 1 2\n 4 5 0 1 2 3 7\n 0 1 2 3\n 1 2 3 6 7]\n\n halo_rows[1] (for neighbor ID 2)\n halo_rows[1].row_offsets = [0 7 12 16 20]\n halo_rows[1].col_indices = [4 5 0 1 2 3 7\n 1 2 3 6 7\n 4 0 1 2\n 0 1 2 3]\n Note, that the ordering of rows corresponds to the list of indices in B2L_maps[i] (and halo_btl[i].B2L_maps[0])\n */\n //this->create_halo_matrices(A, 1, base_index, index_range, my_id, neighbors, B2L_maps, B2L_rings, &(*halo_rows), &(*halo_btl));\n if (this->halo_coloring != LAST) //For multiple rings, we still need 0-num_rings-1 rows, just not the last one.\n {\n //Set up dummy matrix and dummy manager, and exchange the halo rows\n Matrix_d tmpA(0, 0, 0, A.getProps());\n tmpA.set_block_dimx(A.get_block_dimx());\n tmpA.set_block_dimy(A.get_block_dimy());\n tmpA.setResources(A.getResources());\n I64Vector_d tmp_halo_ranges;\n DistributedManager tmp_manager(my_id, base_index, index_range, tmpA, neighbors, tmp_halo_ranges, B2L_maps, B2L_rings, comms, NULL, NULL);\n tmpA.setManager(tmp_manager);\n tmpA.manager->getComms()->exchange_matrix_halo(**halo_rows, **halo_btl, tmpA);\n }\n\n /*\n EXAMPLE\n The partitions exchange their halo_rows and halo_btl fields, partition 1 ends up with:\n halo_btl[0] (received from neighbor ID 0)\n global_id = 0; base_index=0; index_range=6; B2L_rings[0] = [0 2 4]; B2L_maps[0] = [2 3| 0 1] L2H_maps = [4 5]\n halo_rows[0].row_offsets = [0 5 13 17 21]\n halo_rows[0].col_indices = [1 2 3 4 5\n 0 1 2 3 4 5 6 7\n 0 1 3 6\n 0 1 2 3]\n\n halo_btl[1] (received from neighbor ID 2)\n global_id = 2; base_index=0; index_range=8; B2L_rings[0] = [0 2 4]; B2L_maps[0] = [1 2| 0 3] L2H_maps = [6 7]\n halo_rows[1].row_offsets = [0 4 11 16 20]\n halo_rows[1].col_indices = [7 1 2 3\n 5 6 7 0 1 2 3\n 4 5 0 2 3\n 0 1 2 3]\n */\n}\n\ntemplate \nvoid DistributedArranger >::renumber_to_local(Matrix_d &A,\n std::vector &boundary_lists,\n std::vector &halo_lists,\n INDEX_TYPE my_id,\n int64_t base_index, INDEX_TYPE index_range,\n Vector &neighbors,\n I64Vector_h &halo_ranges_h,\n I64Vector_d &halo_ranges,\n IVector &halo_nodes)\n{\n //IVector &halo_nodes = *halo_nodes_p;\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n int num_neighbors = halo_ranges_h.size() / 2;\n int total_rows_of_neighbors = 0;\n std::vector halo_offsets(num_neighbors + 1, 0);\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n total_rows_of_neighbors += halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n halo_offsets[i + 1] = halo_offsets[i] + halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n }\n\n //IVector halo_nodes(total_rows_of_neighbors);\n IVector halo_offsets_d(num_neighbors + 1);\n amgx::thrust::copy(halo_offsets.begin(), halo_offsets.end(), halo_offsets_d.begin());\n int last_halo = halo_nodes[halo_nodes.size() - 1];\n //create renumbering of halo indices (contiguously neighbor - by - neighbor, indexing starting with the number of owned rows\n thrust_wrapper::exclusive_scan(halo_nodes.begin(), halo_nodes.end(), halo_nodes.begin(), size);\n cudaCheckError();\n A.manager->local_to_global_map.resize(halo_nodes[halo_nodes.size() - 1] - size + last_halo);\n renumber_col_indices<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), halo_ranges.raw(), halo_nodes.raw(), halo_offsets_d.raw(), base_index, index_range, num_neighbors, size, A.manager->local_to_global_map.raw());\n cudaCheckError();\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n int size = halo_lists[i]->size();\n int num_blocks = std::min(4096, (size + 127) / 128);\n renumber_halo_lists <<< num_blocks, 128>>>(halo_lists[i]->raw(), halo_nodes.raw() + halo_offsets[i], halo_ranges_h[2 * i], size);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::renumber_to_local(Matrix_d &A,\n std::vector &boundary_lists, std::vector &halo_lists,\n INDEX_TYPE my_id, I64Vector_h &halo_ranges_h, IVector &halo_nodes)\n{\n //IVector &halo_nodes = *halo_nodes_p;\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n int num_neighbors = halo_ranges_h.size() / 2;\n int total_rows_of_neighbors = 0;\n std::vector halo_offsets(num_neighbors + 1, 0);\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n total_rows_of_neighbors += halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n halo_offsets[i + 1] = halo_offsets[i] + halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n }\n\n //IVector halo_nodes(total_rows_of_neighbors+1);\n IVector halo_offsets_d(num_neighbors + 1);\n amgx::thrust::copy(halo_offsets.begin(), halo_offsets.end(), halo_offsets_d.begin());\n int last_halo = (halo_nodes.size() > 0) ? halo_nodes[halo_nodes.size() - 1] : 0;\n //create renumbering of halo indices (contiguously neighbor - by - neighbor, indexing starting with the number of owned rows\n thrust_wrapper::exclusive_scan(halo_nodes.begin(), halo_nodes.end(), halo_nodes.begin(), size);\n cudaCheckError();\n\n if (halo_nodes.size() > 0)\n {\n A.manager->local_to_global_map.resize(halo_nodes[halo_nodes.size() - 1] - size + last_halo);\n }\n\n renumber_col_indices<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), A.manager->halo_ranges.raw(), halo_nodes.raw(), halo_offsets_d.raw(), A.manager->base_index(), A.manager->index_range(), num_neighbors, size, A.manager->local_to_global_map.raw());\n cudaCheckError();\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n int size = halo_lists[i]->size();\n int num_blocks = std::min(4096, (size + 127) / 128);\n renumber_halo_lists <<< num_blocks, 128>>>(halo_lists[i]->raw(), halo_nodes.raw() + halo_offsets[i], halo_ranges_h[2 * i], size);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::prepare_local(Matrix_d &A,\n INDEX_TYPE my_id, INDEX_TYPE rings,\n int64_t base_index, INDEX_TYPE index_range,\n Vector &neighbors,\n I64Vector_h &halo_ranges_h,\n I64Vector_d &halo_ranges,\n std::vector &L2H_maps,\n std::vector &boundary_lists,\n std::vector &halo_lists,\n DistributedComms **comms)\n{\n int num_neighbors = neighbors.size();\n (*comms)->set_neighbors(num_neighbors);\n IVector halo_nodes;\n // copy halo_ranges_h to device and populate boundary_lists\n this->create_boundary_lists(A, base_index, index_range, halo_ranges_h, halo_ranges, boundary_lists, halo_nodes);\n L2H_maps.resize(num_neighbors);\n halo_lists.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n L2H_maps[i].resize(boundary_lists[i].size());\n thrust_wrapper::transform(boundary_lists[i].begin(), boundary_lists[i].end(), amgx::thrust::constant_iterator (halo_ranges_h[2 * i]), L2H_maps[i].begin(), amgx::thrust::plus ());\n halo_lists[i] = &L2H_maps[i];\n }\n\n cudaCheckError();\n // setting up tmp manager in order to use Comm API\n Matrix_d tmpA(0, 0, 0, A.getProps());\n tmpA.set_block_dimx(A.get_block_dimx());\n tmpA.set_block_dimy(A.get_block_dimy());\n tmpA.setResources(A.getResources());\n std::vector B2L_maps;\n std::vector > B2L_rings;\n DistributedManager tmp_manager(my_id, base_index, index_range, tmpA, neighbors, halo_ranges, B2L_maps, B2L_rings, comms, NULL, NULL);\n tmpA.setManager(tmp_manager);\n tmpA.manager->getComms()->exchange_vectors(boundary_lists, tmpA, 0);\n this->renumber_to_local(A, boundary_lists, halo_lists, my_id, base_index, index_range, neighbors, halo_ranges_h, halo_ranges, halo_nodes);\n};\n\ntemplate \nvoid DistributedArranger >::create_B2L_from_neighbors(Matrix_d &A,\n INDEX_TYPE my_id, INDEX_TYPE rings,\n int64_t base_index, INDEX_TYPE index_range,\n Vector &neighbors,\n I64Vector_h &halo_ranges_h,\n I64Vector_d &halo_ranges,\n std::vector &B2L_maps,\n std::vector &L2H_maps,\n std::vector > &B2L_rings,\n DistributedComms **comms,\n std::vector **halo_rows,\n std::vector **halo_btl)\n{\n this->halo_coloring = (*comms)->halo_coloring;\n std::vector boundary_lists;\n std::vector halo_lists;\n this->prepare_local(A, my_id, rings, base_index, index_range, neighbors, halo_ranges_h, halo_ranges, L2H_maps, boundary_lists, halo_lists, comms);\n int num_neighbors = neighbors.size();\n // initialize BTL_maps and BTl_rings\n this->create_maps(A, rings, num_neighbors, B2L_maps, B2L_rings, boundary_lists);\n // initialize halo_btl\n this->create_halo_btl(A, rings, num_neighbors, neighbors, base_index, index_range, my_id, B2L_maps, B2L_rings, halo_btl, halo_lists);\n // initialize halo_rows\n this->create_halo_rows(A, rings, num_neighbors, B2L_maps, B2L_rings, halo_rows, halo_lists);\n\n if (this->halo_coloring != LAST)\n {\n // setting up tmp manager in order to use Comm API\n Matrix_d tmpA(0, 0, 0, A.getProps());\n tmpA.set_block_dimx(A.get_block_dimx());\n tmpA.set_block_dimy(A.get_block_dimy());\n tmpA.setResources(A.getResources());\n DistributedManager tmp_manager(my_id, base_index, index_range, tmpA, neighbors, halo_ranges, B2L_maps, B2L_rings, comms, NULL, NULL);\n tmpA.setManager(tmp_manager);\n tmpA.manager->getComms()->exchange_matrix_halo(**halo_rows, **halo_btl, tmpA);\n }\n}\n\ntemplate \nvoid DistributedArranger >::create_neighbors(Matrix_d &A, INDEX_TYPE my_id,\n int64_t &base_index, INDEX_TYPE &index_range,\n Vector &neighbors,\n I64Vector_h &halo_ranges_h)\n{\n //if (this->num_part>32) FatalError(\"Matrix slice Level 1 module supports < 32 partitions\", AMGX_ERR_NOT_IMPLEMENTED);\n if (this->num_part == 0 || this->part_offsets.size() == 0)\n {\n FatalError(\"Partitioning scheme is not set\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n IVector part_sizes(this->num_part, 0);\n IVector inv_neighbors(this->num_part, 0);\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n base_index = this->part_offsets[my_id];\n index_range = this->part_offsets[my_id + 1] - this->part_offsets[my_id];\n //\n // Step 1: flag neighbors that I have an edge to\n //\n //go through all the column indices and compare it to index ranges of nodes\n calc_num_neighbors<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), this->part_offsets.raw(), part_sizes.raw(), this->num_part, my_id, size);\n cudaCheckError();\n std::vector neighbor_flags(this->num_part);\n amgx::thrust::copy(part_sizes.begin(), part_sizes.end(), neighbor_flags.begin());\n //total number of neighbors\n int num_neighbors = amgx::thrust::reduce(neighbor_flags.begin(), neighbor_flags.end());\n cudaCheckError();\n neighbors.resize(num_neighbors);\n halo_ranges_h.resize(num_neighbors * 2);\n //initialize manager\n int active_part = 0;\n\n for (int i = 0; i < this->num_part; i++)\n {\n if (neighbor_flags[i] > 0)\n {\n neighbors[active_part] = i;\n halo_ranges_h[2 * active_part] = this->part_offsets[i];\n halo_ranges_h[2 * active_part + 1] = this->part_offsets[i + 1];\n active_part++;\n }\n }\n}\n\n\n// need for CUDA 4.2\ntemplate \nvoid DistributedArranger >::create_neighbors(Matrix_d &A, INDEX_TYPE my_id,\n I64Vector_h &halo_ranges_h)\n{\n if (this->num_part == 0 || this->part_offsets.size() == 0)\n {\n FatalError(\"Partitioning scheme is not set\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n IVector part_sizes(this->num_part, 0);\n IVector inv_neighbors(this->num_part, 0);\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n A.manager->set_base_index(this->part_offsets[my_id]);\n A.manager->set_index_range(this->part_offsets[my_id + 1] - this->part_offsets[my_id]);\n //\n // Step 1: flag neighbors that I have an edge to\n //\n //go through all the column indices and compare it to index ranges of nodes\n calc_num_neighbors<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), this->part_offsets.raw(), part_sizes.raw(), this->num_part, my_id, size);\n cudaCheckError();\n std::vector neighbor_flags(this->num_part);\n amgx::thrust::copy(part_sizes.begin(), part_sizes.end(), neighbor_flags.begin());\n //total number of neighbors\n int num_neighbors = amgx::thrust::reduce(neighbor_flags.begin(), neighbor_flags.end());\n cudaCheckError();\n A.manager->neighbors.resize(num_neighbors);\n halo_ranges_h.resize(num_neighbors * 2);\n //initialize manager\n int active_part = 0;\n\n for (int i = 0; i < this->num_part; i++)\n {\n if (neighbor_flags[i] > 0)\n {\n A.manager->neighbors[active_part] = i;\n halo_ranges_h[2 * active_part] = this->part_offsets[i];\n halo_ranges_h[2 * active_part + 1] = this->part_offsets[i + 1];\n active_part++;\n }\n }\n}\n\ntemplate \nvoid DistributedArranger >::update_neighbors_list(\n Matrix_d &A, IVector_h &neighbors, I64Vector_h &halo_ranges_h, I64Vector &halo_ranges, I64Vector_h &part_offsets_h, I64Vector &part_offsets, std::vector &halo_rows_row_offsets, std::vector &halo_rows_col_indices)\n{\n int num_partitions = A.manager->get_num_partitions();\n int my_id = A.manager->global_id();\n int num_neighbors = neighbors.size(); // current number of neighbors\n int total_halo_rows = 0;\n int total_halo_nnz = 0;\n IVector neighbor_flags(num_partitions, 0);\n\n for (int i = 0; i < halo_rows_row_offsets.size(); i++)\n {\n int num_halo_rows = halo_rows_row_offsets[i].size() - 1;\n\n if (num_halo_rows > 0)\n {\n total_halo_rows += num_halo_rows;\n total_halo_nnz += halo_rows_row_offsets[i][num_halo_rows];\n int num_blocks = std::min(4096, (num_halo_rows + 127) / 128);\n calc_num_neighbors_v2_global<16> <<< num_blocks, 128>>>(halo_rows_row_offsets[i].raw(), halo_rows_col_indices[i].raw(),\n part_offsets.raw(), neighbor_flags.raw(), num_partitions, my_id, num_halo_rows);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n IVector_h neighbor_flags_h = neighbor_flags;\n\n // unset 1-ring neighbors & myself\n for (int i = 0; i < num_neighbors; i++)\n {\n neighbor_flags_h[neighbors[i]] = 0;\n }\n\n neighbor_flags_h[my_id] = 0;\n // this will update neighbor list and halo ranges, note that we don't change 1-ring neighbors order\n append_neighbors(A, neighbors, halo_ranges_h, halo_ranges, neighbor_flags_h, part_offsets_h);\n}\n\ntemplate \nvoid DistributedArranger >::append_neighbors(Matrix_d &A, IVector_h &neighbors, I64Vector_h &halo_ranges_h, I64Vector &halo_ranges, IVector_h &neighbor_flags, I64Vector_h &part_offsets_h)\n{\n // This function creates arrays neighbors, halo_ranges_h and halo_ranges\n // base on neighbor_flags\n // Here do an MPI_allgather to figure out which partitions need data from me\n // This is required for non-symmetric matrices\n int num_part = A.manager->get_num_partitions();\n int my_id = A.manager->global_id();\n // pack 0/1 array into array of integers (size/32)\n int packed_size = (num_part + 32 - 1) / 32;\n IVector_h packed_nf(packed_size, 0);\n\n for (int i = 0; i < num_part; i++)\n {\n int packed_pos = i / 32;\n int bit_pos = i % 32;\n packed_nf[packed_pos] += (neighbor_flags[i] << bit_pos);\n }\n\n // exchange packed neighbor flags\n IVector_h gathered_packed_nf;\n A.manager->getComms()->all_gather_v(packed_nf, gathered_packed_nf, num_part);\n // assign neighbors that have edges to me\n int my_id_pos = my_id / 32;\n int my_id_bit = my_id % 32;\n\n for (int i = 0; i < num_part; i++)\n if (gathered_packed_nf[i * packed_size + my_id_pos] & (1 << my_id_bit)) // check my bit\n {\n neighbor_flags[i] = 1;\n }\n\n // compute total number of new neighbors\n int new_neighbors = amgx::thrust::reduce(neighbor_flags.begin(), neighbor_flags.end());\n cudaCheckError();\n // save old size\n int old_neighbors = neighbors.size();\n neighbors.resize(old_neighbors + new_neighbors);\n halo_ranges_h.resize(old_neighbors * 2 + new_neighbors * 2);\n // initialize manager->neighbors and manager->halo_ranges_h for the new nodes\n int active_part = old_neighbors;\n this->num_part = num_part;\n\n for (int i = 0; i < num_part; i++)\n {\n if (neighbor_flags[i] > 0)\n {\n neighbors[active_part] = i;\n halo_ranges_h[2 * active_part] = part_offsets_h[i];\n halo_ranges_h[2 * active_part + 1] = part_offsets_h[i + 1];\n active_part++;\n }\n }\n\n halo_ranges.resize(old_neighbors * 2 + new_neighbors * 2);\n amgx::thrust::copy(halo_ranges_h.begin() + old_neighbors * 2, halo_ranges_h.end(), halo_ranges.begin() + old_neighbors * 2);\n cudaCheckError();\n}\n\n\ntemplate \nvoid DistributedArranger >::create_neighbors_v2(Matrix_d &A)\n{\n // This function counts the number of neighbors, identifies the id of each neighbor,\n // And creates array halo_ranges, which is the range of rows for each neighbor\n // Outputs: A.manager->neighbors\n // A.manager->halo_ranges_h\n // A.manager->halo_ranges\n int num_part = A.manager->get_num_partitions();\n int my_id = A.manager->global_id();\n A.manager->set_base_index(A.manager->part_offsets_h[my_id]);\n A.manager->set_index_range(A.manager->part_offsets_h[my_id + 1] - A.manager->part_offsets_h[my_id]);\n IVector inv_neighbors(num_part, 0);\n IVector part_sizes(num_part, 0);\n int size = A.get_num_rows();\n\n if (size == 0)\n {\n A.manager->neighbors.resize(0);\n }\n\n // TODO: Are these the optimal block_sizes?\n int num_blocks = std::min(4096, (size + 127) / 128);\n\n // -----------------------------------------------\n // Step 1: flag neighbors that I have an edge to\n // -----------------------------------------------\n //go through all the column indices and check if they are halo vertices\n if (num_blocks != 0)\n {\n calc_num_neighbors_v2_local<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), A.manager->part_offsets.raw(), part_sizes.raw(), A.manager->local_to_global_map.raw(), num_part, my_id, size);\n cudaCheckError();\n }\n\n // D2H copy\n IVector_h neighbor_flags(num_part);\n amgx::thrust::copy(part_sizes.begin(), part_sizes.end(), neighbor_flags.begin());\n cudaCheckError();\n append_neighbors(A, A.manager->neighbors, A.manager->halo_ranges_h, A.manager->halo_ranges, neighbor_flags, A.manager->part_offsets_h);\n}\n\ntemplate \nvoid DistributedArranger >::create_neighbors_v2_global(Matrix_d &A, I64Vector_d &A_col_indices_global)\n{\n // This function counts the number of neighbors, identifies the id of each neighbor,\n // And creates array halo_ranges, which is the range of rows for each neighbor\n // Outputs: A.manager->neighbors\n // A.manager->halo_ranges_h\n // A.manager->halo_ranges\n int num_part = A.manager->get_num_partitions();\n int my_id = A.manager->global_id();\n IVector part_sizes(num_part, 0);\n int size = A.get_num_rows();\n\n if (size == 0)\n {\n A.manager->neighbors.resize(0);\n }\n\n // TODO: Are these the optimal block_sizes?\n int num_blocks = std::min(4096, (size + 127) / 128);\n\n // -----------------------------------------------\n // Step 1: flag neighbors that I have an edge to\n // -----------------------------------------------\n //go through all the column indices and check if they are halo vertices\n if (num_blocks != 0)\n {\n calc_num_neighbors_v2_global<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A_col_indices_global.raw(), A.manager->part_offsets.raw(), part_sizes.raw(), num_part, my_id, size);\n cudaCheckError();\n }\n\n // D2H copy\n IVector_h neighbor_flags(num_part);\n amgx::thrust::copy(part_sizes.begin(), part_sizes.end(), neighbor_flags.begin());\n cudaCheckError();\n append_neighbors(A, A.manager->neighbors, A.manager->halo_ranges_h, A.manager->halo_ranges, neighbor_flags, A.manager->part_offsets_h);\n}\n\ntemplate \nvoid DistributedArranger >::create_B2L_one_ring(Matrix_d &A)\n{\n std::vector boundary_lists;\n // create my boundary lists, renumber halos, create L2H maps\n create_boundary_lists_v2(A, boundary_lists);\n // compute halo offsets\n int neighbors = A.manager->num_neighbors();\n A.manager->halo_offsets.resize(neighbors + 1, 0);\n\n for (int i = 0; i < neighbors; i++)\n {\n A.manager->halo_offsets[i] = boundary_lists[i].size();\n }\n\n thrust_wrapper::exclusive_scan(A.manager->halo_offsets.begin(), A.manager->halo_offsets.end(), A.manager->halo_offsets.begin(), A.get_num_rows());\n cudaCheckError();\n // receive neighbors boundary lists for b2l maps\n A.manager->getComms()->exchange_vectors(boundary_lists, A, 0);\n // copy boundary lists to B2L maps\n // release internal boundary lists array\n A.manager->B2L_maps.resize(neighbors);\n\n for (int i = 0; i < neighbors; i++)\n {\n A.manager->B2L_maps[i] = boundary_lists[i];\n }\n\n A.manager->B2L_rings.resize(neighbors);\n\n for (int i = 0; i < neighbors; i++)\n {\n A.manager->B2L_rings[i].resize(2);\n A.manager->B2L_rings[i][0] = 0;\n A.manager->B2L_rings[i][1] = A.manager->B2L_maps[i].size();\n }\n\n this->initialize_B2L_maps_offsets(A, 1);\n A.manager->A = &A;\n}\n\ntemplate \nvoid DistributedArranger >::createRowsLists(Matrix_d &A, bool is_matrix_P)\n{\n int A_num_rows, offset;\n A.getOffsetAndSizeForView(OWNED, &offset, &A_num_rows);\n IVector boundary_flags(A_num_rows + 1, 0);\n IVector interior_flags(A_num_rows + 1, 0);\n\n int num_nz_interior;\n if (!is_matrix_P)\n {\n for (int i = 0; i < A.manager->neighbors.size(); i++)\n {\n int size = A.manager->B2L_rings[i][1];\n int num_blocks = std::min(4096, (size + 127) / 128);\n\n if ( size > 0)\n {\n flag_boundary_nodes <<< num_blocks, 128>>>(A.manager->B2L_maps[i].raw(), boundary_flags.raw(), size);\n cudaCheckError();\n }\n }\n }\n else\n {\n const int cta_size = 256;\n const int nWarps = cta_size / 32;\n int grid_size = std::min( 4096, (int) (A_num_rows + nWarps - 1) / nWarps);\n\n if (A_num_rows > 0)\n {\n IVector nnz_interior(1, 0);\n flag_boundary_nodes_P <<< grid_size, cta_size>>>(A.row_offsets.raw(), A.col_indices.raw(), boundary_flags.raw(), A.manager->halo_offsets[0], A_num_rows, nnz_interior.raw());\n cudaCheckError();\n hipMemcpy(&num_nz_interior, nnz_interior.raw(), sizeof(int), hipMemcpyDefault);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n int size = A_num_rows;\n int num_blocks = std::min(4096, (size + 127) / 128);\n\n if (size > 0)\n {\n flag_interior_nodes <<< num_blocks, 128>>>(boundary_flags.raw(), interior_flags.raw(), size);\n cudaCheckError();\n }\n\n if (boundary_flags.size() > 1)\n {\n thrust_wrapper::exclusive_scan(boundary_flags.begin(), boundary_flags.end(), boundary_flags.begin());\n cudaCheckError();\n }\n\n thrust_wrapper::exclusive_scan(interior_flags.begin(), interior_flags.end(), interior_flags.begin());\n cudaCheckError();\n int num_interior_rows = interior_flags[size];\n int num_boundary_rows = boundary_flags[size];\n A.manager->interior_rows_list.resize(num_interior_rows);\n A.manager->boundary_rows_list.resize(num_boundary_rows);\n\n // Create the list of interior and boundary nodes, using the interior_flag and boundary_flag\n if (size > 0)\n {\n fill_rows_lists <<< num_blocks, 128>>>(A.manager->interior_rows_list.raw(), A.manager->boundary_rows_list.raw(), interior_flags.raw(), boundary_flags.raw(), size);\n cudaCheckError();\n }\n\n // Manually set the view sizes for P\n if(is_matrix_P)\n {\n int num_rows_owned = num_interior_rows + num_boundary_rows;\n int num_nz_owned = A.row_offsets[num_rows_owned];\n A.manager->setViewSizes(num_interior_rows, num_nz_interior, num_rows_owned, num_nz_owned, num_rows_owned, num_nz_owned, A.get_num_rows(), A.get_num_nz());\n }\n}\n\ntemplate \nvoid DistributedArranger >::create_boundary_lists_v3(Matrix_d &A)\n{\n // -----------------------------------------------\n // Flag all nodes in neighbors that I will need (that is I have an edge pointing to it)\n // -----------------------------------------------\n int A_num_rows = A.get_num_rows();\n\n if (A_num_rows == 0)\n {\n return;\n }\n\n int num_neighbors = A.manager->neighbors.size();\n // compute halo offsets\n int total_rows_of_neighbors = 0;\n std::vector halo_offsets(num_neighbors + 1, 0);\n int max_neighbor_size = 0;\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n int num_rows_neighbor = A.manager->halo_ranges_h[2 * i + 1] - A.manager->halo_ranges_h[2 * i];\n total_rows_of_neighbors += num_rows_neighbor;\n halo_offsets[i + 1] = halo_offsets[i] + num_rows_neighbor;\n max_neighbor_size = max_neighbor_size > num_rows_neighbor ? max_neighbor_size : num_rows_neighbor;;\n }\n\n // copy offsets to device\n IVector halo_offsets_d(num_neighbors + 1);\n amgx::thrust::copy(halo_offsets.begin(), halo_offsets.end(), halo_offsets_d.begin());\n // flag neighbors nodes that I need\n IVector halo_nodes(total_rows_of_neighbors + 1);\n thrust_wrapper::fill(halo_nodes.begin(), halo_nodes.end(), 0);\n cudaCheckError();\n int size = A.manager->local_to_global_map.size();\n int num_blocks = std::min(4096, (size + 127) / 128);\n\n if (size > 0)\n {\n flag_halo_nodes_from_local_to_global_map <<< num_blocks, 128>>>(A.manager->halo_ranges.raw(), halo_nodes.raw(), halo_offsets_d.raw(), A.manager->base_index(), A.manager->index_range(), num_neighbors, size, A.manager->local_to_global_map.raw());\n cudaCheckError();\n }\n\n IVector halo_indices(max_neighbor_size + 1);\n A.manager->L2H_maps.resize(num_neighbors);\n A.manager->B2L_maps.resize(num_neighbors);\n is_less_than_zero pred;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n // compute halo indices map [] = 0-based halo number\n int size = halo_offsets[i + 1] - halo_offsets[i];\n amgx::thrust::replace_copy_if(halo_nodes.begin() + halo_offsets[i], halo_nodes.begin() + halo_offsets[i + 1], halo_indices.begin(), pred, 1);\n thrust_wrapper::exclusive_scan(halo_indices.begin(), halo_indices.begin() + size + 1, halo_indices.begin());\n cudaCheckError();\n int num_halo = halo_indices[size];\n A.manager->L2H_maps[i].resize(num_halo);\n A.manager->B2L_maps[i].resize(num_halo);\n // create my boundary lists = list of neighbor inner nodes corresponding to halo numbers 0..M\n // For each node in B2L_map, need to find corresponding local id and store in L2H_map\n int num_blocks = std::min(4096, (size + 127) / 128);\n populate_bdy_list_and_l2h_map <<< num_blocks, 128>>>(A.manager->L2H_maps[i].raw(), halo_nodes.raw() + halo_offsets[i], halo_indices.raw(), A.manager->B2L_maps[i].raw(), size);\n cudaCheckError();\n }\n}\n\n\ntemplate \nvoid DistributedArranger >::initialize_B2L_maps_offsets(Matrix_d &A, int num_rings)\n{\n // Compute the total number of bdy rows in each ring\n std::vector B2L_maps_offsets_h(num_rings);\n A.manager->B2L_rings_sizes.resize(num_rings);\n int num_neighbors = A.manager->neighbors.size();\n\n for (int k = 0; k < num_rings; k++)\n {\n B2L_maps_offsets_h[k].resize(num_neighbors + 1);;\n B2L_maps_offsets_h[k][0] = 0;\n\n for (int j = 0; j < num_neighbors; j++)\n {\n B2L_maps_offsets_h[k][j + 1] = B2L_maps_offsets_h[k][j] + A.manager->B2L_rings[j][k + 1];\n }\n\n A.manager->B2L_rings_sizes[k] = B2L_maps_offsets_h[k][num_neighbors];\n }\n\n // Copy maps_offsets to device\n A.manager->B2L_maps_offsets.resize(num_rings);\n\n for (int i = 0; i < num_rings; i++)\n {\n A.manager->B2L_maps_offsets[i] = B2L_maps_offsets_h[i];\n }\n\n // Store the B2L_maps ptrs on the device\n std::vector B2L_maps_ptrs_h(num_neighbors);\n\n for (int j = 0; j < num_neighbors; j++)\n {\n B2L_maps_ptrs_h[j] = A.manager->B2L_maps[j].raw();\n }\n\n A.manager->B2L_maps_ptrs = B2L_maps_ptrs_h;\n}\n\n\ntemplate \nvoid DistributedArranger >::create_one_ring_halo_rows(Matrix_d &A)\n{\n // what it does:\n // appends halo rows to matrix A\n // creates new B2L_rings, B2L_maps and L2H_maps\n // input:\n // matrix A and 1-ring B2L_maps\n int num_partitions = A.manager->get_num_partitions();\n int my_id = A.manager->global_id();\n int num_neighbors = A.manager->B2L_maps.size();\n int num_ring1_neighbors = num_neighbors;\n int num_ring1_indices = A.manager->local_to_global_map.size();\n std::vector halo_row_offsets(num_neighbors);\n std::vector halo_global_indices(num_neighbors);\n std::vector halo_local_indices(num_neighbors);\n std::vector halo_values(num_neighbors);\n // step 1: setup halo rows with global indices\n create_halo_rows_global_indices(A, halo_row_offsets, halo_global_indices, halo_values);\n // step 2: exchange halo rows\n std::vector dummy_halo_ids(0);\n A.manager->getComms()->exchange_matrix_halo(halo_row_offsets, halo_global_indices, halo_values, dummy_halo_ids, A.manager->neighbors, A.manager->global_id());\n // step 3: append the list f neighbors with the new ring-2 neighbors\n update_neighbors_list(A, A.manager->neighbors, A.manager->halo_ranges_h, A.manager->halo_ranges, A.manager->part_offsets_h, A.manager->part_offsets, halo_row_offsets, halo_global_indices);\n num_neighbors = A.manager->neighbors.size();\n // step 4: mark neighbor nodes, create local halo indices for the 2nd ring\n // this function also creates boundary lists for the 2nd ring if boundary lists array size > 0\n std::vector boundary_lists(num_neighbors);\n compute_local_halo_indices(A.row_offsets, A.col_indices, halo_row_offsets, halo_global_indices, halo_local_indices, A.manager->local_to_global_map, boundary_lists, A.manager->neighbors, A.manager->halo_ranges_h, A.manager->halo_ranges, A.manager->halo_offsets, A.manager->base_index(), A.manager->index_range(), A.get_num_rows(), 1);\n // update renumbering arrays (set identity for 2nd ring)\n //A.manager->halo_P_renumbering.resize(A.manager->local_to_global_map.size());\n //A.manager->halo_R_renumbering.resize(A.manager->local_to_global_map.size());\n //thrust_wrapper::sequence(A.manager->halo_P_renumbering.begin() + num_ring1_indices, A.manager->halo_P_renumbering.end(), num_ring1_indices + A.get_num_rows());\n //thrust_wrapper::sequence(A.manager->halo_R_renumbering.begin() + num_ring1_indices, A.manager->halo_R_renumbering.end(), num_ring1_indices + A.get_num_rows());\n // step 5: update L2H maps = identity\n A.manager->getComms()->set_neighbors(A.manager->neighbors.size());\n // step 6: receive neighbors boundary lists for 2-ring b2l maps\n A.manager->getComms()->exchange_vectors(boundary_lists, A, 0);\n // step 7: update B2L rings & B2L maps\n int rings = 2;\n A.manager->B2L_rings.resize(num_neighbors);\n A.manager->B2L_maps.resize(num_neighbors);\n // modify 1-ring where necessary\n int ring = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n // set size to 0 for 2-ring neighbors only\n if (i >= num_ring1_neighbors)\n {\n A.manager->B2L_maps[i].resize(0);\n }\n\n A.manager->B2L_rings[i].resize(rings + 1);\n A.manager->B2L_rings[i][0] = 0;\n A.manager->B2L_rings[i][ring + 1] = A.manager->B2L_maps[i].size();\n }\n\n // fill up 2-ring maps\n ring = 1;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n // append 2nd ring\n int ring1_size = A.manager->B2L_maps[i].size();\n A.manager->B2L_maps[i].resize(ring1_size + boundary_lists[i].size());\n amgx::thrust::copy(boundary_lists[i].begin(), boundary_lists[i].end(), A.manager->B2L_maps[i].begin() + ring1_size);\n A.manager->B2L_rings[i][ring + 1] = A.manager->B2L_maps[i].size();\n }\n\n cudaCheckError();\n this->initialize_B2L_maps_offsets(A, 2);\n // step 8: append halo_rows to matrix A\n // compute new # of rows & nnz\n this->append_halo_rows(A, halo_row_offsets, halo_local_indices, halo_values);\n // initialize the manager\n A.manager->set_initialized(A.row_offsets);\n // Compute the diagonal\n // TODO: Should only compute diagonal of 1-ring halo rows\n ViewType oldView = A.currentView();\n A.setView(FULL);\n A.set_allow_recompute_diag(true);\n A.computeDiagonal();\n A.setView(oldView);\n // the following steps are necessary only for latency hiding/renumbering, i.e. to use reorder_matrix()\n#if 0\n // step 6: for 2 rings: setup halo_btl\n create_halo_btl_multiple_rings(A, 2);\n#endif\n}\n\n// need for CUDA 4.2\ntemplate \nvoid DistributedArranger >::create_boundary_lists(Matrix_d &A,\n int64_t base_index, INDEX_TYPE index_range,\n I64Vector_h &halo_ranges_h, I64Vector_d &halo_ranges,\n std::vector &boundary_lists, IVector &halo_nodes)\n{\n //\n // Step 2: flag all nodes in neighbors that I will need (that is I have an edge pointing to it)\n //\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n int num_neighbors = halo_ranges_h.size() / 2;\n halo_ranges.resize(num_neighbors * 2);\n amgx::thrust::copy(halo_ranges_h.begin(), halo_ranges_h.end(), halo_ranges.begin());\n cudaCheckError();\n boundary_lists.resize(num_neighbors);\n int total_rows_of_neighbors = 0;\n std::vector halo_offsets(num_neighbors + 1, 0);\n int max_neighbor_size = 0;\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n total_rows_of_neighbors += halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n halo_offsets[i + 1] = halo_offsets[i] + halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n max_neighbor_size = max_neighbor_size > (halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i]) ? max_neighbor_size : (halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i]);;\n }\n\n IVector halo_offsets_d(num_neighbors + 1);\n amgx::thrust::copy(halo_offsets.begin(), halo_offsets.end(), halo_offsets_d.begin());\n halo_nodes.resize(total_rows_of_neighbors);\n thrust_wrapper::fill(halo_nodes.begin(), halo_nodes.end(), 0);\n cudaCheckError();\n flag_halo_nodes<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), halo_ranges.raw(), halo_nodes.raw(), halo_offsets_d.raw(), base_index, index_range, num_neighbors, size);\n cudaCheckError();\n //\n // Step 3: Write the 1-ring halo indices for each neighbor\n //\n IVector halo_indices(max_neighbor_size);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int last = halo_nodes[halo_offsets[i + 1] - 1];\n thrust_wrapper::exclusive_scan(halo_nodes.begin() + halo_offsets[i], halo_nodes.begin() + halo_offsets[i + 1], halo_indices.begin());\n int num_halo = halo_indices[halo_offsets[i + 1] - halo_offsets[i] - 1] + last;\n //A.manager.B2L_maps[i].resize(num_halo);\n boundary_lists[i].resize(num_halo);\n int size = halo_offsets[i + 1] - halo_offsets[i];\n //int num_blocks=min(4096,(size+127)/128);\n populate_B2L <<< num_blocks, 128>>>(halo_indices.raw(), boundary_lists[i].raw(),/*A.manager.B2L_maps[i].raw(),*/ last, size);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::create_boundary_lists(Matrix_d &A,\n I64Vector_h &halo_ranges_h,\n std::vector &boundary_lists, IVector &halo_nodes)\n{\n //\n // Step 2: flag all nodes in neighbors that I will need (that is I have an edge pointing to it)\n //\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n int num_neighbors = halo_ranges_h.size() / 2;\n A.manager->halo_ranges.resize(num_neighbors * 2);\n amgx::thrust::copy(halo_ranges_h.begin(), halo_ranges_h.end(), A.manager->halo_ranges.begin());\n cudaCheckError();\n boundary_lists.resize(num_neighbors);\n int total_rows_of_neighbors = 0;\n std::vector halo_offsets(num_neighbors + 1, 0);\n int max_neighbor_size = 0;\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n total_rows_of_neighbors += halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n halo_offsets[i + 1] = halo_offsets[i] + halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n max_neighbor_size = max_neighbor_size > (halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i]) ? max_neighbor_size : (halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i]);;\n }\n\n IVector halo_offsets_d(num_neighbors + 1);\n amgx::thrust::copy(halo_offsets.begin(), halo_offsets.end(), halo_offsets_d.begin());\n halo_nodes.resize(total_rows_of_neighbors);;\n thrust_wrapper::fill(halo_nodes.begin(), halo_nodes.end(), 0);\n cudaCheckError();\n flag_halo_nodes<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), A.manager->halo_ranges.raw(), halo_nodes.raw(), halo_offsets_d.raw(), A.manager->base_index(), A.manager->index_range(), num_neighbors, size);\n cudaCheckError();\n //\n // Step 3: Write the 1-ring halo indices for each neighbor\n //\n IVector halo_indices(max_neighbor_size);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int last = halo_nodes[halo_offsets[i + 1] - 1];\n thrust_wrapper::exclusive_scan(halo_nodes.begin() + halo_offsets[i], halo_nodes.begin() + halo_offsets[i + 1], halo_indices.begin());\n int num_halo = halo_indices[halo_offsets[i + 1] - halo_offsets[i] - 1] + last;\n //A.manager.B2L_maps[i].resize(num_halo);\n boundary_lists[i].resize(num_halo);\n int size = halo_offsets[i + 1] - halo_offsets[i];\n //int num_blocks=min(4096,(size+127)/128);\n populate_B2L <<< num_blocks, 128>>>(halo_indices.raw(), boundary_lists[i].raw(),/*A.manager.B2L_maps[i].raw(),*/ last, size);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::create_boundary_lists_v2(Matrix_d &A, std::vector &boundary_lists)\n{\n // -----------------------------------------------\n // Flag all nodes in neighbors that I will need (that is I have an edge pointing to it)\n // -----------------------------------------------\n int size = A.get_num_rows();\n\n if (size == 0)\n {\n return;\n }\n\n //TODO: Are these the optimal block_sizes?\n int num_blocks = std::min(4096, (size + 127) / 128);\n int num_neighbors = A.manager->neighbors.size();\n boundary_lists.resize(num_neighbors);\n // compute halo offsets\n int total_rows_of_neighbors = 0;\n std::vector halo_offsets(num_neighbors + 1, 0);\n int max_neighbor_size = 0;\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n int num_rows_neighbor = A.manager->halo_ranges_h[2 * i + 1] - A.manager->halo_ranges_h[2 * i];\n total_rows_of_neighbors += num_rows_neighbor;\n halo_offsets[i + 1] = halo_offsets[i] + num_rows_neighbor;\n max_neighbor_size = max_neighbor_size > num_rows_neighbor ? max_neighbor_size : num_rows_neighbor;;\n }\n\n // copy offsets to device\n IVector halo_offsets_d(num_neighbors + 1);\n amgx::thrust::copy(halo_offsets.begin(), halo_offsets.end(), halo_offsets_d.begin());\n // flag neighbors nodes that I need\n IVector halo_nodes(total_rows_of_neighbors);\n amgx::thrust::fill(halo_nodes.begin(), halo_nodes.end(), 0);\n cudaCheckError();\n flag_halo_nodes_local<16, 1> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), A.manager->halo_ranges.raw(), halo_nodes.raw(), halo_offsets_d.raw(), A.manager->base_index(), A.manager->index_range(), num_neighbors, size, A.manager->local_to_global_map.raw());\n cudaCheckError();\n // store renumbering so later can use on P\n //A.manager->halo_P_renumbering.resize(A.manager->local_to_global_map.size());\n //A.manager->halo_R_renumbering.resize(A.manager->local_to_global_map.size());\n // we need to update local toglobal maps as well because we will be renumbering halo nodes\n I64Vector new_local_to_global_map(A.manager->local_to_global_map.size(), -1);\n // compute new local halo indices for each neighbor node\n // renumbering: compute new column indices sorted by global index\n IVector halo_indices(max_neighbor_size);\n std::vector num_halo(num_neighbors);\n std::vector last(num_neighbors);\n std::vector halo_lists;\n halo_lists.resize(num_neighbors);\n A.manager->L2H_maps.resize(num_neighbors);\n int halo_base = size;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n // compute halo indices map [] = 0-based halo number\n last[i] = halo_nodes[halo_offsets[i + 1] - 1];\n thrust_wrapper::exclusive_scan(halo_nodes.begin() + halo_offsets[i], halo_nodes.begin() + halo_offsets[i + 1], halo_indices.begin());\n num_halo[i] = halo_indices[halo_offsets[i + 1] - halo_offsets[i] - 1] + last[i];\n // renumber my columns for this particular neighbor\n calc_new_halo_mapping<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), A.manager->halo_ranges.raw() + (int64_t) (2 * i), 1, size, A.manager->local_to_global_map.raw(), new_local_to_global_map.raw(), halo_indices.raw(), halo_base, A.manager->global_id());\n cudaCheckError();\n //calc_new_halo_mapping_v2<16><<>>( A.manager->halo_ranges.raw() + (int64_t) (2*i), 1, size, A.manager->local_to_global_map.raw(), new_local_to_global_map.raw(), halo_indices.raw(), halo_base, A.manager->local_to_global_map.size(),A.manager->global_id());\n cudaCheckError();\n // create my boundary lists = list of neighbor inner nodes corresponding to halo numbers 0..M\n boundary_lists[i].resize(num_halo[i]);\n int size = halo_offsets[i + 1] - halo_offsets[i];\n populate_B2L <<< num_blocks, 128>>>(halo_indices.raw(), boundary_lists[i].raw(), last[i], size);\n cudaCheckError();\n // identity L2H\n A.manager->L2H_maps[i].resize(num_halo[i]);\n thrust_wrapper::sequence(A.manager->L2H_maps[i].begin(), A.manager->L2H_maps[i].end(), halo_base);\n halo_lists[i] = &A.manager->L2H_maps[i];\n // continous numbering for halo indices\n halo_base += num_halo[i];\n }\n\n cudaCheckError();\n // Now change the negative col to positive\n flip_sign_col<16> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.col_indices.raw(), size);\n cudaCheckError();\n // copy new local to global mapping\n amgx::thrust::copy(new_local_to_global_map.begin(), new_local_to_global_map.end(), A.manager->local_to_global_map.begin());\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::compute_local_halo_indices( IVector &A_row_offsets, IVector &A_col_indices, std::vector &halo_row_offsets, std::vector &halo_global_indices,\n std::vector &halo_local_indices, I64Vector &local_to_global, std::vector &boundary_lists, IVector_h &neighbors, I64Vector_h &halo_ranges_h, I64Vector &halo_ranges, IVector_h &halo_offsets, int64_t base_index, int64_t index_range, int A_num_rows, int current_num_rings)\n{\n // This function checks the halo_col_indices received from the neighbors, and identifies\n // new halo_indices and updates halo_offsets, local_to_global_map accordingly\n // input: halo row offsets,\n // halo global column indices\n // current local to global map\n // neighbors (new discovered neighbors should already be included)\n // halo_ranges for the neighbors\n //\n // output: halo offsets,\n // halo local column indices,\n // updated local to global map (in place)\n int size = A_num_rows;\n //TODO: Are these the optimal block_sizes?\n int num_blocks = std::min(4096, (size + 127) / 128);\n int num_neighbors = neighbors.size();\n // compute neighbor offsets & total number of neighbor rows\n int total_rows_of_neighbors = 0;\n std::vector neighbor_offsets_h(num_neighbors + 1, 0);\n int max_neighbor_size = 0;\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n total_rows_of_neighbors += halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n neighbor_offsets_h[i + 1] = neighbor_offsets_h[i] + halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i];\n max_neighbor_size = max_neighbor_size > (halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i]) ? max_neighbor_size : (halo_ranges_h[2 * i + 1] - halo_ranges_h[2 * i]);\n }\n\n // copy offsets to device\n IVector neighbor_offsets(num_neighbors + 1);\n amgx::thrust::copy(neighbor_offsets_h.begin(), neighbor_offsets_h.end(), neighbor_offsets.begin());\n // store flags for all neighbor nodes\n IVector neighbor_nodes(total_rows_of_neighbors);\n // flag neighbor nodes that are in the existing rings as -(local_index)-1\n thrust_wrapper::fill(neighbor_nodes.begin(), neighbor_nodes.end(), 0);\n cudaCheckError();\n int local_to_global_size = local_to_global.size();\n int num_blocks2 = std::min(4096, (local_to_global_size + 127) / 128);\n\n if (local_to_global_size != 0)\n {\n flag_halo_nodes_local_v2 <<< num_blocks2, 128>>>(halo_ranges.raw(), neighbor_nodes.raw(), neighbor_offsets.raw(), base_index, index_range, num_neighbors, local_to_global_size, local_to_global.raw());\n cudaCheckError();\n }\n\n // 1) flag NEW neighbors nodes that I need as 1, they will be in the 2nd ring\n // 2) fill out local indices for 1st ring & internal indices\n int num_halos = halo_row_offsets.size();\n\n for (int i = 0; i < num_halos; i++)\n {\n int size = halo_global_indices[i].size();\n\n if (size > 0)\n {\n halo_local_indices[i].resize(size);\n thrust_wrapper::fill(halo_local_indices[i].begin(), halo_local_indices[i].end(), -1); // fill with -1\n // TODO: launch only on halo rows\n flag_halo_nodes_global<16, 1> <<< num_blocks, 128>>>(halo_row_offsets[i].raw(), halo_row_offsets[i].size() - 1, halo_global_indices[i].raw(), halo_ranges.raw(), neighbor_nodes.raw(), neighbor_offsets.raw(), base_index, index_range, num_neighbors, halo_local_indices[i].raw());\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n // replace all negative values with 0 in neighbor flags\n is_less_than_zero pred;\n amgx::thrust::replace_if(neighbor_nodes.begin(), neighbor_nodes.end(), pred, 0);\n cudaCheckError();\n // fill halo offsets for the current number of ring for new neighbors (it will be of size 0)\n int current_num_neighbors = (halo_offsets.size() - 1) / current_num_rings;;\n int num_rings = current_num_rings + 1;\n IVector_h new_halo_offsets(num_rings * num_neighbors + 1);\n\n for (int j = 0; j < current_num_rings; j++)\n {\n for (int i = 0 ; i <= current_num_neighbors; i++)\n {\n new_halo_offsets[num_neighbors * j + i] = halo_offsets[current_num_neighbors * j + i];\n }\n\n for (int i = current_num_neighbors; i < num_neighbors; i++)\n {\n new_halo_offsets[num_neighbors * j + i + 1] = new_halo_offsets[num_neighbors * j + i];\n }\n }\n\n halo_offsets = new_halo_offsets;\n int ring = current_num_rings;\n int current_num_halo_indices = local_to_global.size();\n //int halo_base = index_range + current_num_halo_indices;\n cudaCheckError();\n\n // compute neighbors nodes indices (in-place) for each neighbor\n for (int i = 0; i < num_neighbors; i++)\n {\n int last_node, num_halo;\n\n if (neighbor_offsets_h[i + 1] != neighbor_offsets_h[i])\n {\n last_node = neighbor_nodes[neighbor_offsets_h[i + 1] - 1];\n thrust_wrapper::exclusive_scan(neighbor_nodes.begin() + neighbor_offsets_h[i], neighbor_nodes.begin() + neighbor_offsets_h[i + 1], neighbor_nodes.begin() + neighbor_offsets_h[i]);\n num_halo = neighbor_nodes[neighbor_offsets_h[i + 1] - 1] + last_node;\n }\n else\n {\n num_halo = 0;\n last_node = 0;\n }\n\n // update halo offsets (L2H)\n halo_offsets[ring * num_neighbors + i + 1] = halo_offsets[ring * num_neighbors + i] + num_halo;\n\n // if size = 0 then we don't need to compute it\n if (boundary_lists.size() > 0)\n {\n // create my boundary lists = list of neighbor inner nodes corresponding to halo numbers 0..M\n // basically this will be our 2-ring B2L_maps\n boundary_lists[i].resize(num_halo);\n int size = neighbor_offsets_h[i + 1] - neighbor_offsets_h[i];\n num_blocks = std::min(4096, (size + 127) / 128);\n\n if (size > 0)\n {\n populate_B2L <<< num_blocks, 128>>>(neighbor_nodes.raw() + neighbor_offsets_h[i], boundary_lists[i].raw(), last_node, size);\n cudaCheckError();\n }\n }\n }\n\n cudaCheckError();\n // compute local indices and new local to global mapping\n int new_num_halo_indices = halo_offsets[num_rings * num_neighbors] - halo_offsets[current_num_rings * num_neighbors];\n local_to_global.resize(current_num_halo_indices + new_num_halo_indices);\n IVector halo_offsets_d(halo_offsets.size());\n amgx::thrust::copy(halo_offsets.begin(), halo_offsets.end(), halo_offsets_d.begin());\n cudaCheckError();\n\n // do this for all ring-1 neighbors\n for (int i = 0; i < num_halos; i++)\n {\n int num_neighbor_rows = halo_row_offsets[i].size() - 1;\n num_blocks = std::min(4096, (num_neighbor_rows + 127) / 128);\n\n if (num_blocks > 0)\n {\n calc_new_halo_mapping_ring2<16> <<< num_blocks, 128>>>(\n halo_row_offsets[i].raw(), num_neighbor_rows, halo_global_indices[i].raw(), // halo rows to process\n halo_ranges.raw(), base_index, index_range, num_neighbors, // ranges and # of neighbors\n halo_offsets_d.raw() + current_num_rings * num_neighbors, neighbor_offsets.raw(), neighbor_nodes.raw(), // halo offsets, neighbor offsets and indices\n local_to_global.raw(), halo_local_indices[i].raw()); // output\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n}\n\n\ntemplate \nvoid DistributedArranger >::create_rings(Matrix_d &A, INDEX_TYPE rings,\n INDEX_TYPE neighbors,\n std::vector &B2L_maps,\n std::vector > &B2L_rings)\n{\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n B2L_rings.resize(neighbors);\n\n for (int i = 0; i < neighbors; i ++)\n {\n B2L_rings[i].resize(rings + 1);\n B2L_rings[i][0] = 0;\n }\n\n /*\n EXAMPLE\n For example matrix. partition 1, the matrix look like:\n A.row_offsets = [0 4 11 15 20]\n A.col_indices = [4 0 1 2\n 4 5 0 1 2 3 7\n 0 1 2 3\n 1 2 3 6 7]\n num_neighbors=2; neighbors = [0 2]\n B2L_rings[[0 2][0 2]] B2L_maps [[0 1][1 3]] (Converted from the following input: btl_sizes = [0 2 4] btl_maps = [0 1| 1 3])\n L2H_maps (and halo_lists) [[4 5][6 7]]\n and we want to map further halo rings\n */\n //Scratch space computation\n for (int i = 0; i < neighbors; i++)\n {\n B2L_rings[i][1] = B2L_maps[i].size();\n }\n\n //flag array\n IVector halo_rings(A.get_num_rows());\n\n for (int i = 0; i < neighbors; i++)\n {\n //\n // For each ring: flag nodes that are neighbors of the previous ring, then remove the nodes of the previous rings from the marking\n // count the number of marked nodes, and append their row indices to the B2L_maps\n //\n for (int ring = 1; ring < rings; ring++)\n {\n size = B2L_rings[i][ring] - B2L_rings[i][ring - 1];\n thrust_wrapper::fill(halo_rings.begin(), halo_rings.begin() + size + 1, 0);\n num_blocks = std::min(4096, (size + 127) / 128);\n //Set 1 in flag array for any column index appreaing in the rows that belong to the previous ring\n find_next_ring <<< num_blocks, 128>>>(B2L_maps[i].raw() + B2L_rings[i][ring - 1], A.row_offsets.raw(), A.col_indices.raw(), size, 0/*base_index*/, /*index_range*/size, halo_rings.raw());\n cudaCheckError();\n /*\n EXAMPLE\n for neighbor ID 0, halo_rings array initially [0 0 0 0], find_next_ring look at rows in the first ring [0 1] and sets halo_rings array positions to 1 for each column index in those rows:\n halo_rings: [1 1 1 1]\n */\n //Remove all flags that belong to rows in any previous rings\n remove_previous_rings <<< num_blocks, 128>>>(B2L_maps[i].raw(), halo_rings.raw(), B2L_rings[i][ring]);\n cudaCheckError();\n /*\n EXAMPLE\n for neighbor ID 0, halo_rings array [1 1 1 1], remove_previous_rings sets halo_rings positions to 0 for each index in previous B2L rings ([0 1] in this case):\n halo_rings: [0 0 1 1]\n */\n //Count, index and append them to B2L_maps and update B2L_rings\n thrust_wrapper::inclusive_scan(halo_rings.begin(), halo_rings.begin() + size + 1, halo_rings.begin());\n B2L_rings[i][ring + 1] = B2L_rings[i][ring] + halo_rings[size];\n B2L_maps[i].resize(B2L_rings[i][ring + 1]);\n\n /*\n EXAMPLE\n after inclusive scan: [0 0 1 2], B2L_rings[0][2] = 2+2, and B2L_maps[0] gets resized to 4\n */\n if (B2L_rings[i][ring] == B2L_rings[i][ring + 1])\n {\n for (int fill = ring; fill <= rings; fill++) { B2L_rings[i][fill] = B2L_rings[i][ring]; }\n\n break;\n }\n\n num_blocks = std::min(4096, (A.get_num_rows() + 127) / 128);\n write_ring_rows <<< num_blocks, 128>>>(halo_rings.raw(), B2L_maps[i].raw() + B2L_rings[i][ring], A.get_num_rows());\n cudaCheckError();\n /*\n EXAMPLE\n appends positions i to B2L_maps, where halo_rings[i-1] != halo_rings[i]\n B2L_maps[0] = [0 1| 2 3] with B2L_rings[0] = [0 2 4]\n */\n }\n\n cudaCheckError();\n }\n}\n\n// CUDA 4.2\ntemplate \nvoid DistributedArranger >::create_rings(Matrix_d &A, INDEX_TYPE rings)\n{\n int neighbors = A.manager->num_neighbors();\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n A.manager->B2L_rings.resize(neighbors * rings + 1);\n\n for (int i = 0; i < neighbors; i ++)\n {\n A.manager->B2L_rings[i].resize(rings + 1);\n A.manager->B2L_rings[i][0] = 0;\n }\n\n //Scratch space computation\n for (int i = 0; i < neighbors; i++)\n {\n A.manager->B2L_rings[i][1] = A.manager->B2L_maps[i].size();\n }\n\n IVector halo_rings(A.get_num_rows());\n\n for (int i = 0; i < neighbors; i++)\n {\n //\n // For each ring: flag nodes that are neighbors of the previous ring, then remove the nodes of the previous ring from the marking\n // count the number of marked nodes, and append their row indices to the B2L_maps\n //\n for (int ring = 1; ring < rings; ring++)\n {\n size = A.manager->B2L_rings[i][ring] - A.manager->B2L_rings[i][ring - 1];\n thrust_wrapper::fill(halo_rings.begin(), halo_rings.begin() + size + 1, 0);\n num_blocks = std::min(4096, (size + 127) / 128);\n find_next_ring <<< num_blocks, 128>>>(A.manager->B2L_maps[i].raw() + A.manager->B2L_rings[i][ring - 1], A.row_offsets.raw(), A.col_indices.raw(), size, 0/*base_index*/, /*index_range*/size, halo_rings.raw());\n cudaCheckError();\n remove_previous_rings <<< num_blocks, 128>>>(A.manager->B2L_maps[i].raw(), halo_rings.raw(), A.manager->B2L_rings[i][ring]);\n cudaCheckError();\n thrust_wrapper::inclusive_scan(halo_rings.begin(), halo_rings.begin() + size + 1, halo_rings.begin());\n A.manager->B2L_rings[i][ring + 1] = A.manager->B2L_rings[i][ring] + halo_rings[size];\n A.manager->B2L_maps[i].resize(A.manager->B2L_rings[i][ring + 1]);\n\n if (A.manager->B2L_rings[i][ring] == A.manager->B2L_rings[i][ring + 1])\n {\n for (int fill = ring; fill <= rings; fill++) { A.manager->B2L_rings[i][fill] = A.manager->B2L_rings[i][ring]; }\n\n break;\n }\n\n num_blocks = std::min(4096, (A.get_num_rows() + 127) / 128);\n write_ring_rows <<< num_blocks, 128>>>(halo_rings.raw(), A.manager->B2L_maps[i].raw() + A.manager->B2L_rings[i][ring], A.get_num_rows());\n }\n\n cudaCheckError();\n }\n}\n\ntemplate \nvoid DistributedArranger >::create_maps(Matrix_d &A, INDEX_TYPE rings,\n INDEX_TYPE neighbors,\n std::vector &B2L_maps,\n std::vector > &B2L_rings, std::vector &boundary_lists)\n{\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n B2L_maps.resize(neighbors);\n\n //\n // Step 4: export n-ring halo matrices\n //\n for (int i = 0; i < neighbors; i++)\n {\n B2L_maps[i] = boundary_lists[i];\n }\n\n create_rings(A, rings, neighbors, B2L_maps, B2L_rings);\n //B2L_rings.resize(neighbors*rings+1);\n}\n\n// need for CUDA 4.2\ntemplate \nvoid DistributedArranger >::create_maps(Matrix_d &A, INDEX_TYPE rings, std::vector &boundary_lists)\n{\n int size = A.get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n int neighbors = A.manager->num_neighbors();\n A.manager->B2L_maps.resize(neighbors);\n\n //\n // Step 4: export n-ring halo matrices\n //\n for (int i = 0; i < neighbors; i++)\n {\n A.manager->B2L_maps[i] = boundary_lists[i];\n }\n\n create_rings(A, rings);\n // A.manager->B2L_rings.resize(neighbors*rings+1);\n}\n\ntemplate \nvoid DistributedArranger >::create_halo_rows(Matrix_d &A, INDEX_TYPE rings, INDEX_TYPE neighbors, std::vector &B2L_maps,\n std::vector > &B2L_rings, std::vector **halo_rows_p, std::vector &halo_lists)\n{\n typedef typename MatPrecisionMap::Type ValueTypeA;\n *halo_rows_p = new std::vector >(neighbors);\n std::vector > &halo_rows = **halo_rows_p;\n\n if (this->halo_coloring == LAST) //Create dummy matrices\n {\n for (int i = 0; i < neighbors; i++)\n {\n halo_rows[i].addProps(CSR);\n\n if (A.hasProps(DIAG)) { halo_rows[i].addProps(DIAG); }\n\n halo_rows[i].resize(halo_lists[i]->size(), halo_lists[i]->size(), (A.hasProps(DIAG) ? 0 : halo_lists[i]->size()), A.get_block_dimy(), A.get_block_dimx(), 1);\n\n if (A.hasProps(DIAG))\n {\n thrust_wrapper::fill(halo_rows[i].row_offsets.begin(), halo_rows[i].row_offsets.end(), 0);\n }\n else\n {\n thrust_wrapper::sequence(halo_rows[i].row_offsets.begin(), halo_rows[i].row_offsets.end());\n amgx::thrust::copy(halo_lists[i]->begin(), halo_lists[i]->end(), halo_rows[i].col_indices.begin());\n }\n\n thrust_wrapper::fill(halo_rows[i].values.begin(), halo_rows[i].values.end(), types::util< ValueTypeA >::get_one());\n cudaCheckError();\n }\n\n return;\n }\n\n //Scratch space computation\n int max_size = 0;\n\n for (int i = 0; i < neighbors; i++)\n {\n max_size = max_size > B2L_maps[i].size() ? max_size : B2L_maps[i].size();\n //B2L_rings[i][1]=B2L_maps[i].size(); //I don't know what this was doing here, definitely stupid for rings > 1\n }\n\n IVector matrix_halo_sizes((max_size + 1)*rings);\n\n for (int i = 0; i < neighbors; i++)\n {\n //\n // Write the length of the rows in the order of B2L_maps, then calculate row_offsets\n //\n int size = B2L_rings[i][rings];\n matrix_halo_sizes.resize(size + 1);\n int num_blocks = std::min(4096, (size + 127) / 128);\n write_matrix_rowsize <<< num_blocks, 128>>>(B2L_maps[i].raw(), A.row_offsets.raw(), size, matrix_halo_sizes.raw());\n cudaCheckError();\n thrust_wrapper::exclusive_scan(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, matrix_halo_sizes.begin());\n int nnz_count = matrix_halo_sizes[size];\n //\n // Resize export halo matrix, and copy over the rows\n //\n halo_rows[i].addProps(CSR);\n\n if (A.hasProps(DIAG)) { halo_rows[i].addProps(DIAG); }\n\n halo_rows[i].resize(B2L_maps[i].size(), B2L_maps[i].size(), nnz_count, A.get_block_dimy(), A.get_block_dimx(), 1);\n export_matrix_elements<32> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.values.raw(), A.get_block_size(), B2L_maps[i].raw(), matrix_halo_sizes.raw(), halo_rows[i].values.raw(), A.col_indices.raw(), halo_rows[i].col_indices.raw(), size);\n cudaCheckError();\n amgx::thrust::copy(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, halo_rows[i].row_offsets.begin());\n\n if (A.hasProps(DIAG))\n {\n export_matrix_diagonal <<< num_blocks, 128>>>(A.values.raw() + A.row_offsets[A.get_num_rows()]*A.get_block_size(), A.get_block_size(), B2L_maps[i].raw(), halo_rows[i].values.raw() + halo_rows[i].row_offsets[halo_rows[i].get_num_rows()]*A.get_block_size(), size);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::initialize_manager(Matrix_d &A, Matrix_d &P, int num_owned_coarse_pts)\n{\n // This function initializes the manager of matrix P based on value of num_owned_coarse_pts\n // The following member of the manager are initialized:\n // global_id, num_partitions, comms (copied from manager of matrix A)\n // part_offsets_h, part_offsets, base_index, index_range, num_rows_global, halo_offsets[0] (those are initialized based on value of num_owned_coarse_pts of all partitions in communicator)\n if (P.manager == NULL)\n {\n P.manager = new DistributedManager();\n }\n\n P.setView(OWNED);\n // Copy some information from manager of matrix A\n P.manager->A = &P;\n P.manager->setComms(A.manager->getComms());\n index_type num_parts = A.manager->get_num_partitions();\n index_type my_rank = A.manager->global_id();\n P.manager->set_global_id(my_rank);\n P.manager->set_num_partitions(num_parts);\n // Do a global reduce to obtain part_offsets for RAP\n IVector_h gathered_count;\n A.manager->getComms()->all_gather(num_owned_coarse_pts, gathered_count, num_parts);\n // Compute P.manager->part_offsets\n P.manager->part_offsets_h.resize(num_parts + 1);\n P.manager->part_offsets_h[0] = (int64_t) 0;\n\n for (int i = 1; i < num_parts + 1; i++)\n {\n P.manager->part_offsets_h[i] = P.manager->part_offsets_h[i - 1] + (int64_t) gathered_count[i - 1];\n }\n\n P.manager->part_offsets = P.manager->part_offsets_h;\n // Compute P.manager->halo_ranges\n P.manager->set_base_index(P.manager->part_offsets_h[my_rank]);\n P.manager->set_index_range(num_owned_coarse_pts);\n P.manager->num_rows_global = P.manager->part_offsets_h[num_parts];\n // Set halo_offsets[0]\n P.manager->halo_offsets.resize(1);\n P.manager->halo_offsets[0] = num_owned_coarse_pts;\n}\n\ntemplate \nvoid DistributedArranger >::initialize_manager_from_global_col_indices(Matrix_d &P, I64Vector_d &P_col_indices_global)\n{\n // This function finds the number of neighbors of P\n // fills halo_offsets, local_to_global_map, B2L_maps based on global col indices of matrix P\n int num_owned_coarse_pts = P.manager->halo_offsets[0];\n // By checking the global col indices, identify neighbors\n P.manager->neighbors.resize(0);\n this->create_neighbors_v2_global(P, P_col_indices_global);\n int num_neighbors = P.manager->neighbors.size();\n P.manager->getComms()->set_neighbors(num_neighbors);\n P.manager->halo_offsets.resize(num_neighbors + 1, 0);\n int total_rows_of_neighbors = 0;\n std::vector flag_offsets_h(num_neighbors + 1, 0);\n\n for (int i = 0; i < num_neighbors; i ++)\n {\n int num_rows = (int) (P.manager->halo_ranges_h[2 * i + 1] - P.manager->halo_ranges_h[2 * i]);\n total_rows_of_neighbors += num_rows;\n flag_offsets_h[i + 1] = flag_offsets_h[i] + num_rows;\n }\n\n // Copy to device\n IVector flag_offsets_d(num_neighbors + 1);\n amgx::thrust::copy(flag_offsets_h.begin(), flag_offsets_h.end(), flag_offsets_d.begin());\n // Allocate halo_nodes array that will be used to flag halo nodes\n // Flag_offsets array can be used to determine where each neighbor begins\n IVector halo_nodes(total_rows_of_neighbors + 1);\n thrust_wrapper::fill(halo_nodes.begin(), halo_nodes.end(), 0);\n cudaCheckError();\n int num_blocks = std::min(4096, (P.get_num_rows() + 127) / 128);\n\n if (P.get_num_rows() > 0)\n {\n flag_halo_nodes_global_v2<16> <<< num_blocks, 128>>>(P.row_offsets.raw(), P_col_indices_global.raw(), P.manager->halo_ranges.raw(), halo_nodes.raw(), flag_offsets_d.raw(), P.manager->base_index(), P.manager->index_range(), num_neighbors, P.get_num_rows());\n cudaCheckError();\n }\n\n // Step 3: Write the 1-ring global nodes for each neighbor\n thrust_wrapper::exclusive_scan(halo_nodes.begin(), halo_nodes.end(), halo_nodes.begin());\n cudaCheckError();\n std::vector bdy_lists;\n bdy_lists.resize(num_neighbors);\n int count = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int num_halo = halo_nodes[flag_offsets_h[i + 1]] - count;\n P.manager->halo_offsets[i + 1] = P.manager->halo_offsets[i] + num_halo;\n bdy_lists[i].resize(num_halo);\n\n if (num_halo > 0)\n {\n int size = flag_offsets_h[i + 1] - flag_offsets_h[i];\n int num_blocks = std::min(4096, (size + 127) / 128);\n populate_B2L_v2 <<< num_blocks, 128>>>(halo_nodes.raw() + flag_offsets_h[i], bdy_lists[i].raw(), count, size);\n cudaCheckError();\n }\n\n count += num_halo;\n }\n\n cudaCheckError();\n\n // Fill the col_indices\n if (P.get_num_rows() > 0)\n {\n renumber_col_indices_global<16> <<< num_blocks, 128>>>(P.row_offsets.raw(), P.col_indices.raw(), P_col_indices_global.raw(), P.manager->halo_ranges.raw(), halo_nodes.raw(), flag_offsets_d.raw(), P.manager->base_index(), P.manager->index_range(), num_neighbors, P.get_num_rows(), P.manager->global_id());\n cudaCheckError();\n }\n\n // Create P.local_to_global_map\n P.manager->local_to_global_map.resize(count);\n count = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int num_halo = bdy_lists[i].size();\n\n if (num_halo > 0)\n {\n int num_blocks = std::min(4096, (num_halo + 127) / 128);\n fill_local_to_global_map <<< num_blocks, 128>>>(bdy_lists[i].raw(), P.manager->local_to_global_map.raw() + count, P.manager->halo_ranges_h[2 * i], num_halo);\n cudaCheckError();\n }\n\n count += num_halo;\n }\n\n cudaCheckError();\n P.manager->getComms()->exchange_vectors(bdy_lists, P, 0);\n // copy boundary lists to B2L maps\n // release internal boundary lists array\n P.manager->B2L_maps.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n P.manager->B2L_maps[i] = bdy_lists[i];\n }\n\n P.manager->B2L_rings.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n P.manager->B2L_rings[i].resize(2);\n P.manager->B2L_rings[i][0] = 0;\n P.manager->B2L_rings[i][1] = P.manager->B2L_maps[i].size();\n }\n\n this->initialize_B2L_maps_offsets(P, 1);\n P.manager->_num_interior_nodes = P.get_num_rows();\n P.manager->_num_boundary_nodes = 0;\n P.manager->set_initialized(P.row_offsets);\n int size_one_ring, offset;\n P.manager->getOffsetAndSizeForView(FULL, &offset, &size_one_ring);\n P.setView(OWNED);\n}\n\n\ntemplate \nvoid DistributedArranger >::createTempManager(Matrix_d &B, Matrix_d &A, IVector &row_offsets)\n{\n // Initialize P.manager\n if (B.manager == NULL)\n {\n B.manager = new DistributedManager();\n }\n\n index_type num_parts = A.manager->get_num_partitions();\n index_type my_rank = A.manager->global_id();\n B.manager->A = &B;\n B.manager->setComms(A.manager->getComms());\n B.manager->set_global_id(my_rank);\n B.manager->set_num_partitions(num_parts);\n B.manager->neighbors = A.manager->neighbors;\n int num_neighbors = A.manager->neighbors.size();\n B.manager->getComms()->set_neighbors(num_neighbors);\n // --------------------------------------\n // Create B.manager->B2L_maps\n // --------------------------------------\n B.manager->B2L_maps.resize(num_neighbors);\n int max_B2L_maps_size = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int A_b2l_map_size = A.manager->B2L_rings[i][1];\n\n if (A_b2l_map_size > max_B2L_maps_size)\n {\n max_B2L_maps_size = A_b2l_map_size;\n }\n }\n\n IVector row_sizes(row_offsets.size() - 1);\n thrust_wrapper::adjacent_difference(row_offsets.begin() + 1, row_offsets.end(), row_sizes.begin());\n cudaCheckError();\n IVector B2L_maps_offsets(max_B2L_maps_size + 1);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int A_b2l_map_size = A.manager->B2L_rings[i][1];\n\n if (A_b2l_map_size > 0)\n {\n // For each node in the B2L_map for that neighbor, write the row size in B2L_maps_offsets\n thrust_wrapper::gather(A.manager->B2L_maps[i].begin(), A.manager->B2L_maps[i].begin() + A_b2l_map_size, row_sizes.begin(), B2L_maps_offsets.begin());\n // Do exclusive scan, to go from row length to row offsets\n thrust_wrapper::exclusive_scan(B2L_maps_offsets.begin(), B2L_maps_offsets.begin() + A_b2l_map_size + 1, B2L_maps_offsets.begin());\n B.manager->B2L_maps[i].resize(B2L_maps_offsets[A_b2l_map_size]);\n const int cta_size = 128;\n const int nWarps = cta_size / 32;\n int grid_size = std::min( 4096, (int) (A_b2l_map_size + nWarps - 1) / nWarps);\n // Populate the B2L_map array\n fill_B_B2L_map <<< grid_size, cta_size >>>(\n B2L_maps_offsets.raw(),\n A.manager->B2L_maps[i].raw(),\n row_offsets.raw(),\n B.manager->B2L_maps[i].raw(),\n A_b2l_map_size,\n A.manager->global_id());\n cudaCheckError();\n }\n else\n {\n B.manager->B2L_maps[i].resize(0);\n }\n\n cudaCheckError();\n }\n\n B.manager->B2L_rings.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n B.manager->B2L_rings[i].resize(2);\n B.manager->B2L_rings[i][0] = 0;\n B.manager->B2L_rings[i][1] = B.manager->B2L_maps[i].size();\n }\n\n // Initialize the B2L_maps_offsets (required by gather_v2 and scatter_v2) from the B2L_rings values\n this->initialize_B2L_maps_offsets(B, 1);\n B.manager->halo_offsets.resize(num_neighbors + 1);\n\n // Fill halo_offsets\n for (int i = 0; i < num_neighbors + 1; i++)\n {\n B.manager->halo_offsets[i] = row_offsets[A.manager->halo_offsets[i]];\n }\n}\n\n\ntemplate \nvoid DistributedArranger >::exchange_halo_rows_P(Matrix_d &A, Matrix_d &P, I64Vector &RAP_local_to_global_map, IVector_h &P_neighbors, I64Vector_h &P_halo_ranges_h, I64Vector &P_halo_ranges, IVector_h &P_halo_offsets, I64Vector_h &RAP_part_offsets_h, I64Vector &RAP_part_offsets, index_type num_owned_coarse_pts, int64_t coarse_base_index)\n{\n if (P.hasProps(DIAG) || P.get_block_size() != 1)\n {\n FatalError(\"P with external diagonal or block_size != 1 not supported\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n int num_neighbors = A.manager->num_neighbors() ;\n std::vector halo_rows_P_row_offsets;\n std::vector halo_rows_P_col_indices;\n std::vector halo_rows_P_local_col_indices;\n std::vector halo_rows_P_values;\n // In this function, store in halo_rows_P_row_offsets, halo_rows_P_col_indices and halo_rows_P_values, the rows of P that need to be sent to each neighbors\n // halo_rows_P_col_indices stores global indices\n this->pack_halo_rows_P(A, P, halo_rows_P_row_offsets, halo_rows_P_col_indices, halo_rows_P_values, RAP_local_to_global_map, num_owned_coarse_pts, coarse_base_index);\n // Do the exchange with the neighbors\n // On return, halo_rows_P_rows_offsets, halo_rows_P_col_indices and halo_rows_P_values stores the rows of P received from each neighbor (rows needed to perform A*P)\n std::vector dummy_halo_ids(0);\n A.manager->getComms()->exchange_matrix_halo(halo_rows_P_row_offsets, halo_rows_P_col_indices, halo_rows_P_values, dummy_halo_ids, A.manager->neighbors, A.manager->global_id());\n halo_rows_P_local_col_indices.resize(halo_rows_P_col_indices.size());\n\n for (int i = 0; i < halo_rows_P_col_indices.size(); i++)\n {\n halo_rows_P_local_col_indices[i].resize(halo_rows_P_col_indices[i].size());\n }\n\n // Update the list of neighbors \"P_neighbors\" and the corresponding ranges, offsets\n update_neighbors_list(A, P_neighbors, P_halo_ranges_h, P_halo_ranges, RAP_part_offsets_h, RAP_part_offsets, halo_rows_P_row_offsets, halo_rows_P_col_indices);\n std::vector dummy_boundary_list(0);\n int current_num_rings = 1;\n halo_rows_P_local_col_indices.resize(halo_rows_P_col_indices.size());\n // Convert the global indices in the rows just received to local indices\n // This function updates halo offsets, halo_rows_P_local_col_indices, local_to_global_map\n // This should not modify the existing P manager\n this->compute_local_halo_indices( P.row_offsets, P.col_indices, halo_rows_P_row_offsets, halo_rows_P_col_indices, halo_rows_P_local_col_indices, RAP_local_to_global_map, dummy_boundary_list, P_neighbors, P_halo_ranges_h, P_halo_ranges, P_halo_offsets, coarse_base_index, num_owned_coarse_pts, P.get_num_rows(), current_num_rings);\n // Append the new rows to the matrix P\n P.set_initialized(0);\n this->append_halo_rows(P, halo_rows_P_row_offsets, halo_rows_P_local_col_indices, halo_rows_P_values);\n P.set_initialized(1);\n}\n\n\ntemplate \nvoid DistributedArranger >::append_halo_rows(Matrix_d &A, std::vector &halo_row_offsets, std::vector &halo_local_indices, std::vector &halo_values )\n{\n int new_num_rows = A.get_num_rows();\n int new_num_nnz = A.row_offsets[A.get_num_rows()];\n int cur_row = A.get_num_rows();\n int cur_offset = new_num_nnz;\n int num_neighbors = halo_row_offsets.size();\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = halo_row_offsets[i].size();\n\n if (size != 0)\n {\n new_num_rows += halo_row_offsets[i].size() - 1;\n new_num_nnz += halo_local_indices[i].size();\n }\n }\n\n A.resize(new_num_rows, new_num_rows, new_num_nnz, A.get_block_dimx(), A.get_block_dimy(), 1);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int num_halo_rows = halo_row_offsets[i].size() - 1;\n\n if (num_halo_rows > 0)\n {\n // update halo row offsets in-place\n thrust_wrapper::transform(halo_row_offsets[i].begin(), halo_row_offsets[i].end(), amgx::thrust::constant_iterator(cur_offset), halo_row_offsets[i].begin(), amgx::thrust::plus());\n // insert halo rows\n amgx::thrust::copy(halo_row_offsets[i].begin(), halo_row_offsets[i].end() - 1, A.row_offsets.begin() + cur_row);\n amgx::thrust::copy(halo_local_indices[i].begin(), halo_local_indices[i].end(), A.col_indices.begin() + cur_offset);\n amgx::thrust::copy(halo_values[i].begin(), halo_values[i].end(), A.values.begin() + cur_offset);\n // update counters\n cur_offset = halo_row_offsets[i][num_halo_rows];\n cur_row += num_halo_rows;\n }\n }\n\n cudaCheckError();\n A.row_offsets[A.get_num_rows()] = cur_offset;\n int num_cols = -1;\n num_cols = thrust_wrapper::reduce(A.col_indices.begin(), A.col_indices.end(), num_cols, amgx::thrust::maximum()) + 1;\n cudaCheckError();\n A.set_num_cols(num_cols);\n}\n\n\ntemplate \nvoid DistributedArranger >::exchange_RAP_ext(Matrix_d &RAP, Matrix_d &RAP_full, Matrix_d &A, Matrix_d &P, IVector_h &P_halo_offsets, I64Vector &RAP_local_to_global_map, IVector_h &P_neighbors, I64Vector_h &P_halo_ranges_h, I64Vector &P_halo_ranges, I64Vector_h &RAP_part_offsets_h, I64Vector &RAP_part_offsets, index_type num_owned_coarse_pts, int64_t coarse_base_index, void *wk)\n{\n if (RAP_full.hasProps(DIAG) || RAP_full.get_block_size() != 1)\n {\n FatalError(\"RAP_full with external diagonal or block_size != 1 not supported\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // --------------------------------------------------------------------------------------\n // Pack the rows of RAP_full that need to be sent to neighbors, using the manager of P\n // --------------------------------------------------------------------------------------\n std::vector halo_rows_RAP_row_offsets;\n std::vector halo_rows_RAP_row_ids;\n std::vector halo_rows_RAP_col_indices;\n std::vector halo_rows_RAP_values;\n this->pack_halo_rows_RAP(RAP_full, P, halo_rows_RAP_row_offsets, halo_rows_RAP_col_indices, halo_rows_RAP_values, halo_rows_RAP_row_ids, RAP_local_to_global_map);\n // Send the rows to the neighbors, receive RAP rows from neighbors\n // Col indices received from neighbors are global indices\n // We also receive from neighbors the global row ids\n P.manager->getComms()->exchange_matrix_halo(halo_rows_RAP_row_offsets, halo_rows_RAP_col_indices, halo_rows_RAP_values, halo_rows_RAP_row_ids, P.manager->neighbors, P.manager->global_id());\n std::vector halo_rows_RAP_local_col_indices;\n halo_rows_RAP_local_col_indices.resize(halo_rows_RAP_col_indices.size());\n\n for (int i = 0; i < halo_rows_RAP_col_indices.size(); i++)\n {\n halo_rows_RAP_local_col_indices[i].resize(halo_rows_RAP_col_indices[i].size());\n }\n\n // This function finds the new number of neighbors, renumber halo_cols with local_ids, modifies halo_offsets\n update_neighbors_list(A, P_neighbors, P_halo_ranges_h, P_halo_ranges, RAP_part_offsets_h, RAP_part_offsets, halo_rows_RAP_row_offsets, halo_rows_RAP_col_indices);\n // Now renumber find local indices for global col indices, update local_to_global\n std::vector dummy_boundary_list(0);\n int current_num_rings = 2;\n this->compute_local_halo_indices(RAP_full.row_offsets, RAP_full.col_indices, halo_rows_RAP_row_offsets, halo_rows_RAP_col_indices, halo_rows_RAP_local_col_indices, RAP_local_to_global_map, dummy_boundary_list, P_neighbors, P_halo_ranges_h, P_halo_ranges, P_halo_offsets, coarse_base_index, num_owned_coarse_pts, num_owned_coarse_pts, current_num_rings );\n // Now find local row ids from global ids of rows received\n std::vector halo_rows_RAP_local_row_ids;\n halo_rows_RAP_local_row_ids.resize(halo_rows_RAP_row_ids.size());\n\n for (int i = 0; i < halo_rows_RAP_row_ids.size(); i++)\n {\n int size = halo_rows_RAP_row_ids[i].size();\n halo_rows_RAP_local_row_ids[i].resize(size);\n\n if (size != 0)\n {\n int cta_size = 128;\n int grid_size = std::min( 4096, (size + cta_size - 1) / cta_size);\n fill_halo_rows_row_ids <<< grid_size, cta_size >>>( halo_rows_RAP_row_ids[i].raw(), halo_rows_RAP_local_row_ids[i].raw(), coarse_base_index, num_owned_coarse_pts, size);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n // Add the rows received from neighbors to owned rows in RAP_full and create RAP matrix\n RAP.set_initialized( 0 );\n RAP.set_num_rows(num_owned_coarse_pts);\n RAP.set_num_cols(num_owned_coarse_pts + RAP.manager->local_to_global_map.size());\n RAP.row_offsets.resize(num_owned_coarse_pts + 1);\n RAP_full.setView(OWNED);\n CSR_Multiply::csr_RAP_sparse_add( RAP, RAP_full, halo_rows_RAP_row_offsets, halo_rows_RAP_local_col_indices, halo_rows_RAP_values, halo_rows_RAP_local_row_ids, wk );\n RAP.set_initialized( 1 );\n}\n\n\ntemplate \nvoid DistributedArranger >::pack_halo_rows_P(Matrix_d &A, Matrix_d &P, std::vector &halo_rows_P_row_offsets, std::vector &halo_rows_P_col_indices, std::vector &halo_rows_P_values, I64Vector &RAP_local_to_global_map, index_type num_owned_coarse_pts, int64_t coarse_base_index)\n{\n int num_neighbors = A.manager->num_neighbors();\n halo_rows_P_row_offsets.resize(num_neighbors);\n halo_rows_P_col_indices.resize(num_neighbors);\n halo_rows_P_values.resize(num_neighbors);\n int num_rings_to_send = 1;\n //Scratch space computation\n int max_size = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n max_size = max_size > A.manager->B2L_rings[i][num_rings_to_send] ? max_size : A.manager->B2L_rings[i][num_rings_to_send];\n }\n\n IVector matrix_halo_sizes(max_size + 1);\n\n // Here only using the 1-ring of matrix A\n for (int i = 0; i < num_neighbors; i++)\n {\n // Write the length of the rows in the order of B2L_maps, then calculate row_offsets\n int size = A.manager->B2L_rings[i][num_rings_to_send];\n\n if (size != 0)\n {\n //matrix_halo_sizes.resize(size+1);\n int num_blocks = std::min(4096, (size + 127) / 128);\n write_matrix_rowsize <<< num_blocks, 128>>>(A.manager->B2L_maps[i].raw(), P.row_offsets.raw(), size, matrix_halo_sizes.raw());\n cudaCheckError();\n thrust_wrapper::exclusive_scan(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, matrix_halo_sizes.begin());\n int nnz_count = matrix_halo_sizes[size];\n // Resize export halo matrix, and copy over the rows\n halo_rows_P_row_offsets[i].resize(size + 1);\n halo_rows_P_col_indices[i].resize(nnz_count);\n halo_rows_P_values[i].resize(nnz_count);\n /* WARNING: Since A is reordered (into interior and boundary nodes), while R and P are not reordered,\n you must unreorder A when performing R*A*P product in ordre to obtain the correct result. */\n export_matrix_elements_global<32> <<< num_blocks, 128>>>(P.row_offsets.raw(), P.values.raw(), P.get_block_size(), A.manager->B2L_maps[i].raw(), matrix_halo_sizes.raw(), halo_rows_P_values[i].raw(), P.col_indices.raw(), halo_rows_P_col_indices[i].raw(), size, RAP_local_to_global_map.raw(), NULL /*A.manager->inverse_renumbering.raw()*/, num_owned_coarse_pts, coarse_base_index);\n cudaCheckError();\n amgx::thrust::copy(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, halo_rows_P_row_offsets[i].begin());\n }\n else\n {\n halo_rows_P_row_offsets[i].resize(0);\n halo_rows_P_col_indices[i].resize(0);\n halo_rows_P_values[i].resize(0);\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::pack_halo_rows_RAP(Matrix_d &RAP, Matrix_d &P, std::vector &halo_rows_RAP_row_offsets, std::vector &halo_rows_RAP_col_indices, std::vector &halo_rows_RAP_values, std::vector &halo_rows_RAP_row_ids, I64Vector &RAP_local_to_global_map)\n{\n int num_neighbors = P.manager->neighbors.size();\n int num_owned_coarse_pts = P.manager->halo_offsets[0];\n int64_t P_col_base_index = P.manager->base_index();\n halo_rows_RAP_row_offsets.resize(num_neighbors);\n halo_rows_RAP_row_ids.resize(num_neighbors);\n halo_rows_RAP_col_indices.resize(num_neighbors);\n halo_rows_RAP_values.resize(num_neighbors);\n // Create tmp_B2L_maps that contains the list of rows to send\n std::vector tmp_B2L_maps(num_neighbors);\n\n // Fill B2L_maps with list of halo indices for all rings to send\n for (int i = 0; i < num_neighbors; i++)\n {\n tmp_B2L_maps[i].resize(0);\n halo_rows_RAP_row_ids[i].resize(0);\n int begin = P.manager->halo_offsets[i];\n int end = P.manager->halo_offsets[i + 1];\n\n if (begin > RAP.row_offsets.size() - 1)\n {\n break;\n }\n\n int local_to_global_begin = begin - P.manager->halo_offsets[0];\n int size = end - begin;\n\n if (end > RAP.row_offsets.size() - 1)\n {\n size = RAP.row_offsets.size() - 1 - begin;\n }\n\n if (size > 0)\n {\n tmp_B2L_maps[i].resize(size);\n halo_rows_RAP_row_ids[i].resize(size);\n thrust_wrapper::sequence(tmp_B2L_maps[i].begin(), tmp_B2L_maps[i].begin() + size, begin);\n amgx::thrust::copy(P.manager->local_to_global_map.begin() + (int64_t) local_to_global_begin, P.manager->local_to_global_map.begin() + (int64_t) local_to_global_begin + (int64_t) size, halo_rows_RAP_row_ids[i].begin());\n }\n }\n\n cudaCheckError();\n //Scratch space computation\n int max_size = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = tmp_B2L_maps[i].size();\n max_size = max_size > size ? max_size : size;\n }\n\n IVector matrix_halo_sizes(max_size + 1);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n // Write the length of the rows in the order of B2L_maps, then calculate row_offsets\n int size = tmp_B2L_maps[i].size();\n\n if (size > 0)\n {\n //matrix_halo_sizes.resize(size+1);\n int num_blocks = std::min(4096, (size + 127) / 128);\n write_matrix_rowsize <<< num_blocks, 128>>>(tmp_B2L_maps[i].raw(), RAP.row_offsets.raw(), size, matrix_halo_sizes.raw());\n cudaCheckError();\n thrust_wrapper::exclusive_scan(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, matrix_halo_sizes.begin());\n int nnz_count = matrix_halo_sizes[size];\n // Resize export halo matrix, and copy over the rows\n halo_rows_RAP_row_offsets[i].resize(tmp_B2L_maps[i].size() + 1);\n halo_rows_RAP_col_indices[i].resize(nnz_count);\n halo_rows_RAP_values[i].resize(nnz_count);\n export_matrix_elements_global<32> <<< num_blocks, 128>>>(RAP.row_offsets.raw(), RAP.values.raw(), RAP.get_block_size(), tmp_B2L_maps[i].raw(), matrix_halo_sizes.raw(), halo_rows_RAP_values[i].raw(), RAP.col_indices.raw(), halo_rows_RAP_col_indices[i].raw(), size, RAP_local_to_global_map.raw(), NULL, num_owned_coarse_pts, P_col_base_index);\n cudaCheckError();\n amgx::thrust::copy(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, halo_rows_RAP_row_offsets[i].begin());\n }\n else\n {\n halo_rows_RAP_row_offsets[i].resize(0);\n halo_rows_RAP_col_indices[i].resize(0);\n halo_rows_RAP_values[i].resize(0);\n }\n }\n\n cudaCheckError();\n}\n\n\ntemplate \nvoid DistributedArranger >::sparse_matrix_add(Matrix_d &RAP, std::vector &halo_rows_RAP_row_offsets, std::vector &halo_rows_RAP_col_indices, std::vector &halo_rows_RAP_values, std::vector halo_rows_RAP_row_ids, I64Vector_d &P_local_to_global_map, index_type P_col_num_owned_pts, int64_t P_col_base_index)\n{\n // Look at csr_galerkin_product implementation, use hash\n}\n\n\n// CUDA 4.2\ntemplate \nvoid DistributedArranger >::create_halo_rows(Matrix_d &A,\n std::vector &halo_lists\n )\n{\n int neighbors = A.manager->num_neighbors();\n\n if (neighbors == 0)\n {\n return;\n }\n\n int rings = (A.manager->B2L_rings.size() - 1) / neighbors;\n A.manager->halo_rows = new std::vector >(neighbors);\n std::vector > &halo_rows = *A.manager->halo_rows;\n\n if (this->halo_coloring == LAST) //Create dummy matrices, only containing diagonal elements\n {\n for (int i = 0; i < neighbors; i++)\n {\n halo_rows[i].addProps(CSR);\n\n if (A.hasProps(DIAG)) { halo_rows[i].addProps(DIAG); }\n\n halo_rows[i].resize(halo_lists[i]->size(), halo_lists[i]->size(), (A.hasProps(DIAG) ? 0 : halo_lists[i]->size()), A.get_block_dimy(), A.get_block_dimx(), 1);\n\n if (A.hasProps(DIAG))\n {\n thrust_wrapper::fill(halo_rows[i].row_offsets.begin(), halo_rows[i].row_offsets.end(), 0);\n }\n else\n {\n thrust_wrapper::sequence(halo_rows[i].row_offsets.begin(), halo_rows[i].row_offsets.end());\n amgx::thrust::copy(halo_lists[i]->begin(), halo_lists[i]->end(), halo_rows[i].col_indices.begin());\n }\n\n cudaCheckError();\n }\n\n return;\n }\n\n //Scratch space computation\n int max_size = 0;\n\n for (int i = 0; i < neighbors; i++)\n {\n max_size = max_size > A.manager->B2L_maps[i].size() ? max_size : A.manager->B2L_maps[i].size();\n //B2L_rings[i][1]=B2L_maps[i].size(); //I don't know what this was doing here, definitely stupid for rings > 1\n }\n\n IVector matrix_halo_sizes((max_size + 1)*rings);\n\n for (int i = 0; i < neighbors; i++)\n {\n //\n // Write the length of the rows in the order of B2L_maps, then calculate row_offsets\n //\n int size = A.manager->B2L_rings[i][rings];\n matrix_halo_sizes.resize(size + 1);\n int num_blocks = std::min(4096, (size + 127) / 128);\n write_matrix_rowsize <<< num_blocks, 128>>>(A.manager->B2L_maps[i].raw(), A.row_offsets.raw(), size, matrix_halo_sizes.raw());\n cudaCheckError();\n thrust_wrapper::exclusive_scan(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, matrix_halo_sizes.begin());\n cudaCheckError();\n int nnz_count = matrix_halo_sizes[size];\n //\n // Resize export halo matrix, and copy over the rows\n //\n halo_rows[i].addProps(CSR);\n\n if (A.hasProps(DIAG)) { halo_rows[i].addProps(DIAG); }\n\n halo_rows[i].resize(A.manager->B2L_maps[i].size(), A.manager->B2L_maps[i].size(), nnz_count, A.get_block_dimy(), A.get_block_dimx(), 1);\n export_matrix_elements<32> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.values.raw(), A.get_block_size(), A.manager->B2L_maps[i].raw(), matrix_halo_sizes.raw(), halo_rows[i].values.raw(), A.col_indices.raw(), halo_rows[i].col_indices.raw(), size);\n cudaCheckError();\n amgx::thrust::copy(matrix_halo_sizes.begin(), matrix_halo_sizes.begin() + size + 1, halo_rows[i].row_offsets.begin());\n cudaCheckError();\n\n if (A.hasProps(DIAG))\n {\n export_matrix_diagonal <<< num_blocks, 128>>>(A.values.raw() + A.row_offsets[A.get_num_rows()]*A.get_block_size(), A.get_block_size(), A.manager->B2L_maps[i].raw(), halo_rows[i].values.raw() + halo_rows[i].row_offsets[halo_rows[i].get_num_rows()]*A.get_block_size(), size);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::create_halo_rows_global_indices(Matrix_d &A,\n std::vector &halo_rows_row_offsets,\n std::vector &halo_rows_col_indices,\n std::vector &halo_rows_values)\n{\n // in this function we assume that we only have 1-ring so far\n int neighbors = A.manager->num_neighbors();\n\n for (int i = 0; i < neighbors; i++)\n {\n // compute row offsets and nnz count\n int size = A.manager->B2L_maps[i].size();\n halo_rows_row_offsets[i].resize(size + 1);\n\n if (size > 0)\n {\n int num_blocks = std::min(4096, (size + 127) / 128);\n write_matrix_rowsize <<< num_blocks, 128>>>(A.manager->B2L_maps[i].raw(), A.row_offsets.raw(), size, halo_rows_row_offsets[i].raw());\n cudaCheckError();\n thrust_wrapper::exclusive_scan(halo_rows_row_offsets[i].begin(), halo_rows_row_offsets[i].begin() + size + 1, halo_rows_row_offsets[i].begin());\n // compute global indices\n int nnz_count = halo_rows_row_offsets[i][size];\n halo_rows_col_indices[i].resize(nnz_count);\n halo_rows_values[i].resize(nnz_count);\n export_matrix_elements_global<32> <<< num_blocks, 128>>>(A.row_offsets.raw(), A.values.raw(), A.get_block_size(), A.manager->B2L_maps[i].raw(), halo_rows_row_offsets[i].raw(), halo_rows_values[i].raw(), A.col_indices.raw(), halo_rows_col_indices[i].raw(), size, A.manager->local_to_global_map.raw(), NULL, A.get_num_rows(), A.manager->base_index());\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::create_halo_btl(Matrix_d &A, INDEX_TYPE rings, INDEX_TYPE neighbors,\n Vector &neighbors_list,\n int64_t base_index, INDEX_TYPE index_range,\n INDEX_TYPE global_id, std::vector &B2L_maps,\n std::vector > &B2L_rings,\n std::vector **halo_btl_p,\n\n std::vector &halo_lists\n )\n{\n *halo_btl_p = new std::vector >(neighbors);\n std::vector > &halo_btl = **halo_btl_p;\n /*\n //calling constructors on the allocated memory\n for(int i=0; i();\n */\n\n if (this->halo_coloring != LAST)\n {\n //If we are sending matrices over\n for (int i = 0; i < neighbors; i++)\n {\n //\n // Transfer indexing data too (global indices of exported rows, and sizes of the halo rings)\n //\n /*\n EXAMPLE\n For example partition 1, with 2 rings, and the following B2L and L2H (called halo_lists here) info\n num_neighbors=2; neighbors = [0 2]\n B2L_rings[[0 2 4][0 2 4]] B2L_maps[[0 1| 2 3][1 3| 0 2]]\n L2H_maps (halo_lists) [[4 5][6 7]]\n */\n int size = B2L_rings[i][rings];\n halo_btl[i].set_global_id(global_id);\n halo_btl[i].set_base_index(0);\n halo_btl[i].resize(1, rings);\n halo_btl[i].B2L_maps[0].resize(size);\n halo_btl[i].L2H_maps.resize(1);\n halo_btl[i].L2H_maps[0].resize(halo_lists[i]->size());\n //send over B2L and L2H maps corresponding to current neighbor, so that it will be able to make sense of indices (local to this partition)\n amgx::thrust::copy(B2L_maps[i].begin(), B2L_maps[i].end(), halo_btl[i].B2L_maps[0].begin());\n amgx::thrust::copy(halo_lists[i]->begin(), halo_lists[i]->end(), halo_btl[i].L2H_maps[0].begin());\n int max_index = thrust_wrapper::reduce(halo_lists[i]->begin(), halo_lists[i]->end(), (int)0, amgx::thrust::maximum()) + 1;\n max_index = max_index > A.get_num_rows() ? max_index : A.get_num_rows();\n halo_btl[i].set_index_range(max_index);\n amgx::thrust::copy(B2L_rings[i].begin(), B2L_rings[i].end(), halo_btl[i].B2L_rings[0].begin());\n /*\n EXAMPLE\n These fields are going to be:\n halo_btl[0] (for neighbor ID 0)\n global_id = 1; base_index=0; index_range=6 B2L_rings[0] = [0 2 4] B2L_maps[0] = [0 1| 2 3] L2H_maps = [4 5]\n halo_btl[1] (for neighbor ID 2)\n global_id = 1; base_index=0; index_range=8 B2L_rings[0] = [0 2 4] B2L_maps[0] = [1 3| 0 2] L2H_maps = [6 7]\n */\n }\n }\n else\n {\n if (rings > 1) { FatalError(\"halo_coloring=LAST dummy halo rows creation not implemented\", AMGX_ERR_NOT_IMPLEMENTED); }\n\n //Create dummy halo matrices\n for (int i = 0; i < neighbors; i++)\n {\n int size = halo_lists[i]->size();\n halo_btl[i].set_global_id(neighbors_list[i]);\n halo_btl[i].set_base_index(0);\n halo_btl[i].resize(1, rings);\n halo_btl[i].B2L_maps[0].resize(size);\n halo_btl[i].L2H_maps.resize(1);\n halo_btl[i].L2H_maps[0].resize(B2L_rings[i][rings]);\n //Even though in this case matrices are not being sent over, we create the same data structures, so they can be handled the same way\n amgx::thrust::copy(B2L_maps[i].begin(), B2L_maps[i].end(), halo_btl[i].L2H_maps[0].begin());\n amgx::thrust::copy(halo_lists[i]->begin(), halo_lists[i]->end(), halo_btl[i].B2L_maps[0].begin());\n int max_index = thrust_wrapper::reduce(halo_lists[i]->begin(), halo_lists[i]->end(), (int)0, amgx::thrust::maximum()) + 1;\n max_index = max_index > A.get_num_rows() ? max_index : A.get_num_rows();\n halo_btl[i].set_index_range(max_index);\n halo_btl[i].B2L_rings[0][0] = 0;\n halo_btl[i].B2L_rings[0][1] = size;\n }\n }\n\n cudaCheckError();\n}\n\n// CUDA 4.2\ntemplate \nvoid DistributedArranger >::create_halo_btl(Matrix_d &A,\n std::vector &halo_lists)\n{\n int neighbors = A.manager->num_neighbors();\n\n if (neighbors == 0)\n {\n return;\n }\n\n int rings = (A.manager->B2L_rings.size() - 1) / neighbors;\n A.manager->halo_btl = new std::vector >(neighbors);\n std::vector > &halo_btl = *A.manager->halo_btl;\n\n if (this->halo_coloring != LAST)\n {\n for (int i = 0; i < neighbors; i++)\n {\n //\n // Transfer indexing data too (global indices of exported rows, and sizes of the halo rings)\n //\n int size = A.manager->B2L_rings[i][rings];\n halo_btl[i].set_global_id(A.manager->global_id());\n halo_btl[i].set_base_index(0);\n //Adding one more ring to store halo numbering\n halo_btl[i].resize(1, rings);\n halo_btl[i].B2L_maps[0].resize(size);\n halo_btl[i].L2H_maps.resize(1);\n halo_btl[i].L2H_maps[0].resize(halo_lists[i]->size());\n amgx::thrust::copy(A.manager->B2L_maps[i].begin(), A.manager->B2L_maps[i].end(), halo_btl[i].B2L_maps[0].begin());\n amgx::thrust::copy(halo_lists[i]->begin(), halo_lists[i]->end(), halo_btl[i].L2H_maps[0].begin());\n int max_index = thrust_wrapper::reduce(halo_lists[i]->begin(), halo_lists[i]->end(), (int)0, amgx::thrust::maximum()) + 1;\n max_index = max_index > A.get_num_rows() ? max_index : A.get_num_rows();\n halo_btl[i].set_index_range(max_index);\n amgx::thrust::copy(A.manager->B2L_rings[i].begin(), A.manager->B2L_rings[i].end(), halo_btl[i].B2L_rings[0].begin());\n }\n }\n else\n {\n if (rings > 1) { FatalError(\"halo_coloring=LAST dummy halo rows creation not implemented\", AMGX_ERR_NOT_IMPLEMENTED); }\n\n //Create dummy halo matrices\n for (int i = 0; i < neighbors; i++)\n {\n int size = halo_lists[i]->size();\n halo_btl[i].set_global_id(A.manager->neighbors[i]);\n halo_btl[i].set_base_index(0);\n //Adding one more ring to store halo numbering\n halo_btl[i].resize(1, rings);\n halo_btl[i].B2L_maps[0].resize(size);\n halo_btl[i].L2H_maps.resize(1);\n halo_btl[i].L2H_maps[0].resize(A.manager->B2L_rings[i][rings]);\n amgx::thrust::copy(A.manager->B2L_maps[i].begin(), A.manager->B2L_maps[i].end(), halo_btl[i].L2H_maps[0].begin());\n amgx::thrust::copy(halo_lists[i]->begin(), halo_lists[i]->end(), halo_btl[i].B2L_maps[0].begin());\n int max_index = thrust_wrapper::reduce(halo_lists[i]->begin(), halo_lists[i]->end(), (int)0, amgx::thrust::maximum()) + 1;\n max_index = max_index > A.get_num_rows() ? max_index : A.get_num_rows();\n halo_btl[i].set_index_range(max_index);\n halo_btl[i].B2L_rings[0][0] = 0;\n halo_btl[i].B2L_rings[0][1] = size;\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::create_halo_btl_multiple_rings(Matrix_d &A, int rings)\n{\n // creates halo btl managers\n // required input from A.manager: num_neighbors, B2L_rings, B2L_maps, L2H_maps\n int neighbors = A.manager->num_neighbors();\n A.manager->halo_btl = new std::vector >(neighbors);\n std::vector > &halo_btl = *A.manager->halo_btl;\n\n for (int i = 0; i < neighbors; i++)\n {\n // Transfer indexing data too (global indices of exported rows, and sizes of the halo rings)\n int size = A.manager->B2L_rings[i][rings];\n halo_btl[i].set_global_id(A.manager->global_id());\n halo_btl[i].set_base_index(0);\n // Adding one more ring to store halo numbering\n halo_btl[i].resize(1, rings);\n halo_btl[i].B2L_maps[0].resize(size);\n halo_btl[i].L2H_maps.resize(1);\n halo_btl[i].L2H_maps[0].resize(A.manager->L2H_maps[i].size());\n amgx::thrust::copy(A.manager->B2L_maps[i].begin(), A.manager->B2L_maps[i].end(), halo_btl[i].B2L_maps[0].begin());\n amgx::thrust::copy(A.manager->L2H_maps[i].begin(), A.manager->L2H_maps[i].end(), halo_btl[i].L2H_maps[0].begin());\n amgx::thrust::copy(A.manager->B2L_rings[i].begin(), A.manager->B2L_rings[i].end(), halo_btl[i].B2L_rings[0].begin());\n // compute index range\n int max_index = thrust_wrapper::reduce(A.manager->L2H_maps[i].begin(), A.manager->L2H_maps[i].end(), (int)0, amgx::thrust::maximum()) + 1;\n max_index = max_index > A.get_num_rows() ? max_index : A.get_num_rows();\n halo_btl[i].set_index_range(max_index);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedArranger >::create_halo_matrices(Matrix_d &A, INDEX_TYPE rings,\n int64_t base_index, INDEX_TYPE index_range,\n INDEX_TYPE global_id, Vector &neighbors_,\n std::vector &B2L_maps,\n std::vector > &B2L_rings,\n std::vector **halo_rows_p,\n std::vector **halo_btl_p, std::vector &boundary_lists, std::vector &halo_lists)\n{\n int neighbors = neighbors_.size();\n this->create_maps(A, rings, neighbors, B2L_maps, B2L_rings, boundary_lists);\n this->create_halo_btl(A, rings, neighbors, neighbors_, base_index, index_range, global_id, B2L_maps, B2L_rings, halo_btl_p, halo_lists);\n this->create_halo_rows(A, rings, neighbors, B2L_maps, B2L_rings, halo_rows_p, halo_lists);\n}\n\ntemplate \nvoid DistributedArranger >::create_B2L(Matrix_d &A, INDEX_TYPE my_id, INDEX_TYPE rings)\n{\n this->halo_coloring = A.manager->getComms()->halo_coloring;\n I64Vector_h halo_ranges_h;\n this->create_neighbors(A, my_id, halo_ranges_h);\n A.manager->set_global_id(my_id);\n int num_neighbors = A.manager->neighbors.size();\n A.manager->resize(num_neighbors, rings);\n A.manager->getComms()->set_neighbors(num_neighbors);\n std::vector boundary_lists;\n IVector halo_nodes;\n this->create_boundary_lists(A, halo_ranges_h, boundary_lists, halo_nodes);\n A.manager->L2H_maps.resize(num_neighbors);\n std::vector halo_lists;\n halo_lists.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n A.manager->L2H_maps[i].resize(boundary_lists[i].size());\n thrust_wrapper::transform(boundary_lists[i].begin(), boundary_lists[i].end(), amgx::thrust::constant_iterator (halo_ranges_h[2 * i]), A.manager->L2H_maps[i].begin(), amgx::thrust::plus ());\n halo_lists[i] = &A.manager->L2H_maps[i];\n }\n\n cudaCheckError();\n A.manager->getComms()->exchange_vectors(boundary_lists, A, 0);\n this->renumber_to_local(A, boundary_lists, halo_lists, my_id, halo_ranges_h, halo_nodes);\n this->create_halo_matrices(A, 1, boundary_lists, halo_lists);\n\n if (this->halo_coloring != LAST)\n {\n A.manager->getComms()->exchange_matrix_halo(*A.manager->halo_rows, *A.manager->halo_btl, A);\n }\n\n A.manager->A = &A;\n}\n\ntemplate \nvoid DistributedArranger >::create_halo_matrices(Matrix_d &A, INDEX_TYPE rings, std::vector &boundary_lists, std::vector &halo_lists)\n{\n this->create_maps(A, rings, boundary_lists);\n this->create_halo_btl(A, halo_lists);\n this->create_halo_rows(A, halo_lists);\n}\n\ntemplate \nvoid DistributedArranger >::create_B2L(Matrix_h &A, INDEX_TYPE my_id, INDEX_TYPE rings)\n{\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class DistributedArrangerBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class DistributedArranger::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n\n\n\n\n", "source": "/kaggle/working/amgx/src/distributed/distributed_arranger.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"cuda_runtime.h\"\n#include \"reorder_partition.h\"\n\n#include \"amgx_types/util.h\"\n\n#include \n#include //debug only:\n\nstruct is_my_part\n{\n const int _my_part;\n is_my_part(int my_part) : _my_part(my_part) { }\n\n __host__ __device__\n bool operator()(const int part)\n {\n return (part == _my_part);\n }\n};\n\nnamespace amgx\n{\n\nstatic int insertDiagonals = 1;\n\ntemplate \nstatic __device__ __forceinline__\nindex_type internal_index(index_type i, index_type j, index_type k, index_type nx, index_type ny, index_type nz)\n{\n return k * (nx * ny) + j * nx + i;\n}\n\ntemplate \nstatic __device__ __forceinline__\nint64_t get_global_offset(index_type p, index_type q, index_type r, index_type P, index_type Q, index_type R, index_type num_rows)\n{\n int rank_id = r * (P * Q) + q * P + p;\n return ((int64_t) rank_id) * ((int64_t) num_rows);\n}\n\ntemplate \n__global__\nvoid poisson7pt_count_row_len(index_type *row_len, index_type nx, index_type ny, index_type nz, index_type p, index_type q, index_type r, index_type P, index_type Q, index_type R, index_type num_rows)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_rows ; tidx += blockDim.x * gridDim.x)\n {\n /* compute p,q,r from P,Q,R and myid */\n int i = tidx % nx; // Position in x direction\n int j = (( tidx - i) / nx) % ny; // Position in y\n int k = ( tidx - i - nx * j) / ( nx * ny ); // Position in z\n int substract = ((i == 0) && (p == 0));\n substract += ((i == nx - 1) && (p == P - 1));\n substract += ((j == 0) && (q == 0));\n substract += ((j == ny - 1) && (q == Q - 1));\n substract += ((k == 0) && (r == 0));\n substract += ((k == nz - 1) && (r == R - 1));\n // Store 7 in position (num_rows+1), such that row_len[num_rows+1] = 0\n //substract = (tidx == num_rows+1) ? 7 : substract;\n row_len[tidx] = 7 - substract;\n }\n}\n\ntemplate \n__global__\nvoid poisson7pt_set_col_values(const index_type *__restrict__ row_offsets, index_type *__restrict__ col_indices, mat_value_type *__restrict__ values, index_type *__restrict__ diag, int64_t *__restrict__ local_to_global, index_type nx, index_type ny, index_type nz, index_type p, index_type q, index_type r, index_type P, index_type Q, index_type R, index_type num_rows)\n{\n for (int row = threadIdx.x + blockIdx.x * blockDim.x; row < num_rows ; row += blockDim.x * gridDim.x)\n {\n /* compute p,q,r from P,Q,R and myid */\n int i = row % nx; // Position in x direction\n int j = (( row - i) / nx) % ny; // Position in y\n int k = ( row - i - nx * j) / ( nx * ny ); // Position in z\n int halo_offset = num_rows;\n int pos = row_offsets[row];\n // Diagonal element\n diag[row] = pos;\n col_indices[pos] = row;\n values[pos++] = types::util::get_one() * 6.;\n\n // ----------------------------\n // Neighbor at position i-1\n // ----------------------------\n if (i)\n {\n // Has a i-1 neighbor, which is an internal node at position (i-1,j,k)\n col_indices[pos] = internal_index(i - 1, j, k, nx, ny, nz);\n values[pos++] = types::util::invert(types::util::get_one());\n }\n else if (p)\n {\n // Has a i-1 neighbor, which is a halo node\n int halo_index = halo_offset + k * ny + j;\n col_indices[pos] = halo_index;\n values[pos++] = types::util::invert(types::util::get_one());\n int64_t global_offset = get_global_offset(p - 1, q, r, P, Q, R, num_rows);\n local_to_global[halo_index - num_rows] = global_offset + internal_index(nx - 1, j, k, nx, ny, nz);\n }\n\n if (p)\n {\n halo_offset += ny * nz;\n }\n\n // ----------------------------\n // Neighbor at position i+1\n // ----------------------------\n if (i < nx - 1)\n {\n // Has i+1 neighbor, which is an internal node at position (i+1,j,k)\n col_indices[pos] = internal_index(i + 1, j, k, nx, ny, nz);\n values[pos++] = types::util::invert(types::util::get_one());\n }\n else\n {\n if (p < P - 1)\n {\n // Has i+1 neighbor, which is a halo node\n int halo_index = halo_offset + k * ny + j;\n col_indices[pos] = halo_index;\n values[pos++] = types::util::invert(types::util::get_one());\n int64_t global_offset = get_global_offset(p + 1, q, r, P, Q, R, num_rows);\n local_to_global[halo_index - num_rows] = global_offset + internal_index(0, j, k, nx, ny, nz);\n }\n }\n\n if (p < P - 1)\n {\n halo_offset += ny * nz;\n }\n\n // ----------------------------\n // Neighbor at position j-1\n // ----------------------------\n if (j)\n {\n // Has a j-1 neighbor, which is an internal node at position (i,j-1,k)\n col_indices[pos] = internal_index(i, j - 1, k, nx, ny, nz);\n values[pos++] = types::util::invert(types::util::get_one());\n }\n else if (q)\n {\n // Has a j-1 neighbor, which is a halo node\n int halo_index = halo_offset + k * nx + i;\n col_indices[pos] = halo_index;\n values[pos++] = types::util::invert(types::util::get_one());\n int64_t global_offset = get_global_offset(p, q - 1, r, P, Q, R, num_rows);\n local_to_global[halo_index - num_rows] = global_offset + internal_index(i, ny - 1, k, nx, ny, nz);\n }\n\n if (q)\n {\n halo_offset += nx * nz;\n }\n\n // ----------------------------\n // Neighbor at position j+1\n // ----------------------------\n if (j < ny - 1)\n {\n // Has a j+1 neighbor, which is an internal node at position (i,j+1,k)\n col_indices[pos] = internal_index(i, j + 1, k, nx, ny, nz);\n values[pos++] = types::util::invert(types::util::get_one());\n }\n else\n {\n if (q < Q - 1)\n {\n // Has a j+1 neighbor, which is a halo node\n int halo_index = halo_offset + k * nx + i;\n col_indices[pos] = halo_index;\n values[pos++] = types::util::invert(types::util::get_one());\n int64_t global_offset = get_global_offset(p, q + 1, r, P, Q, R, num_rows);\n local_to_global[halo_index - num_rows] = global_offset + internal_index(i, 0, k, nx, ny, nz);\n }\n }\n\n if (q < Q - 1)\n {\n halo_offset += nx * nz;\n }\n\n // ----------------------------\n // Neighbor at position k-1\n // ----------------------------\n if (k)\n {\n // Has a k-1 neighbor, which is an internal node at position (i,j,k-1)\n col_indices[pos] = internal_index(i, j, k - 1, nx, ny, nz);\n values[pos++] = types::util::invert(types::util::get_one());\n }\n else if (r)\n {\n // Has a k-1 neighbor, which is a halo node\n int halo_index = halo_offset + j * nx + i;\n col_indices[pos] = halo_index;\n values[pos++] = types::util::invert(types::util::get_one());\n int64_t global_offset = get_global_offset(p, q, r - 1, P, Q, R, num_rows);\n local_to_global[halo_index - num_rows] = global_offset + internal_index(i, j, nz - 1, nx, ny, nz);\n }\n\n if (r)\n {\n halo_offset += nx * ny;\n }\n\n // ----------------------------\n // Neighbor at position k+1\n // ----------------------------\n if (k < nz - 1)\n {\n // Has a k+1 neighbor, which is an internal node at position (i,j,k+1)\n col_indices[pos] = internal_index(i, j, k + 1, nx, ny, nz);\n values[pos++] = types::util::invert(types::util::get_one());\n }\n else\n {\n if (r < R - 1)\n {\n // Has a k+1 neighbor, which is a halo node\n int halo_index = halo_offset + j * nx + i;\n col_indices[pos] = halo_index;\n values[pos++] = types::util::invert(types::util::get_one());\n int64_t global_offset = get_global_offset(p, q, r + 1, P, Q, R, num_rows);\n local_to_global[halo_index - num_rows] = global_offset + internal_index(i, j, 0, nx, ny, nz);\n }\n }\n\n if (r < R - 1)\n {\n halo_offset += nx * ny;\n }\n }\n}\n\ntemplate \n__global__\nvoid set_halo_cols_values(int *row_offsets, int *col_indices, mat_value_type *values, int n, int total_rows, int bsize)\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (tid < (total_rows - n) )\n {\n int offset = row_offsets[n + tid];\n col_indices[offset] = n + tid;\n#pragma unroll\n\n for (int i = 0; i < bsize; i++)\n {\n values[offset * bsize + i] = types::util::get_one(); // This is arbitrary\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \n__global__\nvoid zero_copy_row_lengths_ids_offsets(int *d_old_row_offsets, int *root_row_offsets, int *d_row_ids, int n, int total_num_halos, mat_value_type *diag)\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (tid < n + total_num_halos)\n {\n int new_row_id = d_row_ids[tid];\n\n if (tid < n)\n {\n int start = d_old_row_offsets[tid];\n int row_length = d_old_row_offsets[tid + 1] - start; // zero-copy\n\n if (diag != NULL) // will insert the diagonal\n {\n row_length++;\n }\n\n root_row_offsets[new_row_id] = row_length;\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate< typename mat_value_type>\n__global__\nvoid ipc_consolidation_upload_matrix(int num_rows, int *row_ids, const int *old_row_offsets, int *new_row_offsets, const int *h_old_col_indices, int *new_col_indices, const mat_value_type *h_old_values, mat_value_type *new_values, const mat_value_type *h_old_diag, int bsize)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < num_rows)\n {\n int new_row = row_ids[row];\n int src_base = old_row_offsets[row];\n int dst_base = new_row_offsets[new_row];\n\n // Insert the diagonal at the beginning of each row\n if (h_old_diag != NULL)\n {\n new_col_indices[dst_base] = new_row;\n#pragma unroll\n\n for (int j = 0; j < bsize; j++)\n {\n new_values[dst_base * bsize + j] = h_old_diag[row * bsize + j];\n }\n\n // Increment dst_base by one\n dst_base++;\n }\n\n int end = old_row_offsets[row + 1] - src_base;\n\n for (int i = 0; i < end; i++)\n {\n int old_col = h_old_col_indices[src_base + i];\n int new_col = row_ids[old_col];\n new_col_indices[dst_base + i] = new_col;\n#pragma unroll\n\n for (int j = 0; j < bsize; j++)\n {\n new_values[ (dst_base + i)*bsize + j ] = h_old_values[ (src_base + i) * bsize + j ];\n }\n }\n\n row += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate< typename mat_value_type>\n__global__\nvoid ipc_consolidation_replace_values(int num_rows, int *row_ids, const int *old_row_offsets, int *new_row_offsets, const mat_value_type *h_old_values, mat_value_type *new_values, const mat_value_type *h_old_diag, int bsize)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < num_rows)\n {\n int new_row = row_ids[row];\n int src_base = old_row_offsets[row];\n int dst_base = new_row_offsets[new_row];\n\n // Insert the diagonal at the beginning of each row\n if (h_old_diag != NULL)\n {\n#pragma unroll\n\n for (int j = 0; j < bsize; j++)\n {\n new_values[dst_base * bsize + j] = h_old_diag[row * bsize + j];\n }\n\n // Increment dst_base by one\n dst_base++;\n }\n\n int end = old_row_offsets[row + 1] - src_base;\n\n for (int i = 0; i < end; i++)\n {\n#pragma unroll\n\n for (int j = 0; j < bsize; j++)\n {\n new_values[ (dst_base + i)*bsize + j ] = h_old_values[ (src_base + i) * bsize + j ];\n }\n }\n\n row += gridDim.x * blockDim.x;\n }\n}\n\n\n__global__ void flag_halo_ids_kernel(INDEX_TYPE *flags, INDEX_TYPE *ids, INDEX_TYPE offset, INDEX_TYPE size, INDEX_TYPE upper)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n flags[ids[idx] - offset] = 1;\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void read_halo_ids_kernel(INDEX_TYPE *flags, INDEX_TYPE *ids, INDEX_TYPE offset, INDEX_TYPE size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n ids[idx] = flags[ids[idx] - offset];\n idx += blockDim.x * gridDim.x;\n }\n}\n\ntemplate\n__global__ void reorder_vector_values(T *dst, const T *src, const INDEX_TYPE *map, INDEX_TYPE blocksize, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * (blockDim.x / blocksize) + threadIdx.x / blocksize; //vectorised by block size\n int vec_id = threadIdx.x % blocksize;\n\n if (threadIdx.x >= (blockDim.x / blocksize)*blocksize ) { return; }\n\n while (row < num_rows)\n {\n dst[map[row]*blocksize + vec_id] = src[row * blocksize + vec_id];\n row += gridDim.x * (blockDim.x / blocksize);\n }\n}\n\ntemplate\n__global__ void inverse_reorder_vector_values(T *dst, T *src, INDEX_TYPE *map, INDEX_TYPE blocksize, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * (blockDim.x / blocksize) + threadIdx.x / blocksize;\n int vec_id = threadIdx.x % blocksize;\n\n if (threadIdx.x >= (blockDim.x / blocksize)*blocksize ) { return; }\n\n while (row < num_rows)\n {\n dst[row * blocksize + vec_id] = src[map[row] * blocksize + vec_id];\n row += gridDim.x * (blockDim.x / blocksize);\n }\n}\n\n__global__ void remove_boundary_kernel(INDEX_TYPE *flags, INDEX_TYPE *maps, INDEX_TYPE size)\n{\n int element = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (element < size)\n {\n flags[maps[element]] = 0; //this won't be a problem, because we are overwriting the same thing\n element += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void get_unassigned_kernel(INDEX_TYPE *unassigned_flags, INDEX_TYPE *map, INDEX_TYPE *output, INDEX_TYPE part_size, INDEX_TYPE uf_size )\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < part_size)\n {\n if (map[idx] < uf_size)\n {\n if (unassigned_flags[map[idx]] == 0)\n {\n unassigned_flags[map[idx]] = 1;\n output[idx] = 1;\n }\n }\n\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void set_unassigned_kernel(INDEX_TYPE *part_assigned_flags, INDEX_TYPE *part_num, INDEX_TYPE *map, INDEX_TYPE *renum, INDEX_TYPE part_size, INDEX_TYPE max_element, INDEX_TYPE renum_size /*, INDEX_TYPE rank*/)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < part_size)\n {\n if (map[idx] < renum_size)\n {\n if (part_assigned_flags[idx] == 1)\n {\n renum[map[idx]] = max_element + part_num[idx];\n }\n\n //also update the B2L map\n map[idx] = renum[map[idx]];\n }\n\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void renumber_b2l_maps(INDEX_TYPE *map, INDEX_TYPE *renum, INDEX_TYPE part_size, INDEX_TYPE renum_size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < part_size)\n {\n if (map[idx] < renum_size)\n {\n //update the B2L map\n map[idx] = renum[map[idx]];\n idx += blockDim.x * gridDim.x;\n }\n }\n}\n\n__global__ void calc_inverse_renumbering(INDEX_TYPE *renum, INDEX_TYPE *irenum, INDEX_TYPE max_element)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < max_element)\n {\n if (renum[idx] < 0 || renum[idx] >= max_element) { printf(\"Renumbering error: %d %d\\n\", renum[idx], max_element); }\n\n irenum[renum[idx]] = idx;\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void create_halo_mapping(INDEX_TYPE *mapping, INDEX_TYPE *node_list, int64_t base_index, INDEX_TYPE map_offset, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n int idx = node_list[row] - base_index;\n mapping[idx] = map_offset + row;\n row += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void apply_h2l2b_mapping(INDEX_TYPE *mapping, INDEX_TYPE *node_list, int64_t base_index, INDEX_TYPE *b2l_map, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n int idx = node_list[row] - base_index;\n mapping[idx] = b2l_map[row];\n row += blockDim.x * gridDim.x;\n }\n}\n\ntemplate \n__global__ void map_col_indices_and_count_rowlen(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, INDEX_TYPE *row_length,\n INDEX_TYPE *mapping, INDEX_TYPE num_rows, INDEX_TYPE insert_diagonal)\n{\n extern __shared__ volatile int reduction[];\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n int valid = 0;\n\n for (int idx = row_offsets[row] + coopIdx; idx < row_offsets[row + 1]; idx += coop) //this may look horrible, but I expect low branch divergence, because col indices in a row usually belong to the same partition (or at most one more)\n {\n int colIdx = col_indices[idx];\n int new_col_idx = mapping[colIdx];\n\n if (new_col_idx >= 0)\n {\n valid++;\n col_indices[idx] = new_col_idx;\n }\n else\n {\n col_indices[idx] = -1;\n }\n }\n\n reduction[threadIdx.x] = valid;\n\n for (int s = 2; s > 0; s >>= 1)\n {\n if (coopIdx < s)\n {\n reduction[threadIdx.x] += reduction[threadIdx.x + s];\n }\n\n __syncthreads();\n }\n\n if (coopIdx == 0)\n {\n row_length[row] = reduction[threadIdx.x] + insert_diagonal;\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n\n__global__ void renumber_P_col_indices(INDEX_TYPE *__restrict__ col_indices, const INDEX_TYPE *__restrict__ renum, INDEX_TYPE num_owned_coarse_pts, INDEX_TYPE num_owned_fine_pts)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < num_owned_fine_pts )\n {\n INDEX_TYPE col_id = col_indices[idx];\n\n if (col_id < num_owned_coarse_pts)\n {\n col_indices[idx] = renum[col_id];\n }\n\n idx += blockDim.x * gridDim.x;\n }\n}\n\n\ntemplate \n__global__ void reorder_R_matrix(const INDEX_TYPE *old_rows, const INDEX_TYPE *old_cols, const T *old_vals, const INDEX_TYPE *rows, INDEX_TYPE *cols, T *vals, const INDEX_TYPE *renumbering, INDEX_TYPE bsize, INDEX_TYPE num_rows, INDEX_TYPE num_owned_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n INDEX_TYPE src_base = old_rows[row];\n INDEX_TYPE dst_base = row < num_owned_rows ? rows[renumbering[row]] : src_base;\n\n for (int i = coopIdx; i < old_rows[row + 1]*bsize - src_base * bsize; i += coop)\n {\n vals[dst_base * bsize + i] = old_vals[src_base * bsize + i];\n }\n\n for (int i = coopIdx; i < old_rows[row + 1] - src_base; i += coop)\n {\n cols[dst_base + i] = old_cols[src_base + i];\n }\n\n row += blockDim.x * gridDim.x / coop;\n }\n}\n\n\n\ntemplate \n__global__ void reorder_whole_matrix(INDEX_TYPE *old_rows, INDEX_TYPE *old_cols, T *old_vals, INDEX_TYPE *rows, INDEX_TYPE *cols, T *vals, INDEX_TYPE *renumbering, INDEX_TYPE bsize, INDEX_TYPE num_rows, INDEX_TYPE insert_diagonal)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n INDEX_TYPE src_base = old_rows[row];\n INDEX_TYPE dst_base = rows[renumbering[row]];\n\n if (insert_diagonal)\n {\n if (coopIdx == 0) { cols[dst_base] = renumbering[row]; }\n\n for (int i = coopIdx; i < bsize; i += coop)\n {\n vals[dst_base * bsize + i] = old_vals[(old_rows[num_rows] + row) * bsize + i];\n }\n\n dst_base++;\n }\n\n for (int i = coopIdx; i < old_rows[row + 1]*bsize - src_base * bsize; i += coop)\n {\n vals[dst_base * bsize + i] = old_vals[src_base * bsize + i];\n }\n\n for (int i = coopIdx; i < old_rows[row + 1] - src_base; i += coop)\n {\n cols[dst_base + i] = old_cols[src_base + i];\n }\n\n row += blockDim.x * gridDim.x / coop;\n }\n}\n\ntemplate \n__global__ void replace_values_matrix(const T *src_vals_h, const T *src_diag_h, const INDEX_TYPE *old_rows, const INDEX_TYPE *rows, T *vals, const INDEX_TYPE *renumbering, INDEX_TYPE bsize, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n INDEX_TYPE src_base = old_rows[row];\n INDEX_TYPE dst_base = rows[renumbering[row]];\n\n for (int i = coopIdx; i < bsize; i += coop)\n {\n vals[dst_base * bsize + i] = src_diag_h[row * bsize + i];\n }\n\n dst_base++;\n\n for (int i = coopIdx; i < old_rows[row + 1]*bsize - src_base * bsize; i += coop)\n {\n vals[dst_base * bsize + i] = src_vals_h[src_base * bsize + i];\n }\n\n row += blockDim.x * gridDim.x / coop;\n }\n}\n\ntemplate \n__global__ void replace_values_matrix(const T *src_vals_h, const INDEX_TYPE *old_rows, const INDEX_TYPE *rows, T *vals, const INDEX_TYPE *renumbering, INDEX_TYPE bsize, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n INDEX_TYPE src_base = old_rows[row];\n INDEX_TYPE dst_base = rows[renumbering[row]];\n\n for (int i = coopIdx; i < old_rows[row + 1]*bsize - src_base * bsize; i += coop)\n {\n vals[dst_base * bsize + i] = src_vals_h[src_base * bsize + i];\n }\n\n row += blockDim.x * gridDim.x / coop;\n }\n}\n\n//TODO: optimize by vectorizing\ntemplate \n__global__ void reorder_whole_halo_matrix(INDEX_TYPE *old_rows, INDEX_TYPE *old_cols, T *old_vals, INDEX_TYPE *rows, INDEX_TYPE *cols, T *vals,\n INDEX_TYPE *renumbering, INDEX_TYPE bsize, INDEX_TYPE num_rows, INDEX_TYPE insert_diagonal,\n INDEX_TYPE global_offset, INDEX_TYPE local_offset, INDEX_TYPE halo_rows)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < num_rows)\n {\n INDEX_TYPE src_base = old_rows[row];\n INDEX_TYPE dst = rows[row];\n\n if (insert_diagonal)\n {\n cols[dst] = global_offset + row;\n\n for (int j = 0; j < bsize; j++) { vals[dst * bsize + j] = old_vals[(old_rows[halo_rows - local_offset] + local_offset + row) * bsize + j]; }\n\n dst++;\n }\n\n for (int i = 0; i < old_rows[row + 1] - src_base; i++)\n {\n INDEX_TYPE colIdx = old_cols[src_base + i];\n\n if (colIdx >= 0)\n {\n cols[dst] = colIdx;\n\n for (int j = 0; j < bsize; j++) { vals[dst * bsize + j] = old_vals[(src_base + i) * bsize + j]; }\n\n dst++;\n }\n }\n\n row += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void calc_rowlen_reorder(INDEX_TYPE *row_offsets, INDEX_TYPE *row_len, INDEX_TYPE *map, INDEX_TYPE size, INDEX_TYPE insert_diag)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n row_len[map[idx]] = row_offsets[idx + 1] - row_offsets[idx] + insert_diag;\n idx += blockDim.x * gridDim.x;\n }\n}\n\n\ntemplate < class TConfig >\nvoid DistributedManagerBase::remove_boundary(IVector_d &flagArray, IVector_d &B2L_map, int size)\n{\n int num_blocks = std::min(4096, (size + 127) / 128);\n remove_boundary_kernel <<< num_blocks, 128>>>(flagArray.raw(), B2L_map.raw(), size);\n cudaCheckError();\n}\n\ntemplate < class TConfig >\nvoid DistributedManagerBase::get_unassigned(IVector_d &flagArray, IVector_d &B2L_map, IVector_d &partition_flags, int size, int global_size /*, int rank*/)\n{\n int num_blocks = std::min(4096, (size + 191) / 192);\n get_unassigned_kernel <<< num_blocks, 192>>>(flagArray.raw(),\n B2L_map.raw(),\n partition_flags.raw(), size, global_size /*, rank*/);\n cudaCheckError();\n}\n\ntemplate < class TConfig >\nvoid DistributedManagerBase::set_unassigned(IVector_d &partition_flags, IVector_d &partition_renum, IVector_d &B2L_map, IVector_d &renumbering, int size, int max_element, int global_size /*, int rank*/)\n{\n int num_blocks = std::min(4096, (size + 191) / 192);\n set_unassigned_kernel <<< num_blocks, 192>>>(partition_flags.raw(),\n partition_renum.raw(),\n B2L_map.raw(),\n renumbering.raw(),\n size, max_element, global_size /*,rank*/);\n cudaCheckError();\n}\n\n\ntemplate \ninline void DistributedManagerBase::set_initialized(IVector &row_offsets)\n{\n // For P and R sizes the sizes are fixed at creation\n if(m_fixed_view_size)\n {\n return;\n }\n\n if (neighbors.size() > 0)\n {\n //distributed: cache num_rows/num_nz for different views\n _num_rows_interior = _num_interior_nodes;\n _num_nz_interior = row_offsets[_num_rows_interior];\n _num_rows_owned = _num_interior_nodes + _num_boundary_nodes;\n _num_nz_owned = row_offsets[_num_rows_owned];\n _num_rows_full = halo_offsets[neighbors.size()];\n\n if (_num_rows_full >= row_offsets.size())\n {\n _num_nz_full = row_offsets[row_offsets.size() - 1];\n }\n else\n {\n _num_nz_full = row_offsets[_num_rows_full];\n }\n\n _num_rows_all = halo_offsets[halo_offsets.size() - 1];\n _num_nz_all = _num_nz_full;\n }\n else\n {\n _num_rows_interior = _num_interior_nodes;\n _num_nz_interior = row_offsets[_num_rows_interior];\n _num_rows_owned = _num_interior_nodes;\n _num_nz_owned = row_offsets[_num_rows_owned];\n _num_rows_full = _num_rows_owned;\n _num_nz_full = _num_nz_owned;\n _num_rows_all = _num_rows_owned;\n _num_nz_all = _num_nz_owned;\n }\n}\n\ntemplate \nvoid DistributedManagerBase::createAggregatesRenumbering(IVector_h &renumbering, IVector_h_vector &B2L_maps, int size, int num_neighbors, int &num_interior_aggregates, int &num_boundary_aggregates, int num_rings)\n{\n createAggRenumbering(renumbering, B2L_maps, size, num_neighbors, num_interior_aggregates, num_boundary_aggregates, num_rings);\n}\n\ntemplate \nvoid DistributedManagerBase::createAggregatesRenumbering(IVector_d &renumbering, IVector_d_vector &B2L_maps, int size, int num_neighbors, int &num_interior_aggregates, int &num_boundary_aggregates, int num_rings)\n{\n createAggRenumbering(renumbering, B2L_maps, size, num_neighbors, num_interior_aggregates, num_boundary_aggregates, num_rings);\n}\n\ntemplate \ntemplate \nvoid DistributedManagerBase::createAggRenumbering(IVector_hd &renumbering, std::vector &B2L_maps, int size, int num_neighbors, int &num_interior_aggregates, int &num_boundary_aggregates, int num_rings)\n{\n if (num_rings != 1)\n {\n FatalError(\"num_rings > 1 not supported in consolidation\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n //int num_neighbors = this->neighbors.size();\n if (num_neighbors == 0)\n {\n num_boundary_aggregates = 0;\n num_interior_aggregates = size;\n return;\n }\n\n //initial size to size+1 so we have the total size after a scan\n int global_size = size;\n renumbering.resize(size + 1);\n //\n // Step 1 - in the main matrix, separate interior and boundary nodes (1/0 in flagArray), renumber interior ones with an exclusive scan\n //\n IVector_hd flagArray(size + 1);\n thrust_wrapper::fill(flagArray.begin(), flagArray.begin() + size + 1, 1);\n cudaCheckError();\n\n //sets 1 for interior nodes, 0 for boundary node\n for (int i = 0; i < num_neighbors; i++ )\n {\n int size = B2L_maps[i].size();\n remove_boundary(flagArray, B2L_maps[i], size);\n }\n\n //gets the renumbering of interior nodes\n thrust_wrapper::exclusive_scan(flagArray.begin(), flagArray.begin() + size + 1, renumbering.begin());\n cudaCheckError();\n //\n // Step 2 - Renumber nodes that are in the boundary, stepping through each B2L map, and renumbering ones that have not been renumbered yet\n //\n //what is the biggest B2L size\n INDEX_TYPE max_size = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n max_size = max_size > B2L_maps[i].size() ? max_size : B2L_maps[i].size();\n }\n\n //allocate work vectors (should be pretty small)\n IVector_hd partition_flags(max_size);\n IVector_hd partition_renum(max_size);\n //the number of renumbered nodes so far\n int max_element = renumbering[size];\n num_interior_aggregates = max_element;\n num_boundary_aggregates = size - max_element;\n renumbering.resize(size);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n //find nodes that are part of the current boundary and they haven't been renumbered yet\n thrust_wrapper::fill(partition_flags.begin(), partition_flags.begin() + max_size, 0);\n int size = B2L_maps[i].size();\n get_unassigned(flagArray, B2L_maps[i], partition_flags, size, global_size/*,0*/);\n //calculate the local renumbering (within this boundary region) of these nodes\n thrust_wrapper::exclusive_scan(partition_flags.begin(), partition_flags.begin() + max_size, partition_renum.begin());\n //apply renumbering to the big numbering table\n set_unassigned(partition_flags, partition_renum, B2L_maps[i], renumbering, size, max_element, global_size/*,0*/);\n //update the number of renumbered nodes\n max_element += partition_renum[max_size - 1] + partition_flags[max_size - 1];\n }\n\n cudaCheckError();\n}\n\n\ntemplate \ninline DistributedManagerBase::DistributedManagerBase(Matrix &a) :\n m_fine_level_comms(NULL), A(&a), m_pinned_buffer_size(0), m_pinned_buffer(NULL), _num_interior_nodes(0), _num_boundary_nodes(0), _comms(NULL), has_B2L(false),\n neighbors(_neighbors), B2L_maps(_B2L_maps), L2H_maps(_L2H_maps), B2L_rings(_B2L_rings),\n halo_rows_ref_count(0), halo_btl_ref_count(0), halo_ranges(_halo_ranges), halo_ranges_h(_halo_ranges_h), part_offsets(_part_offsets), part_offsets_h(_part_offsets_h), halo_rows(NULL), halo_btl(NULL), m_is_root_partition(false), m_is_glued(false), m_is_fine_level_glued(false), m_is_fine_level_consolidated(false), m_is_fine_level_root_partition(false), m_use_cuda_ipc_consolidation(false), m_fixed_view_size(false)\n{\n cudaEventCreate(&comm_event);\n cudaCheckError();\n cudaStreamCreateWithFlags(&m_int_stream, cudaStreamNonBlocking);\n cudaCheckError();\n cudaStreamCreateWithFlags(&m_bdy_stream, cudaStreamNonBlocking);\n cudaCheckError();\n this->createComms(A->getResources());\n int my_id = this->getComms()->get_global_id();\n int num_parts = this->getComms()->get_num_partitions();\n this->set_global_id(my_id);\n this->set_num_partitions(num_parts);\n}\n\ntemplate \nvoid DistributedManager >::generatePoisson7pt(int nx, int ny, int nz, int P, int Q, int R)\n{\n int my_id = this->getComms()->get_global_id();\n int p, q, r;\n\n if (nx < P || ny < Q || nz < R)\n {\n FatalError(\"(nx < P) or (ny < Q) or (nz < R) not supported\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n /* compute p,q,r from P,Q,R and myid */\n p = my_id % P; // Position in x direction\n q = (( my_id - p) / P) % Q; // Position in y\n r = ( my_id - p - P * q) / ( P * Q ); // Position in z\n // Create A.row_indices, A.col_indices, A.values, A.diag\n int num_rows = nx * ny * nz;\n int num_nonzeros = num_rows * 7; // Ignoring any boundary, 7 nnz per row\n int num_substract = 0;\n\n if (p == 0) { num_substract += ny * nz; }\n\n if (p == P - 1) { num_substract += ny * nz; }\n\n if (q == 0) { num_substract += nx * nz; }\n\n if (q == Q - 1) { num_substract += nx * nz; }\n\n if (r == 0) { num_substract += nx * ny; }\n\n if (r == R - 1) { num_substract += nx * ny; }\n\n num_nonzeros -= num_substract;\n int num_halo_nodes = 2 * (ny * nz + nx * nz + nx * ny) - num_substract;\n this->local_to_global_map.resize(num_halo_nodes);\n this->A->set_initialized(0);\n this->A->resize(0, 0, 0, 1, 1, 1);\n this->A->addProps(CSR);\n this->A->resize(num_rows, num_rows + num_halo_nodes, num_nonzeros, 1, 1, 1);\n const int cta_size = 128;\n const int grid_size = std::min( 4096, (num_rows + cta_size - 1) / cta_size );\n poisson7pt_count_row_len <<< grid_size, cta_size>>>(this->A->row_offsets.raw(), nx, ny, nz, p, q, r, P, Q, R, num_rows);\n cudaCheckError();\n thrust_wrapper::exclusive_scan(this->A->row_offsets.begin(), this->A->row_offsets.end(), this->A->row_offsets.begin());\n cudaCheckError();\n // Now set nonzeros columns and values\n // TODO: vectorize this\n const int grid_size2 = std::min( 4096, (num_rows + cta_size - 1) / cta_size );\n poisson7pt_set_col_values <<< grid_size2, cta_size>>>\n (this->A->row_offsets.raw(),\n this->A->col_indices.raw(),\n this->A->values.raw(),\n this->A->diag.raw(),\n this->local_to_global_map.raw(),\n nx, ny, nz,\n p, q, r,\n P, Q, R,\n num_rows);\n cudaCheckError();\n // fill parts_offsets_h\n // All ranks have same number of nodes\n int num_ranks = P * Q * R;\n this->part_offsets_h.resize(num_ranks + 1);\n this->part_offsets_h[0] = (int64_t) 0;\n\n for (int i = 1; i < num_ranks + 1; i++)\n {\n this->part_offsets_h[i] = this->part_offsets_h[i - 1] + (int64_t) num_rows;\n }\n\n // Device to host copy\n this->part_offsets = this->part_offsets_h;\n this->num_rows_global = P * Q * R * nx * ny * nz;\n// this->A->set_initialized(1);\n}\n\ntemplate \ntemplate \nvoid DistributedManager >::loadDistributed_SetOffsets(\n int num_ranks, int num_rows_global, const t_colIndex* partition_offsets)\n{\n // fill part offsets internal data structures\n this->part_offsets_h.resize(num_ranks + 1);\n\n for (int i = 0; i <= num_ranks; i++)\n {\n this->part_offsets_h[i] = partition_offsets[i];\n }\n // copy to device\n this->part_offsets = this->part_offsets_h;\n // set num of global rows\n this->num_rows_global = num_rows_global;\n cudaCheckError();\n}\n\ntemplate \ntemplate \nstd::map DistributedManager >::loadDistributed_LocalToGlobal(int num_rows, I64Vector_h &off_diag_cols)\n{\n // sort global column indices\n amgx::thrust::sort(off_diag_cols.begin(), off_diag_cols.end());\n // find unique columns and set local <-> global mappings\n // 1) Removed unneeded vector 2) Create map on host first, upload later (less thrust calls)\n I64Vector_h local_to_global_h;\n std::map global_to_local; // temporary\n\n if (off_diag_cols.size() > 0)\n {\n global_to_local[off_diag_cols[0]] = num_rows;\n local_to_global_h.push_back(off_diag_cols[0]);\n }\n\n for (int i = 1; i < off_diag_cols.size(); i++)\n {\n if (off_diag_cols[i] != off_diag_cols[i - 1])\n {\n global_to_local[off_diag_cols[i]] = num_rows + local_to_global_h.size();\n local_to_global_h.push_back(off_diag_cols[i]);\n }\n }\n // Upload finished map in one piece\n this->local_to_global_map.resize(local_to_global_h.size());\n amgx::thrust::copy(local_to_global_h.begin(), local_to_global_h.end(), this->local_to_global_map.begin());\n return global_to_local;\n}\n\ntemplate \nvoid DistributedManager >::loadDistributed_InitLocalMatrix(\n IVector_h local_col_indices,\n int num_rows,\n int num_nonzeros,\n const int block_dimx,\n const int block_dimy,\n const int *row_offsets,\n const mat_value_type *values,\n const void *diag)\n{\n // init local matrix\n this->A->set_initialized(0);\n this->A->resize(0, 0, 0, 1, 1, 1);\n this->A->addProps(CSR);\n\n if (diag)\n {\n this->A->addProps(DIAG);\n }\n\n this->A->resize(num_rows, num_rows + this->local_to_global_map.size(), num_nonzeros, block_dimx, block_dimy, 1);\n cudaCheckError();\n // set local matrix\n amgx::thrust::copy(row_offsets, row_offsets + num_rows + 1, this->A->row_offsets.begin());\n this->A->col_indices = local_col_indices;\n\n amgx::thrust::copy(values, values + num_nonzeros * block_dimx * block_dimy, this->A->values.begin());\n cudaCheckError();\n\n // setup diagonal\n if (diag)\n {\n cudaMemcpy(this->A->values.raw() + this->A->diagOffset()*this->A->get_block_size(), diag, sizeof(mat_value_type) * num_rows * block_dimx * block_dimy, cudaMemcpyDefault);\n cudaCheckError();\n }\n else\n {\n this->A->computeDiagonal();\n }\n cudaCheckError();\n}\n\ntemplate \ntemplate \nvoid DistributedManager >::loadDistributedMatrixPartitionVec(\n int num_rows,\n int num_nonzeros,\n const int block_dimx,\n const int block_dimy,\n const int *row_offsets,\n const t_colIndex *col_indices,\n const mat_value_type *values,\n int num_ranks,\n int num_rows_global,\n const void *diag,\n const int *partition)\n{\n // fetch my rank\n int my_id = this->getComms()->get_global_id();\n // setup partition vector\n IVector_h partitionVec(num_rows_global);\n\n if (partition == NULL)\n {\n IVector_h rowCounts(num_ranks);\n this->getComms()->all_gather(num_rows, rowCounts, 1);\n\n int p = 0;\n for (int i = 0; i < num_ranks; ++i)\n {\n for (int j = 0; j < rowCounts[i]; ++j)\n {\n partitionVec[p++] = i;\n }\n }\n }\n else\n {\n // use existing partition info\n for (int i = 0; i < num_rows_global; i++)\n {\n partitionVec[i] = partition[i];\n }\n }\n\n // compute partition offsets (based on number of elements per partition). Will be modified when calculating partition map.\n t_colIndex *partition_offsets = (t_colIndex *)calloc(num_ranks + 1, sizeof(t_colIndex));\n\n for (int i = 0; i < num_rows_global; i++)\n {\n int pvi = partitionVec[i];\n partition_offsets[pvi + 1]++;\n }\n\n thrust_wrapper::inclusive_scan(partition_offsets, partition_offsets + num_ranks + 1, partition_offsets);\n\n loadDistributed_SetOffsets(num_ranks, num_rows_global, partition_offsets);\n\n // compute partition map (which tells you how the global elements are mapped into the partitions)\n t_colIndex *partition_map = (t_colIndex *)calloc(num_rows_global, sizeof(t_colIndex));\n\n for (int i = 0; i < num_rows_global; i++)\n {\n int pvi = partitionVec[i];\n t_colIndex poi = partition_offsets[pvi];\n partition_map[poi] = i;\n partition_offsets[pvi]++;\n }\n free(partition_offsets);\n\n // compute the inverse partition map\n t_colIndex *ipartition_map = (t_colIndex *)calloc(num_rows_global, sizeof(t_colIndex));\n\n for (int i = 0; i < num_rows_global; i++)\n {\n ipartition_map[partition_map[i]] = i;\n }\n free(partition_map);\n\n int h_cidx_allocated = 0;\n const t_colIndex *h_col_indices_global = (const t_colIndex *)this->getHostPointerForData(col_indices, num_nonzeros * sizeof(t_colIndex), &h_cidx_allocated);\n // gather all off-diag columns\n I64Vector_h off_diag_cols;\n\n for (int i = 0; i < num_nonzeros; i++)\n {\n if (partitionVec[h_col_indices_global[i]] != my_id)\n {\n off_diag_cols.push_back(ipartition_map[h_col_indices_global[i]]);\n }\n }\n\n auto global_to_local = loadDistributed_LocalToGlobal(num_rows, off_diag_cols);\n\n // set 1, then scan to compute local row indices\n IVector_h my_indices(num_rows_global);\n\n for (int i = 0; i < num_nonzeros; i++)\n {\n if (partitionVec[h_col_indices_global[i]] == my_id) // find my local columns and set to 1\n {\n my_indices[ipartition_map[h_col_indices_global[i]]] = 1;\n }\n }\n\n thrust_wrapper::exclusive_scan(my_indices.begin(), my_indices.end(), my_indices.begin());\n // remap colums to local\n IVector_h local_col_indices(num_nonzeros);\n\n for (int i = 0; i < num_nonzeros; i++)\n {\n if (partitionVec[h_col_indices_global[i]] != my_id)\n {\n // off-diag\n local_col_indices[i] = global_to_local[ipartition_map[h_col_indices_global[i]]];\n }\n else\n {\n // diag\n local_col_indices[i] = my_indices[ipartition_map[h_col_indices_global[i]]];\n }\n }\n free(ipartition_map);\n\n loadDistributed_InitLocalMatrix(local_col_indices, num_rows, num_nonzeros, block_dimx, block_dimy, row_offsets, values, diag);\n\n cudaCheckError();\n\n // don't free possibly allocated pinned buffer, since it could be used later. if it would not - it would be deallocated automatically\n /*if (h_cidx_allocated)\n {\n free((void*)h_col_indices_global);\n }*/\n}\n\ntemplate \ntemplate \nvoid DistributedManager >::loadDistributedMatrixPartitionOffsets(\n int num_rows,\n int num_nonzeros,\n const int block_dimx,\n const int block_dimy,\n const int *row_offsets,\n const t_colIndex *col_indices,\n const mat_value_type *values,\n int num_ranks,\n int num_rows_global,\n const void *diag,\n const t_colIndex *partition_offsets)\n{\n // fetch my rank\n int my_id = this->getComms()->get_global_id();\n // sanity check, cheap to perform, and helps prevent harder-to-debug errors later on\n if (!std::is_sorted(partition_offsets, partition_offsets + num_ranks + 1)) {\n FatalError(\"Partition offsets are not sorted.\", AMGX_ERR_BAD_PARAMETERS);\n }\n loadDistributed_SetOffsets(num_ranks, num_rows_global, partition_offsets);\n\n // Create predicate to determine if a column is in the local diagonal block\n t_colIndex my_first_col = this->part_offsets_h[my_id];\n t_colIndex one_past_my_last_col = this->part_offsets_h[my_id + 1];\n auto in_local_diagonal_block = [my_first_col, one_past_my_last_col](const t_colIndex col_index) {\n return col_index >= my_first_col && col_index < one_past_my_last_col;\n };\n\n int h_cidx_allocated = 0;\n const t_colIndex *h_col_indices_global = (const t_colIndex *)this->getHostPointerForData(col_indices, num_nonzeros * sizeof(t_colIndex), &h_cidx_allocated);\n // gather all off-diag columns\n I64Vector_h off_diag_cols;\n for (int i = 0; i < num_nonzeros; i++)\n {\n if (!in_local_diagonal_block(h_col_indices_global[i]))\n {\n off_diag_cols.push_back(h_col_indices_global[i]);\n }\n }\n auto global_to_local = loadDistributed_LocalToGlobal(num_rows, off_diag_cols);\n // set 1, then scan to compute local row indices\n // \"coordinate-shift\" columns so they lie in much smaller range of my diagonal indices\n int diagonal_size = this->part_offsets_h[my_id + 1] - this->part_offsets_h[my_id];\n IVector_h my_indices(diagonal_size);\n for (int i = 0; i < num_nonzeros; i++)\n {\n t_colIndex col_index = h_col_indices_global[i];\n if (in_local_diagonal_block(h_col_indices_global[i])) // find my local columns and set to 1\n {\n // olumns that are on *my* diag partition cannot have an index from 0..num_rows_global\n // instead, part_offsets_h[my_id] <= col_index < part_offsets[my_id+1]\n col_index -= this->part_offsets_h[my_id];\n my_indices[col_index] = 1;\n }\n }\n thrust_wrapper::exclusive_scan(my_indices.begin(), my_indices.end(), my_indices.begin());\n\n // remap colums to local\n IVector_h local_col_indices(num_nonzeros);\n for (int i = 0; i < num_nonzeros; i++)\n {\n t_colIndex col_index = h_col_indices_global[i];\n if (!in_local_diagonal_block(col_index))\n {\n // off-diag\n local_col_indices[i] = global_to_local[col_index];\n }\n else\n {\n // diag\n col_index -= this->part_offsets_h[my_id];\n local_col_indices[i] = my_indices[col_index];\n }\n }\n loadDistributed_InitLocalMatrix(local_col_indices, num_rows, num_nonzeros, block_dimx, block_dimy, row_offsets, values, diag);\n}\n\ntemplate \ntemplate \nvoid DistributedManager >::loadDistributedMatrix(\n int num_rows,\n int num_nonzeros,\n const int block_dimx,\n const int block_dimy,\n const int *row_offsets,\n const t_colIndex *col_indices,\n const mat_value_type *values,\n int num_ranks,\n int num_rows_global,\n const void *diag,\n const MatrixDistribution &dist)\n{\n using PI = MatrixDistribution::PartitionInformation;\n switch (dist.getPartitionInformationStyle()) {\n case PI::PartitionVec:\n loadDistributedMatrixPartitionVec(num_rows, num_nonzeros, block_dimx, block_dimy, \n row_offsets, col_indices, values, num_ranks, num_rows_global, diag, (const int*) dist.getPartitionData());\n break;\n case PI::PartitionOffsets:\n loadDistributedMatrixPartitionOffsets(num_rows, num_nonzeros, block_dimx, block_dimy, \n row_offsets, col_indices, values, num_ranks, num_rows_global, diag, (const t_colIndex*) dist.getPartitionData());\n break;\n default:\n FatalError(\"Unsupported partitioning data format used with loadDistributedMatrix\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::renumberMatrixOneRing(int update_neighbours)\n{\n FatalError(\"Distributed classical AMG not implemented on host\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid DistributedManager >::renumberMatrixOneRing(int update_neighbours)\n{\n // Step 1: Using halo_ranges, flag neighbors and at the same time, flag halo_nodes (flag_halo_nodes_local)\n int my_id = this->global_id();\n int num_parts = this->get_num_partitions();\n this->set_base_index(this->part_offsets_h[my_id]);\n this->set_index_range(this->part_offsets_h[my_id + 1] - this->part_offsets_h[my_id]);\n DistributedArranger *prep = new DistributedArranger;\n\n // Create/update list of neighbors\n if (update_neighbours)\n {\n typedef typename TConfig::template setVecPrec::Type i64vec_value_type;\n typedef Vector I64Vector;\n typedef typename Matrix::MVector MVector;\n std::vector halo_row_offsets(this->neighbors.size());\n std::vector halo_global_indices(this->neighbors.size());\n std::vector halo_values(this->neighbors.size());\n prep->create_halo_rows_global_indices(*(this->A), halo_row_offsets, halo_global_indices, halo_values);\n prep->update_neighbors_list(*(this->A), this->neighbors, this->halo_ranges_h, this->halo_ranges, this->part_offsets_h, this->part_offsets, halo_row_offsets, halo_global_indices);\n }\n else\n {\n prep->create_neighbors_v2(*(this->A));\n }\n\n this->getComms()->set_neighbors(this->neighbors.size());\n // Create B2L_maps and L2H_maps\n prep->create_boundary_lists_v3(*(this->A));\n // halo_offsets\n int neighbors = this->A->manager->num_neighbors();\n int A_num_rows, offset;\n this->A->getOffsetAndSizeForView(OWNED, &offset, &A_num_rows);\n this->halo_offsets.resize(neighbors + 1, 0);\n this->halo_offsets[0] = A_num_rows;\n\n for (int i = 0; i < neighbors; i++)\n {\n this->halo_offsets[i + 1] = this->halo_offsets[i] + this->B2L_maps[i].size();\n }\n\n this->getComms()->exchange_vectors(this->A->manager->B2L_maps, *(this->A), 0);\n // Initialize B2L_rings\n int num_neighbors = this->neighbors.size();\n this->B2L_rings.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n this->B2L_rings[i].resize(2);\n this->B2L_rings[i][0] = 0;\n this->B2L_rings[i][1] = this->B2L_maps[i].size();\n }\n\n prep->initialize_B2L_maps_offsets(*(this->A), 1);\n delete prep;\n //Use the exchanged halo row matrices and the boundary/halo index lists to renumber and consolidate the matrix\n // Step 5: renumber all owned rows and columns\n this->reorder_matrix_owned();\n // Step 6: renumber local_to_global_map\n int num_owned_rows = this->A->manager->halo_offsets[0];\n int size_one_ring;\n this->A->getOffsetAndSizeForView(FULL, &offset, &size_one_ring);\n I64Vector_d global_col_indices(size_one_ring);\n thrust_wrapper::sequence(global_col_indices.begin(), global_col_indices.begin() + num_owned_rows, this->base_index() );\n cudaCheckError();\n global_col_indices.dirtybit = 1;\n this->exchange_halo(global_col_indices, global_col_indices.tag);\n thrust_wrapper::copy(global_col_indices.begin() + num_owned_rows, global_col_indices.begin() + size_one_ring, this->local_to_global_map.begin(), this->get_int_stream(), true);\n cudaCheckError();\n}\n\n\ntemplate \nvoid DistributedManager >::renumber_P_R(Matrix_h &P, Matrix_h &R, Matrix_h &A_fine)\n{\n FatalError(\"Distributed classical AMG not implemented on host\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid DistributedManager >::renumber_P_R(Matrix_d &P, Matrix_d &R, Matrix_d &A_fine)\n{\n int cta_size = 256;\n int num_owned_fine_pts = A_fine.manager->halo_offsets[0];\n int num_owned_coarse_pts, offset;\n // matrix Ac\n this->A->getOffsetAndSizeForView(OWNED, &offset, &num_owned_coarse_pts);\n // Renumber the owned col indices of P (not the halo columns ,since P.manager was created assunming some other numbering)\n int nnz_owned_fine_pts = P.row_offsets[num_owned_fine_pts];\n int num_blocks_fine = std::min(4096, (nnz_owned_fine_pts + cta_size - 1) / cta_size);\n\n if (num_blocks_fine > 0)\n {\n renumber_P_col_indices <<< num_blocks_fine, cta_size>>>(P.col_indices.raw(), this->renumbering.raw(), num_owned_coarse_pts, nnz_owned_fine_pts);\n cudaCheckError();\n }\n\n // Renumber the B2L_maps of P\n for (int i = 0; i < P.manager->neighbors.size(); i++)\n {\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(this->renumbering.begin(), P.manager->B2L_maps[i].begin()),\n amgx::thrust::make_permutation_iterator(this->renumbering.begin(), P.manager->B2L_maps[i].end()),\n P.manager->B2L_maps[i].begin());\n }\n\n cudaCheckError();\n // Don't renumber the L2H_maps or the halo\n // Renumber the local_to_global_map of matrix P (since neighbors renumbered their owned rows)\n // Swap owned rows of R\n IVector new_row_offsets(R.row_offsets.size());\n int insert = 0;\n // Only renumber the owned rows\n int num_blocks_owned = std::min(4096, (num_owned_coarse_pts + cta_size - 1) / cta_size);\n\n if (num_blocks_owned > 0)\n {\n calc_rowlen_reorder <<< num_blocks_owned, cta_size >>>(R.row_offsets.raw(), new_row_offsets.raw(), this->renumbering.raw(), num_owned_coarse_pts, insert);\n cudaCheckError();\n }\n\n thrust_wrapper::exclusive_scan(new_row_offsets.begin(), new_row_offsets.begin() + num_owned_coarse_pts + 1, new_row_offsets.begin());\n cudaCheckError();\n // Copy the row_offsets for halo rows\n amgx::thrust::copy(R.row_offsets.begin() + num_owned_coarse_pts, R.row_offsets.end(), new_row_offsets.begin() + num_owned_coarse_pts);\n cudaCheckError();\n // Reorder the rows of R (no need to reorder the column indices)\n int new_nnz = new_row_offsets[new_row_offsets.size() - 1];\n int halo_offset = new_row_offsets[num_owned_coarse_pts];\n typedef typename MatPrecisionMap::Type ValueTypeA;\n VVector new_values(new_nnz * R.get_block_size(), types::util< ValueTypeA >::get_zero());\n IVector new_col_indices(new_nnz, 0);\n int num_blocks_total = std::min(4096, (R.get_num_rows() + cta_size - 1) / cta_size);\n\n if (num_blocks_total > 0)\n {\n reorder_R_matrix <32> <<< num_blocks_total, 512>>>(R.row_offsets.raw(), R.col_indices.raw(), R.values.raw(), new_row_offsets.raw(), new_col_indices.raw(), new_values.raw(), this->renumbering.raw(), R.get_block_size(), R.get_num_rows(), num_owned_coarse_pts);\n cudaCheckError();\n }\n\n R.col_indices.swap(new_col_indices);\n R.row_offsets.swap(new_row_offsets);\n R.values.swap(new_values);\n\n // Renumber the local_to_global_map (since neighbors have changed their owned numbering)\n if (P.manager->neighbors.size() != 0)\n {\n int size_one_ring = P.manager->halo_offsets[P.manager->neighbors.size()];\n I64Vector_d global_col_indices(size_one_ring);\n thrust_wrapper::sequence(global_col_indices.begin(), global_col_indices.begin() + num_owned_coarse_pts, this->base_index());\n cudaCheckError();\n global_col_indices.dirtybit = 1;\n P.manager->exchange_halo(global_col_indices, global_col_indices.tag);\n thrust_wrapper::copy(global_col_indices.begin() + num_owned_coarse_pts, global_col_indices.begin() + size_one_ring, P.manager->local_to_global_map.begin(), this->get_int_stream(), true);\n cudaCheckError();\n }\n\n DistributedArranger *prep = new DistributedArranger;\n prep->initialize_B2L_maps_offsets(P, 1);\n delete prep;\n}\n\n\ntemplate \nvoid DistributedManager >::createOneRingB2Lmaps()\n{\n // Input:\n // a matrix with N rows, whose column indices are local indices from 0 to N+M-1,\n // where M is a number of 1-ring halo vertices\n // The matrix also contains array \"local_to_global_map\" of size M, which stores the global index of each halo index\n // Ex: assuming a column has index N+K, where 0 <= K < M, then it's global id is local_to_global_map[K]\n // The matrix also contains part_offsets_h and part_offsets array, which stores where each partition begins\n // Output:\n // This function creates all the necessary data to to 1-ring exchanges\n // i.e. list of 1-ring neighbors, B2L_maps for 1-ring, halo_offsets for 1-ring,\n // Also, the function reorders the halo indices, such that 1-ring indices are in the order\n // of neighbors, and therefore, exchange_halo doesn't have to be changed (i.e. L2H = identity)\n // What is does:\n // Based on the global indices of its halo vertices, count the number of neighbors\n // For each neighbor, receive the halo indices that will be needed by neighbor\n // From those, create B2L_maps[0], which contains for all neighbors\n // This function assumes that:\n // part_offset is defined\n\n // B2L_maps\n int my_id = this->global_id();\n int num_parts = this->get_num_partitions();\n this->set_base_index(this->part_offsets_h[my_id]);\n this->set_index_range(this->part_offsets_h[my_id + 1] - this->part_offsets_h[my_id]);\n DistributedArranger *prep = new DistributedArranger;\n // This function creates the array neighbors, which contains a list of partitions to which data\n // needs to be sent and/or received\n prep->create_neighbors_v2(*(this->A));\n // Here change the manager if some partitions have no neighbors\n this->getComms()->set_neighbors(this->neighbors.size());\n prep->create_B2L_one_ring(*(this->A));\n delete prep;\n}\n\ntemplate \nvoid DistributedManager >::createOneRingHaloRows()\n{\n // Input:\n // A matrix with 1-ring B2L_maps, 1-ring halo_offsets\n // Outputs:\n // A matrix with: 1-ring rows,\n // 2-ring B2L_maps,\n // 2-ring halo_offsets\n // 2-ring neighbors\n // Implement here:\n // Look at function create_B2L_from_maps, which calls create_rings, create_halo_btl, create_halo_rows and comms->exchange_matrix_halo\n DistributedArranger *prep = new DistributedArranger;\n prep->create_one_ring_halo_rows(*(this->A));\n // I believe this can be removed since we don't use masked SpMV anymore\n prep->createRowsLists(*(this->A), false);\n delete prep;\n // this is not necessary anymore becasue we don't use latency hiding\n // however in future we might want to get back to this in case we want to use latency hiding\n //this->reorder_matrix();\n}\n\ntemplate \ninline DistributedManagerBase::DistributedManagerBase(\n Matrix &a,\n INDEX_TYPE allocated_halo_depth,\n INDEX_TYPE num_import_rings,\n int num_neighbors,\n const VecInt_t *neighbors_) : m_fine_level_comms(NULL), A(&a), m_pinned_buffer_size(0), m_pinned_buffer(NULL), _num_interior_nodes(0), _num_boundary_nodes(0), _comms(NULL), has_B2L(false), neighbors(_neighbors), halo_rows_ref_count(0), halo_rows(NULL), halo_btl_ref_count(0), halo_btl(NULL), halo_ranges(_halo_ranges), halo_ranges_h(_halo_ranges_h), part_offsets(_part_offsets), part_offsets_h(_part_offsets_h),\n B2L_maps(_B2L_maps), L2H_maps(_L2H_maps), B2L_rings(_B2L_rings), m_is_root_partition(false), m_is_glued(false), m_is_fine_level_glued(false), m_is_fine_level_consolidated(false), m_is_fine_level_root_partition(false), m_use_cuda_ipc_consolidation(false), m_fixed_view_size(false)\n{\n cudaStreamCreateWithFlags(&m_int_stream, cudaStreamNonBlocking);\n cudaCheckError();\n cudaStreamCreateWithFlags(&m_bdy_stream, cudaStreamNonBlocking);\n cudaCheckError();\n\n if (num_import_rings != 1)\n {\n FatalError(\"num_rings > 1 not supported in fine_level consolidation\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (allocated_halo_depth != 1)\n {\n FatalError(\"allocated_halo_depth > 1 not supported in fine_level consolidation\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n this->set_num_halo_rings(num_import_rings);\n neighbors.resize(num_neighbors);\n cudaMemcpy(neighbors.raw(), neighbors_, num_neighbors * sizeof(VecInt_t), cudaMemcpyDefault);\n cudaCheckError();\n}\n\ntemplate \ninline void DistributedManagerBase::cacheMaps(const VecInt_t *b2l_maps, const VecInt_t *b2l_ptrs, const VecInt_t *l2h_maps, const VecInt_t *l2h_ptrs)\n{\n int num_neighbors = this->neighbors.size();\n this->cached_B2L_maps.resize(num_neighbors);\n this->cached_L2H_maps.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = b2l_ptrs[i + 1] - b2l_ptrs[i];\n this->cached_B2L_maps[i].resize(size);\n int count = 0;\n\n for (int j = b2l_ptrs[i]; j < b2l_ptrs[i + 1]; j++)\n {\n this->cached_B2L_maps[i][count] = b2l_maps[j];\n count++;\n }\n\n size = l2h_ptrs[i + 1] - l2h_ptrs[i];\n this->cached_L2H_maps[i].resize(size);\n count = 0;\n\n for (int j = l2h_ptrs[i]; j < l2h_ptrs[i + 1]; j++)\n {\n this->cached_L2H_maps[i][count] = l2h_maps[j];\n count++;\n }\n }\n}\n\ntemplate \ninline void DistributedManagerBase::cacheMapsOneRing()\n{\n int num_neighbors = this->neighbors.size();\n this->cached_B2L_maps.resize(num_neighbors);\n this->cached_L2H_maps.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n this->cached_B2L_maps[i] = this->B2L_maps[i];\n this->cached_L2H_maps[i] = this->L2H_maps[i];\n }\n}\n\ntemplate \ninline void DistributedManagerBase::cacheMapsOneRing(const VecInt_t **b2l_maps, const VecInt_t *b2l_sizes, const VecInt_t **l2h_maps, const VecInt_t *l2h_sizes)\n{\n int num_neighbors = this->neighbors.size();\n this->cached_B2L_maps.resize(num_neighbors);\n this->cached_L2H_maps.resize(num_neighbors);\n // buffering in the case of GPU data. This shouldn't much affect performance\n std::vector b2l_buffer, l2h_buffer;\n std::vector b2l_sizes_buffer, l2h_sizes_buffer;\n b2l_buffer.resize(num_neighbors);\n l2h_buffer.resize(num_neighbors);\n b2l_sizes_buffer.resize(num_neighbors);\n l2h_sizes_buffer.resize(num_neighbors);\n cudaError_t cuda_rc = cudaMemcpy(&(b2l_sizes_buffer[0]), b2l_sizes, sizeof(VecInt_t) * num_neighbors, cudaMemcpyDefault);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy b2l_sizes failed in distributed_manager\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = cudaMemcpy(&(l2h_sizes_buffer[0]), l2h_sizes, sizeof(VecInt_t) * num_neighbors, cudaMemcpyDefault);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy l2h_sizes failed in distributed_manager\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = cudaMemcpy(&(b2l_buffer[0]), b2l_maps, sizeof(VecInt_t *) * num_neighbors, cudaMemcpyDefault);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy b2l_maps failed in distributed_manager\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = cudaMemcpy(&(l2h_buffer[0]), l2h_maps, sizeof(VecInt_t *) * num_neighbors, cudaMemcpyDefault);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy l2h_maps failed in distributed_manager\", AMGX_ERR_CUDA_FAILURE);\n }\n\n // caching all of the maps\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = b2l_sizes_buffer[i];\n this->cached_B2L_maps[i].resize(size);\n cuda_rc = cudaMemcpy(&(this->cached_B2L_maps[i][0]), b2l_buffer[i], sizeof(VecInt_t) * size, cudaMemcpyDefault);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy cached_B2L_maps failed in distributed_manager\", AMGX_ERR_CUDA_FAILURE);\n }\n cudaCheckError();\n size = l2h_sizes_buffer[i];\n this->cached_L2H_maps[i].resize(size);\n cudaMemcpy(&(this->cached_L2H_maps[i][0]), l2h_buffer[i], sizeof(VecInt_t) * size, cudaMemcpyDefault);\n cudaCheckError();\n }\n}\n\ntemplate \nvoid DistributedManagerBase::setAConsolidationFlags( Matrix &in_A)\n{\n this->A = &in_A;\n AMG_Config *rsrc_cfg = this->A->getResources()->getResourcesConfig();\n std::string scope;\n int consolidate_flag, cuda_ipc_flag;\n rsrc_cfg->getParameter(\"fine_level_consolidation\", consolidate_flag, \"default\", scope);\n rsrc_cfg->getParameter(\"use_cuda_ipc_consolidation\", cuda_ipc_flag, \"default\", scope);\n this->m_is_fine_level_consolidated = (consolidate_flag != 0);\n this->m_use_cuda_ipc_consolidation = (cuda_ipc_flag != 0);\n}\n\ntemplate \nvoid DistributedManagerBase::uploadMatrix(int n, int nnz, int block_dimx, int block_dimy, const int *row_ptrs, const int *col_indices, const void *data, const void *diag, Matrix &in_A)\n{\n this->setAConsolidationFlags(in_A);\n\n if (this->m_is_fine_level_consolidated)\n {\n this->A->manager->consolidateAndUploadAll(n, nnz, block_dimx, block_dimy, row_ptrs, col_indices, data, diag, *(this->A));\n }\n else\n {\n this->A->manager->initializeUploadReorderAll(n, nnz, block_dimx, block_dimy, row_ptrs, col_indices, data, diag, *(this->A));\n }\n}\n\ntemplate \nvoid DistributedManagerBase::checkPinnedBuffer(size_t size)\n{\n if ((m_pinned_buffer_size < size) && (m_pinned_buffer != NULL))\n {\n amgx::memory::cudaFreeHost(m_pinned_buffer);\n m_pinned_buffer = NULL;\n m_pinned_buffer_size = 0;\n }\n\n if (m_pinned_buffer == NULL)\n {\n m_pinned_buffer_size = (size_t)(size * 1.1);\n amgx::memory::cudaMallocHost(&m_pinned_buffer, m_pinned_buffer_size);\n }\n}\n\n\ntemplate \nDistributedManagerBase::~DistributedManagerBase()\n{\n if (m_pinned_buffer != NULL)\n {\n amgx::memory::cudaFreeHost(m_pinned_buffer);\n }\n\n destroyComms();\n // from childrens:\n cudaStreamDestroy(this->m_int_stream);\n cudaCheckError();\n cudaStreamDestroy(this->m_bdy_stream);\n cudaCheckError();\n\n if (!this->halo_rows_ref_count && this->halo_rows != NULL)\n {\n delete this->halo_rows;\n this->halo_rows = NULL;\n }\n\n if (!this->halo_btl_ref_count && this->halo_btl != NULL)\n {\n delete this->halo_btl;\n this->halo_btl = NULL;\n }\n}\n\n// if pointer is host pointer - returns data. If it is device pointer - copies it to the m_pinned_buffer and returns pointer to m_pinned_buffer\ntemplate \nvoid *DistributedManagerBase::getHostPointerForData(void *ptr, size_t size, int *allocated)\n{\n cudaError_t rc;\n cudaPointerAttributes att;\n void *ptr_h;\n cudaCheckError();\n /* WARNING: We may accept the following types of allocation for ptr:\n 1. malloc [host memory]\n 2. cudaMalloc [device memory]\n 3. malloc + cudaHostRegister [AMGX_pin_memory/AMGX_unpin_memory host memory]\n 4. cudaHostAlloc [pinned host memory form the beginning]\n\n The correct way to conver these cases is the following:\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2, 3 or 4.\n }\n else{\n //you are in case 1.\n }\n The following pattern of checks should be implemented\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2 or 4.\n }\n else{\n st = cudaHostGetDevicePointer(ptr_on_device, ptr, 0);\n if (st == cudaSuccess){\n //you are in case 3.\n }\n else{\n //you are in case 1.\n }\n }\n The above pattern will be used whenever we need to process input data.\n\n Obs.: parameter size is in bytes and\n parameter allocated indicates whether memory was allocated\n and needs to be release later on. */\n /*\n // original implementation\n cudaPointerGetAttributes(&att, ptr);\n if (att.hostPointer == NULL)\n {\n checkPinnedBuffer(size);\n cudaMemcpy(m_pinned_buffer, ptr, size, cudaMemcpyDefault);\n cudaCheckError();\n return m_pinned_buffer;\n }\n else {\n return ptr;\n }\n */\n *allocated = 0;\n // get pointer to values on the device\n rc = cudaPointerGetAttributes(&att, ptr);\n\n#if CUDART_VERSION >= 11000\n if(rc == cudaSuccess && (att.type == cudaMemoryTypeDevice || att.type == cudaMemoryTypeManaged))\n {\n checkPinnedBuffer(size);\n rc = cudaMemcpy(m_pinned_buffer, ptr, size, cudaMemcpyDefault);\n cudaCheckError();\n\n if (rc != cudaSuccess)\n {\n FatalError(\"Could not copy into the temporary (host) storage. Try pinning the memory to avoid the cudaMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n ptr_h = m_pinned_buffer;\n *allocated = 1;\n cudaGetLastError(); //to reset last error\n return ptr_h;\n }\n else\n {\n cudaGetLastError(); //to reset last error\n return ptr;\n }\n#else\n if (rc == cudaSuccess)\n {\n //you are in case 2 or 4 from the above comment.\n if (att.hostPointer == NULL)\n {\n //you are in case 2\n checkPinnedBuffer(size);\n rc = cudaMemcpy(m_pinned_buffer, ptr, size, cudaMemcpyDefault);\n cudaCheckError();\n\n if (rc != cudaSuccess)\n {\n FatalError(\"Could not copy into the temporary (host) storage. Try pinning the memory to avoid the cudaMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n ptr_h = m_pinned_buffer;\n *allocated = 1;\n }\n else\n {\n //you are in case 4\n ptr_h = ptr;\n }\n }\n else\n {\n //you are in case 1 or 3 from the above comment\n ptr_h = ptr;\n }\n\n cudaGetLastError(); //to reset last error\n\n /* check for null pointers */\n if (ptr_h == NULL)\n {\n FatalError(\"Result of (host) allocation of required temporary storage is NULL. Try pinning the memory to reduce storage requirements.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n return ptr_h;\n#endif\n}\n\n// if pointer is host pointer - returns data. If it is device pointer - copies it to the m_pinned_buffer and returns pointer to m_pinned_buffer\ntemplate \nconst void *DistributedManagerBase::getHostPointerForData(const void *ptr, size_t size, int *allocated)\n{\n cudaError_t rc;\n cudaPointerAttributes att;\n void *ptr_h;\n cudaCheckError();\n /* WARNING: We may accept the following types of allocation for ptr:\n 1. malloc [host memory]\n 2. cudaMalloc [device memory]\n 3. malloc + cudaHostRegister [AMGX_pin_memory/AMGX_unpin_memory host memory]\n 4. cudaHostAlloc [pinned host memory form the beginning]\n\n The correct way to conver these cases is the following:\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2, 3 or 4.\n }\n else{\n //you are in case 1.\n }\n The following pattern of checks should be implemented\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2 or 4.\n }\n else{\n st = cudaHostGetDevicePointer(ptr_on_device, ptr, 0);\n if (st == cudaSuccess){\n //you are in case 3.\n }\n else{\n //you are in case 1.\n }\n }\n The above pattern will be used whenever we need to process input data.\n\n Obs.: parameter size is in bytes and\n parameter allocated indicates whether memory was allocated\n and needs to be release later on. */\n *allocated = 0;\n // get pointer to values on the device\n rc = cudaPointerGetAttributes(&att, ptr);\n\n#if CUDART_VERSION >= 11000\n if(rc == cudaSuccess && (att.type == cudaMemoryTypeDevice || att.type == cudaMemoryTypeManaged))\n {\n checkPinnedBuffer(size);\n rc = cudaMemcpy(m_pinned_buffer, ptr, size, cudaMemcpyDefault);\n cudaCheckError();\n\n if (rc != cudaSuccess)\n {\n FatalError(\"Could not copy into the temporary (host) storage. Try pinning the memory to avoid the cudaMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n ptr_h = m_pinned_buffer;\n *allocated = 1;\n cudaGetLastError(); //to reset last error\n return ptr_h;\n }\n else\n {\n cudaGetLastError(); //to reset last error\n return ptr;\n }\n#else\n if (rc == cudaSuccess)\n {\n //you are in case 2 or 4 from the above comment.\n if (att.hostPointer == NULL)\n {\n //you are in case 2\n checkPinnedBuffer(size);\n rc = cudaMemcpy(m_pinned_buffer, ptr, size, cudaMemcpyDefault);\n cudaCheckError();\n\n if (rc != cudaSuccess)\n {\n FatalError(\"Could not copy into the temporary (host) storage. Try pinning the memory to avoid the cudaMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n ptr_h = m_pinned_buffer;\n *allocated = 1;\n cudaGetLastError(); //to reset last error\n return ptr_h;\n }\n else\n {\n //you are in case 4\n cudaGetLastError(); //to reset last error\n return ptr;\n }\n }\n else\n {\n cudaGetLastError(); //to reset last error\n //you are in case 1 or 3 from the above comment\n return ptr;\n }\n#endif\n}\n\n\ntemplate \nvoid *DistributedManagerBase::getDevicePointerForData(void *ptr, size_t size, int *allocated)\n{\n cudaError_t rc;\n cudaPointerAttributes att;\n void *ptr_d;\n cudaCheckError();\n /* WARNING: We may accept the following types of allocation for ptr:\n 1. malloc [host memory]\n 2. cudaMalloc [device memory]\n 3. malloc + cudaHostRegister [AMGX_pin_memory/AMGX_unpin_memory host memory]\n 4. cudaHostAlloc [pinned host memory form the beginning]\n\n The correct way to conver these cases is the following:\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2, 3 or 4.\n }\n else{\n //you are in case 1.\n }\n The following pattern of checks should be implemented\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2 or 4.\n }\n else{\n st = cudaHostGetDevicePointer(ptr_on_device, ptr, 0);\n if (st == cudaSuccess){\n //you are in case 3.\n }\n else{\n //you are in case 1.\n }\n }\n The above pattern will be used whenever we need to process input data.\n\n Obs.: parameter size is in bytes and\n parameter allocated indicates whether memory was allocated\n and needs to be release later on. */\n *allocated = 0;\n // get pointer to values on the device\n rc = cudaPointerGetAttributes(&att, ptr);\n\n#if CUDART_VERSION >= 11000\n if(rc == cudaSuccess && att.type != cudaMemoryTypeUnregistered)\n {\n // Received device registered data (device, managed, malloc+cudaHostRegister, or cudaMallocHost)\n ptr_d = (void*)att.devicePointer;\n\n if (ptr_d == NULL)\n {\n FatalError(\"Data passed is registered with device but devicePointer null.\", AMGX_ERR_INTERNAL);\n }\n }\n else\n {\n // Received unregistered or host allocated data (malloc)\n rc = amgx::memory::cudaMallocAsync(&ptr_d, size);\n\n if (rc != cudaSuccess)\n {\n FatalError(\"Could not allocate required temporary storage. Try pinning the memory to reduce storage requirements.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n rc = cudaMemcpy(ptr_d, ptr, size, cudaMemcpyDefault);\n cudaCheckError();\n\n if (rc != cudaSuccess)\n {\n FatalError(\"Could not copy into the temporary storage. Try pinning the memory to avoid the cudaMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n *allocated = 1;\n }\n#else\n\n if (rc == cudaSuccess)\n {\n //you are in case 2 or 4 from the above comment.\n ptr_d = (void *)att.devicePointer;\n }\n else\n {\n //you are in case 1 or 3 from the above comment\n rc = cudaHostGetDevicePointer(&ptr_d, ptr, 0);\n\n if (rc != cudaSuccess)\n {\n //you are in case 1\n rc = amgx::memory::cudaMallocAsync(&ptr_d, size);\n\n if (rc != cudaSuccess)\n {\n FatalError(\"Could not allocate required temporary storage. Try pinning the memory to reduce storage requirements.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n rc = cudaMemcpy(ptr_d, ptr, size, cudaMemcpyDefault);\n cudaCheckError();\n\n if (rc != cudaSuccess)\n {\n FatalError(\"Could not copy into the temporary storage. Try pinning the memory to avoid the cudaMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n *allocated = 1;\n }\n }\n\n /* check for null pointers */\n if (ptr_d == NULL)\n {\n FatalError(\"Result of allocation of required temporary storage is NULL. Try pinning the memory to reduce storage requirements.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n#endif\n\n cudaGetLastError(); //to reset last error\n return ptr_d;\n}\n\ntemplate \nconst void *DistributedManagerBase::getDevicePointerForData(const void *ptr, size_t size, int *allocated)\n{\n cudaError_t rc;\n cudaPointerAttributes att;\n void *ptr_d;\n cudaCheckError();\n /* WARNING: We may accept the following types of allocation for ptr:\n 1. malloc [host memory]\n 2. cudaMalloc [device memory]\n 3. malloc + cudaHostRegister [AMGX_pin_memory/AMGX_unpin_memory host memory]\n 4. cudaHostAlloc [pinned host memory form the beginning]\n\n The correct way to conver these cases is the following:\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2, 3 or 4.\n }\n else{\n //you are in case 1.\n }\n The following pattern of checks should be implemented\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2 or 4.\n }\n else{\n st = cudaHostGetDevicePointer(ptr_on_device, ptr, 0);\n if (st == cudaSuccess){\n //you are in case 3.\n }\n else{\n //you are in case 1.\n }\n }\n The above pattern will be used whenever we need to process input data.\n\n Obs.: parameter size is in bytes and\n parameter allocated indicates whether memory was allocated\n and needs to be release later on. */\n *allocated = 0;\n // get pointer to values on the device\n rc = cudaPointerGetAttributes(&att, ptr);\n\n#if CUDART_VERSION >= 11000\n if(rc == cudaSuccess && att.type != cudaMemoryTypeUnregistered)\n {\n // Received device registered data (device, managed, malloc+cudaHostRegister, or cudaMallocHost)\n ptr_d = (void*)att.devicePointer;\n\n if (ptr_d == NULL)\n {\n FatalError(\"Data passed is registered with device but devicePointer null.\", AMGX_ERR_INTERNAL);\n }\n }\n else\n {\n // Received unregistered or host allocated data (malloc)\n rc = amgx::memory::cudaMallocAsync(&ptr_d, size);\n\n if (rc != cudaSuccess)\n {\n FatalError(\"Could not allocate required temporary storage. Try pinning the memory to reduce storage requirements.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n rc = cudaMemcpy(ptr_d, ptr, size, cudaMemcpyDefault);\n cudaCheckError();\n\n if (rc != cudaSuccess)\n {\n FatalError(\"Could not copy into the temporary storage. Try pinning the memory to avoid the cudaMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n *allocated = 1;\n }\n#else\n if (rc == cudaSuccess)\n {\n //you are in case 2 or 4 from the above comment.\n cudaGetLastError(); //to reset last error\n return (const void *)att.devicePointer;\n }\n else\n {\n //you are in case 1 or 3 from the above comment\n rc = cudaHostGetDevicePointer(&ptr_d, (void *)ptr, 0);\n\n if (rc != cudaSuccess)\n {\n //you are in case 1\n rc = amgx::memory::cudaMallocAsync(&ptr_d, size);\n\n if (rc != cudaSuccess)\n {\n FatalError(\"Could not allocate required temporary storage. Try pinning the memory to reduce storage requirements.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n rc = cudaMemcpy(ptr_d, ptr, size, cudaMemcpyDefault);\n cudaCheckError();\n\n if (rc != cudaSuccess)\n {\n FatalError(\"Could not copy into the temporary storage. Try pinning the memory to avoid the cudaMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n *allocated = 1;\n cudaGetLastError(); //to reset last error\n return (const void *)ptr_d;\n }\n }\n\n /* check for null pointers */\n if (ptr_d == NULL)\n {\n FatalError(\"Result of allocation of required temporary storage is NULL. Try pinning the memory to reduce storage requirements.\", AMGX_ERR_BAD_PARAMETERS);\n }\n#endif\n\n // shouldn't get there\n cudaGetLastError(); //to reset last error\n return NULL;\n}\n\ntemplate \nvoid initializeMatrixCopyAll(int n, int nnz, int block_dimx, int block_dimy, const int *row_ptrs, const int *col_indices, const void *data, const void *diag, Matrix *A)\n{\n typedef typename TConfig::MatPrec mat_value_type;\n A->resize( n, n, nnz, block_dimx, block_dimy );\n //Upload the entire matrix\n cudaMemcpy( A->row_offsets.raw(), row_ptrs, (n + 1) * sizeof(int), cudaMemcpyDefault );\n cudaCheckError();\n cudaMemcpy( A->col_indices.raw(), col_indices, (nnz) * sizeof(int), cudaMemcpyDefault );\n cudaCheckError();\n cudaMemcpy( A->values.raw(), (mat_value_type *)data, (nnz * block_dimx * block_dimy) * sizeof(mat_value_type), cudaMemcpyDefault );\n cudaCheckError();\n\n if (diag)\n {\n cudaMemcpy( A->values.raw() + A->diagOffset()*A->get_block_size(), (mat_value_type *)diag, (n * block_dimx * block_dimy) * sizeof(mat_value_type), cudaMemcpyDefault );\n cudaCheckError();\n }\n else\n {\n A->computeDiagonal();\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManagerBase::updateMapsReorder()\n{\n int my_id = this->getComms()->get_global_id();\n DistributedComms *comms_tmp = this->getComms();\n DistributedComms **comms_ = &comms_tmp;\n // Copy B2L_maps in their final place\n int num_neighbors = this->neighbors.size();\n B2L_maps.resize(num_neighbors);\n L2H_maps.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n B2L_maps[i] = this->cached_B2L_maps[i];\n L2H_maps[i] = this->cached_L2H_maps[i];\n }\n\n //Create a DistributedArranger object to map further halo rings and to construct halo row matrices and exchange them (if halo_coloring != LAST)\n DistributedArranger *prep = new DistributedArranger;\n prep->create_B2L_from_maps( (*(this->A)), my_id, this->num_halo_rings(), neighbors,\n B2L_maps, L2H_maps, B2L_rings, comms_, &halo_rows, &halo_btl);\n DistributedManagerBaseInit(my_id, 0, this->A->get_num_rows(), *(this->A), comms_, NULL, NULL);\n //Use the exchanged halo row matrices and the boundary/halo index lists to renumber and consolidate the matrix\n this->reorder_matrix();\n prep->initialize_B2L_maps_offsets(*(this->A), this->num_halo_rings());\n delete prep;\n}\n\ntemplate \nvoid DistributedManagerBase::initializeUploadReorderAll(int n, int nnz, int block_dimx, int block_dimy, const int *row_ptrs, const int *col_indices, const void *data, const void *diag, Matrix &in_A)\n{\n this->A = &in_A;\n initializeMatrixCopyAll(n, nnz, block_dimx, block_dimy, row_ptrs, col_indices, data, diag, this->A);\n this->updateMapsReorder();\n}\n\ntemplate \nvoid DistributedManagerBase::destroyComms()\n{\n if ( (this->_comms) != NULL )\n {\n if (this->_comms->decr_ref_count())\n {\n delete (this->_comms);\n this->_comms = NULL;\n }\n }\n\n if ( (this->m_fine_level_comms) != NULL)\n {\n if (this->m_fine_level_comms->decr_ref_count())\n {\n delete (this->m_fine_level_comms);\n this->m_fine_level_comms = NULL;\n }\n }\n}\n\ntemplate \nvoid DistributedManagerBase::initComms(Resources *rsrc)\n{\n this->createComms(rsrc);\n int my_id = this->getComms()->get_global_id();\n int num_parts = this->getComms()->get_num_partitions();\n this->set_global_id(my_id);\n this->set_num_partitions(num_parts);\n}\n\n\ntemplate \nvoid DistributedManagerBase::createComms(Resources *rsrc)\n{\n // create communicator\n#ifdef AMGX_WITH_MPI\n destroyComms();\n if (rsrc == NULL)\n FatalError(\"Resources should not be NULL\", AMGX_ERR_INTERNAL);\n\n MPI_Comm *mpi_comm = rsrc->getMpiComm();\n AMG_Config *cfg = rsrc->getResourcesConfig();\n std::string comm_value, comm_scope;\n cfg->getParameter(\"communicator\", comm_value, \"default\", comm_scope);\n int rank = -1;\n MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n\n if (comm_value == \"MPI_DIRECT\")\n {\n _comms = new CommsMPIDirect(*cfg, comm_scope, mpi_comm);\n std::string comm_log(\"Using CUDA-Aware MPI (GPU Direct) communicator...\\n\");\n amgx_distributed_output(comm_log.c_str(), comm_log.length());\n }\n else if (comm_value == \"MPI\")\n {\n _comms = new CommsMPIHostBufferStream(*cfg, comm_scope, mpi_comm);\n std::string comm_log(\"Using Normal MPI (Hostbuffer) communicator...\\n\");\n amgx_distributed_output(comm_log.c_str(), comm_log.length());\n }\n else \n { \n FatalError(\"Bad communicator value\", AMGX_ERR_BAD_PARAMETERS); \n }\n\n#endif\n}\n\ntemplate \nvoid DistributedManagerBase::malloc_export_maps(VecInt_t ***b2l_maps_e, VecInt_t **b2l_maps_sizes_e, VecInt_t ***l2h_maps_e, VecInt_t **l2h_maps_sizes_e)\n{\n *b2l_maps_e = (VecInt_t **) malloc(sizeof(VecInt_t *)*this->num_neighbors());\n *l2h_maps_e = (VecInt_t **) malloc(sizeof(VecInt_t *)*this->num_neighbors());\n *b2l_maps_sizes_e = (VecInt_t *) malloc(sizeof(VecInt_t) * (this->num_neighbors()));\n *l2h_maps_sizes_e = (VecInt_t *) malloc(sizeof(VecInt_t) * (this->num_neighbors()));\n\n for (int i = 0; i < this->num_neighbors(); i++)\n {\n (*b2l_maps_sizes_e)[i] = B2L_maps[i].size();\n (*l2h_maps_sizes_e)[i] = L2H_maps[i].size();\n (*b2l_maps_e)[i] = (VecInt_t *) malloc(sizeof(VecInt_t) * ( (*b2l_maps_sizes_e)[i]) );\n\n if (L2H_maps[i].size() != 0)\n {\n (*l2h_maps_e)[i] = (VecInt_t *) malloc(sizeof(VecInt_t) * ( (*l2h_maps_sizes_e)[i]) );\n amgx::thrust::copy(L2H_maps[i].begin(), L2H_maps[i].end(), (*l2h_maps_e)[i]);\n }\n\n amgx::thrust::copy(B2L_maps[i].begin(), B2L_maps[i].end(), (*b2l_maps_e)[i]);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManager >::createRenumbering(IVector &renumbering)\n{\n int num_neighbors = this->neighbors.size();\n // still renumber if the number of neighbors = 0, to support non-symmetric matrices\n // if (num_neighbors == 0) return;\n /*\n EXAMPLE\n Example matrix, partition 1 arrives with state:\n A.row_offsets = [0 4 11 15 20]\n A.col_indices = [4 0 1 2\n 4 5 0 1 2 3 7\n 0 1 2 3\n 1 2 3 6 7]\n num_neighbors=2; neighbors = [0 2]\n B2L_rings[[0 2 4][0 2 4]] B2L_maps[[0 1| 2 3][1 3| 0 2]]\n L2H_maps (and halo_lists) [[4 5][6 7]]\n */\n int size = 0;\n\n if (this->L2H_maps.size())\n {\n size = thrust_wrapper::reduce(this->A->col_indices.begin(), this->A->col_indices.end(), int(0), amgx::thrust::maximum()) + 1; //Sufficient to do reduction on lth maps\n cudaCheckError();\n }\n else\n {\n size = this->A->get_num_rows();\n }\n\n int rings = (this->B2L_rings.size() > 0) ? this->B2L_rings[0].size() - 1 : 0;\n //initial size to size+1 so we have the total size after a scan\n renumbering.resize(size + 1);\n int global_size = size;\n //\n // Step 1 - in the main matrix, separate interior and boundary nodes (1/0 in flagArray), renumber interior ones with an exclusive scan\n //\n IVector flagArray(size + 1);\n thrust_wrapper::fill(flagArray.begin(), flagArray.begin() + size + 1, 1);\n cudaCheckError();\n\n //sets 1 for interior nodes, 0 for boundary node\n for (int i = 0; i < num_neighbors; i++ )\n {\n int size = this->B2L_rings[i][1];\n int num_blocks = std::min(4096, (size + 127) / 128);\n\n if (size > 0)\n {\n remove_boundary_kernel <<< num_blocks, 128>>>(flagArray.raw(), this->B2L_maps[i].raw(), size);\n cudaCheckError();\n }\n\n //If there are any L2H maps\n if (this->L2H_maps.size() && this->L2H_maps[i].size())\n {\n int size = this->L2H_maps[i].size();\n int num_blocks = std::min(4096, (size + 127) / 128);\n remove_boundary_kernel <<< num_blocks, 128>>>(flagArray.raw(), this->L2H_maps[i].raw(), size);\n }\n\n cudaCheckError();\n }\n\n //gets the renumbering of interior nodes\n thrust_wrapper::exclusive_scan(flagArray.begin(), flagArray.begin() + size + 1, renumbering.begin());\n cudaCheckError();\n /*\n EXAMPLE\n After removing 1-ring boundary nodes and halo nodes from flagArray: [0 0 1 0 0 0 0 0]\n After exclusive scan, which gives renumbering for interior nodes (only node #2)\n renumbering: [0 0 0 1 1 1 1 1]\n */\n //\n // Step 2 - Renumber nodes that are in the boundary, stepping through each B2L map, and renumbering ones that have not been renumbered yet\n //\n //what is the biggest B2L size\n INDEX_TYPE max_size = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n max_size = max_size > this->B2L_rings[i][1] ? max_size : this->B2L_rings[i][1];\n\n if (this->L2H_maps.size())\n {\n max_size = max_size > this->L2H_maps[i].size() ? max_size : this->L2H_maps[i].size();\n }\n }\n\n //allocate work vectors (should be pretty small) that are used to renumber boundary nodes\n IVector boundary_renum_flags(max_size);\n IVector boundary_renum(max_size);\n //the number of renumbered nodes so far\n int max_element = renumbering[size];\n this->_num_interior_nodes = max_element;\n this->_num_boundary_nodes = this->A->get_num_rows() - max_element;\n renumbering.resize(size);\n\n /*\n EXAMPLE\n size = 8\n max_size = 2, max_element = 1, num_interior_nodes=1, num_boundary_nodes = 4-1 = 3\n */\n\n for (int i = 0; i < num_neighbors; i++)\n {\n //find nodes that are part of the current boundary and they haven't been renumbered yet\n thrust_wrapper::fill(boundary_renum_flags.begin(), boundary_renum_flags.begin() + max_size, 0);\n int size = this->B2L_rings[i][1];\n int num_blocks = std::min(4096, (size + 191) / 192);\n\n if (size > 0)\n get_unassigned_kernel <<< num_blocks, 192>>>(flagArray.raw(),\n this->B2L_maps[i].raw(),\n boundary_renum_flags.raw(), size, global_size /*,rank*/);\n cudaCheckError();\n\n //calculate the local renumbering (within this boundary region) of these nodes\n thrust_wrapper::exclusive_scan(boundary_renum_flags.begin(), boundary_renum_flags.begin() + max_size, boundary_renum.begin());\n\n //apply renumbering to the big numbering table\n if (size > 0)\n set_unassigned_kernel <<< num_blocks, 192>>>(boundary_renum_flags.raw(),\n boundary_renum.raw(),\n this->B2L_maps[i].raw(),\n renumbering.raw(),\n size, max_element, global_size /*,rank*/);\n cudaCheckError();\n\n //update the number of renumbered nodes\n max_element += boundary_renum[max_size - 1] + boundary_renum_flags[max_size - 1];\n /*\n EXAMPLE\n for neighbor 0 (ID 0)\n boundary_renum_flags = [0 0], size = 2, flagArray [0 0 1 0 0 0 0 0]\n get_unassigned_kernel's output: boundary_renum_flags = [1 1] flagArray [1 1 1 0 0 0 0 0]\n after exclusive scan: boundary_renum [0 1]\n set_unassigned_kernel updates these arrays and renumbers B2L map:\n renumbering = [1 2 0 1 1 1 1 1] B2L_maps[0] = [1 2| 2 3] (note that after element 3 in renumbering and after element 2 we have invalid/not yet updated values)\n max_element = 3\n\n for neighbor 1 (ID 2)\n get_unassigned_kernels's output: boundary_renum_flags [0 1] flagArray [1 1 1 1 0 0 0 0]\n after exclusive scan boundary_renum [0 0]\n set_unassigned_kernel renumbering [1 2 0 3 1 1 1 1] B2L_maps[1] = [2 3| 0 2]\n max_element = 4\n */\n }\n\n cudaCheckError();\n\n //Get renumbering for halo indices\n if (this->L2H_maps.size())\n {\n //TODO: simplify this, we don't need to check whether it has already been renumbered, there is no overlap between halos\n for (int i = 0; i < num_neighbors; i++)\n {\n //find nodes that are part of the current boundary and they haven't been renumbered yet\n thrust_wrapper::fill(boundary_renum_flags.begin(), boundary_renum_flags.begin() + max_size, 0);\n int size = this->L2H_maps[i].size();\n int num_blocks = std::min(4096, (size + 191) / 192);\n\n if (size > 0)\n get_unassigned_kernel <<< num_blocks, 192>>>(flagArray.raw(),\n this->L2H_maps[i].raw(),\n boundary_renum_flags.raw(), size, global_size /*,rank*/);\n cudaCheckError();\n\n //calculate the local renumbering (within this boundary region) of these nodes\n thrust_wrapper::exclusive_scan(boundary_renum_flags.begin(), boundary_renum_flags.begin() + max_size, boundary_renum.begin());\n\n //apply renumbering to the big numbering table\n if (size > 0)\n set_unassigned_kernel <<< num_blocks, 192>>>(boundary_renum_flags.raw(),\n boundary_renum.raw(),\n this->L2H_maps[i].raw(),\n renumbering.raw(),\n size, max_element, global_size /*,rank*/);\n cudaCheckError();\n\n //update the number of renumbered nodes\n max_element += boundary_renum[max_size - 1] + boundary_renum_flags[max_size - 1];\n /*\n EXAMPLE\n for neighbor 0 (ID 0)\n boundary_renum_flags = [0 0], size = 2, flagArray [1 1 1 1 0 0 0 0]\n get_unassigned_kernel's output: boundary_renum_flags = [1 1] flagArray [1 1 1 1 1 1 0 0]\n after exclusive scan: boundary_renum [0 1]\n set_unassigned_kernel updates these arrays and renumbers B2L map:\n renumbering = [1 2 0 3 4 5 1 1] L2H_maps[0] = [4 5]\n max_element = 6\n\n for neighbor 1 (ID 2)\n get_unassigned_kernels's output: boundary_renum_flags [1 1] flagArray [1 1 1 1 1 1 1 1]\n after exclusive scan boundary_renum [0 1]\n set_unassigned_kernel renumbering = [1 2 0 3 4 5 6 7] L2H_maps[1] = [6 7]\n max_element = 8\n */\n }\n\n cudaCheckError();\n }\n\n //apply renumbering to further halo rings too\n if (rings > 1)\n {\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = this->B2L_rings[i][this->B2L_rings[i].size() - 1] - this->B2L_rings[i][1];\n int num_blocks = std::min(4096, (size + 127) / 128);\n renumber_b2l_maps <<< num_blocks, 128>>>(this->B2L_maps[i].raw() + this->B2L_rings[i][1], renumbering.raw(), size, global_size /*, rank*/);\n }\n\n cudaCheckError();\n }\n\n /*\n EXAMPLE\n renumbers further boundary rings as listed in B2L_maps, since they have not been replaced yet with their renumbered values\n B2L_maps [[1 2| 0 3][2 3| 1 0]]\n */\n}\n\ntemplate \nvoid DistributedManager >::reorder_matrix_owned()\n{\n int num_neighbors = this->neighbors.size();\n int size = this->A->get_num_rows();\n int num_blocks = std::min(4096, (size + 511) / 512);\n int rings = (this->B2L_rings.size() > 0) ? this->B2L_rings[0].size() - 1 : 0;\n this->set_num_halo_rings(rings);\n int diag = this->A->hasProps(DIAG);\n\n if (diag)\n {\n FatalError(\"External diag not supported in classical path\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n//\n// Step 1 & 2 - create renumbering\n//\n this->createRenumbering(this->renumbering);\n //now we have the full renumbering table in renum, calculate the inverse\n this->inverse_renumbering.resize(this->renumbering.size());\n\n if (this->renumbering.size() > 1)\n {\n calc_inverse_renumbering <<< std::min(4096, ((int)this->renumbering.size() + 511) / 512), 512 >>> (this->renumbering.raw(), this->inverse_renumbering.raw(), this->renumbering.size());\n cudaCheckError();\n }\n\n//\n// Step 4 - calculate number/offset of nodes in the halos from the neighbors, ring by ring\n//\n this->halo_offsets.resize(num_neighbors + 1);\n this->halo_offsets[0] = size;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n this->halo_offsets[i + 1] = this->halo_offsets[i] + this->L2H_maps[i].size();\n }\n\n this->set_num_halo_rows(this->halo_offsets[this->halo_offsets.size() - 1] - size);\n int nh = this->num_halo_rows();\n int total_rows = size + nh;\n cudaCheckError();\n//\n// Step 6 - renumber halo matrices and calculate row length (to eventually append to the big matrix)\n//\n int insert = 0;\n //recalculate row_offsets\n IVector new_row_offsets(size + 1);\n\n if (num_blocks > 0)\n {\n calc_rowlen_reorder <<< num_blocks, 512>>>(this->A->row_offsets.raw(), new_row_offsets.raw(), this->renumbering.raw(), size, insert);\n cudaCheckError();\n }\n\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(this->renumbering.begin(), this->A->col_indices.begin()),\n amgx::thrust::make_permutation_iterator(this->renumbering.begin(), this->A->col_indices.end()),\n this->A->col_indices.begin());\n cudaCheckError();\n //row_offsets array created by exclusive scan of row sizes\n thrust_wrapper::exclusive_scan(new_row_offsets.begin(), new_row_offsets.begin() + size + 1, new_row_offsets.begin());\n cudaCheckError();\n//\n// Step 7 - consolidate column indices and values\n//\n int new_nnz = new_row_offsets[new_row_offsets.size() - 1];\n typedef typename MatPrecisionMap::Type ValueTypeA;\n VVector new_values((new_nnz + 1 )* this->A->get_block_size(), types::util::get_zero());\n IVector new_col_indices(new_nnz, 0);\n\n //reorder based on row permutation\n if (num_blocks > 0)\n {\n reorder_whole_matrix <32> <<< num_blocks, 512>>>(this->A->row_offsets.raw(), this->A->col_indices.raw(), this->A->values.raw(), new_row_offsets.raw(), new_col_indices.raw(), new_values.raw(), this->renumbering.raw(), this->A->get_block_size(), size, insert);\n cudaCheckError();\n }\n\n //create and append halo rows size\n //create an identity matrix in CSR format\n int nnz = this->A->get_num_nz();\n IVector identity_csr_rows(nh + 1);\n IVector identity_csr_cols(nh);\n VVector identity_csr_vals(nh, types::util::get_one()); //needs to be changed to MVector, but this definition is messed up in the header file (should fix later)\n thrust_wrapper::sequence(identity_csr_rows.begin(), identity_csr_rows.end());\n thrust_wrapper::sequence(identity_csr_cols.begin(), identity_csr_cols.end());\n /*for example, 2x2 identity_csr matrix is created:\n identity_csr_rows = { 0, 1, 2 }\n identity_csr_cols = { 0, 1 }\n identity_csr_vals = { 1.0, 1.0 } */\n //shift identity tmatrix by size = this->A->get_num_rows();\n thrust_wrapper::transform(identity_csr_rows.begin(), identity_csr_rows.end(), amgx::thrust::constant_iterator(nnz), identity_csr_rows.begin(), amgx::thrust::plus());\n thrust_wrapper::transform(identity_csr_cols.begin(), identity_csr_cols.end(), amgx::thrust::constant_iterator(size), identity_csr_cols.begin(), amgx::thrust::plus());\n /*for example, 2x2 identity_csr matrix is created:\n identity_csr_rows = { 0, 1, 2 }\n identity_csr_cols = {size, size+1 }\n identity_csr_vals = { 1.0, 1.0 } */\n /* WARNING: you must be very careful with the view you are setting (cuurently the view coming here by default is ALL = FULL). If\n - classical path is selected then the createOneRingHaloRows -> create_one_ring_halo_rows -> append_halo_rows\n routine will be called. It will overwrite the halo rows setup here (and will use view OWNED, which will ignore the\n halo rows setup here, to determine how the new halo rows should be placed).\n - aggregation path is selected then the extra rows setup here will be used in the R*A*P product, where (in order to match\n dimensions of R and P) it is assumed that (the local partition) matrix A is square, therefore it must be padded by identity\n rows at the bottom to compensate for the \"extra\" columns that are outside of the main square part. The old routines for the\n aggregation path do this padding at the end of the reorder_matrix routine below. */\n //ViewType v = this->A->currentView();\n //this->A->setView(ALL);\n //Approach 1: use existing routine to append the identity matrix to the existing one\n // (seems like too much overhead, also need identity matrix per neighbor)\n //DistributedArranger *prep = new DistributedArranger;\n //prep->append_halo_rows(this->A, identity_csr_rows, identity_csr_cols, identity_csr_vals);\n //delete prep;\n //Approach 2: custom for this routine\n new_row_offsets.resize(total_rows + 1);\n new_col_indices.resize(nnz + nh);\n int bsize = this->A->get_block_size();\n new_values.resize((nnz + nh)*bsize + 1); //extra 1 element stores zero at the end (to follow the original design)\n //new_values[nnz]=-1; //marker to track the last element\n amgx::thrust::copy(identity_csr_rows.begin(), identity_csr_rows.end(), new_row_offsets.begin() + size );\n amgx::thrust::copy(identity_csr_cols.begin(), identity_csr_cols.end(), new_col_indices.begin() + nnz);\n amgx::thrust::copy(new_values.begin() + nnz*bsize, new_values.begin() + (nnz + 1)*bsize, new_values.begin() + (nnz + nh)*bsize);\n amgx::thrust::copy(identity_csr_vals.begin(), identity_csr_vals.end(), new_values.begin() + (nnz*bsize));\n /* WARNING: see above. */\n this->A->set_num_cols(total_rows);\n this->A->set_num_rows(total_rows);\n this->A->col_indices.swap(new_col_indices);\n new_row_offsets.resize(total_rows + 1);\n this->A->row_offsets.swap(new_row_offsets);\n new_row_offsets.swap(this->old_row_offsets);\n this->A->values.swap(new_values);\n this->A->m_seq_offsets.resize(total_rows + 1);\n thrust_wrapper::sequence(this->A->m_seq_offsets.begin(), this->A->m_seq_offsets.end());\n cudaCheckError();\n //TODO: only do this if AMG_Config matrix_halo_exchange!=2\n this->A->delProps(COO);\n if (!insert)\n {\n this->A->computeDiagonal();\n }\n\n this->set_initialized(this->A->row_offsets);\n this->A->setView(OWNED);\n}\n\n\n\ntemplate \nvoid DistributedManager >::reorder_matrix()\n{\n int num_neighbors = this->neighbors.size();\n\n if (num_neighbors == 0) { return; }\n\n int size = this->A->get_num_rows();\n int num_blocks = std::min(4096, (size + 511) / 512);\n int rings = this->B2L_rings[0].size() - 1;\n this->set_num_halo_rings(rings);\n int diag = this->A->hasProps(DIAG);\n std::vector > &halo_rows = *this->halo_rows;\n std::vector > &halo_btl = *this->halo_btl;\n /*\n EXAMPLE\n The example matrix, on partition 1 arrives at this point with the following state:\n num_rings=2\n A.num_rows = 4; A.num_nz = 20\n A.row_offsets = [0 4 11 15 20]\n A.col_indices = [4 0 1 2\n 4 5 0 1 2 3 7\n 0 1 2 3\n 1 2 3 6 7]\n num_neighbors=2; neighbors = [0 2]\n B2L_rings[[0 2 4][0 2 4]] B2L_maps[[0 1| 2 3][1 3| 0 2]]\n L2H_maps (and halo_lists) [[4 5][6 7]]\n\n With the exchange halo rows:\n halo_btl[0] (received from neighbor ID 0)\n global_id = 0; base_index=0; index_range=6; B2L_rings[0] = [0 2 4]; B2L_maps[0] = [2 3| 0 1] L2H_maps = [4 5]\n halo_rows[0].row_offsets = [0 5 13 17 21]\n halo_rows[0].col_indices = [1 2 3 4 5\n 0 1 2 3 4 5 6 7\n 0 1 3 6\n 0 1 2 3]\n\n halo_btl[1] (received from neighbor ID 2)\n global_id = 2; base_index=0; index_range=8; B2L_rings[0] = [0 2 4]; B2L_maps[0] = [1 2| 0 3] L2H_maps = [6 7]\n halo_rows[1].row_offsets = [0 4 11 16 20]\n halo_rows[1].col_indices = [7 1 2 3\n 5 6 7 0 1 2 3\n 4 5 0 2 3\n 0 1 2 3]\n */\n//\n// Step 1 & 2 - create renumbering\n//\n this->createRenumbering(this->renumbering);\n cudaCheckError();\n /*\n EXAMPLE\n this->renumbering = [1 2 0 3 4 5 6 7]\n B2L_maps = [[1 2| 0 3][2 3| 1 0]]\n L2H_maps = [[4 5][6 7]]\n */\n//\n// Step 3 - given a full renumbering of owned nodes, calculate inverse renumbering\n//\n //now we have the full renumbering table in renum, calculate the inverse\n this->inverse_renumbering.resize(this->renumbering.size());\n calc_inverse_renumbering <<< std::min(4096, ((int)this->renumbering.size() + 511) / 512), 512 >>> (this->renumbering.raw(), this->inverse_renumbering.raw(), this->renumbering.size());\n cudaCheckError();\n /*\n EXAMPLE\n this->inverse_renumbering = [2 0 1 3 4 5 6 7]\n */\n//\n// Step 4 - calculate number/offset of nodes in the halos from the neighbors, ring by ring\n//\n this->halo_offsets.resize(rings * num_neighbors + 1, 0);\n\n for (int ring = 0; ring < rings; ring++)\n {\n for (int i = 0; i < num_neighbors; i++)\n {\n this->halo_offsets[ring * num_neighbors + i] = halo_btl[i].B2L_rings[0][ring + 1] - halo_btl[i].B2L_rings[0][ring];\n }\n }\n\n thrust_wrapper::exclusive_scan(this->halo_offsets.begin(), this->halo_offsets.end(), this->halo_offsets.begin(), size);\n cudaCheckError();\n this->set_num_halo_rows(this->halo_offsets[this->halo_offsets.size() - 1] - size);\n int total_rows = size + this->num_halo_rows();\n\n if (total_rows < this->renumbering.size()) { FatalError(\"total rows < renumbering.size(), send/recv maps should cover all matrix halo columns\", AMGX_ERR_NOT_IMPLEMENTED); }\n\n if (total_rows > this->renumbering.size())\n {\n this->A->getResources()->warning(\"# owned nodes + # halo nodes > matrix columns: send/recv maps have some unreferences halo indices, they are not directly connected to our partition and therefore we won't compute them, please use 2-ring comms maps if you want to specify 2nd ring neighbors\");\n }\n\n cudaCheckError();\n /*\n EXAMPLE\n halo_offsets [2 2 2 2]\n after exclusive scan: 4 + [0 2 4 6 8] = [4 6 8 10 12]\n num_halo_rows = 8, total_rows = 12\n */\n//\n// Step 5 - create big mapping table of all halo indices we received (this may use a little too much memory)\n//\n //count number of fine rows of neighbors\n amgx::thrust::host_vector neighbor_rows(num_neighbors + 1);\n int max_num_rows = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n neighbor_rows[i] = halo_btl[i].index_range();\n max_num_rows = max_num_rows > halo_rows[i].get_num_rows() ? max_num_rows : halo_rows[i].get_num_rows();\n }\n\n thrust_wrapper::exclusive_scan(neighbor_rows.begin(), neighbor_rows.end(), neighbor_rows.begin());\n cudaCheckError();\n int total_rows_of_neighbors = neighbor_rows[num_neighbors];\n /*\n EXAMPLE\n neigbor_rows = [0 6 14]\n total_rows_of_neighbors = 14\n */\n IVector halo_mapping(total_rows_of_neighbors);\n thrust_wrapper::fill(halo_mapping.begin(), halo_mapping.end(), -1);\n cudaCheckError();\n\n //ring by ring, neighbor by neighbor assign sequentially increasing numbers for halo nodes\n for (int ring = 0; ring < rings; ring++)\n {\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = halo_btl[i].B2L_rings[0][ring + 1] - halo_btl[i].B2L_rings[0][ring];\n int num_blocks = std::min(4096, (size + 127) / 128);\n //This renumbering has to result in the same renumbering that comes out of L2H renumbering\n create_halo_mapping <<< num_blocks, 128>>>(halo_mapping.raw() + neighbor_rows[i],\n halo_btl[i].B2L_maps[0].raw() + halo_btl[i].B2L_rings[0][ring],\n halo_btl[i].base_index(), this->halo_offsets[ring * num_neighbors + i], size);\n cudaCheckError();\n /*\n EXAMPLE\n ring 0 neighbor 0 - halo_btl[0].B2L_maps[0] = [2 3| 0 1] halo_btl[0].L2H_maps = [4 5]\n halo_mapping = [-1 -1 4 5 -1 -1 |-1 -1 -1 -1 -1 -1 -1 -1]\n ring 0 neighbor 1 - halo_btl[1].B2L_maps[0] = [1 2| 0 3] halo_btl[1].L2H_maps = [6 7]\n halo_mapping = [-1 -1 4 5 -1 -1 |-1 6 7 -1 -1 -1 -1 -1]\n ring 1 neighbor 0 - halo_btl[0].B2L_maps[0] = [2 3| 0 1] halo_btl[0].L2H_maps = [4 5]\n halo_mapping = [8 9 4 5 -1 -1 |-1 6 7 -1 -1 -1 -1 -1]\n ring 1 neighbor 1 - halo_btl[1].B2L_maps[0] = [1 2| 0 3] halo_btl[1].L2H_maps = [6 7]\n halo_mapping = [8 9 4 5 -1 -1 |10 6 7 11 -1 -1 -1 -1]\n\n */\n }\n }\n\n cudaCheckError();\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = halo_btl[i].L2H_maps[0].size();\n int num_blocks = std::min(4096, (size + 127) / 128);\n //Map the column indices of the halo rows that point back to boundary nodes\n apply_h2l2b_mapping <<< num_blocks, 128>>>(halo_mapping.raw() + neighbor_rows[i],\n halo_btl[i].L2H_maps[0].raw(),\n halo_btl[i].base_index(), this->B2L_maps[i].raw(), size);\n cudaCheckError();\n /*\n EXAMPLE\n neighbor 0 - mapping back to our own (boundary) indices\n halo_mapping = [8 9 4 5 1 2 |10 6 7 11 -1 -1 -1 -1]\n neighbor 1 - mapping back to our own (boundary) indices\n halo_mapping = [8 9 4 5 1 2 |10 6 7 11 -1 -1 2 3]\n */\n }\n\n cudaCheckError();\n /*\n EXAMPLE\n neighbor_rows = [0 6 14]\n halo_mapping = [8 9 4 5 1 2 |10 6 7 11 -1 -1 2 3]\n The first part (0-6) of halo_mapping gives a local index for all the indices that we want to know about in halo_btl[0]\n The second part (7-14) gives local indices for halo_btl[1], that is both halo ring there, and the column indices representing vertices in this partition's boundary.\n Note that it does not give indices (-1) for vertices 5 and 6 in neighbor 1 (ID 2), which are column indices connecting it to neighbor 0, hence the two halo regions are not connected\n */\n//\n// Step 6 - renumber halo matrices and calculate row length (to eventually append to the big matrix)\n//\n int insert = 0;\n\n if (this->A->hasProps(DIAG) && insertDiagonals) { insert = 1; }\n\n diag = diag && !insertDiagonals;\n //recalculate row_offsets\n IVector new_row_offsets(size + this->num_halo_rows() + 1);\n calc_rowlen_reorder <<< num_blocks, 512>>>(this->A->row_offsets.raw(), new_row_offsets.raw(), this->renumbering.raw(), size, insert);\n cudaCheckError();\n IVector neighbor_rows_d(num_neighbors + 1);\n amgx::thrust::copy(neighbor_rows.begin(), neighbor_rows.end(), neighbor_rows_d.begin());\n cudaCheckError();\n /*\n EXAMPLE\n get row length according to renumbering\n new_row_offsets = [4 4 7 5 0 0 0 0 0 0 0 0 0]\n */\n //map column indices of my own matrix\n /*map_col_indices<4><<>>(this->A->row_offsets.raw(),\n this->A->col_indices.raw(),\n this->renumbering.raw(),\n this->halo_ranges.raw(),\n halo_mapping.raw(),\n neighbor_rows_d.raw(),\n this->base_index(), num_neighbors, size);*/\n cudaCheckError();\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(this->renumbering.begin(), this->A->col_indices.begin()),\n amgx::thrust::make_permutation_iterator(this->renumbering.begin(), this->A->col_indices.end()),\n this->A->col_indices.begin());\n cudaCheckError();\n /*\n EXAMPLE\n use this->renumbering = [1 2 0 3 4 5 6 7]\n to map old column indices to new column indices (i.e. according to interior - boundary - halo separation), but do not reshuffle them into their place yet\n A.col_indices = [4 0 1 2\n 4 5 0 1 2 3 7\n 0 1 2 3\n 1 2 3 6 7]\n becomes\n A.col_indices = [4 1 2 0\n 4 5 1 2 0 3 7\n 1 2 0 3\n 2 0 3 6 7]\n\n */\n cudaCheckError();\n IVector temp_row_len(max_num_rows);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n //map column indices of halo matrices and count of nonzeros we will keep\n int size = halo_rows[i].get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n map_col_indices_and_count_rowlen<4> <<< num_blocks, 128, 128 * sizeof(INDEX_TYPE)>>>(\n halo_rows[i].row_offsets.raw(),\n halo_rows[i].col_indices.raw(),\n temp_row_len.raw(),\n halo_mapping.raw() + neighbor_rows[i],\n size, insert);\n cudaCheckError();\n\n //number of nonzeros per row copied into big row sizes array\n for (int ring = 0; ring < rings; ring++)\n {\n amgx::thrust::copy(temp_row_len.begin() + halo_btl[i].B2L_rings[0][ring], temp_row_len.begin() + halo_btl[i].B2L_rings[0][ring + 1], new_row_offsets.begin() + this->halo_offsets[ring * num_neighbors + i]);\n }\n\n cudaCheckError();\n /*\n EXAMPLE\n halo_mapping = [8 9 4 5 1 2 |10 6 7 11 -1 -1 2 3]\n look at halo row matrices, and halo_mapping, count column indices that do not map to -1 and map them to their new, local index\n halo_rows[0].col_indices = [1 2 3 4 5\n 0 1 2 3 4 5 6 7\n 0 1 3 6\n 0 1 2 3]\n becomes\n halo_rows[0].col_indices = [9 4 5 1 2\n 8 9 4 5 1 2 -1 -1\n 8 9 5 -1\n 8 9 4 5]\n with temp_row_len = [5 6 3 4] copied into new_row_offsets: [4 4 7 5| 5 6| 0 0| 3 4| 0 0 0]\n\n halo_rows[1].col_indices = [7 1 2 3\n 5 6 7 0 1 2 3\n 4 5 0 2 3\n 0 1 2 3]\n becomes\n halo_rows[1].col_indices = [3 6 7 11\n -1 2 3 10 6 7 11\n -1 -1 10 7 11\n 10 6 7 11]\n with temp_row_len = [4 6 3 4] copied into new_row_offsets: [4 4 7 5| 5 6| 4 6| 3 4| 3 4 0]\n */\n }\n\n cudaCheckError();\n //row_offsets array created by exclusive scan of row sizes\n thrust_wrapper::exclusive_scan(new_row_offsets.begin(), new_row_offsets.begin() + size + this->num_halo_rows() + 1, new_row_offsets.begin());\n cudaCheckError();\n /*\n EXAMPLE\n Exclusive scan to get new_row_offsets array:\n new_row_offsets = [0 4 8 15 20| 25 31| 35 41| 44 48| 51 55]\n */\n//\n// Step 7 - consolidate column indices and values\n//\n int new_nnz = new_row_offsets[new_row_offsets.size() - 1];\n typedef typename MatPrecisionMap::Type ValueTypeA;\n VVector new_values((new_nnz + 1 + diag * (total_rows - 1))* this->A->get_block_size(), types::util::get_zero());\n IVector new_col_indices(new_nnz, 0);\n //reorder based on row permutation\n reorder_whole_matrix <32> <<< num_blocks, 512>>>(this->A->row_offsets.raw(), this->A->col_indices.raw(), this->A->values.raw(), new_row_offsets.raw(), new_col_indices.raw(), new_values.raw(), this->renumbering.raw(), this->A->get_block_size(), size, insert);\n cudaCheckError();\n\n if (diag)\n {\n //reorder based on row permutation\n reorder_vector_values <<< num_blocks, 512>>>(new_values.raw() + new_row_offsets[total_rows]*this->A->get_block_size(),\n this->A->values.raw() + this->A->row_offsets[size]*this->A->get_block_size(),\n this->renumbering.raw(),\n this->A->get_block_size(), size);\n cudaCheckError();\n }\n\n int cumulative_num_rows = size;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n for (int ring = 0; ring < rings; ring++)\n {\n int num_rows = halo_btl[i].B2L_rings[0][ring + 1] - halo_btl[i].B2L_rings[0][ring];\n int num_blocks = std::min(4096, (num_rows + 127) / 128);\n //copy in nonzeros that we are keeping\n //TODO: access pattern - should be implemented with warp-wide scans to decide which nonzeros we are keeping and where the rest is going\n reorder_whole_halo_matrix <<< num_blocks, 128>>>(halo_rows[i].row_offsets.raw() + halo_btl[i].B2L_rings[0][ring],\n halo_rows[i].col_indices.raw(), halo_rows[i].values.raw(),\n new_row_offsets.raw() + this->halo_offsets[ring * num_neighbors + i],\n new_col_indices.raw(), new_values.raw(), NULL, this->A->get_block_size(), num_rows,\n insert, this->halo_offsets[ring * num_neighbors + i], halo_btl[i].B2L_rings[0][ring], halo_btl[i].B2L_rings[0][rings]);\n cudaCheckError();\n\n if (diag)\n {\n amgx::thrust::copy(halo_rows[i].values.begin() + (halo_rows[i].row_offsets[halo_rows[i].get_num_rows()] + halo_btl[i].B2L_rings[0][ring])*this->A->get_block_size(),\n halo_rows[i].values.begin() + (halo_rows[i].row_offsets[halo_rows[i].get_num_rows()] + halo_btl[i].B2L_rings[0][ring + 1])*this->A->get_block_size(),\n new_values.begin() + (new_row_offsets[total_rows] + cumulative_num_rows)*this->A->get_block_size());\n cumulative_num_rows += num_rows;\n }\n }\n }\n\n cudaCheckError();\n /*\n EXAMPLE\n copy everything in place, dropping -1 column indices in the halo and reordering the owned rows\n new_row_offsets = [0 4 8 15 20| 25 31| 35 41| 44 48| 51 55]\n new_col_indices = [1 2 0 3\n 4 1 2 0\n 4 5 1 2 0 3 7\n 2 0 3 6 7 -end of owned\n 9 4 5 1 2\n 8 9 4 5 1 2 - end of neighbor 0 ring 0\n 3 6 7 11\n 2 3 10 6 7 11 - end of neighbor 1 ring 0\n 8 9 5\n 8 9 4 5 - end of neighbor 0 ring 1\n 10 7 11\n 10 6 7 11] - end of neighbor 1 ring 1\n\n */\n this->A->set_num_cols(total_rows);\n this->A->set_num_rows(size);\n this->A->col_indices.swap(new_col_indices);\n new_row_offsets.resize(total_rows + 1);\n this->A->row_offsets.swap(new_row_offsets);\n new_row_offsets.swap(this->old_row_offsets);\n this->A->values.swap(new_values);\n this->A->m_seq_offsets.resize(total_rows + 1);\n thrust_wrapper::sequence(this->A->m_seq_offsets.begin(), this->A->m_seq_offsets.end());\n\n if (insert)\n {\n this->A->delProps(DIAG);\n this->A->diag.resize(total_rows);\n amgx::thrust::copy(this->A->row_offsets.begin(), this->A->row_offsets.end() - 1, this->A->diag.begin());\n }\n\n cudaCheckError();\n delete this->halo_rows;\n delete this->halo_btl;\n //set halo_rows and halo_btl to NULL to avoid a potential double free situation in the future\n this->halo_rows = NULL;\n this->halo_btl = NULL;\n this->A->delProps(COO);\n this->A->set_initialized(1);\n\n //TODO: only do this if AMG_Config matrix_halo_exchange!=2\n if (!insert) { this->A->computeDiagonal(); }\n\n this->A->setView(OWNED);\n}\n\n//function object (functor) for thrust calls (it is a unary operator to add a constant)\ntemplate\nclass add_constant_op\n{\n const T c;\n public:\n add_constant_op(T _c) : c(_c) {}\n __host__ __device__ T operator()(const T &x) const\n {\n return x + c;\n }\n};\n\ntemplate \nvoid DistributedManager >::obtain_shift_l2g_reordering(index_type n, I64Vector_d &l2g, IVector_d &p, IVector_d &q)\n{\n /* WARNING: Exchange halo of the inverse_reordering, which is implicitly based on the local_to_global_map (l2g).\n Notice that it is implicit in the exchange_halo routine, since you are getting exactly the vector\n halo elements, which are exactly the elements you need. They however must be shifted by the partition\n starting points (starting global row indices, which are containe din array part_offsets).\n This allows us to avoid constructing the global vector for inverse permutation,\n as is done in reference MATLAB code. */\n //Recall that part_offsets provide the starting point (global row index) of every partition, in other words,\n //they contain the prefix sum of number of rows assigned to each partition. Also, notice that part_offsets and\n //part_offsets_h have the same values on device and host, respectively. See below few lines for details:\n index_type tag = 1 * 133 + 3 * 7 + 0; //some random number for the tag\n index_type l = p.size();\n q.resize(l);\n amgx::thrust::copy (p.begin(), p.end(), q.begin());\n thrust_wrapper::transform(q.begin(), q.end(), q.begin(), add_constant_op(this->part_offsets[this->global_id()]));\n this->exchange_halo(q, tag);\n thrust_wrapper::sequence (q.begin(), q.begin() + n);\n thrust_wrapper::transform(q.begin(), q.begin() + n, q.begin(), add_constant_op(this->part_offsets[this->global_id()]));\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManager >::unpack_partition(index_type *Bp, index_type *Bc, mat_value_type *Bv)\n{\n index_type l, n, nnz, offset, block_dimx, block_dimy;\n index_type *ir;\n index_type *Ap;\n index_type *Ac;\n mat_value_type *Av;\n IVector q;\n //some initializations\n this->A->getOffsetAndSizeForView(OWNED, &offset, &n);\n this->A->getNnzForView(OWNED, &nnz);\n\n block_dimx = this->A->get_block_dimx();\n block_dimy = this->A->get_block_dimy();\n\n l = this->inverse_renumbering.size();\n ir = this->inverse_renumbering.raw();\n Ap = this->A->row_offsets.raw();\n Ac = this->A->col_indices.raw();\n Av = this->A->values.raw();\n //(i) reorder the matrix back (into mixed interior-boundary nodes)\n //applies to rows and columns (out-of-place)\n reorder_partition\n\t(n, nnz, Ap, Ac, Av, Bp, Bc, Bv, l, ir, block_dimx, block_dimy);\n cudaCheckError();\n //obtain reordering q that combines the shift of the diagonal block with the off-diagonal block indices conversion from local to global\n this->obtain_shift_l2g_reordering(n, this->local_to_global_map, this->inverse_renumbering, q);\n cudaCheckError();\n //(ii) reorder the matrix back (shift the diagonal block and convert off-diagonal block column indices from local to global)\n //applies columns only (in-place)\n reorder_partition\n (n, nnz, Bp, Bc, Bv, Bp, Bc, Bv, q.size(), q.raw(), block_dimx, block_dimy);\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManagerBase::createNeighToDestPartMap(IVector_h &neigh_to_part, IVector_h &neighbors, IVector_h &destination_part, int num_neighbors)\n{\n neigh_to_part.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n neigh_to_part[i] = destination_part[neighbors[i]];\n }\n}\n\ntemplate \nvoid DistributedManagerBase::createConsolidatedNeighToPartMap(IVector_h &cons_neigh_to_part, IVector_h &neigh_to_part, int my_destination_part, IVector_h &destination_part, int &num_cons_neighbors)\n{\n // input: non-initialized cons_neigh_to_part\n // fine_neigh_to_part\n // my_destination_part\n // output: cons_neigh_to_part\n // num_cons_neighbors\n cons_neigh_to_part = neigh_to_part;\n amgx::thrust::sort(cons_neigh_to_part.begin(), cons_neigh_to_part.end());\n cudaCheckError();\n cons_neigh_to_part.erase(amgx::thrust::unique(cons_neigh_to_part.begin(), cons_neigh_to_part.end()), cons_neigh_to_part.end());\n // Remove if fine_neigh maps to same coarse partition\n cons_neigh_to_part.erase(amgx::thrust::remove_if(cons_neigh_to_part.begin(), cons_neigh_to_part.end(), is_my_part(my_destination_part)), cons_neigh_to_part.end());\n num_cons_neighbors = cons_neigh_to_part.size();\n cudaCheckError();\n}\n\n\ntemplate \nvoid DistributedManagerBase::createNeighToConsNeigh(IVector_h &neigh_to_cons_neigh, IVector_h &cons_neigh_to_part, IVector_h &neigh_to_part, int my_destination_part, int &num_neighbors)\n{\n neigh_to_cons_neigh.resize(num_neighbors);\n amgx::thrust::lower_bound(cons_neigh_to_part.begin(), cons_neigh_to_part.end(), neigh_to_part.begin(), neigh_to_part.end(), neigh_to_cons_neigh.begin());\n cudaCheckError();\n\n // Flagging fine neighbors that go to same partition (haven't been found in previous step)\n for (int i = 0; i < num_neighbors; i++)\n {\n if ( neigh_to_part[i] == my_destination_part)\n {\n neigh_to_cons_neigh[i] = -1;\n }\n }\n}\n\ntemplate \ntemplate \nvoid DistributedManagerBase::consB2Lmaps(std::vector &dest_coarse_B2L_maps, std::vector &coarse_B2L_maps, IVector_h &fine_neigh_to_coarse_neigh, int num_coarse_neighbors, int num_fine_neighbors)\n{\n //Merge B2L fine maps per coarse destination\n dest_coarse_B2L_maps.resize(num_coarse_neighbors);\n std::vector dest_coarse_B2L_maps_scratch_sizes(num_coarse_neighbors, 0);\n int my_id = this->global_id();\n\n // Loop over the fine neighbors, to compute size of each dest_coarse_B2L_maps\n for (int i = 0; i < num_fine_neighbors; i++)\n {\n int k = fine_neigh_to_coarse_neigh[i];\n\n if (k != -1)\n {\n dest_coarse_B2L_maps_scratch_sizes[k] += coarse_B2L_maps[i].size();\n }\n }\n\n // Now fill dest_coarse_B2L_maps\n for (int k = 0; k < num_coarse_neighbors; k++)\n {\n dest_coarse_B2L_maps[k].resize( dest_coarse_B2L_maps_scratch_sizes[k] );\n // Reset sizes to 0 (fill use as offset in next loop);\n dest_coarse_B2L_maps_scratch_sizes[k] = 0;\n }\n\n for (int i = 0; i < num_fine_neighbors; i++)\n {\n int k = fine_neigh_to_coarse_neigh[i];\n\n if (k != -1)\n {\n int offset = dest_coarse_B2L_maps_scratch_sizes[k];\n amgx::thrust::copy(coarse_B2L_maps[i].begin(), coarse_B2L_maps[i].end(), dest_coarse_B2L_maps[k].begin() + offset);\n dest_coarse_B2L_maps_scratch_sizes[k] += coarse_B2L_maps[i].size();\n }\n }\n\n cudaCheckError();\n int max_size = 0;\n\n for (int i = 0; i < num_coarse_neighbors; i++)\n {\n int size = dest_coarse_B2L_maps[i].size();\n\n if (size > max_size) { max_size = size; }\n }\n\n // Remove duplicates (aggregates in boundary that go to same merged partition)\n for (int i = 0; i < num_coarse_neighbors; i++)\n {\n int size = dest_coarse_B2L_maps[i].size();\n amgx::thrust::sort(dest_coarse_B2L_maps[i].begin(), dest_coarse_B2L_maps[i].begin() + size);\n index_type num_unique = amgx::thrust::unique(dest_coarse_B2L_maps[i].begin(), dest_coarse_B2L_maps[i].begin() + size) - dest_coarse_B2L_maps[i].begin();\n dest_coarse_B2L_maps[i].erase(dest_coarse_B2L_maps[i].begin() + num_unique, dest_coarse_B2L_maps[i].end());\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManagerBase::computeConsolidatedOffsets(const int my_id, const int my_destination_part, const bool is_root_partition, const int num_interior_rows, const int num_boundary_rows, IVector_h_vector &vertex_counts, const IVector_h &parts_to_consolidate, const int num_parts_to_consolidate, int &interior_offset, int &boundary_offset, int &total_interior_rows_in_merged, int &total_boundary_rows_in_merged, int &total_rows_in_merged, DistributedComms *comms)\n{\n IVector_h my_offsets(4);\n IVector_h my_sizes(2);\n my_sizes[0] = num_interior_rows;\n my_sizes[1] = num_boundary_rows;\n\n if (!is_root_partition)\n {\n //Send number of interior and boundary nodes to root\n comms->send_vector_async(my_sizes, my_destination_part, 777);\n comms->recv_vector(my_offsets, my_destination_part, 778);\n comms->send_vector_wait_all(my_sizes);\n }\n else\n {\n vertex_counts.resize(num_parts_to_consolidate);\n IVector_h child_sizes(2);\n IVector_h offsets_interior(num_parts_to_consolidate);\n IVector_h offsets_boundary(num_parts_to_consolidate);\n int count_int = 0;\n int count_bdy = 0;\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n if (parts_to_consolidate[i] == my_id)\n {\n child_sizes = my_sizes;\n }\n else\n {\n comms->recv_vector(child_sizes, parts_to_consolidate[i], 777);\n }\n\n //Do a simple cumulative sum to determine total number of interior/boundary rows and their offsets on a per contributing partition basis\n offsets_interior[i] = count_int;\n offsets_boundary[i] = count_bdy;\n count_int += child_sizes[0];\n count_bdy += child_sizes[1];\n //Save them\n vertex_counts[i].resize(2);\n vertex_counts[i][0] = child_sizes[0];\n vertex_counts[i][1] = child_sizes[1];\n }\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n //Send back to contributing partitions\n IVector_h offsets_to_send(4);\n offsets_to_send[0] = offsets_interior[i];\n offsets_to_send[1] = offsets_boundary[i];\n offsets_to_send[2] = count_int;\n offsets_to_send[3] = count_bdy;\n\n if (parts_to_consolidate[i] == my_id)\n {\n my_offsets = offsets_to_send;\n }\n else\n {\n comms->send_vector(offsets_to_send, parts_to_consolidate[i], 778); // cannot make async, rewriting internal buffer\n }\n }\n }\n\n interior_offset = my_offsets[0];\n boundary_offset = my_offsets[1] + my_offsets[2] - num_interior_rows;\n total_interior_rows_in_merged = my_offsets[2];\n total_boundary_rows_in_merged = my_offsets[3];\n total_rows_in_merged = my_offsets[2] + my_offsets[3];\n}\n\n\n\n\ntemplate \ntemplate \nvoid DistributedManagerBase::consB2LmapsOnRoot(int &num_consolidated_neighbors, std::vector &consolidated_B2L_maps, IVector_h &consolidated_coarse_ids, std::vector &dest_coarse_B2L_maps, IVector_h &coarse_neigh_to_fine_part, IVector_h &num_bdy_per_coarse_neigh, IVector_h &fine_parts_to_consolidate, int num_fine_parts_to_consolidate, int my_id, int my_destination_part, bool is_root_partition, int num_coarse_neighbors, DistributedComms *comms)\n{\n // TODO: it is possible to optimize exchanges, for example fuse recv_vector in recreating coarse neigbours\n\n // output: num_consolidated_neighbor, consolidated_B2L_maps, consolidated_coarse_ids\n // input: dest_coarse_B2L_maps, is_root_partition, my_id, my_destination_part, num_fine_parts_to_consolidate, num_coarse_neighbors, coarse_neigh_to_fine_part, num_bdy_per_coarse_neigh\n if (my_destination_part != my_id)\n {\n //if not root, send coarse neighbor list using fine indices and the corresponding boundary lists\n IVector_h num_coarse_neigh(1);\n num_coarse_neigh[0] = num_coarse_neighbors;\n comms->send_vector_async(num_coarse_neigh, my_destination_part, 1111);\n comms->send_vector_async(coarse_neigh_to_fine_part, my_destination_part, 2222);\n comms->send_vector_async(num_bdy_per_coarse_neigh, my_destination_part, 3333);\n\n for (int i = 0; i < num_coarse_neighbors; i++)\n {\n comms->send_vector_async(dest_coarse_B2L_maps[i], my_destination_part, 4444 + i) ;\n }\n\n comms->send_vector_wait_all(num_coarse_neigh);\n comms->send_vector_wait_all(coarse_neigh_to_fine_part);\n comms->send_vector_wait_all(num_bdy_per_coarse_neigh);\n\n for (int i = 0; i < num_coarse_neighbors; i++)\n {\n comms->send_vector_wait_all(dest_coarse_B2L_maps[i]) ;\n }\n }\n\n if (is_root_partition)\n {\n IVector_h num_coarse_ids_from_part(fine_parts_to_consolidate);\n IVector_h_vector coarse_ids_from_part(num_fine_parts_to_consolidate);\n IVector_h_vector num_coarse_neigh_bdys_from_part(num_fine_parts_to_consolidate);\n //If root, receive sizes, and resize receive buffers\n int total_num_coarse_ids = 0;\n\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n IVector_h temp(1);\n\n if (current_part != my_id)\n {\n comms->recv_vector(temp, current_part, 1111);\n }\n else\n {\n temp[0] = num_coarse_neighbors;\n }\n\n num_coarse_ids_from_part[i] = temp[0];\n coarse_ids_from_part[i].resize(temp[0]);\n num_coarse_neigh_bdys_from_part[i].resize(temp[0]);\n total_num_coarse_ids += temp[0];\n }\n\n //Create a neighbor list for the consolidated coarse matrix, by merging coarse neighbor lists from partitions that are being merged\n consolidated_coarse_ids.resize(total_num_coarse_ids);\n int count = 0;\n\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n\n // Get from each partition the coarse partition ids in their B2L maps\n if (current_part != my_id)\n {\n comms->recv_vector(coarse_ids_from_part[i], current_part, 2222);\n comms->recv_vector(num_coarse_neigh_bdys_from_part[i], current_part, 3333);\n }\n else\n {\n coarse_ids_from_part[i] = coarse_neigh_to_fine_part;\n num_coarse_neigh_bdys_from_part[i] = num_bdy_per_coarse_neigh;\n }\n\n amgx::thrust::copy(coarse_ids_from_part[i].begin(), coarse_ids_from_part[i].end(), consolidated_coarse_ids.begin() + count);\n count += num_coarse_ids_from_part[i];\n }\n\n cudaCheckError();\n //eliminate duplicates\n amgx::thrust::sort(consolidated_coarse_ids.begin(), consolidated_coarse_ids.end());\n cudaCheckError();\n consolidated_coarse_ids.erase(amgx::thrust::unique(consolidated_coarse_ids.begin(), consolidated_coarse_ids.end()), consolidated_coarse_ids.end());\n cudaCheckError();\n num_consolidated_neighbors = consolidated_coarse_ids.size();\n IVector_h_vector coarse_ids_from_part_to_consolidated_neighbor(num_fine_parts_to_consolidate);;\n\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n coarse_ids_from_part_to_consolidated_neighbor[i].resize(num_coarse_ids_from_part[i]);\n amgx::thrust::lower_bound(consolidated_coarse_ids.begin(), consolidated_coarse_ids.end(), coarse_ids_from_part[i].begin(), coarse_ids_from_part[i].end(), coarse_ids_from_part_to_consolidated_neighbor[i].begin());\n }\n\n cudaCheckError();\n // Map each coarse partition to new coarse ID\n consolidated_B2L_maps.resize(num_consolidated_neighbors);\n IVector_h consolidated_B2L_maps_sizes(num_consolidated_neighbors);\n // Offset in the consolidated_B2L_maps\n IVector_h_vector coarse_ids_offsets(num_fine_parts_to_consolidate);\n\n for (int i = 0; i < num_consolidated_neighbors; i++)\n {\n consolidated_B2L_maps_sizes[i] = 0;\n }\n\n // Compute the size of each consolidated_B2L_maps and offsets into it, where we will receive the parts coming from partitions that are getting merged into this one\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n coarse_ids_offsets[i].resize(num_coarse_ids_from_part[i]);\n\n for (int j = 0; j < num_coarse_ids_from_part[i]; j++)\n {\n int coarse_id = coarse_ids_from_part[i][j];\n int k = num_coarse_neigh_bdys_from_part[i][j];\n coarse_ids_offsets[i][j] = consolidated_B2L_maps_sizes[ coarse_ids_from_part_to_consolidated_neighbor[i][j] ];\n consolidated_B2L_maps_sizes[ coarse_ids_from_part_to_consolidated_neighbor[i][j] ] += k;\n }\n }\n\n for (int i = 0; i < num_consolidated_neighbors; i++)\n {\n consolidated_B2L_maps[i].resize(consolidated_B2L_maps_sizes[i]);\n }\n\n // Receive the B2L maps from each child partition, concatenate them (gets sorted outside)\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n\n for (int j = 0; j < num_coarse_ids_from_part[i]; j++)\n {\n int my_coarse_neigh = coarse_ids_from_part_to_consolidated_neighbor[i][j];\n int offset = coarse_ids_offsets[i][j];\n\n if (current_part != my_id)\n {\n comms->recv_vector( consolidated_B2L_maps[my_coarse_neigh], current_part, 4444 + j, offset, num_coarse_neigh_bdys_from_part[i][j]); //Need to do proper tagging here, otherwise messages from the same source would get mixed up\n }\n else\n {\n amgx::thrust::copy(dest_coarse_B2L_maps[j].begin(), dest_coarse_B2L_maps[j].end(), consolidated_B2L_maps[my_coarse_neigh].begin() + offset);\n }\n }\n }\n\n cudaCheckError();\n }\n}\n\n\ntemplate \nvoid DistributedManagerBase::consolidateAndRenumberHalos(IVector_h &aggregates, const IVector_h &manager_halo_offsets, IVector_h &halo_offsets, const IVector_h &neighbors, int num_fine_neighbors, const IVector_h &consolidated_coarse_ids, int num_consolidated_neighbors, const IVector_h &destination_part, int my_destination_part, bool is_root_partition, IVector_h &fine_parts_to_consolidate, int num_fine_parts_to_consolidate, int num_parts, int my_id, int total_rows_in_merged, int &num_all_aggregates, DistributedComms *comms)\n{\n consAndRenumberHalos(aggregates, manager_halo_offsets, halo_offsets, neighbors, num_fine_neighbors, consolidated_coarse_ids, num_consolidated_neighbors, destination_part, my_destination_part, is_root_partition, fine_parts_to_consolidate, num_fine_parts_to_consolidate, num_parts, my_id, total_rows_in_merged, num_all_aggregates, comms);\n}\n\ntemplate \nvoid DistributedManagerBase::consolidateAndRenumberHalos(IVector_d &aggregates, const IVector_h &manager_halo_offsets, IVector_h &halo_offsets, const IVector_h &neighbors, int num_fine_neighbors, const IVector_h &consolidated_coarse_ids, int num_consolidated_neighbors, const IVector_h &destination_part, int my_destination_part, bool is_root_partition, IVector_h &fine_parts_to_consolidate, int num_fine_parts_to_consolidate, int num_parts, int my_id, int total_rows_in_merged, int &num_all_aggregates, DistributedComms *comms)\n{\n consAndRenumberHalos(aggregates, manager_halo_offsets, halo_offsets, neighbors, num_fine_neighbors, consolidated_coarse_ids, num_consolidated_neighbors, destination_part, my_destination_part, is_root_partition, fine_parts_to_consolidate, num_fine_parts_to_consolidate, num_parts, my_id, total_rows_in_merged, num_all_aggregates, comms);\n}\n\ntemplate \ntemplate \nvoid DistributedManagerBase::consAndRenumberHalos(IVector_hd &aggregates, const IVector_h &manager_halo_offsets, IVector_h &halo_offsets, const IVector_h &neighbors, int num_fine_neighbors, const IVector_h &consolidated_coarse_neigh_to_fine_part, int num_consolidated_neighbors, const IVector_h &destination_part, int my_destination_part, bool is_root_partition, IVector_h &fine_parts_to_consolidate, int num_fine_parts_to_consolidate, int num_parts, int my_id, int total_rows_in_merged, int &num_all_aggregates, DistributedComms *comms)\n{\n /*\n * EXAMPLE 2\n This example is independent from the previous ones.\n Say partition 0 and 1 are merging (into 0) partition 0 is neighbors with 1,2,3 and partition 1 is neighbors with 0,3,4\n Partitions 3 and 4 are merging (into partition 3) and partition 2 is not merging with anyone.\n This example details the renumbering of halo indices on partition 0 and partition 1.\n aggregates on partition 0:\n [(fine interior nodes) (fine boundary nodes) (fine halo from part 1) (fine halo from part 2) (fine halo from part 3)]\n [(fine interior nodes) (fine boundary nodes) (13 13 15) (12 15 17) (14 16 18)]\n aggregates on partition 1:\n [(fine interior nodes) (fine boundary nodes) (fine halo from part 0) (fine halo from part 3) (fine halo from part 4)]\n [(fine interior nodes) (fine boundary nodes) (14 16 17) (18 19 19) (15 15 17)]\n\n manager_halo_offsets on partition 0:\n [22 25 28 31]\n manager_halo_offsets on partition 1:\n [20 23 26 29]\n halo_offsets on both partitions are uninitialised: [0 0 0] and [0 0]\n neighbors on partition 0: [1 2 3] partition 1: [0 3 4]\n num_fine_neighbors partition 0: 3 partition 1: 3\n consolidated_coarse_neigh_to_fine_part partition 0: [2 3] partition 1: [3]\n num_consolidated_neighbors partition 0: 2 partition 1: 1\n destination_part [0 0 2 3 3]\n my_destination_part partition 0: 0 partition 1: 0\n is_root_partition partition 0: true partition 1: false\n fine_parts_to_consolidate partition 0: [0 1]\n num_fine_parts_to_consolidate partition 0: 2\n num_parts 5\n my_id partition 0: 0 partition 1: 1\n total_rows_in_merged partition 0 and 1: 24 (=sum of the two below)\n num_all_aggregates partition partition 0: 13 partition 1: 11 - will be updated with the number of halo aggregates\n */\n //\n // Step 9.2 - com up with nonmerge lists\n //\n int num_fine_nonmerge_neighbors;// = fine_nonmerge_neighbors.size();\n //NUmber of neighbors we are not merging with\n num_fine_nonmerge_neighbors = 0;\n\n for (int i = 0 ; i < num_fine_neighbors; i++)\n {\n if (destination_part[neighbors[i]] != my_destination_part)\n {\n num_fine_nonmerge_neighbors++;\n }\n }\n\n IVector_h halo_sizes(num_fine_nonmerge_neighbors);\n IVector_h fine_nonmerge_neigh_to_cons_fine_part(num_fine_nonmerge_neighbors);\n IVector_h fine_nonmerge_neighbor_to_fine_neighbor(num_fine_nonmerge_neighbors);\n num_fine_nonmerge_neighbors = 0;\n\n for (int i = 0 ; i < num_fine_neighbors; i++)\n {\n if (destination_part[neighbors[i]] != my_destination_part)\n {\n halo_sizes[num_fine_nonmerge_neighbors] = manager_halo_offsets[i + 1] - manager_halo_offsets[i];\n fine_nonmerge_neighbor_to_fine_neighbor[num_fine_nonmerge_neighbors] = i;\n fine_nonmerge_neigh_to_cons_fine_part[num_fine_nonmerge_neighbors] = destination_part[neighbors[i]];\n num_fine_nonmerge_neighbors++;\n }\n }\n\n /*\n * EXAMPLE 2\n num_fine_nonmerge_neighbors partition 0: 2 partition 1: 2\n fine_nonmerge_neighbor_to_fine_neighbor partition 0: [1 2] partition 1: [1 2]\n fine_nonmerge_neigh_to_cons_fine_part partition 0: [2 3] partition 1: [3 3]\n halo_sizes partition 0: [3 3] partition 1: [3 3]\n */\n\n //Send them to root along with the halo parts of the aggregates vector\n if (!is_root_partition)\n {\n IVector_h num_fine_nonmerge_neigh(1);\n num_fine_nonmerge_neigh[0] = num_fine_nonmerge_neighbors;\n // TODO: async? might be faster.\n comms->send_vector(num_fine_nonmerge_neigh, my_destination_part, 1111);\n comms->send_vector(halo_sizes, my_destination_part, 2222);\n comms->send_vector(fine_nonmerge_neigh_to_cons_fine_part, my_destination_part, 3333);\n\n // Here check l2h_identity flag and act accordingly\n for (int i = 0; i < num_fine_nonmerge_neighbors; i++)\n {\n comms->send_vector_async(aggregates, my_destination_part, 4444 + i, manager_halo_offsets[fine_nonmerge_neighbor_to_fine_neighbor[i]], halo_sizes[i]) ;\n }\n\n //comms->send_vector_wait_all(num_fine_nonmerge_neigh);\n //comms->send_vector_wait_all(halo_sizes);\n //comms->send_vector_wait_all(fine_nonmerge_neigh_to_cons_fine_part);\n comms->send_vector_wait_all(aggregates);\n /*\n * EXAMPLE 2\n Partition 1 sends to partition 0:\n num_fine_nonmerge_neigh 2\n halo_sizes [3 3]\n fine_nonmerge_neigh_to_cons_fine_part [3 3]\n for loop: sends two pieces: [(18 19 19)] [(15 15 17)]\n */\n }\n\n if (is_root_partition)\n {\n //\n // Step 9.3 Root receives this info, creates metadata\n //\n std::vector num_fine_nonmerge_neigh_array(num_fine_parts_to_consolidate);\n IVector_h_vector halo_sizes_array(num_fine_parts_to_consolidate);\n IVector_h_vector fine_nonmerge_neigh_to_cons_fine_part_array(num_fine_parts_to_consolidate);\n std::vector > fine_halo_aggregates_to_root_array(num_fine_parts_to_consolidate);\n std::vector min_index_coarse_halo(num_consolidated_neighbors, 0x7FFFFFFF);\n std::vector max_index_coarse_halo(num_consolidated_neighbors, 0);\n std::vector fine_part_to_consolidated_neighbor(num_parts, -1);\n\n for (int i = 0; i < num_consolidated_neighbors; i++)\n {\n fine_part_to_consolidated_neighbor[consolidated_coarse_neigh_to_fine_part[i]] = i;\n }\n\n /*\n * EXAMPLE 2\n everything from here on is for partition 0, since that is the root partition\n fine_part_to_consolidated_neighbor [-1 -1 0 1 -1]\n */\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n IVector_h temp(1);\n\n if (current_part != my_id)\n {\n comms->recv_vector(temp, current_part, 1111);\n }\n else\n {\n temp[0] = num_fine_nonmerge_neighbors;\n }\n\n num_fine_nonmerge_neigh_array[i] = temp[0];\n halo_sizes_array[i].resize(temp[0]);\n fine_nonmerge_neigh_to_cons_fine_part_array[i].resize(temp[0]);\n fine_halo_aggregates_to_root_array[i].resize(temp[0]);\n\n if (current_part != my_id)\n {\n comms->recv_vector(halo_sizes_array[i], current_part, 2222);\n }\n else\n {\n halo_sizes_array[i] = halo_sizes;\n }\n\n if (current_part != my_id)\n {\n comms->recv_vector(fine_nonmerge_neigh_to_cons_fine_part_array[i], current_part, 3333);\n }\n else\n {\n fine_nonmerge_neigh_to_cons_fine_part_array[i] = fine_nonmerge_neigh_to_cons_fine_part;\n }\n\n //Receive the halo regions\n for (int j = 0; j < temp[0]; j++)\n {\n fine_halo_aggregates_to_root_array[i][j].resize(halo_sizes_array[i][j]);\n\n if (current_part != my_id)\n {\n comms->recv_vector(fine_halo_aggregates_to_root_array[i][j], current_part, 4444 + j);\n }\n else\n {\n //HERE\n amgx::thrust::copy(aggregates.begin() + manager_halo_offsets[fine_nonmerge_neighbor_to_fine_neighbor[j]],\n aggregates.begin() + manager_halo_offsets[fine_nonmerge_neighbor_to_fine_neighbor[j]] + halo_sizes[j],\n fine_halo_aggregates_to_root_array[i][j].begin()); //TODO: not do this copying around on the root\n }\n\n#define MIN(a,b) ab?a:b;\n //Find minimum and maximum halo indices as not to allocate too much scratch space later\n int min_index = amgx::thrust::reduce(fine_halo_aggregates_to_root_array[i][j].begin(), fine_halo_aggregates_to_root_array[i][j].end(), int(0x7FFFFFFF), amgx::thrust::minimum());\n int max_index = amgx::thrust::reduce(fine_halo_aggregates_to_root_array[i][j].begin(), fine_halo_aggregates_to_root_array[i][j].end(), int(0), amgx::thrust::maximum());\n min_index_coarse_halo[fine_part_to_consolidated_neighbor[fine_nonmerge_neigh_to_cons_fine_part_array[i][j]]] = MIN((int)min_index_coarse_halo[fine_part_to_consolidated_neighbor[fine_nonmerge_neigh_to_cons_fine_part_array[i][j]]], min_index);\n max_index_coarse_halo[fine_part_to_consolidated_neighbor[fine_nonmerge_neigh_to_cons_fine_part_array[i][j]]] = MAX((int)max_index_coarse_halo[fine_part_to_consolidated_neighbor[fine_nonmerge_neigh_to_cons_fine_part_array[i][j]]], max_index);\n }\n }\n\n cudaCheckError();\n /*\n * EXAMPLE 2\n num_fine_nonmerge_neigh_array = [2 2]\n halo_sizes_array = [[3 3][3 3]]\n fine_nonmerge_neigh_to_cons_fine_part_array[][] = [[2 3][3 3]]\n fine_halo_aggregates_to_root_array[from][to][fine halo vertex] [[[12 15 17][14 16 18]]\n [[18 19 19][15 15 17]]]\n min_index_coarse_halo[12 14]\n max_index_coarse_halo[17 19]\n */\n halo_offsets[0] = total_rows_in_merged;\n //Now we have all the halo nodes, let's renumber them.\n int min_index = amgx::thrust::reduce(min_index_coarse_halo.begin(), min_index_coarse_halo.end(), int(0x7FFFFFFF), amgx::thrust::minimum());\n int max_index = amgx::thrust::reduce(max_index_coarse_halo.begin(), max_index_coarse_halo.end(), int(0), amgx::thrust::maximum());\n cudaCheckError();\n //\n // Step 9.4 compute halo indices on root nodes\n //\n int scratch_size;\n\n if (num_consolidated_neighbors == 0)\n {\n scratch_size = 1;\n }\n else\n {\n scratch_size = max_index - min_index + 2;\n }\n\n IVector scratch(scratch_size);\n\n for (int i = 0; i < num_consolidated_neighbors; i++)\n {\n thrust_wrapper::fill(scratch.begin(), scratch.end(), 0);\n int dest_part = consolidated_coarse_neigh_to_fine_part[i];\n\n //Flag halo indices that occur for a specific coarse neighbor\n for (int j = 0; j < num_fine_parts_to_consolidate; j++)\n {\n for (int k = 0; k < num_fine_nonmerge_neigh_array[j]; k++)\n {\n if (fine_nonmerge_neigh_to_cons_fine_part_array[j][k] == dest_part)\n {\n int size = halo_sizes_array[j][k];\n this->flag_halo_ids(size, scratch, fine_halo_aggregates_to_root_array[j][k], min_index_coarse_halo[i], max_index, min_index);\n }\n }\n }\n\n thrust_wrapper::exclusive_scan(scratch.begin(), scratch.end(), scratch.begin(), halo_offsets[i]); //renumber them with the proper offset into our halo\n halo_offsets[i + 1] = scratch[scratch.size() - 1];\n\n //now read them back\n for (int j = 0; j < num_fine_parts_to_consolidate; j++)\n {\n for (int k = 0; k < num_fine_nonmerge_neigh_array[j]; k++)\n {\n if (fine_nonmerge_neigh_to_cons_fine_part_array[j][k] == dest_part)\n {\n int size = halo_sizes_array[j][k];\n int block_size = 128;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (size - 1) / block_size + 1);\n this->read_halo_ids(size, scratch, fine_halo_aggregates_to_root_array[j][k], min_index_coarse_halo[i]);\n //and send them back to contributing partitions\n cudaDeviceSynchronize(); //TODO: don't need to synchronize when using GPUDirect\n cudaCheckError();\n int current_part = fine_parts_to_consolidate[j];\n int tag = 4444 + dest_part;\n\n if (current_part != my_id)\n {\n comms->send_vector_async(fine_halo_aggregates_to_root_array[j][k], current_part, tag); //!!!!: we are sending them back not in sequential order, need tags!!!!\n }\n else\n {\n amgx::thrust::copy(fine_halo_aggregates_to_root_array[j][k].begin(), fine_halo_aggregates_to_root_array[j][k].end(), aggregates.begin() + manager_halo_offsets[fine_nonmerge_neighbor_to_fine_neighbor[k]]);\n }\n }\n }\n }\n\n /*\n * EXAMPLE 2\n the array that is sent back in pieces:\n fine_halo_aggregates_to_root_array[from][to][fine halo vertex] [[[24 25 26][27 29 31]]\n [[31 32 32][28 28 30]]]\n halo_offsets = [24 27 33]\n */\n } // Loop over consolidated neighbors\n\n cudaCheckError();\n\n // Wait for sends to have completed (this is to prevent fine_halo_aggregates_to_root_array to be destroyed before send has finished)\n for (int i = 0; i < num_consolidated_neighbors; i++)\n {\n int dest_part = consolidated_coarse_neigh_to_fine_part[i];\n\n for (int j = 0; j < num_fine_parts_to_consolidate; j++)\n {\n for (int k = 0; k < num_fine_nonmerge_neigh_array[j]; k++)\n {\n if (fine_nonmerge_neigh_to_cons_fine_part_array[j][k] == dest_part)\n {\n int current_part = fine_parts_to_consolidate[j];\n\n if (current_part != my_id)\n {\n comms->send_vector_wait_all(fine_halo_aggregates_to_root_array[j][k]);\n }\n }\n }\n }\n } // Loop over consolidated neighbors\n\n //Send total number of rows in the aggregated matrix\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n IVector_h total_rows(1);\n total_rows[0] = halo_offsets[num_consolidated_neighbors];\n\n if (current_part != my_id)\n {\n comms->send_vector(total_rows, current_part, 5555);\n }\n else\n {\n num_all_aggregates = total_rows[0];\n }\n }\n } // If is root partition\n\n if (!is_root_partition)\n {\n for (int i = 0; i < num_fine_nonmerge_neighbors; i++)\n {\n int tag = 4444 + fine_nonmerge_neigh_to_cons_fine_part[i];\n comms->recv_vector(aggregates, my_destination_part, tag, manager_halo_offsets[fine_nonmerge_neighbor_to_fine_neighbor[i]], halo_sizes[i]);\n }\n\n IVector_h total_rows(1);\n comms->recv_vector(total_rows, my_destination_part, 5555);\n num_all_aggregates = total_rows[0];\n }\n\n /*\n * EXAMPLE 2\n num_all_aggregates = 33 (both partitions 0 and 1\n */\n}\n\ntemplate \nvoid DistributedManagerBase::ipcExchangePtr(void *&ptr, bool is_root_partition, int num_parts_to_consolidate, IVector_h &parts_to_consolidate, int my_root_partition, int my_id, DistributedComms *comms)\n{\n cudaIpcMemHandle_t handle;\n\n if (is_root_partition)\n {\n cudaIpcGetMemHandle( (cudaIpcMemHandle_t *) &handle, ptr ) ;\n cudaCheckError();\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n int current_part = parts_to_consolidate[i];\n\n if (current_part != my_id)\n {\n comms->send_raw_data(&handle, sizeof(handle), current_part, 456);\n }\n }\n }\n else\n {\n comms->recv_raw_data(&handle, sizeof(handle), my_root_partition, 456);\n cudaError_t err = cudaIpcOpenMemHandle( (void **) &ptr, handle, cudaIpcMemLazyEnablePeerAccess);\n cudaCheckError();\n }\n}\n\ntemplate \nvoid DistributedManagerBase::ipcWaitForChildren(bool is_root_partition, int num_parts_to_consolidate, IVector_h &parts_to_consolidate, int my_destination_part, int my_id, DistributedComms *comms)\n{\n cudaEvent_t event;\n cudaIpcEventHandle_t event_handle;\n cudaEventCreateWithFlags(&event, cudaEventDisableTiming | cudaEventInterprocess);\n cudaCheckError();\n cudaIpcGetEventHandle( &event_handle, event);\n cudaCheckError();\n // Each rank record the event\n cudaEventRecord(event);\n cudaCheckError();\n\n if (is_root_partition)\n {\n std::vector child_events(num_parts_to_consolidate);\n std::vector child_event_handles(num_parts_to_consolidate);\n\n // Root partition receives event_handles from child and stores in child_event_handles\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n int current_part = parts_to_consolidate[i];\n\n if (current_part != my_id)\n {\n comms->recv_raw_data(&(child_event_handles[i]), sizeof(cudaIpcEventHandle_t), current_part, 987 + current_part);\n cudaIpcOpenEventHandle(&child_events[i], child_event_handles[i]);\n cudaCheckError();\n }\n }\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n if (parts_to_consolidate[i] != my_id)\n {\n cudaEventSynchronize(child_events[i]);\n cudaCheckError();\n }\n }\n }\n else\n {\n comms->send_raw_data(&event_handle, sizeof(cudaIpcEventHandle_t), my_destination_part, 987 + my_id);\n }\n}\n\ntemplate \nvoid DistributedManagerBase::ipcWaitForRoot(bool is_root_partition, int num_parts_to_consolidate, IVector_h &parts_to_consolidate, int my_destination_part, int my_id, DistributedComms *comms)\n{\n cudaEvent_t event;\n cudaIpcEventHandle_t event_handle;\n cudaEventCreateWithFlags(&event, cudaEventDisableTiming | cudaEventInterprocess);\n\n if (is_root_partition)\n {\n cudaIpcGetEventHandle( &event_handle, event);\n cudaCheckError();\n // Root records the event\n cudaEventRecord(event);\n cudaCheckError();\n\n // Root partition sends event_handles to child\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n int current_part = parts_to_consolidate[i];\n\n if (current_part != my_id)\n {\n comms->send_raw_data(&event_handle, sizeof(event_handle), current_part, 988 + current_part);\n }\n }\n }\n else\n {\n comms->recv_raw_data(&event_handle, sizeof(event_handle), my_destination_part, 988 + my_id);\n cudaIpcOpenEventHandle(&event, event_handle);\n cudaCheckError();\n cudaEventSynchronize(event);\n cudaCheckError();\n }\n}\n\n\n\ntemplate \nvoid DistributedManagerBase::read_halo_ids(int size, IVector_d &scratch, IVector_d &halo_aggregates, VecInt_t min_index_coarse_halo)\n{\n int block_size = 128;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (size - 1) / block_size + 1);\n read_halo_ids_kernel <<< num_blocks, block_size>>>(scratch.raw(), halo_aggregates.raw(), min_index_coarse_halo, size);\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManagerBase::read_halo_ids(int size, IVector_h &scratch, IVector_h &halo_aggregates, VecInt_t min_index_coarse_halo)\n{\n FatalError(\"read_halo_ids not implemented on host yet\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n\ntemplate \nvoid DistributedManagerBase::flag_halo_ids(int size, IVector_d &scratch, IVector_d &halo_aggregates, VecInt_t min_index_coarse_halo, int max_index, int min_index)\n{\n int block_size = 128;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (size - 1) / block_size + 1);\n flag_halo_ids_kernel <<< num_blocks, block_size>>>(scratch.raw(), halo_aggregates.raw(), min_index_coarse_halo, size, max_index - min_index + 1);\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManagerBase::flag_halo_ids(int size, IVector_h &scratch, IVector_h &halo_aggregates, VecInt_t min_index_coarse_halo, int max_index, int min_index)\n{\n FatalError(\"flag_halo_ids not implemented on host yet\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid DistributedManager >::consolidateAndUploadAll(int n, int nnz, int block_dimx, int block_dimy, const int *row_ptrs, const int *col_indices, const void *data, const void *diag, Matrix &A)\n{\n FatalError(\"Fine level consolidation not implemented on host yet\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate\nvoid DistributedManagerBase::exchangeSolveResultsConsolidation(int &num_iters, std::vector &res_history, AMGX_STATUS &status, bool store_res_history)\n{\n int bsize = this->A->get_block_size();\n PODVector_h res_history_tmp;\n\n if (!m_is_fine_level_consolidated)\n {\n return;\n }\n else\n {\n int my_id = this->getFineLevelComms()->get_global_id();\n IVector_h my_num_iters(1);\n\n if (m_is_fine_level_root_partition)\n {\n my_num_iters[0] = num_iters;\n\n if (store_res_history)\n {\n // Pack the res_history vector into array\n res_history_tmp.resize( (num_iters + 1)*bsize);\n\n for (int i = 0; i < num_iters + 1; i++)\n {\n for (int j = 0; j < bsize; j++)\n {\n res_history_tmp[i * bsize + j] = res_history[i][j];\n }\n }\n }\n\n for (int i = 0; i < m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = m_fine_level_parts_to_consolidate[i];\n\n if (my_id != current_part)\n {\n getFineLevelComms()->send_vector_async(my_num_iters, current_part, 245);\n\n if (store_res_history)\n {\n getFineLevelComms()->send_vector_async(res_history_tmp, current_part, 246);\n }\n }\n }\n\n for (int i = 0; i < m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = m_fine_level_parts_to_consolidate[i];\n\n if (my_id != current_part)\n {\n getFineLevelComms()->send_raw_data(&status, sizeof(status), current_part, 247);\n }\n }\n\n getFineLevelComms()->send_vector_wait_all(my_num_iters);\n\n if (store_res_history)\n {\n getFineLevelComms()->send_vector_wait_all(res_history_tmp);\n }\n }\n else\n {\n // Store num_iters\n getFineLevelComms()->recv_vector(my_num_iters, m_my_fine_level_destination_part, 245);\n num_iters = my_num_iters[0];\n\n if (store_res_history)\n {\n // Fill res_history vector\n res_history.resize(num_iters + 1);\n res_history_tmp.resize( (num_iters + 1)*bsize);\n getFineLevelComms()->recv_vector(res_history_tmp, m_my_fine_level_destination_part, 246);\n\n for (int i = 0; i < num_iters + 1; i++)\n {\n res_history[i].resize(bsize);\n\n for (int j = 0; j < bsize; j++)\n {\n res_history[i][j] = res_history_tmp[i * bsize + j];\n }\n }\n }\n\n getFineLevelComms()->recv_raw_data(&status, sizeof(status), m_my_fine_level_destination_part, 247);\n }\n }\n}\n\ntemplate \nvoid DistributedManager >::consolidateAndUploadAll(int n, int nnz, int block_dimx, int block_dimy, const int *row_ptrs, const int *col_indices, const void *data, const void *diag, Matrix &in_A)\n{\n this->A = &in_A;\n this->createComms(this->A->getResources()); //refresh comms\n DistributedComms *comms = this->getComms();\n int my_id = comms->get_global_id();\n int num_parts = comms->get_num_partitions();\n int num_rings = this->num_halo_rings();\n int num_neighbors = this->neighbors.size();\n // All partitions have to call this, otherwise it fails\n // Step 1: Figure out which partition should be consolidated together based on their host_name and their PCI-E slot ID\n IVector_h destination_part(num_parts);\n this->computeDestinationPartitionsWithCons(my_id, num_parts, destination_part, comms);\n int my_destination_part = destination_part[my_id];\n // Check if I'm root partition and how many msgs I will receive\n bool is_root_partition = false;\n int num_parts_to_consolidate = 0;\n\n for (int i = 0; i < num_parts; i++)\n {\n if (destination_part[i] == my_id)\n {\n is_root_partition = true;\n num_parts_to_consolidate++;\n }\n }\n\n if (my_destination_part >= num_parts)\n {\n FatalError(\"During consolidation, sending data to partition that doesn't exist\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // Create cons_part_to_part map\n IVector_h cons_part_to_part = destination_part;\n amgx::thrust::sort(cons_part_to_part.begin(), cons_part_to_part.end());\n cudaCheckError();\n cons_part_to_part.erase(amgx::thrust::unique(cons_part_to_part.begin(), cons_part_to_part.end()), cons_part_to_part.end());\n cudaCheckError();\n int num_cons_partitions = cons_part_to_part.size();\n\n // If number of consolidated partitions is the same as number of partitions, simply call uploadAll\n if (num_cons_partitions == num_parts)\n {\n this->initializeUploadReorderAll(n, nnz, block_dimx, block_dimy, row_ptrs, col_indices, data, diag, *(this->A));\n this->m_is_fine_level_consolidated = false;\n return;\n }\n\n if (is_root_partition)\n {\n this->A->getResources()->expandRootPool();\n }\n\n this->m_is_fine_level_consolidated = true;\n\n if (num_rings != 1)\n {\n FatalError(\"num_rings > 1 not supported in fine_level consolidation\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // Fill with b2l_maps\n IVector_h_vector B2L_maps_tmp;\n B2L_maps_tmp.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n B2L_maps_tmp[i] = this->cached_B2L_maps[i];\n }\n\n bool useCudaIpc = this->m_use_cuda_ipc_consolidation;\n mat_value_type *data_hd = NULL;\n mat_value_type *diag_hd = NULL;\n int *col_indices_hd = NULL;\n int data_alloc = 0;\n int diag_alloc = 0;\n int col_alloc = 0;\n col_indices_hd = (int *) this->getDevicePointerForData((void *)col_indices, nnz * block_dimx * block_dimy * sizeof(int), &col_alloc);\n data_hd = (mat_value_type *) this->getDevicePointerForData((void *)data, nnz * block_dimx * block_dimy * sizeof(mat_value_type), &data_alloc);\n\n if (diag != NULL)\n {\n diag_hd = (mat_value_type *) this->getDevicePointerForData((void *)diag, nnz * block_dimx * block_dimy * sizeof(mat_value_type), &diag_alloc);\n }\n\n // Copy the original row_offsets array (this is required when replacing coefficients\n this->m_old_row_offsets_CONS.resize(n + 1);\n cudaMemcpy(this->m_old_row_offsets_CONS.raw(), row_ptrs, (n + 1)*sizeof(int), cudaMemcpyDefault);\n cudaCheckError();\n this->m_old_nnz_CONS = nnz;\n // This function:\n // Creates fine level consolidated matrices\n // Modifies the btl_maps, lth_maps\n // Create part_to_cons_part map\n IVector_h part_to_cons_part(num_parts);\n amgx::thrust::lower_bound(cons_part_to_part.begin(), cons_part_to_part.end(), destination_part.begin(), destination_part.end(), part_to_cons_part.begin());\n cudaCheckError();\n IVector_h neigh_to_part;\n this->createNeighToDestPartMap(neigh_to_part, this->neighbors, destination_part, num_neighbors);\n IVector_h cons_neigh_to_part;\n int num_cons_neighbors;\n this->createConsolidatedNeighToPartMap(cons_neigh_to_part, neigh_to_part, my_destination_part, destination_part, num_cons_neighbors);\n IVector_h neigh_to_cons_neigh;\n this->createNeighToConsNeigh( neigh_to_cons_neigh, cons_neigh_to_part, neigh_to_part, my_destination_part, num_neighbors);\n // ---------------------------------------\n // MERGE B2L MAPS BASED ON DEST PARTITION\n // ---------------------------------------\n IVector_h_vector dest_B2L_maps;\n this->consolidateB2Lmaps(dest_B2L_maps, B2L_maps_tmp, neigh_to_cons_neigh, num_cons_neighbors, num_neighbors);\n // ------------------------------------\n // Renumber interior and boundary rows\n // ------------------------------------\n int num_interior_rows;\n int num_boundary_rows;\n IVector_h renumbering;\n this->createAggregatesRenumbering(renumbering, dest_B2L_maps, n, num_cons_neighbors, num_interior_rows, num_boundary_rows, num_rings);\n // --------------------------------------------------\n // Create list of destination parts to consolidate\n // --------------------------------------------------\n // Store whether or not this is a root partition on fine level\n IVector_h parts_to_consolidate;\n parts_to_consolidate.resize(num_parts_to_consolidate);\n int count = 0;\n\n for (int i = 0; i < num_parts; i++)\n {\n if (destination_part[i] == my_id)\n {\n parts_to_consolidate[count] = i;\n count++;\n }\n }\n\n // ---------------------------------------------------------------------\n // Each partition computes its offset for its interior and boundary nodes\n // ---------------------------------------------------------------------\n IVector_h_vector vertex_counts;\n int interior_offset, boundary_offset, total_interior_rows_in_merged, total_boundary_rows_in_merged;\n int total_rows_in_merged;\n this->computeConsolidatedOffsets(my_id, my_destination_part, is_root_partition, num_interior_rows, num_boundary_rows, vertex_counts, parts_to_consolidate, num_parts_to_consolidate, interior_offset, boundary_offset, total_interior_rows_in_merged, total_boundary_rows_in_merged, total_rows_in_merged, comms);\n // -----------------------------------\n // Each partition renumber it's rows\n // -----------------------------------\n int total_num_halos = 0;\n\n // Pack new bdy_ids\n for (int i = 0; i < num_neighbors; i++)\n {\n total_num_halos += this->cached_L2H_maps[i].size();\n }\n\n IVector_h row_ids(n + total_num_halos, -1);\n this->m_row_ids_CONS.resize(n + total_num_halos);\n\n // Renumber the interior and boundary rows\n for (int i = 0; i < n; i++)\n {\n int new_id;\n\n if (renumbering.size() == 0)\n {\n new_id = i;\n }\n else\n {\n new_id = renumbering[i];\n }\n\n new_id += ((new_id >= num_interior_rows) ? boundary_offset : interior_offset);\n row_ids[i] = new_id;\n }\n\n for (int i = 0; i < num_cons_neighbors; i++)\n {\n thrust_wrapper::transform(dest_B2L_maps[i].begin(),\n dest_B2L_maps[i].end(),\n amgx::thrust::constant_iterator(boundary_offset),\n dest_B2L_maps[i].begin(),\n amgx::thrust::plus());\n }\n\n cudaCheckError();\n // -------------------------------------------------\n // Send dest_B2L_maps to root partitions\n // ------------------------------------------------\n IVector_h num_bdy_per_cons_neigh(num_cons_neighbors);\n\n for (int i = 0; i < num_cons_neighbors; i++)\n {\n num_bdy_per_cons_neigh[i] = dest_B2L_maps[i].size();\n }\n\n IVector_h root_cons_neighbors;\n int root_num_cons_neighbors = 0;\n IVector_h_vector cons_B2L_maps;\n this->consolidateB2LmapsOnRoot(root_num_cons_neighbors, cons_B2L_maps, root_cons_neighbors, dest_B2L_maps, cons_neigh_to_part, num_bdy_per_cons_neigh, parts_to_consolidate, num_parts_to_consolidate, my_id, my_destination_part, is_root_partition, num_cons_neighbors, comms);\n IVector_h halo_ids_offsets(num_neighbors + 1);\n IVector_h halo_ids;\n int halo_ids_size = 0;\n halo_ids_offsets[0] = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n halo_ids_size += this->cached_L2H_maps[i].size();\n halo_ids_offsets[i + 1] = halo_ids_size;\n }\n\n halo_ids.resize(halo_ids_size);\n // Do exchange with neighbors\n // Pack new bdy_ids\n IVector_h_vector bdy_ids;\n bdy_ids.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = this->cached_B2L_maps[i].size();\n bdy_ids[i].resize(size);\n\n // Pack buffer\n for (int j = 0; j < size; j++)\n {\n bdy_ids[i][j] = row_ids[this->cached_B2L_maps[i][j]];\n }\n }\n\n for (int i = 0; i < num_neighbors; i++)\n {\n comms->send_vector_async(bdy_ids[i], this->neighbors[i], 6666 + this->neighbors[i]);\n }\n\n for (int i = 0; i < num_neighbors; i++)\n {\n comms->recv_vector(halo_ids, this->neighbors[i], 6666 + my_id, halo_ids_offsets[i], this->cached_L2H_maps[i].size());\n }\n\n for (int i = 0; i < num_neighbors; i++)\n {\n comms->send_vector_wait_all(bdy_ids[i]);\n }\n\n IVector_h halo_offsets(root_num_cons_neighbors + 1, 0);\n int root_num_rows;\n this->consolidateAndRenumberHalos(halo_ids, halo_ids_offsets, halo_offsets, this->neighbors, num_neighbors, root_cons_neighbors, root_num_cons_neighbors, destination_part, my_destination_part, is_root_partition, parts_to_consolidate, num_parts_to_consolidate, num_parts, my_id, total_rows_in_merged, root_num_rows, comms);\n\n if (is_root_partition)\n {\n this->B2L_maps.resize(cons_B2L_maps.size());\n\n for (int i = 0; i < cons_B2L_maps.size(); i++)\n {\n amgx::thrust::sort(cons_B2L_maps[i].begin(), cons_B2L_maps[i].end());\n this->B2L_maps[i].copy(cons_B2L_maps[i]); // H2D copy of B2L maps\n }\n\n cudaCheckError();\n }\n\n // Now renumber the row_ids based on lth_maps\n count = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n for (int j = 0; j < this->cached_L2H_maps[i].size(); j++)\n {\n row_ids[this->cached_L2H_maps[i][j]] = halo_ids[count];\n count++;\n }\n }\n\n cudaMemcpy(this->m_row_ids_CONS.raw(), row_ids.raw(), (n + total_num_halos)*sizeof(int), cudaMemcpyDefault);\n cudaCheckError();\n int bsize = block_dimx * block_dimy;\n\n if (is_root_partition)\n {\n this->A->row_offsets.resize(root_num_rows + 1);\n }\n\n void *root_row_ptr = (void *) this->A->row_offsets.raw();\n\n if (useCudaIpc)\n {\n // ----------------------------------------------------\n // 1. cudaIPC to get pointer to root's row_offset array\n // ----------------------------------------------------\n this->ipcExchangePtr(root_row_ptr, is_root_partition, num_parts_to_consolidate, parts_to_consolidate, my_destination_part, my_id, comms);\n cudaCheckError();\n // -------------------------------------------------------------------\n // 2. each rank copy it's row length on root partition using row_ids\n // -------------------------------------------------------------------\n int cta_size = 128;\n int grid_size = std::min(4096, (n + total_num_halos + cta_size - 1) / cta_size);\n zero_copy_row_lengths_ids_offsets <<< grid_size, cta_size>>>(this->m_old_row_offsets_CONS.raw(), ((int *) root_row_ptr) /* IPC */, this->m_row_ids_CONS.raw(), n, total_num_halos, (mat_value_type *) diag);\n cudaCheckError();\n // Root partition waits for children to be done writing their result\n this->ipcWaitForChildren(is_root_partition, num_parts_to_consolidate, parts_to_consolidate, my_destination_part, my_id, comms);\n cudaCheckError();\n }\n else // CudaIpcNotAvailable\n {\n this->checkPinnedBuffer( std::max( nnz * sizeof(mat_value_type), (n + 1)*std::max(sizeof(index_type), sizeof(value_type)) ) );\n\n if (!is_root_partition)\n {\n IVector_h data_to_send(3);\n data_to_send[0] = n;\n data_to_send[1] = nnz;\n data_to_send[2] = total_num_halos;\n int dummy;\n void *row_ptrs_to_send = this->getHostPointerForData((void *)row_ptrs, sizeof(index_type) * (n + 1), &dummy);\n comms->send_vector(data_to_send, my_destination_part, 10000 + my_id);\n comms->send_raw_data(row_ptrs_to_send, (n + 1)*sizeof(int), my_destination_part, 10001 + my_id);\n comms->send_raw_data(&row_ids[0], (n + total_num_halos)*sizeof(int), my_destination_part, 10002 + my_id);\n }\n else\n {\n cudaEvent_t event;\n cudaEventCreate(&event);\n cudaCheckError();\n //TODO: Could use streams here\n //TODO: Avoid extra device to host copies\n std::vector data_recv(num_parts_to_consolidate);\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n data_recv[i].resize(3);\n int current_part = parts_to_consolidate[i];\n\n if (current_part != my_id)\n {\n comms->recv_vector(data_recv[i], current_part, 10000 + current_part);\n }\n else\n {\n data_recv[i][0] = n;\n data_recv[i][1] = nnz;\n data_recv[i][2] = total_num_halos;\n }\n }\n\n this->m_child_n.resize(num_parts_to_consolidate);\n this->m_child_nnz.resize(num_parts_to_consolidate);\n this->m_child_num_halos.resize(num_parts_to_consolidate);\n this->m_child_row_ids.resize(num_parts_to_consolidate);\n this->m_child_old_row_offsets.resize(num_parts_to_consolidate);\n int max_n = 0;\n int max_nnz = 0;\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n int current_part = parts_to_consolidate[i];\n this->m_child_n[i] = data_recv[i][0];\n this->m_child_nnz[i] = data_recv[i][1];\n this->m_child_num_halos[i] = data_recv[i][2];\n\n if (this->m_child_n[i] > max_n) { max_n = this->m_child_n[i]; }\n\n if (this->m_child_nnz[i] > max_nnz) { max_nnz = this->m_child_nnz[i]; }\n\n this->m_child_row_ids[i].resize(this->m_child_n[i] + this->m_child_num_halos[i]);\n this->m_child_old_row_offsets[i].resize(this->m_child_n[i] + 1);\n }\n\n this->m_child_max_n = max_n;\n this->m_child_max_nnz = max_nnz;\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n int current_part = parts_to_consolidate[i];\n int cta_size = 128;\n int grid_size = std::min(4096, (this->m_child_n[i] + this->m_child_num_halos[i] + cta_size - 1) / cta_size);\n\n if (current_part != my_id)\n {\n comms->recv_vector(this->m_child_old_row_offsets[i], current_part, 10001 + current_part, 0, this->m_child_n[i] + 1);\n comms->recv_vector(this->m_child_row_ids[i], current_part, 10002 + current_part, 0, this->m_child_n[i] + this->m_child_num_halos[i]);\n zero_copy_row_lengths_ids_offsets <<< grid_size, cta_size>>>(this->m_child_old_row_offsets[i].raw(), this->A->row_offsets.raw(), this->m_child_row_ids[i].raw(), this->m_child_n[i], this->m_child_num_halos[i], (mat_value_type *) diag);\n cudaCheckError();\n // Wait for kernel to finish before overwriting host buffer\n cudaEventRecord(event);\n cudaCheckError();\n cudaEventSynchronize(event);\n cudaCheckError();\n }\n else\n {\n zero_copy_row_lengths_ids_offsets <<< grid_size, cta_size>>>(this->m_old_row_offsets_CONS.raw(), this->A->row_offsets.raw(), this->m_row_ids_CONS.raw(), n, total_num_halos, (mat_value_type *) diag);\n cudaCheckError();\n cudaEventRecord(event);\n cudaCheckError();\n cudaEventSynchronize(event);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n cudaEventDestroy(event);\n cudaCheckError();\n } // If root partition\n\n //TODO: is this necessary\n comms->barrier();\n }\n\n //3. root does a exclusive_scan\n if (is_root_partition)\n {\n cudaEvent_t event;\n cudaEventCreate(&event);\n cudaCheckError();\n // Populate the halo rows with diagonal, increase the length of the halo rows\n thrust_wrapper::fill(this->A->row_offsets.begin() + halo_offsets[0], this->A->row_offsets.begin() + halo_offsets[root_num_cons_neighbors], 1);\n thrust_wrapper::exclusive_scan(this->A->row_offsets.begin(), this->A->row_offsets.end(), this->A->row_offsets.begin());\n cudaEventRecord(event);\n cudaCheckError();\n cudaEventSynchronize(event);\n cudaCheckError();\n this->A->set_initialized(0);\n this->A->delProps(DIAG); // We always insert the diagonal\n this->A->delProps(COO); // No COO\n this->A->setColsReorderedByColor(false); // Cols not reordered by color\n int nnz = this->A->row_offsets[root_num_rows]; // This is a device to host copy\n this->A->resize(root_num_rows, root_num_rows, nnz, block_dimx, block_dimy);\n this->A->set_num_nz(nnz); // num_nz doesn't include halo rows\n //this->A->set_initialized(1);\n cudaEventDestroy(event);\n cudaCheckError();\n }\n else\n {\n this->A->set_initialized(0);\n this->A->resize( 0, 0, 0, block_dimx, block_dimy );\n this->A->delProps(DIAG); // We always insert the diagonal\n this->A->delProps(COO); // No COO\n this->A->setColsReorderedByColor(false); // Cols not reordered by color\n //this->A->set_initialized(1);\n }\n\n if (useCudaIpc)\n {\n // ----------------------------------------------\n // 4. Do ipc consolidation of values and columns\n // ----------------------------------------------\n // Child partition waits for parent to create row_offsets\n this->ipcWaitForRoot(is_root_partition, num_parts_to_consolidate, parts_to_consolidate, my_destination_part, my_id, comms);\n void *root_col_ptr = (void *) this->A->col_indices.raw();\n void *root_val_ptr = (void *) this->A->values.raw();\n this->ipcExchangePtr(root_col_ptr, is_root_partition, num_parts_to_consolidate, parts_to_consolidate, my_destination_part, my_id, comms);\n this->ipcExchangePtr(root_val_ptr, is_root_partition, num_parts_to_consolidate, parts_to_consolidate, my_destination_part, my_id, comms);\n int cta_size2 = 128;\n int grid_size2 = std::min(4096, (n + cta_size2 - 1) / cta_size2);\n ipc_consolidation_upload_matrix <<< grid_size2, cta_size2>>>(n, this->m_row_ids_CONS.raw(), this->m_old_row_offsets_CONS.raw(), ( (int *) root_row_ptr ) /*IPC*/, col_indices_hd, ( (int *) root_col_ptr) /*IPC*/, data_hd, ( (mat_value_type *) root_val_ptr ) /*IPC*/, diag_hd, bsize);\n cudaCheckError();\n // Root partition waits for children to upload their matrices\n this->ipcWaitForChildren(is_root_partition, num_parts_to_consolidate, parts_to_consolidate, my_destination_part, my_id, comms);\n cudaCheckError();\n\n // Child partitions close their mem handle (they are done upload data)\n if (!is_root_partition)\n {\n cudaIpcCloseMemHandle(root_row_ptr);\n cudaIpcCloseMemHandle(root_val_ptr);\n cudaIpcCloseMemHandle(root_col_ptr);\n }\n }\n else // If cudaIpcNotAvailable\n {\n if (!is_root_partition)\n {\n int dummy;\n void *col_indices_to_send = this->getHostPointerForData((void *)col_indices, sizeof(index_type) * nnz, &dummy);\n comms->send_raw_data(col_indices_to_send, nnz * sizeof(int), my_destination_part, 10000 + my_id);\n void *data_to_send = this->getHostPointerForData((void *)data, sizeof(mat_value_type) * nnz, &dummy);\n comms->send_raw_data(data_to_send, nnz * bsize * sizeof(mat_value_type), my_destination_part, 10001 + my_id);\n\n if (diag != NULL)\n {\n void *diag_to_send = this->getHostPointerForData((void *)diag, sizeof(mat_value_type) * n, &dummy);\n comms->send_raw_data(diag_to_send, n * bsize * sizeof(mat_value_type), my_destination_part, 10002 + my_id);\n }\n }\n else\n {\n cudaEvent_t event;\n cudaEventCreate(&event);\n cudaCheckError();\n //TODO: Could use streams here\n int *child_col_indices;\n mat_value_type *child_data;\n mat_value_type *child_diag = NULL;\n cudaHostAlloc( (void **) &child_col_indices, this->m_child_max_nnz * sizeof(int), cudaHostAllocMapped);\n cudaHostAlloc( (void **) &child_data, this->m_child_max_nnz * bsize * sizeof(mat_value_type), cudaHostAllocMapped);\n\n if (diag != NULL)\n {\n cudaHostAlloc( (void **) &child_diag, (this->m_child_max_n)*bsize * sizeof(mat_value_type), cudaHostAllocMapped);\n }\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n int current_part = parts_to_consolidate[i];\n int cta_size2 = 128;\n int grid_size2 = std::min(4096, (this->m_child_n[i] + cta_size2 - 1) / cta_size2);\n\n if (current_part != my_id)\n {\n comms->recv_raw_data(child_col_indices, this->m_child_nnz[i]*sizeof(int), current_part, 10000 + current_part);\n comms->recv_raw_data(child_data, this->m_child_nnz[i]*bsize * sizeof(mat_value_type), current_part, 10001 + current_part);\n\n if (diag != NULL)\n {\n comms->recv_raw_data(child_diag, this->m_child_n[i]*bsize * sizeof(mat_value_type), current_part, 10002 + current_part);\n }\n\n int *child_col_indices_hd;\n mat_value_type *child_data_hd;\n mat_value_type *child_diag_hd = NULL;\n cudaHostGetDevicePointer(&child_col_indices_hd, child_col_indices, 0);\n cudaHostGetDevicePointer(&child_data_hd, child_data, 0);\n\n if (diag != NULL)\n {\n cudaHostGetDevicePointer(&child_diag_hd, child_diag, 0);\n }\n\n ipc_consolidation_upload_matrix <<< grid_size2, cta_size2>>>(this->m_child_n[i], this->m_child_row_ids[i].raw(), this->m_child_old_row_offsets[i].raw(), this->A->row_offsets.raw(), child_col_indices_hd, this->A->col_indices.raw(), child_data_hd, this->A->values.raw(), child_diag_hd, bsize);\n cudaCheckError();\n // Wait for kernel to finish before overwriting host buffer\n cudaEventRecord(event);\n cudaCheckError();\n cudaEventSynchronize(event);\n cudaCheckError();\n }\n else\n {\n ipc_consolidation_upload_matrix <<< grid_size2, cta_size2>>>(n, this->m_row_ids_CONS.raw(), this->m_old_row_offsets_CONS.raw(), this->A->row_offsets.raw(), col_indices_hd, this->A->col_indices.raw(), data_hd, this->A->values.raw(), diag_hd, bsize);\n cudaCheckError();\n cudaEventRecord(event);\n cudaCheckError();\n cudaEventSynchronize(event);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n cudaEventDestroy(event);\n cudaCheckError();\n amgx::memory::cudaFreeHost(child_col_indices);\n amgx::memory::cudaFreeHost(child_data);\n\n if (diag != NULL)\n {\n amgx::memory::cudaFreeHost(child_diag);\n }\n } // If root partition\n\n //TODO: is this necessary\n comms->barrier();\n }\n\n // Store the original fine level communicator\n this->m_is_fine_level_root_partition = is_root_partition;\n this->m_my_fine_level_destination_part = my_destination_part;\n // Create a clone of the original communicator\n this->m_fine_level_comms = comms; //this->_comms is the same pointer that this->m_fine_level_comms right now, so we can overwrite this->_comms, but make sure that we release m_fine_level_cons\n this->_comms = this->m_fine_level_comms->CloneSubComm(cons_part_to_part, is_root_partition); // this->_comms will be empty comm for non-root partition and new comm for root ranks only if root partition\n this->m_fine_level_id = my_id;\n\n if (is_root_partition)\n {\n int cta_size = 128;\n int grid_size3 = std::min(4096, ( (root_num_rows - halo_offsets[0]) + cta_size - 1) / cta_size);\n\n if (grid_size3 != 0)\n {\n set_halo_cols_values <<< grid_size3, cta_size>>>(this->A->row_offsets.raw(), this->A->col_indices.raw(), this->A->values.raw(), halo_offsets[0], root_num_rows, bsize);\n cudaCheckError();\n }\n\n int my_cons_id = part_to_cons_part[my_id];\n this->_global_id = my_cons_id;\n this->_num_interior_nodes = total_interior_rows_in_merged;\n this->_num_boundary_nodes = total_boundary_rows_in_merged;\n\n for (int i = 0; i < root_num_cons_neighbors; i++)\n {\n root_cons_neighbors[i] = part_to_cons_part[root_cons_neighbors[i]];\n }\n\n this->_comms->set_neighbors(root_num_cons_neighbors);\n this->neighbors = root_cons_neighbors;\n this->halo_offsets = halo_offsets; // H2D copy of halo offsets\n this->m_num_fine_level_parts_to_consolidate = num_parts_to_consolidate;\n this->m_fine_level_parts_to_consolidate = parts_to_consolidate;\n this->set_num_halo_rings(num_rings);\n this->set_num_halo_rows(halo_offsets[root_num_cons_neighbors] - halo_offsets[0]);\n // B2L_maps has already been copied\n this->B2L_rings.resize(root_num_cons_neighbors);\n\n for (int i = 0; i < root_num_cons_neighbors; i++)\n {\n this->B2L_rings[i].resize(2);\n this->B2L_rings[i][0] = 0;\n this->B2L_rings[i][1] = cons_B2L_maps[i].size();\n }\n\n this->set_initialized(this->A->row_offsets);\n this->A->set_initialized(0);\n this->A->delProps(DIAG);\n this->A->diag.resize(root_num_rows);\n this->A->computeDiagonal(); //\n this->A->setView(OWNED);\n cudaEventCreate(&(this->comm_event));\n cudaCheckError();\n this->A->set_initialized(1);\n }\n else\n {\n this->neighbors.resize(0);\n this->halo_offsets.resize(0);\n }\n\n /* free memory (if needed) */\n if (col_alloc) { amgx::memory::cudaFreeAsync(col_indices_hd); }\n if (data_alloc) { amgx::memory::cudaFreeAsync(data_hd); }\n if (diag_alloc) { amgx::memory::cudaFreeAsync(diag_hd); }\n}\n\ntemplate \nvoid DistributedManager >::replaceMatrixCoefficientsNoCons(int n, int nnz, const mat_value_type *data_pinned, const mat_value_type *diag_pinned)\n{\n //matrix parameters\n //int num_nnz = this->A->get_num_nz();\n int num_rows = this->halo_offsets[0];\n int total_rows = num_rows + this->num_halo_rows();\n int block_size = this->A->get_block_size();\n mat_value_type *data_hd = NULL;\n mat_value_type *diag_hd = NULL;\n int data_alloc = 0;\n int diag_alloc = 0;\n //cuda parameters\n int num_blocks = std::min(4096, (num_rows + 127) / 128);\n\n /* WARNING: the number of non-zero elements (nnz) in the array data_pinned and A->values (num_nnz) might be different at this point.\n 1. If the matrix has CSR property and therefore diagonal is included in the matrix this values will be the same.\n 2. If the matrix has DIAG property and therefore diagonal is originally stored separately, and later appended to the array\n of values, and subsequently inserted into the matrix than num_nnz = nnz + n. We have to account for this fact when replacing the\n coefficients (and use nnz not num_nnz).\n obs.: see calls to computeDiagonal (matrix.cu), AMGX_matrix_upload and AMGX_replace_coefficients (amgx_c.cu), and\n uploadMatrix and replaceMatrixCoefficients[No|With]Cons (distributed_manager.cu) for details. */\n\n /* check early exit */\n if ((this->neighbors.size() == 0 || this->renumbering.size() == 0) && !this->m_is_fine_level_glued)\n {\n return;\n }\n\n cudaCheckError();\n /* allocate if data and diag if they are not pinned */\n data_hd = (mat_value_type *) this->getDevicePointerForData((void *)data_pinned, nnz * block_size * sizeof(mat_value_type), &data_alloc);\n\n if (diag_pinned != NULL)\n {\n diag_hd = (mat_value_type *) this->getDevicePointerForData((void *)diag_pinned, num_rows * block_size * sizeof(mat_value_type), &diag_alloc);\n }\n\n /* replace the values (reordering them if needed) */\n if (insertDiagonals && diag_pinned != NULL)\n {\n replace_values_matrix <32> <<< num_blocks, 512>>>(data_hd, diag_hd, this->old_row_offsets.raw(), this->A->row_offsets.raw(), this->A->values.raw(), this->renumbering.raw(), block_size, num_rows);\n cudaCheckError();\n }\n else\n {\n replace_values_matrix <32> <<< num_blocks, 512>>>(data_hd, this->old_row_offsets.raw(), this->A->row_offsets.raw(), this->A->values.raw(), this->renumbering.raw(), block_size, num_rows);\n cudaCheckError();\n\n if (diag_pinned != NULL)\n {\n reorder_vector_values <<< num_blocks, 512>>>(this->A->values.raw() + this->A->row_offsets[total_rows]*block_size, diag_hd, this->renumbering.raw(), block_size, num_rows);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n\n /* free memory (if needed) */\n if (data_alloc) { amgx::memory::cudaFreeAsync(data_hd); }\n if (diag_alloc) { amgx::memory::cudaFreeAsync(diag_hd); }\n}\n\ntemplate \nvoid DistributedManager >::replaceMatrixCoefficientsWithCons(int n, int nnz, const mat_value_type *data_pinned, const mat_value_type *diag_pinned)\n{\n //matrix parameters\n //int num_nnz = this->A->get_num_nz();\n /* WARNING: in consolidation, for non-root partitions, halo_offsets\n might be NULL due to the call halo_offsets.resize(0); at the end\n of the routine uploadMatrix->consolidateAndUploadAll. We should\n use the parameter n instead this->halo_offsets[0] for num_rows. */\n int num_rows = n;\n int block_size = this->A->get_block_size();\n mat_value_type *data_hd = NULL;\n mat_value_type *diag_hd = NULL;\n int data_alloc = 0;\n int diag_alloc = 0;\n data_hd = (mat_value_type *) this->getDevicePointerForData((void *)data_pinned, nnz * block_size * sizeof(mat_value_type), &data_alloc);\n\n if (diag_pinned != NULL)\n {\n diag_hd = (mat_value_type *) this->getDevicePointerForData((void *)diag_pinned, num_rows * block_size * sizeof(mat_value_type), &diag_alloc);\n }\n\n bool useCudaIpc = this->m_use_cuda_ipc_consolidation;\n\n if (useCudaIpc)\n {\n // Child partitions wait for root to be done\n this->ipcWaitForRoot(this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n cudaCheckError();\n void *root_row_ptr = (void *) this->A->row_offsets.raw();\n void *root_val_ptr = (void *) this->A->values.raw();\n this->ipcExchangePtr(root_row_ptr, this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n this->ipcExchangePtr(root_val_ptr, this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n // replace the values, insert the diagonal\n int ncons = this->m_old_row_offsets_CONS.size() - 1;\n int cta_size = 128;\n int grid_size2 = std::min(4096, (ncons + cta_size - 1) / cta_size);\n ipc_consolidation_replace_values <<< grid_size2, cta_size>>>(ncons, this->m_row_ids_CONS.raw(), this->m_old_row_offsets_CONS.raw(), ( (int *) root_row_ptr )/*IPC*/, data_hd, ( (mat_value_type *) root_val_ptr )/*IPC*/, diag_hd, this->A->get_block_size() );\n cudaCheckError();\n // Root partition wait for child to be done replacing their values\n this->ipcWaitForChildren(this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n cudaCheckError();\n\n if (!this->m_is_fine_level_root_partition)\n {\n cudaIpcCloseMemHandle(root_row_ptr);\n cudaIpcCloseMemHandle(root_val_ptr);\n }\n }\n else // cudaIpcNotAvailable\n {\n if (this->m_is_fine_level_consolidated) // aggregation\n {\n int bsize = this->A->get_block_size();\n int ncons = this->m_old_row_offsets_CONS.size() - 1;\n\n if (!this->m_is_fine_level_root_partition)\n {\n int dummy;\n int nnzcons = this->m_old_nnz_CONS;\n void *data_to_send = this->getHostPointerForData((void *)data_pinned, nnzcons * bsize * sizeof(mat_value_type), &dummy);\n this->getFineLevelComms()->send_raw_data(data_to_send, nnzcons * bsize * sizeof(mat_value_type), this->m_my_fine_level_destination_part, 10001 + this->fine_level_id());\n\n if (diag_pinned != NULL)\n {\n void *diag_to_send = this->getHostPointerForData((void *)diag_pinned, ncons * bsize * sizeof(mat_value_type), &dummy);\n this->getFineLevelComms()->send_raw_data(diag_to_send, ncons * bsize * sizeof(mat_value_type), this->m_my_fine_level_destination_part, 10002 + this->fine_level_id());\n }\n }\n else\n {\n cudaEvent_t event;\n cudaEventCreate(&event);\n cudaCheckError();\n //TODO: Could use streams here\n mat_value_type *child_data;\n mat_value_type *child_diag = NULL;\n cudaHostAlloc( (void **) &child_data, this->m_child_max_nnz * bsize * sizeof(mat_value_type), cudaHostAllocMapped);\n\n if (diag_pinned != NULL)\n {\n cudaHostAlloc( (void **) &child_diag, (this->m_child_max_n)*bsize * sizeof(mat_value_type), cudaHostAllocMapped);\n }\n\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n int cta_size2 = 128;\n int grid_size2 = std::min(4096, (this->m_child_n[i] + cta_size2 - 1) / cta_size2);\n\n if (current_part != this->fine_level_id())\n {\n this->getFineLevelComms()->recv_raw_data(child_data, this->m_child_nnz[i]*bsize * sizeof(mat_value_type), current_part, 10001 + current_part);\n\n if (diag_pinned != NULL)\n {\n this->getFineLevelComms()->recv_raw_data(child_diag, this->m_child_n[i]*bsize * sizeof(mat_value_type), current_part, 10002 + current_part);\n }\n\n mat_value_type *child_data_hd;\n mat_value_type *child_diag_hd = NULL;\n cudaHostGetDevicePointer(&child_data_hd, child_data, 0);\n\n if (diag_pinned != NULL)\n {\n cudaHostGetDevicePointer(&child_diag_hd, child_diag, 0);\n }\n\n ipc_consolidation_replace_values <<< grid_size2, cta_size2>>>(this->m_child_n[i], this->m_child_row_ids[i].raw(), this->m_child_old_row_offsets[i].raw(), this->A->row_offsets.raw(), child_data_hd, this->A->values.raw(), child_diag_hd, bsize);\n cudaCheckError();\n // Wait for kernel to finish before overwriting host buffer\n cudaEventRecord(event);\n cudaCheckError();\n cudaEventSynchronize(event);\n cudaCheckError();\n }\n else\n {\n ipc_consolidation_replace_values <<< grid_size2, cta_size2>>>(ncons, this->m_row_ids_CONS.raw(), this->m_old_row_offsets_CONS.raw(), this->A->row_offsets.raw(), data_hd, this->A->values.raw(), diag_hd, bsize);\n cudaCheckError();\n //cudaEventRecord(event);\n cudaCheckError();\n //cudaEventSynchronize(event);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n cudaEventDestroy(event);\n cudaCheckError();\n amgx::memory::cudaFreeHost(child_data);\n\n if (diag_pinned != NULL)\n {\n amgx::memory::cudaFreeHost(child_diag);\n }\n } // If root partition\n\n //TODO: is this necessary\n this->getFineLevelComms()->barrier();\n } //agg\n else if (this->m_is_fine_level_glued) // classical\n {\n int bsize = this->A->get_block_size();\n int ncons = this->m_old_row_offsets_CONS.size() - 1;\n IVector_h nnz_off;\n nnz_off.resize(this->getConsolidationArrayOffsets().size());\n IVector_h nnz_array;\n nnz_array.resize(this->getConsolidationArrayOffsets().size() - 1);\n this->getFineLevelComms()->all_gather( nnz,\n nnz_array,\n this->getConsolidationArrayOffsets().size() - 1);\n nnz_off[0] = 0;\n\n for (int i = 0; i < nnz_array.size(); i++)\n {\n nnz_off[i + 1] = nnz_off[i] + nnz_array[i];\n }\n\n if (!this->m_is_fine_level_root_partition)\n {\n int dummy;\n void *data_to_send = this->getHostPointerForData((void *)data_pinned, nnz * bsize * sizeof(mat_value_type), &dummy);\n this->getFineLevelComms()->send_raw_data( data_to_send,\n nnz * bsize * sizeof(mat_value_type),\n this->m_my_fine_level_destination_part,\n 10001 + this->fine_level_id());\n\n if (diag_pinned != NULL)\n {\n void *diag_to_send = this->getHostPointerForData((void *)diag_pinned, n * sizeof(mat_value_type), &dummy);\n this->getFineLevelComms()->send_raw_data( diag_to_send,\n n * bsize * sizeof(mat_value_type),\n this->m_my_fine_level_destination_part,\n 10002 + this->fine_level_id());\n //diag.resize(0);\n cudaCheckError();\n }\n\n //values.resize(0);\n cudaCheckError();\n }\n else\n {\n //TODO: Could use streams here\n mat_value_type *child_data;\n mat_value_type *child_diag = NULL;\n // Assumes partions have been glued already\n this->A->getNnzForView(OWNED, &nnz);\n cudaHostAlloc( (void **) &child_data, nnz * bsize * sizeof(mat_value_type), cudaHostAllocMapped);\n\n if (diag_pinned != NULL)\n {\n cudaHostAlloc( (void **) &child_diag, this->halo_offsets[this->neighbors.size()]*bsize * sizeof(mat_value_type), cudaHostAllocMapped);\n }\n\n // roots copy their data\n memcpy ( &child_data[0], data_pinned, nnz_array[this->fine_level_id()]*sizeof(value_type));\n\n if (diag_pinned != NULL)\n {\n memcpy ( &child_diag[0], diag_pinned, n * sizeof(value_type));\n }\n\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n int current_offset = nnz_off[current_part] - nnz_off[this->fine_level_id()] ;\n int current_nnz = nnz_array[current_part];\n\n if (current_part != this->fine_level_id())\n {\n this->getFineLevelComms()->recv_raw_data( &child_data[current_offset],\n current_nnz * bsize * sizeof(mat_value_type),\n current_part,\n 10001 + current_part);\n\n if (diag_pinned != NULL)\n this->getFineLevelComms()->recv_raw_data( &child_diag[this->getConsolidationArrayOffsets()[current_part] - this->getConsolidationArrayOffsets()[this->fine_level_id()]],\n (this->getConsolidationArrayOffsets()[current_part + 1] - this->getConsolidationArrayOffsets()[current_part])*bsize * sizeof(mat_value_type),\n current_part,\n 10002 + current_part);\n }\n }\n\n cudaCheckError();\n // we can follow the usual upload path for raw data now\n // Assumes partions have been glued already\n int os;\n this->A->getOffsetAndSizeForView(OWNED, &os, &n);\n replaceMatrixCoefficientsNoCons( n, nnz, child_data, child_diag);\n cudaCheckError();\n amgx::memory::cudaFreeHost(child_data);\n\n if (diag_pinned != NULL)\n {\n amgx::memory::cudaFreeHost(child_diag);\n }\n } // If root partition\n\n //TODO: is this necessary\n this->getFineLevelComms()->barrier();\n } // cla\n } // not ipc\n\n this->A->setView(OWNED);\n\n /* free memory (if needed) */\n if (data_alloc) { amgx::memory::cudaFreeAsync(data_hd); }\n\n if (diag_alloc) { amgx::memory::cudaFreeAsync(diag_hd); }\n}\n\ntemplate \nvoid DistributedManager >::transformAndUploadVector(VVector_v &v, const void *data, int n, int block_dim)\n{\n if (this->isFineLevelConsolidated() || (this->isFineLevelGlued() && !this->isGlued()))\n {\n transformAndUploadVectorWithCons(v, data, n, block_dim);\n }\n else\n {\n v.resize(n * block_dim);\n cudaCheckError();\n // Upload on host\n cudaMemcpy(v.raw(), (value_type *)data, n * block_dim * sizeof(value_type), cudaMemcpyDefault);\n cudaCheckError();\n // Permute based on renumbering vector\n transformVector(v);\n int tag = 0;\n // Exchange halos\n this->exchange_halo(v, tag);\n }\n}\n\ntemplate \nvoid DistributedManager >::transformAndUploadVectorWithCons(VVector_v &v, const void *data_pinned, int n, int block_dim)\n{\n if (v.get_block_size() != this->A->get_block_dimx()) { printf(\"Blocksize mismatch!\\n\"); }\n\n bool useCudaIpc = this->m_use_cuda_ipc_consolidation;\n this->getFineLevelComms()->barrier();\n void *root_temp_ptr = NULL;\n VVector_v temp;\n\n if (this->m_is_fine_level_root_partition && !this->m_is_fine_level_glued )\n {\n temp.resize(this->halo_offsets[this->neighbors.size()]*v.get_block_size(), types::util::get_zero());\n temp.set_block_dimx(v.get_block_dimx());\n temp.set_block_dimy(v.get_block_dimy());\n root_temp_ptr = (void *) temp.raw();\n }\n\n cudaCheckError();\n int data_alloc = 0;\n value_type *data_hd = NULL;\n\n if (!this->m_is_fine_level_glued )\n {\n data_hd = (value_type *) this->getDevicePointerForData((void *)data_pinned, n * block_dim * sizeof(value_type), &data_alloc);\n }\n\n if (useCudaIpc)\n {\n // Do IPC\n this->ipcExchangePtr(root_temp_ptr, this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n cudaCheckError();\n int num_blocks = std::min(4096, (n + 511) / 512);\n reorder_vector_values <<< num_blocks, 512>>>( (value_type *) root_temp_ptr, data_hd, this->m_row_ids_CONS.raw(), v.get_block_size(), n);\n cudaCheckError();\n // Root partition waits for children to be done\n this->ipcWaitForChildren(this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n cudaCheckError();\n\n if (!this->m_is_fine_level_root_partition)\n {\n cudaIpcCloseMemHandle(root_temp_ptr);\n }\n }\n else // If cudaIpcNotAvail\n {\n if (this->m_is_fine_level_consolidated) // aggregation\n {\n // Exchange the vector between root and child\n if (!this->m_is_fine_level_root_partition)\n {\n IVector_h size(1);\n size[0] = n;\n this->getFineLevelComms()->send_vector(size, this->m_my_fine_level_destination_part, 20000 + this->fine_level_id());\n int dummy;\n void *data_to_send = this->getHostPointerForData((void *)data_pinned, n * v.get_block_size() * sizeof(value_type), &dummy);\n this->getFineLevelComms()->send_raw_data(data_to_send, n * v.get_block_size()*sizeof(value_type), this->m_my_fine_level_destination_part, 20001 + this->fine_level_id());\n }\n else\n {\n cudaEvent_t event;\n cudaEventCreate(&event);\n cudaCheckError();\n IVector_h child_n(this->m_num_fine_level_parts_to_consolidate);\n int max_n = 0;\n\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n\n if (current_part != this->fine_level_id())\n {\n this->getFineLevelComms()->recv_vector(child_n, current_part, 20000 + current_part, i, 1);\n }\n else\n {\n child_n[i] = n;\n }\n\n if (child_n[i] > max_n) { max_n = child_n[i]; }\n }\n\n value_type *child_data;\n cudaHostAlloc( (void **) &child_data, max_n * v.get_block_size()*sizeof(value_type), cudaHostAllocMapped);\n value_type *child_data_hd;\n cudaHostGetDevicePointer(&child_data_hd, child_data, 0);\n\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n int num_blocks = std::min(4096, (child_n[i] + 511) / 512);\n\n if (current_part != this->fine_level_id())\n {\n this->getFineLevelComms()->recv_raw_data(&child_data[0], child_n[i]*v.get_block_size()*sizeof(value_type), current_part, 20001 + current_part);\n reorder_vector_values <<< num_blocks, 512>>>( (value_type *) root_temp_ptr, child_data_hd, this->m_child_row_ids[i].raw(), v.get_block_size(), child_n[i]);\n cudaCheckError();\n cudaEventRecord(event);\n cudaCheckError();\n cudaEventSynchronize(event);\n cudaCheckError();\n }\n else\n {\n reorder_vector_values <<< num_blocks, 512>>>( (value_type *) root_temp_ptr, data_hd, this->m_row_ids_CONS.raw(), v.get_block_size(), n);\n cudaCheckError();\n }\n } // Loop over parts to consolidate\n\n cudaCheckError();\n cudaEventDestroy(event);\n cudaCheckError();\n amgx::memory::cudaFreeHost(child_data);\n } // If root partition\n } //agg\n else if (this->m_is_fine_level_glued) // cla\n {\n value_type *child_data = NULL;\n\n if (!this->m_is_fine_level_root_partition)\n {\n int dummy;\n void *data_to_send = this->getHostPointerForData((void *)data_pinned, n * v.get_block_size() * sizeof(value_type), &dummy);\n this->getFineLevelComms()->send_raw_data( data_to_send,\n n * v.get_block_size()*sizeof(value_type),\n this->m_my_fine_level_destination_part,\n 20001 + this->fine_level_id());\n //v.resize(0); // just in case something resized it betwen iterations\n cudaCheckError();\n }\n else\n {\n cudaHostAlloc( (void **) &child_data, this->halo_offsets[this->neighbors.size()]*v.get_block_size()*sizeof(value_type), cudaHostAllocMapped);\n value_type *child_data_hd;\n cudaHostGetDevicePointer(&child_data_hd, child_data, 0);\n // roots copy their data\n int dummy;\n void *my_data = this->getHostPointerForData((void *)data_pinned, n * v.get_block_size() * sizeof(value_type), &dummy);\n memcpy ( &child_data[0], data_pinned, n * v.get_block_size()*sizeof(value_type));\n\n // Loop over parts to consolidate\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n\n if (current_part != this->fine_level_id())\n {\n this->getFineLevelComms()->recv_raw_data( &child_data[this->getConsolidationArrayOffsets()[current_part] - this->getConsolidationArrayOffsets()[this->fine_level_id()]],\n sizeof(value_type) * (this->getConsolidationArrayOffsets()[current_part + 1] - this->getConsolidationArrayOffsets()[current_part]),\n current_part,\n 20001 + current_part );\n }\n }\n\n // usual path\n // Upload on host\n cudaMemcpy(v.raw(), (value_type *)child_data, v.size()* sizeof(value_type), cudaMemcpyDefault);\n cudaCheckError();\n } // If root partition\n\n // Permute based on renumbering vector\n transformVector(v);\n cudaCheckError();\n // Exchange halos\n int tag = 0;\n this->exchange_halo(v, tag);\n cudaCheckError();\n v.set_unconsolidated_size(n);\n\n // free host\n if (child_data) { amgx::memory::cudaFreeHost(child_data); }\n\n cudaCheckError();\n } //cla\n } // If cudaIpcAvailable\n\n if (!this->m_is_fine_level_glued) // not needed for classcical\n {\n if (this->m_is_fine_level_root_partition)\n {\n v.swap(temp);\n int tag = 0;\n // Root partitions do the exchange\n this->exchange_halo(v, tag);\n }\n\n v.set_unconsolidated_size(n * v.get_block_size());\n v.set_transformed();\n }\n\n /* free memory (if needed) */\n if (data_alloc) { amgx::memory::cudaFreeAsync(data_hd); }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManager >::transformVector(VVector_v &v)\n{\n if (this->neighbors.size() == 0) { return; }\n else if (this->renumbering.size() == 0)\n {\n v.resize(this->halo_offsets[this->neighbors.size()]*v.get_block_size());\n return;\n }\n\n if (v.get_block_size() != this->A->get_block_dimx()) { printf(\"Blocksize mismatch!\\n\"); }\n\n if (v.size() < this->halo_offsets[this->neighbors.size()]*v.get_block_size())\n {\n VVector_v temp(this->halo_offsets[this->neighbors.size()]*v.get_block_size(), types::util::get_zero());\n temp.set_block_dimx(v.get_block_dimx());\n temp.set_block_dimy(v.get_block_dimy());\n\n if (v.size() < this->halo_offsets[0]*this->A->get_block_dimx())\n {\n FatalError(\"Unknown size of input vector - smaller than the number of rows owned by this partition\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n //reorder based on row permutation\n int size = this->halo_offsets[0];\n int num_blocks = std::min(4096, (size + 511) / 512);\n reorder_vector_values <<< num_blocks, 512>>>(temp.raw(), v.raw(), this->renumbering.raw(), v.get_block_size(), size);\n cudaCheckError();\n v.swap(temp);\n }\n else\n {\n VVector_v temp(this->halo_offsets[0]*v.get_block_size());\n int size = this->halo_offsets[0];\n int num_blocks = std::min(4096, (size + 511) / 512);\n reorder_vector_values <<< num_blocks, 512>>>(temp.raw(), v.raw(), this->renumbering.raw(), v.get_block_size(), size);\n cudaCheckError();\n amgx::thrust::copy(temp.begin(), temp.end(), v.begin());\n }\n\n cudaCheckError();\n v.set_transformed();\n}\n\ntemplate \nvoid DistributedManager >::revertVector(VVector_v &v)\n{\n if (this->neighbors.size() == 0 || this->renumbering.size() == 0) { return; }\n\n if (v.get_block_size() != this->A->get_block_dimx()) { printf(\"Blocksize mismatch!\\n\"); }\n\n VVector_v temp(this->halo_offsets[0]*this->A->get_block_dimx());\n\n if (v.size() < this->halo_offsets[0]*v.get_block_size())\n {\n FatalError(\"Unknown size of input vector - smaller than the number of rows owned by this partition\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n //reorder based on row permutation\n int size = this->halo_offsets[0];\n int num_blocks = std::min(4096, (size + 511) / 512);\n inverse_reorder_vector_values <<< num_blocks, 512>>>(temp.raw(), v.raw(), this->renumbering.raw(), v.get_block_size(), size);\n //reorder_vector_values<<>>(temp.raw(), v.raw(), this->inverse_renumbering.raw(), v.get_block_size(), size);\n cudaCheckError();\n v.resize(this->halo_offsets[0]*this->A->get_block_dimx());\n amgx::thrust::copy(temp.begin(), temp.end(), v.begin());\n cudaCheckError();\n}\n\n\ntemplate \nvoid DistributedManagerBase::computeDestinationPartitions(INDEX_TYPE upper_threshold, float avg_size, const int num_parts, int &new_num_parts, bool &wantNeighbors)\n{\n m_destination_partitions.resize(num_parts);\n std::vector dp(num_parts);\n\n if (avg_size < 1.f) { avg_size = 1.f; } // avoid floating point exception\n\n int wanted_num_fine_parts_to_consolidate = ( upper_threshold + (int) avg_size - 1) / ( (int) avg_size );\n new_num_parts = (num_parts + wanted_num_fine_parts_to_consolidate - 1) / wanted_num_fine_parts_to_consolidate;\n\n for (int i = 0; i < num_parts; i++)\n {\n dp[i] = i % new_num_parts;\n }\n\n // example wantNeighbors = true -> destination_part = [0 0 0 0 4 4 4 4 8 8 8 8]\n // example wantNeighbors = false -> destination_part = [0 1 2 3 0 1 2 3 0 1 2 3]\n if (wantNeighbors)\n {\n std::sort (dp.begin(), dp.end());\n m_destination_partitions[0] = 0;\n\n for (int i = 1; i < num_parts; i++)\n {\n if (dp[i - 1] < dp[i])\n {\n m_destination_partitions[i] = i;\n }\n else\n {\n m_destination_partitions[i] = m_destination_partitions[i - 1];\n }\n }\n }\n\n m_my_destination_part = m_destination_partitions[global_id()];\n}\n\ntemplate \nvoid DistributedManagerBase::computeDestinationPartitionsWithCons(int my_id, int num_parts, IVector_h &destination_part, DistributedComms *comms)\n{\n int device_id = this->A->getResources()->getDevice(0);\n std::string my_hostname_tmp;\n comms->get_hostname(my_hostname_tmp);\n // Append PCI-E ID to string\n cudaDeviceProp dev_prop;\n cudaGetDeviceProperties(&dev_prop, device_id);\n std::stringstream s;\n s << my_hostname_tmp << \"_\" << dev_prop.pciBusID << \"_\" << dev_prop.pciDeviceID;\n std::string my_hostname(s.str());\n std::vector hostnames;\n comms->exchange_hostnames(my_hostname, hostnames, num_parts);\n std::vector::iterator low = std::find( hostnames.begin(), hostnames.end(), my_hostname );\n int my_destination_part = low - hostnames.begin();\n // Do a gather into destination_part\n comms->all_gather(my_destination_part, destination_part, num_parts);\n}\n\n\ntemplate \nvoid DistributedManager >::revertAndDownloadVector(VVector_v &v_in, const void *data, int n, int block_dimy)\n{\n if (this->isFineLevelConsolidated() || this->isFineLevelGlued())\n {\n revertAndDownloadVectorWithCons(v_in, data, n, block_dimy);\n }\n else\n {\n if ( n == 0 )\n {\n FatalError(\"Cannot download if size = 0\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (data == NULL )\n {\n FatalError(\"Cannot download to a NULL pointer\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (v_in.size() == 0 )\n {\n FatalError(\"Cannot download an empty vector\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n VVector_v v_out;\n revertVector(v_in, v_out);\n cudaMemcpy((value_type *)data, v_out.raw(), n * block_dimy * sizeof(value_type), cudaMemcpyDefault);\n cudaCheckError();\n }\n}\n\ntemplate \nvoid DistributedManager >::revertVector(VVector_v &v_in, VVector_v &v_out)\n{\n if (!this->isFineLevelGlued() && this->neighbors.size() == 0 || this->renumbering.size() == 0) { return;}\n\n if (v_in.get_block_size() != this->A->get_block_dimx()) { printf(\"Blocksize mismatch!\\n\"); }\n\n if (v_in.size() < this->halo_offsets[0]*v_in.get_block_size())\n {\n FatalError(\"Unknown size of input vector - smaller than the number of rows owned by this partition\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n int size = this->halo_offsets[0];\n\n if (v_out.size() != size * this->A->get_block_dimx())\n {\n v_out.resize(size * this->A->get_block_dimx());\n }\n\n //reorder based on row permutation\n int num_blocks = std::min(4096, (size + 511) / 512);\n inverse_reorder_vector_values <<< num_blocks, 512>>>(v_out.raw(), v_in.raw(), this->renumbering.raw(), v_in.get_block_size(), size);\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManager >::revertAndDownloadVectorWithCons(VVector_v &v_in, const void *data_pinned, int n, int block_dimy)\n{\n if (v_in.get_block_size() != this->A->get_block_dimx()) { printf(\"Blocksize mismatch!\\n\"); }\n\n void *root_v_ptr = NULL;\n int size = v_in.get_unconsolidated_size();\n int num_rows = size / v_in.get_block_size();\n\n if (this->m_is_fine_level_root_partition)\n {\n root_v_ptr = (void *) v_in.raw();\n }\n\n VVector_v temp;\n temp.set_block_dimx(v_in.get_block_dimx());\n temp.set_block_dimy(v_in.get_block_dimy());\n temp.resize(size);\n bool useCudaIpc = this->m_use_cuda_ipc_consolidation;\n\n if (useCudaIpc)\n {\n // Do IPC\n this->ipcExchangePtr(root_v_ptr, this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n // Children partition waits for parent to be done updating vector\n this->ipcWaitForRoot(this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n cudaCheckError();\n //reorder based on row permutation\n int num_blocks = std::min(4096, (num_rows + 511) / 512);\n inverse_reorder_vector_values <<< num_blocks, 512>>>( temp.raw(), (value_type *) root_v_ptr, this->m_row_ids_CONS.raw(), v_in.get_block_size(), num_rows);\n cudaCheckError();\n\n if (!this->m_is_fine_level_root_partition)\n {\n cudaIpcCloseMemHandle(root_v_ptr);\n }\n }\n else\n {\n if (this->m_is_fine_level_consolidated) // aggregation\n {\n if (this->m_is_fine_level_root_partition)\n {\n IVector_h child_n(this->m_num_fine_level_parts_to_consolidate);\n int max_n = 0;\n\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n\n if (current_part != this->fine_level_id())\n {\n this->getFineLevelComms()->recv_vector(child_n, current_part, 30000 + current_part, i, 1);\n }\n else\n {\n child_n[i] = num_rows;\n }\n\n if (child_n[i] > max_n) { max_n = child_n[i]; }\n }\n\n // Resize temp vector\n VVector_v child_temp;;\n child_temp.resize(max_n * v_in.get_block_size());\n\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n // Pack the vector to be sent\n int num_blocks = std::min(4096, (child_n[i] + 511) / 512);\n\n if (current_part != this->fine_level_id())\n {\n inverse_reorder_vector_values <<< num_blocks, 512>>>( child_temp.raw(), (value_type *) root_v_ptr, this->m_child_row_ids[i].raw(), v_in.get_block_size(), child_n[i]);\n cudaCheckError();\n this->getFineLevelComms()->send_vector(child_temp, current_part, 30001 + current_part, 0, child_n[i]*v_in.get_block_size());\n }\n else\n {\n inverse_reorder_vector_values <<< num_blocks, 512>>>( temp.raw(), (value_type *) root_v_ptr, this->m_row_ids_CONS.raw(), v_in.get_block_size(), child_n[i]);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n }\n else\n {\n IVector_h size(1);\n size[0] = num_rows;\n this->getFineLevelComms()->send_vector(size, this->m_my_fine_level_destination_part, 30000 + this->fine_level_id());\n this->getFineLevelComms()->recv_vector(temp, this->m_my_fine_level_destination_part, 30001 + this->fine_level_id());\n }\n }\n else if (this->m_is_fine_level_glued) // classical\n {\n if (this->m_is_fine_level_root_partition)\n {\n temp.resize(v_in.size());\n revertVector(v_in, temp);\n cudaCheckError();\n\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n\n if (current_part != this->fine_level_id())\n {\n this->getFineLevelComms()->send_vector( temp,\n current_part,\n current_part + 30001,\n this->getConsolidationArrayOffsets()[current_part] - this->getConsolidationArrayOffsets()[this->fine_level_id()],\n this->getConsolidationArrayOffsets()[current_part + 1] - this->getConsolidationArrayOffsets()[current_part] );\n cudaCheckError();\n }\n }\n }\n else\n {\n this->getFineLevelComms()->recv_vector(temp, this->m_my_fine_level_destination_part, 30001 + this->fine_level_id());\n cudaCheckError();\n }\n\n temp.resize(this->getConsolidationArrayOffsets()[this->fine_level_id() + 1] - this->getConsolidationArrayOffsets()[this->fine_level_id()]);\n cudaCheckError();\n }\n }\n\n // Copy on host\n cudaMemcpy((value_type *)data_pinned, temp.raw(), temp.size() * sizeof(value_type), cudaMemcpyDefault);\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManager >::transformAndUploadVector(VVector_v &v, const void *data, int n, int block_dim)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n\ntemplate \nvoid DistributedManager >::transformVector(VVector_v &v)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::transformAndUploadVectorWithCons(VVector_v &v, const void *data, int n, int block_dim)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n\ntemplate \nvoid DistributedManager >::revertAndDownloadVector(VVector_v &v, const void *data, int n, int block_dim)\n{\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid DistributedManager >::revertVector(VVector_v &v_in, VVector_v &v_out)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::revertVector(VVector_v &v)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::revertAndDownloadVectorWithCons(VVector_v &v_in, const void *data, int n, int block_dim)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::reorder_matrix()\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::reorder_matrix_owned()\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::obtain_shift_l2g_reordering(index_type n, I64Vector_h &l2g, IVector_h &p, IVector_h &q)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::unpack_partition(index_type *Bp, index_type *Bc, mat_value_type *Bv)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::generatePoisson7pt(int nx, int ny, int nz, int P, int Q, int R)\n{\n FatalError(\"GeneratePoisson7pt only implemented on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \ntemplate \nvoid DistributedManager >::loadDistributedMatrix(\n int num_rows, int num_nonzeros, const int block_dimx, const int block_dimy, const int *row_offsets, \n const t_colIndex *col_indices, const mat_value_type *values, int num_ranks, int num_rows_global, const void *diag, const MatrixDistribution &dist)\n{\n FatalError(\"loadDistributedMatrix only implemented on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid DistributedManager >::createOneRingB2Lmaps()\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::createOneRingHaloRows()\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::replaceMatrixCoefficientsNoCons(int n, int nnz, const mat_value_type *data, const mat_value_type *diag)\n{\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid DistributedManager >::replaceMatrixCoefficientsWithCons(int n, int nnz, const mat_value_type *data, const mat_value_type *diag)\n{\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n\ntemplate \nvoid DistributedManager >::createRenumbering(IVector &renumbering)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate < class TConfig >\nvoid DistributedManagerBase::remove_boundary(IVector_h &flagArray, IVector_h &B2L_map, int size)\n{\n for (int i = 0; i < size; i++)\n {\n flagArray[B2L_map[i]] = 0;\n }\n}\n\ntemplate < class TConfig >\nvoid DistributedManagerBase::get_unassigned(IVector_h &flagArray, IVector_h &B2L_map, IVector_h &partition_flags, int size, int fa_size/*, int rank*/)\n{\n for (int i = 0; i < size; i++)\n {\n if (B2L_map[i] < fa_size)\n {\n if (flagArray[B2L_map[i]] == 0)\n {\n flagArray[B2L_map[i]] = 1;\n partition_flags[i] = 1;\n }\n }\n }\n}\n\ntemplate < class TConfig >\nvoid DistributedManagerBase::set_unassigned(IVector_h &partition_flags, IVector_h &partition_renum, IVector_h &B2L_map, IVector_h &renumbering, int size, int max_element, int renum_size/*, int rank*/)\n{\n for (int i = 0; i < size; i++)\n {\n if (B2L_map[i] < renum_size)\n {\n if (partition_flags[i] == 1)\n {\n renumbering[B2L_map[i]] = max_element + partition_renum[i];\n }\n\n B2L_map[i] = renumbering[B2L_map[i]];\n }\n }\n}\n\n/* print manager for target rank to a file or stdout */\ntemplate\nvoid DistributedManagerBase::print(char *f, char *s, int trank)\n{\n DistributedManagerBase *m = this;\n int rank = 0;\n int level = 0;\n char filename[1024];\n FILE *fid = NULL;\n int i, j, k, t1, t2;\n#ifdef AMGX_WITH_MPI\n MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n#endif\n\n //check target rank\n if (rank == trank)\n {\n //check whether to output to stdout or a file\n if (f == NULL)\n {\n fid = stdout;\n }\n else\n {\n level = m->A->amg_level_index;\n#ifdef _WIN32\n _snprintf_s(filename, 1024, 1024, \"%s_r%d_l%d.mtx\", f, rank, level);\n#else\n snprintf(filename, 1024, \"%s_r%d_l%d.mtx\", f, rank, level);\n#endif\n fid = fopen(filename, \"w\");\n }\n\n cudaDeviceSynchronize();\n cudaCheckError();\n fprintf(fid, \"%s\\n\", s);\n //--- communication info ---\n //compare neighbors\n t1 = m->neighbors.size();\n fprintf(fid, \"neighbors %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->neighbors[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare B2L_rings\n t1 = B2L_rings.size();\n fprintf(fid, \"B2L_rings %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n t2 = m->B2L_rings[i].size();\n fprintf(fid, \"B2L_rings-%d [%d]\\n\", i, t2);\n\n for (j = 0; j < t2; j++)\n {\n k = m->B2L_rings[i][j];\n fprintf(fid, \"%d\\n\", k);\n }\n }\n\n //compare B2L_maps\n t1 = B2L_maps.size();\n fprintf(fid, \"B2L_maps %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n t2 = m->B2L_maps[i].size();\n fprintf(fid, \"B2L_maps-%d [%d]\\n\", i, t2);\n\n for (j = 0; j < t2; j++)\n {\n k = m->B2L_maps[i][j];\n fprintf(fid, \"%d\\n\", k);\n }\n }\n\n //compare L2H_maps\n t1 = L2H_maps.size();\n fprintf(fid, \"L2H_maps %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n t2 = m->L2H_maps[i].size();\n fprintf(fid, \"L2H_maps-%d [%d]\\n\", i, t2);\n\n for (j = 0; j < t2; j++)\n {\n k = m->L2H_maps[i][j];\n fprintf(fid, \"%d\\n\", k);\n }\n }\n\n //--- matrix info ---\n fprintf(fid, \"num_rows_global=%ld\\n\", num_rows_global);\n fprintf(fid, \"_num_rows_interior=%d\\n\", m->_num_rows_interior);\n fprintf(fid, \"_num_rows_owned=%d\\n\", m->_num_rows_owned);\n fprintf(fid, \"_num_rows_full=%d\\n\", m->_num_rows_full);\n fprintf(fid, \"_num_rows_all=%d\\n\", m->_num_rows_all);\n fprintf(fid, \"_num_nz_interior=%d\\n\", m->_num_nz_interior);\n fprintf(fid, \"_num_nz_owned=%d\\n\", m->_num_nz_owned);\n fprintf(fid, \"_num_nz_full=%d\\n\", m->_num_nz_full);\n fprintf(fid, \"_num_nz_all=%d\\n\", m->_num_nz_all);\n //compare # halo rows and halo offsets\n fprintf(fid, \"# halo rings %d and rows %d\\n\", m->num_halo_rings(), m->num_halo_rows());\n t1 = m->halo_offsets.size();\n fprintf(fid, \"halo_offsets %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->halo_offsets[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare halo ranges\n t1 = m->halo_ranges.size();\n fprintf(fid, \"halo_ranges %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->halo_ranges[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare halo ranges (host)\n t1 = m->halo_ranges_h.size();\n fprintf(fid, \"halo_ranges_h %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->halo_ranges_h[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare part offsets\n t1 = m->part_offsets.size();\n fprintf(fid, \"part_offsets %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->part_offsets[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare part offsets (host)\n t1 = m->part_offsets_h.size();\n fprintf(fid, \"part_offsets_h %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->part_offsets_h[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare interior row list\n t1 = m->interior_rows_list.size();\n fprintf(fid, \"interior_rows_list %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->interior_rows_list[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare boundary row list\n t1 = m->boundary_rows_list.size();\n fprintf(fid, \"boundary_rows_list %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->boundary_rows_list[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare halo1 row list\n t1 = m->halo1_rows_list.size();\n fprintf(fid, \"halo1_rows_list %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->halo1_rows_list[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n fprintf(fid, \"pointers halo_rows=%p and halo_btl=%p\\n\", m->halo_rows, m->halo_btl);\n //--- packing info ---\n //compare local to global map\n t1 = m->local_to_global_map.size();\n fprintf(fid, \"local_to_global_map %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->local_to_global_map[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare renumbering\n t1 = m->renumbering.size();\n fprintf(fid, \"renumbering %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->renumbering[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare inverse renumbering\n t1 = m->inverse_renumbering.size();\n fprintf(fid, \"inverse_renumbering %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->inverse_renumbering[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //--- GPU related and miscellaneous info\n //streams\n fprintf(fid, \"streams i=%p, b=%p\\n\", m->get_int_stream(), m->get_bdy_stream());\n //miscellaneous info\n int64_t bi = m->base_index(); //inlined function\n int np = m->get_num_partitions(); //inlined function\n int rp = (int)m->isRootPartition(); //cast from boolean to int\n fprintf(fid, \"gid=%d,bi=%ld,np=%d,rp=%d,ir=%d,in=%d,bn=%d\\n\", m->global_id(), bi, np, rp, m->index_range(), m->num_interior_nodes(), m->num_boundary_nodes());\n cudaDeviceSynchronize();\n cudaCheckError();\n cudaGetLastError();\n\n if (fid != stdout)\n {\n fclose(fid);\n }\n }\n}\n\n/* print manager for target rank to a file or stdout (for all ranks) */\ntemplate\nvoid DistributedManagerBase::printToFile(char *f, char *s)\n{\n DistributedManagerBase *m = this;\n int rank = 0;\n#ifdef AMGX_WITH_MPI\n MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n#endif\n //notice that print will be called with different (target) rank on different ranks/processes\n m->print(f, s, rank);\n}\n\n/* compare two managers */\ntemplate\nint DistributedManagerBase::compare(DistributedManagerBase *m2)\n{\n DistributedManagerBase *m1 = this;\n int i, j, t1, t2;\n //compare neighbors\n t1 = m1->neighbors.size();\n t2 = m2->neighbors.size();\n\n if (t1 != t2)\n {\n return 1;\n }\n\n for (i = 0; i < t1; i++)\n {\n if (m1->neighbors[i] != m2->neighbors[i])\n {\n return 2;\n }\n }\n\n //compare B2L_rings\n for (i = 0; i < (m1->neighbors.size()); i++)\n {\n t1 = m1->B2L_rings[i].size();\n t2 = m2->B2L_rings[i].size();\n\n if (t1 != t2)\n {\n return 3;\n }\n\n for (j = 0; j < t1; j++)\n {\n if (m1->B2L_rings[i][j] != m2->B2L_rings[i][j])\n {\n return 4;\n }\n }\n }\n\n //compare B2L_maps\n t1 = m1->B2L_maps.size();\n t2 = m2->B2L_maps.size();\n\n if (t1 != t2)\n {\n return 5;\n }\n\n for (i = 0; i < t1; i++)\n {\n if (m1->B2L_maps[i] != m2->B2L_maps[i])\n {\n return 6;\n }\n }\n\n //compare L2H_maps\n t1 = m1->L2H_maps.size();\n t2 = m2->L2H_maps.size();\n\n if (t1 != t2)\n {\n return 7;\n }\n\n for (i = 0; i < t1; i++)\n {\n if (m1->L2H_maps[i] != m2->L2H_maps[i])\n {\n return 8;\n }\n }\n\n return 0;\n}\n\ntemplate \nDistributedManager< TemplateConfig >::~DistributedManager< TemplateConfig >()\n{\n}\n\ntemplate \nDistributedManager< TemplateConfig >::~DistributedManager< TemplateConfig >()\n{\n}\n\ntemplate \nvoid DistributedManagerBase::consolidateB2Lmaps(IVector_h_vector &dest_coarse_B2L_maps, IVector_h_vector &coarse_B2L_maps, IVector_h &fine_neigh_to_coarse_neigh, int num_coarse_neighbors, int num_fine_neighbors)\n{\n consB2Lmaps(dest_coarse_B2L_maps, coarse_B2L_maps, fine_neigh_to_coarse_neigh, num_coarse_neighbors, num_fine_neighbors);\n}\n\ntemplate \nvoid DistributedManagerBase::consolidateB2Lmaps(IVector_d_vector &dest_coarse_B2L_maps, IVector_d_vector &coarse_B2L_maps, IVector_h &fine_neigh_to_coarse_neigh, int num_coarse_neighbors, int num_fine_neighbors)\n{\n consB2Lmaps(dest_coarse_B2L_maps, coarse_B2L_maps, fine_neigh_to_coarse_neigh, num_coarse_neighbors, num_fine_neighbors);\n}\n\ntemplate \nvoid DistributedManagerBase::consolidateB2LmapsOnRoot(int &num_consolidated_neighbors, IVector_h_vector &consolidated_B2L_maps, IVector_h &consolidated_coarse_ids, IVector_h_vector &dest_coarse_B2L_maps, IVector_h &coarse_neigh_to_fine_part, IVector_h &num_bdy_per_coarse_neigh, IVector_h &fine_parts_to_consolidate, int num_fine_parts_to_consolidate, int my_id, int my_destination_part, bool is_root_partition, int num_coarse_neighbors, DistributedComms *comms)\n{\n consB2LmapsOnRoot(num_consolidated_neighbors, consolidated_B2L_maps, consolidated_coarse_ids, dest_coarse_B2L_maps, coarse_neigh_to_fine_part, num_bdy_per_coarse_neigh, fine_parts_to_consolidate, num_fine_parts_to_consolidate, my_id, my_destination_part, is_root_partition, num_coarse_neighbors, comms);\n}\n\ntemplate \nvoid DistributedManagerBase::consolidateB2LmapsOnRoot(int &num_consolidated_neighbors, IVector_d_vector &consolidated_B2L_maps, IVector_h &consolidated_coarse_ids, IVector_d_vector &dest_coarse_B2L_maps, IVector_h &coarse_neigh_to_fine_part, IVector_h &num_bdy_per_coarse_neigh, IVector_h &fine_parts_to_consolidate, int num_fine_parts_to_consolidate, int my_id, int my_destination_part, bool is_root_partition, int num_coarse_neighbors, DistributedComms *comms)\n{\n consB2LmapsOnRoot(num_consolidated_neighbors, consolidated_B2L_maps, consolidated_coarse_ids, dest_coarse_B2L_maps, coarse_neigh_to_fine_part, num_bdy_per_coarse_neigh, fine_parts_to_consolidate, num_fine_parts_to_consolidate, my_id, my_destination_part, is_root_partition, num_coarse_neighbors, comms);\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class DistributedManager::Type >;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template void DistributedManager::Type>::loadDistributedMatrix( \\\n int, int, const int, const int, const int*, const int *col_indices, const mat_value_type*, int, int, const void*, const MatrixDistribution &dist); \\\n template void DistributedManager::Type>::loadDistributedMatrix( \\\n int, int, const int, const int, const int*, const int64_t *col_indices, const mat_value_type*, int, int, const void*, const MatrixDistribution &dist);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class DistributedManagerBase::Type >;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"hip/hip_runtime.h\"\n#include \"reorder_partition.h\"\n\n#include \"amgx_types/util.h\"\n\n#include \n#include //debug only:\n\nstruct is_my_part\n{\n const int _my_part;\n is_my_part(int my_part) : _my_part(my_part) { }\n\n __host__ __device__\n bool operator()(const int part)\n {\n return (part == _my_part);\n }\n};\n\nnamespace amgx\n{\n\nstatic int insertDiagonals = 1;\n\ntemplate \nstatic __device__ __forceinline__\nindex_type internal_index(index_type i, index_type j, index_type k, index_type nx, index_type ny, index_type nz)\n{\n return k * (nx * ny) + j * nx + i;\n}\n\ntemplate \nstatic __device__ __forceinline__\nint64_t get_global_offset(index_type p, index_type q, index_type r, index_type P, index_type Q, index_type R, index_type num_rows)\n{\n int rank_id = r * (P * Q) + q * P + p;\n return ((int64_t) rank_id) * ((int64_t) num_rows);\n}\n\ntemplate \n__global__\nvoid poisson7pt_count_row_len(index_type *row_len, index_type nx, index_type ny, index_type nz, index_type p, index_type q, index_type r, index_type P, index_type Q, index_type R, index_type num_rows)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_rows ; tidx += blockDim.x * gridDim.x)\n {\n /* compute p,q,r from P,Q,R and myid */\n int i = tidx % nx; // Position in x direction\n int j = (( tidx - i) / nx) % ny; // Position in y\n int k = ( tidx - i - nx * j) / ( nx * ny ); // Position in z\n int substract = ((i == 0) && (p == 0));\n substract += ((i == nx - 1) && (p == P - 1));\n substract += ((j == 0) && (q == 0));\n substract += ((j == ny - 1) && (q == Q - 1));\n substract += ((k == 0) && (r == 0));\n substract += ((k == nz - 1) && (r == R - 1));\n // Store 7 in position (num_rows+1), such that row_len[num_rows+1] = 0\n //substract = (tidx == num_rows+1) ? 7 : substract;\n row_len[tidx] = 7 - substract;\n }\n}\n\ntemplate \n__global__\nvoid poisson7pt_set_col_values(const index_type *__restrict__ row_offsets, index_type *__restrict__ col_indices, mat_value_type *__restrict__ values, index_type *__restrict__ diag, int64_t *__restrict__ local_to_global, index_type nx, index_type ny, index_type nz, index_type p, index_type q, index_type r, index_type P, index_type Q, index_type R, index_type num_rows)\n{\n for (int row = threadIdx.x + blockIdx.x * blockDim.x; row < num_rows ; row += blockDim.x * gridDim.x)\n {\n /* compute p,q,r from P,Q,R and myid */\n int i = row % nx; // Position in x direction\n int j = (( row - i) / nx) % ny; // Position in y\n int k = ( row - i - nx * j) / ( nx * ny ); // Position in z\n int halo_offset = num_rows;\n int pos = row_offsets[row];\n // Diagonal element\n diag[row] = pos;\n col_indices[pos] = row;\n values[pos++] = types::util::get_one() * 6.;\n\n // ----------------------------\n // Neighbor at position i-1\n // ----------------------------\n if (i)\n {\n // Has a i-1 neighbor, which is an internal node at position (i-1,j,k)\n col_indices[pos] = internal_index(i - 1, j, k, nx, ny, nz);\n values[pos++] = types::util::invert(types::util::get_one());\n }\n else if (p)\n {\n // Has a i-1 neighbor, which is a halo node\n int halo_index = halo_offset + k * ny + j;\n col_indices[pos] = halo_index;\n values[pos++] = types::util::invert(types::util::get_one());\n int64_t global_offset = get_global_offset(p - 1, q, r, P, Q, R, num_rows);\n local_to_global[halo_index - num_rows] = global_offset + internal_index(nx - 1, j, k, nx, ny, nz);\n }\n\n if (p)\n {\n halo_offset += ny * nz;\n }\n\n // ----------------------------\n // Neighbor at position i+1\n // ----------------------------\n if (i < nx - 1)\n {\n // Has i+1 neighbor, which is an internal node at position (i+1,j,k)\n col_indices[pos] = internal_index(i + 1, j, k, nx, ny, nz);\n values[pos++] = types::util::invert(types::util::get_one());\n }\n else\n {\n if (p < P - 1)\n {\n // Has i+1 neighbor, which is a halo node\n int halo_index = halo_offset + k * ny + j;\n col_indices[pos] = halo_index;\n values[pos++] = types::util::invert(types::util::get_one());\n int64_t global_offset = get_global_offset(p + 1, q, r, P, Q, R, num_rows);\n local_to_global[halo_index - num_rows] = global_offset + internal_index(0, j, k, nx, ny, nz);\n }\n }\n\n if (p < P - 1)\n {\n halo_offset += ny * nz;\n }\n\n // ----------------------------\n // Neighbor at position j-1\n // ----------------------------\n if (j)\n {\n // Has a j-1 neighbor, which is an internal node at position (i,j-1,k)\n col_indices[pos] = internal_index(i, j - 1, k, nx, ny, nz);\n values[pos++] = types::util::invert(types::util::get_one());\n }\n else if (q)\n {\n // Has a j-1 neighbor, which is a halo node\n int halo_index = halo_offset + k * nx + i;\n col_indices[pos] = halo_index;\n values[pos++] = types::util::invert(types::util::get_one());\n int64_t global_offset = get_global_offset(p, q - 1, r, P, Q, R, num_rows);\n local_to_global[halo_index - num_rows] = global_offset + internal_index(i, ny - 1, k, nx, ny, nz);\n }\n\n if (q)\n {\n halo_offset += nx * nz;\n }\n\n // ----------------------------\n // Neighbor at position j+1\n // ----------------------------\n if (j < ny - 1)\n {\n // Has a j+1 neighbor, which is an internal node at position (i,j+1,k)\n col_indices[pos] = internal_index(i, j + 1, k, nx, ny, nz);\n values[pos++] = types::util::invert(types::util::get_one());\n }\n else\n {\n if (q < Q - 1)\n {\n // Has a j+1 neighbor, which is a halo node\n int halo_index = halo_offset + k * nx + i;\n col_indices[pos] = halo_index;\n values[pos++] = types::util::invert(types::util::get_one());\n int64_t global_offset = get_global_offset(p, q + 1, r, P, Q, R, num_rows);\n local_to_global[halo_index - num_rows] = global_offset + internal_index(i, 0, k, nx, ny, nz);\n }\n }\n\n if (q < Q - 1)\n {\n halo_offset += nx * nz;\n }\n\n // ----------------------------\n // Neighbor at position k-1\n // ----------------------------\n if (k)\n {\n // Has a k-1 neighbor, which is an internal node at position (i,j,k-1)\n col_indices[pos] = internal_index(i, j, k - 1, nx, ny, nz);\n values[pos++] = types::util::invert(types::util::get_one());\n }\n else if (r)\n {\n // Has a k-1 neighbor, which is a halo node\n int halo_index = halo_offset + j * nx + i;\n col_indices[pos] = halo_index;\n values[pos++] = types::util::invert(types::util::get_one());\n int64_t global_offset = get_global_offset(p, q, r - 1, P, Q, R, num_rows);\n local_to_global[halo_index - num_rows] = global_offset + internal_index(i, j, nz - 1, nx, ny, nz);\n }\n\n if (r)\n {\n halo_offset += nx * ny;\n }\n\n // ----------------------------\n // Neighbor at position k+1\n // ----------------------------\n if (k < nz - 1)\n {\n // Has a k+1 neighbor, which is an internal node at position (i,j,k+1)\n col_indices[pos] = internal_index(i, j, k + 1, nx, ny, nz);\n values[pos++] = types::util::invert(types::util::get_one());\n }\n else\n {\n if (r < R - 1)\n {\n // Has a k+1 neighbor, which is a halo node\n int halo_index = halo_offset + j * nx + i;\n col_indices[pos] = halo_index;\n values[pos++] = types::util::invert(types::util::get_one());\n int64_t global_offset = get_global_offset(p, q, r + 1, P, Q, R, num_rows);\n local_to_global[halo_index - num_rows] = global_offset + internal_index(i, j, 0, nx, ny, nz);\n }\n }\n\n if (r < R - 1)\n {\n halo_offset += nx * ny;\n }\n }\n}\n\ntemplate \n__global__\nvoid set_halo_cols_values(int *row_offsets, int *col_indices, mat_value_type *values, int n, int total_rows, int bsize)\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (tid < (total_rows - n) )\n {\n int offset = row_offsets[n + tid];\n col_indices[offset] = n + tid;\n#pragma unroll\n\n for (int i = 0; i < bsize; i++)\n {\n values[offset * bsize + i] = types::util::get_one(); // This is arbitrary\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \n__global__\nvoid zero_copy_row_lengths_ids_offsets(int *d_old_row_offsets, int *root_row_offsets, int *d_row_ids, int n, int total_num_halos, mat_value_type *diag)\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (tid < n + total_num_halos)\n {\n int new_row_id = d_row_ids[tid];\n\n if (tid < n)\n {\n int start = d_old_row_offsets[tid];\n int row_length = d_old_row_offsets[tid + 1] - start; // zero-copy\n\n if (diag != NULL) // will insert the diagonal\n {\n row_length++;\n }\n\n root_row_offsets[new_row_id] = row_length;\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate< typename mat_value_type>\n__global__\nvoid ipc_consolidation_upload_matrix(int num_rows, int *row_ids, const int *old_row_offsets, int *new_row_offsets, const int *h_old_col_indices, int *new_col_indices, const mat_value_type *h_old_values, mat_value_type *new_values, const mat_value_type *h_old_diag, int bsize)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < num_rows)\n {\n int new_row = row_ids[row];\n int src_base = old_row_offsets[row];\n int dst_base = new_row_offsets[new_row];\n\n // Insert the diagonal at the beginning of each row\n if (h_old_diag != NULL)\n {\n new_col_indices[dst_base] = new_row;\n#pragma unroll\n\n for (int j = 0; j < bsize; j++)\n {\n new_values[dst_base * bsize + j] = h_old_diag[row * bsize + j];\n }\n\n // Increment dst_base by one\n dst_base++;\n }\n\n int end = old_row_offsets[row + 1] - src_base;\n\n for (int i = 0; i < end; i++)\n {\n int old_col = h_old_col_indices[src_base + i];\n int new_col = row_ids[old_col];\n new_col_indices[dst_base + i] = new_col;\n#pragma unroll\n\n for (int j = 0; j < bsize; j++)\n {\n new_values[ (dst_base + i)*bsize + j ] = h_old_values[ (src_base + i) * bsize + j ];\n }\n }\n\n row += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate< typename mat_value_type>\n__global__\nvoid ipc_consolidation_replace_values(int num_rows, int *row_ids, const int *old_row_offsets, int *new_row_offsets, const mat_value_type *h_old_values, mat_value_type *new_values, const mat_value_type *h_old_diag, int bsize)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < num_rows)\n {\n int new_row = row_ids[row];\n int src_base = old_row_offsets[row];\n int dst_base = new_row_offsets[new_row];\n\n // Insert the diagonal at the beginning of each row\n if (h_old_diag != NULL)\n {\n#pragma unroll\n\n for (int j = 0; j < bsize; j++)\n {\n new_values[dst_base * bsize + j] = h_old_diag[row * bsize + j];\n }\n\n // Increment dst_base by one\n dst_base++;\n }\n\n int end = old_row_offsets[row + 1] - src_base;\n\n for (int i = 0; i < end; i++)\n {\n#pragma unroll\n\n for (int j = 0; j < bsize; j++)\n {\n new_values[ (dst_base + i)*bsize + j ] = h_old_values[ (src_base + i) * bsize + j ];\n }\n }\n\n row += gridDim.x * blockDim.x;\n }\n}\n\n\n__global__ void flag_halo_ids_kernel(INDEX_TYPE *flags, INDEX_TYPE *ids, INDEX_TYPE offset, INDEX_TYPE size, INDEX_TYPE upper)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n flags[ids[idx] - offset] = 1;\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void read_halo_ids_kernel(INDEX_TYPE *flags, INDEX_TYPE *ids, INDEX_TYPE offset, INDEX_TYPE size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n ids[idx] = flags[ids[idx] - offset];\n idx += blockDim.x * gridDim.x;\n }\n}\n\ntemplate\n__global__ void reorder_vector_values(T *dst, const T *src, const INDEX_TYPE *map, INDEX_TYPE blocksize, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * (blockDim.x / blocksize) + threadIdx.x / blocksize; //vectorised by block size\n int vec_id = threadIdx.x % blocksize;\n\n if (threadIdx.x >= (blockDim.x / blocksize)*blocksize ) { return; }\n\n while (row < num_rows)\n {\n dst[map[row]*blocksize + vec_id] = src[row * blocksize + vec_id];\n row += gridDim.x * (blockDim.x / blocksize);\n }\n}\n\ntemplate\n__global__ void inverse_reorder_vector_values(T *dst, T *src, INDEX_TYPE *map, INDEX_TYPE blocksize, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * (blockDim.x / blocksize) + threadIdx.x / blocksize;\n int vec_id = threadIdx.x % blocksize;\n\n if (threadIdx.x >= (blockDim.x / blocksize)*blocksize ) { return; }\n\n while (row < num_rows)\n {\n dst[row * blocksize + vec_id] = src[map[row] * blocksize + vec_id];\n row += gridDim.x * (blockDim.x / blocksize);\n }\n}\n\n__global__ void remove_boundary_kernel(INDEX_TYPE *flags, INDEX_TYPE *maps, INDEX_TYPE size)\n{\n int element = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (element < size)\n {\n flags[maps[element]] = 0; //this won't be a problem, because we are overwriting the same thing\n element += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void get_unassigned_kernel(INDEX_TYPE *unassigned_flags, INDEX_TYPE *map, INDEX_TYPE *output, INDEX_TYPE part_size, INDEX_TYPE uf_size )\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < part_size)\n {\n if (map[idx] < uf_size)\n {\n if (unassigned_flags[map[idx]] == 0)\n {\n unassigned_flags[map[idx]] = 1;\n output[idx] = 1;\n }\n }\n\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void set_unassigned_kernel(INDEX_TYPE *part_assigned_flags, INDEX_TYPE *part_num, INDEX_TYPE *map, INDEX_TYPE *renum, INDEX_TYPE part_size, INDEX_TYPE max_element, INDEX_TYPE renum_size /*, INDEX_TYPE rank*/)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < part_size)\n {\n if (map[idx] < renum_size)\n {\n if (part_assigned_flags[idx] == 1)\n {\n renum[map[idx]] = max_element + part_num[idx];\n }\n\n //also update the B2L map\n map[idx] = renum[map[idx]];\n }\n\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void renumber_b2l_maps(INDEX_TYPE *map, INDEX_TYPE *renum, INDEX_TYPE part_size, INDEX_TYPE renum_size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < part_size)\n {\n if (map[idx] < renum_size)\n {\n //update the B2L map\n map[idx] = renum[map[idx]];\n idx += blockDim.x * gridDim.x;\n }\n }\n}\n\n__global__ void calc_inverse_renumbering(INDEX_TYPE *renum, INDEX_TYPE *irenum, INDEX_TYPE max_element)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < max_element)\n {\n if (renum[idx] < 0 || renum[idx] >= max_element) { printf(\"Renumbering error: %d %d\\n\", renum[idx], max_element); }\n\n irenum[renum[idx]] = idx;\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void create_halo_mapping(INDEX_TYPE *mapping, INDEX_TYPE *node_list, int64_t base_index, INDEX_TYPE map_offset, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n int idx = node_list[row] - base_index;\n mapping[idx] = map_offset + row;\n row += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void apply_h2l2b_mapping(INDEX_TYPE *mapping, INDEX_TYPE *node_list, int64_t base_index, INDEX_TYPE *b2l_map, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n int idx = node_list[row] - base_index;\n mapping[idx] = b2l_map[row];\n row += blockDim.x * gridDim.x;\n }\n}\n\ntemplate \n__global__ void map_col_indices_and_count_rowlen(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, INDEX_TYPE *row_length,\n INDEX_TYPE *mapping, INDEX_TYPE num_rows, INDEX_TYPE insert_diagonal)\n{\n extern __shared__ volatile int reduction[];\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n int valid = 0;\n\n for (int idx = row_offsets[row] + coopIdx; idx < row_offsets[row + 1]; idx += coop) //this may look horrible, but I expect low branch divergence, because col indices in a row usually belong to the same partition (or at most one more)\n {\n int colIdx = col_indices[idx];\n int new_col_idx = mapping[colIdx];\n\n if (new_col_idx >= 0)\n {\n valid++;\n col_indices[idx] = new_col_idx;\n }\n else\n {\n col_indices[idx] = -1;\n }\n }\n\n reduction[threadIdx.x] = valid;\n\n for (int s = 2; s > 0; s >>= 1)\n {\n if (coopIdx < s)\n {\n reduction[threadIdx.x] += reduction[threadIdx.x + s];\n }\n\n __syncthreads();\n }\n\n if (coopIdx == 0)\n {\n row_length[row] = reduction[threadIdx.x] + insert_diagonal;\n }\n\n row += gridDim.x * blockDim.x / coop;\n }\n}\n\n\n__global__ void renumber_P_col_indices(INDEX_TYPE *__restrict__ col_indices, const INDEX_TYPE *__restrict__ renum, INDEX_TYPE num_owned_coarse_pts, INDEX_TYPE num_owned_fine_pts)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < num_owned_fine_pts )\n {\n INDEX_TYPE col_id = col_indices[idx];\n\n if (col_id < num_owned_coarse_pts)\n {\n col_indices[idx] = renum[col_id];\n }\n\n idx += blockDim.x * gridDim.x;\n }\n}\n\n\ntemplate \n__global__ void reorder_R_matrix(const INDEX_TYPE *old_rows, const INDEX_TYPE *old_cols, const T *old_vals, const INDEX_TYPE *rows, INDEX_TYPE *cols, T *vals, const INDEX_TYPE *renumbering, INDEX_TYPE bsize, INDEX_TYPE num_rows, INDEX_TYPE num_owned_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n INDEX_TYPE src_base = old_rows[row];\n INDEX_TYPE dst_base = row < num_owned_rows ? rows[renumbering[row]] : src_base;\n\n for (int i = coopIdx; i < old_rows[row + 1]*bsize - src_base * bsize; i += coop)\n {\n vals[dst_base * bsize + i] = old_vals[src_base * bsize + i];\n }\n\n for (int i = coopIdx; i < old_rows[row + 1] - src_base; i += coop)\n {\n cols[dst_base + i] = old_cols[src_base + i];\n }\n\n row += blockDim.x * gridDim.x / coop;\n }\n}\n\n\n\ntemplate \n__global__ void reorder_whole_matrix(INDEX_TYPE *old_rows, INDEX_TYPE *old_cols, T *old_vals, INDEX_TYPE *rows, INDEX_TYPE *cols, T *vals, INDEX_TYPE *renumbering, INDEX_TYPE bsize, INDEX_TYPE num_rows, INDEX_TYPE insert_diagonal)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n INDEX_TYPE src_base = old_rows[row];\n INDEX_TYPE dst_base = rows[renumbering[row]];\n\n if (insert_diagonal)\n {\n if (coopIdx == 0) { cols[dst_base] = renumbering[row]; }\n\n for (int i = coopIdx; i < bsize; i += coop)\n {\n vals[dst_base * bsize + i] = old_vals[(old_rows[num_rows] + row) * bsize + i];\n }\n\n dst_base++;\n }\n\n for (int i = coopIdx; i < old_rows[row + 1]*bsize - src_base * bsize; i += coop)\n {\n vals[dst_base * bsize + i] = old_vals[src_base * bsize + i];\n }\n\n for (int i = coopIdx; i < old_rows[row + 1] - src_base; i += coop)\n {\n cols[dst_base + i] = old_cols[src_base + i];\n }\n\n row += blockDim.x * gridDim.x / coop;\n }\n}\n\ntemplate \n__global__ void replace_values_matrix(const T *src_vals_h, const T *src_diag_h, const INDEX_TYPE *old_rows, const INDEX_TYPE *rows, T *vals, const INDEX_TYPE *renumbering, INDEX_TYPE bsize, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n INDEX_TYPE src_base = old_rows[row];\n INDEX_TYPE dst_base = rows[renumbering[row]];\n\n for (int i = coopIdx; i < bsize; i += coop)\n {\n vals[dst_base * bsize + i] = src_diag_h[row * bsize + i];\n }\n\n dst_base++;\n\n for (int i = coopIdx; i < old_rows[row + 1]*bsize - src_base * bsize; i += coop)\n {\n vals[dst_base * bsize + i] = src_vals_h[src_base * bsize + i];\n }\n\n row += blockDim.x * gridDim.x / coop;\n }\n}\n\ntemplate \n__global__ void replace_values_matrix(const T *src_vals_h, const INDEX_TYPE *old_rows, const INDEX_TYPE *rows, T *vals, const INDEX_TYPE *renumbering, INDEX_TYPE bsize, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop;\n int coopIdx = threadIdx.x % coop;\n\n while (row < num_rows)\n {\n INDEX_TYPE src_base = old_rows[row];\n INDEX_TYPE dst_base = rows[renumbering[row]];\n\n for (int i = coopIdx; i < old_rows[row + 1]*bsize - src_base * bsize; i += coop)\n {\n vals[dst_base * bsize + i] = src_vals_h[src_base * bsize + i];\n }\n\n row += blockDim.x * gridDim.x / coop;\n }\n}\n\n//TODO: optimize by vectorizing\ntemplate \n__global__ void reorder_whole_halo_matrix(INDEX_TYPE *old_rows, INDEX_TYPE *old_cols, T *old_vals, INDEX_TYPE *rows, INDEX_TYPE *cols, T *vals,\n INDEX_TYPE *renumbering, INDEX_TYPE bsize, INDEX_TYPE num_rows, INDEX_TYPE insert_diagonal,\n INDEX_TYPE global_offset, INDEX_TYPE local_offset, INDEX_TYPE halo_rows)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < num_rows)\n {\n INDEX_TYPE src_base = old_rows[row];\n INDEX_TYPE dst = rows[row];\n\n if (insert_diagonal)\n {\n cols[dst] = global_offset + row;\n\n for (int j = 0; j < bsize; j++) { vals[dst * bsize + j] = old_vals[(old_rows[halo_rows - local_offset] + local_offset + row) * bsize + j]; }\n\n dst++;\n }\n\n for (int i = 0; i < old_rows[row + 1] - src_base; i++)\n {\n INDEX_TYPE colIdx = old_cols[src_base + i];\n\n if (colIdx >= 0)\n {\n cols[dst] = colIdx;\n\n for (int j = 0; j < bsize; j++) { vals[dst * bsize + j] = old_vals[(src_base + i) * bsize + j]; }\n\n dst++;\n }\n }\n\n row += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void calc_rowlen_reorder(INDEX_TYPE *row_offsets, INDEX_TYPE *row_len, INDEX_TYPE *map, INDEX_TYPE size, INDEX_TYPE insert_diag)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n row_len[map[idx]] = row_offsets[idx + 1] - row_offsets[idx] + insert_diag;\n idx += blockDim.x * gridDim.x;\n }\n}\n\n\ntemplate < class TConfig >\nvoid DistributedManagerBase::remove_boundary(IVector_d &flagArray, IVector_d &B2L_map, int size)\n{\n int num_blocks = std::min(4096, (size + 127) / 128);\n remove_boundary_kernel <<< num_blocks, 128>>>(flagArray.raw(), B2L_map.raw(), size);\n cudaCheckError();\n}\n\ntemplate < class TConfig >\nvoid DistributedManagerBase::get_unassigned(IVector_d &flagArray, IVector_d &B2L_map, IVector_d &partition_flags, int size, int global_size /*, int rank*/)\n{\n int num_blocks = std::min(4096, (size + 191) / 192);\n get_unassigned_kernel <<< num_blocks, 192>>>(flagArray.raw(),\n B2L_map.raw(),\n partition_flags.raw(), size, global_size /*, rank*/);\n cudaCheckError();\n}\n\ntemplate < class TConfig >\nvoid DistributedManagerBase::set_unassigned(IVector_d &partition_flags, IVector_d &partition_renum, IVector_d &B2L_map, IVector_d &renumbering, int size, int max_element, int global_size /*, int rank*/)\n{\n int num_blocks = std::min(4096, (size + 191) / 192);\n set_unassigned_kernel <<< num_blocks, 192>>>(partition_flags.raw(),\n partition_renum.raw(),\n B2L_map.raw(),\n renumbering.raw(),\n size, max_element, global_size /*,rank*/);\n cudaCheckError();\n}\n\n\ntemplate \ninline void DistributedManagerBase::set_initialized(IVector &row_offsets)\n{\n // For P and R sizes the sizes are fixed at creation\n if(m_fixed_view_size)\n {\n return;\n }\n\n if (neighbors.size() > 0)\n {\n //distributed: cache num_rows/num_nz for different views\n _num_rows_interior = _num_interior_nodes;\n _num_nz_interior = row_offsets[_num_rows_interior];\n _num_rows_owned = _num_interior_nodes + _num_boundary_nodes;\n _num_nz_owned = row_offsets[_num_rows_owned];\n _num_rows_full = halo_offsets[neighbors.size()];\n\n if (_num_rows_full >= row_offsets.size())\n {\n _num_nz_full = row_offsets[row_offsets.size() - 1];\n }\n else\n {\n _num_nz_full = row_offsets[_num_rows_full];\n }\n\n _num_rows_all = halo_offsets[halo_offsets.size() - 1];\n _num_nz_all = _num_nz_full;\n }\n else\n {\n _num_rows_interior = _num_interior_nodes;\n _num_nz_interior = row_offsets[_num_rows_interior];\n _num_rows_owned = _num_interior_nodes;\n _num_nz_owned = row_offsets[_num_rows_owned];\n _num_rows_full = _num_rows_owned;\n _num_nz_full = _num_nz_owned;\n _num_rows_all = _num_rows_owned;\n _num_nz_all = _num_nz_owned;\n }\n}\n\ntemplate \nvoid DistributedManagerBase::createAggregatesRenumbering(IVector_h &renumbering, IVector_h_vector &B2L_maps, int size, int num_neighbors, int &num_interior_aggregates, int &num_boundary_aggregates, int num_rings)\n{\n createAggRenumbering(renumbering, B2L_maps, size, num_neighbors, num_interior_aggregates, num_boundary_aggregates, num_rings);\n}\n\ntemplate \nvoid DistributedManagerBase::createAggregatesRenumbering(IVector_d &renumbering, IVector_d_vector &B2L_maps, int size, int num_neighbors, int &num_interior_aggregates, int &num_boundary_aggregates, int num_rings)\n{\n createAggRenumbering(renumbering, B2L_maps, size, num_neighbors, num_interior_aggregates, num_boundary_aggregates, num_rings);\n}\n\ntemplate \ntemplate \nvoid DistributedManagerBase::createAggRenumbering(IVector_hd &renumbering, std::vector &B2L_maps, int size, int num_neighbors, int &num_interior_aggregates, int &num_boundary_aggregates, int num_rings)\n{\n if (num_rings != 1)\n {\n FatalError(\"num_rings > 1 not supported in consolidation\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n //int num_neighbors = this->neighbors.size();\n if (num_neighbors == 0)\n {\n num_boundary_aggregates = 0;\n num_interior_aggregates = size;\n return;\n }\n\n //initial size to size+1 so we have the total size after a scan\n int global_size = size;\n renumbering.resize(size + 1);\n //\n // Step 1 - in the main matrix, separate interior and boundary nodes (1/0 in flagArray), renumber interior ones with an exclusive scan\n //\n IVector_hd flagArray(size + 1);\n thrust_wrapper::fill(flagArray.begin(), flagArray.begin() + size + 1, 1);\n cudaCheckError();\n\n //sets 1 for interior nodes, 0 for boundary node\n for (int i = 0; i < num_neighbors; i++ )\n {\n int size = B2L_maps[i].size();\n remove_boundary(flagArray, B2L_maps[i], size);\n }\n\n //gets the renumbering of interior nodes\n thrust_wrapper::exclusive_scan(flagArray.begin(), flagArray.begin() + size + 1, renumbering.begin());\n cudaCheckError();\n //\n // Step 2 - Renumber nodes that are in the boundary, stepping through each B2L map, and renumbering ones that have not been renumbered yet\n //\n //what is the biggest B2L size\n INDEX_TYPE max_size = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n max_size = max_size > B2L_maps[i].size() ? max_size : B2L_maps[i].size();\n }\n\n //allocate work vectors (should be pretty small)\n IVector_hd partition_flags(max_size);\n IVector_hd partition_renum(max_size);\n //the number of renumbered nodes so far\n int max_element = renumbering[size];\n num_interior_aggregates = max_element;\n num_boundary_aggregates = size - max_element;\n renumbering.resize(size);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n //find nodes that are part of the current boundary and they haven't been renumbered yet\n thrust_wrapper::fill(partition_flags.begin(), partition_flags.begin() + max_size, 0);\n int size = B2L_maps[i].size();\n get_unassigned(flagArray, B2L_maps[i], partition_flags, size, global_size/*,0*/);\n //calculate the local renumbering (within this boundary region) of these nodes\n thrust_wrapper::exclusive_scan(partition_flags.begin(), partition_flags.begin() + max_size, partition_renum.begin());\n //apply renumbering to the big numbering table\n set_unassigned(partition_flags, partition_renum, B2L_maps[i], renumbering, size, max_element, global_size/*,0*/);\n //update the number of renumbered nodes\n max_element += partition_renum[max_size - 1] + partition_flags[max_size - 1];\n }\n\n cudaCheckError();\n}\n\n\ntemplate \ninline DistributedManagerBase::DistributedManagerBase(Matrix &a) :\n m_fine_level_comms(NULL), A(&a), m_pinned_buffer_size(0), m_pinned_buffer(NULL), _num_interior_nodes(0), _num_boundary_nodes(0), _comms(NULL), has_B2L(false),\n neighbors(_neighbors), B2L_maps(_B2L_maps), L2H_maps(_L2H_maps), B2L_rings(_B2L_rings),\n halo_rows_ref_count(0), halo_btl_ref_count(0), halo_ranges(_halo_ranges), halo_ranges_h(_halo_ranges_h), part_offsets(_part_offsets), part_offsets_h(_part_offsets_h), halo_rows(NULL), halo_btl(NULL), m_is_root_partition(false), m_is_glued(false), m_is_fine_level_glued(false), m_is_fine_level_consolidated(false), m_is_fine_level_root_partition(false), m_use_cuda_ipc_consolidation(false), m_fixed_view_size(false)\n{\n hipEventCreate(&comm_event);\n cudaCheckError();\n hipStreamCreateWithFlags(&m_int_stream, hipStreamNonBlocking);\n cudaCheckError();\n hipStreamCreateWithFlags(&m_bdy_stream, hipStreamNonBlocking);\n cudaCheckError();\n this->createComms(A->getResources());\n int my_id = this->getComms()->get_global_id();\n int num_parts = this->getComms()->get_num_partitions();\n this->set_global_id(my_id);\n this->set_num_partitions(num_parts);\n}\n\ntemplate \nvoid DistributedManager >::generatePoisson7pt(int nx, int ny, int nz, int P, int Q, int R)\n{\n int my_id = this->getComms()->get_global_id();\n int p, q, r;\n\n if (nx < P || ny < Q || nz < R)\n {\n FatalError(\"(nx < P) or (ny < Q) or (nz < R) not supported\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n /* compute p,q,r from P,Q,R and myid */\n p = my_id % P; // Position in x direction\n q = (( my_id - p) / P) % Q; // Position in y\n r = ( my_id - p - P * q) / ( P * Q ); // Position in z\n // Create A.row_indices, A.col_indices, A.values, A.diag\n int num_rows = nx * ny * nz;\n int num_nonzeros = num_rows * 7; // Ignoring any boundary, 7 nnz per row\n int num_substract = 0;\n\n if (p == 0) { num_substract += ny * nz; }\n\n if (p == P - 1) { num_substract += ny * nz; }\n\n if (q == 0) { num_substract += nx * nz; }\n\n if (q == Q - 1) { num_substract += nx * nz; }\n\n if (r == 0) { num_substract += nx * ny; }\n\n if (r == R - 1) { num_substract += nx * ny; }\n\n num_nonzeros -= num_substract;\n int num_halo_nodes = 2 * (ny * nz + nx * nz + nx * ny) - num_substract;\n this->local_to_global_map.resize(num_halo_nodes);\n this->A->set_initialized(0);\n this->A->resize(0, 0, 0, 1, 1, 1);\n this->A->addProps(CSR);\n this->A->resize(num_rows, num_rows + num_halo_nodes, num_nonzeros, 1, 1, 1);\n const int cta_size = 128;\n const int grid_size = std::min( 4096, (num_rows + cta_size - 1) / cta_size );\n poisson7pt_count_row_len <<< grid_size, cta_size>>>(this->A->row_offsets.raw(), nx, ny, nz, p, q, r, P, Q, R, num_rows);\n cudaCheckError();\n thrust_wrapper::exclusive_scan(this->A->row_offsets.begin(), this->A->row_offsets.end(), this->A->row_offsets.begin());\n cudaCheckError();\n // Now set nonzeros columns and values\n // TODO: vectorize this\n const int grid_size2 = std::min( 4096, (num_rows + cta_size - 1) / cta_size );\n poisson7pt_set_col_values <<< grid_size2, cta_size>>>\n (this->A->row_offsets.raw(),\n this->A->col_indices.raw(),\n this->A->values.raw(),\n this->A->diag.raw(),\n this->local_to_global_map.raw(),\n nx, ny, nz,\n p, q, r,\n P, Q, R,\n num_rows);\n cudaCheckError();\n // fill parts_offsets_h\n // All ranks have same number of nodes\n int num_ranks = P * Q * R;\n this->part_offsets_h.resize(num_ranks + 1);\n this->part_offsets_h[0] = (int64_t) 0;\n\n for (int i = 1; i < num_ranks + 1; i++)\n {\n this->part_offsets_h[i] = this->part_offsets_h[i - 1] + (int64_t) num_rows;\n }\n\n // Device to host copy\n this->part_offsets = this->part_offsets_h;\n this->num_rows_global = P * Q * R * nx * ny * nz;\n// this->A->set_initialized(1);\n}\n\ntemplate \ntemplate \nvoid DistributedManager >::loadDistributed_SetOffsets(\n int num_ranks, int num_rows_global, const t_colIndex* partition_offsets)\n{\n // fill part offsets internal data structures\n this->part_offsets_h.resize(num_ranks + 1);\n\n for (int i = 0; i <= num_ranks; i++)\n {\n this->part_offsets_h[i] = partition_offsets[i];\n }\n // copy to device\n this->part_offsets = this->part_offsets_h;\n // set num of global rows\n this->num_rows_global = num_rows_global;\n cudaCheckError();\n}\n\ntemplate \ntemplate \nstd::map DistributedManager >::loadDistributed_LocalToGlobal(int num_rows, I64Vector_h &off_diag_cols)\n{\n // sort global column indices\n amgx::thrust::sort(off_diag_cols.begin(), off_diag_cols.end());\n // find unique columns and set local <-> global mappings\n // 1) Removed unneeded vector 2) Create map on host first, upload later (less thrust calls)\n I64Vector_h local_to_global_h;\n std::map global_to_local; // temporary\n\n if (off_diag_cols.size() > 0)\n {\n global_to_local[off_diag_cols[0]] = num_rows;\n local_to_global_h.push_back(off_diag_cols[0]);\n }\n\n for (int i = 1; i < off_diag_cols.size(); i++)\n {\n if (off_diag_cols[i] != off_diag_cols[i - 1])\n {\n global_to_local[off_diag_cols[i]] = num_rows + local_to_global_h.size();\n local_to_global_h.push_back(off_diag_cols[i]);\n }\n }\n // Upload finished map in one piece\n this->local_to_global_map.resize(local_to_global_h.size());\n amgx::thrust::copy(local_to_global_h.begin(), local_to_global_h.end(), this->local_to_global_map.begin());\n return global_to_local;\n}\n\ntemplate \nvoid DistributedManager >::loadDistributed_InitLocalMatrix(\n IVector_h local_col_indices,\n int num_rows,\n int num_nonzeros,\n const int block_dimx,\n const int block_dimy,\n const int *row_offsets,\n const mat_value_type *values,\n const void *diag)\n{\n // init local matrix\n this->A->set_initialized(0);\n this->A->resize(0, 0, 0, 1, 1, 1);\n this->A->addProps(CSR);\n\n if (diag)\n {\n this->A->addProps(DIAG);\n }\n\n this->A->resize(num_rows, num_rows + this->local_to_global_map.size(), num_nonzeros, block_dimx, block_dimy, 1);\n cudaCheckError();\n // set local matrix\n amgx::thrust::copy(row_offsets, row_offsets + num_rows + 1, this->A->row_offsets.begin());\n this->A->col_indices = local_col_indices;\n\n amgx::thrust::copy(values, values + num_nonzeros * block_dimx * block_dimy, this->A->values.begin());\n cudaCheckError();\n\n // setup diagonal\n if (diag)\n {\n hipMemcpy(this->A->values.raw() + this->A->diagOffset()*this->A->get_block_size(), diag, sizeof(mat_value_type) * num_rows * block_dimx * block_dimy, hipMemcpyDefault);\n cudaCheckError();\n }\n else\n {\n this->A->computeDiagonal();\n }\n cudaCheckError();\n}\n\ntemplate \ntemplate \nvoid DistributedManager >::loadDistributedMatrixPartitionVec(\n int num_rows,\n int num_nonzeros,\n const int block_dimx,\n const int block_dimy,\n const int *row_offsets,\n const t_colIndex *col_indices,\n const mat_value_type *values,\n int num_ranks,\n int num_rows_global,\n const void *diag,\n const int *partition)\n{\n // fetch my rank\n int my_id = this->getComms()->get_global_id();\n // setup partition vector\n IVector_h partitionVec(num_rows_global);\n\n if (partition == NULL)\n {\n IVector_h rowCounts(num_ranks);\n this->getComms()->all_gather(num_rows, rowCounts, 1);\n\n int p = 0;\n for (int i = 0; i < num_ranks; ++i)\n {\n for (int j = 0; j < rowCounts[i]; ++j)\n {\n partitionVec[p++] = i;\n }\n }\n }\n else\n {\n // use existing partition info\n for (int i = 0; i < num_rows_global; i++)\n {\n partitionVec[i] = partition[i];\n }\n }\n\n // compute partition offsets (based on number of elements per partition). Will be modified when calculating partition map.\n t_colIndex *partition_offsets = (t_colIndex *)calloc(num_ranks + 1, sizeof(t_colIndex));\n\n for (int i = 0; i < num_rows_global; i++)\n {\n int pvi = partitionVec[i];\n partition_offsets[pvi + 1]++;\n }\n\n thrust_wrapper::inclusive_scan(partition_offsets, partition_offsets + num_ranks + 1, partition_offsets);\n\n loadDistributed_SetOffsets(num_ranks, num_rows_global, partition_offsets);\n\n // compute partition map (which tells you how the global elements are mapped into the partitions)\n t_colIndex *partition_map = (t_colIndex *)calloc(num_rows_global, sizeof(t_colIndex));\n\n for (int i = 0; i < num_rows_global; i++)\n {\n int pvi = partitionVec[i];\n t_colIndex poi = partition_offsets[pvi];\n partition_map[poi] = i;\n partition_offsets[pvi]++;\n }\n free(partition_offsets);\n\n // compute the inverse partition map\n t_colIndex *ipartition_map = (t_colIndex *)calloc(num_rows_global, sizeof(t_colIndex));\n\n for (int i = 0; i < num_rows_global; i++)\n {\n ipartition_map[partition_map[i]] = i;\n }\n free(partition_map);\n\n int h_cidx_allocated = 0;\n const t_colIndex *h_col_indices_global = (const t_colIndex *)this->getHostPointerForData(col_indices, num_nonzeros * sizeof(t_colIndex), &h_cidx_allocated);\n // gather all off-diag columns\n I64Vector_h off_diag_cols;\n\n for (int i = 0; i < num_nonzeros; i++)\n {\n if (partitionVec[h_col_indices_global[i]] != my_id)\n {\n off_diag_cols.push_back(ipartition_map[h_col_indices_global[i]]);\n }\n }\n\n auto global_to_local = loadDistributed_LocalToGlobal(num_rows, off_diag_cols);\n\n // set 1, then scan to compute local row indices\n IVector_h my_indices(num_rows_global);\n\n for (int i = 0; i < num_nonzeros; i++)\n {\n if (partitionVec[h_col_indices_global[i]] == my_id) // find my local columns and set to 1\n {\n my_indices[ipartition_map[h_col_indices_global[i]]] = 1;\n }\n }\n\n thrust_wrapper::exclusive_scan(my_indices.begin(), my_indices.end(), my_indices.begin());\n // remap colums to local\n IVector_h local_col_indices(num_nonzeros);\n\n for (int i = 0; i < num_nonzeros; i++)\n {\n if (partitionVec[h_col_indices_global[i]] != my_id)\n {\n // off-diag\n local_col_indices[i] = global_to_local[ipartition_map[h_col_indices_global[i]]];\n }\n else\n {\n // diag\n local_col_indices[i] = my_indices[ipartition_map[h_col_indices_global[i]]];\n }\n }\n free(ipartition_map);\n\n loadDistributed_InitLocalMatrix(local_col_indices, num_rows, num_nonzeros, block_dimx, block_dimy, row_offsets, values, diag);\n\n cudaCheckError();\n\n // don't free possibly allocated pinned buffer, since it could be used later. if it would not - it would be deallocated automatically\n /*if (h_cidx_allocated)\n {\n free((void*)h_col_indices_global);\n }*/\n}\n\ntemplate \ntemplate \nvoid DistributedManager >::loadDistributedMatrixPartitionOffsets(\n int num_rows,\n int num_nonzeros,\n const int block_dimx,\n const int block_dimy,\n const int *row_offsets,\n const t_colIndex *col_indices,\n const mat_value_type *values,\n int num_ranks,\n int num_rows_global,\n const void *diag,\n const t_colIndex *partition_offsets)\n{\n // fetch my rank\n int my_id = this->getComms()->get_global_id();\n // sanity check, cheap to perform, and helps prevent harder-to-debug errors later on\n if (!std::is_sorted(partition_offsets, partition_offsets + num_ranks + 1)) {\n FatalError(\"Partition offsets are not sorted.\", AMGX_ERR_BAD_PARAMETERS);\n }\n loadDistributed_SetOffsets(num_ranks, num_rows_global, partition_offsets);\n\n // Create predicate to determine if a column is in the local diagonal block\n t_colIndex my_first_col = this->part_offsets_h[my_id];\n t_colIndex one_past_my_last_col = this->part_offsets_h[my_id + 1];\n auto in_local_diagonal_block = [my_first_col, one_past_my_last_col](const t_colIndex col_index) {\n return col_index >= my_first_col && col_index < one_past_my_last_col;\n };\n\n int h_cidx_allocated = 0;\n const t_colIndex *h_col_indices_global = (const t_colIndex *)this->getHostPointerForData(col_indices, num_nonzeros * sizeof(t_colIndex), &h_cidx_allocated);\n // gather all off-diag columns\n I64Vector_h off_diag_cols;\n for (int i = 0; i < num_nonzeros; i++)\n {\n if (!in_local_diagonal_block(h_col_indices_global[i]))\n {\n off_diag_cols.push_back(h_col_indices_global[i]);\n }\n }\n auto global_to_local = loadDistributed_LocalToGlobal(num_rows, off_diag_cols);\n // set 1, then scan to compute local row indices\n // \"coordinate-shift\" columns so they lie in much smaller range of my diagonal indices\n int diagonal_size = this->part_offsets_h[my_id + 1] - this->part_offsets_h[my_id];\n IVector_h my_indices(diagonal_size);\n for (int i = 0; i < num_nonzeros; i++)\n {\n t_colIndex col_index = h_col_indices_global[i];\n if (in_local_diagonal_block(h_col_indices_global[i])) // find my local columns and set to 1\n {\n // olumns that are on *my* diag partition cannot have an index from 0..num_rows_global\n // instead, part_offsets_h[my_id] <= col_index < part_offsets[my_id+1]\n col_index -= this->part_offsets_h[my_id];\n my_indices[col_index] = 1;\n }\n }\n thrust_wrapper::exclusive_scan(my_indices.begin(), my_indices.end(), my_indices.begin());\n\n // remap colums to local\n IVector_h local_col_indices(num_nonzeros);\n for (int i = 0; i < num_nonzeros; i++)\n {\n t_colIndex col_index = h_col_indices_global[i];\n if (!in_local_diagonal_block(col_index))\n {\n // off-diag\n local_col_indices[i] = global_to_local[col_index];\n }\n else\n {\n // diag\n col_index -= this->part_offsets_h[my_id];\n local_col_indices[i] = my_indices[col_index];\n }\n }\n loadDistributed_InitLocalMatrix(local_col_indices, num_rows, num_nonzeros, block_dimx, block_dimy, row_offsets, values, diag);\n}\n\ntemplate \ntemplate \nvoid DistributedManager >::loadDistributedMatrix(\n int num_rows,\n int num_nonzeros,\n const int block_dimx,\n const int block_dimy,\n const int *row_offsets,\n const t_colIndex *col_indices,\n const mat_value_type *values,\n int num_ranks,\n int num_rows_global,\n const void *diag,\n const MatrixDistribution &dist)\n{\n using PI = MatrixDistribution::PartitionInformation;\n switch (dist.getPartitionInformationStyle()) {\n case PI::PartitionVec:\n loadDistributedMatrixPartitionVec(num_rows, num_nonzeros, block_dimx, block_dimy, \n row_offsets, col_indices, values, num_ranks, num_rows_global, diag, (const int*) dist.getPartitionData());\n break;\n case PI::PartitionOffsets:\n loadDistributedMatrixPartitionOffsets(num_rows, num_nonzeros, block_dimx, block_dimy, \n row_offsets, col_indices, values, num_ranks, num_rows_global, diag, (const t_colIndex*) dist.getPartitionData());\n break;\n default:\n FatalError(\"Unsupported partitioning data format used with loadDistributedMatrix\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::renumberMatrixOneRing(int update_neighbours)\n{\n FatalError(\"Distributed classical AMG not implemented on host\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid DistributedManager >::renumberMatrixOneRing(int update_neighbours)\n{\n // Step 1: Using halo_ranges, flag neighbors and at the same time, flag halo_nodes (flag_halo_nodes_local)\n int my_id = this->global_id();\n int num_parts = this->get_num_partitions();\n this->set_base_index(this->part_offsets_h[my_id]);\n this->set_index_range(this->part_offsets_h[my_id + 1] - this->part_offsets_h[my_id]);\n DistributedArranger *prep = new DistributedArranger;\n\n // Create/update list of neighbors\n if (update_neighbours)\n {\n typedef typename TConfig::template setVecPrec::Type i64vec_value_type;\n typedef Vector I64Vector;\n typedef typename Matrix::MVector MVector;\n std::vector halo_row_offsets(this->neighbors.size());\n std::vector halo_global_indices(this->neighbors.size());\n std::vector halo_values(this->neighbors.size());\n prep->create_halo_rows_global_indices(*(this->A), halo_row_offsets, halo_global_indices, halo_values);\n prep->update_neighbors_list(*(this->A), this->neighbors, this->halo_ranges_h, this->halo_ranges, this->part_offsets_h, this->part_offsets, halo_row_offsets, halo_global_indices);\n }\n else\n {\n prep->create_neighbors_v2(*(this->A));\n }\n\n this->getComms()->set_neighbors(this->neighbors.size());\n // Create B2L_maps and L2H_maps\n prep->create_boundary_lists_v3(*(this->A));\n // halo_offsets\n int neighbors = this->A->manager->num_neighbors();\n int A_num_rows, offset;\n this->A->getOffsetAndSizeForView(OWNED, &offset, &A_num_rows);\n this->halo_offsets.resize(neighbors + 1, 0);\n this->halo_offsets[0] = A_num_rows;\n\n for (int i = 0; i < neighbors; i++)\n {\n this->halo_offsets[i + 1] = this->halo_offsets[i] + this->B2L_maps[i].size();\n }\n\n this->getComms()->exchange_vectors(this->A->manager->B2L_maps, *(this->A), 0);\n // Initialize B2L_rings\n int num_neighbors = this->neighbors.size();\n this->B2L_rings.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n this->B2L_rings[i].resize(2);\n this->B2L_rings[i][0] = 0;\n this->B2L_rings[i][1] = this->B2L_maps[i].size();\n }\n\n prep->initialize_B2L_maps_offsets(*(this->A), 1);\n delete prep;\n //Use the exchanged halo row matrices and the boundary/halo index lists to renumber and consolidate the matrix\n // Step 5: renumber all owned rows and columns\n this->reorder_matrix_owned();\n // Step 6: renumber local_to_global_map\n int num_owned_rows = this->A->manager->halo_offsets[0];\n int size_one_ring;\n this->A->getOffsetAndSizeForView(FULL, &offset, &size_one_ring);\n I64Vector_d global_col_indices(size_one_ring);\n thrust_wrapper::sequence(global_col_indices.begin(), global_col_indices.begin() + num_owned_rows, this->base_index() );\n cudaCheckError();\n global_col_indices.dirtybit = 1;\n this->exchange_halo(global_col_indices, global_col_indices.tag);\n thrust_wrapper::copy(global_col_indices.begin() + num_owned_rows, global_col_indices.begin() + size_one_ring, this->local_to_global_map.begin(), this->get_int_stream(), true);\n cudaCheckError();\n}\n\n\ntemplate \nvoid DistributedManager >::renumber_P_R(Matrix_h &P, Matrix_h &R, Matrix_h &A_fine)\n{\n FatalError(\"Distributed classical AMG not implemented on host\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid DistributedManager >::renumber_P_R(Matrix_d &P, Matrix_d &R, Matrix_d &A_fine)\n{\n int cta_size = 256;\n int num_owned_fine_pts = A_fine.manager->halo_offsets[0];\n int num_owned_coarse_pts, offset;\n // matrix Ac\n this->A->getOffsetAndSizeForView(OWNED, &offset, &num_owned_coarse_pts);\n // Renumber the owned col indices of P (not the halo columns ,since P.manager was created assunming some other numbering)\n int nnz_owned_fine_pts = P.row_offsets[num_owned_fine_pts];\n int num_blocks_fine = std::min(4096, (nnz_owned_fine_pts + cta_size - 1) / cta_size);\n\n if (num_blocks_fine > 0)\n {\n renumber_P_col_indices <<< num_blocks_fine, cta_size>>>(P.col_indices.raw(), this->renumbering.raw(), num_owned_coarse_pts, nnz_owned_fine_pts);\n cudaCheckError();\n }\n\n // Renumber the B2L_maps of P\n for (int i = 0; i < P.manager->neighbors.size(); i++)\n {\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(this->renumbering.begin(), P.manager->B2L_maps[i].begin()),\n amgx::thrust::make_permutation_iterator(this->renumbering.begin(), P.manager->B2L_maps[i].end()),\n P.manager->B2L_maps[i].begin());\n }\n\n cudaCheckError();\n // Don't renumber the L2H_maps or the halo\n // Renumber the local_to_global_map of matrix P (since neighbors renumbered their owned rows)\n // Swap owned rows of R\n IVector new_row_offsets(R.row_offsets.size());\n int insert = 0;\n // Only renumber the owned rows\n int num_blocks_owned = std::min(4096, (num_owned_coarse_pts + cta_size - 1) / cta_size);\n\n if (num_blocks_owned > 0)\n {\n calc_rowlen_reorder <<< num_blocks_owned, cta_size >>>(R.row_offsets.raw(), new_row_offsets.raw(), this->renumbering.raw(), num_owned_coarse_pts, insert);\n cudaCheckError();\n }\n\n thrust_wrapper::exclusive_scan(new_row_offsets.begin(), new_row_offsets.begin() + num_owned_coarse_pts + 1, new_row_offsets.begin());\n cudaCheckError();\n // Copy the row_offsets for halo rows\n amgx::thrust::copy(R.row_offsets.begin() + num_owned_coarse_pts, R.row_offsets.end(), new_row_offsets.begin() + num_owned_coarse_pts);\n cudaCheckError();\n // Reorder the rows of R (no need to reorder the column indices)\n int new_nnz = new_row_offsets[new_row_offsets.size() - 1];\n int halo_offset = new_row_offsets[num_owned_coarse_pts];\n typedef typename MatPrecisionMap::Type ValueTypeA;\n VVector new_values(new_nnz * R.get_block_size(), types::util< ValueTypeA >::get_zero());\n IVector new_col_indices(new_nnz, 0);\n int num_blocks_total = std::min(4096, (R.get_num_rows() + cta_size - 1) / cta_size);\n\n if (num_blocks_total > 0)\n {\n reorder_R_matrix <32> <<< num_blocks_total, 512>>>(R.row_offsets.raw(), R.col_indices.raw(), R.values.raw(), new_row_offsets.raw(), new_col_indices.raw(), new_values.raw(), this->renumbering.raw(), R.get_block_size(), R.get_num_rows(), num_owned_coarse_pts);\n cudaCheckError();\n }\n\n R.col_indices.swap(new_col_indices);\n R.row_offsets.swap(new_row_offsets);\n R.values.swap(new_values);\n\n // Renumber the local_to_global_map (since neighbors have changed their owned numbering)\n if (P.manager->neighbors.size() != 0)\n {\n int size_one_ring = P.manager->halo_offsets[P.manager->neighbors.size()];\n I64Vector_d global_col_indices(size_one_ring);\n thrust_wrapper::sequence(global_col_indices.begin(), global_col_indices.begin() + num_owned_coarse_pts, this->base_index());\n cudaCheckError();\n global_col_indices.dirtybit = 1;\n P.manager->exchange_halo(global_col_indices, global_col_indices.tag);\n thrust_wrapper::copy(global_col_indices.begin() + num_owned_coarse_pts, global_col_indices.begin() + size_one_ring, P.manager->local_to_global_map.begin(), this->get_int_stream(), true);\n cudaCheckError();\n }\n\n DistributedArranger *prep = new DistributedArranger;\n prep->initialize_B2L_maps_offsets(P, 1);\n delete prep;\n}\n\n\ntemplate \nvoid DistributedManager >::createOneRingB2Lmaps()\n{\n // Input:\n // a matrix with N rows, whose column indices are local indices from 0 to N+M-1,\n // where M is a number of 1-ring halo vertices\n // The matrix also contains array \"local_to_global_map\" of size M, which stores the global index of each halo index\n // Ex: assuming a column has index N+K, where 0 <= K < M, then it's global id is local_to_global_map[K]\n // The matrix also contains part_offsets_h and part_offsets array, which stores where each partition begins\n // Output:\n // This function creates all the necessary data to to 1-ring exchanges\n // i.e. list of 1-ring neighbors, B2L_maps for 1-ring, halo_offsets for 1-ring,\n // Also, the function reorders the halo indices, such that 1-ring indices are in the order\n // of neighbors, and therefore, exchange_halo doesn't have to be changed (i.e. L2H = identity)\n // What is does:\n // Based on the global indices of its halo vertices, count the number of neighbors\n // For each neighbor, receive the halo indices that will be needed by neighbor\n // From those, create B2L_maps[0], which contains for all neighbors\n // This function assumes that:\n // part_offset is defined\n\n // B2L_maps\n int my_id = this->global_id();\n int num_parts = this->get_num_partitions();\n this->set_base_index(this->part_offsets_h[my_id]);\n this->set_index_range(this->part_offsets_h[my_id + 1] - this->part_offsets_h[my_id]);\n DistributedArranger *prep = new DistributedArranger;\n // This function creates the array neighbors, which contains a list of partitions to which data\n // needs to be sent and/or received\n prep->create_neighbors_v2(*(this->A));\n // Here change the manager if some partitions have no neighbors\n this->getComms()->set_neighbors(this->neighbors.size());\n prep->create_B2L_one_ring(*(this->A));\n delete prep;\n}\n\ntemplate \nvoid DistributedManager >::createOneRingHaloRows()\n{\n // Input:\n // A matrix with 1-ring B2L_maps, 1-ring halo_offsets\n // Outputs:\n // A matrix with: 1-ring rows,\n // 2-ring B2L_maps,\n // 2-ring halo_offsets\n // 2-ring neighbors\n // Implement here:\n // Look at function create_B2L_from_maps, which calls create_rings, create_halo_btl, create_halo_rows and comms->exchange_matrix_halo\n DistributedArranger *prep = new DistributedArranger;\n prep->create_one_ring_halo_rows(*(this->A));\n // I believe this can be removed since we don't use masked SpMV anymore\n prep->createRowsLists(*(this->A), false);\n delete prep;\n // this is not necessary anymore becasue we don't use latency hiding\n // however in future we might want to get back to this in case we want to use latency hiding\n //this->reorder_matrix();\n}\n\ntemplate \ninline DistributedManagerBase::DistributedManagerBase(\n Matrix &a,\n INDEX_TYPE allocated_halo_depth,\n INDEX_TYPE num_import_rings,\n int num_neighbors,\n const VecInt_t *neighbors_) : m_fine_level_comms(NULL), A(&a), m_pinned_buffer_size(0), m_pinned_buffer(NULL), _num_interior_nodes(0), _num_boundary_nodes(0), _comms(NULL), has_B2L(false), neighbors(_neighbors), halo_rows_ref_count(0), halo_rows(NULL), halo_btl_ref_count(0), halo_btl(NULL), halo_ranges(_halo_ranges), halo_ranges_h(_halo_ranges_h), part_offsets(_part_offsets), part_offsets_h(_part_offsets_h),\n B2L_maps(_B2L_maps), L2H_maps(_L2H_maps), B2L_rings(_B2L_rings), m_is_root_partition(false), m_is_glued(false), m_is_fine_level_glued(false), m_is_fine_level_consolidated(false), m_is_fine_level_root_partition(false), m_use_cuda_ipc_consolidation(false), m_fixed_view_size(false)\n{\n hipStreamCreateWithFlags(&m_int_stream, hipStreamNonBlocking);\n cudaCheckError();\n hipStreamCreateWithFlags(&m_bdy_stream, hipStreamNonBlocking);\n cudaCheckError();\n\n if (num_import_rings != 1)\n {\n FatalError(\"num_rings > 1 not supported in fine_level consolidation\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (allocated_halo_depth != 1)\n {\n FatalError(\"allocated_halo_depth > 1 not supported in fine_level consolidation\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n this->set_num_halo_rings(num_import_rings);\n neighbors.resize(num_neighbors);\n hipMemcpy(neighbors.raw(), neighbors_, num_neighbors * sizeof(VecInt_t), hipMemcpyDefault);\n cudaCheckError();\n}\n\ntemplate \ninline void DistributedManagerBase::cacheMaps(const VecInt_t *b2l_maps, const VecInt_t *b2l_ptrs, const VecInt_t *l2h_maps, const VecInt_t *l2h_ptrs)\n{\n int num_neighbors = this->neighbors.size();\n this->cached_B2L_maps.resize(num_neighbors);\n this->cached_L2H_maps.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = b2l_ptrs[i + 1] - b2l_ptrs[i];\n this->cached_B2L_maps[i].resize(size);\n int count = 0;\n\n for (int j = b2l_ptrs[i]; j < b2l_ptrs[i + 1]; j++)\n {\n this->cached_B2L_maps[i][count] = b2l_maps[j];\n count++;\n }\n\n size = l2h_ptrs[i + 1] - l2h_ptrs[i];\n this->cached_L2H_maps[i].resize(size);\n count = 0;\n\n for (int j = l2h_ptrs[i]; j < l2h_ptrs[i + 1]; j++)\n {\n this->cached_L2H_maps[i][count] = l2h_maps[j];\n count++;\n }\n }\n}\n\ntemplate \ninline void DistributedManagerBase::cacheMapsOneRing()\n{\n int num_neighbors = this->neighbors.size();\n this->cached_B2L_maps.resize(num_neighbors);\n this->cached_L2H_maps.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n this->cached_B2L_maps[i] = this->B2L_maps[i];\n this->cached_L2H_maps[i] = this->L2H_maps[i];\n }\n}\n\ntemplate \ninline void DistributedManagerBase::cacheMapsOneRing(const VecInt_t **b2l_maps, const VecInt_t *b2l_sizes, const VecInt_t **l2h_maps, const VecInt_t *l2h_sizes)\n{\n int num_neighbors = this->neighbors.size();\n this->cached_B2L_maps.resize(num_neighbors);\n this->cached_L2H_maps.resize(num_neighbors);\n // buffering in the case of GPU data. This shouldn't much affect performance\n std::vector b2l_buffer, l2h_buffer;\n std::vector b2l_sizes_buffer, l2h_sizes_buffer;\n b2l_buffer.resize(num_neighbors);\n l2h_buffer.resize(num_neighbors);\n b2l_sizes_buffer.resize(num_neighbors);\n l2h_sizes_buffer.resize(num_neighbors);\n hipError_t cuda_rc = hipMemcpy(&(b2l_sizes_buffer[0]), b2l_sizes, sizeof(VecInt_t) * num_neighbors, hipMemcpyDefault);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy b2l_sizes failed in distributed_manager\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = hipMemcpy(&(l2h_sizes_buffer[0]), l2h_sizes, sizeof(VecInt_t) * num_neighbors, hipMemcpyDefault);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy l2h_sizes failed in distributed_manager\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = hipMemcpy(&(b2l_buffer[0]), b2l_maps, sizeof(VecInt_t *) * num_neighbors, hipMemcpyDefault);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy b2l_maps failed in distributed_manager\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = hipMemcpy(&(l2h_buffer[0]), l2h_maps, sizeof(VecInt_t *) * num_neighbors, hipMemcpyDefault);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy l2h_maps failed in distributed_manager\", AMGX_ERR_CUDA_FAILURE);\n }\n\n // caching all of the maps\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = b2l_sizes_buffer[i];\n this->cached_B2L_maps[i].resize(size);\n cuda_rc = hipMemcpy(&(this->cached_B2L_maps[i][0]), b2l_buffer[i], sizeof(VecInt_t) * size, hipMemcpyDefault);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy cached_B2L_maps failed in distributed_manager\", AMGX_ERR_CUDA_FAILURE);\n }\n cudaCheckError();\n size = l2h_sizes_buffer[i];\n this->cached_L2H_maps[i].resize(size);\n hipMemcpy(&(this->cached_L2H_maps[i][0]), l2h_buffer[i], sizeof(VecInt_t) * size, hipMemcpyDefault);\n cudaCheckError();\n }\n}\n\ntemplate \nvoid DistributedManagerBase::setAConsolidationFlags( Matrix &in_A)\n{\n this->A = &in_A;\n AMG_Config *rsrc_cfg = this->A->getResources()->getResourcesConfig();\n std::string scope;\n int consolidate_flag, cuda_ipc_flag;\n rsrc_cfg->getParameter(\"fine_level_consolidation\", consolidate_flag, \"default\", scope);\n rsrc_cfg->getParameter(\"use_cuda_ipc_consolidation\", cuda_ipc_flag, \"default\", scope);\n this->m_is_fine_level_consolidated = (consolidate_flag != 0);\n this->m_use_cuda_ipc_consolidation = (cuda_ipc_flag != 0);\n}\n\ntemplate \nvoid DistributedManagerBase::uploadMatrix(int n, int nnz, int block_dimx, int block_dimy, const int *row_ptrs, const int *col_indices, const void *data, const void *diag, Matrix &in_A)\n{\n this->setAConsolidationFlags(in_A);\n\n if (this->m_is_fine_level_consolidated)\n {\n this->A->manager->consolidateAndUploadAll(n, nnz, block_dimx, block_dimy, row_ptrs, col_indices, data, diag, *(this->A));\n }\n else\n {\n this->A->manager->initializeUploadReorderAll(n, nnz, block_dimx, block_dimy, row_ptrs, col_indices, data, diag, *(this->A));\n }\n}\n\ntemplate \nvoid DistributedManagerBase::checkPinnedBuffer(size_t size)\n{\n if ((m_pinned_buffer_size < size) && (m_pinned_buffer != NULL))\n {\n amgx::memory::hipHostFree(m_pinned_buffer);\n m_pinned_buffer = NULL;\n m_pinned_buffer_size = 0;\n }\n\n if (m_pinned_buffer == NULL)\n {\n m_pinned_buffer_size = (size_t)(size * 1.1);\n amgx::memory::hipHostMalloc(&m_pinned_buffer, m_pinned_buffer_size);\n }\n}\n\n\ntemplate \nDistributedManagerBase::~DistributedManagerBase()\n{\n if (m_pinned_buffer != NULL)\n {\n amgx::memory::hipHostFree(m_pinned_buffer);\n }\n\n destroyComms();\n // from childrens:\n hipStreamDestroy(this->m_int_stream);\n cudaCheckError();\n hipStreamDestroy(this->m_bdy_stream);\n cudaCheckError();\n\n if (!this->halo_rows_ref_count && this->halo_rows != NULL)\n {\n delete this->halo_rows;\n this->halo_rows = NULL;\n }\n\n if (!this->halo_btl_ref_count && this->halo_btl != NULL)\n {\n delete this->halo_btl;\n this->halo_btl = NULL;\n }\n}\n\n// if pointer is host pointer - returns data. If it is device pointer - copies it to the m_pinned_buffer and returns pointer to m_pinned_buffer\ntemplate \nvoid *DistributedManagerBase::getHostPointerForData(void *ptr, size_t size, int *allocated)\n{\n hipError_t rc;\n hipPointerAttribute_t att;\n void *ptr_h;\n cudaCheckError();\n /* WARNING: We may accept the following types of allocation for ptr:\n 1. malloc [host memory]\n 2. cudaMalloc [device memory]\n 3. malloc + cudaHostRegister [AMGX_pin_memory/AMGX_unpin_memory host memory]\n 4. cudaHostAlloc [pinned host memory form the beginning]\n\n The correct way to conver these cases is the following:\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2, 3 or 4.\n }\n else{\n //you are in case 1.\n }\n The following pattern of checks should be implemented\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2 or 4.\n }\n else{\n st = cudaHostGetDevicePointer(ptr_on_device, ptr, 0);\n if (st == cudaSuccess){\n //you are in case 3.\n }\n else{\n //you are in case 1.\n }\n }\n The above pattern will be used whenever we need to process input data.\n\n Obs.: parameter size is in bytes and\n parameter allocated indicates whether memory was allocated\n and needs to be release later on. */\n /*\n // original implementation\n cudaPointerGetAttributes(&att, ptr);\n if (att.hostPointer == NULL)\n {\n checkPinnedBuffer(size);\n cudaMemcpy(m_pinned_buffer, ptr, size, cudaMemcpyDefault);\n cudaCheckError();\n return m_pinned_buffer;\n }\n else {\n return ptr;\n }\n */\n *allocated = 0;\n // get pointer to values on the device\n rc = hipPointerGetAttributes(&att, ptr);\n\n#if CUDART_VERSION >= 11000\n if(rc == hipSuccess && (att.type == hipMemoryTypeDevice || att.type == hipMemoryTypeManaged))\n {\n checkPinnedBuffer(size);\n rc = hipMemcpy(m_pinned_buffer, ptr, size, hipMemcpyDefault);\n cudaCheckError();\n\n if (rc != hipSuccess)\n {\n FatalError(\"Could not copy into the temporary (host) storage. Try pinning the memory to avoid the hipMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n ptr_h = m_pinned_buffer;\n *allocated = 1;\n hipGetLastError(); //to reset last error\n return ptr_h;\n }\n else\n {\n hipGetLastError(); //to reset last error\n return ptr;\n }\n#else\n if (rc == hipSuccess)\n {\n //you are in case 2 or 4 from the above comment.\n if (att.hostPointer == NULL)\n {\n //you are in case 2\n checkPinnedBuffer(size);\n rc = hipMemcpy(m_pinned_buffer, ptr, size, hipMemcpyDefault);\n cudaCheckError();\n\n if (rc != hipSuccess)\n {\n FatalError(\"Could not copy into the temporary (host) storage. Try pinning the memory to avoid the hipMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n ptr_h = m_pinned_buffer;\n *allocated = 1;\n }\n else\n {\n //you are in case 4\n ptr_h = ptr;\n }\n }\n else\n {\n //you are in case 1 or 3 from the above comment\n ptr_h = ptr;\n }\n\n hipGetLastError(); //to reset last error\n\n /* check for null pointers */\n if (ptr_h == NULL)\n {\n FatalError(\"Result of (host) allocation of required temporary storage is NULL. Try pinning the memory to reduce storage requirements.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n return ptr_h;\n#endif\n}\n\n// if pointer is host pointer - returns data. If it is device pointer - copies it to the m_pinned_buffer and returns pointer to m_pinned_buffer\ntemplate \nconst void *DistributedManagerBase::getHostPointerForData(const void *ptr, size_t size, int *allocated)\n{\n hipError_t rc;\n hipPointerAttribute_t att;\n void *ptr_h;\n cudaCheckError();\n /* WARNING: We may accept the following types of allocation for ptr:\n 1. malloc [host memory]\n 2. cudaMalloc [device memory]\n 3. malloc + cudaHostRegister [AMGX_pin_memory/AMGX_unpin_memory host memory]\n 4. cudaHostAlloc [pinned host memory form the beginning]\n\n The correct way to conver these cases is the following:\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2, 3 or 4.\n }\n else{\n //you are in case 1.\n }\n The following pattern of checks should be implemented\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2 or 4.\n }\n else{\n st = cudaHostGetDevicePointer(ptr_on_device, ptr, 0);\n if (st == cudaSuccess){\n //you are in case 3.\n }\n else{\n //you are in case 1.\n }\n }\n The above pattern will be used whenever we need to process input data.\n\n Obs.: parameter size is in bytes and\n parameter allocated indicates whether memory was allocated\n and needs to be release later on. */\n *allocated = 0;\n // get pointer to values on the device\n rc = hipPointerGetAttributes(&att, ptr);\n\n#if CUDART_VERSION >= 11000\n if(rc == hipSuccess && (att.type == hipMemoryTypeDevice || att.type == hipMemoryTypeManaged))\n {\n checkPinnedBuffer(size);\n rc = hipMemcpy(m_pinned_buffer, ptr, size, hipMemcpyDefault);\n cudaCheckError();\n\n if (rc != hipSuccess)\n {\n FatalError(\"Could not copy into the temporary (host) storage. Try pinning the memory to avoid the hipMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n ptr_h = m_pinned_buffer;\n *allocated = 1;\n hipGetLastError(); //to reset last error\n return ptr_h;\n }\n else\n {\n hipGetLastError(); //to reset last error\n return ptr;\n }\n#else\n if (rc == hipSuccess)\n {\n //you are in case 2 or 4 from the above comment.\n if (att.hostPointer == NULL)\n {\n //you are in case 2\n checkPinnedBuffer(size);\n rc = hipMemcpy(m_pinned_buffer, ptr, size, hipMemcpyDefault);\n cudaCheckError();\n\n if (rc != hipSuccess)\n {\n FatalError(\"Could not copy into the temporary (host) storage. Try pinning the memory to avoid the hipMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n ptr_h = m_pinned_buffer;\n *allocated = 1;\n hipGetLastError(); //to reset last error\n return ptr_h;\n }\n else\n {\n //you are in case 4\n hipGetLastError(); //to reset last error\n return ptr;\n }\n }\n else\n {\n hipGetLastError(); //to reset last error\n //you are in case 1 or 3 from the above comment\n return ptr;\n }\n#endif\n}\n\n\ntemplate \nvoid *DistributedManagerBase::getDevicePointerForData(void *ptr, size_t size, int *allocated)\n{\n hipError_t rc;\n hipPointerAttribute_t att;\n void *ptr_d;\n cudaCheckError();\n /* WARNING: We may accept the following types of allocation for ptr:\n 1. malloc [host memory]\n 2. cudaMalloc [device memory]\n 3. malloc + cudaHostRegister [AMGX_pin_memory/AMGX_unpin_memory host memory]\n 4. cudaHostAlloc [pinned host memory form the beginning]\n\n The correct way to conver these cases is the following:\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2, 3 or 4.\n }\n else{\n //you are in case 1.\n }\n The following pattern of checks should be implemented\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2 or 4.\n }\n else{\n st = cudaHostGetDevicePointer(ptr_on_device, ptr, 0);\n if (st == cudaSuccess){\n //you are in case 3.\n }\n else{\n //you are in case 1.\n }\n }\n The above pattern will be used whenever we need to process input data.\n\n Obs.: parameter size is in bytes and\n parameter allocated indicates whether memory was allocated\n and needs to be release later on. */\n *allocated = 0;\n // get pointer to values on the device\n rc = hipPointerGetAttributes(&att, ptr);\n\n#if CUDART_VERSION >= 11000\n if(rc == hipSuccess && att.type != hipMemoryTypeUnregistered)\n {\n // Received device registered data (device, managed, malloc+cudaHostRegister, or cudaMallocHost)\n ptr_d = (void*)att.devicePointer;\n\n if (ptr_d == NULL)\n {\n FatalError(\"Data passed is registered with device but devicePointer null.\", AMGX_ERR_INTERNAL);\n }\n }\n else\n {\n // Received unregistered or host allocated data (malloc)\n rc = amgx::memory::hipMallocAsync(&ptr_d, size);\n\n if (rc != hipSuccess)\n {\n FatalError(\"Could not allocate required temporary storage. Try pinning the memory to reduce storage requirements.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n rc = hipMemcpy(ptr_d, ptr, size, hipMemcpyDefault);\n cudaCheckError();\n\n if (rc != hipSuccess)\n {\n FatalError(\"Could not copy into the temporary storage. Try pinning the memory to avoid the hipMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n *allocated = 1;\n }\n#else\n\n if (rc == hipSuccess)\n {\n //you are in case 2 or 4 from the above comment.\n ptr_d = (void *)att.devicePointer;\n }\n else\n {\n //you are in case 1 or 3 from the above comment\n rc = hipHostGetDevicePointer(&ptr_d, ptr, 0);\n\n if (rc != hipSuccess)\n {\n //you are in case 1\n rc = amgx::memory::hipMallocAsync(&ptr_d, size);\n\n if (rc != hipSuccess)\n {\n FatalError(\"Could not allocate required temporary storage. Try pinning the memory to reduce storage requirements.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n rc = hipMemcpy(ptr_d, ptr, size, hipMemcpyDefault);\n cudaCheckError();\n\n if (rc != hipSuccess)\n {\n FatalError(\"Could not copy into the temporary storage. Try pinning the memory to avoid the hipMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n *allocated = 1;\n }\n }\n\n /* check for null pointers */\n if (ptr_d == NULL)\n {\n FatalError(\"Result of allocation of required temporary storage is NULL. Try pinning the memory to reduce storage requirements.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n#endif\n\n hipGetLastError(); //to reset last error\n return ptr_d;\n}\n\ntemplate \nconst void *DistributedManagerBase::getDevicePointerForData(const void *ptr, size_t size, int *allocated)\n{\n hipError_t rc;\n hipPointerAttribute_t att;\n void *ptr_d;\n cudaCheckError();\n /* WARNING: We may accept the following types of allocation for ptr:\n 1. malloc [host memory]\n 2. cudaMalloc [device memory]\n 3. malloc + cudaHostRegister [AMGX_pin_memory/AMGX_unpin_memory host memory]\n 4. cudaHostAlloc [pinned host memory form the beginning]\n\n The correct way to conver these cases is the following:\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2, 3 or 4.\n }\n else{\n //you are in case 1.\n }\n The following pattern of checks should be implemented\n cudaPointerAttributes att;\n cudaError_t st = cudaPointerGetAttributes(&att, ptr);\n if (st == cudaSuccess) {\n //you are in case 2 or 4.\n }\n else{\n st = cudaHostGetDevicePointer(ptr_on_device, ptr, 0);\n if (st == cudaSuccess){\n //you are in case 3.\n }\n else{\n //you are in case 1.\n }\n }\n The above pattern will be used whenever we need to process input data.\n\n Obs.: parameter size is in bytes and\n parameter allocated indicates whether memory was allocated\n and needs to be release later on. */\n *allocated = 0;\n // get pointer to values on the device\n rc = hipPointerGetAttributes(&att, ptr);\n\n#if CUDART_VERSION >= 11000\n if(rc == hipSuccess && att.type != hipMemoryTypeUnregistered)\n {\n // Received device registered data (device, managed, malloc+cudaHostRegister, or cudaMallocHost)\n ptr_d = (void*)att.devicePointer;\n\n if (ptr_d == NULL)\n {\n FatalError(\"Data passed is registered with device but devicePointer null.\", AMGX_ERR_INTERNAL);\n }\n }\n else\n {\n // Received unregistered or host allocated data (malloc)\n rc = amgx::memory::hipMallocAsync(&ptr_d, size);\n\n if (rc != hipSuccess)\n {\n FatalError(\"Could not allocate required temporary storage. Try pinning the memory to reduce storage requirements.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n rc = hipMemcpy(ptr_d, ptr, size, hipMemcpyDefault);\n cudaCheckError();\n\n if (rc != hipSuccess)\n {\n FatalError(\"Could not copy into the temporary storage. Try pinning the memory to avoid the hipMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n *allocated = 1;\n }\n#else\n if (rc == hipSuccess)\n {\n //you are in case 2 or 4 from the above comment.\n hipGetLastError(); //to reset last error\n return (const void *)att.devicePointer;\n }\n else\n {\n //you are in case 1 or 3 from the above comment\n rc = hipHostGetDevicePointer(&ptr_d, (void *)ptr, 0);\n\n if (rc != hipSuccess)\n {\n //you are in case 1\n rc = amgx::memory::hipMallocAsync(&ptr_d, size);\n\n if (rc != hipSuccess)\n {\n FatalError(\"Could not allocate required temporary storage. Try pinning the memory to reduce storage requirements.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n rc = hipMemcpy(ptr_d, ptr, size, hipMemcpyDefault);\n cudaCheckError();\n\n if (rc != hipSuccess)\n {\n FatalError(\"Could not copy into the temporary storage. Try pinning the memory to avoid the hipMemcpy.\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n *allocated = 1;\n hipGetLastError(); //to reset last error\n return (const void *)ptr_d;\n }\n }\n\n /* check for null pointers */\n if (ptr_d == NULL)\n {\n FatalError(\"Result of allocation of required temporary storage is NULL. Try pinning the memory to reduce storage requirements.\", AMGX_ERR_BAD_PARAMETERS);\n }\n#endif\n\n // shouldn't get there\n hipGetLastError(); //to reset last error\n return NULL;\n}\n\ntemplate \nvoid initializeMatrixCopyAll(int n, int nnz, int block_dimx, int block_dimy, const int *row_ptrs, const int *col_indices, const void *data, const void *diag, Matrix *A)\n{\n typedef typename TConfig::MatPrec mat_value_type;\n A->resize( n, n, nnz, block_dimx, block_dimy );\n //Upload the entire matrix\n hipMemcpy( A->row_offsets.raw(), row_ptrs, (n + 1) * sizeof(int), hipMemcpyDefault );\n cudaCheckError();\n hipMemcpy( A->col_indices.raw(), col_indices, (nnz) * sizeof(int), hipMemcpyDefault );\n cudaCheckError();\n hipMemcpy( A->values.raw(), (mat_value_type *)data, (nnz * block_dimx * block_dimy) * sizeof(mat_value_type), hipMemcpyDefault );\n cudaCheckError();\n\n if (diag)\n {\n hipMemcpy( A->values.raw() + A->diagOffset()*A->get_block_size(), (mat_value_type *)diag, (n * block_dimx * block_dimy) * sizeof(mat_value_type), hipMemcpyDefault );\n cudaCheckError();\n }\n else\n {\n A->computeDiagonal();\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManagerBase::updateMapsReorder()\n{\n int my_id = this->getComms()->get_global_id();\n DistributedComms *comms_tmp = this->getComms();\n DistributedComms **comms_ = &comms_tmp;\n // Copy B2L_maps in their final place\n int num_neighbors = this->neighbors.size();\n B2L_maps.resize(num_neighbors);\n L2H_maps.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n B2L_maps[i] = this->cached_B2L_maps[i];\n L2H_maps[i] = this->cached_L2H_maps[i];\n }\n\n //Create a DistributedArranger object to map further halo rings and to construct halo row matrices and exchange them (if halo_coloring != LAST)\n DistributedArranger *prep = new DistributedArranger;\n prep->create_B2L_from_maps( (*(this->A)), my_id, this->num_halo_rings(), neighbors,\n B2L_maps, L2H_maps, B2L_rings, comms_, &halo_rows, &halo_btl);\n DistributedManagerBaseInit(my_id, 0, this->A->get_num_rows(), *(this->A), comms_, NULL, NULL);\n //Use the exchanged halo row matrices and the boundary/halo index lists to renumber and consolidate the matrix\n this->reorder_matrix();\n prep->initialize_B2L_maps_offsets(*(this->A), this->num_halo_rings());\n delete prep;\n}\n\ntemplate \nvoid DistributedManagerBase::initializeUploadReorderAll(int n, int nnz, int block_dimx, int block_dimy, const int *row_ptrs, const int *col_indices, const void *data, const void *diag, Matrix &in_A)\n{\n this->A = &in_A;\n initializeMatrixCopyAll(n, nnz, block_dimx, block_dimy, row_ptrs, col_indices, data, diag, this->A);\n this->updateMapsReorder();\n}\n\ntemplate \nvoid DistributedManagerBase::destroyComms()\n{\n if ( (this->_comms) != NULL )\n {\n if (this->_comms->decr_ref_count())\n {\n delete (this->_comms);\n this->_comms = NULL;\n }\n }\n\n if ( (this->m_fine_level_comms) != NULL)\n {\n if (this->m_fine_level_comms->decr_ref_count())\n {\n delete (this->m_fine_level_comms);\n this->m_fine_level_comms = NULL;\n }\n }\n}\n\ntemplate \nvoid DistributedManagerBase::initComms(Resources *rsrc)\n{\n this->createComms(rsrc);\n int my_id = this->getComms()->get_global_id();\n int num_parts = this->getComms()->get_num_partitions();\n this->set_global_id(my_id);\n this->set_num_partitions(num_parts);\n}\n\n\ntemplate \nvoid DistributedManagerBase::createComms(Resources *rsrc)\n{\n // create communicator\n#ifdef AMGX_WITH_MPI\n destroyComms();\n if (rsrc == NULL)\n FatalError(\"Resources should not be NULL\", AMGX_ERR_INTERNAL);\n\n MPI_Comm *mpi_comm = rsrc->getMpiComm();\n AMG_Config *cfg = rsrc->getResourcesConfig();\n std::string comm_value, comm_scope;\n cfg->getParameter(\"communicator\", comm_value, \"default\", comm_scope);\n int rank = -1;\n MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n\n if (comm_value == \"MPI_DIRECT\")\n {\n _comms = new CommsMPIDirect(*cfg, comm_scope, mpi_comm);\n std::string comm_log(\"Using CUDA-Aware MPI (GPU Direct) communicator...\\n\");\n amgx_distributed_output(comm_log.c_str(), comm_log.length());\n }\n else if (comm_value == \"MPI\")\n {\n _comms = new CommsMPIHostBufferStream(*cfg, comm_scope, mpi_comm);\n std::string comm_log(\"Using Normal MPI (Hostbuffer) communicator...\\n\");\n amgx_distributed_output(comm_log.c_str(), comm_log.length());\n }\n else \n { \n FatalError(\"Bad communicator value\", AMGX_ERR_BAD_PARAMETERS); \n }\n\n#endif\n}\n\ntemplate \nvoid DistributedManagerBase::malloc_export_maps(VecInt_t ***b2l_maps_e, VecInt_t **b2l_maps_sizes_e, VecInt_t ***l2h_maps_e, VecInt_t **l2h_maps_sizes_e)\n{\n *b2l_maps_e = (VecInt_t **) malloc(sizeof(VecInt_t *)*this->num_neighbors());\n *l2h_maps_e = (VecInt_t **) malloc(sizeof(VecInt_t *)*this->num_neighbors());\n *b2l_maps_sizes_e = (VecInt_t *) malloc(sizeof(VecInt_t) * (this->num_neighbors()));\n *l2h_maps_sizes_e = (VecInt_t *) malloc(sizeof(VecInt_t) * (this->num_neighbors()));\n\n for (int i = 0; i < this->num_neighbors(); i++)\n {\n (*b2l_maps_sizes_e)[i] = B2L_maps[i].size();\n (*l2h_maps_sizes_e)[i] = L2H_maps[i].size();\n (*b2l_maps_e)[i] = (VecInt_t *) malloc(sizeof(VecInt_t) * ( (*b2l_maps_sizes_e)[i]) );\n\n if (L2H_maps[i].size() != 0)\n {\n (*l2h_maps_e)[i] = (VecInt_t *) malloc(sizeof(VecInt_t) * ( (*l2h_maps_sizes_e)[i]) );\n amgx::thrust::copy(L2H_maps[i].begin(), L2H_maps[i].end(), (*l2h_maps_e)[i]);\n }\n\n amgx::thrust::copy(B2L_maps[i].begin(), B2L_maps[i].end(), (*b2l_maps_e)[i]);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManager >::createRenumbering(IVector &renumbering)\n{\n int num_neighbors = this->neighbors.size();\n // still renumber if the number of neighbors = 0, to support non-symmetric matrices\n // if (num_neighbors == 0) return;\n /*\n EXAMPLE\n Example matrix, partition 1 arrives with state:\n A.row_offsets = [0 4 11 15 20]\n A.col_indices = [4 0 1 2\n 4 5 0 1 2 3 7\n 0 1 2 3\n 1 2 3 6 7]\n num_neighbors=2; neighbors = [0 2]\n B2L_rings[[0 2 4][0 2 4]] B2L_maps[[0 1| 2 3][1 3| 0 2]]\n L2H_maps (and halo_lists) [[4 5][6 7]]\n */\n int size = 0;\n\n if (this->L2H_maps.size())\n {\n size = thrust_wrapper::reduce(this->A->col_indices.begin(), this->A->col_indices.end(), int(0), amgx::thrust::maximum()) + 1; //Sufficient to do reduction on lth maps\n cudaCheckError();\n }\n else\n {\n size = this->A->get_num_rows();\n }\n\n int rings = (this->B2L_rings.size() > 0) ? this->B2L_rings[0].size() - 1 : 0;\n //initial size to size+1 so we have the total size after a scan\n renumbering.resize(size + 1);\n int global_size = size;\n //\n // Step 1 - in the main matrix, separate interior and boundary nodes (1/0 in flagArray), renumber interior ones with an exclusive scan\n //\n IVector flagArray(size + 1);\n thrust_wrapper::fill(flagArray.begin(), flagArray.begin() + size + 1, 1);\n cudaCheckError();\n\n //sets 1 for interior nodes, 0 for boundary node\n for (int i = 0; i < num_neighbors; i++ )\n {\n int size = this->B2L_rings[i][1];\n int num_blocks = std::min(4096, (size + 127) / 128);\n\n if (size > 0)\n {\n remove_boundary_kernel <<< num_blocks, 128>>>(flagArray.raw(), this->B2L_maps[i].raw(), size);\n cudaCheckError();\n }\n\n //If there are any L2H maps\n if (this->L2H_maps.size() && this->L2H_maps[i].size())\n {\n int size = this->L2H_maps[i].size();\n int num_blocks = std::min(4096, (size + 127) / 128);\n remove_boundary_kernel <<< num_blocks, 128>>>(flagArray.raw(), this->L2H_maps[i].raw(), size);\n }\n\n cudaCheckError();\n }\n\n //gets the renumbering of interior nodes\n thrust_wrapper::exclusive_scan(flagArray.begin(), flagArray.begin() + size + 1, renumbering.begin());\n cudaCheckError();\n /*\n EXAMPLE\n After removing 1-ring boundary nodes and halo nodes from flagArray: [0 0 1 0 0 0 0 0]\n After exclusive scan, which gives renumbering for interior nodes (only node #2)\n renumbering: [0 0 0 1 1 1 1 1]\n */\n //\n // Step 2 - Renumber nodes that are in the boundary, stepping through each B2L map, and renumbering ones that have not been renumbered yet\n //\n //what is the biggest B2L size\n INDEX_TYPE max_size = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n max_size = max_size > this->B2L_rings[i][1] ? max_size : this->B2L_rings[i][1];\n\n if (this->L2H_maps.size())\n {\n max_size = max_size > this->L2H_maps[i].size() ? max_size : this->L2H_maps[i].size();\n }\n }\n\n //allocate work vectors (should be pretty small) that are used to renumber boundary nodes\n IVector boundary_renum_flags(max_size);\n IVector boundary_renum(max_size);\n //the number of renumbered nodes so far\n int max_element = renumbering[size];\n this->_num_interior_nodes = max_element;\n this->_num_boundary_nodes = this->A->get_num_rows() - max_element;\n renumbering.resize(size);\n\n /*\n EXAMPLE\n size = 8\n max_size = 2, max_element = 1, num_interior_nodes=1, num_boundary_nodes = 4-1 = 3\n */\n\n for (int i = 0; i < num_neighbors; i++)\n {\n //find nodes that are part of the current boundary and they haven't been renumbered yet\n thrust_wrapper::fill(boundary_renum_flags.begin(), boundary_renum_flags.begin() + max_size, 0);\n int size = this->B2L_rings[i][1];\n int num_blocks = std::min(4096, (size + 191) / 192);\n\n if (size > 0)\n get_unassigned_kernel <<< num_blocks, 192>>>(flagArray.raw(),\n this->B2L_maps[i].raw(),\n boundary_renum_flags.raw(), size, global_size /*,rank*/);\n cudaCheckError();\n\n //calculate the local renumbering (within this boundary region) of these nodes\n thrust_wrapper::exclusive_scan(boundary_renum_flags.begin(), boundary_renum_flags.begin() + max_size, boundary_renum.begin());\n\n //apply renumbering to the big numbering table\n if (size > 0)\n set_unassigned_kernel <<< num_blocks, 192>>>(boundary_renum_flags.raw(),\n boundary_renum.raw(),\n this->B2L_maps[i].raw(),\n renumbering.raw(),\n size, max_element, global_size /*,rank*/);\n cudaCheckError();\n\n //update the number of renumbered nodes\n max_element += boundary_renum[max_size - 1] + boundary_renum_flags[max_size - 1];\n /*\n EXAMPLE\n for neighbor 0 (ID 0)\n boundary_renum_flags = [0 0], size = 2, flagArray [0 0 1 0 0 0 0 0]\n get_unassigned_kernel's output: boundary_renum_flags = [1 1] flagArray [1 1 1 0 0 0 0 0]\n after exclusive scan: boundary_renum [0 1]\n set_unassigned_kernel updates these arrays and renumbers B2L map:\n renumbering = [1 2 0 1 1 1 1 1] B2L_maps[0] = [1 2| 2 3] (note that after element 3 in renumbering and after element 2 we have invalid/not yet updated values)\n max_element = 3\n\n for neighbor 1 (ID 2)\n get_unassigned_kernels's output: boundary_renum_flags [0 1] flagArray [1 1 1 1 0 0 0 0]\n after exclusive scan boundary_renum [0 0]\n set_unassigned_kernel renumbering [1 2 0 3 1 1 1 1] B2L_maps[1] = [2 3| 0 2]\n max_element = 4\n */\n }\n\n cudaCheckError();\n\n //Get renumbering for halo indices\n if (this->L2H_maps.size())\n {\n //TODO: simplify this, we don't need to check whether it has already been renumbered, there is no overlap between halos\n for (int i = 0; i < num_neighbors; i++)\n {\n //find nodes that are part of the current boundary and they haven't been renumbered yet\n thrust_wrapper::fill(boundary_renum_flags.begin(), boundary_renum_flags.begin() + max_size, 0);\n int size = this->L2H_maps[i].size();\n int num_blocks = std::min(4096, (size + 191) / 192);\n\n if (size > 0)\n get_unassigned_kernel <<< num_blocks, 192>>>(flagArray.raw(),\n this->L2H_maps[i].raw(),\n boundary_renum_flags.raw(), size, global_size /*,rank*/);\n cudaCheckError();\n\n //calculate the local renumbering (within this boundary region) of these nodes\n thrust_wrapper::exclusive_scan(boundary_renum_flags.begin(), boundary_renum_flags.begin() + max_size, boundary_renum.begin());\n\n //apply renumbering to the big numbering table\n if (size > 0)\n set_unassigned_kernel <<< num_blocks, 192>>>(boundary_renum_flags.raw(),\n boundary_renum.raw(),\n this->L2H_maps[i].raw(),\n renumbering.raw(),\n size, max_element, global_size /*,rank*/);\n cudaCheckError();\n\n //update the number of renumbered nodes\n max_element += boundary_renum[max_size - 1] + boundary_renum_flags[max_size - 1];\n /*\n EXAMPLE\n for neighbor 0 (ID 0)\n boundary_renum_flags = [0 0], size = 2, flagArray [1 1 1 1 0 0 0 0]\n get_unassigned_kernel's output: boundary_renum_flags = [1 1] flagArray [1 1 1 1 1 1 0 0]\n after exclusive scan: boundary_renum [0 1]\n set_unassigned_kernel updates these arrays and renumbers B2L map:\n renumbering = [1 2 0 3 4 5 1 1] L2H_maps[0] = [4 5]\n max_element = 6\n\n for neighbor 1 (ID 2)\n get_unassigned_kernels's output: boundary_renum_flags [1 1] flagArray [1 1 1 1 1 1 1 1]\n after exclusive scan boundary_renum [0 1]\n set_unassigned_kernel renumbering = [1 2 0 3 4 5 6 7] L2H_maps[1] = [6 7]\n max_element = 8\n */\n }\n\n cudaCheckError();\n }\n\n //apply renumbering to further halo rings too\n if (rings > 1)\n {\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = this->B2L_rings[i][this->B2L_rings[i].size() - 1] - this->B2L_rings[i][1];\n int num_blocks = std::min(4096, (size + 127) / 128);\n renumber_b2l_maps <<< num_blocks, 128>>>(this->B2L_maps[i].raw() + this->B2L_rings[i][1], renumbering.raw(), size, global_size /*, rank*/);\n }\n\n cudaCheckError();\n }\n\n /*\n EXAMPLE\n renumbers further boundary rings as listed in B2L_maps, since they have not been replaced yet with their renumbered values\n B2L_maps [[1 2| 0 3][2 3| 1 0]]\n */\n}\n\ntemplate \nvoid DistributedManager >::reorder_matrix_owned()\n{\n int num_neighbors = this->neighbors.size();\n int size = this->A->get_num_rows();\n int num_blocks = std::min(4096, (size + 511) / 512);\n int rings = (this->B2L_rings.size() > 0) ? this->B2L_rings[0].size() - 1 : 0;\n this->set_num_halo_rings(rings);\n int diag = this->A->hasProps(DIAG);\n\n if (diag)\n {\n FatalError(\"External diag not supported in classical path\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n//\n// Step 1 & 2 - create renumbering\n//\n this->createRenumbering(this->renumbering);\n //now we have the full renumbering table in renum, calculate the inverse\n this->inverse_renumbering.resize(this->renumbering.size());\n\n if (this->renumbering.size() > 1)\n {\n calc_inverse_renumbering <<< std::min(4096, ((int)this->renumbering.size() + 511) / 512), 512 >>> (this->renumbering.raw(), this->inverse_renumbering.raw(), this->renumbering.size());\n cudaCheckError();\n }\n\n//\n// Step 4 - calculate number/offset of nodes in the halos from the neighbors, ring by ring\n//\n this->halo_offsets.resize(num_neighbors + 1);\n this->halo_offsets[0] = size;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n this->halo_offsets[i + 1] = this->halo_offsets[i] + this->L2H_maps[i].size();\n }\n\n this->set_num_halo_rows(this->halo_offsets[this->halo_offsets.size() - 1] - size);\n int nh = this->num_halo_rows();\n int total_rows = size + nh;\n cudaCheckError();\n//\n// Step 6 - renumber halo matrices and calculate row length (to eventually append to the big matrix)\n//\n int insert = 0;\n //recalculate row_offsets\n IVector new_row_offsets(size + 1);\n\n if (num_blocks > 0)\n {\n calc_rowlen_reorder <<< num_blocks, 512>>>(this->A->row_offsets.raw(), new_row_offsets.raw(), this->renumbering.raw(), size, insert);\n cudaCheckError();\n }\n\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(this->renumbering.begin(), this->A->col_indices.begin()),\n amgx::thrust::make_permutation_iterator(this->renumbering.begin(), this->A->col_indices.end()),\n this->A->col_indices.begin());\n cudaCheckError();\n //row_offsets array created by exclusive scan of row sizes\n thrust_wrapper::exclusive_scan(new_row_offsets.begin(), new_row_offsets.begin() + size + 1, new_row_offsets.begin());\n cudaCheckError();\n//\n// Step 7 - consolidate column indices and values\n//\n int new_nnz = new_row_offsets[new_row_offsets.size() - 1];\n typedef typename MatPrecisionMap::Type ValueTypeA;\n VVector new_values((new_nnz + 1 )* this->A->get_block_size(), types::util::get_zero());\n IVector new_col_indices(new_nnz, 0);\n\n //reorder based on row permutation\n if (num_blocks > 0)\n {\n reorder_whole_matrix <32> <<< num_blocks, 512>>>(this->A->row_offsets.raw(), this->A->col_indices.raw(), this->A->values.raw(), new_row_offsets.raw(), new_col_indices.raw(), new_values.raw(), this->renumbering.raw(), this->A->get_block_size(), size, insert);\n cudaCheckError();\n }\n\n //create and append halo rows size\n //create an identity matrix in CSR format\n int nnz = this->A->get_num_nz();\n IVector identity_csr_rows(nh + 1);\n IVector identity_csr_cols(nh);\n VVector identity_csr_vals(nh, types::util::get_one()); //needs to be changed to MVector, but this definition is messed up in the header file (should fix later)\n thrust_wrapper::sequence(identity_csr_rows.begin(), identity_csr_rows.end());\n thrust_wrapper::sequence(identity_csr_cols.begin(), identity_csr_cols.end());\n /*for example, 2x2 identity_csr matrix is created:\n identity_csr_rows = { 0, 1, 2 }\n identity_csr_cols = { 0, 1 }\n identity_csr_vals = { 1.0, 1.0 } */\n //shift identity tmatrix by size = this->A->get_num_rows();\n thrust_wrapper::transform(identity_csr_rows.begin(), identity_csr_rows.end(), amgx::thrust::constant_iterator(nnz), identity_csr_rows.begin(), amgx::thrust::plus());\n thrust_wrapper::transform(identity_csr_cols.begin(), identity_csr_cols.end(), amgx::thrust::constant_iterator(size), identity_csr_cols.begin(), amgx::thrust::plus());\n /*for example, 2x2 identity_csr matrix is created:\n identity_csr_rows = { 0, 1, 2 }\n identity_csr_cols = {size, size+1 }\n identity_csr_vals = { 1.0, 1.0 } */\n /* WARNING: you must be very careful with the view you are setting (cuurently the view coming here by default is ALL = FULL). If\n - classical path is selected then the createOneRingHaloRows -> create_one_ring_halo_rows -> append_halo_rows\n routine will be called. It will overwrite the halo rows setup here (and will use view OWNED, which will ignore the\n halo rows setup here, to determine how the new halo rows should be placed).\n - aggregation path is selected then the extra rows setup here will be used in the R*A*P product, where (in order to match\n dimensions of R and P) it is assumed that (the local partition) matrix A is square, therefore it must be padded by identity\n rows at the bottom to compensate for the \"extra\" columns that are outside of the main square part. The old routines for the\n aggregation path do this padding at the end of the reorder_matrix routine below. */\n //ViewType v = this->A->currentView();\n //this->A->setView(ALL);\n //Approach 1: use existing routine to append the identity matrix to the existing one\n // (seems like too much overhead, also need identity matrix per neighbor)\n //DistributedArranger *prep = new DistributedArranger;\n //prep->append_halo_rows(this->A, identity_csr_rows, identity_csr_cols, identity_csr_vals);\n //delete prep;\n //Approach 2: custom for this routine\n new_row_offsets.resize(total_rows + 1);\n new_col_indices.resize(nnz + nh);\n int bsize = this->A->get_block_size();\n new_values.resize((nnz + nh)*bsize + 1); //extra 1 element stores zero at the end (to follow the original design)\n //new_values[nnz]=-1; //marker to track the last element\n amgx::thrust::copy(identity_csr_rows.begin(), identity_csr_rows.end(), new_row_offsets.begin() + size );\n amgx::thrust::copy(identity_csr_cols.begin(), identity_csr_cols.end(), new_col_indices.begin() + nnz);\n amgx::thrust::copy(new_values.begin() + nnz*bsize, new_values.begin() + (nnz + 1)*bsize, new_values.begin() + (nnz + nh)*bsize);\n amgx::thrust::copy(identity_csr_vals.begin(), identity_csr_vals.end(), new_values.begin() + (nnz*bsize));\n /* WARNING: see above. */\n this->A->set_num_cols(total_rows);\n this->A->set_num_rows(total_rows);\n this->A->col_indices.swap(new_col_indices);\n new_row_offsets.resize(total_rows + 1);\n this->A->row_offsets.swap(new_row_offsets);\n new_row_offsets.swap(this->old_row_offsets);\n this->A->values.swap(new_values);\n this->A->m_seq_offsets.resize(total_rows + 1);\n thrust_wrapper::sequence(this->A->m_seq_offsets.begin(), this->A->m_seq_offsets.end());\n cudaCheckError();\n //TODO: only do this if AMG_Config matrix_halo_exchange!=2\n this->A->delProps(COO);\n if (!insert)\n {\n this->A->computeDiagonal();\n }\n\n this->set_initialized(this->A->row_offsets);\n this->A->setView(OWNED);\n}\n\n\n\ntemplate \nvoid DistributedManager >::reorder_matrix()\n{\n int num_neighbors = this->neighbors.size();\n\n if (num_neighbors == 0) { return; }\n\n int size = this->A->get_num_rows();\n int num_blocks = std::min(4096, (size + 511) / 512);\n int rings = this->B2L_rings[0].size() - 1;\n this->set_num_halo_rings(rings);\n int diag = this->A->hasProps(DIAG);\n std::vector > &halo_rows = *this->halo_rows;\n std::vector > &halo_btl = *this->halo_btl;\n /*\n EXAMPLE\n The example matrix, on partition 1 arrives at this point with the following state:\n num_rings=2\n A.num_rows = 4; A.num_nz = 20\n A.row_offsets = [0 4 11 15 20]\n A.col_indices = [4 0 1 2\n 4 5 0 1 2 3 7\n 0 1 2 3\n 1 2 3 6 7]\n num_neighbors=2; neighbors = [0 2]\n B2L_rings[[0 2 4][0 2 4]] B2L_maps[[0 1| 2 3][1 3| 0 2]]\n L2H_maps (and halo_lists) [[4 5][6 7]]\n\n With the exchange halo rows:\n halo_btl[0] (received from neighbor ID 0)\n global_id = 0; base_index=0; index_range=6; B2L_rings[0] = [0 2 4]; B2L_maps[0] = [2 3| 0 1] L2H_maps = [4 5]\n halo_rows[0].row_offsets = [0 5 13 17 21]\n halo_rows[0].col_indices = [1 2 3 4 5\n 0 1 2 3 4 5 6 7\n 0 1 3 6\n 0 1 2 3]\n\n halo_btl[1] (received from neighbor ID 2)\n global_id = 2; base_index=0; index_range=8; B2L_rings[0] = [0 2 4]; B2L_maps[0] = [1 2| 0 3] L2H_maps = [6 7]\n halo_rows[1].row_offsets = [0 4 11 16 20]\n halo_rows[1].col_indices = [7 1 2 3\n 5 6 7 0 1 2 3\n 4 5 0 2 3\n 0 1 2 3]\n */\n//\n// Step 1 & 2 - create renumbering\n//\n this->createRenumbering(this->renumbering);\n cudaCheckError();\n /*\n EXAMPLE\n this->renumbering = [1 2 0 3 4 5 6 7]\n B2L_maps = [[1 2| 0 3][2 3| 1 0]]\n L2H_maps = [[4 5][6 7]]\n */\n//\n// Step 3 - given a full renumbering of owned nodes, calculate inverse renumbering\n//\n //now we have the full renumbering table in renum, calculate the inverse\n this->inverse_renumbering.resize(this->renumbering.size());\n calc_inverse_renumbering <<< std::min(4096, ((int)this->renumbering.size() + 511) / 512), 512 >>> (this->renumbering.raw(), this->inverse_renumbering.raw(), this->renumbering.size());\n cudaCheckError();\n /*\n EXAMPLE\n this->inverse_renumbering = [2 0 1 3 4 5 6 7]\n */\n//\n// Step 4 - calculate number/offset of nodes in the halos from the neighbors, ring by ring\n//\n this->halo_offsets.resize(rings * num_neighbors + 1, 0);\n\n for (int ring = 0; ring < rings; ring++)\n {\n for (int i = 0; i < num_neighbors; i++)\n {\n this->halo_offsets[ring * num_neighbors + i] = halo_btl[i].B2L_rings[0][ring + 1] - halo_btl[i].B2L_rings[0][ring];\n }\n }\n\n thrust_wrapper::exclusive_scan(this->halo_offsets.begin(), this->halo_offsets.end(), this->halo_offsets.begin(), size);\n cudaCheckError();\n this->set_num_halo_rows(this->halo_offsets[this->halo_offsets.size() - 1] - size);\n int total_rows = size + this->num_halo_rows();\n\n if (total_rows < this->renumbering.size()) { FatalError(\"total rows < renumbering.size(), send/recv maps should cover all matrix halo columns\", AMGX_ERR_NOT_IMPLEMENTED); }\n\n if (total_rows > this->renumbering.size())\n {\n this->A->getResources()->warning(\"# owned nodes + # halo nodes > matrix columns: send/recv maps have some unreferences halo indices, they are not directly connected to our partition and therefore we won't compute them, please use 2-ring comms maps if you want to specify 2nd ring neighbors\");\n }\n\n cudaCheckError();\n /*\n EXAMPLE\n halo_offsets [2 2 2 2]\n after exclusive scan: 4 + [0 2 4 6 8] = [4 6 8 10 12]\n num_halo_rows = 8, total_rows = 12\n */\n//\n// Step 5 - create big mapping table of all halo indices we received (this may use a little too much memory)\n//\n //count number of fine rows of neighbors\n amgx::thrust::host_vector neighbor_rows(num_neighbors + 1);\n int max_num_rows = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n neighbor_rows[i] = halo_btl[i].index_range();\n max_num_rows = max_num_rows > halo_rows[i].get_num_rows() ? max_num_rows : halo_rows[i].get_num_rows();\n }\n\n thrust_wrapper::exclusive_scan(neighbor_rows.begin(), neighbor_rows.end(), neighbor_rows.begin());\n cudaCheckError();\n int total_rows_of_neighbors = neighbor_rows[num_neighbors];\n /*\n EXAMPLE\n neigbor_rows = [0 6 14]\n total_rows_of_neighbors = 14\n */\n IVector halo_mapping(total_rows_of_neighbors);\n thrust_wrapper::fill(halo_mapping.begin(), halo_mapping.end(), -1);\n cudaCheckError();\n\n //ring by ring, neighbor by neighbor assign sequentially increasing numbers for halo nodes\n for (int ring = 0; ring < rings; ring++)\n {\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = halo_btl[i].B2L_rings[0][ring + 1] - halo_btl[i].B2L_rings[0][ring];\n int num_blocks = std::min(4096, (size + 127) / 128);\n //This renumbering has to result in the same renumbering that comes out of L2H renumbering\n create_halo_mapping <<< num_blocks, 128>>>(halo_mapping.raw() + neighbor_rows[i],\n halo_btl[i].B2L_maps[0].raw() + halo_btl[i].B2L_rings[0][ring],\n halo_btl[i].base_index(), this->halo_offsets[ring * num_neighbors + i], size);\n cudaCheckError();\n /*\n EXAMPLE\n ring 0 neighbor 0 - halo_btl[0].B2L_maps[0] = [2 3| 0 1] halo_btl[0].L2H_maps = [4 5]\n halo_mapping = [-1 -1 4 5 -1 -1 |-1 -1 -1 -1 -1 -1 -1 -1]\n ring 0 neighbor 1 - halo_btl[1].B2L_maps[0] = [1 2| 0 3] halo_btl[1].L2H_maps = [6 7]\n halo_mapping = [-1 -1 4 5 -1 -1 |-1 6 7 -1 -1 -1 -1 -1]\n ring 1 neighbor 0 - halo_btl[0].B2L_maps[0] = [2 3| 0 1] halo_btl[0].L2H_maps = [4 5]\n halo_mapping = [8 9 4 5 -1 -1 |-1 6 7 -1 -1 -1 -1 -1]\n ring 1 neighbor 1 - halo_btl[1].B2L_maps[0] = [1 2| 0 3] halo_btl[1].L2H_maps = [6 7]\n halo_mapping = [8 9 4 5 -1 -1 |10 6 7 11 -1 -1 -1 -1]\n\n */\n }\n }\n\n cudaCheckError();\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = halo_btl[i].L2H_maps[0].size();\n int num_blocks = std::min(4096, (size + 127) / 128);\n //Map the column indices of the halo rows that point back to boundary nodes\n apply_h2l2b_mapping <<< num_blocks, 128>>>(halo_mapping.raw() + neighbor_rows[i],\n halo_btl[i].L2H_maps[0].raw(),\n halo_btl[i].base_index(), this->B2L_maps[i].raw(), size);\n cudaCheckError();\n /*\n EXAMPLE\n neighbor 0 - mapping back to our own (boundary) indices\n halo_mapping = [8 9 4 5 1 2 |10 6 7 11 -1 -1 -1 -1]\n neighbor 1 - mapping back to our own (boundary) indices\n halo_mapping = [8 9 4 5 1 2 |10 6 7 11 -1 -1 2 3]\n */\n }\n\n cudaCheckError();\n /*\n EXAMPLE\n neighbor_rows = [0 6 14]\n halo_mapping = [8 9 4 5 1 2 |10 6 7 11 -1 -1 2 3]\n The first part (0-6) of halo_mapping gives a local index for all the indices that we want to know about in halo_btl[0]\n The second part (7-14) gives local indices for halo_btl[1], that is both halo ring there, and the column indices representing vertices in this partition's boundary.\n Note that it does not give indices (-1) for vertices 5 and 6 in neighbor 1 (ID 2), which are column indices connecting it to neighbor 0, hence the two halo regions are not connected\n */\n//\n// Step 6 - renumber halo matrices and calculate row length (to eventually append to the big matrix)\n//\n int insert = 0;\n\n if (this->A->hasProps(DIAG) && insertDiagonals) { insert = 1; }\n\n diag = diag && !insertDiagonals;\n //recalculate row_offsets\n IVector new_row_offsets(size + this->num_halo_rows() + 1);\n calc_rowlen_reorder <<< num_blocks, 512>>>(this->A->row_offsets.raw(), new_row_offsets.raw(), this->renumbering.raw(), size, insert);\n cudaCheckError();\n IVector neighbor_rows_d(num_neighbors + 1);\n amgx::thrust::copy(neighbor_rows.begin(), neighbor_rows.end(), neighbor_rows_d.begin());\n cudaCheckError();\n /*\n EXAMPLE\n get row length according to renumbering\n new_row_offsets = [4 4 7 5 0 0 0 0 0 0 0 0 0]\n */\n //map column indices of my own matrix\n /*map_col_indices<4><<>>(this->A->row_offsets.raw(),\n this->A->col_indices.raw(),\n this->renumbering.raw(),\n this->halo_ranges.raw(),\n halo_mapping.raw(),\n neighbor_rows_d.raw(),\n this->base_index(), num_neighbors, size);*/\n cudaCheckError();\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(this->renumbering.begin(), this->A->col_indices.begin()),\n amgx::thrust::make_permutation_iterator(this->renumbering.begin(), this->A->col_indices.end()),\n this->A->col_indices.begin());\n cudaCheckError();\n /*\n EXAMPLE\n use this->renumbering = [1 2 0 3 4 5 6 7]\n to map old column indices to new column indices (i.e. according to interior - boundary - halo separation), but do not reshuffle them into their place yet\n A.col_indices = [4 0 1 2\n 4 5 0 1 2 3 7\n 0 1 2 3\n 1 2 3 6 7]\n becomes\n A.col_indices = [4 1 2 0\n 4 5 1 2 0 3 7\n 1 2 0 3\n 2 0 3 6 7]\n\n */\n cudaCheckError();\n IVector temp_row_len(max_num_rows);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n //map column indices of halo matrices and count of nonzeros we will keep\n int size = halo_rows[i].get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n map_col_indices_and_count_rowlen<4> <<< num_blocks, 128, 128 * sizeof(INDEX_TYPE)>>>(\n halo_rows[i].row_offsets.raw(),\n halo_rows[i].col_indices.raw(),\n temp_row_len.raw(),\n halo_mapping.raw() + neighbor_rows[i],\n size, insert);\n cudaCheckError();\n\n //number of nonzeros per row copied into big row sizes array\n for (int ring = 0; ring < rings; ring++)\n {\n amgx::thrust::copy(temp_row_len.begin() + halo_btl[i].B2L_rings[0][ring], temp_row_len.begin() + halo_btl[i].B2L_rings[0][ring + 1], new_row_offsets.begin() + this->halo_offsets[ring * num_neighbors + i]);\n }\n\n cudaCheckError();\n /*\n EXAMPLE\n halo_mapping = [8 9 4 5 1 2 |10 6 7 11 -1 -1 2 3]\n look at halo row matrices, and halo_mapping, count column indices that do not map to -1 and map them to their new, local index\n halo_rows[0].col_indices = [1 2 3 4 5\n 0 1 2 3 4 5 6 7\n 0 1 3 6\n 0 1 2 3]\n becomes\n halo_rows[0].col_indices = [9 4 5 1 2\n 8 9 4 5 1 2 -1 -1\n 8 9 5 -1\n 8 9 4 5]\n with temp_row_len = [5 6 3 4] copied into new_row_offsets: [4 4 7 5| 5 6| 0 0| 3 4| 0 0 0]\n\n halo_rows[1].col_indices = [7 1 2 3\n 5 6 7 0 1 2 3\n 4 5 0 2 3\n 0 1 2 3]\n becomes\n halo_rows[1].col_indices = [3 6 7 11\n -1 2 3 10 6 7 11\n -1 -1 10 7 11\n 10 6 7 11]\n with temp_row_len = [4 6 3 4] copied into new_row_offsets: [4 4 7 5| 5 6| 4 6| 3 4| 3 4 0]\n */\n }\n\n cudaCheckError();\n //row_offsets array created by exclusive scan of row sizes\n thrust_wrapper::exclusive_scan(new_row_offsets.begin(), new_row_offsets.begin() + size + this->num_halo_rows() + 1, new_row_offsets.begin());\n cudaCheckError();\n /*\n EXAMPLE\n Exclusive scan to get new_row_offsets array:\n new_row_offsets = [0 4 8 15 20| 25 31| 35 41| 44 48| 51 55]\n */\n//\n// Step 7 - consolidate column indices and values\n//\n int new_nnz = new_row_offsets[new_row_offsets.size() - 1];\n typedef typename MatPrecisionMap::Type ValueTypeA;\n VVector new_values((new_nnz + 1 + diag * (total_rows - 1))* this->A->get_block_size(), types::util::get_zero());\n IVector new_col_indices(new_nnz, 0);\n //reorder based on row permutation\n reorder_whole_matrix <32> <<< num_blocks, 512>>>(this->A->row_offsets.raw(), this->A->col_indices.raw(), this->A->values.raw(), new_row_offsets.raw(), new_col_indices.raw(), new_values.raw(), this->renumbering.raw(), this->A->get_block_size(), size, insert);\n cudaCheckError();\n\n if (diag)\n {\n //reorder based on row permutation\n reorder_vector_values <<< num_blocks, 512>>>(new_values.raw() + new_row_offsets[total_rows]*this->A->get_block_size(),\n this->A->values.raw() + this->A->row_offsets[size]*this->A->get_block_size(),\n this->renumbering.raw(),\n this->A->get_block_size(), size);\n cudaCheckError();\n }\n\n int cumulative_num_rows = size;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n for (int ring = 0; ring < rings; ring++)\n {\n int num_rows = halo_btl[i].B2L_rings[0][ring + 1] - halo_btl[i].B2L_rings[0][ring];\n int num_blocks = std::min(4096, (num_rows + 127) / 128);\n //copy in nonzeros that we are keeping\n //TODO: access pattern - should be implemented with warp-wide scans to decide which nonzeros we are keeping and where the rest is going\n reorder_whole_halo_matrix <<< num_blocks, 128>>>(halo_rows[i].row_offsets.raw() + halo_btl[i].B2L_rings[0][ring],\n halo_rows[i].col_indices.raw(), halo_rows[i].values.raw(),\n new_row_offsets.raw() + this->halo_offsets[ring * num_neighbors + i],\n new_col_indices.raw(), new_values.raw(), NULL, this->A->get_block_size(), num_rows,\n insert, this->halo_offsets[ring * num_neighbors + i], halo_btl[i].B2L_rings[0][ring], halo_btl[i].B2L_rings[0][rings]);\n cudaCheckError();\n\n if (diag)\n {\n amgx::thrust::copy(halo_rows[i].values.begin() + (halo_rows[i].row_offsets[halo_rows[i].get_num_rows()] + halo_btl[i].B2L_rings[0][ring])*this->A->get_block_size(),\n halo_rows[i].values.begin() + (halo_rows[i].row_offsets[halo_rows[i].get_num_rows()] + halo_btl[i].B2L_rings[0][ring + 1])*this->A->get_block_size(),\n new_values.begin() + (new_row_offsets[total_rows] + cumulative_num_rows)*this->A->get_block_size());\n cumulative_num_rows += num_rows;\n }\n }\n }\n\n cudaCheckError();\n /*\n EXAMPLE\n copy everything in place, dropping -1 column indices in the halo and reordering the owned rows\n new_row_offsets = [0 4 8 15 20| 25 31| 35 41| 44 48| 51 55]\n new_col_indices = [1 2 0 3\n 4 1 2 0\n 4 5 1 2 0 3 7\n 2 0 3 6 7 -end of owned\n 9 4 5 1 2\n 8 9 4 5 1 2 - end of neighbor 0 ring 0\n 3 6 7 11\n 2 3 10 6 7 11 - end of neighbor 1 ring 0\n 8 9 5\n 8 9 4 5 - end of neighbor 0 ring 1\n 10 7 11\n 10 6 7 11] - end of neighbor 1 ring 1\n\n */\n this->A->set_num_cols(total_rows);\n this->A->set_num_rows(size);\n this->A->col_indices.swap(new_col_indices);\n new_row_offsets.resize(total_rows + 1);\n this->A->row_offsets.swap(new_row_offsets);\n new_row_offsets.swap(this->old_row_offsets);\n this->A->values.swap(new_values);\n this->A->m_seq_offsets.resize(total_rows + 1);\n thrust_wrapper::sequence(this->A->m_seq_offsets.begin(), this->A->m_seq_offsets.end());\n\n if (insert)\n {\n this->A->delProps(DIAG);\n this->A->diag.resize(total_rows);\n amgx::thrust::copy(this->A->row_offsets.begin(), this->A->row_offsets.end() - 1, this->A->diag.begin());\n }\n\n cudaCheckError();\n delete this->halo_rows;\n delete this->halo_btl;\n //set halo_rows and halo_btl to NULL to avoid a potential double free situation in the future\n this->halo_rows = NULL;\n this->halo_btl = NULL;\n this->A->delProps(COO);\n this->A->set_initialized(1);\n\n //TODO: only do this if AMG_Config matrix_halo_exchange!=2\n if (!insert) { this->A->computeDiagonal(); }\n\n this->A->setView(OWNED);\n}\n\n//function object (functor) for thrust calls (it is a unary operator to add a constant)\ntemplate\nclass add_constant_op\n{\n const T c;\n public:\n add_constant_op(T _c) : c(_c) {}\n __host__ __device__ T operator()(const T &x) const\n {\n return x + c;\n }\n};\n\ntemplate \nvoid DistributedManager >::obtain_shift_l2g_reordering(index_type n, I64Vector_d &l2g, IVector_d &p, IVector_d &q)\n{\n /* WARNING: Exchange halo of the inverse_reordering, which is implicitly based on the local_to_global_map (l2g).\n Notice that it is implicit in the exchange_halo routine, since you are getting exactly the vector\n halo elements, which are exactly the elements you need. They however must be shifted by the partition\n starting points (starting global row indices, which are containe din array part_offsets).\n This allows us to avoid constructing the global vector for inverse permutation,\n as is done in reference MATLAB code. */\n //Recall that part_offsets provide the starting point (global row index) of every partition, in other words,\n //they contain the prefix sum of number of rows assigned to each partition. Also, notice that part_offsets and\n //part_offsets_h have the same values on device and host, respectively. See below few lines for details:\n index_type tag = 1 * 133 + 3 * 7 + 0; //some random number for the tag\n index_type l = p.size();\n q.resize(l);\n amgx::thrust::copy (p.begin(), p.end(), q.begin());\n thrust_wrapper::transform(q.begin(), q.end(), q.begin(), add_constant_op(this->part_offsets[this->global_id()]));\n this->exchange_halo(q, tag);\n thrust_wrapper::sequence (q.begin(), q.begin() + n);\n thrust_wrapper::transform(q.begin(), q.begin() + n, q.begin(), add_constant_op(this->part_offsets[this->global_id()]));\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManager >::unpack_partition(index_type *Bp, index_type *Bc, mat_value_type *Bv)\n{\n index_type l, n, nnz, offset, block_dimx, block_dimy;\n index_type *ir;\n index_type *Ap;\n index_type *Ac;\n mat_value_type *Av;\n IVector q;\n //some initializations\n this->A->getOffsetAndSizeForView(OWNED, &offset, &n);\n this->A->getNnzForView(OWNED, &nnz);\n\n block_dimx = this->A->get_block_dimx();\n block_dimy = this->A->get_block_dimy();\n\n l = this->inverse_renumbering.size();\n ir = this->inverse_renumbering.raw();\n Ap = this->A->row_offsets.raw();\n Ac = this->A->col_indices.raw();\n Av = this->A->values.raw();\n //(i) reorder the matrix back (into mixed interior-boundary nodes)\n //applies to rows and columns (out-of-place)\n reorder_partition\n\t(n, nnz, Ap, Ac, Av, Bp, Bc, Bv, l, ir, block_dimx, block_dimy);\n cudaCheckError();\n //obtain reordering q that combines the shift of the diagonal block with the off-diagonal block indices conversion from local to global\n this->obtain_shift_l2g_reordering(n, this->local_to_global_map, this->inverse_renumbering, q);\n cudaCheckError();\n //(ii) reorder the matrix back (shift the diagonal block and convert off-diagonal block column indices from local to global)\n //applies columns only (in-place)\n reorder_partition\n (n, nnz, Bp, Bc, Bv, Bp, Bc, Bv, q.size(), q.raw(), block_dimx, block_dimy);\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManagerBase::createNeighToDestPartMap(IVector_h &neigh_to_part, IVector_h &neighbors, IVector_h &destination_part, int num_neighbors)\n{\n neigh_to_part.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n neigh_to_part[i] = destination_part[neighbors[i]];\n }\n}\n\ntemplate \nvoid DistributedManagerBase::createConsolidatedNeighToPartMap(IVector_h &cons_neigh_to_part, IVector_h &neigh_to_part, int my_destination_part, IVector_h &destination_part, int &num_cons_neighbors)\n{\n // input: non-initialized cons_neigh_to_part\n // fine_neigh_to_part\n // my_destination_part\n // output: cons_neigh_to_part\n // num_cons_neighbors\n cons_neigh_to_part = neigh_to_part;\n amgx::thrust::sort(cons_neigh_to_part.begin(), cons_neigh_to_part.end());\n cudaCheckError();\n cons_neigh_to_part.erase(amgx::thrust::unique(cons_neigh_to_part.begin(), cons_neigh_to_part.end()), cons_neigh_to_part.end());\n // Remove if fine_neigh maps to same coarse partition\n cons_neigh_to_part.erase(amgx::thrust::remove_if(cons_neigh_to_part.begin(), cons_neigh_to_part.end(), is_my_part(my_destination_part)), cons_neigh_to_part.end());\n num_cons_neighbors = cons_neigh_to_part.size();\n cudaCheckError();\n}\n\n\ntemplate \nvoid DistributedManagerBase::createNeighToConsNeigh(IVector_h &neigh_to_cons_neigh, IVector_h &cons_neigh_to_part, IVector_h &neigh_to_part, int my_destination_part, int &num_neighbors)\n{\n neigh_to_cons_neigh.resize(num_neighbors);\n amgx::thrust::lower_bound(cons_neigh_to_part.begin(), cons_neigh_to_part.end(), neigh_to_part.begin(), neigh_to_part.end(), neigh_to_cons_neigh.begin());\n cudaCheckError();\n\n // Flagging fine neighbors that go to same partition (haven't been found in previous step)\n for (int i = 0; i < num_neighbors; i++)\n {\n if ( neigh_to_part[i] == my_destination_part)\n {\n neigh_to_cons_neigh[i] = -1;\n }\n }\n}\n\ntemplate \ntemplate \nvoid DistributedManagerBase::consB2Lmaps(std::vector &dest_coarse_B2L_maps, std::vector &coarse_B2L_maps, IVector_h &fine_neigh_to_coarse_neigh, int num_coarse_neighbors, int num_fine_neighbors)\n{\n //Merge B2L fine maps per coarse destination\n dest_coarse_B2L_maps.resize(num_coarse_neighbors);\n std::vector dest_coarse_B2L_maps_scratch_sizes(num_coarse_neighbors, 0);\n int my_id = this->global_id();\n\n // Loop over the fine neighbors, to compute size of each dest_coarse_B2L_maps\n for (int i = 0; i < num_fine_neighbors; i++)\n {\n int k = fine_neigh_to_coarse_neigh[i];\n\n if (k != -1)\n {\n dest_coarse_B2L_maps_scratch_sizes[k] += coarse_B2L_maps[i].size();\n }\n }\n\n // Now fill dest_coarse_B2L_maps\n for (int k = 0; k < num_coarse_neighbors; k++)\n {\n dest_coarse_B2L_maps[k].resize( dest_coarse_B2L_maps_scratch_sizes[k] );\n // Reset sizes to 0 (fill use as offset in next loop);\n dest_coarse_B2L_maps_scratch_sizes[k] = 0;\n }\n\n for (int i = 0; i < num_fine_neighbors; i++)\n {\n int k = fine_neigh_to_coarse_neigh[i];\n\n if (k != -1)\n {\n int offset = dest_coarse_B2L_maps_scratch_sizes[k];\n amgx::thrust::copy(coarse_B2L_maps[i].begin(), coarse_B2L_maps[i].end(), dest_coarse_B2L_maps[k].begin() + offset);\n dest_coarse_B2L_maps_scratch_sizes[k] += coarse_B2L_maps[i].size();\n }\n }\n\n cudaCheckError();\n int max_size = 0;\n\n for (int i = 0; i < num_coarse_neighbors; i++)\n {\n int size = dest_coarse_B2L_maps[i].size();\n\n if (size > max_size) { max_size = size; }\n }\n\n // Remove duplicates (aggregates in boundary that go to same merged partition)\n for (int i = 0; i < num_coarse_neighbors; i++)\n {\n int size = dest_coarse_B2L_maps[i].size();\n amgx::thrust::sort(dest_coarse_B2L_maps[i].begin(), dest_coarse_B2L_maps[i].begin() + size);\n index_type num_unique = amgx::thrust::unique(dest_coarse_B2L_maps[i].begin(), dest_coarse_B2L_maps[i].begin() + size) - dest_coarse_B2L_maps[i].begin();\n dest_coarse_B2L_maps[i].erase(dest_coarse_B2L_maps[i].begin() + num_unique, dest_coarse_B2L_maps[i].end());\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManagerBase::computeConsolidatedOffsets(const int my_id, const int my_destination_part, const bool is_root_partition, const int num_interior_rows, const int num_boundary_rows, IVector_h_vector &vertex_counts, const IVector_h &parts_to_consolidate, const int num_parts_to_consolidate, int &interior_offset, int &boundary_offset, int &total_interior_rows_in_merged, int &total_boundary_rows_in_merged, int &total_rows_in_merged, DistributedComms *comms)\n{\n IVector_h my_offsets(4);\n IVector_h my_sizes(2);\n my_sizes[0] = num_interior_rows;\n my_sizes[1] = num_boundary_rows;\n\n if (!is_root_partition)\n {\n //Send number of interior and boundary nodes to root\n comms->send_vector_async(my_sizes, my_destination_part, 777);\n comms->recv_vector(my_offsets, my_destination_part, 778);\n comms->send_vector_wait_all(my_sizes);\n }\n else\n {\n vertex_counts.resize(num_parts_to_consolidate);\n IVector_h child_sizes(2);\n IVector_h offsets_interior(num_parts_to_consolidate);\n IVector_h offsets_boundary(num_parts_to_consolidate);\n int count_int = 0;\n int count_bdy = 0;\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n if (parts_to_consolidate[i] == my_id)\n {\n child_sizes = my_sizes;\n }\n else\n {\n comms->recv_vector(child_sizes, parts_to_consolidate[i], 777);\n }\n\n //Do a simple cumulative sum to determine total number of interior/boundary rows and their offsets on a per contributing partition basis\n offsets_interior[i] = count_int;\n offsets_boundary[i] = count_bdy;\n count_int += child_sizes[0];\n count_bdy += child_sizes[1];\n //Save them\n vertex_counts[i].resize(2);\n vertex_counts[i][0] = child_sizes[0];\n vertex_counts[i][1] = child_sizes[1];\n }\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n //Send back to contributing partitions\n IVector_h offsets_to_send(4);\n offsets_to_send[0] = offsets_interior[i];\n offsets_to_send[1] = offsets_boundary[i];\n offsets_to_send[2] = count_int;\n offsets_to_send[3] = count_bdy;\n\n if (parts_to_consolidate[i] == my_id)\n {\n my_offsets = offsets_to_send;\n }\n else\n {\n comms->send_vector(offsets_to_send, parts_to_consolidate[i], 778); // cannot make async, rewriting internal buffer\n }\n }\n }\n\n interior_offset = my_offsets[0];\n boundary_offset = my_offsets[1] + my_offsets[2] - num_interior_rows;\n total_interior_rows_in_merged = my_offsets[2];\n total_boundary_rows_in_merged = my_offsets[3];\n total_rows_in_merged = my_offsets[2] + my_offsets[3];\n}\n\n\n\n\ntemplate \ntemplate \nvoid DistributedManagerBase::consB2LmapsOnRoot(int &num_consolidated_neighbors, std::vector &consolidated_B2L_maps, IVector_h &consolidated_coarse_ids, std::vector &dest_coarse_B2L_maps, IVector_h &coarse_neigh_to_fine_part, IVector_h &num_bdy_per_coarse_neigh, IVector_h &fine_parts_to_consolidate, int num_fine_parts_to_consolidate, int my_id, int my_destination_part, bool is_root_partition, int num_coarse_neighbors, DistributedComms *comms)\n{\n // TODO: it is possible to optimize exchanges, for example fuse recv_vector in recreating coarse neigbours\n\n // output: num_consolidated_neighbor, consolidated_B2L_maps, consolidated_coarse_ids\n // input: dest_coarse_B2L_maps, is_root_partition, my_id, my_destination_part, num_fine_parts_to_consolidate, num_coarse_neighbors, coarse_neigh_to_fine_part, num_bdy_per_coarse_neigh\n if (my_destination_part != my_id)\n {\n //if not root, send coarse neighbor list using fine indices and the corresponding boundary lists\n IVector_h num_coarse_neigh(1);\n num_coarse_neigh[0] = num_coarse_neighbors;\n comms->send_vector_async(num_coarse_neigh, my_destination_part, 1111);\n comms->send_vector_async(coarse_neigh_to_fine_part, my_destination_part, 2222);\n comms->send_vector_async(num_bdy_per_coarse_neigh, my_destination_part, 3333);\n\n for (int i = 0; i < num_coarse_neighbors; i++)\n {\n comms->send_vector_async(dest_coarse_B2L_maps[i], my_destination_part, 4444 + i) ;\n }\n\n comms->send_vector_wait_all(num_coarse_neigh);\n comms->send_vector_wait_all(coarse_neigh_to_fine_part);\n comms->send_vector_wait_all(num_bdy_per_coarse_neigh);\n\n for (int i = 0; i < num_coarse_neighbors; i++)\n {\n comms->send_vector_wait_all(dest_coarse_B2L_maps[i]) ;\n }\n }\n\n if (is_root_partition)\n {\n IVector_h num_coarse_ids_from_part(fine_parts_to_consolidate);\n IVector_h_vector coarse_ids_from_part(num_fine_parts_to_consolidate);\n IVector_h_vector num_coarse_neigh_bdys_from_part(num_fine_parts_to_consolidate);\n //If root, receive sizes, and resize receive buffers\n int total_num_coarse_ids = 0;\n\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n IVector_h temp(1);\n\n if (current_part != my_id)\n {\n comms->recv_vector(temp, current_part, 1111);\n }\n else\n {\n temp[0] = num_coarse_neighbors;\n }\n\n num_coarse_ids_from_part[i] = temp[0];\n coarse_ids_from_part[i].resize(temp[0]);\n num_coarse_neigh_bdys_from_part[i].resize(temp[0]);\n total_num_coarse_ids += temp[0];\n }\n\n //Create a neighbor list for the consolidated coarse matrix, by merging coarse neighbor lists from partitions that are being merged\n consolidated_coarse_ids.resize(total_num_coarse_ids);\n int count = 0;\n\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n\n // Get from each partition the coarse partition ids in their B2L maps\n if (current_part != my_id)\n {\n comms->recv_vector(coarse_ids_from_part[i], current_part, 2222);\n comms->recv_vector(num_coarse_neigh_bdys_from_part[i], current_part, 3333);\n }\n else\n {\n coarse_ids_from_part[i] = coarse_neigh_to_fine_part;\n num_coarse_neigh_bdys_from_part[i] = num_bdy_per_coarse_neigh;\n }\n\n amgx::thrust::copy(coarse_ids_from_part[i].begin(), coarse_ids_from_part[i].end(), consolidated_coarse_ids.begin() + count);\n count += num_coarse_ids_from_part[i];\n }\n\n cudaCheckError();\n //eliminate duplicates\n amgx::thrust::sort(consolidated_coarse_ids.begin(), consolidated_coarse_ids.end());\n cudaCheckError();\n consolidated_coarse_ids.erase(amgx::thrust::unique(consolidated_coarse_ids.begin(), consolidated_coarse_ids.end()), consolidated_coarse_ids.end());\n cudaCheckError();\n num_consolidated_neighbors = consolidated_coarse_ids.size();\n IVector_h_vector coarse_ids_from_part_to_consolidated_neighbor(num_fine_parts_to_consolidate);;\n\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n coarse_ids_from_part_to_consolidated_neighbor[i].resize(num_coarse_ids_from_part[i]);\n amgx::thrust::lower_bound(consolidated_coarse_ids.begin(), consolidated_coarse_ids.end(), coarse_ids_from_part[i].begin(), coarse_ids_from_part[i].end(), coarse_ids_from_part_to_consolidated_neighbor[i].begin());\n }\n\n cudaCheckError();\n // Map each coarse partition to new coarse ID\n consolidated_B2L_maps.resize(num_consolidated_neighbors);\n IVector_h consolidated_B2L_maps_sizes(num_consolidated_neighbors);\n // Offset in the consolidated_B2L_maps\n IVector_h_vector coarse_ids_offsets(num_fine_parts_to_consolidate);\n\n for (int i = 0; i < num_consolidated_neighbors; i++)\n {\n consolidated_B2L_maps_sizes[i] = 0;\n }\n\n // Compute the size of each consolidated_B2L_maps and offsets into it, where we will receive the parts coming from partitions that are getting merged into this one\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n coarse_ids_offsets[i].resize(num_coarse_ids_from_part[i]);\n\n for (int j = 0; j < num_coarse_ids_from_part[i]; j++)\n {\n int coarse_id = coarse_ids_from_part[i][j];\n int k = num_coarse_neigh_bdys_from_part[i][j];\n coarse_ids_offsets[i][j] = consolidated_B2L_maps_sizes[ coarse_ids_from_part_to_consolidated_neighbor[i][j] ];\n consolidated_B2L_maps_sizes[ coarse_ids_from_part_to_consolidated_neighbor[i][j] ] += k;\n }\n }\n\n for (int i = 0; i < num_consolidated_neighbors; i++)\n {\n consolidated_B2L_maps[i].resize(consolidated_B2L_maps_sizes[i]);\n }\n\n // Receive the B2L maps from each child partition, concatenate them (gets sorted outside)\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n\n for (int j = 0; j < num_coarse_ids_from_part[i]; j++)\n {\n int my_coarse_neigh = coarse_ids_from_part_to_consolidated_neighbor[i][j];\n int offset = coarse_ids_offsets[i][j];\n\n if (current_part != my_id)\n {\n comms->recv_vector( consolidated_B2L_maps[my_coarse_neigh], current_part, 4444 + j, offset, num_coarse_neigh_bdys_from_part[i][j]); //Need to do proper tagging here, otherwise messages from the same source would get mixed up\n }\n else\n {\n amgx::thrust::copy(dest_coarse_B2L_maps[j].begin(), dest_coarse_B2L_maps[j].end(), consolidated_B2L_maps[my_coarse_neigh].begin() + offset);\n }\n }\n }\n\n cudaCheckError();\n }\n}\n\n\ntemplate \nvoid DistributedManagerBase::consolidateAndRenumberHalos(IVector_h &aggregates, const IVector_h &manager_halo_offsets, IVector_h &halo_offsets, const IVector_h &neighbors, int num_fine_neighbors, const IVector_h &consolidated_coarse_ids, int num_consolidated_neighbors, const IVector_h &destination_part, int my_destination_part, bool is_root_partition, IVector_h &fine_parts_to_consolidate, int num_fine_parts_to_consolidate, int num_parts, int my_id, int total_rows_in_merged, int &num_all_aggregates, DistributedComms *comms)\n{\n consAndRenumberHalos(aggregates, manager_halo_offsets, halo_offsets, neighbors, num_fine_neighbors, consolidated_coarse_ids, num_consolidated_neighbors, destination_part, my_destination_part, is_root_partition, fine_parts_to_consolidate, num_fine_parts_to_consolidate, num_parts, my_id, total_rows_in_merged, num_all_aggregates, comms);\n}\n\ntemplate \nvoid DistributedManagerBase::consolidateAndRenumberHalos(IVector_d &aggregates, const IVector_h &manager_halo_offsets, IVector_h &halo_offsets, const IVector_h &neighbors, int num_fine_neighbors, const IVector_h &consolidated_coarse_ids, int num_consolidated_neighbors, const IVector_h &destination_part, int my_destination_part, bool is_root_partition, IVector_h &fine_parts_to_consolidate, int num_fine_parts_to_consolidate, int num_parts, int my_id, int total_rows_in_merged, int &num_all_aggregates, DistributedComms *comms)\n{\n consAndRenumberHalos(aggregates, manager_halo_offsets, halo_offsets, neighbors, num_fine_neighbors, consolidated_coarse_ids, num_consolidated_neighbors, destination_part, my_destination_part, is_root_partition, fine_parts_to_consolidate, num_fine_parts_to_consolidate, num_parts, my_id, total_rows_in_merged, num_all_aggregates, comms);\n}\n\ntemplate \ntemplate \nvoid DistributedManagerBase::consAndRenumberHalos(IVector_hd &aggregates, const IVector_h &manager_halo_offsets, IVector_h &halo_offsets, const IVector_h &neighbors, int num_fine_neighbors, const IVector_h &consolidated_coarse_neigh_to_fine_part, int num_consolidated_neighbors, const IVector_h &destination_part, int my_destination_part, bool is_root_partition, IVector_h &fine_parts_to_consolidate, int num_fine_parts_to_consolidate, int num_parts, int my_id, int total_rows_in_merged, int &num_all_aggregates, DistributedComms *comms)\n{\n /*\n * EXAMPLE 2\n This example is independent from the previous ones.\n Say partition 0 and 1 are merging (into 0) partition 0 is neighbors with 1,2,3 and partition 1 is neighbors with 0,3,4\n Partitions 3 and 4 are merging (into partition 3) and partition 2 is not merging with anyone.\n This example details the renumbering of halo indices on partition 0 and partition 1.\n aggregates on partition 0:\n [(fine interior nodes) (fine boundary nodes) (fine halo from part 1) (fine halo from part 2) (fine halo from part 3)]\n [(fine interior nodes) (fine boundary nodes) (13 13 15) (12 15 17) (14 16 18)]\n aggregates on partition 1:\n [(fine interior nodes) (fine boundary nodes) (fine halo from part 0) (fine halo from part 3) (fine halo from part 4)]\n [(fine interior nodes) (fine boundary nodes) (14 16 17) (18 19 19) (15 15 17)]\n\n manager_halo_offsets on partition 0:\n [22 25 28 31]\n manager_halo_offsets on partition 1:\n [20 23 26 29]\n halo_offsets on both partitions are uninitialised: [0 0 0] and [0 0]\n neighbors on partition 0: [1 2 3] partition 1: [0 3 4]\n num_fine_neighbors partition 0: 3 partition 1: 3\n consolidated_coarse_neigh_to_fine_part partition 0: [2 3] partition 1: [3]\n num_consolidated_neighbors partition 0: 2 partition 1: 1\n destination_part [0 0 2 3 3]\n my_destination_part partition 0: 0 partition 1: 0\n is_root_partition partition 0: true partition 1: false\n fine_parts_to_consolidate partition 0: [0 1]\n num_fine_parts_to_consolidate partition 0: 2\n num_parts 5\n my_id partition 0: 0 partition 1: 1\n total_rows_in_merged partition 0 and 1: 24 (=sum of the two below)\n num_all_aggregates partition partition 0: 13 partition 1: 11 - will be updated with the number of halo aggregates\n */\n //\n // Step 9.2 - com up with nonmerge lists\n //\n int num_fine_nonmerge_neighbors;// = fine_nonmerge_neighbors.size();\n //NUmber of neighbors we are not merging with\n num_fine_nonmerge_neighbors = 0;\n\n for (int i = 0 ; i < num_fine_neighbors; i++)\n {\n if (destination_part[neighbors[i]] != my_destination_part)\n {\n num_fine_nonmerge_neighbors++;\n }\n }\n\n IVector_h halo_sizes(num_fine_nonmerge_neighbors);\n IVector_h fine_nonmerge_neigh_to_cons_fine_part(num_fine_nonmerge_neighbors);\n IVector_h fine_nonmerge_neighbor_to_fine_neighbor(num_fine_nonmerge_neighbors);\n num_fine_nonmerge_neighbors = 0;\n\n for (int i = 0 ; i < num_fine_neighbors; i++)\n {\n if (destination_part[neighbors[i]] != my_destination_part)\n {\n halo_sizes[num_fine_nonmerge_neighbors] = manager_halo_offsets[i + 1] - manager_halo_offsets[i];\n fine_nonmerge_neighbor_to_fine_neighbor[num_fine_nonmerge_neighbors] = i;\n fine_nonmerge_neigh_to_cons_fine_part[num_fine_nonmerge_neighbors] = destination_part[neighbors[i]];\n num_fine_nonmerge_neighbors++;\n }\n }\n\n /*\n * EXAMPLE 2\n num_fine_nonmerge_neighbors partition 0: 2 partition 1: 2\n fine_nonmerge_neighbor_to_fine_neighbor partition 0: [1 2] partition 1: [1 2]\n fine_nonmerge_neigh_to_cons_fine_part partition 0: [2 3] partition 1: [3 3]\n halo_sizes partition 0: [3 3] partition 1: [3 3]\n */\n\n //Send them to root along with the halo parts of the aggregates vector\n if (!is_root_partition)\n {\n IVector_h num_fine_nonmerge_neigh(1);\n num_fine_nonmerge_neigh[0] = num_fine_nonmerge_neighbors;\n // TODO: async? might be faster.\n comms->send_vector(num_fine_nonmerge_neigh, my_destination_part, 1111);\n comms->send_vector(halo_sizes, my_destination_part, 2222);\n comms->send_vector(fine_nonmerge_neigh_to_cons_fine_part, my_destination_part, 3333);\n\n // Here check l2h_identity flag and act accordingly\n for (int i = 0; i < num_fine_nonmerge_neighbors; i++)\n {\n comms->send_vector_async(aggregates, my_destination_part, 4444 + i, manager_halo_offsets[fine_nonmerge_neighbor_to_fine_neighbor[i]], halo_sizes[i]) ;\n }\n\n //comms->send_vector_wait_all(num_fine_nonmerge_neigh);\n //comms->send_vector_wait_all(halo_sizes);\n //comms->send_vector_wait_all(fine_nonmerge_neigh_to_cons_fine_part);\n comms->send_vector_wait_all(aggregates);\n /*\n * EXAMPLE 2\n Partition 1 sends to partition 0:\n num_fine_nonmerge_neigh 2\n halo_sizes [3 3]\n fine_nonmerge_neigh_to_cons_fine_part [3 3]\n for loop: sends two pieces: [(18 19 19)] [(15 15 17)]\n */\n }\n\n if (is_root_partition)\n {\n //\n // Step 9.3 Root receives this info, creates metadata\n //\n std::vector num_fine_nonmerge_neigh_array(num_fine_parts_to_consolidate);\n IVector_h_vector halo_sizes_array(num_fine_parts_to_consolidate);\n IVector_h_vector fine_nonmerge_neigh_to_cons_fine_part_array(num_fine_parts_to_consolidate);\n std::vector > fine_halo_aggregates_to_root_array(num_fine_parts_to_consolidate);\n std::vector min_index_coarse_halo(num_consolidated_neighbors, 0x7FFFFFFF);\n std::vector max_index_coarse_halo(num_consolidated_neighbors, 0);\n std::vector fine_part_to_consolidated_neighbor(num_parts, -1);\n\n for (int i = 0; i < num_consolidated_neighbors; i++)\n {\n fine_part_to_consolidated_neighbor[consolidated_coarse_neigh_to_fine_part[i]] = i;\n }\n\n /*\n * EXAMPLE 2\n everything from here on is for partition 0, since that is the root partition\n fine_part_to_consolidated_neighbor [-1 -1 0 1 -1]\n */\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n IVector_h temp(1);\n\n if (current_part != my_id)\n {\n comms->recv_vector(temp, current_part, 1111);\n }\n else\n {\n temp[0] = num_fine_nonmerge_neighbors;\n }\n\n num_fine_nonmerge_neigh_array[i] = temp[0];\n halo_sizes_array[i].resize(temp[0]);\n fine_nonmerge_neigh_to_cons_fine_part_array[i].resize(temp[0]);\n fine_halo_aggregates_to_root_array[i].resize(temp[0]);\n\n if (current_part != my_id)\n {\n comms->recv_vector(halo_sizes_array[i], current_part, 2222);\n }\n else\n {\n halo_sizes_array[i] = halo_sizes;\n }\n\n if (current_part != my_id)\n {\n comms->recv_vector(fine_nonmerge_neigh_to_cons_fine_part_array[i], current_part, 3333);\n }\n else\n {\n fine_nonmerge_neigh_to_cons_fine_part_array[i] = fine_nonmerge_neigh_to_cons_fine_part;\n }\n\n //Receive the halo regions\n for (int j = 0; j < temp[0]; j++)\n {\n fine_halo_aggregates_to_root_array[i][j].resize(halo_sizes_array[i][j]);\n\n if (current_part != my_id)\n {\n comms->recv_vector(fine_halo_aggregates_to_root_array[i][j], current_part, 4444 + j);\n }\n else\n {\n //HERE\n amgx::thrust::copy(aggregates.begin() + manager_halo_offsets[fine_nonmerge_neighbor_to_fine_neighbor[j]],\n aggregates.begin() + manager_halo_offsets[fine_nonmerge_neighbor_to_fine_neighbor[j]] + halo_sizes[j],\n fine_halo_aggregates_to_root_array[i][j].begin()); //TODO: not do this copying around on the root\n }\n\n#define MIN(a,b) ab?a:b;\n //Find minimum and maximum halo indices as not to allocate too much scratch space later\n int min_index = amgx::thrust::reduce(fine_halo_aggregates_to_root_array[i][j].begin(), fine_halo_aggregates_to_root_array[i][j].end(), int(0x7FFFFFFF), amgx::thrust::minimum());\n int max_index = amgx::thrust::reduce(fine_halo_aggregates_to_root_array[i][j].begin(), fine_halo_aggregates_to_root_array[i][j].end(), int(0), amgx::thrust::maximum());\n min_index_coarse_halo[fine_part_to_consolidated_neighbor[fine_nonmerge_neigh_to_cons_fine_part_array[i][j]]] = MIN((int)min_index_coarse_halo[fine_part_to_consolidated_neighbor[fine_nonmerge_neigh_to_cons_fine_part_array[i][j]]], min_index);\n max_index_coarse_halo[fine_part_to_consolidated_neighbor[fine_nonmerge_neigh_to_cons_fine_part_array[i][j]]] = MAX((int)max_index_coarse_halo[fine_part_to_consolidated_neighbor[fine_nonmerge_neigh_to_cons_fine_part_array[i][j]]], max_index);\n }\n }\n\n cudaCheckError();\n /*\n * EXAMPLE 2\n num_fine_nonmerge_neigh_array = [2 2]\n halo_sizes_array = [[3 3][3 3]]\n fine_nonmerge_neigh_to_cons_fine_part_array[][] = [[2 3][3 3]]\n fine_halo_aggregates_to_root_array[from][to][fine halo vertex] [[[12 15 17][14 16 18]]\n [[18 19 19][15 15 17]]]\n min_index_coarse_halo[12 14]\n max_index_coarse_halo[17 19]\n */\n halo_offsets[0] = total_rows_in_merged;\n //Now we have all the halo nodes, let's renumber them.\n int min_index = amgx::thrust::reduce(min_index_coarse_halo.begin(), min_index_coarse_halo.end(), int(0x7FFFFFFF), amgx::thrust::minimum());\n int max_index = amgx::thrust::reduce(max_index_coarse_halo.begin(), max_index_coarse_halo.end(), int(0), amgx::thrust::maximum());\n cudaCheckError();\n //\n // Step 9.4 compute halo indices on root nodes\n //\n int scratch_size;\n\n if (num_consolidated_neighbors == 0)\n {\n scratch_size = 1;\n }\n else\n {\n scratch_size = max_index - min_index + 2;\n }\n\n IVector scratch(scratch_size);\n\n for (int i = 0; i < num_consolidated_neighbors; i++)\n {\n thrust_wrapper::fill(scratch.begin(), scratch.end(), 0);\n int dest_part = consolidated_coarse_neigh_to_fine_part[i];\n\n //Flag halo indices that occur for a specific coarse neighbor\n for (int j = 0; j < num_fine_parts_to_consolidate; j++)\n {\n for (int k = 0; k < num_fine_nonmerge_neigh_array[j]; k++)\n {\n if (fine_nonmerge_neigh_to_cons_fine_part_array[j][k] == dest_part)\n {\n int size = halo_sizes_array[j][k];\n this->flag_halo_ids(size, scratch, fine_halo_aggregates_to_root_array[j][k], min_index_coarse_halo[i], max_index, min_index);\n }\n }\n }\n\n thrust_wrapper::exclusive_scan(scratch.begin(), scratch.end(), scratch.begin(), halo_offsets[i]); //renumber them with the proper offset into our halo\n halo_offsets[i + 1] = scratch[scratch.size() - 1];\n\n //now read them back\n for (int j = 0; j < num_fine_parts_to_consolidate; j++)\n {\n for (int k = 0; k < num_fine_nonmerge_neigh_array[j]; k++)\n {\n if (fine_nonmerge_neigh_to_cons_fine_part_array[j][k] == dest_part)\n {\n int size = halo_sizes_array[j][k];\n int block_size = 128;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (size - 1) / block_size + 1);\n this->read_halo_ids(size, scratch, fine_halo_aggregates_to_root_array[j][k], min_index_coarse_halo[i]);\n //and send them back to contributing partitions\n hipDeviceSynchronize(); //TODO: don't need to synchronize when using GPUDirect\n cudaCheckError();\n int current_part = fine_parts_to_consolidate[j];\n int tag = 4444 + dest_part;\n\n if (current_part != my_id)\n {\n comms->send_vector_async(fine_halo_aggregates_to_root_array[j][k], current_part, tag); //!!!!: we are sending them back not in sequential order, need tags!!!!\n }\n else\n {\n amgx::thrust::copy(fine_halo_aggregates_to_root_array[j][k].begin(), fine_halo_aggregates_to_root_array[j][k].end(), aggregates.begin() + manager_halo_offsets[fine_nonmerge_neighbor_to_fine_neighbor[k]]);\n }\n }\n }\n }\n\n /*\n * EXAMPLE 2\n the array that is sent back in pieces:\n fine_halo_aggregates_to_root_array[from][to][fine halo vertex] [[[24 25 26][27 29 31]]\n [[31 32 32][28 28 30]]]\n halo_offsets = [24 27 33]\n */\n } // Loop over consolidated neighbors\n\n cudaCheckError();\n\n // Wait for sends to have completed (this is to prevent fine_halo_aggregates_to_root_array to be destroyed before send has finished)\n for (int i = 0; i < num_consolidated_neighbors; i++)\n {\n int dest_part = consolidated_coarse_neigh_to_fine_part[i];\n\n for (int j = 0; j < num_fine_parts_to_consolidate; j++)\n {\n for (int k = 0; k < num_fine_nonmerge_neigh_array[j]; k++)\n {\n if (fine_nonmerge_neigh_to_cons_fine_part_array[j][k] == dest_part)\n {\n int current_part = fine_parts_to_consolidate[j];\n\n if (current_part != my_id)\n {\n comms->send_vector_wait_all(fine_halo_aggregates_to_root_array[j][k]);\n }\n }\n }\n }\n } // Loop over consolidated neighbors\n\n //Send total number of rows in the aggregated matrix\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n IVector_h total_rows(1);\n total_rows[0] = halo_offsets[num_consolidated_neighbors];\n\n if (current_part != my_id)\n {\n comms->send_vector(total_rows, current_part, 5555);\n }\n else\n {\n num_all_aggregates = total_rows[0];\n }\n }\n } // If is root partition\n\n if (!is_root_partition)\n {\n for (int i = 0; i < num_fine_nonmerge_neighbors; i++)\n {\n int tag = 4444 + fine_nonmerge_neigh_to_cons_fine_part[i];\n comms->recv_vector(aggregates, my_destination_part, tag, manager_halo_offsets[fine_nonmerge_neighbor_to_fine_neighbor[i]], halo_sizes[i]);\n }\n\n IVector_h total_rows(1);\n comms->recv_vector(total_rows, my_destination_part, 5555);\n num_all_aggregates = total_rows[0];\n }\n\n /*\n * EXAMPLE 2\n num_all_aggregates = 33 (both partitions 0 and 1\n */\n}\n\ntemplate \nvoid DistributedManagerBase::ipcExchangePtr(void *&ptr, bool is_root_partition, int num_parts_to_consolidate, IVector_h &parts_to_consolidate, int my_root_partition, int my_id, DistributedComms *comms)\n{\n hipIpcMemHandle_t handle;\n\n if (is_root_partition)\n {\n hipIpcGetMemHandle( (hipIpcMemHandle_t *) &handle, ptr ) ;\n cudaCheckError();\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n int current_part = parts_to_consolidate[i];\n\n if (current_part != my_id)\n {\n comms->send_raw_data(&handle, sizeof(handle), current_part, 456);\n }\n }\n }\n else\n {\n comms->recv_raw_data(&handle, sizeof(handle), my_root_partition, 456);\n hipError_t err = hipIpcOpenMemHandle( (void **) &ptr, handle, hipIpcMemLazyEnablePeerAccess);\n cudaCheckError();\n }\n}\n\ntemplate \nvoid DistributedManagerBase::ipcWaitForChildren(bool is_root_partition, int num_parts_to_consolidate, IVector_h &parts_to_consolidate, int my_destination_part, int my_id, DistributedComms *comms)\n{\n hipEvent_t event;\n hipIpcEventHandle_t event_handle;\n hipEventCreateWithFlags(&event, hipEventDisableTiming | hipEventInterprocess);\n cudaCheckError();\n hipIpcGetEventHandle( &event_handle, event);\n cudaCheckError();\n // Each rank record the event\n hipEventRecord(event);\n cudaCheckError();\n\n if (is_root_partition)\n {\n std::vector child_events(num_parts_to_consolidate);\n std::vector child_event_handles(num_parts_to_consolidate);\n\n // Root partition receives event_handles from child and stores in child_event_handles\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n int current_part = parts_to_consolidate[i];\n\n if (current_part != my_id)\n {\n comms->recv_raw_data(&(child_event_handles[i]), sizeof(hipIpcEventHandle_t), current_part, 987 + current_part);\n hipIpcOpenEventHandle(&child_events[i], child_event_handles[i]);\n cudaCheckError();\n }\n }\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n if (parts_to_consolidate[i] != my_id)\n {\n hipEventSynchronize(child_events[i]);\n cudaCheckError();\n }\n }\n }\n else\n {\n comms->send_raw_data(&event_handle, sizeof(hipIpcEventHandle_t), my_destination_part, 987 + my_id);\n }\n}\n\ntemplate \nvoid DistributedManagerBase::ipcWaitForRoot(bool is_root_partition, int num_parts_to_consolidate, IVector_h &parts_to_consolidate, int my_destination_part, int my_id, DistributedComms *comms)\n{\n hipEvent_t event;\n hipIpcEventHandle_t event_handle;\n hipEventCreateWithFlags(&event, hipEventDisableTiming | hipEventInterprocess);\n\n if (is_root_partition)\n {\n hipIpcGetEventHandle( &event_handle, event);\n cudaCheckError();\n // Root records the event\n hipEventRecord(event);\n cudaCheckError();\n\n // Root partition sends event_handles to child\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n int current_part = parts_to_consolidate[i];\n\n if (current_part != my_id)\n {\n comms->send_raw_data(&event_handle, sizeof(event_handle), current_part, 988 + current_part);\n }\n }\n }\n else\n {\n comms->recv_raw_data(&event_handle, sizeof(event_handle), my_destination_part, 988 + my_id);\n hipIpcOpenEventHandle(&event, event_handle);\n cudaCheckError();\n hipEventSynchronize(event);\n cudaCheckError();\n }\n}\n\n\n\ntemplate \nvoid DistributedManagerBase::read_halo_ids(int size, IVector_d &scratch, IVector_d &halo_aggregates, VecInt_t min_index_coarse_halo)\n{\n int block_size = 128;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (size - 1) / block_size + 1);\n read_halo_ids_kernel <<< num_blocks, block_size>>>(scratch.raw(), halo_aggregates.raw(), min_index_coarse_halo, size);\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManagerBase::read_halo_ids(int size, IVector_h &scratch, IVector_h &halo_aggregates, VecInt_t min_index_coarse_halo)\n{\n FatalError(\"read_halo_ids not implemented on host yet\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n\ntemplate \nvoid DistributedManagerBase::flag_halo_ids(int size, IVector_d &scratch, IVector_d &halo_aggregates, VecInt_t min_index_coarse_halo, int max_index, int min_index)\n{\n int block_size = 128;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (size - 1) / block_size + 1);\n flag_halo_ids_kernel <<< num_blocks, block_size>>>(scratch.raw(), halo_aggregates.raw(), min_index_coarse_halo, size, max_index - min_index + 1);\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManagerBase::flag_halo_ids(int size, IVector_h &scratch, IVector_h &halo_aggregates, VecInt_t min_index_coarse_halo, int max_index, int min_index)\n{\n FatalError(\"flag_halo_ids not implemented on host yet\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid DistributedManager >::consolidateAndUploadAll(int n, int nnz, int block_dimx, int block_dimy, const int *row_ptrs, const int *col_indices, const void *data, const void *diag, Matrix &A)\n{\n FatalError(\"Fine level consolidation not implemented on host yet\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate\nvoid DistributedManagerBase::exchangeSolveResultsConsolidation(int &num_iters, std::vector &res_history, AMGX_STATUS &status, bool store_res_history)\n{\n int bsize = this->A->get_block_size();\n PODVector_h res_history_tmp;\n\n if (!m_is_fine_level_consolidated)\n {\n return;\n }\n else\n {\n int my_id = this->getFineLevelComms()->get_global_id();\n IVector_h my_num_iters(1);\n\n if (m_is_fine_level_root_partition)\n {\n my_num_iters[0] = num_iters;\n\n if (store_res_history)\n {\n // Pack the res_history vector into array\n res_history_tmp.resize( (num_iters + 1)*bsize);\n\n for (int i = 0; i < num_iters + 1; i++)\n {\n for (int j = 0; j < bsize; j++)\n {\n res_history_tmp[i * bsize + j] = res_history[i][j];\n }\n }\n }\n\n for (int i = 0; i < m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = m_fine_level_parts_to_consolidate[i];\n\n if (my_id != current_part)\n {\n getFineLevelComms()->send_vector_async(my_num_iters, current_part, 245);\n\n if (store_res_history)\n {\n getFineLevelComms()->send_vector_async(res_history_tmp, current_part, 246);\n }\n }\n }\n\n for (int i = 0; i < m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = m_fine_level_parts_to_consolidate[i];\n\n if (my_id != current_part)\n {\n getFineLevelComms()->send_raw_data(&status, sizeof(status), current_part, 247);\n }\n }\n\n getFineLevelComms()->send_vector_wait_all(my_num_iters);\n\n if (store_res_history)\n {\n getFineLevelComms()->send_vector_wait_all(res_history_tmp);\n }\n }\n else\n {\n // Store num_iters\n getFineLevelComms()->recv_vector(my_num_iters, m_my_fine_level_destination_part, 245);\n num_iters = my_num_iters[0];\n\n if (store_res_history)\n {\n // Fill res_history vector\n res_history.resize(num_iters + 1);\n res_history_tmp.resize( (num_iters + 1)*bsize);\n getFineLevelComms()->recv_vector(res_history_tmp, m_my_fine_level_destination_part, 246);\n\n for (int i = 0; i < num_iters + 1; i++)\n {\n res_history[i].resize(bsize);\n\n for (int j = 0; j < bsize; j++)\n {\n res_history[i][j] = res_history_tmp[i * bsize + j];\n }\n }\n }\n\n getFineLevelComms()->recv_raw_data(&status, sizeof(status), m_my_fine_level_destination_part, 247);\n }\n }\n}\n\ntemplate \nvoid DistributedManager >::consolidateAndUploadAll(int n, int nnz, int block_dimx, int block_dimy, const int *row_ptrs, const int *col_indices, const void *data, const void *diag, Matrix &in_A)\n{\n this->A = &in_A;\n this->createComms(this->A->getResources()); //refresh comms\n DistributedComms *comms = this->getComms();\n int my_id = comms->get_global_id();\n int num_parts = comms->get_num_partitions();\n int num_rings = this->num_halo_rings();\n int num_neighbors = this->neighbors.size();\n // All partitions have to call this, otherwise it fails\n // Step 1: Figure out which partition should be consolidated together based on their host_name and their PCI-E slot ID\n IVector_h destination_part(num_parts);\n this->computeDestinationPartitionsWithCons(my_id, num_parts, destination_part, comms);\n int my_destination_part = destination_part[my_id];\n // Check if I'm root partition and how many msgs I will receive\n bool is_root_partition = false;\n int num_parts_to_consolidate = 0;\n\n for (int i = 0; i < num_parts; i++)\n {\n if (destination_part[i] == my_id)\n {\n is_root_partition = true;\n num_parts_to_consolidate++;\n }\n }\n\n if (my_destination_part >= num_parts)\n {\n FatalError(\"During consolidation, sending data to partition that doesn't exist\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // Create cons_part_to_part map\n IVector_h cons_part_to_part = destination_part;\n amgx::thrust::sort(cons_part_to_part.begin(), cons_part_to_part.end());\n cudaCheckError();\n cons_part_to_part.erase(amgx::thrust::unique(cons_part_to_part.begin(), cons_part_to_part.end()), cons_part_to_part.end());\n cudaCheckError();\n int num_cons_partitions = cons_part_to_part.size();\n\n // If number of consolidated partitions is the same as number of partitions, simply call uploadAll\n if (num_cons_partitions == num_parts)\n {\n this->initializeUploadReorderAll(n, nnz, block_dimx, block_dimy, row_ptrs, col_indices, data, diag, *(this->A));\n this->m_is_fine_level_consolidated = false;\n return;\n }\n\n if (is_root_partition)\n {\n this->A->getResources()->expandRootPool();\n }\n\n this->m_is_fine_level_consolidated = true;\n\n if (num_rings != 1)\n {\n FatalError(\"num_rings > 1 not supported in fine_level consolidation\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // Fill with b2l_maps\n IVector_h_vector B2L_maps_tmp;\n B2L_maps_tmp.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n B2L_maps_tmp[i] = this->cached_B2L_maps[i];\n }\n\n bool useCudaIpc = this->m_use_cuda_ipc_consolidation;\n mat_value_type *data_hd = NULL;\n mat_value_type *diag_hd = NULL;\n int *col_indices_hd = NULL;\n int data_alloc = 0;\n int diag_alloc = 0;\n int col_alloc = 0;\n col_indices_hd = (int *) this->getDevicePointerForData((void *)col_indices, nnz * block_dimx * block_dimy * sizeof(int), &col_alloc);\n data_hd = (mat_value_type *) this->getDevicePointerForData((void *)data, nnz * block_dimx * block_dimy * sizeof(mat_value_type), &data_alloc);\n\n if (diag != NULL)\n {\n diag_hd = (mat_value_type *) this->getDevicePointerForData((void *)diag, nnz * block_dimx * block_dimy * sizeof(mat_value_type), &diag_alloc);\n }\n\n // Copy the original row_offsets array (this is required when replacing coefficients\n this->m_old_row_offsets_CONS.resize(n + 1);\n hipMemcpy(this->m_old_row_offsets_CONS.raw(), row_ptrs, (n + 1)*sizeof(int), hipMemcpyDefault);\n cudaCheckError();\n this->m_old_nnz_CONS = nnz;\n // This function:\n // Creates fine level consolidated matrices\n // Modifies the btl_maps, lth_maps\n // Create part_to_cons_part map\n IVector_h part_to_cons_part(num_parts);\n amgx::thrust::lower_bound(cons_part_to_part.begin(), cons_part_to_part.end(), destination_part.begin(), destination_part.end(), part_to_cons_part.begin());\n cudaCheckError();\n IVector_h neigh_to_part;\n this->createNeighToDestPartMap(neigh_to_part, this->neighbors, destination_part, num_neighbors);\n IVector_h cons_neigh_to_part;\n int num_cons_neighbors;\n this->createConsolidatedNeighToPartMap(cons_neigh_to_part, neigh_to_part, my_destination_part, destination_part, num_cons_neighbors);\n IVector_h neigh_to_cons_neigh;\n this->createNeighToConsNeigh( neigh_to_cons_neigh, cons_neigh_to_part, neigh_to_part, my_destination_part, num_neighbors);\n // ---------------------------------------\n // MERGE B2L MAPS BASED ON DEST PARTITION\n // ---------------------------------------\n IVector_h_vector dest_B2L_maps;\n this->consolidateB2Lmaps(dest_B2L_maps, B2L_maps_tmp, neigh_to_cons_neigh, num_cons_neighbors, num_neighbors);\n // ------------------------------------\n // Renumber interior and boundary rows\n // ------------------------------------\n int num_interior_rows;\n int num_boundary_rows;\n IVector_h renumbering;\n this->createAggregatesRenumbering(renumbering, dest_B2L_maps, n, num_cons_neighbors, num_interior_rows, num_boundary_rows, num_rings);\n // --------------------------------------------------\n // Create list of destination parts to consolidate\n // --------------------------------------------------\n // Store whether or not this is a root partition on fine level\n IVector_h parts_to_consolidate;\n parts_to_consolidate.resize(num_parts_to_consolidate);\n int count = 0;\n\n for (int i = 0; i < num_parts; i++)\n {\n if (destination_part[i] == my_id)\n {\n parts_to_consolidate[count] = i;\n count++;\n }\n }\n\n // ---------------------------------------------------------------------\n // Each partition computes its offset for its interior and boundary nodes\n // ---------------------------------------------------------------------\n IVector_h_vector vertex_counts;\n int interior_offset, boundary_offset, total_interior_rows_in_merged, total_boundary_rows_in_merged;\n int total_rows_in_merged;\n this->computeConsolidatedOffsets(my_id, my_destination_part, is_root_partition, num_interior_rows, num_boundary_rows, vertex_counts, parts_to_consolidate, num_parts_to_consolidate, interior_offset, boundary_offset, total_interior_rows_in_merged, total_boundary_rows_in_merged, total_rows_in_merged, comms);\n // -----------------------------------\n // Each partition renumber it's rows\n // -----------------------------------\n int total_num_halos = 0;\n\n // Pack new bdy_ids\n for (int i = 0; i < num_neighbors; i++)\n {\n total_num_halos += this->cached_L2H_maps[i].size();\n }\n\n IVector_h row_ids(n + total_num_halos, -1);\n this->m_row_ids_CONS.resize(n + total_num_halos);\n\n // Renumber the interior and boundary rows\n for (int i = 0; i < n; i++)\n {\n int new_id;\n\n if (renumbering.size() == 0)\n {\n new_id = i;\n }\n else\n {\n new_id = renumbering[i];\n }\n\n new_id += ((new_id >= num_interior_rows) ? boundary_offset : interior_offset);\n row_ids[i] = new_id;\n }\n\n for (int i = 0; i < num_cons_neighbors; i++)\n {\n thrust_wrapper::transform(dest_B2L_maps[i].begin(),\n dest_B2L_maps[i].end(),\n amgx::thrust::constant_iterator(boundary_offset),\n dest_B2L_maps[i].begin(),\n amgx::thrust::plus());\n }\n\n cudaCheckError();\n // -------------------------------------------------\n // Send dest_B2L_maps to root partitions\n // ------------------------------------------------\n IVector_h num_bdy_per_cons_neigh(num_cons_neighbors);\n\n for (int i = 0; i < num_cons_neighbors; i++)\n {\n num_bdy_per_cons_neigh[i] = dest_B2L_maps[i].size();\n }\n\n IVector_h root_cons_neighbors;\n int root_num_cons_neighbors = 0;\n IVector_h_vector cons_B2L_maps;\n this->consolidateB2LmapsOnRoot(root_num_cons_neighbors, cons_B2L_maps, root_cons_neighbors, dest_B2L_maps, cons_neigh_to_part, num_bdy_per_cons_neigh, parts_to_consolidate, num_parts_to_consolidate, my_id, my_destination_part, is_root_partition, num_cons_neighbors, comms);\n IVector_h halo_ids_offsets(num_neighbors + 1);\n IVector_h halo_ids;\n int halo_ids_size = 0;\n halo_ids_offsets[0] = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n halo_ids_size += this->cached_L2H_maps[i].size();\n halo_ids_offsets[i + 1] = halo_ids_size;\n }\n\n halo_ids.resize(halo_ids_size);\n // Do exchange with neighbors\n // Pack new bdy_ids\n IVector_h_vector bdy_ids;\n bdy_ids.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = this->cached_B2L_maps[i].size();\n bdy_ids[i].resize(size);\n\n // Pack buffer\n for (int j = 0; j < size; j++)\n {\n bdy_ids[i][j] = row_ids[this->cached_B2L_maps[i][j]];\n }\n }\n\n for (int i = 0; i < num_neighbors; i++)\n {\n comms->send_vector_async(bdy_ids[i], this->neighbors[i], 6666 + this->neighbors[i]);\n }\n\n for (int i = 0; i < num_neighbors; i++)\n {\n comms->recv_vector(halo_ids, this->neighbors[i], 6666 + my_id, halo_ids_offsets[i], this->cached_L2H_maps[i].size());\n }\n\n for (int i = 0; i < num_neighbors; i++)\n {\n comms->send_vector_wait_all(bdy_ids[i]);\n }\n\n IVector_h halo_offsets(root_num_cons_neighbors + 1, 0);\n int root_num_rows;\n this->consolidateAndRenumberHalos(halo_ids, halo_ids_offsets, halo_offsets, this->neighbors, num_neighbors, root_cons_neighbors, root_num_cons_neighbors, destination_part, my_destination_part, is_root_partition, parts_to_consolidate, num_parts_to_consolidate, num_parts, my_id, total_rows_in_merged, root_num_rows, comms);\n\n if (is_root_partition)\n {\n this->B2L_maps.resize(cons_B2L_maps.size());\n\n for (int i = 0; i < cons_B2L_maps.size(); i++)\n {\n amgx::thrust::sort(cons_B2L_maps[i].begin(), cons_B2L_maps[i].end());\n this->B2L_maps[i].copy(cons_B2L_maps[i]); // H2D copy of B2L maps\n }\n\n cudaCheckError();\n }\n\n // Now renumber the row_ids based on lth_maps\n count = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n for (int j = 0; j < this->cached_L2H_maps[i].size(); j++)\n {\n row_ids[this->cached_L2H_maps[i][j]] = halo_ids[count];\n count++;\n }\n }\n\n hipMemcpy(this->m_row_ids_CONS.raw(), row_ids.raw(), (n + total_num_halos)*sizeof(int), hipMemcpyDefault);\n cudaCheckError();\n int bsize = block_dimx * block_dimy;\n\n if (is_root_partition)\n {\n this->A->row_offsets.resize(root_num_rows + 1);\n }\n\n void *root_row_ptr = (void *) this->A->row_offsets.raw();\n\n if (useCudaIpc)\n {\n // ----------------------------------------------------\n // 1. cudaIPC to get pointer to root's row_offset array\n // ----------------------------------------------------\n this->ipcExchangePtr(root_row_ptr, is_root_partition, num_parts_to_consolidate, parts_to_consolidate, my_destination_part, my_id, comms);\n cudaCheckError();\n // -------------------------------------------------------------------\n // 2. each rank copy it's row length on root partition using row_ids\n // -------------------------------------------------------------------\n int cta_size = 128;\n int grid_size = std::min(4096, (n + total_num_halos + cta_size - 1) / cta_size);\n zero_copy_row_lengths_ids_offsets <<< grid_size, cta_size>>>(this->m_old_row_offsets_CONS.raw(), ((int *) root_row_ptr) /* IPC */, this->m_row_ids_CONS.raw(), n, total_num_halos, (mat_value_type *) diag);\n cudaCheckError();\n // Root partition waits for children to be done writing their result\n this->ipcWaitForChildren(is_root_partition, num_parts_to_consolidate, parts_to_consolidate, my_destination_part, my_id, comms);\n cudaCheckError();\n }\n else // CudaIpcNotAvailable\n {\n this->checkPinnedBuffer( std::max( nnz * sizeof(mat_value_type), (n + 1)*std::max(sizeof(index_type), sizeof(value_type)) ) );\n\n if (!is_root_partition)\n {\n IVector_h data_to_send(3);\n data_to_send[0] = n;\n data_to_send[1] = nnz;\n data_to_send[2] = total_num_halos;\n int dummy;\n void *row_ptrs_to_send = this->getHostPointerForData((void *)row_ptrs, sizeof(index_type) * (n + 1), &dummy);\n comms->send_vector(data_to_send, my_destination_part, 10000 + my_id);\n comms->send_raw_data(row_ptrs_to_send, (n + 1)*sizeof(int), my_destination_part, 10001 + my_id);\n comms->send_raw_data(&row_ids[0], (n + total_num_halos)*sizeof(int), my_destination_part, 10002 + my_id);\n }\n else\n {\n hipEvent_t event;\n hipEventCreate(&event);\n cudaCheckError();\n //TODO: Could use streams here\n //TODO: Avoid extra device to host copies\n std::vector data_recv(num_parts_to_consolidate);\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n data_recv[i].resize(3);\n int current_part = parts_to_consolidate[i];\n\n if (current_part != my_id)\n {\n comms->recv_vector(data_recv[i], current_part, 10000 + current_part);\n }\n else\n {\n data_recv[i][0] = n;\n data_recv[i][1] = nnz;\n data_recv[i][2] = total_num_halos;\n }\n }\n\n this->m_child_n.resize(num_parts_to_consolidate);\n this->m_child_nnz.resize(num_parts_to_consolidate);\n this->m_child_num_halos.resize(num_parts_to_consolidate);\n this->m_child_row_ids.resize(num_parts_to_consolidate);\n this->m_child_old_row_offsets.resize(num_parts_to_consolidate);\n int max_n = 0;\n int max_nnz = 0;\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n int current_part = parts_to_consolidate[i];\n this->m_child_n[i] = data_recv[i][0];\n this->m_child_nnz[i] = data_recv[i][1];\n this->m_child_num_halos[i] = data_recv[i][2];\n\n if (this->m_child_n[i] > max_n) { max_n = this->m_child_n[i]; }\n\n if (this->m_child_nnz[i] > max_nnz) { max_nnz = this->m_child_nnz[i]; }\n\n this->m_child_row_ids[i].resize(this->m_child_n[i] + this->m_child_num_halos[i]);\n this->m_child_old_row_offsets[i].resize(this->m_child_n[i] + 1);\n }\n\n this->m_child_max_n = max_n;\n this->m_child_max_nnz = max_nnz;\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n int current_part = parts_to_consolidate[i];\n int cta_size = 128;\n int grid_size = std::min(4096, (this->m_child_n[i] + this->m_child_num_halos[i] + cta_size - 1) / cta_size);\n\n if (current_part != my_id)\n {\n comms->recv_vector(this->m_child_old_row_offsets[i], current_part, 10001 + current_part, 0, this->m_child_n[i] + 1);\n comms->recv_vector(this->m_child_row_ids[i], current_part, 10002 + current_part, 0, this->m_child_n[i] + this->m_child_num_halos[i]);\n zero_copy_row_lengths_ids_offsets <<< grid_size, cta_size>>>(this->m_child_old_row_offsets[i].raw(), this->A->row_offsets.raw(), this->m_child_row_ids[i].raw(), this->m_child_n[i], this->m_child_num_halos[i], (mat_value_type *) diag);\n cudaCheckError();\n // Wait for kernel to finish before overwriting host buffer\n hipEventRecord(event);\n cudaCheckError();\n hipEventSynchronize(event);\n cudaCheckError();\n }\n else\n {\n zero_copy_row_lengths_ids_offsets <<< grid_size, cta_size>>>(this->m_old_row_offsets_CONS.raw(), this->A->row_offsets.raw(), this->m_row_ids_CONS.raw(), n, total_num_halos, (mat_value_type *) diag);\n cudaCheckError();\n hipEventRecord(event);\n cudaCheckError();\n hipEventSynchronize(event);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n hipEventDestroy(event);\n cudaCheckError();\n } // If root partition\n\n //TODO: is this necessary\n comms->barrier();\n }\n\n //3. root does a exclusive_scan\n if (is_root_partition)\n {\n hipEvent_t event;\n hipEventCreate(&event);\n cudaCheckError();\n // Populate the halo rows with diagonal, increase the length of the halo rows\n thrust_wrapper::fill(this->A->row_offsets.begin() + halo_offsets[0], this->A->row_offsets.begin() + halo_offsets[root_num_cons_neighbors], 1);\n thrust_wrapper::exclusive_scan(this->A->row_offsets.begin(), this->A->row_offsets.end(), this->A->row_offsets.begin());\n hipEventRecord(event);\n cudaCheckError();\n hipEventSynchronize(event);\n cudaCheckError();\n this->A->set_initialized(0);\n this->A->delProps(DIAG); // We always insert the diagonal\n this->A->delProps(COO); // No COO\n this->A->setColsReorderedByColor(false); // Cols not reordered by color\n int nnz = this->A->row_offsets[root_num_rows]; // This is a device to host copy\n this->A->resize(root_num_rows, root_num_rows, nnz, block_dimx, block_dimy);\n this->A->set_num_nz(nnz); // num_nz doesn't include halo rows\n //this->A->set_initialized(1);\n hipEventDestroy(event);\n cudaCheckError();\n }\n else\n {\n this->A->set_initialized(0);\n this->A->resize( 0, 0, 0, block_dimx, block_dimy );\n this->A->delProps(DIAG); // We always insert the diagonal\n this->A->delProps(COO); // No COO\n this->A->setColsReorderedByColor(false); // Cols not reordered by color\n //this->A->set_initialized(1);\n }\n\n if (useCudaIpc)\n {\n // ----------------------------------------------\n // 4. Do ipc consolidation of values and columns\n // ----------------------------------------------\n // Child partition waits for parent to create row_offsets\n this->ipcWaitForRoot(is_root_partition, num_parts_to_consolidate, parts_to_consolidate, my_destination_part, my_id, comms);\n void *root_col_ptr = (void *) this->A->col_indices.raw();\n void *root_val_ptr = (void *) this->A->values.raw();\n this->ipcExchangePtr(root_col_ptr, is_root_partition, num_parts_to_consolidate, parts_to_consolidate, my_destination_part, my_id, comms);\n this->ipcExchangePtr(root_val_ptr, is_root_partition, num_parts_to_consolidate, parts_to_consolidate, my_destination_part, my_id, comms);\n int cta_size2 = 128;\n int grid_size2 = std::min(4096, (n + cta_size2 - 1) / cta_size2);\n ipc_consolidation_upload_matrix <<< grid_size2, cta_size2>>>(n, this->m_row_ids_CONS.raw(), this->m_old_row_offsets_CONS.raw(), ( (int *) root_row_ptr ) /*IPC*/, col_indices_hd, ( (int *) root_col_ptr) /*IPC*/, data_hd, ( (mat_value_type *) root_val_ptr ) /*IPC*/, diag_hd, bsize);\n cudaCheckError();\n // Root partition waits for children to upload their matrices\n this->ipcWaitForChildren(is_root_partition, num_parts_to_consolidate, parts_to_consolidate, my_destination_part, my_id, comms);\n cudaCheckError();\n\n // Child partitions close their mem handle (they are done upload data)\n if (!is_root_partition)\n {\n hipIpcCloseMemHandle(root_row_ptr);\n hipIpcCloseMemHandle(root_val_ptr);\n hipIpcCloseMemHandle(root_col_ptr);\n }\n }\n else // If cudaIpcNotAvailable\n {\n if (!is_root_partition)\n {\n int dummy;\n void *col_indices_to_send = this->getHostPointerForData((void *)col_indices, sizeof(index_type) * nnz, &dummy);\n comms->send_raw_data(col_indices_to_send, nnz * sizeof(int), my_destination_part, 10000 + my_id);\n void *data_to_send = this->getHostPointerForData((void *)data, sizeof(mat_value_type) * nnz, &dummy);\n comms->send_raw_data(data_to_send, nnz * bsize * sizeof(mat_value_type), my_destination_part, 10001 + my_id);\n\n if (diag != NULL)\n {\n void *diag_to_send = this->getHostPointerForData((void *)diag, sizeof(mat_value_type) * n, &dummy);\n comms->send_raw_data(diag_to_send, n * bsize * sizeof(mat_value_type), my_destination_part, 10002 + my_id);\n }\n }\n else\n {\n hipEvent_t event;\n hipEventCreate(&event);\n cudaCheckError();\n //TODO: Could use streams here\n int *child_col_indices;\n mat_value_type *child_data;\n mat_value_type *child_diag = NULL;\n hipHostAlloc( (void **) &child_col_indices, this->m_child_max_nnz * sizeof(int), hipHostMallocMapped);\n hipHostAlloc( (void **) &child_data, this->m_child_max_nnz * bsize * sizeof(mat_value_type), hipHostMallocMapped);\n\n if (diag != NULL)\n {\n hipHostAlloc( (void **) &child_diag, (this->m_child_max_n)*bsize * sizeof(mat_value_type), hipHostMallocMapped);\n }\n\n for (int i = 0; i < num_parts_to_consolidate; i++)\n {\n int current_part = parts_to_consolidate[i];\n int cta_size2 = 128;\n int grid_size2 = std::min(4096, (this->m_child_n[i] + cta_size2 - 1) / cta_size2);\n\n if (current_part != my_id)\n {\n comms->recv_raw_data(child_col_indices, this->m_child_nnz[i]*sizeof(int), current_part, 10000 + current_part);\n comms->recv_raw_data(child_data, this->m_child_nnz[i]*bsize * sizeof(mat_value_type), current_part, 10001 + current_part);\n\n if (diag != NULL)\n {\n comms->recv_raw_data(child_diag, this->m_child_n[i]*bsize * sizeof(mat_value_type), current_part, 10002 + current_part);\n }\n\n int *child_col_indices_hd;\n mat_value_type *child_data_hd;\n mat_value_type *child_diag_hd = NULL;\n hipHostGetDevicePointer(&child_col_indices_hd, child_col_indices, 0);\n hipHostGetDevicePointer(&child_data_hd, child_data, 0);\n\n if (diag != NULL)\n {\n hipHostGetDevicePointer(&child_diag_hd, child_diag, 0);\n }\n\n ipc_consolidation_upload_matrix <<< grid_size2, cta_size2>>>(this->m_child_n[i], this->m_child_row_ids[i].raw(), this->m_child_old_row_offsets[i].raw(), this->A->row_offsets.raw(), child_col_indices_hd, this->A->col_indices.raw(), child_data_hd, this->A->values.raw(), child_diag_hd, bsize);\n cudaCheckError();\n // Wait for kernel to finish before overwriting host buffer\n hipEventRecord(event);\n cudaCheckError();\n hipEventSynchronize(event);\n cudaCheckError();\n }\n else\n {\n ipc_consolidation_upload_matrix <<< grid_size2, cta_size2>>>(n, this->m_row_ids_CONS.raw(), this->m_old_row_offsets_CONS.raw(), this->A->row_offsets.raw(), col_indices_hd, this->A->col_indices.raw(), data_hd, this->A->values.raw(), diag_hd, bsize);\n cudaCheckError();\n hipEventRecord(event);\n cudaCheckError();\n hipEventSynchronize(event);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n hipEventDestroy(event);\n cudaCheckError();\n amgx::memory::hipHostFree(child_col_indices);\n amgx::memory::hipHostFree(child_data);\n\n if (diag != NULL)\n {\n amgx::memory::hipHostFree(child_diag);\n }\n } // If root partition\n\n //TODO: is this necessary\n comms->barrier();\n }\n\n // Store the original fine level communicator\n this->m_is_fine_level_root_partition = is_root_partition;\n this->m_my_fine_level_destination_part = my_destination_part;\n // Create a clone of the original communicator\n this->m_fine_level_comms = comms; //this->_comms is the same pointer that this->m_fine_level_comms right now, so we can overwrite this->_comms, but make sure that we release m_fine_level_cons\n this->_comms = this->m_fine_level_comms->CloneSubComm(cons_part_to_part, is_root_partition); // this->_comms will be empty comm for non-root partition and new comm for root ranks only if root partition\n this->m_fine_level_id = my_id;\n\n if (is_root_partition)\n {\n int cta_size = 128;\n int grid_size3 = std::min(4096, ( (root_num_rows - halo_offsets[0]) + cta_size - 1) / cta_size);\n\n if (grid_size3 != 0)\n {\n set_halo_cols_values <<< grid_size3, cta_size>>>(this->A->row_offsets.raw(), this->A->col_indices.raw(), this->A->values.raw(), halo_offsets[0], root_num_rows, bsize);\n cudaCheckError();\n }\n\n int my_cons_id = part_to_cons_part[my_id];\n this->_global_id = my_cons_id;\n this->_num_interior_nodes = total_interior_rows_in_merged;\n this->_num_boundary_nodes = total_boundary_rows_in_merged;\n\n for (int i = 0; i < root_num_cons_neighbors; i++)\n {\n root_cons_neighbors[i] = part_to_cons_part[root_cons_neighbors[i]];\n }\n\n this->_comms->set_neighbors(root_num_cons_neighbors);\n this->neighbors = root_cons_neighbors;\n this->halo_offsets = halo_offsets; // H2D copy of halo offsets\n this->m_num_fine_level_parts_to_consolidate = num_parts_to_consolidate;\n this->m_fine_level_parts_to_consolidate = parts_to_consolidate;\n this->set_num_halo_rings(num_rings);\n this->set_num_halo_rows(halo_offsets[root_num_cons_neighbors] - halo_offsets[0]);\n // B2L_maps has already been copied\n this->B2L_rings.resize(root_num_cons_neighbors);\n\n for (int i = 0; i < root_num_cons_neighbors; i++)\n {\n this->B2L_rings[i].resize(2);\n this->B2L_rings[i][0] = 0;\n this->B2L_rings[i][1] = cons_B2L_maps[i].size();\n }\n\n this->set_initialized(this->A->row_offsets);\n this->A->set_initialized(0);\n this->A->delProps(DIAG);\n this->A->diag.resize(root_num_rows);\n this->A->computeDiagonal(); //\n this->A->setView(OWNED);\n hipEventCreate(&(this->comm_event));\n cudaCheckError();\n this->A->set_initialized(1);\n }\n else\n {\n this->neighbors.resize(0);\n this->halo_offsets.resize(0);\n }\n\n /* free memory (if needed) */\n if (col_alloc) { amgx::memory::hipFreeAsync(col_indices_hd); }\n if (data_alloc) { amgx::memory::hipFreeAsync(data_hd); }\n if (diag_alloc) { amgx::memory::hipFreeAsync(diag_hd); }\n}\n\ntemplate \nvoid DistributedManager >::replaceMatrixCoefficientsNoCons(int n, int nnz, const mat_value_type *data_pinned, const mat_value_type *diag_pinned)\n{\n //matrix parameters\n //int num_nnz = this->A->get_num_nz();\n int num_rows = this->halo_offsets[0];\n int total_rows = num_rows + this->num_halo_rows();\n int block_size = this->A->get_block_size();\n mat_value_type *data_hd = NULL;\n mat_value_type *diag_hd = NULL;\n int data_alloc = 0;\n int diag_alloc = 0;\n //cuda parameters\n int num_blocks = std::min(4096, (num_rows + 127) / 128);\n\n /* WARNING: the number of non-zero elements (nnz) in the array data_pinned and A->values (num_nnz) might be different at this point.\n 1. If the matrix has CSR property and therefore diagonal is included in the matrix this values will be the same.\n 2. If the matrix has DIAG property and therefore diagonal is originally stored separately, and later appended to the array\n of values, and subsequently inserted into the matrix than num_nnz = nnz + n. We have to account for this fact when replacing the\n coefficients (and use nnz not num_nnz).\n obs.: see calls to computeDiagonal (matrix.cu), AMGX_matrix_upload and AMGX_replace_coefficients (amgx_c.cu), and\n uploadMatrix and replaceMatrixCoefficients[No|With]Cons (distributed_manager.cu) for details. */\n\n /* check early exit */\n if ((this->neighbors.size() == 0 || this->renumbering.size() == 0) && !this->m_is_fine_level_glued)\n {\n return;\n }\n\n cudaCheckError();\n /* allocate if data and diag if they are not pinned */\n data_hd = (mat_value_type *) this->getDevicePointerForData((void *)data_pinned, nnz * block_size * sizeof(mat_value_type), &data_alloc);\n\n if (diag_pinned != NULL)\n {\n diag_hd = (mat_value_type *) this->getDevicePointerForData((void *)diag_pinned, num_rows * block_size * sizeof(mat_value_type), &diag_alloc);\n }\n\n /* replace the values (reordering them if needed) */\n if (insertDiagonals && diag_pinned != NULL)\n {\n replace_values_matrix <32> <<< num_blocks, 512>>>(data_hd, diag_hd, this->old_row_offsets.raw(), this->A->row_offsets.raw(), this->A->values.raw(), this->renumbering.raw(), block_size, num_rows);\n cudaCheckError();\n }\n else\n {\n replace_values_matrix <32> <<< num_blocks, 512>>>(data_hd, this->old_row_offsets.raw(), this->A->row_offsets.raw(), this->A->values.raw(), this->renumbering.raw(), block_size, num_rows);\n cudaCheckError();\n\n if (diag_pinned != NULL)\n {\n reorder_vector_values <<< num_blocks, 512>>>(this->A->values.raw() + this->A->row_offsets[total_rows]*block_size, diag_hd, this->renumbering.raw(), block_size, num_rows);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n\n /* free memory (if needed) */\n if (data_alloc) { amgx::memory::hipFreeAsync(data_hd); }\n if (diag_alloc) { amgx::memory::hipFreeAsync(diag_hd); }\n}\n\ntemplate \nvoid DistributedManager >::replaceMatrixCoefficientsWithCons(int n, int nnz, const mat_value_type *data_pinned, const mat_value_type *diag_pinned)\n{\n //matrix parameters\n //int num_nnz = this->A->get_num_nz();\n /* WARNING: in consolidation, for non-root partitions, halo_offsets\n might be NULL due to the call halo_offsets.resize(0); at the end\n of the routine uploadMatrix->consolidateAndUploadAll. We should\n use the parameter n instead this->halo_offsets[0] for num_rows. */\n int num_rows = n;\n int block_size = this->A->get_block_size();\n mat_value_type *data_hd = NULL;\n mat_value_type *diag_hd = NULL;\n int data_alloc = 0;\n int diag_alloc = 0;\n data_hd = (mat_value_type *) this->getDevicePointerForData((void *)data_pinned, nnz * block_size * sizeof(mat_value_type), &data_alloc);\n\n if (diag_pinned != NULL)\n {\n diag_hd = (mat_value_type *) this->getDevicePointerForData((void *)diag_pinned, num_rows * block_size * sizeof(mat_value_type), &diag_alloc);\n }\n\n bool useCudaIpc = this->m_use_cuda_ipc_consolidation;\n\n if (useCudaIpc)\n {\n // Child partitions wait for root to be done\n this->ipcWaitForRoot(this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n cudaCheckError();\n void *root_row_ptr = (void *) this->A->row_offsets.raw();\n void *root_val_ptr = (void *) this->A->values.raw();\n this->ipcExchangePtr(root_row_ptr, this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n this->ipcExchangePtr(root_val_ptr, this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n // replace the values, insert the diagonal\n int ncons = this->m_old_row_offsets_CONS.size() - 1;\n int cta_size = 128;\n int grid_size2 = std::min(4096, (ncons + cta_size - 1) / cta_size);\n ipc_consolidation_replace_values <<< grid_size2, cta_size>>>(ncons, this->m_row_ids_CONS.raw(), this->m_old_row_offsets_CONS.raw(), ( (int *) root_row_ptr )/*IPC*/, data_hd, ( (mat_value_type *) root_val_ptr )/*IPC*/, diag_hd, this->A->get_block_size() );\n cudaCheckError();\n // Root partition wait for child to be done replacing their values\n this->ipcWaitForChildren(this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n cudaCheckError();\n\n if (!this->m_is_fine_level_root_partition)\n {\n hipIpcCloseMemHandle(root_row_ptr);\n hipIpcCloseMemHandle(root_val_ptr);\n }\n }\n else // cudaIpcNotAvailable\n {\n if (this->m_is_fine_level_consolidated) // aggregation\n {\n int bsize = this->A->get_block_size();\n int ncons = this->m_old_row_offsets_CONS.size() - 1;\n\n if (!this->m_is_fine_level_root_partition)\n {\n int dummy;\n int nnzcons = this->m_old_nnz_CONS;\n void *data_to_send = this->getHostPointerForData((void *)data_pinned, nnzcons * bsize * sizeof(mat_value_type), &dummy);\n this->getFineLevelComms()->send_raw_data(data_to_send, nnzcons * bsize * sizeof(mat_value_type), this->m_my_fine_level_destination_part, 10001 + this->fine_level_id());\n\n if (diag_pinned != NULL)\n {\n void *diag_to_send = this->getHostPointerForData((void *)diag_pinned, ncons * bsize * sizeof(mat_value_type), &dummy);\n this->getFineLevelComms()->send_raw_data(diag_to_send, ncons * bsize * sizeof(mat_value_type), this->m_my_fine_level_destination_part, 10002 + this->fine_level_id());\n }\n }\n else\n {\n hipEvent_t event;\n hipEventCreate(&event);\n cudaCheckError();\n //TODO: Could use streams here\n mat_value_type *child_data;\n mat_value_type *child_diag = NULL;\n hipHostAlloc( (void **) &child_data, this->m_child_max_nnz * bsize * sizeof(mat_value_type), hipHostMallocMapped);\n\n if (diag_pinned != NULL)\n {\n hipHostAlloc( (void **) &child_diag, (this->m_child_max_n)*bsize * sizeof(mat_value_type), hipHostMallocMapped);\n }\n\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n int cta_size2 = 128;\n int grid_size2 = std::min(4096, (this->m_child_n[i] + cta_size2 - 1) / cta_size2);\n\n if (current_part != this->fine_level_id())\n {\n this->getFineLevelComms()->recv_raw_data(child_data, this->m_child_nnz[i]*bsize * sizeof(mat_value_type), current_part, 10001 + current_part);\n\n if (diag_pinned != NULL)\n {\n this->getFineLevelComms()->recv_raw_data(child_diag, this->m_child_n[i]*bsize * sizeof(mat_value_type), current_part, 10002 + current_part);\n }\n\n mat_value_type *child_data_hd;\n mat_value_type *child_diag_hd = NULL;\n hipHostGetDevicePointer(&child_data_hd, child_data, 0);\n\n if (diag_pinned != NULL)\n {\n hipHostGetDevicePointer(&child_diag_hd, child_diag, 0);\n }\n\n ipc_consolidation_replace_values <<< grid_size2, cta_size2>>>(this->m_child_n[i], this->m_child_row_ids[i].raw(), this->m_child_old_row_offsets[i].raw(), this->A->row_offsets.raw(), child_data_hd, this->A->values.raw(), child_diag_hd, bsize);\n cudaCheckError();\n // Wait for kernel to finish before overwriting host buffer\n hipEventRecord(event);\n cudaCheckError();\n hipEventSynchronize(event);\n cudaCheckError();\n }\n else\n {\n ipc_consolidation_replace_values <<< grid_size2, cta_size2>>>(ncons, this->m_row_ids_CONS.raw(), this->m_old_row_offsets_CONS.raw(), this->A->row_offsets.raw(), data_hd, this->A->values.raw(), diag_hd, bsize);\n cudaCheckError();\n //cudaEventRecord(event);\n cudaCheckError();\n //cudaEventSynchronize(event);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n hipEventDestroy(event);\n cudaCheckError();\n amgx::memory::hipHostFree(child_data);\n\n if (diag_pinned != NULL)\n {\n amgx::memory::hipHostFree(child_diag);\n }\n } // If root partition\n\n //TODO: is this necessary\n this->getFineLevelComms()->barrier();\n } //agg\n else if (this->m_is_fine_level_glued) // classical\n {\n int bsize = this->A->get_block_size();\n int ncons = this->m_old_row_offsets_CONS.size() - 1;\n IVector_h nnz_off;\n nnz_off.resize(this->getConsolidationArrayOffsets().size());\n IVector_h nnz_array;\n nnz_array.resize(this->getConsolidationArrayOffsets().size() - 1);\n this->getFineLevelComms()->all_gather( nnz,\n nnz_array,\n this->getConsolidationArrayOffsets().size() - 1);\n nnz_off[0] = 0;\n\n for (int i = 0; i < nnz_array.size(); i++)\n {\n nnz_off[i + 1] = nnz_off[i] + nnz_array[i];\n }\n\n if (!this->m_is_fine_level_root_partition)\n {\n int dummy;\n void *data_to_send = this->getHostPointerForData((void *)data_pinned, nnz * bsize * sizeof(mat_value_type), &dummy);\n this->getFineLevelComms()->send_raw_data( data_to_send,\n nnz * bsize * sizeof(mat_value_type),\n this->m_my_fine_level_destination_part,\n 10001 + this->fine_level_id());\n\n if (diag_pinned != NULL)\n {\n void *diag_to_send = this->getHostPointerForData((void *)diag_pinned, n * sizeof(mat_value_type), &dummy);\n this->getFineLevelComms()->send_raw_data( diag_to_send,\n n * bsize * sizeof(mat_value_type),\n this->m_my_fine_level_destination_part,\n 10002 + this->fine_level_id());\n //diag.resize(0);\n cudaCheckError();\n }\n\n //values.resize(0);\n cudaCheckError();\n }\n else\n {\n //TODO: Could use streams here\n mat_value_type *child_data;\n mat_value_type *child_diag = NULL;\n // Assumes partions have been glued already\n this->A->getNnzForView(OWNED, &nnz);\n hipHostAlloc( (void **) &child_data, nnz * bsize * sizeof(mat_value_type), hipHostMallocMapped);\n\n if (diag_pinned != NULL)\n {\n hipHostAlloc( (void **) &child_diag, this->halo_offsets[this->neighbors.size()]*bsize * sizeof(mat_value_type), hipHostMallocMapped);\n }\n\n // roots copy their data\n memcpy ( &child_data[0], data_pinned, nnz_array[this->fine_level_id()]*sizeof(value_type));\n\n if (diag_pinned != NULL)\n {\n memcpy ( &child_diag[0], diag_pinned, n * sizeof(value_type));\n }\n\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n int current_offset = nnz_off[current_part] - nnz_off[this->fine_level_id()] ;\n int current_nnz = nnz_array[current_part];\n\n if (current_part != this->fine_level_id())\n {\n this->getFineLevelComms()->recv_raw_data( &child_data[current_offset],\n current_nnz * bsize * sizeof(mat_value_type),\n current_part,\n 10001 + current_part);\n\n if (diag_pinned != NULL)\n this->getFineLevelComms()->recv_raw_data( &child_diag[this->getConsolidationArrayOffsets()[current_part] - this->getConsolidationArrayOffsets()[this->fine_level_id()]],\n (this->getConsolidationArrayOffsets()[current_part + 1] - this->getConsolidationArrayOffsets()[current_part])*bsize * sizeof(mat_value_type),\n current_part,\n 10002 + current_part);\n }\n }\n\n cudaCheckError();\n // we can follow the usual upload path for raw data now\n // Assumes partions have been glued already\n int os;\n this->A->getOffsetAndSizeForView(OWNED, &os, &n);\n replaceMatrixCoefficientsNoCons( n, nnz, child_data, child_diag);\n cudaCheckError();\n amgx::memory::hipHostFree(child_data);\n\n if (diag_pinned != NULL)\n {\n amgx::memory::hipHostFree(child_diag);\n }\n } // If root partition\n\n //TODO: is this necessary\n this->getFineLevelComms()->barrier();\n } // cla\n } // not ipc\n\n this->A->setView(OWNED);\n\n /* free memory (if needed) */\n if (data_alloc) { amgx::memory::hipFreeAsync(data_hd); }\n\n if (diag_alloc) { amgx::memory::hipFreeAsync(diag_hd); }\n}\n\ntemplate \nvoid DistributedManager >::transformAndUploadVector(VVector_v &v, const void *data, int n, int block_dim)\n{\n if (this->isFineLevelConsolidated() || (this->isFineLevelGlued() && !this->isGlued()))\n {\n transformAndUploadVectorWithCons(v, data, n, block_dim);\n }\n else\n {\n v.resize(n * block_dim);\n cudaCheckError();\n // Upload on host\n hipMemcpy(v.raw(), (value_type *)data, n * block_dim * sizeof(value_type), hipMemcpyDefault);\n cudaCheckError();\n // Permute based on renumbering vector\n transformVector(v);\n int tag = 0;\n // Exchange halos\n this->exchange_halo(v, tag);\n }\n}\n\ntemplate \nvoid DistributedManager >::transformAndUploadVectorWithCons(VVector_v &v, const void *data_pinned, int n, int block_dim)\n{\n if (v.get_block_size() != this->A->get_block_dimx()) { printf(\"Blocksize mismatch!\\n\"); }\n\n bool useCudaIpc = this->m_use_cuda_ipc_consolidation;\n this->getFineLevelComms()->barrier();\n void *root_temp_ptr = NULL;\n VVector_v temp;\n\n if (this->m_is_fine_level_root_partition && !this->m_is_fine_level_glued )\n {\n temp.resize(this->halo_offsets[this->neighbors.size()]*v.get_block_size(), types::util::get_zero());\n temp.set_block_dimx(v.get_block_dimx());\n temp.set_block_dimy(v.get_block_dimy());\n root_temp_ptr = (void *) temp.raw();\n }\n\n cudaCheckError();\n int data_alloc = 0;\n value_type *data_hd = NULL;\n\n if (!this->m_is_fine_level_glued )\n {\n data_hd = (value_type *) this->getDevicePointerForData((void *)data_pinned, n * block_dim * sizeof(value_type), &data_alloc);\n }\n\n if (useCudaIpc)\n {\n // Do IPC\n this->ipcExchangePtr(root_temp_ptr, this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n cudaCheckError();\n int num_blocks = std::min(4096, (n + 511) / 512);\n reorder_vector_values <<< num_blocks, 512>>>( (value_type *) root_temp_ptr, data_hd, this->m_row_ids_CONS.raw(), v.get_block_size(), n);\n cudaCheckError();\n // Root partition waits for children to be done\n this->ipcWaitForChildren(this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n cudaCheckError();\n\n if (!this->m_is_fine_level_root_partition)\n {\n hipIpcCloseMemHandle(root_temp_ptr);\n }\n }\n else // If cudaIpcNotAvail\n {\n if (this->m_is_fine_level_consolidated) // aggregation\n {\n // Exchange the vector between root and child\n if (!this->m_is_fine_level_root_partition)\n {\n IVector_h size(1);\n size[0] = n;\n this->getFineLevelComms()->send_vector(size, this->m_my_fine_level_destination_part, 20000 + this->fine_level_id());\n int dummy;\n void *data_to_send = this->getHostPointerForData((void *)data_pinned, n * v.get_block_size() * sizeof(value_type), &dummy);\n this->getFineLevelComms()->send_raw_data(data_to_send, n * v.get_block_size()*sizeof(value_type), this->m_my_fine_level_destination_part, 20001 + this->fine_level_id());\n }\n else\n {\n hipEvent_t event;\n hipEventCreate(&event);\n cudaCheckError();\n IVector_h child_n(this->m_num_fine_level_parts_to_consolidate);\n int max_n = 0;\n\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n\n if (current_part != this->fine_level_id())\n {\n this->getFineLevelComms()->recv_vector(child_n, current_part, 20000 + current_part, i, 1);\n }\n else\n {\n child_n[i] = n;\n }\n\n if (child_n[i] > max_n) { max_n = child_n[i]; }\n }\n\n value_type *child_data;\n hipHostAlloc( (void **) &child_data, max_n * v.get_block_size()*sizeof(value_type), hipHostMallocMapped);\n value_type *child_data_hd;\n hipHostGetDevicePointer(&child_data_hd, child_data, 0);\n\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n int num_blocks = std::min(4096, (child_n[i] + 511) / 512);\n\n if (current_part != this->fine_level_id())\n {\n this->getFineLevelComms()->recv_raw_data(&child_data[0], child_n[i]*v.get_block_size()*sizeof(value_type), current_part, 20001 + current_part);\n reorder_vector_values <<< num_blocks, 512>>>( (value_type *) root_temp_ptr, child_data_hd, this->m_child_row_ids[i].raw(), v.get_block_size(), child_n[i]);\n cudaCheckError();\n hipEventRecord(event);\n cudaCheckError();\n hipEventSynchronize(event);\n cudaCheckError();\n }\n else\n {\n reorder_vector_values <<< num_blocks, 512>>>( (value_type *) root_temp_ptr, data_hd, this->m_row_ids_CONS.raw(), v.get_block_size(), n);\n cudaCheckError();\n }\n } // Loop over parts to consolidate\n\n cudaCheckError();\n hipEventDestroy(event);\n cudaCheckError();\n amgx::memory::hipHostFree(child_data);\n } // If root partition\n } //agg\n else if (this->m_is_fine_level_glued) // cla\n {\n value_type *child_data = NULL;\n\n if (!this->m_is_fine_level_root_partition)\n {\n int dummy;\n void *data_to_send = this->getHostPointerForData((void *)data_pinned, n * v.get_block_size() * sizeof(value_type), &dummy);\n this->getFineLevelComms()->send_raw_data( data_to_send,\n n * v.get_block_size()*sizeof(value_type),\n this->m_my_fine_level_destination_part,\n 20001 + this->fine_level_id());\n //v.resize(0); // just in case something resized it betwen iterations\n cudaCheckError();\n }\n else\n {\n hipHostAlloc( (void **) &child_data, this->halo_offsets[this->neighbors.size()]*v.get_block_size()*sizeof(value_type), hipHostMallocMapped);\n value_type *child_data_hd;\n hipHostGetDevicePointer(&child_data_hd, child_data, 0);\n // roots copy their data\n int dummy;\n void *my_data = this->getHostPointerForData((void *)data_pinned, n * v.get_block_size() * sizeof(value_type), &dummy);\n memcpy ( &child_data[0], data_pinned, n * v.get_block_size()*sizeof(value_type));\n\n // Loop over parts to consolidate\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n\n if (current_part != this->fine_level_id())\n {\n this->getFineLevelComms()->recv_raw_data( &child_data[this->getConsolidationArrayOffsets()[current_part] - this->getConsolidationArrayOffsets()[this->fine_level_id()]],\n sizeof(value_type) * (this->getConsolidationArrayOffsets()[current_part + 1] - this->getConsolidationArrayOffsets()[current_part]),\n current_part,\n 20001 + current_part );\n }\n }\n\n // usual path\n // Upload on host\n hipMemcpy(v.raw(), (value_type *)child_data, v.size()* sizeof(value_type), hipMemcpyDefault);\n cudaCheckError();\n } // If root partition\n\n // Permute based on renumbering vector\n transformVector(v);\n cudaCheckError();\n // Exchange halos\n int tag = 0;\n this->exchange_halo(v, tag);\n cudaCheckError();\n v.set_unconsolidated_size(n);\n\n // free host\n if (child_data) { amgx::memory::hipHostFree(child_data); }\n\n cudaCheckError();\n } //cla\n } // If cudaIpcAvailable\n\n if (!this->m_is_fine_level_glued) // not needed for classcical\n {\n if (this->m_is_fine_level_root_partition)\n {\n v.swap(temp);\n int tag = 0;\n // Root partitions do the exchange\n this->exchange_halo(v, tag);\n }\n\n v.set_unconsolidated_size(n * v.get_block_size());\n v.set_transformed();\n }\n\n /* free memory (if needed) */\n if (data_alloc) { amgx::memory::hipFreeAsync(data_hd); }\n\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManager >::transformVector(VVector_v &v)\n{\n if (this->neighbors.size() == 0) { return; }\n else if (this->renumbering.size() == 0)\n {\n v.resize(this->halo_offsets[this->neighbors.size()]*v.get_block_size());\n return;\n }\n\n if (v.get_block_size() != this->A->get_block_dimx()) { printf(\"Blocksize mismatch!\\n\"); }\n\n if (v.size() < this->halo_offsets[this->neighbors.size()]*v.get_block_size())\n {\n VVector_v temp(this->halo_offsets[this->neighbors.size()]*v.get_block_size(), types::util::get_zero());\n temp.set_block_dimx(v.get_block_dimx());\n temp.set_block_dimy(v.get_block_dimy());\n\n if (v.size() < this->halo_offsets[0]*this->A->get_block_dimx())\n {\n FatalError(\"Unknown size of input vector - smaller than the number of rows owned by this partition\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n //reorder based on row permutation\n int size = this->halo_offsets[0];\n int num_blocks = std::min(4096, (size + 511) / 512);\n reorder_vector_values <<< num_blocks, 512>>>(temp.raw(), v.raw(), this->renumbering.raw(), v.get_block_size(), size);\n cudaCheckError();\n v.swap(temp);\n }\n else\n {\n VVector_v temp(this->halo_offsets[0]*v.get_block_size());\n int size = this->halo_offsets[0];\n int num_blocks = std::min(4096, (size + 511) / 512);\n reorder_vector_values <<< num_blocks, 512>>>(temp.raw(), v.raw(), this->renumbering.raw(), v.get_block_size(), size);\n cudaCheckError();\n amgx::thrust::copy(temp.begin(), temp.end(), v.begin());\n }\n\n cudaCheckError();\n v.set_transformed();\n}\n\ntemplate \nvoid DistributedManager >::revertVector(VVector_v &v)\n{\n if (this->neighbors.size() == 0 || this->renumbering.size() == 0) { return; }\n\n if (v.get_block_size() != this->A->get_block_dimx()) { printf(\"Blocksize mismatch!\\n\"); }\n\n VVector_v temp(this->halo_offsets[0]*this->A->get_block_dimx());\n\n if (v.size() < this->halo_offsets[0]*v.get_block_size())\n {\n FatalError(\"Unknown size of input vector - smaller than the number of rows owned by this partition\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n //reorder based on row permutation\n int size = this->halo_offsets[0];\n int num_blocks = std::min(4096, (size + 511) / 512);\n inverse_reorder_vector_values <<< num_blocks, 512>>>(temp.raw(), v.raw(), this->renumbering.raw(), v.get_block_size(), size);\n //reorder_vector_values<<>>(temp.raw(), v.raw(), this->inverse_renumbering.raw(), v.get_block_size(), size);\n cudaCheckError();\n v.resize(this->halo_offsets[0]*this->A->get_block_dimx());\n amgx::thrust::copy(temp.begin(), temp.end(), v.begin());\n cudaCheckError();\n}\n\n\ntemplate \nvoid DistributedManagerBase::computeDestinationPartitions(INDEX_TYPE upper_threshold, float avg_size, const int num_parts, int &new_num_parts, bool &wantNeighbors)\n{\n m_destination_partitions.resize(num_parts);\n std::vector dp(num_parts);\n\n if (avg_size < 1.f) { avg_size = 1.f; } // avoid floating point exception\n\n int wanted_num_fine_parts_to_consolidate = ( upper_threshold + (int) avg_size - 1) / ( (int) avg_size );\n new_num_parts = (num_parts + wanted_num_fine_parts_to_consolidate - 1) / wanted_num_fine_parts_to_consolidate;\n\n for (int i = 0; i < num_parts; i++)\n {\n dp[i] = i % new_num_parts;\n }\n\n // example wantNeighbors = true -> destination_part = [0 0 0 0 4 4 4 4 8 8 8 8]\n // example wantNeighbors = false -> destination_part = [0 1 2 3 0 1 2 3 0 1 2 3]\n if (wantNeighbors)\n {\n std::sort (dp.begin(), dp.end());\n m_destination_partitions[0] = 0;\n\n for (int i = 1; i < num_parts; i++)\n {\n if (dp[i - 1] < dp[i])\n {\n m_destination_partitions[i] = i;\n }\n else\n {\n m_destination_partitions[i] = m_destination_partitions[i - 1];\n }\n }\n }\n\n m_my_destination_part = m_destination_partitions[global_id()];\n}\n\ntemplate \nvoid DistributedManagerBase::computeDestinationPartitionsWithCons(int my_id, int num_parts, IVector_h &destination_part, DistributedComms *comms)\n{\n int device_id = this->A->getResources()->getDevice(0);\n std::string my_hostname_tmp;\n comms->get_hostname(my_hostname_tmp);\n // Append PCI-E ID to string\n hipDeviceProp_t dev_prop;\n hipGetDeviceProperties(&dev_prop, device_id);\n std::stringstream s;\n s << my_hostname_tmp << \"_\" << dev_prop.pciBusID << \"_\" << dev_prop.pciDeviceID;\n std::string my_hostname(s.str());\n std::vector hostnames;\n comms->exchange_hostnames(my_hostname, hostnames, num_parts);\n std::vector::iterator low = std::find( hostnames.begin(), hostnames.end(), my_hostname );\n int my_destination_part = low - hostnames.begin();\n // Do a gather into destination_part\n comms->all_gather(my_destination_part, destination_part, num_parts);\n}\n\n\ntemplate \nvoid DistributedManager >::revertAndDownloadVector(VVector_v &v_in, const void *data, int n, int block_dimy)\n{\n if (this->isFineLevelConsolidated() || this->isFineLevelGlued())\n {\n revertAndDownloadVectorWithCons(v_in, data, n, block_dimy);\n }\n else\n {\n if ( n == 0 )\n {\n FatalError(\"Cannot download if size = 0\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (data == NULL )\n {\n FatalError(\"Cannot download to a NULL pointer\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (v_in.size() == 0 )\n {\n FatalError(\"Cannot download an empty vector\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n VVector_v v_out;\n revertVector(v_in, v_out);\n hipMemcpy((value_type *)data, v_out.raw(), n * block_dimy * sizeof(value_type), hipMemcpyDefault);\n cudaCheckError();\n }\n}\n\ntemplate \nvoid DistributedManager >::revertVector(VVector_v &v_in, VVector_v &v_out)\n{\n if (!this->isFineLevelGlued() && this->neighbors.size() == 0 || this->renumbering.size() == 0) { return;}\n\n if (v_in.get_block_size() != this->A->get_block_dimx()) { printf(\"Blocksize mismatch!\\n\"); }\n\n if (v_in.size() < this->halo_offsets[0]*v_in.get_block_size())\n {\n FatalError(\"Unknown size of input vector - smaller than the number of rows owned by this partition\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n int size = this->halo_offsets[0];\n\n if (v_out.size() != size * this->A->get_block_dimx())\n {\n v_out.resize(size * this->A->get_block_dimx());\n }\n\n //reorder based on row permutation\n int num_blocks = std::min(4096, (size + 511) / 512);\n inverse_reorder_vector_values <<< num_blocks, 512>>>(v_out.raw(), v_in.raw(), this->renumbering.raw(), v_in.get_block_size(), size);\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManager >::revertAndDownloadVectorWithCons(VVector_v &v_in, const void *data_pinned, int n, int block_dimy)\n{\n if (v_in.get_block_size() != this->A->get_block_dimx()) { printf(\"Blocksize mismatch!\\n\"); }\n\n void *root_v_ptr = NULL;\n int size = v_in.get_unconsolidated_size();\n int num_rows = size / v_in.get_block_size();\n\n if (this->m_is_fine_level_root_partition)\n {\n root_v_ptr = (void *) v_in.raw();\n }\n\n VVector_v temp;\n temp.set_block_dimx(v_in.get_block_dimx());\n temp.set_block_dimy(v_in.get_block_dimy());\n temp.resize(size);\n bool useCudaIpc = this->m_use_cuda_ipc_consolidation;\n\n if (useCudaIpc)\n {\n // Do IPC\n this->ipcExchangePtr(root_v_ptr, this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n // Children partition waits for parent to be done updating vector\n this->ipcWaitForRoot(this->m_is_fine_level_root_partition, this->m_num_fine_level_parts_to_consolidate, this->m_fine_level_parts_to_consolidate, this->m_my_fine_level_destination_part, this->fine_level_id(), this->getFineLevelComms());\n cudaCheckError();\n //reorder based on row permutation\n int num_blocks = std::min(4096, (num_rows + 511) / 512);\n inverse_reorder_vector_values <<< num_blocks, 512>>>( temp.raw(), (value_type *) root_v_ptr, this->m_row_ids_CONS.raw(), v_in.get_block_size(), num_rows);\n cudaCheckError();\n\n if (!this->m_is_fine_level_root_partition)\n {\n hipIpcCloseMemHandle(root_v_ptr);\n }\n }\n else\n {\n if (this->m_is_fine_level_consolidated) // aggregation\n {\n if (this->m_is_fine_level_root_partition)\n {\n IVector_h child_n(this->m_num_fine_level_parts_to_consolidate);\n int max_n = 0;\n\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n\n if (current_part != this->fine_level_id())\n {\n this->getFineLevelComms()->recv_vector(child_n, current_part, 30000 + current_part, i, 1);\n }\n else\n {\n child_n[i] = num_rows;\n }\n\n if (child_n[i] > max_n) { max_n = child_n[i]; }\n }\n\n // Resize temp vector\n VVector_v child_temp;;\n child_temp.resize(max_n * v_in.get_block_size());\n\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n // Pack the vector to be sent\n int num_blocks = std::min(4096, (child_n[i] + 511) / 512);\n\n if (current_part != this->fine_level_id())\n {\n inverse_reorder_vector_values <<< num_blocks, 512>>>( child_temp.raw(), (value_type *) root_v_ptr, this->m_child_row_ids[i].raw(), v_in.get_block_size(), child_n[i]);\n cudaCheckError();\n this->getFineLevelComms()->send_vector(child_temp, current_part, 30001 + current_part, 0, child_n[i]*v_in.get_block_size());\n }\n else\n {\n inverse_reorder_vector_values <<< num_blocks, 512>>>( temp.raw(), (value_type *) root_v_ptr, this->m_row_ids_CONS.raw(), v_in.get_block_size(), child_n[i]);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n }\n else\n {\n IVector_h size(1);\n size[0] = num_rows;\n this->getFineLevelComms()->send_vector(size, this->m_my_fine_level_destination_part, 30000 + this->fine_level_id());\n this->getFineLevelComms()->recv_vector(temp, this->m_my_fine_level_destination_part, 30001 + this->fine_level_id());\n }\n }\n else if (this->m_is_fine_level_glued) // classical\n {\n if (this->m_is_fine_level_root_partition)\n {\n temp.resize(v_in.size());\n revertVector(v_in, temp);\n cudaCheckError();\n\n for (int i = 0; i < this->m_num_fine_level_parts_to_consolidate; i++)\n {\n int current_part = this->m_fine_level_parts_to_consolidate[i];\n\n if (current_part != this->fine_level_id())\n {\n this->getFineLevelComms()->send_vector( temp,\n current_part,\n current_part + 30001,\n this->getConsolidationArrayOffsets()[current_part] - this->getConsolidationArrayOffsets()[this->fine_level_id()],\n this->getConsolidationArrayOffsets()[current_part + 1] - this->getConsolidationArrayOffsets()[current_part] );\n cudaCheckError();\n }\n }\n }\n else\n {\n this->getFineLevelComms()->recv_vector(temp, this->m_my_fine_level_destination_part, 30001 + this->fine_level_id());\n cudaCheckError();\n }\n\n temp.resize(this->getConsolidationArrayOffsets()[this->fine_level_id() + 1] - this->getConsolidationArrayOffsets()[this->fine_level_id()]);\n cudaCheckError();\n }\n }\n\n // Copy on host\n hipMemcpy((value_type *)data_pinned, temp.raw(), temp.size() * sizeof(value_type), hipMemcpyDefault);\n cudaCheckError();\n}\n\ntemplate \nvoid DistributedManager >::transformAndUploadVector(VVector_v &v, const void *data, int n, int block_dim)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n\ntemplate \nvoid DistributedManager >::transformVector(VVector_v &v)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::transformAndUploadVectorWithCons(VVector_v &v, const void *data, int n, int block_dim)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n\ntemplate \nvoid DistributedManager >::revertAndDownloadVector(VVector_v &v, const void *data, int n, int block_dim)\n{\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid DistributedManager >::revertVector(VVector_v &v_in, VVector_v &v_out)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::revertVector(VVector_v &v)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::revertAndDownloadVectorWithCons(VVector_v &v_in, const void *data, int n, int block_dim)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::reorder_matrix()\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::reorder_matrix_owned()\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::obtain_shift_l2g_reordering(index_type n, I64Vector_h &l2g, IVector_h &p, IVector_h &q)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::unpack_partition(index_type *Bp, index_type *Bc, mat_value_type *Bv)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::generatePoisson7pt(int nx, int ny, int nz, int P, int Q, int R)\n{\n FatalError(\"GeneratePoisson7pt only implemented on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \ntemplate \nvoid DistributedManager >::loadDistributedMatrix(\n int num_rows, int num_nonzeros, const int block_dimx, const int block_dimy, const int *row_offsets, \n const t_colIndex *col_indices, const mat_value_type *values, int num_ranks, int num_rows_global, const void *diag, const MatrixDistribution &dist)\n{\n FatalError(\"loadDistributedMatrix only implemented on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid DistributedManager >::createOneRingB2Lmaps()\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::createOneRingHaloRows()\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid DistributedManager >::replaceMatrixCoefficientsNoCons(int n, int nnz, const mat_value_type *data, const mat_value_type *diag)\n{\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid DistributedManager >::replaceMatrixCoefficientsWithCons(int n, int nnz, const mat_value_type *data, const mat_value_type *diag)\n{\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n\ntemplate \nvoid DistributedManager >::createRenumbering(IVector &renumbering)\n{\n if (this->neighbors.size() > 0)\n {\n FatalError(\"Distributed solve only supported on devices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate < class TConfig >\nvoid DistributedManagerBase::remove_boundary(IVector_h &flagArray, IVector_h &B2L_map, int size)\n{\n for (int i = 0; i < size; i++)\n {\n flagArray[B2L_map[i]] = 0;\n }\n}\n\ntemplate < class TConfig >\nvoid DistributedManagerBase::get_unassigned(IVector_h &flagArray, IVector_h &B2L_map, IVector_h &partition_flags, int size, int fa_size/*, int rank*/)\n{\n for (int i = 0; i < size; i++)\n {\n if (B2L_map[i] < fa_size)\n {\n if (flagArray[B2L_map[i]] == 0)\n {\n flagArray[B2L_map[i]] = 1;\n partition_flags[i] = 1;\n }\n }\n }\n}\n\ntemplate < class TConfig >\nvoid DistributedManagerBase::set_unassigned(IVector_h &partition_flags, IVector_h &partition_renum, IVector_h &B2L_map, IVector_h &renumbering, int size, int max_element, int renum_size/*, int rank*/)\n{\n for (int i = 0; i < size; i++)\n {\n if (B2L_map[i] < renum_size)\n {\n if (partition_flags[i] == 1)\n {\n renumbering[B2L_map[i]] = max_element + partition_renum[i];\n }\n\n B2L_map[i] = renumbering[B2L_map[i]];\n }\n }\n}\n\n/* print manager for target rank to a file or stdout */\ntemplate\nvoid DistributedManagerBase::print(char *f, char *s, int trank)\n{\n DistributedManagerBase *m = this;\n int rank = 0;\n int level = 0;\n char filename[1024];\n FILE *fid = NULL;\n int i, j, k, t1, t2;\n#ifdef AMGX_WITH_MPI\n MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n#endif\n\n //check target rank\n if (rank == trank)\n {\n //check whether to output to stdout or a file\n if (f == NULL)\n {\n fid = stdout;\n }\n else\n {\n level = m->A->amg_level_index;\n#ifdef _WIN32\n _snprintf_s(filename, 1024, 1024, \"%s_r%d_l%d.mtx\", f, rank, level);\n#else\n snprintf(filename, 1024, \"%s_r%d_l%d.mtx\", f, rank, level);\n#endif\n fid = fopen(filename, \"w\");\n }\n\n hipDeviceSynchronize();\n cudaCheckError();\n fprintf(fid, \"%s\\n\", s);\n //--- communication info ---\n //compare neighbors\n t1 = m->neighbors.size();\n fprintf(fid, \"neighbors %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->neighbors[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare B2L_rings\n t1 = B2L_rings.size();\n fprintf(fid, \"B2L_rings %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n t2 = m->B2L_rings[i].size();\n fprintf(fid, \"B2L_rings-%d [%d]\\n\", i, t2);\n\n for (j = 0; j < t2; j++)\n {\n k = m->B2L_rings[i][j];\n fprintf(fid, \"%d\\n\", k);\n }\n }\n\n //compare B2L_maps\n t1 = B2L_maps.size();\n fprintf(fid, \"B2L_maps %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n t2 = m->B2L_maps[i].size();\n fprintf(fid, \"B2L_maps-%d [%d]\\n\", i, t2);\n\n for (j = 0; j < t2; j++)\n {\n k = m->B2L_maps[i][j];\n fprintf(fid, \"%d\\n\", k);\n }\n }\n\n //compare L2H_maps\n t1 = L2H_maps.size();\n fprintf(fid, \"L2H_maps %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n t2 = m->L2H_maps[i].size();\n fprintf(fid, \"L2H_maps-%d [%d]\\n\", i, t2);\n\n for (j = 0; j < t2; j++)\n {\n k = m->L2H_maps[i][j];\n fprintf(fid, \"%d\\n\", k);\n }\n }\n\n //--- matrix info ---\n fprintf(fid, \"num_rows_global=%ld\\n\", num_rows_global);\n fprintf(fid, \"_num_rows_interior=%d\\n\", m->_num_rows_interior);\n fprintf(fid, \"_num_rows_owned=%d\\n\", m->_num_rows_owned);\n fprintf(fid, \"_num_rows_full=%d\\n\", m->_num_rows_full);\n fprintf(fid, \"_num_rows_all=%d\\n\", m->_num_rows_all);\n fprintf(fid, \"_num_nz_interior=%d\\n\", m->_num_nz_interior);\n fprintf(fid, \"_num_nz_owned=%d\\n\", m->_num_nz_owned);\n fprintf(fid, \"_num_nz_full=%d\\n\", m->_num_nz_full);\n fprintf(fid, \"_num_nz_all=%d\\n\", m->_num_nz_all);\n //compare # halo rows and halo offsets\n fprintf(fid, \"# halo rings %d and rows %d\\n\", m->num_halo_rings(), m->num_halo_rows());\n t1 = m->halo_offsets.size();\n fprintf(fid, \"halo_offsets %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->halo_offsets[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare halo ranges\n t1 = m->halo_ranges.size();\n fprintf(fid, \"halo_ranges %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->halo_ranges[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare halo ranges (host)\n t1 = m->halo_ranges_h.size();\n fprintf(fid, \"halo_ranges_h %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->halo_ranges_h[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare part offsets\n t1 = m->part_offsets.size();\n fprintf(fid, \"part_offsets %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->part_offsets[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare part offsets (host)\n t1 = m->part_offsets_h.size();\n fprintf(fid, \"part_offsets_h %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->part_offsets_h[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare interior row list\n t1 = m->interior_rows_list.size();\n fprintf(fid, \"interior_rows_list %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->interior_rows_list[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare boundary row list\n t1 = m->boundary_rows_list.size();\n fprintf(fid, \"boundary_rows_list %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->boundary_rows_list[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare halo1 row list\n t1 = m->halo1_rows_list.size();\n fprintf(fid, \"halo1_rows_list %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->halo1_rows_list[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n fprintf(fid, \"pointers halo_rows=%p and halo_btl=%p\\n\", m->halo_rows, m->halo_btl);\n //--- packing info ---\n //compare local to global map\n t1 = m->local_to_global_map.size();\n fprintf(fid, \"local_to_global_map %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->local_to_global_map[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare renumbering\n t1 = m->renumbering.size();\n fprintf(fid, \"renumbering %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->renumbering[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //compare inverse renumbering\n t1 = m->inverse_renumbering.size();\n fprintf(fid, \"inverse_renumbering %d\\n\", t1);\n\n for (i = 0; i < t1; i++)\n {\n k = m->inverse_renumbering[i];\n fprintf(fid, \"%d\\n\", k);\n }\n\n //--- GPU related and miscellaneous info\n //streams\n fprintf(fid, \"streams i=%p, b=%p\\n\", m->get_int_stream(), m->get_bdy_stream());\n //miscellaneous info\n int64_t bi = m->base_index(); //inlined function\n int np = m->get_num_partitions(); //inlined function\n int rp = (int)m->isRootPartition(); //cast from boolean to int\n fprintf(fid, \"gid=%d,bi=%ld,np=%d,rp=%d,ir=%d,in=%d,bn=%d\\n\", m->global_id(), bi, np, rp, m->index_range(), m->num_interior_nodes(), m->num_boundary_nodes());\n hipDeviceSynchronize();\n cudaCheckError();\n hipGetLastError();\n\n if (fid != stdout)\n {\n fclose(fid);\n }\n }\n}\n\n/* print manager for target rank to a file or stdout (for all ranks) */\ntemplate\nvoid DistributedManagerBase::printToFile(char *f, char *s)\n{\n DistributedManagerBase *m = this;\n int rank = 0;\n#ifdef AMGX_WITH_MPI\n MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n#endif\n //notice that print will be called with different (target) rank on different ranks/processes\n m->print(f, s, rank);\n}\n\n/* compare two managers */\ntemplate\nint DistributedManagerBase::compare(DistributedManagerBase *m2)\n{\n DistributedManagerBase *m1 = this;\n int i, j, t1, t2;\n //compare neighbors\n t1 = m1->neighbors.size();\n t2 = m2->neighbors.size();\n\n if (t1 != t2)\n {\n return 1;\n }\n\n for (i = 0; i < t1; i++)\n {\n if (m1->neighbors[i] != m2->neighbors[i])\n {\n return 2;\n }\n }\n\n //compare B2L_rings\n for (i = 0; i < (m1->neighbors.size()); i++)\n {\n t1 = m1->B2L_rings[i].size();\n t2 = m2->B2L_rings[i].size();\n\n if (t1 != t2)\n {\n return 3;\n }\n\n for (j = 0; j < t1; j++)\n {\n if (m1->B2L_rings[i][j] != m2->B2L_rings[i][j])\n {\n return 4;\n }\n }\n }\n\n //compare B2L_maps\n t1 = m1->B2L_maps.size();\n t2 = m2->B2L_maps.size();\n\n if (t1 != t2)\n {\n return 5;\n }\n\n for (i = 0; i < t1; i++)\n {\n if (m1->B2L_maps[i] != m2->B2L_maps[i])\n {\n return 6;\n }\n }\n\n //compare L2H_maps\n t1 = m1->L2H_maps.size();\n t2 = m2->L2H_maps.size();\n\n if (t1 != t2)\n {\n return 7;\n }\n\n for (i = 0; i < t1; i++)\n {\n if (m1->L2H_maps[i] != m2->L2H_maps[i])\n {\n return 8;\n }\n }\n\n return 0;\n}\n\ntemplate \nDistributedManager< TemplateConfig >::~DistributedManager< TemplateConfig >()\n{\n}\n\ntemplate \nDistributedManager< TemplateConfig >::~DistributedManager< TemplateConfig >()\n{\n}\n\ntemplate \nvoid DistributedManagerBase::consolidateB2Lmaps(IVector_h_vector &dest_coarse_B2L_maps, IVector_h_vector &coarse_B2L_maps, IVector_h &fine_neigh_to_coarse_neigh, int num_coarse_neighbors, int num_fine_neighbors)\n{\n consB2Lmaps(dest_coarse_B2L_maps, coarse_B2L_maps, fine_neigh_to_coarse_neigh, num_coarse_neighbors, num_fine_neighbors);\n}\n\ntemplate \nvoid DistributedManagerBase::consolidateB2Lmaps(IVector_d_vector &dest_coarse_B2L_maps, IVector_d_vector &coarse_B2L_maps, IVector_h &fine_neigh_to_coarse_neigh, int num_coarse_neighbors, int num_fine_neighbors)\n{\n consB2Lmaps(dest_coarse_B2L_maps, coarse_B2L_maps, fine_neigh_to_coarse_neigh, num_coarse_neighbors, num_fine_neighbors);\n}\n\ntemplate \nvoid DistributedManagerBase::consolidateB2LmapsOnRoot(int &num_consolidated_neighbors, IVector_h_vector &consolidated_B2L_maps, IVector_h &consolidated_coarse_ids, IVector_h_vector &dest_coarse_B2L_maps, IVector_h &coarse_neigh_to_fine_part, IVector_h &num_bdy_per_coarse_neigh, IVector_h &fine_parts_to_consolidate, int num_fine_parts_to_consolidate, int my_id, int my_destination_part, bool is_root_partition, int num_coarse_neighbors, DistributedComms *comms)\n{\n consB2LmapsOnRoot(num_consolidated_neighbors, consolidated_B2L_maps, consolidated_coarse_ids, dest_coarse_B2L_maps, coarse_neigh_to_fine_part, num_bdy_per_coarse_neigh, fine_parts_to_consolidate, num_fine_parts_to_consolidate, my_id, my_destination_part, is_root_partition, num_coarse_neighbors, comms);\n}\n\ntemplate \nvoid DistributedManagerBase::consolidateB2LmapsOnRoot(int &num_consolidated_neighbors, IVector_d_vector &consolidated_B2L_maps, IVector_h &consolidated_coarse_ids, IVector_d_vector &dest_coarse_B2L_maps, IVector_h &coarse_neigh_to_fine_part, IVector_h &num_bdy_per_coarse_neigh, IVector_h &fine_parts_to_consolidate, int num_fine_parts_to_consolidate, int my_id, int my_destination_part, bool is_root_partition, int num_coarse_neighbors, DistributedComms *comms)\n{\n consB2LmapsOnRoot(num_consolidated_neighbors, consolidated_B2L_maps, consolidated_coarse_ids, dest_coarse_B2L_maps, coarse_neigh_to_fine_part, num_bdy_per_coarse_neigh, fine_parts_to_consolidate, num_fine_parts_to_consolidate, my_id, my_destination_part, is_root_partition, num_coarse_neighbors, comms);\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class DistributedManager::Type >;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template void DistributedManager::Type>::loadDistributedMatrix( \\\n int, int, const int, const int, const int*, const int *col_indices, const mat_value_type*, int, int, const void*, const MatrixDistribution &dist); \\\n template void DistributedManager::Type>::loadDistributedMatrix( \\\n int, int, const int, const int, const int*, const int64_t *col_indices, const mat_value_type*, int, int, const void*, const MatrixDistribution &dist);\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class DistributedManagerBase::Type >;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n\n", "source": "/kaggle/working/amgx/src/distributed/distributed_manager.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n\nnamespace amgx\n{\n\n/***************************************\n * Source Definitions\n ***************************************/\n\ntemplate \n__global__\nvoid copy_buffer(value_type *__restrict__ dest, const value_type *__restrict__ src, int size)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n dest[tidx] = src[tidx];\n }\n}\n\ntemplate \nstatic void error_vector_too_small(const Vector &v, int required_size)\n{\n std::stringstream ss;\n ss << \"Vector size too small: not enough space for halo elements.\" << std::endl;\n ss << \"Vector: {tag = \" << v.tag << \", \" << \"size = \" << v.size() << \"}\" << std::endl;\n ss << \"Required size: \" << required_size << std::endl;\n FatalError(ss.str(), AMGX_ERR_INTERNAL);\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::printString(const std::string &str)\n{\n#ifdef AMGX_WITH_MPI\n int total = 0;\n int part = 0;\n MPI_Comm_size( mpi_comm, &total );\n MPI_Comm_rank( mpi_comm, &part);\n\n for (int i = 0; i < total; i++)\n {\n MPI_Barrier(mpi_comm);\n\n if (i == part)\n {\n std::cout << \"my_id=\" << i << std::endl;\n std::cout << str << std::endl;\n }\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_reduction(T &arrays, T2 &data, const Operator &m)\n{\n#ifdef AMGX_WITH_MPI\n int numRanks = 0;\n MPI_Comm_size( mpi_comm, &numRanks );\n typedef typename T2::value_type value_type;\n std::vector linear_buffer(numRanks * data.size());\n // Allgather() copy data from all ranks into a 1D array (linear_buffer)\n MPI_Allgather(&data[0], sizeof(value_type)*data.size(), MPI_BYTE,\n &linear_buffer[0], sizeof(value_type)*data.size(), MPI_BYTE, mpi_comm);\n // then it is copied from 1D array to the 2D array (called arrays)\n arrays.resize(numRanks);\n\n // first copy data from other ranks\n for (int i = 0; i < numRanks; i++)\n {\n if (i != m.getManager()->global_id())\n {\n arrays[i].assign(linear_buffer.begin() + i * data.size(), linear_buffer.begin() + (i + 1)*data.size());\n }\n }\n\n // then copy data from its won rank\n arrays[m.getManager()->global_id()] = data;\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate\n__global__\nvoid print_data(T **data, int size)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n T *loc = data[0];\n printf(\"data = %d\\n\", loc[tid]);\n}\n\n\ntemplate\n__global__\nvoid print_data_f(T **data, int size)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n T *loc = data[0];\n printf(\"data = %f\\n\", loc[tid]);\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_setup(T &b, const Matrix &m, int num_rings)\n{\n#ifdef AMGX_WITH_MPI\n int bsize = b.get_block_size();\n int num_cols = b.get_num_cols();\n\n if (bsize != 1 && num_cols != 1)\n FatalError(\"Error: vector cannot have block size and subspace size.\",\n AMGX_ERR_INTERNAL);\n\n // set num neighbors = size of b2l_rings\n // need to do this because comms might have more neighbors than our matrix knows about\n set_neighbors(m.manager->B2L_rings.size());\n\n if (b.in_transfer & SENDING)\n {\n b.in_transfer = IDLE;\n }\n\n b.requests.resize(2 * neighbors); //first part is sends second is receives\n b.statuses.resize(2 * neighbors);\n\n for (int i = 0; i < 2 * neighbors; i++)\n {\n b.requests[i] = MPI_REQUEST_NULL;\n }\n\n typedef typename T::value_type vtyp;\n\n if (b.linear_buffers_size < neighbors)\n {\n if (b.linear_buffers_size != 0) { amgx::memory::cudaFreeHost(b.linear_buffers); }\n cudaCheckError();\n\n amgx::memory::cudaMallocHost((void **) & (b.linear_buffers), neighbors * sizeof(vtyp *));\n cudaCheckError();\n b.linear_buffers_size = neighbors;\n b.linear_buffers_ptrs.resize(neighbors);\n }\n\n cudaCheckError();\n\n int send_size = 0;\n for (int i = 0; i < neighbors; i++)\n {\n send_size += m.manager->B2L_rings[i][num_rings] * bsize * num_cols;\n }\n\n int total_size = send_size + (m.manager->halo_offsets[num_rings * neighbors] - m.manager->halo_offsets[0]) * bsize * num_cols;\n\n b.buffer_size = send_size;\n\n if (b.buffer == NULL)\n {\n b.buffer = new T(total_size);\n }\n else if (total_size > b.buffer->size())\n {\n b.buffer->resize(total_size);\n\n // It is more efficient to synchronise only when linear buffers change\n cudaStreamSynchronize(0);\n cudaCheckError();\n }\n\n int offset = 0;\n for (int i = 0; i < neighbors; i++)\n {\n b.linear_buffers[i] = b.buffer->raw() + offset;\n offset += m.manager->B2L_rings[i][num_rings] * bsize * num_cols;\n }\n\n // Copy to device\n {\n cudaMemcpy(amgx::thrust::raw_pointer_cast(&b.linear_buffers_ptrs[0]), &(b.linear_buffers[0]), neighbors * sizeof(vtyp *), cudaMemcpyDefault);\n cudaCheckError();\n }\n\n if (total_size != 0)\n {\n if (b.explicit_host_buffer == NULL)\n {\n b.host_buffer.resize(1);\n cudaEventCreateWithFlags(&b.mpi_event, cudaEventDisableTiming);\n amgx::memory::cudaMallocHost((void **)&b.explicit_host_buffer, total_size * sizeof(vtyp));\n cudaCheckError();\n }\n else if (total_size > b.explicit_buffer_size)\n {\n amgx::memory::cudaFreeHost(b.explicit_host_buffer);\n cudaCheckError();\n amgx::memory::cudaMallocHost((void **)&b.explicit_host_buffer, total_size * sizeof(vtyp));\n }\n\n cudaCheckError();\n b.explicit_buffer_size = total_size;\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_setup_L2H(T &b, Matrix &m, int num_rings)\n{\n#ifdef AMGX_WITH_MPI\n int num_neighbors = m.manager->neighbors.size();\n int bsize = b.get_block_size();\n b.requests.resize(2 * num_neighbors);\n b.statuses.resize(2 * num_neighbors);\n\n for (int i = 0; i < 2 * num_neighbors; i++)\n {\n b.requests[i] = MPI_REQUEST_NULL;\n }\n\n // compute new sizes\n int send_size = (m.manager->halo_offsets[num_rings * num_neighbors] - m.manager->halo_offsets[0]) * bsize;\n int recv_size = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n recv_size += m.manager->B2L_rings[i][num_rings] * bsize;\n }\n\n int size = send_size + recv_size;\n // resize device and host buffers\n b.buffer_size = size;\n\n if (b.buffer == NULL)\n {\n b.buffer = new T(size);\n }\n else if (size > b.buffer->size())\n {\n b.buffer->resize(size);\n cudaStreamSynchronize(0);\n cudaCheckError();\n }\n\n b.host_buffer.resize(send_size + recv_size);\n typedef typename T::value_type vtyp;\n\n if (size != 0)\n {\n if (b.explicit_host_buffer == NULL)\n {\n b.host_buffer.resize(1);\n cudaEventCreateWithFlags(&b.mpi_event, cudaEventDisableTiming);\n amgx::memory::cudaMallocHost((void **)&b.explicit_host_buffer, size * sizeof(vtyp));\n cudaCheckError();\n }\n else if (size > b.explicit_buffer_size)\n {\n amgx::memory::cudaFreeHost(b.explicit_host_buffer);\n cudaCheckError();\n amgx::memory::cudaMallocHost((void **)&b.explicit_host_buffer, size * sizeof(vtyp));\n }\n\n cudaCheckError();\n b.explicit_buffer_size = size;\n }\n\n // setup linear buffers\n if (b.linear_buffers_size < neighbors)\n {\n if (b.linear_buffers_size != 0) { amgx::memory::cudaFreeHost(b.linear_buffers); }\n cudaCheckError();\n\n amgx::memory::cudaMallocHost((void **) & (b.linear_buffers), neighbors * sizeof(vtyp *));\n cudaCheckError();\n b.linear_buffers_size = neighbors;\n }\n\n int total = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n b.linear_buffers[i] = b.buffer->raw() + total;\n\n for (int j = 0; j < num_rings; j++)\n {\n total += (m.manager->halo_offsets[j * num_neighbors + i + 1] - m.manager->halo_offsets[j * num_neighbors + i]) * bsize;\n }\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_vec_exchange(std::vector > > &b, const Matrix &m)\n{\n#ifdef AMGX_WITH_MPI\n using VecHost = Vector>;\n using VecType = Vector>;\n std::vector host_send(b.size());\n std::vector host_recv(b.size());\n std::vector send_sizes(b.size());\n std::vector recv_sizes(b.size());\n\n for (int i = 0; i < neighbors; i++)\n {\n send_sizes[i] = b[i].size();\n MPI_Irecv(&recv_sizes[i], 1, MPI_INT, m.manager->neighbors[i], m.manager->neighbors[i], mpi_comm, &requests[neighbors + i]);\n }\n\n for (int i = 0; i < neighbors; i++ )\n {\n MPI_Isend(&send_sizes[i], 1, MPI_INT, m.manager->neighbors[i], m.manager->global_id(), mpi_comm, &requests[i]);\n }\n\n SynchronizerVisitor syncV(Direction::DeviceToHost, b, host_send, recv_sizes, m);\n Accept(syncV);\n MPI_Waitall(2 * neighbors, &requests[0], MPI_STATUSES_IGNORE); //I wait for data to be sent too, as I will deallocate it upon exit\n ReceiverVisitor receiverV(b, host_recv, recv_sizes, m);\n Accept(receiverV);\n SenderVisitor senderV(b, host_send, send_sizes, m);\n Accept(senderV);\n MPI_Waitall(2 * neighbors, &requests[0], MPI_STATUSES_IGNORE); //I wait for data to be sent too, as I will deallocate it upon exit\n syncV.switch_direction();\n syncV.set_local(receiverV.get_local());//only CUDA-aware MPI needs this\n syncV.set_host_vec(host_recv);//only Normal MPI needs this\n Accept(syncV);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::send_vec(T &b, int destination, int tag, int offset, int size)\n{\n#ifdef AMGX_WITH_MPI\n\n if (size == -1) { size = b.size(); }\n\n if (size == 0) { return; }\n\n {\n // Wait for any outstanding receive calls\n int num_recv_requests = b.recv_requests.size();\n\n if (num_recv_requests != 0)\n {\n MPI_Waitall(num_recv_requests, &b.recv_requests[0], MPI_STATUSES_IGNORE);\n b.recv_requests.clear();\n b.recv_statuses.clear();\n }\n }\n\n SynchSendVecVisitor sv(b, destination, tag, offset, size);\n Accept(sv);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::send_vec_async(T &b, int destination, int tag, int offset, int size)\n{\n#ifdef AMGX_WITH_MPI\n\n if (size == -1) { size = b.size(); }\n\n if (size == 0) { return; }\n\n // Wait for any outstanding receive calls\n int num_recv_requests = b.recv_requests.size();\n\n if (num_recv_requests != 0)\n {\n // Wait for any outstanding sends, receives to finish\n MPI_Waitall(num_recv_requests, &b.recv_requests[0], MPI_STATUSES_IGNORE);\n b.recv_requests.clear();\n b.recv_statuses.clear();\n }\n\n MPI_Request new_req;\n MPI_Status new_status;\n b.send_requests.push_back(new_req);\n b.send_statuses.push_back(new_status);\n int num_send_requests = b.send_requests.size();\n AsynchSendVecVisitor sv(b, destination, tag, offset, size);\n Accept(sv);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::send_vec_wait_all(T &b)\n{\n#ifdef AMGX_WITH_MPI\n int num_send_requests = b.send_requests.size();\n MPI_Waitall(num_send_requests, &b.send_requests[0], MPI_STATUSES_IGNORE);\n b.send_requests.clear();\n b.send_statuses.clear();\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::recv_vec(T &b, int source, int tag, int offset, int size)\n{\n#ifdef AMGX_WITH_MPI\n\n if (size == -1) { size = b.size(); }\n\n if (size == 0) { return; }\n\n // Wait for any outstanding sendcalls\n int num_send_requests = b.send_requests.size();\n\n if (num_send_requests != 0)\n {\n // Wait for any outstanding sends, receives to finish\n MPI_Waitall(num_send_requests, &b.send_requests[0], MPI_STATUSES_IGNORE);\n b.send_requests.clear();\n b.send_statuses.clear();\n }\n\n SynchRecvVecVisitor rv(b, source, tag, offset, size);\n Accept(rv);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::recv_vec_async(T &b, int source, int tag, int offset, int size)\n{\n#ifdef AMGX_WITH_MPI\n\n if (size == -1) { size = b.size(); }\n\n if (size == 0) { return; }\n\n // Wait for any outstanding sendcalls\n int num_send_requests = b.send_requests.size();\n\n if (num_send_requests != 0)\n {\n // Wait for any outstanding sends, receives to finish\n MPI_Waitall(num_send_requests, &b.send_requests[0], MPI_STATUSES_IGNORE);\n b.send_requests.clear();\n b.send_statuses.clear();\n }\n\n MPI_Request new_req;\n MPI_Status new_status;\n b.recv_requests.push_back(new_req);\n b.recv_statuses.push_back(new_status);\n int num_recv_requests = b.recv_requests.size();\n AsynchRecvVecVisitor rv(b, source, tag, offset, size);\n Accept(rv);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::recv_vec_wait_all(T &b)\n{\n#ifdef AMGX_WITH_MPI\n int num_recv_requests = b.recv_requests.size();\n MPI_Waitall(num_recv_requests, &b.recv_requests[0], MPI_STATUSES_IGNORE);\n b.recv_requests.clear();\n b.recv_statuses.clear();\n CopyHostFunctor cpyf(b);\n LightVisitor lvcpy(cpyf);\n Accept(lvcpy);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_gather_L2H(T &b, const Matrix &m, int num_rings, cudaStream_t stream)\n{\n#ifdef AMGX_WITH_MPI\n // Gather L2H\n int num_neighbors = m.manager->neighbors.size();\n int bsize = b.get_block_size();\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int total = 0;\n\n for (int j = 0; j < num_rings; j++)\n {\n int size = (m.manager->halo_offsets[j * num_neighbors + i + 1] - m.manager->halo_offsets[j * num_neighbors + i]) * bsize;\n\n if (size != 0)\n {\n // we need to use new indices after renumbering - these are stored in L2H_maps\n cudaMemcpyAsync(b.linear_buffers[i] + total, b.raw() + m.manager->halo_offsets[j * num_neighbors + i], size*sizeof(typename T::value_type), cudaMemcpyDefault, stream);\n cudaCheckError();\n total += size;\n }\n }\n\n cudaCheckError();\n }\n cudaStreamSynchronize(stream);\n cudaCheckError();\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_gather_L2H_v2(T &b, const Matrix &m, int num_rings, cudaStream_t stream)\n{\n#ifdef AMGX_WITH_MPI\n\n if (num_rings != 1)\n {\n FatalError(\"Halo exchange with ring > 1 not implemented in do_gather_L2H\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // Gather L2H\n int num_neighbors = m.manager->neighbors.size();\n int bsize = b.get_block_size();\n int total = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n total += (m.manager->halo_offsets[i + 1] - m.manager->halo_offsets[i]) * bsize;\n }\n\n if (total != 0)\n {\n cudaMemcpyAsync(b.linear_buffers[0], b.raw() + m.manager->halo_offsets[0], total*sizeof(typename T::value_type), cudaMemcpyDefault, stream);\n cudaCheckError();\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_add_from_halo(T &b, const Matrix &m, int num_rings, cudaStream_t &stream)\n{\n#ifdef AMGX_WITH_MPI\n typedef typename T::value_type vtyp;\n AddFromHalo1Functor ex1(b, m, num_rings, 0, stream);\n AddFromHalo2Functor ex2(b, m, num_rings, 0, stream);\n AddFromHalo3Functor ex3(b, m, num_rings, 0, stream);\n FSMVisitor fsmV;\n fsmV.get_functors().push_back(&ex1);\n fsmV.get_functors().push_back(&ex2);\n fsmV.get_functors().push_back(&ex3);\n Accept(fsmV);\n ex2.get_send_size() = ex1.get_send_size();//pass relevant info from one state to another\n fsmV.next();//advance FSM\n b.in_transfer = RECEIVING | SENDING;\n\n Accept(fsmV);\n //pass relevant info from one state to another:\n //\n ex3.get_recv_size() = ex2.get_offset();\n ex3.get_send_size() = ex2.get_send_size();\n fsmV.next();//advance FSM\n\n // I only wait to receive data, I can start working before all my buffers were sent\n MPI_Waitall(2 * neighbors, &b.requests[0], MPI_STATUSES_IGNORE);\n b.dirtybit = 0;\n Accept(fsmV);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_exchange_halo(T &b, const Matrix &m, int num_rings)\n{\n#ifdef AMGX_WITH_MPI\n typedef typename T::value_type value_type;\n ExcHalo1Functor ex1(b, m, num_rings, 0);\n ExcHalo2Functor ex2(b, m, num_rings, 0);\n ExcHalo3Functor ex3(b, m, num_rings, 0);\n FSMVisitor fsmV;\n fsmV.get_functors().push_back(&ex1);\n fsmV.get_functors().push_back(&ex2);\n fsmV.get_functors().push_back(&ex3);\n Accept(fsmV);\n fsmV.next();//advance FSM\n b.in_transfer = RECEIVING | SENDING;\n Accept(fsmV);\n ex3.get_offset() = ex2.get_offset();//pass relevant info from one state to another\n fsmV.next();//advance FSM\n MPI_Waitall(2 * neighbors, &b.requests[0], MPI_STATUSES_IGNORE); //I only wait to receive data, I can start working before all my buffers were sent\n b.dirtybit = 0;\n //FSM step 3:\n Accept(fsmV);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_exchange_halo_async(T &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream)\n{\n#ifdef AMGX_WITH_MPI\n\n if (tag == -1) { FatalError(\"Untagged vector\\n\", AMGX_ERR_NOT_IMPLEMENTED); }\n\n cudaEventSynchronize(event);\n cudaCheckError();\n ExcHalo1Functor ex1(b, m, 1, 0, tag, stream);\n ExcHalo2AsyncFunctor ex2(b, m, 1, 0, tag);\n ExcHalo3AsyncFunctor ex3(b, m, 1, 0, tag, stream);\n FSMVisitor fsmV;\n fsmV.get_functors().push_back(&ex1);\n fsmV.get_functors().push_back(&ex2);\n fsmV.get_functors().push_back(&ex3);\n Accept(fsmV);\n fsmV.next();//advance FSM\n b.in_transfer = RECEIVING | SENDING;\n Accept(fsmV);\n fsmV.next();//advance FSM\n\n if (min_rows_latency_hiding < 0 || m.manager->num_rows_all() < min_rows_latency_hiding)\n {\n MPI_Waitall(2 * neighbors, &b.requests[0], /*&b.statuses[0]*/ MPI_STATUSES_IGNORE); //I only wait to receive data, I can start working before all my buffers were sent\n b.dirtybit = 0;\n Accept(fsmV);\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_exchange_halo_wait(T &b, const Matrix &m, cudaStream_t stream)\n{\n#ifdef AMGX_WITH_MPI\n int bsize = b.get_block_size();\n\n if (!(min_rows_latency_hiding < 0 || m.manager->num_rows_all() < min_rows_latency_hiding))\n {\n MPI_Waitall(2 * neighbors, &b.requests[0], MPI_STATUSES_IGNORE);\n b.dirtybit = 0;\n b.in_transfer = IDLE;\n HalloWaitCopyFunctor hcpyf(b, m, stream);\n LightVisitor lvcpy(hcpyf);//Visitor branching off: no-op for GPU Direct\n Accept(lvcpy);\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_send_receive_wait(T &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t &stream)\n{\n#ifdef AMGX_WITH_MPI\n\n if (tag == -1)\n {\n FatalError(\"Untagged vector\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n SendRecvWait1Functor ex1(b, m, 1, 0, tag, stream);\n SendRecvWait2Functor ex2(b, m, 1, 0, tag, stream);\n SendRecvWait3Functor ex3(b, m, 1, 0, tag, stream);\n FSMVisitor fsmV;\n fsmV.get_functors().push_back(&ex1);\n fsmV.get_functors().push_back(&ex2);\n fsmV.get_functors().push_back(&ex3);\n cudaEventSynchronize(event);\n cudaCheckError();\n Accept(fsmV);\n fsmV.next();//advance FSM\n b.in_transfer = RECEIVING | SENDING;\n Accept(fsmV);\n fsmV.next();//advance FSM\n\n if (neighbors > 0)\n {\n MPI_Waitall(2 * neighbors, &b.requests[0], MPI_STATUSES_IGNORE);\n }\n\n Accept(fsmV);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nbool CommsMPIHostBufferStream::do_exchange_halo_query(T &b, const Matrix &m)\n{\n#ifdef AMGX_WITH_MPI\n\n if (b.dirtybit == 0) { return true; }\n\n if ((b.in_transfer & RECEIVING) == 0) { return false; }\n\n int finished = 1;\n MPI_Status status;\n\n for (int i = 0; i < neighbors; i++)\n {\n int res = MPI_Test(&b.requests[neighbors + i], &finished, &status);\n\n if (finished == 0 || res != MPI_SUCCESS)\n {\n return false;\n }\n }\n\n b.dirtybit = 0;\n b.in_transfer = SENDING;\n HalloWaitCopyFunctor hcpyf(b, m, NULL);\n LightVisitor lvcpy(hcpyf);//Visitor branching off: no-op for GPU Direct\n Accept(lvcpy);\n return true;\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::exchange_matrix_halo(IVector_Array &row_offsets, I64Vector_Array &col_indices, MVector_Array &values, I64Vector_Array &halo_row_ids, IVector_h &neighbors_list, int global_id)\n{\n#ifdef AMGX_WITH_MPI\n int total = 0;\n MPI_Comm_size( mpi_comm, &total );\n int num_neighbors = neighbors_list.size();\n std::vector local_row_offsets(num_neighbors);\n std::vector local_col_indices(num_neighbors);\n std::vector local_values(num_neighbors);\n std::vector send_row_offsets(num_neighbors);\n std::vector send_col_indices(num_neighbors);\n std::vector send_values(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n send_row_offsets[i] = row_offsets[i];\n send_col_indices[i] = col_indices[i];\n send_values[i] = values[i];\n }\n\n std::vector local_row_ids(0);\n std::vector send_row_ids(0);\n\n if (halo_row_ids.size() != 0)\n {\n local_row_ids.resize(num_neighbors);\n send_row_ids.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n send_row_ids[i] = halo_row_ids[i];\n }\n }\n\n // send metadata\n std::vector metadata(num_neighbors * 2); // num_rows+1, num_nz\n\n for (int i = 0; i < num_neighbors; i++)\n {\n metadata[i * 2 + 0] = row_offsets[i].size();\n metadata[i * 2 + 1] = col_indices[i].size();\n MPI_Isend(&metadata[i * 2 + 0], 2, MPI_INT, neighbors_list[i], 0, mpi_comm, &requests[i]);\n }\n\n // receive metadata\n std::vector metadata_recv(2);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n MPI_Recv(&metadata_recv[0], 2, MPI_INT, neighbors_list[i], 0, mpi_comm, MPI_STATUSES_IGNORE);\n local_row_offsets[i].resize(metadata_recv[0]);\n local_col_indices[i].resize(metadata_recv[1]);\n local_values[i].resize(metadata_recv[1]);\n\n if (local_row_ids.size() != 0)\n {\n if (metadata_recv[0] - 1 > 0)\n {\n local_row_ids[i].resize(metadata_recv[0] - 1); // row_ids is one smaller than row_offsets\n }\n }\n }\n\n MPI_Waitall(num_neighbors, &requests[0], MPI_STATUSES_IGNORE); // data is already received, just closing the handles\n // receive matrix data\n typedef typename T_Config::MatPrec mvalue;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n MPI_Irecv(local_row_offsets[i].raw(), local_row_offsets[i].size(), MPI_INT, neighbors_list[i], 10 * neighbors_list[i] + 0, mpi_comm, &requests[3 * num_neighbors + i]);\n MPI_Irecv(local_col_indices[i].raw(), local_col_indices[i].size()*sizeof(int64_t), MPI_BYTE, neighbors_list[i], 10 * neighbors_list[i] + 1, mpi_comm, &requests[4 * num_neighbors + i]);\n MPI_Irecv(local_values[i].raw(), local_values[i].size()*sizeof(mvalue), MPI_BYTE, neighbors_list[i], 10 * neighbors_list[i] + 2, mpi_comm, &requests[5 * num_neighbors + i]);\n\n if (send_row_ids.size() != 0)\n {\n MPI_Irecv(local_row_ids[i].raw(), local_row_ids[i].size()*sizeof(int64_t), MPI_BYTE, neighbors_list[i], 10 * neighbors_list[i] + 3, mpi_comm, &requests[7 * num_neighbors + i]);\n }\n }\n //Note: GPU Direct should use row_offsets[], col_indices[], values[] directly in here:\n // send matrix: row offsets, col indices, values\n for (int i = 0; i < num_neighbors; i++)\n {\n MPI_Isend(send_row_offsets[i].raw(), send_row_offsets[i].size(), MPI_INT, neighbors_list[i], 10 * global_id + 0, mpi_comm, &requests[i]);\n MPI_Isend(send_col_indices[i].raw(), send_col_indices[i].size()*sizeof(int64_t), MPI_BYTE, neighbors_list[i], 10 * global_id + 1, mpi_comm, &requests[num_neighbors + i]);\n MPI_Isend(send_values[i].raw(), send_values[i].size()*sizeof(mvalue), MPI_BYTE, neighbors_list[i], 10 * global_id + 2, mpi_comm, &requests[2 * num_neighbors + i]);\n\n if (send_row_ids.size() != 0)\n {\n MPI_Isend(send_row_ids[i].raw(), send_row_ids[i].size()*sizeof(int64_t), MPI_BYTE, neighbors_list[i], 10 * global_id + 3, mpi_comm, &requests[6 * num_neighbors + i]);\n }\n }\n\n if (halo_row_ids.size() != 0)\n {\n MPI_Waitall(8 * num_neighbors, &requests[0], MPI_STATUSES_IGNORE); //I have to wait for my stuff to be sent too, because I deallocate those matrices upon exditing this function\n }\n else\n {\n MPI_Waitall(6 * num_neighbors, &requests[0], MPI_STATUSES_IGNORE); //I have to wait for my stuff to be sent too, because I deallocate those matrices upon exditing this function\n }\n\n //Note: GPU Direct should swap here\n for (int i = 0; i < num_neighbors; i++)\n {\n row_offsets[i] = local_row_offsets[i];\n col_indices[i] = local_col_indices[i];\n values[i] = local_values[i];\n\n if (halo_row_ids.size() != 0)\n {\n halo_row_ids[i] = local_row_ids[i];\n }\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::exchange_matrix_halo(Matrix_Array &halo_rows, DistributedManager_Array &halo_btl, const Matrix &m)\n{\n#ifdef AMGX_WITH_MPI\n int total = 0;\n MPI_Comm_size( mpi_comm, &total );\n int bsize = m.get_block_size();\n int rings = m.manager->B2L_rings[0].size() - 1;\n int diag = m.hasProps(DIAG);\n std::vector local_copy(halo_rows.size());\n std::vector send(halo_rows.size());\n std::vector local_copy_manager(halo_rows.size());\n std::vector send_manager(halo_rows.size());\n\n //Note: GPU Direct does *NOT* do that:\n for (int i = 0; i < halo_rows.size(); i++)\n {\n send_manager[i] = halo_btl[i];\n send[i] = halo_rows[i];\n }\n\n {\n // there shouldn't be any uncompleted requests, because we don't want to rewrite them\n int completed;\n MPI_Testall(requests.size(), &requests[0], &completed, MPI_STATUSES_IGNORE);\n\n if (!completed)\n {\n MPI_Waitall(requests.size(), &requests[0], MPI_STATUSES_IGNORE);\n }\n }\n\n std::vector metadata(neighbors * (rings + 1 + 5)); //ring offsets (rings+1), num_rows, num_nz, base_index, index_range\n\n for (int i = 0; i < neighbors; i++)\n {\n for (int j = 0; j <= rings; j++) { metadata[i * (rings + 1 + 5) + j] = halo_btl[i].B2L_rings[0][j]; }\n\n metadata[i * (rings + 1 + 5) + rings + 1] = halo_rows[i].get_num_rows();\n metadata[i * (rings + 1 + 5) + rings + 2] = halo_rows[i].get_num_nz();\n metadata[i * (rings + 1 + 5) + rings + 3] = halo_btl[i].base_index();\n metadata[i * (rings + 1 + 5) + rings + 4] = halo_btl[i].index_range();\n metadata[i * (rings + 1 + 5) + rings + 5] = halo_btl[i].L2H_maps[0].size();\n MPI_Isend(&metadata[i * (rings + 1 + 5)], rings + 6, MPI_INT, m.manager->neighbors[i], 0, mpi_comm, &requests[i]);\n }\n\n std::vector metadata_recv(rings + 1 + 5);\n\n for (int i = 0; i < neighbors; i++)\n {\n MPI_Recv(&metadata_recv[0], rings + 6, MPI_INT, m.manager->neighbors[i], 0, mpi_comm, /*&status*/MPI_STATUSES_IGNORE);\n local_copy[i].addProps(CSR);\n\n if (diag) { local_copy[i].addProps(DIAG); }\n\n local_copy[i].resize(metadata_recv[rings + 1], metadata_recv[rings + 1], metadata_recv[rings + 2], m.get_block_dimy(), m.get_block_dimx(), 1);\n local_copy_manager[i].set_base_index(metadata_recv[rings + 3]);\n local_copy_manager[i].set_index_range(metadata_recv[rings + 4]);\n local_copy_manager[i].B2L_rings.resize(1);\n local_copy_manager[i].B2L_rings[0].resize(rings + 1);\n local_copy_manager[i].B2L_maps.resize(1);\n local_copy_manager[i].B2L_maps[0].resize(local_copy[i].get_num_rows());\n local_copy_manager[i].L2H_maps.resize(1);\n local_copy_manager[i].L2H_maps[0].resize(metadata_recv[rings + 5]);\n\n for (int j = 0; j <= rings; j++) { local_copy_manager[i].B2L_rings[0][j] = metadata_recv[j]; }\n }\n\n MPI_Waitall(neighbors, &requests[0], MPI_STATUSES_IGNORE); //I have to wait for my stuff to be sent too, because I deallocate those matrices upon exditing this function\n typedef typename T_Config::MatPrec mvalue;\n\n for (int i = 0; i < neighbors; i++)\n {\n MPI_Irecv(local_copy[i].row_offsets.raw(), local_copy[i].row_offsets.size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->neighbors[i] + 0, mpi_comm, &requests[5 * neighbors + 5 * i]);\n MPI_Irecv(local_copy[i].col_indices.raw(), local_copy[i].col_indices.size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->neighbors[i] + 1, mpi_comm, &requests[5 * neighbors + 5 * i + 1]);\n MPI_Irecv(local_copy_manager[i].B2L_maps[0].raw(), local_copy_manager[i].B2L_maps[0].size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->neighbors[i] + 2, mpi_comm, &requests[5 * neighbors + 5 * i + 2]);\n MPI_Irecv(local_copy_manager[i].L2H_maps[0].raw(), local_copy_manager[i].L2H_maps[0].size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->neighbors[i] + 3, mpi_comm, &requests[5 * neighbors + 5 * i + 3]);\n MPI_Irecv(local_copy[i].values.raw(), local_copy[i].values.size()*sizeof(mvalue), MPI_BYTE, m.manager->neighbors[i], 10 * m.manager->neighbors[i] + 4, mpi_comm, &requests[5 * neighbors + 5 * i + 4]);\n }\n\n //Note: GPU Direct uses halo_rows and halo_btl in here:\n for (int i = 0; i < neighbors; i++)\n {\n MPI_Isend(send[i].row_offsets.raw(), send[i].row_offsets.size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->global_id() + 0, mpi_comm, &requests[5 * i]);\n MPI_Isend(send[i].col_indices.raw(), send[i].col_indices.size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->global_id() + 1, mpi_comm, &requests[5 * i + 1]);\n MPI_Isend(send_manager[i].B2L_maps[0].raw(), send_manager[i].B2L_maps[0].size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->global_id() + 2, mpi_comm, &requests[5 * i + 2]);\n MPI_Isend(send_manager[i].L2H_maps[0].raw(), send_manager[i].L2H_maps[0].size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->global_id() + 3, mpi_comm, &requests[5 * i + 3]);\n MPI_Isend(send[i].values.raw(), send[i].values.size()*sizeof(mvalue), MPI_BYTE, m.manager->neighbors[i], 10 * m.manager->global_id() + 4, mpi_comm, &requests[5 * i + 4]);\n }\n\n MPI_Waitall(2 * 5 * neighbors, &requests[0], MPI_STATUSES_IGNORE); //I have to wait for my stuff to be sent too, because I deallocate those matrices upon exditing this function\n\n //Note: GPU Direct swaps here:\n for (int i = 0; i < halo_rows.size(); i++)\n {\n halo_btl[i] = local_copy_manager[i];\n halo_rows[i] = local_copy[i];\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::send_raw_data(const void *ptr, int size, int destination, int tag )\n{\n#ifdef AMGX_WITH_MPI\n MPI_Send(const_cast(ptr), size, MPI_BYTE, destination, tag, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::recv_raw_data(void *ptr, int size, int source, int tag)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Recv(ptr, size, MPI_BYTE, source, tag, mpi_comm, MPI_STATUSES_IGNORE);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::setup(DVector &b, const Matrix &m, int tag, int num_rings) { do_setup(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo(DVector &b, const Matrix &m, int tag, int num_rings) { do_exchange_halo(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_async(DVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream) { do_exchange_halo_async(b, m, event, tag, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_wait(DVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream) { do_exchange_halo_wait(b, m, stream);}\n\ntemplate \nvoid CommsMPIHostBufferStream::setup(FVector &b, const Matrix &m, int tag, int num_rings) { do_setup(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo(FVector &b, const Matrix &m, int tag, int num_rings) { do_exchange_halo(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_async(FVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream) { do_exchange_halo_async(b, m, event, tag, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_wait(FVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream) { do_exchange_halo_wait(b, m, stream);}\n\n\ntemplate \nvoid CommsMPIHostBufferStream::setup(ZVector &b, const Matrix &m, int tag, int num_rings) { do_setup(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo(ZVector &b, const Matrix &m, int tag, int num_rings) { do_exchange_halo(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_async(ZVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream) { do_exchange_halo_async(b, m, event, tag, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_wait(ZVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream) { do_exchange_halo_wait(b, m, stream);}\n\ntemplate \nvoid CommsMPIHostBufferStream::setup(CVector &b, const Matrix &m, int tag, int num_rings) { do_setup(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo(CVector &b, const Matrix &m, int tag, int num_rings) { do_exchange_halo(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_async(CVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream) { do_exchange_halo_async(b, m, event, tag, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_wait(CVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream) { do_exchange_halo_wait(b, m, stream);}\n\ntemplate \nvoid CommsMPIHostBufferStream::setup(IVector &b, const Matrix &m, int tag, int num_rings) { do_setup(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo(IVector &b, const Matrix &m, int tag, int num_rings) { do_exchange_halo(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_async(IVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream) { do_exchange_halo_async(b, m, event, tag, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_wait(IVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream) { do_exchange_halo_wait(b, m, stream);}\n\ntemplate \nvoid CommsMPIHostBufferStream::setup(BVector &b, const Matrix &m, int tag, int num_rings) { do_setup(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo(BVector &b, const Matrix &m, int tag, int num_rings) { FatalError(\"MPI Comms boolean exchange not implemented\", AMGX_ERR_NOT_IMPLEMENTED); /*do_exchange_halo(b, m);*/}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_async(BVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream) { FatalError(\"MPI Comms boolean exchange not implemented\", AMGX_ERR_NOT_IMPLEMENTED);/*do_exchange_halo_async(b,m, event, tag, stream);*/}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_wait(BVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream) { FatalError(\"MPI Comms boolean exchange not implemented\", AMGX_ERR_NOT_IMPLEMENTED);/*do_exchange_halo_wait(b, m, stream);*/}\n\ntemplate \nvoid CommsMPIHostBufferStream::setup(I64Vector &b, const Matrix &m, int tag, int num_rings) { do_setup(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo(I64Vector &b, const Matrix &m, int tag, int num_rings) { do_exchange_halo(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_async(I64Vector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream) { do_exchange_halo_async(b, m, event, tag, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_wait(I64Vector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t stream) { do_exchange_halo_wait(b, m, stream);}\n\ntemplate \nvoid CommsMPIHostBufferStream::setup_L2H(DVector &b, Matrix &m, int num_rings) { do_setup_L2H(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::setup_L2H(IVector &b, Matrix &m, int num_rings) { do_setup_L2H(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::setup_L2H(FVector &b, Matrix &m, int num_rings) { do_setup_L2H(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::setup_L2H(CVector &b, Matrix &m, int num_rings) { do_setup_L2H(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::setup_L2H(ZVector &b, Matrix &m, int num_rings) { do_setup_L2H(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::setup_L2H(BVector &b, Matrix &m, int num_rings) { do_setup_L2H(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::setup_L2H(I64Vector &b, Matrix &m, int num_rings) { do_setup_L2H(b, m, num_rings);}\n\ntemplate \nvoid CommsMPIHostBufferStream::send_receive_wait(DVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t &stream)\n{\n do_send_receive_wait(b, m, event, tag, stream);\n}\ntemplate \nvoid CommsMPIHostBufferStream::send_receive_wait(IVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t &stream)\n{\n do_send_receive_wait(b, m, event, tag, stream);\n}\ntemplate \nvoid CommsMPIHostBufferStream::send_receive_wait(FVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t &stream)\n{\n do_send_receive_wait(b, m, event, tag, stream);\n}\ntemplate \nvoid CommsMPIHostBufferStream::send_receive_wait(CVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t &stream)\n{\n do_send_receive_wait(b, m, event, tag, stream);\n}\ntemplate \nvoid CommsMPIHostBufferStream::send_receive_wait(ZVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t &stream)\n{\n do_send_receive_wait(b, m, event, tag, stream);\n}\ntemplate \nvoid CommsMPIHostBufferStream::send_receive_wait(BVector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t &stream)\n{\n FatalError(\"MPI Comms send-receive with streams is not implemented\", AMGX_ERR_NOT_IMPLEMENTED);\n}\ntemplate \nvoid CommsMPIHostBufferStream::send_receive_wait(I64Vector &b, const Matrix &m, cudaEvent_t event, int tag, cudaStream_t &stream)\n{\n do_send_receive_wait(b, m, event, tag, stream);\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::add_from_halo(DVector &b, const Matrix &m, int tag, int num_rings, cudaStream_t &stream) { do_add_from_halo(b, m, num_rings, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::add_from_halo(FVector &b, const Matrix &m, int tag, int num_rings, cudaStream_t &stream) { do_add_from_halo(b, m, num_rings, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::add_from_halo(CVector &b, const Matrix &m, int tag, int num_rings, cudaStream_t &stream) { do_add_from_halo(b, m, num_rings, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::add_from_halo(ZVector &b, const Matrix &m, int tag, int num_rings, cudaStream_t &stream) { do_add_from_halo(b, m, num_rings, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::add_from_halo(IVector &b, const Matrix &m, int tag, int num_rings, cudaStream_t &stream) { do_add_from_halo(b, m, num_rings, stream);}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H(DVector &b, const Matrix &m, int num_rings, cudaStream_t stream)\n{\n do_gather_L2H(b, m, num_rings, stream);\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H(FVector &b, const Matrix &m, int num_rings, cudaStream_t stream)\n{\n do_gather_L2H(b, m, num_rings, stream);\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H(CVector &b, const Matrix &m, int num_rings, cudaStream_t stream)\n{\n do_gather_L2H(b, m, num_rings, stream);\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H(ZVector &b, const Matrix &m, int num_rings, cudaStream_t stream)\n{\n do_gather_L2H(b, m, num_rings, stream);\n}\n\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H(IVector &b, const Matrix &m, int num_rings, cudaStream_t stream)\n{\n do_gather_L2H(b, m, num_rings, stream);\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H_v2(DVector &b, const Matrix &m, int num_rings, cudaStream_t stream)\n{\n do_gather_L2H_v2(b, m, num_rings, stream);;\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H_v2(FVector &b, const Matrix &m, int num_rings, cudaStream_t stream)\n{\n do_gather_L2H_v2(b, m, num_rings, stream);;\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H_v2(CVector &b, const Matrix &m, int num_rings, cudaStream_t stream)\n{\n do_gather_L2H_v2(b, m, num_rings, stream);;\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H_v2(ZVector &b, const Matrix &m, int num_rings, cudaStream_t stream)\n{\n do_gather_L2H_v2(b, m, num_rings, stream);;\n}\n\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H_v2(IVector &b, const Matrix &m, int num_rings, cudaStream_t stream)\n{\n do_gather_L2H_v2(b, m, num_rings, stream);;\n}\n\n\ntemplate \nbool CommsMPIHostBufferStream::exchange_halo_query(DVector &b, const Matrix &m, cudaEvent_t event) { return do_exchange_halo_query(b, m); }\ntemplate \nbool CommsMPIHostBufferStream::exchange_halo_query(FVector &b, const Matrix &m, cudaEvent_t event) { return do_exchange_halo_query(b, m); }\ntemplate \nbool CommsMPIHostBufferStream::exchange_halo_query(CVector &b, const Matrix &m, cudaEvent_t event) { return do_exchange_halo_query(b, m); }\ntemplate \nbool CommsMPIHostBufferStream::exchange_halo_query(ZVector &b, const Matrix &m, cudaEvent_t event) { return do_exchange_halo_query(b, m); }\ntemplate \nbool CommsMPIHostBufferStream::exchange_halo_query(IVector &b, const Matrix &m, cudaEvent_t event) { return do_exchange_halo_query(b, m); }\ntemplate \nbool CommsMPIHostBufferStream::exchange_halo_query(BVector &b, const Matrix &m, cudaEvent_t event) { return false;}\ntemplate \nbool CommsMPIHostBufferStream::exchange_halo_query(I64Vector &b, const Matrix &m, cudaEvent_t event) { return false;}\n\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce(HDVector_Array &a, HDVector &b, const Operator &m, int tag) { do_reduction(a, b, m);};\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce(HFVector_Array &a, HFVector &b, const Operator &m, int tag) { do_reduction(a, b, m);};\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce(HCVector_Array &a, HCVector &b, const Operator &m, int tag) { do_reduction(a, b, m);};\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce(HZVector_Array &a, HZVector &b, const Operator &m, int tag) { do_reduction(a, b, m);};\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce(HIVector_Array &a, HIVector &b, const Operator &m, int tag) { do_reduction(a, b, m);};\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce(HBVector_Array &a, HBVector &b, const Operator &m, int tag) { FatalError(\"MPI Comms boolean reduction not implemented\", AMGX_ERR_NOT_IMPLEMENTED);/*do_reduction(a,b,m);*/};\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce(HI64Vector_Array &a, HI64Vector &b, const Operator &m, int tag) { do_reduction(a, b, m);};\n\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce_sum(HDVector &a, HDVector &b, const Matrix &m, int tag)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allreduce(&b[0], &a[0], b.size(), MPI_DOUBLE, MPI_SUM, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n};\n\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce_sum(HFVector &a, HFVector &b, const Matrix &m, int tag)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allreduce(&b[0], &a[0], b.size(), MPI_FLOAT, MPI_SUM, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n};\n\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce_sum(HCVector &a, HCVector &b, const Matrix &m, int tag)\n{\n FatalError(\"MPI sum reduce is now supported with complex types\", AMGX_ERR_NOT_IMPLEMENTED);\n};\n\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce_sum(HZVector &a, HZVector &b, const Matrix &m, int tag)\n{\n FatalError(\"MPI sum reduce is now supported with complex types\", AMGX_ERR_NOT_IMPLEMENTED);\n};\n\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce_sum(HIVector &a, HIVector &b, const Matrix &m, int tag)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allreduce(&b[0], &a[0], b.size(), MPI_INT, MPI_SUM, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n};\n\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce_sum(HI64Vector &a, HI64Vector &b, const Matrix &m, int tag)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allreduce(&b[0], &a[0], b.size(), MPI_LONG_LONG_INT, MPI_SUM, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n};\n\ntemplate \nvoid CommsMPIHostBufferStream::exchange_vectors(DVector_Array &a, const Matrix &m, int tag) { do_vec_exchange(a, m);};\ntemplate \nvoid CommsMPIHostBufferStream::exchange_vectors(FVector_Array &a, const Matrix &m, int tag) { do_vec_exchange(a, m);};\ntemplate \nvoid CommsMPIHostBufferStream::exchange_vectors(CVector_Array &a, const Matrix &m, int tag) { do_vec_exchange(a, m);};\ntemplate \nvoid CommsMPIHostBufferStream::exchange_vectors(ZVector_Array &a, const Matrix &m, int tag) { do_vec_exchange(a, m);};\ntemplate \nvoid CommsMPIHostBufferStream::exchange_vectors(IVector_Array &a, const Matrix &m, int tag) { do_vec_exchange(a, m);};\ntemplate \nvoid CommsMPIHostBufferStream::exchange_vectors(BVector_Array &a, const Matrix &m, int tag) { do_vec_exchange(a, m);};\n\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(DIVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(HIVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(DDVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(HDVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(DFVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(HFVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(DCVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(HCVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(DZVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(HZVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\n\n\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(DIVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(HIVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(DDVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(HDVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(DFVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(HFVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(DCVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(HCVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(DZVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(HZVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\n\n\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(DIVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(HIVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(DDVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(HDVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(DFVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(HFVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(DCVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(HCVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(DZVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(HZVector &a) { send_vec_wait_all(a);};\n\n\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(DIVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(HIVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(DDVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(HDVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(DFVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(HFVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(DCVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(HCVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(DZVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(HZVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\n\n\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(DIVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(HIVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(DDVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(HDVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(DFVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(HFVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(DCVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(HCVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(DZVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(HZVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\n\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(DIVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(HIVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(DDVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(HDVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(DFVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(HFVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(DCVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(HCVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(DZVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(HZVector &a) { recv_vec_wait_all(a);};\n\ntemplate \nint CommsMPIHostBufferStream::get_num_partitions()\n{\n#ifdef AMGX_WITH_MPI\n int total = 0;\n MPI_Comm_size( mpi_comm, &total );\n return total;\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nint CommsMPIHostBufferStream::get_global_id()\n{\n#ifdef AMGX_WITH_MPI\n int rank = 0;\n MPI_Comm_rank( mpi_comm, &rank);\n return rank;\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\n\ntemplate \nvoid CommsMPIHostBufferStream::barrier()\n{\n#ifdef AMGX_WITH_MPI\n MPI_Barrier(mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\n#define MAX_HOSTNAME_LENGTH 256\ntemplate \nvoid CommsMPIHostBufferStream::get_hostname(std::string &my_hostname)\n{\n#ifdef AMGX_WITH_MPI\n my_hostname.resize(MAX_HOSTNAME_LENGTH);\n int hostname_length;\n MPI_Get_processor_name(&my_hostname[0], &hostname_length);\n my_hostname.resize(hostname_length);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\n\ntemplate \nvoid CommsMPIHostBufferStream::exchange_hostnames(std::string &my_hostname, std::vector &hostnames, int num_parts)\n{\n#ifdef AMGX_WITH_MPI\n\n if (num_parts != this->get_num_partitions())\n {\n FatalError(\"Inconsistent number of partitions\", AMGX_ERR_CORE);\n }\n\n hostnames.resize(num_parts);\n std::vector hostname_lengths(num_parts);\n std::vector displs(num_parts + 1);\n int my_hostname_length = my_hostname.length();\n MPI_Allgather( &my_hostname_length, 1, MPI_INT, hostname_lengths.data(), 1, MPI_INT, mpi_comm);\n // Compute displacement for each hostname\n displs[0] = 0;\n int totlen = 0;\n\n for (int i = 0; i < num_parts; i++)\n {\n displs[i + 1] = displs[i] + hostname_lengths[i];\n totlen += hostname_lengths[i];\n }\n\n std::string all_names;\n all_names.resize(totlen);\n MPI_Allgatherv( &my_hostname[0], my_hostname_length, MPI_CHAR,\n (void *)(all_names.data()), hostname_lengths.data(), displs.data(), MPI_CHAR,\n mpi_comm);\n\n for (int i = 0; i < num_parts; i++)\n {\n hostnames[i] = std::string(&all_names[displs[i]], hostname_lengths[i]);\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather(const int &my_data, HIVector &gathered_data, int num_parts) { all_gather_templated(my_data, gathered_data, num_parts); }\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather(const int64_t &my_data, HI64Vector &gathered_data, int num_parts) { all_gather_templated(my_data, gathered_data, num_parts); }\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather_v(HIVector &my_data, HIVector &gathered_data, int num_parts) { all_gather_v_templated(my_data[0], my_data.size(), gathered_data, num_parts); }\n\ntemplate \nvoid CommsMPIHostBufferStream::all_reduce_max(IndexType_h &my_data, IndexType_h &result_data)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allreduce(&my_data, &result_data, 1, MPI_INT, MPI_MAX, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::all_gather_templated(const T &my_data, T2 &gathered_data, int num_parts)\n{\n#ifdef AMGX_WITH_MPI\n gathered_data.resize(num_parts);\n MPI_Allgather(&my_data, sizeof(T), MPI_BYTE,\n &gathered_data[0], sizeof(T), MPI_BYTE,\n mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::all_gather_v_templated(T &my_data, int num_elems, T2 &gathered_data, int num_parts)\n{\n#ifdef AMGX_WITH_MPI\n gathered_data.resize(num_parts * num_elems);\n MPI_Allgather(&my_data, num_elems * sizeof(T), MPI_BYTE,\n &gathered_data[0], num_elems * sizeof(T), MPI_BYTE,\n mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather_v(HDVector& data, int num_elems, HDVector& gathered_data, HIVector counts, HIVector displs)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allgatherv(data.raw(), num_elems, MPI_DOUBLE, gathered_data.raw(), counts.raw(), displs.raw(), MPI_DOUBLE, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather_v(HFVector& data, int num_elems, HFVector& gathered_data, HIVector counts, HIVector displs)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allgatherv(data.raw(), num_elems, MPI_FLOAT, gathered_data.raw(), counts.raw(), displs.raw(), MPI_FLOAT, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather_v(HCVector& data, int num_elems, HCVector& gathered_data, HIVector counts, HIVector displs)\n{\n#ifdef AMGX_WITH_MPI\n FatalError(\"AllgatherV with complex data.\", AMGX_ERR_NOT_IMPLEMENTED);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather_v(HZVector& data, int num_elems, HZVector& gathered_data, HIVector counts, HIVector displs)\n{\n#ifdef AMGX_WITH_MPI\n FatalError(\"AllgatherV with complex data.\", AMGX_ERR_NOT_IMPLEMENTED);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather_v(HIVector& data, int num_elems, HIVector& gathered_data, HIVector counts, HIVector displs)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allgatherv(data.raw(), num_elems, MPI_INT, gathered_data.raw(), counts.raw(), displs.raw(), MPI_INT, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class CommsMPIHostBufferStream::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\nnamespace amgx\n{\n\n/***************************************\n * Source Definitions\n ***************************************/\n\ntemplate \n__global__\nvoid copy_buffer(value_type *__restrict__ dest, const value_type *__restrict__ src, int size)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n dest[tidx] = src[tidx];\n }\n}\n\ntemplate \nstatic void error_vector_too_small(const Vector &v, int required_size)\n{\n std::stringstream ss;\n ss << \"Vector size too small: not enough space for halo elements.\" << std::endl;\n ss << \"Vector: {tag = \" << v.tag << \", \" << \"size = \" << v.size() << \"}\" << std::endl;\n ss << \"Required size: \" << required_size << std::endl;\n FatalError(ss.str(), AMGX_ERR_INTERNAL);\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::printString(const std::string &str)\n{\n#ifdef AMGX_WITH_MPI\n int total = 0;\n int part = 0;\n MPI_Comm_size( mpi_comm, &total );\n MPI_Comm_rank( mpi_comm, &part);\n\n for (int i = 0; i < total; i++)\n {\n MPI_Barrier(mpi_comm);\n\n if (i == part)\n {\n std::cout << \"my_id=\" << i << std::endl;\n std::cout << str << std::endl;\n }\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_reduction(T &arrays, T2 &data, const Operator &m)\n{\n#ifdef AMGX_WITH_MPI\n int numRanks = 0;\n MPI_Comm_size( mpi_comm, &numRanks );\n typedef typename T2::value_type value_type;\n std::vector linear_buffer(numRanks * data.size());\n // Allgather() copy data from all ranks into a 1D array (linear_buffer)\n MPI_Allgather(&data[0], sizeof(value_type)*data.size(), MPI_BYTE,\n &linear_buffer[0], sizeof(value_type)*data.size(), MPI_BYTE, mpi_comm);\n // then it is copied from 1D array to the 2D array (called arrays)\n arrays.resize(numRanks);\n\n // first copy data from other ranks\n for (int i = 0; i < numRanks; i++)\n {\n if (i != m.getManager()->global_id())\n {\n arrays[i].assign(linear_buffer.begin() + i * data.size(), linear_buffer.begin() + (i + 1)*data.size());\n }\n }\n\n // then copy data from its won rank\n arrays[m.getManager()->global_id()] = data;\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate\n__global__\nvoid print_data(T **data, int size)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n T *loc = data[0];\n printf(\"data = %d\\n\", loc[tid]);\n}\n\n\ntemplate\n__global__\nvoid print_data_f(T **data, int size)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n T *loc = data[0];\n printf(\"data = %f\\n\", loc[tid]);\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_setup(T &b, const Matrix &m, int num_rings)\n{\n#ifdef AMGX_WITH_MPI\n int bsize = b.get_block_size();\n int num_cols = b.get_num_cols();\n\n if (bsize != 1 && num_cols != 1)\n FatalError(\"Error: vector cannot have block size and subspace size.\",\n AMGX_ERR_INTERNAL);\n\n // set num neighbors = size of b2l_rings\n // need to do this because comms might have more neighbors than our matrix knows about\n set_neighbors(m.manager->B2L_rings.size());\n\n if (b.in_transfer & SENDING)\n {\n b.in_transfer = IDLE;\n }\n\n b.requests.resize(2 * neighbors); //first part is sends second is receives\n b.statuses.resize(2 * neighbors);\n\n for (int i = 0; i < 2 * neighbors; i++)\n {\n b.requests[i] = MPI_REQUEST_NULL;\n }\n\n typedef typename T::value_type vtyp;\n\n if (b.linear_buffers_size < neighbors)\n {\n if (b.linear_buffers_size != 0) { amgx::memory::hipHostFree(b.linear_buffers); }\n cudaCheckError();\n\n amgx::memory::hipHostMalloc((void **) & (b.linear_buffers), neighbors * sizeof(vtyp *));\n cudaCheckError();\n b.linear_buffers_size = neighbors;\n b.linear_buffers_ptrs.resize(neighbors);\n }\n\n cudaCheckError();\n\n int send_size = 0;\n for (int i = 0; i < neighbors; i++)\n {\n send_size += m.manager->B2L_rings[i][num_rings] * bsize * num_cols;\n }\n\n int total_size = send_size + (m.manager->halo_offsets[num_rings * neighbors] - m.manager->halo_offsets[0]) * bsize * num_cols;\n\n b.buffer_size = send_size;\n\n if (b.buffer == NULL)\n {\n b.buffer = new T(total_size);\n }\n else if (total_size > b.buffer->size())\n {\n b.buffer->resize(total_size);\n\n // It is more efficient to synchronise only when linear buffers change\n hipStreamSynchronize(0);\n cudaCheckError();\n }\n\n int offset = 0;\n for (int i = 0; i < neighbors; i++)\n {\n b.linear_buffers[i] = b.buffer->raw() + offset;\n offset += m.manager->B2L_rings[i][num_rings] * bsize * num_cols;\n }\n\n // Copy to device\n {\n hipMemcpy(amgx::thrust::raw_pointer_cast(&b.linear_buffers_ptrs[0]), &(b.linear_buffers[0]), neighbors * sizeof(vtyp *), hipMemcpyDefault);\n cudaCheckError();\n }\n\n if (total_size != 0)\n {\n if (b.explicit_host_buffer == NULL)\n {\n b.host_buffer.resize(1);\n hipEventCreateWithFlags(&b.mpi_event, hipEventDisableTiming);\n amgx::memory::hipHostMalloc((void **)&b.explicit_host_buffer, total_size * sizeof(vtyp));\n cudaCheckError();\n }\n else if (total_size > b.explicit_buffer_size)\n {\n amgx::memory::hipHostFree(b.explicit_host_buffer);\n cudaCheckError();\n amgx::memory::hipHostMalloc((void **)&b.explicit_host_buffer, total_size * sizeof(vtyp));\n }\n\n cudaCheckError();\n b.explicit_buffer_size = total_size;\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_setup_L2H(T &b, Matrix &m, int num_rings)\n{\n#ifdef AMGX_WITH_MPI\n int num_neighbors = m.manager->neighbors.size();\n int bsize = b.get_block_size();\n b.requests.resize(2 * num_neighbors);\n b.statuses.resize(2 * num_neighbors);\n\n for (int i = 0; i < 2 * num_neighbors; i++)\n {\n b.requests[i] = MPI_REQUEST_NULL;\n }\n\n // compute new sizes\n int send_size = (m.manager->halo_offsets[num_rings * num_neighbors] - m.manager->halo_offsets[0]) * bsize;\n int recv_size = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n recv_size += m.manager->B2L_rings[i][num_rings] * bsize;\n }\n\n int size = send_size + recv_size;\n // resize device and host buffers\n b.buffer_size = size;\n\n if (b.buffer == NULL)\n {\n b.buffer = new T(size);\n }\n else if (size > b.buffer->size())\n {\n b.buffer->resize(size);\n hipStreamSynchronize(0);\n cudaCheckError();\n }\n\n b.host_buffer.resize(send_size + recv_size);\n typedef typename T::value_type vtyp;\n\n if (size != 0)\n {\n if (b.explicit_host_buffer == NULL)\n {\n b.host_buffer.resize(1);\n hipEventCreateWithFlags(&b.mpi_event, hipEventDisableTiming);\n amgx::memory::hipHostMalloc((void **)&b.explicit_host_buffer, size * sizeof(vtyp));\n cudaCheckError();\n }\n else if (size > b.explicit_buffer_size)\n {\n amgx::memory::hipHostFree(b.explicit_host_buffer);\n cudaCheckError();\n amgx::memory::hipHostMalloc((void **)&b.explicit_host_buffer, size * sizeof(vtyp));\n }\n\n cudaCheckError();\n b.explicit_buffer_size = size;\n }\n\n // setup linear buffers\n if (b.linear_buffers_size < neighbors)\n {\n if (b.linear_buffers_size != 0) { amgx::memory::hipHostFree(b.linear_buffers); }\n cudaCheckError();\n\n amgx::memory::hipHostMalloc((void **) & (b.linear_buffers), neighbors * sizeof(vtyp *));\n cudaCheckError();\n b.linear_buffers_size = neighbors;\n }\n\n int total = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n b.linear_buffers[i] = b.buffer->raw() + total;\n\n for (int j = 0; j < num_rings; j++)\n {\n total += (m.manager->halo_offsets[j * num_neighbors + i + 1] - m.manager->halo_offsets[j * num_neighbors + i]) * bsize;\n }\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_vec_exchange(std::vector > > &b, const Matrix &m)\n{\n#ifdef AMGX_WITH_MPI\n using VecHost = Vector>;\n using VecType = Vector>;\n std::vector host_send(b.size());\n std::vector host_recv(b.size());\n std::vector send_sizes(b.size());\n std::vector recv_sizes(b.size());\n\n for (int i = 0; i < neighbors; i++)\n {\n send_sizes[i] = b[i].size();\n MPI_Irecv(&recv_sizes[i], 1, MPI_INT, m.manager->neighbors[i], m.manager->neighbors[i], mpi_comm, &requests[neighbors + i]);\n }\n\n for (int i = 0; i < neighbors; i++ )\n {\n MPI_Isend(&send_sizes[i], 1, MPI_INT, m.manager->neighbors[i], m.manager->global_id(), mpi_comm, &requests[i]);\n }\n\n SynchronizerVisitor syncV(Direction::DeviceToHost, b, host_send, recv_sizes, m);\n Accept(syncV);\n MPI_Waitall(2 * neighbors, &requests[0], MPI_STATUSES_IGNORE); //I wait for data to be sent too, as I will deallocate it upon exit\n ReceiverVisitor receiverV(b, host_recv, recv_sizes, m);\n Accept(receiverV);\n SenderVisitor senderV(b, host_send, send_sizes, m);\n Accept(senderV);\n MPI_Waitall(2 * neighbors, &requests[0], MPI_STATUSES_IGNORE); //I wait for data to be sent too, as I will deallocate it upon exit\n syncV.switch_direction();\n syncV.set_local(receiverV.get_local());//only CUDA-aware MPI needs this\n syncV.set_host_vec(host_recv);//only Normal MPI needs this\n Accept(syncV);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::send_vec(T &b, int destination, int tag, int offset, int size)\n{\n#ifdef AMGX_WITH_MPI\n\n if (size == -1) { size = b.size(); }\n\n if (size == 0) { return; }\n\n {\n // Wait for any outstanding receive calls\n int num_recv_requests = b.recv_requests.size();\n\n if (num_recv_requests != 0)\n {\n MPI_Waitall(num_recv_requests, &b.recv_requests[0], MPI_STATUSES_IGNORE);\n b.recv_requests.clear();\n b.recv_statuses.clear();\n }\n }\n\n SynchSendVecVisitor sv(b, destination, tag, offset, size);\n Accept(sv);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::send_vec_async(T &b, int destination, int tag, int offset, int size)\n{\n#ifdef AMGX_WITH_MPI\n\n if (size == -1) { size = b.size(); }\n\n if (size == 0) { return; }\n\n // Wait for any outstanding receive calls\n int num_recv_requests = b.recv_requests.size();\n\n if (num_recv_requests != 0)\n {\n // Wait for any outstanding sends, receives to finish\n MPI_Waitall(num_recv_requests, &b.recv_requests[0], MPI_STATUSES_IGNORE);\n b.recv_requests.clear();\n b.recv_statuses.clear();\n }\n\n MPI_Request new_req;\n MPI_Status new_status;\n b.send_requests.push_back(new_req);\n b.send_statuses.push_back(new_status);\n int num_send_requests = b.send_requests.size();\n AsynchSendVecVisitor sv(b, destination, tag, offset, size);\n Accept(sv);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::send_vec_wait_all(T &b)\n{\n#ifdef AMGX_WITH_MPI\n int num_send_requests = b.send_requests.size();\n MPI_Waitall(num_send_requests, &b.send_requests[0], MPI_STATUSES_IGNORE);\n b.send_requests.clear();\n b.send_statuses.clear();\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::recv_vec(T &b, int source, int tag, int offset, int size)\n{\n#ifdef AMGX_WITH_MPI\n\n if (size == -1) { size = b.size(); }\n\n if (size == 0) { return; }\n\n // Wait for any outstanding sendcalls\n int num_send_requests = b.send_requests.size();\n\n if (num_send_requests != 0)\n {\n // Wait for any outstanding sends, receives to finish\n MPI_Waitall(num_send_requests, &b.send_requests[0], MPI_STATUSES_IGNORE);\n b.send_requests.clear();\n b.send_statuses.clear();\n }\n\n SynchRecvVecVisitor rv(b, source, tag, offset, size);\n Accept(rv);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::recv_vec_async(T &b, int source, int tag, int offset, int size)\n{\n#ifdef AMGX_WITH_MPI\n\n if (size == -1) { size = b.size(); }\n\n if (size == 0) { return; }\n\n // Wait for any outstanding sendcalls\n int num_send_requests = b.send_requests.size();\n\n if (num_send_requests != 0)\n {\n // Wait for any outstanding sends, receives to finish\n MPI_Waitall(num_send_requests, &b.send_requests[0], MPI_STATUSES_IGNORE);\n b.send_requests.clear();\n b.send_statuses.clear();\n }\n\n MPI_Request new_req;\n MPI_Status new_status;\n b.recv_requests.push_back(new_req);\n b.recv_statuses.push_back(new_status);\n int num_recv_requests = b.recv_requests.size();\n AsynchRecvVecVisitor rv(b, source, tag, offset, size);\n Accept(rv);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::recv_vec_wait_all(T &b)\n{\n#ifdef AMGX_WITH_MPI\n int num_recv_requests = b.recv_requests.size();\n MPI_Waitall(num_recv_requests, &b.recv_requests[0], MPI_STATUSES_IGNORE);\n b.recv_requests.clear();\n b.recv_statuses.clear();\n CopyHostFunctor cpyf(b);\n LightVisitor lvcpy(cpyf);\n Accept(lvcpy);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_gather_L2H(T &b, const Matrix &m, int num_rings, hipStream_t stream)\n{\n#ifdef AMGX_WITH_MPI\n // Gather L2H\n int num_neighbors = m.manager->neighbors.size();\n int bsize = b.get_block_size();\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int total = 0;\n\n for (int j = 0; j < num_rings; j++)\n {\n int size = (m.manager->halo_offsets[j * num_neighbors + i + 1] - m.manager->halo_offsets[j * num_neighbors + i]) * bsize;\n\n if (size != 0)\n {\n // we need to use new indices after renumbering - these are stored in L2H_maps\n hipMemcpyAsync(b.linear_buffers[i] + total, b.raw() + m.manager->halo_offsets[j * num_neighbors + i], size*sizeof(typename T::value_type), hipMemcpyDefault, stream);\n cudaCheckError();\n total += size;\n }\n }\n\n cudaCheckError();\n }\n hipStreamSynchronize(stream);\n cudaCheckError();\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_gather_L2H_v2(T &b, const Matrix &m, int num_rings, hipStream_t stream)\n{\n#ifdef AMGX_WITH_MPI\n\n if (num_rings != 1)\n {\n FatalError(\"Halo exchange with ring > 1 not implemented in do_gather_L2H\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // Gather L2H\n int num_neighbors = m.manager->neighbors.size();\n int bsize = b.get_block_size();\n int total = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n total += (m.manager->halo_offsets[i + 1] - m.manager->halo_offsets[i]) * bsize;\n }\n\n if (total != 0)\n {\n hipMemcpyAsync(b.linear_buffers[0], b.raw() + m.manager->halo_offsets[0], total*sizeof(typename T::value_type), hipMemcpyDefault, stream);\n cudaCheckError();\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_add_from_halo(T &b, const Matrix &m, int num_rings, hipStream_t &stream)\n{\n#ifdef AMGX_WITH_MPI\n typedef typename T::value_type vtyp;\n AddFromHalo1Functor ex1(b, m, num_rings, 0, stream);\n AddFromHalo2Functor ex2(b, m, num_rings, 0, stream);\n AddFromHalo3Functor ex3(b, m, num_rings, 0, stream);\n FSMVisitor fsmV;\n fsmV.get_functors().push_back(&ex1);\n fsmV.get_functors().push_back(&ex2);\n fsmV.get_functors().push_back(&ex3);\n Accept(fsmV);\n ex2.get_send_size() = ex1.get_send_size();//pass relevant info from one state to another\n fsmV.next();//advance FSM\n b.in_transfer = RECEIVING | SENDING;\n\n Accept(fsmV);\n //pass relevant info from one state to another:\n //\n ex3.get_recv_size() = ex2.get_offset();\n ex3.get_send_size() = ex2.get_send_size();\n fsmV.next();//advance FSM\n\n // I only wait to receive data, I can start working before all my buffers were sent\n MPI_Waitall(2 * neighbors, &b.requests[0], MPI_STATUSES_IGNORE);\n b.dirtybit = 0;\n Accept(fsmV);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_exchange_halo(T &b, const Matrix &m, int num_rings)\n{\n#ifdef AMGX_WITH_MPI\n typedef typename T::value_type value_type;\n ExcHalo1Functor ex1(b, m, num_rings, 0);\n ExcHalo2Functor ex2(b, m, num_rings, 0);\n ExcHalo3Functor ex3(b, m, num_rings, 0);\n FSMVisitor fsmV;\n fsmV.get_functors().push_back(&ex1);\n fsmV.get_functors().push_back(&ex2);\n fsmV.get_functors().push_back(&ex3);\n Accept(fsmV);\n fsmV.next();//advance FSM\n b.in_transfer = RECEIVING | SENDING;\n Accept(fsmV);\n ex3.get_offset() = ex2.get_offset();//pass relevant info from one state to another\n fsmV.next();//advance FSM\n MPI_Waitall(2 * neighbors, &b.requests[0], MPI_STATUSES_IGNORE); //I only wait to receive data, I can start working before all my buffers were sent\n b.dirtybit = 0;\n //FSM step 3:\n Accept(fsmV);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_exchange_halo_async(T &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream)\n{\n#ifdef AMGX_WITH_MPI\n\n if (tag == -1) { FatalError(\"Untagged vector\\n\", AMGX_ERR_NOT_IMPLEMENTED); }\n\n hipEventSynchronize(event);\n cudaCheckError();\n ExcHalo1Functor ex1(b, m, 1, 0, tag, stream);\n ExcHalo2AsyncFunctor ex2(b, m, 1, 0, tag);\n ExcHalo3AsyncFunctor ex3(b, m, 1, 0, tag, stream);\n FSMVisitor fsmV;\n fsmV.get_functors().push_back(&ex1);\n fsmV.get_functors().push_back(&ex2);\n fsmV.get_functors().push_back(&ex3);\n Accept(fsmV);\n fsmV.next();//advance FSM\n b.in_transfer = RECEIVING | SENDING;\n Accept(fsmV);\n fsmV.next();//advance FSM\n\n if (min_rows_latency_hiding < 0 || m.manager->num_rows_all() < min_rows_latency_hiding)\n {\n MPI_Waitall(2 * neighbors, &b.requests[0], /*&b.statuses[0]*/ MPI_STATUSES_IGNORE); //I only wait to receive data, I can start working before all my buffers were sent\n b.dirtybit = 0;\n Accept(fsmV);\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_exchange_halo_wait(T &b, const Matrix &m, hipStream_t stream)\n{\n#ifdef AMGX_WITH_MPI\n int bsize = b.get_block_size();\n\n if (!(min_rows_latency_hiding < 0 || m.manager->num_rows_all() < min_rows_latency_hiding))\n {\n MPI_Waitall(2 * neighbors, &b.requests[0], MPI_STATUSES_IGNORE);\n b.dirtybit = 0;\n b.in_transfer = IDLE;\n HalloWaitCopyFunctor hcpyf(b, m, stream);\n LightVisitor lvcpy(hcpyf);//Visitor branching off: no-op for GPU Direct\n Accept(lvcpy);\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::do_send_receive_wait(T &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t &stream)\n{\n#ifdef AMGX_WITH_MPI\n\n if (tag == -1)\n {\n FatalError(\"Untagged vector\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n SendRecvWait1Functor ex1(b, m, 1, 0, tag, stream);\n SendRecvWait2Functor ex2(b, m, 1, 0, tag, stream);\n SendRecvWait3Functor ex3(b, m, 1, 0, tag, stream);\n FSMVisitor fsmV;\n fsmV.get_functors().push_back(&ex1);\n fsmV.get_functors().push_back(&ex2);\n fsmV.get_functors().push_back(&ex3);\n hipEventSynchronize(event);\n cudaCheckError();\n Accept(fsmV);\n fsmV.next();//advance FSM\n b.in_transfer = RECEIVING | SENDING;\n Accept(fsmV);\n fsmV.next();//advance FSM\n\n if (neighbors > 0)\n {\n MPI_Waitall(2 * neighbors, &b.requests[0], MPI_STATUSES_IGNORE);\n }\n\n Accept(fsmV);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nbool CommsMPIHostBufferStream::do_exchange_halo_query(T &b, const Matrix &m)\n{\n#ifdef AMGX_WITH_MPI\n\n if (b.dirtybit == 0) { return true; }\n\n if ((b.in_transfer & RECEIVING) == 0) { return false; }\n\n int finished = 1;\n MPI_Status status;\n\n for (int i = 0; i < neighbors; i++)\n {\n int res = MPI_Test(&b.requests[neighbors + i], &finished, &status);\n\n if (finished == 0 || res != MPI_SUCCESS)\n {\n return false;\n }\n }\n\n b.dirtybit = 0;\n b.in_transfer = SENDING;\n HalloWaitCopyFunctor hcpyf(b, m, NULL);\n LightVisitor lvcpy(hcpyf);//Visitor branching off: no-op for GPU Direct\n Accept(lvcpy);\n return true;\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::exchange_matrix_halo(IVector_Array &row_offsets, I64Vector_Array &col_indices, MVector_Array &values, I64Vector_Array &halo_row_ids, IVector_h &neighbors_list, int global_id)\n{\n#ifdef AMGX_WITH_MPI\n int total = 0;\n MPI_Comm_size( mpi_comm, &total );\n int num_neighbors = neighbors_list.size();\n std::vector local_row_offsets(num_neighbors);\n std::vector local_col_indices(num_neighbors);\n std::vector local_values(num_neighbors);\n std::vector send_row_offsets(num_neighbors);\n std::vector send_col_indices(num_neighbors);\n std::vector send_values(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n send_row_offsets[i] = row_offsets[i];\n send_col_indices[i] = col_indices[i];\n send_values[i] = values[i];\n }\n\n std::vector local_row_ids(0);\n std::vector send_row_ids(0);\n\n if (halo_row_ids.size() != 0)\n {\n local_row_ids.resize(num_neighbors);\n send_row_ids.resize(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n send_row_ids[i] = halo_row_ids[i];\n }\n }\n\n // send metadata\n std::vector metadata(num_neighbors * 2); // num_rows+1, num_nz\n\n for (int i = 0; i < num_neighbors; i++)\n {\n metadata[i * 2 + 0] = row_offsets[i].size();\n metadata[i * 2 + 1] = col_indices[i].size();\n MPI_Isend(&metadata[i * 2 + 0], 2, MPI_INT, neighbors_list[i], 0, mpi_comm, &requests[i]);\n }\n\n // receive metadata\n std::vector metadata_recv(2);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n MPI_Recv(&metadata_recv[0], 2, MPI_INT, neighbors_list[i], 0, mpi_comm, MPI_STATUSES_IGNORE);\n local_row_offsets[i].resize(metadata_recv[0]);\n local_col_indices[i].resize(metadata_recv[1]);\n local_values[i].resize(metadata_recv[1]);\n\n if (local_row_ids.size() != 0)\n {\n if (metadata_recv[0] - 1 > 0)\n {\n local_row_ids[i].resize(metadata_recv[0] - 1); // row_ids is one smaller than row_offsets\n }\n }\n }\n\n MPI_Waitall(num_neighbors, &requests[0], MPI_STATUSES_IGNORE); // data is already received, just closing the handles\n // receive matrix data\n typedef typename T_Config::MatPrec mvalue;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n MPI_Irecv(local_row_offsets[i].raw(), local_row_offsets[i].size(), MPI_INT, neighbors_list[i], 10 * neighbors_list[i] + 0, mpi_comm, &requests[3 * num_neighbors + i]);\n MPI_Irecv(local_col_indices[i].raw(), local_col_indices[i].size()*sizeof(int64_t), MPI_BYTE, neighbors_list[i], 10 * neighbors_list[i] + 1, mpi_comm, &requests[4 * num_neighbors + i]);\n MPI_Irecv(local_values[i].raw(), local_values[i].size()*sizeof(mvalue), MPI_BYTE, neighbors_list[i], 10 * neighbors_list[i] + 2, mpi_comm, &requests[5 * num_neighbors + i]);\n\n if (send_row_ids.size() != 0)\n {\n MPI_Irecv(local_row_ids[i].raw(), local_row_ids[i].size()*sizeof(int64_t), MPI_BYTE, neighbors_list[i], 10 * neighbors_list[i] + 3, mpi_comm, &requests[7 * num_neighbors + i]);\n }\n }\n //Note: GPU Direct should use row_offsets[], col_indices[], values[] directly in here:\n // send matrix: row offsets, col indices, values\n for (int i = 0; i < num_neighbors; i++)\n {\n MPI_Isend(send_row_offsets[i].raw(), send_row_offsets[i].size(), MPI_INT, neighbors_list[i], 10 * global_id + 0, mpi_comm, &requests[i]);\n MPI_Isend(send_col_indices[i].raw(), send_col_indices[i].size()*sizeof(int64_t), MPI_BYTE, neighbors_list[i], 10 * global_id + 1, mpi_comm, &requests[num_neighbors + i]);\n MPI_Isend(send_values[i].raw(), send_values[i].size()*sizeof(mvalue), MPI_BYTE, neighbors_list[i], 10 * global_id + 2, mpi_comm, &requests[2 * num_neighbors + i]);\n\n if (send_row_ids.size() != 0)\n {\n MPI_Isend(send_row_ids[i].raw(), send_row_ids[i].size()*sizeof(int64_t), MPI_BYTE, neighbors_list[i], 10 * global_id + 3, mpi_comm, &requests[6 * num_neighbors + i]);\n }\n }\n\n if (halo_row_ids.size() != 0)\n {\n MPI_Waitall(8 * num_neighbors, &requests[0], MPI_STATUSES_IGNORE); //I have to wait for my stuff to be sent too, because I deallocate those matrices upon exditing this function\n }\n else\n {\n MPI_Waitall(6 * num_neighbors, &requests[0], MPI_STATUSES_IGNORE); //I have to wait for my stuff to be sent too, because I deallocate those matrices upon exditing this function\n }\n\n //Note: GPU Direct should swap here\n for (int i = 0; i < num_neighbors; i++)\n {\n row_offsets[i] = local_row_offsets[i];\n col_indices[i] = local_col_indices[i];\n values[i] = local_values[i];\n\n if (halo_row_ids.size() != 0)\n {\n halo_row_ids[i] = local_row_ids[i];\n }\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::exchange_matrix_halo(Matrix_Array &halo_rows, DistributedManager_Array &halo_btl, const Matrix &m)\n{\n#ifdef AMGX_WITH_MPI\n int total = 0;\n MPI_Comm_size( mpi_comm, &total );\n int bsize = m.get_block_size();\n int rings = m.manager->B2L_rings[0].size() - 1;\n int diag = m.hasProps(DIAG);\n std::vector local_copy(halo_rows.size());\n std::vector send(halo_rows.size());\n std::vector local_copy_manager(halo_rows.size());\n std::vector send_manager(halo_rows.size());\n\n //Note: GPU Direct does *NOT* do that:\n for (int i = 0; i < halo_rows.size(); i++)\n {\n send_manager[i] = halo_btl[i];\n send[i] = halo_rows[i];\n }\n\n {\n // there shouldn't be any uncompleted requests, because we don't want to rewrite them\n int completed;\n MPI_Testall(requests.size(), &requests[0], &completed, MPI_STATUSES_IGNORE);\n\n if (!completed)\n {\n MPI_Waitall(requests.size(), &requests[0], MPI_STATUSES_IGNORE);\n }\n }\n\n std::vector metadata(neighbors * (rings + 1 + 5)); //ring offsets (rings+1), num_rows, num_nz, base_index, index_range\n\n for (int i = 0; i < neighbors; i++)\n {\n for (int j = 0; j <= rings; j++) { metadata[i * (rings + 1 + 5) + j] = halo_btl[i].B2L_rings[0][j]; }\n\n metadata[i * (rings + 1 + 5) + rings + 1] = halo_rows[i].get_num_rows();\n metadata[i * (rings + 1 + 5) + rings + 2] = halo_rows[i].get_num_nz();\n metadata[i * (rings + 1 + 5) + rings + 3] = halo_btl[i].base_index();\n metadata[i * (rings + 1 + 5) + rings + 4] = halo_btl[i].index_range();\n metadata[i * (rings + 1 + 5) + rings + 5] = halo_btl[i].L2H_maps[0].size();\n MPI_Isend(&metadata[i * (rings + 1 + 5)], rings + 6, MPI_INT, m.manager->neighbors[i], 0, mpi_comm, &requests[i]);\n }\n\n std::vector metadata_recv(rings + 1 + 5);\n\n for (int i = 0; i < neighbors; i++)\n {\n MPI_Recv(&metadata_recv[0], rings + 6, MPI_INT, m.manager->neighbors[i], 0, mpi_comm, /*&status*/MPI_STATUSES_IGNORE);\n local_copy[i].addProps(CSR);\n\n if (diag) { local_copy[i].addProps(DIAG); }\n\n local_copy[i].resize(metadata_recv[rings + 1], metadata_recv[rings + 1], metadata_recv[rings + 2], m.get_block_dimy(), m.get_block_dimx(), 1);\n local_copy_manager[i].set_base_index(metadata_recv[rings + 3]);\n local_copy_manager[i].set_index_range(metadata_recv[rings + 4]);\n local_copy_manager[i].B2L_rings.resize(1);\n local_copy_manager[i].B2L_rings[0].resize(rings + 1);\n local_copy_manager[i].B2L_maps.resize(1);\n local_copy_manager[i].B2L_maps[0].resize(local_copy[i].get_num_rows());\n local_copy_manager[i].L2H_maps.resize(1);\n local_copy_manager[i].L2H_maps[0].resize(metadata_recv[rings + 5]);\n\n for (int j = 0; j <= rings; j++) { local_copy_manager[i].B2L_rings[0][j] = metadata_recv[j]; }\n }\n\n MPI_Waitall(neighbors, &requests[0], MPI_STATUSES_IGNORE); //I have to wait for my stuff to be sent too, because I deallocate those matrices upon exditing this function\n typedef typename T_Config::MatPrec mvalue;\n\n for (int i = 0; i < neighbors; i++)\n {\n MPI_Irecv(local_copy[i].row_offsets.raw(), local_copy[i].row_offsets.size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->neighbors[i] + 0, mpi_comm, &requests[5 * neighbors + 5 * i]);\n MPI_Irecv(local_copy[i].col_indices.raw(), local_copy[i].col_indices.size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->neighbors[i] + 1, mpi_comm, &requests[5 * neighbors + 5 * i + 1]);\n MPI_Irecv(local_copy_manager[i].B2L_maps[0].raw(), local_copy_manager[i].B2L_maps[0].size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->neighbors[i] + 2, mpi_comm, &requests[5 * neighbors + 5 * i + 2]);\n MPI_Irecv(local_copy_manager[i].L2H_maps[0].raw(), local_copy_manager[i].L2H_maps[0].size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->neighbors[i] + 3, mpi_comm, &requests[5 * neighbors + 5 * i + 3]);\n MPI_Irecv(local_copy[i].values.raw(), local_copy[i].values.size()*sizeof(mvalue), MPI_BYTE, m.manager->neighbors[i], 10 * m.manager->neighbors[i] + 4, mpi_comm, &requests[5 * neighbors + 5 * i + 4]);\n }\n\n //Note: GPU Direct uses halo_rows and halo_btl in here:\n for (int i = 0; i < neighbors; i++)\n {\n MPI_Isend(send[i].row_offsets.raw(), send[i].row_offsets.size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->global_id() + 0, mpi_comm, &requests[5 * i]);\n MPI_Isend(send[i].col_indices.raw(), send[i].col_indices.size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->global_id() + 1, mpi_comm, &requests[5 * i + 1]);\n MPI_Isend(send_manager[i].B2L_maps[0].raw(), send_manager[i].B2L_maps[0].size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->global_id() + 2, mpi_comm, &requests[5 * i + 2]);\n MPI_Isend(send_manager[i].L2H_maps[0].raw(), send_manager[i].L2H_maps[0].size(), MPI_INT, m.manager->neighbors[i], 10 * m.manager->global_id() + 3, mpi_comm, &requests[5 * i + 3]);\n MPI_Isend(send[i].values.raw(), send[i].values.size()*sizeof(mvalue), MPI_BYTE, m.manager->neighbors[i], 10 * m.manager->global_id() + 4, mpi_comm, &requests[5 * i + 4]);\n }\n\n MPI_Waitall(2 * 5 * neighbors, &requests[0], MPI_STATUSES_IGNORE); //I have to wait for my stuff to be sent too, because I deallocate those matrices upon exditing this function\n\n //Note: GPU Direct swaps here:\n for (int i = 0; i < halo_rows.size(); i++)\n {\n halo_btl[i] = local_copy_manager[i];\n halo_rows[i] = local_copy[i];\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::send_raw_data(const void *ptr, int size, int destination, int tag )\n{\n#ifdef AMGX_WITH_MPI\n MPI_Send(const_cast(ptr), size, MPI_BYTE, destination, tag, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::recv_raw_data(void *ptr, int size, int source, int tag)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Recv(ptr, size, MPI_BYTE, source, tag, mpi_comm, MPI_STATUSES_IGNORE);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::setup(DVector &b, const Matrix &m, int tag, int num_rings) { do_setup(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo(DVector &b, const Matrix &m, int tag, int num_rings) { do_exchange_halo(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_async(DVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream) { do_exchange_halo_async(b, m, event, tag, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_wait(DVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream) { do_exchange_halo_wait(b, m, stream);}\n\ntemplate \nvoid CommsMPIHostBufferStream::setup(FVector &b, const Matrix &m, int tag, int num_rings) { do_setup(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo(FVector &b, const Matrix &m, int tag, int num_rings) { do_exchange_halo(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_async(FVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream) { do_exchange_halo_async(b, m, event, tag, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_wait(FVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream) { do_exchange_halo_wait(b, m, stream);}\n\n\ntemplate \nvoid CommsMPIHostBufferStream::setup(ZVector &b, const Matrix &m, int tag, int num_rings) { do_setup(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo(ZVector &b, const Matrix &m, int tag, int num_rings) { do_exchange_halo(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_async(ZVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream) { do_exchange_halo_async(b, m, event, tag, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_wait(ZVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream) { do_exchange_halo_wait(b, m, stream);}\n\ntemplate \nvoid CommsMPIHostBufferStream::setup(CVector &b, const Matrix &m, int tag, int num_rings) { do_setup(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo(CVector &b, const Matrix &m, int tag, int num_rings) { do_exchange_halo(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_async(CVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream) { do_exchange_halo_async(b, m, event, tag, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_wait(CVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream) { do_exchange_halo_wait(b, m, stream);}\n\ntemplate \nvoid CommsMPIHostBufferStream::setup(IVector &b, const Matrix &m, int tag, int num_rings) { do_setup(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo(IVector &b, const Matrix &m, int tag, int num_rings) { do_exchange_halo(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_async(IVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream) { do_exchange_halo_async(b, m, event, tag, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_wait(IVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream) { do_exchange_halo_wait(b, m, stream);}\n\ntemplate \nvoid CommsMPIHostBufferStream::setup(BVector &b, const Matrix &m, int tag, int num_rings) { do_setup(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo(BVector &b, const Matrix &m, int tag, int num_rings) { FatalError(\"MPI Comms boolean exchange not implemented\", AMGX_ERR_NOT_IMPLEMENTED); /*do_exchange_halo(b, m);*/}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_async(BVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream) { FatalError(\"MPI Comms boolean exchange not implemented\", AMGX_ERR_NOT_IMPLEMENTED);/*do_exchange_halo_async(b,m, event, tag, stream);*/}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_wait(BVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream) { FatalError(\"MPI Comms boolean exchange not implemented\", AMGX_ERR_NOT_IMPLEMENTED);/*do_exchange_halo_wait(b, m, stream);*/}\n\ntemplate \nvoid CommsMPIHostBufferStream::setup(I64Vector &b, const Matrix &m, int tag, int num_rings) { do_setup(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo(I64Vector &b, const Matrix &m, int tag, int num_rings) { do_exchange_halo(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_async(I64Vector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream) { do_exchange_halo_async(b, m, event, tag, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::exchange_halo_wait(I64Vector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t stream) { do_exchange_halo_wait(b, m, stream);}\n\ntemplate \nvoid CommsMPIHostBufferStream::setup_L2H(DVector &b, Matrix &m, int num_rings) { do_setup_L2H(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::setup_L2H(IVector &b, Matrix &m, int num_rings) { do_setup_L2H(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::setup_L2H(FVector &b, Matrix &m, int num_rings) { do_setup_L2H(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::setup_L2H(CVector &b, Matrix &m, int num_rings) { do_setup_L2H(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::setup_L2H(ZVector &b, Matrix &m, int num_rings) { do_setup_L2H(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::setup_L2H(BVector &b, Matrix &m, int num_rings) { do_setup_L2H(b, m, num_rings);}\ntemplate \nvoid CommsMPIHostBufferStream::setup_L2H(I64Vector &b, Matrix &m, int num_rings) { do_setup_L2H(b, m, num_rings);}\n\ntemplate \nvoid CommsMPIHostBufferStream::send_receive_wait(DVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t &stream)\n{\n do_send_receive_wait(b, m, event, tag, stream);\n}\ntemplate \nvoid CommsMPIHostBufferStream::send_receive_wait(IVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t &stream)\n{\n do_send_receive_wait(b, m, event, tag, stream);\n}\ntemplate \nvoid CommsMPIHostBufferStream::send_receive_wait(FVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t &stream)\n{\n do_send_receive_wait(b, m, event, tag, stream);\n}\ntemplate \nvoid CommsMPIHostBufferStream::send_receive_wait(CVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t &stream)\n{\n do_send_receive_wait(b, m, event, tag, stream);\n}\ntemplate \nvoid CommsMPIHostBufferStream::send_receive_wait(ZVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t &stream)\n{\n do_send_receive_wait(b, m, event, tag, stream);\n}\ntemplate \nvoid CommsMPIHostBufferStream::send_receive_wait(BVector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t &stream)\n{\n FatalError(\"MPI Comms send-receive with streams is not implemented\", AMGX_ERR_NOT_IMPLEMENTED);\n}\ntemplate \nvoid CommsMPIHostBufferStream::send_receive_wait(I64Vector &b, const Matrix &m, hipEvent_t event, int tag, hipStream_t &stream)\n{\n do_send_receive_wait(b, m, event, tag, stream);\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::add_from_halo(DVector &b, const Matrix &m, int tag, int num_rings, hipStream_t &stream) { do_add_from_halo(b, m, num_rings, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::add_from_halo(FVector &b, const Matrix &m, int tag, int num_rings, hipStream_t &stream) { do_add_from_halo(b, m, num_rings, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::add_from_halo(CVector &b, const Matrix &m, int tag, int num_rings, hipStream_t &stream) { do_add_from_halo(b, m, num_rings, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::add_from_halo(ZVector &b, const Matrix &m, int tag, int num_rings, hipStream_t &stream) { do_add_from_halo(b, m, num_rings, stream);}\ntemplate \nvoid CommsMPIHostBufferStream::add_from_halo(IVector &b, const Matrix &m, int tag, int num_rings, hipStream_t &stream) { do_add_from_halo(b, m, num_rings, stream);}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H(DVector &b, const Matrix &m, int num_rings, hipStream_t stream)\n{\n do_gather_L2H(b, m, num_rings, stream);\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H(FVector &b, const Matrix &m, int num_rings, hipStream_t stream)\n{\n do_gather_L2H(b, m, num_rings, stream);\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H(CVector &b, const Matrix &m, int num_rings, hipStream_t stream)\n{\n do_gather_L2H(b, m, num_rings, stream);\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H(ZVector &b, const Matrix &m, int num_rings, hipStream_t stream)\n{\n do_gather_L2H(b, m, num_rings, stream);\n}\n\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H(IVector &b, const Matrix &m, int num_rings, hipStream_t stream)\n{\n do_gather_L2H(b, m, num_rings, stream);\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H_v2(DVector &b, const Matrix &m, int num_rings, hipStream_t stream)\n{\n do_gather_L2H_v2(b, m, num_rings, stream);;\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H_v2(FVector &b, const Matrix &m, int num_rings, hipStream_t stream)\n{\n do_gather_L2H_v2(b, m, num_rings, stream);;\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H_v2(CVector &b, const Matrix &m, int num_rings, hipStream_t stream)\n{\n do_gather_L2H_v2(b, m, num_rings, stream);;\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H_v2(ZVector &b, const Matrix &m, int num_rings, hipStream_t stream)\n{\n do_gather_L2H_v2(b, m, num_rings, stream);;\n}\n\n\ntemplate \nvoid CommsMPIHostBufferStream::gather_L2H_v2(IVector &b, const Matrix &m, int num_rings, hipStream_t stream)\n{\n do_gather_L2H_v2(b, m, num_rings, stream);;\n}\n\n\ntemplate \nbool CommsMPIHostBufferStream::exchange_halo_query(DVector &b, const Matrix &m, hipEvent_t event) { return do_exchange_halo_query(b, m); }\ntemplate \nbool CommsMPIHostBufferStream::exchange_halo_query(FVector &b, const Matrix &m, hipEvent_t event) { return do_exchange_halo_query(b, m); }\ntemplate \nbool CommsMPIHostBufferStream::exchange_halo_query(CVector &b, const Matrix &m, hipEvent_t event) { return do_exchange_halo_query(b, m); }\ntemplate \nbool CommsMPIHostBufferStream::exchange_halo_query(ZVector &b, const Matrix &m, hipEvent_t event) { return do_exchange_halo_query(b, m); }\ntemplate \nbool CommsMPIHostBufferStream::exchange_halo_query(IVector &b, const Matrix &m, hipEvent_t event) { return do_exchange_halo_query(b, m); }\ntemplate \nbool CommsMPIHostBufferStream::exchange_halo_query(BVector &b, const Matrix &m, hipEvent_t event) { return false;}\ntemplate \nbool CommsMPIHostBufferStream::exchange_halo_query(I64Vector &b, const Matrix &m, hipEvent_t event) { return false;}\n\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce(HDVector_Array &a, HDVector &b, const Operator &m, int tag) { do_reduction(a, b, m);};\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce(HFVector_Array &a, HFVector &b, const Operator &m, int tag) { do_reduction(a, b, m);};\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce(HCVector_Array &a, HCVector &b, const Operator &m, int tag) { do_reduction(a, b, m);};\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce(HZVector_Array &a, HZVector &b, const Operator &m, int tag) { do_reduction(a, b, m);};\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce(HIVector_Array &a, HIVector &b, const Operator &m, int tag) { do_reduction(a, b, m);};\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce(HBVector_Array &a, HBVector &b, const Operator &m, int tag) { FatalError(\"MPI Comms boolean reduction not implemented\", AMGX_ERR_NOT_IMPLEMENTED);/*do_reduction(a,b,m);*/};\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce(HI64Vector_Array &a, HI64Vector &b, const Operator &m, int tag) { do_reduction(a, b, m);};\n\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce_sum(HDVector &a, HDVector &b, const Matrix &m, int tag)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allreduce(&b[0], &a[0], b.size(), MPI_DOUBLE, MPI_SUM, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n};\n\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce_sum(HFVector &a, HFVector &b, const Matrix &m, int tag)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allreduce(&b[0], &a[0], b.size(), MPI_FLOAT, MPI_SUM, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n};\n\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce_sum(HCVector &a, HCVector &b, const Matrix &m, int tag)\n{\n FatalError(\"MPI sum reduce is now supported with complex types\", AMGX_ERR_NOT_IMPLEMENTED);\n};\n\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce_sum(HZVector &a, HZVector &b, const Matrix &m, int tag)\n{\n FatalError(\"MPI sum reduce is now supported with complex types\", AMGX_ERR_NOT_IMPLEMENTED);\n};\n\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce_sum(HIVector &a, HIVector &b, const Matrix &m, int tag)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allreduce(&b[0], &a[0], b.size(), MPI_INT, MPI_SUM, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n};\n\ntemplate \nvoid CommsMPIHostBufferStream::global_reduce_sum(HI64Vector &a, HI64Vector &b, const Matrix &m, int tag)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allreduce(&b[0], &a[0], b.size(), MPI_LONG_LONG_INT, MPI_SUM, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n};\n\ntemplate \nvoid CommsMPIHostBufferStream::exchange_vectors(DVector_Array &a, const Matrix &m, int tag) { do_vec_exchange(a, m);};\ntemplate \nvoid CommsMPIHostBufferStream::exchange_vectors(FVector_Array &a, const Matrix &m, int tag) { do_vec_exchange(a, m);};\ntemplate \nvoid CommsMPIHostBufferStream::exchange_vectors(CVector_Array &a, const Matrix &m, int tag) { do_vec_exchange(a, m);};\ntemplate \nvoid CommsMPIHostBufferStream::exchange_vectors(ZVector_Array &a, const Matrix &m, int tag) { do_vec_exchange(a, m);};\ntemplate \nvoid CommsMPIHostBufferStream::exchange_vectors(IVector_Array &a, const Matrix &m, int tag) { do_vec_exchange(a, m);};\ntemplate \nvoid CommsMPIHostBufferStream::exchange_vectors(BVector_Array &a, const Matrix &m, int tag) { do_vec_exchange(a, m);};\n\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(DIVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(HIVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(DDVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(HDVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(DFVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(HFVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(DCVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(HCVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(DZVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector(HZVector &a, int destination, int tag, int offset, int size) { send_vec(a, destination, tag, offset, size);};\n\n\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(DIVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(HIVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(DDVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(HDVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(DFVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(HFVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(DCVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(HCVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(DZVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_async(HZVector &a, int destination, int tag, int offset, int size) { send_vec_async(a, destination, tag, offset, size);};\n\n\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(DIVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(HIVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(DDVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(HDVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(DFVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(HFVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(DCVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(HCVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(DZVector &a) { send_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::send_vector_wait_all(HZVector &a) { send_vec_wait_all(a);};\n\n\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(DIVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(HIVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(DDVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(HDVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(DFVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(HFVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(DCVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(HCVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(DZVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector(HZVector &a, int source, int tag, int offset, int size) { recv_vec(a, source, tag, offset, size);};\n\n\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(DIVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(HIVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(DDVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(HDVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(DFVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(HFVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(DCVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(HCVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(DZVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_async(HZVector &a, int source, int tag, int offset, int size) { recv_vec_async(a, source, tag, offset, size);};\n\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(DIVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(HIVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(DDVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(HDVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(DFVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(HFVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(DCVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(HCVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(DZVector &a) { recv_vec_wait_all(a);};\ntemplate \nvoid CommsMPIHostBufferStream::recv_vector_wait_all(HZVector &a) { recv_vec_wait_all(a);};\n\ntemplate \nint CommsMPIHostBufferStream::get_num_partitions()\n{\n#ifdef AMGX_WITH_MPI\n int total = 0;\n MPI_Comm_size( mpi_comm, &total );\n return total;\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nint CommsMPIHostBufferStream::get_global_id()\n{\n#ifdef AMGX_WITH_MPI\n int rank = 0;\n MPI_Comm_rank( mpi_comm, &rank);\n return rank;\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\n\ntemplate \nvoid CommsMPIHostBufferStream::barrier()\n{\n#ifdef AMGX_WITH_MPI\n MPI_Barrier(mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\n#define MAX_HOSTNAME_LENGTH 256\ntemplate \nvoid CommsMPIHostBufferStream::get_hostname(std::string &my_hostname)\n{\n#ifdef AMGX_WITH_MPI\n my_hostname.resize(MAX_HOSTNAME_LENGTH);\n int hostname_length;\n MPI_Get_processor_name(&my_hostname[0], &hostname_length);\n my_hostname.resize(hostname_length);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\n\ntemplate \nvoid CommsMPIHostBufferStream::exchange_hostnames(std::string &my_hostname, std::vector &hostnames, int num_parts)\n{\n#ifdef AMGX_WITH_MPI\n\n if (num_parts != this->get_num_partitions())\n {\n FatalError(\"Inconsistent number of partitions\", AMGX_ERR_CORE);\n }\n\n hostnames.resize(num_parts);\n std::vector hostname_lengths(num_parts);\n std::vector displs(num_parts + 1);\n int my_hostname_length = my_hostname.length();\n MPI_Allgather( &my_hostname_length, 1, MPI_INT, hostname_lengths.data(), 1, MPI_INT, mpi_comm);\n // Compute displacement for each hostname\n displs[0] = 0;\n int totlen = 0;\n\n for (int i = 0; i < num_parts; i++)\n {\n displs[i + 1] = displs[i] + hostname_lengths[i];\n totlen += hostname_lengths[i];\n }\n\n std::string all_names;\n all_names.resize(totlen);\n MPI_Allgatherv( &my_hostname[0], my_hostname_length, MPI_CHAR,\n (void *)(all_names.data()), hostname_lengths.data(), displs.data(), MPI_CHAR,\n mpi_comm);\n\n for (int i = 0; i < num_parts; i++)\n {\n hostnames[i] = std::string(&all_names[displs[i]], hostname_lengths[i]);\n }\n\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather(const int &my_data, HIVector &gathered_data, int num_parts) { all_gather_templated(my_data, gathered_data, num_parts); }\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather(const int64_t &my_data, HI64Vector &gathered_data, int num_parts) { all_gather_templated(my_data, gathered_data, num_parts); }\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather_v(HIVector &my_data, HIVector &gathered_data, int num_parts) { all_gather_v_templated(my_data[0], my_data.size(), gathered_data, num_parts); }\n\ntemplate \nvoid CommsMPIHostBufferStream::all_reduce_max(IndexType_h &my_data, IndexType_h &result_data)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allreduce(&my_data, &result_data, 1, MPI_INT, MPI_MAX, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::all_gather_templated(const T &my_data, T2 &gathered_data, int num_parts)\n{\n#ifdef AMGX_WITH_MPI\n gathered_data.resize(num_parts);\n MPI_Allgather(&my_data, sizeof(T), MPI_BYTE,\n &gathered_data[0], sizeof(T), MPI_BYTE,\n mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \ntemplate \nvoid CommsMPIHostBufferStream::all_gather_v_templated(T &my_data, int num_elems, T2 &gathered_data, int num_parts)\n{\n#ifdef AMGX_WITH_MPI\n gathered_data.resize(num_parts * num_elems);\n MPI_Allgather(&my_data, num_elems * sizeof(T), MPI_BYTE,\n &gathered_data[0], num_elems * sizeof(T), MPI_BYTE,\n mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather_v(HDVector& data, int num_elems, HDVector& gathered_data, HIVector counts, HIVector displs)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allgatherv(data.raw(), num_elems, MPI_DOUBLE, gathered_data.raw(), counts.raw(), displs.raw(), MPI_DOUBLE, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather_v(HFVector& data, int num_elems, HFVector& gathered_data, HIVector counts, HIVector displs)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allgatherv(data.raw(), num_elems, MPI_FLOAT, gathered_data.raw(), counts.raw(), displs.raw(), MPI_FLOAT, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather_v(HCVector& data, int num_elems, HCVector& gathered_data, HIVector counts, HIVector displs)\n{\n#ifdef AMGX_WITH_MPI\n FatalError(\"AllgatherV with complex data.\", AMGX_ERR_NOT_IMPLEMENTED);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather_v(HZVector& data, int num_elems, HZVector& gathered_data, HIVector counts, HIVector displs)\n{\n#ifdef AMGX_WITH_MPI\n FatalError(\"AllgatherV with complex data.\", AMGX_ERR_NOT_IMPLEMENTED);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\ntemplate \nvoid CommsMPIHostBufferStream::all_gather_v(HIVector& data, int num_elems, HIVector& gathered_data, HIVector counts, HIVector displs)\n{\n#ifdef AMGX_WITH_MPI\n MPI_Allgatherv(data.raw(), num_elems, MPI_INT, gathered_data.raw(), counts.raw(), displs.raw(), MPI_INT, mpi_comm);\n#else\n FatalError(\"MPI Comms module requires compiling with MPI\", AMGX_ERR_NOT_IMPLEMENTED);\n#endif\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class CommsMPIHostBufferStream::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/distributed/comms_mpi_hostbuffer_stream.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"sm_utils.inl\"\n\nnamespace amgx\n{\nnamespace multicolor_gauss_seidel_solver\n{\n\n// -------------------------\n// Kernels\n// -------------------------\n\ntemplate\n__global__\nvoid setupBlockGSSmooth1x1(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_rows)\n{\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (tid < num_rows)\n {\n //epsilon (val) is a small number with the same sign as val\n ValueTypeA val = values[dia_indices[tid]];\n Dinv[tid] = isNotCloseToZero(val) ? ValueTypeA(1) / val : ValueTypeA(1) / epsilon(val);\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate\n__global__\nvoid setupBlockGSSmooth3by3BlockDiaCsrKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows)\n\n{\n int warp_id = threadIdx.x / 32;\n int warp_thread_id = threadIdx.x & 31;\n\n // padding row blocks to fit in a single warp\n if ( warp_thread_id >= blockrows_per_warp * bsize_sq ) { return; }\n\n // new thread id with padding\n int tid = warp_id * blockrows_per_warp * bsize_sq + warp_thread_id;\n int cta_blockrow_id = tid / bsize_sq;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int mat_entry_index = tid - cta_blockrow_id * bsize_sq;\n const int i_ind = mat_entry_index / bsize;\n const int j_ind = mat_entry_index - i_ind * bsize;\n __shared__ ValueTypeA s_Amat[bsize_sq * blockrows_per_cta ];\n ValueTypeA e_out;\n\n while (blockrow_id < num_block_rows)\n {\n int offset = blockrow_id * bsize_sq + mat_entry_index;\n // Store the diagonal\n e_out = values[bsize_sq * dia_indices[blockrow_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n s_Amat[tid] = e_out;\n compute_block_inverse_row_major\n ( s_Amat, cta_blockrow_id * bsize_sq, offset, i_ind, j_ind, Dinv );\n blockrow_id += gridDim.x * blockrows_per_cta;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockGSSmooth4by4BlockDiaCsrKernel_V2(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows)\n\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int halfwarp_id = tid >> log_bsize_sq;\n const int block_halfwarp_id = threadIdx.x >> log_bsize_sq;\n const int mat_entry_index = threadIdx.x & (bsize_sq - 1);\n const int i_ind = mat_entry_index >> log_bsize;\n const int j_ind = mat_entry_index & (bsize - 1);\n __shared__ ValueTypeA s_Amat[bsize_sq * halfwarps_per_block ];\n ValueTypeA e_out;\n\n while (halfwarp_id < num_block_rows)\n {\n int offset = halfwarp_id * bsize_sq + mat_entry_index;\n // Store the diagonal\n e_out = values[bsize_sq * dia_indices[halfwarp_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n s_Amat[threadIdx.x] = e_out;\n compute_block_inverse_row_major\n ( s_Amat, block_halfwarp_id * bsize_sq, offset, i_ind, j_ind, Dinv );\n halfwarp_id += gridDim.x * blockDim.x >> log_bsize_sq;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockGSSmoothbBybBlockDiaCsrKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows, int bsize, int bsize_sq, ValueTypeA *temp1)\n\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int halfwarp_id = tid >> 4;\n const int block_halfwarp_id = threadIdx.x >> 4;\n const int mat_entry_index = threadIdx.x & (16 - 1);\n const int i_ind = mat_entry_index >> 2;\n const int j_ind = mat_entry_index & 3;\n extern __shared__ char sharedc[];\n ValueTypeA *s_Amat;\n s_Amat = (ValueTypeA *)&sharedc[0];\n int tile_num = (bsize - 1) / 4 + 1;\n ValueTypeA *e_out = &temp1[(blockIdx.x * blockDim.x + threadIdx.x) * tile_num * tile_num];\n\n while (halfwarp_id < num_block_rows)\n {\n int offset = halfwarp_id * bsize_sq + i_ind * bsize + j_ind;\n int s_offset = block_halfwarp_id * bsize_sq;\n\n // Store the diagonal\n for (int t1 = 0; t1 < tile_num; t1++)\n for (int t2 = 0; t2 < tile_num; t2++)\n if ((t1 * 4 + i_ind) < bsize && (t2 * 4 + j_ind) < bsize)\n {\n e_out[t1 * tile_num + t2] = values[bsize_sq * dia_indices[halfwarp_id] + (t1 * 4 + i_ind) * bsize + t2 * 4 + j_ind];\n }\n\n // Each thread stores its entry in s_Amat\n for (int t1 = 0; t1 < tile_num; t1++)\n for (int t2 = 0; t2 < tile_num; t2++)\n if ((t1 * 4 + i_ind) < bsize && (t2 * 4 + j_ind) < bsize)\n {\n s_Amat[s_offset + (t1 * 4 + i_ind) * bsize + t2 * 4 + j_ind] = e_out[t1 * tile_num + t2];\n }\n\n\n compute_block_inverse2\n ( s_Amat, s_offset, offset, i_ind, j_ind, Dinv, tile_num, bsize, bsize_sq );\n halfwarp_id += gridDim.x * halfwarps_per_block;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockGSSmoothBlockDiaCsrKernel_V2(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices,\n ValueTypeA *Dinv, const int num_block_rows)\n{\n int cta_blockrow_id = threadIdx.x / bsize_sq;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int mat_entry_index = threadIdx.x - cta_blockrow_id * bsize_sq;\n const int i_ind = mat_entry_index / bsize;\n const int j_ind = mat_entry_index - i_ind * bsize;\n __shared__ ValueTypeA s_Amat[bsize_sq * blockrows_per_cta];\n int offset, s_offset;\n ValueTypeA e_out;\n\n while (blockrow_id < num_block_rows && cta_blockrow_id < blockrows_per_cta)\n {\n // Store the diagonal\n offset = blockrow_id * bsize_sq + mat_entry_index;\n e_out = values[bsize_sq * dia_indices[blockrow_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n s_Amat[threadIdx.x] = e_out;\n\n s_offset = cta_blockrow_id * bsize_sq;\n// THIS IS DEFINED IN BLOCK COMMON\n//#define s_A(ROW,COL) s_Amat[s_offset+ROW*bsize+COL]\n ValueTypeA diag;\n\n for (int row = 0; row < bsize; row++)\n {\n diag = isNotCloseToZero(s_A(row, row)) ? ValueTypeA(1) / s_A(row, row) : ValueTypeA(1) / epsilon(s_A(row, row));\n //This condition is a subset of the condition below for i_ind==0\n //if ((i_ind == 0) && !(j_ind == row))\n // s_A(row,j_ind) = s_A(row,j_ind)*diag;\n\n if ((i_ind != row) && !(j_ind == row))\n {\n s_A(i_ind, j_ind) = -(s_A(i_ind, row) * s_A(row, j_ind)) + s_A(i_ind, j_ind);\n }\n\n if (i_ind == 0)\n {\n s_A(j_ind, row) = (j_ind == row) ? diag : -(s_A(j_ind, row) * diag);\n }\n }\n\n Dinv[offset] = s_A(i_ind, j_ind);\n blockrow_id += (gridDim.x * blockDim.x) / bsize_sq;\n }\n}\n\n// Kernel to smooth, NAIVE implementation with texture\ntemplate\n__global__\nvoid multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *dia_indices, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int *sorted_rows_by_color, const int num_rows_per_color, const int num_block_rows, ValueTypeB *xout)\n{\n const int vec_entry_index = threadIdx.x & (bsize - 1);\n const int block_eighthwarp_id = threadIdx.x >> log_bsize;\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int eighthwarp_id = tid >> log_bsize;\n volatile __shared__ ValueTypeB s_xtemp[ bsize * eighthwarps_per_block ];\n ValueTypeB bmAx, xin;\n ValueTypeB temp[bsize];\n int offset, i, s_offset;\n\n while (eighthwarp_id < num_rows_per_color)\n {\n i = sorted_rows_by_color[eighthwarp_id];\n offset = i * bsize + vec_entry_index;\n bmAx = b[offset];\n // Contribution from diagonal\n xin = x[offset];\n s_xtemp[threadIdx.x] = xin;\n // Load dia_values and do matrix multiply\n s_offset = block_eighthwarp_id * bsize;\n loadAsVector(nonzero_values + dia_indices[i]*bsize * bsize + vec_entry_index * bsize, temp);\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n\n if (jcol != i)\n {\n offset = jcol * bsize + vec_entry_index;\n s_xtemp[threadIdx.x] = x[offset];\n // Load nonzero_values\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n s_offset = block_eighthwarp_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n }\n }\n\n s_xtemp[threadIdx.x] = bmAx;\n bmAx = 0.;\n // Load Dinv and multiply to RHS\n offset = i * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(Dinv + offset, temp);\n s_offset = block_eighthwarp_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx += temp[m] * s_xtemp[s_offset + m];\n }\n\n xout[i * bsize + vec_entry_index] = xin + weight * bmAx;\n eighthwarp_id += (gridDim.x * blockDim.x >> log_bsize);\n }\n}\n\n// Kernel to smooth, NAIVE implementation with texture\ntemplate\n__global__\nvoid multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *dia_indices, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int *sorted_rows_by_color, const int num_rows_per_color, const int num_block_rows, ValueTypeB *xout)\n{\n //const int vec_entry_index = threadIdx.x & (bsize-1);\n //const int block_eighthwarp_id = threadIdx.x >> log_bsize;\n //const int tid = blockDim.x*blockIdx.x+ threadIdx.x;\n //int eighthwarp_id = tid >> log_bsize;\n //volatile __shared__ ValueType s_xtemp[ bsize*eighthwarps_per_block ];\n int cta_blockrow_id = (threadIdx.x) / bsize;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int vec_entry_index = threadIdx.x - cta_blockrow_id * bsize;\n volatile __shared__ ValueTypeB s_xtemp[ bsize * blockrows_per_cta];\n ValueTypeB bmAx, xin;\n ValueTypeB temp[bsize];\n int offset, i, s_offset;\n\n while (blockrow_id < num_rows_per_color && cta_blockrow_id < blockrows_per_cta)\n {\n i = sorted_rows_by_color[blockrow_id];\n offset = i * bsize + vec_entry_index;\n bmAx = b[offset];\n // Contribution from diagonal\n xin = x[offset];\n s_xtemp[threadIdx.x] = xin;\n // Load dia_values and do matrix multiply\n s_offset = cta_blockrow_id * bsize;\n offset = dia_indices[i] * bsize * bsize + vec_entry_index * bsize;\n#pragma unroll\n\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = nonzero_values[offset + k];\n }\n\n //loadAsVector(dia_values+offset,temp);\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n\n if (jcol != i)\n {\n offset = jcol * bsize + vec_entry_index;\n s_xtemp[threadIdx.x] = x[offset];\n // Load nonzero_values\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n#pragma unroll\n\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = nonzero_values[offset + k];\n }\n\n //loadAsVector(nonzero_values+offset,temp);\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n }\n }\n\n s_xtemp[threadIdx.x] = bmAx;\n bmAx = 0.;\n // Load Dinv and multiply to RHS\n offset = i * bsize * bsize + vec_entry_index * bsize;\n#pragma unroll\n\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = Dinv[offset + k];\n }\n\n //loadAsVector(Dinv+offset,temp);\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx += temp[m] * s_xtemp[s_offset + m];\n }\n\n xout[i * bsize + vec_entry_index] = xin + weight * bmAx;\n blockrow_id += (gridDim.x * blockDim.x) / bsize;\n }\n}\n\n\n\n// Kernel to smooth, NAIVE implementation with texture\n// This is TERRIBLE in terms of memory access pattern\ntemplate\n__global__\nvoid multicolorGSSmoothCsrKernel_NAIVE_tex(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *diag, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int *sorted_rows_by_color, const int num_rows_per_color, const int num_block_rows, ValueTypeB *xout)\n{\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int i;\n ValueTypeB bmAx, xin;\n ValueTypeB temp, s_xtemp;\n ValueTypeA dia;\n\n while (tid < num_rows_per_color)\n {\n i = sorted_rows_by_color[tid];\n bmAx = b[i];\n xin = x[i];\n dia = isNotCloseToZero(nonzero_values[diag[i]]) ? nonzero_values[diag[i]] : epsilon( nonzero_values[diag[i]]);\n bmAx -= xin * dia;\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n\n if (jcol != i)\n {\n s_xtemp = x[jcol];\n temp = nonzero_values[jind];\n bmAx -= temp * s_xtemp;\n }\n }\n\n bmAx /= dia;\n xout[i] = xin + weight * bmAx;\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel to smooth, warp per row version\n// Each warp processes a row; contributions of threads are reduced and stored to x by lane 0\ntemplate\n__global__\nvoid multicolorGSSmoothCsrKernel_WarpPerRow(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *diag, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int *sorted_rows_by_color, const int num_rows_per_color, const int num_block_rows, ValueTypeB *xout)\n{\n int i;\n ValueTypeB bmAx, xin;\n ValueTypeA dia, diatemp;\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n \n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n\n constexpr unsigned full_mask = 0xffffffff;\n \n int row_id = blockIdx.x * NUM_WARPS + warp_id;\n int a_col_tmp;\n\n while (row_id < num_rows_per_color)\n {\n i = sorted_rows_by_color[row_id];\n\n if ( lane_id < 2 )\n {\n a_col_tmp = row_offsets[i + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n xin = x[i];\n bmAx = amgx::types::util::get_zero(); //utils::Ld::load(&b[i]);\n\n for (a_col_it += lane_id; utils::any(a_col_it < a_col_end); a_col_it += WARP_SIZE )\n {\n const bool is_active = a_col_it < a_col_end; \n \n int x_row_id;\n ValueTypeB a_value;\n ValueTypeB x_value;\n\n if (is_active)\n {\n x_row_id = column_indices[a_col_it];\n a_value = nonzero_values[a_col_it];\n x_value = x[x_row_id];\n\n bmAx -= a_value*x_value;\n }\n }\n\n #pragma unroll\n for (int offset = 32 >> 1 ; offset > 0; offset = offset >> 1) \n {\n bmAx += __shfl_down_sync(full_mask, bmAx, offset);\n }\n \n if (lane_id == 0)\n {\n bmAx += b[i];\n diatemp = nonzero_values[diag[i]];\n dia = isNotCloseToZero(diatemp) ? diatemp : epsilon(diatemp);\n bmAx /= dia;\n xout[i] = xin + weight * bmAx;\n }\n\n row_id += gridDim.x * NUM_WARPS;\n }\n}\n\n// Kernel to smooth, N per row version\n// Each row is processed by N threads; contributions of threads are reduced and stored to x by first thread processing each row\ntemplate\n__global__\nvoid multicolorGSSmoothCsrKernel_nPerRow(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *diag, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int *sorted_rows_by_color, const int num_rows_per_color, const int num_block_rows, ValueTypeB *xout)\n{\n constexpr unsigned full_mask = 0xffffffff;\n\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n\n const int me_in_row = lane_id % n_per_row;\n const int batch_row = lane_id / n_per_row;\n\n int row_id = blockIdx.x * CTA_SIZE/n_per_row + (warp_id*WARP_SIZE)/n_per_row + batch_row;\n const int row_id_increment = gridDim.x * blockDim.x/n_per_row;\n\n int i;\n ValueTypeB bmAx, temp, s_xtemp, diatemp;\n ValueTypeA dia;\n IndexType jcol;\n int jit, jmax;\n\n while (row_id < num_rows_per_color)\n {\n i = sorted_rows_by_color[row_id];\n bmAx = amgx::types::util::get_zero();\n jit = row_offsets[i];\n jmax = row_offsets[i + 1]; \n\n for (jit = jit + me_in_row; jit < jmax; jit += n_per_row)\n {\n jcol = column_indices[jit];\n temp = nonzero_values[jit]; \n s_xtemp = x[jcol];\n\n bmAx -= temp * s_xtemp;\n }\n\n #pragma unroll\n for (int offset = n_per_row >> 1 ; offset > 0; offset = offset >> 1) \n {\n bmAx += __shfl_down_sync(full_mask, bmAx, offset);\n }\n\n if (me_in_row == 0)\n {\n diatemp = nonzero_values[diag[i]];\n dia = isNotCloseToZero(diatemp) ? diatemp : epsilon(diatemp);\n bmAx += b[i];\n bmAx /= dia;\n \n xout[i] = x[i] + weight * bmAx;\n }\n\n row_id += row_id_increment;\n }\n}\n\n// Kernel to smooth, NAIVE implementation with texture\n// Batch version\ntemplate\n__global__\nvoid multicolorGSSmoothCsrKernel_NAIVE_tex_batched(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *diag, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv, const int batch_sz,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int *sorted_rows_by_color, const int num_rows_per_color, const int num_block_rows, ValueTypeB *xout)\n{\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int bid = tid / batch_sz;\n int vid = tid % batch_sz;\n int i;\n ValueTypeB bmAx, xin;\n ValueTypeB temp, s_xtemp;\n ValueTypeA dia;\n\n while (bid < num_rows_per_color)\n {\n i = sorted_rows_by_color[bid];\n bmAx = b[batch_sz * i + vid];\n xin = x[batch_sz * i + vid];\n dia = isNotCloseToZero(nonzero_values[diag[i]]) ? nonzero_values[diag[i]] : epsilon( nonzero_values[diag[i]]);\n bmAx -= xin * dia;\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n\n if (jcol != i)\n {\n s_xtemp = x[batch_sz * jcol + vid];\n temp = nonzero_values[jind];\n bmAx -= temp * s_xtemp;\n }\n }\n\n bmAx /= dia;\n xout[batch_sz * i + vid] = xin + weight * bmAx;\n tid += gridDim.x * blockDim.x;\n bid = tid / batch_sz;\n }\n}\n\n// -------------------\n// Methods\n// -------------------\n\n// Constructor\ntemplate\nMulticolorGaussSeidelSolver_Base::MulticolorGaussSeidelSolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope)\n{\n this->weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n this->symFlag = cfg.AMG_Config::template getParameter(\"symmetric_GS\", cfg_scope);\n this->m_reorder_cols_by_color_desired = (cfg.AMG_Config::template getParameter(\"reorder_cols_by_color\", cfg_scope) != 0);\n this->m_insert_diagonal_desired = (cfg.AMG_Config::template getParameter(\"insert_diag_while_reordering\", cfg_scope) != 0);\n\n if (cfg.AMG_Config::template getParameter(\"use_bsrxmv\", cfg_scope))\n {\n this->use_bsrxmv = 1;\n }\n else\n {\n this->use_bsrxmv = 0;\n }\n\n if (this->weight == 0)\n {\n this->weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Multicolor GaussSeidel smoother\\n\");\n }\n\n this->gs_method = KernelMethod::DEFAULT;\n\n cudaEventCreateWithFlags(&m_start, cudaEventDisableTiming);\n cudaCheckError();\n cudaEventCreateWithFlags(&m_end, cudaEventDisableTiming);\n cudaCheckError();\n}\n\n// Destructor\ntemplate\nMulticolorGaussSeidelSolver_Base::~MulticolorGaussSeidelSolver_Base()\n{\n this->Dinv.resize(0);\n}\n\ntemplate\nvoid MulticolorGaussSeidelSolver_Base::computeDinv(Matrix &A)\n{\n ViewType oldView = A.currentView();\n A.setViewExterior();\n\n if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n computeDinv_4x4(A);\n }\n else if (A.get_block_dimx() == 3 && A.get_block_dimy() == 3)\n {\n computeDinv_3x3(A);\n }\n else if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n computeDinv_1x1(A);\n }\n else if (A.get_block_dimx() == 2 && A.get_block_dimy() == 2)\n {\n computeDinv_2x2(A);\n }\n else if (A.get_block_dimx() == 5 && A.get_block_dimy() == 5)\n {\n computeDinv_5x5(A);\n }\n else if (A.get_block_dimx() == A.get_block_dimy())\n {\n computeDinv_bxb(A, A.get_block_dimx());\n }\n else\n {\n FatalError(\"Unsupported block size for MulticolorGaussSeidelSolver computeEinv\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n}\n\n// Method to compute the diagonal matrix E in DILU smoother\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_3x3(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_4x4(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_2x2(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_5x5(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_1x1(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_bxb(const Matrix_h &A, const int bsize)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_4x4(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_cols()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 512;\n const int halfwarps_per_block = threads_per_block / 16;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / halfwarps_per_block + 1);\n cudaFuncSetCacheConfig(setupBlockGSSmooth4by4BlockDiaCsrKernel_V2, cudaFuncCachePreferL1);\n setupBlockGSSmooth4by4BlockDiaCsrKernel_V2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_bxb(const Matrix_d &A, const int bsize)\n{\n//both DIAG supported\n //TODO: check correctness with blocksizes > 5\n this->Dinv.resize(A.get_num_cols()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr= this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n\n MVector temp(AMGX_GRID_MAX_SIZE * ((bsize-1)/4+1)*((bsize-1)/4+1));\n ValueTypeA *temp_ptr= temp.raw();\n\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 512;\n const int halfwarps_per_block = threads_per_block/16;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows()-1)/halfwarps_per_block + 1);\n\n cudaFuncSetCacheConfig(setupBlockGSSmoothbBybBlockDiaCsrKernel,cudaFuncCachePreferL1);\n setupBlockGSSmoothbBybBlockDiaCsrKernel <<>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows(), bsize, bsize*bsize, temp_ptr);\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_3x3(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_cols()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int bsize = 3;\n const int bsize_sq = bsize * bsize;\n const int WARP_SIZE = 32;\n const int threads_per_block = 256;\n const int blockrows_per_warp = WARP_SIZE / bsize_sq;\n const int blockrows_per_cta = (threads_per_block / WARP_SIZE) * blockrows_per_warp;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n cudaFuncSetCacheConfig(setupBlockGSSmooth3by3BlockDiaCsrKernel, cudaFuncCachePreferL1);\n setupBlockGSSmooth3by3BlockDiaCsrKernel <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_2x2(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_cols()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int bsize = 2;\n const int bsize_sq = bsize * bsize;\n const int WARP_SIZE = 32;\n const int threads_per_block = 256;\n const int blockrows_per_warp = WARP_SIZE / bsize_sq;\n const int blockrows_per_cta = (threads_per_block / WARP_SIZE) * blockrows_per_warp;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n cudaFuncSetCacheConfig(setupBlockGSSmooth3by3BlockDiaCsrKernel, cudaFuncCachePreferL1);\n setupBlockGSSmooth3by3BlockDiaCsrKernel <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_5x5(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_cols()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int bsize = 5;\n const int bsize_sq = bsize * bsize;\n const int WARP_SIZE = 32;\n const int threads_per_block = 256;\n const int blockrows_per_warp = WARP_SIZE / bsize_sq;\n const int blockrows_per_cta = (threads_per_block / WARP_SIZE) * blockrows_per_warp;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n cudaFuncSetCacheConfig(setupBlockGSSmooth3by3BlockDiaCsrKernel, cudaFuncCachePreferL1);\n setupBlockGSSmooth3by3BlockDiaCsrKernel <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_1x1(const Matrix_d &A)\n{\n if (this->use_bsrxmv && A.get_block_dimx() == A.get_block_dimy())\n {\n this->Dinv.resize(A.get_num_cols()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const int threads_per_block = 256;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() + threads_per_block - 1) / threads_per_block);\n cudaFuncSetCacheConfig(setupBlockGSSmooth1x1, cudaFuncCachePreferL1);\n setupBlockGSSmooth1x1 <<< num_blocks, threads_per_block>>>(A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), A.diag.raw(), this->Dinv.raw(), A.get_num_rows());\n cudaCheckError();\n }\n}\n\ntemplate\nvoid\nMulticolorGaussSeidelSolver_Base::printSolverParameters() const\n{\n std::cout << \"relaxation_factor= \" << this->weight << std::endl;\n std::cout << \"use_bsrxmv = \" << this->use_bsrxmv << std::endl;\n std::cout << \"symmetric_GS = \" << this->symFlag << std::endl;\n}\n\n// Solver setup\ntemplate\nvoid\nMulticolorGaussSeidelSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n this->m_explicit_A = dynamic_cast*>(this->m_A);\n\n if (!this->m_explicit_A)\n {\n FatalError(\"MulticolorGaussSeidelSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n if ( this->m_explicit_A->getColoringLevel() < 1)\n {\n FatalError(\"Matrix must be colored to use multicolor gauss-seidel solver. Try setting: coloring_level=1 in the configuration file\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if ( this->m_explicit_A->getBlockFormat() != ROW_MAJOR )\n {\n FatalError(\"Multicolor DILU solver only supports row major format for the blocks\", AMGX_ERR_CONFIGURATION);\n }\n\n computeDinv( *this->m_explicit_A );\n}\n\n//\ntemplate\nvoid\nMulticolorGaussSeidelSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_1x1(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_1x1_naive(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_3x3(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_4x4(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_BxB(Matrix_h &A, VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_1x1_naive(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const IndexType *A_diag_ptr = A.diag.raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n const int threads_per_block = 64;\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int threads_per_block = 512;\n //const int eightwarps_per_block = threads_per_block/4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n cudaFuncSetCacheConfig(multicolorGSSmoothCsrKernel_NAIVE_tex, cudaFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n cudaFuncSetCacheConfig(multicolorGSSmoothCsrKernel_NAIVE_tex, cudaFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_1x1(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const IndexType *A_diag_ptr = A.diag.raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n\n cudaStream_t stream = 0; // Stream where previous work is done. Currently - default stream\n cudaStream_t work_stream = stream;\n int use_l2_hint = 0;\n int use_aux_stream = 0; // use separate async stream to do smoothing\n cudaDeviceProp props = getDeviceProperties();\n int arch = 10 * props.major + props.minor;\n\n const int WARP_SIZE = 32;\n const int GRID_SIZE = 1024;\n const int CTA_SIZE = 672;\n \n KernelMethod selected_method = this->gs_method;\n if (this->gs_method == KernelMethod::DEFAULT)\n {\n // Use 4 threads per row\n // If >500K rows, use L2 hint to keep \u2018x\u2019 cached (Ampere and newer)\n // If <500K NNZ per color, use T32_PER_ROW (WARP_PER_ROW might work as well or better)\n // If >20 NNZ per row, use T32_PER_ROW (WARP_PER_ROW might work as well or better)\n selected_method = KernelMethod::T4_PER_ROW;\n if (A.get_num_rows() > 500000) use_l2_hint = 1;\n if (A.get_num_nz() / A.get_num_rows() > 20) selected_method = KernelMethod::T32_PER_ROW;\n if (A.get_num_nz() / this->m_explicit_A->getMatrixColoring().getNumColors() < 500000) selected_method = KernelMethod::T32_PER_ROW;\n }\n\n if (use_l2_hint) use_aux_stream = 1; // l2 hint can be enabled only on async stream\n if (use_aux_stream) work_stream = this->get_aux_stream();\n\n if (use_aux_stream)\n {\n cudaEventRecord(this->m_start, stream);\n cudaCheckError();\n cudaStreamWaitEvent(work_stream, this->m_start, 0); // work_stream to wait for stream \n }\n\n // try to keep 'x' in L2 cache, if at least Ampere & CUDA 11\n#if CUDART_VERSION >= 11000\n cudaStreamAttrValue stream_attribute;\n if (arch >= 80 && use_l2_hint == 1)\n {\n cudaDeviceProp prop = getDeviceProperties();\n size_t x_size = min( A.get_num_rows()*8 , prop.persistingL2CacheMaxSize ); // set-aside length of 'x' (number of rows in A)*8 bytes\n cudaDeviceSetLimit( cudaLimitPersistingL2CacheSize, x_size); // for persisting accesses or the max allowed\n stream_attribute.accessPolicyWindow.base_ptr = reinterpret_cast(x_ptr); \n stream_attribute.accessPolicyWindow.num_bytes = x_size;\n stream_attribute.accessPolicyWindow.hitRatio = 1.0; // Hint for cache hit ratio\n stream_attribute.accessPolicyWindow.hitProp = cudaAccessPropertyPersisting; // Type of access property on cache hit\n stream_attribute.accessPolicyWindow.missProp = cudaAccessPropertyStreaming; // Type of access property on cache miss.\n cudaStreamSetAttribute(work_stream, cudaStreamAttributeAccessPolicyWindow, &stream_attribute); \n }\n#endif\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n if (selected_method == KernelMethod::NAIVE)\n {\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / CTA_SIZE + 1 );\n cudaFuncSetCacheConfig(multicolorGSSmoothCsrKernel_NAIVE_tex, cudaFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n }\n else if (selected_method == KernelMethod::WARP_PER_ROW)\n {\n cudaFuncSetCacheConfig(multicolorGSSmoothCsrKernel_WarpPerRow, cudaFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_WarpPerRow <<< GRID_SIZE, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n }\n else if (selected_method == KernelMethod::T32_PER_ROW)\n {\n const int N_PER_ROW = 32;\n cudaFuncSetCacheConfig(multicolorGSSmoothCsrKernel_nPerRow, cudaFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_nPerRow <<< GRID_SIZE, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr); \n }\n else if (selected_method == KernelMethod::T4_PER_ROW)\n {\n const int N_PER_ROW = 4;\n cudaFuncSetCacheConfig(multicolorGSSmoothCsrKernel_nPerRow, cudaFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_nPerRow <<< GRID_SIZE, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr); \n }\n else\n {\n FatalError(\"Incorrect GS configuration (gs_method value?)\", AMGX_ERR_CONFIGURATION);\n }\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n if (selected_method == KernelMethod::NAIVE)\n {\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / CTA_SIZE + 1 );\n cudaFuncSetCacheConfig(multicolorGSSmoothCsrKernel_NAIVE_tex, cudaFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n }\n else if (selected_method == KernelMethod::WARP_PER_ROW)\n {\n cudaFuncSetCacheConfig(multicolorGSSmoothCsrKernel_WarpPerRow, cudaFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_WarpPerRow <<< GRID_SIZE, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n }\n else if (selected_method == KernelMethod::T32_PER_ROW)\n {\n const int N_PER_ROW = 32;\n cudaFuncSetCacheConfig(multicolorGSSmoothCsrKernel_nPerRow, cudaFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_nPerRow <<< GRID_SIZE, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr); \n }\n else if (selected_method == KernelMethod::T4_PER_ROW)\n {\n const int N_PER_ROW = 4;\n cudaFuncSetCacheConfig(multicolorGSSmoothCsrKernel_nPerRow, cudaFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_nPerRow <<< GRID_SIZE, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr); \n }\n else\n {\n FatalError(\"Incorrect GS configuration (gs_method value?)\", AMGX_ERR_CONFIGURATION);\n }\n cudaCheckError(); \n }\n }\n\n // reset persisting L2 cache\n#if CUDART_VERSION >= 11000\n if (arch >= 80 && use_l2_hint == 1)\n { \n stream_attribute.accessPolicyWindow.num_bytes = 0; // Setting the window size to 0 disable it\n cudaStreamSetAttribute(work_stream, cudaStreamAttributeAccessPolicyWindow, &stream_attribute); // Overwrite the access policy attribute to a CUDA Stream\n cudaCtxResetPersistingL2Cache(); // Remove any persistent lines in L2 \n }\n#endif\n\n if (use_aux_stream)\n {\n cudaEventRecord(this->m_start, work_stream);\n cudaCheckError();\n cudaStreamWaitEvent(stream, this->m_start, 0); // stream to wait for work_stream\n }\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::batch_smooth_1x1(const Matrix_d &A, int batch_sz, const VVector &b, VVector &x)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const IndexType *A_diag_ptr = A.diag.raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const int threads_per_block = 64;\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n\n for (int ivec = 0; ivec < batch_sz; ivec++)\n {\n const ValueTypeB *b_ptr = b.raw() + ivec * A.get_num_rows();\n ValueTypeB *x_ptr = x.raw() + ivec * A.get_num_rows();\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int threads_per_block = 512;\n //const int eightwarps_per_block = threads_per_block/4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n cudaFuncSetCacheConfig(multicolorGSSmoothCsrKernel_NAIVE_tex, cudaFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n cudaFuncSetCacheConfig(multicolorGSSmoothCsrKernel_NAIVE_tex, cudaFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::batch_smooth_1x1_fast(const Matrix_d &A, int batch_sz, const VVector &b, VVector &x)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const IndexType *A_diag_ptr = A.diag.raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const int threads_per_block = 64;\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n const ValueTypeB *b_ptr = b.raw();\n ValueTypeB *x_ptr = x.raw() ;\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int threads_per_block = 512;\n //const int eightwarps_per_block = threads_per_block/4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n cudaFuncSetCacheConfig(multicolorGSSmoothCsrKernel_NAIVE_tex_batched, cudaFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex_batched <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr, batch_sz,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n cudaFuncSetCacheConfig(multicolorGSSmoothCsrKernel_NAIVE_tex_batched, cudaFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex_batched <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr, batch_sz,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_4x4(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int threads_per_block = 512;\n const int eightwarps_per_block = threads_per_block / 4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / eightwarps_per_block + 1 );\n cudaFuncSetCacheConfig(multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2, cudaFuncCachePreferL1);\n multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int threads_per_block = 512;\n const int eightwarps_per_block = threads_per_block / 4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / eightwarps_per_block + 1 );\n cudaFuncSetCacheConfig(multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2, cudaFuncCachePreferL1);\n multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n } // End of loop over colors\n } // End of if symFlag\n\n //cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_3x3(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n const int threads_per_block = 64;\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int threads_per_block = 512;\n const int blockrows_per_cta = threads_per_block / 3;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows_per_color - 1) / blockrows_per_cta + 1);\n cudaFuncSetCacheConfig(multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2, cudaFuncCachePreferL1);\n multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int blockrows_per_cta = threads_per_block / 3;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows_per_color - 1) / blockrows_per_cta + 1);\n cudaFuncSetCacheConfig(multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2, cudaFuncCachePreferL1);\n multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n } // End of loop over colors\n } // End of if symFlag\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_BxB(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flag)\n{\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n int x_dirty_status = x.dirtybit;\n int x_delayed = x.delayed_send;\n int b_delayed = b.delayed_send;\n VVector y(b.size());\n y.set_block_dimx(b.get_block_dimx());\n y.set_block_dimy(b.get_block_dimy());\n y.tag = this->tag * 100 + 3;\n amgx::thrust::copy(b.begin(), b.end(), y.begin());\n cudaCheckError();\n y.dirtybit = 0;\n y.delayed_send = 1;\n x.delayed_send = 1;\n b.delayed_send = 1;\n\n for (int color = 0; color < num_colors; color++)\n {\n Cusparse::bsrmv(Cusparse::ALL_COLORS, color, (ValueTypeA) - 1.0f, A, x, (ValueTypeA)1.0f, y, separation_flag); // y= -A.x + y for current color\n\n if (separation_flag == A.getViewExterior()) { y.dirtybit = 0; }\n\n Cusparse::bsrmv(color, (ValueTypeA)this->weight, A, this->Dinv, y, (ValueTypeA)1.0f, x, separation_flag); // x= w*Dinv.y + x for current color\n\n if (separation_flag == A.getViewExterior()) { x.dirtybit = 0; }\n }\n\n if (this->symFlag)\n {\n amgx::thrust::copy(b.begin(), b.end(), y.begin());\n\n for (int color = num_colors - 1; color >= 0; color--)\n {\n Cusparse::bsrmv(Cusparse::ALL_COLORS, color, (ValueTypeA) - 1.0f, A, x, (ValueTypeA)1.0f, y, separation_flag); // y= -A.x + y for current color\n\n if (separation_flag == A.getViewExterior()) { y.dirtybit = 0; }\n\n Cusparse::bsrmv(color, (ValueTypeA)this->weight, A, this->Dinv, y, (ValueTypeA)1.0f, x, separation_flag); // x= w*Dinv.y + x for current color\n\n if (separation_flag == A.getViewExterior()) { x.dirtybit = 0; }\n }\n\n cudaCheckError();\n }\n\n x.delayed_send = x_delayed;\n b.delayed_send = b_delayed;\n x.dirtybit = x_dirty_status;\n}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nMulticolorGaussSeidelSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n if (xIsZero) { x.dirtybit = 0; }\n\n if (!this->m_explicit_A->is_matrix_singleGPU())\n {\n this->m_explicit_A->manager->exchange_halo_async(x, x.tag);\n this->m_explicit_A->manager->exchange_halo_async(b, b.tag);\n }\n\n if (this->m_explicit_A->getViewExterior() == this->m_explicit_A->getViewInterior())\n {\n if (!this->m_explicit_A->is_matrix_singleGPU())\n {\n this->m_explicit_A->manager->exchange_halo_wait(x, x.tag);\n this->m_explicit_A->manager->exchange_halo_wait(b, b.tag);\n }\n }\n\n ViewType oldView = this->m_explicit_A->currentView();\n bool latencyHiding = true;\n ViewType flags;\n\n if (this->m_explicit_A->is_matrix_singleGPU() || (x.dirtybit == 0 && b.dirtybit == 0))\n {\n latencyHiding = false;\n this->m_explicit_A->setViewExterior();\n flags = this->m_explicit_A->getViewExterior();\n }\n else\n {\n this->m_explicit_A->setViewInterior();\n flags = this->m_explicit_A->getViewInterior();\n }\n\n if (xIsZero)\n {\n thrust_wrapper::fill(x.begin(), x.end(), 0.);\n cudaCheckError();\n }\n\n if ( this->m_explicit_A->get_block_dimx() == 1 && this->m_explicit_A->get_block_dimy() == 1 )\n {\n if (this->use_bsrxmv)\n {\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n }\n else\n {\n smooth_1x1(*this->m_explicit_A, b, x, flags);\n }\n }\n else if ( this->m_explicit_A->get_block_dimx() == 3 && this->m_explicit_A->get_block_dimy() == 3 )\n {\n //if (this->use_bsrxmv)\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n //else\n // smooth_3x3(*this->m_explicit_A, b, x, flags);\n }\n else if ( this->m_explicit_A->get_block_dimx() == 4 && this->m_explicit_A->get_block_dimy() == 4 )\n {\n if (this->use_bsrxmv)\n {\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n }\n else\n {\n smooth_4x4(*this->m_explicit_A, b, x, flags);\n }\n }\n else if (this->m_explicit_A->get_block_dimx() == this->m_explicit_A->get_block_dimy())\n {\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n }\n else\n {\n FatalError(\"Unsupported block size for MulticolorGaussSeidelSolver smooth\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if (latencyHiding)\n {\n if (!this->m_explicit_A->is_matrix_singleGPU())\n {\n this->m_explicit_A->manager->exchange_halo_wait(x, x.tag);\n this->m_explicit_A->manager->exchange_halo_wait(b, b.tag);\n }\n\n this->m_explicit_A->setViewExterior();\n flags = (ViewType)(~(this->m_explicit_A->getViewInterior()) & this->m_explicit_A->getViewExterior());\n\n if (flags != 0)\n {\n if ( this->m_explicit_A->get_block_dimx() == 1 && this->m_explicit_A->get_block_dimy() == 1 )\n {\n if (this->use_bsrxmv)\n {\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n }\n else\n {\n smooth_1x1(*this->m_explicit_A, b, x, flags);\n }\n }\n else if ( this->m_explicit_A->get_block_dimx() == 3 && this->m_explicit_A->get_block_dimy() == 3 )\n {\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n }\n else if ( this->m_explicit_A->get_block_dimx() == 4 && this->m_explicit_A->get_block_dimy() == 4 )\n {\n if (this->use_bsrxmv)\n {\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n }\n else\n {\n smooth_4x4(*this->m_explicit_A, b, x, flags);\n }\n }\n else if (this->m_explicit_A->get_block_dimx() == this->m_explicit_A->get_block_dimy())\n {\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n }\n else\n {\n FatalError(\"Unsupported block size for MulticolorGaussSeidelSolver smooth\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n }\n }\n\n x.dirtybit = 1;\n //if (!this->m_explicit_A->is_matrix_singleGPU() && x.delayed_send==0)\n // this->m_explicit_A->manager->exchange_halo_async(x, x.tag);\n this->m_explicit_A->setView(oldView);\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nMulticolorGaussSeidelSolver_Base::solve_finalize( VVector &b, VVector &x )\n{}\n\ntemplate\ncudaStream_t\nMulticolorGaussSeidelSolver_Base::get_aux_stream()\n{\n static std::shared_ptr aux_stream;\n if (!aux_stream)\n {\n aux_stream = std::make_shared();\n }\n\n return aux_stream->get();\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class MulticolorGaussSeidelSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class MulticolorGaussSeidelSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"sm_utils.inl\"\n\nnamespace amgx\n{\nnamespace multicolor_gauss_seidel_solver\n{\n\n// -------------------------\n// Kernels\n// -------------------------\n\ntemplate\n__global__\nvoid setupBlockGSSmooth1x1(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_rows)\n{\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (tid < num_rows)\n {\n //epsilon (val) is a small number with the same sign as val\n ValueTypeA val = values[dia_indices[tid]];\n Dinv[tid] = isNotCloseToZero(val) ? ValueTypeA(1) / val : ValueTypeA(1) / epsilon(val);\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate\n__global__\nvoid setupBlockGSSmooth3by3BlockDiaCsrKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows)\n\n{\n int warp_id = threadIdx.x / 32;\n int warp_thread_id = threadIdx.x & 31;\n\n // padding row blocks to fit in a single warp\n if ( warp_thread_id >= blockrows_per_warp * bsize_sq ) { return; }\n\n // new thread id with padding\n int tid = warp_id * blockrows_per_warp * bsize_sq + warp_thread_id;\n int cta_blockrow_id = tid / bsize_sq;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int mat_entry_index = tid - cta_blockrow_id * bsize_sq;\n const int i_ind = mat_entry_index / bsize;\n const int j_ind = mat_entry_index - i_ind * bsize;\n __shared__ ValueTypeA s_Amat[bsize_sq * blockrows_per_cta ];\n ValueTypeA e_out;\n\n while (blockrow_id < num_block_rows)\n {\n int offset = blockrow_id * bsize_sq + mat_entry_index;\n // Store the diagonal\n e_out = values[bsize_sq * dia_indices[blockrow_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n s_Amat[tid] = e_out;\n compute_block_inverse_row_major\n ( s_Amat, cta_blockrow_id * bsize_sq, offset, i_ind, j_ind, Dinv );\n blockrow_id += gridDim.x * blockrows_per_cta;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockGSSmooth4by4BlockDiaCsrKernel_V2(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows)\n\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int halfwarp_id = tid >> log_bsize_sq;\n const int block_halfwarp_id = threadIdx.x >> log_bsize_sq;\n const int mat_entry_index = threadIdx.x & (bsize_sq - 1);\n const int i_ind = mat_entry_index >> log_bsize;\n const int j_ind = mat_entry_index & (bsize - 1);\n __shared__ ValueTypeA s_Amat[bsize_sq * halfwarps_per_block ];\n ValueTypeA e_out;\n\n while (halfwarp_id < num_block_rows)\n {\n int offset = halfwarp_id * bsize_sq + mat_entry_index;\n // Store the diagonal\n e_out = values[bsize_sq * dia_indices[halfwarp_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n s_Amat[threadIdx.x] = e_out;\n compute_block_inverse_row_major\n ( s_Amat, block_halfwarp_id * bsize_sq, offset, i_ind, j_ind, Dinv );\n halfwarp_id += gridDim.x * blockDim.x >> log_bsize_sq;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockGSSmoothbBybBlockDiaCsrKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows, int bsize, int bsize_sq, ValueTypeA *temp1)\n\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int halfwarp_id = tid >> 4;\n const int block_halfwarp_id = threadIdx.x >> 4;\n const int mat_entry_index = threadIdx.x & (16 - 1);\n const int i_ind = mat_entry_index >> 2;\n const int j_ind = mat_entry_index & 3;\n extern __shared__ char sharedc[];\n ValueTypeA *s_Amat;\n s_Amat = (ValueTypeA *)&sharedc[0];\n int tile_num = (bsize - 1) / 4 + 1;\n ValueTypeA *e_out = &temp1[(blockIdx.x * blockDim.x + threadIdx.x) * tile_num * tile_num];\n\n while (halfwarp_id < num_block_rows)\n {\n int offset = halfwarp_id * bsize_sq + i_ind * bsize + j_ind;\n int s_offset = block_halfwarp_id * bsize_sq;\n\n // Store the diagonal\n for (int t1 = 0; t1 < tile_num; t1++)\n for (int t2 = 0; t2 < tile_num; t2++)\n if ((t1 * 4 + i_ind) < bsize && (t2 * 4 + j_ind) < bsize)\n {\n e_out[t1 * tile_num + t2] = values[bsize_sq * dia_indices[halfwarp_id] + (t1 * 4 + i_ind) * bsize + t2 * 4 + j_ind];\n }\n\n // Each thread stores its entry in s_Amat\n for (int t1 = 0; t1 < tile_num; t1++)\n for (int t2 = 0; t2 < tile_num; t2++)\n if ((t1 * 4 + i_ind) < bsize && (t2 * 4 + j_ind) < bsize)\n {\n s_Amat[s_offset + (t1 * 4 + i_ind) * bsize + t2 * 4 + j_ind] = e_out[t1 * tile_num + t2];\n }\n\n\n compute_block_inverse2\n ( s_Amat, s_offset, offset, i_ind, j_ind, Dinv, tile_num, bsize, bsize_sq );\n halfwarp_id += gridDim.x * halfwarps_per_block;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockGSSmoothBlockDiaCsrKernel_V2(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices,\n ValueTypeA *Dinv, const int num_block_rows)\n{\n int cta_blockrow_id = threadIdx.x / bsize_sq;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int mat_entry_index = threadIdx.x - cta_blockrow_id * bsize_sq;\n const int i_ind = mat_entry_index / bsize;\n const int j_ind = mat_entry_index - i_ind * bsize;\n __shared__ ValueTypeA s_Amat[bsize_sq * blockrows_per_cta];\n int offset, s_offset;\n ValueTypeA e_out;\n\n while (blockrow_id < num_block_rows && cta_blockrow_id < blockrows_per_cta)\n {\n // Store the diagonal\n offset = blockrow_id * bsize_sq + mat_entry_index;\n e_out = values[bsize_sq * dia_indices[blockrow_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n s_Amat[threadIdx.x] = e_out;\n\n s_offset = cta_blockrow_id * bsize_sq;\n// THIS IS DEFINED IN BLOCK COMMON\n//#define s_A(ROW,COL) s_Amat[s_offset+ROW*bsize+COL]\n ValueTypeA diag;\n\n for (int row = 0; row < bsize; row++)\n {\n diag = isNotCloseToZero(s_A(row, row)) ? ValueTypeA(1) / s_A(row, row) : ValueTypeA(1) / epsilon(s_A(row, row));\n //This condition is a subset of the condition below for i_ind==0\n //if ((i_ind == 0) && !(j_ind == row))\n // s_A(row,j_ind) = s_A(row,j_ind)*diag;\n\n if ((i_ind != row) && !(j_ind == row))\n {\n s_A(i_ind, j_ind) = -(s_A(i_ind, row) * s_A(row, j_ind)) + s_A(i_ind, j_ind);\n }\n\n if (i_ind == 0)\n {\n s_A(j_ind, row) = (j_ind == row) ? diag : -(s_A(j_ind, row) * diag);\n }\n }\n\n Dinv[offset] = s_A(i_ind, j_ind);\n blockrow_id += (gridDim.x * blockDim.x) / bsize_sq;\n }\n}\n\n// Kernel to smooth, NAIVE implementation with texture\ntemplate\n__global__\nvoid multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *dia_indices, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int *sorted_rows_by_color, const int num_rows_per_color, const int num_block_rows, ValueTypeB *xout)\n{\n const int vec_entry_index = threadIdx.x & (bsize - 1);\n const int block_eighthwarp_id = threadIdx.x >> log_bsize;\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int eighthwarp_id = tid >> log_bsize;\n volatile __shared__ ValueTypeB s_xtemp[ bsize * eighthwarps_per_block ];\n ValueTypeB bmAx, xin;\n ValueTypeB temp[bsize];\n int offset, i, s_offset;\n\n while (eighthwarp_id < num_rows_per_color)\n {\n i = sorted_rows_by_color[eighthwarp_id];\n offset = i * bsize + vec_entry_index;\n bmAx = b[offset];\n // Contribution from diagonal\n xin = x[offset];\n s_xtemp[threadIdx.x] = xin;\n // Load dia_values and do matrix multiply\n s_offset = block_eighthwarp_id * bsize;\n loadAsVector(nonzero_values + dia_indices[i]*bsize * bsize + vec_entry_index * bsize, temp);\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n\n if (jcol != i)\n {\n offset = jcol * bsize + vec_entry_index;\n s_xtemp[threadIdx.x] = x[offset];\n // Load nonzero_values\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n s_offset = block_eighthwarp_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n }\n }\n\n s_xtemp[threadIdx.x] = bmAx;\n bmAx = 0.;\n // Load Dinv and multiply to RHS\n offset = i * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(Dinv + offset, temp);\n s_offset = block_eighthwarp_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx += temp[m] * s_xtemp[s_offset + m];\n }\n\n xout[i * bsize + vec_entry_index] = xin + weight * bmAx;\n eighthwarp_id += (gridDim.x * blockDim.x >> log_bsize);\n }\n}\n\n// Kernel to smooth, NAIVE implementation with texture\ntemplate\n__global__\nvoid multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *dia_indices, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int *sorted_rows_by_color, const int num_rows_per_color, const int num_block_rows, ValueTypeB *xout)\n{\n //const int vec_entry_index = threadIdx.x & (bsize-1);\n //const int block_eighthwarp_id = threadIdx.x >> log_bsize;\n //const int tid = blockDim.x*blockIdx.x+ threadIdx.x;\n //int eighthwarp_id = tid >> log_bsize;\n //volatile __shared__ ValueType s_xtemp[ bsize*eighthwarps_per_block ];\n int cta_blockrow_id = (threadIdx.x) / bsize;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int vec_entry_index = threadIdx.x - cta_blockrow_id * bsize;\n volatile __shared__ ValueTypeB s_xtemp[ bsize * blockrows_per_cta];\n ValueTypeB bmAx, xin;\n ValueTypeB temp[bsize];\n int offset, i, s_offset;\n\n while (blockrow_id < num_rows_per_color && cta_blockrow_id < blockrows_per_cta)\n {\n i = sorted_rows_by_color[blockrow_id];\n offset = i * bsize + vec_entry_index;\n bmAx = b[offset];\n // Contribution from diagonal\n xin = x[offset];\n s_xtemp[threadIdx.x] = xin;\n // Load dia_values and do matrix multiply\n s_offset = cta_blockrow_id * bsize;\n offset = dia_indices[i] * bsize * bsize + vec_entry_index * bsize;\n#pragma unroll\n\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = nonzero_values[offset + k];\n }\n\n //loadAsVector(dia_values+offset,temp);\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n\n if (jcol != i)\n {\n offset = jcol * bsize + vec_entry_index;\n s_xtemp[threadIdx.x] = x[offset];\n // Load nonzero_values\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n#pragma unroll\n\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = nonzero_values[offset + k];\n }\n\n //loadAsVector(nonzero_values+offset,temp);\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n }\n }\n\n s_xtemp[threadIdx.x] = bmAx;\n bmAx = 0.;\n // Load Dinv and multiply to RHS\n offset = i * bsize * bsize + vec_entry_index * bsize;\n#pragma unroll\n\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = Dinv[offset + k];\n }\n\n //loadAsVector(Dinv+offset,temp);\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx += temp[m] * s_xtemp[s_offset + m];\n }\n\n xout[i * bsize + vec_entry_index] = xin + weight * bmAx;\n blockrow_id += (gridDim.x * blockDim.x) / bsize;\n }\n}\n\n\n\n// Kernel to smooth, NAIVE implementation with texture\n// This is TERRIBLE in terms of memory access pattern\ntemplate\n__global__\nvoid multicolorGSSmoothCsrKernel_NAIVE_tex(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *diag, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int *sorted_rows_by_color, const int num_rows_per_color, const int num_block_rows, ValueTypeB *xout)\n{\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int i;\n ValueTypeB bmAx, xin;\n ValueTypeB temp, s_xtemp;\n ValueTypeA dia;\n\n while (tid < num_rows_per_color)\n {\n i = sorted_rows_by_color[tid];\n bmAx = b[i];\n xin = x[i];\n dia = isNotCloseToZero(nonzero_values[diag[i]]) ? nonzero_values[diag[i]] : epsilon( nonzero_values[diag[i]]);\n bmAx -= xin * dia;\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n\n if (jcol != i)\n {\n s_xtemp = x[jcol];\n temp = nonzero_values[jind];\n bmAx -= temp * s_xtemp;\n }\n }\n\n bmAx /= dia;\n xout[i] = xin + weight * bmAx;\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel to smooth, warp per row version\n// Each warp processes a row; contributions of threads are reduced and stored to x by lane 0\ntemplate\n__global__\nvoid multicolorGSSmoothCsrKernel_WarpPerRow(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *diag, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int *sorted_rows_by_color, const int num_rows_per_color, const int num_block_rows, ValueTypeB *xout)\n{\n int i;\n ValueTypeB bmAx, xin;\n ValueTypeA dia, diatemp;\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n \n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n\n constexpr unsigned full_mask = 0xffffffff;\n \n int row_id = blockIdx.x * NUM_WARPS + warp_id;\n int a_col_tmp;\n\n while (row_id < num_rows_per_color)\n {\n i = sorted_rows_by_color[row_id];\n\n if ( lane_id < 2 )\n {\n a_col_tmp = row_offsets[i + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n xin = x[i];\n bmAx = amgx::types::util::get_zero(); //utils::Ld::load(&b[i]);\n\n for (a_col_it += lane_id; utils::any(a_col_it < a_col_end); a_col_it += WARP_SIZE )\n {\n const bool is_active = a_col_it < a_col_end; \n \n int x_row_id;\n ValueTypeB a_value;\n ValueTypeB x_value;\n\n if (is_active)\n {\n x_row_id = column_indices[a_col_it];\n a_value = nonzero_values[a_col_it];\n x_value = x[x_row_id];\n\n bmAx -= a_value*x_value;\n }\n }\n\n #pragma unroll\n for (int offset = 32 >> 1 ; offset > 0; offset = offset >> 1) \n {\n bmAx += __shfl_down_sync(full_mask, bmAx, offset);\n }\n \n if (lane_id == 0)\n {\n bmAx += b[i];\n diatemp = nonzero_values[diag[i]];\n dia = isNotCloseToZero(diatemp) ? diatemp : epsilon(diatemp);\n bmAx /= dia;\n xout[i] = xin + weight * bmAx;\n }\n\n row_id += gridDim.x * NUM_WARPS;\n }\n}\n\n// Kernel to smooth, N per row version\n// Each row is processed by N threads; contributions of threads are reduced and stored to x by first thread processing each row\ntemplate\n__global__\nvoid multicolorGSSmoothCsrKernel_nPerRow(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *diag, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int *sorted_rows_by_color, const int num_rows_per_color, const int num_block_rows, ValueTypeB *xout)\n{\n constexpr unsigned full_mask = 0xffffffff;\n\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n\n const int me_in_row = lane_id % n_per_row;\n const int batch_row = lane_id / n_per_row;\n\n int row_id = blockIdx.x * CTA_SIZE/n_per_row + (warp_id*WARP_SIZE)/n_per_row + batch_row;\n const int row_id_increment = gridDim.x * blockDim.x/n_per_row;\n\n int i;\n ValueTypeB bmAx, temp, s_xtemp, diatemp;\n ValueTypeA dia;\n IndexType jcol;\n int jit, jmax;\n\n while (row_id < num_rows_per_color)\n {\n i = sorted_rows_by_color[row_id];\n bmAx = amgx::types::util::get_zero();\n jit = row_offsets[i];\n jmax = row_offsets[i + 1]; \n\n for (jit = jit + me_in_row; jit < jmax; jit += n_per_row)\n {\n jcol = column_indices[jit];\n temp = nonzero_values[jit]; \n s_xtemp = x[jcol];\n\n bmAx -= temp * s_xtemp;\n }\n\n #pragma unroll\n for (int offset = n_per_row >> 1 ; offset > 0; offset = offset >> 1) \n {\n bmAx += __shfl_down_sync(full_mask, bmAx, offset);\n }\n\n if (me_in_row == 0)\n {\n diatemp = nonzero_values[diag[i]];\n dia = isNotCloseToZero(diatemp) ? diatemp : epsilon(diatemp);\n bmAx += b[i];\n bmAx /= dia;\n \n xout[i] = x[i] + weight * bmAx;\n }\n\n row_id += row_id_increment;\n }\n}\n\n// Kernel to smooth, NAIVE implementation with texture\n// Batch version\ntemplate\n__global__\nvoid multicolorGSSmoothCsrKernel_NAIVE_tex_batched(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *diag, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv, const int batch_sz,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int *sorted_rows_by_color, const int num_rows_per_color, const int num_block_rows, ValueTypeB *xout)\n{\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int bid = tid / batch_sz;\n int vid = tid % batch_sz;\n int i;\n ValueTypeB bmAx, xin;\n ValueTypeB temp, s_xtemp;\n ValueTypeA dia;\n\n while (bid < num_rows_per_color)\n {\n i = sorted_rows_by_color[bid];\n bmAx = b[batch_sz * i + vid];\n xin = x[batch_sz * i + vid];\n dia = isNotCloseToZero(nonzero_values[diag[i]]) ? nonzero_values[diag[i]] : epsilon( nonzero_values[diag[i]]);\n bmAx -= xin * dia;\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n\n if (jcol != i)\n {\n s_xtemp = x[batch_sz * jcol + vid];\n temp = nonzero_values[jind];\n bmAx -= temp * s_xtemp;\n }\n }\n\n bmAx /= dia;\n xout[batch_sz * i + vid] = xin + weight * bmAx;\n tid += gridDim.x * blockDim.x;\n bid = tid / batch_sz;\n }\n}\n\n// -------------------\n// Methods\n// -------------------\n\n// Constructor\ntemplate\nMulticolorGaussSeidelSolver_Base::MulticolorGaussSeidelSolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope)\n{\n this->weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n this->symFlag = cfg.AMG_Config::template getParameter(\"symmetric_GS\", cfg_scope);\n this->m_reorder_cols_by_color_desired = (cfg.AMG_Config::template getParameter(\"reorder_cols_by_color\", cfg_scope) != 0);\n this->m_insert_diagonal_desired = (cfg.AMG_Config::template getParameter(\"insert_diag_while_reordering\", cfg_scope) != 0);\n\n if (cfg.AMG_Config::template getParameter(\"use_bsrxmv\", cfg_scope))\n {\n this->use_bsrxmv = 1;\n }\n else\n {\n this->use_bsrxmv = 0;\n }\n\n if (this->weight == 0)\n {\n this->weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Multicolor GaussSeidel smoother\\n\");\n }\n\n this->gs_method = KernelMethod::DEFAULT;\n\n hipEventCreateWithFlags(&m_start, hipEventDisableTiming);\n cudaCheckError();\n hipEventCreateWithFlags(&m_end, hipEventDisableTiming);\n cudaCheckError();\n}\n\n// Destructor\ntemplate\nMulticolorGaussSeidelSolver_Base::~MulticolorGaussSeidelSolver_Base()\n{\n this->Dinv.resize(0);\n}\n\ntemplate\nvoid MulticolorGaussSeidelSolver_Base::computeDinv(Matrix &A)\n{\n ViewType oldView = A.currentView();\n A.setViewExterior();\n\n if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n computeDinv_4x4(A);\n }\n else if (A.get_block_dimx() == 3 && A.get_block_dimy() == 3)\n {\n computeDinv_3x3(A);\n }\n else if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n computeDinv_1x1(A);\n }\n else if (A.get_block_dimx() == 2 && A.get_block_dimy() == 2)\n {\n computeDinv_2x2(A);\n }\n else if (A.get_block_dimx() == 5 && A.get_block_dimy() == 5)\n {\n computeDinv_5x5(A);\n }\n else if (A.get_block_dimx() == A.get_block_dimy())\n {\n computeDinv_bxb(A, A.get_block_dimx());\n }\n else\n {\n FatalError(\"Unsupported block size for MulticolorGaussSeidelSolver computeEinv\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n}\n\n// Method to compute the diagonal matrix E in DILU smoother\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_3x3(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_4x4(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_2x2(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_5x5(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_1x1(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_bxb(const Matrix_h &A, const int bsize)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_4x4(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_cols()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 512;\n const int halfwarps_per_block = threads_per_block / 16;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / halfwarps_per_block + 1);\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockGSSmooth4by4BlockDiaCsrKernel_V2), hipFuncCachePreferL1);\n setupBlockGSSmooth4by4BlockDiaCsrKernel_V2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_bxb(const Matrix_d &A, const int bsize)\n{\n//both DIAG supported\n //TODO: check correctness with blocksizes > 5\n this->Dinv.resize(A.get_num_cols()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr= this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n\n MVector temp(AMGX_GRID_MAX_SIZE * ((bsize-1)/4+1)*((bsize-1)/4+1));\n ValueTypeA *temp_ptr= temp.raw();\n\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 512;\n const int halfwarps_per_block = threads_per_block/16;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows()-1)/halfwarps_per_block + 1);\n\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockGSSmoothbBybBlockDiaCsrKernel),hipFuncCachePreferL1);\n setupBlockGSSmoothbBybBlockDiaCsrKernel <<>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows(), bsize, bsize*bsize, temp_ptr);\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_3x3(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_cols()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int bsize = 3;\n const int bsize_sq = bsize * bsize;\n const int WARP_SIZE = 32;\n const int threads_per_block = 256;\n const int blockrows_per_warp = WARP_SIZE / bsize_sq;\n const int blockrows_per_cta = (threads_per_block / WARP_SIZE) * blockrows_per_warp;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockGSSmooth3by3BlockDiaCsrKernel), hipFuncCachePreferL1);\n setupBlockGSSmooth3by3BlockDiaCsrKernel <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_2x2(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_cols()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int bsize = 2;\n const int bsize_sq = bsize * bsize;\n const int WARP_SIZE = 32;\n const int threads_per_block = 256;\n const int blockrows_per_warp = WARP_SIZE / bsize_sq;\n const int blockrows_per_cta = (threads_per_block / WARP_SIZE) * blockrows_per_warp;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockGSSmooth3by3BlockDiaCsrKernel), hipFuncCachePreferL1);\n setupBlockGSSmooth3by3BlockDiaCsrKernel <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_5x5(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_cols()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int bsize = 5;\n const int bsize_sq = bsize * bsize;\n const int WARP_SIZE = 32;\n const int threads_per_block = 256;\n const int blockrows_per_warp = WARP_SIZE / bsize_sq;\n const int blockrows_per_cta = (threads_per_block / WARP_SIZE) * blockrows_per_warp;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockGSSmooth3by3BlockDiaCsrKernel), hipFuncCachePreferL1);\n setupBlockGSSmooth3by3BlockDiaCsrKernel <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::computeDinv_1x1(const Matrix_d &A)\n{\n if (this->use_bsrxmv && A.get_block_dimx() == A.get_block_dimy())\n {\n this->Dinv.resize(A.get_num_cols()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const int threads_per_block = 256;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() + threads_per_block - 1) / threads_per_block);\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockGSSmooth1x1), hipFuncCachePreferL1);\n setupBlockGSSmooth1x1 <<< num_blocks, threads_per_block>>>(A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), A.diag.raw(), this->Dinv.raw(), A.get_num_rows());\n cudaCheckError();\n }\n}\n\ntemplate\nvoid\nMulticolorGaussSeidelSolver_Base::printSolverParameters() const\n{\n std::cout << \"relaxation_factor= \" << this->weight << std::endl;\n std::cout << \"use_bsrxmv = \" << this->use_bsrxmv << std::endl;\n std::cout << \"symmetric_GS = \" << this->symFlag << std::endl;\n}\n\n// Solver setup\ntemplate\nvoid\nMulticolorGaussSeidelSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n this->m_explicit_A = dynamic_cast*>(this->m_A);\n\n if (!this->m_explicit_A)\n {\n FatalError(\"MulticolorGaussSeidelSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n if ( this->m_explicit_A->getColoringLevel() < 1)\n {\n FatalError(\"Matrix must be colored to use multicolor gauss-seidel solver. Try setting: coloring_level=1 in the configuration file\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if ( this->m_explicit_A->getBlockFormat() != ROW_MAJOR )\n {\n FatalError(\"Multicolor DILU solver only supports row major format for the blocks\", AMGX_ERR_CONFIGURATION);\n }\n\n computeDinv( *this->m_explicit_A );\n}\n\n//\ntemplate\nvoid\nMulticolorGaussSeidelSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_1x1(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_1x1_naive(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_3x3(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_4x4(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_BxB(Matrix_h &A, VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_1x1_naive(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const IndexType *A_diag_ptr = A.diag.raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n const int threads_per_block = 64;\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int threads_per_block = 512;\n //const int eightwarps_per_block = threads_per_block/4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothCsrKernel_NAIVE_tex), hipFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothCsrKernel_NAIVE_tex), hipFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_1x1(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const IndexType *A_diag_ptr = A.diag.raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n\n hipStream_t stream = 0; // Stream where previous work is done. Currently - default stream\n hipStream_t work_stream = stream;\n int use_l2_hint = 0;\n int use_aux_stream = 0; // use separate async stream to do smoothing\n hipDeviceProp_t props = getDeviceProperties();\n int arch = 10 * props.major + props.minor;\n\n const int WARP_SIZE = 32;\n const int GRID_SIZE = 1024;\n const int CTA_SIZE = 672;\n \n KernelMethod selected_method = this->gs_method;\n if (this->gs_method == KernelMethod::DEFAULT)\n {\n // Use 4 threads per row\n // If >500K rows, use L2 hint to keep \u2018x\u2019 cached (Ampere and newer)\n // If <500K NNZ per color, use T32_PER_ROW (WARP_PER_ROW might work as well or better)\n // If >20 NNZ per row, use T32_PER_ROW (WARP_PER_ROW might work as well or better)\n selected_method = KernelMethod::T4_PER_ROW;\n if (A.get_num_rows() > 500000) use_l2_hint = 1;\n if (A.get_num_nz() / A.get_num_rows() > 20) selected_method = KernelMethod::T32_PER_ROW;\n if (A.get_num_nz() / this->m_explicit_A->getMatrixColoring().getNumColors() < 500000) selected_method = KernelMethod::T32_PER_ROW;\n }\n\n if (use_l2_hint) use_aux_stream = 1; // l2 hint can be enabled only on async stream\n if (use_aux_stream) work_stream = this->get_aux_stream();\n\n if (use_aux_stream)\n {\n hipEventRecord(this->m_start, stream);\n cudaCheckError();\n hipStreamWaitEvent(work_stream, this->m_start, 0); // work_stream to wait for stream \n }\n\n // try to keep 'x' in L2 cache, if at least Ampere & CUDA 11\n#if CUDART_VERSION >= 11000\n hipLaunchAttributeValue stream_attribute;\n if (arch >= 80 && use_l2_hint == 1)\n {\n hipDeviceProp_t prop = getDeviceProperties();\n size_t x_size = min( A.get_num_rows()*8 , prop.persistingL2CacheMaxSize ); // set-aside length of 'x' (number of rows in A)*8 bytes\n hipDeviceSetLimit( cudaLimitPersistingL2CacheSize, x_size); // for persisting accesses or the max allowed\n stream_attribute.accessPolicyWindow.base_ptr = reinterpret_cast(x_ptr); \n stream_attribute.accessPolicyWindow.num_bytes = x_size;\n stream_attribute.accessPolicyWindow.hitRatio = 1.0; // Hint for cache hit ratio\n stream_attribute.accessPolicyWindow.hitProp = hipAccessPropertyPersisting; // Type of access property on cache hit\n stream_attribute.accessPolicyWindow.missProp = hipAccessPropertyStreaming; // Type of access property on cache miss.\n hipStreamSetAttribute(work_stream, hipLaunchAttributeAccessPolicyWindow, &stream_attribute); \n }\n#endif\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n if (selected_method == KernelMethod::NAIVE)\n {\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / CTA_SIZE + 1 );\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothCsrKernel_NAIVE_tex), hipFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n }\n else if (selected_method == KernelMethod::WARP_PER_ROW)\n {\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothCsrKernel_WarpPerRow), hipFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_WarpPerRow <<< GRID_SIZE, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n }\n else if (selected_method == KernelMethod::T32_PER_ROW)\n {\n const int N_PER_ROW = 32;\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothCsrKernel_nPerRow), hipFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_nPerRow <<< GRID_SIZE, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr); \n }\n else if (selected_method == KernelMethod::T4_PER_ROW)\n {\n const int N_PER_ROW = 4;\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothCsrKernel_nPerRow), hipFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_nPerRow <<< GRID_SIZE, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr); \n }\n else\n {\n FatalError(\"Incorrect GS configuration (gs_method value?)\", AMGX_ERR_CONFIGURATION);\n }\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n if (selected_method == KernelMethod::NAIVE)\n {\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / CTA_SIZE + 1 );\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothCsrKernel_NAIVE_tex), hipFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n }\n else if (selected_method == KernelMethod::WARP_PER_ROW)\n {\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothCsrKernel_WarpPerRow), hipFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_WarpPerRow <<< GRID_SIZE, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n }\n else if (selected_method == KernelMethod::T32_PER_ROW)\n {\n const int N_PER_ROW = 32;\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothCsrKernel_nPerRow), hipFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_nPerRow <<< GRID_SIZE, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr); \n }\n else if (selected_method == KernelMethod::T4_PER_ROW)\n {\n const int N_PER_ROW = 4;\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothCsrKernel_nPerRow), hipFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_nPerRow <<< GRID_SIZE, CTA_SIZE, 0, work_stream>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr); \n }\n else\n {\n FatalError(\"Incorrect GS configuration (gs_method value?)\", AMGX_ERR_CONFIGURATION);\n }\n cudaCheckError(); \n }\n }\n\n // reset persisting L2 cache\n#if CUDART_VERSION >= 11000\n if (arch >= 80 && use_l2_hint == 1)\n { \n stream_attribute.accessPolicyWindow.num_bytes = 0; // Setting the window size to 0 disable it\n hipStreamSetAttribute(work_stream, hipLaunchAttributeAccessPolicyWindow, &stream_attribute); // Overwrite the access policy attribute to a CUDA Stream\n cudaCtxResetPersistingL2Cache(); // Remove any persistent lines in L2 \n }\n#endif\n\n if (use_aux_stream)\n {\n hipEventRecord(this->m_start, work_stream);\n cudaCheckError();\n hipStreamWaitEvent(stream, this->m_start, 0); // stream to wait for work_stream\n }\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::batch_smooth_1x1(const Matrix_d &A, int batch_sz, const VVector &b, VVector &x)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const IndexType *A_diag_ptr = A.diag.raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const int threads_per_block = 64;\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n\n for (int ivec = 0; ivec < batch_sz; ivec++)\n {\n const ValueTypeB *b_ptr = b.raw() + ivec * A.get_num_rows();\n ValueTypeB *x_ptr = x.raw() + ivec * A.get_num_rows();\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int threads_per_block = 512;\n //const int eightwarps_per_block = threads_per_block/4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothCsrKernel_NAIVE_tex), hipFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothCsrKernel_NAIVE_tex), hipFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::batch_smooth_1x1_fast(const Matrix_d &A, int batch_sz, const VVector &b, VVector &x)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const IndexType *A_diag_ptr = A.diag.raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const int threads_per_block = 64;\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n const ValueTypeB *b_ptr = b.raw();\n ValueTypeB *x_ptr = x.raw() ;\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int threads_per_block = 512;\n //const int eightwarps_per_block = threads_per_block/4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothCsrKernel_NAIVE_tex_batched), hipFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex_batched <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr, batch_sz,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothCsrKernel_NAIVE_tex_batched), hipFuncCachePreferL1);\n multicolorGSSmoothCsrKernel_NAIVE_tex_batched <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr, batch_sz,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_4x4(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int threads_per_block = 512;\n const int eightwarps_per_block = threads_per_block / 4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / eightwarps_per_block + 1 );\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2), hipFuncCachePreferL1);\n multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int threads_per_block = 512;\n const int eightwarps_per_block = threads_per_block / 4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / eightwarps_per_block + 1 );\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2), hipFuncCachePreferL1);\n multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n } // End of loop over colors\n } // End of if symFlag\n\n //cudaCheckError();\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_3x3(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n const int threads_per_block = 64;\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int threads_per_block = 512;\n const int blockrows_per_cta = threads_per_block / 3;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows_per_color - 1) / blockrows_per_cta + 1);\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2), hipFuncCachePreferL1);\n multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int blockrows_per_cta = threads_per_block / 3;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows_per_color - 1) / blockrows_per_cta + 1);\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2), hipFuncCachePreferL1);\n multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n } // End of loop over colors\n } // End of if symFlag\n}\n\ntemplate \nvoid MulticolorGaussSeidelSolver >::smooth_BxB(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flag)\n{\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n int x_dirty_status = x.dirtybit;\n int x_delayed = x.delayed_send;\n int b_delayed = b.delayed_send;\n VVector y(b.size());\n y.set_block_dimx(b.get_block_dimx());\n y.set_block_dimy(b.get_block_dimy());\n y.tag = this->tag * 100 + 3;\n amgx::thrust::copy(b.begin(), b.end(), y.begin());\n cudaCheckError();\n y.dirtybit = 0;\n y.delayed_send = 1;\n x.delayed_send = 1;\n b.delayed_send = 1;\n\n for (int color = 0; color < num_colors; color++)\n {\n Cusparse::bsrmv(Cusparse::ALL_COLORS, color, (ValueTypeA) - 1.0f, A, x, (ValueTypeA)1.0f, y, separation_flag); // y= -A.x + y for current color\n\n if (separation_flag == A.getViewExterior()) { y.dirtybit = 0; }\n\n Cusparse::bsrmv(color, (ValueTypeA)this->weight, A, this->Dinv, y, (ValueTypeA)1.0f, x, separation_flag); // x= w*Dinv.y + x for current color\n\n if (separation_flag == A.getViewExterior()) { x.dirtybit = 0; }\n }\n\n if (this->symFlag)\n {\n amgx::thrust::copy(b.begin(), b.end(), y.begin());\n\n for (int color = num_colors - 1; color >= 0; color--)\n {\n Cusparse::bsrmv(Cusparse::ALL_COLORS, color, (ValueTypeA) - 1.0f, A, x, (ValueTypeA)1.0f, y, separation_flag); // y= -A.x + y for current color\n\n if (separation_flag == A.getViewExterior()) { y.dirtybit = 0; }\n\n Cusparse::bsrmv(color, (ValueTypeA)this->weight, A, this->Dinv, y, (ValueTypeA)1.0f, x, separation_flag); // x= w*Dinv.y + x for current color\n\n if (separation_flag == A.getViewExterior()) { x.dirtybit = 0; }\n }\n\n cudaCheckError();\n }\n\n x.delayed_send = x_delayed;\n b.delayed_send = b_delayed;\n x.dirtybit = x_dirty_status;\n}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nMulticolorGaussSeidelSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n if (xIsZero) { x.dirtybit = 0; }\n\n if (!this->m_explicit_A->is_matrix_singleGPU())\n {\n this->m_explicit_A->manager->exchange_halo_async(x, x.tag);\n this->m_explicit_A->manager->exchange_halo_async(b, b.tag);\n }\n\n if (this->m_explicit_A->getViewExterior() == this->m_explicit_A->getViewInterior())\n {\n if (!this->m_explicit_A->is_matrix_singleGPU())\n {\n this->m_explicit_A->manager->exchange_halo_wait(x, x.tag);\n this->m_explicit_A->manager->exchange_halo_wait(b, b.tag);\n }\n }\n\n ViewType oldView = this->m_explicit_A->currentView();\n bool latencyHiding = true;\n ViewType flags;\n\n if (this->m_explicit_A->is_matrix_singleGPU() || (x.dirtybit == 0 && b.dirtybit == 0))\n {\n latencyHiding = false;\n this->m_explicit_A->setViewExterior();\n flags = this->m_explicit_A->getViewExterior();\n }\n else\n {\n this->m_explicit_A->setViewInterior();\n flags = this->m_explicit_A->getViewInterior();\n }\n\n if (xIsZero)\n {\n thrust_wrapper::fill(x.begin(), x.end(), 0.);\n cudaCheckError();\n }\n\n if ( this->m_explicit_A->get_block_dimx() == 1 && this->m_explicit_A->get_block_dimy() == 1 )\n {\n if (this->use_bsrxmv)\n {\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n }\n else\n {\n smooth_1x1(*this->m_explicit_A, b, x, flags);\n }\n }\n else if ( this->m_explicit_A->get_block_dimx() == 3 && this->m_explicit_A->get_block_dimy() == 3 )\n {\n //if (this->use_bsrxmv)\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n //else\n // smooth_3x3(*this->m_explicit_A, b, x, flags);\n }\n else if ( this->m_explicit_A->get_block_dimx() == 4 && this->m_explicit_A->get_block_dimy() == 4 )\n {\n if (this->use_bsrxmv)\n {\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n }\n else\n {\n smooth_4x4(*this->m_explicit_A, b, x, flags);\n }\n }\n else if (this->m_explicit_A->get_block_dimx() == this->m_explicit_A->get_block_dimy())\n {\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n }\n else\n {\n FatalError(\"Unsupported block size for MulticolorGaussSeidelSolver smooth\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if (latencyHiding)\n {\n if (!this->m_explicit_A->is_matrix_singleGPU())\n {\n this->m_explicit_A->manager->exchange_halo_wait(x, x.tag);\n this->m_explicit_A->manager->exchange_halo_wait(b, b.tag);\n }\n\n this->m_explicit_A->setViewExterior();\n flags = (ViewType)(~(this->m_explicit_A->getViewInterior()) & this->m_explicit_A->getViewExterior());\n\n if (flags != 0)\n {\n if ( this->m_explicit_A->get_block_dimx() == 1 && this->m_explicit_A->get_block_dimy() == 1 )\n {\n if (this->use_bsrxmv)\n {\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n }\n else\n {\n smooth_1x1(*this->m_explicit_A, b, x, flags);\n }\n }\n else if ( this->m_explicit_A->get_block_dimx() == 3 && this->m_explicit_A->get_block_dimy() == 3 )\n {\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n }\n else if ( this->m_explicit_A->get_block_dimx() == 4 && this->m_explicit_A->get_block_dimy() == 4 )\n {\n if (this->use_bsrxmv)\n {\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n }\n else\n {\n smooth_4x4(*this->m_explicit_A, b, x, flags);\n }\n }\n else if (this->m_explicit_A->get_block_dimx() == this->m_explicit_A->get_block_dimy())\n {\n smooth_BxB(*this->m_explicit_A, b, x, flags);\n }\n else\n {\n FatalError(\"Unsupported block size for MulticolorGaussSeidelSolver smooth\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n }\n }\n\n x.dirtybit = 1;\n //if (!this->m_explicit_A->is_matrix_singleGPU() && x.delayed_send==0)\n // this->m_explicit_A->manager->exchange_halo_async(x, x.tag);\n this->m_explicit_A->setView(oldView);\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nMulticolorGaussSeidelSolver_Base::solve_finalize( VVector &b, VVector &x )\n{}\n\ntemplate\nhipStream_t\nMulticolorGaussSeidelSolver_Base::get_aux_stream()\n{\n static std::shared_ptr aux_stream;\n if (!aux_stream)\n {\n aux_stream = std::make_shared();\n }\n\n return aux_stream->get();\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class MulticolorGaussSeidelSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class MulticolorGaussSeidelSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/solvers/multicolor_gauss_seidel_solver.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\nnamespace fixcolor_gauss_seidel_solver\n{\n\n// -------------------------\n// Kernels\n// -------------------------\n\ntemplate\n__global__\nvoid setupBlockGSSmooth3by3BlockDiaCsrKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows)\n{\n int warp_id = threadIdx.x / 32;\n int warp_thread_id = threadIdx.x & 31;\n\n // padding row blocks to fit in a single warp\n if ( warp_thread_id >= blockrows_per_warp * bsize_sq ) { return; }\n\n // new thread id with padding\n int tid = warp_id * blockrows_per_warp * bsize_sq + warp_thread_id;\n int cta_blockrow_id = tid / bsize_sq;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int mat_entry_index = tid - cta_blockrow_id * bsize_sq;\n const int i_ind = mat_entry_index / bsize;\n const int j_ind = mat_entry_index - i_ind * bsize;\n volatile __shared__ ValueTypeA s_Amat[bsize_sq * blockrows_per_cta ];\n ValueTypeA e_out;\n\n while (blockrow_id < num_block_rows)\n {\n int offset = blockrow_id * bsize_sq + mat_entry_index;\n // Store the diagonal\n e_out = values[bsize_sq * dia_indices[blockrow_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n s_Amat[tid] = e_out;\n compute_block_inverse_row_major\n ( s_Amat, cta_blockrow_id * bsize_sq, offset, i_ind, j_ind, Dinv );\n blockrow_id += gridDim.x * blockrows_per_cta;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockGSSmooth4by4BlockDiaCsrKernel_V2(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows)\n\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int halfwarp_id = tid >> log_bsize_sq;\n const int block_halfwarp_id = threadIdx.x >> log_bsize_sq;\n const int mat_entry_index = threadIdx.x & (bsize_sq - 1);\n const int i_ind = mat_entry_index >> log_bsize;\n const int j_ind = mat_entry_index & (bsize - 1);\n volatile __shared__ ValueTypeA s_Amat[bsize_sq * halfwarps_per_block ];\n ValueTypeA e_out;\n\n while (halfwarp_id < num_block_rows)\n {\n int offset = halfwarp_id * bsize_sq + mat_entry_index;\n // Store the diagonal\n e_out = values[bsize_sq * dia_indices[halfwarp_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n s_Amat[threadIdx.x] = e_out;\n compute_block_inverse_row_major\n ( s_Amat, block_halfwarp_id * bsize_sq, offset, i_ind, j_ind, Dinv );\n halfwarp_id += gridDim.x * blockDim.x >> log_bsize_sq;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockGSSmoothbBybBlockDiaCsrKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows, int bsize, int bsize_sq, ValueTypeA *temp1)\n\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int halfwarp_id = tid >> 4;\n const int block_halfwarp_id = threadIdx.x >> 4;\n const int mat_entry_index = threadIdx.x & (16 - 1);\n const int i_ind = mat_entry_index >> 2;\n const int j_ind = mat_entry_index & 3;\n extern __shared__ volatile char sharedc[];\n volatile ValueTypeA *s_Amat;\n s_Amat = (ValueTypeA *)&sharedc[0];\n int tile_num = (bsize - 1) / 4 + 1;\n ValueTypeA *e_out = &temp1[(blockIdx.x * blockDim.x + threadIdx.x) * tile_num * tile_num];\n\n while (halfwarp_id < num_block_rows)\n {\n int offset = halfwarp_id * bsize_sq + i_ind * bsize + j_ind;\n int s_offset = block_halfwarp_id * bsize_sq;\n\n // Store the diagonal\n for (int t1 = 0; t1 < tile_num; t1++)\n for (int t2 = 0; t2 < tile_num; t2++)\n if ((t1 * 4 + i_ind) < bsize && (t2 * 4 + j_ind) < bsize)\n {\n e_out[t1 * tile_num + t2] = values[bsize_sq * dia_indices[halfwarp_id] + (t1 * 4 + i_ind) * bsize + t2 * 4 + j_ind];\n }\n\n // Each thread stores its entry in s_Amat\n for (int t1 = 0; t1 < tile_num; t1++)\n for (int t2 = 0; t2 < tile_num; t2++)\n if ((t1 * 4 + i_ind) < bsize && (t2 * 4 + j_ind) < bsize)\n {\n s_Amat[s_offset + (t1 * 4 + i_ind) * bsize + t2 * 4 + j_ind] = e_out[t1 * tile_num + t2];\n }\n\n compute_block_inverse2\n ( s_Amat, s_offset, offset, i_ind, j_ind, Dinv, tile_num, bsize, bsize_sq );\n halfwarp_id += gridDim.x * halfwarps_per_block;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockGSSmoothBlockDiaCsrKernel_V2(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices,\n ValueTypeA *Dinv, const int num_block_rows)\n{\n int cta_blockrow_id = threadIdx.x / bsize_sq;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int mat_entry_index = threadIdx.x - cta_blockrow_id * bsize_sq;\n const int i_ind = mat_entry_index / bsize;\n const int j_ind = mat_entry_index - i_ind * bsize;\n volatile __shared__ ValueTypeA s_Amat[bsize_sq * blockrows_per_cta];\n int offset, s_offset;\n ValueTypeA e_out;\n\n while (blockrow_id < num_block_rows && cta_blockrow_id < blockrows_per_cta)\n {\n // Store the diagonal\n offset = blockrow_id * bsize_sq + mat_entry_index;\n e_out = values[bsize_sq * dia_indices[blockrow_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n s_Amat[threadIdx.x] = e_out;\n s_offset = cta_blockrow_id * bsize_sq;\n#define s_A(ROW,COL) s_Amat[s_offset+ROW*bsize+COL]\n ValueTypeA diag;\n\n for (int row = 0; row < bsize; row++)\n {\n diag = 1.0 / s_A(row, row);\n\n if ((i_ind == 0) && !(j_ind == row))\n {\n s_A(row, j_ind) = s_A(row, j_ind) * diag;\n }\n\n if ((i_ind != row) && !(j_ind == row))\n {\n s_A(i_ind, j_ind) = -(s_A(i_ind, row) * s_A(row, j_ind)) + s_A(i_ind, j_ind);\n }\n\n if (i_ind == 0)\n {\n s_A(j_ind, row) = (j_ind == row) ? diag : -(s_A(j_ind, row) * diag);\n }\n }\n\n Dinv[offset] = s_A(i_ind, j_ind);\n blockrow_id += (gridDim.x * blockDim.x) / bsize_sq;\n }\n}\n\n\n\n\n// Kernel to smooth, NAIVE implementation with texture\ntemplate\n__global__\nvoid multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *dia_indices, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int num_rows_per_line, const int color_num, const int num_block_rows, ValueTypeB *xout)\n{\n const int vec_entry_index = threadIdx.x & (bsize - 1);\n const int block_eighthwarp_id = threadIdx.x >> log_bsize;\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int eighthwarp_id = tid >> log_bsize;\n volatile __shared__ ValueTypeB s_xtemp[ bsize * eighthwarps_per_block ];\n ValueTypeB bmAx, xin;\n ValueTypeB temp[bsize];\n int offset, i, s_offset;\n int n1d2 = num_rows_per_line * num_rows_per_line;\n int pln, tid2, row, col;\n int idx_1 = color_num & 1;\n int idx_2 = (color_num & 2) >> 1;\n int idx_3 = (color_num & 4) >> 2;\n\n while (true)\n {\n pln = eighthwarp_id / n1d2;\n tid2 = eighthwarp_id % n1d2;\n row = tid2 / num_rows_per_line;\n col = tid2 % num_rows_per_line;\n i = 4 * n1d2 * (2 * pln + idx_1) + 2 * num_rows_per_line * (2 * row + idx_2) + 2 * col + idx_3;\n\n //if (pln >= num_rows_per_line) return;\n //if (row >= num_rows_per_line) return;\n if (i > num_block_rows) { return; }\n\n offset = i * bsize + vec_entry_index;\n bmAx = b[offset];\n // Contribution from diagonal\n xin = x[offset];\n s_xtemp[threadIdx.x] = xin;\n // Load dia_values and do matrix multiply\n s_offset = block_eighthwarp_id * bsize;\n loadAsVector(nonzero_values + dia_indices[i]*bsize * bsize + vec_entry_index * bsize, temp);\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n\n if (jcol != i)\n {\n offset = jcol * bsize + vec_entry_index;\n s_xtemp[threadIdx.x] = x[offset];\n // Load nonzero_values\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n s_offset = block_eighthwarp_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n }\n }\n\n s_xtemp[threadIdx.x] = bmAx;\n bmAx = 0.;\n // Load Dinv and multiply to RHS\n offset = i * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(Dinv + offset, temp);\n s_offset = block_eighthwarp_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx += temp[m] * s_xtemp[s_offset + m];\n }\n\n xout[i * bsize + vec_entry_index] = xin + weight * bmAx;\n eighthwarp_id += (gridDim.x * blockDim.x >> log_bsize);\n }\n}\n\n// Kernel to smooth, NAIVE implementation with texture\ntemplate\n__global__\nvoid multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *dia_indices, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int *sorted_rows_by_color, const int num_rows_per_color, const int num_block_rows, ValueTypeB *xout)\n{\n int cta_blockrow_id = (threadIdx.x) / bsize;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int vec_entry_index = threadIdx.x - cta_blockrow_id * bsize;\n volatile __shared__ ValueTypeB s_xtemp[ bsize * blockrows_per_cta];\n ValueTypeB bmAx, xin;\n ValueTypeB temp[bsize];\n int offset, i, s_offset;\n\n while (blockrow_id < num_rows_per_color && cta_blockrow_id < blockrows_per_cta)\n {\n i = sorted_rows_by_color[blockrow_id];\n offset = i * bsize + vec_entry_index;\n bmAx = b[offset];\n // Contribution from diagonal\n xin = x[offset];\n s_xtemp[threadIdx.x] = xin;\n // Load dia_values and do matrix multiply\n s_offset = cta_blockrow_id * bsize;\n offset = dia_indices[i] * bsize * bsize + vec_entry_index * bsize;\n#pragma unroll\n\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = nonzero_values[offset + k];\n }\n\n //loadAsVector(dia_values+offset,temp);\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n\n if (jcol != i)\n {\n offset = jcol * bsize + vec_entry_index;\n s_xtemp[threadIdx.x] = x[offset];\n // Load nonzero_values\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n#pragma unroll\n\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = nonzero_values[offset + k];\n }\n\n //loadAsVector(nonzero_values+offset,temp);\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n }\n }\n\n s_xtemp[threadIdx.x] = bmAx;\n bmAx = 0.;\n // Load Dinv and multiply to RHS\n offset = i * bsize * bsize + vec_entry_index * bsize;\n#pragma unroll\n\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = Dinv[offset + k];\n }\n\n //loadAsVector(Dinv+offset,temp);\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx += temp[m] * s_xtemp[s_offset + m];\n }\n\n xout[i * bsize + vec_entry_index] = xin + weight * bmAx;\n blockrow_id += (gridDim.x * blockDim.x) / bsize;\n }\n}\n\n\n\n// Kernel to smooth, NAIVE implementation with texture\n// This is TERRIBLE in terms of memory access pattern\ntemplate\n__global__\nvoid fixcolorGSSmoothCsrKernel_NAIVE_tex(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *diag, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int num_rows_per_line, const int color_num, const int num_block_rows, ValueTypeB *xout)\n{\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int n1d2 = num_rows_per_line * num_rows_per_line;\n int pln, tid2, row, col;\n int idx_1 = color_num & 1;\n int idx_2 = (color_num & 2) >> 1;\n int idx_3 = (color_num & 4) >> 2;\n int i;\n ValueTypeB bmAx, xin;\n ValueTypeB temp, s_xtemp;\n ValueTypeB dia;\n\n while (true)\n {\n pln = tid / n1d2;\n tid2 = tid % n1d2;\n row = tid2 / num_rows_per_line;\n col = tid2 % num_rows_per_line;\n i = 4 * n1d2 * (2 * pln + idx_1) + 2 * num_rows_per_line * (2 * row + idx_2) + 2 * col + idx_3;\n\n //if (pln >= num_rows_per_line) return;\n //if (row >= num_rows_per_line) return;\n if (i > num_block_rows) { return; }\n\n bmAx = b[i];\n xin = x[i];\n dia = nonzero_values[diag[i]];\n bmAx -= xin * dia;\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n\n if (jcol != i)\n {\n s_xtemp = x[jcol];\n temp = nonzero_values[jind];\n bmAx -= temp * s_xtemp;\n }\n }\n\n bmAx /= dia;\n xout[i] = xin + weight * bmAx;\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\n// -------------------\n// Methods\n// -------------------\n\n\n// Constructor\ntemplate\nFixcolorGaussSeidelSolver_Base::FixcolorGaussSeidelSolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope)\n{\n this->weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n symFlag = cfg.AMG_Config::template getParameter(\"symmetric_GS\", cfg_scope);\n\n if (cfg.AMG_Config::template getParameter(\"use_bsrxmv\", cfg_scope))\n {\n this->use_bsrxmv = 1;\n }\n else\n {\n this->use_bsrxmv = 0;\n }\n\n if (weight == 0)\n {\n weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Multicolor GaussSeidel smoother\\n\");\n }\n\n num_colors = 8;\n}\n\n// Destructor\ntemplate\nFixcolorGaussSeidelSolver_Base::~FixcolorGaussSeidelSolver_Base()\n{\n}\n\ntemplate\nvoid\nFixcolorGaussSeidelSolver_Base::printSolverParameters() const\n{\n std::cout << \"relaxation_factor= \" << this->weight << std::endl;\n std::cout << \"use_bsrxmv = \" << this->use_bsrxmv << std::endl;\n std::cout << \"symmetric_GS = \" << this->symFlag << std::endl;\n}\n\n\ntemplate\nvoid FixcolorGaussSeidelSolver_Base::computeDinv(Matrix &A)\n{\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n ViewType oldView = A.currentView();\n A.setView(A_as_matrix->getViewExterior());\n\n if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n computeDinv_4x4(A);\n }\n else if (A.get_block_dimx() == 3 && A.get_block_dimy() == 3)\n {\n computeDinv_3x3(A);\n }\n else if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n computeDinv_1x1(A);\n }\n else if (A.get_block_dimx() == 2 && A.get_block_dimy() == 2)\n {\n computeDinv_2x2(A);\n }\n else if (A.get_block_dimx() == 5 && A.get_block_dimy() == 5)\n {\n computeDinv_5x5(A);\n }\n else if (A.get_block_dimx() == A.get_block_dimy())\n {\n computeDinv_bxb(A, A.get_block_dimx());\n }\n else\n {\n FatalError(\"Unsupported block size for FixcolorGaussSeidelSolver computeEinv\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n}\n\n// Method to compute the diagonal matrix E in DILU smoother\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_3x3(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_4x4(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_2x2(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_5x5(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_1x1(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_bxb(const Matrix_h &A, const int bsize)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_4x4(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 512;\n const int halfwarps_per_block = threads_per_block / 16;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / halfwarps_per_block + 1);\n cudaFuncSetCacheConfig(setupBlockGSSmooth4by4BlockDiaCsrKernel_V2, cudaFuncCachePreferL1);\n setupBlockGSSmooth4by4BlockDiaCsrKernel_V2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_bxb(const Matrix_d &A, const int bsize)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n MVector temp(AMGX_GRID_MAX_SIZE * ((bsize - 1) / 4 + 1) * ((bsize - 1) / 4 + 1));\n ValueTypeA *temp_ptr = temp.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 512;\n const int halfwarps_per_block = threads_per_block / 16;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / halfwarps_per_block + 1);\n cudaFuncSetCacheConfig(setupBlockGSSmoothbBybBlockDiaCsrKernel, cudaFuncCachePreferL1);\n setupBlockGSSmoothbBybBlockDiaCsrKernel <<< num_blocks, threads_per_block, sizeof(ValueTypeA)*bsize *bsize *halfwarps_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows(), bsize, bsize * bsize, temp_ptr);\n cudaCheckError();\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_3x3(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 256;\n const int blockrows_per_cta = threads_per_block / 9;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n cudaFuncSetCacheConfig(setupBlockGSSmooth3by3BlockDiaCsrKernel < IndexType, ValueTypeA, blockrows_per_cta, 32 / 9, 3, 9 >, cudaFuncCachePreferL1);\n setupBlockGSSmooth3by3BlockDiaCsrKernel < IndexType, ValueTypeA, blockrows_per_cta, 32 / 9, 3, 9 > <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_2x2(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 256;\n const int blockrows_per_cta = threads_per_block / 4;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n cudaFuncSetCacheConfig(setupBlockGSSmooth3by3BlockDiaCsrKernel < IndexType, ValueTypeA, blockrows_per_cta, 32 / 4, 2, 4 >, cudaFuncCachePreferL1);\n setupBlockGSSmooth3by3BlockDiaCsrKernel < IndexType, ValueTypeA, blockrows_per_cta, 32 / 4, 2, 4 > <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_5x5(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 256;\n const int blockrows_per_cta = threads_per_block / 25;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n cudaFuncSetCacheConfig(setupBlockGSSmooth3by3BlockDiaCsrKernel < IndexType, ValueTypeA, blockrows_per_cta, 32 / 25, 5, 25 >, cudaFuncCachePreferL1);\n setupBlockGSSmooth3by3BlockDiaCsrKernel < IndexType, ValueTypeA, blockrows_per_cta, 32 / 25, 5, 25 > <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_1x1(const Matrix_d &A)\n{\n}\n\n// Solver setup\ntemplate\nvoid\nFixcolorGaussSeidelSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n Matrix *A_as_matrix = dynamic_cast*>(Base::m_A);\n\n if (!A_as_matrix)\n {\n FatalError(\"FixcolorGaussSeidelSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n if (!this->use_bsrxmv)\n {\n FatalError(\"Use bsrxmv implementation, old implementation is buggy for diagonal matrix\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n computeDinv(*A_as_matrix);\n}\n\n//\ntemplate\nvoid\nFixcolorGaussSeidelSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_1x1(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_3x3(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_4x4(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_BxB(Matrix_h &A, VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_1x1(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n const IndexType *A_diag_ptr = A.diag.raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n int n = A.get_num_rows();\n int num_rows_per_line;\n num_rows_per_line = (int) ceil(cbrt((double)n) / 2);\n\n for (int i = 0; i < this->num_colors; i++)\n {\n int num_rows_per_color = n / this->num_colors + 1;\n const int threads_per_block = 512;\n //const int eightwarps_per_block = threads_per_block/4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n cudaFuncSetCacheConfig(fixcolorGSSmoothCsrKernel_NAIVE_tex, cudaFuncCachePreferL1);\n fixcolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, num_rows_per_line, i, n, x_ptr);\n }\n\n if (this->symFlag == 1)\n {\n int num_rows_per_color = n / this->num_colors + 1;\n\n for (int i = this->num_colors - 1; i >= 0; i--)\n {\n const int threads_per_block = 512;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n cudaFuncSetCacheConfig(fixcolorGSSmoothCsrKernel_NAIVE_tex, cudaFuncCachePreferL1);\n fixcolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, num_rows_per_line, i, n, x_ptr);\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_4x4(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n //const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n int n = A.get_num_rows();\n int num_rows_per_line;\n num_rows_per_line = (int) ceil(cbrt((double)n) / 2);\n\n for (int i = 0; i < this->num_colors; i++)\n {\n const IndexType num_rows_per_color = n / this->num_colors + 1;\n const int threads_per_block = 512;\n const int eightwarps_per_block = threads_per_block / 4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / eightwarps_per_block + 1 );\n cudaFuncSetCacheConfig(multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2, cudaFuncCachePreferL1);\n multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, num_rows_per_line, i, n, x_ptr);\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = this->num_colors - 1; i >= 0; i--)\n {\n const IndexType num_rows_per_color = n / this->num_colors + 1;\n const int threads_per_block = 512;\n const int eightwarps_per_block = threads_per_block / 4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / eightwarps_per_block + 1 );\n cudaFuncSetCacheConfig(multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2, cudaFuncCachePreferL1);\n multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, num_rows_per_line, i, n, x_ptr);\n cudaCheckError();\n } // End of loop over colors\n } // End of if symFlag\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_3x3(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n Matrix_d *A_as_matrix = (Matrix_d *) this->m_A;\n\n if (!A.hasProps(COLORING)) { FatalError(\"Matrix is not colored, exiting\", AMGX_ERR_BAD_PARAMETERS); }\n\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n //TODO: generate required indices on the fly since coloring is not requested for this solver\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n const int threads_per_block = 64;\n int num_colors = this->num_colors;\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == A_as_matrix->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color <= 0) { continue; }\n\n const int threads_per_block = 512;\n const int blockrows_per_cta = threads_per_block / 3;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows_per_color - 1) / blockrows_per_cta + 1);\n cudaFuncSetCacheConfig(multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2, cudaFuncCachePreferL1);\n multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == A_as_matrix->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int blockrows_per_cta = threads_per_block / 3;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows_per_color - 1) / blockrows_per_cta + 1);\n cudaFuncSetCacheConfig(multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2, cudaFuncCachePreferL1);\n multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n } // End of loop over colors\n } // End of if symFlag\n}\n\ntemplate <>\nvoid FixcolorGaussSeidelSolver >::smooth_BxB(Matrix > &A, Vector > &b, Vector > &x, ViewType separation_flag)\n{\n FatalError(\"Mixed precision is not supported\", AMGX_ERR_NOT_IMPLEMENTED);\n};\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_BxB(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flag)\n{\n if (!A.hasProps(COLORING)) { FatalError(\"Matrix is not colored, exiting\", AMGX_ERR_BAD_PARAMETERS); }\n\n int num_colors = this->num_colors;\n VVector y(b.size());\n amgx::thrust::copy(b.begin(), b.end(), y.begin()); // y= b for all colors\n cudaCheckError();\n\n for (int color = 0; color < num_colors; color++)\n {\n Cusparse::bsrmv(Cusparse::ALL_COLORS, color, (ValueTypeB) - 1.0f, A, x, (ValueTypeB)1.0f, y, separation_flag); // y= -A.x + y for current color\n Cusparse::bsrmv(color, (ValueTypeB)this->weight, A, this->Dinv, y, (ValueTypeB)1.0f, x, separation_flag); // x= w*Dinv.y + x for current color\n }\n\n cudaCheckError();\n\n if (this->symFlag)\n {\n y = b; // y= b for all colors\n\n for (int color = num_colors - 1; color >= 0; color--)\n {\n Cusparse::bsrmv(Cusparse::ALL_COLORS, color, (ValueTypeB) - 1.0f, A, x, (ValueTypeB)1.0f, y, separation_flag); // y= -A.x + y for current color\n Cusparse::bsrmv(color, (ValueTypeB)this->weight, A, this->Dinv, y, (ValueTypeB)1.0f, x, separation_flag); // x= w*Dinv.y + x for current color\n }\n }\n\n cudaCheckError();\n}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nFixcolorGaussSeidelSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero)\n{\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n\n if (xIsZero) { x.dirtybit = 0; }\n\n if (!A_as_matrix->is_matrix_singleGPU())\n {\n A_as_matrix->manager->exchange_halo_async(x, x.tag);\n A_as_matrix->manager->exchange_halo_async(b, b.tag);\n }\n\n if (A_as_matrix->getViewExterior() == A_as_matrix->getViewInterior())\n {\n if (!A_as_matrix->is_matrix_singleGPU())\n {\n A_as_matrix->manager->exchange_halo_wait(x, x.tag);\n A_as_matrix->manager->exchange_halo_wait(b, b.tag);\n }\n }\n\n if ((A_as_matrix->get_block_dimx() == 4 && A_as_matrix->get_block_dimy() == 4) || (A_as_matrix->get_block_dimx() == 1 && A_as_matrix->get_block_dimy() == 1))\n {\n if (!A_as_matrix->is_matrix_singleGPU())\n {\n A_as_matrix->manager->exchange_halo_wait(x, x.tag);\n A_as_matrix->manager->exchange_halo_wait(b, b.tag);\n }\n }\n\n ViewType oldView = A_as_matrix->currentView();\n bool latencyHiding = true;\n ViewType flags;\n\n if (A_as_matrix->is_matrix_singleGPU() || (x.dirtybit == 0 && b.dirtybit == 0))\n {\n latencyHiding = false;\n A_as_matrix->setViewExterior();\n flags = A_as_matrix->getViewExterior();\n }\n else\n {\n flags = A_as_matrix->getViewInterior();\n A_as_matrix->setViewInterior();\n }\n\n if (xIsZero)\n {\n thrust_wrapper::fill(x.begin(), x.end(), 0.);\n cudaCheckError();\n }\n\n if (this->use_bsrxmv && A_as_matrix->get_block_dimx() == A_as_matrix->get_block_dimy())\n {\n smooth_BxB(*A_as_matrix, b, x, flags);\n }\n else if ( A_as_matrix->get_block_dimx() == 4 && A_as_matrix->get_block_dimy() == 4 )\n {\n smooth_4x4(*A_as_matrix, b, x, flags);\n }\n else if ( A_as_matrix->get_block_dimx() == 3 && A_as_matrix->get_block_dimy() == 3 )\n {\n smooth_3x3(*A_as_matrix, b, x, flags);\n }\n else if ( A_as_matrix->get_block_dimx() == 1 && A_as_matrix->get_block_dimy() == 1 )\n {\n smooth_1x1(*A_as_matrix, b, x, flags);\n }\n else\n {\n FatalError(\"Unsupported block size for MulticolorGaussSeidelSolver smooth\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if (latencyHiding)\n {\n if (!A_as_matrix->is_matrix_singleGPU())\n {\n A_as_matrix->manager->exchange_halo_wait(x, x.tag);\n A_as_matrix->manager->exchange_halo_wait(b, b.tag);\n }\n\n A_as_matrix->setViewExterior();\n flags = (ViewType)(~(A_as_matrix->getViewInterior()) & A_as_matrix->getViewExterior());\n\n if (flags != 0)\n {\n if (this->use_bsrxmv && A_as_matrix->get_block_dimx() == A_as_matrix->get_block_dimy())\n {\n smooth_BxB(*A_as_matrix, b, x, flags);\n }\n else if ( A_as_matrix->get_block_dimx() == 4 && A_as_matrix->get_block_dimy() == 4 )\n {\n smooth_4x4(*A_as_matrix, b, x, flags);\n }\n else if ( A_as_matrix->get_block_dimx() == 3 && A_as_matrix->get_block_dimy() == 3 )\n {\n smooth_3x3(*A_as_matrix, b, x, flags);\n }\n else if ( A_as_matrix->get_block_dimx() == 1 && A_as_matrix->get_block_dimy() == 1 )\n {\n smooth_1x1(*A_as_matrix, b, x, flags);\n }\n }\n }\n\n x.dirtybit = 1;\n A_as_matrix->setView(oldView);\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nFixcolorGaussSeidelSolver_Base::solve_finalize(VVector &b, VVector &x)\n{\n}\n\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class FixcolorGaussSeidelSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class FixcolorGaussSeidelSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\nnamespace fixcolor_gauss_seidel_solver\n{\n\n// -------------------------\n// Kernels\n// -------------------------\n\ntemplate\n__global__\nvoid setupBlockGSSmooth3by3BlockDiaCsrKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows)\n{\n int warp_id = threadIdx.x / 32;\n int warp_thread_id = threadIdx.x & 31;\n\n // padding row blocks to fit in a single warp\n if ( warp_thread_id >= blockrows_per_warp * bsize_sq ) { return; }\n\n // new thread id with padding\n int tid = warp_id * blockrows_per_warp * bsize_sq + warp_thread_id;\n int cta_blockrow_id = tid / bsize_sq;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int mat_entry_index = tid - cta_blockrow_id * bsize_sq;\n const int i_ind = mat_entry_index / bsize;\n const int j_ind = mat_entry_index - i_ind * bsize;\n volatile __shared__ ValueTypeA s_Amat[bsize_sq * blockrows_per_cta ];\n ValueTypeA e_out;\n\n while (blockrow_id < num_block_rows)\n {\n int offset = blockrow_id * bsize_sq + mat_entry_index;\n // Store the diagonal\n e_out = values[bsize_sq * dia_indices[blockrow_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n s_Amat[tid] = e_out;\n compute_block_inverse_row_major\n ( s_Amat, cta_blockrow_id * bsize_sq, offset, i_ind, j_ind, Dinv );\n blockrow_id += gridDim.x * blockrows_per_cta;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockGSSmooth4by4BlockDiaCsrKernel_V2(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows)\n\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int halfwarp_id = tid >> log_bsize_sq;\n const int block_halfwarp_id = threadIdx.x >> log_bsize_sq;\n const int mat_entry_index = threadIdx.x & (bsize_sq - 1);\n const int i_ind = mat_entry_index >> log_bsize;\n const int j_ind = mat_entry_index & (bsize - 1);\n volatile __shared__ ValueTypeA s_Amat[bsize_sq * halfwarps_per_block ];\n ValueTypeA e_out;\n\n while (halfwarp_id < num_block_rows)\n {\n int offset = halfwarp_id * bsize_sq + mat_entry_index;\n // Store the diagonal\n e_out = values[bsize_sq * dia_indices[halfwarp_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n s_Amat[threadIdx.x] = e_out;\n compute_block_inverse_row_major\n ( s_Amat, block_halfwarp_id * bsize_sq, offset, i_ind, j_ind, Dinv );\n halfwarp_id += gridDim.x * blockDim.x >> log_bsize_sq;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockGSSmoothbBybBlockDiaCsrKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows, int bsize, int bsize_sq, ValueTypeA *temp1)\n\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int halfwarp_id = tid >> 4;\n const int block_halfwarp_id = threadIdx.x >> 4;\n const int mat_entry_index = threadIdx.x & (16 - 1);\n const int i_ind = mat_entry_index >> 2;\n const int j_ind = mat_entry_index & 3;\n extern __shared__ volatile char sharedc[];\n volatile ValueTypeA *s_Amat;\n s_Amat = (ValueTypeA *)&sharedc[0];\n int tile_num = (bsize - 1) / 4 + 1;\n ValueTypeA *e_out = &temp1[(blockIdx.x * blockDim.x + threadIdx.x) * tile_num * tile_num];\n\n while (halfwarp_id < num_block_rows)\n {\n int offset = halfwarp_id * bsize_sq + i_ind * bsize + j_ind;\n int s_offset = block_halfwarp_id * bsize_sq;\n\n // Store the diagonal\n for (int t1 = 0; t1 < tile_num; t1++)\n for (int t2 = 0; t2 < tile_num; t2++)\n if ((t1 * 4 + i_ind) < bsize && (t2 * 4 + j_ind) < bsize)\n {\n e_out[t1 * tile_num + t2] = values[bsize_sq * dia_indices[halfwarp_id] + (t1 * 4 + i_ind) * bsize + t2 * 4 + j_ind];\n }\n\n // Each thread stores its entry in s_Amat\n for (int t1 = 0; t1 < tile_num; t1++)\n for (int t2 = 0; t2 < tile_num; t2++)\n if ((t1 * 4 + i_ind) < bsize && (t2 * 4 + j_ind) < bsize)\n {\n s_Amat[s_offset + (t1 * 4 + i_ind) * bsize + t2 * 4 + j_ind] = e_out[t1 * tile_num + t2];\n }\n\n compute_block_inverse2\n ( s_Amat, s_offset, offset, i_ind, j_ind, Dinv, tile_num, bsize, bsize_sq );\n halfwarp_id += gridDim.x * halfwarps_per_block;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockGSSmoothBlockDiaCsrKernel_V2(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices,\n ValueTypeA *Dinv, const int num_block_rows)\n{\n int cta_blockrow_id = threadIdx.x / bsize_sq;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int mat_entry_index = threadIdx.x - cta_blockrow_id * bsize_sq;\n const int i_ind = mat_entry_index / bsize;\n const int j_ind = mat_entry_index - i_ind * bsize;\n volatile __shared__ ValueTypeA s_Amat[bsize_sq * blockrows_per_cta];\n int offset, s_offset;\n ValueTypeA e_out;\n\n while (blockrow_id < num_block_rows && cta_blockrow_id < blockrows_per_cta)\n {\n // Store the diagonal\n offset = blockrow_id * bsize_sq + mat_entry_index;\n e_out = values[bsize_sq * dia_indices[blockrow_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n s_Amat[threadIdx.x] = e_out;\n s_offset = cta_blockrow_id * bsize_sq;\n#define s_A(ROW,COL) s_Amat[s_offset+ROW*bsize+COL]\n ValueTypeA diag;\n\n for (int row = 0; row < bsize; row++)\n {\n diag = 1.0 / s_A(row, row);\n\n if ((i_ind == 0) && !(j_ind == row))\n {\n s_A(row, j_ind) = s_A(row, j_ind) * diag;\n }\n\n if ((i_ind != row) && !(j_ind == row))\n {\n s_A(i_ind, j_ind) = -(s_A(i_ind, row) * s_A(row, j_ind)) + s_A(i_ind, j_ind);\n }\n\n if (i_ind == 0)\n {\n s_A(j_ind, row) = (j_ind == row) ? diag : -(s_A(j_ind, row) * diag);\n }\n }\n\n Dinv[offset] = s_A(i_ind, j_ind);\n blockrow_id += (gridDim.x * blockDim.x) / bsize_sq;\n }\n}\n\n\n\n\n// Kernel to smooth, NAIVE implementation with texture\ntemplate\n__global__\nvoid multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *dia_indices, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int num_rows_per_line, const int color_num, const int num_block_rows, ValueTypeB *xout)\n{\n const int vec_entry_index = threadIdx.x & (bsize - 1);\n const int block_eighthwarp_id = threadIdx.x >> log_bsize;\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int eighthwarp_id = tid >> log_bsize;\n volatile __shared__ ValueTypeB s_xtemp[ bsize * eighthwarps_per_block ];\n ValueTypeB bmAx, xin;\n ValueTypeB temp[bsize];\n int offset, i, s_offset;\n int n1d2 = num_rows_per_line * num_rows_per_line;\n int pln, tid2, row, col;\n int idx_1 = color_num & 1;\n int idx_2 = (color_num & 2) >> 1;\n int idx_3 = (color_num & 4) >> 2;\n\n while (true)\n {\n pln = eighthwarp_id / n1d2;\n tid2 = eighthwarp_id % n1d2;\n row = tid2 / num_rows_per_line;\n col = tid2 % num_rows_per_line;\n i = 4 * n1d2 * (2 * pln + idx_1) + 2 * num_rows_per_line * (2 * row + idx_2) + 2 * col + idx_3;\n\n //if (pln >= num_rows_per_line) return;\n //if (row >= num_rows_per_line) return;\n if (i > num_block_rows) { return; }\n\n offset = i * bsize + vec_entry_index;\n bmAx = b[offset];\n // Contribution from diagonal\n xin = x[offset];\n s_xtemp[threadIdx.x] = xin;\n // Load dia_values and do matrix multiply\n s_offset = block_eighthwarp_id * bsize;\n loadAsVector(nonzero_values + dia_indices[i]*bsize * bsize + vec_entry_index * bsize, temp);\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n\n if (jcol != i)\n {\n offset = jcol * bsize + vec_entry_index;\n s_xtemp[threadIdx.x] = x[offset];\n // Load nonzero_values\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n s_offset = block_eighthwarp_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n }\n }\n\n s_xtemp[threadIdx.x] = bmAx;\n bmAx = 0.;\n // Load Dinv and multiply to RHS\n offset = i * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(Dinv + offset, temp);\n s_offset = block_eighthwarp_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx += temp[m] * s_xtemp[s_offset + m];\n }\n\n xout[i * bsize + vec_entry_index] = xin + weight * bmAx;\n eighthwarp_id += (gridDim.x * blockDim.x >> log_bsize);\n }\n}\n\n// Kernel to smooth, NAIVE implementation with texture\ntemplate\n__global__\nvoid multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *dia_indices, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int *sorted_rows_by_color, const int num_rows_per_color, const int num_block_rows, ValueTypeB *xout)\n{\n int cta_blockrow_id = (threadIdx.x) / bsize;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int vec_entry_index = threadIdx.x - cta_blockrow_id * bsize;\n volatile __shared__ ValueTypeB s_xtemp[ bsize * blockrows_per_cta];\n ValueTypeB bmAx, xin;\n ValueTypeB temp[bsize];\n int offset, i, s_offset;\n\n while (blockrow_id < num_rows_per_color && cta_blockrow_id < blockrows_per_cta)\n {\n i = sorted_rows_by_color[blockrow_id];\n offset = i * bsize + vec_entry_index;\n bmAx = b[offset];\n // Contribution from diagonal\n xin = x[offset];\n s_xtemp[threadIdx.x] = xin;\n // Load dia_values and do matrix multiply\n s_offset = cta_blockrow_id * bsize;\n offset = dia_indices[i] * bsize * bsize + vec_entry_index * bsize;\n#pragma unroll\n\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = nonzero_values[offset + k];\n }\n\n //loadAsVector(dia_values+offset,temp);\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n\n if (jcol != i)\n {\n offset = jcol * bsize + vec_entry_index;\n s_xtemp[threadIdx.x] = x[offset];\n // Load nonzero_values\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n#pragma unroll\n\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = nonzero_values[offset + k];\n }\n\n //loadAsVector(nonzero_values+offset,temp);\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset + m];\n }\n }\n }\n\n s_xtemp[threadIdx.x] = bmAx;\n bmAx = 0.;\n // Load Dinv and multiply to RHS\n offset = i * bsize * bsize + vec_entry_index * bsize;\n#pragma unroll\n\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = Dinv[offset + k];\n }\n\n //loadAsVector(Dinv+offset,temp);\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx += temp[m] * s_xtemp[s_offset + m];\n }\n\n xout[i * bsize + vec_entry_index] = xin + weight * bmAx;\n blockrow_id += (gridDim.x * blockDim.x) / bsize;\n }\n}\n\n\n\n// Kernel to smooth, NAIVE implementation with texture\n// This is TERRIBLE in terms of memory access pattern\ntemplate\n__global__\nvoid fixcolorGSSmoothCsrKernel_NAIVE_tex(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *diag, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, ValueTypeB weight, const int num_rows_per_line, const int color_num, const int num_block_rows, ValueTypeB *xout)\n{\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int n1d2 = num_rows_per_line * num_rows_per_line;\n int pln, tid2, row, col;\n int idx_1 = color_num & 1;\n int idx_2 = (color_num & 2) >> 1;\n int idx_3 = (color_num & 4) >> 2;\n int i;\n ValueTypeB bmAx, xin;\n ValueTypeB temp, s_xtemp;\n ValueTypeB dia;\n\n while (true)\n {\n pln = tid / n1d2;\n tid2 = tid % n1d2;\n row = tid2 / num_rows_per_line;\n col = tid2 % num_rows_per_line;\n i = 4 * n1d2 * (2 * pln + idx_1) + 2 * num_rows_per_line * (2 * row + idx_2) + 2 * col + idx_3;\n\n //if (pln >= num_rows_per_line) return;\n //if (row >= num_rows_per_line) return;\n if (i > num_block_rows) { return; }\n\n bmAx = b[i];\n xin = x[i];\n dia = nonzero_values[diag[i]];\n bmAx -= xin * dia;\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n\n if (jcol != i)\n {\n s_xtemp = x[jcol];\n temp = nonzero_values[jind];\n bmAx -= temp * s_xtemp;\n }\n }\n\n bmAx /= dia;\n xout[i] = xin + weight * bmAx;\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\n// -------------------\n// Methods\n// -------------------\n\n\n// Constructor\ntemplate\nFixcolorGaussSeidelSolver_Base::FixcolorGaussSeidelSolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope)\n{\n this->weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n symFlag = cfg.AMG_Config::template getParameter(\"symmetric_GS\", cfg_scope);\n\n if (cfg.AMG_Config::template getParameter(\"use_bsrxmv\", cfg_scope))\n {\n this->use_bsrxmv = 1;\n }\n else\n {\n this->use_bsrxmv = 0;\n }\n\n if (weight == 0)\n {\n weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Multicolor GaussSeidel smoother\\n\");\n }\n\n num_colors = 8;\n}\n\n// Destructor\ntemplate\nFixcolorGaussSeidelSolver_Base::~FixcolorGaussSeidelSolver_Base()\n{\n}\n\ntemplate\nvoid\nFixcolorGaussSeidelSolver_Base::printSolverParameters() const\n{\n std::cout << \"relaxation_factor= \" << this->weight << std::endl;\n std::cout << \"use_bsrxmv = \" << this->use_bsrxmv << std::endl;\n std::cout << \"symmetric_GS = \" << this->symFlag << std::endl;\n}\n\n\ntemplate\nvoid FixcolorGaussSeidelSolver_Base::computeDinv(Matrix &A)\n{\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n ViewType oldView = A.currentView();\n A.setView(A_as_matrix->getViewExterior());\n\n if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n computeDinv_4x4(A);\n }\n else if (A.get_block_dimx() == 3 && A.get_block_dimy() == 3)\n {\n computeDinv_3x3(A);\n }\n else if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n computeDinv_1x1(A);\n }\n else if (A.get_block_dimx() == 2 && A.get_block_dimy() == 2)\n {\n computeDinv_2x2(A);\n }\n else if (A.get_block_dimx() == 5 && A.get_block_dimy() == 5)\n {\n computeDinv_5x5(A);\n }\n else if (A.get_block_dimx() == A.get_block_dimy())\n {\n computeDinv_bxb(A, A.get_block_dimx());\n }\n else\n {\n FatalError(\"Unsupported block size for FixcolorGaussSeidelSolver computeEinv\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n}\n\n// Method to compute the diagonal matrix E in DILU smoother\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_3x3(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_4x4(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_2x2(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_5x5(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_1x1(const Matrix_h &A)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_bxb(const Matrix_h &A, const int bsize)\n{\n FatalError(\"Multicolor GS smoother not implemented with host format, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_4x4(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 512;\n const int halfwarps_per_block = threads_per_block / 16;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / halfwarps_per_block + 1);\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockGSSmooth4by4BlockDiaCsrKernel_V2), hipFuncCachePreferL1);\n setupBlockGSSmooth4by4BlockDiaCsrKernel_V2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_bxb(const Matrix_d &A, const int bsize)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n MVector temp(AMGX_GRID_MAX_SIZE * ((bsize - 1) / 4 + 1) * ((bsize - 1) / 4 + 1));\n ValueTypeA *temp_ptr = temp.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 512;\n const int halfwarps_per_block = threads_per_block / 16;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / halfwarps_per_block + 1);\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockGSSmoothbBybBlockDiaCsrKernel), hipFuncCachePreferL1);\n setupBlockGSSmoothbBybBlockDiaCsrKernel <<< num_blocks, threads_per_block, sizeof(ValueTypeA)*bsize *bsize *halfwarps_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows(), bsize, bsize * bsize, temp_ptr);\n cudaCheckError();\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_3x3(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 256;\n const int blockrows_per_cta = threads_per_block / 9;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockGSSmooth3by3BlockDiaCsrKernel < IndexType, ValueTypeA, blockrows_per_cta, 32 / 9, 3, 9 >), hipFuncCachePreferL1);\n setupBlockGSSmooth3by3BlockDiaCsrKernel < IndexType, ValueTypeA, blockrows_per_cta, 32 / 9, 3, 9 > <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_2x2(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 256;\n const int blockrows_per_cta = threads_per_block / 4;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockGSSmooth3by3BlockDiaCsrKernel < IndexType, ValueTypeA, blockrows_per_cta, 32 / 4, 2, 4 >), hipFuncCachePreferL1);\n setupBlockGSSmooth3by3BlockDiaCsrKernel < IndexType, ValueTypeA, blockrows_per_cta, 32 / 4, 2, 4 > <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_5x5(const Matrix_d &A)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 256;\n const int blockrows_per_cta = threads_per_block / 25;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockGSSmooth3by3BlockDiaCsrKernel < IndexType, ValueTypeA, blockrows_per_cta, 32 / 25, 5, 25 >), hipFuncCachePreferL1);\n setupBlockGSSmooth3by3BlockDiaCsrKernel < IndexType, ValueTypeA, blockrows_per_cta, 32 / 25, 5, 25 > <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::computeDinv_1x1(const Matrix_d &A)\n{\n}\n\n// Solver setup\ntemplate\nvoid\nFixcolorGaussSeidelSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n Matrix *A_as_matrix = dynamic_cast*>(Base::m_A);\n\n if (!A_as_matrix)\n {\n FatalError(\"FixcolorGaussSeidelSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n if (!this->use_bsrxmv)\n {\n FatalError(\"Use bsrxmv implementation, old implementation is buggy for diagonal matrix\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n computeDinv(*A_as_matrix);\n}\n\n//\ntemplate\nvoid\nFixcolorGaussSeidelSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_1x1(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_3x3(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_4x4(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_BxB(Matrix_h &A, VVector &b, VVector &x, ViewType separation_flag)\n{\n FatalError(\"Haven't implemented Block Multicolor Gauss-Seidel smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_1x1(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n const IndexType *A_diag_ptr = A.diag.raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n int n = A.get_num_rows();\n int num_rows_per_line;\n num_rows_per_line = (int) ceil(cbrt((double)n) / 2);\n\n for (int i = 0; i < this->num_colors; i++)\n {\n int num_rows_per_color = n / this->num_colors + 1;\n const int threads_per_block = 512;\n //const int eightwarps_per_block = threads_per_block/4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n hipFuncSetCacheConfig(reinterpret_cast(fixcolorGSSmoothCsrKernel_NAIVE_tex), hipFuncCachePreferL1);\n fixcolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, num_rows_per_line, i, n, x_ptr);\n }\n\n if (this->symFlag == 1)\n {\n int num_rows_per_color = n / this->num_colors + 1;\n\n for (int i = this->num_colors - 1; i >= 0; i--)\n {\n const int threads_per_block = 512;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / threads_per_block + 1 );\n hipFuncSetCacheConfig(reinterpret_cast(fixcolorGSSmoothCsrKernel_NAIVE_tex), hipFuncCachePreferL1);\n fixcolorGSSmoothCsrKernel_NAIVE_tex <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_diag_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, num_rows_per_line, i, n, x_ptr);\n }\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_4x4(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n //const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n int n = A.get_num_rows();\n int num_rows_per_line;\n num_rows_per_line = (int) ceil(cbrt((double)n) / 2);\n\n for (int i = 0; i < this->num_colors; i++)\n {\n const IndexType num_rows_per_color = n / this->num_colors + 1;\n const int threads_per_block = 512;\n const int eightwarps_per_block = threads_per_block / 4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / eightwarps_per_block + 1 );\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2), hipFuncCachePreferL1);\n multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, num_rows_per_line, i, n, x_ptr);\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = this->num_colors - 1; i >= 0; i--)\n {\n const IndexType num_rows_per_color = n / this->num_colors + 1;\n const int threads_per_block = 512;\n const int eightwarps_per_block = threads_per_block / 4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int)(num_rows_per_color - 1) / eightwarps_per_block + 1 );\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2), hipFuncCachePreferL1);\n multicolorGSSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, num_rows_per_line, i, n, x_ptr);\n cudaCheckError();\n } // End of loop over colors\n } // End of if symFlag\n}\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_3x3(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flag)\n{\n Matrix_d *A_as_matrix = (Matrix_d *) this->m_A;\n\n if (!A.hasProps(COLORING)) { FatalError(\"Matrix is not colored, exiting\", AMGX_ERR_BAD_PARAMETERS); }\n\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n //TODO: generate required indices on the fly since coloring is not requested for this solver\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n const int threads_per_block = 64;\n int num_colors = this->num_colors;\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == A_as_matrix->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color <= 0) { continue; }\n\n const int threads_per_block = 512;\n const int blockrows_per_cta = threads_per_block / 3;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows_per_color - 1) / blockrows_per_cta + 1);\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2), hipFuncCachePreferL1);\n multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n }\n\n if (this->symFlag == 1)\n {\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = ((separation_flag & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flag == A_as_matrix->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int blockrows_per_cta = threads_per_block / 3;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows_per_color - 1) / blockrows_per_cta + 1);\n hipFuncSetCacheConfig(reinterpret_cast(multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2), hipFuncCachePreferL1);\n multicolorGSSmoothBlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color, A.get_num_rows(), x_ptr);\n cudaCheckError();\n } // End of loop over colors\n } // End of if symFlag\n}\n\ntemplate <>\nvoid FixcolorGaussSeidelSolver >::smooth_BxB(Matrix > &A, Vector > &b, Vector > &x, ViewType separation_flag)\n{\n FatalError(\"Mixed precision is not supported\", AMGX_ERR_NOT_IMPLEMENTED);\n};\n\ntemplate \nvoid FixcolorGaussSeidelSolver >::smooth_BxB(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flag)\n{\n if (!A.hasProps(COLORING)) { FatalError(\"Matrix is not colored, exiting\", AMGX_ERR_BAD_PARAMETERS); }\n\n int num_colors = this->num_colors;\n VVector y(b.size());\n amgx::thrust::copy(b.begin(), b.end(), y.begin()); // y= b for all colors\n cudaCheckError();\n\n for (int color = 0; color < num_colors; color++)\n {\n Cusparse::bsrmv(Cusparse::ALL_COLORS, color, (ValueTypeB) - 1.0f, A, x, (ValueTypeB)1.0f, y, separation_flag); // y= -A.x + y for current color\n Cusparse::bsrmv(color, (ValueTypeB)this->weight, A, this->Dinv, y, (ValueTypeB)1.0f, x, separation_flag); // x= w*Dinv.y + x for current color\n }\n\n cudaCheckError();\n\n if (this->symFlag)\n {\n y = b; // y= b for all colors\n\n for (int color = num_colors - 1; color >= 0; color--)\n {\n Cusparse::bsrmv(Cusparse::ALL_COLORS, color, (ValueTypeB) - 1.0f, A, x, (ValueTypeB)1.0f, y, separation_flag); // y= -A.x + y for current color\n Cusparse::bsrmv(color, (ValueTypeB)this->weight, A, this->Dinv, y, (ValueTypeB)1.0f, x, separation_flag); // x= w*Dinv.y + x for current color\n }\n }\n\n cudaCheckError();\n}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nFixcolorGaussSeidelSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero)\n{\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n\n if (xIsZero) { x.dirtybit = 0; }\n\n if (!A_as_matrix->is_matrix_singleGPU())\n {\n A_as_matrix->manager->exchange_halo_async(x, x.tag);\n A_as_matrix->manager->exchange_halo_async(b, b.tag);\n }\n\n if (A_as_matrix->getViewExterior() == A_as_matrix->getViewInterior())\n {\n if (!A_as_matrix->is_matrix_singleGPU())\n {\n A_as_matrix->manager->exchange_halo_wait(x, x.tag);\n A_as_matrix->manager->exchange_halo_wait(b, b.tag);\n }\n }\n\n if ((A_as_matrix->get_block_dimx() == 4 && A_as_matrix->get_block_dimy() == 4) || (A_as_matrix->get_block_dimx() == 1 && A_as_matrix->get_block_dimy() == 1))\n {\n if (!A_as_matrix->is_matrix_singleGPU())\n {\n A_as_matrix->manager->exchange_halo_wait(x, x.tag);\n A_as_matrix->manager->exchange_halo_wait(b, b.tag);\n }\n }\n\n ViewType oldView = A_as_matrix->currentView();\n bool latencyHiding = true;\n ViewType flags;\n\n if (A_as_matrix->is_matrix_singleGPU() || (x.dirtybit == 0 && b.dirtybit == 0))\n {\n latencyHiding = false;\n A_as_matrix->setViewExterior();\n flags = A_as_matrix->getViewExterior();\n }\n else\n {\n flags = A_as_matrix->getViewInterior();\n A_as_matrix->setViewInterior();\n }\n\n if (xIsZero)\n {\n thrust_wrapper::fill(x.begin(), x.end(), 0.);\n cudaCheckError();\n }\n\n if (this->use_bsrxmv && A_as_matrix->get_block_dimx() == A_as_matrix->get_block_dimy())\n {\n smooth_BxB(*A_as_matrix, b, x, flags);\n }\n else if ( A_as_matrix->get_block_dimx() == 4 && A_as_matrix->get_block_dimy() == 4 )\n {\n smooth_4x4(*A_as_matrix, b, x, flags);\n }\n else if ( A_as_matrix->get_block_dimx() == 3 && A_as_matrix->get_block_dimy() == 3 )\n {\n smooth_3x3(*A_as_matrix, b, x, flags);\n }\n else if ( A_as_matrix->get_block_dimx() == 1 && A_as_matrix->get_block_dimy() == 1 )\n {\n smooth_1x1(*A_as_matrix, b, x, flags);\n }\n else\n {\n FatalError(\"Unsupported block size for MulticolorGaussSeidelSolver smooth\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if (latencyHiding)\n {\n if (!A_as_matrix->is_matrix_singleGPU())\n {\n A_as_matrix->manager->exchange_halo_wait(x, x.tag);\n A_as_matrix->manager->exchange_halo_wait(b, b.tag);\n }\n\n A_as_matrix->setViewExterior();\n flags = (ViewType)(~(A_as_matrix->getViewInterior()) & A_as_matrix->getViewExterior());\n\n if (flags != 0)\n {\n if (this->use_bsrxmv && A_as_matrix->get_block_dimx() == A_as_matrix->get_block_dimy())\n {\n smooth_BxB(*A_as_matrix, b, x, flags);\n }\n else if ( A_as_matrix->get_block_dimx() == 4 && A_as_matrix->get_block_dimy() == 4 )\n {\n smooth_4x4(*A_as_matrix, b, x, flags);\n }\n else if ( A_as_matrix->get_block_dimx() == 3 && A_as_matrix->get_block_dimy() == 3 )\n {\n smooth_3x3(*A_as_matrix, b, x, flags);\n }\n else if ( A_as_matrix->get_block_dimx() == 1 && A_as_matrix->get_block_dimy() == 1 )\n {\n smooth_1x1(*A_as_matrix, b, x, flags);\n }\n }\n }\n\n x.dirtybit = 1;\n A_as_matrix->setView(oldView);\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nFixcolorGaussSeidelSolver_Base::solve_finalize(VVector &b, VVector &x)\n{\n}\n\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class FixcolorGaussSeidelSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class FixcolorGaussSeidelSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/solvers/fixcolor_gauss_seidel_solver.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace amgx\n{\nnamespace block_jacobi_solver\n{\n\ntemplate \nstruct jacobi_presmooth_functor\n{\n double omega;\n jacobi_presmooth_functor( double omega ) : omega( omega ) {}\n __host__ __device__ ValueTypeB operator()( const ValueTypeB &b, const ValueTypeA &d ) const { return isNotCloseToZero(d) ? b * omega / d : b * omega / epsilon(d); }\n};\n\ntemplate \nstruct jacobi_postsmooth_functor\n{\n double omega;\n jacobi_postsmooth_functor( double omega ) : omega( omega ) {}\n template __host__ __device__ ValueTypeB operator( )( const Tuple &t ) const\n {\n ValueTypeB x = amgx::thrust::get<0>(t);\n ValueTypeA d = amgx::thrust::get<1>(t);\n ValueTypeB b = amgx::thrust::get<2>(t);\n ValueTypeB y = amgx::thrust::get<3>(t);\n // return x + omega * (b - y) / d.\n d = isNotCloseToZero(d) ? d : epsilon(d);\n d = types::util::get_one() / d;\n b = b - y;\n b = b * omega;\n return b * d + x;\n }\n};\n\ntemplate \nstruct add_functor\n{\n __host__ __device__ ValueTypeB operator()( const ValueTypeB &x, const ValueTypeB &y )const { return x + y; }\n};\n\ntemplate\n__device__ __forceinline__ T fmnaOp (T a, T b, T c)\n{\n return -(a * b) + c;\n}\ntemplate\n__device__ __forceinline__ T mulOp (T a, T b)\n{\n return a * b;\n}\ntemplate\n__device__ __forceinline__ T rcpOp (T a)\n{\n return 1.0 / (isNotCloseToZero(a) ? a : epsilon(a));\n}\ntemplate\n__device__ __forceinline__ T absOp (T a)\n{\n return fabs(a);\n}\n\n// -----------------------------------\n// KERNELS\n// -----------------------------------\n\ntemplate\n__global__ void matinv_matrix_per_thread_pivot (const T1 *A, T2 *Ainv, int batch)\n{\n#define A(row,col) A[(col)*N+(row)]\n#define AA(row,col) AA[(col)*N+(row)]\n#define Ainv(row,col) Ainv[(col)*N+(row)]\n const int blkNum = blockIdx.x * blockDim.x + threadIdx.x;\n int perm0, perm1, perm2, perm3;\n int icol0, icol1, icol2, icol3;\n T2 AA00, AA01, AA02, AA03, AA10, AA11, AA12, AA13;\n T2 AA20, AA21, AA22, AA23, AA30, AA31, AA32, AA33;\n T2 p, t;\n int i, pvt;\n A += blkNum * N * N;\n Ainv += blkNum * N * N;\n\n if (blkNum < batch)\n {\n AA00 = A[0];\n AA10 = A[1];\n AA20 = A[2];\n AA30 = A[3];\n AA01 = A[4];\n AA11 = A[5];\n AA21 = A[6];\n AA31 = A[7];\n AA02 = A[8];\n AA12 = A[9];\n AA22 = A[10];\n AA32 = A[11];\n AA03 = A[12];\n AA13 = A[13];\n AA23 = A[14];\n AA33 = A[15];\n perm0 = 0;\n perm1 = 1;\n perm2 = 2;\n perm3 = 3;\n /****************** iteration 0 ***********/\n /* search pivot row */\n p = absOp (AA00);\n pvt = 0;\n t = absOp (AA10);\n\n if (t > p)\n {\n p = t;\n pvt = 1;\n }\n\n t = absOp (AA20);\n\n if (t > p)\n {\n p = t;\n pvt = 2;\n }\n\n t = absOp (AA30);\n\n if (t > p)\n {\n p = t;\n pvt = 3;\n }\n\n /* swap pivot row with row 0 */\n if (pvt == 1)\n {\n t = AA00;\n AA00 = AA10;\n AA10 = t;\n t = AA01;\n AA01 = AA11;\n AA11 = t;\n t = AA02;\n AA02 = AA12;\n AA12 = t;\n t = AA03;\n AA03 = AA13;\n AA13 = t;\n /* update permutation vector based on row swap */\n i = perm0;\n perm0 = perm1;\n perm1 = i;\n }\n\n if (pvt == 2)\n {\n t = AA00;\n AA00 = AA20;\n AA20 = t;\n t = AA01;\n AA01 = AA21;\n AA21 = t;\n t = AA02;\n AA02 = AA22;\n AA22 = t;\n t = AA03;\n AA03 = AA23;\n AA23 = t;\n /* update permutation vector based on row swap */\n i = perm0;\n perm0 = perm2;\n perm2 = i;\n }\n\n if (pvt == 3)\n {\n t = AA00;\n AA00 = AA30;\n AA30 = t;\n t = AA01;\n AA01 = AA31;\n AA31 = t;\n t = AA02;\n AA02 = AA32;\n AA32 = t;\n t = AA03;\n AA03 = AA33;\n AA33 = t;\n /* update permutation vector based on row swap */\n i = perm0;\n perm0 = perm3;\n perm3 = i;\n }\n\n /* scale current row */\n t = rcpOp (AA00);\n icol0 = perm0;\n AA00 = t;\n AA01 = mulOp (t, AA01);\n AA02 = mulOp (t, AA02);\n AA03 = mulOp (t, AA03);\n /* eliminate above and below current row */\n t = AA10;\n AA10 = mulOp (-t, AA00);\n AA11 = fmnaOp (t, AA01, AA11);\n AA12 = fmnaOp (t, AA02, AA12);\n AA13 = fmnaOp (t, AA03, AA13);\n t = AA20;\n AA20 = mulOp (-t, AA00);\n AA21 = fmnaOp (t, AA01, AA21);\n AA22 = fmnaOp (t, AA02, AA22);\n AA23 = fmnaOp (t, AA03, AA23);\n t = AA30;\n AA30 = mulOp (-t, AA00);\n AA31 = fmnaOp (t, AA01, AA31);\n AA32 = fmnaOp (t, AA02, AA32);\n AA33 = fmnaOp (t, AA03, AA33);\n /****************** iteration 1 ***********/\n /* search pivot row */\n p = absOp (AA11);\n pvt = 1;\n t = absOp (AA21);\n\n if (t > p)\n {\n p = t;\n pvt = 2;\n }\n\n t = absOp (AA31);\n\n if (t > p)\n {\n p = t;\n pvt = 3;\n }\n\n /* swap pivot row with row 1 */\n if (pvt == 2)\n {\n t = AA10;\n AA10 = AA20;\n AA20 = t;\n t = AA11;\n AA11 = AA21;\n AA21 = t;\n t = AA12;\n AA12 = AA22;\n AA22 = t;\n t = AA13;\n AA13 = AA23;\n AA23 = t;\n /* update permutation vector based on row swap */\n i = perm1;\n perm1 = perm2;\n perm2 = i;\n }\n else if (pvt == 3)\n {\n t = AA10;\n AA10 = AA30;\n AA30 = t;\n t = AA11;\n AA11 = AA31;\n AA31 = t;\n t = AA12;\n AA12 = AA32;\n AA32 = t;\n t = AA13;\n AA13 = AA33;\n AA33 = t;\n /* update permutation vector based on row swap */\n i = perm1;\n perm1 = perm3;\n perm3 = i;\n }\n\n /* scale current row */\n t = rcpOp (AA11);\n icol1 = perm1;\n AA10 = mulOp (t, AA10);\n AA11 = t;\n AA12 = mulOp (t, AA12);\n AA13 = mulOp (t, AA13);\n /* eliminate above and below current row */\n t = AA01;\n AA00 = fmnaOp (t, AA10, AA00);\n AA01 = mulOp (-t, AA11);\n AA02 = fmnaOp (t, AA12, AA02);\n AA03 = fmnaOp (t, AA13, AA03);\n t = AA21;\n AA20 = fmnaOp (t, AA10, AA20);\n AA21 = mulOp (-t, AA11);\n AA22 = fmnaOp (t, AA12, AA22);\n AA23 = fmnaOp (t, AA13, AA23);\n t = AA31;\n AA30 = fmnaOp (t, AA10, AA30);\n AA31 = mulOp (-t, AA11);\n AA32 = fmnaOp (t, AA12, AA32);\n AA33 = fmnaOp (t, AA13, AA33);\n /****************** iteration 2 ****************/\n /* search pivot row */\n p = absOp (AA22);\n pvt = 2;\n t = absOp (AA32);\n\n if (t > p)\n {\n p = t;\n pvt = 3;\n }\n\n /* swap pivot row with row 2 */\n if (pvt == 3)\n {\n t = AA20;\n AA20 = AA30;\n AA30 = t;\n t = AA21;\n AA21 = AA31;\n AA31 = t;\n t = AA22;\n AA22 = AA32;\n AA32 = t;\n t = AA23;\n AA23 = AA33;\n AA33 = t;\n /* update permutation vector based on row swap */\n i = perm2;\n perm2 = perm3;\n perm3 = i;\n }\n\n /* scale current row */\n t = rcpOp (AA22);\n icol2 = perm2;\n AA20 = mulOp (t, AA20);\n AA21 = mulOp (t, AA21);\n AA22 = t;\n AA23 = mulOp (t, AA23);\n /* eliminate above and below current row */\n t = AA02;\n AA00 = fmnaOp (t, AA20, AA00);\n AA01 = fmnaOp (t, AA21, AA01);\n AA02 = mulOp (-t, AA22);\n AA03 = fmnaOp (t, AA23, AA03);\n t = AA12;\n AA10 = fmnaOp (t, AA20, AA10);\n AA11 = fmnaOp (t, AA21, AA11);\n AA12 = mulOp (-t, AA22);\n AA13 = fmnaOp (t, AA23, AA13);\n t = AA32;\n AA30 = fmnaOp (t, AA20, AA30);\n AA31 = fmnaOp (t, AA21, AA31);\n AA32 = mulOp (-t, AA22);\n AA33 = fmnaOp (t, AA23, AA33);\n /****************** iteration 3 ****************/\n /* scale current row */\n t = rcpOp (AA33);\n icol3 = perm3;\n AA30 = mulOp (t, AA30);\n AA31 = mulOp (t, AA31);\n AA32 = mulOp (t, AA32);\n AA33 = t;\n /* eliminate above and below current row */\n t = AA03;\n AA00 = fmnaOp (t, AA30, AA00);\n AA01 = fmnaOp (t, AA31, AA01);\n AA02 = fmnaOp (t, AA32, AA02);\n AA03 = mulOp (-t, AA33);\n t = AA13;\n AA10 = fmnaOp (t, AA30, AA10);\n AA11 = fmnaOp (t, AA31, AA11);\n AA12 = fmnaOp (t, AA32, AA12);\n AA13 = mulOp (-t, AA33);\n t = AA23;\n AA20 = fmnaOp (t, AA30, AA20);\n AA21 = fmnaOp (t, AA31, AA21);\n AA22 = fmnaOp (t, AA32, AA22);\n AA23 = mulOp (-t, AA33);\n /* sort columns into the correct order */\n Ainv(0, icol0) = AA00;\n Ainv(1, icol0) = AA10;\n Ainv(2, icol0) = AA20;\n Ainv(3, icol0) = AA30;\n Ainv(0, icol1) = AA01;\n Ainv(1, icol1) = AA11;\n Ainv(2, icol1) = AA21;\n Ainv(3, icol1) = AA31;\n Ainv(0, icol2) = AA02;\n Ainv(1, icol2) = AA12;\n Ainv(2, icol2) = AA22;\n Ainv(3, icol2) = AA32;\n Ainv(0, icol3) = AA03;\n Ainv(1, icol3) = AA13;\n Ainv(2, icol3) = AA23;\n Ainv(3, icol3) = AA33;\n }\n}\n\n\ntemplate\n__global__ void matinv_matrix_per_thread_no_pivot (const T *A, T *Ainv, int batch)\n{\n#define A(row,col) A[(col)*N+(row)]\n#define AA(row,col) AA[(col)*N+(row)]\n#define Ainv(row,col) Ainv[(col)*N+(row)]\n const int blkNum = blockIdx.x * blockDim.x + threadIdx.x;\n T AA00, AA01, AA02, AA03, AA10, AA11, AA12, AA13;\n T AA20, AA21, AA22, AA23, AA30, AA31, AA32, AA33;\n T t;\n A += blkNum * N * N;\n Ainv += blkNum * N * N;\n\n if (blkNum < batch)\n {\n AA00 = A[0];\n AA10 = A[1];\n AA20 = A[2];\n AA30 = A[3];\n AA01 = A[4];\n AA11 = A[5];\n AA21 = A[6];\n AA31 = A[7];\n AA02 = A[8];\n AA12 = A[9];\n AA22 = A[10];\n AA32 = A[11];\n AA03 = A[12];\n AA13 = A[13];\n AA23 = A[14];\n AA33 = A[15];\n /****************** iteration 0 ***********/\n /* search pivot row */\n t = 1.0 / (AA00);\n AA00 = t;\n AA01 = t * AA01;\n AA02 = t * AA02;\n AA03 = t * AA03;\n /* eliminate above and below current row */\n t = AA10;\n AA10 = -t * AA00;\n AA11 = fmnaOp (t, AA01, AA11);\n AA12 = fmnaOp (t, AA02, AA12);\n AA13 = fmnaOp (t, AA03, AA13);\n t = AA20;\n AA20 = -t * AA00;\n AA21 = fmnaOp (t, AA01, AA21);\n AA22 = fmnaOp (t, AA02, AA22);\n AA23 = fmnaOp (t, AA03, AA23);\n t = AA30;\n AA30 = -t * AA00;\n AA31 = fmnaOp (t, AA01, AA31);\n AA32 = fmnaOp (t, AA02, AA32);\n AA33 = fmnaOp (t, AA03, AA33);\n /****************** iteration 1 ***********/\n /* scale current row */\n t = 1.0 / (AA11);\n AA10 = t * AA10;\n AA11 = t;\n AA12 = t * AA12;\n AA13 = t * AA13;\n /* eliminate above and below current row */\n t = AA01;\n AA00 = fmnaOp (t, AA10, AA00);\n AA01 = -t * AA11;\n AA02 = fmnaOp (t, AA12, AA02);\n AA03 = fmnaOp (t, AA13, AA03);\n t = AA21;\n AA20 = fmnaOp (t, AA10, AA20);\n AA21 = -t * AA11;\n AA22 = fmnaOp (t, AA12, AA22);\n AA23 = fmnaOp (t, AA13, AA23);\n t = AA31;\n AA30 = fmnaOp (t, AA10, AA30);\n AA31 = -t * AA11;\n AA32 = fmnaOp (t, AA12, AA32);\n AA33 = fmnaOp (t, AA13, AA33);\n /****************** iteration 2 ****************/\n /* scale current row */\n t = 1.0 / (AA22);\n AA20 = t * AA20;\n AA21 = t * AA21;\n AA22 = t;\n AA23 = t * AA23;\n /* eliminate above and below current row */\n t = AA02;\n AA00 = fmnaOp (t, AA20, AA00);\n AA01 = fmnaOp (t, AA21, AA01);\n AA02 = -t * AA22;\n AA03 = fmnaOp (t, AA23, AA03);\n t = AA12;\n AA10 = fmnaOp (t, AA20, AA10);\n AA11 = fmnaOp (t, AA21, AA11);\n AA12 = -t * AA22;\n AA13 = fmnaOp (t, AA23, AA13);\n t = AA32;\n AA30 = fmnaOp (t, AA20, AA30);\n AA31 = fmnaOp (t, AA21, AA31);\n AA32 = -t * AA22;\n AA33 = fmnaOp (t, AA23, AA33);\n /****************** iteration 3 ****************/\n /* scale current row */\n t = 1.0 / (AA33);\n AA30 = t * AA30;\n AA31 = t * AA31;\n AA32 = t * AA32;\n AA33 = t;\n /* eliminate above and below current row */\n t = AA03;\n AA00 = fmnaOp (t, AA30, AA00);\n AA01 = fmnaOp (t, AA31, AA01);\n AA02 = fmnaOp (t, AA32, AA02);\n AA03 = -t * AA33;\n t = AA13;\n AA10 = fmnaOp (t, AA30, AA10);\n AA11 = fmnaOp (t, AA31, AA11);\n AA12 = fmnaOp (t, AA32, AA12);\n AA13 = -t * AA33;\n t = AA23;\n AA20 = fmnaOp (t, AA30, AA20);\n AA21 = fmnaOp (t, AA31, AA21);\n AA22 = fmnaOp (t, AA32, AA22);\n AA23 = -t * AA33;\n /* sort columns into the correct order */\n Ainv(0, 0) = AA00;\n Ainv(1, 0) = AA10;\n Ainv(2, 0) = AA20;\n Ainv(3, 0) = AA30;\n Ainv(0, 1) = AA01;\n Ainv(1, 1) = AA11;\n Ainv(2, 1) = AA21;\n Ainv(3, 1) = AA31;\n Ainv(0, 2) = AA02;\n Ainv(1, 2) = AA12;\n Ainv(2, 2) = AA22;\n Ainv(3, 2) = AA32;\n Ainv(0, 3) = AA03;\n Ainv(1, 3) = AA13;\n Ainv(2, 3) = AA23;\n Ainv(3, 3) = AA33;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockJacobiSmoothbBigBlockDiaCsrKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows, int bsize, int bsize_sq, ValueTypeA *temp1)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int halfwarp_id = tid >> 4;\n const int block_halfwarp_id = threadIdx.x >> 4;\n const int mat_entry_index = threadIdx.x & (16 - 1);\n const int i_ind = mat_entry_index >> 2;\n const int j_ind = mat_entry_index & 3;\n extern __shared__ volatile char schar[];\n volatile ValueTypeA *s_Amat;\n s_Amat = (ValueTypeA *)&schar[0];\n int tile_num = (bsize - 1) / 4 + 1;\n ValueTypeA *e_out = &temp1[(blockIdx.x * blockDim.x + threadIdx.x) * tile_num * tile_num];\n\n while (halfwarp_id < num_block_rows)\n {\n int offset = halfwarp_id * bsize_sq + i_ind * bsize + j_ind;\n int s_offset = block_halfwarp_id * bsize_sq;\n\n // Store the diagonal\n for (int t1 = 0; t1 < tile_num; t1++)\n for (int t2 = 0; t2 < tile_num; t2++)\n if ((t1 * 4 + i_ind) < bsize && (t2 * 4 + j_ind) < bsize)\n {\n e_out[t1 * tile_num + t2] = values[bsize_sq * dia_indices[halfwarp_id] + (t1 * 4 + i_ind) * bsize + t2 * 4 + j_ind];\n }\n\n // Each thread stores its entry in s_Amat\n for (int t1 = 0; t1 < tile_num; t1++)\n for (int t2 = 0; t2 < tile_num; t2++)\n if ((t1 * 4 + i_ind) < bsize && (t2 * 4 + j_ind) < bsize)\n {\n types::util::volcast( e_out[t1 * tile_num + t2], s_Amat + (s_offset + (t1 * 4 + i_ind) * bsize + t2 * 4 + j_ind) );\n }\n\n compute_block_inverse2\n ( s_Amat, s_offset, offset, i_ind, j_ind, Dinv, tile_num, bsize, bsize_sq );\n halfwarp_id += gridDim.x * halfwarps_per_block;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockJacobiSmoothBbyBBlockDiaCsrKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows)\n\n{\n int warp_id = threadIdx.x / 32;\n int warp_thread_id = threadIdx.x & 31;\n\n // padding row blocks to fit in a single warp\n if ( warp_thread_id >= blockrows_per_warp * bsize_sq ) { return; }\n\n // new thread id with padding\n int tid = warp_id * blockrows_per_warp * bsize_sq + warp_thread_id;\n int cta_blockrow_id = tid / bsize_sq;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int mat_entry_index = tid - cta_blockrow_id * bsize_sq;\n const int i_ind = mat_entry_index / bsize;\n const int j_ind = mat_entry_index - i_ind * bsize;\n volatile __shared__ ValueTypeA s_Amat[bsize_sq * blockrows_per_cta ];\n ValueTypeA e_out;\n\n while (blockrow_id < num_block_rows)\n {\n int offset = blockrow_id * bsize_sq + mat_entry_index;\n // Store the diagonal\n e_out = values[bsize_sq * dia_indices[blockrow_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n types::util::volcast(e_out, s_Amat + tid);\n compute_block_inverse_row_major\n ( s_Amat, cta_blockrow_id * bsize_sq, offset, i_ind, j_ind, Dinv );\n blockrow_id += gridDim.x * blockrows_per_cta;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockJacobiSmooth4by4BlockDiaCsrKernel_V2(const IndexType *dia_indices, const ValueTypeA *A_values, ValueTypeA *Dinv, const int num_block_rows)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int halfwarp_id = tid >> log_bsize_sq;\n const int block_halfwarp_id = threadIdx.x >> log_bsize_sq;\n const int mat_entry_index = threadIdx.x & (bsize_sq - 1);\n const int i_ind = mat_entry_index >> log_bsize;\n const int j_ind = mat_entry_index & (bsize - 1);\n volatile __shared__ ValueTypeA s_Amat[bsize_sq * halfwarps_per_block ];\n int offset;\n ValueTypeA e_out;\n\n while (halfwarp_id < num_block_rows)\n {\n // Store the diagonal\n offset = halfwarp_id * bsize_sq;\n e_out = A_values[bsize_sq * dia_indices[halfwarp_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n types::util::volcast(e_out, s_Amat + threadIdx.x);\n compute_block_inverse_row_major\n ( s_Amat, block_halfwarp_id * bsize_sq, offset + mat_entry_index, i_ind, j_ind, Dinv );\n halfwarp_id += gridDim.x * blockDim.x >> log_bsize_sq;\n }\n}\n\n// Kernel to smooth with Jacobi smoother, Dinv assumed to be computed\ntemplate\n__global__\nvoid jacobiSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *dia_indices, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, double weight, const int num_block_rows, ValueTypeB *xout, const int row_offset)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int eighthwarp_id = row_offset + (tid >> log_bsize);\n const int block_eighthwarp_id = threadIdx.x >> log_bsize;\n const int vec_entry_index = threadIdx.x & (bsize - 1);\n volatile __shared__ ValueTypeB s_xtemp[ bsize * eighthwarps_per_block ];\n ValueTypeB bmAx, xin;\n ValueTypeB temp[bsize];\n int offset, i, s_offset;\n\n while (eighthwarp_id < num_block_rows)\n {\n i = eighthwarp_id;\n offset = i * bsize + vec_entry_index;\n // 1. COMPUTING b-Ax\n bmAx = b[offset];\n // Contribution from diagonal\n xin = x[offset];\n types::util::volcast(xin, s_xtemp + threadIdx.x);\n // Load dia_values and do matrix multiply\n loadAsVector(nonzero_values + bsize_sq * dia_indices[i] + vec_entry_index * bsize, temp);\n s_offset = block_eighthwarp_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx = bmAx - temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = __cachingLoad(&column_indices[jind]);\n\n if (jcol != i)\n {\n offset = jcol * bsize + vec_entry_index;\n types::util::volcast(x[offset], s_xtemp + threadIdx.x);\n // Load nonzero_values\n s_offset = block_eighthwarp_id * bsize;\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx = bmAx - temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n }\n }\n\n types::util::volcast(bmAx, s_xtemp + threadIdx.x);\n bmAx = types::util::get_zero();\n // 2. Multiply by Dinv\n // Load Dinv and multiply to RHS\n offset = i * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(Dinv + offset, temp);\n s_offset = block_eighthwarp_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx = bmAx + temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n\n xout[i * bsize + vec_entry_index] = xin + bmAx * weight;\n eighthwarp_id += gridDim.x * blockDim.x >> log_bsize;\n }\n}\n\n// Kernel to smooth with jacobi smoother, zero initial guess\ntemplate\n__global__\nvoid jacobiSmooth4by4ZeroBlockDiaCsrKernel_NAIVE_tex_readDinv2(const ValueTypeA *Dinv, const ValueTypeB *b, double weight, const int num_block_rows, ValueTypeB *xout, const int row_offset)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n const int block_eighthwarp_id = threadIdx.x >> log_bsize;\n const int vec_entry_index = threadIdx.x & (bsize - 1);\n volatile __shared__ ValueTypeB s_xtemp[ bsize * eighthwarps_per_block ];\n ValueTypeB bmAx;\n ValueTypeB temp[bsize];\n int offset, i, s_offset;\n\n for (int eighthwarp_id = row_offset + (tid >> log_bsize); eighthwarp_id < num_block_rows; eighthwarp_id += (gridDim.x * blockDim.x >> log_bsize))\n {\n i = eighthwarp_id;\n offset = i * bsize + vec_entry_index;\n types::util::volcast(b[offset], s_xtemp + threadIdx.x);\n bmAx = types::util::get_zero();\n // Load Dinv and multiply to RHS\n offset = i * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(Dinv + offset, temp);\n s_offset = block_eighthwarp_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx = bmAx + temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n\n xout[i * bsize + vec_entry_index] = bmAx * weight;\n }\n}\n\n//--------------------------------\n// Methods\n//--------------------------------\n\n// Constructor\ntemplate\nBlockJacobiSolver_Base::BlockJacobiSolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope)\n{\n weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n\n if (weight == 0)\n {\n weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Block Jacobi smoother\\n\");\n }\n}\n\n// Destructor\ntemplate\nBlockJacobiSolver_Base::~BlockJacobiSolver_Base()\n{\n this->Dinv.resize(0);\n}\n\ntemplate\nvoid\nBlockJacobiSolver_Base::printSolverParameters() const\n{\n std::cout << \"relaxation_factor= \" << this->weight << std::endl;\n}\n\n// Solver setup\ntemplate\nvoid\nBlockJacobiSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n Matrix *A_as_matrix = dynamic_cast*>(this->m_A);\n\n if (!A_as_matrix)\n {\n FatalError(\"BlockJacobiSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n computeDinv( *A_as_matrix );\n\n if ( A_as_matrix->getBlockFormat() != ROW_MAJOR )\n {\n FatalError(\"Block Jacobi solver only supports row major format\", AMGX_ERR_CONFIGURATION);\n }\n}\n\n//\ntemplate\nvoid\nBlockJacobiSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nBlockJacobiSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n //bool done = false;\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n\n if (xIsZero) { x.dirtybit = 0; }\n\n ViewType oldView = A_as_matrix->currentView();\n\n A_as_matrix->setViewExterior();\n ViewType flags = (ViewType)(A_as_matrix->getViewInterior() | A_as_matrix->getViewExterior());\n\n if (A_as_matrix->get_block_dimx() == 1 && A_as_matrix->get_block_dimy() == 1)\n {\n if (xIsZero)\n {\n smooth_with_0_initial_guess_1x1(*A_as_matrix, b, x, flags);\n }\n else\n {\n smooth_1x1(*A_as_matrix, b, x, flags);\n }\n }\n else if (A_as_matrix->get_block_dimx() == 4 && A_as_matrix->get_block_dimy() == 4)\n {\n if (xIsZero)\n {\n smooth_with_0_initial_guess_4x4(*A_as_matrix, b, x, flags);\n }\n else\n {\n smooth_4x4(*A_as_matrix, b, x, flags);\n }\n }\n else if (A_as_matrix->get_block_dimx() == A_as_matrix->get_block_dimy())\n {\n if (xIsZero)\n {\n thrust_wrapper::fill(x.begin(), x.end(), types::util::get_zero());\n cudaCheckError();\n }\n\n smooth_BxB(*A_as_matrix, b, x, true, flags);\n }\n else\n {\n FatalError(\"Unsupported block size for BlockJacobi_Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if (A_as_matrix->get_block_dimx() == 4 && A_as_matrix->get_block_dimy() == 4)\n {\n if (!xIsZero) // we write to t_res vector to avoid race condition in the kernel\n {\n x.swap(this->t_res);\n }\n }\n\n x.dirtybit = 1;\n A_as_matrix->setView(oldView);\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nBlockJacobiSolver_Base::solve_finalize( VVector &b, VVector &x )\n{}\n\ntemplate\nvoid BlockJacobiSolver_Base::computeDinv( Matrix &A)\n{\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n ViewType oldView = A.currentView();\n A.setView(A_as_matrix->getViewExterior());\n\n if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n this->computeDinv_1x1(A);\n }\n else if (A.get_block_dimx() == 2 && A.get_block_dimy() == 2)\n {\n this->computeDinv_bxb<2>(A);\n }\n else if (A.get_block_dimx() == 3 && A.get_block_dimy() == 3)\n {\n this->computeDinv_3x3(A);\n }\n else if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n this->computeDinv_4x4(A);\n }\n else if (A.get_block_dimx() == 5 && A.get_block_dimy() == 5)\n {\n this->computeDinv_bxb<5>(A);\n }\n else if (A.get_block_dimx() == A.get_block_dimy() && A.get_block_dimy() > 5)\n {\n this->computeDinv_Big(A, A.get_block_dimx());\n }\n\n A.setView(oldView);\n}\n\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid BlockJacobiSolver >::computeDinv_1x1(const Matrix_d &A)\n{\n Matrix_d *A_as_matrix = (Matrix_d *) this->m_A;\n// supports both diag\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), types::util::get_zero());\n\n if ( A_as_matrix->hasProps(DIAG) )\n {\n const int num_values = A_as_matrix->diagOffset() * A_as_matrix->get_block_size();\n amgx::thrust::copy( A_as_matrix->values.begin() + num_values, A_as_matrix->values.begin() + num_values + A_as_matrix->get_num_rows()*A_as_matrix->get_block_size(), this->Dinv.begin() );\n cudaCheckError();\n }\n else\n {\n find_diag( *A_as_matrix );\n }\n}\n\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid BlockJacobiSolver >::computeDinv_1x1(const Matrix_h &A)\n{\n // Do nothing\n}\n\n\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid BlockJacobiSolver >::computeDinv_4x4(const Matrix_h &A)\n{\n //FatalError(\"Block Jacobi smoother not implemented with this format, exiting\");\n //std::cout << \"Warning, CPU version doesn't store the inverse of the blocks, like the GPU version\" << std::endl;\n}\n\ntemplate \nvoid BlockJacobiSolver >::computeDinv_3x3(const Matrix_h &A)\n{\n FatalError(\"3*3 Block Jacobi smoother not implemented with host format, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n //std::cout << \"Warning, CPU version doesn't store the inverse of the blocks, like the GPU version\" << std::endl;\n}\n\ntemplate \nvoid BlockJacobiSolver >::computeDinv_Big(const Matrix_h &A, const int bsize)\n{\n FatalError(\"Big Block Jacobi smoother not implemented with host format, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n //std::cout << \"Warning, CPU version doesn't store the inverse of the blocks, like the GPU version\" << std::endl;\n}\n\n\n// Finding diag on device, CSR format\ntemplate \nvoid BlockJacobiSolver >::find_diag( const Matrix_h &A )\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n if (A.col_indices[j] == i)\n {\n this->Dinv[i] = A.values[j];\n break;\n }\n\n if (j == A.row_offsets[i + 1] - 1)\n {\n FatalError(\"Could not find a diagonal value\", AMGX_ERR_BAD_PARAMETERS);\n }\n }\n }\n}\n\n\n// Finding diag on device, CSR format\ntemplate \nvoid BlockJacobiSolver >::find_diag( const Matrix_d &A )\n{\n AMGX_CPU_PROFILER( \"JacobiSolver::find_diag \" );\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((float)(A.get_num_rows()) / (float)(THREADS_PER_BLOCK)));\n find_diag_kernel_indexed_dia <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK>>>(\n A.get_num_rows(),\n A.diag.raw(),\n A.values.raw(),\n this->Dinv.raw());\n cudaCheckError();\n}\n\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid BlockJacobiSolver >::computeDinv_4x4(const Matrix_d &A)\n{\n// supports both diag\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), types::util::get_zero());\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_indices_ptr = A.diag.raw();\n const ValueTypeA *A_values = A.values.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n#if 1\n const int threads_per_block = 512;\n const int halfwarps_per_block = threads_per_block / 16;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / halfwarps_per_block + 1);\n cudaFuncSetCacheConfig(setupBlockJacobiSmooth4by4BlockDiaCsrKernel_V2, cudaFuncCachePreferL1);\n setupBlockJacobiSmooth4by4BlockDiaCsrKernel_V2 <<< num_blocks, threads_per_block>>>\n (A_dia_indices_ptr, A_values, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n#else\n cudaFuncSetCacheConfig(matinv_matrix_per_thread_pivot, cudaFuncCachePreferL1);\n matinv_matrix_per_thread_pivot <<< (A.num_block_rows + 127) / 128, 128 >>> (A_dia_values_ptr, Dinv_ptr, A.num_block_rows);\n cudaCheckError();\n#endif\n}\n\ntemplate \nvoid BlockJacobiSolver >::computeDinv_3x3(const Matrix_d &A)\n{\n const int bsize = 3;\n// supports both diag\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), types::util::get_zero());\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 256;\n const int blockrows_per_warp = 32 / (bsize * bsize);\n // blockrows per cta = blockrows_per_warp * number_of_warps_per_cta\n const int blockrows_per_cta = (threads_per_block / 32) * blockrows_per_warp ;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n cudaFuncSetCacheConfig(setupBlockJacobiSmoothBbyBBlockDiaCsrKernel, cudaFuncCachePreferL1);\n setupBlockJacobiSmoothBbyBBlockDiaCsrKernel <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid BlockJacobiSolver >::computeDinv_Big(const Matrix_d &A, const int bsize)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), types::util::get_zero());\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n MVector temp(AMGX_GRID_MAX_SIZE * ((bsize - 1) / 4 + 1) * ((bsize - 1) / 4 + 1));\n ValueTypeA *temp_ptr = temp.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 512;\n const int halfwarps_per_block = threads_per_block / 16;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / halfwarps_per_block + 1);\n cudaFuncSetCacheConfig(setupBlockJacobiSmoothbBigBlockDiaCsrKernel, cudaFuncCachePreferL1);\n setupBlockJacobiSmoothbBigBlockDiaCsrKernel <<< num_blocks, threads_per_block, sizeof(ValueTypeA)*bsize *bsize *halfwarps_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows(), bsize, bsize * bsize, temp_ptr);\n cudaCheckError();\n}\n\n\ntemplate\ntemplate\nvoid BlockJacobiSolver_Base::computeDinv_bxb(const Matrix &A)\n{\n if (TConfig::memSpace == AMGX_host)\n {\n FatalError(\"BlockJacobiSmooth Not implemented for host\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n // supports both diag\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), types::util::get_zero());\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 256;\n const int blockrows_per_warp = 32 / (bsize * bsize);\n const int blockrows_per_cta = (threads_per_block / 32) * blockrows_per_warp;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n cudaFuncSetCacheConfig(setupBlockJacobiSmoothBbyBBlockDiaCsrKernel, cudaFuncCachePreferL1);\n setupBlockJacobiSmoothBbyBBlockDiaCsrKernel <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n }\n}\n\ntemplate \nvoid BlockJacobiSolver >::smooth_BxB(Matrix_h &A, VVector &b, VVector &x, bool firstStep, ViewType separation_flags)\n{\n FatalError(\"M*M Block Jacobi smoother not implemented with host format, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n\ntemplate \nvoid BlockJacobiSolver >::smooth_4x4(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n VVector newx(x.size());\n int bsize = A.get_block_dimx();\n // Allocate space for block_matrix\n ValueTypeA **E = new ValueTypeA* [bsize];\n\n for ( int i = 0; i < bsize; i++)\n {\n E[i] = new ValueTypeA[bsize];\n }\n\n ValueTypeB *bmAx = new ValueTypeB[bsize];\n ValueTypeB *temp = new ValueTypeB[bsize];\n\n //for each block row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n // Compute b - sum_j A_j x_j (denoted by bmAx) for block_row i\n // Load diagonal\n for (int m = 0; m < bsize; m++)\n {\n for (int n = 0; n < bsize; n++)\n {\n E[m][n] = A.values[A.diag[i * bsize * bsize + bsize * m + n]];\n }\n\n bmAx[m] = types::util::get_zero();\n temp[m] = types::util::get_zero();\n }\n\n // Contribution from diagonal\n for (int m = 0; m < bsize; m++)\n for (int n = 0; n < bsize; n++)\n {\n bmAx[m] = bmAx[m] - E[m][n] * x[i * bsize + n];\n }\n\n // Contribution from each nonzero column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n IndexType jcol = A.col_indices[j];\n\n for (int m = 0; m < bsize; m++)\n for (int n = 0; n < bsize; n++)\n {\n bmAx[m] = bmAx[m] - A.values[j * bsize * bsize + bsize * m + n] * x[jcol * bsize + n];\n }\n }\n\n // Add b\n for (int m = 0; m < bsize; m++)\n {\n bmAx[m] = bmAx[m] + b[i * bsize + m];\n }\n\n gaussianEliminationRowMajor(E, temp, bmAx, bsize);\n\n // Compute new value of x\n for (int m = 0; m < bsize; m++)\n {\n newx[i * bsize + m] = x[i * bsize + m] + temp[m] * this->weight;\n }\n }\n\n x.swap(newx);\n}\n\n\n\ntemplate \nvoid BlockJacobiSolver >::smooth_with_0_initial_guess_4x4(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n IndexType bsize = A.get_block_dimy();\n ValueTypeA **E = new ValueTypeA* [bsize];\n\n for ( int i = 0; i < bsize; i++)\n {\n E[i] = new ValueTypeA[bsize];\n }\n\n ValueTypeB *rhs = new ValueTypeB[bsize];\n ValueTypeB *temp = new ValueTypeB[bsize];\n\n //for each block row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n // Load diagonal\n for (int m = 0; m < bsize; m++)\n {\n for (int n = 0; n < bsize; n++)\n {\n E[m][n] = A.values[A.diag[i * bsize * bsize + bsize * m + n]];\n }\n\n rhs[m] = types::util::get_zero();\n }\n\n //rhs\n for (int m = 0; m < bsize; m++)\n {\n rhs[m] = rhs[m] + b[i * bsize + m];\n }\n\n // Solve for temp\n gaussianEliminationRowMajor(E, temp, rhs, bsize);\n\n for (int m = 0; m < bsize; m++)\n {\n x[i * bsize + m] = temp[m] * this->weight;\n }\n }\n}\n\ntemplate \nvoid BlockJacobiSolver >::smooth_1x1_const(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n VVector newx(x.size());\n\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeB Axi = types::util::get_zero();\n ValueTypeB d = types::util::get_one() * A.values[A.diag[i]];\n ValueTypeB mydiaginv = types::util::get_one() * this->weight / (isNotCloseToZero(d) ? d : epsilon(d) );\n\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n Axi = Axi + A.values[j] * x[A.col_indices[j]];\n }\n\n newx[i] = x[i] + (b[i] - Axi) * mydiaginv ;\n }\n\n x.swap(newx);\n}\n\ntemplate \nvoid BlockJacobiSolver >::smooth_1x1(Matrix_h &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n this->smooth_1x1_const(A, b, x, separation_flags);\n}\n\ntemplate \nvoid BlockJacobiSolver >::smooth_1x1(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n AMGX_CPU_PROFILER( \"JacobiSolver::smooth_1x1 \" );\n\n if (this->t_res.size() != x.size())\n {\n this->t_res.resize(x.size());\n }\n\n if (this->y.size() != b.size())\n {\n this->y.resize(b.size());\n this->y.tag = this->tag * 100 + 3;\n this->y.set_block_dimx(b.get_block_dimx());\n this->y.set_block_dimy(b.get_block_dimy());\n }\n\n int num_rows = A.get_num_rows();\n int offset = 0;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n this->y.dirtybit = 0;\n\n multiply( A, x, this->y, separation_flags );\n\n thrust_wrapper::transform( \n amgx::thrust::make_zip_iterator(\n amgx::thrust::make_tuple(x.begin() + offset, this->Dinv.begin() + offset, b.begin() + offset, this->y.begin() + offset)), \n amgx::thrust::make_zip_iterator(\n amgx::thrust::make_tuple(x.begin() + A.get_num_rows(), this->Dinv.begin() + A.get_num_rows(), b.begin() + A.get_num_rows(), this->y.begin() + A.get_num_rows())),\n x.begin() + offset,\n jacobi_postsmooth_functor( this->weight ));\n cudaCheckError();\n}\n\n\ntemplate \nvoid BlockJacobiSolver >::smooth_with_0_initial_guess_1x1(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeB d = types::util::get_one() * A.values[A.diag[i]];\n ValueTypeB mydiag = types::util::get_one() * this->weight / (isNotCloseToZero(d) ? d : epsilon(d));\n x[i] = b[i] * mydiag;\n }\n}\n\ntemplate \nvoid BlockJacobiSolver >::smooth_with_0_initial_guess_1x1(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n AMGX_CPU_PROFILER( \"JacobiSolver::smooth_with_0_initial_guess_1x1 \" );\n int num_rows = A.get_num_rows();\n int offset = 0;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n thrust_wrapper::transform( b.begin( ) + offset,\n b.begin( ) + A.get_num_rows(),\n this->Dinv.begin( ) + offset,\n x.begin( ) + offset,\n jacobi_presmooth_functor( this->weight ));\n cudaCheckError();\n}\n\ntemplate \nvoid BlockJacobiSolver >::smooth_BxB(Matrix_d &A, VVector &b, VVector &x, bool firstStep, ViewType separation_flags)\n{\n IndexType num_rows;\n IndexType offset;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n\n // aux vector initialization\n if (this->y.size() != b.size())\n {\n this->y.resize(b.size());\n this->y.tag = this->tag * 100 + 3;\n this->y.set_block_dimx(b.get_block_dimx());\n this->y.set_block_dimy(b.get_block_dimy());\n }\n\n amgx::thrust::copy(b.begin(), b.end(), this->y.begin()); // copy of vector b\n cudaCheckError();\n Cusparse::bsrmv(types::util::get_minus_one(), A, x, types::util::get_one(), this->y, separation_flags); // y= -1.0f*(A.x) + y\n cudaCheckError();\n\n Cusparse::bsrmv(types::util::get_one() * this->weight, A, this->Dinv, this->y, types::util::get_one(), x, separation_flags); // t_res = t_res + w*(Dinv.y) @ view\n cudaCheckError();\n}\n\ntemplate \nvoid BlockJacobiSolver >::smooth_4x4(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n if (this->t_res.size() != x.size())\n {\n this->t_res.resize(x.size());\n }\n\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n ValueTypeB *x_ptr = x.raw();\n ValueTypeB *xout_ptr = this->t_res.raw(); // always store original x\n IndexType num_rows = A.get_num_rows();\n IndexType offset = 0;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n const int threads_per_block = 512;\n const int eightwarps_per_block = threads_per_block / 4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows - 1) / eightwarps_per_block + 1);\n\n // XXX We need to add latency hiding here\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo(x, x.tag);\n }\n\n cudaFuncSetCacheConfig(jacobiSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2, cudaFuncCachePreferL1);\n jacobiSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, offset + num_rows, xout_ptr, offset);\n cudaCheckError();\n}\n\n\ntemplate \nvoid BlockJacobiSolver >::smooth_with_0_initial_guess_4x4(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n cudaCheckError();\n const ValueTypeA *A_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n IndexType num_rows = A.get_num_rows();\n IndexType offset = 0;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n const int threads_per_block = 512;\n const int eightwarps_per_block = threads_per_block / 4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows - 1) / eightwarps_per_block + 1);\n cudaFuncSetCacheConfig(jacobiSmooth4by4ZeroBlockDiaCsrKernel_NAIVE_tex_readDinv2, cudaFuncCachePreferL1);\n jacobiSmooth4by4ZeroBlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (Dinv_ptr, b_ptr, this->weight, offset + num_rows, x_ptr, offset);\n cudaCheckError();\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class BlockJacobiSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class BlockJacobiSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace block_jacobi\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace amgx\n{\nnamespace block_jacobi_solver\n{\n\ntemplate \nstruct jacobi_presmooth_functor\n{\n double omega;\n jacobi_presmooth_functor( double omega ) : omega( omega ) {}\n __host__ __device__ ValueTypeB operator()( const ValueTypeB &b, const ValueTypeA &d ) const { return isNotCloseToZero(d) ? b * omega / d : b * omega / epsilon(d); }\n};\n\ntemplate \nstruct jacobi_postsmooth_functor\n{\n double omega;\n jacobi_postsmooth_functor( double omega ) : omega( omega ) {}\n template __host__ __device__ ValueTypeB operator( )( const Tuple &t ) const\n {\n ValueTypeB x = amgx::thrust::get<0>(t);\n ValueTypeA d = amgx::thrust::get<1>(t);\n ValueTypeB b = amgx::thrust::get<2>(t);\n ValueTypeB y = amgx::thrust::get<3>(t);\n // return x + omega * (b - y) / d.\n d = isNotCloseToZero(d) ? d : epsilon(d);\n d = types::util::get_one() / d;\n b = b - y;\n b = b * omega;\n return b * d + x;\n }\n};\n\ntemplate \nstruct add_functor\n{\n __host__ __device__ ValueTypeB operator()( const ValueTypeB &x, const ValueTypeB &y )const { return x + y; }\n};\n\ntemplate\n__device__ __forceinline__ T fmnaOp (T a, T b, T c)\n{\n return -(a * b) + c;\n}\ntemplate\n__device__ __forceinline__ T mulOp (T a, T b)\n{\n return a * b;\n}\ntemplate\n__device__ __forceinline__ T rcpOp (T a)\n{\n return 1.0 / (isNotCloseToZero(a) ? a : epsilon(a));\n}\ntemplate\n__device__ __forceinline__ T absOp (T a)\n{\n return fabs(a);\n}\n\n// -----------------------------------\n// KERNELS\n// -----------------------------------\n\ntemplate\n__global__ void matinv_matrix_per_thread_pivot (const T1 *A, T2 *Ainv, int batch)\n{\n#define A(row,col) A[(col)*N+(row)]\n#define AA(row,col) AA[(col)*N+(row)]\n#define Ainv(row,col) Ainv[(col)*N+(row)]\n const int blkNum = blockIdx.x * blockDim.x + threadIdx.x;\n int perm0, perm1, perm2, perm3;\n int icol0, icol1, icol2, icol3;\n T2 AA00, AA01, AA02, AA03, AA10, AA11, AA12, AA13;\n T2 AA20, AA21, AA22, AA23, AA30, AA31, AA32, AA33;\n T2 p, t;\n int i, pvt;\n A += blkNum * N * N;\n Ainv += blkNum * N * N;\n\n if (blkNum < batch)\n {\n AA00 = A[0];\n AA10 = A[1];\n AA20 = A[2];\n AA30 = A[3];\n AA01 = A[4];\n AA11 = A[5];\n AA21 = A[6];\n AA31 = A[7];\n AA02 = A[8];\n AA12 = A[9];\n AA22 = A[10];\n AA32 = A[11];\n AA03 = A[12];\n AA13 = A[13];\n AA23 = A[14];\n AA33 = A[15];\n perm0 = 0;\n perm1 = 1;\n perm2 = 2;\n perm3 = 3;\n /****************** iteration 0 ***********/\n /* search pivot row */\n p = absOp (AA00);\n pvt = 0;\n t = absOp (AA10);\n\n if (t > p)\n {\n p = t;\n pvt = 1;\n }\n\n t = absOp (AA20);\n\n if (t > p)\n {\n p = t;\n pvt = 2;\n }\n\n t = absOp (AA30);\n\n if (t > p)\n {\n p = t;\n pvt = 3;\n }\n\n /* swap pivot row with row 0 */\n if (pvt == 1)\n {\n t = AA00;\n AA00 = AA10;\n AA10 = t;\n t = AA01;\n AA01 = AA11;\n AA11 = t;\n t = AA02;\n AA02 = AA12;\n AA12 = t;\n t = AA03;\n AA03 = AA13;\n AA13 = t;\n /* update permutation vector based on row swap */\n i = perm0;\n perm0 = perm1;\n perm1 = i;\n }\n\n if (pvt == 2)\n {\n t = AA00;\n AA00 = AA20;\n AA20 = t;\n t = AA01;\n AA01 = AA21;\n AA21 = t;\n t = AA02;\n AA02 = AA22;\n AA22 = t;\n t = AA03;\n AA03 = AA23;\n AA23 = t;\n /* update permutation vector based on row swap */\n i = perm0;\n perm0 = perm2;\n perm2 = i;\n }\n\n if (pvt == 3)\n {\n t = AA00;\n AA00 = AA30;\n AA30 = t;\n t = AA01;\n AA01 = AA31;\n AA31 = t;\n t = AA02;\n AA02 = AA32;\n AA32 = t;\n t = AA03;\n AA03 = AA33;\n AA33 = t;\n /* update permutation vector based on row swap */\n i = perm0;\n perm0 = perm3;\n perm3 = i;\n }\n\n /* scale current row */\n t = rcpOp (AA00);\n icol0 = perm0;\n AA00 = t;\n AA01 = mulOp (t, AA01);\n AA02 = mulOp (t, AA02);\n AA03 = mulOp (t, AA03);\n /* eliminate above and below current row */\n t = AA10;\n AA10 = mulOp (-t, AA00);\n AA11 = fmnaOp (t, AA01, AA11);\n AA12 = fmnaOp (t, AA02, AA12);\n AA13 = fmnaOp (t, AA03, AA13);\n t = AA20;\n AA20 = mulOp (-t, AA00);\n AA21 = fmnaOp (t, AA01, AA21);\n AA22 = fmnaOp (t, AA02, AA22);\n AA23 = fmnaOp (t, AA03, AA23);\n t = AA30;\n AA30 = mulOp (-t, AA00);\n AA31 = fmnaOp (t, AA01, AA31);\n AA32 = fmnaOp (t, AA02, AA32);\n AA33 = fmnaOp (t, AA03, AA33);\n /****************** iteration 1 ***********/\n /* search pivot row */\n p = absOp (AA11);\n pvt = 1;\n t = absOp (AA21);\n\n if (t > p)\n {\n p = t;\n pvt = 2;\n }\n\n t = absOp (AA31);\n\n if (t > p)\n {\n p = t;\n pvt = 3;\n }\n\n /* swap pivot row with row 1 */\n if (pvt == 2)\n {\n t = AA10;\n AA10 = AA20;\n AA20 = t;\n t = AA11;\n AA11 = AA21;\n AA21 = t;\n t = AA12;\n AA12 = AA22;\n AA22 = t;\n t = AA13;\n AA13 = AA23;\n AA23 = t;\n /* update permutation vector based on row swap */\n i = perm1;\n perm1 = perm2;\n perm2 = i;\n }\n else if (pvt == 3)\n {\n t = AA10;\n AA10 = AA30;\n AA30 = t;\n t = AA11;\n AA11 = AA31;\n AA31 = t;\n t = AA12;\n AA12 = AA32;\n AA32 = t;\n t = AA13;\n AA13 = AA33;\n AA33 = t;\n /* update permutation vector based on row swap */\n i = perm1;\n perm1 = perm3;\n perm3 = i;\n }\n\n /* scale current row */\n t = rcpOp (AA11);\n icol1 = perm1;\n AA10 = mulOp (t, AA10);\n AA11 = t;\n AA12 = mulOp (t, AA12);\n AA13 = mulOp (t, AA13);\n /* eliminate above and below current row */\n t = AA01;\n AA00 = fmnaOp (t, AA10, AA00);\n AA01 = mulOp (-t, AA11);\n AA02 = fmnaOp (t, AA12, AA02);\n AA03 = fmnaOp (t, AA13, AA03);\n t = AA21;\n AA20 = fmnaOp (t, AA10, AA20);\n AA21 = mulOp (-t, AA11);\n AA22 = fmnaOp (t, AA12, AA22);\n AA23 = fmnaOp (t, AA13, AA23);\n t = AA31;\n AA30 = fmnaOp (t, AA10, AA30);\n AA31 = mulOp (-t, AA11);\n AA32 = fmnaOp (t, AA12, AA32);\n AA33 = fmnaOp (t, AA13, AA33);\n /****************** iteration 2 ****************/\n /* search pivot row */\n p = absOp (AA22);\n pvt = 2;\n t = absOp (AA32);\n\n if (t > p)\n {\n p = t;\n pvt = 3;\n }\n\n /* swap pivot row with row 2 */\n if (pvt == 3)\n {\n t = AA20;\n AA20 = AA30;\n AA30 = t;\n t = AA21;\n AA21 = AA31;\n AA31 = t;\n t = AA22;\n AA22 = AA32;\n AA32 = t;\n t = AA23;\n AA23 = AA33;\n AA33 = t;\n /* update permutation vector based on row swap */\n i = perm2;\n perm2 = perm3;\n perm3 = i;\n }\n\n /* scale current row */\n t = rcpOp (AA22);\n icol2 = perm2;\n AA20 = mulOp (t, AA20);\n AA21 = mulOp (t, AA21);\n AA22 = t;\n AA23 = mulOp (t, AA23);\n /* eliminate above and below current row */\n t = AA02;\n AA00 = fmnaOp (t, AA20, AA00);\n AA01 = fmnaOp (t, AA21, AA01);\n AA02 = mulOp (-t, AA22);\n AA03 = fmnaOp (t, AA23, AA03);\n t = AA12;\n AA10 = fmnaOp (t, AA20, AA10);\n AA11 = fmnaOp (t, AA21, AA11);\n AA12 = mulOp (-t, AA22);\n AA13 = fmnaOp (t, AA23, AA13);\n t = AA32;\n AA30 = fmnaOp (t, AA20, AA30);\n AA31 = fmnaOp (t, AA21, AA31);\n AA32 = mulOp (-t, AA22);\n AA33 = fmnaOp (t, AA23, AA33);\n /****************** iteration 3 ****************/\n /* scale current row */\n t = rcpOp (AA33);\n icol3 = perm3;\n AA30 = mulOp (t, AA30);\n AA31 = mulOp (t, AA31);\n AA32 = mulOp (t, AA32);\n AA33 = t;\n /* eliminate above and below current row */\n t = AA03;\n AA00 = fmnaOp (t, AA30, AA00);\n AA01 = fmnaOp (t, AA31, AA01);\n AA02 = fmnaOp (t, AA32, AA02);\n AA03 = mulOp (-t, AA33);\n t = AA13;\n AA10 = fmnaOp (t, AA30, AA10);\n AA11 = fmnaOp (t, AA31, AA11);\n AA12 = fmnaOp (t, AA32, AA12);\n AA13 = mulOp (-t, AA33);\n t = AA23;\n AA20 = fmnaOp (t, AA30, AA20);\n AA21 = fmnaOp (t, AA31, AA21);\n AA22 = fmnaOp (t, AA32, AA22);\n AA23 = mulOp (-t, AA33);\n /* sort columns into the correct order */\n Ainv(0, icol0) = AA00;\n Ainv(1, icol0) = AA10;\n Ainv(2, icol0) = AA20;\n Ainv(3, icol0) = AA30;\n Ainv(0, icol1) = AA01;\n Ainv(1, icol1) = AA11;\n Ainv(2, icol1) = AA21;\n Ainv(3, icol1) = AA31;\n Ainv(0, icol2) = AA02;\n Ainv(1, icol2) = AA12;\n Ainv(2, icol2) = AA22;\n Ainv(3, icol2) = AA32;\n Ainv(0, icol3) = AA03;\n Ainv(1, icol3) = AA13;\n Ainv(2, icol3) = AA23;\n Ainv(3, icol3) = AA33;\n }\n}\n\n\ntemplate\n__global__ void matinv_matrix_per_thread_no_pivot (const T *A, T *Ainv, int batch)\n{\n#define A(row,col) A[(col)*N+(row)]\n#define AA(row,col) AA[(col)*N+(row)]\n#define Ainv(row,col) Ainv[(col)*N+(row)]\n const int blkNum = blockIdx.x * blockDim.x + threadIdx.x;\n T AA00, AA01, AA02, AA03, AA10, AA11, AA12, AA13;\n T AA20, AA21, AA22, AA23, AA30, AA31, AA32, AA33;\n T t;\n A += blkNum * N * N;\n Ainv += blkNum * N * N;\n\n if (blkNum < batch)\n {\n AA00 = A[0];\n AA10 = A[1];\n AA20 = A[2];\n AA30 = A[3];\n AA01 = A[4];\n AA11 = A[5];\n AA21 = A[6];\n AA31 = A[7];\n AA02 = A[8];\n AA12 = A[9];\n AA22 = A[10];\n AA32 = A[11];\n AA03 = A[12];\n AA13 = A[13];\n AA23 = A[14];\n AA33 = A[15];\n /****************** iteration 0 ***********/\n /* search pivot row */\n t = 1.0 / (AA00);\n AA00 = t;\n AA01 = t * AA01;\n AA02 = t * AA02;\n AA03 = t * AA03;\n /* eliminate above and below current row */\n t = AA10;\n AA10 = -t * AA00;\n AA11 = fmnaOp (t, AA01, AA11);\n AA12 = fmnaOp (t, AA02, AA12);\n AA13 = fmnaOp (t, AA03, AA13);\n t = AA20;\n AA20 = -t * AA00;\n AA21 = fmnaOp (t, AA01, AA21);\n AA22 = fmnaOp (t, AA02, AA22);\n AA23 = fmnaOp (t, AA03, AA23);\n t = AA30;\n AA30 = -t * AA00;\n AA31 = fmnaOp (t, AA01, AA31);\n AA32 = fmnaOp (t, AA02, AA32);\n AA33 = fmnaOp (t, AA03, AA33);\n /****************** iteration 1 ***********/\n /* scale current row */\n t = 1.0 / (AA11);\n AA10 = t * AA10;\n AA11 = t;\n AA12 = t * AA12;\n AA13 = t * AA13;\n /* eliminate above and below current row */\n t = AA01;\n AA00 = fmnaOp (t, AA10, AA00);\n AA01 = -t * AA11;\n AA02 = fmnaOp (t, AA12, AA02);\n AA03 = fmnaOp (t, AA13, AA03);\n t = AA21;\n AA20 = fmnaOp (t, AA10, AA20);\n AA21 = -t * AA11;\n AA22 = fmnaOp (t, AA12, AA22);\n AA23 = fmnaOp (t, AA13, AA23);\n t = AA31;\n AA30 = fmnaOp (t, AA10, AA30);\n AA31 = -t * AA11;\n AA32 = fmnaOp (t, AA12, AA32);\n AA33 = fmnaOp (t, AA13, AA33);\n /****************** iteration 2 ****************/\n /* scale current row */\n t = 1.0 / (AA22);\n AA20 = t * AA20;\n AA21 = t * AA21;\n AA22 = t;\n AA23 = t * AA23;\n /* eliminate above and below current row */\n t = AA02;\n AA00 = fmnaOp (t, AA20, AA00);\n AA01 = fmnaOp (t, AA21, AA01);\n AA02 = -t * AA22;\n AA03 = fmnaOp (t, AA23, AA03);\n t = AA12;\n AA10 = fmnaOp (t, AA20, AA10);\n AA11 = fmnaOp (t, AA21, AA11);\n AA12 = -t * AA22;\n AA13 = fmnaOp (t, AA23, AA13);\n t = AA32;\n AA30 = fmnaOp (t, AA20, AA30);\n AA31 = fmnaOp (t, AA21, AA31);\n AA32 = -t * AA22;\n AA33 = fmnaOp (t, AA23, AA33);\n /****************** iteration 3 ****************/\n /* scale current row */\n t = 1.0 / (AA33);\n AA30 = t * AA30;\n AA31 = t * AA31;\n AA32 = t * AA32;\n AA33 = t;\n /* eliminate above and below current row */\n t = AA03;\n AA00 = fmnaOp (t, AA30, AA00);\n AA01 = fmnaOp (t, AA31, AA01);\n AA02 = fmnaOp (t, AA32, AA02);\n AA03 = -t * AA33;\n t = AA13;\n AA10 = fmnaOp (t, AA30, AA10);\n AA11 = fmnaOp (t, AA31, AA11);\n AA12 = fmnaOp (t, AA32, AA12);\n AA13 = -t * AA33;\n t = AA23;\n AA20 = fmnaOp (t, AA30, AA20);\n AA21 = fmnaOp (t, AA31, AA21);\n AA22 = fmnaOp (t, AA32, AA22);\n AA23 = -t * AA33;\n /* sort columns into the correct order */\n Ainv(0, 0) = AA00;\n Ainv(1, 0) = AA10;\n Ainv(2, 0) = AA20;\n Ainv(3, 0) = AA30;\n Ainv(0, 1) = AA01;\n Ainv(1, 1) = AA11;\n Ainv(2, 1) = AA21;\n Ainv(3, 1) = AA31;\n Ainv(0, 2) = AA02;\n Ainv(1, 2) = AA12;\n Ainv(2, 2) = AA22;\n Ainv(3, 2) = AA32;\n Ainv(0, 3) = AA03;\n Ainv(1, 3) = AA13;\n Ainv(2, 3) = AA23;\n Ainv(3, 3) = AA33;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockJacobiSmoothbBigBlockDiaCsrKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows, int bsize, int bsize_sq, ValueTypeA *temp1)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int halfwarp_id = tid >> 4;\n const int block_halfwarp_id = threadIdx.x >> 4;\n const int mat_entry_index = threadIdx.x & (16 - 1);\n const int i_ind = mat_entry_index >> 2;\n const int j_ind = mat_entry_index & 3;\n extern __shared__ volatile char schar[];\n volatile ValueTypeA *s_Amat;\n s_Amat = (ValueTypeA *)&schar[0];\n int tile_num = (bsize - 1) / 4 + 1;\n ValueTypeA *e_out = &temp1[(blockIdx.x * blockDim.x + threadIdx.x) * tile_num * tile_num];\n\n while (halfwarp_id < num_block_rows)\n {\n int offset = halfwarp_id * bsize_sq + i_ind * bsize + j_ind;\n int s_offset = block_halfwarp_id * bsize_sq;\n\n // Store the diagonal\n for (int t1 = 0; t1 < tile_num; t1++)\n for (int t2 = 0; t2 < tile_num; t2++)\n if ((t1 * 4 + i_ind) < bsize && (t2 * 4 + j_ind) < bsize)\n {\n e_out[t1 * tile_num + t2] = values[bsize_sq * dia_indices[halfwarp_id] + (t1 * 4 + i_ind) * bsize + t2 * 4 + j_ind];\n }\n\n // Each thread stores its entry in s_Amat\n for (int t1 = 0; t1 < tile_num; t1++)\n for (int t2 = 0; t2 < tile_num; t2++)\n if ((t1 * 4 + i_ind) < bsize && (t2 * 4 + j_ind) < bsize)\n {\n types::util::volcast( e_out[t1 * tile_num + t2], s_Amat + (s_offset + (t1 * 4 + i_ind) * bsize + t2 * 4 + j_ind) );\n }\n\n compute_block_inverse2\n ( s_Amat, s_offset, offset, i_ind, j_ind, Dinv, tile_num, bsize, bsize_sq );\n halfwarp_id += gridDim.x * halfwarps_per_block;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockJacobiSmoothBbyBBlockDiaCsrKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, ValueTypeA *Dinv, const int num_block_rows)\n\n{\n int warp_id = threadIdx.x / 32;\n int warp_thread_id = threadIdx.x & 31;\n\n // padding row blocks to fit in a single warp\n if ( warp_thread_id >= blockrows_per_warp * bsize_sq ) { return; }\n\n // new thread id with padding\n int tid = warp_id * blockrows_per_warp * bsize_sq + warp_thread_id;\n int cta_blockrow_id = tid / bsize_sq;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int mat_entry_index = tid - cta_blockrow_id * bsize_sq;\n const int i_ind = mat_entry_index / bsize;\n const int j_ind = mat_entry_index - i_ind * bsize;\n volatile __shared__ ValueTypeA s_Amat[bsize_sq * blockrows_per_cta ];\n ValueTypeA e_out;\n\n while (blockrow_id < num_block_rows)\n {\n int offset = blockrow_id * bsize_sq + mat_entry_index;\n // Store the diagonal\n e_out = values[bsize_sq * dia_indices[blockrow_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n types::util::volcast(e_out, s_Amat + tid);\n compute_block_inverse_row_major\n ( s_Amat, cta_blockrow_id * bsize_sq, offset, i_ind, j_ind, Dinv );\n blockrow_id += gridDim.x * blockrows_per_cta;\n }\n}\n\ntemplate\n__global__\nvoid setupBlockJacobiSmooth4by4BlockDiaCsrKernel_V2(const IndexType *dia_indices, const ValueTypeA *A_values, ValueTypeA *Dinv, const int num_block_rows)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int halfwarp_id = tid >> log_bsize_sq;\n const int block_halfwarp_id = threadIdx.x >> log_bsize_sq;\n const int mat_entry_index = threadIdx.x & (bsize_sq - 1);\n const int i_ind = mat_entry_index >> log_bsize;\n const int j_ind = mat_entry_index & (bsize - 1);\n volatile __shared__ ValueTypeA s_Amat[bsize_sq * halfwarps_per_block ];\n int offset;\n ValueTypeA e_out;\n\n while (halfwarp_id < num_block_rows)\n {\n // Store the diagonal\n offset = halfwarp_id * bsize_sq;\n e_out = A_values[bsize_sq * dia_indices[halfwarp_id] + mat_entry_index];\n // Each thread stores its entry in s_Amat\n types::util::volcast(e_out, s_Amat + threadIdx.x);\n compute_block_inverse_row_major\n ( s_Amat, block_halfwarp_id * bsize_sq, offset + mat_entry_index, i_ind, j_ind, Dinv );\n halfwarp_id += gridDim.x * blockDim.x >> log_bsize_sq;\n }\n}\n\n// Kernel to smooth with Jacobi smoother, Dinv assumed to be computed\ntemplate\n__global__\nvoid jacobiSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *dia_indices, const ValueTypeA *nonzero_values, const ValueTypeA *Dinv,\n const ValueTypeB *b, const ValueTypeB *x, double weight, const int num_block_rows, ValueTypeB *xout, const int row_offset)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n int eighthwarp_id = row_offset + (tid >> log_bsize);\n const int block_eighthwarp_id = threadIdx.x >> log_bsize;\n const int vec_entry_index = threadIdx.x & (bsize - 1);\n volatile __shared__ ValueTypeB s_xtemp[ bsize * eighthwarps_per_block ];\n ValueTypeB bmAx, xin;\n ValueTypeB temp[bsize];\n int offset, i, s_offset;\n\n while (eighthwarp_id < num_block_rows)\n {\n i = eighthwarp_id;\n offset = i * bsize + vec_entry_index;\n // 1. COMPUTING b-Ax\n bmAx = b[offset];\n // Contribution from diagonal\n xin = x[offset];\n types::util::volcast(xin, s_xtemp + threadIdx.x);\n // Load dia_values and do matrix multiply\n loadAsVector(nonzero_values + bsize_sq * dia_indices[i] + vec_entry_index * bsize, temp);\n s_offset = block_eighthwarp_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx = bmAx - temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = __cachingLoad(&column_indices[jind]);\n\n if (jcol != i)\n {\n offset = jcol * bsize + vec_entry_index;\n types::util::volcast(x[offset], s_xtemp + threadIdx.x);\n // Load nonzero_values\n s_offset = block_eighthwarp_id * bsize;\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx = bmAx - temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n }\n }\n\n types::util::volcast(bmAx, s_xtemp + threadIdx.x);\n bmAx = types::util::get_zero();\n // 2. Multiply by Dinv\n // Load Dinv and multiply to RHS\n offset = i * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(Dinv + offset, temp);\n s_offset = block_eighthwarp_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx = bmAx + temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n\n xout[i * bsize + vec_entry_index] = xin + bmAx * weight;\n eighthwarp_id += gridDim.x * blockDim.x >> log_bsize;\n }\n}\n\n// Kernel to smooth with jacobi smoother, zero initial guess\ntemplate\n__global__\nvoid jacobiSmooth4by4ZeroBlockDiaCsrKernel_NAIVE_tex_readDinv2(const ValueTypeA *Dinv, const ValueTypeB *b, double weight, const int num_block_rows, ValueTypeB *xout, const int row_offset)\n{\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n const int block_eighthwarp_id = threadIdx.x >> log_bsize;\n const int vec_entry_index = threadIdx.x & (bsize - 1);\n volatile __shared__ ValueTypeB s_xtemp[ bsize * eighthwarps_per_block ];\n ValueTypeB bmAx;\n ValueTypeB temp[bsize];\n int offset, i, s_offset;\n\n for (int eighthwarp_id = row_offset + (tid >> log_bsize); eighthwarp_id < num_block_rows; eighthwarp_id += (gridDim.x * blockDim.x >> log_bsize))\n {\n i = eighthwarp_id;\n offset = i * bsize + vec_entry_index;\n types::util::volcast(b[offset], s_xtemp + threadIdx.x);\n bmAx = types::util::get_zero();\n // Load Dinv and multiply to RHS\n offset = i * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(Dinv + offset, temp);\n s_offset = block_eighthwarp_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx = bmAx + temp[m] * types::util::volcast(s_xtemp[s_offset + m]);\n }\n\n xout[i * bsize + vec_entry_index] = bmAx * weight;\n }\n}\n\n//--------------------------------\n// Methods\n//--------------------------------\n\n// Constructor\ntemplate\nBlockJacobiSolver_Base::BlockJacobiSolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope)\n{\n weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n\n if (weight == 0)\n {\n weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Block Jacobi smoother\\n\");\n }\n}\n\n// Destructor\ntemplate\nBlockJacobiSolver_Base::~BlockJacobiSolver_Base()\n{\n this->Dinv.resize(0);\n}\n\ntemplate\nvoid\nBlockJacobiSolver_Base::printSolverParameters() const\n{\n std::cout << \"relaxation_factor= \" << this->weight << std::endl;\n}\n\n// Solver setup\ntemplate\nvoid\nBlockJacobiSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n Matrix *A_as_matrix = dynamic_cast*>(this->m_A);\n\n if (!A_as_matrix)\n {\n FatalError(\"BlockJacobiSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n computeDinv( *A_as_matrix );\n\n if ( A_as_matrix->getBlockFormat() != ROW_MAJOR )\n {\n FatalError(\"Block Jacobi solver only supports row major format\", AMGX_ERR_CONFIGURATION);\n }\n}\n\n//\ntemplate\nvoid\nBlockJacobiSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nBlockJacobiSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n //bool done = false;\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n\n if (xIsZero) { x.dirtybit = 0; }\n\n ViewType oldView = A_as_matrix->currentView();\n\n A_as_matrix->setViewExterior();\n ViewType flags = (ViewType)(A_as_matrix->getViewInterior() | A_as_matrix->getViewExterior());\n\n if (A_as_matrix->get_block_dimx() == 1 && A_as_matrix->get_block_dimy() == 1)\n {\n if (xIsZero)\n {\n smooth_with_0_initial_guess_1x1(*A_as_matrix, b, x, flags);\n }\n else\n {\n smooth_1x1(*A_as_matrix, b, x, flags);\n }\n }\n else if (A_as_matrix->get_block_dimx() == 4 && A_as_matrix->get_block_dimy() == 4)\n {\n if (xIsZero)\n {\n smooth_with_0_initial_guess_4x4(*A_as_matrix, b, x, flags);\n }\n else\n {\n smooth_4x4(*A_as_matrix, b, x, flags);\n }\n }\n else if (A_as_matrix->get_block_dimx() == A_as_matrix->get_block_dimy())\n {\n if (xIsZero)\n {\n thrust_wrapper::fill(x.begin(), x.end(), types::util::get_zero());\n cudaCheckError();\n }\n\n smooth_BxB(*A_as_matrix, b, x, true, flags);\n }\n else\n {\n FatalError(\"Unsupported block size for BlockJacobi_Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if (A_as_matrix->get_block_dimx() == 4 && A_as_matrix->get_block_dimy() == 4)\n {\n if (!xIsZero) // we write to t_res vector to avoid race condition in the kernel\n {\n x.swap(this->t_res);\n }\n }\n\n x.dirtybit = 1;\n A_as_matrix->setView(oldView);\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nBlockJacobiSolver_Base::solve_finalize( VVector &b, VVector &x )\n{}\n\ntemplate\nvoid BlockJacobiSolver_Base::computeDinv( Matrix &A)\n{\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n ViewType oldView = A.currentView();\n A.setView(A_as_matrix->getViewExterior());\n\n if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n this->computeDinv_1x1(A);\n }\n else if (A.get_block_dimx() == 2 && A.get_block_dimy() == 2)\n {\n this->computeDinv_bxb<2>(A);\n }\n else if (A.get_block_dimx() == 3 && A.get_block_dimy() == 3)\n {\n this->computeDinv_3x3(A);\n }\n else if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n this->computeDinv_4x4(A);\n }\n else if (A.get_block_dimx() == 5 && A.get_block_dimy() == 5)\n {\n this->computeDinv_bxb<5>(A);\n }\n else if (A.get_block_dimx() == A.get_block_dimy() && A.get_block_dimy() > 5)\n {\n this->computeDinv_Big(A, A.get_block_dimx());\n }\n\n A.setView(oldView);\n}\n\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid BlockJacobiSolver >::computeDinv_1x1(const Matrix_d &A)\n{\n Matrix_d *A_as_matrix = (Matrix_d *) this->m_A;\n// supports both diag\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), types::util::get_zero());\n\n if ( A_as_matrix->hasProps(DIAG) )\n {\n const int num_values = A_as_matrix->diagOffset() * A_as_matrix->get_block_size();\n amgx::thrust::copy( A_as_matrix->values.begin() + num_values, A_as_matrix->values.begin() + num_values + A_as_matrix->get_num_rows()*A_as_matrix->get_block_size(), this->Dinv.begin() );\n cudaCheckError();\n }\n else\n {\n find_diag( *A_as_matrix );\n }\n}\n\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid BlockJacobiSolver >::computeDinv_1x1(const Matrix_h &A)\n{\n // Do nothing\n}\n\n\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid BlockJacobiSolver >::computeDinv_4x4(const Matrix_h &A)\n{\n //FatalError(\"Block Jacobi smoother not implemented with this format, exiting\");\n //std::cout << \"Warning, CPU version doesn't store the inverse of the blocks, like the GPU version\" << std::endl;\n}\n\ntemplate \nvoid BlockJacobiSolver >::computeDinv_3x3(const Matrix_h &A)\n{\n FatalError(\"3*3 Block Jacobi smoother not implemented with host format, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n //std::cout << \"Warning, CPU version doesn't store the inverse of the blocks, like the GPU version\" << std::endl;\n}\n\ntemplate \nvoid BlockJacobiSolver >::computeDinv_Big(const Matrix_h &A, const int bsize)\n{\n FatalError(\"Big Block Jacobi smoother not implemented with host format, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n //std::cout << \"Warning, CPU version doesn't store the inverse of the blocks, like the GPU version\" << std::endl;\n}\n\n\n// Finding diag on device, CSR format\ntemplate \nvoid BlockJacobiSolver >::find_diag( const Matrix_h &A )\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n if (A.col_indices[j] == i)\n {\n this->Dinv[i] = A.values[j];\n break;\n }\n\n if (j == A.row_offsets[i + 1] - 1)\n {\n FatalError(\"Could not find a diagonal value\", AMGX_ERR_BAD_PARAMETERS);\n }\n }\n }\n}\n\n\n// Finding diag on device, CSR format\ntemplate \nvoid BlockJacobiSolver >::find_diag( const Matrix_d &A )\n{\n AMGX_CPU_PROFILER( \"JacobiSolver::find_diag \" );\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((float)(A.get_num_rows()) / (float)(THREADS_PER_BLOCK)));\n find_diag_kernel_indexed_dia <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK>>>(\n A.get_num_rows(),\n A.diag.raw(),\n A.values.raw(),\n this->Dinv.raw());\n cudaCheckError();\n}\n\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid BlockJacobiSolver >::computeDinv_4x4(const Matrix_d &A)\n{\n// supports both diag\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), types::util::get_zero());\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_indices_ptr = A.diag.raw();\n const ValueTypeA *A_values = A.values.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n#if 1\n const int threads_per_block = 512;\n const int halfwarps_per_block = threads_per_block / 16;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / halfwarps_per_block + 1);\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockJacobiSmooth4by4BlockDiaCsrKernel_V2), hipFuncCachePreferL1);\n setupBlockJacobiSmooth4by4BlockDiaCsrKernel_V2 <<< num_blocks, threads_per_block>>>\n (A_dia_indices_ptr, A_values, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n#else\n hipFuncSetCacheConfig(reinterpret_cast(matinv_matrix_per_thread_pivot), hipFuncCachePreferL1);\n matinv_matrix_per_thread_pivot <<< (A.num_block_rows + 127) / 128, 128 >>> (A_dia_values_ptr, Dinv_ptr, A.num_block_rows);\n cudaCheckError();\n#endif\n}\n\ntemplate \nvoid BlockJacobiSolver >::computeDinv_3x3(const Matrix_d &A)\n{\n const int bsize = 3;\n// supports both diag\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), types::util::get_zero());\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 256;\n const int blockrows_per_warp = 32 / (bsize * bsize);\n // blockrows per cta = blockrows_per_warp * number_of_warps_per_cta\n const int blockrows_per_cta = (threads_per_block / 32) * blockrows_per_warp ;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockJacobiSmoothBbyBBlockDiaCsrKernel), hipFuncCachePreferL1);\n setupBlockJacobiSmoothBbyBBlockDiaCsrKernel <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n}\n\ntemplate \nvoid BlockJacobiSolver >::computeDinv_Big(const Matrix_d &A, const int bsize)\n{\n//both DIAG supported\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), types::util::get_zero());\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n MVector temp(AMGX_GRID_MAX_SIZE * ((bsize - 1) / 4 + 1) * ((bsize - 1) / 4 + 1));\n ValueTypeA *temp_ptr = temp.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 512;\n const int halfwarps_per_block = threads_per_block / 16;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / halfwarps_per_block + 1);\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockJacobiSmoothbBigBlockDiaCsrKernel), hipFuncCachePreferL1);\n setupBlockJacobiSmoothbBigBlockDiaCsrKernel <<< num_blocks, threads_per_block, sizeof(ValueTypeA)*bsize *bsize *halfwarps_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows(), bsize, bsize * bsize, temp_ptr);\n cudaCheckError();\n}\n\n\ntemplate\ntemplate\nvoid BlockJacobiSolver_Base::computeDinv_bxb(const Matrix &A)\n{\n if (TConfig::memSpace == AMGX_host)\n {\n FatalError(\"BlockJacobiSmooth Not implemented for host\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n // supports both diag\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), types::util::get_zero());\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n // MUST BE MULTIPLE OF 16\n const int threads_per_block = 256;\n const int blockrows_per_warp = 32 / (bsize * bsize);\n const int blockrows_per_cta = (threads_per_block / 32) * blockrows_per_warp;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n hipFuncSetCacheConfig(reinterpret_cast(setupBlockJacobiSmoothBbyBBlockDiaCsrKernel), hipFuncCachePreferL1);\n setupBlockJacobiSmoothBbyBBlockDiaCsrKernel <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, Dinv_ptr, A.get_num_rows());\n cudaCheckError();\n }\n}\n\ntemplate \nvoid BlockJacobiSolver >::smooth_BxB(Matrix_h &A, VVector &b, VVector &x, bool firstStep, ViewType separation_flags)\n{\n FatalError(\"M*M Block Jacobi smoother not implemented with host format, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n\ntemplate \nvoid BlockJacobiSolver >::smooth_4x4(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n VVector newx(x.size());\n int bsize = A.get_block_dimx();\n // Allocate space for block_matrix\n ValueTypeA **E = new ValueTypeA* [bsize];\n\n for ( int i = 0; i < bsize; i++)\n {\n E[i] = new ValueTypeA[bsize];\n }\n\n ValueTypeB *bmAx = new ValueTypeB[bsize];\n ValueTypeB *temp = new ValueTypeB[bsize];\n\n //for each block row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n // Compute b - sum_j A_j x_j (denoted by bmAx) for block_row i\n // Load diagonal\n for (int m = 0; m < bsize; m++)\n {\n for (int n = 0; n < bsize; n++)\n {\n E[m][n] = A.values[A.diag[i * bsize * bsize + bsize * m + n]];\n }\n\n bmAx[m] = types::util::get_zero();\n temp[m] = types::util::get_zero();\n }\n\n // Contribution from diagonal\n for (int m = 0; m < bsize; m++)\n for (int n = 0; n < bsize; n++)\n {\n bmAx[m] = bmAx[m] - E[m][n] * x[i * bsize + n];\n }\n\n // Contribution from each nonzero column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n IndexType jcol = A.col_indices[j];\n\n for (int m = 0; m < bsize; m++)\n for (int n = 0; n < bsize; n++)\n {\n bmAx[m] = bmAx[m] - A.values[j * bsize * bsize + bsize * m + n] * x[jcol * bsize + n];\n }\n }\n\n // Add b\n for (int m = 0; m < bsize; m++)\n {\n bmAx[m] = bmAx[m] + b[i * bsize + m];\n }\n\n gaussianEliminationRowMajor(E, temp, bmAx, bsize);\n\n // Compute new value of x\n for (int m = 0; m < bsize; m++)\n {\n newx[i * bsize + m] = x[i * bsize + m] + temp[m] * this->weight;\n }\n }\n\n x.swap(newx);\n}\n\n\n\ntemplate \nvoid BlockJacobiSolver >::smooth_with_0_initial_guess_4x4(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n IndexType bsize = A.get_block_dimy();\n ValueTypeA **E = new ValueTypeA* [bsize];\n\n for ( int i = 0; i < bsize; i++)\n {\n E[i] = new ValueTypeA[bsize];\n }\n\n ValueTypeB *rhs = new ValueTypeB[bsize];\n ValueTypeB *temp = new ValueTypeB[bsize];\n\n //for each block row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n // Load diagonal\n for (int m = 0; m < bsize; m++)\n {\n for (int n = 0; n < bsize; n++)\n {\n E[m][n] = A.values[A.diag[i * bsize * bsize + bsize * m + n]];\n }\n\n rhs[m] = types::util::get_zero();\n }\n\n //rhs\n for (int m = 0; m < bsize; m++)\n {\n rhs[m] = rhs[m] + b[i * bsize + m];\n }\n\n // Solve for temp\n gaussianEliminationRowMajor(E, temp, rhs, bsize);\n\n for (int m = 0; m < bsize; m++)\n {\n x[i * bsize + m] = temp[m] * this->weight;\n }\n }\n}\n\ntemplate \nvoid BlockJacobiSolver >::smooth_1x1_const(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n VVector newx(x.size());\n\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeB Axi = types::util::get_zero();\n ValueTypeB d = types::util::get_one() * A.values[A.diag[i]];\n ValueTypeB mydiaginv = types::util::get_one() * this->weight / (isNotCloseToZero(d) ? d : epsilon(d) );\n\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n Axi = Axi + A.values[j] * x[A.col_indices[j]];\n }\n\n newx[i] = x[i] + (b[i] - Axi) * mydiaginv ;\n }\n\n x.swap(newx);\n}\n\ntemplate \nvoid BlockJacobiSolver >::smooth_1x1(Matrix_h &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n this->smooth_1x1_const(A, b, x, separation_flags);\n}\n\ntemplate \nvoid BlockJacobiSolver >::smooth_1x1(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n AMGX_CPU_PROFILER( \"JacobiSolver::smooth_1x1 \" );\n\n if (this->t_res.size() != x.size())\n {\n this->t_res.resize(x.size());\n }\n\n if (this->y.size() != b.size())\n {\n this->y.resize(b.size());\n this->y.tag = this->tag * 100 + 3;\n this->y.set_block_dimx(b.get_block_dimx());\n this->y.set_block_dimy(b.get_block_dimy());\n }\n\n int num_rows = A.get_num_rows();\n int offset = 0;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n this->y.dirtybit = 0;\n\n multiply( A, x, this->y, separation_flags );\n\n thrust_wrapper::transform( \n amgx::thrust::make_zip_iterator(\n amgx::thrust::make_tuple(x.begin() + offset, this->Dinv.begin() + offset, b.begin() + offset, this->y.begin() + offset)), \n amgx::thrust::make_zip_iterator(\n amgx::thrust::make_tuple(x.begin() + A.get_num_rows(), this->Dinv.begin() + A.get_num_rows(), b.begin() + A.get_num_rows(), this->y.begin() + A.get_num_rows())),\n x.begin() + offset,\n jacobi_postsmooth_functor( this->weight ));\n cudaCheckError();\n}\n\n\ntemplate \nvoid BlockJacobiSolver >::smooth_with_0_initial_guess_1x1(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeB d = types::util::get_one() * A.values[A.diag[i]];\n ValueTypeB mydiag = types::util::get_one() * this->weight / (isNotCloseToZero(d) ? d : epsilon(d));\n x[i] = b[i] * mydiag;\n }\n}\n\ntemplate \nvoid BlockJacobiSolver >::smooth_with_0_initial_guess_1x1(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n AMGX_CPU_PROFILER( \"JacobiSolver::smooth_with_0_initial_guess_1x1 \" );\n int num_rows = A.get_num_rows();\n int offset = 0;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n thrust_wrapper::transform( b.begin( ) + offset,\n b.begin( ) + A.get_num_rows(),\n this->Dinv.begin( ) + offset,\n x.begin( ) + offset,\n jacobi_presmooth_functor( this->weight ));\n cudaCheckError();\n}\n\ntemplate \nvoid BlockJacobiSolver >::smooth_BxB(Matrix_d &A, VVector &b, VVector &x, bool firstStep, ViewType separation_flags)\n{\n IndexType num_rows;\n IndexType offset;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n\n // aux vector initialization\n if (this->y.size() != b.size())\n {\n this->y.resize(b.size());\n this->y.tag = this->tag * 100 + 3;\n this->y.set_block_dimx(b.get_block_dimx());\n this->y.set_block_dimy(b.get_block_dimy());\n }\n\n amgx::thrust::copy(b.begin(), b.end(), this->y.begin()); // copy of vector b\n cudaCheckError();\n Cusparse::bsrmv(types::util::get_minus_one(), A, x, types::util::get_one(), this->y, separation_flags); // y= -1.0f*(A.x) + y\n cudaCheckError();\n\n Cusparse::bsrmv(types::util::get_one() * this->weight, A, this->Dinv, this->y, types::util::get_one(), x, separation_flags); // t_res = t_res + w*(Dinv.y) @ view\n cudaCheckError();\n}\n\ntemplate \nvoid BlockJacobiSolver >::smooth_4x4(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n if (this->t_res.size() != x.size())\n {\n this->t_res.resize(x.size());\n }\n\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n ValueTypeB *x_ptr = x.raw();\n ValueTypeB *xout_ptr = this->t_res.raw(); // always store original x\n IndexType num_rows = A.get_num_rows();\n IndexType offset = 0;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n const int threads_per_block = 512;\n const int eightwarps_per_block = threads_per_block / 4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows - 1) / eightwarps_per_block + 1);\n\n // XXX We need to add latency hiding here\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo(x, x.tag);\n }\n\n hipFuncSetCacheConfig(reinterpret_cast(jacobiSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2), hipFuncCachePreferL1);\n jacobiSmooth4by4BlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, Dinv_ptr,\n b_ptr, x_ptr, this->weight, offset + num_rows, xout_ptr, offset);\n cudaCheckError();\n}\n\n\ntemplate \nvoid BlockJacobiSolver >::smooth_with_0_initial_guess_4x4(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n cudaCheckError();\n const ValueTypeA *A_values_ptr = A.values.raw();\n const ValueTypeB *b_ptr = b.raw();\n const ValueTypeA *Dinv_ptr = this->Dinv.raw();\n ValueTypeB *x_ptr = x.raw();\n IndexType num_rows = A.get_num_rows();\n IndexType offset = 0;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n const int threads_per_block = 512;\n const int eightwarps_per_block = threads_per_block / 4;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows - 1) / eightwarps_per_block + 1);\n hipFuncSetCacheConfig(reinterpret_cast(jacobiSmooth4by4ZeroBlockDiaCsrKernel_NAIVE_tex_readDinv2), hipFuncCachePreferL1);\n jacobiSmooth4by4ZeroBlockDiaCsrKernel_NAIVE_tex_readDinv2 <<< num_blocks, threads_per_block>>>\n (Dinv_ptr, b_ptr, this->weight, offset + num_rows, x_ptr, offset);\n cudaCheckError();\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class BlockJacobiSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class BlockJacobiSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace block_jacobi\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/solvers/block_jacobi_solver.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\n// -----------\n// Kernels\n// -----------\n\n/*************************************************************************\n* \"random\" hash function for both device and host\n************************************************************************/\n__host__ __device__ static int ourHash(const int i, const int max)\n{\n unsigned int a = i;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return int(((a ^ 0x4a51e590) / (float)UINT_MAX) * max);\n}\n\nstruct prg\n{\n float a, b;\n int max_int;\n\n __host__ __device__\n prg(int _max_int, float _a = 0.f, float _b = 1.f) : a(_a), b(_b), max_int(_max_int) {};\n\n __host__ __device__\n int operator()(const unsigned int n) const\n {\n int ru = ourHash(n, max_int);\n return (ru);\n }\n};\n\ntemplate \nvoid initRandom(Vector &vec, int size, int max_int)\n{\n vec.resize(size);\n amgx::thrust::counting_iterator index_sequence_begin(0);\n amgx::thrust::transform(index_sequence_begin,\n index_sequence_begin + size,\n vec.begin(),\n prg(max_int));\n}\n\ntemplate\n__global__ void compute_anorm_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n ValueTypeA *d)\n{\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n ValueTypeB d_ = 0;\n IndexType row_start = Ap[ridx];\n IndexType row_end = Ap[ridx + 1];\n\n for (int j = row_start; j < row_end; j++)\n {\n ValueTypeB Aij = Ax[j];\n d_ += Aij * Aij;\n }\n\n // Store L2-norm\n d[ridx] = d_;\n }\n}\n\ntemplate\n__global__ void compute_multicolor_anorm_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const ValueTypeA *d,\n const int *sorted_rows_by_color,\n const int num_rows_per_color)\n{\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = tidx; ridx < num_rows_per_color; ridx += blockDim.x * gridDim.x)\n {\n int i = sorted_rows_by_color[ridx];\n ValueTypeB d_ = 0;\n IndexType row_start = Ap[i];\n IndexType row_end = Ap[i + 1];\n\n for (int j = row_start; j < row_end; j++)\n {\n ValueTypeB Aij = Ax[j];\n d_ += Aij * Aij;\n }\n\n // Store L2-norm\n d[ridx] = d_;\n }\n}\n\ntemplate\n__global__ void compute_cumul_inv_kernel(const IndexType a_cum_num_rows,\n ValueTypeA *a_cum,\n ValueTypeB d_inv,\n int c_inv_sz,\n IndexType *c_inv)\n{\n const int tidx = blockDim.x * blockIdx.x + threadIdx.x;\n int idx, idx1, idx2;\n\n for (int ridx = tidx; ridx < a_cum_num_rows; ridx += blockDim.x * gridDim.x)\n {\n // printf(\"%d %f %f\\n\", ridx, d_inv, a_cum[ridx]);\n //printf(\"%f\\n\", a_cum[ridx]);\n double a = a_cum[ridx];\n // if (ridx < 0 || ridx >= a_cum_num_rows)\n // printf(\"!! %d %d\\n\", ridx, idx);\n idx1 = int(a / d_inv) - 1; // get index in inverse table (floor - 1)\n\n if (ridx < a_cum_num_rows - 1)\n {\n idx2 = a_cum[ridx + 1] / d_inv - 1; // get index in inverse table (floor - 1)\n }\n else\n {\n idx2 = c_inv_sz;\n }\n\n // printf(\"%d %d\\n\", idx1, idx2);\n for ( idx = idx1; idx < idx2; idx++)\n {\n if (idx >= c_inv_sz || idx < 0)\n {\n printf(\"Ai! %d %d\\n\", idx, ridx);\n }\n\n c_inv[idx] = ridx;\n }\n }\n}\n\ntemplate\n__global__ void compute_amax_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n IndexType *amax_idx)\n{\n ValueTypeA maxVal(0), avalue;\n IndexType jmax;\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n IndexType row_start = Ap[ridx];\n IndexType row_end = Ap[ridx + 1];\n jmax = row_start;\n\n for (int j = row_start; j < row_end; j++)\n {\n avalue = Ax[j];\n\n if (avalue > maxVal)\n {\n maxVal = avalue;\n jmax = j;\n }\n }\n\n // Store position of maxvalue\n amax_idx[ridx] = jmax;\n }\n}\n\ntemplate\n__global__ void kaczmarz_smooth_kernel_naive_atomics(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const ValueTypeA *d,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB *xout,\n const IndexType row_offset)\n\n{\n // Naive implementation, needs x copy in xout at the very beginning\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = row_offset + tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n IndexType row_start = Ap[ridx];\n IndexType row_end = Ap[ridx + 1];\n ValueTypeB Axi = 0.0;\n ValueTypeB r;\n\n for (int j = row_start; j < row_end; j++)\n {\n Axi += Ax[j] * xout[Aj[j]];\n }\n\n r = (b[ridx] - Axi) / ( isNotCloseToZero( d[ridx]) ? d[ridx] : epsilon(d[ridx]) );\n\n for (int j = row_start; j < row_end; j++)\n {\n //xout[Aj[j]] += r*Ax[j];\n utils::atomic_add(&xout[Aj[j]], r * Ax[j]);\n }\n }\n}\n\ntemplate\n__global__ void kaczmarz_smooth_kernel_warp_atomics(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const ValueTypeA *d,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB *xout,\n const IndexType row_offset)\n\n{\n const int num_warps = kCtaSize / 32;\n const int num_rows_per_iter = num_warps * gridDim.x;\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n for ( int ridx = blockIdx.x * num_warps + warpId ; ridx < num_rows ;\n ridx += num_rows_per_iter )\n {\n IndexType row_start = Ap[ridx];\n IndexType row_end = Ap[ridx + 1];\n ValueTypeB Axi = 0.0;\n ValueTypeB r;\n\n for (int j = row_start + laneId; utils::any( j < row_end) ; j += 32)\n {\n ValueTypeB aValue = j < row_end ? Ax[j] : ValueTypeB(0);\n ValueTypeB xValue = j < row_end ? xout[Aj[j]] : ValueTypeB(0);\n r = utils::warp_reduce<1, utils::Add>(aValue * xValue);\n Axi += r;\n }\n\n r = (b[ridx] - Axi) / ( isNotCloseToZero( d[ridx]) ? d[ridx] : epsilon(d[ridx]) );\n\n for (int j = row_start + laneId; utils::any( j < row_end) ; j += 32)\n {\n //ValueTypeB dx = j < row_end ? r*Ax[j] : ValueTypeB(0);\n //int aj = j < row_end ? r*Ax[j] : ValueTypeB(0);\n if (j < row_end)\n {\n utils::atomic_add(&xout[Aj[j]], r * Ax[j]);\n }\n }\n }\n}\n\ntemplate\n__global__ void randomized_kaczmarz_smooth_kernel_warp_atomics(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const IndexType *c_inv,\n const IndexType *rnd_rows,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB *xout,\n const IndexType row_offset)\n\n{\n const int num_warps = kCtaSize / 32;\n const int num_rows_per_iter = num_warps * gridDim.x;\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n for ( int ridx = blockIdx.x * num_warps + warpId ; ridx < num_rows ;\n ridx += num_rows_per_iter )\n {\n int irow = c_inv[rnd_rows[ridx]];\n IndexType row_start = Ap[irow];\n IndexType row_end = Ap[irow + 1];\n ValueTypeB Axi = 0.0;\n ValueTypeB r;\n ValueTypeA aa;\n ValueTypeA AA = 0.0;\n\n for (int j = row_start + laneId; utils::any( j < row_end) ; j += 32)\n {\n ValueTypeB aValue = j < row_end ? Ax[j] : ValueTypeB(0);\n ValueTypeB xValue = j < row_end ? xout[Aj[j]] : ValueTypeB(0);\n r = utils::warp_reduce<1, utils::Add>(aValue * xValue);\n aa = utils::warp_reduce<1, utils::Add>(aValue * aValue);\n Axi += r;\n AA += aa;\n }\n\n r = (b[ridx] - Axi) / ( isNotCloseToZero( AA) ? AA : epsilon(AA) );\n\n for (int j = row_start + laneId; utils::any( j < row_end) ; j += 32)\n {\n //ValueTypeB dx = j < row_end ? r*Ax[j] : ValueTypeB(0);\n //int aj = j < row_end ? r*Ax[j] : ValueTypeB(0);\n if (j < row_end)\n {\n utils::atomic_add(&xout[Aj[j]], r * Ax[j]);\n }\n }\n }\n}\n\ntemplate\n__global__ void kaczmarz_smooth_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const IndexType *amax,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB *xout,\n const IndexType row_offset)\n\n{\n // Naive implementation, needs x copy in xout at the very beginning\n //IndexType tidx = blockDim.x*blockIdx.x + threadIdx.x;\n IndexType i, t;\n const int num_warps = kCtaSize / 32;\n const int num_rows_per_iter = num_warps * gridDim.x;\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n for ( int ridx = blockIdx.x * num_warps + warpId ; ridx < num_rows ;\n ridx += num_rows_per_iter )\n {\n ValueTypeB Axi = 0.0;\n ValueTypeB r;\n i = ourHash(ridx, num_rows);\n IndexType row_start = Ap[i];\n IndexType row_end = Ap[i + 1];\n\n for (int j = row_start + laneId; utils::any( j < row_end) ; j += 32)\n {\n ValueTypeB aValue = j < row_end ? Ax[j] : ValueTypeB(0);\n ValueTypeB xValue = j < row_end ? xout[Aj[j]] : ValueTypeB(0);\n r = utils::warp_reduce<1, utils::Add>(aValue * xValue);\n Axi += r;\n //Axi += utils::Warp_reduce_linear<1,32>::execute(aValue * xValue);\n //Axi += Ax[j] * xout[Aj[j]];\n printf(\"j = %d, r = %f\\n\", j, r);\n }\n\n if (laneId == 0)\n {\n r = (b[i] - Axi);// / ( isNotCloseToZero( d[ridx]) ? d[ridx] : epsilon(d[ridx]) );\n t = row_start + ourHash(ridx, row_end - row_start);\n printf(\"ridx=%d, i=%d, t=%d, Aj[t]=%d, r=%f\\n\", ridx, i, t, Aj[t], r);\n xout[Aj[t]] += r * ((row_end - row_start) * Ax[t]) * 0.5;\n }\n }\n}\n\ntemplate\n__global__ void multicolor_kaczmarz_smooth_kernel_naive(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const ValueTypeA *d,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB weight,\n const int *sorted_rows_by_color,\n const int num_rows_per_color,\n ValueTypeB *xout)\n\n{\n int i;\n // Naive implementation, needs x copy in xout at the very beginning\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = tidx; ridx < num_rows_per_color; ridx += blockDim.x * gridDim.x)\n {\n i = sorted_rows_by_color[ridx];\n IndexType row_start = Ap[i];\n IndexType row_end = Ap[i + 1];\n ValueTypeB Axi = 0.0;\n ValueTypeB r;\n\n for (int j = row_start; j < row_end; j++)\n {\n Axi += Ax[j] * xout[Aj[j]];\n }\n\n r = (b[i] - Axi) / ( isNotCloseToZero( d[i]) ? d[i] : epsilon(d[i]) );\n\n for (int j = row_start; j < row_end; j++)\n {\n utils::atomic_add(&xout[Aj[j]], r * Ax[j]);\n }\n }\n}\n\ntemplate\n__global__ void multicolor_kaczmarz_smooth_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const ValueTypeA *d,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB weight,\n const int *sorted_rows_by_color,\n const int num_rows_per_color,\n ValueTypeB *xout)\n\n{\n const int num_warps = kCtaSize / 32;\n const int num_rows_per_iter = num_warps * gridDim.x;\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n int i;\n\n for ( int ridx = blockIdx.x * num_warps + warpId ; ridx < num_rows_per_color ;\n ridx += num_rows_per_iter )\n {\n i = sorted_rows_by_color[ridx];\n IndexType row_start = Ap[i];\n IndexType row_end = Ap[i + 1];\n ValueTypeB Axi = 0.0;\n ValueTypeB r;\n\n for (int j = row_start + laneId; utils::any( j < row_end) ; j += 32)\n {\n ValueTypeB aValue = j < row_end ? Ax[j] : ValueTypeB(0);\n ValueTypeB xValue = j < row_end ? xout[Aj[j]] : ValueTypeB(0);\n r = utils::warp_reduce<1, utils::Add>(aValue * xValue);\n Axi += r;\n }\n\n r = (b[i] - Axi) / ( isNotCloseToZero( d[i]) ? d[i] : epsilon(d[i]) );\n\n for (int j = row_start + laneId; utils::any( j < row_end) ; j += 32)\n {\n if (j < row_end)\n {\n utils::atomic_add(&xout[Aj[j]], r * Ax[j]);\n }\n }\n }\n}\n\ntemplate\n__global__ void jacobi_smooth_with_0_initial_guess_kernel(const IndexType num_rows,\n const ValueTypeA *d,\n const ValueTypeB *b,\n ValueTypeB *x,\n ValueTypeB weight,\n const IndexType row_offset)\n\n{\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = row_offset + tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n x[ridx] = weight * b[ridx] / ( isNotCloseToZero( d[ridx]) ? d[ridx] : epsilon(d[ridx]) );\n }\n}\n\n// -----------------\n// Methods\n// -----------------\n\n// Constructor\ntemplate\nKaczmarzSolver_Base::KaczmarzSolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope), m_an(0), m_amax(0), m_c_inv(0)\n{\n weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n this->m_coloring_needed = (cfg.AMG_Config::template getParameter(\"kaczmarz_coloring_needed\", cfg_scope) != 0);\n this->m_reorder_cols_by_color_desired = (cfg.AMG_Config::template getParameter(\"reorder_cols_by_color\", cfg_scope) != 0);\n this->m_randomized = true;\n\n if (weight == 0)\n {\n weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Block Jacobi smoother\\n\");;\n }\n}\n\n// Destructor\ntemplate\nKaczmarzSolver_Base::~KaczmarzSolver_Base()\n{\n}\n\n// Solver setup\ntemplate\nvoid\nKaczmarzSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n m_explicit_A = dynamic_cast*>(Base::m_A);\n\n if (!m_explicit_A)\n {\n FatalError(\"Kaczmarz solver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n compute_anorm( *this->m_explicit_A );\n\n if (m_randomized) // MC RK is not supported here\n {\n if (m_coloring_needed)\n {\n //FatalError(\"Randomized Kaczmarz solver does not support coloring\", AMGX_ERR_INTERNAL);\n m_coloring_needed = false;\n }\n\n double d_inv = this->m_an[0];\n int c_sz = this->m_an.size();\n d_inv = amgx::thrust::reduce(this->m_an.begin(), this->m_an.end(), d_inv, amgx::thrust::minimum());\n amgx::thrust::inclusive_scan(this->m_an.begin(), this->m_an.end(), this->m_an.begin()); // in-place scan\n int c_inv_sz = (this->m_an[c_sz - 1] + d_inv - 1 ) / d_inv;\n this->m_c_inv.resize(c_inv_sz, -1);\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((ValueTypeB)c_sz / (ValueTypeB)THREADS_PER_BLOCK));\n\n if (c_sz > 0)\n {\n device_vector_alloc aa(c_sz, 1);\n compute_cumul_inv_kernel <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK>>>\n (c_sz,\n this->m_an.raw(),\n d_inv,\n c_inv_sz,\n this->m_c_inv.raw());\n cudaCheckError();\n }\n\n cudaDeviceSynchronize();\n cudaCheckError();\n }\n}\n\ntemplate\nvoid KaczmarzSolver_Base::compute_anorm( Matrix &A)\n{\n this->m_an.resize(A.get_num_rows()*A.get_block_dimx());\n ViewType oldView = A.currentView();\n A.setView(this->m_explicit_A->getViewExterior());\n\n if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n compute_anorm_1x1(A);\n }\n else\n {\n FatalError(\"Unsupported block size for KaczmarzSolver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n}\n\ntemplate\nvoid KaczmarzSolver_Base::compute_amax( Matrix &A)\n{\n this->m_amax.resize(A.get_num_rows()*A.get_block_dimx());\n ViewType oldView = A.currentView();\n A.setView(this->m_explicit_A->getViewExterior());\n\n if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n compute_amax_1x1(A);\n }\n else\n {\n FatalError(\"Unsupported block size for KaczmarzSolver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n}\n\ntemplate \nvoid KaczmarzSolver >::compute_anorm_1x1(const Matrix_d &A)\n{\n//DIAG: starnge issues trying to add DIAG property handling\n// now leaving !DIAG only\n if (A.hasProps(DIAG))\n {\n FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((ValueTypeB)A.get_num_rows() / (ValueTypeB)THREADS_PER_BLOCK));\n\n if (A.get_num_rows() > 0)\n {\n compute_anorm_kernel <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK>>>\n ((int)A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->m_an.raw());\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid KaczmarzSolver >::compute_amax_1x1(const Matrix_d &A)\n{\n//DIAG: starnge issues trying to add DIAG property handling\n// now leaving !DIAG only\n if (A.hasProps(DIAG))\n {\n FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((ValueTypeB)A.get_num_rows() / (ValueTypeB)THREADS_PER_BLOCK));\n\n if (A.get_num_rows() > 0)\n {\n compute_amax_kernel <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK>>>\n ((int)A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->m_amax.raw());\n }\n\n cudaCheckError();\n}\n\ntemplate\nvoid\nKaczmarzSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nKaczmarzSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n if (xIsZero) { x.dirtybit = 0; }\n\n if (!this->m_explicit_A->is_matrix_singleGPU())\n {\n this->m_explicit_A->manager->exchange_halo_async(x, x.tag);\n\n if (this->m_explicit_A->getViewExterior() == this->m_explicit_A->getViewInterior())\n {\n this->m_explicit_A->manager->exchange_halo_wait(x, x.tag);\n }\n }\n\n ViewType oldView = this->m_explicit_A->currentView();\n ViewType flags;\n bool latencyHiding = true;\n\n if (this->m_explicit_A->is_matrix_singleGPU() || (x.dirtybit == 0))\n {\n latencyHiding = false;\n this->m_explicit_A->setViewExterior();\n flags = this->m_explicit_A->getViewExterior();\n }\n else\n {\n flags = this->m_explicit_A->getViewInterior();\n this->m_explicit_A->setViewInterior();\n }\n\n if (this->m_explicit_A->get_block_dimx() == 1 && this->m_explicit_A->get_block_dimy() == 1)\n {\n if (xIsZero)\n {\n //smooth_with_0_initial_guess_1x1(*this->m_explicit_A, b, x, flags);\n smooth_1x1(*this->m_explicit_A, b, x, flags, latencyHiding);\n }\n else\n {\n smooth_1x1(*this->m_explicit_A, b, x, flags, latencyHiding);\n }\n }\n else\n {\n FatalError(\"Unsupported block size for Kaczmarz_Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n x.dirtybit = 1;\n this->m_explicit_A->setView(oldView);\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nKaczmarzSolver_Base::solve_finalize( VVector &b, VVector &x )\n{\n}\n\n// Multicolor version\ntemplate \nvoid KaczmarzSolver >::smooth_1x1_MC(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags, bool latency_hiding)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n ValueTypeB *x_ptr = x.raw();\n IndexType num_rows = A.get_num_rows();\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = ((separation_flags & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flags == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int threads_per_block = 128;\n const int blockrows_per_warp = 1;\n const int blockrows_per_cta = (threads_per_block / 32) * blockrows_per_warp;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows_per_color / blockrows_per_cta + 1));\n multicolor_kaczmarz_smooth_kernel <<< num_blocks, threads_per_block >>>\n (A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->m_an.raw(),\n b.raw(),\n x_ptr,\n this->weight,\n A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color,\n x_ptr);\n cudaCheckError();\n }\n}\n\ntemplate \nvoid KaczmarzSolver >::smooth_1x1_naive(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags, bool latency_hiding)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n ValueTypeB *x_ptr = x.raw();\n IndexType num_rows = A.get_num_rows();\n IndexType offset = 0;\n // Skipping Multi-GPU logic for now\n // Current. Will be exact only with one warp per grid\n const int threads_per_block = 32;\n const int num_blocks = 1;\n\n if (this->m_randomized)\n {\n IVector rnd_rows;\n int c_inv_sz = this->m_c_inv.size();\n initRandom(rnd_rows, A.get_num_rows(), c_inv_sz);\n randomized_kaczmarz_smooth_kernel_warp_atomics <<< num_blocks, threads_per_block >>>\n (A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->m_c_inv.raw(),\n rnd_rows.raw(),\n b.raw(),\n x_ptr,\n x_ptr,\n offset);\n cudaCheckError();\n }\n else\n {\n kaczmarz_smooth_kernel_warp_atomics <<< num_blocks, threads_per_block >>>\n (A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->m_an.raw(),\n b.raw(),\n x_ptr,\n x_ptr,\n offset);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid KaczmarzSolver >::smooth_1x1(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags, bool latency_hiding)\n{\n if (this->m_coloring_needed)\n {\n smooth_1x1_MC(A, b, x, separation_flags, latency_hiding);\n }\n else\n {\n smooth_1x1_naive(A, b, x, separation_flags, latency_hiding);\n }\n}\n\ntemplate \nvoid KaczmarzSolver >::smooth_with_0_initial_guess_1x1(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n ViewType oldView = A.currentView();\n // Process all rows\n A.setViewExterior();\n ViewType flags = A.getViewExterior();\n int offset, num_rows;\n A.getOffsetAndSizeForView(flags, &offset, &num_rows);\n A.setView(oldView);\n cudaCheckError();\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class KaczmarzSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class KaczmarzSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\n// -----------\n// Kernels\n// -----------\n\n/*************************************************************************\n* \"random\" hash function for both device and host\n************************************************************************/\n__host__ __device__ static int ourHash(const int i, const int max)\n{\n unsigned int a = i;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return int(((a ^ 0x4a51e590) / (float)UINT_MAX) * max);\n}\n\nstruct prg\n{\n float a, b;\n int max_int;\n\n __host__ __device__\n prg(int _max_int, float _a = 0.f, float _b = 1.f) : a(_a), b(_b), max_int(_max_int) {};\n\n __host__ __device__\n int operator()(const unsigned int n) const\n {\n int ru = ourHash(n, max_int);\n return (ru);\n }\n};\n\ntemplate \nvoid initRandom(Vector &vec, int size, int max_int)\n{\n vec.resize(size);\n amgx::thrust::counting_iterator index_sequence_begin(0);\n amgx::thrust::transform(index_sequence_begin,\n index_sequence_begin + size,\n vec.begin(),\n prg(max_int));\n}\n\ntemplate\n__global__ void compute_anorm_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n ValueTypeA *d)\n{\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n ValueTypeB d_ = 0;\n IndexType row_start = Ap[ridx];\n IndexType row_end = Ap[ridx + 1];\n\n for (int j = row_start; j < row_end; j++)\n {\n ValueTypeB Aij = Ax[j];\n d_ += Aij * Aij;\n }\n\n // Store L2-norm\n d[ridx] = d_;\n }\n}\n\ntemplate\n__global__ void compute_multicolor_anorm_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const ValueTypeA *d,\n const int *sorted_rows_by_color,\n const int num_rows_per_color)\n{\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = tidx; ridx < num_rows_per_color; ridx += blockDim.x * gridDim.x)\n {\n int i = sorted_rows_by_color[ridx];\n ValueTypeB d_ = 0;\n IndexType row_start = Ap[i];\n IndexType row_end = Ap[i + 1];\n\n for (int j = row_start; j < row_end; j++)\n {\n ValueTypeB Aij = Ax[j];\n d_ += Aij * Aij;\n }\n\n // Store L2-norm\n d[ridx] = d_;\n }\n}\n\ntemplate\n__global__ void compute_cumul_inv_kernel(const IndexType a_cum_num_rows,\n ValueTypeA *a_cum,\n ValueTypeB d_inv,\n int c_inv_sz,\n IndexType *c_inv)\n{\n const int tidx = blockDim.x * blockIdx.x + threadIdx.x;\n int idx, idx1, idx2;\n\n for (int ridx = tidx; ridx < a_cum_num_rows; ridx += blockDim.x * gridDim.x)\n {\n // printf(\"%d %f %f\\n\", ridx, d_inv, a_cum[ridx]);\n //printf(\"%f\\n\", a_cum[ridx]);\n double a = a_cum[ridx];\n // if (ridx < 0 || ridx >= a_cum_num_rows)\n // printf(\"!! %d %d\\n\", ridx, idx);\n idx1 = int(a / d_inv) - 1; // get index in inverse table (floor - 1)\n\n if (ridx < a_cum_num_rows - 1)\n {\n idx2 = a_cum[ridx + 1] / d_inv - 1; // get index in inverse table (floor - 1)\n }\n else\n {\n idx2 = c_inv_sz;\n }\n\n // printf(\"%d %d\\n\", idx1, idx2);\n for ( idx = idx1; idx < idx2; idx++)\n {\n if (idx >= c_inv_sz || idx < 0)\n {\n printf(\"Ai! %d %d\\n\", idx, ridx);\n }\n\n c_inv[idx] = ridx;\n }\n }\n}\n\ntemplate\n__global__ void compute_amax_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n IndexType *amax_idx)\n{\n ValueTypeA maxVal(0), avalue;\n IndexType jmax;\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n IndexType row_start = Ap[ridx];\n IndexType row_end = Ap[ridx + 1];\n jmax = row_start;\n\n for (int j = row_start; j < row_end; j++)\n {\n avalue = Ax[j];\n\n if (avalue > maxVal)\n {\n maxVal = avalue;\n jmax = j;\n }\n }\n\n // Store position of maxvalue\n amax_idx[ridx] = jmax;\n }\n}\n\ntemplate\n__global__ void kaczmarz_smooth_kernel_naive_atomics(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const ValueTypeA *d,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB *xout,\n const IndexType row_offset)\n\n{\n // Naive implementation, needs x copy in xout at the very beginning\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = row_offset + tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n IndexType row_start = Ap[ridx];\n IndexType row_end = Ap[ridx + 1];\n ValueTypeB Axi = 0.0;\n ValueTypeB r;\n\n for (int j = row_start; j < row_end; j++)\n {\n Axi += Ax[j] * xout[Aj[j]];\n }\n\n r = (b[ridx] - Axi) / ( isNotCloseToZero( d[ridx]) ? d[ridx] : epsilon(d[ridx]) );\n\n for (int j = row_start; j < row_end; j++)\n {\n //xout[Aj[j]] += r*Ax[j];\n utils::atomic_add(&xout[Aj[j]], r * Ax[j]);\n }\n }\n}\n\ntemplate\n__global__ void kaczmarz_smooth_kernel_warp_atomics(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const ValueTypeA *d,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB *xout,\n const IndexType row_offset)\n\n{\n const int num_warps = kCtaSize / 32;\n const int num_rows_per_iter = num_warps * gridDim.x;\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n for ( int ridx = blockIdx.x * num_warps + warpId ; ridx < num_rows ;\n ridx += num_rows_per_iter )\n {\n IndexType row_start = Ap[ridx];\n IndexType row_end = Ap[ridx + 1];\n ValueTypeB Axi = 0.0;\n ValueTypeB r;\n\n for (int j = row_start + laneId; utils::any( j < row_end) ; j += 32)\n {\n ValueTypeB aValue = j < row_end ? Ax[j] : ValueTypeB(0);\n ValueTypeB xValue = j < row_end ? xout[Aj[j]] : ValueTypeB(0);\n r = utils::warp_reduce<1, utils::Add>(aValue * xValue);\n Axi += r;\n }\n\n r = (b[ridx] - Axi) / ( isNotCloseToZero( d[ridx]) ? d[ridx] : epsilon(d[ridx]) );\n\n for (int j = row_start + laneId; utils::any( j < row_end) ; j += 32)\n {\n //ValueTypeB dx = j < row_end ? r*Ax[j] : ValueTypeB(0);\n //int aj = j < row_end ? r*Ax[j] : ValueTypeB(0);\n if (j < row_end)\n {\n utils::atomic_add(&xout[Aj[j]], r * Ax[j]);\n }\n }\n }\n}\n\ntemplate\n__global__ void randomized_kaczmarz_smooth_kernel_warp_atomics(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const IndexType *c_inv,\n const IndexType *rnd_rows,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB *xout,\n const IndexType row_offset)\n\n{\n const int num_warps = kCtaSize / 32;\n const int num_rows_per_iter = num_warps * gridDim.x;\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n for ( int ridx = blockIdx.x * num_warps + warpId ; ridx < num_rows ;\n ridx += num_rows_per_iter )\n {\n int irow = c_inv[rnd_rows[ridx]];\n IndexType row_start = Ap[irow];\n IndexType row_end = Ap[irow + 1];\n ValueTypeB Axi = 0.0;\n ValueTypeB r;\n ValueTypeA aa;\n ValueTypeA AA = 0.0;\n\n for (int j = row_start + laneId; utils::any( j < row_end) ; j += 32)\n {\n ValueTypeB aValue = j < row_end ? Ax[j] : ValueTypeB(0);\n ValueTypeB xValue = j < row_end ? xout[Aj[j]] : ValueTypeB(0);\n r = utils::warp_reduce<1, utils::Add>(aValue * xValue);\n aa = utils::warp_reduce<1, utils::Add>(aValue * aValue);\n Axi += r;\n AA += aa;\n }\n\n r = (b[ridx] - Axi) / ( isNotCloseToZero( AA) ? AA : epsilon(AA) );\n\n for (int j = row_start + laneId; utils::any( j < row_end) ; j += 32)\n {\n //ValueTypeB dx = j < row_end ? r*Ax[j] : ValueTypeB(0);\n //int aj = j < row_end ? r*Ax[j] : ValueTypeB(0);\n if (j < row_end)\n {\n utils::atomic_add(&xout[Aj[j]], r * Ax[j]);\n }\n }\n }\n}\n\ntemplate\n__global__ void kaczmarz_smooth_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const IndexType *amax,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB *xout,\n const IndexType row_offset)\n\n{\n // Naive implementation, needs x copy in xout at the very beginning\n //IndexType tidx = blockDim.x*blockIdx.x + threadIdx.x;\n IndexType i, t;\n const int num_warps = kCtaSize / 32;\n const int num_rows_per_iter = num_warps * gridDim.x;\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n for ( int ridx = blockIdx.x * num_warps + warpId ; ridx < num_rows ;\n ridx += num_rows_per_iter )\n {\n ValueTypeB Axi = 0.0;\n ValueTypeB r;\n i = ourHash(ridx, num_rows);\n IndexType row_start = Ap[i];\n IndexType row_end = Ap[i + 1];\n\n for (int j = row_start + laneId; utils::any( j < row_end) ; j += 32)\n {\n ValueTypeB aValue = j < row_end ? Ax[j] : ValueTypeB(0);\n ValueTypeB xValue = j < row_end ? xout[Aj[j]] : ValueTypeB(0);\n r = utils::warp_reduce<1, utils::Add>(aValue * xValue);\n Axi += r;\n //Axi += utils::Warp_reduce_linear<1,32>::execute(aValue * xValue);\n //Axi += Ax[j] * xout[Aj[j]];\n printf(\"j = %d, r = %f\\n\", j, r);\n }\n\n if (laneId == 0)\n {\n r = (b[i] - Axi);// / ( isNotCloseToZero( d[ridx]) ? d[ridx] : epsilon(d[ridx]) );\n t = row_start + ourHash(ridx, row_end - row_start);\n printf(\"ridx=%d, i=%d, t=%d, Aj[t]=%d, r=%f\\n\", ridx, i, t, Aj[t], r);\n xout[Aj[t]] += r * ((row_end - row_start) * Ax[t]) * 0.5;\n }\n }\n}\n\ntemplate\n__global__ void multicolor_kaczmarz_smooth_kernel_naive(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const ValueTypeA *d,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB weight,\n const int *sorted_rows_by_color,\n const int num_rows_per_color,\n ValueTypeB *xout)\n\n{\n int i;\n // Naive implementation, needs x copy in xout at the very beginning\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = tidx; ridx < num_rows_per_color; ridx += blockDim.x * gridDim.x)\n {\n i = sorted_rows_by_color[ridx];\n IndexType row_start = Ap[i];\n IndexType row_end = Ap[i + 1];\n ValueTypeB Axi = 0.0;\n ValueTypeB r;\n\n for (int j = row_start; j < row_end; j++)\n {\n Axi += Ax[j] * xout[Aj[j]];\n }\n\n r = (b[i] - Axi) / ( isNotCloseToZero( d[i]) ? d[i] : epsilon(d[i]) );\n\n for (int j = row_start; j < row_end; j++)\n {\n utils::atomic_add(&xout[Aj[j]], r * Ax[j]);\n }\n }\n}\n\ntemplate\n__global__ void multicolor_kaczmarz_smooth_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const ValueTypeA *d,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB weight,\n const int *sorted_rows_by_color,\n const int num_rows_per_color,\n ValueTypeB *xout)\n\n{\n const int num_warps = kCtaSize / 32;\n const int num_rows_per_iter = num_warps * gridDim.x;\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n int i;\n\n for ( int ridx = blockIdx.x * num_warps + warpId ; ridx < num_rows_per_color ;\n ridx += num_rows_per_iter )\n {\n i = sorted_rows_by_color[ridx];\n IndexType row_start = Ap[i];\n IndexType row_end = Ap[i + 1];\n ValueTypeB Axi = 0.0;\n ValueTypeB r;\n\n for (int j = row_start + laneId; utils::any( j < row_end) ; j += 32)\n {\n ValueTypeB aValue = j < row_end ? Ax[j] : ValueTypeB(0);\n ValueTypeB xValue = j < row_end ? xout[Aj[j]] : ValueTypeB(0);\n r = utils::warp_reduce<1, utils::Add>(aValue * xValue);\n Axi += r;\n }\n\n r = (b[i] - Axi) / ( isNotCloseToZero( d[i]) ? d[i] : epsilon(d[i]) );\n\n for (int j = row_start + laneId; utils::any( j < row_end) ; j += 32)\n {\n if (j < row_end)\n {\n utils::atomic_add(&xout[Aj[j]], r * Ax[j]);\n }\n }\n }\n}\n\ntemplate\n__global__ void jacobi_smooth_with_0_initial_guess_kernel(const IndexType num_rows,\n const ValueTypeA *d,\n const ValueTypeB *b,\n ValueTypeB *x,\n ValueTypeB weight,\n const IndexType row_offset)\n\n{\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = row_offset + tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n x[ridx] = weight * b[ridx] / ( isNotCloseToZero( d[ridx]) ? d[ridx] : epsilon(d[ridx]) );\n }\n}\n\n// -----------------\n// Methods\n// -----------------\n\n// Constructor\ntemplate\nKaczmarzSolver_Base::KaczmarzSolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope), m_an(0), m_amax(0), m_c_inv(0)\n{\n weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n this->m_coloring_needed = (cfg.AMG_Config::template getParameter(\"kaczmarz_coloring_needed\", cfg_scope) != 0);\n this->m_reorder_cols_by_color_desired = (cfg.AMG_Config::template getParameter(\"reorder_cols_by_color\", cfg_scope) != 0);\n this->m_randomized = true;\n\n if (weight == 0)\n {\n weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Block Jacobi smoother\\n\");;\n }\n}\n\n// Destructor\ntemplate\nKaczmarzSolver_Base::~KaczmarzSolver_Base()\n{\n}\n\n// Solver setup\ntemplate\nvoid\nKaczmarzSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n m_explicit_A = dynamic_cast*>(Base::m_A);\n\n if (!m_explicit_A)\n {\n FatalError(\"Kaczmarz solver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n compute_anorm( *this->m_explicit_A );\n\n if (m_randomized) // MC RK is not supported here\n {\n if (m_coloring_needed)\n {\n //FatalError(\"Randomized Kaczmarz solver does not support coloring\", AMGX_ERR_INTERNAL);\n m_coloring_needed = false;\n }\n\n double d_inv = this->m_an[0];\n int c_sz = this->m_an.size();\n d_inv = amgx::thrust::reduce(this->m_an.begin(), this->m_an.end(), d_inv, amgx::thrust::minimum());\n amgx::thrust::inclusive_scan(this->m_an.begin(), this->m_an.end(), this->m_an.begin()); // in-place scan\n int c_inv_sz = (this->m_an[c_sz - 1] + d_inv - 1 ) / d_inv;\n this->m_c_inv.resize(c_inv_sz, -1);\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((ValueTypeB)c_sz / (ValueTypeB)THREADS_PER_BLOCK));\n\n if (c_sz > 0)\n {\n device_vector_alloc aa(c_sz, 1);\n compute_cumul_inv_kernel <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK>>>\n (c_sz,\n this->m_an.raw(),\n d_inv,\n c_inv_sz,\n this->m_c_inv.raw());\n cudaCheckError();\n }\n\n hipDeviceSynchronize();\n cudaCheckError();\n }\n}\n\ntemplate\nvoid KaczmarzSolver_Base::compute_anorm( Matrix &A)\n{\n this->m_an.resize(A.get_num_rows()*A.get_block_dimx());\n ViewType oldView = A.currentView();\n A.setView(this->m_explicit_A->getViewExterior());\n\n if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n compute_anorm_1x1(A);\n }\n else\n {\n FatalError(\"Unsupported block size for KaczmarzSolver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n}\n\ntemplate\nvoid KaczmarzSolver_Base::compute_amax( Matrix &A)\n{\n this->m_amax.resize(A.get_num_rows()*A.get_block_dimx());\n ViewType oldView = A.currentView();\n A.setView(this->m_explicit_A->getViewExterior());\n\n if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n compute_amax_1x1(A);\n }\n else\n {\n FatalError(\"Unsupported block size for KaczmarzSolver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n}\n\ntemplate \nvoid KaczmarzSolver >::compute_anorm_1x1(const Matrix_d &A)\n{\n//DIAG: starnge issues trying to add DIAG property handling\n// now leaving !DIAG only\n if (A.hasProps(DIAG))\n {\n FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((ValueTypeB)A.get_num_rows() / (ValueTypeB)THREADS_PER_BLOCK));\n\n if (A.get_num_rows() > 0)\n {\n compute_anorm_kernel <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK>>>\n ((int)A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->m_an.raw());\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid KaczmarzSolver >::compute_amax_1x1(const Matrix_d &A)\n{\n//DIAG: starnge issues trying to add DIAG property handling\n// now leaving !DIAG only\n if (A.hasProps(DIAG))\n {\n FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((ValueTypeB)A.get_num_rows() / (ValueTypeB)THREADS_PER_BLOCK));\n\n if (A.get_num_rows() > 0)\n {\n compute_amax_kernel <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK>>>\n ((int)A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->m_amax.raw());\n }\n\n cudaCheckError();\n}\n\ntemplate\nvoid\nKaczmarzSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nKaczmarzSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n if (xIsZero) { x.dirtybit = 0; }\n\n if (!this->m_explicit_A->is_matrix_singleGPU())\n {\n this->m_explicit_A->manager->exchange_halo_async(x, x.tag);\n\n if (this->m_explicit_A->getViewExterior() == this->m_explicit_A->getViewInterior())\n {\n this->m_explicit_A->manager->exchange_halo_wait(x, x.tag);\n }\n }\n\n ViewType oldView = this->m_explicit_A->currentView();\n ViewType flags;\n bool latencyHiding = true;\n\n if (this->m_explicit_A->is_matrix_singleGPU() || (x.dirtybit == 0))\n {\n latencyHiding = false;\n this->m_explicit_A->setViewExterior();\n flags = this->m_explicit_A->getViewExterior();\n }\n else\n {\n flags = this->m_explicit_A->getViewInterior();\n this->m_explicit_A->setViewInterior();\n }\n\n if (this->m_explicit_A->get_block_dimx() == 1 && this->m_explicit_A->get_block_dimy() == 1)\n {\n if (xIsZero)\n {\n //smooth_with_0_initial_guess_1x1(*this->m_explicit_A, b, x, flags);\n smooth_1x1(*this->m_explicit_A, b, x, flags, latencyHiding);\n }\n else\n {\n smooth_1x1(*this->m_explicit_A, b, x, flags, latencyHiding);\n }\n }\n else\n {\n FatalError(\"Unsupported block size for Kaczmarz_Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n x.dirtybit = 1;\n this->m_explicit_A->setView(oldView);\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nKaczmarzSolver_Base::solve_finalize( VVector &b, VVector &x )\n{\n}\n\n// Multicolor version\ntemplate \nvoid KaczmarzSolver >::smooth_1x1_MC(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags, bool latency_hiding)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n ValueTypeB *x_ptr = x.raw();\n IndexType num_rows = A.get_num_rows();\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n const IndexType *A_sorted_rows_by_color_ptr = A.getMatrixColoring().getSortedRowsByColor().raw();\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = ((separation_flags & INTERIOR) == 0) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = ((separation_flags == this->m_explicit_A->getViewInterior()) ? A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i] : A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1]) - color_offset;\n\n if (num_rows_per_color == 0) { continue; }\n\n const int threads_per_block = 128;\n const int blockrows_per_warp = 1;\n const int blockrows_per_cta = (threads_per_block / 32) * blockrows_per_warp;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows_per_color / blockrows_per_cta + 1));\n multicolor_kaczmarz_smooth_kernel <<< num_blocks, threads_per_block >>>\n (A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->m_an.raw(),\n b.raw(),\n x_ptr,\n this->weight,\n A_sorted_rows_by_color_ptr + color_offset, num_rows_per_color,\n x_ptr);\n cudaCheckError();\n }\n}\n\ntemplate \nvoid KaczmarzSolver >::smooth_1x1_naive(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags, bool latency_hiding)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n ValueTypeB *x_ptr = x.raw();\n IndexType num_rows = A.get_num_rows();\n IndexType offset = 0;\n // Skipping Multi-GPU logic for now\n // Current. Will be exact only with one warp per grid\n const int threads_per_block = 32;\n const int num_blocks = 1;\n\n if (this->m_randomized)\n {\n IVector rnd_rows;\n int c_inv_sz = this->m_c_inv.size();\n initRandom(rnd_rows, A.get_num_rows(), c_inv_sz);\n randomized_kaczmarz_smooth_kernel_warp_atomics <<< num_blocks, threads_per_block >>>\n (A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->m_c_inv.raw(),\n rnd_rows.raw(),\n b.raw(),\n x_ptr,\n x_ptr,\n offset);\n cudaCheckError();\n }\n else\n {\n kaczmarz_smooth_kernel_warp_atomics <<< num_blocks, threads_per_block >>>\n (A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->m_an.raw(),\n b.raw(),\n x_ptr,\n x_ptr,\n offset);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid KaczmarzSolver >::smooth_1x1(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags, bool latency_hiding)\n{\n if (this->m_coloring_needed)\n {\n smooth_1x1_MC(A, b, x, separation_flags, latency_hiding);\n }\n else\n {\n smooth_1x1_naive(A, b, x, separation_flags, latency_hiding);\n }\n}\n\ntemplate \nvoid KaczmarzSolver >::smooth_with_0_initial_guess_1x1(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n ViewType oldView = A.currentView();\n // Process all rows\n A.setViewExterior();\n ViewType flags = A.getViewExterior();\n int offset, num_rows;\n A.getOffsetAndSizeForView(flags, &offset, &num_rows);\n A.setView(oldView);\n cudaCheckError();\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class KaczmarzSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class KaczmarzSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/solvers/kaczmarz_solver.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\n// Constructor\ntemplate\nGaussSeidelSolver_Base::GaussSeidelSolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope)\n{\n this->weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n\n if (this->weight == 0)\n {\n this->weight = 1.;\n }\n}\n\n// Destructor\ntemplate\nGaussSeidelSolver_Base::~GaussSeidelSolver_Base()\n{\n}\n\ntemplate\nvoid\nGaussSeidelSolver_Base::printSolverParameters() const\n{\n std::cout << \"relaxation_factor= \" << this->weight << std::endl;\n}\n\n// Solver setup\ntemplate\nvoid\nGaussSeidelSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n Matrix *A_as_matrix = dynamic_cast*>(this->m_A);\n\n if (!A_as_matrix)\n {\n FatalError(\"GaussSeidelSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n if ( A_as_matrix->get_block_size() != 1 )\n {\n FatalError(\"Unsupported block size for GaussSeidelSolver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n this->diag.resize((int) A_as_matrix->get_num_rows());\n\n if (A_as_matrix->hasProps(DIAG))\n {\n amgx::thrust::copy( A_as_matrix->values.begin() + A_as_matrix->diagOffset()*A_as_matrix->get_block_size()/*block_size == 1*/, A_as_matrix->values.end(), this->diag.begin());\n cudaCheckError();\n }\n else\n {\n find_diag(*A_as_matrix);\n }\n}\n\n//\ntemplate\nvoid\nGaussSeidelSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nGaussSeidelSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n\n if ( A_as_matrix->get_block_size() == 1)\n {\n if (xIsZero)\n {\n smooth_with_0_initial_guess_1x1( *A_as_matrix, b, x );\n }\n else\n {\n smooth_1x1( *A_as_matrix, b, x );\n }\n }\n else\n {\n FatalError(\"Unsupported block size for GaussSeidel_Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nGaussSeidelSolver_Base::solve_finalize( VVector &b, VVector &x )\n{\n}\n\ntemplate\nvoid\nGaussSeidelSolver_Base::find_diag( const Matrix &in_A )\n{\n if (in_A.get_block_size() == 1)\n {\n find_diag_1x1(in_A);\n }\n else\n {\n FatalError(\"Unsupported block size for GaussSeidelSolver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\n/********************************************/\n\n// Finding diag on host, CSR format\ntemplate \nvoid GaussSeidelSolver >::find_diag_1x1(const Matrix_h &A)\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n if (A.col_indices[j] == i)\n {\n this->diag[i] = A.values[j];\n break;\n }\n\n if (j == A.row_offsets[i + 1] - 1)\n {\n std::string error = \"Could not find a diagonal value at row \" + std::to_string(i);\n FatalError(error.c_str(), AMGX_ERR_BAD_PARAMETERS);\n }\n }\n }\n}\n\ntemplate \nvoid GaussSeidelSolver >::smooth_1x1(const Matrix_h &A, const VVector &b, VVector &x)\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeB Axi = 0.0;\n ValueTypeA diag = isNotCloseToZero(this->diag[i]) ? this->diag[i] : epsilon(this->diag[i]);\n\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n Axi += A.values[j] * x[A.col_indices[j]];\n }\n\n x[i] = x[i] + this->weight * (b[i] - Axi) / diag;\n }\n}\n\ntemplate \nvoid GaussSeidelSolver >::smooth_with_0_initial_guess_1x1(const Matrix_h &A, const VVector &b, VVector &x)\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeA diag = isNotCloseToZero(this->diag[i]) ? this->diag[i] : epsilon(this->diag[i]);\n x[i] = this->weight * b[i] / diag;\n }\n}\n\ntemplate \nvoid GaussSeidelSolver >::find_diag_1x1(const Matrix_d &A)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((ValueType)A.get_num_rows() / (ValueType)THREADS_PER_BLOCK));\n this->diag.resize(A.get_num_rows());\n find_diag_kernel <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK >>>\n ((int)A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->diag.raw());\n cudaCheckError();\n}\n\ntemplate\n__global__ void GS_smooth_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const ValueTypeB *diag,\n const ValueTypeB *b,\n const ValueTypeB weight,\n ValueTypeB *x)\n\n{\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n IndexType row_start = Ap[ridx];\n IndexType row_end = Ap[ridx + 1];\n ValueTypeB Axi = 0.0;\n ValueTypeA mydiag = isNotCloseToZero(diag[ridx]) ? diag[ridx] : epsilon(diag[ridx]);\n\n for (int j = row_start; j < row_end; j++)\n {\n Axi += Ax[j] * x[Aj[j]];\n }\n\n ValueTypeB tmp = x[ridx] + weight * (b[ridx] - Axi) / mydiag;\n x[ridx] = tmp;\n }\n}\n\n\n\ntemplate \nvoid GaussSeidelSolver >::smooth_1x1(const Matrix_d &A, const VVector &b, VVector &x)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n typedef typename VVector::value_type ValueTypeB;\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((ValueTypeB)A.get_num_rows() / (ValueTypeB)THREADS_PER_BLOCK));\n GS_smooth_kernel <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK >>>\n ((int)A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->diag.raw(),\n b.raw(),\n this->weight,\n x.raw());\n cudaCheckError();\n}\n\ntemplate \nvoid GaussSeidelSolver >::smooth_with_0_initial_guess_1x1(const Matrix_d &A, const VVector &b, VVector &x)\n{\n thrust_wrapper::fill(x.begin(), x.end(), ValueTypeB(0));\n cudaCheckError();\n smooth_1x1(A, b, x);\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class GaussSeidelSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class GaussSeidelSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\n// Constructor\ntemplate\nGaussSeidelSolver_Base::GaussSeidelSolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope)\n{\n this->weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n\n if (this->weight == 0)\n {\n this->weight = 1.;\n }\n}\n\n// Destructor\ntemplate\nGaussSeidelSolver_Base::~GaussSeidelSolver_Base()\n{\n}\n\ntemplate\nvoid\nGaussSeidelSolver_Base::printSolverParameters() const\n{\n std::cout << \"relaxation_factor= \" << this->weight << std::endl;\n}\n\n// Solver setup\ntemplate\nvoid\nGaussSeidelSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n Matrix *A_as_matrix = dynamic_cast*>(this->m_A);\n\n if (!A_as_matrix)\n {\n FatalError(\"GaussSeidelSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n if ( A_as_matrix->get_block_size() != 1 )\n {\n FatalError(\"Unsupported block size for GaussSeidelSolver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n this->diag.resize((int) A_as_matrix->get_num_rows());\n\n if (A_as_matrix->hasProps(DIAG))\n {\n amgx::thrust::copy( A_as_matrix->values.begin() + A_as_matrix->diagOffset()*A_as_matrix->get_block_size()/*block_size == 1*/, A_as_matrix->values.end(), this->diag.begin());\n cudaCheckError();\n }\n else\n {\n find_diag(*A_as_matrix);\n }\n}\n\n//\ntemplate\nvoid\nGaussSeidelSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nGaussSeidelSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n\n if ( A_as_matrix->get_block_size() == 1)\n {\n if (xIsZero)\n {\n smooth_with_0_initial_guess_1x1( *A_as_matrix, b, x );\n }\n else\n {\n smooth_1x1( *A_as_matrix, b, x );\n }\n }\n else\n {\n FatalError(\"Unsupported block size for GaussSeidel_Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nGaussSeidelSolver_Base::solve_finalize( VVector &b, VVector &x )\n{\n}\n\ntemplate\nvoid\nGaussSeidelSolver_Base::find_diag( const Matrix &in_A )\n{\n if (in_A.get_block_size() == 1)\n {\n find_diag_1x1(in_A);\n }\n else\n {\n FatalError(\"Unsupported block size for GaussSeidelSolver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\n/********************************************/\n\n// Finding diag on host, CSR format\ntemplate \nvoid GaussSeidelSolver >::find_diag_1x1(const Matrix_h &A)\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n if (A.col_indices[j] == i)\n {\n this->diag[i] = A.values[j];\n break;\n }\n\n if (j == A.row_offsets[i + 1] - 1)\n {\n std::string error = \"Could not find a diagonal value at row \" + std::to_string(i);\n FatalError(error.c_str(), AMGX_ERR_BAD_PARAMETERS);\n }\n }\n }\n}\n\ntemplate \nvoid GaussSeidelSolver >::smooth_1x1(const Matrix_h &A, const VVector &b, VVector &x)\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeB Axi = 0.0;\n ValueTypeA diag = isNotCloseToZero(this->diag[i]) ? this->diag[i] : epsilon(this->diag[i]);\n\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n Axi += A.values[j] * x[A.col_indices[j]];\n }\n\n x[i] = x[i] + this->weight * (b[i] - Axi) / diag;\n }\n}\n\ntemplate \nvoid GaussSeidelSolver >::smooth_with_0_initial_guess_1x1(const Matrix_h &A, const VVector &b, VVector &x)\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeA diag = isNotCloseToZero(this->diag[i]) ? this->diag[i] : epsilon(this->diag[i]);\n x[i] = this->weight * b[i] / diag;\n }\n}\n\ntemplate \nvoid GaussSeidelSolver >::find_diag_1x1(const Matrix_d &A)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((ValueType)A.get_num_rows() / (ValueType)THREADS_PER_BLOCK));\n this->diag.resize(A.get_num_rows());\n find_diag_kernel <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK >>>\n ((int)A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->diag.raw());\n cudaCheckError();\n}\n\ntemplate\n__global__ void GS_smooth_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const ValueTypeB *diag,\n const ValueTypeB *b,\n const ValueTypeB weight,\n ValueTypeB *x)\n\n{\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n IndexType row_start = Ap[ridx];\n IndexType row_end = Ap[ridx + 1];\n ValueTypeB Axi = 0.0;\n ValueTypeA mydiag = isNotCloseToZero(diag[ridx]) ? diag[ridx] : epsilon(diag[ridx]);\n\n for (int j = row_start; j < row_end; j++)\n {\n Axi += Ax[j] * x[Aj[j]];\n }\n\n ValueTypeB tmp = x[ridx] + weight * (b[ridx] - Axi) / mydiag;\n x[ridx] = tmp;\n }\n}\n\n\n\ntemplate \nvoid GaussSeidelSolver >::smooth_1x1(const Matrix_d &A, const VVector &b, VVector &x)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n typedef typename VVector::value_type ValueTypeB;\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((ValueTypeB)A.get_num_rows() / (ValueTypeB)THREADS_PER_BLOCK));\n GS_smooth_kernel <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK >>>\n ((int)A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->diag.raw(),\n b.raw(),\n this->weight,\n x.raw());\n cudaCheckError();\n}\n\ntemplate \nvoid GaussSeidelSolver >::smooth_with_0_initial_guess_1x1(const Matrix_d &A, const VVector &b, VVector &x)\n{\n thrust_wrapper::fill(x.begin(), x.end(), ValueTypeB(0));\n cudaCheckError();\n smooth_1x1(A, b, x);\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class GaussSeidelSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class GaussSeidelSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/solvers/gauss_seidel_solver.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#ifdef _WIN32\n#pragma warning (push)\n#pragma warning (disable : 4244 4267 4521)\n#endif\n#include \n#include \n#ifdef _WIN32\n#pragma warning (pop)\n#endif\n\n#include \n\n\nnamespace amgx\n{\nnamespace polynomial_solver\n{\n\n\ntemplate \n__global__\nvoid aux_norm1_csr(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *dia_values, const ValueTypeA *nonzero_values,\n const int num_rows, int bsize, ValueTypeB *row_sum)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int bsize_sq = bsize * bsize;\n\n while (tid < num_rows)\n {\n ValueTypeB tmp_sum = 0;\n int idx_i = tid / bsize;\n int offset_i = tid % bsize;\n\n for (int i = 0; i < bsize; i++)\n {\n tmp_sum += fabs(nonzero_values[dia_values[idx_i] * bsize_sq + offset_i * bsize + i]);\n }\n\n for (int j = row_offsets[idx_i]; j < row_offsets[idx_i + 1]; j++)\n {\n // Compute edge weight\n for (int i = 0; i < bsize; i++)\n {\n tmp_sum += fabs(nonzero_values[j * bsize_sq + offset_i * bsize + i]);\n }\n }\n\n row_sum[tid] = tmp_sum / fabs(nonzero_values[dia_values[idx_i] * bsize_sq + offset_i * bsize + offset_i]);\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \n__global__\nvoid get_diaginv(const IndexType *dia_idx, const ValueTypeA *nonzero_values,\n const int num_rows, int bsize, ValueTypeB *Dinv)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int bsize_sq = bsize * bsize;\n\n while (tid < num_rows)\n {\n int idx_i = tid / bsize;\n int offset_i = tid % bsize;\n ValueTypeB diag = nonzero_values[dia_idx[idx_i] * bsize_sq + offset_i * bsize + offset_i];\n Dinv[tid] = ValueTypeB(1) / (isNotCloseToZero(diag) ? diag : epsilon(diag) );\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\n// Constructor\ntemplate\nPolynomialSolverBase::PolynomialSolverBase( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope), R(0)\n{\n ndeg0 = cfg.AMG_Config::template getParameter(\"kpz_order\", cfg_scope);\n}\n\n// Destructor\ntemplate\nPolynomialSolverBase::~PolynomialSolverBase()\n{\n}\n\ntemplate\nvoid\nPolynomialSolverBase::printSolverParameters() const\n{\n std::cout << \"kpz_order = \" << this->ndeg0 << std::endl;\n}\n\n\n// Solver setup\ntemplate\nvoid\nPolynomialSolverBase::solver_setup(bool reuse_matrix_structure)\n{\n this->m_explicit_A = dynamic_cast*>(this->m_A);\n\n if (!this->m_explicit_A)\n {\n FatalError(\"PolynomialSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n if (ndeg0 == 0) { ndeg0 = 6; }\n\n int N = this->m_explicit_A->get_num_rows() * this->m_explicit_A->get_block_dimy();\n ValueTypeA mu0, mu1, smu0, smu1;\n const IndexType *A_row_offsets_ptr = this->m_explicit_A->row_offsets.raw();\n const IndexType *A_column_indices_ptr = this->m_explicit_A->col_indices.raw();\n const IndexType *A_dia_ptr = this->m_explicit_A->diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = this->m_explicit_A->values.raw();\n VVector row_sum(N);\n ValueTypeB *row_sum_ptr = row_sum.raw();\n const int threads_per_block = 512;\n const int num_blocks = std::min((N - 1) / threads_per_block + 1, AMGX_GRID_MAX_SIZE);\n aux_norm1_csr <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ptr, A_nonzero_values_ptr,\n N, this->m_explicit_A->get_block_dimy(), row_sum_ptr);\n cudaCheckError();\n mu0 = nrmmax(row_sum);\n cudaCheckError();\n\n if (mu0 == 0) { mu0 = nrmmax(this->m_explicit_A->values); }\n\n cudaCheckError();\n mu0 = 1.0 / mu0;\n mu1 = 4.0 * mu0; // default set 8;\n smu0 = sqrt(mu0);\n smu1 = sqrt(mu1);\n k[1] = (mu0 + mu1) / 2.0;\n k[2] = (smu0 + smu1) * (smu0 + smu1) / 2.0;\n k[3] = mu0 * mu1;\n k[4] = 2.0 * k[3] / k[2]; // 4.0*mu0*mu1/(sqrt(mu0)+sqrt(mu1))/(sqrt(mu0)+sqrt(mu1));\n k[5] = (mu1 - 2.0 * smu0 * smu1 + mu0) / (mu1 + 2.0 * smu0 * smu1 + mu0); // square of (sqrt(kappa)-1)/(sqrt(kappa)+1);\n\n if (this->m_explicit_A->get_block_size() == 1)\n {\n MatrixCusp wA((Matrix *) &*this->m_explicit_A);\n ValueTypeA rho = cusp::detail::ritz_spectral_radius_symmetric(wA, 8);\n cudaCheckError();\n cusp::array1d coeffs;\n cusp::relaxation::detail::chebyshev_polynomial_coefficients(rho, coeffs);\n cudaCheckError();\n poly = cusp::relaxation::polynomial::memory_space > (wA, coeffs);\n cudaCheckError();\n }\n\n R.resize(N);\n V0.resize(N);\n V.resize(N);\n Rbar.resize(N);\n Sn.resize(N);\n Dinv.resize(N);\n R.set_block_dimy(this->m_explicit_A->get_block_dimy());\n R.set_block_dimx(1);\n V0.set_block_dimy(this->m_explicit_A->get_block_dimy());\n V0.set_block_dimx(1);\n V.set_block_dimy(this->m_explicit_A->get_block_dimy());\n V.set_block_dimx(1);\n Rbar.set_block_dimy(this->m_explicit_A->get_block_dimy());\n Rbar.set_block_dimx(1);\n Sn.set_block_dimy(this->m_explicit_A->get_block_dimy());\n Sn.set_block_dimx(1);\n Dinv.set_block_dimy(this->m_explicit_A->get_block_dimy());\n Dinv.set_block_dimx(1);\n ValueTypeB *Dinv_ptr = Dinv.raw();\n get_diaginv <<< num_blocks, threads_per_block>>>(A_dia_ptr, A_nonzero_values_ptr, N, this->m_explicit_A->get_block_dimy(), Dinv_ptr);\n cudaCheckError();\n}\n\n//\ntemplate\nvoid\nPolynomialSolverBase::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nPolynomialSolverBase::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n smooth_common_sqblocks( *this->m_explicit_A, b, x );\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nPolynomialSolverBase::solve_finalize( VVector &b, VVector &x )\n{\n}\n\ntemplate\nstruct poly_smooth\n{\n static void poly_postsmooth(const Matrix &A, const Vector &B, Vector &C, cusp::relaxation::polynomial &poly)\n {\n FatalError(\"Mixed precision is not supported for scalar matrix type\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n static void poly_presmooth(const Matrix &A, const Vector &B, Vector &C, cusp::relaxation::polynomial &poly)\n {\n FatalError(\"Mixed precision is not supported for scalar matrix type\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n};\n\ntemplate\nstruct poly_smooth>\n{\n static void poly_postsmooth(const Matrix > &A, const Vector > &B, Vector > &C, cusp::relaxation::polynomial::Type> &poly)\n {\n MatrixCusp, cusp::csr_format> wA((Matrix > *) &A);\n poly.postsmooth(wA, B, C);\n cudaCheckError();\n }\n\n static void poly_presmooth(const Matrix > &A, const Vector > &B, Vector > &C, cusp::relaxation::polynomial::Type> &poly)\n {\n MatrixCusp, cusp::csr_format> wA((Matrix > *) &A);\n poly.presmooth(wA, B, C);\n cudaCheckError();\n }\n};\n\ntemplate\nstruct poly_smooth>\n{\n static void poly_postsmooth(const Matrix > &A, const Vector > &B, Vector > &C, cusp::relaxation::polynomial::Type> &poly)\n {\n MatrixCusp, cusp::csr_format> wA((Matrix > *) &A);\n poly.postsmooth(wA, B, C);\n cudaCheckError();\n }\n\n static void poly_presmooth(const Matrix > &A, const Vector > &B, Vector > &C, cusp::relaxation::polynomial::Type> &poly)\n {\n MatrixCusp, cusp::csr_format> wA((Matrix > *) &A);\n poly.presmooth(wA, B, C);\n cudaCheckError();\n }\n};\n\ntemplate \nvoid PolynomialSolver >::smooth_1x1(const Matrix_h &A, const VVector &B, VVector &C)\n{\n FatalError(\"Unsupported on host for Polynomial smoother\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\ntemplate \nvoid PolynomialSolver >::smooth_1x1(const Matrix_d &A, const VVector &B, VVector &C)\n{\n if (A.hasProps(DIAG))\n {\n FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (A.get_block_size() != 1)\n {\n FatalError(\"Unsupported block size for PolynomialSolver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n poly_smooth>::poly_postsmooth(A, B, C, this->poly);\n}\n\ntemplate \nvoid PolynomialSolver >::smooth_common_sqblocks(Matrix_h &A, const VVector &b, VVector &x)\n{\n FatalError(\"Unsupported on host for Polynomial smoother\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\ntemplate \nvoid PolynomialSolver >::smooth_common_sqblocks(Matrix_d &A, const VVector &b, VVector &x)\n{\n //0, get residule : R = b - Ax\n multiply( A, x, this->R );\n cusp::blas::axpby( b, this->R, this->R, ValueTypeB(1), ValueTypeB(-1) );\n //1, set up rbar\n cusp::blas::xmy(this->Dinv, this->R, this->Rbar);\n //2, set up V and V0\n multiply( A, this->Rbar, this->V );\n cusp::blas::xmy(this->Dinv, this->V, this->V);\n this->V0 = this->Rbar;\n cusp::blas::scal(this->V0, this->k[1]);\n cusp::blas::axpby(this->Rbar, this->V, this->V, this->k[2], this->k[3]);\n\n //3, iterate to get v_{j+1}\n for (int i = 0; i < this->ndeg0; i++)\n {\n multiply( A, this->V, this->Rbar );\n cusp::blas::axpby( this->R, this->Rbar, this->Rbar, ValueTypeB(1), ValueTypeB(-1) );\n cusp::blas::xmy(this->Dinv, this->Rbar, this->Rbar);\n cusp::blas::axpbypcz(this->Rbar, this->V, this->V0, this->Sn,\n this->k[4], this->k[5] + 1.0, -1.0 * this->k[5]);\n // V0 = V\n this->V0 = this->V;\n // V = V+Sn\n cusp::blas::copy(this->Sn, this->V);\n }\n\n //4, update solution\n cusp::blas::axpy(this->V, x, ValueTypeB(1));\n}\n\n\ntemplate \nvoid PolynomialSolver >::smooth_with_0_initial_guess_1x1(const Matrix_h &A, const VVector &b, VVector &x)\n{\n FatalError(\"Unsupported on host for Polynomial smoother\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\ntemplate \nvoid PolynomialSolver >::smooth_with_0_initial_guess_1x1(const Matrix_d &A, const VVector &b, VVector &x)\n{\n if (A.hasProps(DIAG))\n {\n FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (A.get_block_size() != 1)\n {\n FatalError(\"Unsupported block size for PolynomialSolver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n poly_smooth>::poly_presmooth(A, b, x, this->poly);\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class PolynomialSolverBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class PolynomialSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n\n\n\n} // namespace polynomial_solver\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#ifdef _WIN32\n#pragma warning (push)\n#pragma warning (disable : 4244 4267 4521)\n#endif\n#include \n#include \n#ifdef _WIN32\n#pragma warning (pop)\n#endif\n\n#include \n\n\nnamespace amgx\n{\nnamespace polynomial_solver\n{\n\n\ntemplate \n__global__\nvoid aux_norm1_csr(const IndexType *row_offsets, const IndexType *column_indices, const IndexType *dia_values, const ValueTypeA *nonzero_values,\n const int num_rows, int bsize, ValueTypeB *row_sum)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int bsize_sq = bsize * bsize;\n\n while (tid < num_rows)\n {\n ValueTypeB tmp_sum = 0;\n int idx_i = tid / bsize;\n int offset_i = tid % bsize;\n\n for (int i = 0; i < bsize; i++)\n {\n tmp_sum += fabs(nonzero_values[dia_values[idx_i] * bsize_sq + offset_i * bsize + i]);\n }\n\n for (int j = row_offsets[idx_i]; j < row_offsets[idx_i + 1]; j++)\n {\n // Compute edge weight\n for (int i = 0; i < bsize; i++)\n {\n tmp_sum += fabs(nonzero_values[j * bsize_sq + offset_i * bsize + i]);\n }\n }\n\n row_sum[tid] = tmp_sum / fabs(nonzero_values[dia_values[idx_i] * bsize_sq + offset_i * bsize + offset_i]);\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \n__global__\nvoid get_diaginv(const IndexType *dia_idx, const ValueTypeA *nonzero_values,\n const int num_rows, int bsize, ValueTypeB *Dinv)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int bsize_sq = bsize * bsize;\n\n while (tid < num_rows)\n {\n int idx_i = tid / bsize;\n int offset_i = tid % bsize;\n ValueTypeB diag = nonzero_values[dia_idx[idx_i] * bsize_sq + offset_i * bsize + offset_i];\n Dinv[tid] = ValueTypeB(1) / (isNotCloseToZero(diag) ? diag : epsilon(diag) );\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\n// Constructor\ntemplate\nPolynomialSolverBase::PolynomialSolverBase( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope), R(0)\n{\n ndeg0 = cfg.AMG_Config::template getParameter(\"kpz_order\", cfg_scope);\n}\n\n// Destructor\ntemplate\nPolynomialSolverBase::~PolynomialSolverBase()\n{\n}\n\ntemplate\nvoid\nPolynomialSolverBase::printSolverParameters() const\n{\n std::cout << \"kpz_order = \" << this->ndeg0 << std::endl;\n}\n\n\n// Solver setup\ntemplate\nvoid\nPolynomialSolverBase::solver_setup(bool reuse_matrix_structure)\n{\n this->m_explicit_A = dynamic_cast*>(this->m_A);\n\n if (!this->m_explicit_A)\n {\n FatalError(\"PolynomialSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n if (ndeg0 == 0) { ndeg0 = 6; }\n\n int N = this->m_explicit_A->get_num_rows() * this->m_explicit_A->get_block_dimy();\n ValueTypeA mu0, mu1, smu0, smu1;\n const IndexType *A_row_offsets_ptr = this->m_explicit_A->row_offsets.raw();\n const IndexType *A_column_indices_ptr = this->m_explicit_A->col_indices.raw();\n const IndexType *A_dia_ptr = this->m_explicit_A->diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = this->m_explicit_A->values.raw();\n VVector row_sum(N);\n ValueTypeB *row_sum_ptr = row_sum.raw();\n const int threads_per_block = 512;\n const int num_blocks = std::min((N - 1) / threads_per_block + 1, AMGX_GRID_MAX_SIZE);\n aux_norm1_csr <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_ptr, A_nonzero_values_ptr,\n N, this->m_explicit_A->get_block_dimy(), row_sum_ptr);\n cudaCheckError();\n mu0 = nrmmax(row_sum);\n cudaCheckError();\n\n if (mu0 == 0) { mu0 = nrmmax(this->m_explicit_A->values); }\n\n cudaCheckError();\n mu0 = 1.0 / mu0;\n mu1 = 4.0 * mu0; // default set 8;\n smu0 = sqrt(mu0);\n smu1 = sqrt(mu1);\n k[1] = (mu0 + mu1) / 2.0;\n k[2] = (smu0 + smu1) * (smu0 + smu1) / 2.0;\n k[3] = mu0 * mu1;\n k[4] = 2.0 * k[3] / k[2]; // 4.0*mu0*mu1/(sqrt(mu0)+sqrt(mu1))/(sqrt(mu0)+sqrt(mu1));\n k[5] = (mu1 - 2.0 * smu0 * smu1 + mu0) / (mu1 + 2.0 * smu0 * smu1 + mu0); // square of (sqrt(kappa)-1)/(sqrt(kappa)+1);\n\n if (this->m_explicit_A->get_block_size() == 1)\n {\n MatrixCusp wA((Matrix *) &*this->m_explicit_A);\n ValueTypeA rho = cusp::detail::ritz_spectral_radius_symmetric(wA, 8);\n cudaCheckError();\n cusp::array1d coeffs;\n cusp::relaxation::detail::chebyshev_polynomial_coefficients(rho, coeffs);\n cudaCheckError();\n poly = cusp::relaxation::polynomial::memory_space > (wA, coeffs);\n cudaCheckError();\n }\n\n R.resize(N);\n V0.resize(N);\n V.resize(N);\n Rbar.resize(N);\n Sn.resize(N);\n Dinv.resize(N);\n R.set_block_dimy(this->m_explicit_A->get_block_dimy());\n R.set_block_dimx(1);\n V0.set_block_dimy(this->m_explicit_A->get_block_dimy());\n V0.set_block_dimx(1);\n V.set_block_dimy(this->m_explicit_A->get_block_dimy());\n V.set_block_dimx(1);\n Rbar.set_block_dimy(this->m_explicit_A->get_block_dimy());\n Rbar.set_block_dimx(1);\n Sn.set_block_dimy(this->m_explicit_A->get_block_dimy());\n Sn.set_block_dimx(1);\n Dinv.set_block_dimy(this->m_explicit_A->get_block_dimy());\n Dinv.set_block_dimx(1);\n ValueTypeB *Dinv_ptr = Dinv.raw();\n get_diaginv <<< num_blocks, threads_per_block>>>(A_dia_ptr, A_nonzero_values_ptr, N, this->m_explicit_A->get_block_dimy(), Dinv_ptr);\n cudaCheckError();\n}\n\n//\ntemplate\nvoid\nPolynomialSolverBase::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nPolynomialSolverBase::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n smooth_common_sqblocks( *this->m_explicit_A, b, x );\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nPolynomialSolverBase::solve_finalize( VVector &b, VVector &x )\n{\n}\n\ntemplate\nstruct poly_smooth\n{\n static void poly_postsmooth(const Matrix &A, const Vector &B, Vector &C, cusp::relaxation::polynomial &poly)\n {\n FatalError(\"Mixed precision is not supported for scalar matrix type\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n static void poly_presmooth(const Matrix &A, const Vector &B, Vector &C, cusp::relaxation::polynomial &poly)\n {\n FatalError(\"Mixed precision is not supported for scalar matrix type\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n};\n\ntemplate\nstruct poly_smooth>\n{\n static void poly_postsmooth(const Matrix > &A, const Vector > &B, Vector > &C, cusp::relaxation::polynomial::Type> &poly)\n {\n MatrixCusp, cusp::csr_format> wA((Matrix > *) &A);\n poly.postsmooth(wA, B, C);\n cudaCheckError();\n }\n\n static void poly_presmooth(const Matrix > &A, const Vector > &B, Vector > &C, cusp::relaxation::polynomial::Type> &poly)\n {\n MatrixCusp, cusp::csr_format> wA((Matrix > *) &A);\n poly.presmooth(wA, B, C);\n cudaCheckError();\n }\n};\n\ntemplate\nstruct poly_smooth>\n{\n static void poly_postsmooth(const Matrix > &A, const Vector > &B, Vector > &C, cusp::relaxation::polynomial::Type> &poly)\n {\n MatrixCusp, cusp::csr_format> wA((Matrix > *) &A);\n poly.postsmooth(wA, B, C);\n cudaCheckError();\n }\n\n static void poly_presmooth(const Matrix > &A, const Vector > &B, Vector > &C, cusp::relaxation::polynomial::Type> &poly)\n {\n MatrixCusp, cusp::csr_format> wA((Matrix > *) &A);\n poly.presmooth(wA, B, C);\n cudaCheckError();\n }\n};\n\ntemplate \nvoid PolynomialSolver >::smooth_1x1(const Matrix_h &A, const VVector &B, VVector &C)\n{\n FatalError(\"Unsupported on host for Polynomial smoother\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\ntemplate \nvoid PolynomialSolver >::smooth_1x1(const Matrix_d &A, const VVector &B, VVector &C)\n{\n if (A.hasProps(DIAG))\n {\n FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (A.get_block_size() != 1)\n {\n FatalError(\"Unsupported block size for PolynomialSolver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n poly_smooth>::poly_postsmooth(A, B, C, this->poly);\n}\n\ntemplate \nvoid PolynomialSolver >::smooth_common_sqblocks(Matrix_h &A, const VVector &b, VVector &x)\n{\n FatalError(\"Unsupported on host for Polynomial smoother\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\ntemplate \nvoid PolynomialSolver >::smooth_common_sqblocks(Matrix_d &A, const VVector &b, VVector &x)\n{\n //0, get residule : R = b - Ax\n multiply( A, x, this->R );\n cusp::blas::axpby( b, this->R, this->R, ValueTypeB(1), ValueTypeB(-1) );\n //1, set up rbar\n cusp::blas::xmy(this->Dinv, this->R, this->Rbar);\n //2, set up V and V0\n multiply( A, this->Rbar, this->V );\n cusp::blas::xmy(this->Dinv, this->V, this->V);\n this->V0 = this->Rbar;\n cusp::blas::scal(this->V0, this->k[1]);\n cusp::blas::axpby(this->Rbar, this->V, this->V, this->k[2], this->k[3]);\n\n //3, iterate to get v_{j+1}\n for (int i = 0; i < this->ndeg0; i++)\n {\n multiply( A, this->V, this->Rbar );\n cusp::blas::axpby( this->R, this->Rbar, this->Rbar, ValueTypeB(1), ValueTypeB(-1) );\n cusp::blas::xmy(this->Dinv, this->Rbar, this->Rbar);\n cusp::blas::axpbypcz(this->Rbar, this->V, this->V0, this->Sn,\n this->k[4], this->k[5] + 1.0, -1.0 * this->k[5]);\n // V0 = V\n this->V0 = this->V;\n // V = V+Sn\n cusp::blas::copy(this->Sn, this->V);\n }\n\n //4, update solution\n cusp::blas::axpy(this->V, x, ValueTypeB(1));\n}\n\n\ntemplate \nvoid PolynomialSolver >::smooth_with_0_initial_guess_1x1(const Matrix_h &A, const VVector &b, VVector &x)\n{\n FatalError(\"Unsupported on host for Polynomial smoother\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\ntemplate \nvoid PolynomialSolver >::smooth_with_0_initial_guess_1x1(const Matrix_d &A, const VVector &b, VVector &x)\n{\n if (A.hasProps(DIAG))\n {\n FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (A.get_block_size() != 1)\n {\n FatalError(\"Unsupported block size for PolynomialSolver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n poly_smooth>::poly_presmooth(A, b, x, this->poly);\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class PolynomialSolverBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class PolynomialSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n\n\n\n} // namespace polynomial_solver\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/solvers/polynomial_solver.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace amgx\n{\n\n// -----------\n// Kernels\n// -----------\n\ntemplate \n__global__ void jacobi_l1_postsmooth(\n int n, ValueTypeB omega, ValueTypeB* x, ValueTypeA* d_in, ValueTypeB* b_in, ValueTypeB* y_in)\n{\n int gid = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int i = gid; i < n; i += blockDim.x * gridDim.x)\n {\n ValueTypeA d = d_in[i];\n ValueTypeB b = b_in[i];\n ValueTypeB y = y_in[i];\n d = ValueTypeA( 1 ) / (isNotCloseToZero(d) ? d : epsilon(d));\n b -= y;\n b *= omega;\n x[i] = b * d + x[i];\n }\n}\n\ntemplate \n__global__ void jacobi_l1_postsmooth_zero(\n int n, ValueTypeB omega, ValueTypeB* x, ValueTypeA* d_in, ValueTypeB* b_in)\n{\n int gid = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int i = gid; i < n; i += blockDim.x * gridDim.x)\n {\n ValueTypeA d = d_in[i];\n ValueTypeB b = b_in[i];\n x[i] = omega * b / (isNotCloseToZero(d) ? d : epsilon(d));\n }\n}\n\ntemplate\n__global__ void compute_d_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n ValueTypeA *d)\n{\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n ValueTypeB d_ = 0;\n IndexType row_start = Ap[ridx];\n IndexType row_end = Ap[ridx + 1];\n // check if we need +ve or -ve d\n bool is_npd = false;\n\n for (int j = row_start; j < row_end; j++)\n {\n ValueTypeB Aij = Ax[j];\n\n if (Aj[j] == ridx && Aij < 0.) { is_npd = true; }\n\n //if not the diagonal then compute the absolute value\n //if(Aj[j]!=ridxa) JE: must include diagonal here or else we can get cancellation and Nan results\n Aij = fabs(Aij);\n d_ += Aij;\n }\n\n // set sign of L1-norm appropriately\n d[ridx] = (is_npd) ? -d_ : d_;\n }\n}\n\ntemplate\n__global__ void compute_d_4x4_kernel(const IndexType num_rows,\n const IndexType *A_row_offsets,\n const IndexType *A_col_indices,\n const IndexType *A_dia_indices,\n const ValueTypeA *A_nonzero_values,\n ValueTypeA *d)\n{\n const int nWarps = CtaSize / 32; // Number of warps per block\n const int warpId = utils::warp_id();\n const int laneId = utils::lane_id();\n // Lane ID in the 2 16-wide segments.\n const int lane_id_div_16 = laneId / 16;\n const int lane_id_mod_16 = laneId % 16;\n // Coordinates inside a 4x4 block of the matrix.\n const int idx_i = lane_id_mod_16 / 4;\n const int idx_j = lane_id_mod_16 % 4;\n __shared__ volatile ValueTypeA s_diag[nWarps][32];\n int aRowId = blockIdx.x * nWarps + warpId;\n\n while (aRowId < num_rows )\n {\n int aColBeg = A_row_offsets[aRowId];\n int aColEnd = A_row_offsets[aRowId + 1];\n s_diag[warpId][laneId] = 0.;\n\n //TODO: Should prefetch the column indices\n for (int i = aColBeg; i < aColEnd; i += 2)\n {\n const int my_i = i + lane_id_div_16;\n int aColId = my_i < aColEnd ? A_col_indices[my_i] : int (-1);\n s_diag[warpId][laneId] += (aRowId != aColId && aColId != -1) ? fabs(A_nonzero_values[16 * my_i + lane_id_mod_16]) : 0.;\n }\n\n // Threads 0 to 15 compute inverse\n if (lane_id_div_16 == 0)\n {\n s_diag[warpId][laneId] += s_diag[warpId][laneId + 16];\n s_diag[warpId][laneId] += fabs( A_nonzero_values[A_dia_indices[aRowId] * 16 + laneId]);\n const int offset = 16 * aRowId;\n compute_block_inverse_row_major4x4_formula2\n (s_diag[warpId], 16 * lane_id_div_16, offset, idx_i, idx_j, d);\n }\n\n aRowId += nWarps * gridDim.x;\n }\n}\n\n\ntemplate\n__global__ void jacobi_smooth_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const ValueTypeA *d,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB *xout,\n const IndexType row_offset)\n\n{\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = row_offset + tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n IndexType row_start = Ap[ridx];\n IndexType row_end = Ap[ridx + 1];\n ValueTypeB Axi = 0.0;\n\n for (int j = row_start; j < row_end; j++)\n {\n Axi += Ax[j] * x[Aj[j]];\n }\n\n xout[ridx] = x[ridx] + (b[ridx] - Axi) / ( isNotCloseToZero( d[ridx]) ? d[ridx] : epsilon(d[ridx]) );\n }\n}\n\n\ntemplate\n__global__ void jacobi_smooth_4x4_kernel(const IndexType num_rows,\n const IndexType *row_offsets,\n const IndexType *column_indices,\n const ValueTypeA *nonzero_values,\n const ValueTypeA *d,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB *xout,\n const IndexType row_offset)\n\n{\n int warp_id = threadIdx.x / 32;\n int warp_thread_id = threadIdx.x & 31;\n\n // padding row blocks to fit in a single warp\n if ( warp_thread_id >= blockrows_per_warp * bsize ) { return; }\n\n // new thread id with padding\n int tid = warp_id * blockrows_per_warp * bsize + warp_thread_id;\n // Here we use one thread per row (not block row)\n int cta_blockrow_id = (tid) / bsize;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int vec_entry_index = tid - cta_blockrow_id * bsize;\n volatile __shared__ ValueTypeB s_xtemp[ bsize * blockrows_per_cta];\n int offset, s_offset, i;\n ValueTypeB bmAx, temp[bsize];\n\n while (blockrow_id < num_rows && cta_blockrow_id < blockrows_per_cta)\n {\n i = blockrow_id;\n // Load RHS and x\n // groups of 4 threads load 32-byte, through texture, no wasted bandwidth\n offset = i * bsize + vec_entry_index;\n bmAx = b[offset];\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n offset = jcol * bsize + vec_entry_index;\n s_xtemp[tid] = x[offset] * (jcol != i);\n //for some reason we do better in perf if we write to volatile shared memory and then read this back...\n //s_xtemp[threadIdx.x] += delta[offset]*(row_colors[jcol] < current_color && current_color!= 0);\n // Load nonzero_values\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n // Do matrix multiply\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset++];\n }\n }\n\n // Each thread stores bmAx in xtemp\n s_xtemp[tid] = bmAx;\n bmAx = 0.;\n // Load Dinv\n offset = i * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(d + offset, temp);\n // Do matrix-vector multiply\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx += temp[m] * s_xtemp[s_offset++];\n }\n\n xout[i * bsize + vec_entry_index] = bmAx;\n blockrow_id += blockrows_per_cta * gridDim.x;\n }\n}\n\n\n\ntemplate\n__global__ void jacobi_smooth_with_0_initial_guess_kernel(const IndexType num_rows,\n const ValueTypeA *d,\n const ValueTypeB *b,\n ValueTypeB *x,\n ValueTypeB weight,\n const IndexType row_offset)\n\n{\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = row_offset + tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n x[ridx] = weight * b[ridx] / ( isNotCloseToZero( d[ridx]) ? d[ridx] : epsilon(d[ridx]) );\n }\n}\n\n// -----------------\n// Methods\n// -----------------\n\n// Constructor\ntemplate\nJacobiL1Solver_Base::JacobiL1Solver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope), m_d(0)\n{\n weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n\n if (weight == 0)\n {\n weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Block Jacobi smoother\\n\");;\n }\n}\n\n// Destructor\ntemplate\nJacobiL1Solver_Base::~JacobiL1Solver_Base()\n{\n}\n\n// Solver setup\ntemplate\nvoid\nJacobiL1Solver_Base::solver_setup(bool reuse_matrix_structure)\n{\n m_explicit_A = dynamic_cast*>(Base::m_A);\n\n if (!m_explicit_A)\n {\n FatalError(\"JacobiL1Solver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n compute_d( *this->m_explicit_A );\n int N = this->m_explicit_A->get_num_cols() * this->m_explicit_A->get_block_dimy();\n this->y_tmp.resize(N);\n}\n\ntemplate\nvoid JacobiL1Solver_Base::compute_d( Matrix &A)\n{\n this->m_d.resize(A.get_num_rows()*A.get_block_size());\n ViewType oldView = A.currentView();\n A.setView(this->m_explicit_A->getViewExterior());\n\n if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n compute_d_1x1(A);\n }\n else if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n compute_d_4x4(A);\n }\n else\n {\n FatalError(\"Unsupported block size for JacobiL1Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n}\n\n\ntemplate \nvoid JacobiL1Solver >::compute_d_1x1(const Matrix_h &A)\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeB d_ = 0;\n\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n ValueTypeB Aij = A.values[j];\n //if not the diagonal then compute the absolute value\n //if(A.col_indices[j]!=i) JE have to include diagonal or risk cancellation if diagonal is negative\n Aij = fabs(Aij);\n d_ += Aij;\n }\n\n this->m_d[i] = d_;\n }\n}\n\n\ntemplate \nvoid JacobiL1Solver >::compute_d_4x4(const Matrix_h &A)\n{\n FatalError(\"4x4 block matrices not supported on host\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\n\ntemplate \nvoid JacobiL1Solver >::compute_d_1x1(const Matrix_d &A)\n{\n//DIAG: starnge issues trying to add DIAG property handling\n// now leaving !DIAG only\n if (A.hasProps(DIAG))\n {\n FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((ValueTypeB)A.get_num_rows() / (ValueTypeB)THREADS_PER_BLOCK));\n\n if (A.get_num_rows() > 0)\n {\n compute_d_kernel <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK>>>\n ((int)A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->m_d.raw());\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid JacobiL1Solver >::compute_d_4x4(const Matrix_d &A)\n{\n // now leaving !DIAG only\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n const int cta_size = 128;\n const int nWarps = cta_size / 32;\n int grid_size = std::min( 1024, ( A.get_num_rows() + nWarps - 1 ) / nWarps );\n compute_d_4x4_kernel <<< grid_size, cta_size >>>\n (A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->m_d.raw());\n cudaCheckError();\n}\n\n\n\n\n\n//\ntemplate\nvoid\nJacobiL1Solver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nJacobiL1Solver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n if (xIsZero) { x.dirtybit = 0; }\n\n ViewType oldView = this->m_explicit_A->currentView();\n this->m_explicit_A->setViewExterior();\n\n if (this->m_explicit_A->get_block_dimx() == 1 && this->m_explicit_A->get_block_dimy() == 1)\n {\n if (xIsZero)\n {\n smooth_with_0_initial_guess_1x1(*this->m_explicit_A, b, x, this->m_explicit_A->getViewExterior());\n }\n else\n {\n smooth_1x1(*this->m_explicit_A, b, x, this->m_explicit_A->getViewExterior(), false);\n }\n }\n else\n {\n FatalError(\"Unsupported block size for JacobiL1_Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n x.dirtybit = 1;\n this->m_explicit_A->setView(oldView);\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nJacobiL1Solver_Base::solve_finalize( VVector &b, VVector &x )\n{\n}\n\n\n\ntemplate \nvoid JacobiL1Solver >::smooth_1x1(Matrix_h &A, VVector &b, VVector &x, ViewType separation_flags, bool latency_hiding)\n{\n VVector newx((int)x.size());\n\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeB Axi = 0.0;\n\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n Axi += A.values[j] * x[A.col_indices[j]];\n }\n\n ValueTypeA d = this->m_d[i];\n newx[i] = x[i] + (b[i] - Axi) / ( isNotCloseToZero( d) ? d : epsilon(d) );\n }\n\n x.swap(newx);\n}\n\n\ntemplate \nvoid JacobiL1Solver >::smooth_4x4(Matrix_h &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n FatalError(\"4x4 block size not supported on host\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid JacobiL1Solver >::smooth_with_0_initial_guess_1x1(Matrix_h &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeA d = this->m_d[i];\n x[i] = b[i] / ( isNotCloseToZero( d) ? d : epsilon(d) );\n }\n}\n\ntemplate \nvoid JacobiL1Solver >::smooth_1x1(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags, bool latency_hiding)\n{\n this->y_tmp.set_block_dimx(b.get_block_dimx());\n this->y_tmp.set_block_dimy(b.get_block_dimy());\n\n multiply(A, x, this->y_tmp, A.getViewExterior());\n\n int offset, num_rows;\n A.getOffsetAndSizeForView(A.getViewExterior(), &offset, &num_rows);\n\n int nthreads_per_block = 128;\n int n = num_rows - offset;\n int nblocks = n / nthreads_per_block + 1;\n jacobi_l1_postsmooth<<>>(n, this->weight, x.raw() + offset, this->m_d.raw() + offset, b.raw() + offset, this->y_tmp.raw() + offset);\n\n cudaCheckError();\n}\n\n\ntemplate \nvoid JacobiL1Solver >::smooth_4x4(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n VVector xout(x.size());\n ValueTypeB *x_ptr = x.raw();\n ValueTypeB *xout_ptr = xout.raw();\n IndexType num_rows = A.get_num_rows();\n IndexType offset = 0;\n\n //It is only safe to swap these vectors if there is no halo exchange in process\n if (separation_flags != this->m_explicit_A->getViewExterior())\n {\n amgx::thrust::copy(x.begin(), x.end(), xout.begin()); //TODO: only interior+bndry part\n cudaCheckError();\n x_ptr = xout.raw();\n xout_ptr = x.raw();\n\n if (separation_flags & INTERIOR)\n {\n }\n else\n {\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n }\n }\n\n const int threads_per_block = 128;\n const int blockrows_per_warp = 32 / 4;\n const int blockrows_per_cta = (threads_per_block / 32) * blockrows_per_warp;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n jacobi_smooth_4x4_kernel <<< num_blocks, threads_per_block >>>\n ((int)A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->m_d.raw(),\n b.raw(),\n x_ptr,\n xout_ptr, offset);\n cudaCheckError();\n\n if ((separation_flags != this->m_explicit_A->getViewExterior())/* && (separation_flags & INTERIOR)*/) {}\n else { x.swap(xout); }\n}\n\ntemplate \nvoid JacobiL1Solver >::smooth_with_0_initial_guess_1x1(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n ViewType oldView = A.currentView();\n // Process all rows\n A.setViewExterior();\n\n int offset, num_rows;\n A.getOffsetAndSizeForView(A.getViewExterior(), &offset, &num_rows);\n\n int nthreads_per_block = 128;\n int n = num_rows - offset;\n int nblocks = n / nthreads_per_block + 1;\n jacobi_l1_postsmooth_zero<<>>(n, this->weight, x.raw() + offset, this->m_d.raw() + offset, b.raw() + offset);\n cudaCheckError();\n\n A.setView(oldView);\n cudaCheckError();\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class JacobiL1Solver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class JacobiL1Solver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace amgx\n{\n\n// -----------\n// Kernels\n// -----------\n\ntemplate \n__global__ void jacobi_l1_postsmooth(\n int n, ValueTypeB omega, ValueTypeB* x, ValueTypeA* d_in, ValueTypeB* b_in, ValueTypeB* y_in)\n{\n int gid = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int i = gid; i < n; i += blockDim.x * gridDim.x)\n {\n ValueTypeA d = d_in[i];\n ValueTypeB b = b_in[i];\n ValueTypeB y = y_in[i];\n d = ValueTypeA( 1 ) / (isNotCloseToZero(d) ? d : epsilon(d));\n b -= y;\n b *= omega;\n x[i] = b * d + x[i];\n }\n}\n\ntemplate \n__global__ void jacobi_l1_postsmooth_zero(\n int n, ValueTypeB omega, ValueTypeB* x, ValueTypeA* d_in, ValueTypeB* b_in)\n{\n int gid = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int i = gid; i < n; i += blockDim.x * gridDim.x)\n {\n ValueTypeA d = d_in[i];\n ValueTypeB b = b_in[i];\n x[i] = omega * b / (isNotCloseToZero(d) ? d : epsilon(d));\n }\n}\n\ntemplate\n__global__ void compute_d_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n ValueTypeA *d)\n{\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n ValueTypeB d_ = 0;\n IndexType row_start = Ap[ridx];\n IndexType row_end = Ap[ridx + 1];\n // check if we need +ve or -ve d\n bool is_npd = false;\n\n for (int j = row_start; j < row_end; j++)\n {\n ValueTypeB Aij = Ax[j];\n\n if (Aj[j] == ridx && Aij < 0.) { is_npd = true; }\n\n //if not the diagonal then compute the absolute value\n //if(Aj[j]!=ridxa) JE: must include diagonal here or else we can get cancellation and Nan results\n Aij = fabs(Aij);\n d_ += Aij;\n }\n\n // set sign of L1-norm appropriately\n d[ridx] = (is_npd) ? -d_ : d_;\n }\n}\n\ntemplate\n__global__ void compute_d_4x4_kernel(const IndexType num_rows,\n const IndexType *A_row_offsets,\n const IndexType *A_col_indices,\n const IndexType *A_dia_indices,\n const ValueTypeA *A_nonzero_values,\n ValueTypeA *d)\n{\n const int nWarps = CtaSize / 32; // Number of warps per block\n const int warpId = utils::warp_id();\n const int laneId = utils::lane_id();\n // Lane ID in the 2 16-wide segments.\n const int lane_id_div_16 = laneId / 16;\n const int lane_id_mod_16 = laneId % 16;\n // Coordinates inside a 4x4 block of the matrix.\n const int idx_i = lane_id_mod_16 / 4;\n const int idx_j = lane_id_mod_16 % 4;\n __shared__ volatile ValueTypeA s_diag[nWarps][32];\n int aRowId = blockIdx.x * nWarps + warpId;\n\n while (aRowId < num_rows )\n {\n int aColBeg = A_row_offsets[aRowId];\n int aColEnd = A_row_offsets[aRowId + 1];\n s_diag[warpId][laneId] = 0.;\n\n //TODO: Should prefetch the column indices\n for (int i = aColBeg; i < aColEnd; i += 2)\n {\n const int my_i = i + lane_id_div_16;\n int aColId = my_i < aColEnd ? A_col_indices[my_i] : int (-1);\n s_diag[warpId][laneId] += (aRowId != aColId && aColId != -1) ? fabs(A_nonzero_values[16 * my_i + lane_id_mod_16]) : 0.;\n }\n\n // Threads 0 to 15 compute inverse\n if (lane_id_div_16 == 0)\n {\n s_diag[warpId][laneId] += s_diag[warpId][laneId + 16];\n s_diag[warpId][laneId] += fabs( A_nonzero_values[A_dia_indices[aRowId] * 16 + laneId]);\n const int offset = 16 * aRowId;\n compute_block_inverse_row_major4x4_formula2\n (s_diag[warpId], 16 * lane_id_div_16, offset, idx_i, idx_j, d);\n }\n\n aRowId += nWarps * gridDim.x;\n }\n}\n\n\ntemplate\n__global__ void jacobi_smooth_kernel(const IndexType num_rows,\n const IndexType *Ap,\n const IndexType *Aj,\n const ValueTypeA *Ax,\n const ValueTypeA *d,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB *xout,\n const IndexType row_offset)\n\n{\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = row_offset + tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n IndexType row_start = Ap[ridx];\n IndexType row_end = Ap[ridx + 1];\n ValueTypeB Axi = 0.0;\n\n for (int j = row_start; j < row_end; j++)\n {\n Axi += Ax[j] * x[Aj[j]];\n }\n\n xout[ridx] = x[ridx] + (b[ridx] - Axi) / ( isNotCloseToZero( d[ridx]) ? d[ridx] : epsilon(d[ridx]) );\n }\n}\n\n\ntemplate\n__global__ void jacobi_smooth_4x4_kernel(const IndexType num_rows,\n const IndexType *row_offsets,\n const IndexType *column_indices,\n const ValueTypeA *nonzero_values,\n const ValueTypeA *d,\n const ValueTypeB *b,\n const ValueTypeB *x,\n ValueTypeB *xout,\n const IndexType row_offset)\n\n{\n int warp_id = threadIdx.x / 32;\n int warp_thread_id = threadIdx.x & 31;\n\n // padding row blocks to fit in a single warp\n if ( warp_thread_id >= blockrows_per_warp * bsize ) { return; }\n\n // new thread id with padding\n int tid = warp_id * blockrows_per_warp * bsize + warp_thread_id;\n // Here we use one thread per row (not block row)\n int cta_blockrow_id = (tid) / bsize;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int vec_entry_index = tid - cta_blockrow_id * bsize;\n volatile __shared__ ValueTypeB s_xtemp[ bsize * blockrows_per_cta];\n int offset, s_offset, i;\n ValueTypeB bmAx, temp[bsize];\n\n while (blockrow_id < num_rows && cta_blockrow_id < blockrows_per_cta)\n {\n i = blockrow_id;\n // Load RHS and x\n // groups of 4 threads load 32-byte, through texture, no wasted bandwidth\n offset = i * bsize + vec_entry_index;\n bmAx = b[offset];\n // Contribution from each nonzero column\n int jmax = row_offsets[i + 1];\n\n for (int jind = row_offsets[i]; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n offset = jcol * bsize + vec_entry_index;\n s_xtemp[tid] = x[offset] * (jcol != i);\n //for some reason we do better in perf if we write to volatile shared memory and then read this back...\n //s_xtemp[threadIdx.x] += delta[offset]*(row_colors[jcol] < current_color && current_color!= 0);\n // Load nonzero_values\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n // Do matrix multiply\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_xtemp[s_offset++];\n }\n }\n\n // Each thread stores bmAx in xtemp\n s_xtemp[tid] = bmAx;\n bmAx = 0.;\n // Load Dinv\n offset = i * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(d + offset, temp);\n // Do matrix-vector multiply\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n bmAx += temp[m] * s_xtemp[s_offset++];\n }\n\n xout[i * bsize + vec_entry_index] = bmAx;\n blockrow_id += blockrows_per_cta * gridDim.x;\n }\n}\n\n\n\ntemplate\n__global__ void jacobi_smooth_with_0_initial_guess_kernel(const IndexType num_rows,\n const ValueTypeA *d,\n const ValueTypeB *b,\n ValueTypeB *x,\n ValueTypeB weight,\n const IndexType row_offset)\n\n{\n IndexType tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int ridx = row_offset + tidx; ridx < num_rows; ridx += blockDim.x * gridDim.x)\n {\n x[ridx] = weight * b[ridx] / ( isNotCloseToZero( d[ridx]) ? d[ridx] : epsilon(d[ridx]) );\n }\n}\n\n// -----------------\n// Methods\n// -----------------\n\n// Constructor\ntemplate\nJacobiL1Solver_Base::JacobiL1Solver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope), m_d(0)\n{\n weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n\n if (weight == 0)\n {\n weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Block Jacobi smoother\\n\");;\n }\n}\n\n// Destructor\ntemplate\nJacobiL1Solver_Base::~JacobiL1Solver_Base()\n{\n}\n\n// Solver setup\ntemplate\nvoid\nJacobiL1Solver_Base::solver_setup(bool reuse_matrix_structure)\n{\n m_explicit_A = dynamic_cast*>(Base::m_A);\n\n if (!m_explicit_A)\n {\n FatalError(\"JacobiL1Solver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n compute_d( *this->m_explicit_A );\n int N = this->m_explicit_A->get_num_cols() * this->m_explicit_A->get_block_dimy();\n this->y_tmp.resize(N);\n}\n\ntemplate\nvoid JacobiL1Solver_Base::compute_d( Matrix &A)\n{\n this->m_d.resize(A.get_num_rows()*A.get_block_size());\n ViewType oldView = A.currentView();\n A.setView(this->m_explicit_A->getViewExterior());\n\n if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n compute_d_1x1(A);\n }\n else if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n compute_d_4x4(A);\n }\n else\n {\n FatalError(\"Unsupported block size for JacobiL1Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n}\n\n\ntemplate \nvoid JacobiL1Solver >::compute_d_1x1(const Matrix_h &A)\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeB d_ = 0;\n\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n ValueTypeB Aij = A.values[j];\n //if not the diagonal then compute the absolute value\n //if(A.col_indices[j]!=i) JE have to include diagonal or risk cancellation if diagonal is negative\n Aij = fabs(Aij);\n d_ += Aij;\n }\n\n this->m_d[i] = d_;\n }\n}\n\n\ntemplate \nvoid JacobiL1Solver >::compute_d_4x4(const Matrix_h &A)\n{\n FatalError(\"4x4 block matrices not supported on host\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\n\ntemplate \nvoid JacobiL1Solver >::compute_d_1x1(const Matrix_d &A)\n{\n//DIAG: starnge issues trying to add DIAG property handling\n// now leaving !DIAG only\n if (A.hasProps(DIAG))\n {\n FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((ValueTypeB)A.get_num_rows() / (ValueTypeB)THREADS_PER_BLOCK));\n\n if (A.get_num_rows() > 0)\n {\n compute_d_kernel <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK>>>\n ((int)A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->m_d.raw());\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid JacobiL1Solver >::compute_d_4x4(const Matrix_d &A)\n{\n // now leaving !DIAG only\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n const int cta_size = 128;\n const int nWarps = cta_size / 32;\n int grid_size = std::min( 1024, ( A.get_num_rows() + nWarps - 1 ) / nWarps );\n compute_d_4x4_kernel <<< grid_size, cta_size >>>\n (A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->m_d.raw());\n cudaCheckError();\n}\n\n\n\n\n\n//\ntemplate\nvoid\nJacobiL1Solver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nJacobiL1Solver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n if (xIsZero) { x.dirtybit = 0; }\n\n ViewType oldView = this->m_explicit_A->currentView();\n this->m_explicit_A->setViewExterior();\n\n if (this->m_explicit_A->get_block_dimx() == 1 && this->m_explicit_A->get_block_dimy() == 1)\n {\n if (xIsZero)\n {\n smooth_with_0_initial_guess_1x1(*this->m_explicit_A, b, x, this->m_explicit_A->getViewExterior());\n }\n else\n {\n smooth_1x1(*this->m_explicit_A, b, x, this->m_explicit_A->getViewExterior(), false);\n }\n }\n else\n {\n FatalError(\"Unsupported block size for JacobiL1_Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n x.dirtybit = 1;\n this->m_explicit_A->setView(oldView);\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nJacobiL1Solver_Base::solve_finalize( VVector &b, VVector &x )\n{\n}\n\n\n\ntemplate \nvoid JacobiL1Solver >::smooth_1x1(Matrix_h &A, VVector &b, VVector &x, ViewType separation_flags, bool latency_hiding)\n{\n VVector newx((int)x.size());\n\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeB Axi = 0.0;\n\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n Axi += A.values[j] * x[A.col_indices[j]];\n }\n\n ValueTypeA d = this->m_d[i];\n newx[i] = x[i] + (b[i] - Axi) / ( isNotCloseToZero( d) ? d : epsilon(d) );\n }\n\n x.swap(newx);\n}\n\n\ntemplate \nvoid JacobiL1Solver >::smooth_4x4(Matrix_h &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n FatalError(\"4x4 block size not supported on host\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate \nvoid JacobiL1Solver >::smooth_with_0_initial_guess_1x1(Matrix_h &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeA d = this->m_d[i];\n x[i] = b[i] / ( isNotCloseToZero( d) ? d : epsilon(d) );\n }\n}\n\ntemplate \nvoid JacobiL1Solver >::smooth_1x1(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags, bool latency_hiding)\n{\n this->y_tmp.set_block_dimx(b.get_block_dimx());\n this->y_tmp.set_block_dimy(b.get_block_dimy());\n\n multiply(A, x, this->y_tmp, A.getViewExterior());\n\n int offset, num_rows;\n A.getOffsetAndSizeForView(A.getViewExterior(), &offset, &num_rows);\n\n int nthreads_per_block = 128;\n int n = num_rows - offset;\n int nblocks = n / nthreads_per_block + 1;\n jacobi_l1_postsmooth<<>>(n, this->weight, x.raw() + offset, this->m_d.raw() + offset, b.raw() + offset, this->y_tmp.raw() + offset);\n\n cudaCheckError();\n}\n\n\ntemplate \nvoid JacobiL1Solver >::smooth_4x4(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n VVector xout(x.size());\n ValueTypeB *x_ptr = x.raw();\n ValueTypeB *xout_ptr = xout.raw();\n IndexType num_rows = A.get_num_rows();\n IndexType offset = 0;\n\n //It is only safe to swap these vectors if there is no halo exchange in process\n if (separation_flags != this->m_explicit_A->getViewExterior())\n {\n amgx::thrust::copy(x.begin(), x.end(), xout.begin()); //TODO: only interior+bndry part\n cudaCheckError();\n x_ptr = xout.raw();\n xout_ptr = x.raw();\n\n if (separation_flags & INTERIOR)\n {\n }\n else\n {\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n }\n }\n\n const int threads_per_block = 128;\n const int blockrows_per_warp = 32 / 4;\n const int blockrows_per_cta = (threads_per_block / 32) * blockrows_per_warp;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n jacobi_smooth_4x4_kernel <<< num_blocks, threads_per_block >>>\n ((int)A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n this->m_d.raw(),\n b.raw(),\n x_ptr,\n xout_ptr, offset);\n cudaCheckError();\n\n if ((separation_flags != this->m_explicit_A->getViewExterior())/* && (separation_flags & INTERIOR)*/) {}\n else { x.swap(xout); }\n}\n\ntemplate \nvoid JacobiL1Solver >::smooth_with_0_initial_guess_1x1(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n ViewType oldView = A.currentView();\n // Process all rows\n A.setViewExterior();\n\n int offset, num_rows;\n A.getOffsetAndSizeForView(A.getViewExterior(), &offset, &num_rows);\n\n int nthreads_per_block = 128;\n int n = num_rows - offset;\n int nblocks = n / nthreads_per_block + 1;\n jacobi_l1_postsmooth_zero<<>>(n, this->weight, x.raw() + offset, this->m_d.raw() + offset, b.raw() + offset);\n cudaCheckError();\n\n A.setView(oldView);\n cudaCheckError();\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class JacobiL1Solver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class JacobiL1Solver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/solvers/jacobi_l1_solver.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#define _USE_MATH_DEFINES\n#include \"math.h\"\n#include \n#include \n\n#include \n\n#undef DEBUG_CHEBYSHEV_OUTPUT\n\nnamespace amgx\n{\nnamespace chebyshev_poly_smoother\n{\n\ntemplate \nstruct jacobi_presmooth_functor\n{\n ValueTypeB omega;\n jacobi_presmooth_functor( ValueTypeB omega ) : omega( omega ) {}\n __host__ __device__ ValueTypeB operator()( const ValueTypeB &b, const ValueTypeA &d ) const { return isNotCloseToZero(d) ? omega * b / d : omega * b / epsilon(d); }\n};\n\ntemplate \nstruct jacobi_postsmooth_functor\n{\n ValueTypeB omega;\n jacobi_postsmooth_functor( ValueTypeB omega ) : omega( omega ) {}\n template __host__ __device__ ValueTypeB operator( )( const Tuple &t ) const\n {\n ValueTypeB x = amgx::thrust::get<0>(t);\n ValueTypeA d = amgx::thrust::get<1>(t);\n ValueTypeB b = amgx::thrust::get<2>(t);\n ValueTypeB y = amgx::thrust::get<3>(t);\n // return x + omega * (b - y) / d.\n d = isNotCloseToZero(d) ? d : epsilon(d);\n d = ValueTypeA( 1 ) / d;\n b -= y;\n b *= omega;\n return b * d + x;\n }\n};\n\ntemplate \nstruct add_functor\n{\n __host__ __device__ ValueTypeB operator()( const ValueTypeB &x, const ValueTypeB &y )const { return x + y; }\n};\n\n\n// -----------------------------------\n// KERNELS\n// -----------------------------------\n\ntemplate \n__host__ __device__ ValueTypeB magicDampBeta(int m)\n{\n return M_PI / (4 * (ValueTypeB)m + 2);\n}\n\ntemplate \n__host__ __device__ ValueTypeB magicDamp(int n, ValueTypeB beta)\n{\n return cos(beta) * cos(beta) / (cos(beta * (2 * n + 1)) * cos(beta * (2 * n + 1)) - sin(beta) * sin(beta));\n}\n\ntemplate\n__global__\nvoid getLambdaEstimate(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, const int num_rows, ValueTypeB *out)\n\n{\n int row_id = blockDim.x * blockIdx.x + threadIdx.x;\n ValueTypeB max_sum = (ValueTypeB)0.0;\n\n while (row_id < num_rows)\n {\n ValueTypeB cur_sum = (ValueTypeB)0.0;\n\n for (int j = row_offsets[row_id]; j < row_offsets[row_id + 1]; j++)\n {\n cur_sum += abs(values[j]);\n }\n\n if (diag)\n {\n cur_sum += abs(values[dia_indices[row_id]]);\n }\n\n max_sum = max(max_sum, cur_sum);\n row_id += gridDim.x * blockDim.x;\n }\n\n out[blockDim.x * blockIdx.x + threadIdx.x] = max_sum;\n}\n\n//--------------------------------\n// Methods\n//--------------------------------\n\n// Constructor\ntemplate\nChebyshevPolySolver_Base::ChebyshevPolySolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope)\n{\n poly_order = cfg.AMG_Config::template getParameter(\"chebyshev_polynomial_order\", cfg_scope);\n poly_order = min(10, max(poly_order, 1));\n tau.resize(poly_order);\n}\n\n// Destructor\ntemplate\nChebyshevPolySolver_Base::~ChebyshevPolySolver_Base()\n{\n this->tau.resize(0);\n}\n\ntemplate\nvoid\nChebyshevPolySolver_Base::printSolverParameters() const\n{\n std::cout << \"chebyshev_polynomial_order= \" << this->poly_order << std::endl;\n}\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid ChebyshevPolySolver >::compute_eigenmax_estimate(const Matrix_d &A, ValueTypeB &lambda)\n{\n#define LAMBDA_BLOCK_SIZE 256\n VVector tsum(A.get_num_rows());\n const int threads_per_block = 256;\n const int blockrows_per_cta = threads_per_block;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n\n if (A.hasProps(DIAG))\n {\n cudaFuncSetCacheConfig(getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, true >, cudaFuncCachePreferL1);\n getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, true > <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, A.get_num_rows(), tsum.raw());\n cudaCheckError();\n }\n else\n {\n cudaFuncSetCacheConfig(getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, false >, cudaFuncCachePreferL1);\n getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, false > <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, A.get_num_rows(), tsum.raw());\n cudaCheckError();\n }\n\n lambda = *(amgx::thrust::max_element(tsum.begin(), tsum.end()));\n#ifdef DEBUG_CHEBYSHEV_OUTPUT\n printf(\"Lambda for A on row %lu is: %f\\n\", amgx::thrust::max_element(tsum.begin(), tsum.end()) - tsum.begin(), lambda);\n#endif\n}\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid ChebyshevPolySolver >::compute_eigenmax_estimate(const Matrix_h &A, ValueTypeB &lambda)\n{\n FatalError(\"compute_eigenmax_estimate is not implemented for host\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n// Solver setup\ntemplate\nvoid\nChebyshevPolySolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n Matrix *A_as_matrix = dynamic_cast*>(this->m_A);\n\n if (!A_as_matrix)\n {\n FatalError(\"ChebyshevPolySolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n ValueTypeB lambda = 0.0;\n compute_eigenmax_estimate( *A_as_matrix, lambda );\n ValueTypeB beta = magicDampBeta(poly_order);\n#ifdef DEBUG_CHEBYSHEV_OUTPUT\n int lvl = 0;\n A_as_matrix->getParameter(\"level\", lvl);\n printf(\"Tau values for level %d : \", lvl);\n#endif\n\n for (int i = 0; i < poly_order; i++)\n {\n tau[i] = magicDamp(i, beta) / lambda;\n#ifdef DEBUG_CHEBYSHEV_OUTPUT\n printf(\"%f%s\", tau[i], (i == poly_order - 1) ? \"\\n\" : \" \");\n#endif\n }\n}\n\n//\ntemplate\nvoid\nChebyshevPolySolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nChebyshevPolySolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n\n if (xIsZero) { x.dirtybit = 0; }\n\n if (!A_as_matrix->is_matrix_singleGPU())\n {\n A_as_matrix->manager->exchange_halo_async(x, x.tag);\n\n if (A_as_matrix->getViewExterior() == A_as_matrix->getViewInterior())\n {\n A_as_matrix->manager->exchange_halo_wait(x, x.tag);\n }\n }\n\n ViewType oldView = A_as_matrix->currentView();\n ViewType flags;\n bool latencyHiding = true;\n\n if (A_as_matrix->is_matrix_singleGPU() || (x.dirtybit == 0))\n {\n latencyHiding = false;\n A_as_matrix->setViewExterior();\n flags = (ViewType)(A_as_matrix->getViewInterior() | A_as_matrix->getViewExterior());\n }\n else\n {\n flags = A_as_matrix->getViewInterior();\n A_as_matrix->setViewInterior();\n }\n\n if (A_as_matrix->get_block_dimx() == 1 && A_as_matrix->get_block_dimy() == 1)\n {\n smooth_1x1(*A_as_matrix, b, x, flags);\n }\n else\n {\n FatalError(\"Unsupported block size for BlockJacobi_Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if (latencyHiding)\n {\n A_as_matrix->manager->exchange_halo_wait(x, x.tag);\n A_as_matrix->setViewExterior();\n flags = (ViewType)(~(A_as_matrix->getViewInterior()) & A_as_matrix->getViewExterior());\n\n if (flags != 0)\n {\n if (A_as_matrix->get_block_dimx() == 1 && A_as_matrix->get_block_dimy() == 1)\n {\n smooth_1x1(*A_as_matrix, b, x, flags);\n }\n else\n {\n FatalError(\"Unsupported block size for BlockJacobi_Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n }\n }\n\n x.dirtybit = 1;\n A_as_matrix->setView(oldView);\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nChebyshevPolySolver_Base::solve_finalize( VVector &b, VVector &x )\n{}\n\n\ntemplate \nvoid ChebyshevPolySolver >::smooth_1x1(Matrix_h &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n FatalError(\"chebyshev poly smoother not implemented with host format, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nstruct chebyshev_poly_functor\n{\n ValueTypeB tau;\n chebyshev_poly_functor( ValueTypeB tau ) : tau( tau ) {}\n template __host__ __device__ ValueTypeB operator( )( const Tuple &t ) const\n {\n ValueTypeB x = amgx::thrust::get<0>(t);\n ValueTypeB b = amgx::thrust::get<1>(t);\n ValueTypeB y = amgx::thrust::get<2>(t);\n return x + tau * (b - y);\n }\n};\n\ntemplate \nvoid ChebyshevPolySolver >::smooth_1x1(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n if (this->y.size() != b.size())\n {\n this->y.resize(b.size());\n this->y.tag = this->tag * 100 + 3;\n this->y.set_block_dimx(b.get_block_dimx());\n this->y.set_block_dimy(b.get_block_dimy());\n }\n\n int num_rows = A.get_num_rows();\n int offset = 0;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n bool latencyHiding = separation_flags != A.getViewIntExt();// we change view only when do latency hiding, maybe it's better to use some explicit flag attached to matrix?\n\n for (int i = 0; i < this->poly_order; i++)\n {\n this->y.dirtybit = 0;\n //y = A*x\n multiply( A, x, this->y, separation_flags );\n //x += tau_i * (b - y)\n amgx::thrust::transform( amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple( x.begin() + offset, b.begin() + offset, this->y.begin() + offset)),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple( x.begin() + num_rows, b.begin() + num_rows, this->y.begin() + num_rows)),\n x.begin() + offset,\n chebyshev_poly_functor( this->tau[i] ));\n //Cublas::axpy(num_rows, (ValueTypeB)(this->tau[i]), b.raw(), 1, x.raw(), 1);\n //Cublas::axpy(num_rows, (ValueTypeB)(-1.0*this->tau[i]), this->y.raw(), 1, x.raw(), 1);\n }\n}\n\ntemplate \nvoid ChebyshevPolySolver >::smooth_with_0_initial_guess_1x1(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n FatalError(\"chebyshev poly smoother not implemented with host format, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid ChebyshevPolySolver >::smooth_with_0_initial_guess_1x1(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n}\n\ntemplate \nvoid ChebyshevPolySolver >::smooth_BxB(Matrix_h &A, VVector &b, VVector &x, bool firstStep, ViewType separation_flags)\n{\n FatalError(\"M*M chebyshev poly smoother not implemented with host format, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid ChebyshevPolySolver >::smooth_BxB(Matrix_d &A, VVector &b, VVector &x, bool firstStep, ViewType separation_flags)\n{\n FatalError(\"M*M chebyshev poly smoother not implemented with host format, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class ChebyshevPolySolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class ChebyshevPolySolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace chebyshev_poly_smoother\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#define _USE_MATH_DEFINES\n#include \"math.h\"\n#include \n#include \n\n#include \n\n#undef DEBUG_CHEBYSHEV_OUTPUT\n\nnamespace amgx\n{\nnamespace chebyshev_poly_smoother\n{\n\ntemplate \nstruct jacobi_presmooth_functor\n{\n ValueTypeB omega;\n jacobi_presmooth_functor( ValueTypeB omega ) : omega( omega ) {}\n __host__ __device__ ValueTypeB operator()( const ValueTypeB &b, const ValueTypeA &d ) const { return isNotCloseToZero(d) ? omega * b / d : omega * b / epsilon(d); }\n};\n\ntemplate \nstruct jacobi_postsmooth_functor\n{\n ValueTypeB omega;\n jacobi_postsmooth_functor( ValueTypeB omega ) : omega( omega ) {}\n template __host__ __device__ ValueTypeB operator( )( const Tuple &t ) const\n {\n ValueTypeB x = amgx::thrust::get<0>(t);\n ValueTypeA d = amgx::thrust::get<1>(t);\n ValueTypeB b = amgx::thrust::get<2>(t);\n ValueTypeB y = amgx::thrust::get<3>(t);\n // return x + omega * (b - y) / d.\n d = isNotCloseToZero(d) ? d : epsilon(d);\n d = ValueTypeA( 1 ) / d;\n b -= y;\n b *= omega;\n return b * d + x;\n }\n};\n\ntemplate \nstruct add_functor\n{\n __host__ __device__ ValueTypeB operator()( const ValueTypeB &x, const ValueTypeB &y )const { return x + y; }\n};\n\n\n// -----------------------------------\n// KERNELS\n// -----------------------------------\n\ntemplate \n__host__ __device__ ValueTypeB magicDampBeta(int m)\n{\n return M_PI / (4 * (ValueTypeB)m + 2);\n}\n\ntemplate \n__host__ __device__ ValueTypeB magicDamp(int n, ValueTypeB beta)\n{\n return cos(beta) * cos(beta) / (cos(beta * (2 * n + 1)) * cos(beta * (2 * n + 1)) - sin(beta) * sin(beta));\n}\n\ntemplate\n__global__\nvoid getLambdaEstimate(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, const int num_rows, ValueTypeB *out)\n\n{\n int row_id = blockDim.x * blockIdx.x + threadIdx.x;\n ValueTypeB max_sum = (ValueTypeB)0.0;\n\n while (row_id < num_rows)\n {\n ValueTypeB cur_sum = (ValueTypeB)0.0;\n\n for (int j = row_offsets[row_id]; j < row_offsets[row_id + 1]; j++)\n {\n cur_sum += abs(values[j]);\n }\n\n if (diag)\n {\n cur_sum += abs(values[dia_indices[row_id]]);\n }\n\n max_sum = max(max_sum, cur_sum);\n row_id += gridDim.x * blockDim.x;\n }\n\n out[blockDim.x * blockIdx.x + threadIdx.x] = max_sum;\n}\n\n//--------------------------------\n// Methods\n//--------------------------------\n\n// Constructor\ntemplate\nChebyshevPolySolver_Base::ChebyshevPolySolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope)\n{\n poly_order = cfg.AMG_Config::template getParameter(\"chebyshev_polynomial_order\", cfg_scope);\n poly_order = min(10, max(poly_order, 1));\n tau.resize(poly_order);\n}\n\n// Destructor\ntemplate\nChebyshevPolySolver_Base::~ChebyshevPolySolver_Base()\n{\n this->tau.resize(0);\n}\n\ntemplate\nvoid\nChebyshevPolySolver_Base::printSolverParameters() const\n{\n std::cout << \"chebyshev_polynomial_order= \" << this->poly_order << std::endl;\n}\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid ChebyshevPolySolver >::compute_eigenmax_estimate(const Matrix_d &A, ValueTypeB &lambda)\n{\n#define LAMBDA_BLOCK_SIZE 256\n VVector tsum(A.get_num_rows());\n const int threads_per_block = 256;\n const int blockrows_per_cta = threads_per_block;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n\n if (A.hasProps(DIAG))\n {\n hipFuncSetCacheConfig(reinterpret_cast(getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, true >), hipFuncCachePreferL1);\n getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, true > <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, A.get_num_rows(), tsum.raw());\n cudaCheckError();\n }\n else\n {\n hipFuncSetCacheConfig(reinterpret_cast(getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, false >), hipFuncCachePreferL1);\n getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, false > <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, A.get_num_rows(), tsum.raw());\n cudaCheckError();\n }\n\n lambda = *(amgx::thrust::max_element(tsum.begin(), tsum.end()));\n#ifdef DEBUG_CHEBYSHEV_OUTPUT\n printf(\"Lambda for A on row %lu is: %f\\n\", amgx::thrust::max_element(tsum.begin(), tsum.end()) - tsum.begin(), lambda);\n#endif\n}\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid ChebyshevPolySolver >::compute_eigenmax_estimate(const Matrix_h &A, ValueTypeB &lambda)\n{\n FatalError(\"compute_eigenmax_estimate is not implemented for host\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n// Solver setup\ntemplate\nvoid\nChebyshevPolySolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n Matrix *A_as_matrix = dynamic_cast*>(this->m_A);\n\n if (!A_as_matrix)\n {\n FatalError(\"ChebyshevPolySolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n ValueTypeB lambda = 0.0;\n compute_eigenmax_estimate( *A_as_matrix, lambda );\n ValueTypeB beta = magicDampBeta(poly_order);\n#ifdef DEBUG_CHEBYSHEV_OUTPUT\n int lvl = 0;\n A_as_matrix->getParameter(\"level\", lvl);\n printf(\"Tau values for level %d : \", lvl);\n#endif\n\n for (int i = 0; i < poly_order; i++)\n {\n tau[i] = magicDamp(i, beta) / lambda;\n#ifdef DEBUG_CHEBYSHEV_OUTPUT\n printf(\"%f%s\", tau[i], (i == poly_order - 1) ? \"\\n\" : \" \");\n#endif\n }\n}\n\n//\ntemplate\nvoid\nChebyshevPolySolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nChebyshevPolySolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n\n if (xIsZero) { x.dirtybit = 0; }\n\n if (!A_as_matrix->is_matrix_singleGPU())\n {\n A_as_matrix->manager->exchange_halo_async(x, x.tag);\n\n if (A_as_matrix->getViewExterior() == A_as_matrix->getViewInterior())\n {\n A_as_matrix->manager->exchange_halo_wait(x, x.tag);\n }\n }\n\n ViewType oldView = A_as_matrix->currentView();\n ViewType flags;\n bool latencyHiding = true;\n\n if (A_as_matrix->is_matrix_singleGPU() || (x.dirtybit == 0))\n {\n latencyHiding = false;\n A_as_matrix->setViewExterior();\n flags = (ViewType)(A_as_matrix->getViewInterior() | A_as_matrix->getViewExterior());\n }\n else\n {\n flags = A_as_matrix->getViewInterior();\n A_as_matrix->setViewInterior();\n }\n\n if (A_as_matrix->get_block_dimx() == 1 && A_as_matrix->get_block_dimy() == 1)\n {\n smooth_1x1(*A_as_matrix, b, x, flags);\n }\n else\n {\n FatalError(\"Unsupported block size for BlockJacobi_Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if (latencyHiding)\n {\n A_as_matrix->manager->exchange_halo_wait(x, x.tag);\n A_as_matrix->setViewExterior();\n flags = (ViewType)(~(A_as_matrix->getViewInterior()) & A_as_matrix->getViewExterior());\n\n if (flags != 0)\n {\n if (A_as_matrix->get_block_dimx() == 1 && A_as_matrix->get_block_dimy() == 1)\n {\n smooth_1x1(*A_as_matrix, b, x, flags);\n }\n else\n {\n FatalError(\"Unsupported block size for BlockJacobi_Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n }\n }\n\n x.dirtybit = 1;\n A_as_matrix->setView(oldView);\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nChebyshevPolySolver_Base::solve_finalize( VVector &b, VVector &x )\n{}\n\n\ntemplate \nvoid ChebyshevPolySolver >::smooth_1x1(Matrix_h &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n FatalError(\"chebyshev poly smoother not implemented with host format, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nstruct chebyshev_poly_functor\n{\n ValueTypeB tau;\n chebyshev_poly_functor( ValueTypeB tau ) : tau( tau ) {}\n template __host__ __device__ ValueTypeB operator( )( const Tuple &t ) const\n {\n ValueTypeB x = amgx::thrust::get<0>(t);\n ValueTypeB b = amgx::thrust::get<1>(t);\n ValueTypeB y = amgx::thrust::get<2>(t);\n return x + tau * (b - y);\n }\n};\n\ntemplate \nvoid ChebyshevPolySolver >::smooth_1x1(Matrix_d &A, VVector &b, VVector &x, ViewType separation_flags)\n{\n if (this->y.size() != b.size())\n {\n this->y.resize(b.size());\n this->y.tag = this->tag * 100 + 3;\n this->y.set_block_dimx(b.get_block_dimx());\n this->y.set_block_dimy(b.get_block_dimy());\n }\n\n int num_rows = A.get_num_rows();\n int offset = 0;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n bool latencyHiding = separation_flags != A.getViewIntExt();// we change view only when do latency hiding, maybe it's better to use some explicit flag attached to matrix?\n\n for (int i = 0; i < this->poly_order; i++)\n {\n this->y.dirtybit = 0;\n //y = A*x\n multiply( A, x, this->y, separation_flags );\n //x += tau_i * (b - y)\n amgx::thrust::transform( amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple( x.begin() + offset, b.begin() + offset, this->y.begin() + offset)),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple( x.begin() + num_rows, b.begin() + num_rows, this->y.begin() + num_rows)),\n x.begin() + offset,\n chebyshev_poly_functor( this->tau[i] ));\n //Cublas::axpy(num_rows, (ValueTypeB)(this->tau[i]), b.raw(), 1, x.raw(), 1);\n //Cublas::axpy(num_rows, (ValueTypeB)(-1.0*this->tau[i]), this->y.raw(), 1, x.raw(), 1);\n }\n}\n\ntemplate \nvoid ChebyshevPolySolver >::smooth_with_0_initial_guess_1x1(const Matrix_h &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n FatalError(\"chebyshev poly smoother not implemented with host format, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid ChebyshevPolySolver >::smooth_with_0_initial_guess_1x1(const Matrix_d &A, const VVector &b, VVector &x, ViewType separation_flags)\n{\n}\n\ntemplate \nvoid ChebyshevPolySolver >::smooth_BxB(Matrix_h &A, VVector &b, VVector &x, bool firstStep, ViewType separation_flags)\n{\n FatalError(\"M*M chebyshev poly smoother not implemented with host format, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid ChebyshevPolySolver >::smooth_BxB(Matrix_d &A, VVector &b, VVector &x, bool firstStep, ViewType separation_flags)\n{\n FatalError(\"M*M chebyshev poly smoother not implemented with host format, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class ChebyshevPolySolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class ChebyshevPolySolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace chebyshev_poly_smoother\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/solvers/chebyshev_poly.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n// TODO: Have 2 groups of 16 threads collaborate\n// TODO: Add support for outside diagonal\n// TODO: Add support for unsorted rows\n\n#define EXPERIMENTAL_LU_FACTORS\n#define EXPERIMENTAL_LU_FORWARD\n#define EXPERIMENTAL_LU_BACKWARD\n\nnamespace amgx\n{\n\nnamespace multicolor_ilu_solver\n{\n\n// -----------\n// Kernels\n// -----------\n\n#ifdef EXPERIMENTAL_LU_FORWARD\n\ntemplate\n__global__\n__launch_bounds__( CtaSize, 16 )\nvoid LU_forward_4x4_kernel_warp( const IndexType *LU_row_offsets,\n const IndexType *LU_smaller_color_offsets,\n const IndexType *LU_column_indices,\n const ValueTypeA *LU_nonzero_values,\n const IndexType *A_row_offsets,\n const IndexType *A_column_indices,\n const ValueTypeA *A_nonzero_values,\n const IndexType *A_dia_indices,\n const ValueTypeB *x,\n const ValueTypeB *b,\n ValueTypeB *delta,\n const int *sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n bool xIsZero )\n{\n const int nHalfWarps = CtaSize / 16; // Number of half warps per Cta\n const int warpId = utils::warp_id();\n const int laneId = utils::lane_id();\n const int halfWarpId = threadIdx.x / 16;\n const int halfLaneId = threadIdx.x % 16;\n const int halfLaneId_div_4 = halfLaneId / 4;\n const int halfLaneId_mod_4 = halfLaneId % 4;\n const int upperHalf = 16 * (laneId / 16);\n // Shared memory needed to exchange X and delta.\n __shared__ volatile ValueTypeB s_mem[CtaSize];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile ValueTypeB *my_s_mem = &s_mem[16 * halfWarpId];\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( int aRowIt = blockIdx.x * nHalfWarps + halfWarpId ; aRowIt < num_rows_per_color ; aRowIt += gridDim.x * nHalfWarps )\n {\n int aRowId = sorted_rows_by_color[aRowIt];\n // Load one block of B.\n ValueTypeB my_bmAx(0);\n\n unsigned int active_mask = utils::activemask();\n\n if ( ROW_MAJOR )\n {\n if ( halfLaneId_mod_4 == 0 )\n {\n my_bmAx = __cachingLoad(&b[4 * aRowId + halfLaneId_div_4]);\n }\n }\n else\n {\n if ( halfLaneId_div_4 == 0 )\n {\n my_bmAx = __cachingLoad(&b[4 * aRowId + halfLaneId_mod_4]);\n }\n }\n\n // Don't do anything if X is zero.\n if ( !xIsZero )\n {\n int aColBegin = A_row_offsets[aRowId ];\n int aColEnd = A_row_offsets[aRowId + 1];\n int aColMax = aColEnd;\n\n if ( hasDiag )\n {\n ++aColMax;\n }\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( aColBegin < aColMax, active_mask ) ; aColBegin += 16 )\n {\n int aColIt = aColBegin + halfLaneId;\n // Get the ID of the column.\n int aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColId = A_column_indices[aColIt];\n }\n\n if ( hasDiag && aColIt == aColEnd )\n {\n aColId = aRowId;\n }\n\n // Count the number of active columns.\n int vote = utils::ballot(aColId != -1, active_mask);\n // The number of iterations.\n int nCols = max( __popc( vote & 0x0000ffff ), __popc( vote & 0xffff0000 ) );\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < nCols ; k += 4 )\n {\n int my_k = k + halfLaneId_div_4;\n // Load 8 blocks of X.\n int waColId = utils::shfl( aColId, upperHalf + my_k, warpSize, active_mask );\n ValueTypeB my_x(0);\n\n if ( waColId != -1 )\n {\n my_x = __cachingLoad(&x[4 * waColId + halfLaneId_mod_4]);\n }\n\n my_s_mem[halfLaneId] = my_x;\n // Load 8 blocks of A.\n#pragma unroll\n for ( int i = 0 ; i < 4 ; ++i )\n {\n int w_aColTmp = aColBegin + k + i, w_aColIt = -1;\n\n if ( w_aColTmp < aColEnd )\n {\n w_aColIt = w_aColTmp;\n }\n\n if ( hasDiag && w_aColTmp == aColEnd )\n {\n w_aColIt = A_dia_indices[aRowId];\n }\n\n ValueTypeA my_val(0);\n\n if ( w_aColIt != -1 )\n {\n my_val = A_nonzero_values[16 * w_aColIt + halfLaneId];\n }\n\n if ( ROW_MAJOR )\n {\n my_bmAx -= my_val * my_s_mem[4 * i + halfLaneId_mod_4];\n }\n else\n {\n my_bmAx -= my_val * my_s_mem[4 * i + halfLaneId_div_4];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n } // if xIsZero\n\n // Contribution from each nonzero column that has color less than yours\n if ( current_color != 0 )\n {\n // TODO: Use constant or texture here\n int aColBegin = LU_row_offsets[aRowId];\n int aColEnd = LU_smaller_color_offsets[aRowId];\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( aColBegin < aColEnd, active_mask ) ; aColBegin += 16 )\n {\n int aColIt = aColBegin + halfLaneId;\n int aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColId = LU_column_indices[aColIt];\n }\n\n // Count the number of active columns.\n int vote = utils::ballot(aColId != -1, active_mask);\n // The number of iterations.\n int nCols = max( __popc( vote & 0x0000ffff ), __popc( vote & 0xffff0000 ) );\n\n for ( int k = 0 ; k < nCols ; k += 4 )\n {\n int my_k = k + halfLaneId_div_4;\n // Load 8 blocks of X.\n int waColId = utils::shfl( aColId, upperHalf + my_k, warpSize, active_mask );\n ValueTypeB my_delta(0);\n\n if ( waColId != -1 )\n {\n my_delta = delta[4 * waColId + halfLaneId_mod_4];\n }\n\n my_s_mem[halfLaneId] = my_delta;\n utils::syncwarp(); // making sure smem write propagated\n // Update b-Ax.\n#pragma unroll\n for ( int i = 0 ; i < 4 ; ++i )\n {\n int w_aColTmp = aColBegin + k + i, w_aColIt = -1;\n\n if ( w_aColTmp < aColEnd )\n {\n w_aColIt = w_aColTmp;\n }\n\n ValueTypeA my_val(0);\n\n if ( w_aColIt != -1 )\n {\n my_val = LU_nonzero_values[16 * w_aColIt + halfLaneId];\n }\n\n if ( ROW_MAJOR )\n {\n my_bmAx -= my_val * my_s_mem[4 * i + halfLaneId_mod_4];\n }\n else\n {\n my_bmAx -= my_val * my_s_mem[4 * i + halfLaneId_div_4];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n } // If current_color != 0\n\n // Reduce bmAx terms.\n if ( ROW_MAJOR )\n {\n my_bmAx += utils::shfl_xor( my_bmAx, 1, warpSize, active_mask );\n my_bmAx += utils::shfl_xor( my_bmAx, 2, warpSize, active_mask );\n }\n else\n {\n my_bmAx += utils::shfl_xor( my_bmAx, 4, warpSize, active_mask );\n my_bmAx += utils::shfl_xor( my_bmAx, 8, warpSize, active_mask );\n }\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n if ( halfLaneId_mod_4 == 0 )\n {\n delta[4 * aRowId + halfLaneId_div_4] = my_bmAx;\n }\n }\n else\n {\n if ( halfLaneId_div_4 == 0 )\n {\n delta[4 * aRowId + halfLaneId_mod_4] = my_bmAx;\n }\n }\n }\n}\n\n#else\ntemplate\n__global__\nvoid LU_forward_4x4_kernel(const IndexType *LU_row_offsets, const IndexType *LU_smaller_color_offsets, const IndexType *LU_column_indices, const ValueTypeA *LU_nonzero_values, const IndexType *A_row_offsets, const IndexType *A_column_indices, const ValueTypeA *A_nonzero_values,\n const ValueTypeB *x, const ValueTypeB *b, ValueTypeB *delta, const int *sorted_rows_by_color,\n const int num_rows_per_color, const int current_color, bool xIsZero)\n\n\n{\n int warp_id = threadIdx.x / 32;\n int warp_thread_id = threadIdx.x & 31;\n\n // padding row blocks to fit in a single warp\n if ( warp_thread_id >= blockrows_per_warp * bsize ) { return; }\n\n // new thread id with padding\n int tid = warp_id * blockrows_per_warp * bsize + warp_thread_id;\n // Here we use one thread per row (not block row)\n int cta_blockrow_id = (tid) / bsize;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int vec_entry_index = tid - cta_blockrow_id * bsize;\n volatile __shared__ ValueTypeB s_delta_temp[ bsize * blockrows_per_cta];\n int offset, s_offset, i;\n ValueTypeB bmAx, temp[bsize];\n\n while (blockrow_id < num_rows_per_color && cta_blockrow_id < blockrows_per_cta)\n {\n i = sorted_rows_by_color[blockrow_id];\n // Load RHS and x\n offset = i * bsize + vec_entry_index;\n bmAx = b[offset];\n\n if (!xIsZero)\n {\n int jmin = A_row_offsets[i];\n int jmax = A_row_offsets[i + 1];\n\n //TODO: Assumes inside diagonal\n for (int jind = jmin; jind < jmax; jind++)\n {\n IndexType jcol = A_column_indices[jind];\n offset = jcol * bsize + vec_entry_index;\n s_delta_temp[tid] = x[offset];\n\n // Load nonzero_values\n if (ROW_MAJOR)\n {\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(A_nonzero_values + offset, temp);\n }\n else\n {\n offset = jind * bsize * bsize + vec_entry_index;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n temp[m] = A_nonzero_values[offset + bsize * m];\n }\n }\n\n // Do matrix multiply\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_delta_temp[s_offset++];\n }\n }\n }\n\n // Contribution from each nonzero column that has color less than yours\n if (current_color != 0)\n {\n int jmin = LU_row_offsets[i];\n int jmax = LU_smaller_color_offsets[i];\n\n for (int jind = jmin; jind < jmax; jind++)\n {\n IndexType jcol = LU_column_indices[jind];\n offset = jcol * bsize + vec_entry_index;\n s_delta_temp[tid] = __ldcg(delta + offset);\n\n // Load nonzero_values\n if (ROW_MAJOR)\n {\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(LU_nonzero_values + offset, temp);\n }\n else\n {\n offset = jind * bsize * bsize + vec_entry_index;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n temp[m] = LU_nonzero_values[offset + bsize * m];\n }\n }\n\n // Do matrix multiply\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_delta_temp[s_offset++];\n }\n }\n }\n\n delta[i * bsize + vec_entry_index] = bmAx;\n blockrow_id += blockrows_per_cta * gridDim.x;\n }\n}\n#endif\n\n#ifdef EXPERIMENTAL_LU_BACKWARD\n\ntemplate< typename IndexType, typename ValueTypeA, typename ValueTypeB, int CtaSize, bool ROW_MAJOR >\n__global__\n__launch_bounds__( CtaSize, 16 )\nvoid LU_backward_4x4_kernel_warp( const IndexType *row_offsets,\n const IndexType *larger_color_offsets,\n const IndexType *column_indices,\n const IndexType *dia_indices,\n const ValueTypeA *nonzero_values,\n const ValueTypeB *delta,\n ValueTypeB *Delta,\n ValueTypeB *x,\n const int *sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int num_colors,\n const ValueTypeB weight,\n bool xIsZero )\n{\n const int nHalfWarps = CtaSize / 16; // Number of half warps per CTA.\n const int warpId = utils::warp_id();\n const int laneId = utils::lane_id();\n const int halfWarpId = threadIdx.x / 16;\n const int halfLaneId = threadIdx.x % 16;\n const int halfLaneId_div_4 = halfLaneId / 4;\n const int halfLaneId_mod_4 = halfLaneId % 4;\n const int upperHalf = 16 * (laneId / 16);\n // Shared memory needed to exchange X and delta.\n __shared__ volatile ValueTypeB s_mem[CtaSize];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile ValueTypeB *my_s_mem = &s_mem[16 * halfWarpId];\n\n // Iterate over the rows of the matrix. One warp per two rows.\n for ( int aRowIt = blockIdx.x * nHalfWarps + halfWarpId ; aRowIt < num_rows_per_color ; aRowIt += gridDim.x * nHalfWarps )\n {\n int aRowId = sorted_rows_by_color[aRowIt];\n unsigned int active_mask = utils::activemask();\n // Load one block of B.\n ValueTypeB my_bmAx(0);\n\n if ( ROW_MAJOR )\n {\n if ( halfLaneId_mod_4 == 0 )\n {\n my_bmAx = delta[4 * aRowId + halfLaneId_div_4];\n }\n }\n else\n {\n if ( halfLaneId_div_4 == 0 )\n {\n my_bmAx = delta[4 * aRowId + halfLaneId_mod_4];\n }\n }\n\n // Don't do anything if the color is not the interesting one.\n if ( current_color != num_colors - 1 )\n {\n // The range of the rows.\n int aColBegin = larger_color_offsets[aRowId], aColEnd = row_offsets[aRowId + 1];\n\n // Each warp load column indices of 16 nonzero blocks\n for ( ; utils::any( aColBegin < aColEnd, active_mask ) ; aColBegin += 16 )\n {\n int aColIt = aColBegin + halfLaneId;\n // Get the ID of the column.\n int aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColId = column_indices[aColIt];\n }\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < 16 ; k += 4 )\n {\n int my_k = k + halfLaneId_div_4;\n // Exchange column indices.\n int waColId = utils::shfl( aColId, upperHalf + my_k, warpSize, active_mask );\n // Load 8 blocks of X if needed.\n ValueTypeB *my_ptr = Delta;\n\n if ( xIsZero )\n {\n my_ptr = x;\n }\n\n ValueTypeB my_x(0);\n\n if ( waColId != -1 )\n {\n my_x = my_ptr[4 * waColId + halfLaneId_mod_4];\n }\n\n my_s_mem[halfLaneId] = my_x;\n utils::syncwarp();\n // Load 8 blocks of A.\n#pragma unroll\n for ( int i = 0 ; i < 4 ; ++i )\n {\n int w_aColTmp = aColBegin + k + i, w_aColIt = -1;\n\n if ( w_aColTmp < aColEnd )\n {\n w_aColIt = w_aColTmp;\n }\n\n ValueTypeA my_val(0);\n\n if ( w_aColIt != -1 )\n {\n my_val = nonzero_values[16 * w_aColIt + halfLaneId];\n }\n\n if ( ROW_MAJOR )\n {\n my_bmAx -= my_val * my_s_mem[4 * i + halfLaneId_mod_4];\n }\n else\n {\n my_bmAx -= my_val * my_s_mem[4 * i + halfLaneId_div_4];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Reduce bmAx terms.\n if ( ROW_MAJOR )\n {\n my_bmAx += utils::shfl_xor( my_bmAx, 1, warpSize, active_mask );\n my_bmAx += utils::shfl_xor( my_bmAx, 2, warpSize, active_mask );\n }\n else\n {\n my_bmAx += utils::shfl_xor( my_bmAx, 4, warpSize, active_mask );\n my_bmAx += utils::shfl_xor( my_bmAx, 8, warpSize, active_mask );\n }\n } // if current_color != num_colors-1\n\n // Update the shared terms.\n if ( ROW_MAJOR )\n {\n if ( halfLaneId_mod_4 == 0 )\n {\n my_s_mem[halfLaneId_div_4] = my_bmAx;\n }\n }\n else\n {\n if ( halfLaneId_div_4 == 0 )\n {\n my_s_mem[halfLaneId_mod_4] = my_bmAx;\n }\n }\n\n // Update the diagonal term.\n int w_aColIt = dia_indices[aRowId];\n ValueTypeA my_val(0);\n utils::syncwarp();\n\n if ( w_aColIt != -1 )\n {\n my_val = nonzero_values[16 * w_aColIt + halfLaneId];\n }\n\n if ( ROW_MAJOR )\n {\n my_bmAx = my_val * my_s_mem[halfLaneId_mod_4];\n }\n else\n {\n my_bmAx = my_val * my_s_mem[halfLaneId_div_4];\n }\n\n // Regroup results.\n if ( ROW_MAJOR )\n {\n my_bmAx += utils::shfl_xor( my_bmAx, 1 );\n my_bmAx += utils::shfl_xor( my_bmAx, 2 );\n }\n else\n {\n my_bmAx += utils::shfl_xor( my_bmAx, 4 );\n my_bmAx += utils::shfl_xor( my_bmAx, 8 );\n }\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n ValueTypeB my_x(0);\n\n if ( !xIsZero && halfLaneId_mod_4 == 0 )\n {\n my_x = x[4 * aRowId + halfLaneId_div_4];\n }\n\n my_x += weight * my_bmAx;\n\n if ( !xIsZero && halfLaneId_mod_4 == 0 )\n {\n Delta[4 * aRowId + halfLaneId_div_4] = my_bmAx;\n }\n\n if ( halfLaneId_mod_4 == 0 )\n {\n x[4 * aRowId + halfLaneId_div_4] = my_x;\n }\n }\n else\n {\n ValueTypeB my_x(0);\n\n if ( !xIsZero && halfLaneId_div_4 == 0 )\n {\n my_x = x[4 * aRowId + halfLaneId_mod_4];\n }\n\n my_x += weight * my_bmAx;\n\n if ( !xIsZero && halfLaneId_div_4 == 0 )\n {\n Delta[4 * aRowId + halfLaneId_mod_4] = my_bmAx;\n }\n\n if ( halfLaneId_div_4 == 0 )\n {\n x[4 * aRowId + halfLaneId_mod_4] = my_x;\n }\n }\n }\n}\n\n#else\n\ntemplate\n__global__\nvoid LU_backward_4x4_kernel(const IndexType *row_offsets, const IndexType *larger_color_offsets, const IndexType *column_indices, const IndexType *dia_indices, const ValueTypeA *nonzero_values,\n const ValueTypeB *delta, ValueTypeB *Delta, ValueTypeB *x, const int *sorted_rows_by_color,\n const int num_rows_per_color, const int current_color, const int num_colors, const ValueTypeB weight, bool xIsZero)\n\n\n{\n int warp_id = threadIdx.x / 32;\n int warp_thread_id = threadIdx.x & 31;\n\n // padding row blocks to fit in a single warp\n if ( warp_thread_id >= blockrows_per_warp * bsize ) { return; }\n\n // new thread id with padding\n int tid = warp_id * blockrows_per_warp * bsize + warp_thread_id;\n // Here we use one thread per row (not block row)\n int cta_blockrow_id = (tid) / bsize;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int vec_entry_index = tid - cta_blockrow_id * bsize;\n volatile __shared__ ValueTypeB s_x_temp[ bsize * blockrows_per_cta];\n int offset, s_offset, i;\n ValueTypeB bmAx, temp[bsize];\n\n while (blockrow_id < num_rows_per_color && cta_blockrow_id < blockrows_per_cta)\n {\n i = sorted_rows_by_color[blockrow_id];\n // Load RHS and x\n offset = i * bsize + vec_entry_index;\n bmAx = delta[offset];\n\n // Contribution from each nonzero column that has color less than yours\n if (current_color != num_colors)\n {\n int jmin = larger_color_offsets[i];\n int jmax = row_offsets[i + 1];\n\n for (int jind = jmin; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n offset = jcol * bsize + vec_entry_index;\n\n if (xIsZero)\n {\n s_x_temp[tid] = __ldcg(x + offset);\n }\n else\n {\n s_x_temp[tid] = __ldcg(Delta + offset);\n }\n\n // Load nonzero_values\n if (ROW_MAJOR)\n {\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n }\n else\n {\n offset = jind * bsize * bsize + vec_entry_index;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n temp[m] = nonzero_values[offset + bsize * m];\n }\n }\n\n // Do matrix multiply\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_x_temp[s_offset++];\n }\n }\n }\n\n s_x_temp[tid] = bmAx;\n bmAx = 0.;\n\n // Load diagonals (which store the inverse)\n if (ROW_MAJOR)\n {\n offset = dia_indices[i] * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n }\n else\n {\n offset = dia_indices[i] * bsize * bsize + vec_entry_index;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n temp[m] = nonzero_values[offset + bsize * m];\n }\n }\n\n // Do matrix-vector multiply\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n bmAx += temp[m] * s_x_temp[s_offset++];\n }\n\n offset = i * bsize + vec_entry_index;\n\n if (xIsZero)\n {\n x[offset] = weight * bmAx;\n }\n else\n {\n Delta[offset] = bmAx;\n x[offset] += weight * bmAx ;\n }\n\n blockrow_id += blockrows_per_cta * gridDim.x;\n }\n}\n\n#endif\n\n// Assumptions:\n// CtaSize must be multiple of 32\n// SMemSize should be larger than the maximum number of columns in the matrix\n// Matrix B is superset of matrix A\n\ntemplate< int CtaSize, int SMemSize>\n__global__ __launch_bounds__( CtaSize )\nvoid\ncomputeAtoLUmapping_kernel( int A_nRows,\n const int *__restrict A_row_offsets,\n const int *__restrict A_col_indices,\n const int *__restrict B_row_offsets,\n const int *__restrict B_col_indices,\n int *__restrict AtoBmapping,\n int *wk_returnValue )\n{\n const int nWarps = CtaSize / 32; // Number of warps per Cta\n const int warpId = utils::warp_id();\n const int laneId = utils::lane_id();\n // Rows are stored in SMEM. Linear storage.\n __shared__ volatile int s_colInd[nWarps][SMemSize];\n // The row this warp is responsible for\n int aRowId = blockIdx.x * nWarps + warpId;\n\n // Loop over rows of A.\n for ( ; aRowId < A_nRows ; aRowId += nWarps * gridDim.x )\n {\n // Insert all the column indices of matrix B in the shared memory table\n int bColBeg = B_row_offsets[aRowId];\n int bColEnd = B_row_offsets[aRowId + 1];\n // The number of columns.\n const int nCols = bColEnd - bColBeg;\n\n //TODO: Add fallback for cases where number of nonzeros exceed SMemSize\n if ( nCols > SMemSize )\n {\n wk_returnValue[0] = 1;\n return;\n }\n\n // Fill-in the local table.\n const int NUM_STEPS = SMemSize / 32;\n#pragma unroll\n for ( int step = 0, k = laneId ; step < NUM_STEPS ; ++step, k += 32 )\n {\n int bColIt = bColBeg + k;\n int bColId = -1;\n\n if ( bColIt < bColEnd )\n {\n bColId = B_col_indices[bColIt];\n }\n\n s_colInd[warpId][k] = bColId;\n }\n\n // Now load column indices of current row of A\n int aColIt = A_row_offsets[aRowId];\n int aColEnd = A_row_offsets[aRowId + 1];\n\n for ( aColIt += laneId ; utils::any(aColIt < aColEnd) ; aColIt += 32 )\n {\n // The column.\n int aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColId = A_col_indices[aColIt];\n }\n\n // Each thread searches for its column id, and gets the corresponding bColIt\n // TODO: Try binary search or using hash table\n int foundOffset = -1;\n\n if ( aColId == -1 )\n {\n foundOffset = -2;\n }\n\n for ( int i = 0 ; i < nCols && utils::any(foundOffset == -1) ; ++i )\n if ( foundOffset == -1 && s_colInd[warpId][i] == aColId )\n {\n foundOffset = i;\n }\n\n // Store the result.\n if ( aColIt < aColEnd )\n {\n AtoBmapping[aColIt] = bColBeg + foundOffset;\n }\n }\n } // if RowId < A_nRows;\n}\n\ntemplate< int CtaSize, int SMemSize>\n__global__ __launch_bounds__( CtaSize )\nvoid\ncomputeAtoLUmappingExtDiag_kernel( int A_nRows,\n const int *__restrict A_row_offsets,\n const int *__restrict A_col_indices,\n const int *__restrict A_dia_indices,\n const int *__restrict B_row_offsets,\n const int *__restrict B_col_indices,\n int *__restrict AtoBmapping,\n int *wk_returnValue )\n{\n const int nWarps = CtaSize / 32; // Number of warps per Cta\n const int warpId = utils::warp_id();\n const int laneId = utils::lane_id();\n // Rows are stored in SMEM. Linear storage.\n __shared__ volatile int s_colInd[nWarps][SMemSize];\n // The row this warp is responsible for\n int aRowId = blockIdx.x * nWarps + warpId;\n\n // Loop over rows of A.\n for ( ; aRowId < A_nRows ; aRowId += nWarps * gridDim.x )\n {\n // Insert all the column indices of matrix B in the shared memory table\n int bColBeg = B_row_offsets[aRowId];\n int bColEnd = B_row_offsets[aRowId + 1];\n // The number of columns.\n const int nCols = bColEnd - bColBeg;\n\n //TODO: Add fallback for cases where number of nonzeros exceed SMemSize\n if ( nCols > SMemSize )\n {\n wk_returnValue[0] = 1;\n return;\n }\n\n // Fill-in the local table.\n const int NUM_STEPS = SMemSize / 32;\n#pragma unroll\n for ( int step = 0, k = laneId ; step < NUM_STEPS ; ++step, k += 32 )\n {\n int bColIt = bColBeg + k;\n int bColId = -1;\n\n if ( bColIt < bColEnd )\n {\n bColId = B_col_indices[bColIt];\n }\n\n s_colInd[warpId][k] = bColId;\n }\n\n // Now load column indices of current row of A\n int aColIt = A_row_offsets[aRowId];\n int aColEnd = A_row_offsets[aRowId + 1];\n\n for ( aColIt += laneId ; utils::any(aColIt <= aColEnd) ; aColIt += 32 )\n {\n // The column.\n int aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColId = A_col_indices[aColIt];\n }\n\n if ( aColIt == aColEnd )\n {\n aColId = aRowId;\n }\n\n // Each thread searches for its column id, and gets the corresponding bColIt\n // TODO: Try binary search or using hash table\n int foundOffset = -1;\n\n if ( aColId == -1 )\n {\n foundOffset = -2;\n }\n\n for ( int i = 0 ; i < nCols && utils::any(foundOffset == -1) ; ++i )\n if ( foundOffset == -1 && s_colInd[warpId][i] == aColId )\n {\n foundOffset = i;\n }\n\n // Store the result.\n int aDst = -1;\n\n if ( aColIt < aColEnd )\n {\n aDst = aColIt;\n }\n\n if ( aColIt == aColEnd )\n {\n aDst = A_dia_indices[aRowId];\n }\n\n if ( aDst != -1 )\n {\n AtoBmapping[aDst] = bColBeg + foundOffset;\n }\n }\n }\n}\n\n#ifdef EXPERIMENTAL_LU_FACTORS\n\ntemplate< typename ValueTypeA, int CtaSize, int SMemSize, bool ROW_MAJOR >\n__global__\n__launch_bounds__( CtaSize, 12 )\nvoid\ncompute_LU_factors_4x4_kernel_warp( int A_nRows,\n const int *__restrict A_row_offsets,\n const int *__restrict A_col_indices,\n const int *__restrict A_dia_indices,\n ValueTypeA *__restrict A_nonzero_values,\n const int *__restrict A_smaller_color_offsets,\n const int *__restrict A_larger_color_offsets,\n const int *sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n int *wk_returnValue )\n{\n const int nWarps = CtaSize / 32; // Number of warps per Cta\n const int warpId = utils::warp_id();\n const int laneId = utils::lane_id();\n // Lane ID in the 2 16-wide segments.\n const int lane_id_div_16 = laneId / 16;\n const int lane_id_mod_16 = laneId % 16;\n // Coordinates inside a 4x4 block of the matrix.\n const int idx_i = lane_id_mod_16 / 4;\n const int idx_j = lane_id_mod_16 % 4;\n int globalWarpId = blockIdx.x * nWarps + warpId;\n // Shared memory to store the blocks to process\n __shared__ volatile ValueTypeA s_C_mtx[nWarps][32];\n __shared__ volatile ValueTypeA s_F_mtx[nWarps][16];\n // Shared memory to store the proposed column to load\n __shared__ volatile int s_aColSrc[nWarps][32];\n // Shared memory to store the column indices of the current row\n __shared__ volatile int s_keys[nWarps][SMemSize];\n\n while (globalWarpId < num_rows_per_color)\n {\n int storedRowId[2];\n int I = 0;\n\n for (; I < 2 && globalWarpId < num_rows_per_color ; I++)\n {\n int aRowId = sorted_rows_by_color[globalWarpId];\n storedRowId[I] = aRowId;\n int aColBeg = A_row_offsets[aRowId + 0];\n int aColEnd = A_row_offsets[aRowId + 1];\n int aColSmaller = A_smaller_color_offsets[aRowId];\n // The number of columns.\n const int nCols = aColEnd - aColBeg;\n\n //TODO: Add fallback for cases where number of nonzeros exceed SMemSize\n if ( nCols > SMemSize )\n {\n wk_returnValue[0] = 1;\n return;\n }\n\n // Fill-in the local table.\n const int NUM_STEPS = SMemSize / 32;\n\n#pragma unroll\n for ( int step = 0, k = laneId ; step < NUM_STEPS ; ++step, k += 32 )\n {\n int aColIt = aColBeg + k;\n int aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColId = A_col_indices[aColIt];\n }\n\n s_keys[warpId][k] = aColId;\n }\n\n // Now load all column indices of neighbours that have colors smaller than yours\n for ( int aColIt = aColBeg; aColIt < aColSmaller ; aColIt++)\n {\n unsigned int active_mask = utils::activemask();\n // Read the row to process, should be a broadcast\n int waRowId = s_keys[warpId][aColIt - aColBeg];\n // Compute multiplicative factor, load C_jj in first half, C_ij in second half\n int aColIdx = aColIt;\n\n if ( lane_id_div_16 == 0 )\n {\n aColIdx = A_dia_indices[waRowId];\n }\n\n s_C_mtx[warpId][laneId] = A_nonzero_values[16 * aColIdx + lane_id_mod_16];\n // Threads 0-15 perform the matrix product\n ValueTypeA tmp(0);\n\n if (ROW_MAJOR)\n {\n\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_C_mtx[warpId][16 + 4 * idx_i + m] * s_C_mtx[warpId][4 * m + idx_j];\n }\n }\n else\n {\n\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_C_mtx[warpId][16 + 4 * m + idx_j] * s_C_mtx[warpId][4 * idx_i + m];\n }\n }\n\n if ( lane_id_div_16 == 0 )\n {\n s_F_mtx[warpId][laneId] = tmp;\n A_nonzero_values[16 * aColIt + laneId] = tmp;\n }\n\n int waColIt = __ldcg(A_larger_color_offsets + waRowId);\n int waColEnd = __ldcg(A_row_offsets + waRowId + 1);\n\n // Load the first 32 columns of waRowId\n for (waColIt += laneId ; utils::any(waColIt < waColEnd, active_mask ); waColIt += 32 )\n {\n // Each thread loads its column id\n int waColId = -1;\n\n if ( waColIt < waColEnd )\n {\n waColId = A_col_indices[waColIt];\n }\n\n // Find the right column.\n int found_aColIt = -1;\n\n#pragma unroll 4\n for ( int i = 0, num_keys = aColEnd - aColBeg ; i < num_keys ; ++i )\n if ( s_keys[warpId][i] == waColId )\n {\n found_aColIt = i;\n }\n\n if ( found_aColIt != -1 )\n {\n found_aColIt += aColBeg;\n }\n\n // Store all the columns that have been found\n const int pred = found_aColIt != -1;\n int vote = utils::ballot( pred, active_mask );\n const int idst = __popc(vote & utils::lane_mask_lt());\n\n if (pred)\n {\n s_aColSrc[warpId][idst] = laneId;\n }\n utils::syncwarp(active_mask);\n\n const int n_cols = __popc( vote );\n\n // Process all columns that have been found\n for ( int k = 0 ; k < n_cols ; k += 2 )\n {\n const int my_k = k + lane_id_div_16;\n // Where to get columns from.\n int a_col_it = -1, w_col_it = -1;\n // Load column to load\n a_col_it = utils::shfl(found_aColIt, s_aColSrc[warpId][my_k], warpSize, active_mask);\n w_col_it = utils::shfl(waColIt, s_aColSrc[warpId][my_k], warpSize, active_mask);\n\n if ( my_k >= n_cols )\n {\n a_col_it = -1;\n w_col_it = -1;\n }\n\n ValueTypeA my_C(0);\n\n if ( w_col_it != -1 )\n {\n my_C = A_nonzero_values[16 * w_col_it + lane_id_mod_16];\n }\n\n s_C_mtx[warpId][laneId] = my_C;\n // Run the matrix-matrix product.\n ValueTypeA tmp(0);\n utils::syncwarp( active_mask );\n\n if (ROW_MAJOR)\n {\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_F_mtx[warpId][4 * idx_i + m] * s_C_mtx[warpId][16 * lane_id_div_16 + 4 * m + idx_j];\n }\n }\n else\n {\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_F_mtx[warpId][4 * m + idx_j] * s_C_mtx[warpId][16 * lane_id_div_16 + 4 * idx_i + m];\n }\n }\n\n if ( a_col_it != -1 )\n {\n A_nonzero_values[16 * a_col_it + lane_id_mod_16] -= tmp;\n }\n } // Loop over columns that have a match (for k=0;k( s_C_mtx[warpId], 16 * lane_id_div_16, offset, idx_i, idx_j, A_nonzero_values );\n }\n else\n {\n compute_block_inverse_col_major4x4_formula2( s_C_mtx[warpId], 16 * lane_id_div_16, offset, idx_i, idx_j, A_nonzero_values );\n }\n } // End of if statement\n } // End of while loop\n}\n\n#else\n\ntemplate< typename ValueTypeA, int CtaSize, int SMemSize, bool ROW_MAJOR>\n__global__ __launch_bounds__( CtaSize )\nvoid\ncomputeLUFactors_4x4_kernel( int A_nRows,\n const int *__restrict A_row_offsets,\n const int *__restrict A_col_indices,\n const int *__restrict A_dia_indices,\n ValueTypeA *__restrict A_nonzero_values,\n const int *__restrict A_smaller_color_offsets,\n const int *__restrict A_larger_color_offsets,\n const int *sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n int *wk_returnValue )\n{\n const int nWarps = CtaSize / 32; // Number of warps per Cta\n const int warpId = utils::warp_id();\n const int laneId = utils::lane_id();\n int lane_mask_lt = utils::lane_mask_lt();\n\n // Lane ID in the 2 16-wide segments.\n const int lane_id_div_16 = laneId / 16;\n const int lane_id_mod_16 = laneId % 16;\n // Coordinates inside a 4x4 block of the matrix.\n const int idx_i = lane_id_mod_16 / 4;\n const int idx_j = lane_id_mod_16 % 4;\n int globalWarpId = blockIdx.x * nWarps + warpId;\n // Shared memory to store the blocks to process\n __shared__ volatile ValueTypeA s_C_mtx[nWarps][32];\n // Shared memory to store the proposed column to load\n __shared__ volatile int s_aColItToLoad[nWarps][32];\n __shared__ volatile int s_waColItToLoad[nWarps][32];\n // Shared memory to store the proposed column to load\n __shared__ volatile unsigned s_aColIds[nWarps][32];\n // The size of the hash table (one per warp - shared memory).\n __shared__ volatile int s_size[nWarps][2];\n // Shared memory to store the column indices of the current row\n __shared__ volatile int s_keys[nWarps][SMemSize];\n\n while (globalWarpId < num_rows_per_color)\n {\n int aRowId = sorted_rows_by_color[globalWarpId];\n // Insert all the column indices in shared memory\n // TODO: Use texture here\n int aColBeg = A_row_offsets[aRowId];\n int aColEnd = A_row_offsets[aRowId + 1];\n int aColIt = aColBeg;\n\n // Check if number of nonzeros will fit in shared memory\n if ( (aColEnd - aColBeg) > SMemSize )\n {\n wk_returnValue[0] = 1;\n return;\n }\n\n // Load the all the column indices of row into shared memory\n for ( aColIt += laneId ; utils::any( aColIt < aColEnd ) ; aColIt += 32 )\n {\n int aColId = aColIt < aColEnd ? (int) A_col_indices[aColIt] : -1;\n s_keys[warpId][aColIt - aColBeg] = aColId;\n }\n\n // Now load all column indices of neighbours that have colors smaller than yours\n aColIt = aColBeg;\n int aColSmaller = A_smaller_color_offsets[aRowId];\n\n for ( ; utils::any( (aColIt + laneId) < aColSmaller ) ; aColIt += 32 )\n {\n int aColId = (aColIt + laneId) < aColSmaller ? (int) A_col_indices[aColIt + laneId] : -1;\n // Each thread pushes its column\n s_aColIds[warpId][laneId] = aColId;\n\n // Have warp collaborate to load each row\n for ( int j = 0; j < 32; j++)\n {\n // Check if row to load is valid\n if ( ( aColIt + j ) >= aColSmaller ) { break; }\n\n // Read the row to process, should be a broadcast\n int waRowId = s_aColIds[warpId][j];\n\n // Compute multiplicative factor, load C_jj in first half, C_ij in second half\n if (lane_id_div_16 == 0)\n {\n s_C_mtx[warpId][laneId] = A_nonzero_values[ 16 * A_dia_indices[waRowId] + lane_id_mod_16 ];\n }\n else\n {\n s_C_mtx[warpId][laneId] = A_nonzero_values[ 16 * (aColIt + j) + lane_id_mod_16 ];\n }\n\n // Threads 0-15 perform the matrix product\n utils::syncwarp();\n if (lane_id_div_16 == 0)\n {\n ValueTypeA tmp(0);\n\n if (ROW_MAJOR)\n {\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_C_mtx[warpId][16 + 4 * idx_i + m] * s_C_mtx[warpId][4 * m + idx_j];\n }\n }\n else\n {\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_C_mtx[warpId][16 + 4 * m + idx_j] * s_C_mtx[warpId][4 * idx_i + m];\n }\n }\n\n s_C_mtx[warpId][laneId] = tmp;\n A_nonzero_values[16 * (aColIt + j) + laneId] = tmp;\n }\n\n int waColIt = A_larger_color_offsets[waRowId];\n int waColEnd = A_row_offsets[waRowId + 1];\n\n //// Load the first 32 columns of waRowId\n for (waColIt += laneId ; utils::any(waColIt < waColEnd ); waColIt += 32 )\n {\n // Each thread loads its column id\n int waColId = waColIt < waColEnd ? A_col_indices[waColIt] : int (-1);\n // TODO: Try binary search if columns are ordered\n int found_aColIt = -1;\n\n //TODO: if invalid waColId, don't search\n for (int i = 0 ; utils::any(found_aColIt == -1) && i < aColEnd - aColBeg ; i++)\n {\n if (s_keys[warpId][i] == waColId) { found_aColIt = aColBeg + i; }\n }\n\n // Store all the columns that have been found\n const int pred = found_aColIt != -1;\n const int vote = utils::ballot( pred );\n const int idst = __popc(vote & lane_mask_lt);\n\n if (pred)\n {\n s_aColItToLoad [warpId][idst] = found_aColIt;\n s_waColItToLoad[warpId][idst] = waColIt;\n }\n\n const int n_cols = __popc( vote );\n\n // Process all columns that have been found\n for ( int k = 0 ; k < n_cols ; k++ )\n {\n // Load column to load\n const int a_col_it = k < n_cols ? s_aColItToLoad [warpId][k] : -1;\n const int w_col_it = k < n_cols ? s_waColItToLoad[warpId][k] : -1;\n\n if (lane_id_div_16 == 1)\n {\n s_C_mtx[warpId][laneId] = A_nonzero_values[16 * w_col_it + lane_id_mod_16];\n // Run the matrix-matrix product.\n ValueTypeA tmp(0);\n utils::syncwarp(utils::activemask());\n\n if (ROW_MAJOR)\n {\n\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_C_mtx[warpId][4 * idx_i + m] * s_C_mtx[warpId][16 + 4 * m + idx_j];\n }\n }\n else\n {\n\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_C_mtx[warpId][4 * m + idx_j] * s_C_mtx[warpId][16 + 4 * idx_i + m];\n }\n }\n\n A_nonzero_values[16 * a_col_it + lane_id_mod_16] -= tmp;\n }\n } // Loop over columns that have a match (for k=0;k\n (s_C_mtx[warpId], 0, offset, idx_i, idx_j, A_nonzero_values);\n }\n else\n {\n compute_block_inverse_col_major\n (s_C_mtx[warpId], 0, offset, idx_i, idx_j, A_nonzero_values);\n }\n }\n\n globalWarpId += nWarps * gridDim.x;\n } // if RowId < A_nRows;\n}\n#endif\n// ----------\n// Methods\n// ----------\n\n// Constructor\ntemplate\nMulticolorILUSolver_Base::MulticolorILUSolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope)\n{\n m_sparsity_level = cfg.AMG_Config::template getParameter(\"ilu_sparsity_level\", cfg_scope);\n m_weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n this->m_reorder_cols_by_color_desired = (cfg.AMG_Config::template getParameter(\"reorder_cols_by_color\", cfg_scope) != 0);\n this->m_insert_diagonal_desired = (cfg.AMG_Config::template getParameter(\"insert_diag_while_reordering\", cfg_scope) != 0);\n\n if (cfg.AMG_Config::template getParameter(\"use_bsrxmv\", cfg_scope))\n {\n this->m_use_bsrxmv = 1;\n }\n else\n {\n this->m_use_bsrxmv = 0;\n }\n\n if (m_weight == ValueTypeB(0.))\n {\n m_weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Multicolor DILU smoother\\n\");\n }\n}\n\n// Destructor\ntemplate\nMulticolorILUSolver_Base::~MulticolorILUSolver_Base()\n{\n m_LU.set_initialized(0);\n m_A_to_LU_mapping.clear();\n m_A_to_LU_mapping.shrink_to_fit();\n m_LU.resize(0, 0, 0, 1);\n}\n\ntemplate \nvoid MulticolorILUSolver >::computeAtoLUmapping()\n{\n FatalError(\"Haven't implemented Multicolor ILU smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n\ntemplate \nvoid MulticolorILUSolver >::computeAtoLUmapping()\n{\n Matrix &m_A = *this->m_explicit_A;\n const int CtaSize = 128; // Number of threads per CTA\n const int SMemSize = 128; // per warp\n const int nWarps = CtaSize / 32;\n int GridSize = std::min( AMGX_GRID_MAX_SIZE, ( this->m_explicit_A->get_num_rows( ) + nWarps - 1 ) / nWarps );\n // Global memory workspaces\n device_vector_alloc returnValue(1);\n returnValue[0] = 0;\n\n if (this->m_explicit_A->hasProps(DIAG))\n {\n computeAtoLUmappingExtDiag_kernel <<< GridSize, CtaSize >>> (\n m_A.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &m_A.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &m_A.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &m_A.diag[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_A_to_LU_mapping[0] ),\n amgx::thrust::raw_pointer_cast( &returnValue[0] ));\n }\n else\n {\n computeAtoLUmapping_kernel <<< GridSize, CtaSize >>> (\n m_A.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &m_A.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &m_A.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_A_to_LU_mapping[0] ),\n amgx::thrust::raw_pointer_cast( &returnValue[0] ));\n }\n\n cudaCheckError();\n\n // fallback path that allows 1024 nonzeros per row\n if (returnValue[0] == 1)\n {\n returnValue[0] = 0;\n const int SMemSize2 = 1024 ; // per warp\n\n if (this->m_explicit_A->hasProps(DIAG))\n {\n computeAtoLUmappingExtDiag_kernel <<< GridSize, CtaSize >>> (\n m_A.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &m_A.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &m_A.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &m_A.diag[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_A_to_LU_mapping[0] ),\n amgx::thrust::raw_pointer_cast( &returnValue[0] ));\n }\n else\n {\n computeAtoLUmapping_kernel <<< GridSize, CtaSize >>> (\n m_A.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &m_A.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &m_A.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_A_to_LU_mapping[0] ),\n amgx::thrust::raw_pointer_cast( &returnValue[0] ));\n }\n\n cudaCheckError();\n }\n\n if (returnValue[0] == 1)\n {\n FatalError( \"Number of nonzeros per row exceeds allocated shared memory\", AMGX_ERR_NO_MEMORY);\n }\n}\n\ntemplate \nvoid MulticolorILUSolver >::fillLUValuesWithAValues()\n{\n FatalError(\"Haven't implemented Multicolor ILU smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n\ntemplate \nvoid MulticolorILUSolver >::fillLUValuesWithAValues()\n{\n if (this->m_sparsity_level == 0)\n {\n this->m_LU.values = this->m_explicit_A->values;\n }\n else\n {\n // TODO: Should probably store the inverse mapping of AtoLUmapping instead\n // This will allow to use unpermuteVector and have coalesced writes\n // instead of coalesced reads\n thrust_wrapper::fill(this->m_LU.values.begin(), this->m_LU.values.end(), 0.);\n cudaCheckError();\n\n if (this->m_explicit_A->hasProps(DIAG))\n {\n amgx::permuteVector(this->m_explicit_A->values, this->m_LU.values, this->m_A_to_LU_mapping, (this->m_explicit_A->get_num_nz() + this->m_explicit_A->get_num_rows())*this->m_explicit_A->get_block_size());\n }\n else\n {\n amgx::permuteVector(this->m_explicit_A->values, this->m_LU.values, this->m_A_to_LU_mapping, this->m_explicit_A->get_num_nz()*this->m_explicit_A->get_block_size());\n }\n }\n}\n\n\ntemplate \nvoid MulticolorILUSolver >::computeLUSparsityPattern()\n{\n FatalError(\"Haven't implemented Multicolor ILU smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorILUSolver >::computeLUSparsityPattern()\n{\n // ILU0\n if (this->m_sparsity_level == 0)\n {\n // Copy everything except the values\n this->m_LU.copy_structure(*this->m_explicit_A);\n }\n // ILU1\n else if (this->m_sparsity_level == 1)\n {\n this->sparsity_wk = CSR_Multiply::csr_workspace_create( *this->m_cfg, \"default\" );\n CSR_Multiply::csr_sparsity_ilu1( *this->m_explicit_A, this->m_LU, this->sparsity_wk );\n CSR_Multiply::csr_workspace_delete( this->sparsity_wk );\n\n if (this->m_use_bsrxmv)\n {\n this->m_LU.set_initialized(0);\n this->m_LU.computeDiagonal();\n this->m_LU.set_initialized(1);\n }\n\n this->m_LU.setMatrixColoring(&(this->m_explicit_A->getMatrixColoring()));\n }\n else\n {\n FatalError(\"Haven't implemented Multicolor ILU smoother for this sparsity level. \", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid MulticolorILUSolver >::computeLUFactors()\n{\n FatalError(\"Haven't implemented Multicolor ILU smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorILUSolver >::computeLUFactors()\n{\n const int CtaSize = 128; // Number of threads per CTA\n const int SMemSize = 128;\n const int nWarps = CtaSize / 32;\n device_vector_alloc returnValue(1);\n returnValue[0] = 0;\n int num_colors = this->m_LU.getMatrixColoring().getNumColors();\n const IndexType *LU_sorted_rows_by_color_ptr = this->m_LU.getMatrixColoring().getSortedRowsByColor().raw();\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = this->m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = this->m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n#ifdef EXPERIMENTAL_LU_FACTORS\n int GridSize = std::min( 2048, ( num_rows_per_color + nWarps - 1 ) / nWarps );\n\n if ( GridSize == 0 )\n {\n continue; // if perfect coloring (color 0 has no vertices)\n }\n\n if ( this->m_LU.get_block_dimx() == 4 && this->m_LU.get_block_dimy() == 4 )\n {\n if ( this->m_explicit_A->getBlockFormat() == ROW_MAJOR )\n {\n compute_LU_factors_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n this->m_LU.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.diag[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.values[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_smaller_color_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_larger_color_offsets[0] ),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n amgx::thrust::raw_pointer_cast( &returnValue[0] ) );\n }\n else\n {\n compute_LU_factors_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n this->m_LU.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.diag[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.values[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_smaller_color_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_larger_color_offsets[0] ),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n amgx::thrust::raw_pointer_cast( &returnValue[0] ) );\n }\n\n cudaCheckError();\n }\n else\n {\n FatalError(\"Unsupported block size for Multicolor ILU solver, computeLUFactors\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n#else\n int GridSize = std::min( AMGX_GRID_MAX_SIZE, ( num_rows_per_color + nWarps - 1 ) / nWarps );\n\n if ( GridSize == 0 )\n {\n continue; // if perfect coloring (color 0 has no vertices)\n }\n\n if ( this->m_LU.get_block_dimx() == 4 && this->m_LU.get_block_dimy() == 4 )\n {\n //computeLUFactors_4x4_kernel <<< GridSize, CtaSize>>> (\n if (this->m_explicit_A->getBlockFormat() == ROW_MAJOR)\n {\n computeLUFactors_4x4_kernel <<< GridSize, CtaSize>>> (\n this->m_LU.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.diag[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.values[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_smaller_color_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_larger_color_offsets[0] ),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n amgx::thrust::raw_pointer_cast( &returnValue[0] ) );\n }\n else\n {\n computeLUFactors_4x4_kernel <<< GridSize, CtaSize>>> (\n this->m_LU.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.diag[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.values[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_smaller_color_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_larger_color_offsets[0] ),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n amgx::thrust::raw_pointer_cast( &returnValue[0] ) );\n }\n\n cudaCheckError();\n }\n else\n {\n FatalError(\"Unsupported block size for Multicolor ILU solver, computeLUFactors\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n#endif\n }\n\n // Check returnValue flag\n if ( returnValue[0] == 1 )\n {\n FatalError( \"Number of nonzeros per row exceeds allocated shared memory\", AMGX_ERR_NO_MEMORY);\n }\n}\n\n\n// Solver pre-setup\ntemplate\nvoid\nMulticolorILUSolver_Base::pre_setup()\n{\n // Check if matrix is colored\n if (this->m_explicit_A->getColoringLevel() < m_sparsity_level + 1)\n {\n FatalError(\"Matrix must be colored with coloring_level > sparsity_level for the multicolorILUsolver\", AMGX_ERR_CONFIGURATION);\n }\n\n // Compute extended sparsity pattern based on coloring and matrix A\n computeLUSparsityPattern();\n\n if (this->m_LU.hasProps(DIAG))\n {\n FatalError(\"Multicolor ILU smoother does not support outside diagonal. Try setting reorder_cols_by_color=1 and insert_diag_while_reordering=1 in the multicolor_ilu solver scope in configuration file\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (m_sparsity_level == 0 && !this->m_LU.getColsReorderedByColor())\n {\n FatalError(\"Multicolor ILU smoother requires matrix to be reordered by color with ILU0 solver. Try setting reorder_cols_by_color=1 and insert_diag_while_reordering=1 in the multicolor_ilu solver scope in configuration file\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // Reorder the columns of LU by color\n if (m_sparsity_level != 0)\n {\n // Reorder columns of LU by color\n m_LU.reorderColumnsByColor(false);\n\n // Compute mapping between entries in A and entries in LU\n if (this->m_explicit_A->hasProps(DIAG))\n {\n m_A_to_LU_mapping.resize(this->m_explicit_A->get_num_nz() + this->m_explicit_A->get_num_rows());\n }\n else\n {\n m_A_to_LU_mapping.resize(this->m_explicit_A->get_num_nz());\n }\n\n computeAtoLUmapping();\n }\n\n int N = this->m_LU.get_num_rows() * this->m_LU.get_block_dimy();\n m_delta.resize(N);\n m_Delta.resize(N);\n m_Delta.set_block_dimy(this->m_explicit_A->get_block_dimy());\n m_Delta.set_block_dimx(1);\n m_delta.set_block_dimy(this->m_explicit_A->get_block_dimy());\n m_delta.set_block_dimx(1);\n}\n\ntemplate\nvoid\nMulticolorILUSolver_Base::printSolverParameters() const\n{\n std::cout << \"relaxation_factor = \" << this->m_weight << std::endl;\n std::cout << \"use_bsrxmv = \" << this->m_use_bsrxmv << std::endl;\n std::cout << \"ilu_sparsity_level = \" << this->m_sparsity_level << std::endl;\n}\n\n\n// Solver setup\ntemplate\nvoid\nMulticolorILUSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n this->m_explicit_A = dynamic_cast*>(this->m_A);\n\n if (!this->m_explicit_A)\n {\n FatalError(\"MulticolorILUSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n if (this->m_explicit_A->getColoringLevel() < 1)\n {\n FatalError(\"Matrix must be colored to use multicolor ilu solver. Try setting: coloring_level=1 or coloring_level=2 in the configuration file\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (!reuse_matrix_structure)\n {\n this->pre_setup();\n }\n\n // Fill LU sparsity pattern\n fillLUValuesWithAValues();\n // Compute LU factors in place (update LU.values)\n computeLUFactors();\n}\n\n//\ntemplate\nvoid\nMulticolorILUSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nMulticolorILUSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n if ( !m_use_bsrxmv && (this->m_LU.get_block_dimx() == 4 && this->m_LU.get_block_dimy() == 4) )\n {\n smooth_4x4(b, x, xIsZero);\n }\n else\n {\n smooth_bxb(b, x, xIsZero);\n }\n\n // Do we converge ?\n return this->converged(b, x);\n}\n\ntemplate\nvoid\nMulticolorILUSolver_Base::solve_finalize( VVector &b, VVector &x )\n{}\n\ntemplate \nvoid MulticolorILUSolver >::smooth_4x4(const VVector &b, VVector &x, bool xIsZero)\n{\n FatalError(\"Haven't implemented Multicolor DILU smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorILUSolver >::smooth_4x4(const VVector &b, VVector &x, bool xIsZero)\n{\n Matrix &m_LU = this->m_LU;\n Matrix &m_A = *this->m_explicit_A;\n int N = m_LU.get_num_rows() * m_LU.get_block_dimy();\n cudaCheckError();\n\n if (!m_LU.getColsReorderedByColor())\n {\n FatalError(\"ILU solver currently only works if columns are reordered by color. Try setting reordering_cols_by_color=1 in the multicolor_ilu solver scope in the configuration file\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // ---------------------------------------------------------\n // Solving Lower triangular system, with identity diagonal\n // ---------------------------------------------------------\n const IndexType *LU_sorted_rows_by_color_ptr = m_LU.getMatrixColoring().getSortedRowsByColor().raw();\n int num_colors = this->m_LU.getMatrixColoring().getNumColors();\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n#ifdef EXPERIMENTAL_LU_FORWARD\n const int CtaSize = 128; // Number of threads per CTA\n const int nHalfWarps = CtaSize / 16;\n int GridSize = std::min( 2048, ( num_rows_per_color + nHalfWarps - 1 ) / nHalfWarps );\n\n if ( GridSize == 0 )\n {\n continue; // if perfect coloring (color 0 has no vertices)\n }\n\n if ( this->m_explicit_A->getBlockFormat() == ROW_MAJOR )\n {\n if (m_A.hasProps(DIAG))\n {\n LU_forward_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n m_LU.row_offsets.raw(),\n m_LU.m_smaller_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.values.raw(),\n m_A.row_offsets.raw(),\n m_A.col_indices.raw(),\n m_A.values.raw(),\n m_A.diag.raw(),\n x.raw(),\n b.raw(),\n this->m_delta.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n xIsZero );\n }\n else\n {\n LU_forward_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n m_LU.row_offsets.raw(),\n m_LU.m_smaller_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.values.raw(),\n m_A.row_offsets.raw(),\n m_A.col_indices.raw(),\n m_A.values.raw(),\n m_A.diag.raw(),\n x.raw(),\n b.raw(),\n this->m_delta.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n xIsZero );\n }\n }\n else\n {\n // COL_MAJOR\n if (m_A.hasProps(DIAG))\n {\n LU_forward_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n m_LU.row_offsets.raw(),\n m_LU.m_smaller_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.values.raw(),\n m_A.row_offsets.raw(),\n m_A.col_indices.raw(),\n m_A.values.raw(),\n m_A.diag.raw(),\n x.raw(),\n b.raw(),\n this->m_delta.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n xIsZero );\n }\n else\n {\n LU_forward_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n m_LU.row_offsets.raw(),\n m_LU.m_smaller_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.values.raw(),\n m_A.row_offsets.raw(),\n m_A.col_indices.raw(),\n m_A.values.raw(),\n m_A.diag.raw(),\n x.raw(),\n b.raw(),\n this->m_delta.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n xIsZero );\n }\n }\n\n#else\n const int CtaSize = 128;\n const int blockrows_per_cta = CtaSize / 4;\n const int GridSize = min( AMGX_GRID_MAX_SIZE, (int) (num_rows_per_color + blockrows_per_cta - 1) / blockrows_per_cta);\n\n if ( GridSize == 0 )\n {\n continue; // if perfect coloring (color 0 has no vertices)\n }\n\n if (this->m_explicit_A->hasProps(DIAG))\n {\n FatalError(\"this implementation of LU forward solve does not support A with external diagonal\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (this->m_explicit_A->getBlockFormat() == ROW_MAJOR)\n {\n LU_forward_4x4_kernel <<< GridSize, CtaSize>>>\n (m_LU.row_offsets.raw(),\n m_LU.m_smaller_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.values.raw(),\n m_A.row_offsets.raw(),\n m_A.col_indices.raw(),\n m_A.values.raw(),\n x.raw(),\n b.raw(),\n delta.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n xIsZero);\n }\n else\n {\n LU_forward_4x4_kernel <<< GridSize, CtaSize>>>\n (m_LU.row_offsets.raw(),\n m_LU.m_smaller_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.values.raw(),\n m_A.row_offsets.raw(),\n m_A.col_indices.raw(),\n m_A.values.raw(),\n x.raw(),\n b.raw(),\n delta.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n xIsZero);\n }\n\n#endif\n cudaCheckError();\n }\n\n // --------------------\n // Backward Sweep\n // --------------------\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n#ifdef EXPERIMENTAL_LU_BACKWARD\n const int CtaSize = 128; // Number of threads per CTA\n const int nHalfWarps = CtaSize / 16;\n int GridSize = std::min( 2048, ( num_rows_per_color + nHalfWarps - 1 ) / nHalfWarps );\n\n if ( GridSize == 0 )\n {\n continue; // if perfect coloring (color 0 has no vertices)\n }\n\n if (this->m_explicit_A->getBlockFormat() == ROW_MAJOR)\n {\n LU_backward_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n m_LU.row_offsets.raw(),\n m_LU.m_larger_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.diag.raw(),\n m_LU.values.raw(),\n this->m_delta.raw(),\n this->m_Delta.raw(),\n x.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n num_colors,\n this->m_weight,\n xIsZero);\n }\n else\n {\n LU_backward_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n m_LU.row_offsets.raw(),\n m_LU.m_larger_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.diag.raw(),\n m_LU.values.raw(),\n this->m_delta.raw(),\n this->m_Delta.raw(),\n x.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n num_colors,\n this->m_weight,\n xIsZero);\n }\n\n#else\n const int CtaSize = 128;\n const int blockrows_per_cta = CtaSize / 4;\n const int GridSize = min( AMGX_GRID_MAX_SIZE, (int) (num_rows_per_color + blockrows_per_cta - 1) / blockrows_per_cta);\n\n if ( GridSize == 0 )\n {\n continue; // if perfect coloring (color 0 has no vertices)\n }\n\n if (this->m_explicit_A->getBlockFormat() == ROW_MAJOR)\n {\n LU_backward_4x4_kernel <<< GridSize, CtaSize>>>\n (m_LU.row_offsets.raw(),\n m_LU.m_larger_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.diag.raw(),\n m_LU.values.raw(),\n this->m_delta.raw(),\n this->m_Delta.raw(),\n x.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n num_colors,\n this->m_weight,\n xIsZero);\n }\n else\n {\n LU_backward_4x4_kernel <<< GridSize, CtaSize>>>\n (m_LU.row_offsets.raw(),\n m_LU.m_larger_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.diag.raw(),\n m_LU.values.raw(),\n this->m_delta.raw(),\n this->m_Delta.raw(),\n x.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n num_colors,\n this->m_weight,\n xIsZero);\n }\n\n#endif\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorILUSolver >::smooth_bxb(const VVector &b, VVector &x, bool xIsZero)\n{\n FatalError(\"Haven't implemented Multicolor DILU smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorILUSolver >::smooth_bxb(const VVector &b, VVector &x, bool xIsZero)\n{\n Matrix &m_LU = this->m_LU;\n Matrix &m_A = *this->m_explicit_A;\n int N = m_LU.get_num_rows() * m_LU.get_block_dimy();\n\n if (!m_LU.getColsReorderedByColor())\n {\n FatalError(\"ILU solver currently only works if columns are reordered by color. Try setting reorder_cols_by_color=1 in the multicolor_ilu solver scope in the configuration file\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (this->m_explicit_A->getBlockFormat() == COL_MAJOR)\n {\n FatalError(\"ILU solver for arbitrary block sizes only works with ROW_MAJOR matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // ---------------------------------------------------------\n // Solving Lower triangular system, with identity diagonal\n // ---------------------------------------------------------\n const IndexType *LU_sorted_rows_by_color_ptr = m_LU.getMatrixColoring().getSortedRowsByColor().raw();\n int num_colors = this->m_LU.getMatrixColoring().getNumColors();\n //delta = b;\n amgx::thrust::copy(b.begin(), b.end(), this->m_delta.begin());\n //delta = delta - Ax;\n Cusparse::bsrmv((ValueTypeA) - 1.0, m_A, x, (ValueTypeA)1.0, this->m_delta);\n cudaCheckError();\n // Setting Delta to zero\n thrust_wrapper::fill(this->m_Delta.begin(), this->m_Delta.end(), (ValueTypeB)0.0f);\n cudaCheckError();\n bool skipped_end = false;\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n\n if (num_rows_per_color == 0) { continue; } // if perfect coloring (color 0 has no vertices)\n\n if (skipped_end)\n {\n // delta = delta - LU*Delta smaller colors\n Cusparse::bsrmv(Cusparse::SMALLER_COLORS, i, (ValueTypeA) - 1.0f, m_LU, this->m_delta, (ValueTypeA)1.0f, this->m_delta);\n }\n\n if (num_rows_per_color > 0)\n {\n skipped_end = true;\n }\n }\n\n cudaCheckError();\n skipped_end = false;\n\n // --------------------\n // Backward Sweep\n // --------------------\n for (int i = num_colors - 1; i >= 0; i--)\n {\n // delta = delta - LU*Delta larger colors\n Cusparse::bsrmv(Cusparse::LARGER_COLORS, i, (ValueTypeA) - 1.0f, m_LU, this->m_Delta, (ValueTypeA)1.0f, this->m_delta);\n // Multiple by inverse stored on diagonal\n Cusparse::bsrmv(Cusparse::DIAG_COL, i, (ValueTypeA) 1.0f, m_LU, this->m_delta, 0.0f, this->m_Delta);\n }\n\n cudaCheckError();\n axpy(this->m_Delta, x, this->m_weight, 0, x.size());\n cudaCheckError();\n}\n\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class MulticolorILUSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class MulticolorILUSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n// TODO: Have 2 groups of 16 threads collaborate\n// TODO: Add support for outside diagonal\n// TODO: Add support for unsorted rows\n\n#define EXPERIMENTAL_LU_FACTORS\n#define EXPERIMENTAL_LU_FORWARD\n#define EXPERIMENTAL_LU_BACKWARD\n\nnamespace amgx\n{\n\nnamespace multicolor_ilu_solver\n{\n\n// -----------\n// Kernels\n// -----------\n\n#ifdef EXPERIMENTAL_LU_FORWARD\n\ntemplate\n__global__\n__launch_bounds__( CtaSize, 16 )\nvoid LU_forward_4x4_kernel_warp( const IndexType *LU_row_offsets,\n const IndexType *LU_smaller_color_offsets,\n const IndexType *LU_column_indices,\n const ValueTypeA *LU_nonzero_values,\n const IndexType *A_row_offsets,\n const IndexType *A_column_indices,\n const ValueTypeA *A_nonzero_values,\n const IndexType *A_dia_indices,\n const ValueTypeB *x,\n const ValueTypeB *b,\n ValueTypeB *delta,\n const int *sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n bool xIsZero )\n{\n const int nHalfWarps = CtaSize / 16; // Number of half warps per Cta\n const int warpId = utils::warp_id();\n const int laneId = utils::lane_id();\n const int halfWarpId = threadIdx.x / 16;\n const int halfLaneId = threadIdx.x % 16;\n const int halfLaneId_div_4 = halfLaneId / 4;\n const int halfLaneId_mod_4 = halfLaneId % 4;\n const int upperHalf = 16 * (laneId / 16);\n // Shared memory needed to exchange X and delta.\n __shared__ volatile ValueTypeB s_mem[CtaSize];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile ValueTypeB *my_s_mem = &s_mem[16 * halfWarpId];\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( int aRowIt = blockIdx.x * nHalfWarps + halfWarpId ; aRowIt < num_rows_per_color ; aRowIt += gridDim.x * nHalfWarps )\n {\n int aRowId = sorted_rows_by_color[aRowIt];\n // Load one block of B.\n ValueTypeB my_bmAx(0);\n\n unsigned int active_mask = utils::activemask();\n\n if ( ROW_MAJOR )\n {\n if ( halfLaneId_mod_4 == 0 )\n {\n my_bmAx = __cachingLoad(&b[4 * aRowId + halfLaneId_div_4]);\n }\n }\n else\n {\n if ( halfLaneId_div_4 == 0 )\n {\n my_bmAx = __cachingLoad(&b[4 * aRowId + halfLaneId_mod_4]);\n }\n }\n\n // Don't do anything if X is zero.\n if ( !xIsZero )\n {\n int aColBegin = A_row_offsets[aRowId ];\n int aColEnd = A_row_offsets[aRowId + 1];\n int aColMax = aColEnd;\n\n if ( hasDiag )\n {\n ++aColMax;\n }\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( aColBegin < aColMax, active_mask ) ; aColBegin += 16 )\n {\n int aColIt = aColBegin + halfLaneId;\n // Get the ID of the column.\n int aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColId = A_column_indices[aColIt];\n }\n\n if ( hasDiag && aColIt == aColEnd )\n {\n aColId = aRowId;\n }\n\n // Count the number of active columns.\n int vote = utils::ballot(aColId != -1, active_mask);\n // The number of iterations.\n int nCols = max( __popc( vote & 0x0000ffff ), __popc( vote & 0xffff0000 ) );\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < nCols ; k += 4 )\n {\n int my_k = k + halfLaneId_div_4;\n // Load 8 blocks of X.\n int waColId = utils::shfl( aColId, upperHalf + my_k, warpSize, active_mask );\n ValueTypeB my_x(0);\n\n if ( waColId != -1 )\n {\n my_x = __cachingLoad(&x[4 * waColId + halfLaneId_mod_4]);\n }\n\n my_s_mem[halfLaneId] = my_x;\n // Load 8 blocks of A.\n#pragma unroll\n for ( int i = 0 ; i < 4 ; ++i )\n {\n int w_aColTmp = aColBegin + k + i, w_aColIt = -1;\n\n if ( w_aColTmp < aColEnd )\n {\n w_aColIt = w_aColTmp;\n }\n\n if ( hasDiag && w_aColTmp == aColEnd )\n {\n w_aColIt = A_dia_indices[aRowId];\n }\n\n ValueTypeA my_val(0);\n\n if ( w_aColIt != -1 )\n {\n my_val = A_nonzero_values[16 * w_aColIt + halfLaneId];\n }\n\n if ( ROW_MAJOR )\n {\n my_bmAx -= my_val * my_s_mem[4 * i + halfLaneId_mod_4];\n }\n else\n {\n my_bmAx -= my_val * my_s_mem[4 * i + halfLaneId_div_4];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n } // if xIsZero\n\n // Contribution from each nonzero column that has color less than yours\n if ( current_color != 0 )\n {\n // TODO: Use constant or texture here\n int aColBegin = LU_row_offsets[aRowId];\n int aColEnd = LU_smaller_color_offsets[aRowId];\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( aColBegin < aColEnd, active_mask ) ; aColBegin += 16 )\n {\n int aColIt = aColBegin + halfLaneId;\n int aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColId = LU_column_indices[aColIt];\n }\n\n // Count the number of active columns.\n int vote = utils::ballot(aColId != -1, active_mask);\n // The number of iterations.\n int nCols = max( __popc( vote & 0x0000ffff ), __popc( vote & 0xffff0000 ) );\n\n for ( int k = 0 ; k < nCols ; k += 4 )\n {\n int my_k = k + halfLaneId_div_4;\n // Load 8 blocks of X.\n int waColId = utils::shfl( aColId, upperHalf + my_k, warpSize, active_mask );\n ValueTypeB my_delta(0);\n\n if ( waColId != -1 )\n {\n my_delta = delta[4 * waColId + halfLaneId_mod_4];\n }\n\n my_s_mem[halfLaneId] = my_delta;\n utils::syncwarp(); // making sure smem write propagated\n // Update b-Ax.\n#pragma unroll\n for ( int i = 0 ; i < 4 ; ++i )\n {\n int w_aColTmp = aColBegin + k + i, w_aColIt = -1;\n\n if ( w_aColTmp < aColEnd )\n {\n w_aColIt = w_aColTmp;\n }\n\n ValueTypeA my_val(0);\n\n if ( w_aColIt != -1 )\n {\n my_val = LU_nonzero_values[16 * w_aColIt + halfLaneId];\n }\n\n if ( ROW_MAJOR )\n {\n my_bmAx -= my_val * my_s_mem[4 * i + halfLaneId_mod_4];\n }\n else\n {\n my_bmAx -= my_val * my_s_mem[4 * i + halfLaneId_div_4];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n } // If current_color != 0\n\n // Reduce bmAx terms.\n if ( ROW_MAJOR )\n {\n my_bmAx += utils::shfl_xor( my_bmAx, 1, warpSize, active_mask );\n my_bmAx += utils::shfl_xor( my_bmAx, 2, warpSize, active_mask );\n }\n else\n {\n my_bmAx += utils::shfl_xor( my_bmAx, 4, warpSize, active_mask );\n my_bmAx += utils::shfl_xor( my_bmAx, 8, warpSize, active_mask );\n }\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n if ( halfLaneId_mod_4 == 0 )\n {\n delta[4 * aRowId + halfLaneId_div_4] = my_bmAx;\n }\n }\n else\n {\n if ( halfLaneId_div_4 == 0 )\n {\n delta[4 * aRowId + halfLaneId_mod_4] = my_bmAx;\n }\n }\n }\n}\n\n#else\ntemplate\n__global__\nvoid LU_forward_4x4_kernel(const IndexType *LU_row_offsets, const IndexType *LU_smaller_color_offsets, const IndexType *LU_column_indices, const ValueTypeA *LU_nonzero_values, const IndexType *A_row_offsets, const IndexType *A_column_indices, const ValueTypeA *A_nonzero_values,\n const ValueTypeB *x, const ValueTypeB *b, ValueTypeB *delta, const int *sorted_rows_by_color,\n const int num_rows_per_color, const int current_color, bool xIsZero)\n\n\n{\n int warp_id = threadIdx.x / 32;\n int warp_thread_id = threadIdx.x & 31;\n\n // padding row blocks to fit in a single warp\n if ( warp_thread_id >= blockrows_per_warp * bsize ) { return; }\n\n // new thread id with padding\n int tid = warp_id * blockrows_per_warp * bsize + warp_thread_id;\n // Here we use one thread per row (not block row)\n int cta_blockrow_id = (tid) / bsize;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int vec_entry_index = tid - cta_blockrow_id * bsize;\n volatile __shared__ ValueTypeB s_delta_temp[ bsize * blockrows_per_cta];\n int offset, s_offset, i;\n ValueTypeB bmAx, temp[bsize];\n\n while (blockrow_id < num_rows_per_color && cta_blockrow_id < blockrows_per_cta)\n {\n i = sorted_rows_by_color[blockrow_id];\n // Load RHS and x\n offset = i * bsize + vec_entry_index;\n bmAx = b[offset];\n\n if (!xIsZero)\n {\n int jmin = A_row_offsets[i];\n int jmax = A_row_offsets[i + 1];\n\n //TODO: Assumes inside diagonal\n for (int jind = jmin; jind < jmax; jind++)\n {\n IndexType jcol = A_column_indices[jind];\n offset = jcol * bsize + vec_entry_index;\n s_delta_temp[tid] = x[offset];\n\n // Load nonzero_values\n if (ROW_MAJOR)\n {\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(A_nonzero_values + offset, temp);\n }\n else\n {\n offset = jind * bsize * bsize + vec_entry_index;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n temp[m] = A_nonzero_values[offset + bsize * m];\n }\n }\n\n // Do matrix multiply\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_delta_temp[s_offset++];\n }\n }\n }\n\n // Contribution from each nonzero column that has color less than yours\n if (current_color != 0)\n {\n int jmin = LU_row_offsets[i];\n int jmax = LU_smaller_color_offsets[i];\n\n for (int jind = jmin; jind < jmax; jind++)\n {\n IndexType jcol = LU_column_indices[jind];\n offset = jcol * bsize + vec_entry_index;\n s_delta_temp[tid] = __ldcg(delta + offset);\n\n // Load nonzero_values\n if (ROW_MAJOR)\n {\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(LU_nonzero_values + offset, temp);\n }\n else\n {\n offset = jind * bsize * bsize + vec_entry_index;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n temp[m] = LU_nonzero_values[offset + bsize * m];\n }\n }\n\n // Do matrix multiply\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_delta_temp[s_offset++];\n }\n }\n }\n\n delta[i * bsize + vec_entry_index] = bmAx;\n blockrow_id += blockrows_per_cta * gridDim.x;\n }\n}\n#endif\n\n#ifdef EXPERIMENTAL_LU_BACKWARD\n\ntemplate< typename IndexType, typename ValueTypeA, typename ValueTypeB, int CtaSize, bool ROW_MAJOR >\n__global__\n__launch_bounds__( CtaSize, 16 )\nvoid LU_backward_4x4_kernel_warp( const IndexType *row_offsets,\n const IndexType *larger_color_offsets,\n const IndexType *column_indices,\n const IndexType *dia_indices,\n const ValueTypeA *nonzero_values,\n const ValueTypeB *delta,\n ValueTypeB *Delta,\n ValueTypeB *x,\n const int *sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int num_colors,\n const ValueTypeB weight,\n bool xIsZero )\n{\n const int nHalfWarps = CtaSize / 16; // Number of half warps per CTA.\n const int warpId = utils::warp_id();\n const int laneId = utils::lane_id();\n const int halfWarpId = threadIdx.x / 16;\n const int halfLaneId = threadIdx.x % 16;\n const int halfLaneId_div_4 = halfLaneId / 4;\n const int halfLaneId_mod_4 = halfLaneId % 4;\n const int upperHalf = 16 * (laneId / 16);\n // Shared memory needed to exchange X and delta.\n __shared__ volatile ValueTypeB s_mem[CtaSize];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile ValueTypeB *my_s_mem = &s_mem[16 * halfWarpId];\n\n // Iterate over the rows of the matrix. One warp per two rows.\n for ( int aRowIt = blockIdx.x * nHalfWarps + halfWarpId ; aRowIt < num_rows_per_color ; aRowIt += gridDim.x * nHalfWarps )\n {\n int aRowId = sorted_rows_by_color[aRowIt];\n unsigned int active_mask = utils::activemask();\n // Load one block of B.\n ValueTypeB my_bmAx(0);\n\n if ( ROW_MAJOR )\n {\n if ( halfLaneId_mod_4 == 0 )\n {\n my_bmAx = delta[4 * aRowId + halfLaneId_div_4];\n }\n }\n else\n {\n if ( halfLaneId_div_4 == 0 )\n {\n my_bmAx = delta[4 * aRowId + halfLaneId_mod_4];\n }\n }\n\n // Don't do anything if the color is not the interesting one.\n if ( current_color != num_colors - 1 )\n {\n // The range of the rows.\n int aColBegin = larger_color_offsets[aRowId], aColEnd = row_offsets[aRowId + 1];\n\n // Each warp load column indices of 16 nonzero blocks\n for ( ; utils::any( aColBegin < aColEnd, active_mask ) ; aColBegin += 16 )\n {\n int aColIt = aColBegin + halfLaneId;\n // Get the ID of the column.\n int aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColId = column_indices[aColIt];\n }\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < 16 ; k += 4 )\n {\n int my_k = k + halfLaneId_div_4;\n // Exchange column indices.\n int waColId = utils::shfl( aColId, upperHalf + my_k, warpSize, active_mask );\n // Load 8 blocks of X if needed.\n ValueTypeB *my_ptr = Delta;\n\n if ( xIsZero )\n {\n my_ptr = x;\n }\n\n ValueTypeB my_x(0);\n\n if ( waColId != -1 )\n {\n my_x = my_ptr[4 * waColId + halfLaneId_mod_4];\n }\n\n my_s_mem[halfLaneId] = my_x;\n utils::syncwarp();\n // Load 8 blocks of A.\n#pragma unroll\n for ( int i = 0 ; i < 4 ; ++i )\n {\n int w_aColTmp = aColBegin + k + i, w_aColIt = -1;\n\n if ( w_aColTmp < aColEnd )\n {\n w_aColIt = w_aColTmp;\n }\n\n ValueTypeA my_val(0);\n\n if ( w_aColIt != -1 )\n {\n my_val = nonzero_values[16 * w_aColIt + halfLaneId];\n }\n\n if ( ROW_MAJOR )\n {\n my_bmAx -= my_val * my_s_mem[4 * i + halfLaneId_mod_4];\n }\n else\n {\n my_bmAx -= my_val * my_s_mem[4 * i + halfLaneId_div_4];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Reduce bmAx terms.\n if ( ROW_MAJOR )\n {\n my_bmAx += utils::shfl_xor( my_bmAx, 1, warpSize, active_mask );\n my_bmAx += utils::shfl_xor( my_bmAx, 2, warpSize, active_mask );\n }\n else\n {\n my_bmAx += utils::shfl_xor( my_bmAx, 4, warpSize, active_mask );\n my_bmAx += utils::shfl_xor( my_bmAx, 8, warpSize, active_mask );\n }\n } // if current_color != num_colors-1\n\n // Update the shared terms.\n if ( ROW_MAJOR )\n {\n if ( halfLaneId_mod_4 == 0 )\n {\n my_s_mem[halfLaneId_div_4] = my_bmAx;\n }\n }\n else\n {\n if ( halfLaneId_div_4 == 0 )\n {\n my_s_mem[halfLaneId_mod_4] = my_bmAx;\n }\n }\n\n // Update the diagonal term.\n int w_aColIt = dia_indices[aRowId];\n ValueTypeA my_val(0);\n utils::syncwarp();\n\n if ( w_aColIt != -1 )\n {\n my_val = nonzero_values[16 * w_aColIt + halfLaneId];\n }\n\n if ( ROW_MAJOR )\n {\n my_bmAx = my_val * my_s_mem[halfLaneId_mod_4];\n }\n else\n {\n my_bmAx = my_val * my_s_mem[halfLaneId_div_4];\n }\n\n // Regroup results.\n if ( ROW_MAJOR )\n {\n my_bmAx += utils::shfl_xor( my_bmAx, 1 );\n my_bmAx += utils::shfl_xor( my_bmAx, 2 );\n }\n else\n {\n my_bmAx += utils::shfl_xor( my_bmAx, 4 );\n my_bmAx += utils::shfl_xor( my_bmAx, 8 );\n }\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n ValueTypeB my_x(0);\n\n if ( !xIsZero && halfLaneId_mod_4 == 0 )\n {\n my_x = x[4 * aRowId + halfLaneId_div_4];\n }\n\n my_x += weight * my_bmAx;\n\n if ( !xIsZero && halfLaneId_mod_4 == 0 )\n {\n Delta[4 * aRowId + halfLaneId_div_4] = my_bmAx;\n }\n\n if ( halfLaneId_mod_4 == 0 )\n {\n x[4 * aRowId + halfLaneId_div_4] = my_x;\n }\n }\n else\n {\n ValueTypeB my_x(0);\n\n if ( !xIsZero && halfLaneId_div_4 == 0 )\n {\n my_x = x[4 * aRowId + halfLaneId_mod_4];\n }\n\n my_x += weight * my_bmAx;\n\n if ( !xIsZero && halfLaneId_div_4 == 0 )\n {\n Delta[4 * aRowId + halfLaneId_mod_4] = my_bmAx;\n }\n\n if ( halfLaneId_div_4 == 0 )\n {\n x[4 * aRowId + halfLaneId_mod_4] = my_x;\n }\n }\n }\n}\n\n#else\n\ntemplate\n__global__\nvoid LU_backward_4x4_kernel(const IndexType *row_offsets, const IndexType *larger_color_offsets, const IndexType *column_indices, const IndexType *dia_indices, const ValueTypeA *nonzero_values,\n const ValueTypeB *delta, ValueTypeB *Delta, ValueTypeB *x, const int *sorted_rows_by_color,\n const int num_rows_per_color, const int current_color, const int num_colors, const ValueTypeB weight, bool xIsZero)\n\n\n{\n int warp_id = threadIdx.x / 32;\n int warp_thread_id = threadIdx.x & 31;\n\n // padding row blocks to fit in a single warp\n if ( warp_thread_id >= blockrows_per_warp * bsize ) { return; }\n\n // new thread id with padding\n int tid = warp_id * blockrows_per_warp * bsize + warp_thread_id;\n // Here we use one thread per row (not block row)\n int cta_blockrow_id = (tid) / bsize;\n int blockrow_id = blockIdx.x * blockrows_per_cta + cta_blockrow_id;\n const int vec_entry_index = tid - cta_blockrow_id * bsize;\n volatile __shared__ ValueTypeB s_x_temp[ bsize * blockrows_per_cta];\n int offset, s_offset, i;\n ValueTypeB bmAx, temp[bsize];\n\n while (blockrow_id < num_rows_per_color && cta_blockrow_id < blockrows_per_cta)\n {\n i = sorted_rows_by_color[blockrow_id];\n // Load RHS and x\n offset = i * bsize + vec_entry_index;\n bmAx = delta[offset];\n\n // Contribution from each nonzero column that has color less than yours\n if (current_color != num_colors)\n {\n int jmin = larger_color_offsets[i];\n int jmax = row_offsets[i + 1];\n\n for (int jind = jmin; jind < jmax; jind++)\n {\n IndexType jcol = column_indices[jind];\n offset = jcol * bsize + vec_entry_index;\n\n if (xIsZero)\n {\n s_x_temp[tid] = __ldcg(x + offset);\n }\n else\n {\n s_x_temp[tid] = __ldcg(Delta + offset);\n }\n\n // Load nonzero_values\n if (ROW_MAJOR)\n {\n offset = jind * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n }\n else\n {\n offset = jind * bsize * bsize + vec_entry_index;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n temp[m] = nonzero_values[offset + bsize * m];\n }\n }\n\n // Do matrix multiply\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n bmAx -= temp[m] * s_x_temp[s_offset++];\n }\n }\n }\n\n s_x_temp[tid] = bmAx;\n bmAx = 0.;\n\n // Load diagonals (which store the inverse)\n if (ROW_MAJOR)\n {\n offset = dia_indices[i] * bsize * bsize + vec_entry_index * bsize;\n loadAsVector(nonzero_values + offset, temp);\n }\n else\n {\n offset = dia_indices[i] * bsize * bsize + vec_entry_index;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n temp[m] = nonzero_values[offset + bsize * m];\n }\n }\n\n // Do matrix-vector multiply\n s_offset = cta_blockrow_id * bsize;\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n bmAx += temp[m] * s_x_temp[s_offset++];\n }\n\n offset = i * bsize + vec_entry_index;\n\n if (xIsZero)\n {\n x[offset] = weight * bmAx;\n }\n else\n {\n Delta[offset] = bmAx;\n x[offset] += weight * bmAx ;\n }\n\n blockrow_id += blockrows_per_cta * gridDim.x;\n }\n}\n\n#endif\n\n// Assumptions:\n// CtaSize must be multiple of 32\n// SMemSize should be larger than the maximum number of columns in the matrix\n// Matrix B is superset of matrix A\n\ntemplate< int CtaSize, int SMemSize>\n__global__ __launch_bounds__( CtaSize )\nvoid\ncomputeAtoLUmapping_kernel( int A_nRows,\n const int *__restrict A_row_offsets,\n const int *__restrict A_col_indices,\n const int *__restrict B_row_offsets,\n const int *__restrict B_col_indices,\n int *__restrict AtoBmapping,\n int *wk_returnValue )\n{\n const int nWarps = CtaSize / 32; // Number of warps per Cta\n const int warpId = utils::warp_id();\n const int laneId = utils::lane_id();\n // Rows are stored in SMEM. Linear storage.\n __shared__ volatile int s_colInd[nWarps][SMemSize];\n // The row this warp is responsible for\n int aRowId = blockIdx.x * nWarps + warpId;\n\n // Loop over rows of A.\n for ( ; aRowId < A_nRows ; aRowId += nWarps * gridDim.x )\n {\n // Insert all the column indices of matrix B in the shared memory table\n int bColBeg = B_row_offsets[aRowId];\n int bColEnd = B_row_offsets[aRowId + 1];\n // The number of columns.\n const int nCols = bColEnd - bColBeg;\n\n //TODO: Add fallback for cases where number of nonzeros exceed SMemSize\n if ( nCols > SMemSize )\n {\n wk_returnValue[0] = 1;\n return;\n }\n\n // Fill-in the local table.\n const int NUM_STEPS = SMemSize / 32;\n#pragma unroll\n for ( int step = 0, k = laneId ; step < NUM_STEPS ; ++step, k += 32 )\n {\n int bColIt = bColBeg + k;\n int bColId = -1;\n\n if ( bColIt < bColEnd )\n {\n bColId = B_col_indices[bColIt];\n }\n\n s_colInd[warpId][k] = bColId;\n }\n\n // Now load column indices of current row of A\n int aColIt = A_row_offsets[aRowId];\n int aColEnd = A_row_offsets[aRowId + 1];\n\n for ( aColIt += laneId ; utils::any(aColIt < aColEnd) ; aColIt += 32 )\n {\n // The column.\n int aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColId = A_col_indices[aColIt];\n }\n\n // Each thread searches for its column id, and gets the corresponding bColIt\n // TODO: Try binary search or using hash table\n int foundOffset = -1;\n\n if ( aColId == -1 )\n {\n foundOffset = -2;\n }\n\n for ( int i = 0 ; i < nCols && utils::any(foundOffset == -1) ; ++i )\n if ( foundOffset == -1 && s_colInd[warpId][i] == aColId )\n {\n foundOffset = i;\n }\n\n // Store the result.\n if ( aColIt < aColEnd )\n {\n AtoBmapping[aColIt] = bColBeg + foundOffset;\n }\n }\n } // if RowId < A_nRows;\n}\n\ntemplate< int CtaSize, int SMemSize>\n__global__ __launch_bounds__( CtaSize )\nvoid\ncomputeAtoLUmappingExtDiag_kernel( int A_nRows,\n const int *__restrict A_row_offsets,\n const int *__restrict A_col_indices,\n const int *__restrict A_dia_indices,\n const int *__restrict B_row_offsets,\n const int *__restrict B_col_indices,\n int *__restrict AtoBmapping,\n int *wk_returnValue )\n{\n const int nWarps = CtaSize / 32; // Number of warps per Cta\n const int warpId = utils::warp_id();\n const int laneId = utils::lane_id();\n // Rows are stored in SMEM. Linear storage.\n __shared__ volatile int s_colInd[nWarps][SMemSize];\n // The row this warp is responsible for\n int aRowId = blockIdx.x * nWarps + warpId;\n\n // Loop over rows of A.\n for ( ; aRowId < A_nRows ; aRowId += nWarps * gridDim.x )\n {\n // Insert all the column indices of matrix B in the shared memory table\n int bColBeg = B_row_offsets[aRowId];\n int bColEnd = B_row_offsets[aRowId + 1];\n // The number of columns.\n const int nCols = bColEnd - bColBeg;\n\n //TODO: Add fallback for cases where number of nonzeros exceed SMemSize\n if ( nCols > SMemSize )\n {\n wk_returnValue[0] = 1;\n return;\n }\n\n // Fill-in the local table.\n const int NUM_STEPS = SMemSize / 32;\n#pragma unroll\n for ( int step = 0, k = laneId ; step < NUM_STEPS ; ++step, k += 32 )\n {\n int bColIt = bColBeg + k;\n int bColId = -1;\n\n if ( bColIt < bColEnd )\n {\n bColId = B_col_indices[bColIt];\n }\n\n s_colInd[warpId][k] = bColId;\n }\n\n // Now load column indices of current row of A\n int aColIt = A_row_offsets[aRowId];\n int aColEnd = A_row_offsets[aRowId + 1];\n\n for ( aColIt += laneId ; utils::any(aColIt <= aColEnd) ; aColIt += 32 )\n {\n // The column.\n int aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColId = A_col_indices[aColIt];\n }\n\n if ( aColIt == aColEnd )\n {\n aColId = aRowId;\n }\n\n // Each thread searches for its column id, and gets the corresponding bColIt\n // TODO: Try binary search or using hash table\n int foundOffset = -1;\n\n if ( aColId == -1 )\n {\n foundOffset = -2;\n }\n\n for ( int i = 0 ; i < nCols && utils::any(foundOffset == -1) ; ++i )\n if ( foundOffset == -1 && s_colInd[warpId][i] == aColId )\n {\n foundOffset = i;\n }\n\n // Store the result.\n int aDst = -1;\n\n if ( aColIt < aColEnd )\n {\n aDst = aColIt;\n }\n\n if ( aColIt == aColEnd )\n {\n aDst = A_dia_indices[aRowId];\n }\n\n if ( aDst != -1 )\n {\n AtoBmapping[aDst] = bColBeg + foundOffset;\n }\n }\n }\n}\n\n#ifdef EXPERIMENTAL_LU_FACTORS\n\ntemplate< typename ValueTypeA, int CtaSize, int SMemSize, bool ROW_MAJOR >\n__global__\n__launch_bounds__( CtaSize, 12 )\nvoid\ncompute_LU_factors_4x4_kernel_warp( int A_nRows,\n const int *__restrict A_row_offsets,\n const int *__restrict A_col_indices,\n const int *__restrict A_dia_indices,\n ValueTypeA *__restrict A_nonzero_values,\n const int *__restrict A_smaller_color_offsets,\n const int *__restrict A_larger_color_offsets,\n const int *sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n int *wk_returnValue )\n{\n const int nWarps = CtaSize / 32; // Number of warps per Cta\n const int warpId = utils::warp_id();\n const int laneId = utils::lane_id();\n // Lane ID in the 2 16-wide segments.\n const int lane_id_div_16 = laneId / 16;\n const int lane_id_mod_16 = laneId % 16;\n // Coordinates inside a 4x4 block of the matrix.\n const int idx_i = lane_id_mod_16 / 4;\n const int idx_j = lane_id_mod_16 % 4;\n int globalWarpId = blockIdx.x * nWarps + warpId;\n // Shared memory to store the blocks to process\n __shared__ volatile ValueTypeA s_C_mtx[nWarps][32];\n __shared__ volatile ValueTypeA s_F_mtx[nWarps][16];\n // Shared memory to store the proposed column to load\n __shared__ volatile int s_aColSrc[nWarps][32];\n // Shared memory to store the column indices of the current row\n __shared__ volatile int s_keys[nWarps][SMemSize];\n\n while (globalWarpId < num_rows_per_color)\n {\n int storedRowId[2];\n int I = 0;\n\n for (; I < 2 && globalWarpId < num_rows_per_color ; I++)\n {\n int aRowId = sorted_rows_by_color[globalWarpId];\n storedRowId[I] = aRowId;\n int aColBeg = A_row_offsets[aRowId + 0];\n int aColEnd = A_row_offsets[aRowId + 1];\n int aColSmaller = A_smaller_color_offsets[aRowId];\n // The number of columns.\n const int nCols = aColEnd - aColBeg;\n\n //TODO: Add fallback for cases where number of nonzeros exceed SMemSize\n if ( nCols > SMemSize )\n {\n wk_returnValue[0] = 1;\n return;\n }\n\n // Fill-in the local table.\n const int NUM_STEPS = SMemSize / 32;\n\n#pragma unroll\n for ( int step = 0, k = laneId ; step < NUM_STEPS ; ++step, k += 32 )\n {\n int aColIt = aColBeg + k;\n int aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColId = A_col_indices[aColIt];\n }\n\n s_keys[warpId][k] = aColId;\n }\n\n // Now load all column indices of neighbours that have colors smaller than yours\n for ( int aColIt = aColBeg; aColIt < aColSmaller ; aColIt++)\n {\n unsigned int active_mask = utils::activemask();\n // Read the row to process, should be a broadcast\n int waRowId = s_keys[warpId][aColIt - aColBeg];\n // Compute multiplicative factor, load C_jj in first half, C_ij in second half\n int aColIdx = aColIt;\n\n if ( lane_id_div_16 == 0 )\n {\n aColIdx = A_dia_indices[waRowId];\n }\n\n s_C_mtx[warpId][laneId] = A_nonzero_values[16 * aColIdx + lane_id_mod_16];\n // Threads 0-15 perform the matrix product\n ValueTypeA tmp(0);\n\n if (ROW_MAJOR)\n {\n\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_C_mtx[warpId][16 + 4 * idx_i + m] * s_C_mtx[warpId][4 * m + idx_j];\n }\n }\n else\n {\n\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_C_mtx[warpId][16 + 4 * m + idx_j] * s_C_mtx[warpId][4 * idx_i + m];\n }\n }\n\n if ( lane_id_div_16 == 0 )\n {\n s_F_mtx[warpId][laneId] = tmp;\n A_nonzero_values[16 * aColIt + laneId] = tmp;\n }\n\n int waColIt = __ldcg(A_larger_color_offsets + waRowId);\n int waColEnd = __ldcg(A_row_offsets + waRowId + 1);\n\n // Load the first 32 columns of waRowId\n for (waColIt += laneId ; utils::any(waColIt < waColEnd, active_mask ); waColIt += 32 )\n {\n // Each thread loads its column id\n int waColId = -1;\n\n if ( waColIt < waColEnd )\n {\n waColId = A_col_indices[waColIt];\n }\n\n // Find the right column.\n int found_aColIt = -1;\n\n#pragma unroll 4\n for ( int i = 0, num_keys = aColEnd - aColBeg ; i < num_keys ; ++i )\n if ( s_keys[warpId][i] == waColId )\n {\n found_aColIt = i;\n }\n\n if ( found_aColIt != -1 )\n {\n found_aColIt += aColBeg;\n }\n\n // Store all the columns that have been found\n const int pred = found_aColIt != -1;\n int vote = utils::ballot( pred, active_mask );\n const int idst = __popc(vote & utils::lane_mask_lt());\n\n if (pred)\n {\n s_aColSrc[warpId][idst] = laneId;\n }\n utils::syncwarp(active_mask);\n\n const int n_cols = __popc( vote );\n\n // Process all columns that have been found\n for ( int k = 0 ; k < n_cols ; k += 2 )\n {\n const int my_k = k + lane_id_div_16;\n // Where to get columns from.\n int a_col_it = -1, w_col_it = -1;\n // Load column to load\n a_col_it = utils::shfl(found_aColIt, s_aColSrc[warpId][my_k], warpSize, active_mask);\n w_col_it = utils::shfl(waColIt, s_aColSrc[warpId][my_k], warpSize, active_mask);\n\n if ( my_k >= n_cols )\n {\n a_col_it = -1;\n w_col_it = -1;\n }\n\n ValueTypeA my_C(0);\n\n if ( w_col_it != -1 )\n {\n my_C = A_nonzero_values[16 * w_col_it + lane_id_mod_16];\n }\n\n s_C_mtx[warpId][laneId] = my_C;\n // Run the matrix-matrix product.\n ValueTypeA tmp(0);\n utils::syncwarp( active_mask );\n\n if (ROW_MAJOR)\n {\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_F_mtx[warpId][4 * idx_i + m] * s_C_mtx[warpId][16 * lane_id_div_16 + 4 * m + idx_j];\n }\n }\n else\n {\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_F_mtx[warpId][4 * m + idx_j] * s_C_mtx[warpId][16 * lane_id_div_16 + 4 * idx_i + m];\n }\n }\n\n if ( a_col_it != -1 )\n {\n A_nonzero_values[16 * a_col_it + lane_id_mod_16] -= tmp;\n }\n } // Loop over columns that have a match (for k=0;k( s_C_mtx[warpId], 16 * lane_id_div_16, offset, idx_i, idx_j, A_nonzero_values );\n }\n else\n {\n compute_block_inverse_col_major4x4_formula2( s_C_mtx[warpId], 16 * lane_id_div_16, offset, idx_i, idx_j, A_nonzero_values );\n }\n } // End of if statement\n } // End of while loop\n}\n\n#else\n\ntemplate< typename ValueTypeA, int CtaSize, int SMemSize, bool ROW_MAJOR>\n__global__ __launch_bounds__( CtaSize )\nvoid\ncomputeLUFactors_4x4_kernel( int A_nRows,\n const int *__restrict A_row_offsets,\n const int *__restrict A_col_indices,\n const int *__restrict A_dia_indices,\n ValueTypeA *__restrict A_nonzero_values,\n const int *__restrict A_smaller_color_offsets,\n const int *__restrict A_larger_color_offsets,\n const int *sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n int *wk_returnValue )\n{\n const int nWarps = CtaSize / 32; // Number of warps per Cta\n const int warpId = utils::warp_id();\n const int laneId = utils::lane_id();\n int lane_mask_lt = utils::lane_mask_lt();\n\n // Lane ID in the 2 16-wide segments.\n const int lane_id_div_16 = laneId / 16;\n const int lane_id_mod_16 = laneId % 16;\n // Coordinates inside a 4x4 block of the matrix.\n const int idx_i = lane_id_mod_16 / 4;\n const int idx_j = lane_id_mod_16 % 4;\n int globalWarpId = blockIdx.x * nWarps + warpId;\n // Shared memory to store the blocks to process\n __shared__ volatile ValueTypeA s_C_mtx[nWarps][32];\n // Shared memory to store the proposed column to load\n __shared__ volatile int s_aColItToLoad[nWarps][32];\n __shared__ volatile int s_waColItToLoad[nWarps][32];\n // Shared memory to store the proposed column to load\n __shared__ volatile unsigned s_aColIds[nWarps][32];\n // The size of the hash table (one per warp - shared memory).\n __shared__ volatile int s_size[nWarps][2];\n // Shared memory to store the column indices of the current row\n __shared__ volatile int s_keys[nWarps][SMemSize];\n\n while (globalWarpId < num_rows_per_color)\n {\n int aRowId = sorted_rows_by_color[globalWarpId];\n // Insert all the column indices in shared memory\n // TODO: Use texture here\n int aColBeg = A_row_offsets[aRowId];\n int aColEnd = A_row_offsets[aRowId + 1];\n int aColIt = aColBeg;\n\n // Check if number of nonzeros will fit in shared memory\n if ( (aColEnd - aColBeg) > SMemSize )\n {\n wk_returnValue[0] = 1;\n return;\n }\n\n // Load the all the column indices of row into shared memory\n for ( aColIt += laneId ; utils::any( aColIt < aColEnd ) ; aColIt += 32 )\n {\n int aColId = aColIt < aColEnd ? (int) A_col_indices[aColIt] : -1;\n s_keys[warpId][aColIt - aColBeg] = aColId;\n }\n\n // Now load all column indices of neighbours that have colors smaller than yours\n aColIt = aColBeg;\n int aColSmaller = A_smaller_color_offsets[aRowId];\n\n for ( ; utils::any( (aColIt + laneId) < aColSmaller ) ; aColIt += 32 )\n {\n int aColId = (aColIt + laneId) < aColSmaller ? (int) A_col_indices[aColIt + laneId] : -1;\n // Each thread pushes its column\n s_aColIds[warpId][laneId] = aColId;\n\n // Have warp collaborate to load each row\n for ( int j = 0; j < 32; j++)\n {\n // Check if row to load is valid\n if ( ( aColIt + j ) >= aColSmaller ) { break; }\n\n // Read the row to process, should be a broadcast\n int waRowId = s_aColIds[warpId][j];\n\n // Compute multiplicative factor, load C_jj in first half, C_ij in second half\n if (lane_id_div_16 == 0)\n {\n s_C_mtx[warpId][laneId] = A_nonzero_values[ 16 * A_dia_indices[waRowId] + lane_id_mod_16 ];\n }\n else\n {\n s_C_mtx[warpId][laneId] = A_nonzero_values[ 16 * (aColIt + j) + lane_id_mod_16 ];\n }\n\n // Threads 0-15 perform the matrix product\n utils::syncwarp();\n if (lane_id_div_16 == 0)\n {\n ValueTypeA tmp(0);\n\n if (ROW_MAJOR)\n {\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_C_mtx[warpId][16 + 4 * idx_i + m] * s_C_mtx[warpId][4 * m + idx_j];\n }\n }\n else\n {\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_C_mtx[warpId][16 + 4 * m + idx_j] * s_C_mtx[warpId][4 * idx_i + m];\n }\n }\n\n s_C_mtx[warpId][laneId] = tmp;\n A_nonzero_values[16 * (aColIt + j) + laneId] = tmp;\n }\n\n int waColIt = A_larger_color_offsets[waRowId];\n int waColEnd = A_row_offsets[waRowId + 1];\n\n //// Load the first 32 columns of waRowId\n for (waColIt += laneId ; utils::any(waColIt < waColEnd ); waColIt += 32 )\n {\n // Each thread loads its column id\n int waColId = waColIt < waColEnd ? A_col_indices[waColIt] : int (-1);\n // TODO: Try binary search if columns are ordered\n int found_aColIt = -1;\n\n //TODO: if invalid waColId, don't search\n for (int i = 0 ; utils::any(found_aColIt == -1) && i < aColEnd - aColBeg ; i++)\n {\n if (s_keys[warpId][i] == waColId) { found_aColIt = aColBeg + i; }\n }\n\n // Store all the columns that have been found\n const int pred = found_aColIt != -1;\n const int vote = utils::ballot( pred );\n const int idst = __popc(vote & lane_mask_lt);\n\n if (pred)\n {\n s_aColItToLoad [warpId][idst] = found_aColIt;\n s_waColItToLoad[warpId][idst] = waColIt;\n }\n\n const int n_cols = __popc( vote );\n\n // Process all columns that have been found\n for ( int k = 0 ; k < n_cols ; k++ )\n {\n // Load column to load\n const int a_col_it = k < n_cols ? s_aColItToLoad [warpId][k] : -1;\n const int w_col_it = k < n_cols ? s_waColItToLoad[warpId][k] : -1;\n\n if (lane_id_div_16 == 1)\n {\n s_C_mtx[warpId][laneId] = A_nonzero_values[16 * w_col_it + lane_id_mod_16];\n // Run the matrix-matrix product.\n ValueTypeA tmp(0);\n utils::syncwarp(utils::activemask());\n\n if (ROW_MAJOR)\n {\n\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_C_mtx[warpId][4 * idx_i + m] * s_C_mtx[warpId][16 + 4 * m + idx_j];\n }\n }\n else\n {\n\n#pragma unroll\n for ( int m = 0 ; m < 4 ; ++m )\n {\n tmp += s_C_mtx[warpId][4 * m + idx_j] * s_C_mtx[warpId][16 + 4 * idx_i + m];\n }\n }\n\n A_nonzero_values[16 * a_col_it + lane_id_mod_16] -= tmp;\n }\n } // Loop over columns that have a match (for k=0;k\n (s_C_mtx[warpId], 0, offset, idx_i, idx_j, A_nonzero_values);\n }\n else\n {\n compute_block_inverse_col_major\n (s_C_mtx[warpId], 0, offset, idx_i, idx_j, A_nonzero_values);\n }\n }\n\n globalWarpId += nWarps * gridDim.x;\n } // if RowId < A_nRows;\n}\n#endif\n// ----------\n// Methods\n// ----------\n\n// Constructor\ntemplate\nMulticolorILUSolver_Base::MulticolorILUSolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope)\n{\n m_sparsity_level = cfg.AMG_Config::template getParameter(\"ilu_sparsity_level\", cfg_scope);\n m_weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n this->m_reorder_cols_by_color_desired = (cfg.AMG_Config::template getParameter(\"reorder_cols_by_color\", cfg_scope) != 0);\n this->m_insert_diagonal_desired = (cfg.AMG_Config::template getParameter(\"insert_diag_while_reordering\", cfg_scope) != 0);\n\n if (cfg.AMG_Config::template getParameter(\"use_bsrxmv\", cfg_scope))\n {\n this->m_use_bsrxmv = 1;\n }\n else\n {\n this->m_use_bsrxmv = 0;\n }\n\n if (m_weight == ValueTypeB(0.))\n {\n m_weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Multicolor DILU smoother\\n\");\n }\n}\n\n// Destructor\ntemplate\nMulticolorILUSolver_Base::~MulticolorILUSolver_Base()\n{\n m_LU.set_initialized(0);\n m_A_to_LU_mapping.clear();\n m_A_to_LU_mapping.shrink_to_fit();\n m_LU.resize(0, 0, 0, 1);\n}\n\ntemplate \nvoid MulticolorILUSolver >::computeAtoLUmapping()\n{\n FatalError(\"Haven't implemented Multicolor ILU smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n\ntemplate \nvoid MulticolorILUSolver >::computeAtoLUmapping()\n{\n Matrix &m_A = *this->m_explicit_A;\n const int CtaSize = 128; // Number of threads per CTA\n const int SMemSize = 128; // per warp\n const int nWarps = CtaSize / 32;\n int GridSize = std::min( AMGX_GRID_MAX_SIZE, ( this->m_explicit_A->get_num_rows( ) + nWarps - 1 ) / nWarps );\n // Global memory workspaces\n device_vector_alloc returnValue(1);\n returnValue[0] = 0;\n\n if (this->m_explicit_A->hasProps(DIAG))\n {\n computeAtoLUmappingExtDiag_kernel <<< GridSize, CtaSize >>> (\n m_A.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &m_A.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &m_A.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &m_A.diag[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_A_to_LU_mapping[0] ),\n amgx::thrust::raw_pointer_cast( &returnValue[0] ));\n }\n else\n {\n computeAtoLUmapping_kernel <<< GridSize, CtaSize >>> (\n m_A.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &m_A.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &m_A.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_A_to_LU_mapping[0] ),\n amgx::thrust::raw_pointer_cast( &returnValue[0] ));\n }\n\n cudaCheckError();\n\n // fallback path that allows 1024 nonzeros per row\n if (returnValue[0] == 1)\n {\n returnValue[0] = 0;\n const int SMemSize2 = 1024 ; // per warp\n\n if (this->m_explicit_A->hasProps(DIAG))\n {\n computeAtoLUmappingExtDiag_kernel <<< GridSize, CtaSize >>> (\n m_A.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &m_A.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &m_A.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &m_A.diag[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_A_to_LU_mapping[0] ),\n amgx::thrust::raw_pointer_cast( &returnValue[0] ));\n }\n else\n {\n computeAtoLUmapping_kernel <<< GridSize, CtaSize >>> (\n m_A.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &m_A.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &m_A.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_A_to_LU_mapping[0] ),\n amgx::thrust::raw_pointer_cast( &returnValue[0] ));\n }\n\n cudaCheckError();\n }\n\n if (returnValue[0] == 1)\n {\n FatalError( \"Number of nonzeros per row exceeds allocated shared memory\", AMGX_ERR_NO_MEMORY);\n }\n}\n\ntemplate \nvoid MulticolorILUSolver >::fillLUValuesWithAValues()\n{\n FatalError(\"Haven't implemented Multicolor ILU smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n\ntemplate \nvoid MulticolorILUSolver >::fillLUValuesWithAValues()\n{\n if (this->m_sparsity_level == 0)\n {\n this->m_LU.values = this->m_explicit_A->values;\n }\n else\n {\n // TODO: Should probably store the inverse mapping of AtoLUmapping instead\n // This will allow to use unpermuteVector and have coalesced writes\n // instead of coalesced reads\n thrust_wrapper::fill(this->m_LU.values.begin(), this->m_LU.values.end(), 0.);\n cudaCheckError();\n\n if (this->m_explicit_A->hasProps(DIAG))\n {\n amgx::permuteVector(this->m_explicit_A->values, this->m_LU.values, this->m_A_to_LU_mapping, (this->m_explicit_A->get_num_nz() + this->m_explicit_A->get_num_rows())*this->m_explicit_A->get_block_size());\n }\n else\n {\n amgx::permuteVector(this->m_explicit_A->values, this->m_LU.values, this->m_A_to_LU_mapping, this->m_explicit_A->get_num_nz()*this->m_explicit_A->get_block_size());\n }\n }\n}\n\n\ntemplate \nvoid MulticolorILUSolver >::computeLUSparsityPattern()\n{\n FatalError(\"Haven't implemented Multicolor ILU smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorILUSolver >::computeLUSparsityPattern()\n{\n // ILU0\n if (this->m_sparsity_level == 0)\n {\n // Copy everything except the values\n this->m_LU.copy_structure(*this->m_explicit_A);\n }\n // ILU1\n else if (this->m_sparsity_level == 1)\n {\n this->sparsity_wk = CSR_Multiply::csr_workspace_create( *this->m_cfg, \"default\" );\n CSR_Multiply::csr_sparsity_ilu1( *this->m_explicit_A, this->m_LU, this->sparsity_wk );\n CSR_Multiply::csr_workspace_delete( this->sparsity_wk );\n\n if (this->m_use_bsrxmv)\n {\n this->m_LU.set_initialized(0);\n this->m_LU.computeDiagonal();\n this->m_LU.set_initialized(1);\n }\n\n this->m_LU.setMatrixColoring(&(this->m_explicit_A->getMatrixColoring()));\n }\n else\n {\n FatalError(\"Haven't implemented Multicolor ILU smoother for this sparsity level. \", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid MulticolorILUSolver >::computeLUFactors()\n{\n FatalError(\"Haven't implemented Multicolor ILU smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorILUSolver >::computeLUFactors()\n{\n const int CtaSize = 128; // Number of threads per CTA\n const int SMemSize = 128;\n const int nWarps = CtaSize / 32;\n device_vector_alloc returnValue(1);\n returnValue[0] = 0;\n int num_colors = this->m_LU.getMatrixColoring().getNumColors();\n const IndexType *LU_sorted_rows_by_color_ptr = this->m_LU.getMatrixColoring().getSortedRowsByColor().raw();\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = this->m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = this->m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n#ifdef EXPERIMENTAL_LU_FACTORS\n int GridSize = std::min( 2048, ( num_rows_per_color + nWarps - 1 ) / nWarps );\n\n if ( GridSize == 0 )\n {\n continue; // if perfect coloring (color 0 has no vertices)\n }\n\n if ( this->m_LU.get_block_dimx() == 4 && this->m_LU.get_block_dimy() == 4 )\n {\n if ( this->m_explicit_A->getBlockFormat() == ROW_MAJOR )\n {\n compute_LU_factors_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n this->m_LU.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.diag[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.values[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_smaller_color_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_larger_color_offsets[0] ),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n amgx::thrust::raw_pointer_cast( &returnValue[0] ) );\n }\n else\n {\n compute_LU_factors_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n this->m_LU.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.diag[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.values[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_smaller_color_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_larger_color_offsets[0] ),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n amgx::thrust::raw_pointer_cast( &returnValue[0] ) );\n }\n\n cudaCheckError();\n }\n else\n {\n FatalError(\"Unsupported block size for Multicolor ILU solver, computeLUFactors\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n#else\n int GridSize = std::min( AMGX_GRID_MAX_SIZE, ( num_rows_per_color + nWarps - 1 ) / nWarps );\n\n if ( GridSize == 0 )\n {\n continue; // if perfect coloring (color 0 has no vertices)\n }\n\n if ( this->m_LU.get_block_dimx() == 4 && this->m_LU.get_block_dimy() == 4 )\n {\n //computeLUFactors_4x4_kernel <<< GridSize, CtaSize>>> (\n if (this->m_explicit_A->getBlockFormat() == ROW_MAJOR)\n {\n computeLUFactors_4x4_kernel <<< GridSize, CtaSize>>> (\n this->m_LU.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.diag[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.values[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_smaller_color_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_larger_color_offsets[0] ),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n amgx::thrust::raw_pointer_cast( &returnValue[0] ) );\n }\n else\n {\n computeLUFactors_4x4_kernel <<< GridSize, CtaSize>>> (\n this->m_LU.get_num_rows( ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.row_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.col_indices[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.diag[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.values[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_smaller_color_offsets[0] ),\n amgx::thrust::raw_pointer_cast( &this->m_LU.m_larger_color_offsets[0] ),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n amgx::thrust::raw_pointer_cast( &returnValue[0] ) );\n }\n\n cudaCheckError();\n }\n else\n {\n FatalError(\"Unsupported block size for Multicolor ILU solver, computeLUFactors\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n#endif\n }\n\n // Check returnValue flag\n if ( returnValue[0] == 1 )\n {\n FatalError( \"Number of nonzeros per row exceeds allocated shared memory\", AMGX_ERR_NO_MEMORY);\n }\n}\n\n\n// Solver pre-setup\ntemplate\nvoid\nMulticolorILUSolver_Base::pre_setup()\n{\n // Check if matrix is colored\n if (this->m_explicit_A->getColoringLevel() < m_sparsity_level + 1)\n {\n FatalError(\"Matrix must be colored with coloring_level > sparsity_level for the multicolorILUsolver\", AMGX_ERR_CONFIGURATION);\n }\n\n // Compute extended sparsity pattern based on coloring and matrix A\n computeLUSparsityPattern();\n\n if (this->m_LU.hasProps(DIAG))\n {\n FatalError(\"Multicolor ILU smoother does not support outside diagonal. Try setting reorder_cols_by_color=1 and insert_diag_while_reordering=1 in the multicolor_ilu solver scope in configuration file\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (m_sparsity_level == 0 && !this->m_LU.getColsReorderedByColor())\n {\n FatalError(\"Multicolor ILU smoother requires matrix to be reordered by color with ILU0 solver. Try setting reorder_cols_by_color=1 and insert_diag_while_reordering=1 in the multicolor_ilu solver scope in configuration file\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // Reorder the columns of LU by color\n if (m_sparsity_level != 0)\n {\n // Reorder columns of LU by color\n m_LU.reorderColumnsByColor(false);\n\n // Compute mapping between entries in A and entries in LU\n if (this->m_explicit_A->hasProps(DIAG))\n {\n m_A_to_LU_mapping.resize(this->m_explicit_A->get_num_nz() + this->m_explicit_A->get_num_rows());\n }\n else\n {\n m_A_to_LU_mapping.resize(this->m_explicit_A->get_num_nz());\n }\n\n computeAtoLUmapping();\n }\n\n int N = this->m_LU.get_num_rows() * this->m_LU.get_block_dimy();\n m_delta.resize(N);\n m_Delta.resize(N);\n m_Delta.set_block_dimy(this->m_explicit_A->get_block_dimy());\n m_Delta.set_block_dimx(1);\n m_delta.set_block_dimy(this->m_explicit_A->get_block_dimy());\n m_delta.set_block_dimx(1);\n}\n\ntemplate\nvoid\nMulticolorILUSolver_Base::printSolverParameters() const\n{\n std::cout << \"relaxation_factor = \" << this->m_weight << std::endl;\n std::cout << \"use_bsrxmv = \" << this->m_use_bsrxmv << std::endl;\n std::cout << \"ilu_sparsity_level = \" << this->m_sparsity_level << std::endl;\n}\n\n\n// Solver setup\ntemplate\nvoid\nMulticolorILUSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n this->m_explicit_A = dynamic_cast*>(this->m_A);\n\n if (!this->m_explicit_A)\n {\n FatalError(\"MulticolorILUSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n if (this->m_explicit_A->getColoringLevel() < 1)\n {\n FatalError(\"Matrix must be colored to use multicolor ilu solver. Try setting: coloring_level=1 or coloring_level=2 in the configuration file\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (!reuse_matrix_structure)\n {\n this->pre_setup();\n }\n\n // Fill LU sparsity pattern\n fillLUValuesWithAValues();\n // Compute LU factors in place (update LU.values)\n computeLUFactors();\n}\n\n//\ntemplate\nvoid\nMulticolorILUSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nMulticolorILUSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n if ( !m_use_bsrxmv && (this->m_LU.get_block_dimx() == 4 && this->m_LU.get_block_dimy() == 4) )\n {\n smooth_4x4(b, x, xIsZero);\n }\n else\n {\n smooth_bxb(b, x, xIsZero);\n }\n\n // Do we converge ?\n return this->converged(b, x);\n}\n\ntemplate\nvoid\nMulticolorILUSolver_Base::solve_finalize( VVector &b, VVector &x )\n{}\n\ntemplate \nvoid MulticolorILUSolver >::smooth_4x4(const VVector &b, VVector &x, bool xIsZero)\n{\n FatalError(\"Haven't implemented Multicolor DILU smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorILUSolver >::smooth_4x4(const VVector &b, VVector &x, bool xIsZero)\n{\n Matrix &m_LU = this->m_LU;\n Matrix &m_A = *this->m_explicit_A;\n int N = m_LU.get_num_rows() * m_LU.get_block_dimy();\n cudaCheckError();\n\n if (!m_LU.getColsReorderedByColor())\n {\n FatalError(\"ILU solver currently only works if columns are reordered by color. Try setting reordering_cols_by_color=1 in the multicolor_ilu solver scope in the configuration file\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // ---------------------------------------------------------\n // Solving Lower triangular system, with identity diagonal\n // ---------------------------------------------------------\n const IndexType *LU_sorted_rows_by_color_ptr = m_LU.getMatrixColoring().getSortedRowsByColor().raw();\n int num_colors = this->m_LU.getMatrixColoring().getNumColors();\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n#ifdef EXPERIMENTAL_LU_FORWARD\n const int CtaSize = 128; // Number of threads per CTA\n const int nHalfWarps = CtaSize / 16;\n int GridSize = std::min( 2048, ( num_rows_per_color + nHalfWarps - 1 ) / nHalfWarps );\n\n if ( GridSize == 0 )\n {\n continue; // if perfect coloring (color 0 has no vertices)\n }\n\n if ( this->m_explicit_A->getBlockFormat() == ROW_MAJOR )\n {\n if (m_A.hasProps(DIAG))\n {\n LU_forward_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n m_LU.row_offsets.raw(),\n m_LU.m_smaller_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.values.raw(),\n m_A.row_offsets.raw(),\n m_A.col_indices.raw(),\n m_A.values.raw(),\n m_A.diag.raw(),\n x.raw(),\n b.raw(),\n this->m_delta.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n xIsZero );\n }\n else\n {\n LU_forward_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n m_LU.row_offsets.raw(),\n m_LU.m_smaller_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.values.raw(),\n m_A.row_offsets.raw(),\n m_A.col_indices.raw(),\n m_A.values.raw(),\n m_A.diag.raw(),\n x.raw(),\n b.raw(),\n this->m_delta.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n xIsZero );\n }\n }\n else\n {\n // COL_MAJOR\n if (m_A.hasProps(DIAG))\n {\n LU_forward_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n m_LU.row_offsets.raw(),\n m_LU.m_smaller_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.values.raw(),\n m_A.row_offsets.raw(),\n m_A.col_indices.raw(),\n m_A.values.raw(),\n m_A.diag.raw(),\n x.raw(),\n b.raw(),\n this->m_delta.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n xIsZero );\n }\n else\n {\n LU_forward_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n m_LU.row_offsets.raw(),\n m_LU.m_smaller_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.values.raw(),\n m_A.row_offsets.raw(),\n m_A.col_indices.raw(),\n m_A.values.raw(),\n m_A.diag.raw(),\n x.raw(),\n b.raw(),\n this->m_delta.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n xIsZero );\n }\n }\n\n#else\n const int CtaSize = 128;\n const int blockrows_per_cta = CtaSize / 4;\n const int GridSize = min( AMGX_GRID_MAX_SIZE, (int) (num_rows_per_color + blockrows_per_cta - 1) / blockrows_per_cta);\n\n if ( GridSize == 0 )\n {\n continue; // if perfect coloring (color 0 has no vertices)\n }\n\n if (this->m_explicit_A->hasProps(DIAG))\n {\n FatalError(\"this implementation of LU forward solve does not support A with external diagonal\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (this->m_explicit_A->getBlockFormat() == ROW_MAJOR)\n {\n LU_forward_4x4_kernel <<< GridSize, CtaSize>>>\n (m_LU.row_offsets.raw(),\n m_LU.m_smaller_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.values.raw(),\n m_A.row_offsets.raw(),\n m_A.col_indices.raw(),\n m_A.values.raw(),\n x.raw(),\n b.raw(),\n delta.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n xIsZero);\n }\n else\n {\n LU_forward_4x4_kernel <<< GridSize, CtaSize>>>\n (m_LU.row_offsets.raw(),\n m_LU.m_smaller_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.values.raw(),\n m_A.row_offsets.raw(),\n m_A.col_indices.raw(),\n m_A.values.raw(),\n x.raw(),\n b.raw(),\n delta.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n xIsZero);\n }\n\n#endif\n cudaCheckError();\n }\n\n // --------------------\n // Backward Sweep\n // --------------------\n for (int i = num_colors - 1; i >= 0; i--)\n {\n const IndexType color_offset = m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n#ifdef EXPERIMENTAL_LU_BACKWARD\n const int CtaSize = 128; // Number of threads per CTA\n const int nHalfWarps = CtaSize / 16;\n int GridSize = std::min( 2048, ( num_rows_per_color + nHalfWarps - 1 ) / nHalfWarps );\n\n if ( GridSize == 0 )\n {\n continue; // if perfect coloring (color 0 has no vertices)\n }\n\n if (this->m_explicit_A->getBlockFormat() == ROW_MAJOR)\n {\n LU_backward_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n m_LU.row_offsets.raw(),\n m_LU.m_larger_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.diag.raw(),\n m_LU.values.raw(),\n this->m_delta.raw(),\n this->m_Delta.raw(),\n x.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n num_colors,\n this->m_weight,\n xIsZero);\n }\n else\n {\n LU_backward_4x4_kernel_warp <<< GridSize, CtaSize>>>(\n m_LU.row_offsets.raw(),\n m_LU.m_larger_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.diag.raw(),\n m_LU.values.raw(),\n this->m_delta.raw(),\n this->m_Delta.raw(),\n x.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n num_colors,\n this->m_weight,\n xIsZero);\n }\n\n#else\n const int CtaSize = 128;\n const int blockrows_per_cta = CtaSize / 4;\n const int GridSize = min( AMGX_GRID_MAX_SIZE, (int) (num_rows_per_color + blockrows_per_cta - 1) / blockrows_per_cta);\n\n if ( GridSize == 0 )\n {\n continue; // if perfect coloring (color 0 has no vertices)\n }\n\n if (this->m_explicit_A->getBlockFormat() == ROW_MAJOR)\n {\n LU_backward_4x4_kernel <<< GridSize, CtaSize>>>\n (m_LU.row_offsets.raw(),\n m_LU.m_larger_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.diag.raw(),\n m_LU.values.raw(),\n this->m_delta.raw(),\n this->m_Delta.raw(),\n x.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n num_colors,\n this->m_weight,\n xIsZero);\n }\n else\n {\n LU_backward_4x4_kernel <<< GridSize, CtaSize>>>\n (m_LU.row_offsets.raw(),\n m_LU.m_larger_color_offsets.raw(),\n m_LU.col_indices.raw(),\n m_LU.diag.raw(),\n m_LU.values.raw(),\n this->m_delta.raw(),\n this->m_Delta.raw(),\n x.raw(),\n LU_sorted_rows_by_color_ptr + color_offset,\n num_rows_per_color,\n i,\n num_colors,\n this->m_weight,\n xIsZero);\n }\n\n#endif\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid MulticolorILUSolver >::smooth_bxb(const VVector &b, VVector &x, bool xIsZero)\n{\n FatalError(\"Haven't implemented Multicolor DILU smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MulticolorILUSolver >::smooth_bxb(const VVector &b, VVector &x, bool xIsZero)\n{\n Matrix &m_LU = this->m_LU;\n Matrix &m_A = *this->m_explicit_A;\n int N = m_LU.get_num_rows() * m_LU.get_block_dimy();\n\n if (!m_LU.getColsReorderedByColor())\n {\n FatalError(\"ILU solver currently only works if columns are reordered by color. Try setting reorder_cols_by_color=1 in the multicolor_ilu solver scope in the configuration file\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (this->m_explicit_A->getBlockFormat() == COL_MAJOR)\n {\n FatalError(\"ILU solver for arbitrary block sizes only works with ROW_MAJOR matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // ---------------------------------------------------------\n // Solving Lower triangular system, with identity diagonal\n // ---------------------------------------------------------\n const IndexType *LU_sorted_rows_by_color_ptr = m_LU.getMatrixColoring().getSortedRowsByColor().raw();\n int num_colors = this->m_LU.getMatrixColoring().getNumColors();\n //delta = b;\n amgx::thrust::copy(b.begin(), b.end(), this->m_delta.begin());\n //delta = delta - Ax;\n Cusparse::bsrmv((ValueTypeA) - 1.0, m_A, x, (ValueTypeA)1.0, this->m_delta);\n cudaCheckError();\n // Setting Delta to zero\n thrust_wrapper::fill(this->m_Delta.begin(), this->m_Delta.end(), (ValueTypeB)0.0f);\n cudaCheckError();\n bool skipped_end = false;\n\n for (int i = 0; i < num_colors; i++)\n {\n const IndexType color_offset = m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const IndexType num_rows_per_color = m_LU.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n\n if (num_rows_per_color == 0) { continue; } // if perfect coloring (color 0 has no vertices)\n\n if (skipped_end)\n {\n // delta = delta - LU*Delta smaller colors\n Cusparse::bsrmv(Cusparse::SMALLER_COLORS, i, (ValueTypeA) - 1.0f, m_LU, this->m_delta, (ValueTypeA)1.0f, this->m_delta);\n }\n\n if (num_rows_per_color > 0)\n {\n skipped_end = true;\n }\n }\n\n cudaCheckError();\n skipped_end = false;\n\n // --------------------\n // Backward Sweep\n // --------------------\n for (int i = num_colors - 1; i >= 0; i--)\n {\n // delta = delta - LU*Delta larger colors\n Cusparse::bsrmv(Cusparse::LARGER_COLORS, i, (ValueTypeA) - 1.0f, m_LU, this->m_Delta, (ValueTypeA)1.0f, this->m_delta);\n // Multiple by inverse stored on diagonal\n Cusparse::bsrmv(Cusparse::DIAG_COL, i, (ValueTypeA) 1.0f, m_LU, this->m_delta, 0.0f, this->m_Delta);\n }\n\n cudaCheckError();\n axpy(this->m_Delta, x, this->m_weight, 0, x.size());\n cudaCheckError();\n}\n\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class MulticolorILUSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class MulticolorILUSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/solvers/multicolor_ilu_solver.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \n\n#define AMGX_ILU_COLORING\n\nnamespace amgx\n{\nnamespace multicolor_dilu_solver\n{\n\nenum { CTA_SIZE = 128, WARP_SIZE = 32 };\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, int N, int CTA_SIZE, int WARP_SIZE, int NUM_WARP_ITERS_PER_BLOCK >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_setup_NxN_kernel_large( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict A_diag,\n const Matrix_type *__restrict A_vals,\n Matrix_type *__restrict Einv,\n const int *sorted_rows_by_color,\n const int *row_colors,\n const int num_rows_per_color,\n const int current_color )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items computer per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_WARPS_PER_CTA;\n // Number of items per grid.\n const int NUM_ITEMS_PER_GRID = gridDim.x * NUM_ITEMS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Shared memory to broadcast column IDs.\n __shared__ volatile int s_a_col_ids[CTA_SIZE];\n __shared__ volatile int s_a_col_its[CTA_SIZE];\n // Each thread keeps its own pointer.\n volatile int *my_s_a_col_ids = &s_a_col_ids[threadIdx.x - lane_id];\n volatile int *my_s_a_col_its = &s_a_col_its[threadIdx.x - lane_id];\n // Shared memory to store the matrices.\n __shared__ volatile Vector_type s_A_mtx[CTA_SIZE * NUM_WARP_ITERS_PER_BLOCK];\n __shared__ volatile Vector_type s_B_mtx[CTA_SIZE * NUM_WARP_ITERS_PER_BLOCK];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_A_mtx = &s_A_mtx[warp_id * NUM_WARP_ITERS_PER_BLOCK * WARP_SIZE];\n volatile Vector_type *my_s_B_mtx = &s_B_mtx[warp_id * NUM_WARP_ITERS_PER_BLOCK * WARP_SIZE];\n // Shared memory to store the index of the element Aji.\n __shared__ volatile int s_A_ji[NUM_WARPS_PER_CTA];\n // Each thread keeps its own pointer.\n volatile int *my_s_A_ji = &s_A_ji[warp_id];\n // Precomputing some stuff\n int idx[NUM_WARP_ITERS_PER_BLOCK];\n int idy[NUM_WARP_ITERS_PER_BLOCK];\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n const int id = (WARP_SIZE * wb + lane_id) % NxN;\n idx[wb] = id / N;\n idy[wb] = id % N;\n }\n\n // Determine which NxN block the threads work with.\n int a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; utils::any( a_row_it < num_rows_per_color ) ; a_row_it += NUM_ITEMS_PER_GRID )\n {\n int a_row_id = -1;\n\n if ( a_row_it < num_rows_per_color )\n {\n a_row_id = sorted_rows_by_color[a_row_it];\n }\n\n // Load the diagonal.\n Vector_type e_out[NUM_WARP_ITERS_PER_BLOCK];\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n e_out[wb] = (Vector_type)0.0;\n }\n\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n if ( a_row_id != -1 && (wb * WARP_SIZE + lane_id) < NxN)\n {\n e_out[wb] = A_vals[NxN * A_diag[a_row_id] + wb * WARP_SIZE + lane_id];\n }\n\n // Skip the 1st iteration of the outer-loop (that loop runs on the host).\n if ( current_color != 0 )\n {\n // Ranges of the rows.\n int a_col_begin(0), a_col_end(0);\n\n if ( a_row_id != -1 )\n {\n a_col_begin = A_rows[a_row_id ];\n a_col_end = A_rows[a_row_id + 1];\n }\n\n // Iterate over the elements in the columns.\n for ( ; a_col_begin < a_col_end ; a_col_begin += NxN )\n {\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id;\n // The identifier of the column if the iterator is valid.\n int a_col_tmp = -1, a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_tmp = A_cols[a_col_it];\n }\n\n if ( a_col_tmp != -1 && row_colors[a_col_tmp] < current_color )\n {\n a_col_id = a_col_tmp;\n }\n\n // When the diagonal is stored inside the matrix, we have to reject it. We\n // could be using a template parameter but it's not needed since that\n // rejection is really cheap (a couple of extra cycles -- CMP+MOV).\n if ( a_col_id == a_row_id )\n {\n a_col_id = -1;\n }\n\n // We partition valid and invalid column ids. Valid ones come first.\n int vote = utils::ballot( a_col_id != -1 );\n int ones = __popc( vote );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( a_col_id == -1 )\n {\n dest = ones + lane_id - dest;\n }\n\n my_s_a_col_ids[dest] = a_col_id;\n my_s_a_col_its[dest] = a_col_it;\n // Temporary storage with zeros for OOB\n Vector_type my_A[NUM_WARP_ITERS_PER_BLOCK], my_B[NUM_WARP_ITERS_PER_BLOCK];\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n my_A[wb] = (Vector_type)0.0;\n my_B[wb] = (Vector_type)0.0;\n }\n\n // Threads collaborate to load the rows.\n for ( int k = 0 ; k < WARP_SIZE ; ++k )\n {\n // Exchange column indices.\n const int uniform_a_col_id = my_s_a_col_ids[k];\n\n // Early exit.\n if ( uniform_a_col_id == -1 )\n {\n break;\n }\n\n // Load the iterator.\n const int uniform_a_col_it = my_s_a_col_its[k];\n // Load the two matrices.\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n if ((wb * WARP_SIZE + lane_id) < NxN)\n {\n my_A[wb] = A_vals[NxN * uniform_a_col_it + wb * WARP_SIZE + lane_id];\n my_B[wb] = Einv [NxN * uniform_a_col_id + wb * WARP_SIZE + lane_id];\n }\n\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n my_s_A_mtx[lane_id + wb * WARP_SIZE] = my_A[wb];\n my_s_B_mtx[lane_id + wb * WARP_SIZE] = my_B[wb];\n }\n\n // Compute the product of matrices.\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n my_A[wb] = (Vector_type)0.0;\n#pragma unroll\n\n for ( int m = 0 ; m < N ; ++m )\n {\n my_A[wb] += my_s_A_mtx[N * idx[wb] + m] * my_s_B_mtx[N * m + idy[wb]];\n }\n }\n\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n if ((wb * WARP_SIZE + lane_id) < NxN)\n {\n my_s_A_mtx[lane_id + wb * WARP_SIZE] = my_A[wb];\n }\n\n // We looking for columns in the two rows we're interested in.\n int b_col_it = A_rows[uniform_a_col_id ];\n int b_col_end = A_rows[uniform_a_col_id + 1];\n\n // Init the marker to -1.\n if ( lane_id == 0 )\n {\n *my_s_A_ji = -1;\n }\n\n // Run the loop.\n b_col_it += lane_id;\n int shared_found = utils::ballot( lane_id == 0 && uniform_a_col_id == -1 );\n\n do\n {\n bool found = b_col_it < b_col_end && A_cols[b_col_it] == a_row_id;\n\n if ( found )\n {\n *my_s_A_ji = b_col_it;\n }\n\n shared_found = shared_found | utils::ballot(found);\n b_col_it += NxN;\n }\n while ( __popc( shared_found ) == 0 && utils::any( b_col_it < b_col_end ) );\n\n // Load the blocks.\n const int w_aji = *my_s_A_ji;\n Vector_type my_C[NUM_WARP_ITERS_PER_BLOCK];\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n my_C[wb] = (Vector_type)0.0;\n\n if ( w_aji != -1 && (wb * WARP_SIZE + lane_id) < NxN)\n {\n my_C[wb] = A_vals[NxN * w_aji + wb * WARP_SIZE + lane_id];\n }\n\n my_s_B_mtx[wb * WARP_SIZE + lane_id] = my_C[wb];\n }\n\n // Update e_out.\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n#pragma unroll\n\n for ( int m = 0 ; m < N ; ++m )\n {\n e_out[wb] -= my_s_A_mtx[N * idx[wb] + m] * my_s_B_mtx[N * m + idy[wb]];\n }\n }\n }\n } // a_col_begin < a_col_end\n } // current_color != 0\n\n // Store e_out in A\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n my_s_B_mtx[wb * WARP_SIZE + lane_id] = my_s_A_mtx[wb * WARP_SIZE + lane_id] = e_out[wb];\n }\n\n // Invert the matrices.\n#pragma unroll\n\n for ( int row = 0 ; row < N ; ++row )\n {\n Vector_type diag(0), diag_tmp = my_s_A_mtx[N * row + row];\n\n if ( isNotCloseToZero(diag_tmp) )\n {\n diag = Vector_type(1) / diag_tmp;\n }\n else\n {\n diag = Vector_type(1) / epsilon(diag_tmp);\n }\n\n if ( lane_id < N && lane_id != row)\n {\n my_s_A_mtx[N * row + lane_id] = my_s_B_mtx[N * row + lane_id] = my_s_B_mtx[N * row + lane_id] * diag;\n }\n\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n if ( idx[wb] != row && idy[wb] != row)\n {\n my_s_A_mtx[wb * WARP_SIZE + lane_id] = my_s_B_mtx[wb * WARP_SIZE + lane_id] - my_s_B_mtx[N * idx[wb] + row] * my_s_B_mtx[N * row + idy[wb]];\n }\n\n if ( lane_id < N )\n {\n Vector_type tmp = diag;\n\n if ( lane_id != row )\n {\n tmp = -my_s_A_mtx[N * lane_id + row] * diag;\n }\n\n my_s_A_mtx[N * lane_id + row] = tmp;\n }\n\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n my_s_B_mtx[wb * WARP_SIZE + lane_id] = my_s_A_mtx[wb * WARP_SIZE + lane_id];\n }\n }\n\n // Store the results to Einv.\n if ( a_row_id != -1 )\n#pragma unroll\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n if (wb * WARP_SIZE + lane_id < NxN)\n {\n Einv[NxN * a_row_id + wb * WARP_SIZE + lane_id] = my_s_A_mtx[wb * WARP_SIZE + lane_id];\n }\n }\n}\n\n\ntemplate< typename Matrix_type, typename Vector_type, int N, int CTA_SIZE, int WARP_SIZE >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_setup_NxN_kernel( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict A_diag,\n const Matrix_type *__restrict A_vals,\n Matrix_type *__restrict Einv,\n const int *sorted_rows_by_color,\n const int *row_colors,\n const int num_rows_per_color,\n const int current_color )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int NUM_ITEMS_PER_WARP = WARP_SIZE / NxN;\n // Upper-bound on the number of items per warp.\n const int NUM_ITEMS_PER_WARP_CEIL = (WARP_SIZE + NxN - 1) / NxN;\n // Number of items computer per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_ITEMS_PER_WARP * NUM_WARPS_PER_CTA;\n // Number of items per grid.\n const int NUM_ITEMS_PER_GRID = gridDim.x * NUM_ITEMS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_NxN = lane_id / NxN;\n const int lane_id_mod_NxN = lane_id % NxN;\n // Useful index to compute matrix products.\n const int lane_id_mod_NxN_div_N = lane_id_mod_NxN / N;\n const int lane_id_mod_NxN_mod_N = lane_id_mod_NxN % N;\n // We need NxN to compute a NxN block. Encode a mask for the first block.\n int mask_tmp = utils::ballot( lane_id_div_NxN == 0 );\n // Mask for ballots. We shift the mask with NxN active bits by the needed number of bits.\n const int mask_NxN = mask_tmp << (lane_id_div_NxN * __popc(mask_tmp));\n // Shared memory to broadcast column IDs.\n __shared__ volatile int s_a_col_ids[CTA_SIZE];\n __shared__ volatile int s_a_col_its[CTA_SIZE];\n // Each thread keeps its own pointer.\n volatile int *my_s_a_col_ids = &s_a_col_ids[threadIdx.x - lane_id_mod_NxN];\n volatile int *my_s_a_col_its = &s_a_col_its[threadIdx.x - lane_id_mod_NxN];\n // Shared memory to store the matrices.\n __shared__ volatile Vector_type s_A_mtx[CTA_SIZE];\n __shared__ volatile Vector_type s_B_mtx[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_A_mtx = &s_A_mtx[threadIdx.x - lane_id_mod_NxN];\n volatile Vector_type *my_s_B_mtx = &s_B_mtx[threadIdx.x - lane_id_mod_NxN];\n // Shared memory to store the index of the element Aji.\n __shared__ volatile int s_A_ji[NUM_WARPS_PER_CTA * NUM_ITEMS_PER_WARP_CEIL];\n // Each thread keeps its own pointer.\n volatile int *my_s_A_ji = &s_A_ji[warp_id * NUM_ITEMS_PER_WARP_CEIL + lane_id_div_NxN];\n // Determine which NxN block the threads work with.\n int a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + warp_id * NUM_ITEMS_PER_WARP + lane_id_div_NxN;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; utils::any( a_row_it < num_rows_per_color ) ; a_row_it += NUM_ITEMS_PER_GRID )\n {\n // Is the thread active? For example, for 5x5 only the first 25 threads are active per warp.\n // At compile time, the compiler will see is_active == true for 2x2 (since NxN & (NxN-1) evals\n // to false ; that's the common trick to determine if a number is a power of 2).\n int is_active = true;\n\n if ( NxN & (NxN - 1) )\n {\n is_active = lane_id_div_NxN < NUM_ITEMS_PER_WARP;\n }\n\n int a_row_id = -1;\n\n if ( is_active && a_row_it < num_rows_per_color )\n {\n a_row_id = sorted_rows_by_color[a_row_it];\n }\n\n // Load the diagonal.\n Vector_type e_out(0);\n\n if ( a_row_id != -1 )\n {\n e_out = A_vals[NxN * A_diag[a_row_id] + lane_id_mod_NxN];\n }\n\n // Skip the 1st iteration of the outer-loop (that loop runs on the host).\n if ( current_color != 0 )\n {\n // Ranges of the rows.\n int a_col_begin(0), a_col_end(0);\n\n if ( a_row_id != -1 )\n {\n a_col_begin = A_rows[a_row_id ];\n a_col_end = A_rows[a_row_id + 1];\n }\n\n // Iterate over the elements in the columns.\n for ( ; a_col_begin < a_col_end ; a_col_begin += NxN )\n {\n unsigned int active_mask = utils::activemask();\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id_mod_NxN;\n // The identifier of the column if the iterator is valid.\n int a_col_tmp = -1, a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_tmp = A_cols[a_col_it];\n }\n\n if ( a_col_tmp != -1 && row_colors[a_col_tmp] < current_color )\n {\n a_col_id = a_col_tmp;\n }\n\n // When the diagonal is stored inside the matrix, we have to reject it. We\n // could be using a template parameter but it's not needed since that\n // rejection is really cheap (a couple of extra cycles -- CMP+MOV).\n if ( a_col_id == a_row_id )\n {\n a_col_id = -1;\n }\n\n // We partition valid and invalid column ids. Valid ones come first.\n int vote = utils::ballot( a_col_id != -1, active_mask ) & mask_NxN;\n int ones = __popc( vote );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( a_col_id == -1 )\n {\n dest = ones + lane_id_mod_NxN - dest;\n }\n\n my_s_a_col_ids[dest] = a_col_id;\n my_s_a_col_its[dest] = a_col_it;\n\n // Threads collaborate to load the rows.\n for ( int k = 0 ; k < NxN ; ++k )\n {\n // Exchange column indices.\n const int uniform_a_col_id = my_s_a_col_ids[k];\n\n // Early exit.\n if ( utils::all( uniform_a_col_id == -1, active_mask ) )\n {\n break;\n }\n\n // Load the iterator.\n const int uniform_a_col_it = my_s_a_col_its[k];\n // Load the two matrices.\n Vector_type my_A(0), my_B(0);\n\n if ( uniform_a_col_id != -1 )\n {\n my_A = A_vals[NxN * uniform_a_col_it + lane_id_mod_NxN];\n my_B = Einv [NxN * uniform_a_col_id + lane_id_mod_NxN];\n }\n\n my_s_A_mtx[lane_id_mod_NxN] = my_A;\n my_s_B_mtx[lane_id_mod_NxN] = my_B;\n utils::syncwarp(active_mask);\n // Compute the product of matrices.\n Vector_type tmp(0);\n#pragma unroll\n\n for ( int m = 0 ; m < N ; ++m )\n {\n tmp += my_s_A_mtx[N * lane_id_mod_NxN_div_N + m] * my_s_B_mtx[N * m + lane_id_mod_NxN_mod_N];\n }\n\n my_s_A_mtx[lane_id_mod_NxN] = tmp;\n // We looking for columns in the two rows we're interested in.\n int b_col_it(0), b_col_end(0);\n\n if ( is_active && uniform_a_col_id != -1 )\n {\n b_col_it = A_rows[uniform_a_col_id ];\n b_col_end = A_rows[uniform_a_col_id + 1];\n }\n\n // Init the marker to -1.\n if ( lane_id_mod_NxN == 0 )\n {\n *my_s_A_ji = -1;\n }\n\n // Run the loop.\n b_col_it += lane_id_mod_NxN;\n int shared_found = utils::ballot( lane_id_mod_NxN == 0 && uniform_a_col_id == -1, active_mask );\n\n do\n {\n bool found = b_col_it < b_col_end && A_cols[b_col_it] == a_row_id;\n\n if ( found )\n {\n *my_s_A_ji = b_col_it;\n }\n\n shared_found = shared_found | utils::ballot(found, active_mask);\n b_col_it += NxN;\n }\n while ( __popc( shared_found ) < NUM_ITEMS_PER_WARP && utils::any( b_col_it < b_col_end, active_mask ) );\n\n // Load the blocks.\n const int w_aji = *my_s_A_ji;\n Vector_type my_C(0);\n\n if ( w_aji != -1 )\n {\n my_C = A_vals[NxN * w_aji + lane_id_mod_NxN];\n }\n\n my_s_B_mtx[lane_id_mod_NxN] = my_C;\n // Update e_out.\n#pragma unroll\n\n for ( int m = 0 ; m < N ; ++m )\n {\n e_out -= my_s_A_mtx[N * lane_id_mod_NxN_div_N + m] * my_s_B_mtx[N * m + lane_id_mod_NxN_mod_N];\n }\n }\n } // a_col_begin < a_col_end\n } // current_color != 0\n\n // Store e_out in A\n my_s_A_mtx[lane_id_mod_NxN] = e_out;\n utils::syncwarp();\n\n // Invert the matrices.\n#pragma unroll\n\n for ( int row = 0 ; row < N ; ++row )\n {\n Vector_type diag(0), diag_tmp = my_s_A_mtx[N * row + row];\n\n if ( isNotCloseToZero(diag_tmp) )\n {\n diag = Vector_type(1) / diag_tmp;\n }\n else\n {\n diag = Vector_type(1) / epsilon(diag_tmp);\n }\n\n if ( is_active && lane_id_mod_NxN_div_N == 0 && lane_id_mod_NxN_mod_N != row )\n {\n my_s_A_mtx[N * row + lane_id_mod_NxN_mod_N] *= diag;\n }\n utils::syncwarp();\n\n if ( is_active && lane_id_mod_NxN_div_N != row && lane_id_mod_NxN_mod_N != row )\n {\n my_s_A_mtx[lane_id_mod_NxN] -= my_s_A_mtx[N * lane_id_mod_NxN_div_N + row] * my_s_A_mtx[N * row + lane_id_mod_NxN_mod_N];\n }\n utils::syncwarp();\n\n if ( is_active && lane_id_mod_NxN_div_N == 0 )\n {\n Vector_type tmp = diag;\n\n if ( lane_id_mod_NxN_mod_N != row )\n {\n tmp = -my_s_A_mtx[N * lane_id_mod_NxN_mod_N + row] * diag;\n }\n\n my_s_A_mtx[N * lane_id_mod_NxN_mod_N + row] = tmp;\n }\n utils::syncwarp();\n }\n\n // Store the results to Einv.\n if ( a_row_id != -1 )\n {\n Einv[NxN * a_row_id + lane_id_mod_NxN] = my_s_A_mtx[lane_id_mod_NxN];\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, int NUM_THREADS_PER_ROW, int CTA_SIZE, int WARP_SIZE >\n__global__\n__launch_bounds__( CTA_SIZE, 16 )\nvoid DILU_setup_1x1_kernel( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict A_diag,\n const Matrix_type *__restrict A_vals,\n Matrix_type *__restrict Einv,\n const int *sorted_rows_by_color,\n const int *row_colors,\n const int num_rows_per_color,\n const int current_color )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Number of items per grid.\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_NTPR = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_NTPR = lane_id % NUM_THREADS_PER_ROW;\n // Shared memory to broadcast column IDs.\n __shared__ int s_a_col_ids[CTA_SIZE];\n // Each thread keeps its own pointer.\n int *my_s_a_col_ids = &s_a_col_ids[warp_id * WARP_SIZE];\n // Shared memory to store the matrices.\n __shared__ int s_A_ji[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n int *my_s_A_ji = &s_A_ji[warp_id * WARP_SIZE];\n // Determine which NxN block the threads work with.\n int a_row_it = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_WARPS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load the diagonal.\n Vector_type e_out(0);\n\n // Skip the 1st iteration of the outer-loop (that loop runs on the host).\n if ( current_color != 0 )\n {\n // Ranges of the row.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Iterate over the elements in the columns.\n for ( ; a_col_begin < a_col_end ; a_col_begin += WARP_SIZE )\n {\n // Each thread loads a single element.\n int a_col_it = a_col_begin + lane_id;\n // The identifier of the column if the iterator is valid.\n int a_col_tmp = -1, a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_tmp = A_cols[a_col_it];\n }\n\n if ( a_col_tmp != -1 && row_colors[a_col_tmp] < current_color )\n {\n a_col_id = a_col_tmp;\n }\n\n // When the diagonal is stored inside the matrix, we have to reject it. We\n // could be using a template parameter but it's not needed since that\n // rejection is really cheap (a couple of extra cycles -- CMP+MOV).\n if ( a_col_id == a_row_id )\n {\n a_col_id = -1;\n }\n\n // We partition valid and invalid column ids. Valid ones come first.\n int vote = utils::ballot( a_col_id != -1 );\n int ones = __popc( vote );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( a_col_id == -1 )\n {\n dest = ones + lane_id - dest;\n }\n\n my_s_a_col_ids[dest] = a_col_id;\n // Reset A_jis.\n my_s_A_ji[lane_id] = -1;\n\n __syncwarp();\n\n // Threads collaborate to load the rows.\n for ( int k = 0 ; k < ones ; k += WARP_SIZE / NUM_THREADS_PER_ROW )\n {\n const int local_k = k + lane_id_div_NTPR;\n // Exchange column indices.\n int uniform_a_col_id = -1;\n\n if ( local_k < ones )\n {\n uniform_a_col_id = my_s_a_col_ids[local_k];\n }\n\n // We look for columns in the rows we're interested in.\n int b_col_it(0), b_col_end(0);\n\n if ( uniform_a_col_id != -1 )\n {\n b_col_it = A_rows[uniform_a_col_id ];\n b_col_end = A_rows[uniform_a_col_id + 1];\n }\n\n // Run the loop.\n b_col_it += lane_id_mod_NTPR;\n int shared_found = utils::ballot( lane_id_mod_NTPR == 0 && uniform_a_col_id == -1 );\n\n do\n {\n bool found = b_col_it < b_col_end && A_cols[b_col_it] == a_row_id;\n\n if ( found )\n {\n my_s_A_ji[local_k] = b_col_it;\n }\n\n shared_found = shared_found | utils::ballot(found);\n b_col_it += NUM_THREADS_PER_ROW;\n }\n while ( __popc( shared_found ) < WARP_SIZE / NUM_THREADS_PER_ROW && utils::any( b_col_it < b_col_end ) );\n }\n\n __syncwarp();\n\n // Where to get my A_ji from (if any).\n int a_ji_it = my_s_A_ji[dest];\n // Grab A_jis.\n Matrix_type a_ji(0);\n\n if ( a_ji_it != -1 )\n {\n a_ji = A_vals[a_ji_it];\n }\n\n // Update e_out.\n if ( a_col_id != -1 )\n {\n e_out += a_ji * Einv[a_col_id] * A_vals[a_col_it];\n }\n } // a_col_begin < a_col_end\n } // current_color != 0\n\n // Reduce the e_outs in one value.\n#pragma unroll\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n e_out += utils::shfl_xor( e_out, mask );\n }\n\n // Store the result.\n if ( lane_id == 0 )\n {\n Matrix_type res = A_vals[A_diag[a_row_id]] - e_out;\n\n if ( res != Matrix_type(0) )\n {\n res = Matrix_type(1) / res;\n }\n\n Einv[a_row_id] = static_cast(res);\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int N, bool ROW_MAJOR, int WARP_SIZE, typename Value_type >\nstatic __device__ __forceinline__\nValue_type reduce_distributed_vectors( Value_type x, int is_leader )\n{\n if ( N & (N - 1) )\n {\n#pragma unroll\n\n for ( int i = 1 ; i < N ; ++i )\n {\n Value_type other_x = utils::shfl_down( x, ROW_MAJOR ? i : N * i );\n\n if ( is_leader )\n {\n x += other_x;\n }\n }\n }\n else\n {\n#pragma unroll\n\n for ( int i = 1 ; i < N ; i <<= 1 )\n {\n x += utils::shfl_xor( x, ROW_MAJOR ? i : N * i );\n }\n }\n\n return x;\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, int N, int CTA_SIZE, int WARP_SIZE, bool ROW_MAJOR, bool HAS_EXTERNAL_DIAG >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_forward_NxN_kernel( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int NUM_ITEMS_PER_WARP = WARP_SIZE / NxN;\n // Number of items computer per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_ITEMS_PER_WARP * NUM_WARPS_PER_CTA;\n // Number of items per grid.\n const int NUM_ITEMS_PER_GRID = gridDim.x * NUM_ITEMS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_NxN = lane_id / NxN;\n const int lane_id_mod_NxN = lane_id % NxN;\n // Useful index to compute matrix products.\n const int lane_id_mod_NxN_div_N = lane_id_mod_NxN / N;\n const int lane_id_mod_NxN_mod_N = lane_id_mod_NxN % N;\n // We to get my data from when I use SHFL.\n const int shfl_offset = lane_id - lane_id_mod_NxN;\n // Shared memory needed to exchange X and delta.\n __shared__ volatile Vector_type s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_mem = &s_mem[threadIdx.x - lane_id_mod_NxN];\n // Is the thread active? For example, for 5x5 only the first 25 threads are active per warp.\n // At compile time, the compiler will see is_active == true for 2x2 (since NxN & (NxN-1) evals\n // to false ; that's the common trick to determine if a number is a power of 2).\n int is_active = true;\n\n if ( NxN & (NxN - 1) )\n {\n is_active = lane_id_div_NxN < NUM_ITEMS_PER_WARP;\n }\n\n // Determine which NxN block the threads work with.\n int a_row_it = num_rows_per_color;\n\n if ( is_active )\n {\n a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + warp_id * NUM_ITEMS_PER_WARP + lane_id_div_NxN;\n }\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_ITEMS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_bmAx(0);\n\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n my_bmAx = __cachingLoad(&b[N * a_row_id + lane_id_mod_NxN_div_N]);\n }\n }\n else\n {\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n my_bmAx = b[N * a_row_id + lane_id_mod_NxN_mod_N];\n }\n }\n\n // Don't do anything if X is zero.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // If the diagonal is stored separately, we have a special treatment.\n int a_col_max = a_col_end;\n\n if ( HAS_EXTERNAL_DIAG )\n {\n ++a_col_max;\n }\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( a_col_begin < a_col_max ) ; a_col_begin += NxN )\n {\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id_mod_NxN;\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n if ( HAS_EXTERNAL_DIAG && a_col_it == a_col_end )\n {\n a_col_id = a_row_id;\n }\n\n // Determine if the color is valid.\n int a_col_is_valid = false;\n#ifdef AMGX_ILU_COLORING\n if ( a_col_id != -1 && current_color != 0 )\n {\n if ( boundary_coloring == FIRST )\n {\n a_col_is_valid = a_col_id >= boundary_index;\n }\n else\n {\n a_col_is_valid = a_col_id < boundary_index && row_colors[a_col_id] < current_color;\n }\n }\n\n#else\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n a_col_is_valid = row_colors[a_col_id] < current_color;\n }\n\n#endif\n // Count the number of active columns.\n // int vote = utils::ballot(aColId != -1);\n // The number of iterations.\n // int nCols = max( __popc( vote & 0x0000ffff ), __popc( vote & 0xffff0000 ) );\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < NxN ; k += N )\n {\n int my_k = k + lane_id_mod_NxN_div_N;\n // Load N blocks of X.\n int uniform_a_col_id = utils::shfl( a_col_id, shfl_offset + my_k );\n int uniform_a_col_is_valid = utils::shfl( a_col_is_valid, shfl_offset + my_k );\n Vector_type my_x(0);\n\n if ( uniform_a_col_id != -1 )\n {\n my_x = __cachingLoad(&x[N * uniform_a_col_id + lane_id_mod_NxN_mod_N]);\n }\n\n if ( uniform_a_col_id != -1 && uniform_a_col_is_valid )\n {\n my_x += delta[N * uniform_a_col_id + lane_id_mod_NxN_mod_N];\n }\n\n my_s_mem[lane_id_mod_NxN] = my_x;\n // Load N blocks of A.\n#pragma unroll\n\n for ( int i = 0 ; i < N ; ++i )\n {\n int uniform_a_col_tmp = a_col_begin + k + i, uniform_a_col_it = -1;\n\n if ( uniform_a_col_tmp < a_col_end )\n {\n uniform_a_col_it = uniform_a_col_tmp;\n }\n\n if ( HAS_EXTERNAL_DIAG && is_active && uniform_a_col_tmp == a_col_end )\n {\n uniform_a_col_it = A_diag[a_row_id];\n }\n\n Matrix_type my_val(0);\n\n if ( uniform_a_col_it != -1 )\n {\n my_val = A_vals[NxN * uniform_a_col_it + lane_id_mod_NxN];\n }\n\n if ( ROW_MAJOR )\n {\n my_bmAx -= my_val * my_s_mem[N * i + lane_id_mod_NxN_mod_N];\n }\n else\n {\n my_bmAx -= my_val * my_s_mem[N * i + lane_id_mod_NxN_div_N];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Load Einvs.\n Vector_type my_Einv = Einv[NxN * a_row_id + lane_id_mod_NxN];\n // Reduce bmAx terms.\n int is_leader = lane_id_mod_NxN_div_N == 0;\n\n if ( ROW_MAJOR )\n {\n is_leader = lane_id_mod_NxN_mod_N == 0;\n }\n\n my_bmAx = reduce_distributed_vectors( my_bmAx, is_leader );\n\n // Update the shared terms.\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n my_s_mem[lane_id_mod_NxN_div_N] = my_bmAx;\n }\n }\n else\n {\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n my_s_mem[lane_id_mod_NxN_mod_N] = my_bmAx;\n }\n }\n\n // Update the diagonal term.\n if ( ROW_MAJOR )\n {\n my_bmAx = my_Einv * my_s_mem[lane_id_mod_NxN_mod_N];\n }\n else\n {\n my_bmAx = my_Einv * my_s_mem[lane_id_mod_NxN_div_N];\n }\n\n // Reduce bmAx terms.\n my_bmAx = reduce_distributed_vectors( my_bmAx, is_leader );\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n delta[N * a_row_id + lane_id_mod_NxN_div_N] = my_bmAx;\n }\n }\n else\n {\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n delta[N * a_row_id + lane_id_mod_NxN_mod_N] = my_bmAx;\n }\n }\n }\n}\n\ntemplate< typename Matrix_type, typename Vector_type, int N, int CTA_SIZE, int WARP_SIZE, bool HAS_EXTERNAL_DIAG, int NUM_WARP_ITERS_PER_BLOCK >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_forward_NxN_kernel_large( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of rows computed per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_WARPS_PER_CTA;\n // Number of rows? per grid.\n const int NUM_ITEMS_PER_GRID = CTA_SIZE;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n // Useful index to compute matrix products.\n const int lane_id_div_N = lane_id / N;\n const int lane_id_mod_N = lane_id % N; // id of a lane inside the block\n const int blocks_per_warp = WARP_SIZE / N; // we process this cols per warp per row\n const int row_elems_per_warp = blocks_per_warp * N;\n // Shared to store bmAx\n __shared__ volatile Vector_type bmAx[CTA_SIZE * NUM_WARP_ITERS_PER_BLOCK];\n volatile Vector_type *my_bmAx_s = &bmAx[warp_id * NUM_WARP_ITERS_PER_BLOCK * WARP_SIZE];\n // Determine which NxN block the threads work with.\n int a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_ITEMS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_bmAx(0);\n\n if ( lane_id < N )\n {\n my_bmAx = __cachingLoad(&b[N * a_row_id + lane_id]);\n }\n\n#pragma unroll\n\n for (int i = 0; i < NUM_WARP_ITERS_PER_BLOCK; i++)\n {\n my_bmAx_s[WARP_SIZE * i + lane_id] = 0.0;\n }\n\n // Don't do anything if X is zero.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // If the diagonal is stored separately, we have a special treatment.\n int a_col_max = a_col_end;\n\n if ( HAS_EXTERNAL_DIAG )\n {\n ++a_col_max;\n }\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( a_col_begin < a_col_max ) ; a_col_begin += WARP_SIZE ) // NxN\n {\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id;\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n if ( HAS_EXTERNAL_DIAG && a_col_it == a_col_end )\n {\n a_col_id = a_row_id;\n }\n\n // Determine if the color is valid.\n int a_col_is_valid = false;\n#ifdef AMGX_ILU_COLORING\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n if ( boundary_coloring == FIRST )\n {\n a_col_is_valid = a_col_id >= boundary_index;\n }\n else\n {\n a_col_is_valid = a_col_id < boundary_index && row_colors[a_col_id] < current_color;\n }\n }\n\n#else\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n a_col_is_valid = row_colors[a_col_id] < current_color;\n }\n\n#endif\n\n // Loop over columns. We compute blocks_per_warp columns per iteration.\n for ( int k = 0 ; k < WARP_SIZE ; k += blocks_per_warp )\n {\n // id of the processed block by this thread\n int my_k = k + lane_id_div_N;\n // Load N blocks of X (if valid)\n int uniform_a_col_id = utils::shfl( a_col_id, my_k );\n int uniform_a_col_is_valid = utils::shfl( a_col_is_valid, my_k );\n Vector_type my_x(0);\n\n if ( uniform_a_col_id != -1 && lane_id < row_elems_per_warp)\n {\n my_x = __cachingLoad(&x[N * uniform_a_col_id + lane_id_mod_N]);\n }\n\n if ( uniform_a_col_id != -1 && uniform_a_col_is_valid && lane_id < row_elems_per_warp)\n {\n my_x += delta[N * uniform_a_col_id + lane_id_mod_N];\n }\n\n //my_s_mem[lane_id] = my_x;\n#pragma unroll\n\n for ( int i = 0 ; i < blocks_per_warp ; ++i )\n {\n // k-th batch of blocks, i-th block. each thread process a column/row of a_it = uniform_a_col_tmp\n int uniform_a_col_tmp = a_col_begin + k + i, uniform_a_col_it = -1;\n\n // check if we are going out of bounds/color\n if ( uniform_a_col_tmp < a_col_end )\n {\n uniform_a_col_it = uniform_a_col_tmp;\n }\n\n if ( HAS_EXTERNAL_DIAG && uniform_a_col_tmp == a_col_end )\n {\n uniform_a_col_it = A_diag[a_row_id];\n }\n\n // swipe with the whole warp\n if (uniform_a_col_it != -1)\n {\n int block_inside_id = lane_id;\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n Matrix_type my_val(0);\n\n if ( uniform_a_col_it != -1 && block_inside_id < NxN)\n {\n my_val = A_vals[NxN * uniform_a_col_it + block_inside_id];\n }\n\n my_bmAx_s[block_inside_id] -= my_val * utils::shfl(my_x, N * i + block_inside_id % N); // MOD IS SLOW!\n block_inside_id += WARP_SIZE;\n }\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Load Einvs.\n Vector_type my_Einv[NUM_WARP_ITERS_PER_BLOCK];\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n my_Einv[j] = 0.0;\n }\n\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n if ((WARP_SIZE * j + lane_id) < NxN)\n {\n my_Einv[j] = Einv[NxN * a_row_id + WARP_SIZE * j + lane_id];\n }\n }\n\n // Reduce bmAx terms.\n {\n#pragma unroll\n\n for ( int i = 0 ; i < N ; ++i )\n {\n if ( lane_id < N )\n {\n my_bmAx += my_bmAx_s[N * lane_id + i];\n }\n }\n }\n // Update the diagonal term.\n int block_inside_id = lane_id;\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n my_bmAx_s[block_inside_id] = my_Einv[j] * utils::shfl(my_bmAx, block_inside_id % N);\n block_inside_id += WARP_SIZE;\n }\n\n // Reduce bmAx terms.\n {\n my_bmAx = 0.0;\n#pragma unroll\n\n for ( int i = 0 ; i < N ; ++i )\n {\n int idx = N * lane_id + i;\n\n if ( lane_id < N )\n {\n my_bmAx += my_bmAx_s[idx];\n }\n }\n }\n\n // Store the results.\n if ( lane_id < N )\n {\n delta[N * a_row_id + lane_id] = my_bmAx;\n }\n }\n}\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, int CTA_SIZE, int WARP_SIZE, bool ROW_MAJOR, bool HAS_EXTERNAL_DIAG >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_forward_4x4_kernel( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Number of items per warp.\n const int NUM_ITEMS_PER_WARP = WARP_SIZE / 16;\n // Number of items computer per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_ITEMS_PER_WARP * NUM_WARPS_PER_CTA;\n // Number of items per grid.\n const int NUM_ITEMS_PER_GRID = gridDim.x * NUM_ITEMS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_mod_16 = lane_id % 16;\n // Useful index to compute matrix products.\n const int lane_id_mod_16_div_4 = lane_id_mod_16 / 4;\n const int lane_id_mod_16_mod_4 = lane_id_mod_16 % 4;\n // We to get my data from when I use SHFL.\n const int shfl_offset = lane_id - lane_id_mod_16;\n // Shared memory needed to exchange X and delta.\n __shared__ volatile Vector_type s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_mem = &s_mem[threadIdx.x - lane_id_mod_16];\n // Determine which 16 block the threads work with.\n int a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + threadIdx.x / 16;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_ITEMS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_bmAx(0);\n\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_16_mod_4 == 0 )\n {\n my_bmAx = __cachingLoad(&b[4 * a_row_id + lane_id_mod_16_div_4]);\n }\n }\n else\n {\n if ( lane_id_mod_16_div_4 == 0 )\n {\n my_bmAx = b[4 * a_row_id + lane_id_mod_16_mod_4];\n }\n }\n\n // Don't do anything if X is zero.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // If the diagonal is stored separately, we have a special treatment.\n int a_col_max = a_col_end;\n\n if ( HAS_EXTERNAL_DIAG )\n {\n ++a_col_max;\n }\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( a_col_begin < a_col_max ) ; a_col_begin += 16 )\n {\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id_mod_16;\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n if ( HAS_EXTERNAL_DIAG && a_col_it == a_col_end )\n {\n a_col_id = a_row_id;\n }\n\n // Determine if the color is valid.\n int a_col_is_valid = false;\n#ifdef AMGX_ILU_COLORING\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n if ( boundary_coloring == FIRST )\n {\n a_col_is_valid = a_col_id >= boundary_index;\n }\n else\n {\n a_col_is_valid = a_col_id < boundary_index && row_colors[a_col_id] < current_color;\n }\n }\n\n#else\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n a_col_is_valid = row_colors[a_col_id] < current_color;\n }\n\n#endif\n // Count the number of active columns.\n // int vote = utils::ballot(aColId != -1);\n // The number of iterations.\n // int nCols = max( __popc( vote & 0x0000ffff ), __popc( vote & 0xffff0000 ) );\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < 16 ; k += 4 )\n {\n int my_k = k + lane_id_mod_16_div_4;\n // Load N blocks of X.\n int uniform_a_col_id = utils::shfl( a_col_id, shfl_offset + my_k );\n int uniform_a_col_is_valid = utils::shfl( a_col_is_valid, shfl_offset + my_k );\n Vector_type my_x(0);\n\n if ( uniform_a_col_id != -1 )\n {\n my_x = __cachingLoad(&x[4 * uniform_a_col_id + lane_id_mod_16_mod_4]);\n }\n\n if ( uniform_a_col_id != -1 && uniform_a_col_is_valid )\n {\n my_x += delta[4 * uniform_a_col_id + lane_id_mod_16_mod_4];\n }\n\n my_s_mem[lane_id_mod_16] = my_x;\n // Load N blocks of A.\n#pragma unroll\n\n for ( int i = 0 ; i < 4 ; ++i )\n {\n int uniform_a_col_tmp = a_col_begin + k + i, uniform_a_col_it = -1;\n\n if ( uniform_a_col_tmp < a_col_end )\n {\n uniform_a_col_it = uniform_a_col_tmp;\n }\n\n if ( HAS_EXTERNAL_DIAG && uniform_a_col_tmp == a_col_end )\n {\n uniform_a_col_it = A_diag[a_row_id];\n }\n\n Matrix_type my_val(0);\n\n if ( uniform_a_col_it != -1 )\n {\n my_val = A_vals[16 * uniform_a_col_it + lane_id_mod_16];\n }\n\n if ( ROW_MAJOR )\n {\n my_bmAx -= my_val * my_s_mem[4 * i + lane_id_mod_16_mod_4];\n }\n else\n {\n my_bmAx -= my_val * my_s_mem[4 * i + lane_id_mod_16_div_4];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Load Einvs.\n Vector_type my_Einv = Einv[16 * a_row_id + lane_id_mod_16];\n // Reduce bmAx terms.\n int is_leader = lane_id_mod_16_div_4 == 0;\n\n if ( ROW_MAJOR )\n {\n is_leader = lane_id_mod_16_mod_4 == 0;\n }\n\n my_bmAx = reduce_distributed_vectors<4, ROW_MAJOR, WARP_SIZE>( my_bmAx, is_leader );\n\n // Update the shared terms.\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_16_mod_4 == 0 )\n {\n my_s_mem[lane_id_mod_16_div_4] = my_bmAx;\n }\n }\n else\n {\n if ( lane_id_mod_16_div_4 == 0 )\n {\n my_s_mem[lane_id_mod_16_mod_4] = my_bmAx;\n }\n }\n\n // Update the diagonal term.\n if ( ROW_MAJOR )\n {\n my_bmAx = my_Einv * my_s_mem[lane_id_mod_16_mod_4];\n }\n else\n {\n my_bmAx = my_Einv * my_s_mem[lane_id_mod_16_div_4];\n }\n\n // Reduce bmAx terms.\n my_bmAx = reduce_distributed_vectors<4, ROW_MAJOR, WARP_SIZE>( my_bmAx, is_leader );\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_16_mod_4 == 0 )\n {\n delta[4 * a_row_id + lane_id_mod_16_div_4] = my_bmAx;\n }\n }\n else\n {\n if ( lane_id_mod_16_div_4 == 0 )\n {\n delta[4 * a_row_id + lane_id_mod_16_mod_4] = my_bmAx;\n }\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, int CTA_SIZE, bool HAS_EXTERNAL_DIAG >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_forward_4x4_kernel_row_major_vec4( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *Einv,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n // Number of half warps per CTA.\n const int NUM_HALF_WARPS = CTA_SIZE / 16;\n // Coordinates of the thread.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Coordinates of the thread in the CTA.\n const int thread_id_div_16 = threadIdx.x / 16;\n const int thread_id_mod_16 = threadIdx.x % 16;\n // Useful constants.\n const int thread_id_mod_16_div_4 = thread_id_mod_16 / 4;\n const int thread_id_mod_16_mod_4 = thread_id_mod_16 % 4;\n const int shfl_offset = 16 * (lane_id / 16);\n // Shared memory needed to exchange X and delta.\n __shared__ volatile Vector_type s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_mem = &s_mem[16 * thread_id_div_16];\n // The iterator over rows.\n int a_row_it = blockIdx.x * NUM_HALF_WARPS + thread_id_div_16;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += gridDim.x * NUM_HALF_WARPS )\n {\n unsigned int active_mask = utils::activemask();\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_bmAx(0);\n\n if ( thread_id_mod_16_div_4 == 0 )\n {\n my_bmAx = __cachingLoad(&b[4 * a_row_id + thread_id_mod_16_mod_4]);\n }\n\n // The range of the row.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // If it has an external diagonal, we need one more item to put the diag.\n int a_col_max = a_col_end;\n\n if ( HAS_EXTERNAL_DIAG )\n {\n ++a_col_max;\n }\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; a_col_begin < a_col_max ; a_col_begin += 16 )\n {\n unsigned int active_mask_inner = utils::activemask();\n int a_col_it = a_col_begin + thread_id_mod_16;\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = __cachingLoad(&A_cols[a_col_it]);\n }\n\n if ( HAS_EXTERNAL_DIAG && a_col_it == a_col_end )\n {\n a_col_id = a_row_id;\n }\n\n // Determine if the color is valid.\n int a_col_is_valid = false;\n#ifdef AMGX_ILU_COLORING\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n if ( boundary_coloring == FIRST )\n {\n a_col_is_valid = a_col_id >= boundary_index;\n }\n else\n {\n a_col_is_valid = a_col_id < boundary_index && __cachingLoad(&row_colors[a_col_id]) < current_color;\n }\n }\n\n#else\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n a_col_is_valid = row_colors[a_col_id] < current_color;\n }\n\n#endif\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < 16 ; k += 4 )\n {\n int my_k = k + thread_id_mod_16_div_4;\n // Load 8 blocks of X.\n int uniform_a_col_id = utils::shfl( a_col_id, shfl_offset + my_k, warpSize, active_mask_inner );\n int uniform_a_col_is_valid = utils::shfl( a_col_is_valid, shfl_offset + my_k, warpSize, active_mask_inner );\n Vector_type my_x(0);\n\n if ( uniform_a_col_id != -1 )\n {\n my_x = __cachingLoad(&x[4 * uniform_a_col_id + thread_id_mod_16_mod_4]);\n }\n\n if ( uniform_a_col_id != -1 && uniform_a_col_is_valid )\n {\n my_x += delta[4 * uniform_a_col_id + thread_id_mod_16_mod_4];\n }\n\n my_s_mem[thread_id_mod_16] = my_x;\n int uniform_a_col_tmp = a_col_begin + my_k, uniform_a_col_it = -1;\n\n if ( uniform_a_col_tmp < a_col_end )\n {\n uniform_a_col_it = uniform_a_col_tmp;\n }\n\n if ( HAS_EXTERNAL_DIAG && uniform_a_col_tmp == a_col_end )\n {\n uniform_a_col_it = A_diag[a_row_id];\n }\n\n Matrix_type my_vals[4] = { Matrix_type(0) };\n\n if ( uniform_a_col_it != -1 )\n {\n utils::load_vec4( my_vals, &A_vals[16 * uniform_a_col_it + 4 * thread_id_mod_16_mod_4] );\n }\n\n my_bmAx -= my_vals[0] * my_s_mem[4 * thread_id_mod_16_div_4 + 0];\n my_bmAx -= my_vals[1] * my_s_mem[4 * thread_id_mod_16_div_4 + 1];\n my_bmAx -= my_vals[2] * my_s_mem[4 * thread_id_mod_16_div_4 + 2];\n my_bmAx -= my_vals[3] * my_s_mem[4 * thread_id_mod_16_div_4 + 3];\n }\n }\n\n // Load Einvs.\n Matrix_type my_Einv = Einv[16 * a_row_id + thread_id_mod_16];\n // Reduce bmAx terms.\n my_bmAx += utils::shfl_xor( my_bmAx, 4, warpSize, active_mask );\n my_bmAx += utils::shfl_xor( my_bmAx, 8, warpSize, active_mask );\n\n // Update the shared terms.\n if ( thread_id_mod_16_div_4 == 0 )\n {\n my_s_mem[thread_id_mod_16_mod_4] = my_bmAx;\n }\n\n // Update the diagonal term.\n my_bmAx = my_Einv * my_s_mem[thread_id_mod_16_mod_4];\n // Reduce bmAx terms.\n my_bmAx += utils::shfl_xor( my_bmAx, 1, warpSize, active_mask );\n my_bmAx += utils::shfl_xor( my_bmAx, 2, warpSize, active_mask );\n\n // Store the results.\n if ( thread_id_mod_16_mod_4 == 0 )\n {\n delta[4 * a_row_id + thread_id_mod_16_div_4] = my_bmAx;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, int NUM_THREADS_PER_ROW, int CTA_SIZE, int WARP_SIZE, bool HAS_EXTERNAL_DIAG >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_forward_1x1_kernel( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n // Number of items per CTA.\n const int NUM_ROWS_PER_CTA = CTA_SIZE / NUM_THREADS_PER_ROW;\n // Number of items per grid.\n const int NUM_ROWS_PER_GRID = gridDim.x * NUM_ROWS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_mod_NTPR = lane_id % NUM_THREADS_PER_ROW;\n // Determine which NxN block the threads work with.\n int a_row_it = blockIdx.x * NUM_ROWS_PER_CTA + (threadIdx.x / NUM_THREADS_PER_ROW);\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_ROWS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_bmAx = amgx::types::util::get_zero();\n\n if ( lane_id_mod_NTPR == 0 )\n {\n my_bmAx = __cachingLoad(&b[a_row_id]);\n }\n\n // If it has an external diag.\n if ( HAS_EXTERNAL_DIAG && lane_id_mod_NTPR == 0 )\n {\n my_bmAx -= A_vals[A_diag[a_row_id]] * x[a_row_id];\n }\n\n // Don't do anything if X is zero.\n int a_col_it = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n\n // If the diagonal is stored separately, we have a special treatment.\n //if( HAS_EXTERNAL_DIAG )\n // ++a_col_end;\n\n // Each warp load column indices of 32 nonzero blocks\n for ( a_col_it += lane_id_mod_NTPR ; utils::any( a_col_it < a_col_end ) ; a_col_it += NUM_THREADS_PER_ROW )\n {\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Ignore the diagonal element since its color is smaller, and been accounted for above\n if (HAS_EXTERNAL_DIAG && a_col_id == a_row_id)\n {\n a_col_id = -1;\n }\n\n // Load x.\n Vector_type my_x(0);\n\n if ( a_col_id != -1 )\n {\n my_x = __cachingLoad(&x[a_col_id]);\n }\n\n // Is it really a valid column (due to coloring).\n int valid = false;\n#ifdef AMGX_ILU_COLORING\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n if ( boundary_coloring == FIRST )\n {\n valid = a_col_id >= boundary_index;\n }\n else\n {\n valid = a_col_id < boundary_index && row_colors[a_col_id] < current_color;\n }\n }\n\n#else\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n valid = row_colors[a_col_tmp] < current_color;\n }\n\n#endif\n\n // Load my x value.\n if ( valid )\n {\n my_x += delta[a_col_id];\n }\n\n // Load my item from A.\n Matrix_type my_val(0);\n\n if ( a_col_it < a_col_end )\n {\n my_val = A_vals[a_col_it];\n }\n\n // Update bmAx.\n my_bmAx -= my_val * my_x;\n }\n\n // Reduce bmAx terms.\n#pragma unroll\n\n for ( int mask = NUM_THREADS_PER_ROW / 2 ; mask > 0 ; mask >>= 1 )\n {\n my_bmAx += utils::shfl_xor( my_bmAx, mask );\n }\n\n // Store the results.\n if ( lane_id_mod_NTPR == 0 )\n {\n delta[a_row_id] = Einv[a_row_id] * my_bmAx;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType, int N, int CTA_SIZE, int WARP_SIZE, bool ROW_MAJOR >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_backward_NxN_kernel( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int NUM_ITEMS_PER_WARP = WARP_SIZE / NxN;\n // Number of items computer per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_ITEMS_PER_WARP * NUM_WARPS_PER_CTA;\n // Number of items per grid.\n const int NUM_ITEMS_PER_GRID = gridDim.x * NUM_ITEMS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_NxN = lane_id / NxN;\n const int lane_id_mod_NxN = lane_id % NxN;\n // Useful index to compute matrix products.\n const int lane_id_mod_NxN_div_N = lane_id_mod_NxN / N;\n const int lane_id_mod_NxN_mod_N = lane_id_mod_NxN % N;\n // We to get my data from when I use SHFL.\n const int shfl_offset = lane_id - lane_id_mod_NxN;\n // Shared memory needed to exchange X and delta.\n __shared__ volatile Vector_type s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_mem = &s_mem[threadIdx.x - lane_id_mod_NxN];\n // Is the thread active? For example, for 5x5 only the first 25 threads are active per warp.\n // At compile time, the compiler will see is_active == true for 2x2 (since NxN & (NxN-1) evals\n // to false ; that's the common trick to determine if a number is a power of 2).\n int is_active = true;\n\n if ( NxN & (NxN - 1) )\n {\n is_active = lane_id_div_NxN < NUM_ITEMS_PER_WARP;\n }\n\n // Determine which NxN block the threads work with.\n int a_row_it = num_rows_per_color;\n\n if ( is_active )\n {\n a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + warp_id * NUM_ITEMS_PER_WARP + lane_id_div_NxN;\n }\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_ITEMS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_delta(0);\n // Don't do anything if X is zero.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( a_col_begin < a_col_end ) ; a_col_begin += NxN )\n {\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id_mod_NxN;\n // Get the ID of the column.\n int a_col_tmp = -1, a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_tmp = A_cols[a_col_it];\n }\n\n // Make sure the column is interesting.\n#ifdef AMGX_ILU_COLORING\n int valid = false;\n\n if ( a_col_tmp != -1 && current_color != 0 )\n {\n if ( boundary_coloring == LAST )\n {\n valid = a_col_tmp >= boundary_index;\n }\n else\n {\n valid = a_col_tmp < boundary_index && row_colors[a_col_tmp] > current_color;\n }\n }\n\n#else\n int valid = false;\n\n if ( a_col_tmp != -1 && row_colors[a_col_tmp] > current_color )\n {\n valid = true;\n }\n\n#endif\n\n // Set the column id.\n if ( valid )\n {\n a_col_id = a_col_tmp;\n }\n\n // Count the number of active columns.\n // int vote = utils::ballot(aColId != -1);\n // The number of iterations.\n // int nCols = max( __popc( vote & 0x0000ffff ), __popc( vote & 0xffff0000 ) );\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < NxN ; k += N )\n {\n int my_k = k + lane_id_mod_NxN_div_N;\n // Load N blocks of X.\n int uniform_a_col_id = utils::shfl( a_col_id, shfl_offset + my_k );\n Vector_type my_x(0);\n\n if ( uniform_a_col_id != -1 )\n {\n my_x = Delta[N * uniform_a_col_id + lane_id_mod_NxN_mod_N];\n }\n\n my_s_mem[lane_id_mod_NxN] = my_x;\n // Load N blocks of A.\n#pragma unroll\n\n for ( int i = 0 ; i < N ; ++i )\n {\n //if( uniform_a_col_id == -1 )\n // break;\n int uniform_a_col_tmp = a_col_begin + k + i, uniform_a_col_it = -1;\n\n if ( uniform_a_col_tmp < a_col_end )\n {\n uniform_a_col_it = uniform_a_col_tmp;\n }\n\n Matrix_type my_val(0);\n\n if ( uniform_a_col_it != -1 )\n {\n my_val = A_vals[NxN * uniform_a_col_it + lane_id_mod_NxN];\n }\n\n if ( ROW_MAJOR )\n {\n my_delta += my_val * my_s_mem[N * i + lane_id_mod_NxN_mod_N];\n }\n else\n {\n my_delta += my_val * my_s_mem[N * i + lane_id_mod_NxN_div_N];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Load Einvs.\n Matrix_type my_Einv = Einv[NxN * a_row_id + lane_id_mod_NxN];\n // Reduce bmAx terms.\n int is_leader = lane_id_mod_NxN_div_N == 0;\n\n if ( ROW_MAJOR )\n {\n is_leader = lane_id_mod_NxN_mod_N == 0;\n }\n\n my_delta = reduce_distributed_vectors( my_delta, is_leader );\n\n // Update the shared terms.\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n my_s_mem[lane_id_mod_NxN_div_N] = my_delta;\n }\n }\n else\n {\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n my_s_mem[lane_id_mod_NxN_mod_N] = my_delta;\n }\n }\n\n // Update the diagonal term.\n if ( ROW_MAJOR )\n {\n my_delta = my_Einv * my_s_mem[lane_id_mod_NxN_mod_N];\n }\n else\n {\n my_delta = my_Einv * my_s_mem[lane_id_mod_NxN_div_N];\n }\n\n // Reduce bmAx terms.\n my_delta = reduce_distributed_vectors( my_delta, is_leader );\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n const int offset = N * a_row_id + lane_id_mod_NxN_div_N;\n Vector_type my_b(0), my_x(0);\n\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n my_b = __cachingLoad(&delta[offset]);\n my_x = x [offset];\n }\n\n my_delta = my_b - my_delta;\n\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n x [offset] = my_x + weight * my_delta;\n Delta[offset] = my_delta;\n }\n }\n else\n {\n const int offset = N * a_row_id + lane_id_mod_NxN_mod_N;\n Vector_type my_b(0), my_x(0);\n\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n my_b = __cachingLoad(&delta[offset]);\n my_x = x [offset];\n }\n\n my_delta = my_b - my_delta;\n\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n x [offset] = my_x + weight * my_delta;\n Delta[offset] = my_delta;\n }\n }\n }\n}\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType, int N, int CTA_SIZE, int WARP_SIZE, bool ROW_MAJOR, int NUM_WARP_ITERS_PER_BLOCK >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_backward_NxN_kernel_large( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items computer per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_WARPS_PER_CTA;\n // Number of items per grid.\n const int NUM_ITEMS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_N = lane_id / N;\n const int lane_id_mod_N = lane_id % N; // id of a lane inside the block\n const int blocks_per_warp = WARP_SIZE / N; // we process this cols per warp per row\n const int row_elems_per_warp = blocks_per_warp * N;\n // Shared to store t_delta\n __shared__ volatile Vector_type delta_s[CTA_SIZE * NUM_WARP_ITERS_PER_BLOCK];\n volatile Vector_type *my_delta_s = &delta_s[warp_id * NUM_WARP_ITERS_PER_BLOCK * WARP_SIZE];\n // Determine which NxN block the threads work with.\n int a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_ITEMS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Accumulator\n Vector_type my_delta(0);\n //Vector_type mAx[NUM_WARP_ITERS_PER_BLOCK];\n#pragma unroll\n\n for (int i = 0; i < NUM_WARP_ITERS_PER_BLOCK; i++)\n {\n my_delta_s[WARP_SIZE * i + lane_id] = 0.0;\n }\n\n // Don't do anything if X is zero.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( a_col_begin < a_col_end ) ; a_col_begin += WARP_SIZE )\n {\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id;\n // Get the ID of the column.\n int a_col_tmp = -1, a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_tmp = A_cols[a_col_it];\n }\n\n // Make sure the column is interesting.\n#ifdef AMGX_ILU_COLORING\n int valid = false;\n\n if ( a_col_tmp != -1 && current_color != 0 )\n {\n if ( boundary_coloring == LAST )\n {\n valid = a_col_tmp >= boundary_index;\n }\n else\n {\n valid = a_col_tmp < boundary_index && row_colors[a_col_tmp] > current_color;\n }\n }\n\n#else\n int valid = false;\n\n if ( a_col_tmp != -1 && row_colors[a_col_tmp] > current_color )\n {\n valid = true;\n }\n\n#endif\n\n // Set the column id.\n if ( valid )\n {\n a_col_id = a_col_tmp;\n }\n\n // Loop over columns. We compute blocks_per_warp columns per iteration.\n for ( int k = 0 ; k < WARP_SIZE ; k += blocks_per_warp )\n {\n // id of the processed block by this thread\n int my_k = k + lane_id_div_N;\n // Load N blocks of X (if valid)\n int uniform_a_col_id = utils::shfl( a_col_id, my_k );\n Vector_type my_x(0);\n\n if ( uniform_a_col_id != -1 && lane_id < row_elems_per_warp)\n {\n my_x = Delta[N * uniform_a_col_id + lane_id_mod_N];\n }\n\n // Load blocks of A.\n // for each block in a batch\n#pragma unroll\n\n for ( int i = 0 ; i < blocks_per_warp ; ++i )\n {\n // k-th batch of blocks, i-th block. each thread process a column/row of a_it = uniform_a_col_tmp\n int uniform_a_col_tmp = a_col_begin + k + i, uniform_a_col_it = -1;\n\n // check if we are going out of bounds/color\n if ( uniform_a_col_tmp < a_col_end )\n {\n uniform_a_col_it = uniform_a_col_tmp;\n }\n\n // swipe with the whole warp\n if (uniform_a_col_it != -1)\n {\n int block_inside_id = lane_id;\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n Matrix_type my_val(0);\n\n if ( uniform_a_col_it != -1 && block_inside_id < NxN)\n {\n my_val = A_vals[NxN * uniform_a_col_it + block_inside_id];\n }\n\n my_delta_s[block_inside_id] -= my_val * utils::shfl(my_x, N * i + block_inside_id % N); //my_s_mem[N*i + block_inside_id % N]; // MOD IS SLOW!\n block_inside_id += WARP_SIZE;\n }\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Load Einvs.\n Vector_type my_Einv[NUM_WARP_ITERS_PER_BLOCK];\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n my_Einv[j] = 0.0;\n }\n\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n if ((WARP_SIZE * j + lane_id) < NxN)\n {\n my_Einv[j] = Einv[NxN * a_row_id + WARP_SIZE * j + lane_id];\n }\n }\n\n // Reduce bmAx terms.\n {\n#pragma unroll\n\n for ( int i = 0 ; i < N ; ++i )\n {\n if ( lane_id < N )\n {\n my_delta += my_delta_s[N * lane_id + i];\n }\n }\n }\n\n // Update the diagonal term.\n if ( ROW_MAJOR )\n {\n int block_inside_id = lane_id;\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n my_delta_s[block_inside_id] = my_Einv[j] * utils::shfl(my_delta, block_inside_id % N);\n block_inside_id += WARP_SIZE;\n }\n }\n\n // Reduce bmAx terms.\n {\n my_delta = 0.0;\n#pragma unroll\n\n for ( int i = 0 ; i < N ; ++i )\n {\n if ( lane_id < N )\n {\n my_delta += my_delta_s[N * lane_id + i];\n }\n }\n }\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n const int offset = N * a_row_id + lane_id;\n Vector_type my_b(0), my_x(0);\n\n if ( lane_id < N )\n {\n my_b = __cachingLoad(&delta[offset]);\n my_x = x [offset];\n }\n\n my_delta = my_b - my_delta;\n\n if ( lane_id < N )\n {\n x [offset] = my_x + weight * my_delta;\n Delta[offset] = my_delta;\n }\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename IndexType, typename ValueTypeA, typename ValueTypeB, typename WeightType, int CTA_SIZE, bool ROW_MAJOR >\n__global__\n__launch_bounds__( CTA_SIZE, 16 )\nvoid DILU_backward_4x4_kernel( const IndexType *row_offsets,\n const IndexType *column_indices,\n const ValueTypeA *nonzero_values,\n ValueTypeB *x,\n const WeightType weight,\n const int *sorted_rows_by_color,\n const int *__restrict row_colors,\n const ValueTypeA *Einv,\n const ValueTypeB *delta,\n ValueTypeB *Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const IndexType boundary_index)\n{\n const int nHalfWarps = CTA_SIZE / 16; // Number of half warps per CTA.\n const int laneId = utils::lane_id();\n const int halfWarpId = threadIdx.x / 16;\n const int halfLaneId = threadIdx.x % 16;\n const int halfLaneId_div_4 = halfLaneId / 4;\n const int halfLaneId_mod_4 = halfLaneId % 4;\n const int upperHalf = 16 * (laneId / 16);\n // Shared memory needed to exchange X and delta.\n __shared__ volatile ValueTypeB s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile ValueTypeB *my_s_mem = &s_mem[16 * halfWarpId];\n\n // Iterate over the rows of the matrix. One warp per two rows.\n for ( int aRowIt = blockIdx.x * nHalfWarps + halfWarpId ; aRowIt < num_rows_per_color ; aRowIt += gridDim.x * nHalfWarps )\n {\n int aRowId = sorted_rows_by_color[aRowIt];\n // Load one block of B.\n ValueTypeB my_delta(0);\n // The range of the rows.\n int aColBegin = row_offsets[aRowId ];\n int aColEnd = row_offsets[aRowId + 1];\n\n // Each warp load column indices of 16 nonzero blocks\n for ( ; aColBegin < aColEnd ; aColBegin += 16 )\n {\n int aColIt = aColBegin + halfLaneId;\n // Get the ID of the column.\n int aColTmp = -1, aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColTmp = column_indices[aColIt];\n }\n\n#ifdef AMGX_ILU_COLORING\n bool valid = (((aColTmp < boundary_index || boundary_coloring == SYNC_COLORS) && (row_colors[aColTmp] > current_color)) || (aColTmp >= boundary_index && boundary_coloring == LAST));\n\n if ( aColTmp != -1 && valid )\n {\n aColId = aColTmp;\n }\n\n#else\n\n if ( aColTmp != -1 && row_colors[aColTmp] > current_color )\n {\n aColId = aColTmp;\n }\n\n#endif\n for ( int k = 0 ; k < 16 ; k += 4 )\n {\n int my_k = k + halfLaneId_div_4;\n // Exchange column indices.\n int waColId = utils::shfl( aColId, upperHalf + my_k );\n // Load 8 blocks of X if needed.\n ValueTypeB my_x(0);\n\n if ( waColId != -1 )\n {\n my_x = Delta[4 * waColId + halfLaneId_mod_4];\n }\n\n my_s_mem[halfLaneId] = my_x;\n // Load 8 blocks of A.\n#pragma unroll\n\n for ( int i = 0 ; i < 4 ; ++i )\n {\n const int k_i = k + i;\n int w_aColTmp = aColBegin + k_i, w_aColIt = -1;\n if ( utils::shfl( aColId, upperHalf + k_i ) != -1 && w_aColTmp < aColEnd )\n w_aColIt = w_aColTmp;\n\n ValueTypeA my_val(0);\n\n if ( w_aColIt != -1 )\n {\n my_val = nonzero_values[16 * w_aColIt + halfLaneId];\n }\n\n if ( ROW_MAJOR )\n {\n my_delta += my_val * my_s_mem[4 * i + halfLaneId_mod_4];\n }\n else\n {\n my_delta += my_val * my_s_mem[4 * i + halfLaneId_div_4];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Load EINV values.\n ValueTypeA my_Einv = Einv[16 * aRowId + halfLaneId];\n // Reduce delta terms.\n if ( ROW_MAJOR )\n {\n my_delta += utils::shfl_xor( my_delta, 1 );\n my_delta += utils::shfl_xor( my_delta, 2 );\n }\n else\n {\n my_delta += utils::shfl_xor( my_delta, 4 );\n my_delta += utils::shfl_xor( my_delta, 8 );\n }\n\n // Update the shared terms.\n if ( ROW_MAJOR )\n {\n if ( halfLaneId_mod_4 == 0 )\n {\n my_s_mem[halfLaneId_div_4] = my_delta;\n }\n }\n else\n {\n if ( halfLaneId_div_4 == 0 )\n {\n my_s_mem[halfLaneId_mod_4] = my_delta;\n }\n }\n\n // Update the diagonal term.\n if ( ROW_MAJOR )\n {\n my_delta = my_Einv * my_s_mem[halfLaneId_mod_4];\n }\n else\n {\n my_delta = my_Einv * my_s_mem[halfLaneId_div_4];\n }\n\n // Regroup results.\n if ( ROW_MAJOR )\n {\n my_delta += utils::shfl_xor( my_delta, 1 );\n my_delta += utils::shfl_xor( my_delta, 2 );\n }\n else\n {\n my_delta += utils::shfl_xor( my_delta, 4 );\n my_delta += utils::shfl_xor( my_delta, 8 );\n }\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n int offset = 4 * aRowId + halfLaneId_div_4;\n ValueTypeB my_b(0), my_x(0);\n\n if ( halfLaneId_mod_4 == 0 )\n {\n my_b = __cachingLoad(&delta[offset]);\n my_x = x[offset];\n }\n\n my_delta = my_b - my_delta;\n\n if ( halfLaneId_mod_4 == 0 )\n {\n x[offset] = my_x + weight * my_delta;\n Delta[offset] = my_delta;\n }\n }\n else\n {\n int offset = 4 * aRowId + halfLaneId_mod_4;\n ValueTypeB my_b(0), my_x(0);\n\n if ( halfLaneId_div_4 == 0 )\n {\n my_b = __cachingLoad(&delta[offset]);\n my_x = x[offset];\n }\n\n my_delta = my_b - my_delta;\n\n if ( halfLaneId_div_4 == 0 )\n {\n x[offset] = my_x + weight * my_delta;\n Delta[offset] = my_delta;\n }\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType, int CTA_SIZE >\n__global__\n__launch_bounds__( CTA_SIZE, 16 )\nvoid DILU_backward_4x4_kernel_row_major_vec4( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n // Number of half warps per CTA.\n const int NUM_HALF_WARPS = CTA_SIZE / 16;\n // Coordinates of the thread.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Coordinates of the thread in the CTA.\n const int thread_id_div_16 = threadIdx.x / 16;\n const int thread_id_mod_16 = threadIdx.x % 16;\n // Useful constants.\n const int thread_id_mod_16_div_4 = thread_id_mod_16 / 4;\n const int thread_id_mod_16_mod_4 = thread_id_mod_16 % 4;\n const int shfl_offset = 16 * (lane_id / 16);\n // Shared memory needed to exchange X and delta.\n __shared__ volatile Vector_type s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_mem = &s_mem[16 * thread_id_div_16];\n // The iterator over rows.\n int a_row_it = blockIdx.x * NUM_HALF_WARPS + thread_id_div_16;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += gridDim.x * NUM_HALF_WARPS )\n {\n unsigned int active_mask = utils::activemask();\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_delta(0);\n // Don't do anything if X is zero.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; a_col_begin < a_col_end ; a_col_begin += 16 )\n {\n unsigned int active_mask_inner = utils::activemask();\n int a_col_it = a_col_begin + thread_id_mod_16;\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = __cachingLoad(&A_cols[a_col_it]);\n }\n\n#ifdef AMGX_ILU_COLORING\n int valid = false;\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n if ( boundary_coloring == LAST )\n {\n valid = a_col_id >= boundary_index;\n }\n else\n {\n valid = a_col_id < boundary_index && __cachingLoad(&row_colors[a_col_id]) > current_color;\n }\n }\n\n#else\n int valid = false;\n\n if ( a_col_id != -1 && row_colors[a_col_id] > current_color )\n {\n valid = true;\n }\n\n#endif\n\n // Set the column id.\n if ( !valid )\n {\n a_col_id = -1;\n }\n\n // Loop over columns. We compute 8 columns per iteration.\n#pragma unroll 2\n\n for ( int k = 0 ; k < 16 ; k += 4 )\n {\n int my_k = k + thread_id_mod_16_div_4;\n // Load 8 blocks of X.\n int uniform_a_col_id = utils::shfl( a_col_id, shfl_offset + my_k, warpSize, active_mask_inner );\n Vector_type my_Delta(0);\n\n if ( uniform_a_col_id != -1 )\n {\n my_Delta = Delta[4 * uniform_a_col_id + thread_id_mod_16_mod_4];\n }\n\n my_s_mem[thread_id_mod_16] = my_Delta;\n int uniform_a_col_it = a_col_begin + my_k;\n\n if ( uniform_a_col_id == -1 || uniform_a_col_it >= a_col_end )\n {\n uniform_a_col_it = -1;\n }\n\n Matrix_type my_vals[4] = { Matrix_type(0) };\n\n if ( uniform_a_col_it != -1 )\n {\n utils::load_vec4( my_vals, &A_vals[16 * uniform_a_col_it + 4 * thread_id_mod_16_mod_4] );\n }\n\n my_delta += my_vals[0] * my_s_mem[4 * thread_id_mod_16_div_4 + 0];\n my_delta += my_vals[1] * my_s_mem[4 * thread_id_mod_16_div_4 + 1];\n my_delta += my_vals[2] * my_s_mem[4 * thread_id_mod_16_div_4 + 2];\n my_delta += my_vals[3] * my_s_mem[4 * thread_id_mod_16_div_4 + 3];\n } // Loop over k\n } // Loop over aColIt\n\n // Load EINV values.\n Matrix_type my_Einv = Einv[16 * a_row_id + thread_id_mod_16];\n // Reduce delta terms.\n my_delta += utils::shfl_xor( my_delta, 4, warpSize, active_mask );\n my_delta += utils::shfl_xor( my_delta, 8, warpSize, active_mask );\n\n // Update the shared terms.\n if ( thread_id_mod_16_div_4 == 0 )\n {\n my_s_mem[thread_id_mod_16_mod_4] = my_delta;\n }\n\n // Update the diagonal term.\n my_delta = my_Einv * my_s_mem[thread_id_mod_16_mod_4];\n // Regroup results.\n my_delta += utils::shfl_xor( my_delta, 1, warpSize, active_mask );\n my_delta += utils::shfl_xor( my_delta, 2, warpSize, active_mask );\n // Store the results.\n int offset = 4 * a_row_id + thread_id_mod_16_div_4;\n Vector_type my_b(0), my_x(0);\n\n if ( thread_id_mod_16_mod_4 == 0 )\n {\n my_b = __cachingLoad(&delta[offset]);\n my_x = x [offset];\n }\n\n my_delta = my_b - my_delta;\n\n if ( thread_id_mod_16_mod_4 == 0 )\n {\n x [offset] = my_x + weight * my_delta;\n Delta[offset] = my_delta;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType, int NUM_THREADS_PER_ROW, int CTA_SIZE, int WARP_SIZE >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_backward_1x1_kernel( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n // Number of items per CTA.\n const int NUM_ROWS_PER_CTA = CTA_SIZE / NUM_THREADS_PER_ROW;\n // Number of items per grid.\n const int NUM_ROWS_PER_GRID = gridDim.x * NUM_ROWS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_mod_NTPR = lane_id % NUM_THREADS_PER_ROW;\n // Determine which NxN block the threads work with.\n int a_row_it = blockIdx.x * NUM_ROWS_PER_CTA + (threadIdx.x / NUM_THREADS_PER_ROW);\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_ROWS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_delta(0);\n // Don't do anything if X is zero.\n int a_col_it = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Each warp load column indices of 32 nonzero blocks\n for ( a_col_it += lane_id_mod_NTPR ; utils::any( a_col_it < a_col_end ) ; a_col_it += NUM_THREADS_PER_ROW )\n {\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it really a valid column (due to coloring).\n int valid = false;\n#ifdef AMGX_ILU_COLORING\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n //if( boundary_coloring == LAST )\n // valid = a_col_id >= boundary_index;\n //else\n // valid = a_col_id < boundary_index && row_colors[a_col_id] > current_color;\n valid = (((a_col_id < boundary_index || boundary_coloring == SYNC_COLORS) && (row_colors[a_col_id] > current_color)) || (a_col_id >= boundary_index && boundary_coloring == LAST));\n }\n\n#else\n\n //if( a_col_id != -1 && current_color != 0 )\n if ( a_col_id != -1 )\n {\n valid = row_colors[a_col_id] > current_color;\n }\n\n#endif\n // Load my Delta value.\n Vector_type my_Delta(0);\n\n if ( valid )\n {\n my_Delta = Delta[a_col_id];\n }\n\n // Load my item from A.\n Matrix_type my_val(0);\n\n if ( valid )\n {\n my_val = A_vals[a_col_it];\n }\n\n // Update bmAx.\n my_delta += my_val * my_Delta;\n }\n\n // Reduce bmAx terms.\n#pragma unroll\n\n for ( int mask = NUM_THREADS_PER_ROW / 2 ; mask > 0 ; mask >>= 1 )\n {\n my_delta += utils::shfl_xor( my_delta, mask );\n }\n\n // Store the results.\n if ( lane_id_mod_NTPR == 0 )\n {\n Vector_type my_x = __cachingLoad(&delta[a_row_id]) - Einv[a_row_id] * my_delta;\n x [a_row_id] += weight * my_x;\n Delta[a_row_id] = my_x;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType, int N, int CTA_SIZE >\n__global__\n__launch_bounds__( CTA_SIZE, 16 )\nvoid DILU_backward_NxN_kernel_skip( Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color )\n{\n const int NUM_ITEMS_PER_CTA = CTA_SIZE / N; // Number of updated block items per CTA\n const int ITEM_ID = threadIdx.x / N;\n const int ITEM_BLOCK_OFFSET = threadIdx.x % N;\n const int is_active = ITEM_ID < NUM_ITEMS_PER_CTA;\n // The first row.\n int a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + ITEM_ID;\n\n // Iterate over the rows of the matrix. One warp per two rows.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += gridDim.x * NUM_ITEMS_PER_CTA )\n {\n if ( is_active )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n const int idx = N * a_row_id + ITEM_BLOCK_OFFSET;\n Vector_type my_b = __cachingLoad(&delta[idx]);\n Vector_type my_x = x[idx];\n x[idx] = my_x + weight * my_b;\n Delta[idx] = my_b;\n }\n }\n}\n\n// ----------\n// Methods\n// ----------\n\ntemplate< typename Matrix_type, typename Vector_type, int N >\nvoid DILU_forward_NxN_dispatch( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const ColoringType boundary_coloring,\n const int boundary_index,\n const int row_major,\n const int has_external_diag )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int NUM_ROWS_PER_WARP = std::max(WARP_SIZE / NxN, 1);\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = NUM_ROWS_PER_WARP * NUM_WARPS_PER_CTA;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n // Branch to the correct kernel call.\n int code = 2 * (row_major ? 1 : 0) + (has_external_diag ? 1 : 0);\n\n switch ( code )\n {\n case 0: // Column-major, no external diagonal.\n DILU_forward_NxN_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n break;\n\n case 1: // Column-major, external diagonal.\n DILU_forward_NxN_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n break;\n\n case 2: // Row-major, no external diagonal.\n DILU_forward_NxN_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n break;\n\n case 3: // Row-major, external diagonal.\n DILU_forward_NxN_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n break;\n\n default:\n FatalError( \"Internal error\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n cudaCheckError();\n}\n\ntemplate< typename Matrix_type, typename Vector_type, int N >\nvoid DILU_forward_NxN_dispatch_large( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const ColoringType boundary_coloring,\n const int boundary_index,\n const int row_major,\n const int has_external_diag )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = NUM_WARPS_PER_CTA;\n // Each warp is going to sweep through bloock this many times\n const int NUM_WARP_ITERS_PER_BLOCK = ((NxN - 1) / WARP_SIZE) + 1;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n\n // Branch to the correct kernel call.\n if (!row_major)\n {\n FatalError(\"COL MAJOR is not supported for this large block_size\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n switch ( has_external_diag )\n {\n case 0: // Row-major, no external diagonal.\n DILU_forward_NxN_kernel_large <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n break;\n\n case 1: // Row-major, external diagonal.\n DILU_forward_NxN_kernel_large <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n break;\n\n default:\n FatalError( \"Internal error\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n cudaCheckError();\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type >\nvoid DILU_forward_NxN_dispatch( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const ColoringType boundary_coloring,\n const int boundary_index,\n const int block_size,\n const int row_major,\n const int has_external_diag )\n{\n switch ( block_size )\n {\n case 1:\n {\n const int NUM_THREADS_PER_ROW = 8;\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = CTA_SIZE / NUM_THREADS_PER_ROW;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n\n if ( has_external_diag )\n {\n DILU_forward_1x1_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n }\n else\n {\n DILU_forward_1x1_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n }\n\n cudaCheckError();\n }\n break;\n\n case 2:\n DILU_forward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index,\n row_major,\n has_external_diag );\n break;\n\n case 3:\n DILU_forward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index,\n row_major,\n has_external_diag );\n break;\n\n case 4:\n if ( row_major )\n {\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = CTA_SIZE / 16;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n\n if ( has_external_diag )\n //DILU_forward_4x4_kernel<<>>(\n DILU_forward_4x4_kernel_row_major_vec4 <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n else\n DILU_forward_4x4_kernel_row_major_vec4 <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n\n cudaCheckError();\n }\n else\n DILU_forward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index,\n row_major,\n has_external_diag );\n\n break;\n\n case 5:\n DILU_forward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index,\n row_major,\n has_external_diag );\n break;\n\n case 8:\n DILU_forward_NxN_dispatch_large(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index,\n row_major,\n has_external_diag );\n break;\n\n case 10:\n DILU_forward_NxN_dispatch_large(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index,\n row_major,\n has_external_diag );\n break;\n\n default:\n FatalError( \"Internal error\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n cudaCheckError();\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType, int N >\nvoid DILU_backward_NxN_dispatch( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const int boundary_index,\n const int row_major )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int NUM_ROWS_PER_WARP = std::max(WARP_SIZE / NxN, 1);\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = NUM_ROWS_PER_WARP * NUM_WARPS_PER_CTA;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n\n // Branch to the correct kernel call.\n if ( row_major )\n {\n DILU_backward_NxN_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index );\n }\n else\n {\n DILU_backward_NxN_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index );\n }\n\n cudaCheckError();\n}\n\n\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType, int N, int NUM_WARP_ITERS_PER_BLOCK >\nvoid DILU_backward_NxN_dispatch_large( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const int boundary_index,\n const int row_major )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = NUM_WARPS_PER_CTA;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n\n // Branch to the correct kernel call.\n if ( row_major )\n {\n DILU_backward_NxN_kernel_large <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index );\n }\n else\n {\n FatalError(\"col major is not supported for this blocksize in multicolor DILU solver\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n cudaCheckError();\n}\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType>\nvoid DILU_backward_NxN_dispatch( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const int boundary_index,\n const int block_size,\n const int row_major )\n{\n switch ( block_size )\n {\n case 1:\n {\n const int NUM_THREADS_PER_ROW = 8;\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = CTA_SIZE / NUM_THREADS_PER_ROW;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n DILU_backward_1x1_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index );\n cudaCheckError();\n }\n break;\n\n case 2:\n DILU_backward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index,\n row_major );\n break;\n\n case 3:\n DILU_backward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index,\n row_major );\n break;\n\n case 4:\n\n //if( false )\n if ( row_major )\n {\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = CTA_SIZE / 16;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n //DILU_backward_NxN_kernel<<>>(\n DILU_backward_4x4_kernel_row_major_vec4 <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index );\n cudaCheckError();\n }\n else\n DILU_backward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index,\n row_major );\n\n break;\n\n case 5:\n DILU_backward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index,\n row_major );\n break;\n\n case 8:\n DILU_backward_NxN_dispatch_large(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index,\n row_major );\n break;\n\n case 10:\n DILU_backward_NxN_dispatch_large(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index,\n row_major );\n break;\n\n default:\n FatalError( \"Internal error\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n cudaCheckError();\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< class T_Config >\nMulticolorDILUSolver_Base::MulticolorDILUSolver_Base( AMG_Config &cfg,\n const std::string &cfg_scope,\n ThreadManager *tmng ) :\n Solver( cfg, cfg_scope, tmng )\n{\n this->weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n this->m_reorder_cols_by_color_desired = (cfg.AMG_Config::template getParameter(\"reorder_cols_by_color\", cfg_scope) != 0);\n this->m_insert_diagonal_desired = (cfg.AMG_Config::template getParameter(\"insert_diag_while_reordering\", cfg_scope) != 0);\n this->m_boundary_coloring = cfg.AMG_Config::template getParameter(\"boundary_coloring\", cfg_scope);\n this->always_obey_coloring = 0;\n\n if (weight == 0)\n {\n weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Multicolor DILU smoother\\n\");\n }\n}\n\n// Destructor\ntemplate\nMulticolorDILUSolver_Base::~MulticolorDILUSolver_Base()\n{\n Einv.clear();\n Einv.shrink_to_fit();\n}\n\ntemplate\nvoid MulticolorDILUSolver_Base::computeEinv(Matrix &A)\n{\n ViewType oldView = A.currentView();\n A.setView(this->m_explicit_A->getViewExterior());\n\n if ( A.get_block_dimx() != A.get_block_dimy() )\n {\n FatalError(\"DILU implemented only for squared blocks\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if ( A.get_block_dimx() > 32) // actually much more less than 32 doe to register file limitations, but...\n {\n FatalError(\"DILU implemented only for squared blocks of size <= 32\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n computeEinv_NxN( A, A.get_block_dimx() );\n A.setView(oldView);\n}\n\ntemplate< class T_Config >\nvoid\nMulticolorDILUSolver_Base::printSolverParameters() const\n{\n std::cout << \"relaxation_factor = \" << this->weight << std::endl;\n}\n\n// Solver setup\ntemplate< class T_Config >\nvoid\nMulticolorDILUSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n m_explicit_A = dynamic_cast*>(this->m_A);\n\n if (!this->m_explicit_A)\n {\n FatalError(\"MulticolorDILUSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n int N = this->m_explicit_A->get_num_cols() * this->m_explicit_A->get_block_dimy();\n\n if (this->m_explicit_A->getColoringLevel() < 1)\n {\n FatalError(\"Matrix must be colored to use multicolor dilu solver. Try setting: coloring_level=1 in the configuration file\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n m_delta.resize(N);\n m_Delta.resize(N);\n m_delta.set_block_dimy(this->m_explicit_A->get_block_dimy());\n m_Delta.set_block_dimy(this->m_explicit_A->get_block_dimy());\n m_delta.set_block_dimx(1);\n m_Delta.set_block_dimx(1);\n\n if ( this->m_explicit_A->getBlockFormat() != ROW_MAJOR )\n {\n FatalError(\"Multicolor DILU solver only supports row major format for the blocks\", AMGX_ERR_CONFIGURATION);\n }\n\n computeEinv( *this->m_explicit_A );\n}\n\n//\ntemplate< class T_Config >\nvoid\nMulticolorDILUSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nMulticolorDILUSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n if ( this->m_explicit_A->get_block_dimx() != this->m_explicit_A->get_block_dimy() )\n {\n FatalError(\"DILU implemented only for squared blocks\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if ( this->m_explicit_A->get_block_dimx() > 32) // actually much more less than 32 doe to register file limitations, but...\n {\n FatalError(\"DILU implemented only for squared blocks of size <= 32\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if (xIsZero)\n {\n x.dirtybit = 0;\n }\n\n if (!this->m_explicit_A->is_matrix_singleGPU())\n {\n this->m_explicit_A->manager->exchange_halo_async(x, x.tag);\n this->m_explicit_A->manager->exchange_halo_async(b, b.tag);\n }\n\n if (this->m_explicit_A->getViewExterior() == this->m_explicit_A->getViewInterior())\n {\n if (!this->m_explicit_A->is_matrix_singleGPU())\n {\n this->m_explicit_A->manager->exchange_halo_wait(x, x.tag);\n this->m_explicit_A->manager->exchange_halo_wait(b, b.tag);\n }\n }\n\n ViewType oldView = this->m_explicit_A->currentView();\n ViewType flags;\n bool latencyHiding = true;\n\n if (this->m_explicit_A->is_matrix_singleGPU() || (x.dirtybit == 0 && b.dirtybit == 0))\n {\n latencyHiding = false;\n this->m_explicit_A->setViewExterior();\n flags = (ViewType)(this->m_explicit_A->getViewExterior());\n }\n else\n {\n flags = (ViewType)(this->m_explicit_A->getViewInterior());\n this->m_explicit_A->setViewInterior();\n }\n\n if (xIsZero)\n {\n thrust_wrapper::fill(x.begin(), x.end(), types::util::get_zero());\n cudaCheckError();\n }\n\n this->smooth_NxN(*this->m_explicit_A, b, x, flags);\n\n if (latencyHiding)\n {\n if (!this->m_explicit_A->is_matrix_singleGPU())\n {\n this->m_explicit_A->manager->exchange_halo_wait(x, x.tag);\n this->m_explicit_A->manager->exchange_halo_wait(b, b.tag);\n }\n\n this->m_explicit_A->setViewExterior();\n flags = (ViewType)(~(this->m_explicit_A->getViewInterior()) & this->m_explicit_A->getViewExterior());\n\n if (flags != 0)\n {\n this->smooth_NxN(*this->m_explicit_A, b, x, flags);\n }\n }\n\n x.dirtybit = 1;\n this->m_explicit_A->setView(oldView);\n return (this->converged(b, x));\n}\n\ntemplate\nvoid\nMulticolorDILUSolver_Base::solve_finalize( VVector &b, VVector &x )\n{}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid MulticolorDILUSolver >::computeEinv_NxN(const Matrix_h &A, const int bsize)\n{\n FatalError(\"Multicolor DILU smoother not implemented for host format, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid MulticolorDILUSolver >::smooth_NxN( const Matrix_h &A, VVector &b, VVector &x, ViewType separation_flag )\n{\n FatalError(\"Haven't implemented Multicolor DILU smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nMulticolorDILUSolver >::MulticolorDILUSolver(\n AMG_Config &cfg,\n const std::string &cfg_scope,\n ThreadManager *tmng ) :\n MulticolorDILUSolver_Base >( cfg, cfg_scope, tmng )\n{\n int device = 0;\n cudaGetDevice( &device );\n cudaCheckError();\n cudaDeviceProp properties;\n cudaGetDeviceProperties( &properties, device );\n cudaCheckError();\n m_is_kepler = properties.major >= 3;\n}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid MulticolorDILUSolver >::computeEinv_NxN(const Matrix_d &A, const int bsize)\n{\n const int bsize_sq = bsize * bsize;\n this->Einv.resize( A.get_num_cols()*bsize_sq, 0.0 );\n\n // sol::prof_start();\n for ( int i = 0, num_colors = A.getMatrixColoring().getNumColors() ; i < num_colors ; ++i )\n {\n const int color_offset = A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const int num_rows_per_color = A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n\n if ( num_rows_per_color == 0 )\n {\n continue;\n }\n\n const int CTA_SIZE = 128;\n const int NUM_WARPS_PER_CTA = CTA_SIZE / 32;\n int ROWS_PER_WARP = 1;\n\n if ( bsize_sq > 1 && bsize_sq < 6 )\n {\n ROWS_PER_WARP = 32 / bsize_sq;\n }\n\n const int ROWS_PER_CTA = ROWS_PER_WARP * NUM_WARPS_PER_CTA;\n const int GRID_SIZE = std::min( 4096, (num_rows_per_color + ROWS_PER_CTA - 1) / ROWS_PER_CTA );\n cudaStream_t stream = amgx::thrust::global_thread_handle::get_stream();\n\n switch ( bsize )\n {\n case 1:\n DILU_setup_1x1_kernel <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->Einv.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n num_rows_per_color,\n i );\n break;\n\n case 2:\n DILU_setup_NxN_kernel <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->Einv.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n num_rows_per_color,\n i );\n break;\n\n case 3:\n DILU_setup_NxN_kernel <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->Einv.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n num_rows_per_color,\n i );\n break;\n\n case 4:\n DILU_setup_NxN_kernel <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->Einv.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n num_rows_per_color,\n i );\n break;\n\n case 5:\n DILU_setup_NxN_kernel <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->Einv.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n num_rows_per_color,\n i );\n break;\n\n case 8:\n DILU_setup_NxN_kernel_large <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->Einv.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n num_rows_per_color,\n i );\n break;\n\n case 10:\n DILU_setup_NxN_kernel_large <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->Einv.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n num_rows_per_color,\n i );\n break;\n\n default:\n FatalError( \"Multicolor-DILU Setup: block size was not enabled in the code, contact AMGX developers.\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE );\n }\n\n cudaCheckError();\n }\n}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid MulticolorDILUSolver >::smooth_NxN( const Matrix_d &A, VVector &b, VVector &x, ViewType separation_flag )\n{\n AMGX_CPU_PROFILER( \"MulticolorDILUSolver::smooth_NxN \" );\n int offset = 0, separation = 0;\n A.getOffsetAndSizeForView(INTERIOR, &offset, &separation);\n\n // Only have separation=num interior rows if we are only working on the interior\n // and the boundary coloring is FIRST or LAST, otherwise set separation offset to\n // total number of rows\n if ( separation_flag != this->m_explicit_A->getViewInterior() ||\n this->m_explicit_A->getViewExterior() == this->m_explicit_A->getViewInterior() ||\n this->m_boundary_coloring != LAST && this->m_boundary_coloring != FIRST )\n {\n separation = A.row_offsets.size() - 1;\n }\n else\n {\n amgx_printf(\"separation active\\n\");\n }\n\n // --------------------\n // Forward Sweep\n // --------------------\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n\n for ( int i = 0 ; i < num_colors ; ++i )\n {\n int color_offset(0);\n\n if ( separation_flag & INTERIOR )\n {\n color_offset = A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n }\n else\n {\n color_offset = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i];\n }\n\n int num_rows_per_color(0);\n\n if ( separation_flag == this->m_explicit_A->getViewInterior() )\n {\n num_rows_per_color = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i];\n }\n else\n {\n num_rows_per_color = A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1];\n }\n\n num_rows_per_color -= color_offset;\n\n if ( num_rows_per_color == 0 )\n {\n continue;\n }\n\n int boundary_index = separation;\n\n if ( this->m_boundary_coloring == SYNC_COLORS )\n {\n boundary_index = A.get_num_rows();\n }\n\n DILU_forward_NxN_dispatch(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n A.diag.raw(),\n x.raw(),\n b.raw(),\n this->m_delta.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n num_rows_per_color,\n i,\n A.getMatrixColoring().getRowColors().raw(),\n this->Einv.raw(),\n this->m_boundary_coloring,\n boundary_index,\n A.get_block_dimy(),\n A.getBlockFormat() == ROW_MAJOR,\n A.hasProps(DIAG) );\n cudaCheckError();\n }\n\n // --------------------\n // Backward Sweep\n // --------------------\n for ( int i = num_colors - 1 ; i >= 0 ; --i )\n {\n int color_offset(0);\n\n if ( separation_flag & INTERIOR )\n {\n color_offset = A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n }\n else\n {\n color_offset = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i];\n }\n\n int num_rows_per_color(0);\n\n if ( separation_flag == this->m_explicit_A->getViewInterior() )\n {\n num_rows_per_color = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i];\n }\n else\n {\n num_rows_per_color = A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1];\n }\n\n num_rows_per_color -= color_offset;\n\n if ( num_rows_per_color == 0 )\n {\n continue;\n }\n\n if ( i == num_colors - 1 )\n {\n const int NUM_ROWS_PER_CTA = CTA_SIZE / A.get_block_dimy();\n const int GRID_SIZE = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n\n switch ( A.get_block_dimy() )\n {\n case 1:\n DILU_backward_NxN_kernel_skip <<< GRID_SIZE, CTA_SIZE>>>(\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color );\n break;\n\n case 2:\n DILU_backward_NxN_kernel_skip <<< GRID_SIZE, CTA_SIZE>>>(\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color );\n break;\n\n case 3:\n DILU_backward_NxN_kernel_skip <<< GRID_SIZE, CTA_SIZE>>>(\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color );\n break;\n\n case 4:\n DILU_backward_NxN_kernel_skip <<< GRID_SIZE, CTA_SIZE>>>(\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color );\n break;\n\n case 5:\n DILU_backward_NxN_kernel_skip <<< GRID_SIZE, CTA_SIZE>>>(\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color );\n break;\n\n case 8:\n DILU_backward_NxN_kernel_skip <<< GRID_SIZE, CTA_SIZE>>>(\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color );\n break;\n\n case 10:\n DILU_backward_NxN_kernel_skip <<< GRID_SIZE, CTA_SIZE>>>(\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color );\n break;\n }\n\n cudaCheckError();\n }\n else\n {\n DILU_backward_NxN_dispatch(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n this->Einv.raw(),\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color,\n i,\n this->m_boundary_coloring,\n separation,\n A.get_block_dimy(),\n A.getBlockFormat() == ROW_MAJOR );\n cudaCheckError();\n }\n }\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class MulticolorDILUSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n// AMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class MulticolorDILUSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n// AMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n} // namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \n\n#define AMGX_ILU_COLORING\n\nnamespace amgx\n{\nnamespace multicolor_dilu_solver\n{\n\nenum { CTA_SIZE = 128, WARP_SIZE = 32 };\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, int N, int CTA_SIZE, int WARP_SIZE, int NUM_WARP_ITERS_PER_BLOCK >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_setup_NxN_kernel_large( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict A_diag,\n const Matrix_type *__restrict A_vals,\n Matrix_type *__restrict Einv,\n const int *sorted_rows_by_color,\n const int *row_colors,\n const int num_rows_per_color,\n const int current_color )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items computer per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_WARPS_PER_CTA;\n // Number of items per grid.\n const int NUM_ITEMS_PER_GRID = gridDim.x * NUM_ITEMS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Shared memory to broadcast column IDs.\n __shared__ volatile int s_a_col_ids[CTA_SIZE];\n __shared__ volatile int s_a_col_its[CTA_SIZE];\n // Each thread keeps its own pointer.\n volatile int *my_s_a_col_ids = &s_a_col_ids[threadIdx.x - lane_id];\n volatile int *my_s_a_col_its = &s_a_col_its[threadIdx.x - lane_id];\n // Shared memory to store the matrices.\n __shared__ volatile Vector_type s_A_mtx[CTA_SIZE * NUM_WARP_ITERS_PER_BLOCK];\n __shared__ volatile Vector_type s_B_mtx[CTA_SIZE * NUM_WARP_ITERS_PER_BLOCK];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_A_mtx = &s_A_mtx[warp_id * NUM_WARP_ITERS_PER_BLOCK * WARP_SIZE];\n volatile Vector_type *my_s_B_mtx = &s_B_mtx[warp_id * NUM_WARP_ITERS_PER_BLOCK * WARP_SIZE];\n // Shared memory to store the index of the element Aji.\n __shared__ volatile int s_A_ji[NUM_WARPS_PER_CTA];\n // Each thread keeps its own pointer.\n volatile int *my_s_A_ji = &s_A_ji[warp_id];\n // Precomputing some stuff\n int idx[NUM_WARP_ITERS_PER_BLOCK];\n int idy[NUM_WARP_ITERS_PER_BLOCK];\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n const int id = (WARP_SIZE * wb + lane_id) % NxN;\n idx[wb] = id / N;\n idy[wb] = id % N;\n }\n\n // Determine which NxN block the threads work with.\n int a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; utils::any( a_row_it < num_rows_per_color ) ; a_row_it += NUM_ITEMS_PER_GRID )\n {\n int a_row_id = -1;\n\n if ( a_row_it < num_rows_per_color )\n {\n a_row_id = sorted_rows_by_color[a_row_it];\n }\n\n // Load the diagonal.\n Vector_type e_out[NUM_WARP_ITERS_PER_BLOCK];\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n e_out[wb] = (Vector_type)0.0;\n }\n\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n if ( a_row_id != -1 && (wb * WARP_SIZE + lane_id) < NxN)\n {\n e_out[wb] = A_vals[NxN * A_diag[a_row_id] + wb * WARP_SIZE + lane_id];\n }\n\n // Skip the 1st iteration of the outer-loop (that loop runs on the host).\n if ( current_color != 0 )\n {\n // Ranges of the rows.\n int a_col_begin(0), a_col_end(0);\n\n if ( a_row_id != -1 )\n {\n a_col_begin = A_rows[a_row_id ];\n a_col_end = A_rows[a_row_id + 1];\n }\n\n // Iterate over the elements in the columns.\n for ( ; a_col_begin < a_col_end ; a_col_begin += NxN )\n {\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id;\n // The identifier of the column if the iterator is valid.\n int a_col_tmp = -1, a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_tmp = A_cols[a_col_it];\n }\n\n if ( a_col_tmp != -1 && row_colors[a_col_tmp] < current_color )\n {\n a_col_id = a_col_tmp;\n }\n\n // When the diagonal is stored inside the matrix, we have to reject it. We\n // could be using a template parameter but it's not needed since that\n // rejection is really cheap (a couple of extra cycles -- CMP+MOV).\n if ( a_col_id == a_row_id )\n {\n a_col_id = -1;\n }\n\n // We partition valid and invalid column ids. Valid ones come first.\n int vote = utils::ballot( a_col_id != -1 );\n int ones = __popc( vote );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( a_col_id == -1 )\n {\n dest = ones + lane_id - dest;\n }\n\n my_s_a_col_ids[dest] = a_col_id;\n my_s_a_col_its[dest] = a_col_it;\n // Temporary storage with zeros for OOB\n Vector_type my_A[NUM_WARP_ITERS_PER_BLOCK], my_B[NUM_WARP_ITERS_PER_BLOCK];\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n my_A[wb] = (Vector_type)0.0;\n my_B[wb] = (Vector_type)0.0;\n }\n\n // Threads collaborate to load the rows.\n for ( int k = 0 ; k < WARP_SIZE ; ++k )\n {\n // Exchange column indices.\n const int uniform_a_col_id = my_s_a_col_ids[k];\n\n // Early exit.\n if ( uniform_a_col_id == -1 )\n {\n break;\n }\n\n // Load the iterator.\n const int uniform_a_col_it = my_s_a_col_its[k];\n // Load the two matrices.\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n if ((wb * WARP_SIZE + lane_id) < NxN)\n {\n my_A[wb] = A_vals[NxN * uniform_a_col_it + wb * WARP_SIZE + lane_id];\n my_B[wb] = Einv [NxN * uniform_a_col_id + wb * WARP_SIZE + lane_id];\n }\n\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n my_s_A_mtx[lane_id + wb * WARP_SIZE] = my_A[wb];\n my_s_B_mtx[lane_id + wb * WARP_SIZE] = my_B[wb];\n }\n\n // Compute the product of matrices.\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n my_A[wb] = (Vector_type)0.0;\n#pragma unroll\n\n for ( int m = 0 ; m < N ; ++m )\n {\n my_A[wb] += my_s_A_mtx[N * idx[wb] + m] * my_s_B_mtx[N * m + idy[wb]];\n }\n }\n\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n if ((wb * WARP_SIZE + lane_id) < NxN)\n {\n my_s_A_mtx[lane_id + wb * WARP_SIZE] = my_A[wb];\n }\n\n // We looking for columns in the two rows we're interested in.\n int b_col_it = A_rows[uniform_a_col_id ];\n int b_col_end = A_rows[uniform_a_col_id + 1];\n\n // Init the marker to -1.\n if ( lane_id == 0 )\n {\n *my_s_A_ji = -1;\n }\n\n // Run the loop.\n b_col_it += lane_id;\n int shared_found = utils::ballot( lane_id == 0 && uniform_a_col_id == -1 );\n\n do\n {\n bool found = b_col_it < b_col_end && A_cols[b_col_it] == a_row_id;\n\n if ( found )\n {\n *my_s_A_ji = b_col_it;\n }\n\n shared_found = shared_found | utils::ballot(found);\n b_col_it += NxN;\n }\n while ( __popc( shared_found ) == 0 && utils::any( b_col_it < b_col_end ) );\n\n // Load the blocks.\n const int w_aji = *my_s_A_ji;\n Vector_type my_C[NUM_WARP_ITERS_PER_BLOCK];\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n my_C[wb] = (Vector_type)0.0;\n\n if ( w_aji != -1 && (wb * WARP_SIZE + lane_id) < NxN)\n {\n my_C[wb] = A_vals[NxN * w_aji + wb * WARP_SIZE + lane_id];\n }\n\n my_s_B_mtx[wb * WARP_SIZE + lane_id] = my_C[wb];\n }\n\n // Update e_out.\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n#pragma unroll\n\n for ( int m = 0 ; m < N ; ++m )\n {\n e_out[wb] -= my_s_A_mtx[N * idx[wb] + m] * my_s_B_mtx[N * m + idy[wb]];\n }\n }\n }\n } // a_col_begin < a_col_end\n } // current_color != 0\n\n // Store e_out in A\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n my_s_B_mtx[wb * WARP_SIZE + lane_id] = my_s_A_mtx[wb * WARP_SIZE + lane_id] = e_out[wb];\n }\n\n // Invert the matrices.\n#pragma unroll\n\n for ( int row = 0 ; row < N ; ++row )\n {\n Vector_type diag(0), diag_tmp = my_s_A_mtx[N * row + row];\n\n if ( isNotCloseToZero(diag_tmp) )\n {\n diag = Vector_type(1) / diag_tmp;\n }\n else\n {\n diag = Vector_type(1) / epsilon(diag_tmp);\n }\n\n if ( lane_id < N && lane_id != row)\n {\n my_s_A_mtx[N * row + lane_id] = my_s_B_mtx[N * row + lane_id] = my_s_B_mtx[N * row + lane_id] * diag;\n }\n\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n if ( idx[wb] != row && idy[wb] != row)\n {\n my_s_A_mtx[wb * WARP_SIZE + lane_id] = my_s_B_mtx[wb * WARP_SIZE + lane_id] - my_s_B_mtx[N * idx[wb] + row] * my_s_B_mtx[N * row + idy[wb]];\n }\n\n if ( lane_id < N )\n {\n Vector_type tmp = diag;\n\n if ( lane_id != row )\n {\n tmp = -my_s_A_mtx[N * lane_id + row] * diag;\n }\n\n my_s_A_mtx[N * lane_id + row] = tmp;\n }\n\n#pragma unroll\n\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n {\n my_s_B_mtx[wb * WARP_SIZE + lane_id] = my_s_A_mtx[wb * WARP_SIZE + lane_id];\n }\n }\n\n // Store the results to Einv.\n if ( a_row_id != -1 )\n#pragma unroll\n for (int wb = 0; wb < NUM_WARP_ITERS_PER_BLOCK; wb++)\n if (wb * WARP_SIZE + lane_id < NxN)\n {\n Einv[NxN * a_row_id + wb * WARP_SIZE + lane_id] = my_s_A_mtx[wb * WARP_SIZE + lane_id];\n }\n }\n}\n\n\ntemplate< typename Matrix_type, typename Vector_type, int N, int CTA_SIZE, int WARP_SIZE >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_setup_NxN_kernel( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict A_diag,\n const Matrix_type *__restrict A_vals,\n Matrix_type *__restrict Einv,\n const int *sorted_rows_by_color,\n const int *row_colors,\n const int num_rows_per_color,\n const int current_color )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int NUM_ITEMS_PER_WARP = WARP_SIZE / NxN;\n // Upper-bound on the number of items per warp.\n const int NUM_ITEMS_PER_WARP_CEIL = (WARP_SIZE + NxN - 1) / NxN;\n // Number of items computer per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_ITEMS_PER_WARP * NUM_WARPS_PER_CTA;\n // Number of items per grid.\n const int NUM_ITEMS_PER_GRID = gridDim.x * NUM_ITEMS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_NxN = lane_id / NxN;\n const int lane_id_mod_NxN = lane_id % NxN;\n // Useful index to compute matrix products.\n const int lane_id_mod_NxN_div_N = lane_id_mod_NxN / N;\n const int lane_id_mod_NxN_mod_N = lane_id_mod_NxN % N;\n // We need NxN to compute a NxN block. Encode a mask for the first block.\n int mask_tmp = utils::ballot( lane_id_div_NxN == 0 );\n // Mask for ballots. We shift the mask with NxN active bits by the needed number of bits.\n const int mask_NxN = mask_tmp << (lane_id_div_NxN * __popc(mask_tmp));\n // Shared memory to broadcast column IDs.\n __shared__ volatile int s_a_col_ids[CTA_SIZE];\n __shared__ volatile int s_a_col_its[CTA_SIZE];\n // Each thread keeps its own pointer.\n volatile int *my_s_a_col_ids = &s_a_col_ids[threadIdx.x - lane_id_mod_NxN];\n volatile int *my_s_a_col_its = &s_a_col_its[threadIdx.x - lane_id_mod_NxN];\n // Shared memory to store the matrices.\n __shared__ volatile Vector_type s_A_mtx[CTA_SIZE];\n __shared__ volatile Vector_type s_B_mtx[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_A_mtx = &s_A_mtx[threadIdx.x - lane_id_mod_NxN];\n volatile Vector_type *my_s_B_mtx = &s_B_mtx[threadIdx.x - lane_id_mod_NxN];\n // Shared memory to store the index of the element Aji.\n __shared__ volatile int s_A_ji[NUM_WARPS_PER_CTA * NUM_ITEMS_PER_WARP_CEIL];\n // Each thread keeps its own pointer.\n volatile int *my_s_A_ji = &s_A_ji[warp_id * NUM_ITEMS_PER_WARP_CEIL + lane_id_div_NxN];\n // Determine which NxN block the threads work with.\n int a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + warp_id * NUM_ITEMS_PER_WARP + lane_id_div_NxN;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; utils::any( a_row_it < num_rows_per_color ) ; a_row_it += NUM_ITEMS_PER_GRID )\n {\n // Is the thread active? For example, for 5x5 only the first 25 threads are active per warp.\n // At compile time, the compiler will see is_active == true for 2x2 (since NxN & (NxN-1) evals\n // to false ; that's the common trick to determine if a number is a power of 2).\n int is_active = true;\n\n if ( NxN & (NxN - 1) )\n {\n is_active = lane_id_div_NxN < NUM_ITEMS_PER_WARP;\n }\n\n int a_row_id = -1;\n\n if ( is_active && a_row_it < num_rows_per_color )\n {\n a_row_id = sorted_rows_by_color[a_row_it];\n }\n\n // Load the diagonal.\n Vector_type e_out(0);\n\n if ( a_row_id != -1 )\n {\n e_out = A_vals[NxN * A_diag[a_row_id] + lane_id_mod_NxN];\n }\n\n // Skip the 1st iteration of the outer-loop (that loop runs on the host).\n if ( current_color != 0 )\n {\n // Ranges of the rows.\n int a_col_begin(0), a_col_end(0);\n\n if ( a_row_id != -1 )\n {\n a_col_begin = A_rows[a_row_id ];\n a_col_end = A_rows[a_row_id + 1];\n }\n\n // Iterate over the elements in the columns.\n for ( ; a_col_begin < a_col_end ; a_col_begin += NxN )\n {\n unsigned int active_mask = utils::activemask();\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id_mod_NxN;\n // The identifier of the column if the iterator is valid.\n int a_col_tmp = -1, a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_tmp = A_cols[a_col_it];\n }\n\n if ( a_col_tmp != -1 && row_colors[a_col_tmp] < current_color )\n {\n a_col_id = a_col_tmp;\n }\n\n // When the diagonal is stored inside the matrix, we have to reject it. We\n // could be using a template parameter but it's not needed since that\n // rejection is really cheap (a couple of extra cycles -- CMP+MOV).\n if ( a_col_id == a_row_id )\n {\n a_col_id = -1;\n }\n\n // We partition valid and invalid column ids. Valid ones come first.\n int vote = utils::ballot( a_col_id != -1, active_mask ) & mask_NxN;\n int ones = __popc( vote );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( a_col_id == -1 )\n {\n dest = ones + lane_id_mod_NxN - dest;\n }\n\n my_s_a_col_ids[dest] = a_col_id;\n my_s_a_col_its[dest] = a_col_it;\n\n // Threads collaborate to load the rows.\n for ( int k = 0 ; k < NxN ; ++k )\n {\n // Exchange column indices.\n const int uniform_a_col_id = my_s_a_col_ids[k];\n\n // Early exit.\n if ( utils::all( uniform_a_col_id == -1, active_mask ) )\n {\n break;\n }\n\n // Load the iterator.\n const int uniform_a_col_it = my_s_a_col_its[k];\n // Load the two matrices.\n Vector_type my_A(0), my_B(0);\n\n if ( uniform_a_col_id != -1 )\n {\n my_A = A_vals[NxN * uniform_a_col_it + lane_id_mod_NxN];\n my_B = Einv [NxN * uniform_a_col_id + lane_id_mod_NxN];\n }\n\n my_s_A_mtx[lane_id_mod_NxN] = my_A;\n my_s_B_mtx[lane_id_mod_NxN] = my_B;\n utils::syncwarp(active_mask);\n // Compute the product of matrices.\n Vector_type tmp(0);\n#pragma unroll\n\n for ( int m = 0 ; m < N ; ++m )\n {\n tmp += my_s_A_mtx[N * lane_id_mod_NxN_div_N + m] * my_s_B_mtx[N * m + lane_id_mod_NxN_mod_N];\n }\n\n my_s_A_mtx[lane_id_mod_NxN] = tmp;\n // We looking for columns in the two rows we're interested in.\n int b_col_it(0), b_col_end(0);\n\n if ( is_active && uniform_a_col_id != -1 )\n {\n b_col_it = A_rows[uniform_a_col_id ];\n b_col_end = A_rows[uniform_a_col_id + 1];\n }\n\n // Init the marker to -1.\n if ( lane_id_mod_NxN == 0 )\n {\n *my_s_A_ji = -1;\n }\n\n // Run the loop.\n b_col_it += lane_id_mod_NxN;\n int shared_found = utils::ballot( lane_id_mod_NxN == 0 && uniform_a_col_id == -1, active_mask );\n\n do\n {\n bool found = b_col_it < b_col_end && A_cols[b_col_it] == a_row_id;\n\n if ( found )\n {\n *my_s_A_ji = b_col_it;\n }\n\n shared_found = shared_found | utils::ballot(found, active_mask);\n b_col_it += NxN;\n }\n while ( __popc( shared_found ) < NUM_ITEMS_PER_WARP && utils::any( b_col_it < b_col_end, active_mask ) );\n\n // Load the blocks.\n const int w_aji = *my_s_A_ji;\n Vector_type my_C(0);\n\n if ( w_aji != -1 )\n {\n my_C = A_vals[NxN * w_aji + lane_id_mod_NxN];\n }\n\n my_s_B_mtx[lane_id_mod_NxN] = my_C;\n // Update e_out.\n#pragma unroll\n\n for ( int m = 0 ; m < N ; ++m )\n {\n e_out -= my_s_A_mtx[N * lane_id_mod_NxN_div_N + m] * my_s_B_mtx[N * m + lane_id_mod_NxN_mod_N];\n }\n }\n } // a_col_begin < a_col_end\n } // current_color != 0\n\n // Store e_out in A\n my_s_A_mtx[lane_id_mod_NxN] = e_out;\n utils::syncwarp();\n\n // Invert the matrices.\n#pragma unroll\n\n for ( int row = 0 ; row < N ; ++row )\n {\n Vector_type diag(0), diag_tmp = my_s_A_mtx[N * row + row];\n\n if ( isNotCloseToZero(diag_tmp) )\n {\n diag = Vector_type(1) / diag_tmp;\n }\n else\n {\n diag = Vector_type(1) / epsilon(diag_tmp);\n }\n\n if ( is_active && lane_id_mod_NxN_div_N == 0 && lane_id_mod_NxN_mod_N != row )\n {\n my_s_A_mtx[N * row + lane_id_mod_NxN_mod_N] *= diag;\n }\n utils::syncwarp();\n\n if ( is_active && lane_id_mod_NxN_div_N != row && lane_id_mod_NxN_mod_N != row )\n {\n my_s_A_mtx[lane_id_mod_NxN] -= my_s_A_mtx[N * lane_id_mod_NxN_div_N + row] * my_s_A_mtx[N * row + lane_id_mod_NxN_mod_N];\n }\n utils::syncwarp();\n\n if ( is_active && lane_id_mod_NxN_div_N == 0 )\n {\n Vector_type tmp = diag;\n\n if ( lane_id_mod_NxN_mod_N != row )\n {\n tmp = -my_s_A_mtx[N * lane_id_mod_NxN_mod_N + row] * diag;\n }\n\n my_s_A_mtx[N * lane_id_mod_NxN_mod_N + row] = tmp;\n }\n utils::syncwarp();\n }\n\n // Store the results to Einv.\n if ( a_row_id != -1 )\n {\n Einv[NxN * a_row_id + lane_id_mod_NxN] = my_s_A_mtx[lane_id_mod_NxN];\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, int NUM_THREADS_PER_ROW, int CTA_SIZE, int WARP_SIZE >\n__global__\n__launch_bounds__( CTA_SIZE, 16 )\nvoid DILU_setup_1x1_kernel( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict A_diag,\n const Matrix_type *__restrict A_vals,\n Matrix_type *__restrict Einv,\n const int *sorted_rows_by_color,\n const int *row_colors,\n const int num_rows_per_color,\n const int current_color )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Number of items per grid.\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_NTPR = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_NTPR = lane_id % NUM_THREADS_PER_ROW;\n // Shared memory to broadcast column IDs.\n __shared__ int s_a_col_ids[CTA_SIZE];\n // Each thread keeps its own pointer.\n int *my_s_a_col_ids = &s_a_col_ids[warp_id * WARP_SIZE];\n // Shared memory to store the matrices.\n __shared__ int s_A_ji[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n int *my_s_A_ji = &s_A_ji[warp_id * WARP_SIZE];\n // Determine which NxN block the threads work with.\n int a_row_it = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_WARPS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load the diagonal.\n Vector_type e_out(0);\n\n // Skip the 1st iteration of the outer-loop (that loop runs on the host).\n if ( current_color != 0 )\n {\n // Ranges of the row.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Iterate over the elements in the columns.\n for ( ; a_col_begin < a_col_end ; a_col_begin += WARP_SIZE )\n {\n // Each thread loads a single element.\n int a_col_it = a_col_begin + lane_id;\n // The identifier of the column if the iterator is valid.\n int a_col_tmp = -1, a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_tmp = A_cols[a_col_it];\n }\n\n if ( a_col_tmp != -1 && row_colors[a_col_tmp] < current_color )\n {\n a_col_id = a_col_tmp;\n }\n\n // When the diagonal is stored inside the matrix, we have to reject it. We\n // could be using a template parameter but it's not needed since that\n // rejection is really cheap (a couple of extra cycles -- CMP+MOV).\n if ( a_col_id == a_row_id )\n {\n a_col_id = -1;\n }\n\n // We partition valid and invalid column ids. Valid ones come first.\n int vote = utils::ballot( a_col_id != -1 );\n int ones = __popc( vote );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( a_col_id == -1 )\n {\n dest = ones + lane_id - dest;\n }\n\n my_s_a_col_ids[dest] = a_col_id;\n // Reset A_jis.\n my_s_A_ji[lane_id] = -1;\n\n __syncwarp();\n\n // Threads collaborate to load the rows.\n for ( int k = 0 ; k < ones ; k += WARP_SIZE / NUM_THREADS_PER_ROW )\n {\n const int local_k = k + lane_id_div_NTPR;\n // Exchange column indices.\n int uniform_a_col_id = -1;\n\n if ( local_k < ones )\n {\n uniform_a_col_id = my_s_a_col_ids[local_k];\n }\n\n // We look for columns in the rows we're interested in.\n int b_col_it(0), b_col_end(0);\n\n if ( uniform_a_col_id != -1 )\n {\n b_col_it = A_rows[uniform_a_col_id ];\n b_col_end = A_rows[uniform_a_col_id + 1];\n }\n\n // Run the loop.\n b_col_it += lane_id_mod_NTPR;\n int shared_found = utils::ballot( lane_id_mod_NTPR == 0 && uniform_a_col_id == -1 );\n\n do\n {\n bool found = b_col_it < b_col_end && A_cols[b_col_it] == a_row_id;\n\n if ( found )\n {\n my_s_A_ji[local_k] = b_col_it;\n }\n\n shared_found = shared_found | utils::ballot(found);\n b_col_it += NUM_THREADS_PER_ROW;\n }\n while ( __popc( shared_found ) < WARP_SIZE / NUM_THREADS_PER_ROW && utils::any( b_col_it < b_col_end ) );\n }\n\n __syncwarp();\n\n // Where to get my A_ji from (if any).\n int a_ji_it = my_s_A_ji[dest];\n // Grab A_jis.\n Matrix_type a_ji(0);\n\n if ( a_ji_it != -1 )\n {\n a_ji = A_vals[a_ji_it];\n }\n\n // Update e_out.\n if ( a_col_id != -1 )\n {\n e_out += a_ji * Einv[a_col_id] * A_vals[a_col_it];\n }\n } // a_col_begin < a_col_end\n } // current_color != 0\n\n // Reduce the e_outs in one value.\n#pragma unroll\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n e_out += utils::shfl_xor( e_out, mask );\n }\n\n // Store the result.\n if ( lane_id == 0 )\n {\n Matrix_type res = A_vals[A_diag[a_row_id]] - e_out;\n\n if ( res != Matrix_type(0) )\n {\n res = Matrix_type(1) / res;\n }\n\n Einv[a_row_id] = static_cast(res);\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int N, bool ROW_MAJOR, int WARP_SIZE, typename Value_type >\nstatic __device__ __forceinline__\nValue_type reduce_distributed_vectors( Value_type x, int is_leader )\n{\n if ( N & (N - 1) )\n {\n#pragma unroll\n\n for ( int i = 1 ; i < N ; ++i )\n {\n Value_type other_x = utils::shfl_down( x, ROW_MAJOR ? i : N * i );\n\n if ( is_leader )\n {\n x += other_x;\n }\n }\n }\n else\n {\n#pragma unroll\n\n for ( int i = 1 ; i < N ; i <<= 1 )\n {\n x += utils::shfl_xor( x, ROW_MAJOR ? i : N * i );\n }\n }\n\n return x;\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, int N, int CTA_SIZE, int WARP_SIZE, bool ROW_MAJOR, bool HAS_EXTERNAL_DIAG >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_forward_NxN_kernel( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int NUM_ITEMS_PER_WARP = WARP_SIZE / NxN;\n // Number of items computer per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_ITEMS_PER_WARP * NUM_WARPS_PER_CTA;\n // Number of items per grid.\n const int NUM_ITEMS_PER_GRID = gridDim.x * NUM_ITEMS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_NxN = lane_id / NxN;\n const int lane_id_mod_NxN = lane_id % NxN;\n // Useful index to compute matrix products.\n const int lane_id_mod_NxN_div_N = lane_id_mod_NxN / N;\n const int lane_id_mod_NxN_mod_N = lane_id_mod_NxN % N;\n // We to get my data from when I use SHFL.\n const int shfl_offset = lane_id - lane_id_mod_NxN;\n // Shared memory needed to exchange X and delta.\n __shared__ volatile Vector_type s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_mem = &s_mem[threadIdx.x - lane_id_mod_NxN];\n // Is the thread active? For example, for 5x5 only the first 25 threads are active per warp.\n // At compile time, the compiler will see is_active == true for 2x2 (since NxN & (NxN-1) evals\n // to false ; that's the common trick to determine if a number is a power of 2).\n int is_active = true;\n\n if ( NxN & (NxN - 1) )\n {\n is_active = lane_id_div_NxN < NUM_ITEMS_PER_WARP;\n }\n\n // Determine which NxN block the threads work with.\n int a_row_it = num_rows_per_color;\n\n if ( is_active )\n {\n a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + warp_id * NUM_ITEMS_PER_WARP + lane_id_div_NxN;\n }\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_ITEMS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_bmAx(0);\n\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n my_bmAx = __cachingLoad(&b[N * a_row_id + lane_id_mod_NxN_div_N]);\n }\n }\n else\n {\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n my_bmAx = b[N * a_row_id + lane_id_mod_NxN_mod_N];\n }\n }\n\n // Don't do anything if X is zero.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // If the diagonal is stored separately, we have a special treatment.\n int a_col_max = a_col_end;\n\n if ( HAS_EXTERNAL_DIAG )\n {\n ++a_col_max;\n }\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( a_col_begin < a_col_max ) ; a_col_begin += NxN )\n {\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id_mod_NxN;\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n if ( HAS_EXTERNAL_DIAG && a_col_it == a_col_end )\n {\n a_col_id = a_row_id;\n }\n\n // Determine if the color is valid.\n int a_col_is_valid = false;\n#ifdef AMGX_ILU_COLORING\n if ( a_col_id != -1 && current_color != 0 )\n {\n if ( boundary_coloring == FIRST )\n {\n a_col_is_valid = a_col_id >= boundary_index;\n }\n else\n {\n a_col_is_valid = a_col_id < boundary_index && row_colors[a_col_id] < current_color;\n }\n }\n\n#else\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n a_col_is_valid = row_colors[a_col_id] < current_color;\n }\n\n#endif\n // Count the number of active columns.\n // int vote = utils::ballot(aColId != -1);\n // The number of iterations.\n // int nCols = max( __popc( vote & 0x0000ffff ), __popc( vote & 0xffff0000 ) );\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < NxN ; k += N )\n {\n int my_k = k + lane_id_mod_NxN_div_N;\n // Load N blocks of X.\n int uniform_a_col_id = utils::shfl( a_col_id, shfl_offset + my_k );\n int uniform_a_col_is_valid = utils::shfl( a_col_is_valid, shfl_offset + my_k );\n Vector_type my_x(0);\n\n if ( uniform_a_col_id != -1 )\n {\n my_x = __cachingLoad(&x[N * uniform_a_col_id + lane_id_mod_NxN_mod_N]);\n }\n\n if ( uniform_a_col_id != -1 && uniform_a_col_is_valid )\n {\n my_x += delta[N * uniform_a_col_id + lane_id_mod_NxN_mod_N];\n }\n\n my_s_mem[lane_id_mod_NxN] = my_x;\n // Load N blocks of A.\n#pragma unroll\n\n for ( int i = 0 ; i < N ; ++i )\n {\n int uniform_a_col_tmp = a_col_begin + k + i, uniform_a_col_it = -1;\n\n if ( uniform_a_col_tmp < a_col_end )\n {\n uniform_a_col_it = uniform_a_col_tmp;\n }\n\n if ( HAS_EXTERNAL_DIAG && is_active && uniform_a_col_tmp == a_col_end )\n {\n uniform_a_col_it = A_diag[a_row_id];\n }\n\n Matrix_type my_val(0);\n\n if ( uniform_a_col_it != -1 )\n {\n my_val = A_vals[NxN * uniform_a_col_it + lane_id_mod_NxN];\n }\n\n if ( ROW_MAJOR )\n {\n my_bmAx -= my_val * my_s_mem[N * i + lane_id_mod_NxN_mod_N];\n }\n else\n {\n my_bmAx -= my_val * my_s_mem[N * i + lane_id_mod_NxN_div_N];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Load Einvs.\n Vector_type my_Einv = Einv[NxN * a_row_id + lane_id_mod_NxN];\n // Reduce bmAx terms.\n int is_leader = lane_id_mod_NxN_div_N == 0;\n\n if ( ROW_MAJOR )\n {\n is_leader = lane_id_mod_NxN_mod_N == 0;\n }\n\n my_bmAx = reduce_distributed_vectors( my_bmAx, is_leader );\n\n // Update the shared terms.\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n my_s_mem[lane_id_mod_NxN_div_N] = my_bmAx;\n }\n }\n else\n {\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n my_s_mem[lane_id_mod_NxN_mod_N] = my_bmAx;\n }\n }\n\n // Update the diagonal term.\n if ( ROW_MAJOR )\n {\n my_bmAx = my_Einv * my_s_mem[lane_id_mod_NxN_mod_N];\n }\n else\n {\n my_bmAx = my_Einv * my_s_mem[lane_id_mod_NxN_div_N];\n }\n\n // Reduce bmAx terms.\n my_bmAx = reduce_distributed_vectors( my_bmAx, is_leader );\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n delta[N * a_row_id + lane_id_mod_NxN_div_N] = my_bmAx;\n }\n }\n else\n {\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n delta[N * a_row_id + lane_id_mod_NxN_mod_N] = my_bmAx;\n }\n }\n }\n}\n\ntemplate< typename Matrix_type, typename Vector_type, int N, int CTA_SIZE, int WARP_SIZE, bool HAS_EXTERNAL_DIAG, int NUM_WARP_ITERS_PER_BLOCK >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_forward_NxN_kernel_large( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of rows computed per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_WARPS_PER_CTA;\n // Number of rows? per grid.\n const int NUM_ITEMS_PER_GRID = CTA_SIZE;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n // Useful index to compute matrix products.\n const int lane_id_div_N = lane_id / N;\n const int lane_id_mod_N = lane_id % N; // id of a lane inside the block\n const int blocks_per_warp = WARP_SIZE / N; // we process this cols per warp per row\n const int row_elems_per_warp = blocks_per_warp * N;\n // Shared to store bmAx\n __shared__ volatile Vector_type bmAx[CTA_SIZE * NUM_WARP_ITERS_PER_BLOCK];\n volatile Vector_type *my_bmAx_s = &bmAx[warp_id * NUM_WARP_ITERS_PER_BLOCK * WARP_SIZE];\n // Determine which NxN block the threads work with.\n int a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_ITEMS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_bmAx(0);\n\n if ( lane_id < N )\n {\n my_bmAx = __cachingLoad(&b[N * a_row_id + lane_id]);\n }\n\n#pragma unroll\n\n for (int i = 0; i < NUM_WARP_ITERS_PER_BLOCK; i++)\n {\n my_bmAx_s[WARP_SIZE * i + lane_id] = 0.0;\n }\n\n // Don't do anything if X is zero.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // If the diagonal is stored separately, we have a special treatment.\n int a_col_max = a_col_end;\n\n if ( HAS_EXTERNAL_DIAG )\n {\n ++a_col_max;\n }\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( a_col_begin < a_col_max ) ; a_col_begin += WARP_SIZE ) // NxN\n {\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id;\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n if ( HAS_EXTERNAL_DIAG && a_col_it == a_col_end )\n {\n a_col_id = a_row_id;\n }\n\n // Determine if the color is valid.\n int a_col_is_valid = false;\n#ifdef AMGX_ILU_COLORING\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n if ( boundary_coloring == FIRST )\n {\n a_col_is_valid = a_col_id >= boundary_index;\n }\n else\n {\n a_col_is_valid = a_col_id < boundary_index && row_colors[a_col_id] < current_color;\n }\n }\n\n#else\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n a_col_is_valid = row_colors[a_col_id] < current_color;\n }\n\n#endif\n\n // Loop over columns. We compute blocks_per_warp columns per iteration.\n for ( int k = 0 ; k < WARP_SIZE ; k += blocks_per_warp )\n {\n // id of the processed block by this thread\n int my_k = k + lane_id_div_N;\n // Load N blocks of X (if valid)\n int uniform_a_col_id = utils::shfl( a_col_id, my_k );\n int uniform_a_col_is_valid = utils::shfl( a_col_is_valid, my_k );\n Vector_type my_x(0);\n\n if ( uniform_a_col_id != -1 && lane_id < row_elems_per_warp)\n {\n my_x = __cachingLoad(&x[N * uniform_a_col_id + lane_id_mod_N]);\n }\n\n if ( uniform_a_col_id != -1 && uniform_a_col_is_valid && lane_id < row_elems_per_warp)\n {\n my_x += delta[N * uniform_a_col_id + lane_id_mod_N];\n }\n\n //my_s_mem[lane_id] = my_x;\n#pragma unroll\n\n for ( int i = 0 ; i < blocks_per_warp ; ++i )\n {\n // k-th batch of blocks, i-th block. each thread process a column/row of a_it = uniform_a_col_tmp\n int uniform_a_col_tmp = a_col_begin + k + i, uniform_a_col_it = -1;\n\n // check if we are going out of bounds/color\n if ( uniform_a_col_tmp < a_col_end )\n {\n uniform_a_col_it = uniform_a_col_tmp;\n }\n\n if ( HAS_EXTERNAL_DIAG && uniform_a_col_tmp == a_col_end )\n {\n uniform_a_col_it = A_diag[a_row_id];\n }\n\n // swipe with the whole warp\n if (uniform_a_col_it != -1)\n {\n int block_inside_id = lane_id;\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n Matrix_type my_val(0);\n\n if ( uniform_a_col_it != -1 && block_inside_id < NxN)\n {\n my_val = A_vals[NxN * uniform_a_col_it + block_inside_id];\n }\n\n my_bmAx_s[block_inside_id] -= my_val * utils::shfl(my_x, N * i + block_inside_id % N); // MOD IS SLOW!\n block_inside_id += WARP_SIZE;\n }\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Load Einvs.\n Vector_type my_Einv[NUM_WARP_ITERS_PER_BLOCK];\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n my_Einv[j] = 0.0;\n }\n\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n if ((WARP_SIZE * j + lane_id) < NxN)\n {\n my_Einv[j] = Einv[NxN * a_row_id + WARP_SIZE * j + lane_id];\n }\n }\n\n // Reduce bmAx terms.\n {\n#pragma unroll\n\n for ( int i = 0 ; i < N ; ++i )\n {\n if ( lane_id < N )\n {\n my_bmAx += my_bmAx_s[N * lane_id + i];\n }\n }\n }\n // Update the diagonal term.\n int block_inside_id = lane_id;\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n my_bmAx_s[block_inside_id] = my_Einv[j] * utils::shfl(my_bmAx, block_inside_id % N);\n block_inside_id += WARP_SIZE;\n }\n\n // Reduce bmAx terms.\n {\n my_bmAx = 0.0;\n#pragma unroll\n\n for ( int i = 0 ; i < N ; ++i )\n {\n int idx = N * lane_id + i;\n\n if ( lane_id < N )\n {\n my_bmAx += my_bmAx_s[idx];\n }\n }\n }\n\n // Store the results.\n if ( lane_id < N )\n {\n delta[N * a_row_id + lane_id] = my_bmAx;\n }\n }\n}\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, int CTA_SIZE, int WARP_SIZE, bool ROW_MAJOR, bool HAS_EXTERNAL_DIAG >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_forward_4x4_kernel( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Number of items per warp.\n const int NUM_ITEMS_PER_WARP = WARP_SIZE / 16;\n // Number of items computer per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_ITEMS_PER_WARP * NUM_WARPS_PER_CTA;\n // Number of items per grid.\n const int NUM_ITEMS_PER_GRID = gridDim.x * NUM_ITEMS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_mod_16 = lane_id % 16;\n // Useful index to compute matrix products.\n const int lane_id_mod_16_div_4 = lane_id_mod_16 / 4;\n const int lane_id_mod_16_mod_4 = lane_id_mod_16 % 4;\n // We to get my data from when I use SHFL.\n const int shfl_offset = lane_id - lane_id_mod_16;\n // Shared memory needed to exchange X and delta.\n __shared__ volatile Vector_type s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_mem = &s_mem[threadIdx.x - lane_id_mod_16];\n // Determine which 16 block the threads work with.\n int a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + threadIdx.x / 16;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_ITEMS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_bmAx(0);\n\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_16_mod_4 == 0 )\n {\n my_bmAx = __cachingLoad(&b[4 * a_row_id + lane_id_mod_16_div_4]);\n }\n }\n else\n {\n if ( lane_id_mod_16_div_4 == 0 )\n {\n my_bmAx = b[4 * a_row_id + lane_id_mod_16_mod_4];\n }\n }\n\n // Don't do anything if X is zero.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // If the diagonal is stored separately, we have a special treatment.\n int a_col_max = a_col_end;\n\n if ( HAS_EXTERNAL_DIAG )\n {\n ++a_col_max;\n }\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( a_col_begin < a_col_max ) ; a_col_begin += 16 )\n {\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id_mod_16;\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n if ( HAS_EXTERNAL_DIAG && a_col_it == a_col_end )\n {\n a_col_id = a_row_id;\n }\n\n // Determine if the color is valid.\n int a_col_is_valid = false;\n#ifdef AMGX_ILU_COLORING\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n if ( boundary_coloring == FIRST )\n {\n a_col_is_valid = a_col_id >= boundary_index;\n }\n else\n {\n a_col_is_valid = a_col_id < boundary_index && row_colors[a_col_id] < current_color;\n }\n }\n\n#else\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n a_col_is_valid = row_colors[a_col_id] < current_color;\n }\n\n#endif\n // Count the number of active columns.\n // int vote = utils::ballot(aColId != -1);\n // The number of iterations.\n // int nCols = max( __popc( vote & 0x0000ffff ), __popc( vote & 0xffff0000 ) );\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < 16 ; k += 4 )\n {\n int my_k = k + lane_id_mod_16_div_4;\n // Load N blocks of X.\n int uniform_a_col_id = utils::shfl( a_col_id, shfl_offset + my_k );\n int uniform_a_col_is_valid = utils::shfl( a_col_is_valid, shfl_offset + my_k );\n Vector_type my_x(0);\n\n if ( uniform_a_col_id != -1 )\n {\n my_x = __cachingLoad(&x[4 * uniform_a_col_id + lane_id_mod_16_mod_4]);\n }\n\n if ( uniform_a_col_id != -1 && uniform_a_col_is_valid )\n {\n my_x += delta[4 * uniform_a_col_id + lane_id_mod_16_mod_4];\n }\n\n my_s_mem[lane_id_mod_16] = my_x;\n // Load N blocks of A.\n#pragma unroll\n\n for ( int i = 0 ; i < 4 ; ++i )\n {\n int uniform_a_col_tmp = a_col_begin + k + i, uniform_a_col_it = -1;\n\n if ( uniform_a_col_tmp < a_col_end )\n {\n uniform_a_col_it = uniform_a_col_tmp;\n }\n\n if ( HAS_EXTERNAL_DIAG && uniform_a_col_tmp == a_col_end )\n {\n uniform_a_col_it = A_diag[a_row_id];\n }\n\n Matrix_type my_val(0);\n\n if ( uniform_a_col_it != -1 )\n {\n my_val = A_vals[16 * uniform_a_col_it + lane_id_mod_16];\n }\n\n if ( ROW_MAJOR )\n {\n my_bmAx -= my_val * my_s_mem[4 * i + lane_id_mod_16_mod_4];\n }\n else\n {\n my_bmAx -= my_val * my_s_mem[4 * i + lane_id_mod_16_div_4];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Load Einvs.\n Vector_type my_Einv = Einv[16 * a_row_id + lane_id_mod_16];\n // Reduce bmAx terms.\n int is_leader = lane_id_mod_16_div_4 == 0;\n\n if ( ROW_MAJOR )\n {\n is_leader = lane_id_mod_16_mod_4 == 0;\n }\n\n my_bmAx = reduce_distributed_vectors<4, ROW_MAJOR, WARP_SIZE>( my_bmAx, is_leader );\n\n // Update the shared terms.\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_16_mod_4 == 0 )\n {\n my_s_mem[lane_id_mod_16_div_4] = my_bmAx;\n }\n }\n else\n {\n if ( lane_id_mod_16_div_4 == 0 )\n {\n my_s_mem[lane_id_mod_16_mod_4] = my_bmAx;\n }\n }\n\n // Update the diagonal term.\n if ( ROW_MAJOR )\n {\n my_bmAx = my_Einv * my_s_mem[lane_id_mod_16_mod_4];\n }\n else\n {\n my_bmAx = my_Einv * my_s_mem[lane_id_mod_16_div_4];\n }\n\n // Reduce bmAx terms.\n my_bmAx = reduce_distributed_vectors<4, ROW_MAJOR, WARP_SIZE>( my_bmAx, is_leader );\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_16_mod_4 == 0 )\n {\n delta[4 * a_row_id + lane_id_mod_16_div_4] = my_bmAx;\n }\n }\n else\n {\n if ( lane_id_mod_16_div_4 == 0 )\n {\n delta[4 * a_row_id + lane_id_mod_16_mod_4] = my_bmAx;\n }\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, int CTA_SIZE, bool HAS_EXTERNAL_DIAG >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_forward_4x4_kernel_row_major_vec4( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *Einv,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n // Number of half warps per CTA.\n const int NUM_HALF_WARPS = CTA_SIZE / 16;\n // Coordinates of the thread.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Coordinates of the thread in the CTA.\n const int thread_id_div_16 = threadIdx.x / 16;\n const int thread_id_mod_16 = threadIdx.x % 16;\n // Useful constants.\n const int thread_id_mod_16_div_4 = thread_id_mod_16 / 4;\n const int thread_id_mod_16_mod_4 = thread_id_mod_16 % 4;\n const int shfl_offset = 16 * (lane_id / 16);\n // Shared memory needed to exchange X and delta.\n __shared__ volatile Vector_type s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_mem = &s_mem[16 * thread_id_div_16];\n // The iterator over rows.\n int a_row_it = blockIdx.x * NUM_HALF_WARPS + thread_id_div_16;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += gridDim.x * NUM_HALF_WARPS )\n {\n unsigned int active_mask = utils::activemask();\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_bmAx(0);\n\n if ( thread_id_mod_16_div_4 == 0 )\n {\n my_bmAx = __cachingLoad(&b[4 * a_row_id + thread_id_mod_16_mod_4]);\n }\n\n // The range of the row.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // If it has an external diagonal, we need one more item to put the diag.\n int a_col_max = a_col_end;\n\n if ( HAS_EXTERNAL_DIAG )\n {\n ++a_col_max;\n }\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; a_col_begin < a_col_max ; a_col_begin += 16 )\n {\n unsigned int active_mask_inner = utils::activemask();\n int a_col_it = a_col_begin + thread_id_mod_16;\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = __cachingLoad(&A_cols[a_col_it]);\n }\n\n if ( HAS_EXTERNAL_DIAG && a_col_it == a_col_end )\n {\n a_col_id = a_row_id;\n }\n\n // Determine if the color is valid.\n int a_col_is_valid = false;\n#ifdef AMGX_ILU_COLORING\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n if ( boundary_coloring == FIRST )\n {\n a_col_is_valid = a_col_id >= boundary_index;\n }\n else\n {\n a_col_is_valid = a_col_id < boundary_index && __cachingLoad(&row_colors[a_col_id]) < current_color;\n }\n }\n\n#else\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n a_col_is_valid = row_colors[a_col_id] < current_color;\n }\n\n#endif\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < 16 ; k += 4 )\n {\n int my_k = k + thread_id_mod_16_div_4;\n // Load 8 blocks of X.\n int uniform_a_col_id = utils::shfl( a_col_id, shfl_offset + my_k, warpSize, active_mask_inner );\n int uniform_a_col_is_valid = utils::shfl( a_col_is_valid, shfl_offset + my_k, warpSize, active_mask_inner );\n Vector_type my_x(0);\n\n if ( uniform_a_col_id != -1 )\n {\n my_x = __cachingLoad(&x[4 * uniform_a_col_id + thread_id_mod_16_mod_4]);\n }\n\n if ( uniform_a_col_id != -1 && uniform_a_col_is_valid )\n {\n my_x += delta[4 * uniform_a_col_id + thread_id_mod_16_mod_4];\n }\n\n my_s_mem[thread_id_mod_16] = my_x;\n int uniform_a_col_tmp = a_col_begin + my_k, uniform_a_col_it = -1;\n\n if ( uniform_a_col_tmp < a_col_end )\n {\n uniform_a_col_it = uniform_a_col_tmp;\n }\n\n if ( HAS_EXTERNAL_DIAG && uniform_a_col_tmp == a_col_end )\n {\n uniform_a_col_it = A_diag[a_row_id];\n }\n\n Matrix_type my_vals[4] = { Matrix_type(0) };\n\n if ( uniform_a_col_it != -1 )\n {\n utils::load_vec4( my_vals, &A_vals[16 * uniform_a_col_it + 4 * thread_id_mod_16_mod_4] );\n }\n\n my_bmAx -= my_vals[0] * my_s_mem[4 * thread_id_mod_16_div_4 + 0];\n my_bmAx -= my_vals[1] * my_s_mem[4 * thread_id_mod_16_div_4 + 1];\n my_bmAx -= my_vals[2] * my_s_mem[4 * thread_id_mod_16_div_4 + 2];\n my_bmAx -= my_vals[3] * my_s_mem[4 * thread_id_mod_16_div_4 + 3];\n }\n }\n\n // Load Einvs.\n Matrix_type my_Einv = Einv[16 * a_row_id + thread_id_mod_16];\n // Reduce bmAx terms.\n my_bmAx += utils::shfl_xor( my_bmAx, 4, warpSize, active_mask );\n my_bmAx += utils::shfl_xor( my_bmAx, 8, warpSize, active_mask );\n\n // Update the shared terms.\n if ( thread_id_mod_16_div_4 == 0 )\n {\n my_s_mem[thread_id_mod_16_mod_4] = my_bmAx;\n }\n\n // Update the diagonal term.\n my_bmAx = my_Einv * my_s_mem[thread_id_mod_16_mod_4];\n // Reduce bmAx terms.\n my_bmAx += utils::shfl_xor( my_bmAx, 1, warpSize, active_mask );\n my_bmAx += utils::shfl_xor( my_bmAx, 2, warpSize, active_mask );\n\n // Store the results.\n if ( thread_id_mod_16_mod_4 == 0 )\n {\n delta[4 * a_row_id + thread_id_mod_16_div_4] = my_bmAx;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, int NUM_THREADS_PER_ROW, int CTA_SIZE, int WARP_SIZE, bool HAS_EXTERNAL_DIAG >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_forward_1x1_kernel( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n // Number of items per CTA.\n const int NUM_ROWS_PER_CTA = CTA_SIZE / NUM_THREADS_PER_ROW;\n // Number of items per grid.\n const int NUM_ROWS_PER_GRID = gridDim.x * NUM_ROWS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_mod_NTPR = lane_id % NUM_THREADS_PER_ROW;\n // Determine which NxN block the threads work with.\n int a_row_it = blockIdx.x * NUM_ROWS_PER_CTA + (threadIdx.x / NUM_THREADS_PER_ROW);\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_ROWS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_bmAx = amgx::types::util::get_zero();\n\n if ( lane_id_mod_NTPR == 0 )\n {\n my_bmAx = __cachingLoad(&b[a_row_id]);\n }\n\n // If it has an external diag.\n if ( HAS_EXTERNAL_DIAG && lane_id_mod_NTPR == 0 )\n {\n my_bmAx -= A_vals[A_diag[a_row_id]] * x[a_row_id];\n }\n\n // Don't do anything if X is zero.\n int a_col_it = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n\n // If the diagonal is stored separately, we have a special treatment.\n //if( HAS_EXTERNAL_DIAG )\n // ++a_col_end;\n\n // Each warp load column indices of 32 nonzero blocks\n for ( a_col_it += lane_id_mod_NTPR ; utils::any( a_col_it < a_col_end ) ; a_col_it += NUM_THREADS_PER_ROW )\n {\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Ignore the diagonal element since its color is smaller, and been accounted for above\n if (HAS_EXTERNAL_DIAG && a_col_id == a_row_id)\n {\n a_col_id = -1;\n }\n\n // Load x.\n Vector_type my_x(0);\n\n if ( a_col_id != -1 )\n {\n my_x = __cachingLoad(&x[a_col_id]);\n }\n\n // Is it really a valid column (due to coloring).\n int valid = false;\n#ifdef AMGX_ILU_COLORING\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n if ( boundary_coloring == FIRST )\n {\n valid = a_col_id >= boundary_index;\n }\n else\n {\n valid = a_col_id < boundary_index && row_colors[a_col_id] < current_color;\n }\n }\n\n#else\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n valid = row_colors[a_col_tmp] < current_color;\n }\n\n#endif\n\n // Load my x value.\n if ( valid )\n {\n my_x += delta[a_col_id];\n }\n\n // Load my item from A.\n Matrix_type my_val(0);\n\n if ( a_col_it < a_col_end )\n {\n my_val = A_vals[a_col_it];\n }\n\n // Update bmAx.\n my_bmAx -= my_val * my_x;\n }\n\n // Reduce bmAx terms.\n#pragma unroll\n\n for ( int mask = NUM_THREADS_PER_ROW / 2 ; mask > 0 ; mask >>= 1 )\n {\n my_bmAx += utils::shfl_xor( my_bmAx, mask );\n }\n\n // Store the results.\n if ( lane_id_mod_NTPR == 0 )\n {\n delta[a_row_id] = Einv[a_row_id] * my_bmAx;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType, int N, int CTA_SIZE, int WARP_SIZE, bool ROW_MAJOR >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_backward_NxN_kernel( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int NUM_ITEMS_PER_WARP = WARP_SIZE / NxN;\n // Number of items computer per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_ITEMS_PER_WARP * NUM_WARPS_PER_CTA;\n // Number of items per grid.\n const int NUM_ITEMS_PER_GRID = gridDim.x * NUM_ITEMS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_NxN = lane_id / NxN;\n const int lane_id_mod_NxN = lane_id % NxN;\n // Useful index to compute matrix products.\n const int lane_id_mod_NxN_div_N = lane_id_mod_NxN / N;\n const int lane_id_mod_NxN_mod_N = lane_id_mod_NxN % N;\n // We to get my data from when I use SHFL.\n const int shfl_offset = lane_id - lane_id_mod_NxN;\n // Shared memory needed to exchange X and delta.\n __shared__ volatile Vector_type s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_mem = &s_mem[threadIdx.x - lane_id_mod_NxN];\n // Is the thread active? For example, for 5x5 only the first 25 threads are active per warp.\n // At compile time, the compiler will see is_active == true for 2x2 (since NxN & (NxN-1) evals\n // to false ; that's the common trick to determine if a number is a power of 2).\n int is_active = true;\n\n if ( NxN & (NxN - 1) )\n {\n is_active = lane_id_div_NxN < NUM_ITEMS_PER_WARP;\n }\n\n // Determine which NxN block the threads work with.\n int a_row_it = num_rows_per_color;\n\n if ( is_active )\n {\n a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + warp_id * NUM_ITEMS_PER_WARP + lane_id_div_NxN;\n }\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_ITEMS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_delta(0);\n // Don't do anything if X is zero.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( a_col_begin < a_col_end ) ; a_col_begin += NxN )\n {\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id_mod_NxN;\n // Get the ID of the column.\n int a_col_tmp = -1, a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_tmp = A_cols[a_col_it];\n }\n\n // Make sure the column is interesting.\n#ifdef AMGX_ILU_COLORING\n int valid = false;\n\n if ( a_col_tmp != -1 && current_color != 0 )\n {\n if ( boundary_coloring == LAST )\n {\n valid = a_col_tmp >= boundary_index;\n }\n else\n {\n valid = a_col_tmp < boundary_index && row_colors[a_col_tmp] > current_color;\n }\n }\n\n#else\n int valid = false;\n\n if ( a_col_tmp != -1 && row_colors[a_col_tmp] > current_color )\n {\n valid = true;\n }\n\n#endif\n\n // Set the column id.\n if ( valid )\n {\n a_col_id = a_col_tmp;\n }\n\n // Count the number of active columns.\n // int vote = utils::ballot(aColId != -1);\n // The number of iterations.\n // int nCols = max( __popc( vote & 0x0000ffff ), __popc( vote & 0xffff0000 ) );\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < NxN ; k += N )\n {\n int my_k = k + lane_id_mod_NxN_div_N;\n // Load N blocks of X.\n int uniform_a_col_id = utils::shfl( a_col_id, shfl_offset + my_k );\n Vector_type my_x(0);\n\n if ( uniform_a_col_id != -1 )\n {\n my_x = Delta[N * uniform_a_col_id + lane_id_mod_NxN_mod_N];\n }\n\n my_s_mem[lane_id_mod_NxN] = my_x;\n // Load N blocks of A.\n#pragma unroll\n\n for ( int i = 0 ; i < N ; ++i )\n {\n //if( uniform_a_col_id == -1 )\n // break;\n int uniform_a_col_tmp = a_col_begin + k + i, uniform_a_col_it = -1;\n\n if ( uniform_a_col_tmp < a_col_end )\n {\n uniform_a_col_it = uniform_a_col_tmp;\n }\n\n Matrix_type my_val(0);\n\n if ( uniform_a_col_it != -1 )\n {\n my_val = A_vals[NxN * uniform_a_col_it + lane_id_mod_NxN];\n }\n\n if ( ROW_MAJOR )\n {\n my_delta += my_val * my_s_mem[N * i + lane_id_mod_NxN_mod_N];\n }\n else\n {\n my_delta += my_val * my_s_mem[N * i + lane_id_mod_NxN_div_N];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Load Einvs.\n Matrix_type my_Einv = Einv[NxN * a_row_id + lane_id_mod_NxN];\n // Reduce bmAx terms.\n int is_leader = lane_id_mod_NxN_div_N == 0;\n\n if ( ROW_MAJOR )\n {\n is_leader = lane_id_mod_NxN_mod_N == 0;\n }\n\n my_delta = reduce_distributed_vectors( my_delta, is_leader );\n\n // Update the shared terms.\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n my_s_mem[lane_id_mod_NxN_div_N] = my_delta;\n }\n }\n else\n {\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n my_s_mem[lane_id_mod_NxN_mod_N] = my_delta;\n }\n }\n\n // Update the diagonal term.\n if ( ROW_MAJOR )\n {\n my_delta = my_Einv * my_s_mem[lane_id_mod_NxN_mod_N];\n }\n else\n {\n my_delta = my_Einv * my_s_mem[lane_id_mod_NxN_div_N];\n }\n\n // Reduce bmAx terms.\n my_delta = reduce_distributed_vectors( my_delta, is_leader );\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n const int offset = N * a_row_id + lane_id_mod_NxN_div_N;\n Vector_type my_b(0), my_x(0);\n\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n my_b = __cachingLoad(&delta[offset]);\n my_x = x [offset];\n }\n\n my_delta = my_b - my_delta;\n\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n x [offset] = my_x + weight * my_delta;\n Delta[offset] = my_delta;\n }\n }\n else\n {\n const int offset = N * a_row_id + lane_id_mod_NxN_mod_N;\n Vector_type my_b(0), my_x(0);\n\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n my_b = __cachingLoad(&delta[offset]);\n my_x = x [offset];\n }\n\n my_delta = my_b - my_delta;\n\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n x [offset] = my_x + weight * my_delta;\n Delta[offset] = my_delta;\n }\n }\n }\n}\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType, int N, int CTA_SIZE, int WARP_SIZE, bool ROW_MAJOR, int NUM_WARP_ITERS_PER_BLOCK >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_backward_NxN_kernel_large( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items computer per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_WARPS_PER_CTA;\n // Number of items per grid.\n const int NUM_ITEMS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_N = lane_id / N;\n const int lane_id_mod_N = lane_id % N; // id of a lane inside the block\n const int blocks_per_warp = WARP_SIZE / N; // we process this cols per warp per row\n const int row_elems_per_warp = blocks_per_warp * N;\n // Shared to store t_delta\n __shared__ volatile Vector_type delta_s[CTA_SIZE * NUM_WARP_ITERS_PER_BLOCK];\n volatile Vector_type *my_delta_s = &delta_s[warp_id * NUM_WARP_ITERS_PER_BLOCK * WARP_SIZE];\n // Determine which NxN block the threads work with.\n int a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_ITEMS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Accumulator\n Vector_type my_delta(0);\n //Vector_type mAx[NUM_WARP_ITERS_PER_BLOCK];\n#pragma unroll\n\n for (int i = 0; i < NUM_WARP_ITERS_PER_BLOCK; i++)\n {\n my_delta_s[WARP_SIZE * i + lane_id] = 0.0;\n }\n\n // Don't do anything if X is zero.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( a_col_begin < a_col_end ) ; a_col_begin += WARP_SIZE )\n {\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id;\n // Get the ID of the column.\n int a_col_tmp = -1, a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_tmp = A_cols[a_col_it];\n }\n\n // Make sure the column is interesting.\n#ifdef AMGX_ILU_COLORING\n int valid = false;\n\n if ( a_col_tmp != -1 && current_color != 0 )\n {\n if ( boundary_coloring == LAST )\n {\n valid = a_col_tmp >= boundary_index;\n }\n else\n {\n valid = a_col_tmp < boundary_index && row_colors[a_col_tmp] > current_color;\n }\n }\n\n#else\n int valid = false;\n\n if ( a_col_tmp != -1 && row_colors[a_col_tmp] > current_color )\n {\n valid = true;\n }\n\n#endif\n\n // Set the column id.\n if ( valid )\n {\n a_col_id = a_col_tmp;\n }\n\n // Loop over columns. We compute blocks_per_warp columns per iteration.\n for ( int k = 0 ; k < WARP_SIZE ; k += blocks_per_warp )\n {\n // id of the processed block by this thread\n int my_k = k + lane_id_div_N;\n // Load N blocks of X (if valid)\n int uniform_a_col_id = utils::shfl( a_col_id, my_k );\n Vector_type my_x(0);\n\n if ( uniform_a_col_id != -1 && lane_id < row_elems_per_warp)\n {\n my_x = Delta[N * uniform_a_col_id + lane_id_mod_N];\n }\n\n // Load blocks of A.\n // for each block in a batch\n#pragma unroll\n\n for ( int i = 0 ; i < blocks_per_warp ; ++i )\n {\n // k-th batch of blocks, i-th block. each thread process a column/row of a_it = uniform_a_col_tmp\n int uniform_a_col_tmp = a_col_begin + k + i, uniform_a_col_it = -1;\n\n // check if we are going out of bounds/color\n if ( uniform_a_col_tmp < a_col_end )\n {\n uniform_a_col_it = uniform_a_col_tmp;\n }\n\n // swipe with the whole warp\n if (uniform_a_col_it != -1)\n {\n int block_inside_id = lane_id;\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n Matrix_type my_val(0);\n\n if ( uniform_a_col_it != -1 && block_inside_id < NxN)\n {\n my_val = A_vals[NxN * uniform_a_col_it + block_inside_id];\n }\n\n my_delta_s[block_inside_id] -= my_val * utils::shfl(my_x, N * i + block_inside_id % N); //my_s_mem[N*i + block_inside_id % N]; // MOD IS SLOW!\n block_inside_id += WARP_SIZE;\n }\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Load Einvs.\n Vector_type my_Einv[NUM_WARP_ITERS_PER_BLOCK];\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n my_Einv[j] = 0.0;\n }\n\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n if ((WARP_SIZE * j + lane_id) < NxN)\n {\n my_Einv[j] = Einv[NxN * a_row_id + WARP_SIZE * j + lane_id];\n }\n }\n\n // Reduce bmAx terms.\n {\n#pragma unroll\n\n for ( int i = 0 ; i < N ; ++i )\n {\n if ( lane_id < N )\n {\n my_delta += my_delta_s[N * lane_id + i];\n }\n }\n }\n\n // Update the diagonal term.\n if ( ROW_MAJOR )\n {\n int block_inside_id = lane_id;\n#pragma unroll\n\n for (int j = 0; j < NUM_WARP_ITERS_PER_BLOCK; j++)\n {\n my_delta_s[block_inside_id] = my_Einv[j] * utils::shfl(my_delta, block_inside_id % N);\n block_inside_id += WARP_SIZE;\n }\n }\n\n // Reduce bmAx terms.\n {\n my_delta = 0.0;\n#pragma unroll\n\n for ( int i = 0 ; i < N ; ++i )\n {\n if ( lane_id < N )\n {\n my_delta += my_delta_s[N * lane_id + i];\n }\n }\n }\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n const int offset = N * a_row_id + lane_id;\n Vector_type my_b(0), my_x(0);\n\n if ( lane_id < N )\n {\n my_b = __cachingLoad(&delta[offset]);\n my_x = x [offset];\n }\n\n my_delta = my_b - my_delta;\n\n if ( lane_id < N )\n {\n x [offset] = my_x + weight * my_delta;\n Delta[offset] = my_delta;\n }\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename IndexType, typename ValueTypeA, typename ValueTypeB, typename WeightType, int CTA_SIZE, bool ROW_MAJOR >\n__global__\n__launch_bounds__( CTA_SIZE, 16 )\nvoid DILU_backward_4x4_kernel( const IndexType *row_offsets,\n const IndexType *column_indices,\n const ValueTypeA *nonzero_values,\n ValueTypeB *x,\n const WeightType weight,\n const int *sorted_rows_by_color,\n const int *__restrict row_colors,\n const ValueTypeA *Einv,\n const ValueTypeB *delta,\n ValueTypeB *Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const IndexType boundary_index)\n{\n const int nHalfWarps = CTA_SIZE / 16; // Number of half warps per CTA.\n const int laneId = utils::lane_id();\n const int halfWarpId = threadIdx.x / 16;\n const int halfLaneId = threadIdx.x % 16;\n const int halfLaneId_div_4 = halfLaneId / 4;\n const int halfLaneId_mod_4 = halfLaneId % 4;\n const int upperHalf = 16 * (laneId / 16);\n // Shared memory needed to exchange X and delta.\n __shared__ volatile ValueTypeB s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile ValueTypeB *my_s_mem = &s_mem[16 * halfWarpId];\n\n // Iterate over the rows of the matrix. One warp per two rows.\n for ( int aRowIt = blockIdx.x * nHalfWarps + halfWarpId ; aRowIt < num_rows_per_color ; aRowIt += gridDim.x * nHalfWarps )\n {\n int aRowId = sorted_rows_by_color[aRowIt];\n // Load one block of B.\n ValueTypeB my_delta(0);\n // The range of the rows.\n int aColBegin = row_offsets[aRowId ];\n int aColEnd = row_offsets[aRowId + 1];\n\n // Each warp load column indices of 16 nonzero blocks\n for ( ; aColBegin < aColEnd ; aColBegin += 16 )\n {\n int aColIt = aColBegin + halfLaneId;\n // Get the ID of the column.\n int aColTmp = -1, aColId = -1;\n\n if ( aColIt < aColEnd )\n {\n aColTmp = column_indices[aColIt];\n }\n\n#ifdef AMGX_ILU_COLORING\n bool valid = (((aColTmp < boundary_index || boundary_coloring == SYNC_COLORS) && (row_colors[aColTmp] > current_color)) || (aColTmp >= boundary_index && boundary_coloring == LAST));\n\n if ( aColTmp != -1 && valid )\n {\n aColId = aColTmp;\n }\n\n#else\n\n if ( aColTmp != -1 && row_colors[aColTmp] > current_color )\n {\n aColId = aColTmp;\n }\n\n#endif\n for ( int k = 0 ; k < 16 ; k += 4 )\n {\n int my_k = k + halfLaneId_div_4;\n // Exchange column indices.\n int waColId = utils::shfl( aColId, upperHalf + my_k );\n // Load 8 blocks of X if needed.\n ValueTypeB my_x(0);\n\n if ( waColId != -1 )\n {\n my_x = Delta[4 * waColId + halfLaneId_mod_4];\n }\n\n my_s_mem[halfLaneId] = my_x;\n // Load 8 blocks of A.\n#pragma unroll\n\n for ( int i = 0 ; i < 4 ; ++i )\n {\n const int k_i = k + i;\n int w_aColTmp = aColBegin + k_i, w_aColIt = -1;\n if ( utils::shfl( aColId, upperHalf + k_i ) != -1 && w_aColTmp < aColEnd )\n w_aColIt = w_aColTmp;\n\n ValueTypeA my_val(0);\n\n if ( w_aColIt != -1 )\n {\n my_val = nonzero_values[16 * w_aColIt + halfLaneId];\n }\n\n if ( ROW_MAJOR )\n {\n my_delta += my_val * my_s_mem[4 * i + halfLaneId_mod_4];\n }\n else\n {\n my_delta += my_val * my_s_mem[4 * i + halfLaneId_div_4];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n // Load EINV values.\n ValueTypeA my_Einv = Einv[16 * aRowId + halfLaneId];\n // Reduce delta terms.\n if ( ROW_MAJOR )\n {\n my_delta += utils::shfl_xor( my_delta, 1 );\n my_delta += utils::shfl_xor( my_delta, 2 );\n }\n else\n {\n my_delta += utils::shfl_xor( my_delta, 4 );\n my_delta += utils::shfl_xor( my_delta, 8 );\n }\n\n // Update the shared terms.\n if ( ROW_MAJOR )\n {\n if ( halfLaneId_mod_4 == 0 )\n {\n my_s_mem[halfLaneId_div_4] = my_delta;\n }\n }\n else\n {\n if ( halfLaneId_div_4 == 0 )\n {\n my_s_mem[halfLaneId_mod_4] = my_delta;\n }\n }\n\n // Update the diagonal term.\n if ( ROW_MAJOR )\n {\n my_delta = my_Einv * my_s_mem[halfLaneId_mod_4];\n }\n else\n {\n my_delta = my_Einv * my_s_mem[halfLaneId_div_4];\n }\n\n // Regroup results.\n if ( ROW_MAJOR )\n {\n my_delta += utils::shfl_xor( my_delta, 1 );\n my_delta += utils::shfl_xor( my_delta, 2 );\n }\n else\n {\n my_delta += utils::shfl_xor( my_delta, 4 );\n my_delta += utils::shfl_xor( my_delta, 8 );\n }\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n int offset = 4 * aRowId + halfLaneId_div_4;\n ValueTypeB my_b(0), my_x(0);\n\n if ( halfLaneId_mod_4 == 0 )\n {\n my_b = __cachingLoad(&delta[offset]);\n my_x = x[offset];\n }\n\n my_delta = my_b - my_delta;\n\n if ( halfLaneId_mod_4 == 0 )\n {\n x[offset] = my_x + weight * my_delta;\n Delta[offset] = my_delta;\n }\n }\n else\n {\n int offset = 4 * aRowId + halfLaneId_mod_4;\n ValueTypeB my_b(0), my_x(0);\n\n if ( halfLaneId_div_4 == 0 )\n {\n my_b = __cachingLoad(&delta[offset]);\n my_x = x[offset];\n }\n\n my_delta = my_b - my_delta;\n\n if ( halfLaneId_div_4 == 0 )\n {\n x[offset] = my_x + weight * my_delta;\n Delta[offset] = my_delta;\n }\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType, int CTA_SIZE >\n__global__\n__launch_bounds__( CTA_SIZE, 16 )\nvoid DILU_backward_4x4_kernel_row_major_vec4( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n // Number of half warps per CTA.\n const int NUM_HALF_WARPS = CTA_SIZE / 16;\n // Coordinates of the thread.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Coordinates of the thread in the CTA.\n const int thread_id_div_16 = threadIdx.x / 16;\n const int thread_id_mod_16 = threadIdx.x % 16;\n // Useful constants.\n const int thread_id_mod_16_div_4 = thread_id_mod_16 / 4;\n const int thread_id_mod_16_mod_4 = thread_id_mod_16 % 4;\n const int shfl_offset = 16 * (lane_id / 16);\n // Shared memory needed to exchange X and delta.\n __shared__ volatile Vector_type s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_mem = &s_mem[16 * thread_id_div_16];\n // The iterator over rows.\n int a_row_it = blockIdx.x * NUM_HALF_WARPS + thread_id_div_16;\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += gridDim.x * NUM_HALF_WARPS )\n {\n unsigned int active_mask = utils::activemask();\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_delta(0);\n // Don't do anything if X is zero.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; a_col_begin < a_col_end ; a_col_begin += 16 )\n {\n unsigned int active_mask_inner = utils::activemask();\n int a_col_it = a_col_begin + thread_id_mod_16;\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = __cachingLoad(&A_cols[a_col_it]);\n }\n\n#ifdef AMGX_ILU_COLORING\n int valid = false;\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n if ( boundary_coloring == LAST )\n {\n valid = a_col_id >= boundary_index;\n }\n else\n {\n valid = a_col_id < boundary_index && __cachingLoad(&row_colors[a_col_id]) > current_color;\n }\n }\n\n#else\n int valid = false;\n\n if ( a_col_id != -1 && row_colors[a_col_id] > current_color )\n {\n valid = true;\n }\n\n#endif\n\n // Set the column id.\n if ( !valid )\n {\n a_col_id = -1;\n }\n\n // Loop over columns. We compute 8 columns per iteration.\n#pragma unroll 2\n\n for ( int k = 0 ; k < 16 ; k += 4 )\n {\n int my_k = k + thread_id_mod_16_div_4;\n // Load 8 blocks of X.\n int uniform_a_col_id = utils::shfl( a_col_id, shfl_offset + my_k, warpSize, active_mask_inner );\n Vector_type my_Delta(0);\n\n if ( uniform_a_col_id != -1 )\n {\n my_Delta = Delta[4 * uniform_a_col_id + thread_id_mod_16_mod_4];\n }\n\n my_s_mem[thread_id_mod_16] = my_Delta;\n int uniform_a_col_it = a_col_begin + my_k;\n\n if ( uniform_a_col_id == -1 || uniform_a_col_it >= a_col_end )\n {\n uniform_a_col_it = -1;\n }\n\n Matrix_type my_vals[4] = { Matrix_type(0) };\n\n if ( uniform_a_col_it != -1 )\n {\n utils::load_vec4( my_vals, &A_vals[16 * uniform_a_col_it + 4 * thread_id_mod_16_mod_4] );\n }\n\n my_delta += my_vals[0] * my_s_mem[4 * thread_id_mod_16_div_4 + 0];\n my_delta += my_vals[1] * my_s_mem[4 * thread_id_mod_16_div_4 + 1];\n my_delta += my_vals[2] * my_s_mem[4 * thread_id_mod_16_div_4 + 2];\n my_delta += my_vals[3] * my_s_mem[4 * thread_id_mod_16_div_4 + 3];\n } // Loop over k\n } // Loop over aColIt\n\n // Load EINV values.\n Matrix_type my_Einv = Einv[16 * a_row_id + thread_id_mod_16];\n // Reduce delta terms.\n my_delta += utils::shfl_xor( my_delta, 4, warpSize, active_mask );\n my_delta += utils::shfl_xor( my_delta, 8, warpSize, active_mask );\n\n // Update the shared terms.\n if ( thread_id_mod_16_div_4 == 0 )\n {\n my_s_mem[thread_id_mod_16_mod_4] = my_delta;\n }\n\n // Update the diagonal term.\n my_delta = my_Einv * my_s_mem[thread_id_mod_16_mod_4];\n // Regroup results.\n my_delta += utils::shfl_xor( my_delta, 1, warpSize, active_mask );\n my_delta += utils::shfl_xor( my_delta, 2, warpSize, active_mask );\n // Store the results.\n int offset = 4 * a_row_id + thread_id_mod_16_div_4;\n Vector_type my_b(0), my_x(0);\n\n if ( thread_id_mod_16_mod_4 == 0 )\n {\n my_b = __cachingLoad(&delta[offset]);\n my_x = x [offset];\n }\n\n my_delta = my_b - my_delta;\n\n if ( thread_id_mod_16_mod_4 == 0 )\n {\n x [offset] = my_x + weight * my_delta;\n Delta[offset] = my_delta;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType, int NUM_THREADS_PER_ROW, int CTA_SIZE, int WARP_SIZE >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid DILU_backward_1x1_kernel( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const int boundary_index )\n{\n // Number of items per CTA.\n const int NUM_ROWS_PER_CTA = CTA_SIZE / NUM_THREADS_PER_ROW;\n // Number of items per grid.\n const int NUM_ROWS_PER_GRID = gridDim.x * NUM_ROWS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_mod_NTPR = lane_id % NUM_THREADS_PER_ROW;\n // Determine which NxN block the threads work with.\n int a_row_it = blockIdx.x * NUM_ROWS_PER_CTA + (threadIdx.x / NUM_THREADS_PER_ROW);\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += NUM_ROWS_PER_GRID )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n // Load one block of B.\n Vector_type my_delta(0);\n // Don't do anything if X is zero.\n int a_col_it = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Each warp load column indices of 32 nonzero blocks\n for ( a_col_it += lane_id_mod_NTPR ; utils::any( a_col_it < a_col_end ) ; a_col_it += NUM_THREADS_PER_ROW )\n {\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it really a valid column (due to coloring).\n int valid = false;\n#ifdef AMGX_ILU_COLORING\n\n if ( a_col_id != -1 && current_color != 0 )\n {\n //if( boundary_coloring == LAST )\n // valid = a_col_id >= boundary_index;\n //else\n // valid = a_col_id < boundary_index && row_colors[a_col_id] > current_color;\n valid = (((a_col_id < boundary_index || boundary_coloring == SYNC_COLORS) && (row_colors[a_col_id] > current_color)) || (a_col_id >= boundary_index && boundary_coloring == LAST));\n }\n\n#else\n\n //if( a_col_id != -1 && current_color != 0 )\n if ( a_col_id != -1 )\n {\n valid = row_colors[a_col_id] > current_color;\n }\n\n#endif\n // Load my Delta value.\n Vector_type my_Delta(0);\n\n if ( valid )\n {\n my_Delta = Delta[a_col_id];\n }\n\n // Load my item from A.\n Matrix_type my_val(0);\n\n if ( valid )\n {\n my_val = A_vals[a_col_it];\n }\n\n // Update bmAx.\n my_delta += my_val * my_Delta;\n }\n\n // Reduce bmAx terms.\n#pragma unroll\n\n for ( int mask = NUM_THREADS_PER_ROW / 2 ; mask > 0 ; mask >>= 1 )\n {\n my_delta += utils::shfl_xor( my_delta, mask );\n }\n\n // Store the results.\n if ( lane_id_mod_NTPR == 0 )\n {\n Vector_type my_x = __cachingLoad(&delta[a_row_id]) - Einv[a_row_id] * my_delta;\n x [a_row_id] += weight * my_x;\n Delta[a_row_id] = my_x;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType, int N, int CTA_SIZE >\n__global__\n__launch_bounds__( CTA_SIZE, 16 )\nvoid DILU_backward_NxN_kernel_skip( Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color )\n{\n const int NUM_ITEMS_PER_CTA = CTA_SIZE / N; // Number of updated block items per CTA\n const int ITEM_ID = threadIdx.x / N;\n const int ITEM_BLOCK_OFFSET = threadIdx.x % N;\n const int is_active = ITEM_ID < NUM_ITEMS_PER_CTA;\n // The first row.\n int a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + ITEM_ID;\n\n // Iterate over the rows of the matrix. One warp per two rows.\n for ( ; a_row_it < num_rows_per_color ; a_row_it += gridDim.x * NUM_ITEMS_PER_CTA )\n {\n if ( is_active )\n {\n int a_row_id = sorted_rows_by_color[a_row_it];\n const int idx = N * a_row_id + ITEM_BLOCK_OFFSET;\n Vector_type my_b = __cachingLoad(&delta[idx]);\n Vector_type my_x = x[idx];\n x[idx] = my_x + weight * my_b;\n Delta[idx] = my_b;\n }\n }\n}\n\n// ----------\n// Methods\n// ----------\n\ntemplate< typename Matrix_type, typename Vector_type, int N >\nvoid DILU_forward_NxN_dispatch( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const ColoringType boundary_coloring,\n const int boundary_index,\n const int row_major,\n const int has_external_diag )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int NUM_ROWS_PER_WARP = std::max(WARP_SIZE / NxN, 1);\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = NUM_ROWS_PER_WARP * NUM_WARPS_PER_CTA;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n // Branch to the correct kernel call.\n int code = 2 * (row_major ? 1 : 0) + (has_external_diag ? 1 : 0);\n\n switch ( code )\n {\n case 0: // Column-major, no external diagonal.\n DILU_forward_NxN_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n break;\n\n case 1: // Column-major, external diagonal.\n DILU_forward_NxN_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n break;\n\n case 2: // Row-major, no external diagonal.\n DILU_forward_NxN_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n break;\n\n case 3: // Row-major, external diagonal.\n DILU_forward_NxN_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n break;\n\n default:\n FatalError( \"Internal error\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n cudaCheckError();\n}\n\ntemplate< typename Matrix_type, typename Vector_type, int N >\nvoid DILU_forward_NxN_dispatch_large( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const ColoringType boundary_coloring,\n const int boundary_index,\n const int row_major,\n const int has_external_diag )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = NUM_WARPS_PER_CTA;\n // Each warp is going to sweep through bloock this many times\n const int NUM_WARP_ITERS_PER_BLOCK = ((NxN - 1) / WARP_SIZE) + 1;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n\n // Branch to the correct kernel call.\n if (!row_major)\n {\n FatalError(\"COL MAJOR is not supported for this large block_size\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n switch ( has_external_diag )\n {\n case 0: // Row-major, no external diagonal.\n DILU_forward_NxN_kernel_large <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n break;\n\n case 1: // Row-major, external diagonal.\n DILU_forward_NxN_kernel_large <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n break;\n\n default:\n FatalError( \"Internal error\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n cudaCheckError();\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type >\nvoid DILU_forward_NxN_dispatch( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *__restrict delta,\n const int *__restrict sorted_rows_by_color,\n const int num_rows_per_color,\n const int current_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const ColoringType boundary_coloring,\n const int boundary_index,\n const int block_size,\n const int row_major,\n const int has_external_diag )\n{\n switch ( block_size )\n {\n case 1:\n {\n const int NUM_THREADS_PER_ROW = 8;\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = CTA_SIZE / NUM_THREADS_PER_ROW;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n\n if ( has_external_diag )\n {\n DILU_forward_1x1_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n }\n else\n {\n DILU_forward_1x1_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n }\n\n cudaCheckError();\n }\n break;\n\n case 2:\n DILU_forward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index,\n row_major,\n has_external_diag );\n break;\n\n case 3:\n DILU_forward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index,\n row_major,\n has_external_diag );\n break;\n\n case 4:\n if ( row_major )\n {\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = CTA_SIZE / 16;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n\n if ( has_external_diag )\n //DILU_forward_4x4_kernel<<>>(\n DILU_forward_4x4_kernel_row_major_vec4 <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n else\n DILU_forward_4x4_kernel_row_major_vec4 <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index );\n\n cudaCheckError();\n }\n else\n DILU_forward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index,\n row_major,\n has_external_diag );\n\n break;\n\n case 5:\n DILU_forward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index,\n row_major,\n has_external_diag );\n break;\n\n case 8:\n DILU_forward_NxN_dispatch_large(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index,\n row_major,\n has_external_diag );\n break;\n\n case 10:\n DILU_forward_NxN_dispatch_large(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n delta,\n sorted_rows_by_color,\n num_rows_per_color,\n current_color,\n row_colors,\n Einv,\n boundary_coloring,\n boundary_index,\n row_major,\n has_external_diag );\n break;\n\n default:\n FatalError( \"Internal error\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n cudaCheckError();\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType, int N >\nvoid DILU_backward_NxN_dispatch( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const int boundary_index,\n const int row_major )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int NUM_ROWS_PER_WARP = std::max(WARP_SIZE / NxN, 1);\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = NUM_ROWS_PER_WARP * NUM_WARPS_PER_CTA;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n\n // Branch to the correct kernel call.\n if ( row_major )\n {\n DILU_backward_NxN_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index );\n }\n else\n {\n DILU_backward_NxN_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index );\n }\n\n cudaCheckError();\n}\n\n\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType, int N, int NUM_WARP_ITERS_PER_BLOCK >\nvoid DILU_backward_NxN_dispatch_large( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const int boundary_index,\n const int row_major )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = NUM_WARPS_PER_CTA;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n\n // Branch to the correct kernel call.\n if ( row_major )\n {\n DILU_backward_NxN_kernel_large <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index );\n }\n else\n {\n FatalError(\"col major is not supported for this blocksize in multicolor DILU solver\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n cudaCheckError();\n}\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Matrix_type, typename Vector_type, typename WeightType>\nvoid DILU_backward_NxN_dispatch( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n Vector_type *__restrict x,\n const WeightType weight,\n const int *__restrict sorted_rows_by_color,\n const int *__restrict row_colors,\n const Matrix_type *__restrict Einv,\n const Vector_type *delta,\n Vector_type *__restrict Delta,\n const int num_rows_per_color,\n const int current_color,\n const ColoringType boundary_coloring,\n const int boundary_index,\n const int block_size,\n const int row_major )\n{\n switch ( block_size )\n {\n case 1:\n {\n const int NUM_THREADS_PER_ROW = 8;\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = CTA_SIZE / NUM_THREADS_PER_ROW;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n DILU_backward_1x1_kernel <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index );\n cudaCheckError();\n }\n break;\n\n case 2:\n DILU_backward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index,\n row_major );\n break;\n\n case 3:\n DILU_backward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index,\n row_major );\n break;\n\n case 4:\n\n //if( false )\n if ( row_major )\n {\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = CTA_SIZE / 16;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n //DILU_backward_NxN_kernel<<>>(\n DILU_backward_4x4_kernel_row_major_vec4 <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index );\n cudaCheckError();\n }\n else\n DILU_backward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index,\n row_major );\n\n break;\n\n case 5:\n DILU_backward_NxN_dispatch(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index,\n row_major );\n break;\n\n case 8:\n DILU_backward_NxN_dispatch_large(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index,\n row_major );\n break;\n\n case 10:\n DILU_backward_NxN_dispatch_large(\n A_rows,\n A_cols,\n A_vals,\n x,\n weight,\n sorted_rows_by_color,\n row_colors,\n Einv,\n delta,\n Delta,\n num_rows_per_color,\n current_color,\n boundary_coloring,\n boundary_index,\n row_major );\n break;\n\n default:\n FatalError( \"Internal error\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n cudaCheckError();\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< class T_Config >\nMulticolorDILUSolver_Base::MulticolorDILUSolver_Base( AMG_Config &cfg,\n const std::string &cfg_scope,\n ThreadManager *tmng ) :\n Solver( cfg, cfg_scope, tmng )\n{\n this->weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n this->m_reorder_cols_by_color_desired = (cfg.AMG_Config::template getParameter(\"reorder_cols_by_color\", cfg_scope) != 0);\n this->m_insert_diagonal_desired = (cfg.AMG_Config::template getParameter(\"insert_diag_while_reordering\", cfg_scope) != 0);\n this->m_boundary_coloring = cfg.AMG_Config::template getParameter(\"boundary_coloring\", cfg_scope);\n this->always_obey_coloring = 0;\n\n if (weight == 0)\n {\n weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Multicolor DILU smoother\\n\");\n }\n}\n\n// Destructor\ntemplate\nMulticolorDILUSolver_Base::~MulticolorDILUSolver_Base()\n{\n Einv.clear();\n Einv.shrink_to_fit();\n}\n\ntemplate\nvoid MulticolorDILUSolver_Base::computeEinv(Matrix &A)\n{\n ViewType oldView = A.currentView();\n A.setView(this->m_explicit_A->getViewExterior());\n\n if ( A.get_block_dimx() != A.get_block_dimy() )\n {\n FatalError(\"DILU implemented only for squared blocks\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if ( A.get_block_dimx() > 32) // actually much more less than 32 doe to register file limitations, but...\n {\n FatalError(\"DILU implemented only for squared blocks of size <= 32\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n computeEinv_NxN( A, A.get_block_dimx() );\n A.setView(oldView);\n}\n\ntemplate< class T_Config >\nvoid\nMulticolorDILUSolver_Base::printSolverParameters() const\n{\n std::cout << \"relaxation_factor = \" << this->weight << std::endl;\n}\n\n// Solver setup\ntemplate< class T_Config >\nvoid\nMulticolorDILUSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n m_explicit_A = dynamic_cast*>(this->m_A);\n\n if (!this->m_explicit_A)\n {\n FatalError(\"MulticolorDILUSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n int N = this->m_explicit_A->get_num_cols() * this->m_explicit_A->get_block_dimy();\n\n if (this->m_explicit_A->getColoringLevel() < 1)\n {\n FatalError(\"Matrix must be colored to use multicolor dilu solver. Try setting: coloring_level=1 in the configuration file\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n m_delta.resize(N);\n m_Delta.resize(N);\n m_delta.set_block_dimy(this->m_explicit_A->get_block_dimy());\n m_Delta.set_block_dimy(this->m_explicit_A->get_block_dimy());\n m_delta.set_block_dimx(1);\n m_Delta.set_block_dimx(1);\n\n if ( this->m_explicit_A->getBlockFormat() != ROW_MAJOR )\n {\n FatalError(\"Multicolor DILU solver only supports row major format for the blocks\", AMGX_ERR_CONFIGURATION);\n }\n\n computeEinv( *this->m_explicit_A );\n}\n\n//\ntemplate< class T_Config >\nvoid\nMulticolorDILUSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nMulticolorDILUSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n if ( this->m_explicit_A->get_block_dimx() != this->m_explicit_A->get_block_dimy() )\n {\n FatalError(\"DILU implemented only for squared blocks\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if ( this->m_explicit_A->get_block_dimx() > 32) // actually much more less than 32 doe to register file limitations, but...\n {\n FatalError(\"DILU implemented only for squared blocks of size <= 32\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if (xIsZero)\n {\n x.dirtybit = 0;\n }\n\n if (!this->m_explicit_A->is_matrix_singleGPU())\n {\n this->m_explicit_A->manager->exchange_halo_async(x, x.tag);\n this->m_explicit_A->manager->exchange_halo_async(b, b.tag);\n }\n\n if (this->m_explicit_A->getViewExterior() == this->m_explicit_A->getViewInterior())\n {\n if (!this->m_explicit_A->is_matrix_singleGPU())\n {\n this->m_explicit_A->manager->exchange_halo_wait(x, x.tag);\n this->m_explicit_A->manager->exchange_halo_wait(b, b.tag);\n }\n }\n\n ViewType oldView = this->m_explicit_A->currentView();\n ViewType flags;\n bool latencyHiding = true;\n\n if (this->m_explicit_A->is_matrix_singleGPU() || (x.dirtybit == 0 && b.dirtybit == 0))\n {\n latencyHiding = false;\n this->m_explicit_A->setViewExterior();\n flags = (ViewType)(this->m_explicit_A->getViewExterior());\n }\n else\n {\n flags = (ViewType)(this->m_explicit_A->getViewInterior());\n this->m_explicit_A->setViewInterior();\n }\n\n if (xIsZero)\n {\n thrust_wrapper::fill(x.begin(), x.end(), types::util::get_zero());\n cudaCheckError();\n }\n\n this->smooth_NxN(*this->m_explicit_A, b, x, flags);\n\n if (latencyHiding)\n {\n if (!this->m_explicit_A->is_matrix_singleGPU())\n {\n this->m_explicit_A->manager->exchange_halo_wait(x, x.tag);\n this->m_explicit_A->manager->exchange_halo_wait(b, b.tag);\n }\n\n this->m_explicit_A->setViewExterior();\n flags = (ViewType)(~(this->m_explicit_A->getViewInterior()) & this->m_explicit_A->getViewExterior());\n\n if (flags != 0)\n {\n this->smooth_NxN(*this->m_explicit_A, b, x, flags);\n }\n }\n\n x.dirtybit = 1;\n this->m_explicit_A->setView(oldView);\n return (this->converged(b, x));\n}\n\ntemplate\nvoid\nMulticolorDILUSolver_Base::solve_finalize( VVector &b, VVector &x )\n{}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid MulticolorDILUSolver >::computeEinv_NxN(const Matrix_h &A, const int bsize)\n{\n FatalError(\"Multicolor DILU smoother not implemented for host format, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid MulticolorDILUSolver >::smooth_NxN( const Matrix_h &A, VVector &b, VVector &x, ViewType separation_flag )\n{\n FatalError(\"Haven't implemented Multicolor DILU smoother for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nMulticolorDILUSolver >::MulticolorDILUSolver(\n AMG_Config &cfg,\n const std::string &cfg_scope,\n ThreadManager *tmng ) :\n MulticolorDILUSolver_Base >( cfg, cfg_scope, tmng )\n{\n int device = 0;\n hipGetDevice( &device );\n cudaCheckError();\n hipDeviceProp_t properties;\n hipGetDeviceProperties( &properties, device );\n cudaCheckError();\n m_is_kepler = properties.major >= 3;\n}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid MulticolorDILUSolver >::computeEinv_NxN(const Matrix_d &A, const int bsize)\n{\n const int bsize_sq = bsize * bsize;\n this->Einv.resize( A.get_num_cols()*bsize_sq, 0.0 );\n\n // sol::prof_start();\n for ( int i = 0, num_colors = A.getMatrixColoring().getNumColors() ; i < num_colors ; ++i )\n {\n const int color_offset = A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n const int num_rows_per_color = A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1] - color_offset;\n\n if ( num_rows_per_color == 0 )\n {\n continue;\n }\n\n const int CTA_SIZE = 128;\n const int NUM_WARPS_PER_CTA = CTA_SIZE / 32;\n int ROWS_PER_WARP = 1;\n\n if ( bsize_sq > 1 && bsize_sq < 6 )\n {\n ROWS_PER_WARP = 32 / bsize_sq;\n }\n\n const int ROWS_PER_CTA = ROWS_PER_WARP * NUM_WARPS_PER_CTA;\n const int GRID_SIZE = std::min( 4096, (num_rows_per_color + ROWS_PER_CTA - 1) / ROWS_PER_CTA );\n hipStream_t stream = amgx::thrust::global_thread_handle::get_stream();\n\n switch ( bsize )\n {\n case 1:\n DILU_setup_1x1_kernel <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->Einv.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n num_rows_per_color,\n i );\n break;\n\n case 2:\n DILU_setup_NxN_kernel <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->Einv.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n num_rows_per_color,\n i );\n break;\n\n case 3:\n DILU_setup_NxN_kernel <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->Einv.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n num_rows_per_color,\n i );\n break;\n\n case 4:\n DILU_setup_NxN_kernel <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->Einv.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n num_rows_per_color,\n i );\n break;\n\n case 5:\n DILU_setup_NxN_kernel <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->Einv.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n num_rows_per_color,\n i );\n break;\n\n case 8:\n DILU_setup_NxN_kernel_large <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->Einv.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n num_rows_per_color,\n i );\n break;\n\n case 10:\n DILU_setup_NxN_kernel_large <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n this->Einv.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n num_rows_per_color,\n i );\n break;\n\n default:\n FatalError( \"Multicolor-DILU Setup: block size was not enabled in the code, contact AMGX developers.\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE );\n }\n\n cudaCheckError();\n }\n}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid MulticolorDILUSolver >::smooth_NxN( const Matrix_d &A, VVector &b, VVector &x, ViewType separation_flag )\n{\n AMGX_CPU_PROFILER( \"MulticolorDILUSolver::smooth_NxN \" );\n int offset = 0, separation = 0;\n A.getOffsetAndSizeForView(INTERIOR, &offset, &separation);\n\n // Only have separation=num interior rows if we are only working on the interior\n // and the boundary coloring is FIRST or LAST, otherwise set separation offset to\n // total number of rows\n if ( separation_flag != this->m_explicit_A->getViewInterior() ||\n this->m_explicit_A->getViewExterior() == this->m_explicit_A->getViewInterior() ||\n this->m_boundary_coloring != LAST && this->m_boundary_coloring != FIRST )\n {\n separation = A.row_offsets.size() - 1;\n }\n else\n {\n amgx_printf(\"separation active\\n\");\n }\n\n // --------------------\n // Forward Sweep\n // --------------------\n const int num_colors = this->m_explicit_A->getMatrixColoring().getNumColors();\n\n for ( int i = 0 ; i < num_colors ; ++i )\n {\n int color_offset(0);\n\n if ( separation_flag & INTERIOR )\n {\n color_offset = A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n }\n else\n {\n color_offset = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i];\n }\n\n int num_rows_per_color(0);\n\n if ( separation_flag == this->m_explicit_A->getViewInterior() )\n {\n num_rows_per_color = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i];\n }\n else\n {\n num_rows_per_color = A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1];\n }\n\n num_rows_per_color -= color_offset;\n\n if ( num_rows_per_color == 0 )\n {\n continue;\n }\n\n int boundary_index = separation;\n\n if ( this->m_boundary_coloring == SYNC_COLORS )\n {\n boundary_index = A.get_num_rows();\n }\n\n DILU_forward_NxN_dispatch(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n A.diag.raw(),\n x.raw(),\n b.raw(),\n this->m_delta.raw(),\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n num_rows_per_color,\n i,\n A.getMatrixColoring().getRowColors().raw(),\n this->Einv.raw(),\n this->m_boundary_coloring,\n boundary_index,\n A.get_block_dimy(),\n A.getBlockFormat() == ROW_MAJOR,\n A.hasProps(DIAG) );\n cudaCheckError();\n }\n\n // --------------------\n // Backward Sweep\n // --------------------\n for ( int i = num_colors - 1 ; i >= 0 ; --i )\n {\n int color_offset(0);\n\n if ( separation_flag & INTERIOR )\n {\n color_offset = A.getMatrixColoring().getOffsetsRowsPerColor()[i];\n }\n else\n {\n color_offset = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i];\n }\n\n int num_rows_per_color(0);\n\n if ( separation_flag == this->m_explicit_A->getViewInterior() )\n {\n num_rows_per_color = A.getMatrixColoring().getSeparationOffsetsRowsPerColor()[i];\n }\n else\n {\n num_rows_per_color = A.getMatrixColoring().getOffsetsRowsPerColor()[i + 1];\n }\n\n num_rows_per_color -= color_offset;\n\n if ( num_rows_per_color == 0 )\n {\n continue;\n }\n\n if ( i == num_colors - 1 )\n {\n const int NUM_ROWS_PER_CTA = CTA_SIZE / A.get_block_dimy();\n const int GRID_SIZE = std::min( 4096, (num_rows_per_color + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n\n switch ( A.get_block_dimy() )\n {\n case 1:\n DILU_backward_NxN_kernel_skip <<< GRID_SIZE, CTA_SIZE>>>(\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color );\n break;\n\n case 2:\n DILU_backward_NxN_kernel_skip <<< GRID_SIZE, CTA_SIZE>>>(\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color );\n break;\n\n case 3:\n DILU_backward_NxN_kernel_skip <<< GRID_SIZE, CTA_SIZE>>>(\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color );\n break;\n\n case 4:\n DILU_backward_NxN_kernel_skip <<< GRID_SIZE, CTA_SIZE>>>(\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color );\n break;\n\n case 5:\n DILU_backward_NxN_kernel_skip <<< GRID_SIZE, CTA_SIZE>>>(\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color );\n break;\n\n case 8:\n DILU_backward_NxN_kernel_skip <<< GRID_SIZE, CTA_SIZE>>>(\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color );\n break;\n\n case 10:\n DILU_backward_NxN_kernel_skip <<< GRID_SIZE, CTA_SIZE>>>(\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color );\n break;\n }\n\n cudaCheckError();\n }\n else\n {\n DILU_backward_NxN_dispatch(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n x.raw(),\n this->weight,\n A.getMatrixColoring().getSortedRowsByColor().raw() + color_offset,\n A.getMatrixColoring().getRowColors().raw(),\n this->Einv.raw(),\n this->m_delta.raw(),\n this->m_Delta.raw(),\n num_rows_per_color,\n i,\n this->m_boundary_coloring,\n separation,\n A.get_block_dimy(),\n A.getBlockFormat() == ROW_MAJOR );\n cudaCheckError();\n }\n }\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class MulticolorDILUSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n// AMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class MulticolorDILUSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n// AMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n} // namespace amgx\n\n", "source": "/kaggle/working/amgx/src/solvers/multicolor_dilu_solver.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace amgx\n{\nnamespace cf_jacobi_solver\n{\n\nstruct is_coarse\n{\n __host__ __device__\n int operator()(const int &x)\n {\n return (int) (x == COARSE);\n }\n};\n\nstruct is_fine\n{\n __host__ __device__\n int operator()(const int &x)\n {\n return (int) (x == FINE);\n }\n};\n\nstruct is_eq_minus_one\n{\n __host__ __device__\n int operator()(const int &x)\n {\n return (int) (x == -1);\n }\n};\n\n\ntemplate \nstruct jacobi_presmooth_functor\n{\n ValueTypeB omega;\n jacobi_presmooth_functor( ValueTypeB omega ) : omega( omega ) {}\n __host__ __device__ ValueTypeB operator()( const ValueTypeB &b, const ValueTypeA &d ) const { return isNotCloseToZero(d) ? omega * b / d : omega * b / epsilon(d); }\n};\n\ntemplate \nstruct jacobi_postsmooth_functor\n{\n ValueTypeB omega;\n jacobi_postsmooth_functor( ValueTypeB omega ) : omega( omega ) {}\n template __host__ __device__ ValueTypeB operator( )( const Tuple &t ) const\n {\n ValueTypeB x = amgx::thrust::get<0>(t);\n ValueTypeA d = amgx::thrust::get<1>(t);\n ValueTypeB b = amgx::thrust::get<2>(t);\n ValueTypeB y = amgx::thrust::get<3>(t);\n // return x + omega * (b - y) / d.\n d = isNotCloseToZero(d) ? d : epsilon(d);\n d = ValueTypeA( 1 ) / d;\n b -= y;\n b *= omega;\n return b * d + x;\n }\n};\n\ntemplate \nstruct add_functor\n{\n __host__ __device__ ValueTypeB operator()( const ValueTypeB &x, const ValueTypeB &y )const { return x + y; }\n};\n\ntemplate\n__device__ __forceinline__ T fmnaOp (T a, T b, T c)\n{\n return -(a * b) + c;\n}\ntemplate\n__device__ __forceinline__ T mulOp (T a, T b)\n{\n return a * b;\n}\ntemplate\n__device__ __forceinline__ T rcpOp (T a)\n{\n return 1.0 / (isNotCloseToZero(a) ? a : epsilon(a));\n}\ntemplate\n__device__ __forceinline__ T absOp (T a)\n{\n return fabs(a);\n}\n\n// -----------------------------------\n// KERNELS\n// -----------------------------------\n\ntemplate\n__global__\nvoid jacobi_zero_ini_masked_step(const int *row_ids, const int nrows_to_process, const ValueTypeA *Dinv, const ValueTypeB *b, const ValueTypeB relaxation_factor, ValueTypeB *x)\n{\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (tid < nrows_to_process)\n {\n int i = row_ids[tid];\n ValueTypeB d = Dinv[i];\n d = isNotCloseToZero(d) ? d : epsilon(d);\n x[i] = relaxation_factor * b[i] / d;\n tid += (blockDim.x * gridDim.x);\n }\n}\n\ntemplate\n__global__\nvoid jacobi_masked_step(const int *row_ids, const int nrows_to_process, const ValueTypeA *Dinv, const ValueTypeB *b, const ValueTypeB *y, const ValueTypeB relaxation_factor, ValueTypeB *x)\n{\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (tid < nrows_to_process)\n {\n int i = row_ids[tid];\n ValueTypeB d = Dinv[i];\n d = isNotCloseToZero(d) ? d : epsilon(d);\n x[i] += relaxation_factor * (b[i] - y[i]) / d;\n tid += (blockDim.x * gridDim.x);\n }\n}\n\n__global__\nvoid agg_write_agg(const int *agg_map, const int nrows, int *dst)\n{\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (tid < nrows)\n {\n dst[agg_map[tid]] = tid;\n tid += (blockDim.x * gridDim.x);\n }\n}\n\n\n//--------------------------------\n// Methods\n//--------------------------------\n\n// Constructor\ntemplate\nCFJacobiSolver_Base::CFJacobiSolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope)\n{\n weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n int param_mode = cfg.AMG_Config::template getParameter(\"cf_smoothing_mode\", cfg_scope);\n\n switch (param_mode)\n {\n case 0:\n this->mode = CF_CF;\n break;\n\n case 1:\n this->mode = CF_FC;\n break;\n\n case 2:\n this->mode = CF_FCF;\n break;\n\n case 3:\n this->mode = CF_CFC;\n break;\n\n default:\n this->mode = CF_CF;\n }\n\n if (weight == 0)\n {\n weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Block Jacobi smoother\\n\");\n }\n}\n\n// Destructor\ntemplate\nCFJacobiSolver_Base::~CFJacobiSolver_Base()\n{\n this->Dinv.resize(0);\n}\n\ntemplate\nvoid\nCFJacobiSolver_Base::printSolverParameters() const\n{\n std::cout << \"relaxation_factor= \" << this->weight << std::endl;\n}\n\n// Solver setup\ntemplate\nvoid\nCFJacobiSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n Matrix *A_as_matrix = dynamic_cast*>(this->m_A);\n\n if (!A_as_matrix)\n {\n FatalError(\"CFJacobiSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n computeDinv( *A_as_matrix );\n\n if ( A_as_matrix->getBlockFormat() != ROW_MAJOR )\n {\n FatalError(\" CFJacobiSolver only supports row major format\", AMGX_ERR_CONFIGURATION);\n }\n\n if (A_as_matrix->hasParameter(\"cf_map\")) // classical case\n {\n IVector *cf_map = A_as_matrix->template getParameterPtr< IVector >(\"cf_map\");\n int nrows = A_as_matrix->get_num_rows();\n this->num_coarse = amgx::thrust::count(cf_map->begin(), cf_map->end(), (int)COARSE);\n this->c_rows.resize(this->num_coarse);\n this->f_rows.resize(nrows - this->num_coarse);\n amgx::thrust::counting_iterator zero(0);\n amgx::thrust::counting_iterator zero_plus_nrows = zero + nrows;\n amgx::thrust::copy_if(zero, zero_plus_nrows, cf_map->begin(), this->c_rows.begin(), is_coarse());\n amgx::thrust::copy_if(zero, zero_plus_nrows, cf_map->begin(), this->f_rows.begin(), is_fine());\n cudaCheckError();\n // partitioning check\n /*\n {\n typedef typename TConfig::template setMemSpace::Type TConfig_h;\n typedef typename TConfig_h::template setVecPrec::Type ivec_value_type_h;\n typedef Vector IVector_h;\n\n IVector_h cf_map_h = *cf_map;\n IVector_h cf_rows = this->cf_rows;\n for (int i = 0; i < cf_map_h.size(); i++)\n {\n\n if ((cf_map_h[cf_rows[i]] == FINE && i < this->num_coarse) || (cf_map_h[cf_rows[i]] == COARSE && i >= this->num_coarse))\n printf(\"CFJ FAIL at i==%d: cf_rows[i]==%d, cf_map[row]==%d, num_coarse==%d\\n\", i, cf_rows[i], cf_map_h[cf_rows[i]], this->num_coarse);\n }\n }\n */\n }\n else if (A_as_matrix->hasParameter(\"aggregates_map\")) // aggregation case\n {\n IVector *agg_map = A_as_matrix->template getParameterPtr< IVector >(\"aggregates_map\");\n int agg_num = A_as_matrix->template getParameter< int >(\"aggregates_num\");\n this->num_coarse = agg_num;\n const int threads_per_block = 256;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A_as_matrix->get_num_rows() + threads_per_block - 1) / threads_per_block);\n int nrows = A_as_matrix->get_num_rows();\n this->c_rows.resize(this->num_coarse);\n this->f_rows.resize(nrows - this->num_coarse);\n IVector tmap(nrows, -1);\n // can use thrust permutation operator here.\n agg_write_agg <<< num_blocks, threads_per_block>>>(agg_map->raw(), agg_map->size(), this->c_rows.raw());\n cudaCheckError();\n agg_write_agg <<< num_blocks, threads_per_block>>>(this->c_rows.raw(), this->num_coarse, tmap.raw());\n cudaCheckError();\n amgx::thrust::counting_iterator zero(0);\n amgx::thrust::counting_iterator zero_plus_nrows = zero + nrows;\n amgx::thrust::copy_if(zero, zero_plus_nrows, tmap.begin(), this->f_rows.begin(), is_eq_minus_one());\n cudaCheckError();\n }\n else\n {\n FatalError(\"No info from AMG level was found for C-F separation, use different smoother or drink 1 beer\", AMGX_ERR_BAD_PARAMETERS);\n }\n}\n\n//\ntemplate\nvoid\nCFJacobiSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nCFJacobiSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n //bool done = false;\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n int *smoothing_direction = A_as_matrix->template getParameterPtr (\"smoothing_direction\");\n // smoothing direction == 0 : presmoothing\n // smoothing direction == 1 : postsmoothing\n SmoothingOrder current_order;\n\n switch (this->mode)\n {\n case CF_CF:\n current_order = (*smoothing_direction == 0) ? CF_CF : CF_FC;\n break;\n\n case CF_FC:\n current_order = (*smoothing_direction == 1) ? CF_CF : CF_FC;\n break;\n\n case CF_FCF:\n current_order = (*smoothing_direction == 0) ? CF_FCF : CF_CFC;\n break;\n\n case CF_CFC:\n current_order = (*smoothing_direction == 1) ? CF_FCF : CF_CFC;\n break;\n }\n\n// no multi-gpu for now\n ViewType flags = OWNED;\n\n if (xIsZero) { x.dirtybit = 0; }\n\n if (A_as_matrix->get_block_dimx() == 1 && A_as_matrix->get_block_dimy() == 1)\n {\n if (xIsZero)\n {\n smooth_with_0_initial_guess_1x1(*A_as_matrix, b, x, current_order, flags);\n }\n else\n {\n smooth_1x1(*A_as_matrix, b, x, current_order, flags);\n }\n }\n else\n {\n FatalError(\"Unsupported block size for BlockJacobi_Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nCFJacobiSolver_Base::solve_finalize( VVector &b, VVector &x )\n{}\n\ntemplate\nvoid CFJacobiSolver_Base::computeDinv( Matrix &A)\n{\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n ViewType oldView = A.currentView();\n A.setView(A_as_matrix->getViewExterior());\n\n if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n this->computeDinv_1x1(A);\n }\n\n A.setView(oldView);\n}\n\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid CFJacobiSolver >::computeDinv_1x1(const Matrix_d &A)\n{\n Matrix_d *A_as_matrix = (Matrix_d *) this->m_A;\n// supports both diag\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n\n if ( A_as_matrix->hasProps(DIAG) )\n {\n const int num_values = A_as_matrix->diagOffset() * A_as_matrix->get_block_size();\n amgx::thrust::copy( A_as_matrix->values.begin() + num_values, A_as_matrix->values.begin() + num_values + A_as_matrix->get_num_rows()*A_as_matrix->get_block_size(), this->Dinv.begin() );\n cudaCheckError();\n }\n else\n {\n find_diag( *A_as_matrix );\n }\n}\n\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid CFJacobiSolver >::computeDinv_1x1(const Matrix_h &A)\n{\n // Do nothing\n}\n\n\n// Finding diag on device, CSR format\ntemplate \nvoid CFJacobiSolver >::find_diag( const Matrix_h &A )\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n if (A.col_indices[j] == i)\n {\n this->Dinv[i] = A.values[j];\n break;\n }\n\n if (j == A.row_offsets[i + 1] - 1)\n {\n FatalError(\"Could not find a diagonal value\", AMGX_ERR_BAD_PARAMETERS);\n }\n }\n }\n}\n\n\n// Finding diag on device, CSR format\ntemplate \nvoid CFJacobiSolver >::find_diag( const Matrix_d &A )\n{\n AMGX_CPU_PROFILER( \"JacobiSolver::find_diag \" );\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((ValueTypeB)A.get_num_rows() / (ValueTypeB)THREADS_PER_BLOCK));\n find_diag_kernel_indexed_dia <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK>>>(\n A.get_num_rows(),\n A.diag.raw(),\n A.values.raw(),\n this->Dinv.raw());\n cudaCheckError();\n}\n\ntemplate \nvoid CFJacobiSolver >::smooth_1x1(Matrix_h &A, VVector &b, VVector &x, SmoothingOrder order, ViewType separation_flags)\n{\n VVector newx(x.size());\n\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeB Axi = 0.0;\n ValueTypeB d = A.values[A.diag[i]];\n ValueTypeB mydiaginv = this->weight / (isNotCloseToZero(d) ? d : epsilon(d) );\n\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n Axi += A.values[j] * x[A.col_indices[j]];\n }\n\n newx[i] = x[i] + (b[i] - Axi) * mydiaginv ;\n }\n\n x.swap(newx);\n}\n\ntemplate \nvoid CFJacobiSolver >::smooth_1x1(Matrix_d &A, VVector &b, VVector &x, SmoothingOrder order, ViewType separation_flags)\n{\n AMGX_CPU_PROFILER( \"JacobiSolver::smooth_1x1 \" );\n\n if (this->y.size() != b.size())\n {\n this->y.resize(b.size());\n this->y.tag = this->tag * 100 + 3;\n this->y.set_block_dimx(b.get_block_dimx());\n this->y.set_block_dimy(b.get_block_dimy());\n }\n\n int num_rows = A.get_num_rows();\n int offset = 0;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n const int threads_per_block = 256;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() + threads_per_block - 1) / threads_per_block);\n\n if (order == CF_CF || order == CF_FC)\n {\n IVector &rows_first = (order == CF_CF) ? this->c_rows : this->f_rows;\n IVector &rows_second = (order == CF_CF) ? this->f_rows : this->c_rows;\n const int threads_per_block = 256;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() + threads_per_block - 1) / threads_per_block);\n // if use transform + permutation iterator it will yield into two separate permutation reads - for src and dst, so using simple kernel here\n multiply_masked( A, x, this->y, rows_first, separation_flags );\n jacobi_masked_step <<< num_blocks, threads_per_block>>>(\n rows_first.raw(),\n rows_first.size(),\n this->Dinv.raw(),\n b.raw(),\n this->y.raw(),\n this->weight,\n x.raw());\n cudaCheckError();\n multiply_masked( A, x, this->y, rows_second, separation_flags );\n jacobi_masked_step <<< num_blocks, threads_per_block>>>(\n rows_second.raw(),\n rows_second.size(),\n this->Dinv.raw(),\n b.raw(),\n this->y.raw(),\n this->weight,\n x.raw());\n cudaCheckError();\n }\n else\n {\n FatalError(\"CF_FCF and CF_CFC is not yet done\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n\ntemplate \nvoid CFJacobiSolver >::smooth_with_0_initial_guess_1x1( Matrix_h &A, VVector &b, VVector &x, SmoothingOrder order, ViewType separation_flags)\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeB d = A.values[A.diag[i]];\n ValueTypeB mydiag = this->weight / (isNotCloseToZero(d) ? d : epsilon(d));\n x[i] = b[i] * mydiag;\n }\n}\n\ntemplate \nvoid CFJacobiSolver >::smooth_with_0_initial_guess_1x1(Matrix_d &A, VVector &b, VVector &x, SmoothingOrder order, ViewType separation_flags)\n{\n AMGX_CPU_PROFILER( \"JacobiSolver::smooth_with_0_initial_guess_1x1 \" );\n int num_rows = A.get_num_rows();\n int offset = 0;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n const int threads_per_block = 256;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() + threads_per_block - 1) / threads_per_block);\n\n if (this->y.size() != b.size())\n {\n this->y.resize(b.size());\n this->y.tag = this->tag * 100 + 3;\n this->y.set_block_dimx(b.get_block_dimx());\n this->y.set_block_dimy(b.get_block_dimy());\n }\n\n if (order == CF_CF || order == CF_FC)\n {\n// IVector& rows_first = (order == CF_CF) ? this->c_rows : this->f_rows;\n IVector &rows_second = (order == CF_CF) ? this->f_rows : this->c_rows;\n const int threads_per_block = 256;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() + threads_per_block - 1) / threads_per_block);\n /*jacobi_zero_ini_masked_step<<>>(\n rows_first.raw(),\n rows_first.size(),\n this->Dinv.raw(),\n b.raw(),\n this->weight,\n x.raw());*/\n cudaCheckError();\n // it is not so much harder to initialize whole vector instead of just C or F points\n amgx::thrust::transform(b.begin( ),\n b.begin( ) + A.get_num_rows(),\n this->Dinv.begin( ),\n x.begin( ),\n jacobi_presmooth_functor( this->weight ));\n cudaCheckError();\n multiply_masked( A, x, this->y, rows_second, separation_flags );\n cudaCheckError();\n jacobi_masked_step <<< num_blocks, threads_per_block>>>(\n rows_second.raw(),\n rows_second.size(),\n this->Dinv.raw(),\n b.raw(),\n this->y.raw(),\n this->weight,\n x.raw());\n cudaCheckError();\n }\n else\n {\n FatalError(\"CF_FCF and CF_CFC is not yet done\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n cudaCheckError();\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class CFJacobiSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class CFJacobiSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace block_jacobi\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace amgx\n{\nnamespace cf_jacobi_solver\n{\n\nstruct is_coarse\n{\n __host__ __device__\n int operator()(const int &x)\n {\n return (int) (x == COARSE);\n }\n};\n\nstruct is_fine\n{\n __host__ __device__\n int operator()(const int &x)\n {\n return (int) (x == FINE);\n }\n};\n\nstruct is_eq_minus_one\n{\n __host__ __device__\n int operator()(const int &x)\n {\n return (int) (x == -1);\n }\n};\n\n\ntemplate \nstruct jacobi_presmooth_functor\n{\n ValueTypeB omega;\n jacobi_presmooth_functor( ValueTypeB omega ) : omega( omega ) {}\n __host__ __device__ ValueTypeB operator()( const ValueTypeB &b, const ValueTypeA &d ) const { return isNotCloseToZero(d) ? omega * b / d : omega * b / epsilon(d); }\n};\n\ntemplate \nstruct jacobi_postsmooth_functor\n{\n ValueTypeB omega;\n jacobi_postsmooth_functor( ValueTypeB omega ) : omega( omega ) {}\n template __host__ __device__ ValueTypeB operator( )( const Tuple &t ) const\n {\n ValueTypeB x = amgx::thrust::get<0>(t);\n ValueTypeA d = amgx::thrust::get<1>(t);\n ValueTypeB b = amgx::thrust::get<2>(t);\n ValueTypeB y = amgx::thrust::get<3>(t);\n // return x + omega * (b - y) / d.\n d = isNotCloseToZero(d) ? d : epsilon(d);\n d = ValueTypeA( 1 ) / d;\n b -= y;\n b *= omega;\n return b * d + x;\n }\n};\n\ntemplate \nstruct add_functor\n{\n __host__ __device__ ValueTypeB operator()( const ValueTypeB &x, const ValueTypeB &y )const { return x + y; }\n};\n\ntemplate\n__device__ __forceinline__ T fmnaOp (T a, T b, T c)\n{\n return -(a * b) + c;\n}\ntemplate\n__device__ __forceinline__ T mulOp (T a, T b)\n{\n return a * b;\n}\ntemplate\n__device__ __forceinline__ T rcpOp (T a)\n{\n return 1.0 / (isNotCloseToZero(a) ? a : epsilon(a));\n}\ntemplate\n__device__ __forceinline__ T absOp (T a)\n{\n return fabs(a);\n}\n\n// -----------------------------------\n// KERNELS\n// -----------------------------------\n\ntemplate\n__global__\nvoid jacobi_zero_ini_masked_step(const int *row_ids, const int nrows_to_process, const ValueTypeA *Dinv, const ValueTypeB *b, const ValueTypeB relaxation_factor, ValueTypeB *x)\n{\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (tid < nrows_to_process)\n {\n int i = row_ids[tid];\n ValueTypeB d = Dinv[i];\n d = isNotCloseToZero(d) ? d : epsilon(d);\n x[i] = relaxation_factor * b[i] / d;\n tid += (blockDim.x * gridDim.x);\n }\n}\n\ntemplate\n__global__\nvoid jacobi_masked_step(const int *row_ids, const int nrows_to_process, const ValueTypeA *Dinv, const ValueTypeB *b, const ValueTypeB *y, const ValueTypeB relaxation_factor, ValueTypeB *x)\n{\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (tid < nrows_to_process)\n {\n int i = row_ids[tid];\n ValueTypeB d = Dinv[i];\n d = isNotCloseToZero(d) ? d : epsilon(d);\n x[i] += relaxation_factor * (b[i] - y[i]) / d;\n tid += (blockDim.x * gridDim.x);\n }\n}\n\n__global__\nvoid agg_write_agg(const int *agg_map, const int nrows, int *dst)\n{\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (tid < nrows)\n {\n dst[agg_map[tid]] = tid;\n tid += (blockDim.x * gridDim.x);\n }\n}\n\n\n//--------------------------------\n// Methods\n//--------------------------------\n\n// Constructor\ntemplate\nCFJacobiSolver_Base::CFJacobiSolver_Base( AMG_Config &cfg, const std::string &cfg_scope) : Solver( cfg, cfg_scope)\n{\n weight = cfg.AMG_Config::template getParameter(\"relaxation_factor\", cfg_scope);\n int param_mode = cfg.AMG_Config::template getParameter(\"cf_smoothing_mode\", cfg_scope);\n\n switch (param_mode)\n {\n case 0:\n this->mode = CF_CF;\n break;\n\n case 1:\n this->mode = CF_FC;\n break;\n\n case 2:\n this->mode = CF_FCF;\n break;\n\n case 3:\n this->mode = CF_CFC;\n break;\n\n default:\n this->mode = CF_CF;\n }\n\n if (weight == 0)\n {\n weight = 1.;\n amgx_printf(\"Warning, setting weight to 1 instead of estimating largest_eigen_value in Block Jacobi smoother\\n\");\n }\n}\n\n// Destructor\ntemplate\nCFJacobiSolver_Base::~CFJacobiSolver_Base()\n{\n this->Dinv.resize(0);\n}\n\ntemplate\nvoid\nCFJacobiSolver_Base::printSolverParameters() const\n{\n std::cout << \"relaxation_factor= \" << this->weight << std::endl;\n}\n\n// Solver setup\ntemplate\nvoid\nCFJacobiSolver_Base::solver_setup(bool reuse_matrix_structure)\n{\n Matrix *A_as_matrix = dynamic_cast*>(this->m_A);\n\n if (!A_as_matrix)\n {\n FatalError(\"CFJacobiSolver only works with explicit matrices\", AMGX_ERR_INTERNAL);\n }\n\n computeDinv( *A_as_matrix );\n\n if ( A_as_matrix->getBlockFormat() != ROW_MAJOR )\n {\n FatalError(\" CFJacobiSolver only supports row major format\", AMGX_ERR_CONFIGURATION);\n }\n\n if (A_as_matrix->hasParameter(\"cf_map\")) // classical case\n {\n IVector *cf_map = A_as_matrix->template getParameterPtr< IVector >(\"cf_map\");\n int nrows = A_as_matrix->get_num_rows();\n this->num_coarse = amgx::thrust::count(cf_map->begin(), cf_map->end(), (int)COARSE);\n this->c_rows.resize(this->num_coarse);\n this->f_rows.resize(nrows - this->num_coarse);\n amgx::thrust::counting_iterator zero(0);\n amgx::thrust::counting_iterator zero_plus_nrows = zero + nrows;\n amgx::thrust::copy_if(zero, zero_plus_nrows, cf_map->begin(), this->c_rows.begin(), is_coarse());\n amgx::thrust::copy_if(zero, zero_plus_nrows, cf_map->begin(), this->f_rows.begin(), is_fine());\n cudaCheckError();\n // partitioning check\n /*\n {\n typedef typename TConfig::template setMemSpace::Type TConfig_h;\n typedef typename TConfig_h::template setVecPrec::Type ivec_value_type_h;\n typedef Vector IVector_h;\n\n IVector_h cf_map_h = *cf_map;\n IVector_h cf_rows = this->cf_rows;\n for (int i = 0; i < cf_map_h.size(); i++)\n {\n\n if ((cf_map_h[cf_rows[i]] == FINE && i < this->num_coarse) || (cf_map_h[cf_rows[i]] == COARSE && i >= this->num_coarse))\n printf(\"CFJ FAIL at i==%d: cf_rows[i]==%d, cf_map[row]==%d, num_coarse==%d\\n\", i, cf_rows[i], cf_map_h[cf_rows[i]], this->num_coarse);\n }\n }\n */\n }\n else if (A_as_matrix->hasParameter(\"aggregates_map\")) // aggregation case\n {\n IVector *agg_map = A_as_matrix->template getParameterPtr< IVector >(\"aggregates_map\");\n int agg_num = A_as_matrix->template getParameter< int >(\"aggregates_num\");\n this->num_coarse = agg_num;\n const int threads_per_block = 256;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A_as_matrix->get_num_rows() + threads_per_block - 1) / threads_per_block);\n int nrows = A_as_matrix->get_num_rows();\n this->c_rows.resize(this->num_coarse);\n this->f_rows.resize(nrows - this->num_coarse);\n IVector tmap(nrows, -1);\n // can use thrust permutation operator here.\n agg_write_agg <<< num_blocks, threads_per_block>>>(agg_map->raw(), agg_map->size(), this->c_rows.raw());\n cudaCheckError();\n agg_write_agg <<< num_blocks, threads_per_block>>>(this->c_rows.raw(), this->num_coarse, tmap.raw());\n cudaCheckError();\n amgx::thrust::counting_iterator zero(0);\n amgx::thrust::counting_iterator zero_plus_nrows = zero + nrows;\n amgx::thrust::copy_if(zero, zero_plus_nrows, tmap.begin(), this->f_rows.begin(), is_eq_minus_one());\n cudaCheckError();\n }\n else\n {\n FatalError(\"No info from AMG level was found for C-F separation, use different smoother or drink 1 beer\", AMGX_ERR_BAD_PARAMETERS);\n }\n}\n\n//\ntemplate\nvoid\nCFJacobiSolver_Base::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n}\n\n// Solve one iteration\ntemplate\nAMGX_STATUS\nCFJacobiSolver_Base::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n //bool done = false;\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n int *smoothing_direction = A_as_matrix->template getParameterPtr (\"smoothing_direction\");\n // smoothing direction == 0 : presmoothing\n // smoothing direction == 1 : postsmoothing\n SmoothingOrder current_order;\n\n switch (this->mode)\n {\n case CF_CF:\n current_order = (*smoothing_direction == 0) ? CF_CF : CF_FC;\n break;\n\n case CF_FC:\n current_order = (*smoothing_direction == 1) ? CF_CF : CF_FC;\n break;\n\n case CF_FCF:\n current_order = (*smoothing_direction == 0) ? CF_FCF : CF_CFC;\n break;\n\n case CF_CFC:\n current_order = (*smoothing_direction == 1) ? CF_FCF : CF_CFC;\n break;\n }\n\n// no multi-gpu for now\n ViewType flags = OWNED;\n\n if (xIsZero) { x.dirtybit = 0; }\n\n if (A_as_matrix->get_block_dimx() == 1 && A_as_matrix->get_block_dimy() == 1)\n {\n if (xIsZero)\n {\n smooth_with_0_initial_guess_1x1(*A_as_matrix, b, x, current_order, flags);\n }\n else\n {\n smooth_1x1(*A_as_matrix, b, x, current_order, flags);\n }\n }\n else\n {\n FatalError(\"Unsupported block size for BlockJacobi_Solver\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n return this->converged( b, x );\n}\n\ntemplate\nvoid\nCFJacobiSolver_Base::solve_finalize( VVector &b, VVector &x )\n{}\n\ntemplate\nvoid CFJacobiSolver_Base::computeDinv( Matrix &A)\n{\n Matrix *A_as_matrix = (Matrix *) this->m_A;\n ViewType oldView = A.currentView();\n A.setView(A_as_matrix->getViewExterior());\n\n if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n this->computeDinv_1x1(A);\n }\n\n A.setView(oldView);\n}\n\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid CFJacobiSolver >::computeDinv_1x1(const Matrix_d &A)\n{\n Matrix_d *A_as_matrix = (Matrix_d *) this->m_A;\n// supports both diag\n this->Dinv.resize(A.get_num_rows()*A.get_block_dimx()*A.get_block_dimy(), 0.0);\n\n if ( A_as_matrix->hasProps(DIAG) )\n {\n const int num_values = A_as_matrix->diagOffset() * A_as_matrix->get_block_size();\n amgx::thrust::copy( A_as_matrix->values.begin() + num_values, A_as_matrix->values.begin() + num_values + A_as_matrix->get_num_rows()*A_as_matrix->get_block_size(), this->Dinv.begin() );\n cudaCheckError();\n }\n else\n {\n find_diag( *A_as_matrix );\n }\n}\n\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid CFJacobiSolver >::computeDinv_1x1(const Matrix_h &A)\n{\n // Do nothing\n}\n\n\n// Finding diag on device, CSR format\ntemplate \nvoid CFJacobiSolver >::find_diag( const Matrix_h &A )\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n if (A.col_indices[j] == i)\n {\n this->Dinv[i] = A.values[j];\n break;\n }\n\n if (j == A.row_offsets[i + 1] - 1)\n {\n FatalError(\"Could not find a diagonal value\", AMGX_ERR_BAD_PARAMETERS);\n }\n }\n }\n}\n\n\n// Finding diag on device, CSR format\ntemplate \nvoid CFJacobiSolver >::find_diag( const Matrix_d &A )\n{\n AMGX_CPU_PROFILER( \"JacobiSolver::find_diag \" );\n const size_t THREADS_PER_BLOCK = 128;\n const size_t NUM_BLOCKS = std::min(AMGX_GRID_MAX_SIZE, (int)ceil((ValueTypeB)A.get_num_rows() / (ValueTypeB)THREADS_PER_BLOCK));\n find_diag_kernel_indexed_dia <<< (unsigned int)NUM_BLOCKS, (unsigned int)THREADS_PER_BLOCK>>>(\n A.get_num_rows(),\n A.diag.raw(),\n A.values.raw(),\n this->Dinv.raw());\n cudaCheckError();\n}\n\ntemplate \nvoid CFJacobiSolver >::smooth_1x1(Matrix_h &A, VVector &b, VVector &x, SmoothingOrder order, ViewType separation_flags)\n{\n VVector newx(x.size());\n\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeB Axi = 0.0;\n ValueTypeB d = A.values[A.diag[i]];\n ValueTypeB mydiaginv = this->weight / (isNotCloseToZero(d) ? d : epsilon(d) );\n\n //for each column\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n Axi += A.values[j] * x[A.col_indices[j]];\n }\n\n newx[i] = x[i] + (b[i] - Axi) * mydiaginv ;\n }\n\n x.swap(newx);\n}\n\ntemplate \nvoid CFJacobiSolver >::smooth_1x1(Matrix_d &A, VVector &b, VVector &x, SmoothingOrder order, ViewType separation_flags)\n{\n AMGX_CPU_PROFILER( \"JacobiSolver::smooth_1x1 \" );\n\n if (this->y.size() != b.size())\n {\n this->y.resize(b.size());\n this->y.tag = this->tag * 100 + 3;\n this->y.set_block_dimx(b.get_block_dimx());\n this->y.set_block_dimy(b.get_block_dimy());\n }\n\n int num_rows = A.get_num_rows();\n int offset = 0;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n const int threads_per_block = 256;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() + threads_per_block - 1) / threads_per_block);\n\n if (order == CF_CF || order == CF_FC)\n {\n IVector &rows_first = (order == CF_CF) ? this->c_rows : this->f_rows;\n IVector &rows_second = (order == CF_CF) ? this->f_rows : this->c_rows;\n const int threads_per_block = 256;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() + threads_per_block - 1) / threads_per_block);\n // if use transform + permutation iterator it will yield into two separate permutation reads - for src and dst, so using simple kernel here\n multiply_masked( A, x, this->y, rows_first, separation_flags );\n jacobi_masked_step <<< num_blocks, threads_per_block>>>(\n rows_first.raw(),\n rows_first.size(),\n this->Dinv.raw(),\n b.raw(),\n this->y.raw(),\n this->weight,\n x.raw());\n cudaCheckError();\n multiply_masked( A, x, this->y, rows_second, separation_flags );\n jacobi_masked_step <<< num_blocks, threads_per_block>>>(\n rows_second.raw(),\n rows_second.size(),\n this->Dinv.raw(),\n b.raw(),\n this->y.raw(),\n this->weight,\n x.raw());\n cudaCheckError();\n }\n else\n {\n FatalError(\"CF_FCF and CF_CFC is not yet done\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n\ntemplate \nvoid CFJacobiSolver >::smooth_with_0_initial_guess_1x1( Matrix_h &A, VVector &b, VVector &x, SmoothingOrder order, ViewType separation_flags)\n{\n //for each row\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n ValueTypeB d = A.values[A.diag[i]];\n ValueTypeB mydiag = this->weight / (isNotCloseToZero(d) ? d : epsilon(d));\n x[i] = b[i] * mydiag;\n }\n}\n\ntemplate \nvoid CFJacobiSolver >::smooth_with_0_initial_guess_1x1(Matrix_d &A, VVector &b, VVector &x, SmoothingOrder order, ViewType separation_flags)\n{\n AMGX_CPU_PROFILER( \"JacobiSolver::smooth_with_0_initial_guess_1x1 \" );\n int num_rows = A.get_num_rows();\n int offset = 0;\n A.getOffsetAndSizeForView(separation_flags, &offset, &num_rows);\n const int threads_per_block = 256;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() + threads_per_block - 1) / threads_per_block);\n\n if (this->y.size() != b.size())\n {\n this->y.resize(b.size());\n this->y.tag = this->tag * 100 + 3;\n this->y.set_block_dimx(b.get_block_dimx());\n this->y.set_block_dimy(b.get_block_dimy());\n }\n\n if (order == CF_CF || order == CF_FC)\n {\n// IVector& rows_first = (order == CF_CF) ? this->c_rows : this->f_rows;\n IVector &rows_second = (order == CF_CF) ? this->f_rows : this->c_rows;\n const int threads_per_block = 256;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() + threads_per_block - 1) / threads_per_block);\n /*jacobi_zero_ini_masked_step<<>>(\n rows_first.raw(),\n rows_first.size(),\n this->Dinv.raw(),\n b.raw(),\n this->weight,\n x.raw());*/\n cudaCheckError();\n // it is not so much harder to initialize whole vector instead of just C or F points\n amgx::thrust::transform(b.begin( ),\n b.begin( ) + A.get_num_rows(),\n this->Dinv.begin( ),\n x.begin( ),\n jacobi_presmooth_functor( this->weight ));\n cudaCheckError();\n multiply_masked( A, x, this->y, rows_second, separation_flags );\n cudaCheckError();\n jacobi_masked_step <<< num_blocks, threads_per_block>>>(\n rows_second.raw(),\n rows_second.size(),\n this->Dinv.raw(),\n b.raw(),\n this->y.raw(),\n this->weight,\n x.raw());\n cudaCheckError();\n }\n else\n {\n FatalError(\"CF_FCF and CF_CFC is not yet done\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n cudaCheckError();\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class CFJacobiSolver_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class CFJacobiSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace block_jacobi\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/solvers/cf_jacobi_solver.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n\n#include // for amgx::thrust::max_element\n\nnamespace amgx\n{\n\ntemplate\n__global__\nvoid getLambdaEstimate(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, const int num_rows, ValueTypeB *out)\n\n{\n int row_id = blockDim.x * blockIdx.x + threadIdx.x;\n ValueTypeB max_sum = (ValueTypeB)0.0;\n\n while (row_id < num_rows)\n {\n ValueTypeB cur_sum = (ValueTypeB)0.0;\n\n for (int j = row_offsets[row_id]; j < row_offsets[row_id + 1]; j++)\n {\n cur_sum += abs(values[j]);\n }\n\n if (diag)\n {\n cur_sum += abs(values[dia_indices[row_id]]);\n }\n\n max_sum = max(max_sum, cur_sum);\n row_id += gridDim.x * blockDim.x;\n }\n\n out[blockDim.x * blockIdx.x + threadIdx.x] = max_sum;\n}\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid Chebyshev_Solver::compute_eigenmax_estimate(const Matrix &A, ValueTypeB &lambda)\n{\n#define LAMBDA_BLOCK_SIZE 256\n VVector tsum(A.get_num_rows());\n const int threads_per_block = 256;\n const int blockrows_per_cta = threads_per_block;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n\n if (A.hasProps(DIAG))\n {\n cudaFuncSetCacheConfig(getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, true >, cudaFuncCachePreferL1);\n getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, true > <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, A.get_num_rows(), tsum.raw());\n cudaCheckError();\n }\n else\n {\n cudaFuncSetCacheConfig(getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, false >, cudaFuncCachePreferL1);\n getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, false > <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, A.get_num_rows(), tsum.raw());\n cudaCheckError();\n }\n\n lambda = *(amgx::thrust::max_element(tsum.begin(), tsum.end()));\n}\n\n// Constructor\ntemplate< class T_Config>\nChebyshev_Solver::Chebyshev_Solver( AMG_Config &cfg, const std::string &cfg_scope) :\n Solver( cfg, cfg_scope),\n m_buffer_N(0), m_eigsolver(NULL)\n{\n std::string solverName, new_scope, tmp_scope;\n cfg.getParameter( \"preconditioner\", solverName, cfg_scope, new_scope );\n m_lambda_mode = cfg.AMG_Config::template getParameter(\"chebyshev_lambda_estimate_mode\", cfg_scope);\n m_cheby_order = cfg.AMG_Config::template getParameter(\"chebyshev_polynomial_order\", cfg_scope);\n // 0 - use eigensolver to get BOTH estimates\n // 1 - use eigensolver to get maximum estimate\n // 2 - use max sum of abs values as a rough estimate for maximum eigenvalue\n // 3 - use user provided cheby_max_lambda and cheby_min_lambda\n\n if (m_lambda_mode == 3)\n {\n m_user_max_lambda = cfg.AMG_Config::template getParameter(\"cheby_max_lambda\", cfg_scope);\n m_user_min_lambda = cfg.AMG_Config::template getParameter(\"cheby_min_lambda\", cfg_scope);\n }\n\n if (solverName.compare(\"NOSOLVER\") == 0)\n {\n no_preconditioner = true;\n m_preconditioner = NULL;\n }\n else\n {\n no_preconditioner = false;\n m_preconditioner = SolverFactory::allocate( cfg, cfg_scope, \"preconditioner\" );\n }\n\n std::string eig_cfg_string = \"algorithm=AGGREGATION,\\n\"\n \"eig_solver=LANCZOS,\\n\"\n \"verbosity_level=0,\\n\"\n \"eig_max_iters=128,\\n\"\n \"eig_tolerance=1e-4,\\n\"\n \"eig_which=largest,\\n\"\n \"eig_eigenvector=0,\\n\"\n \"eig_eigenvector_solver=default\";\n\n if (m_lambda_mode < 2)\n {\n AMG_Configuration eig_cfg;\n eig_cfg.parseParameterString(eig_cfg_string.c_str());\n m_eigsolver = EigenSolverFactory::allocate(*eig_cfg.getConfigObject(), \"default\", \"eig_solver\");\n }\n}\n\ntemplate\nChebyshev_Solver::~Chebyshev_Solver()\n{\n if (!no_preconditioner) { delete m_preconditioner; }\n\n if (!m_eigsolver) { delete m_eigsolver; }\n}\n\n\ntemplate\nvoid\nChebyshev_Solver::solver_setup(bool reuse_matrix_structure)\n{\n AMGX_CPU_PROFILER( \"Chebyshev_Solver::solver_setup \" );\n ViewType oldView = this->m_A->currentView();\n this->m_A->setViewExterior();\n\n // Setup the preconditionner\n if (!no_preconditioner)\n {\n m_preconditioner->setup(*this->m_A, reuse_matrix_structure);\n }\n\n // The number of elements in temporary vectors.\n this->m_buffer_N = static_cast( this->m_A->get_num_cols() * this->m_A->get_block_dimy() );\n Matrix *mtx_A = dynamic_cast*>(this->m_A);\n VVector eig_solver_t_x;\n\n // m_lambda_mode:\n // 0: use eigensolver to get lmin and lmax estimate\n // 1: use eigensolver to get lmax estimate, set lmin = lmax/8\n // 2: use max row sum as lmax estimate, set lmin = lmax/8\n if (m_lambda_mode < 2)\n {\n if (!no_preconditioner)\n {\n SolverOperator *MA = new SolverOperator (this->m_A, m_preconditioner);\n m_eigsolver->setup(*MA);\n m_eigsolver->solve(eig_solver_t_x);\n delete MA;\n }\n else\n {\n m_eigsolver->setup(*mtx_A);\n m_eigsolver->solve(eig_solver_t_x);\n }\n\n const std::vector &lambdas = m_eigsolver->get_eigenvalues();\n this->m_lmax = lambdas[0] * 1.05;\n\n if (m_lambda_mode == 0)\n {\n this->m_lmin = lambdas[lambdas.size() - 1] * 0.95;\n }\n else\n {\n this->m_lmin = this->m_lmax * 0.125;\n }\n }\n else if (m_lambda_mode == 2)\n {\n if (no_preconditioner)\n {\n Matrix *pA = dynamic_cast< Matrix* > (this->m_A);\n compute_eigenmax_estimate(*pA, this->m_lmax);\n this->m_lmin = this->m_lmax * 0.125;\n }\n else\n {\n // assuming that this preconditioner would be good enough to reduce spectrum to the largest eigen value = 1.0\n this->m_lmax = 0.9;\n this->m_lmin = this->m_lmax * 0.125;\n }\n }\n else if (m_lambda_mode == 3)\n {\n if (no_preconditioner)\n {\n Matrix *pA = dynamic_cast< Matrix* > (this->m_A);\n compute_eigenmax_estimate(*pA, this->m_lmax);\n this->m_lmin = this->m_lmax * 0.125;\n }\n else\n {\n // Use user input estimates\n this->m_lmax = this->m_user_max_lambda;\n this->m_lmin = this->m_user_min_lambda;\n }\n }\n else\n {\n FatalError(\"Not supported chebyshev_lambda_estimate_mode.\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n // Allocate memory needed for iterating.\n m_p.resize( this->m_buffer_N );\n m_z.resize( this->m_buffer_N );\n m_Ap.resize( this->m_buffer_N );\n m_xp.resize( this->m_buffer_N );\n m_rp.resize( this->m_buffer_N );\n m_p.set_block_dimy(this->m_A->get_block_dimy());\n m_p.set_block_dimx(1);\n m_p.dirtybit = 1;\n m_p.delayed_send = 1;\n m_p.tag = this->tag * 100 + 1;\n m_Ap.set_block_dimy(this->m_A->get_block_dimy());\n m_Ap.set_block_dimx(1);\n m_Ap.dirtybit = 1;\n m_Ap.delayed_send = 1;\n m_Ap.tag = this->tag * 100 + 2;\n m_z.set_block_dimy(this->m_A->get_block_dimy());\n m_z.set_block_dimx(1);\n m_z.dirtybit = 1;\n m_z.delayed_send = 1;\n m_z.tag = this->tag * 100 + 3;\n m_xp.set_block_dimy(this->m_A->get_block_dimy());\n m_xp.set_block_dimx(1);\n m_xp.dirtybit = 1;\n m_xp.delayed_send = 1;\n m_xp.tag = this->tag * 100 + 4;\n this->m_A->setView(oldView);\n}\n\ntemplate\nvoid\nChebyshev_Solver::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n AMGX_CPU_PROFILER( \"Chebyshev_Solver::solve_init \" );\n Operator &A = *this->m_A;\n ViewType oldView = A.currentView();\n A.setViewExterior();\n int offset, size;\n A.getOffsetAndSizeForView(A.getViewExterior(), &offset, &size);\n\n // Run one iteration of preconditioner with zero initial guess\n if (no_preconditioner)\n {\n copy(*this->m_r, m_z, offset, size);\n }\n else\n {\n m_z.delayed_send = 1;\n this->m_r->delayed_send = 1;\n m_preconditioner->solve( *this->m_r, m_z, true );\n m_z.delayed_send = 1;\n this->m_r->delayed_send = 1;\n }\n\n // m_p - res after precond\n copy( m_z, m_p, offset, size );\n A.setView(oldView);\n m_gamma = 0.;\n m_beta = 0.;\n first_iter = 0;\n}\n\ntemplate\nAMGX_STATUS\nChebyshev_Solver::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n AMGX_STATUS conv_stat = AMGX_ST_NOT_CONVERGED;\n\n AMGX_CPU_PROFILER( \"Chebyshev_Solver::solve_iteration \" );\n Operator &A = *this->m_A;\n ViewType oldView = A.currentView();\n A.setViewExterior();\n int offset, size;\n A.getOffsetAndSizeForView(A.getViewExterior(), &offset, &size);\n ValueTypeB a = (this->m_lmax + this->m_lmin) / 2;\n ValueTypeB c = (this->m_lmax - this->m_lmin) / 2;\n\n for (int i = 0; i < m_cheby_order; i++)\n {\n // apply precond\n if (no_preconditioner)\n {\n copy(*this->m_r, m_z, offset, size);\n }\n else\n {\n m_z.delayed_send = 1;\n this->m_r->delayed_send = 1;\n m_preconditioner->solve( *this->m_r, m_z, true );\n m_z.delayed_send = 1;\n this->m_r->delayed_send = 1;\n }\n\n if (first_iter == 0)\n {\n m_gamma = 1. / a;\n first_iter = 1;\n }\n else\n {\n m_beta = c * c * m_gamma * m_gamma / 4.;\n\n if (m_gamma != ValueTypeB(0) && (a - (m_beta / m_gamma)) != ValueTypeB(0))\n {\n m_gamma = 1. / (a - m_beta / m_gamma);\n }\n\n axpby( m_z, m_p, m_p, ValueTypeB( 1 ), m_beta, offset, size);\n }\n\n axpy( m_p, x, m_gamma, offset, size );\n this->compute_residual( b, x);\n }\n\n // Do we converge ?\n if ( this->m_monitor_convergence &&\n isDone( ( conv_stat = this->compute_norm_and_converged() ) ) )\n {\n A.setView(oldView);\n return conv_stat;\n }\n\n // No convergence so far.\n A.setView(oldView);\n return this->m_monitor_convergence ? AMGX_ST_NOT_CONVERGED : AMGX_ST_CONVERGED;\n}\n\ntemplate\nvoid\nChebyshev_Solver::solve_finalize( VVector &b, VVector &x )\n{}\n\ntemplate\nvoid\nChebyshev_Solver::printSolverParameters() const\n{\n if (!no_preconditioner)\n {\n std::cout << \"preconditioner: \" << this->m_preconditioner->getName()\n << \" with scope name: \"\n << this->m_preconditioner->getScope() << std::endl;\n }\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class Chebyshev_Solver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n#include // for amgx::thrust::max_element\n\nnamespace amgx\n{\n\ntemplate\n__global__\nvoid getLambdaEstimate(const IndexType *row_offsets, const IndexType *column_indices, const ValueTypeA *values, const IndexType *dia_indices, const int num_rows, ValueTypeB *out)\n\n{\n int row_id = blockDim.x * blockIdx.x + threadIdx.x;\n ValueTypeB max_sum = (ValueTypeB)0.0;\n\n while (row_id < num_rows)\n {\n ValueTypeB cur_sum = (ValueTypeB)0.0;\n\n for (int j = row_offsets[row_id]; j < row_offsets[row_id + 1]; j++)\n {\n cur_sum += abs(values[j]);\n }\n\n if (diag)\n {\n cur_sum += abs(values[dia_indices[row_id]]);\n }\n\n max_sum = max(max_sum, cur_sum);\n row_id += gridDim.x * blockDim.x;\n }\n\n out[blockDim.x * blockIdx.x + threadIdx.x] = max_sum;\n}\n\n// Method to compute the inverse of the diagonal blocks\ntemplate \nvoid Chebyshev_Solver::compute_eigenmax_estimate(const Matrix &A, ValueTypeB &lambda)\n{\n#define LAMBDA_BLOCK_SIZE 256\n VVector tsum(A.get_num_rows());\n const int threads_per_block = 256;\n const int blockrows_per_cta = threads_per_block;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() - 1) / blockrows_per_cta + 1);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueTypeA *A_nonzero_values_ptr = A.values.raw();\n\n if (A.hasProps(DIAG))\n {\n hipFuncSetCacheConfig(reinterpret_cast(getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, true >), hipFuncCachePreferL1);\n getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, true > <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, A.get_num_rows(), tsum.raw());\n cudaCheckError();\n }\n else\n {\n hipFuncSetCacheConfig(reinterpret_cast(getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, false >), hipFuncCachePreferL1);\n getLambdaEstimate < IndexType, ValueTypeA, ValueTypeB, LAMBDA_BLOCK_SIZE, LAMBDA_BLOCK_SIZE / 32, false > <<< num_blocks, threads_per_block>>>\n (A_row_offsets_ptr, A_column_indices_ptr, A_nonzero_values_ptr, A_dia_idx_ptr, A.get_num_rows(), tsum.raw());\n cudaCheckError();\n }\n\n lambda = *(amgx::thrust::max_element(tsum.begin(), tsum.end()));\n}\n\n// Constructor\ntemplate< class T_Config>\nChebyshev_Solver::Chebyshev_Solver( AMG_Config &cfg, const std::string &cfg_scope) :\n Solver( cfg, cfg_scope),\n m_buffer_N(0), m_eigsolver(NULL)\n{\n std::string solverName, new_scope, tmp_scope;\n cfg.getParameter( \"preconditioner\", solverName, cfg_scope, new_scope );\n m_lambda_mode = cfg.AMG_Config::template getParameter(\"chebyshev_lambda_estimate_mode\", cfg_scope);\n m_cheby_order = cfg.AMG_Config::template getParameter(\"chebyshev_polynomial_order\", cfg_scope);\n // 0 - use eigensolver to get BOTH estimates\n // 1 - use eigensolver to get maximum estimate\n // 2 - use max sum of abs values as a rough estimate for maximum eigenvalue\n // 3 - use user provided cheby_max_lambda and cheby_min_lambda\n\n if (m_lambda_mode == 3)\n {\n m_user_max_lambda = cfg.AMG_Config::template getParameter(\"cheby_max_lambda\", cfg_scope);\n m_user_min_lambda = cfg.AMG_Config::template getParameter(\"cheby_min_lambda\", cfg_scope);\n }\n\n if (solverName.compare(\"NOSOLVER\") == 0)\n {\n no_preconditioner = true;\n m_preconditioner = NULL;\n }\n else\n {\n no_preconditioner = false;\n m_preconditioner = SolverFactory::allocate( cfg, cfg_scope, \"preconditioner\" );\n }\n\n std::string eig_cfg_string = \"algorithm=AGGREGATION,\\n\"\n \"eig_solver=LANCZOS,\\n\"\n \"verbosity_level=0,\\n\"\n \"eig_max_iters=128,\\n\"\n \"eig_tolerance=1e-4,\\n\"\n \"eig_which=largest,\\n\"\n \"eig_eigenvector=0,\\n\"\n \"eig_eigenvector_solver=default\";\n\n if (m_lambda_mode < 2)\n {\n AMG_Configuration eig_cfg;\n eig_cfg.parseParameterString(eig_cfg_string.c_str());\n m_eigsolver = EigenSolverFactory::allocate(*eig_cfg.getConfigObject(), \"default\", \"eig_solver\");\n }\n}\n\ntemplate\nChebyshev_Solver::~Chebyshev_Solver()\n{\n if (!no_preconditioner) { delete m_preconditioner; }\n\n if (!m_eigsolver) { delete m_eigsolver; }\n}\n\n\ntemplate\nvoid\nChebyshev_Solver::solver_setup(bool reuse_matrix_structure)\n{\n AMGX_CPU_PROFILER( \"Chebyshev_Solver::solver_setup \" );\n ViewType oldView = this->m_A->currentView();\n this->m_A->setViewExterior();\n\n // Setup the preconditionner\n if (!no_preconditioner)\n {\n m_preconditioner->setup(*this->m_A, reuse_matrix_structure);\n }\n\n // The number of elements in temporary vectors.\n this->m_buffer_N = static_cast( this->m_A->get_num_cols() * this->m_A->get_block_dimy() );\n Matrix *mtx_A = dynamic_cast*>(this->m_A);\n VVector eig_solver_t_x;\n\n // m_lambda_mode:\n // 0: use eigensolver to get lmin and lmax estimate\n // 1: use eigensolver to get lmax estimate, set lmin = lmax/8\n // 2: use max row sum as lmax estimate, set lmin = lmax/8\n if (m_lambda_mode < 2)\n {\n if (!no_preconditioner)\n {\n SolverOperator *MA = new SolverOperator (this->m_A, m_preconditioner);\n m_eigsolver->setup(*MA);\n m_eigsolver->solve(eig_solver_t_x);\n delete MA;\n }\n else\n {\n m_eigsolver->setup(*mtx_A);\n m_eigsolver->solve(eig_solver_t_x);\n }\n\n const std::vector &lambdas = m_eigsolver->get_eigenvalues();\n this->m_lmax = lambdas[0] * 1.05;\n\n if (m_lambda_mode == 0)\n {\n this->m_lmin = lambdas[lambdas.size() - 1] * 0.95;\n }\n else\n {\n this->m_lmin = this->m_lmax * 0.125;\n }\n }\n else if (m_lambda_mode == 2)\n {\n if (no_preconditioner)\n {\n Matrix *pA = dynamic_cast< Matrix* > (this->m_A);\n compute_eigenmax_estimate(*pA, this->m_lmax);\n this->m_lmin = this->m_lmax * 0.125;\n }\n else\n {\n // assuming that this preconditioner would be good enough to reduce spectrum to the largest eigen value = 1.0\n this->m_lmax = 0.9;\n this->m_lmin = this->m_lmax * 0.125;\n }\n }\n else if (m_lambda_mode == 3)\n {\n if (no_preconditioner)\n {\n Matrix *pA = dynamic_cast< Matrix* > (this->m_A);\n compute_eigenmax_estimate(*pA, this->m_lmax);\n this->m_lmin = this->m_lmax * 0.125;\n }\n else\n {\n // Use user input estimates\n this->m_lmax = this->m_user_max_lambda;\n this->m_lmin = this->m_user_min_lambda;\n }\n }\n else\n {\n FatalError(\"Not supported chebyshev_lambda_estimate_mode.\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n // Allocate memory needed for iterating.\n m_p.resize( this->m_buffer_N );\n m_z.resize( this->m_buffer_N );\n m_Ap.resize( this->m_buffer_N );\n m_xp.resize( this->m_buffer_N );\n m_rp.resize( this->m_buffer_N );\n m_p.set_block_dimy(this->m_A->get_block_dimy());\n m_p.set_block_dimx(1);\n m_p.dirtybit = 1;\n m_p.delayed_send = 1;\n m_p.tag = this->tag * 100 + 1;\n m_Ap.set_block_dimy(this->m_A->get_block_dimy());\n m_Ap.set_block_dimx(1);\n m_Ap.dirtybit = 1;\n m_Ap.delayed_send = 1;\n m_Ap.tag = this->tag * 100 + 2;\n m_z.set_block_dimy(this->m_A->get_block_dimy());\n m_z.set_block_dimx(1);\n m_z.dirtybit = 1;\n m_z.delayed_send = 1;\n m_z.tag = this->tag * 100 + 3;\n m_xp.set_block_dimy(this->m_A->get_block_dimy());\n m_xp.set_block_dimx(1);\n m_xp.dirtybit = 1;\n m_xp.delayed_send = 1;\n m_xp.tag = this->tag * 100 + 4;\n this->m_A->setView(oldView);\n}\n\ntemplate\nvoid\nChebyshev_Solver::solve_init( VVector &b, VVector &x, bool xIsZero )\n{\n AMGX_CPU_PROFILER( \"Chebyshev_Solver::solve_init \" );\n Operator &A = *this->m_A;\n ViewType oldView = A.currentView();\n A.setViewExterior();\n int offset, size;\n A.getOffsetAndSizeForView(A.getViewExterior(), &offset, &size);\n\n // Run one iteration of preconditioner with zero initial guess\n if (no_preconditioner)\n {\n copy(*this->m_r, m_z, offset, size);\n }\n else\n {\n m_z.delayed_send = 1;\n this->m_r->delayed_send = 1;\n m_preconditioner->solve( *this->m_r, m_z, true );\n m_z.delayed_send = 1;\n this->m_r->delayed_send = 1;\n }\n\n // m_p - res after precond\n copy( m_z, m_p, offset, size );\n A.setView(oldView);\n m_gamma = 0.;\n m_beta = 0.;\n first_iter = 0;\n}\n\ntemplate\nAMGX_STATUS\nChebyshev_Solver::solve_iteration( VVector &b, VVector &x, bool xIsZero )\n{\n AMGX_STATUS conv_stat = AMGX_ST_NOT_CONVERGED;\n\n AMGX_CPU_PROFILER( \"Chebyshev_Solver::solve_iteration \" );\n Operator &A = *this->m_A;\n ViewType oldView = A.currentView();\n A.setViewExterior();\n int offset, size;\n A.getOffsetAndSizeForView(A.getViewExterior(), &offset, &size);\n ValueTypeB a = (this->m_lmax + this->m_lmin) / 2;\n ValueTypeB c = (this->m_lmax - this->m_lmin) / 2;\n\n for (int i = 0; i < m_cheby_order; i++)\n {\n // apply precond\n if (no_preconditioner)\n {\n copy(*this->m_r, m_z, offset, size);\n }\n else\n {\n m_z.delayed_send = 1;\n this->m_r->delayed_send = 1;\n m_preconditioner->solve( *this->m_r, m_z, true );\n m_z.delayed_send = 1;\n this->m_r->delayed_send = 1;\n }\n\n if (first_iter == 0)\n {\n m_gamma = 1. / a;\n first_iter = 1;\n }\n else\n {\n m_beta = c * c * m_gamma * m_gamma / 4.;\n\n if (m_gamma != ValueTypeB(0) && (a - (m_beta / m_gamma)) != ValueTypeB(0))\n {\n m_gamma = 1. / (a - m_beta / m_gamma);\n }\n\n axpby( m_z, m_p, m_p, ValueTypeB( 1 ), m_beta, offset, size);\n }\n\n axpy( m_p, x, m_gamma, offset, size );\n this->compute_residual( b, x);\n }\n\n // Do we converge ?\n if ( this->m_monitor_convergence &&\n isDone( ( conv_stat = this->compute_norm_and_converged() ) ) )\n {\n A.setView(oldView);\n return conv_stat;\n }\n\n // No convergence so far.\n A.setView(oldView);\n return this->m_monitor_convergence ? AMGX_ST_NOT_CONVERGED : AMGX_ST_CONVERGED;\n}\n\ntemplate\nvoid\nChebyshev_Solver::solve_finalize( VVector &b, VVector &x )\n{}\n\ntemplate\nvoid\nChebyshev_Solver::printSolverParameters() const\n{\n if (!no_preconditioner)\n {\n std::cout << \"preconditioner: \" << this->m_preconditioner->getName()\n << \" with scope name: \"\n << this->m_preconditioner->getScope() << std::endl;\n }\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class Chebyshev_Solver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/solvers/cheb_solver.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"solvers/dense_lu_solver.h\"\n#include \"solvers/block_common_solver.h\"\n\n#include \"amgx_types/util.h\"\n\n#include \n\n//trick to get nvcc to print a defined value\n//(__CUDA_ARCH__) @ compile time:\n//\nnamespace //unnamed namespace\n{\nstruct X\n{\n private:\n X(void) {}\n};\n\ntemplate\nstruct Print_N_atCompile\n{\n X array[N];\n};\n}//end unnamed namespace\n\nnamespace amgx\n{\nnamespace dense_lu_solver\n{\n\nenum { WARP_SIZE = 32, CTA_SIZE = 128 };\n\n//\n// supporting kernels\n//\ntemplate< typename T, typename Tv, int WARP_SIZE >\n__global__\nvoid csr_to_dense_kernel(\n const int num_rows,\n const int num_cols,\n const int block_num_rows,\n const int block_num_cols,\n const int *__restrict A_csr_rows,\n const int *__restrict A_csr_cols,\n const T *__restrict A_csr_vals,\n const int *__restrict A_csr_diag,\n T *__restrict A_dense,\n const int lda)\n{\n // Note:\n // To handle block CSR, the sparsity pattern csr_rows and csr_cols only store\n // the typical csr info assuming each block is a scalar.\n // The values in csr_vals has all entries in the blocks, using row major to\n // store the block. So we need the number of entries in each block as stride.\n const int block_mxn = block_num_rows * block_num_cols;\n // Each lane copies one entry in a block and iterate through row sparsity pattern.\n // Essentially one warp per row-block. For 4x4, we have 16 working threads per warp.\n const int lane_id = threadIdx.x % WARP_SIZE;\n // find the (row,col) local to a block\n const int block_row = lane_id / block_num_cols;\n const int block_col = lane_id % block_num_cols;\n\n // These are wasted threads per warp\n if ( block_row >= block_num_rows ) { return; }\n\n // The first row to consider. One row per warp.\n int row = (blockIdx.x * blockDim.x + threadIdx.x) / WARP_SIZE;\n const int row_offset = blockDim.x * gridDim.x / WARP_SIZE;\n\n for ( ; row < num_rows ; row += row_offset )\n {\n int dense_row = row * block_num_rows + block_row;\n // Iterate over each row and copy the elements into col-major A_dense\n int row_end = A_csr_rows[row + 1];\n\n for (int row_it = A_csr_rows[row]; row_it < row_end ; ++row_it )\n {\n int col = A_csr_cols[row_it];\n\n if ( col >= num_rows ) { continue; } // Skip entries corresponding to halo\n\n int dense_col = col * block_num_cols + block_col;\n A_dense[dense_col * lda + dense_row] = A_csr_vals[block_mxn * row_it + lane_id];\n }\n\n // copy diagonal block\n if ( A_csr_diag )\n {\n int diag_it = A_csr_diag[row];\n int dense_col = row * block_num_cols + block_col; // diag means row=col\n A_dense[dense_col * lda + dense_row] = A_csr_vals[block_mxn * diag_it + lane_id];\n }\n }\n}\n\n\ntemplate< int N, bool ROW_MAJOR, int WARP_SIZE, typename Value_type >\nstatic __device__ __forceinline__\nValue_type reduce_distributed_vectors( Value_type x, int is_leader, unsigned int active_mask )\n{\n if ( N & (N - 1) )\n {\n#pragma unroll\n\n for ( int i = 1 ; i < N ; ++i )\n {\n Value_type other_x = utils::shfl_down( x, ROW_MAJOR ? i : N * i, WARP_SIZE, active_mask );\n\n if ( is_leader )\n {\n x += other_x;\n }\n }\n }\n else\n {\n#pragma unroll\n\n for ( int i = 1 ; i < N ; i <<= 1 )\n {\n x += utils::shfl_xor( x, ROW_MAJOR ? i : N * i, WARP_SIZE, active_mask );\n }\n }\n\n return x;\n}\n\n\n\ntemplate< typename Matrix_type, typename Vector_type, int N, int CTA_SIZE,\n int WARP_SIZE, bool ROW_MAJOR, bool HAS_EXTERNAL_DIAG >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid b_minus_A_halo_x( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *new_rhs,\n const int num_owned_rows)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int NUM_ITEMS_PER_WARP = WARP_SIZE / NxN;\n // Number of items computer per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_ITEMS_PER_WARP * NUM_WARPS_PER_CTA;\n // Number of items per grid.\n const int NUM_ITEMS_PER_GRID = gridDim.x * NUM_ITEMS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_NxN = lane_id / NxN;\n const int lane_id_mod_NxN = lane_id % NxN;\n // Useful index to compute matrix products.\n const int lane_id_mod_NxN_div_N = lane_id_mod_NxN / N;\n const int lane_id_mod_NxN_mod_N = lane_id_mod_NxN % N;\n // We to get my data from when I use SHFL.\n const int shfl_offset = lane_id - lane_id_mod_NxN;\n // Shared memory needed to exchange X and delta.\n __shared__ volatile Vector_type s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_mem = &s_mem[threadIdx.x - lane_id_mod_NxN];\n // Is the thread active? For example, for 5x5 only the first 25 threads are active per warp.\n // At compile time, the compiler will see is_active == true for 2x2 (since NxN & (NxN-1) evals\n // to false ; that's the common trick to determine if a number is a power of 2).\n int is_active = true;\n\n if ( NxN & (NxN - 1) )\n {\n is_active = lane_id_div_NxN < NUM_ITEMS_PER_WARP;\n }\n\n // Determine which NxN block the threads work with.\n int a_row_it = num_owned_rows;\n\n if ( is_active )\n {\n a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + warp_id * NUM_ITEMS_PER_WARP + lane_id_div_NxN;\n }\n\n unsigned int active_mask = utils::ballot(is_active);\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_owned_rows; a_row_it += NUM_ITEMS_PER_GRID )\n {\n int a_row_id = a_row_it;\n // Load one block of B.\n Vector_type my_bmAx(0);\n\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n my_bmAx = __cachingLoad(&b[N * a_row_id + lane_id_mod_NxN_div_N]);\n }\n }\n else\n {\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n my_bmAx = b[N * a_row_id + lane_id_mod_NxN_mod_N];\n }\n }\n\n // Don't do anything if X is zero.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // If the diagonal is stored separately, we have a special treatment.\n int a_col_max = a_col_end;\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( a_col_begin < a_col_max ) ; a_col_begin += NxN )\n {\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id_mod_NxN;\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Determine if the column is halo column\n int a_col_is_valid = (a_col_id != -1) && (a_col_id >= num_owned_rows);\n // Count the number of active columns.\n // int vote = __ballot(aColId != -1);\n // The number of iterations.\n // int nCols = max( __popc( vote & 0x0000ffff ), __popc( vote & 0xffff0000 ) );\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < NxN ; k += N )\n {\n int my_k = k + lane_id_mod_NxN_div_N;\n // Load N blocks of X.\n int uniform_a_col_id = utils::shfl( a_col_id, shfl_offset + my_k, WARP_SIZE, active_mask);\n int uniform_a_col_is_valid = utils::shfl( a_col_is_valid, shfl_offset + my_k, WARP_SIZE, active_mask );\n Vector_type my_x(0);\n\n if ( uniform_a_col_id != -1 && uniform_a_col_is_valid)\n {\n my_x = __cachingLoad(&x[N * uniform_a_col_id + lane_id_mod_NxN_mod_N]);\n //printf(\"loading entry %d, num_rows = %d, my_x = %f\\n\",uniform_a_col_id,num_owned_rows,my_x);\n }\n\n my_s_mem[lane_id_mod_NxN] = my_x;\n // Load N blocks of A.\n#pragma unroll\n\n for ( int i = 0 ; i < N ; ++i )\n {\n int uniform_a_col_tmp = a_col_begin + k + i, uniform_a_col_it = -1;\n\n if ( uniform_a_col_tmp < a_col_end )\n {\n uniform_a_col_it = uniform_a_col_tmp;\n }\n\n Matrix_type my_val(0);\n\n if ( uniform_a_col_it != -1 )\n {\n my_val = A_vals[NxN * uniform_a_col_it + lane_id_mod_NxN];\n }\n\n if ( ROW_MAJOR )\n {\n my_bmAx -= my_val * my_s_mem[N * i + lane_id_mod_NxN_mod_N];\n }\n else\n {\n my_bmAx -= my_val * my_s_mem[N * i + lane_id_mod_NxN_div_N];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n Vector_type my_Einv = (lane_id_mod_NxN == 0 || lane_id_mod_NxN == 5 || lane_id_mod_NxN == 10 || lane_id_mod_NxN == 15) ? 1. : 0.;\n // Reduce bmAx terms.\n int is_leader = lane_id_mod_NxN_div_N == 0;\n\n if ( ROW_MAJOR )\n {\n is_leader = lane_id_mod_NxN_mod_N == 0;\n }\n\n my_bmAx = reduce_distributed_vectors( my_bmAx, is_leader, active_mask );\n\n // Update the shared terms.\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n my_s_mem[lane_id_mod_NxN_div_N] = my_bmAx;\n }\n }\n else\n {\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n my_s_mem[lane_id_mod_NxN_mod_N] = my_bmAx;\n }\n }\n\n // Update the diagonal term.\n if ( ROW_MAJOR )\n {\n my_bmAx = my_Einv * my_s_mem[lane_id_mod_NxN_mod_N];\n }\n else\n {\n my_bmAx = my_Einv * my_s_mem[lane_id_mod_NxN_div_N];\n }\n\n // Reduce bmAx terms.\n my_bmAx = reduce_distributed_vectors( my_bmAx, is_leader, active_mask );\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n new_rhs[N * a_row_id + lane_id_mod_NxN_div_N] = my_bmAx;\n }\n }\n else\n {\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n new_rhs[N * a_row_id + lane_id_mod_NxN_mod_N] = my_bmAx;\n }\n }\n }\n}\n\n\n////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n\n\ntemplate< typename Vector_type, typename Matrix_type, int N>\nvoid distributed_rhs_mod_dispatch( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *new_rhs,\n const int num_owned_rows,\n const int row_major,\n const int has_external_diag\n )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int NUM_ROWS_PER_WARP = WARP_SIZE / NxN;\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = NUM_ROWS_PER_WARP * NUM_WARPS_PER_CTA;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_owned_rows + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n // Branch to the correct kernel call.\n int code = 2 * (row_major ? 1 : 0) + (has_external_diag ? 1 : 0);\n\n switch ( code )\n {\n case 0: // Column-major, no external diagonal.\n b_minus_A_halo_x <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n new_rhs,\n num_owned_rows\n );\n cudaCheckError();\n break;\n\n case 1: // Column-major, external diagonal.\n b_minus_A_halo_x <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n new_rhs,\n num_owned_rows\n );\n cudaCheckError();\n break;\n\n case 2: // Row-major, no external diagonal.\n b_minus_A_halo_x <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n new_rhs,\n num_owned_rows\n );\n cudaCheckError();\n break;\n\n case 3: // Row-major, external diagonal.\n b_minus_A_halo_x <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n new_rhs,\n num_owned_rows\n );\n cudaCheckError();\n break;\n\n default:\n FatalError( \"Internal error\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n cudaCheckError();\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I>\nvoid DenseLUSolver >::\ndistributed_rhs_mod(\n const Vector_d &x,\n const Vector_d &b,\n Vector_d &new_rhs)\n{\n const Matrix_d *A = dynamic_cast(Base::m_A);\n\n switch (A->get_block_dimy())\n {\n case 1:\n distributed_rhs_mod_dispatch(A->row_offsets.raw(), A->col_indices.raw(), A->values.raw(), A->diag.raw(), x.raw(), b.raw(), new_rhs.raw(), A->get_num_rows(), A->getBlockFormat() == ROW_MAJOR, A->hasProps(DIAG));\n break;\n\n case 2:\n distributed_rhs_mod_dispatch(A->row_offsets.raw(), A->col_indices.raw(), A->values.raw(), A->diag.raw(), x.raw(), b.raw(), new_rhs.raw(), A->get_num_rows(), A->getBlockFormat() == ROW_MAJOR, A->hasProps(DIAG));\n break;\n\n case 3:\n distributed_rhs_mod_dispatch(A->row_offsets.raw(), A->col_indices.raw(), A->values.raw(), A->diag.raw(), x.raw(), b.raw(), new_rhs.raw(), A->get_num_rows(), A->getBlockFormat() == ROW_MAJOR, A->hasProps(DIAG));\n break;\n\n case 4:\n distributed_rhs_mod_dispatch(A->row_offsets.raw(), A->col_indices.raw(), A->values.raw(), A->diag.raw(), x.raw(), b.raw(), new_rhs.raw(), A->get_num_rows(), A->getBlockFormat() == ROW_MAJOR, A->hasProps(DIAG));\n break;\n\n case 5:\n distributed_rhs_mod_dispatch(A->row_offsets.raw(), A->col_indices.raw(), A->values.raw(), A->diag.raw(), x.raw(), b.raw(), new_rhs.raw(), A->get_num_rows(), A->getBlockFormat() == ROW_MAJOR, A->hasProps(DIAG));\n break;\n }\n}\n\n// copy non zero elements only\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid DenseLUSolver >::\ncsr_to_dense(\n)\n{\n const Matrix_d *A = dynamic_cast(Base::m_A);\n const int block_size = 256;\n const int num_warps = block_size / WARP_SIZE;\n const int grid_size = (A->get_num_rows() + num_warps - 1) / num_warps;\n cudaStream_t stream = amgx::thrust::global_thread_handle::get_stream();\n csr_to_dense_kernel <<< grid_size, block_size, 0, stream>>>(\n A->get_num_rows(),\n A->get_num_cols(),\n A->get_block_dimx(),\n A->get_block_dimy(),\n A->row_offsets.raw(),\n A->col_indices.raw(),\n A->values.raw(),\n A->hasProps(DIAG) ? A->diag.raw() : NULL,\n m_dense_A,\n m_lda\n );\n cudaCheckError();\n}\n\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid DenseLUSolver >::cudense_getrf()\n{\n int wsize = 0;\n cusolverStatus_t status1 = cusolverDnXgetrf_bufferSize(m_cuds_handle,\n m_num_rows,\n m_num_cols,\n m_dense_A,\n m_lda,\n &wsize);\n\n if ( status1 != CUSOLVER_STATUS_SUCCESS)\n {\n FatalError( \"Failed kernel in DenseLU\", AMGX_ERR_INTERNAL);\n }\n\n cudaCheckError();\n allocMem(m_trf_wspace, wsize, false);\n status1 = cusolverDnXgetrf(m_cuds_handle, m_num_rows, m_num_cols,\n m_dense_A, m_lda, m_trf_wspace, m_ipiv, m_cuds_info);\n cudaCheckError();\n\n if ( status1 != CUSOLVER_STATUS_SUCCESS)\n {\n FatalError( \"failed kernel in Dense LU is likely due to invalid input parameters\",\n AMGX_ERR_INTERNAL);\n }\n else\n {\n int t_info;\n cudaMemcpy(&t_info, m_cuds_info, sizeof(int), cudaMemcpyDefault);\n cudaCheckError();\n\n if (t_info != 0)\n {\n FatalError( \"Fail to get info from cudense\", AMGX_ERR_INTERNAL);\n }\n else\n {\n // We follow the standard established by Lapack and used in cudense.\n if (t_info > 0)\n {\n FatalError( \"Dense LU factorization failed due to a singular matrix\", AMGX_ERR_INTERNAL);\n }\n else if (t_info < 0)\n {\n FatalError( \"Invalid input parameter(s) to dense LU\", AMGX_ERR_INTERNAL);\n }\n }\n }\n\n cudaCheckError();\n\n if (m_trf_wspace)\n {\n amgx::memory::cudaFreeAsync(m_trf_wspace);\n m_trf_wspace = 0;\n }\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid DenseLUSolver >::cudense_getrs( Vector_d &x )\n{\n //Solve L*X = RHS\n cusolverStatus_t status = cusolverDnXgetrs(m_cuds_handle,\n CUBLAS_OP_N,\n m_num_rows,\n 1,\n m_dense_A,\n m_lda,\n m_ipiv,\n (Matrix_data *)(x.raw()),\n m_num_rows,\n m_cuds_info);\n\n if (status != CUSOLVER_STATUS_SUCCESS)\n {\n FatalError( \"cuSolver trsv failed to solve Lx=rhs\", AMGX_ERR_INTERNAL);\n }\n\n cudaCheckError();\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\ntemplate< class DataType, class IndexType >\nvoid\nDenseLUSolver >::\nallocMem(DataType *&ptr,\n IndexType numEntry,\n bool initToZero)\n{\n if ( ptr != NULL ) { amgx::memory::cudaFreeAsync(ptr); }\n cudaCheckError();\n\n size_t sz = numEntry * sizeof(DataType);\n amgx::memory::cudaMallocAsync((void **)&ptr, sz);\n cudaCheckError();\n\n if (initToZero)\n {\n cudaMemset(ptr, 0x0, sz);\n cudaCheckError();\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nDenseLUSolver >::\nDenseLUSolver(AMG_Config &cfg,\n const std::string &cfg_scope,\n ThreadManager *tmng)\n : Solver(cfg, cfg_scope, tmng),\n m_cuds_handle(0),\n m_num_rows(0),\n m_num_cols(0),\n m_lda(0),\n m_dense_A(0),\n m_ipiv(0),\n m_cuds_info(0),\n m_trf_wspace(0)\n{\n // Allocate a handle for cudense\n cusolverStatus_t status = cusolverDnCreate(&m_cuds_handle);\n\n if ( status != CUSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not create the CUDENSE handle\", AMGX_ERR_CUDA_FAILURE );\n }\n\n // Allocate a handle for cublas\n cublasStatus_t cublasStatus = cublasCreate(&m_cublas_handle);\n\n if ( cublasStatus != CUBLAS_STATUS_SUCCESS )\n {\n FatalError( \"Could not create the CUBLAS handle\", AMGX_ERR_CUDA_FAILURE );\n }\n\n // Define the cudense stream.\n status = cusolverDnSetStream(m_cuds_handle, amgx::thrust::global_thread_handle::get_stream());\n\n if ( status != CUSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not set the stream for CUDENSE\", AMGX_ERR_CUDA_FAILURE );\n }\n\n // Make sure we don't run more than 1 iteration.\n this->set_max_iters(1);\n allocMem(m_cuds_info, sizeof(int), false);\n\n // Determine if the scalable coarse solve optimisation is enabled\n m_enable_exact_solve = (\n cfg.getParameter( \"algorithm\", cfg_scope) == CLASSICAL &&\n cfg.getParameter( \"exact_coarse_solve\", cfg_scope) == 1);\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nDenseLUSolver >::~DenseLUSolver()\n{\n if (m_cuds_handle) { cusolverDnDestroy(m_cuds_handle); }\n\n if (m_cublas_handle) { cublasDestroy(m_cublas_handle); }\n\n if (m_dense_A)\n {\n amgx::memory::cudaFreeAsync(m_dense_A);\n }\n\n if (m_ipiv)\n {\n amgx::memory::cudaFreeAsync(m_ipiv);\n }\n\n if (m_trf_wspace)\n {\n amgx::memory::cudaFreeAsync(m_trf_wspace);\n }\n\n if (m_cuds_info)\n {\n amgx::memory::cudaFreeAsync(m_cuds_info);\n }\n\n cudaCheckError();\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid DenseLUSolver >::solve_init(Vector_d &, Vector_d &, bool)\n{}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid DenseLUSolver >::solve_finalize(Vector_d &, Vector_d &)\n{}\n\n// Offset the local row offsets to global row offsets\ntemplate \n__global__ void local_row_offsets_to_global(\n int num_rows, int offset,\n IndexType* local_Arows)\n{\n int i = threadIdx.x + blockIdx.x*blockDim.x;\n if(i >= num_rows) { return; }\n\n local_Arows[i] += offset;\n}\n\n// Offset local packed column indices to global unpacked indices\ntemplate \n__global__ void local_col_indices_to_global(\n int nnz, int num_rows, int offset, IndexType* local_Acols, L2GType* l2g)\n{\n int i = threadIdx.x + blockIdx.x*blockDim.x;\n if(i >= nnz) { return; }\n\n if(local_Acols[i] >= num_rows)\n {\n local_Acols[i] = l2g[local_Acols[i] - num_rows];\n }\n else\n {\n local_Acols[i] += offset;\n }\n}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid DenseLUSolver >::\nsolver_setup(bool reuse_matrix_structure)\n{\n Matrix_d *A = dynamic_cast(Base::m_A);\n\n if (!A)\n {\n FatalError(\"DenseLUSolver only works with explicit matrices\",\n AMGX_ERR_INTERNAL);\n }\n\n ViewType oldView = A->currentView();\n A->setViewExterior();\n\n if(A->is_matrix_distributed() && m_enable_exact_solve)\n {\n#ifdef AMGX_WITH_MPI\n int rank = A->manager->global_id();\n int nranks = A->manager->get_num_partitions();\n MPI_Comm comm = A->manager->getComms()->get_mpi_comm();\n\n int offset, num_rows, nnz;\n A->getOffsetAndSizeForView(OWNED, &offset, &num_rows);\n A->getNnzForView(OWNED, &nnz);\n\n m_num_rows = A->manager->num_rows_global * A->get_block_dimx();\n m_num_cols = A->manager->num_rows_global * A->get_block_dimy();\n m_lda = m_num_rows; // col-major\n\n // Allocate mem for cudense pivoting sequence.\n allocMem(m_ipiv, m_num_rows, false);\n\n // Allocate memory to store the dense A and initialize to zero.\n allocMem(m_dense_A, static_cast(m_num_cols) * m_num_rows, true);\n\n // Much of the data can be reused if we are performing a resetup\n if (!reuse_matrix_structure)\n {\n // Gather the number of non zeros on each rank\n A->manager->getComms()->all_gather(nnz, nz_all, nranks);\n\n // Gather the number of rows on each rank\n A->manager->getComms()->all_gather(num_rows, row_all, nranks);\n\n // Get the number of non zeros on all ranks\n m_nnz_global = thrust_wrapper::reduce(nz_all.begin(), nz_all.end());\n\n // Turn the non-zero counts into displacements\n nz_displs.resize(nranks);\n thrust_wrapper::exclusive_scan(nz_all.begin(), nz_all.end(), nz_displs.begin());\n\n // Turn the number of rows into displacements\n row_displs.resize(nranks);\n thrust_wrapper::exclusive_scan(row_all.begin(), row_all.end(), row_displs.begin());\n\n IVector_d local_Acols_d(nnz);\n IVector_d local_Arows_d(num_rows);\n\n amgx::thrust::copy(A->col_indices.begin(), A->col_indices.begin() + nnz, local_Acols_d.begin());\n amgx::thrust::copy(A->row_offsets.begin(), A->row_offsets.begin() + num_rows, local_Arows_d.begin());\n\n // XXX Local to global map is the current limiting factor to enabling this\n // code for the aggregation based path. It's not clear whether there is\n // a structure that provides the same inverse mapping with aggregation.\n // Note that at one point inverse_renumbering was tested for aggregation\n // but didn't appear to work in all cases.\n\n // Convert the local column indices and row offsets to the global index space\n constexpr int nthreads = 128;\n int nblocks = nnz / nthreads + 1;\n local_col_indices_to_global<<>>(nnz, num_rows, row_displs[rank], local_Acols_d.raw(), A->manager->local_to_global_map.raw());\n cudaCheckError();\n\n nblocks = num_rows / nthreads + 1;\n local_row_offsets_to_global<<>>(num_rows, nz_displs[rank], local_Arows_d.raw());\n cudaCheckError();\n\n // Copy the transformed indices to the host\n IVector_h local_Acols_h(nnz);\n IVector_h local_Arows_h(num_rows);\n amgx::thrust::copy(local_Acols_d.begin(), local_Acols_d.end(), local_Acols_h.begin());\n amgx::thrust::copy(local_Arows_d.begin(), local_Arows_d.end(), local_Arows_h.begin());\n\n // Gather the local matrix structure redundantly to every rank\n IVector_h Acols_global_h(m_nnz_global);\n A->manager->getComms()->all_gather_v(local_Acols_h, nnz, Acols_global_h, nz_all, nz_displs);\n\n // Note: Copy the local data to global without guard value\n IVector_h Arows_global_h(m_num_rows + 1);\n A->manager->getComms()->all_gather_v(local_Arows_h, num_rows, Arows_global_h, row_all, row_displs);\n\n // Manually set the guard value on the global matrix\n Arows_global_h[m_num_rows] = m_nnz_global;\n Acols_global.resize(m_nnz_global);\n Arows_global.resize(m_num_rows + 1);\n amgx::thrust::copy(Acols_global_h.begin(), Acols_global_h.end(), Acols_global.begin());\n amgx::thrust::copy(Arows_global_h.begin(), Arows_global_h.end(), Arows_global.begin());\n }\n\n // Fetch to the host a copy of the local sparse matrix\n MVector_h local_Avals_h(nnz);\n amgx::thrust::copy(A->values.begin(), A->values.begin() + nnz, local_Avals_h.begin());\n\n // Gather the matrix values to all ranks\n MVector_h Avals_global_h(m_nnz_global);\n A->manager->getComms()->all_gather_v(local_Avals_h, nnz, Avals_global_h, nz_all, nz_displs);\n\n allocMem(m_dense_A, static_cast(m_num_cols) * m_lda, true);\n\n MVector_d Avals_global(m_nnz_global);\n amgx::thrust::copy(Avals_global_h.begin(), Avals_global_h.end(), Avals_global.begin());\n\n const int block_size = 256;\n const int num_warps = block_size / WARP_SIZE;\n const int grid_size = (m_nnz_global + num_warps - 1) / num_warps;\n cudaStream_t stream = amgx::thrust::global_thread_handle::get_stream();\n csr_to_dense_kernel<<>>(\n m_num_rows,\n m_num_cols,\n A->get_block_dimx(),\n A->get_block_dimy(),\n Arows_global.raw(),\n Acols_global.raw(),\n Avals_global.raw(),\n A->hasProps(DIAG) ? A->diag.raw() : NULL,\n m_dense_A,\n m_lda);\n cudaCheckError();\n\n cudaStreamSynchronize(stream);\n cudaCheckError();\n#endif\n }\n else\n {\n ViewType oldView = A->currentView();\n A->setViewExterior();\n m_num_rows = A->get_num_rows() * A->get_block_dimx();\n // don't use A->get_num_cols() because A is rectangular.\n // Only the diagonal block owned by this rank is factored.\n m_num_cols = A->get_num_rows() * A->get_block_dimy();\n m_lda = m_num_rows; // col-major\n // Allocate mem for cudense pivoting sequence.\n allocMem(m_ipiv, m_num_rows, false);\n // Allocate memory to store the dense A and initialize to zero.\n allocMem(m_dense_A, static_cast(m_num_cols) * m_lda, true);\n csr_to_dense(); // copy sparse A to dense_A\n }\n\n cudense_getrf(); // do LU factor\n A->setView(oldView);\n}\n\n// There is one subtle point here (for inexact solve):\n// We only do LU on the diagonal blocks associated with each rank.\n// Halo is used to update the right-hand-side (RHS) vector.\n// For multi GPU cases, this is essentially block Jacobi. Since the block size\n// is the size of each partition, this is OK trade-off between accuracy and runtime.\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nAMGX_STATUS\nDenseLUSolver >::\nsolve_iteration(Vector_d &rhs,\n Vector_d &x,\n bool xIsZero)\n{\n Matrix_d *A = dynamic_cast(Base::m_A);\n ViewType oldView = A->currentView();\n A->setViewExterior();\n\n if(A->is_matrix_distributed() && m_enable_exact_solve)\n {\n#ifdef AMGX_WITH_MPI\n int offset, num_rows;\n A->getOffsetAndSizeForView(OWNED, &offset, &num_rows);\n\n int rank = A->manager->global_id();\n int nranks = A->manager->get_num_partitions();\n MPI_Comm comm = A->manager->getComms()->get_mpi_comm();\n\n // Make host copy of the RHS\n MVector_h rhs_local_h(num_rows);\n amgx::thrust::copy(rhs.begin(), rhs.begin() + num_rows, rhs_local_h.begin());\n\n // Gather the local RHS from all ranks to global vectors on all ranks\n MVector_h rhs_global_h(m_num_rows);\n A->manager->getComms()->all_gather_v(rhs_local_h, num_rows, rhs_global_h, row_all, row_displs);\n\n //Solve L*X = RHS\n MVector_d x_global(m_num_rows);\n amgx::thrust::copy(rhs_global_h.begin(), rhs_global_h.end(), x_global.begin());\n cusolverStatus_t status = \n cusolverDnXgetrs(m_cuds_handle,\n CUBLAS_OP_N,\n m_num_rows,\n 1,\n m_dense_A,\n m_lda,\n m_ipiv,\n x_global.raw(),\n m_num_rows,\n m_cuds_info);\n\n // Copy the local portion of the solution back into x\n amgx::thrust::copy(x_global.begin() + row_displs[rank], x_global.begin() + row_displs[rank] + num_rows, x.begin());\n\n if (status != CUSOLVER_STATUS_SUCCESS)\n {\n FatalError(\"cuSolver trsv failed to solve Lx=rhs\", AMGX_ERR_INTERNAL);\n }\n#endif\n }\n else\n {\n Matrix_d *A = dynamic_cast(Base::m_A);\n ViewType oldView = A->currentView();\n A->setViewExterior();\n\n if ((!A->is_matrix_singleGPU()) && (!xIsZero))\n {\n // Modify rhs to include contribution from halo nodes\n // i.e. new_rhs = b - A_halo*x;\n // Note: dense_lu solver doesn't support latency hiding\n A->manager->exchange_halo_async(x, x.tag);\n A->manager->exchange_halo_wait(x, x.tag);\n Vector_d new_rhs(rhs.size());\n distributed_rhs_mod(x, rhs, new_rhs);\n amgx::thrust::copy(new_rhs.begin(), new_rhs.begin() + m_num_rows, x.begin());\n cudaCheckError();\n }\n else\n {\n x.copy(rhs);\n }\n\n cudense_getrs(x); // triangular solves\n }\n\n //Speculative send of x vector\n x.dirtybit = 1;\n A->setView(oldView);\n return AMGX_ST_CONVERGED; // direct solver always converges\n}\n\n\n#define AMGX_CASE_LINE(CASE) template class DenseLUSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n// AMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace dense_lu\n} // namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"solvers/dense_lu_solver.h\"\n#include \"solvers/block_common_solver.h\"\n\n#include \"amgx_types/util.h\"\n\n#include \n\n//trick to get nvcc to print a defined value\n//(__CUDA_ARCH__) @ compile time:\n//\nnamespace //unnamed namespace\n{\nstruct X\n{\n private:\n X(void) {}\n};\n\ntemplate\nstruct Print_N_atCompile\n{\n X array[N];\n};\n}//end unnamed namespace\n\nnamespace amgx\n{\nnamespace dense_lu_solver\n{\n\nenum { WARP_SIZE = 32, CTA_SIZE = 128 };\n\n//\n// supporting kernels\n//\ntemplate< typename T, typename Tv, int WARP_SIZE >\n__global__\nvoid csr_to_dense_kernel(\n const int num_rows,\n const int num_cols,\n const int block_num_rows,\n const int block_num_cols,\n const int *__restrict A_csr_rows,\n const int *__restrict A_csr_cols,\n const T *__restrict A_csr_vals,\n const int *__restrict A_csr_diag,\n T *__restrict A_dense,\n const int lda)\n{\n // Note:\n // To handle block CSR, the sparsity pattern csr_rows and csr_cols only store\n // the typical csr info assuming each block is a scalar.\n // The values in csr_vals has all entries in the blocks, using row major to\n // store the block. So we need the number of entries in each block as stride.\n const int block_mxn = block_num_rows * block_num_cols;\n // Each lane copies one entry in a block and iterate through row sparsity pattern.\n // Essentially one warp per row-block. For 4x4, we have 16 working threads per warp.\n const int lane_id = threadIdx.x % WARP_SIZE;\n // find the (row,col) local to a block\n const int block_row = lane_id / block_num_cols;\n const int block_col = lane_id % block_num_cols;\n\n // These are wasted threads per warp\n if ( block_row >= block_num_rows ) { return; }\n\n // The first row to consider. One row per warp.\n int row = (blockIdx.x * blockDim.x + threadIdx.x) / WARP_SIZE;\n const int row_offset = blockDim.x * gridDim.x / WARP_SIZE;\n\n for ( ; row < num_rows ; row += row_offset )\n {\n int dense_row = row * block_num_rows + block_row;\n // Iterate over each row and copy the elements into col-major A_dense\n int row_end = A_csr_rows[row + 1];\n\n for (int row_it = A_csr_rows[row]; row_it < row_end ; ++row_it )\n {\n int col = A_csr_cols[row_it];\n\n if ( col >= num_rows ) { continue; } // Skip entries corresponding to halo\n\n int dense_col = col * block_num_cols + block_col;\n A_dense[dense_col * lda + dense_row] = A_csr_vals[block_mxn * row_it + lane_id];\n }\n\n // copy diagonal block\n if ( A_csr_diag )\n {\n int diag_it = A_csr_diag[row];\n int dense_col = row * block_num_cols + block_col; // diag means row=col\n A_dense[dense_col * lda + dense_row] = A_csr_vals[block_mxn * diag_it + lane_id];\n }\n }\n}\n\n\ntemplate< int N, bool ROW_MAJOR, int WARP_SIZE, typename Value_type >\nstatic __device__ __forceinline__\nValue_type reduce_distributed_vectors( Value_type x, int is_leader, unsigned int active_mask )\n{\n if ( N & (N - 1) )\n {\n#pragma unroll\n\n for ( int i = 1 ; i < N ; ++i )\n {\n Value_type other_x = utils::shfl_down( x, ROW_MAJOR ? i : N * i, WARP_SIZE, active_mask );\n\n if ( is_leader )\n {\n x += other_x;\n }\n }\n }\n else\n {\n#pragma unroll\n\n for ( int i = 1 ; i < N ; i <<= 1 )\n {\n x += utils::shfl_xor( x, ROW_MAJOR ? i : N * i, WARP_SIZE, active_mask );\n }\n }\n\n return x;\n}\n\n\n\ntemplate< typename Matrix_type, typename Vector_type, int N, int CTA_SIZE,\n int WARP_SIZE, bool ROW_MAJOR, bool HAS_EXTERNAL_DIAG >\n__global__\n__launch_bounds__( CTA_SIZE, 12 )\nvoid b_minus_A_halo_x( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *new_rhs,\n const int num_owned_rows)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int NUM_ITEMS_PER_WARP = WARP_SIZE / NxN;\n // Number of items computer per CTA.\n const int NUM_ITEMS_PER_CTA = NUM_ITEMS_PER_WARP * NUM_WARPS_PER_CTA;\n // Number of items per grid.\n const int NUM_ITEMS_PER_GRID = gridDim.x * NUM_ITEMS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_NxN = lane_id / NxN;\n const int lane_id_mod_NxN = lane_id % NxN;\n // Useful index to compute matrix products.\n const int lane_id_mod_NxN_div_N = lane_id_mod_NxN / N;\n const int lane_id_mod_NxN_mod_N = lane_id_mod_NxN % N;\n // We to get my data from when I use SHFL.\n const int shfl_offset = lane_id - lane_id_mod_NxN;\n // Shared memory needed to exchange X and delta.\n __shared__ volatile Vector_type s_mem[CTA_SIZE];\n // Each thread keeps its own pointer to shared memory to avoid some extra computations.\n volatile Vector_type *my_s_mem = &s_mem[threadIdx.x - lane_id_mod_NxN];\n // Is the thread active? For example, for 5x5 only the first 25 threads are active per warp.\n // At compile time, the compiler will see is_active == true for 2x2 (since NxN & (NxN-1) evals\n // to false ; that's the common trick to determine if a number is a power of 2).\n int is_active = true;\n\n if ( NxN & (NxN - 1) )\n {\n is_active = lane_id_div_NxN < NUM_ITEMS_PER_WARP;\n }\n\n // Determine which NxN block the threads work with.\n int a_row_it = num_owned_rows;\n\n if ( is_active )\n {\n a_row_it = blockIdx.x * NUM_ITEMS_PER_CTA + warp_id * NUM_ITEMS_PER_WARP + lane_id_div_NxN;\n }\n\n unsigned int active_mask = utils::ballot(is_active);\n\n // Iterate over the rows of the matrix. One warp per row.\n for ( ; a_row_it < num_owned_rows; a_row_it += NUM_ITEMS_PER_GRID )\n {\n int a_row_id = a_row_it;\n // Load one block of B.\n Vector_type my_bmAx(0);\n\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n my_bmAx = __cachingLoad(&b[N * a_row_id + lane_id_mod_NxN_div_N]);\n }\n }\n else\n {\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n my_bmAx = b[N * a_row_id + lane_id_mod_NxN_mod_N];\n }\n }\n\n // Don't do anything if X is zero.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // If the diagonal is stored separately, we have a special treatment.\n int a_col_max = a_col_end;\n\n // Each warp load column indices of 32 nonzero blocks\n for ( ; utils::any( a_col_begin < a_col_max ) ; a_col_begin += NxN )\n {\n // Each thread loads a single element. If !is_active, a_col_end == 0.\n int a_col_it = a_col_begin + lane_id_mod_NxN;\n // Get the ID of the column.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Determine if the column is halo column\n int a_col_is_valid = (a_col_id != -1) && (a_col_id >= num_owned_rows);\n // Count the number of active columns.\n // int vote = __ballot(aColId != -1);\n // The number of iterations.\n // int nCols = max( __popc( vote & 0x0000ffff ), __popc( vote & 0xffff0000 ) );\n\n // Loop over columns. We compute 8 columns per iteration.\n for ( int k = 0 ; k < NxN ; k += N )\n {\n int my_k = k + lane_id_mod_NxN_div_N;\n // Load N blocks of X.\n int uniform_a_col_id = utils::shfl( a_col_id, shfl_offset + my_k, WARP_SIZE, active_mask);\n int uniform_a_col_is_valid = utils::shfl( a_col_is_valid, shfl_offset + my_k, WARP_SIZE, active_mask );\n Vector_type my_x(0);\n\n if ( uniform_a_col_id != -1 && uniform_a_col_is_valid)\n {\n my_x = __cachingLoad(&x[N * uniform_a_col_id + lane_id_mod_NxN_mod_N]);\n //printf(\"loading entry %d, num_rows = %d, my_x = %f\\n\",uniform_a_col_id,num_owned_rows,my_x);\n }\n\n my_s_mem[lane_id_mod_NxN] = my_x;\n // Load N blocks of A.\n#pragma unroll\n\n for ( int i = 0 ; i < N ; ++i )\n {\n int uniform_a_col_tmp = a_col_begin + k + i, uniform_a_col_it = -1;\n\n if ( uniform_a_col_tmp < a_col_end )\n {\n uniform_a_col_it = uniform_a_col_tmp;\n }\n\n Matrix_type my_val(0);\n\n if ( uniform_a_col_it != -1 )\n {\n my_val = A_vals[NxN * uniform_a_col_it + lane_id_mod_NxN];\n }\n\n if ( ROW_MAJOR )\n {\n my_bmAx -= my_val * my_s_mem[N * i + lane_id_mod_NxN_mod_N];\n }\n else\n {\n my_bmAx -= my_val * my_s_mem[N * i + lane_id_mod_NxN_div_N];\n }\n }\n } // Loop over k\n } // Loop over aColIt\n\n Vector_type my_Einv = (lane_id_mod_NxN == 0 || lane_id_mod_NxN == 5 || lane_id_mod_NxN == 10 || lane_id_mod_NxN == 15) ? 1. : 0.;\n // Reduce bmAx terms.\n int is_leader = lane_id_mod_NxN_div_N == 0;\n\n if ( ROW_MAJOR )\n {\n is_leader = lane_id_mod_NxN_mod_N == 0;\n }\n\n my_bmAx = reduce_distributed_vectors( my_bmAx, is_leader, active_mask );\n\n // Update the shared terms.\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n my_s_mem[lane_id_mod_NxN_div_N] = my_bmAx;\n }\n }\n else\n {\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n my_s_mem[lane_id_mod_NxN_mod_N] = my_bmAx;\n }\n }\n\n // Update the diagonal term.\n if ( ROW_MAJOR )\n {\n my_bmAx = my_Einv * my_s_mem[lane_id_mod_NxN_mod_N];\n }\n else\n {\n my_bmAx = my_Einv * my_s_mem[lane_id_mod_NxN_div_N];\n }\n\n // Reduce bmAx terms.\n my_bmAx = reduce_distributed_vectors( my_bmAx, is_leader, active_mask );\n\n // Store the results.\n if ( ROW_MAJOR )\n {\n if ( lane_id_mod_NxN_mod_N == 0 )\n {\n new_rhs[N * a_row_id + lane_id_mod_NxN_div_N] = my_bmAx;\n }\n }\n else\n {\n if ( lane_id_mod_NxN_div_N == 0 )\n {\n new_rhs[N * a_row_id + lane_id_mod_NxN_mod_N] = my_bmAx;\n }\n }\n }\n}\n\n\n////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n\n\ntemplate< typename Vector_type, typename Matrix_type, int N>\nvoid distributed_rhs_mod_dispatch( const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Matrix_type *__restrict A_vals,\n const int *__restrict A_diag,\n const Vector_type *x,\n const Vector_type *b,\n Vector_type *new_rhs,\n const int num_owned_rows,\n const int row_major,\n const int has_external_diag\n )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int NUM_ROWS_PER_WARP = WARP_SIZE / NxN;\n // Number of items computer per CTA.\n const int NUM_ROWS_PER_CTA = NUM_ROWS_PER_WARP * NUM_WARPS_PER_CTA;\n // The number of threads to launch.\n const int grid_size = std::min( 4096, (num_owned_rows + NUM_ROWS_PER_CTA - 1) / NUM_ROWS_PER_CTA );\n // Branch to the correct kernel call.\n int code = 2 * (row_major ? 1 : 0) + (has_external_diag ? 1 : 0);\n\n switch ( code )\n {\n case 0: // Column-major, no external diagonal.\n b_minus_A_halo_x <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n new_rhs,\n num_owned_rows\n );\n cudaCheckError();\n break;\n\n case 1: // Column-major, external diagonal.\n b_minus_A_halo_x <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n new_rhs,\n num_owned_rows\n );\n cudaCheckError();\n break;\n\n case 2: // Row-major, no external diagonal.\n b_minus_A_halo_x <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n new_rhs,\n num_owned_rows\n );\n cudaCheckError();\n break;\n\n case 3: // Row-major, external diagonal.\n b_minus_A_halo_x <<< grid_size, CTA_SIZE>>>(\n A_rows,\n A_cols,\n A_vals,\n A_diag,\n x,\n b,\n new_rhs,\n num_owned_rows\n );\n cudaCheckError();\n break;\n\n default:\n FatalError( \"Internal error\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n cudaCheckError();\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I>\nvoid DenseLUSolver >::\ndistributed_rhs_mod(\n const Vector_d &x,\n const Vector_d &b,\n Vector_d &new_rhs)\n{\n const Matrix_d *A = dynamic_cast(Base::m_A);\n\n switch (A->get_block_dimy())\n {\n case 1:\n distributed_rhs_mod_dispatch(A->row_offsets.raw(), A->col_indices.raw(), A->values.raw(), A->diag.raw(), x.raw(), b.raw(), new_rhs.raw(), A->get_num_rows(), A->getBlockFormat() == ROW_MAJOR, A->hasProps(DIAG));\n break;\n\n case 2:\n distributed_rhs_mod_dispatch(A->row_offsets.raw(), A->col_indices.raw(), A->values.raw(), A->diag.raw(), x.raw(), b.raw(), new_rhs.raw(), A->get_num_rows(), A->getBlockFormat() == ROW_MAJOR, A->hasProps(DIAG));\n break;\n\n case 3:\n distributed_rhs_mod_dispatch(A->row_offsets.raw(), A->col_indices.raw(), A->values.raw(), A->diag.raw(), x.raw(), b.raw(), new_rhs.raw(), A->get_num_rows(), A->getBlockFormat() == ROW_MAJOR, A->hasProps(DIAG));\n break;\n\n case 4:\n distributed_rhs_mod_dispatch(A->row_offsets.raw(), A->col_indices.raw(), A->values.raw(), A->diag.raw(), x.raw(), b.raw(), new_rhs.raw(), A->get_num_rows(), A->getBlockFormat() == ROW_MAJOR, A->hasProps(DIAG));\n break;\n\n case 5:\n distributed_rhs_mod_dispatch(A->row_offsets.raw(), A->col_indices.raw(), A->values.raw(), A->diag.raw(), x.raw(), b.raw(), new_rhs.raw(), A->get_num_rows(), A->getBlockFormat() == ROW_MAJOR, A->hasProps(DIAG));\n break;\n }\n}\n\n// copy non zero elements only\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid DenseLUSolver >::\ncsr_to_dense(\n)\n{\n const Matrix_d *A = dynamic_cast(Base::m_A);\n const int block_size = 256;\n const int num_warps = block_size / WARP_SIZE;\n const int grid_size = (A->get_num_rows() + num_warps - 1) / num_warps;\n hipStream_t stream = amgx::thrust::global_thread_handle::get_stream();\n csr_to_dense_kernel <<< grid_size, block_size, 0, stream>>>(\n A->get_num_rows(),\n A->get_num_cols(),\n A->get_block_dimx(),\n A->get_block_dimy(),\n A->row_offsets.raw(),\n A->col_indices.raw(),\n A->values.raw(),\n A->hasProps(DIAG) ? A->diag.raw() : NULL,\n m_dense_A,\n m_lda\n );\n cudaCheckError();\n}\n\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid DenseLUSolver >::cudense_getrf()\n{\n int wsize = 0;\n hipsolverStatus_t status1 = hipsolverDnXgetrf_bufferSize(m_cuds_handle,\n m_num_rows,\n m_num_cols,\n m_dense_A,\n m_lda,\n &wsize);\n\n if ( status1 != HIPSOLVER_STATUS_SUCCESS)\n {\n FatalError( \"Failed kernel in DenseLU\", AMGX_ERR_INTERNAL);\n }\n\n cudaCheckError();\n allocMem(m_trf_wspace, wsize, false);\n status1 = hipsolverDnXgetrf(m_cuds_handle, m_num_rows, m_num_cols,\n m_dense_A, m_lda, m_trf_wspace, m_ipiv, m_cuds_info);\n cudaCheckError();\n\n if ( status1 != HIPSOLVER_STATUS_SUCCESS)\n {\n FatalError( \"failed kernel in Dense LU is likely due to invalid input parameters\",\n AMGX_ERR_INTERNAL);\n }\n else\n {\n int t_info;\n hipMemcpy(&t_info, m_cuds_info, sizeof(int), hipMemcpyDefault);\n cudaCheckError();\n\n if (t_info != 0)\n {\n FatalError( \"Fail to get info from cudense\", AMGX_ERR_INTERNAL);\n }\n else\n {\n // We follow the standard established by Lapack and used in cudense.\n if (t_info > 0)\n {\n FatalError( \"Dense LU factorization failed due to a singular matrix\", AMGX_ERR_INTERNAL);\n }\n else if (t_info < 0)\n {\n FatalError( \"Invalid input parameter(s) to dense LU\", AMGX_ERR_INTERNAL);\n }\n }\n }\n\n cudaCheckError();\n\n if (m_trf_wspace)\n {\n amgx::memory::hipFreeAsync(m_trf_wspace);\n m_trf_wspace = 0;\n }\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid DenseLUSolver >::cudense_getrs( Vector_d &x )\n{\n //Solve L*X = RHS\n hipsolverStatus_t status = hipsolverDnXgetrs(m_cuds_handle,\n HIPBLAS_OP_N,\n m_num_rows,\n 1,\n m_dense_A,\n m_lda,\n m_ipiv,\n (Matrix_data *)(x.raw()),\n m_num_rows,\n m_cuds_info);\n\n if (status != HIPSOLVER_STATUS_SUCCESS)\n {\n FatalError( \"cuSolver trsv failed to solve Lx=rhs\", AMGX_ERR_INTERNAL);\n }\n\n cudaCheckError();\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\ntemplate< class DataType, class IndexType >\nvoid\nDenseLUSolver >::\nallocMem(DataType *&ptr,\n IndexType numEntry,\n bool initToZero)\n{\n if ( ptr != NULL ) { amgx::memory::hipFreeAsync(ptr); }\n cudaCheckError();\n\n size_t sz = numEntry * sizeof(DataType);\n amgx::memory::hipMallocAsync((void **)&ptr, sz);\n cudaCheckError();\n\n if (initToZero)\n {\n hipMemset(ptr, 0x0, sz);\n cudaCheckError();\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nDenseLUSolver >::\nDenseLUSolver(AMG_Config &cfg,\n const std::string &cfg_scope,\n ThreadManager *tmng)\n : Solver(cfg, cfg_scope, tmng),\n m_cuds_handle(0),\n m_num_rows(0),\n m_num_cols(0),\n m_lda(0),\n m_dense_A(0),\n m_ipiv(0),\n m_cuds_info(0),\n m_trf_wspace(0)\n{\n // Allocate a handle for cudense\n hipsolverStatus_t status = hipsolverDnCreate(&m_cuds_handle);\n\n if ( status != HIPSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not create the CUDENSE handle\", AMGX_ERR_CUDA_FAILURE );\n }\n\n // Allocate a handle for cublas\n hipblasStatus_t hipblasStatus_t = hipblasCreate(&m_cublas_handle);\n\n if ( hipblasStatus_t != HIPBLAS_STATUS_SUCCESS )\n {\n FatalError( \"Could not create the CUBLAS handle\", AMGX_ERR_CUDA_FAILURE );\n }\n\n // Define the cudense stream.\n status = hipsolverSetStream(m_cuds_handle, amgx::thrust::global_thread_handle::get_stream());\n\n if ( status != HIPSOLVER_STATUS_SUCCESS )\n {\n FatalError( \"Could not set the stream for CUDENSE\", AMGX_ERR_CUDA_FAILURE );\n }\n\n // Make sure we don't run more than 1 iteration.\n this->set_max_iters(1);\n allocMem(m_cuds_info, sizeof(int), false);\n\n // Determine if the scalable coarse solve optimisation is enabled\n m_enable_exact_solve = (\n cfg.getParameter( \"algorithm\", cfg_scope) == CLASSICAL &&\n cfg.getParameter( \"exact_coarse_solve\", cfg_scope) == 1);\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nDenseLUSolver >::~DenseLUSolver()\n{\n if (m_cuds_handle) { hipsolverDnDestroy(m_cuds_handle); }\n\n if (m_cublas_handle) { hipblasDestroy(m_cublas_handle); }\n\n if (m_dense_A)\n {\n amgx::memory::hipFreeAsync(m_dense_A);\n }\n\n if (m_ipiv)\n {\n amgx::memory::hipFreeAsync(m_ipiv);\n }\n\n if (m_trf_wspace)\n {\n amgx::memory::hipFreeAsync(m_trf_wspace);\n }\n\n if (m_cuds_info)\n {\n amgx::memory::hipFreeAsync(m_cuds_info);\n }\n\n cudaCheckError();\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid DenseLUSolver >::solve_init(Vector_d &, Vector_d &, bool)\n{}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid DenseLUSolver >::solve_finalize(Vector_d &, Vector_d &)\n{}\n\n// Offset the local row offsets to global row offsets\ntemplate \n__global__ void local_row_offsets_to_global(\n int num_rows, int offset,\n IndexType* local_Arows)\n{\n int i = threadIdx.x + blockIdx.x*blockDim.x;\n if(i >= num_rows) { return; }\n\n local_Arows[i] += offset;\n}\n\n// Offset local packed column indices to global unpacked indices\ntemplate \n__global__ void local_col_indices_to_global(\n int nnz, int num_rows, int offset, IndexType* local_Acols, L2GType* l2g)\n{\n int i = threadIdx.x + blockIdx.x*blockDim.x;\n if(i >= nnz) { return; }\n\n if(local_Acols[i] >= num_rows)\n {\n local_Acols[i] = l2g[local_Acols[i] - num_rows];\n }\n else\n {\n local_Acols[i] += offset;\n }\n}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid DenseLUSolver >::\nsolver_setup(bool reuse_matrix_structure)\n{\n Matrix_d *A = dynamic_cast(Base::m_A);\n\n if (!A)\n {\n FatalError(\"DenseLUSolver only works with explicit matrices\",\n AMGX_ERR_INTERNAL);\n }\n\n ViewType oldView = A->currentView();\n A->setViewExterior();\n\n if(A->is_matrix_distributed() && m_enable_exact_solve)\n {\n#ifdef AMGX_WITH_MPI\n int rank = A->manager->global_id();\n int nranks = A->manager->get_num_partitions();\n MPI_Comm comm = A->manager->getComms()->get_mpi_comm();\n\n int offset, num_rows, nnz;\n A->getOffsetAndSizeForView(OWNED, &offset, &num_rows);\n A->getNnzForView(OWNED, &nnz);\n\n m_num_rows = A->manager->num_rows_global * A->get_block_dimx();\n m_num_cols = A->manager->num_rows_global * A->get_block_dimy();\n m_lda = m_num_rows; // col-major\n\n // Allocate mem for cudense pivoting sequence.\n allocMem(m_ipiv, m_num_rows, false);\n\n // Allocate memory to store the dense A and initialize to zero.\n allocMem(m_dense_A, static_cast(m_num_cols) * m_num_rows, true);\n\n // Much of the data can be reused if we are performing a resetup\n if (!reuse_matrix_structure)\n {\n // Gather the number of non zeros on each rank\n A->manager->getComms()->all_gather(nnz, nz_all, nranks);\n\n // Gather the number of rows on each rank\n A->manager->getComms()->all_gather(num_rows, row_all, nranks);\n\n // Get the number of non zeros on all ranks\n m_nnz_global = thrust_wrapper::reduce(nz_all.begin(), nz_all.end());\n\n // Turn the non-zero counts into displacements\n nz_displs.resize(nranks);\n thrust_wrapper::exclusive_scan(nz_all.begin(), nz_all.end(), nz_displs.begin());\n\n // Turn the number of rows into displacements\n row_displs.resize(nranks);\n thrust_wrapper::exclusive_scan(row_all.begin(), row_all.end(), row_displs.begin());\n\n IVector_d local_Acols_d(nnz);\n IVector_d local_Arows_d(num_rows);\n\n amgx::thrust::copy(A->col_indices.begin(), A->col_indices.begin() + nnz, local_Acols_d.begin());\n amgx::thrust::copy(A->row_offsets.begin(), A->row_offsets.begin() + num_rows, local_Arows_d.begin());\n\n // XXX Local to global map is the current limiting factor to enabling this\n // code for the aggregation based path. It's not clear whether there is\n // a structure that provides the same inverse mapping with aggregation.\n // Note that at one point inverse_renumbering was tested for aggregation\n // but didn't appear to work in all cases.\n\n // Convert the local column indices and row offsets to the global index space\n constexpr int nthreads = 128;\n int nblocks = nnz / nthreads + 1;\n local_col_indices_to_global<<>>(nnz, num_rows, row_displs[rank], local_Acols_d.raw(), A->manager->local_to_global_map.raw());\n cudaCheckError();\n\n nblocks = num_rows / nthreads + 1;\n local_row_offsets_to_global<<>>(num_rows, nz_displs[rank], local_Arows_d.raw());\n cudaCheckError();\n\n // Copy the transformed indices to the host\n IVector_h local_Acols_h(nnz);\n IVector_h local_Arows_h(num_rows);\n amgx::thrust::copy(local_Acols_d.begin(), local_Acols_d.end(), local_Acols_h.begin());\n amgx::thrust::copy(local_Arows_d.begin(), local_Arows_d.end(), local_Arows_h.begin());\n\n // Gather the local matrix structure redundantly to every rank\n IVector_h Acols_global_h(m_nnz_global);\n A->manager->getComms()->all_gather_v(local_Acols_h, nnz, Acols_global_h, nz_all, nz_displs);\n\n // Note: Copy the local data to global without guard value\n IVector_h Arows_global_h(m_num_rows + 1);\n A->manager->getComms()->all_gather_v(local_Arows_h, num_rows, Arows_global_h, row_all, row_displs);\n\n // Manually set the guard value on the global matrix\n Arows_global_h[m_num_rows] = m_nnz_global;\n Acols_global.resize(m_nnz_global);\n Arows_global.resize(m_num_rows + 1);\n amgx::thrust::copy(Acols_global_h.begin(), Acols_global_h.end(), Acols_global.begin());\n amgx::thrust::copy(Arows_global_h.begin(), Arows_global_h.end(), Arows_global.begin());\n }\n\n // Fetch to the host a copy of the local sparse matrix\n MVector_h local_Avals_h(nnz);\n amgx::thrust::copy(A->values.begin(), A->values.begin() + nnz, local_Avals_h.begin());\n\n // Gather the matrix values to all ranks\n MVector_h Avals_global_h(m_nnz_global);\n A->manager->getComms()->all_gather_v(local_Avals_h, nnz, Avals_global_h, nz_all, nz_displs);\n\n allocMem(m_dense_A, static_cast(m_num_cols) * m_lda, true);\n\n MVector_d Avals_global(m_nnz_global);\n amgx::thrust::copy(Avals_global_h.begin(), Avals_global_h.end(), Avals_global.begin());\n\n const int block_size = 256;\n const int num_warps = block_size / WARP_SIZE;\n const int grid_size = (m_nnz_global + num_warps - 1) / num_warps;\n hipStream_t stream = amgx::thrust::global_thread_handle::get_stream();\n csr_to_dense_kernel<<>>(\n m_num_rows,\n m_num_cols,\n A->get_block_dimx(),\n A->get_block_dimy(),\n Arows_global.raw(),\n Acols_global.raw(),\n Avals_global.raw(),\n A->hasProps(DIAG) ? A->diag.raw() : NULL,\n m_dense_A,\n m_lda);\n cudaCheckError();\n\n hipStreamSynchronize(stream);\n cudaCheckError();\n#endif\n }\n else\n {\n ViewType oldView = A->currentView();\n A->setViewExterior();\n m_num_rows = A->get_num_rows() * A->get_block_dimx();\n // don't use A->get_num_cols() because A is rectangular.\n // Only the diagonal block owned by this rank is factored.\n m_num_cols = A->get_num_rows() * A->get_block_dimy();\n m_lda = m_num_rows; // col-major\n // Allocate mem for cudense pivoting sequence.\n allocMem(m_ipiv, m_num_rows, false);\n // Allocate memory to store the dense A and initialize to zero.\n allocMem(m_dense_A, static_cast(m_num_cols) * m_lda, true);\n csr_to_dense(); // copy sparse A to dense_A\n }\n\n cudense_getrf(); // do LU factor\n A->setView(oldView);\n}\n\n// There is one subtle point here (for inexact solve):\n// We only do LU on the diagonal blocks associated with each rank.\n// Halo is used to update the right-hand-side (RHS) vector.\n// For multi GPU cases, this is essentially block Jacobi. Since the block size\n// is the size of each partition, this is OK trade-off between accuracy and runtime.\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nAMGX_STATUS\nDenseLUSolver >::\nsolve_iteration(Vector_d &rhs,\n Vector_d &x,\n bool xIsZero)\n{\n Matrix_d *A = dynamic_cast(Base::m_A);\n ViewType oldView = A->currentView();\n A->setViewExterior();\n\n if(A->is_matrix_distributed() && m_enable_exact_solve)\n {\n#ifdef AMGX_WITH_MPI\n int offset, num_rows;\n A->getOffsetAndSizeForView(OWNED, &offset, &num_rows);\n\n int rank = A->manager->global_id();\n int nranks = A->manager->get_num_partitions();\n MPI_Comm comm = A->manager->getComms()->get_mpi_comm();\n\n // Make host copy of the RHS\n MVector_h rhs_local_h(num_rows);\n amgx::thrust::copy(rhs.begin(), rhs.begin() + num_rows, rhs_local_h.begin());\n\n // Gather the local RHS from all ranks to global vectors on all ranks\n MVector_h rhs_global_h(m_num_rows);\n A->manager->getComms()->all_gather_v(rhs_local_h, num_rows, rhs_global_h, row_all, row_displs);\n\n //Solve L*X = RHS\n MVector_d x_global(m_num_rows);\n amgx::thrust::copy(rhs_global_h.begin(), rhs_global_h.end(), x_global.begin());\n hipsolverStatus_t status = \n hipsolverDnXgetrs(m_cuds_handle,\n HIPBLAS_OP_N,\n m_num_rows,\n 1,\n m_dense_A,\n m_lda,\n m_ipiv,\n x_global.raw(),\n m_num_rows,\n m_cuds_info);\n\n // Copy the local portion of the solution back into x\n amgx::thrust::copy(x_global.begin() + row_displs[rank], x_global.begin() + row_displs[rank] + num_rows, x.begin());\n\n if (status != HIPSOLVER_STATUS_SUCCESS)\n {\n FatalError(\"cuSolver trsv failed to solve Lx=rhs\", AMGX_ERR_INTERNAL);\n }\n#endif\n }\n else\n {\n Matrix_d *A = dynamic_cast(Base::m_A);\n ViewType oldView = A->currentView();\n A->setViewExterior();\n\n if ((!A->is_matrix_singleGPU()) && (!xIsZero))\n {\n // Modify rhs to include contribution from halo nodes\n // i.e. new_rhs = b - A_halo*x;\n // Note: dense_lu solver doesn't support latency hiding\n A->manager->exchange_halo_async(x, x.tag);\n A->manager->exchange_halo_wait(x, x.tag);\n Vector_d new_rhs(rhs.size());\n distributed_rhs_mod(x, rhs, new_rhs);\n amgx::thrust::copy(new_rhs.begin(), new_rhs.begin() + m_num_rows, x.begin());\n cudaCheckError();\n }\n else\n {\n x.copy(rhs);\n }\n\n cudense_getrs(x); // triangular solves\n }\n\n //Speculative send of x vector\n x.dirtybit = 1;\n A->setView(oldView);\n return AMGX_ST_CONVERGED; // direct solver always converges\n}\n\n\n#define AMGX_CASE_LINE(CASE) template class DenseLUSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n// AMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace dense_lu\n} // namespace amgx\n\n", "source": "/kaggle/working/amgx/src/solvers/dense_lu_solver.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\ntemplate class BinormalizationScaler;\n\n/**********************************************************************\n * HOST FUNCTIONS\n *********************************************************************/\ntemplate \nvoid computeBetaGammaHost(int rows, IndexType *offsets, IndexType *indices, MatrixType *vals,\n VectorType *x, VectorType *y, VectorType *beta, VectorType *gamma)\n{\n for (int i = 0; i < rows; i++) { gamma[i] = 0.; }\n\n for (int i = 0; i < rows; i++)\n {\n VectorType bi = 0.;\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int col = indices[jj];\n VectorType val = vals[jj];\n bi += (val * val) * y[col];\n gamma[col] += (val * val) * x[i];\n }\n\n beta[i] = bi;\n }\n}\n\n// compute Gamma on its own\ntemplate \nvoid computeGammaHost(int rows, IndexType *offsets, IndexType *indices, MatrixType *vals,\n VectorType *x, VectorType *gamma)\n{\n for (int i = 0; i < rows; i++) { gamma[i] = 0.; }\n\n for (int i = 0; i < rows; i++)\n {\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int col = indices[jj];\n VectorType val = vals[jj];\n gamma[col] += (val * val) * x[i];\n }\n }\n}\n\n// compute Beta on its own\ntemplate \nvoid computeBetaHost(int nrows, IndexType *offsets, IndexType *indices, MatrixType *vals,\n VectorType *y, VectorType *beta)\n{\n for (int i = 0; i < nrows; i++)\n {\n VectorType bi = 0.;\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int col = indices[jj];\n VectorType val = vals[jj];\n bi += (val * val) * y[col];\n }\n\n beta[i] = bi;\n }\n}\n\ntemplate \nvoid scaleMatrixHost(int nrows, IndexType *offsets, IndexType *indices, MatrixType *values,\n VectorType *x, VectorType *y)\n{\n for (int i = 0; i < nrows; i++)\n {\n VectorType fi = sqrt(fabs(x[i]));\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int j = indices[jj];\n VectorType gj = sqrt(fabs(y[j]));\n values[jj] *= fi * gj;\n }\n }\n}\n\n/**********************************************************************\n * DEVICE FUNCTIONS\n *********************************************************************/\n\n// compute initial beta, which is B*[1,...,1]'\ntemplate \n__global__\nvoid computeBetaIniDevice(int nrows, IndexType *offsets, IndexType *indices, MatrixValue *values, VectorValue *beta)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < nrows; i += gridDim.x * blockDim.x)\n {\n VectorValue rowsum = 0.0;\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n rowsum += values[jj] * values[jj];\n }\n\n beta[i] = rowsum;\n }\n}\n\ntemplate \n__global__\nvoid grabDiagonalVector(int nrows, IndexType *offsets, IndexType *indices, MatrixType *values, VectorType *diag)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < nrows; i += gridDim.x * blockDim.x)\n {\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n IndexType j = indices[jj];\n\n if (i == j) { diag[i] = values[jj]; }\n }\n }\n}\n\n// functor to generate stddev of vectors\ntemplate \nstruct std_f\n{\n std_f(T x) : v(x) {};\n T v;\n\n __host__ __device__\n T operator()(const T &x1, const T &x2) const\n {\n return (x1 * x2 - v) * (x1 * x2 - v);\n }\n};\n\n// scaled the matrix using diag(F)*A*diag(G), f = sqrt(fabs(x)), g = sqrt(fabs(y))\ntemplate \n__global__\nvoid scaleMatrixDevice(int rows, IndexType *offsets, IndexType *indices, MatrixType *values,\n VectorType *x)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n VectorType fi = fabs(x[i]);\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int j = indices[jj];\n VectorType fj = fabs(x[j]);\n values[jj] *= sqrt(fabs(fi * fj));\n }\n }\n}\n\n\ntemplate \n__global__\nvoid getColRowNorms(int rows, IndexType *offsets, IndexType *indices, ValueTypeA *values,\n ValueTypeB *rownorms, ValueTypeB *colnorms)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n ValueTypeB rownorm = 0.;\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int j = indices[jj];\n ValueTypeB curval = values[jj] * values[jj];\n rownorm += curval;\n utils::atomic_add(colnorms + j, curval);\n }\n\n rownorms[i] = rownorm;\n }\n}\n\n// these warp reductions should be able to be replaced with amgx:: functions\ntemplate \n__device__ __inline__ T warpReduceSum(T val)\n{\n if (warpSize > 16) { val += utils::shfl_down(val, 16, warpSize); }\n\n if (warpSize > 8) { val += utils::shfl_down(val, 8, warpSize); }\n\n if (warpSize > 4) { val += utils::shfl_down(val, 4, warpSize); }\n\n if (warpSize > 2) { val += utils::shfl_down(val, 2, warpSize); }\n\n if (warpSize > 1) { val += utils::shfl_down(val, 1, warpSize); }\n\n return val;\n}\n\ntemplate \n__device__ T warpReduceSumShared(volatile T *vals, const int lane_id)\n{\n if (warpSize > 16) { vals[lane_id] += vals[lane_id + 16]; }\n\n if (warpSize > 8) { vals[lane_id] += vals[lane_id + 8]; }\n\n if (warpSize > 4) { vals[lane_id] += vals[lane_id + 4]; }\n\n if (warpSize > 2) { vals[lane_id] += vals[lane_id + 2]; }\n\n if (warpSize > 1) { vals[lane_id] += vals[lane_id + 1]; }\n\n return vals[lane_id];\n}\n\n// compute gamma = B^T*x (B = A.^2)\ntemplate \n__global__\nvoid computeBetaDevice(const int nrows, IndexType *offsets, IndexType *indices, ValueTypeA *values,\n ValueTypeB *diag, ValueTypeB *x, ValueTypeB *xn, ValueTypeB *beta, const ValueTypeB avg, ValueTypeB *avg_vec)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < nrows; i += gridDim.x * blockDim.x)\n {\n ValueTypeB xi = x[i];\n ValueTypeB bi = beta[i];\n ValueTypeB di = diag[i];\n ValueTypeB c0 = -di * xi * xi + 2 * bi * xi - nrows * avg;\n ValueTypeB c1 = (nrows - 2) * (bi - di * xi);\n ValueTypeB c2 = (nrows - 1) * di;\n assert(c0 > epsilon(c0)); //\n // delta = xi - x(i)\n ValueTypeB dx = (2 * c0) / (-c1 - sqrt(c1 * c1 - 4 * c2 * c0)) - x[i];\n ValueTypeB davg = 0.;\n ValueTypeB dbeta = 0.;\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int j = indices[jj];\n ValueTypeA Bij = values[jj] * values[jj];\n davg += Bij * x[j]; // += x' * B(:, i) == B(i, :)*x, because B is symmetric\n dbeta += dx * Bij;\n }\n\n beta[i] = bi + dbeta;\n avg_vec[i] = dx * (davg + bi + di * dx) / nrows;\n //utils::atomic_add(avg, davg);\n }\n}\n\n\ntemplate\nstruct square_value\n{\n __host__ __device__ T operator()(const T &x) const\n {\n return x * x;\n }\n};\n\n// vector constant scale operand\ntemplate \nstruct vmul_scale_const\n{\n T _alpha;\n\n vmul_scale_const(T alpha): _alpha(alpha) {};\n\n __host__ __device__\n T operator()(const T &vec) const\n {\n return vec * _alpha;\n }\n};\n\n// vector scale operand\ntemplate \nstruct vmul_scale\n{\n\n vmul_scale() {};\n\n __host__ __device__\n T operator()(const T &vec, const T &alpha) const\n {\n return (vec * sqrt(fabs(alpha)));\n }\n};\n\n// vector unscale operand\ntemplate \nstruct vmul_unscale\n{\n\n vmul_unscale() {};\n\n __host__ __device__\n T operator()(const T &vec, const T &alpha) const\n {\n return (vec / sqrt(fabs(alpha)));\n }\n};\n\n\n// Setup on Device\ntemplate \nvoid BinormalizationScaler >::setup(Matrix_d &A)\n{\n if (A.is_matrix_distributed())\n {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // move these out to config parameters?\n const int max_iters = 10;\n const ValueTypeB tolerance = 1e-10;\n VVector diag(A.get_num_rows());\n grabDiagonalVector <<< 4096, 128>>>(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), diag.raw());\n cudaCheckError();\n int nrows = A.get_num_rows();\n // temporary vectors\n VVector x(nrows, 1), xn(nrows), davg(nrows), beta(nrows, 0);\n computeBetaIniDevice <<< 4096, 256>>>(nrows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), beta.raw());\n cudaCheckError();\n ValueTypeB avg = thrust_wrapper::reduce(beta.begin(), beta.end()) / nrows;\n // calculate initial std1 and std2\n amgx::thrust::device_ptr x_ptr(x.raw()), beta_ptr(beta.raw());\n ValueTypeB stdx = sqrt(amgx::thrust::inner_product(x_ptr, x_ptr + nrows, beta_ptr, ValueTypeB(0.), amgx::thrust::plus(), std_f(avg)) / nrows) / avg;\n\n for (int t = 0; t < max_iters; t++)\n {\n if (fabs(stdx) < tolerance) { break; } // finished\n\n computeBetaDevice <<< 4096, 256>>>(nrows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(),\n diag.raw(), x.raw(), xn.raw(), beta.raw(), avg, davg.raw());\n cudaCheckError();\n avg += thrust_wrapper::reduce(davg.begin(), davg.end());\n // ValueTypeB stdx_old = stdx;\n stdx = sqrt(amgx::thrust::inner_product(x_ptr, x_ptr + nrows, beta_ptr, ValueTypeB(0.), amgx::thrust::plus(), std_f(avg)) / nrows) / avg;\n // print it #, current error, convergence rate\n // printf(\"ITER: %d %.3e %.3e %.4lg\\n\",t, stdx, stdx_old, stdx / stdx_old);\n }\n\n //Save scaling vectors for later user, setup complete\n scale_vector = VVector(x);\n}\n\n\n// Matrix Scaling on Device\ntemplate \nvoid BinormalizationScaler >::scaleMatrix(Matrix_d &A, ScaleDirection scaleOrUnscale)\n{\n if (scale_vector.size() != A.get_num_rows())\n {\n FatalError(\"Must call setup(A) before binormalization scaling can scale matrix\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (A.is_matrix_distributed())\n {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n int nrows = A.get_num_rows();\n /*VVector rownorms(nrows, 0.0);\n VVector colnorms(nrows, 0.0);\n getColRowNorms<<<4096,256>>>(nrows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), rownorms.raw(), colnorms.raw());\n cudaCheckError();\n ValueTypeB row_max = *(amgx::thrust::max_element(rownorms.begin(), rownorms.end()));\n ValueTypeB row_min = *(amgx::thrust::min_element(rownorms.begin(), rownorms.end()));\n ValueTypeB col_max = *(amgx::thrust::max_element(colnorms.begin(), colnorms.end()));\n ValueTypeB col_min = *(amgx::thrust::min_element(colnorms.begin(), colnorms.end()));\n cudaCheckError();\n printf(\"Original Matrix: rowmax: %e, rowmin: %e, colmax: %e, colmin: %e\\n\", row_max, row_min, col_max, col_min);fflush(stdout);*/\n scaleMatrixDevice <<< 4096, 256>>>(nrows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), scale_vector.raw());\n cudaCheckError();\n ValueTypeB C_norm = sqrt(thrust_wrapper::transform_reduce(A.values.begin(), A.values.begin() + A.get_num_nz() * A.get_block_size(), square_value(), 0., amgx::thrust::plus()) / nrows);\n thrust_wrapper::transform(A.values.begin(), A.values.begin() + A.get_num_nz()*A.get_block_size(), A.values.begin(), vmul_scale_const(1. / C_norm) );\n thrust_wrapper::transform(scale_vector.begin(), scale_vector.end(), scale_vector.begin(), vmul_scale_const(sqrt(1. / C_norm)) );\n cudaCheckError();\n /*thrust_wrapper::fill(rownorms.begin(), rownorms.end(), 0.);\n thrust_wrapper::fill(colnorms.begin(), colnorms.end(), 0.);\n getColRowNorms<<<4096,256>>>(nrows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), rownorms.raw(), colnorms.raw());\n cudaCheckError();\n row_max = *(amgx::thrust::max_element(rownorms.begin(), rownorms.end()));\n row_min = *(amgx::thrust::min_element(rownorms.begin(), rownorms.end()));\n col_max = *(amgx::thrust::max_element(colnorms.begin(), colnorms.end()));\n col_min = *(amgx::thrust::min_element(colnorms.begin(), colnorms.end()));\n cudaCheckError();\n printf(\"Scaled Matrix: rowmax: %e, rowmin: %e, colmax: %e, colmin: %e\\n\", row_max, row_min, col_max, col_min);fflush(stdout);*/\n exit(0);\n}\n\n// Setup on Host\ntemplate \nvoid BinormalizationScaler >::setup(Matrix_h &A)\n{\n FatalError(\"Host not supported\", AMGX_ERR_NOT_IMPLEMENTED);\n /*if (A.is_matrix_distributed()) {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n // move these out to config parameters?\n const int max_iters = 10;\n const ValueTypeB tolerance = 1e-10;\n\n int rows = A.get_num_rows(), cols = A.get_num_cols();\n // temporary vectors\n VVector x(rows, 1), y(cols, 1), beta(rows, 0), gamma(cols, 0);\n\n // perform matvecs to get beta and gamma (spmv for beta, spmvT for gamma)\n computeBetaGammaHost(rows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(),\n x.raw(), y.raw(), beta.raw(), gamma.raw());\n\n double std1 = 0., std2 = 0., sum1 = cols, sum2 = rows;\n // calculate initial std1 and std2\n for (int i=0; i\nvoid BinormalizationScaler >::scaleMatrix(Matrix_h &A, ScaleDirection scaleOrUnscale)\n{\n FatalError(\"Host not supported\", AMGX_ERR_NOT_IMPLEMENTED);\n /*if (A.is_matrix_distributed()) {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n if (left_scale.size() != A.get_num_rows()) {\n FatalError(\"Must call setup(A) before binormalization scaling can scale matrix\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // A_scaled = F*A*G (f = diag(F) = sqrt(fabs(x)), g = diag(G) = sqrt(fabs(y))\n // A_ij = f_i * A_ij * g_j\n scaleMatrixHost(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), left_scale.raw(), right_scale.raw());*/\n}\n\n\ntemplate \nvoid BinormalizationScaler >::scaleVector(VVector &v, ScaleDirection scaleOrUnscale, ScaleSide leftOrRight)\n{\n if (scaleOrUnscale == amgx::SCALE)\n {\n thrust_wrapper::transform(v.begin(), v.end(), this->scale_vector.begin(), v.begin(), vmul_scale() );\n }\n else\n {\n thrust_wrapper::transform(v.begin(), v.end(), this->scale_vector.begin(), v.begin(), vmul_unscale() );\n }\n}\n\ntemplate \nvoid BinormalizationScaler >::scaleVector(VVector &v, ScaleDirection scaleOrUnscale, ScaleSide leftOrRight)\n{\n FatalError(\"4x4 block size not supported\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class BinormalizationScaler_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class BinormalizationScaler::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n\n} // namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\ntemplate class BinormalizationScaler;\n\n/**********************************************************************\n * HOST FUNCTIONS\n *********************************************************************/\ntemplate \nvoid computeBetaGammaHost(int rows, IndexType *offsets, IndexType *indices, MatrixType *vals,\n VectorType *x, VectorType *y, VectorType *beta, VectorType *gamma)\n{\n for (int i = 0; i < rows; i++) { gamma[i] = 0.; }\n\n for (int i = 0; i < rows; i++)\n {\n VectorType bi = 0.;\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int col = indices[jj];\n VectorType val = vals[jj];\n bi += (val * val) * y[col];\n gamma[col] += (val * val) * x[i];\n }\n\n beta[i] = bi;\n }\n}\n\n// compute Gamma on its own\ntemplate \nvoid computeGammaHost(int rows, IndexType *offsets, IndexType *indices, MatrixType *vals,\n VectorType *x, VectorType *gamma)\n{\n for (int i = 0; i < rows; i++) { gamma[i] = 0.; }\n\n for (int i = 0; i < rows; i++)\n {\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int col = indices[jj];\n VectorType val = vals[jj];\n gamma[col] += (val * val) * x[i];\n }\n }\n}\n\n// compute Beta on its own\ntemplate \nvoid computeBetaHost(int nrows, IndexType *offsets, IndexType *indices, MatrixType *vals,\n VectorType *y, VectorType *beta)\n{\n for (int i = 0; i < nrows; i++)\n {\n VectorType bi = 0.;\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int col = indices[jj];\n VectorType val = vals[jj];\n bi += (val * val) * y[col];\n }\n\n beta[i] = bi;\n }\n}\n\ntemplate \nvoid scaleMatrixHost(int nrows, IndexType *offsets, IndexType *indices, MatrixType *values,\n VectorType *x, VectorType *y)\n{\n for (int i = 0; i < nrows; i++)\n {\n VectorType fi = sqrt(fabs(x[i]));\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int j = indices[jj];\n VectorType gj = sqrt(fabs(y[j]));\n values[jj] *= fi * gj;\n }\n }\n}\n\n/**********************************************************************\n * DEVICE FUNCTIONS\n *********************************************************************/\n\n// compute initial beta, which is B*[1,...,1]'\ntemplate \n__global__\nvoid computeBetaIniDevice(int nrows, IndexType *offsets, IndexType *indices, MatrixValue *values, VectorValue *beta)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < nrows; i += gridDim.x * blockDim.x)\n {\n VectorValue rowsum = 0.0;\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n rowsum += values[jj] * values[jj];\n }\n\n beta[i] = rowsum;\n }\n}\n\ntemplate \n__global__\nvoid grabDiagonalVector(int nrows, IndexType *offsets, IndexType *indices, MatrixType *values, VectorType *diag)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < nrows; i += gridDim.x * blockDim.x)\n {\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n IndexType j = indices[jj];\n\n if (i == j) { diag[i] = values[jj]; }\n }\n }\n}\n\n// functor to generate stddev of vectors\ntemplate \nstruct std_f\n{\n std_f(T x) : v(x) {};\n T v;\n\n __host__ __device__\n T operator()(const T &x1, const T &x2) const\n {\n return (x1 * x2 - v) * (x1 * x2 - v);\n }\n};\n\n// scaled the matrix using diag(F)*A*diag(G), f = sqrt(fabs(x)), g = sqrt(fabs(y))\ntemplate \n__global__\nvoid scaleMatrixDevice(int rows, IndexType *offsets, IndexType *indices, MatrixType *values,\n VectorType *x)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n VectorType fi = fabs(x[i]);\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int j = indices[jj];\n VectorType fj = fabs(x[j]);\n values[jj] *= sqrt(fabs(fi * fj));\n }\n }\n}\n\n\ntemplate \n__global__\nvoid getColRowNorms(int rows, IndexType *offsets, IndexType *indices, ValueTypeA *values,\n ValueTypeB *rownorms, ValueTypeB *colnorms)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n ValueTypeB rownorm = 0.;\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int j = indices[jj];\n ValueTypeB curval = values[jj] * values[jj];\n rownorm += curval;\n utils::atomic_add(colnorms + j, curval);\n }\n\n rownorms[i] = rownorm;\n }\n}\n\n// these warp reductions should be able to be replaced with amgx:: functions\ntemplate \n__device__ __inline__ T warpReduceSum(T val)\n{\n if (warpSize > 16) { val += utils::shfl_down(val, 16, warpSize); }\n\n if (warpSize > 8) { val += utils::shfl_down(val, 8, warpSize); }\n\n if (warpSize > 4) { val += utils::shfl_down(val, 4, warpSize); }\n\n if (warpSize > 2) { val += utils::shfl_down(val, 2, warpSize); }\n\n if (warpSize > 1) { val += utils::shfl_down(val, 1, warpSize); }\n\n return val;\n}\n\ntemplate \n__device__ T warpReduceSumShared(volatile T *vals, const int lane_id)\n{\n if (warpSize > 16) { vals[lane_id] += vals[lane_id + 16]; }\n\n if (warpSize > 8) { vals[lane_id] += vals[lane_id + 8]; }\n\n if (warpSize > 4) { vals[lane_id] += vals[lane_id + 4]; }\n\n if (warpSize > 2) { vals[lane_id] += vals[lane_id + 2]; }\n\n if (warpSize > 1) { vals[lane_id] += vals[lane_id + 1]; }\n\n return vals[lane_id];\n}\n\n// compute gamma = B^T*x (B = A.^2)\ntemplate \n__global__\nvoid computeBetaDevice(const int nrows, IndexType *offsets, IndexType *indices, ValueTypeA *values,\n ValueTypeB *diag, ValueTypeB *x, ValueTypeB *xn, ValueTypeB *beta, const ValueTypeB avg, ValueTypeB *avg_vec)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < nrows; i += gridDim.x * blockDim.x)\n {\n ValueTypeB xi = x[i];\n ValueTypeB bi = beta[i];\n ValueTypeB di = diag[i];\n ValueTypeB c0 = -di * xi * xi + 2 * bi * xi - nrows * avg;\n ValueTypeB c1 = (nrows - 2) * (bi - di * xi);\n ValueTypeB c2 = (nrows - 1) * di;\n assert(c0 > epsilon(c0)); //\n // delta = xi - x(i)\n ValueTypeB dx = (2 * c0) / (-c1 - sqrt(c1 * c1 - 4 * c2 * c0)) - x[i];\n ValueTypeB davg = 0.;\n ValueTypeB dbeta = 0.;\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int j = indices[jj];\n ValueTypeA Bij = values[jj] * values[jj];\n davg += Bij * x[j]; // += x' * B(:, i) == B(i, :)*x, because B is symmetric\n dbeta += dx * Bij;\n }\n\n beta[i] = bi + dbeta;\n avg_vec[i] = dx * (davg + bi + di * dx) / nrows;\n //utils::atomic_add(avg, davg);\n }\n}\n\n\ntemplate\nstruct square_value\n{\n __host__ __device__ T operator()(const T &x) const\n {\n return x * x;\n }\n};\n\n// vector constant scale operand\ntemplate \nstruct vmul_scale_const\n{\n T _alpha;\n\n vmul_scale_const(T alpha): _alpha(alpha) {};\n\n __host__ __device__\n T operator()(const T &vec) const\n {\n return vec * _alpha;\n }\n};\n\n// vector scale operand\ntemplate \nstruct vmul_scale\n{\n\n vmul_scale() {};\n\n __host__ __device__\n T operator()(const T &vec, const T &alpha) const\n {\n return (vec * sqrt(fabs(alpha)));\n }\n};\n\n// vector unscale operand\ntemplate \nstruct vmul_unscale\n{\n\n vmul_unscale() {};\n\n __host__ __device__\n T operator()(const T &vec, const T &alpha) const\n {\n return (vec / sqrt(fabs(alpha)));\n }\n};\n\n\n// Setup on Device\ntemplate \nvoid BinormalizationScaler >::setup(Matrix_d &A)\n{\n if (A.is_matrix_distributed())\n {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // move these out to config parameters?\n const int max_iters = 10;\n const ValueTypeB tolerance = 1e-10;\n VVector diag(A.get_num_rows());\n grabDiagonalVector <<< 4096, 128>>>(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), diag.raw());\n cudaCheckError();\n int nrows = A.get_num_rows();\n // temporary vectors\n VVector x(nrows, 1), xn(nrows), davg(nrows), beta(nrows, 0);\n computeBetaIniDevice <<< 4096, 256>>>(nrows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), beta.raw());\n cudaCheckError();\n ValueTypeB avg = thrust_wrapper::reduce(beta.begin(), beta.end()) / nrows;\n // calculate initial std1 and std2\n amgx::thrust::device_ptr x_ptr(x.raw()), beta_ptr(beta.raw());\n ValueTypeB stdx = sqrt(amgx::thrust::inner_product(x_ptr, x_ptr + nrows, beta_ptr, ValueTypeB(0.), amgx::thrust::plus(), std_f(avg)) / nrows) / avg;\n\n for (int t = 0; t < max_iters; t++)\n {\n if (fabs(stdx) < tolerance) { break; } // finished\n\n computeBetaDevice <<< 4096, 256>>>(nrows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(),\n diag.raw(), x.raw(), xn.raw(), beta.raw(), avg, davg.raw());\n cudaCheckError();\n avg += thrust_wrapper::reduce(davg.begin(), davg.end());\n // ValueTypeB stdx_old = stdx;\n stdx = sqrt(amgx::thrust::inner_product(x_ptr, x_ptr + nrows, beta_ptr, ValueTypeB(0.), amgx::thrust::plus(), std_f(avg)) / nrows) / avg;\n // print it #, current error, convergence rate\n // printf(\"ITER: %d %.3e %.3e %.4lg\\n\",t, stdx, stdx_old, stdx / stdx_old);\n }\n\n //Save scaling vectors for later user, setup complete\n scale_vector = VVector(x);\n}\n\n\n// Matrix Scaling on Device\ntemplate \nvoid BinormalizationScaler >::scaleMatrix(Matrix_d &A, ScaleDirection scaleOrUnscale)\n{\n if (scale_vector.size() != A.get_num_rows())\n {\n FatalError(\"Must call setup(A) before binormalization scaling can scale matrix\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (A.is_matrix_distributed())\n {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n int nrows = A.get_num_rows();\n /*VVector rownorms(nrows, 0.0);\n VVector colnorms(nrows, 0.0);\n getColRowNorms<<<4096,256>>>(nrows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), rownorms.raw(), colnorms.raw());\n cudaCheckError();\n ValueTypeB row_max = *(amgx::thrust::max_element(rownorms.begin(), rownorms.end()));\n ValueTypeB row_min = *(amgx::thrust::min_element(rownorms.begin(), rownorms.end()));\n ValueTypeB col_max = *(amgx::thrust::max_element(colnorms.begin(), colnorms.end()));\n ValueTypeB col_min = *(amgx::thrust::min_element(colnorms.begin(), colnorms.end()));\n cudaCheckError();\n printf(\"Original Matrix: rowmax: %e, rowmin: %e, colmax: %e, colmin: %e\\n\", row_max, row_min, col_max, col_min);fflush(stdout);*/\n scaleMatrixDevice <<< 4096, 256>>>(nrows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), scale_vector.raw());\n cudaCheckError();\n ValueTypeB C_norm = sqrt(thrust_wrapper::transform_reduce(A.values.begin(), A.values.begin() + A.get_num_nz() * A.get_block_size(), square_value(), 0., amgx::thrust::plus()) / nrows);\n thrust_wrapper::transform(A.values.begin(), A.values.begin() + A.get_num_nz()*A.get_block_size(), A.values.begin(), vmul_scale_const(1. / C_norm) );\n thrust_wrapper::transform(scale_vector.begin(), scale_vector.end(), scale_vector.begin(), vmul_scale_const(sqrt(1. / C_norm)) );\n cudaCheckError();\n /*thrust_wrapper::fill(rownorms.begin(), rownorms.end(), 0.);\n thrust_wrapper::fill(colnorms.begin(), colnorms.end(), 0.);\n getColRowNorms<<<4096,256>>>(nrows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), rownorms.raw(), colnorms.raw());\n cudaCheckError();\n row_max = *(amgx::thrust::max_element(rownorms.begin(), rownorms.end()));\n row_min = *(amgx::thrust::min_element(rownorms.begin(), rownorms.end()));\n col_max = *(amgx::thrust::max_element(colnorms.begin(), colnorms.end()));\n col_min = *(amgx::thrust::min_element(colnorms.begin(), colnorms.end()));\n cudaCheckError();\n printf(\"Scaled Matrix: rowmax: %e, rowmin: %e, colmax: %e, colmin: %e\\n\", row_max, row_min, col_max, col_min);fflush(stdout);*/\n exit(0);\n}\n\n// Setup on Host\ntemplate \nvoid BinormalizationScaler >::setup(Matrix_h &A)\n{\n FatalError(\"Host not supported\", AMGX_ERR_NOT_IMPLEMENTED);\n /*if (A.is_matrix_distributed()) {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n // move these out to config parameters?\n const int max_iters = 10;\n const ValueTypeB tolerance = 1e-10;\n\n int rows = A.get_num_rows(), cols = A.get_num_cols();\n // temporary vectors\n VVector x(rows, 1), y(cols, 1), beta(rows, 0), gamma(cols, 0);\n\n // perform matvecs to get beta and gamma (spmv for beta, spmvT for gamma)\n computeBetaGammaHost(rows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(),\n x.raw(), y.raw(), beta.raw(), gamma.raw());\n\n double std1 = 0., std2 = 0., sum1 = cols, sum2 = rows;\n // calculate initial std1 and std2\n for (int i=0; i\nvoid BinormalizationScaler >::scaleMatrix(Matrix_h &A, ScaleDirection scaleOrUnscale)\n{\n FatalError(\"Host not supported\", AMGX_ERR_NOT_IMPLEMENTED);\n /*if (A.is_matrix_distributed()) {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n if (left_scale.size() != A.get_num_rows()) {\n FatalError(\"Must call setup(A) before binormalization scaling can scale matrix\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // A_scaled = F*A*G (f = diag(F) = sqrt(fabs(x)), g = diag(G) = sqrt(fabs(y))\n // A_ij = f_i * A_ij * g_j\n scaleMatrixHost(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), left_scale.raw(), right_scale.raw());*/\n}\n\n\ntemplate \nvoid BinormalizationScaler >::scaleVector(VVector &v, ScaleDirection scaleOrUnscale, ScaleSide leftOrRight)\n{\n if (scaleOrUnscale == amgx::SCALE)\n {\n thrust_wrapper::transform(v.begin(), v.end(), this->scale_vector.begin(), v.begin(), vmul_scale() );\n }\n else\n {\n thrust_wrapper::transform(v.begin(), v.end(), this->scale_vector.begin(), v.begin(), vmul_unscale() );\n }\n}\n\ntemplate \nvoid BinormalizationScaler >::scaleVector(VVector &v, ScaleDirection scaleOrUnscale, ScaleSide leftOrRight)\n{\n FatalError(\"4x4 block size not supported\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class BinormalizationScaler_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class BinormalizationScaler::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n\n} // namespace amgx\n\n", "source": "/kaggle/working/amgx/src/scalers/binormalization.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n\nnamespace amgx\n{\n\n/**********************************************************************\n * DEVICE FUNCTIONS\n *********************************************************************/\ntemplate \n__global__\nvoid ScaleMatrix(int rows, IndexType *offsets, IndexType *indices, MatrixType *values,\n VectorType *scale, ScaleDirection scaleOrUnscale)\n{\n if ( scaleOrUnscale == SCALE)\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n VectorType si = scale[i];\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n IndexType j = indices[jj];\n VectorType sj = scale[j];\n // scale in-place\n values[jj] *= si * sj;\n }\n }\n\n if ( scaleOrUnscale == UNSCALE)\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n VectorType si = scale[i];\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n IndexType j = indices[jj];\n VectorType sj = scale[j];\n // scale in-place\n values[jj] /= si * sj;\n }\n }\n}\n\ntemplate \n__global__\nvoid scaleVectorKernel(int rows, VectorType *values, VectorType *diag, ScaleDirection scaleOrUnscale, ScaleSide leftOrRight)\n{\n //Ignore the scaleside, it is the same for both in the diagonal symmetric case\n if (scaleOrUnscale == SCALE)\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n values[i] *= diag[i];\n }\n\n if (scaleOrUnscale == UNSCALE)\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n values[i] /= diag[i];\n }\n}\n\ntemplate \n__global__\nvoid checkPositiveVector(int N, T *v, bool *positive)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < N; i += gridDim.x * blockDim.x)\n {\n if (v[i] < T(0)) { positive[0] = false; }\n }\n}\n\ntemplate \n__global__\nvoid sqrtVector(int N, T *v)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < N; i += gridDim.x * blockDim.x)\n {\n // we know by this point that the diagonal is +ve\n v[i] = 1. / sqrt(v[i]);\n }\n}\n\ntemplate \n__global__\nvoid grabDiagonalVector(int rows, IndexType *offsets, IndexType *indices, MatrixType *values, VectorType *diag)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n IndexType j = indices[jj];\n\n if (i == j) { diag[i] = values[jj]; }\n }\n }\n}\n\n// Setup scaling vector\n// Device version\ntemplate \nvoid DiagonalSymmetricScaler >::setup(Matrix_d &A)\n{\n if (A.is_matrix_distributed())\n {\n FatalError(\"Diagonal Symmetric scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // grab diagonal\n diag = VVector(A.get_num_rows());\n grabDiagonalVector <<< 4096, 128>>>(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), diag.raw());\n cudaCheckError();\n // check diagonal +ve\n bool positive = true, *d_positive;\n amgx::memory::cudaMallocAsync((void**)&d_positive, sizeof(bool));\n cudaCheckError();\n cudaMemcpy(d_positive, &positive, sizeof(bool), cudaMemcpyHostToDevice);\n cudaCheckError();\n checkPositiveVector <<< 4096, 256>>>(diag.size(), diag.raw(), d_positive);\n cudaCheckError();\n cudaMemcpy(&positive, d_positive, sizeof(bool), cudaMemcpyDeviceToHost);\n cudaCheckError();\n\n // fail out if necessary\n if (!positive)\n {\n FatalError(\"Diagonal symmetric scaling only applies to SPD systems with positive diagonal entries\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // set scale = 1./sqrt(diag) , reciprocal square root\n sqrtVector <<< 4096, 256>>>(diag.size(), diag.raw());\n cudaCheckError();\n}\n\n//Host version\ntemplate \nvoid DiagonalSymmetricScaler >::setup(Matrix_h &A)\n{\n if (A.is_matrix_distributed())\n {\n FatalError(\"Diagonal Symmetric scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // grab diagonal\n diag = VVector(A.get_num_rows());\n //grabDiagonalVector<<<4096, 128>>>(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), diag.raw());\n cudaCheckError();\n IndexType *offsets = A.row_offsets.raw();\n IndexType *indices = A.col_indices.raw();\n ValueTypeA *values = A.values.raw();\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n IndexType j = indices[jj];\n\n if (i == j) { diag[i] = values[jj]; }\n }\n }\n\n // check diagonal +ve\n bool positive = true;\n\n //checkPositiveVector<<<4096, 256>>>(diag.size(), diag.raw(), d_positive);\n cudaCheckError();\n for (int i = 0; i < diag.size(); i++)\n {\n if (diag[i] < ValueTypeB(0)) { positive = false; }\n }\n\n // fail out if necessary\n if (!positive)\n {\n FatalError(\"Diagonal symmetric scaling only applies to SPD systems with positive diagonal entries\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // set scale = 1./sqrt(diag) , reciprocal square root\n // sqrtVector<<<4096, 256>>>(diag.size(), diag.raw());\n cudaCheckError();\n for (int i = 0; i < diag.size(); i++)\n {\n // we know by this point that the diagonal is +ve\n diag[i] = 1. / sqrt(diag[i]);\n }\n}\n\n// 1x1 Device\ntemplate \nvoid DiagonalSymmetricScaler >::scaleMatrix(Matrix_d &A, ScaleDirection scaleOrUnscale)\n{\n // A_ij *= scale[i] * scale[j];\n ScaleMatrix <<< 4096, 256>>>(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(),\n A.values.raw(), diag.raw(), scaleOrUnscale);\n cudaCheckError();\n}\n\n// 1x1 Host\ntemplate \nvoid DiagonalSymmetricScaler >::scaleMatrix(Matrix_h &A, ScaleDirection scaleOrUnscale)\n{\n IndexType *offsets = A.row_offsets.raw();\n IndexType *indices = A.col_indices.raw();\n ValueTypeA *values = A.values.raw();\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n if ( scaleOrUnscale == SCALE)\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n IndexType j = indices[jj];\n values[jj] *= diag[i] * diag[j];\n }\n\n if ( scaleOrUnscale == UNSCALE)\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n IndexType j = indices[jj];\n values[jj] /= diag[i] * diag[j];\n }\n }\n}\n\n\ntemplate \nvoid DiagonalSymmetricScaler >::scaleVector(VVector &v, ScaleDirection scaleOrUnscale, ScaleSide leftOrRight)\n{\n scaleVectorKernel <<< 4096, 256>>>(v.size(), v.raw(), diag.raw(), scaleOrUnscale, leftOrRight);\n cudaCheckError();\n}\n\n// 4x4 Host\ntemplate \nvoid DiagonalSymmetricScaler >::scaleVector(VVector &v, ScaleDirection scaleOrUnscale, ScaleSide leftOrRight)\n{\n //ignore leftOrRIght for diagonal scale, they are the same\n if ( scaleOrUnscale == SCALE)\n for (int i = 0; i < diag.size(); i++)\n {\n v[i] *= diag[i];\n }\n\n if ( scaleOrUnscale == UNSCALE)\n for (int i = 0; i < diag.size(); i++)\n {\n v[i] /= diag[i];\n }\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class DiagonalSymmetricScaler_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class DiagonalSymmetricScaler::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\nnamespace amgx\n{\n\n/**********************************************************************\n * DEVICE FUNCTIONS\n *********************************************************************/\ntemplate \n__global__\nvoid ScaleMatrix(int rows, IndexType *offsets, IndexType *indices, MatrixType *values,\n VectorType *scale, ScaleDirection scaleOrUnscale)\n{\n if ( scaleOrUnscale == SCALE)\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n VectorType si = scale[i];\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n IndexType j = indices[jj];\n VectorType sj = scale[j];\n // scale in-place\n values[jj] *= si * sj;\n }\n }\n\n if ( scaleOrUnscale == UNSCALE)\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n VectorType si = scale[i];\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n IndexType j = indices[jj];\n VectorType sj = scale[j];\n // scale in-place\n values[jj] /= si * sj;\n }\n }\n}\n\ntemplate \n__global__\nvoid scaleVectorKernel(int rows, VectorType *values, VectorType *diag, ScaleDirection scaleOrUnscale, ScaleSide leftOrRight)\n{\n //Ignore the scaleside, it is the same for both in the diagonal symmetric case\n if (scaleOrUnscale == SCALE)\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n values[i] *= diag[i];\n }\n\n if (scaleOrUnscale == UNSCALE)\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n values[i] /= diag[i];\n }\n}\n\ntemplate \n__global__\nvoid checkPositiveVector(int N, T *v, bool *positive)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < N; i += gridDim.x * blockDim.x)\n {\n if (v[i] < T(0)) { positive[0] = false; }\n }\n}\n\ntemplate \n__global__\nvoid sqrtVector(int N, T *v)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < N; i += gridDim.x * blockDim.x)\n {\n // we know by this point that the diagonal is +ve\n v[i] = 1. / sqrt(v[i]);\n }\n}\n\ntemplate \n__global__\nvoid grabDiagonalVector(int rows, IndexType *offsets, IndexType *indices, MatrixType *values, VectorType *diag)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n IndexType j = indices[jj];\n\n if (i == j) { diag[i] = values[jj]; }\n }\n }\n}\n\n// Setup scaling vector\n// Device version\ntemplate \nvoid DiagonalSymmetricScaler >::setup(Matrix_d &A)\n{\n if (A.is_matrix_distributed())\n {\n FatalError(\"Diagonal Symmetric scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // grab diagonal\n diag = VVector(A.get_num_rows());\n grabDiagonalVector <<< 4096, 128>>>(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), diag.raw());\n cudaCheckError();\n // check diagonal +ve\n bool positive = true, *d_positive;\n amgx::memory::hipMallocAsync((void**)&d_positive, sizeof(bool));\n cudaCheckError();\n hipMemcpy(d_positive, &positive, sizeof(bool), hipMemcpyHostToDevice);\n cudaCheckError();\n checkPositiveVector <<< 4096, 256>>>(diag.size(), diag.raw(), d_positive);\n cudaCheckError();\n hipMemcpy(&positive, d_positive, sizeof(bool), hipMemcpyDeviceToHost);\n cudaCheckError();\n\n // fail out if necessary\n if (!positive)\n {\n FatalError(\"Diagonal symmetric scaling only applies to SPD systems with positive diagonal entries\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // set scale = 1./sqrt(diag) , reciprocal square root\n sqrtVector <<< 4096, 256>>>(diag.size(), diag.raw());\n cudaCheckError();\n}\n\n//Host version\ntemplate \nvoid DiagonalSymmetricScaler >::setup(Matrix_h &A)\n{\n if (A.is_matrix_distributed())\n {\n FatalError(\"Diagonal Symmetric scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // grab diagonal\n diag = VVector(A.get_num_rows());\n //grabDiagonalVector<<<4096, 128>>>(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), diag.raw());\n cudaCheckError();\n IndexType *offsets = A.row_offsets.raw();\n IndexType *indices = A.col_indices.raw();\n ValueTypeA *values = A.values.raw();\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n IndexType j = indices[jj];\n\n if (i == j) { diag[i] = values[jj]; }\n }\n }\n\n // check diagonal +ve\n bool positive = true;\n\n //checkPositiveVector<<<4096, 256>>>(diag.size(), diag.raw(), d_positive);\n cudaCheckError();\n for (int i = 0; i < diag.size(); i++)\n {\n if (diag[i] < ValueTypeB(0)) { positive = false; }\n }\n\n // fail out if necessary\n if (!positive)\n {\n FatalError(\"Diagonal symmetric scaling only applies to SPD systems with positive diagonal entries\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // set scale = 1./sqrt(diag) , reciprocal square root\n // sqrtVector<<<4096, 256>>>(diag.size(), diag.raw());\n cudaCheckError();\n for (int i = 0; i < diag.size(); i++)\n {\n // we know by this point that the diagonal is +ve\n diag[i] = 1. / sqrt(diag[i]);\n }\n}\n\n// 1x1 Device\ntemplate \nvoid DiagonalSymmetricScaler >::scaleMatrix(Matrix_d &A, ScaleDirection scaleOrUnscale)\n{\n // A_ij *= scale[i] * scale[j];\n ScaleMatrix <<< 4096, 256>>>(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(),\n A.values.raw(), diag.raw(), scaleOrUnscale);\n cudaCheckError();\n}\n\n// 1x1 Host\ntemplate \nvoid DiagonalSymmetricScaler >::scaleMatrix(Matrix_h &A, ScaleDirection scaleOrUnscale)\n{\n IndexType *offsets = A.row_offsets.raw();\n IndexType *indices = A.col_indices.raw();\n ValueTypeA *values = A.values.raw();\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n if ( scaleOrUnscale == SCALE)\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n IndexType j = indices[jj];\n values[jj] *= diag[i] * diag[j];\n }\n\n if ( scaleOrUnscale == UNSCALE)\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n IndexType j = indices[jj];\n values[jj] /= diag[i] * diag[j];\n }\n }\n}\n\n\ntemplate \nvoid DiagonalSymmetricScaler >::scaleVector(VVector &v, ScaleDirection scaleOrUnscale, ScaleSide leftOrRight)\n{\n scaleVectorKernel <<< 4096, 256>>>(v.size(), v.raw(), diag.raw(), scaleOrUnscale, leftOrRight);\n cudaCheckError();\n}\n\n// 4x4 Host\ntemplate \nvoid DiagonalSymmetricScaler >::scaleVector(VVector &v, ScaleDirection scaleOrUnscale, ScaleSide leftOrRight)\n{\n //ignore leftOrRIght for diagonal scale, they are the same\n if ( scaleOrUnscale == SCALE)\n for (int i = 0; i < diag.size(); i++)\n {\n v[i] *= diag[i];\n }\n\n if ( scaleOrUnscale == UNSCALE)\n for (int i = 0; i < diag.size(); i++)\n {\n v[i] /= diag[i];\n }\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class DiagonalSymmetricScaler_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class DiagonalSymmetricScaler::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n\n", "source": "/kaggle/working/amgx/src/scalers/diagonal_symmetric.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\ntemplate class NBinormalizationScaler;\n\n/**********************************************************************\n * HOST FUNCTIONS\n *********************************************************************/\ntemplate \nvoid computeBetaGammaHost(int rows, IndexType *offsets, IndexType *indices, MatrixType *vals,\n VectorType *x, VectorType *y, VectorType *beta, VectorType *gamma)\n{\n for (int i = 0; i < rows; i++) { gamma[i] = 0.; }\n\n for (int i = 0; i < rows; i++)\n {\n VectorType bi = 0.;\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int col = indices[jj];\n VectorType val = vals[jj];\n bi += (val * val) * y[col];\n gamma[col] += (val * val) * x[i];\n }\n\n beta[i] = bi;\n }\n}\n\n// compute Gamma on its own\ntemplate \nvoid computeGammaHost(int rows, IndexType *offsets, IndexType *indices, MatrixType *vals,\n VectorType *x, VectorType *gamma)\n{\n for (int i = 0; i < rows; i++) { gamma[i] = 0.; }\n\n for (int i = 0; i < rows; i++)\n {\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int col = indices[jj];\n VectorType val = vals[jj];\n gamma[col] += (val * val) * x[i];\n }\n }\n}\n\n// compute Beta on its own\ntemplate \nvoid computeBetaHost(int rows, IndexType *offsets, IndexType *indices, MatrixType *vals,\n VectorType *y, VectorType *beta)\n{\n for (int i = 0; i < rows; i++)\n {\n VectorType bi = 0.;\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int col = indices[jj];\n VectorType val = vals[jj];\n bi += (val * val) * y[col];\n }\n\n beta[i] = bi;\n }\n}\n\ntemplate \nvoid scaleMatrixHost(int rows, IndexType *offsets, IndexType *indices, MatrixType *values,\n VectorType *x, VectorType *y)\n{\n for (int i = 0; i < rows; i++)\n {\n VectorType fi = sqrt(fabs(x[i]));\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int j = indices[jj];\n VectorType gj = sqrt(fabs(y[j]));\n values[jj] *= fi * gj;\n }\n }\n}\n\n/**********************************************************************\n * DEVICE FUNCTIONS\n *********************************************************************/\n\n// these warp reductions should be able to be replaced with amgx:: functions\ntemplate \n__device__ __inline__ T warpReduceSum(T val)\n{\n if (warpSize > 16) { val += utils::shfl_down(val, 16, warpSize); }\n\n utils::syncwarp();\n\n if (warpSize > 8) { val += utils::shfl_down(val, 8, warpSize); }\n\n utils::syncwarp();\n\n if (warpSize > 4) { val += utils::shfl_down(val, 4, warpSize); }\n\n utils::syncwarp();\n\n if (warpSize > 2) { val += utils::shfl_down(val, 2, warpSize); }\n\n utils::syncwarp();\n\n if (warpSize > 1) { val += utils::shfl_down(val, 1, warpSize); }\n\n return val;\n}\n\ntemplate \n__device__ T warpReduceSumShared(volatile T *vals, const int lane_id)\n{\n if (warpSize > 16) { vals[lane_id] += vals[lane_id + 16]; }\n\n if (warpSize > 8) { vals[lane_id] += vals[lane_id + 8]; }\n\n if (warpSize > 4) { vals[lane_id] += vals[lane_id + 4]; }\n\n if (warpSize > 2) { vals[lane_id] += vals[lane_id + 2]; }\n\n if (warpSize > 1) { vals[lane_id] += vals[lane_id + 1]; }\n\n return vals[lane_id];\n}\n\n// compute beta = B*y, gamma = C*x (B = A.^2, C = B^T)\ntemplate \n__global__\nvoid computeBetaGammaDevice(IndexType rows, IndexType *offsets, IndexType *indices, MatrixValue *values,\n VectorValue *x, VectorValue *y, VectorValue *beta, VectorValue *gamma)\n{\n const int vectors_per_block = VectorsPerCTA;\n const int vector_id = threadIdx.x / VectorSize;\n const int lane_id = threadIdx.x % VectorSize;\n\n for (int i = vectors_per_block * blockIdx.x + vector_id; i < rows; i += vectors_per_block * gridDim.x)\n {\n // load start + end pointers\n int row_tmp;\n\n if (lane_id < 2)\n {\n row_tmp = offsets[i + lane_id];\n }\n\n // distribute to all other threads in warp\n int row_begin = utils::shfl(row_tmp, vector_id * VectorSize, warpSize, utils::activemask());\n int row_end = utils::shfl(row_tmp, vector_id * VectorSize + 1, warpSize, utils::activemask());\n VectorValue bi(0.);\n\n for (int jj = row_begin + lane_id; utils::any(jj < row_end, utils::activemask()); jj += VectorSize)\n {\n int col = -1;\n VectorValue val(0.);\n\n if (jj < row_end)\n {\n col = indices[jj];\n val = values[jj];\n bi += (val * val) * y[col];\n utils::atomic_add(&gamma[col], (val * val) * x[i]);\n }\n }\n\n // reduce over bi\n VectorValue bi_s = warpReduceSum(bi);\n\n if (lane_id == 0)\n {\n beta[i] = bi_s;\n }\n }\n}\n\n// compute gamma = B^T*x (B = A.^2)\ntemplate \n__global__\nvoid computeGammaDevice(int rows, IndexType *offsets, IndexType *indices, MatrixValue *values,\n VectorValue *x, VectorValue *gamma)\n{\n const int vectors_per_block = CTASize / VectorSize;\n const int vector_id = threadIdx.x / VectorSize;\n const int lane_id = threadIdx.x % VectorSize;\n\n for (int i = vectors_per_block * blockIdx.x + vector_id; i < rows; i += vectors_per_block * gridDim.x)\n {\n // load start + end pointers\n int row_tmp;\n\n if (lane_id < 2)\n {\n row_tmp = offsets[i + lane_id];\n }\n\n // distribute to all other threads in warp\n int row_begin = utils::shfl(row_tmp, vector_id * VectorSize, warpSize, utils::activemask());\n int row_end = utils::shfl(row_tmp, vector_id * VectorSize + 1, warpSize, utils::activemask());\n\n for (int jj = row_begin + lane_id; utils::any(jj < row_end, utils::activemask()); jj += VectorSize)\n {\n int col = -1;\n VectorValue val = 0.;\n\n if (jj < row_end)\n {\n col = indices[jj];\n val = values[jj];\n utils::atomic_add(&gamma[col], (val * val) * x[i]);\n }\n }\n }\n}\n\n// compute beta = B*y (B = A.^2)\ntemplate \n__global__\nvoid computeBetaDevice(int rows, IndexType *offsets, IndexType *indices, MatrixValue *values,\n VectorValue *y, VectorValue *beta)\n{\n const int vectors_per_block = CTASize / VectorSize;\n const int vector_id = threadIdx.x / VectorSize;\n const int lane_id = threadIdx.x % VectorSize;\n\n for (int i = vectors_per_block * blockIdx.x + vector_id; i < rows; i += vectors_per_block * gridDim.x)\n {\n // load start + end pointers\n int row_tmp;\n\n if (lane_id < 2)\n {\n row_tmp = offsets[i + lane_id];\n }\n\n // distribute to all other threads in warp\n int row_begin = utils::shfl(row_tmp, vector_id * VectorSize, warpSize, utils::activemask());\n int row_end = utils::shfl(row_tmp, vector_id * VectorSize + 1, warpSize, utils::activemask());\n VectorValue bi = 0.;\n\n for (int jj = row_begin + lane_id; utils::any(jj < row_end, utils::activemask()); jj += VectorSize)\n {\n int col = -1;\n VectorValue val = 0.;\n\n if (jj < row_end)\n {\n col = indices[jj];\n val = values[jj];\n bi += (val * val) * y[col];\n }\n }\n\n // reduce over bi\n VectorValue bi_s = warpReduceSum(bi);\n\n if (lane_id == 0)\n {\n beta[i] = bi_s;\n }\n }\n}\n\ntemplate \n__global__\nvoid setOneOverVector(int N, ValueType *x, ValueType sum1, ValueType *beta)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < N; i += gridDim.x * blockDim.x)\n {\n //x[i] = ( isNotCloseToZero(beta[i]) ? sum1 / beta[i] : sum1 / epsilon(beta[i]) );\n x[i] = ( isNotCloseToZero(beta[i]) ? (sum1 / beta[i]) : (ValueType)1. );\n }\n}\n\ntemplate\nstruct square_value\n{\n __host__ __device__ T operator()(const T &x) const\n {\n return x * x;\n }\n};\n\n// functor to generate stddev of vectors\ntemplate \nstruct std_f\n{\n std_f(T x) : v(x) {};\n T v;\n\n __host__ __device__\n T operator()(const T &x1, const T &x2) const\n {\n return (x1 * x2 - v) * (x1 * x2 - v);\n }\n};\n\n// scaled the matrix using diag(F)*A*diag(G), f = sqrt(fabs(x)), g = sqrt(fabs(y))\ntemplate \n__global__\nvoid scaleMatrixDevice(int rows, IndexType *offsets, IndexType *indices, MatrixType *values,\n VectorType *x, VectorType *y)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n VectorType fi = sqrt(fabs(x[i]));\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int j = indices[jj];\n VectorType gj = sqrt(fabs(y[j]));\n\n // scale matrix value in place\n if (direction == amgx::SCALE)\n {\n values[jj] *= fi * gj;\n }\n else\n {\n values[jj] /= fi * gj;\n }\n }\n }\n}\n\n\ntemplate \n__global__\nvoid getColRowNorms(int rows, IndexType *offsets, IndexType *indices, MatrixType *values,\n VectorType *rownorms, VectorType *colnorms)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int j = indices[jj];\n VectorType curval = values[jj] * values[jj];\n rownorms[i] += curval;\n utils::atomic_add(colnorms + j, curval);\n }\n }\n}\n\n// vector constant scale operand\ntemplate \nstruct vmul_scale_const\n{\n T _alpha;\n\n vmul_scale_const(T alpha): _alpha(alpha) {};\n\n __host__ __device__\n T operator()(const T &vec) const\n {\n return vec * _alpha;\n }\n};\n\n\n// vector scale operand\ntemplate \nstruct vmul_scale\n{\n\n vmul_scale() {};\n\n __host__ __device__\n T operator()(const T &vec, const T &alpha) const\n {\n return (vec * sqrt(fabs(alpha)));\n }\n};\n\n// vector unscale operand\ntemplate \nstruct vmul_unscale\n{\n\n vmul_unscale() {};\n\n __host__ __device__\n T operator()(const T &vec, const T &alpha) const\n {\n return (vec / sqrt(fabs(alpha)));\n }\n};\n\n\n// Setup on Device\ntemplate \nvoid NBinormalizationScaler >::setup(Matrix_d &A)\n{\n if (A.is_matrix_distributed())\n {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // move these out to config parameters?\n const int max_iters = 50;\n const ValueTypeB tolerance = 1e-10;\n int rows = A.get_num_rows(), cols = A.get_num_cols();\n // temporary vectors\n VVector x(rows, 1), y(cols, 1), beta(rows, 0), gamma(cols, 0);\n // perform matvecs to get beta and gamma (spmv for beta, spmvT for gamma)\n computeBetaGammaDevice<256, 8, 32> <<< 4096, 256>>>(rows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(),\n x.raw(), y.raw(), beta.raw(), gamma.raw());\n cudaCheckError();\n ValueTypeB sum1 = cols, sum2 = rows, std1, std2;\n // calculate initial std1 and std2\n amgx::thrust::device_ptr x_ptr(x.raw()), y_ptr(y.raw()), beta_ptr(beta.raw()), gamma_ptr(gamma.raw());\n std1 = sqrt(amgx::thrust::inner_product(x_ptr, x_ptr + rows, beta_ptr, ValueTypeB(0.), amgx::thrust::plus(), std_f(sum1)) / rows) / sum1;\n std2 = sqrt(amgx::thrust::inner_product(y_ptr, y_ptr + cols, gamma_ptr, ValueTypeB(0.), thrust::plus(), std_f(sum2)) / cols) / sum2;\n ValueTypeB std = sqrt(std1 * std1 + std2 * std2);\n\n for (int t = 0; t < max_iters; t++)\n {\n if (std < tolerance) { break; } // finished\n\n // x = sum1 ./ beta\n setOneOverVector <<< 4096, 256>>>(rows, x.raw(), sum1, beta.raw());\n cudaCheckError();\n // gamma = C*x := B'*x\n thrust_wrapper::fill(gamma.begin(), gamma.end(), ValueTypeB(0.));\n computeGammaDevice<256, 8> <<< 4096, 256>>>(rows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), x.raw(), gamma.raw());\n cudaCheckError();\n // gamma = 1 ./ beta\n setOneOverVector <<< 4096, 256>>>(cols, y.raw(), sum2, gamma.raw());\n cudaCheckError();\n // beta = B*y\n computeBetaDevice<256, 8> <<< 4096, 256>>>(rows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), y.raw(), beta.raw());\n cudaCheckError();\n //ValueTypeB std_old = std;\n std = sqrt(amgx::thrust::inner_product(x_ptr, x_ptr + rows, beta_ptr, ValueTypeB(0.), amgx::thrust::plus(), std_f(sum1)) / rows) / sum1;\n // print it #, current error, convergence rate\n //printf(\"ITER: %d %.3e %.4lg\\n\",t, std, std / std_old);\n }\n\n //Save scaling vectors for later user, setup complete\n left_scale = VVector(beta);\n right_scale = VVector(gamma);\n this->scaled_before = false;\n}\n\n\n// Matrix Scaling on Device\ntemplate \nvoid NBinormalizationScaler >::scaleMatrix(Matrix_d &A, ScaleDirection scaleOrUnscale)\n{\n if (left_scale.size() != A.get_num_rows())\n {\n FatalError(\"Must call setup(A) before binormalization scaling can scale matrix\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (A.is_matrix_distributed())\n {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n int nrows = A.get_num_rows();\n\n if (scaleOrUnscale == amgx::SCALE)\n {\n /*VVector rownorms(nrows, 0.0);\n VVector colnorms(nrows, 0.0);\n getColRowNorms<<<4096,256>>>(nrows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), rownorms.raw(), colnorms.raw());\n cudaCheckError();\n ValueTypeB row_max = *(amgx::thrust::max_element(rownorms.begin(), rownorms.end()));\n ValueTypeB row_min = *(amgx::thrust::min_element(rownorms.begin(), rownorms.end()));\n ValueTypeB col_max = *(amgx::thrust::max_element(colnorms.begin(), colnorms.end()));\n ValueTypeB col_min = *(amgx::thrust::min_element(colnorms.begin(), colnorms.end()));\n cudaCheckError();\n printf(\"Original Matrix: rowmax: %e, rowmin: %e, colmax: %e, colmin: %e\\n\", row_max, row_min, col_max, col_min);fflush(stdout);*/\n // A_scaled = F*A*G (f = diag(F) = sqrt(fabs(x)), g = diag(G) = sqrt(fabs(y))\n // A_ij = f_i * A_ij * g_j\n scaleMatrixDevice <<< 4096, 256>>>(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), left_scale.raw(), right_scale.raw());\n cudaCheckError();\n\n if (!scaled_before)\n {\n this->norm_coef = sqrt(thrust_wrapper::transform_reduce(A.values.begin(), A.values.begin() + A.get_num_nz() * A.get_block_size(), square_value(), 0., amgx::thrust::plus()) / A.get_num_rows());\n cudaCheckError();\n thrust_wrapper::transform(A.values.begin(), A.values.begin() + A.get_num_nz()*A.get_block_size(), A.values.begin(), vmul_scale_const(1. / this->norm_coef) );\n thrust_wrapper::transform(left_scale.begin(), left_scale.end(), left_scale.begin(), vmul_scale_const(sqrt(1. / this->norm_coef)) );\n thrust_wrapper::transform(right_scale.begin(), right_scale.end(), right_scale.begin(), vmul_scale_const(sqrt(1. / this->norm_coef)) );\n cudaCheckError();\n /*thrust_wrapper::fill(rownorms.begin(), rownorms.end(), 0.);\n thrust_wrapper::fill(colnorms.begin(), colnorms.end(), 0.);\n getColRowNorms<<<4096,256>>>(nrows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), rownorms.raw(), colnorms.raw());\n cudaCheckError();\n row_max = *(amgx::thrust::max_element(rownorms.begin(), rownorms.end()));\n row_min = *(amgx::thrust::min_element(rownorms.begin(), rownorms.end()));\n col_max = *(amgx::thrust::max_element(colnorms.begin(), colnorms.end()));\n col_min = *(amgx::thrust::min_element(colnorms.begin(), colnorms.end()));\n cudaCheckError();\n printf(\"Scaled Matrix: rowmax: %e, rowmin: %e, colmax: %e, colmin: %e\\n\", row_max, row_min, col_max, col_min);fflush(stdout);*/\n }\n\n this->scaled_before = true;\n }\n else\n {\n scaleMatrixDevice <<< 4096, 256>>>(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), left_scale.raw(), right_scale.raw());\n cudaCheckError();\n }\n}\n\n// Setup on Host\ntemplate \nvoid NBinormalizationScaler >::setup(Matrix_h &A)\n{\n if (A.is_matrix_distributed())\n {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // move these out to config parameters?\n const int max_iters = 10;\n const ValueTypeB tolerance = 1e-10;\n int rows = A.get_num_rows(), cols = A.get_num_cols();\n // temporary vectors\n VVector x(rows, 1), y(cols, 1), beta(rows, 0), gamma(cols, 0);\n // perform matvecs to get beta and gamma (spmv for beta, spmvT for gamma)\n computeBetaGammaHost(rows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(),\n x.raw(), y.raw(), beta.raw(), gamma.raw());\n double std1 = 0., std2 = 0., sum1 = cols, sum2 = rows;\n\n // calculate initial std1 and std2\n for (int i = 0; i < rows; i++)\n {\n std1 += pow(x[i] * beta[i] - sum1, 2.0);\n }\n\n std1 = sqrt(std1 / rows) / sum1;\n\n for (int i = 0; i < cols; i++)\n {\n std2 += pow(y[i] * gamma[i] - sum2, 2.0);\n }\n\n std2 = sqrt(std2 / cols) / sum2;\n //printf(\"std1: %lg, std2: %lg\\n\",std1, std2);\n double std_initial = sqrt((std1 * std1) + (std2 * std2));\n double std = std_initial;\n\n for (int t = 0; t < max_iters; t++)\n {\n if (std < tolerance) { break; } // finished\n\n // x = sum1 ./ beta\n for (int i = 0; i < rows; i++) { x[i] = ( isNotCloseToZero(beta[i]) ? sum1 / beta[i] : sum1 / epsilon(beta[i]) ); }\n\n // gamma = C*x\n computeGammaHost(rows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), x.raw(), gamma.raw());\n\n // gamma = 1 ./ beta\n for (int i = 0; i < cols; i++) { y[i] = ( isNotCloseToZero(gamma[i]) ? sum2 / gamma[i] : sum2 / epsilon(gamma[i]) ); }\n\n // beta = B*y\n computeBetaHost(rows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), y.raw(), beta.raw());\n //ValueTypeB std_old = std;\n std = 0.;\n\n for (int i = 0; i < rows; i++)\n {\n std += pow(x[i] * beta[i] - sum1, 2.0);\n }\n\n std = sqrt(std / rows) / sum1;\n // print it #, current error, convergence rate\n //printf(\"ITER: %d %.3e %.4lg\\n\",t, std, std / std_old);\n }\n\n //Save scaling vectors for later user, setup complete\n left_scale = VVector(beta);\n right_scale = VVector(gamma);\n}\n\n\ntemplate \nvoid NBinormalizationScaler >::scaleMatrix(Matrix_h &A, ScaleDirection scaleOrUnscale)\n{\n if (A.is_matrix_distributed())\n {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (left_scale.size() != A.get_num_rows())\n {\n FatalError(\"Must call setup(A) before binormalization scaling can scale matrix\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // A_scaled = F*A*G (f = diag(F) = sqrt(fabs(x)), g = diag(G) = sqrt(fabs(y))\n // A_ij = f_i * A_ij * g_j\n scaleMatrixHost(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), left_scale.raw(), right_scale.raw());\n}\n\n\ntemplate \nvoid NBinormalizationScaler >::scaleVector(VVector &v, ScaleDirection scaleOrUnscale, ScaleSide leftOrRight)\n{\n VVector *scale_vector = (leftOrRight == amgx::LEFT) ? &this->left_scale : &this->right_scale;\n\n //thrust_wrapper::transform(v.begin(), v.end(), scale_vector->begin(), v.begin(), (scaleOrUnscale == amgx::SCALE) ? vmul_scale() : vmul_unscale() );\n if (scaleOrUnscale == amgx::SCALE)\n {\n thrust_wrapper::transform(v.begin(), v.end(), scale_vector->begin(), v.begin(), vmul_scale() );\n }\n else\n {\n thrust_wrapper::transform(v.begin(), v.end(), scale_vector->begin(), v.begin(), vmul_unscale() );\n }\n}\n\ntemplate \nvoid NBinormalizationScaler >::scaleVector(VVector &v, ScaleDirection scaleOrUnscale, ScaleSide leftOrRight)\n{\n FatalError(\"4x4 block size not supported\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class NBinormalizationScaler_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class NBinormalizationScaler::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n\n} // namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\ntemplate class NBinormalizationScaler;\n\n/**********************************************************************\n * HOST FUNCTIONS\n *********************************************************************/\ntemplate \nvoid computeBetaGammaHost(int rows, IndexType *offsets, IndexType *indices, MatrixType *vals,\n VectorType *x, VectorType *y, VectorType *beta, VectorType *gamma)\n{\n for (int i = 0; i < rows; i++) { gamma[i] = 0.; }\n\n for (int i = 0; i < rows; i++)\n {\n VectorType bi = 0.;\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int col = indices[jj];\n VectorType val = vals[jj];\n bi += (val * val) * y[col];\n gamma[col] += (val * val) * x[i];\n }\n\n beta[i] = bi;\n }\n}\n\n// compute Gamma on its own\ntemplate \nvoid computeGammaHost(int rows, IndexType *offsets, IndexType *indices, MatrixType *vals,\n VectorType *x, VectorType *gamma)\n{\n for (int i = 0; i < rows; i++) { gamma[i] = 0.; }\n\n for (int i = 0; i < rows; i++)\n {\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int col = indices[jj];\n VectorType val = vals[jj];\n gamma[col] += (val * val) * x[i];\n }\n }\n}\n\n// compute Beta on its own\ntemplate \nvoid computeBetaHost(int rows, IndexType *offsets, IndexType *indices, MatrixType *vals,\n VectorType *y, VectorType *beta)\n{\n for (int i = 0; i < rows; i++)\n {\n VectorType bi = 0.;\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int col = indices[jj];\n VectorType val = vals[jj];\n bi += (val * val) * y[col];\n }\n\n beta[i] = bi;\n }\n}\n\ntemplate \nvoid scaleMatrixHost(int rows, IndexType *offsets, IndexType *indices, MatrixType *values,\n VectorType *x, VectorType *y)\n{\n for (int i = 0; i < rows; i++)\n {\n VectorType fi = sqrt(fabs(x[i]));\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int j = indices[jj];\n VectorType gj = sqrt(fabs(y[j]));\n values[jj] *= fi * gj;\n }\n }\n}\n\n/**********************************************************************\n * DEVICE FUNCTIONS\n *********************************************************************/\n\n// these warp reductions should be able to be replaced with amgx:: functions\ntemplate \n__device__ __inline__ T warpReduceSum(T val)\n{\n if (warpSize > 16) { val += utils::shfl_down(val, 16, warpSize); }\n\n utils::syncwarp();\n\n if (warpSize > 8) { val += utils::shfl_down(val, 8, warpSize); }\n\n utils::syncwarp();\n\n if (warpSize > 4) { val += utils::shfl_down(val, 4, warpSize); }\n\n utils::syncwarp();\n\n if (warpSize > 2) { val += utils::shfl_down(val, 2, warpSize); }\n\n utils::syncwarp();\n\n if (warpSize > 1) { val += utils::shfl_down(val, 1, warpSize); }\n\n return val;\n}\n\ntemplate \n__device__ T warpReduceSumShared(volatile T *vals, const int lane_id)\n{\n if (warpSize > 16) { vals[lane_id] += vals[lane_id + 16]; }\n\n if (warpSize > 8) { vals[lane_id] += vals[lane_id + 8]; }\n\n if (warpSize > 4) { vals[lane_id] += vals[lane_id + 4]; }\n\n if (warpSize > 2) { vals[lane_id] += vals[lane_id + 2]; }\n\n if (warpSize > 1) { vals[lane_id] += vals[lane_id + 1]; }\n\n return vals[lane_id];\n}\n\n// compute beta = B*y, gamma = C*x (B = A.^2, C = B^T)\ntemplate \n__global__\nvoid computeBetaGammaDevice(IndexType rows, IndexType *offsets, IndexType *indices, MatrixValue *values,\n VectorValue *x, VectorValue *y, VectorValue *beta, VectorValue *gamma)\n{\n const int vectors_per_block = VectorsPerCTA;\n const int vector_id = threadIdx.x / VectorSize;\n const int lane_id = threadIdx.x % VectorSize;\n\n for (int i = vectors_per_block * blockIdx.x + vector_id; i < rows; i += vectors_per_block * gridDim.x)\n {\n // load start + end pointers\n int row_tmp;\n\n if (lane_id < 2)\n {\n row_tmp = offsets[i + lane_id];\n }\n\n // distribute to all other threads in warp\n int row_begin = utils::shfl(row_tmp, vector_id * VectorSize, warpSize, utils::activemask());\n int row_end = utils::shfl(row_tmp, vector_id * VectorSize + 1, warpSize, utils::activemask());\n VectorValue bi(0.);\n\n for (int jj = row_begin + lane_id; utils::any(jj < row_end, utils::activemask()); jj += VectorSize)\n {\n int col = -1;\n VectorValue val(0.);\n\n if (jj < row_end)\n {\n col = indices[jj];\n val = values[jj];\n bi += (val * val) * y[col];\n utils::atomic_add(&gamma[col], (val * val) * x[i]);\n }\n }\n\n // reduce over bi\n VectorValue bi_s = warpReduceSum(bi);\n\n if (lane_id == 0)\n {\n beta[i] = bi_s;\n }\n }\n}\n\n// compute gamma = B^T*x (B = A.^2)\ntemplate \n__global__\nvoid computeGammaDevice(int rows, IndexType *offsets, IndexType *indices, MatrixValue *values,\n VectorValue *x, VectorValue *gamma)\n{\n const int vectors_per_block = CTASize / VectorSize;\n const int vector_id = threadIdx.x / VectorSize;\n const int lane_id = threadIdx.x % VectorSize;\n\n for (int i = vectors_per_block * blockIdx.x + vector_id; i < rows; i += vectors_per_block * gridDim.x)\n {\n // load start + end pointers\n int row_tmp;\n\n if (lane_id < 2)\n {\n row_tmp = offsets[i + lane_id];\n }\n\n // distribute to all other threads in warp\n int row_begin = utils::shfl(row_tmp, vector_id * VectorSize, warpSize, utils::activemask());\n int row_end = utils::shfl(row_tmp, vector_id * VectorSize + 1, warpSize, utils::activemask());\n\n for (int jj = row_begin + lane_id; utils::any(jj < row_end, utils::activemask()); jj += VectorSize)\n {\n int col = -1;\n VectorValue val = 0.;\n\n if (jj < row_end)\n {\n col = indices[jj];\n val = values[jj];\n utils::atomic_add(&gamma[col], (val * val) * x[i]);\n }\n }\n }\n}\n\n// compute beta = B*y (B = A.^2)\ntemplate \n__global__\nvoid computeBetaDevice(int rows, IndexType *offsets, IndexType *indices, MatrixValue *values,\n VectorValue *y, VectorValue *beta)\n{\n const int vectors_per_block = CTASize / VectorSize;\n const int vector_id = threadIdx.x / VectorSize;\n const int lane_id = threadIdx.x % VectorSize;\n\n for (int i = vectors_per_block * blockIdx.x + vector_id; i < rows; i += vectors_per_block * gridDim.x)\n {\n // load start + end pointers\n int row_tmp;\n\n if (lane_id < 2)\n {\n row_tmp = offsets[i + lane_id];\n }\n\n // distribute to all other threads in warp\n int row_begin = utils::shfl(row_tmp, vector_id * VectorSize, warpSize, utils::activemask());\n int row_end = utils::shfl(row_tmp, vector_id * VectorSize + 1, warpSize, utils::activemask());\n VectorValue bi = 0.;\n\n for (int jj = row_begin + lane_id; utils::any(jj < row_end, utils::activemask()); jj += VectorSize)\n {\n int col = -1;\n VectorValue val = 0.;\n\n if (jj < row_end)\n {\n col = indices[jj];\n val = values[jj];\n bi += (val * val) * y[col];\n }\n }\n\n // reduce over bi\n VectorValue bi_s = warpReduceSum(bi);\n\n if (lane_id == 0)\n {\n beta[i] = bi_s;\n }\n }\n}\n\ntemplate \n__global__\nvoid setOneOverVector(int N, ValueType *x, ValueType sum1, ValueType *beta)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < N; i += gridDim.x * blockDim.x)\n {\n //x[i] = ( isNotCloseToZero(beta[i]) ? sum1 / beta[i] : sum1 / epsilon(beta[i]) );\n x[i] = ( isNotCloseToZero(beta[i]) ? (sum1 / beta[i]) : (ValueType)1. );\n }\n}\n\ntemplate\nstruct square_value\n{\n __host__ __device__ T operator()(const T &x) const\n {\n return x * x;\n }\n};\n\n// functor to generate stddev of vectors\ntemplate \nstruct std_f\n{\n std_f(T x) : v(x) {};\n T v;\n\n __host__ __device__\n T operator()(const T &x1, const T &x2) const\n {\n return (x1 * x2 - v) * (x1 * x2 - v);\n }\n};\n\n// scaled the matrix using diag(F)*A*diag(G), f = sqrt(fabs(x)), g = sqrt(fabs(y))\ntemplate \n__global__\nvoid scaleMatrixDevice(int rows, IndexType *offsets, IndexType *indices, MatrixType *values,\n VectorType *x, VectorType *y)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n VectorType fi = sqrt(fabs(x[i]));\n\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int j = indices[jj];\n VectorType gj = sqrt(fabs(y[j]));\n\n // scale matrix value in place\n if (direction == amgx::SCALE)\n {\n values[jj] *= fi * gj;\n }\n else\n {\n values[jj] /= fi * gj;\n }\n }\n }\n}\n\n\ntemplate \n__global__\nvoid getColRowNorms(int rows, IndexType *offsets, IndexType *indices, MatrixType *values,\n VectorType *rownorms, VectorType *colnorms)\n{\n for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < rows; i += gridDim.x * blockDim.x)\n {\n for (int jj = offsets[i]; jj < offsets[i + 1]; jj++)\n {\n int j = indices[jj];\n VectorType curval = values[jj] * values[jj];\n rownorms[i] += curval;\n utils::atomic_add(colnorms + j, curval);\n }\n }\n}\n\n// vector constant scale operand\ntemplate \nstruct vmul_scale_const\n{\n T _alpha;\n\n vmul_scale_const(T alpha): _alpha(alpha) {};\n\n __host__ __device__\n T operator()(const T &vec) const\n {\n return vec * _alpha;\n }\n};\n\n\n// vector scale operand\ntemplate \nstruct vmul_scale\n{\n\n vmul_scale() {};\n\n __host__ __device__\n T operator()(const T &vec, const T &alpha) const\n {\n return (vec * sqrt(fabs(alpha)));\n }\n};\n\n// vector unscale operand\ntemplate \nstruct vmul_unscale\n{\n\n vmul_unscale() {};\n\n __host__ __device__\n T operator()(const T &vec, const T &alpha) const\n {\n return (vec / sqrt(fabs(alpha)));\n }\n};\n\n\n// Setup on Device\ntemplate \nvoid NBinormalizationScaler >::setup(Matrix_d &A)\n{\n if (A.is_matrix_distributed())\n {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // move these out to config parameters?\n const int max_iters = 50;\n const ValueTypeB tolerance = 1e-10;\n int rows = A.get_num_rows(), cols = A.get_num_cols();\n // temporary vectors\n VVector x(rows, 1), y(cols, 1), beta(rows, 0), gamma(cols, 0);\n // perform matvecs to get beta and gamma (spmv for beta, spmvT for gamma)\n computeBetaGammaDevice<256, 8, 32> <<< 4096, 256>>>(rows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(),\n x.raw(), y.raw(), beta.raw(), gamma.raw());\n cudaCheckError();\n ValueTypeB sum1 = cols, sum2 = rows, std1, std2;\n // calculate initial std1 and std2\n amgx::thrust::device_ptr x_ptr(x.raw()), y_ptr(y.raw()), beta_ptr(beta.raw()), gamma_ptr(gamma.raw());\n std1 = sqrt(amgx::thrust::inner_product(x_ptr, x_ptr + rows, beta_ptr, ValueTypeB(0.), amgx::thrust::plus(), std_f(sum1)) / rows) / sum1;\n std2 = sqrt(amgx::thrust::inner_product(y_ptr, y_ptr + cols, gamma_ptr, ValueTypeB(0.), thrust::plus(), std_f(sum2)) / cols) / sum2;\n ValueTypeB std = sqrt(std1 * std1 + std2 * std2);\n\n for (int t = 0; t < max_iters; t++)\n {\n if (std < tolerance) { break; } // finished\n\n // x = sum1 ./ beta\n setOneOverVector <<< 4096, 256>>>(rows, x.raw(), sum1, beta.raw());\n cudaCheckError();\n // gamma = C*x := B'*x\n thrust_wrapper::fill(gamma.begin(), gamma.end(), ValueTypeB(0.));\n computeGammaDevice<256, 8> <<< 4096, 256>>>(rows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), x.raw(), gamma.raw());\n cudaCheckError();\n // gamma = 1 ./ beta\n setOneOverVector <<< 4096, 256>>>(cols, y.raw(), sum2, gamma.raw());\n cudaCheckError();\n // beta = B*y\n computeBetaDevice<256, 8> <<< 4096, 256>>>(rows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), y.raw(), beta.raw());\n cudaCheckError();\n //ValueTypeB std_old = std;\n std = sqrt(amgx::thrust::inner_product(x_ptr, x_ptr + rows, beta_ptr, ValueTypeB(0.), amgx::thrust::plus(), std_f(sum1)) / rows) / sum1;\n // print it #, current error, convergence rate\n //printf(\"ITER: %d %.3e %.4lg\\n\",t, std, std / std_old);\n }\n\n //Save scaling vectors for later user, setup complete\n left_scale = VVector(beta);\n right_scale = VVector(gamma);\n this->scaled_before = false;\n}\n\n\n// Matrix Scaling on Device\ntemplate \nvoid NBinormalizationScaler >::scaleMatrix(Matrix_d &A, ScaleDirection scaleOrUnscale)\n{\n if (left_scale.size() != A.get_num_rows())\n {\n FatalError(\"Must call setup(A) before binormalization scaling can scale matrix\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (A.is_matrix_distributed())\n {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n int nrows = A.get_num_rows();\n\n if (scaleOrUnscale == amgx::SCALE)\n {\n /*VVector rownorms(nrows, 0.0);\n VVector colnorms(nrows, 0.0);\n getColRowNorms<<<4096,256>>>(nrows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), rownorms.raw(), colnorms.raw());\n cudaCheckError();\n ValueTypeB row_max = *(amgx::thrust::max_element(rownorms.begin(), rownorms.end()));\n ValueTypeB row_min = *(amgx::thrust::min_element(rownorms.begin(), rownorms.end()));\n ValueTypeB col_max = *(amgx::thrust::max_element(colnorms.begin(), colnorms.end()));\n ValueTypeB col_min = *(amgx::thrust::min_element(colnorms.begin(), colnorms.end()));\n cudaCheckError();\n printf(\"Original Matrix: rowmax: %e, rowmin: %e, colmax: %e, colmin: %e\\n\", row_max, row_min, col_max, col_min);fflush(stdout);*/\n // A_scaled = F*A*G (f = diag(F) = sqrt(fabs(x)), g = diag(G) = sqrt(fabs(y))\n // A_ij = f_i * A_ij * g_j\n scaleMatrixDevice <<< 4096, 256>>>(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), left_scale.raw(), right_scale.raw());\n cudaCheckError();\n\n if (!scaled_before)\n {\n this->norm_coef = sqrt(thrust_wrapper::transform_reduce(A.values.begin(), A.values.begin() + A.get_num_nz() * A.get_block_size(), square_value(), 0., amgx::thrust::plus()) / A.get_num_rows());\n cudaCheckError();\n thrust_wrapper::transform(A.values.begin(), A.values.begin() + A.get_num_nz()*A.get_block_size(), A.values.begin(), vmul_scale_const(1. / this->norm_coef) );\n thrust_wrapper::transform(left_scale.begin(), left_scale.end(), left_scale.begin(), vmul_scale_const(sqrt(1. / this->norm_coef)) );\n thrust_wrapper::transform(right_scale.begin(), right_scale.end(), right_scale.begin(), vmul_scale_const(sqrt(1. / this->norm_coef)) );\n cudaCheckError();\n /*thrust_wrapper::fill(rownorms.begin(), rownorms.end(), 0.);\n thrust_wrapper::fill(colnorms.begin(), colnorms.end(), 0.);\n getColRowNorms<<<4096,256>>>(nrows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), rownorms.raw(), colnorms.raw());\n cudaCheckError();\n row_max = *(amgx::thrust::max_element(rownorms.begin(), rownorms.end()));\n row_min = *(amgx::thrust::min_element(rownorms.begin(), rownorms.end()));\n col_max = *(amgx::thrust::max_element(colnorms.begin(), colnorms.end()));\n col_min = *(amgx::thrust::min_element(colnorms.begin(), colnorms.end()));\n cudaCheckError();\n printf(\"Scaled Matrix: rowmax: %e, rowmin: %e, colmax: %e, colmin: %e\\n\", row_max, row_min, col_max, col_min);fflush(stdout);*/\n }\n\n this->scaled_before = true;\n }\n else\n {\n scaleMatrixDevice <<< 4096, 256>>>(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), left_scale.raw(), right_scale.raw());\n cudaCheckError();\n }\n}\n\n// Setup on Host\ntemplate \nvoid NBinormalizationScaler >::setup(Matrix_h &A)\n{\n if (A.is_matrix_distributed())\n {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // move these out to config parameters?\n const int max_iters = 10;\n const ValueTypeB tolerance = 1e-10;\n int rows = A.get_num_rows(), cols = A.get_num_cols();\n // temporary vectors\n VVector x(rows, 1), y(cols, 1), beta(rows, 0), gamma(cols, 0);\n // perform matvecs to get beta and gamma (spmv for beta, spmvT for gamma)\n computeBetaGammaHost(rows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(),\n x.raw(), y.raw(), beta.raw(), gamma.raw());\n double std1 = 0., std2 = 0., sum1 = cols, sum2 = rows;\n\n // calculate initial std1 and std2\n for (int i = 0; i < rows; i++)\n {\n std1 += pow(x[i] * beta[i] - sum1, 2.0);\n }\n\n std1 = sqrt(std1 / rows) / sum1;\n\n for (int i = 0; i < cols; i++)\n {\n std2 += pow(y[i] * gamma[i] - sum2, 2.0);\n }\n\n std2 = sqrt(std2 / cols) / sum2;\n //printf(\"std1: %lg, std2: %lg\\n\",std1, std2);\n double std_initial = sqrt((std1 * std1) + (std2 * std2));\n double std = std_initial;\n\n for (int t = 0; t < max_iters; t++)\n {\n if (std < tolerance) { break; } // finished\n\n // x = sum1 ./ beta\n for (int i = 0; i < rows; i++) { x[i] = ( isNotCloseToZero(beta[i]) ? sum1 / beta[i] : sum1 / epsilon(beta[i]) ); }\n\n // gamma = C*x\n computeGammaHost(rows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), x.raw(), gamma.raw());\n\n // gamma = 1 ./ beta\n for (int i = 0; i < cols; i++) { y[i] = ( isNotCloseToZero(gamma[i]) ? sum2 / gamma[i] : sum2 / epsilon(gamma[i]) ); }\n\n // beta = B*y\n computeBetaHost(rows, A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), y.raw(), beta.raw());\n //ValueTypeB std_old = std;\n std = 0.;\n\n for (int i = 0; i < rows; i++)\n {\n std += pow(x[i] * beta[i] - sum1, 2.0);\n }\n\n std = sqrt(std / rows) / sum1;\n // print it #, current error, convergence rate\n //printf(\"ITER: %d %.3e %.4lg\\n\",t, std, std / std_old);\n }\n\n //Save scaling vectors for later user, setup complete\n left_scale = VVector(beta);\n right_scale = VVector(gamma);\n}\n\n\ntemplate \nvoid NBinormalizationScaler >::scaleMatrix(Matrix_h &A, ScaleDirection scaleOrUnscale)\n{\n if (A.is_matrix_distributed())\n {\n FatalError(\"Binormalization scaling not supported for distributed matrices\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (left_scale.size() != A.get_num_rows())\n {\n FatalError(\"Must call setup(A) before binormalization scaling can scale matrix\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // A_scaled = F*A*G (f = diag(F) = sqrt(fabs(x)), g = diag(G) = sqrt(fabs(y))\n // A_ij = f_i * A_ij * g_j\n scaleMatrixHost(A.get_num_rows(), A.row_offsets.raw(), A.col_indices.raw(), A.values.raw(), left_scale.raw(), right_scale.raw());\n}\n\n\ntemplate \nvoid NBinormalizationScaler >::scaleVector(VVector &v, ScaleDirection scaleOrUnscale, ScaleSide leftOrRight)\n{\n VVector *scale_vector = (leftOrRight == amgx::LEFT) ? &this->left_scale : &this->right_scale;\n\n //thrust_wrapper::transform(v.begin(), v.end(), scale_vector->begin(), v.begin(), (scaleOrUnscale == amgx::SCALE) ? vmul_scale() : vmul_unscale() );\n if (scaleOrUnscale == amgx::SCALE)\n {\n thrust_wrapper::transform(v.begin(), v.end(), scale_vector->begin(), v.begin(), vmul_scale() );\n }\n else\n {\n thrust_wrapper::transform(v.begin(), v.end(), scale_vector->begin(), v.begin(), vmul_unscale() );\n }\n}\n\ntemplate \nvoid NBinormalizationScaler >::scaleVector(VVector &v, ScaleDirection scaleOrUnscale, ScaleSide leftOrRight)\n{\n FatalError(\"4x4 block size not supported\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class NBinormalizationScaler_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class NBinormalizationScaler::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n\n} // namespace amgx\n\n", "source": "/kaggle/working/amgx/src/scalers/nbinormalization.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"unit_test.h\"\n#include \"vector.h\"\n#include \"test_utils.h\"\n#include \"util.h\"\n#include \"time.h\"\n\nnamespace amgx\n\n{\n\ntemplate\nstatic __global__\nvoid set_raw(T *v, int N)\n{\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < N)\n {\n v[i] = i;\n }\n}\ntemplate\nstatic __global__\nvoid set_flat(T v, int N)\n{\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < N)\n {\n v[i] = i;\n }\n}\n\ntemplate\nstatic __global__\nvoid set_block(T v, int K, int I, int J)\n{\n for (int k = threadIdx.x; k < K; k += blockDim.x)\n for (int i = 0; i < I; i++)\n for (int j = 0; j < J; j++)\n {\n v(k, i, j) = k * I * J + i * J + j;\n }\n}\n\n// parameter is used as test name\nDECLARE_UNITTEST_BEGIN(VectorTests);\n\n\n\nvoid run()\n{\n typedef typename Vector_h::index_type index_type;\n typedef typename Vector_h::value_type value_type;\n Vector_h h_v1;\n Vector_h h_v2;\n Vector_d d_v1;\n Vector_d d_v2;\n UNITTEST_ASSERT_TRUE(h_v1.size() == 0);\n UNITTEST_ASSERT_TRUE(h_v1.raw() == 0);\n UNITTEST_ASSERT_TRUE(d_v1.size() == 0);\n UNITTEST_ASSERT_TRUE(d_v1.raw() == 0);\n h_v1.resize(4);\n UNITTEST_ASSERT_TRUE_DESC(\"vector resize\", h_v1.size() == 4);\n h_v1[3] = 1;\n UNITTEST_ASSERT_TRUE_DESC(\"vector assign value\", h_v1[3] == 1);\n h_v1.clear();\n UNITTEST_ASSERT_TRUE_DESC(\"vector clear\", h_v1.empty());\n UNITTEST_ASSERT_TRUE_DESC(\"vector clear\", h_v1.size() == 0);\n d_v1.resize(4);\n UNITTEST_ASSERT_TRUE_DESC(\"vector resize\", d_v1.size() == 4);\n d_v1[3] = 1;\n UNITTEST_ASSERT_TRUE_DESC(\"vector assign value\", d_v1[3] == 1);\n d_v1.clear();\n UNITTEST_ASSERT_TRUE_DESC(\"vector clear\", d_v1.empty());\n UNITTEST_ASSERT_TRUE_DESC(\"vector clear\", d_v1.size() == 0);\n const Vector_d d_one(64, 1);\n const Vector_d d_zero(64, 0);\n const Vector_d h_one(64, 1);\n const Vector_d h_zero(64, 0);\n UNITTEST_ASSERT_TRUE_DESC(\"device vector copy constructor size\", d_one.size() == 64);\n UNITTEST_ASSERT_TRUE_DESC(\"device vector copy constructor value\", d_one[5] == 1);\n UNITTEST_ASSERT_TRUE_DESC(\"host vector copy constructor size\", h_one.size() == 64);\n UNITTEST_ASSERT_TRUE_DESC(\"host vector copy constructor value\", h_one[5] == 1);\n d_v1 = h_one;\n UNITTEST_ASSERT_EQUAL_DESC(\"vector host to device\", d_v1, d_one);\n h_v2 = d_one;\n UNITTEST_ASSERT_EQUAL_DESC(\"vector device to host\", h_v2, h_one);\n d_v1 = d_one;\n d_v2 = d_one;\n h_v1 = h_zero;\n h_v2 = h_zero;\n h_v2.copy_async(d_v1);\n d_v2.copy_async(h_v1);\n h_v2.sync();\n d_v2.sync();\n UNITTEST_ASSERT_EQUAL_DESC(\"vector device to host async\", h_v2, h_one);\n UNITTEST_ASSERT_EQUAL_DESC(\"vector host to device async\", h_v1, d_zero);\n d_v1 = d_zero;\n d_v2 = d_one;\n d_v1.swap(d_v2);\n UNITTEST_ASSERT_EQUAL_DESC(\"vector swap1\", d_v1, d_one);\n UNITTEST_ASSERT_EQUAL_DESC(\"vector swap2\", d_v2, d_zero);\n //set vector to 0,1,2,3,...,N using the raw pointer\n set_raw <<< 1, 64>>>(d_v1.raw(), 64);\n cudaCheckError();\n d_v2 = d_zero;\n d_v2.set_block_dimx(4);\n d_v2.set_block_dimy(4);\n //set vector to 0,1,2,3,...,N using operator[]\n set_flat <<< 1, 64>>>(d_v2.pod(), 64);\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL_DESC(\"block vector operator[]\", d_v1, d_v2);\n d_v2 = d_zero;\n d_v2.set_block_dimx(4);\n d_v2.set_block_dimy(4);\n //set vector to 0,1,2,3,...,N using operator()\n set_block <<< 1, 64>>>(d_v2.pod(), 4, 4, 4);\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL_DESC(\"block vector operator() 4x4\", d_v1, d_v2);\n d_v2 = d_zero;\n d_v2.set_block_dimx(8);\n d_v2.set_block_dimy(2);\n //set vector to 0,1,2,3,...,N using operator()\n set_block <<< 1, 64>>>(d_v2.pod(), 4, 2, 8);\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL_DESC(\"block vector operator() 2x8\", d_v1, d_v2);\n d_v2 = d_zero;\n d_v2.set_block_dimx(2);\n d_v2.set_block_dimy(8);\n //set vector to 0,1,2,3,...,N using operator()\n set_block <<< 1, 64>>>(d_v2.pod(), 4, 8, 2);\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL_DESC(\"block vector operator() 8x2\", d_v1, d_v2);\n}\n\nDECLARE_UNITTEST_END(VectorTests);\n\n\n// if you want to be able run this test for all available configs you can write this:\n#define AMGX_CASE_LINE(CASE) VectorTests ::Type> VectorTests_instance_mode##CASE;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n// or you can specify several desired configs\n//VectorTests ::Type> VectorTests_instance_mode_dDDI;\n\n} //namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \"unit_test.h\"\n#include \"vector.h\"\n#include \"test_utils.h\"\n#include \"util.h\"\n#include \"time.h\"\n\nnamespace amgx\n\n{\n\ntemplate\nstatic __global__\nvoid set_raw(T *v, int N)\n{\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < N)\n {\n v[i] = i;\n }\n}\ntemplate\nstatic __global__\nvoid set_flat(T v, int N)\n{\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (i < N)\n {\n v[i] = i;\n }\n}\n\ntemplate\nstatic __global__\nvoid set_block(T v, int K, int I, int J)\n{\n for (int k = threadIdx.x; k < K; k += blockDim.x)\n for (int i = 0; i < I; i++)\n for (int j = 0; j < J; j++)\n {\n v(k, i, j) = k * I * J + i * J + j;\n }\n}\n\n// parameter is used as test name\nDECLARE_UNITTEST_BEGIN(VectorTests);\n\n\n\nvoid run()\n{\n typedef typename Vector_h::index_type index_type;\n typedef typename Vector_h::value_type value_type;\n Vector_h h_v1;\n Vector_h h_v2;\n Vector_d d_v1;\n Vector_d d_v2;\n UNITTEST_ASSERT_TRUE(h_v1.size() == 0);\n UNITTEST_ASSERT_TRUE(h_v1.raw() == 0);\n UNITTEST_ASSERT_TRUE(d_v1.size() == 0);\n UNITTEST_ASSERT_TRUE(d_v1.raw() == 0);\n h_v1.resize(4);\n UNITTEST_ASSERT_TRUE_DESC(\"vector resize\", h_v1.size() == 4);\n h_v1[3] = 1;\n UNITTEST_ASSERT_TRUE_DESC(\"vector assign value\", h_v1[3] == 1);\n h_v1.clear();\n UNITTEST_ASSERT_TRUE_DESC(\"vector clear\", h_v1.empty());\n UNITTEST_ASSERT_TRUE_DESC(\"vector clear\", h_v1.size() == 0);\n d_v1.resize(4);\n UNITTEST_ASSERT_TRUE_DESC(\"vector resize\", d_v1.size() == 4);\n d_v1[3] = 1;\n UNITTEST_ASSERT_TRUE_DESC(\"vector assign value\", d_v1[3] == 1);\n d_v1.clear();\n UNITTEST_ASSERT_TRUE_DESC(\"vector clear\", d_v1.empty());\n UNITTEST_ASSERT_TRUE_DESC(\"vector clear\", d_v1.size() == 0);\n const Vector_d d_one(64, 1);\n const Vector_d d_zero(64, 0);\n const Vector_d h_one(64, 1);\n const Vector_d h_zero(64, 0);\n UNITTEST_ASSERT_TRUE_DESC(\"device vector copy constructor size\", d_one.size() == 64);\n UNITTEST_ASSERT_TRUE_DESC(\"device vector copy constructor value\", d_one[5] == 1);\n UNITTEST_ASSERT_TRUE_DESC(\"host vector copy constructor size\", h_one.size() == 64);\n UNITTEST_ASSERT_TRUE_DESC(\"host vector copy constructor value\", h_one[5] == 1);\n d_v1 = h_one;\n UNITTEST_ASSERT_EQUAL_DESC(\"vector host to device\", d_v1, d_one);\n h_v2 = d_one;\n UNITTEST_ASSERT_EQUAL_DESC(\"vector device to host\", h_v2, h_one);\n d_v1 = d_one;\n d_v2 = d_one;\n h_v1 = h_zero;\n h_v2 = h_zero;\n h_v2.copy_async(d_v1);\n d_v2.copy_async(h_v1);\n h_v2.sync();\n d_v2.sync();\n UNITTEST_ASSERT_EQUAL_DESC(\"vector device to host async\", h_v2, h_one);\n UNITTEST_ASSERT_EQUAL_DESC(\"vector host to device async\", h_v1, d_zero);\n d_v1 = d_zero;\n d_v2 = d_one;\n d_v1.swap(d_v2);\n UNITTEST_ASSERT_EQUAL_DESC(\"vector swap1\", d_v1, d_one);\n UNITTEST_ASSERT_EQUAL_DESC(\"vector swap2\", d_v2, d_zero);\n //set vector to 0,1,2,3,...,N using the raw pointer\n set_raw <<< 1, 64>>>(d_v1.raw(), 64);\n cudaCheckError();\n d_v2 = d_zero;\n d_v2.set_block_dimx(4);\n d_v2.set_block_dimy(4);\n //set vector to 0,1,2,3,...,N using operator[]\n set_flat <<< 1, 64>>>(d_v2.pod(), 64);\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL_DESC(\"block vector operator[]\", d_v1, d_v2);\n d_v2 = d_zero;\n d_v2.set_block_dimx(4);\n d_v2.set_block_dimy(4);\n //set vector to 0,1,2,3,...,N using operator()\n set_block <<< 1, 64>>>(d_v2.pod(), 4, 4, 4);\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL_DESC(\"block vector operator() 4x4\", d_v1, d_v2);\n d_v2 = d_zero;\n d_v2.set_block_dimx(8);\n d_v2.set_block_dimy(2);\n //set vector to 0,1,2,3,...,N using operator()\n set_block <<< 1, 64>>>(d_v2.pod(), 4, 2, 8);\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL_DESC(\"block vector operator() 2x8\", d_v1, d_v2);\n d_v2 = d_zero;\n d_v2.set_block_dimx(2);\n d_v2.set_block_dimy(8);\n //set vector to 0,1,2,3,...,N using operator()\n set_block <<< 1, 64>>>(d_v2.pod(), 4, 8, 2);\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL_DESC(\"block vector operator() 8x2\", d_v1, d_v2);\n}\n\nDECLARE_UNITTEST_END(VectorTests);\n\n\n// if you want to be able run this test for all available configs you can write this:\n#define AMGX_CASE_LINE(CASE) VectorTests ::Type> VectorTests_instance_mode##CASE;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n// or you can specify several desired configs\n//VectorTests ::Type> VectorTests_instance_mode_dDDI;\n\n} //namespace amgx\n", "source": "/kaggle/working/amgx/src/tests/vector_tests.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \"unit_test.h\"\n#include \"matrix.h\"\n#include \"blas.h\"\n#include \"multiply.h\"\n#include \"solvers/dense_lu_solver.h\"\n\nnamespace amgx\n{\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename T >\n__global__ void split_l_and_u(int n, const T *lu, int lda, T *l, T *u )\n{\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n int j = blockIdx.y * blockDim.y + threadIdx.y;\n\n if (i >= n || j >= n)\n {\n return;\n }\n\n T l_ij = i == j ? T(1) : T(0), u_ij = T(0);\n T lu_ij = lu[i * lda + j];\n\n if (i <= j)\n {\n u_ij = lu_ij;\n }\n else\n {\n l_ij = lu_ij;\n }\n\n l[i * lda + j] = l_ij;\n u[i * lda + j] = u_ij;\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDECLARE_UNITTEST_BEGIN(DenseLUSolverTest_Base);\n\nstd::string base_keywords()\n{\n return \"dense_lu\";\n}\n\ntemplate< typename Matrix >\nvoid make_identity(Matrix &A)\n{\n typedef typename Matrix::TConfig Matrix_config;\n typedef typename Matrix_config::template setMemSpace::Type Config_h;\n typedef typename Config_h::template setVecPrec::Type IVector_config_h;\n typedef Vector FVector_h;\n typedef Vector IVector_h;\n const int num_rows = A.get_num_rows();\n IVector_h row_offsets(num_rows + 1), col_indices(num_rows);\n\n for ( int i = 0 ; i < num_rows ; ++i )\n {\n row_offsets[i] = i;\n col_indices[i] = i;\n }\n\n row_offsets.back() = num_rows;\n FVector_h values(num_rows, 1.0);\n A.row_offsets.copy(row_offsets);\n A.col_indices.copy(col_indices);\n A.values. copy(values);\n}\n\ntemplate< typename Matrix_h, typename Matrix_data >\nvoid csr_to_dense(const Matrix_h &A_h, Matrix_data *dense_A_h, int lda)\n{\n for ( int i = 0 ; i < A_h.get_num_rows() ; ++i )\n for ( int j = A_h.row_offsets[i] ; j < A_h.row_offsets[i + 1] ; ++j )\n {\n dense_A_h[i * lda + A_h.col_indices[j]] = A_h.values[j];\n }\n}\n\ntemplate< typename Matrix_data >\nvoid l_times_u(int n, Matrix_data *lu_d, int lda)\n{\n Matrix_data *l_d, *u_d;\n amgx::memory::cudaMallocAsync((void **) &l_d, n * lda * sizeof(Matrix_data), 0);\n UNITTEST_ASSERT_EQUAL(cudaGetLastError(), cudaSuccess);\n UNITTEST_ASSERT_EQUAL(cudaStreamSynchronize(0), cudaSuccess);\n cudaCheckError();\n amgx::memory::cudaMallocAsync((void **) &u_d, n * lda * sizeof(Matrix_data), 0);\n UNITTEST_ASSERT_EQUAL(cudaGetLastError(), cudaSuccess);\n UNITTEST_ASSERT_EQUAL(cudaStreamSynchronize(0), cudaSuccess);\n cudaCheckError();\n // Split LU.\n dim3 block_dim(16, 16);\n dim3 grid_dim((n + block_dim.x - 1) / block_dim.x, (n + block_dim.y - 1) / block_dim.y);\n split_l_and_u <<< grid_dim, block_dim>>>(n, lu_d, lda, l_d, u_d);\n cudaCheckError();\n cudaError_t status = cudaDeviceSynchronize();\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL(status, cudaSuccess);\n // LxU = LU.\n Matrix_data one(1), zero(0);\n cublasHandle_t cublas_handle;\n cublasStatus_t cublas_status;\n cublas_status = cublasCreate(&cublas_handle);\n UNITTEST_ASSERT_EQUAL(cublas_status, CUBLAS_STATUS_SUCCESS);\n cublas_status = cublasGemm(cublas_handle,\n CUBLAS_OP_N,\n CUBLAS_OP_N,\n n, n, n,\n &one, l_d, lda,\n u_d, lda,\n &zero, lu_d, lda);\n UNITTEST_ASSERT_EQUAL(cublas_status, CUBLAS_STATUS_SUCCESS);\n cublas_status = cublasDestroy(cublas_handle);\n UNITTEST_ASSERT_EQUAL(cublas_status, CUBLAS_STATUS_SUCCESS);\n cudaFreeAsync(l_d, 0);\n UNITTEST_ASSERT_EQUAL(cudaGetLastError(), cudaSuccess);\n UNITTEST_ASSERT_EQUAL(cudaStreamSynchronize(0), cudaSuccess);\n cudaCheckError();\n cudaFreeAsync(u_d, 0);\n UNITTEST_ASSERT_EQUAL(cudaGetLastError(), cudaSuccess);\n UNITTEST_ASSERT_EQUAL(cudaStreamSynchronize(0), cudaSuccess);\n cudaCheckError();\n}\n\ntemplate< typename Matrix, typename Matrix_data >\nvoid check_lu_product(const Matrix &A_h, int n, Matrix_data *lu_d, int lda)\n{\n // Compute LxU.\n l_times_u(n, lu_d, lda);\n // Copy LxU to the host.\n Matrix_data *lu_h = new Matrix_data[n * lda];\n cudaMemcpy(lu_h, lu_d, n * lda * sizeof(Matrix_data), cudaMemcpyDeviceToHost);\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL(cudaGetLastError(), cudaSuccess);\n // Make sure LxU equals A.\n Matrix_data *dense_A_h = new Matrix_data[n * lda];\n std::memset(dense_A_h, 0x0, n * lda * sizeof(Matrix_data));\n csr_to_dense(A_h, dense_A_h, lda);\n\n // Compare the matrices.\n for ( int i = 0 ; i < n ; ++i )\n for ( int j = 0 ; j < n ; ++j )\n {\n UNITTEST_ASSERT_EQUAL(lu_h[i * lda + j], dense_A_h[i * lda + j]);\n }\n\n delete[] dense_A_h;\n delete[] lu_h;\n}\n\nDECLARE_UNITTEST_END(DenseLUSolverTest_Base);\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDECLARE_UNITTEST_BEGIN_EXTD(DenseLUSolverTest_Factorization_Id_32, DenseLUSolverTest_Base);\n\nvoid run()\n{\n // Make sure the solver produces A=LU where A is the identity matrix.\n typedef typename T_Config::template setMemSpace::Type Config_h;\n typedef Vector FVector_h;\n typedef typename TConfig::MatPrec Matrix_data;\n const int N = 32;\n Matrix A(N, N, N, CSR);\n A.set_initialized(0);\n this->make_identity(A);\n A.set_initialized(1);\n AMG_Config cfg;\n dense_lu_solver::DenseLUSolver solver(cfg, \"\", NULL);\n solver.setup(A, false);\n FVector_h m(N * N);\n cudaMemcpy(m.raw(), solver.get_dense_A(), N * N * sizeof(Matrix_data), cudaMemcpyDeviceToHost);\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL(cudaGetLastError(), cudaSuccess);\n\n for ( int i = 0 ; i < N ; ++i )\n for ( int j = 0 ; j < N ; ++j )\n {\n UNITTEST_ASSERT_EQUAL(m[i * solver.get_lda() + j], i == j ? Matrix_data(1) : Matrix_data(0));\n }\n}\n\nDECLARE_UNITTEST_END(DenseLUSolverTest_Factorization_Id_32)\n\nDenseLUSolverTest_Factorization_Id_32::Type> DenseLUSolverTest_Factorization_Id_32_dDDI;\nDenseLUSolverTest_Factorization_Id_32::Type> DenseLUSolverTest_Factorization_Id_32_dFFI;\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDECLARE_UNITTEST_BEGIN_EXTD(DenseLUSolverTest_Factorization_Id_256, DenseLUSolverTest_Base);\n\nvoid run()\n{\n // Make sure the solver produces A=LU where A is the identity matrix.\n typedef typename T_Config::template setMemSpace::Type Config_h;\n typedef Vector FVector_h;\n typedef typename TConfig::MatPrec Matrix_data;\n const int N = 256;\n Matrix A(N, N, N, CSR);\n A.set_initialized(0);\n this->make_identity(A);\n A.set_initialized(1);\n AMG_Config cfg;\n dense_lu_solver::DenseLUSolver solver(cfg, \"\", NULL);\n solver.setup(A, false);\n FVector_h m(N * N);\n cudaMemcpy(m.raw(), solver.get_dense_A(), N * N * sizeof(Matrix_data), cudaMemcpyDeviceToHost);\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL(cudaGetLastError(), cudaSuccess);\n\n for ( int i = 0 ; i < N ; ++i )\n for ( int j = 0 ; j < N ; ++j )\n {\n UNITTEST_ASSERT_EQUAL(m[i * solver.get_lda() + j], i == j ? Matrix_data(1) : Matrix_data(0));\n }\n}\n\nDECLARE_UNITTEST_END(DenseLUSolverTest_Factorization_Id_256)\n\nDenseLUSolverTest_Factorization_Id_256::Type> DenseLUSolverTest_Factorization_Id_256_dDDI;\nDenseLUSolverTest_Factorization_Id_256::Type> DenseLUSolverTest_Factorization_Id_256_dFFI;\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDECLARE_UNITTEST_BEGIN_EXTD(DenseLUSolverTest_Solve_Poisson3D, DenseLUSolverTest_Base);\n\nvoid run()\n{\n typedef typename T_Config::MatPrec Matrix_data;\n typedef typename T_Config::template setMemSpace::Type Config_d;\n typedef typename T_Config::template setMemSpace ::Type Config_h;\n typedef Matrix Matrix_d;\n typedef Vector Vector_d;\n typedef Matrix Matrix_h;\n typedef Vector Vector_h;\n Matrix_h A_h;\n Vector_h x_h, b_h;\n \n A_h.set_initialized(0);\n generatePoissonForTest(A_h, 1, 0, 27, 16, 16, 16);\n A_h.set_initialized(1);\n \n AMG_Config cfg;\n cfg.parseParameterString(\"monitor_residual=1\");\n Matrix_d A_d(A_h);\n dense_lu_solver::DenseLUSolver solver(cfg, \"default\", NULL);\n solver.setup(A_d, false);\n const int n = A_h.get_num_rows();\n Vector_d b_d(n), x_d(n), r_d(n);\n thrust_wrapper::fill(b_d.begin(), b_d.end(), Matrix_data(1));\n solver.solve(b_d, x_d, false);\n solver.compute_residual(b_d, x_d, r_d);\n Vector_h resid_nrm(1);\n solver.compute_norm(r_d, resid_nrm);\n Vector_h rhs_nrm(1);\n solver.compute_norm(b_d, rhs_nrm);\n double relError = resid_nrm[0] / rhs_nrm[0];\n\n if ( T_Config::matPrec == AMGX_matDouble )\n {\n UNITTEST_ASSERT_EQUAL_TOL(resid_nrm[0], 0.0, 1.0e-12);\n }\n else\n {\n UNITTEST_ASSERT_EQUAL_TOL(resid_nrm[0], 0.0f, 1.0e-6f);\n }\n}\n\nDECLARE_UNITTEST_END(DenseLUSolverTest_Solve_Poisson3D)\n\nDenseLUSolverTest_Solve_Poisson3D::Type> DenseLUSolverTest_Solve_Poisson3D_dDDI;\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDECLARE_UNITTEST_BEGIN_EXTD(DenseLUSolverTest_Solve_Id_32, DenseLUSolverTest_Base);\n\nvoid run()\n{\n typedef typename T_Config::template setMemSpace::Type Config_h;\n typedef Vector FVector_h;\n typedef typename TConfig::MatPrec Matrix_data;\n const int N = 32;\n Matrix A(N, N, N, CSR);\n A.set_initialized(0);\n this->make_identity(A);\n A.set_initialized(1);\n AMG_Config cfg;\n dense_lu_solver::DenseLUSolver solver(cfg, \"\", NULL);\n solver.setup(A, false);\n FVector_h b_h(N);\n\n for ( int i = 0 ; i < N ; ++i )\n {\n b_h[i] = Matrix_data(rand()) / RAND_MAX;\n }\n\n Vector b(b_h), x(N), r(N);\n solver.solve(b, x, false);\n UNITTEST_ASSERT_EQUAL(x, b);\n}\n\nDECLARE_UNITTEST_END(DenseLUSolverTest_Solve_Id_32)\n\nDenseLUSolverTest_Solve_Id_32::Type> DenseLUSolverTest_Solve_Id_32_dDDI;\nDenseLUSolverTest_Solve_Id_32::Type> DenseLUSolverTest_Solve_Id_32_dFFI;\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDECLARE_UNITTEST_BEGIN_EXTD(DenseLUSolverTest_Solve_Id_256, DenseLUSolverTest_Base);\n\nvoid run()\n{\n typedef typename T_Config::template setMemSpace::Type Config_h;\n typedef Vector FVector_h;\n typedef typename TConfig::MatPrec Matrix_data;\n const int N = 256;\n Matrix A(N, N, N, CSR);\n A.set_initialized(0);\n this->make_identity(A);\n A.set_initialized(1);\n AMG_Config cfg;\n dense_lu_solver::DenseLUSolver solver(cfg, \"\", NULL);\n solver.setup(A, false);\n FVector_h b_h(N);\n\n for ( int i = 0 ; i < N ; ++i )\n {\n b_h[i] = Matrix_data(rand()) / RAND_MAX;\n }\n\n Vector b(b_h), x(N);\n solver.solve(b, x, false);\n UNITTEST_ASSERT_EQUAL(x, b);\n}\n\nDECLARE_UNITTEST_END(DenseLUSolverTest_Solve_Id_256)\n\nDenseLUSolverTest_Solve_Id_256::Type> DenseLUSolverTest_Solve_Id_256_dDDI;\nDenseLUSolverTest_Solve_Id_256::Type> DenseLUSolverTest_Solve_Id_256_dFFI;\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \"unit_test.h\"\n#include \"matrix.h\"\n#include \"blas.h\"\n#include \"multiply.h\"\n#include \"solvers/dense_lu_solver.h\"\n\nnamespace amgx\n{\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename T >\n__global__ void split_l_and_u(int n, const T *lu, int lda, T *l, T *u )\n{\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n int j = blockIdx.y * blockDim.y + threadIdx.y;\n\n if (i >= n || j >= n)\n {\n return;\n }\n\n T l_ij = i == j ? T(1) : T(0), u_ij = T(0);\n T lu_ij = lu[i * lda + j];\n\n if (i <= j)\n {\n u_ij = lu_ij;\n }\n else\n {\n l_ij = lu_ij;\n }\n\n l[i * lda + j] = l_ij;\n u[i * lda + j] = u_ij;\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDECLARE_UNITTEST_BEGIN(DenseLUSolverTest_Base);\n\nstd::string base_keywords()\n{\n return \"dense_lu\";\n}\n\ntemplate< typename Matrix >\nvoid make_identity(Matrix &A)\n{\n typedef typename Matrix::TConfig Matrix_config;\n typedef typename Matrix_config::template setMemSpace::Type Config_h;\n typedef typename Config_h::template setVecPrec::Type IVector_config_h;\n typedef Vector FVector_h;\n typedef Vector IVector_h;\n const int num_rows = A.get_num_rows();\n IVector_h row_offsets(num_rows + 1), col_indices(num_rows);\n\n for ( int i = 0 ; i < num_rows ; ++i )\n {\n row_offsets[i] = i;\n col_indices[i] = i;\n }\n\n row_offsets.back() = num_rows;\n FVector_h values(num_rows, 1.0);\n A.row_offsets.copy(row_offsets);\n A.col_indices.copy(col_indices);\n A.values. copy(values);\n}\n\ntemplate< typename Matrix_h, typename Matrix_data >\nvoid csr_to_dense(const Matrix_h &A_h, Matrix_data *dense_A_h, int lda)\n{\n for ( int i = 0 ; i < A_h.get_num_rows() ; ++i )\n for ( int j = A_h.row_offsets[i] ; j < A_h.row_offsets[i + 1] ; ++j )\n {\n dense_A_h[i * lda + A_h.col_indices[j]] = A_h.values[j];\n }\n}\n\ntemplate< typename Matrix_data >\nvoid l_times_u(int n, Matrix_data *lu_d, int lda)\n{\n Matrix_data *l_d, *u_d;\n amgx::memory::hipMallocAsync((void **) &l_d, n * lda * sizeof(Matrix_data), 0);\n UNITTEST_ASSERT_EQUAL(hipGetLastError(), hipSuccess);\n UNITTEST_ASSERT_EQUAL(hipStreamSynchronize(0), hipSuccess);\n cudaCheckError();\n amgx::memory::hipMallocAsync((void **) &u_d, n * lda * sizeof(Matrix_data), 0);\n UNITTEST_ASSERT_EQUAL(hipGetLastError(), hipSuccess);\n UNITTEST_ASSERT_EQUAL(hipStreamSynchronize(0), hipSuccess);\n cudaCheckError();\n // Split LU.\n dim3 block_dim(16, 16);\n dim3 grid_dim((n + block_dim.x - 1) / block_dim.x, (n + block_dim.y - 1) / block_dim.y);\n split_l_and_u <<< grid_dim, block_dim>>>(n, lu_d, lda, l_d, u_d);\n cudaCheckError();\n hipError_t status = hipDeviceSynchronize();\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL(status, hipSuccess);\n // LxU = LU.\n Matrix_data one(1), zero(0);\n hipblasHandle_t cublas_handle;\n hipblasStatus_t cublas_status;\n cublas_status = hipblasCreate(&cublas_handle);\n UNITTEST_ASSERT_EQUAL(cublas_status, HIPBLAS_STATUS_SUCCESS);\n cublas_status = cublasGemm(cublas_handle,\n HIPBLAS_OP_N,\n HIPBLAS_OP_N,\n n, n, n,\n &one, l_d, lda,\n u_d, lda,\n &zero, lu_d, lda);\n UNITTEST_ASSERT_EQUAL(cublas_status, HIPBLAS_STATUS_SUCCESS);\n cublas_status = hipblasDestroy(cublas_handle);\n UNITTEST_ASSERT_EQUAL(cublas_status, HIPBLAS_STATUS_SUCCESS);\n hipFreeAsync(l_d, 0);\n UNITTEST_ASSERT_EQUAL(hipGetLastError(), hipSuccess);\n UNITTEST_ASSERT_EQUAL(hipStreamSynchronize(0), hipSuccess);\n cudaCheckError();\n hipFreeAsync(u_d, 0);\n UNITTEST_ASSERT_EQUAL(hipGetLastError(), hipSuccess);\n UNITTEST_ASSERT_EQUAL(hipStreamSynchronize(0), hipSuccess);\n cudaCheckError();\n}\n\ntemplate< typename Matrix, typename Matrix_data >\nvoid check_lu_product(const Matrix &A_h, int n, Matrix_data *lu_d, int lda)\n{\n // Compute LxU.\n l_times_u(n, lu_d, lda);\n // Copy LxU to the host.\n Matrix_data *lu_h = new Matrix_data[n * lda];\n hipMemcpy(lu_h, lu_d, n * lda * sizeof(Matrix_data), hipMemcpyDeviceToHost);\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL(hipGetLastError(), hipSuccess);\n // Make sure LxU equals A.\n Matrix_data *dense_A_h = new Matrix_data[n * lda];\n std::memset(dense_A_h, 0x0, n * lda * sizeof(Matrix_data));\n csr_to_dense(A_h, dense_A_h, lda);\n\n // Compare the matrices.\n for ( int i = 0 ; i < n ; ++i )\n for ( int j = 0 ; j < n ; ++j )\n {\n UNITTEST_ASSERT_EQUAL(lu_h[i * lda + j], dense_A_h[i * lda + j]);\n }\n\n delete[] dense_A_h;\n delete[] lu_h;\n}\n\nDECLARE_UNITTEST_END(DenseLUSolverTest_Base);\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDECLARE_UNITTEST_BEGIN_EXTD(DenseLUSolverTest_Factorization_Id_32, DenseLUSolverTest_Base);\n\nvoid run()\n{\n // Make sure the solver produces A=LU where A is the identity matrix.\n typedef typename T_Config::template setMemSpace::Type Config_h;\n typedef Vector FVector_h;\n typedef typename TConfig::MatPrec Matrix_data;\n const int N = 32;\n Matrix A(N, N, N, CSR);\n A.set_initialized(0);\n this->make_identity(A);\n A.set_initialized(1);\n AMG_Config cfg;\n dense_lu_solver::DenseLUSolver solver(cfg, \"\", NULL);\n solver.setup(A, false);\n FVector_h m(N * N);\n hipMemcpy(m.raw(), solver.get_dense_A(), N * N * sizeof(Matrix_data), hipMemcpyDeviceToHost);\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL(hipGetLastError(), hipSuccess);\n\n for ( int i = 0 ; i < N ; ++i )\n for ( int j = 0 ; j < N ; ++j )\n {\n UNITTEST_ASSERT_EQUAL(m[i * solver.get_lda() + j], i == j ? Matrix_data(1) : Matrix_data(0));\n }\n}\n\nDECLARE_UNITTEST_END(DenseLUSolverTest_Factorization_Id_32)\n\nDenseLUSolverTest_Factorization_Id_32::Type> DenseLUSolverTest_Factorization_Id_32_dDDI;\nDenseLUSolverTest_Factorization_Id_32::Type> DenseLUSolverTest_Factorization_Id_32_dFFI;\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDECLARE_UNITTEST_BEGIN_EXTD(DenseLUSolverTest_Factorization_Id_256, DenseLUSolverTest_Base);\n\nvoid run()\n{\n // Make sure the solver produces A=LU where A is the identity matrix.\n typedef typename T_Config::template setMemSpace::Type Config_h;\n typedef Vector FVector_h;\n typedef typename TConfig::MatPrec Matrix_data;\n const int N = 256;\n Matrix A(N, N, N, CSR);\n A.set_initialized(0);\n this->make_identity(A);\n A.set_initialized(1);\n AMG_Config cfg;\n dense_lu_solver::DenseLUSolver solver(cfg, \"\", NULL);\n solver.setup(A, false);\n FVector_h m(N * N);\n hipMemcpy(m.raw(), solver.get_dense_A(), N * N * sizeof(Matrix_data), hipMemcpyDeviceToHost);\n cudaCheckError();\n UNITTEST_ASSERT_EQUAL(hipGetLastError(), hipSuccess);\n\n for ( int i = 0 ; i < N ; ++i )\n for ( int j = 0 ; j < N ; ++j )\n {\n UNITTEST_ASSERT_EQUAL(m[i * solver.get_lda() + j], i == j ? Matrix_data(1) : Matrix_data(0));\n }\n}\n\nDECLARE_UNITTEST_END(DenseLUSolverTest_Factorization_Id_256)\n\nDenseLUSolverTest_Factorization_Id_256::Type> DenseLUSolverTest_Factorization_Id_256_dDDI;\nDenseLUSolverTest_Factorization_Id_256::Type> DenseLUSolverTest_Factorization_Id_256_dFFI;\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDECLARE_UNITTEST_BEGIN_EXTD(DenseLUSolverTest_Solve_Poisson3D, DenseLUSolverTest_Base);\n\nvoid run()\n{\n typedef typename T_Config::MatPrec Matrix_data;\n typedef typename T_Config::template setMemSpace::Type Config_d;\n typedef typename T_Config::template setMemSpace ::Type Config_h;\n typedef Matrix Matrix_d;\n typedef Vector Vector_d;\n typedef Matrix Matrix_h;\n typedef Vector Vector_h;\n Matrix_h A_h;\n Vector_h x_h, b_h;\n \n A_h.set_initialized(0);\n generatePoissonForTest(A_h, 1, 0, 27, 16, 16, 16);\n A_h.set_initialized(1);\n \n AMG_Config cfg;\n cfg.parseParameterString(\"monitor_residual=1\");\n Matrix_d A_d(A_h);\n dense_lu_solver::DenseLUSolver solver(cfg, \"default\", NULL);\n solver.setup(A_d, false);\n const int n = A_h.get_num_rows();\n Vector_d b_d(n), x_d(n), r_d(n);\n thrust_wrapper::fill(b_d.begin(), b_d.end(), Matrix_data(1));\n solver.solve(b_d, x_d, false);\n solver.compute_residual(b_d, x_d, r_d);\n Vector_h resid_nrm(1);\n solver.compute_norm(r_d, resid_nrm);\n Vector_h rhs_nrm(1);\n solver.compute_norm(b_d, rhs_nrm);\n double relError = resid_nrm[0] / rhs_nrm[0];\n\n if ( T_Config::matPrec == AMGX_matDouble )\n {\n UNITTEST_ASSERT_EQUAL_TOL(resid_nrm[0], 0.0, 1.0e-12);\n }\n else\n {\n UNITTEST_ASSERT_EQUAL_TOL(resid_nrm[0], 0.0f, 1.0e-6f);\n }\n}\n\nDECLARE_UNITTEST_END(DenseLUSolverTest_Solve_Poisson3D)\n\nDenseLUSolverTest_Solve_Poisson3D::Type> DenseLUSolverTest_Solve_Poisson3D_dDDI;\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDECLARE_UNITTEST_BEGIN_EXTD(DenseLUSolverTest_Solve_Id_32, DenseLUSolverTest_Base);\n\nvoid run()\n{\n typedef typename T_Config::template setMemSpace::Type Config_h;\n typedef Vector FVector_h;\n typedef typename TConfig::MatPrec Matrix_data;\n const int N = 32;\n Matrix A(N, N, N, CSR);\n A.set_initialized(0);\n this->make_identity(A);\n A.set_initialized(1);\n AMG_Config cfg;\n dense_lu_solver::DenseLUSolver solver(cfg, \"\", NULL);\n solver.setup(A, false);\n FVector_h b_h(N);\n\n for ( int i = 0 ; i < N ; ++i )\n {\n b_h[i] = Matrix_data(rand()) / RAND_MAX;\n }\n\n Vector b(b_h), x(N), r(N);\n solver.solve(b, x, false);\n UNITTEST_ASSERT_EQUAL(x, b);\n}\n\nDECLARE_UNITTEST_END(DenseLUSolverTest_Solve_Id_32)\n\nDenseLUSolverTest_Solve_Id_32::Type> DenseLUSolverTest_Solve_Id_32_dDDI;\nDenseLUSolverTest_Solve_Id_32::Type> DenseLUSolverTest_Solve_Id_32_dFFI;\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDECLARE_UNITTEST_BEGIN_EXTD(DenseLUSolverTest_Solve_Id_256, DenseLUSolverTest_Base);\n\nvoid run()\n{\n typedef typename T_Config::template setMemSpace::Type Config_h;\n typedef Vector FVector_h;\n typedef typename TConfig::MatPrec Matrix_data;\n const int N = 256;\n Matrix A(N, N, N, CSR);\n A.set_initialized(0);\n this->make_identity(A);\n A.set_initialized(1);\n AMG_Config cfg;\n dense_lu_solver::DenseLUSolver solver(cfg, \"\", NULL);\n solver.setup(A, false);\n FVector_h b_h(N);\n\n for ( int i = 0 ; i < N ; ++i )\n {\n b_h[i] = Matrix_data(rand()) / RAND_MAX;\n }\n\n Vector b(b_h), x(N);\n solver.solve(b, x, false);\n UNITTEST_ASSERT_EQUAL(x, b);\n}\n\nDECLARE_UNITTEST_END(DenseLUSolverTest_Solve_Id_256)\n\nDenseLUSolverTest_Solve_Id_256::Type> DenseLUSolverTest_Solve_Id_256_dDDI;\nDenseLUSolverTest_Solve_Id_256::Type> DenseLUSolverTest_Solve_Id_256_dFFI;\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace amgx\n\n", "source": "/kaggle/working/amgx/src/tests/dense_lu.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"unit_test.h\"\n#include \"amg_solver.h\"\n#include \"test_utils.h\"\n#include \"matrix_coloring/min_max.h\"\n#include \"matrix_coloring/multi_hash.h\"\n#include \"matrix_coloring/parallel_greedy.h\"\n#include \n#include \"matrix_analysis.h\"\n\nnamespace amgx\n{\n\n__global__ void start_coloring_mark() {}\n__global__ void end_coloring_mark() {}\n\ntemplate\nstruct neighbor_color_validity_checker\n{\n __host__ __device__ static int check(const int *A_row_offsets, const int *A_col_indices, const int *color, const int row_start, const int row_end, const int my_color, const int self)\n {\n for (int row_it = row_start; row_it < row_end; row_it++)\n {\n int c = A_col_indices[row_it];\n int col = color[c];\n\n if (my_color == col && c != self)\n {\n return 0;\n }\n\n return neighbor_color_validity_checker < LEVEL - 1 >::check(A_row_offsets, A_col_indices, color, A_row_offsets[c], A_row_offsets[c + 1], my_color, self);\n }\n\n return 1; //empty row: ok\n }\n\n};\n\ntemplate<>\nstruct neighbor_color_validity_checker<0>\n{\n __host__ __device__ static int check(const int *A_row_offsets, const int *A_col_indices, const int *color, const int row_start, const int row_end, const int my_color, const int self)\n {\n //end compile time recursion\n return 1;\n }\n};\n\ntemplate\nstruct CheckColoring\n{\n const IndexType *_A_offsets;\n const IndexType *_A_cols;\n const IndexType *_colors;\n const int _coloring_level;\n\n CheckColoring(const IndexType *A_offsets, const IndexType *A_cols, const IndexType *colors, const int coloring_level) : _A_offsets(A_offsets), _A_cols(A_cols), _colors(colors), _coloring_level(coloring_level) { }\n\n __host__ __device__ IndexType operator()(IndexType i)\n {\n IndexType col = _colors[i];\n // do not count uncolored nodes\n\n if (col == 0)\n {\n return 0;\n }\n\n IndexType row_start = _A_offsets[i];\n IndexType row_end = _A_offsets[i + 1];\n\n if (_coloring_level == 4)\n {\n return neighbor_color_validity_checker<4>::check(_A_offsets, _A_cols, _colors, row_start, row_end, col, i) == 0;\n }\n else if (_coloring_level == 5)\n {\n return neighbor_color_validity_checker<5>::check(_A_offsets, _A_cols, _colors, row_start, row_end, col, i) == 0;\n }\n else if (_coloring_level == 6)\n {\n return neighbor_color_validity_checker<6>::check(_A_offsets, _A_cols, _colors, row_start, row_end, col, i) == 0;\n }\n else if (_coloring_level == 7)\n {\n return neighbor_color_validity_checker<7>::check(_A_offsets, _A_cols, _colors, row_start, row_end, col, i) == 0;\n }\n else\n {\n for (IndexType r = row_start; r < row_end; r++)\n {\n IndexType j = _A_cols[r];\n\n // skip diagonal\n if (j == i) { continue; }\n\n IndexType jcol = _colors[j];\n\n if (jcol == 0) { continue; }\n\n // if 2 colors are adjacent, return 1\n if (jcol == col)\n {\n //printf(\"invalid coloring, row=%d, col=%d, row_color=%d, col_color=%d\\n\",i,j,col,jcol);\n return 1;\n }\n\n if (_coloring_level >= 2)\n {\n IndexType row_start = _A_offsets[j];\n IndexType row_end = _A_offsets[j + 1];\n\n for (IndexType r = row_start; r < row_end; r++)\n {\n IndexType k = _A_cols[r];\n\n // skip diagonal\n if (k == i) { continue; }\n\n IndexType jcol = _colors[k];\n\n if (jcol == 0) { continue; }\n\n // if 2 colors are adjacent, return 1\n if (jcol == col)\n {\n //printf(\"invalid 2 coloring, row=%d, col=%d, row_color=%d, col_color=%d\\n\",i,j2,col,jcol);\n return 1;\n }\n\n if (_coloring_level >= 3)\n {\n IndexType row_start = _A_offsets[k];\n IndexType row_end = _A_offsets[k + 1];\n\n for (IndexType r = row_start; r < row_end; r++)\n {\n IndexType l = _A_cols[r];\n\n // skip diagonal\n if (l == i) { continue; }\n\n IndexType jcol = _colors[l];\n\n if (jcol == 0) { continue; }\n\n // if 2 colors are adjacent, return 1\n if (jcol == col)\n {\n //printf(\"invalid 2 coloring, row=%d, col=%d, row_color=%d, col_color=%d\\n\",i,j2,col,jcol);\n return 1;\n }\n }\n }\n }\n }\n }\n } // loop over first ring neighbours\n\n // no conflict => return 0\n return 0;\n }\n};\n\nstatic std::map total_runtime; //for global aggregate scoring across tests\nstatic std::map coloring_score;\nstatic std::map coloring_runs;\n\nstatic FILE *f_csv_file = 0; //once per test rewrite\nstruct is_zero\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x == 0;\n }\n};\n\nDECLARE_UNITTEST_BEGIN(MatrixColoringTest_Base);\n\nstd::string base_keywords()\n{\n return \"coloring\";\n}\n\n// dense histogram using binary search\ntemplate \nvoid color_histogram(const Vector1 &row_colors, Vector2 &histogram)\n{\n typedef typename Vector1::value_type ValueType; // input value type\n typedef typename Vector2::value_type IndexType; // histogram index type\n // copy input data (could be skipped if input is allowed to be modified)\n device_vector_alloc data(row_colors);\n // sort data to bring equal elements together\n thrust_wrapper::sort(data.begin(), data.end());\n // number of histogram bins is equal to the maximum value plus one\n IndexType num_bins = data.back() + 1;\n // resize histogram storage\n histogram.resize(num_bins);\n // find the end of each bin of values\n amgx::thrust::counting_iterator search_begin(0);\n amgx::thrust::upper_bound(data.begin(), data.end(),\n search_begin, search_begin + num_bins,\n histogram.begin());\n // compute the histogram by taking differences of the cumulative histogram\n amgx::thrust::adjacent_difference(histogram.begin(), histogram.end(),\n histogram.begin());\n}\n\n\n\n\ntemplate< AMGX_VecPrecision VecPrecision, AMGX_MatPrecision MatPrecision >\nvoid\ncolor_matrix_file(const std::string &filename, int max_coloring_level = 3)\n{\n AMG_Config cfg;\n size_t pool_size = cfg.getParameter(\"device_mem_pool_size\", \"default\");\n size_t max_alloc_size = cfg.getParameter(\"device_mem_pool_max_alloc_size\", \"default\");\n\n if ( memory::hasPinnedMemoryPool() )\n {\n memory::setPinnedMemoryPool( new memory::PinnedMemoryPool() );\n }\n\n if ( memory::hasDeviceMemoryPool() )\n {\n memory::setDeviceMemoryPool( new memory::DeviceMemoryPool(pool_size, max_alloc_size, 0) );\n }\n\n cudaDeviceSynchronize(); // why ????\n cudaCheckError();\n bool verbose_output = 0;\n const std::string cfg_scope = \"default\";\n std::stringstream config_string_base;\n config_string_base << \"max_uncolored_percentage=0, coloring_level=1\";\n std::vector coloring_schemes;\n#define ALL 1\n#if ALL\n coloring_schemes.push_back(\"GREEDY_RECOLOR\");\n coloring_schemes.push_back(\"MIN_MAX\");\n coloring_schemes.push_back(\"MIN_MAX_2RING\");\n coloring_schemes.push_back(\"PARALLEL_GREEDY\");\n coloring_schemes.push_back(\"MULTI_HASH,max_uncolored_percentage=0.0\");\n coloring_schemes.push_back(\"GREEDY_RECOLOR,coloring_custom_arg=FIRSTSTEP\");\n coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING\");\n //coloring_schemes.push_back(\"SERIAL_GREEDY_BFS, coloring_custom_arg=IDO\");\n coloring_schemes.push_back(\"SERIAL_GREEDY_BFS\");\n#else\n //max_coloring_level = 2;\n#endif\n\n if (max_coloring_level >= 2)\n {\n#if 1\n coloring_schemes.push_back(\"\"); //placeholder, newline in output if enabled\n coloring_schemes.push_back(\"MIN_MAX, coloring_level=2\");\n coloring_schemes.push_back(\"PARALLEL_GREEDY, coloring_level=2\");\n coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=2\");\n //coloring_schemes.push_back(\"GREEDY_RECOLOR, coloring_level=2\");\n //coloring_schemes.push_back(\"GREEDY_RECOLOR, coloring_level=2,coloring_custom_arg=FIRSTSTEP\");\n //coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=2, coloring_custom_arg=FULL64\");\n //coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=2, coloring_custom_arg=FULL256\");\n //coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=2, coloring_custom_arg=CN1\");\n //coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=2, coloring_custom_arg=CN2\");\n coloring_schemes.push_back(\"SERIAL_GREEDY_BFS, coloring_level=2\");\n /*coloring_schemes.push_back(\"SERIAL_GREEDY_BFS, coloring_custom_arg=IDO, coloring_level=2\");\n */\n#else\n coloring_schemes.push_back(\"PARALLEL_GREEDY, coloring_level=2\");\n coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=2\");\n#endif\n }\n\n if (max_coloring_level >= 3)\n {\n coloring_schemes.push_back(\"\"); //placeholder, newline in output if enabled\n coloring_schemes.push_back(\"PARALLEL_GREEDY, coloring_level=3\");\n coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=3\");\n coloring_schemes.push_back(\"SERIAL_GREEDY_BFS, coloring_level=3\");\n }\n\n if (max_coloring_level >= 4)\n {\n coloring_schemes.push_back(\"\"); //placeholder, newline in output if enabled\n coloring_schemes.push_back(\"PARALLEL_GREEDY, coloring_level=4\");\n coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=4\");\n coloring_schemes.push_back(\"SERIAL_GREEDY_BFS, coloring_level=4\");\n }\n\n if (max_coloring_level >= 5)\n {\n coloring_schemes.push_back(\"\"); //placeholder, newline in output if enabled\n coloring_schemes.push_back(\"PARALLEL_GREEDY, coloring_level=5\");\n coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=5\");\n coloring_schemes.push_back(\"SERIAL_GREEDY_BFS, coloring_level=5\");\n }\n\n if (verbose_output) //enable for writing a csv with all the info\n {\n if (!f_csv_file)\n {\n f_csv_file = fopen(\"matrix_coloring_test_out.txt\", \"w\");\n }\n else\n {\n f_csv_file = fopen(\"matrix_coloring_test_out.txt\", \"a\");\n }\n }\n\n if (verbose_output) { std::cout << std::endl; }\n\n std::vector text_outputs(coloring_schemes.size());\n int min_time_i = -1;\n float min_time = 0;\n int min_colors_i = -1;\n int min_colors = 0;\n int min_colors_iX = -1;\n int min_colorsX = 0;\n typedef TemplateConfig Config_h;\n typedef TemplateConfig Config_d;\n typedef Matrix Matrix_h;\n typedef Vector Vector_h;\n typedef Matrix Matrix_d;\n typedef Vector Vector_d;\n Matrix_h A_h;\n Vector_h x_h, b_h;\n A_h.set_initialized(0);\n A_h.addProps(CSR);\n std::string fail_msg = \"Cannot open \" + filename;\n this->PrintOnFail(fail_msg.c_str());\n //UNITTEST_ASSERT_TRUE(this->read_system( filename.c_str(), A_h, b_h, x_h ));\n UNITTEST_ASSERT_TRUE(MatrixIO::readSystem( filename.c_str(), A_h, b_h, x_h ) == AMGX_OK);;\n A_h.set_initialized(1);\n MatrixAnalysis ana(&A_h, &b_h);\n bool str_sym, sym;\n bool verbose = false;\n ana.checkSymmetry(str_sym, sym, verbose);\n Matrix_d A_d( A_h );\n\n for (int i = 0; i < coloring_schemes.size(); i++)\n {\n if (coloring_schemes[i] == \"\")\n {\n if (verbose_output) { text_outputs[i] = \"\\n\"; }\n\n if (verbose_output) { std::cout << \"\\n\"; }\n\n min_colorsX = 0;\n min_colors_iX = -1;\n continue;\n }\n\n std::stringstream config_string;\n std::stringstream text_output;\n config_string << config_string_base.str() << \", matrix_coloring_scheme=\" << coloring_schemes[i] << std::endl;\n AMG_Config cfg;\n UNITTEST_ASSERT_TRUE ( cfg.parseParameterString(config_string.str().c_str()) == AMGX_OK);\n UNITTEST_ASSERT_TRUE_DESC(\"Matrix does not have structural symmetry, unit test invalid\", str_sym);;\n A_d.set_initialized(0);\n A_d.addProps(CSR);\n// cudaDeviceSynchronize();\ncudaCheckError();\n cudaEvent_t color_start, color_stop;\n cudaEventCreate(&color_start);\n cudaCheckError();\n cudaEventCreate(&color_stop);\n cudaCheckError();\n //A_d.colorMatrix(cfg,cfg_scope);\n //A_d.set_initialized(0);\n float elapsed_time;\n cudaEventRecord( color_start);\n cudaCheckError();\n A_d.colorMatrix(cfg, cfg_scope);\n //amgx::testing_tools::hash_path_determinism_checker::singleton()->checkpoint(\"colors\", (void*)A_d.getMatrixColoring().getRowColors().raw(), A_d.getMatrixColoring().getRowColors().size()*4);\n //amgx::testing_tools::hash_path_determinism_checker::singleton()->checkpoint(\"cols\", (void*)A_d.col_indices.raw(), A_d.col_indices.size()*4);\n //amgx::testing_tools::hash_path_determinism_checker::singleton()->checkpoint(\"rows\", (void*)A_d.row_offsets.raw(), A_d.row_offsets.size()*4);\n cudaEventRecord( color_stop);\n cudaCheckError();\n cudaEventSynchronize( color_stop);\n cudaCheckError();\n cudaEventElapsedTime( &elapsed_time, color_start, color_stop);\n elapsed_time *= 1e-3f;\n A_d.set_initialized(1);\n // check coloring\n const IndexType *A_row_offsets_ptr = A_d.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A_d.col_indices.raw();\n const IndexType *row_colors_ptr = A_d.getMatrixColoring().getRowColors().raw();\n IndexType num_bad, num_bad_plus;\n {\n //order +1\n CheckColoring checker(A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, 1 + A_d.getMatrixColoring().getColoringLevel());\n num_bad_plus =\n thrust_wrapper::transform_reduce(amgx::thrust::counting_iterator(0), amgx::thrust::counting_iterator(A_d.get_num_rows()),\n checker, (IndexType)0, amgx::thrust::plus());\n }\n {\n CheckColoring checker(A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, A_d.getMatrixColoring().getColoringLevel());\n num_bad =\n thrust_wrapper::transform_reduce(amgx::thrust::counting_iterator(0), amgx::thrust::counting_iterator(A_d.get_num_rows()),\n checker, (IndexType)0, amgx::thrust::plus());\n }\n int num_uncolored = (int) amgx::thrust::count_if( A_d.getMatrixColoring().getRowColors().begin(), A_d.getMatrixColoring().getRowColors().begin() + A_d.get_num_rows(), is_zero() );\n int mincol = (int) amgx::thrust::reduce( A_d.getMatrixColoring().getRowColors().begin(), A_d.getMatrixColoring().getRowColors().begin() + A_d.get_num_rows(), 0, amgx::thrust::minimum() );\n int maxcol = (int) amgx::thrust::reduce( A_d.getMatrixColoring().getRowColors().begin(), A_d.getMatrixColoring().getRowColors().begin() + A_d.get_num_rows(), 0, amgx::thrust::maximum() );\n std::stringstream error_str;\n error_str << \"Coloring scheme \" << coloring_schemes[i] << \" with \" << num_bad << \" invalid rows, \" << num_uncolored << \"(\" << mincol << \",\" << maxcol << \")\" << \" uncolored rows and \" << A_d.getMatrixColoring().getNumColors() << \" used colors and bonus = \" << num_bad_plus;\n\n if (verbose_output)\n {\n text_output << std::setw(28) << coloring_schemes[i] << \" : num_colors=\" << std::setw(6) << A_d.getMatrixColoring().getNumColors() << \" , runtime= \" << std::setw(8) << std::fixed << std::setprecision(5) << elapsed_time << \" mM(\" << mincol << \",\" << maxcol << \")\"; // << std::endl;;\n }\n\n if (num_bad == 0 && num_uncolored == 0)\n {\n if (elapsed_time < min_time || min_time_i == -1)\n {\n min_time_i = i;\n min_time = elapsed_time;\n }\n\n if (A_d.getMatrixColoring().getNumColors() < min_colors || min_colors_i == -1)\n {\n min_colors_i = i;\n min_colors = A_d.getMatrixColoring().getNumColors();\n }\n\n if (A_d.getMatrixColoring().getNumColors() < min_colorsX || min_colors_iX == -1)\n {\n min_colors_iX = i;\n min_colorsX = A_d.getMatrixColoring().getNumColors();\n }\n }\n\n int degree;\n {\n device_vector_alloc hist;\n amgx::thrust::host_vector h_hist;\n device_vector_alloc columns(A_d.get_num_rows());\n amgx::thrust::adjacent_difference(A_d.row_offsets.begin(), A_d.row_offsets.end(), columns.begin());\n degree = amgx::thrust::reduce(columns.begin(), columns.end(), -1, amgx::thrust::maximum());\n color_histogram(A_d.getMatrixColoring().getRowColors(), hist);\n h_hist = hist;\n float total = 0, m = h_hist[1], M = h_hist[1]; //0 is 0\n\n for (int j = 1; j < h_hist.size(); j++)\n {\n if (h_hist[j] > M) { M = h_hist[j]; }\n\n if (h_hist[j] < m) { m = h_hist[j]; }\n\n total += h_hist[j];\n }\n\n float avg = total / (h_hist.size() - 1);\n float rms = 0;\n\n for (int j = 0; j < h_hist.size(); j++)\n {\n rms = (h_hist[j] - avg) * (h_hist[j] - avg);\n }\n\n rms = sqrt(rms / (h_hist.size() - 1));\n\n if (verbose_output) { text_output << \" [\" << \" min=\" << m << \", max=\" << M << \", avg=\" << avg << \", rms=\" << rms << \", deg=\" << degree << \", uncol=\" << num_uncolored << \" bonus=\" << num_bad_plus << \"]\"; }\n\n if (f_csv_file) { fprintf(f_csv_file, \"%s;%d;%d;%d;%.3f;%.3f;%.3f;%.3f;%.3f\\n\", coloring_schemes[i].c_str(), A_d.getMatrixColoring().getNumColors(), num_bad, num_uncolored, elapsed_time * 1e3, m, M, avg, rms); }\n }\n total_runtime[coloring_schemes[i]] += elapsed_time;\n coloring_score[coloring_schemes[i]] += double(A_d.getMatrixColoring().getNumColors() - 1) / pow(static_cast(degree), static_cast(A_d.getMatrixColoring().getColoringLevel()));\n coloring_runs[coloring_schemes[i]] += 1.0;\n\n if (verbose_output) { text_output << std::endl; }\n\n if (verbose_output) { printf(\"%s\", text_output.str().c_str()); }\n\n text_outputs[i] = text_output.str();\n UNITTEST_ASSERT_EQUAL_DESC(error_str.str().c_str(), num_bad, 0);\n }\n\n if (verbose_output && 0)\n {\n for (int i = 0; i < text_outputs.size(); i++)\n {\n if (i == min_colors_i) { std::cout << \"*\"; }\n\n if (i == min_time_i) { std::cout << \"#\"; }\n\n std::cout << text_outputs[i];\n }\n }\n\n if (f_csv_file)\n {\n fprintf(f_csv_file, \"\\n\");\n fclose(f_csv_file);\n }\n\n if (verbose_output)\n {\n printf(\"\\n\\n\\n\");\n\n for (int i = 0; i < text_outputs.size(); i++)\n {\n if (coloring_schemes[i] == \"\")\n {\n std::cout << \"\\n\";\n continue;\n }\n\n std::string name = coloring_schemes[i];\n std::cout << std::setw(80) << coloring_schemes[i];\n std::cout << \" runtime=\" << std::setw(8) << std::fixed << std::setprecision(5) << total_runtime[name];\n std::cout << \" inefficiency=\" << std::setw(6) << std::fixed << std::setprecision(3) << coloring_score[name] / coloring_runs[name]; // << std::endl;;\n std::cout << \"\\n\";\n }\n\n printf(\"\\n\\n\\n\");\n }\n\n cudaDeviceSynchronize();\n cudaCheckError();\n}\n\nDECLARE_UNITTEST_END(MatrixColoringTest_Base);\n\n\nDECLARE_UNITTEST_BEGIN_EXTD(MatrixColoringTestatmosmodj, MatrixColoringTest_Base);\n\nvoid run()\n{\n //std::cout << std::endl;\n MatrixColoringTest_Base::template color_matrix_file( UnitTest::get_configuration().data_folder + \"Public/Florida/atmosmodj.mtx.bin\");\n}\n\nDECLARE_UNITTEST_END(MatrixColoringTestatmosmodj)\nMatrixColoringTestatmosmodj::Type> MatrixColoringTestatmosmodj_dDDI;\n\n\n//--------------------------------------------------------------------------------------------------\nDECLARE_UNITTEST_BEGIN_EXTD(MatrixColoringTestatmosmodl, MatrixColoringTest_Base);\n\nvoid run()\n{\n //std::cout << std::endl;\n MatrixColoringTest_Base::template color_matrix_file( UnitTest::get_configuration().data_folder + \"Public/Florida/atmosmodl.mtx.bin\");\n}\n\nDECLARE_UNITTEST_END(MatrixColoringTestatmosmodl)\nMatrixColoringTestatmosmodl::Type> MatrixColoringTestatmosmodl_dDDI;\n\n\n//--------------------------------------------------------------------------------------------------\nDECLARE_UNITTEST_BEGIN_EXTD(MatrixColoringTestpoisson, MatrixColoringTest_Base);\n\nvoid run()\n{\n //std::cout << std::endl;\n MatrixColoringTest_Base::template color_matrix_file( UnitTest::get_configuration().data_folder + \"Internal/poisson/poisson27x50x50x50.mtx.bin\");\n}\n\nDECLARE_UNITTEST_END(MatrixColoringTestpisson)\nMatrixColoringTestpoisson::Type> MatrixColoringTestpoisson_dDDI;\n\n} //namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \"unit_test.h\"\n#include \"amg_solver.h\"\n#include \"test_utils.h\"\n#include \"matrix_coloring/min_max.h\"\n#include \"matrix_coloring/multi_hash.h\"\n#include \"matrix_coloring/parallel_greedy.h\"\n#include \n#include \"matrix_analysis.h\"\n\nnamespace amgx\n{\n\n__global__ void start_coloring_mark() {}\n__global__ void end_coloring_mark() {}\n\ntemplate\nstruct neighbor_color_validity_checker\n{\n __host__ __device__ static int check(const int *A_row_offsets, const int *A_col_indices, const int *color, const int row_start, const int row_end, const int my_color, const int self)\n {\n for (int row_it = row_start; row_it < row_end; row_it++)\n {\n int c = A_col_indices[row_it];\n int col = color[c];\n\n if (my_color == col && c != self)\n {\n return 0;\n }\n\n return neighbor_color_validity_checker < LEVEL - 1 >::check(A_row_offsets, A_col_indices, color, A_row_offsets[c], A_row_offsets[c + 1], my_color, self);\n }\n\n return 1; //empty row: ok\n }\n\n};\n\ntemplate<>\nstruct neighbor_color_validity_checker<0>\n{\n __host__ __device__ static int check(const int *A_row_offsets, const int *A_col_indices, const int *color, const int row_start, const int row_end, const int my_color, const int self)\n {\n //end compile time recursion\n return 1;\n }\n};\n\ntemplate\nstruct CheckColoring\n{\n const IndexType *_A_offsets;\n const IndexType *_A_cols;\n const IndexType *_colors;\n const int _coloring_level;\n\n CheckColoring(const IndexType *A_offsets, const IndexType *A_cols, const IndexType *colors, const int coloring_level) : _A_offsets(A_offsets), _A_cols(A_cols), _colors(colors), _coloring_level(coloring_level) { }\n\n __host__ __device__ IndexType operator()(IndexType i)\n {\n IndexType col = _colors[i];\n // do not count uncolored nodes\n\n if (col == 0)\n {\n return 0;\n }\n\n IndexType row_start = _A_offsets[i];\n IndexType row_end = _A_offsets[i + 1];\n\n if (_coloring_level == 4)\n {\n return neighbor_color_validity_checker<4>::check(_A_offsets, _A_cols, _colors, row_start, row_end, col, i) == 0;\n }\n else if (_coloring_level == 5)\n {\n return neighbor_color_validity_checker<5>::check(_A_offsets, _A_cols, _colors, row_start, row_end, col, i) == 0;\n }\n else if (_coloring_level == 6)\n {\n return neighbor_color_validity_checker<6>::check(_A_offsets, _A_cols, _colors, row_start, row_end, col, i) == 0;\n }\n else if (_coloring_level == 7)\n {\n return neighbor_color_validity_checker<7>::check(_A_offsets, _A_cols, _colors, row_start, row_end, col, i) == 0;\n }\n else\n {\n for (IndexType r = row_start; r < row_end; r++)\n {\n IndexType j = _A_cols[r];\n\n // skip diagonal\n if (j == i) { continue; }\n\n IndexType jcol = _colors[j];\n\n if (jcol == 0) { continue; }\n\n // if 2 colors are adjacent, return 1\n if (jcol == col)\n {\n //printf(\"invalid coloring, row=%d, col=%d, row_color=%d, col_color=%d\\n\",i,j,col,jcol);\n return 1;\n }\n\n if (_coloring_level >= 2)\n {\n IndexType row_start = _A_offsets[j];\n IndexType row_end = _A_offsets[j + 1];\n\n for (IndexType r = row_start; r < row_end; r++)\n {\n IndexType k = _A_cols[r];\n\n // skip diagonal\n if (k == i) { continue; }\n\n IndexType jcol = _colors[k];\n\n if (jcol == 0) { continue; }\n\n // if 2 colors are adjacent, return 1\n if (jcol == col)\n {\n //printf(\"invalid 2 coloring, row=%d, col=%d, row_color=%d, col_color=%d\\n\",i,j2,col,jcol);\n return 1;\n }\n\n if (_coloring_level >= 3)\n {\n IndexType row_start = _A_offsets[k];\n IndexType row_end = _A_offsets[k + 1];\n\n for (IndexType r = row_start; r < row_end; r++)\n {\n IndexType l = _A_cols[r];\n\n // skip diagonal\n if (l == i) { continue; }\n\n IndexType jcol = _colors[l];\n\n if (jcol == 0) { continue; }\n\n // if 2 colors are adjacent, return 1\n if (jcol == col)\n {\n //printf(\"invalid 2 coloring, row=%d, col=%d, row_color=%d, col_color=%d\\n\",i,j2,col,jcol);\n return 1;\n }\n }\n }\n }\n }\n }\n } // loop over first ring neighbours\n\n // no conflict => return 0\n return 0;\n }\n};\n\nstatic std::map total_runtime; //for global aggregate scoring across tests\nstatic std::map coloring_score;\nstatic std::map coloring_runs;\n\nstatic FILE *f_csv_file = 0; //once per test rewrite\nstruct is_zero\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x == 0;\n }\n};\n\nDECLARE_UNITTEST_BEGIN(MatrixColoringTest_Base);\n\nstd::string base_keywords()\n{\n return \"coloring\";\n}\n\n// dense histogram using binary search\ntemplate \nvoid color_histogram(const Vector1 &row_colors, Vector2 &histogram)\n{\n typedef typename Vector1::value_type ValueType; // input value type\n typedef typename Vector2::value_type IndexType; // histogram index type\n // copy input data (could be skipped if input is allowed to be modified)\n device_vector_alloc data(row_colors);\n // sort data to bring equal elements together\n thrust_wrapper::sort(data.begin(), data.end());\n // number of histogram bins is equal to the maximum value plus one\n IndexType num_bins = data.back() + 1;\n // resize histogram storage\n histogram.resize(num_bins);\n // find the end of each bin of values\n amgx::thrust::counting_iterator search_begin(0);\n amgx::thrust::upper_bound(data.begin(), data.end(),\n search_begin, search_begin + num_bins,\n histogram.begin());\n // compute the histogram by taking differences of the cumulative histogram\n amgx::thrust::adjacent_difference(histogram.begin(), histogram.end(),\n histogram.begin());\n}\n\n\n\n\ntemplate< AMGX_VecPrecision VecPrecision, AMGX_MatPrecision MatPrecision >\nvoid\ncolor_matrix_file(const std::string &filename, int max_coloring_level = 3)\n{\n AMG_Config cfg;\n size_t pool_size = cfg.getParameter(\"device_mem_pool_size\", \"default\");\n size_t max_alloc_size = cfg.getParameter(\"device_mem_pool_max_alloc_size\", \"default\");\n\n if ( memory::hasPinnedMemoryPool() )\n {\n memory::setPinnedMemoryPool( new memory::PinnedMemoryPool() );\n }\n\n if ( memory::hasDeviceMemoryPool() )\n {\n memory::setDeviceMemoryPool( new memory::DeviceMemoryPool(pool_size, max_alloc_size, 0) );\n }\n\n hipDeviceSynchronize(); // why ????\n cudaCheckError();\n bool verbose_output = 0;\n const std::string cfg_scope = \"default\";\n std::stringstream config_string_base;\n config_string_base << \"max_uncolored_percentage=0, coloring_level=1\";\n std::vector coloring_schemes;\n#define ALL 1\n#if ALL\n coloring_schemes.push_back(\"GREEDY_RECOLOR\");\n coloring_schemes.push_back(\"MIN_MAX\");\n coloring_schemes.push_back(\"MIN_MAX_2RING\");\n coloring_schemes.push_back(\"PARALLEL_GREEDY\");\n coloring_schemes.push_back(\"MULTI_HASH,max_uncolored_percentage=0.0\");\n coloring_schemes.push_back(\"GREEDY_RECOLOR,coloring_custom_arg=FIRSTSTEP\");\n coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING\");\n //coloring_schemes.push_back(\"SERIAL_GREEDY_BFS, coloring_custom_arg=IDO\");\n coloring_schemes.push_back(\"SERIAL_GREEDY_BFS\");\n#else\n //max_coloring_level = 2;\n#endif\n\n if (max_coloring_level >= 2)\n {\n#if 1\n coloring_schemes.push_back(\"\"); //placeholder, newline in output if enabled\n coloring_schemes.push_back(\"MIN_MAX, coloring_level=2\");\n coloring_schemes.push_back(\"PARALLEL_GREEDY, coloring_level=2\");\n coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=2\");\n //coloring_schemes.push_back(\"GREEDY_RECOLOR, coloring_level=2\");\n //coloring_schemes.push_back(\"GREEDY_RECOLOR, coloring_level=2,coloring_custom_arg=FIRSTSTEP\");\n //coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=2, coloring_custom_arg=FULL64\");\n //coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=2, coloring_custom_arg=FULL256\");\n //coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=2, coloring_custom_arg=CN1\");\n //coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=2, coloring_custom_arg=CN2\");\n coloring_schemes.push_back(\"SERIAL_GREEDY_BFS, coloring_level=2\");\n /*coloring_schemes.push_back(\"SERIAL_GREEDY_BFS, coloring_custom_arg=IDO, coloring_level=2\");\n */\n#else\n coloring_schemes.push_back(\"PARALLEL_GREEDY, coloring_level=2\");\n coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=2\");\n#endif\n }\n\n if (max_coloring_level >= 3)\n {\n coloring_schemes.push_back(\"\"); //placeholder, newline in output if enabled\n coloring_schemes.push_back(\"PARALLEL_GREEDY, coloring_level=3\");\n coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=3\");\n coloring_schemes.push_back(\"SERIAL_GREEDY_BFS, coloring_level=3\");\n }\n\n if (max_coloring_level >= 4)\n {\n coloring_schemes.push_back(\"\"); //placeholder, newline in output if enabled\n coloring_schemes.push_back(\"PARALLEL_GREEDY, coloring_level=4\");\n coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=4\");\n coloring_schemes.push_back(\"SERIAL_GREEDY_BFS, coloring_level=4\");\n }\n\n if (max_coloring_level >= 5)\n {\n coloring_schemes.push_back(\"\"); //placeholder, newline in output if enabled\n coloring_schemes.push_back(\"PARALLEL_GREEDY, coloring_level=5\");\n coloring_schemes.push_back(\"GREEDY_MIN_MAX_2RING, coloring_level=5\");\n coloring_schemes.push_back(\"SERIAL_GREEDY_BFS, coloring_level=5\");\n }\n\n if (verbose_output) //enable for writing a csv with all the info\n {\n if (!f_csv_file)\n {\n f_csv_file = fopen(\"matrix_coloring_test_out.txt\", \"w\");\n }\n else\n {\n f_csv_file = fopen(\"matrix_coloring_test_out.txt\", \"a\");\n }\n }\n\n if (verbose_output) { std::cout << std::endl; }\n\n std::vector text_outputs(coloring_schemes.size());\n int min_time_i = -1;\n float min_time = 0;\n int min_colors_i = -1;\n int min_colors = 0;\n int min_colors_iX = -1;\n int min_colorsX = 0;\n typedef TemplateConfig Config_h;\n typedef TemplateConfig Config_d;\n typedef Matrix Matrix_h;\n typedef Vector Vector_h;\n typedef Matrix Matrix_d;\n typedef Vector Vector_d;\n Matrix_h A_h;\n Vector_h x_h, b_h;\n A_h.set_initialized(0);\n A_h.addProps(CSR);\n std::string fail_msg = \"Cannot open \" + filename;\n this->PrintOnFail(fail_msg.c_str());\n //UNITTEST_ASSERT_TRUE(this->read_system( filename.c_str(), A_h, b_h, x_h ));\n UNITTEST_ASSERT_TRUE(MatrixIO::readSystem( filename.c_str(), A_h, b_h, x_h ) == AMGX_OK);;\n A_h.set_initialized(1);\n MatrixAnalysis ana(&A_h, &b_h);\n bool str_sym, sym;\n bool verbose = false;\n ana.checkSymmetry(str_sym, sym, verbose);\n Matrix_d A_d( A_h );\n\n for (int i = 0; i < coloring_schemes.size(); i++)\n {\n if (coloring_schemes[i] == \"\")\n {\n if (verbose_output) { text_outputs[i] = \"\\n\"; }\n\n if (verbose_output) { std::cout << \"\\n\"; }\n\n min_colorsX = 0;\n min_colors_iX = -1;\n continue;\n }\n\n std::stringstream config_string;\n std::stringstream text_output;\n config_string << config_string_base.str() << \", matrix_coloring_scheme=\" << coloring_schemes[i] << std::endl;\n AMG_Config cfg;\n UNITTEST_ASSERT_TRUE ( cfg.parseParameterString(config_string.str().c_str()) == AMGX_OK);\n UNITTEST_ASSERT_TRUE_DESC(\"Matrix does not have structural symmetry, unit test invalid\", str_sym);;\n A_d.set_initialized(0);\n A_d.addProps(CSR);\n// cudaDeviceSynchronize();\ncudaCheckError();\n hipEvent_t color_start, color_stop;\n hipEventCreate(&color_start);\n cudaCheckError();\n hipEventCreate(&color_stop);\n cudaCheckError();\n //A_d.colorMatrix(cfg,cfg_scope);\n //A_d.set_initialized(0);\n float elapsed_time;\n hipEventRecord( color_start);\n cudaCheckError();\n A_d.colorMatrix(cfg, cfg_scope);\n //amgx::testing_tools::hash_path_determinism_checker::singleton()->checkpoint(\"colors\", (void*)A_d.getMatrixColoring().getRowColors().raw(), A_d.getMatrixColoring().getRowColors().size()*4);\n //amgx::testing_tools::hash_path_determinism_checker::singleton()->checkpoint(\"cols\", (void*)A_d.col_indices.raw(), A_d.col_indices.size()*4);\n //amgx::testing_tools::hash_path_determinism_checker::singleton()->checkpoint(\"rows\", (void*)A_d.row_offsets.raw(), A_d.row_offsets.size()*4);\n hipEventRecord( color_stop);\n cudaCheckError();\n hipEventSynchronize( color_stop);\n cudaCheckError();\n hipEventElapsedTime( &elapsed_time, color_start, color_stop);\n elapsed_time *= 1e-3f;\n A_d.set_initialized(1);\n // check coloring\n const IndexType *A_row_offsets_ptr = A_d.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A_d.col_indices.raw();\n const IndexType *row_colors_ptr = A_d.getMatrixColoring().getRowColors().raw();\n IndexType num_bad, num_bad_plus;\n {\n //order +1\n CheckColoring checker(A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, 1 + A_d.getMatrixColoring().getColoringLevel());\n num_bad_plus =\n thrust_wrapper::transform_reduce(amgx::thrust::counting_iterator(0), amgx::thrust::counting_iterator(A_d.get_num_rows()),\n checker, (IndexType)0, amgx::thrust::plus());\n }\n {\n CheckColoring checker(A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, A_d.getMatrixColoring().getColoringLevel());\n num_bad =\n thrust_wrapper::transform_reduce(amgx::thrust::counting_iterator(0), amgx::thrust::counting_iterator(A_d.get_num_rows()),\n checker, (IndexType)0, amgx::thrust::plus());\n }\n int num_uncolored = (int) amgx::thrust::count_if( A_d.getMatrixColoring().getRowColors().begin(), A_d.getMatrixColoring().getRowColors().begin() + A_d.get_num_rows(), is_zero() );\n int mincol = (int) amgx::thrust::reduce( A_d.getMatrixColoring().getRowColors().begin(), A_d.getMatrixColoring().getRowColors().begin() + A_d.get_num_rows(), 0, amgx::thrust::minimum() );\n int maxcol = (int) amgx::thrust::reduce( A_d.getMatrixColoring().getRowColors().begin(), A_d.getMatrixColoring().getRowColors().begin() + A_d.get_num_rows(), 0, amgx::thrust::maximum() );\n std::stringstream error_str;\n error_str << \"Coloring scheme \" << coloring_schemes[i] << \" with \" << num_bad << \" invalid rows, \" << num_uncolored << \"(\" << mincol << \",\" << maxcol << \")\" << \" uncolored rows and \" << A_d.getMatrixColoring().getNumColors() << \" used colors and bonus = \" << num_bad_plus;\n\n if (verbose_output)\n {\n text_output << std::setw(28) << coloring_schemes[i] << \" : num_colors=\" << std::setw(6) << A_d.getMatrixColoring().getNumColors() << \" , runtime= \" << std::setw(8) << std::fixed << std::setprecision(5) << elapsed_time << \" mM(\" << mincol << \",\" << maxcol << \")\"; // << std::endl;;\n }\n\n if (num_bad == 0 && num_uncolored == 0)\n {\n if (elapsed_time < min_time || min_time_i == -1)\n {\n min_time_i = i;\n min_time = elapsed_time;\n }\n\n if (A_d.getMatrixColoring().getNumColors() < min_colors || min_colors_i == -1)\n {\n min_colors_i = i;\n min_colors = A_d.getMatrixColoring().getNumColors();\n }\n\n if (A_d.getMatrixColoring().getNumColors() < min_colorsX || min_colors_iX == -1)\n {\n min_colors_iX = i;\n min_colorsX = A_d.getMatrixColoring().getNumColors();\n }\n }\n\n int degree;\n {\n device_vector_alloc hist;\n amgx::thrust::host_vector h_hist;\n device_vector_alloc columns(A_d.get_num_rows());\n amgx::thrust::adjacent_difference(A_d.row_offsets.begin(), A_d.row_offsets.end(), columns.begin());\n degree = amgx::thrust::reduce(columns.begin(), columns.end(), -1, amgx::thrust::maximum());\n color_histogram(A_d.getMatrixColoring().getRowColors(), hist);\n h_hist = hist;\n float total = 0, m = h_hist[1], M = h_hist[1]; //0 is 0\n\n for (int j = 1; j < h_hist.size(); j++)\n {\n if (h_hist[j] > M) { M = h_hist[j]; }\n\n if (h_hist[j] < m) { m = h_hist[j]; }\n\n total += h_hist[j];\n }\n\n float avg = total / (h_hist.size() - 1);\n float rms = 0;\n\n for (int j = 0; j < h_hist.size(); j++)\n {\n rms = (h_hist[j] - avg) * (h_hist[j] - avg);\n }\n\n rms = sqrt(rms / (h_hist.size() - 1));\n\n if (verbose_output) { text_output << \" [\" << \" min=\" << m << \", max=\" << M << \", avg=\" << avg << \", rms=\" << rms << \", deg=\" << degree << \", uncol=\" << num_uncolored << \" bonus=\" << num_bad_plus << \"]\"; }\n\n if (f_csv_file) { fprintf(f_csv_file, \"%s;%d;%d;%d;%.3f;%.3f;%.3f;%.3f;%.3f\\n\", coloring_schemes[i].c_str(), A_d.getMatrixColoring().getNumColors(), num_bad, num_uncolored, elapsed_time * 1e3, m, M, avg, rms); }\n }\n total_runtime[coloring_schemes[i]] += elapsed_time;\n coloring_score[coloring_schemes[i]] += double(A_d.getMatrixColoring().getNumColors() - 1) / pow(static_cast(degree), static_cast(A_d.getMatrixColoring().getColoringLevel()));\n coloring_runs[coloring_schemes[i]] += 1.0;\n\n if (verbose_output) { text_output << std::endl; }\n\n if (verbose_output) { printf(\"%s\", text_output.str().c_str()); }\n\n text_outputs[i] = text_output.str();\n UNITTEST_ASSERT_EQUAL_DESC(error_str.str().c_str(), num_bad, 0);\n }\n\n if (verbose_output && 0)\n {\n for (int i = 0; i < text_outputs.size(); i++)\n {\n if (i == min_colors_i) { std::cout << \"*\"; }\n\n if (i == min_time_i) { std::cout << \"#\"; }\n\n std::cout << text_outputs[i];\n }\n }\n\n if (f_csv_file)\n {\n fprintf(f_csv_file, \"\\n\");\n fclose(f_csv_file);\n }\n\n if (verbose_output)\n {\n printf(\"\\n\\n\\n\");\n\n for (int i = 0; i < text_outputs.size(); i++)\n {\n if (coloring_schemes[i] == \"\")\n {\n std::cout << \"\\n\";\n continue;\n }\n\n std::string name = coloring_schemes[i];\n std::cout << std::setw(80) << coloring_schemes[i];\n std::cout << \" runtime=\" << std::setw(8) << std::fixed << std::setprecision(5) << total_runtime[name];\n std::cout << \" inefficiency=\" << std::setw(6) << std::fixed << std::setprecision(3) << coloring_score[name] / coloring_runs[name]; // << std::endl;;\n std::cout << \"\\n\";\n }\n\n printf(\"\\n\\n\\n\");\n }\n\n hipDeviceSynchronize();\n cudaCheckError();\n}\n\nDECLARE_UNITTEST_END(MatrixColoringTest_Base);\n\n\nDECLARE_UNITTEST_BEGIN_EXTD(MatrixColoringTestatmosmodj, MatrixColoringTest_Base);\n\nvoid run()\n{\n //std::cout << std::endl;\n MatrixColoringTest_Base::template color_matrix_file( UnitTest::get_configuration().data_folder + \"Public/Florida/atmosmodj.mtx.bin\");\n}\n\nDECLARE_UNITTEST_END(MatrixColoringTestatmosmodj)\nMatrixColoringTestatmosmodj::Type> MatrixColoringTestatmosmodj_dDDI;\n\n\n//--------------------------------------------------------------------------------------------------\nDECLARE_UNITTEST_BEGIN_EXTD(MatrixColoringTestatmosmodl, MatrixColoringTest_Base);\n\nvoid run()\n{\n //std::cout << std::endl;\n MatrixColoringTest_Base::template color_matrix_file( UnitTest::get_configuration().data_folder + \"Public/Florida/atmosmodl.mtx.bin\");\n}\n\nDECLARE_UNITTEST_END(MatrixColoringTestatmosmodl)\nMatrixColoringTestatmosmodl::Type> MatrixColoringTestatmosmodl_dDDI;\n\n\n//--------------------------------------------------------------------------------------------------\nDECLARE_UNITTEST_BEGIN_EXTD(MatrixColoringTestpoisson, MatrixColoringTest_Base);\n\nvoid run()\n{\n //std::cout << std::endl;\n MatrixColoringTest_Base::template color_matrix_file( UnitTest::get_configuration().data_folder + \"Internal/poisson/poisson27x50x50x50.mtx.bin\");\n}\n\nDECLARE_UNITTEST_END(MatrixColoringTestpisson)\nMatrixColoringTestpoisson::Type> MatrixColoringTestpoisson_dDDI;\n\n} //namespace amgx\n", "source": "/kaggle/working/amgx/src/tests/matrix_coloring_test.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"unit_test.h\"\n#include \"amg_solver.h\"\n#include \"aggregation/coarseAgenerators/low_deg_coarse_A_generator.h\"\n#include \n#include \n#include \"test_utils.h\"\n#include \"util.h\"\n#include \"time.h\"\n#include \n#include \n\nnamespace amgx\n{\n\n\nstruct Not_empty_row\n{\n typedef bool result_type;\n template< typename Tuple >\n inline __device__ __host__ bool operator()( const Tuple &t ) const\n {\n return amgx::thrust::get<0>(t) != amgx::thrust::get<1>(t);\n }\n};\n\ntemplate< typename Matrix, typename Vector >\nstatic\nvoid build_sort_permutation( const Matrix &M, Vector &permutation )\n{\n int num_nz = M.get_num_nz();\n Vector row_indices( num_nz, 0 );\n amgx::thrust::scatter_if( amgx::thrust::counting_iterator(0),\n amgx::thrust::counting_iterator(M.row_offsets.size() - 1),\n M.row_offsets.begin(),\n amgx::thrust::make_transform_iterator(\n amgx::thrust::make_zip_iterator( amgx::thrust::make_tuple( M.row_offsets.begin(), M.row_offsets.begin() + 1 ) ),\n Not_empty_row()),\n row_indices.begin());\n cudaCheckError();\n amgx::thrust::inclusive_scan( row_indices.begin(), row_indices.begin() + M.get_num_nz(), row_indices.begin(), amgx::thrust::maximum() );\n cudaCheckError();\n permutation.resize( num_nz );\n thrust_wrapper::sequence( permutation.begin(), permutation.end() );\n cudaCheckError();\n Vector tmp( M.col_indices );\n amgx::thrust::stable_sort_by_key( tmp.begin(), tmp.end(), permutation.begin() );\n cudaCheckError();\n tmp = row_indices;\n amgx::thrust::gather( permutation.begin(), permutation.end(), tmp.begin(), row_indices.begin() );\n cudaCheckError();\n amgx::thrust::stable_sort_by_key( row_indices.begin(), row_indices.end(), permutation.begin() );\n cudaCheckError();\n}\n\ntemplate< typename Value_type >\nstatic __device__ __inline__\nbool equal( Value_type x, Value_type y, Value_type epsilon, Value_type max_abs_error )\n{\n if ( x == y )\n {\n return true;\n }\n\n if ( abs(x - y) < max_abs_error )\n {\n return true;\n }\n\n if ( abs(x - y) <= epsilon )\n {\n return true;\n }\n\n return false;\n}\n\nstatic __device__ __inline__\nbool equal( float x, float y )\n{\n return equal( x, y, 1.0e-6f, 1.0e10f * FLT_MIN );\n}\n\nstatic __device__ __inline__\nbool equal( double x, double y )\n{\n return equal( x, y, 1.0e-12, 1.0e10 * DBL_MIN );\n}\n\ntemplate< typename Value_type >\n__global__\nvoid compare_matrices_kernel( const int num_rows,\n const int block_size_sq,\n const int has_diag,\n const int *A_rows,\n const int *A_cols,\n const int *A_diag,\n const Value_type *A_vals,\n const int *A_perm,\n const int *B_rows,\n const int *B_cols,\n const int *B_diag,\n const Value_type *B_vals,\n const int *B_perm,\n int *ok )\n{\n int row_id = blockIdx.x * blockDim.x + threadIdx.x;\n\n for ( ; row_id < num_rows ; row_id += gridDim.x * blockDim.x )\n {\n if ( has_diag )\n {\n int a_it = A_diag[row_id];\n int b_it = B_diag[row_id];\n Value_type a_val = A_vals[a_it];\n Value_type b_val = A_vals[b_it];\n\n if ( !equal( a_val, b_val ) )\n {\n printf( \"ERROR: row=%d: Diag A[%d]=%f and B[%d]=%f are different!!!\\n\", row_id, a_it, a_val, b_it, b_val );\n ok[0] = 0;\n return;\n }\n }\n\n int row_it = A_rows[row_id + 0];\n int row_end = A_rows[row_id + 1];\n\n if ( row_it != B_rows[row_id + 0] || row_end != B_rows[row_id + 1] )\n {\n printf( \"ERROR: Rows A[%d] and B[%d] have different lenghts!!!\\n\", row_id, row_id );\n ok[0] = 0;\n return;\n }\n\n for ( ; row_it < row_end ; ++row_it )\n {\n const int a_it = A_perm[row_it];\n const int b_it = B_perm[row_it];\n const int a_col = A_cols[a_it];\n const int b_col = B_cols[b_it];\n\n if ( a_col != b_col )\n {\n printf( \"ERROR: row=%d row_it=%d: Cols A[%d]=%d and B[%d]=%d are different!!!\\n\", row_id, row_it, a_it, a_col, b_it, b_col );\n ok[0] = 0;\n return;\n }\n\n for ( int k = 0 ; k < block_size_sq ; ++k )\n {\n Value_type a_val = A_vals[block_size_sq * a_it + k];\n Value_type b_val = B_vals[block_size_sq * b_it + k];\n\n if ( !equal( a_val, b_val ) )\n {\n printf( \"ERROR: row=%d row_it=%d: Vals A[%d]=%f and B[%d]=%f are different!!!\\n\", row_id, row_it, a_it, a_val, b_it, b_val );\n ok[0] = 0;\n return;\n }\n }\n }\n }\n}\n\n\n\n\n\n\n// parameter is used as test name\nDECLARE_UNITTEST_BEGIN(AggregatesCoarseGeneratorTest);\n\n// setup restriction on HOST\nvoid fillRowOffsetsAndColIndices(const int num_aggregates,\n typename Matrix_h::IVector aggregates,\n const int R_num_cols,\n typename Matrix_h::IVector &R_row_offsets,\n typename Matrix_h::IVector &R_col_indices)\n{\n for (int i = 0; i < num_aggregates + 1; i++)\n {\n R_row_offsets[i] = 0;\n }\n\n // Count number of neighbors for each row\n for (int i = 0; i < R_num_cols; i++)\n {\n int I = aggregates[i];\n R_row_offsets[I]++;\n } \n\n R_row_offsets[num_aggregates] = R_num_cols;\n\n for (int i = num_aggregates - 1; i >= 0; i--)\n {\n R_row_offsets[i] = R_row_offsets[i + 1] - R_row_offsets[i];\n }\n\n /* Set column indices. */\n for (int i = 0; i < R_num_cols; i++)\n {\n int I = aggregates[i];\n int Ip = R_row_offsets[I]++;\n R_col_indices[Ip] = i;\n }\n\n /* Reset r[i] to start of row memory. */\n for (int i = num_aggregates - 1; i > 0; i--)\n {\n R_row_offsets[i] = R_row_offsets[i - 1];\n }\n\n R_row_offsets[0] = 0;\n}\n\ntemplate< typename Matrix_type >\nvoid compare_matrices( const Matrix_type &A, const Matrix_type &B )\n{\n typedef typename Matrix_type::IVector IVector;\n IVector A_perm;\n build_sort_permutation( A, A_perm );\n IVector B_perm;\n build_sort_permutation( B, B_perm );\n IVector ok( 1, 1 );\n compare_matrices_kernel <<< 2048, 256>>>( A.get_num_rows(),\n A.get_block_dimx() * A.get_block_dimy(),\n A.hasProps(DIAG),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n A_perm.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n B.diag.raw(),\n B.values.raw(),\n B_perm.raw(),\n ok.raw() );\n cudaCheckError();\n UNITTEST_ASSERT_TRUE( ok[0] == 1 );\n}\n\nvoid run()\n{\n cudaCheckError();\n this->randomize( 375139 );\n AMG_Config cfg;\n const std::string &cfg_scope = \"default\";\n // setup generators\n cudaCheckError();\n aggregation::CoarseAGenerator *d_generator = new aggregation::LowDegCoarseAGenerator(cfg, cfg_scope);\n this->PrintOnFail(\"Device generator creation\");\n UNITTEST_ASSERT_TRUE(d_generator != NULL);\n cudaCheckError();\n aggregation::CoarseAGenerator *h_generator = new aggregation::LowDegCoarseAGenerator(cfg, cfg_scope);\n this->PrintOnFail(\"Host generator creation\");\n UNITTEST_ASSERT_TRUE(h_generator != NULL);\n cudaCheckError();\n MatrixA A, Ac;\n int bsizes[] = {1, 2, 3, 4, 5, 8, 10};\n\n for (int diag_prop = 0; diag_prop < 2; diag_prop++)\n {\n for (int bs = 0; bs < sizeof(bsizes)/sizeof(bsizes[0]); bs++)\n {\n int b = bsizes[bs];\n // setup matrix A\n cudaCheckError();\n generateMatrixRandomStruct::generate(A, 128, diag_prop, b, false);\n cudaCheckError();\n random_fill(A);\n cudaCheckError();\n Matrix_h h_A = A;\n cudaCheckError();\n Matrix_h h_Ac;\n cudaCheckError();\n // setup aggregates on CPU\n int num_aggregates = (A.get_num_rows() - 1) / 2 + 1; //A.get_num_rows();\n typename Matrix_h::IVector h_aggregates;\n h_aggregates.resize( A.get_num_rows() );\n\n for ( int i = 0; i < h_aggregates.size(); i++ )\n {\n h_aggregates[i] = i / 2; //i;\n }\n\n // setup R matrix on CPU\n typename Matrix_h::IVector h_R_row_offsets;\n typename Matrix_h::IVector h_R_col_indices;\n h_R_row_offsets.resize( num_aggregates + 1 );\n h_R_col_indices.resize( A.get_num_rows() );\n fillRowOffsetsAndColIndices( num_aggregates, h_aggregates, A.get_num_rows(), h_R_row_offsets, h_R_col_indices );\n // assign GPU vectors\n IVector aggregates = h_aggregates;\n IVector R_row_offsets = h_R_row_offsets;\n IVector R_col_indices = h_R_col_indices;\n // compute Galerkin product on CPU and GPU\n h_generator->computeAOperator(h_A, h_Ac, h_aggregates, h_R_row_offsets, h_R_col_indices, num_aggregates);\n d_generator->computeAOperator(A, Ac, aggregates, R_row_offsets, R_col_indices, num_aggregates);\n // simple check on matrix size\n this->PrintOnFail(\"Coarse matrix has wrong size %i != num aggregates %i\", Ac.get_num_rows(), num_aggregates);\n UNITTEST_ASSERT_TRUE( Ac.get_num_rows() == num_aggregates );\n // dump matrix to file\n VVector v;\n v.resize(Ac.get_num_rows() * Ac.get_block_dimy());\n random_fill(v);\n // compare structure\n this->PrintOnFail(\"Coarse matrix has incorrect structure, diag prop %i, block size %i, num rows %i, num aggregates %i\", diag_prop, b, A.get_num_rows(), num_aggregates);\n compare_matrices(Ac, MatrixA(h_Ac) /*, b==1 && !diag_prop */ );\n }\n }\n\n delete d_generator;\n delete h_generator;\n}\n\nDECLARE_UNITTEST_END(AggregatesCoarseGeneratorTest);\n\n// run for all device configs\n#define AMGX_CASE_LINE(CASE) AggregatesCoarseGeneratorTest::Type> AggregatesCoarseGeneratorTest_##CASE;\nAMGX_FORALL_BUILDS_DEVICE(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n//AggregatesCoarseGeneratorTest::Type> AggregatesCoarseGeneratorTest_dDDI;\n\n\n} //namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \"unit_test.h\"\n#include \"amg_solver.h\"\n#include \"aggregation/coarseAgenerators/low_deg_coarse_A_generator.h\"\n#include \n#include \n#include \"test_utils.h\"\n#include \"util.h\"\n#include \"time.h\"\n#include \n#include \n\nnamespace amgx\n{\n\n\nstruct Not_empty_row\n{\n typedef bool result_type;\n template< typename Tuple >\n inline __device__ __host__ bool operator()( const Tuple &t ) const\n {\n return amgx::thrust::get<0>(t) != amgx::thrust::get<1>(t);\n }\n};\n\ntemplate< typename Matrix, typename Vector >\nstatic\nvoid build_sort_permutation( const Matrix &M, Vector &permutation )\n{\n int num_nz = M.get_num_nz();\n Vector row_indices( num_nz, 0 );\n amgx::thrust::scatter_if( amgx::thrust::counting_iterator(0),\n amgx::thrust::counting_iterator(M.row_offsets.size() - 1),\n M.row_offsets.begin(),\n amgx::thrust::make_transform_iterator(\n amgx::thrust::make_zip_iterator( amgx::thrust::make_tuple( M.row_offsets.begin(), M.row_offsets.begin() + 1 ) ),\n Not_empty_row()),\n row_indices.begin());\n cudaCheckError();\n amgx::thrust::inclusive_scan( row_indices.begin(), row_indices.begin() + M.get_num_nz(), row_indices.begin(), amgx::thrust::maximum() );\n cudaCheckError();\n permutation.resize( num_nz );\n thrust_wrapper::sequence( permutation.begin(), permutation.end() );\n cudaCheckError();\n Vector tmp( M.col_indices );\n amgx::thrust::stable_sort_by_key( tmp.begin(), tmp.end(), permutation.begin() );\n cudaCheckError();\n tmp = row_indices;\n amgx::thrust::gather( permutation.begin(), permutation.end(), tmp.begin(), row_indices.begin() );\n cudaCheckError();\n amgx::thrust::stable_sort_by_key( row_indices.begin(), row_indices.end(), permutation.begin() );\n cudaCheckError();\n}\n\ntemplate< typename Value_type >\nstatic __device__ __inline__\nbool equal( Value_type x, Value_type y, Value_type epsilon, Value_type max_abs_error )\n{\n if ( x == y )\n {\n return true;\n }\n\n if ( abs(x - y) < max_abs_error )\n {\n return true;\n }\n\n if ( abs(x - y) <= epsilon )\n {\n return true;\n }\n\n return false;\n}\n\nstatic __device__ __inline__\nbool equal( float x, float y )\n{\n return equal( x, y, 1.0e-6f, 1.0e10f * FLT_MIN );\n}\n\nstatic __device__ __inline__\nbool equal( double x, double y )\n{\n return equal( x, y, 1.0e-12, 1.0e10 * DBL_MIN );\n}\n\ntemplate< typename Value_type >\n__global__\nvoid compare_matrices_kernel( const int num_rows,\n const int block_size_sq,\n const int has_diag,\n const int *A_rows,\n const int *A_cols,\n const int *A_diag,\n const Value_type *A_vals,\n const int *A_perm,\n const int *B_rows,\n const int *B_cols,\n const int *B_diag,\n const Value_type *B_vals,\n const int *B_perm,\n int *ok )\n{\n int row_id = blockIdx.x * blockDim.x + threadIdx.x;\n\n for ( ; row_id < num_rows ; row_id += gridDim.x * blockDim.x )\n {\n if ( has_diag )\n {\n int a_it = A_diag[row_id];\n int b_it = B_diag[row_id];\n Value_type a_val = A_vals[a_it];\n Value_type b_val = A_vals[b_it];\n\n if ( !equal( a_val, b_val ) )\n {\n printf( \"ERROR: row=%d: Diag A[%d]=%f and B[%d]=%f are different!!!\\n\", row_id, a_it, a_val, b_it, b_val );\n ok[0] = 0;\n return;\n }\n }\n\n int row_it = A_rows[row_id + 0];\n int row_end = A_rows[row_id + 1];\n\n if ( row_it != B_rows[row_id + 0] || row_end != B_rows[row_id + 1] )\n {\n printf( \"ERROR: Rows A[%d] and B[%d] have different lenghts!!!\\n\", row_id, row_id );\n ok[0] = 0;\n return;\n }\n\n for ( ; row_it < row_end ; ++row_it )\n {\n const int a_it = A_perm[row_it];\n const int b_it = B_perm[row_it];\n const int a_col = A_cols[a_it];\n const int b_col = B_cols[b_it];\n\n if ( a_col != b_col )\n {\n printf( \"ERROR: row=%d row_it=%d: Cols A[%d]=%d and B[%d]=%d are different!!!\\n\", row_id, row_it, a_it, a_col, b_it, b_col );\n ok[0] = 0;\n return;\n }\n\n for ( int k = 0 ; k < block_size_sq ; ++k )\n {\n Value_type a_val = A_vals[block_size_sq * a_it + k];\n Value_type b_val = B_vals[block_size_sq * b_it + k];\n\n if ( !equal( a_val, b_val ) )\n {\n printf( \"ERROR: row=%d row_it=%d: Vals A[%d]=%f and B[%d]=%f are different!!!\\n\", row_id, row_it, a_it, a_val, b_it, b_val );\n ok[0] = 0;\n return;\n }\n }\n }\n }\n}\n\n\n\n\n\n\n// parameter is used as test name\nDECLARE_UNITTEST_BEGIN(AggregatesCoarseGeneratorTest);\n\n// setup restriction on HOST\nvoid fillRowOffsetsAndColIndices(const int num_aggregates,\n typename Matrix_h::IVector aggregates,\n const int R_num_cols,\n typename Matrix_h::IVector &R_row_offsets,\n typename Matrix_h::IVector &R_col_indices)\n{\n for (int i = 0; i < num_aggregates + 1; i++)\n {\n R_row_offsets[i] = 0;\n }\n\n // Count number of neighbors for each row\n for (int i = 0; i < R_num_cols; i++)\n {\n int I = aggregates[i];\n R_row_offsets[I]++;\n } \n\n R_row_offsets[num_aggregates] = R_num_cols;\n\n for (int i = num_aggregates - 1; i >= 0; i--)\n {\n R_row_offsets[i] = R_row_offsets[i + 1] - R_row_offsets[i];\n }\n\n /* Set column indices. */\n for (int i = 0; i < R_num_cols; i++)\n {\n int I = aggregates[i];\n int Ip = R_row_offsets[I]++;\n R_col_indices[Ip] = i;\n }\n\n /* Reset r[i] to start of row memory. */\n for (int i = num_aggregates - 1; i > 0; i--)\n {\n R_row_offsets[i] = R_row_offsets[i - 1];\n }\n\n R_row_offsets[0] = 0;\n}\n\ntemplate< typename Matrix_type >\nvoid compare_matrices( const Matrix_type &A, const Matrix_type &B )\n{\n typedef typename Matrix_type::IVector IVector;\n IVector A_perm;\n build_sort_permutation( A, A_perm );\n IVector B_perm;\n build_sort_permutation( B, B_perm );\n IVector ok( 1, 1 );\n compare_matrices_kernel <<< 2048, 256>>>( A.get_num_rows(),\n A.get_block_dimx() * A.get_block_dimy(),\n A.hasProps(DIAG),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n A_perm.raw(),\n B.row_offsets.raw(),\n B.col_indices.raw(),\n B.diag.raw(),\n B.values.raw(),\n B_perm.raw(),\n ok.raw() );\n cudaCheckError();\n UNITTEST_ASSERT_TRUE( ok[0] == 1 );\n}\n\nvoid run()\n{\n cudaCheckError();\n this->randomize( 375139 );\n AMG_Config cfg;\n const std::string &cfg_scope = \"default\";\n // setup generators\n cudaCheckError();\n aggregation::CoarseAGenerator *d_generator = new aggregation::LowDegCoarseAGenerator(cfg, cfg_scope);\n this->PrintOnFail(\"Device generator creation\");\n UNITTEST_ASSERT_TRUE(d_generator != NULL);\n cudaCheckError();\n aggregation::CoarseAGenerator *h_generator = new aggregation::LowDegCoarseAGenerator(cfg, cfg_scope);\n this->PrintOnFail(\"Host generator creation\");\n UNITTEST_ASSERT_TRUE(h_generator != NULL);\n cudaCheckError();\n MatrixA A, Ac;\n int bsizes[] = {1, 2, 3, 4, 5, 8, 10};\n\n for (int diag_prop = 0; diag_prop < 2; diag_prop++)\n {\n for (int bs = 0; bs < sizeof(bsizes)/sizeof(bsizes[0]); bs++)\n {\n int b = bsizes[bs];\n // setup matrix A\n cudaCheckError();\n generateMatrixRandomStruct::generate(A, 128, diag_prop, b, false);\n cudaCheckError();\n random_fill(A);\n cudaCheckError();\n Matrix_h h_A = A;\n cudaCheckError();\n Matrix_h h_Ac;\n cudaCheckError();\n // setup aggregates on CPU\n int num_aggregates = (A.get_num_rows() - 1) / 2 + 1; //A.get_num_rows();\n typename Matrix_h::IVector h_aggregates;\n h_aggregates.resize( A.get_num_rows() );\n\n for ( int i = 0; i < h_aggregates.size(); i++ )\n {\n h_aggregates[i] = i / 2; //i;\n }\n\n // setup R matrix on CPU\n typename Matrix_h::IVector h_R_row_offsets;\n typename Matrix_h::IVector h_R_col_indices;\n h_R_row_offsets.resize( num_aggregates + 1 );\n h_R_col_indices.resize( A.get_num_rows() );\n fillRowOffsetsAndColIndices( num_aggregates, h_aggregates, A.get_num_rows(), h_R_row_offsets, h_R_col_indices );\n // assign GPU vectors\n IVector aggregates = h_aggregates;\n IVector R_row_offsets = h_R_row_offsets;\n IVector R_col_indices = h_R_col_indices;\n // compute Galerkin product on CPU and GPU\n h_generator->computeAOperator(h_A, h_Ac, h_aggregates, h_R_row_offsets, h_R_col_indices, num_aggregates);\n d_generator->computeAOperator(A, Ac, aggregates, R_row_offsets, R_col_indices, num_aggregates);\n // simple check on matrix size\n this->PrintOnFail(\"Coarse matrix has wrong size %i != num aggregates %i\", Ac.get_num_rows(), num_aggregates);\n UNITTEST_ASSERT_TRUE( Ac.get_num_rows() == num_aggregates );\n // dump matrix to file\n VVector v;\n v.resize(Ac.get_num_rows() * Ac.get_block_dimy());\n random_fill(v);\n // compare structure\n this->PrintOnFail(\"Coarse matrix has incorrect structure, diag prop %i, block size %i, num rows %i, num aggregates %i\", diag_prop, b, A.get_num_rows(), num_aggregates);\n compare_matrices(Ac, MatrixA(h_Ac) /*, b==1 && !diag_prop */ );\n }\n }\n\n delete d_generator;\n delete h_generator;\n}\n\nDECLARE_UNITTEST_END(AggregatesCoarseGeneratorTest);\n\n// run for all device configs\n#define AMGX_CASE_LINE(CASE) AggregatesCoarseGeneratorTest::Type> AggregatesCoarseGeneratorTest_##CASE;\nAMGX_FORALL_BUILDS_DEVICE(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n//AggregatesCoarseGeneratorTest::Type> AggregatesCoarseGeneratorTest_dDDI;\n\n\n} //namespace amgx\n", "source": "/kaggle/working/amgx/src/tests/aggregates_coarse_generator.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#define AMGX_EXPLICIT_SHIFT 1\nnamespace amgx\n{\n\ntemplate \n__global__ void shift_diagonal(index_type num_rows,\n const index_type *row_offsets,\n const index_type *col_indices,\n value_type *values,\n value_type shift)\n{\n index_type tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int r = tidx; r < num_rows; r += blockDim.x * gridDim.x)\n {\n index_type row_start = row_offsets[r];\n index_type row_end = row_offsets[r + 1];\n\n for (int j = row_start; j < row_end; j++)\n {\n if (col_indices[j] == r)\n {\n values[j] += shift;\n continue;\n }\n }\n }\n}\n// OBSOLETE compute a when the input is H only\ntemplate \n__global__ void dangling_nodes(index_type num_rows,\n const index_type *row_offsets,\n value_type *aa)\n{\n index_type tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int r = tidx; r < num_rows; r += blockDim.x * gridDim.x)\n {\n index_type row_start = row_offsets[r];\n index_type row_end = row_offsets[r + 1];\n\n // NOTE 1 : a = alpha*a + (1-alpha)e\n // NOTE 2 : a is initialized to (1-alpha)\n if (row_start == row_end)\n {\n aa[r] = 1.0; // NOTE 3 : alpha*1 + (1-alpha)*1 = 1.0\n }\n }\n}\n// used when a is given as input and the matrix is H^T\ntemplate \n__global__ void update_a(index_type num_rows,\n value_type *aa,\n value_type beta)\n{\n index_type tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int r = tidx; r < num_rows; r += blockDim.x * gridDim.x)\n {\n // NOTE 1 : a = alpha*a + (1-alpha)e\n if (aa[r] == 0.0)\n {\n aa[r] = beta; // NOTE 2 : alpha*0 + (1-alpha)*1 = (1-alpha)\n }\n }\n}\n\ntemplate \nSingleIteration_EigenSolver::SingleIteration_EigenSolver(AMG_Config &cfg,\n const std::string &cfg_scope)\n : Base(cfg, cfg_scope), m_cfg(cfg), m_operator(NULL)\n{\n m_convergence_check_freq = cfg.getParameter(\"eig_convergence_check_freq\", cfg_scope);\n}\n\ntemplate \nSingleIteration_EigenSolver::~SingleIteration_EigenSolver()\n{\n /*if (m_operator)\n delete m_operator;*/\n if (this->m_which == EIG_SMALLEST)\n {\n delete m_operator;\n }\n\n if (this->m_which == EIG_PAGERANK)\n {\n delete m_operator;\n }\n\n free_allocated();\n}\n\ntemplate \nvoid SingleIteration_EigenSolver::free_allocated()\n{\n}\n\ntemplate \nvoid SingleIteration_EigenSolver::shift_matrix()\n{\n ValueTypeMat shift = this->m_shift;\n\n if (shift == 0)\n {\n return;\n }\n\n Matrix *pA = dynamic_cast< Matrix* > (this->m_A);\n Matrix &A = *pA;\n int num_threads = 128;\n int max_grid_size = 4096;\n int num_rows = A.get_num_rows();\n int num_blocks = std::min(max_grid_size, (num_rows + num_threads - 1) / num_rows);\n shift_diagonal <<< num_blocks, num_threads>>>(num_rows, A.row_offsets.raw(),\n A.col_indices.raw(), A.values.raw(), -shift);\n cudaCheckError();\n}\n// OBSOLETE compute a when the input is H only\ntemplate \nvoid SingleIteration_EigenSolver::get_dangling_nodes()\n{\n Matrix *pA = dynamic_cast< Matrix* > (this->m_A);\n Matrix &A = *pA;\n int num_threads = 128;\n int max_grid_size = 4096;\n int num_rows = A.get_num_rows();\n int num_blocks = std::min(max_grid_size, (num_rows + num_threads - 1) / num_rows);\n dangling_nodes <<< num_blocks, num_threads>>>(num_rows, A.row_offsets.raw(), m_a.raw());\n cudaCheckError();\n // CPU CODE : you shouldn't use it, it is very slow\n /* for (int i = 0; i < num_rows; ++i)\n {\n if (A.row_offsets[i] == A.row_offsets[i+1])\n m_a[i] = 1.0; // alpha*1 + (1-alpha)*1 = 1.0\n else\n m_a[i] = beta;\n } */\n}\n\n// 0 are replaced by 1-alpha\ntemplate \nvoid SingleIteration_EigenSolver::update_dangling_nodes()\n{\n Operator &A = *this->m_A;\n ValueTypeVec beta = 1.0 - this->m_damping_factor;\n int num_rows = A.get_num_rows();\n\n // CPU\n if (TConfig::memSpace == AMGX_host)\n {\n for (int i = 0; i < num_rows; i++)\n {\n // NOTE 1 : a = alpha*a + (1-alpha)e\n if (m_a[i] == 0.0)\n {\n m_a[i] = beta; // NOTE 2 : alpha*0 + (1-alpha)*1 = (1-alpha)\n }\n }\n }\n //GPU\n else\n {\n int num_threads = 128;\n int max_grid_size = 4096;\n int num_blocks = std::min(max_grid_size, (num_rows + num_threads - 1) / num_rows);\n update_a <<< num_blocks, num_threads>>>(num_rows, m_a.raw(), beta );\n cudaCheckError();\n }\n}\n\ntemplate \nvoid SingleIteration_EigenSolver::solver_setup()\n{\n Operator &A = *this->m_A;\n ViewType oldView = A.currentView();\n A.setViewExterior();\n\n if (this->m_which == EIG_PAGERANK)\n {\n PagerankOperator *op = new PagerankOperator(A, &m_a, &m_b, this->m_damping_factor);\n m_operator = op;\n }\n else if (this->m_which == EIG_SMALLEST)\n {\n Solver *solver = SolverFactory::allocate(m_cfg, \"default\", \"solver\");\n#ifdef AMGX_EXPLICIT_SHIFT\n shift_matrix();\n SolveOperator *solve_op = new SolveOperator(A, *solver);\n#else\n ShiftedOperator *op = new ShiftedOperator(A, -this->m_shift);\n SolveOperator *solve_op = new SolveOperator(*op, *solver);\n#endif\n solve_op->setup();\n m_operator = solve_op;\n }\n else\n {\n m_operator = &A;\n }\n\n const int N = static_cast(A.get_num_cols() * A.get_block_dimy());\n // Allocate two vectors.\n m_v.resize(N);\n m_x.resize(N);\n m_allocated_vectors.push_back(&m_v);\n m_allocated_vectors.push_back(&m_x);\n\n // Vectors \"a\" and \"b\" are needed only for PR\n if (this->m_which == EIG_PAGERANK)\n {\n m_a.resize(N);\n m_b.resize(N);\n m_allocated_vectors.push_back(&m_a);\n m_allocated_vectors.push_back(&m_b);\n }\n\n int start_tag = 100;\n\n for (int i = 0; i < m_allocated_vectors.size(); ++i)\n {\n VVector *v = m_allocated_vectors[i];\n v->tag = start_tag + i;\n v->set_block_dimy(A.get_block_dimy());\n v->set_block_dimx(1);\n v->dirtybit = 1;\n v->delayed_send = 1;\n }\n\n A.setView(oldView);\n}\n\ntemplate \nvoid SingleIteration_EigenSolver::solver_pagerank_setup(VVector &a)\n{\n if (this->m_which == EIG_PAGERANK)\n {\n Matrix *pA = dynamic_cast< Matrix* > (this->m_A);\n Matrix &A = *pA;\n ViewType oldView = A.currentView();\n A.setViewExterior();\n int offset, size;\n A.getOffsetAndSizeForView(A.getViewExterior(), &offset, &size);\n copy(a, m_a, offset, size);\n //Get the number of rows of the matrix\n int num_rows_loc = A.get_num_rows();\n int num_rows_glob = num_rows_loc;\n // MPI?\n#ifdef AMGX_WITH_MPI\n int mpi_initialized = 0;\n MPI_Initialized(&mpi_initialized);\n\n if (mpi_initialized)\n {\n if (A.is_matrix_distributed())\n {\n A.getManager()->global_reduce_sum(&num_rows_glob);\n }\n }\n\n#endif\n // Requiered to compute b\n // a = alpha*a + (1-alpha)e\n update_dangling_nodes();\n // b is a constant and uniform vector\n ValueTypeMat tmp = 1.0 / num_rows_glob;\n fill(m_b, tmp);\n A.setView(oldView);\n }\n}\n\ntemplate \nvoid SingleIteration_EigenSolver::solve_init(VVector &x)\n{\n Operator &A = *this->m_A;\n ViewType oldView = A.currentView();\n A.setViewExterior();\n int offset, size;\n A.getOffsetAndSizeForView(A.getViewExterior(), &offset, &size);\n copy(x, m_x, offset, size);\n std::swap(m_x, m_v);\n A.setView(oldView);\n}\n\n// One iteration of the single iteration.\n// MATLAB code of the algorithm:\n// while true\n// V = X / norm(X);\n// X = linsolve(A, V);\n// lambda = V' * X;\n// R = X - lambda * V;\n// if norm(R) < tol * abs(lambda)\n// break;\n// end\n\ntemplate \nbool SingleIteration_EigenSolver::solve_iteration(VVector &x)\n{\n Operator &A = *this->m_A;\n ViewType oldView = A.currentView();\n A.setViewExterior();\n int offset, size;\n A.getOffsetAndSizeForView(A.getViewExterior(), &offset, &size);\n // V = X / norm(X)\n scal(m_v, ValueTypeVec(1) / get_norm(A, m_v, this->m_norm_type), offset, size);\n // X = linsolve(A, v)\n m_operator->apply(m_v, m_x);\n\n if ((this->m_curr_iter % this->m_convergence_check_freq) == 0)\n {\n ValueTypeVec lambda;\n\n if (this->m_which == EIG_PAGERANK)\n {\n // The maximum eigenvalue of G is 1.0\n lambda = 1.0;\n // v = x - v\n // axpy(m_x, m_v, ValueTypeVec(-1),offset, size);\n axpby(m_v, m_x, m_v, ValueTypeVec(-1.0), ValueTypeVec(1), offset, size);\n }\n else\n {\n // lambda = v.x\n lambda = dot(A, m_v, m_x);\n // v = x - lambda * v\n axpby(m_v, m_x, m_v, -lambda, ValueTypeVec(1), offset, size);\n }\n\n ValueTypeVec residual_norm = get_norm(A, m_v, this->m_norm_type);\n this->m_residuals.push_back(residual_norm / fabs(lambda));\n\n // Check convergence.\n if (residual_norm < this->m_tolerance * fabs(lambda))\n {\n // Normalize eigenvector.\n if (this->m_which == EIG_PAGERANK)\n {\n //Norm L1 is more fited for the output of PageRank\n ValueTypeVec norm = get_norm(A, m_x, L1);\n scal(m_x, ValueTypeVec(1) / norm, offset, size);\n }\n else\n {\n ValueTypeVec norm = get_norm(A, m_x, this->m_norm_type);\n scal(m_x, ValueTypeVec(1) / norm, offset, size);\n }\n\n // With inverse iteration we need to scale the eigenvector by the inverse of the eigenvalue,\n // but this doesn't seems to be needed since we already normalized the eigenvector just above.\n this->m_eigenvectors.push_back(m_x);\n copy(m_x, x, offset, size);\n this->m_eigenvalues.push_back(lambda);\n this->postprocess_eigenpairs();\n return true;\n }\n }\n\n std::swap(m_x, m_v);\n A.setView(oldView);\n return false;\n}\n\ntemplate \nvoid SingleIteration_EigenSolver::solve_finalize()\n{\n}\n\n// Explicit template instantiation.\n#define AMGX_CASE_LINE(CASE) template class SingleIteration_EigenSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n};\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#define AMGX_EXPLICIT_SHIFT 1\nnamespace amgx\n{\n\ntemplate \n__global__ void shift_diagonal(index_type num_rows,\n const index_type *row_offsets,\n const index_type *col_indices,\n value_type *values,\n value_type shift)\n{\n index_type tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int r = tidx; r < num_rows; r += blockDim.x * gridDim.x)\n {\n index_type row_start = row_offsets[r];\n index_type row_end = row_offsets[r + 1];\n\n for (int j = row_start; j < row_end; j++)\n {\n if (col_indices[j] == r)\n {\n values[j] += shift;\n continue;\n }\n }\n }\n}\n// OBSOLETE compute a when the input is H only\ntemplate \n__global__ void dangling_nodes(index_type num_rows,\n const index_type *row_offsets,\n value_type *aa)\n{\n index_type tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int r = tidx; r < num_rows; r += blockDim.x * gridDim.x)\n {\n index_type row_start = row_offsets[r];\n index_type row_end = row_offsets[r + 1];\n\n // NOTE 1 : a = alpha*a + (1-alpha)e\n // NOTE 2 : a is initialized to (1-alpha)\n if (row_start == row_end)\n {\n aa[r] = 1.0; // NOTE 3 : alpha*1 + (1-alpha)*1 = 1.0\n }\n }\n}\n// used when a is given as input and the matrix is H^T\ntemplate \n__global__ void update_a(index_type num_rows,\n value_type *aa,\n value_type beta)\n{\n index_type tidx = blockDim.x * blockIdx.x + threadIdx.x;\n\n for (int r = tidx; r < num_rows; r += blockDim.x * gridDim.x)\n {\n // NOTE 1 : a = alpha*a + (1-alpha)e\n if (aa[r] == 0.0)\n {\n aa[r] = beta; // NOTE 2 : alpha*0 + (1-alpha)*1 = (1-alpha)\n }\n }\n}\n\ntemplate \nSingleIteration_EigenSolver::SingleIteration_EigenSolver(AMG_Config &cfg,\n const std::string &cfg_scope)\n : Base(cfg, cfg_scope), m_cfg(cfg), m_operator(NULL)\n{\n m_convergence_check_freq = cfg.getParameter(\"eig_convergence_check_freq\", cfg_scope);\n}\n\ntemplate \nSingleIteration_EigenSolver::~SingleIteration_EigenSolver()\n{\n /*if (m_operator)\n delete m_operator;*/\n if (this->m_which == EIG_SMALLEST)\n {\n delete m_operator;\n }\n\n if (this->m_which == EIG_PAGERANK)\n {\n delete m_operator;\n }\n\n free_allocated();\n}\n\ntemplate \nvoid SingleIteration_EigenSolver::free_allocated()\n{\n}\n\ntemplate \nvoid SingleIteration_EigenSolver::shift_matrix()\n{\n ValueTypeMat shift = this->m_shift;\n\n if (shift == 0)\n {\n return;\n }\n\n Matrix *pA = dynamic_cast< Matrix* > (this->m_A);\n Matrix &A = *pA;\n int num_threads = 128;\n int max_grid_size = 4096;\n int num_rows = A.get_num_rows();\n int num_blocks = std::min(max_grid_size, (num_rows + num_threads - 1) / num_rows);\n shift_diagonal <<< num_blocks, num_threads>>>(num_rows, A.row_offsets.raw(),\n A.col_indices.raw(), A.values.raw(), -shift);\n cudaCheckError();\n}\n// OBSOLETE compute a when the input is H only\ntemplate \nvoid SingleIteration_EigenSolver::get_dangling_nodes()\n{\n Matrix *pA = dynamic_cast< Matrix* > (this->m_A);\n Matrix &A = *pA;\n int num_threads = 128;\n int max_grid_size = 4096;\n int num_rows = A.get_num_rows();\n int num_blocks = std::min(max_grid_size, (num_rows + num_threads - 1) / num_rows);\n dangling_nodes <<< num_blocks, num_threads>>>(num_rows, A.row_offsets.raw(), m_a.raw());\n cudaCheckError();\n // CPU CODE : you shouldn't use it, it is very slow\n /* for (int i = 0; i < num_rows; ++i)\n {\n if (A.row_offsets[i] == A.row_offsets[i+1])\n m_a[i] = 1.0; // alpha*1 + (1-alpha)*1 = 1.0\n else\n m_a[i] = beta;\n } */\n}\n\n// 0 are replaced by 1-alpha\ntemplate \nvoid SingleIteration_EigenSolver::update_dangling_nodes()\n{\n Operator &A = *this->m_A;\n ValueTypeVec beta = 1.0 - this->m_damping_factor;\n int num_rows = A.get_num_rows();\n\n // CPU\n if (TConfig::memSpace == AMGX_host)\n {\n for (int i = 0; i < num_rows; i++)\n {\n // NOTE 1 : a = alpha*a + (1-alpha)e\n if (m_a[i] == 0.0)\n {\n m_a[i] = beta; // NOTE 2 : alpha*0 + (1-alpha)*1 = (1-alpha)\n }\n }\n }\n //GPU\n else\n {\n int num_threads = 128;\n int max_grid_size = 4096;\n int num_blocks = std::min(max_grid_size, (num_rows + num_threads - 1) / num_rows);\n update_a <<< num_blocks, num_threads>>>(num_rows, m_a.raw(), beta );\n cudaCheckError();\n }\n}\n\ntemplate \nvoid SingleIteration_EigenSolver::solver_setup()\n{\n Operator &A = *this->m_A;\n ViewType oldView = A.currentView();\n A.setViewExterior();\n\n if (this->m_which == EIG_PAGERANK)\n {\n PagerankOperator *op = new PagerankOperator(A, &m_a, &m_b, this->m_damping_factor);\n m_operator = op;\n }\n else if (this->m_which == EIG_SMALLEST)\n {\n Solver *solver = SolverFactory::allocate(m_cfg, \"default\", \"solver\");\n#ifdef AMGX_EXPLICIT_SHIFT\n shift_matrix();\n SolveOperator *solve_op = new SolveOperator(A, *solver);\n#else\n ShiftedOperator *op = new ShiftedOperator(A, -this->m_shift);\n SolveOperator *solve_op = new SolveOperator(*op, *solver);\n#endif\n solve_op->setup();\n m_operator = solve_op;\n }\n else\n {\n m_operator = &A;\n }\n\n const int N = static_cast(A.get_num_cols() * A.get_block_dimy());\n // Allocate two vectors.\n m_v.resize(N);\n m_x.resize(N);\n m_allocated_vectors.push_back(&m_v);\n m_allocated_vectors.push_back(&m_x);\n\n // Vectors \"a\" and \"b\" are needed only for PR\n if (this->m_which == EIG_PAGERANK)\n {\n m_a.resize(N);\n m_b.resize(N);\n m_allocated_vectors.push_back(&m_a);\n m_allocated_vectors.push_back(&m_b);\n }\n\n int start_tag = 100;\n\n for (int i = 0; i < m_allocated_vectors.size(); ++i)\n {\n VVector *v = m_allocated_vectors[i];\n v->tag = start_tag + i;\n v->set_block_dimy(A.get_block_dimy());\n v->set_block_dimx(1);\n v->dirtybit = 1;\n v->delayed_send = 1;\n }\n\n A.setView(oldView);\n}\n\ntemplate \nvoid SingleIteration_EigenSolver::solver_pagerank_setup(VVector &a)\n{\n if (this->m_which == EIG_PAGERANK)\n {\n Matrix *pA = dynamic_cast< Matrix* > (this->m_A);\n Matrix &A = *pA;\n ViewType oldView = A.currentView();\n A.setViewExterior();\n int offset, size;\n A.getOffsetAndSizeForView(A.getViewExterior(), &offset, &size);\n copy(a, m_a, offset, size);\n //Get the number of rows of the matrix\n int num_rows_loc = A.get_num_rows();\n int num_rows_glob = num_rows_loc;\n // MPI?\n#ifdef AMGX_WITH_MPI\n int mpi_initialized = 0;\n MPI_Initialized(&mpi_initialized);\n\n if (mpi_initialized)\n {\n if (A.is_matrix_distributed())\n {\n A.getManager()->global_reduce_sum(&num_rows_glob);\n }\n }\n\n#endif\n // Requiered to compute b\n // a = alpha*a + (1-alpha)e\n update_dangling_nodes();\n // b is a constant and uniform vector\n ValueTypeMat tmp = 1.0 / num_rows_glob;\n fill(m_b, tmp);\n A.setView(oldView);\n }\n}\n\ntemplate \nvoid SingleIteration_EigenSolver::solve_init(VVector &x)\n{\n Operator &A = *this->m_A;\n ViewType oldView = A.currentView();\n A.setViewExterior();\n int offset, size;\n A.getOffsetAndSizeForView(A.getViewExterior(), &offset, &size);\n copy(x, m_x, offset, size);\n std::swap(m_x, m_v);\n A.setView(oldView);\n}\n\n// One iteration of the single iteration.\n// MATLAB code of the algorithm:\n// while true\n// V = X / norm(X);\n// X = linsolve(A, V);\n// lambda = V' * X;\n// R = X - lambda * V;\n// if norm(R) < tol * abs(lambda)\n// break;\n// end\n\ntemplate \nbool SingleIteration_EigenSolver::solve_iteration(VVector &x)\n{\n Operator &A = *this->m_A;\n ViewType oldView = A.currentView();\n A.setViewExterior();\n int offset, size;\n A.getOffsetAndSizeForView(A.getViewExterior(), &offset, &size);\n // V = X / norm(X)\n scal(m_v, ValueTypeVec(1) / get_norm(A, m_v, this->m_norm_type), offset, size);\n // X = linsolve(A, v)\n m_operator->apply(m_v, m_x);\n\n if ((this->m_curr_iter % this->m_convergence_check_freq) == 0)\n {\n ValueTypeVec lambda;\n\n if (this->m_which == EIG_PAGERANK)\n {\n // The maximum eigenvalue of G is 1.0\n lambda = 1.0;\n // v = x - v\n // axpy(m_x, m_v, ValueTypeVec(-1),offset, size);\n axpby(m_v, m_x, m_v, ValueTypeVec(-1.0), ValueTypeVec(1), offset, size);\n }\n else\n {\n // lambda = v.x\n lambda = dot(A, m_v, m_x);\n // v = x - lambda * v\n axpby(m_v, m_x, m_v, -lambda, ValueTypeVec(1), offset, size);\n }\n\n ValueTypeVec residual_norm = get_norm(A, m_v, this->m_norm_type);\n this->m_residuals.push_back(residual_norm / fabs(lambda));\n\n // Check convergence.\n if (residual_norm < this->m_tolerance * fabs(lambda))\n {\n // Normalize eigenvector.\n if (this->m_which == EIG_PAGERANK)\n {\n //Norm L1 is more fited for the output of PageRank\n ValueTypeVec norm = get_norm(A, m_x, L1);\n scal(m_x, ValueTypeVec(1) / norm, offset, size);\n }\n else\n {\n ValueTypeVec norm = get_norm(A, m_x, this->m_norm_type);\n scal(m_x, ValueTypeVec(1) / norm, offset, size);\n }\n\n // With inverse iteration we need to scale the eigenvector by the inverse of the eigenvalue,\n // but this doesn't seems to be needed since we already normalized the eigenvector just above.\n this->m_eigenvectors.push_back(m_x);\n copy(m_x, x, offset, size);\n this->m_eigenvalues.push_back(lambda);\n this->postprocess_eigenpairs();\n return true;\n }\n }\n\n std::swap(m_x, m_v);\n A.setView(oldView);\n return false;\n}\n\ntemplate \nvoid SingleIteration_EigenSolver::solve_finalize()\n{\n}\n\n// Explicit template instantiation.\n#define AMGX_CASE_LINE(CASE) template class SingleIteration_EigenSolver::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n};\n", "source": "/kaggle/working/amgx/src/eigensolvers/single_iteration_eigensolver.cu"} {"input": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\nnamespace\n{\n\ntemplate \n__global__\nvoid kernel_clear_lower_triangular(T *V, int dim, int lda)\n{\n int r = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (r < dim)\n {\n for (int c = 0; c < r; ++c)\n {\n V[c * lda + r] = 0.;\n }\n }\n}\n\ntemplate \nvoid clear_lower_triangular(Vector &V)\n{\n int num_threads = 128;\n int max_grid_size = 4096;\n int num_rows = V.get_num_rows();\n int num_blocks = std::min(max_grid_size, (num_rows + num_threads - 1) / num_rows);\n kernel_clear_lower_triangular <<< num_blocks, num_threads>>>(V.raw(), num_rows, V.get_lda());\n cudaCheckError();\n}\n\n\nstruct TreeNode\n{\n TreeNode(int v)\n : left(0), right(0), parent(0), value(v)\n {\n }\n\n bool is_leaf() const\n {\n return left == right;\n }\n\n bool is_root() const\n {\n return !parent;\n }\n\n TreeNode *left;\n TreeNode *right;\n TreeNode *parent;\n int value;\n};\n\n// Recursive function to print a reduction tree to an ofstream in dot\n// format.\nvoid print_tree_rec(TreeNode *n, std::ofstream &f)\n{\n TreeNode *left = n->left;\n TreeNode *right = n->right;\n TreeNode *parent = n->parent;\n // We use the address of the nodes as labels in dot format to\n // avoid having to renumber the nodes.\n size_t address = reinterpret_cast(n);\n size_t parent_address = reinterpret_cast(parent);\n size_t left_address = reinterpret_cast(left);\n size_t right_address = reinterpret_cast(right);\n f << address << \"[label=\" << n->value << \"];\" << std::endl;\n\n if (parent)\n {\n f << address << \" -> \" << parent_address << std::endl;\n }\n\n if (left)\n {\n f << address << \" -> \" << left_address << std::endl;\n print_tree_rec(left, f);\n }\n\n if (right)\n {\n f << address << \" -> \" << right_address << std::endl;\n print_tree_rec(right, f);\n }\n}\n\n#ifdef DEBUG\n// Debug function: print a reduction tree to a file in dot format.\nvoid print_tree(TreeNode *root, const char *filename)\n{\n std::ofstream f;\n f.open(filename);\n f << \"digraph A {\" << std::endl;\n print_tree_rec(root, f);\n f << \"}\" << std::endl;\n}\n\n// Create a degenerate tree of size N, this is not an efficient\n// reduction tree for a parallel operation (rank 0 is involved in all\n// operations) but it is used for testing.\n// Example with N == 4:\n// 0\n// / \\\n// 0 \\\n// / \\ \\\n// 0 \\ \\\n// / \\ \\ \\\n// 0 1 2 3\n// The label on each node is the active process for this operation.\nTreeNode *build_degenerate_tree(int N)\n{\n std::vector leaves;\n\n // Allocate one leaf node for each rank.\n for (int i = 0; i < N; ++i)\n {\n leaves.push_back(new TreeNode(i));\n }\n\n TreeNode *previous_node = leaves[0];\n\n // For each leaf except the first one, create the new root of the\n // tree labeled with 0.\n // The left child of this node is the previous root and the right\n // child is the current leaf.\n for (int i = 0; i < N - 1; ++i)\n {\n TreeNode *left = previous_node;\n TreeNode *right = leaves[i + 1];\n TreeNode *ancestor = new TreeNode(previous_node->value);\n left->parent = ancestor;\n right->parent = ancestor;\n ancestor->left = left;\n ancestor->right = right;\n previous_node = ancestor;\n }\n\n TreeNode *root = previous_node;\n return root;\n}\n\n#endif\n// Create a full tree of size N, N must be a power of 2.\n// For instance if N == 4:\n// 0\n// / \\\n// / \\\n// 0 2\n// / \\ / \\\n// 0 1 2 3\n// The label on each node is the active process for this operation.\nTreeNode *build_full_tree(int N)\n{\n // Nodes of the previous (deepest) level of the tree.\n std::vector previous_level;\n std::vector current_level;\n\n // Allocate N leaves nodes, one for each process.\n for (int i = 0; i < N; ++i)\n {\n previous_level.push_back(new TreeNode(i));\n }\n\n // Level i of the tree will have 2^i nodes, we construct the tree\n // from leaves to root. At each level we allocate new nodes and\n // connect these nodes to the previous level.\n for (int stride = 2; stride <= N; stride <<= 1)\n {\n current_level.clear();\n int children_index = 0;\n\n for (int j = 0; j < N; j += stride)\n {\n TreeNode *left = previous_level[children_index++];\n TreeNode *right = previous_level[children_index++];\n TreeNode *ancestor = new TreeNode(j);\n left->parent = ancestor;\n right->parent = ancestor;\n ancestor->left = left;\n ancestor->right = right;\n current_level.push_back(ancestor);\n }\n\n previous_level = current_level;\n }\n\n TreeNode *root = current_level.front();\n return root;\n}\n\nvoid delete_tree(TreeNode *root)\n{\n if (!root)\n {\n return;\n }\n\n delete_tree(root->left);\n delete_tree(root->right);\n delete root;\n}\n\nint previous_power_2(int v)\n{\n unsigned p = 1;\n\n while (p <= v)\n {\n p <<= 1;\n }\n\n return p >> 1;\n}\n\n// Build a reduction tree of size N. We build a full tree of size M,\n// the nearest lower power of 2 and we add remaining nodes in a\n// similar process as build_degenerate tree: rank 0 is involved in all\n// operations.\n\nTreeNode *build_reduction_tree(int N)\n{\n int nearest_pow2 = previous_power_2(N);\n TreeNode *root = build_full_tree(nearest_pow2);\n\n if (nearest_pow2 == N)\n {\n return root;\n }\n\n // Number of ranks is not a power of two, create additional nodes for the remaining ranks.\n TreeNode *previous = root;\n int root_value = root->value;\n\n for (int i = nearest_pow2; i < N; ++i)\n {\n TreeNode *left = previous;\n TreeNode *right = new TreeNode(i);\n TreeNode *ancestor = new TreeNode(root_value);\n left->parent = ancestor;\n right->parent = ancestor;\n ancestor->left = left;\n ancestor->right = right;\n previous = ancestor;\n }\n\n return previous;\n}\n\n// Given a reduction tree, we have to compute for a given rank what\n// are the actions to perform, i.e. we need to find the list of nodes\n// in the tree whose value is the current rank.\n// For the up-sweep order, we need the list of these nodes ordered\n// from bottom to top (root). This order is given by performing a\n// breadth-first traversal of the tree and pushing encountered nodes\n// to a stack.\nvoid compute_up_sweep_order(TreeNode *root, int v, std::stack &order)\n{\n std::queue to_visit;\n to_visit.push(root);\n\n while (!to_visit.empty())\n {\n TreeNode *current = to_visit.front();\n to_visit.pop();\n\n if (current->value == v)\n {\n order.push(current);\n }\n\n if (current->left)\n {\n to_visit.push(current->left);\n }\n\n if (current->right)\n {\n to_visit.push(current->right);\n }\n }\n}\n\ntemplate \nvoid\nvstack(Vector &dst, const Vector &top, const Vector &bottom)\n{\n typedef typename Vector::value_type value_type;\n dst.resize((top.get_num_rows() + bottom.get_num_rows()) * top.get_num_cols());\n dst.set_num_rows(top.get_num_rows() + bottom.get_num_rows());\n dst.set_num_cols(top.get_num_cols());\n dst.set_lda(top.get_num_rows() + bottom.get_num_rows());\n int dpitch = dst.get_lda() * sizeof(value_type);\n int spitch = top.get_lda() * sizeof(value_type);\n int width = top.get_num_cols() * sizeof(value_type);\n int height = top.get_num_rows();\n cudaMemcpy2D(dst.raw(), dpitch, top.raw(), spitch, width, height,\n cudaMemcpyDeviceToDevice);\n cudaMemcpy2D(dst.raw() + top.get_num_cols(), dpitch, bottom.raw(), spitch, width, height,\n cudaMemcpyDeviceToDevice);\n}\n\ntemplate \nvoid vslice(const Vector &src, Vector &dst, int start, int end)\n{\n typedef typename Vector::value_type value_type;\n int cols = src.get_num_cols();\n int dst_rows = end - start;\n dst.resize(dst_rows * cols);\n dst.set_num_rows(dst_rows);\n dst.set_num_cols(cols);\n dst.set_lda(dst_rows);\n int dpitch = dst.get_lda() * sizeof(value_type);\n int spitch = src.get_lda() * sizeof(value_type);\n int width = dst.get_num_cols() * sizeof(value_type);\n int height = dst.get_num_rows();\n cudaMemcpy2D(dst.raw(), dpitch, src.raw() + start, spitch, width, height,\n cudaMemcpyDeviceToDevice);\n}\n\n} // end anonymous namespace.\n\ntemplate \nHouseholderQR::HouseholderQR(TMatrix &A)\n : m_A(A), m_use_R_inverse(true)\n{\n}\n\ntemplate \nvoid HouseholderQR::QR(TVector &V)\n{\n int rows = V.get_num_rows();\n int cols = V.get_num_cols();\n int lda = V.get_lda();\n Lapack::geqrf(V, m_tau, m_work);\n Lapack::orgqr(V, m_tau, m_work);\n}\n\ntemplate \nvoid HouseholderQR::QR(TVector &V, TVector &R)\n{\n int rows = V.get_num_rows();\n int cols = V.get_num_cols();\n int lda = V.get_lda();\n R.resize(cols * cols);\n R.set_num_rows(cols);\n R.set_num_cols(cols);\n R.set_lda(cols);\n Lapack::geqrf(V, m_tau, m_work);\n vslice(V, R, 0, cols);\n clear_lower_triangular(R);\n Lapack::orgqr(V, m_tau, m_work);\n}\n\n\n// Send a vector to rank 'destination'.\n// If 'destination' is the current rank, we don't perform any call to\n// MPI and instead save the vector in a stack.\ntemplate \nvoid\nHouseholderQR::send_vector(TVector &V, int destination)\n{\n TMatrix &A = m_A;\n int current_rank = A.manager->global_id();\n\n if (destination == current_rank)\n {\n m_local_comms_stack.push(V);\n }\n else\n {\n A.manager->getComms()->send_vector(V, destination, 0);\n }\n}\n\n// Receive a vector from rank 'source'.\n// If 'source' is the current rank, we just pop the vector from our\n// local stack.\ntemplate \nvoid\nHouseholderQR::receive_vector(TVector &V, int source)\n{\n TMatrix &A = m_A;\n int current_rank = A.manager->global_id();\n\n if (source == current_rank)\n {\n V = m_local_comms_stack.top();\n m_local_comms_stack.pop();\n }\n else\n {\n A.manager->getComms()->recv_vector(V, source, 0);\n }\n}\n\n// Instead of performing the down-sweep phase of the TSQR algorithm to\n// get the Q matrix, we can compute Q = A * R ^ -1. This alternative\n// use less communications.\ntemplate \nvoid HouseholderQR::inverse_phase(TVector &V, TVector &R, int root)\n{\n TMatrix &A = m_A;\n int num_partitions = A.manager->getComms()->get_num_partitions();\n int current_partition = A.manager->global_id();\n\n if (current_partition == root)\n {\n for (int i = 0; i < num_partitions; ++i)\n {\n send_vector(R, i);\n }\n }\n\n receive_vector(R, root);\n Vector_h h_R = R;\n Lapack::trtri(h_R);\n R = h_R;\n // Q = V * R^-1\n Cublas::gemm(1, V, R, 0, V);\n}\n\ntemplate \nvoid HouseholderQR::QR_decomposition(TVector &V)\n{\n int rows = V.get_num_rows();\n int cols = V.get_num_cols();\n int lda = V.get_lda();\n TMatrix &A = m_A;\n\n if (cols == 1)\n {\n ValueTypeVec norm = get_norm(A, V, L2);\n scal(V, 1 / norm, 0, rows);\n return;\n }\n\n m_tau.resize(cols);\n m_work.resize(rows * cols);\n\n if (A.is_matrix_singleGPU())\n {\n QR(V);\n return;\n }\n\n // Multi-GPU path.\n int num_partitions = A.manager->getComms()->get_num_partitions();\n int current_partition = A.manager->global_id();\n // Stack of the Q matrices compute during reduction.\n std::stack Q_stack;\n TVector B1(cols * cols);\n B1.set_num_rows(cols);\n B1.set_num_cols(cols);\n B1.set_lda(cols);\n TVector B2(cols * cols);\n B2.set_num_rows(cols);\n B2.set_num_cols(cols);\n B2.set_lda(cols);\n TVector B;\n TVector R;\n TreeNode *reduce_tree = build_reduction_tree(num_partitions);\n // Debug: uncomment to print reduction tree.\n /*\n if (current_partition == 0)\n print_tree(reduce_tree, \"tree.dot\");\n */\n std::stack up_sweep_order;\n compute_up_sweep_order(reduce_tree, current_partition, up_sweep_order);\n std::stack down_sweep_order;\n std::stack order = up_sweep_order;\n\n while (!order.empty())\n {\n down_sweep_order.push(order.top());\n order.pop();\n }\n\n // Up-sweep phase:\n // - At each leaf node: compute a QR on its part of the vector then send the R matrix to its parent.\n // - At each non-leaf node: receives a R matrix from both\n // children, compute a QR and send the result R to the parent.\n // Q matrices are saved in a stack.\n while (!up_sweep_order.empty())\n {\n TreeNode *current = up_sweep_order.top();\n up_sweep_order.pop();\n TreeNode *left_child = current->left;\n TreeNode *right_child = current->right;\n TreeNode *parent = current->parent;\n\n if (current->is_leaf())\n {\n B = V;\n }\n else\n {\n // Receive R matrices from children and stack the two one above the other.\n receive_vector(B1, left_child->value);\n receive_vector(B2, right_child->value);\n vstack(B, B1, B2);\n }\n\n // costly for first iteration since we have to copy the vector to the stack.\n QR(B, R);\n Q_stack.push(B);\n\n if (!current->is_root())\n {\n send_vector(R, parent->value);\n }\n }\n\n if (m_use_R_inverse)\n {\n int reduce_root = reduce_tree->value;\n delete_tree(reduce_tree);\n inverse_phase(V, R, reduce_root);\n return;\n }\n\n // Down-sweep phase:\n // - At each node: receive a matrix from your parent (root of the\n // tree uses the identity matrix) and multiply this matrix with\n // the saved Q matrix for this node (computed during the\n // up-sweep phase).\n // Split this vector vertically and send the first half to your\n // left child and the second half to your right child.\n // Leaves nodes do not send anything, after multiplication the\n // result is the local part of the global Q matrix.\n while (!down_sweep_order.empty())\n {\n TreeNode *current = down_sweep_order.top();\n down_sweep_order.pop();\n TreeNode *left_child = current->left;\n TreeNode *right_child = current->right;\n TreeNode *parent = current->parent;\n TVector Q = Q_stack.top();\n Q_stack.pop();\n\n if (current->is_root())\n {\n B = Q;\n }\n else\n {\n // If leaf, write directly to output vector.\n TVector &output = current->is_leaf() ? V : B;\n receive_vector(B1, parent->value);\n Cublas::gemm(1, Q, B1, 0, output);\n }\n\n if (current->is_leaf())\n {\n continue;\n }\n\n int rows = B.get_num_rows();\n vslice(B, B1, 0, rows / 2);\n vslice(B, B2, rows / 2, rows);\n send_vector(B1, left_child->value);\n send_vector(B2, right_child->value);\n }\n\n delete_tree(reduce_tree);\n}\n\n#define AMGX_CASE_LINE(CASE) template class HouseholderQR::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n\n", "output": "// SPDX-FileCopyrightText: 2013 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\nnamespace\n{\n\ntemplate \n__global__\nvoid kernel_clear_lower_triangular(T *V, int dim, int lda)\n{\n int r = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (r < dim)\n {\n for (int c = 0; c < r; ++c)\n {\n V[c * lda + r] = 0.;\n }\n }\n}\n\ntemplate \nvoid clear_lower_triangular(Vector &V)\n{\n int num_threads = 128;\n int max_grid_size = 4096;\n int num_rows = V.get_num_rows();\n int num_blocks = std::min(max_grid_size, (num_rows + num_threads - 1) / num_rows);\n kernel_clear_lower_triangular <<< num_blocks, num_threads>>>(V.raw(), num_rows, V.get_lda());\n cudaCheckError();\n}\n\n\nstruct TreeNode\n{\n TreeNode(int v)\n : left(0), right(0), parent(0), value(v)\n {\n }\n\n bool is_leaf() const\n {\n return left == right;\n }\n\n bool is_root() const\n {\n return !parent;\n }\n\n TreeNode *left;\n TreeNode *right;\n TreeNode *parent;\n int value;\n};\n\n// Recursive function to print a reduction tree to an ofstream in dot\n// format.\nvoid print_tree_rec(TreeNode *n, std::ofstream &f)\n{\n TreeNode *left = n->left;\n TreeNode *right = n->right;\n TreeNode *parent = n->parent;\n // We use the address of the nodes as labels in dot format to\n // avoid having to renumber the nodes.\n size_t address = reinterpret_cast(n);\n size_t parent_address = reinterpret_cast(parent);\n size_t left_address = reinterpret_cast(left);\n size_t right_address = reinterpret_cast(right);\n f << address << \"[label=\" << n->value << \"];\" << std::endl;\n\n if (parent)\n {\n f << address << \" -> \" << parent_address << std::endl;\n }\n\n if (left)\n {\n f << address << \" -> \" << left_address << std::endl;\n print_tree_rec(left, f);\n }\n\n if (right)\n {\n f << address << \" -> \" << right_address << std::endl;\n print_tree_rec(right, f);\n }\n}\n\n#ifdef DEBUG\n// Debug function: print a reduction tree to a file in dot format.\nvoid print_tree(TreeNode *root, const char *filename)\n{\n std::ofstream f;\n f.open(filename);\n f << \"digraph A {\" << std::endl;\n print_tree_rec(root, f);\n f << \"}\" << std::endl;\n}\n\n// Create a degenerate tree of size N, this is not an efficient\n// reduction tree for a parallel operation (rank 0 is involved in all\n// operations) but it is used for testing.\n// Example with N == 4:\n// 0\n// / \\\n// 0 \\\n// / \\ \\\n// 0 \\ \\\n// / \\ \\ \\\n// 0 1 2 3\n// The label on each node is the active process for this operation.\nTreeNode *build_degenerate_tree(int N)\n{\n std::vector leaves;\n\n // Allocate one leaf node for each rank.\n for (int i = 0; i < N; ++i)\n {\n leaves.push_back(new TreeNode(i));\n }\n\n TreeNode *previous_node = leaves[0];\n\n // For each leaf except the first one, create the new root of the\n // tree labeled with 0.\n // The left child of this node is the previous root and the right\n // child is the current leaf.\n for (int i = 0; i < N - 1; ++i)\n {\n TreeNode *left = previous_node;\n TreeNode *right = leaves[i + 1];\n TreeNode *ancestor = new TreeNode(previous_node->value);\n left->parent = ancestor;\n right->parent = ancestor;\n ancestor->left = left;\n ancestor->right = right;\n previous_node = ancestor;\n }\n\n TreeNode *root = previous_node;\n return root;\n}\n\n#endif\n// Create a full tree of size N, N must be a power of 2.\n// For instance if N == 4:\n// 0\n// / \\\n// / \\\n// 0 2\n// / \\ / \\\n// 0 1 2 3\n// The label on each node is the active process for this operation.\nTreeNode *build_full_tree(int N)\n{\n // Nodes of the previous (deepest) level of the tree.\n std::vector previous_level;\n std::vector current_level;\n\n // Allocate N leaves nodes, one for each process.\n for (int i = 0; i < N; ++i)\n {\n previous_level.push_back(new TreeNode(i));\n }\n\n // Level i of the tree will have 2^i nodes, we construct the tree\n // from leaves to root. At each level we allocate new nodes and\n // connect these nodes to the previous level.\n for (int stride = 2; stride <= N; stride <<= 1)\n {\n current_level.clear();\n int children_index = 0;\n\n for (int j = 0; j < N; j += stride)\n {\n TreeNode *left = previous_level[children_index++];\n TreeNode *right = previous_level[children_index++];\n TreeNode *ancestor = new TreeNode(j);\n left->parent = ancestor;\n right->parent = ancestor;\n ancestor->left = left;\n ancestor->right = right;\n current_level.push_back(ancestor);\n }\n\n previous_level = current_level;\n }\n\n TreeNode *root = current_level.front();\n return root;\n}\n\nvoid delete_tree(TreeNode *root)\n{\n if (!root)\n {\n return;\n }\n\n delete_tree(root->left);\n delete_tree(root->right);\n delete root;\n}\n\nint previous_power_2(int v)\n{\n unsigned p = 1;\n\n while (p <= v)\n {\n p <<= 1;\n }\n\n return p >> 1;\n}\n\n// Build a reduction tree of size N. We build a full tree of size M,\n// the nearest lower power of 2 and we add remaining nodes in a\n// similar process as build_degenerate tree: rank 0 is involved in all\n// operations.\n\nTreeNode *build_reduction_tree(int N)\n{\n int nearest_pow2 = previous_power_2(N);\n TreeNode *root = build_full_tree(nearest_pow2);\n\n if (nearest_pow2 == N)\n {\n return root;\n }\n\n // Number of ranks is not a power of two, create additional nodes for the remaining ranks.\n TreeNode *previous = root;\n int root_value = root->value;\n\n for (int i = nearest_pow2; i < N; ++i)\n {\n TreeNode *left = previous;\n TreeNode *right = new TreeNode(i);\n TreeNode *ancestor = new TreeNode(root_value);\n left->parent = ancestor;\n right->parent = ancestor;\n ancestor->left = left;\n ancestor->right = right;\n previous = ancestor;\n }\n\n return previous;\n}\n\n// Given a reduction tree, we have to compute for a given rank what\n// are the actions to perform, i.e. we need to find the list of nodes\n// in the tree whose value is the current rank.\n// For the up-sweep order, we need the list of these nodes ordered\n// from bottom to top (root). This order is given by performing a\n// breadth-first traversal of the tree and pushing encountered nodes\n// to a stack.\nvoid compute_up_sweep_order(TreeNode *root, int v, std::stack &order)\n{\n std::queue to_visit;\n to_visit.push(root);\n\n while (!to_visit.empty())\n {\n TreeNode *current = to_visit.front();\n to_visit.pop();\n\n if (current->value == v)\n {\n order.push(current);\n }\n\n if (current->left)\n {\n to_visit.push(current->left);\n }\n\n if (current->right)\n {\n to_visit.push(current->right);\n }\n }\n}\n\ntemplate \nvoid\nvstack(Vector &dst, const Vector &top, const Vector &bottom)\n{\n typedef typename Vector::value_type value_type;\n dst.resize((top.get_num_rows() + bottom.get_num_rows()) * top.get_num_cols());\n dst.set_num_rows(top.get_num_rows() + bottom.get_num_rows());\n dst.set_num_cols(top.get_num_cols());\n dst.set_lda(top.get_num_rows() + bottom.get_num_rows());\n int dpitch = dst.get_lda() * sizeof(value_type);\n int spitch = top.get_lda() * sizeof(value_type);\n int width = top.get_num_cols() * sizeof(value_type);\n int height = top.get_num_rows();\n hipMemcpy2D(dst.raw(), dpitch, top.raw(), spitch, width, height,\n hipMemcpyDeviceToDevice);\n hipMemcpy2D(dst.raw() + top.get_num_cols(), dpitch, bottom.raw(), spitch, width, height,\n hipMemcpyDeviceToDevice);\n}\n\ntemplate \nvoid vslice(const Vector &src, Vector &dst, int start, int end)\n{\n typedef typename Vector::value_type value_type;\n int cols = src.get_num_cols();\n int dst_rows = end - start;\n dst.resize(dst_rows * cols);\n dst.set_num_rows(dst_rows);\n dst.set_num_cols(cols);\n dst.set_lda(dst_rows);\n int dpitch = dst.get_lda() * sizeof(value_type);\n int spitch = src.get_lda() * sizeof(value_type);\n int width = dst.get_num_cols() * sizeof(value_type);\n int height = dst.get_num_rows();\n hipMemcpy2D(dst.raw(), dpitch, src.raw() + start, spitch, width, height,\n hipMemcpyDeviceToDevice);\n}\n\n} // end anonymous namespace.\n\ntemplate \nHouseholderQR::HouseholderQR(TMatrix &A)\n : m_A(A), m_use_R_inverse(true)\n{\n}\n\ntemplate \nvoid HouseholderQR::QR(TVector &V)\n{\n int rows = V.get_num_rows();\n int cols = V.get_num_cols();\n int lda = V.get_lda();\n Lapack::geqrf(V, m_tau, m_work);\n Lapack::orgqr(V, m_tau, m_work);\n}\n\ntemplate \nvoid HouseholderQR::QR(TVector &V, TVector &R)\n{\n int rows = V.get_num_rows();\n int cols = V.get_num_cols();\n int lda = V.get_lda();\n R.resize(cols * cols);\n R.set_num_rows(cols);\n R.set_num_cols(cols);\n R.set_lda(cols);\n Lapack::geqrf(V, m_tau, m_work);\n vslice(V, R, 0, cols);\n clear_lower_triangular(R);\n Lapack::orgqr(V, m_tau, m_work);\n}\n\n\n// Send a vector to rank 'destination'.\n// If 'destination' is the current rank, we don't perform any call to\n// MPI and instead save the vector in a stack.\ntemplate \nvoid\nHouseholderQR::send_vector(TVector &V, int destination)\n{\n TMatrix &A = m_A;\n int current_rank = A.manager->global_id();\n\n if (destination == current_rank)\n {\n m_local_comms_stack.push(V);\n }\n else\n {\n A.manager->getComms()->send_vector(V, destination, 0);\n }\n}\n\n// Receive a vector from rank 'source'.\n// If 'source' is the current rank, we just pop the vector from our\n// local stack.\ntemplate \nvoid\nHouseholderQR::receive_vector(TVector &V, int source)\n{\n TMatrix &A = m_A;\n int current_rank = A.manager->global_id();\n\n if (source == current_rank)\n {\n V = m_local_comms_stack.top();\n m_local_comms_stack.pop();\n }\n else\n {\n A.manager->getComms()->recv_vector(V, source, 0);\n }\n}\n\n// Instead of performing the down-sweep phase of the TSQR algorithm to\n// get the Q matrix, we can compute Q = A * R ^ -1. This alternative\n// use less communications.\ntemplate \nvoid HouseholderQR::inverse_phase(TVector &V, TVector &R, int root)\n{\n TMatrix &A = m_A;\n int num_partitions = A.manager->getComms()->get_num_partitions();\n int current_partition = A.manager->global_id();\n\n if (current_partition == root)\n {\n for (int i = 0; i < num_partitions; ++i)\n {\n send_vector(R, i);\n }\n }\n\n receive_vector(R, root);\n Vector_h h_R = R;\n Lapack::trtri(h_R);\n R = h_R;\n // Q = V * R^-1\n Cublas::gemm(1, V, R, 0, V);\n}\n\ntemplate \nvoid HouseholderQR::QR_decomposition(TVector &V)\n{\n int rows = V.get_num_rows();\n int cols = V.get_num_cols();\n int lda = V.get_lda();\n TMatrix &A = m_A;\n\n if (cols == 1)\n {\n ValueTypeVec norm = get_norm(A, V, L2);\n scal(V, 1 / norm, 0, rows);\n return;\n }\n\n m_tau.resize(cols);\n m_work.resize(rows * cols);\n\n if (A.is_matrix_singleGPU())\n {\n QR(V);\n return;\n }\n\n // Multi-GPU path.\n int num_partitions = A.manager->getComms()->get_num_partitions();\n int current_partition = A.manager->global_id();\n // Stack of the Q matrices compute during reduction.\n std::stack Q_stack;\n TVector B1(cols * cols);\n B1.set_num_rows(cols);\n B1.set_num_cols(cols);\n B1.set_lda(cols);\n TVector B2(cols * cols);\n B2.set_num_rows(cols);\n B2.set_num_cols(cols);\n B2.set_lda(cols);\n TVector B;\n TVector R;\n TreeNode *reduce_tree = build_reduction_tree(num_partitions);\n // Debug: uncomment to print reduction tree.\n /*\n if (current_partition == 0)\n print_tree(reduce_tree, \"tree.dot\");\n */\n std::stack up_sweep_order;\n compute_up_sweep_order(reduce_tree, current_partition, up_sweep_order);\n std::stack down_sweep_order;\n std::stack order = up_sweep_order;\n\n while (!order.empty())\n {\n down_sweep_order.push(order.top());\n order.pop();\n }\n\n // Up-sweep phase:\n // - At each leaf node: compute a QR on its part of the vector then send the R matrix to its parent.\n // - At each non-leaf node: receives a R matrix from both\n // children, compute a QR and send the result R to the parent.\n // Q matrices are saved in a stack.\n while (!up_sweep_order.empty())\n {\n TreeNode *current = up_sweep_order.top();\n up_sweep_order.pop();\n TreeNode *left_child = current->left;\n TreeNode *right_child = current->right;\n TreeNode *parent = current->parent;\n\n if (current->is_leaf())\n {\n B = V;\n }\n else\n {\n // Receive R matrices from children and stack the two one above the other.\n receive_vector(B1, left_child->value);\n receive_vector(B2, right_child->value);\n vstack(B, B1, B2);\n }\n\n // costly for first iteration since we have to copy the vector to the stack.\n QR(B, R);\n Q_stack.push(B);\n\n if (!current->is_root())\n {\n send_vector(R, parent->value);\n }\n }\n\n if (m_use_R_inverse)\n {\n int reduce_root = reduce_tree->value;\n delete_tree(reduce_tree);\n inverse_phase(V, R, reduce_root);\n return;\n }\n\n // Down-sweep phase:\n // - At each node: receive a matrix from your parent (root of the\n // tree uses the identity matrix) and multiply this matrix with\n // the saved Q matrix for this node (computed during the\n // up-sweep phase).\n // Split this vector vertically and send the first half to your\n // left child and the second half to your right child.\n // Leaves nodes do not send anything, after multiplication the\n // result is the local part of the global Q matrix.\n while (!down_sweep_order.empty())\n {\n TreeNode *current = down_sweep_order.top();\n down_sweep_order.pop();\n TreeNode *left_child = current->left;\n TreeNode *right_child = current->right;\n TreeNode *parent = current->parent;\n TVector Q = Q_stack.top();\n Q_stack.pop();\n\n if (current->is_root())\n {\n B = Q;\n }\n else\n {\n // If leaf, write directly to output vector.\n TVector &output = current->is_leaf() ? V : B;\n receive_vector(B1, parent->value);\n Cublas::gemm(1, Q, B1, 0, output);\n }\n\n if (current->is_leaf())\n {\n continue;\n }\n\n int rows = B.get_num_rows();\n vslice(B, B1, 0, rows / 2);\n vslice(B, B2, rows / 2, rows);\n send_vector(B1, left_child->value);\n send_vector(B2, right_child->value);\n }\n\n delete_tree(reduce_tree);\n}\n\n#define AMGX_CASE_LINE(CASE) template class HouseholderQR::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n\n", "source": "/kaggle/working/amgx/src/eigensolvers/qr.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\nnamespace amgx\n{\n\n// Methods to check if Ax=b is easily invertible\ntemplate \nbool Cycle_Base::isASolvable(const Matrix &A)\n{\n // currently we have to disable this because in the distributed solver we have to launch coarse solver to exchange halos\n // otherwise in some cases nodes with > 1 coarse rows will be stuck waiting on halos from the node with 1 coarse row\n return false;\n /*return (A.get_num_rows()==1) && (A.get_block_dimy()==A.get_block_dimx());*/\n}\n\ntemplate\n__global__\nvoid direct_1x1_solve(const ValueTypeA *values, const ValueTypeB *rhs, ValueTypeB *x)\n\n{\n // it is always 1 thread because 1x1, but whatever\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (tid == 0) { x[tid] = rhs[tid] / values[tid]; }\n}\n\ntemplate \nvoid Cycle >::solveExactly_1x1(const Matrix &A, Vector &x, Vector &b)\n{\n x[0] = b[0] / A.values[0];\n}\n\n\ntemplate \nvoid Cycle >::solveExactly_1x1(const Matrix &A, Vector &x, Vector &b)\n{\n // only 1 thread needed though\n direct_1x1_solve <<< 1, 32>>>(\n A.values.raw(),\n b.raw(),\n x.raw());\n cudaCheckError();\n}\n\ntemplate \nvoid Cycle >::solveExactly_4x4(const Matrix &A, Vector &x, Vector &b)\n{\n GaussianElimination::gaussianElimination(A, x, b);\n}\n\n\ntemplate \nvoid Cycle >::solveExactly_4x4(const Matrix &A, Vector &x, Vector &b)\n{\n GaussianElimination::gaussianElimination(A, x, b);\n}\n\ntemplate \nvoid Cycle_Base::solveExactly(const Matrix &A, VVector &x, VVector &b)\n{\n if (A.get_block_size() == 1)\n {\n solveExactly_1x1(A, x, b);\n }\n else if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n solveExactly_4x4(A, x, b);\n }\n else\n {\n FatalError(\"Unsupported dimension for aggregation amg level\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\ntemplate\nstd::map*> &\nCycleFactory::getFactories( )\n{\n static std::map*> s_factories;\n return s_factories;\n}\n\ntemplate\nvoid CycleFactory::registerFactory(std::string name, CycleFactory *f)\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::const_iterator it = factories.find(name);\n\n if (it != factories.end())\n {\n std::string error = \"CycleFactory '\" + name + \"' has already been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n factories[name] = f;\n}\n\ntemplate\nvoid CycleFactory::unregisterFactory(std::string name)\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.find(name);\n\n if (it == factories.end())\n {\n std::string error = \"CycleFactory '\" + name + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n CycleFactory *factory = it->second;\n assert( factory != NULL );\n delete factory;\n factories.erase(it);\n}\n\ntemplate\nvoid CycleFactory::unregisterFactories( )\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.begin( );\n\n for ( ; it != factories.end( ) ; )\n {\n CycleFactory *factory = it->second;\n assert( factory != NULL );\n it++;\n delete factory;\n }\n\n factories.clear( );\n}\n\ntemplate\nCycle *CycleFactory::allocate(AMG_Class *amg, AMG_Level *level, VVector &b, VVector &c)\n{\n std::map*> &factories = getFactories( );\n std::string cycle_name = amg->m_cfg->AMG_Config::template getParameter(\"cycle\", amg->m_cfg_scope);\n typename std::map *>::const_iterator it = factories.find(cycle_name);\n\n if (it == factories.end())\n {\n std::string error = \"CycleFactory '\" + cycle_name + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n return it->second->create(amg, level, b, c);\n};\n\ntemplate\nvoid CycleFactory::generate(AMG_Class *amg, AMG_Level *level, VVector &b, VVector &c)\n{\n Cycle *cycle = allocate( amg, level, b, c );\n delete cycle;\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n\n#define AMGX_CASE_LINE(CASE) template class Cycle_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Cycle::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class CycleFactory::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\nnamespace amgx\n{\n\n// Methods to check if Ax=b is easily invertible\ntemplate \nbool Cycle_Base::isASolvable(const Matrix &A)\n{\n // currently we have to disable this because in the distributed solver we have to launch coarse solver to exchange halos\n // otherwise in some cases nodes with > 1 coarse rows will be stuck waiting on halos from the node with 1 coarse row\n return false;\n /*return (A.get_num_rows()==1) && (A.get_block_dimy()==A.get_block_dimx());*/\n}\n\ntemplate\n__global__\nvoid direct_1x1_solve(const ValueTypeA *values, const ValueTypeB *rhs, ValueTypeB *x)\n\n{\n // it is always 1 thread because 1x1, but whatever\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n\n if (tid == 0) { x[tid] = rhs[tid] / values[tid]; }\n}\n\ntemplate \nvoid Cycle >::solveExactly_1x1(const Matrix &A, Vector &x, Vector &b)\n{\n x[0] = b[0] / A.values[0];\n}\n\n\ntemplate \nvoid Cycle >::solveExactly_1x1(const Matrix &A, Vector &x, Vector &b)\n{\n // only 1 thread needed though\n direct_1x1_solve <<< 1, 32>>>(\n A.values.raw(),\n b.raw(),\n x.raw());\n cudaCheckError();\n}\n\ntemplate \nvoid Cycle >::solveExactly_4x4(const Matrix &A, Vector &x, Vector &b)\n{\n GaussianElimination::gaussianElimination(A, x, b);\n}\n\n\ntemplate \nvoid Cycle >::solveExactly_4x4(const Matrix &A, Vector &x, Vector &b)\n{\n GaussianElimination::gaussianElimination(A, x, b);\n}\n\ntemplate \nvoid Cycle_Base::solveExactly(const Matrix &A, VVector &x, VVector &b)\n{\n if (A.get_block_size() == 1)\n {\n solveExactly_1x1(A, x, b);\n }\n else if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n solveExactly_4x4(A, x, b);\n }\n else\n {\n FatalError(\"Unsupported dimension for aggregation amg level\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\ntemplate\nstd::map*> &\nCycleFactory::getFactories( )\n{\n static std::map*> s_factories;\n return s_factories;\n}\n\ntemplate\nvoid CycleFactory::registerFactory(std::string name, CycleFactory *f)\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::const_iterator it = factories.find(name);\n\n if (it != factories.end())\n {\n std::string error = \"CycleFactory '\" + name + \"' has already been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n factories[name] = f;\n}\n\ntemplate\nvoid CycleFactory::unregisterFactory(std::string name)\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.find(name);\n\n if (it == factories.end())\n {\n std::string error = \"CycleFactory '\" + name + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n CycleFactory *factory = it->second;\n assert( factory != NULL );\n delete factory;\n factories.erase(it);\n}\n\ntemplate\nvoid CycleFactory::unregisterFactories( )\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.begin( );\n\n for ( ; it != factories.end( ) ; )\n {\n CycleFactory *factory = it->second;\n assert( factory != NULL );\n it++;\n delete factory;\n }\n\n factories.clear( );\n}\n\ntemplate\nCycle *CycleFactory::allocate(AMG_Class *amg, AMG_Level *level, VVector &b, VVector &c)\n{\n std::map*> &factories = getFactories( );\n std::string cycle_name = amg->m_cfg->AMG_Config::template getParameter(\"cycle\", amg->m_cfg_scope);\n typename std::map *>::const_iterator it = factories.find(cycle_name);\n\n if (it == factories.end())\n {\n std::string error = \"CycleFactory '\" + cycle_name + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n return it->second->create(amg, level, b, c);\n};\n\ntemplate\nvoid CycleFactory::generate(AMG_Class *amg, AMG_Level *level, VVector &b, VVector &c)\n{\n Cycle *cycle = allocate( amg, level, b, c );\n delete cycle;\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n\n#define AMGX_CASE_LINE(CASE) template class Cycle_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Cycle::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class CycleFactory::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/cycles/cycle.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\nnamespace locally_downwind_kernels\n{\n// ---------------------------\n// Kernels\n// ---------------------------\n\n// Kernel to color the rows of the matrix, using min-max approach\ntemplate \n__global__\nvoid colorRowsKernel(IndexType *row_colors, const int num_colors, const int num_rows)\n{\n for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < num_rows; i += gridDim.x * blockDim.x)\n {\n row_colors[i] = i % num_colors;\n }\n}\n\n\n//is this edge and outgoing edge?\ntemplate \n__host__ __device__\nbool outgoing_edge( const IndexType *ia, const IndexType *ja, const ValueType *aa, IndexType i, IndexType ii, int blocksize )\n{\n IndexType j = ja[ii];\n ValueType weight = 0.0;\n\n for (IndexType iii = ii * blocksize; iii < (ii + 1)*blocksize; iii++)\n {\n weight += (aa[iii] * aa[iii]);\n }\n\n for (IndexType jj = ia[j]; jj < ia[j + 1]; jj++)\n {\n if ( ja[jj] == i )\n {\n ValueType counter_weight = 0.0;\n\n for (IndexType jjj = jj * blocksize; jjj < (jj + 1)*blocksize; jjj++)\n {\n counter_weight += (aa[jjj] * aa[jjj]);\n }\n\n return (weight >= counter_weight);\n }\n }\n\n return true;\n}\n\ntemplate \nstruct Buffer\n{\n IndexType *buffer;\n int buffersize;\n int loptr;\n int hiptr;\n int offset;\n\n __device__ int size();\n __device__ IndexType pop();\n __device__ void push( IndexType node );\n};\n\ntemplate \n__device__\nIndexType Buffer::pop()\n{\n int pos = loptr % buffersize;\n loptr++;\n return buffer[offset + pos];\n}\n\ntemplate \n__device__\nvoid Buffer::push(IndexType element )\n{\n buffer[offset + hiptr] = element;\n hiptr = (hiptr + 1) % buffersize;\n}\n\ntemplate \n__device__\nint Buffer::size()\n{\n return (hiptr - loptr) % buffersize;\n}\n\n\ntemplate \n__global__\nvoid traverse( const IndexType *ia, const IndexType *ja, const ValueType *aa, const IndexType *ria, const IndexType *rja, const IndexType *agg, IndexType *color, IndexType numAggregates, int buffersize, int blocksize )\n{\n extern __shared__ IndexType tr_smem[];\n Buffer ring;\n ring.buffer = &tr_smem[0];\n ring.buffersize = buffersize;\n ring.loptr = 0;\n ring.hiptr = 0;\n ring.offset = threadIdx.x * buffersize;\n int aggregate = threadIdx.x + blockDim.x * blockIdx.x;\n\n if ( aggregate < numAggregates )\n {\n bool nodesLeft = true;\n\n for (int allowed_incoming_edges = 0; nodesLeft; allowed_incoming_edges++)\n {\n nodesLeft = false;\n\n for (IndexType ii = ria[aggregate]; ii < ria[aggregate + 1]; ii++)\n {\n //find possible root node for modified BFS coloring\n IndexType node = rja[ii];\n int found_edges = 0;\n\n //not colored yet\n if ( color[node] == -1 )\n {\n for (IndexType jj = ia[node]; jj < ia[node + 1]; jj++)\n {\n if ( agg[ja[jj]] == aggregate && !outgoing_edge( ia, ja, aa, node, jj, blocksize ) )\n {\n found_edges++;\n }\n }\n\n nodesLeft = true;\n }\n else\n {\n found_edges = allowed_incoming_edges + 1;\n }\n\n //start modified BFS\n if ( found_edges <= allowed_incoming_edges )\n {\n color[node] = 0; //TODO: make the node look around for already set neighbors first\n ring.push(node);\n\n while ( ring.size() > 0 )\n {\n //this node is already colored.\n node = ring.pop();\n int myInitialColor = color[node];\n //traverse all neighbors to determin minimum own color\n int myColor = myInitialColor;\n\n for (IndexType jj = ia[node]; jj < ia[node + 1]; jj++)\n {\n IndexType j = ja[jj];\n\n if ( color[j] == myColor && j != node)\n {\n //try next color\n myColor++;\n jj = ia[node];\n }\n }\n\n //repair own color\n //color[node] = myColor;\n\n //traverse all children\n for (IndexType jj = ia[node]; jj < ia[node + 1]; jj++)\n {\n IndexType child = ja[jj];\n\n if ( agg[child] != aggregate || //only set colors in own aggregate\n !outgoing_edge( ia, ja, aa, node, jj, blocksize )) //only process outgoing edges\n {\n continue;\n }\n\n //unset -> set and push\n if ( color[child] == -1)\n {\n color[child] = myColor + 1;\n ring.push( child );\n }\n\n if ( color[child] >= myInitialColor )\n {\n color[child] = myColor + 1; //max( myColor+1, color[child] );\n ring.push( child );\n }\n\n /*\n //same color: reset and reset children, no push\n if( color[child] == myColor )\n {\n color[child] = myColor+1;\n for(IndexType kk = ia[child]; kk < ia[child+1]; kk++)\n {\n //same restrictions for setting colors apply here:\n //stay in aggregate, only set outgoing edges, only set uncolored or same level\n IndexType grandchild = ja[kk];\n if( agg[grandchild] != aggregate || !outgoing_edge( ia, ja, aa, child, kk, blocksize ) )\n continue;\n\n if( color[grandchild] == -1 )\n {\n color[grandchild] = myColor+2;\n ring.push( grandchild );\n }\n if( color[grandchild] == myColor+1 )\n color[grandchild] = myColor+2;\n }\n }\n */\n }\n }\n }\n }\n }\n }\n}\n\ntemplate \n__global__\nvoid repair( const IndexType *ia, const IndexType *ja, const ValueType *aa, const IndexType *ria, const IndexType *rja, const IndexType *agg, IndexType *color, IndexType numAggregates, int buffersize, int blocksize )\n{\n extern __shared__ IndexType tr_smem[];\n Buffer ring;\n ring.buffer = &tr_smem[0];\n ring.buffersize = buffersize;\n ring.loptr = 0;\n ring.hiptr = 0;\n ring.offset = threadIdx.x * buffersize;\n int aggregate = threadIdx.x + blockDim.x * blockIdx.x;\n\n if ( aggregate < numAggregates )\n {\n bool nodesLeft = true;\n\n for (int allowed_incoming_edges = 0; nodesLeft; allowed_incoming_edges++)\n {\n nodesLeft = false;\n\n for (IndexType ii = ria[aggregate]; ii < ria[aggregate + 1]; ii++)\n {\n //find possible root node for modified BFS coloring\n IndexType node = rja[ii];\n int found_edges = 0;\n\n //not colored yet\n if ( color[node] == -1 )\n {\n for (IndexType jj = ia[node]; jj < ia[node + 1]; jj++)\n {\n if ( agg[ja[jj]] == aggregate && !outgoing_edge( ia, ja, aa, node, jj, blocksize ) )\n {\n found_edges++;\n }\n }\n\n nodesLeft = true;\n }\n else\n {\n found_edges = allowed_incoming_edges + 1;\n }\n\n //start modified BFS\n if ( found_edges <= allowed_incoming_edges )\n {\n color[node] = 0; //TODO: make the node look around for already set neighbors first\n ring.push(node);\n\n while ( ring.size() > 0 )\n {\n //this node is already colored.\n node = ring.pop();\n int myInitialColor = color[node];\n //traverse all neighbors to determin minimum own color\n int myColor = myInitialColor;\n\n for (IndexType jj = ia[node]; jj < ia[node + 1]; jj++)\n {\n IndexType j = ja[jj];\n\n if ( color[j] == myColor && j != node)\n {\n //try next color\n myColor++;\n jj = ia[node];\n }\n }\n\n //traverse all children\n for (IndexType jj = ia[node]; jj < ia[node + 1]; jj++)\n {\n IndexType child = ja[jj];\n\n if ( agg[child] != aggregate || //only set colors in own aggregate\n !outgoing_edge( ia, ja, aa, node, jj, blocksize )) //only process outgoing edges\n {\n continue;\n }\n\n //unset -> set and push\n if ( color[child] == -1)\n {\n color[child] = myColor + 1;\n ring.push( child );\n }\n\n if ( color[child] >= myInitialColor )\n {\n color[child] = myColor + 1; //max( myColor+1, color[child] );\n ring.push( child );\n }\n\n /*\n //same color: reset and reset children, no push\n if( color[child] == myColor )\n {\n color[child] = myColor+1;\n for(IndexType kk = ia[child]; kk < ia[child+1]; kk++)\n {\n //same restrictions for setting colors apply here:\n //stay in aggregate, only set outgoing edges, only set uncolored or same level\n IndexType grandchild = ja[kk];\n if( agg[grandchild] != aggregate || !outgoing_edge( ia, ja, aa, child, kk, blocksize ) )\n continue;\n\n if( color[grandchild] == -1 )\n {\n color[grandchild] = myColor+2;\n ring.push( grandchild );\n }\n if( color[grandchild] == myColor+1 )\n color[grandchild] = myColor+2;\n }\n }\n */\n }\n }\n }\n }\n }\n }\n}\n\n\n\n}//locally_downwind_kernels namepsace\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate\nLocallyDownwindColoringBase::LocallyDownwindColoringBase(AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n this->m_num_colors = cfg.AMG_Config::template getParameter(\"num_colors\", cfg_scope);\n}\n\ntemplate\nvoid LocallyDownwindColoringBase::colorMatrix(Matrix &A)\n{\n //wait for colorMatrixUsingAggregates to be called\n std::cout << \"coloring denied\" << std::endl;\n return;\n}\n\ntemplate\nvoid LocallyDownwindColoringBase::colorMatrixUsingAggregates(Matrix &A, IVector &R_row_offsets, IVector &R_col_indices, IVector &aggregates )\n{\n#define CPU_VERSION\n#ifdef CPU_VERSION\n IndexType numRows = A.get_num_rows();\n IndexType nnz = A.get_num_nz();\n int blockdim = A.get_block_dimx() * A.get_block_dimy();\n //allocate memory on host\n IndexType *ia = new IndexType[numRows + 1];\n IndexType *ja = new IndexType[A.col_indices.size()];\n ValueType *aa = new ValueType[nnz * A.get_block_dimx()*A.get_block_dimy()];\n IndexType *ria = new IndexType[R_row_offsets.size()];\n IndexType *rja = new IndexType[R_col_indices.size()];\n IndexType *agg = new IndexType[aggregates.size()];\n IndexType *color = new IndexType[numRows];\n //copy data from device to host\n cudaMemcpy(ia, A.row_offsets.raw(), (numRows + 1)*sizeof(IndexType), cudaMemcpyDeviceToHost );\n cudaCheckError();\n cudaMemcpy(ja, A.col_indices.raw(), A.col_indices.size()*sizeof(IndexType), cudaMemcpyDeviceToHost );\n cudaCheckError();\n cudaMemcpy(aa, A.values.raw(), nnz * blockdim * sizeof(ValueType), cudaMemcpyDeviceToHost );\n cudaCheckError();\n cudaMemcpy(ria, R_row_offsets.raw(), R_row_offsets.size()*sizeof(IndexType), cudaMemcpyDeviceToHost );\n cudaCheckError();\n cudaMemcpy(rja, R_col_indices.raw(), R_col_indices.size()*sizeof(IndexType), cudaMemcpyDeviceToHost );\n cudaCheckError();\n cudaMemcpy(agg, aggregates.raw(), aggregates.size()*sizeof(IndexType), cudaMemcpyDeviceToHost );\n cudaCheckError();\n\n for (IndexType i = 0; i < numRows; i++)\n {\n color[i] = -1;\n }\n\n //color aggregate by aggregate\n for (IndexType aggregate = 0; aggregate < R_row_offsets.size() - 1; aggregate++)\n {\n std::queue q;\n bool nodesLeft = true;\n\n while ( nodesLeft )\n {\n //find uncolored node to start with (i.e. with minimum in degree)\n IndexType min_in_degree = numRows;\n IndexType next_node = -1;\n nodesLeft = false;\n\n for (IndexType node_index = ria[aggregate]; node_index < ria[aggregate + 1]; node_index++)\n {\n IndexType node = rja[node_index];\n int in_degree = 0;\n\n if ( color[node] == -1 )\n {\n nodesLeft = true;\n\n for (IndexType ii = ia[node]; ii < ia[node + 1]; ii++)\n {\n if ( agg[ja[ii]] == aggregate && !locally_downwind_kernels::outgoing_edge( ia, ja, aa, node, ii, blockdim ) )\n {\n in_degree++;\n }\n }\n\n if ( in_degree < min_in_degree )\n {\n min_in_degree = in_degree;\n next_node = node;\n }\n }\n }\n\n if (!nodesLeft)\n {\n break;\n }\n\n //start modified BFS\n color[next_node] = 0;\n q.push( next_node );\n\n while ( q.size() > 0 )\n {\n IndexType node = q.front();\n q.pop();\n int myInitialColor = color[node];\n int myColor = myInitialColor;\n\n //find valid color for this node\n for (IndexType ii = ia[node]; ii < ia[node + 1]; ii++)\n {\n if ( color[ja[ii]] == myColor && ja[ii] != node && (agg[ja[ii]] != aggregate || !locally_downwind_kernels::outgoing_edge( ia, ja, aa, node, ii, blockdim )) )\n {\n myColor++;\n ii = ia[node] - 1;\n }\n }\n\n //set color\n color[node] = myColor;\n\n //update children\n for (IndexType jj = ia[node]; jj < ia[node + 1]; jj++)\n {\n IndexType child = ja[jj];\n\n if ( agg[child] != aggregate || //only set colors in own aggregate\n !locally_downwind_kernels::outgoing_edge( ia, ja, aa, node, jj, blockdim ) ||//only process outgoing edges\n child == node) // don't mess with yourself\n {\n continue;\n }\n\n //unset -> set and push\n if ( color[child] >= myInitialColor )\n {\n color[child] = max( myColor + 1, color[child] );\n q.push( child );\n }\n\n if ( color[child] == -1)\n {\n color[child] = myColor + 1;\n q.push( child );\n }\n }\n }\n }\n }\n\n //copy back results\n this->m_row_colors.resize(A.row_offsets.size() - 1, 0);\n cudaMemcpy( this->m_row_colors.raw(), color, numRows * sizeof(IndexType), cudaMemcpyHostToDevice );\n cudaCheckError();\n //free all the others\n delete [] ia;\n delete [] ja;\n delete [] aa;\n delete [] agg;\n delete [] ria;\n delete [] rja;\n delete [] color;\n#else\n std::cout << \"coloring with aggregate information\" << std::endl;\n ViewType oldView = A.currentView();\n this->m_row_colors.resize(A.row_offsets.size() - 1, 0);\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n if (this->m_coloring_level == 0)\n {\n FatalError(\"Calling coloring scheme but coloring level==0\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else if (this->m_coloring_level == 1)\n {\n IndexType numRows = A.get_num_rows();\n IndexType numAggregates = R_row_offsets.size() - 1;\n int blocksize = A.get_block_dimx() * A.get_block_dimy();\n int max_aggregate_size = 100;\n const int threads_per_block = 64;\n const int num_blocks = (numAggregates - 1) / threads_per_block + 1;\n const int smem_size = max_aggregate_size * threads_per_block * sizeof(IndexType);\n this->m_row_colors.resize( numRows );\n thrust_wrapper::fill( this->m_row_colors.begin(), this->m_row_colors.end(), -1 );\n cudaCheckError();\n std::cout << \"start coloring kernel\" << std::endl;\n locally_downwind_kernels::traverse <<< num_blocks, threads_per_block, smem_size>>>( A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n R_row_offsets.raw(),\n R_col_indices.raw(),\n aggregates.raw(),\n this->m_row_colors.raw(),\n numAggregates,\n max_aggregate_size,\n blocksize);\n cudaCheckError();\n cudaDeviceSynchronize();\n cudaCheckError();\n std::cout << \"uncolored nodes: \" << amgx::thrust::count( this->m_row_colors.begin(), this->m_row_colors.end(), -1 ) << std::endl;\n }\n else\n {\n FatalError(\"Locally Downwind coloring algorithm can only do one ring coloring\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n A.setView(oldView);\n#endif\n}\n\n// Block version\ntemplate \nvoid LocallyDownwindColoring >::colorMatrixOneRing(Matrix_d &A)\n{\n FatalError(\"This method is no longer needed\", AMGX_ERR_NOT_IMPLEMENTED);\n /*\n // One thread per row\n const int num_rows = A.get_num_rows();\n IndexType *row_colors_ptr = this->m_row_colors.raw();\n\n const int threads_per_block = 64;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows-1)/threads_per_block + 1);\n\n locally_downwind_kernels::colorRowsKernel <<>>(row_colors_ptr, this->m_num_colors, num_rows);\n cudaCheckError();\n */\n}\n\ntemplate \nvoid LocallyDownwindColoring >::colorMatrixOneRing(Matrix_h &A)\n{\n FatalError(\"Haven't implemented locally downwind coloring for host\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n#define AMGX_CASE_LINE(CASE) template class LocallyDownwindColoringBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class LocallyDownwindColoring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // end namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\nnamespace locally_downwind_kernels\n{\n// ---------------------------\n// Kernels\n// ---------------------------\n\n// Kernel to color the rows of the matrix, using min-max approach\ntemplate \n__global__\nvoid colorRowsKernel(IndexType *row_colors, const int num_colors, const int num_rows)\n{\n for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < num_rows; i += gridDim.x * blockDim.x)\n {\n row_colors[i] = i % num_colors;\n }\n}\n\n\n//is this edge and outgoing edge?\ntemplate \n__host__ __device__\nbool outgoing_edge( const IndexType *ia, const IndexType *ja, const ValueType *aa, IndexType i, IndexType ii, int blocksize )\n{\n IndexType j = ja[ii];\n ValueType weight = 0.0;\n\n for (IndexType iii = ii * blocksize; iii < (ii + 1)*blocksize; iii++)\n {\n weight += (aa[iii] * aa[iii]);\n }\n\n for (IndexType jj = ia[j]; jj < ia[j + 1]; jj++)\n {\n if ( ja[jj] == i )\n {\n ValueType counter_weight = 0.0;\n\n for (IndexType jjj = jj * blocksize; jjj < (jj + 1)*blocksize; jjj++)\n {\n counter_weight += (aa[jjj] * aa[jjj]);\n }\n\n return (weight >= counter_weight);\n }\n }\n\n return true;\n}\n\ntemplate \nstruct Buffer\n{\n IndexType *buffer;\n int buffersize;\n int loptr;\n int hiptr;\n int offset;\n\n __device__ int size();\n __device__ IndexType pop();\n __device__ void push( IndexType node );\n};\n\ntemplate \n__device__\nIndexType Buffer::pop()\n{\n int pos = loptr % buffersize;\n loptr++;\n return buffer[offset + pos];\n}\n\ntemplate \n__device__\nvoid Buffer::push(IndexType element )\n{\n buffer[offset + hiptr] = element;\n hiptr = (hiptr + 1) % buffersize;\n}\n\ntemplate \n__device__\nint Buffer::size()\n{\n return (hiptr - loptr) % buffersize;\n}\n\n\ntemplate \n__global__\nvoid traverse( const IndexType *ia, const IndexType *ja, const ValueType *aa, const IndexType *ria, const IndexType *rja, const IndexType *agg, IndexType *color, IndexType numAggregates, int buffersize, int blocksize )\n{\n extern __shared__ IndexType tr_smem[];\n Buffer ring;\n ring.buffer = &tr_smem[0];\n ring.buffersize = buffersize;\n ring.loptr = 0;\n ring.hiptr = 0;\n ring.offset = threadIdx.x * buffersize;\n int aggregate = threadIdx.x + blockDim.x * blockIdx.x;\n\n if ( aggregate < numAggregates )\n {\n bool nodesLeft = true;\n\n for (int allowed_incoming_edges = 0; nodesLeft; allowed_incoming_edges++)\n {\n nodesLeft = false;\n\n for (IndexType ii = ria[aggregate]; ii < ria[aggregate + 1]; ii++)\n {\n //find possible root node for modified BFS coloring\n IndexType node = rja[ii];\n int found_edges = 0;\n\n //not colored yet\n if ( color[node] == -1 )\n {\n for (IndexType jj = ia[node]; jj < ia[node + 1]; jj++)\n {\n if ( agg[ja[jj]] == aggregate && !outgoing_edge( ia, ja, aa, node, jj, blocksize ) )\n {\n found_edges++;\n }\n }\n\n nodesLeft = true;\n }\n else\n {\n found_edges = allowed_incoming_edges + 1;\n }\n\n //start modified BFS\n if ( found_edges <= allowed_incoming_edges )\n {\n color[node] = 0; //TODO: make the node look around for already set neighbors first\n ring.push(node);\n\n while ( ring.size() > 0 )\n {\n //this node is already colored.\n node = ring.pop();\n int myInitialColor = color[node];\n //traverse all neighbors to determin minimum own color\n int myColor = myInitialColor;\n\n for (IndexType jj = ia[node]; jj < ia[node + 1]; jj++)\n {\n IndexType j = ja[jj];\n\n if ( color[j] == myColor && j != node)\n {\n //try next color\n myColor++;\n jj = ia[node];\n }\n }\n\n //repair own color\n //color[node] = myColor;\n\n //traverse all children\n for (IndexType jj = ia[node]; jj < ia[node + 1]; jj++)\n {\n IndexType child = ja[jj];\n\n if ( agg[child] != aggregate || //only set colors in own aggregate\n !outgoing_edge( ia, ja, aa, node, jj, blocksize )) //only process outgoing edges\n {\n continue;\n }\n\n //unset -> set and push\n if ( color[child] == -1)\n {\n color[child] = myColor + 1;\n ring.push( child );\n }\n\n if ( color[child] >= myInitialColor )\n {\n color[child] = myColor + 1; //max( myColor+1, color[child] );\n ring.push( child );\n }\n\n /*\n //same color: reset and reset children, no push\n if( color[child] == myColor )\n {\n color[child] = myColor+1;\n for(IndexType kk = ia[child]; kk < ia[child+1]; kk++)\n {\n //same restrictions for setting colors apply here:\n //stay in aggregate, only set outgoing edges, only set uncolored or same level\n IndexType grandchild = ja[kk];\n if( agg[grandchild] != aggregate || !outgoing_edge( ia, ja, aa, child, kk, blocksize ) )\n continue;\n\n if( color[grandchild] == -1 )\n {\n color[grandchild] = myColor+2;\n ring.push( grandchild );\n }\n if( color[grandchild] == myColor+1 )\n color[grandchild] = myColor+2;\n }\n }\n */\n }\n }\n }\n }\n }\n }\n}\n\ntemplate \n__global__\nvoid repair( const IndexType *ia, const IndexType *ja, const ValueType *aa, const IndexType *ria, const IndexType *rja, const IndexType *agg, IndexType *color, IndexType numAggregates, int buffersize, int blocksize )\n{\n extern __shared__ IndexType tr_smem[];\n Buffer ring;\n ring.buffer = &tr_smem[0];\n ring.buffersize = buffersize;\n ring.loptr = 0;\n ring.hiptr = 0;\n ring.offset = threadIdx.x * buffersize;\n int aggregate = threadIdx.x + blockDim.x * blockIdx.x;\n\n if ( aggregate < numAggregates )\n {\n bool nodesLeft = true;\n\n for (int allowed_incoming_edges = 0; nodesLeft; allowed_incoming_edges++)\n {\n nodesLeft = false;\n\n for (IndexType ii = ria[aggregate]; ii < ria[aggregate + 1]; ii++)\n {\n //find possible root node for modified BFS coloring\n IndexType node = rja[ii];\n int found_edges = 0;\n\n //not colored yet\n if ( color[node] == -1 )\n {\n for (IndexType jj = ia[node]; jj < ia[node + 1]; jj++)\n {\n if ( agg[ja[jj]] == aggregate && !outgoing_edge( ia, ja, aa, node, jj, blocksize ) )\n {\n found_edges++;\n }\n }\n\n nodesLeft = true;\n }\n else\n {\n found_edges = allowed_incoming_edges + 1;\n }\n\n //start modified BFS\n if ( found_edges <= allowed_incoming_edges )\n {\n color[node] = 0; //TODO: make the node look around for already set neighbors first\n ring.push(node);\n\n while ( ring.size() > 0 )\n {\n //this node is already colored.\n node = ring.pop();\n int myInitialColor = color[node];\n //traverse all neighbors to determin minimum own color\n int myColor = myInitialColor;\n\n for (IndexType jj = ia[node]; jj < ia[node + 1]; jj++)\n {\n IndexType j = ja[jj];\n\n if ( color[j] == myColor && j != node)\n {\n //try next color\n myColor++;\n jj = ia[node];\n }\n }\n\n //traverse all children\n for (IndexType jj = ia[node]; jj < ia[node + 1]; jj++)\n {\n IndexType child = ja[jj];\n\n if ( agg[child] != aggregate || //only set colors in own aggregate\n !outgoing_edge( ia, ja, aa, node, jj, blocksize )) //only process outgoing edges\n {\n continue;\n }\n\n //unset -> set and push\n if ( color[child] == -1)\n {\n color[child] = myColor + 1;\n ring.push( child );\n }\n\n if ( color[child] >= myInitialColor )\n {\n color[child] = myColor + 1; //max( myColor+1, color[child] );\n ring.push( child );\n }\n\n /*\n //same color: reset and reset children, no push\n if( color[child] == myColor )\n {\n color[child] = myColor+1;\n for(IndexType kk = ia[child]; kk < ia[child+1]; kk++)\n {\n //same restrictions for setting colors apply here:\n //stay in aggregate, only set outgoing edges, only set uncolored or same level\n IndexType grandchild = ja[kk];\n if( agg[grandchild] != aggregate || !outgoing_edge( ia, ja, aa, child, kk, blocksize ) )\n continue;\n\n if( color[grandchild] == -1 )\n {\n color[grandchild] = myColor+2;\n ring.push( grandchild );\n }\n if( color[grandchild] == myColor+1 )\n color[grandchild] = myColor+2;\n }\n }\n */\n }\n }\n }\n }\n }\n }\n}\n\n\n\n}//locally_downwind_kernels namepsace\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate\nLocallyDownwindColoringBase::LocallyDownwindColoringBase(AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n this->m_num_colors = cfg.AMG_Config::template getParameter(\"num_colors\", cfg_scope);\n}\n\ntemplate\nvoid LocallyDownwindColoringBase::colorMatrix(Matrix &A)\n{\n //wait for colorMatrixUsingAggregates to be called\n std::cout << \"coloring denied\" << std::endl;\n return;\n}\n\ntemplate\nvoid LocallyDownwindColoringBase::colorMatrixUsingAggregates(Matrix &A, IVector &R_row_offsets, IVector &R_col_indices, IVector &aggregates )\n{\n#define CPU_VERSION\n#ifdef CPU_VERSION\n IndexType numRows = A.get_num_rows();\n IndexType nnz = A.get_num_nz();\n int blockdim = A.get_block_dimx() * A.get_block_dimy();\n //allocate memory on host\n IndexType *ia = new IndexType[numRows + 1];\n IndexType *ja = new IndexType[A.col_indices.size()];\n ValueType *aa = new ValueType[nnz * A.get_block_dimx()*A.get_block_dimy()];\n IndexType *ria = new IndexType[R_row_offsets.size()];\n IndexType *rja = new IndexType[R_col_indices.size()];\n IndexType *agg = new IndexType[aggregates.size()];\n IndexType *color = new IndexType[numRows];\n //copy data from device to host\n hipMemcpy(ia, A.row_offsets.raw(), (numRows + 1)*sizeof(IndexType), hipMemcpyDeviceToHost );\n cudaCheckError();\n hipMemcpy(ja, A.col_indices.raw(), A.col_indices.size()*sizeof(IndexType), hipMemcpyDeviceToHost );\n cudaCheckError();\n hipMemcpy(aa, A.values.raw(), nnz * blockdim * sizeof(ValueType), hipMemcpyDeviceToHost );\n cudaCheckError();\n hipMemcpy(ria, R_row_offsets.raw(), R_row_offsets.size()*sizeof(IndexType), hipMemcpyDeviceToHost );\n cudaCheckError();\n hipMemcpy(rja, R_col_indices.raw(), R_col_indices.size()*sizeof(IndexType), hipMemcpyDeviceToHost );\n cudaCheckError();\n hipMemcpy(agg, aggregates.raw(), aggregates.size()*sizeof(IndexType), hipMemcpyDeviceToHost );\n cudaCheckError();\n\n for (IndexType i = 0; i < numRows; i++)\n {\n color[i] = -1;\n }\n\n //color aggregate by aggregate\n for (IndexType aggregate = 0; aggregate < R_row_offsets.size() - 1; aggregate++)\n {\n std::queue q;\n bool nodesLeft = true;\n\n while ( nodesLeft )\n {\n //find uncolored node to start with (i.e. with minimum in degree)\n IndexType min_in_degree = numRows;\n IndexType next_node = -1;\n nodesLeft = false;\n\n for (IndexType node_index = ria[aggregate]; node_index < ria[aggregate + 1]; node_index++)\n {\n IndexType node = rja[node_index];\n int in_degree = 0;\n\n if ( color[node] == -1 )\n {\n nodesLeft = true;\n\n for (IndexType ii = ia[node]; ii < ia[node + 1]; ii++)\n {\n if ( agg[ja[ii]] == aggregate && !locally_downwind_kernels::outgoing_edge( ia, ja, aa, node, ii, blockdim ) )\n {\n in_degree++;\n }\n }\n\n if ( in_degree < min_in_degree )\n {\n min_in_degree = in_degree;\n next_node = node;\n }\n }\n }\n\n if (!nodesLeft)\n {\n break;\n }\n\n //start modified BFS\n color[next_node] = 0;\n q.push( next_node );\n\n while ( q.size() > 0 )\n {\n IndexType node = q.front();\n q.pop();\n int myInitialColor = color[node];\n int myColor = myInitialColor;\n\n //find valid color for this node\n for (IndexType ii = ia[node]; ii < ia[node + 1]; ii++)\n {\n if ( color[ja[ii]] == myColor && ja[ii] != node && (agg[ja[ii]] != aggregate || !locally_downwind_kernels::outgoing_edge( ia, ja, aa, node, ii, blockdim )) )\n {\n myColor++;\n ii = ia[node] - 1;\n }\n }\n\n //set color\n color[node] = myColor;\n\n //update children\n for (IndexType jj = ia[node]; jj < ia[node + 1]; jj++)\n {\n IndexType child = ja[jj];\n\n if ( agg[child] != aggregate || //only set colors in own aggregate\n !locally_downwind_kernels::outgoing_edge( ia, ja, aa, node, jj, blockdim ) ||//only process outgoing edges\n child == node) // don't mess with yourself\n {\n continue;\n }\n\n //unset -> set and push\n if ( color[child] >= myInitialColor )\n {\n color[child] = max( myColor + 1, color[child] );\n q.push( child );\n }\n\n if ( color[child] == -1)\n {\n color[child] = myColor + 1;\n q.push( child );\n }\n }\n }\n }\n }\n\n //copy back results\n this->m_row_colors.resize(A.row_offsets.size() - 1, 0);\n hipMemcpy( this->m_row_colors.raw(), color, numRows * sizeof(IndexType), hipMemcpyHostToDevice );\n cudaCheckError();\n //free all the others\n delete [] ia;\n delete [] ja;\n delete [] aa;\n delete [] agg;\n delete [] ria;\n delete [] rja;\n delete [] color;\n#else\n std::cout << \"coloring with aggregate information\" << std::endl;\n ViewType oldView = A.currentView();\n this->m_row_colors.resize(A.row_offsets.size() - 1, 0);\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n if (this->m_coloring_level == 0)\n {\n FatalError(\"Calling coloring scheme but coloring level==0\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else if (this->m_coloring_level == 1)\n {\n IndexType numRows = A.get_num_rows();\n IndexType numAggregates = R_row_offsets.size() - 1;\n int blocksize = A.get_block_dimx() * A.get_block_dimy();\n int max_aggregate_size = 100;\n const int threads_per_block = 64;\n const int num_blocks = (numAggregates - 1) / threads_per_block + 1;\n const int smem_size = max_aggregate_size * threads_per_block * sizeof(IndexType);\n this->m_row_colors.resize( numRows );\n thrust_wrapper::fill( this->m_row_colors.begin(), this->m_row_colors.end(), -1 );\n cudaCheckError();\n std::cout << \"start coloring kernel\" << std::endl;\n locally_downwind_kernels::traverse <<< num_blocks, threads_per_block, smem_size>>>( A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n R_row_offsets.raw(),\n R_col_indices.raw(),\n aggregates.raw(),\n this->m_row_colors.raw(),\n numAggregates,\n max_aggregate_size,\n blocksize);\n cudaCheckError();\n hipDeviceSynchronize();\n cudaCheckError();\n std::cout << \"uncolored nodes: \" << amgx::thrust::count( this->m_row_colors.begin(), this->m_row_colors.end(), -1 ) << std::endl;\n }\n else\n {\n FatalError(\"Locally Downwind coloring algorithm can only do one ring coloring\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n A.setView(oldView);\n#endif\n}\n\n// Block version\ntemplate \nvoid LocallyDownwindColoring >::colorMatrixOneRing(Matrix_d &A)\n{\n FatalError(\"This method is no longer needed\", AMGX_ERR_NOT_IMPLEMENTED);\n /*\n // One thread per row\n const int num_rows = A.get_num_rows();\n IndexType *row_colors_ptr = this->m_row_colors.raw();\n\n const int threads_per_block = 64;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows-1)/threads_per_block + 1);\n\n locally_downwind_kernels::colorRowsKernel <<>>(row_colors_ptr, this->m_num_colors, num_rows);\n cudaCheckError();\n */\n}\n\ntemplate \nvoid LocallyDownwindColoring >::colorMatrixOneRing(Matrix_h &A)\n{\n FatalError(\"Haven't implemented locally downwind coloring for host\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n#define AMGX_CASE_LINE(CASE) template class LocallyDownwindColoringBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class LocallyDownwindColoring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // end namespace amgx\n\n", "source": "/kaggle/working/amgx/src/matrix_coloring/locally_downwind.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n// Pseudo-random number generator\nnamespace amgx\n{\n\n__host__ __device__ unsigned int hash(unsigned int a, unsigned int seed)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n\nstruct is_less_than_zero\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x < 0;\n }\n};\n\n// ---------------------------\n// Kernels\n// ---------------------------\ntemplate \n__global__\nvoid colorRowsMultiHashKernel_1step(const IndexType *A_offsets, const IndexType *A_column_indices, IndexType *row_colors, const int num_rows, const int next_color)\n{\n int hash_j;\n int hash_i;\n int my_colors[num_hash];\n int my_row_color;\n\n for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < num_rows; i += blockDim.x * gridDim.x)\n {\n int num_possible_colors = 0;\n bool max_i, min_i;\n int row_start = A_offsets[i];\n int row_end = A_offsets[i + 1];\n\n //#pragma unroll //can't unroll because it is not the innermost loop\n for (int k = 0; k < num_hash; k++)\n {\n max_i = true;\n min_i = true;\n\n for (int r = row_start; r < row_end; r++)\n {\n int j = A_column_indices[r];\n\n if (j >= num_rows) { continue; }\n\n hash_j = hash(j, k);\n hash_i = hash(i, k);\n\n // There is an uncolored neighbour that is greater\n if ( hash_j > hash_i)\n {\n max_i = false;\n }\n\n // There is an uncolored neighbour that is smaller\n if (hash_j < hash_i)\n {\n min_i = false;\n }\n }\n\n // If not colored or colored but coin flip decides color should be changed\n if (max_i)\n {\n my_colors[num_possible_colors++] = 2 * k;\n }\n else if (min_i)\n {\n my_colors[num_possible_colors++] = 2 * k + 1;\n }\n }\n\n if (num_possible_colors)\n {\n int rand_pick = hash(i, 0) % num_possible_colors;\n my_row_color = my_colors[rand_pick];\n }\n else\n {\n my_row_color = 2 * num_hash;\n }\n\n row_colors[i] = my_row_color;\n }\n}\n\ntemplate \n__global__\nvoid colorRowsMultiHashKernel(const IndexType *A_offsets, const IndexType *A_column_indices, IndexType *row_colors, const int num_rows, const int next_color, const int seed)\n{\n unsigned int i_rand[num_hash];\n\n for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < num_rows; i += blockDim.x * gridDim.x)\n {\n // skip if previously matched\n if (row_colors[i] != -1) { continue; }\n\n int not_min = 0, not_max = 0;\n int t;\n\n#pragma unroll\n for (t = 0; t < num_hash; t++)\n {\n i_rand[t] = hash(i, seed + 1043 * t);\n }\n\n // have we been proved to be not min or max\n int row_start = A_offsets[i];\n int row_end = A_offsets[i + 1];\n int possible_colors = 2 * num_hash;\n\n for (int r = row_start; r < row_end; r++)\n {\n int j = A_column_indices[r];\n\n // skip diagonal\n if (j == i || j >= num_rows)\n {\n continue;\n }\n\n int j_color = row_colors[j];\n\n if (j_color != -1 && j_color < next_color)\n {\n continue;\n }\n\n#pragma unroll\n for (t = 0; t < num_hash; t++)\n {\n unsigned int j_rand = hash(j, seed + 1043 * t);\n\n // bail if any neighbor is greater\n if (i_rand[t] <= j_rand && !(not_max & (0x1 << t) ))\n {\n not_max |= (0x1 << t);\n possible_colors--;\n }\n\n if (i_rand[t] >= j_rand && !(not_min & (0x1 << t) ))\n {\n not_min |= (0x1 << t);\n possible_colors--;\n }\n }\n\n if (possible_colors == 0)\n {\n break;\n }\n }\n\n if (possible_colors == 0) { continue; }\n\n // made it here, so possible_colors > 0\n // pick one of the possible colors pseudo-randomly\n int col_id = i % possible_colors;\n int this_col_id = 0;\n\n for (t = 0; t < num_hash; t++)\n {\n if (!(not_min & (0x1 << t) ) && col_id == this_col_id)\n {\n row_colors[i] = 2 * t + next_color;\n return;\n }\n\n this_col_id += !(not_min & (0x1 << t));\n\n if (!(not_max & (0x1 << t) ) && col_id == this_col_id)\n {\n row_colors[i] = 2 * t + 1 + next_color;\n return;\n }\n\n this_col_id += !(not_max & (0x1 << t));\n }\n }\n}\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate\nMultiHashMatrixColoringBase::MultiHashMatrixColoringBase(AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n if (cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\"))\n {\n m_uncolored_fraction = 0.;\n }\n else\n {\n m_uncolored_fraction = cfg.AMG_Config::template getParameter(\"max_uncolored_percentage\", cfg_scope);\n }\n\n max_num_hash = cfg.AMG_Config::template getParameter(\"max_num_hash\", cfg_scope);\n}\n\ntemplate\nvoid MultiHashMatrixColoringBase::colorMatrix(Matrix &A)\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize(A.row_offsets.size() - 1, 0);\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n if (this->m_coloring_level == 0)\n {\n FatalError(\"Calling coloring scheme but coloring level==0\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n }\n else if (this->m_coloring_level == 1)\n {\n this->colorMatrixOneRing(A);\n }\n else\n {\n FatalError(\"Multi-hash coloring algorithm can only do one ring coloring\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n A.setView(oldView);\n}\n\n\n// Block version\ntemplate \nvoid MultiHashMatrixColoring >::colorMatrixOneRing(Matrix_d &A)\n{\n // One thread per row\n const int num_rows = A.get_num_rows();\n int max_uncolored_rows = (int) (this->m_uncolored_fraction * ((ValueType) num_rows));\n thrust_wrapper::fill(this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, -1);\n cudaCheckError();\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n IndexType *row_colors_ptr = this->m_row_colors.raw();\n const int threads_per_block = 256;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (num_rows - 1) / threads_per_block + 1);\n this->m_num_colors = 0;\n // Heuristic for setting the number of hash function to use\n int avg_nonzero = 1.5 * A.row_offsets[num_rows] / num_rows;\n this->num_hash = min(avg_nonzero, this->max_num_hash);\n int next_color = 0;\n int it = 0;\n int seed = 1012;\n\n if (avg_nonzero != 0)\n {\n for ( int num_uncolored = num_rows; num_uncolored > max_uncolored_rows ; )\n {\n it++;\n\n // Assign all nodes to 0 colors by default\n if (this->num_hash == 1)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 2)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 3)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 4)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 5)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 6)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 7)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 8)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 9)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 10)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 11)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 12)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 13)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 14)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 15)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 16)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 17)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 18)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 19)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 20)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 21)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 22)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 23)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 24)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 25)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash > 25)\n {\n FatalError(\"Multi-hash coloring algorithm currently can't handle more than 25 hash functions\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n cudaCheckError();\n seed = hash(seed, 0);\n next_color += 2 * this->num_hash;\n num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_less_than_zero() );\n cudaCheckError();\n }\n }\n else\n {\n thrust_wrapper::fill(this->m_row_colors.begin(), this->m_row_colors.end(), 0);\n cudaCheckError();\n }\n\n this->m_num_colors = *amgx::thrust::max_element(this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows) + 1;\n cudaCheckError();\n}\n\ntemplate \nvoid MultiHashMatrixColoring >::colorMatrixOneRing(Matrix_h &A)\n{\n FatalError(\"Haven't implemented MultiHash matrix coloring for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n#define AMGX_CASE_LINE(CASE) template class MultiHashMatrixColoringBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class MultiHashMatrixColoring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // end namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n// Pseudo-random number generator\nnamespace amgx\n{\n\n__host__ __device__ unsigned int hash(unsigned int a, unsigned int seed)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n\nstruct is_less_than_zero\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x < 0;\n }\n};\n\n// ---------------------------\n// Kernels\n// ---------------------------\ntemplate \n__global__\nvoid colorRowsMultiHashKernel_1step(const IndexType *A_offsets, const IndexType *A_column_indices, IndexType *row_colors, const int num_rows, const int next_color)\n{\n int hash_j;\n int hash_i;\n int my_colors[num_hash];\n int my_row_color;\n\n for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < num_rows; i += blockDim.x * gridDim.x)\n {\n int num_possible_colors = 0;\n bool max_i, min_i;\n int row_start = A_offsets[i];\n int row_end = A_offsets[i + 1];\n\n //#pragma unroll //can't unroll because it is not the innermost loop\n for (int k = 0; k < num_hash; k++)\n {\n max_i = true;\n min_i = true;\n\n for (int r = row_start; r < row_end; r++)\n {\n int j = A_column_indices[r];\n\n if (j >= num_rows) { continue; }\n\n hash_j = hash(j, k);\n hash_i = hash(i, k);\n\n // There is an uncolored neighbour that is greater\n if ( hash_j > hash_i)\n {\n max_i = false;\n }\n\n // There is an uncolored neighbour that is smaller\n if (hash_j < hash_i)\n {\n min_i = false;\n }\n }\n\n // If not colored or colored but coin flip decides color should be changed\n if (max_i)\n {\n my_colors[num_possible_colors++] = 2 * k;\n }\n else if (min_i)\n {\n my_colors[num_possible_colors++] = 2 * k + 1;\n }\n }\n\n if (num_possible_colors)\n {\n int rand_pick = hash(i, 0) % num_possible_colors;\n my_row_color = my_colors[rand_pick];\n }\n else\n {\n my_row_color = 2 * num_hash;\n }\n\n row_colors[i] = my_row_color;\n }\n}\n\ntemplate \n__global__\nvoid colorRowsMultiHashKernel(const IndexType *A_offsets, const IndexType *A_column_indices, IndexType *row_colors, const int num_rows, const int next_color, const int seed)\n{\n unsigned int i_rand[num_hash];\n\n for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < num_rows; i += blockDim.x * gridDim.x)\n {\n // skip if previously matched\n if (row_colors[i] != -1) { continue; }\n\n int not_min = 0, not_max = 0;\n int t;\n\n#pragma unroll\n for (t = 0; t < num_hash; t++)\n {\n i_rand[t] = hash(i, seed + 1043 * t);\n }\n\n // have we been proved to be not min or max\n int row_start = A_offsets[i];\n int row_end = A_offsets[i + 1];\n int possible_colors = 2 * num_hash;\n\n for (int r = row_start; r < row_end; r++)\n {\n int j = A_column_indices[r];\n\n // skip diagonal\n if (j == i || j >= num_rows)\n {\n continue;\n }\n\n int j_color = row_colors[j];\n\n if (j_color != -1 && j_color < next_color)\n {\n continue;\n }\n\n#pragma unroll\n for (t = 0; t < num_hash; t++)\n {\n unsigned int j_rand = hash(j, seed + 1043 * t);\n\n // bail if any neighbor is greater\n if (i_rand[t] <= j_rand && !(not_max & (0x1 << t) ))\n {\n not_max |= (0x1 << t);\n possible_colors--;\n }\n\n if (i_rand[t] >= j_rand && !(not_min & (0x1 << t) ))\n {\n not_min |= (0x1 << t);\n possible_colors--;\n }\n }\n\n if (possible_colors == 0)\n {\n break;\n }\n }\n\n if (possible_colors == 0) { continue; }\n\n // made it here, so possible_colors > 0\n // pick one of the possible colors pseudo-randomly\n int col_id = i % possible_colors;\n int this_col_id = 0;\n\n for (t = 0; t < num_hash; t++)\n {\n if (!(not_min & (0x1 << t) ) && col_id == this_col_id)\n {\n row_colors[i] = 2 * t + next_color;\n return;\n }\n\n this_col_id += !(not_min & (0x1 << t));\n\n if (!(not_max & (0x1 << t) ) && col_id == this_col_id)\n {\n row_colors[i] = 2 * t + 1 + next_color;\n return;\n }\n\n this_col_id += !(not_max & (0x1 << t));\n }\n }\n}\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate\nMultiHashMatrixColoringBase::MultiHashMatrixColoringBase(AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n if (cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\"))\n {\n m_uncolored_fraction = 0.;\n }\n else\n {\n m_uncolored_fraction = cfg.AMG_Config::template getParameter(\"max_uncolored_percentage\", cfg_scope);\n }\n\n max_num_hash = cfg.AMG_Config::template getParameter(\"max_num_hash\", cfg_scope);\n}\n\ntemplate\nvoid MultiHashMatrixColoringBase::colorMatrix(Matrix &A)\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize(A.row_offsets.size() - 1, 0);\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n if (this->m_coloring_level == 0)\n {\n FatalError(\"Calling coloring scheme but coloring level==0\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n }\n else if (this->m_coloring_level == 1)\n {\n this->colorMatrixOneRing(A);\n }\n else\n {\n FatalError(\"Multi-hash coloring algorithm can only do one ring coloring\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n A.setView(oldView);\n}\n\n\n// Block version\ntemplate \nvoid MultiHashMatrixColoring >::colorMatrixOneRing(Matrix_d &A)\n{\n // One thread per row\n const int num_rows = A.get_num_rows();\n int max_uncolored_rows = (int) (this->m_uncolored_fraction * ((ValueType) num_rows));\n thrust_wrapper::fill(this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, -1);\n cudaCheckError();\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n IndexType *row_colors_ptr = this->m_row_colors.raw();\n const int threads_per_block = 256;\n const int num_blocks = std::min(AMGX_GRID_MAX_SIZE, (int) (num_rows - 1) / threads_per_block + 1);\n this->m_num_colors = 0;\n // Heuristic for setting the number of hash function to use\n int avg_nonzero = 1.5 * A.row_offsets[num_rows] / num_rows;\n this->num_hash = min(avg_nonzero, this->max_num_hash);\n int next_color = 0;\n int it = 0;\n int seed = 1012;\n\n if (avg_nonzero != 0)\n {\n for ( int num_uncolored = num_rows; num_uncolored > max_uncolored_rows ; )\n {\n it++;\n\n // Assign all nodes to 0 colors by default\n if (this->num_hash == 1)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 2)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 3)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 4)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 5)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 6)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 7)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 8)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 9)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 10)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 11)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 12)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 13)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 14)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 15)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 16)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 17)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 18)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 19)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 20)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 21)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 22)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 23)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 24)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash == 25)\n {\n colorRowsMultiHashKernel <<< num_blocks, threads_per_block>>> (A_row_offsets_ptr, A_column_indices_ptr, row_colors_ptr, num_rows, next_color, seed);\n }\n else if (this->num_hash > 25)\n {\n FatalError(\"Multi-hash coloring algorithm currently can't handle more than 25 hash functions\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n cudaCheckError();\n seed = hash(seed, 0);\n next_color += 2 * this->num_hash;\n num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_less_than_zero() );\n cudaCheckError();\n }\n }\n else\n {\n thrust_wrapper::fill(this->m_row_colors.begin(), this->m_row_colors.end(), 0);\n cudaCheckError();\n }\n\n this->m_num_colors = *amgx::thrust::max_element(this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows) + 1;\n cudaCheckError();\n}\n\ntemplate \nvoid MultiHashMatrixColoring >::colorMatrixOneRing(Matrix_h &A)\n{\n FatalError(\"Haven't implemented MultiHash matrix coloring for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n#define AMGX_CASE_LINE(CASE) template class MultiHashMatrixColoringBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class MultiHashMatrixColoring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // end namespace amgx\n\n", "source": "/kaggle/working/amgx/src/matrix_coloring/multi_hash.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace amgx\n{\n\nvoid coloring_histogram(int *out_hist_host, int num_rows, int max_color, int *row_colors)\n{\n amgx::thrust::device_ptr data(row_colors);\n\n if (false) //TODO: enable, since almost always faster\n {\n device_vector_alloc histogram(num_rows);\n thrust_wrapper::sort(data, data + num_rows);\n cudaCheckError();\n amgx::thrust::counting_iterator search_begin(0);\n amgx::thrust::upper_bound(data, data + num_rows,\n search_begin, search_begin + num_rows,\n histogram.begin());\n cudaCheckError();\n amgx::thrust::adjacent_difference(histogram.begin(), histogram.end(), histogram.begin());\n cudaCheckError();\n amgx::thrust::host_vector h(max_color + 1);\n amgx::thrust::copy(histogram.begin(), histogram.begin() + max_color + 1, h.begin());\n cudaCheckError();\n\n for (int i = 0; i <= max_color; i++)\n {\n out_hist_host[i] = h[i]; //TODO avoid this copy\n }\n }\n else\n {\n for (int i = 0; i <= max_color; i++)\n {\n out_hist_host[i] = amgx::thrust::count(data, data + num_rows, i);\n }\n\n cudaCheckError();\n }\n}\n\n\n__global__ void permute_colors_kernel(int num_rows, int *row_colors, int *color_permutation)\n{\n int row_id = blockIdx.x * blockDim.x + threadIdx.x;\n\n for ( ; row_id < num_rows ; row_id += blockDim.x * gridDim.x )\n {\n int color = row_colors[row_id];\n color = __ldg(color_permutation + color);\n row_colors[row_id] = color;\n }\n}\n\n__global__ void reverse_colors_kernel(int num_rows, int max_color, int *row_colors)\n{\n int row_id = blockIdx.x * blockDim.x + threadIdx.x;\n\n for ( ; row_id < num_rows ; row_id += blockDim.x * gridDim.x )\n {\n int color = row_colors[row_id];\n\n if (color > 0)\n {\n //1 -> max_color\n //max_color -> 1\n color = max_color - color + 1;\n }\n\n row_colors[row_id] = color;\n }\n}\n\nint reverse_colors(int num_rows, int max_color, int *row_colors)\n{\n const int GRID_SIZE_1 = std::min( 2048, (num_rows + 256 - 1) / 256 );\n reverse_colors_kernel <<< GRID_SIZE_1, 256>>>(num_rows, max_color, row_colors);\n cudaCheckError();\n return max_color;\n}\n\nint eliminate_null_colors(int num_rows, int max_color, int *row_colors)\n{\n amgx::thrust::host_vector hist(max_color + 1, 0);\n amgx::thrust::host_vector perm(max_color + 1, 0);\n device_vector_alloc perm_d(max_color + 1, 0);\n coloring_histogram(amgx::thrust::raw_pointer_cast(hist.data()), num_rows, max_color, row_colors);\n int nonempty_color = 1;\n\n for (int i = 1; i <= max_color; i++) //0 is blank\n {\n if (hist[i] > 0) //keep it if not empty\n {\n perm[i] = nonempty_color++;\n }\n }\n\n perm_d = perm;\n const int GRID_SIZE_1 = std::min( 2048, (num_rows + 256 - 1) / 256 );\n permute_colors_kernel <<< GRID_SIZE_1, 256>>>(num_rows, row_colors, amgx::thrust::raw_pointer_cast(perm_d.data()));\n cudaCheckError();\n return nonempty_color - 1;\n}\nint reorder_colors_by_frequency(int num_rows, int max_color, int *row_colors)\n{\n amgx::thrust::host_vector hist(max_color + 1, 0);\n amgx::thrust::host_vector perm(max_color + 1, 0);\n device_vector_alloc perm_d(max_color + 1, 0);\n coloring_histogram(amgx::thrust::raw_pointer_cast(hist.data()), num_rows, max_color, row_colors);\n int nonempty_color = 1;\n\n for (int i = 1; i <= max_color; i++) //0 is blank\n {\n if (hist[i] > 0) //keep it if not empty\n {\n perm[i] = nonempty_color++;\n }\n }\n\n perm_d = perm;\n const int GRID_SIZE_1 = std::min( 2048, (num_rows + 256 - 1) / 256 );\n permute_colors_kernel <<< GRID_SIZE_1, 256>>>(num_rows, row_colors, amgx::thrust::raw_pointer_cast(perm_d.data()));\n cudaCheckError();\n return nonempty_color - 1;\n //return 0;\n}\n\n}\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace amgx\n{\n\nvoid coloring_histogram(int *out_hist_host, int num_rows, int max_color, int *row_colors)\n{\n amgx::thrust::device_ptr data(row_colors);\n\n if (false) //TODO: enable, since almost always faster\n {\n device_vector_alloc histogram(num_rows);\n thrust_wrapper::sort(data, data + num_rows);\n cudaCheckError();\n amgx::thrust::counting_iterator search_begin(0);\n amgx::thrust::upper_bound(data, data + num_rows,\n search_begin, search_begin + num_rows,\n histogram.begin());\n cudaCheckError();\n amgx::thrust::adjacent_difference(histogram.begin(), histogram.end(), histogram.begin());\n cudaCheckError();\n amgx::thrust::host_vector h(max_color + 1);\n amgx::thrust::copy(histogram.begin(), histogram.begin() + max_color + 1, h.begin());\n cudaCheckError();\n\n for (int i = 0; i <= max_color; i++)\n {\n out_hist_host[i] = h[i]; //TODO avoid this copy\n }\n }\n else\n {\n for (int i = 0; i <= max_color; i++)\n {\n out_hist_host[i] = amgx::thrust::count(data, data + num_rows, i);\n }\n\n cudaCheckError();\n }\n}\n\n\n__global__ void permute_colors_kernel(int num_rows, int *row_colors, int *color_permutation)\n{\n int row_id = blockIdx.x * blockDim.x + threadIdx.x;\n\n for ( ; row_id < num_rows ; row_id += blockDim.x * gridDim.x )\n {\n int color = row_colors[row_id];\n color = __ldg(color_permutation + color);\n row_colors[row_id] = color;\n }\n}\n\n__global__ void reverse_colors_kernel(int num_rows, int max_color, int *row_colors)\n{\n int row_id = blockIdx.x * blockDim.x + threadIdx.x;\n\n for ( ; row_id < num_rows ; row_id += blockDim.x * gridDim.x )\n {\n int color = row_colors[row_id];\n\n if (color > 0)\n {\n //1 -> max_color\n //max_color -> 1\n color = max_color - color + 1;\n }\n\n row_colors[row_id] = color;\n }\n}\n\nint reverse_colors(int num_rows, int max_color, int *row_colors)\n{\n const int GRID_SIZE_1 = std::min( 2048, (num_rows + 256 - 1) / 256 );\n reverse_colors_kernel <<< GRID_SIZE_1, 256>>>(num_rows, max_color, row_colors);\n cudaCheckError();\n return max_color;\n}\n\nint eliminate_null_colors(int num_rows, int max_color, int *row_colors)\n{\n amgx::thrust::host_vector hist(max_color + 1, 0);\n amgx::thrust::host_vector perm(max_color + 1, 0);\n device_vector_alloc perm_d(max_color + 1, 0);\n coloring_histogram(amgx::thrust::raw_pointer_cast(hist.data()), num_rows, max_color, row_colors);\n int nonempty_color = 1;\n\n for (int i = 1; i <= max_color; i++) //0 is blank\n {\n if (hist[i] > 0) //keep it if not empty\n {\n perm[i] = nonempty_color++;\n }\n }\n\n perm_d = perm;\n const int GRID_SIZE_1 = std::min( 2048, (num_rows + 256 - 1) / 256 );\n permute_colors_kernel <<< GRID_SIZE_1, 256>>>(num_rows, row_colors, amgx::thrust::raw_pointer_cast(perm_d.data()));\n cudaCheckError();\n return nonempty_color - 1;\n}\nint reorder_colors_by_frequency(int num_rows, int max_color, int *row_colors)\n{\n amgx::thrust::host_vector hist(max_color + 1, 0);\n amgx::thrust::host_vector perm(max_color + 1, 0);\n device_vector_alloc perm_d(max_color + 1, 0);\n coloring_histogram(amgx::thrust::raw_pointer_cast(hist.data()), num_rows, max_color, row_colors);\n int nonempty_color = 1;\n\n for (int i = 1; i <= max_color; i++) //0 is blank\n {\n if (hist[i] > 0) //keep it if not empty\n {\n perm[i] = nonempty_color++;\n }\n }\n\n perm_d = perm;\n const int GRID_SIZE_1 = std::min( 2048, (num_rows + 256 - 1) / 256 );\n permute_colors_kernel <<< GRID_SIZE_1, 256>>>(num_rows, row_colors, amgx::thrust::raw_pointer_cast(perm_d.data()));\n cudaCheckError();\n return nonempty_color - 1;\n //return 0;\n}\n\n}\n", "source": "/kaggle/working/amgx/src/matrix_coloring/coloring_utils.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\nnamespace amgx\n{\n/***************************************\n * Source Definitions\n ***************************************/\n\ntemplate\nMatrixColoring::MatrixColoring(AMG_Config &cfg, const std::string &cfg_scope) :\n m_num_colors(0), m_row_colors(0), m_sorted_rows_by_color(0), m_offsets_rows_per_color(0), m_ref_count(1), m_boundary_coloring(SYNC_COLORS), m_halo_coloring(FIRST)\n{\n m_coloring_level = cfg.getParameter(\"coloring_level\", cfg_scope);\n m_boundary_coloring = cfg.getParameter(\"boundary_coloring\", cfg_scope);\n m_halo_coloring = cfg.getParameter(\"halo_coloring\", cfg_scope);\n}\n\ntemplate\nMatrixColoring::~MatrixColoring()\n{\n}\n\n__global__ void findSeparation(INDEX_TYPE *rows_by_color, INDEX_TYPE *offsets_by_color, INDEX_TYPE *separation, INDEX_TYPE boundary, INDEX_TYPE num_colors, INDEX_TYPE num_rows)\n{\n int block_offset = blockIdx.x * (blockDim.x / 32) * 31; //each warp does 31 rows, 1 on the left side is redundant\n int lane = threadIdx.x % 32;\n int element = block_offset + (threadIdx.x / 32) * 31 + lane - 1;\n\n while (element < num_rows)\n {\n int color = 0;\n int row = -1;\n\n if (element != -1)\n {\n row = rows_by_color[element];\n\n while ((color < num_colors) && ((element < offsets_by_color[color]) || (element >= offsets_by_color[color + 1]))) { color ++; }\n\n if ((element == offsets_by_color[color]) && (row >= boundary)) { separation[color] = element; } //special case when first row of color is immediately a boundary node\n\n if ((element == offsets_by_color[color + 1] - 1) && (row < boundary)) { separation[color] = element + 1; } //special case when I am the last, but I am still not a boundary\n }\n\n unsigned int result = utils::ballot(row >= boundary, utils::activemask());\n\n //if (result>0) printf(\"%x\\n\", result);\n if (lane > 0 && row >= boundary && ((result >> (lane - 1)) & 1) == 0) { separation[color] = element; }\n\n element += gridDim.x * (blockDim.x / 32) * 31;\n }\n}\n\n//prints how many edges fail to obey coloring property\n//if the optional aggregates parameter is specified, it also measures the downwind coloring property:\n//for each incoming edge (j,i), where i and j share the same aggregate, holds: color(j) < color(i)\ntemplate \nvoid MatrixColoring::assertColoring( Matrix &A, IVector &aggregates )\n{\n IndexType numRows = A.get_num_rows();\n IndexType nnz = A.get_num_nz();\n IndexType blocksize = A.get_block_dimx() * A.get_block_dimy();\n IVector &coloring = this->m_row_colors;\n bool check_downwind = aggregates.size() == A.get_num_rows();\n //allocate host memory\n IndexType *ia = new IndexType[numRows + 1];\n IndexType *ja = new IndexType[nnz];\n ValueType *aa = new ValueType[nnz * blocksize];\n IndexType *color = new IndexType[numRows];\n //copy to host\n cudaMemcpy( ia, A.row_offsets.raw(), sizeof(IndexType) * (numRows + 1), cudaMemcpyDeviceToHost );\n cudaCheckError();\n cudaMemcpy( ja, A.col_indices.raw(), sizeof(IndexType)*nnz, cudaMemcpyDeviceToHost );\n cudaCheckError();\n cudaMemcpy( aa, A.values.raw(), sizeof(ValueType)*blocksize * nnz, cudaMemcpyDeviceToHost );\n cudaCheckError();\n cudaMemcpy( color, coloring.raw(), sizeof(IndexType)*numRows, cudaMemcpyDeviceToHost );\n cudaCheckError();\n IndexType *agg = new IndexType[numRows];\n\n if ( check_downwind )\n {\n cudaMemcpy( agg, aggregates.raw(), sizeof(IndexType)*numRows, cudaMemcpyDeviceToHost );\n cudaCheckError();\n }\n\n //count how many nodes have a color\n IndexType *color_used = new IndexType[numRows];\n\n for (IndexType i = 0; i < numRows; i++)\n {\n color_used[i] = 0;\n }\n\n for (IndexType i = 0; i < numRows; i++)\n {\n if ( color[i] >= 0 && color[i] < numRows )\n {\n color_used[color[i]]++;\n }\n else\n {\n std::cout << \"color out of range: color[\" << i << \"] = \" << color[i] << std::endl;\n }\n }\n\n // count violations of these two properties:\n // 1. locally downwind: for incoming edges (j,i) in same aggregate: color(j) < color(i)\n // 2. valid coloring: for neighbors j: color(j) != color(i)\n int violation_1 = 0;\n int property_1 = 0;\n int violation_2 = 0;\n int property_2 = 0;\n int inner_edges = 0;\n\n //note: property 1 cannot be enforeced all the time. Each cycle for example will violate it regardless of the coloring.\n for (IndexType i = 0; i < numRows; i++)\n {\n for (IndexType ii = ia[i]; ii < ia[i + 1]; ii++)\n {\n IndexType j = ja[ii];\n\n if ( j == i )\n {\n continue;\n }\n\n //check coloring property\n if ( color[j] == color[i] )\n {\n violation_2++;\n }\n\n property_2++;\n\n if ( check_downwind && agg[j] == agg[i] )\n {\n //look for transpose edge to decide outgoing or not\n bool outgoing = true;\n\n for (IndexType jj = ia[j]; jj < ia[j + 1]; jj++)\n {\n //found\n if ( ja[jj] == i )\n {\n ValueType weight = types::util::get_zero();\n\n for (IndexType iii = ii * blocksize; iii < (ii + 1)*blocksize; iii++)\n {\n weight = weight + aa[iii] * aa[iii];\n }\n\n ValueType counter_weight = types::util::get_zero();\n\n for (IndexType jjj = jj * blocksize; jjj < (jj + 1)*blocksize; jjj++)\n {\n counter_weight = counter_weight + aa[jjj] * aa[jjj];\n }\n\n outgoing = types::util::abs(weight) > types::util::abs(counter_weight);\n break;\n }\n }\n\n //outgoing -> check downwind property\n if ( outgoing )\n {\n if ( color[j] <= color[i] )\n {\n violation_1++;\n }\n\n property_1++;\n }\n\n inner_edges++;\n }\n }\n }\n\n //tell results\n if ( check_downwind )\n {\n std::cout << 200 * property_1 / double(inner_edges) << \"% of all edges inside an aggregate are directed\" << std::endl;\n\n if ( property_1 > 0 )\n {\n std::cout << 100 * violation_1 / double(property_1) << \"% of all outgoing edges inside an aggregate are not colored downwind\" << std::endl;\n }\n }\n\n std::cout << 100 * violation_2 / double(property_2) << \"% of all edges violated coloring property\" << std::endl;\n std::cout << \"number of nodes that use this color:\" << std::endl;\n\n for (IndexType i = 0; i < numRows; i++)\n if ( color_used[i] > 0 )\n {\n std::cout << i << \": \" << color_used[i] << std::endl;\n }\n\n //free!\n delete [] ia;\n delete [] ja;\n delete [] aa;\n delete [] agg;\n delete [] color;\n delete [] color_used;\n}\ntemplate\nvoid MatrixColoring::createColorArrays(Matrix &A)\n{\n ViewType old = A.currentView();\n A.setViewExterior();\n int num_rows = A.get_num_rows();\n\n //Disabled since currently we are not doing halo exchanges during colored execution\n /*typedef TemplateConfig hvector_type;\n typedef Vector HVector;\n if (!A.is_matrix_singleGPU()) {\n HVector num_colors(1);\n std::vector partition_num_colors(0);\n num_colors[0] = m_num_colors;\n A.manager->getComms()->global_reduce(partition_num_colors, num_colors, A, 6332);\n int max_partition_colors = 0;\n for (int i = 0; i < partition_num_colors.size(); i++)\n max_partition_colors = std::max(partition_num_colors[i][0],max_partition_colors);\n m_num_colors = max_partition_colors;\n }*/\n\n if (m_halo_coloring == LAST)\n {\n thrust_wrapper::fill(m_row_colors.begin() + num_rows, m_row_colors.end(), m_num_colors);\n cudaCheckError();\n }\n\n IVector offsets_rows_per_color;\n\n if (m_offsets_rows_per_color.size() == 0)\n {\n // Sort the vertices based o their color\n m_sorted_rows_by_color.resize(num_rows);\n // Copy row colors\n IVector row_colors(m_row_colors);\n thrust_wrapper::sequence(m_sorted_rows_by_color.begin(), m_sorted_rows_by_color.end());\n amgx::thrust::sort_by_key(row_colors.begin(), row_colors.begin() + num_rows, m_sorted_rows_by_color.begin());\n cudaCheckError();\n // Compute the offset for each color\n offsets_rows_per_color.resize(m_num_colors + 1);\n m_offsets_rows_per_color.resize(m_num_colors + 1);\n // Compute interior-exterior separation for every color\n m_offsets_rows_per_color_separation.resize(m_num_colors);\n //m_offsets_rows_per_color_separation_halo.resize(m_num_colors);\n amgx::thrust::lower_bound(row_colors.begin(),\n row_colors.begin() + num_rows,\n amgx::thrust::counting_iterator(0),\n amgx::thrust::counting_iterator(offsets_rows_per_color.size()),\n offsets_rows_per_color.begin());\n // Copy from device to host\n m_offsets_rows_per_color = offsets_rows_per_color;\n cudaCheckError();\n }\n else\n {\n m_offsets_rows_per_color_separation.resize(m_num_colors);\n }\n\n cudaCheckError();\n\n if (!A.is_matrix_singleGPU() && (A.getViewExterior() != A.getViewInterior()))\n {\n A.setViewInterior();\n int separation = A.get_num_rows();\n\n if (TConfig::memSpace == AMGX_host)\n {\n for (int i = 0; i < m_num_colors; i++)\n {\n m_offsets_rows_per_color_separation[i] = m_offsets_rows_per_color[i]\n + (amgx::thrust::lower_bound(m_sorted_rows_by_color.begin() + m_offsets_rows_per_color[i],\n m_sorted_rows_by_color.begin() + m_offsets_rows_per_color[i + 1],\n separation)\n - (m_sorted_rows_by_color.begin() + m_offsets_rows_per_color[i]));\n }\n\n cudaCheckError();\n }\n // this is not a proper search, rather we look at every single element. But it is still a lot faster than the above (~10*)\n else\n {\n IVector separation_offsets_rows_per_color(m_num_colors);\n int size = num_rows;\n int num_blocks = std::min(4096, (size + 123) / 124);\n findSeparation <<< num_blocks, 128>>>(m_sorted_rows_by_color.raw(), offsets_rows_per_color.raw(), separation_offsets_rows_per_color.raw(), separation, m_num_colors, num_rows);\n cudaCheckError();\n amgx::thrust::copy(separation_offsets_rows_per_color.begin(), separation_offsets_rows_per_color.end(), m_offsets_rows_per_color_separation.begin());\n cudaCheckError();\n\n for (int i = 0; i < m_num_colors; i++)\n {\n if (this->m_offsets_rows_per_color[i] == this->m_offsets_rows_per_color[i + 1])\n {\n this->m_offsets_rows_per_color_separation[i] = this->m_offsets_rows_per_color[i + 1];\n }\n }\n }\n }\n else\n {\n amgx::thrust::copy(m_offsets_rows_per_color.begin() + 1, m_offsets_rows_per_color.end(), m_offsets_rows_per_color_separation.begin());\n cudaCheckError();\n }\n\n A.setView(old);\n}\n\ntemplate\nstd::map*> &\nMatrixColoringFactory::getFactories( )\n{\n static std::map *> s_factories;\n return s_factories;\n}\n\ntemplate\nvoid MatrixColoringFactory::registerFactory(std::string name, MatrixColoringFactory *f)\n{\n std::map *> &factories = getFactories( );\n typename std::map *>::iterator it = factories.find(name);\n\n if (it != factories.end())\n {\n std::string error = \"MatrixColoringFactory '\" + name + \"' has already been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n factories[name] = f;\n}\n\ntemplate\nvoid MatrixColoringFactory::unregisterFactory(std::string name)\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.find(name);\n\n if (it == factories.end())\n {\n std::string error = \"MatrixColoringFactory '\" + name + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n MatrixColoringFactory *factory = it->second;\n assert( factory != NULL );\n delete factory;\n factories.erase(it);\n}\n\ntemplate\nvoid MatrixColoringFactory::unregisterFactories( )\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.begin( );\n\n for ( ; it != factories.end( ) ; )\n {\n MatrixColoringFactory *factory = it->second;\n assert( factory != NULL );\n it++;\n delete factory;\n }\n\n factories.clear( );\n}\n\ntemplate\nMatrixColoring *MatrixColoringFactory::allocate(AMG_Config &cfg, const std::string &cfg_scope)\n{\n std::map *> &factories = getFactories( );\n std::string matrix_coloring_scheme = cfg.getParameter(\"matrix_coloring_scheme\", cfg_scope);\n typename std::map *>::const_iterator it = factories.find(matrix_coloring_scheme);\n\n if (it == factories.end())\n {\n std::string error = \"MatrixColoringFactory '\" + matrix_coloring_scheme + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n return it->second->create(cfg, cfg_scope);\n};\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class MatrixColoring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class MatrixColoringFactory::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n}\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\nnamespace amgx\n{\n/***************************************\n * Source Definitions\n ***************************************/\n\ntemplate\nMatrixColoring::MatrixColoring(AMG_Config &cfg, const std::string &cfg_scope) :\n m_num_colors(0), m_row_colors(0), m_sorted_rows_by_color(0), m_offsets_rows_per_color(0), m_ref_count(1), m_boundary_coloring(SYNC_COLORS), m_halo_coloring(FIRST)\n{\n m_coloring_level = cfg.getParameter(\"coloring_level\", cfg_scope);\n m_boundary_coloring = cfg.getParameter(\"boundary_coloring\", cfg_scope);\n m_halo_coloring = cfg.getParameter(\"halo_coloring\", cfg_scope);\n}\n\ntemplate\nMatrixColoring::~MatrixColoring()\n{\n}\n\n__global__ void findSeparation(INDEX_TYPE *rows_by_color, INDEX_TYPE *offsets_by_color, INDEX_TYPE *separation, INDEX_TYPE boundary, INDEX_TYPE num_colors, INDEX_TYPE num_rows)\n{\n int block_offset = blockIdx.x * (blockDim.x / 32) * 31; //each warp does 31 rows, 1 on the left side is redundant\n int lane = threadIdx.x % 32;\n int element = block_offset + (threadIdx.x / 32) * 31 + lane - 1;\n\n while (element < num_rows)\n {\n int color = 0;\n int row = -1;\n\n if (element != -1)\n {\n row = rows_by_color[element];\n\n while ((color < num_colors) && ((element < offsets_by_color[color]) || (element >= offsets_by_color[color + 1]))) { color ++; }\n\n if ((element == offsets_by_color[color]) && (row >= boundary)) { separation[color] = element; } //special case when first row of color is immediately a boundary node\n\n if ((element == offsets_by_color[color + 1] - 1) && (row < boundary)) { separation[color] = element + 1; } //special case when I am the last, but I am still not a boundary\n }\n\n unsigned int result = utils::ballot(row >= boundary, utils::activemask());\n\n //if (result>0) printf(\"%x\\n\", result);\n if (lane > 0 && row >= boundary && ((result >> (lane - 1)) & 1) == 0) { separation[color] = element; }\n\n element += gridDim.x * (blockDim.x / 32) * 31;\n }\n}\n\n//prints how many edges fail to obey coloring property\n//if the optional aggregates parameter is specified, it also measures the downwind coloring property:\n//for each incoming edge (j,i), where i and j share the same aggregate, holds: color(j) < color(i)\ntemplate \nvoid MatrixColoring::assertColoring( Matrix &A, IVector &aggregates )\n{\n IndexType numRows = A.get_num_rows();\n IndexType nnz = A.get_num_nz();\n IndexType blocksize = A.get_block_dimx() * A.get_block_dimy();\n IVector &coloring = this->m_row_colors;\n bool check_downwind = aggregates.size() == A.get_num_rows();\n //allocate host memory\n IndexType *ia = new IndexType[numRows + 1];\n IndexType *ja = new IndexType[nnz];\n ValueType *aa = new ValueType[nnz * blocksize];\n IndexType *color = new IndexType[numRows];\n //copy to host\n hipMemcpy( ia, A.row_offsets.raw(), sizeof(IndexType) * (numRows + 1), hipMemcpyDeviceToHost );\n cudaCheckError();\n hipMemcpy( ja, A.col_indices.raw(), sizeof(IndexType)*nnz, hipMemcpyDeviceToHost );\n cudaCheckError();\n hipMemcpy( aa, A.values.raw(), sizeof(ValueType)*blocksize * nnz, hipMemcpyDeviceToHost );\n cudaCheckError();\n hipMemcpy( color, coloring.raw(), sizeof(IndexType)*numRows, hipMemcpyDeviceToHost );\n cudaCheckError();\n IndexType *agg = new IndexType[numRows];\n\n if ( check_downwind )\n {\n hipMemcpy( agg, aggregates.raw(), sizeof(IndexType)*numRows, hipMemcpyDeviceToHost );\n cudaCheckError();\n }\n\n //count how many nodes have a color\n IndexType *color_used = new IndexType[numRows];\n\n for (IndexType i = 0; i < numRows; i++)\n {\n color_used[i] = 0;\n }\n\n for (IndexType i = 0; i < numRows; i++)\n {\n if ( color[i] >= 0 && color[i] < numRows )\n {\n color_used[color[i]]++;\n }\n else\n {\n std::cout << \"color out of range: color[\" << i << \"] = \" << color[i] << std::endl;\n }\n }\n\n // count violations of these two properties:\n // 1. locally downwind: for incoming edges (j,i) in same aggregate: color(j) < color(i)\n // 2. valid coloring: for neighbors j: color(j) != color(i)\n int violation_1 = 0;\n int property_1 = 0;\n int violation_2 = 0;\n int property_2 = 0;\n int inner_edges = 0;\n\n //note: property 1 cannot be enforeced all the time. Each cycle for example will violate it regardless of the coloring.\n for (IndexType i = 0; i < numRows; i++)\n {\n for (IndexType ii = ia[i]; ii < ia[i + 1]; ii++)\n {\n IndexType j = ja[ii];\n\n if ( j == i )\n {\n continue;\n }\n\n //check coloring property\n if ( color[j] == color[i] )\n {\n violation_2++;\n }\n\n property_2++;\n\n if ( check_downwind && agg[j] == agg[i] )\n {\n //look for transpose edge to decide outgoing or not\n bool outgoing = true;\n\n for (IndexType jj = ia[j]; jj < ia[j + 1]; jj++)\n {\n //found\n if ( ja[jj] == i )\n {\n ValueType weight = types::util::get_zero();\n\n for (IndexType iii = ii * blocksize; iii < (ii + 1)*blocksize; iii++)\n {\n weight = weight + aa[iii] * aa[iii];\n }\n\n ValueType counter_weight = types::util::get_zero();\n\n for (IndexType jjj = jj * blocksize; jjj < (jj + 1)*blocksize; jjj++)\n {\n counter_weight = counter_weight + aa[jjj] * aa[jjj];\n }\n\n outgoing = types::util::abs(weight) > types::util::abs(counter_weight);\n break;\n }\n }\n\n //outgoing -> check downwind property\n if ( outgoing )\n {\n if ( color[j] <= color[i] )\n {\n violation_1++;\n }\n\n property_1++;\n }\n\n inner_edges++;\n }\n }\n }\n\n //tell results\n if ( check_downwind )\n {\n std::cout << 200 * property_1 / double(inner_edges) << \"% of all edges inside an aggregate are directed\" << std::endl;\n\n if ( property_1 > 0 )\n {\n std::cout << 100 * violation_1 / double(property_1) << \"% of all outgoing edges inside an aggregate are not colored downwind\" << std::endl;\n }\n }\n\n std::cout << 100 * violation_2 / double(property_2) << \"% of all edges violated coloring property\" << std::endl;\n std::cout << \"number of nodes that use this color:\" << std::endl;\n\n for (IndexType i = 0; i < numRows; i++)\n if ( color_used[i] > 0 )\n {\n std::cout << i << \": \" << color_used[i] << std::endl;\n }\n\n //free!\n delete [] ia;\n delete [] ja;\n delete [] aa;\n delete [] agg;\n delete [] color;\n delete [] color_used;\n}\ntemplate\nvoid MatrixColoring::createColorArrays(Matrix &A)\n{\n ViewType old = A.currentView();\n A.setViewExterior();\n int num_rows = A.get_num_rows();\n\n //Disabled since currently we are not doing halo exchanges during colored execution\n /*typedef TemplateConfig hvector_type;\n typedef Vector HVector;\n if (!A.is_matrix_singleGPU()) {\n HVector num_colors(1);\n std::vector partition_num_colors(0);\n num_colors[0] = m_num_colors;\n A.manager->getComms()->global_reduce(partition_num_colors, num_colors, A, 6332);\n int max_partition_colors = 0;\n for (int i = 0; i < partition_num_colors.size(); i++)\n max_partition_colors = std::max(partition_num_colors[i][0],max_partition_colors);\n m_num_colors = max_partition_colors;\n }*/\n\n if (m_halo_coloring == LAST)\n {\n thrust_wrapper::fill(m_row_colors.begin() + num_rows, m_row_colors.end(), m_num_colors);\n cudaCheckError();\n }\n\n IVector offsets_rows_per_color;\n\n if (m_offsets_rows_per_color.size() == 0)\n {\n // Sort the vertices based o their color\n m_sorted_rows_by_color.resize(num_rows);\n // Copy row colors\n IVector row_colors(m_row_colors);\n thrust_wrapper::sequence(m_sorted_rows_by_color.begin(), m_sorted_rows_by_color.end());\n amgx::thrust::sort_by_key(row_colors.begin(), row_colors.begin() + num_rows, m_sorted_rows_by_color.begin());\n cudaCheckError();\n // Compute the offset for each color\n offsets_rows_per_color.resize(m_num_colors + 1);\n m_offsets_rows_per_color.resize(m_num_colors + 1);\n // Compute interior-exterior separation for every color\n m_offsets_rows_per_color_separation.resize(m_num_colors);\n //m_offsets_rows_per_color_separation_halo.resize(m_num_colors);\n amgx::thrust::lower_bound(row_colors.begin(),\n row_colors.begin() + num_rows,\n amgx::thrust::counting_iterator(0),\n amgx::thrust::counting_iterator(offsets_rows_per_color.size()),\n offsets_rows_per_color.begin());\n // Copy from device to host\n m_offsets_rows_per_color = offsets_rows_per_color;\n cudaCheckError();\n }\n else\n {\n m_offsets_rows_per_color_separation.resize(m_num_colors);\n }\n\n cudaCheckError();\n\n if (!A.is_matrix_singleGPU() && (A.getViewExterior() != A.getViewInterior()))\n {\n A.setViewInterior();\n int separation = A.get_num_rows();\n\n if (TConfig::memSpace == AMGX_host)\n {\n for (int i = 0; i < m_num_colors; i++)\n {\n m_offsets_rows_per_color_separation[i] = m_offsets_rows_per_color[i]\n + (amgx::thrust::lower_bound(m_sorted_rows_by_color.begin() + m_offsets_rows_per_color[i],\n m_sorted_rows_by_color.begin() + m_offsets_rows_per_color[i + 1],\n separation)\n - (m_sorted_rows_by_color.begin() + m_offsets_rows_per_color[i]));\n }\n\n cudaCheckError();\n }\n // this is not a proper search, rather we look at every single element. But it is still a lot faster than the above (~10*)\n else\n {\n IVector separation_offsets_rows_per_color(m_num_colors);\n int size = num_rows;\n int num_blocks = std::min(4096, (size + 123) / 124);\n findSeparation <<< num_blocks, 128>>>(m_sorted_rows_by_color.raw(), offsets_rows_per_color.raw(), separation_offsets_rows_per_color.raw(), separation, m_num_colors, num_rows);\n cudaCheckError();\n amgx::thrust::copy(separation_offsets_rows_per_color.begin(), separation_offsets_rows_per_color.end(), m_offsets_rows_per_color_separation.begin());\n cudaCheckError();\n\n for (int i = 0; i < m_num_colors; i++)\n {\n if (this->m_offsets_rows_per_color[i] == this->m_offsets_rows_per_color[i + 1])\n {\n this->m_offsets_rows_per_color_separation[i] = this->m_offsets_rows_per_color[i + 1];\n }\n }\n }\n }\n else\n {\n amgx::thrust::copy(m_offsets_rows_per_color.begin() + 1, m_offsets_rows_per_color.end(), m_offsets_rows_per_color_separation.begin());\n cudaCheckError();\n }\n\n A.setView(old);\n}\n\ntemplate\nstd::map*> &\nMatrixColoringFactory::getFactories( )\n{\n static std::map *> s_factories;\n return s_factories;\n}\n\ntemplate\nvoid MatrixColoringFactory::registerFactory(std::string name, MatrixColoringFactory *f)\n{\n std::map *> &factories = getFactories( );\n typename std::map *>::iterator it = factories.find(name);\n\n if (it != factories.end())\n {\n std::string error = \"MatrixColoringFactory '\" + name + \"' has already been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n factories[name] = f;\n}\n\ntemplate\nvoid MatrixColoringFactory::unregisterFactory(std::string name)\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.find(name);\n\n if (it == factories.end())\n {\n std::string error = \"MatrixColoringFactory '\" + name + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n MatrixColoringFactory *factory = it->second;\n assert( factory != NULL );\n delete factory;\n factories.erase(it);\n}\n\ntemplate\nvoid MatrixColoringFactory::unregisterFactories( )\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.begin( );\n\n for ( ; it != factories.end( ) ; )\n {\n MatrixColoringFactory *factory = it->second;\n assert( factory != NULL );\n it++;\n delete factory;\n }\n\n factories.clear( );\n}\n\ntemplate\nMatrixColoring *MatrixColoringFactory::allocate(AMG_Config &cfg, const std::string &cfg_scope)\n{\n std::map *> &factories = getFactories( );\n std::string matrix_coloring_scheme = cfg.getParameter(\"matrix_coloring_scheme\", cfg_scope);\n typename std::map *>::const_iterator it = factories.find(matrix_coloring_scheme);\n\n if (it == factories.end())\n {\n std::string error = \"MatrixColoringFactory '\" + matrix_coloring_scheme + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n return it->second->create(cfg, cfg_scope);\n};\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class MatrixColoring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class MatrixColoringFactory::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n}\n\n", "source": "/kaggle/working/amgx/src/matrix_coloring/matrix_coloring.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n// Pseudo-random number generator\n\nnamespace amgx\n{\n\nstatic __host__ __device__ unsigned int hash_function(unsigned int a, unsigned int seed)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n\nstruct is_zero\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x == 0;\n }\n};\n\n// ---------------------------\n// Kernels\n// ---------------------------\n\ntemplate< int CTA_SIZE>\n__global__\nvoid color_kernel_warp( const int A_num_rows, const int *A_rows, const int *A_cols, int num_colors, int *new_color, int *A_colors )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / 32;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n volatile __shared__ unsigned long long colors_used[NUM_WARPS_PER_CTA][32];\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n colors_used[warp_id][lane_id] = 0ull;\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 ) // Already colored!!!\n {\n continue;\n }\n\n int row_hash = hash_function(row_id, 0);\n bool max_row = true;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += 32)\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, 0 );\n // Get the color of the column.\n int col_color = -1;\n\n if ( row_it < row_end && col_id < A_num_rows )\n {\n col_color = A_colors[col_id];\n }\n\n if (col_color != -1)\n {\n // Set the bit corresponding to that color to 1\n colors_used[warp_id][lane_id] |= (1 << col_color);\n }\n\n max_row &= (row_hash > col_hash || col_color != 0 || col_color == -1);\n }\n\n max_row &= utils::all( max_row );\n\n if (max_row) // If I'm a max\n {\n unsigned int active_mask = utils::activemask();\n // warp reduction to figure out what color to assign\n if (lane_id < 16) { colors_used[warp_id][lane_id] |= colors_used[warp_id][lane_id + 16]; } utils::syncwarp(active_mask);\n\n if (lane_id < 8) { colors_used[warp_id][lane_id] |= colors_used[warp_id][lane_id + 8]; } utils::syncwarp(active_mask);\n\n if (lane_id < 4) { colors_used[warp_id][lane_id] |= colors_used[warp_id][lane_id + 4]; } utils::syncwarp(active_mask);\n\n if (lane_id < 2) { colors_used[warp_id][lane_id] |= colors_used[warp_id][lane_id + 2]; } utils::syncwarp(active_mask);\n\n if (lane_id < 1) { colors_used[warp_id][lane_id] |= colors_used[warp_id][lane_id + 1]; } utils::syncwarp(active_mask);\n\n if (lane_id == 0)\n {\n // Find the first unassigned color, starting at 1\n#pragma unroll\n for (int i = 1; i < 64; i++)\n {\n unsigned long long mask = 1 << i;\n\n if ( !(colors_used[warp_id][0] & mask) )\n {\n A_colors[row_id] = i;\n\n if (i >= num_colors) { new_color[0] = 1; }\n\n break;\n }\n }\n }\n }\n }\n}\n\ntemplate< int CTA_SIZE>\n__global__\nvoid color_kernel_thread( const int A_num_rows, const int *A_rows, const int *A_cols, int *d_num_colors, int *d_new_color, int *A_colors, int *d_num_uncolored_block, int *d_done)\n{\n const int NUM_THREADS_PER_GRID = gridDim.x * CTA_SIZE;\n // Row identifier.\n int row_id = blockIdx.x * CTA_SIZE + threadIdx.x;\n\n if (d_done[0] == 1)\n {\n return;\n }\n\n int num_uncolored_thread = 0;\n\n // Iterate over rows.\n for ( ; row_id < A_num_rows ; row_id += NUM_THREADS_PER_GRID )\n {\n unsigned long long used_colors = 0ull;\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 ) // Already colored!!!\n {\n continue;\n }\n else\n {\n num_uncolored_thread++;\n }\n\n int row_hash = hash_function(row_id, 0);\n //int row_hash = row_id;\n bool max_row = true;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( int row_it = row_begin; row_it < row_end ; row_it++)\n {\n // Get the column index (if the iterator is valid).\n int col_id = A_cols[row_it];\n\n if (col_id >= A_num_rows || col_id == row_id) { continue; }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, 0 );\n // Get the color of the column.\n int col_color = A_colors[col_id];\n\n // Update the color.\n if ( col_color > 0)\n {\n used_colors |= 1ull << (64 - col_color);\n }\n\n // Is it still a max row?\n max_row &= (row_hash > col_hash || col_color != 0);\n }\n\n if ( max_row )\n {\n int my_color = 64 - utils::bfind( ~used_colors );\n\n if ( my_color > 0 )\n {\n A_colors[row_id] = my_color;\n }\n\n if ( my_color >= d_num_colors[0] )\n {\n d_new_color[0] = 1;\n }\n }\n }\n\n // Do per block reduction\n __shared__ volatile int smem[CTA_SIZE];\n smem[threadIdx.x] = num_uncolored_thread;\n __syncthreads();\n\n for ( int off = blockDim.x / 2; off >= 32; off = off / 2 )\n {\n if ( threadIdx.x < off )\n {\n smem[threadIdx.x] += smem[threadIdx.x + off];\n }\n\n __syncthreads();\n }\n\n // warp reduce\n if ( threadIdx.x < 32 )\n {\n smem[threadIdx.x] += smem[threadIdx.x + 16];\n smem[threadIdx.x] += smem[threadIdx.x + 8];\n smem[threadIdx.x] += smem[threadIdx.x + 4];\n smem[threadIdx.x] += smem[threadIdx.x + 2];\n smem[threadIdx.x] += smem[threadIdx.x + 1];\n }\n\n if ( threadIdx.x == 0 )\n {\n d_num_uncolored_block[blockIdx.x] = smem[0];\n }\n}\n\ntemplate\n__global__\nvoid count_uncolored_and_num_colors_kernel(int *d_new_color, int *d_num_colors, int max_uncolored_rows, int *d_num_uncolored_block, int *d_num_uncolored, int *d_done, int *h_done, int *d_start_color, int CTA_SIZE)\n{\n // Do block reduction\n __shared__ volatile int smem[MAX_CTA_SIZE];\n\n if (threadIdx.x >= CTA_SIZE)\n {\n smem[threadIdx.x] = 0;\n }\n else\n {\n smem[threadIdx.x] = d_num_uncolored_block[threadIdx.x];\n }\n\n __syncthreads();\n\n for ( int off = blockDim.x / 2; off >= 32; off = off / 2 )\n {\n if ( threadIdx.x < off )\n {\n smem[threadIdx.x] += smem[threadIdx.x + off];\n }\n\n __syncthreads();\n }\n\n // warp reduce\n if ( threadIdx.x < 32 )\n {\n smem[threadIdx.x] += smem[threadIdx.x + 16];\n smem[threadIdx.x] += smem[threadIdx.x + 8];\n smem[threadIdx.x] += smem[threadIdx.x + 4];\n smem[threadIdx.x] += smem[threadIdx.x + 2];\n smem[threadIdx.x] += smem[threadIdx.x + 1];\n }\n\n if ( threadIdx.x == 0 )\n {\n int prev_num_uncolored = d_num_uncolored[0];\n int num_uncolored = smem[0];\n int num_colors = d_num_colors[0];\n\n if (d_new_color[0])\n {\n num_colors++; // New color has been added\n d_new_color[0] = 0; // Reset to 0\n }\n\n // Check for termination\n if (TERMINATE_IF_FULL_COLORS_USED)\n {\n // Check for termination\n if (num_colors == 64 || prev_num_uncolored == num_uncolored || num_uncolored <= max_uncolored_rows)\n {\n h_done[0] = 1;\n d_done[0] = 1;\n }\n }\n else\n {\n if (prev_num_uncolored == num_uncolored || num_uncolored <= max_uncolored_rows)\n {\n if (num_uncolored <= max_uncolored_rows)\n {\n h_done[0] = 1;\n d_done[0] = 1;\n }\n else\n {\n d_start_color[0] += 32;\n }\n }\n }\n\n // Write to gmem\n d_num_uncolored[0] = num_uncolored;\n d_num_colors[0] = num_colors;\n }\n}\n\n\n//////////////////////////////\n// ANYRING version: begin\n// This unrolled recursion is used to visit neighbors.\n//////////////////////////////\n\ntemplate\nstruct parallel_greedy_neighborhood\n{\n __device__ __inline__ static void visit(int row_id, int row_hash, int visit_id, const int A_num_rows, const int *A_rows, const int *A_cols, int *A_colors,\n const int &start_color, bool &max_row, unsigned long long &used_colors\n )\n {\n int row_begin = A_rows[visit_id ];\n int row_end = A_rows[visit_id + 1];\n\n for ( int row_it = row_begin; row_it < row_end ; row_it++)\n {\n // Get the column index (if the iterator is valid).\n int col_id = A_cols[row_it];\n\n if (col_id >= A_num_rows || col_id == visit_id) { continue; }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, 0 );\n // Get the color of the column.\n int col_color = A_colors[col_id];\n\n // Update the color.\n if ( col_color - start_color > 0)\n {\n used_colors |= 1ull << (64 - col_color + start_color);\n }\n\n // Is it still a max row?\n max_row &= (row_hash > col_hash || (row_hash == col_hash && col_id >= row_id) || col_color != 0);\n\n if (!max_row) { break; }\n\n parallel_greedy_neighborhood < COLORING_LEVEL - 1 >::visit(row_id, row_hash, col_id, A_num_rows, A_rows, A_cols, A_colors, start_color, max_row, used_colors);\n\n if (!max_row) { break; }\n }\n }\n};\n\ntemplate<>\nstruct parallel_greedy_neighborhood<0>\n{\n __device__ __inline__ static void visit(int row_id, int row_hash, int visit_id, const int A_num_rows, const int *A_rows, const int *A_cols, int *A_colors,\n const int &start_color, bool &max_row, unsigned long long &used_colors\n )\n {\n ;\n }\n};\n\n\n//////////////////////////////\n// This kernel is equivalent to color_kernel_thread, but if COLORING_LEVEL>=2 it recurses to neighbors'+ neighbors.\n// it also uses start_color_ to continue coloring with parallel greedy instead of falling back to min_max if more than 64 colors are used\n//////////////////////////////\n\ntemplate< int CTA_SIZE, int COLORING_LEVEL >\n__global__\nvoid color_kernel_thread_anyring( const int A_num_rows, const int *A_rows, const int *A_cols, int *d_num_colors, int *d_new_color, int *A_colors, int *d_num_uncolored_block, int *d_done,\n const int *start_color_)\n{\n const int NUM_THREADS_PER_GRID = gridDim.x * CTA_SIZE;\n // Row identifier.\n int row_id = blockIdx.x * CTA_SIZE + threadIdx.x;\n\n if (d_done[0] == 1)\n {\n return;\n }\n\n const int start_color = start_color_[0];// (d_num_colors[0] / 64)*64;\n int num_uncolored_thread = 0;\n\n // Iterate over rows.\n for ( ; row_id < A_num_rows ; row_id += NUM_THREADS_PER_GRID )\n {\n unsigned long long used_colors = 0ull;\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 ) // Already colored!!!\n {\n continue;\n }\n else\n {\n num_uncolored_thread++;\n }\n\n int row_hash = hash_function(row_id, 0);\n //int row_hash = row_id;\n bool max_row = true;\n\n if (COLORING_LEVEL == 1)\n {\n // This path is equivalent to color_kernel_thread, but it uses start_color.\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( int row_it = row_begin; row_it < row_end ; row_it++)\n {\n // Get the column index (if the iterator is valid).\n int col_id = A_cols[row_it];\n\n if (col_id >= A_num_rows || col_id == row_id) { continue; }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, 0 );\n // Get the color of the column.\n int col_color = A_colors[col_id];\n\n // Update the color.\n if ( col_color - start_color > 0)\n {\n used_colors |= 1ull << (64 - col_color + start_color);\n }\n\n // Is it still a max row?\n max_row &= (row_hash > col_hash || (row_hash == col_hash && col_id >= row_id) || col_color != 0);\n\n if (!max_row) { break; }\n }\n }\n else\n {\n //recurse\n parallel_greedy_neighborhood::visit(row_id, row_hash, row_id, A_num_rows, A_rows, A_cols, A_colors, start_color, max_row, used_colors);\n }\n\n if ( max_row )\n {\n int my_color = 64 - utils::bfind( ~used_colors );\n\n if ( my_color > 0 && my_color < 64)\n {\n A_colors[row_id] = my_color + start_color;\n\n if ( my_color + start_color >= d_num_colors[0] )\n {\n d_new_color[0] = 1;\n }\n }\n }\n }\n\n // Do per block reduction\n __shared__ volatile int smem[CTA_SIZE];\n smem[threadIdx.x] = num_uncolored_thread;\n __syncthreads();\n\n for ( int off = blockDim.x / 2; off >= 32; off = off / 2 )\n {\n if ( threadIdx.x < off )\n {\n smem[threadIdx.x] += smem[threadIdx.x + off];\n }\n\n __syncthreads();\n }\n\n // warp reduce\n if ( threadIdx.x < 32 )\n {\n smem[threadIdx.x] += smem[threadIdx.x + 16];\n smem[threadIdx.x] += smem[threadIdx.x + 8];\n smem[threadIdx.x] += smem[threadIdx.x + 4];\n smem[threadIdx.x] += smem[threadIdx.x + 2];\n smem[threadIdx.x] += smem[threadIdx.x + 1];\n }\n\n if ( threadIdx.x == 0 )\n {\n d_num_uncolored_block[blockIdx.x] = smem[0];\n }\n}\n\n//////////////////////////////\n// ANYRING version: end\n//////////////////////////////\n\ntemplate< int CTA_SIZE>\n__global__\nvoid color_kernel_thread_fallback( const int A_num_rows, const int *A_rows, const int *A_cols, int num_colors, int *new_color, int *A_colors )\n{\n const int NUM_THREADS_PER_GRID = gridDim.x * CTA_SIZE;\n // Row identifier.\n int row_id = blockIdx.x * CTA_SIZE + threadIdx.x;\n\n // Iterate over rows.\n for ( ; row_id < A_num_rows ; row_id += NUM_THREADS_PER_GRID )\n {\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 ) // Already colored!!!\n {\n continue;\n }\n\n int row_hash = hash_function(row_id, 0);\n bool max_row = true;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n int max_color = 0;\n\n //if (row_id == 2) printf(\"row_id=%d\\n\",row_id);\n for ( int row_it = row_begin; row_it < row_end ; row_it++)\n {\n // Get the column index (if the iterator is valid).\n int col_id = A_cols[row_it];\n\n if (col_id >= A_num_rows) { continue; }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, 0 );\n // Get the color of the column.\n int col_color = A_colors[col_id];\n\n // Update the color.\n if ( col_color > max_color)\n {\n max_color = col_color;\n }\n\n // Is it still a max row?\n max_row &= (row_hash > col_hash || col_color != 0);\n }\n\n if ( max_row )\n {\n int my_color = max_color + 1;\n\n if ( my_color > 0 )\n {\n A_colors[row_id] = my_color;\n }\n\n if ( my_color >= num_colors )\n {\n new_color[0] = 1;\n }\n }\n }\n}\n\n\n\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid\ndbg_check_coloring_kernel( const int A_num_rows, const int *A_rows, const int *A_cols, const int *A_colors, const int *A_gtlt_count, int *error_found )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = A_colors[row_id];\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n // Get the color of the column.\n int col_color = -1;\n\n if ( row_it < row_end && col_id < A_num_rows )\n {\n col_color = A_colors[col_id];\n }\n\n // Is there something wrong ??\n if ( row_id != col_id && row_color == col_color )\n {\n if ( A_gtlt_count != NULL && !error_found[0] )\n {\n printf( \"row_id=%d, row_color=%d, col_id=%d, col_color=%d\\n\", row_id, row_color, col_id, col_color );\n }\n\n error_found[0] = 1;\n }\n }\n }\n}\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate< class T_Config >\nParallel_Greedy_Matrix_Coloring_Base::Parallel_Greedy_Matrix_Coloring_Base( AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n if ( this->m_coloring_level > 5 )\n {\n FatalError( \"Not implemented for coloring_level > 5\", AMGX_ERR_NOT_SUPPORTED_TARGET );\n }\n\n if ( cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\"))\n {\n m_uncolored_fraction = 0.0;\n }\n else\n {\n m_uncolored_fraction = cfg.AMG_Config::template getParameter(\"max_uncolored_percentage\", cfg_scope);\n }\n}\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nParallel_Greedy_Matrix_Coloring >::colorMatrix( Matrix_d &A )\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize( A.row_offsets.size() - 1 );\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n const int num_rows = A.get_num_rows();\n const int max_uncolored_rows = static_cast( this->m_uncolored_fraction * num_rows );\n const int CTA_SIZE = 128;\n const int MAX_GRID_SIZE = 1024;\n\n if (MAX_GRID_SIZE > 1024)\n {\n FatalError(\"Grid size must be less than 1024\\n\", AMGX_ERR_CORE);\n }\n\n const int GRID_SIZE = std::min( MAX_GRID_SIZE, (num_rows + CTA_SIZE - 1) / CTA_SIZE);\n thrust_wrapper::fill( this->m_row_colors.begin(), this->m_row_colors.end(), 0 );\n cudaCheckError();\n typedef typename Matrix_d::IVector IVector_d;\n cudaStream_t stream = amgx::thrust::global_thread_handle::get_stream();\n IVector_d d_new_color(1);\n cudaMemsetAsync(d_new_color.raw(), 0, sizeof(int), stream);\n cudaCheckError();\n this->m_num_colors = 1;\n IVector_d d_num_colors(1);\n IVector_d d_start_color(1); //for coloring safely without fallback\n cudaMemcpyAsync(d_num_colors.raw(), &this->m_num_colors, sizeof(int), cudaMemcpyHostToDevice, stream);\n cudaCheckError();\n int tmp = 0;\n cudaMemcpyAsync(d_start_color.raw(), &tmp, sizeof(int), cudaMemcpyHostToDevice, stream);\n cudaCheckError();\n IVector_d d_num_uncolored(1);\n IVector_d d_num_uncolored_block(MAX_GRID_SIZE);\n int *h_done = NULL;\n amgx::memory::cudaMallocHost( (void **) &h_done, sizeof(int));\n cudaCheckError();\n IVector_d d_done(1);\n d_done[0] = 0;\n *h_done = 0;\n cudaEvent_t throttle_event = 0;\n cudaEventCreateWithFlags(&throttle_event, cudaEventDisableTiming);\n int iteration = 0;\n\n while (h_done[0] == 0)\n {\n#define PG_ARGS num_rows,\\\n A.row_offsets.raw(),\\\n A.col_indices.raw(),\\\n d_num_colors.raw(),\\\n d_new_color.raw(),\\\n this->m_row_colors.raw(),\\\n d_num_uncolored_block.raw(),\\\n d_done.raw()\n\n if (this->m_coloring_level == 1)\n {\n color_kernel_thread <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(PG_ARGS);\n cudaCheckError();\n count_uncolored_and_num_colors_kernel <<< 1, MAX_GRID_SIZE, 0, stream>>>(d_new_color.raw(), d_num_colors.raw(), max_uncolored_rows, d_num_uncolored_block.raw(), d_num_uncolored.raw(), d_done.raw(), h_done, d_start_color.raw(), GRID_SIZE);\n cudaCheckError();\n }\n else\n {\n if (this->m_coloring_level == 2)\n {\n color_kernel_thread_anyring <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(PG_ARGS, d_start_color.raw());\n cudaCheckError();\n }\n else if (this->m_coloring_level == 3)\n {\n color_kernel_thread_anyring <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(PG_ARGS, d_start_color.raw());\n cudaCheckError();\n }\n else if (this->m_coloring_level == 4)\n {\n color_kernel_thread_anyring <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(PG_ARGS, d_start_color.raw());\n cudaCheckError();\n }\n else if (this->m_coloring_level == 5)\n {\n color_kernel_thread_anyring <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(PG_ARGS, d_start_color.raw());\n cudaCheckError();\n }\n\n count_uncolored_and_num_colors_kernel <<< 1, MAX_GRID_SIZE, 0, stream>>>(d_new_color.raw(), d_num_colors.raw(), max_uncolored_rows, d_num_uncolored_block.raw(), d_num_uncolored.raw(), d_done.raw(), h_done, d_start_color.raw(), GRID_SIZE);\n cudaCheckError();\n }\n\n // Throttle every 4 iteration\n if (iteration % 4 == 0)\n {\n cudaEventRecord(throttle_event);\n cudaCheckError();\n }\n else\n {\n cudaEventSynchronize(throttle_event);\n cudaCheckError();\n };\n\n iteration++;\n\n cudaCheckError();\n }\n\n typedef typename Matrix_h::IVector IVector_h;\n IVector_h new_color(1);\n cudaEventDestroy(throttle_event);\n cudaCheckError();\n this->m_num_colors = d_num_colors[0];\n int num_uncolored = d_num_uncolored[0];\n int prev_num_uncolored = 0;\n\n // Fallback path if # colors exceeds 64, instead of picking the smallest unassigned color,\n // it sets its color to max neighbor color + 1\n if (this->m_num_colors == 64)\n {\n // Choose highest unassigned color among the neighbors\n for ( int num_uncolored = num_rows ; num_uncolored > max_uncolored_rows && prev_num_uncolored != num_uncolored ; )\n {\n prev_num_uncolored = num_uncolored;\n new_color[0] = 0; // Host to device transfer\n color_kernel_thread_fallback <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_num_colors,\n amgx::thrust::raw_pointer_cast (&new_color.front()),\n this->m_row_colors.raw() );\n cudaCheckError();\n num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_zero() );\n\n if (new_color[0]) { this->m_num_colors++; }\n\n cudaCheckError();\n }\n }\n\n#if 0\n device_vector_alloc error_found( 1, 0 );\n dbg_check_coloring_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n NULL,\n amgx::thrust::raw_pointer_cast( &error_found.front() ) );\n cudaCheckError();\n\n if ( error_found[0] != 0 )\n {\n std::cout << \"INVALID COLORING !!! Two neighbors have the same color!!!\" << std::endl;\n }\n\n#endif\n amgx::memory::cudaFreeHost(h_done);\n cudaCheckError();\n A.setView(oldView);\n}\n\n\n#define AMGX_CASE_LINE(CASE) template class Parallel_Greedy_Matrix_Coloring_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Parallel_Greedy_Matrix_Coloring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // end namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n// Pseudo-random number generator\n\nnamespace amgx\n{\n\nstatic __host__ __device__ unsigned int hash_function(unsigned int a, unsigned int seed)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n\nstruct is_zero\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x == 0;\n }\n};\n\n// ---------------------------\n// Kernels\n// ---------------------------\n\ntemplate< int CTA_SIZE>\n__global__\nvoid color_kernel_warp( const int A_num_rows, const int *A_rows, const int *A_cols, int num_colors, int *new_color, int *A_colors )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / 32;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n volatile __shared__ unsigned long long colors_used[NUM_WARPS_PER_CTA][32];\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n colors_used[warp_id][lane_id] = 0ull;\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 ) // Already colored!!!\n {\n continue;\n }\n\n int row_hash = hash_function(row_id, 0);\n bool max_row = true;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += 32)\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, 0 );\n // Get the color of the column.\n int col_color = -1;\n\n if ( row_it < row_end && col_id < A_num_rows )\n {\n col_color = A_colors[col_id];\n }\n\n if (col_color != -1)\n {\n // Set the bit corresponding to that color to 1\n colors_used[warp_id][lane_id] |= (1 << col_color);\n }\n\n max_row &= (row_hash > col_hash || col_color != 0 || col_color == -1);\n }\n\n max_row &= utils::all( max_row );\n\n if (max_row) // If I'm a max\n {\n unsigned int active_mask = utils::activemask();\n // warp reduction to figure out what color to assign\n if (lane_id < 16) { colors_used[warp_id][lane_id] |= colors_used[warp_id][lane_id + 16]; } utils::syncwarp(active_mask);\n\n if (lane_id < 8) { colors_used[warp_id][lane_id] |= colors_used[warp_id][lane_id + 8]; } utils::syncwarp(active_mask);\n\n if (lane_id < 4) { colors_used[warp_id][lane_id] |= colors_used[warp_id][lane_id + 4]; } utils::syncwarp(active_mask);\n\n if (lane_id < 2) { colors_used[warp_id][lane_id] |= colors_used[warp_id][lane_id + 2]; } utils::syncwarp(active_mask);\n\n if (lane_id < 1) { colors_used[warp_id][lane_id] |= colors_used[warp_id][lane_id + 1]; } utils::syncwarp(active_mask);\n\n if (lane_id == 0)\n {\n // Find the first unassigned color, starting at 1\n#pragma unroll\n for (int i = 1; i < 64; i++)\n {\n unsigned long long mask = 1 << i;\n\n if ( !(colors_used[warp_id][0] & mask) )\n {\n A_colors[row_id] = i;\n\n if (i >= num_colors) { new_color[0] = 1; }\n\n break;\n }\n }\n }\n }\n }\n}\n\ntemplate< int CTA_SIZE>\n__global__\nvoid color_kernel_thread( const int A_num_rows, const int *A_rows, const int *A_cols, int *d_num_colors, int *d_new_color, int *A_colors, int *d_num_uncolored_block, int *d_done)\n{\n const int NUM_THREADS_PER_GRID = gridDim.x * CTA_SIZE;\n // Row identifier.\n int row_id = blockIdx.x * CTA_SIZE + threadIdx.x;\n\n if (d_done[0] == 1)\n {\n return;\n }\n\n int num_uncolored_thread = 0;\n\n // Iterate over rows.\n for ( ; row_id < A_num_rows ; row_id += NUM_THREADS_PER_GRID )\n {\n unsigned long long used_colors = 0ull;\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 ) // Already colored!!!\n {\n continue;\n }\n else\n {\n num_uncolored_thread++;\n }\n\n int row_hash = hash_function(row_id, 0);\n //int row_hash = row_id;\n bool max_row = true;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( int row_it = row_begin; row_it < row_end ; row_it++)\n {\n // Get the column index (if the iterator is valid).\n int col_id = A_cols[row_it];\n\n if (col_id >= A_num_rows || col_id == row_id) { continue; }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, 0 );\n // Get the color of the column.\n int col_color = A_colors[col_id];\n\n // Update the color.\n if ( col_color > 0)\n {\n used_colors |= 1ull << (64 - col_color);\n }\n\n // Is it still a max row?\n max_row &= (row_hash > col_hash || col_color != 0);\n }\n\n if ( max_row )\n {\n int my_color = 64 - utils::bfind( ~used_colors );\n\n if ( my_color > 0 )\n {\n A_colors[row_id] = my_color;\n }\n\n if ( my_color >= d_num_colors[0] )\n {\n d_new_color[0] = 1;\n }\n }\n }\n\n // Do per block reduction\n __shared__ volatile int smem[CTA_SIZE];\n smem[threadIdx.x] = num_uncolored_thread;\n __syncthreads();\n\n for ( int off = blockDim.x / 2; off >= 32; off = off / 2 )\n {\n if ( threadIdx.x < off )\n {\n smem[threadIdx.x] += smem[threadIdx.x + off];\n }\n\n __syncthreads();\n }\n\n // warp reduce\n if ( threadIdx.x < 32 )\n {\n smem[threadIdx.x] += smem[threadIdx.x + 16];\n smem[threadIdx.x] += smem[threadIdx.x + 8];\n smem[threadIdx.x] += smem[threadIdx.x + 4];\n smem[threadIdx.x] += smem[threadIdx.x + 2];\n smem[threadIdx.x] += smem[threadIdx.x + 1];\n }\n\n if ( threadIdx.x == 0 )\n {\n d_num_uncolored_block[blockIdx.x] = smem[0];\n }\n}\n\ntemplate\n__global__\nvoid count_uncolored_and_num_colors_kernel(int *d_new_color, int *d_num_colors, int max_uncolored_rows, int *d_num_uncolored_block, int *d_num_uncolored, int *d_done, int *h_done, int *d_start_color, int CTA_SIZE)\n{\n // Do block reduction\n __shared__ volatile int smem[MAX_CTA_SIZE];\n\n if (threadIdx.x >= CTA_SIZE)\n {\n smem[threadIdx.x] = 0;\n }\n else\n {\n smem[threadIdx.x] = d_num_uncolored_block[threadIdx.x];\n }\n\n __syncthreads();\n\n for ( int off = blockDim.x / 2; off >= 32; off = off / 2 )\n {\n if ( threadIdx.x < off )\n {\n smem[threadIdx.x] += smem[threadIdx.x + off];\n }\n\n __syncthreads();\n }\n\n // warp reduce\n if ( threadIdx.x < 32 )\n {\n smem[threadIdx.x] += smem[threadIdx.x + 16];\n smem[threadIdx.x] += smem[threadIdx.x + 8];\n smem[threadIdx.x] += smem[threadIdx.x + 4];\n smem[threadIdx.x] += smem[threadIdx.x + 2];\n smem[threadIdx.x] += smem[threadIdx.x + 1];\n }\n\n if ( threadIdx.x == 0 )\n {\n int prev_num_uncolored = d_num_uncolored[0];\n int num_uncolored = smem[0];\n int num_colors = d_num_colors[0];\n\n if (d_new_color[0])\n {\n num_colors++; // New color has been added\n d_new_color[0] = 0; // Reset to 0\n }\n\n // Check for termination\n if (TERMINATE_IF_FULL_COLORS_USED)\n {\n // Check for termination\n if (num_colors == 64 || prev_num_uncolored == num_uncolored || num_uncolored <= max_uncolored_rows)\n {\n h_done[0] = 1;\n d_done[0] = 1;\n }\n }\n else\n {\n if (prev_num_uncolored == num_uncolored || num_uncolored <= max_uncolored_rows)\n {\n if (num_uncolored <= max_uncolored_rows)\n {\n h_done[0] = 1;\n d_done[0] = 1;\n }\n else\n {\n d_start_color[0] += 32;\n }\n }\n }\n\n // Write to gmem\n d_num_uncolored[0] = num_uncolored;\n d_num_colors[0] = num_colors;\n }\n}\n\n\n//////////////////////////////\n// ANYRING version: begin\n// This unrolled recursion is used to visit neighbors.\n//////////////////////////////\n\ntemplate\nstruct parallel_greedy_neighborhood\n{\n __device__ __inline__ static void visit(int row_id, int row_hash, int visit_id, const int A_num_rows, const int *A_rows, const int *A_cols, int *A_colors,\n const int &start_color, bool &max_row, unsigned long long &used_colors\n )\n {\n int row_begin = A_rows[visit_id ];\n int row_end = A_rows[visit_id + 1];\n\n for ( int row_it = row_begin; row_it < row_end ; row_it++)\n {\n // Get the column index (if the iterator is valid).\n int col_id = A_cols[row_it];\n\n if (col_id >= A_num_rows || col_id == visit_id) { continue; }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, 0 );\n // Get the color of the column.\n int col_color = A_colors[col_id];\n\n // Update the color.\n if ( col_color - start_color > 0)\n {\n used_colors |= 1ull << (64 - col_color + start_color);\n }\n\n // Is it still a max row?\n max_row &= (row_hash > col_hash || (row_hash == col_hash && col_id >= row_id) || col_color != 0);\n\n if (!max_row) { break; }\n\n parallel_greedy_neighborhood < COLORING_LEVEL - 1 >::visit(row_id, row_hash, col_id, A_num_rows, A_rows, A_cols, A_colors, start_color, max_row, used_colors);\n\n if (!max_row) { break; }\n }\n }\n};\n\ntemplate<>\nstruct parallel_greedy_neighborhood<0>\n{\n __device__ __inline__ static void visit(int row_id, int row_hash, int visit_id, const int A_num_rows, const int *A_rows, const int *A_cols, int *A_colors,\n const int &start_color, bool &max_row, unsigned long long &used_colors\n )\n {\n ;\n }\n};\n\n\n//////////////////////////////\n// This kernel is equivalent to color_kernel_thread, but if COLORING_LEVEL>=2 it recurses to neighbors'+ neighbors.\n// it also uses start_color_ to continue coloring with parallel greedy instead of falling back to min_max if more than 64 colors are used\n//////////////////////////////\n\ntemplate< int CTA_SIZE, int COLORING_LEVEL >\n__global__\nvoid color_kernel_thread_anyring( const int A_num_rows, const int *A_rows, const int *A_cols, int *d_num_colors, int *d_new_color, int *A_colors, int *d_num_uncolored_block, int *d_done,\n const int *start_color_)\n{\n const int NUM_THREADS_PER_GRID = gridDim.x * CTA_SIZE;\n // Row identifier.\n int row_id = blockIdx.x * CTA_SIZE + threadIdx.x;\n\n if (d_done[0] == 1)\n {\n return;\n }\n\n const int start_color = start_color_[0];// (d_num_colors[0] / 64)*64;\n int num_uncolored_thread = 0;\n\n // Iterate over rows.\n for ( ; row_id < A_num_rows ; row_id += NUM_THREADS_PER_GRID )\n {\n unsigned long long used_colors = 0ull;\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 ) // Already colored!!!\n {\n continue;\n }\n else\n {\n num_uncolored_thread++;\n }\n\n int row_hash = hash_function(row_id, 0);\n //int row_hash = row_id;\n bool max_row = true;\n\n if (COLORING_LEVEL == 1)\n {\n // This path is equivalent to color_kernel_thread, but it uses start_color.\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( int row_it = row_begin; row_it < row_end ; row_it++)\n {\n // Get the column index (if the iterator is valid).\n int col_id = A_cols[row_it];\n\n if (col_id >= A_num_rows || col_id == row_id) { continue; }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, 0 );\n // Get the color of the column.\n int col_color = A_colors[col_id];\n\n // Update the color.\n if ( col_color - start_color > 0)\n {\n used_colors |= 1ull << (64 - col_color + start_color);\n }\n\n // Is it still a max row?\n max_row &= (row_hash > col_hash || (row_hash == col_hash && col_id >= row_id) || col_color != 0);\n\n if (!max_row) { break; }\n }\n }\n else\n {\n //recurse\n parallel_greedy_neighborhood::visit(row_id, row_hash, row_id, A_num_rows, A_rows, A_cols, A_colors, start_color, max_row, used_colors);\n }\n\n if ( max_row )\n {\n int my_color = 64 - utils::bfind( ~used_colors );\n\n if ( my_color > 0 && my_color < 64)\n {\n A_colors[row_id] = my_color + start_color;\n\n if ( my_color + start_color >= d_num_colors[0] )\n {\n d_new_color[0] = 1;\n }\n }\n }\n }\n\n // Do per block reduction\n __shared__ volatile int smem[CTA_SIZE];\n smem[threadIdx.x] = num_uncolored_thread;\n __syncthreads();\n\n for ( int off = blockDim.x / 2; off >= 32; off = off / 2 )\n {\n if ( threadIdx.x < off )\n {\n smem[threadIdx.x] += smem[threadIdx.x + off];\n }\n\n __syncthreads();\n }\n\n // warp reduce\n if ( threadIdx.x < 32 )\n {\n smem[threadIdx.x] += smem[threadIdx.x + 16];\n smem[threadIdx.x] += smem[threadIdx.x + 8];\n smem[threadIdx.x] += smem[threadIdx.x + 4];\n smem[threadIdx.x] += smem[threadIdx.x + 2];\n smem[threadIdx.x] += smem[threadIdx.x + 1];\n }\n\n if ( threadIdx.x == 0 )\n {\n d_num_uncolored_block[blockIdx.x] = smem[0];\n }\n}\n\n//////////////////////////////\n// ANYRING version: end\n//////////////////////////////\n\ntemplate< int CTA_SIZE>\n__global__\nvoid color_kernel_thread_fallback( const int A_num_rows, const int *A_rows, const int *A_cols, int num_colors, int *new_color, int *A_colors )\n{\n const int NUM_THREADS_PER_GRID = gridDim.x * CTA_SIZE;\n // Row identifier.\n int row_id = blockIdx.x * CTA_SIZE + threadIdx.x;\n\n // Iterate over rows.\n for ( ; row_id < A_num_rows ; row_id += NUM_THREADS_PER_GRID )\n {\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 ) // Already colored!!!\n {\n continue;\n }\n\n int row_hash = hash_function(row_id, 0);\n bool max_row = true;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n int max_color = 0;\n\n //if (row_id == 2) printf(\"row_id=%d\\n\",row_id);\n for ( int row_it = row_begin; row_it < row_end ; row_it++)\n {\n // Get the column index (if the iterator is valid).\n int col_id = A_cols[row_it];\n\n if (col_id >= A_num_rows) { continue; }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, 0 );\n // Get the color of the column.\n int col_color = A_colors[col_id];\n\n // Update the color.\n if ( col_color > max_color)\n {\n max_color = col_color;\n }\n\n // Is it still a max row?\n max_row &= (row_hash > col_hash || col_color != 0);\n }\n\n if ( max_row )\n {\n int my_color = max_color + 1;\n\n if ( my_color > 0 )\n {\n A_colors[row_id] = my_color;\n }\n\n if ( my_color >= num_colors )\n {\n new_color[0] = 1;\n }\n }\n }\n}\n\n\n\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid\ndbg_check_coloring_kernel( const int A_num_rows, const int *A_rows, const int *A_cols, const int *A_colors, const int *A_gtlt_count, int *error_found )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = A_colors[row_id];\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n // Get the color of the column.\n int col_color = -1;\n\n if ( row_it < row_end && col_id < A_num_rows )\n {\n col_color = A_colors[col_id];\n }\n\n // Is there something wrong ??\n if ( row_id != col_id && row_color == col_color )\n {\n if ( A_gtlt_count != NULL && !error_found[0] )\n {\n printf( \"row_id=%d, row_color=%d, col_id=%d, col_color=%d\\n\", row_id, row_color, col_id, col_color );\n }\n\n error_found[0] = 1;\n }\n }\n }\n}\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate< class T_Config >\nParallel_Greedy_Matrix_Coloring_Base::Parallel_Greedy_Matrix_Coloring_Base( AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n if ( this->m_coloring_level > 5 )\n {\n FatalError( \"Not implemented for coloring_level > 5\", AMGX_ERR_NOT_SUPPORTED_TARGET );\n }\n\n if ( cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\"))\n {\n m_uncolored_fraction = 0.0;\n }\n else\n {\n m_uncolored_fraction = cfg.AMG_Config::template getParameter(\"max_uncolored_percentage\", cfg_scope);\n }\n}\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nParallel_Greedy_Matrix_Coloring >::colorMatrix( Matrix_d &A )\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize( A.row_offsets.size() - 1 );\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n const int num_rows = A.get_num_rows();\n const int max_uncolored_rows = static_cast( this->m_uncolored_fraction * num_rows );\n const int CTA_SIZE = 128;\n const int MAX_GRID_SIZE = 1024;\n\n if (MAX_GRID_SIZE > 1024)\n {\n FatalError(\"Grid size must be less than 1024\\n\", AMGX_ERR_CORE);\n }\n\n const int GRID_SIZE = std::min( MAX_GRID_SIZE, (num_rows + CTA_SIZE - 1) / CTA_SIZE);\n thrust_wrapper::fill( this->m_row_colors.begin(), this->m_row_colors.end(), 0 );\n cudaCheckError();\n typedef typename Matrix_d::IVector IVector_d;\n hipStream_t stream = amgx::thrust::global_thread_handle::get_stream();\n IVector_d d_new_color(1);\n hipMemsetAsync(d_new_color.raw(), 0, sizeof(int), stream);\n cudaCheckError();\n this->m_num_colors = 1;\n IVector_d d_num_colors(1);\n IVector_d d_start_color(1); //for coloring safely without fallback\n hipMemcpyAsync(d_num_colors.raw(), &this->m_num_colors, sizeof(int), hipMemcpyHostToDevice, stream);\n cudaCheckError();\n int tmp = 0;\n hipMemcpyAsync(d_start_color.raw(), &tmp, sizeof(int), hipMemcpyHostToDevice, stream);\n cudaCheckError();\n IVector_d d_num_uncolored(1);\n IVector_d d_num_uncolored_block(MAX_GRID_SIZE);\n int *h_done = NULL;\n amgx::memory::hipHostMalloc( (void **) &h_done, sizeof(int));\n cudaCheckError();\n IVector_d d_done(1);\n d_done[0] = 0;\n *h_done = 0;\n hipEvent_t throttle_event = 0;\n hipEventCreateWithFlags(&throttle_event, hipEventDisableTiming);\n int iteration = 0;\n\n while (h_done[0] == 0)\n {\n#define PG_ARGS num_rows,\\\n A.row_offsets.raw(),\\\n A.col_indices.raw(),\\\n d_num_colors.raw(),\\\n d_new_color.raw(),\\\n this->m_row_colors.raw(),\\\n d_num_uncolored_block.raw(),\\\n d_done.raw()\n\n if (this->m_coloring_level == 1)\n {\n color_kernel_thread <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(PG_ARGS);\n cudaCheckError();\n count_uncolored_and_num_colors_kernel <<< 1, MAX_GRID_SIZE, 0, stream>>>(d_new_color.raw(), d_num_colors.raw(), max_uncolored_rows, d_num_uncolored_block.raw(), d_num_uncolored.raw(), d_done.raw(), h_done, d_start_color.raw(), GRID_SIZE);\n cudaCheckError();\n }\n else\n {\n if (this->m_coloring_level == 2)\n {\n color_kernel_thread_anyring <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(PG_ARGS, d_start_color.raw());\n cudaCheckError();\n }\n else if (this->m_coloring_level == 3)\n {\n color_kernel_thread_anyring <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(PG_ARGS, d_start_color.raw());\n cudaCheckError();\n }\n else if (this->m_coloring_level == 4)\n {\n color_kernel_thread_anyring <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(PG_ARGS, d_start_color.raw());\n cudaCheckError();\n }\n else if (this->m_coloring_level == 5)\n {\n color_kernel_thread_anyring <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(PG_ARGS, d_start_color.raw());\n cudaCheckError();\n }\n\n count_uncolored_and_num_colors_kernel <<< 1, MAX_GRID_SIZE, 0, stream>>>(d_new_color.raw(), d_num_colors.raw(), max_uncolored_rows, d_num_uncolored_block.raw(), d_num_uncolored.raw(), d_done.raw(), h_done, d_start_color.raw(), GRID_SIZE);\n cudaCheckError();\n }\n\n // Throttle every 4 iteration\n if (iteration % 4 == 0)\n {\n hipEventRecord(throttle_event);\n cudaCheckError();\n }\n else\n {\n hipEventSynchronize(throttle_event);\n cudaCheckError();\n };\n\n iteration++;\n\n cudaCheckError();\n }\n\n typedef typename Matrix_h::IVector IVector_h;\n IVector_h new_color(1);\n hipEventDestroy(throttle_event);\n cudaCheckError();\n this->m_num_colors = d_num_colors[0];\n int num_uncolored = d_num_uncolored[0];\n int prev_num_uncolored = 0;\n\n // Fallback path if # colors exceeds 64, instead of picking the smallest unassigned color,\n // it sets its color to max neighbor color + 1\n if (this->m_num_colors == 64)\n {\n // Choose highest unassigned color among the neighbors\n for ( int num_uncolored = num_rows ; num_uncolored > max_uncolored_rows && prev_num_uncolored != num_uncolored ; )\n {\n prev_num_uncolored = num_uncolored;\n new_color[0] = 0; // Host to device transfer\n color_kernel_thread_fallback <<< GRID_SIZE, CTA_SIZE, 0, stream>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_num_colors,\n amgx::thrust::raw_pointer_cast (&new_color.front()),\n this->m_row_colors.raw() );\n cudaCheckError();\n num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_zero() );\n\n if (new_color[0]) { this->m_num_colors++; }\n\n cudaCheckError();\n }\n }\n\n#if 0\n device_vector_alloc error_found( 1, 0 );\n dbg_check_coloring_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n NULL,\n amgx::thrust::raw_pointer_cast( &error_found.front() ) );\n cudaCheckError();\n\n if ( error_found[0] != 0 )\n {\n std::cout << \"INVALID COLORING !!! Two neighbors have the same color!!!\" << std::endl;\n }\n\n#endif\n amgx::memory::hipHostFree(h_done);\n cudaCheckError();\n A.setView(oldView);\n}\n\n\n#define AMGX_CASE_LINE(CASE) template class Parallel_Greedy_Matrix_Coloring_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Parallel_Greedy_Matrix_Coloring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // end namespace amgx\n\n", "source": "/kaggle/working/amgx/src/matrix_coloring/parallel_greedy.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\n// ---------------------------\n// Kernels\n// ---------------------------\n\ntemplate \n__global__\nvoid color3d(IndexType *row_colors, const int n, const int n1d)\n{\n unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int n1d2 = n1d * n1d;\n int x, y, z;\n\n for (; idx < n; idx += gridDim.x * blockDim.x)\n {\n int label = idx % n1d2;\n x = (label % n1d) % 2;\n y = (label / n1d) % 2;\n z = (idx / n1d2) % 2 ;\n row_colors[idx] = z * 4 + y * 2 + x;\n }\n}\n\ntemplate \n__global__\nvoid color2d(IndexType *row_colors, const int n, const int n1d)\n{\n unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int x, y;\n\n for (; idx < n; idx += gridDim.x * blockDim.x)\n {\n x = (idx % n1d) % 2;\n y = (idx / n1d) % 2;\n row_colors[idx] = 2 * y + x;\n }\n}\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate\nUniformMatrixColoringBase::UniformMatrixColoringBase(AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n}\n\n\n// Block version\ntemplate \nvoid UniformMatrixColoringBase::colorMatrix(Matrix &A)\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize(A.row_offsets.size() - 1, 0);\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n if (this->m_coloring_level == 0)\n {\n FatalError(\"Calling coloring scheme but coloring level==0\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n }\n else\n {\n int num_rows = A.get_num_rows();\n int dim = 2;\n this->m_num_colors = 2 << dim;\n\n if (A.get_num_rows() < this->m_num_colors)\n {\n thrust_wrapper::sequence(this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows);\n cudaCheckError();\n this->m_num_colors = A.get_num_rows();\n }\n else\n {\n IndexType *row_colors_ptr = this->m_row_colors.raw();\n const int threads_per_block = 512;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_rows - 1) / threads_per_block + 1 );\n\n if (dim == 3)\n {\n int n1d = round(cbrt((double)num_rows));\n //printf(\"Uniform coloring: n==%d, dim == %d, n1d == %d\\n\", num_rows, dim, n1d);\n color3d <<< num_blocks, threads_per_block>>>(row_colors_ptr, num_rows, n1d);\n cudaCheckError();\n }\n else\n {\n int n1d = round(sqrt((double)num_rows));\n //printf(\"Uniform coloring: n==%d, dim == %d, n1d == %d\\n\", num_rows, dim, n1d);\n color2d <<< num_blocks, threads_per_block>>>(row_colors_ptr, num_rows, n1d);\n }\n\n cudaCheckError();\n }\n }\n\n A.setView(oldView);\n}\n\n\n#define AMGX_CASE_LINE(CASE) template class UniformMatrixColoringBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class UniformMatrixColoring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n\n} // end namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\n// ---------------------------\n// Kernels\n// ---------------------------\n\ntemplate \n__global__\nvoid color3d(IndexType *row_colors, const int n, const int n1d)\n{\n unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int n1d2 = n1d * n1d;\n int x, y, z;\n\n for (; idx < n; idx += gridDim.x * blockDim.x)\n {\n int label = idx % n1d2;\n x = (label % n1d) % 2;\n y = (label / n1d) % 2;\n z = (idx / n1d2) % 2 ;\n row_colors[idx] = z * 4 + y * 2 + x;\n }\n}\n\ntemplate \n__global__\nvoid color2d(IndexType *row_colors, const int n, const int n1d)\n{\n unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int x, y;\n\n for (; idx < n; idx += gridDim.x * blockDim.x)\n {\n x = (idx % n1d) % 2;\n y = (idx / n1d) % 2;\n row_colors[idx] = 2 * y + x;\n }\n}\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate\nUniformMatrixColoringBase::UniformMatrixColoringBase(AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n}\n\n\n// Block version\ntemplate \nvoid UniformMatrixColoringBase::colorMatrix(Matrix &A)\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize(A.row_offsets.size() - 1, 0);\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n if (this->m_coloring_level == 0)\n {\n FatalError(\"Calling coloring scheme but coloring level==0\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n }\n else\n {\n int num_rows = A.get_num_rows();\n int dim = 2;\n this->m_num_colors = 2 << dim;\n\n if (A.get_num_rows() < this->m_num_colors)\n {\n thrust_wrapper::sequence(this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows);\n cudaCheckError();\n this->m_num_colors = A.get_num_rows();\n }\n else\n {\n IndexType *row_colors_ptr = this->m_row_colors.raw();\n const int threads_per_block = 512;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_rows - 1) / threads_per_block + 1 );\n\n if (dim == 3)\n {\n int n1d = round(cbrt((double)num_rows));\n //printf(\"Uniform coloring: n==%d, dim == %d, n1d == %d\\n\", num_rows, dim, n1d);\n color3d <<< num_blocks, threads_per_block>>>(row_colors_ptr, num_rows, n1d);\n cudaCheckError();\n }\n else\n {\n int n1d = round(sqrt((double)num_rows));\n //printf(\"Uniform coloring: n==%d, dim == %d, n1d == %d\\n\", num_rows, dim, n1d);\n color2d <<< num_blocks, threads_per_block>>>(row_colors_ptr, num_rows, n1d);\n }\n\n cudaCheckError();\n }\n }\n\n A.setView(oldView);\n}\n\n\n#define AMGX_CASE_LINE(CASE) template class UniformMatrixColoringBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class UniformMatrixColoring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n\n} // end namespace amgx\n\n", "source": "/kaggle/working/amgx/src/matrix_coloring/uniform.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n// Pseudo-random number generator\nnamespace amgx\n{\n\nstatic __host__ __device__ __forceinline__ unsigned int hash_function(unsigned int a, unsigned int seed, unsigned int rows = 0)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n\n// ---------------------------\n// Kernels: Any ring parallel greedy with color buffers and max id propagation.\n// ---------------------------\n\n\n/*\n * This kernel computes the index of the node with maximum hash in a neighborhood and the colors used in a vertex neighborhood.\n * Since both operation are associative (max and or), it is possible to iterate the operation propagating the values across neigborhoods (first, second and so on rings).\n *\n * The N-th iteration computes, for each vertex:\n * 1. The ID of the uncolored vertex with ID of the vertex having max hash inside the N-ring neighborhood.\n * If all neighbors in the N-hood are colored, saves -1 (this allows work efficiency).\n * 2. The buffer of used colors in the n-ring neighborhood.\n *\n * And for that, it uses data from previous iteration.\n *\n * For avoiding the waste of work, once any n-ring neighborhood is all colored, it starts getting ignored:\n * no new MAX HASH can be found there and the set of used colors for that ring will not change.\n */\n\n\ntemplate< bool USE_EXT, int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid find_max_neighbor_kernel_and_propagate_used_colors( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const int *__restrict A_colors,\n\n int *level_n0_max_id_in, //for multiple iterations. load the ID of the neighbor with max hash in n-1 ring hood\n int *level_n1_out_max_id,//for storing the ID of the n-ring neighbor with max hash (in the n-ring hood)\n\n unsigned long long int *level_n0_used_colors_in, //for multiple iterations: load used color buffers from n-1 level\n unsigned long long int *level_n1_used_colors_out,//for storing the used color buffers of the n-level neighborhood of each vertex\n\n //int* interesting,\n const int currentSafeColor, //starts assigning colors from this. multiple of 64\n const int seed)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n const int n1_max_id = level_n1_out_max_id[row_id];\n\n if (n1_max_id == -1)\n {\n //-1 means that all the vertices in this n-ring\n //level neighborhood were colored during the previous coloring epochs.\n //Since nothing could have changed, just skip:\n //save work and use the color buffers computed in last step.\n continue;\n }\n\n //If the old max 'n1_max_id' is still uncolored, it still will be the max in this iteration.\n //Plus, nothing else would get colored in its n1-hood, since 'n1_max_id' is the max.\n if (A_colors[n1_max_id] <= 0)\n {\n continue;\n }\n\n int row_color = A_colors[row_id];\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n unsigned long long used_colors = 0ull;\n\n //For multiple iterations: load the used color buffer from previous iter.\n if (USE_EXT)\n {\n used_colors = level_n0_used_colors_in[row_id];\n }\n\n if (row_color - currentSafeColor > 0) //if the color belongs to the current buffer [currentSafeColor,currentSafeColor+64)\n {\n used_colors |= (1ull << (64 - row_color + currentSafeColor)); //add my color to buffer\n }\n\n // Hash my row id.\n int max_hash = hash_function( row_id, seed, 0 );\n int max_hash_id = row_id; //id of the neighbor with max hash, initialized with myself\n\n //If already colored, don't count me for finding max\n if (row_color != 0)\n {\n max_hash_id = -1; //-1 means no max\n max_hash = 0;//LONG_LONG_MIN;\n }\n\n //iterate the neighbors\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, seed, 0);\n // Get the color of the column.\n int col_color = -1;\n\n if ( row_it < row_end && col_id < A_num_rows)\n {\n col_color = A_colors[col_id];\n\n if ( col_color - currentSafeColor > 0 ) //if in the current buffer..\n {\n used_colors |= (1ull << (64 - col_color + currentSafeColor)); //..add it to buffer\n }\n\n if (USE_EXT)\n {\n //combine used_colors buffer with the ones computed in previous iterations\n used_colors |= level_n0_used_colors_in[col_id];\n col_id = level_n0_max_id_in[col_id];\n\n if (col_id < 0) //row and its k-1 level neigbhors all colored\n {\n col_hash = 0;//LONG_LONG_MIN;\n col_color = -1;\n }\n else\n {\n col_hash = hash_function(col_id, seed, 0); //col hash is the max hash of its neighbors in n-1 level\n col_color = 0; //there is an uncolored neighbor\n }\n }\n }\n\n //check if any neighbor has a greater max id.\n if (col_color == 0 && (max_hash_id < 0 || col_hash > max_hash || (col_hash == max_hash && col_id >= max_hash_id)))\n {\n max_hash_id = col_id;\n max_hash = col_hash;\n }\n }\n\n //BEGIN: Reduce used_colors/max_hashes amongst subwarps\n#pragma unroll\n\n for (int i = WARP_SIZE / 2; i >= 1; i /= 2)\n {\n int tmp = utils::shfl_xor(max_hash, i, WARP_SIZE); //ss_max_hash[threadIdx.x+i];\n int tmp_id = utils::shfl_xor(max_hash_id, i, WARP_SIZE);\n\n if (tmp_id >= 0 && (max_hash_id < 0 || tmp > max_hash || (tmp == max_hash && tmp_id >= max_hash_id)))\n {\n max_hash = tmp;\n max_hash_id = tmp_id;\n }\n\n //\n int tmp_hi = __double2hiint( __longlong_as_double( used_colors ) );\n int tmp_lo = __double2loint( __longlong_as_double( used_colors ) );\n tmp_hi = utils::shfl_xor(tmp_hi, i, WARP_SIZE);\n tmp_lo = utils::shfl_xor(tmp_lo, i, WARP_SIZE);\n long long tmpu = __double_as_longlong(__hiloint2double(tmp_hi, tmp_lo));\n used_colors |= tmpu;\n }\n\n //END: Reduce used_colors/max_hashes amongst subwarps\n\n //The subwarp leader stores the result.\n if ( lane_id == 0 )\n {\n //if max_hash_id=1, will skip work next time: every n-ring neighbor is colored\n level_n1_out_max_id[row_id] = max_hash_id;\n level_n1_used_colors_out[row_id] = used_colors;\n }\n }\n}\n\n\n/*\n * This kernel colors rows using a greedy strategy and the used colors buffer from propagate iterations.\n */\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid color_kernel_greedy_onlymax(\n const int A_num_rows, const int *A_rows, const int *A_cols,\n const int *max_hash_id_in_neighborhood,\n\n const int seed,\n int *A_colors,\n const unsigned long long *used_colors_in,\n int *finished_colors_flag,\n const int currentSafeColor)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n //int dice = hash_function(row_id, 19881988, 0);\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 ) // Already colored!!!\n {\n continue;\n }\n\n //int row_hash = hash_function(row_id,seed,0);\n // The number of vertices that are greater/smaller than me.\n int max_id = max_hash_id_in_neighborhood[row_id];\n\n if (max_id != row_id) //was not a max in its own neighborhood\n {\n continue;\n }\n\n int max_hash = hash_function(max_id, seed, 0);\n //A_colors[row_id] = 1;\n //continue;\n bool is_max_vertex = true;\n //bool all_neighbors_colored = true;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n unsigned long long used_colors = 0ull;\n used_colors |= used_colors_in[row_id];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n // Filter out used colors\n if ( row_it < row_end && col_id < A_num_rows)\n {\n //int col_color = A_colors[col_id];\n //if(col_color == 0) all_neighbors_colored = false;\n used_colors |= used_colors_in[col_id];\n int nn_max_id = max_hash_id_in_neighborhood[col_id];\n int nn_max = hash_function(nn_max_id, seed, 0);\n is_max_vertex &= (max_hash > nn_max || (max_hash == nn_max && row_id >= nn_max_id));\n }\n\n if (!is_max_vertex) { break; }\n }\n\n //reduce used colors bit by bit.\n#pragma unroll\n\n for (int i = WARP_SIZE / 2; i >= 1; i /= 2)\n {\n int tmp_hi = __double2hiint( __longlong_as_double( used_colors ) );\n int tmp_lo = __double2loint( __longlong_as_double( used_colors ) );\n tmp_hi = utils::shfl_xor(tmp_hi, i, WARP_SIZE);\n tmp_lo = utils::shfl_xor(tmp_lo, i, WARP_SIZE);\n long long tmp = __double_as_longlong(__hiloint2double(tmp_hi, tmp_lo));\n used_colors |= tmp;\n }\n\n int my_color_1 = 64 - utils::bfind( ~used_colors );\n\n if (__popc(used_colors) >= 64 || my_color_1 > 64 || my_color_1 <= 0)\n {\n if (lane_id == 0)\n {\n finished_colors_flag[0] = 1;\n }\n\n is_max_vertex = 0;\n }\n\n my_color_1 += currentSafeColor;\n\n if ( subwarp::all( is_max_vertex ) )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = my_color_1;\n }\n\n continue;\n }\n }\n}\n\n\n/*\n *\n * One ring parallel greedy with greater hash count metric instead of raw hash metric. Similar to MIN_MAX_2RING\n *\n * */\ntemplate< int CTA_SIZE, int WARP_SIZE, bool USE_EXTERNAL_METRIC >\n__global__\nvoid count_gtlt_kernel( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const int *__restrict A_colors,\n int *A_gtlt_count,\n const int seed,\n const int *metric_in)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 )\n {\n if ( lane_id == 0 )\n {\n A_gtlt_count[row_id] = -1;\n }\n\n continue;\n }\n\n // The number of vertices that are greater/smaller than me.\n int gt_count = 0, lt_count = 0;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n // Hash my row id.\n int row_cols = row_end - row_begin;\n int row_hash = hash_function( row_id, seed, row_cols );\n\n if (USE_EXTERNAL_METRIC)\n {\n row_hash = metric_in[row_id] >> 16; //gt count\n }\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n int col_cols = 1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n#if USE_ROWS\n col_cols = A_rows[col_id + 1] - A_rows[col_id];\n#endif\n }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, seed, col_cols);\n // Get the color of the column.\n int col_color = -1;\n\n if ( row_it < row_end && col_id < A_num_rows)\n {\n col_color = A_colors[col_id];\n\n if (USE_EXTERNAL_METRIC)\n {\n col_hash = metric_in[col_id] >> 16; //gt count\n }\n }\n\n // Threads determine if they are greater than the row hash.\n int gt_pred = col_color == 0 && col_hash > row_hash;\n int lt_pred = col_color == 0 && col_hash < row_hash;\n // Count greater/smaller neighbors.\n gt_count += __popc( subwarp::masked_ballot( gt_pred ) );//subwarp::masked_ballot\n lt_count += __popc( subwarp::masked_ballot( lt_pred ) );\n }\n\n // The warp leader stores the result.\n int my_gtlt_count = (gt_count << 16) | lt_count;\n\n if ( lane_id == 0 )\n {\n A_gtlt_count[row_id] = my_gtlt_count;\n }\n }\n}\ntemplate< int CTA_SIZE, int WARP_SIZE, bool LOAD_EXTERNAL_USED_COLORS >\n__global__\nvoid color_kernel_greedy_gtlt(\n const int A_num_rows, const int *A_rows, const int *A_cols,\n const int *A_gtlt_count, const int current_color,\n const int weakness_bound, const int max_color_to_use,\n int *A_colors,\n const int min_degree, const unsigned long long *used_colors_in)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n //int dice = hash_function(row_id, 19881988, 0);\n int row_color = A_colors[row_id];\n\n if ( row_color > 0 ) // Already colored!!!\n {\n continue;\n }\n\n if ( row_color < 0 )\n {\n A_colors[row_id] = 0; //discovered\n\n if (row_id != 0) { continue; } //process next time\n }\n\n // The number of vertices that are greater/smaller than me.\n int row_gtlt_count = A_gtlt_count[row_id];\n // Split gtlt_count into 2.\n int row_gt_count = row_gtlt_count >> 16;\n int row_lt_count = row_gtlt_count & 0xffff;\n // Do we skip it.\n int candidate = 1;\n // Predicates. Is a vertex min/max.\n int is_max_vertex = row_gt_count <= weakness_bound;\n int is_min_vertex = row_lt_count <= weakness_bound;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n int degree = row_end - row_begin;\n\n if (degree < min_degree && min_degree > 0)\n {\n continue;\n }\n\n unsigned long long used_colors = 0ull;\n\n if (LOAD_EXTERNAL_USED_COLORS)\n {\n used_colors |= used_colors_in[row_id];\n }\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n if ( row_it < row_end && col_id < A_num_rows )\n {\n int col_color = A_colors[col_id];\n\n if ( col_color > 0 )\n {\n used_colors |= (1ull << (64 - col_color));\n }\n }\n\n if (LOAD_EXTERNAL_USED_COLORS)\n {\n used_colors |= used_colors_in[col_id];\n }\n\n // Get the gt/lt count.\n int col_gtlt_count = -1;\n\n if ( row_it < row_end && col_id < A_num_rows )\n {\n col_gtlt_count = A_gtlt_count[col_id];\n }\n\n // Split gtlt_count into 2. col_gt/lt_count == -1 if already colored.\n int col_gt_count = col_gtlt_count >> 16;\n int col_lt_count = col_gtlt_count & 0xffff;\n\n // Threads determine if they are greater than the row hash.\n //if( col_gt_count != -1 )\n // is_max_vertex &= col_gt_count < row_gt_count || (col_gt_count == row_gt_count && col_id < row_id);\n //if( col_lt_count != -1 )\n // is_min_vertex &= col_lt_count /*>*/ < row_lt_count || (col_lt_count == row_lt_count && col_id > row_id);\n if ( col_gtlt_count != -1 )\n {\n is_max_vertex &= col_gt_count > row_gt_count || (col_gt_count == row_gt_count && col_id >= row_id);\n is_min_vertex &= col_lt_count > row_lt_count || (col_lt_count == row_lt_count && col_id >= row_id);\n }\n }\n\n //is_max_vertex = row_gt_count==0;\n //is_min_vertex = false;\n //reduce used colors bit by bit.\n#pragma unroll\n\n for (int i = WARP_SIZE / 2; i >= 1; i /= 2)\n {\n int tmp_hi = __double2hiint( __longlong_as_double( used_colors ) );\n int tmp_lo = __double2loint( __longlong_as_double( used_colors ) );\n tmp_hi = utils::shfl_xor(tmp_hi, i, WARP_SIZE);\n tmp_lo = utils::shfl_xor(tmp_lo, i, WARP_SIZE);\n long long tmp = __double_as_longlong(__hiloint2double(tmp_hi, tmp_lo));\n used_colors |= tmp;\n }\n\n int my_color_1 = 0;\n int my_color_2 = 0;\n int free_colors = __popc(used_colors);\n my_color_1 = 64 - utils::bfind( ~used_colors );\n used_colors |= (1ull << (64 - my_color_1));\n my_color_2 = 64 - utils::bfind( ~used_colors );\n free_colors = 0;\n\n if ( candidate && subwarp::all( is_max_vertex ) && my_color_1 < max_color_to_use)\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = my_color_1;\n }\n\n continue;\n }\n\n if ( candidate && subwarp::all( is_min_vertex ) && free_colors >= 2)\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = my_color_2;\n }\n\n continue;\n }\n }\n}\n\nstruct is_zero\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x == 0;\n }\n};\nstruct is_minus_one\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x == -1;\n }\n};\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate< class T_Config >\nGreedy_Min_Max_2Ring_Matrix_Coloring_Base::Greedy_Min_Max_2Ring_Matrix_Coloring_Base( AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n if ( this->m_coloring_level > 5 || this->m_coloring_level < 0)\n {\n FatalError( \"Not implemented for coloring_level != 1\", AMGX_ERR_NOT_SUPPORTED_TARGET );\n }\n\n if ( cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\") == 1 )\n {\n m_uncolored_fraction = 0.0;\n }\n else\n {\n m_uncolored_fraction = cfg.AMG_Config::template getParameter(\"max_uncolored_percentage\", cfg_scope);\n }\n\n m_weakness_bound = cfg.AMG_Config::template getParameter( \"weakness_bound\", cfg_scope );\n m_late_rejection = cfg.AMG_Config::template getParameter( \"late_rejection\", cfg_scope ) != 0;\n m_coloring_try_remove_last_color = cfg.AMG_Config::template getParameter( \"coloring_try_remove_last_colors\", cfg_scope );\n m_coloring_custom_arg = cfg.AMG_Config::template getParameter( \"coloring_custom_arg\", cfg_scope );\n}\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\ntemplate\nvoid\nGreedy_Min_Max_2Ring_Matrix_Coloring >::color_step(\n Matrix_d &A, int seed, int max_color )\n{\n const int num_rows = A.get_num_rows();\n const int max_uncolored_rows = static_cast( this->m_uncolored_fraction * num_rows );\n const int WARP_SIZE = NNZ;\n const int CTA_SIZE = 256;\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int GRID_SIZE = std::min( 2048, (num_rows + NUM_WARPS_PER_CTA - 1) / NUM_WARPS_PER_CTA );\n const int NUM_WARPS_PER_CTA_1 = CTA_SIZE;\n const int GRID_SIZE_1 = std::min( 2048, (num_rows + NUM_WARPS_PER_CTA_1 - 1) / NUM_WARPS_PER_CTA_1 );\n device_vector_alloc finished_colors_flag(1);\n finished_colors_flag[0] = 0;\n int *finished_colors_flag_ptr = amgx::thrust::raw_pointer_cast( &finished_colors_flag.front() );\n device_vector_alloc gtlt_count;\n std::vector > max_id_per_level(this->m_coloring_level);\n std::vector > used_colors_per_level(this->m_coloring_level);\n\n if (this->m_coloring_level == 1)\n {\n gtlt_count.resize(num_rows);\n }\n else\n {\n for (int i = 1; i < this->m_coloring_level; i++)\n {\n max_id_per_level[i].resize(num_rows);\n used_colors_per_level[i].resize(num_rows);\n }\n }\n\n int prev_uncolored = num_rows;\n int currentSafeColor = 0;\n int iter = 0;\n int changes_seed = 0;\n int num_color_box = 64;\n\n for ( int num_uncolored = num_rows ; num_uncolored > max_uncolored_rows; iter++)\n {\n if (this->m_coloring_level == 1)\n {\n //parallel greedy with greater hash count metric instead of raw hash\n count_gtlt_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n seed\n , 0);\n cudaCheckError();\n color_kernel_greedy_gtlt <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n this->m_num_colors,\n this->m_weakness_bound, max_color,\n this->m_row_colors.raw(), 0, 0);\n cudaCheckError();\n }\n else\n {\n //Any ring parallel greedy with propagation\n for (int i = 0; i < this->m_coloring_level - 1; i++)\n {\n if (i == 0)\n find_max_neighbor_kernel_and_propagate_used_colors <<< GRID_SIZE, CTA_SIZE>>>(num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n 0, // in\n amgx::thrust::raw_pointer_cast(max_id_per_level[i + 1].data()), //out\n 0, // in\n amgx::thrust::raw_pointer_cast(used_colors_per_level[i + 1].data()), //out\n //amgx::thrust::raw_pointer_cast(interesting_in.data()),\n currentSafeColor,\n seed);\n else\n find_max_neighbor_kernel_and_propagate_used_colors <<< GRID_SIZE, CTA_SIZE>>>(num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n amgx::thrust::raw_pointer_cast(max_id_per_level[i ].data()), // in\n amgx::thrust::raw_pointer_cast(max_id_per_level[i + 1].data()), //out\n amgx::thrust::raw_pointer_cast(used_colors_per_level[i ].data()), // in\n amgx::thrust::raw_pointer_cast(used_colors_per_level[i + 1].data()), //out\n //amgx::thrust::raw_pointer_cast(interesting_in.data()),\n currentSafeColor,\n seed);\n\n cudaCheckError();\n }\n\n color_kernel_greedy_onlymax <<< GRID_SIZE, CTA_SIZE>>>( num_rows,\n A.row_offsets.raw(), A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( max_id_per_level[this->m_coloring_level - 1].data()), //in: max_id_tmp\n seed,\n this->m_row_colors.raw(),\n (unsigned long long *)amgx::thrust::raw_pointer_cast( used_colors_per_level[this->m_coloring_level - 1].data() ),\n finished_colors_flag_ptr,\n currentSafeColor);\n cudaCheckError();\n }\n\n int nnz_per_row = float(A.get_num_nz()) / A.get_num_rows();\n\n if (iter > std::min((int)std::pow(nnz_per_row, this->m_coloring_level), 63)) //estimate number of epochs required for coloring\n {\n num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_zero() );\n cudaCheckError();\n\n if (num_uncolored == prev_uncolored)\n {\n max_color++;\n\n if (this->m_coloring_level >= 2)\n {\n //printf(\"%d Fail with numcolor %d %d\\n\",seed,max_color,num_uncolored);\n if (changes_seed < 0)\n {\n seed++;\n changes_seed++;\n }\n else\n {\n changes_seed = 0;\n currentSafeColor += num_color_box;\n\n //reset buffers\n for (int i = 1; i < this->m_coloring_level; i++)\n {\n thrust_wrapper::fill(max_id_per_level[i].begin(), max_id_per_level[i] .end(), 0);\n thrust_wrapper::fill(used_colors_per_level[i].begin(), used_colors_per_level[i].end(), 0);\n }\n\n cudaCheckError();\n //printf(\"opening new color box %d\\n\",currentSafeColor);\n }\n }\n }\n\n prev_uncolored = num_uncolored;\n //printf(\"%d\\t%d\\t%d\\n\", iter, 0, num_uncolored);\n }\n }\n\n this->m_num_colors = thrust_wrapper::reduce( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, 0, amgx::thrust::maximum() ) + 1;\n cudaCheckError();\n}\n\n\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nGreedy_Min_Max_2Ring_Matrix_Coloring >::colorMatrix( Matrix_d &A )\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize( A.row_offsets.size() - 1 );\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n this->m_num_colors = 1;\n thrust_wrapper::fill( this->m_row_colors.begin(), this->m_row_colors.end(), 0 );\n cudaCheckError();\n float avg_nnz_per_row = A.get_num_nz() / float(A.get_num_rows());\n //Choose subwarp (group of threads processing a row) size\n {\n if ( avg_nnz_per_row < 4.0f )\n {\n color_step<1>(A, 0, 64);\n }\n else if ( avg_nnz_per_row < 8.0f )\n {\n color_step<1>(A, 0, 64);\n }\n else if (avg_nnz_per_row < 32)\n {\n color_step<4>(A, 0, 64);\n }\n else\n {\n color_step<32>(A, 0, 64);\n }\n }\n A.setView(oldView);\n}\n\n\n#define AMGX_CASE_LINE(CASE) template class Greedy_Min_Max_2Ring_Matrix_Coloring_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Greedy_Min_Max_2Ring_Matrix_Coloring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // end namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n// Pseudo-random number generator\nnamespace amgx\n{\n\nstatic __host__ __device__ __forceinline__ unsigned int hash_function(unsigned int a, unsigned int seed, unsigned int rows = 0)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n\n// ---------------------------\n// Kernels: Any ring parallel greedy with color buffers and max id propagation.\n// ---------------------------\n\n\n/*\n * This kernel computes the index of the node with maximum hash in a neighborhood and the colors used in a vertex neighborhood.\n * Since both operation are associative (max and or), it is possible to iterate the operation propagating the values across neigborhoods (first, second and so on rings).\n *\n * The N-th iteration computes, for each vertex:\n * 1. The ID of the uncolored vertex with ID of the vertex having max hash inside the N-ring neighborhood.\n * If all neighbors in the N-hood are colored, saves -1 (this allows work efficiency).\n * 2. The buffer of used colors in the n-ring neighborhood.\n *\n * And for that, it uses data from previous iteration.\n *\n * For avoiding the waste of work, once any n-ring neighborhood is all colored, it starts getting ignored:\n * no new MAX HASH can be found there and the set of used colors for that ring will not change.\n */\n\n\ntemplate< bool USE_EXT, int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid find_max_neighbor_kernel_and_propagate_used_colors( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const int *__restrict A_colors,\n\n int *level_n0_max_id_in, //for multiple iterations. load the ID of the neighbor with max hash in n-1 ring hood\n int *level_n1_out_max_id,//for storing the ID of the n-ring neighbor with max hash (in the n-ring hood)\n\n unsigned long long int *level_n0_used_colors_in, //for multiple iterations: load used color buffers from n-1 level\n unsigned long long int *level_n1_used_colors_out,//for storing the used color buffers of the n-level neighborhood of each vertex\n\n //int* interesting,\n const int currentSafeColor, //starts assigning colors from this. multiple of 64\n const int seed)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n const int n1_max_id = level_n1_out_max_id[row_id];\n\n if (n1_max_id == -1)\n {\n //-1 means that all the vertices in this n-ring\n //level neighborhood were colored during the previous coloring epochs.\n //Since nothing could have changed, just skip:\n //save work and use the color buffers computed in last step.\n continue;\n }\n\n //If the old max 'n1_max_id' is still uncolored, it still will be the max in this iteration.\n //Plus, nothing else would get colored in its n1-hood, since 'n1_max_id' is the max.\n if (A_colors[n1_max_id] <= 0)\n {\n continue;\n }\n\n int row_color = A_colors[row_id];\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n unsigned long long used_colors = 0ull;\n\n //For multiple iterations: load the used color buffer from previous iter.\n if (USE_EXT)\n {\n used_colors = level_n0_used_colors_in[row_id];\n }\n\n if (row_color - currentSafeColor > 0) //if the color belongs to the current buffer [currentSafeColor,currentSafeColor+64)\n {\n used_colors |= (1ull << (64 - row_color + currentSafeColor)); //add my color to buffer\n }\n\n // Hash my row id.\n int max_hash = hash_function( row_id, seed, 0 );\n int max_hash_id = row_id; //id of the neighbor with max hash, initialized with myself\n\n //If already colored, don't count me for finding max\n if (row_color != 0)\n {\n max_hash_id = -1; //-1 means no max\n max_hash = 0;//LONG_LONG_MIN;\n }\n\n //iterate the neighbors\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, seed, 0);\n // Get the color of the column.\n int col_color = -1;\n\n if ( row_it < row_end && col_id < A_num_rows)\n {\n col_color = A_colors[col_id];\n\n if ( col_color - currentSafeColor > 0 ) //if in the current buffer..\n {\n used_colors |= (1ull << (64 - col_color + currentSafeColor)); //..add it to buffer\n }\n\n if (USE_EXT)\n {\n //combine used_colors buffer with the ones computed in previous iterations\n used_colors |= level_n0_used_colors_in[col_id];\n col_id = level_n0_max_id_in[col_id];\n\n if (col_id < 0) //row and its k-1 level neigbhors all colored\n {\n col_hash = 0;//LONG_LONG_MIN;\n col_color = -1;\n }\n else\n {\n col_hash = hash_function(col_id, seed, 0); //col hash is the max hash of its neighbors in n-1 level\n col_color = 0; //there is an uncolored neighbor\n }\n }\n }\n\n //check if any neighbor has a greater max id.\n if (col_color == 0 && (max_hash_id < 0 || col_hash > max_hash || (col_hash == max_hash && col_id >= max_hash_id)))\n {\n max_hash_id = col_id;\n max_hash = col_hash;\n }\n }\n\n //BEGIN: Reduce used_colors/max_hashes amongst subwarps\n#pragma unroll\n\n for (int i = WARP_SIZE / 2; i >= 1; i /= 2)\n {\n int tmp = utils::shfl_xor(max_hash, i, WARP_SIZE); //ss_max_hash[threadIdx.x+i];\n int tmp_id = utils::shfl_xor(max_hash_id, i, WARP_SIZE);\n\n if (tmp_id >= 0 && (max_hash_id < 0 || tmp > max_hash || (tmp == max_hash && tmp_id >= max_hash_id)))\n {\n max_hash = tmp;\n max_hash_id = tmp_id;\n }\n\n //\n int tmp_hi = __double2hiint( __longlong_as_double( used_colors ) );\n int tmp_lo = __double2loint( __longlong_as_double( used_colors ) );\n tmp_hi = utils::shfl_xor(tmp_hi, i, WARP_SIZE);\n tmp_lo = utils::shfl_xor(tmp_lo, i, WARP_SIZE);\n long long tmpu = __double_as_longlong(__hiloint2double(tmp_hi, tmp_lo));\n used_colors |= tmpu;\n }\n\n //END: Reduce used_colors/max_hashes amongst subwarps\n\n //The subwarp leader stores the result.\n if ( lane_id == 0 )\n {\n //if max_hash_id=1, will skip work next time: every n-ring neighbor is colored\n level_n1_out_max_id[row_id] = max_hash_id;\n level_n1_used_colors_out[row_id] = used_colors;\n }\n }\n}\n\n\n/*\n * This kernel colors rows using a greedy strategy and the used colors buffer from propagate iterations.\n */\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid color_kernel_greedy_onlymax(\n const int A_num_rows, const int *A_rows, const int *A_cols,\n const int *max_hash_id_in_neighborhood,\n\n const int seed,\n int *A_colors,\n const unsigned long long *used_colors_in,\n int *finished_colors_flag,\n const int currentSafeColor)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n //int dice = hash_function(row_id, 19881988, 0);\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 ) // Already colored!!!\n {\n continue;\n }\n\n //int row_hash = hash_function(row_id,seed,0);\n // The number of vertices that are greater/smaller than me.\n int max_id = max_hash_id_in_neighborhood[row_id];\n\n if (max_id != row_id) //was not a max in its own neighborhood\n {\n continue;\n }\n\n int max_hash = hash_function(max_id, seed, 0);\n //A_colors[row_id] = 1;\n //continue;\n bool is_max_vertex = true;\n //bool all_neighbors_colored = true;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n unsigned long long used_colors = 0ull;\n used_colors |= used_colors_in[row_id];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n // Filter out used colors\n if ( row_it < row_end && col_id < A_num_rows)\n {\n //int col_color = A_colors[col_id];\n //if(col_color == 0) all_neighbors_colored = false;\n used_colors |= used_colors_in[col_id];\n int nn_max_id = max_hash_id_in_neighborhood[col_id];\n int nn_max = hash_function(nn_max_id, seed, 0);\n is_max_vertex &= (max_hash > nn_max || (max_hash == nn_max && row_id >= nn_max_id));\n }\n\n if (!is_max_vertex) { break; }\n }\n\n //reduce used colors bit by bit.\n#pragma unroll\n\n for (int i = WARP_SIZE / 2; i >= 1; i /= 2)\n {\n int tmp_hi = __double2hiint( __longlong_as_double( used_colors ) );\n int tmp_lo = __double2loint( __longlong_as_double( used_colors ) );\n tmp_hi = utils::shfl_xor(tmp_hi, i, WARP_SIZE);\n tmp_lo = utils::shfl_xor(tmp_lo, i, WARP_SIZE);\n long long tmp = __double_as_longlong(__hiloint2double(tmp_hi, tmp_lo));\n used_colors |= tmp;\n }\n\n int my_color_1 = 64 - utils::bfind( ~used_colors );\n\n if (__popc(used_colors) >= 64 || my_color_1 > 64 || my_color_1 <= 0)\n {\n if (lane_id == 0)\n {\n finished_colors_flag[0] = 1;\n }\n\n is_max_vertex = 0;\n }\n\n my_color_1 += currentSafeColor;\n\n if ( subwarp::all( is_max_vertex ) )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = my_color_1;\n }\n\n continue;\n }\n }\n}\n\n\n/*\n *\n * One ring parallel greedy with greater hash count metric instead of raw hash metric. Similar to MIN_MAX_2RING\n *\n * */\ntemplate< int CTA_SIZE, int WARP_SIZE, bool USE_EXTERNAL_METRIC >\n__global__\nvoid count_gtlt_kernel( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const int *__restrict A_colors,\n int *A_gtlt_count,\n const int seed,\n const int *metric_in)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 )\n {\n if ( lane_id == 0 )\n {\n A_gtlt_count[row_id] = -1;\n }\n\n continue;\n }\n\n // The number of vertices that are greater/smaller than me.\n int gt_count = 0, lt_count = 0;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n // Hash my row id.\n int row_cols = row_end - row_begin;\n int row_hash = hash_function( row_id, seed, row_cols );\n\n if (USE_EXTERNAL_METRIC)\n {\n row_hash = metric_in[row_id] >> 16; //gt count\n }\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n int col_cols = 1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n#if USE_ROWS\n col_cols = A_rows[col_id + 1] - A_rows[col_id];\n#endif\n }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, seed, col_cols);\n // Get the color of the column.\n int col_color = -1;\n\n if ( row_it < row_end && col_id < A_num_rows)\n {\n col_color = A_colors[col_id];\n\n if (USE_EXTERNAL_METRIC)\n {\n col_hash = metric_in[col_id] >> 16; //gt count\n }\n }\n\n // Threads determine if they are greater than the row hash.\n int gt_pred = col_color == 0 && col_hash > row_hash;\n int lt_pred = col_color == 0 && col_hash < row_hash;\n // Count greater/smaller neighbors.\n gt_count += __popc( subwarp::masked_ballot( gt_pred ) );//subwarp::masked_ballot\n lt_count += __popc( subwarp::masked_ballot( lt_pred ) );\n }\n\n // The warp leader stores the result.\n int my_gtlt_count = (gt_count << 16) | lt_count;\n\n if ( lane_id == 0 )\n {\n A_gtlt_count[row_id] = my_gtlt_count;\n }\n }\n}\ntemplate< int CTA_SIZE, int WARP_SIZE, bool LOAD_EXTERNAL_USED_COLORS >\n__global__\nvoid color_kernel_greedy_gtlt(\n const int A_num_rows, const int *A_rows, const int *A_cols,\n const int *A_gtlt_count, const int current_color,\n const int weakness_bound, const int max_color_to_use,\n int *A_colors,\n const int min_degree, const unsigned long long *used_colors_in)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n //int dice = hash_function(row_id, 19881988, 0);\n int row_color = A_colors[row_id];\n\n if ( row_color > 0 ) // Already colored!!!\n {\n continue;\n }\n\n if ( row_color < 0 )\n {\n A_colors[row_id] = 0; //discovered\n\n if (row_id != 0) { continue; } //process next time\n }\n\n // The number of vertices that are greater/smaller than me.\n int row_gtlt_count = A_gtlt_count[row_id];\n // Split gtlt_count into 2.\n int row_gt_count = row_gtlt_count >> 16;\n int row_lt_count = row_gtlt_count & 0xffff;\n // Do we skip it.\n int candidate = 1;\n // Predicates. Is a vertex min/max.\n int is_max_vertex = row_gt_count <= weakness_bound;\n int is_min_vertex = row_lt_count <= weakness_bound;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n int degree = row_end - row_begin;\n\n if (degree < min_degree && min_degree > 0)\n {\n continue;\n }\n\n unsigned long long used_colors = 0ull;\n\n if (LOAD_EXTERNAL_USED_COLORS)\n {\n used_colors |= used_colors_in[row_id];\n }\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n if ( row_it < row_end && col_id < A_num_rows )\n {\n int col_color = A_colors[col_id];\n\n if ( col_color > 0 )\n {\n used_colors |= (1ull << (64 - col_color));\n }\n }\n\n if (LOAD_EXTERNAL_USED_COLORS)\n {\n used_colors |= used_colors_in[col_id];\n }\n\n // Get the gt/lt count.\n int col_gtlt_count = -1;\n\n if ( row_it < row_end && col_id < A_num_rows )\n {\n col_gtlt_count = A_gtlt_count[col_id];\n }\n\n // Split gtlt_count into 2. col_gt/lt_count == -1 if already colored.\n int col_gt_count = col_gtlt_count >> 16;\n int col_lt_count = col_gtlt_count & 0xffff;\n\n // Threads determine if they are greater than the row hash.\n //if( col_gt_count != -1 )\n // is_max_vertex &= col_gt_count < row_gt_count || (col_gt_count == row_gt_count && col_id < row_id);\n //if( col_lt_count != -1 )\n // is_min_vertex &= col_lt_count /*>*/ < row_lt_count || (col_lt_count == row_lt_count && col_id > row_id);\n if ( col_gtlt_count != -1 )\n {\n is_max_vertex &= col_gt_count > row_gt_count || (col_gt_count == row_gt_count && col_id >= row_id);\n is_min_vertex &= col_lt_count > row_lt_count || (col_lt_count == row_lt_count && col_id >= row_id);\n }\n }\n\n //is_max_vertex = row_gt_count==0;\n //is_min_vertex = false;\n //reduce used colors bit by bit.\n#pragma unroll\n\n for (int i = WARP_SIZE / 2; i >= 1; i /= 2)\n {\n int tmp_hi = __double2hiint( __longlong_as_double( used_colors ) );\n int tmp_lo = __double2loint( __longlong_as_double( used_colors ) );\n tmp_hi = utils::shfl_xor(tmp_hi, i, WARP_SIZE);\n tmp_lo = utils::shfl_xor(tmp_lo, i, WARP_SIZE);\n long long tmp = __double_as_longlong(__hiloint2double(tmp_hi, tmp_lo));\n used_colors |= tmp;\n }\n\n int my_color_1 = 0;\n int my_color_2 = 0;\n int free_colors = __popc(used_colors);\n my_color_1 = 64 - utils::bfind( ~used_colors );\n used_colors |= (1ull << (64 - my_color_1));\n my_color_2 = 64 - utils::bfind( ~used_colors );\n free_colors = 0;\n\n if ( candidate && subwarp::all( is_max_vertex ) && my_color_1 < max_color_to_use)\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = my_color_1;\n }\n\n continue;\n }\n\n if ( candidate && subwarp::all( is_min_vertex ) && free_colors >= 2)\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = my_color_2;\n }\n\n continue;\n }\n }\n}\n\nstruct is_zero\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x == 0;\n }\n};\nstruct is_minus_one\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x == -1;\n }\n};\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate< class T_Config >\nGreedy_Min_Max_2Ring_Matrix_Coloring_Base::Greedy_Min_Max_2Ring_Matrix_Coloring_Base( AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n if ( this->m_coloring_level > 5 || this->m_coloring_level < 0)\n {\n FatalError( \"Not implemented for coloring_level != 1\", AMGX_ERR_NOT_SUPPORTED_TARGET );\n }\n\n if ( cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\") == 1 )\n {\n m_uncolored_fraction = 0.0;\n }\n else\n {\n m_uncolored_fraction = cfg.AMG_Config::template getParameter(\"max_uncolored_percentage\", cfg_scope);\n }\n\n m_weakness_bound = cfg.AMG_Config::template getParameter( \"weakness_bound\", cfg_scope );\n m_late_rejection = cfg.AMG_Config::template getParameter( \"late_rejection\", cfg_scope ) != 0;\n m_coloring_try_remove_last_color = cfg.AMG_Config::template getParameter( \"coloring_try_remove_last_colors\", cfg_scope );\n m_coloring_custom_arg = cfg.AMG_Config::template getParameter( \"coloring_custom_arg\", cfg_scope );\n}\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\ntemplate\nvoid\nGreedy_Min_Max_2Ring_Matrix_Coloring >::color_step(\n Matrix_d &A, int seed, int max_color )\n{\n const int num_rows = A.get_num_rows();\n const int max_uncolored_rows = static_cast( this->m_uncolored_fraction * num_rows );\n const int WARP_SIZE = NNZ;\n const int CTA_SIZE = 256;\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int GRID_SIZE = std::min( 2048, (num_rows + NUM_WARPS_PER_CTA - 1) / NUM_WARPS_PER_CTA );\n const int NUM_WARPS_PER_CTA_1 = CTA_SIZE;\n const int GRID_SIZE_1 = std::min( 2048, (num_rows + NUM_WARPS_PER_CTA_1 - 1) / NUM_WARPS_PER_CTA_1 );\n device_vector_alloc finished_colors_flag(1);\n finished_colors_flag[0] = 0;\n int *finished_colors_flag_ptr = amgx::thrust::raw_pointer_cast( &finished_colors_flag.front() );\n device_vector_alloc gtlt_count;\n std::vector > max_id_per_level(this->m_coloring_level);\n std::vector > used_colors_per_level(this->m_coloring_level);\n\n if (this->m_coloring_level == 1)\n {\n gtlt_count.resize(num_rows);\n }\n else\n {\n for (int i = 1; i < this->m_coloring_level; i++)\n {\n max_id_per_level[i].resize(num_rows);\n used_colors_per_level[i].resize(num_rows);\n }\n }\n\n int prev_uncolored = num_rows;\n int currentSafeColor = 0;\n int iter = 0;\n int changes_seed = 0;\n int num_color_box = 64;\n\n for ( int num_uncolored = num_rows ; num_uncolored > max_uncolored_rows; iter++)\n {\n if (this->m_coloring_level == 1)\n {\n //parallel greedy with greater hash count metric instead of raw hash\n count_gtlt_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n seed\n , 0);\n cudaCheckError();\n color_kernel_greedy_gtlt <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n this->m_num_colors,\n this->m_weakness_bound, max_color,\n this->m_row_colors.raw(), 0, 0);\n cudaCheckError();\n }\n else\n {\n //Any ring parallel greedy with propagation\n for (int i = 0; i < this->m_coloring_level - 1; i++)\n {\n if (i == 0)\n find_max_neighbor_kernel_and_propagate_used_colors <<< GRID_SIZE, CTA_SIZE>>>(num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n 0, // in\n amgx::thrust::raw_pointer_cast(max_id_per_level[i + 1].data()), //out\n 0, // in\n amgx::thrust::raw_pointer_cast(used_colors_per_level[i + 1].data()), //out\n //amgx::thrust::raw_pointer_cast(interesting_in.data()),\n currentSafeColor,\n seed);\n else\n find_max_neighbor_kernel_and_propagate_used_colors <<< GRID_SIZE, CTA_SIZE>>>(num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n amgx::thrust::raw_pointer_cast(max_id_per_level[i ].data()), // in\n amgx::thrust::raw_pointer_cast(max_id_per_level[i + 1].data()), //out\n amgx::thrust::raw_pointer_cast(used_colors_per_level[i ].data()), // in\n amgx::thrust::raw_pointer_cast(used_colors_per_level[i + 1].data()), //out\n //amgx::thrust::raw_pointer_cast(interesting_in.data()),\n currentSafeColor,\n seed);\n\n cudaCheckError();\n }\n\n color_kernel_greedy_onlymax <<< GRID_SIZE, CTA_SIZE>>>( num_rows,\n A.row_offsets.raw(), A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( max_id_per_level[this->m_coloring_level - 1].data()), //in: max_id_tmp\n seed,\n this->m_row_colors.raw(),\n (unsigned long long *)amgx::thrust::raw_pointer_cast( used_colors_per_level[this->m_coloring_level - 1].data() ),\n finished_colors_flag_ptr,\n currentSafeColor);\n cudaCheckError();\n }\n\n int nnz_per_row = float(A.get_num_nz()) / A.get_num_rows();\n\n if (iter > std::min((int)std::pow(nnz_per_row, this->m_coloring_level), 63)) //estimate number of epochs required for coloring\n {\n num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_zero() );\n cudaCheckError();\n\n if (num_uncolored == prev_uncolored)\n {\n max_color++;\n\n if (this->m_coloring_level >= 2)\n {\n //printf(\"%d Fail with numcolor %d %d\\n\",seed,max_color,num_uncolored);\n if (changes_seed < 0)\n {\n seed++;\n changes_seed++;\n }\n else\n {\n changes_seed = 0;\n currentSafeColor += num_color_box;\n\n //reset buffers\n for (int i = 1; i < this->m_coloring_level; i++)\n {\n thrust_wrapper::fill(max_id_per_level[i].begin(), max_id_per_level[i] .end(), 0);\n thrust_wrapper::fill(used_colors_per_level[i].begin(), used_colors_per_level[i].end(), 0);\n }\n\n cudaCheckError();\n //printf(\"opening new color box %d\\n\",currentSafeColor);\n }\n }\n }\n\n prev_uncolored = num_uncolored;\n //printf(\"%d\\t%d\\t%d\\n\", iter, 0, num_uncolored);\n }\n }\n\n this->m_num_colors = thrust_wrapper::reduce( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, 0, amgx::thrust::maximum() ) + 1;\n cudaCheckError();\n}\n\n\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nGreedy_Min_Max_2Ring_Matrix_Coloring >::colorMatrix( Matrix_d &A )\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize( A.row_offsets.size() - 1 );\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n this->m_num_colors = 1;\n thrust_wrapper::fill( this->m_row_colors.begin(), this->m_row_colors.end(), 0 );\n cudaCheckError();\n float avg_nnz_per_row = A.get_num_nz() / float(A.get_num_rows());\n //Choose subwarp (group of threads processing a row) size\n {\n if ( avg_nnz_per_row < 4.0f )\n {\n color_step<1>(A, 0, 64);\n }\n else if ( avg_nnz_per_row < 8.0f )\n {\n color_step<1>(A, 0, 64);\n }\n else if (avg_nnz_per_row < 32)\n {\n color_step<4>(A, 0, 64);\n }\n else\n {\n color_step<32>(A, 0, 64);\n }\n }\n A.setView(oldView);\n}\n\n\n#define AMGX_CASE_LINE(CASE) template class Greedy_Min_Max_2Ring_Matrix_Coloring_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Greedy_Min_Max_2Ring_Matrix_Coloring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // end namespace amgx\n", "source": "/kaggle/working/amgx/src/matrix_coloring/greedy_min_max_2ring.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include //std::pow\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#define GET_COLORS_THRUST 0\n#define GET_COLORS_TWOPASS 1\n#define GET_COLORS_ATOMIC 2\n#define GET_COLORS_GATHER 3\n#define GETMAX_THRUST 0\n#define GETMAX_TWOPASS 1\n\n//CONFIG\n\n#define GET_COLORS GET_COLORS_TWOPASS\n//#define GET_COLORS GET_COLORS_THRUST\n\n#define GETMAX_COLOR GETMAX_TWOPASS\n//#define GETMAX_COLOR GETMAX_THRUST\n\n//#define DO_SORT 1\n#define USE_GTLT 0\n#define DISCARD_COLORED (1 && !USE_GTLT)\n#define ONLY_FIRSTSTEP 0\n\nnamespace amgx\n{\n//normal hash function\n__host__ __device__ __forceinline__ unsigned int hash2(unsigned int a, const unsigned int seed)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n//rehashes a given hash (h), rotating its bits.\n__host__ __device__ __forceinline__ unsigned int rehash(const unsigned int &a, const unsigned int h, const unsigned int k, const unsigned int seed)\n{\n return (h << k) ^ (h >> k);\n}\n\n//compute hashes for a column, compare with the row's hashes and the updated bitfield is_max/is_min for the row.\ntemplate\n__device__ __forceinline__ void compute_hashes(const int &row_id, const int &col_id, const int &base_row_hash, const int &base_col_hash, const int &seed, unsigned long long int &is_min, unsigned long long int &is_max)\n{\n int col_hash_k = base_col_hash;\n int row_hash_k = base_row_hash;\n#pragma unroll\n\n for (int k = 0; k < HASHES; k++)\n {\n col_hash_k = rehash(col_id, base_col_hash, k, seed);\n row_hash_k = rehash(row_id, base_row_hash, k, seed);\n unsigned long long q = (~(1ull << (63 - k))); //bit of hash k in bitfields\n\n if (col_hash_k > row_hash_k || (col_hash_k == row_hash_k && row_id < col_id))\n {\n is_max &= q;\n }\n\n //unsigned long long x2 = col_hash_k < row_hash_k;\n if (col_hash_k < row_hash_k || (col_hash_k == row_hash_k && row_id > col_id))\n {\n is_min &= q;\n }\n }\n}\n\ntemplate\nstruct neigbor_visitor\n{\n\n template\n static __device__ __forceinline__ void visit(const int &row_id, const int &base_row_hash, const int &explore_row, const int *A_rows, const int *A_cols, const int *A_colors_in, const int lane_id, const int seed, unsigned long long int &is_min, unsigned long long int &is_max)\n {\n int row_begin = A_rows[explore_row ];\n int row_end = A_rows[explore_row + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n bool do_compute_hashes = true;\n col_id = A_cols[row_it];//__load_global(A_cols+row_it);\n#if DISCARD_COLORED\n\n if (DISCARD_COLORED_T)\n {\n int col_color = __ldg(A_colors_in + col_id);\n\n if (col_color != 0) { do_compute_hashes = false; }\n }\n\n#endif\n\n if (do_compute_hashes)\n {\n //#pragma unroll\n //for(int k=0; k<(HASHES/16>0)?HASHES/16:1; k++)\n //{\n int base_col_hash = hash2(col_id, seed/*+k*12358*/);\n compute_hashes(row_id, col_id, base_row_hash, base_col_hash, seed, is_min, is_max);\n //}\n }\n\n neigbor_visitor < COLORING_LEVEL - 1 >::template visit(row_id, base_row_hash, col_id, A_rows, A_cols, A_colors_in, lane_id, seed, is_min, is_max);\n }\n }\n\n};\ntemplate<>\nstruct neigbor_visitor<0>\n{\n\n template\n static __device__ __forceinline__ void visit(const int &row_id, const int &base_row_hash, const int &explore_row, const int *A_rows, const int *A_cols, const int *A_colors_in, const int lane_id, const int seed, unsigned long long int &is_min, unsigned long long int &is_max)\n {\n //is_min |= 0ull;\n }\n\n};\n\n\ntemplate< bool DISCARD_COLORED_T, int HASHES, int COLORING_LEVEL, int CTA_SIZE, int WARP_SIZE>\n__global__\nvoid fast_multihash_kernel(\n const int A_num_rows, const int *A_rows, const int *A_cols,\n int *A_colors, const int *A_colors_in,\n int c0, int seed,\n int *sets_per_block)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n int warp_count = 0;\n\n for ( ; utils::any(row_id < A_num_rows) ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = -1;\n int available_color = 0;\n\n if (row_id < A_num_rows)\n {\n#if DISCARD_COLORED\n\n if (DISCARD_COLORED_T)\n {\n row_color = A_colors_in[row_id];\n A_colors[row_id] = row_color;\n }\n else\n#endif\n {\n row_color = A_colors[row_id];\n }\n }\n\n if (row_color == 0)\n {\n unsigned long long int is_min = ~0ull;\n unsigned long long int is_max = ~0ull;\n int base_row_hash = hash2(row_id, seed);\n#if 1 //COLORING_LEVELR>1\n neigbor_visitor::template visit(row_id, base_row_hash, row_id, A_rows, A_cols, A_colors_in, lane_id, seed, is_min, is_max);\n#else\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n col_id = A_cols[row_it];//__load_global(A_cols+row_it);\n#if DISCARD_COLORED\n\n if (DISCARD_COLORED_T)\n {\n int col_color = __ldg(A_colors_in + col_id);\n\n if (col_color != 0) { continue; }\n }\n\n#endif\n int base_col_hash = hash2(col_id, seed);\n compute_hashes(row_id, col_id, base_row_hash, base_col_hash, seed, is_min, is_max);\n }\n\n#endif\n\n if (WARP_SIZE > 1)\n {\n using namespace amgx::strided_reduction;\n amgx::strided_reduction::op_and op;\n warp_reduce<1, CTA_SIZE, WARP_SIZE>(is_max, op);\n warp_reduce<1, CTA_SIZE, WARP_SIZE>(is_min, op);\n }\n\n#if 1\n int cmax = 63 - utils::bfind(is_max);\n int cmin = 63 - utils::bfind(is_min);\n\n if (cmin < HASHES)\n {\n available_color = cmin * 2 + 2;\n }\n\n if (cmax < HASHES)\n {\n available_color = cmax * 2 + 1;\n }\n\n#else\n#pragma unroll\n\n for (int k = 0; k < HASHES; k++)\n {\n if (is_max & (1ull << 63 - k))\n {\n available_color = 2 * k + 1;\n break;\n }\n else if (is_min & (1ull << 63 - k))\n {\n available_color = 2 * k + 2;\n break;\n }\n }\n\n#endif\n }\n\n bool has_set_color = (lane_id == 0) && (available_color != 0);\n\n if (has_set_color)\n {\n A_colors[row_id] = available_color + c0;\n }\n\n warp_count += amgx::strided_reduction::warp_binary_count(has_set_color);\n }\n\n#if GET_COLORS == GET_COLORS_ATOMIC\n amgx::strided_reduction::block_count_atomic<1, CTA_SIZE, 32, int>(warp_count, sets_per_block);\n //__threadfence_system();\n#elif GET_COLORS == GET_COLORS_TWOPASS\n amgx::strided_reduction::block_count<1, CTA_SIZE, 32, int>(warp_count, sets_per_block);\n#elif GET_COLORS == GET_COLORS_GATHER\n amgx::strided_reduction::block_count<1, CTA_SIZE, 32, int>(warp_count, sets_per_block);\n amgx::strided_reduction::block_count_gatherer global_gather(sets_per_block);\n#endif\n}\n\n////////////////////////////\n\ntemplate< int HASHES, int CTA_SIZE, int WARP_SIZE>\n__global__\nvoid fast_multihash_kernel_gtlt_kernel(\n const int A_num_rows, const int *A_rows, const int *A_cols,\n int *A_colors, unsigned long long int *gtlt_out,\n int seed)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n //int warp_count = 0;\n\n for ( ; utils::any(row_id < A_num_rows) ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = -1;\n\n if (row_id < A_num_rows)\n {\n row_color = A_colors[row_id];\n }\n\n if (row_color == 0)\n {\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n int base_row_hash = hash2(row_id, seed);\n unsigned char count_gt[HASHES];\n unsigned char count_lt[HASHES];\n#pragma unroll\n\n for (int k = 0; k < HASHES; k++)\n {\n count_gt[k] = 0;\n count_lt[k] = 0;\n }\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n col_id = A_cols[row_it];\n int col_color = A_colors[col_id];\n\n if (col_color != 0) { continue; }\n\n int base_col_hash = hash2(col_id, seed);\n int col_hash_k = base_col_hash;\n int row_hash_k = base_row_hash;\n#pragma unroll\n\n for (int k = 0; k < HASHES; k++)\n {\n col_hash_k = rehash(col_id, base_col_hash, k, seed);\n row_hash_k = rehash(row_id, base_row_hash, k, seed);\n\n if (col_hash_k > row_hash_k || (col_hash_k == row_hash_k && row_id < col_id))\n {\n if (count_gt[k] < 255) //avoid overflow\n {\n ++count_gt[k];\n }\n }\n\n if (col_hash_k < row_hash_k || (col_hash_k == row_hash_k && row_id > col_id))\n {\n if (count_lt[k] < 255)\n {\n ++count_lt[k];\n }\n }\n }\n }\n\n //compress count_gt,lt in 2 bits and pack in an integer\n unsigned long long int packed_gtlt = 0;\n#pragma unroll\n\n for (int k = 0; k < HASHES; k++)\n {\n unsigned long long int max3_gt = count_gt[k];\n unsigned long long int max3_lt = count_lt[k];\n\n if (max3_gt > 3) { max3_gt = 3; } //TODO use only bits above\n\n if (max3_lt > 3) { max3_lt = 3; }\n\n //TODO: use BFI\n packed_gtlt |= (max3_gt << (k * 2));\n packed_gtlt |= (max3_lt << (k * 2 + 32));\n }\n\n if (lane_id == 0)\n {\n gtlt_out[row_id] = packed_gtlt;\n }\n }\n }\n}\n\ntemplate< int HASHES, int CTA_SIZE, int WARP_SIZE>\n__global__\nvoid fast_multihash_kernel_gtlt_assign_kernel(\n const int A_num_rows, const int *A_rows, const int *A_cols,\n int *A_colors, unsigned long long int *gtlt_in,\n int c0)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n for ( ; utils::any(row_id < A_num_rows) ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = -1;\n unsigned long long int row_gtlt = 0ull;\n\n if (row_id < A_num_rows)\n {\n row_color = A_colors[row_id];\n row_gtlt = gtlt_in[row_id];\n }\n\n int available_color = 0;\n\n if (row_color == 0)\n {\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n unsigned long long int is_min = ~0ull;\n unsigned long long int is_max = ~0ull;\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n unsigned long long int col_gtlt = 0ull;\n\n if (row_it < row_end)\n {\n col_id = A_cols[row_it];\n col_gtlt = gtlt_in[col_id];\n#pragma unroll\n\n for (int k = 0; k < HASHES; k++)\n {\n int row_gt_k = (row_gtlt >> (k * 2 + 0)) & 3ull;\n int row_lt_k = (row_gtlt >> (k * 2 + 32)) & 3ull;\n int col_gt_k = (col_gtlt >> (k * 2 + 0)) & 3ull;\n int col_lt_k = (col_gtlt >> (k * 2 + 32)) & 3ull;\n unsigned long long q = (~(1ull << (63 - k)));\n\n if ((row_gt_k > col_gt_k) || (col_gt_k == row_gt_k && row_id < col_id))\n {\n is_max &= q; //xor with k bit: if was 1, becomes 0, if was 0 stays 0\n }\n\n if ((row_lt_k > col_lt_k) || (col_lt_k == row_lt_k && row_id > col_id))\n {\n is_min &= q;\n }\n }\n }\n }\n\n int cmax = 63 - utils::bfind(is_max);\n int cmin = 63 - utils::bfind(is_min);\n\n //if(cmin < HASHES)\n // available_color = cmin*2+2;\n if (cmax < HASHES)\n {\n available_color = cmax * 2 + 1;\n }\n }\n\n bool has_set_color = (lane_id == 0) && (available_color != 0);\n\n if (has_set_color)\n {\n A_colors[row_id] = available_color + c0;\n }\n }\n}\n\n////////////////////////////\n\ntemplate< bool SORTED_ROWS, int COLORING_LEVEL, int CTA_SIZE, int WARP_SIZE>\n__global__\nvoid recolor_greedy_kernel(\n const int A_num_rows, const int *__restrict A_rows,\n const int *__restrict A_cols,\n int color_to_reassign,\n int *__restrict A_colors,\n int *sorted_rows_by_color, int offset_start, int offset_end\n , int *maxs_per_block\n)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n int job_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n int last_job = A_num_rows;\n\n if (SORTED_ROWS)\n {\n job_id = job_id + offset_start;\n last_job = offset_end;\n }\n\n int thread_max = 0;\n\n //const int NUM_THREADS_PER_GRID = gridDim.x * CTA_SIZE;\n //int row_id = blockIdx.x * CTA_SIZE + threadIdx.x;\n\n // Iterate over rows.\n for ( ; utils::any(job_id < last_job) ; job_id += NUM_WARPS_PER_GRID )\n {\n int row_color = -1;\n int row_id = -1;\n\n if (job_id < last_job)\n {\n if (SORTED_ROWS)\n {\n row_id = sorted_rows_by_color[job_id];\n }\n else\n {\n row_id = job_id;\n }\n\n if (row_id < A_num_rows)\n {\n row_color = A_colors[row_id];\n }\n }\n\n used_color_structure_64_bit used_colors;\n used_color_structure_64_bit::aux_tmp aux;\n //used_color_structure_64_bit_colorbox<5> used_colors; //still buggy\n\n if (row_color == color_to_reassign)\n {\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n //used_colors.use_color(color_to_reassign);\n //int num_greater = 0;\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n\n if (row_it < row_end)\n {\n col_id = A_cols[row_it];\n }\n\n if (row_it < row_end)// && col_id < A_num_rows)\n {\n int col_color = A_colors[col_id];\n\n //if(col_color>40) printf(\"%d\\n\",col_color);\n if (col_color != 0)\n {\n used_colors.use_color(col_color, aux);\n }\n }\n\n if (COLORING_LEVEL == 2)\n {\n int row_begin = A_rows[col_id ];\n int row_end = A_rows[col_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n\n if (row_it < row_end)\n {\n col_id = A_cols[row_it];\n }\n\n if (row_it < row_end)// && col_id < A_num_rows)\n {\n int col_color = A_colors[col_id];\n\n //if(col_color>40) printf(\"%d\\n\",col_color);\n if (col_color != 0)\n {\n used_colors.use_color(col_color, aux);\n }\n }\n }\n }\n }\n }\n\n used_colors.sync_subwarp(lane_id);\n\n /*if(row_color != color_to_reassign && row_id != -1)\n {\n printf(\"ERR %d %d!\\n\", row_color, color_to_reassign);\n }*/\n if (row_color == color_to_reassign && lane_id == 0)\n {\n int new_color = used_colors.first_available_color(aux);\n\n if (new_color != row_color && new_color > 0)\n {\n //printf(\"changing %d from %d to %d\\n\",row_id,row_color,new_color);\n A_colors[row_id] = new_color;\n }\n\n#if GETMAX_COLOR == GETMAX_TWOPASS\n\n if (new_color > thread_max)\n {\n thread_max = new_color;\n }\n\n#endif\n /*if(new_color >= 7)\n {\n printf(\"SEVEN %d %d\\n\", new_color, thread_max);\n }*/\n }\n }\n\n#if GETMAX_COLOR == GETMAX_TWOPASS\n using namespace amgx::strided_reduction;\n amgx::strided_reduction::op_max op;\n warp_reduce<1, CTA_SIZE, 32>(thread_max, op);\n block_reduce<1, CTA_SIZE, 32, true>(thread_max, maxs_per_block, op);\n#endif\n}\n\n////////////////////////////\n\n\n\ntemplate< bool SORTED_ROWS, int COLORING_LEVEL, int CTA_SIZE, int WARP_SIZE>\n__global__\nvoid recolor_greedy_kernel_anyring(\n const int A_num_rows, const int *A_rows, const int *A_cols,\n int color_to_reassign, int *A_colors,\n int *sorted_rows_by_color, int offset_start, int offset_end, int *maxs_per_block,\n\n used_color_structure_64_bit *used_colors_inout\n)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n int job_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n int last_job = A_num_rows;\n\n if (SORTED_ROWS)\n {\n job_id = job_id + offset_start;\n last_job = offset_end;\n }\n\n int thread_max = 0;\n\n //const int NUM_THREADS_PER_GRID = gridDim.x * CTA_SIZE;\n //int row_id = blockIdx.x * CTA_SIZE + threadIdx.x;\n\n // Iterate over rows.\n for ( ; utils::any(job_id < last_job) ; job_id += NUM_WARPS_PER_GRID )\n {\n int row_color = -1;\n int row_id = -1;\n\n if (job_id < last_job)\n {\n if (SORTED_ROWS)\n {\n row_id = sorted_rows_by_color[job_id];\n }\n else\n {\n row_id = job_id;\n }\n\n if (row_id < A_num_rows)\n {\n row_color = A_colors[row_id];\n }\n }\n\n used_color_structure_64_bit used_colors;\n used_color_structure_64_bit::aux_tmp aux;\n used_colors = used_colors_inout[row_id];\n //used_color_structure_64_bit_colorbox<5> used_colors; //still buggy\n\n if (row_color == color_to_reassign)\n {\n int new_color = used_colors.first_available_color(aux);\n\n if (lane_id == 0)\n {\n if (new_color != row_color && new_color > 0)\n {\n //printf(\"changing %d from %d to %d\\n\",row_id,row_color,new_color);\n A_colors[row_id] = new_color;\n }\n\n#if GETMAX_COLOR == GETMAX_TWOPASS\n\n if (new_color > thread_max)\n {\n thread_max = new_color;\n }\n\n#endif\n }\n\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n\n if (row_it < row_end)\n {\n col_id = A_cols[row_it];\n }\n\n if (row_it < row_end)// && col_id < A_num_rows)\n {\n int col_color = A_colors[col_id];\n\n if (col_color == 0)\n {\n used_color_structure_64_bit used_colors_col = used_colors_inout[col_id];\n used_colors_col.use_color(new_color, aux);\n used_colors_inout[col_id] = used_colors_col;\n }\n\n if (COLORING_LEVEL == 2)\n {\n int row_begin = A_rows[col_id ];\n int row_end = A_rows[col_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n\n if (row_it < row_end)\n {\n col_id = A_cols[row_it];\n }\n\n if (row_it < row_end)// && col_id < A_num_rows)\n {\n int col_color = A_colors[col_id];\n\n if (col_color == 0)\n {\n used_color_structure_64_bit used_colors_col = used_colors_inout[col_id];\n used_colors_col.use_color(new_color, aux);\n used_colors_inout[col_id] = used_colors_col;\n }\n }\n }\n }\n }\n }\n }\n }\n\n#if GETMAX_COLOR == GETMAX_TWOPASS\n using namespace amgx::strided_reduction;\n amgx::strided_reduction::op_max op;\n warp_reduce<1, CTA_SIZE, 32>(thread_max, op);\n block_reduce<1, CTA_SIZE, 32, true>(thread_max, maxs_per_block, op);\n#endif\n}\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate< class T_Config >\nGreedy_Recolor_MatrixColoring_Base::Greedy_Recolor_MatrixColoring_Base( AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n first_pass_config = cfg;\n first_pass_config_scope = cfg_scope;\n //MULTI_HASH: it would be faster but it leaves some uncolored vertices: uncolored vertices cannot be processed in parallel\n first_pass_config.setParameter(\"matrix_coloring_scheme\", std::string(\"MULTI_HASH\"), first_pass_config_scope);\n first_pass_config.setParameter(\"max_uncolored_percentage\", 0.0, first_pass_config_scope);\n m_coloring_custom_arg = cfg.AMG_Config::template getParameter( \"coloring_custom_arg\", cfg_scope );\n m_coloring_try_remove_last_color_ = cfg.AMG_Config::template getParameter( \"coloring_try_remove_last_colors\", cfg_scope );\n}\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nGreedy_Recolor_MatrixColoring >::colorMatrix( Matrix_d &A )\n{\n const int num_rows = A.get_num_rows();\n const int num_nz = A.get_num_nz();\n int avg_nz = num_nz / num_rows;\n //printf(\"AVG %d \",avg_nz);\n //P(single) = 1-((avg-1)/avg)^(K)\n //P(all) = P(single)^(N)\n // = (1-((avg-1)/avg)^(K))^(N) == 0.99\n // = (1-((avg-1)/avg)^(K)) == Pow(0.99,1/N)\n //P(single == 0.99) K = Log(1-0.99)/Log(1/avg)\n //P(all == 0.99) K = Log(1-Pow(0.99,1/N))/Log(1/avg)\n //99% probability\n //table here http://www.wolframalpha.com/input/?i=Log%281-0.99%29%2FLog%28%28a-1%29%2Fa%29+for+a+in+[2%2C10]\n avg_nz = pow(static_cast(avg_nz), static_cast(this->m_coloring_level));\n\n if (avg_nz > 16) { avg_nz = 16; }\n\n if (this->m_coloring_custom_arg != \"\")\n {\n int x = atoi(this->m_coloring_custom_arg.c_str());\n\n if (x > 0)\n {\n avg_nz = x;\n }\n }\n\n#define CASENZ(X,ORDER,DO_SORT)\\\n else if(avg_nz <= X)\\\n {\\\n color_matrix_specialized_per_numhash(A);\\\n }\n#define CASEORDER(ORDER,DO_SORT)\\\n if(0)\\\n {\\\n\\\n }\\\n CASENZ(2,ORDER,DO_SORT)\\\n CASENZ(4,ORDER,DO_SORT)\\\n CASENZ(8,ORDER,DO_SORT)\\\n CASENZ(12,ORDER,DO_SORT)\\\n CASENZ(16,ORDER,DO_SORT)\\\n CASENZ(20,ORDER,DO_SORT)\\\n CASENZ(24,ORDER,DO_SORT)\\\n CASENZ(32,ORDER,DO_SORT)\\\n CASENZ(48,ORDER,DO_SORT)\\\n else\\\n {\\\n color_matrix_specialized_per_numhash<16,ORDER,DO_SORT>(A);\\\n }\n\n if (0)\n {\n }\n else if (this->m_coloring_level == 1)\n {\n CASEORDER(1, 0);\n }\n else if (this->m_coloring_level == 2)\n {\n CASEORDER(2, 1);\n }\n else if (this->m_coloring_level == 3)\n {\n CASEORDER(3, 1);\n }\n else if (this->m_coloring_level == 4)\n {\n CASEORDER(4, 1);\n }\n else if (this->m_coloring_level == 5)\n {\n CASEORDER(5, 1);\n }\n}\n\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\ntemplate\nvoid\nGreedy_Recolor_MatrixColoring >::color_matrix_specialized_per_numhash( Matrix_d &A )\n{\n const int num_rows = A.get_num_rows();\n const int n_blocks = 26 * 4; //TODO: choose optimal blocks number\n const int block_size = 256;\n ViewType oldView = A.currentView();\n#if 1 //use fast multihash above as precolorer\n this->m_row_colors.resize( A.row_offsets.size() - 1, 0 );\n typedef typename Matrix > ::IVector IVector_h;\n IVector sorted_rows_by_color;\n IVector_h offsets_rows_per_color;\n {\n device_vector_alloc sets_per_block(n_blocks + 1);\n int *sets_per_block_ptr = amgx::thrust::raw_pointer_cast(sets_per_block.data());\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n int num_uncolored = num_rows;\n int i = 0;\n#if DISCARD_COLORED\n device_vector_alloc color_in(num_rows);\n#else\n device_vector_alloc color_in;\n#endif\n#if USE_GTLT\n device_vector_alloc gtlt(num_rows);\n const int K = K1 * 0.5;\n#else\n const int K = K1;\n#endif\n int prev_num_uncolored = num_uncolored;\n unsigned int seed = 0;\n\n while (num_uncolored > 0)\n {\n //int mx = amgx::thrust::reduce( this->m_row_colors.begin(), this->m_row_colors.begin()+num_rows, 0, amgx::thrust::maximum() );\n //if(mx > 2*K*i) printf(\"ERR UX %d %d %d\\n\", num_uncolored, mx, 2*K*i);\n#if USE_GTLT\n fast_multihash_kernel_gtlt_kernel <<< n_blocks, block_size>>>\n (num_rows, A.row_offsets.raw(), A.col_indices.raw(),\n this->m_row_colors.raw(), amgx::thrust::raw_pointer_cast(gtlt.data()),\n seed);\n cudaCheckError();\n fast_multihash_kernel_gtlt_assign_kernel <<< n_blocks, block_size>>>\n (num_rows, A.row_offsets.raw(), A.col_indices.raw(),\n this->m_row_colors.raw(), amgx::thrust::raw_pointer_cast(gtlt.data()),\n 2 * K * i);\n cudaCheckError();\n#else\n\n if (num_uncolored == num_rows || (!DISCARD_COLORED))\n {\n fast_multihash_kernel <<< n_blocks, block_size>>>(num_rows, A.row_offsets.raw(), A.col_indices.raw(),\n this->m_row_colors.raw(), amgx::thrust::raw_pointer_cast(color_in.data()),\n 2 * K * i, seed,\n sets_per_block_ptr);\n cudaCheckError();\n }\n else\n {\n fast_multihash_kernel <<< n_blocks, block_size>>>(num_rows, A.row_offsets.raw(), A.col_indices.raw(),\n this->m_row_colors.raw(), amgx::thrust::raw_pointer_cast(color_in.data()),\n 2 * K * i, seed,\n sets_per_block_ptr);\n }\n\n cudaCheckError();\n#endif\n seed = hash2(seed, 0);\n#if GET_COLORS == GET_COLORS_THRUST || USE_GTLT\n num_uncolored = amgx::thrust::count( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, 0);\n cudaCheckError();\n#elif GET_COLORS == GET_COLORS_TWOPASS\n num_uncolored -= amgx::strided_reduction::count_block_results_1(n_blocks, sets_per_block_ptr, amgx::strided_reduction::op_sum());\n#elif (GET_COLORS == GET_COLORS_ATOMIC) || (GET_COLORS == GET_COLORS_GATHER)\n num_uncolored -= sets_per_block[0];\n#endif\n //printf(\"%d -> %d\\n\", prev_num_uncolored,num_uncolored,prev_num_uncolored-num_uncolored);\n i++;\n\n if (prev_num_uncolored == num_uncolored)\n {\n //printf(\"ERR %d %d %d\\n\", prev_num_uncolored,num_uncolored,num_rows);\n break;\n }\n\n prev_num_uncolored = num_uncolored;\n#if DISCARD_COLORED\n color_in.swap(this->m_row_colors);\n#endif\n /*if (iteration++%4==0) {\n cudaEventRecord(throttle_event);\n cudaCheckError();\n } else {\n cudaEventSynchronize(throttle_event);\n cudaCheckError();\n }*/\n }\n\n#if DISCARD_COLORED\n color_in.swap(this->m_row_colors);\n#endif\n int max_color = i * K * 2; //available_color+c0 = 2*K*i + 2*(K-1)+2=2K(i+1), but i++\n //max_color = amgx::thrust::reduce( this->m_row_colors.begin(), this->m_row_colors.begin()+num_rows, 0, amgx::thrust::maximum() );\n //printf(\"CC %d<%d\\n\",max_color,i*K*2+2);\n this->m_num_colors = max_color + 1;\n#if ONLY_FIRSTSTEP\n\n if (1)\n#else\n if (this->m_coloring_custom_arg == \"FIRSTSTEP\")\n#endif\n {\n A.setView(oldView);\n return;\n }\n\n if (DO_SORT)\n {\n //return;\n //this->createColorArrays(A)\n IVector row_colors(this->m_row_colors); //useless copy\n sorted_rows_by_color.resize(num_rows);\n //this->m_sorted_rows_by_color.resize(num_rows);\n thrust_wrapper::sequence(sorted_rows_by_color.begin(), sorted_rows_by_color.end()); //useless sequence\n thrust_wrapper::sort_by_key(row_colors.begin(), row_colors.begin() + num_rows, sorted_rows_by_color.begin()); //useless read from sequence\n cudaCheckError();\n IVector offsets_rows_per_color_d(this->m_num_colors + 1);\n amgx::thrust::lower_bound(row_colors.begin(),\n row_colors.begin() + num_rows,\n amgx::thrust::counting_iterator(0),\n amgx::thrust::counting_iterator(offsets_rows_per_color_d.size()),\n offsets_rows_per_color_d.begin());\n // Copy from device to host\n offsets_rows_per_color = offsets_rows_per_color_d;\n cudaCheckError();\n /*for(int i=0;im_halo_coloring == SYNC_COLORS) A.setView(ALL);\n else A.setViewExterior();*/\n }\n }\n#else\n //1st pass: color with a fast method\n MatrixColoring *first_pass_matrix_coloring = MatrixColoringFactory::allocate(this->first_pass_config, this->first_pass_config_scope);\n first_pass_matrix_coloring->colorMatrix(A);\n#if DO_SORT\n first_pass_matrix_coloring->createColorArrays(A);\n#endif\n this->m_row_colors = first_pass_matrix_coloring->getRowColors();\n this->m_num_colors = first_pass_matrix_coloring->getNumColors();\n sorted_rows_by_color = first_pass_matrix_coloring->getSortedRowsByColor();\n offsets_rows_per_color = first_pass_matrix_coloring->getOffsetsRowsPerColor();\n delete first_pass_matrix_coloring;\n //ViewType oldView = A.currentView();\n //if (this->m_halo_coloring == SYNC_COLORS) A.setView(ALL);\n //else A.setViewExterior();\n#endif\n cudaCheckError();\n //2nd pass: optimize using parallel greedy for each color.\n //If 1st pass coloring is valid all elements of same color can be processed in parallel by definition\n#if GETMAX_COLOR == GETMAX_TWOPASS\n device_vector_alloc max_color_per_block(n_blocks + 1, 1);\n int *max_color_per_block_ptr = amgx::thrust::raw_pointer_cast(max_color_per_block.data());\n#else\n int *max_color_per_block_ptr = 0;\n#endif\n float avg_nnz_per_row = A.get_num_nz() / float(A.get_num_rows());\n int WARP_SIZE = 2;\n\n if ( avg_nnz_per_row < 4.0f )\n {\n WARP_SIZE = 2;\n }\n else if ( avg_nnz_per_row < 8.0f )\n {\n WARP_SIZE = 4;\n }\n else if ( avg_nnz_per_row < 32.0f )\n {\n WARP_SIZE = 8;\n }\n else\n {\n WARP_SIZE = 32;\n }\n\n if (this->m_coloring_level >= 2)\n {\n WARP_SIZE = 1;\n }\n\n //;\n //int max_colorx = amgx::thrust::reduce( this->m_row_colors.begin(), this->m_row_colors.begin()+num_rows, 0, amgx::thrust::maximum() );\n //printf(\"MAX %d\",max_colorx);\n#define RECOLOR_CALL_CASE(__WSIZE) if(WARP_SIZE==__WSIZE){ \\\n recolor_greedy_kernel<<>>(\\\n num_rows,\\\n A.row_offsets.raw(),\\\n A.col_indices.raw(),\\\n color, this->m_row_colors.raw(),\\\n sorted_rows_by_color.raw(),start,end,\\\n max_color_per_block_ptr\\\n ); cudaCheckError(); }\n cudaCheckError();\n#if 1 //straight order: generates a color histogram similar to parallel greedy\n //start from color=2: 1st color would be assigned to 1, no change\n\n //iteration = 0;\n\n for (int color = 2; color <= this->m_num_colors; color++)\n {\n const int GRID_SIZE = n_blocks;//\n int start, end;\n\n if (DO_SORT)\n {\n start = offsets_rows_per_color[color];\n end = num_rows;\n\n if (color < this->m_num_colors) { end = offsets_rows_per_color[color + 1]; }\n\n if (end >= num_rows) { end = num_rows; }\n\n //printf(\"F %d %d %d < num_rows: %d\\n\",color,start,end,num_rows,(int)end<=start);\n if (end <= start) { continue; }\n\n //int NUM_WARPS_PER_CTA = block_size/WARP_SIZE;\n //const int GRID_SIZE = std::min( 2048, ((end-start) + NUM_WARPS_PER_CTA-1) / NUM_WARPS_PER_CTA );\n //printf(\"%d - %d > %d %d %d\\n\",end-start,color,start,end,num_rows);\n }\n else\n {\n start = 0;\n end = num_rows;\n }\n\n //printf(\"%d : %d %d\\n\", color, start, end);\n RECOLOR_CALL_CASE(1);\n RECOLOR_CALL_CASE(2);\n RECOLOR_CALL_CASE(4);\n RECOLOR_CALL_CASE(8);\n RECOLOR_CALL_CASE(16);\n RECOLOR_CALL_CASE(32);\n cudaCheckError();\n /*if (iteration++%4==0) {\n cudaEventRecord(throttle_event);\n cudaCheckError();\n } else {\n cudaEventSynchronize(throttle_event);\n cudaCheckError();\n }*/\n }\n\n#else //reverse order: generates a different kind of histogram\n\n for (int color = this->m_num_colors; color >= 0; color--)\n {\n RECOLOR_CALL_CASE(1);\n RECOLOR_CALL_CASE(2);\n RECOLOR_CALL_CASE(4);\n RECOLOR_CALL_CASE(8);\n RECOLOR_CALL_CASE(16);\n RECOLOR_CALL_CASE(32);\n }\n\n#endif\n cudaCheckError();\n#if GETMAX_COLOR == GETMAX_TWOPASS\n int max_color = amgx::strided_reduction::count_block_results_1(n_blocks, max_color_per_block_ptr, amgx::strided_reduction::op_max());\n /*int max_colorg = amgx::thrust::reduce( this->m_row_colors.begin(), this->m_row_colors.begin()+num_rows, 0, amgx::thrust::maximum() );\n\n max_color=0;\n for(int i=0;i max_color) max_color = max_color_per_block[i];\n }\n printf(\"MAXCOLOR ref=%d %d\\n\",max_colorg,max_color);\n if(max_colorg!=max_color)exit(1);*/\n#else\n int max_color = thrust_wrapper::reduce( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, 0, amgx::thrust::maximum() );\n#endif\n //printf(\"MAXCOLOR %d %d\\n\",max_color_gold,max_color);\n cudaCheckError();\n this->m_num_colors = max_color + 1;\n A.setView(oldView);\n}\n\n\n\n#define AMGX_CASE_LINE(CASE) template class Greedy_Recolor_MatrixColoring_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Greedy_Recolor_MatrixColoring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // end namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include //std::pow\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#define GET_COLORS_THRUST 0\n#define GET_COLORS_TWOPASS 1\n#define GET_COLORS_ATOMIC 2\n#define GET_COLORS_GATHER 3\n#define GETMAX_THRUST 0\n#define GETMAX_TWOPASS 1\n\n//CONFIG\n\n#define GET_COLORS GET_COLORS_TWOPASS\n//#define GET_COLORS GET_COLORS_THRUST\n\n#define GETMAX_COLOR GETMAX_TWOPASS\n//#define GETMAX_COLOR GETMAX_THRUST\n\n//#define DO_SORT 1\n#define USE_GTLT 0\n#define DISCARD_COLORED (1 && !USE_GTLT)\n#define ONLY_FIRSTSTEP 0\n\nnamespace amgx\n{\n//normal hash function\n__host__ __device__ __forceinline__ unsigned int hash2(unsigned int a, const unsigned int seed)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n//rehashes a given hash (h), rotating its bits.\n__host__ __device__ __forceinline__ unsigned int rehash(const unsigned int &a, const unsigned int h, const unsigned int k, const unsigned int seed)\n{\n return (h << k) ^ (h >> k);\n}\n\n//compute hashes for a column, compare with the row's hashes and the updated bitfield is_max/is_min for the row.\ntemplate\n__device__ __forceinline__ void compute_hashes(const int &row_id, const int &col_id, const int &base_row_hash, const int &base_col_hash, const int &seed, unsigned long long int &is_min, unsigned long long int &is_max)\n{\n int col_hash_k = base_col_hash;\n int row_hash_k = base_row_hash;\n#pragma unroll\n\n for (int k = 0; k < HASHES; k++)\n {\n col_hash_k = rehash(col_id, base_col_hash, k, seed);\n row_hash_k = rehash(row_id, base_row_hash, k, seed);\n unsigned long long q = (~(1ull << (63 - k))); //bit of hash k in bitfields\n\n if (col_hash_k > row_hash_k || (col_hash_k == row_hash_k && row_id < col_id))\n {\n is_max &= q;\n }\n\n //unsigned long long x2 = col_hash_k < row_hash_k;\n if (col_hash_k < row_hash_k || (col_hash_k == row_hash_k && row_id > col_id))\n {\n is_min &= q;\n }\n }\n}\n\ntemplate\nstruct neigbor_visitor\n{\n\n template\n static __device__ __forceinline__ void visit(const int &row_id, const int &base_row_hash, const int &explore_row, const int *A_rows, const int *A_cols, const int *A_colors_in, const int lane_id, const int seed, unsigned long long int &is_min, unsigned long long int &is_max)\n {\n int row_begin = A_rows[explore_row ];\n int row_end = A_rows[explore_row + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n bool do_compute_hashes = true;\n col_id = A_cols[row_it];//__load_global(A_cols+row_it);\n#if DISCARD_COLORED\n\n if (DISCARD_COLORED_T)\n {\n int col_color = __ldg(A_colors_in + col_id);\n\n if (col_color != 0) { do_compute_hashes = false; }\n }\n\n#endif\n\n if (do_compute_hashes)\n {\n //#pragma unroll\n //for(int k=0; k<(HASHES/16>0)?HASHES/16:1; k++)\n //{\n int base_col_hash = hash2(col_id, seed/*+k*12358*/);\n compute_hashes(row_id, col_id, base_row_hash, base_col_hash, seed, is_min, is_max);\n //}\n }\n\n neigbor_visitor < COLORING_LEVEL - 1 >::template visit(row_id, base_row_hash, col_id, A_rows, A_cols, A_colors_in, lane_id, seed, is_min, is_max);\n }\n }\n\n};\ntemplate<>\nstruct neigbor_visitor<0>\n{\n\n template\n static __device__ __forceinline__ void visit(const int &row_id, const int &base_row_hash, const int &explore_row, const int *A_rows, const int *A_cols, const int *A_colors_in, const int lane_id, const int seed, unsigned long long int &is_min, unsigned long long int &is_max)\n {\n //is_min |= 0ull;\n }\n\n};\n\n\ntemplate< bool DISCARD_COLORED_T, int HASHES, int COLORING_LEVEL, int CTA_SIZE, int WARP_SIZE>\n__global__\nvoid fast_multihash_kernel(\n const int A_num_rows, const int *A_rows, const int *A_cols,\n int *A_colors, const int *A_colors_in,\n int c0, int seed,\n int *sets_per_block)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n int warp_count = 0;\n\n for ( ; utils::any(row_id < A_num_rows) ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = -1;\n int available_color = 0;\n\n if (row_id < A_num_rows)\n {\n#if DISCARD_COLORED\n\n if (DISCARD_COLORED_T)\n {\n row_color = A_colors_in[row_id];\n A_colors[row_id] = row_color;\n }\n else\n#endif\n {\n row_color = A_colors[row_id];\n }\n }\n\n if (row_color == 0)\n {\n unsigned long long int is_min = ~0ull;\n unsigned long long int is_max = ~0ull;\n int base_row_hash = hash2(row_id, seed);\n#if 1 //COLORING_LEVELR>1\n neigbor_visitor::template visit(row_id, base_row_hash, row_id, A_rows, A_cols, A_colors_in, lane_id, seed, is_min, is_max);\n#else\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n col_id = A_cols[row_it];//__load_global(A_cols+row_it);\n#if DISCARD_COLORED\n\n if (DISCARD_COLORED_T)\n {\n int col_color = __ldg(A_colors_in + col_id);\n\n if (col_color != 0) { continue; }\n }\n\n#endif\n int base_col_hash = hash2(col_id, seed);\n compute_hashes(row_id, col_id, base_row_hash, base_col_hash, seed, is_min, is_max);\n }\n\n#endif\n\n if (WARP_SIZE > 1)\n {\n using namespace amgx::strided_reduction;\n amgx::strided_reduction::op_and op;\n warp_reduce<1, CTA_SIZE, WARP_SIZE>(is_max, op);\n warp_reduce<1, CTA_SIZE, WARP_SIZE>(is_min, op);\n }\n\n#if 1\n int cmax = 63 - utils::bfind(is_max);\n int cmin = 63 - utils::bfind(is_min);\n\n if (cmin < HASHES)\n {\n available_color = cmin * 2 + 2;\n }\n\n if (cmax < HASHES)\n {\n available_color = cmax * 2 + 1;\n }\n\n#else\n#pragma unroll\n\n for (int k = 0; k < HASHES; k++)\n {\n if (is_max & (1ull << 63 - k))\n {\n available_color = 2 * k + 1;\n break;\n }\n else if (is_min & (1ull << 63 - k))\n {\n available_color = 2 * k + 2;\n break;\n }\n }\n\n#endif\n }\n\n bool has_set_color = (lane_id == 0) && (available_color != 0);\n\n if (has_set_color)\n {\n A_colors[row_id] = available_color + c0;\n }\n\n warp_count += amgx::strided_reduction::warp_binary_count(has_set_color);\n }\n\n#if GET_COLORS == GET_COLORS_ATOMIC\n amgx::strided_reduction::block_count_atomic<1, CTA_SIZE, 32, int>(warp_count, sets_per_block);\n //__threadfence_system();\n#elif GET_COLORS == GET_COLORS_TWOPASS\n amgx::strided_reduction::block_count<1, CTA_SIZE, 32, int>(warp_count, sets_per_block);\n#elif GET_COLORS == GET_COLORS_GATHER\n amgx::strided_reduction::block_count<1, CTA_SIZE, 32, int>(warp_count, sets_per_block);\n amgx::strided_reduction::block_count_gatherer global_gather(sets_per_block);\n#endif\n}\n\n////////////////////////////\n\ntemplate< int HASHES, int CTA_SIZE, int WARP_SIZE>\n__global__\nvoid fast_multihash_kernel_gtlt_kernel(\n const int A_num_rows, const int *A_rows, const int *A_cols,\n int *A_colors, unsigned long long int *gtlt_out,\n int seed)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n //int warp_count = 0;\n\n for ( ; utils::any(row_id < A_num_rows) ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = -1;\n\n if (row_id < A_num_rows)\n {\n row_color = A_colors[row_id];\n }\n\n if (row_color == 0)\n {\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n int base_row_hash = hash2(row_id, seed);\n unsigned char count_gt[HASHES];\n unsigned char count_lt[HASHES];\n#pragma unroll\n\n for (int k = 0; k < HASHES; k++)\n {\n count_gt[k] = 0;\n count_lt[k] = 0;\n }\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n col_id = A_cols[row_it];\n int col_color = A_colors[col_id];\n\n if (col_color != 0) { continue; }\n\n int base_col_hash = hash2(col_id, seed);\n int col_hash_k = base_col_hash;\n int row_hash_k = base_row_hash;\n#pragma unroll\n\n for (int k = 0; k < HASHES; k++)\n {\n col_hash_k = rehash(col_id, base_col_hash, k, seed);\n row_hash_k = rehash(row_id, base_row_hash, k, seed);\n\n if (col_hash_k > row_hash_k || (col_hash_k == row_hash_k && row_id < col_id))\n {\n if (count_gt[k] < 255) //avoid overflow\n {\n ++count_gt[k];\n }\n }\n\n if (col_hash_k < row_hash_k || (col_hash_k == row_hash_k && row_id > col_id))\n {\n if (count_lt[k] < 255)\n {\n ++count_lt[k];\n }\n }\n }\n }\n\n //compress count_gt,lt in 2 bits and pack in an integer\n unsigned long long int packed_gtlt = 0;\n#pragma unroll\n\n for (int k = 0; k < HASHES; k++)\n {\n unsigned long long int max3_gt = count_gt[k];\n unsigned long long int max3_lt = count_lt[k];\n\n if (max3_gt > 3) { max3_gt = 3; } //TODO use only bits above\n\n if (max3_lt > 3) { max3_lt = 3; }\n\n //TODO: use BFI\n packed_gtlt |= (max3_gt << (k * 2));\n packed_gtlt |= (max3_lt << (k * 2 + 32));\n }\n\n if (lane_id == 0)\n {\n gtlt_out[row_id] = packed_gtlt;\n }\n }\n }\n}\n\ntemplate< int HASHES, int CTA_SIZE, int WARP_SIZE>\n__global__\nvoid fast_multihash_kernel_gtlt_assign_kernel(\n const int A_num_rows, const int *A_rows, const int *A_cols,\n int *A_colors, unsigned long long int *gtlt_in,\n int c0)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n for ( ; utils::any(row_id < A_num_rows) ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = -1;\n unsigned long long int row_gtlt = 0ull;\n\n if (row_id < A_num_rows)\n {\n row_color = A_colors[row_id];\n row_gtlt = gtlt_in[row_id];\n }\n\n int available_color = 0;\n\n if (row_color == 0)\n {\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n unsigned long long int is_min = ~0ull;\n unsigned long long int is_max = ~0ull;\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n unsigned long long int col_gtlt = 0ull;\n\n if (row_it < row_end)\n {\n col_id = A_cols[row_it];\n col_gtlt = gtlt_in[col_id];\n#pragma unroll\n\n for (int k = 0; k < HASHES; k++)\n {\n int row_gt_k = (row_gtlt >> (k * 2 + 0)) & 3ull;\n int row_lt_k = (row_gtlt >> (k * 2 + 32)) & 3ull;\n int col_gt_k = (col_gtlt >> (k * 2 + 0)) & 3ull;\n int col_lt_k = (col_gtlt >> (k * 2 + 32)) & 3ull;\n unsigned long long q = (~(1ull << (63 - k)));\n\n if ((row_gt_k > col_gt_k) || (col_gt_k == row_gt_k && row_id < col_id))\n {\n is_max &= q; //xor with k bit: if was 1, becomes 0, if was 0 stays 0\n }\n\n if ((row_lt_k > col_lt_k) || (col_lt_k == row_lt_k && row_id > col_id))\n {\n is_min &= q;\n }\n }\n }\n }\n\n int cmax = 63 - utils::bfind(is_max);\n int cmin = 63 - utils::bfind(is_min);\n\n //if(cmin < HASHES)\n // available_color = cmin*2+2;\n if (cmax < HASHES)\n {\n available_color = cmax * 2 + 1;\n }\n }\n\n bool has_set_color = (lane_id == 0) && (available_color != 0);\n\n if (has_set_color)\n {\n A_colors[row_id] = available_color + c0;\n }\n }\n}\n\n////////////////////////////\n\ntemplate< bool SORTED_ROWS, int COLORING_LEVEL, int CTA_SIZE, int WARP_SIZE>\n__global__\nvoid recolor_greedy_kernel(\n const int A_num_rows, const int *__restrict A_rows,\n const int *__restrict A_cols,\n int color_to_reassign,\n int *__restrict A_colors,\n int *sorted_rows_by_color, int offset_start, int offset_end\n , int *maxs_per_block\n)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n int job_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n int last_job = A_num_rows;\n\n if (SORTED_ROWS)\n {\n job_id = job_id + offset_start;\n last_job = offset_end;\n }\n\n int thread_max = 0;\n\n //const int NUM_THREADS_PER_GRID = gridDim.x * CTA_SIZE;\n //int row_id = blockIdx.x * CTA_SIZE + threadIdx.x;\n\n // Iterate over rows.\n for ( ; utils::any(job_id < last_job) ; job_id += NUM_WARPS_PER_GRID )\n {\n int row_color = -1;\n int row_id = -1;\n\n if (job_id < last_job)\n {\n if (SORTED_ROWS)\n {\n row_id = sorted_rows_by_color[job_id];\n }\n else\n {\n row_id = job_id;\n }\n\n if (row_id < A_num_rows)\n {\n row_color = A_colors[row_id];\n }\n }\n\n used_color_structure_64_bit used_colors;\n used_color_structure_64_bit::aux_tmp aux;\n //used_color_structure_64_bit_colorbox<5> used_colors; //still buggy\n\n if (row_color == color_to_reassign)\n {\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n //used_colors.use_color(color_to_reassign);\n //int num_greater = 0;\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n\n if (row_it < row_end)\n {\n col_id = A_cols[row_it];\n }\n\n if (row_it < row_end)// && col_id < A_num_rows)\n {\n int col_color = A_colors[col_id];\n\n //if(col_color>40) printf(\"%d\\n\",col_color);\n if (col_color != 0)\n {\n used_colors.use_color(col_color, aux);\n }\n }\n\n if (COLORING_LEVEL == 2)\n {\n int row_begin = A_rows[col_id ];\n int row_end = A_rows[col_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n\n if (row_it < row_end)\n {\n col_id = A_cols[row_it];\n }\n\n if (row_it < row_end)// && col_id < A_num_rows)\n {\n int col_color = A_colors[col_id];\n\n //if(col_color>40) printf(\"%d\\n\",col_color);\n if (col_color != 0)\n {\n used_colors.use_color(col_color, aux);\n }\n }\n }\n }\n }\n }\n\n used_colors.sync_subwarp(lane_id);\n\n /*if(row_color != color_to_reassign && row_id != -1)\n {\n printf(\"ERR %d %d!\\n\", row_color, color_to_reassign);\n }*/\n if (row_color == color_to_reassign && lane_id == 0)\n {\n int new_color = used_colors.first_available_color(aux);\n\n if (new_color != row_color && new_color > 0)\n {\n //printf(\"changing %d from %d to %d\\n\",row_id,row_color,new_color);\n A_colors[row_id] = new_color;\n }\n\n#if GETMAX_COLOR == GETMAX_TWOPASS\n\n if (new_color > thread_max)\n {\n thread_max = new_color;\n }\n\n#endif\n /*if(new_color >= 7)\n {\n printf(\"SEVEN %d %d\\n\", new_color, thread_max);\n }*/\n }\n }\n\n#if GETMAX_COLOR == GETMAX_TWOPASS\n using namespace amgx::strided_reduction;\n amgx::strided_reduction::op_max op;\n warp_reduce<1, CTA_SIZE, 32>(thread_max, op);\n block_reduce<1, CTA_SIZE, 32, true>(thread_max, maxs_per_block, op);\n#endif\n}\n\n////////////////////////////\n\n\n\ntemplate< bool SORTED_ROWS, int COLORING_LEVEL, int CTA_SIZE, int WARP_SIZE>\n__global__\nvoid recolor_greedy_kernel_anyring(\n const int A_num_rows, const int *A_rows, const int *A_cols,\n int color_to_reassign, int *A_colors,\n int *sorted_rows_by_color, int offset_start, int offset_end, int *maxs_per_block,\n\n used_color_structure_64_bit *used_colors_inout\n)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n const int warp_id = subwarp::subwarp_id();\n const int lane_id = subwarp::sublane_id();\n int job_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n int last_job = A_num_rows;\n\n if (SORTED_ROWS)\n {\n job_id = job_id + offset_start;\n last_job = offset_end;\n }\n\n int thread_max = 0;\n\n //const int NUM_THREADS_PER_GRID = gridDim.x * CTA_SIZE;\n //int row_id = blockIdx.x * CTA_SIZE + threadIdx.x;\n\n // Iterate over rows.\n for ( ; utils::any(job_id < last_job) ; job_id += NUM_WARPS_PER_GRID )\n {\n int row_color = -1;\n int row_id = -1;\n\n if (job_id < last_job)\n {\n if (SORTED_ROWS)\n {\n row_id = sorted_rows_by_color[job_id];\n }\n else\n {\n row_id = job_id;\n }\n\n if (row_id < A_num_rows)\n {\n row_color = A_colors[row_id];\n }\n }\n\n used_color_structure_64_bit used_colors;\n used_color_structure_64_bit::aux_tmp aux;\n used_colors = used_colors_inout[row_id];\n //used_color_structure_64_bit_colorbox<5> used_colors; //still buggy\n\n if (row_color == color_to_reassign)\n {\n int new_color = used_colors.first_available_color(aux);\n\n if (lane_id == 0)\n {\n if (new_color != row_color && new_color > 0)\n {\n //printf(\"changing %d from %d to %d\\n\",row_id,row_color,new_color);\n A_colors[row_id] = new_color;\n }\n\n#if GETMAX_COLOR == GETMAX_TWOPASS\n\n if (new_color > thread_max)\n {\n thread_max = new_color;\n }\n\n#endif\n }\n\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n\n if (row_it < row_end)\n {\n col_id = A_cols[row_it];\n }\n\n if (row_it < row_end)// && col_id < A_num_rows)\n {\n int col_color = A_colors[col_id];\n\n if (col_color == 0)\n {\n used_color_structure_64_bit used_colors_col = used_colors_inout[col_id];\n used_colors_col.use_color(new_color, aux);\n used_colors_inout[col_id] = used_colors_col;\n }\n\n if (COLORING_LEVEL == 2)\n {\n int row_begin = A_rows[col_id ];\n int row_end = A_rows[col_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE)\n {\n int row_it = row_begin + lane_id;\n int col_id = -1;\n\n if (row_it < row_end)\n {\n col_id = A_cols[row_it];\n }\n\n if (row_it < row_end)// && col_id < A_num_rows)\n {\n int col_color = A_colors[col_id];\n\n if (col_color == 0)\n {\n used_color_structure_64_bit used_colors_col = used_colors_inout[col_id];\n used_colors_col.use_color(new_color, aux);\n used_colors_inout[col_id] = used_colors_col;\n }\n }\n }\n }\n }\n }\n }\n }\n\n#if GETMAX_COLOR == GETMAX_TWOPASS\n using namespace amgx::strided_reduction;\n amgx::strided_reduction::op_max op;\n warp_reduce<1, CTA_SIZE, 32>(thread_max, op);\n block_reduce<1, CTA_SIZE, 32, true>(thread_max, maxs_per_block, op);\n#endif\n}\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate< class T_Config >\nGreedy_Recolor_MatrixColoring_Base::Greedy_Recolor_MatrixColoring_Base( AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n first_pass_config = cfg;\n first_pass_config_scope = cfg_scope;\n //MULTI_HASH: it would be faster but it leaves some uncolored vertices: uncolored vertices cannot be processed in parallel\n first_pass_config.setParameter(\"matrix_coloring_scheme\", std::string(\"MULTI_HASH\"), first_pass_config_scope);\n first_pass_config.setParameter(\"max_uncolored_percentage\", 0.0, first_pass_config_scope);\n m_coloring_custom_arg = cfg.AMG_Config::template getParameter( \"coloring_custom_arg\", cfg_scope );\n m_coloring_try_remove_last_color_ = cfg.AMG_Config::template getParameter( \"coloring_try_remove_last_colors\", cfg_scope );\n}\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nGreedy_Recolor_MatrixColoring >::colorMatrix( Matrix_d &A )\n{\n const int num_rows = A.get_num_rows();\n const int num_nz = A.get_num_nz();\n int avg_nz = num_nz / num_rows;\n //printf(\"AVG %d \",avg_nz);\n //P(single) = 1-((avg-1)/avg)^(K)\n //P(all) = P(single)^(N)\n // = (1-((avg-1)/avg)^(K))^(N) == 0.99\n // = (1-((avg-1)/avg)^(K)) == Pow(0.99,1/N)\n //P(single == 0.99) K = Log(1-0.99)/Log(1/avg)\n //P(all == 0.99) K = Log(1-Pow(0.99,1/N))/Log(1/avg)\n //99% probability\n //table here http://www.wolframalpha.com/input/?i=Log%281-0.99%29%2FLog%28%28a-1%29%2Fa%29+for+a+in+[2%2C10]\n avg_nz = pow(static_cast(avg_nz), static_cast(this->m_coloring_level));\n\n if (avg_nz > 16) { avg_nz = 16; }\n\n if (this->m_coloring_custom_arg != \"\")\n {\n int x = atoi(this->m_coloring_custom_arg.c_str());\n\n if (x > 0)\n {\n avg_nz = x;\n }\n }\n\n#define CASENZ(X,ORDER,DO_SORT)\\\n else if(avg_nz <= X)\\\n {\\\n color_matrix_specialized_per_numhash(A);\\\n }\n#define CASEORDER(ORDER,DO_SORT)\\\n if(0)\\\n {\\\n\\\n }\\\n CASENZ(2,ORDER,DO_SORT)\\\n CASENZ(4,ORDER,DO_SORT)\\\n CASENZ(8,ORDER,DO_SORT)\\\n CASENZ(12,ORDER,DO_SORT)\\\n CASENZ(16,ORDER,DO_SORT)\\\n CASENZ(20,ORDER,DO_SORT)\\\n CASENZ(24,ORDER,DO_SORT)\\\n CASENZ(32,ORDER,DO_SORT)\\\n CASENZ(48,ORDER,DO_SORT)\\\n else\\\n {\\\n color_matrix_specialized_per_numhash<16,ORDER,DO_SORT>(A);\\\n }\n\n if (0)\n {\n }\n else if (this->m_coloring_level == 1)\n {\n CASEORDER(1, 0);\n }\n else if (this->m_coloring_level == 2)\n {\n CASEORDER(2, 1);\n }\n else if (this->m_coloring_level == 3)\n {\n CASEORDER(3, 1);\n }\n else if (this->m_coloring_level == 4)\n {\n CASEORDER(4, 1);\n }\n else if (this->m_coloring_level == 5)\n {\n CASEORDER(5, 1);\n }\n}\n\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\ntemplate\nvoid\nGreedy_Recolor_MatrixColoring >::color_matrix_specialized_per_numhash( Matrix_d &A )\n{\n const int num_rows = A.get_num_rows();\n const int n_blocks = 26 * 4; //TODO: choose optimal blocks number\n const int block_size = 256;\n ViewType oldView = A.currentView();\n#if 1 //use fast multihash above as precolorer\n this->m_row_colors.resize( A.row_offsets.size() - 1, 0 );\n typedef typename Matrix > ::IVector IVector_h;\n IVector sorted_rows_by_color;\n IVector_h offsets_rows_per_color;\n {\n device_vector_alloc sets_per_block(n_blocks + 1);\n int *sets_per_block_ptr = amgx::thrust::raw_pointer_cast(sets_per_block.data());\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n int num_uncolored = num_rows;\n int i = 0;\n#if DISCARD_COLORED\n device_vector_alloc color_in(num_rows);\n#else\n device_vector_alloc color_in;\n#endif\n#if USE_GTLT\n device_vector_alloc gtlt(num_rows);\n const int K = K1 * 0.5;\n#else\n const int K = K1;\n#endif\n int prev_num_uncolored = num_uncolored;\n unsigned int seed = 0;\n\n while (num_uncolored > 0)\n {\n //int mx = amgx::thrust::reduce( this->m_row_colors.begin(), this->m_row_colors.begin()+num_rows, 0, amgx::thrust::maximum() );\n //if(mx > 2*K*i) printf(\"ERR UX %d %d %d\\n\", num_uncolored, mx, 2*K*i);\n#if USE_GTLT\n fast_multihash_kernel_gtlt_kernel <<< n_blocks, block_size>>>\n (num_rows, A.row_offsets.raw(), A.col_indices.raw(),\n this->m_row_colors.raw(), amgx::thrust::raw_pointer_cast(gtlt.data()),\n seed);\n cudaCheckError();\n fast_multihash_kernel_gtlt_assign_kernel <<< n_blocks, block_size>>>\n (num_rows, A.row_offsets.raw(), A.col_indices.raw(),\n this->m_row_colors.raw(), amgx::thrust::raw_pointer_cast(gtlt.data()),\n 2 * K * i);\n cudaCheckError();\n#else\n\n if (num_uncolored == num_rows || (!DISCARD_COLORED))\n {\n fast_multihash_kernel <<< n_blocks, block_size>>>(num_rows, A.row_offsets.raw(), A.col_indices.raw(),\n this->m_row_colors.raw(), amgx::thrust::raw_pointer_cast(color_in.data()),\n 2 * K * i, seed,\n sets_per_block_ptr);\n cudaCheckError();\n }\n else\n {\n fast_multihash_kernel <<< n_blocks, block_size>>>(num_rows, A.row_offsets.raw(), A.col_indices.raw(),\n this->m_row_colors.raw(), amgx::thrust::raw_pointer_cast(color_in.data()),\n 2 * K * i, seed,\n sets_per_block_ptr);\n }\n\n cudaCheckError();\n#endif\n seed = hash2(seed, 0);\n#if GET_COLORS == GET_COLORS_THRUST || USE_GTLT\n num_uncolored = amgx::thrust::count( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, 0);\n cudaCheckError();\n#elif GET_COLORS == GET_COLORS_TWOPASS\n num_uncolored -= amgx::strided_reduction::count_block_results_1(n_blocks, sets_per_block_ptr, amgx::strided_reduction::op_sum());\n#elif (GET_COLORS == GET_COLORS_ATOMIC) || (GET_COLORS == GET_COLORS_GATHER)\n num_uncolored -= sets_per_block[0];\n#endif\n //printf(\"%d -> %d\\n\", prev_num_uncolored,num_uncolored,prev_num_uncolored-num_uncolored);\n i++;\n\n if (prev_num_uncolored == num_uncolored)\n {\n //printf(\"ERR %d %d %d\\n\", prev_num_uncolored,num_uncolored,num_rows);\n break;\n }\n\n prev_num_uncolored = num_uncolored;\n#if DISCARD_COLORED\n color_in.swap(this->m_row_colors);\n#endif\n /*if (iteration++%4==0) {\n cudaEventRecord(throttle_event);\n cudaCheckError();\n } else {\n cudaEventSynchronize(throttle_event);\n cudaCheckError();\n }*/\n }\n\n#if DISCARD_COLORED\n color_in.swap(this->m_row_colors);\n#endif\n int max_color = i * K * 2; //available_color+c0 = 2*K*i + 2*(K-1)+2=2K(i+1), but i++\n //max_color = amgx::thrust::reduce( this->m_row_colors.begin(), this->m_row_colors.begin()+num_rows, 0, amgx::thrust::maximum() );\n //printf(\"CC %d<%d\\n\",max_color,i*K*2+2);\n this->m_num_colors = max_color + 1;\n#if ONLY_FIRSTSTEP\n\n if (1)\n#else\n if (this->m_coloring_custom_arg == \"FIRSTSTEP\")\n#endif\n {\n A.setView(oldView);\n return;\n }\n\n if (DO_SORT)\n {\n //return;\n //this->createColorArrays(A)\n IVector row_colors(this->m_row_colors); //useless copy\n sorted_rows_by_color.resize(num_rows);\n //this->m_sorted_rows_by_color.resize(num_rows);\n thrust_wrapper::sequence(sorted_rows_by_color.begin(), sorted_rows_by_color.end()); //useless sequence\n thrust_wrapper::sort_by_key(row_colors.begin(), row_colors.begin() + num_rows, sorted_rows_by_color.begin()); //useless read from sequence\n cudaCheckError();\n IVector offsets_rows_per_color_d(this->m_num_colors + 1);\n amgx::thrust::lower_bound(row_colors.begin(),\n row_colors.begin() + num_rows,\n amgx::thrust::counting_iterator(0),\n amgx::thrust::counting_iterator(offsets_rows_per_color_d.size()),\n offsets_rows_per_color_d.begin());\n // Copy from device to host\n offsets_rows_per_color = offsets_rows_per_color_d;\n cudaCheckError();\n /*for(int i=0;im_halo_coloring == SYNC_COLORS) A.setView(ALL);\n else A.setViewExterior();*/\n }\n }\n#else\n //1st pass: color with a fast method\n MatrixColoring *first_pass_matrix_coloring = MatrixColoringFactory::allocate(this->first_pass_config, this->first_pass_config_scope);\n first_pass_matrix_coloring->colorMatrix(A);\n#if DO_SORT\n first_pass_matrix_coloring->createColorArrays(A);\n#endif\n this->m_row_colors = first_pass_matrix_coloring->getRowColors();\n this->m_num_colors = first_pass_matrix_coloring->getNumColors();\n sorted_rows_by_color = first_pass_matrix_coloring->getSortedRowsByColor();\n offsets_rows_per_color = first_pass_matrix_coloring->getOffsetsRowsPerColor();\n delete first_pass_matrix_coloring;\n //ViewType oldView = A.currentView();\n //if (this->m_halo_coloring == SYNC_COLORS) A.setView(ALL);\n //else A.setViewExterior();\n#endif\n cudaCheckError();\n //2nd pass: optimize using parallel greedy for each color.\n //If 1st pass coloring is valid all elements of same color can be processed in parallel by definition\n#if GETMAX_COLOR == GETMAX_TWOPASS\n device_vector_alloc max_color_per_block(n_blocks + 1, 1);\n int *max_color_per_block_ptr = amgx::thrust::raw_pointer_cast(max_color_per_block.data());\n#else\n int *max_color_per_block_ptr = 0;\n#endif\n float avg_nnz_per_row = A.get_num_nz() / float(A.get_num_rows());\n int WARP_SIZE = 2;\n\n if ( avg_nnz_per_row < 4.0f )\n {\n WARP_SIZE = 2;\n }\n else if ( avg_nnz_per_row < 8.0f )\n {\n WARP_SIZE = 4;\n }\n else if ( avg_nnz_per_row < 32.0f )\n {\n WARP_SIZE = 8;\n }\n else\n {\n WARP_SIZE = 32;\n }\n\n if (this->m_coloring_level >= 2)\n {\n WARP_SIZE = 1;\n }\n\n //;\n //int max_colorx = amgx::thrust::reduce( this->m_row_colors.begin(), this->m_row_colors.begin()+num_rows, 0, amgx::thrust::maximum() );\n //printf(\"MAX %d\",max_colorx);\n#define RECOLOR_CALL_CASE(__WSIZE) if(WARP_SIZE==__WSIZE){ \\\n recolor_greedy_kernel<<>>(\\\n num_rows,\\\n A.row_offsets.raw(),\\\n A.col_indices.raw(),\\\n color, this->m_row_colors.raw(),\\\n sorted_rows_by_color.raw(),start,end,\\\n max_color_per_block_ptr\\\n ); cudaCheckError(); }\n cudaCheckError();\n#if 1 //straight order: generates a color histogram similar to parallel greedy\n //start from color=2: 1st color would be assigned to 1, no change\n\n //iteration = 0;\n\n for (int color = 2; color <= this->m_num_colors; color++)\n {\n const int GRID_SIZE = n_blocks;//\n int start, end;\n\n if (DO_SORT)\n {\n start = offsets_rows_per_color[color];\n end = num_rows;\n\n if (color < this->m_num_colors) { end = offsets_rows_per_color[color + 1]; }\n\n if (end >= num_rows) { end = num_rows; }\n\n //printf(\"F %d %d %d < num_rows: %d\\n\",color,start,end,num_rows,(int)end<=start);\n if (end <= start) { continue; }\n\n //int NUM_WARPS_PER_CTA = block_size/WARP_SIZE;\n //const int GRID_SIZE = std::min( 2048, ((end-start) + NUM_WARPS_PER_CTA-1) / NUM_WARPS_PER_CTA );\n //printf(\"%d - %d > %d %d %d\\n\",end-start,color,start,end,num_rows);\n }\n else\n {\n start = 0;\n end = num_rows;\n }\n\n //printf(\"%d : %d %d\\n\", color, start, end);\n RECOLOR_CALL_CASE(1);\n RECOLOR_CALL_CASE(2);\n RECOLOR_CALL_CASE(4);\n RECOLOR_CALL_CASE(8);\n RECOLOR_CALL_CASE(16);\n RECOLOR_CALL_CASE(32);\n cudaCheckError();\n /*if (iteration++%4==0) {\n cudaEventRecord(throttle_event);\n cudaCheckError();\n } else {\n cudaEventSynchronize(throttle_event);\n cudaCheckError();\n }*/\n }\n\n#else //reverse order: generates a different kind of histogram\n\n for (int color = this->m_num_colors; color >= 0; color--)\n {\n RECOLOR_CALL_CASE(1);\n RECOLOR_CALL_CASE(2);\n RECOLOR_CALL_CASE(4);\n RECOLOR_CALL_CASE(8);\n RECOLOR_CALL_CASE(16);\n RECOLOR_CALL_CASE(32);\n }\n\n#endif\n cudaCheckError();\n#if GETMAX_COLOR == GETMAX_TWOPASS\n int max_color = amgx::strided_reduction::count_block_results_1(n_blocks, max_color_per_block_ptr, amgx::strided_reduction::op_max());\n /*int max_colorg = amgx::thrust::reduce( this->m_row_colors.begin(), this->m_row_colors.begin()+num_rows, 0, amgx::thrust::maximum() );\n\n max_color=0;\n for(int i=0;i max_color) max_color = max_color_per_block[i];\n }\n printf(\"MAXCOLOR ref=%d %d\\n\",max_colorg,max_color);\n if(max_colorg!=max_color)exit(1);*/\n#else\n int max_color = thrust_wrapper::reduce( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, 0, amgx::thrust::maximum() );\n#endif\n //printf(\"MAXCOLOR %d %d\\n\",max_color_gold,max_color);\n cudaCheckError();\n this->m_num_colors = max_color + 1;\n A.setView(oldView);\n}\n\n\n\n#define AMGX_CASE_LINE(CASE) template class Greedy_Recolor_MatrixColoring_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Greedy_Recolor_MatrixColoring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // end namespace amgx\n", "source": "/kaggle/working/amgx/src/matrix_coloring/greedy_recolor.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#define MAX(a, b) (((a) > (b)) ? (a) : (b))\n\n//#define USE_EXPERIMENTAL_MIN_MAX\n#define USE_EXPERIMENTAL_NEIGHBORS\n\n// Pseudo-random number generator\nnamespace amgx\n{\n\n__host__ __device__ unsigned int hash_function(unsigned int a, unsigned int seed)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n\nstruct is_zero\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x == 0;\n }\n};\n\n// ---------------------------\n// Kernels\n// ---------------------------\n\n// Kernel to color the rows of the matrix, using min-max approach\n#ifdef USE_EXPERIMENTAL_MIN_MAX\ntemplate< typename IndexType, int kCtaSize >\n__global__\nvoid colorRowsKernel(const IndexType *A_offsets,\n const IndexType *A_column_indices,\n IndexType *row_colors,\n const int current_color,\n const int num_rows)\n{\n const int num_rows_per_grid = gridDim.x * ( kCtaSize / 32 );\n const int lane_id = threadIdx.x % 32;\n\n for ( int row_id = ( blockIdx.x * kCtaSize + threadIdx.x ) / 32 ; row_id < num_rows ; row_id += num_rows_per_grid )\n {\n int my_row_color = row_colors[row_id];\n\n if ( my_row_color != 0 )\n {\n continue;\n }\n\n int hash_i = hash_function(row_id, 0);\n bool max_i = true;\n bool min_i = true;\n int row_it = A_offsets[row_id ];\n int row_end = A_offsets[row_id + 1];\n\n for ( row_it += lane_id ; utils::any(row_it < row_end) ; row_it += 32 )\n {\n int j = row_it < row_end ? A_column_indices[row_it] : -1;\n int hash_j = hash_function(j, 0);\n int row_color_j = -1;\n\n if ( row_it < row_end && hash_i != hash_j && j < num_rows ) //TODO: is this correct? (ignore edges thru bndry\n {\n row_color_j = row_colors[j];\n }\n\n // There is an uncolored neighbour that is greater\n max_i &= utils::all( hash_i > hash_j || ( row_color_j != 0 && row_color_j != current_color ) );\n // There is an uncolored neighbour that is smaller\n min_i &= utils::all( hash_i < hash_j || ( row_color_j != 0 && row_color_j != current_color + 1 ) );\n }\n\n if ( lane_id == 0 && (max_i || min_i) )\n {\n row_colors[row_id] = current_color + (min_i ? 1 : 0);\n }\n }\n}\n#else\ntemplate \n__global__\nvoid colorRowsKernel(const IndexType *A_offsets, const IndexType *A_column_indices, IndexType *row_colors,\n const int current_color, const int num_rows)\n{\n for ( int i = threadIdx.x + blockDim.x * blockIdx.x; i < num_rows ; i += gridDim.x * blockDim.x )\n {\n int my_row_color = 0;\n\n // skip if already colored\n if (row_colors[i] == 0)\n {\n int hash_i = hash_function(i, 0);\n bool max_i, min_i;\n max_i = true;\n min_i = true;\n int row_start = A_offsets[i];\n int row_end = A_offsets[i + 1];\n\n for (int r = row_start; r < row_end; r++)\n {\n int j = A_column_indices[r];\n\n if (j >= num_rows) { continue; }\n\n int hash_j = hash_function(j, 0);\n int row_color_j = row_colors[j];\n\n // There is an uncolored neighbour that is greater\n if ( hash_j > hash_i && (row_color_j == 0 || row_color_j == current_color) )\n {\n max_i = false;\n }\n\n // There is an uncolored neighbour that is smaller\n if (hash_j < hash_i && (row_color_j == 0 || row_color_j == current_color + 1) )\n {\n min_i = false;\n }\n }\n\n if (max_i)\n {\n my_row_color = current_color;\n }\n else if (min_i)\n {\n my_row_color = current_color + 1;\n }\n\n if (my_row_color != 0)\n {\n row_colors[i] = my_row_color;\n }\n }\n }\n}\n#endif\n\n#ifdef USE_EXPERIMENTAL_NEIGHBORS\n\ntemplate< int CTA_SIZE, int NUM_THREADS_PER_ROW >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nfind_min_max_neighbors_kernel( const int *__restrict A_offsets,\n const int *__restrict A_column_indices,\n const int *__restrict row_colors,\n int *__restrict max_hash_array,\n int *__restrict min_hash_array,\n const int current_color,\n const int num_rows )\n{\n const int NUM_ROWS_PER_CTA = CTA_SIZE / NUM_THREADS_PER_ROW;\n const int warp_id = threadIdx.x / NUM_THREADS_PER_ROW;\n const int lane_id = threadIdx.x % NUM_THREADS_PER_ROW;\n\n for ( int row_id = blockIdx.x * NUM_ROWS_PER_CTA + warp_id ; utils::any(row_id < num_rows) ; row_id += gridDim.x * NUM_ROWS_PER_CTA )\n {\n int max_hash = INT_MIN;\n int min_hash = INT_MAX;\n\n if ( lane_id == 0 && row_id < num_rows )\n {\n int hash_i = hash_function(row_id, 0);\n\n if ( row_colors[row_id] == 0 )\n {\n max_hash = hash_i;\n min_hash = hash_i;\n }\n }\n\n int row_it = 0, row_end = 0;\n\n if ( row_id < num_rows )\n {\n row_it = A_offsets[row_id + 0];\n row_end = A_offsets[row_id + 1];\n }\n\n for ( row_it += lane_id ; utils::any(row_it < row_end) ; row_it += NUM_THREADS_PER_ROW )\n {\n int a_col_id = -1;\n\n if ( row_it < row_end )\n {\n a_col_id = A_column_indices[row_it];\n }\n\n int hash_j = 0;\n int col_color = -1;\n\n if ( a_col_id != -1 && a_col_id < num_rows )\n {\n hash_j = hash_function(a_col_id, 0);\n col_color = row_colors[a_col_id];\n }\n\n // There is an uncolored neighbour that is greater\n if ( col_color == 0 && hash_j > max_hash )\n {\n max_hash = hash_j;\n }\n\n // There is an uncolored neighbour that is smaller\n if ( col_color == 0 && hash_j < min_hash )\n {\n min_hash = hash_j;\n }\n }\n\n#pragma unroll\n for ( int mask = NUM_THREADS_PER_ROW / 2 ; mask > 0 ; mask >>= 1 )\n {\n min_hash = min( min_hash, utils::shfl_xor( min_hash, mask ) );\n max_hash = max( max_hash, utils::shfl_xor( max_hash, mask ) );\n }\n\n if ( row_id < num_rows && lane_id == 0 )\n {\n max_hash_array[row_id] = max_hash;\n min_hash_array[row_id] = min_hash;\n }\n }\n}\n\n#else\n\ntemplate \n__global__\nvoid FindMaxMinNeighboursKernel(const IndexType *A_offsets, const IndexType *A_column_indices, IndexType *row_colors, IndexType *max_hash_array, IndexType *min_hash_array,\n const int current_color, const int num_rows)\n{\n for ( int i = threadIdx.x + blockDim.x * blockIdx.x; i < num_rows ; i += gridDim.x * blockDim.x )\n {\n int max_hash;\n int min_hash;\n int hash_i = hash_function(i, 0);\n int row_col = row_colors[i];\n max_hash = (row_col == 0) ? hash_i : INT_MIN;\n min_hash = (row_col == 0) ? hash_i : INT_MAX;\n int row_start = A_offsets[i];\n int row_end = A_offsets[i + 1];\n\n for (int r = row_start; r < row_end; r++)\n {\n int j = A_column_indices[r];\n\n if (j >= num_rows) { continue; }\n\n int hash_j = hash_function(j, 0);\n int row_color_j = row_colors[j];\n // Only consider if neighbour is not colored\n // There is an uncolored neighbour that is greater\n max_hash = (row_color_j == 0 && hash_j > max_hash) ? hash_j : max_hash;\n // There is an uncolored neighbour that is smaller\n min_hash = (row_color_j == 0 && hash_j < min_hash) ? hash_j : min_hash;\n }\n\n max_hash_array[i] = max_hash;\n min_hash_array[i] = min_hash;\n }\n}\n\n#endif\n\ntemplate \n__global__\nvoid colorRowsRingTwoKernel(const IndexType *A_offsets, const IndexType *A_column_indices, IndexType *row_colors, const IndexType *max_hash_array, const IndexType *min_hash_array,\n const int current_color, const int num_rows)\n{\n for ( int i = threadIdx.x + blockDim.x * blockIdx.x; i < num_rows ; i += gridDim.x * blockDim.x )\n {\n int my_row_color = 0;\n\n // skip if already colored\n if (row_colors[i] == 0)\n {\n int hash_i = hash_function(i, 0);\n int max_hash = max_hash_array[i];\n int min_hash = min_hash_array[i];\n bool max_i, min_i;\n max_i = true;\n min_i = true;\n\n // Check if I was a max or min in previous step\n if ( hash_i != max_hash && hash_i != min_hash)\n {\n continue;\n }\n\n if (hash_i != min_hash) { min_i = false; }\n\n if (hash_i != max_hash) { max_i = false; }\n\n int row_start = A_offsets[i];\n int row_end = A_offsets[i + 1];\n\n for (int r = row_start; r < row_end; r++)\n {\n int j = A_column_indices[r];\n\n if (j >= num_rows) { continue; }\n\n int min_hash_j = min_hash_array[j];\n int max_hash_j = max_hash_array[j];\n // There is a neighbour that has a uncolored neighbour with larger hash\n max_i = (max_hash_j > max_hash && max_hash_j != INT_MIN) ? false : max_i;\n // There is a neighbour that has a uncolored neighbour with smaller hash\n min_i = (min_hash_j < min_hash && min_hash_j != INT_MAX) ? false : min_i;\n }\n\n if (max_i && min_i)\n {\n my_row_color = (i % 2) ? current_color : current_color + 1;\n }\n else if (max_i)\n {\n my_row_color = current_color;\n }\n else if (min_i)\n {\n my_row_color = current_color + 1;\n }\n\n if (my_row_color != 0)\n {\n row_colors[i] = my_row_color;\n }\n }\n }\n}\n\n\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate\nMinMaxMatrixColoringBase::MinMaxMatrixColoringBase(AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n if (cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\"))\n {\n m_uncolored_fraction = 0;\n }\n else\n {\n m_uncolored_fraction = cfg.AMG_Config::template getParameter(\"max_uncolored_percentage\", cfg_scope);\n }\n}\n\n\n\n// Block version\ntemplate \nvoid MinMaxMatrixColoring >::colorMatrixOneRing(Matrix_d &A)\n{\n const int num_rows = A.get_num_rows();\n int max_uncolored_rows = (int) (this->m_uncolored_fraction * ((ValueType) num_rows));\n IndexType *row_colors_ptr = this->m_row_colors.raw();\n const int threads_per_block = 256;\n const int warps_per_block = threads_per_block / 32;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_rows + warps_per_block - 1) / warps_per_block );\n this->m_num_colors = 1;\n thrust_wrapper::fill(this->m_row_colors.begin(), this->m_row_colors.end(), 0);\n cudaCheckError();\n\n for ( int num_uncolored = num_rows; num_uncolored > max_uncolored_rows ; )\n {\n#ifdef USE_EXPERIMENTAL_MIN_MAX\n colorRowsKernel <<< num_blocks, threads_per_block, 0, amgx::thrust::global_thread_handle::get_stream()>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n row_colors_ptr,\n this->m_num_colors,\n num_rows);\n cudaCheckError();\n#else\n colorRowsKernel <<< num_blocks, threads_per_block, 0, amgx::thrust::global_thread_handle::get_stream()>>>(A.row_offsets.raw(), A.col_indices.raw(), row_colors_ptr, this->m_num_colors, num_rows);\n cudaCheckError();\n#endif\n cudaCheckError();\n this->m_num_colors += 2;\n num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_zero() );\n cudaCheckError();\n }\n\n this->m_num_colors = thrust_wrapper::reduce( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, 0, amgx::thrust::maximum() ) + 1;\n cudaCheckError();\n}\n\n// Block version\ntemplate \nvoid MinMaxMatrixColoringBase::colorMatrix(Matrix &A)\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize(A.row_offsets.size() - 1);\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n if (this->m_coloring_level == 0)\n {\n FatalError(\"Calling coloring scheme but coloring level==0\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n }\n else if (this->m_coloring_level == 1)\n {\n this->colorMatrixOneRing(A);\n }\n else if (this->m_coloring_level == 2)\n {\n this->colorMatrixTwoRing(A);\n }\n else\n {\n FatalError(\"Min-max coloring does not support coloring_level > 2\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n }\n\n A.setView(oldView);\n}\n\ntemplate \nvoid MinMaxMatrixColoring >::colorMatrixTwoRing(Matrix_d &A)\n{\n const int num_rows = A.get_num_rows();\n const int max_uncolored_rows = (int) (this->m_uncolored_fraction * ((ValueType) num_rows));\n const int threads_per_block = 256;\n#ifdef USE_EXPERIMENTAL_NEIGHBORS\n const int half_warps_per_block = threads_per_block / 16;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_rows + half_warps_per_block - 1) / half_warps_per_block );\n#else\n const int warps_per_block = threads_per_block / 32;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_rows + warps_per_block - 1) / warps_per_block );\n#endif\n this->m_num_colors = 1;\n thrust_wrapper::fill(this->m_row_colors.begin(), this->m_row_colors.end(), 0);\n cudaCheckError();\n IVector max_hash_array(num_rows);\n IVector min_hash_array(num_rows);\n float avg_nnz_per_row = A.get_num_nz() / float(num_rows);\n\n for ( int num_uncolored = num_rows; num_uncolored > max_uncolored_rows ; )\n {\n // Each vertex checks its neighbours and store the max and min values of its neighbours\n#ifdef USE_EXPERIMENTAL_NEIGHBORS\n if ( avg_nnz_per_row < 8.0f )\n {\n find_min_max_neighbors_kernel <<< num_blocks, threads_per_block>>>( A.row_offsets.raw(), A.col_indices.raw(), this->m_row_colors.raw(), max_hash_array.raw(), min_hash_array.raw(), this->m_num_colors, num_rows );\n cudaCheckError();\n }\n else if ( avg_nnz_per_row < 16.0f )\n {\n find_min_max_neighbors_kernel <<< num_blocks, threads_per_block>>>( A.row_offsets.raw(), A.col_indices.raw(), this->m_row_colors.raw(), max_hash_array.raw(), min_hash_array.raw(), this->m_num_colors, num_rows );\n cudaCheckError();\n }\n else // if( avg_nnz_per_row < 32.0f )\n {\n find_min_max_neighbors_kernel <<< num_blocks, threads_per_block>>>( A.row_offsets.raw(), A.col_indices.raw(), this->m_row_colors.raw(), max_hash_array.raw(), min_hash_array.raw(), this->m_num_colors, num_rows );\n cudaCheckError();\n }\n\n#else\n FindMaxMinNeighboursKernel <<< num_blocks, threads_per_block>>>(A.row_offsets.raw(), A.col_indices.raw(), this->m_row_colors.raw(), max_hash_array.raw(), min_hash_array.raw(), this->m_num_colors, num_rows);\n cudaCheckError();\n#endif\n // Each vertex checks if its still the min or max\n colorRowsRingTwoKernel <<< num_blocks, threads_per_block>>>(A.row_offsets.raw(), A.col_indices.raw(), this->m_row_colors.raw(), max_hash_array.raw(), min_hash_array.raw(), this->m_num_colors, num_rows);\n cudaCheckError();\n this->m_num_colors += 2;\n num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_zero() );\n cudaCheckError();\n }\n\n this->m_num_colors = thrust_wrapper::reduce( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, 0, amgx::thrust::maximum() ) + 1;\n cudaCheckError();\n#if 0\n std::cout << \"Num colors=\" << this->m_num_colors << std::endl;\n // Print coloring statistics\n //for (int i=0;im_num_colors;i++)\n // std::cout << \"color= \" << i << \", num_rows = \" << A.offsets_rows_per_color[i+1] - A.offsets_rows_per_color[i] << std::endl;\n#endif\n}\n\ntemplate \nvoid MinMaxMatrixColoring >::colorMatrixOneRing(Matrix_h &A)\n{\n FatalError(\"Haven't implemented MinMax matrix coloring for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MinMaxMatrixColoring >::colorMatrixTwoRing(Matrix_h &A)\n{\n FatalError(\"Haven't implemented MinMax matrix coloring for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n\n#define AMGX_CASE_LINE(CASE) template class MinMaxMatrixColoringBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class MinMaxMatrixColoring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n\n} // end namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#define MAX(a, b) (((a) > (b)) ? (a) : (b))\n\n//#define USE_EXPERIMENTAL_MIN_MAX\n#define USE_EXPERIMENTAL_NEIGHBORS\n\n// Pseudo-random number generator\nnamespace amgx\n{\n\n__host__ __device__ unsigned int hash_function(unsigned int a, unsigned int seed)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n\nstruct is_zero\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x == 0;\n }\n};\n\n// ---------------------------\n// Kernels\n// ---------------------------\n\n// Kernel to color the rows of the matrix, using min-max approach\n#ifdef USE_EXPERIMENTAL_MIN_MAX\ntemplate< typename IndexType, int kCtaSize >\n__global__\nvoid colorRowsKernel(const IndexType *A_offsets,\n const IndexType *A_column_indices,\n IndexType *row_colors,\n const int current_color,\n const int num_rows)\n{\n const int num_rows_per_grid = gridDim.x * ( kCtaSize / 32 );\n const int lane_id = threadIdx.x % 32;\n\n for ( int row_id = ( blockIdx.x * kCtaSize + threadIdx.x ) / 32 ; row_id < num_rows ; row_id += num_rows_per_grid )\n {\n int my_row_color = row_colors[row_id];\n\n if ( my_row_color != 0 )\n {\n continue;\n }\n\n int hash_i = hash_function(row_id, 0);\n bool max_i = true;\n bool min_i = true;\n int row_it = A_offsets[row_id ];\n int row_end = A_offsets[row_id + 1];\n\n for ( row_it += lane_id ; utils::any(row_it < row_end) ; row_it += 32 )\n {\n int j = row_it < row_end ? A_column_indices[row_it] : -1;\n int hash_j = hash_function(j, 0);\n int row_color_j = -1;\n\n if ( row_it < row_end && hash_i != hash_j && j < num_rows ) //TODO: is this correct? (ignore edges thru bndry\n {\n row_color_j = row_colors[j];\n }\n\n // There is an uncolored neighbour that is greater\n max_i &= utils::all( hash_i > hash_j || ( row_color_j != 0 && row_color_j != current_color ) );\n // There is an uncolored neighbour that is smaller\n min_i &= utils::all( hash_i < hash_j || ( row_color_j != 0 && row_color_j != current_color + 1 ) );\n }\n\n if ( lane_id == 0 && (max_i || min_i) )\n {\n row_colors[row_id] = current_color + (min_i ? 1 : 0);\n }\n }\n}\n#else\ntemplate \n__global__\nvoid colorRowsKernel(const IndexType *A_offsets, const IndexType *A_column_indices, IndexType *row_colors,\n const int current_color, const int num_rows)\n{\n for ( int i = threadIdx.x + blockDim.x * blockIdx.x; i < num_rows ; i += gridDim.x * blockDim.x )\n {\n int my_row_color = 0;\n\n // skip if already colored\n if (row_colors[i] == 0)\n {\n int hash_i = hash_function(i, 0);\n bool max_i, min_i;\n max_i = true;\n min_i = true;\n int row_start = A_offsets[i];\n int row_end = A_offsets[i + 1];\n\n for (int r = row_start; r < row_end; r++)\n {\n int j = A_column_indices[r];\n\n if (j >= num_rows) { continue; }\n\n int hash_j = hash_function(j, 0);\n int row_color_j = row_colors[j];\n\n // There is an uncolored neighbour that is greater\n if ( hash_j > hash_i && (row_color_j == 0 || row_color_j == current_color) )\n {\n max_i = false;\n }\n\n // There is an uncolored neighbour that is smaller\n if (hash_j < hash_i && (row_color_j == 0 || row_color_j == current_color + 1) )\n {\n min_i = false;\n }\n }\n\n if (max_i)\n {\n my_row_color = current_color;\n }\n else if (min_i)\n {\n my_row_color = current_color + 1;\n }\n\n if (my_row_color != 0)\n {\n row_colors[i] = my_row_color;\n }\n }\n }\n}\n#endif\n\n#ifdef USE_EXPERIMENTAL_NEIGHBORS\n\ntemplate< int CTA_SIZE, int NUM_THREADS_PER_ROW >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nfind_min_max_neighbors_kernel( const int *__restrict A_offsets,\n const int *__restrict A_column_indices,\n const int *__restrict row_colors,\n int *__restrict max_hash_array,\n int *__restrict min_hash_array,\n const int current_color,\n const int num_rows )\n{\n const int NUM_ROWS_PER_CTA = CTA_SIZE / NUM_THREADS_PER_ROW;\n const int warp_id = threadIdx.x / NUM_THREADS_PER_ROW;\n const int lane_id = threadIdx.x % NUM_THREADS_PER_ROW;\n\n for ( int row_id = blockIdx.x * NUM_ROWS_PER_CTA + warp_id ; utils::any(row_id < num_rows) ; row_id += gridDim.x * NUM_ROWS_PER_CTA )\n {\n int max_hash = INT_MIN;\n int min_hash = INT_MAX;\n\n if ( lane_id == 0 && row_id < num_rows )\n {\n int hash_i = hash_function(row_id, 0);\n\n if ( row_colors[row_id] == 0 )\n {\n max_hash = hash_i;\n min_hash = hash_i;\n }\n }\n\n int row_it = 0, row_end = 0;\n\n if ( row_id < num_rows )\n {\n row_it = A_offsets[row_id + 0];\n row_end = A_offsets[row_id + 1];\n }\n\n for ( row_it += lane_id ; utils::any(row_it < row_end) ; row_it += NUM_THREADS_PER_ROW )\n {\n int a_col_id = -1;\n\n if ( row_it < row_end )\n {\n a_col_id = A_column_indices[row_it];\n }\n\n int hash_j = 0;\n int col_color = -1;\n\n if ( a_col_id != -1 && a_col_id < num_rows )\n {\n hash_j = hash_function(a_col_id, 0);\n col_color = row_colors[a_col_id];\n }\n\n // There is an uncolored neighbour that is greater\n if ( col_color == 0 && hash_j > max_hash )\n {\n max_hash = hash_j;\n }\n\n // There is an uncolored neighbour that is smaller\n if ( col_color == 0 && hash_j < min_hash )\n {\n min_hash = hash_j;\n }\n }\n\n#pragma unroll\n for ( int mask = NUM_THREADS_PER_ROW / 2 ; mask > 0 ; mask >>= 1 )\n {\n min_hash = min( min_hash, utils::shfl_xor( min_hash, mask ) );\n max_hash = max( max_hash, utils::shfl_xor( max_hash, mask ) );\n }\n\n if ( row_id < num_rows && lane_id == 0 )\n {\n max_hash_array[row_id] = max_hash;\n min_hash_array[row_id] = min_hash;\n }\n }\n}\n\n#else\n\ntemplate \n__global__\nvoid FindMaxMinNeighboursKernel(const IndexType *A_offsets, const IndexType *A_column_indices, IndexType *row_colors, IndexType *max_hash_array, IndexType *min_hash_array,\n const int current_color, const int num_rows)\n{\n for ( int i = threadIdx.x + blockDim.x * blockIdx.x; i < num_rows ; i += gridDim.x * blockDim.x )\n {\n int max_hash;\n int min_hash;\n int hash_i = hash_function(i, 0);\n int row_col = row_colors[i];\n max_hash = (row_col == 0) ? hash_i : INT_MIN;\n min_hash = (row_col == 0) ? hash_i : INT_MAX;\n int row_start = A_offsets[i];\n int row_end = A_offsets[i + 1];\n\n for (int r = row_start; r < row_end; r++)\n {\n int j = A_column_indices[r];\n\n if (j >= num_rows) { continue; }\n\n int hash_j = hash_function(j, 0);\n int row_color_j = row_colors[j];\n // Only consider if neighbour is not colored\n // There is an uncolored neighbour that is greater\n max_hash = (row_color_j == 0 && hash_j > max_hash) ? hash_j : max_hash;\n // There is an uncolored neighbour that is smaller\n min_hash = (row_color_j == 0 && hash_j < min_hash) ? hash_j : min_hash;\n }\n\n max_hash_array[i] = max_hash;\n min_hash_array[i] = min_hash;\n }\n}\n\n#endif\n\ntemplate \n__global__\nvoid colorRowsRingTwoKernel(const IndexType *A_offsets, const IndexType *A_column_indices, IndexType *row_colors, const IndexType *max_hash_array, const IndexType *min_hash_array,\n const int current_color, const int num_rows)\n{\n for ( int i = threadIdx.x + blockDim.x * blockIdx.x; i < num_rows ; i += gridDim.x * blockDim.x )\n {\n int my_row_color = 0;\n\n // skip if already colored\n if (row_colors[i] == 0)\n {\n int hash_i = hash_function(i, 0);\n int max_hash = max_hash_array[i];\n int min_hash = min_hash_array[i];\n bool max_i, min_i;\n max_i = true;\n min_i = true;\n\n // Check if I was a max or min in previous step\n if ( hash_i != max_hash && hash_i != min_hash)\n {\n continue;\n }\n\n if (hash_i != min_hash) { min_i = false; }\n\n if (hash_i != max_hash) { max_i = false; }\n\n int row_start = A_offsets[i];\n int row_end = A_offsets[i + 1];\n\n for (int r = row_start; r < row_end; r++)\n {\n int j = A_column_indices[r];\n\n if (j >= num_rows) { continue; }\n\n int min_hash_j = min_hash_array[j];\n int max_hash_j = max_hash_array[j];\n // There is a neighbour that has a uncolored neighbour with larger hash\n max_i = (max_hash_j > max_hash && max_hash_j != INT_MIN) ? false : max_i;\n // There is a neighbour that has a uncolored neighbour with smaller hash\n min_i = (min_hash_j < min_hash && min_hash_j != INT_MAX) ? false : min_i;\n }\n\n if (max_i && min_i)\n {\n my_row_color = (i % 2) ? current_color : current_color + 1;\n }\n else if (max_i)\n {\n my_row_color = current_color;\n }\n else if (min_i)\n {\n my_row_color = current_color + 1;\n }\n\n if (my_row_color != 0)\n {\n row_colors[i] = my_row_color;\n }\n }\n }\n}\n\n\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate\nMinMaxMatrixColoringBase::MinMaxMatrixColoringBase(AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n if (cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\"))\n {\n m_uncolored_fraction = 0;\n }\n else\n {\n m_uncolored_fraction = cfg.AMG_Config::template getParameter(\"max_uncolored_percentage\", cfg_scope);\n }\n}\n\n\n\n// Block version\ntemplate \nvoid MinMaxMatrixColoring >::colorMatrixOneRing(Matrix_d &A)\n{\n const int num_rows = A.get_num_rows();\n int max_uncolored_rows = (int) (this->m_uncolored_fraction * ((ValueType) num_rows));\n IndexType *row_colors_ptr = this->m_row_colors.raw();\n const int threads_per_block = 256;\n const int warps_per_block = threads_per_block / 32;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_rows + warps_per_block - 1) / warps_per_block );\n this->m_num_colors = 1;\n thrust_wrapper::fill(this->m_row_colors.begin(), this->m_row_colors.end(), 0);\n cudaCheckError();\n\n for ( int num_uncolored = num_rows; num_uncolored > max_uncolored_rows ; )\n {\n#ifdef USE_EXPERIMENTAL_MIN_MAX\n colorRowsKernel <<< num_blocks, threads_per_block, 0, amgx::thrust::global_thread_handle::get_stream()>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n row_colors_ptr,\n this->m_num_colors,\n num_rows);\n cudaCheckError();\n#else\n colorRowsKernel <<< num_blocks, threads_per_block, 0, amgx::thrust::global_thread_handle::get_stream()>>>(A.row_offsets.raw(), A.col_indices.raw(), row_colors_ptr, this->m_num_colors, num_rows);\n cudaCheckError();\n#endif\n cudaCheckError();\n this->m_num_colors += 2;\n num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_zero() );\n cudaCheckError();\n }\n\n this->m_num_colors = thrust_wrapper::reduce( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, 0, amgx::thrust::maximum() ) + 1;\n cudaCheckError();\n}\n\n// Block version\ntemplate \nvoid MinMaxMatrixColoringBase::colorMatrix(Matrix &A)\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize(A.row_offsets.size() - 1);\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n if (this->m_coloring_level == 0)\n {\n FatalError(\"Calling coloring scheme but coloring level==0\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n }\n else if (this->m_coloring_level == 1)\n {\n this->colorMatrixOneRing(A);\n }\n else if (this->m_coloring_level == 2)\n {\n this->colorMatrixTwoRing(A);\n }\n else\n {\n FatalError(\"Min-max coloring does not support coloring_level > 2\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n }\n\n A.setView(oldView);\n}\n\ntemplate \nvoid MinMaxMatrixColoring >::colorMatrixTwoRing(Matrix_d &A)\n{\n const int num_rows = A.get_num_rows();\n const int max_uncolored_rows = (int) (this->m_uncolored_fraction * ((ValueType) num_rows));\n const int threads_per_block = 256;\n#ifdef USE_EXPERIMENTAL_NEIGHBORS\n const int half_warps_per_block = threads_per_block / 16;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_rows + half_warps_per_block - 1) / half_warps_per_block );\n#else\n const int warps_per_block = threads_per_block / 32;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_rows + warps_per_block - 1) / warps_per_block );\n#endif\n this->m_num_colors = 1;\n thrust_wrapper::fill(this->m_row_colors.begin(), this->m_row_colors.end(), 0);\n cudaCheckError();\n IVector max_hash_array(num_rows);\n IVector min_hash_array(num_rows);\n float avg_nnz_per_row = A.get_num_nz() / float(num_rows);\n\n for ( int num_uncolored = num_rows; num_uncolored > max_uncolored_rows ; )\n {\n // Each vertex checks its neighbours and store the max and min values of its neighbours\n#ifdef USE_EXPERIMENTAL_NEIGHBORS\n if ( avg_nnz_per_row < 8.0f )\n {\n find_min_max_neighbors_kernel <<< num_blocks, threads_per_block>>>( A.row_offsets.raw(), A.col_indices.raw(), this->m_row_colors.raw(), max_hash_array.raw(), min_hash_array.raw(), this->m_num_colors, num_rows );\n cudaCheckError();\n }\n else if ( avg_nnz_per_row < 16.0f )\n {\n find_min_max_neighbors_kernel <<< num_blocks, threads_per_block>>>( A.row_offsets.raw(), A.col_indices.raw(), this->m_row_colors.raw(), max_hash_array.raw(), min_hash_array.raw(), this->m_num_colors, num_rows );\n cudaCheckError();\n }\n else // if( avg_nnz_per_row < 32.0f )\n {\n find_min_max_neighbors_kernel <<< num_blocks, threads_per_block>>>( A.row_offsets.raw(), A.col_indices.raw(), this->m_row_colors.raw(), max_hash_array.raw(), min_hash_array.raw(), this->m_num_colors, num_rows );\n cudaCheckError();\n }\n\n#else\n FindMaxMinNeighboursKernel <<< num_blocks, threads_per_block>>>(A.row_offsets.raw(), A.col_indices.raw(), this->m_row_colors.raw(), max_hash_array.raw(), min_hash_array.raw(), this->m_num_colors, num_rows);\n cudaCheckError();\n#endif\n // Each vertex checks if its still the min or max\n colorRowsRingTwoKernel <<< num_blocks, threads_per_block>>>(A.row_offsets.raw(), A.col_indices.raw(), this->m_row_colors.raw(), max_hash_array.raw(), min_hash_array.raw(), this->m_num_colors, num_rows);\n cudaCheckError();\n this->m_num_colors += 2;\n num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_zero() );\n cudaCheckError();\n }\n\n this->m_num_colors = thrust_wrapper::reduce( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, 0, amgx::thrust::maximum() ) + 1;\n cudaCheckError();\n#if 0\n std::cout << \"Num colors=\" << this->m_num_colors << std::endl;\n // Print coloring statistics\n //for (int i=0;im_num_colors;i++)\n // std::cout << \"color= \" << i << \", num_rows = \" << A.offsets_rows_per_color[i+1] - A.offsets_rows_per_color[i] << std::endl;\n#endif\n}\n\ntemplate \nvoid MinMaxMatrixColoring >::colorMatrixOneRing(Matrix_h &A)\n{\n FatalError(\"Haven't implemented MinMax matrix coloring for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MinMaxMatrixColoring >::colorMatrixTwoRing(Matrix_h &A)\n{\n FatalError(\"Haven't implemented MinMax matrix coloring for host format\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n\n#define AMGX_CASE_LINE(CASE) template class MinMaxMatrixColoringBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class MinMaxMatrixColoring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n\n} // end namespace amgx\n\n", "source": "/kaggle/working/amgx/src/matrix_coloring/min_max.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n// ---------------------------\n// Kernels\n// ---------------------------\n\n// Kernel to color the rows of the matrix, using min-max approach\ntemplate \n__global__\nvoid colorRowsKernel(IndexType *row_colors, const int num_colors, const int num_rows)\n{\n for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < num_rows; i += gridDim.x * blockDim.x)\n {\n row_colors[i] = i % num_colors;\n }\n}\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate\nRoundRobinMatrixColoringBase::RoundRobinMatrixColoringBase(AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n if (cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\"))\n {\n FatalError(\"Current implementation of the round-robin coloring does not permit an exact coloring, and therefore cannot lead to deterministic results. Implementation of a deterministic round-robin coloring algorithm still pending\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n this->m_num_colors = cfg.AMG_Config::template getParameter(\"num_colors\", cfg_scope);\n}\n\ntemplate\nvoid RoundRobinMatrixColoringBase::colorMatrix(Matrix &A)\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize(A.row_offsets.size() - 1, 0);\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n if (this->m_coloring_level == 0)\n {\n FatalError(\"Callingn coloring scheme but coloring level==0\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else if (this->m_coloring_level == 1)\n {\n this->colorMatrixOneRing(A);\n }\n else\n {\n FatalError(\"Round robin coloring algorithm can only do one ring coloring\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n A.setView(oldView);\n}\n\n// Block version\ntemplate \nvoid RoundRobinMatrixColoring >::colorMatrixOneRing(Matrix_d &A)\n{\n // One thread per row\n const int num_rows = A.get_num_rows();\n IndexType *row_colors_ptr = this->m_row_colors.raw();\n const int threads_per_block = 64;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows - 1) / threads_per_block + 1);\n colorRowsKernel <<< num_blocks, threads_per_block>>>(row_colors_ptr, this->m_num_colors, num_rows);\n cudaCheckError();\n /*\n // Sort the vertices based on their color\n A.sorted_rows_by_color.resize(num_owned_rows);\n\n // Copy row colors\n IVector row_colors(num_owned_rows);\n amgx::thrust::copy(A.row_colors.begin(), A.row_colors.begin()+num_owned_rows, row_colors.begin());\n\n thrust_wrapper::sequence(A.sorted_rows_by_color.begin(),A.sorted_rows_by_color.end());\n amgx::thrust::sort_by_key(row_colors.begin(),row_colors.end(),A.sorted_rows_by_color.begin());\n cudaCheckError();\n\n // Compute the offset for each color\n IVector d_offsets_rows_per_color(A.num_colors+1);\n A.offsets_rows_per_color.resize(A.num_colors+1);\n\n amgx::thrust::lower_bound(row_colors.begin(),\n row_colors.end(),\n amgx::thrust::counting_iterator(0),\n amgx::thrust::counting_iterator(d_offsets_rows_per_color.size()),\n d_offsets_rows_per_color.begin());\n\n cudaCheckError();\n\n // Copy from device to host\n A.offsets_rows_per_color = d_offsets_rows_per_color;\n */\n}\n\ntemplate \nvoid RoundRobinMatrixColoring >::colorMatrixOneRing(Matrix_h &A)\n{\n FatalError(\"Haven't implemented RoundRobin matrix coloring for host\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n#define AMGX_CASE_LINE(CASE) template class RoundRobinMatrixColoringBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class RoundRobinMatrixColoring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // end namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n// ---------------------------\n// Kernels\n// ---------------------------\n\n// Kernel to color the rows of the matrix, using min-max approach\ntemplate \n__global__\nvoid colorRowsKernel(IndexType *row_colors, const int num_colors, const int num_rows)\n{\n for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < num_rows; i += gridDim.x * blockDim.x)\n {\n row_colors[i] = i % num_colors;\n }\n}\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate\nRoundRobinMatrixColoringBase::RoundRobinMatrixColoringBase(AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n if (cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\"))\n {\n FatalError(\"Current implementation of the round-robin coloring does not permit an exact coloring, and therefore cannot lead to deterministic results. Implementation of a deterministic round-robin coloring algorithm still pending\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n this->m_num_colors = cfg.AMG_Config::template getParameter(\"num_colors\", cfg_scope);\n}\n\ntemplate\nvoid RoundRobinMatrixColoringBase::colorMatrix(Matrix &A)\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize(A.row_offsets.size() - 1, 0);\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n if (this->m_coloring_level == 0)\n {\n FatalError(\"Callingn coloring scheme but coloring level==0\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else if (this->m_coloring_level == 1)\n {\n this->colorMatrixOneRing(A);\n }\n else\n {\n FatalError(\"Round robin coloring algorithm can only do one ring coloring\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n A.setView(oldView);\n}\n\n// Block version\ntemplate \nvoid RoundRobinMatrixColoring >::colorMatrixOneRing(Matrix_d &A)\n{\n // One thread per row\n const int num_rows = A.get_num_rows();\n IndexType *row_colors_ptr = this->m_row_colors.raw();\n const int threads_per_block = 64;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (int) (num_rows - 1) / threads_per_block + 1);\n colorRowsKernel <<< num_blocks, threads_per_block>>>(row_colors_ptr, this->m_num_colors, num_rows);\n cudaCheckError();\n /*\n // Sort the vertices based on their color\n A.sorted_rows_by_color.resize(num_owned_rows);\n\n // Copy row colors\n IVector row_colors(num_owned_rows);\n amgx::thrust::copy(A.row_colors.begin(), A.row_colors.begin()+num_owned_rows, row_colors.begin());\n\n thrust_wrapper::sequence(A.sorted_rows_by_color.begin(),A.sorted_rows_by_color.end());\n amgx::thrust::sort_by_key(row_colors.begin(),row_colors.end(),A.sorted_rows_by_color.begin());\n cudaCheckError();\n\n // Compute the offset for each color\n IVector d_offsets_rows_per_color(A.num_colors+1);\n A.offsets_rows_per_color.resize(A.num_colors+1);\n\n amgx::thrust::lower_bound(row_colors.begin(),\n row_colors.end(),\n amgx::thrust::counting_iterator(0),\n amgx::thrust::counting_iterator(d_offsets_rows_per_color.size()),\n d_offsets_rows_per_color.begin());\n\n cudaCheckError();\n\n // Copy from device to host\n A.offsets_rows_per_color = d_offsets_rows_per_color;\n */\n}\n\ntemplate \nvoid RoundRobinMatrixColoring >::colorMatrixOneRing(Matrix_h &A)\n{\n FatalError(\"Haven't implemented RoundRobin matrix coloring for host\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n#define AMGX_CASE_LINE(CASE) template class RoundRobinMatrixColoringBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class RoundRobinMatrixColoring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // end namespace amgx\n\n", "source": "/kaggle/working/amgx/src/matrix_coloring/round_robin.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#define COLORING_DEBUG 1\n\n\n// Pseudo-random number generator\nnamespace amgx\n{\n\nstatic __host__ __device__ unsigned int hash_function(unsigned int a, unsigned int seed, unsigned int rows = 0)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n\nstruct is_zero\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x == 0;\n }\n};\n\n// ---------------------------\n// Kernels\n// ---------------------------\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid count_gtlt_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict A_colors,\n int *A_gtlt_count,\n const int seed)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 )\n {\n if ( lane_id == 0 )\n {\n A_gtlt_count[row_id] = -1;\n }\n\n continue;\n }\n\n // Hash my row id.\n int row_hash = hash_function( row_id, seed );\n // The number of vertices that are greater/smaller than me.\n int gt_count = 0, lt_count = 0;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, seed );\n // Get the color of the column.\n int col_color = -1;\n\n if ( row_it < row_end && col_id < A_num_rows)\n {\n col_color = A_colors[col_id];\n }\n\n // Threads determine if they are greater than the row hash.\n int gt_pred = col_color == 0 && col_hash > row_hash;\n int lt_pred = col_color == 0 && col_hash < row_hash;\n // Count greater/smaller neighbors.\n gt_count += __popc( utils::ballot( gt_pred ) );\n lt_count += __popc( utils::ballot( lt_pred ) );\n }\n\n // The warp leader stores the result.\n int my_gtlt_count = (gt_count << 16) | lt_count;\n\n if ( lane_id == 0 )\n {\n A_gtlt_count[row_id] = my_gtlt_count;\n }\n }\n}\n\n\ntemplate< int CTA_SIZE, int WARP_SIZE, bool LATE_REJECTION >\n__global__\nvoid color_kernel( const int A_num_rows, const int *A_rows, const int *A_cols, const int *A_gtlt_count, const int current_color, const int weakness_bound, int *A_colors )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 ) // Already colored!!!\n {\n continue;\n }\n\n // The number of vertices that are greater/smaller than me.\n int row_gtlt_count = A_gtlt_count[row_id];\n // Split gtlt_count into 2.\n int row_gt_count = row_gtlt_count >> 16;\n int row_lt_count = row_gtlt_count & 0xffff;\n\n // Min-max algorithm.\n if ( row_gt_count == 0 && row_lt_count == 0 )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = current_color + (row_id & 1);\n }\n\n continue;\n }\n\n if ( row_gt_count == 0 )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = current_color;\n }\n\n continue;\n }\n\n if ( row_lt_count == 0 )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = current_color + 1;\n }\n\n continue;\n }\n\n // Do we skip it.\n int candidate = 1;\n // Predicates. Is a vertex min/max.\n int is_max_vertex = row_gt_count <= weakness_bound;\n int is_min_vertex = row_lt_count <= weakness_bound;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n // Get the color of the column (it could help late rejection).\n if ( LATE_REJECTION )\n {\n int col_color = -1;\n\n if ( row_it < row_end && col_id < A_num_rows)\n {\n col_color = A_colors[col_id];\n }\n\n // Late rejection test.\n if ( col_color == current_color )\n {\n is_max_vertex = 0;\n }\n\n if ( col_color == current_color + 1 )\n {\n is_min_vertex = 0;\n }\n }\n\n // Get the gt/lt count.\n int col_gtlt_count = -1;\n\n if ( row_it < row_end && col_id < A_num_rows )\n {\n col_gtlt_count = A_gtlt_count[col_id];\n }\n\n // Split gtlt_count into 2. col_gt/lt_count == -1 if already colored.\n int col_gt_count = col_gtlt_count >> 16;\n int col_lt_count = col_gtlt_count & 0xffff;\n // Threads determine if they are greater than the row hash.\n\n if ( col_gtlt_count != -1 )\n {\n is_max_vertex &= col_gt_count > row_gt_count || (col_gt_count == row_gt_count && col_id <= row_id);\n is_min_vertex &= col_lt_count > row_lt_count || (col_lt_count == row_lt_count && col_id >= row_id);\n }\n }\n\n // The warp leader stores the result.\n if ( candidate && utils::all( is_max_vertex ) )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = current_color;\n }\n\n continue;\n }\n\n if ( candidate && utils::all( is_min_vertex ) )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = current_color + 1;\n }\n\n continue;\n }\n }\n}\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid\ndbg_check_coloring_kernel( const int A_num_rows, const int *A_rows, const int *A_cols, const int *A_colors, const int *A_gtlt_count, int *error_found )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = A_colors[row_id];\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n // Get the color of the column.\n int col_color = -1;\n\n if ( row_it < row_end && col_id < A_num_rows)\n {\n col_color = A_colors[col_id];\n }\n\n // Is there something wrong ??\n if ( row_id != col_id && row_color == col_color && row_color != 0)\n {\n if ( A_gtlt_count != NULL && !error_found[0] )\n {\n //printf( \"row_id=%d, row_color=%d, col_id=%d, col_color=%d\\n\", row_id, row_color, col_id, col_color );\n /*\n // The number of vertices that are greater/smaller than me.\n int row_gtlt_count = A_gtlt_count[row_id];\n int row_gt_count = row_gtlt_count >> 16;\n int row_lt_count = row_gtlt_count & 0xffff;\n printf( \"row_gt_count=%d, row_gt_count=%d\\n\", row_gt_count, row_lt_count );\n\n int col_gtlt_count = A_gtlt_count[col_id];\n int col_gt_count = col_gtlt_count >> 16;\n int col_lt_count = col_gtlt_count & 0xffff;\n printf( \"col_gt_count=%d, col_gt_count=%d\\n\", col_gt_count, col_lt_count ); */\n }\n\n error_found[0] = 1;\n atomicAdd(error_found + 1, 1);\n }\n }\n }\n}\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate< class T_Config >\nMin_Max_2Ring_Matrix_Coloring_Base::Min_Max_2Ring_Matrix_Coloring_Base( AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n if ( this->m_coloring_level != 1 )\n {\n FatalError( \"Not implemented for coloring_level != 1\", AMGX_ERR_NOT_SUPPORTED_TARGET );\n }\n\n if ( cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\") )\n {\n m_uncolored_fraction = 0.0;\n }\n else\n {\n m_uncolored_fraction = cfg.AMG_Config::template getParameter(\"max_uncolored_percentage\", cfg_scope);\n }\n\n m_weakness_bound = cfg.AMG_Config::template getParameter( \"weakness_bound\", cfg_scope );\n m_late_rejection = cfg.AMG_Config::template getParameter( \"late_rejection\", cfg_scope ) != 0;\n}\n\n\n\n#if !NEW_COLORER_TESTS\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nMin_Max_2Ring_Matrix_Coloring >::colorMatrix( Matrix_d &A )\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize( A.row_offsets.size() - 1 );\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n const int num_rows = A.get_num_rows();\n const int max_uncolored_rows = static_cast( this->m_uncolored_fraction * num_rows );\n const int CTA_SIZE = 128;\n const int NUM_WARPS_PER_CTA = CTA_SIZE / 32;\n const int GRID_SIZE = std::min( 2048, (num_rows + NUM_WARPS_PER_CTA - 1) / NUM_WARPS_PER_CTA );\n this->m_num_colors = 1;\n thrust_wrapper::fill( this->m_row_colors.begin(), this->m_row_colors.end(), 0 );\n cudaCheckError();\n device_vector_alloc gtlt_count( num_rows );\n\n for ( int num_uncolored = num_rows ; num_uncolored > max_uncolored_rows ; )\n {\n count_gtlt_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ), 0);\n cudaCheckError();\n\n if ( this->m_late_rejection )\n color_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n this->m_num_colors,\n this->m_weakness_bound,\n this->m_row_colors.raw() );\n else\n color_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n this->m_num_colors,\n this->m_weakness_bound,\n this->m_row_colors.raw() );\n\n cudaCheckError();\n#if 0\n device_vector_alloc error_found( 1, 0 );\n dbg_check_coloring_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n amgx::thrust::raw_pointer_cast( &error_found.front() ) );\n cudaCheckError();\n#endif\n this->m_num_colors += 2;\n num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_zero() );\n cudaCheckError();\n }\n\n this->m_num_colors = thrust_wrapper::reduce( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, 0, amgx::thrust::maximum() ) + 1;\n cudaCheckError();\n#if 0\n device_vector_alloc error_found( 1, 0 );\n dbg_check_coloring_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n NULL,\n amgx::thrust::raw_pointer_cast( &error_found.front() ) );\n cudaCheckError();\n\n if ( error_found[0] != 0 )\n {\n std::cout << \"INVALID COLORING !!! Two neighbors have the same color!!!\" << std::endl;\n }\n\n#endif\n A.setView(oldView);\n}\n\n\n#else\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid\ndbg_coloring_histogram_kernel( int *colors_count, const int A_num_rows, const int *A_colors, const int n_colors )\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n __shared__ unsigned int color_counts[256];\n\n if (threadIdx.x < n_colors)\n {\n color_counts[threadIdx.x] = 0;\n }\n\n __syncthreads();\n\n for (; tid < A_num_rows; tid += blockDim.x * gridDim.x)\n {\n int col = A_colors[tid];\n atomicAdd(color_counts + col, 1);\n }\n\n __syncthreads();\n\n if (threadIdx.x < n_colors)\n {\n atomicAdd(colors_count + threadIdx.x, color_counts[threadIdx.x]);\n }\n}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nMin_Max_2Ring_Matrix_Coloring >::debug_coloring(Matrix_d &A, int step)\n{\n#if COLORING_DEBUG\n const int num_rows = A.get_num_rows();\n const int CTA_SIZE = 128;\n const int NUM_WARPS_PER_CTA = CTA_SIZE / 32;\n const int GRID_SIZE = std::min( 2048, (num_rows + NUM_WARPS_PER_CTA - 1) / NUM_WARPS_PER_CTA );\n int num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_zero() );\n cudaCheckError();\n int maxr = A.row_offsets[1] - A.row_offsets[0];\n /*for(int i=2;i<=num_rows;i++)\n {\n int d=A.row_offsets[i]-A.row_offsets[i-1];\n if(d>maxr)\n {\n maxr=d;\n }\n }*/\n device_vector_alloc error_found( 2, 0 );\n error_found[0] = 0;\n error_found[1] = 0;\n dbg_check_coloring_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n 0,//amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n amgx::thrust::raw_pointer_cast( &error_found.front() ) );\n cudaCheckError();\n {\n device_vector_alloc color_histogram(this->m_num_colors + 1);\n\n for (int i = 0; i < color_histogram.size(); i++)\n {\n color_histogram[i] = 0;\n }\n\n dbg_coloring_histogram_kernel <<< GRID_SIZE, CTA_SIZE>>>( amgx::thrust::raw_pointer_cast(&color_histogram.front()), num_rows, this->m_row_colors.raw(), this->m_num_colors + 1);\n cudaCheckError();\n\n for (int i = 0; i < color_histogram.size(); i++)\n {\n std::cout << step << \"\\t\" << \"H[\" << i << \"] = \" << color_histogram[i] << std::endl;\n }\n\n std::cout << step << \"\\t\" << \"Errors=\" << error_found[1] << std::endl;\n std::cout << step << \"\\t\" << \"Uncolored=\" << num_uncolored << std::endl;\n std::cout << step << \"\\t\" << \"Num colors=\" << this->m_num_colors << \"/\" << maxr << std::endl;\n }\n#endif\n}\n\ntemplate< int CTA_SIZE, int WARP_SIZE, bool LATE_REJECTION >\n__global__\nvoid color_kernel_greedy( const int A_num_rows, const int *A_rows, const int *A_cols, const int *A_gtlt_count, const int current_color, const int weakness_bound, int *A_colors )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n int dice = hash_function(row_id, 19881988, 0);\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 ) // Already colored!!!\n {\n continue;\n }\n\n // The number of vertices that are greater/smaller than me.\n int row_gtlt_count = A_gtlt_count[row_id];\n // Split gtlt_count into 2.\n int row_gt_count = row_gtlt_count >> 16;\n int row_lt_count = row_gtlt_count & 0xffff;\n // Do we skip it.\n int candidate = 1;\n // Predicates. Is a vertex min/max.\n int is_max_vertex = true;//row_gt_count <= weakness_bound;\n int is_min_vertex = true;//row_lt_count <= weakness_bound;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n unsigned long long used_colors = 0ull;\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n if ( row_it < row_end && col_id < A_num_rows )\n {\n int col_color = A_colors[col_id];\n\n if ( col_color > 0 )\n {\n used_colors |= (1ull << (64 - col_color));\n }\n }\n\n // Get the gt/lt count.\n int col_gtlt_count = -1;\n\n if ( row_it < row_end && col_id < A_num_rows )\n {\n col_gtlt_count = A_gtlt_count[col_id];\n }\n\n // Split gtlt_count into 2. col_gt/lt_count == -1 if already colored.\n int col_gt_count = col_gtlt_count >> 16;\n int col_lt_count = col_gtlt_count & 0xffff;\n\n // Threads determine if they are greater than the row hash.\n //if( col_gt_count != -1 )\n // is_max_vertex &= col_gt_count < row_gt_count || (col_gt_count == row_gt_count && col_id < row_id);\n //if( col_lt_count != -1 )\n // is_min_vertex &= col_lt_count /*>*/ < row_lt_count || (col_lt_count == row_lt_count && col_id > row_id);\n if ( col_gtlt_count != -1 )\n {\n is_max_vertex &= col_gt_count > row_gt_count || (col_gt_count == row_gt_count && col_id >= row_id);\n is_min_vertex &= col_lt_count > row_lt_count || (col_lt_count == row_lt_count && col_id >= row_id);\n }\n }\n\n is_min_vertex = false;\n //reduce used colors bit by bit.\n#pragma unroll\n for (int i = WARP_SIZE / 2; i >= 1; i /= 2)\n {\n int tmp_hi = __double2hiint( __longlong_as_double( used_colors ) );\n int tmp_lo = __double2hiint( __longlong_as_double( used_colors ) );\n tmp_hi = utils::shfl_xor(tmp_hi, i, WARP_SIZE);\n tmp_lo = utils::shfl_xor(tmp_lo, i, WARP_SIZE);\n long long tmp = __double_as_longlong(__hiloint2double(tmp_hi, tmp_lo));\n used_colors |= tmp;\n }\n\n int my_color = 64 - utils::bfind( ~used_colors );\n\n if (my_color <= 0) { my_color = 1; }\n\n /*int sets=0;\n for(int c=1; c<=64; c++)\n {\n unsigned long long int b = (1ull << (64-c));\n if((~used_colors & b) && sets < (dice % 3)-1)\n {\n sets++;\n my_color = c;\n }\n }*/\n\n // The warp leader stores the result.\n if ( candidate && utils::all( is_max_vertex ) )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = my_color;\n }\n\n continue;\n }\n\n if ( candidate && utils::all( is_min_vertex ) )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = my_color + 1;\n }\n\n continue;\n }\n }\n}\n\n\ntemplate< int CTA_SIZE, int WARP_SIZE>\n__global__\nvoid color_kernel_reassign_tail_thread( const int A_num_rows, const int *A_rows, const int *A_cols, int num_colors, int first_tail_color, int *A_colors )\n{\n const int NUM_THREADS_PER_GRID = gridDim.x * CTA_SIZE;\n // Row identifier.\n int row_id = blockIdx.x * CTA_SIZE + threadIdx.x;\n\n // Iterate over rows.\n for ( ; row_id < A_num_rows ; row_id += NUM_THREADS_PER_GRID )\n {\n unsigned long long used_colors = 0ull;\n int row_color = A_colors[row_id];\n\n if ( row_color < first_tail_color ) // Already colored!!!\n {\n continue;\n }\n\n int row_hash = hash_function(row_id, 198719871987L, 0);\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( int row_it = row_begin; row_it < row_end ; row_it++)\n {\n int col_id = A_cols[row_it];\n\n if (col_id >= A_num_rows) { continue; }\n\n int col_hash = hash_function( col_id, 1987, 0 );\n int col_color = A_colors[col_id];\n }\n }\n}\n\n__global__ void unassign_color(int *A_colors, int num_rows)\n{\n for (int tid = threadIdx.x + blockIdx.x * blockDim.x; tid < num_rows; tid += gridDim.x * blockDim.x)\n {\n if (tid % 9 == 0)\n {\n A_colors[tid] = 0;\n }\n }\n}\n\n\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nMin_Max_2Ring_Matrix_Coloring >::color_step( Matrix_d &A, int seed )\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize( A.row_offsets.size() - 1 );\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n const int num_rows = A.get_num_rows();\n const int max_uncolored_rows = static_cast( this->m_uncolored_fraction * num_rows );\n const int CTA_SIZE = 128;\n const int NUM_WARPS_PER_CTA = CTA_SIZE / 32;\n const int GRID_SIZE = std::min( 2048, (num_rows + NUM_WARPS_PER_CTA - 1) / NUM_WARPS_PER_CTA );\n device_vector_alloc gtlt_count( num_rows );\n\n for ( int num_uncolored = num_rows ; num_uncolored > max_uncolored_rows ; )\n {\n count_gtlt_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n seed);\n cudaCheckError();\n\n if ( this->m_late_rejection )\n color_kernel_greedy <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n this->m_num_colors,\n this->m_weakness_bound,\n this->m_row_colors.raw() );\n else\n color_kernel_greedy <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n this->m_num_colors,\n this->m_weakness_bound,\n this->m_row_colors.raw() );\n\n cudaCheckError();\n this->m_num_colors += 2;\n num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_zero() );\n cudaCheckError();\n }\n\n this->m_num_colors = thrust_wrapper::reduce( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, 0, amgx::thrust::maximum() ) + 1;\n}\n\n\n\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nMin_Max_2Ring_Matrix_Coloring >::color_matrix( Matrix_d &A )\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize( A.row_offsets.size() - 1 );\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n this->m_num_colors = 1;\n thrust_wrapper::fill( this->m_row_colors.begin(), this->m_row_colors.end(), 0 );\n cudaCheckError();\n color_step(A, 0);\n\n A.setView(oldView);\n}\n\n\n#endif\n\n\n#define AMGX_CASE_LINE(CASE) template class Min_Max_2Ring_Matrix_Coloring_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Min_Max_2Ring_Matrix_Coloring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // end namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#define COLORING_DEBUG 1\n\n\n// Pseudo-random number generator\nnamespace amgx\n{\n\nstatic __host__ __device__ unsigned int hash_function(unsigned int a, unsigned int seed, unsigned int rows = 0)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n\nstruct is_zero\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x == 0;\n }\n};\n\n// ---------------------------\n// Kernels\n// ---------------------------\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid count_gtlt_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict A_colors,\n int *A_gtlt_count,\n const int seed)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 )\n {\n if ( lane_id == 0 )\n {\n A_gtlt_count[row_id] = -1;\n }\n\n continue;\n }\n\n // Hash my row id.\n int row_hash = hash_function( row_id, seed );\n // The number of vertices that are greater/smaller than me.\n int gt_count = 0, lt_count = 0;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n // Each thread hashes its column id.\n int col_hash = hash_function( col_id, seed );\n // Get the color of the column.\n int col_color = -1;\n\n if ( row_it < row_end && col_id < A_num_rows)\n {\n col_color = A_colors[col_id];\n }\n\n // Threads determine if they are greater than the row hash.\n int gt_pred = col_color == 0 && col_hash > row_hash;\n int lt_pred = col_color == 0 && col_hash < row_hash;\n // Count greater/smaller neighbors.\n gt_count += __popc( utils::ballot( gt_pred ) );\n lt_count += __popc( utils::ballot( lt_pred ) );\n }\n\n // The warp leader stores the result.\n int my_gtlt_count = (gt_count << 16) | lt_count;\n\n if ( lane_id == 0 )\n {\n A_gtlt_count[row_id] = my_gtlt_count;\n }\n }\n}\n\n\ntemplate< int CTA_SIZE, int WARP_SIZE, bool LATE_REJECTION >\n__global__\nvoid color_kernel( const int A_num_rows, const int *A_rows, const int *A_cols, const int *A_gtlt_count, const int current_color, const int weakness_bound, int *A_colors )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 ) // Already colored!!!\n {\n continue;\n }\n\n // The number of vertices that are greater/smaller than me.\n int row_gtlt_count = A_gtlt_count[row_id];\n // Split gtlt_count into 2.\n int row_gt_count = row_gtlt_count >> 16;\n int row_lt_count = row_gtlt_count & 0xffff;\n\n // Min-max algorithm.\n if ( row_gt_count == 0 && row_lt_count == 0 )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = current_color + (row_id & 1);\n }\n\n continue;\n }\n\n if ( row_gt_count == 0 )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = current_color;\n }\n\n continue;\n }\n\n if ( row_lt_count == 0 )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = current_color + 1;\n }\n\n continue;\n }\n\n // Do we skip it.\n int candidate = 1;\n // Predicates. Is a vertex min/max.\n int is_max_vertex = row_gt_count <= weakness_bound;\n int is_min_vertex = row_lt_count <= weakness_bound;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n // Get the color of the column (it could help late rejection).\n if ( LATE_REJECTION )\n {\n int col_color = -1;\n\n if ( row_it < row_end && col_id < A_num_rows)\n {\n col_color = A_colors[col_id];\n }\n\n // Late rejection test.\n if ( col_color == current_color )\n {\n is_max_vertex = 0;\n }\n\n if ( col_color == current_color + 1 )\n {\n is_min_vertex = 0;\n }\n }\n\n // Get the gt/lt count.\n int col_gtlt_count = -1;\n\n if ( row_it < row_end && col_id < A_num_rows )\n {\n col_gtlt_count = A_gtlt_count[col_id];\n }\n\n // Split gtlt_count into 2. col_gt/lt_count == -1 if already colored.\n int col_gt_count = col_gtlt_count >> 16;\n int col_lt_count = col_gtlt_count & 0xffff;\n // Threads determine if they are greater than the row hash.\n\n if ( col_gtlt_count != -1 )\n {\n is_max_vertex &= col_gt_count > row_gt_count || (col_gt_count == row_gt_count && col_id <= row_id);\n is_min_vertex &= col_lt_count > row_lt_count || (col_lt_count == row_lt_count && col_id >= row_id);\n }\n }\n\n // The warp leader stores the result.\n if ( candidate && utils::all( is_max_vertex ) )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = current_color;\n }\n\n continue;\n }\n\n if ( candidate && utils::all( is_min_vertex ) )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = current_color + 1;\n }\n\n continue;\n }\n }\n}\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid\ndbg_check_coloring_kernel( const int A_num_rows, const int *A_rows, const int *A_cols, const int *A_colors, const int *A_gtlt_count, int *error_found )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n int row_color = A_colors[row_id];\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n // Get the color of the column.\n int col_color = -1;\n\n if ( row_it < row_end && col_id < A_num_rows)\n {\n col_color = A_colors[col_id];\n }\n\n // Is there something wrong ??\n if ( row_id != col_id && row_color == col_color && row_color != 0)\n {\n if ( A_gtlt_count != NULL && !error_found[0] )\n {\n //printf( \"row_id=%d, row_color=%d, col_id=%d, col_color=%d\\n\", row_id, row_color, col_id, col_color );\n /*\n // The number of vertices that are greater/smaller than me.\n int row_gtlt_count = A_gtlt_count[row_id];\n int row_gt_count = row_gtlt_count >> 16;\n int row_lt_count = row_gtlt_count & 0xffff;\n printf( \"row_gt_count=%d, row_gt_count=%d\\n\", row_gt_count, row_lt_count );\n\n int col_gtlt_count = A_gtlt_count[col_id];\n int col_gt_count = col_gtlt_count >> 16;\n int col_lt_count = col_gtlt_count & 0xffff;\n printf( \"col_gt_count=%d, col_gt_count=%d\\n\", col_gt_count, col_lt_count ); */\n }\n\n error_found[0] = 1;\n atomicAdd(error_found + 1, 1);\n }\n }\n }\n}\n\n// ---------------------------\n// Methods\n// ---------------------------\n\ntemplate< class T_Config >\nMin_Max_2Ring_Matrix_Coloring_Base::Min_Max_2Ring_Matrix_Coloring_Base( AMG_Config &cfg, const std::string &cfg_scope) : MatrixColoring(cfg, cfg_scope)\n{\n if ( this->m_coloring_level != 1 )\n {\n FatalError( \"Not implemented for coloring_level != 1\", AMGX_ERR_NOT_SUPPORTED_TARGET );\n }\n\n if ( cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\") )\n {\n m_uncolored_fraction = 0.0;\n }\n else\n {\n m_uncolored_fraction = cfg.AMG_Config::template getParameter(\"max_uncolored_percentage\", cfg_scope);\n }\n\n m_weakness_bound = cfg.AMG_Config::template getParameter( \"weakness_bound\", cfg_scope );\n m_late_rejection = cfg.AMG_Config::template getParameter( \"late_rejection\", cfg_scope ) != 0;\n}\n\n\n\n#if !NEW_COLORER_TESTS\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nMin_Max_2Ring_Matrix_Coloring >::colorMatrix( Matrix_d &A )\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize( A.row_offsets.size() - 1 );\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n const int num_rows = A.get_num_rows();\n const int max_uncolored_rows = static_cast( this->m_uncolored_fraction * num_rows );\n const int CTA_SIZE = 128;\n const int NUM_WARPS_PER_CTA = CTA_SIZE / 32;\n const int GRID_SIZE = std::min( 2048, (num_rows + NUM_WARPS_PER_CTA - 1) / NUM_WARPS_PER_CTA );\n this->m_num_colors = 1;\n thrust_wrapper::fill( this->m_row_colors.begin(), this->m_row_colors.end(), 0 );\n cudaCheckError();\n device_vector_alloc gtlt_count( num_rows );\n\n for ( int num_uncolored = num_rows ; num_uncolored > max_uncolored_rows ; )\n {\n count_gtlt_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ), 0);\n cudaCheckError();\n\n if ( this->m_late_rejection )\n color_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n this->m_num_colors,\n this->m_weakness_bound,\n this->m_row_colors.raw() );\n else\n color_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n this->m_num_colors,\n this->m_weakness_bound,\n this->m_row_colors.raw() );\n\n cudaCheckError();\n#if 0\n device_vector_alloc error_found( 1, 0 );\n dbg_check_coloring_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n amgx::thrust::raw_pointer_cast( &error_found.front() ) );\n cudaCheckError();\n#endif\n this->m_num_colors += 2;\n num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_zero() );\n cudaCheckError();\n }\n\n this->m_num_colors = thrust_wrapper::reduce( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, 0, amgx::thrust::maximum() ) + 1;\n cudaCheckError();\n#if 0\n device_vector_alloc error_found( 1, 0 );\n dbg_check_coloring_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n NULL,\n amgx::thrust::raw_pointer_cast( &error_found.front() ) );\n cudaCheckError();\n\n if ( error_found[0] != 0 )\n {\n std::cout << \"INVALID COLORING !!! Two neighbors have the same color!!!\" << std::endl;\n }\n\n#endif\n A.setView(oldView);\n}\n\n\n#else\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid\ndbg_coloring_histogram_kernel( int *colors_count, const int A_num_rows, const int *A_colors, const int n_colors )\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n __shared__ unsigned int color_counts[256];\n\n if (threadIdx.x < n_colors)\n {\n color_counts[threadIdx.x] = 0;\n }\n\n __syncthreads();\n\n for (; tid < A_num_rows; tid += blockDim.x * gridDim.x)\n {\n int col = A_colors[tid];\n atomicAdd(color_counts + col, 1);\n }\n\n __syncthreads();\n\n if (threadIdx.x < n_colors)\n {\n atomicAdd(colors_count + threadIdx.x, color_counts[threadIdx.x]);\n }\n}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nMin_Max_2Ring_Matrix_Coloring >::debug_coloring(Matrix_d &A, int step)\n{\n#if COLORING_DEBUG\n const int num_rows = A.get_num_rows();\n const int CTA_SIZE = 128;\n const int NUM_WARPS_PER_CTA = CTA_SIZE / 32;\n const int GRID_SIZE = std::min( 2048, (num_rows + NUM_WARPS_PER_CTA - 1) / NUM_WARPS_PER_CTA );\n int num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_zero() );\n cudaCheckError();\n int maxr = A.row_offsets[1] - A.row_offsets[0];\n /*for(int i=2;i<=num_rows;i++)\n {\n int d=A.row_offsets[i]-A.row_offsets[i-1];\n if(d>maxr)\n {\n maxr=d;\n }\n }*/\n device_vector_alloc error_found( 2, 0 );\n error_found[0] = 0;\n error_found[1] = 0;\n dbg_check_coloring_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n 0,//amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n amgx::thrust::raw_pointer_cast( &error_found.front() ) );\n cudaCheckError();\n {\n device_vector_alloc color_histogram(this->m_num_colors + 1);\n\n for (int i = 0; i < color_histogram.size(); i++)\n {\n color_histogram[i] = 0;\n }\n\n dbg_coloring_histogram_kernel <<< GRID_SIZE, CTA_SIZE>>>( amgx::thrust::raw_pointer_cast(&color_histogram.front()), num_rows, this->m_row_colors.raw(), this->m_num_colors + 1);\n cudaCheckError();\n\n for (int i = 0; i < color_histogram.size(); i++)\n {\n std::cout << step << \"\\t\" << \"H[\" << i << \"] = \" << color_histogram[i] << std::endl;\n }\n\n std::cout << step << \"\\t\" << \"Errors=\" << error_found[1] << std::endl;\n std::cout << step << \"\\t\" << \"Uncolored=\" << num_uncolored << std::endl;\n std::cout << step << \"\\t\" << \"Num colors=\" << this->m_num_colors << \"/\" << maxr << std::endl;\n }\n#endif\n}\n\ntemplate< int CTA_SIZE, int WARP_SIZE, bool LATE_REJECTION >\n__global__\nvoid color_kernel_greedy( const int A_num_rows, const int *A_rows, const int *A_cols, const int *A_gtlt_count, const int current_color, const int weakness_bound, int *A_colors )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // Thread coordinates.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Row identifier.\n int row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Iterate over the rows of the matrix.\n for ( ; row_id < A_num_rows ; row_id += NUM_WARPS_PER_GRID )\n {\n int dice = hash_function(row_id, 19881988, 0);\n int row_color = A_colors[row_id];\n\n if ( row_color != 0 ) // Already colored!!!\n {\n continue;\n }\n\n // The number of vertices that are greater/smaller than me.\n int row_gtlt_count = A_gtlt_count[row_id];\n // Split gtlt_count into 2.\n int row_gt_count = row_gtlt_count >> 16;\n int row_lt_count = row_gtlt_count & 0xffff;\n // Do we skip it.\n int candidate = 1;\n // Predicates. Is a vertex min/max.\n int is_max_vertex = true;//row_gt_count <= weakness_bound;\n int is_min_vertex = true;//row_lt_count <= weakness_bound;\n // Iterators over my row.\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n unsigned long long used_colors = 0ull;\n\n for ( ; row_begin < row_end ; row_begin += WARP_SIZE )\n {\n // Iterator.\n int row_it = row_begin + lane_id;\n // Get the column index (if the iterator is valid).\n int col_id = -1;\n\n if ( row_it < row_end )\n {\n col_id = A_cols[row_it];\n }\n\n if ( row_it < row_end && col_id < A_num_rows )\n {\n int col_color = A_colors[col_id];\n\n if ( col_color > 0 )\n {\n used_colors |= (1ull << (64 - col_color));\n }\n }\n\n // Get the gt/lt count.\n int col_gtlt_count = -1;\n\n if ( row_it < row_end && col_id < A_num_rows )\n {\n col_gtlt_count = A_gtlt_count[col_id];\n }\n\n // Split gtlt_count into 2. col_gt/lt_count == -1 if already colored.\n int col_gt_count = col_gtlt_count >> 16;\n int col_lt_count = col_gtlt_count & 0xffff;\n\n // Threads determine if they are greater than the row hash.\n //if( col_gt_count != -1 )\n // is_max_vertex &= col_gt_count < row_gt_count || (col_gt_count == row_gt_count && col_id < row_id);\n //if( col_lt_count != -1 )\n // is_min_vertex &= col_lt_count /*>*/ < row_lt_count || (col_lt_count == row_lt_count && col_id > row_id);\n if ( col_gtlt_count != -1 )\n {\n is_max_vertex &= col_gt_count > row_gt_count || (col_gt_count == row_gt_count && col_id >= row_id);\n is_min_vertex &= col_lt_count > row_lt_count || (col_lt_count == row_lt_count && col_id >= row_id);\n }\n }\n\n is_min_vertex = false;\n //reduce used colors bit by bit.\n#pragma unroll\n for (int i = WARP_SIZE / 2; i >= 1; i /= 2)\n {\n int tmp_hi = __double2hiint( __longlong_as_double( used_colors ) );\n int tmp_lo = __double2hiint( __longlong_as_double( used_colors ) );\n tmp_hi = utils::shfl_xor(tmp_hi, i, WARP_SIZE);\n tmp_lo = utils::shfl_xor(tmp_lo, i, WARP_SIZE);\n long long tmp = __double_as_longlong(__hiloint2double(tmp_hi, tmp_lo));\n used_colors |= tmp;\n }\n\n int my_color = 64 - utils::bfind( ~used_colors );\n\n if (my_color <= 0) { my_color = 1; }\n\n /*int sets=0;\n for(int c=1; c<=64; c++)\n {\n unsigned long long int b = (1ull << (64-c));\n if((~used_colors & b) && sets < (dice % 3)-1)\n {\n sets++;\n my_color = c;\n }\n }*/\n\n // The warp leader stores the result.\n if ( candidate && utils::all( is_max_vertex ) )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = my_color;\n }\n\n continue;\n }\n\n if ( candidate && utils::all( is_min_vertex ) )\n {\n if ( lane_id == 0 )\n {\n A_colors[row_id] = my_color + 1;\n }\n\n continue;\n }\n }\n}\n\n\ntemplate< int CTA_SIZE, int WARP_SIZE>\n__global__\nvoid color_kernel_reassign_tail_thread( const int A_num_rows, const int *A_rows, const int *A_cols, int num_colors, int first_tail_color, int *A_colors )\n{\n const int NUM_THREADS_PER_GRID = gridDim.x * CTA_SIZE;\n // Row identifier.\n int row_id = blockIdx.x * CTA_SIZE + threadIdx.x;\n\n // Iterate over rows.\n for ( ; row_id < A_num_rows ; row_id += NUM_THREADS_PER_GRID )\n {\n unsigned long long used_colors = 0ull;\n int row_color = A_colors[row_id];\n\n if ( row_color < first_tail_color ) // Already colored!!!\n {\n continue;\n }\n\n int row_hash = hash_function(row_id, 198719871987L, 0);\n int row_begin = A_rows[row_id ];\n int row_end = A_rows[row_id + 1];\n\n for ( int row_it = row_begin; row_it < row_end ; row_it++)\n {\n int col_id = A_cols[row_it];\n\n if (col_id >= A_num_rows) { continue; }\n\n int col_hash = hash_function( col_id, 1987, 0 );\n int col_color = A_colors[col_id];\n }\n }\n}\n\n__global__ void unassign_color(int *A_colors, int num_rows)\n{\n for (int tid = threadIdx.x + blockIdx.x * blockDim.x; tid < num_rows; tid += gridDim.x * blockDim.x)\n {\n if (tid % 9 == 0)\n {\n A_colors[tid] = 0;\n }\n }\n}\n\n\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nMin_Max_2Ring_Matrix_Coloring >::color_step( Matrix_d &A, int seed )\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize( A.row_offsets.size() - 1 );\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n const int num_rows = A.get_num_rows();\n const int max_uncolored_rows = static_cast( this->m_uncolored_fraction * num_rows );\n const int CTA_SIZE = 128;\n const int NUM_WARPS_PER_CTA = CTA_SIZE / 32;\n const int GRID_SIZE = std::min( 2048, (num_rows + NUM_WARPS_PER_CTA - 1) / NUM_WARPS_PER_CTA );\n device_vector_alloc gtlt_count( num_rows );\n\n for ( int num_uncolored = num_rows ; num_uncolored > max_uncolored_rows ; )\n {\n count_gtlt_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n this->m_row_colors.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n seed);\n cudaCheckError();\n\n if ( this->m_late_rejection )\n color_kernel_greedy <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n this->m_num_colors,\n this->m_weakness_bound,\n this->m_row_colors.raw() );\n else\n color_kernel_greedy <<< GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( >lt_count.front() ),\n this->m_num_colors,\n this->m_weakness_bound,\n this->m_row_colors.raw() );\n\n cudaCheckError();\n this->m_num_colors += 2;\n num_uncolored = (int) amgx::thrust::count_if( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, is_zero() );\n cudaCheckError();\n }\n\n this->m_num_colors = thrust_wrapper::reduce( this->m_row_colors.begin(), this->m_row_colors.begin() + num_rows, 0, amgx::thrust::maximum() ) + 1;\n}\n\n\n\n\n// Block version\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nMin_Max_2Ring_Matrix_Coloring >::color_matrix( Matrix_d &A )\n{\n ViewType oldView = A.currentView();\n this->m_row_colors.resize( A.row_offsets.size() - 1 );\n\n if (this->m_halo_coloring == SYNC_COLORS) { A.setView(ALL); }\n else { A.setViewExterior(); }\n\n this->m_num_colors = 1;\n thrust_wrapper::fill( this->m_row_colors.begin(), this->m_row_colors.end(), 0 );\n cudaCheckError();\n color_step(A, 0);\n\n A.setView(oldView);\n}\n\n\n#endif\n\n\n#define AMGX_CASE_LINE(CASE) template class Min_Max_2Ring_Matrix_Coloring_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Min_Max_2Ring_Matrix_Coloring::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // end namespace amgx\n\n", "source": "/kaggle/working/amgx/src/matrix_coloring/min_max_2ring.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n\n#ifdef _WIN32\n#pragma warning (push)\n#pragma warning (disable : 4244 4267 4521)\n#endif\n#ifdef _WIN32\n#pragma warning (pop)\n#endif\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\nnamespace aggregation\n{\n\n\n// ----------------------\n// Kernels\n// ----------------------\n\ntemplate \n__global__\nvoid set_to_one_kernel(IndexType start, IndexType end, IndexType *ind, ValueType *v)\n{\n for (int tid = start + blockDim.x * blockIdx.x + threadIdx.x; tid < end; tid += gridDim.x * blockDim.x)\n {\n v[ind[tid]] = types::util::get_one();\n }\n}\n\ntemplate \n__global__\nvoid renumberAggregatesKernel(const IndexType *renumbering, const int interior_offset, const int bdy_offset, IndexType *aggregates, const int num_aggregates, const int n_interior, const int renumbering_size)\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (tid < num_aggregates)\n {\n IndexType new_agg_id;\n\n if (renumbering_size == 0)\n {\n new_agg_id = aggregates[tid];\n }\n else\n {\n new_agg_id = renumbering[aggregates[tid]];\n }\n\n //if (aggregates[tid] > num_aggregates)\n //{\n //printf(\"ID %d old %d + %d = %d\\n\", tid, new_agg_id, ((new_agg_id >= n_interior) ? bdy_offset : interior_offset), new_agg_id + ((new_agg_id >= n_interior) ? bdy_offset : interior_offset));\n //}\n new_agg_id += ((new_agg_id >= n_interior) ? bdy_offset : interior_offset);\n aggregates[tid] = new_agg_id;\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel to restrict residual using csr_format\ntemplate \n__global__\nvoid restrictResidualKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueType *r, ValueType *rr, const int num_aggregates)\n{\n int jmin, jmax;\n\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_aggregates; tid += gridDim.x * blockDim.x)\n {\n ValueType temp(types::util::get_zero());\n jmin = row_offsets[tid];\n jmax = row_offsets[tid + 1];\n\n for (int j = jmin; j < jmax; j++)\n {\n int j_col = column_indices[j];\n temp = temp + r[j_col];\n }\n\n rr[tid] = temp;\n }\n}\n\n// Kernel to restrict residual using block_dia_csr_format\ntemplate \n__global__\nvoid restrictResidualBlockDiaCsrKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueType *r, ValueType *rr, const int num_aggregates)\n{\n ValueType rr_temp[bsize];\n int offset, jmin, jmax;\n\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_aggregates; tid += gridDim.x * blockDim.x)\n {\n // Initialize to zero\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n rr_temp[m] = types::util::get_zero();\n }\n\n jmin = row_offsets[tid];\n jmax = row_offsets[tid + 1];\n\n for (int j = jmin; j < jmax; j++)\n {\n int jcol = column_indices[j];\n offset = jcol * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n rr_temp[m] = rr_temp[m] + r[offset + m];\n }\n }\n\n offset = tid * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n rr[offset + m] = rr_temp[m];\n };\n }\n}\n\n// Kernel to prolongate and apply the correction for csr format\ntemplate \n__global__\nvoid prolongateAndApplyCorrectionKernel(const ValueType alpha, const int num_rows, ValueType *x, const ValueType *e, const IndexType *aggregates, IndexType num_aggregates)\n{\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_rows; tid += gridDim.x * blockDim.x)\n {\n IndexType I = aggregates[tid];\n x[tid] = x[tid] + alpha * e[I];\n }\n}\n\n// Kernel to prolongate and apply the correction for block-dia-csr format\ntemplate \n__global__\nvoid prolongateAndApplyCorrectionBlockDiaCsrKernel(const ValueType alpha, const int num_block_rows, ValueType *x, const ValueType *e, const IndexType *aggregates, IndexType num_aggregates, const int bsize)\n{\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_block_rows; tid += gridDim.x * blockDim.x)\n {\n IndexType I = aggregates[tid];\n\n for (int m = 0; m < bsize; m++)\n {\n x[tid * bsize + m] = x[tid * bsize + m] + alpha * e[I * bsize + m];\n }\n }\n}\n\ntemplate \n__global__\nvoid prolongateVector(const IndexType *aggregates, const ValueType *in, ValueType *out, IndexType fine_rows, IndexType coarse_rows, int blocksize)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < fine_rows * blocksize )\n {\n int i = tid / blocksize;\n int e = tid % blocksize;\n IndexType I = aggregates[i];\n out[tid] = in[ I * blocksize + e ];\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \n__global__\nvoid applyCorrection(ValueType lambda, const ValueType *e, ValueType *x, IndexType numRows )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < numRows )\n {\n x[tid] = x[tid] + lambda * e[tid];\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// -------------------------------\n// Methods\n// ------------------------------\n\ntemplate \nvoid Aggregation_AMG_Level_Base::transfer_level(AMG_Level *ref_lvl)\n{\n Aggregation_AMG_Level_Base *ref_agg_lvl = dynamic_cast*>(ref_lvl);\n this->scale_counter = ref_agg_lvl->scale_counter;\n this->scale = ref_agg_lvl->scale;\n this->m_R_row_offsets.copy(ref_agg_lvl->m_R_row_offsets);\n this->m_R_column_indices.copy(ref_agg_lvl->m_R_column_indices);\n this->m_aggregates.copy(ref_agg_lvl->m_aggregates);\n this->m_aggregates_fine_idx.copy(ref_agg_lvl->m_aggregates_fine_idx);\n this->m_num_aggregates = ref_agg_lvl->m_num_aggregates;\n this->m_num_all_aggregates = ref_agg_lvl->m_num_all_aggregates;\n}\n\n\ntypedef std::pair mypair;\nbool comparator ( const mypair &l, const mypair &r) { return l.first < r.first; }\n\n// Method to compute R\n// General path\n// TODO: this could be merged with selector to save some computations\ntemplate \nvoid Aggregation_AMG_Level_Base::computeRestrictionOperator_common()\n{\n m_R_row_offsets.resize(m_num_all_aggregates + 1); //create one more row for the pseudo aggregate\n IVector R_row_indices(m_aggregates);\n#if AMGX_ASYNCCPU_PROOF_OF_CONCEPT\n bool use_cpu = m_aggregates.size() < 4096;\n\n if (use_cpu)\n {\n struct computeRestrictionTask : public task\n {\n Aggregation_AMG_Level_Base *self;\n IVector *R_row_indices;\n\n void run()\n {\n int N = self->m_aggregates.size();\n IVector_h R_row_indices_host(self->m_aggregates);\n std::vector pairs(N);\n\n for (int i = 0; i < N; i++)\n {\n pairs[i].first = R_row_indices_host[i];\n pairs[i].second = i;\n }\n\n std::stable_sort(pairs.begin(), pairs.end(), comparator);\n IVector_h R_column_indices(self->A->get_num_rows());\n\n for (int i = 0; i < N; i++)\n {\n R_column_indices[i] = pairs[i].second;\n R_row_indices_host[i] = pairs[i].first;\n }\n\n self->m_R_column_indices = R_column_indices;\n *R_row_indices = R_row_indices_host;\n }\n };\n computeRestrictionTask *t = new computeRestrictionTask();\n t->self = this;\n t->R_row_indices = &R_row_indices;\n t->run();\n delete t;\n }\n else\n#endif\n {\n m_R_column_indices.resize(this->A->get_num_rows());\n thrust_wrapper::sequence(m_R_column_indices.begin(), m_R_column_indices.end());\n cudaCheckError();\n amgx::thrust::sort_by_key(R_row_indices.begin(), R_row_indices.end(), m_R_column_indices.begin());\n cudaCheckError();\n }\n\n amgx::thrust::lower_bound(R_row_indices.begin(),\n R_row_indices.end(),\n amgx::thrust::counting_iterator(0),\n amgx::thrust::counting_iterator(m_R_row_offsets.size()),\n m_R_row_offsets.begin());\n cudaCheckError();\n}\n\n\n// two methods below could be merged\n// Method to compute R on HOST using csr format\ntemplate \nvoid Aggregation_AMG_Level >::computeRestrictionOperator_1x1()\n{\n this->m_R_row_offsets.resize(this->m_num_all_aggregates + 1);\n this->m_R_column_indices.resize(this->A->get_num_rows());\n this->fillRowOffsetsAndColIndices(this->A->get_num_rows());\n}\n\n// Method to compute R on HOST using block dia-csr format\ntemplate \nvoid Aggregation_AMG_Level >::computeRestrictionOperator_4x4()\n{\n this->m_R_row_offsets.resize(this->m_num_all_aggregates + 1);\n this->m_R_column_indices.resize(this->A->get_num_rows());\n this->fillRowOffsetsAndColIndices(this->A->get_num_rows());\n}\n\n// Method to create R_row_offsest and R_column_indices array on HOST using csr or block dia-csr format\ntemplate \nvoid Aggregation_AMG_Level_Base::fillRowOffsetsAndColIndices(const int R_num_cols)\n{\n for (int i = 0; i < m_num_all_aggregates + 1; i++)\n {\n m_R_row_offsets[i] = 0;\n }\n\n // Count number of neighbors for each row\n for (int i = 0; i < R_num_cols; i++)\n {\n int I = m_aggregates[i];\n m_R_row_offsets[I]++;\n }\n\n m_R_row_offsets[m_num_all_aggregates] = R_num_cols;\n\n for (int i = m_num_all_aggregates - 1; i >= 0; i--)\n {\n m_R_row_offsets[i] = m_R_row_offsets[i + 1] - m_R_row_offsets[i];\n }\n\n /* Set column indices. */\n for (int i = 0; i < R_num_cols; i++)\n {\n int I = m_aggregates[i];\n int Ip = m_R_row_offsets[I]++;\n m_R_column_indices[Ip] = i;\n }\n\n /* Reset r[i] to start of row memory. */\n for (int i = m_num_all_aggregates - 1; i > 0; i--)\n {\n m_R_row_offsets[i] = m_R_row_offsets[i - 1];\n }\n\n m_R_row_offsets[0] = 0;\n}\n\n// Method to compute R on DEVICE using block dia-csr format\ntemplate \nvoid Aggregation_AMG_Level >::computeRestrictionOperator_4x4()\n{\n this->computeRestrictionOperator_common();\n}\n\ntemplate \nvoid Aggregation_AMG_Level >::computeRestrictionOperator_1x1()\n{\n this->computeRestrictionOperator_common();\n}\n\n// Method to restrict Residual on host using csr_matrix format\ntemplate \nvoid Aggregation_AMG_Level >::restrictResidual_1x1(const VVector &r, VVector &rr)\n{\n ValueTypeB temp;\n\n for (int i = 0; i < this->m_num_aggregates; i++)\n {\n temp = types::util::get_zero();\n\n for (int j = this->m_R_row_offsets[i]; j < this->m_R_row_offsets[i + 1]; j++)\n {\n int j_col = this->m_R_column_indices[j];\n temp = temp + r[j_col];\n }\n\n rr[i] = temp;\n }\n}\n\n// Method to restrict Residual on host using block_dia_csr_matrix format\ntemplate \nvoid Aggregation_AMG_Level >::restrictResidual_4x4(const VVector &r, VVector &rr)\n{\n IndexType bsize = this->A->get_block_dimy();\n ValueTypeB *temp = new ValueTypeB[bsize];\n\n for (int i = 0; i < this->m_num_aggregates; i++)\n {\n // Initialize temp to 0\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = types::util::get_zero();\n }\n\n // Add contributions from each fine point\n for (int j = this->m_R_row_offsets[i]; j < this->m_R_row_offsets[i + 1]; j++)\n {\n int j_col = this->m_R_column_indices[j];\n\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = temp[k] + r[j_col * bsize + k];\n }\n }\n\n // Store result\n for (int k = 0; k < bsize; k++)\n {\n rr[i * bsize + k] = temp[k];\n }\n }\n}\n\n// Method to restrict Residual on device using csr_matrix format\ntemplate \nvoid Aggregation_AMG_Level >::restrictResidual_1x1(const VVector &r, VVector &rr)\n{\n int block_size = 128;\n int max_threads;\n\n if (!this->isConsolidationLevel())\n {\n max_threads = this->m_num_aggregates;\n }\n else\n {\n max_threads = this->m_num_all_aggregates;\n }\n\n int num_blocks = max_threads / block_size + 1;\n const IndexType *R_row_offsets_ptr = this->m_R_row_offsets.raw();\n const IndexType *R_column_indices_ptr = this->m_R_column_indices.raw();\n const ValueTypeB *r_ptr = r.raw();\n ValueTypeB *rr_ptr = rr.raw();\n restrictResidualKernel <<< num_blocks, block_size>>>(R_row_offsets_ptr, R_column_indices_ptr, r_ptr, rr_ptr, max_threads);\n cudaCheckError();\n}\n\n// Method to restrict Residual on device using block_dia_csr_matrix format\ntemplate \nvoid Aggregation_AMG_Level >::restrictResidual_4x4(const VVector &r, VVector &rr)\n{\n int block_size = 128;\n int max_threads;\n\n if (!this->isConsolidationLevel())\n {\n max_threads = this->m_num_aggregates;\n }\n else\n {\n max_threads = this->m_num_all_aggregates;\n };\n\n const int num_blocks = max_threads / block_size + 1;\n const IndexType *R_row_offsets_ptr = this->m_R_row_offsets.raw();\n\n const IndexType *R_column_indices_ptr = this->m_R_column_indices.raw();\n\n const ValueTypeB *r_ptr = r.raw();\n\n ValueTypeB *rr_ptr = rr.raw();\n\n cudaCheckError();\n\n switch ( this->getA().get_block_dimy() )\n {\n case 2:\n restrictResidualBlockDiaCsrKernel <<< num_blocks, block_size>>>(R_row_offsets_ptr, R_column_indices_ptr, r_ptr, rr_ptr, max_threads);\n cudaCheckError();\n break;\n\n case 3:\n restrictResidualBlockDiaCsrKernel <<< num_blocks, block_size>>>(R_row_offsets_ptr, R_column_indices_ptr, r_ptr, rr_ptr, max_threads);\n cudaCheckError();\n break;\n\n case 4:\n restrictResidualBlockDiaCsrKernel <<< num_blocks, block_size>>>(R_row_offsets_ptr, R_column_indices_ptr, r_ptr, rr_ptr, max_threads);\n cudaCheckError();\n break;\n\n case 5:\n restrictResidualBlockDiaCsrKernel <<< num_blocks, block_size>>>(R_row_offsets_ptr, R_column_indices_ptr, r_ptr, rr_ptr, max_threads);\n cudaCheckError();\n break;\n\n case 8:\n restrictResidualBlockDiaCsrKernel <<< num_blocks, block_size>>>(R_row_offsets_ptr, R_column_indices_ptr, r_ptr, rr_ptr, max_threads);\n cudaCheckError();\n break;\n\n case 10:\n restrictResidualBlockDiaCsrKernel <<< num_blocks, block_size>>>(R_row_offsets_ptr, R_column_indices_ptr, r_ptr, rr_ptr, max_threads);\n cudaCheckError();\n break;\n\n default:\n FatalError( \"Unsupported block size in restrictResidual_4x4!!!\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE );\n }\n\n cudaCheckError();\n}\n\n__inline__ float getAlpha(float &nom, float &denom)\n{\n float alpha;\n\n if (nom * denom <= 0. || std::abs(nom) < std::abs(denom))\n {\n alpha = 1.;\n }\n else if (std::abs(nom) > 2.*std::abs(denom))\n {\n alpha = 2.;\n }\n else\n {\n alpha = nom / denom;\n }\n\n return alpha;\n}\n\n__inline__ double getAlpha(double &nom, double &denom)\n{\n double alpha;\n\n if (nom * denom <= 0. || std::abs(nom) < std::abs(denom))\n {\n alpha = 1.;\n }\n else if (std::abs(nom) > 2.*std::abs(denom))\n {\n alpha = 2.;\n }\n else\n {\n alpha = nom / denom;\n }\n\n return alpha;\n}\n\n__inline__ cuComplex getAlpha(cuComplex &nom, cuComplex &denom)\n{\n cuComplex alpha;\n\n if (types::util::abs(nom) < types::util::abs(denom))\n {\n alpha = make_cuComplex(1.f, 0.f);\n }\n else if (types::util::abs(nom) > 2.*types::util::abs(denom))\n {\n alpha = make_cuComplex(2.f, 0.f);\n }\n else\n {\n alpha = nom / denom;\n }\n\n return alpha;\n}\n\n__inline__ cuDoubleComplex getAlpha(cuDoubleComplex &nom, cuDoubleComplex &denom)\n{\n cuDoubleComplex alpha;\n\n if (types::util::abs(nom) < types::util::abs(denom))\n {\n alpha = make_cuDoubleComplex(1., 0.);\n }\n else if (types::util::abs(nom) > 2.*types::util::abs(denom))\n {\n alpha = make_cuDoubleComplex(2., 0.);\n }\n else\n {\n alpha = nom / denom;\n }\n\n return alpha;\n}\n\ntemplate< class T_Config>\ntypename T_Config::VecPrec Aggregation_AMG_Level_Base::computeAlpha(const Vector &e, const Vector &bc, const Vector &tmp)\n{\n typename T_Config::VecPrec alpha = types::util::get_one();\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n int size = Ac.get_num_rows();\n VVector v(2, types::util::get_zero());\n v[0] = amgx::thrust::inner_product(e.begin(), e.begin() + size, bc.begin(), types::util::get_zero());\n v[1] = amgx::thrust::inner_product(e.begin(), e.begin() + size, tmp.begin(), types::util::get_zero());\n cudaCheckError();\n return getAlpha(v[0], v[1]);\n}\n\n// Method to prolongate the error on HOST using csr format\ntemplate \nvoid Aggregation_AMG_Level >::prolongateAndApplyCorrection_1x1(Vector > &e, Vector > &bc, Vector > &x, Vector > &tmp)\n{\n Matrix > &A = this->getA();\n Matrix > &C = this->next_h->getA();\n\n if ( this->m_error_scaling >= 2 )\n {\n FatalError(\"error_scaling=2,3 is not implemented on host\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n ValueTypeB alpha = types::util::get_one();\n\n if (this->m_error_scaling)\n {\n multiply(this->next_h->getA(), e, tmp);\n alpha = this->computeAlpha (e, bc, tmp);\n }\n\n // Apply correction on all (interior and exterior) equations.\n for (int i = 0; i < A.get_num_cols(); i++)\n {\n int I = this->m_aggregates[i];\n x[i] = x[i] + alpha * e[I];\n }\n}\n\n// Method to prolongate the error on HOST using block_dia_csr format\ntemplate \nvoid Aggregation_AMG_Level >::prolongateAndApplyCorrection_4x4(Vector > &e, Vector > &bc, Vector > &x, Vector > &tmp)\n{\n if (this->A->get_block_dimy() != this->A->get_block_dimx())\n {\n FatalError(\"Aggregation_AMG_Level not implemented for non square blocks, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if ( this->m_error_scaling >= 2 )\n {\n FatalError(\"error_scaling=2,3 is not implemented on host\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n Matrix &C = this->next_h->getA();\n ValueTypeB alpha = types::util::get_one();\n\n if (this->m_error_scaling)\n {\n multiply(this->next_h->getA(), e, tmp);\n alpha = this->computeAlpha (e, bc, tmp);\n }\n\n // Apply correction on all equations.\n for (int i = 0; i < this->A->get_num_rows(); i++)\n {\n int I = this->m_aggregates[i];\n\n for (int k = 0; k < this->A->get_block_dimy(); k++)\n {\n x[i * this->A->get_block_dimy() + k] = x[i * this->A->get_block_dimy() + k] + alpha * e[I * this->A->get_block_dimy() + k];\n }\n }\n}\n\n// Prolongate the error on DEVICE using csr format\ntemplate \nvoid Aggregation_AMG_Level >::prolongateAndApplyCorrection_1x1(Vector > &e, Vector > &bc, Vector > &x, Vector > &tmp)\n{\n ValueTypeB alpha = types::util::get_one();\n const int block_size = 128;\n const int num_blocks = this->A->get_num_rows() / block_size + 1;\n const IndexType *aggregates_ptr = this->m_aggregates.raw();\n ValueTypeB *x_ptr = x.raw();\n const ValueTypeB *e_ptr = e.raw();\n\n if (this->m_error_scaling)\n {\n FatalError(\"error_scaling=1 is deprecated\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n prolongateAndApplyCorrectionKernel <<< num_blocks, block_size>>>(alpha, (int)this->A->get_num_rows(), x_ptr, e_ptr, aggregates_ptr, this->m_num_aggregates);\n cudaCheckError();\n}\n\n// Prolongate the error on DEVICE using block dia-csr format\ntemplate \nvoid Aggregation_AMG_Level >::prolongateAndApplyCorrection_4x4(Vector > &ec,\n Vector > &bf,\n Vector > &xf,\n Vector > &rf)\n{\n if ( this->m_error_scaling >= 2 )\n {\n if ( this->scale_counter > 0 )\n {\n const IndexType *aggregates_ptr = this->m_aggregates.raw();\n ValueTypeB *x_ptr = xf.raw();\n const ValueTypeB *e_ptr = ec.raw();\n const int block_size = 128;\n const int num_blocks = (this->A->get_num_rows() - 1) / block_size + 1;\n prolongateAndApplyCorrectionBlockDiaCsrKernel <<< num_blocks, block_size>>>(this->scale, (int)this->getA().get_num_rows(), x_ptr, e_ptr, aggregates_ptr, this->m_num_aggregates, this->getA().get_block_dimy());\n cudaCheckError();\n this->scale_counter--;\n return;\n }\n\n bool vanek_scaling = this->m_error_scaling > 3;\n IndexType numRowsCoarse = this->next_d->getA().get_num_rows();\n IndexType numRowsFine = this->A->get_num_rows();\n IndexType blockdim = this->A->get_block_dimx();\n\n if ( blockdim != this->A->get_block_dimy() )\n {\n FatalError(\"Unsupported dimension for aggregation amg level\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n VVector ef( rf.size() );\n VVector Aef( rf.size() );\n ef.set_block_dimy( blockdim );\n Aef.set_block_dimy( blockdim );\n // prolongate e\n const int threads_per_block = 256;\n const int num_block_values = std::min( AMGX_GRID_MAX_SIZE, (numRowsFine * blockdim - 1) / threads_per_block + 1);\n const cudaStream_t stream = nullptr;\n prolongateVector <<< num_block_values, threads_per_block, 0, stream>>>( this->m_aggregates.raw(), ec.raw(), ef.raw(), numRowsFine, numRowsCoarse, blockdim );\n cudaCheckError();\n ef.dirtybit = 1;\n cudaStreamSynchronize(stream);\n cudaCheckError();\n int preSmooth;\n\n if ( vanek_scaling )\n {\n preSmooth = this->amg->getNumPostsweeps();\n }\n else\n {\n preSmooth = this->scaling_smoother_steps;\n }\n\n //smooth error\n this->smoother->setTolerance( 0.0 );\n this->smoother->set_max_iters( preSmooth );\n\n if ( vanek_scaling )\n {\n thrust_wrapper::fill( Aef.begin(), Aef.end(), types::util::get_zero() );\n cudaCheckError();\n this->smoother->solve( Aef, ef, false ); //smooth correction with rhs 0\n this->smoother->solve( bf, xf, false ); // smooth x with rhs residual\n //recompute residual\n int offset, size;\n this->getA().getOffsetAndSizeForView(OWNED, &offset, &size);\n axmb( this->getA(), xf, bf, rf, offset, size );\n }\n else\n {\n this->smoother->solve( rf, ef, false ); //smooth correction with rhs residual\n }\n\n // multiply for lambda computation\n multiply(this->getA(), ef, Aef, OWNED);\n ValueTypeB nominator, denominator;\n int offset = 0, size = 0;\n this->A->getOffsetAndSizeForView(OWNED, &offset, &size);\n\n if ( this->m_error_scaling == 2 || this->m_error_scaling == 4 )\n {\n // compute lambda=/\n nominator = amgx::thrust::inner_product( rf.begin(), rf.end(), Aef.begin(), types::util::get_zero() );\n denominator = amgx::thrust::inner_product( Aef.begin(), Aef.end(), Aef.begin(), types::util::get_zero() );\n cudaCheckError();\n }\n\n if ( this->m_error_scaling == 3 || this->m_error_scaling == 5)\n {\n // compute lambda=/\n nominator = amgx::thrust::inner_product( rf.begin(), rf.begin() + size * blockdim, ef.begin(), types::util::get_zero() );\n denominator = amgx::thrust::inner_product( ef.begin(), ef.begin() + size * blockdim, Aef.begin(), types::util::get_zero() );\n\n if (!this->A->is_matrix_singleGPU())\n {\n this->A->getManager()->global_reduce_sum(&nominator);\n this->A->getManager()->global_reduce_sum(&denominator);\n }\n\n cudaCheckError();\n }\n\n if (types::util::abs(denominator) == 0.0)\n {\n nominator = denominator = types::util::get_one();\n }\n\n // apply correction x <- x + lambda*e\n const int num_block_fine = std::min( AMGX_GRID_MAX_SIZE, (numRowsFine * blockdim - 1) / threads_per_block + 1 );\n ValueTypeB alpha = nominator / denominator;\n\n if ( types::util::abs(alpha) < .3 )\n {\n alpha = (alpha / types::util::abs(alpha)) * .3; // it was this before: alpha = .3, which is not 100% equal\n }\n\n if ( types::util::abs(alpha) > 10 )\n {\n alpha = (alpha / types::util::abs(alpha)) * 10.; // it was this before: alpha = 10., which is not 100% equal\n }\n\n applyCorrection <<< num_block_fine, threads_per_block, 0, stream>>>( alpha, ef.raw(), xf.raw(), numRowsFine * blockdim );\n cudaCheckError();\n this->scale_counter = this->reuse_scale; //reuse this scale scale_counter times\n this->scale = alpha;\n return;\n }\n\n ValueTypeB alpha = types::util::get_one();\n const int block_size = 128;\n const int num_blocks = this->A->get_num_rows() / block_size + 1;\n const IndexType *aggregates_ptr = this->m_aggregates.raw();\n ValueTypeB *x_ptr = xf.raw();\n const ValueTypeB *e_ptr = ec.raw();\n\n if (this->m_error_scaling == 1)\n {\n FatalError(\"error_scaling=1 is deprecated\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n prolongateAndApplyCorrectionBlockDiaCsrKernel <<< num_blocks, block_size>>>(alpha, (int)this->A->get_num_rows(), x_ptr, e_ptr, aggregates_ptr, this->m_num_aggregates, this->A->get_block_dimy());\n cudaCheckError();\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::prolongateAndApplyCorrection(VVector &e, VVector &bf, VVector &x, VVector &tmp)\n{\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n\n //this is dirty, but error scaling 2 and 3 do not have a specialized version. Instead, the general version sits in the 4x4 function\n if ( this->m_error_scaling >= 2 )\n {\n prolongateAndApplyCorrection_4x4(e, bf, x, tmp);\n }\n else if (this->A->get_block_size() == 1)\n {\n prolongateAndApplyCorrection_1x1(e, bf, x, tmp);\n }\n else if (this->A->get_block_dimx() == this->A->get_block_dimy() )\n {\n prolongateAndApplyCorrection_4x4(e, bf, x, tmp);\n }\n else\n {\n FatalError(\"Unsupported dimension for aggregation amg level\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n x.dirtybit = 1;\n\n if (!this->A->is_matrix_singleGPU() && x.delayed_send == 0)\n {\n if (x.in_transfer & RECEIVING) { this->A->manager->exchange_halo_wait(x, x.tag); }\n\n this->A->manager->exchange_halo_async(x, x.tag);\n }\n}\n\n\ntemplate \nvoid Aggregation_AMG_Level_Base::restrictResidual(VVector &r, VVector &rr)\n{\n if (this->A->get_block_size() == 1)\n {\n restrictResidual_1x1(r, rr);\n }\n else if (this->A->get_block_dimx() == this->A->get_block_dimy() )\n {\n restrictResidual_4x4(r, rr);\n }\n else\n {\n FatalError(\"Unsupported dimension for aggregation amg level\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n //TODO: check level transfer between host and device for multiGPU\n if (!this->A->is_matrix_singleGPU())\n {\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n rr.dirtybit = 1;\n\n if (!Ac.is_matrix_singleGPU() && !this->isConsolidationLevel() && rr.delayed_send == 0)\n {\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA(); //TODO problem in memoryspace transfer is here\n\n if (rr.in_transfer & RECEIVING) { Ac.manager->exchange_halo_wait(rr, rr.tag); }\n\n Ac.manager->exchange_halo_async(rr, rr.tag);\n }\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::computeRestrictionOperator()\n{\n if (this->A->get_block_size() == 1)\n {\n computeRestrictionOperator_1x1();\n }\n else if (this->A->get_block_dimx() == 4 && this->A->get_block_dimy() == 4)\n {\n computeRestrictionOperator_4x4();\n }\n else\n {\n this->computeRestrictionOperator_common();\n }\n}\n\ntemplate \n__global__ void coarse_to_global(IndexType *aggregates, IndexType *aggregates_global, IndexType *renumbering, IndexType num_elements, int64_t offset)\n{\n int element = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (element < num_elements)\n {\n renumbering[aggregates[element]] = aggregates_global[element] + offset; //this won't be a problem, because we are overwriting the same thing\n element += blockDim.x * gridDim.x;\n }\n}\n\ntemplate \n__global__ void export_matrix_elements(IndexType *row_offsets, IndexType *col_indices, T *values, IndexType *maps, IndexType *renumbering, IndexType *new_row_offsets, IndexType *new_col_indices, T *new_values, IndexType bsize, IndexType size)\n{\n int idx = blockIdx.x * blockDim.x / 32 + threadIdx.x / 32;\n int coopIdx = threadIdx.x % 32;\n\n while (idx < size)\n {\n int row = maps[idx];\n INDEX_TYPE src_base = row_offsets[row];\n INDEX_TYPE dst_base = new_row_offsets[idx];\n\n for (int m = coopIdx; m < row_offsets[row + 1]*bsize - src_base * bsize; m += 32)\n {\n new_values[dst_base * bsize + m] = values[src_base * bsize + m];\n }\n\n for (int m = coopIdx; m < row_offsets[row + 1] - src_base; m += 32)\n {\n new_col_indices[dst_base + m] = renumbering[col_indices[src_base + m]];\n }\n\n idx += gridDim.x * blockDim.x / 32;\n }\n}\n\ntemplate \n__global__ void export_matrix_diagonal(T *values, INDEX_TYPE bsize, INDEX_TYPE *maps, T *output, INDEX_TYPE size)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (idx < size)\n {\n int row = maps[idx];\n INDEX_TYPE src_base = row;\n INDEX_TYPE dst_base = idx;\n\n for (int m = 0; m < bsize; m++)\n {\n output[dst_base * bsize + m] = values[src_base * bsize + m];\n }\n\n idx += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void remove_boundary(INDEX_TYPE *flags, INDEX_TYPE *maps, INDEX_TYPE size)\n{\n int element = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (element < size)\n {\n flags[maps[element]] = 0; //this won't be a problem, because we are overwriting the same thing\n element += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void calc_inverse_renumbering(INDEX_TYPE *renum, INDEX_TYPE *irenum, INDEX_TYPE *renum_gbl, INDEX_TYPE base_index, INDEX_TYPE max_element)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < max_element)\n {\n irenum[renum[idx]] = renum_gbl[idx] - base_index;\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void create_halo_mapping(INDEX_TYPE *mapping, INDEX_TYPE *node_list, INDEX_TYPE base_index, INDEX_TYPE map_offset, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n int idx = node_list[row] - base_index;\n mapping[idx] = map_offset + row;\n row += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void map_col_indices_and_count_rowlen(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, INDEX_TYPE *row_length, INDEX_TYPE *renumbering, INDEX_TYPE *mapping, INDEX_TYPE *map_offsets, int64_t *index_ranges, INDEX_TYPE part_id, INDEX_TYPE my_id, INDEX_TYPE base_index, INDEX_TYPE my_range, INDEX_TYPE num_neighbors, INDEX_TYPE num_rows)\n{\n extern __shared__ volatile int reduction[];\n int row = blockIdx.x * blockDim.x / 4 + threadIdx.x / 4;\n int coopIdx = threadIdx.x % 4;\n\n while (row < num_rows)\n {\n int valid = 0;\n\n for (int idx = row_offsets[row] + coopIdx; idx < row_offsets[row + 1]; idx += 4) //this may look horrible, but I expect low branch divergence, because col indices in a row usually belong to the same partition (or at most one more)\n {\n int colIdx = col_indices[idx];\n int part = -2;\n\n if (colIdx >= index_ranges[2 * part_id] && colIdx < index_ranges[2 * part_id + 1]) //the col index probably belongs to the partition I am working on\n {\n part = part_id;\n }\n else if (colIdx >= base_index && colIdx < base_index + my_range) //or points back to the owned partition\n {\n part = -1;\n }\n else //or else it points to a third partition\n {\n for (int i = 0; i < num_neighbors; i++)\n {\n if (colIdx >= index_ranges[2 * i] && colIdx < index_ranges[2 * i + 1])\n {\n part = i;\n }\n }\n }\n\n if (part == -2)\n {\n col_indices[idx] = -1;\n#ifdef DEBUG\n printf(\"Column index encountered that does not belong to any of my neighbors!! %d\\n\", colIdx);\n#endif\n }\n else\n {\n if (part == -1)\n {\n col_indices[idx] = renumbering[colIdx - base_index];\n valid++;\n }\n else\n {\n int new_col_idx = mapping[map_offsets[part] + colIdx - index_ranges[2 * part]];\n\n if (new_col_idx >= 0)\n {\n valid++;\n col_indices[idx] = new_col_idx;\n }\n else\n {\n col_indices[idx] = -1;\n }\n }\n }\n }\n\n reduction[threadIdx.x] = valid;\n\n for (int s = 2; s > 0; s >>= 1)\n {\n if (coopIdx < s)\n {\n reduction[threadIdx.x] += reduction[threadIdx.x + s];\n }\n\n __syncthreads();\n }\n\n if (coopIdx == 0)\n {\n row_length[row] = reduction[threadIdx.x];\n }\n\n row += gridDim.x * blockDim.x / 4;\n }\n}\n\n__global__ void map_col_indices(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *halo_ranges, INDEX_TYPE *halo_renumbering, INDEX_TYPE *halo_rows, INDEX_TYPE *global_renumbering, INDEX_TYPE num_neighbors, INDEX_TYPE num_rows, INDEX_TYPE num_rows_processed)\n{\n int row = blockIdx.x * blockDim.x / 4 + threadIdx.x / 4;\n int coopIdx = threadIdx.x % 4;\n\n while (row < num_rows_processed)\n {\n for (int idx = row_offsets[row] + coopIdx; idx < row_offsets[row + 1]; idx += 4)\n {\n int colIdx = col_indices[idx];\n int part = 0;\n\n if (colIdx < num_rows)\n {\n part = -1;\n }\n else\n {\n colIdx = global_renumbering[colIdx];\n\n for (int i = 0; i < num_neighbors; i++)\n {\n if (colIdx >= halo_ranges[2 * i] && colIdx < halo_ranges[2 * i + 1])\n {\n part = i;\n break;\n }\n }\n }\n\n if (part == -1)\n {\n col_indices[idx] = colIdx;\n }\n else\n {\n col_indices[idx] = halo_renumbering[halo_rows[part] + colIdx - halo_ranges[2 * part]];\n }\n }\n\n row += gridDim.x * blockDim.x / 4;\n }\n}\n\ntemplate \n__global__ void reorder_whole_matrix(INDEX_TYPE *old_rows, INDEX_TYPE *old_cols, T *old_vals, INDEX_TYPE *rows, INDEX_TYPE *cols, T *vals, INDEX_TYPE bsize, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < num_rows)\n {\n INDEX_TYPE dst_row = row;\n INDEX_TYPE src_base = old_rows[row];\n INDEX_TYPE dst = rows[dst_row];\n\n for (int i = 0; i < old_rows[row + 1] - src_base; i++)\n {\n INDEX_TYPE colIdx = old_cols[src_base + i];\n\n if (colIdx >= 0)\n {\n cols[dst] = colIdx;\n\n for (int j = 0; j < bsize; j++) { vals[dst * bsize + j] = old_vals[(src_base + i) * bsize + j]; }\n\n dst++;\n }\n }\n\n row += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void calc_gbl_renumbering(INDEX_TYPE *inv_renum, INDEX_TYPE *gbl_renum, INDEX_TYPE size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n gbl_renum[inv_renum[idx]] = idx;\n idx += blockDim.x * gridDim.x;\n }\n}\n\ntemplate \n__global__ void write_diagonals(ValueType *values, INDEX_TYPE *diag, INDEX_TYPE *map, ValueType *output, INDEX_TYPE bsize, INDEX_TYPE size)\n{\n int nzPerBlock = blockDim.x / bsize;\n int row = blockIdx.x * nzPerBlock + threadIdx.x / bsize;\n int vecIdx = threadIdx.x % bsize;\n\n if (threadIdx.x >= (blockDim.x / bsize)*bsize) { return; }\n\n while (row < size)\n {\n output[row * bsize + vecIdx] = values[diag[map[row]] * bsize + vecIdx];\n row += gridDim.x * nzPerBlock;\n }\n}\n\ntemplate \n__global__ void write_diagonals_back(ValueType *values, INDEX_TYPE *diag, ValueType *source, INDEX_TYPE bsize, INDEX_TYPE size)\n{\n int nzPerBlock = blockDim.x / bsize;\n int row = blockIdx.x * nzPerBlock + threadIdx.x / bsize;\n int vecIdx = threadIdx.x % bsize;\n\n if (threadIdx.x >= (blockDim.x / bsize)*bsize) { return; }\n\n while (row < size)\n {\n values[diag[row]*bsize + vecIdx] = source[row * bsize + vecIdx];\n row += gridDim.x * nzPerBlock;\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::prepareNextLevelMatrix_full(const Matrix &A, Matrix &Ac)\n{\n if (A.is_matrix_singleGPU()) { return; }\n\n int num_neighbors = A.manager->neighbors.size();\n\n if (TConfig::memSpace == AMGX_host)\n {\n FatalError(\"Aggregation AMG Not implemented for host\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n int c_size = Ac.get_num_rows();\n int f_size = A.get_num_rows();\n int diag = Ac.hasProps(DIAG);\n\n if (A.manager->B2L_rings[0].size() > 2) { FatalError(\"Aggregation_AMG_Level prepareNextLevelMatrix not implemented >1 halo rings\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE); }\n\n //get coarse -> fine global renumbering\n IVector renumbering(c_size);\n int num_blocks = std::min(4096, (c_size + 127) / 128);\n coarse_to_global <<< num_blocks, 128>>>(this->m_aggregates.raw(), this->m_aggregates_fine_idx.raw(), renumbering.raw(), f_size, 0);\n cudaCheckError();\n //\n // Step 0 - form halo matrices that are exported to neighbors\n //\n std::vector > halo_rows(num_neighbors);\n std::vector > halo_btl(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++ )\n {\n int num_unique = Ac.manager->B2L_rings[i][1];\n //prepare export halo matrices\n halo_btl[i].resize(1, 1);\n halo_btl[i].set_global_id(Ac.manager->global_id());\n halo_btl[i].B2L_maps[0].resize(num_unique);\n halo_btl[i].B2L_rings[0].resize(2);\n halo_btl[i].B2L_rings[0][0] = 0;\n halo_btl[i].B2L_rings[0][1] = num_unique;\n halo_btl[i].set_index_range(A.manager->index_range());\n halo_btl[i].set_base_index(A.manager->base_index());\n //global indices of rows of the halo matrix\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator( renumbering.begin(), Ac.manager->B2L_maps[i].begin()),\n amgx::thrust::make_permutation_iterator( renumbering.begin(), Ac.manager->B2L_maps[i].begin() + num_unique),\n halo_btl[i].B2L_maps[0].begin());\n cudaCheckError();\n halo_rows[i].addProps(CSR);\n\n if (diag) { halo_rows[i].addProps(DIAG); }\n\n //calculate row length and row_offsets\n halo_rows[i].row_offsets.resize(num_unique + 1);\n thrust_wrapper::transform(amgx::thrust::make_permutation_iterator(Ac.row_offsets.begin() + 1, Ac.manager->B2L_maps[i].begin()),\n amgx::thrust::make_permutation_iterator(Ac.row_offsets.begin() + 1, Ac.manager->B2L_maps[i].end()),\n amgx::thrust::make_permutation_iterator(Ac.row_offsets.begin(), Ac.manager->B2L_maps[i].begin()),\n halo_rows[i].row_offsets.begin(),\n amgx::thrust::minus());\n cudaCheckError();\n thrust_wrapper::exclusive_scan(halo_rows[i].row_offsets.begin(), halo_rows[i].row_offsets.end(), halo_rows[i].row_offsets.begin());\n cudaCheckError();\n //resize halo matrix\n IndexType num_nz = halo_rows[i].row_offsets[num_unique];\n halo_rows[i].resize(num_unique, num_unique, num_nz, Ac.get_block_dimy(), Ac.get_block_dimx(), 1);\n //copy relevant rows and renumber their column indices\n num_blocks = std::min(4096, (num_unique + 127) / 128);\n export_matrix_elements <<< num_blocks, 128>>>(Ac.row_offsets.raw(), Ac.col_indices.raw(), Ac.values.raw(), Ac.manager->B2L_maps[i].raw(), renumbering.raw(), halo_rows[i].row_offsets.raw(), halo_rows[i].col_indices.raw(), halo_rows[i].values.raw(), A.get_block_size(), num_unique);\n cudaCheckError();\n\n if (diag)\n {\n export_matrix_diagonal <<< num_blocks, 128>>>(Ac.values.raw() + Ac.row_offsets[Ac.get_num_rows()]*Ac.get_block_size(), Ac.get_block_size(), Ac.manager->B2L_maps[i].raw(), halo_rows[i].values.raw() + halo_rows[i].row_offsets[halo_rows[i].get_num_rows()]*Ac.get_block_size(), num_unique);\n cudaCheckError();\n }\n }\n\n Ac.manager->getComms()->exchange_matrix_halo(halo_rows, halo_btl, Ac);\n //--------------------- renumbering/reordering matrix, integrating halo -----------------------------\n Ac.set_initialized(0);\n //number of owned rows\n c_size = Ac.manager->halo_offsets[0];\n f_size = A.manager->halo_offsets[0];\n num_blocks = std::min(4096, (c_size + 511) / 512);\n int rings = 1;\n //\n // Step 1 - calculate inverse renumbering (to global indices - base_index)\n //\n Ac.manager->inverse_renumbering.resize(c_size);\n thrust_wrapper::transform(renumbering.begin(),\n renumbering.begin() + c_size,\n amgx::thrust::constant_iterator(A.manager->base_index()),\n Ac.manager->inverse_renumbering.begin(),\n amgx::thrust::minus());\n cudaCheckError();\n //big renumbering table for going from global index to owned local index\n IVector global_to_coarse_local(Ac.manager->index_range());\n thrust_wrapper::fill(global_to_coarse_local.begin(), global_to_coarse_local.begin() + Ac.manager->index_range(), -1);\n cudaCheckError();\n calc_gbl_renumbering <<< num_blocks, 512>>>(Ac.manager->inverse_renumbering.raw(), global_to_coarse_local.raw(), c_size);\n cudaCheckError();\n Ac.manager->set_num_halo_rows(Ac.manager->halo_offsets[Ac.manager->halo_offsets.size() - 1] - c_size);\n cudaCheckError();\n //\n // Step 2 - create big mapping table of all halo indices we received (this may use a little too much memory sum(fine nodes per neighbor)\n //\n amgx::thrust::host_vector neighbor_rows(num_neighbors + 1);\n int max_num_rows = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n neighbor_rows[i] = halo_rows[i].manager->index_range();\n max_num_rows = max_num_rows > halo_rows[i].get_num_rows() ? max_num_rows : halo_rows[i].get_num_rows();\n }\n\n thrust_wrapper::exclusive_scan(neighbor_rows.begin(), neighbor_rows.end(), neighbor_rows.begin());\n cudaCheckError();\n int total_rows_of_neighbors = neighbor_rows[num_neighbors];\n IVector halo_mapping(total_rows_of_neighbors);\n thrust_wrapper::fill(halo_mapping.begin(), halo_mapping.end(), -1);\n cudaCheckError();\n\n for (int ring = 0; ring < rings; ring++)\n {\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = halo_btl[i].B2L_rings[0][ring + 1] - halo_btl[i].B2L_rings[0][ring];\n int num_blocks = std::min(4096, (size + 127) / 128);\n create_halo_mapping <<< num_blocks, 128>>>(halo_mapping.raw() + neighbor_rows[i],\n halo_btl[i].B2L_maps[0].raw() + halo_btl[i].B2L_rings[0][ring],\n halo_btl[i].base_index(),\n Ac.manager->halo_offsets[ring * num_neighbors + i], size);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n //\n // Step 3 - renumber halo matrices and calculate row length (to eventually append to the big matrix)\n //\n INDEX_TYPE owned_nnz = Ac.row_offsets[c_size];\n IVector neighbor_rows_d(num_neighbors + 1);\n amgx::thrust::copy(neighbor_rows.begin(), neighbor_rows.end(), neighbor_rows_d.begin());\n cudaCheckError();\n //map column indices of my own matrix (the ones that point outward)\n map_col_indices <<< num_blocks, 512>>>(Ac.row_offsets.raw() + Ac.manager->num_interior_nodes(),\n Ac.col_indices.raw(),\n Ac.manager->halo_ranges.raw(),\n halo_mapping.raw(),\n neighbor_rows_d.raw(),\n renumbering.raw(),\n num_neighbors, c_size, c_size - Ac.manager->num_interior_nodes());\n cudaCheckError();\n IVector temp_row_len(max_num_rows);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n //map column indices of halo matrices\n int size = halo_rows[i].get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n map_col_indices_and_count_rowlen <<< num_blocks, 128, 128 * sizeof(INDEX_TYPE)>>>(\n halo_rows[i].row_offsets.raw(),\n halo_rows[i].col_indices.raw(),\n temp_row_len.raw(),\n global_to_coarse_local.raw(),\n halo_mapping.raw(),\n neighbor_rows_d.raw(),\n Ac.manager->halo_ranges.raw(),\n i,\n Ac.manager->global_id(),\n Ac.manager->base_index(),\n Ac.manager->index_range(),\n num_neighbors,\n size);\n cudaCheckError();\n\n for (int ring = 0; ring < rings; ring++)\n {\n amgx::thrust::copy(temp_row_len.begin() + halo_btl[i].B2L_rings[0][ring], temp_row_len.begin() + halo_btl[i].B2L_rings[0][ring + 1], Ac.row_offsets.begin() + Ac.manager->halo_offsets[ring * num_neighbors + i]);\n }\n }\n\n cudaCheckError();\n INDEX_TYPE old_nnz = Ac.row_offsets[Ac.row_offsets.size() - 1];\n thrust_wrapper::exclusive_scan(Ac.row_offsets.begin() + c_size, Ac.row_offsets.end(), Ac.row_offsets.begin() + c_size, owned_nnz);\n cudaCheckError();\n //\n // Step 4 - consolidate column indices and values\n //\n int new_nnz = Ac.row_offsets[Ac.row_offsets.size() - 1];\n\n Ac.col_indices.resize(new_nnz);\n Ac.values.resize((new_nnz + 1 + diag * (Ac.row_offsets.size() - 2)) * A.get_block_size());\n\n if (diag)\n {\n MVector diags(c_size * Ac.get_block_size());\n amgx::thrust::copy(Ac.values.begin() + old_nnz * Ac.get_block_size(),\n Ac.values.begin() + old_nnz * Ac.get_block_size() + c_size * Ac.get_block_size(),\n diags.begin());\n amgx::thrust::copy(diags.begin(), diags.begin() + c_size * Ac.get_block_size(),\n Ac.values.begin() + Ac.row_offsets[Ac.get_num_rows()]*Ac.get_block_size());\n cudaCheckError();\n }\n\n int cumulative_num_rows = c_size;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n for (int ring = 0; ring < rings; ring++)\n {\n int num_rows = halo_btl[i].B2L_rings[0][ring + 1] - halo_btl[i].B2L_rings[0][ring];\n int num_blocks = std::min(4096, (num_rows + 127) / 128);\n reorder_whole_matrix <<< num_blocks, 128>>>(halo_rows[i].row_offsets.raw() + halo_btl[i].B2L_rings[0][ring], halo_rows[i].col_indices.raw(), halo_rows[i].values.raw(), Ac.row_offsets.raw() + Ac.manager->halo_offsets[ring * num_neighbors + i], Ac.col_indices.raw(), Ac.values.raw(), Ac.get_block_size(), num_rows);\n cudaCheckError();\n\n if (diag)\n {\n amgx::thrust::copy(halo_rows[i].values.begin() + (halo_rows[i].row_offsets[halo_rows[i].get_num_rows()] + halo_btl[i].B2L_rings[0][ring])*Ac.get_block_size(),\n halo_rows[i].values.begin() + (halo_rows[i].row_offsets[halo_rows[i].get_num_rows()] + halo_btl[i].B2L_rings[0][ring + 1])*Ac.get_block_size(),\n Ac.values.begin() + (Ac.row_offsets[Ac.get_num_rows()] + cumulative_num_rows)*Ac.get_block_size());\n cumulative_num_rows += num_rows;\n }\n }\n }\n\n cudaCheckError();\n Ac.set_num_cols(Ac.manager->halo_offsets[Ac.manager->halo_offsets.size() - 1]);\n Ac.set_num_rows(Ac.get_num_cols());\n Ac.set_num_nz(new_nnz);\n Ac.delProps(COO);\n Ac.set_initialized(1);\n Ac.computeDiagonal();\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::prepareNextLevelMatrix_diag(const Matrix &A, Matrix &Ac)\n{\n if (A.is_matrix_singleGPU()) { return; }\n\n int num_neighbors = A.manager->neighbors.size();\n\n if (TConfig::memSpace == AMGX_host)\n {\n FatalError(\"Aggregation AMG Not implemented for host\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n int c_size = Ac.manager->halo_offsets[0];\n int f_size = A.manager->halo_offsets[0];\n int diag = Ac.hasProps(DIAG);\n Ac.manager->inverse_renumbering.resize(c_size);\n //get coarse -> fine renumbering\n int num_blocks = std::min(4096, (c_size + 127) / 128);\n coarse_to_global <<< num_blocks, 128>>>(this->m_aggregates.raw(), this->m_aggregates_fine_idx.raw(), Ac.manager->inverse_renumbering.raw(), f_size, -1 * A.manager->base_index());\n cudaCheckError();\n Ac.manager->set_num_halo_rows(Ac.manager->halo_offsets[Ac.manager->halo_offsets.size() - 1] - c_size);\n\n if (!diag) { Ac.computeDiagonal(); }\n\n Ac.set_initialized(1);\n std::vector diagonals(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = Ac.manager->B2L_rings[i][Ac.manager->B2L_rings.size() - 1];\n diagonals[i].resize(Ac.get_block_size()*size);\n int num_blocks = std::min(4096, (size + 127) / 128);\n write_diagonals <<< num_blocks, 128>>>(Ac.values.raw(), Ac.diag.raw(), Ac.manager->B2L_maps[i].raw(), diagonals[i].raw(), Ac.get_block_size(), size);\n }\n\n cudaCheckError();\n Ac.manager->getComms()->exchange_vectors(diagonals, Ac, this->tag * 100 + 10 + 2);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = Ac.manager->halo_offsets[i + 1] - Ac.manager->halo_offsets[i];\n\n if (Ac.hasProps(DIAG)) { amgx::thrust::copy(diagonals[i].begin(), diagonals[i].begin() + Ac.get_block_size()*size, Ac.values.begin() + Ac.get_block_size() * (Ac.diagOffset() + Ac.manager->halo_offsets[i])); }\n else\n {\n int num_blocks = std::min(4096, (size + 127) / 128);\n write_diagonals_back <<< num_blocks, 128>>>(Ac.values.raw(), Ac.diag.raw() + Ac.manager->halo_offsets[i], diagonals[i].raw(), Ac.get_block_size(), size);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::prepareNextLevelMatrix_none(const Matrix &A, Matrix &Ac)\n{\n if (A.is_matrix_singleGPU()) { return; }\n\n int num_neighbors = A.manager->neighbors.size();\n\n if (TConfig::memSpace == AMGX_host)\n {\n FatalError(\"Aggregation AMG Not implemented for host\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n int c_size = Ac.manager->halo_offsets[0];\n int f_size = A.manager->halo_offsets[0];\n int diag = Ac.hasProps(DIAG);\n Ac.manager->inverse_renumbering.resize(c_size);\n //get coarse -> fine renumbering\n int num_blocks = std::min(4096, (c_size + 127) / 128);\n coarse_to_global <<< num_blocks, 128>>>(this->m_aggregates.raw(), this->m_aggregates_fine_idx.raw(), Ac.manager->inverse_renumbering.raw(), f_size, 0);\n cudaCheckError();\n Ac.manager->set_num_halo_rows(Ac.manager->halo_offsets[Ac.manager->halo_offsets.size() - 1] - c_size);\n Ac.set_initialized(1);\n\n if (!diag) { Ac.computeDiagonal(); }\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::prepareNextLevelMatrix(const Matrix &A, Matrix &Ac)\n{\n if (m_matrix_halo_exchange == 0)\n {\n this->prepareNextLevelMatrix_none(A, Ac);\n }\n else if (m_matrix_halo_exchange == 1)\n {\n this->prepareNextLevelMatrix_diag(A, Ac);\n }\n else if (m_matrix_halo_exchange == 2)\n {\n this->prepareNextLevelMatrix_full(A, Ac);\n }\n else\n {\n FatalError(\"Invalid Aggregation matrix_halo_exchange parameter\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n\n__global__ void set_halo_rowlen(INDEX_TYPE *work, INDEX_TYPE *output, INDEX_TYPE size, INDEX_TYPE diag)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n if (work[idx + 1] - work[idx] > 0)\n {\n output[idx] += work[idx + 1] - work[idx] - (1 - diag);\n }\n\n idx += blockDim.x * gridDim.x;\n }\n}\n\ntemplate \n__global__ void append_halo_nz(INDEX_TYPE *row_offsets, INDEX_TYPE *new_row_offsets, INDEX_TYPE *col_indices, INDEX_TYPE *new_col_indices, T *values, T *new_values, INDEX_TYPE size, INDEX_TYPE diag, INDEX_TYPE halo_offset, INDEX_TYPE block_size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n int add_diag = !diag;\n\n if (!diag && new_col_indices[new_row_offsets[idx]] != -1) { add_diag = 0; } //if diag or there is already soimething in the row, then don't add diagonal nonzero (inside diag)\n\n int append_offset = -1;\n\n for (int i = new_row_offsets[idx]; i < new_row_offsets[idx + 1]; i++)\n {\n if (new_col_indices[i] == -1) {append_offset = i; break;}\n }\n\n for (int i = row_offsets[idx]; i < row_offsets[idx + 1]; i++)\n {\n if (diag && i == row_offsets[idx]) //if outside diag and this is the first nonzero in a non-empty row, overwrite diagonal value\n {\n for (int j = 0; j < block_size; j++)\n {\n new_values[(new_row_offsets[size] + halo_offset + idx)*block_size + j] = values[(row_offsets[size] + halo_offset + idx) * block_size + j];\n }\n }\n\n int col_idx = col_indices[i];\n\n if (append_offset == -1 && (col_idx != halo_offset + idx)) {printf(\"ERROR: append offset is -1 but row has nonzeros in it old %d to %d new %d to %d\\n\", row_offsets[idx], row_offsets[idx + 1], new_row_offsets[idx], new_row_offsets[idx + 1]); append_offset = 0;}\n\n if (col_idx != halo_offset + idx || add_diag)\n {\n new_col_indices[append_offset] = col_idx;\n\n for (int j = 0; j < block_size; j++)\n {\n new_values[append_offset * block_size + j] = values[i * block_size + j];\n }\n\n append_offset++;\n }\n }\n\n idx += blockDim.x * gridDim.x;\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::createCoarseB2LMaps(std::vector &in_coarse_B2L_maps)\n{\n Matrix &A = this->getA();\n m_num_all_aggregates = m_num_aggregates;\n int num_neighbors = A.manager->neighbors.size();\n IndexType max_b2l = 0;\n\n for (int i = 0; i < num_neighbors; i++ ) { max_b2l = max_b2l > A.manager->B2L_rings[i][1] ? max_b2l : A.manager->B2L_rings[i][1]; }\n\n IVector B2L_aggregates(max_b2l);\n IVector indices(max_b2l);\n\n for (int i = 0; i < num_neighbors; i++ )\n {\n int size = A.manager->B2L_rings[i][1];\n thrust_wrapper::fill(B2L_aggregates.begin(), B2L_aggregates.begin() + size, 0);\n thrust_wrapper::sequence(indices.begin(), indices.begin() + size);\n //substitute coarse aggregate indices for fine boundary nodes\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(this->m_aggregates.begin(), A.manager->B2L_maps[i].begin()),\n amgx::thrust::make_permutation_iterator(this->m_aggregates.begin(), A.manager->B2L_maps[i].begin() + size),\n B2L_aggregates.begin());\n //find the unique ones\n amgx::thrust::sort_by_key(B2L_aggregates.begin(), B2L_aggregates.begin() + size, indices.begin());\n IndexType num_unique = amgx::thrust::unique_by_key(B2L_aggregates.begin(), B2L_aggregates.begin() + size, indices.begin()).first - B2L_aggregates.begin();\n in_coarse_B2L_maps[i].resize(num_unique);\n //sort it back so we have the original ordering\n amgx::thrust::sort_by_key(indices.begin(), indices.begin() + num_unique, B2L_aggregates.begin());\n amgx::thrust::copy(B2L_aggregates.begin(), B2L_aggregates.begin() + num_unique, in_coarse_B2L_maps[i].begin());\n }\n\n cudaCheckError();\n}\n\n\n__global__ void populate_coarse_boundary(INDEX_TYPE *flags, INDEX_TYPE *indices, INDEX_TYPE *maps, INDEX_TYPE *output, INDEX_TYPE size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n output[flags[maps[indices[idx]]]] = maps[indices[idx]];\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void flag_coarse_boundary(INDEX_TYPE *flags, INDEX_TYPE *indices, INDEX_TYPE *maps, INDEX_TYPE size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n flags[maps[indices[idx]]] = 1;\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void flag_halo_indices(INDEX_TYPE *flags, INDEX_TYPE *indices, INDEX_TYPE offset, INDEX_TYPE size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n flags[indices[idx] - offset] = 1;\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void apply_halo_aggregate_indices(INDEX_TYPE *flags, INDEX_TYPE *indices, INDEX_TYPE *output, INDEX_TYPE offset, INDEX_TYPE aggregates_offset, INDEX_TYPE size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n output[idx] = flags[indices[idx] - offset] + aggregates_offset;\n idx += blockDim.x * gridDim.x;\n }\n}\n\n// renumbering the aggregates/communicationg with neighbors\ntemplate \nvoid Aggregation_AMG_Level_Base::setNeighborAggregates()\n{\n Matrix &A = this->getA();\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n m_num_all_aggregates = m_num_aggregates;\n\n /* WARNING: the matrix reordering always happens inside createRenumbering routine. There are three ways to get to this routine\n 1. matrix_upload_all -> uploadMatrix -> initializeUploadReorderAll -> reorder_matrix -> createRenumbering\n 2. read_system_distributed -> renumberMatrixOneRing -> reorder_matrix_owned -> createRenumbering\n 3. solver_setup -> ... -> AMG_Level::setup -> createCoarseMatrices -> setNeighborAggregates -> createRenumbering\n If you are reading the renumbering from file you might need to add intercept code in if statement below,\n otherwise this routine will exit before calling createRenumbering routine (in case of single or disjoint partitions).\n */\n if (this->getA().is_matrix_singleGPU()) { return; }\n\n int num_neighbors = A.manager->neighbors.size();\n\n //\n // Step 0 - set up coarse matrix metadata\n //\n if (Ac.manager == NULL) { Ac.manager = new DistributedManager(); }\n\n Ac.manager->resize(A.manager->neighbors.size(), 1);\n Ac.manager->A = &Ac;\n int f_size = A.get_num_rows();\n Ac.manager->setComms(A.manager->getComms());\n Ac.manager->set_global_id(A.manager->global_id());\n Ac.manager->neighbors = A.manager->neighbors;\n Ac.manager->set_base_index(A.manager->base_index());\n Ac.manager->halo_ranges = A.manager->halo_ranges;\n Ac.manager->set_index_range(A.manager->index_range());\n //-------------------------------------- Section 1 - renumbering -----------------------------------------------------------\n //\n // Step 1 - calculate coarse level B2L maps - any aggregate that has a fine boundary node, becomes a coarse boundary node\n //\n m_num_all_aggregates = m_num_aggregates;\n int vec_size = m_num_aggregates + 1; //A.manager->num_boundary_nodes()+1;\n IVector B2L_aggregates(vec_size);\n\n for (int i = 0; i < A.manager->neighbors.size(); i++)\n {\n thrust_wrapper::fill(B2L_aggregates.begin(), B2L_aggregates.begin() + vec_size, 0);\n int size = A.manager->B2L_rings[i][1];\n int block_size = 128;\n int grid_size = std::min( 4096, ( size + block_size - 1 ) / block_size);\n flag_coarse_boundary <<< grid_size, block_size>>>(B2L_aggregates.raw(), A.manager->B2L_maps[i].raw(), this->m_aggregates.raw(), size);\n cudaCheckError();\n thrust_wrapper::exclusive_scan(B2L_aggregates.begin(), B2L_aggregates.begin() + vec_size, B2L_aggregates.begin());\n (Ac.manager->B2L_maps)[i].resize(B2L_aggregates[vec_size - 1]);\n populate_coarse_boundary <<< grid_size, block_size>>>(B2L_aggregates.raw(), A.manager->B2L_maps[i].raw(), this->m_aggregates.raw(), Ac.manager->B2L_maps[i].raw(), size);\n }\n\n cudaCheckError();\n\n for (int i = 0; i < num_neighbors; i++)\n {\n Ac.manager->B2L_rings[i].resize(2);\n Ac.manager->B2L_rings[i][0] = 0;\n Ac.manager->B2L_rings[i][1] = Ac.manager->B2L_maps[i].size();\n }\n\n DistributedArranger *prep = new DistributedArranger;\n prep->initialize_B2L_maps_offsets(Ac, 1);\n delete prep;\n Ac.set_num_rows(m_num_aggregates);\n IVector renumbering(m_num_aggregates + 1); /* +1 is actually not needed, it will be resized in createRenumbering */\n Ac.manager->createRenumbering(renumbering);\n //\n // Step 2 - renumber aggregates, so boundary nodes will have higher index than interior ones (based on the renumberiong we have been calculating)\n //\n /* WARNING: 1. Thrust scatter and gather routines seem more appropriate here, but they implicitly assume that the input\n and output have certain size correlation, which is not matched by vectors in our case. The only remaining option\n is to use make_permutation as is done below. Example of Thrust scatter and gather calls\n IVector ttt(f_size,-1);\n amgx::thrust::scatter(this->m_aggregates.begin(), this->m_aggregates.begin()+f_size, renumbering.begin(), ttt.begin());\n amgx::thrust::gather(renumbering.begin(), renumbering.end(), this->m_aggregates.begin(), ttt.begin());\n amgx::thrust::copy(ttt.begin(), ttt.end(), this->m_aggregates.begin());\n\n 2. The original thrust composite call is illegal because it uses the same array (m_aggregates) for input and output.\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(renumbering.begin(), this->m_aggregates.begin()),\n amgx::thrust::make_permutation_iterator(renumbering.begin(), this->m_aggregates.begin()+f_size),\n this->m_aggregates.begin());\n Although it somehow still works, it is much safer to use explicit temporary storage for the intermediate result.\n */\n /* WARNING: must save unreordered aggregates for later use before reordering them. */\n IVector unreordered_aggregates(this->m_aggregates);\n /* WARNING: change Thrust call to explicitly use temporary storage for the intermediate result. The earlier version is illegal, but somehow still works. */\n IVector ttt(f_size, -1);\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(renumbering.begin(), this->m_aggregates.begin()),\n amgx::thrust::make_permutation_iterator(renumbering.begin(), this->m_aggregates.begin() + f_size),\n ttt.begin());\n amgx::thrust::copy(ttt.begin(), ttt.end(), this->m_aggregates.begin());\n cudaCheckError();\n\n //we don't need renumbering anymore, it will be identity on the coarse level\n\n //-------------------------------------- Section 2 - communication -----------------------------------------------------------\n\n //\n // Step 3 - populate aggregates_fine_idx, which stores for every fine node the original global index of the aggregate (which is lowest global index of nodes aggregated together)\n //\n\n //\n // These are different when we do /don't do matrix halo exchanges - when we do we need global indices to match nodes,\n // and in this case Ac after computeA will not have the same ordering of halo nodes as after prepareNextLevel_full.\n // However when we do not do matrix halo exchange we are only interested in the ordering of halo nodes on the coarse level,\n // and we can get that by exchanging the (already renumbered) aggregates vector.\n //\n if (m_matrix_halo_exchange == 2)\n {\n //Find original global indices of nodes that have the minimum id in the aggregates.\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(A.manager->inverse_renumbering.begin(), this->m_aggregates_fine_idx.begin()),\n amgx::thrust::make_permutation_iterator(A.manager->inverse_renumbering.begin(), this->m_aggregates_fine_idx.begin() + f_size),\n this->m_aggregates_fine_idx.begin());\n thrust_wrapper::transform(this->m_aggregates_fine_idx.begin(),\n this->m_aggregates_fine_idx.begin() + f_size,\n amgx::thrust::constant_iterator(A.manager->base_index()),\n this->m_aggregates_fine_idx.begin(),\n amgx::thrust::plus());\n //communicate\n this->m_aggregates_fine_idx.set_block_dimx(1);\n this->m_aggregates_fine_idx.set_block_dimy(1);\n m_aggregates_fine_idx.dirtybit = 1;\n A.manager->exchange_halo(m_aggregates_fine_idx, this->tag * 100 + 1 * 10 + 0);\n }\n else\n {\n //communicate\n this->m_aggregates.set_block_dimx(1);\n this->m_aggregates.set_block_dimy(1);\n m_aggregates.dirtybit = 1;\n /* WARNING: you should exchange unreordered aggregates, and append them to your own reordered aggregates, to conform to asusmptions done by distributed_mamanger. */\n //A.manager->exchange_halo(m_aggregates, this->tag*100+1*10+0); //wrong\n A.manager->exchange_halo(unreordered_aggregates, this->tag * 100 + 1 * 10 + 0);\n amgx::thrust::copy(unreordered_aggregates.begin() + f_size, unreordered_aggregates.end(), this->m_aggregates.begin() + f_size);\n }\n\n cudaCheckError();\n //\n // Step 4 - consolidate neighbors' aggregates into own list to be able to perform Galerkin product with the n-ring halo\n //\n IVector &exchanged_aggregates = m_matrix_halo_exchange == 2 ? this->m_aggregates_fine_idx : this->m_aggregates;\n int min_index = amgx::thrust::reduce(exchanged_aggregates.begin() + A.manager->halo_offsets[0], exchanged_aggregates.begin() + A.manager->halo_offsets[num_neighbors], (int)0xFFFFFFF, amgx::thrust::minimum());\n int max_index = amgx::thrust::reduce(exchanged_aggregates.begin() + A.manager->halo_offsets[0], exchanged_aggregates.begin() + A.manager->halo_offsets[num_neighbors], (int)0, amgx::thrust::maximum());\n cudaCheckError();\n int s_size = max_index - min_index + 2;\n IVector scratch(s_size);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = A.manager->halo_offsets[i + 1] - A.manager->halo_offsets[i];\n //Could also use local minimums to perform the same operation. The results are the same.\n //int min_local = amgx::thrust::reduce(exchanged_aggregates.begin()+A.manager->halo_offsets[i], exchanged_aggregates.begin()+A.manager->halo_offsets[i+1], (int)0xFFFFFFF, amgx::thrust::minimum());\n thrust_wrapper::fill(scratch.begin(), scratch.begin() + s_size, 0);\n int block_size = 128;\n int grid_size = std::min( 4096, ( size + block_size - 1 ) / block_size);\n flag_halo_indices <<< grid_size, block_size>>>(scratch.raw(), exchanged_aggregates.raw() + A.manager->halo_offsets[i], min_index /*min_local*/, size);\n cudaCheckError();\n thrust_wrapper::exclusive_scan(scratch.begin(), scratch.begin() + s_size, scratch.begin());\n apply_halo_aggregate_indices <<< grid_size, block_size>>>(scratch.raw(), exchanged_aggregates.raw() + A.manager->halo_offsets[i], this->m_aggregates.raw() + A.manager->halo_offsets[i], min_index /*min_local*/, m_num_all_aggregates, size);\n cudaCheckError();\n Ac.manager->halo_offsets[i] = m_num_all_aggregates;\n m_num_all_aggregates += scratch[s_size - 1];\n }\n\n cudaCheckError();\n Ac.manager->halo_offsets[num_neighbors] = m_num_all_aggregates;\n}\n\n//TODO: The consolidate and unconsolidate parts could be made more efficient by only sending the\n// nonzero values\ntemplate \nvoid Aggregation_AMG_Level_Base::consolidateVector(VVector &x)\n{\n int my_id = this->getA().manager->global_id();\n\n if (this->getA().manager->isRootPartition())\n {\n // Here all partitions being consolidated should have same vector size, see TODO above\n INDEX_TYPE num_parts = this->getA().manager->getNumPartsToConsolidate();\n\n for (int i = 0; i < num_parts; i++)\n {\n int current_part = this->getA().manager->getPartsToConsolidate()[i];\n\n // Vector has been set to correct size\n if (current_part != my_id)\n {\n //printf(\"Root partition %d receiving %d -> %d and %d -> %d (total %d)\\n\", this->getA().manager->global_id(), this->getA().manager->getConsolidationArrayOffsets()[i], this->getA().manager->getConsolidationArrayOffsets()[i+1], this->getA().manager->getConsolidationArrayOffsets()[num_parts+i], this->getA().manager->getConsolidationArrayOffsets()[num_parts+i+1], (int)x.size()/x.get_block_size());\n this->getA().manager->getComms()->recv_vector(x, current_part, 10000 + current_part, x.get_block_size()*this->getA().manager->getConsolidationArrayOffsets()[i], x.get_block_size() * (this->getA().manager->getConsolidationArrayOffsets()[i + 1] - this->getA().manager->getConsolidationArrayOffsets()[i]));\n this->getA().manager->getComms()->recv_vector(x, current_part, 20000 + current_part, x.get_block_size()*this->getA().manager->getConsolidationArrayOffsets()[num_parts + i], x.get_block_size() * (this->getA().manager->getConsolidationArrayOffsets()[num_parts + i + 1] - this->getA().manager->getConsolidationArrayOffsets()[num_parts + i]));\n }\n }\n }\n else\n {\n int my_destination_part = this->getA().manager->getMyDestinationPartition();\n int i_off, i_size, b_off, b_size;\n this->getA().manager->getConsolidationOffsets(&i_off, &i_size, &b_off, &b_size);\n // Here all partitions being consolidated should have same vector size, see TODO above\n this->getA().manager->getComms()->send_vector_async(x, my_destination_part, 10000 + my_id, i_off * x.get_block_size(), i_size * x.get_block_size());\n this->getA().manager->getComms()->send_vector_async(x, my_destination_part, 20000 + my_id, b_off * x.get_block_size(), b_size * x.get_block_size());\n }\n}\n\n//TODO: The consolidate and unconsolidate parts could be made more efficient by only sending the\n// nonzero values\ntemplate \nvoid Aggregation_AMG_Level_Base::unconsolidateVector(VVector &x)\n{\n if (this->getA().manager->isRootPartition())\n {\n INDEX_TYPE num_parts = this->getA().manager->getNumPartsToConsolidate();\n\n for (int i = 0; i < num_parts; i++)\n {\n int current_part = this->getA().manager->getPartsToConsolidate()[i];\n\n // Vector has been set to correct size\n if (current_part != this->getA().manager->global_id())\n {\n this->getA().manager->getComms()->send_vector_async(x, current_part, 30000 + current_part, x.get_block_size()*this->getA().manager->getConsolidationArrayOffsets()[i], x.get_block_size() * (this->getA().manager->getConsolidationArrayOffsets()[i + 1] - this->getA().manager->getConsolidationArrayOffsets()[i]));\n this->getA().manager->getComms()->send_vector_async(x, current_part, 40000 + current_part, x.get_block_size()*this->getA().manager->getConsolidationArrayOffsets()[num_parts + i], x.get_block_size() * (this->getA().manager->getConsolidationArrayOffsets()[num_parts + i + 1] - this->getA().manager->getConsolidationArrayOffsets()[num_parts + i]));\n }\n }\n }\n else\n {\n int my_destination_part = this->getA().manager->getMyDestinationPartition();\n // Vector x is of unknown size\n int i_off, i_size, b_off, b_size;\n this->getA().manager->getConsolidationOffsets(&i_off, &i_size, &b_off, &b_size);\n this->getA().manager->getComms()->recv_vector(x, my_destination_part, 30000 + this->getA().manager->global_id(), i_off * x.get_block_size(), i_size * x.get_block_size());\n this->getA().manager->getComms()->recv_vector(x, my_destination_part, 40000 + this->getA().manager->global_id(), b_off * x.get_block_size(), b_size * x.get_block_size());\n }\n}\n\n\ntemplate \nvoid Aggregation_AMG_Level_Base::createCoarseVertices()\n{\n //Set the aggregates\n this->m_selector->setAggregates(this->getA(), this->m_aggregates, this->m_aggregates_fine_idx, this->m_num_aggregates);\n\n if ( this->m_print_aggregation_info )\n {\n this->m_selector->printAggregationInfo( this->m_aggregates, this->m_aggregates_fine_idx, this->m_num_aggregates );\n }\n\n this->getA().template setParameter< int > (\"aggregates_num\", this->m_num_aggregates); // ptr to aaggregates\n}\n\n// Creating the next level\ntemplate \nvoid Aggregation_AMG_Level_Base::createCoarseMatrices()\n{\n Matrix &A = this->getA();\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n\n /* WARNING: do not recompute prolongation (P) and restriction (R) when you\n are reusing the level structure (structure_reuse_levels > 0).\n Notice that in aggregation path, prolongation P is implicit,\n and is used through the aggregates array. */\n\n bool const consolidation_level = !A.is_matrix_singleGPU() && this->isConsolidationLevel();\n\n // bookkeeping for the coarse grid: renumber aggregates,\n // if consolidation compute consolidated halo-offsets, etc\n if (!this->isReuseLevel())\n {\n if (consolidation_level)\n {\n // Consolidation-path steps 1-9\n this->consolidationBookKeeping();\n }\n else\n {\n this->setNeighborAggregates();\n }\n }\n\n this->getA().setView(ALL);\n\n // Compute restriction operator\n // TODO: computing the restriction operator could be merged with the selector to save some work\n // If we reuse the level we keep the previous restriction operator\n if (this->isReuseLevel() == false)\n {\n computeRestrictionOperator();\n }\n\n Ac.set_initialized(0);\n Ac.copyAuxData(&A);\n this->m_coarseAGenerator->computeAOperator(A, Ac, this->m_aggregates, this->m_R_row_offsets, this->m_R_column_indices, this->m_num_all_aggregates);\n Ac.setColsReorderedByColor(false);\n Ac.setView(FULL);\n\n if (consolidation_level)\n {\n // Consolidation-path Steps 11-12, send matrices to root, consolidate, final bookkeeping\n this->consolidateCoarseGridMatrix();\n }\n else\n {\n this->prepareNextLevelMatrix(A, Ac);\n }\n\n A.setView(OWNED);\n Ac.setView(OWNED);\n\n this->m_next_level_size = this->m_num_all_aggregates * Ac.get_block_dimy();\n\n if (this->m_print_aggregation_info)\n {\n MatrixAnalysis ana(&Ac);\n ana.aggregatesQuality2(this->m_aggregates, this->m_num_aggregates, A);\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::consolidationBookKeeping()\n{\n Matrix &A = this->getA();\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n\n\n int num_parts, num_fine_neighbors, my_id;\n\n if (!A.is_matrix_singleGPU())\n {\n num_parts = A.manager->getComms()->get_num_partitions();\n num_fine_neighbors = A.manager->neighbors.size();\n my_id = A.manager->global_id();\n }\n else\n {\n num_parts = 1;\n num_fine_neighbors = 0;\n my_id = 0;\n }\n\n // ----------------------------------------------------\n // Consolidate multiple fine matrices into one coarse matrix\n // ----------------------------------------------------\n // ----------------\n // Step 1\n // Decide which partitions should be merged together, store in destination_partitions vector\n // ---------------\n IVector_h &destination_part = A.manager->getDestinationPartitions();\n int my_destination_part = A.manager->getMyDestinationPartition();\n\n if (my_destination_part >= num_parts)\n {\n FatalError(\"During consolidation, sending data to partition that doesn't exist\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // Create mapping from coarse partition indices (ranks on the coarse consolidated level) to partition indices on the fine level (ranks on the fine level)\n IVector_h coarse_part_to_fine_part = destination_part;\n amgx::thrust::sort(coarse_part_to_fine_part.begin(), coarse_part_to_fine_part.end());\n cudaCheckError();\n coarse_part_to_fine_part.erase(thrust::unique(coarse_part_to_fine_part.begin(), coarse_part_to_fine_part.end()), coarse_part_to_fine_part.end());\n cudaCheckError();\n //Then, the number of coarse partitions is simply the size of this vector\n int num_coarse_partitions = coarse_part_to_fine_part.size();\n // Create mapping from fine partition indices to coarse partition indices, with fine partitions that are merging together having the same coarse indices\n IVector_h fine_part_to_coarse_part(num_parts);\n amgx::thrust::lower_bound(coarse_part_to_fine_part.begin(), coarse_part_to_fine_part.end(), destination_part.begin(), destination_part.end(), fine_part_to_coarse_part.begin());\n cudaCheckError();\n // Create mapping from this specific partition's neighbors to consolidated coarse neighbors, but using their fine index (aka. destination partition indices for my neighbors)\n IVector_h fine_neigh_to_fine_part;\n A.manager->createNeighToDestPartMap(fine_neigh_to_fine_part, A.manager->neighbors, destination_part, num_fine_neighbors);\n // Create mapping from consolidated coarse neighbors to fine partition indices (even if the current partition is not going to be a root)\n IVector_h coarse_neigh_to_fine_part;\n int num_coarse_neighbors;\n A.manager->createConsolidatedNeighToPartMap(coarse_neigh_to_fine_part, fine_neigh_to_fine_part, my_destination_part, destination_part, num_coarse_neighbors);\n // Create mapping from fine neighbors to coarse neighbors, with fine neighbors this partition is merging with labeled with -1\n IVector_h fine_neigh_to_coarse_neigh;\n A.manager->createNeighToConsNeigh(fine_neigh_to_coarse_neigh, coarse_neigh_to_fine_part, fine_neigh_to_fine_part, my_destination_part, num_fine_neighbors);\n /*\n EXAMPLE\n Take the following partition graph (that describes connections between partitions, vertices are the partitions themselves), this is the same graph that is used in the setup example\n number of partitions num_parts=12\n CSR row_offsets [0 4 8 13 21 25 32 36 41 46 50 57 61]\n CSR col_indices [0 1 3 8\n 0 1 2 3\n 1 2 3 4 5\n 0 1 2 3 4 5 8 10\n 2 4 5 6\n 2 3 4 5 6 7 10\n 4 5 6 7\n 5 6 7 9 10\n 0 3 8 10 11\n 7 9 10 11\n 3 5 7 8 9 10 11\n 8 9 10 11]\n destination_part = [0 0 0 0 4 4 4 4 8 8 8 8]\n coarse_part_to_fine_part = [0 4 8] num_coarse_partitions = 3\n fine_part_to_coarse_part = [0 0 0 0 1 1 1 1 2 2 2 2]\n original neighbor lists correspond to the rows of the matrix, minus the diagonal elements: (part 0)[1 3 8] (part 3)[0 1 2 4 5 8 10] (part 10)[3 5 7 8 9 11]\n fine_neigh_to_fine_part (part 0)[0 0 2] (part 3)[0 0 0 0 1 2 2] (part 10)[0 1 1 2 2 2]\n coarse_neigh_to_fine_part (part 0)[8] (part 3)[4 8] (part 10)[0 4]\n fine_neigh_to_coarse_neigh (part 0)[-1 -1 0] (part 3)[-1 -1 -1 0 0 1 1] (part 10)[0 1 1 -1 -1 -1]\n */\n // --------------------------\n // Step 2\n // Create coarse B2L_maps, by mapping fine B2L maps to coarse indices using this->m_aggregates and eliminating duplicates\n // --------------------------\n std::vector coarse_B2L_maps(num_fine_neighbors);\n m_num_all_aggregates = m_num_aggregates;\n int num_neighbors_temp = A.manager->neighbors.size();\n int num_rings = A.manager->B2L_rings[0].size() - 1;\n\n if (num_rings != 1)\n {\n FatalError(\"num_rings > 1 not supported in consolidation\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n IndexType max_b2l = 0;\n\n\n for (int i = 0; i < num_neighbors_temp; i++ ) { max_b2l = max_b2l > A.manager->B2L_rings[i][1] ? max_b2l : A.manager->B2L_rings[i][1]; }\n\n IVector B2L_aggregates(max_b2l);\n IVector indices(max_b2l);\n\n //TODO: use the algorithm from setNeighborAggregates()\n for (int i = 0; i < num_neighbors_temp; i++ )\n {\n int size = A.manager->B2L_rings[i][1];\n thrust_wrapper::fill(B2L_aggregates.begin(), B2L_aggregates.begin() + size, 0);\n thrust_wrapper::sequence(indices.begin(), indices.begin() + size);\n //substitute coarse aggregate indices for fine boundary nodes\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(this->m_aggregates.begin(), A.manager->B2L_maps[i].begin()),\n amgx::thrust::make_permutation_iterator(this->m_aggregates.begin(), A.manager->B2L_maps[i].begin() + size),\n B2L_aggregates.begin());\n //find the unique ones\n amgx::thrust::sort_by_key(B2L_aggregates.begin(), B2L_aggregates.begin() + size, indices.begin());\n IndexType num_unique = amgx::thrust::unique_by_key(B2L_aggregates.begin(), B2L_aggregates.begin() + size, indices.begin()).first - B2L_aggregates.begin();\n coarse_B2L_maps[i].resize(num_unique);\n //sort it back so we have the original ordering\n amgx::thrust::sort_by_key(indices.begin(), indices.begin() + num_unique, B2L_aggregates.begin());\n amgx::thrust::copy(B2L_aggregates.begin(), B2L_aggregates.begin() + num_unique, coarse_B2L_maps[i].begin());\n }\n\n cudaCheckError();\n /*\n * EXAMPLE\n say, partition 3 has the following coarse B2L_maps:\n neighbors [0 1 2 4 5 8 10]\n B2L_maps[0(=0)] = [6 7 8]\n B2L_maps[1(=1)] = [8 9 10]\n B2L_maps[2(=2)] = [10 11 12 13]\n B2L_maps[3(=4)] = [13 14 15]\n B2L_maps[4(=5)] = [15 16 17]\n B2L_maps[5(=8)] = [6 18 19]\n B2L_maps[6(=10)] = [17 20 19]\n */\n // ---------------------------------------------------\n // Step 3\n // create new B2L maps for each merged destination neighbor and drop B2L maps to neighbors we are merging with\n // ---------------------------------------------------\n std::vector dest_coarse_B2L_maps;\n A.manager->consolidateB2Lmaps(dest_coarse_B2L_maps, coarse_B2L_maps, fine_neigh_to_coarse_neigh, num_coarse_neighbors, num_fine_neighbors);\n /*\n * EXAMPLE\n Then, merging the coarse B2L maps on partition 3, we get:\n coarse_neigh_to_fine_part [4 8]\n dest_coarse_B2L_maps[0(=4)] = [13 14 15 16 17]\n dest_coarse_B2L_maps[1(=8)] = [6 17 18 19 20]\n */\n // -----------------------\n // Step 4\n // Create interior-boundary renumbering of aggregates according to dest_coarse_B2L_maps\n // -----------------------\n // Now renumber the aggregates with all interior aggregates first, boundary aggregates second\n int num_interior_aggregates; //returned by createAggregatesRenumbering\n int num_boundary_aggregates; //returned by createAggregatesRenumbering\n IVector renumbering; //returned by createAggregatesRenumbering\n // Following calls create renumbering array and modifies B2L_maps\n A.manager->createAggregatesRenumbering(renumbering, dest_coarse_B2L_maps, this->m_num_aggregates, num_coarse_neighbors, num_interior_aggregates, num_boundary_aggregates, num_rings);\n /*\n * EXAMPLE\n Partition 3 will get a renumbering vector of size 21, for the 21 owned agggregates:\n [0 1 2 3 4 5 17 6 7 8 9 10 11 12 13 14 15 16 18 19 20]\n num_interior_aggregates = 12\n num_boundary_aggregates = 9\n */\n // -------------------------------------------------\n // Step 5\n // Determine whether root partition, make list of partitions merged into one\n // ------------------------------------------------\n // Check if I'm root partition and how fine partitions (including myself) are merging into me\n // bool is_root_partition = false;\n bool &is_root_partition = this->m_is_root_partition;\n is_root_partition = false; \n int num_fine_parts_to_consolidate = 0;\n // IVector_h fine_parts_to_consolidate;\n IVector_h &fine_parts_to_consolidate = this->m_fine_parts_to_consolidate;\n\n for (int i = 0; i < num_parts; i++)\n {\n if (destination_part[i] == my_id)\n {\n is_root_partition = true;\n num_fine_parts_to_consolidate++;\n }\n }\n\n fine_parts_to_consolidate.resize(num_fine_parts_to_consolidate);\n int count = 0;\n\n for (int i = 0; i < num_parts; i++)\n {\n if (destination_part[i] == my_id)\n {\n fine_parts_to_consolidate[count] = i;\n count++;\n }\n }\n\n //save this information as state, as this will also be required during solve for restriction/prolongation\n A.manager->setIsRootPartition(is_root_partition);\n A.manager->setNumPartsToConsolidate(num_fine_parts_to_consolidate);\n A.manager->setPartsToConsolidate(fine_parts_to_consolidate);\n\n // Create a new distributed communicator for coarse levels that only contains active partitions\n if (Ac.manager == NULL)\n {\n Ac.manager = new DistributedManager();\n }\n\n Ac.manager->setComms(A.manager->getComms()->Clone());\n Ac.manager->getComms()->createSubComm(coarse_part_to_fine_part, is_root_partition);\n\n\n /*\n * EXAMPLE\n isRootPartition is true for partitions 0,4,8 false for others\n num_fine_parts_to_consolidate = 4 for partitions 0,4,8\n fine_parts_to_consolidate (part 0)[0 1 2 3] (part 4)[4 5 6 7] (part 8)[8 9 10 11]\n */\n // ----------------------\n // Step 6\n // Compute number of interior, boundary and total nodes in the consolidated coarse matrix. Create offsets so that partitions being merged together will have their aggregate indices ordered like this:\n // [num_interior(fine_parts_to_consolidate[0]] num_interior(fine_parts_to_consolidate[1]] ... num_interior(fine_parts_to_consolidate[num_fine_parts_to_consolidate]\n // num_boundary(fine_parts_to_consolidate[0]] num_boundary(fine_parts_to_consolidate[1]] ... num_boundary(fine_parts_to_consolidate[num_fine_parts_to_consolidate] ]\n // ----------------------\n // Gather to get number of interior/boundary aggregates of neighbors I will merge with\n // std::vector vertex_counts;\n std::vector &vertex_counts = this->m_vertex_counts;\n // int interior_offset, boundary_offset, total_interior_rows_in_merged, total_boundary_rows_in_merged;\n int interior_offset, boundary_offset;\n int &total_interior_rows_in_merged = this->m_total_interior_rows_in_merged;\n int &total_boundary_rows_in_merged = this->m_total_boundary_rows_in_merged;\n int total_rows_in_merged;\n //Computes these offsets on the root, sends them back\n A.manager->computeConsolidatedOffsets(my_id, my_destination_part, is_root_partition, num_interior_aggregates, num_boundary_aggregates, vertex_counts, fine_parts_to_consolidate, num_fine_parts_to_consolidate, interior_offset, boundary_offset, total_interior_rows_in_merged, total_boundary_rows_in_merged, total_rows_in_merged, A.manager->getComms());\n //Partitions save these offsets, as it will be required during solve restriction/prolongation\n A.manager->setConsolidationOffsets(interior_offset, num_interior_aggregates, boundary_offset + num_interior_aggregates, num_boundary_aggregates);\n /*\n * EXAMPLE\n For root partition 0, say we have the following interior/boundary counts (note that partition 1 has 0 boundary, as it is only connected to partitions it is merging with)\n part 0 - interior: 10 boundary 3\n part 1 - interior: 18\n part 2 - interior: 10 boundary 16\n part 3 - interior: 12 boundary 9\n interior_offset for partitions 0,1,2,3: 0 10 28 38 (total_interior_rows_in_merged 50)\n boundary_offset for partitions 0,1,2,3: 0 3 3 19 (total_boundary_rows_in_merged 28)\n */\n // ----------------------\n // Step 7\n // Each partition renumbers its aggregates and dest_coarse_B2L_maps using offsets computed in Step 6 and permutation in Step 4\n // ----------------------\n // Kernel to renumber the aggregates\n int block_size = 128;\n int grid_size = std::min( 4096, ( A.manager->halo_offsets[0] + block_size - 1 ) / block_size);\n renumberAggregatesKernel <<< grid_size, block_size >>>(renumbering.raw(), interior_offset, boundary_offset, this->m_aggregates.raw(), A.manager->halo_offsets[0], num_interior_aggregates, renumbering.size());\n cudaCheckError();\n\n for (int i = 0; i < num_coarse_neighbors; i++)\n {\n thrust_wrapper::transform(dest_coarse_B2L_maps[i].begin(),\n dest_coarse_B2L_maps[i].end(),\n amgx::thrust::constant_iterator(boundary_offset),\n dest_coarse_B2L_maps[i].begin(),\n amgx::thrust::plus());\n }\n\n cudaCheckError();\n /*\n * EXAMPLE\n Partition 3 had a renumbering vector:\n [0 1 2 3 4 5 17 6 7 8 9 10 11 12 13 14 15 16 18 19 20]\n which is now adjusted to account for the consolidated coarse matrices' indices:\n [38 39 40 41 42 43 74 44 45 46 47 48 49 69 70 71 72 73 75 76 77]\n And the dest_coarse_B2L_maps, which looked like:\n dest_coarse_B2L_maps[0(=4)] = [13 14 15 16 17]\n dest_coarse_B2L_maps[1(=8)] = [6 17 18 19 20]\n is now:\n dest_coarse_B2L_maps[0(=4)] = [69 70 71 72 73]\n dest_coarse_B2L_maps[1(=8)] = [74 73 75 76 77]\n */\n // -------------------------------------------------\n // Step 8\n // Send dest_coarse_B2L_maps to root partitions\n // ------------------------------------------------\n // Each fine partition sends to its root the number of coarse neighbors it has, their ids, and the number of boundary nodes for each coarse neighbor\n IVector_h num_bdy_per_coarse_neigh(num_coarse_neighbors);\n\n for (int i = 0; i < num_coarse_neighbors; i++)\n {\n num_bdy_per_coarse_neigh[i] = dest_coarse_B2L_maps[i].size();\n }\n\n IVector_h consolidated_coarse_neigh_to_fine_part; //consolidated list of coarse neighbors for the root partition, using fine partition indices\n int num_consolidated_neighbors = 0;\n // std::vector consolidated_B2L_maps; //concatenates dest_coarse_B2L_maps received from partitions that are merging into the same root and pointing to the same destination coarse neighbor\n std::vector &consolidated_B2L_maps = this->m_consolidated_B2L_maps;\n A.manager->consolidateB2LmapsOnRoot(num_consolidated_neighbors, consolidated_B2L_maps, consolidated_coarse_neigh_to_fine_part, dest_coarse_B2L_maps, coarse_neigh_to_fine_part, num_bdy_per_coarse_neigh, fine_parts_to_consolidate, num_fine_parts_to_consolidate, my_id, my_destination_part, is_root_partition, num_coarse_neighbors, A.manager->getComms());\n //\n // Step 9 - figuring out halo aggregate IDs\n //\n //Now we need to update halo aggregate IDs - this is just a halo exchange on this->m_aggregates between partitions\n //that are being merged together, but we need to send other halos to the root to come up with the halo renumbering\n //TODO: separate transactions, send \"real halo\" to the root nodes (coarse neighbors) immediately\n //Step 9.1: takes care of synchronizing the aggregate IDs between partitions we are merging together and got consistent halo aggregate IDs for neighbor we are not merging with (which are going to be sent to the root in 9.2)\n A.manager->exchange_halo(this->m_aggregates, 6666);\n /*\n * EXAMPLE 2\n This example is independent from the previous ones.\n Say partition 0 and 1 are merging (into 0) partition 0 is neighbors with 1,2,3 and partition 1 is neighbors with 0,3,4\n Partitions 3 and 4 are merging (into partition 3) and partition 2 is not merging with anyone.\n This example details the renumbering of halo indices on partition 0 and partition 1.\n After the exchange halo, we have:\n this->m_aggregates on partition 0:\n [(fine interior nodes) (fine boundary nodes) (fine halo from part 1) (fine halo from part 2) (fine halo from part 3)]\n [(fine interior nodes) (fine boundary nodes) (13 13 15) (12 15 17) (14 16 18)]\n aggregates on partition 1:\n [(fine interior nodes) (fine boundary nodes) (fine halo from part 0) (fine halo from part 3) (fine halo from part 4)]\n [(fine interior nodes) (fine boundary nodes) (14 16 17) (18 19 19) (15 15 17)]\n indices in (fine halo from part 0) and (fine halo from part 1) actually contain interior aggregate indices (if they are not connected to partitions 2,3 or 4), because the boundary is disappearing there.\n Indices in halo regions contain remote-local indices.\n\n This example is used throughout consolidateAndRenumberHalos\n */\n //Step 9.2 - 9.5\n // IVector_h halo_offsets(num_consolidated_neighbors + 1, 0);\n IVector_h &halo_offsets = this->m_consolidated_halo_offsets;\n halo_offsets = IVector_h(num_consolidated_neighbors + 1, 0);\n A.manager->consolidateAndRenumberHalos(this->m_aggregates, A.manager->halo_offsets, halo_offsets, A.manager->neighbors, num_fine_neighbors, consolidated_coarse_neigh_to_fine_part, num_consolidated_neighbors, destination_part, my_destination_part, is_root_partition, fine_parts_to_consolidate, num_fine_parts_to_consolidate, num_parts, my_id, total_rows_in_merged, this->m_num_all_aggregates, A.manager->getComms());\n\n if (is_root_partition)\n {\n for (int i = 0; i < consolidated_B2L_maps.size(); i++)\n {\n amgx::thrust::sort(consolidated_B2L_maps[i].begin(), consolidated_B2L_maps[i].end());\n }\n\n this->m_consolidated_neighbors.resize(num_consolidated_neighbors);\n for (int i = 0; i < num_consolidated_neighbors; i++)\n {\n this->m_consolidated_neighbors[i] = fine_part_to_coarse_part[consolidated_coarse_neigh_to_fine_part[i]];\n }\n \n cudaCheckError();\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::consolidateCoarseGridMatrix()\n{\n Matrix &A = this->getA();\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n\n int my_id = A.manager->global_id();\n \n IVector_h &destination_part = A.manager->getDestinationPartitions();\n int my_destination_part = A.manager->getMyDestinationPartition();\n\n // bookkeeping stored in AMG_Level_Base\n std::vector &vertex_counts = this->m_vertex_counts;\n IVector_h &fine_parts_to_consolidate = this->m_fine_parts_to_consolidate;\n\n // bookkeeping stored in either AMG_Level_Base or Acs' DistributedManager\n IVector_h &halo_offsets = this->isReuseLevel() ? Ac.manager->getHaloOffsets() : this->m_consolidated_halo_offsets;\n std::vector &consolidated_B2L_maps = this->isReuseLevel() ? Ac.manager->getB2Lmaps() : this->m_consolidated_B2L_maps;\n\n int num_consolidated_neighbors = this->isRootPartition() ? this->m_consolidated_neighbors.size() : 0;\n\n if (!this->isRootPartition())\n {\n A.manager->getComms()->send_vector_async(Ac.row_offsets, my_destination_part, 1111);\n A.manager->getComms()->send_vector_async(Ac.col_indices, my_destination_part, 1112);\n A.manager->getComms()->send_vector_async(Ac.values, my_destination_part, 1113);\n }\n else\n {\n int num_fine_parts_to_consolidate = fine_parts_to_consolidate.size();\n\n int total_num_rows = this->m_num_all_aggregates;\n IVector new_row_offsets(total_num_rows + 1, 0);\n\n //if diags are inside then we won't be counting those twice when computing halo row length\n if (!Ac.hasProps(DIAG))\n {\n thrust_wrapper::fill(new_row_offsets.begin() + halo_offsets[0], new_row_offsets.begin() + halo_offsets[num_consolidated_neighbors], 1);\n cudaCheckError();\n }\n\n std::vector recv_row_offsets(num_fine_parts_to_consolidate);\n std::vector num_nz(num_fine_parts_to_consolidate);\n IVector *work_row_offsets;\n std::vector index_offset_array(2 * num_fine_parts_to_consolidate + 1);\n int interior_offset = 0;\n int boundary_offset = 0;\n\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n boundary_offset += vertex_counts[i][0];\n }\n\n int max_num_nz = 0;\n\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n\n //receive row offsets\n if (current_part != my_id)\n {\n recv_row_offsets[i].resize(total_num_rows + 1);\n A.manager->getComms()->recv_vector(recv_row_offsets[i], current_part, 1111);\n work_row_offsets = &(recv_row_offsets[i]);\n num_nz[i] = (*work_row_offsets)[work_row_offsets->size() - 1];\n max_num_nz = max_num_nz > num_nz[i] ? max_num_nz : num_nz[i];\n }\n else\n {\n work_row_offsets = &(Ac.row_offsets);\n num_nz[i] = Ac.get_num_nz();\n }\n\n //Get interior row length\n thrust_wrapper::transform(work_row_offsets->begin() + interior_offset + 1,\n work_row_offsets->begin() + interior_offset + vertex_counts[i][0] + 1,\n work_row_offsets->begin() + interior_offset,\n new_row_offsets.begin() + interior_offset,\n amgx::thrust::minus());\n cudaCheckError();\n //Get boundary row length\n thrust_wrapper::transform(work_row_offsets->begin() + boundary_offset + 1,\n work_row_offsets->begin() + boundary_offset + vertex_counts[i][1] + 1,\n work_row_offsets->begin() + boundary_offset,\n new_row_offsets.begin() + boundary_offset,\n amgx::thrust::minus());\n cudaCheckError();\n //Increment halo row length by one for every nonzero that is an edge from the halo into this partition\n int size = halo_offsets[num_consolidated_neighbors] - halo_offsets[0];\n const int block_size = 128;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (size - 1) / block_size + 1);\n set_halo_rowlen <<< num_blocks, block_size>>>(work_row_offsets->raw() + halo_offsets[0], new_row_offsets.raw() + halo_offsets[0], size, Ac.hasProps(DIAG));\n cudaCheckError();\n index_offset_array[i] = interior_offset;\n index_offset_array[num_fine_parts_to_consolidate + i] = boundary_offset;\n interior_offset += vertex_counts[i][0];\n boundary_offset += vertex_counts[i][1];\n index_offset_array[i + 1] = interior_offset;\n index_offset_array[num_fine_parts_to_consolidate + i + 1] = boundary_offset;\n }\n\n A.manager->setConsolidationArrayOffsets(index_offset_array);\n //Exclusive scan row length array to get row offsets\n thrust_wrapper::exclusive_scan(new_row_offsets.begin(), new_row_offsets.end(), new_row_offsets.begin());\n cudaCheckError();\n //Prepare to receive column indices and values\n int num_nz_consolidated = new_row_offsets[new_row_offsets.size() - 1];\n IVector recv_col_indices(max_num_nz);\n IVector new_col_indices(num_nz_consolidated);\n MVector recv_values((max_num_nz + 1 + Ac.hasProps(DIAG) * (halo_offsets[num_consolidated_neighbors] - 1))*Ac.get_block_size());\n MVector new_values((num_nz_consolidated + 1 + Ac.hasProps(DIAG) * (halo_offsets[num_consolidated_neighbors] - 1))*Ac.get_block_size());\n thrust_wrapper::fill(new_col_indices.begin() + new_row_offsets[halo_offsets[0]], new_col_indices.end(), -1); //Set all the halo col indices to -1\n\n\n if (!Ac.hasProps(DIAG)) { thrust_wrapper::fill(new_values.begin() + num_nz_consolidated * Ac.get_block_size(), new_values.end(), types::util::get_zero()); }\n\n cudaCheckError();\n IVector *work_col_indices;\n MVector *work_values;\n interior_offset = 0;\n boundary_offset = 0;\n\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n boundary_offset += vertex_counts[i][0];\n }\n\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n\n if (current_part != my_id)\n {\n A.manager->getComms()->recv_vector(recv_col_indices, current_part, 1112, 0, num_nz[i]);\n A.manager->getComms()->recv_vector(recv_values, current_part, 1113, 0, (num_nz[i] + 1 + Ac.hasProps(DIAG) * (halo_offsets[num_consolidated_neighbors] - 1))*Ac.get_block_size());\n work_col_indices = &(recv_col_indices);\n work_row_offsets = &(recv_row_offsets[i]);\n work_values = &(recv_values);\n }\n else\n {\n work_row_offsets = &(Ac.row_offsets);\n work_col_indices = &(Ac.col_indices);\n work_values = &(Ac.values);\n }\n\n //Put interior rows in place\n amgx::thrust::copy(work_col_indices->begin() + (*work_row_offsets)[interior_offset],\n work_col_indices->begin() + (*work_row_offsets)[interior_offset + vertex_counts[i][0]],\n new_col_indices.begin() + new_row_offsets[interior_offset]);\n cudaCheckError();\n amgx::thrust::copy(work_values->begin() + (*work_row_offsets)[interior_offset]*Ac.get_block_size(),\n work_values->begin() + ((*work_row_offsets)[interior_offset + vertex_counts[i][0]])*Ac.get_block_size(),\n new_values.begin() + new_row_offsets[interior_offset]*Ac.get_block_size());\n cudaCheckError();\n //Put boundary rows in place\n amgx::thrust::copy(work_col_indices->begin() + (*work_row_offsets)[boundary_offset],\n work_col_indices->begin() + (*work_row_offsets)[boundary_offset + vertex_counts[i][1]],\n new_col_indices.begin() + new_row_offsets[boundary_offset]);\n cudaCheckError();\n amgx::thrust::copy(work_values->begin() + (*work_row_offsets)[boundary_offset]*Ac.get_block_size(),\n work_values->begin() + ((*work_row_offsets)[boundary_offset + vertex_counts[i][1]])*Ac.get_block_size(),\n new_values.begin() + new_row_offsets[boundary_offset]*Ac.get_block_size());\n cudaCheckError();\n //Process halo rows (merge)\n int size = halo_offsets[num_consolidated_neighbors] - halo_offsets[0];\n const int block_size = 128;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (size - 1) / block_size + 1);\n //TODO: vectorise this kernel, will be inefficient for larger block sizes\n append_halo_nz <<< num_blocks, block_size>>>(work_row_offsets->raw() + halo_offsets[0],\n new_row_offsets.raw() + halo_offsets[0],\n work_col_indices->raw(),\n new_col_indices.raw(),\n work_values->raw(),\n new_values.raw(),\n size, Ac.hasProps(DIAG), halo_offsets[0], Ac.get_block_size());\n cudaCheckError();\n\n // Diagonals\n if (Ac.hasProps(DIAG))\n {\n // Diagonal corresponding to interior rows\n amgx::thrust::copy(work_values->begin() + (num_nz[i] + interior_offset)*Ac.get_block_size(),\n work_values->begin() + (num_nz[i] + interior_offset + vertex_counts[i][0])*Ac.get_block_size(),\n new_values.begin() + (new_row_offsets[halo_offsets[halo_offsets.size() - 1]] + interior_offset)*Ac.get_block_size());\n // Diagonal corresponding to boundary rows\n amgx::thrust::copy(work_values->begin() + (num_nz[i] + boundary_offset)*Ac.get_block_size(),\n work_values->begin() + (num_nz[i] + boundary_offset + vertex_counts[i][1])*Ac.get_block_size(),\n new_values.begin() + (new_row_offsets[halo_offsets[halo_offsets.size() - 1]] + boundary_offset)*Ac.get_block_size());\n cudaCheckError();\n }\n\n interior_offset += vertex_counts[i][0];\n boundary_offset += vertex_counts[i][1];\n }\n\n Ac.set_initialized(0);\n Ac.row_offsets = new_row_offsets;\n Ac.col_indices = new_col_indices;\n Ac.values = new_values;\n }\n\n // A new distributed communicator for coarse levels that only contains active partitions\n // has already been created in consolidatedBookKeeping!\n\n //\n // Step 12 - finalizing, bookkeping\n //\n if (this->isRootPartition())\n {\n // int my_consolidated_id = fine_part_to_coarse_part[my_id];\n int my_consolidated_id = Ac.manager->getComms()->get_global_id();\n\n if (!this->isReuseLevel())\n {\n Ac.manager->initializeAfterConsolidation(\n my_consolidated_id,\n Ac,\n this->m_consolidated_neighbors,\n this->m_total_interior_rows_in_merged,\n this->m_total_boundary_rows_in_merged,\n this->m_num_all_aggregates,\n this->m_consolidated_halo_offsets,\n this->m_consolidated_B2L_maps,\n 1,\n true);\n\n // this is now stored in Acs DistributedManager\n this->m_consolidated_neighbors.resize(0);\n this->m_consolidated_halo_offsets.resize(0);\n this->m_consolidated_B2L_maps.resize(0);\n\n Ac.manager->B2L_rings.resize(num_consolidated_neighbors + 1);\n\n for (int i = 0; i < num_consolidated_neighbors; i++)\n {\n Ac.manager->B2L_rings[i].resize(2);\n Ac.manager->B2L_rings[i][0] = 0;\n Ac.manager->B2L_rings[i][1] = consolidated_B2L_maps[i].size();\n }\n }\n\n Ac.manager->set_initialized(Ac.row_offsets);\n Ac.manager->getComms()->set_neighbors(num_consolidated_neighbors);\n int new_nnz = Ac.row_offsets[Ac.row_offsets.size() - 1];\n\n Ac.set_num_nz(new_nnz);\n Ac.set_num_cols(Ac.manager->halo_offsets[Ac.manager->halo_offsets.size() - 1]);\n Ac.set_num_rows(Ac.get_num_cols());\n\n if (A.hasProps(DIAG)) { Ac.addProps(DIAG); }\n\n Ac.computeDiagonal();\n Ac.set_initialized(1);\n }\n else\n {\n this->getA().manager->getComms()->send_vector_wait_all(Ac.row_offsets);\n this->getA().manager->getComms()->send_vector_wait_all(Ac.col_indices);\n this->getA().manager->getComms()->send_vector_wait_all(Ac.values);\n\n Ac.set_initialized(0);\n // set size of Ac to be zero\n Ac.resize(0, 0, 0, 1);\n Ac.set_initialized(1);\n }\n}\n\n// -------------------------------------------------------------\n// Explicit instantiations\n// -------------------------------------------------------------\n\n#define AMGX_CASE_LINE(CASE) template class Aggregation_AMG_Level::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n}\n\n}\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n#ifdef _WIN32\n#pragma warning (push)\n#pragma warning (disable : 4244 4267 4521)\n#endif\n#ifdef _WIN32\n#pragma warning (pop)\n#endif\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\nnamespace aggregation\n{\n\n\n// ----------------------\n// Kernels\n// ----------------------\n\ntemplate \n__global__\nvoid set_to_one_kernel(IndexType start, IndexType end, IndexType *ind, ValueType *v)\n{\n for (int tid = start + blockDim.x * blockIdx.x + threadIdx.x; tid < end; tid += gridDim.x * blockDim.x)\n {\n v[ind[tid]] = types::util::get_one();\n }\n}\n\ntemplate \n__global__\nvoid renumberAggregatesKernel(const IndexType *renumbering, const int interior_offset, const int bdy_offset, IndexType *aggregates, const int num_aggregates, const int n_interior, const int renumbering_size)\n{\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (tid < num_aggregates)\n {\n IndexType new_agg_id;\n\n if (renumbering_size == 0)\n {\n new_agg_id = aggregates[tid];\n }\n else\n {\n new_agg_id = renumbering[aggregates[tid]];\n }\n\n //if (aggregates[tid] > num_aggregates)\n //{\n //printf(\"ID %d old %d + %d = %d\\n\", tid, new_agg_id, ((new_agg_id >= n_interior) ? bdy_offset : interior_offset), new_agg_id + ((new_agg_id >= n_interior) ? bdy_offset : interior_offset));\n //}\n new_agg_id += ((new_agg_id >= n_interior) ? bdy_offset : interior_offset);\n aggregates[tid] = new_agg_id;\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel to restrict residual using csr_format\ntemplate \n__global__\nvoid restrictResidualKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueType *r, ValueType *rr, const int num_aggregates)\n{\n int jmin, jmax;\n\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_aggregates; tid += gridDim.x * blockDim.x)\n {\n ValueType temp(types::util::get_zero());\n jmin = row_offsets[tid];\n jmax = row_offsets[tid + 1];\n\n for (int j = jmin; j < jmax; j++)\n {\n int j_col = column_indices[j];\n temp = temp + r[j_col];\n }\n\n rr[tid] = temp;\n }\n}\n\n// Kernel to restrict residual using block_dia_csr_format\ntemplate \n__global__\nvoid restrictResidualBlockDiaCsrKernel(const IndexType *row_offsets, const IndexType *column_indices, const ValueType *r, ValueType *rr, const int num_aggregates)\n{\n ValueType rr_temp[bsize];\n int offset, jmin, jmax;\n\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_aggregates; tid += gridDim.x * blockDim.x)\n {\n // Initialize to zero\n#pragma unroll\n for (int m = 0; m < bsize; m++)\n {\n rr_temp[m] = types::util::get_zero();\n }\n\n jmin = row_offsets[tid];\n jmax = row_offsets[tid + 1];\n\n for (int j = jmin; j < jmax; j++)\n {\n int jcol = column_indices[j];\n offset = jcol * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n rr_temp[m] = rr_temp[m] + r[offset + m];\n }\n }\n\n offset = tid * bsize;\n#pragma unroll\n\n for (int m = 0; m < bsize; m++)\n {\n rr[offset + m] = rr_temp[m];\n };\n }\n}\n\n// Kernel to prolongate and apply the correction for csr format\ntemplate \n__global__\nvoid prolongateAndApplyCorrectionKernel(const ValueType alpha, const int num_rows, ValueType *x, const ValueType *e, const IndexType *aggregates, IndexType num_aggregates)\n{\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_rows; tid += gridDim.x * blockDim.x)\n {\n IndexType I = aggregates[tid];\n x[tid] = x[tid] + alpha * e[I];\n }\n}\n\n// Kernel to prolongate and apply the correction for block-dia-csr format\ntemplate \n__global__\nvoid prolongateAndApplyCorrectionBlockDiaCsrKernel(const ValueType alpha, const int num_block_rows, ValueType *x, const ValueType *e, const IndexType *aggregates, IndexType num_aggregates, const int bsize)\n{\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_block_rows; tid += gridDim.x * blockDim.x)\n {\n IndexType I = aggregates[tid];\n\n for (int m = 0; m < bsize; m++)\n {\n x[tid * bsize + m] = x[tid * bsize + m] + alpha * e[I * bsize + m];\n }\n }\n}\n\ntemplate \n__global__\nvoid prolongateVector(const IndexType *aggregates, const ValueType *in, ValueType *out, IndexType fine_rows, IndexType coarse_rows, int blocksize)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < fine_rows * blocksize )\n {\n int i = tid / blocksize;\n int e = tid % blocksize;\n IndexType I = aggregates[i];\n out[tid] = in[ I * blocksize + e ];\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \n__global__\nvoid applyCorrection(ValueType lambda, const ValueType *e, ValueType *x, IndexType numRows )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < numRows )\n {\n x[tid] = x[tid] + lambda * e[tid];\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// -------------------------------\n// Methods\n// ------------------------------\n\ntemplate \nvoid Aggregation_AMG_Level_Base::transfer_level(AMG_Level *ref_lvl)\n{\n Aggregation_AMG_Level_Base *ref_agg_lvl = dynamic_cast*>(ref_lvl);\n this->scale_counter = ref_agg_lvl->scale_counter;\n this->scale = ref_agg_lvl->scale;\n this->m_R_row_offsets.copy(ref_agg_lvl->m_R_row_offsets);\n this->m_R_column_indices.copy(ref_agg_lvl->m_R_column_indices);\n this->m_aggregates.copy(ref_agg_lvl->m_aggregates);\n this->m_aggregates_fine_idx.copy(ref_agg_lvl->m_aggregates_fine_idx);\n this->m_num_aggregates = ref_agg_lvl->m_num_aggregates;\n this->m_num_all_aggregates = ref_agg_lvl->m_num_all_aggregates;\n}\n\n\ntypedef std::pair mypair;\nbool comparator ( const mypair &l, const mypair &r) { return l.first < r.first; }\n\n// Method to compute R\n// General path\n// TODO: this could be merged with selector to save some computations\ntemplate \nvoid Aggregation_AMG_Level_Base::computeRestrictionOperator_common()\n{\n m_R_row_offsets.resize(m_num_all_aggregates + 1); //create one more row for the pseudo aggregate\n IVector R_row_indices(m_aggregates);\n#if AMGX_ASYNCCPU_PROOF_OF_CONCEPT\n bool use_cpu = m_aggregates.size() < 4096;\n\n if (use_cpu)\n {\n struct computeRestrictionTask : public task\n {\n Aggregation_AMG_Level_Base *self;\n IVector *R_row_indices;\n\n void run()\n {\n int N = self->m_aggregates.size();\n IVector_h R_row_indices_host(self->m_aggregates);\n std::vector pairs(N);\n\n for (int i = 0; i < N; i++)\n {\n pairs[i].first = R_row_indices_host[i];\n pairs[i].second = i;\n }\n\n std::stable_sort(pairs.begin(), pairs.end(), comparator);\n IVector_h R_column_indices(self->A->get_num_rows());\n\n for (int i = 0; i < N; i++)\n {\n R_column_indices[i] = pairs[i].second;\n R_row_indices_host[i] = pairs[i].first;\n }\n\n self->m_R_column_indices = R_column_indices;\n *R_row_indices = R_row_indices_host;\n }\n };\n computeRestrictionTask *t = new computeRestrictionTask();\n t->self = this;\n t->R_row_indices = &R_row_indices;\n t->run();\n delete t;\n }\n else\n#endif\n {\n m_R_column_indices.resize(this->A->get_num_rows());\n thrust_wrapper::sequence(m_R_column_indices.begin(), m_R_column_indices.end());\n cudaCheckError();\n amgx::thrust::sort_by_key(R_row_indices.begin(), R_row_indices.end(), m_R_column_indices.begin());\n cudaCheckError();\n }\n\n amgx::thrust::lower_bound(R_row_indices.begin(),\n R_row_indices.end(),\n amgx::thrust::counting_iterator(0),\n amgx::thrust::counting_iterator(m_R_row_offsets.size()),\n m_R_row_offsets.begin());\n cudaCheckError();\n}\n\n\n// two methods below could be merged\n// Method to compute R on HOST using csr format\ntemplate \nvoid Aggregation_AMG_Level >::computeRestrictionOperator_1x1()\n{\n this->m_R_row_offsets.resize(this->m_num_all_aggregates + 1);\n this->m_R_column_indices.resize(this->A->get_num_rows());\n this->fillRowOffsetsAndColIndices(this->A->get_num_rows());\n}\n\n// Method to compute R on HOST using block dia-csr format\ntemplate \nvoid Aggregation_AMG_Level >::computeRestrictionOperator_4x4()\n{\n this->m_R_row_offsets.resize(this->m_num_all_aggregates + 1);\n this->m_R_column_indices.resize(this->A->get_num_rows());\n this->fillRowOffsetsAndColIndices(this->A->get_num_rows());\n}\n\n// Method to create R_row_offsest and R_column_indices array on HOST using csr or block dia-csr format\ntemplate \nvoid Aggregation_AMG_Level_Base::fillRowOffsetsAndColIndices(const int R_num_cols)\n{\n for (int i = 0; i < m_num_all_aggregates + 1; i++)\n {\n m_R_row_offsets[i] = 0;\n }\n\n // Count number of neighbors for each row\n for (int i = 0; i < R_num_cols; i++)\n {\n int I = m_aggregates[i];\n m_R_row_offsets[I]++;\n }\n\n m_R_row_offsets[m_num_all_aggregates] = R_num_cols;\n\n for (int i = m_num_all_aggregates - 1; i >= 0; i--)\n {\n m_R_row_offsets[i] = m_R_row_offsets[i + 1] - m_R_row_offsets[i];\n }\n\n /* Set column indices. */\n for (int i = 0; i < R_num_cols; i++)\n {\n int I = m_aggregates[i];\n int Ip = m_R_row_offsets[I]++;\n m_R_column_indices[Ip] = i;\n }\n\n /* Reset r[i] to start of row memory. */\n for (int i = m_num_all_aggregates - 1; i > 0; i--)\n {\n m_R_row_offsets[i] = m_R_row_offsets[i - 1];\n }\n\n m_R_row_offsets[0] = 0;\n}\n\n// Method to compute R on DEVICE using block dia-csr format\ntemplate \nvoid Aggregation_AMG_Level >::computeRestrictionOperator_4x4()\n{\n this->computeRestrictionOperator_common();\n}\n\ntemplate \nvoid Aggregation_AMG_Level >::computeRestrictionOperator_1x1()\n{\n this->computeRestrictionOperator_common();\n}\n\n// Method to restrict Residual on host using csr_matrix format\ntemplate \nvoid Aggregation_AMG_Level >::restrictResidual_1x1(const VVector &r, VVector &rr)\n{\n ValueTypeB temp;\n\n for (int i = 0; i < this->m_num_aggregates; i++)\n {\n temp = types::util::get_zero();\n\n for (int j = this->m_R_row_offsets[i]; j < this->m_R_row_offsets[i + 1]; j++)\n {\n int j_col = this->m_R_column_indices[j];\n temp = temp + r[j_col];\n }\n\n rr[i] = temp;\n }\n}\n\n// Method to restrict Residual on host using block_dia_csr_matrix format\ntemplate \nvoid Aggregation_AMG_Level >::restrictResidual_4x4(const VVector &r, VVector &rr)\n{\n IndexType bsize = this->A->get_block_dimy();\n ValueTypeB *temp = new ValueTypeB[bsize];\n\n for (int i = 0; i < this->m_num_aggregates; i++)\n {\n // Initialize temp to 0\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = types::util::get_zero();\n }\n\n // Add contributions from each fine point\n for (int j = this->m_R_row_offsets[i]; j < this->m_R_row_offsets[i + 1]; j++)\n {\n int j_col = this->m_R_column_indices[j];\n\n for (int k = 0; k < bsize; k++)\n {\n temp[k] = temp[k] + r[j_col * bsize + k];\n }\n }\n\n // Store result\n for (int k = 0; k < bsize; k++)\n {\n rr[i * bsize + k] = temp[k];\n }\n }\n}\n\n// Method to restrict Residual on device using csr_matrix format\ntemplate \nvoid Aggregation_AMG_Level >::restrictResidual_1x1(const VVector &r, VVector &rr)\n{\n int block_size = 128;\n int max_threads;\n\n if (!this->isConsolidationLevel())\n {\n max_threads = this->m_num_aggregates;\n }\n else\n {\n max_threads = this->m_num_all_aggregates;\n }\n\n int num_blocks = max_threads / block_size + 1;\n const IndexType *R_row_offsets_ptr = this->m_R_row_offsets.raw();\n const IndexType *R_column_indices_ptr = this->m_R_column_indices.raw();\n const ValueTypeB *r_ptr = r.raw();\n ValueTypeB *rr_ptr = rr.raw();\n restrictResidualKernel <<< num_blocks, block_size>>>(R_row_offsets_ptr, R_column_indices_ptr, r_ptr, rr_ptr, max_threads);\n cudaCheckError();\n}\n\n// Method to restrict Residual on device using block_dia_csr_matrix format\ntemplate \nvoid Aggregation_AMG_Level >::restrictResidual_4x4(const VVector &r, VVector &rr)\n{\n int block_size = 128;\n int max_threads;\n\n if (!this->isConsolidationLevel())\n {\n max_threads = this->m_num_aggregates;\n }\n else\n {\n max_threads = this->m_num_all_aggregates;\n };\n\n const int num_blocks = max_threads / block_size + 1;\n const IndexType *R_row_offsets_ptr = this->m_R_row_offsets.raw();\n\n const IndexType *R_column_indices_ptr = this->m_R_column_indices.raw();\n\n const ValueTypeB *r_ptr = r.raw();\n\n ValueTypeB *rr_ptr = rr.raw();\n\n cudaCheckError();\n\n switch ( this->getA().get_block_dimy() )\n {\n case 2:\n restrictResidualBlockDiaCsrKernel <<< num_blocks, block_size>>>(R_row_offsets_ptr, R_column_indices_ptr, r_ptr, rr_ptr, max_threads);\n cudaCheckError();\n break;\n\n case 3:\n restrictResidualBlockDiaCsrKernel <<< num_blocks, block_size>>>(R_row_offsets_ptr, R_column_indices_ptr, r_ptr, rr_ptr, max_threads);\n cudaCheckError();\n break;\n\n case 4:\n restrictResidualBlockDiaCsrKernel <<< num_blocks, block_size>>>(R_row_offsets_ptr, R_column_indices_ptr, r_ptr, rr_ptr, max_threads);\n cudaCheckError();\n break;\n\n case 5:\n restrictResidualBlockDiaCsrKernel <<< num_blocks, block_size>>>(R_row_offsets_ptr, R_column_indices_ptr, r_ptr, rr_ptr, max_threads);\n cudaCheckError();\n break;\n\n case 8:\n restrictResidualBlockDiaCsrKernel <<< num_blocks, block_size>>>(R_row_offsets_ptr, R_column_indices_ptr, r_ptr, rr_ptr, max_threads);\n cudaCheckError();\n break;\n\n case 10:\n restrictResidualBlockDiaCsrKernel <<< num_blocks, block_size>>>(R_row_offsets_ptr, R_column_indices_ptr, r_ptr, rr_ptr, max_threads);\n cudaCheckError();\n break;\n\n default:\n FatalError( \"Unsupported block size in restrictResidual_4x4!!!\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE );\n }\n\n cudaCheckError();\n}\n\n__inline__ float getAlpha(float &nom, float &denom)\n{\n float alpha;\n\n if (nom * denom <= 0. || std::abs(nom) < std::abs(denom))\n {\n alpha = 1.;\n }\n else if (std::abs(nom) > 2.*std::abs(denom))\n {\n alpha = 2.;\n }\n else\n {\n alpha = nom / denom;\n }\n\n return alpha;\n}\n\n__inline__ double getAlpha(double &nom, double &denom)\n{\n double alpha;\n\n if (nom * denom <= 0. || std::abs(nom) < std::abs(denom))\n {\n alpha = 1.;\n }\n else if (std::abs(nom) > 2.*std::abs(denom))\n {\n alpha = 2.;\n }\n else\n {\n alpha = nom / denom;\n }\n\n return alpha;\n}\n\n__inline__ hipComplex getAlpha(hipComplex &nom, hipComplex &denom)\n{\n hipComplex alpha;\n\n if (types::util::abs(nom) < types::util::abs(denom))\n {\n alpha = make_hipComplex(1.f, 0.f);\n }\n else if (types::util::abs(nom) > 2.*types::util::abs(denom))\n {\n alpha = make_hipComplex(2.f, 0.f);\n }\n else\n {\n alpha = nom / denom;\n }\n\n return alpha;\n}\n\n__inline__ hipDoubleComplex getAlpha(hipDoubleComplex &nom, hipDoubleComplex &denom)\n{\n hipDoubleComplex alpha;\n\n if (types::util::abs(nom) < types::util::abs(denom))\n {\n alpha = make_hipDoubleComplex(1., 0.);\n }\n else if (types::util::abs(nom) > 2.*types::util::abs(denom))\n {\n alpha = make_hipDoubleComplex(2., 0.);\n }\n else\n {\n alpha = nom / denom;\n }\n\n return alpha;\n}\n\ntemplate< class T_Config>\ntypename T_Config::VecPrec Aggregation_AMG_Level_Base::computeAlpha(const Vector &e, const Vector &bc, const Vector &tmp)\n{\n typename T_Config::VecPrec alpha = types::util::get_one();\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n int size = Ac.get_num_rows();\n VVector v(2, types::util::get_zero());\n v[0] = amgx::thrust::inner_product(e.begin(), e.begin() + size, bc.begin(), types::util::get_zero());\n v[1] = amgx::thrust::inner_product(e.begin(), e.begin() + size, tmp.begin(), types::util::get_zero());\n cudaCheckError();\n return getAlpha(v[0], v[1]);\n}\n\n// Method to prolongate the error on HOST using csr format\ntemplate \nvoid Aggregation_AMG_Level >::prolongateAndApplyCorrection_1x1(Vector > &e, Vector > &bc, Vector > &x, Vector > &tmp)\n{\n Matrix > &A = this->getA();\n Matrix > &C = this->next_h->getA();\n\n if ( this->m_error_scaling >= 2 )\n {\n FatalError(\"error_scaling=2,3 is not implemented on host\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n ValueTypeB alpha = types::util::get_one();\n\n if (this->m_error_scaling)\n {\n multiply(this->next_h->getA(), e, tmp);\n alpha = this->computeAlpha (e, bc, tmp);\n }\n\n // Apply correction on all (interior and exterior) equations.\n for (int i = 0; i < A.get_num_cols(); i++)\n {\n int I = this->m_aggregates[i];\n x[i] = x[i] + alpha * e[I];\n }\n}\n\n// Method to prolongate the error on HOST using block_dia_csr format\ntemplate \nvoid Aggregation_AMG_Level >::prolongateAndApplyCorrection_4x4(Vector > &e, Vector > &bc, Vector > &x, Vector > &tmp)\n{\n if (this->A->get_block_dimy() != this->A->get_block_dimx())\n {\n FatalError(\"Aggregation_AMG_Level not implemented for non square blocks, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if ( this->m_error_scaling >= 2 )\n {\n FatalError(\"error_scaling=2,3 is not implemented on host\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n Matrix &C = this->next_h->getA();\n ValueTypeB alpha = types::util::get_one();\n\n if (this->m_error_scaling)\n {\n multiply(this->next_h->getA(), e, tmp);\n alpha = this->computeAlpha (e, bc, tmp);\n }\n\n // Apply correction on all equations.\n for (int i = 0; i < this->A->get_num_rows(); i++)\n {\n int I = this->m_aggregates[i];\n\n for (int k = 0; k < this->A->get_block_dimy(); k++)\n {\n x[i * this->A->get_block_dimy() + k] = x[i * this->A->get_block_dimy() + k] + alpha * e[I * this->A->get_block_dimy() + k];\n }\n }\n}\n\n// Prolongate the error on DEVICE using csr format\ntemplate \nvoid Aggregation_AMG_Level >::prolongateAndApplyCorrection_1x1(Vector > &e, Vector > &bc, Vector > &x, Vector > &tmp)\n{\n ValueTypeB alpha = types::util::get_one();\n const int block_size = 128;\n const int num_blocks = this->A->get_num_rows() / block_size + 1;\n const IndexType *aggregates_ptr = this->m_aggregates.raw();\n ValueTypeB *x_ptr = x.raw();\n const ValueTypeB *e_ptr = e.raw();\n\n if (this->m_error_scaling)\n {\n FatalError(\"error_scaling=1 is deprecated\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n prolongateAndApplyCorrectionKernel <<< num_blocks, block_size>>>(alpha, (int)this->A->get_num_rows(), x_ptr, e_ptr, aggregates_ptr, this->m_num_aggregates);\n cudaCheckError();\n}\n\n// Prolongate the error on DEVICE using block dia-csr format\ntemplate \nvoid Aggregation_AMG_Level >::prolongateAndApplyCorrection_4x4(Vector > &ec,\n Vector > &bf,\n Vector > &xf,\n Vector > &rf)\n{\n if ( this->m_error_scaling >= 2 )\n {\n if ( this->scale_counter > 0 )\n {\n const IndexType *aggregates_ptr = this->m_aggregates.raw();\n ValueTypeB *x_ptr = xf.raw();\n const ValueTypeB *e_ptr = ec.raw();\n const int block_size = 128;\n const int num_blocks = (this->A->get_num_rows() - 1) / block_size + 1;\n prolongateAndApplyCorrectionBlockDiaCsrKernel <<< num_blocks, block_size>>>(this->scale, (int)this->getA().get_num_rows(), x_ptr, e_ptr, aggregates_ptr, this->m_num_aggregates, this->getA().get_block_dimy());\n cudaCheckError();\n this->scale_counter--;\n return;\n }\n\n bool vanek_scaling = this->m_error_scaling > 3;\n IndexType numRowsCoarse = this->next_d->getA().get_num_rows();\n IndexType numRowsFine = this->A->get_num_rows();\n IndexType blockdim = this->A->get_block_dimx();\n\n if ( blockdim != this->A->get_block_dimy() )\n {\n FatalError(\"Unsupported dimension for aggregation amg level\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n VVector ef( rf.size() );\n VVector Aef( rf.size() );\n ef.set_block_dimy( blockdim );\n Aef.set_block_dimy( blockdim );\n // prolongate e\n const int threads_per_block = 256;\n const int num_block_values = std::min( AMGX_GRID_MAX_SIZE, (numRowsFine * blockdim - 1) / threads_per_block + 1);\n const hipStream_t stream = nullptr;\n prolongateVector <<< num_block_values, threads_per_block, 0, stream>>>( this->m_aggregates.raw(), ec.raw(), ef.raw(), numRowsFine, numRowsCoarse, blockdim );\n cudaCheckError();\n ef.dirtybit = 1;\n hipStreamSynchronize(stream);\n cudaCheckError();\n int preSmooth;\n\n if ( vanek_scaling )\n {\n preSmooth = this->amg->getNumPostsweeps();\n }\n else\n {\n preSmooth = this->scaling_smoother_steps;\n }\n\n //smooth error\n this->smoother->setTolerance( 0.0 );\n this->smoother->set_max_iters( preSmooth );\n\n if ( vanek_scaling )\n {\n thrust_wrapper::fill( Aef.begin(), Aef.end(), types::util::get_zero() );\n cudaCheckError();\n this->smoother->solve( Aef, ef, false ); //smooth correction with rhs 0\n this->smoother->solve( bf, xf, false ); // smooth x with rhs residual\n //recompute residual\n int offset, size;\n this->getA().getOffsetAndSizeForView(OWNED, &offset, &size);\n axmb( this->getA(), xf, bf, rf, offset, size );\n }\n else\n {\n this->smoother->solve( rf, ef, false ); //smooth correction with rhs residual\n }\n\n // multiply for lambda computation\n multiply(this->getA(), ef, Aef, OWNED);\n ValueTypeB nominator, denominator;\n int offset = 0, size = 0;\n this->A->getOffsetAndSizeForView(OWNED, &offset, &size);\n\n if ( this->m_error_scaling == 2 || this->m_error_scaling == 4 )\n {\n // compute lambda=/\n nominator = amgx::thrust::inner_product( rf.begin(), rf.end(), Aef.begin(), types::util::get_zero() );\n denominator = amgx::thrust::inner_product( Aef.begin(), Aef.end(), Aef.begin(), types::util::get_zero() );\n cudaCheckError();\n }\n\n if ( this->m_error_scaling == 3 || this->m_error_scaling == 5)\n {\n // compute lambda=/\n nominator = amgx::thrust::inner_product( rf.begin(), rf.begin() + size * blockdim, ef.begin(), types::util::get_zero() );\n denominator = amgx::thrust::inner_product( ef.begin(), ef.begin() + size * blockdim, Aef.begin(), types::util::get_zero() );\n\n if (!this->A->is_matrix_singleGPU())\n {\n this->A->getManager()->global_reduce_sum(&nominator);\n this->A->getManager()->global_reduce_sum(&denominator);\n }\n\n cudaCheckError();\n }\n\n if (types::util::abs(denominator) == 0.0)\n {\n nominator = denominator = types::util::get_one();\n }\n\n // apply correction x <- x + lambda*e\n const int num_block_fine = std::min( AMGX_GRID_MAX_SIZE, (numRowsFine * blockdim - 1) / threads_per_block + 1 );\n ValueTypeB alpha = nominator / denominator;\n\n if ( types::util::abs(alpha) < .3 )\n {\n alpha = (alpha / types::util::abs(alpha)) * .3; // it was this before: alpha = .3, which is not 100% equal\n }\n\n if ( types::util::abs(alpha) > 10 )\n {\n alpha = (alpha / types::util::abs(alpha)) * 10.; // it was this before: alpha = 10., which is not 100% equal\n }\n\n applyCorrection <<< num_block_fine, threads_per_block, 0, stream>>>( alpha, ef.raw(), xf.raw(), numRowsFine * blockdim );\n cudaCheckError();\n this->scale_counter = this->reuse_scale; //reuse this scale scale_counter times\n this->scale = alpha;\n return;\n }\n\n ValueTypeB alpha = types::util::get_one();\n const int block_size = 128;\n const int num_blocks = this->A->get_num_rows() / block_size + 1;\n const IndexType *aggregates_ptr = this->m_aggregates.raw();\n ValueTypeB *x_ptr = xf.raw();\n const ValueTypeB *e_ptr = ec.raw();\n\n if (this->m_error_scaling == 1)\n {\n FatalError(\"error_scaling=1 is deprecated\", AMGX_ERR_NOT_IMPLEMENTED );\n }\n\n prolongateAndApplyCorrectionBlockDiaCsrKernel <<< num_blocks, block_size>>>(alpha, (int)this->A->get_num_rows(), x_ptr, e_ptr, aggregates_ptr, this->m_num_aggregates, this->A->get_block_dimy());\n cudaCheckError();\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::prolongateAndApplyCorrection(VVector &e, VVector &bf, VVector &x, VVector &tmp)\n{\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n\n //this is dirty, but error scaling 2 and 3 do not have a specialized version. Instead, the general version sits in the 4x4 function\n if ( this->m_error_scaling >= 2 )\n {\n prolongateAndApplyCorrection_4x4(e, bf, x, tmp);\n }\n else if (this->A->get_block_size() == 1)\n {\n prolongateAndApplyCorrection_1x1(e, bf, x, tmp);\n }\n else if (this->A->get_block_dimx() == this->A->get_block_dimy() )\n {\n prolongateAndApplyCorrection_4x4(e, bf, x, tmp);\n }\n else\n {\n FatalError(\"Unsupported dimension for aggregation amg level\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n x.dirtybit = 1;\n\n if (!this->A->is_matrix_singleGPU() && x.delayed_send == 0)\n {\n if (x.in_transfer & RECEIVING) { this->A->manager->exchange_halo_wait(x, x.tag); }\n\n this->A->manager->exchange_halo_async(x, x.tag);\n }\n}\n\n\ntemplate \nvoid Aggregation_AMG_Level_Base::restrictResidual(VVector &r, VVector &rr)\n{\n if (this->A->get_block_size() == 1)\n {\n restrictResidual_1x1(r, rr);\n }\n else if (this->A->get_block_dimx() == this->A->get_block_dimy() )\n {\n restrictResidual_4x4(r, rr);\n }\n else\n {\n FatalError(\"Unsupported dimension for aggregation amg level\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n //TODO: check level transfer between host and device for multiGPU\n if (!this->A->is_matrix_singleGPU())\n {\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n rr.dirtybit = 1;\n\n if (!Ac.is_matrix_singleGPU() && !this->isConsolidationLevel() && rr.delayed_send == 0)\n {\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA(); //TODO problem in memoryspace transfer is here\n\n if (rr.in_transfer & RECEIVING) { Ac.manager->exchange_halo_wait(rr, rr.tag); }\n\n Ac.manager->exchange_halo_async(rr, rr.tag);\n }\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::computeRestrictionOperator()\n{\n if (this->A->get_block_size() == 1)\n {\n computeRestrictionOperator_1x1();\n }\n else if (this->A->get_block_dimx() == 4 && this->A->get_block_dimy() == 4)\n {\n computeRestrictionOperator_4x4();\n }\n else\n {\n this->computeRestrictionOperator_common();\n }\n}\n\ntemplate \n__global__ void coarse_to_global(IndexType *aggregates, IndexType *aggregates_global, IndexType *renumbering, IndexType num_elements, int64_t offset)\n{\n int element = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (element < num_elements)\n {\n renumbering[aggregates[element]] = aggregates_global[element] + offset; //this won't be a problem, because we are overwriting the same thing\n element += blockDim.x * gridDim.x;\n }\n}\n\ntemplate \n__global__ void export_matrix_elements(IndexType *row_offsets, IndexType *col_indices, T *values, IndexType *maps, IndexType *renumbering, IndexType *new_row_offsets, IndexType *new_col_indices, T *new_values, IndexType bsize, IndexType size)\n{\n int idx = blockIdx.x * blockDim.x / 32 + threadIdx.x / 32;\n int coopIdx = threadIdx.x % 32;\n\n while (idx < size)\n {\n int row = maps[idx];\n INDEX_TYPE src_base = row_offsets[row];\n INDEX_TYPE dst_base = new_row_offsets[idx];\n\n for (int m = coopIdx; m < row_offsets[row + 1]*bsize - src_base * bsize; m += 32)\n {\n new_values[dst_base * bsize + m] = values[src_base * bsize + m];\n }\n\n for (int m = coopIdx; m < row_offsets[row + 1] - src_base; m += 32)\n {\n new_col_indices[dst_base + m] = renumbering[col_indices[src_base + m]];\n }\n\n idx += gridDim.x * blockDim.x / 32;\n }\n}\n\ntemplate \n__global__ void export_matrix_diagonal(T *values, INDEX_TYPE bsize, INDEX_TYPE *maps, T *output, INDEX_TYPE size)\n{\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (idx < size)\n {\n int row = maps[idx];\n INDEX_TYPE src_base = row;\n INDEX_TYPE dst_base = idx;\n\n for (int m = 0; m < bsize; m++)\n {\n output[dst_base * bsize + m] = values[src_base * bsize + m];\n }\n\n idx += gridDim.x * blockDim.x;\n }\n}\n\n__global__ void remove_boundary(INDEX_TYPE *flags, INDEX_TYPE *maps, INDEX_TYPE size)\n{\n int element = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (element < size)\n {\n flags[maps[element]] = 0; //this won't be a problem, because we are overwriting the same thing\n element += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void calc_inverse_renumbering(INDEX_TYPE *renum, INDEX_TYPE *irenum, INDEX_TYPE *renum_gbl, INDEX_TYPE base_index, INDEX_TYPE max_element)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < max_element)\n {\n irenum[renum[idx]] = renum_gbl[idx] - base_index;\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void create_halo_mapping(INDEX_TYPE *mapping, INDEX_TYPE *node_list, INDEX_TYPE base_index, INDEX_TYPE map_offset, INDEX_TYPE size)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < size)\n {\n int idx = node_list[row] - base_index;\n mapping[idx] = map_offset + row;\n row += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void map_col_indices_and_count_rowlen(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, INDEX_TYPE *row_length, INDEX_TYPE *renumbering, INDEX_TYPE *mapping, INDEX_TYPE *map_offsets, int64_t *index_ranges, INDEX_TYPE part_id, INDEX_TYPE my_id, INDEX_TYPE base_index, INDEX_TYPE my_range, INDEX_TYPE num_neighbors, INDEX_TYPE num_rows)\n{\n extern __shared__ volatile int reduction[];\n int row = blockIdx.x * blockDim.x / 4 + threadIdx.x / 4;\n int coopIdx = threadIdx.x % 4;\n\n while (row < num_rows)\n {\n int valid = 0;\n\n for (int idx = row_offsets[row] + coopIdx; idx < row_offsets[row + 1]; idx += 4) //this may look horrible, but I expect low branch divergence, because col indices in a row usually belong to the same partition (or at most one more)\n {\n int colIdx = col_indices[idx];\n int part = -2;\n\n if (colIdx >= index_ranges[2 * part_id] && colIdx < index_ranges[2 * part_id + 1]) //the col index probably belongs to the partition I am working on\n {\n part = part_id;\n }\n else if (colIdx >= base_index && colIdx < base_index + my_range) //or points back to the owned partition\n {\n part = -1;\n }\n else //or else it points to a third partition\n {\n for (int i = 0; i < num_neighbors; i++)\n {\n if (colIdx >= index_ranges[2 * i] && colIdx < index_ranges[2 * i + 1])\n {\n part = i;\n }\n }\n }\n\n if (part == -2)\n {\n col_indices[idx] = -1;\n#ifdef DEBUG\n printf(\"Column index encountered that does not belong to any of my neighbors!! %d\\n\", colIdx);\n#endif\n }\n else\n {\n if (part == -1)\n {\n col_indices[idx] = renumbering[colIdx - base_index];\n valid++;\n }\n else\n {\n int new_col_idx = mapping[map_offsets[part] + colIdx - index_ranges[2 * part]];\n\n if (new_col_idx >= 0)\n {\n valid++;\n col_indices[idx] = new_col_idx;\n }\n else\n {\n col_indices[idx] = -1;\n }\n }\n }\n }\n\n reduction[threadIdx.x] = valid;\n\n for (int s = 2; s > 0; s >>= 1)\n {\n if (coopIdx < s)\n {\n reduction[threadIdx.x] += reduction[threadIdx.x + s];\n }\n\n __syncthreads();\n }\n\n if (coopIdx == 0)\n {\n row_length[row] = reduction[threadIdx.x];\n }\n\n row += gridDim.x * blockDim.x / 4;\n }\n}\n\n__global__ void map_col_indices(INDEX_TYPE *row_offsets, INDEX_TYPE *col_indices, int64_t *halo_ranges, INDEX_TYPE *halo_renumbering, INDEX_TYPE *halo_rows, INDEX_TYPE *global_renumbering, INDEX_TYPE num_neighbors, INDEX_TYPE num_rows, INDEX_TYPE num_rows_processed)\n{\n int row = blockIdx.x * blockDim.x / 4 + threadIdx.x / 4;\n int coopIdx = threadIdx.x % 4;\n\n while (row < num_rows_processed)\n {\n for (int idx = row_offsets[row] + coopIdx; idx < row_offsets[row + 1]; idx += 4)\n {\n int colIdx = col_indices[idx];\n int part = 0;\n\n if (colIdx < num_rows)\n {\n part = -1;\n }\n else\n {\n colIdx = global_renumbering[colIdx];\n\n for (int i = 0; i < num_neighbors; i++)\n {\n if (colIdx >= halo_ranges[2 * i] && colIdx < halo_ranges[2 * i + 1])\n {\n part = i;\n break;\n }\n }\n }\n\n if (part == -1)\n {\n col_indices[idx] = colIdx;\n }\n else\n {\n col_indices[idx] = halo_renumbering[halo_rows[part] + colIdx - halo_ranges[2 * part]];\n }\n }\n\n row += gridDim.x * blockDim.x / 4;\n }\n}\n\ntemplate \n__global__ void reorder_whole_matrix(INDEX_TYPE *old_rows, INDEX_TYPE *old_cols, T *old_vals, INDEX_TYPE *rows, INDEX_TYPE *cols, T *vals, INDEX_TYPE bsize, INDEX_TYPE num_rows)\n{\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n\n while (row < num_rows)\n {\n INDEX_TYPE dst_row = row;\n INDEX_TYPE src_base = old_rows[row];\n INDEX_TYPE dst = rows[dst_row];\n\n for (int i = 0; i < old_rows[row + 1] - src_base; i++)\n {\n INDEX_TYPE colIdx = old_cols[src_base + i];\n\n if (colIdx >= 0)\n {\n cols[dst] = colIdx;\n\n for (int j = 0; j < bsize; j++) { vals[dst * bsize + j] = old_vals[(src_base + i) * bsize + j]; }\n\n dst++;\n }\n }\n\n row += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void calc_gbl_renumbering(INDEX_TYPE *inv_renum, INDEX_TYPE *gbl_renum, INDEX_TYPE size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n gbl_renum[inv_renum[idx]] = idx;\n idx += blockDim.x * gridDim.x;\n }\n}\n\ntemplate \n__global__ void write_diagonals(ValueType *values, INDEX_TYPE *diag, INDEX_TYPE *map, ValueType *output, INDEX_TYPE bsize, INDEX_TYPE size)\n{\n int nzPerBlock = blockDim.x / bsize;\n int row = blockIdx.x * nzPerBlock + threadIdx.x / bsize;\n int vecIdx = threadIdx.x % bsize;\n\n if (threadIdx.x >= (blockDim.x / bsize)*bsize) { return; }\n\n while (row < size)\n {\n output[row * bsize + vecIdx] = values[diag[map[row]] * bsize + vecIdx];\n row += gridDim.x * nzPerBlock;\n }\n}\n\ntemplate \n__global__ void write_diagonals_back(ValueType *values, INDEX_TYPE *diag, ValueType *source, INDEX_TYPE bsize, INDEX_TYPE size)\n{\n int nzPerBlock = blockDim.x / bsize;\n int row = blockIdx.x * nzPerBlock + threadIdx.x / bsize;\n int vecIdx = threadIdx.x % bsize;\n\n if (threadIdx.x >= (blockDim.x / bsize)*bsize) { return; }\n\n while (row < size)\n {\n values[diag[row]*bsize + vecIdx] = source[row * bsize + vecIdx];\n row += gridDim.x * nzPerBlock;\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::prepareNextLevelMatrix_full(const Matrix &A, Matrix &Ac)\n{\n if (A.is_matrix_singleGPU()) { return; }\n\n int num_neighbors = A.manager->neighbors.size();\n\n if (TConfig::memSpace == AMGX_host)\n {\n FatalError(\"Aggregation AMG Not implemented for host\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n int c_size = Ac.get_num_rows();\n int f_size = A.get_num_rows();\n int diag = Ac.hasProps(DIAG);\n\n if (A.manager->B2L_rings[0].size() > 2) { FatalError(\"Aggregation_AMG_Level prepareNextLevelMatrix not implemented >1 halo rings\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE); }\n\n //get coarse -> fine global renumbering\n IVector renumbering(c_size);\n int num_blocks = std::min(4096, (c_size + 127) / 128);\n coarse_to_global <<< num_blocks, 128>>>(this->m_aggregates.raw(), this->m_aggregates_fine_idx.raw(), renumbering.raw(), f_size, 0);\n cudaCheckError();\n //\n // Step 0 - form halo matrices that are exported to neighbors\n //\n std::vector > halo_rows(num_neighbors);\n std::vector > halo_btl(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++ )\n {\n int num_unique = Ac.manager->B2L_rings[i][1];\n //prepare export halo matrices\n halo_btl[i].resize(1, 1);\n halo_btl[i].set_global_id(Ac.manager->global_id());\n halo_btl[i].B2L_maps[0].resize(num_unique);\n halo_btl[i].B2L_rings[0].resize(2);\n halo_btl[i].B2L_rings[0][0] = 0;\n halo_btl[i].B2L_rings[0][1] = num_unique;\n halo_btl[i].set_index_range(A.manager->index_range());\n halo_btl[i].set_base_index(A.manager->base_index());\n //global indices of rows of the halo matrix\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator( renumbering.begin(), Ac.manager->B2L_maps[i].begin()),\n amgx::thrust::make_permutation_iterator( renumbering.begin(), Ac.manager->B2L_maps[i].begin() + num_unique),\n halo_btl[i].B2L_maps[0].begin());\n cudaCheckError();\n halo_rows[i].addProps(CSR);\n\n if (diag) { halo_rows[i].addProps(DIAG); }\n\n //calculate row length and row_offsets\n halo_rows[i].row_offsets.resize(num_unique + 1);\n thrust_wrapper::transform(amgx::thrust::make_permutation_iterator(Ac.row_offsets.begin() + 1, Ac.manager->B2L_maps[i].begin()),\n amgx::thrust::make_permutation_iterator(Ac.row_offsets.begin() + 1, Ac.manager->B2L_maps[i].end()),\n amgx::thrust::make_permutation_iterator(Ac.row_offsets.begin(), Ac.manager->B2L_maps[i].begin()),\n halo_rows[i].row_offsets.begin(),\n amgx::thrust::minus());\n cudaCheckError();\n thrust_wrapper::exclusive_scan(halo_rows[i].row_offsets.begin(), halo_rows[i].row_offsets.end(), halo_rows[i].row_offsets.begin());\n cudaCheckError();\n //resize halo matrix\n IndexType num_nz = halo_rows[i].row_offsets[num_unique];\n halo_rows[i].resize(num_unique, num_unique, num_nz, Ac.get_block_dimy(), Ac.get_block_dimx(), 1);\n //copy relevant rows and renumber their column indices\n num_blocks = std::min(4096, (num_unique + 127) / 128);\n export_matrix_elements <<< num_blocks, 128>>>(Ac.row_offsets.raw(), Ac.col_indices.raw(), Ac.values.raw(), Ac.manager->B2L_maps[i].raw(), renumbering.raw(), halo_rows[i].row_offsets.raw(), halo_rows[i].col_indices.raw(), halo_rows[i].values.raw(), A.get_block_size(), num_unique);\n cudaCheckError();\n\n if (diag)\n {\n export_matrix_diagonal <<< num_blocks, 128>>>(Ac.values.raw() + Ac.row_offsets[Ac.get_num_rows()]*Ac.get_block_size(), Ac.get_block_size(), Ac.manager->B2L_maps[i].raw(), halo_rows[i].values.raw() + halo_rows[i].row_offsets[halo_rows[i].get_num_rows()]*Ac.get_block_size(), num_unique);\n cudaCheckError();\n }\n }\n\n Ac.manager->getComms()->exchange_matrix_halo(halo_rows, halo_btl, Ac);\n //--------------------- renumbering/reordering matrix, integrating halo -----------------------------\n Ac.set_initialized(0);\n //number of owned rows\n c_size = Ac.manager->halo_offsets[0];\n f_size = A.manager->halo_offsets[0];\n num_blocks = std::min(4096, (c_size + 511) / 512);\n int rings = 1;\n //\n // Step 1 - calculate inverse renumbering (to global indices - base_index)\n //\n Ac.manager->inverse_renumbering.resize(c_size);\n thrust_wrapper::transform(renumbering.begin(),\n renumbering.begin() + c_size,\n amgx::thrust::constant_iterator(A.manager->base_index()),\n Ac.manager->inverse_renumbering.begin(),\n amgx::thrust::minus());\n cudaCheckError();\n //big renumbering table for going from global index to owned local index\n IVector global_to_coarse_local(Ac.manager->index_range());\n thrust_wrapper::fill(global_to_coarse_local.begin(), global_to_coarse_local.begin() + Ac.manager->index_range(), -1);\n cudaCheckError();\n calc_gbl_renumbering <<< num_blocks, 512>>>(Ac.manager->inverse_renumbering.raw(), global_to_coarse_local.raw(), c_size);\n cudaCheckError();\n Ac.manager->set_num_halo_rows(Ac.manager->halo_offsets[Ac.manager->halo_offsets.size() - 1] - c_size);\n cudaCheckError();\n //\n // Step 2 - create big mapping table of all halo indices we received (this may use a little too much memory sum(fine nodes per neighbor)\n //\n amgx::thrust::host_vector neighbor_rows(num_neighbors + 1);\n int max_num_rows = 0;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n neighbor_rows[i] = halo_rows[i].manager->index_range();\n max_num_rows = max_num_rows > halo_rows[i].get_num_rows() ? max_num_rows : halo_rows[i].get_num_rows();\n }\n\n thrust_wrapper::exclusive_scan(neighbor_rows.begin(), neighbor_rows.end(), neighbor_rows.begin());\n cudaCheckError();\n int total_rows_of_neighbors = neighbor_rows[num_neighbors];\n IVector halo_mapping(total_rows_of_neighbors);\n thrust_wrapper::fill(halo_mapping.begin(), halo_mapping.end(), -1);\n cudaCheckError();\n\n for (int ring = 0; ring < rings; ring++)\n {\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = halo_btl[i].B2L_rings[0][ring + 1] - halo_btl[i].B2L_rings[0][ring];\n int num_blocks = std::min(4096, (size + 127) / 128);\n create_halo_mapping <<< num_blocks, 128>>>(halo_mapping.raw() + neighbor_rows[i],\n halo_btl[i].B2L_maps[0].raw() + halo_btl[i].B2L_rings[0][ring],\n halo_btl[i].base_index(),\n Ac.manager->halo_offsets[ring * num_neighbors + i], size);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n //\n // Step 3 - renumber halo matrices and calculate row length (to eventually append to the big matrix)\n //\n INDEX_TYPE owned_nnz = Ac.row_offsets[c_size];\n IVector neighbor_rows_d(num_neighbors + 1);\n amgx::thrust::copy(neighbor_rows.begin(), neighbor_rows.end(), neighbor_rows_d.begin());\n cudaCheckError();\n //map column indices of my own matrix (the ones that point outward)\n map_col_indices <<< num_blocks, 512>>>(Ac.row_offsets.raw() + Ac.manager->num_interior_nodes(),\n Ac.col_indices.raw(),\n Ac.manager->halo_ranges.raw(),\n halo_mapping.raw(),\n neighbor_rows_d.raw(),\n renumbering.raw(),\n num_neighbors, c_size, c_size - Ac.manager->num_interior_nodes());\n cudaCheckError();\n IVector temp_row_len(max_num_rows);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n //map column indices of halo matrices\n int size = halo_rows[i].get_num_rows();\n int num_blocks = std::min(4096, (size + 127) / 128);\n map_col_indices_and_count_rowlen <<< num_blocks, 128, 128 * sizeof(INDEX_TYPE)>>>(\n halo_rows[i].row_offsets.raw(),\n halo_rows[i].col_indices.raw(),\n temp_row_len.raw(),\n global_to_coarse_local.raw(),\n halo_mapping.raw(),\n neighbor_rows_d.raw(),\n Ac.manager->halo_ranges.raw(),\n i,\n Ac.manager->global_id(),\n Ac.manager->base_index(),\n Ac.manager->index_range(),\n num_neighbors,\n size);\n cudaCheckError();\n\n for (int ring = 0; ring < rings; ring++)\n {\n amgx::thrust::copy(temp_row_len.begin() + halo_btl[i].B2L_rings[0][ring], temp_row_len.begin() + halo_btl[i].B2L_rings[0][ring + 1], Ac.row_offsets.begin() + Ac.manager->halo_offsets[ring * num_neighbors + i]);\n }\n }\n\n cudaCheckError();\n INDEX_TYPE old_nnz = Ac.row_offsets[Ac.row_offsets.size() - 1];\n thrust_wrapper::exclusive_scan(Ac.row_offsets.begin() + c_size, Ac.row_offsets.end(), Ac.row_offsets.begin() + c_size, owned_nnz);\n cudaCheckError();\n //\n // Step 4 - consolidate column indices and values\n //\n int new_nnz = Ac.row_offsets[Ac.row_offsets.size() - 1];\n\n Ac.col_indices.resize(new_nnz);\n Ac.values.resize((new_nnz + 1 + diag * (Ac.row_offsets.size() - 2)) * A.get_block_size());\n\n if (diag)\n {\n MVector diags(c_size * Ac.get_block_size());\n amgx::thrust::copy(Ac.values.begin() + old_nnz * Ac.get_block_size(),\n Ac.values.begin() + old_nnz * Ac.get_block_size() + c_size * Ac.get_block_size(),\n diags.begin());\n amgx::thrust::copy(diags.begin(), diags.begin() + c_size * Ac.get_block_size(),\n Ac.values.begin() + Ac.row_offsets[Ac.get_num_rows()]*Ac.get_block_size());\n cudaCheckError();\n }\n\n int cumulative_num_rows = c_size;\n\n for (int i = 0; i < num_neighbors; i++)\n {\n for (int ring = 0; ring < rings; ring++)\n {\n int num_rows = halo_btl[i].B2L_rings[0][ring + 1] - halo_btl[i].B2L_rings[0][ring];\n int num_blocks = std::min(4096, (num_rows + 127) / 128);\n reorder_whole_matrix <<< num_blocks, 128>>>(halo_rows[i].row_offsets.raw() + halo_btl[i].B2L_rings[0][ring], halo_rows[i].col_indices.raw(), halo_rows[i].values.raw(), Ac.row_offsets.raw() + Ac.manager->halo_offsets[ring * num_neighbors + i], Ac.col_indices.raw(), Ac.values.raw(), Ac.get_block_size(), num_rows);\n cudaCheckError();\n\n if (diag)\n {\n amgx::thrust::copy(halo_rows[i].values.begin() + (halo_rows[i].row_offsets[halo_rows[i].get_num_rows()] + halo_btl[i].B2L_rings[0][ring])*Ac.get_block_size(),\n halo_rows[i].values.begin() + (halo_rows[i].row_offsets[halo_rows[i].get_num_rows()] + halo_btl[i].B2L_rings[0][ring + 1])*Ac.get_block_size(),\n Ac.values.begin() + (Ac.row_offsets[Ac.get_num_rows()] + cumulative_num_rows)*Ac.get_block_size());\n cumulative_num_rows += num_rows;\n }\n }\n }\n\n cudaCheckError();\n Ac.set_num_cols(Ac.manager->halo_offsets[Ac.manager->halo_offsets.size() - 1]);\n Ac.set_num_rows(Ac.get_num_cols());\n Ac.set_num_nz(new_nnz);\n Ac.delProps(COO);\n Ac.set_initialized(1);\n Ac.computeDiagonal();\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::prepareNextLevelMatrix_diag(const Matrix &A, Matrix &Ac)\n{\n if (A.is_matrix_singleGPU()) { return; }\n\n int num_neighbors = A.manager->neighbors.size();\n\n if (TConfig::memSpace == AMGX_host)\n {\n FatalError(\"Aggregation AMG Not implemented for host\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n int c_size = Ac.manager->halo_offsets[0];\n int f_size = A.manager->halo_offsets[0];\n int diag = Ac.hasProps(DIAG);\n Ac.manager->inverse_renumbering.resize(c_size);\n //get coarse -> fine renumbering\n int num_blocks = std::min(4096, (c_size + 127) / 128);\n coarse_to_global <<< num_blocks, 128>>>(this->m_aggregates.raw(), this->m_aggregates_fine_idx.raw(), Ac.manager->inverse_renumbering.raw(), f_size, -1 * A.manager->base_index());\n cudaCheckError();\n Ac.manager->set_num_halo_rows(Ac.manager->halo_offsets[Ac.manager->halo_offsets.size() - 1] - c_size);\n\n if (!diag) { Ac.computeDiagonal(); }\n\n Ac.set_initialized(1);\n std::vector diagonals(num_neighbors);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = Ac.manager->B2L_rings[i][Ac.manager->B2L_rings.size() - 1];\n diagonals[i].resize(Ac.get_block_size()*size);\n int num_blocks = std::min(4096, (size + 127) / 128);\n write_diagonals <<< num_blocks, 128>>>(Ac.values.raw(), Ac.diag.raw(), Ac.manager->B2L_maps[i].raw(), diagonals[i].raw(), Ac.get_block_size(), size);\n }\n\n cudaCheckError();\n Ac.manager->getComms()->exchange_vectors(diagonals, Ac, this->tag * 100 + 10 + 2);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = Ac.manager->halo_offsets[i + 1] - Ac.manager->halo_offsets[i];\n\n if (Ac.hasProps(DIAG)) { amgx::thrust::copy(diagonals[i].begin(), diagonals[i].begin() + Ac.get_block_size()*size, Ac.values.begin() + Ac.get_block_size() * (Ac.diagOffset() + Ac.manager->halo_offsets[i])); }\n else\n {\n int num_blocks = std::min(4096, (size + 127) / 128);\n write_diagonals_back <<< num_blocks, 128>>>(Ac.values.raw(), Ac.diag.raw() + Ac.manager->halo_offsets[i], diagonals[i].raw(), Ac.get_block_size(), size);\n cudaCheckError();\n }\n }\n\n cudaCheckError();\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::prepareNextLevelMatrix_none(const Matrix &A, Matrix &Ac)\n{\n if (A.is_matrix_singleGPU()) { return; }\n\n int num_neighbors = A.manager->neighbors.size();\n\n if (TConfig::memSpace == AMGX_host)\n {\n FatalError(\"Aggregation AMG Not implemented for host\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n else\n {\n int c_size = Ac.manager->halo_offsets[0];\n int f_size = A.manager->halo_offsets[0];\n int diag = Ac.hasProps(DIAG);\n Ac.manager->inverse_renumbering.resize(c_size);\n //get coarse -> fine renumbering\n int num_blocks = std::min(4096, (c_size + 127) / 128);\n coarse_to_global <<< num_blocks, 128>>>(this->m_aggregates.raw(), this->m_aggregates_fine_idx.raw(), Ac.manager->inverse_renumbering.raw(), f_size, 0);\n cudaCheckError();\n Ac.manager->set_num_halo_rows(Ac.manager->halo_offsets[Ac.manager->halo_offsets.size() - 1] - c_size);\n Ac.set_initialized(1);\n\n if (!diag) { Ac.computeDiagonal(); }\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::prepareNextLevelMatrix(const Matrix &A, Matrix &Ac)\n{\n if (m_matrix_halo_exchange == 0)\n {\n this->prepareNextLevelMatrix_none(A, Ac);\n }\n else if (m_matrix_halo_exchange == 1)\n {\n this->prepareNextLevelMatrix_diag(A, Ac);\n }\n else if (m_matrix_halo_exchange == 2)\n {\n this->prepareNextLevelMatrix_full(A, Ac);\n }\n else\n {\n FatalError(\"Invalid Aggregation matrix_halo_exchange parameter\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n\n__global__ void set_halo_rowlen(INDEX_TYPE *work, INDEX_TYPE *output, INDEX_TYPE size, INDEX_TYPE diag)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n if (work[idx + 1] - work[idx] > 0)\n {\n output[idx] += work[idx + 1] - work[idx] - (1 - diag);\n }\n\n idx += blockDim.x * gridDim.x;\n }\n}\n\ntemplate \n__global__ void append_halo_nz(INDEX_TYPE *row_offsets, INDEX_TYPE *new_row_offsets, INDEX_TYPE *col_indices, INDEX_TYPE *new_col_indices, T *values, T *new_values, INDEX_TYPE size, INDEX_TYPE diag, INDEX_TYPE halo_offset, INDEX_TYPE block_size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n int add_diag = !diag;\n\n if (!diag && new_col_indices[new_row_offsets[idx]] != -1) { add_diag = 0; } //if diag or there is already soimething in the row, then don't add diagonal nonzero (inside diag)\n\n int append_offset = -1;\n\n for (int i = new_row_offsets[idx]; i < new_row_offsets[idx + 1]; i++)\n {\n if (new_col_indices[i] == -1) {append_offset = i; break;}\n }\n\n for (int i = row_offsets[idx]; i < row_offsets[idx + 1]; i++)\n {\n if (diag && i == row_offsets[idx]) //if outside diag and this is the first nonzero in a non-empty row, overwrite diagonal value\n {\n for (int j = 0; j < block_size; j++)\n {\n new_values[(new_row_offsets[size] + halo_offset + idx)*block_size + j] = values[(row_offsets[size] + halo_offset + idx) * block_size + j];\n }\n }\n\n int col_idx = col_indices[i];\n\n if (append_offset == -1 && (col_idx != halo_offset + idx)) {printf(\"ERROR: append offset is -1 but row has nonzeros in it old %d to %d new %d to %d\\n\", row_offsets[idx], row_offsets[idx + 1], new_row_offsets[idx], new_row_offsets[idx + 1]); append_offset = 0;}\n\n if (col_idx != halo_offset + idx || add_diag)\n {\n new_col_indices[append_offset] = col_idx;\n\n for (int j = 0; j < block_size; j++)\n {\n new_values[append_offset * block_size + j] = values[i * block_size + j];\n }\n\n append_offset++;\n }\n }\n\n idx += blockDim.x * gridDim.x;\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::createCoarseB2LMaps(std::vector &in_coarse_B2L_maps)\n{\n Matrix &A = this->getA();\n m_num_all_aggregates = m_num_aggregates;\n int num_neighbors = A.manager->neighbors.size();\n IndexType max_b2l = 0;\n\n for (int i = 0; i < num_neighbors; i++ ) { max_b2l = max_b2l > A.manager->B2L_rings[i][1] ? max_b2l : A.manager->B2L_rings[i][1]; }\n\n IVector B2L_aggregates(max_b2l);\n IVector indices(max_b2l);\n\n for (int i = 0; i < num_neighbors; i++ )\n {\n int size = A.manager->B2L_rings[i][1];\n thrust_wrapper::fill(B2L_aggregates.begin(), B2L_aggregates.begin() + size, 0);\n thrust_wrapper::sequence(indices.begin(), indices.begin() + size);\n //substitute coarse aggregate indices for fine boundary nodes\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(this->m_aggregates.begin(), A.manager->B2L_maps[i].begin()),\n amgx::thrust::make_permutation_iterator(this->m_aggregates.begin(), A.manager->B2L_maps[i].begin() + size),\n B2L_aggregates.begin());\n //find the unique ones\n amgx::thrust::sort_by_key(B2L_aggregates.begin(), B2L_aggregates.begin() + size, indices.begin());\n IndexType num_unique = amgx::thrust::unique_by_key(B2L_aggregates.begin(), B2L_aggregates.begin() + size, indices.begin()).first - B2L_aggregates.begin();\n in_coarse_B2L_maps[i].resize(num_unique);\n //sort it back so we have the original ordering\n amgx::thrust::sort_by_key(indices.begin(), indices.begin() + num_unique, B2L_aggregates.begin());\n amgx::thrust::copy(B2L_aggregates.begin(), B2L_aggregates.begin() + num_unique, in_coarse_B2L_maps[i].begin());\n }\n\n cudaCheckError();\n}\n\n\n__global__ void populate_coarse_boundary(INDEX_TYPE *flags, INDEX_TYPE *indices, INDEX_TYPE *maps, INDEX_TYPE *output, INDEX_TYPE size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n output[flags[maps[indices[idx]]]] = maps[indices[idx]];\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void flag_coarse_boundary(INDEX_TYPE *flags, INDEX_TYPE *indices, INDEX_TYPE *maps, INDEX_TYPE size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n flags[maps[indices[idx]]] = 1;\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void flag_halo_indices(INDEX_TYPE *flags, INDEX_TYPE *indices, INDEX_TYPE offset, INDEX_TYPE size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n flags[indices[idx] - offset] = 1;\n idx += blockDim.x * gridDim.x;\n }\n}\n\n__global__ void apply_halo_aggregate_indices(INDEX_TYPE *flags, INDEX_TYPE *indices, INDEX_TYPE *output, INDEX_TYPE offset, INDEX_TYPE aggregates_offset, INDEX_TYPE size)\n{\n int idx = blockDim.x * blockIdx.x + threadIdx.x;\n\n while (idx < size)\n {\n output[idx] = flags[indices[idx] - offset] + aggregates_offset;\n idx += blockDim.x * gridDim.x;\n }\n}\n\n// renumbering the aggregates/communicationg with neighbors\ntemplate \nvoid Aggregation_AMG_Level_Base::setNeighborAggregates()\n{\n Matrix &A = this->getA();\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n m_num_all_aggregates = m_num_aggregates;\n\n /* WARNING: the matrix reordering always happens inside createRenumbering routine. There are three ways to get to this routine\n 1. matrix_upload_all -> uploadMatrix -> initializeUploadReorderAll -> reorder_matrix -> createRenumbering\n 2. read_system_distributed -> renumberMatrixOneRing -> reorder_matrix_owned -> createRenumbering\n 3. solver_setup -> ... -> AMG_Level::setup -> createCoarseMatrices -> setNeighborAggregates -> createRenumbering\n If you are reading the renumbering from file you might need to add intercept code in if statement below,\n otherwise this routine will exit before calling createRenumbering routine (in case of single or disjoint partitions).\n */\n if (this->getA().is_matrix_singleGPU()) { return; }\n\n int num_neighbors = A.manager->neighbors.size();\n\n //\n // Step 0 - set up coarse matrix metadata\n //\n if (Ac.manager == NULL) { Ac.manager = new DistributedManager(); }\n\n Ac.manager->resize(A.manager->neighbors.size(), 1);\n Ac.manager->A = &Ac;\n int f_size = A.get_num_rows();\n Ac.manager->setComms(A.manager->getComms());\n Ac.manager->set_global_id(A.manager->global_id());\n Ac.manager->neighbors = A.manager->neighbors;\n Ac.manager->set_base_index(A.manager->base_index());\n Ac.manager->halo_ranges = A.manager->halo_ranges;\n Ac.manager->set_index_range(A.manager->index_range());\n //-------------------------------------- Section 1 - renumbering -----------------------------------------------------------\n //\n // Step 1 - calculate coarse level B2L maps - any aggregate that has a fine boundary node, becomes a coarse boundary node\n //\n m_num_all_aggregates = m_num_aggregates;\n int vec_size = m_num_aggregates + 1; //A.manager->num_boundary_nodes()+1;\n IVector B2L_aggregates(vec_size);\n\n for (int i = 0; i < A.manager->neighbors.size(); i++)\n {\n thrust_wrapper::fill(B2L_aggregates.begin(), B2L_aggregates.begin() + vec_size, 0);\n int size = A.manager->B2L_rings[i][1];\n int block_size = 128;\n int grid_size = std::min( 4096, ( size + block_size - 1 ) / block_size);\n flag_coarse_boundary <<< grid_size, block_size>>>(B2L_aggregates.raw(), A.manager->B2L_maps[i].raw(), this->m_aggregates.raw(), size);\n cudaCheckError();\n thrust_wrapper::exclusive_scan(B2L_aggregates.begin(), B2L_aggregates.begin() + vec_size, B2L_aggregates.begin());\n (Ac.manager->B2L_maps)[i].resize(B2L_aggregates[vec_size - 1]);\n populate_coarse_boundary <<< grid_size, block_size>>>(B2L_aggregates.raw(), A.manager->B2L_maps[i].raw(), this->m_aggregates.raw(), Ac.manager->B2L_maps[i].raw(), size);\n }\n\n cudaCheckError();\n\n for (int i = 0; i < num_neighbors; i++)\n {\n Ac.manager->B2L_rings[i].resize(2);\n Ac.manager->B2L_rings[i][0] = 0;\n Ac.manager->B2L_rings[i][1] = Ac.manager->B2L_maps[i].size();\n }\n\n DistributedArranger *prep = new DistributedArranger;\n prep->initialize_B2L_maps_offsets(Ac, 1);\n delete prep;\n Ac.set_num_rows(m_num_aggregates);\n IVector renumbering(m_num_aggregates + 1); /* +1 is actually not needed, it will be resized in createRenumbering */\n Ac.manager->createRenumbering(renumbering);\n //\n // Step 2 - renumber aggregates, so boundary nodes will have higher index than interior ones (based on the renumberiong we have been calculating)\n //\n /* WARNING: 1. Thrust scatter and gather routines seem more appropriate here, but they implicitly assume that the input\n and output have certain size correlation, which is not matched by vectors in our case. The only remaining option\n is to use make_permutation as is done below. Example of Thrust scatter and gather calls\n IVector ttt(f_size,-1);\n amgx::thrust::scatter(this->m_aggregates.begin(), this->m_aggregates.begin()+f_size, renumbering.begin(), ttt.begin());\n amgx::thrust::gather(renumbering.begin(), renumbering.end(), this->m_aggregates.begin(), ttt.begin());\n amgx::thrust::copy(ttt.begin(), ttt.end(), this->m_aggregates.begin());\n\n 2. The original thrust composite call is illegal because it uses the same array (m_aggregates) for input and output.\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(renumbering.begin(), this->m_aggregates.begin()),\n amgx::thrust::make_permutation_iterator(renumbering.begin(), this->m_aggregates.begin()+f_size),\n this->m_aggregates.begin());\n Although it somehow still works, it is much safer to use explicit temporary storage for the intermediate result.\n */\n /* WARNING: must save unreordered aggregates for later use before reordering them. */\n IVector unreordered_aggregates(this->m_aggregates);\n /* WARNING: change Thrust call to explicitly use temporary storage for the intermediate result. The earlier version is illegal, but somehow still works. */\n IVector ttt(f_size, -1);\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(renumbering.begin(), this->m_aggregates.begin()),\n amgx::thrust::make_permutation_iterator(renumbering.begin(), this->m_aggregates.begin() + f_size),\n ttt.begin());\n amgx::thrust::copy(ttt.begin(), ttt.end(), this->m_aggregates.begin());\n cudaCheckError();\n\n //we don't need renumbering anymore, it will be identity on the coarse level\n\n //-------------------------------------- Section 2 - communication -----------------------------------------------------------\n\n //\n // Step 3 - populate aggregates_fine_idx, which stores for every fine node the original global index of the aggregate (which is lowest global index of nodes aggregated together)\n //\n\n //\n // These are different when we do /don't do matrix halo exchanges - when we do we need global indices to match nodes,\n // and in this case Ac after computeA will not have the same ordering of halo nodes as after prepareNextLevel_full.\n // However when we do not do matrix halo exchange we are only interested in the ordering of halo nodes on the coarse level,\n // and we can get that by exchanging the (already renumbered) aggregates vector.\n //\n if (m_matrix_halo_exchange == 2)\n {\n //Find original global indices of nodes that have the minimum id in the aggregates.\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(A.manager->inverse_renumbering.begin(), this->m_aggregates_fine_idx.begin()),\n amgx::thrust::make_permutation_iterator(A.manager->inverse_renumbering.begin(), this->m_aggregates_fine_idx.begin() + f_size),\n this->m_aggregates_fine_idx.begin());\n thrust_wrapper::transform(this->m_aggregates_fine_idx.begin(),\n this->m_aggregates_fine_idx.begin() + f_size,\n amgx::thrust::constant_iterator(A.manager->base_index()),\n this->m_aggregates_fine_idx.begin(),\n amgx::thrust::plus());\n //communicate\n this->m_aggregates_fine_idx.set_block_dimx(1);\n this->m_aggregates_fine_idx.set_block_dimy(1);\n m_aggregates_fine_idx.dirtybit = 1;\n A.manager->exchange_halo(m_aggregates_fine_idx, this->tag * 100 + 1 * 10 + 0);\n }\n else\n {\n //communicate\n this->m_aggregates.set_block_dimx(1);\n this->m_aggregates.set_block_dimy(1);\n m_aggregates.dirtybit = 1;\n /* WARNING: you should exchange unreordered aggregates, and append them to your own reordered aggregates, to conform to asusmptions done by distributed_mamanger. */\n //A.manager->exchange_halo(m_aggregates, this->tag*100+1*10+0); //wrong\n A.manager->exchange_halo(unreordered_aggregates, this->tag * 100 + 1 * 10 + 0);\n amgx::thrust::copy(unreordered_aggregates.begin() + f_size, unreordered_aggregates.end(), this->m_aggregates.begin() + f_size);\n }\n\n cudaCheckError();\n //\n // Step 4 - consolidate neighbors' aggregates into own list to be able to perform Galerkin product with the n-ring halo\n //\n IVector &exchanged_aggregates = m_matrix_halo_exchange == 2 ? this->m_aggregates_fine_idx : this->m_aggregates;\n int min_index = amgx::thrust::reduce(exchanged_aggregates.begin() + A.manager->halo_offsets[0], exchanged_aggregates.begin() + A.manager->halo_offsets[num_neighbors], (int)0xFFFFFFF, amgx::thrust::minimum());\n int max_index = amgx::thrust::reduce(exchanged_aggregates.begin() + A.manager->halo_offsets[0], exchanged_aggregates.begin() + A.manager->halo_offsets[num_neighbors], (int)0, amgx::thrust::maximum());\n cudaCheckError();\n int s_size = max_index - min_index + 2;\n IVector scratch(s_size);\n\n for (int i = 0; i < num_neighbors; i++)\n {\n int size = A.manager->halo_offsets[i + 1] - A.manager->halo_offsets[i];\n //Could also use local minimums to perform the same operation. The results are the same.\n //int min_local = amgx::thrust::reduce(exchanged_aggregates.begin()+A.manager->halo_offsets[i], exchanged_aggregates.begin()+A.manager->halo_offsets[i+1], (int)0xFFFFFFF, amgx::thrust::minimum());\n thrust_wrapper::fill(scratch.begin(), scratch.begin() + s_size, 0);\n int block_size = 128;\n int grid_size = std::min( 4096, ( size + block_size - 1 ) / block_size);\n flag_halo_indices <<< grid_size, block_size>>>(scratch.raw(), exchanged_aggregates.raw() + A.manager->halo_offsets[i], min_index /*min_local*/, size);\n cudaCheckError();\n thrust_wrapper::exclusive_scan(scratch.begin(), scratch.begin() + s_size, scratch.begin());\n apply_halo_aggregate_indices <<< grid_size, block_size>>>(scratch.raw(), exchanged_aggregates.raw() + A.manager->halo_offsets[i], this->m_aggregates.raw() + A.manager->halo_offsets[i], min_index /*min_local*/, m_num_all_aggregates, size);\n cudaCheckError();\n Ac.manager->halo_offsets[i] = m_num_all_aggregates;\n m_num_all_aggregates += scratch[s_size - 1];\n }\n\n cudaCheckError();\n Ac.manager->halo_offsets[num_neighbors] = m_num_all_aggregates;\n}\n\n//TODO: The consolidate and unconsolidate parts could be made more efficient by only sending the\n// nonzero values\ntemplate \nvoid Aggregation_AMG_Level_Base::consolidateVector(VVector &x)\n{\n int my_id = this->getA().manager->global_id();\n\n if (this->getA().manager->isRootPartition())\n {\n // Here all partitions being consolidated should have same vector size, see TODO above\n INDEX_TYPE num_parts = this->getA().manager->getNumPartsToConsolidate();\n\n for (int i = 0; i < num_parts; i++)\n {\n int current_part = this->getA().manager->getPartsToConsolidate()[i];\n\n // Vector has been set to correct size\n if (current_part != my_id)\n {\n //printf(\"Root partition %d receiving %d -> %d and %d -> %d (total %d)\\n\", this->getA().manager->global_id(), this->getA().manager->getConsolidationArrayOffsets()[i], this->getA().manager->getConsolidationArrayOffsets()[i+1], this->getA().manager->getConsolidationArrayOffsets()[num_parts+i], this->getA().manager->getConsolidationArrayOffsets()[num_parts+i+1], (int)x.size()/x.get_block_size());\n this->getA().manager->getComms()->recv_vector(x, current_part, 10000 + current_part, x.get_block_size()*this->getA().manager->getConsolidationArrayOffsets()[i], x.get_block_size() * (this->getA().manager->getConsolidationArrayOffsets()[i + 1] - this->getA().manager->getConsolidationArrayOffsets()[i]));\n this->getA().manager->getComms()->recv_vector(x, current_part, 20000 + current_part, x.get_block_size()*this->getA().manager->getConsolidationArrayOffsets()[num_parts + i], x.get_block_size() * (this->getA().manager->getConsolidationArrayOffsets()[num_parts + i + 1] - this->getA().manager->getConsolidationArrayOffsets()[num_parts + i]));\n }\n }\n }\n else\n {\n int my_destination_part = this->getA().manager->getMyDestinationPartition();\n int i_off, i_size, b_off, b_size;\n this->getA().manager->getConsolidationOffsets(&i_off, &i_size, &b_off, &b_size);\n // Here all partitions being consolidated should have same vector size, see TODO above\n this->getA().manager->getComms()->send_vector_async(x, my_destination_part, 10000 + my_id, i_off * x.get_block_size(), i_size * x.get_block_size());\n this->getA().manager->getComms()->send_vector_async(x, my_destination_part, 20000 + my_id, b_off * x.get_block_size(), b_size * x.get_block_size());\n }\n}\n\n//TODO: The consolidate and unconsolidate parts could be made more efficient by only sending the\n// nonzero values\ntemplate \nvoid Aggregation_AMG_Level_Base::unconsolidateVector(VVector &x)\n{\n if (this->getA().manager->isRootPartition())\n {\n INDEX_TYPE num_parts = this->getA().manager->getNumPartsToConsolidate();\n\n for (int i = 0; i < num_parts; i++)\n {\n int current_part = this->getA().manager->getPartsToConsolidate()[i];\n\n // Vector has been set to correct size\n if (current_part != this->getA().manager->global_id())\n {\n this->getA().manager->getComms()->send_vector_async(x, current_part, 30000 + current_part, x.get_block_size()*this->getA().manager->getConsolidationArrayOffsets()[i], x.get_block_size() * (this->getA().manager->getConsolidationArrayOffsets()[i + 1] - this->getA().manager->getConsolidationArrayOffsets()[i]));\n this->getA().manager->getComms()->send_vector_async(x, current_part, 40000 + current_part, x.get_block_size()*this->getA().manager->getConsolidationArrayOffsets()[num_parts + i], x.get_block_size() * (this->getA().manager->getConsolidationArrayOffsets()[num_parts + i + 1] - this->getA().manager->getConsolidationArrayOffsets()[num_parts + i]));\n }\n }\n }\n else\n {\n int my_destination_part = this->getA().manager->getMyDestinationPartition();\n // Vector x is of unknown size\n int i_off, i_size, b_off, b_size;\n this->getA().manager->getConsolidationOffsets(&i_off, &i_size, &b_off, &b_size);\n this->getA().manager->getComms()->recv_vector(x, my_destination_part, 30000 + this->getA().manager->global_id(), i_off * x.get_block_size(), i_size * x.get_block_size());\n this->getA().manager->getComms()->recv_vector(x, my_destination_part, 40000 + this->getA().manager->global_id(), b_off * x.get_block_size(), b_size * x.get_block_size());\n }\n}\n\n\ntemplate \nvoid Aggregation_AMG_Level_Base::createCoarseVertices()\n{\n //Set the aggregates\n this->m_selector->setAggregates(this->getA(), this->m_aggregates, this->m_aggregates_fine_idx, this->m_num_aggregates);\n\n if ( this->m_print_aggregation_info )\n {\n this->m_selector->printAggregationInfo( this->m_aggregates, this->m_aggregates_fine_idx, this->m_num_aggregates );\n }\n\n this->getA().template setParameter< int > (\"aggregates_num\", this->m_num_aggregates); // ptr to aaggregates\n}\n\n// Creating the next level\ntemplate \nvoid Aggregation_AMG_Level_Base::createCoarseMatrices()\n{\n Matrix &A = this->getA();\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n\n /* WARNING: do not recompute prolongation (P) and restriction (R) when you\n are reusing the level structure (structure_reuse_levels > 0).\n Notice that in aggregation path, prolongation P is implicit,\n and is used through the aggregates array. */\n\n bool const consolidation_level = !A.is_matrix_singleGPU() && this->isConsolidationLevel();\n\n // bookkeeping for the coarse grid: renumber aggregates,\n // if consolidation compute consolidated halo-offsets, etc\n if (!this->isReuseLevel())\n {\n if (consolidation_level)\n {\n // Consolidation-path steps 1-9\n this->consolidationBookKeeping();\n }\n else\n {\n this->setNeighborAggregates();\n }\n }\n\n this->getA().setView(ALL);\n\n // Compute restriction operator\n // TODO: computing the restriction operator could be merged with the selector to save some work\n // If we reuse the level we keep the previous restriction operator\n if (this->isReuseLevel() == false)\n {\n computeRestrictionOperator();\n }\n\n Ac.set_initialized(0);\n Ac.copyAuxData(&A);\n this->m_coarseAGenerator->computeAOperator(A, Ac, this->m_aggregates, this->m_R_row_offsets, this->m_R_column_indices, this->m_num_all_aggregates);\n Ac.setColsReorderedByColor(false);\n Ac.setView(FULL);\n\n if (consolidation_level)\n {\n // Consolidation-path Steps 11-12, send matrices to root, consolidate, final bookkeeping\n this->consolidateCoarseGridMatrix();\n }\n else\n {\n this->prepareNextLevelMatrix(A, Ac);\n }\n\n A.setView(OWNED);\n Ac.setView(OWNED);\n\n this->m_next_level_size = this->m_num_all_aggregates * Ac.get_block_dimy();\n\n if (this->m_print_aggregation_info)\n {\n MatrixAnalysis ana(&Ac);\n ana.aggregatesQuality2(this->m_aggregates, this->m_num_aggregates, A);\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::consolidationBookKeeping()\n{\n Matrix &A = this->getA();\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n\n\n int num_parts, num_fine_neighbors, my_id;\n\n if (!A.is_matrix_singleGPU())\n {\n num_parts = A.manager->getComms()->get_num_partitions();\n num_fine_neighbors = A.manager->neighbors.size();\n my_id = A.manager->global_id();\n }\n else\n {\n num_parts = 1;\n num_fine_neighbors = 0;\n my_id = 0;\n }\n\n // ----------------------------------------------------\n // Consolidate multiple fine matrices into one coarse matrix\n // ----------------------------------------------------\n // ----------------\n // Step 1\n // Decide which partitions should be merged together, store in destination_partitions vector\n // ---------------\n IVector_h &destination_part = A.manager->getDestinationPartitions();\n int my_destination_part = A.manager->getMyDestinationPartition();\n\n if (my_destination_part >= num_parts)\n {\n FatalError(\"During consolidation, sending data to partition that doesn't exist\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // Create mapping from coarse partition indices (ranks on the coarse consolidated level) to partition indices on the fine level (ranks on the fine level)\n IVector_h coarse_part_to_fine_part = destination_part;\n amgx::thrust::sort(coarse_part_to_fine_part.begin(), coarse_part_to_fine_part.end());\n cudaCheckError();\n coarse_part_to_fine_part.erase(thrust::unique(coarse_part_to_fine_part.begin(), coarse_part_to_fine_part.end()), coarse_part_to_fine_part.end());\n cudaCheckError();\n //Then, the number of coarse partitions is simply the size of this vector\n int num_coarse_partitions = coarse_part_to_fine_part.size();\n // Create mapping from fine partition indices to coarse partition indices, with fine partitions that are merging together having the same coarse indices\n IVector_h fine_part_to_coarse_part(num_parts);\n amgx::thrust::lower_bound(coarse_part_to_fine_part.begin(), coarse_part_to_fine_part.end(), destination_part.begin(), destination_part.end(), fine_part_to_coarse_part.begin());\n cudaCheckError();\n // Create mapping from this specific partition's neighbors to consolidated coarse neighbors, but using their fine index (aka. destination partition indices for my neighbors)\n IVector_h fine_neigh_to_fine_part;\n A.manager->createNeighToDestPartMap(fine_neigh_to_fine_part, A.manager->neighbors, destination_part, num_fine_neighbors);\n // Create mapping from consolidated coarse neighbors to fine partition indices (even if the current partition is not going to be a root)\n IVector_h coarse_neigh_to_fine_part;\n int num_coarse_neighbors;\n A.manager->createConsolidatedNeighToPartMap(coarse_neigh_to_fine_part, fine_neigh_to_fine_part, my_destination_part, destination_part, num_coarse_neighbors);\n // Create mapping from fine neighbors to coarse neighbors, with fine neighbors this partition is merging with labeled with -1\n IVector_h fine_neigh_to_coarse_neigh;\n A.manager->createNeighToConsNeigh(fine_neigh_to_coarse_neigh, coarse_neigh_to_fine_part, fine_neigh_to_fine_part, my_destination_part, num_fine_neighbors);\n /*\n EXAMPLE\n Take the following partition graph (that describes connections between partitions, vertices are the partitions themselves), this is the same graph that is used in the setup example\n number of partitions num_parts=12\n CSR row_offsets [0 4 8 13 21 25 32 36 41 46 50 57 61]\n CSR col_indices [0 1 3 8\n 0 1 2 3\n 1 2 3 4 5\n 0 1 2 3 4 5 8 10\n 2 4 5 6\n 2 3 4 5 6 7 10\n 4 5 6 7\n 5 6 7 9 10\n 0 3 8 10 11\n 7 9 10 11\n 3 5 7 8 9 10 11\n 8 9 10 11]\n destination_part = [0 0 0 0 4 4 4 4 8 8 8 8]\n coarse_part_to_fine_part = [0 4 8] num_coarse_partitions = 3\n fine_part_to_coarse_part = [0 0 0 0 1 1 1 1 2 2 2 2]\n original neighbor lists correspond to the rows of the matrix, minus the diagonal elements: (part 0)[1 3 8] (part 3)[0 1 2 4 5 8 10] (part 10)[3 5 7 8 9 11]\n fine_neigh_to_fine_part (part 0)[0 0 2] (part 3)[0 0 0 0 1 2 2] (part 10)[0 1 1 2 2 2]\n coarse_neigh_to_fine_part (part 0)[8] (part 3)[4 8] (part 10)[0 4]\n fine_neigh_to_coarse_neigh (part 0)[-1 -1 0] (part 3)[-1 -1 -1 0 0 1 1] (part 10)[0 1 1 -1 -1 -1]\n */\n // --------------------------\n // Step 2\n // Create coarse B2L_maps, by mapping fine B2L maps to coarse indices using this->m_aggregates and eliminating duplicates\n // --------------------------\n std::vector coarse_B2L_maps(num_fine_neighbors);\n m_num_all_aggregates = m_num_aggregates;\n int num_neighbors_temp = A.manager->neighbors.size();\n int num_rings = A.manager->B2L_rings[0].size() - 1;\n\n if (num_rings != 1)\n {\n FatalError(\"num_rings > 1 not supported in consolidation\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n IndexType max_b2l = 0;\n\n\n for (int i = 0; i < num_neighbors_temp; i++ ) { max_b2l = max_b2l > A.manager->B2L_rings[i][1] ? max_b2l : A.manager->B2L_rings[i][1]; }\n\n IVector B2L_aggregates(max_b2l);\n IVector indices(max_b2l);\n\n //TODO: use the algorithm from setNeighborAggregates()\n for (int i = 0; i < num_neighbors_temp; i++ )\n {\n int size = A.manager->B2L_rings[i][1];\n thrust_wrapper::fill(B2L_aggregates.begin(), B2L_aggregates.begin() + size, 0);\n thrust_wrapper::sequence(indices.begin(), indices.begin() + size);\n //substitute coarse aggregate indices for fine boundary nodes\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(this->m_aggregates.begin(), A.manager->B2L_maps[i].begin()),\n amgx::thrust::make_permutation_iterator(this->m_aggregates.begin(), A.manager->B2L_maps[i].begin() + size),\n B2L_aggregates.begin());\n //find the unique ones\n amgx::thrust::sort_by_key(B2L_aggregates.begin(), B2L_aggregates.begin() + size, indices.begin());\n IndexType num_unique = amgx::thrust::unique_by_key(B2L_aggregates.begin(), B2L_aggregates.begin() + size, indices.begin()).first - B2L_aggregates.begin();\n coarse_B2L_maps[i].resize(num_unique);\n //sort it back so we have the original ordering\n amgx::thrust::sort_by_key(indices.begin(), indices.begin() + num_unique, B2L_aggregates.begin());\n amgx::thrust::copy(B2L_aggregates.begin(), B2L_aggregates.begin() + num_unique, coarse_B2L_maps[i].begin());\n }\n\n cudaCheckError();\n /*\n * EXAMPLE\n say, partition 3 has the following coarse B2L_maps:\n neighbors [0 1 2 4 5 8 10]\n B2L_maps[0(=0)] = [6 7 8]\n B2L_maps[1(=1)] = [8 9 10]\n B2L_maps[2(=2)] = [10 11 12 13]\n B2L_maps[3(=4)] = [13 14 15]\n B2L_maps[4(=5)] = [15 16 17]\n B2L_maps[5(=8)] = [6 18 19]\n B2L_maps[6(=10)] = [17 20 19]\n */\n // ---------------------------------------------------\n // Step 3\n // create new B2L maps for each merged destination neighbor and drop B2L maps to neighbors we are merging with\n // ---------------------------------------------------\n std::vector dest_coarse_B2L_maps;\n A.manager->consolidateB2Lmaps(dest_coarse_B2L_maps, coarse_B2L_maps, fine_neigh_to_coarse_neigh, num_coarse_neighbors, num_fine_neighbors);\n /*\n * EXAMPLE\n Then, merging the coarse B2L maps on partition 3, we get:\n coarse_neigh_to_fine_part [4 8]\n dest_coarse_B2L_maps[0(=4)] = [13 14 15 16 17]\n dest_coarse_B2L_maps[1(=8)] = [6 17 18 19 20]\n */\n // -----------------------\n // Step 4\n // Create interior-boundary renumbering of aggregates according to dest_coarse_B2L_maps\n // -----------------------\n // Now renumber the aggregates with all interior aggregates first, boundary aggregates second\n int num_interior_aggregates; //returned by createAggregatesRenumbering\n int num_boundary_aggregates; //returned by createAggregatesRenumbering\n IVector renumbering; //returned by createAggregatesRenumbering\n // Following calls create renumbering array and modifies B2L_maps\n A.manager->createAggregatesRenumbering(renumbering, dest_coarse_B2L_maps, this->m_num_aggregates, num_coarse_neighbors, num_interior_aggregates, num_boundary_aggregates, num_rings);\n /*\n * EXAMPLE\n Partition 3 will get a renumbering vector of size 21, for the 21 owned agggregates:\n [0 1 2 3 4 5 17 6 7 8 9 10 11 12 13 14 15 16 18 19 20]\n num_interior_aggregates = 12\n num_boundary_aggregates = 9\n */\n // -------------------------------------------------\n // Step 5\n // Determine whether root partition, make list of partitions merged into one\n // ------------------------------------------------\n // Check if I'm root partition and how fine partitions (including myself) are merging into me\n // bool is_root_partition = false;\n bool &is_root_partition = this->m_is_root_partition;\n is_root_partition = false; \n int num_fine_parts_to_consolidate = 0;\n // IVector_h fine_parts_to_consolidate;\n IVector_h &fine_parts_to_consolidate = this->m_fine_parts_to_consolidate;\n\n for (int i = 0; i < num_parts; i++)\n {\n if (destination_part[i] == my_id)\n {\n is_root_partition = true;\n num_fine_parts_to_consolidate++;\n }\n }\n\n fine_parts_to_consolidate.resize(num_fine_parts_to_consolidate);\n int count = 0;\n\n for (int i = 0; i < num_parts; i++)\n {\n if (destination_part[i] == my_id)\n {\n fine_parts_to_consolidate[count] = i;\n count++;\n }\n }\n\n //save this information as state, as this will also be required during solve for restriction/prolongation\n A.manager->setIsRootPartition(is_root_partition);\n A.manager->setNumPartsToConsolidate(num_fine_parts_to_consolidate);\n A.manager->setPartsToConsolidate(fine_parts_to_consolidate);\n\n // Create a new distributed communicator for coarse levels that only contains active partitions\n if (Ac.manager == NULL)\n {\n Ac.manager = new DistributedManager();\n }\n\n Ac.manager->setComms(A.manager->getComms()->Clone());\n Ac.manager->getComms()->createSubComm(coarse_part_to_fine_part, is_root_partition);\n\n\n /*\n * EXAMPLE\n isRootPartition is true for partitions 0,4,8 false for others\n num_fine_parts_to_consolidate = 4 for partitions 0,4,8\n fine_parts_to_consolidate (part 0)[0 1 2 3] (part 4)[4 5 6 7] (part 8)[8 9 10 11]\n */\n // ----------------------\n // Step 6\n // Compute number of interior, boundary and total nodes in the consolidated coarse matrix. Create offsets so that partitions being merged together will have their aggregate indices ordered like this:\n // [num_interior(fine_parts_to_consolidate[0]] num_interior(fine_parts_to_consolidate[1]] ... num_interior(fine_parts_to_consolidate[num_fine_parts_to_consolidate]\n // num_boundary(fine_parts_to_consolidate[0]] num_boundary(fine_parts_to_consolidate[1]] ... num_boundary(fine_parts_to_consolidate[num_fine_parts_to_consolidate] ]\n // ----------------------\n // Gather to get number of interior/boundary aggregates of neighbors I will merge with\n // std::vector vertex_counts;\n std::vector &vertex_counts = this->m_vertex_counts;\n // int interior_offset, boundary_offset, total_interior_rows_in_merged, total_boundary_rows_in_merged;\n int interior_offset, boundary_offset;\n int &total_interior_rows_in_merged = this->m_total_interior_rows_in_merged;\n int &total_boundary_rows_in_merged = this->m_total_boundary_rows_in_merged;\n int total_rows_in_merged;\n //Computes these offsets on the root, sends them back\n A.manager->computeConsolidatedOffsets(my_id, my_destination_part, is_root_partition, num_interior_aggregates, num_boundary_aggregates, vertex_counts, fine_parts_to_consolidate, num_fine_parts_to_consolidate, interior_offset, boundary_offset, total_interior_rows_in_merged, total_boundary_rows_in_merged, total_rows_in_merged, A.manager->getComms());\n //Partitions save these offsets, as it will be required during solve restriction/prolongation\n A.manager->setConsolidationOffsets(interior_offset, num_interior_aggregates, boundary_offset + num_interior_aggregates, num_boundary_aggregates);\n /*\n * EXAMPLE\n For root partition 0, say we have the following interior/boundary counts (note that partition 1 has 0 boundary, as it is only connected to partitions it is merging with)\n part 0 - interior: 10 boundary 3\n part 1 - interior: 18\n part 2 - interior: 10 boundary 16\n part 3 - interior: 12 boundary 9\n interior_offset for partitions 0,1,2,3: 0 10 28 38 (total_interior_rows_in_merged 50)\n boundary_offset for partitions 0,1,2,3: 0 3 3 19 (total_boundary_rows_in_merged 28)\n */\n // ----------------------\n // Step 7\n // Each partition renumbers its aggregates and dest_coarse_B2L_maps using offsets computed in Step 6 and permutation in Step 4\n // ----------------------\n // Kernel to renumber the aggregates\n int block_size = 128;\n int grid_size = std::min( 4096, ( A.manager->halo_offsets[0] + block_size - 1 ) / block_size);\n renumberAggregatesKernel <<< grid_size, block_size >>>(renumbering.raw(), interior_offset, boundary_offset, this->m_aggregates.raw(), A.manager->halo_offsets[0], num_interior_aggregates, renumbering.size());\n cudaCheckError();\n\n for (int i = 0; i < num_coarse_neighbors; i++)\n {\n thrust_wrapper::transform(dest_coarse_B2L_maps[i].begin(),\n dest_coarse_B2L_maps[i].end(),\n amgx::thrust::constant_iterator(boundary_offset),\n dest_coarse_B2L_maps[i].begin(),\n amgx::thrust::plus());\n }\n\n cudaCheckError();\n /*\n * EXAMPLE\n Partition 3 had a renumbering vector:\n [0 1 2 3 4 5 17 6 7 8 9 10 11 12 13 14 15 16 18 19 20]\n which is now adjusted to account for the consolidated coarse matrices' indices:\n [38 39 40 41 42 43 74 44 45 46 47 48 49 69 70 71 72 73 75 76 77]\n And the dest_coarse_B2L_maps, which looked like:\n dest_coarse_B2L_maps[0(=4)] = [13 14 15 16 17]\n dest_coarse_B2L_maps[1(=8)] = [6 17 18 19 20]\n is now:\n dest_coarse_B2L_maps[0(=4)] = [69 70 71 72 73]\n dest_coarse_B2L_maps[1(=8)] = [74 73 75 76 77]\n */\n // -------------------------------------------------\n // Step 8\n // Send dest_coarse_B2L_maps to root partitions\n // ------------------------------------------------\n // Each fine partition sends to its root the number of coarse neighbors it has, their ids, and the number of boundary nodes for each coarse neighbor\n IVector_h num_bdy_per_coarse_neigh(num_coarse_neighbors);\n\n for (int i = 0; i < num_coarse_neighbors; i++)\n {\n num_bdy_per_coarse_neigh[i] = dest_coarse_B2L_maps[i].size();\n }\n\n IVector_h consolidated_coarse_neigh_to_fine_part; //consolidated list of coarse neighbors for the root partition, using fine partition indices\n int num_consolidated_neighbors = 0;\n // std::vector consolidated_B2L_maps; //concatenates dest_coarse_B2L_maps received from partitions that are merging into the same root and pointing to the same destination coarse neighbor\n std::vector &consolidated_B2L_maps = this->m_consolidated_B2L_maps;\n A.manager->consolidateB2LmapsOnRoot(num_consolidated_neighbors, consolidated_B2L_maps, consolidated_coarse_neigh_to_fine_part, dest_coarse_B2L_maps, coarse_neigh_to_fine_part, num_bdy_per_coarse_neigh, fine_parts_to_consolidate, num_fine_parts_to_consolidate, my_id, my_destination_part, is_root_partition, num_coarse_neighbors, A.manager->getComms());\n //\n // Step 9 - figuring out halo aggregate IDs\n //\n //Now we need to update halo aggregate IDs - this is just a halo exchange on this->m_aggregates between partitions\n //that are being merged together, but we need to send other halos to the root to come up with the halo renumbering\n //TODO: separate transactions, send \"real halo\" to the root nodes (coarse neighbors) immediately\n //Step 9.1: takes care of synchronizing the aggregate IDs between partitions we are merging together and got consistent halo aggregate IDs for neighbor we are not merging with (which are going to be sent to the root in 9.2)\n A.manager->exchange_halo(this->m_aggregates, 6666);\n /*\n * EXAMPLE 2\n This example is independent from the previous ones.\n Say partition 0 and 1 are merging (into 0) partition 0 is neighbors with 1,2,3 and partition 1 is neighbors with 0,3,4\n Partitions 3 and 4 are merging (into partition 3) and partition 2 is not merging with anyone.\n This example details the renumbering of halo indices on partition 0 and partition 1.\n After the exchange halo, we have:\n this->m_aggregates on partition 0:\n [(fine interior nodes) (fine boundary nodes) (fine halo from part 1) (fine halo from part 2) (fine halo from part 3)]\n [(fine interior nodes) (fine boundary nodes) (13 13 15) (12 15 17) (14 16 18)]\n aggregates on partition 1:\n [(fine interior nodes) (fine boundary nodes) (fine halo from part 0) (fine halo from part 3) (fine halo from part 4)]\n [(fine interior nodes) (fine boundary nodes) (14 16 17) (18 19 19) (15 15 17)]\n indices in (fine halo from part 0) and (fine halo from part 1) actually contain interior aggregate indices (if they are not connected to partitions 2,3 or 4), because the boundary is disappearing there.\n Indices in halo regions contain remote-local indices.\n\n This example is used throughout consolidateAndRenumberHalos\n */\n //Step 9.2 - 9.5\n // IVector_h halo_offsets(num_consolidated_neighbors + 1, 0);\n IVector_h &halo_offsets = this->m_consolidated_halo_offsets;\n halo_offsets = IVector_h(num_consolidated_neighbors + 1, 0);\n A.manager->consolidateAndRenumberHalos(this->m_aggregates, A.manager->halo_offsets, halo_offsets, A.manager->neighbors, num_fine_neighbors, consolidated_coarse_neigh_to_fine_part, num_consolidated_neighbors, destination_part, my_destination_part, is_root_partition, fine_parts_to_consolidate, num_fine_parts_to_consolidate, num_parts, my_id, total_rows_in_merged, this->m_num_all_aggregates, A.manager->getComms());\n\n if (is_root_partition)\n {\n for (int i = 0; i < consolidated_B2L_maps.size(); i++)\n {\n amgx::thrust::sort(consolidated_B2L_maps[i].begin(), consolidated_B2L_maps[i].end());\n }\n\n this->m_consolidated_neighbors.resize(num_consolidated_neighbors);\n for (int i = 0; i < num_consolidated_neighbors; i++)\n {\n this->m_consolidated_neighbors[i] = fine_part_to_coarse_part[consolidated_coarse_neigh_to_fine_part[i]];\n }\n \n cudaCheckError();\n }\n}\n\ntemplate \nvoid Aggregation_AMG_Level_Base::consolidateCoarseGridMatrix()\n{\n Matrix &A = this->getA();\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n\n int my_id = A.manager->global_id();\n \n IVector_h &destination_part = A.manager->getDestinationPartitions();\n int my_destination_part = A.manager->getMyDestinationPartition();\n\n // bookkeeping stored in AMG_Level_Base\n std::vector &vertex_counts = this->m_vertex_counts;\n IVector_h &fine_parts_to_consolidate = this->m_fine_parts_to_consolidate;\n\n // bookkeeping stored in either AMG_Level_Base or Acs' DistributedManager\n IVector_h &halo_offsets = this->isReuseLevel() ? Ac.manager->getHaloOffsets() : this->m_consolidated_halo_offsets;\n std::vector &consolidated_B2L_maps = this->isReuseLevel() ? Ac.manager->getB2Lmaps() : this->m_consolidated_B2L_maps;\n\n int num_consolidated_neighbors = this->isRootPartition() ? this->m_consolidated_neighbors.size() : 0;\n\n if (!this->isRootPartition())\n {\n A.manager->getComms()->send_vector_async(Ac.row_offsets, my_destination_part, 1111);\n A.manager->getComms()->send_vector_async(Ac.col_indices, my_destination_part, 1112);\n A.manager->getComms()->send_vector_async(Ac.values, my_destination_part, 1113);\n }\n else\n {\n int num_fine_parts_to_consolidate = fine_parts_to_consolidate.size();\n\n int total_num_rows = this->m_num_all_aggregates;\n IVector new_row_offsets(total_num_rows + 1, 0);\n\n //if diags are inside then we won't be counting those twice when computing halo row length\n if (!Ac.hasProps(DIAG))\n {\n thrust_wrapper::fill(new_row_offsets.begin() + halo_offsets[0], new_row_offsets.begin() + halo_offsets[num_consolidated_neighbors], 1);\n cudaCheckError();\n }\n\n std::vector recv_row_offsets(num_fine_parts_to_consolidate);\n std::vector num_nz(num_fine_parts_to_consolidate);\n IVector *work_row_offsets;\n std::vector index_offset_array(2 * num_fine_parts_to_consolidate + 1);\n int interior_offset = 0;\n int boundary_offset = 0;\n\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n boundary_offset += vertex_counts[i][0];\n }\n\n int max_num_nz = 0;\n\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n\n //receive row offsets\n if (current_part != my_id)\n {\n recv_row_offsets[i].resize(total_num_rows + 1);\n A.manager->getComms()->recv_vector(recv_row_offsets[i], current_part, 1111);\n work_row_offsets = &(recv_row_offsets[i]);\n num_nz[i] = (*work_row_offsets)[work_row_offsets->size() - 1];\n max_num_nz = max_num_nz > num_nz[i] ? max_num_nz : num_nz[i];\n }\n else\n {\n work_row_offsets = &(Ac.row_offsets);\n num_nz[i] = Ac.get_num_nz();\n }\n\n //Get interior row length\n thrust_wrapper::transform(work_row_offsets->begin() + interior_offset + 1,\n work_row_offsets->begin() + interior_offset + vertex_counts[i][0] + 1,\n work_row_offsets->begin() + interior_offset,\n new_row_offsets.begin() + interior_offset,\n amgx::thrust::minus());\n cudaCheckError();\n //Get boundary row length\n thrust_wrapper::transform(work_row_offsets->begin() + boundary_offset + 1,\n work_row_offsets->begin() + boundary_offset + vertex_counts[i][1] + 1,\n work_row_offsets->begin() + boundary_offset,\n new_row_offsets.begin() + boundary_offset,\n amgx::thrust::minus());\n cudaCheckError();\n //Increment halo row length by one for every nonzero that is an edge from the halo into this partition\n int size = halo_offsets[num_consolidated_neighbors] - halo_offsets[0];\n const int block_size = 128;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (size - 1) / block_size + 1);\n set_halo_rowlen <<< num_blocks, block_size>>>(work_row_offsets->raw() + halo_offsets[0], new_row_offsets.raw() + halo_offsets[0], size, Ac.hasProps(DIAG));\n cudaCheckError();\n index_offset_array[i] = interior_offset;\n index_offset_array[num_fine_parts_to_consolidate + i] = boundary_offset;\n interior_offset += vertex_counts[i][0];\n boundary_offset += vertex_counts[i][1];\n index_offset_array[i + 1] = interior_offset;\n index_offset_array[num_fine_parts_to_consolidate + i + 1] = boundary_offset;\n }\n\n A.manager->setConsolidationArrayOffsets(index_offset_array);\n //Exclusive scan row length array to get row offsets\n thrust_wrapper::exclusive_scan(new_row_offsets.begin(), new_row_offsets.end(), new_row_offsets.begin());\n cudaCheckError();\n //Prepare to receive column indices and values\n int num_nz_consolidated = new_row_offsets[new_row_offsets.size() - 1];\n IVector recv_col_indices(max_num_nz);\n IVector new_col_indices(num_nz_consolidated);\n MVector recv_values((max_num_nz + 1 + Ac.hasProps(DIAG) * (halo_offsets[num_consolidated_neighbors] - 1))*Ac.get_block_size());\n MVector new_values((num_nz_consolidated + 1 + Ac.hasProps(DIAG) * (halo_offsets[num_consolidated_neighbors] - 1))*Ac.get_block_size());\n thrust_wrapper::fill(new_col_indices.begin() + new_row_offsets[halo_offsets[0]], new_col_indices.end(), -1); //Set all the halo col indices to -1\n\n\n if (!Ac.hasProps(DIAG)) { thrust_wrapper::fill(new_values.begin() + num_nz_consolidated * Ac.get_block_size(), new_values.end(), types::util::get_zero()); }\n\n cudaCheckError();\n IVector *work_col_indices;\n MVector *work_values;\n interior_offset = 0;\n boundary_offset = 0;\n\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n boundary_offset += vertex_counts[i][0];\n }\n\n for (int i = 0; i < num_fine_parts_to_consolidate; i++)\n {\n int current_part = fine_parts_to_consolidate[i];\n\n if (current_part != my_id)\n {\n A.manager->getComms()->recv_vector(recv_col_indices, current_part, 1112, 0, num_nz[i]);\n A.manager->getComms()->recv_vector(recv_values, current_part, 1113, 0, (num_nz[i] + 1 + Ac.hasProps(DIAG) * (halo_offsets[num_consolidated_neighbors] - 1))*Ac.get_block_size());\n work_col_indices = &(recv_col_indices);\n work_row_offsets = &(recv_row_offsets[i]);\n work_values = &(recv_values);\n }\n else\n {\n work_row_offsets = &(Ac.row_offsets);\n work_col_indices = &(Ac.col_indices);\n work_values = &(Ac.values);\n }\n\n //Put interior rows in place\n amgx::thrust::copy(work_col_indices->begin() + (*work_row_offsets)[interior_offset],\n work_col_indices->begin() + (*work_row_offsets)[interior_offset + vertex_counts[i][0]],\n new_col_indices.begin() + new_row_offsets[interior_offset]);\n cudaCheckError();\n amgx::thrust::copy(work_values->begin() + (*work_row_offsets)[interior_offset]*Ac.get_block_size(),\n work_values->begin() + ((*work_row_offsets)[interior_offset + vertex_counts[i][0]])*Ac.get_block_size(),\n new_values.begin() + new_row_offsets[interior_offset]*Ac.get_block_size());\n cudaCheckError();\n //Put boundary rows in place\n amgx::thrust::copy(work_col_indices->begin() + (*work_row_offsets)[boundary_offset],\n work_col_indices->begin() + (*work_row_offsets)[boundary_offset + vertex_counts[i][1]],\n new_col_indices.begin() + new_row_offsets[boundary_offset]);\n cudaCheckError();\n amgx::thrust::copy(work_values->begin() + (*work_row_offsets)[boundary_offset]*Ac.get_block_size(),\n work_values->begin() + ((*work_row_offsets)[boundary_offset + vertex_counts[i][1]])*Ac.get_block_size(),\n new_values.begin() + new_row_offsets[boundary_offset]*Ac.get_block_size());\n cudaCheckError();\n //Process halo rows (merge)\n int size = halo_offsets[num_consolidated_neighbors] - halo_offsets[0];\n const int block_size = 128;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (size - 1) / block_size + 1);\n //TODO: vectorise this kernel, will be inefficient for larger block sizes\n append_halo_nz <<< num_blocks, block_size>>>(work_row_offsets->raw() + halo_offsets[0],\n new_row_offsets.raw() + halo_offsets[0],\n work_col_indices->raw(),\n new_col_indices.raw(),\n work_values->raw(),\n new_values.raw(),\n size, Ac.hasProps(DIAG), halo_offsets[0], Ac.get_block_size());\n cudaCheckError();\n\n // Diagonals\n if (Ac.hasProps(DIAG))\n {\n // Diagonal corresponding to interior rows\n amgx::thrust::copy(work_values->begin() + (num_nz[i] + interior_offset)*Ac.get_block_size(),\n work_values->begin() + (num_nz[i] + interior_offset + vertex_counts[i][0])*Ac.get_block_size(),\n new_values.begin() + (new_row_offsets[halo_offsets[halo_offsets.size() - 1]] + interior_offset)*Ac.get_block_size());\n // Diagonal corresponding to boundary rows\n amgx::thrust::copy(work_values->begin() + (num_nz[i] + boundary_offset)*Ac.get_block_size(),\n work_values->begin() + (num_nz[i] + boundary_offset + vertex_counts[i][1])*Ac.get_block_size(),\n new_values.begin() + (new_row_offsets[halo_offsets[halo_offsets.size() - 1]] + boundary_offset)*Ac.get_block_size());\n cudaCheckError();\n }\n\n interior_offset += vertex_counts[i][0];\n boundary_offset += vertex_counts[i][1];\n }\n\n Ac.set_initialized(0);\n Ac.row_offsets = new_row_offsets;\n Ac.col_indices = new_col_indices;\n Ac.values = new_values;\n }\n\n // A new distributed communicator for coarse levels that only contains active partitions\n // has already been created in consolidatedBookKeeping!\n\n //\n // Step 12 - finalizing, bookkeping\n //\n if (this->isRootPartition())\n {\n // int my_consolidated_id = fine_part_to_coarse_part[my_id];\n int my_consolidated_id = Ac.manager->getComms()->get_global_id();\n\n if (!this->isReuseLevel())\n {\n Ac.manager->initializeAfterConsolidation(\n my_consolidated_id,\n Ac,\n this->m_consolidated_neighbors,\n this->m_total_interior_rows_in_merged,\n this->m_total_boundary_rows_in_merged,\n this->m_num_all_aggregates,\n this->m_consolidated_halo_offsets,\n this->m_consolidated_B2L_maps,\n 1,\n true);\n\n // this is now stored in Acs DistributedManager\n this->m_consolidated_neighbors.resize(0);\n this->m_consolidated_halo_offsets.resize(0);\n this->m_consolidated_B2L_maps.resize(0);\n\n Ac.manager->B2L_rings.resize(num_consolidated_neighbors + 1);\n\n for (int i = 0; i < num_consolidated_neighbors; i++)\n {\n Ac.manager->B2L_rings[i].resize(2);\n Ac.manager->B2L_rings[i][0] = 0;\n Ac.manager->B2L_rings[i][1] = consolidated_B2L_maps[i].size();\n }\n }\n\n Ac.manager->set_initialized(Ac.row_offsets);\n Ac.manager->getComms()->set_neighbors(num_consolidated_neighbors);\n int new_nnz = Ac.row_offsets[Ac.row_offsets.size() - 1];\n\n Ac.set_num_nz(new_nnz);\n Ac.set_num_cols(Ac.manager->halo_offsets[Ac.manager->halo_offsets.size() - 1]);\n Ac.set_num_rows(Ac.get_num_cols());\n\n if (A.hasProps(DIAG)) { Ac.addProps(DIAG); }\n\n Ac.computeDiagonal();\n Ac.set_initialized(1);\n }\n else\n {\n this->getA().manager->getComms()->send_vector_wait_all(Ac.row_offsets);\n this->getA().manager->getComms()->send_vector_wait_all(Ac.col_indices);\n this->getA().manager->getComms()->send_vector_wait_all(Ac.values);\n\n Ac.set_initialized(0);\n // set size of Ac to be zero\n Ac.resize(0, 0, 0, 1);\n Ac.set_initialized(1);\n }\n}\n\n// -------------------------------------------------------------\n// Explicit instantiations\n// -------------------------------------------------------------\n\n#define AMGX_CASE_LINE(CASE) template class Aggregation_AMG_Level::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n}\n\n}\n", "source": "/kaggle/working/amgx/src/aggregation/aggregation_amg_level.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include //count\n#include //sort\n#include //lower_bound\n#include //unique\n#include //offsets_to_indices\n\n#define EXPERIMENTAL_ITERATIVE_MATCHING\n\nnamespace amgx\n{\nnamespace aggregation\n{\nnamespace size2_selector\n{\n\n// include common routines for all selectors\n#include \n\n// ------------------------\n// Kernels\n// ------------------------\n\n#ifndef DELETE\n// Kernel to compute the weight of the edges with block_dia_csr format\ntemplate \n__global__\nvoid computeEdgeWeightsBlockDiaCsr(const IndexType *row_offsets, const IndexType *column_indices,\n const ValueType *dia_values, const ValueType *nonzero_values, const IndexType num_block_rows, float *edge_weights, int bsize)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int jcol;\n int bsize_sq = bsize * bsize;\n\n while (tid < num_block_rows)\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (jcol >= num_block_rows) { continue; }\n\n // Compute edge weight\n for (int k = row_offsets[jcol]; k < row_offsets[jcol + 1]; k++)\n {\n if (column_indices[k] == tid)\n {\n edge_weights[j] = (float) 0.5 * (fabs(nonzero_values[j * bsize_sq]) + fabs(nonzero_values[k * bsize_sq]))\n / max( fabs(dia_values[tid * bsize_sq]), fabs(dia_values[jcol * bsize_sq]));\n break;\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n#endif\n\n// findStrongestNeighbour kernel for csr_matrix format\n// computes weight on the fly\ntemplate \n__global__\nvoid findStrongestNeighbourCsr(const IndexType *row_offsets, const IndexType *column_indices,\n const ValueType *values, const ValueType *diag, const IndexType num_rows, IndexType *aggregates, int *strongest_neighbour)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n ValueType weight;\n int jcol;\n\n while (tid < num_rows)\n {\n ValueType max_weight_unaggregated = 0.;\n ValueType max_weight_aggregated = 0.;\n int strongest_unaggregated = -1;\n int strongest_aggregated = -1;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (tid == jcol || jcol >= num_rows) { continue; }\n\n // Compute edge weight\n //weight = fabs(values[j])/max( fabs(diag[tid]),fabs(diag[jcol]));\n for (int k = row_offsets[jcol]; k < row_offsets[jcol + 1]; k++)\n {\n if (column_indices[k] == tid)\n {\n weight = 0.5 * (fabs(values[j]) + fabs(values[k])) / max( fabs(diag[tid]), fabs(diag[jcol]));\n break;\n }\n }\n\n // Identify strongest aggregated and unaggregated neighbours\n if (aggregates[jcol] == -1 && (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && jcol > strongest_unaggregated))) // unaggregated\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n else if (aggregates[jcol] != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) // aggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_unaggregated == -1 && strongest_aggregated != -1) // All neighbours are aggregated\n {\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n else if (strongest_unaggregated != -1)\n {\n strongest_neighbour[tid] = strongest_unaggregated;\n }\n else\n {\n strongest_neighbour[tid] = tid;\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// findStrongestNeighbour kernel for block_dia_csr_matrix format\n// computes weight on the fly\ntemplate \n__global__\nvoid findStrongestNeighbourBlockDiaCsr(const IndexType *row_offsets, const IndexType *column_indices,\n const ValueType *dia_values, const ValueType *nonzero_values, const IndexType num_block_rows, IndexType *aggregates, int *strongest_neighbour, int bsize)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n ValueType weight;\n int jcol;\n int bsize_sq = bsize * bsize;\n\n while (tid < num_block_rows)\n {\n ValueType max_weight_unaggregated = 0.;\n ValueType max_weight_aggregated = 0.;\n int strongest_unaggregated = -1;\n int strongest_aggregated = -1;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (jcol >= num_block_rows) { continue; }\n\n // Compute edge weight\n for (int k = row_offsets[jcol]; k < row_offsets[jcol + 1]; k++)\n {\n if (column_indices[k] == tid)\n {\n weight = 0.5 * (fabs(nonzero_values[j * bsize_sq]) + fabs(nonzero_values[k * bsize_sq]))\n / max( fabs(dia_values[tid * bsize_sq]), fabs(dia_values[jcol * bsize_sq]));\n break;\n }\n }\n\n // Identify strongest aggregated and unaggregated neighbours\n if (aggregates[jcol] == -1 && (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && jcol > strongest_unaggregated))) // unaggregated\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n else if (aggregates[jcol] != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) // aggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_unaggregated == -1 && strongest_aggregated != -1) // All neighbours are aggregated\n // Put in same aggregate as strongest neighbour\n {\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n else if (strongest_unaggregated != -1)\n {\n strongest_neighbour[tid] = strongest_unaggregated;\n }\n else\n {\n strongest_neighbour[tid] = tid;\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\n__device__\nfloat random_weight2(int i, int j)\n{\n#define RAND_MULTIPLIER 1145637293\n unsigned long i_min = (min(i, j) * RAND_MULTIPLIER);\n unsigned long i_max = (max(i, j) * RAND_MULTIPLIER);\n return ((float)i_min / i_max);\n}\n\n\n// findStrongestNeighbour kernel for block_dia_csr_matrix format\n// Reads the weight from edge_weights array\ntemplate \n__global__\nvoid findStrongestNeighbourBlockDiaCsr_V2(const IndexType *row_offsets, const IndexType *column_indices,\n float *edge_weights, IndexType num_block_rows, IndexType *aggregates,\n IndexType *strongest_neighbour_1phase, IndexType *strongest_neighbour,\n const size_t bsize, int phase, bool merge_singletons)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n float weight;\n int jcol;\n\n while (tid < num_block_rows)\n {\n int strongest_unaggregated = -1;\n int strongest_aggregated = -1;\n float max_weight_unaggregated = 0.;\n float max_weight_aggregated = 0.;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (phase == 1) { weight = edge_weights[j]; }\n else { weight = random_weight2(tid, jcol); }\n\n if (tid == jcol || jcol >= num_block_rows) { continue; } // skip diagonal and halo\n\n if (phase == 2 && strongest_neighbour_1phase[jcol] != tid) { continue; } // if 2nd phase only accept those who gave a hand on the 1st phase\n\n // Identify strongest aggregated and unaggregated neighbours\n if (aggregates[jcol] == -1 && (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && jcol > strongest_unaggregated))) // unaggregated\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n else if (aggregates[jcol] != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) // aggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_unaggregated == -1 && strongest_aggregated != -1) // All neighbours are aggregated\n {\n if ( merge_singletons )\n // Put in same aggregate as strongest neighbour\n {\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n else\n {\n aggregates[tid] = tid;\n }\n }\n else if (strongest_unaggregated != -1)\n {\n if (phase == 2)\n {\n float rand_w1 = random_weight2(tid, strongest_neighbour_1phase[tid]);\n strongest_neighbour[tid] = max_weight_unaggregated > rand_w1 ? strongest_unaggregated : strongest_neighbour_1phase[tid];\n }\n else { strongest_neighbour_1phase[tid] = strongest_unaggregated; }\n }\n else\n {\n if (phase == 2) { strongest_neighbour[tid] = strongest_neighbour_1phase[tid]; }\n else { strongest_neighbour_1phase[tid] = tid; }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel that checks if perfect matchs exist\ntemplate \n__global__\nvoid matchEdges(const IndexType num_rows, IndexType *aggregates, int *strongest_neighbour)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int potential_match, potential_match_neighbour;\n\n while (tid < num_rows)\n {\n if (aggregates[tid] == -1) // Unaggregated row\n {\n potential_match = strongest_neighbour[tid];\n potential_match_neighbour = strongest_neighbour[potential_match];\n\n if (potential_match != -1 && potential_match_neighbour == tid) // we have a match\n {\n aggregates[tid] = ( potential_match > tid ) ? tid : potential_match;\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \n__global__\nvoid countAggregates(const IndexType num_rows, IndexType *aggregates, int *num_unaggregated)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int c = 0;\n int i = tid;\n\n while ( i < num_rows )\n {\n c += ( aggregates[i] == -1 );\n i += gridDim.x * blockDim.x;\n }\n\n __shared__ int smem[block_size];\n smem[threadIdx.x] = c;\n __syncthreads();\n\n for ( int off = blockDim.x / 2; off >= 1; off = off / 2 )\n {\n if ( threadIdx.x < off )\n {\n smem[threadIdx.x] += smem[threadIdx.x + off];\n }\n __syncthreads();\n }\n\n if ( threadIdx.x == 0 )\n {\n atomicAdd(num_unaggregated, smem[0]);\n }\n}\n\n// Kernel that merges unaggregated vertices its strongest aggregated neighbour\n// For csr_matrix_format\ntemplate \n__global__\nvoid mergeWithExistingAggregatesCsr(const IndexType *row_offsets, const IndexType *column_indices, const ValueType *values,\n const ValueType *diag, const int num_rows, IndexType *aggregates, int deterministic, IndexType *aggregates_candidate)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int jcol;\n ValueType weight;\n\n while (tid < num_rows)\n {\n ValueType max_weight_aggregated = 0.;\n int strongest_aggregated = -1;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (tid == jcol || jcol >= num_rows) { continue; }\n\n // Compute edge weight\n weight = fabs(values[j]) / max( fabs(diag[tid]), fabs(diag[jcol]));\n\n // Identify strongest aggregated neighbour\n if (aggregates[jcol] != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) // aggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_aggregated != -1) // Found a neighbour to aggregate to\n {\n if (deterministic)\n {\n aggregates_candidate[tid] = aggregates[strongest_aggregated];\n }\n else\n {\n // Put in same aggregate as strongest neighbour\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n }\n else // All neighbours are unaggregated, leave alone\n {\n if (deterministic)\n {\n aggregates_candidate[tid] = tid;\n }\n else\n {\n aggregates[tid] = tid;\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate \n__global__\nvoid joinExistingAggregates(IndexType num_rows, IndexType *aggregates, IndexType *aggregates_candidate)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while (tid < num_rows)\n {\n if (aggregates[tid] == -1 && aggregates_candidate[tid] != -1) // Unaggregated row\n {\n aggregates[tid] = aggregates_candidate[tid];\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel that merges unaggregated vertices its strongest aggregated neighbour\n// Edge weights are computed on the fly\n// For block_dia_csr_matrix_format\ntemplate \n__global__\nvoid mergeWithExistingAggregatesBlockDiaCsr(const IndexType *row_offsets, const IndexType *column_indices, const ValueType *dia_values, const ValueType *nonzero_values,\n const int num_block_rows, IndexType *aggregates, int bsize, int deterministic, IndexType *aggregates_candidate)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int jcol;\n ValueType weight;\n int bsize_sq = bsize * bsize;\n\n while (tid < num_block_rows)\n {\n int strongest_aggregated = -1;\n ValueType max_weight_aggregated = 0.;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (jcol >= num_block_rows) { continue; }\n\n // Compute edge weight\n weight = fabs(nonzero_values[j * bsize_sq]) / max( fabs(dia_values[tid * bsize_sq]), fabs(dia_values[jcol * bsize_sq]));\n\n // Identify strongest aggregated neighbour\n if (aggregates[jcol] != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) // aggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_aggregated != -1) // Found a neighbour to aggregate to\n {\n if (deterministic)\n {\n aggregates_candidate[tid] = aggregates[strongest_aggregated];\n }\n else\n {\n // Put in same aggregate as strongest neighbour\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n }\n else // All neighbours are unaggregated, leave alone\n {\n if (deterministic)\n {\n aggregates_candidate[tid] = tid;\n }\n else\n {\n aggregates[tid] = tid;\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel that merges unaggregated vertices its strongest aggregated neighbour\n// Weights are read from edge_weights array\n// For block_dia_csr_matrix_format\ntemplate \n__global__\nvoid mergeWithExistingAggregatesBlockDiaCsr_V2(const IndexType *row_offsets, const IndexType *column_indices, const float *edge_weights,\n const int num_block_rows, IndexType *aggregates, int bsize, const int deterministic, IndexType *aggregates_candidate)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int jcol;\n float weight;\n\n while (tid < num_block_rows)\n {\n float max_weight_aggregated = 0.;\n int strongest_aggregated = -1;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n // Compute edge weight\n weight = edge_weights[j];\n jcol = column_indices[j];\n\n if (jcol == tid || jcol >= num_block_rows) { continue; } // skip diagonal\n\n // Identify strongest aggregated neighbour\n if (aggregates[jcol] != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) //\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_aggregated != -1) // Found a neighbour to aggregate to\n {\n if (deterministic)\n {\n aggregates_candidate[tid] = aggregates[strongest_aggregated];\n }\n else\n {\n // Put in same aggregate as strongest neighbour\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n }\n else // All neighbours are unaggregated, leave alone\n {\n if (deterministic)\n {\n aggregates_candidate[tid] = tid;\n }\n else\n {\n aggregates[tid] = tid;\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel to extract diagonal for csr_matrix format\ntemplate \n__global__\nvoid getDiagonalKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, const IndexType numRows, ValueType *diagonal)\n{\n int tIdx = threadIdx.x + blockDim.x * blockIdx.x;\n\n while (tIdx < numRows)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n for (int j = offset; j < offset + numj; j++)\n {\n int jcol = column_indices[j];\n\n if (tIdx == jcol)\n {\n diagonal[tIdx] = values[j];\n }\n }\n\n tIdx += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel to extract diagonal for csr_matrix format\ntemplate \n__global__\nvoid getDiagonalKernelNoDiaProp(const IndexType *dia_idx, const ValueType *values, const IndexType numRows, ValueType *diagonal)\n{\n int tIdx = threadIdx.x + blockDim.x * blockIdx.x;\n\n while (tIdx < numRows)\n {\n diagonal[tIdx] = values[dia_idx[tIdx]];\n tIdx += gridDim.x * blockDim.x;\n }\n}\n\n// -----------------\n// Methods\n// ----------------\n\n// Constructor\ntemplate\nSize2SelectorBase::Size2SelectorBase(AMG_Config &cfg, const std::string &cfg_scope)\n{\n deterministic = cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\");\n max_iterations = cfg.AMG_Config::template getParameter(\"max_matching_iterations\", cfg_scope);\n numUnassigned_tol = cfg.AMG_Config::template getParameter(\"max_unassigned_percentage\", cfg_scope);\n two_phase = cfg.AMG_Config::template getParameter(\"handshaking_phases\", cfg_scope) == 2;\n m_aggregation_edge_weight_component = cfg.AMG_Config::template getParameter(\"aggregation_edge_weight_component\", cfg_scope);\n merge_singletons = cfg.AMG_Config::template getParameter(\"merge_singletons\", cfg_scope) == 1;\n weight_formula = cfg.AMG_Config::template getParameter(\"weight_formula\", cfg_scope);\n}\n\n// setAggregates for csr_matrix_h format\ntemplate \nvoid Size2Selector >::setAggregates_1x1(const Matrix_h &A,\n typename Matrix_h::IVector &aggregates, typename Matrix_h::IVector &aggregates_global, int &num_aggregates)\n{\n FatalError(\"Size2 selector: setAggregates not implemented on CPU, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n// setAggregates for block_dia_csr_matrix_h format\ntemplate \nvoid Size2Selector >::setAggregates_common_sqblocks(const Matrix_h &A,\n typename Matrix_h::IVector &aggregates, typename Matrix_h::IVector &aggregates_global, int &num_aggregates)\n{\n FatalError(\"Size2 selector: setAggregates not implemented on CPU, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n#ifndef DELETE\n// setAggregates for csr_matrix_d format\ntemplate \nvoid Size2Selector >::setAggregates_1x1(\n const Matrix_d &A, typename Matrix_d::IVector &aggregates, typename Matrix_d::IVector &aggregates_global, int &num_aggregates)\n\n{\n if (!A.is_matrix_singleGPU())\n {\n aggregates.resize(A.manager->halo_offset(A.manager->num_neighbors()));\n }\n else\n {\n aggregates.resize(A.get_num_rows());\n }\n\n thrust_wrapper::fill(aggregates.begin(), aggregates.end(), -1);\n cudaCheckError();\n //typedef typename csr_matrix_d::index_type IndexType;\n //typedef typename csr_matrix_d::value_type ValueType;\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const ValueType *A_values_ptr = A.values.raw();\n const IndexType *A_dia_ptr = A.diag.raw();\n IndexType num_rows = (int)A.get_num_rows();\n typename Matrix_d::IVector strongest_neighbour(num_rows, -1);\n typename Matrix_d::MVector diag(num_rows, 0);\n Vector > aggregated(num_rows, 0);\n IndexType *strongest_neighbour_ptr = strongest_neighbour.raw();\n ValueType *diag_ptr = diag.raw();\n IndexType *aggregates_ptr = aggregates.raw();\n const int threads_per_block = 256;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_rows - 1) / threads_per_block + 1 );\n getDiagonalKernelNoDiaProp <<< num_blocks, threads_per_block>>>(A_dia_ptr, A_values_ptr, num_rows, diag_ptr);\n cudaCheckError();\n int numUnassigned = num_rows;\n int numUnassigned_previous = numUnassigned;\n int icount = 0;\n\n do\n {\n // For each row, find the strongest neighbour who hasn't been assigned\n findStrongestNeighbourCsr <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_values_ptr, diag_ptr, num_rows, aggregates_ptr, strongest_neighbour_ptr);\n cudaCheckError();\n // Look for perfect matches. Also, for nodes without unaggregated neighbours, merge with aggregate containing strongest neighbour\n matchEdges <<< num_blocks, threads_per_block>>>(num_rows, aggregates_ptr, strongest_neighbour_ptr);\n cudaCheckError();\n numUnassigned_previous = numUnassigned;\n numUnassigned = (int)thrust_wrapper::count(aggregates.begin(), aggregates.begin() + num_rows, -1);\n cudaCheckError();\n icount++;\n }\n while (!(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_rows < this->numUnassigned_tol || numUnassigned_previous == numUnassigned));\n\n // Merge remaining vertices with current aggregates\n if (!this->deterministic)\n {\n while (numUnassigned != 0 )\n {\n mergeWithExistingAggregatesCsr <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_values_ptr,\n diag_ptr, num_rows, aggregates_ptr, this->deterministic, (IndexType *) NULL);\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(aggregates.begin(), aggregates.begin() + num_rows, -1);\n cudaCheckError();\n };\n }\n else\n {\n typename Matrix_d::IVector aggregates_candidate(num_rows, -1);\n\n while (numUnassigned != 0 )\n {\n mergeWithExistingAggregatesCsr <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_values_ptr,\n diag_ptr, num_rows, aggregates_ptr, this->deterministic, aggregates_candidate.raw());\n cudaCheckError();\n // Sync here\n joinExistingAggregates <<< num_blocks, threads_per_block>>>(num_rows, aggregates_ptr, aggregates_candidate.raw());\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(aggregates.begin(), aggregates.begin() + num_rows, -1);\n cudaCheckError();\n };\n\n aggregates_candidate.resize(0);\n }\n\n this->renumberAndCountAggregates(aggregates, aggregates_global, num_rows, num_aggregates);\n}\n#endif\n\n// setAggregates for block_dia_csr_matrix_d format\ntemplate \nvoid Size2Selector >::setAggregates_common_sqblocks(\n const Matrix_d &A, typename Matrix_d::IVector &aggregates, typename Matrix_d::IVector &aggregates_global, int &num_aggregates)\n{\n IndexType num_block_rows = (int) A.get_num_rows();\n IndexType num_nonzero_blocks = (int) A.get_num_nz();\n // both ways are supported\n IndexType total_nz = (A.is_matrix_singleGPU()) ? num_nonzero_blocks : A.manager->num_nz_all();\n typename Matrix_d::IVector row_indices(total_nz);\n cusp::detail::offsets_to_indices(A.row_offsets, row_indices);\n IndexType total_rows = (A.is_matrix_singleGPU()) ? A.get_num_rows() : A.manager->num_rows_all();\n aggregates.resize(total_rows);\n thrust_wrapper::fill(aggregates.begin(), aggregates.end(), -1);\n cudaCheckError();\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_row_indices_ptr = row_indices.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueType *A_nonzero_values_ptr = A.values.raw();\n typename Matrix_d::IVector strongest_neighbour(num_block_rows, -1);\n typename Matrix_d::IVector strongest_neighbour_1phase(num_block_rows, -1);\n Vector > aggregated(num_block_rows, 0);\n IndexType *strongest_neighbour_ptr = strongest_neighbour.raw();\n IndexType *strongest_neighbour_1phase_ptr = strongest_neighbour_1phase.raw();\n IndexType *aggregates_ptr = aggregates.raw();\n const int threads_per_block = 256;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_block_rows - 1) / threads_per_block + 1 );\n const int num_blocks_V2 = std::min( AMGX_GRID_MAX_SIZE, (num_nonzero_blocks - 1) / threads_per_block + 1);\n int numUnassigned = num_block_rows;\n int numUnassigned_previous = numUnassigned;\n Vector > edge_weights(num_nonzero_blocks, -1);\n float *edge_weights_ptr = edge_weights.raw();\n float *rand_edge_weights_ptr = NULL;//(this->two_phase ? rand_edge_weights.raw() : NULL);\n // Compute the edge weights\n cudaFuncSetCacheConfig(computeEdgeWeightsBlockDiaCsr_V2, cudaFuncCachePreferL1);\n computeEdgeWeightsBlockDiaCsr_V2 <<< num_blocks_V2, threads_per_block, 0, amgx::thrust::global_thread_handle::get_stream()>>>(A_row_offsets_ptr, A_row_indices_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, num_nonzero_blocks, edge_weights_ptr, rand_edge_weights_ptr, num_block_rows, A.get_block_dimy(), this->m_aggregation_edge_weight_component, this->weight_formula);\n cudaCheckError();\n cudaStream_t str = amgx::thrust::global_thread_handle::get_stream();\n#ifdef EXPERIMENTAL_ITERATIVE_MATCHING\n AsyncEvent *throttle_event = new AsyncEvent;\n throttle_event->create();\n // TODO: pinned host memory\n typename Matrix_h::IVector h_unagg_vec(1);\n typename Matrix_d::IVector d_unagg_vec(1);\n int *unaggregated = h_unagg_vec.raw();\n int *d_unaggregated = d_unagg_vec.raw();\n#endif\n int icount, s = 1;\n {\n icount = 0;\n float *weights_ptr = edge_weights_ptr;\n\n do\n {\n if ( !this->two_phase )\n {\n // 1-phase handshaking\n findStrongestNeighbourBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, weights_ptr, num_block_rows, aggregates_ptr, strongest_neighbour_ptr, strongest_neighbour_ptr, A.get_block_dimy(), 1, this->merge_singletons);\n cudaCheckError();\n }\n else\n {\n // 2-phase handshaking\n findStrongestNeighbourBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, weights_ptr, num_block_rows, aggregates_ptr, strongest_neighbour_1phase_ptr, strongest_neighbour_ptr, A.get_block_dimy(), 1, this->merge_singletons);\n cudaCheckError();\n // 2nd phase: for each block_row, find the strongest neighbour among those who gave hand on 1st phase\n findStrongestNeighbourBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, weights_ptr, num_block_rows, aggregates_ptr, strongest_neighbour_1phase_ptr, strongest_neighbour_ptr, A.get_block_dimy(), 2, this->merge_singletons);\n cudaCheckError();\n }\n\n // Look for perfect matches. Also, for nodes without unaggregated neighbours, merge with aggregate containing strongest neighbour\n matchEdges <<< num_blocks, threads_per_block, 0, str>>>(num_block_rows, aggregates_ptr, strongest_neighbour_ptr);\n cudaCheckError();\n#ifdef EXPERIMENTAL_ITERATIVE_MATCHING\n s = (icount & 1);\n\n if ( s == 0 )\n {\n // count unaggregated vertices\n cudaError_t cuda_rc = cudaMemsetAsync(d_unaggregated, 0, sizeof(int), str);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemsetAsync failed in size2_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n countAggregates <<< num_blocks, threads_per_block, 0, str>>>(num_block_rows, aggregates_ptr, d_unaggregated);\n cudaCheckError();\n cuda_rc = cudaMemcpyAsync(unaggregated, d_unaggregated, sizeof(int), cudaMemcpyDeviceToHost, str);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpyAsync failed in size2_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n throttle_event->record(str);\n }\n else\n {\n throttle_event->sync();\n numUnassigned_previous = numUnassigned;\n numUnassigned = *unaggregated;\n }\n\n#else\n cudaError_t cuda_rc = cudaStreamSynchronize(str);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaStreamSynchronize failed in size2_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n numUnassigned_previous = numUnassigned;\n numUnassigned = (int)thrust_wrapper::count(aggregates.begin(), aggregates.begin() + num_block_rows, -1);\n cudaCheckError();\n#endif\n icount++;\n }\n while ( (s == 0) || !(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_block_rows < this->numUnassigned_tol || numUnassigned == numUnassigned_previous));\n }\n#ifdef EXPERIMENTAL_ITERATIVE_MATCHING\n delete throttle_event;\n#endif\n\n if ( this->merge_singletons )\n {\n // Merge remaining vertices with current aggregates\n if (!this->deterministic)\n {\n while (numUnassigned != 0)\n {\n mergeWithExistingAggregatesBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, A.get_block_dimy(), this->deterministic, (IndexType *) NULL);\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(aggregates.begin(), aggregates.begin() + num_block_rows, -1);\n cudaCheckError();\n }\n }\n else\n {\n typename Matrix_d::IVector aggregates_candidate(num_block_rows, -1);\n\n while (numUnassigned != 0)\n {\n mergeWithExistingAggregatesBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, A.get_block_dimy(), this->deterministic, aggregates_candidate.raw());\n cudaCheckError();\n joinExistingAggregates <<< num_blocks, threads_per_block, 0, str>>>(num_block_rows, aggregates_ptr, aggregates_candidate.raw());\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(aggregates.begin(), aggregates.begin() + num_block_rows, -1);\n cudaCheckError();\n }\n\n aggregates_candidate.resize(0);\n }\n }\n else\n {\n //make singletons\n aggregateSingletons <<< num_blocks, threads_per_block, 0, str>>>( aggregates_ptr, num_block_rows );\n cudaCheckError();\n }\n\n this->renumberAndCountAggregates(aggregates, aggregates_global, num_block_rows, num_aggregates);\n}\n\ntemplate\nvoid Size2SelectorBase::setAggregates(Matrix &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n if (A.get_block_dimx() == A.get_block_dimy())\n {\n setAggregates_common_sqblocks( A, aggregates, aggregates_global, num_aggregates );\n }\n else\n {\n FatalError(\"Unsupported block size for Size2\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\n\n// -------------------------\n// Explict instantiations\n// -------------------------\n#define AMGX_CASE_LINE(CASE) template class Size2SelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class Size2Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n}\n}\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include //count\n#include //sort\n#include //lower_bound\n#include //unique\n#include //offsets_to_indices\n\n#define EXPERIMENTAL_ITERATIVE_MATCHING\n\nnamespace amgx\n{\nnamespace aggregation\n{\nnamespace size2_selector\n{\n\n// include common routines for all selectors\n#include \n\n// ------------------------\n// Kernels\n// ------------------------\n\n#ifndef DELETE\n// Kernel to compute the weight of the edges with block_dia_csr format\ntemplate \n__global__\nvoid computeEdgeWeightsBlockDiaCsr(const IndexType *row_offsets, const IndexType *column_indices,\n const ValueType *dia_values, const ValueType *nonzero_values, const IndexType num_block_rows, float *edge_weights, int bsize)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int jcol;\n int bsize_sq = bsize * bsize;\n\n while (tid < num_block_rows)\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (jcol >= num_block_rows) { continue; }\n\n // Compute edge weight\n for (int k = row_offsets[jcol]; k < row_offsets[jcol + 1]; k++)\n {\n if (column_indices[k] == tid)\n {\n edge_weights[j] = (float) 0.5 * (fabs(nonzero_values[j * bsize_sq]) + fabs(nonzero_values[k * bsize_sq]))\n / max( fabs(dia_values[tid * bsize_sq]), fabs(dia_values[jcol * bsize_sq]));\n break;\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n#endif\n\n// findStrongestNeighbour kernel for csr_matrix format\n// computes weight on the fly\ntemplate \n__global__\nvoid findStrongestNeighbourCsr(const IndexType *row_offsets, const IndexType *column_indices,\n const ValueType *values, const ValueType *diag, const IndexType num_rows, IndexType *aggregates, int *strongest_neighbour)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n ValueType weight;\n int jcol;\n\n while (tid < num_rows)\n {\n ValueType max_weight_unaggregated = 0.;\n ValueType max_weight_aggregated = 0.;\n int strongest_unaggregated = -1;\n int strongest_aggregated = -1;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (tid == jcol || jcol >= num_rows) { continue; }\n\n // Compute edge weight\n //weight = fabs(values[j])/max( fabs(diag[tid]),fabs(diag[jcol]));\n for (int k = row_offsets[jcol]; k < row_offsets[jcol + 1]; k++)\n {\n if (column_indices[k] == tid)\n {\n weight = 0.5 * (fabs(values[j]) + fabs(values[k])) / max( fabs(diag[tid]), fabs(diag[jcol]));\n break;\n }\n }\n\n // Identify strongest aggregated and unaggregated neighbours\n if (aggregates[jcol] == -1 && (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && jcol > strongest_unaggregated))) // unaggregated\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n else if (aggregates[jcol] != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) // aggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_unaggregated == -1 && strongest_aggregated != -1) // All neighbours are aggregated\n {\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n else if (strongest_unaggregated != -1)\n {\n strongest_neighbour[tid] = strongest_unaggregated;\n }\n else\n {\n strongest_neighbour[tid] = tid;\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// findStrongestNeighbour kernel for block_dia_csr_matrix format\n// computes weight on the fly\ntemplate \n__global__\nvoid findStrongestNeighbourBlockDiaCsr(const IndexType *row_offsets, const IndexType *column_indices,\n const ValueType *dia_values, const ValueType *nonzero_values, const IndexType num_block_rows, IndexType *aggregates, int *strongest_neighbour, int bsize)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n ValueType weight;\n int jcol;\n int bsize_sq = bsize * bsize;\n\n while (tid < num_block_rows)\n {\n ValueType max_weight_unaggregated = 0.;\n ValueType max_weight_aggregated = 0.;\n int strongest_unaggregated = -1;\n int strongest_aggregated = -1;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (jcol >= num_block_rows) { continue; }\n\n // Compute edge weight\n for (int k = row_offsets[jcol]; k < row_offsets[jcol + 1]; k++)\n {\n if (column_indices[k] == tid)\n {\n weight = 0.5 * (fabs(nonzero_values[j * bsize_sq]) + fabs(nonzero_values[k * bsize_sq]))\n / max( fabs(dia_values[tid * bsize_sq]), fabs(dia_values[jcol * bsize_sq]));\n break;\n }\n }\n\n // Identify strongest aggregated and unaggregated neighbours\n if (aggregates[jcol] == -1 && (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && jcol > strongest_unaggregated))) // unaggregated\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n else if (aggregates[jcol] != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) // aggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_unaggregated == -1 && strongest_aggregated != -1) // All neighbours are aggregated\n // Put in same aggregate as strongest neighbour\n {\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n else if (strongest_unaggregated != -1)\n {\n strongest_neighbour[tid] = strongest_unaggregated;\n }\n else\n {\n strongest_neighbour[tid] = tid;\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\n__device__\nfloat random_weight2(int i, int j)\n{\n#define RAND_MULTIPLIER 1145637293\n unsigned long i_min = (min(i, j) * RAND_MULTIPLIER);\n unsigned long i_max = (max(i, j) * RAND_MULTIPLIER);\n return ((float)i_min / i_max);\n}\n\n\n// findStrongestNeighbour kernel for block_dia_csr_matrix format\n// Reads the weight from edge_weights array\ntemplate \n__global__\nvoid findStrongestNeighbourBlockDiaCsr_V2(const IndexType *row_offsets, const IndexType *column_indices,\n float *edge_weights, IndexType num_block_rows, IndexType *aggregates,\n IndexType *strongest_neighbour_1phase, IndexType *strongest_neighbour,\n const size_t bsize, int phase, bool merge_singletons)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n float weight;\n int jcol;\n\n while (tid < num_block_rows)\n {\n int strongest_unaggregated = -1;\n int strongest_aggregated = -1;\n float max_weight_unaggregated = 0.;\n float max_weight_aggregated = 0.;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (phase == 1) { weight = edge_weights[j]; }\n else { weight = random_weight2(tid, jcol); }\n\n if (tid == jcol || jcol >= num_block_rows) { continue; } // skip diagonal and halo\n\n if (phase == 2 && strongest_neighbour_1phase[jcol] != tid) { continue; } // if 2nd phase only accept those who gave a hand on the 1st phase\n\n // Identify strongest aggregated and unaggregated neighbours\n if (aggregates[jcol] == -1 && (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && jcol > strongest_unaggregated))) // unaggregated\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n else if (aggregates[jcol] != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) // aggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_unaggregated == -1 && strongest_aggregated != -1) // All neighbours are aggregated\n {\n if ( merge_singletons )\n // Put in same aggregate as strongest neighbour\n {\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n else\n {\n aggregates[tid] = tid;\n }\n }\n else if (strongest_unaggregated != -1)\n {\n if (phase == 2)\n {\n float rand_w1 = random_weight2(tid, strongest_neighbour_1phase[tid]);\n strongest_neighbour[tid] = max_weight_unaggregated > rand_w1 ? strongest_unaggregated : strongest_neighbour_1phase[tid];\n }\n else { strongest_neighbour_1phase[tid] = strongest_unaggregated; }\n }\n else\n {\n if (phase == 2) { strongest_neighbour[tid] = strongest_neighbour_1phase[tid]; }\n else { strongest_neighbour_1phase[tid] = tid; }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel that checks if perfect matchs exist\ntemplate \n__global__\nvoid matchEdges(const IndexType num_rows, IndexType *aggregates, int *strongest_neighbour)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int potential_match, potential_match_neighbour;\n\n while (tid < num_rows)\n {\n if (aggregates[tid] == -1) // Unaggregated row\n {\n potential_match = strongest_neighbour[tid];\n potential_match_neighbour = strongest_neighbour[potential_match];\n\n if (potential_match != -1 && potential_match_neighbour == tid) // we have a match\n {\n aggregates[tid] = ( potential_match > tid ) ? tid : potential_match;\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \n__global__\nvoid countAggregates(const IndexType num_rows, IndexType *aggregates, int *num_unaggregated)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int c = 0;\n int i = tid;\n\n while ( i < num_rows )\n {\n c += ( aggregates[i] == -1 );\n i += gridDim.x * blockDim.x;\n }\n\n __shared__ int smem[block_size];\n smem[threadIdx.x] = c;\n __syncthreads();\n\n for ( int off = blockDim.x / 2; off >= 1; off = off / 2 )\n {\n if ( threadIdx.x < off )\n {\n smem[threadIdx.x] += smem[threadIdx.x + off];\n }\n __syncthreads();\n }\n\n if ( threadIdx.x == 0 )\n {\n atomicAdd(num_unaggregated, smem[0]);\n }\n}\n\n// Kernel that merges unaggregated vertices its strongest aggregated neighbour\n// For csr_matrix_format\ntemplate \n__global__\nvoid mergeWithExistingAggregatesCsr(const IndexType *row_offsets, const IndexType *column_indices, const ValueType *values,\n const ValueType *diag, const int num_rows, IndexType *aggregates, int deterministic, IndexType *aggregates_candidate)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int jcol;\n ValueType weight;\n\n while (tid < num_rows)\n {\n ValueType max_weight_aggregated = 0.;\n int strongest_aggregated = -1;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (tid == jcol || jcol >= num_rows) { continue; }\n\n // Compute edge weight\n weight = fabs(values[j]) / max( fabs(diag[tid]), fabs(diag[jcol]));\n\n // Identify strongest aggregated neighbour\n if (aggregates[jcol] != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) // aggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_aggregated != -1) // Found a neighbour to aggregate to\n {\n if (deterministic)\n {\n aggregates_candidate[tid] = aggregates[strongest_aggregated];\n }\n else\n {\n // Put in same aggregate as strongest neighbour\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n }\n else // All neighbours are unaggregated, leave alone\n {\n if (deterministic)\n {\n aggregates_candidate[tid] = tid;\n }\n else\n {\n aggregates[tid] = tid;\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate \n__global__\nvoid joinExistingAggregates(IndexType num_rows, IndexType *aggregates, IndexType *aggregates_candidate)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while (tid < num_rows)\n {\n if (aggregates[tid] == -1 && aggregates_candidate[tid] != -1) // Unaggregated row\n {\n aggregates[tid] = aggregates_candidate[tid];\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel that merges unaggregated vertices its strongest aggregated neighbour\n// Edge weights are computed on the fly\n// For block_dia_csr_matrix_format\ntemplate \n__global__\nvoid mergeWithExistingAggregatesBlockDiaCsr(const IndexType *row_offsets, const IndexType *column_indices, const ValueType *dia_values, const ValueType *nonzero_values,\n const int num_block_rows, IndexType *aggregates, int bsize, int deterministic, IndexType *aggregates_candidate)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int jcol;\n ValueType weight;\n int bsize_sq = bsize * bsize;\n\n while (tid < num_block_rows)\n {\n int strongest_aggregated = -1;\n ValueType max_weight_aggregated = 0.;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (jcol >= num_block_rows) { continue; }\n\n // Compute edge weight\n weight = fabs(nonzero_values[j * bsize_sq]) / max( fabs(dia_values[tid * bsize_sq]), fabs(dia_values[jcol * bsize_sq]));\n\n // Identify strongest aggregated neighbour\n if (aggregates[jcol] != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) // aggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_aggregated != -1) // Found a neighbour to aggregate to\n {\n if (deterministic)\n {\n aggregates_candidate[tid] = aggregates[strongest_aggregated];\n }\n else\n {\n // Put in same aggregate as strongest neighbour\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n }\n else // All neighbours are unaggregated, leave alone\n {\n if (deterministic)\n {\n aggregates_candidate[tid] = tid;\n }\n else\n {\n aggregates[tid] = tid;\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel that merges unaggregated vertices its strongest aggregated neighbour\n// Weights are read from edge_weights array\n// For block_dia_csr_matrix_format\ntemplate \n__global__\nvoid mergeWithExistingAggregatesBlockDiaCsr_V2(const IndexType *row_offsets, const IndexType *column_indices, const float *edge_weights,\n const int num_block_rows, IndexType *aggregates, int bsize, const int deterministic, IndexType *aggregates_candidate)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int jcol;\n float weight;\n\n while (tid < num_block_rows)\n {\n float max_weight_aggregated = 0.;\n int strongest_aggregated = -1;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n // Compute edge weight\n weight = edge_weights[j];\n jcol = column_indices[j];\n\n if (jcol == tid || jcol >= num_block_rows) { continue; } // skip diagonal\n\n // Identify strongest aggregated neighbour\n if (aggregates[jcol] != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) //\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_aggregated != -1) // Found a neighbour to aggregate to\n {\n if (deterministic)\n {\n aggregates_candidate[tid] = aggregates[strongest_aggregated];\n }\n else\n {\n // Put in same aggregate as strongest neighbour\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n }\n else // All neighbours are unaggregated, leave alone\n {\n if (deterministic)\n {\n aggregates_candidate[tid] = tid;\n }\n else\n {\n aggregates[tid] = tid;\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel to extract diagonal for csr_matrix format\ntemplate \n__global__\nvoid getDiagonalKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, const IndexType numRows, ValueType *diagonal)\n{\n int tIdx = threadIdx.x + blockDim.x * blockIdx.x;\n\n while (tIdx < numRows)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n for (int j = offset; j < offset + numj; j++)\n {\n int jcol = column_indices[j];\n\n if (tIdx == jcol)\n {\n diagonal[tIdx] = values[j];\n }\n }\n\n tIdx += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel to extract diagonal for csr_matrix format\ntemplate \n__global__\nvoid getDiagonalKernelNoDiaProp(const IndexType *dia_idx, const ValueType *values, const IndexType numRows, ValueType *diagonal)\n{\n int tIdx = threadIdx.x + blockDim.x * blockIdx.x;\n\n while (tIdx < numRows)\n {\n diagonal[tIdx] = values[dia_idx[tIdx]];\n tIdx += gridDim.x * blockDim.x;\n }\n}\n\n// -----------------\n// Methods\n// ----------------\n\n// Constructor\ntemplate\nSize2SelectorBase::Size2SelectorBase(AMG_Config &cfg, const std::string &cfg_scope)\n{\n deterministic = cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\");\n max_iterations = cfg.AMG_Config::template getParameter(\"max_matching_iterations\", cfg_scope);\n numUnassigned_tol = cfg.AMG_Config::template getParameter(\"max_unassigned_percentage\", cfg_scope);\n two_phase = cfg.AMG_Config::template getParameter(\"handshaking_phases\", cfg_scope) == 2;\n m_aggregation_edge_weight_component = cfg.AMG_Config::template getParameter(\"aggregation_edge_weight_component\", cfg_scope);\n merge_singletons = cfg.AMG_Config::template getParameter(\"merge_singletons\", cfg_scope) == 1;\n weight_formula = cfg.AMG_Config::template getParameter(\"weight_formula\", cfg_scope);\n}\n\n// setAggregates for csr_matrix_h format\ntemplate \nvoid Size2Selector >::setAggregates_1x1(const Matrix_h &A,\n typename Matrix_h::IVector &aggregates, typename Matrix_h::IVector &aggregates_global, int &num_aggregates)\n{\n FatalError(\"Size2 selector: setAggregates not implemented on CPU, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n// setAggregates for block_dia_csr_matrix_h format\ntemplate \nvoid Size2Selector >::setAggregates_common_sqblocks(const Matrix_h &A,\n typename Matrix_h::IVector &aggregates, typename Matrix_h::IVector &aggregates_global, int &num_aggregates)\n{\n FatalError(\"Size2 selector: setAggregates not implemented on CPU, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n#ifndef DELETE\n// setAggregates for csr_matrix_d format\ntemplate \nvoid Size2Selector >::setAggregates_1x1(\n const Matrix_d &A, typename Matrix_d::IVector &aggregates, typename Matrix_d::IVector &aggregates_global, int &num_aggregates)\n\n{\n if (!A.is_matrix_singleGPU())\n {\n aggregates.resize(A.manager->halo_offset(A.manager->num_neighbors()));\n }\n else\n {\n aggregates.resize(A.get_num_rows());\n }\n\n thrust_wrapper::fill(aggregates.begin(), aggregates.end(), -1);\n cudaCheckError();\n //typedef typename csr_matrix_d::index_type IndexType;\n //typedef typename csr_matrix_d::value_type ValueType;\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const ValueType *A_values_ptr = A.values.raw();\n const IndexType *A_dia_ptr = A.diag.raw();\n IndexType num_rows = (int)A.get_num_rows();\n typename Matrix_d::IVector strongest_neighbour(num_rows, -1);\n typename Matrix_d::MVector diag(num_rows, 0);\n Vector > aggregated(num_rows, 0);\n IndexType *strongest_neighbour_ptr = strongest_neighbour.raw();\n ValueType *diag_ptr = diag.raw();\n IndexType *aggregates_ptr = aggregates.raw();\n const int threads_per_block = 256;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_rows - 1) / threads_per_block + 1 );\n getDiagonalKernelNoDiaProp <<< num_blocks, threads_per_block>>>(A_dia_ptr, A_values_ptr, num_rows, diag_ptr);\n cudaCheckError();\n int numUnassigned = num_rows;\n int numUnassigned_previous = numUnassigned;\n int icount = 0;\n\n do\n {\n // For each row, find the strongest neighbour who hasn't been assigned\n findStrongestNeighbourCsr <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_values_ptr, diag_ptr, num_rows, aggregates_ptr, strongest_neighbour_ptr);\n cudaCheckError();\n // Look for perfect matches. Also, for nodes without unaggregated neighbours, merge with aggregate containing strongest neighbour\n matchEdges <<< num_blocks, threads_per_block>>>(num_rows, aggregates_ptr, strongest_neighbour_ptr);\n cudaCheckError();\n numUnassigned_previous = numUnassigned;\n numUnassigned = (int)thrust_wrapper::count(aggregates.begin(), aggregates.begin() + num_rows, -1);\n cudaCheckError();\n icount++;\n }\n while (!(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_rows < this->numUnassigned_tol || numUnassigned_previous == numUnassigned));\n\n // Merge remaining vertices with current aggregates\n if (!this->deterministic)\n {\n while (numUnassigned != 0 )\n {\n mergeWithExistingAggregatesCsr <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_values_ptr,\n diag_ptr, num_rows, aggregates_ptr, this->deterministic, (IndexType *) NULL);\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(aggregates.begin(), aggregates.begin() + num_rows, -1);\n cudaCheckError();\n };\n }\n else\n {\n typename Matrix_d::IVector aggregates_candidate(num_rows, -1);\n\n while (numUnassigned != 0 )\n {\n mergeWithExistingAggregatesCsr <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_values_ptr,\n diag_ptr, num_rows, aggregates_ptr, this->deterministic, aggregates_candidate.raw());\n cudaCheckError();\n // Sync here\n joinExistingAggregates <<< num_blocks, threads_per_block>>>(num_rows, aggregates_ptr, aggregates_candidate.raw());\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(aggregates.begin(), aggregates.begin() + num_rows, -1);\n cudaCheckError();\n };\n\n aggregates_candidate.resize(0);\n }\n\n this->renumberAndCountAggregates(aggregates, aggregates_global, num_rows, num_aggregates);\n}\n#endif\n\n// setAggregates for block_dia_csr_matrix_d format\ntemplate \nvoid Size2Selector >::setAggregates_common_sqblocks(\n const Matrix_d &A, typename Matrix_d::IVector &aggregates, typename Matrix_d::IVector &aggregates_global, int &num_aggregates)\n{\n IndexType num_block_rows = (int) A.get_num_rows();\n IndexType num_nonzero_blocks = (int) A.get_num_nz();\n // both ways are supported\n IndexType total_nz = (A.is_matrix_singleGPU()) ? num_nonzero_blocks : A.manager->num_nz_all();\n typename Matrix_d::IVector row_indices(total_nz);\n cusp::detail::offsets_to_indices(A.row_offsets, row_indices);\n IndexType total_rows = (A.is_matrix_singleGPU()) ? A.get_num_rows() : A.manager->num_rows_all();\n aggregates.resize(total_rows);\n thrust_wrapper::fill(aggregates.begin(), aggregates.end(), -1);\n cudaCheckError();\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_row_indices_ptr = row_indices.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueType *A_nonzero_values_ptr = A.values.raw();\n typename Matrix_d::IVector strongest_neighbour(num_block_rows, -1);\n typename Matrix_d::IVector strongest_neighbour_1phase(num_block_rows, -1);\n Vector > aggregated(num_block_rows, 0);\n IndexType *strongest_neighbour_ptr = strongest_neighbour.raw();\n IndexType *strongest_neighbour_1phase_ptr = strongest_neighbour_1phase.raw();\n IndexType *aggregates_ptr = aggregates.raw();\n const int threads_per_block = 256;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_block_rows - 1) / threads_per_block + 1 );\n const int num_blocks_V2 = std::min( AMGX_GRID_MAX_SIZE, (num_nonzero_blocks - 1) / threads_per_block + 1);\n int numUnassigned = num_block_rows;\n int numUnassigned_previous = numUnassigned;\n Vector > edge_weights(num_nonzero_blocks, -1);\n float *edge_weights_ptr = edge_weights.raw();\n float *rand_edge_weights_ptr = NULL;//(this->two_phase ? rand_edge_weights.raw() : NULL);\n // Compute the edge weights\n hipFuncSetCacheConfig(reinterpret_cast(computeEdgeWeightsBlockDiaCsr_V2), hipFuncCachePreferL1);\n computeEdgeWeightsBlockDiaCsr_V2 <<< num_blocks_V2, threads_per_block, 0, amgx::thrust::global_thread_handle::get_stream()>>>(A_row_offsets_ptr, A_row_indices_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, num_nonzero_blocks, edge_weights_ptr, rand_edge_weights_ptr, num_block_rows, A.get_block_dimy(), this->m_aggregation_edge_weight_component, this->weight_formula);\n cudaCheckError();\n hipStream_t str = amgx::thrust::global_thread_handle::get_stream();\n#ifdef EXPERIMENTAL_ITERATIVE_MATCHING\n AsyncEvent *throttle_event = new AsyncEvent;\n throttle_event->create();\n // TODO: pinned host memory\n typename Matrix_h::IVector h_unagg_vec(1);\n typename Matrix_d::IVector d_unagg_vec(1);\n int *unaggregated = h_unagg_vec.raw();\n int *d_unaggregated = d_unagg_vec.raw();\n#endif\n int icount, s = 1;\n {\n icount = 0;\n float *weights_ptr = edge_weights_ptr;\n\n do\n {\n if ( !this->two_phase )\n {\n // 1-phase handshaking\n findStrongestNeighbourBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, weights_ptr, num_block_rows, aggregates_ptr, strongest_neighbour_ptr, strongest_neighbour_ptr, A.get_block_dimy(), 1, this->merge_singletons);\n cudaCheckError();\n }\n else\n {\n // 2-phase handshaking\n findStrongestNeighbourBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, weights_ptr, num_block_rows, aggregates_ptr, strongest_neighbour_1phase_ptr, strongest_neighbour_ptr, A.get_block_dimy(), 1, this->merge_singletons);\n cudaCheckError();\n // 2nd phase: for each block_row, find the strongest neighbour among those who gave hand on 1st phase\n findStrongestNeighbourBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, weights_ptr, num_block_rows, aggregates_ptr, strongest_neighbour_1phase_ptr, strongest_neighbour_ptr, A.get_block_dimy(), 2, this->merge_singletons);\n cudaCheckError();\n }\n\n // Look for perfect matches. Also, for nodes without unaggregated neighbours, merge with aggregate containing strongest neighbour\n matchEdges <<< num_blocks, threads_per_block, 0, str>>>(num_block_rows, aggregates_ptr, strongest_neighbour_ptr);\n cudaCheckError();\n#ifdef EXPERIMENTAL_ITERATIVE_MATCHING\n s = (icount & 1);\n\n if ( s == 0 )\n {\n // count unaggregated vertices\n hipError_t cuda_rc = hipMemsetAsync(d_unaggregated, 0, sizeof(int), str);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemsetAsync failed in size2_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n countAggregates <<< num_blocks, threads_per_block, 0, str>>>(num_block_rows, aggregates_ptr, d_unaggregated);\n cudaCheckError();\n cuda_rc = hipMemcpyAsync(unaggregated, d_unaggregated, sizeof(int), hipMemcpyDeviceToHost, str);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpyAsync failed in size2_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n throttle_event->record(str);\n }\n else\n {\n throttle_event->sync();\n numUnassigned_previous = numUnassigned;\n numUnassigned = *unaggregated;\n }\n\n#else\n hipError_t cuda_rc = hipStreamSynchronize(str);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipStreamSynchronize failed in size2_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n numUnassigned_previous = numUnassigned;\n numUnassigned = (int)thrust_wrapper::count(aggregates.begin(), aggregates.begin() + num_block_rows, -1);\n cudaCheckError();\n#endif\n icount++;\n }\n while ( (s == 0) || !(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_block_rows < this->numUnassigned_tol || numUnassigned == numUnassigned_previous));\n }\n#ifdef EXPERIMENTAL_ITERATIVE_MATCHING\n delete throttle_event;\n#endif\n\n if ( this->merge_singletons )\n {\n // Merge remaining vertices with current aggregates\n if (!this->deterministic)\n {\n while (numUnassigned != 0)\n {\n mergeWithExistingAggregatesBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, A.get_block_dimy(), this->deterministic, (IndexType *) NULL);\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(aggregates.begin(), aggregates.begin() + num_block_rows, -1);\n cudaCheckError();\n }\n }\n else\n {\n typename Matrix_d::IVector aggregates_candidate(num_block_rows, -1);\n\n while (numUnassigned != 0)\n {\n mergeWithExistingAggregatesBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, A.get_block_dimy(), this->deterministic, aggregates_candidate.raw());\n cudaCheckError();\n joinExistingAggregates <<< num_blocks, threads_per_block, 0, str>>>(num_block_rows, aggregates_ptr, aggregates_candidate.raw());\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(aggregates.begin(), aggregates.begin() + num_block_rows, -1);\n cudaCheckError();\n }\n\n aggregates_candidate.resize(0);\n }\n }\n else\n {\n //make singletons\n aggregateSingletons <<< num_blocks, threads_per_block, 0, str>>>( aggregates_ptr, num_block_rows );\n cudaCheckError();\n }\n\n this->renumberAndCountAggregates(aggregates, aggregates_global, num_block_rows, num_aggregates);\n}\n\ntemplate\nvoid Size2SelectorBase::setAggregates(Matrix &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n if (A.get_block_dimx() == A.get_block_dimy())\n {\n setAggregates_common_sqblocks( A, aggregates, aggregates_global, num_aggregates );\n }\n else\n {\n FatalError(\"Unsupported block size for Size2\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\n\n// -------------------------\n// Explict instantiations\n// -------------------------\n#define AMGX_CASE_LINE(CASE) template class Size2SelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class Size2Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n}\n}\n", "source": "/kaggle/working/amgx/src/aggregation/selectors/size2_selector.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n\nnamespace amgx\n{\nnamespace aggregation\n{\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n__host__ __device__ int hash_function(int a, int seed = 17)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int NUM_THREADS_PER_ROW, int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid compute_edge_weights_diag( const int A_num_rows,\n const int block_size,\n const int num_owned,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict A_diag,\n const Value_type *__restrict A_vals,\n Value_type *__restrict A_edge_weights )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Shared memory to broadcast columns.\n __shared__ volatile int s_bcast_col[CTA_SIZE];\n // Shared memory to store columns ji.\n __shared__ volatile Value_type s_bcast_ji[CTA_SIZE];\n // Constants.\n // const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n // const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += gridDim.x * NUM_WARPS_PER_CTA )\n {\n // The diagonal of (i,i).\n Value_type row_diag = A_vals[block_size * A_diag[a_row_id]];\n // Load the range of the row.\n int a_col_begin = A_rows[a_row_id + 0];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Iterate over the columns of A.\n for ( ; a_col_begin < a_col_end ; a_col_begin += WARP_SIZE )\n {\n // Each lane works on a different column.\n const int a_col_it = a_col_begin + lane_id;\n // Is it an active thread.\n int is_active = a_col_it < a_col_end;\n // Active nodes load column ids.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n s_bcast_col[threadIdx.x] = a_col_id;\n // Reset ji coefficients.\n s_bcast_ji[threadIdx.x] = Value_type(0);\n\n // Iterate over columns to find ji coefficients if they exist.\n for ( int k = 0, num_cols = __popc(utils::ballot(is_active)) ; k < num_cols ; ++k )\n {\n const int uniform_col_id = s_bcast_col[warp_id * WARP_SIZE + k];\n\n // Skip columns whose id is greater than num_owned.\n if ( uniform_col_id == -1 || uniform_col_id == a_row_id || uniform_col_id >= num_owned )\n {\n continue;\n }\n\n // The bounds of the row.\n int b_row_begin = A_rows[uniform_col_id + 0];\n int b_row_end = A_rows[uniform_col_id + 1];\n\n // Iterate over the row.\n for ( int not_found = 1 ; not_found && b_row_begin < b_row_end ; b_row_begin += WARP_SIZE )\n {\n const int b_row_it = b_row_begin + lane_id;\n // Load the column id.\n int b_col_id = -1;\n\n if ( b_row_it < b_row_end )\n {\n b_col_id = A_cols[b_row_it];\n }\n\n // Has anyone found the column.\n not_found = utils::all( b_col_id != a_row_id );\n\n // If someone found the column, ask it to load the value.\n if ( b_col_id == a_row_id )\n {\n s_bcast_ji[warp_id * WARP_SIZE + k] = A_vals[block_size * b_row_it];\n }\n }\n }\n\n // Deactivate \"invalid\" threads.\n is_active = is_active && a_col_id < num_owned;\n // The diagonal value associated with the column.\n Value_type col_diag(0);\n\n if ( is_active )\n {\n col_diag = A_vals[block_size * A_diag[a_col_id]];\n }\n\n // The value of the column.\n Value_type col_val(0);\n\n if ( is_active )\n {\n col_val = A_vals[block_size * a_col_it];\n }\n\n // Compute the denominator.\n Value_type den = max( abs(row_diag), abs(col_diag) );\n // Compute the weight of the edge.\n Value_type weight(0);\n\n if ( den != Value_type(0) )\n {\n weight = Value_type(0.5) * (abs(col_val) + abs(s_bcast_ji[threadIdx.x])) / den;\n }\n\n if ( is_active )\n {\n A_edge_weights[a_col_it] = weight;\n }\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid compute_ring_leader( const int A_num_rows,\n const int num_owned,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict is_aggregated,\n const int *__restrict in_leader_id,\n const int *__restrict in_leader_hash,\n int *__restrict out_leader_id,\n int *__restrict out_leader_hash )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += gridDim.x * NUM_WARPS_PER_CTA )\n {\n // Skip already aggregated vertices.\n if ( is_aggregated[a_row_id] )\n {\n continue;\n }\n\n // Load the range of the row.\n int a_col_begin = A_rows[a_row_id + 0];\n int a_col_end = A_rows[a_row_id + 1];\n // The max hash in my ith ring.\n int my_min_id = A_num_rows;\n int my_max_hash = INT_MIN;\n\n if ( in_leader_id )\n {\n my_min_id = in_leader_id [a_row_id];\n my_max_hash = in_leader_hash[a_row_id];\n }\n else\n {\n my_min_id = a_row_id;\n my_max_hash = hash_function(a_row_id);\n }\n\n // Iterate over the columns of A.\n for ( ; a_col_begin < a_col_end ; a_col_begin += WARP_SIZE )\n {\n // Each lane works on a different column.\n const int a_col_it = a_col_begin + lane_id;\n // Is it an active thread.\n int is_active = a_col_it < a_col_end;\n // Active nodes load column ids.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Compute the hash value if needed.\n int col_min_id = A_num_rows;\n int col_max_hash = INT_MIN;\n\n if ( in_leader_id )\n {\n if ( is_active && !is_aggregated[a_col_id] )\n {\n col_min_id = in_leader_id [a_col_id];\n col_max_hash = in_leader_hash[a_col_id];\n }\n }\n else\n {\n if ( is_active && !is_aggregated[a_col_id] )\n {\n col_min_id = a_col_id;\n col_max_hash = hash_function(a_col_id);\n }\n }\n\n // Update the max_hash if needed.\n if ( is_active && (my_max_hash < col_max_hash || (my_max_hash == col_max_hash && my_min_id >= col_min_id)) )\n {\n my_min_id = col_min_id;\n my_max_hash = col_max_hash;\n }\n }\n\n // Reduce the max hash.\n#pragma unroll\n\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n int other_min_id = utils::shfl_xor(my_min_id, mask);\n int other_max_hash = utils::shfl_xor(my_max_hash, mask);\n\n if ( other_max_hash > my_max_hash || (other_max_hash == my_max_hash && other_min_id < my_min_id) )\n {\n my_min_id = other_min_id;\n my_max_hash = other_max_hash;\n }\n }\n\n // The 1st thread stores the result.\n if ( lane_id == 0 )\n {\n out_leader_id [a_row_id] = my_min_id;\n out_leader_hash[a_row_id] = my_max_hash;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid build_aggregates( const int num_rings,\n const int A_num_rows,\n const int num_owned,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n int *__restrict is_aggregated,\n const int *__restrict in_leader_id,\n const int *__restrict in_leader_hash,\n const Value_type *__restrict edge_weights,\n int *__restrict num_aggregates,\n int *__restrict aggregates,\n int *__restrict num_unaggregated)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += gridDim.x * NUM_WARPS_PER_CTA )\n {\n // Skip already aggregated vertices.\n if ( is_aggregated[a_row_id] )\n {\n continue;\n }\n\n // The max hash in my ith ring.\n int my_hash = hash_function(a_row_id), my_max_hash = in_leader_hash[a_row_id];\n\n // Skip if I'm not a local king.\n if ( my_hash < my_max_hash || (my_hash == my_max_hash && a_row_id > in_leader_id[a_row_id]) )\n {\n continue;\n }\n\n // We start at the row.\n int curr_row = a_row_id;\n\n // Iterate to form the aggregate.\n for ( int aggregate_size = 1, aggregate_id = -1 ; aggregate_size <= num_rings ; ++aggregate_size )\n {\n // Load the range of the row.\n int a_col_begin = A_rows[curr_row + 0];\n int a_col_end = A_rows[curr_row + 1];\n // The max weight.\n int max_id = -1;\n float max_weight = -1.0f;\n // The max aggregated weight.\n int max_aggregated_id = -1;\n float max_aggregated_weight = -1.0f;\n\n // Iterate over the columns of A.\n for ( ; a_col_begin < a_col_end ; a_col_begin += WARP_SIZE )\n {\n // Each lane works on a different column.\n const int a_col_it = a_col_begin + lane_id;\n // Is it an active thread.\n int is_active = a_col_it < a_col_end;\n // Active nodes load column ids.\n int a_col_id = -1;\n float a_col_weight = -1.0f;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n a_col_weight = static_cast(edge_weights[a_col_it]);\n }\n\n // Is the column aggregated?\n int is_col_aggregated = is_active && is_aggregated[a_col_id];\n\n // Select the column if it's not aggregated and its weight is better.\n if ( is_active && !is_col_aggregated && a_col_weight > max_weight )\n {\n max_id = a_col_id;\n max_weight = a_col_weight;\n }\n\n // Update max aggregated weight if needed.\n if ( is_active && is_col_aggregated && a_col_weight > max_aggregated_weight )\n {\n max_aggregated_id = a_col_id;\n max_aggregated_weight = a_col_weight;\n }\n }\n\n // Is there a valid max_id.\n int valid_max_id = utils::any(max_id != -1, utils::activemask());\n\n // No valid max id?\n if ( !valid_max_id && aggregate_size > 1 )\n {\n break;\n }\n\n // Find the max id.\n if ( !valid_max_id )\n {\n max_id = max_aggregated_id;\n max_weight = max_aggregated_weight;\n }\n\n#pragma unroll\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n int other_max_id = utils::shfl_xor(max_id, mask);\n float other_max_weight = utils::shfl_xor(max_weight, mask);\n\n if ( other_max_weight > max_weight )\n {\n max_id = other_max_id;\n max_weight = other_max_weight;\n }\n }\n\n // We know it's a singleton so merge with an existing aggregate.\n if ( !valid_max_id )\n {\n if ( lane_id == 0 && max_id != -1 )\n {\n is_aggregated[curr_row] = 1;\n aggregates[curr_row] = aggregates[max_id];\n atomicAdd(num_unaggregated, -1);\n }\n\n break;\n }\n\n // It's not a singleton but the 1st vertex in the aggregate.\n if ( lane_id == 0 && aggregate_size == 1 )\n {\n aggregate_id = atomicAdd(num_aggregates, 1);\n is_aggregated[a_row_id] = 1;\n aggregates[a_row_id] = aggregate_id;\n atomicAdd(num_unaggregated, -1);\n }\n\n // Set the aggregate of the winner.\n if ( lane_id == 0 )\n {\n aggregates[max_id] = aggregate_id;\n is_aggregated[max_id] = 1;\n atomicAdd(num_unaggregated, -1);\n }\n\n // Set the next row to consider.\n curr_row = max_id;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nenum { WARP_SIZE = 32 };\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nParallelGreedySelector >::setAggregates_1x1( const MatrixType &A, IVector &aggregates, IVector &aggregates_global, int &num_aggregates )\n{}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nParallelGreedySelector >::setAggregates_common_sqblocks( const MatrixType &A, IVector &aggregates, IVector &aggregates_global, int &num_aggregates )\n{\n const int num_rings = 4;\n // The number of rows of A.\n const int num_rows = A.get_num_rows();\n // The size of the block.\n const int block_size = A.get_block_dimx() * A.get_block_dimy();\n // The number of threads per CTA.\n const int CTA_SIZE = 128;\n // The number of warps per CTA.\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // The number of CTAs in a grid where each thread is independent.\n const int THREAD_GRID_SIZE = (A.get_num_rows() + CTA_SIZE - 1) / CTA_SIZE;\n // The number of CTAs in a grid where each warp proceeds a row.\n const int WARP_GRID_SIZE = (A.get_num_rows() + NUM_WARPS_PER_CTA - 1) / NUM_WARPS_PER_CTA;\n // Edge weights.\n device_vector_alloc edge_weights(A.get_num_nz());\n // Compute edge weights.\n compute_edge_weights_diag <<< WARP_GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n block_size,\n num_rows, // It should be num_owned!!!\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n amgx::thrust::raw_pointer_cast( &edge_weights.front() ) );\n cudaCheckError();\n // Make sure there's enough room to store aggregates.\n aggregates.resize(num_rows);\n // The number of aggregates.\n device_vector_alloc dev_num_aggregates(1, 0), dev_num_unaggregated(1, num_rows);\n // Is a vertex already aggregated.\n device_vector_alloc is_aggregated(num_rows, 0);\n // Hash values.\n device_vector_alloc ring_leader_id0_array(num_rows);\n device_vector_alloc ring_leader_id1_array(num_rows);\n device_vector_alloc ring_leader_hash0_array(num_rows);\n device_vector_alloc ring_leader_hash1_array(num_rows);\n // Is there any unaggregated vertex?\n int num_unaggregated = num_rows;\n\n // Iterate until all vertices are aggregated.\n while ( num_unaggregated > 0 )\n {\n int *ring_leader_id0 = amgx::thrust::raw_pointer_cast(&ring_leader_id0_array.front());\n int *ring_leader_id1 = amgx::thrust::raw_pointer_cast(&ring_leader_id1_array.front());\n int *ring_leader_hash0 = amgx::thrust::raw_pointer_cast(&ring_leader_hash0_array.front());\n int *ring_leader_hash1 = amgx::thrust::raw_pointer_cast(&ring_leader_hash1_array.front());\n\n // Count N-ring of vertices.\n for ( int i = 0 ; i < 2 * num_rings ; ++i )\n {\n compute_ring_leader <<< WARP_GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n num_rows, // It should be num_owned!!!\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast(&is_aggregated.front()),\n i == 0 ? NULL : ring_leader_id0,\n i == 0 ? NULL : ring_leader_hash0,\n ring_leader_id1,\n ring_leader_hash1);\n cudaCheckError();\n std::swap(ring_leader_id0, ring_leader_id1);\n std::swap(ring_leader_hash0, ring_leader_hash1);\n }\n\n // Perform the assignments to aggregates.\n build_aggregates <<< WARP_GRID_SIZE, CTA_SIZE>>>(\n num_rings,\n num_rows,\n num_rows, // It should be num_owned!!!\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast(&is_aggregated.front()),\n ring_leader_id0,\n ring_leader_hash0,\n amgx::thrust::raw_pointer_cast(&edge_weights.front()),\n amgx::thrust::raw_pointer_cast(&dev_num_aggregates.front()),\n aggregates.raw(),\n amgx::thrust::raw_pointer_cast(&dev_num_unaggregated.front()));\n cudaCheckError();\n // Number of aggregated vertices.\n num_unaggregated = dev_num_unaggregated[0];\n }\n\n // The number of aggregates.\n num_aggregates = dev_num_aggregates[0];\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nParallelGreedySelector >::setAggregates_1x1( const MatrixType &A, IVector &aggregates, IVector &aggregates_global, int &num_aggregates )\n{}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nParallelGreedySelector >::setAggregates_common_sqblocks( const MatrixType &A, IVector &aggregates, IVector &aggregates_global, int &num_aggregates )\n{}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< class T_Config >\nParallelGreedySelector_Base::ParallelGreedySelector_Base( AMG_Config &cfg, const std::string &cfg_scope ) :\n Selector()\n{\n}\n\ntemplate< class T_Config >\nvoid ParallelGreedySelector_Base::setAggregates( MatrixType &A, IVector &aggregates, IVector &aggregates_global, int &num_aggregates )\n{\n setAggregates_common_sqblocks( A, aggregates, aggregates_global, num_aggregates );\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define AMGX_CASE_LINE(CASE) template class ParallelGreedySelector_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class ParallelGreedySelector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace aggregation\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\nnamespace amgx\n{\nnamespace aggregation\n{\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n__host__ __device__ int hash_function(int a, int seed = 17)\n{\n a ^= seed;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int NUM_THREADS_PER_ROW, int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid compute_edge_weights_diag( const int A_num_rows,\n const int block_size,\n const int num_owned,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict A_diag,\n const Value_type *__restrict A_vals,\n Value_type *__restrict A_edge_weights )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Shared memory to broadcast columns.\n __shared__ volatile int s_bcast_col[CTA_SIZE];\n // Shared memory to store columns ji.\n __shared__ volatile Value_type s_bcast_ji[CTA_SIZE];\n // Constants.\n // const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n // const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += gridDim.x * NUM_WARPS_PER_CTA )\n {\n // The diagonal of (i,i).\n Value_type row_diag = A_vals[block_size * A_diag[a_row_id]];\n // Load the range of the row.\n int a_col_begin = A_rows[a_row_id + 0];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Iterate over the columns of A.\n for ( ; a_col_begin < a_col_end ; a_col_begin += WARP_SIZE )\n {\n // Each lane works on a different column.\n const int a_col_it = a_col_begin + lane_id;\n // Is it an active thread.\n int is_active = a_col_it < a_col_end;\n // Active nodes load column ids.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n s_bcast_col[threadIdx.x] = a_col_id;\n // Reset ji coefficients.\n s_bcast_ji[threadIdx.x] = Value_type(0);\n\n // Iterate over columns to find ji coefficients if they exist.\n for ( int k = 0, num_cols = __popc(utils::ballot(is_active)) ; k < num_cols ; ++k )\n {\n const int uniform_col_id = s_bcast_col[warp_id * WARP_SIZE + k];\n\n // Skip columns whose id is greater than num_owned.\n if ( uniform_col_id == -1 || uniform_col_id == a_row_id || uniform_col_id >= num_owned )\n {\n continue;\n }\n\n // The bounds of the row.\n int b_row_begin = A_rows[uniform_col_id + 0];\n int b_row_end = A_rows[uniform_col_id + 1];\n\n // Iterate over the row.\n for ( int not_found = 1 ; not_found && b_row_begin < b_row_end ; b_row_begin += WARP_SIZE )\n {\n const int b_row_it = b_row_begin + lane_id;\n // Load the column id.\n int b_col_id = -1;\n\n if ( b_row_it < b_row_end )\n {\n b_col_id = A_cols[b_row_it];\n }\n\n // Has anyone found the column.\n not_found = utils::all( b_col_id != a_row_id );\n\n // If someone found the column, ask it to load the value.\n if ( b_col_id == a_row_id )\n {\n s_bcast_ji[warp_id * WARP_SIZE + k] = A_vals[block_size * b_row_it];\n }\n }\n }\n\n // Deactivate \"invalid\" threads.\n is_active = is_active && a_col_id < num_owned;\n // The diagonal value associated with the column.\n Value_type col_diag(0);\n\n if ( is_active )\n {\n col_diag = A_vals[block_size * A_diag[a_col_id]];\n }\n\n // The value of the column.\n Value_type col_val(0);\n\n if ( is_active )\n {\n col_val = A_vals[block_size * a_col_it];\n }\n\n // Compute the denominator.\n Value_type den = max( abs(row_diag), abs(col_diag) );\n // Compute the weight of the edge.\n Value_type weight(0);\n\n if ( den != Value_type(0) )\n {\n weight = Value_type(0.5) * (abs(col_val) + abs(s_bcast_ji[threadIdx.x])) / den;\n }\n\n if ( is_active )\n {\n A_edge_weights[a_col_it] = weight;\n }\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid compute_ring_leader( const int A_num_rows,\n const int num_owned,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict is_aggregated,\n const int *__restrict in_leader_id,\n const int *__restrict in_leader_hash,\n int *__restrict out_leader_id,\n int *__restrict out_leader_hash )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += gridDim.x * NUM_WARPS_PER_CTA )\n {\n // Skip already aggregated vertices.\n if ( is_aggregated[a_row_id] )\n {\n continue;\n }\n\n // Load the range of the row.\n int a_col_begin = A_rows[a_row_id + 0];\n int a_col_end = A_rows[a_row_id + 1];\n // The max hash in my ith ring.\n int my_min_id = A_num_rows;\n int my_max_hash = INT_MIN;\n\n if ( in_leader_id )\n {\n my_min_id = in_leader_id [a_row_id];\n my_max_hash = in_leader_hash[a_row_id];\n }\n else\n {\n my_min_id = a_row_id;\n my_max_hash = hash_function(a_row_id);\n }\n\n // Iterate over the columns of A.\n for ( ; a_col_begin < a_col_end ; a_col_begin += WARP_SIZE )\n {\n // Each lane works on a different column.\n const int a_col_it = a_col_begin + lane_id;\n // Is it an active thread.\n int is_active = a_col_it < a_col_end;\n // Active nodes load column ids.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Compute the hash value if needed.\n int col_min_id = A_num_rows;\n int col_max_hash = INT_MIN;\n\n if ( in_leader_id )\n {\n if ( is_active && !is_aggregated[a_col_id] )\n {\n col_min_id = in_leader_id [a_col_id];\n col_max_hash = in_leader_hash[a_col_id];\n }\n }\n else\n {\n if ( is_active && !is_aggregated[a_col_id] )\n {\n col_min_id = a_col_id;\n col_max_hash = hash_function(a_col_id);\n }\n }\n\n // Update the max_hash if needed.\n if ( is_active && (my_max_hash < col_max_hash || (my_max_hash == col_max_hash && my_min_id >= col_min_id)) )\n {\n my_min_id = col_min_id;\n my_max_hash = col_max_hash;\n }\n }\n\n // Reduce the max hash.\n#pragma unroll\n\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n int other_min_id = utils::shfl_xor(my_min_id, mask);\n int other_max_hash = utils::shfl_xor(my_max_hash, mask);\n\n if ( other_max_hash > my_max_hash || (other_max_hash == my_max_hash && other_min_id < my_min_id) )\n {\n my_min_id = other_min_id;\n my_max_hash = other_max_hash;\n }\n }\n\n // The 1st thread stores the result.\n if ( lane_id == 0 )\n {\n out_leader_id [a_row_id] = my_min_id;\n out_leader_hash[a_row_id] = my_max_hash;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int CTA_SIZE, int WARP_SIZE >\n__global__\nvoid build_aggregates( const int num_rings,\n const int A_num_rows,\n const int num_owned,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n int *__restrict is_aggregated,\n const int *__restrict in_leader_id,\n const int *__restrict in_leader_hash,\n const Value_type *__restrict edge_weights,\n int *__restrict num_aggregates,\n int *__restrict aggregates,\n int *__restrict num_unaggregated)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += gridDim.x * NUM_WARPS_PER_CTA )\n {\n // Skip already aggregated vertices.\n if ( is_aggregated[a_row_id] )\n {\n continue;\n }\n\n // The max hash in my ith ring.\n int my_hash = hash_function(a_row_id), my_max_hash = in_leader_hash[a_row_id];\n\n // Skip if I'm not a local king.\n if ( my_hash < my_max_hash || (my_hash == my_max_hash && a_row_id > in_leader_id[a_row_id]) )\n {\n continue;\n }\n\n // We start at the row.\n int curr_row = a_row_id;\n\n // Iterate to form the aggregate.\n for ( int aggregate_size = 1, aggregate_id = -1 ; aggregate_size <= num_rings ; ++aggregate_size )\n {\n // Load the range of the row.\n int a_col_begin = A_rows[curr_row + 0];\n int a_col_end = A_rows[curr_row + 1];\n // The max weight.\n int max_id = -1;\n float max_weight = -1.0f;\n // The max aggregated weight.\n int max_aggregated_id = -1;\n float max_aggregated_weight = -1.0f;\n\n // Iterate over the columns of A.\n for ( ; a_col_begin < a_col_end ; a_col_begin += WARP_SIZE )\n {\n // Each lane works on a different column.\n const int a_col_it = a_col_begin + lane_id;\n // Is it an active thread.\n int is_active = a_col_it < a_col_end;\n // Active nodes load column ids.\n int a_col_id = -1;\n float a_col_weight = -1.0f;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n a_col_weight = static_cast(edge_weights[a_col_it]);\n }\n\n // Is the column aggregated?\n int is_col_aggregated = is_active && is_aggregated[a_col_id];\n\n // Select the column if it's not aggregated and its weight is better.\n if ( is_active && !is_col_aggregated && a_col_weight > max_weight )\n {\n max_id = a_col_id;\n max_weight = a_col_weight;\n }\n\n // Update max aggregated weight if needed.\n if ( is_active && is_col_aggregated && a_col_weight > max_aggregated_weight )\n {\n max_aggregated_id = a_col_id;\n max_aggregated_weight = a_col_weight;\n }\n }\n\n // Is there a valid max_id.\n int valid_max_id = utils::any(max_id != -1, utils::activemask());\n\n // No valid max id?\n if ( !valid_max_id && aggregate_size > 1 )\n {\n break;\n }\n\n // Find the max id.\n if ( !valid_max_id )\n {\n max_id = max_aggregated_id;\n max_weight = max_aggregated_weight;\n }\n\n#pragma unroll\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n int other_max_id = utils::shfl_xor(max_id, mask);\n float other_max_weight = utils::shfl_xor(max_weight, mask);\n\n if ( other_max_weight > max_weight )\n {\n max_id = other_max_id;\n max_weight = other_max_weight;\n }\n }\n\n // We know it's a singleton so merge with an existing aggregate.\n if ( !valid_max_id )\n {\n if ( lane_id == 0 && max_id != -1 )\n {\n is_aggregated[curr_row] = 1;\n aggregates[curr_row] = aggregates[max_id];\n atomicAdd(num_unaggregated, -1);\n }\n\n break;\n }\n\n // It's not a singleton but the 1st vertex in the aggregate.\n if ( lane_id == 0 && aggregate_size == 1 )\n {\n aggregate_id = atomicAdd(num_aggregates, 1);\n is_aggregated[a_row_id] = 1;\n aggregates[a_row_id] = aggregate_id;\n atomicAdd(num_unaggregated, -1);\n }\n\n // Set the aggregate of the winner.\n if ( lane_id == 0 )\n {\n aggregates[max_id] = aggregate_id;\n is_aggregated[max_id] = 1;\n atomicAdd(num_unaggregated, -1);\n }\n\n // Set the next row to consider.\n curr_row = max_id;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nenum { WARP_SIZE = 32 };\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nParallelGreedySelector >::setAggregates_1x1( const MatrixType &A, IVector &aggregates, IVector &aggregates_global, int &num_aggregates )\n{}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nParallelGreedySelector >::setAggregates_common_sqblocks( const MatrixType &A, IVector &aggregates, IVector &aggregates_global, int &num_aggregates )\n{\n const int num_rings = 4;\n // The number of rows of A.\n const int num_rows = A.get_num_rows();\n // The size of the block.\n const int block_size = A.get_block_dimx() * A.get_block_dimy();\n // The number of threads per CTA.\n const int CTA_SIZE = 128;\n // The number of warps per CTA.\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // The number of CTAs in a grid where each thread is independent.\n const int THREAD_GRID_SIZE = (A.get_num_rows() + CTA_SIZE - 1) / CTA_SIZE;\n // The number of CTAs in a grid where each warp proceeds a row.\n const int WARP_GRID_SIZE = (A.get_num_rows() + NUM_WARPS_PER_CTA - 1) / NUM_WARPS_PER_CTA;\n // Edge weights.\n device_vector_alloc edge_weights(A.get_num_nz());\n // Compute edge weights.\n compute_edge_weights_diag <<< WARP_GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n block_size,\n num_rows, // It should be num_owned!!!\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n amgx::thrust::raw_pointer_cast( &edge_weights.front() ) );\n cudaCheckError();\n // Make sure there's enough room to store aggregates.\n aggregates.resize(num_rows);\n // The number of aggregates.\n device_vector_alloc dev_num_aggregates(1, 0), dev_num_unaggregated(1, num_rows);\n // Is a vertex already aggregated.\n device_vector_alloc is_aggregated(num_rows, 0);\n // Hash values.\n device_vector_alloc ring_leader_id0_array(num_rows);\n device_vector_alloc ring_leader_id1_array(num_rows);\n device_vector_alloc ring_leader_hash0_array(num_rows);\n device_vector_alloc ring_leader_hash1_array(num_rows);\n // Is there any unaggregated vertex?\n int num_unaggregated = num_rows;\n\n // Iterate until all vertices are aggregated.\n while ( num_unaggregated > 0 )\n {\n int *ring_leader_id0 = amgx::thrust::raw_pointer_cast(&ring_leader_id0_array.front());\n int *ring_leader_id1 = amgx::thrust::raw_pointer_cast(&ring_leader_id1_array.front());\n int *ring_leader_hash0 = amgx::thrust::raw_pointer_cast(&ring_leader_hash0_array.front());\n int *ring_leader_hash1 = amgx::thrust::raw_pointer_cast(&ring_leader_hash1_array.front());\n\n // Count N-ring of vertices.\n for ( int i = 0 ; i < 2 * num_rings ; ++i )\n {\n compute_ring_leader <<< WARP_GRID_SIZE, CTA_SIZE>>>(\n num_rows,\n num_rows, // It should be num_owned!!!\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast(&is_aggregated.front()),\n i == 0 ? NULL : ring_leader_id0,\n i == 0 ? NULL : ring_leader_hash0,\n ring_leader_id1,\n ring_leader_hash1);\n cudaCheckError();\n std::swap(ring_leader_id0, ring_leader_id1);\n std::swap(ring_leader_hash0, ring_leader_hash1);\n }\n\n // Perform the assignments to aggregates.\n build_aggregates <<< WARP_GRID_SIZE, CTA_SIZE>>>(\n num_rings,\n num_rows,\n num_rows, // It should be num_owned!!!\n A.row_offsets.raw(),\n A.col_indices.raw(),\n amgx::thrust::raw_pointer_cast(&is_aggregated.front()),\n ring_leader_id0,\n ring_leader_hash0,\n amgx::thrust::raw_pointer_cast(&edge_weights.front()),\n amgx::thrust::raw_pointer_cast(&dev_num_aggregates.front()),\n aggregates.raw(),\n amgx::thrust::raw_pointer_cast(&dev_num_unaggregated.front()));\n cudaCheckError();\n // Number of aggregated vertices.\n num_unaggregated = dev_num_unaggregated[0];\n }\n\n // The number of aggregates.\n num_aggregates = dev_num_aggregates[0];\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nParallelGreedySelector >::setAggregates_1x1( const MatrixType &A, IVector &aggregates, IVector &aggregates_global, int &num_aggregates )\n{}\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nParallelGreedySelector >::setAggregates_common_sqblocks( const MatrixType &A, IVector &aggregates, IVector &aggregates_global, int &num_aggregates )\n{}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< class T_Config >\nParallelGreedySelector_Base::ParallelGreedySelector_Base( AMG_Config &cfg, const std::string &cfg_scope ) :\n Selector()\n{\n}\n\ntemplate< class T_Config >\nvoid ParallelGreedySelector_Base::setAggregates( MatrixType &A, IVector &aggregates, IVector &aggregates_global, int &num_aggregates )\n{\n setAggregates_common_sqblocks( A, aggregates, aggregates_global, num_aggregates );\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define AMGX_CASE_LINE(CASE) template class ParallelGreedySelector_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class ParallelGreedySelector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace aggregation\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/aggregation/selectors/parallel_greedy_selector.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include //count\n#include //sort\n#include //unique\n#include //remove\n#include //transform_inclusive_scan\n\n#include //offsets_to_indices\n\n#include \n#include \nnamespace amgx\n{\n\n\nnamespace strided_reduction\n{\ntemplate\nvoid count_block_results_pinned_memory(scalar_t *out_host, const int n_blocks, scalar_t *out_d, const OP &op = OP(), cudaStream_t stream = 0)\n{\n strided_reduction_collect_partials <<< 1, 32, 0, stream>>>(out_host, out_d, n_blocks);\n cudaCheckError();\n}\n\n\ntemplate\nscalar_t count_block_results_pinned_memory(const int a, const int i, const int n_blocks, scalar_t *out_d, const OP &op = OP(), cudaStream_t stream = 0) //STRIDE=1 case\n{\n static scalar_t *ret = 0;\n static cudaEvent_t throttle_event = 0;\n const int buffers = 1;\n\n if (ret == 0)\n {\n cudaError_t cuda_rc = amgx::memory::cudaMallocHost((void **)&ret, buffers * sizeof(scalar_t));\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMallocHost failed in size8_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n ret[0] = 0;\n cudaEventCreateWithFlags(&throttle_event, cudaEventDisableTiming);\n }\n\n int ib = i % buffers;\n count_block_results_pinned_memory<1, scalar_t, OP>(ret + ib, n_blocks, out_d, op, stream);\n\n if (ib == buffers - 1)\n {\n cudaError_t cuda_rc = cudaEventRecord(throttle_event);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaEventRecord failed in size8_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = cudaEventSynchronize(throttle_event);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaEventSynchronize failed in size8_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n scalar_t tot = 0;\n\n for (int j = 0; j < buffers; j++)\n {\n tot += ret[j];\n }\n\n return tot + buffers - 1;\n }\n else\n {\n return -1;\n }\n}\n\n\n}\n\nvoid analyze_coloring(device_vector_alloc aggregates_d, device_vector_alloc colors_d);\n\nnamespace aggregation\n{\nnamespace size8_selector\n{\n\n#include \n\ntemplate \n__global__ __launch_bounds__(256, 4)\nvoid my_findStrongestNeighbourBlockDiaCsr_NoMergeClean(\n const IndexType *row_offsets, const IndexType *column_indices,\n const float *edge_weights, const IndexType num_block_rows, const IndexType num_nonzero,\n IndexType *aggregated, IndexType *aggregates, int *strongest_neighbour,\n IndexType *partner_index, float *weight_strongest_neighbour, int deterministic,\n const IndexType *unassigned_rows,\n const int num_unassigned_row)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n const int lane_id = utils::lane_id();\n bool valid_tid = false;\n\n for (; utils::any(valid_tid = tid < num_block_rows); tid += gridDim.x * blockDim.x)\n {\n int jmin = -NUM_COLS * 2, jmax = -NUM_COLS * 4;\n float weight;\n int jcol;\n float max_weight_unaggregated = 0;\n int strongest_unaggregated = -1;\n bool is_unassigned = false;\n\n if (valid_tid)\n {\n is_unassigned = (__ldcs(partner_index + tid) == -1);\n }\n\n if (is_unassigned) // Unaggregated row\n {\n jmin = __ldca(row_offsets + tid);\n jmax = __ldca(row_offsets + tid + 1);\n }\n\n if (utils::any(is_unassigned))\n {\n int jj = jmin - amgx::strided_reduction::warp_loader::align_shift(jmin);\n\n for (; utils::any(jj < jmax, utils::activemask()); jj += NUM_COLS)\n {\n int I[NUM_COLS];\n float W[NUM_COLS];\n int P[NUM_COLS];\n amgx::strided_reduction::warp_loader::load(column_indices, jj, num_nonzero, I);\n\n#pragma unroll\n for (int i = 0; i < NUM_COLS; i++)\n {\n int j = jj + i;\n jcol = I[i];\n\n if (j >= jmin && j < jmax)\n {\n P[i] = __ldg(partner_index + jcol); //make this load ASAP\n }\n }\n\n amgx::strided_reduction::warp_loader::load(edge_weights, jj, num_nonzero, W);\n\n#pragma unroll\n for (int i = 0; i < NUM_COLS; i++)\n {\n weight = W[i];\n jcol = I[i];\n int j = jj + i;\n\n if (j >= jmin && j < jmax)\n {\n if (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && jcol > strongest_unaggregated)) // unaggregated\n {\n if (tid != jcol && P[i] == -1)\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n }\n }\n }\n }\n\n if (strongest_unaggregated == -1) // All neighbours are aggregated\n {\n // Put in its own aggregate\n if (!deterministic && is_unassigned)\n {\n partner_index[tid] = tid;\n }\n }\n else\n {\n strongest_neighbour[tid] = strongest_unaggregated;\n }\n }\n }\n}\n\n\n\n#define ALGORITHM_NOMERGE 0\n#define ALGORITHM_STOREWEIGHTS 1\n#define ALGORITHM_STOREWEIGHTS_2 2\n\n\ntemplate \n__global__ __launch_bounds__(256, 4)\nvoid my_findStrongestNeighbourBlockDiaCsr_NoMerge(\n const IndexType *row_offsets, const IndexType *column_indices,\n const float *edge_weights, const IndexType num_block_rows, const IndexType num_nonzero,\n IndexType *aggregated, IndexType *aggregates, int *strongest_neighbour,\n IndexType *partner_index, float *weight_strongest_neighbour, int deterministic,\n const IndexType *n_unassigned_per_block, const IndexType *unassigned_per_block\n //const int num_unassigned_row\n //const IndexType *unassigned_rows,\n)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int bid = blockIdx.x;\n const int lane_id = utils::lane_id();\n bool valid_tid = false;\n\n for (; utils::any( valid_tid = tid < num_block_rows); tid += gridDim.x * blockDim.x)\n {\n int jmin = -NUM_COLS * 2, jmax = -NUM_COLS * 4;\n float weight;\n int jcol;\n float max_weight_unaggregated = 0;\n int strongest_unaggregated = -1;\n float max_weight_aggregated = 0.;\n int strongest_aggregated = -1;\n int partner = -1;\n int partner0, partner1, partner2;\n int agg_jcol;\n bool is_unassigned = false;\n int rowi = -1;\n\n if (LOAD_ONLY_UNASSIGNED)\n {\n if (valid_tid)\n {\n rowi = unassigned_per_block[tid];//unassigned_per_block+bid*256+threadIdx.x);\n is_unassigned = (__ldg(partner_index + rowi) == -1);\n }\n\n if (is_unassigned)\n {\n jmin = __ldg(row_offsets + rowi);\n jmax = __ldg(row_offsets + rowi + 1);\n }\n }\n else\n {\n rowi = tid;\n\n if (ALGORITHM == ALGORITHM_NOMERGE)\n {\n if (valid_tid) { is_unassigned = (__ldcs(partner_index + tid) == -1); }\n }\n else //ALGORITHM_STOREWEIGHTS or ALGORITHM_STOREWEIGHTS_2\n {\n if (valid_tid) { is_unassigned = (__ldcs(aggregated + tid) == -1); }\n }\n\n if (is_unassigned) // mind the else above\n {\n jmin = __ldg(row_offsets + rowi);\n jmax = __ldlu(row_offsets + rowi + 1);\n }\n }\n\n if (utils::any(is_unassigned))\n {\n if (is_unassigned) // Unaggregated row\n {\n if (ALGORITHM == ALGORITHM_STOREWEIGHTS)\n {\n partner = partner_index[rowi];\n }\n else if (ALGORITHM == ALGORITHM_STOREWEIGHTS_2)\n {\n partner0 = partner_index[rowi];\n partner1 = partner_index[num_block_rows + rowi];\n partner2 = partner_index[2 * num_block_rows + rowi];\n }\n }\n\n int jj = jmin - amgx::strided_reduction::warp_loader::align_shift(jmin);\n\n for (; utils::any(jj < jmax && jmax >= 0); jj += NUM_COLS)\n {\n int I[NUM_COLS];\n float W[NUM_COLS];\n int P[NUM_COLS];\n int jj_ok = (jj >= 0 && jj < jmax && jmax >= 0) ? jj : 0;\n\n amgx::strided_reduction::warp_loader::load(column_indices, jj_ok, num_nonzero, I);\n#pragma unroll\n\n for (int i = 0; i < NUM_COLS; i++)\n {\n int j = jj + i;\n jcol = I[i];\n\n if (j >= jmin && j < jmax)\n {\n if (ALGORITHM == ALGORITHM_NOMERGE)\n {\n P[i] = __ldg(partner_index + jcol); //make this load ASAP\n }\n else\n {\n P[i] = __ldg(aggregated + jcol);\n }\n }\n }\n\n amgx::strided_reduction::warp_loader::load(edge_weights, jj_ok, num_nonzero, W);\n#pragma unroll\n\n for (int i = 0; i < NUM_COLS; i++)\n {\n weight = W[i];\n jcol = I[i];\n int j = jj + i;\n\n if (j >= jmin && j < jmax)\n {\n if (ALGORITHM == ALGORITHM_NOMERGE)\n {\n if (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && jcol > strongest_unaggregated)) // unaggregated\n {\n if (rowi != jcol && P[i] == -1)\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n }\n }\n else\n {\n bool partner_condition;\n\n if (ALGORITHM == ALGORITHM_STOREWEIGHTS)\n {\n partner_condition = jcol != partner;\n }\n else if (ALGORITHM == ALGORITHM_STOREWEIGHTS_2)\n {\n partner_condition = jcol != partner0 && jcol != partner1 && jcol != partner2;\n }\n\n agg_jcol = P[i];\n\n if (partner_condition && rowi != jcol)\n {\n if (agg_jcol == -1 && (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && jcol > strongest_unaggregated))) // unaggregated\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n else if (agg_jcol != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) // unaggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n }\n }\n }\n }\n\n if (valid_tid && is_unassigned)\n {\n if (ALGORITHM == ALGORITHM_NOMERGE)\n {\n if (strongest_unaggregated == -1) // All neighbours are aggregated\n {\n // Put in its own aggregate\n if (!deterministic)\n {\n partner_index[rowi] = rowi;\n }\n }\n else\n {\n strongest_neighbour[rowi] = strongest_unaggregated;\n }\n }\n else if (ALGORITHM == ALGORITHM_STOREWEIGHTS)\n {\n if (strongest_unaggregated == -1) // All neighbours are aggregated\n {\n if (!deterministic)\n {\n if (strongest_aggregated != -1)\n {\n aggregates[tid] = aggregates[strongest_aggregated];\n aggregated[tid] = 1;\n aggregates[partner] = aggregates[strongest_aggregated];\n aggregated[partner] = 1;\n }\n else // leave in its own aggregate\n {\n aggregated[partner] = 1;\n aggregated[tid] = 1;\n }\n }\n }\n else // Found an unaggregated aggregate\n {\n weight_strongest_neighbour[tid] = max_weight_unaggregated;\n strongest_neighbour[tid] = aggregates[strongest_unaggregated];\n }\n }\n else if (ALGORITHM == ALGORITHM_STOREWEIGHTS_2)\n {\n if (strongest_unaggregated == -1 && strongest_aggregated != -1) // all neighbours are aggregated, store the strongest aggregated\n {\n weight_strongest_neighbour[tid] = -max_weight_aggregated;\n strongest_neighbour[tid] = aggregates[strongest_aggregated];\n }\n else if (strongest_unaggregated != -1)\n {\n weight_strongest_neighbour[tid] = max_weight_unaggregated;\n strongest_neighbour[tid] = aggregates[strongest_unaggregated];\n }\n }\n }\n }\n\n bid += gridDim.x;\n }\n}\n\n\n\n\n#define INST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(numcols, algo,assume,c) template __global__ void my_findStrongestNeighbourBlockDiaCsr_NoMerge(\\\nconst int *row_offsets, const int *column_indices,\\\nconst float *edge_weights, const int num_block_rows,const int num_nonzero,\\\nint *aggregated, int *aggregates, int *strongest_neighbour,\\\nint *partner_index, float *weight_strongest_neighbour, int deterministic,\\\nconst int* n_unassigned_per_block, const int * unassigned_per_block);\n\n#define INSTmy_findStrongestNeighbourBlockDiaCsr_NoMergeClean(numcols) template __global__ void my_findStrongestNeighbourBlockDiaCsr_NoMergeClean(\\\nconst int *row_offsets, const int *column_indices,\\\nconst float *edge_weights, const int num_block_rows,const int num_nonzero,\\\nint *aggregated, int *aggregates, int *strongest_neighbour,\\\nint *partner_index, float *weight_strongest_neighbour, int deterministic,\\\nconst int *unassigned_rows,\\\nconst int num_unassigned_row);\n\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(1, ALGORITHM_NOMERGE, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(1, ALGORITHM_STOREWEIGHTS, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(1, ALGORITHM_STOREWEIGHTS_2, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(1, ALGORITHM_NOMERGE, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(1, ALGORITHM_STOREWEIGHTS, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(1, ALGORITHM_STOREWEIGHTS_2, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(1, ALGORITHM_NOMERGE, 0, 1) //load only unassigned\n\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(2, ALGORITHM_NOMERGE, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(2, ALGORITHM_STOREWEIGHTS, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(2, ALGORITHM_STOREWEIGHTS_2, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(2, ALGORITHM_NOMERGE, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(2, ALGORITHM_STOREWEIGHTS, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(2, ALGORITHM_STOREWEIGHTS_2, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(2, ALGORITHM_NOMERGE, 0, 1) //load only unassigned\n\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(4, ALGORITHM_NOMERGE, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(4, ALGORITHM_STOREWEIGHTS, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(4, ALGORITHM_STOREWEIGHTS_2, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(4, ALGORITHM_NOMERGE, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(4, ALGORITHM_STOREWEIGHTS, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(4, ALGORITHM_STOREWEIGHTS_2, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(4, ALGORITHM_NOMERGE, 0, 1) //load only unassigned\n\n\nINSTmy_findStrongestNeighbourBlockDiaCsr_NoMergeClean(1)\nINSTmy_findStrongestNeighbourBlockDiaCsr_NoMergeClean(2)\nINSTmy_findStrongestNeighbourBlockDiaCsr_NoMergeClean(4)\n\ntemplate\n__global__ void my_blockCompact(\n int *partner_index, const int num_rows,\n int *unassigned_per_block_in,\n int *n_unassigned_per_block, int *unassigned_per_block)\n{\n int bid = blockIdx.x; //RMV\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n for (; bid < num_rows / 256 + 1/*__any(tid < num_rows)*/; tid += gridDim.x * blockDim.x)\n {\n int row = tid;\n bool no_partner = 0; //RMV\n\n if (tid < num_rows)\n {\n if (ALREADY_COMPACT)\n {\n row = unassigned_per_block_in[tid];\n }\n\n if (partner_index[row] == -1) // Unaggregated row\n {\n no_partner = 1;\n }\n }\n\n amgx::strided_reduction::block_binary_compaction<256, 32, 1>(\n n_unassigned_per_block, unassigned_per_block, bid,\n no_partner, row);\n bid += gridDim.x;\n }\n}\n\ntemplate __global__ void my_blockCompact<0>(\n int *partner_index, const int num_rows,\n int *unassigned_per_block_in,\n int *n_unassigned_per_block, int *unassigned_per_block);\ntemplate __global__ void my_blockCompact<1>(\n int *partner_index, const int num_rows,\n int *unassigned_per_block_in,\n int *n_unassigned_per_block, int *unassigned_per_block);\n\n__global__ void my_MatchEdgesWithCompaction(const int num_rows, int *partner_index, int *aggregates, const int *strongest_neighbour, int *sets_per_block,\n int *unassigned_per_block_in, int *n_unassigned_per_block, int *unassigned_per_block\n )\n{\n int potential_match, potential_match_neighbour;\n int warp_count = 0;\n int bid = blockIdx.x; //RMV\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n for (; bid < num_rows / 1024 + 1/*__any(tid < num_rows)*/; tid += gridDim.x * blockDim.x)\n {\n int row = tid;\n bool no_partner = 0;\n\n if (tid < num_rows)\n {\n if (partner_index[row] == -1) // Unaggregated row\n {\n no_partner = 1;\n potential_match = strongest_neighbour[row];\n\n if (potential_match != -1)\n {\n potential_match_neighbour = __ldcs(strongest_neighbour + potential_match);\n\n if ( potential_match_neighbour == row ) // we have a match\n {\n no_partner = 0;\n //partner_notnull = 1;//RMV\n partner_index[row] = potential_match;\n aggregates[row] = ( potential_match > row ) ? row : potential_match;\n }\n }\n }\n }\n\n amgx::strided_reduction::block_binary_compaction<1024, 32, 1>(\n n_unassigned_per_block, unassigned_per_block, bid,\n no_partner, row);\n warp_count += amgx::strided_reduction::warp_binary_count(no_partner);\n bid += gridDim.x;\n }\n\n amgx::strided_reduction::block_count<1, 1024, 32, int>(warp_count, sets_per_block);\n}\n\n__global__ void my_MatchEdges(const int num_rows, int *partner_index, int *aggregates, const int *strongest_neighbour, int *sets_per_block)\n{\n int potential_match, potential_match_neighbour;\n int warp_count = 0;\n\n for (int tid = threadIdx.x + blockDim.x * blockIdx.x; utils::any(tid < num_rows); tid += gridDim.x * blockDim.x)\n {\n bool has_set_partner_index = 0;\n\n if (tid < num_rows)\n {\n if (partner_index[tid] == -1) // Unaggregated row\n {\n potential_match = strongest_neighbour[tid];\n\n if (potential_match != -1)\n {\n potential_match_neighbour = __ldcs(strongest_neighbour + potential_match);\n\n if ( potential_match_neighbour == tid ) // we have a match\n {\n has_set_partner_index = 1;\n partner_index[tid] = potential_match;\n aggregates[tid] = ( potential_match > tid) ? tid : potential_match;\n }\n }\n }\n }\n\n warp_count += amgx::strided_reduction::warp_binary_count(has_set_partner_index);\n }\n\n amgx::strided_reduction::block_count<1, 1024, 32, int>(warp_count, sets_per_block);\n}\n\n\n// matchEdges\n__global__ void my_joinExistingAggregates(int num_rows, int *aggregates, int *aggregated, int *aggregates_candidate, int *sets_per_block)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int warp_count = 0;\n\n while (utils::any(tid < num_rows))\n {\n bool has_set_partner_index = 0;\n\n if (tid < num_rows)\n {\n if (aggregated[tid] == -1 && aggregates_candidate[tid] != -1) // Unaggregated row\n {\n aggregates[tid] = aggregates_candidate[tid];\n aggregated[tid] = 1;\n has_set_partner_index = 1;\n }\n }\n\n warp_count += amgx::strided_reduction::warp_binary_count(has_set_partner_index);\n tid += gridDim.x * blockDim.x;\n }\n\n amgx::strided_reduction::block_count<1, 1024, 32, int>(warp_count, sets_per_block);\n}\n\n\n// Kernel that checks if perfect matchs exist\n__global__ void my_matchAggregates(int *aggregates, int *aggregated, int *strongest_neighbour, const int num_rows, int *sets_per_block)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int potential_match, potential_match_neighbour, my_aggregate;\n int warp_count = 0;\n\n while (utils::any(tid < num_rows))\n {\n bool has_set_partner_index = 0;\n\n if (tid < num_rows) if (aggregated[tid] == -1) // Unaggregated row\n {\n potential_match = strongest_neighbour[tid];\n\n if (potential_match != -1)\n {\n potential_match_neighbour = __ldcs(strongest_neighbour + potential_match); //or global\n my_aggregate = aggregates[tid];\n\n if (potential_match_neighbour == my_aggregate) // we have a match\n {\n has_set_partner_index = 1;\n aggregated[tid] = 1;\n aggregates[tid] = ( potential_match > my_aggregate) ? my_aggregate : potential_match;\n }\n }\n }\n\n warp_count += amgx::strided_reduction::warp_binary_count(has_set_partner_index);\n tid += gridDim.x * blockDim.x;\n }\n\n amgx::strided_reduction::block_count<1, 1024, 32, int>(warp_count, sets_per_block);\n}\n\n\n// Kernel that checks if perfect matchs exist\n__global__ void my_matchAggregatesSize4(int *aggregates, int *aggregated, int *strongest_neighbour, int *partner_index, const int num_rows, int *sets_per_block)\n{\n int potential_match, potential_match_neighbour, my_aggregate;\n int warp_count = 0;\n\n for (int tid = threadIdx.x + blockDim.x * blockIdx.x; utils::any(tid < num_rows); tid += blockDim.x * gridDim.x)\n {\n bool has_set_partner_index = 0;\n\n if (tid < num_rows) if (aggregated[tid] == -1) // Unaggregated row\n {\n potential_match = strongest_neighbour[tid];\n\n if (potential_match != -1)\n {\n potential_match_neighbour = __ldcs(strongest_neighbour + potential_match);\n my_aggregate = aggregates[tid];\n\n if (potential_match_neighbour == my_aggregate) // we have a match\n {\n has_set_partner_index = 1;\n aggregated[tid] = 1;\n aggregates[tid] = ( potential_match > my_aggregate) ? my_aggregate : potential_match;\n partner_index[tid + num_rows] = potential_match;\n partner_index[tid + 2 * num_rows] = partner_index[potential_match];\n }\n }\n }\n\n warp_count += amgx::strided_reduction::warp_binary_count(has_set_partner_index);\n }\n\n amgx::strided_reduction::block_count<1, 1024, 32, int>(warp_count, sets_per_block);\n}\n\n// ------------------------\n// Kernels\n// ------------------------\n\n// findStrongestNeighbour kernel for block_dia_csr_matrix format\n// Reads the weight from edge_weights array\ntemplate \n__global__\nvoid findStrongestNeighbourBlockDiaCsr_StoreWeight_2(const IndexType *row_offsets, const IndexType *column_indices,\n const float *edge_weights, const IndexType num_block_rows, IndexType *aggregated, IndexType *aggregates, int *strongest_neighbour, IndexType *partner_index, float *weight_strongest_neighbour)\n{\n float weight;\n int jcol, jmin, jmax;\n int partner0, partner1, partner2;\n int agg_jcol;\n\n for (int tid = threadIdx.x + blockDim.x * blockIdx.x; tid < num_block_rows; tid += blockDim.x * gridDim.x)\n {\n float max_weight_unaggregated = 0.;\n float max_weight_aggregated = 0.;\n int strongest_unaggregated = -1;\n int strongest_aggregated = -1;\n\n if (aggregated[tid] == -1) // Unaggregated row\n {\n partner0 = partner_index[tid];\n partner1 = partner_index[num_block_rows + tid];\n partner2 = partner_index[2 * num_block_rows + tid];\n jmin = row_offsets[tid];\n jmax = row_offsets[tid + 1];\n\n for (int j = jmin; j < jmax; j++)\n {\n jcol = column_indices[j];\n\n if (jcol == tid || jcol >= num_block_rows) { continue; }\n\n weight = edge_weights[j];\n agg_jcol = aggregated[jcol];\n\n if (jcol != partner0 && jcol != partner1 && jcol != partner2)\n {\n if (agg_jcol == -1 && (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && jcol > strongest_unaggregated))) // unaggregated\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n else if (agg_jcol != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) // unaggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n }\n\n if (strongest_unaggregated == -1 && strongest_aggregated != -1) // all neighbours are aggregated, store the strongest aggregated\n {\n weight_strongest_neighbour[tid] = -max_weight_aggregated;\n strongest_neighbour[tid] = aggregates[strongest_aggregated];\n }\n else if (strongest_unaggregated != -1)\n {\n weight_strongest_neighbour[tid] = max_weight_unaggregated;\n strongest_neighbour[tid] = aggregates[strongest_unaggregated];\n }\n }\n }\n}\n\n// findStrongestNeighbour kernel for block_dia_csr_matrix format\n// Reads the weight from edge_weights array\ntemplate \n__global__\nvoid agreeOnProposal_2(const IndexType *row_offsets, const IndexType *column_indices,\n IndexType num_block_rows, IndexType *aggregated, int *strongest_neighbour, float *weight_strongest_neighbour, IndexType *partner_index, int *aggregates, int deterministic)\n{\n int partner[3];\n float weight[3];\n\n for (int tid = threadIdx.x + blockDim.x * blockIdx.x; tid < num_block_rows; tid += blockDim.x * gridDim.x)\n {\n int strongest_partner = -1;\n float my_weight = 0.;\n\n if (aggregated[tid] == -1)\n {\n my_weight = weight_strongest_neighbour[tid];\n#pragma unroll\n\n for (int m = 0; m < 3; m++)\n {\n partner[m] = partner_index[tid + m * num_block_rows];\n }\n\n#pragma unroll\n\n for (int m = 0; m < 3; m++)\n {\n weight[m] = weight_strongest_neighbour[partner[m]];\n }\n\n#pragma unroll\n\n for (int m = 0; m < 3; m++)\n {\n //if (weight[m] > my_weight && weight[m] > 0.) // there is a partner that has an unaggregated neighbour\n if (weight[m] > my_weight)\n {\n if (weight[m] > 0.) // there is a partner that has an unaggregated neighbour\n {\n strongest_partner = m;\n my_weight = weight[m];\n }\n }\n else if (weight[m] < my_weight) // there is a partner without an unaggregated neighbour, whose neighbour is stronger than mine\n {\n if (my_weight < 0.)\n {\n strongest_partner = m;\n my_weight = weight[m];\n }\n }\n }\n\n if (my_weight < 0.) // means all neighbours of vertices in aggregate are aggregated, merge to another aggregate\n {\n if (!deterministic)\n {\n aggregated[tid] = 1;\n aggregates[tid] = strongest_partner != -1 ? strongest_neighbour[partner[strongest_partner]] : strongest_neighbour[tid];\n }\n }\n else if (strongest_partner != -1) // store my partner's pick\n {\n strongest_neighbour[tid] = strongest_neighbour[partner[strongest_partner]];\n }\n }\n }\n}\n\ntemplate \n__global__\nvoid agreeOnProposal_2_deterministic(int *strongest_neighbour_out, const IndexType *row_offsets, const IndexType *column_indices,\n IndexType num_block_rows, IndexType *aggregated, int *strongest_neighbour, float *weight_strongest_neighbour, IndexType *partner_index, int *aggregates, int deterministic)\n{\n int partner[3];\n float weight[3];\n\n for (int tid = threadIdx.x + blockDim.x * blockIdx.x; tid < num_block_rows; tid += blockDim.x * gridDim.x)\n {\n int strongest_partner = -1;\n float my_weight = 0.;\n //copy here to avoid redundant copies before launching the kernel\n int new_strongest_neighbour_out = strongest_neighbour[tid];\n\n if (aggregated[tid] == -1)\n {\n my_weight = weight_strongest_neighbour[tid];\n#pragma unroll\n\n for (int m = 0; m < 3; m++)\n {\n partner[m] = partner_index[tid + m * num_block_rows];\n }\n\n#pragma unroll\n\n for (int m = 0; m < 3; m++)\n {\n weight[m] = weight_strongest_neighbour[partner[m]];\n }\n\n#pragma unroll\n\n for (int m = 0; m < 3; m++)\n {\n if (weight[m] > my_weight)\n {\n if (weight[m] > 0.) // there is a partner that has an unaggregated neighbour\n {\n strongest_partner = m;\n my_weight = weight[m];\n }\n }\n else if (weight[m] < my_weight) // there is a partner without an unaggregated neighbour, whose neighbour is stronger than mine\n {\n if (my_weight < 0.)\n {\n strongest_partner = m;\n my_weight = weight[m];\n }\n }\n }\n\n if (my_weight < 0.) // means all neighbours of vertices in aggregate are aggregated, merge to another aggregate\n {\n if (!deterministic)\n {\n aggregated[tid] = 1;\n aggregates[tid] = strongest_partner != -1 ? strongest_neighbour[partner[strongest_partner]] : strongest_neighbour[tid];\n }\n }\n else if (strongest_partner != -1) // store my partner's pick\n {\n new_strongest_neighbour_out = strongest_neighbour[partner[strongest_partner]];\n }\n }\n\n //copy here to avoid redundant copies before launching the kernel\n strongest_neighbour_out[tid] = new_strongest_neighbour_out;\n }\n}\n\n\n// Kernel that checks if perfect matchs exist\ntemplate \n__global__\nvoid matchAggregatesSize4(IndexType *aggregates, IndexType *aggregated, IndexType *strongest_neighbour, IndexType *partner_index, const IndexType num_rows)\n{\n int potential_match, potential_match_neighbour, my_aggregate;\n\n for (int tid = threadIdx.x + blockDim.x * blockIdx.x; tid < num_rows; tid += blockDim.x * gridDim.x)\n {\n if (aggregated[tid] == -1) // Unaggregated row\n {\n potential_match = strongest_neighbour[tid];\n\n if (potential_match != -1)\n {\n potential_match_neighbour = strongest_neighbour[potential_match];\n my_aggregate = aggregates[tid];\n\n if (potential_match_neighbour == my_aggregate) // we have a match\n {\n aggregated[tid] = 1;\n aggregates[tid] = ( potential_match > my_aggregate) ? my_aggregate : potential_match;\n partner_index[tid + num_rows] = potential_match;\n partner_index[tid + 2 * num_rows] = partner_index[potential_match];\n }\n }\n }\n }\n}\n\n\ntemplate \n__global__\nvoid assignUnassignedVertices_2(IndexType *partner_index, const IndexType num_rows)\n{\n for (int tid = threadIdx.x + blockDim.x * blockIdx.x; tid < num_rows; tid += blockDim.x * gridDim.x)\n {\n if (partner_index[num_rows + tid] == -1) // Unaggregated row\n {\n partner_index[num_rows + tid] = tid;\n }\n\n if (partner_index[2 * num_rows + tid] == -1) // Unaggregated row\n {\n partner_index[2 * num_rows + tid] = tid;\n }\n }\n}\n\n\n// -----------------\n// Methods\n// ----------------\n\n// Constructor\ntemplate\nSize8SelectorBase::Size8SelectorBase(AMG_Config &cfg, const std::string &cfg_scope)\n{\n deterministic = cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\");\n max_iterations = cfg.AMG_Config::template getParameter(\"max_matching_iterations\", cfg_scope);\n numUnassigned_tol = cfg.AMG_Config::template getParameter(\"max_unassigned_percentage\", cfg_scope);\n m_aggregation_edge_weight_component = cfg.AMG_Config::template getParameter(\"aggregation_edge_weight_component\", cfg_scope);\n weight_formula = cfg.AMG_Config::template getParameter(\"weight_formula\", cfg_scope);\n}\n\n\n// setAggregates for block_dia_csr_matrix_h format\ntemplate \nvoid Size8Selector >::setAggregates_common_sqblock(const Matrix_h &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n FatalError(\"Size8 selector: setAggregates not implemented on CPU, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n\ntemplate \n__global__ //__launch_bounds__(256,2)\nvoid computeEdgeWeightsBlockDiaCsr_V2_1(\n const IndexType *row_offsets,\n //const int* __myrestrict row_offsets,\n const IndexType *row_indices,\n const IndexType *column_indices,\n const IndexType *dia_values,\n const ValueType *nonzero_values,\n //const ValueType* __restrict nonzero_values,\n const IndexType num_nonzero_blocks,\n float *str_edge_weights, float *rand_edge_weights, int num_owned, int bsize, int component)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int i, j;\n int bsize_sq = bsize * bsize;\n float kvalue;\n int matrix_weight_entry = component * bsize + component;\n bool valid_tid;\n\n while (utils::any(valid_tid = tid < num_nonzero_blocks))\n {\n i = -1;\n double d1, d2, w1;\n\n if (valid_tid)\n {\n if ( rand_edge_weights != NULL )\n {\n rand_edge_weights[tid] = random_weight(i, j, num_owned);\n }\n\n i = row_indices[tid];\n j = column_indices[tid];\n d1 = types::util::abs(__cachingLoad(&nonzero_values[__ldg(dia_values + i) * bsize_sq + matrix_weight_entry]));\n d2 = types::util::abs(__cachingLoad(&nonzero_values[__ldg(dia_values + j) * bsize_sq + matrix_weight_entry]));\n }\n\n const bool valid_j = valid_tid && i != j && j < num_owned;\n int ki = -1; //my transpose index, initialized to not found\n //int diag_j = -1; //j diagonal index\n\n if (!utils::any(valid_j))\n {\n continue;\n }\n\n int kmin = 0, kmax = 0;\n\n if (valid_j)\n {\n kmin = __cachingLoad(&row_offsets[j ]);\n kmax = __cachingLoad(&row_offsets[j + 1]);\n }\n\n for ( int k = kmin ; k < kmax ; ++k )\n {\n const int idx = __ldg(column_indices + k);\n\n if (idx == i)\n {\n ki = k; //find the transpose ji\n }\n }\n\n kvalue = 0.0f;\n\n if (ki > -1)\n {\n kvalue = types::util::abs(__cachingLoad(&nonzero_values[ki * bsize_sq + matrix_weight_entry]));\n }\n\n if (valid_tid)\n {\n w1 = types::util::abs(__cachingLoad(&nonzero_values[tid * bsize_sq + matrix_weight_entry]));\n str_edge_weights[tid] = 0.5 * (w1 + kvalue) / ( (float) max(d1, d2) ) * valid_j;\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n#if AMGX_ASYNCCPU_PROOF_OF_CONCEPT\n\ntemplate\nstruct async_size8_task : public task_cuda\n{\n T *size8;\n M *A;\n V *aggregates;\n V *aggregates_global;\n int *num_aggregates;\n\n typedef typename T::IndexType IndexType;\n\n void run()\n {\n const IndexType nnz_per_row = (*A).get_num_nz() / (*A).get_num_rows();\n\n if (0)\n {\n }\n else if (nnz_per_row > 2)\n {\n size8->template setAggregates_common_sqblock_avg_specialized<4>(*A, *aggregates, *aggregates_global, *num_aggregates);\n }\n else if (nnz_per_row > 1)\n {\n size8->template setAggregates_common_sqblock_avg_specialized<4>(*A, *aggregates, *aggregates_global, *num_aggregates);\n }\n else\n {\n size8->template setAggregates_common_sqblock_avg_specialized<2>(*A, *aggregates, *aggregates_global, *num_aggregates);\n }\n }\n};\n\n\ntemplate\nasync_size8_task *make_async_size8_task(T *size8, M &A, V &aggregates, V &aggregates_global, int &num_aggregates)\n{\n async_size8_task *ret = new async_size8_task;\n ret->size8 = size8;\n ret->A = &A;\n ret->aggregates = &aggregates;\n ret->aggregates_global = &aggregates_global;\n ret->num_aggregates = &num_aggregates;\n static task_chain_cuda_streamset *ss = new task_chain_cuda_streamset(1);\n task_chain_cuda *cr = new task_chain_cuda(ss);\n cr->append(ret, asyncmanager::singleton()->main_thread_queue(2));\n return ret;\n}\n\n#endif\n\ntemplate \nvoid Size8Selector >::setAggregates_common_sqblock(const Matrix_d &A,\n typename Matrix_d::IVector &aggregates,\n typename Matrix_d::IVector &aggregates_global,\n int &num_aggregates)\n{\n#if AMGX_ASYNCCPU_PROOF_OF_CONCEPT\n bool task = false;\n bool push = false;\n\n if (task)\n {\n task_cuda *t = make_async_size8_task(this, A, aggregates, aggregates_global, num_aggregates);\n enqueue_async_get_receipt(asyncmanager::singleton()->global_parallel_queue, t)->wait();\n return;\n }\n\n cudaStream_t stream_old;\n static cudaStream_t stream = 0;\n\n if (push)\n {\n stream_old = amgx::thrust::global_thread_handle::threadStream[getCurrentThreadId()];\n cudaError_t cuda_rc = cudaStreamSynchronize(amgx::thrust::global_thread_handle::threadStream[getCurrentThreadId()]);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaStreamSynchronize failed in size8_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n\n if (stream == 0) { cudaStreamCreate(&stream); cudaCheckError(); }\n\n amgx::thrust::global_thread_handle::threadStream[getCurrentThreadId()] = stream;\n }\n\n#endif\n const IndexType nnz_per_row = A.get_num_nz() / A.get_num_rows();\n\n if (0)\n {\n }\n else if (nnz_per_row > 2)\n {\n setAggregates_common_sqblock_avg_specialized<4>(A, aggregates, aggregates_global, num_aggregates);\n }\n else if (nnz_per_row > 1)\n {\n setAggregates_common_sqblock_avg_specialized<4>(A, aggregates, aggregates_global, num_aggregates);\n }\n else\n {\n setAggregates_common_sqblock_avg_specialized<2>(A, aggregates, aggregates_global, num_aggregates);\n }\n\n#if AMGX_ASYNCCPU_PROOF_OF_CONCEPT\n\n if (push)\n {\n //cudaStreamSynchronize(stream);\n cudaCheckError();\n //amgx::thrust::global_thread_handle::threadStream[getCurrentThreadId()] = stream_old;\n }\n\n#endif\n}\n\n// setAggregates for block_dia_csr_matrix_d format\ntemplate \ntemplate\nvoid Size8Selector >:: setAggregates_common_sqblock_avg_specialized(const Matrix_d &A,\n typename Matrix_d::IVector &aggregates, typename Matrix_d::IVector &aggregates_global, int &num_aggregates)\n{\n#if AMGX_ASYNCCPU_PROOF_OF_CONCEPT\n cudaStream_t stream = amgx::thrust::global_thread_handle::threadStream[getCurrentThreadId()];\n#else\n cudaStream_t stream = amgx::thrust::global_thread_handle::get_stream();\n#endif\n const IndexType num_block_rows = A.get_num_rows();\n const IndexType num_nonzero_blocks = A.get_num_nz();\n\n if (!A.is_matrix_singleGPU())\n {\n aggregates.resize(A.manager->halo_offset(A.manager->num_neighbors()));\n }\n else\n {\n aggregates.resize(num_block_rows);\n }\n\n // Initially, put each vertex in its own aggregate\n thrust_wrapper::sequence(aggregates.begin(), aggregates.begin() + num_block_rows);\n cudaCheckError();\n IndexType *aggregates_ptr = aggregates.raw();\n // Create row_indices array\n IndexType total_nz = (A.is_matrix_singleGPU()) ? num_nonzero_blocks : A.manager->num_nz_all();\n typename Matrix_d::IVector row_indices(total_nz);\n cusp::detail::offsets_to_indices(A.row_offsets, row_indices);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_row_indices_ptr = row_indices.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n //const ValueType *A_dia_val_ptr = amgx::thrust::raw_pointer_cast(&A.values[A.get_block_size()*A.diagOffset()]);\n const ValueType *A_nonzero_values_ptr = A.values.raw();\n typename Matrix_d::IVector strongest_neighbour(num_block_rows, -1);\n typename Matrix_d::IVector partner_index(3 * num_block_rows, -1);\n typename Matrix_d::IVector strongest_neighbour_tmp(num_block_rows);\n IndexType *strongest_neighbour_ptr = strongest_neighbour.raw();\n IndexType *partner_index_ptr = partner_index.raw();\n const int threads_per_block = 256;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_block_rows - 1) / threads_per_block + 1);\n int numUnassigned = num_block_rows;\n int numUnassigned_previous = numUnassigned;\n Vector > edge_weights(num_nonzero_blocks + 8); //8-padded\n float *edge_weights_ptr = edge_weights.raw();\n float *rand_edge_weights_ptr = NULL;\n const int num_blocks_V2 = std::min( AMGX_GRID_MAX_SIZE, (num_nonzero_blocks - 1) / threads_per_block + 1);\n // Compute the edge weights\n#if AMGX_ASYNCCPU_PROOF_OF_CONCEPT\n int avoid_thrust_count = 1;//0;\n int newFindStrongest = 1;//0;\n int newWeights = 1;//0;\n#else\n int avoid_thrust_count = 0;//0;\n int newFindStrongest = 0;//0;\n int newWeights = 0;//0;\n#endif\n int usenumassignedassumption = false;\n\n if (newWeights == 0)\n {\n cudaFuncSetCacheConfig(computeEdgeWeightsBlockDiaCsr_V2, cudaFuncCachePreferL1);\n computeEdgeWeightsBlockDiaCsr_V2 <<< num_blocks_V2, threads_per_block, 0, stream>>>(\n A_row_offsets_ptr, A_row_indices_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, num_nonzero_blocks, edge_weights_ptr, rand_edge_weights_ptr, num_block_rows, A.get_block_dimy(), this->m_aggregation_edge_weight_component, this->weight_formula);\n cudaCheckError();\n }\n else\n {\n cudaFuncSetCacheConfig(computeEdgeWeightsBlockDiaCsr_V2_1, cudaFuncCachePreferL1);\n computeEdgeWeightsBlockDiaCsr_V2_1 <<< num_blocks_V2, threads_per_block, 0, stream>>>(\n A_row_offsets_ptr, A_row_indices_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, num_nonzero_blocks, edge_weights_ptr, rand_edge_weights_ptr, num_block_rows, A.get_block_dimy(), this->m_aggregation_edge_weight_component);\n cudaCheckError();\n }\n\n // -------------------------------------------------\n // First create aggregates of size 2\n // -------------------------------------------------\n int icount = 0;\n const int num_blocks_1024 = std::min( 13 * 2, (num_block_rows - 1) / 1024 + 1);\n device_vector_alloc sets_per_block_t(num_blocks_1024);\n int *sets_per_block = amgx::thrust::raw_pointer_cast(sets_per_block_t.data());\n cudaCheckError();\n\n do\n {\n if (newFindStrongest)\n {\n if (numUnassigned == num_block_rows && usenumassignedassumption)\n {\n cudaFuncSetCacheConfig(my_findStrongestNeighbourBlockDiaCsr_NoMerge, cudaFuncCachePreferL1);\n my_findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block, 0, stream>>>(\n A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, num_nonzero_blocks, 0, 0, strongest_neighbour_ptr, partner_index_ptr, 0, this->deterministic, 0, 0);\n cudaCheckError();\n }\n else\n {\n cudaFuncSetCacheConfig(my_findStrongestNeighbourBlockDiaCsr_NoMerge, cudaFuncCachePreferL1);\n my_findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block, 0, stream>>>(\n A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, num_nonzero_blocks, 0, 0, strongest_neighbour_ptr, partner_index_ptr, 0, this->deterministic, 0, 0);\n cudaCheckError();\n }\n }\n else\n {\n findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, partner_index_ptr, strongest_neighbour_ptr, this->deterministic);\n }\n\n cudaCheckError();\n // Look for perfect matches\n numUnassigned_previous = numUnassigned;\n\n if (avoid_thrust_count == 0)\n {\n matchEdges <<< num_blocks, threads_per_block, 0, stream>>>(num_block_rows, partner_index_ptr, aggregates_ptr, strongest_neighbour_ptr);\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(partner_index.begin(), partner_index.begin() + num_block_rows, -1);\n cudaCheckError();\n }\n else\n {\n my_MatchEdges <<< num_blocks_1024, 1024, 0, stream>>>(num_block_rows, partner_index_ptr, aggregates_ptr, strongest_neighbour_ptr, sets_per_block);\n cudaCheckError();\n numUnassigned = numUnassigned_previous - amgx::strided_reduction::count_block_results_pinned_memory(0, icount, num_blocks_1024, sets_per_block, amgx::strided_reduction::op_sum(), stream);\n }\n\n icount++;\n }\n while (!(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_block_rows < this->numUnassigned_tol || numUnassigned_previous == numUnassigned));\n\n assignUnassignedVertices <<< num_blocks, threads_per_block, 0, stream>>>(partner_index_ptr, num_block_rows);\n cudaCheckError();\n // -------------------------------------------------\n // Merge aggregates to create aggregates of size 4\n // -------------------------------------------------\n Vector > weight_strongest_neighbour(num_block_rows, -1);\n float *weight_strongest_neighbour_ptr = weight_strongest_neighbour.raw();\n // At this point, partner index contain either your index or your neighbours index, depending on weither you're matched or not\n // aggregates contain the largest vertex index of vertices in aggregate\n typename Matrix_d::IVector aggregated(num_block_rows, -1);\n IndexType *aggregated_ptr = aggregated.raw();\n // now used as flag to check if aggregated or not\n icount = 0;\n numUnassigned = num_block_rows;\n numUnassigned_previous = numUnassigned;\n\n do\n {\n // Each vertex stores in strongest_neighbour the aggregates number of strongest neighbour and the weight of connection\n if (newFindStrongest)\n {\n if (numUnassigned == num_block_rows && usenumassignedassumption)\n {\n cudaFuncSetCacheConfig(my_findStrongestNeighbourBlockDiaCsr_NoMerge, cudaFuncCachePreferL1);\n my_findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, num_nonzero_blocks, aggregated_ptr, aggregates_ptr, strongest_neighbour_ptr, partner_index_ptr, weight_strongest_neighbour_ptr, this->deterministic, 0, 0);\n cudaCheckError();\n }\n else\n {\n cudaFuncSetCacheConfig(my_findStrongestNeighbourBlockDiaCsr_NoMerge, cudaFuncCachePreferL1);\n my_findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, num_nonzero_blocks, aggregated_ptr, aggregates_ptr, strongest_neighbour_ptr, partner_index_ptr, weight_strongest_neighbour_ptr, this->deterministic, 0, 0);\n cudaCheckError();\n }\n }\n else\n {\n findStrongestNeighbourBlockDiaCsr_StoreWeight <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregated_ptr, aggregates_ptr, strongest_neighbour_ptr, partner_index_ptr, weight_strongest_neighbour_ptr, this->deterministic);\n }\n\n cudaCheckError();\n // Each vertex in same aggregates will agree on aggregates to propose too, and both store the aggregate number they want to match with\n agreeOnProposal <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, num_block_rows, aggregated_ptr, strongest_neighbour_ptr, weight_strongest_neighbour_ptr, partner_index_ptr, aggregates_ptr);\n cudaCheckError();\n numUnassigned_previous = numUnassigned;\n\n if (avoid_thrust_count == 0)\n {\n matchAggregatesSize4 <<< num_blocks, threads_per_block, 0, stream>>>(aggregates_ptr, aggregated_ptr, strongest_neighbour_ptr, partner_index_ptr, num_block_rows);\n cudaCheckError();\n numUnassigned = thrust_wrapper::count(aggregated.begin(), aggregated.end(), -1);\n }\n else\n {\n my_matchAggregatesSize4 <<< num_blocks_1024, 1024, 0, stream>>>(aggregates_ptr, aggregated_ptr, strongest_neighbour_ptr, partner_index_ptr, num_block_rows, sets_per_block);\n cudaCheckError();\n numUnassigned = numUnassigned_previous - amgx::strided_reduction::count_block_results_pinned_memory(1, icount, num_blocks_1024, sets_per_block, amgx::strided_reduction::op_sum(), stream);\n }\n\n cudaCheckError();\n icount++;\n }\n while (!(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_block_rows < this->numUnassigned_tol || numUnassigned == numUnassigned_previous) );\n\n assignUnassignedVertices_2 <<< num_blocks, threads_per_block, 0, stream>>>(partner_index_ptr, num_block_rows);\n cudaCheckError();\n // -------------------------------------------------\n // Merge aggregates to create aggregates of size 8\n // -------------------------------------------------\n thrust_wrapper::fill(aggregated.begin(), aggregated.end(), -1);\n cudaCheckError();\n thrust_wrapper::fill(weight_strongest_neighbour.begin(), weight_strongest_neighbour.end(), -1.);\n cudaCheckError();\n icount = 0;\n numUnassigned = num_block_rows;\n numUnassigned_previous = numUnassigned;\n\n do\n {\n // Each vertex stores in strongest_neighbour the aggregates number of strongest neighbour and the weight of connection\n if (newFindStrongest)\n {\n if (numUnassigned == num_block_rows && usenumassignedassumption)\n {\n cudaFuncSetCacheConfig(my_findStrongestNeighbourBlockDiaCsr_NoMerge, cudaFuncCachePreferL1);\n my_findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, num_nonzero_blocks, aggregated_ptr, aggregates_ptr, strongest_neighbour_ptr, partner_index_ptr, weight_strongest_neighbour_ptr, this->deterministic, 0, 0);\n cudaCheckError();\n }\n else\n {\n cudaFuncSetCacheConfig(my_findStrongestNeighbourBlockDiaCsr_NoMerge, cudaFuncCachePreferL1);\n my_findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, num_nonzero_blocks, aggregated_ptr, aggregates_ptr, strongest_neighbour_ptr, partner_index_ptr, weight_strongest_neighbour_ptr, this->deterministic, 0, 0);\n cudaCheckError();\n }\n }\n else\n {\n findStrongestNeighbourBlockDiaCsr_StoreWeight_2 <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregated_ptr, aggregates_ptr, strongest_neighbour_ptr, partner_index_ptr, weight_strongest_neighbour_ptr);\n }\n\n cudaCheckError();\n\n // Each vertex in same aggregates will agree on aggregates to propose too, and both store the aggregate number they want to match with\n if (!this->deterministic)\n {\n agreeOnProposal_2 <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, num_block_rows, aggregated_ptr, strongest_neighbour_ptr, weight_strongest_neighbour_ptr, partner_index_ptr, aggregates_ptr, this->deterministic);\n cudaCheckError();\n }\n else\n {\n //strongest_neighbour_tmp = strongest_neighbour; // copied that directly in the kernel\n agreeOnProposal_2_deterministic <<< num_blocks, threads_per_block, 0, stream>>>(\n strongest_neighbour_tmp.raw(),\n A_row_offsets_ptr,\n A_column_indices_ptr, num_block_rows,\n aggregated_ptr, strongest_neighbour_ptr,\n weight_strongest_neighbour_ptr, partner_index_ptr, aggregates_ptr, this->deterministic);\n cudaCheckError();\n strongest_neighbour_tmp.swap(strongest_neighbour);\n strongest_neighbour_ptr = strongest_neighbour.raw(); //re-saving the correct pointer..\n }\n\n cudaCheckError();\n numUnassigned_previous = numUnassigned;\n\n if (avoid_thrust_count == 0)\n {\n matchAggregates <<< num_blocks, threads_per_block, 0, stream>>>(aggregates_ptr, aggregated_ptr, strongest_neighbour_ptr, num_block_rows);\n cudaCheckError();\n numUnassigned = thrust_wrapper::count(aggregated.begin(), aggregated.end(), -1);\n }\n else\n {\n my_matchAggregates <<< num_blocks_1024, 1024, 0, stream>>>(aggregates_ptr, aggregated_ptr, strongest_neighbour_ptr, num_block_rows, sets_per_block);\n cudaCheckError();\n numUnassigned = numUnassigned_previous - amgx::strided_reduction::count_block_results_pinned_memory(2, icount, num_blocks_1024, sets_per_block, amgx::strided_reduction::op_sum(), stream);\n }\n\n cudaCheckError();\n icount++;\n }\n while (!(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_block_rows < this->numUnassigned_tol || numUnassigned_previous == numUnassigned) );\n\n // Merge remaining vertices with current aggregates\n int local_iter = 0;\n\n if (!this->deterministic)\n {\n while (numUnassigned != 0)\n {\n mergeWithExistingAggregatesBlockDiaCsr <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, aggregated_ptr, this->deterministic, (IndexType *) NULL, local_iter > 1);\n cudaCheckError();\n numUnassigned_previous = numUnassigned;\n numUnassigned = (int)thrust_wrapper::count(aggregated.begin(), aggregated.end(), -1);\n cudaCheckError();\n local_iter++;\n }\n }\n else\n {\n typename Matrix_d::IVector aggregates_candidate(num_block_rows, -1);\n\n while (numUnassigned != 0)\n {\n // allow singletons only from the 2nd local iteration\n mergeWithExistingAggregatesBlockDiaCsr <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, aggregated_ptr, this->deterministic, aggregates_candidate.raw(), local_iter > 1);\n cudaCheckError();\n numUnassigned_previous = numUnassigned;\n\n if (avoid_thrust_count == 0)\n {\n joinExistingAggregates <<< num_blocks, threads_per_block, 0, stream>>>(num_block_rows, aggregates_ptr, aggregated_ptr, aggregates_candidate.raw());\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(aggregated.begin(), aggregated.end(), -1);\n }\n else\n {\n my_joinExistingAggregates <<< num_blocks_1024, 1024, 0, stream>>>(num_block_rows, aggregates_ptr, aggregated_ptr, aggregates_candidate.raw(), sets_per_block);\n cudaCheckError();\n numUnassigned = numUnassigned_previous - amgx::strided_reduction::count_block_results_pinned_memory(3, local_iter, num_blocks_1024, sets_per_block, amgx::strided_reduction::op_sum(), stream);\n }\n\n cudaCheckError();\n local_iter++;\n }\n\n aggregates_candidate.resize(0);\n }\n\n this->renumberAndCountAggregates(aggregates, aggregates_global, num_block_rows, num_aggregates);\n //analyze_coloring(aggregates, A.getMatrixColoring().getRowColors());\n}\n\ntemplate \nvoid Size8SelectorBase::setAggregates(Matrix &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n if (A.get_block_dimx() == A.get_block_dimy())\n {\n setAggregates_common_sqblock( A, aggregates, aggregates_global, num_aggregates );\n }\n else\n {\n FatalError(\"Unsupported block size for Size8Selector\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\n// -------------------------\n// Explict instantiations\n// -------------------------\n#define AMGX_CASE_LINE(CASE) template class Size8SelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class Size8Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n}\n}\n}\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include //count\n#include //sort\n#include //unique\n#include //remove\n#include //transform_inclusive_scan\n\n#include //offsets_to_indices\n\n#include \n#include \nnamespace amgx\n{\n\n\nnamespace strided_reduction\n{\ntemplate\nvoid count_block_results_pinned_memory(scalar_t *out_host, const int n_blocks, scalar_t *out_d, const OP &op = OP(), hipStream_t stream = 0)\n{\n strided_reduction_collect_partials <<< 1, 32, 0, stream>>>(out_host, out_d, n_blocks);\n cudaCheckError();\n}\n\n\ntemplate\nscalar_t count_block_results_pinned_memory(const int a, const int i, const int n_blocks, scalar_t *out_d, const OP &op = OP(), hipStream_t stream = 0) //STRIDE=1 case\n{\n static scalar_t *ret = 0;\n static hipEvent_t throttle_event = 0;\n const int buffers = 1;\n\n if (ret == 0)\n {\n hipError_t cuda_rc = amgx::memory::hipHostMalloc((void **)&ret, buffers * sizeof(scalar_t));\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipHostMalloc failed in size8_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n ret[0] = 0;\n hipEventCreateWithFlags(&throttle_event, hipEventDisableTiming);\n }\n\n int ib = i % buffers;\n count_block_results_pinned_memory<1, scalar_t, OP>(ret + ib, n_blocks, out_d, op, stream);\n\n if (ib == buffers - 1)\n {\n hipError_t cuda_rc = hipEventRecord(throttle_event);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipEventRecord failed in size8_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = hipEventSynchronize(throttle_event);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipEventSynchronize failed in size8_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n scalar_t tot = 0;\n\n for (int j = 0; j < buffers; j++)\n {\n tot += ret[j];\n }\n\n return tot + buffers - 1;\n }\n else\n {\n return -1;\n }\n}\n\n\n}\n\nvoid analyze_coloring(device_vector_alloc aggregates_d, device_vector_alloc colors_d);\n\nnamespace aggregation\n{\nnamespace size8_selector\n{\n\n#include \n\ntemplate \n__global__ __launch_bounds__(256, 4)\nvoid my_findStrongestNeighbourBlockDiaCsr_NoMergeClean(\n const IndexType *row_offsets, const IndexType *column_indices,\n const float *edge_weights, const IndexType num_block_rows, const IndexType num_nonzero,\n IndexType *aggregated, IndexType *aggregates, int *strongest_neighbour,\n IndexType *partner_index, float *weight_strongest_neighbour, int deterministic,\n const IndexType *unassigned_rows,\n const int num_unassigned_row)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n const int lane_id = utils::lane_id();\n bool valid_tid = false;\n\n for (; utils::any(valid_tid = tid < num_block_rows); tid += gridDim.x * blockDim.x)\n {\n int jmin = -NUM_COLS * 2, jmax = -NUM_COLS * 4;\n float weight;\n int jcol;\n float max_weight_unaggregated = 0;\n int strongest_unaggregated = -1;\n bool is_unassigned = false;\n\n if (valid_tid)\n {\n is_unassigned = (__ldcs(partner_index + tid) == -1);\n }\n\n if (is_unassigned) // Unaggregated row\n {\n jmin = __ldca(row_offsets + tid);\n jmax = __ldca(row_offsets + tid + 1);\n }\n\n if (utils::any(is_unassigned))\n {\n int jj = jmin - amgx::strided_reduction::warp_loader::align_shift(jmin);\n\n for (; utils::any(jj < jmax, utils::activemask()); jj += NUM_COLS)\n {\n int I[NUM_COLS];\n float W[NUM_COLS];\n int P[NUM_COLS];\n amgx::strided_reduction::warp_loader::load(column_indices, jj, num_nonzero, I);\n\n#pragma unroll\n for (int i = 0; i < NUM_COLS; i++)\n {\n int j = jj + i;\n jcol = I[i];\n\n if (j >= jmin && j < jmax)\n {\n P[i] = __ldg(partner_index + jcol); //make this load ASAP\n }\n }\n\n amgx::strided_reduction::warp_loader::load(edge_weights, jj, num_nonzero, W);\n\n#pragma unroll\n for (int i = 0; i < NUM_COLS; i++)\n {\n weight = W[i];\n jcol = I[i];\n int j = jj + i;\n\n if (j >= jmin && j < jmax)\n {\n if (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && jcol > strongest_unaggregated)) // unaggregated\n {\n if (tid != jcol && P[i] == -1)\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n }\n }\n }\n }\n\n if (strongest_unaggregated == -1) // All neighbours are aggregated\n {\n // Put in its own aggregate\n if (!deterministic && is_unassigned)\n {\n partner_index[tid] = tid;\n }\n }\n else\n {\n strongest_neighbour[tid] = strongest_unaggregated;\n }\n }\n }\n}\n\n\n\n#define ALGORITHM_NOMERGE 0\n#define ALGORITHM_STOREWEIGHTS 1\n#define ALGORITHM_STOREWEIGHTS_2 2\n\n\ntemplate \n__global__ __launch_bounds__(256, 4)\nvoid my_findStrongestNeighbourBlockDiaCsr_NoMerge(\n const IndexType *row_offsets, const IndexType *column_indices,\n const float *edge_weights, const IndexType num_block_rows, const IndexType num_nonzero,\n IndexType *aggregated, IndexType *aggregates, int *strongest_neighbour,\n IndexType *partner_index, float *weight_strongest_neighbour, int deterministic,\n const IndexType *n_unassigned_per_block, const IndexType *unassigned_per_block\n //const int num_unassigned_row\n //const IndexType *unassigned_rows,\n)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int bid = blockIdx.x;\n const int lane_id = utils::lane_id();\n bool valid_tid = false;\n\n for (; utils::any( valid_tid = tid < num_block_rows); tid += gridDim.x * blockDim.x)\n {\n int jmin = -NUM_COLS * 2, jmax = -NUM_COLS * 4;\n float weight;\n int jcol;\n float max_weight_unaggregated = 0;\n int strongest_unaggregated = -1;\n float max_weight_aggregated = 0.;\n int strongest_aggregated = -1;\n int partner = -1;\n int partner0, partner1, partner2;\n int agg_jcol;\n bool is_unassigned = false;\n int rowi = -1;\n\n if (LOAD_ONLY_UNASSIGNED)\n {\n if (valid_tid)\n {\n rowi = unassigned_per_block[tid];//unassigned_per_block+bid*256+threadIdx.x);\n is_unassigned = (__ldg(partner_index + rowi) == -1);\n }\n\n if (is_unassigned)\n {\n jmin = __ldg(row_offsets + rowi);\n jmax = __ldg(row_offsets + rowi + 1);\n }\n }\n else\n {\n rowi = tid;\n\n if (ALGORITHM == ALGORITHM_NOMERGE)\n {\n if (valid_tid) { is_unassigned = (__ldcs(partner_index + tid) == -1); }\n }\n else //ALGORITHM_STOREWEIGHTS or ALGORITHM_STOREWEIGHTS_2\n {\n if (valid_tid) { is_unassigned = (__ldcs(aggregated + tid) == -1); }\n }\n\n if (is_unassigned) // mind the else above\n {\n jmin = __ldg(row_offsets + rowi);\n jmax = __ldlu(row_offsets + rowi + 1);\n }\n }\n\n if (utils::any(is_unassigned))\n {\n if (is_unassigned) // Unaggregated row\n {\n if (ALGORITHM == ALGORITHM_STOREWEIGHTS)\n {\n partner = partner_index[rowi];\n }\n else if (ALGORITHM == ALGORITHM_STOREWEIGHTS_2)\n {\n partner0 = partner_index[rowi];\n partner1 = partner_index[num_block_rows + rowi];\n partner2 = partner_index[2 * num_block_rows + rowi];\n }\n }\n\n int jj = jmin - amgx::strided_reduction::warp_loader::align_shift(jmin);\n\n for (; utils::any(jj < jmax && jmax >= 0); jj += NUM_COLS)\n {\n int I[NUM_COLS];\n float W[NUM_COLS];\n int P[NUM_COLS];\n int jj_ok = (jj >= 0 && jj < jmax && jmax >= 0) ? jj : 0;\n\n amgx::strided_reduction::warp_loader::load(column_indices, jj_ok, num_nonzero, I);\n#pragma unroll\n\n for (int i = 0; i < NUM_COLS; i++)\n {\n int j = jj + i;\n jcol = I[i];\n\n if (j >= jmin && j < jmax)\n {\n if (ALGORITHM == ALGORITHM_NOMERGE)\n {\n P[i] = __ldg(partner_index + jcol); //make this load ASAP\n }\n else\n {\n P[i] = __ldg(aggregated + jcol);\n }\n }\n }\n\n amgx::strided_reduction::warp_loader::load(edge_weights, jj_ok, num_nonzero, W);\n#pragma unroll\n\n for (int i = 0; i < NUM_COLS; i++)\n {\n weight = W[i];\n jcol = I[i];\n int j = jj + i;\n\n if (j >= jmin && j < jmax)\n {\n if (ALGORITHM == ALGORITHM_NOMERGE)\n {\n if (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && jcol > strongest_unaggregated)) // unaggregated\n {\n if (rowi != jcol && P[i] == -1)\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n }\n }\n else\n {\n bool partner_condition;\n\n if (ALGORITHM == ALGORITHM_STOREWEIGHTS)\n {\n partner_condition = jcol != partner;\n }\n else if (ALGORITHM == ALGORITHM_STOREWEIGHTS_2)\n {\n partner_condition = jcol != partner0 && jcol != partner1 && jcol != partner2;\n }\n\n agg_jcol = P[i];\n\n if (partner_condition && rowi != jcol)\n {\n if (agg_jcol == -1 && (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && jcol > strongest_unaggregated))) // unaggregated\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n else if (agg_jcol != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) // unaggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n }\n }\n }\n }\n\n if (valid_tid && is_unassigned)\n {\n if (ALGORITHM == ALGORITHM_NOMERGE)\n {\n if (strongest_unaggregated == -1) // All neighbours are aggregated\n {\n // Put in its own aggregate\n if (!deterministic)\n {\n partner_index[rowi] = rowi;\n }\n }\n else\n {\n strongest_neighbour[rowi] = strongest_unaggregated;\n }\n }\n else if (ALGORITHM == ALGORITHM_STOREWEIGHTS)\n {\n if (strongest_unaggregated == -1) // All neighbours are aggregated\n {\n if (!deterministic)\n {\n if (strongest_aggregated != -1)\n {\n aggregates[tid] = aggregates[strongest_aggregated];\n aggregated[tid] = 1;\n aggregates[partner] = aggregates[strongest_aggregated];\n aggregated[partner] = 1;\n }\n else // leave in its own aggregate\n {\n aggregated[partner] = 1;\n aggregated[tid] = 1;\n }\n }\n }\n else // Found an unaggregated aggregate\n {\n weight_strongest_neighbour[tid] = max_weight_unaggregated;\n strongest_neighbour[tid] = aggregates[strongest_unaggregated];\n }\n }\n else if (ALGORITHM == ALGORITHM_STOREWEIGHTS_2)\n {\n if (strongest_unaggregated == -1 && strongest_aggregated != -1) // all neighbours are aggregated, store the strongest aggregated\n {\n weight_strongest_neighbour[tid] = -max_weight_aggregated;\n strongest_neighbour[tid] = aggregates[strongest_aggregated];\n }\n else if (strongest_unaggregated != -1)\n {\n weight_strongest_neighbour[tid] = max_weight_unaggregated;\n strongest_neighbour[tid] = aggregates[strongest_unaggregated];\n }\n }\n }\n }\n\n bid += gridDim.x;\n }\n}\n\n\n\n\n#define INST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(numcols, algo,assume,c) template __global__ void my_findStrongestNeighbourBlockDiaCsr_NoMerge(\\\nconst int *row_offsets, const int *column_indices,\\\nconst float *edge_weights, const int num_block_rows,const int num_nonzero,\\\nint *aggregated, int *aggregates, int *strongest_neighbour,\\\nint *partner_index, float *weight_strongest_neighbour, int deterministic,\\\nconst int* n_unassigned_per_block, const int * unassigned_per_block);\n\n#define INSTmy_findStrongestNeighbourBlockDiaCsr_NoMergeClean(numcols) template __global__ void my_findStrongestNeighbourBlockDiaCsr_NoMergeClean(\\\nconst int *row_offsets, const int *column_indices,\\\nconst float *edge_weights, const int num_block_rows,const int num_nonzero,\\\nint *aggregated, int *aggregates, int *strongest_neighbour,\\\nint *partner_index, float *weight_strongest_neighbour, int deterministic,\\\nconst int *unassigned_rows,\\\nconst int num_unassigned_row);\n\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(1, ALGORITHM_NOMERGE, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(1, ALGORITHM_STOREWEIGHTS, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(1, ALGORITHM_STOREWEIGHTS_2, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(1, ALGORITHM_NOMERGE, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(1, ALGORITHM_STOREWEIGHTS, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(1, ALGORITHM_STOREWEIGHTS_2, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(1, ALGORITHM_NOMERGE, 0, 1) //load only unassigned\n\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(2, ALGORITHM_NOMERGE, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(2, ALGORITHM_STOREWEIGHTS, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(2, ALGORITHM_STOREWEIGHTS_2, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(2, ALGORITHM_NOMERGE, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(2, ALGORITHM_STOREWEIGHTS, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(2, ALGORITHM_STOREWEIGHTS_2, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(2, ALGORITHM_NOMERGE, 0, 1) //load only unassigned\n\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(4, ALGORITHM_NOMERGE, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(4, ALGORITHM_STOREWEIGHTS, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(4, ALGORITHM_STOREWEIGHTS_2, 0, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(4, ALGORITHM_NOMERGE, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(4, ALGORITHM_STOREWEIGHTS, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(4, ALGORITHM_STOREWEIGHTS_2, 1, 0)\nINST_my_findStrongestNeighbourBlockDiaCsr_NoMerge(4, ALGORITHM_NOMERGE, 0, 1) //load only unassigned\n\n\nINSTmy_findStrongestNeighbourBlockDiaCsr_NoMergeClean(1)\nINSTmy_findStrongestNeighbourBlockDiaCsr_NoMergeClean(2)\nINSTmy_findStrongestNeighbourBlockDiaCsr_NoMergeClean(4)\n\ntemplate\n__global__ void my_blockCompact(\n int *partner_index, const int num_rows,\n int *unassigned_per_block_in,\n int *n_unassigned_per_block, int *unassigned_per_block)\n{\n int bid = blockIdx.x; //RMV\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n for (; bid < num_rows / 256 + 1/*__any(tid < num_rows)*/; tid += gridDim.x * blockDim.x)\n {\n int row = tid;\n bool no_partner = 0; //RMV\n\n if (tid < num_rows)\n {\n if (ALREADY_COMPACT)\n {\n row = unassigned_per_block_in[tid];\n }\n\n if (partner_index[row] == -1) // Unaggregated row\n {\n no_partner = 1;\n }\n }\n\n amgx::strided_reduction::block_binary_compaction<256, 32, 1>(\n n_unassigned_per_block, unassigned_per_block, bid,\n no_partner, row);\n bid += gridDim.x;\n }\n}\n\ntemplate __global__ void my_blockCompact<0>(\n int *partner_index, const int num_rows,\n int *unassigned_per_block_in,\n int *n_unassigned_per_block, int *unassigned_per_block);\ntemplate __global__ void my_blockCompact<1>(\n int *partner_index, const int num_rows,\n int *unassigned_per_block_in,\n int *n_unassigned_per_block, int *unassigned_per_block);\n\n__global__ void my_MatchEdgesWithCompaction(const int num_rows, int *partner_index, int *aggregates, const int *strongest_neighbour, int *sets_per_block,\n int *unassigned_per_block_in, int *n_unassigned_per_block, int *unassigned_per_block\n )\n{\n int potential_match, potential_match_neighbour;\n int warp_count = 0;\n int bid = blockIdx.x; //RMV\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n for (; bid < num_rows / 1024 + 1/*__any(tid < num_rows)*/; tid += gridDim.x * blockDim.x)\n {\n int row = tid;\n bool no_partner = 0;\n\n if (tid < num_rows)\n {\n if (partner_index[row] == -1) // Unaggregated row\n {\n no_partner = 1;\n potential_match = strongest_neighbour[row];\n\n if (potential_match != -1)\n {\n potential_match_neighbour = __ldcs(strongest_neighbour + potential_match);\n\n if ( potential_match_neighbour == row ) // we have a match\n {\n no_partner = 0;\n //partner_notnull = 1;//RMV\n partner_index[row] = potential_match;\n aggregates[row] = ( potential_match > row ) ? row : potential_match;\n }\n }\n }\n }\n\n amgx::strided_reduction::block_binary_compaction<1024, 32, 1>(\n n_unassigned_per_block, unassigned_per_block, bid,\n no_partner, row);\n warp_count += amgx::strided_reduction::warp_binary_count(no_partner);\n bid += gridDim.x;\n }\n\n amgx::strided_reduction::block_count<1, 1024, 32, int>(warp_count, sets_per_block);\n}\n\n__global__ void my_MatchEdges(const int num_rows, int *partner_index, int *aggregates, const int *strongest_neighbour, int *sets_per_block)\n{\n int potential_match, potential_match_neighbour;\n int warp_count = 0;\n\n for (int tid = threadIdx.x + blockDim.x * blockIdx.x; utils::any(tid < num_rows); tid += gridDim.x * blockDim.x)\n {\n bool has_set_partner_index = 0;\n\n if (tid < num_rows)\n {\n if (partner_index[tid] == -1) // Unaggregated row\n {\n potential_match = strongest_neighbour[tid];\n\n if (potential_match != -1)\n {\n potential_match_neighbour = __ldcs(strongest_neighbour + potential_match);\n\n if ( potential_match_neighbour == tid ) // we have a match\n {\n has_set_partner_index = 1;\n partner_index[tid] = potential_match;\n aggregates[tid] = ( potential_match > tid) ? tid : potential_match;\n }\n }\n }\n }\n\n warp_count += amgx::strided_reduction::warp_binary_count(has_set_partner_index);\n }\n\n amgx::strided_reduction::block_count<1, 1024, 32, int>(warp_count, sets_per_block);\n}\n\n\n// matchEdges\n__global__ void my_joinExistingAggregates(int num_rows, int *aggregates, int *aggregated, int *aggregates_candidate, int *sets_per_block)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int warp_count = 0;\n\n while (utils::any(tid < num_rows))\n {\n bool has_set_partner_index = 0;\n\n if (tid < num_rows)\n {\n if (aggregated[tid] == -1 && aggregates_candidate[tid] != -1) // Unaggregated row\n {\n aggregates[tid] = aggregates_candidate[tid];\n aggregated[tid] = 1;\n has_set_partner_index = 1;\n }\n }\n\n warp_count += amgx::strided_reduction::warp_binary_count(has_set_partner_index);\n tid += gridDim.x * blockDim.x;\n }\n\n amgx::strided_reduction::block_count<1, 1024, 32, int>(warp_count, sets_per_block);\n}\n\n\n// Kernel that checks if perfect matchs exist\n__global__ void my_matchAggregates(int *aggregates, int *aggregated, int *strongest_neighbour, const int num_rows, int *sets_per_block)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int potential_match, potential_match_neighbour, my_aggregate;\n int warp_count = 0;\n\n while (utils::any(tid < num_rows))\n {\n bool has_set_partner_index = 0;\n\n if (tid < num_rows) if (aggregated[tid] == -1) // Unaggregated row\n {\n potential_match = strongest_neighbour[tid];\n\n if (potential_match != -1)\n {\n potential_match_neighbour = __ldcs(strongest_neighbour + potential_match); //or global\n my_aggregate = aggregates[tid];\n\n if (potential_match_neighbour == my_aggregate) // we have a match\n {\n has_set_partner_index = 1;\n aggregated[tid] = 1;\n aggregates[tid] = ( potential_match > my_aggregate) ? my_aggregate : potential_match;\n }\n }\n }\n\n warp_count += amgx::strided_reduction::warp_binary_count(has_set_partner_index);\n tid += gridDim.x * blockDim.x;\n }\n\n amgx::strided_reduction::block_count<1, 1024, 32, int>(warp_count, sets_per_block);\n}\n\n\n// Kernel that checks if perfect matchs exist\n__global__ void my_matchAggregatesSize4(int *aggregates, int *aggregated, int *strongest_neighbour, int *partner_index, const int num_rows, int *sets_per_block)\n{\n int potential_match, potential_match_neighbour, my_aggregate;\n int warp_count = 0;\n\n for (int tid = threadIdx.x + blockDim.x * blockIdx.x; utils::any(tid < num_rows); tid += blockDim.x * gridDim.x)\n {\n bool has_set_partner_index = 0;\n\n if (tid < num_rows) if (aggregated[tid] == -1) // Unaggregated row\n {\n potential_match = strongest_neighbour[tid];\n\n if (potential_match != -1)\n {\n potential_match_neighbour = __ldcs(strongest_neighbour + potential_match);\n my_aggregate = aggregates[tid];\n\n if (potential_match_neighbour == my_aggregate) // we have a match\n {\n has_set_partner_index = 1;\n aggregated[tid] = 1;\n aggregates[tid] = ( potential_match > my_aggregate) ? my_aggregate : potential_match;\n partner_index[tid + num_rows] = potential_match;\n partner_index[tid + 2 * num_rows] = partner_index[potential_match];\n }\n }\n }\n\n warp_count += amgx::strided_reduction::warp_binary_count(has_set_partner_index);\n }\n\n amgx::strided_reduction::block_count<1, 1024, 32, int>(warp_count, sets_per_block);\n}\n\n// ------------------------\n// Kernels\n// ------------------------\n\n// findStrongestNeighbour kernel for block_dia_csr_matrix format\n// Reads the weight from edge_weights array\ntemplate \n__global__\nvoid findStrongestNeighbourBlockDiaCsr_StoreWeight_2(const IndexType *row_offsets, const IndexType *column_indices,\n const float *edge_weights, const IndexType num_block_rows, IndexType *aggregated, IndexType *aggregates, int *strongest_neighbour, IndexType *partner_index, float *weight_strongest_neighbour)\n{\n float weight;\n int jcol, jmin, jmax;\n int partner0, partner1, partner2;\n int agg_jcol;\n\n for (int tid = threadIdx.x + blockDim.x * blockIdx.x; tid < num_block_rows; tid += blockDim.x * gridDim.x)\n {\n float max_weight_unaggregated = 0.;\n float max_weight_aggregated = 0.;\n int strongest_unaggregated = -1;\n int strongest_aggregated = -1;\n\n if (aggregated[tid] == -1) // Unaggregated row\n {\n partner0 = partner_index[tid];\n partner1 = partner_index[num_block_rows + tid];\n partner2 = partner_index[2 * num_block_rows + tid];\n jmin = row_offsets[tid];\n jmax = row_offsets[tid + 1];\n\n for (int j = jmin; j < jmax; j++)\n {\n jcol = column_indices[j];\n\n if (jcol == tid || jcol >= num_block_rows) { continue; }\n\n weight = edge_weights[j];\n agg_jcol = aggregated[jcol];\n\n if (jcol != partner0 && jcol != partner1 && jcol != partner2)\n {\n if (agg_jcol == -1 && (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && jcol > strongest_unaggregated))) // unaggregated\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n else if (agg_jcol != -1 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && jcol > strongest_aggregated))) // unaggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n }\n\n if (strongest_unaggregated == -1 && strongest_aggregated != -1) // all neighbours are aggregated, store the strongest aggregated\n {\n weight_strongest_neighbour[tid] = -max_weight_aggregated;\n strongest_neighbour[tid] = aggregates[strongest_aggregated];\n }\n else if (strongest_unaggregated != -1)\n {\n weight_strongest_neighbour[tid] = max_weight_unaggregated;\n strongest_neighbour[tid] = aggregates[strongest_unaggregated];\n }\n }\n }\n}\n\n// findStrongestNeighbour kernel for block_dia_csr_matrix format\n// Reads the weight from edge_weights array\ntemplate \n__global__\nvoid agreeOnProposal_2(const IndexType *row_offsets, const IndexType *column_indices,\n IndexType num_block_rows, IndexType *aggregated, int *strongest_neighbour, float *weight_strongest_neighbour, IndexType *partner_index, int *aggregates, int deterministic)\n{\n int partner[3];\n float weight[3];\n\n for (int tid = threadIdx.x + blockDim.x * blockIdx.x; tid < num_block_rows; tid += blockDim.x * gridDim.x)\n {\n int strongest_partner = -1;\n float my_weight = 0.;\n\n if (aggregated[tid] == -1)\n {\n my_weight = weight_strongest_neighbour[tid];\n#pragma unroll\n\n for (int m = 0; m < 3; m++)\n {\n partner[m] = partner_index[tid + m * num_block_rows];\n }\n\n#pragma unroll\n\n for (int m = 0; m < 3; m++)\n {\n weight[m] = weight_strongest_neighbour[partner[m]];\n }\n\n#pragma unroll\n\n for (int m = 0; m < 3; m++)\n {\n //if (weight[m] > my_weight && weight[m] > 0.) // there is a partner that has an unaggregated neighbour\n if (weight[m] > my_weight)\n {\n if (weight[m] > 0.) // there is a partner that has an unaggregated neighbour\n {\n strongest_partner = m;\n my_weight = weight[m];\n }\n }\n else if (weight[m] < my_weight) // there is a partner without an unaggregated neighbour, whose neighbour is stronger than mine\n {\n if (my_weight < 0.)\n {\n strongest_partner = m;\n my_weight = weight[m];\n }\n }\n }\n\n if (my_weight < 0.) // means all neighbours of vertices in aggregate are aggregated, merge to another aggregate\n {\n if (!deterministic)\n {\n aggregated[tid] = 1;\n aggregates[tid] = strongest_partner != -1 ? strongest_neighbour[partner[strongest_partner]] : strongest_neighbour[tid];\n }\n }\n else if (strongest_partner != -1) // store my partner's pick\n {\n strongest_neighbour[tid] = strongest_neighbour[partner[strongest_partner]];\n }\n }\n }\n}\n\ntemplate \n__global__\nvoid agreeOnProposal_2_deterministic(int *strongest_neighbour_out, const IndexType *row_offsets, const IndexType *column_indices,\n IndexType num_block_rows, IndexType *aggregated, int *strongest_neighbour, float *weight_strongest_neighbour, IndexType *partner_index, int *aggregates, int deterministic)\n{\n int partner[3];\n float weight[3];\n\n for (int tid = threadIdx.x + blockDim.x * blockIdx.x; tid < num_block_rows; tid += blockDim.x * gridDim.x)\n {\n int strongest_partner = -1;\n float my_weight = 0.;\n //copy here to avoid redundant copies before launching the kernel\n int new_strongest_neighbour_out = strongest_neighbour[tid];\n\n if (aggregated[tid] == -1)\n {\n my_weight = weight_strongest_neighbour[tid];\n#pragma unroll\n\n for (int m = 0; m < 3; m++)\n {\n partner[m] = partner_index[tid + m * num_block_rows];\n }\n\n#pragma unroll\n\n for (int m = 0; m < 3; m++)\n {\n weight[m] = weight_strongest_neighbour[partner[m]];\n }\n\n#pragma unroll\n\n for (int m = 0; m < 3; m++)\n {\n if (weight[m] > my_weight)\n {\n if (weight[m] > 0.) // there is a partner that has an unaggregated neighbour\n {\n strongest_partner = m;\n my_weight = weight[m];\n }\n }\n else if (weight[m] < my_weight) // there is a partner without an unaggregated neighbour, whose neighbour is stronger than mine\n {\n if (my_weight < 0.)\n {\n strongest_partner = m;\n my_weight = weight[m];\n }\n }\n }\n\n if (my_weight < 0.) // means all neighbours of vertices in aggregate are aggregated, merge to another aggregate\n {\n if (!deterministic)\n {\n aggregated[tid] = 1;\n aggregates[tid] = strongest_partner != -1 ? strongest_neighbour[partner[strongest_partner]] : strongest_neighbour[tid];\n }\n }\n else if (strongest_partner != -1) // store my partner's pick\n {\n new_strongest_neighbour_out = strongest_neighbour[partner[strongest_partner]];\n }\n }\n\n //copy here to avoid redundant copies before launching the kernel\n strongest_neighbour_out[tid] = new_strongest_neighbour_out;\n }\n}\n\n\n// Kernel that checks if perfect matchs exist\ntemplate \n__global__\nvoid matchAggregatesSize4(IndexType *aggregates, IndexType *aggregated, IndexType *strongest_neighbour, IndexType *partner_index, const IndexType num_rows)\n{\n int potential_match, potential_match_neighbour, my_aggregate;\n\n for (int tid = threadIdx.x + blockDim.x * blockIdx.x; tid < num_rows; tid += blockDim.x * gridDim.x)\n {\n if (aggregated[tid] == -1) // Unaggregated row\n {\n potential_match = strongest_neighbour[tid];\n\n if (potential_match != -1)\n {\n potential_match_neighbour = strongest_neighbour[potential_match];\n my_aggregate = aggregates[tid];\n\n if (potential_match_neighbour == my_aggregate) // we have a match\n {\n aggregated[tid] = 1;\n aggregates[tid] = ( potential_match > my_aggregate) ? my_aggregate : potential_match;\n partner_index[tid + num_rows] = potential_match;\n partner_index[tid + 2 * num_rows] = partner_index[potential_match];\n }\n }\n }\n }\n}\n\n\ntemplate \n__global__\nvoid assignUnassignedVertices_2(IndexType *partner_index, const IndexType num_rows)\n{\n for (int tid = threadIdx.x + blockDim.x * blockIdx.x; tid < num_rows; tid += blockDim.x * gridDim.x)\n {\n if (partner_index[num_rows + tid] == -1) // Unaggregated row\n {\n partner_index[num_rows + tid] = tid;\n }\n\n if (partner_index[2 * num_rows + tid] == -1) // Unaggregated row\n {\n partner_index[2 * num_rows + tid] = tid;\n }\n }\n}\n\n\n// -----------------\n// Methods\n// ----------------\n\n// Constructor\ntemplate\nSize8SelectorBase::Size8SelectorBase(AMG_Config &cfg, const std::string &cfg_scope)\n{\n deterministic = cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\");\n max_iterations = cfg.AMG_Config::template getParameter(\"max_matching_iterations\", cfg_scope);\n numUnassigned_tol = cfg.AMG_Config::template getParameter(\"max_unassigned_percentage\", cfg_scope);\n m_aggregation_edge_weight_component = cfg.AMG_Config::template getParameter(\"aggregation_edge_weight_component\", cfg_scope);\n weight_formula = cfg.AMG_Config::template getParameter(\"weight_formula\", cfg_scope);\n}\n\n\n// setAggregates for block_dia_csr_matrix_h format\ntemplate \nvoid Size8Selector >::setAggregates_common_sqblock(const Matrix_h &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n FatalError(\"Size8 selector: setAggregates not implemented on CPU, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n\ntemplate \n__global__ //__launch_bounds__(256,2)\nvoid computeEdgeWeightsBlockDiaCsr_V2_1(\n const IndexType *row_offsets,\n //const int* __myrestrict row_offsets,\n const IndexType *row_indices,\n const IndexType *column_indices,\n const IndexType *dia_values,\n const ValueType *nonzero_values,\n //const ValueType* __restrict nonzero_values,\n const IndexType num_nonzero_blocks,\n float *str_edge_weights, float *rand_edge_weights, int num_owned, int bsize, int component)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int i, j;\n int bsize_sq = bsize * bsize;\n float kvalue;\n int matrix_weight_entry = component * bsize + component;\n bool valid_tid;\n\n while (utils::any(valid_tid = tid < num_nonzero_blocks))\n {\n i = -1;\n double d1, d2, w1;\n\n if (valid_tid)\n {\n if ( rand_edge_weights != NULL )\n {\n rand_edge_weights[tid] = random_weight(i, j, num_owned);\n }\n\n i = row_indices[tid];\n j = column_indices[tid];\n d1 = types::util::abs(__cachingLoad(&nonzero_values[__ldg(dia_values + i) * bsize_sq + matrix_weight_entry]));\n d2 = types::util::abs(__cachingLoad(&nonzero_values[__ldg(dia_values + j) * bsize_sq + matrix_weight_entry]));\n }\n\n const bool valid_j = valid_tid && i != j && j < num_owned;\n int ki = -1; //my transpose index, initialized to not found\n //int diag_j = -1; //j diagonal index\n\n if (!utils::any(valid_j))\n {\n continue;\n }\n\n int kmin = 0, kmax = 0;\n\n if (valid_j)\n {\n kmin = __cachingLoad(&row_offsets[j ]);\n kmax = __cachingLoad(&row_offsets[j + 1]);\n }\n\n for ( int k = kmin ; k < kmax ; ++k )\n {\n const int idx = __ldg(column_indices + k);\n\n if (idx == i)\n {\n ki = k; //find the transpose ji\n }\n }\n\n kvalue = 0.0f;\n\n if (ki > -1)\n {\n kvalue = types::util::abs(__cachingLoad(&nonzero_values[ki * bsize_sq + matrix_weight_entry]));\n }\n\n if (valid_tid)\n {\n w1 = types::util::abs(__cachingLoad(&nonzero_values[tid * bsize_sq + matrix_weight_entry]));\n str_edge_weights[tid] = 0.5 * (w1 + kvalue) / ( (float) max(d1, d2) ) * valid_j;\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n#if AMGX_ASYNCCPU_PROOF_OF_CONCEPT\n\ntemplate\nstruct async_size8_task : public task_cuda\n{\n T *size8;\n M *A;\n V *aggregates;\n V *aggregates_global;\n int *num_aggregates;\n\n typedef typename T::IndexType IndexType;\n\n void run()\n {\n const IndexType nnz_per_row = (*A).get_num_nz() / (*A).get_num_rows();\n\n if (0)\n {\n }\n else if (nnz_per_row > 2)\n {\n size8->template setAggregates_common_sqblock_avg_specialized<4>(*A, *aggregates, *aggregates_global, *num_aggregates);\n }\n else if (nnz_per_row > 1)\n {\n size8->template setAggregates_common_sqblock_avg_specialized<4>(*A, *aggregates, *aggregates_global, *num_aggregates);\n }\n else\n {\n size8->template setAggregates_common_sqblock_avg_specialized<2>(*A, *aggregates, *aggregates_global, *num_aggregates);\n }\n }\n};\n\n\ntemplate\nasync_size8_task *make_async_size8_task(T *size8, M &A, V &aggregates, V &aggregates_global, int &num_aggregates)\n{\n async_size8_task *ret = new async_size8_task;\n ret->size8 = size8;\n ret->A = &A;\n ret->aggregates = &aggregates;\n ret->aggregates_global = &aggregates_global;\n ret->num_aggregates = &num_aggregates;\n static task_chain_cuda_streamset *ss = new task_chain_cuda_streamset(1);\n task_chain_cuda *cr = new task_chain_cuda(ss);\n cr->append(ret, asyncmanager::singleton()->main_thread_queue(2));\n return ret;\n}\n\n#endif\n\ntemplate \nvoid Size8Selector >::setAggregates_common_sqblock(const Matrix_d &A,\n typename Matrix_d::IVector &aggregates,\n typename Matrix_d::IVector &aggregates_global,\n int &num_aggregates)\n{\n#if AMGX_ASYNCCPU_PROOF_OF_CONCEPT\n bool task = false;\n bool push = false;\n\n if (task)\n {\n task_cuda *t = make_async_size8_task(this, A, aggregates, aggregates_global, num_aggregates);\n enqueue_async_get_receipt(asyncmanager::singleton()->global_parallel_queue, t)->wait();\n return;\n }\n\n hipStream_t stream_old;\n static hipStream_t stream = 0;\n\n if (push)\n {\n stream_old = amgx::thrust::global_thread_handle::threadStream[getCurrentThreadId()];\n hipError_t cuda_rc = hipStreamSynchronize(amgx::thrust::global_thread_handle::threadStream[getCurrentThreadId()]);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipStreamSynchronize failed in size8_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n\n if (stream == 0) { hipStreamCreate(&stream); cudaCheckError(); }\n\n amgx::thrust::global_thread_handle::threadStream[getCurrentThreadId()] = stream;\n }\n\n#endif\n const IndexType nnz_per_row = A.get_num_nz() / A.get_num_rows();\n\n if (0)\n {\n }\n else if (nnz_per_row > 2)\n {\n setAggregates_common_sqblock_avg_specialized<4>(A, aggregates, aggregates_global, num_aggregates);\n }\n else if (nnz_per_row > 1)\n {\n setAggregates_common_sqblock_avg_specialized<4>(A, aggregates, aggregates_global, num_aggregates);\n }\n else\n {\n setAggregates_common_sqblock_avg_specialized<2>(A, aggregates, aggregates_global, num_aggregates);\n }\n\n#if AMGX_ASYNCCPU_PROOF_OF_CONCEPT\n\n if (push)\n {\n //cudaStreamSynchronize(stream);\n cudaCheckError();\n //amgx::thrust::global_thread_handle::threadStream[getCurrentThreadId()] = stream_old;\n }\n\n#endif\n}\n\n// setAggregates for block_dia_csr_matrix_d format\ntemplate \ntemplate\nvoid Size8Selector >:: setAggregates_common_sqblock_avg_specialized(const Matrix_d &A,\n typename Matrix_d::IVector &aggregates, typename Matrix_d::IVector &aggregates_global, int &num_aggregates)\n{\n#if AMGX_ASYNCCPU_PROOF_OF_CONCEPT\n hipStream_t stream = amgx::thrust::global_thread_handle::threadStream[getCurrentThreadId()];\n#else\n hipStream_t stream = amgx::thrust::global_thread_handle::get_stream();\n#endif\n const IndexType num_block_rows = A.get_num_rows();\n const IndexType num_nonzero_blocks = A.get_num_nz();\n\n if (!A.is_matrix_singleGPU())\n {\n aggregates.resize(A.manager->halo_offset(A.manager->num_neighbors()));\n }\n else\n {\n aggregates.resize(num_block_rows);\n }\n\n // Initially, put each vertex in its own aggregate\n thrust_wrapper::sequence(aggregates.begin(), aggregates.begin() + num_block_rows);\n cudaCheckError();\n IndexType *aggregates_ptr = aggregates.raw();\n // Create row_indices array\n IndexType total_nz = (A.is_matrix_singleGPU()) ? num_nonzero_blocks : A.manager->num_nz_all();\n typename Matrix_d::IVector row_indices(total_nz);\n cusp::detail::offsets_to_indices(A.row_offsets, row_indices);\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_row_indices_ptr = row_indices.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n //const ValueType *A_dia_val_ptr = amgx::thrust::raw_pointer_cast(&A.values[A.get_block_size()*A.diagOffset()]);\n const ValueType *A_nonzero_values_ptr = A.values.raw();\n typename Matrix_d::IVector strongest_neighbour(num_block_rows, -1);\n typename Matrix_d::IVector partner_index(3 * num_block_rows, -1);\n typename Matrix_d::IVector strongest_neighbour_tmp(num_block_rows);\n IndexType *strongest_neighbour_ptr = strongest_neighbour.raw();\n IndexType *partner_index_ptr = partner_index.raw();\n const int threads_per_block = 256;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_block_rows - 1) / threads_per_block + 1);\n int numUnassigned = num_block_rows;\n int numUnassigned_previous = numUnassigned;\n Vector > edge_weights(num_nonzero_blocks + 8); //8-padded\n float *edge_weights_ptr = edge_weights.raw();\n float *rand_edge_weights_ptr = NULL;\n const int num_blocks_V2 = std::min( AMGX_GRID_MAX_SIZE, (num_nonzero_blocks - 1) / threads_per_block + 1);\n // Compute the edge weights\n#if AMGX_ASYNCCPU_PROOF_OF_CONCEPT\n int avoid_thrust_count = 1;//0;\n int newFindStrongest = 1;//0;\n int newWeights = 1;//0;\n#else\n int avoid_thrust_count = 0;//0;\n int newFindStrongest = 0;//0;\n int newWeights = 0;//0;\n#endif\n int usenumassignedassumption = false;\n\n if (newWeights == 0)\n {\n hipFuncSetCacheConfig(reinterpret_cast(computeEdgeWeightsBlockDiaCsr_V2), hipFuncCachePreferL1);\n computeEdgeWeightsBlockDiaCsr_V2 <<< num_blocks_V2, threads_per_block, 0, stream>>>(\n A_row_offsets_ptr, A_row_indices_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, num_nonzero_blocks, edge_weights_ptr, rand_edge_weights_ptr, num_block_rows, A.get_block_dimy(), this->m_aggregation_edge_weight_component, this->weight_formula);\n cudaCheckError();\n }\n else\n {\n hipFuncSetCacheConfig(reinterpret_cast(computeEdgeWeightsBlockDiaCsr_V2_1), hipFuncCachePreferL1);\n computeEdgeWeightsBlockDiaCsr_V2_1 <<< num_blocks_V2, threads_per_block, 0, stream>>>(\n A_row_offsets_ptr, A_row_indices_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, num_nonzero_blocks, edge_weights_ptr, rand_edge_weights_ptr, num_block_rows, A.get_block_dimy(), this->m_aggregation_edge_weight_component);\n cudaCheckError();\n }\n\n // -------------------------------------------------\n // First create aggregates of size 2\n // -------------------------------------------------\n int icount = 0;\n const int num_blocks_1024 = std::min( 13 * 2, (num_block_rows - 1) / 1024 + 1);\n device_vector_alloc sets_per_block_t(num_blocks_1024);\n int *sets_per_block = amgx::thrust::raw_pointer_cast(sets_per_block_t.data());\n cudaCheckError();\n\n do\n {\n if (newFindStrongest)\n {\n if (numUnassigned == num_block_rows && usenumassignedassumption)\n {\n hipFuncSetCacheConfig(reinterpret_cast(my_findStrongestNeighbourBlockDiaCsr_NoMerge), hipFuncCachePreferL1);\n my_findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block, 0, stream>>>(\n A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, num_nonzero_blocks, 0, 0, strongest_neighbour_ptr, partner_index_ptr, 0, this->deterministic, 0, 0);\n cudaCheckError();\n }\n else\n {\n hipFuncSetCacheConfig(reinterpret_cast(my_findStrongestNeighbourBlockDiaCsr_NoMerge), hipFuncCachePreferL1);\n my_findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block, 0, stream>>>(\n A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, num_nonzero_blocks, 0, 0, strongest_neighbour_ptr, partner_index_ptr, 0, this->deterministic, 0, 0);\n cudaCheckError();\n }\n }\n else\n {\n findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, partner_index_ptr, strongest_neighbour_ptr, this->deterministic);\n }\n\n cudaCheckError();\n // Look for perfect matches\n numUnassigned_previous = numUnassigned;\n\n if (avoid_thrust_count == 0)\n {\n matchEdges <<< num_blocks, threads_per_block, 0, stream>>>(num_block_rows, partner_index_ptr, aggregates_ptr, strongest_neighbour_ptr);\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(partner_index.begin(), partner_index.begin() + num_block_rows, -1);\n cudaCheckError();\n }\n else\n {\n my_MatchEdges <<< num_blocks_1024, 1024, 0, stream>>>(num_block_rows, partner_index_ptr, aggregates_ptr, strongest_neighbour_ptr, sets_per_block);\n cudaCheckError();\n numUnassigned = numUnassigned_previous - amgx::strided_reduction::count_block_results_pinned_memory(0, icount, num_blocks_1024, sets_per_block, amgx::strided_reduction::op_sum(), stream);\n }\n\n icount++;\n }\n while (!(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_block_rows < this->numUnassigned_tol || numUnassigned_previous == numUnassigned));\n\n assignUnassignedVertices <<< num_blocks, threads_per_block, 0, stream>>>(partner_index_ptr, num_block_rows);\n cudaCheckError();\n // -------------------------------------------------\n // Merge aggregates to create aggregates of size 4\n // -------------------------------------------------\n Vector > weight_strongest_neighbour(num_block_rows, -1);\n float *weight_strongest_neighbour_ptr = weight_strongest_neighbour.raw();\n // At this point, partner index contain either your index or your neighbours index, depending on weither you're matched or not\n // aggregates contain the largest vertex index of vertices in aggregate\n typename Matrix_d::IVector aggregated(num_block_rows, -1);\n IndexType *aggregated_ptr = aggregated.raw();\n // now used as flag to check if aggregated or not\n icount = 0;\n numUnassigned = num_block_rows;\n numUnassigned_previous = numUnassigned;\n\n do\n {\n // Each vertex stores in strongest_neighbour the aggregates number of strongest neighbour and the weight of connection\n if (newFindStrongest)\n {\n if (numUnassigned == num_block_rows && usenumassignedassumption)\n {\n hipFuncSetCacheConfig(reinterpret_cast(my_findStrongestNeighbourBlockDiaCsr_NoMerge), hipFuncCachePreferL1);\n my_findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, num_nonzero_blocks, aggregated_ptr, aggregates_ptr, strongest_neighbour_ptr, partner_index_ptr, weight_strongest_neighbour_ptr, this->deterministic, 0, 0);\n cudaCheckError();\n }\n else\n {\n hipFuncSetCacheConfig(reinterpret_cast(my_findStrongestNeighbourBlockDiaCsr_NoMerge), hipFuncCachePreferL1);\n my_findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, num_nonzero_blocks, aggregated_ptr, aggregates_ptr, strongest_neighbour_ptr, partner_index_ptr, weight_strongest_neighbour_ptr, this->deterministic, 0, 0);\n cudaCheckError();\n }\n }\n else\n {\n findStrongestNeighbourBlockDiaCsr_StoreWeight <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregated_ptr, aggregates_ptr, strongest_neighbour_ptr, partner_index_ptr, weight_strongest_neighbour_ptr, this->deterministic);\n }\n\n cudaCheckError();\n // Each vertex in same aggregates will agree on aggregates to propose too, and both store the aggregate number they want to match with\n agreeOnProposal <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, num_block_rows, aggregated_ptr, strongest_neighbour_ptr, weight_strongest_neighbour_ptr, partner_index_ptr, aggregates_ptr);\n cudaCheckError();\n numUnassigned_previous = numUnassigned;\n\n if (avoid_thrust_count == 0)\n {\n matchAggregatesSize4 <<< num_blocks, threads_per_block, 0, stream>>>(aggregates_ptr, aggregated_ptr, strongest_neighbour_ptr, partner_index_ptr, num_block_rows);\n cudaCheckError();\n numUnassigned = thrust_wrapper::count(aggregated.begin(), aggregated.end(), -1);\n }\n else\n {\n my_matchAggregatesSize4 <<< num_blocks_1024, 1024, 0, stream>>>(aggregates_ptr, aggregated_ptr, strongest_neighbour_ptr, partner_index_ptr, num_block_rows, sets_per_block);\n cudaCheckError();\n numUnassigned = numUnassigned_previous - amgx::strided_reduction::count_block_results_pinned_memory(1, icount, num_blocks_1024, sets_per_block, amgx::strided_reduction::op_sum(), stream);\n }\n\n cudaCheckError();\n icount++;\n }\n while (!(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_block_rows < this->numUnassigned_tol || numUnassigned == numUnassigned_previous) );\n\n assignUnassignedVertices_2 <<< num_blocks, threads_per_block, 0, stream>>>(partner_index_ptr, num_block_rows);\n cudaCheckError();\n // -------------------------------------------------\n // Merge aggregates to create aggregates of size 8\n // -------------------------------------------------\n thrust_wrapper::fill(aggregated.begin(), aggregated.end(), -1);\n cudaCheckError();\n thrust_wrapper::fill(weight_strongest_neighbour.begin(), weight_strongest_neighbour.end(), -1.);\n cudaCheckError();\n icount = 0;\n numUnassigned = num_block_rows;\n numUnassigned_previous = numUnassigned;\n\n do\n {\n // Each vertex stores in strongest_neighbour the aggregates number of strongest neighbour and the weight of connection\n if (newFindStrongest)\n {\n if (numUnassigned == num_block_rows && usenumassignedassumption)\n {\n hipFuncSetCacheConfig(reinterpret_cast(my_findStrongestNeighbourBlockDiaCsr_NoMerge), hipFuncCachePreferL1);\n my_findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, num_nonzero_blocks, aggregated_ptr, aggregates_ptr, strongest_neighbour_ptr, partner_index_ptr, weight_strongest_neighbour_ptr, this->deterministic, 0, 0);\n cudaCheckError();\n }\n else\n {\n hipFuncSetCacheConfig(reinterpret_cast(my_findStrongestNeighbourBlockDiaCsr_NoMerge), hipFuncCachePreferL1);\n my_findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, num_nonzero_blocks, aggregated_ptr, aggregates_ptr, strongest_neighbour_ptr, partner_index_ptr, weight_strongest_neighbour_ptr, this->deterministic, 0, 0);\n cudaCheckError();\n }\n }\n else\n {\n findStrongestNeighbourBlockDiaCsr_StoreWeight_2 <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregated_ptr, aggregates_ptr, strongest_neighbour_ptr, partner_index_ptr, weight_strongest_neighbour_ptr);\n }\n\n cudaCheckError();\n\n // Each vertex in same aggregates will agree on aggregates to propose too, and both store the aggregate number they want to match with\n if (!this->deterministic)\n {\n agreeOnProposal_2 <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, num_block_rows, aggregated_ptr, strongest_neighbour_ptr, weight_strongest_neighbour_ptr, partner_index_ptr, aggregates_ptr, this->deterministic);\n cudaCheckError();\n }\n else\n {\n //strongest_neighbour_tmp = strongest_neighbour; // copied that directly in the kernel\n agreeOnProposal_2_deterministic <<< num_blocks, threads_per_block, 0, stream>>>(\n strongest_neighbour_tmp.raw(),\n A_row_offsets_ptr,\n A_column_indices_ptr, num_block_rows,\n aggregated_ptr, strongest_neighbour_ptr,\n weight_strongest_neighbour_ptr, partner_index_ptr, aggregates_ptr, this->deterministic);\n cudaCheckError();\n strongest_neighbour_tmp.swap(strongest_neighbour);\n strongest_neighbour_ptr = strongest_neighbour.raw(); //re-saving the correct pointer..\n }\n\n cudaCheckError();\n numUnassigned_previous = numUnassigned;\n\n if (avoid_thrust_count == 0)\n {\n matchAggregates <<< num_blocks, threads_per_block, 0, stream>>>(aggregates_ptr, aggregated_ptr, strongest_neighbour_ptr, num_block_rows);\n cudaCheckError();\n numUnassigned = thrust_wrapper::count(aggregated.begin(), aggregated.end(), -1);\n }\n else\n {\n my_matchAggregates <<< num_blocks_1024, 1024, 0, stream>>>(aggregates_ptr, aggregated_ptr, strongest_neighbour_ptr, num_block_rows, sets_per_block);\n cudaCheckError();\n numUnassigned = numUnassigned_previous - amgx::strided_reduction::count_block_results_pinned_memory(2, icount, num_blocks_1024, sets_per_block, amgx::strided_reduction::op_sum(), stream);\n }\n\n cudaCheckError();\n icount++;\n }\n while (!(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_block_rows < this->numUnassigned_tol || numUnassigned_previous == numUnassigned) );\n\n // Merge remaining vertices with current aggregates\n int local_iter = 0;\n\n if (!this->deterministic)\n {\n while (numUnassigned != 0)\n {\n mergeWithExistingAggregatesBlockDiaCsr <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, aggregated_ptr, this->deterministic, (IndexType *) NULL, local_iter > 1);\n cudaCheckError();\n numUnassigned_previous = numUnassigned;\n numUnassigned = (int)thrust_wrapper::count(aggregated.begin(), aggregated.end(), -1);\n cudaCheckError();\n local_iter++;\n }\n }\n else\n {\n typename Matrix_d::IVector aggregates_candidate(num_block_rows, -1);\n\n while (numUnassigned != 0)\n {\n // allow singletons only from the 2nd local iteration\n mergeWithExistingAggregatesBlockDiaCsr <<< num_blocks, threads_per_block, 0, stream>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, aggregated_ptr, this->deterministic, aggregates_candidate.raw(), local_iter > 1);\n cudaCheckError();\n numUnassigned_previous = numUnassigned;\n\n if (avoid_thrust_count == 0)\n {\n joinExistingAggregates <<< num_blocks, threads_per_block, 0, stream>>>(num_block_rows, aggregates_ptr, aggregated_ptr, aggregates_candidate.raw());\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(aggregated.begin(), aggregated.end(), -1);\n }\n else\n {\n my_joinExistingAggregates <<< num_blocks_1024, 1024, 0, stream>>>(num_block_rows, aggregates_ptr, aggregated_ptr, aggregates_candidate.raw(), sets_per_block);\n cudaCheckError();\n numUnassigned = numUnassigned_previous - amgx::strided_reduction::count_block_results_pinned_memory(3, local_iter, num_blocks_1024, sets_per_block, amgx::strided_reduction::op_sum(), stream);\n }\n\n cudaCheckError();\n local_iter++;\n }\n\n aggregates_candidate.resize(0);\n }\n\n this->renumberAndCountAggregates(aggregates, aggregates_global, num_block_rows, num_aggregates);\n //analyze_coloring(aggregates, A.getMatrixColoring().getRowColors());\n}\n\ntemplate \nvoid Size8SelectorBase::setAggregates(Matrix &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n if (A.get_block_dimx() == A.get_block_dimy())\n {\n setAggregates_common_sqblock( A, aggregates, aggregates_global, num_aggregates );\n }\n else\n {\n FatalError(\"Unsupported block size for Size8Selector\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\n// -------------------------\n// Explict instantiations\n// -------------------------\n#define AMGX_CASE_LINE(CASE) template class Size8SelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class Size8Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n}\n}\n}\n", "source": "/kaggle/working/amgx/src/aggregation/selectors/size8_selector.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include //count\n#include //sort\n#include //lower_bound\n#include //unique\n#include \n#include //offsets_to_indices\n#include \n#include \n\nnamespace amgx\n{\nnamespace aggregation\n{\nnamespace adaptive\n{\n\ntemplate \n__global__\nvoid checkDiagonalDominance( const IndexType *row_offsets, const IndexType *col_indices, const ValueType *values, const IndexType *diag, IndexType *aggregates, int num_owned, int bsize, double diag_dom )\n{\n //this time tid = i\n int i = threadIdx.x + blockDim.x * blockIdx.x;\n int bsize_sq = bsize * bsize;\n\n while ( i < num_owned )\n {\n bool dd = false;\n\n for ( int row = 0; row < bsize; row++ )\n {\n ValueType sum = 0.0;\n int iimax = row_offsets[i + 1];\n\n for ( int ii = row_offsets[i]; ii < iimax; ii++ )\n {\n //dont add the diagonal\n if ( col_indices[ii] == i )\n {\n for ( int col = 0; col < bsize; col++ )\n {\n sum += fabs( values[ii * bsize_sq + bsize * row + col] );\n\n if ( col == row )\n {\n continue;\n }\n }\n }\n else\n {\n for ( int col = 0; col < bsize; col++ )\n {\n sum += fabs( values[ii * bsize_sq + bsize * row + col] );\n }\n }\n }\n\n //check for DD condition and sort out\n dd = sum <= diag_dom * fabs( values[diag[i] * bsize_sq + row * bsize + row] );\n\n if ( !dd )\n {\n break;\n }\n }\n\n if ( dd )\n {\n aggregates[i] = num_owned;\n }\n\n i += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate \n__global__\nvoid assignNodeToBin( const ValueType *x, IndexType *aggregates, ValueType min, ValueType max, IndexType numRows, IndexType numAggregates )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n double range = (max - min) / numAggregates;\n\n while ( tid < numRows )\n {\n aggregates[tid] = floor( (x[tid] - min) / range );\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate \n__global__\nvoid rescaleVector( ValueType *x, IndexType numRows )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < numRows )\n {\n x[tid] = 2 * x[tid] - 1;\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// -----------------\n// Methods\n// ----------------\n\n// Constructor\ntemplate\nAdaptiveSelectorBase::AdaptiveSelectorBase(AMG_Config &cfg, const std::string &cfg_scope)\n{\n smoother = SolverFactory::allocate( cfg, cfg_scope, \"smoother\" );\n}\n/*\ntemplate \nvoid curandGenerateUniformWrapper( curandGenerator_t generator, ValueType *outputPtr, size_t num );\ntemplate<> void curandGenerateUniformWrapper( curandGenerator_t generator, float *outputPtr, size_t num )\n{\n curandGenerateUniform( generator, outputPtr, num);\n}\ntemplate<> void curandGenerateUniformWrapper( curandGenerator_t generator, double *outputPtr, size_t num )\n{\n curandGenerateUniformDouble( generator, outputPtr, num);\n}\n*/\n\ntemplate\nvoid AdaptiveSelectorBase::setAggregates(Matrix &A,\n IVector &aggregates, IVector &aggregates_global, int &numAggregates)\n{\n FatalError(\"Adaptive selector is still in development\", AMGX_ERR_NOT_IMPLEMENTED);\n /*\n if (A.get_block_dimx() == A.get_block_dimy())\n {\n IndexType numRows = A.get_num_rows();\n\n\n // This is how adaptive aggregation works:\n // 1. init random x vector (random values would preferably be between -1 and 1) and 0 rhs\n // 2. apply smoother a certain number of steps\n // 3. solve \"binning problem\":\n // 1. find min and max of x\n // 2. create N*coarseningrate bins with linear range\n // 3. assign nodes to bin, this is the final aggregates vector\n\n // allocate\n VVector& x = *Allocator::allocate( numRows );\n VVector& rhs = *Allocator::allocate( numRows );\n\n // initialize rhs\n thrust_wrapper::fill( rhs.begin(), rhs.end(), 0.0 );\n\n //compute random numbers\n curandGenerator_t generator;\n curandCreateGenerator(&generator, CURAND_RNG_PSEUDO_DEFAULT);\n curandGenerateUniformWrapper(generator, x.raw(), numRows);\n\n cudaStream_t str = 0;\n const int threads_per_block = 256;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (numRows-1)/threads_per_block + 1 );\n rescaleVector<<>>( x.raw(), numRows );\n cudaCheckError();\n\n\n // smooth\n smoother->setup( A, false );\n smoother->set_max_iters( 15 ); //TODO: make this a parameter\n smoother->solve( x, rhs, false );\n\n //free rhs\n Allocator::free( &rhs, numRows );\n\n //find min, max\n ValueType min = amgx::thrust::reduce( x.begin(), x.end(), 0.0, amgx::thrust::minimum() );\n ValueType max = amgx::thrust::reduce( x.begin(), x.end(), 0.0, amgx::thrust::maximum() );\n\n //bin\n aggregates.resize( numRows );\n numAggregates = numRows / 4;\n assignNodeToBin<<>>( x.raw(), aggregates.raw(), min, max, numRows, numAggregates );\n cudaCheckError();\n\n //sync\n cudaStreamSynchronize( str );\n cudaCheckError();\n\n //free\n Allocator::free( &x, numRows );\n\n //fill \"gaps\" of empty bins\n this->renumberAndCountAggregates(aggregates, aggregates_global, numRows, numAggregates);\n }\n else\n FatalError(\"Unsupported block size for Adaptive Aggregation\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n */\n}\n\n// -------------------------\n// Explict instantiations\n// -------------------------\n#define AMGX_CASE_LINE(CASE) template class AdaptiveSelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class AdaptiveSelector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n}\n}\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include //count\n#include //sort\n#include //lower_bound\n#include //unique\n#include \n#include //offsets_to_indices\n#include \n#include \n\nnamespace amgx\n{\nnamespace aggregation\n{\nnamespace adaptive\n{\n\ntemplate \n__global__\nvoid checkDiagonalDominance( const IndexType *row_offsets, const IndexType *col_indices, const ValueType *values, const IndexType *diag, IndexType *aggregates, int num_owned, int bsize, double diag_dom )\n{\n //this time tid = i\n int i = threadIdx.x + blockDim.x * blockIdx.x;\n int bsize_sq = bsize * bsize;\n\n while ( i < num_owned )\n {\n bool dd = false;\n\n for ( int row = 0; row < bsize; row++ )\n {\n ValueType sum = 0.0;\n int iimax = row_offsets[i + 1];\n\n for ( int ii = row_offsets[i]; ii < iimax; ii++ )\n {\n //dont add the diagonal\n if ( col_indices[ii] == i )\n {\n for ( int col = 0; col < bsize; col++ )\n {\n sum += fabs( values[ii * bsize_sq + bsize * row + col] );\n\n if ( col == row )\n {\n continue;\n }\n }\n }\n else\n {\n for ( int col = 0; col < bsize; col++ )\n {\n sum += fabs( values[ii * bsize_sq + bsize * row + col] );\n }\n }\n }\n\n //check for DD condition and sort out\n dd = sum <= diag_dom * fabs( values[diag[i] * bsize_sq + row * bsize + row] );\n\n if ( !dd )\n {\n break;\n }\n }\n\n if ( dd )\n {\n aggregates[i] = num_owned;\n }\n\n i += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate \n__global__\nvoid assignNodeToBin( const ValueType *x, IndexType *aggregates, ValueType min, ValueType max, IndexType numRows, IndexType numAggregates )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n double range = (max - min) / numAggregates;\n\n while ( tid < numRows )\n {\n aggregates[tid] = floor( (x[tid] - min) / range );\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate \n__global__\nvoid rescaleVector( ValueType *x, IndexType numRows )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < numRows )\n {\n x[tid] = 2 * x[tid] - 1;\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// -----------------\n// Methods\n// ----------------\n\n// Constructor\ntemplate\nAdaptiveSelectorBase::AdaptiveSelectorBase(AMG_Config &cfg, const std::string &cfg_scope)\n{\n smoother = SolverFactory::allocate( cfg, cfg_scope, \"smoother\" );\n}\n/*\ntemplate \nvoid curandGenerateUniformWrapper( curandGenerator_t generator, ValueType *outputPtr, size_t num );\ntemplate<> void curandGenerateUniformWrapper( curandGenerator_t generator, float *outputPtr, size_t num )\n{\n curandGenerateUniform( generator, outputPtr, num);\n}\ntemplate<> void curandGenerateUniformWrapper( curandGenerator_t generator, double *outputPtr, size_t num )\n{\n curandGenerateUniformDouble( generator, outputPtr, num);\n}\n*/\n\ntemplate\nvoid AdaptiveSelectorBase::setAggregates(Matrix &A,\n IVector &aggregates, IVector &aggregates_global, int &numAggregates)\n{\n FatalError(\"Adaptive selector is still in development\", AMGX_ERR_NOT_IMPLEMENTED);\n /*\n if (A.get_block_dimx() == A.get_block_dimy())\n {\n IndexType numRows = A.get_num_rows();\n\n\n // This is how adaptive aggregation works:\n // 1. init random x vector (random values would preferably be between -1 and 1) and 0 rhs\n // 2. apply smoother a certain number of steps\n // 3. solve \"binning problem\":\n // 1. find min and max of x\n // 2. create N*coarseningrate bins with linear range\n // 3. assign nodes to bin, this is the final aggregates vector\n\n // allocate\n VVector& x = *Allocator::allocate( numRows );\n VVector& rhs = *Allocator::allocate( numRows );\n\n // initialize rhs\n thrust_wrapper::fill( rhs.begin(), rhs.end(), 0.0 );\n\n //compute random numbers\n curandGenerator_t generator;\n curandCreateGenerator(&generator, CURAND_RNG_PSEUDO_DEFAULT);\n curandGenerateUniformWrapper(generator, x.raw(), numRows);\n\n cudaStream_t str = 0;\n const int threads_per_block = 256;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (numRows-1)/threads_per_block + 1 );\n rescaleVector<<>>( x.raw(), numRows );\n cudaCheckError();\n\n\n // smooth\n smoother->setup( A, false );\n smoother->set_max_iters( 15 ); //TODO: make this a parameter\n smoother->solve( x, rhs, false );\n\n //free rhs\n Allocator::free( &rhs, numRows );\n\n //find min, max\n ValueType min = amgx::thrust::reduce( x.begin(), x.end(), 0.0, amgx::thrust::minimum() );\n ValueType max = amgx::thrust::reduce( x.begin(), x.end(), 0.0, amgx::thrust::maximum() );\n\n //bin\n aggregates.resize( numRows );\n numAggregates = numRows / 4;\n assignNodeToBin<<>>( x.raw(), aggregates.raw(), min, max, numRows, numAggregates );\n cudaCheckError();\n\n //sync\n cudaStreamSynchronize( str );\n cudaCheckError();\n\n //free\n Allocator::free( &x, numRows );\n\n //fill \"gaps\" of empty bins\n this->renumberAndCountAggregates(aggregates, aggregates_global, numRows, numAggregates);\n }\n else\n FatalError(\"Unsupported block size for Adaptive Aggregation\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n */\n}\n\n// -------------------------\n// Explict instantiations\n// -------------------------\n#define AMGX_CASE_LINE(CASE) template class AdaptiveSelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class AdaptiveSelector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n}\n}\n", "source": "/kaggle/working/amgx/src/aggregation/selectors/adaptive.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"util.h\"\n\n#include \n#include \n#include \n#include \n\n#include \n\n#include //count\n#include //sort\n#include //lower_bound\n#include //unique\n#include //offsets_to_indices\n\n#define epsilon 1e-10\nnamespace amgx\n{\nnamespace aggregation\n{\n\n// ------------------------\n// Kernels\n// ------------------------\n\n//generate the aggregates labels\ntemplate \n__global__\nvoid generatelabel1d(IndexType *aggregation, const ValueType *points, const double cord_min, const double cord_max, const int npoint, const int nlevel)\n{\n unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;\n double distance = 1.01 * (cord_max - cord_min);\n\n for (; idx < npoint; idx += gridDim.x * blockDim.x)\n {\n aggregation[idx] = (int)((points[idx] - cord_min) / distance * nlevel);\n }\n}\n\ntemplate \n__global__\nvoid generatelabel1dpxb(IndexType *aggregation, const int alpha, const ValueType *points, const double cord_min, const double cord_max, const int npoint, const int nlevel)\n{\n unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;\n double distance = 1.01 * (cord_max - cord_min);\n\n for (; idx < npoint; idx += gridDim.x * blockDim.x)\n {\n aggregation[idx] = aggregation[idx] + alpha * (int)((points[idx] - cord_min) / distance * nlevel);\n }\n}\n\ntemplate \n__global__\nvoid generatelabel3d(IndexType *aggregation, const int n, const int n1d, const int n1d_orig)\n{\n unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int n1d2 = n1d_orig * n1d_orig;\n int x, y, z;\n\n for (; idx < n; idx += gridDim.x * blockDim.x)\n {\n int label = idx % n1d2;\n x = (label % n1d_orig) / 2;\n y = (label / n1d_orig) / 2;\n z = idx / n1d2 / 2;\n aggregation[idx] = z * n1d * n1d + y * n1d + x;\n }\n}\n\ntemplate \n__global__\nvoid generatelabel2d(IndexType *aggregation, const int n, const int n1d, const int n1d_orig)\n{\n unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int x, y;\n\n for (; idx < n; idx += gridDim.x * blockDim.x)\n {\n x = (idx % n1d_orig) / 2;\n y = (idx / n1d_orig) / 2;\n aggregation[idx] = y * n1d + x;\n }\n}\n\n\ntemplate \n__global__\nvoid aggregateLevel(IndexType *indices, const IndexType *next_level, const int num_rows)\n{\n unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n for (; idx < num_rows; idx += gridDim.x * blockDim.x)\n {\n indices[idx] = next_level[indices[idx]];\n }\n}\n\n\n\n// -----------------\n// Methods\n// ----------------\n\n// Constructor\ntemplate\nGEO_SelectorBase::GEO_SelectorBase(AMG_Config &cfg, const std::string &cfg_scope)\n{\n}\n\ntemplate \nbool compx (typename GEO_SelectorBase::p3d i, typename GEO_SelectorBase::p3d j) { return (i.x < j.x); }\n\ntemplate \nbool compy (typename GEO_SelectorBase::p3d i, typename GEO_SelectorBase::p3d j) { return (i.y < j.y); }\n\ntemplate \nbool compz (typename GEO_SelectorBase::p3d i, typename GEO_SelectorBase::p3d j) { return (i.z < j.z); }\n\ntemplate \nstruct reduce_functor2\n{\n __host__ __device__\n Tuple operator()(const Tuple x, const Tuple y)\n {\n return amgx::thrust::make_tuple(amgx::thrust::get<0>(x) + amgx::thrust::get<0>(y),\n amgx::thrust::get<1>(x) + amgx::thrust::get<1>(y),\n amgx::thrust::get<2>(x) + amgx::thrust::get<2>(y)\n );\n }\n};\n\ntemplate \nstruct reduce_functor3\n{\n __host__ __device__\n Tuple operator()(const Tuple x, const Tuple y)\n {\n return amgx::thrust::make_tuple(amgx::thrust::get<0>(x) + amgx::thrust::get<0>(y),\n amgx::thrust::get<1>(x) + amgx::thrust::get<1>(y),\n amgx::thrust::get<2>(x) + amgx::thrust::get<2>(y),\n amgx::thrust::get<3>(x) + amgx::thrust::get<3>(y)\n );\n }\n};\n\n// interpolates geometric info from the matrix using the matrix aggregation\ntemplate\nvoid GEO_SelectorBase::interpolateGeoinfo( Matrix &A )\n{\n const int threads_per_block = 512;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (A.get_num_rows() - 1) / threads_per_block + 1 );\n ngeo_x.resize(A.get_num_rows());\n ngeo_y.resize(A.get_num_rows());\n\n if (this->dimension == 3)\n {\n ngeo_z.resize(A.get_num_rows());\n }\n\n if (!(A.hasParameter(\"aggregates_info\")) )\n {\n FatalError(\"Cannot find information about previous aggregations for GEO selector\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n std::vector *agg_info = A.template getParameterPtr< std::vector >(\"aggregates_info\");\n IVector cur_geo_idx = *(agg_info->at(0));\n IVector v_counter(this->num_origonal, 1), v_new_counter(A.get_num_rows(), 1), r_coord(A.get_num_rows());\n\n if ( this->num_origonal != cur_geo_idx.size() )\n {\n FatalError(\"GEO size doesn't match original matrix dimension\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n for (size_t i = 1; i < agg_info->size(); i++)\n {\n aggregateLevel <<< num_blocks, threads_per_block>>>(cur_geo_idx.raw(), agg_info->at(i)->raw(), this->num_origonal);\n cudaCheckError();\n }\n\n if (this->dimension == 2)\n {\n amgx::thrust::sort_by_key(cur_geo_idx.begin(), cur_geo_idx.end(), amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(this->cord_x->begin(), this->cord_y->begin())));\n amgx::thrust::reduce_by_key(\n cur_geo_idx.begin(),\n cur_geo_idx.end(),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(this->cord_x->begin(), this->cord_y->begin(), v_counter.begin())),\n r_coord.begin(),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(ngeo_x.begin(), ngeo_y.begin(), v_new_counter.begin())),\n amgx::thrust::equal_to(),\n reduce_functor2< amgx::thrust::tuple< ValueType, ValueType, IndexType>, ValueType, IndexType > ()\n );\n }\n else\n {\n amgx::thrust::sort_by_key(cur_geo_idx.begin(), cur_geo_idx.end(), amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(this->cord_x->begin(), this->cord_y->begin(), this->cord_z->begin())));\n amgx::thrust::reduce_by_key(\n cur_geo_idx.begin(),\n cur_geo_idx.end(),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(this->cord_x->begin(), this->cord_y->begin(), this->cord_z->begin(), v_counter.begin())),\n r_coord.begin(),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(ngeo_x.begin(), ngeo_y.begin(), ngeo_z.begin(), v_new_counter.begin())),\n amgx::thrust::equal_to(),\n reduce_functor3< amgx::thrust::tuple< ValueType, ValueType, ValueType, IndexType>, ValueType, IndexType > ()\n );\n }\n\n cudaCheckError();\n}\n\n// Choose the aggregates for csr_matrix_d format\ntemplate \nvoid GEO_Selector >::setAggregates_1x1( Matrix_d &A,\n typename Matrix_d::IVector &aggregates, typename Matrix_d::IVector &aggregates_global, int &num_aggregates)\n{\n if (this->dimension == 0)\n {\n FatalError(\"No input geometric information, exiting\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n //initialize the aggregates vector\n int n = A.get_num_rows();\n aggregates.resize(n);\n thrust_wrapper::fill(aggregates.begin(), aggregates.end(), -1);\n cudaCheckError();\n int nlevel;\n\n if (this->dimension == 2) { nlevel = floor(log(sqrt((double)n)) / log(2.0)); }\n\n if (this->dimension == 3) { nlevel = ceil(log(cbrt((double)n)) / log(2.0)); }\n\n int num_per_row = (int)std::pow(2, nlevel - 1);\n const int threads_per_block = 512;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (n - 1) / threads_per_block + 1 );\n // generate aggregation index in 1d\n IndexType *aggregation_ptr = amgx::thrust::raw_pointer_cast(&aggregates[0]);\n\n if (A.hasParameter(\"uniform_based\"))\n {\n printf(\"GEOAggregating %d-rows matrix as uniform\\n\", A.get_num_rows());\n int npr_orig = 2 * num_per_row;\n\n if (this->dimension == 3)\n {\n generatelabel3d <<< num_blocks, threads_per_block>>>(aggregation_ptr, n, num_per_row, npr_orig);\n cudaCheckError();\n }\n else\n {\n generatelabel2d <<< num_blocks, threads_per_block>>>(aggregation_ptr, n, num_per_row, npr_orig);\n cudaCheckError();\n }\n }\n else\n {\n std::vector geo_ptrs(3);\n // real geometry source\n geo_ptrs[0] = this->cord_x;\n geo_ptrs[1] = this->cord_y;\n geo_ptrs[2] = this->cord_z;\n\n // do we need to interpolate from the finest level?\n if (A.get_num_rows() != this->num_origonal)\n {\n this->interpolateGeoinfo(A);\n\n if (this->dimension == 3)\n {\n geo_ptrs[2] = &this->ngeo_z;\n }\n\n geo_ptrs[0] = &this->ngeo_x;\n geo_ptrs[1] = &this->ngeo_y;\n }\n\n //Find the boundary coordinates\n this->xmax = *amgx::thrust::max_element(geo_ptrs[0]->begin(), geo_ptrs[0]->end());\n this->xmin = *amgx::thrust::min_element(geo_ptrs[0]->begin(), geo_ptrs[0]->end());\n this->ymax = *amgx::thrust::max_element(geo_ptrs[1]->begin(), geo_ptrs[1]->end());\n this->ymin = *amgx::thrust::min_element(geo_ptrs[1]->begin(), geo_ptrs[1]->end());\n\n if (this->dimension == 3)\n {\n this->zmax = *amgx::thrust::max_element(geo_ptrs[2]->begin(), geo_ptrs[2]->end());\n this->zmin = *amgx::thrust::min_element(geo_ptrs[2]->begin(), geo_ptrs[2]->end());\n }\n\n cudaCheckError();\n ValueType *point_ptr = geo_ptrs[0]->raw();\n generatelabel1d <<< num_blocks, threads_per_block>>>(aggregation_ptr, point_ptr, this->xmin, this->xmax, n, num_per_row);\n cudaCheckError();\n point_ptr = geo_ptrs[1]->raw();\n generatelabel1dpxb <<< num_blocks, threads_per_block>>>(aggregation_ptr, num_per_row, point_ptr, this->ymin, this->ymax, n, num_per_row);\n cudaCheckError();\n\n if (this->dimension > 2)\n {\n point_ptr = geo_ptrs[2]->raw();\n generatelabel1dpxb <<< num_blocks, threads_per_block>>>(aggregation_ptr, num_per_row * num_per_row, point_ptr, this->zmin, this->zmax, n, num_per_row);\n cudaCheckError();\n }\n }\n\n num_aggregates = num_per_row * num_per_row;\n\n if (this->dimension == 3) { num_aggregates *= num_per_row; }\n\n A.template setParameter < int > (\"uniform_based\", 1); // not exactly correct. But this is workarond for now, since parameters are instanlty copied for the coarser matrix in aggregation level\n}\n\ntemplate \nvoid GEO_Selector >::setAggregates_common_sqblocks(const Matrix_d &A,\n typename Matrix_d::IVector &aggregates, typename Matrix_d::IVector &aggregates_global, int &num_aggregates)\n{\n FatalError(\"GEO_Selector not implemented for non square blocks, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\n// Choose aggregates on host\ntemplate \nvoid GEO_Selector >::setAggregates_1x1( Matrix_h &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n FatalError(\"setAggregates not implemented on CPU, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n// Choose aggregates on host\ntemplate \nvoid GEO_Selector >::setAggregates_common_sqblocks(const Matrix_h &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n FatalError(\"setAggregates not implemented on CPU, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid GEO_SelectorBase::setAggregates(Matrix &A, IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n //printf(\"Selector: %d processing with the GEO\\n\", A.get_num_rows());\n this->dimension = A.template getParameter(\"dim\");\n cord_x = A.template getParameterPtr< VVector >(\"geo.x\");\n cord_y = A.template getParameterPtr< VVector >(\"geo.y\");\n num_origonal = A.template getParameter(\"geo_size\");\n\n if (this->dimension == 3)\n {\n cord_z = A.template getParameterPtr< VVector >(\"geo.z\");\n }\n\n if (this->dimension == 0) { FatalError(\"GEO_SELECTOR: Problem dimension is not valid.\", AMGX_ERR_BAD_PARAMETERS); }\n\n if (num_origonal == 0) { FatalError(\"Problem size is not valid.\", AMGX_ERR_BAD_PARAMETERS); }\n\n if (A.get_block_size() == 1 && !A.hasProps(DIAG) )\n {\n setAggregates_1x1( A, aggregates, aggregates_global, num_aggregates );\n }\n else if (A.get_block_dimx() == A.get_block_dimy())\n {\n setAggregates_1x1( A, aggregates, aggregates_global, num_aggregates );\n }\n else\n {\n FatalError(\"Unsupported block size for GEO_Selector\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\n// ---------------------------\n// Explict instantiations\n// ---------------------------\n#define AMGX_CASE_LINE(CASE) template class GEO_SelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class GEO_Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n}\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"util.h\"\n\n#include \n#include \n#include \n#include \n\n#include \n\n#include //count\n#include //sort\n#include //lower_bound\n#include //unique\n#include //offsets_to_indices\n\n#define epsilon 1e-10\nnamespace amgx\n{\nnamespace aggregation\n{\n\n// ------------------------\n// Kernels\n// ------------------------\n\n//generate the aggregates labels\ntemplate \n__global__\nvoid generatelabel1d(IndexType *aggregation, const ValueType *points, const double cord_min, const double cord_max, const int npoint, const int nlevel)\n{\n unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;\n double distance = 1.01 * (cord_max - cord_min);\n\n for (; idx < npoint; idx += gridDim.x * blockDim.x)\n {\n aggregation[idx] = (int)((points[idx] - cord_min) / distance * nlevel);\n }\n}\n\ntemplate \n__global__\nvoid generatelabel1dpxb(IndexType *aggregation, const int alpha, const ValueType *points, const double cord_min, const double cord_max, const int npoint, const int nlevel)\n{\n unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;\n double distance = 1.01 * (cord_max - cord_min);\n\n for (; idx < npoint; idx += gridDim.x * blockDim.x)\n {\n aggregation[idx] = aggregation[idx] + alpha * (int)((points[idx] - cord_min) / distance * nlevel);\n }\n}\n\ntemplate \n__global__\nvoid generatelabel3d(IndexType *aggregation, const int n, const int n1d, const int n1d_orig)\n{\n unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int n1d2 = n1d_orig * n1d_orig;\n int x, y, z;\n\n for (; idx < n; idx += gridDim.x * blockDim.x)\n {\n int label = idx % n1d2;\n x = (label % n1d_orig) / 2;\n y = (label / n1d_orig) / 2;\n z = idx / n1d2 / 2;\n aggregation[idx] = z * n1d * n1d + y * n1d + x;\n }\n}\n\ntemplate \n__global__\nvoid generatelabel2d(IndexType *aggregation, const int n, const int n1d, const int n1d_orig)\n{\n unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int x, y;\n\n for (; idx < n; idx += gridDim.x * blockDim.x)\n {\n x = (idx % n1d_orig) / 2;\n y = (idx / n1d_orig) / 2;\n aggregation[idx] = y * n1d + x;\n }\n}\n\n\ntemplate \n__global__\nvoid aggregateLevel(IndexType *indices, const IndexType *next_level, const int num_rows)\n{\n unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n for (; idx < num_rows; idx += gridDim.x * blockDim.x)\n {\n indices[idx] = next_level[indices[idx]];\n }\n}\n\n\n\n// -----------------\n// Methods\n// ----------------\n\n// Constructor\ntemplate\nGEO_SelectorBase::GEO_SelectorBase(AMG_Config &cfg, const std::string &cfg_scope)\n{\n}\n\ntemplate \nbool compx (typename GEO_SelectorBase::p3d i, typename GEO_SelectorBase::p3d j) { return (i.x < j.x); }\n\ntemplate \nbool compy (typename GEO_SelectorBase::p3d i, typename GEO_SelectorBase::p3d j) { return (i.y < j.y); }\n\ntemplate \nbool compz (typename GEO_SelectorBase::p3d i, typename GEO_SelectorBase::p3d j) { return (i.z < j.z); }\n\ntemplate \nstruct reduce_functor2\n{\n __host__ __device__\n Tuple operator()(const Tuple x, const Tuple y)\n {\n return amgx::thrust::make_tuple(amgx::thrust::get<0>(x) + amgx::thrust::get<0>(y),\n amgx::thrust::get<1>(x) + amgx::thrust::get<1>(y),\n amgx::thrust::get<2>(x) + amgx::thrust::get<2>(y)\n );\n }\n};\n\ntemplate \nstruct reduce_functor3\n{\n __host__ __device__\n Tuple operator()(const Tuple x, const Tuple y)\n {\n return amgx::thrust::make_tuple(amgx::thrust::get<0>(x) + amgx::thrust::get<0>(y),\n amgx::thrust::get<1>(x) + amgx::thrust::get<1>(y),\n amgx::thrust::get<2>(x) + amgx::thrust::get<2>(y),\n amgx::thrust::get<3>(x) + amgx::thrust::get<3>(y)\n );\n }\n};\n\n// interpolates geometric info from the matrix using the matrix aggregation\ntemplate\nvoid GEO_SelectorBase::interpolateGeoinfo( Matrix &A )\n{\n const int threads_per_block = 512;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (A.get_num_rows() - 1) / threads_per_block + 1 );\n ngeo_x.resize(A.get_num_rows());\n ngeo_y.resize(A.get_num_rows());\n\n if (this->dimension == 3)\n {\n ngeo_z.resize(A.get_num_rows());\n }\n\n if (!(A.hasParameter(\"aggregates_info\")) )\n {\n FatalError(\"Cannot find information about previous aggregations for GEO selector\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n std::vector *agg_info = A.template getParameterPtr< std::vector >(\"aggregates_info\");\n IVector cur_geo_idx = *(agg_info->at(0));\n IVector v_counter(this->num_origonal, 1), v_new_counter(A.get_num_rows(), 1), r_coord(A.get_num_rows());\n\n if ( this->num_origonal != cur_geo_idx.size() )\n {\n FatalError(\"GEO size doesn't match original matrix dimension\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n for (size_t i = 1; i < agg_info->size(); i++)\n {\n aggregateLevel <<< num_blocks, threads_per_block>>>(cur_geo_idx.raw(), agg_info->at(i)->raw(), this->num_origonal);\n cudaCheckError();\n }\n\n if (this->dimension == 2)\n {\n amgx::thrust::sort_by_key(cur_geo_idx.begin(), cur_geo_idx.end(), amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(this->cord_x->begin(), this->cord_y->begin())));\n amgx::thrust::reduce_by_key(\n cur_geo_idx.begin(),\n cur_geo_idx.end(),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(this->cord_x->begin(), this->cord_y->begin(), v_counter.begin())),\n r_coord.begin(),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(ngeo_x.begin(), ngeo_y.begin(), v_new_counter.begin())),\n amgx::thrust::equal_to(),\n reduce_functor2< amgx::thrust::tuple< ValueType, ValueType, IndexType>, ValueType, IndexType > ()\n );\n }\n else\n {\n amgx::thrust::sort_by_key(cur_geo_idx.begin(), cur_geo_idx.end(), amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(this->cord_x->begin(), this->cord_y->begin(), this->cord_z->begin())));\n amgx::thrust::reduce_by_key(\n cur_geo_idx.begin(),\n cur_geo_idx.end(),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(this->cord_x->begin(), this->cord_y->begin(), this->cord_z->begin(), v_counter.begin())),\n r_coord.begin(),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(ngeo_x.begin(), ngeo_y.begin(), ngeo_z.begin(), v_new_counter.begin())),\n amgx::thrust::equal_to(),\n reduce_functor3< amgx::thrust::tuple< ValueType, ValueType, ValueType, IndexType>, ValueType, IndexType > ()\n );\n }\n\n cudaCheckError();\n}\n\n// Choose the aggregates for csr_matrix_d format\ntemplate \nvoid GEO_Selector >::setAggregates_1x1( Matrix_d &A,\n typename Matrix_d::IVector &aggregates, typename Matrix_d::IVector &aggregates_global, int &num_aggregates)\n{\n if (this->dimension == 0)\n {\n FatalError(\"No input geometric information, exiting\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n //initialize the aggregates vector\n int n = A.get_num_rows();\n aggregates.resize(n);\n thrust_wrapper::fill(aggregates.begin(), aggregates.end(), -1);\n cudaCheckError();\n int nlevel;\n\n if (this->dimension == 2) { nlevel = floor(log(sqrt((double)n)) / log(2.0)); }\n\n if (this->dimension == 3) { nlevel = ceil(log(cbrt((double)n)) / log(2.0)); }\n\n int num_per_row = (int)std::pow(2, nlevel - 1);\n const int threads_per_block = 512;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (n - 1) / threads_per_block + 1 );\n // generate aggregation index in 1d\n IndexType *aggregation_ptr = amgx::thrust::raw_pointer_cast(&aggregates[0]);\n\n if (A.hasParameter(\"uniform_based\"))\n {\n printf(\"GEOAggregating %d-rows matrix as uniform\\n\", A.get_num_rows());\n int npr_orig = 2 * num_per_row;\n\n if (this->dimension == 3)\n {\n generatelabel3d <<< num_blocks, threads_per_block>>>(aggregation_ptr, n, num_per_row, npr_orig);\n cudaCheckError();\n }\n else\n {\n generatelabel2d <<< num_blocks, threads_per_block>>>(aggregation_ptr, n, num_per_row, npr_orig);\n cudaCheckError();\n }\n }\n else\n {\n std::vector geo_ptrs(3);\n // real geometry source\n geo_ptrs[0] = this->cord_x;\n geo_ptrs[1] = this->cord_y;\n geo_ptrs[2] = this->cord_z;\n\n // do we need to interpolate from the finest level?\n if (A.get_num_rows() != this->num_origonal)\n {\n this->interpolateGeoinfo(A);\n\n if (this->dimension == 3)\n {\n geo_ptrs[2] = &this->ngeo_z;\n }\n\n geo_ptrs[0] = &this->ngeo_x;\n geo_ptrs[1] = &this->ngeo_y;\n }\n\n //Find the boundary coordinates\n this->xmax = *amgx::thrust::max_element(geo_ptrs[0]->begin(), geo_ptrs[0]->end());\n this->xmin = *amgx::thrust::min_element(geo_ptrs[0]->begin(), geo_ptrs[0]->end());\n this->ymax = *amgx::thrust::max_element(geo_ptrs[1]->begin(), geo_ptrs[1]->end());\n this->ymin = *amgx::thrust::min_element(geo_ptrs[1]->begin(), geo_ptrs[1]->end());\n\n if (this->dimension == 3)\n {\n this->zmax = *amgx::thrust::max_element(geo_ptrs[2]->begin(), geo_ptrs[2]->end());\n this->zmin = *amgx::thrust::min_element(geo_ptrs[2]->begin(), geo_ptrs[2]->end());\n }\n\n cudaCheckError();\n ValueType *point_ptr = geo_ptrs[0]->raw();\n generatelabel1d <<< num_blocks, threads_per_block>>>(aggregation_ptr, point_ptr, this->xmin, this->xmax, n, num_per_row);\n cudaCheckError();\n point_ptr = geo_ptrs[1]->raw();\n generatelabel1dpxb <<< num_blocks, threads_per_block>>>(aggregation_ptr, num_per_row, point_ptr, this->ymin, this->ymax, n, num_per_row);\n cudaCheckError();\n\n if (this->dimension > 2)\n {\n point_ptr = geo_ptrs[2]->raw();\n generatelabel1dpxb <<< num_blocks, threads_per_block>>>(aggregation_ptr, num_per_row * num_per_row, point_ptr, this->zmin, this->zmax, n, num_per_row);\n cudaCheckError();\n }\n }\n\n num_aggregates = num_per_row * num_per_row;\n\n if (this->dimension == 3) { num_aggregates *= num_per_row; }\n\n A.template setParameter < int > (\"uniform_based\", 1); // not exactly correct. But this is workarond for now, since parameters are instanlty copied for the coarser matrix in aggregation level\n}\n\ntemplate \nvoid GEO_Selector >::setAggregates_common_sqblocks(const Matrix_d &A,\n typename Matrix_d::IVector &aggregates, typename Matrix_d::IVector &aggregates_global, int &num_aggregates)\n{\n FatalError(\"GEO_Selector not implemented for non square blocks, exiting\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\n// Choose aggregates on host\ntemplate \nvoid GEO_Selector >::setAggregates_1x1( Matrix_h &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n FatalError(\"setAggregates not implemented on CPU, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n// Choose aggregates on host\ntemplate \nvoid GEO_Selector >::setAggregates_common_sqblocks(const Matrix_h &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n FatalError(\"setAggregates not implemented on CPU, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid GEO_SelectorBase::setAggregates(Matrix &A, IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n //printf(\"Selector: %d processing with the GEO\\n\", A.get_num_rows());\n this->dimension = A.template getParameter(\"dim\");\n cord_x = A.template getParameterPtr< VVector >(\"geo.x\");\n cord_y = A.template getParameterPtr< VVector >(\"geo.y\");\n num_origonal = A.template getParameter(\"geo_size\");\n\n if (this->dimension == 3)\n {\n cord_z = A.template getParameterPtr< VVector >(\"geo.z\");\n }\n\n if (this->dimension == 0) { FatalError(\"GEO_SELECTOR: Problem dimension is not valid.\", AMGX_ERR_BAD_PARAMETERS); }\n\n if (num_origonal == 0) { FatalError(\"Problem size is not valid.\", AMGX_ERR_BAD_PARAMETERS); }\n\n if (A.get_block_size() == 1 && !A.hasProps(DIAG) )\n {\n setAggregates_1x1( A, aggregates, aggregates_global, num_aggregates );\n }\n else if (A.get_block_dimx() == A.get_block_dimy())\n {\n setAggregates_1x1( A, aggregates, aggregates_global, num_aggregates );\n }\n else\n {\n FatalError(\"Unsupported block size for GEO_Selector\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\n// ---------------------------\n// Explict instantiations\n// ---------------------------\n#define AMGX_CASE_LINE(CASE) template class GEO_SelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class GEO_Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n}\n", "source": "/kaggle/working/amgx/src/aggregation/selectors/geo_selector.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include //count\n#include //sort\n#include //unique\n#include //offsets_to_indices\n\nnamespace amgx\n{\nnamespace aggregation\n{\nnamespace size4_selector\n{\n\n// include common routines for all selectors\n#include \n\n// ------------------------\n// Kernels\n// ------------------------\n\n#ifndef DELETE\n// Kernel to compute the weight of the edges with block_dia_csr format\ntemplate \n__global__\nvoid computeEdgeWeightsBlockDiaCsr(const IndexType *row_offsets, const IndexType *column_indices,\n const IndexType *dia_idx, const ValueType *nonzero_values, const IndexType num_block_rows, float *edge_weights, int bsize)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int jcol;\n int bsize_sq = bsize * bsize;\n\n while (tid < num_block_rows)\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (jcol >= num_block_rows) { continue; }\n\n // Compute edge weight\n for (int k = row_offsets[jcol]; k < row_offsets[jcol + 1]; k++)\n {\n if (column_indices[k] == tid)\n {\n edge_weights[j] = 0.5 * (fabs(nonzero_values[j * bsize_sq]) + fabs(nonzero_values[k * bsize_sq]))\n / max( fabs(nonzero_values[bsize_sq * dia_idx[tid]]), fabs(nonzero_values[bsize_sq * dia_idx[jcol]]));\n break;\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n#endif\n\n// -----------------\n// Methods\n// ----------------\n\n// Constructor\ntemplate\nSize4SelectorBase::Size4SelectorBase(AMG_Config &cfg, const std::string &cfg_scope)\n{\n deterministic = cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\");\n max_iterations = cfg.AMG_Config::template getParameter(\"max_matching_iterations\", cfg_scope);\n numUnassigned_tol = cfg.AMG_Config::template getParameter(\"max_unassigned_percentage\", cfg_scope);\n m_aggregation_edge_weight_component = cfg.AMG_Config::template getParameter(\"aggregation_edge_weight_component\", cfg_scope);\n weight_formula = cfg.AMG_Config::template getParameter(\"weight_formula\", cfg_scope);\n}\n\ntemplate \nvoid Size4Selector >::setAggregates_common_sqblock(const Matrix_h &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n FatalError(\"OnePhaseHandshaking selector: setAggregates not implemented on CPU, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n\ntemplate \nvoid Size4Selector >::setAggregates_common_sqblock(const Matrix_d &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n // both ways are supported\n const IndexType num_block_rows = A.get_num_rows();\n const IndexType num_nonzero_blocks = A.get_num_nz();\n\n if (!A.is_matrix_singleGPU())\n {\n aggregates.resize(A.manager->halo_offset(A.manager->num_neighbors()));\n }\n else\n {\n aggregates.resize(num_block_rows);\n }\n\n // Initially, put each vertex in its own aggregate\n thrust_wrapper::sequence(aggregates.begin(), aggregates.begin() + num_block_rows);\n cudaCheckError();\n IndexType *aggregates_ptr = aggregates.raw();\n // Create row_indices array\n IndexType total_nz = A.is_matrix_singleGPU() ? num_nonzero_blocks : A.manager->num_nz_all();\n typename Matrix_d::IVector row_indices(total_nz);\n\n if (total_nz > 0)\n {\n cusp::detail::offsets_to_indices(A.row_offsets, row_indices);\n }\n\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_row_indices_ptr = row_indices.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueType *A_nonzero_values_ptr = A.values.raw();\n typename Matrix_d::IVector strongest_neighbour(num_block_rows, -1);\n typename Matrix_d::IVector partner_index(num_block_rows * 3, -1);\n IndexType *strongest_neighbour_ptr = strongest_neighbour.raw();\n IndexType *partner_index_ptr = partner_index.raw();\n const int threads_per_block = 256;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_block_rows - 1) / threads_per_block + 1);\n int numUnassigned = num_block_rows;\n int numUnassigned_previous = numUnassigned;\n Vector > edge_weights(num_nonzero_blocks, -1);\n float *edge_weights_ptr = edge_weights.raw();\n float *rand_edge_weights_ptr = NULL;\n // Compute the edge weights\n const int num_blocks_V2 = std::min( AMGX_GRID_MAX_SIZE, (num_nonzero_blocks - 1) / threads_per_block + 1);\n\n if (num_nonzero_blocks > 0)\n {\n cudaFuncSetCacheConfig(computeEdgeWeightsBlockDiaCsr_V2, cudaFuncCachePreferL1);\n computeEdgeWeightsBlockDiaCsr_V2 <<< num_blocks_V2, threads_per_block>>>(A_row_offsets_ptr, A_row_indices_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, num_nonzero_blocks, edge_weights_ptr, rand_edge_weights_ptr, num_block_rows, A.get_block_dimy(), this->m_aggregation_edge_weight_component, this->weight_formula);\n cudaCheckError();\n }\n\n // -------------------------------------------------\n // First create aggregates of size 2\n // -------------------------------------------------\n int icount = 0;\n\n do\n {\n // For each block_row, find the strongest neighbour who hasn't been assigned\n findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, partner_index_ptr, strongest_neighbour_ptr, this->deterministic);\n cudaCheckError();\n // Look for perfect matches\n matchEdges <<< num_blocks, threads_per_block>>>(num_block_rows, partner_index_ptr, aggregates_ptr, strongest_neighbour_ptr);\n cudaCheckError();\n numUnassigned_previous = numUnassigned;\n numUnassigned = (int)thrust_wrapper::count(partner_index.begin(), partner_index.end(), -1);\n cudaCheckError();\n icount++;\n }\n while (!(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_block_rows < this->numUnassigned_tol || numUnassigned_previous == numUnassigned));\n\n assignUnassignedVertices <<< num_blocks, threads_per_block>>>(partner_index_ptr, num_block_rows);\n cudaCheckError();\n // -------------------------------------------------\n // Merge aggregates to create aggregates of size 4\n // -------------------------------------------------\n Vector > weight_strongest_neighbour(num_block_rows, -1);\n float *weight_strongest_neighbour_ptr = weight_strongest_neighbour.raw();\n // now used as flag to check if aggregated or not\n typename Matrix_d::IVector aggregated(num_block_rows, -1);\n IndexType *aggregated_ptr = aggregated.raw();\n icount = 0;\n numUnassigned = num_block_rows;\n numUnassigned_previous = numUnassigned;\n\n do\n {\n // Each vertex stores in strongest_neighbour the aggregates number of strongest neighbour and the weight of connection\n findStrongestNeighbourBlockDiaCsr_StoreWeight <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregated_ptr, aggregates_ptr, strongest_neighbour_ptr, partner_index_ptr, weight_strongest_neighbour_ptr, this->deterministic);\n cudaCheckError();\n // Each vertex in same aggregates will agree on aggregates to propose too, and both store the aggregate number they want to match with\n agreeOnProposal <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, num_block_rows, aggregated_ptr, strongest_neighbour_ptr, weight_strongest_neighbour_ptr, partner_index_ptr, aggregates_ptr);\n cudaCheckError();\n matchAggregates <<< num_blocks, threads_per_block>>>(aggregates_ptr, aggregated_ptr, strongest_neighbour_ptr, num_block_rows);\n cudaCheckError();\n numUnassigned_previous = numUnassigned;\n numUnassigned = thrust_wrapper::count(aggregated.begin(), aggregated.end(), -1);\n cudaCheckError();\n icount++;\n }\n while (!(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_block_rows < this->numUnassigned_tol || numUnassigned_previous == numUnassigned));\n\n // Merge remaining vertices with current aggregates\n\n if (!this->deterministic)\n {\n while (numUnassigned != 0)\n {\n mergeWithExistingAggregatesBlockDiaCsr <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, aggregated_ptr, this->deterministic, (IndexType *) NULL);\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(aggregated.begin(), aggregated.end(), -1);\n cudaCheckError();\n };\n }\n else\n {\n typename Matrix_d::IVector aggregates_candidate(num_block_rows, -1);\n\n while (numUnassigned != 0)\n {\n mergeWithExistingAggregatesBlockDiaCsr <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, aggregated_ptr, this->deterministic, aggregates_candidate.raw());\n cudaCheckError();\n joinExistingAggregates <<< num_blocks, threads_per_block>>>(num_block_rows, aggregates_ptr, aggregated_ptr, aggregates_candidate.raw());\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(aggregated.begin(), aggregated.end(), -1);\n cudaCheckError();\n };\n }\n\n this->renumberAndCountAggregates(aggregates, aggregates_global, num_block_rows, num_aggregates);\n}\n\ntemplate \nvoid Size4SelectorBase::setAggregates(Matrix &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n if (A.get_block_dimx() == A.get_block_dimy())\n {\n setAggregates_common_sqblock( A, aggregates, aggregates_global, num_aggregates );\n }\n else\n {\n FatalError(\"Unsupported block size for Size4Selector\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\n// -------------------------\n// Explict instantiations\n// -------------------------\n#define AMGX_CASE_LINE(CASE) template class Size4SelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class Size4Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n}\n}\n}\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include //count\n#include //sort\n#include //unique\n#include //offsets_to_indices\n\nnamespace amgx\n{\nnamespace aggregation\n{\nnamespace size4_selector\n{\n\n// include common routines for all selectors\n#include \n\n// ------------------------\n// Kernels\n// ------------------------\n\n#ifndef DELETE\n// Kernel to compute the weight of the edges with block_dia_csr format\ntemplate \n__global__\nvoid computeEdgeWeightsBlockDiaCsr(const IndexType *row_offsets, const IndexType *column_indices,\n const IndexType *dia_idx, const ValueType *nonzero_values, const IndexType num_block_rows, float *edge_weights, int bsize)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int jcol;\n int bsize_sq = bsize * bsize;\n\n while (tid < num_block_rows)\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (jcol >= num_block_rows) { continue; }\n\n // Compute edge weight\n for (int k = row_offsets[jcol]; k < row_offsets[jcol + 1]; k++)\n {\n if (column_indices[k] == tid)\n {\n edge_weights[j] = 0.5 * (fabs(nonzero_values[j * bsize_sq]) + fabs(nonzero_values[k * bsize_sq]))\n / max( fabs(nonzero_values[bsize_sq * dia_idx[tid]]), fabs(nonzero_values[bsize_sq * dia_idx[jcol]]));\n break;\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n#endif\n\n// -----------------\n// Methods\n// ----------------\n\n// Constructor\ntemplate\nSize4SelectorBase::Size4SelectorBase(AMG_Config &cfg, const std::string &cfg_scope)\n{\n deterministic = cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\");\n max_iterations = cfg.AMG_Config::template getParameter(\"max_matching_iterations\", cfg_scope);\n numUnassigned_tol = cfg.AMG_Config::template getParameter(\"max_unassigned_percentage\", cfg_scope);\n m_aggregation_edge_weight_component = cfg.AMG_Config::template getParameter(\"aggregation_edge_weight_component\", cfg_scope);\n weight_formula = cfg.AMG_Config::template getParameter(\"weight_formula\", cfg_scope);\n}\n\ntemplate \nvoid Size4Selector >::setAggregates_common_sqblock(const Matrix_h &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n FatalError(\"OnePhaseHandshaking selector: setAggregates not implemented on CPU, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n\ntemplate \nvoid Size4Selector >::setAggregates_common_sqblock(const Matrix_d &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n // both ways are supported\n const IndexType num_block_rows = A.get_num_rows();\n const IndexType num_nonzero_blocks = A.get_num_nz();\n\n if (!A.is_matrix_singleGPU())\n {\n aggregates.resize(A.manager->halo_offset(A.manager->num_neighbors()));\n }\n else\n {\n aggregates.resize(num_block_rows);\n }\n\n // Initially, put each vertex in its own aggregate\n thrust_wrapper::sequence(aggregates.begin(), aggregates.begin() + num_block_rows);\n cudaCheckError();\n IndexType *aggregates_ptr = aggregates.raw();\n // Create row_indices array\n IndexType total_nz = A.is_matrix_singleGPU() ? num_nonzero_blocks : A.manager->num_nz_all();\n typename Matrix_d::IVector row_indices(total_nz);\n\n if (total_nz > 0)\n {\n cusp::detail::offsets_to_indices(A.row_offsets, row_indices);\n }\n\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_row_indices_ptr = row_indices.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueType *A_nonzero_values_ptr = A.values.raw();\n typename Matrix_d::IVector strongest_neighbour(num_block_rows, -1);\n typename Matrix_d::IVector partner_index(num_block_rows * 3, -1);\n IndexType *strongest_neighbour_ptr = strongest_neighbour.raw();\n IndexType *partner_index_ptr = partner_index.raw();\n const int threads_per_block = 256;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_block_rows - 1) / threads_per_block + 1);\n int numUnassigned = num_block_rows;\n int numUnassigned_previous = numUnassigned;\n Vector > edge_weights(num_nonzero_blocks, -1);\n float *edge_weights_ptr = edge_weights.raw();\n float *rand_edge_weights_ptr = NULL;\n // Compute the edge weights\n const int num_blocks_V2 = std::min( AMGX_GRID_MAX_SIZE, (num_nonzero_blocks - 1) / threads_per_block + 1);\n\n if (num_nonzero_blocks > 0)\n {\n hipFuncSetCacheConfig(reinterpret_cast(computeEdgeWeightsBlockDiaCsr_V2), hipFuncCachePreferL1);\n computeEdgeWeightsBlockDiaCsr_V2 <<< num_blocks_V2, threads_per_block>>>(A_row_offsets_ptr, A_row_indices_ptr, A_column_indices_ptr, A_dia_idx_ptr, A_nonzero_values_ptr, num_nonzero_blocks, edge_weights_ptr, rand_edge_weights_ptr, num_block_rows, A.get_block_dimy(), this->m_aggregation_edge_weight_component, this->weight_formula);\n cudaCheckError();\n }\n\n // -------------------------------------------------\n // First create aggregates of size 2\n // -------------------------------------------------\n int icount = 0;\n\n do\n {\n // For each block_row, find the strongest neighbour who hasn't been assigned\n findStrongestNeighbourBlockDiaCsr_NoMerge <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, partner_index_ptr, strongest_neighbour_ptr, this->deterministic);\n cudaCheckError();\n // Look for perfect matches\n matchEdges <<< num_blocks, threads_per_block>>>(num_block_rows, partner_index_ptr, aggregates_ptr, strongest_neighbour_ptr);\n cudaCheckError();\n numUnassigned_previous = numUnassigned;\n numUnassigned = (int)thrust_wrapper::count(partner_index.begin(), partner_index.end(), -1);\n cudaCheckError();\n icount++;\n }\n while (!(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_block_rows < this->numUnassigned_tol || numUnassigned_previous == numUnassigned));\n\n assignUnassignedVertices <<< num_blocks, threads_per_block>>>(partner_index_ptr, num_block_rows);\n cudaCheckError();\n // -------------------------------------------------\n // Merge aggregates to create aggregates of size 4\n // -------------------------------------------------\n Vector > weight_strongest_neighbour(num_block_rows, -1);\n float *weight_strongest_neighbour_ptr = weight_strongest_neighbour.raw();\n // now used as flag to check if aggregated or not\n typename Matrix_d::IVector aggregated(num_block_rows, -1);\n IndexType *aggregated_ptr = aggregated.raw();\n icount = 0;\n numUnassigned = num_block_rows;\n numUnassigned_previous = numUnassigned;\n\n do\n {\n // Each vertex stores in strongest_neighbour the aggregates number of strongest neighbour and the weight of connection\n findStrongestNeighbourBlockDiaCsr_StoreWeight <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregated_ptr, aggregates_ptr, strongest_neighbour_ptr, partner_index_ptr, weight_strongest_neighbour_ptr, this->deterministic);\n cudaCheckError();\n // Each vertex in same aggregates will agree on aggregates to propose too, and both store the aggregate number they want to match with\n agreeOnProposal <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, num_block_rows, aggregated_ptr, strongest_neighbour_ptr, weight_strongest_neighbour_ptr, partner_index_ptr, aggregates_ptr);\n cudaCheckError();\n matchAggregates <<< num_blocks, threads_per_block>>>(aggregates_ptr, aggregated_ptr, strongest_neighbour_ptr, num_block_rows);\n cudaCheckError();\n numUnassigned_previous = numUnassigned;\n numUnassigned = thrust_wrapper::count(aggregated.begin(), aggregated.end(), -1);\n cudaCheckError();\n icount++;\n }\n while (!(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_block_rows < this->numUnassigned_tol || numUnassigned_previous == numUnassigned));\n\n // Merge remaining vertices with current aggregates\n\n if (!this->deterministic)\n {\n while (numUnassigned != 0)\n {\n mergeWithExistingAggregatesBlockDiaCsr <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, aggregated_ptr, this->deterministic, (IndexType *) NULL);\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(aggregated.begin(), aggregated.end(), -1);\n cudaCheckError();\n };\n }\n else\n {\n typename Matrix_d::IVector aggregates_candidate(num_block_rows, -1);\n\n while (numUnassigned != 0)\n {\n mergeWithExistingAggregatesBlockDiaCsr <<< num_blocks, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, aggregated_ptr, this->deterministic, aggregates_candidate.raw());\n cudaCheckError();\n joinExistingAggregates <<< num_blocks, threads_per_block>>>(num_block_rows, aggregates_ptr, aggregated_ptr, aggregates_candidate.raw());\n cudaCheckError();\n numUnassigned = (int)thrust_wrapper::count(aggregated.begin(), aggregated.end(), -1);\n cudaCheckError();\n };\n }\n\n this->renumberAndCountAggregates(aggregates, aggregates_global, num_block_rows, num_aggregates);\n}\n\ntemplate \nvoid Size4SelectorBase::setAggregates(Matrix &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n if (A.get_block_dimx() == A.get_block_dimy())\n {\n setAggregates_common_sqblock( A, aggregates, aggregates_global, num_aggregates );\n }\n else\n {\n FatalError(\"Unsupported block size for Size4Selector\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\n// -------------------------\n// Explict instantiations\n// -------------------------\n#define AMGX_CASE_LINE(CASE) template class Size4SelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class Size4Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n}\n}\n}\n", "source": "/kaggle/working/amgx/src/aggregation/selectors/size4_selector.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n\n#include \n\n#include \n\nnamespace amgx\n{\n\nnamespace aggregation\n{\n\ntemplate\nvoid Selector::renumberAndCountAggregates(IVector &aggregates, IVector &aggregates_global, const IndexType num_block_rows, IndexType &num_aggregates)\n{\n // renumber aggregates\n IVector scratch(num_block_rows + 1);\n\n if ( num_block_rows != aggregates.size() )\n {\n // we are in a distributed environment\n aggregates_global.resize(aggregates.size());\n amgx::thrust::copy(aggregates.begin(), aggregates.begin() + num_block_rows, aggregates_global.begin());\n }\n\n // set scratch[aggregates[i]] = 1\n thrust::fill(amgx::thrust::make_permutation_iterator(scratch.begin(), aggregates.begin()),\n amgx::thrust::make_permutation_iterator(scratch.begin(), aggregates.begin() + num_block_rows), 1);\n // do prefix sum on scratch\n thrust_wrapper::exclusive_scan(scratch.begin(), scratch.end(), scratch.begin());\n // aggregates[i] = scratch[aggregates[i]]\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(scratch.begin(), aggregates.begin()),\n amgx::thrust::make_permutation_iterator(scratch.begin(), aggregates.begin() + num_block_rows),\n aggregates.begin());\n // update number of aggregates\n num_aggregates = scratch[scratch.size() - 1];\n cudaCheckError();\n}\n\n\ntemplate\n__global__\nvoid countRowSizes( const IndexType *R_row_offsets, IndexType *size_array, int num_aggregates)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < num_aggregates )\n {\n int size = R_row_offsets[tid + 1] - R_row_offsets[tid];\n atomicAdd( &size_array[size], 1 );\n tid += gridDim.x * blockDim.x;\n }\n}\n\n//prints aggregate info:\n//counts aggregates grouped by size\ntemplate\nvoid Selector::printAggregationInfo(const IVector &aggregates, const IVector &aggregates_global, const IndexType num_aggregates) const\n{\n //build restriction row offsets\n Matrix P;\n P.set_initialized(0);\n P.addProps(CSR);\n P.delProps(COO);\n P.delProps(DIAG);\n P.setColsReorderedByColor(false);\n P.resize( 0, 0, 0, 1, 1); //make sure the matrix is empty\n Matrix R;\n R.set_initialized(0);\n R.addProps(CSR);\n R.delProps(COO);\n R.delProps(DIAG);\n R.setColsReorderedByColor(false);\n R.resize( 0, 0, 0, 1, 1);\n int num_rows = aggregates.size();\n //set correct sizes.\n P.row_offsets.resize( num_rows + 1 );\n P.values.resize( num_rows, types::util::get_one());\n P.col_indices.resize( num_rows );\n //setup offset array.\n thrust_wrapper::sequence( P.row_offsets.begin(), P.row_offsets.end() );\n //swap in aggregates\n amgx::thrust::copy( aggregates.begin(), aggregates.end(), P.col_indices.begin() );\n cudaCheckError();\n //inform P about its size\n P.resize( num_rows, num_aggregates, num_rows, 1, 1, false ); //declare scalar, otherwise transpose won't work\n //P is ready\n P.set_initialized(1);\n //do the transpose\n R.set_initialized(0);\n transpose( P, R ); //note: this won't change P\n //ok now we have row offsets.\n //we are ready to count the row sizes\n IVector size_array;\n size_array.resize(num_rows + 1, 0 );\n const int threads_per_block = 256;\n const int num_blocks = std::min( (int)AMGX_GRID_MAX_SIZE, (int)(num_rows - 1) / threads_per_block + 1 );\n //count\n countRowSizes <<< num_blocks, threads_per_block, 0, 0>>>( R.row_offsets.raw(), size_array.raw(), num_aggregates );\n cudaCheckError();\n cudaDeviceSynchronize();\n cudaCheckError();\n //copy to host and print\n amgx::thrust::host_vector size_array_host;\n size_array_host.resize(size_array.size() );\n amgx::thrust::copy( size_array.begin(), size_array.end(), size_array_host.begin() );\n std::cout << \"number of nodes \" << num_rows << std::endl;\n std::cout << \"number of aggregates by size\" << std::endl;\n\n for (int i = 1; i < size_array_host.size(); i++)\n {\n if ( size_array_host[i] > 0 )\n {\n std::cout << \"size \" << i << \": \" << size_array_host[i] << std::endl;\n }\n }\n\n std::cout << \"total \" << (num_aggregates - size_array_host[0]) << std::endl;\n}\n\ntemplate\nvoid Selector::assertAggregates( const IVector &aggregates, int numAggregates )\n{\n int *res_sizes = new int[numAggregates];\n\n for ( int i = 0; i < numAggregates; i++ )\n {\n res_sizes[i] = 0;\n }\n\n int *h_aggregates = new int[aggregates.size()];\n cudaMemcpy( h_aggregates, aggregates.raw(), sizeof(int)*aggregates.size(), cudaMemcpyDeviceToHost );\n cudaCheckError();\n\n for ( int i = 0; i < aggregates.size(); i++ )\n {\n if ( h_aggregates[i] < 0 || h_aggregates[i] >= numAggregates )\n {\n std::cout << \"Error: aggregates[\" << i << \"]=\" << h_aggregates[i] << \" but numAggregates=\" << numAggregates << std::endl;\n }\n else\n {\n res_sizes[h_aggregates[i]]++;\n }\n }\n\n for ( int i = 0; i < numAggregates; i++ )\n {\n if ( res_sizes[i] == 0 )\n {\n std::cout << \"Error: aggregate \" << i << \" is empty\" << std::endl;\n }\n }\n\n std::cout << \"assertAggregates done.\" << std::endl;\n}\n\ntemplate\nvoid Selector::assertRestriction( const IVector &R_row_offsets, const IVector &R_col_indices, const IVector &aggregates )\n{\n int *r_ia = new int[R_row_offsets.size()];\n int *r_ja = new int[R_col_indices.size()];\n int *agg = new int[aggregates.size()];\n int *used_col = new int[aggregates.size()];\n\n for ( int i = 0; i < aggregates.size(); i++ )\n {\n used_col[i] = 0;\n }\n\n cudaError_t cuda_rc = cudaMemcpy( r_ia, R_row_offsets.raw(), sizeof(int)*R_row_offsets.size(), cudaMemcpyDeviceToHost );\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy R_row_offsets D2H failed in agg_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = cudaMemcpy( r_ja, R_col_indices.raw(), sizeof(int)*R_col_indices.size(), cudaMemcpyDeviceToHost );\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy R_col_indices D2H failed in agg_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = cudaMemcpy( agg, aggregates.raw(), sizeof(int)*aggregates.size(), cudaMemcpyDeviceToHost );\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy aggregates D2H failed in agg_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n\n for ( int i = 0; i < R_row_offsets.size() - 1; i++ )\n {\n for ( int ii = r_ia[i]; ii < r_ia[i + 1]; ii++ )\n {\n int j = r_ja[ii];\n used_col[j]++;\n\n if ( used_col[j] > 1 )\n {\n std::cout << \"column \" << j << \" is present at least \" << used_col[j] << \" times\" << std::endl;\n }\n\n if ( j < 0 || j >= aggregates.size() )\n {\n std::cout << \"Error: j out of bounds, j = \" << j << \" and numRows = \" << aggregates.size() << std::endl;\n }\n else if ( agg[j] != i )\n {\n std::cout << \"Error: agg[\" << j << \"] = \" << agg[j] << \" != \" << i << std::endl;\n }\n }\n }\n\n std::cout << \"assert restriction done\" << std::endl;\n}\n\n\n\ntemplate\nstd::map*> &\nSelectorFactory::getFactories( )\n{\n static std::map*> s_factories;\n return s_factories;\n}\n\ntemplate\nvoid SelectorFactory::registerFactory(std::string name, SelectorFactory *f)\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::const_iterator it = factories.find(name);\n\n if (it != factories.end())\n {\n std::string error = \"SelectorFactory '\" + name + \"' has already been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n factories[name] = f;\n}\n\ntemplate\nvoid SelectorFactory::unregisterFactory(std::string name)\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.find(name);\n\n if (it == factories.end())\n {\n std::string error = \"SelectorFactory '\" + name + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n SelectorFactory *factory = it->second;\n assert( factory != NULL );\n delete factory;\n factories.erase(it);\n}\n\ntemplate\nvoid SelectorFactory::unregisterFactories( )\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.begin( );\n\n for ( ; it != factories.end( ) ; )\n {\n SelectorFactory *factory = it->second;\n assert( factory != NULL );\n it++;\n delete factory;\n }\n\n factories.clear( );\n}\n\ntemplate\nSelector *SelectorFactory::allocate(AMG_Config &cfg, const std::string ¤t_scope)\n{\n std::map*> &factories = getFactories( );\n int agg_lvl_change = cfg.AMG_Config::template getParameter(\"fine_levels\", current_scope);\n std::string selector;\n selector = cfg.getParameter(\"selector\", current_scope);\n typename std::map *>::const_iterator it = factories.find(selector);\n\n if (it == factories.end())\n {\n std::string error = \"SelectorFactory '\" + selector + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n return it->second->create(cfg, current_scope);\n};\n\n\n\n//****************************************\n// Explict instantiations\n// **************************************\n#define AMGX_CASE_LINE(CASE) template class Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class SelectorFactory::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n\n}\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n#include \n\n#include \n\nnamespace amgx\n{\n\nnamespace aggregation\n{\n\ntemplate\nvoid Selector::renumberAndCountAggregates(IVector &aggregates, IVector &aggregates_global, const IndexType num_block_rows, IndexType &num_aggregates)\n{\n // renumber aggregates\n IVector scratch(num_block_rows + 1);\n\n if ( num_block_rows != aggregates.size() )\n {\n // we are in a distributed environment\n aggregates_global.resize(aggregates.size());\n amgx::thrust::copy(aggregates.begin(), aggregates.begin() + num_block_rows, aggregates_global.begin());\n }\n\n // set scratch[aggregates[i]] = 1\n thrust::fill(amgx::thrust::make_permutation_iterator(scratch.begin(), aggregates.begin()),\n amgx::thrust::make_permutation_iterator(scratch.begin(), aggregates.begin() + num_block_rows), 1);\n // do prefix sum on scratch\n thrust_wrapper::exclusive_scan(scratch.begin(), scratch.end(), scratch.begin());\n // aggregates[i] = scratch[aggregates[i]]\n amgx::thrust::copy(amgx::thrust::make_permutation_iterator(scratch.begin(), aggregates.begin()),\n amgx::thrust::make_permutation_iterator(scratch.begin(), aggregates.begin() + num_block_rows),\n aggregates.begin());\n // update number of aggregates\n num_aggregates = scratch[scratch.size() - 1];\n cudaCheckError();\n}\n\n\ntemplate\n__global__\nvoid countRowSizes( const IndexType *R_row_offsets, IndexType *size_array, int num_aggregates)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < num_aggregates )\n {\n int size = R_row_offsets[tid + 1] - R_row_offsets[tid];\n atomicAdd( &size_array[size], 1 );\n tid += gridDim.x * blockDim.x;\n }\n}\n\n//prints aggregate info:\n//counts aggregates grouped by size\ntemplate\nvoid Selector::printAggregationInfo(const IVector &aggregates, const IVector &aggregates_global, const IndexType num_aggregates) const\n{\n //build restriction row offsets\n Matrix P;\n P.set_initialized(0);\n P.addProps(CSR);\n P.delProps(COO);\n P.delProps(DIAG);\n P.setColsReorderedByColor(false);\n P.resize( 0, 0, 0, 1, 1); //make sure the matrix is empty\n Matrix R;\n R.set_initialized(0);\n R.addProps(CSR);\n R.delProps(COO);\n R.delProps(DIAG);\n R.setColsReorderedByColor(false);\n R.resize( 0, 0, 0, 1, 1);\n int num_rows = aggregates.size();\n //set correct sizes.\n P.row_offsets.resize( num_rows + 1 );\n P.values.resize( num_rows, types::util::get_one());\n P.col_indices.resize( num_rows );\n //setup offset array.\n thrust_wrapper::sequence( P.row_offsets.begin(), P.row_offsets.end() );\n //swap in aggregates\n amgx::thrust::copy( aggregates.begin(), aggregates.end(), P.col_indices.begin() );\n cudaCheckError();\n //inform P about its size\n P.resize( num_rows, num_aggregates, num_rows, 1, 1, false ); //declare scalar, otherwise transpose won't work\n //P is ready\n P.set_initialized(1);\n //do the transpose\n R.set_initialized(0);\n transpose( P, R ); //note: this won't change P\n //ok now we have row offsets.\n //we are ready to count the row sizes\n IVector size_array;\n size_array.resize(num_rows + 1, 0 );\n const int threads_per_block = 256;\n const int num_blocks = std::min( (int)AMGX_GRID_MAX_SIZE, (int)(num_rows - 1) / threads_per_block + 1 );\n //count\n countRowSizes <<< num_blocks, threads_per_block, 0, 0>>>( R.row_offsets.raw(), size_array.raw(), num_aggregates );\n cudaCheckError();\n hipDeviceSynchronize();\n cudaCheckError();\n //copy to host and print\n amgx::thrust::host_vector size_array_host;\n size_array_host.resize(size_array.size() );\n amgx::thrust::copy( size_array.begin(), size_array.end(), size_array_host.begin() );\n std::cout << \"number of nodes \" << num_rows << std::endl;\n std::cout << \"number of aggregates by size\" << std::endl;\n\n for (int i = 1; i < size_array_host.size(); i++)\n {\n if ( size_array_host[i] > 0 )\n {\n std::cout << \"size \" << i << \": \" << size_array_host[i] << std::endl;\n }\n }\n\n std::cout << \"total \" << (num_aggregates - size_array_host[0]) << std::endl;\n}\n\ntemplate\nvoid Selector::assertAggregates( const IVector &aggregates, int numAggregates )\n{\n int *res_sizes = new int[numAggregates];\n\n for ( int i = 0; i < numAggregates; i++ )\n {\n res_sizes[i] = 0;\n }\n\n int *h_aggregates = new int[aggregates.size()];\n hipMemcpy( h_aggregates, aggregates.raw(), sizeof(int)*aggregates.size(), hipMemcpyDeviceToHost );\n cudaCheckError();\n\n for ( int i = 0; i < aggregates.size(); i++ )\n {\n if ( h_aggregates[i] < 0 || h_aggregates[i] >= numAggregates )\n {\n std::cout << \"Error: aggregates[\" << i << \"]=\" << h_aggregates[i] << \" but numAggregates=\" << numAggregates << std::endl;\n }\n else\n {\n res_sizes[h_aggregates[i]]++;\n }\n }\n\n for ( int i = 0; i < numAggregates; i++ )\n {\n if ( res_sizes[i] == 0 )\n {\n std::cout << \"Error: aggregate \" << i << \" is empty\" << std::endl;\n }\n }\n\n std::cout << \"assertAggregates done.\" << std::endl;\n}\n\ntemplate\nvoid Selector::assertRestriction( const IVector &R_row_offsets, const IVector &R_col_indices, const IVector &aggregates )\n{\n int *r_ia = new int[R_row_offsets.size()];\n int *r_ja = new int[R_col_indices.size()];\n int *agg = new int[aggregates.size()];\n int *used_col = new int[aggregates.size()];\n\n for ( int i = 0; i < aggregates.size(); i++ )\n {\n used_col[i] = 0;\n }\n\n hipError_t cuda_rc = hipMemcpy( r_ia, R_row_offsets.raw(), sizeof(int)*R_row_offsets.size(), hipMemcpyDeviceToHost );\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy R_row_offsets D2H failed in agg_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = hipMemcpy( r_ja, R_col_indices.raw(), sizeof(int)*R_col_indices.size(), hipMemcpyDeviceToHost );\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy R_col_indices D2H failed in agg_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = hipMemcpy( agg, aggregates.raw(), sizeof(int)*aggregates.size(), hipMemcpyDeviceToHost );\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy aggregates D2H failed in agg_selector\", AMGX_ERR_CUDA_FAILURE);\n }\n\n for ( int i = 0; i < R_row_offsets.size() - 1; i++ )\n {\n for ( int ii = r_ia[i]; ii < r_ia[i + 1]; ii++ )\n {\n int j = r_ja[ii];\n used_col[j]++;\n\n if ( used_col[j] > 1 )\n {\n std::cout << \"column \" << j << \" is present at least \" << used_col[j] << \" times\" << std::endl;\n }\n\n if ( j < 0 || j >= aggregates.size() )\n {\n std::cout << \"Error: j out of bounds, j = \" << j << \" and numRows = \" << aggregates.size() << std::endl;\n }\n else if ( agg[j] != i )\n {\n std::cout << \"Error: agg[\" << j << \"] = \" << agg[j] << \" != \" << i << std::endl;\n }\n }\n }\n\n std::cout << \"assert restriction done\" << std::endl;\n}\n\n\n\ntemplate\nstd::map*> &\nSelectorFactory::getFactories( )\n{\n static std::map*> s_factories;\n return s_factories;\n}\n\ntemplate\nvoid SelectorFactory::registerFactory(std::string name, SelectorFactory *f)\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::const_iterator it = factories.find(name);\n\n if (it != factories.end())\n {\n std::string error = \"SelectorFactory '\" + name + \"' has already been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n factories[name] = f;\n}\n\ntemplate\nvoid SelectorFactory::unregisterFactory(std::string name)\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.find(name);\n\n if (it == factories.end())\n {\n std::string error = \"SelectorFactory '\" + name + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n SelectorFactory *factory = it->second;\n assert( factory != NULL );\n delete factory;\n factories.erase(it);\n}\n\ntemplate\nvoid SelectorFactory::unregisterFactories( )\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.begin( );\n\n for ( ; it != factories.end( ) ; )\n {\n SelectorFactory *factory = it->second;\n assert( factory != NULL );\n it++;\n delete factory;\n }\n\n factories.clear( );\n}\n\ntemplate\nSelector *SelectorFactory::allocate(AMG_Config &cfg, const std::string ¤t_scope)\n{\n std::map*> &factories = getFactories( );\n int agg_lvl_change = cfg.AMG_Config::template getParameter(\"fine_levels\", current_scope);\n std::string selector;\n selector = cfg.getParameter(\"selector\", current_scope);\n typename std::map *>::const_iterator it = factories.find(selector);\n\n if (it == factories.end())\n {\n std::string error = \"SelectorFactory '\" + selector + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n return it->second->create(cfg, current_scope);\n};\n\n\n\n//****************************************\n// Explict instantiations\n// **************************************\n#define AMGX_CASE_LINE(CASE) template class Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class SelectorFactory::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n\n}\n", "source": "/kaggle/working/amgx/src/aggregation/selectors/agg_selector.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include //count\n#include //sort\n#include //lower_bound\n#include //unique\n#include \n#include //offsets_to_indices\n#include \n#include \n\n#include \n#include \n\n#define EXPERIMENTAL_ITERATIVE_MATCHING\n\nnamespace amgx\n{\nnamespace aggregation\n{\nnamespace multi_pairwise\n{\n\n// include common routines for all selectors\n#include \n\n// ------------------------\n// Kernels\n// ------------------------\n\n__device__\nfloat random_weight2(int i, int j)\n{\n#define RAND_MULTIPLIER 1145637293\n unsigned long i_min = (min(i, j) * RAND_MULTIPLIER);\n unsigned long i_max = (max(i, j) * RAND_MULTIPLIER);\n return ((float)i_min / i_max);\n}\n\n__device__\nunsigned long random_weight3(int i, int j)\n{\n unsigned long a;\n a = (i + j) ^ 8;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n\n\n// findStrongestNeighbour kernel for block_dia_csr_matrix format\n// Reads the weight from edge_weights array\ntemplate \n__global__\nvoid findStrongestNeighbourBlockDiaCsr_V2(const IndexType *row_offsets, const IndexType *column_indices,\n MatrixValueType *edge_weights, IndexType num_block_rows, IndexType *aggregates,\n IndexType *strongest_neighbour_1phase, IndexType *strongest_neighbour,\n const size_t bsize, int phase, int merge_singletons)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n MatrixValueType weight;\n int jcol;\n\n while (tid < num_block_rows)\n {\n int strongest_unaggregated = -1;\n int strongest_aggregated = -1;\n MatrixValueType max_weight_unaggregated = 0.;\n MatrixValueType max_weight_aggregated = 0.;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (phase == 1) { weight = edge_weights[j]; }\n else { weight = random_weight2(tid, jcol); }\n\n if (tid == jcol || jcol >= num_block_rows) { continue; } // skip diagonal and halo\n\n if (phase == 2 && strongest_neighbour_1phase[jcol] != tid) { continue; } // if 2nd phase only accept those who gave a hand on the 1st phase\n\n // Identify strongest aggregated and unaggregated neighbours (method by multi_pairwise)\n if (aggregates[jcol] == -1 && weight > 0.0 && (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && random_weight3(tid, jcol) > random_weight3(tid, strongest_unaggregated)))) // unaggregated\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n else if (aggregates[jcol] != -1 && weight > 0.0 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && random_weight3(tid, jcol) > random_weight3(tid, strongest_aggregated)))) // aggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_unaggregated == -1 && strongest_aggregated != -1) // All neighbours are aggregated\n {\n if ( merge_singletons == 1 )\n // Put in same aggregate as strongest neighbour\n {\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n else\n // create singleton\n {\n aggregates[tid] = tid;\n }\n }\n else if (strongest_unaggregated != -1)\n {\n if (phase == 2)\n {\n MatrixValueType rand_w1 = random_weight2(tid, strongest_neighbour_1phase[tid]);\n strongest_neighbour[tid] = max_weight_unaggregated > rand_w1 ? strongest_unaggregated : strongest_neighbour_1phase[tid];\n }\n else { strongest_neighbour_1phase[tid] = strongest_unaggregated; }\n }\n else\n {\n if (phase == 2) { strongest_neighbour[tid] = strongest_neighbour_1phase[tid]; }\n else { strongest_neighbour_1phase[tid] = tid; }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate \n__device__\nbool atomicJoin( IndexType node, IndexType aggregate, IndexType *aggregates, int *sizes, int allowed )\n{\n int mySize = sizes[node];\n int theirSize = sizes[aggregate];\n int theirSizeOld = theirSize;\n\n do\n {\n int newSize = mySize + theirSize;\n\n if ( newSize > allowed )\n {\n return false;\n }\n\n theirSizeOld = theirSize;\n theirSize = atomicCAS( &sizes[aggregate], theirSize, newSize );\n }\n while ( theirSize != theirSizeOld );\n\n aggregates[node] = aggregate;\n return true;\n}\n\ntemplate \n__global__\nvoid findStrongestNeighbourBlockDiaCsr_V3(const IndexType *row_offsets,\n const IndexType *column_indices,\n MatrixValueType *edge_weights,\n IndexType num_block_rows,\n IndexType *aggregates,\n IndexType *strongest_neighbour,\n int *sizes,\n int *degree,\n const size_t bsize,\n int max_aggregate_size,\n int merge_singletons)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n MatrixValueType weight;\n int jcol;\n\n while (tid < num_block_rows)\n {\n int strongest_unaggregated = -1;\n int strongest_aggregated = -1;\n int lowest_degree;\n\n if ( use_degree )\n {\n lowest_degree = degree[tid]; //only interested in finding lower degree than self\n }\n else\n {\n lowest_degree = 0; //if we decide to not use degree than just propose to the strongest edge\n }\n\n int lowest_degree_neighbor = tid;\n MatrixValueType lowest_degree_weight = 1e100; //high start value, so that same degree neighbor won't become lowest degree neighbor\n MatrixValueType max_weight_unaggregated = 0.;\n MatrixValueType max_weight_aggregated = 0.;\n int mySize;\n\n if ( merge_singletons == 2 )\n {\n mySize = sizes[tid];\n }\n else\n {\n mySize = 0;\n }\n\n if ( merge_singletons != 2 )\n {\n max_aggregate_size = 100000;\n }\n\n //this aggregate is already full\n if (mySize == max_aggregate_size)\n {\n aggregates[tid] = tid;\n }\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (tid == jcol || jcol >= num_block_rows) { continue; } // skip diagonal and halo\n\n weight = edge_weights[j];\n\n if (weight <= 0.0) { continue; }\n\n if ( aggregates[jcol] != -1 ) //aggregated neighbor\n {\n int theirSize;\n\n if ( merge_singletons == 2 )\n {\n theirSize = aggregates[sizes[jcol]];\n }\n else\n {\n theirSize = 0;\n }\n\n //if all neighbors are aggregated, find the strongest edge to neighbor aggregate that is not full yet\n if (mySize + theirSize <= max_aggregate_size &&\n (weight > max_weight_aggregated)) // aggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n else //unaggregated neighbor\n {\n if ( use_degree && merge_singletons == 2 )\n {\n int theirSize = sizes[jcol];\n //get lowest degree neighbor or find out that there is no lower degree neighbor\n int current_degree = degree[jcol];\n\n if ( mySize + theirSize <= max_aggregate_size && (current_degree < lowest_degree || current_degree == lowest_degree && weight > lowest_degree_weight) )\n {\n lowest_degree = current_degree;\n lowest_degree_weight = weight;\n lowest_degree_neighbor = jcol;\n }\n\n //get highest weight neighbor\n if ( mySize + theirSize <= max_aggregate_size && (weight > max_weight_unaggregated) )\n {\n strongest_unaggregated = jcol;\n max_weight_unaggregated = weight;\n }\n }\n\n if ( use_degree && merge_singletons != 2 ) //same as above but ignore sizes\n {\n //get lowest degree neighbor or find out that there is no lower degree neighbor\n int current_degree = degree[jcol];\n\n if ( current_degree < lowest_degree || current_degree == lowest_degree && weight > lowest_degree_weight)\n {\n lowest_degree = current_degree;\n lowest_degree_weight = weight;\n lowest_degree_neighbor = jcol;\n }\n\n //get highest weight neighbor\n if (weight > max_weight_unaggregated)\n {\n strongest_unaggregated = jcol;\n max_weight_unaggregated = weight;\n }\n }\n\n if ( !use_degree && merge_singletons == 2 )\n {\n //get highest weight neighbor only but pay attention to the aggregate sizes\n int theirSize = sizes[jcol]; //get highest weight neighbor\n\n if ( mySize + theirSize <= max_aggregate_size && (weight > max_weight_unaggregated) )\n {\n strongest_unaggregated = jcol;\n max_weight_unaggregated = weight;\n }\n }\n\n if ( !use_degree && merge_singletons != 2 )\n {\n //just highest weight\n if (weight > max_weight_unaggregated)\n {\n strongest_unaggregated = jcol;\n max_weight_unaggregated = weight;\n }\n }\n }\n }\n\n //prefer lowest degree neighbor\n if ( lowest_degree_neighbor != tid )\n {\n strongest_unaggregated = lowest_degree_neighbor;\n }\n\n if (strongest_unaggregated != -1) //Unaggregated neighbor exists\n {\n strongest_neighbour[tid] = strongest_unaggregated; //assign strongest aggregated\n }\n\n if (strongest_unaggregated == -1 && strongest_aggregated != -1) // All neighbours are aggregated but small enough aggregated neighbors exist\n {\n if ( merge_singletons == 0 )\n {\n aggregates[tid] = tid;\n }\n\n if ( merge_singletons == 1 )\n {\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n\n if ( merge_singletons == 2)\n {\n atomicJoin( tid, aggregates[strongest_aggregated], aggregates, sizes, max_aggregate_size ); //try to join, can fail. maybe it works next round.\n }\n }\n\n if (strongest_unaggregated == -1 && strongest_aggregated == -1) //no feasable neighbor at all, become singleton\n {\n strongest_neighbour[tid] = tid; //become singleton\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \n__global__\nvoid computeDegree( const IndexType *ia, const IndexType *ja, const ValueType *weights, IndexType *aggregates, IndexType *sizes, IndexType *degree, IndexType numRows, IndexType max_aggregate_size)\n{\n int i = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( i < numRows )\n {\n int myDegree = 0;\n IndexType ia_ip1 = ia[i + 1];\n\n for ( IndexType ii = ia[i]; ii < ia_ip1; ii++ )\n {\n IndexType j = ja[ii];\n\n if ( j == i )\n {\n continue;\n }\n\n int mySize, theirSize;\n\n if ( sizes != NULL )\n {\n mySize = sizes[i];\n theirSize = sizes[j];\n }\n else\n {\n mySize = theirSize = 0;\n }\n\n if ( weights[ii] > 0.0 && aggregates[j] == -1 && mySize + theirSize <= max_aggregate_size )\n {\n myDegree++;\n }\n }\n\n degree[i] = myDegree;\n i += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate \n__global__\nvoid mergeSingletonsSmart(const IndexType *ia, const IndexType *ja, const ValueType *weights, IndexType *aggregates, IndexType *sizes, IndexType numRows, int max_aggregate_size)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < numRows )\n {\n //unaggregated nodes try to join or create their own aggregate\n if ( aggregates[tid] == -1 )\n {\n bool joined = false;\n\n while ( !joined )\n {\n int neighbor_aggregate = -1;\n ValueType max_weight = 0.0;\n IndexType mySize = sizes[tid];\n\n for (IndexType ii = ia[tid]; ii < ia[tid + 1]; ii++)\n {\n IndexType j = ja[ii];\n\n if (j == tid || j >= numRows) { continue; }\n\n if ( aggregates[j] != -1 && sizes[aggregates[j]] + mySize <= max_aggregate_size && weights[ii] > max_weight )\n {\n neighbor_aggregate = aggregates[j];\n max_weight = weights[ii];\n }\n }\n\n //no possible neighbor found\n if ( neighbor_aggregate == -1 )\n {\n //create own aggregate\n aggregates[tid] = tid;\n joined = true;\n }\n else\n {\n //try to join\n joined = atomicJoin( tid, neighbor_aggregate, aggregates, sizes, max_aggregate_size );\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate \n__global__\nvoid updateAggregateSizes( IndexType *sizesSource, IndexType *sizes, IndexType *aggregates, IndexType numRows )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < numRows )\n {\n IndexType agg = aggregates[tid];\n IndexType aggregateSize = sizes[agg];\n IndexType mySize = sizesSource[tid];\n\n while ( mySize > aggregateSize )\n {\n aggregateSize = atomicCAS( &sizes[agg], aggregateSize, mySize );\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\n// Kernel that checks if perfect matchs exist\ntemplate \n__global__\nvoid matchEdges(const IndexType num_rows, IndexType *aggregates, int *strongest_neighbour, IndexType *sizes)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int potential_match, potential_match_neighbour;\n\n while (tid < num_rows)\n {\n if (aggregates[tid] == -1) // Unaggregated row\n {\n potential_match = strongest_neighbour[tid];\n potential_match_neighbour = strongest_neighbour[potential_match];\n\n if ( potential_match == tid )\n {\n aggregates[tid] = tid;\n }\n else if (potential_match != -1 && potential_match_neighbour == tid && tid < potential_match) // we have a match\n {\n aggregates[tid] = tid;\n aggregates[potential_match] = tid;\n\n if ( sizes != NULL)\n {\n sizes[tid] += sizes[potential_match];\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \n__global__\nvoid countAggregates(const IndexType num_rows, IndexType *aggregates, int *num_unaggregated)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int c = 0;\n int i = tid;\n\n while ( i < num_rows )\n {\n c += ( aggregates[i] == -1 );\n i += gridDim.x * blockDim.x;\n }\n\n __shared__ volatile int smem[block_size];\n smem[threadIdx.x] = c;\n __syncthreads();\n\n for ( int off = blockDim.x / 2; off >= 32; off = off / 2 )\n {\n if ( threadIdx.x < off )\n {\n smem[threadIdx.x] += smem[threadIdx.x + off];\n }\n\n __syncthreads();\n }\n\n // warp reduce\n if ( threadIdx.x < 32 )\n {\n smem[threadIdx.x] += smem[threadIdx.x + 16];\n smem[threadIdx.x] += smem[threadIdx.x + 8];\n smem[threadIdx.x] += smem[threadIdx.x + 4];\n smem[threadIdx.x] += smem[threadIdx.x + 2];\n smem[threadIdx.x] += smem[threadIdx.x + 1];\n }\n\n if ( threadIdx.x == 0 )\n {\n atomicAdd(num_unaggregated, smem[0]);\n }\n}\n\ntemplate \n__global__\nvoid joinExistingAggregates(IndexType num_rows, IndexType *aggregates, IndexType *aggregates_candidate)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while (tid < num_rows)\n {\n if (aggregates[tid] == -1 && aggregates_candidate[tid] != -1) // Unaggregated row\n {\n aggregates[tid] = aggregates_candidate[tid];\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel that merges unaggregated vertices its strongest aggregated neighbour\n// Weights are read from edge_weights array\n// For block_dia_csr_matrix_format\ntemplate \n__global__\nvoid mergeWithExistingAggregatesBlockDiaCsr_V2(const IndexType *row_offsets, const IndexType *column_indices, const MatrixValueType *edge_weights,\n const int num_block_rows, IndexType *aggregates, int bsize, const int deterministic, IndexType *aggregates_candidate)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int jcol;\n MatrixValueType weight;\n\n while (tid < num_block_rows)\n {\n MatrixValueType max_weight_aggregated = 0.;\n int strongest_aggregated = -1;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n // Compute edge weight\n weight = edge_weights[j];\n jcol = column_indices[j];\n\n if (jcol == tid || jcol >= num_block_rows) { continue; } // skip diagonal\n\n if ( aggregates[jcol] == num_block_rows ) { continue; } // skip dd rows\n\n // Identify strongest aggregated neighbour\n if (aggregates[jcol] != -1 && weight > 0 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && random_weight3( tid, jcol ) > random_weight3( tid, strongest_aggregated )))) //\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_aggregated != -1) // Found a neighbour to aggregate to\n {\n if (deterministic == 1)\n {\n aggregates_candidate[tid] = aggregates[strongest_aggregated];\n }\n else\n {\n // Put in same aggregate as strongest neighbour\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n }\n else // All neighbours are unaggregated, leave alone\n {\n if (deterministic == 1)\n {\n aggregates_candidate[tid] = tid;\n }\n else\n {\n aggregates[tid] = tid;\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel to extract diagonal for csr_matrix format\ntemplate \n__global__\nvoid getDiagonalKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, const IndexType numRows, ValueType *diagonal)\n{\n int tIdx = threadIdx.x + blockDim.x * blockIdx.x;\n\n while (tIdx < numRows)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n for (int j = offset; j < offset + numj; j++)\n {\n int jcol = column_indices[j];\n\n if (tIdx == jcol)\n {\n diagonal[tIdx] = values[j];\n }\n }\n\n tIdx += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel to extract diagonal for csr_matrix format\ntemplate \n__global__\nvoid getDiagonalKernelNoDiaProp(const IndexType *dia_idx, const ValueType *values, const IndexType numRows, ValueType *diagonal)\n{\n int tIdx = threadIdx.x + blockDim.x * blockIdx.x;\n\n while (tIdx < numRows)\n {\n diagonal[tIdx] = values[dia_idx[tIdx]];\n tIdx += gridDim.x * blockDim.x;\n }\n}\n\n// filter edge weights like this:\n// set w_ij = 0 iff\n// w_ij < alpha * sqrt( max_k{w_ik} * max_l{w_jl} )\n// alpha is some constant, 0.25 or 0.5 should work fine\ntemplate\n__global__\nvoid filterWeights( const IndexType *row_offsets, const IndexType *row_indices, const IndexType *col_indices, const IndexType *diag, const ValueType *old_weights, ValueType *new_weights, IndexType num_nonzero_blocks, IndexType num_owned, ValueType alpha )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int i, j, kmin, kmax;\n ValueType max_ik, max_jl;\n\n while ( tid < num_nonzero_blocks )\n {\n i = row_indices[tid];\n j = col_indices[tid];\n\n if ( i != j && j < num_owned )\n {\n //find max_k{w_ik}\n kmin = row_offsets[i];\n kmax = row_offsets[i + 1];\n max_ik = 0.0;\n\n for (int k = kmin; k < kmax; k++)\n {\n if ( col_indices[k] != i && old_weights[k] > max_ik )\n {\n max_ik = old_weights[k];\n }\n }\n\n //find max_l{w_jl}\n kmin = row_offsets[j];\n kmax = row_offsets[j + 1];\n max_jl = 0.0;\n\n for (int l = kmin; l < kmax; l++)\n {\n if ( col_indices[l] != j && old_weights[l] > max_jl )\n {\n max_jl = old_weights[l];\n }\n }\n\n //test squared inequality\n if ( old_weights[tid] * old_weights[tid] < alpha * alpha * max_ik * max_jl )\n {\n new_weights[tid] = 0.0;\n }\n else //rescale to relative importance. this should also increase the chance of a handshake\n {\n new_weights[tid] = old_weights[tid];\n }\n\n // new_weights[tid] = old_weights[tid] / sqrt(max_ik*max_jl);\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate\n__global__\nvoid gatherValuesInterleaved( const ValueType *inValues, ValueType *outValues, IndexType nnz, int sq_blocksize, int index_offset )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < nnz )\n {\n //at least the write is coalesced\n outValues[tid] = inValues[tid * sq_blocksize + index_offset];\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate\n__global__\nvoid addToWeights( ValueTypeM *edge_weights, const ValueTypeV *x, const IndexType *row_indices, IndexType *col_indices, IndexType nnz, double scale )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < nnz )\n {\n int i = row_indices[tid];\n int j = col_indices[tid];\n edge_weights[tid] -= static_cast( scale * fabs( x[i] - x[j] ) );\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \n__global__\nvoid rescaleVector( ValueType *x, IndexType numRows )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < numRows )\n {\n x[tid] = 2 * x[tid] - 1;\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// -----------------\n// Methods\n// ----------------\n\n// Constructor\ntemplate\nMultiPairwiseSelectorBase::MultiPairwiseSelectorBase(AMG_Config &cfg, const std::string &cfg_scope)\n{\n deterministic = cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\");\n max_iterations = cfg.AMG_Config::template getParameter(\"max_matching_iterations\", cfg_scope);\n numUnassigned_tol = cfg.AMG_Config::template getParameter(\"max_unassigned_percentage\", cfg_scope);\n two_phase = cfg.AMG_Config::template getParameter(\"handshaking_phases\", cfg_scope) == 2;\n m_aggregation_edge_weight_component = cfg.AMG_Config::template getParameter(\"aggregation_edge_weight_component\", cfg_scope);\n aggregation_passes = cfg.AMG_Config::template getParameter(\"aggregation_passes\", cfg_scope); //default to size 8 aggregates. maybe its more convenient to have that as a config parameter\n filter_weights = cfg.AMG_Config::template getParameter(\"filter_weights\", cfg_scope); //by default: no filtering\n filter_weights_alpha = cfg.AMG_Config::template getParameter( \"filter_weights_alpha\", cfg_scope ); //default to 0.25\n full_ghost_level = cfg.AMG_Config::template getParameter( \"full_ghost_level\", cfg_scope ); //defaults to 0\n notay_weights = cfg.AMG_Config::template getParameter( \"notay_weights\", cfg_scope ); //defaults to 0\n ghost_offdiag_limit = cfg.AMG_Config::template getParameter( \"ghost_offdiag_limit\", cfg_scope ); //defaults to 0\n merge_singletons = cfg.AMG_Config::template getParameter( \"merge_singletons\", cfg_scope ); //defaults to 1\n weight_formula = cfg.AMG_Config::template getParameter( \"weight_formula\", cfg_scope ); //wheight formula defaults to 0\n serial_matching = cfg.AMG_Config::template getParameter( \"serial_matching\", cfg_scope ) != 0; //will use a serial matching algorithm instead of handshake\n modified_handshake = cfg.AMG_Config::template getParameter(\"modified_handshake\", cfg_scope ) == 1;\n //passes = 1 -> max = 3\n //passes = 2 -> max = 5\n //passes = 3 -> max = 10\n //passes = 4 -> max = 18\n max_aggregate_size = 2;\n\n for (int i = 1; i < aggregation_passes; i ++)\n {\n max_aggregate_size *= 2;\n }\n\n max_aggregate_size += aggregation_passes - (aggregation_passes / 2);\n mCfg = cfg;\n mCfg_scope = cfg_scope;\n}\n\n// setAggregates for block_dia_csr_matrix_h format\ntemplate \nvoid MultiPairwiseSelector >::setAggregates_common_sqblocks(Matrix_h &A,\n typename Matrix_h::IVector &aggregates, typename Matrix_h::IVector &aggregates_global, int &num_aggregates, MVector &edge_weights, IVector &sizes)\n{\n FatalError(\"MultiPairwise selector: setAggregates not implemented on CPU, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n// device specialization\n\n//edge_weights is an in/out parameter:\n//if its size is zero, the edge_weights will be computed from A and stored into edge_weights\n//else the edge_weights will not be computed and assumed to be valid for the given A. the value array of A is not used in this case\ntemplate \nvoid MultiPairwiseSelector >::setAggregates_common_sqblocks(Matrix_d &A,\n typename Matrix_d::IVector &aggregates,\n typename Matrix_d::IVector &aggregates_global,\n int &num_aggregates,\n MVector &edge_weights,\n IVector &sizes)\n{\n IndexType num_block_rows = (int) A.get_num_rows();\n IndexType num_nonzero_blocks = (int) A.get_num_nz();\n // both ways are supported\n IndexType total_nz = (A.is_matrix_singleGPU()) ? num_nonzero_blocks : A.manager->num_nz_all();\n typename Matrix_d::IVector &row_indices = A.row_indices;\n row_indices.resize( total_nz);\n cusp::detail::offsets_to_indices(A.row_offsets, row_indices);\n IndexType total_rows = (A.is_matrix_singleGPU()) ? A.get_num_rows() : A.manager->num_rows_all();\n aggregates.resize(total_rows);\n thrust_wrapper::fill(aggregates.begin(), aggregates.end(), -1);\n cudaCheckError();\n\n if ( this->merge_singletons == 2 && sizes.size() == 0 )\n {\n sizes.resize( total_rows, 1 ); //init with all ones\n }\n\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_row_indices_ptr = row_indices.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueType *A_nonzero_values_ptr = A.values.raw();\n typename Matrix_d::IVector strongest_neighbour(num_block_rows, -1);\n typename Matrix_d::IVector strongest_neighbour_1phase(num_block_rows, -1);\n Vector > aggregated(num_block_rows, 0);\n IndexType *strongest_neighbour_ptr = strongest_neighbour.raw();\n IndexType *strongest_neighbour_1phase_ptr = strongest_neighbour_1phase.raw();\n IndexType *aggregates_ptr = aggregates.raw();\n const int threads_per_block = 256;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_block_rows - 1) / threads_per_block + 1 );\n int numUnassigned = num_block_rows;\n int numUnassigned_previous = numUnassigned;\n bool computeWeights = ( edge_weights.size() == 0 );\n\n if (computeWeights)\n {\n if ( A.hasProps( DIAG ) )\n {\n edge_weights.resize( num_nonzero_blocks + num_block_rows, 0.0 );\n }\n else\n {\n edge_weights.resize( num_nonzero_blocks + 1, -1 ); //+1 is important to some algorithms\n }\n }\n\n ValueType *edge_weights_ptr = edge_weights.raw();\n ValueType *rand_edge_weights_ptr = NULL;\n cudaStream_t str = amgx::thrust::global_thread_handle::get_stream();\n\n // Compute the edge weights\n if ( computeWeights )\n {\n const int num_blocks_V2 = std::min( AMGX_GRID_MAX_SIZE, (num_nonzero_blocks - 1) / threads_per_block + 1);\n //compute with std formula\n cudaFuncSetCacheConfig(computeEdgeWeightsBlockDiaCsr_V2, cudaFuncCachePreferL1);\n computeEdgeWeightsBlockDiaCsr_V2 <<< num_blocks_V2, threads_per_block, 0, str>>>(A_row_offsets_ptr,\n A_row_indices_ptr,\n A_column_indices_ptr,\n A_dia_idx_ptr,\n A_nonzero_values_ptr,\n num_nonzero_blocks,\n edge_weights_ptr,\n rand_edge_weights_ptr,\n num_block_rows,\n A.get_block_dimy(),\n this->m_aggregation_edge_weight_component,\n this->weight_formula);\n cudaCheckError();\n }\n\n //filter weights if desired\n if ( this->filter_weights == 1 )\n {\n MVector tmp( edge_weights.size() );\n const int num_blocks_filter = std::min( AMGX_GRID_MAX_SIZE, (num_nonzero_blocks - 1) / threads_per_block + 1);\n cudaStreamSynchronize(str);\n cudaCheckError();\n filterWeights <<< num_blocks_filter, threads_per_block, 0, str>>>( A_row_offsets_ptr,\n A_row_indices_ptr,\n A_column_indices_ptr,\n A_dia_idx_ptr,\n edge_weights_ptr,\n tmp.raw(),\n num_nonzero_blocks,\n num_block_rows,\n this->filter_weights_alpha);\n cudaCheckError();\n cudaStreamSynchronize(str);\n cudaCheckError();\n tmp.swap( edge_weights );\n edge_weights_ptr = edge_weights.raw();\n }\n\n// compute matching\n if ( !this->serial_matching )\n {\n IVector degree;\n\n if ( this->modified_handshake )\n {\n degree.resize( num_block_rows );\n }\n\n#ifdef EXPERIMENTAL_ITERATIVE_MATCHING\n // TODO: allocate host pinned memory\n AsyncEvent *throttle_event = new AsyncEvent;\n throttle_event->create();\n typename Matrix_h::IVector h_unagg_vec(1);\n typename Matrix_d::IVector d_unagg_vec(1);\n int *unaggregated = h_unagg_vec.raw();\n int *d_unaggregated = d_unagg_vec.raw();\n#endif\n int icount, s = 1;\n {\n icount = 0;\n ValueType *weights_ptr = edge_weights_ptr;\n\n do\n {\n if ( !this->two_phase )\n {\n if ( this->modified_handshake )\n computeDegree <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr,\n A_column_indices_ptr,\n weights_ptr,\n aggregates_ptr,\n sizes.raw(),\n degree.raw(),\n num_block_rows,\n this->max_aggregate_size );\n cudaCheckError();\n\n // 1-phase handshaking\n if ( this->modified_handshake )\n findStrongestNeighbourBlockDiaCsr_V3\n <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr,\n A_column_indices_ptr,\n weights_ptr,\n num_block_rows,\n aggregates_ptr,\n strongest_neighbour_ptr,\n sizes.raw(),\n degree.raw(),\n A.get_block_dimy(),\n this->max_aggregate_size,\n this->merge_singletons);\n else\n findStrongestNeighbourBlockDiaCsr_V3\n <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr,\n A_column_indices_ptr,\n weights_ptr,\n num_block_rows,\n aggregates_ptr,\n strongest_neighbour_ptr,\n sizes.raw(),\n degree.raw(),\n A.get_block_dimy(),\n this->max_aggregate_size,\n this->merge_singletons);\n\n cudaCheckError();\n }\n else\n {\n // 2-phase handshaking\n findStrongestNeighbourBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, weights_ptr, num_block_rows, aggregates_ptr, strongest_neighbour_1phase_ptr, strongest_neighbour_ptr, A.get_block_dimy(), 1, this->merge_singletons);\n cudaCheckError();\n // 2nd phase: for each block_row, find the strongest neighbour among those who gave hand on 1st phase\n findStrongestNeighbourBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, weights_ptr, num_block_rows, aggregates_ptr, strongest_neighbour_1phase_ptr, strongest_neighbour_ptr, A.get_block_dimy(), 2, this->merge_singletons);\n cudaCheckError();\n }\n\n // Look for perfect matches. Also, for nodes without unaggregated neighbours, merge with aggregate containing strongest neighbour\n if ( this->merge_singletons == 2 )\n {\n matchEdges <<< num_blocks, threads_per_block, 0, str>>>(num_block_rows, aggregates_ptr, strongest_neighbour_ptr, sizes.raw());\n cudaCheckError();\n }\n else\n {\n matchEdges <<< num_blocks, threads_per_block, 0, str>>>(num_block_rows, aggregates_ptr, strongest_neighbour_ptr, (int *)NULL);\n }\n\n cudaCheckError();\n#ifdef EXPERIMENTAL_ITERATIVE_MATCHING\n s = (icount & 1);\n\n if ( s == 0 )\n {\n // count unaggregated vertices\n cudaError_t cuda_rc = cudaMemsetAsync(d_unaggregated, 0, sizeof(int), str);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemsetAsync failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n countAggregates <<< num_blocks, threads_per_block, 0, str>>>(num_block_rows, aggregates_ptr, d_unaggregated);\n cudaCheckError();\n cuda_rc = cudaMemcpyAsync(unaggregated, d_unaggregated, sizeof(int), cudaMemcpyDeviceToHost, str);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpyAsync failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n throttle_event->record(str);\n }\n else\n {\n throttle_event->sync();\n numUnassigned_previous = numUnassigned;\n numUnassigned = *unaggregated;\n }\n\n#else\n cudaError_t cuda_rc = cudaStreamSynchronize(str);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaStreamSynchronize failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n numUnassigned_previous = numUnassigned;\n numUnassigned = (int)amgx::thrust::count(aggregates.begin(), aggregates.begin() + num_block_rows, -1);\n cudaCheckError();\n#endif\n icount++;\n }\n while ( (s == 0) || !(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_block_rows < this->numUnassigned_tol || numUnassigned == numUnassigned_previous));\n\n }\n // printf(\"%i,\\n\", icount);\n#ifdef EXPERIMENTAL_ITERATIVE_MATCHING\n delete throttle_event;\n#endif\n }\n else\n {\n computeMatchingSerialGreedy( A, aggregates, num_aggregates, edge_weights );\n }\n\n if ( this->merge_singletons == 1 )\n {\n // Merge remaining vertices with current aggregates\n if (this->deterministic != 1)\n {\n while (numUnassigned != 0)\n {\n mergeWithExistingAggregatesBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, A.get_block_dimy(), this->deterministic, (IndexType *) NULL);\n cudaCheckError();\n numUnassigned = (int)amgx::thrust::count(aggregates.begin(), aggregates.begin() + num_block_rows, -1);\n cudaCheckError();\n }\n }\n else\n {\n typename Matrix_d::IVector aggregates_candidate(num_block_rows, -1);\n\n while (numUnassigned != 0)\n {\n mergeWithExistingAggregatesBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, A.get_block_dimy(), this->deterministic, aggregates_candidate.raw());\n cudaCheckError();\n joinExistingAggregates <<< num_blocks, threads_per_block, 0, str>>>(num_block_rows, aggregates_ptr, aggregates_candidate.raw());\n cudaCheckError();\n numUnassigned = (int)amgx::thrust::count(aggregates.begin(), aggregates.begin() + num_block_rows, -1);\n cudaCheckError();\n }\n\n aggregates_candidate.resize(0);\n }\n }\n else if (this->merge_singletons == 0 )\n {\n //make singletons\n aggregateSingletons <<< num_blocks, threads_per_block, 0, str>>>( aggregates_ptr, num_block_rows );\n cudaCheckError();\n }\n else if ( this->merge_singletons == 2 )\n {\n //merges all remaining singletons into adequate neighbors if possible\n mergeSingletonsSmart <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr,\n A_column_indices_ptr,\n edge_weights_ptr,\n aggregates_ptr,\n sizes.raw(),\n num_block_rows,\n this->max_aggregate_size);\n cudaCheckError();\n }\n\n //This will assign num_aggregates to the pseudo aggregate without counting it. Perfect!\n this->renumberAndCountAggregates(aggregates, aggregates_global, num_block_rows, num_aggregates);\n\n if ( this->merge_singletons == 2 )\n {\n //udpate the sizes vector, so it matches the renumbered aggregates size\n IVector sizesSource;\n sizesSource.swap( sizes );\n sizes.resize( num_aggregates, 1 );\n updateAggregateSizes <<< num_blocks, threads_per_block, 0, str>>>( sizesSource.raw(), sizes.raw(), aggregates_ptr, num_block_rows );\n cudaCheckError();\n }\n\n}\n\n\n//instead of a handshake, we use a serial greedy algorithm to compute a better matching\n//the algorithm:\n// 1. compute degree of every node and sort nodes by degree into double linked list\n// 2. while non-isolated nodes left:\n// take node with minimum degree > 0\n// find strongest edge to unaggregated node and assign to new aggregate\n// remove both nodes from linked list\n// decrease degree of each neighbor by one for each of the two nodes\n// update list\ntemplate \nvoid MultiPairwiseSelector >::computeMatchingSerialGreedy( const Matrix_d &A, IVector &aggregates, int &numAggregates, MVector &edge_weights)\n{\n IndexType numRows = A.row_offsets.size() - 1;\n IndexType nnz = A.col_indices.size();\n //allocate memory on host\n IndexType *ia = new IndexType[numRows + 1];\n IndexType *ja = new IndexType[nnz];\n ValueType *w = new ValueType[nnz];\n IndexType *agg = new IndexType[numRows];\n IndexType *deg = new IndexType[numRows];\n //copy\n cudaError_t cuda_rc = cudaMemcpy( ia, A.row_offsets.raw(), sizeof(IndexType) * (numRows + 1), cudaMemcpyDeviceToHost );\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy row_offsets D2H failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = cudaMemcpy( ja, A.col_indices.raw(), sizeof(IndexType)*nnz, cudaMemcpyDeviceToHost );\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy col_indices D2H failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = cudaMemcpy( w, edge_weights.raw(), sizeof(ValueType)*nnz, cudaMemcpyDeviceToHost );\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy edge_weights D2H failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n //init agg and compute the degree of each aggregate\n int max_degree = 0;\n\n for (IndexType i = 0; i < numRows; i++)\n {\n agg[i] = -1;\n int degree = 0;\n\n for (IndexType ii = ia[i]; ii < ia[i + 1]; ii++) //only care for positive weights\n {\n if ( ja[ii] != i && w[ii] > 0.0 )\n {\n degree++;\n }\n }\n\n if ( degree > max_degree )\n {\n max_degree = degree;\n }\n\n deg[i] = degree;\n }\n\n if ( max_degree >= numRows )\n {\n FatalError( \"max degree is greater than numRows.\", AMGX_ERR_UNKNOWN );\n }\n\n //init double linked list\n IndexType *fwd = new IndexType[numRows + max_degree + 1];\n IndexType *bwd = new IndexType[numRows + max_degree + 1];\n\n for (IndexType i = 0; i < numRows + max_degree + 1; i++)\n {\n fwd[i] = i;\n bwd[i] = i;\n }\n\n IndexType nodesLeft = numRows;\n numAggregates = 0;\n\n //insert nodes into list\n for (IndexType i = numRows - 1; i >= 0; i--) //inserting in backward order the nodes will be sorted by index in case of same degree\n {\n //insert forward following root\n fwd[i] = fwd[numRows + deg[i]];\n fwd[numRows + deg[i]] = i;\n //insert backward\n bwd[i] = numRows + deg[i];\n bwd[fwd[i]] = i;\n\n //isolated nodes cannot be aggregated\n if ( deg[i] == 0 )\n {\n nodesLeft--;\n }\n }\n\n while ( nodesLeft > 0 )\n {\n IndexType node = numRows;\n int degree;\n\n for (degree = 1; degree <= max_degree; degree++)\n {\n //list not empty -> select node\n if ( fwd[numRows + degree] < numRows ) //selecting the first node will select the most recently inserted one or the one with lowest index. both is preferable\n {\n node = fwd[numRows + degree];\n }\n\n if ( node < numRows )\n {\n break;\n }\n }\n\n //no node with degree > 1 found even though nodesLeft > 0\n if ( node == numRows )\n {\n FatalError(\"nodeLeft counting or list invalid\", AMGX_ERR_UNKNOWN );\n }\n\n if ( agg[node] != -1 )\n {\n FatalError(\"node is already aggregated\", AMGX_ERR_UNKNOWN );\n }\n\n //find strongest edge\n ValueType max_weight = 0.0;\n IndexType max_node = numRows; //use this as gatekeeper, so if weight == 0 the node index will not be greater than this\n\n for (IndexType ii = ia[node]; ii < ia[node + 1]; ii++)\n {\n IndexType j = ja[ii];\n\n if ( agg[j] != -1 || j == node)\n {\n continue;\n }\n\n if ( w[ii] > 0.0 )\n {\n degree--;\n }\n\n //deterministic, doesn't selects 0 weight.\n if ( w[ii] > max_weight || (w[ii] == max_weight && j > max_node) ) //always taking the edge pointing to the max node can give good alignment if numbering is structured\n {\n max_node = j;\n max_weight = w[ii];\n }\n } //Note that there has to be at least one neighbor node because degree of node is at least 1.\n\n if ( max_node == numRows )\n {\n FatalError( \"node has no neighbor although degree of node is at least 1\", AMGX_ERR_UNKNOWN );\n }\n\n if ( degree != 0 )\n {\n FatalError( \"node degree corrupted\", AMGX_ERR_UNKNOWN );\n }\n\n //aggregate\n agg[node] = node;\n agg[max_node] = node;\n numAggregates++;\n nodesLeft -= 2;\n //remove from list\n fwd[bwd[node]] = fwd[node];\n bwd[fwd[node]] = bwd[node];\n fwd[bwd[max_node]] = fwd[max_node];\n bwd[fwd[max_node]] = bwd[max_node];\n\n //update neighbors and list\n //max_node first\n for (IndexType ii = ia[max_node]; ii < ia[max_node + 1]; ii++)\n {\n IndexType j = ja[ii];\n\n if ( agg[j] != -1 || w[ii] <= 0.0)\n {\n continue;\n }\n\n //remove j from list\n fwd[bwd[j]] = fwd[j];\n bwd[fwd[j]] = bwd[j];\n //update degree of j\n deg[j]--;\n //add j back to start of the list\n fwd[j] = fwd[numRows + deg[j]];\n bwd[j] = numRows + deg[j];\n bwd[fwd[j]] = j;\n fwd[bwd[j]] = j;\n\n if (deg[j] == 0)\n {\n nodesLeft--;\n }\n }\n\n //node second, this will prefer nodes neighbors over max_nodes neighbors when choosing the next node\n for (IndexType ii = ia[node]; ii < ia[node + 1]; ii++)\n {\n IndexType j = ja[ii];\n\n if ( agg[j] != -1 || w[ii] <= 0.0)\n {\n continue;\n }\n\n //remove j from list\n fwd[bwd[j]] = fwd[j];\n bwd[fwd[j]] = bwd[j];\n //update degree of j\n deg[j]--;\n //add j back to start of the list\n fwd[j] = fwd[numRows + deg[j]];\n bwd[j] = numRows + deg[j];\n bwd[fwd[j]] = j;\n fwd[bwd[j]] = j;\n\n if (deg[j] == 0)\n {\n nodesLeft--;\n }\n }\n }\n\n //copy result back to device\n cuda_rc = cudaMemcpy( aggregates.raw(), agg, sizeof(IndexType)*numRows, cudaMemcpyHostToDevice );\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy aggregates H2D failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n\n //assert matching\n for (IndexType node = 0; node < numRows; node++)\n {\n if ( agg[node] == -1 )\n {\n continue;\n }\n\n for ( IndexType partner = 0; partner < numRows; partner++)\n {\n if ( agg[partner] == agg[node] )\n {\n if ( partner == node )\n {\n continue;\n }\n\n bool neighbor = false;\n\n for (IndexType ii = ia[node]; ii < ia[node + 1]; ii++)\n if ( ja[ii] == partner )\n {\n neighbor = true;\n break;\n }\n\n if ( !neighbor )\n {\n for (IndexType ii = ia[partner]; ii < ia[partner + 1]; ii++)\n if ( ja[ii] == node )\n {\n neighbor = true;\n break;\n }\n }\n\n if ( !neighbor )\n {\n FatalError(\"Internal error in aggregation selector\", AMGX_ERR_INTERNAL);\n }\n }\n }\n }\n\n //you shall not leak memory\n delete[] ia;\n delete[] ja;\n delete[] w;\n delete[] agg;\n delete[] deg;\n delete[] fwd;\n delete[] bwd;\n}\n\n\n\n\n//this kernel merges aggregate2 into aggregate1\ntemplate\n__global__\nvoid mergeAggregates(IndexType *aggregate1, const IndexType *aggregate2, IndexType sizeAggregate1, IndexType sizeAggregate2, IndexType sizeAggregate3)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < sizeAggregate1 )\n {\n if ( aggregate1[tid] == sizeAggregate2 )\n {\n aggregate1[tid] = sizeAggregate3;\n }\n else\n {\n aggregate1[tid] = aggregate2[aggregate1[tid]];\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\n\ntemplate \nvoid MultiPairwiseSelector >::computeIncompleteGalerkin( const Matrix_h &A,\n Matrix_h &Ac,\n const typename Matrix_h::IVector &aggregates,\n const typename Matrix_h::IVector &R_row_offsets,\n const typename Matrix_h::IVector &R_column_indices,\n const int num_aggregates )\n{\n FatalError(\"computeIncomlpetegalerkin is not supported on host. Run with ghost_offdiag_limit=0 instead.\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MultiPairwiseSelector >::computeIncompleteGalerkin( const Matrix_d &A,\n Matrix_d &Ac,\n const typename Matrix_d::IVector &aggregates,\n const typename Matrix_d::IVector &R_row_offsets,\n const typename Matrix_d::IVector &R_column_indices,\n const int num_aggregates )\n{\n FatalError(\"computeIncomlpetegalerkin is not implemented yet. run with ghost_offdiag_limit=0 instead.\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate\nvoid MultiPairwiseSelectorBase::assertRestriction( const IVector &R_row_offsets, const IVector &R_col_indices, const IVector &aggregates )\n{\n int *r_ia = new int[R_row_offsets.size()];\n int *r_ja = new int[R_col_indices.size()];\n int *agg = new int[aggregates.size()];\n int *used_col = new int[aggregates.size()];\n\n for ( int i = 0; i < aggregates.size(); i++ )\n {\n used_col[i] = 0;\n }\n\n cudaError_t cuda_rc = cudaMemcpy( r_ia, R_row_offsets.raw(), sizeof(int)*R_row_offsets.size(), cudaMemcpyDeviceToHost );\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy R_row_offsets D2H failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = cudaMemcpy( r_ja, R_col_indices.raw(), sizeof(int)*R_col_indices.size(), cudaMemcpyDeviceToHost );\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy R_col_indices D2H failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = cudaMemcpy( agg, aggregates.raw(), sizeof(int)*aggregates.size(), cudaMemcpyDeviceToHost );\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMemcpy aggregates D2H failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n\n for ( int i = 0; i < R_row_offsets.size() - 1; i++ )\n {\n for ( int ii = r_ia[i]; ii < r_ia[i + 1]; ii++ )\n {\n int j = r_ja[ii];\n used_col[j]++;\n\n if ( used_col[j] > 1 )\n {\n std::cout << \"column \" << j << \" is present at least \" << used_col[j] << \" times\" << std::endl;\n }\n\n if ( j < 0 || j >= aggregates.size() )\n {\n std::cout << \"Error: j out of bounds, j = \" << j << \" and numRows = \" << aggregates.size() << std::endl;\n }\n else if ( agg[j] != i )\n {\n std::cout << \"Error: agg[\" << j << \"] = \" << agg[j] << \" != \" << i << std::endl;\n }\n }\n }\n\n std::cout << \"assert restriction done\" << std::endl;\n}\n\ntemplate\nvoid MultiPairwiseSelectorBase::setAggregates(Matrix &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n if (A.get_block_dimx() == A.get_block_dimy())\n {\n //ghost level matrix. this is a probably a weight matrix\n Matrix ghostA;\n ghostA.values.resize(0);\n //prolongation and restriction operator. this is only needed in when LowDegCoarseAGenerator is used\n IVector R_row_offsets;\n IVector R_col_indices;\n //holds the size of each aggregate\n IVector sizes;\n sizes.resize(0);\n //aggregates for ghost level\n IVector aggregates_current;\n IVector aggregates_global_current;\n bool aggregates_initialized = true;\n\n if (aggregates.size() == 0)\n {\n aggregates_initialized = false;\n\n if (!A.is_matrix_singleGPU())\n {\n aggregates.resize(A.manager->halo_offset(A.manager->num_neighbors()));\n }\n else\n {\n aggregates.resize(A.get_num_rows());\n }\n }\n\n //for mergeAggregates kernel\n const int threads_per_block = 256;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (A.get_num_rows() - 1) / threads_per_block + 1 );\n cudaStream_t stream = amgx::thrust::global_thread_handle::get_stream();\n //initialize and prepare weight matrix\n Matrix w;\n w.set_initialized(0);\n w.addProps(CSR);\n w.delProps(COO);\n w.setColsReorderedByColor(false);\n w.resize( 0, 0, 0, 1, 1, true ); //empty scalar 0x0 matrix\n w.values.resize(0); //matrix resize sets the values array to nnz+1 for no apparent reason\n IndexType targetSize = 1;\n //initialize coarse A generator\n CoarseAGenerator *cag;\n const bool use_restriction = true;\n const bool shrink_ghost_level = false;\n cag = new LowDegCoarseAGenerator(mCfg, mCfg_scope);\n // This will make coarseAGenerator to allocate more memory inside of galerkin\n ghostA.manager = new DistributedManager();\n w.manager = new DistributedManager();\n Matrix *curA = &A;\n\n //foreach pass do:\n // 1. build aggregates and weights\n // 2. create weight matrix (in full_ghost_level mode this is the input matrix or the last ghostlevel matrix)\n // 3. if in full ghost level mode, build R\n // 4. compute next level\n for (int current_pass = 1; true; current_pass++)\n {\n const IndexType numRows = curA->get_num_rows();\n const IndexType nnz = curA->get_num_nz();\n targetSize *= 2;\n\n if ( full_ghost_level )\n {\n w.values.resize(0); //compute weights from curA\n }\n else\n {\n w.values.swap( ghostA.values ); //use the weights computed with the galerkin operator (this will do nothing in the first pass, both values have size 0)\n }\n\n // create aggregates from correct input matrix\n setAggregates_common_sqblocks( *curA, aggregates_current, aggregates_global_current, num_aggregates, w.values, sizes );\n\n\n if ( current_pass > 1 )\n {\n //merge original aggregate with the newly created ones\n mergeAggregates <<< num_blocks, threads_per_block, 0, stream >>>( aggregates.raw(), aggregates_current.raw(), A.get_num_rows(), numRows, num_aggregates );\n cudaCheckError();\n //mergeAggregates<<< num_blocks, threads_per_block, 0, stream >>>( aggregates_global.raw(), aggregates_global_current.raw(), A.get_num_rows() );\n //cudaCheckError();\n }\n\n //try to free memory\n if ( full_ghost_level )\n {\n //then we don't need to save the weights, only for original level to do post processing\n w.values.resize(0);\n }\n else\n {\n //save edge weights for original level later\n //in that case we can throw away the values of ghostA as we will use the values to compute the next ghost level\n ghostA.values.resize(0);\n }\n\n // this is the break condition for the loop\n if ( current_pass >= aggregation_passes || num_aggregates <= 1 || num_aggregates == numRows)\n {\n //this means, aggregates has not been initialized yet\n if ( !aggregates_initialized )\n {\n aggregates.swap( aggregates_current );\n }\n\n if ( !aggregates_initialized )\n {\n aggregates_global.swap( aggregates_global_current );\n }\n\n cudaStreamSynchronize( stream );\n cudaCheckError();\n\n break;\n }\n\n //prepare A to be corrupted\n curA->set_initialized(0);\n //swap in ia, ja from curA\n w.row_offsets.swap( curA->row_offsets );\n w.col_indices.swap( curA->col_indices );\n\n if ( full_ghost_level )\n {\n if ( shrink_ghost_level && curA->get_block_dimx() > 1)\n {\n //set w to correct size\n w.values.resize( nnz );\n //define grid and offsets\n const int num_blocks_inter = std::min( (int)AMGX_GRID_MAX_SIZE, (int)(nnz - 1) / threads_per_block + 1 );\n const int sq_blocksize = A.get_block_dimx() * A.get_block_dimy();\n const int index_offset = A.get_block_dimy() * m_aggregation_edge_weight_component + m_aggregation_edge_weight_component;\n //do the interleaved copy\n gatherValuesInterleaved <<< num_blocks_inter, threads_per_block, 0, stream>>>( A.values.raw(), w.values.raw(), nnz, sq_blocksize, index_offset );\n cudaCheckError();\n cudaStreamSynchronize( stream );\n cudaCheckError();\n }\n else\n {\n w.values.swap( curA->values );\n }\n }\n\n w.diag.swap( curA->diag );\n\n //resize to inform the matrix of its new size\n if ( full_ghost_level && !shrink_ghost_level )\n {\n w.set_block_dimx( A.get_block_dimx() );\n w.set_block_dimy( A.get_block_dimy() );\n }\n else\n {\n w.set_block_dimx( 1 );\n w.set_block_dimy( 1 );\n }\n\n w.set_num_rows( numRows );\n w.set_num_cols( numRows );\n w.set_num_nz( nnz );\n w.set_allow_recompute_diag( false );\n\n if ( curA->hasProps( DIAG ) )\n {\n w.addProps( DIAG );\n }\n\n //ready to use\n w.set_initialized(1);\n\n //compute restriction operator\n if ( use_restriction )\n {\n IVector R_row_indices(aggregates_current);\n R_row_offsets.resize(num_aggregates + 2);\n R_col_indices.resize(numRows);\n thrust_wrapper::sequence(R_col_indices.begin(), R_col_indices.end());\n cudaCheckError();\n amgx::thrust::sort_by_key(R_row_indices.begin(), R_row_indices.end(), R_col_indices.begin());\n cudaCheckError();\n cusp::detail::indices_to_offsets(R_row_indices, R_row_offsets);\n cudaCheckError();\n //delete last row, which holds the pseudo aggregate\n R_row_offsets.resize( num_aggregates + 1);\n R_col_indices.resize( R_row_offsets[num_aggregates] );\n }\n\n // 3. compute galerkin ghost level\n if ( ghost_offdiag_limit == 0 )\n {\n //compute full galerkin\n cag->computeAOperator(w,\n ghostA,\n aggregates_current,\n R_row_offsets,\n R_col_indices,\n num_aggregates);\n }\n else\n {\n //compute incomplete galerkin\n computeIncompleteGalerkin(w,\n ghostA,\n aggregates_current,\n R_row_offsets,\n R_col_indices,\n num_aggregates);\n }\n\n //from now on w will be destroyed again.\n w.set_initialized(0);\n\n //repair the original A matrix. its ia and ja are in w\n if ( current_pass == 1 )\n {\n //swap back\n w.row_offsets.swap( A.row_offsets );\n w.col_indices.swap( A.col_indices );\n\n //only in that case we have swapped the values\n if ( full_ghost_level && !shrink_ghost_level )\n {\n w.values.swap( A.values );\n }\n\n //save the edge weights of the original level\n A.diag.swap( w.diag );\n A.set_initialized(1); //A is repaired now\n //save the first aggregates into the original aggregate vector so we can merge them later\n aggregates.swap( aggregates_current );\n aggregates_global.swap( aggregates_global_current );\n aggregates_initialized = true;\n curA = &ghostA;\n }\n }\n\n delete cag;\n }\n else\n {\n FatalError(\"Unsupported block size for MultiPairwise\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\n// -------------------------\n// Explict instantiations\n// -------------------------\n#define AMGX_CASE_LINE(CASE) template class MultiPairwiseSelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class MultiPairwiseSelector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n}\n}\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include //count\n#include //sort\n#include //lower_bound\n#include //unique\n#include \n#include //offsets_to_indices\n#include \n#include \n\n#include \n#include \n\n#define EXPERIMENTAL_ITERATIVE_MATCHING\n\nnamespace amgx\n{\nnamespace aggregation\n{\nnamespace multi_pairwise\n{\n\n// include common routines for all selectors\n#include \n\n// ------------------------\n// Kernels\n// ------------------------\n\n__device__\nfloat random_weight2(int i, int j)\n{\n#define RAND_MULTIPLIER 1145637293\n unsigned long i_min = (min(i, j) * RAND_MULTIPLIER);\n unsigned long i_max = (max(i, j) * RAND_MULTIPLIER);\n return ((float)i_min / i_max);\n}\n\n__device__\nunsigned long random_weight3(int i, int j)\n{\n unsigned long a;\n a = (i + j) ^ 8;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return a;\n}\n\n\n// findStrongestNeighbour kernel for block_dia_csr_matrix format\n// Reads the weight from edge_weights array\ntemplate \n__global__\nvoid findStrongestNeighbourBlockDiaCsr_V2(const IndexType *row_offsets, const IndexType *column_indices,\n MatrixValueType *edge_weights, IndexType num_block_rows, IndexType *aggregates,\n IndexType *strongest_neighbour_1phase, IndexType *strongest_neighbour,\n const size_t bsize, int phase, int merge_singletons)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n MatrixValueType weight;\n int jcol;\n\n while (tid < num_block_rows)\n {\n int strongest_unaggregated = -1;\n int strongest_aggregated = -1;\n MatrixValueType max_weight_unaggregated = 0.;\n MatrixValueType max_weight_aggregated = 0.;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (phase == 1) { weight = edge_weights[j]; }\n else { weight = random_weight2(tid, jcol); }\n\n if (tid == jcol || jcol >= num_block_rows) { continue; } // skip diagonal and halo\n\n if (phase == 2 && strongest_neighbour_1phase[jcol] != tid) { continue; } // if 2nd phase only accept those who gave a hand on the 1st phase\n\n // Identify strongest aggregated and unaggregated neighbours (method by multi_pairwise)\n if (aggregates[jcol] == -1 && weight > 0.0 && (weight > max_weight_unaggregated || (weight == max_weight_unaggregated && random_weight3(tid, jcol) > random_weight3(tid, strongest_unaggregated)))) // unaggregated\n {\n max_weight_unaggregated = weight;\n strongest_unaggregated = jcol;\n }\n else if (aggregates[jcol] != -1 && weight > 0.0 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && random_weight3(tid, jcol) > random_weight3(tid, strongest_aggregated)))) // aggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_unaggregated == -1 && strongest_aggregated != -1) // All neighbours are aggregated\n {\n if ( merge_singletons == 1 )\n // Put in same aggregate as strongest neighbour\n {\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n else\n // create singleton\n {\n aggregates[tid] = tid;\n }\n }\n else if (strongest_unaggregated != -1)\n {\n if (phase == 2)\n {\n MatrixValueType rand_w1 = random_weight2(tid, strongest_neighbour_1phase[tid]);\n strongest_neighbour[tid] = max_weight_unaggregated > rand_w1 ? strongest_unaggregated : strongest_neighbour_1phase[tid];\n }\n else { strongest_neighbour_1phase[tid] = strongest_unaggregated; }\n }\n else\n {\n if (phase == 2) { strongest_neighbour[tid] = strongest_neighbour_1phase[tid]; }\n else { strongest_neighbour_1phase[tid] = tid; }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate \n__device__\nbool atomicJoin( IndexType node, IndexType aggregate, IndexType *aggregates, int *sizes, int allowed )\n{\n int mySize = sizes[node];\n int theirSize = sizes[aggregate];\n int theirSizeOld = theirSize;\n\n do\n {\n int newSize = mySize + theirSize;\n\n if ( newSize > allowed )\n {\n return false;\n }\n\n theirSizeOld = theirSize;\n theirSize = atomicCAS( &sizes[aggregate], theirSize, newSize );\n }\n while ( theirSize != theirSizeOld );\n\n aggregates[node] = aggregate;\n return true;\n}\n\ntemplate \n__global__\nvoid findStrongestNeighbourBlockDiaCsr_V3(const IndexType *row_offsets,\n const IndexType *column_indices,\n MatrixValueType *edge_weights,\n IndexType num_block_rows,\n IndexType *aggregates,\n IndexType *strongest_neighbour,\n int *sizes,\n int *degree,\n const size_t bsize,\n int max_aggregate_size,\n int merge_singletons)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n MatrixValueType weight;\n int jcol;\n\n while (tid < num_block_rows)\n {\n int strongest_unaggregated = -1;\n int strongest_aggregated = -1;\n int lowest_degree;\n\n if ( use_degree )\n {\n lowest_degree = degree[tid]; //only interested in finding lower degree than self\n }\n else\n {\n lowest_degree = 0; //if we decide to not use degree than just propose to the strongest edge\n }\n\n int lowest_degree_neighbor = tid;\n MatrixValueType lowest_degree_weight = 1e100; //high start value, so that same degree neighbor won't become lowest degree neighbor\n MatrixValueType max_weight_unaggregated = 0.;\n MatrixValueType max_weight_aggregated = 0.;\n int mySize;\n\n if ( merge_singletons == 2 )\n {\n mySize = sizes[tid];\n }\n else\n {\n mySize = 0;\n }\n\n if ( merge_singletons != 2 )\n {\n max_aggregate_size = 100000;\n }\n\n //this aggregate is already full\n if (mySize == max_aggregate_size)\n {\n aggregates[tid] = tid;\n }\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n jcol = column_indices[j];\n\n if (tid == jcol || jcol >= num_block_rows) { continue; } // skip diagonal and halo\n\n weight = edge_weights[j];\n\n if (weight <= 0.0) { continue; }\n\n if ( aggregates[jcol] != -1 ) //aggregated neighbor\n {\n int theirSize;\n\n if ( merge_singletons == 2 )\n {\n theirSize = aggregates[sizes[jcol]];\n }\n else\n {\n theirSize = 0;\n }\n\n //if all neighbors are aggregated, find the strongest edge to neighbor aggregate that is not full yet\n if (mySize + theirSize <= max_aggregate_size &&\n (weight > max_weight_aggregated)) // aggregated\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n else //unaggregated neighbor\n {\n if ( use_degree && merge_singletons == 2 )\n {\n int theirSize = sizes[jcol];\n //get lowest degree neighbor or find out that there is no lower degree neighbor\n int current_degree = degree[jcol];\n\n if ( mySize + theirSize <= max_aggregate_size && (current_degree < lowest_degree || current_degree == lowest_degree && weight > lowest_degree_weight) )\n {\n lowest_degree = current_degree;\n lowest_degree_weight = weight;\n lowest_degree_neighbor = jcol;\n }\n\n //get highest weight neighbor\n if ( mySize + theirSize <= max_aggregate_size && (weight > max_weight_unaggregated) )\n {\n strongest_unaggregated = jcol;\n max_weight_unaggregated = weight;\n }\n }\n\n if ( use_degree && merge_singletons != 2 ) //same as above but ignore sizes\n {\n //get lowest degree neighbor or find out that there is no lower degree neighbor\n int current_degree = degree[jcol];\n\n if ( current_degree < lowest_degree || current_degree == lowest_degree && weight > lowest_degree_weight)\n {\n lowest_degree = current_degree;\n lowest_degree_weight = weight;\n lowest_degree_neighbor = jcol;\n }\n\n //get highest weight neighbor\n if (weight > max_weight_unaggregated)\n {\n strongest_unaggregated = jcol;\n max_weight_unaggregated = weight;\n }\n }\n\n if ( !use_degree && merge_singletons == 2 )\n {\n //get highest weight neighbor only but pay attention to the aggregate sizes\n int theirSize = sizes[jcol]; //get highest weight neighbor\n\n if ( mySize + theirSize <= max_aggregate_size && (weight > max_weight_unaggregated) )\n {\n strongest_unaggregated = jcol;\n max_weight_unaggregated = weight;\n }\n }\n\n if ( !use_degree && merge_singletons != 2 )\n {\n //just highest weight\n if (weight > max_weight_unaggregated)\n {\n strongest_unaggregated = jcol;\n max_weight_unaggregated = weight;\n }\n }\n }\n }\n\n //prefer lowest degree neighbor\n if ( lowest_degree_neighbor != tid )\n {\n strongest_unaggregated = lowest_degree_neighbor;\n }\n\n if (strongest_unaggregated != -1) //Unaggregated neighbor exists\n {\n strongest_neighbour[tid] = strongest_unaggregated; //assign strongest aggregated\n }\n\n if (strongest_unaggregated == -1 && strongest_aggregated != -1) // All neighbours are aggregated but small enough aggregated neighbors exist\n {\n if ( merge_singletons == 0 )\n {\n aggregates[tid] = tid;\n }\n\n if ( merge_singletons == 1 )\n {\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n\n if ( merge_singletons == 2)\n {\n atomicJoin( tid, aggregates[strongest_aggregated], aggregates, sizes, max_aggregate_size ); //try to join, can fail. maybe it works next round.\n }\n }\n\n if (strongest_unaggregated == -1 && strongest_aggregated == -1) //no feasable neighbor at all, become singleton\n {\n strongest_neighbour[tid] = tid; //become singleton\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \n__global__\nvoid computeDegree( const IndexType *ia, const IndexType *ja, const ValueType *weights, IndexType *aggregates, IndexType *sizes, IndexType *degree, IndexType numRows, IndexType max_aggregate_size)\n{\n int i = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( i < numRows )\n {\n int myDegree = 0;\n IndexType ia_ip1 = ia[i + 1];\n\n for ( IndexType ii = ia[i]; ii < ia_ip1; ii++ )\n {\n IndexType j = ja[ii];\n\n if ( j == i )\n {\n continue;\n }\n\n int mySize, theirSize;\n\n if ( sizes != NULL )\n {\n mySize = sizes[i];\n theirSize = sizes[j];\n }\n else\n {\n mySize = theirSize = 0;\n }\n\n if ( weights[ii] > 0.0 && aggregates[j] == -1 && mySize + theirSize <= max_aggregate_size )\n {\n myDegree++;\n }\n }\n\n degree[i] = myDegree;\n i += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate \n__global__\nvoid mergeSingletonsSmart(const IndexType *ia, const IndexType *ja, const ValueType *weights, IndexType *aggregates, IndexType *sizes, IndexType numRows, int max_aggregate_size)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < numRows )\n {\n //unaggregated nodes try to join or create their own aggregate\n if ( aggregates[tid] == -1 )\n {\n bool joined = false;\n\n while ( !joined )\n {\n int neighbor_aggregate = -1;\n ValueType max_weight = 0.0;\n IndexType mySize = sizes[tid];\n\n for (IndexType ii = ia[tid]; ii < ia[tid + 1]; ii++)\n {\n IndexType j = ja[ii];\n\n if (j == tid || j >= numRows) { continue; }\n\n if ( aggregates[j] != -1 && sizes[aggregates[j]] + mySize <= max_aggregate_size && weights[ii] > max_weight )\n {\n neighbor_aggregate = aggregates[j];\n max_weight = weights[ii];\n }\n }\n\n //no possible neighbor found\n if ( neighbor_aggregate == -1 )\n {\n //create own aggregate\n aggregates[tid] = tid;\n joined = true;\n }\n else\n {\n //try to join\n joined = atomicJoin( tid, neighbor_aggregate, aggregates, sizes, max_aggregate_size );\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate \n__global__\nvoid updateAggregateSizes( IndexType *sizesSource, IndexType *sizes, IndexType *aggregates, IndexType numRows )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < numRows )\n {\n IndexType agg = aggregates[tid];\n IndexType aggregateSize = sizes[agg];\n IndexType mySize = sizesSource[tid];\n\n while ( mySize > aggregateSize )\n {\n aggregateSize = atomicCAS( &sizes[agg], aggregateSize, mySize );\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\n// Kernel that checks if perfect matchs exist\ntemplate \n__global__\nvoid matchEdges(const IndexType num_rows, IndexType *aggregates, int *strongest_neighbour, IndexType *sizes)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int potential_match, potential_match_neighbour;\n\n while (tid < num_rows)\n {\n if (aggregates[tid] == -1) // Unaggregated row\n {\n potential_match = strongest_neighbour[tid];\n potential_match_neighbour = strongest_neighbour[potential_match];\n\n if ( potential_match == tid )\n {\n aggregates[tid] = tid;\n }\n else if (potential_match != -1 && potential_match_neighbour == tid && tid < potential_match) // we have a match\n {\n aggregates[tid] = tid;\n aggregates[potential_match] = tid;\n\n if ( sizes != NULL)\n {\n sizes[tid] += sizes[potential_match];\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \n__global__\nvoid countAggregates(const IndexType num_rows, IndexType *aggregates, int *num_unaggregated)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int c = 0;\n int i = tid;\n\n while ( i < num_rows )\n {\n c += ( aggregates[i] == -1 );\n i += gridDim.x * blockDim.x;\n }\n\n __shared__ volatile int smem[block_size];\n smem[threadIdx.x] = c;\n __syncthreads();\n\n for ( int off = blockDim.x / 2; off >= 32; off = off / 2 )\n {\n if ( threadIdx.x < off )\n {\n smem[threadIdx.x] += smem[threadIdx.x + off];\n }\n\n __syncthreads();\n }\n\n // warp reduce\n if ( threadIdx.x < 32 )\n {\n smem[threadIdx.x] += smem[threadIdx.x + 16];\n smem[threadIdx.x] += smem[threadIdx.x + 8];\n smem[threadIdx.x] += smem[threadIdx.x + 4];\n smem[threadIdx.x] += smem[threadIdx.x + 2];\n smem[threadIdx.x] += smem[threadIdx.x + 1];\n }\n\n if ( threadIdx.x == 0 )\n {\n atomicAdd(num_unaggregated, smem[0]);\n }\n}\n\ntemplate \n__global__\nvoid joinExistingAggregates(IndexType num_rows, IndexType *aggregates, IndexType *aggregates_candidate)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while (tid < num_rows)\n {\n if (aggregates[tid] == -1 && aggregates_candidate[tid] != -1) // Unaggregated row\n {\n aggregates[tid] = aggregates_candidate[tid];\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel that merges unaggregated vertices its strongest aggregated neighbour\n// Weights are read from edge_weights array\n// For block_dia_csr_matrix_format\ntemplate \n__global__\nvoid mergeWithExistingAggregatesBlockDiaCsr_V2(const IndexType *row_offsets, const IndexType *column_indices, const MatrixValueType *edge_weights,\n const int num_block_rows, IndexType *aggregates, int bsize, const int deterministic, IndexType *aggregates_candidate)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int jcol;\n MatrixValueType weight;\n\n while (tid < num_block_rows)\n {\n MatrixValueType max_weight_aggregated = 0.;\n int strongest_aggregated = -1;\n\n if (aggregates[tid] == -1) // Unaggregated row\n {\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n // Compute edge weight\n weight = edge_weights[j];\n jcol = column_indices[j];\n\n if (jcol == tid || jcol >= num_block_rows) { continue; } // skip diagonal\n\n if ( aggregates[jcol] == num_block_rows ) { continue; } // skip dd rows\n\n // Identify strongest aggregated neighbour\n if (aggregates[jcol] != -1 && weight > 0 && (weight > max_weight_aggregated || (weight == max_weight_aggregated && random_weight3( tid, jcol ) > random_weight3( tid, strongest_aggregated )))) //\n {\n max_weight_aggregated = weight;\n strongest_aggregated = jcol;\n }\n }\n\n if (strongest_aggregated != -1) // Found a neighbour to aggregate to\n {\n if (deterministic == 1)\n {\n aggregates_candidate[tid] = aggregates[strongest_aggregated];\n }\n else\n {\n // Put in same aggregate as strongest neighbour\n aggregates[tid] = aggregates[strongest_aggregated];\n }\n }\n else // All neighbours are unaggregated, leave alone\n {\n if (deterministic == 1)\n {\n aggregates_candidate[tid] = tid;\n }\n else\n {\n aggregates[tid] = tid;\n }\n }\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel to extract diagonal for csr_matrix format\ntemplate \n__global__\nvoid getDiagonalKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, const IndexType numRows, ValueType *diagonal)\n{\n int tIdx = threadIdx.x + blockDim.x * blockIdx.x;\n\n while (tIdx < numRows)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n for (int j = offset; j < offset + numj; j++)\n {\n int jcol = column_indices[j];\n\n if (tIdx == jcol)\n {\n diagonal[tIdx] = values[j];\n }\n }\n\n tIdx += gridDim.x * blockDim.x;\n }\n}\n\n// Kernel to extract diagonal for csr_matrix format\ntemplate \n__global__\nvoid getDiagonalKernelNoDiaProp(const IndexType *dia_idx, const ValueType *values, const IndexType numRows, ValueType *diagonal)\n{\n int tIdx = threadIdx.x + blockDim.x * blockIdx.x;\n\n while (tIdx < numRows)\n {\n diagonal[tIdx] = values[dia_idx[tIdx]];\n tIdx += gridDim.x * blockDim.x;\n }\n}\n\n// filter edge weights like this:\n// set w_ij = 0 iff\n// w_ij < alpha * sqrt( max_k{w_ik} * max_l{w_jl} )\n// alpha is some constant, 0.25 or 0.5 should work fine\ntemplate\n__global__\nvoid filterWeights( const IndexType *row_offsets, const IndexType *row_indices, const IndexType *col_indices, const IndexType *diag, const ValueType *old_weights, ValueType *new_weights, IndexType num_nonzero_blocks, IndexType num_owned, ValueType alpha )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n int i, j, kmin, kmax;\n ValueType max_ik, max_jl;\n\n while ( tid < num_nonzero_blocks )\n {\n i = row_indices[tid];\n j = col_indices[tid];\n\n if ( i != j && j < num_owned )\n {\n //find max_k{w_ik}\n kmin = row_offsets[i];\n kmax = row_offsets[i + 1];\n max_ik = 0.0;\n\n for (int k = kmin; k < kmax; k++)\n {\n if ( col_indices[k] != i && old_weights[k] > max_ik )\n {\n max_ik = old_weights[k];\n }\n }\n\n //find max_l{w_jl}\n kmin = row_offsets[j];\n kmax = row_offsets[j + 1];\n max_jl = 0.0;\n\n for (int l = kmin; l < kmax; l++)\n {\n if ( col_indices[l] != j && old_weights[l] > max_jl )\n {\n max_jl = old_weights[l];\n }\n }\n\n //test squared inequality\n if ( old_weights[tid] * old_weights[tid] < alpha * alpha * max_ik * max_jl )\n {\n new_weights[tid] = 0.0;\n }\n else //rescale to relative importance. this should also increase the chance of a handshake\n {\n new_weights[tid] = old_weights[tid];\n }\n\n // new_weights[tid] = old_weights[tid] / sqrt(max_ik*max_jl);\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate\n__global__\nvoid gatherValuesInterleaved( const ValueType *inValues, ValueType *outValues, IndexType nnz, int sq_blocksize, int index_offset )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < nnz )\n {\n //at least the write is coalesced\n outValues[tid] = inValues[tid * sq_blocksize + index_offset];\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\ntemplate\n__global__\nvoid addToWeights( ValueTypeM *edge_weights, const ValueTypeV *x, const IndexType *row_indices, IndexType *col_indices, IndexType nnz, double scale )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < nnz )\n {\n int i = row_indices[tid];\n int j = col_indices[tid];\n edge_weights[tid] -= static_cast( scale * fabs( x[i] - x[j] ) );\n tid += gridDim.x * blockDim.x;\n }\n}\n\ntemplate \n__global__\nvoid rescaleVector( ValueType *x, IndexType numRows )\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < numRows )\n {\n x[tid] = 2 * x[tid] - 1;\n tid += gridDim.x * blockDim.x;\n }\n}\n\n// -----------------\n// Methods\n// ----------------\n\n// Constructor\ntemplate\nMultiPairwiseSelectorBase::MultiPairwiseSelectorBase(AMG_Config &cfg, const std::string &cfg_scope)\n{\n deterministic = cfg.AMG_Config::template getParameter(\"determinism_flag\", \"default\");\n max_iterations = cfg.AMG_Config::template getParameter(\"max_matching_iterations\", cfg_scope);\n numUnassigned_tol = cfg.AMG_Config::template getParameter(\"max_unassigned_percentage\", cfg_scope);\n two_phase = cfg.AMG_Config::template getParameter(\"handshaking_phases\", cfg_scope) == 2;\n m_aggregation_edge_weight_component = cfg.AMG_Config::template getParameter(\"aggregation_edge_weight_component\", cfg_scope);\n aggregation_passes = cfg.AMG_Config::template getParameter(\"aggregation_passes\", cfg_scope); //default to size 8 aggregates. maybe its more convenient to have that as a config parameter\n filter_weights = cfg.AMG_Config::template getParameter(\"filter_weights\", cfg_scope); //by default: no filtering\n filter_weights_alpha = cfg.AMG_Config::template getParameter( \"filter_weights_alpha\", cfg_scope ); //default to 0.25\n full_ghost_level = cfg.AMG_Config::template getParameter( \"full_ghost_level\", cfg_scope ); //defaults to 0\n notay_weights = cfg.AMG_Config::template getParameter( \"notay_weights\", cfg_scope ); //defaults to 0\n ghost_offdiag_limit = cfg.AMG_Config::template getParameter( \"ghost_offdiag_limit\", cfg_scope ); //defaults to 0\n merge_singletons = cfg.AMG_Config::template getParameter( \"merge_singletons\", cfg_scope ); //defaults to 1\n weight_formula = cfg.AMG_Config::template getParameter( \"weight_formula\", cfg_scope ); //wheight formula defaults to 0\n serial_matching = cfg.AMG_Config::template getParameter( \"serial_matching\", cfg_scope ) != 0; //will use a serial matching algorithm instead of handshake\n modified_handshake = cfg.AMG_Config::template getParameter(\"modified_handshake\", cfg_scope ) == 1;\n //passes = 1 -> max = 3\n //passes = 2 -> max = 5\n //passes = 3 -> max = 10\n //passes = 4 -> max = 18\n max_aggregate_size = 2;\n\n for (int i = 1; i < aggregation_passes; i ++)\n {\n max_aggregate_size *= 2;\n }\n\n max_aggregate_size += aggregation_passes - (aggregation_passes / 2);\n mCfg = cfg;\n mCfg_scope = cfg_scope;\n}\n\n// setAggregates for block_dia_csr_matrix_h format\ntemplate \nvoid MultiPairwiseSelector >::setAggregates_common_sqblocks(Matrix_h &A,\n typename Matrix_h::IVector &aggregates, typename Matrix_h::IVector &aggregates_global, int &num_aggregates, MVector &edge_weights, IVector &sizes)\n{\n FatalError(\"MultiPairwise selector: setAggregates not implemented on CPU, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n// device specialization\n\n//edge_weights is an in/out parameter:\n//if its size is zero, the edge_weights will be computed from A and stored into edge_weights\n//else the edge_weights will not be computed and assumed to be valid for the given A. the value array of A is not used in this case\ntemplate \nvoid MultiPairwiseSelector >::setAggregates_common_sqblocks(Matrix_d &A,\n typename Matrix_d::IVector &aggregates,\n typename Matrix_d::IVector &aggregates_global,\n int &num_aggregates,\n MVector &edge_weights,\n IVector &sizes)\n{\n IndexType num_block_rows = (int) A.get_num_rows();\n IndexType num_nonzero_blocks = (int) A.get_num_nz();\n // both ways are supported\n IndexType total_nz = (A.is_matrix_singleGPU()) ? num_nonzero_blocks : A.manager->num_nz_all();\n typename Matrix_d::IVector &row_indices = A.row_indices;\n row_indices.resize( total_nz);\n cusp::detail::offsets_to_indices(A.row_offsets, row_indices);\n IndexType total_rows = (A.is_matrix_singleGPU()) ? A.get_num_rows() : A.manager->num_rows_all();\n aggregates.resize(total_rows);\n thrust_wrapper::fill(aggregates.begin(), aggregates.end(), -1);\n cudaCheckError();\n\n if ( this->merge_singletons == 2 && sizes.size() == 0 )\n {\n sizes.resize( total_rows, 1 ); //init with all ones\n }\n\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_row_indices_ptr = row_indices.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_idx_ptr = A.diag.raw();\n const ValueType *A_nonzero_values_ptr = A.values.raw();\n typename Matrix_d::IVector strongest_neighbour(num_block_rows, -1);\n typename Matrix_d::IVector strongest_neighbour_1phase(num_block_rows, -1);\n Vector > aggregated(num_block_rows, 0);\n IndexType *strongest_neighbour_ptr = strongest_neighbour.raw();\n IndexType *strongest_neighbour_1phase_ptr = strongest_neighbour_1phase.raw();\n IndexType *aggregates_ptr = aggregates.raw();\n const int threads_per_block = 256;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (num_block_rows - 1) / threads_per_block + 1 );\n int numUnassigned = num_block_rows;\n int numUnassigned_previous = numUnassigned;\n bool computeWeights = ( edge_weights.size() == 0 );\n\n if (computeWeights)\n {\n if ( A.hasProps( DIAG ) )\n {\n edge_weights.resize( num_nonzero_blocks + num_block_rows, 0.0 );\n }\n else\n {\n edge_weights.resize( num_nonzero_blocks + 1, -1 ); //+1 is important to some algorithms\n }\n }\n\n ValueType *edge_weights_ptr = edge_weights.raw();\n ValueType *rand_edge_weights_ptr = NULL;\n hipStream_t str = amgx::thrust::global_thread_handle::get_stream();\n\n // Compute the edge weights\n if ( computeWeights )\n {\n const int num_blocks_V2 = std::min( AMGX_GRID_MAX_SIZE, (num_nonzero_blocks - 1) / threads_per_block + 1);\n //compute with std formula\n hipFuncSetCacheConfig(reinterpret_cast(computeEdgeWeightsBlockDiaCsr_V2), hipFuncCachePreferL1);\n computeEdgeWeightsBlockDiaCsr_V2 <<< num_blocks_V2, threads_per_block, 0, str>>>(A_row_offsets_ptr,\n A_row_indices_ptr,\n A_column_indices_ptr,\n A_dia_idx_ptr,\n A_nonzero_values_ptr,\n num_nonzero_blocks,\n edge_weights_ptr,\n rand_edge_weights_ptr,\n num_block_rows,\n A.get_block_dimy(),\n this->m_aggregation_edge_weight_component,\n this->weight_formula);\n cudaCheckError();\n }\n\n //filter weights if desired\n if ( this->filter_weights == 1 )\n {\n MVector tmp( edge_weights.size() );\n const int num_blocks_filter = std::min( AMGX_GRID_MAX_SIZE, (num_nonzero_blocks - 1) / threads_per_block + 1);\n hipStreamSynchronize(str);\n cudaCheckError();\n filterWeights <<< num_blocks_filter, threads_per_block, 0, str>>>( A_row_offsets_ptr,\n A_row_indices_ptr,\n A_column_indices_ptr,\n A_dia_idx_ptr,\n edge_weights_ptr,\n tmp.raw(),\n num_nonzero_blocks,\n num_block_rows,\n this->filter_weights_alpha);\n cudaCheckError();\n hipStreamSynchronize(str);\n cudaCheckError();\n tmp.swap( edge_weights );\n edge_weights_ptr = edge_weights.raw();\n }\n\n// compute matching\n if ( !this->serial_matching )\n {\n IVector degree;\n\n if ( this->modified_handshake )\n {\n degree.resize( num_block_rows );\n }\n\n#ifdef EXPERIMENTAL_ITERATIVE_MATCHING\n // TODO: allocate host pinned memory\n AsyncEvent *throttle_event = new AsyncEvent;\n throttle_event->create();\n typename Matrix_h::IVector h_unagg_vec(1);\n typename Matrix_d::IVector d_unagg_vec(1);\n int *unaggregated = h_unagg_vec.raw();\n int *d_unaggregated = d_unagg_vec.raw();\n#endif\n int icount, s = 1;\n {\n icount = 0;\n ValueType *weights_ptr = edge_weights_ptr;\n\n do\n {\n if ( !this->two_phase )\n {\n if ( this->modified_handshake )\n computeDegree <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr,\n A_column_indices_ptr,\n weights_ptr,\n aggregates_ptr,\n sizes.raw(),\n degree.raw(),\n num_block_rows,\n this->max_aggregate_size );\n cudaCheckError();\n\n // 1-phase handshaking\n if ( this->modified_handshake )\n findStrongestNeighbourBlockDiaCsr_V3\n <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr,\n A_column_indices_ptr,\n weights_ptr,\n num_block_rows,\n aggregates_ptr,\n strongest_neighbour_ptr,\n sizes.raw(),\n degree.raw(),\n A.get_block_dimy(),\n this->max_aggregate_size,\n this->merge_singletons);\n else\n findStrongestNeighbourBlockDiaCsr_V3\n <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr,\n A_column_indices_ptr,\n weights_ptr,\n num_block_rows,\n aggregates_ptr,\n strongest_neighbour_ptr,\n sizes.raw(),\n degree.raw(),\n A.get_block_dimy(),\n this->max_aggregate_size,\n this->merge_singletons);\n\n cudaCheckError();\n }\n else\n {\n // 2-phase handshaking\n findStrongestNeighbourBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, weights_ptr, num_block_rows, aggregates_ptr, strongest_neighbour_1phase_ptr, strongest_neighbour_ptr, A.get_block_dimy(), 1, this->merge_singletons);\n cudaCheckError();\n // 2nd phase: for each block_row, find the strongest neighbour among those who gave hand on 1st phase\n findStrongestNeighbourBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, weights_ptr, num_block_rows, aggregates_ptr, strongest_neighbour_1phase_ptr, strongest_neighbour_ptr, A.get_block_dimy(), 2, this->merge_singletons);\n cudaCheckError();\n }\n\n // Look for perfect matches. Also, for nodes without unaggregated neighbours, merge with aggregate containing strongest neighbour\n if ( this->merge_singletons == 2 )\n {\n matchEdges <<< num_blocks, threads_per_block, 0, str>>>(num_block_rows, aggregates_ptr, strongest_neighbour_ptr, sizes.raw());\n cudaCheckError();\n }\n else\n {\n matchEdges <<< num_blocks, threads_per_block, 0, str>>>(num_block_rows, aggregates_ptr, strongest_neighbour_ptr, (int *)NULL);\n }\n\n cudaCheckError();\n#ifdef EXPERIMENTAL_ITERATIVE_MATCHING\n s = (icount & 1);\n\n if ( s == 0 )\n {\n // count unaggregated vertices\n hipError_t cuda_rc = hipMemsetAsync(d_unaggregated, 0, sizeof(int), str);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemsetAsync failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n countAggregates <<< num_blocks, threads_per_block, 0, str>>>(num_block_rows, aggregates_ptr, d_unaggregated);\n cudaCheckError();\n cuda_rc = hipMemcpyAsync(unaggregated, d_unaggregated, sizeof(int), hipMemcpyDeviceToHost, str);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpyAsync failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n throttle_event->record(str);\n }\n else\n {\n throttle_event->sync();\n numUnassigned_previous = numUnassigned;\n numUnassigned = *unaggregated;\n }\n\n#else\n hipError_t cuda_rc = hipStreamSynchronize(str);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipStreamSynchronize failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n numUnassigned_previous = numUnassigned;\n numUnassigned = (int)amgx::thrust::count(aggregates.begin(), aggregates.begin() + num_block_rows, -1);\n cudaCheckError();\n#endif\n icount++;\n }\n while ( (s == 0) || !(numUnassigned == 0 || icount > this->max_iterations || 1.0 * numUnassigned / num_block_rows < this->numUnassigned_tol || numUnassigned == numUnassigned_previous));\n\n }\n // printf(\"%i,\\n\", icount);\n#ifdef EXPERIMENTAL_ITERATIVE_MATCHING\n delete throttle_event;\n#endif\n }\n else\n {\n computeMatchingSerialGreedy( A, aggregates, num_aggregates, edge_weights );\n }\n\n if ( this->merge_singletons == 1 )\n {\n // Merge remaining vertices with current aggregates\n if (this->deterministic != 1)\n {\n while (numUnassigned != 0)\n {\n mergeWithExistingAggregatesBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, A.get_block_dimy(), this->deterministic, (IndexType *) NULL);\n cudaCheckError();\n numUnassigned = (int)amgx::thrust::count(aggregates.begin(), aggregates.begin() + num_block_rows, -1);\n cudaCheckError();\n }\n }\n else\n {\n typename Matrix_d::IVector aggregates_candidate(num_block_rows, -1);\n\n while (numUnassigned != 0)\n {\n mergeWithExistingAggregatesBlockDiaCsr_V2 <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr, A_column_indices_ptr, edge_weights_ptr, num_block_rows, aggregates_ptr, A.get_block_dimy(), this->deterministic, aggregates_candidate.raw());\n cudaCheckError();\n joinExistingAggregates <<< num_blocks, threads_per_block, 0, str>>>(num_block_rows, aggregates_ptr, aggregates_candidate.raw());\n cudaCheckError();\n numUnassigned = (int)amgx::thrust::count(aggregates.begin(), aggregates.begin() + num_block_rows, -1);\n cudaCheckError();\n }\n\n aggregates_candidate.resize(0);\n }\n }\n else if (this->merge_singletons == 0 )\n {\n //make singletons\n aggregateSingletons <<< num_blocks, threads_per_block, 0, str>>>( aggregates_ptr, num_block_rows );\n cudaCheckError();\n }\n else if ( this->merge_singletons == 2 )\n {\n //merges all remaining singletons into adequate neighbors if possible\n mergeSingletonsSmart <<< num_blocks, threads_per_block, 0, str>>>(A_row_offsets_ptr,\n A_column_indices_ptr,\n edge_weights_ptr,\n aggregates_ptr,\n sizes.raw(),\n num_block_rows,\n this->max_aggregate_size);\n cudaCheckError();\n }\n\n //This will assign num_aggregates to the pseudo aggregate without counting it. Perfect!\n this->renumberAndCountAggregates(aggregates, aggregates_global, num_block_rows, num_aggregates);\n\n if ( this->merge_singletons == 2 )\n {\n //udpate the sizes vector, so it matches the renumbered aggregates size\n IVector sizesSource;\n sizesSource.swap( sizes );\n sizes.resize( num_aggregates, 1 );\n updateAggregateSizes <<< num_blocks, threads_per_block, 0, str>>>( sizesSource.raw(), sizes.raw(), aggregates_ptr, num_block_rows );\n cudaCheckError();\n }\n\n}\n\n\n//instead of a handshake, we use a serial greedy algorithm to compute a better matching\n//the algorithm:\n// 1. compute degree of every node and sort nodes by degree into double linked list\n// 2. while non-isolated nodes left:\n// take node with minimum degree > 0\n// find strongest edge to unaggregated node and assign to new aggregate\n// remove both nodes from linked list\n// decrease degree of each neighbor by one for each of the two nodes\n// update list\ntemplate \nvoid MultiPairwiseSelector >::computeMatchingSerialGreedy( const Matrix_d &A, IVector &aggregates, int &numAggregates, MVector &edge_weights)\n{\n IndexType numRows = A.row_offsets.size() - 1;\n IndexType nnz = A.col_indices.size();\n //allocate memory on host\n IndexType *ia = new IndexType[numRows + 1];\n IndexType *ja = new IndexType[nnz];\n ValueType *w = new ValueType[nnz];\n IndexType *agg = new IndexType[numRows];\n IndexType *deg = new IndexType[numRows];\n //copy\n hipError_t cuda_rc = hipMemcpy( ia, A.row_offsets.raw(), sizeof(IndexType) * (numRows + 1), hipMemcpyDeviceToHost );\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy row_offsets D2H failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = hipMemcpy( ja, A.col_indices.raw(), sizeof(IndexType)*nnz, hipMemcpyDeviceToHost );\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy col_indices D2H failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = hipMemcpy( w, edge_weights.raw(), sizeof(ValueType)*nnz, hipMemcpyDeviceToHost );\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy edge_weights D2H failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n //init agg and compute the degree of each aggregate\n int max_degree = 0;\n\n for (IndexType i = 0; i < numRows; i++)\n {\n agg[i] = -1;\n int degree = 0;\n\n for (IndexType ii = ia[i]; ii < ia[i + 1]; ii++) //only care for positive weights\n {\n if ( ja[ii] != i && w[ii] > 0.0 )\n {\n degree++;\n }\n }\n\n if ( degree > max_degree )\n {\n max_degree = degree;\n }\n\n deg[i] = degree;\n }\n\n if ( max_degree >= numRows )\n {\n FatalError( \"max degree is greater than numRows.\", AMGX_ERR_UNKNOWN );\n }\n\n //init double linked list\n IndexType *fwd = new IndexType[numRows + max_degree + 1];\n IndexType *bwd = new IndexType[numRows + max_degree + 1];\n\n for (IndexType i = 0; i < numRows + max_degree + 1; i++)\n {\n fwd[i] = i;\n bwd[i] = i;\n }\n\n IndexType nodesLeft = numRows;\n numAggregates = 0;\n\n //insert nodes into list\n for (IndexType i = numRows - 1; i >= 0; i--) //inserting in backward order the nodes will be sorted by index in case of same degree\n {\n //insert forward following root\n fwd[i] = fwd[numRows + deg[i]];\n fwd[numRows + deg[i]] = i;\n //insert backward\n bwd[i] = numRows + deg[i];\n bwd[fwd[i]] = i;\n\n //isolated nodes cannot be aggregated\n if ( deg[i] == 0 )\n {\n nodesLeft--;\n }\n }\n\n while ( nodesLeft > 0 )\n {\n IndexType node = numRows;\n int degree;\n\n for (degree = 1; degree <= max_degree; degree++)\n {\n //list not empty -> select node\n if ( fwd[numRows + degree] < numRows ) //selecting the first node will select the most recently inserted one or the one with lowest index. both is preferable\n {\n node = fwd[numRows + degree];\n }\n\n if ( node < numRows )\n {\n break;\n }\n }\n\n //no node with degree > 1 found even though nodesLeft > 0\n if ( node == numRows )\n {\n FatalError(\"nodeLeft counting or list invalid\", AMGX_ERR_UNKNOWN );\n }\n\n if ( agg[node] != -1 )\n {\n FatalError(\"node is already aggregated\", AMGX_ERR_UNKNOWN );\n }\n\n //find strongest edge\n ValueType max_weight = 0.0;\n IndexType max_node = numRows; //use this as gatekeeper, so if weight == 0 the node index will not be greater than this\n\n for (IndexType ii = ia[node]; ii < ia[node + 1]; ii++)\n {\n IndexType j = ja[ii];\n\n if ( agg[j] != -1 || j == node)\n {\n continue;\n }\n\n if ( w[ii] > 0.0 )\n {\n degree--;\n }\n\n //deterministic, doesn't selects 0 weight.\n if ( w[ii] > max_weight || (w[ii] == max_weight && j > max_node) ) //always taking the edge pointing to the max node can give good alignment if numbering is structured\n {\n max_node = j;\n max_weight = w[ii];\n }\n } //Note that there has to be at least one neighbor node because degree of node is at least 1.\n\n if ( max_node == numRows )\n {\n FatalError( \"node has no neighbor although degree of node is at least 1\", AMGX_ERR_UNKNOWN );\n }\n\n if ( degree != 0 )\n {\n FatalError( \"node degree corrupted\", AMGX_ERR_UNKNOWN );\n }\n\n //aggregate\n agg[node] = node;\n agg[max_node] = node;\n numAggregates++;\n nodesLeft -= 2;\n //remove from list\n fwd[bwd[node]] = fwd[node];\n bwd[fwd[node]] = bwd[node];\n fwd[bwd[max_node]] = fwd[max_node];\n bwd[fwd[max_node]] = bwd[max_node];\n\n //update neighbors and list\n //max_node first\n for (IndexType ii = ia[max_node]; ii < ia[max_node + 1]; ii++)\n {\n IndexType j = ja[ii];\n\n if ( agg[j] != -1 || w[ii] <= 0.0)\n {\n continue;\n }\n\n //remove j from list\n fwd[bwd[j]] = fwd[j];\n bwd[fwd[j]] = bwd[j];\n //update degree of j\n deg[j]--;\n //add j back to start of the list\n fwd[j] = fwd[numRows + deg[j]];\n bwd[j] = numRows + deg[j];\n bwd[fwd[j]] = j;\n fwd[bwd[j]] = j;\n\n if (deg[j] == 0)\n {\n nodesLeft--;\n }\n }\n\n //node second, this will prefer nodes neighbors over max_nodes neighbors when choosing the next node\n for (IndexType ii = ia[node]; ii < ia[node + 1]; ii++)\n {\n IndexType j = ja[ii];\n\n if ( agg[j] != -1 || w[ii] <= 0.0)\n {\n continue;\n }\n\n //remove j from list\n fwd[bwd[j]] = fwd[j];\n bwd[fwd[j]] = bwd[j];\n //update degree of j\n deg[j]--;\n //add j back to start of the list\n fwd[j] = fwd[numRows + deg[j]];\n bwd[j] = numRows + deg[j];\n bwd[fwd[j]] = j;\n fwd[bwd[j]] = j;\n\n if (deg[j] == 0)\n {\n nodesLeft--;\n }\n }\n }\n\n //copy result back to device\n cuda_rc = hipMemcpy( aggregates.raw(), agg, sizeof(IndexType)*numRows, hipMemcpyHostToDevice );\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy aggregates H2D failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n\n //assert matching\n for (IndexType node = 0; node < numRows; node++)\n {\n if ( agg[node] == -1 )\n {\n continue;\n }\n\n for ( IndexType partner = 0; partner < numRows; partner++)\n {\n if ( agg[partner] == agg[node] )\n {\n if ( partner == node )\n {\n continue;\n }\n\n bool neighbor = false;\n\n for (IndexType ii = ia[node]; ii < ia[node + 1]; ii++)\n if ( ja[ii] == partner )\n {\n neighbor = true;\n break;\n }\n\n if ( !neighbor )\n {\n for (IndexType ii = ia[partner]; ii < ia[partner + 1]; ii++)\n if ( ja[ii] == node )\n {\n neighbor = true;\n break;\n }\n }\n\n if ( !neighbor )\n {\n FatalError(\"Internal error in aggregation selector\", AMGX_ERR_INTERNAL);\n }\n }\n }\n }\n\n //you shall not leak memory\n delete[] ia;\n delete[] ja;\n delete[] w;\n delete[] agg;\n delete[] deg;\n delete[] fwd;\n delete[] bwd;\n}\n\n\n\n\n//this kernel merges aggregate2 into aggregate1\ntemplate\n__global__\nvoid mergeAggregates(IndexType *aggregate1, const IndexType *aggregate2, IndexType sizeAggregate1, IndexType sizeAggregate2, IndexType sizeAggregate3)\n{\n int tid = threadIdx.x + blockDim.x * blockIdx.x;\n\n while ( tid < sizeAggregate1 )\n {\n if ( aggregate1[tid] == sizeAggregate2 )\n {\n aggregate1[tid] = sizeAggregate3;\n }\n else\n {\n aggregate1[tid] = aggregate2[aggregate1[tid]];\n }\n\n tid += gridDim.x * blockDim.x;\n }\n}\n\n\n\ntemplate \nvoid MultiPairwiseSelector >::computeIncompleteGalerkin( const Matrix_h &A,\n Matrix_h &Ac,\n const typename Matrix_h::IVector &aggregates,\n const typename Matrix_h::IVector &R_row_offsets,\n const typename Matrix_h::IVector &R_column_indices,\n const int num_aggregates )\n{\n FatalError(\"computeIncomlpetegalerkin is not supported on host. Run with ghost_offdiag_limit=0 instead.\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\ntemplate \nvoid MultiPairwiseSelector >::computeIncompleteGalerkin( const Matrix_d &A,\n Matrix_d &Ac,\n const typename Matrix_d::IVector &aggregates,\n const typename Matrix_d::IVector &R_row_offsets,\n const typename Matrix_d::IVector &R_column_indices,\n const int num_aggregates )\n{\n FatalError(\"computeIncomlpetegalerkin is not implemented yet. run with ghost_offdiag_limit=0 instead.\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate\nvoid MultiPairwiseSelectorBase::assertRestriction( const IVector &R_row_offsets, const IVector &R_col_indices, const IVector &aggregates )\n{\n int *r_ia = new int[R_row_offsets.size()];\n int *r_ja = new int[R_col_indices.size()];\n int *agg = new int[aggregates.size()];\n int *used_col = new int[aggregates.size()];\n\n for ( int i = 0; i < aggregates.size(); i++ )\n {\n used_col[i] = 0;\n }\n\n hipError_t cuda_rc = hipMemcpy( r_ia, R_row_offsets.raw(), sizeof(int)*R_row_offsets.size(), hipMemcpyDeviceToHost );\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy R_row_offsets D2H failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = hipMemcpy( r_ja, R_col_indices.raw(), sizeof(int)*R_col_indices.size(), hipMemcpyDeviceToHost );\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy R_col_indices D2H failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = hipMemcpy( agg, aggregates.raw(), sizeof(int)*aggregates.size(), hipMemcpyDeviceToHost );\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipMemcpy aggregates D2H failed in multi_pairwise\", AMGX_ERR_CUDA_FAILURE);\n }\n\n for ( int i = 0; i < R_row_offsets.size() - 1; i++ )\n {\n for ( int ii = r_ia[i]; ii < r_ia[i + 1]; ii++ )\n {\n int j = r_ja[ii];\n used_col[j]++;\n\n if ( used_col[j] > 1 )\n {\n std::cout << \"column \" << j << \" is present at least \" << used_col[j] << \" times\" << std::endl;\n }\n\n if ( j < 0 || j >= aggregates.size() )\n {\n std::cout << \"Error: j out of bounds, j = \" << j << \" and numRows = \" << aggregates.size() << std::endl;\n }\n else if ( agg[j] != i )\n {\n std::cout << \"Error: agg[\" << j << \"] = \" << agg[j] << \" != \" << i << std::endl;\n }\n }\n }\n\n std::cout << \"assert restriction done\" << std::endl;\n}\n\ntemplate\nvoid MultiPairwiseSelectorBase::setAggregates(Matrix &A,\n IVector &aggregates, IVector &aggregates_global, int &num_aggregates)\n{\n if (A.get_block_dimx() == A.get_block_dimy())\n {\n //ghost level matrix. this is a probably a weight matrix\n Matrix ghostA;\n ghostA.values.resize(0);\n //prolongation and restriction operator. this is only needed in when LowDegCoarseAGenerator is used\n IVector R_row_offsets;\n IVector R_col_indices;\n //holds the size of each aggregate\n IVector sizes;\n sizes.resize(0);\n //aggregates for ghost level\n IVector aggregates_current;\n IVector aggregates_global_current;\n bool aggregates_initialized = true;\n\n if (aggregates.size() == 0)\n {\n aggregates_initialized = false;\n\n if (!A.is_matrix_singleGPU())\n {\n aggregates.resize(A.manager->halo_offset(A.manager->num_neighbors()));\n }\n else\n {\n aggregates.resize(A.get_num_rows());\n }\n }\n\n //for mergeAggregates kernel\n const int threads_per_block = 256;\n const int num_blocks = std::min( AMGX_GRID_MAX_SIZE, (A.get_num_rows() - 1) / threads_per_block + 1 );\n hipStream_t stream = amgx::thrust::global_thread_handle::get_stream();\n //initialize and prepare weight matrix\n Matrix w;\n w.set_initialized(0);\n w.addProps(CSR);\n w.delProps(COO);\n w.setColsReorderedByColor(false);\n w.resize( 0, 0, 0, 1, 1, true ); //empty scalar 0x0 matrix\n w.values.resize(0); //matrix resize sets the values array to nnz+1 for no apparent reason\n IndexType targetSize = 1;\n //initialize coarse A generator\n CoarseAGenerator *cag;\n const bool use_restriction = true;\n const bool shrink_ghost_level = false;\n cag = new LowDegCoarseAGenerator(mCfg, mCfg_scope);\n // This will make coarseAGenerator to allocate more memory inside of galerkin\n ghostA.manager = new DistributedManager();\n w.manager = new DistributedManager();\n Matrix *curA = &A;\n\n //foreach pass do:\n // 1. build aggregates and weights\n // 2. create weight matrix (in full_ghost_level mode this is the input matrix or the last ghostlevel matrix)\n // 3. if in full ghost level mode, build R\n // 4. compute next level\n for (int current_pass = 1; true; current_pass++)\n {\n const IndexType numRows = curA->get_num_rows();\n const IndexType nnz = curA->get_num_nz();\n targetSize *= 2;\n\n if ( full_ghost_level )\n {\n w.values.resize(0); //compute weights from curA\n }\n else\n {\n w.values.swap( ghostA.values ); //use the weights computed with the galerkin operator (this will do nothing in the first pass, both values have size 0)\n }\n\n // create aggregates from correct input matrix\n setAggregates_common_sqblocks( *curA, aggregates_current, aggregates_global_current, num_aggregates, w.values, sizes );\n\n\n if ( current_pass > 1 )\n {\n //merge original aggregate with the newly created ones\n mergeAggregates <<< num_blocks, threads_per_block, 0, stream >>>( aggregates.raw(), aggregates_current.raw(), A.get_num_rows(), numRows, num_aggregates );\n cudaCheckError();\n //mergeAggregates<<< num_blocks, threads_per_block, 0, stream >>>( aggregates_global.raw(), aggregates_global_current.raw(), A.get_num_rows() );\n //cudaCheckError();\n }\n\n //try to free memory\n if ( full_ghost_level )\n {\n //then we don't need to save the weights, only for original level to do post processing\n w.values.resize(0);\n }\n else\n {\n //save edge weights for original level later\n //in that case we can throw away the values of ghostA as we will use the values to compute the next ghost level\n ghostA.values.resize(0);\n }\n\n // this is the break condition for the loop\n if ( current_pass >= aggregation_passes || num_aggregates <= 1 || num_aggregates == numRows)\n {\n //this means, aggregates has not been initialized yet\n if ( !aggregates_initialized )\n {\n aggregates.swap( aggregates_current );\n }\n\n if ( !aggregates_initialized )\n {\n aggregates_global.swap( aggregates_global_current );\n }\n\n hipStreamSynchronize( stream );\n cudaCheckError();\n\n break;\n }\n\n //prepare A to be corrupted\n curA->set_initialized(0);\n //swap in ia, ja from curA\n w.row_offsets.swap( curA->row_offsets );\n w.col_indices.swap( curA->col_indices );\n\n if ( full_ghost_level )\n {\n if ( shrink_ghost_level && curA->get_block_dimx() > 1)\n {\n //set w to correct size\n w.values.resize( nnz );\n //define grid and offsets\n const int num_blocks_inter = std::min( (int)AMGX_GRID_MAX_SIZE, (int)(nnz - 1) / threads_per_block + 1 );\n const int sq_blocksize = A.get_block_dimx() * A.get_block_dimy();\n const int index_offset = A.get_block_dimy() * m_aggregation_edge_weight_component + m_aggregation_edge_weight_component;\n //do the interleaved copy\n gatherValuesInterleaved <<< num_blocks_inter, threads_per_block, 0, stream>>>( A.values.raw(), w.values.raw(), nnz, sq_blocksize, index_offset );\n cudaCheckError();\n hipStreamSynchronize( stream );\n cudaCheckError();\n }\n else\n {\n w.values.swap( curA->values );\n }\n }\n\n w.diag.swap( curA->diag );\n\n //resize to inform the matrix of its new size\n if ( full_ghost_level && !shrink_ghost_level )\n {\n w.set_block_dimx( A.get_block_dimx() );\n w.set_block_dimy( A.get_block_dimy() );\n }\n else\n {\n w.set_block_dimx( 1 );\n w.set_block_dimy( 1 );\n }\n\n w.set_num_rows( numRows );\n w.set_num_cols( numRows );\n w.set_num_nz( nnz );\n w.set_allow_recompute_diag( false );\n\n if ( curA->hasProps( DIAG ) )\n {\n w.addProps( DIAG );\n }\n\n //ready to use\n w.set_initialized(1);\n\n //compute restriction operator\n if ( use_restriction )\n {\n IVector R_row_indices(aggregates_current);\n R_row_offsets.resize(num_aggregates + 2);\n R_col_indices.resize(numRows);\n thrust_wrapper::sequence(R_col_indices.begin(), R_col_indices.end());\n cudaCheckError();\n amgx::thrust::sort_by_key(R_row_indices.begin(), R_row_indices.end(), R_col_indices.begin());\n cudaCheckError();\n cusp::detail::indices_to_offsets(R_row_indices, R_row_offsets);\n cudaCheckError();\n //delete last row, which holds the pseudo aggregate\n R_row_offsets.resize( num_aggregates + 1);\n R_col_indices.resize( R_row_offsets[num_aggregates] );\n }\n\n // 3. compute galerkin ghost level\n if ( ghost_offdiag_limit == 0 )\n {\n //compute full galerkin\n cag->computeAOperator(w,\n ghostA,\n aggregates_current,\n R_row_offsets,\n R_col_indices,\n num_aggregates);\n }\n else\n {\n //compute incomplete galerkin\n computeIncompleteGalerkin(w,\n ghostA,\n aggregates_current,\n R_row_offsets,\n R_col_indices,\n num_aggregates);\n }\n\n //from now on w will be destroyed again.\n w.set_initialized(0);\n\n //repair the original A matrix. its ia and ja are in w\n if ( current_pass == 1 )\n {\n //swap back\n w.row_offsets.swap( A.row_offsets );\n w.col_indices.swap( A.col_indices );\n\n //only in that case we have swapped the values\n if ( full_ghost_level && !shrink_ghost_level )\n {\n w.values.swap( A.values );\n }\n\n //save the edge weights of the original level\n A.diag.swap( w.diag );\n A.set_initialized(1); //A is repaired now\n //save the first aggregates into the original aggregate vector so we can merge them later\n aggregates.swap( aggregates_current );\n aggregates_global.swap( aggregates_global_current );\n aggregates_initialized = true;\n curA = &ghostA;\n }\n }\n\n delete cag;\n }\n else\n {\n FatalError(\"Unsupported block size for MultiPairwise\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n}\n\n// -------------------------\n// Explict instantiations\n// -------------------------\n#define AMGX_CASE_LINE(CASE) template class MultiPairwiseSelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class MultiPairwiseSelector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n}\n}\n", "source": "/kaggle/working/amgx/src/aggregation/selectors/multi_pairwise.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n//-----------------------------------------------------\n// Method to compute the Galerkin product: A_c=R*A*P\n// The number of non-zeros per row is computed using Thrust\n// The coarse matrix is created by custom kernel\n// --------------------------------------------------------\n\n#include \n#include \n#include //indices_to_offsets\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\nnamespace aggregation\n{\n\ntypedef amgx::thrust::tuple tuple_t;\n\nstruct isDiagonal\n{\n __host__ __device__\n bool operator()(tuple_t IJ)\n {\n return ( amgx::thrust::get<0>(IJ) == amgx::thrust::get<1>(IJ) );\n }\n\n};\n\n// --------------------\n// Kernels\n// --------------------\n\n// Kernel to store aggregate I of each fine point index i\ntemplate \n__global__\nvoid iToIKernel(const IndexType *row_offsets, const IndexType *aggregates, IndexType *I, const int num_rows)\n{\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_rows; tid += gridDim.x * blockDim.x)\n {\n int agg = aggregates[tid];\n\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n I[j] = agg;\n }\n }\n}\n\n// Kernel to store aggregate J of each fine point index j\ntemplate \n__global__\nvoid jToJKernel(const IndexType *column_indices, const IndexType *aggregates, IndexType *J, const int num_entries)\n{\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_entries; tid += gridDim.x * blockDim.x)\n {\n int j = column_indices[tid];\n J[tid] = aggregates[j];\n }\n}\n\n// Kernel to fill matrix Ac for 4x4 block matrices\ntemplate \n__global__\nvoid fillA_4by4Blocks_Thrust_Kernel(const IndexType *A_row_offsets, const IndexType *A_column_indices, const IndexType *A_dia_values, const ValueType *A_nonzero_values, const IndexType *R_row_offsets, const IndexType *R_column_indices, const IndexType *Ac_row_offsets, const IndexType *Ac_column_indices, IndexType *Ac_dia_values, ValueType *Ac_nonzero_values, const IndexType *aggregates, const int num_aggregates, const int num_threads)\n{\n const int tid = threadIdx.x;\n const int I = blockIdx.x * threads_per_block + tid;\n const int halfwarp_id = tid >> log_bsize_sq;\n const int mat_entry_index = tid & (bsize_sq - 1);\n volatile __shared__ IndexType current_cell_to_read[threads_per_block / bsize_sq];\n volatile __shared__ IndexType current_cell_to_write[threads_per_block / bsize_sq];\n volatile __shared__ IndexType max_nonzeros_i[threads_per_block];\n volatile __shared__ IndexType max_nonzeros_j[threads_per_block];\n const int size_hash_reg = (2. / 3.) * max_nonzeros_per_row; // 2/3 is arbitrary\n const int size_hash_shared = max_nonzeros_per_row - size_hash_reg;\n IndexType shared_offset = threadIdx.x * size_hash_shared;\n volatile __shared__ IndexType hash_shared[size_hash_shared * threads_per_block];\n IndexType hash_reg[size_hash_reg];\n int Ac_offset_min;\n int Ac_offset_max;\n#pragma unroll\n\n for (int i = 0; i < size_hash_reg; i++)\n {\n hash_reg[i] = -1;\n }\n\n for (int i = 0; i < size_hash_shared; i++)\n {\n hash_shared[shared_offset + i] = -1;\n }\n\n // Find maximum number of nonzeros for threads in same halfwarp\n int i_offset_min, i_offset_max;\n\n if (I < num_aggregates)\n {\n i_offset_min = R_row_offsets[I];\n i_offset_max = R_row_offsets[I + 1];\n max_nonzeros_i[tid] = i_offset_max - i_offset_min;\n Ac_offset_min = Ac_row_offsets[I];\n Ac_offset_max = Ac_row_offsets[I + 1];\n\n for (int k = Ac_offset_min, icount = 0; k < Ac_offset_max; k++, icount++)\n {\n if (icount < size_hash_reg)\n {\n hash_reg[icount] = Ac_column_indices[k];\n }\n else\n {\n hash_shared[shared_offset + icount - size_hash_reg] = Ac_column_indices[k];\n }\n }\n }\n else\n {\n i_offset_min = 0;\n i_offset_max = 0;\n max_nonzeros_i[tid] = 0;\n }\n\n int max_loop_i = 0;\n\n for (int m = 0; m < bsize_sq; m++)\n {\n int max_nonzeros = max_nonzeros_i[halfwarp_id * bsize_sq + m];\n\n if (max_nonzeros > max_loop_i)\n {\n max_loop_i = max_nonzeros;\n }\n }\n\n // For all threads that could do useful work\n if (I < num_threads)\n {\n for (int iloop = 0, i_offset = i_offset_min; iloop < max_loop_i; iloop++, i_offset++)\n {\n // Contribution from A_dia_values of fine point i\n int i = (i_offset < i_offset_max) ? R_column_indices[i_offset] : -1 ;\n\n // Have threads collaborate to load in coalesced fashion\n for (int m = 0; m < bsize_sq; m++)\n {\n // Write which cell to load in shared memory buffer\n if (mat_entry_index == m)\n {\n current_cell_to_read[halfwarp_id] = i;\n current_cell_to_write[halfwarp_id] = I;\n }\n\n // All threads read from shared which cell to read and write\n int cell_to_read = current_cell_to_read[halfwarp_id];\n int cell_to_write = current_cell_to_write[halfwarp_id];\n\n // Here all threads in half-warp will take same path\n if (cell_to_read != -1)\n {\n Ac_nonzero_values[ bsize_sq * Ac_dia_values[cell_to_write] + mat_entry_index] += A_nonzero_values[bsize_sq * A_dia_values[cell_to_read] + mat_entry_index];\n }\n }\n\n // Contribution from A_nonzero_values of fine point i\n // Find maximum number of nonzeros for threads in same halfwarp_id\n int j_offset_min, j_offset_max;\n\n if (i != -1)\n {\n j_offset_min = A_row_offsets[i];\n j_offset_max = A_row_offsets[i + 1];\n max_nonzeros_j[tid] = j_offset_max - j_offset_min;\n }\n else\n {\n j_offset_min = 0;\n j_offset_max = 0;\n max_nonzeros_j[tid] = 0;\n }\n\n int max_loop_j = 0;\n\n for (int m = 0; m < bsize_sq; m++)\n {\n int max_nonzeros = max_nonzeros_j[halfwarp_id * bsize_sq + m];\n\n if (max_nonzeros > max_loop_j)\n {\n max_loop_j = max_nonzeros;\n }\n }\n\n // Add contribution from nonzero_values of A\n int j_offset_to_read;\n\n for (int jloop = 0, j_offset = j_offset_min; jloop < max_loop_j; jloop++, j_offset++)\n {\n int J, j, k;\n\n if (j_offset < j_offset_max)\n {\n j_offset_to_read = j_offset;\n j = A_column_indices[j_offset];\n J = aggregates[j];\n\n // Find index k where to store the data and create A_column_indices\n if (I != J)\n {\n // This weird construct is to allow loop unrolling and avoid register spilling (see original version below)\n int found = 0;\n k = Ac_offset_min;\n#pragma unroll\n\n for (int icount = 0; icount < max_nonzeros_per_row; icount++)\n {\n if (k < Ac_offset_max)\n {\n if (found == 0)\n {\n int Jtemp = (icount < size_hash_reg) ? hash_reg[icount] : hash_shared[shared_offset + icount - size_hash_reg];\n\n if (J == Jtemp)\n {\n found = 1;\n }\n\n if (found == 0) { k++; }\n }\n }\n }\n } // if I != J\n } // if j_offset < j_offset_max\n else\n {\n j_offset_to_read = -1;\n }\n\n // Have threads collaborate to load in coalesced fashion\n for (int m = 0; m < bsize_sq; m++)\n {\n // Write which cell to load in shared memory buffer\n if (mat_entry_index == m)\n {\n current_cell_to_read[halfwarp_id] = j_offset_to_read;\n\n if (I != J)\n {\n current_cell_to_write[halfwarp_id] = -k;\n }\n else\n {\n current_cell_to_write[halfwarp_id] = I + 1;\n }\n }\n\n // All threads read from shared which cell to read and write\n int cell_to_read = current_cell_to_read[halfwarp_id];\n int cell_to_write = current_cell_to_write[halfwarp_id];\n\n if (cell_to_read != -1)\n {\n if (cell_to_write <= 0)\n {\n Ac_nonzero_values[(-cell_to_write)*bsize_sq + mat_entry_index] += A_nonzero_values[cell_to_read * bsize_sq + mat_entry_index];\n }\n else\n {\n Ac_nonzero_values[ Ac_dia_values[ cell_to_write - 1 ] * bsize_sq + mat_entry_index] += A_nonzero_values[cell_to_read * bsize_sq + mat_entry_index];\n }\n }\n }\n } // j_offset_loop\n } // i_offset_loop\n } // if tid < num_threads\n}\n\n// Constructor\ntemplate\nHybridCoarseAGeneratorBase::HybridCoarseAGeneratorBase()\n{\n}\n\n//-----------------------------------------------------\n// Method to compute the Galerkin product: A_c=R*A*P\n// The number of non-zeros per row is computed using Thrust\n// The coarse matrix is created by custom kernel\n//-----------------------------------------------------\n\ntemplate \nvoid HybridCoarseAGenerator >::computeAOperator_4x4(const Matrix > &A, Matrix > &Ac, const typename Matrix_d::IVector &aggregates, const typename Matrix_d::IVector &R_row_offsets, const typename Matrix_d::IVector &R_column_indices, const int num_aggregates)\n{\n// supports both DIA properties\n if (!A.hasProps(DIAG))\n {\n FatalError(\"Hybryd coarser does not support inside diagonal yet\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (A.get_num_nz() == 0)\n {\n FatalError(\"Hybryd coarser does not work correctly with diagonal matrices\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typename Matrix_d::IVector I(A.get_num_nz(), -1);\n typename Matrix_d::IVector J(A.get_num_nz(), -1);\n typedef device_vector_alloc IntVector;\n typedef typename IntVector::iterator IntIterator;\n typedef amgx::thrust::tuple< IntIterator, IntIterator> IntIteratorTuple;\n typedef amgx::thrust::zip_iterator ZipIterator;\n ZipIterator new_end;\n const int block_size_I = 128;\n const int block_size_J = 256;\n const int num_blocks_I = std::min( AMGX_GRID_MAX_SIZE, (int) ((A.get_num_rows() - 1) / block_size_I + 1));\n const int num_blocks_J = std::min( AMGX_GRID_MAX_SIZE, (int) ((A.get_num_nz() - 1) / block_size_J + 1));\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_values_ptr = A.diag.raw();\n const ValueType *A_nonzero_values_ptr = A.values.raw();\n const IndexType *aggregates_ptr = aggregates.raw();\n IndexType *I_ptr = I.raw();\n IndexType *J_ptr = J.raw();\n // Kernel to fill array I with aggregates number for fine points i\n iToIKernel <<< num_blocks_I, block_size_I>>>(A_row_offsets_ptr, aggregates_ptr, I_ptr, (int) A.get_num_rows());\n cudaCheckError();\n // Kernel to fill array J with aggregates number for fine points j\n jToJKernel <<< num_blocks_J, block_size_J>>>(A_column_indices_ptr, aggregates_ptr, J_ptr, (int) A.get_num_nz());\n cudaCheckError();\n // Sort (I,J) by rows and columns (I,J)\n IVector permutation(A.get_num_nz());\n thrust_wrapper::sequence(permutation.begin(), permutation.end());\n cudaCheckError();\n // compute permutation and sort by (I,J)\n {\n IVector temp(J);\n amgx::thrust::stable_sort_by_key(temp.begin(), temp.end(), permutation.begin());\n cudaCheckError();\n temp = I;\n //I = temp;\n thrust_wrapper::gather(permutation.begin(), permutation.end(), temp.begin(), I.begin());\n cudaCheckError();\n amgx::thrust::stable_sort_by_key(I.begin(), I.end(), permutation.begin());\n cudaCheckError();\n temp = J;\n //J = temp;\n thrust_wrapper::gather(permutation.begin(), permutation.end(), temp.begin(), J.begin());\n cudaCheckError();\n }\n // Remove duplicate tuples\n new_end = amgx::thrust::unique(amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.end(), J.end())), amgx::thrust::equal_to < amgx::thrust::tuple >());\n cudaCheckError();\n IntIteratorTuple endTuple = new_end.get_iterator_tuple();\n I.erase(amgx::thrust::get<0>(endTuple), I.end());\n J.erase(amgx::thrust::get<1>(endTuple), J.end());\n // Remove diagonal terms\n new_end = amgx::thrust::remove_if(amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.end(), J.end())), isDiagonal() );\n cudaCheckError();\n endTuple = new_end.get_iterator_tuple();\n I.erase(amgx::thrust::get<0>(endTuple), I.end());\n J.erase(amgx::thrust::get<1>(endTuple), J.end());\n int nonzero_blocks = J.size();\n // Resize Ac\n Ac.addProps(CSR);\n\n if (A.hasProps(DIAG)) { Ac.addProps(DIAG); }\n\n Ac.resize(num_aggregates, num_aggregates, nonzero_blocks, A.get_block_dimy(), A.get_block_dimx());\n //Ac.resize(num_aggregates,num_aggregates,nonzero_blocks,A.get_block_dimy(),A.get_block_dimx(), 1);\n // Ac.column_indices\n Ac.col_indices = J;\n J.clear();\n J.shrink_to_fit();\n // Convert array new_row_indices to offsets\n cusp::detail::indices_to_offsets(I, Ac.row_offsets);\n cudaCheckError();\n I.resize(Ac.row_offsets.size());\n // Compute the maximum number of nonzeros\n amgx::thrust::adjacent_difference(Ac.row_offsets.begin(), Ac.row_offsets.end(), I.begin());\n cudaCheckError();\n const IndexType max_nonzero_per_row = *amgx::thrust::max_element(I.begin() + 1, I.end());\n cudaCheckError();\n //std::cout << \"max_nonzero_per_row\" << max_nonzero_per_row << std::endl;\n I.clear();\n I.shrink_to_fit();\n const IndexType *R_row_offsets_ptr = R_row_offsets.raw();\n const IndexType *R_column_indices_ptr = R_column_indices.raw();\n // Get max_nonzero_per_row\n //const IndexType max_nonzero_per_row = 32;\n const int threads_per_block = 128;\n // Store the column_indices in a register\n // Option to print the number of nonzeros distribution\n //printNonzeroStats(Ac_row_offsets_temp,num_aggregates);\n // Resize Ac and doing exclusive scan on Ac_row_offsets_temp\n IndexType *Ac_row_offsets_ptr = Ac.row_offsets.raw();\n IndexType *Ac_column_indices_ptr = Ac.col_indices.raw();\n IndexType *Ac_dia_values_ptr = Ac.diag.raw();\n ValueType *Ac_nonzero_values_ptr = Ac.values.raw();\n // Now create Ac.dia_values and Ac.nonzero_values\n //thrust_wrapper::fill(Ac.diag.begin(),Ac.diag.end(),0.);\n thrust_wrapper::fill(Ac.values.begin(), Ac.values.end(), 0.);\n cudaCheckError();\n // Coalesced version of kernel to fill A\n const int num_threads = ( ( num_aggregates + 15) / 16 ) * 16;\n const int num_blocks2 = ( ( num_threads + threads_per_block - 1) / threads_per_block );\n\n if (max_nonzero_per_row < 16)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 20)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 24)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 28)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 32)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 36)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 40)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 44)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 48)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 52)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 56)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 60)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 64)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 128)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 256)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 512)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else\n {\n FatalError(\"Maximum number of nonzeros is too large\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid HybridCoarseAGenerator >::computeAOperator_4x4(const Matrix > &A, Matrix > &Ac, const IVector &aggregates, const IVector &R_row_offsets, const IVector &R_column_indices, const int num_aggregates)\n{\n FatalError(\"Host is unsupported for HybridCoarseAGenerator\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n// ------------------------------------------------\ntemplate \nvoid HybridCoarseAGeneratorBase::computeAOperator(const Matrix &A, Matrix &Ac, const IVector &aggregates, const IVector &R_row_offsets, const IVector &R_column_indices, const int num_aggregates)\n{\n Ac.set_initialized(0);\n\n if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n computeAOperator_4x4( A, Ac, aggregates, R_row_offsets, R_column_indices, num_aggregates );\n }\n else\n {\n FatalError(\"Unsupported block size for HybridCoarseAGenerator\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n Ac.set_initialized(1);\n}\n\n// ---------------------------\n// Explict instantiations\n// ---------------------------\n#define AMGX_CASE_LINE(CASE) template class HybridCoarseAGeneratorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class HybridCoarseAGenerator::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n}\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n//-----------------------------------------------------\n// Method to compute the Galerkin product: A_c=R*A*P\n// The number of non-zeros per row is computed using Thrust\n// The coarse matrix is created by custom kernel\n// --------------------------------------------------------\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include //indices_to_offsets\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\nnamespace aggregation\n{\n\ntypedef amgx::thrust::tuple tuple_t;\n\nstruct isDiagonal\n{\n __host__ __device__\n bool operator()(tuple_t IJ)\n {\n return ( amgx::thrust::get<0>(IJ) == amgx::thrust::get<1>(IJ) );\n }\n\n};\n\n// --------------------\n// Kernels\n// --------------------\n\n// Kernel to store aggregate I of each fine point index i\ntemplate \n__global__\nvoid iToIKernel(const IndexType *row_offsets, const IndexType *aggregates, IndexType *I, const int num_rows)\n{\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_rows; tid += gridDim.x * blockDim.x)\n {\n int agg = aggregates[tid];\n\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n I[j] = agg;\n }\n }\n}\n\n// Kernel to store aggregate J of each fine point index j\ntemplate \n__global__\nvoid jToJKernel(const IndexType *column_indices, const IndexType *aggregates, IndexType *J, const int num_entries)\n{\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_entries; tid += gridDim.x * blockDim.x)\n {\n int j = column_indices[tid];\n J[tid] = aggregates[j];\n }\n}\n\n// Kernel to fill matrix Ac for 4x4 block matrices\ntemplate \n__global__\nvoid fillA_4by4Blocks_Thrust_Kernel(const IndexType *A_row_offsets, const IndexType *A_column_indices, const IndexType *A_dia_values, const ValueType *A_nonzero_values, const IndexType *R_row_offsets, const IndexType *R_column_indices, const IndexType *Ac_row_offsets, const IndexType *Ac_column_indices, IndexType *Ac_dia_values, ValueType *Ac_nonzero_values, const IndexType *aggregates, const int num_aggregates, const int num_threads)\n{\n const int tid = threadIdx.x;\n const int I = blockIdx.x * threads_per_block + tid;\n const int halfwarp_id = tid >> log_bsize_sq;\n const int mat_entry_index = tid & (bsize_sq - 1);\n volatile __shared__ IndexType current_cell_to_read[threads_per_block / bsize_sq];\n volatile __shared__ IndexType current_cell_to_write[threads_per_block / bsize_sq];\n volatile __shared__ IndexType max_nonzeros_i[threads_per_block];\n volatile __shared__ IndexType max_nonzeros_j[threads_per_block];\n const int size_hash_reg = (2. / 3.) * max_nonzeros_per_row; // 2/3 is arbitrary\n const int size_hash_shared = max_nonzeros_per_row - size_hash_reg;\n IndexType shared_offset = threadIdx.x * size_hash_shared;\n volatile __shared__ IndexType hash_shared[size_hash_shared * threads_per_block];\n IndexType hash_reg[size_hash_reg];\n int Ac_offset_min;\n int Ac_offset_max;\n#pragma unroll\n\n for (int i = 0; i < size_hash_reg; i++)\n {\n hash_reg[i] = -1;\n }\n\n for (int i = 0; i < size_hash_shared; i++)\n {\n hash_shared[shared_offset + i] = -1;\n }\n\n // Find maximum number of nonzeros for threads in same halfwarp\n int i_offset_min, i_offset_max;\n\n if (I < num_aggregates)\n {\n i_offset_min = R_row_offsets[I];\n i_offset_max = R_row_offsets[I + 1];\n max_nonzeros_i[tid] = i_offset_max - i_offset_min;\n Ac_offset_min = Ac_row_offsets[I];\n Ac_offset_max = Ac_row_offsets[I + 1];\n\n for (int k = Ac_offset_min, icount = 0; k < Ac_offset_max; k++, icount++)\n {\n if (icount < size_hash_reg)\n {\n hash_reg[icount] = Ac_column_indices[k];\n }\n else\n {\n hash_shared[shared_offset + icount - size_hash_reg] = Ac_column_indices[k];\n }\n }\n }\n else\n {\n i_offset_min = 0;\n i_offset_max = 0;\n max_nonzeros_i[tid] = 0;\n }\n\n int max_loop_i = 0;\n\n for (int m = 0; m < bsize_sq; m++)\n {\n int max_nonzeros = max_nonzeros_i[halfwarp_id * bsize_sq + m];\n\n if (max_nonzeros > max_loop_i)\n {\n max_loop_i = max_nonzeros;\n }\n }\n\n // For all threads that could do useful work\n if (I < num_threads)\n {\n for (int iloop = 0, i_offset = i_offset_min; iloop < max_loop_i; iloop++, i_offset++)\n {\n // Contribution from A_dia_values of fine point i\n int i = (i_offset < i_offset_max) ? R_column_indices[i_offset] : -1 ;\n\n // Have threads collaborate to load in coalesced fashion\n for (int m = 0; m < bsize_sq; m++)\n {\n // Write which cell to load in shared memory buffer\n if (mat_entry_index == m)\n {\n current_cell_to_read[halfwarp_id] = i;\n current_cell_to_write[halfwarp_id] = I;\n }\n\n // All threads read from shared which cell to read and write\n int cell_to_read = current_cell_to_read[halfwarp_id];\n int cell_to_write = current_cell_to_write[halfwarp_id];\n\n // Here all threads in half-warp will take same path\n if (cell_to_read != -1)\n {\n Ac_nonzero_values[ bsize_sq * Ac_dia_values[cell_to_write] + mat_entry_index] += A_nonzero_values[bsize_sq * A_dia_values[cell_to_read] + mat_entry_index];\n }\n }\n\n // Contribution from A_nonzero_values of fine point i\n // Find maximum number of nonzeros for threads in same halfwarp_id\n int j_offset_min, j_offset_max;\n\n if (i != -1)\n {\n j_offset_min = A_row_offsets[i];\n j_offset_max = A_row_offsets[i + 1];\n max_nonzeros_j[tid] = j_offset_max - j_offset_min;\n }\n else\n {\n j_offset_min = 0;\n j_offset_max = 0;\n max_nonzeros_j[tid] = 0;\n }\n\n int max_loop_j = 0;\n\n for (int m = 0; m < bsize_sq; m++)\n {\n int max_nonzeros = max_nonzeros_j[halfwarp_id * bsize_sq + m];\n\n if (max_nonzeros > max_loop_j)\n {\n max_loop_j = max_nonzeros;\n }\n }\n\n // Add contribution from nonzero_values of A\n int j_offset_to_read;\n\n for (int jloop = 0, j_offset = j_offset_min; jloop < max_loop_j; jloop++, j_offset++)\n {\n int J, j, k;\n\n if (j_offset < j_offset_max)\n {\n j_offset_to_read = j_offset;\n j = A_column_indices[j_offset];\n J = aggregates[j];\n\n // Find index k where to store the data and create A_column_indices\n if (I != J)\n {\n // This weird construct is to allow loop unrolling and avoid register spilling (see original version below)\n int found = 0;\n k = Ac_offset_min;\n#pragma unroll\n\n for (int icount = 0; icount < max_nonzeros_per_row; icount++)\n {\n if (k < Ac_offset_max)\n {\n if (found == 0)\n {\n int Jtemp = (icount < size_hash_reg) ? hash_reg[icount] : hash_shared[shared_offset + icount - size_hash_reg];\n\n if (J == Jtemp)\n {\n found = 1;\n }\n\n if (found == 0) { k++; }\n }\n }\n }\n } // if I != J\n } // if j_offset < j_offset_max\n else\n {\n j_offset_to_read = -1;\n }\n\n // Have threads collaborate to load in coalesced fashion\n for (int m = 0; m < bsize_sq; m++)\n {\n // Write which cell to load in shared memory buffer\n if (mat_entry_index == m)\n {\n current_cell_to_read[halfwarp_id] = j_offset_to_read;\n\n if (I != J)\n {\n current_cell_to_write[halfwarp_id] = -k;\n }\n else\n {\n current_cell_to_write[halfwarp_id] = I + 1;\n }\n }\n\n // All threads read from shared which cell to read and write\n int cell_to_read = current_cell_to_read[halfwarp_id];\n int cell_to_write = current_cell_to_write[halfwarp_id];\n\n if (cell_to_read != -1)\n {\n if (cell_to_write <= 0)\n {\n Ac_nonzero_values[(-cell_to_write)*bsize_sq + mat_entry_index] += A_nonzero_values[cell_to_read * bsize_sq + mat_entry_index];\n }\n else\n {\n Ac_nonzero_values[ Ac_dia_values[ cell_to_write - 1 ] * bsize_sq + mat_entry_index] += A_nonzero_values[cell_to_read * bsize_sq + mat_entry_index];\n }\n }\n }\n } // j_offset_loop\n } // i_offset_loop\n } // if tid < num_threads\n}\n\n// Constructor\ntemplate\nHybridCoarseAGeneratorBase::HybridCoarseAGeneratorBase()\n{\n}\n\n//-----------------------------------------------------\n// Method to compute the Galerkin product: A_c=R*A*P\n// The number of non-zeros per row is computed using Thrust\n// The coarse matrix is created by custom kernel\n//-----------------------------------------------------\n\ntemplate \nvoid HybridCoarseAGenerator >::computeAOperator_4x4(const Matrix > &A, Matrix > &Ac, const typename Matrix_d::IVector &aggregates, const typename Matrix_d::IVector &R_row_offsets, const typename Matrix_d::IVector &R_column_indices, const int num_aggregates)\n{\n// supports both DIA properties\n if (!A.hasProps(DIAG))\n {\n FatalError(\"Hybryd coarser does not support inside diagonal yet\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n if (A.get_num_nz() == 0)\n {\n FatalError(\"Hybryd coarser does not work correctly with diagonal matrices\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typename Matrix_d::IVector I(A.get_num_nz(), -1);\n typename Matrix_d::IVector J(A.get_num_nz(), -1);\n typedef device_vector_alloc IntVector;\n typedef typename IntVector::iterator IntIterator;\n typedef amgx::thrust::tuple< IntIterator, IntIterator> IntIteratorTuple;\n typedef amgx::thrust::zip_iterator ZipIterator;\n ZipIterator new_end;\n const int block_size_I = 128;\n const int block_size_J = 256;\n const int num_blocks_I = std::min( AMGX_GRID_MAX_SIZE, (int) ((A.get_num_rows() - 1) / block_size_I + 1));\n const int num_blocks_J = std::min( AMGX_GRID_MAX_SIZE, (int) ((A.get_num_nz() - 1) / block_size_J + 1));\n const IndexType *A_row_offsets_ptr = A.row_offsets.raw();\n const IndexType *A_column_indices_ptr = A.col_indices.raw();\n const IndexType *A_dia_values_ptr = A.diag.raw();\n const ValueType *A_nonzero_values_ptr = A.values.raw();\n const IndexType *aggregates_ptr = aggregates.raw();\n IndexType *I_ptr = I.raw();\n IndexType *J_ptr = J.raw();\n // Kernel to fill array I with aggregates number for fine points i\n iToIKernel <<< num_blocks_I, block_size_I>>>(A_row_offsets_ptr, aggregates_ptr, I_ptr, (int) A.get_num_rows());\n cudaCheckError();\n // Kernel to fill array J with aggregates number for fine points j\n jToJKernel <<< num_blocks_J, block_size_J>>>(A_column_indices_ptr, aggregates_ptr, J_ptr, (int) A.get_num_nz());\n cudaCheckError();\n // Sort (I,J) by rows and columns (I,J)\n IVector permutation(A.get_num_nz());\n thrust_wrapper::sequence(permutation.begin(), permutation.end());\n cudaCheckError();\n // compute permutation and sort by (I,J)\n {\n IVector temp(J);\n amgx::thrust::stable_sort_by_key(temp.begin(), temp.end(), permutation.begin());\n cudaCheckError();\n temp = I;\n //I = temp;\n thrust_wrapper::gather(permutation.begin(), permutation.end(), temp.begin(), I.begin());\n cudaCheckError();\n amgx::thrust::stable_sort_by_key(I.begin(), I.end(), permutation.begin());\n cudaCheckError();\n temp = J;\n //J = temp;\n thrust_wrapper::gather(permutation.begin(), permutation.end(), temp.begin(), J.begin());\n cudaCheckError();\n }\n // Remove duplicate tuples\n new_end = amgx::thrust::unique(amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.end(), J.end())), amgx::thrust::equal_to < amgx::thrust::tuple >());\n cudaCheckError();\n IntIteratorTuple endTuple = new_end.get_iterator_tuple();\n I.erase(amgx::thrust::get<0>(endTuple), I.end());\n J.erase(amgx::thrust::get<1>(endTuple), J.end());\n // Remove diagonal terms\n new_end = amgx::thrust::remove_if(amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.end(), J.end())), isDiagonal() );\n cudaCheckError();\n endTuple = new_end.get_iterator_tuple();\n I.erase(amgx::thrust::get<0>(endTuple), I.end());\n J.erase(amgx::thrust::get<1>(endTuple), J.end());\n int nonzero_blocks = J.size();\n // Resize Ac\n Ac.addProps(CSR);\n\n if (A.hasProps(DIAG)) { Ac.addProps(DIAG); }\n\n Ac.resize(num_aggregates, num_aggregates, nonzero_blocks, A.get_block_dimy(), A.get_block_dimx());\n //Ac.resize(num_aggregates,num_aggregates,nonzero_blocks,A.get_block_dimy(),A.get_block_dimx(), 1);\n // Ac.column_indices\n Ac.col_indices = J;\n J.clear();\n J.shrink_to_fit();\n // Convert array new_row_indices to offsets\n cusp::detail::indices_to_offsets(I, Ac.row_offsets);\n cudaCheckError();\n I.resize(Ac.row_offsets.size());\n // Compute the maximum number of nonzeros\n amgx::thrust::adjacent_difference(Ac.row_offsets.begin(), Ac.row_offsets.end(), I.begin());\n cudaCheckError();\n const IndexType max_nonzero_per_row = *amgx::thrust::max_element(I.begin() + 1, I.end());\n cudaCheckError();\n //std::cout << \"max_nonzero_per_row\" << max_nonzero_per_row << std::endl;\n I.clear();\n I.shrink_to_fit();\n const IndexType *R_row_offsets_ptr = R_row_offsets.raw();\n const IndexType *R_column_indices_ptr = R_column_indices.raw();\n // Get max_nonzero_per_row\n //const IndexType max_nonzero_per_row = 32;\n const int threads_per_block = 128;\n // Store the column_indices in a register\n // Option to print the number of nonzeros distribution\n //printNonzeroStats(Ac_row_offsets_temp,num_aggregates);\n // Resize Ac and doing exclusive scan on Ac_row_offsets_temp\n IndexType *Ac_row_offsets_ptr = Ac.row_offsets.raw();\n IndexType *Ac_column_indices_ptr = Ac.col_indices.raw();\n IndexType *Ac_dia_values_ptr = Ac.diag.raw();\n ValueType *Ac_nonzero_values_ptr = Ac.values.raw();\n // Now create Ac.dia_values and Ac.nonzero_values\n //thrust_wrapper::fill(Ac.diag.begin(),Ac.diag.end(),0.);\n thrust_wrapper::fill(Ac.values.begin(), Ac.values.end(), 0.);\n cudaCheckError();\n // Coalesced version of kernel to fill A\n const int num_threads = ( ( num_aggregates + 15) / 16 ) * 16;\n const int num_blocks2 = ( ( num_threads + threads_per_block - 1) / threads_per_block );\n\n if (max_nonzero_per_row < 16)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 20)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 24)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 28)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 32)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 36)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 40)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 44)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 48)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 52)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 56)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 60)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 64)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 128)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 256)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else if (max_nonzero_per_row < 512)\n {\n fillA_4by4Blocks_Thrust_Kernel <<< num_blocks2, threads_per_block>>>(A_row_offsets_ptr, A_column_indices_ptr, A_dia_values_ptr, A_nonzero_values_ptr, R_row_offsets_ptr, R_column_indices_ptr, Ac_row_offsets_ptr, Ac_column_indices_ptr, Ac_dia_values_ptr, Ac_nonzero_values_ptr, aggregates_ptr, num_aggregates, num_threads);\n }\n else\n {\n FatalError(\"Maximum number of nonzeros is too large\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n cudaCheckError();\n}\n\ntemplate \nvoid HybridCoarseAGenerator >::computeAOperator_4x4(const Matrix > &A, Matrix > &Ac, const IVector &aggregates, const IVector &R_row_offsets, const IVector &R_column_indices, const int num_aggregates)\n{\n FatalError(\"Host is unsupported for HybridCoarseAGenerator\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n}\n\n// ------------------------------------------------\ntemplate \nvoid HybridCoarseAGeneratorBase::computeAOperator(const Matrix &A, Matrix &Ac, const IVector &aggregates, const IVector &R_row_offsets, const IVector &R_column_indices, const int num_aggregates)\n{\n Ac.set_initialized(0);\n\n if (A.get_block_dimx() == 4 && A.get_block_dimy() == 4)\n {\n computeAOperator_4x4( A, Ac, aggregates, R_row_offsets, R_column_indices, num_aggregates );\n }\n else\n {\n FatalError(\"Unsupported block size for HybridCoarseAGenerator\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n Ac.set_initialized(1);\n}\n\n// ---------------------------\n// Explict instantiations\n// ---------------------------\n#define AMGX_CASE_LINE(CASE) template class HybridCoarseAGeneratorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class HybridCoarseAGenerator::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n}\n", "source": "/kaggle/working/amgx/src/aggregation/coarseAgenerators/hybrid_coarse_A_generator.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n// --------------------------------------------------------\n// --------------------------------------------------------\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\nnamespace aggregation\n{\n\ntypedef amgx::thrust::tuple tuple_t;\n\n// --------------------\n// Kernels\n// --------------------\n\n// Kernel to store aggregate I of each fine point index i\ntemplate \n__global__\nvoid iToIKernel(const IndexType *row_offsets, const IndexType *aggregates, IndexType *I, const int num_rows)\n{\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_rows; tid += gridDim.x * blockDim.x)\n {\n int agg = aggregates[tid];\n\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n I[j] = agg;\n }\n }\n}\n\n// Kernel to store aggregate J of each fine point index j\ntemplate \n__global__\nvoid jToJKernel(const IndexType *column_indices, const IndexType *aggregates, IndexType *J, const int num_entries)\n{\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_entries; tid += gridDim.x * blockDim.x)\n {\n int j = column_indices[tid];\n J[tid] = aggregates[j];\n }\n}\n\n// Constructor\ntemplate\nThrustCoarseAGeneratorBase::ThrustCoarseAGeneratorBase()\n{\n}\n\n//-----------------------------------------------------\n// Method to compute the Galerkin product: A_c=R*A*P\n//-----------------------------------------------------\n\n// Method to compute A on DEVICE using csr format\ntemplate \nvoid ThrustCoarseAGenerator >::computeAOperator_1x1(const Matrix_d &A, Matrix_d &Ac, const IVector &aggregates, const IVector &R_row_offsets, const IVector &R_column_indices, const int num_aggregates)\n{\n if (A.hasProps(DIAG))\n {\n FatalError(\"ThrustCoarseAGenerator: unsupported diagonal\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n cudaCheckError();\n IVector I(A.get_num_nz(), -1);\n IVector J(A.get_num_nz(), -1);\n VVector V(A.get_num_nz(), -1);\n const int block_size_I = 128;\n const int block_size_J = 256;\n const int num_blocks_I = std::min( AMGX_GRID_MAX_SIZE, (int) ((A.get_num_rows() - 1) / block_size_I + 1) );\n const int num_blocks_J = std::min( AMGX_GRID_MAX_SIZE, (int) ((A.get_num_nz() - 1) / block_size_J + 1) );\n const IndexType *row_offsets_ptr = A.row_offsets.raw();\n const IndexType *column_indices_ptr = A.col_indices.raw();\n const IndexType *aggregates_ptr = aggregates.raw();\n IndexType *I_ptr = I.raw();\n IndexType *J_ptr = J.raw();\n // Kernel to fill array I with aggregates number for fine points i\n iToIKernel <<< num_blocks_I, block_size_I>>>(row_offsets_ptr, aggregates_ptr, I_ptr, (int)A.get_num_rows());\n cudaCheckError();\n // Kernel to fill array J with aggregates number for fine points j\n jToJKernel <<< num_blocks_J, block_size_J>>>(column_indices_ptr, aggregates_ptr, J_ptr, (int)A.get_num_nz());\n cudaCheckError();\n // Copy A.values to V array\n amgx::thrust::copy(A.values.begin(), A.values.begin() + A.get_num_nz()*A.get_block_size(), V.begin());\n cudaCheckError();\n // Sort (I,J,V) by rows and columns (I,J)\n cusp::detail::sort_by_row_and_column(I, J, V);\n cudaCheckError();\n // compute unique number of nonzeros in the output\n IndexType NNZ = amgx::thrust::inner_product(amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.end (), J.end())) - 1,\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())) + 1,\n IndexType(0),\n amgx::thrust::plus(),\n amgx::thrust::not_equal_to< amgx::thrust::tuple >()) + 1;\n cudaCheckError();\n // allocate space for coarse matrix Ac\n Ac.addProps(CSR);\n\n if (A.hasProps(DIAG)) { Ac.addProps(DIAG); }\n\n if (A.is_matrix_singleGPU())\n {\n Ac.resize(num_aggregates, num_aggregates, NNZ, 1);\n }\n else\n {\n Ac.resize_spare(num_aggregates, num_aggregates, NNZ, A.get_block_dimy(), A.get_block_dimx(), 1.0);\n\n if (A.hasProps(DIAG)) { Ac.computeDiagonal(); }\n }\n\n // Reduce by key to fill in Ac.column_indices and Ac.values\n IVector new_row_indices(NNZ, 0);\n amgx::thrust::reduce_by_key(amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.end(), J.end())),\n V.begin(),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(new_row_indices.begin(), Ac.col_indices.begin())),\n Ac.values.begin(),\n amgx::thrust::equal_to< amgx::thrust::tuple >(),\n amgx::thrust::plus());\n cudaCheckError();\n // Convert array new_row_indices to offsets\n cusp::detail::indices_to_offsets(new_row_indices, Ac.row_offsets);\n cudaCheckError();\n I.clear();\n I.shrink_to_fit();\n J.clear();\n J.shrink_to_fit();\n V.clear();\n V.shrink_to_fit();\n}\n\n// Method to compute A on HOST using csr format\ntemplate \nvoid ThrustCoarseAGenerator >::computeAOperator_1x1(const Matrix_h &A, Matrix_h &Ac, const IVector &aggregates, const IVector &R_row_offsets, const IVector &R_column_indices, const int num_aggregates)\n{\n if (A.hasProps(DIAG))\n {\n FatalError(\"ThrustCoarseAGenerator: unsupported diagonal\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n IVector I(A.get_num_nz(), -1);\n IVector J(A.get_num_nz(), -1);\n VVector V(A.get_num_nz(), -1);\n const IndexType *row_offsets_ptr = A.row_offsets.raw();\n const IndexType *column_indices_ptr = A.col_indices.raw();\n const IndexType *aggregates_ptr = aggregates.raw();\n IndexType *I_ptr = I.raw();\n IndexType *J_ptr = J.raw();\n\n // Kernel to fill array I with aggregates number for fine points i\n for ( int tid = 0; tid < (int)A.get_num_rows(); tid++ )\n {\n int agg = aggregates_ptr[tid];\n\n for (int j = row_offsets_ptr[tid]; j < row_offsets_ptr[tid + 1]; j++)\n {\n I_ptr[j] = agg;\n }\n }\n\n // Kernel to fill array J with aggregates number for fine points j\n for ( int tid = 0; tid < (int)A.get_num_nz(); tid++ )\n {\n int j = column_indices_ptr[tid];\n J_ptr[tid] = aggregates_ptr[j];\n }\n\n // Copy A.values to V array\n amgx::thrust::copy(A.values.begin(), A.values.begin() + A.get_num_nz()*A.get_block_size(), V.begin());\n cudaCheckError();\n // Sort (I,J,V) by rows and columns (I,J)\n cusp::detail::sort_by_row_and_column(I, J, V);\n cudaCheckError();\n // compute unique number of nonzeros in the output\n IndexType NNZ = amgx::thrust::inner_product(amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.end (), J.end())) - 1,\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())) + 1,\n IndexType(0),\n amgx::thrust::plus(),\n amgx::thrust::not_equal_to< amgx::thrust::tuple >()) + 1;\n cudaCheckError();\n // allocate space for coarse matrix Ac\n Ac.addProps(CSR);\n\n if (A.hasProps(DIAG)) { Ac.addProps(DIAG); }\n\n if (A.is_matrix_singleGPU())\n {\n Ac.resize(num_aggregates, num_aggregates, NNZ, 1);\n }\n else\n {\n Ac.resize_spare(num_aggregates, num_aggregates, NNZ, A.get_block_dimy(), A.get_block_dimx(), 1.0);\n\n if (A.hasProps(DIAG)) { Ac.computeDiagonal(); }\n }\n\n // Reduce by key to fill in Ac.column_indices and Ac.values\n typename Matrix_h::IVector new_row_indices(NNZ, 0);\n amgx::thrust::reduce_by_key(amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.end(), J.end())),\n V.begin(),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(new_row_indices.begin(), Ac.col_indices.begin())),\n Ac.values.begin(),\n amgx::thrust::equal_to< amgx::thrust::tuple >(),\n amgx::thrust::plus());\n cudaCheckError();\n // Convert array new_row_indices to offsets\n cusp::detail::indices_to_offsets(new_row_indices, Ac.row_offsets);\n cudaCheckError();\n I.clear();\n I.shrink_to_fit();\n J.clear();\n J.shrink_to_fit();\n V.clear();\n V.shrink_to_fit();\n}\n\n// ------------------------------------------------\ntemplate \nvoid ThrustCoarseAGeneratorBase::computeAOperator(const Matrix &A, Matrix &Ac, const IVector &aggregates, const IVector &R_row_offsets, const IVector &R_column_indices, const int num_aggregates)\n{\n Ac.set_initialized(0);\n\n if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n computeAOperator_1x1( A, Ac, aggregates, R_row_offsets, R_column_indices, num_aggregates );\n }\n else\n {\n FatalError(\"Unsupported block size for ThrustCoarseAGenerator\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if (Ac.is_matrix_singleGPU()) { Ac.computeDiagonal(); }\n\n Ac.set_initialized(1);\n}\n\n// ---------------------------\n// Explict instantiations\n// ---------------------------\n#define AMGX_CASE_LINE(CASE) template class ThrustCoarseAGeneratorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class ThrustCoarseAGenerator::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n}\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n// --------------------------------------------------------\n// --------------------------------------------------------\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\nnamespace aggregation\n{\n\ntypedef amgx::thrust::tuple tuple_t;\n\n// --------------------\n// Kernels\n// --------------------\n\n// Kernel to store aggregate I of each fine point index i\ntemplate \n__global__\nvoid iToIKernel(const IndexType *row_offsets, const IndexType *aggregates, IndexType *I, const int num_rows)\n{\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_rows; tid += gridDim.x * blockDim.x)\n {\n int agg = aggregates[tid];\n\n for (int j = row_offsets[tid]; j < row_offsets[tid + 1]; j++)\n {\n I[j] = agg;\n }\n }\n}\n\n// Kernel to store aggregate J of each fine point index j\ntemplate \n__global__\nvoid jToJKernel(const IndexType *column_indices, const IndexType *aggregates, IndexType *J, const int num_entries)\n{\n for (int tid = blockDim.x * blockIdx.x + threadIdx.x; tid < num_entries; tid += gridDim.x * blockDim.x)\n {\n int j = column_indices[tid];\n J[tid] = aggregates[j];\n }\n}\n\n// Constructor\ntemplate\nThrustCoarseAGeneratorBase::ThrustCoarseAGeneratorBase()\n{\n}\n\n//-----------------------------------------------------\n// Method to compute the Galerkin product: A_c=R*A*P\n//-----------------------------------------------------\n\n// Method to compute A on DEVICE using csr format\ntemplate \nvoid ThrustCoarseAGenerator >::computeAOperator_1x1(const Matrix_d &A, Matrix_d &Ac, const IVector &aggregates, const IVector &R_row_offsets, const IVector &R_column_indices, const int num_aggregates)\n{\n if (A.hasProps(DIAG))\n {\n FatalError(\"ThrustCoarseAGenerator: unsupported diagonal\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n cudaCheckError();\n IVector I(A.get_num_nz(), -1);\n IVector J(A.get_num_nz(), -1);\n VVector V(A.get_num_nz(), -1);\n const int block_size_I = 128;\n const int block_size_J = 256;\n const int num_blocks_I = std::min( AMGX_GRID_MAX_SIZE, (int) ((A.get_num_rows() - 1) / block_size_I + 1) );\n const int num_blocks_J = std::min( AMGX_GRID_MAX_SIZE, (int) ((A.get_num_nz() - 1) / block_size_J + 1) );\n const IndexType *row_offsets_ptr = A.row_offsets.raw();\n const IndexType *column_indices_ptr = A.col_indices.raw();\n const IndexType *aggregates_ptr = aggregates.raw();\n IndexType *I_ptr = I.raw();\n IndexType *J_ptr = J.raw();\n // Kernel to fill array I with aggregates number for fine points i\n iToIKernel <<< num_blocks_I, block_size_I>>>(row_offsets_ptr, aggregates_ptr, I_ptr, (int)A.get_num_rows());\n cudaCheckError();\n // Kernel to fill array J with aggregates number for fine points j\n jToJKernel <<< num_blocks_J, block_size_J>>>(column_indices_ptr, aggregates_ptr, J_ptr, (int)A.get_num_nz());\n cudaCheckError();\n // Copy A.values to V array\n amgx::thrust::copy(A.values.begin(), A.values.begin() + A.get_num_nz()*A.get_block_size(), V.begin());\n cudaCheckError();\n // Sort (I,J,V) by rows and columns (I,J)\n cusp::detail::sort_by_row_and_column(I, J, V);\n cudaCheckError();\n // compute unique number of nonzeros in the output\n IndexType NNZ = amgx::thrust::inner_product(amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.end (), J.end())) - 1,\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())) + 1,\n IndexType(0),\n amgx::thrust::plus(),\n amgx::thrust::not_equal_to< amgx::thrust::tuple >()) + 1;\n cudaCheckError();\n // allocate space for coarse matrix Ac\n Ac.addProps(CSR);\n\n if (A.hasProps(DIAG)) { Ac.addProps(DIAG); }\n\n if (A.is_matrix_singleGPU())\n {\n Ac.resize(num_aggregates, num_aggregates, NNZ, 1);\n }\n else\n {\n Ac.resize_spare(num_aggregates, num_aggregates, NNZ, A.get_block_dimy(), A.get_block_dimx(), 1.0);\n\n if (A.hasProps(DIAG)) { Ac.computeDiagonal(); }\n }\n\n // Reduce by key to fill in Ac.column_indices and Ac.values\n IVector new_row_indices(NNZ, 0);\n amgx::thrust::reduce_by_key(amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.end(), J.end())),\n V.begin(),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(new_row_indices.begin(), Ac.col_indices.begin())),\n Ac.values.begin(),\n amgx::thrust::equal_to< amgx::thrust::tuple >(),\n amgx::thrust::plus());\n cudaCheckError();\n // Convert array new_row_indices to offsets\n cusp::detail::indices_to_offsets(new_row_indices, Ac.row_offsets);\n cudaCheckError();\n I.clear();\n I.shrink_to_fit();\n J.clear();\n J.shrink_to_fit();\n V.clear();\n V.shrink_to_fit();\n}\n\n// Method to compute A on HOST using csr format\ntemplate \nvoid ThrustCoarseAGenerator >::computeAOperator_1x1(const Matrix_h &A, Matrix_h &Ac, const IVector &aggregates, const IVector &R_row_offsets, const IVector &R_column_indices, const int num_aggregates)\n{\n if (A.hasProps(DIAG))\n {\n FatalError(\"ThrustCoarseAGenerator: unsupported diagonal\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n IVector I(A.get_num_nz(), -1);\n IVector J(A.get_num_nz(), -1);\n VVector V(A.get_num_nz(), -1);\n const IndexType *row_offsets_ptr = A.row_offsets.raw();\n const IndexType *column_indices_ptr = A.col_indices.raw();\n const IndexType *aggregates_ptr = aggregates.raw();\n IndexType *I_ptr = I.raw();\n IndexType *J_ptr = J.raw();\n\n // Kernel to fill array I with aggregates number for fine points i\n for ( int tid = 0; tid < (int)A.get_num_rows(); tid++ )\n {\n int agg = aggregates_ptr[tid];\n\n for (int j = row_offsets_ptr[tid]; j < row_offsets_ptr[tid + 1]; j++)\n {\n I_ptr[j] = agg;\n }\n }\n\n // Kernel to fill array J with aggregates number for fine points j\n for ( int tid = 0; tid < (int)A.get_num_nz(); tid++ )\n {\n int j = column_indices_ptr[tid];\n J_ptr[tid] = aggregates_ptr[j];\n }\n\n // Copy A.values to V array\n amgx::thrust::copy(A.values.begin(), A.values.begin() + A.get_num_nz()*A.get_block_size(), V.begin());\n cudaCheckError();\n // Sort (I,J,V) by rows and columns (I,J)\n cusp::detail::sort_by_row_and_column(I, J, V);\n cudaCheckError();\n // compute unique number of nonzeros in the output\n IndexType NNZ = amgx::thrust::inner_product(amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.end (), J.end())) - 1,\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())) + 1,\n IndexType(0),\n amgx::thrust::plus(),\n amgx::thrust::not_equal_to< amgx::thrust::tuple >()) + 1;\n cudaCheckError();\n // allocate space for coarse matrix Ac\n Ac.addProps(CSR);\n\n if (A.hasProps(DIAG)) { Ac.addProps(DIAG); }\n\n if (A.is_matrix_singleGPU())\n {\n Ac.resize(num_aggregates, num_aggregates, NNZ, 1);\n }\n else\n {\n Ac.resize_spare(num_aggregates, num_aggregates, NNZ, A.get_block_dimy(), A.get_block_dimx(), 1.0);\n\n if (A.hasProps(DIAG)) { Ac.computeDiagonal(); }\n }\n\n // Reduce by key to fill in Ac.column_indices and Ac.values\n typename Matrix_h::IVector new_row_indices(NNZ, 0);\n amgx::thrust::reduce_by_key(amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.begin(), J.begin())),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(I.end(), J.end())),\n V.begin(),\n amgx::thrust::make_zip_iterator(amgx::thrust::make_tuple(new_row_indices.begin(), Ac.col_indices.begin())),\n Ac.values.begin(),\n amgx::thrust::equal_to< amgx::thrust::tuple >(),\n amgx::thrust::plus());\n cudaCheckError();\n // Convert array new_row_indices to offsets\n cusp::detail::indices_to_offsets(new_row_indices, Ac.row_offsets);\n cudaCheckError();\n I.clear();\n I.shrink_to_fit();\n J.clear();\n J.shrink_to_fit();\n V.clear();\n V.shrink_to_fit();\n}\n\n// ------------------------------------------------\ntemplate \nvoid ThrustCoarseAGeneratorBase::computeAOperator(const Matrix &A, Matrix &Ac, const IVector &aggregates, const IVector &R_row_offsets, const IVector &R_column_indices, const int num_aggregates)\n{\n Ac.set_initialized(0);\n\n if (A.get_block_dimx() == 1 && A.get_block_dimy() == 1)\n {\n computeAOperator_1x1( A, Ac, aggregates, R_row_offsets, R_column_indices, num_aggregates );\n }\n else\n {\n FatalError(\"Unsupported block size for ThrustCoarseAGenerator\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n if (Ac.is_matrix_singleGPU()) { Ac.computeDiagonal(); }\n\n Ac.set_initialized(1);\n}\n\n// ---------------------------\n// Explict instantiations\n// ---------------------------\n#define AMGX_CASE_LINE(CASE) template class ThrustCoarseAGeneratorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n#define AMGX_CASE_LINE(CASE) template class ThrustCoarseAGenerator::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n}\n}\n", "source": "/kaggle/working/amgx/src/aggregation/coarseAgenerators/thrust_coarse_A_generator.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\nnamespace aggregation\n{\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#include \n#include \n\nstatic __device__ __forceinline__ int get_work( int *queue, int warp_id, int count = 1 )\n{\n int offset = -1;\n\n if ( utils::lane_id() == 0 )\n {\n offset = atomicAdd( queue, count );\n }\n\n return utils::shfl( offset, 0 );\n}\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool HAS_DIAG, bool COUNT_ONLY >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_sparsity_kernel( const int R_num_rows, // same as num_aggregates.\n const int *R_rows,\n const int *R_cols,\n const int *A_rows,\n const int *A_cols,\n const int *aggregates,\n int *Ac_rows,\n int *Ac_cols,\n int *Ac_pos,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int r_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[r_row_id * gmem_size], gmem_size );\n // Loop over rows of R.\n for ( ; r_row_id < R_num_rows ; r_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Make sure we have to proceed.\n if ( COUNT_ONLY )\n {\n volatile int *status = reinterpret_cast( wk_status );\n\n if ( set.has_failed() || *status != 0 )\n {\n return;\n }\n }\n\n // Clear the set.\n set.clear();\n // Load the range of the row.\n int r_col_it = R_rows[r_row_id + 0];\n int r_col_end = R_rows[r_row_id + 1];\n\n // Iterate over the columns of R.\n for ( r_col_it += lane_id ; utils::any(r_col_it < r_col_end) ; r_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = r_col_it < r_col_end;\n // Columns of R map to rows of A. Each thread of the warp loads its R-col/A-row ID.\n int a_row_id = -1;\n\n if ( is_active )\n {\n a_row_id = R_cols[r_col_it];\n }\n\n const int num_rows = __popc( utils::ballot(is_active) );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n const int uniform_a_row_id = utils::shfl( a_row_id, local_k );\n // Load the range of the row of B.\n int a_col_it = 0, a_col_end = 0;\n\n if ( is_active_k )\n {\n a_col_it = A_rows[uniform_a_row_id + 0];\n a_col_end = A_rows[uniform_a_row_id + 1];\n }\n\n // Iterate over the range of columns of B.\n for ( a_col_it += lane_id_mod_num_threads ; utils::any(a_col_it < a_col_end) ; a_col_it += NUM_THREADS_PER_ROW )\n {\n int a_col_id = -1, a_agg_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n a_agg_id = aggregates[a_col_id];\n }\n\n if ( HAS_DIAG && a_agg_id == r_row_id )\n {\n a_agg_id = -1;\n }\n\n set.insert( a_agg_id, COUNT_ONLY ? wk_status : NULL );\n }\n }\n }\n\n // Store the results.\n if ( COUNT_ONLY )\n {\n int count = set.compute_size();\n\n if ( lane_id == 0 )\n {\n Ac_rows[r_row_id] = count;\n }\n }\n else\n {\n int ac_col_it = Ac_rows[r_row_id];\n set.store_with_positions( &Ac_cols[ac_col_it], &Ac_pos[ac_col_it] );\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int NUM_THREADS_PER_ROW, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool HAS_DIAG >\n__global__\n__launch_bounds__( CTA_SIZE, 8 )\nvoid fill_A_kernel_1x1( const int R_num_rows,\n const int *R_rows,\n const int *R_cols,\n const int *A_rows,\n const int *A_cols,\n const int *A_diag,\n const Value_type *A_vals,\n const int *aggregates,\n const int *Ac_rows,\n const int *Ac_cols,\n const int *Ac_pos,\n const int *Ac_diag,\n Value_type *Ac_vals,\n int gmem_size,\n int *g_keys,\n Value_type *g_vals,\n int *wk_work_queue )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The hash values stored in shared memory.\n __shared__ Value_type s_vals[NUM_WARPS * SMEM_SIZE]; \n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int r_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_map map( &s_keys[warp_id * SMEM_SIZE ],\n &g_keys[r_row_id * gmem_size ],\n &s_vals[warp_id * SMEM_SIZE],\n &g_vals[r_row_id * gmem_size ], gmem_size );\n // Loop over rows of A.\n for ( ; r_row_id < R_num_rows ; r_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // The indices of the output row.\n int ac_col_it = Ac_rows[r_row_id + 0];\n int ac_col_end = Ac_rows[r_row_id + 1];\n // Clear the set first. TODO: Make sure it's needed. I don't think it is!!!!\n map.clear();\n // Populate the map.\n map.load( ac_col_end - ac_col_it, &Ac_cols[ac_col_it], &Ac_pos[ac_col_it] );\n int r_col_it = R_rows[r_row_id + 0];\n int r_col_end = R_rows[r_row_id + 1];\n // The diagonal.\n Value_type r_diag(types::util::get_zero());\n\n // _iterate over the columns of A to build C_hat.\n for ( r_col_it += lane_id ; utils::any(r_col_it < r_col_end) ; r_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = r_col_it < r_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_row_id = -1;\n\n if ( is_active )\n {\n a_row_id = R_cols[r_col_it];\n }\n\n // Update the diagonal (if needed).\n if ( HAS_DIAG && is_active )\n {\n r_diag = r_diag + A_vals[A_diag[a_row_id]];\n }\n\n const int num_rows = __popc( utils::ballot(is_active) );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n const int uniform_a_row_id = utils::shfl( a_row_id, local_k );\n // The range of the row of B.\n int a_col_it = 0, a_col_end = 0;\n\n if ( local_k < num_rows )\n {\n a_col_it = utils::Ld::load( &A_rows[uniform_a_row_id + 0] );\n a_col_end = utils::Ld::load( &A_rows[uniform_a_row_id + 1] );\n }\n\n // Iterate over the range of columns of B.\n for ( a_col_it += lane_id_mod_num_threads ; utils::any(a_col_it < a_col_end) ; a_col_it += NUM_THREADS_PER_ROW )\n {\n // Load columns and values.\n int a_col_id = -1;\n Value_type a_value(types::util::get_zero());\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n a_value = A_vals[a_col_it];\n }\n\n // Find the aggregate.\n int a_agg_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_agg_id = aggregates[a_col_id];\n }\n\n // Update the diag/hash map.\n if ( HAS_DIAG && a_agg_id == r_row_id )\n {\n r_diag = r_diag + a_value;\n a_agg_id = -1;\n }\n\n map.insert( a_agg_id, a_value, NULL ); // It won't insert. Only update.\n }\n }\n }\n\n // Update the diagonal.\n if ( HAS_DIAG )\n {\n r_diag = utils::warp_reduce<1, utils::Add>( r_diag );\n\n if ( lane_id == 0 )\n {\n Ac_vals[Ac_diag[r_row_id]] = r_diag;\n }\n }\n\n // Store the results.\n int count = ac_col_end - ac_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n map.store( count, &Ac_vals[ac_col_it] );\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool HAS_DIAG >\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700\n__global__ __launch_bounds__( CTA_SIZE, 8 )\n#else \n__global__ __launch_bounds__( CTA_SIZE, 8 )\n#endif\nvoid fill_A_kernel_4x4( const int R_num_rows, // same as num_aggregates.\n const int *R_rows,\n const int *R_cols,\n const int *A_rows,\n const int *A_cols,\n const int *A_diag,\n const Value_type *A_vals,\n const int *aggregates,\n const int *Ac_rows,\n const int *Ac_cols,\n const int *Ac_pos,\n const int *Ac_diag,\n Value_type *Ac_vals,\n const int gmem_size,\n int *g_keys,\n int *g_idx,\n int *wk_work_queue )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_16 = lane_id / 16;\n const int lane_id_mod_16 = lane_id % 16;\n const int warp_offset = 16 * lane_id_div_16;\n // First threads load the row IDs of A needed by the CTA...\n int r_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // My index.\n Hash_index index( &g_idx[r_row_id * gmem_size] );\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[r_row_id * gmem_size], gmem_size );\n // Loop over rows of R.\n for ( ; r_row_id < R_num_rows ; r_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // The indices of the row.\n int ac_col_it = Ac_rows[r_row_id + 0];\n int ac_col_end = Ac_rows[r_row_id + 1];\n // Clear the set first.\n set.clear(true);\n // Populate the index.\n set.load_index( ac_col_end - ac_col_it, &Ac_cols[ac_col_it], &Ac_pos[ac_col_it], index, false );\n // Load the range of the row.\n int r_col_it = R_rows[r_row_id + 0];\n int r_col_end = R_rows[r_row_id + 1];\n // Diagonal value (each half warp stores a diagonal element).\n Value_type ac_diag(types::util::get_zero());\n\n // Iterate over the columns of R.\n for ( r_col_it += lane_id_div_16 ; utils::any(r_col_it < r_col_end) ; r_col_it += 2 )\n {\n // Is it an active thread.\n const bool is_active = r_col_it < r_col_end;\n // Columns of R map to rows of A. Each thread of the warp loads its R-col/A-row ID.\n int a_row_id = -1;\n\n if ( is_active )\n {\n a_row_id = R_cols[r_col_it];\n }\n\n // Update the diagonal if needed.\n if ( HAS_DIAG && is_active )\n {\n ac_diag = ac_diag + A_vals[16 * A_diag[a_row_id] + lane_id_mod_16];\n }\n\n // Load the range of the row of A.\n int a_col_begin = 0, a_col_end = 0;\n\n if ( is_active )\n {\n a_col_begin = A_rows[a_row_id + 0];\n a_col_end = A_rows[a_row_id + 1];\n }\n\n // Iterate over the range of columns of B.\n for ( ; utils::any(a_col_begin < a_col_end) ; a_col_begin += 16 )\n {\n int a_col_it = a_col_begin + lane_id_mod_16;\n // Each thread loads a column-ID and an aggregate.\n int a_col_id = -1, ac_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols [a_col_it];\n ac_col_id = aggregates[a_col_id];\n }\n\n // Each thread uses the hashed index to find the position associated with the aggregate.\n int key = ac_col_id;\n\n if ( HAS_DIAG && ac_col_id == r_row_id )\n {\n key = -1;\n }\n\n int ac_idx = ac_col_it + set.find_index( key, index, false );\n\n // Iterate over the 16 items.\n for ( int k = 0 ; k < 16 ; ++k )\n {\n int uniform_ac_col = utils::shfl( ac_col_id, warp_offset + k );\n int uniform_ac_idx = utils::shfl( ac_idx, warp_offset + k );\n\n // Early loop exit.\n if ( utils::all( uniform_ac_col == -1 ) )\n {\n break;\n }\n\n // The index of the item.\n const int uniform_a_col_it = a_col_begin + k;\n // Load the value if possible.\n Value_type a_value(types::util::get_zero());\n\n if ( uniform_a_col_it < a_col_end )\n {\n a_value = A_vals[16 * uniform_a_col_it + lane_id_mod_16];\n }\n\n // Proceed diagonal if needed.\n if ( HAS_DIAG && uniform_ac_col == r_row_id )\n {\n ac_diag = ac_diag + a_value;\n uniform_ac_col = -1;\n }\n\n // Get the id of the column computed by the other half warp.\n int other_ac_col = utils::shfl_xor( uniform_ac_col, 16 );\n // If both half warps want to write to the same location, we have a conflict!!!\n int are_fighting = uniform_ac_col == other_ac_col;\n\n // Reduce the two values to a single one.\n if ( uniform_ac_col != -1 && are_fighting )\n {\n a_value = a_value + utils::shfl_xor( a_value, 16 );\n }\n\n // If the two half warps fight, only one can be the winner... It's the 1st half!!!\n int is_winner = !are_fighting || lane_id_div_16 == 0;\n\n // Update the value.\n if ( uniform_ac_col != -1 && is_winner )\n {\n Ac_vals[16 * uniform_ac_idx + lane_id_mod_16] = Ac_vals[16 * uniform_ac_idx + lane_id_mod_16] + a_value;\n }\n }\n }\n }\n\n if ( HAS_DIAG )\n {\n ac_diag = ac_diag + utils::shfl_xor( ac_diag, 16 );\n\n if ( lane_id_div_16 == 0 )\n {\n Ac_vals[16 * Ac_diag[r_row_id] + lane_id_mod_16] = ac_diag;\n }\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int N, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool HAS_DIAG, bool FORCE_DETERMINISM >\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700\n__global__ __launch_bounds__( CTA_SIZE, 8 )\n#else\n__global__ __launch_bounds__( CTA_SIZE, 8 )\n#endif\nvoid fill_A_kernel_NxN( const int R_num_rows, // same as num_aggregates.\n const int *R_rows,\n const int *R_cols,\n const int *A_rows,\n const int *A_cols,\n const int *A_diag,\n const Value_type *A_vals,\n const int *aggregates,\n const int *Ac_rows,\n const int *Ac_cols,\n const int *Ac_pos,\n const int *Ac_diag,\n Value_type *Ac_vals,\n const int gmem_size,\n int *g_keys,\n int *g_idx,\n int *wk_work_queue )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int T_WARP = FORCE_DETERMINISM ? 1 : WARP_SIZE / NxN;\n const int NUM_ITEMS_PER_WARP = T_WARP == 0 ? 1 : T_WARP;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_NxN = lane_id / NxN;\n const int lane_id_mod_NxN = lane_id % NxN;\n const int warp_offset = NxN * lane_id_div_NxN;\n // First threads load the row IDs of A needed by the CTA...\n int r_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // My index.\n Hash_index index( &g_idx[r_row_id * gmem_size] );\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[r_row_id * gmem_size], gmem_size );\n // Loop over rows of R.\n\n for ( ; r_row_id < R_num_rows ; r_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // The indices of the row.\n int ac_col_it = Ac_rows[r_row_id + 0];\n int ac_col_end = Ac_rows[r_row_id + 1];\n // Clear the set first.\n set.clear(true);\n // Populate the index.\n set.load_index( ac_col_end - ac_col_it, &Ac_cols[ac_col_it], &Ac_pos[ac_col_it], index, false );\n // Load the range of the row.\n int r_col_it = R_rows[r_row_id + 0];\n int r_col_end = R_rows[r_row_id + 1];\n // Diagonal value (each half warp stores a diagonal element).\n Value_type ac_diag(types::util::get_zero());\n\n // Iterate over the columns of R.\n for ( r_col_it += lane_id_div_NxN ; utils::any(r_col_it < r_col_end) ; r_col_it += NUM_ITEMS_PER_WARP )\n {\n // Is it an active thread.\n const bool is_active = r_col_it < r_col_end && lane_id_div_NxN < NUM_ITEMS_PER_WARP;\n // Columns of R map to rows of A. Each thread of the warp loads its R-col/A-row ID.\n int a_row_id = -1;\n\n if ( is_active )\n {\n a_row_id = R_cols[r_col_it];\n }\n\n // Update the diagonal if needed.\n if ( HAS_DIAG && is_active )\n {\n ac_diag = ac_diag + A_vals[NxN * A_diag[a_row_id] + lane_id_mod_NxN];\n }\n\n // Load the range of the row of A.\n int a_col_begin = 0, a_col_end = 0;\n\n if ( is_active )\n {\n a_col_begin = A_rows[a_row_id + 0];\n a_col_end = A_rows[a_row_id + 1];\n }\n\n // Iterate over the range of columns of B.\n for ( ; utils::any(a_col_begin < a_col_end) ; a_col_begin += NxN )\n {\n int a_col_it = a_col_begin + lane_id_mod_NxN;\n // Is it active.\n const bool is_active_k = a_col_it < a_col_end && lane_id_div_NxN < NUM_ITEMS_PER_WARP;\n // Each thread loads a column-ID and an aggregate.\n int a_col_id = -1, ac_col_id = -1;\n\n if ( is_active_k )\n {\n a_col_id = A_cols [a_col_it];\n ac_col_id = aggregates[a_col_id];\n }\n\n // Each thread uses the hashed index to find the position associated with the aggregate.\n int key = ac_col_id;\n\n if ( HAS_DIAG && ac_col_id == r_row_id )\n {\n key = -1;\n }\n\n int ac_idx = ac_col_it + set.find_index( key, index, false );\n\n // Iterate over the NxN items.\n #pragma unroll 1\n for ( int k = 0 ; k < NxN ; ++k )\n {\n int uniform_ac_col = utils::shfl( ac_col_id, warp_offset + k );\n int uniform_ac_idx = utils::shfl( ac_idx, warp_offset + k );\n\n if ( lane_id_div_NxN >= NUM_ITEMS_PER_WARP )\n {\n uniform_ac_col = -1;\n uniform_ac_idx = -1;\n }\n\n // Early loop exit.\n if ( utils::all( uniform_ac_col == -1 ) )\n {\n break;\n }\n\n // The index of the item.\n const int uniform_a_col_it = a_col_begin + k;\n // Load the value if possible.\n Value_type a_value(types::util::get_zero());\n\n if ( uniform_a_col_it < a_col_end && lane_id_div_NxN < NUM_ITEMS_PER_WARP )\n {\n a_value = A_vals[NxN * uniform_a_col_it + lane_id_mod_NxN];\n }\n\n // Update the diagonal if it is a diagonal term.\n if ( HAS_DIAG && uniform_ac_col == r_row_id )\n {\n ac_diag = ac_diag + a_value;\n uniform_ac_col = -1;\n }\n\n // Update the value.\n if ( uniform_ac_col != -1 )\n {\n utils::atomic_add( &Ac_vals[NxN * uniform_ac_idx + lane_id_mod_NxN], a_value );\n }\n }\n }\n }\n\n if ( HAS_DIAG )\n {\n if ( !FORCE_DETERMINISM )\n {\n ac_diag = utils::warp_reduce( ac_diag );\n }\n\n if ( lane_id_div_NxN == 0 )\n {\n Ac_vals[NxN * Ac_diag[r_row_id] + lane_id_mod_NxN] = ac_diag;\n }\n }\n }\n}\n\n// when blocksize is larger than warp size\ntemplate< typename Value_type, int N, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool HAS_DIAG, bool FORCE_DETERMINISM, int NUM_BLOCK_ITERS_PER_WARP>\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700\n__global__ __launch_bounds__( CTA_SIZE, 8 )\n#else\n__global__ __launch_bounds__( CTA_SIZE, 8 )\n#endif\nvoid fill_A_kernel_NxN_large( const int R_num_rows, // same as num_aggregates.\n const int *R_rows,\n const int *R_cols,\n const int *A_rows,\n const int *A_cols,\n const int *A_diag,\n const Value_type *A_vals,\n const int *aggregates,\n const int *Ac_rows,\n const int *Ac_cols,\n const int *Ac_pos,\n const int *Ac_diag,\n Value_type *Ac_vals,\n const int gmem_size,\n int *g_keys,\n int *g_idx,\n int *wk_work_queue )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp. Let's be chill here and take 1 per warp for large blocks\n const int NUM_ITEMS_PER_WARP = 1;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int r_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // My index.\n Hash_index index( &g_idx[r_row_id * gmem_size] );\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[r_row_id * gmem_size], gmem_size );\n // Loop over rows of R.\n for ( ; r_row_id < R_num_rows ; r_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // The indices of the row.\n int ac_col_it = Ac_rows[r_row_id + 0];\n int ac_col_end = Ac_rows[r_row_id + 1];\n // Clear the set first.\n set.clear(true);\n // Populate the index.\n set.load_index( ac_col_end - ac_col_it, &Ac_cols[ac_col_it], &Ac_pos[ac_col_it], index, false );\n // Load the range of the row.\n int r_col_it = R_rows[r_row_id + 0];\n int r_col_end = R_rows[r_row_id + 1];\n // Diagonal value (each half warp stores a diagonal element).\n Value_type ac_diag(types::util::get_zero());\n\n // Iterate over the columns of R.\n for ( ; utils::any(r_col_it < r_col_end) ; r_col_it += NUM_ITEMS_PER_WARP )\n {\n // Columns of R map to rows of A. Each thread of the warp loads its R-col/A-row ID.\n int a_row_id = R_cols[r_col_it];\n\n // Update the diagonal if needed.\n if ( HAS_DIAG )\n {\n ac_diag = ac_diag + A_vals[NxN * A_diag[a_row_id] + lane_id];\n }\n\n // Load the range of the row of A.\n int a_col_begin = A_rows[a_row_id + 0];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Iterate over the range of columns of B.\n for ( ; utils::any(a_col_begin < a_col_end) ; a_col_begin += NxN )\n {\n int a_col_it = a_col_begin + lane_id;\n // Is it active.\n const bool is_active_k = a_col_it < a_col_end;\n // Each thread loads a column-ID and an aggregate.\n int a_col_id = -1, ac_col_id = -1;\n\n if ( is_active_k )\n {\n a_col_id = A_cols [a_col_it];\n ac_col_id = aggregates[a_col_id];\n }\n\n // Each thread uses the hashed index to find the position associated with the aggregate.\n int key = ac_col_id;\n\n if ( HAS_DIAG && ac_col_id == r_row_id )\n {\n key = -1;\n }\n\n int ac_idx = ac_col_it + set.find_index( key, index, false );\n\n // Iterate over the NxN items.\n for ( int k = 0 ; k < NxN ; ++k )\n {\n int uniform_ac_col = utils::shfl( ac_col_id, k );\n int uniform_ac_idx = utils::shfl( ac_idx, k );\n\n // Early loop exit.\n if ( utils::all( uniform_ac_col == -1 ) )\n {\n break;\n }\n\n // The index of the item.\n const int uniform_a_col_it = a_col_begin + k;\n // iterate through the block\n#pragma unroll\n\n for (int i = 0; i < NUM_BLOCK_ITERS_PER_WARP; i++)\n {\n // Load the value if possible.\n Value_type a_value(types::util::get_zero());\n\n if ( uniform_a_col_it < a_col_end && (WARP_SIZE * i + lane_id) < NxN )\n {\n a_value = A_vals[NxN * uniform_a_col_it + WARP_SIZE * i + lane_id];\n }\n\n // Update the diagonal if it is a diagonal term.\n if ( HAS_DIAG && uniform_ac_col == r_row_id )\n {\n ac_diag = ac_diag + a_value;\n uniform_ac_col = -1;\n }\n\n // Update the value.\n if ( uniform_ac_col != -1 && (WARP_SIZE * i + lane_id) < NxN)\n {\n utils::atomic_add( &Ac_vals[NxN * uniform_ac_idx + WARP_SIZE * i + lane_id], a_value );\n }\n }\n }\n }\n }\n\n if ( HAS_DIAG )\n {\n if ( !FORCE_DETERMINISM )\n {\n ac_diag = utils::warp_reduce( ac_diag );\n }\n\n Ac_vals[NxN * Ac_diag[r_row_id] + lane_id] = ac_diag;\n }\n }\n}\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nenum { WARP_SIZE = 32, SMEM_SIZE = 128 };\n\ntemplate< int CTA_SIZE, bool HAS_DIAG, bool COUNT_ONLY, typename Workspace >\nstatic\nvoid compute_sparsity_dispatch( Workspace &hash_wk,\n const int R_num_rows,\n const int *R_rows,\n const int *R_cols,\n const int *A_rows,\n const int *A_cols,\n const int *aggregates,\n int *Ac_rows,\n int *Ac_cols,\n int *Ac_pos )\n{\n cudaDeviceProp props = getDeviceProperties();\n int GRID_SIZE = (props.major >= 7) ? 1024 : 128;\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int *h_status;\n cudaError_t cuda_rc = amgx::memory::cudaMallocHost((void **) &h_status, sizeof(int));\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMallocHost failed for h_status in low_deg_coarse_A_generator\", AMGX_ERR_CUDA_FAILURE);\n }\n int *h_work_offset;\n cuda_rc = amgx::memory::cudaMallocHost((void **) &h_work_offset, sizeof(int));\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaMallocHost failed for h_work_offset in low_deg_coarse_A_generator\", AMGX_ERR_CUDA_FAILURE);\n }\n int attempt = 0;\n bool warning_printed = 0;\n\n for ( bool done = false ; !done && attempt < 10 ; ++attempt )\n {\n // Double the amount of GMEM (if needed).\n if ( attempt > 0 )\n {\n if (!warning_printed)\n {\n amgx_printf(\"WARNING: Used settings might result in degraded performance for the MG coarsener for this matrix.\\n\");\n amgx_printf(\"WARNING: You might want to try different selector or MG algorithm for better performance.\\n\");\n warning_printed = 1;\n }\n\n hash_wk.expand();\n }\n\n // Reset the status.\n int *p_status = h_status;\n *p_status = 0;\n cudaMemcpyAsync( hash_wk.get_status(), p_status, sizeof(int), cudaMemcpyHostToDevice, amgx::thrust::global_thread_handle::get_stream() );\n cudaCheckError();\n // Reset the work queue.\n int *p_work_offset = h_work_offset;\n *p_work_offset = GRID_SIZE * NUM_WARPS;\n cudaMemcpyAsync( hash_wk.get_work_queue(), p_work_offset, sizeof(int), cudaMemcpyHostToDevice, amgx::thrust::global_thread_handle::get_stream() );\n cudaCheckError();\n // Launch the kernel.\n compute_sparsity_kernel<8, CTA_SIZE, SMEM_SIZE, WARP_SIZE, HAS_DIAG, COUNT_ONLY> <<< GRID_SIZE, CTA_SIZE, 0, amgx::thrust::global_thread_handle::get_stream()>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n hash_wk.get_work_queue(),\n hash_wk.get_status() );\n cudaCheckError();\n // Read the result from count_non_zeroes.\n cudaMemcpyAsync( p_status, hash_wk.get_status(), sizeof(int), cudaMemcpyDeviceToHost, amgx::thrust::global_thread_handle::get_stream() );\n cudaCheckError();\n cuda_rc = cudaStreamSynchronize(amgx::thrust::global_thread_handle::get_stream());\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaStreamSynchronize failed in low_deg_coarse_A_generator\", AMGX_ERR_CUDA_FAILURE);\n }\n done = (*p_status == 0);\n cudaCheckError();\n }\n\n cuda_rc = amgx::memory::cudaFreeHost(h_status);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaFreeHost failed for h_status in low_deg_coarse_A_generator\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = amgx::memory::cudaFreeHost(h_work_offset);\n if (cuda_rc != cudaSuccess)\n {\n FatalError(\"cudaFreeHost failed for h_work_offset in low_deg_coarse_A_generator\", AMGX_ERR_CUDA_FAILURE);\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int CTA_SIZE, bool HAS_DIAG, typename Workspace, typename Value_type >\nstatic\nvoid fill_A_dispatch( Workspace &hash_wk,\n const int block_size,\n const int R_num_rows, // same as num_aggregates.\n const int *R_rows,\n const int *R_cols,\n const int *A_rows,\n const int *A_cols,\n const int *A_diag,\n const Value_type *A_vals,\n const int *aggregates,\n const int *Ac_rows,\n const int *Ac_cols,\n const int *Ac_pos,\n const int *Ac_diag,\n Value_type *Ac_vals,\n bool force_determinism )\n{\n cudaDeviceProp props = getDeviceProperties();\n int GRID_SIZE = (props.major >= 7) ? 1024 : 128;\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int work_offset = GRID_SIZE * NUM_WARPS;\n cudaMemcpyAsync( hash_wk.get_work_queue(), &work_offset, sizeof(int), cudaMemcpyHostToDevice, amgx::thrust::global_thread_handle::get_stream() );\n cudaCheckError();\n\n // Launch the kernel.\n switch ( block_size )\n {\n case 1:\n fill_A_kernel_1x1 <<< GRID_SIZE, CTA_SIZE>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n hash_wk.get_vals(),\n hash_wk.get_work_queue() );\n cudaCheckError();\n break;\n\n case 2:\n fill_A_kernel_NxN <<< GRID_SIZE, CTA_SIZE>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n reinterpret_cast( hash_wk.get_vals() ),\n hash_wk.get_work_queue() );\n cudaCheckError();\n break;\n\n case 3:\n fill_A_kernel_NxN <<< GRID_SIZE, CTA_SIZE>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n reinterpret_cast( hash_wk.get_vals() ),\n hash_wk.get_work_queue() );\n cudaCheckError();\n break;\n\n case 4:\n if ( force_determinism )\n fill_A_kernel_NxN <<< GRID_SIZE, CTA_SIZE>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n reinterpret_cast( hash_wk.get_vals() ),\n hash_wk.get_work_queue() );\n else\n fill_A_kernel_4x4 <<< GRID_SIZE, CTA_SIZE, 0, amgx::thrust::global_thread_handle::get_stream()>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n reinterpret_cast( hash_wk.get_vals() ),\n hash_wk.get_work_queue() );\n \n cudaCheckError();\n break;\n\n case 5:\n fill_A_kernel_NxN <<< GRID_SIZE, CTA_SIZE>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n reinterpret_cast( hash_wk.get_vals() ),\n hash_wk.get_work_queue() );\n cudaCheckError();\n break;\n\n case 8:\n fill_A_kernel_NxN_large <<< GRID_SIZE, CTA_SIZE>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n reinterpret_cast( hash_wk.get_vals() ),\n hash_wk.get_work_queue() );\n cudaCheckError();\n break;\n\n case 10:\n fill_A_kernel_NxN_large <<< GRID_SIZE, CTA_SIZE>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n reinterpret_cast( hash_wk.get_vals() ),\n hash_wk.get_work_queue() );\n cudaCheckError();\n break;\n\n default:\n FatalError( \"LOW_DEG not implemented for this block size\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE );\n }\n\n cudaCheckError();\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nLowDegCoarseAGenerator >::computeAOperator( const Matrix_d &A,\n Matrix_d &Ac,\n const IVector &aggregates,\n const IVector &R_row_offsets,\n const IVector &R_column_indices,\n const int num_aggregates )\n{\n if ( A.get_block_dimx() != A.get_block_dimy() )\n {\n FatalError( \"LowDegCoarseAGenerator implemented for squared blocks\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE );\n }\n\n // The matrix Ac will be modified.\n Ac.set_initialized(0);\n // Is the diagonal stored separately??\n const int diag_prop = A.hasProps(DIAG);\n // Allocate a workspace for hashing.\n typedef TemplateConfig TConfig_d;\n\n cudaDeviceProp props = getDeviceProperties();\n int grid_size = (props.major >= 7) ? 1024 : 128;\n\n Hash_Workspace hash_wk(true, grid_size);\n\n // Compute row offsets of Ac.\n Ac.addProps(CSR);\n Ac.set_num_rows( num_aggregates );\n Ac.set_num_cols( num_aggregates );\n Ac.row_offsets.resize( num_aggregates + 1 );\n // Compute the number of non-zero elements per row of Ac.\n const int CTA_SIZE = 128;\n\n if ( diag_prop )\n compute_sparsity_dispatch(\n hash_wk,\n num_aggregates,\n R_row_offsets.raw(),\n R_column_indices.raw(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n aggregates.raw(),\n Ac.row_offsets.raw(),\n NULL,\n NULL );\n else\n compute_sparsity_dispatch(\n hash_wk,\n num_aggregates,\n R_row_offsets.raw(),\n R_column_indices.raw(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n aggregates.raw(),\n Ac.row_offsets.raw(),\n NULL,\n NULL );\n\n cudaCheckError();\n // Compute the number of non-zeroes.\n thrust_wrapper::exclusive_scan( Ac.row_offsets.begin(), Ac.row_offsets.end(), Ac.row_offsets.begin() );\n cudaCheckError();\n int nonzero_blocks = Ac.row_offsets[num_aggregates];\n\n if ( diag_prop )\n {\n Ac.addProps(DIAG);\n }\n\n if ( A.is_matrix_singleGPU() )\n {\n Ac.resize( num_aggregates, num_aggregates, nonzero_blocks, A.get_block_dimy(), A.get_block_dimx(), !diag_prop );\n }\n else\n {\n //have 3% more nz for storage\n Ac.resize_spare( num_aggregates, num_aggregates, nonzero_blocks, A.get_block_dimy(), A.get_block_dimx(), 1.0 );\n\n if ( diag_prop )\n {\n Ac.computeDiagonal();\n }\n }\n\n // Vector to store the positions in the hash table.\n device_vector_alloc Ac_pos(nonzero_blocks);\n\n // Compute the sparsity pattern of the rows of Ac.\n if ( diag_prop )\n compute_sparsity_dispatch(\n hash_wk,\n num_aggregates,\n R_row_offsets.raw(),\n R_column_indices.raw(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n aggregates.raw(),\n Ac.row_offsets.raw(),\n Ac.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &Ac_pos.front() ));\n else\n compute_sparsity_dispatch(\n hash_wk,\n num_aggregates,\n R_row_offsets.raw(),\n R_column_indices.raw(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n aggregates.raw(),\n Ac.row_offsets.raw(),\n Ac.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &Ac_pos.front() ));\n\n cudaCheckError();\n\n // Reset values if needed.\n if ( A.get_block_dimy() != 1 )\n {\n thrust_wrapper::fill( Ac.values.begin(), Ac.values.end(), types::util::get_zero() );\n cudaCheckError();\n }\n\n // Compute values.\n if ( diag_prop )\n {\n fill_A_dispatch(\n hash_wk,\n A.get_block_dimy(),\n num_aggregates,\n R_row_offsets.raw(),\n R_column_indices.raw(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n aggregates.raw(),\n Ac.row_offsets.raw(),\n Ac.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &Ac_pos.front() ),\n Ac.diag.raw(),\n Ac.values.raw(),\n this->m_force_determinism );\n }\n else\n {\n fill_A_dispatch(\n hash_wk,\n A.get_block_dimy(),\n num_aggregates,\n R_row_offsets.raw(),\n R_column_indices.raw(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n aggregates.raw(),\n Ac.row_offsets.raw(),\n Ac.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &Ac_pos.front() ),\n Ac.diag.raw(),\n Ac.values.raw(),\n this->m_force_determinism );\n }\n\n cudaCheckError();\n\n // Update the diagonal if needed.\n if ( Ac.is_matrix_singleGPU() )\n {\n Ac.computeDiagonal();\n }\n\n cudaCheckError();\n // Finalize the modification.\n Ac.set_initialized(1);\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nLowDegCoarseAGenerator >::computeAOperator( const Matrix_h &h_A,\n Matrix_h &h_Ac,\n const IVector &h_aggregates,\n const IVector &h_R_row_offsets,\n const IVector &h_R_column_indices,\n const int num_aggregates )\n{\n h_Ac.set_initialized(0);\n IVector rows;\n IVector inds;\n typename Matrix_h::MVector vals;\n typename Matrix_h::MVector diag;\n int num_nnz = 0;\n int diag_prop = h_A.hasProps(DIAG);\n\n for ( int row = 0; row < num_aggregates; row++ )\n {\n for ( int col = 0; col < num_aggregates; col++ )\n {\n int fill = 0;\n typename Matrix_h::MVector cur(h_A.get_block_size(), types::util::get_zero());\n\n for ( int rc = h_R_row_offsets[row]; rc < h_R_row_offsets[row + 1]; rc++ )\n {\n int j = h_R_column_indices[rc];\n\n for ( int ac = h_A.row_offsets[j]; ac < h_A.row_offsets[j + 1] + diag_prop; ac++ )\n {\n int k = (ac == h_A.row_offsets[j + 1]) ? j : h_A.col_indices[ac];\n\n for ( int q = h_R_row_offsets[col]; q < h_R_row_offsets[col + 1]; q++ )\n if ( k == h_R_column_indices[q] )\n {\n fill = 1;\n int val_idx = (ac == h_A.row_offsets[j + 1]) ? h_A.get_num_nz() + j : ac;\n\n for ( int v = 0; v < h_A.get_block_size(); v++)\n {\n cur[v] = cur[v] + h_A.values[val_idx * h_A.get_block_size() + v];\n }\n }\n }\n }\n\n if ( fill )\n {\n if ( row != col || !diag_prop )\n {\n inds.push_back(col);\n rows.push_back(row);\n num_nnz++;\n\n for ( int v = 0; v < h_A.get_block_size(); v++ )\n {\n vals.push_back(cur[v]);\n }\n }\n else\n {\n for ( int v = 0; v < h_A.get_block_size(); v++ )\n {\n diag.push_back(cur[v]);\n }\n }\n }\n }\n }\n\n rows.push_back(-1);\n\n // add diagonal to the end\n if ( diag_prop )\n {\n for ( int v = 0; v < num_aggregates * h_A.get_block_size(); v++ )\n {\n vals.push_back(diag[v]);\n }\n }\n else\n {\n // Add a zero at the end\n for (int v = 0; v < h_A.get_block_size(); v++)\n {\n vals.push_back(types::util::get_zero());\n }\n }\n\n h_Ac.resize(num_aggregates, num_aggregates, num_nnz, h_A.get_block_dimx(), h_A.get_block_dimy(), 1);\n h_Ac.row_indices = rows;\n h_Ac.col_indices = inds;\n h_Ac.values = vals;\n h_Ac.addProps( CSR | ( diag_prop ? DIAG : 0 ) );\n h_Ac.computeDiagonal();\n h_Ac.set_initialized(1);\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define AMGX_CASE_LINE(CASE) template class LowDegCoarseAGenerator::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace aggregation\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\nnamespace aggregation\n{\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#include \n#include \n\nstatic __device__ __forceinline__ int get_work( int *queue, int warp_id, int count = 1 )\n{\n int offset = -1;\n\n if ( utils::lane_id() == 0 )\n {\n offset = atomicAdd( queue, count );\n }\n\n return utils::shfl( offset, 0 );\n}\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool HAS_DIAG, bool COUNT_ONLY >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_sparsity_kernel( const int R_num_rows, // same as num_aggregates.\n const int *R_rows,\n const int *R_cols,\n const int *A_rows,\n const int *A_cols,\n const int *aggregates,\n int *Ac_rows,\n int *Ac_cols,\n int *Ac_pos,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int r_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[r_row_id * gmem_size], gmem_size );\n // Loop over rows of R.\n for ( ; r_row_id < R_num_rows ; r_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Make sure we have to proceed.\n if ( COUNT_ONLY )\n {\n volatile int *status = reinterpret_cast( wk_status );\n\n if ( set.has_failed() || *status != 0 )\n {\n return;\n }\n }\n\n // Clear the set.\n set.clear();\n // Load the range of the row.\n int r_col_it = R_rows[r_row_id + 0];\n int r_col_end = R_rows[r_row_id + 1];\n\n // Iterate over the columns of R.\n for ( r_col_it += lane_id ; utils::any(r_col_it < r_col_end) ; r_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = r_col_it < r_col_end;\n // Columns of R map to rows of A. Each thread of the warp loads its R-col/A-row ID.\n int a_row_id = -1;\n\n if ( is_active )\n {\n a_row_id = R_cols[r_col_it];\n }\n\n const int num_rows = __popc( utils::ballot(is_active) );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n const int uniform_a_row_id = utils::shfl( a_row_id, local_k );\n // Load the range of the row of B.\n int a_col_it = 0, a_col_end = 0;\n\n if ( is_active_k )\n {\n a_col_it = A_rows[uniform_a_row_id + 0];\n a_col_end = A_rows[uniform_a_row_id + 1];\n }\n\n // Iterate over the range of columns of B.\n for ( a_col_it += lane_id_mod_num_threads ; utils::any(a_col_it < a_col_end) ; a_col_it += NUM_THREADS_PER_ROW )\n {\n int a_col_id = -1, a_agg_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n a_agg_id = aggregates[a_col_id];\n }\n\n if ( HAS_DIAG && a_agg_id == r_row_id )\n {\n a_agg_id = -1;\n }\n\n set.insert( a_agg_id, COUNT_ONLY ? wk_status : NULL );\n }\n }\n }\n\n // Store the results.\n if ( COUNT_ONLY )\n {\n int count = set.compute_size();\n\n if ( lane_id == 0 )\n {\n Ac_rows[r_row_id] = count;\n }\n }\n else\n {\n int ac_col_it = Ac_rows[r_row_id];\n set.store_with_positions( &Ac_cols[ac_col_it], &Ac_pos[ac_col_it] );\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int NUM_THREADS_PER_ROW, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool HAS_DIAG >\n__global__\n__launch_bounds__( CTA_SIZE, 8 )\nvoid fill_A_kernel_1x1( const int R_num_rows,\n const int *R_rows,\n const int *R_cols,\n const int *A_rows,\n const int *A_cols,\n const int *A_diag,\n const Value_type *A_vals,\n const int *aggregates,\n const int *Ac_rows,\n const int *Ac_cols,\n const int *Ac_pos,\n const int *Ac_diag,\n Value_type *Ac_vals,\n int gmem_size,\n int *g_keys,\n Value_type *g_vals,\n int *wk_work_queue )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The hash values stored in shared memory.\n __shared__ Value_type s_vals[NUM_WARPS * SMEM_SIZE]; \n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int r_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_map map( &s_keys[warp_id * SMEM_SIZE ],\n &g_keys[r_row_id * gmem_size ],\n &s_vals[warp_id * SMEM_SIZE],\n &g_vals[r_row_id * gmem_size ], gmem_size );\n // Loop over rows of A.\n for ( ; r_row_id < R_num_rows ; r_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // The indices of the output row.\n int ac_col_it = Ac_rows[r_row_id + 0];\n int ac_col_end = Ac_rows[r_row_id + 1];\n // Clear the set first. TODO: Make sure it's needed. I don't think it is!!!!\n map.clear();\n // Populate the map.\n map.load( ac_col_end - ac_col_it, &Ac_cols[ac_col_it], &Ac_pos[ac_col_it] );\n int r_col_it = R_rows[r_row_id + 0];\n int r_col_end = R_rows[r_row_id + 1];\n // The diagonal.\n Value_type r_diag(types::util::get_zero());\n\n // _iterate over the columns of A to build C_hat.\n for ( r_col_it += lane_id ; utils::any(r_col_it < r_col_end) ; r_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = r_col_it < r_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_row_id = -1;\n\n if ( is_active )\n {\n a_row_id = R_cols[r_col_it];\n }\n\n // Update the diagonal (if needed).\n if ( HAS_DIAG && is_active )\n {\n r_diag = r_diag + A_vals[A_diag[a_row_id]];\n }\n\n const int num_rows = __popc( utils::ballot(is_active) );\n\n // Uniform loop: threads collaborate to load other elements.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n const int uniform_a_row_id = utils::shfl( a_row_id, local_k );\n // The range of the row of B.\n int a_col_it = 0, a_col_end = 0;\n\n if ( local_k < num_rows )\n {\n a_col_it = utils::Ld::load( &A_rows[uniform_a_row_id + 0] );\n a_col_end = utils::Ld::load( &A_rows[uniform_a_row_id + 1] );\n }\n\n // Iterate over the range of columns of B.\n for ( a_col_it += lane_id_mod_num_threads ; utils::any(a_col_it < a_col_end) ; a_col_it += NUM_THREADS_PER_ROW )\n {\n // Load columns and values.\n int a_col_id = -1;\n Value_type a_value(types::util::get_zero());\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n a_value = A_vals[a_col_it];\n }\n\n // Find the aggregate.\n int a_agg_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_agg_id = aggregates[a_col_id];\n }\n\n // Update the diag/hash map.\n if ( HAS_DIAG && a_agg_id == r_row_id )\n {\n r_diag = r_diag + a_value;\n a_agg_id = -1;\n }\n\n map.insert( a_agg_id, a_value, NULL ); // It won't insert. Only update.\n }\n }\n }\n\n // Update the diagonal.\n if ( HAS_DIAG )\n {\n r_diag = utils::warp_reduce<1, utils::Add>( r_diag );\n\n if ( lane_id == 0 )\n {\n Ac_vals[Ac_diag[r_row_id]] = r_diag;\n }\n }\n\n // Store the results.\n int count = ac_col_end - ac_col_it;\n\n if ( count == 0 )\n {\n continue;\n }\n\n map.store( count, &Ac_vals[ac_col_it] );\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool HAS_DIAG >\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700\n__global__ __launch_bounds__( CTA_SIZE, 8 )\n#else \n__global__ __launch_bounds__( CTA_SIZE, 8 )\n#endif\nvoid fill_A_kernel_4x4( const int R_num_rows, // same as num_aggregates.\n const int *R_rows,\n const int *R_cols,\n const int *A_rows,\n const int *A_cols,\n const int *A_diag,\n const Value_type *A_vals,\n const int *aggregates,\n const int *Ac_rows,\n const int *Ac_cols,\n const int *Ac_pos,\n const int *Ac_diag,\n Value_type *Ac_vals,\n const int gmem_size,\n int *g_keys,\n int *g_idx,\n int *wk_work_queue )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_16 = lane_id / 16;\n const int lane_id_mod_16 = lane_id % 16;\n const int warp_offset = 16 * lane_id_div_16;\n // First threads load the row IDs of A needed by the CTA...\n int r_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // My index.\n Hash_index index( &g_idx[r_row_id * gmem_size] );\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[r_row_id * gmem_size], gmem_size );\n // Loop over rows of R.\n for ( ; r_row_id < R_num_rows ; r_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // The indices of the row.\n int ac_col_it = Ac_rows[r_row_id + 0];\n int ac_col_end = Ac_rows[r_row_id + 1];\n // Clear the set first.\n set.clear(true);\n // Populate the index.\n set.load_index( ac_col_end - ac_col_it, &Ac_cols[ac_col_it], &Ac_pos[ac_col_it], index, false );\n // Load the range of the row.\n int r_col_it = R_rows[r_row_id + 0];\n int r_col_end = R_rows[r_row_id + 1];\n // Diagonal value (each half warp stores a diagonal element).\n Value_type ac_diag(types::util::get_zero());\n\n // Iterate over the columns of R.\n for ( r_col_it += lane_id_div_16 ; utils::any(r_col_it < r_col_end) ; r_col_it += 2 )\n {\n // Is it an active thread.\n const bool is_active = r_col_it < r_col_end;\n // Columns of R map to rows of A. Each thread of the warp loads its R-col/A-row ID.\n int a_row_id = -1;\n\n if ( is_active )\n {\n a_row_id = R_cols[r_col_it];\n }\n\n // Update the diagonal if needed.\n if ( HAS_DIAG && is_active )\n {\n ac_diag = ac_diag + A_vals[16 * A_diag[a_row_id] + lane_id_mod_16];\n }\n\n // Load the range of the row of A.\n int a_col_begin = 0, a_col_end = 0;\n\n if ( is_active )\n {\n a_col_begin = A_rows[a_row_id + 0];\n a_col_end = A_rows[a_row_id + 1];\n }\n\n // Iterate over the range of columns of B.\n for ( ; utils::any(a_col_begin < a_col_end) ; a_col_begin += 16 )\n {\n int a_col_it = a_col_begin + lane_id_mod_16;\n // Each thread loads a column-ID and an aggregate.\n int a_col_id = -1, ac_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols [a_col_it];\n ac_col_id = aggregates[a_col_id];\n }\n\n // Each thread uses the hashed index to find the position associated with the aggregate.\n int key = ac_col_id;\n\n if ( HAS_DIAG && ac_col_id == r_row_id )\n {\n key = -1;\n }\n\n int ac_idx = ac_col_it + set.find_index( key, index, false );\n\n // Iterate over the 16 items.\n for ( int k = 0 ; k < 16 ; ++k )\n {\n int uniform_ac_col = utils::shfl( ac_col_id, warp_offset + k );\n int uniform_ac_idx = utils::shfl( ac_idx, warp_offset + k );\n\n // Early loop exit.\n if ( utils::all( uniform_ac_col == -1 ) )\n {\n break;\n }\n\n // The index of the item.\n const int uniform_a_col_it = a_col_begin + k;\n // Load the value if possible.\n Value_type a_value(types::util::get_zero());\n\n if ( uniform_a_col_it < a_col_end )\n {\n a_value = A_vals[16 * uniform_a_col_it + lane_id_mod_16];\n }\n\n // Proceed diagonal if needed.\n if ( HAS_DIAG && uniform_ac_col == r_row_id )\n {\n ac_diag = ac_diag + a_value;\n uniform_ac_col = -1;\n }\n\n // Get the id of the column computed by the other half warp.\n int other_ac_col = utils::shfl_xor( uniform_ac_col, 16 );\n // If both half warps want to write to the same location, we have a conflict!!!\n int are_fighting = uniform_ac_col == other_ac_col;\n\n // Reduce the two values to a single one.\n if ( uniform_ac_col != -1 && are_fighting )\n {\n a_value = a_value + utils::shfl_xor( a_value, 16 );\n }\n\n // If the two half warps fight, only one can be the winner... It's the 1st half!!!\n int is_winner = !are_fighting || lane_id_div_16 == 0;\n\n // Update the value.\n if ( uniform_ac_col != -1 && is_winner )\n {\n Ac_vals[16 * uniform_ac_idx + lane_id_mod_16] = Ac_vals[16 * uniform_ac_idx + lane_id_mod_16] + a_value;\n }\n }\n }\n }\n\n if ( HAS_DIAG )\n {\n ac_diag = ac_diag + utils::shfl_xor( ac_diag, 16 );\n\n if ( lane_id_div_16 == 0 )\n {\n Ac_vals[16 * Ac_diag[r_row_id] + lane_id_mod_16] = ac_diag;\n }\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int N, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool HAS_DIAG, bool FORCE_DETERMINISM >\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700\n__global__ __launch_bounds__( CTA_SIZE, 8 )\n#else\n__global__ __launch_bounds__( CTA_SIZE, 8 )\n#endif\nvoid fill_A_kernel_NxN( const int R_num_rows, // same as num_aggregates.\n const int *R_rows,\n const int *R_cols,\n const int *A_rows,\n const int *A_cols,\n const int *A_diag,\n const Value_type *A_vals,\n const int *aggregates,\n const int *Ac_rows,\n const int *Ac_cols,\n const int *Ac_pos,\n const int *Ac_diag,\n Value_type *Ac_vals,\n const int gmem_size,\n int *g_keys,\n int *g_idx,\n int *wk_work_queue )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp.\n const int T_WARP = FORCE_DETERMINISM ? 1 : WARP_SIZE / NxN;\n const int NUM_ITEMS_PER_WARP = T_WARP == 0 ? 1 : T_WARP;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_NxN = lane_id / NxN;\n const int lane_id_mod_NxN = lane_id % NxN;\n const int warp_offset = NxN * lane_id_div_NxN;\n // First threads load the row IDs of A needed by the CTA...\n int r_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // My index.\n Hash_index index( &g_idx[r_row_id * gmem_size] );\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[r_row_id * gmem_size], gmem_size );\n // Loop over rows of R.\n\n for ( ; r_row_id < R_num_rows ; r_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // The indices of the row.\n int ac_col_it = Ac_rows[r_row_id + 0];\n int ac_col_end = Ac_rows[r_row_id + 1];\n // Clear the set first.\n set.clear(true);\n // Populate the index.\n set.load_index( ac_col_end - ac_col_it, &Ac_cols[ac_col_it], &Ac_pos[ac_col_it], index, false );\n // Load the range of the row.\n int r_col_it = R_rows[r_row_id + 0];\n int r_col_end = R_rows[r_row_id + 1];\n // Diagonal value (each half warp stores a diagonal element).\n Value_type ac_diag(types::util::get_zero());\n\n // Iterate over the columns of R.\n for ( r_col_it += lane_id_div_NxN ; utils::any(r_col_it < r_col_end) ; r_col_it += NUM_ITEMS_PER_WARP )\n {\n // Is it an active thread.\n const bool is_active = r_col_it < r_col_end && lane_id_div_NxN < NUM_ITEMS_PER_WARP;\n // Columns of R map to rows of A. Each thread of the warp loads its R-col/A-row ID.\n int a_row_id = -1;\n\n if ( is_active )\n {\n a_row_id = R_cols[r_col_it];\n }\n\n // Update the diagonal if needed.\n if ( HAS_DIAG && is_active )\n {\n ac_diag = ac_diag + A_vals[NxN * A_diag[a_row_id] + lane_id_mod_NxN];\n }\n\n // Load the range of the row of A.\n int a_col_begin = 0, a_col_end = 0;\n\n if ( is_active )\n {\n a_col_begin = A_rows[a_row_id + 0];\n a_col_end = A_rows[a_row_id + 1];\n }\n\n // Iterate over the range of columns of B.\n for ( ; utils::any(a_col_begin < a_col_end) ; a_col_begin += NxN )\n {\n int a_col_it = a_col_begin + lane_id_mod_NxN;\n // Is it active.\n const bool is_active_k = a_col_it < a_col_end && lane_id_div_NxN < NUM_ITEMS_PER_WARP;\n // Each thread loads a column-ID and an aggregate.\n int a_col_id = -1, ac_col_id = -1;\n\n if ( is_active_k )\n {\n a_col_id = A_cols [a_col_it];\n ac_col_id = aggregates[a_col_id];\n }\n\n // Each thread uses the hashed index to find the position associated with the aggregate.\n int key = ac_col_id;\n\n if ( HAS_DIAG && ac_col_id == r_row_id )\n {\n key = -1;\n }\n\n int ac_idx = ac_col_it + set.find_index( key, index, false );\n\n // Iterate over the NxN items.\n #pragma unroll 1\n for ( int k = 0 ; k < NxN ; ++k )\n {\n int uniform_ac_col = utils::shfl( ac_col_id, warp_offset + k );\n int uniform_ac_idx = utils::shfl( ac_idx, warp_offset + k );\n\n if ( lane_id_div_NxN >= NUM_ITEMS_PER_WARP )\n {\n uniform_ac_col = -1;\n uniform_ac_idx = -1;\n }\n\n // Early loop exit.\n if ( utils::all( uniform_ac_col == -1 ) )\n {\n break;\n }\n\n // The index of the item.\n const int uniform_a_col_it = a_col_begin + k;\n // Load the value if possible.\n Value_type a_value(types::util::get_zero());\n\n if ( uniform_a_col_it < a_col_end && lane_id_div_NxN < NUM_ITEMS_PER_WARP )\n {\n a_value = A_vals[NxN * uniform_a_col_it + lane_id_mod_NxN];\n }\n\n // Update the diagonal if it is a diagonal term.\n if ( HAS_DIAG && uniform_ac_col == r_row_id )\n {\n ac_diag = ac_diag + a_value;\n uniform_ac_col = -1;\n }\n\n // Update the value.\n if ( uniform_ac_col != -1 )\n {\n utils::atomic_add( &Ac_vals[NxN * uniform_ac_idx + lane_id_mod_NxN], a_value );\n }\n }\n }\n }\n\n if ( HAS_DIAG )\n {\n if ( !FORCE_DETERMINISM )\n {\n ac_diag = utils::warp_reduce( ac_diag );\n }\n\n if ( lane_id_div_NxN == 0 )\n {\n Ac_vals[NxN * Ac_diag[r_row_id] + lane_id_mod_NxN] = ac_diag;\n }\n }\n }\n}\n\n// when blocksize is larger than warp size\ntemplate< typename Value_type, int N, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, bool HAS_DIAG, bool FORCE_DETERMINISM, int NUM_BLOCK_ITERS_PER_WARP>\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700\n__global__ __launch_bounds__( CTA_SIZE, 8 )\n#else\n__global__ __launch_bounds__( CTA_SIZE, 8 )\n#endif\nvoid fill_A_kernel_NxN_large( const int R_num_rows, // same as num_aggregates.\n const int *R_rows,\n const int *R_cols,\n const int *A_rows,\n const int *A_cols,\n const int *A_diag,\n const Value_type *A_vals,\n const int *aggregates,\n const int *Ac_rows,\n const int *Ac_cols,\n const int *Ac_pos,\n const int *Ac_diag,\n Value_type *Ac_vals,\n const int gmem_size,\n int *g_keys,\n int *g_idx,\n int *wk_work_queue )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Squared N.\n const int NxN = N * N;\n // Number of items per warp. Let's be chill here and take 1 per warp for large blocks\n const int NUM_ITEMS_PER_WARP = 1;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int r_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // My index.\n Hash_index index( &g_idx[r_row_id * gmem_size] );\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[r_row_id * gmem_size], gmem_size );\n // Loop over rows of R.\n for ( ; r_row_id < R_num_rows ; r_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // The indices of the row.\n int ac_col_it = Ac_rows[r_row_id + 0];\n int ac_col_end = Ac_rows[r_row_id + 1];\n // Clear the set first.\n set.clear(true);\n // Populate the index.\n set.load_index( ac_col_end - ac_col_it, &Ac_cols[ac_col_it], &Ac_pos[ac_col_it], index, false );\n // Load the range of the row.\n int r_col_it = R_rows[r_row_id + 0];\n int r_col_end = R_rows[r_row_id + 1];\n // Diagonal value (each half warp stores a diagonal element).\n Value_type ac_diag(types::util::get_zero());\n\n // Iterate over the columns of R.\n for ( ; utils::any(r_col_it < r_col_end) ; r_col_it += NUM_ITEMS_PER_WARP )\n {\n // Columns of R map to rows of A. Each thread of the warp loads its R-col/A-row ID.\n int a_row_id = R_cols[r_col_it];\n\n // Update the diagonal if needed.\n if ( HAS_DIAG )\n {\n ac_diag = ac_diag + A_vals[NxN * A_diag[a_row_id] + lane_id];\n }\n\n // Load the range of the row of A.\n int a_col_begin = A_rows[a_row_id + 0];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Iterate over the range of columns of B.\n for ( ; utils::any(a_col_begin < a_col_end) ; a_col_begin += NxN )\n {\n int a_col_it = a_col_begin + lane_id;\n // Is it active.\n const bool is_active_k = a_col_it < a_col_end;\n // Each thread loads a column-ID and an aggregate.\n int a_col_id = -1, ac_col_id = -1;\n\n if ( is_active_k )\n {\n a_col_id = A_cols [a_col_it];\n ac_col_id = aggregates[a_col_id];\n }\n\n // Each thread uses the hashed index to find the position associated with the aggregate.\n int key = ac_col_id;\n\n if ( HAS_DIAG && ac_col_id == r_row_id )\n {\n key = -1;\n }\n\n int ac_idx = ac_col_it + set.find_index( key, index, false );\n\n // Iterate over the NxN items.\n for ( int k = 0 ; k < NxN ; ++k )\n {\n int uniform_ac_col = utils::shfl( ac_col_id, k );\n int uniform_ac_idx = utils::shfl( ac_idx, k );\n\n // Early loop exit.\n if ( utils::all( uniform_ac_col == -1 ) )\n {\n break;\n }\n\n // The index of the item.\n const int uniform_a_col_it = a_col_begin + k;\n // iterate through the block\n#pragma unroll\n\n for (int i = 0; i < NUM_BLOCK_ITERS_PER_WARP; i++)\n {\n // Load the value if possible.\n Value_type a_value(types::util::get_zero());\n\n if ( uniform_a_col_it < a_col_end && (WARP_SIZE * i + lane_id) < NxN )\n {\n a_value = A_vals[NxN * uniform_a_col_it + WARP_SIZE * i + lane_id];\n }\n\n // Update the diagonal if it is a diagonal term.\n if ( HAS_DIAG && uniform_ac_col == r_row_id )\n {\n ac_diag = ac_diag + a_value;\n uniform_ac_col = -1;\n }\n\n // Update the value.\n if ( uniform_ac_col != -1 && (WARP_SIZE * i + lane_id) < NxN)\n {\n utils::atomic_add( &Ac_vals[NxN * uniform_ac_idx + WARP_SIZE * i + lane_id], a_value );\n }\n }\n }\n }\n }\n\n if ( HAS_DIAG )\n {\n if ( !FORCE_DETERMINISM )\n {\n ac_diag = utils::warp_reduce( ac_diag );\n }\n\n Ac_vals[NxN * Ac_diag[r_row_id] + lane_id] = ac_diag;\n }\n }\n}\n\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nenum { WARP_SIZE = 32, SMEM_SIZE = 128 };\n\ntemplate< int CTA_SIZE, bool HAS_DIAG, bool COUNT_ONLY, typename Workspace >\nstatic\nvoid compute_sparsity_dispatch( Workspace &hash_wk,\n const int R_num_rows,\n const int *R_rows,\n const int *R_cols,\n const int *A_rows,\n const int *A_cols,\n const int *aggregates,\n int *Ac_rows,\n int *Ac_cols,\n int *Ac_pos )\n{\n hipDeviceProp_t props = getDeviceProperties();\n int GRID_SIZE = (props.major >= 7) ? 1024 : 128;\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int *h_status;\n hipError_t cuda_rc = amgx::memory::hipHostMalloc((void **) &h_status, sizeof(int));\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipHostMalloc failed for h_status in low_deg_coarse_A_generator\", AMGX_ERR_CUDA_FAILURE);\n }\n int *h_work_offset;\n cuda_rc = amgx::memory::hipHostMalloc((void **) &h_work_offset, sizeof(int));\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipHostMalloc failed for h_work_offset in low_deg_coarse_A_generator\", AMGX_ERR_CUDA_FAILURE);\n }\n int attempt = 0;\n bool warning_printed = 0;\n\n for ( bool done = false ; !done && attempt < 10 ; ++attempt )\n {\n // Double the amount of GMEM (if needed).\n if ( attempt > 0 )\n {\n if (!warning_printed)\n {\n amgx_printf(\"WARNING: Used settings might result in degraded performance for the MG coarsener for this matrix.\\n\");\n amgx_printf(\"WARNING: You might want to try different selector or MG algorithm for better performance.\\n\");\n warning_printed = 1;\n }\n\n hash_wk.expand();\n }\n\n // Reset the status.\n int *p_status = h_status;\n *p_status = 0;\n hipMemcpyAsync( hash_wk.get_status(), p_status, sizeof(int), hipMemcpyHostToDevice, amgx::thrust::global_thread_handle::get_stream() );\n cudaCheckError();\n // Reset the work queue.\n int *p_work_offset = h_work_offset;\n *p_work_offset = GRID_SIZE * NUM_WARPS;\n hipMemcpyAsync( hash_wk.get_work_queue(), p_work_offset, sizeof(int), hipMemcpyHostToDevice, amgx::thrust::global_thread_handle::get_stream() );\n cudaCheckError();\n // Launch the kernel.\n compute_sparsity_kernel<8, CTA_SIZE, SMEM_SIZE, WARP_SIZE, HAS_DIAG, COUNT_ONLY> <<< GRID_SIZE, CTA_SIZE, 0, amgx::thrust::global_thread_handle::get_stream()>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n hash_wk.get_work_queue(),\n hash_wk.get_status() );\n cudaCheckError();\n // Read the result from count_non_zeroes.\n hipMemcpyAsync( p_status, hash_wk.get_status(), sizeof(int), hipMemcpyDeviceToHost, amgx::thrust::global_thread_handle::get_stream() );\n cudaCheckError();\n cuda_rc = hipStreamSynchronize(amgx::thrust::global_thread_handle::get_stream());\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipStreamSynchronize failed in low_deg_coarse_A_generator\", AMGX_ERR_CUDA_FAILURE);\n }\n done = (*p_status == 0);\n cudaCheckError();\n }\n\n cuda_rc = amgx::memory::hipHostFree(h_status);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipHostFree failed for h_status in low_deg_coarse_A_generator\", AMGX_ERR_CUDA_FAILURE);\n }\n cuda_rc = amgx::memory::hipHostFree(h_work_offset);\n if (cuda_rc != hipSuccess)\n {\n FatalError(\"hipHostFree failed for h_work_offset in low_deg_coarse_A_generator\", AMGX_ERR_CUDA_FAILURE);\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int CTA_SIZE, bool HAS_DIAG, typename Workspace, typename Value_type >\nstatic\nvoid fill_A_dispatch( Workspace &hash_wk,\n const int block_size,\n const int R_num_rows, // same as num_aggregates.\n const int *R_rows,\n const int *R_cols,\n const int *A_rows,\n const int *A_cols,\n const int *A_diag,\n const Value_type *A_vals,\n const int *aggregates,\n const int *Ac_rows,\n const int *Ac_cols,\n const int *Ac_pos,\n const int *Ac_diag,\n Value_type *Ac_vals,\n bool force_determinism )\n{\n hipDeviceProp_t props = getDeviceProperties();\n int GRID_SIZE = (props.major >= 7) ? 1024 : 128;\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int work_offset = GRID_SIZE * NUM_WARPS;\n hipMemcpyAsync( hash_wk.get_work_queue(), &work_offset, sizeof(int), hipMemcpyHostToDevice, amgx::thrust::global_thread_handle::get_stream() );\n cudaCheckError();\n\n // Launch the kernel.\n switch ( block_size )\n {\n case 1:\n fill_A_kernel_1x1 <<< GRID_SIZE, CTA_SIZE>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n hash_wk.get_vals(),\n hash_wk.get_work_queue() );\n cudaCheckError();\n break;\n\n case 2:\n fill_A_kernel_NxN <<< GRID_SIZE, CTA_SIZE>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n reinterpret_cast( hash_wk.get_vals() ),\n hash_wk.get_work_queue() );\n cudaCheckError();\n break;\n\n case 3:\n fill_A_kernel_NxN <<< GRID_SIZE, CTA_SIZE>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n reinterpret_cast( hash_wk.get_vals() ),\n hash_wk.get_work_queue() );\n cudaCheckError();\n break;\n\n case 4:\n if ( force_determinism )\n fill_A_kernel_NxN <<< GRID_SIZE, CTA_SIZE>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n reinterpret_cast( hash_wk.get_vals() ),\n hash_wk.get_work_queue() );\n else\n fill_A_kernel_4x4 <<< GRID_SIZE, CTA_SIZE, 0, amgx::thrust::global_thread_handle::get_stream()>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n reinterpret_cast( hash_wk.get_vals() ),\n hash_wk.get_work_queue() );\n \n cudaCheckError();\n break;\n\n case 5:\n fill_A_kernel_NxN <<< GRID_SIZE, CTA_SIZE>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n reinterpret_cast( hash_wk.get_vals() ),\n hash_wk.get_work_queue() );\n cudaCheckError();\n break;\n\n case 8:\n fill_A_kernel_NxN_large <<< GRID_SIZE, CTA_SIZE>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n reinterpret_cast( hash_wk.get_vals() ),\n hash_wk.get_work_queue() );\n cudaCheckError();\n break;\n\n case 10:\n fill_A_kernel_NxN_large <<< GRID_SIZE, CTA_SIZE>>>(\n R_num_rows,\n R_rows,\n R_cols,\n A_rows,\n A_cols,\n A_diag,\n A_vals,\n aggregates,\n Ac_rows,\n Ac_cols,\n Ac_pos,\n Ac_diag,\n Ac_vals,\n hash_wk.get_gmem_size(),\n hash_wk.get_keys(),\n reinterpret_cast( hash_wk.get_vals() ),\n hash_wk.get_work_queue() );\n cudaCheckError();\n break;\n\n default:\n FatalError( \"LOW_DEG not implemented for this block size\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE );\n }\n\n cudaCheckError();\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nLowDegCoarseAGenerator >::computeAOperator( const Matrix_d &A,\n Matrix_d &Ac,\n const IVector &aggregates,\n const IVector &R_row_offsets,\n const IVector &R_column_indices,\n const int num_aggregates )\n{\n if ( A.get_block_dimx() != A.get_block_dimy() )\n {\n FatalError( \"LowDegCoarseAGenerator implemented for squared blocks\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE );\n }\n\n // The matrix Ac will be modified.\n Ac.set_initialized(0);\n // Is the diagonal stored separately??\n const int diag_prop = A.hasProps(DIAG);\n // Allocate a workspace for hashing.\n typedef TemplateConfig TConfig_d;\n\n hipDeviceProp_t props = getDeviceProperties();\n int grid_size = (props.major >= 7) ? 1024 : 128;\n\n Hash_Workspace hash_wk(true, grid_size);\n\n // Compute row offsets of Ac.\n Ac.addProps(CSR);\n Ac.set_num_rows( num_aggregates );\n Ac.set_num_cols( num_aggregates );\n Ac.row_offsets.resize( num_aggregates + 1 );\n // Compute the number of non-zero elements per row of Ac.\n const int CTA_SIZE = 128;\n\n if ( diag_prop )\n compute_sparsity_dispatch(\n hash_wk,\n num_aggregates,\n R_row_offsets.raw(),\n R_column_indices.raw(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n aggregates.raw(),\n Ac.row_offsets.raw(),\n NULL,\n NULL );\n else\n compute_sparsity_dispatch(\n hash_wk,\n num_aggregates,\n R_row_offsets.raw(),\n R_column_indices.raw(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n aggregates.raw(),\n Ac.row_offsets.raw(),\n NULL,\n NULL );\n\n cudaCheckError();\n // Compute the number of non-zeroes.\n thrust_wrapper::exclusive_scan( Ac.row_offsets.begin(), Ac.row_offsets.end(), Ac.row_offsets.begin() );\n cudaCheckError();\n int nonzero_blocks = Ac.row_offsets[num_aggregates];\n\n if ( diag_prop )\n {\n Ac.addProps(DIAG);\n }\n\n if ( A.is_matrix_singleGPU() )\n {\n Ac.resize( num_aggregates, num_aggregates, nonzero_blocks, A.get_block_dimy(), A.get_block_dimx(), !diag_prop );\n }\n else\n {\n //have 3% more nz for storage\n Ac.resize_spare( num_aggregates, num_aggregates, nonzero_blocks, A.get_block_dimy(), A.get_block_dimx(), 1.0 );\n\n if ( diag_prop )\n {\n Ac.computeDiagonal();\n }\n }\n\n // Vector to store the positions in the hash table.\n device_vector_alloc Ac_pos(nonzero_blocks);\n\n // Compute the sparsity pattern of the rows of Ac.\n if ( diag_prop )\n compute_sparsity_dispatch(\n hash_wk,\n num_aggregates,\n R_row_offsets.raw(),\n R_column_indices.raw(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n aggregates.raw(),\n Ac.row_offsets.raw(),\n Ac.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &Ac_pos.front() ));\n else\n compute_sparsity_dispatch(\n hash_wk,\n num_aggregates,\n R_row_offsets.raw(),\n R_column_indices.raw(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n aggregates.raw(),\n Ac.row_offsets.raw(),\n Ac.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &Ac_pos.front() ));\n\n cudaCheckError();\n\n // Reset values if needed.\n if ( A.get_block_dimy() != 1 )\n {\n thrust_wrapper::fill( Ac.values.begin(), Ac.values.end(), types::util::get_zero() );\n cudaCheckError();\n }\n\n // Compute values.\n if ( diag_prop )\n {\n fill_A_dispatch(\n hash_wk,\n A.get_block_dimy(),\n num_aggregates,\n R_row_offsets.raw(),\n R_column_indices.raw(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n aggregates.raw(),\n Ac.row_offsets.raw(),\n Ac.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &Ac_pos.front() ),\n Ac.diag.raw(),\n Ac.values.raw(),\n this->m_force_determinism );\n }\n else\n {\n fill_A_dispatch(\n hash_wk,\n A.get_block_dimy(),\n num_aggregates,\n R_row_offsets.raw(),\n R_column_indices.raw(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.diag.raw(),\n A.values.raw(),\n aggregates.raw(),\n Ac.row_offsets.raw(),\n Ac.col_indices.raw(),\n amgx::thrust::raw_pointer_cast( &Ac_pos.front() ),\n Ac.diag.raw(),\n Ac.values.raw(),\n this->m_force_determinism );\n }\n\n cudaCheckError();\n\n // Update the diagonal if needed.\n if ( Ac.is_matrix_singleGPU() )\n {\n Ac.computeDiagonal();\n }\n\n cudaCheckError();\n // Finalize the modification.\n Ac.set_initialized(1);\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid\nLowDegCoarseAGenerator >::computeAOperator( const Matrix_h &h_A,\n Matrix_h &h_Ac,\n const IVector &h_aggregates,\n const IVector &h_R_row_offsets,\n const IVector &h_R_column_indices,\n const int num_aggregates )\n{\n h_Ac.set_initialized(0);\n IVector rows;\n IVector inds;\n typename Matrix_h::MVector vals;\n typename Matrix_h::MVector diag;\n int num_nnz = 0;\n int diag_prop = h_A.hasProps(DIAG);\n\n for ( int row = 0; row < num_aggregates; row++ )\n {\n for ( int col = 0; col < num_aggregates; col++ )\n {\n int fill = 0;\n typename Matrix_h::MVector cur(h_A.get_block_size(), types::util::get_zero());\n\n for ( int rc = h_R_row_offsets[row]; rc < h_R_row_offsets[row + 1]; rc++ )\n {\n int j = h_R_column_indices[rc];\n\n for ( int ac = h_A.row_offsets[j]; ac < h_A.row_offsets[j + 1] + diag_prop; ac++ )\n {\n int k = (ac == h_A.row_offsets[j + 1]) ? j : h_A.col_indices[ac];\n\n for ( int q = h_R_row_offsets[col]; q < h_R_row_offsets[col + 1]; q++ )\n if ( k == h_R_column_indices[q] )\n {\n fill = 1;\n int val_idx = (ac == h_A.row_offsets[j + 1]) ? h_A.get_num_nz() + j : ac;\n\n for ( int v = 0; v < h_A.get_block_size(); v++)\n {\n cur[v] = cur[v] + h_A.values[val_idx * h_A.get_block_size() + v];\n }\n }\n }\n }\n\n if ( fill )\n {\n if ( row != col || !diag_prop )\n {\n inds.push_back(col);\n rows.push_back(row);\n num_nnz++;\n\n for ( int v = 0; v < h_A.get_block_size(); v++ )\n {\n vals.push_back(cur[v]);\n }\n }\n else\n {\n for ( int v = 0; v < h_A.get_block_size(); v++ )\n {\n diag.push_back(cur[v]);\n }\n }\n }\n }\n }\n\n rows.push_back(-1);\n\n // add diagonal to the end\n if ( diag_prop )\n {\n for ( int v = 0; v < num_aggregates * h_A.get_block_size(); v++ )\n {\n vals.push_back(diag[v]);\n }\n }\n else\n {\n // Add a zero at the end\n for (int v = 0; v < h_A.get_block_size(); v++)\n {\n vals.push_back(types::util::get_zero());\n }\n }\n\n h_Ac.resize(num_aggregates, num_aggregates, num_nnz, h_A.get_block_dimx(), h_A.get_block_dimy(), 1);\n h_Ac.row_indices = rows;\n h_Ac.col_indices = inds;\n h_Ac.values = vals;\n h_Ac.addProps( CSR | ( diag_prop ? DIAG : 0 ) );\n h_Ac.computeDiagonal();\n h_Ac.set_initialized(1);\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define AMGX_CASE_LINE(CASE) template class LowDegCoarseAGenerator::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace aggregation\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/aggregation/coarseAgenerators/low_deg_coarse_A_generator.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#define COARSE_CLA_CONSO 0\n\n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include // for minmax_element\n\n#include \n#include \n#include \n\n#include \n\n#include \n#include \n#include \n\n#include \nnamespace amgx\n{\n\nnamespace classical\n{\n\nstruct is_zero\n{\n __host__ __device__\n bool operator()(const double &v)\n {\n return fabs(v) < 1e-10;\n }\n};\n\n#define AMGX_CAL_BLOCK_SIZE 256\n\n/* There might be a situation where not all local_to_global_map columns are present in the matrix (because some rows were removed\n and the columns in these rows are therefore no longer present. This kernel creates the flags array that marks existing columns. */\ntemplate\n__global__ __launch_bounds__( AMGX_CAL_BLOCK_SIZE )\nvoid flag_existing_local_to_global_columns(ind_t n, ind_t *row_offsets, ind_t *col_indices, ind_t *flags)\n{\n ind_t i, j, s, e, col;\n\n //go through the matrix\n for (i = threadIdx.x + blockIdx.x * blockDim.x; i < n; i += blockDim.x * gridDim.x)\n {\n s = row_offsets[i];\n e = row_offsets[i + 1];\n\n for (j = s; j < e; j++)\n {\n col = col_indices[j];\n\n //flag columns outside of the square part (which correspond to local_to_global_map)\n if (col >= n)\n {\n flags[col - n] = 1;\n }\n }\n }\n}\n\n/* Renumber the indices based on the prefix-scan/sum of the flags array */\ntemplate\n__global__ __launch_bounds__( AMGX_CAL_BLOCK_SIZE )\nvoid compress_existing_local_columns(ind_t n, ind_t *row_offsets, ind_t *col_indices, ind_t *flags)\n{\n ind_t i, j, s, e, col;\n\n //go through the matrix\n for (i = threadIdx.x + blockIdx.x * blockDim.x; i < n; i += blockDim.x * gridDim.x)\n {\n s = row_offsets[i];\n e = row_offsets[i + 1];\n\n for (j = s; j < e; j++)\n {\n col = col_indices[j];\n\n //flag columns outside of the square part (which correspond to local_to_global_map)\n if (col >= n)\n {\n col_indices[j] = n + flags[col - n];\n }\n }\n }\n}\n\n/* compress the local to global columns indices based on the prefix-scan/sum of the flags array */\ntemplate\n__global__ __launch_bounds__( AMGX_CAL_BLOCK_SIZE )\nvoid compress_existing_local_to_global_columns(ind_t n, ind64_t *l2g_in, ind64_t *l2g_out, ind_t *flags)\n{\n ind_t i;\n\n //go through the arrays (and copy the updated indices when needed)\n for (i = threadIdx.x + blockIdx.x * blockDim.x; i < n; i += blockDim.x * gridDim.x)\n {\n if (flags[i] != flags[i + 1])\n {\n l2g_out[flags[i]] = l2g_in[i];\n }\n }\n}\n\n\ntemplate \nSelector *chooseAggressiveSelector(AMG_Config *m_cfg, std::string std_scope)\n{\n AMG_Config cfg;\n std::string cfg_string(\"\");\n cfg_string += \"default:\";\n // if necessary, allocate aggressive selector + interpolator\n bool use_pmis = false, use_hmis = false;\n // default argument - use the same selector as normal coarsening\n std::string agg_selector = m_cfg->AMG_Config::template getParameter(\"aggressive_selector\", std_scope);\n\n if (agg_selector == \"DEFAULT\")\n {\n std::string std_selector = m_cfg->AMG_Config::template getParameter(\"selector\", std_scope);\n\n if (std_selector == \"PMIS\") { cfg_string += \"selector=AGGRESSIVE_PMIS\"; use_pmis = true; }\n else if (std_selector == \"HMIS\") { cfg_string += \"selector=AGGRESSIVE_HMIS\"; use_hmis = true; }\n else\n {\n FatalError(\"Must use either PMIS or HMIS algorithms with aggressive coarsening\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n // otherwise use specified selector\n else if (agg_selector == \"PMIS\") { cfg_string += \"selector=AGGRESSIVE_PMIS\"; use_pmis = true; }\n else if (agg_selector == \"HMIS\") { cfg_string += \"selector=AGGRESSIVE_HMIS\"; use_hmis = true; }\n else\n {\n FatalError(\"Invalid aggressive coarsener selected\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // check a selector has been selected\n if (!use_pmis && !use_hmis)\n {\n FatalError(\"No aggressive selector chosen\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n cfg.parseParameterString(cfg_string.c_str());\n // now allocate the selector and interpolator\n return classical::SelectorFactory::allocate(cfg, \"default\" /*std_scope*/);\n}\n\ntemplate \nInterpolator *chooseAggressiveInterpolator(AMG_Config *m_cfg, std::string std_scope)\n{\n // temporary config and pointer to main config\n AMG_Config cfg(*m_cfg);\n std::string cfg_string(\"default:\");\n // Set the interpolator\n cfg_string += \"interpolator=\";\n cfg_string += m_cfg->AMG_Config::template getParameter(\"aggressive_interpolator\", std_scope);\n cfg.parseParameterString(cfg_string.c_str());\n // now allocate the selector and interpolator\n return InterpolatorFactory::allocate(cfg, \"default\");\n}\n\ntemplate \nClassical_AMG_Level_Base::Classical_AMG_Level_Base(AMG_Class *amg) : AMG_Level(amg)\n{\n strength = StrengthFactory::allocate(*(amg->m_cfg), amg->m_cfg_scope);\n selector = classical::SelectorFactory::allocate(*(amg->m_cfg), amg->m_cfg_scope);\n interpolator = InterpolatorFactory::allocate(*(amg->m_cfg), amg->m_cfg_scope);\n trunc_factor = amg->m_cfg->AMG_Config::template getParameter(\"interp_truncation_factor\", amg->m_cfg_scope);\n max_elmts = amg->m_cfg->AMG_Config::template getParameter(\"interp_max_elements\", amg->m_cfg_scope);\n max_row_sum = amg->m_cfg->AMG_Config::template getParameter(\"max_row_sum\", amg->m_cfg_scope);\n num_aggressive_levels = amg->m_cfg->AMG_Config::template getParameter(\"aggressive_levels\", amg->m_cfg_scope);\n}\n\ntemplate \nClassical_AMG_Level_Base::~Classical_AMG_Level_Base()\n{\n delete strength;\n delete selector;\n delete interpolator;\n}\n\ntemplate \nvoid Classical_AMG_Level_Base::transfer_level(AMG_Level *ref_lvl)\n{\n Classical_AMG_Level_Base *ref_cla_lvl = dynamic_cast*>(ref_lvl);\n this->P.copy(ref_cla_lvl->P);\n this->R.copy(ref_cla_lvl->R);\n this->m_s_con.copy(ref_cla_lvl->m_s_con);\n this->m_scratch.copy(ref_cla_lvl->m_scratch);\n this->m_cf_map.copy(ref_cla_lvl->m_cf_map);\n}\n\n/****************************************\n * Computes the A, P, and R operators\n ***************************************/\ntemplate \nvoid Classical_AMG_Level_Base::createCoarseVertices()\n{\n if (AMG_Level::getLevelIndex() < this->num_aggressive_levels)\n {\n if (selector) { delete selector; }\n\n selector = chooseAggressiveSelector(AMG_Level::amg->m_cfg, AMG_Level::amg->m_cfg_scope);\n }\n\n Matrix &RAP = this->getNextLevel( typename Matrix::memory_space( ) )->getA( );\n Matrix &A = this->getA();\n int size_all, size_full, nnz_full;\n\n if (!A.is_matrix_singleGPU())\n {\n int offset;\n // Need to get number of 2-ring rows\n A.getOffsetAndSizeForView(ALL, &offset, &size_all);\n A.getOffsetAndSizeForView(FULL, &offset, &size_full);\n A.getNnzForView(FULL, &nnz_full);\n }\n else\n {\n size_all = A.get_num_rows();\n size_full = A.get_num_rows();\n nnz_full = A.get_num_nz();\n }\n\n this->m_cf_map.resize(size_all);\n this->m_s_con.resize(nnz_full);\n this->m_scratch.resize(size_full);\n thrust_wrapper::fill(this->m_cf_map.begin(), this->m_cf_map.end(), 0);\n cudaCheckError();\n thrust_wrapper::fill(this->m_s_con.begin(), this->m_s_con.end(), false);\n cudaCheckError();\n thrust_wrapper::fill(this->m_scratch.begin(), this->m_scratch.end(), 0);\n cudaCheckError();\n markCoarseFinePoints();\n}\n\ntemplate \nvoid Classical_AMG_Level_Base::createCoarseMatrices()\n{\n // allocate aggressive interpolator if needed\n if (AMG_Level::getLevelIndex() < this->num_aggressive_levels)\n {\n if (interpolator) { delete interpolator; }\n\n interpolator = chooseAggressiveInterpolator(AMG_Level::amg->m_cfg, AMG_Level::amg->m_cfg_scope);\n }\n\n Matrix &RAP = this->getNextLevel( typename Matrix::memory_space( ) )->getA( );\n Matrix &A = this->getA();\n /* WARNING: exit if D1 interpolator is selected in distributed setting */\n std::string s(\"\");\n s += AMG_Level::amg->m_cfg->AMG_Config::template getParameter(\"interpolator\", AMG_Level::amg->m_cfg_scope);\n\n if (A.is_matrix_distributed() && (s.compare(\"D1\") == 0))\n {\n FatalError(\"D1 interpolation is not supported in distributed settings\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n /* WARNING: do not recompute prolongation (P) and restriction (R) when you\n are reusing the level structure (structure_reuse_levels > 0) */\n if (this->isReuseLevel() == false)\n {\n computeProlongationOperator();\n }\n\n // Compute Restriction operator and coarse matrix Ac\n if (!this->A->is_matrix_distributed() || this->A->manager->get_num_partitions() == 1)\n {\n /* WARNING: see above warning. */\n if (this->isReuseLevel() == false)\n {\n computeRestrictionOperator();\n }\n\n computeAOperator();\n }\n else\n {\n /* WARNING: notice that in this case the computeRestructionOperator() is called\n inside computeAOperator_distributed() routine. */\n computeAOperator_distributed();\n }\n\n// we also need to renumber columns of P and rows or R correspondingly since we changed RAP halo columns\n// for R we just keep track of renumbering in and exchange proper vectors in restriction\n// for P we actually need to renumber columns for prolongation:\n if (A.is_matrix_distributed() && this->A->manager->get_num_partitions() > 1)\n {\n RAP.set_initialized(0);\n // Renumber the owned nodes as interior and boundary (renumber rows and columns)\n // We are passing reuse flag to not create neighbours list from scratch, but rather update based on new halos\n RAP.manager->renumberMatrixOneRing(this->isReuseLevel());\n // Renumber the column indices of P and shuffle rows of P\n RAP.manager->renumber_P_R(this->P, this->R, A);\n // Create the B2L_maps for RAP\n RAP.manager->createOneRingHaloRows();\n RAP.manager->getComms()->set_neighbors(RAP.manager->num_neighbors());\n RAP.setView(OWNED);\n RAP.set_initialized(1);\n // update # of columns in P - this is necessary for correct CSR multiply\n P.set_initialized(0);\n int new_num_cols = thrust_wrapper::reduce(P.col_indices.begin(), P.col_indices.end(), int(0), amgx::thrust::maximum()) + 1;\n cudaCheckError();\n P.set_num_cols(new_num_cols);\n P.set_initialized(1);\n }\n\n RAP.copyAuxData(&A);\n\n if (!A.is_matrix_singleGPU() && RAP.manager == NULL)\n {\n RAP.manager = new DistributedManager();\n }\n\n if (this->getA().is_matrix_singleGPU())\n {\n this->m_next_level_size = this->getNextLevel(typename Matrix::memory_space() )->getA().get_num_rows() * this->getNextLevel(typename Matrix::memory_space() )->getA().get_block_dimy();\n }\n else\n {\n // m_next_level_size is the size that will be used to allocate xc, bc vectors\n int size, offset;\n this->getNextLevel(typename Matrix::memory_space())->getA().getOffsetAndSizeForView(FULL, &offset, &size);\n this->m_next_level_size = size * this->getNextLevel(typename Matrix::memory_space() )->getA().get_block_dimy();\n }\n}\n\ntemplate \nvoid Classical_AMG_Level_Base::markCoarseFinePoints()\n{\n Matrix &A = this->getA();\n //allocate necessary memory\n typedef Vector::Type> IVector;\n typedef Vector::Type> BVector;\n typedef Vector::Type> FVector;\n FVector weights;\n\n if (!A.is_matrix_singleGPU())\n {\n int size, offset;\n A.getOffsetAndSizeForView(FULL, &offset, &size);\n // size should now contain the number of 1-ring rows\n weights.resize(size);\n }\n else\n {\n weights.resize(A.get_num_rows());\n }\n\n thrust_wrapper::fill(weights.begin(), weights.end(), 0.0);\n cudaCheckError();\n\n // extend A to include 1st ring nodes\n // compute strong connections and weights\n if (!A.is_matrix_singleGPU())\n {\n ViewType oldView = A.currentView();\n A.setView(FULL);\n strength->computeStrongConnectionsAndWeights(A, this->m_s_con, weights, this->max_row_sum);\n A.setView(oldView);\n }\n else\n {\n strength->computeStrongConnectionsAndWeights(A, this->m_s_con, weights, this->max_row_sum);\n }\n\n // Exchange the one-ring of the weights\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo(weights, weights.tag);\n }\n\n //mark coarse and fine points\n selector->markCoarseFinePoints(A, weights, this->m_s_con, this->m_cf_map, this->m_scratch);\n // we do resize cf_map to zero later, so we are saving separate copy\n this->m_cf_map.dirtybit = 1;\n\n // Do a two ring exchange of cf_map\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_2ring(this->m_cf_map, m_cf_map.tag);\n }\n\n // Modify cf_map array such that coarse points are assigned a local index, while fine points entries are not touched\n selector->renumberAndCountCoarsePoints(this->m_cf_map, this->m_num_coarse_vertices, A.get_num_rows());\n}\n\n\ntemplate \nvoid Classical_AMG_Level_Base::computeProlongationOperator()\n{\n this->Profile.tic(\"computeP\");\n Matrix &A = this->getA();\n //allocate necessary memory\n typedef Vector::Type> IVector;\n typedef Vector::Type> BVector;\n typedef Vector::Type> FVector;\n //generate the interpolation matrix\n interpolator->generateInterpolationMatrix(A, this->m_cf_map, this->m_s_con, this->m_scratch, P);\n this->m_cf_map.clear();\n this->m_cf_map.shrink_to_fit();\n this->m_scratch.clear();\n this->m_scratch.shrink_to_fit();\n this->m_s_con.clear();\n this->m_s_con.shrink_to_fit();\n\n // truncate based on max # of elements if desired\n if (this->max_elmts > 0 && P.get_num_rows() > 0)\n {\n Truncate::truncateByMaxElements(P, this->max_elmts);\n }\n\n if (!P.isLatencyHidingEnabled(*this->amg->m_cfg))\n {\n // This will cause bsrmv_with_mask to not do latency hiding\n P.setInteriorView(OWNED);\n P.setExteriorView(OWNED);\n }\n}\n\n/**********************************************\n * computes R=P^T\n **********************************************/\ntemplate \nvoid Classical_AMG_Level_Base::computeRestrictionOperator()\n{\n this->Profile.tic(\"computeR\");\n R.set_initialized(0);\n P.setView(OWNED);\n transpose(P, R, P.get_num_rows());\n\n if (!R.isLatencyHidingEnabled(*this->amg->m_cfg))\n {\n // This will cause bsrmv_with_mask_restriction to not do latency hiding\n R.setInteriorView(OWNED);\n R.setExteriorView(OWNED);\n }\n\n if(P.is_matrix_distributed())\n {\n // Setup the number of non-zeros in R using stub DistributedManager\n R.manager = new DistributedManager();\n int nrows_owned = P.manager->halo_offsets[0];\n int nrows_full = P.manager->halo_offsets[P.manager->neighbors.size()];\n int nz_full = R.row_offsets[nrows_full];\n int nz_owned = R.row_offsets[nrows_owned];\n R.manager->setViewSizes(nrows_owned, nz_owned, nrows_owned, nz_owned, nrows_full, nz_full, R.get_num_rows(), R.get_num_nz());\n }\n\n R.set_initialized(1);\n this->Profile.toc(\"computeR\");\n}\n\n/**********************************************\n * computes the Galerkin product: A_c=R*A*P\n **********************************************/\n\ntemplate \nvoid Classical_AMG_Level >::computeAOperator_1x1()\n{\n this->Profile.tic(\"computeA\");\n Matrix RA;\n RA.addProps(CSR);\n RA.set_block_dimx(this->getA().get_block_dimx());\n RA.set_block_dimy(this->getA().get_block_dimy());\n Matrix &RAP = this->getNextLevel( typename Matrix::memory_space( ) )->getA( );\n RAP.addProps(CSR);\n RAP.set_block_dimx(this->getA().get_block_dimx());\n RAP.set_block_dimy(this->getA().get_block_dimy());\n Matrix &Atmp = this->getA();\n multiplyMM(this->R, this->getA(), RA);\n multiplyMM(RA, this->P, RAP);\n RAP.sortByRowAndColumn();\n RAP.set_initialized(1);\n this->Profile.toc(\"computeA\");\n}\n\ntemplate \nvoid Classical_AMG_Level >::computeAOperator_1x1_distributed()\n{\n FatalError(\"Distributed classical AMG not implemented for host\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n\ntemplate \nvoid Classical_AMG_Level >::computeAOperator_1x1()\n{\n this->Profile.tic(\"computeA\");\n Matrix &RAP = this->getNextLevel( device_memory( ) )->getA( );\n RAP.addProps(CSR);\n RAP.set_block_dimx(this->getA().get_block_dimx());\n RAP.set_block_dimy(this->getA().get_block_dimy());\n this->R.set_initialized( 0 );\n this->R.addProps( CSR );\n this->R.set_initialized( 1 );\n this->P.set_initialized( 0 );\n this->P.addProps( CSR );\n this->P.set_initialized( 1 );\n void *wk = AMG_Level::amg->getCsrWorkspace();\n\n if ( wk == NULL )\n {\n wk = CSR_Multiply::csr_workspace_create( *(AMG_Level::amg->m_cfg), AMG_Level::amg->m_cfg_scope );\n AMG_Level::amg->setCsrWorkspace( wk );\n }\n\n int spmm_verbose = this->amg->m_cfg->AMG_Config::template getParameter(\"spmm_verbose\", this->amg->m_cfg_scope);\n\n if ( spmm_verbose )\n {\n typedef typename Matrix::IVector::const_iterator Iterator;\n typedef amgx::thrust::pair Result;\n std::ostringstream buffer;\n buffer << \"SPMM: Level \" << this->getLevelIndex() << std::endl;\n\n if ( this->getLevelIndex() == 0 )\n {\n device_vector_alloc num_nz( this->getA().row_offsets.size() );\n amgx::thrust::adjacent_difference( this->getA().row_offsets.begin(), this->getA().row_offsets.end(), num_nz.begin() );\n cudaCheckError();\n Result result = amgx::thrust::minmax_element( num_nz.begin() + 1, num_nz.end() );\n cudaCheckError();\n int min_size = *result.first;\n int max_size = *result.second;\n int sum = thrust_wrapper::reduce( num_nz.begin() + 1, num_nz.end() );\n cudaCheckError();\n double avg_size = double(sum) / this->getA().get_num_rows();\n buffer << \"SPMM: A: \" << std::endl;\n buffer << \"SPMM: Matrix avg row size: \" << avg_size << std::endl;\n buffer << \"SPMM: Matrix min row size: \" << min_size << std::endl;\n buffer << \"SPMM: Matrix max row size: \" << max_size << std::endl;\n }\n\n device_vector_alloc num_nz( this->P.row_offsets.size() );\n amgx::thrust::adjacent_difference( this->P.row_offsets.begin(), this->P.row_offsets.end(), num_nz.begin() );\n cudaCheckError();\n Result result = amgx::thrust::minmax_element( num_nz.begin() + 1, num_nz.end() );\n cudaCheckError();\n int min_size = *result.first;\n int max_size = *result.second;\n int sum = thrust_wrapper::reduce( num_nz.begin() + 1, num_nz.end() );\n cudaCheckError();\n double avg_size = double(sum) / this->P.get_num_rows();\n buffer << \"SPMM: P: \" << std::endl;\n buffer << \"SPMM: Matrix avg row size: \" << avg_size << std::endl;\n buffer << \"SPMM: Matrix min row size: \" << min_size << std::endl;\n buffer << \"SPMM: Matrix max row size: \" << max_size << std::endl;\n num_nz.resize( this->R.row_offsets.size() );\n amgx::thrust::adjacent_difference( this->R.row_offsets.begin(), this->R.row_offsets.end(), num_nz.begin() );\n cudaCheckError();\n result = amgx::thrust::minmax_element( num_nz.begin() + 1, num_nz.end() );\n cudaCheckError();\n min_size = *result.first;\n max_size = *result.second;\n sum = thrust_wrapper::reduce( num_nz.begin() + 1, num_nz.end() );\n cudaCheckError();\n avg_size = double(sum) / this->R.get_num_rows();\n buffer << \"SPMM: R: \" << std::endl;\n buffer << \"SPMM: Matrix avg row size: \" << avg_size << std::endl;\n buffer << \"SPMM: Matrix min row size: \" << min_size << std::endl;\n buffer << \"SPMM: Matrix max row size: \" << max_size << std::endl;\n amgx_output( buffer.str().c_str(), static_cast( buffer.str().length() ) );\n }\n\n RAP.set_initialized( 0 );\n CSR_Multiply::csr_galerkin_product( this->R, this->getA(), this->P, RAP, NULL, NULL, NULL, NULL, NULL, NULL, wk );\n RAP.set_initialized( 1 );\n int spmm_no_sort = this->amg->m_cfg->AMG_Config::template getParameter(\"spmm_no_sort\", this->amg->m_cfg_scope);\n this->Profile.toc(\"computeA\");\n}\n/**********************************************\n * computes the restriction: rr=R*r\n **********************************************/\ntemplate \nvoid Classical_AMG_Level_Base::restrictResidual(VVector &r, VVector &rr)\n{\n// we need to resize residual vector to make sure it can store halo rows to be sent\n if (!P.is_matrix_singleGPU())\n {\n typedef typename TConfig::MemSpace MemorySpace;\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n#if COARSE_CLA_CONSO\n int desired_size ;\n\n if (this->getNextLevel(MemorySpace())->isConsolidationLevel())\n {\n desired_size = std::max(P.manager->halo_offsets[P.manager->neighbors.size()], Ac.manager->halo_offsets_before_glue[Ac.manager->neighbors_before_glue.size()] * rr.get_block_size());\n }\n else\n {\n desired_size = std::max(P.manager->halo_offsets[P.manager->neighbors.size()], Ac.manager->halo_offsets[Ac.manager->neighbors.size()] * rr.get_block_size());\n }\n\n#else\n int desired_size = std::max(P.manager->halo_offsets[P.manager->neighbors.size()], Ac.manager->halo_offsets[Ac.manager->neighbors.size()] * rr.get_block_size());\n#endif\n rr.resize(desired_size);\n }\n\n#if 1\n this->Profile.tic(\"restrictRes\");\n\n // Disable speculative send of rr\n if (P.is_matrix_singleGPU())\n {\n multiply( R, r, rr);\n }\n else\n {\n multiply_with_mask_restriction( R, r, rr, P);\n }\n\n#endif\n // exchange halo residuals & add residual contribution from neighbors\n rr.dirtybit = 1;\n\n if (!P.is_matrix_singleGPU())\n {\n int desired_size = P.manager->halo_offsets[P.manager->neighbors.size()] * rr.get_block_size();\n\n if (rr.size() < desired_size)\n {\n rr.resize(desired_size);\n }\n }\n\n this->Profile.toc(\"restrictRes\");\n}\n\nstruct is_minus_one\n{\n __host__ __device__\n bool operator()(const int &x)\n {\n return x == -1;\n }\n};\n\n\ntemplate \nvoid Classical_AMG_Level >::computeAOperator_1x1_distributed()\n{\n Matrix &A = this->getA();\n Matrix &P = this->P;\n Matrix &RAP = this->getNextLevel( device_memory( ) )->getA( );\n RAP.addProps(CSR);\n RAP.set_block_dimx(this->getA().get_block_dimx());\n RAP.set_block_dimy(this->getA().get_block_dimy());\n IndexType num_parts = A.manager->get_num_partitions();\n IndexType num_neighbors = A.manager->num_neighbors();\n IndexType my_rank = A.manager->global_id();\n // OWNED includes interior and boundary\n A.setView(OWNED);\n int num_owned_coarse_pts = P.manager->halo_offsets[0];\n int num_owned_fine_pts = A.manager->halo_offsets[0];\n\n // Initialize RAP.manager\n if (RAP.manager == NULL)\n {\n RAP.manager = new DistributedManager();\n }\n\n RAP.manager->A = &RAP;\n RAP.manager->setComms(A.manager->getComms());\n RAP.manager->set_global_id(my_rank);\n RAP.manager->set_num_partitions(num_parts);\n RAP.manager->part_offsets_h = P.manager->part_offsets_h;\n RAP.manager->part_offsets = P.manager->part_offsets;\n RAP.manager->set_base_index(RAP.manager->part_offsets_h[my_rank]);\n RAP.manager->set_index_range(num_owned_coarse_pts);\n RAP.manager->num_rows_global = RAP.manager->part_offsets_h[num_parts];\n // --------------------------------------------------------------------\n // Using the B2L_maps of matrix A, identify the rows of P that need to be sent to neighbors,\n // so that they can compute A*P\n // Once rows of P are identified, convert the column indices to global indices, and send them to neighbors\n // ---------------------------------------------------------------------------\n // Copy some information about the manager of P, since we don't want to modify those\n IVector_h P_neighbors = P.manager->neighbors;\n I64Vector_h P_halo_ranges_h = P.manager->halo_ranges_h;\n I64Vector_d P_halo_ranges = P.manager->halo_ranges;\n RAP.manager->local_to_global_map = P.manager->local_to_global_map;\n IVector_h P_halo_offsets = P.manager->halo_offsets;\n // Create a temporary distributed arranger\n DistributedArranger *prep = new DistributedArranger;\n prep->exchange_halo_rows_P(A, this->P, RAP.manager->local_to_global_map, P_neighbors, P_halo_ranges_h, P_halo_ranges, P_halo_offsets, RAP.manager->part_offsets_h, RAP.manager->part_offsets, num_owned_coarse_pts, RAP.manager->part_offsets_h[my_rank]);\n cudaCheckError();\n // At this point, we can compute RAP_full which contains some rows that will need to be sent to neighbors\n // i.e. RAP_full = [ RAP_int ]\n // [ RAP_ext ]\n // RAP is [ RAP_int ] + [RAP_ext_received_from_neighbors]\n // We can reuse the serial galerkin product since R, A and P use local indices\n // TODO: latency hiding (i.e. compute RAP_ext, exchange_matrix_halo, then do RAP_int)\n /* WARNING: do not recompute prolongation (P) and restriction (R) when you\n are reusing the level structure (structure_reuse_levels > 0) */\n /* We force for matrix P to have only owned rows to be seen for the correct galerkin product computation*/\n this->P.set_initialized(0);\n this->P.set_num_rows(num_owned_fine_pts);\n this->P.addProps( CSR );\n this->P.set_initialized(1);\n\n if (this->isReuseLevel() == false)\n {\n this->R.set_initialized( 0 );\n this->R.addProps( CSR );\n // Take the tranpose of P to get R\n // Single-GPU transpose, no mpi exchange\n this->computeRestrictionOperator();\n this->R.set_initialized( 1 );\n }\n\n this->Profile.tic(\"computeA\");\n Matrix RAP_full;\n // Initialize the workspace needed for galerkin product\n void *wk = AMG_Level::amg->getCsrWorkspace();\n\n if ( wk == NULL )\n {\n wk = CSR_Multiply::csr_workspace_create( *(AMG_Level::amg->m_cfg), AMG_Level::amg->m_cfg_scope );\n AMG_Level::amg->setCsrWorkspace( wk );\n }\n\n // Single-GPU RAP, no mpi exchange\n RAP_full.set_initialized( 0 );\n /* WARNING: Since A is reordered (into interior and boundary nodes), while R and P are not reordered,\n you must unreorder A when performing R*A*P product in ordre to obtain the correct result. */\n CSR_Multiply::csr_galerkin_product( this->R, this->getA(), this->P, RAP_full,\n /* permutation for rows of R, A and P */ NULL, NULL /*&(this->getA().manager->renumbering)*/, NULL,\n /* permutation for cols of R, A and P */ NULL, NULL /*&(this->getA().manager->inverse_renumbering)*/, NULL,\n wk );\n RAP_full.set_initialized( 1 );\n this->Profile.toc(\"computeA\");\n // ----------------------------------------------------------------------------------------------\n // Now, send rows of RAP_full requireq by neighbors, received rows from neighbors and create RAP\n // ----------------------------------------------------------------------------------------------\n prep->exchange_RAP_ext(RAP, RAP_full, A, this->P, P_halo_offsets, RAP.manager->local_to_global_map, P_neighbors, P_halo_ranges_h, P_halo_ranges, RAP.manager->part_offsets_h, RAP.manager->part_offsets, num_owned_coarse_pts, RAP.manager->part_offsets_h[my_rank], wk);\n // Delete temporary distributed arranger\n delete prep;\n /* WARNING: The RAP matrix generated at this point contains extra rows (that correspond to rows of R,\n that was obtained by locally transposing P). This rows are ignored by setting the # of matrix rows\n to be smaller, so that they correspond to number of owned coarse nodes. This should be fine, but\n it leaves holes in the matrix as there might be columns that belong to the extra rows that now do not\n belong to the smaller matrix with number of owned coarse nodes rows. The same is trued about the\n local_to_global_map. These two data structures match at this point. However, in the next calls\n local_to_global (exclusively) will be used to geberate B2L_maps (wihtout going through column indices)\n which creates extra elements in the B2L that simply do not exist in the new matrices. I strongly suspect\n this is the reason fore the bug. The below fix simply compresses the matrix so that there are no holes\n in it, or in the local_2_global_map. */\n //mark local_to_global_columns that exist in the owned coarse nodes rows.\n IndexType nrow = RAP.get_num_rows();\n IndexType ncol = RAP.get_num_cols();\n IndexType nl2g = ncol - nrow;\n\n if (nl2g > 0)\n {\n IVector l2g_p(nl2g + 1, 0); //+1 is needed for prefix_sum/exclusive_scan\n I64Vector l2g_t(nl2g, 0);\n IndexType nblocks = (nrow + AMGX_CAL_BLOCK_SIZE - 1) / AMGX_CAL_BLOCK_SIZE;\n\n if (nblocks > 0)\n flag_existing_local_to_global_columns <<< nblocks, AMGX_CAL_BLOCK_SIZE>>>\n (nrow, RAP.row_offsets.raw(), RAP.col_indices.raw(), l2g_p.raw());\n\n cudaCheckError();\n /*\n //slow version of the above kernel\n for(int ii=0; ii=nrow){\n int kk = col-RAP.get_num_rows();\n l2g_p[kk] = 1;\n }\n }\n }\n cudaCheckError();\n */\n //create a pointer map for their location using prefix sum\n thrust_wrapper::exclusive_scan(l2g_p.begin(), l2g_p.end(), l2g_p.begin());\n int new_nl2g = l2g_p[nl2g];\n\n //compress the columns using the pointer map\n if (nblocks > 0)\n compress_existing_local_columns <<< nblocks, AMGX_CAL_BLOCK_SIZE>>>\n (nrow, RAP.row_offsets.raw(), RAP.col_indices.raw(), l2g_p.raw());\n\n cudaCheckError();\n /*\n //slow version of the above kernel\n for(int ii=0; ii=nrow){\n int kk = col-RAP.get_num_rows();\n RAP.col_indices[jj] = nrow+l2g_p[kk];\n }\n }\n }\n cudaCheckError();\n */\n //adjust matrix size (number of columns) accordingly\n RAP.set_initialized(0);\n RAP.set_num_cols(nrow + new_nl2g);\n RAP.set_initialized(1);\n //compress local_to_global_map using the pointer map\n nblocks = (nl2g + AMGX_CAL_BLOCK_SIZE - 1) / AMGX_CAL_BLOCK_SIZE;\n\n if (nblocks > 0)\n compress_existing_local_to_global_columns <<< nblocks, AMGX_CAL_BLOCK_SIZE>>>\n (nl2g, RAP.manager->local_to_global_map.raw(), l2g_t.raw(), l2g_p.raw());\n\n cudaCheckError();\n amgx::thrust::copy(l2g_t.begin(), l2g_t.begin() + new_nl2g, RAP.manager->local_to_global_map.begin());\n cudaCheckError();\n /*\n //slow version of the above kernel (through Thrust)\n for(int ii=0; ii<(l2g_p.size()-1); ii++){\n if (l2g_p[ii] != l2g_p[ii+1]){\n RAP.manager->local_to_global_map[l2g_p[ii]] = RAP.manager->local_to_global_map[ii];\n }\n }\n cudaCheckError();\n */\n //adjust local_to_global_map size accordingly\n RAP.manager->local_to_global_map.resize(new_nl2g);\n }\n}\n\n/**********************************************\n * prolongates the error: x+=P*e\n **********************************************/\ntemplate \nvoid Classical_AMG_Level_Base::prolongateAndApplyCorrection(VVector &e, VVector &bc, VVector &x, VVector &tmp)\n{\n this->Profile.tic(\"proCorr\");\n // Use P.manager to exchange halo of e before doing P\n // (since P has columns belonging to one of P.neighbors)\n e.dirtybit = 1;\n\n if (!P.is_matrix_singleGPU())\n {\n // get coarse matrix\n typedef typename TConfig::MemSpace MemorySpace;\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n#if COARSE_CLA_CONSO\n int e_size;\n\n if (this->getNextLevel(MemorySpace())->isConsolidationLevel())\n {\n e_size = std::max(P.manager->halo_offsets[P.manager->neighbors.size()], Ac.manager->halo_offsets_before_glue[Ac.manager->neighbors_before_glue.size()]) * e.get_block_size();\n }\n else\n {\n e_size = std::max(P.manager->halo_offsets[P.manager->neighbors.size()], Ac.manager->halo_offsets[Ac.manager->neighbors.size()]) * e.get_block_size();\n }\n\n if (e.size() < e_size) { e.resize(e_size); }\n\n#else\n int e_size = std::max(P.manager->halo_offsets[P.manager->neighbors.size()], Ac.manager->halo_offsets[Ac.manager->neighbors.size()]) * e.get_block_size();\n e.resize(e_size);\n#endif\n }\n\n if (P.is_matrix_singleGPU())\n {\n if (e.size() > 0)\n {\n multiply( P, e, tmp);\n }\n }\n else\n {\n multiply_with_mask( P, e, tmp);\n }\n\n // get owned num rows for fine matrix\n int owned_size;\n\n if (this->A->is_matrix_distributed())\n {\n int owned_offset;\n P.manager->getOffsetAndSizeForView(OWNED, &owned_offset, &owned_size);\n }\n else\n {\n owned_size = x.size();\n }\n\n //apply\n axpby(x, tmp, x, ValueType(1), ValueType(1), 0, owned_size);\n this->Profile.toc(\"proCorr\");\n x.dirtybit = 1;\n}\n\ntemplate \nvoid Classical_AMG_Level_Base::computeAOperator()\n{\n if (this->A->get_block_size() == 1)\n {\n computeAOperator_1x1();\n }\n else\n {\n FatalError(\"Classical AMG not implemented for block_size != 1\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid Classical_AMG_Level_Base::computeAOperator_distributed()\n{\n if (this->A->get_block_size() == 1)\n {\n computeAOperator_1x1_distributed();\n }\n else\n {\n FatalError(\"Classical AMG not implemented for block_size != 1\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n\ntemplate \nvoid Classical_AMG_Level_Base::consolidateVector(VVector &x)\n{\n#ifdef AMGX_WITH_MPI\n#if COARSE_CLA_CONSO\n typedef typename TConfig::MemSpace MemorySpace;\n Matrix &A = this->getA();\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n MPI_Comm comm, temp_com;\n comm = Ac.manager->getComms()->get_mpi_comm();\n temp_com = compute_glue_matrices_communicator(Ac);\n glue_vector(Ac, comm, x, temp_com);\n#endif\n#endif\n}\n\ntemplate \nvoid Classical_AMG_Level_Base::unconsolidateVector(VVector &x)\n{\n#ifdef AMGX_WITH_MPI\n#if COARSE_CLA_CONSO\n typedef typename TConfig::MemSpace MemorySpace;\n Matrix &A = this->getA();\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n MPI_Comm comm, temp_com;\n comm = Ac.manager->getComms()->get_mpi_comm();\n temp_com = compute_glue_matrices_communicator(Ac);\n unglue_vector(Ac, comm, x, temp_com, x);\n#endif\n#endif\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class Classical_AMG_Level_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Classical_AMG_Level::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace classical\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#define COARSE_CLA_CONSO 0\n\n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include // for minmax_element\n\n#include \n#include \n#include \n\n#include \n\n#include \n#include \n#include \n\n#include \nnamespace amgx\n{\n\nnamespace classical\n{\n\nstruct is_zero\n{\n __host__ __device__\n bool operator()(const double &v)\n {\n return fabs(v) < 1e-10;\n }\n};\n\n#define AMGX_CAL_BLOCK_SIZE 256\n\n/* There might be a situation where not all local_to_global_map columns are present in the matrix (because some rows were removed\n and the columns in these rows are therefore no longer present. This kernel creates the flags array that marks existing columns. */\ntemplate\n__global__ __launch_bounds__( AMGX_CAL_BLOCK_SIZE )\nvoid flag_existing_local_to_global_columns(ind_t n, ind_t *row_offsets, ind_t *col_indices, ind_t *flags)\n{\n ind_t i, j, s, e, col;\n\n //go through the matrix\n for (i = threadIdx.x + blockIdx.x * blockDim.x; i < n; i += blockDim.x * gridDim.x)\n {\n s = row_offsets[i];\n e = row_offsets[i + 1];\n\n for (j = s; j < e; j++)\n {\n col = col_indices[j];\n\n //flag columns outside of the square part (which correspond to local_to_global_map)\n if (col >= n)\n {\n flags[col - n] = 1;\n }\n }\n }\n}\n\n/* Renumber the indices based on the prefix-scan/sum of the flags array */\ntemplate\n__global__ __launch_bounds__( AMGX_CAL_BLOCK_SIZE )\nvoid compress_existing_local_columns(ind_t n, ind_t *row_offsets, ind_t *col_indices, ind_t *flags)\n{\n ind_t i, j, s, e, col;\n\n //go through the matrix\n for (i = threadIdx.x + blockIdx.x * blockDim.x; i < n; i += blockDim.x * gridDim.x)\n {\n s = row_offsets[i];\n e = row_offsets[i + 1];\n\n for (j = s; j < e; j++)\n {\n col = col_indices[j];\n\n //flag columns outside of the square part (which correspond to local_to_global_map)\n if (col >= n)\n {\n col_indices[j] = n + flags[col - n];\n }\n }\n }\n}\n\n/* compress the local to global columns indices based on the prefix-scan/sum of the flags array */\ntemplate\n__global__ __launch_bounds__( AMGX_CAL_BLOCK_SIZE )\nvoid compress_existing_local_to_global_columns(ind_t n, ind64_t *l2g_in, ind64_t *l2g_out, ind_t *flags)\n{\n ind_t i;\n\n //go through the arrays (and copy the updated indices when needed)\n for (i = threadIdx.x + blockIdx.x * blockDim.x; i < n; i += blockDim.x * gridDim.x)\n {\n if (flags[i] != flags[i + 1])\n {\n l2g_out[flags[i]] = l2g_in[i];\n }\n }\n}\n\n\ntemplate \nSelector *chooseAggressiveSelector(AMG_Config *m_cfg, std::string std_scope)\n{\n AMG_Config cfg;\n std::string cfg_string(\"\");\n cfg_string += \"default:\";\n // if necessary, allocate aggressive selector + interpolator\n bool use_pmis = false, use_hmis = false;\n // default argument - use the same selector as normal coarsening\n std::string agg_selector = m_cfg->AMG_Config::template getParameter(\"aggressive_selector\", std_scope);\n\n if (agg_selector == \"DEFAULT\")\n {\n std::string std_selector = m_cfg->AMG_Config::template getParameter(\"selector\", std_scope);\n\n if (std_selector == \"PMIS\") { cfg_string += \"selector=AGGRESSIVE_PMIS\"; use_pmis = true; }\n else if (std_selector == \"HMIS\") { cfg_string += \"selector=AGGRESSIVE_HMIS\"; use_hmis = true; }\n else\n {\n FatalError(\"Must use either PMIS or HMIS algorithms with aggressive coarsening\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n // otherwise use specified selector\n else if (agg_selector == \"PMIS\") { cfg_string += \"selector=AGGRESSIVE_PMIS\"; use_pmis = true; }\n else if (agg_selector == \"HMIS\") { cfg_string += \"selector=AGGRESSIVE_HMIS\"; use_hmis = true; }\n else\n {\n FatalError(\"Invalid aggressive coarsener selected\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n // check a selector has been selected\n if (!use_pmis && !use_hmis)\n {\n FatalError(\"No aggressive selector chosen\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n cfg.parseParameterString(cfg_string.c_str());\n // now allocate the selector and interpolator\n return classical::SelectorFactory::allocate(cfg, \"default\" /*std_scope*/);\n}\n\ntemplate \nInterpolator *chooseAggressiveInterpolator(AMG_Config *m_cfg, std::string std_scope)\n{\n // temporary config and pointer to main config\n AMG_Config cfg(*m_cfg);\n std::string cfg_string(\"default:\");\n // Set the interpolator\n cfg_string += \"interpolator=\";\n cfg_string += m_cfg->AMG_Config::template getParameter(\"aggressive_interpolator\", std_scope);\n cfg.parseParameterString(cfg_string.c_str());\n // now allocate the selector and interpolator\n return InterpolatorFactory::allocate(cfg, \"default\");\n}\n\ntemplate \nClassical_AMG_Level_Base::Classical_AMG_Level_Base(AMG_Class *amg) : AMG_Level(amg)\n{\n strength = StrengthFactory::allocate(*(amg->m_cfg), amg->m_cfg_scope);\n selector = classical::SelectorFactory::allocate(*(amg->m_cfg), amg->m_cfg_scope);\n interpolator = InterpolatorFactory::allocate(*(amg->m_cfg), amg->m_cfg_scope);\n trunc_factor = amg->m_cfg->AMG_Config::template getParameter(\"interp_truncation_factor\", amg->m_cfg_scope);\n max_elmts = amg->m_cfg->AMG_Config::template getParameter(\"interp_max_elements\", amg->m_cfg_scope);\n max_row_sum = amg->m_cfg->AMG_Config::template getParameter(\"max_row_sum\", amg->m_cfg_scope);\n num_aggressive_levels = amg->m_cfg->AMG_Config::template getParameter(\"aggressive_levels\", amg->m_cfg_scope);\n}\n\ntemplate \nClassical_AMG_Level_Base::~Classical_AMG_Level_Base()\n{\n delete strength;\n delete selector;\n delete interpolator;\n}\n\ntemplate \nvoid Classical_AMG_Level_Base::transfer_level(AMG_Level *ref_lvl)\n{\n Classical_AMG_Level_Base *ref_cla_lvl = dynamic_cast*>(ref_lvl);\n this->P.copy(ref_cla_lvl->P);\n this->R.copy(ref_cla_lvl->R);\n this->m_s_con.copy(ref_cla_lvl->m_s_con);\n this->m_scratch.copy(ref_cla_lvl->m_scratch);\n this->m_cf_map.copy(ref_cla_lvl->m_cf_map);\n}\n\n/****************************************\n * Computes the A, P, and R operators\n ***************************************/\ntemplate \nvoid Classical_AMG_Level_Base::createCoarseVertices()\n{\n if (AMG_Level::getLevelIndex() < this->num_aggressive_levels)\n {\n if (selector) { delete selector; }\n\n selector = chooseAggressiveSelector(AMG_Level::amg->m_cfg, AMG_Level::amg->m_cfg_scope);\n }\n\n Matrix &RAP = this->getNextLevel( typename Matrix::memory_space( ) )->getA( );\n Matrix &A = this->getA();\n int size_all, size_full, nnz_full;\n\n if (!A.is_matrix_singleGPU())\n {\n int offset;\n // Need to get number of 2-ring rows\n A.getOffsetAndSizeForView(ALL, &offset, &size_all);\n A.getOffsetAndSizeForView(FULL, &offset, &size_full);\n A.getNnzForView(FULL, &nnz_full);\n }\n else\n {\n size_all = A.get_num_rows();\n size_full = A.get_num_rows();\n nnz_full = A.get_num_nz();\n }\n\n this->m_cf_map.resize(size_all);\n this->m_s_con.resize(nnz_full);\n this->m_scratch.resize(size_full);\n thrust_wrapper::fill(this->m_cf_map.begin(), this->m_cf_map.end(), 0);\n cudaCheckError();\n thrust_wrapper::fill(this->m_s_con.begin(), this->m_s_con.end(), false);\n cudaCheckError();\n thrust_wrapper::fill(this->m_scratch.begin(), this->m_scratch.end(), 0);\n cudaCheckError();\n markCoarseFinePoints();\n}\n\ntemplate \nvoid Classical_AMG_Level_Base::createCoarseMatrices()\n{\n // allocate aggressive interpolator if needed\n if (AMG_Level::getLevelIndex() < this->num_aggressive_levels)\n {\n if (interpolator) { delete interpolator; }\n\n interpolator = chooseAggressiveInterpolator(AMG_Level::amg->m_cfg, AMG_Level::amg->m_cfg_scope);\n }\n\n Matrix &RAP = this->getNextLevel( typename Matrix::memory_space( ) )->getA( );\n Matrix &A = this->getA();\n /* WARNING: exit if D1 interpolator is selected in distributed setting */\n std::string s(\"\");\n s += AMG_Level::amg->m_cfg->AMG_Config::template getParameter(\"interpolator\", AMG_Level::amg->m_cfg_scope);\n\n if (A.is_matrix_distributed() && (s.compare(\"D1\") == 0))\n {\n FatalError(\"D1 interpolation is not supported in distributed settings\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n /* WARNING: do not recompute prolongation (P) and restriction (R) when you\n are reusing the level structure (structure_reuse_levels > 0) */\n if (this->isReuseLevel() == false)\n {\n computeProlongationOperator();\n }\n\n // Compute Restriction operator and coarse matrix Ac\n if (!this->A->is_matrix_distributed() || this->A->manager->get_num_partitions() == 1)\n {\n /* WARNING: see above warning. */\n if (this->isReuseLevel() == false)\n {\n computeRestrictionOperator();\n }\n\n computeAOperator();\n }\n else\n {\n /* WARNING: notice that in this case the computeRestructionOperator() is called\n inside computeAOperator_distributed() routine. */\n computeAOperator_distributed();\n }\n\n// we also need to renumber columns of P and rows or R correspondingly since we changed RAP halo columns\n// for R we just keep track of renumbering in and exchange proper vectors in restriction\n// for P we actually need to renumber columns for prolongation:\n if (A.is_matrix_distributed() && this->A->manager->get_num_partitions() > 1)\n {\n RAP.set_initialized(0);\n // Renumber the owned nodes as interior and boundary (renumber rows and columns)\n // We are passing reuse flag to not create neighbours list from scratch, but rather update based on new halos\n RAP.manager->renumberMatrixOneRing(this->isReuseLevel());\n // Renumber the column indices of P and shuffle rows of P\n RAP.manager->renumber_P_R(this->P, this->R, A);\n // Create the B2L_maps for RAP\n RAP.manager->createOneRingHaloRows();\n RAP.manager->getComms()->set_neighbors(RAP.manager->num_neighbors());\n RAP.setView(OWNED);\n RAP.set_initialized(1);\n // update # of columns in P - this is necessary for correct CSR multiply\n P.set_initialized(0);\n int new_num_cols = thrust_wrapper::reduce(P.col_indices.begin(), P.col_indices.end(), int(0), amgx::thrust::maximum()) + 1;\n cudaCheckError();\n P.set_num_cols(new_num_cols);\n P.set_initialized(1);\n }\n\n RAP.copyAuxData(&A);\n\n if (!A.is_matrix_singleGPU() && RAP.manager == NULL)\n {\n RAP.manager = new DistributedManager();\n }\n\n if (this->getA().is_matrix_singleGPU())\n {\n this->m_next_level_size = this->getNextLevel(typename Matrix::memory_space() )->getA().get_num_rows() * this->getNextLevel(typename Matrix::memory_space() )->getA().get_block_dimy();\n }\n else\n {\n // m_next_level_size is the size that will be used to allocate xc, bc vectors\n int size, offset;\n this->getNextLevel(typename Matrix::memory_space())->getA().getOffsetAndSizeForView(FULL, &offset, &size);\n this->m_next_level_size = size * this->getNextLevel(typename Matrix::memory_space() )->getA().get_block_dimy();\n }\n}\n\ntemplate \nvoid Classical_AMG_Level_Base::markCoarseFinePoints()\n{\n Matrix &A = this->getA();\n //allocate necessary memory\n typedef Vector::Type> IVector;\n typedef Vector::Type> BVector;\n typedef Vector::Type> FVector;\n FVector weights;\n\n if (!A.is_matrix_singleGPU())\n {\n int size, offset;\n A.getOffsetAndSizeForView(FULL, &offset, &size);\n // size should now contain the number of 1-ring rows\n weights.resize(size);\n }\n else\n {\n weights.resize(A.get_num_rows());\n }\n\n thrust_wrapper::fill(weights.begin(), weights.end(), 0.0);\n cudaCheckError();\n\n // extend A to include 1st ring nodes\n // compute strong connections and weights\n if (!A.is_matrix_singleGPU())\n {\n ViewType oldView = A.currentView();\n A.setView(FULL);\n strength->computeStrongConnectionsAndWeights(A, this->m_s_con, weights, this->max_row_sum);\n A.setView(oldView);\n }\n else\n {\n strength->computeStrongConnectionsAndWeights(A, this->m_s_con, weights, this->max_row_sum);\n }\n\n // Exchange the one-ring of the weights\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo(weights, weights.tag);\n }\n\n //mark coarse and fine points\n selector->markCoarseFinePoints(A, weights, this->m_s_con, this->m_cf_map, this->m_scratch);\n // we do resize cf_map to zero later, so we are saving separate copy\n this->m_cf_map.dirtybit = 1;\n\n // Do a two ring exchange of cf_map\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_2ring(this->m_cf_map, m_cf_map.tag);\n }\n\n // Modify cf_map array such that coarse points are assigned a local index, while fine points entries are not touched\n selector->renumberAndCountCoarsePoints(this->m_cf_map, this->m_num_coarse_vertices, A.get_num_rows());\n}\n\n\ntemplate \nvoid Classical_AMG_Level_Base::computeProlongationOperator()\n{\n this->Profile.tic(\"computeP\");\n Matrix &A = this->getA();\n //allocate necessary memory\n typedef Vector::Type> IVector;\n typedef Vector::Type> BVector;\n typedef Vector::Type> FVector;\n //generate the interpolation matrix\n interpolator->generateInterpolationMatrix(A, this->m_cf_map, this->m_s_con, this->m_scratch, P);\n this->m_cf_map.clear();\n this->m_cf_map.shrink_to_fit();\n this->m_scratch.clear();\n this->m_scratch.shrink_to_fit();\n this->m_s_con.clear();\n this->m_s_con.shrink_to_fit();\n\n // truncate based on max # of elements if desired\n if (this->max_elmts > 0 && P.get_num_rows() > 0)\n {\n Truncate::truncateByMaxElements(P, this->max_elmts);\n }\n\n if (!P.isLatencyHidingEnabled(*this->amg->m_cfg))\n {\n // This will cause bsrmv_with_mask to not do latency hiding\n P.setInteriorView(OWNED);\n P.setExteriorView(OWNED);\n }\n}\n\n/**********************************************\n * computes R=P^T\n **********************************************/\ntemplate \nvoid Classical_AMG_Level_Base::computeRestrictionOperator()\n{\n this->Profile.tic(\"computeR\");\n R.set_initialized(0);\n P.setView(OWNED);\n transpose(P, R, P.get_num_rows());\n\n if (!R.isLatencyHidingEnabled(*this->amg->m_cfg))\n {\n // This will cause bsrmv_with_mask_restriction to not do latency hiding\n R.setInteriorView(OWNED);\n R.setExteriorView(OWNED);\n }\n\n if(P.is_matrix_distributed())\n {\n // Setup the number of non-zeros in R using stub DistributedManager\n R.manager = new DistributedManager();\n int nrows_owned = P.manager->halo_offsets[0];\n int nrows_full = P.manager->halo_offsets[P.manager->neighbors.size()];\n int nz_full = R.row_offsets[nrows_full];\n int nz_owned = R.row_offsets[nrows_owned];\n R.manager->setViewSizes(nrows_owned, nz_owned, nrows_owned, nz_owned, nrows_full, nz_full, R.get_num_rows(), R.get_num_nz());\n }\n\n R.set_initialized(1);\n this->Profile.toc(\"computeR\");\n}\n\n/**********************************************\n * computes the Galerkin product: A_c=R*A*P\n **********************************************/\n\ntemplate \nvoid Classical_AMG_Level >::computeAOperator_1x1()\n{\n this->Profile.tic(\"computeA\");\n Matrix RA;\n RA.addProps(CSR);\n RA.set_block_dimx(this->getA().get_block_dimx());\n RA.set_block_dimy(this->getA().get_block_dimy());\n Matrix &RAP = this->getNextLevel( typename Matrix::memory_space( ) )->getA( );\n RAP.addProps(CSR);\n RAP.set_block_dimx(this->getA().get_block_dimx());\n RAP.set_block_dimy(this->getA().get_block_dimy());\n Matrix &Atmp = this->getA();\n multiplyMM(this->R, this->getA(), RA);\n multiplyMM(RA, this->P, RAP);\n RAP.sortByRowAndColumn();\n RAP.set_initialized(1);\n this->Profile.toc(\"computeA\");\n}\n\ntemplate \nvoid Classical_AMG_Level >::computeAOperator_1x1_distributed()\n{\n FatalError(\"Distributed classical AMG not implemented for host\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n\ntemplate \nvoid Classical_AMG_Level >::computeAOperator_1x1()\n{\n this->Profile.tic(\"computeA\");\n Matrix &RAP = this->getNextLevel( device_memory( ) )->getA( );\n RAP.addProps(CSR);\n RAP.set_block_dimx(this->getA().get_block_dimx());\n RAP.set_block_dimy(this->getA().get_block_dimy());\n this->R.set_initialized( 0 );\n this->R.addProps( CSR );\n this->R.set_initialized( 1 );\n this->P.set_initialized( 0 );\n this->P.addProps( CSR );\n this->P.set_initialized( 1 );\n void *wk = AMG_Level::amg->getCsrWorkspace();\n\n if ( wk == NULL )\n {\n wk = CSR_Multiply::csr_workspace_create( *(AMG_Level::amg->m_cfg), AMG_Level::amg->m_cfg_scope );\n AMG_Level::amg->setCsrWorkspace( wk );\n }\n\n int spmm_verbose = this->amg->m_cfg->AMG_Config::template getParameter(\"spmm_verbose\", this->amg->m_cfg_scope);\n\n if ( spmm_verbose )\n {\n typedef typename Matrix::IVector::const_iterator Iterator;\n typedef amgx::thrust::pair Result;\n std::ostringstream buffer;\n buffer << \"SPMM: Level \" << this->getLevelIndex() << std::endl;\n\n if ( this->getLevelIndex() == 0 )\n {\n device_vector_alloc num_nz( this->getA().row_offsets.size() );\n amgx::thrust::adjacent_difference( this->getA().row_offsets.begin(), this->getA().row_offsets.end(), num_nz.begin() );\n cudaCheckError();\n Result result = amgx::thrust::minmax_element( num_nz.begin() + 1, num_nz.end() );\n cudaCheckError();\n int min_size = *result.first;\n int max_size = *result.second;\n int sum = thrust_wrapper::reduce( num_nz.begin() + 1, num_nz.end() );\n cudaCheckError();\n double avg_size = double(sum) / this->getA().get_num_rows();\n buffer << \"SPMM: A: \" << std::endl;\n buffer << \"SPMM: Matrix avg row size: \" << avg_size << std::endl;\n buffer << \"SPMM: Matrix min row size: \" << min_size << std::endl;\n buffer << \"SPMM: Matrix max row size: \" << max_size << std::endl;\n }\n\n device_vector_alloc num_nz( this->P.row_offsets.size() );\n amgx::thrust::adjacent_difference( this->P.row_offsets.begin(), this->P.row_offsets.end(), num_nz.begin() );\n cudaCheckError();\n Result result = amgx::thrust::minmax_element( num_nz.begin() + 1, num_nz.end() );\n cudaCheckError();\n int min_size = *result.first;\n int max_size = *result.second;\n int sum = thrust_wrapper::reduce( num_nz.begin() + 1, num_nz.end() );\n cudaCheckError();\n double avg_size = double(sum) / this->P.get_num_rows();\n buffer << \"SPMM: P: \" << std::endl;\n buffer << \"SPMM: Matrix avg row size: \" << avg_size << std::endl;\n buffer << \"SPMM: Matrix min row size: \" << min_size << std::endl;\n buffer << \"SPMM: Matrix max row size: \" << max_size << std::endl;\n num_nz.resize( this->R.row_offsets.size() );\n amgx::thrust::adjacent_difference( this->R.row_offsets.begin(), this->R.row_offsets.end(), num_nz.begin() );\n cudaCheckError();\n result = amgx::thrust::minmax_element( num_nz.begin() + 1, num_nz.end() );\n cudaCheckError();\n min_size = *result.first;\n max_size = *result.second;\n sum = thrust_wrapper::reduce( num_nz.begin() + 1, num_nz.end() );\n cudaCheckError();\n avg_size = double(sum) / this->R.get_num_rows();\n buffer << \"SPMM: R: \" << std::endl;\n buffer << \"SPMM: Matrix avg row size: \" << avg_size << std::endl;\n buffer << \"SPMM: Matrix min row size: \" << min_size << std::endl;\n buffer << \"SPMM: Matrix max row size: \" << max_size << std::endl;\n amgx_output( buffer.str().c_str(), static_cast( buffer.str().length() ) );\n }\n\n RAP.set_initialized( 0 );\n CSR_Multiply::csr_galerkin_product( this->R, this->getA(), this->P, RAP, NULL, NULL, NULL, NULL, NULL, NULL, wk );\n RAP.set_initialized( 1 );\n int spmm_no_sort = this->amg->m_cfg->AMG_Config::template getParameter(\"spmm_no_sort\", this->amg->m_cfg_scope);\n this->Profile.toc(\"computeA\");\n}\n/**********************************************\n * computes the restriction: rr=R*r\n **********************************************/\ntemplate \nvoid Classical_AMG_Level_Base::restrictResidual(VVector &r, VVector &rr)\n{\n// we need to resize residual vector to make sure it can store halo rows to be sent\n if (!P.is_matrix_singleGPU())\n {\n typedef typename TConfig::MemSpace MemorySpace;\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n#if COARSE_CLA_CONSO\n int desired_size ;\n\n if (this->getNextLevel(MemorySpace())->isConsolidationLevel())\n {\n desired_size = std::max(P.manager->halo_offsets[P.manager->neighbors.size()], Ac.manager->halo_offsets_before_glue[Ac.manager->neighbors_before_glue.size()] * rr.get_block_size());\n }\n else\n {\n desired_size = std::max(P.manager->halo_offsets[P.manager->neighbors.size()], Ac.manager->halo_offsets[Ac.manager->neighbors.size()] * rr.get_block_size());\n }\n\n#else\n int desired_size = std::max(P.manager->halo_offsets[P.manager->neighbors.size()], Ac.manager->halo_offsets[Ac.manager->neighbors.size()] * rr.get_block_size());\n#endif\n rr.resize(desired_size);\n }\n\n#if 1\n this->Profile.tic(\"restrictRes\");\n\n // Disable speculative send of rr\n if (P.is_matrix_singleGPU())\n {\n multiply( R, r, rr);\n }\n else\n {\n multiply_with_mask_restriction( R, r, rr, P);\n }\n\n#endif\n // exchange halo residuals & add residual contribution from neighbors\n rr.dirtybit = 1;\n\n if (!P.is_matrix_singleGPU())\n {\n int desired_size = P.manager->halo_offsets[P.manager->neighbors.size()] * rr.get_block_size();\n\n if (rr.size() < desired_size)\n {\n rr.resize(desired_size);\n }\n }\n\n this->Profile.toc(\"restrictRes\");\n}\n\nstruct is_minus_one\n{\n __host__ __device__\n bool operator()(const int &x)\n {\n return x == -1;\n }\n};\n\n\ntemplate \nvoid Classical_AMG_Level >::computeAOperator_1x1_distributed()\n{\n Matrix &A = this->getA();\n Matrix &P = this->P;\n Matrix &RAP = this->getNextLevel( device_memory( ) )->getA( );\n RAP.addProps(CSR);\n RAP.set_block_dimx(this->getA().get_block_dimx());\n RAP.set_block_dimy(this->getA().get_block_dimy());\n IndexType num_parts = A.manager->get_num_partitions();\n IndexType num_neighbors = A.manager->num_neighbors();\n IndexType my_rank = A.manager->global_id();\n // OWNED includes interior and boundary\n A.setView(OWNED);\n int num_owned_coarse_pts = P.manager->halo_offsets[0];\n int num_owned_fine_pts = A.manager->halo_offsets[0];\n\n // Initialize RAP.manager\n if (RAP.manager == NULL)\n {\n RAP.manager = new DistributedManager();\n }\n\n RAP.manager->A = &RAP;\n RAP.manager->setComms(A.manager->getComms());\n RAP.manager->set_global_id(my_rank);\n RAP.manager->set_num_partitions(num_parts);\n RAP.manager->part_offsets_h = P.manager->part_offsets_h;\n RAP.manager->part_offsets = P.manager->part_offsets;\n RAP.manager->set_base_index(RAP.manager->part_offsets_h[my_rank]);\n RAP.manager->set_index_range(num_owned_coarse_pts);\n RAP.manager->num_rows_global = RAP.manager->part_offsets_h[num_parts];\n // --------------------------------------------------------------------\n // Using the B2L_maps of matrix A, identify the rows of P that need to be sent to neighbors,\n // so that they can compute A*P\n // Once rows of P are identified, convert the column indices to global indices, and send them to neighbors\n // ---------------------------------------------------------------------------\n // Copy some information about the manager of P, since we don't want to modify those\n IVector_h P_neighbors = P.manager->neighbors;\n I64Vector_h P_halo_ranges_h = P.manager->halo_ranges_h;\n I64Vector_d P_halo_ranges = P.manager->halo_ranges;\n RAP.manager->local_to_global_map = P.manager->local_to_global_map;\n IVector_h P_halo_offsets = P.manager->halo_offsets;\n // Create a temporary distributed arranger\n DistributedArranger *prep = new DistributedArranger;\n prep->exchange_halo_rows_P(A, this->P, RAP.manager->local_to_global_map, P_neighbors, P_halo_ranges_h, P_halo_ranges, P_halo_offsets, RAP.manager->part_offsets_h, RAP.manager->part_offsets, num_owned_coarse_pts, RAP.manager->part_offsets_h[my_rank]);\n cudaCheckError();\n // At this point, we can compute RAP_full which contains some rows that will need to be sent to neighbors\n // i.e. RAP_full = [ RAP_int ]\n // [ RAP_ext ]\n // RAP is [ RAP_int ] + [RAP_ext_received_from_neighbors]\n // We can reuse the serial galerkin product since R, A and P use local indices\n // TODO: latency hiding (i.e. compute RAP_ext, exchange_matrix_halo, then do RAP_int)\n /* WARNING: do not recompute prolongation (P) and restriction (R) when you\n are reusing the level structure (structure_reuse_levels > 0) */\n /* We force for matrix P to have only owned rows to be seen for the correct galerkin product computation*/\n this->P.set_initialized(0);\n this->P.set_num_rows(num_owned_fine_pts);\n this->P.addProps( CSR );\n this->P.set_initialized(1);\n\n if (this->isReuseLevel() == false)\n {\n this->R.set_initialized( 0 );\n this->R.addProps( CSR );\n // Take the tranpose of P to get R\n // Single-GPU transpose, no mpi exchange\n this->computeRestrictionOperator();\n this->R.set_initialized( 1 );\n }\n\n this->Profile.tic(\"computeA\");\n Matrix RAP_full;\n // Initialize the workspace needed for galerkin product\n void *wk = AMG_Level::amg->getCsrWorkspace();\n\n if ( wk == NULL )\n {\n wk = CSR_Multiply::csr_workspace_create( *(AMG_Level::amg->m_cfg), AMG_Level::amg->m_cfg_scope );\n AMG_Level::amg->setCsrWorkspace( wk );\n }\n\n // Single-GPU RAP, no mpi exchange\n RAP_full.set_initialized( 0 );\n /* WARNING: Since A is reordered (into interior and boundary nodes), while R and P are not reordered,\n you must unreorder A when performing R*A*P product in ordre to obtain the correct result. */\n CSR_Multiply::csr_galerkin_product( this->R, this->getA(), this->P, RAP_full,\n /* permutation for rows of R, A and P */ NULL, NULL /*&(this->getA().manager->renumbering)*/, NULL,\n /* permutation for cols of R, A and P */ NULL, NULL /*&(this->getA().manager->inverse_renumbering)*/, NULL,\n wk );\n RAP_full.set_initialized( 1 );\n this->Profile.toc(\"computeA\");\n // ----------------------------------------------------------------------------------------------\n // Now, send rows of RAP_full requireq by neighbors, received rows from neighbors and create RAP\n // ----------------------------------------------------------------------------------------------\n prep->exchange_RAP_ext(RAP, RAP_full, A, this->P, P_halo_offsets, RAP.manager->local_to_global_map, P_neighbors, P_halo_ranges_h, P_halo_ranges, RAP.manager->part_offsets_h, RAP.manager->part_offsets, num_owned_coarse_pts, RAP.manager->part_offsets_h[my_rank], wk);\n // Delete temporary distributed arranger\n delete prep;\n /* WARNING: The RAP matrix generated at this point contains extra rows (that correspond to rows of R,\n that was obtained by locally transposing P). This rows are ignored by setting the # of matrix rows\n to be smaller, so that they correspond to number of owned coarse nodes. This should be fine, but\n it leaves holes in the matrix as there might be columns that belong to the extra rows that now do not\n belong to the smaller matrix with number of owned coarse nodes rows. The same is trued about the\n local_to_global_map. These two data structures match at this point. However, in the next calls\n local_to_global (exclusively) will be used to geberate B2L_maps (wihtout going through column indices)\n which creates extra elements in the B2L that simply do not exist in the new matrices. I strongly suspect\n this is the reason fore the bug. The below fix simply compresses the matrix so that there are no holes\n in it, or in the local_2_global_map. */\n //mark local_to_global_columns that exist in the owned coarse nodes rows.\n IndexType nrow = RAP.get_num_rows();\n IndexType ncol = RAP.get_num_cols();\n IndexType nl2g = ncol - nrow;\n\n if (nl2g > 0)\n {\n IVector l2g_p(nl2g + 1, 0); //+1 is needed for prefix_sum/exclusive_scan\n I64Vector l2g_t(nl2g, 0);\n IndexType nblocks = (nrow + AMGX_CAL_BLOCK_SIZE - 1) / AMGX_CAL_BLOCK_SIZE;\n\n if (nblocks > 0)\n flag_existing_local_to_global_columns <<< nblocks, AMGX_CAL_BLOCK_SIZE>>>\n (nrow, RAP.row_offsets.raw(), RAP.col_indices.raw(), l2g_p.raw());\n\n cudaCheckError();\n /*\n //slow version of the above kernel\n for(int ii=0; ii=nrow){\n int kk = col-RAP.get_num_rows();\n l2g_p[kk] = 1;\n }\n }\n }\n cudaCheckError();\n */\n //create a pointer map for their location using prefix sum\n thrust_wrapper::exclusive_scan(l2g_p.begin(), l2g_p.end(), l2g_p.begin());\n int new_nl2g = l2g_p[nl2g];\n\n //compress the columns using the pointer map\n if (nblocks > 0)\n compress_existing_local_columns <<< nblocks, AMGX_CAL_BLOCK_SIZE>>>\n (nrow, RAP.row_offsets.raw(), RAP.col_indices.raw(), l2g_p.raw());\n\n cudaCheckError();\n /*\n //slow version of the above kernel\n for(int ii=0; ii=nrow){\n int kk = col-RAP.get_num_rows();\n RAP.col_indices[jj] = nrow+l2g_p[kk];\n }\n }\n }\n cudaCheckError();\n */\n //adjust matrix size (number of columns) accordingly\n RAP.set_initialized(0);\n RAP.set_num_cols(nrow + new_nl2g);\n RAP.set_initialized(1);\n //compress local_to_global_map using the pointer map\n nblocks = (nl2g + AMGX_CAL_BLOCK_SIZE - 1) / AMGX_CAL_BLOCK_SIZE;\n\n if (nblocks > 0)\n compress_existing_local_to_global_columns <<< nblocks, AMGX_CAL_BLOCK_SIZE>>>\n (nl2g, RAP.manager->local_to_global_map.raw(), l2g_t.raw(), l2g_p.raw());\n\n cudaCheckError();\n amgx::thrust::copy(l2g_t.begin(), l2g_t.begin() + new_nl2g, RAP.manager->local_to_global_map.begin());\n cudaCheckError();\n /*\n //slow version of the above kernel (through Thrust)\n for(int ii=0; ii<(l2g_p.size()-1); ii++){\n if (l2g_p[ii] != l2g_p[ii+1]){\n RAP.manager->local_to_global_map[l2g_p[ii]] = RAP.manager->local_to_global_map[ii];\n }\n }\n cudaCheckError();\n */\n //adjust local_to_global_map size accordingly\n RAP.manager->local_to_global_map.resize(new_nl2g);\n }\n}\n\n/**********************************************\n * prolongates the error: x+=P*e\n **********************************************/\ntemplate \nvoid Classical_AMG_Level_Base::prolongateAndApplyCorrection(VVector &e, VVector &bc, VVector &x, VVector &tmp)\n{\n this->Profile.tic(\"proCorr\");\n // Use P.manager to exchange halo of e before doing P\n // (since P has columns belonging to one of P.neighbors)\n e.dirtybit = 1;\n\n if (!P.is_matrix_singleGPU())\n {\n // get coarse matrix\n typedef typename TConfig::MemSpace MemorySpace;\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n#if COARSE_CLA_CONSO\n int e_size;\n\n if (this->getNextLevel(MemorySpace())->isConsolidationLevel())\n {\n e_size = std::max(P.manager->halo_offsets[P.manager->neighbors.size()], Ac.manager->halo_offsets_before_glue[Ac.manager->neighbors_before_glue.size()]) * e.get_block_size();\n }\n else\n {\n e_size = std::max(P.manager->halo_offsets[P.manager->neighbors.size()], Ac.manager->halo_offsets[Ac.manager->neighbors.size()]) * e.get_block_size();\n }\n\n if (e.size() < e_size) { e.resize(e_size); }\n\n#else\n int e_size = std::max(P.manager->halo_offsets[P.manager->neighbors.size()], Ac.manager->halo_offsets[Ac.manager->neighbors.size()]) * e.get_block_size();\n e.resize(e_size);\n#endif\n }\n\n if (P.is_matrix_singleGPU())\n {\n if (e.size() > 0)\n {\n multiply( P, e, tmp);\n }\n }\n else\n {\n multiply_with_mask( P, e, tmp);\n }\n\n // get owned num rows for fine matrix\n int owned_size;\n\n if (this->A->is_matrix_distributed())\n {\n int owned_offset;\n P.manager->getOffsetAndSizeForView(OWNED, &owned_offset, &owned_size);\n }\n else\n {\n owned_size = x.size();\n }\n\n //apply\n axpby(x, tmp, x, ValueType(1), ValueType(1), 0, owned_size);\n this->Profile.toc(\"proCorr\");\n x.dirtybit = 1;\n}\n\ntemplate \nvoid Classical_AMG_Level_Base::computeAOperator()\n{\n if (this->A->get_block_size() == 1)\n {\n computeAOperator_1x1();\n }\n else\n {\n FatalError(\"Classical AMG not implemented for block_size != 1\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\ntemplate \nvoid Classical_AMG_Level_Base::computeAOperator_distributed()\n{\n if (this->A->get_block_size() == 1)\n {\n computeAOperator_1x1_distributed();\n }\n else\n {\n FatalError(\"Classical AMG not implemented for block_size != 1\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n}\n\n\ntemplate \nvoid Classical_AMG_Level_Base::consolidateVector(VVector &x)\n{\n#ifdef AMGX_WITH_MPI\n#if COARSE_CLA_CONSO\n typedef typename TConfig::MemSpace MemorySpace;\n Matrix &A = this->getA();\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n MPI_Comm comm, temp_com;\n comm = Ac.manager->getComms()->get_mpi_comm();\n temp_com = compute_glue_matrices_communicator(Ac);\n glue_vector(Ac, comm, x, temp_com);\n#endif\n#endif\n}\n\ntemplate \nvoid Classical_AMG_Level_Base::unconsolidateVector(VVector &x)\n{\n#ifdef AMGX_WITH_MPI\n#if COARSE_CLA_CONSO\n typedef typename TConfig::MemSpace MemorySpace;\n Matrix &A = this->getA();\n Matrix &Ac = this->getNextLevel( MemorySpace( ) )->getA();\n MPI_Comm comm, temp_com;\n comm = Ac.manager->getComms()->get_mpi_comm();\n temp_com = compute_glue_matrices_communicator(Ac);\n unglue_vector(Ac, comm, x, temp_com, x);\n#endif\n#endif\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class Classical_AMG_Level_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Classical_AMG_Level::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace classical\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/classical/classical_amg_level.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"solvers/solver.h\"\n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n\nnamespace amgx\n{\n\nnamespace classical\n{\n\n///////////////////////////////////////////////////////////////////////////////////////////////////\n\n/*************************************************************************\n * Generate uniform random vector\n ************************************************************************/\nstruct prg\n{\n float a, b;\n\n __host__ __device__\n prg(float _a = 0.f, float _b = 1.f) : a(_a), b(_b) {};\n\n __host__ __device__\n float operator()(const unsigned int n) const\n {\n amgx::thrust::default_random_engine rng;\n amgx::thrust::uniform_real_distribution dist(a, b);\n rng.discard(n);\n return dist(rng);\n }\n};\n\n\n// Initialization of test error vector with random entries\ntemplate \nvoid initRandom(Vector &vec)\n{\n const unsigned int size = vec.size();\n amgx::thrust::counting_iterator index_sequence_begin(0);\n thrust_wrapper::transform(index_sequence_begin, index_sequence_begin + size,\n vec.begin(), prg(0.f, 1.f));\n}\n\n// Zero out the entries of a vector corresponding to the coarse points\ntemplate \n__global__\nvoid zero_coarse_points_kernel(const int *cf_map, const int numRows, ValueType *vec)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n // if point is coarse, zero out the corresponding entry of vec\n if (cf_map[tIdx] == COARSE)\n {\n vec[tIdx] = 0.;\n }\n }\n}\n\n// Create a vector which extracts fine points from the error vector and is\n// zero at the coarse points\ntemplate \n__global__\nvoid get_fine_error_kernel( const int *cf_map, const int numRows,\n const ValueType *v_err, ValueType *vec )\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n // if point is coarse, zero out the corresponding entry of vec\n if (cf_map[tIdx] == COARSE)\n {\n vec[tIdx] = 0.;\n }\n else\n {\n vec[tIdx] = v_err[tIdx];\n }\n }\n}\n\n\n// Update only the fine points of the error vector (v_err) based on input vec.\ntemplate \n__global__\nvoid update_fine_zero_coarse_error_kernel(const int *cf_map, const int numRows,\n ValueType *v_err, const ValueType *vec)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n // if point is coarse, zero out the corresponding entry of vec\n if (cf_map[tIdx] == FINE)\n {\n v_err[tIdx] = vec[tIdx];\n }\n else\n {\n v_err[tIdx] = (ValueType) 0;\n }\n }\n}\n\n\n// Update coarse and fine points from the matrix coloring taking the first\n// levels colors to form the independent set.\ntemplate \n__global__\nvoid update_cf_map_kernel(IndexType *cf_map, const int numRows, const int levels, const IndexType *A_row_colors)\n{\n //coarse = [coarse; find(independent)];\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n // go over each row\n if (A_row_colors[tIdx] <= levels) // double check that this is what we want\n {\n cf_map[tIdx] = COARSE;\n }\n }\n}\n\n\n\n#define EXPERIMENTAL_CR\n\n#ifdef EXPERIMENTAL_CR\n\n\n\n// Update only the fine points of the error vector (v_err) based on input vec.\ntemplate \n__global__\nvoid compute_Asc_nnzPerRow_kernel2(const IndexType *ArowOffsets, const IndexType *AcolInd,\n const ValueType *Avalues, const ValueType *AdiagValues,\n const IndexType AnumRows, const IndexType *cf_map,\n IndexType *Asc_nnzPerRow, ValueType *row_thresh, const ValueType thresh)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AnumRows; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on one \"fine\" row of A at a time\n const int ArowBegin = ArowOffsets[tIdx];\n const int ArowEnd = ArowOffsets[tIdx + 1];\n\n if (cf_map[tIdx] == FINE)\n {\n // work with fine pts only\n // Matlab pseudo-code:\n // Asc(:,i) = A(:,i)/A(i,i)<-thresh;\n // but this thresh is dynamic\n ValueType fine_avg_thresh = 0.;\n int fine_count = 0;\n\n for (int j = ArowBegin; j < ArowEnd; j++)\n {\n int AjColIndex = AcolInd[j];\n\n if ((AjColIndex != tIdx) && (cf_map[AjColIndex] == FINE))\n {\n fine_avg_thresh += Avalues[j];\n fine_count ++;\n }\n }\n\n fine_avg_thresh *= thresh / fine_count;\n row_thresh[tIdx] = fine_avg_thresh;\n\n for (int j = ArowBegin; j < ArowEnd; j++)\n {\n int AjColIndex = AcolInd[j];\n\n if (AjColIndex == tIdx)\n {\n /*Asc_nnzPerRow[tIdx]++;*/\n continue;\n }\n\n if (cf_map[AjColIndex] == FINE)\n {\n if (Avalues[j] / AdiagValues[tIdx] < fine_avg_thresh / AdiagValues[tIdx] || Avalues[j] / AdiagValues[AjColIndex] < fine_avg_thresh / AdiagValues[AjColIndex])\n {\n Asc_nnzPerRow[tIdx]++;\n }\n }\n }\n }\n }\n}\n\n\n\n// Update only the fine points of the error vector (v_err) based on input vec.\ntemplate \n__global__\nvoid compute_AscColInd_kernel(const IndexType *ArowOffsets, const IndexType *AcolInd,\n const ValueType *Avalues, const ValueType *AdiagValues,\n const IndexType AnumRows, const IndexType *cf_map,\n const IndexType *AscRowOffsets, IndexType *AscColInd,\n const ValueType *row_thresh)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AnumRows; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on one \"coarse\" row of A at a time\n const int ArowBegin = ArowOffsets[tIdx];\n const int ArowEnd = ArowOffsets[tIdx + 1];\n const int AscRowEnd = AscRowOffsets[tIdx + 1];\n ValueType rowMax(0);\n\n if (cf_map[tIdx] == FINE) // tIdx = j (in matlab) i.e. row of A\n {\n // work with fine pts only\n // Matlab pseudo-code:\n // Asc(:,i) = A(:,i)/A(i,i)<-thresh;\n int AscEntryNum = AscRowOffsets[tIdx]; // starting location of the current column of P\n\n // calculate max(-rowentries)\n for (int j = ArowBegin; j < ArowEnd && AscEntryNum < AscRowEnd; j++)\n {\n rowMax = max( rowMax, -Avalues[j]);\n }\n\n for (int j = ArowBegin; j < ArowEnd && AscEntryNum < AscRowEnd; j++)\n {\n // go thru each \"fine\" column of A\n int AjColIndex = AcolInd[j];\n\n if (AjColIndex == tIdx) { /*AscColInd[AscEntryNum] = AjColIndex; AscEntryNum++;*/ continue; }\n\n if (cf_map[AjColIndex] == FINE)\n {\n if (-Avalues[j] >= row_thresh[tIdx]*rowMax)\n {\n AscColInd[AscEntryNum] = AjColIndex;\n AscEntryNum++;\n }\n }\n }\n }\n }\n}\n\n#else\n\n\n// Update only the fine points of the error vector (v_err) based on input vec.\ntemplate \n__global__\nvoid compute_Asc_nnzPerRow_kernel(const IndexType *ArowOffsets, const IndexType *AcolInd,\n const ValueType *Avalues, const ValueType *AdiagValues,\n const IndexType AnumRows, const IndexType *cf_map,\n IndexType *Asc_nnzPerRow, const ValueType thresh)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AnumRows; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on one \"fine\" row of A at a time\n const int ArowBegin = ArowOffsets[tIdx];\n const int ArowEnd = ArowOffsets[tIdx + 1];\n\n if (cf_map[tIdx] == FINE)\n {\n // work with fine pts only\n // Matlab pseudo-code:\n // Asc(:,i) = A(:,i)/A(i,i)<-thresh;\n for (int j = ArowBegin; j < ArowEnd; j++)\n {\n // go thru each \"fine\" column of A\n int AjColIndex = AcolInd[j];\n\n if (AjColIndex == tIdx)\n {\n /*Asc_nnzPerRow[tIdx]++;*/\n continue;\n }\n\n if (cf_map[AjColIndex] == FINE)\n {\n if (Avalues[j] / AdiagValues[tIdx] < -thresh || Avalues[j] / AdiagValues[AjColIndex] < -thresh)\n {\n Asc_nnzPerRow[tIdx]++;\n }\n }\n }\n }\n }\n}\n\n\n\n// Update only the fine points of the error vector (v_err) based on input vec.\ntemplate \n__global__\nvoid compute_AscColInd_kernel(const IndexType *ArowOffsets, const IndexType *AcolInd,\n const ValueType *Avalues, const ValueType *AdiagValues,\n const IndexType AnumRows, const IndexType *cf_map,\n const IndexType *AscRowOffsets, IndexType *AscColInd,\n const ValueType thresh)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AnumRows; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on one \"coarse\" row of A at a time\n const int ArowBegin = ArowOffsets[tIdx];\n const int ArowEnd = ArowOffsets[tIdx + 1];\n const int AscRowEnd = AscRowOffsets[tIdx + 1];\n\n if (cf_map[tIdx] == FINE) // tIdx = j (in matlab) i.e. row of A\n {\n // work with fine pts only\n // Matlab pseudo-code:\n // Asc(:,i) = A(:,i)/A(i,i)<-thresh;\n int AscEntryNum = AscRowOffsets[tIdx]; // starting location of the current column of P\n\n for (int j = ArowBegin; j < ArowEnd && AscEntryNum < AscRowEnd; j++)\n {\n // go thru each \"fine\" column of A\n int AjColIndex = AcolInd[j];\n\n if (AjColIndex == tIdx) { /*AscColInd[AscEntryNum] = AjColIndex; AscEntryNum++;*/ continue; }\n\n if (cf_map[AjColIndex] == FINE)\n {\n if ((Avalues[j] / AdiagValues[tIdx] < -thresh) || (Avalues[j] / AdiagValues[AjColIndex] < -thresh))\n {\n AscColInd[AscEntryNum] = AjColIndex;\n AscEntryNum++;\n }\n }\n }\n }\n }\n}\n\n#endif\n///////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nCR_Selector >::\nCR_Selector(AMG_Config &cfg, const std::string &cfg_scope) : \n CR_SelectorBase(cfg, cfg_scope),\n m_smoother(0) {}\n\n\ntemplate \nCR_Selector >::\n~CR_Selector()\n{\n if (m_smoother) { delete m_smoother; }\n}\n\n\n// Private functions for the implementation on device\n\n// Presmooth error\ntemplate \nvoid CR_Selector >\n::presmoothFineError( Matrix_d &A, const VVector &AdiagValues, const IVector &cf_map,\n Vector_d &v_u, Vector_d &v_tmp, Vector_d &v_z,\n Solver *smoother, ValueType &norm0,\n const ValueType rho_thresh, const int pre )\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Vector_d::value_type ValueTypeB;\n // Matrix A\n const IndexType AnumRows = (IndexType) A.get_num_rows();\n // Coarse/Fine map\n const int *cf_map_ptr = cf_map.raw();\n // Vector with the initial 'fine' error\n ValueTypeB *v_u_ptr = v_u.raw();\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (AnumRows / blocksize + 1));\n // Compute 'energy norm' of v_u (normalize error before smoothing)\n // v_tmp = diag(Aff).*v_u\n thrust_wrapper::transform(AdiagValues.begin(), AdiagValues.end(),\n v_u.begin(), v_tmp.begin(),\n amgx::thrust::multiplies());\n cudaCheckError();\n // norm0 = sqrt(v_tmp' * v_u)\n norm0 = sqrt( amgx::thrust::inner_product(v_u.begin(), v_u.end(),\n v_tmp.begin(), ValueTypeB(0.)) );\n cudaCheckError();\n // normalize: v_u = v_u / (norm0 + 1e-12);\n thrust_wrapper::transform(v_u.begin(), v_u.end(),\n amgx::thrust::make_constant_iterator(norm0 + 1.0e-12), v_u.begin(),\n amgx::thrust::divides());\n cudaCheckError();\n\n // ---------------------------- begin error presmoothing --------------------------\n for (int k = 1; k <= 5; k++)\n {\n // Presmooth with Symmetric Gauss-Seidel (or ilu) for pre number of sm(ooth) s(teps)\n // WARNING: zeroing out technique used here is not equivalent to the fine matrix pre-smoothing!\n for (int sms = 1; sms <= pre; sms++)\n {\n smoother->solve_iteration( v_z, v_u, false );\n zero_coarse_points_kernel\n <<< numBlocks, blocksize>>>(cf_map_ptr, AnumRows, v_u_ptr);\n cudaCheckError();\n }\n\n // -------------------- compute 'energy norm' of v_u ------------------------------\n // v_tmp = diag(Aff).*v_u\n thrust_wrapper::transform(AdiagValues.begin(), AdiagValues.end(),\n v_u.begin(), v_tmp.begin(),\n amgx::thrust::multiplies());\n cudaCheckError();\n // norm0 = sqrt(v_tmp' * v_u)\n norm0 = sqrt( amgx::thrust::inner_product(v_u.begin(), v_u.end(),\n v_tmp.begin(), ValueTypeB(0.)) );\n cudaCheckError();\n // -------------------- done with 'energy norm' of v_u ----------------------------\n // normalize: v_u = v_u / (norm0 + 1e-12);\n thrust_wrapper::transform(v_u.begin(), v_u.end(),\n amgx::thrust::make_constant_iterator(norm0 + 1.0e-12), v_u.begin(),\n amgx::thrust::divides());\n cudaCheckError();\n\n // norm0 := rho\n if (norm0 > 5 || norm0 < rho_thresh) { break; }\n }\n\n // ---------------------------- end error presmoothing ------------------------------\n}\n\n\n// A single iteration of CR while loop.\ntemplate \nvoid CR_Selector >\n::CR_iteration( Matrix_d &A, IVector &cf_map, int &numFine,\n Vector_d &v_err, ValueType &norm0,\n const ValueType maxit )\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n typedef typename Vector_d::value_type ValueTypeB;\n // Matrix A\n const IndexType AnumRows = (IndexType) A.get_num_rows();\n // Coarse/Fine map\n IndexType *cf_map_ptr = cf_map.raw();\n // Raw ptr for the row colors from the multi-color smoother\n const IndexType *A_row_colors_ptr;\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (AnumRows / blocksize + 1));\n const int pre = 10;\n const ValueType alpha = 0.7;\n const ValueType rho_thresh = 1.0e-2;\n const IndexType *ArowOffsets_ptr = A.row_offsets.raw();\n const IndexType *AcolInd_ptr = A.col_indices.raw();\n const ValueType *Avalues_ptr = A.values.raw();\n // temporary vectors - diagonal and non-zero offsets\n VVector AdiagValues(AnumRows, 0.0);\n // Extract the diagonal of A\n find_diag_kernel_indexed_dia <<< numBlocks, blocksize>>>(AnumRows, A.diag.raw(), Avalues_ptr, AdiagValues.raw());\n cudaCheckError();\n // ----------------------------------- define smoother ------------------------------------------\n // Set the smoother to Gauss-Seidel (json config format)\n AMG_Config cfg;\n std::string cfg_string = \"{ \\\"config_version\\\": 2\";\n cfg_string += \", \\\"solver\\\": { \\\"scope\\\": \\\"presmoothError\\\", \\\"solver\\\": \\\"MULTICOLOR_GS\\\"\";\n cfg_string += \" ,\\\"matrix_coloring_scheme\\\": \\\"PARALLEL_GREEDY\\\"\";\n cfg_string += \" } }\";\n cfg.parseParameterString(cfg_string.c_str());\n // Allocate smoother\n m_smoother = SolverFactory::allocate(cfg, \"default\", \"solver\");\n // --------------------------------- end define smoother ----------------------------------------\n Vector_d v_z(AnumRows, 0.0);\n Vector_d v_u(AnumRows);\n Vector_d v_tmp(v_u.size());\n ValueTypeB *v_err_ptr = v_err.raw();\n ValueTypeB *v_u_ptr = v_u.raw();\n // Store the number of non-zeros in each column of the interpolation matrix P\n IVector Asc_nnzPerRow(AnumRows, 0);\n#ifdef EXPERIMENTAL_CR\n VVector Asc_thresh_per_row(AnumRows, 0);\n ValueTypeB CR_P_trunc = 0.9; // truncation coefficient for CR P matrix, [0.0 ... 1.0]\n ValueTypeB Asc_thresh = (CR_P_trunc < epsilon(CR_P_trunc)) ? (-1. / epsilon(CR_P_trunc)) : (-1. / CR_P_trunc + 1);\n compute_Asc_nnzPerRow_kernel2\n <<< numBlocks, blocksize>>>(ArowOffsets_ptr, AcolInd_ptr,\n Avalues_ptr, AdiagValues.raw(),\n AnumRows, cf_map_ptr,\n Asc_nnzPerRow.raw(), Asc_thresh_per_row.raw(), Asc_thresh);\n cudaCheckError();\n#else\n compute_Asc_nnzPerRow_kernel\n <<< numBlocks, blocksize>>>(ArowOffsets_ptr, AcolInd_ptr,\n Avalues_ptr, AdiagValues.raw(),\n AnumRows, cf_map_ptr,\n Asc_nnzPerRow.raw(), 0.1);\n cudaCheckError();\n#endif\n cudaCheckError();\n // get the offsets in Asc with an inclusive scan\n thrust_wrapper::inclusive_scan(Asc_nnzPerRow.begin(), Asc_nnzPerRow.end(), Asc_nnzPerRow.begin());\n cudaCheckError();\n // get total num of non-zeros in P\n const int Asc_nnz = Asc_nnzPerRow[AnumRows - 1];\n // Declare Asc (strong connections of A)\n Matrix_d Asc;\n // resize Asc\n Asc.resize(0, 0, 0, 1);\n Asc.addProps(CSR);\n Asc.resize(AnumRows, AnumRows, Asc_nnz, 1);\n Asc.row_offsets[0] = 0;\n // set P offsets (P column offsets or P^T row offsets)\n amgx::thrust::copy(Asc_nnzPerRow.begin(), Asc_nnzPerRow.end(), Asc.row_offsets.begin() + 1);\n cudaCheckError();\n#ifdef EXPERIMENTAL_CR\n compute_AscColInd_kernel\n <<< numBlocks, blocksize>>>(ArowOffsets_ptr, AcolInd_ptr,\n Avalues_ptr, AdiagValues.raw(),\n AnumRows, cf_map_ptr,\n Asc.row_offsets.raw(), Asc.col_indices.raw(), Asc_thresh_per_row.raw());\n cudaCheckError();\n#else\n compute_AscColInd_kernel\n <<< numBlocks, blocksize>>>(ArowOffsets_ptr, AcolInd_ptr,\n Avalues_ptr, AdiagValues.raw(),\n AnumRows, cf_map_ptr,\n Asc.row_offsets.raw(), Asc.col_indices.raw(), 0.25);\n cudaCheckError();\n#endif\n cudaCheckError();\n Asc.set_initialized(0);\n Asc.colorMatrix(cfg, \"presmoothError\");\n Asc.set_initialized(1);\n // get the coloring of matrix\n A_row_colors_ptr = Asc.getMatrixColoring().getRowColors().raw();\n m_smoother->setup(A, false);\n int levels = 1; // number of CR iterations\n\n while (levels < 5)\n {\n // Here, we normally (i.e. as is done in Brannick's matlab code) we need to smooth using Aff,\n // however, since we are using Gauss-Seidel (symmetric or not) we can instead use A and\n // simply zero out v_err at the coarse points of each iteration.\n // v_u = v_err(fine)\n get_fine_error_kernel\n <<< numBlocks, blocksize>>>(cf_map_ptr, AnumRows, v_err_ptr, v_u_ptr);\n cudaCheckError();\n // Presmooth error at the current FINE points.\n presmoothFineError( A, AdiagValues, cf_map, v_u, v_tmp, v_z,\n m_smoother, norm0, rho_thresh, pre );\n\n // check for convergence and break if it's fast enough\n if (norm0 <= alpha && levels > 1)\n {\n break;\n }\n\n // v_err(fine) = v_u;\n update_fine_zero_coarse_error_kernel\n <<< numBlocks, blocksize>>>(cf_map_ptr, AnumRows, v_err_ptr, v_u_ptr);\n cudaCheckError();\n // TODO: add aggressive coarsening here.\n // if (AggressiveCoarsening) { A2 = Al*Al; }\n // else { A2 = Al; }\n // Add nodes with lowest colors (<=levels) to the coarse set.\n update_cf_map_kernel\n <<< numBlocks, blocksize>>>(cf_map_ptr, AnumRows, levels - 1, A_row_colors_ptr);\n cudaCheckError();\n // Get the new FINE points count.\n numFine = (int) amgx::thrust::count(cf_map.begin(), cf_map.end(), (int)FINE);\n cudaCheckError();\n\n if (numFine == 0) { break; } // no fine points left, just exit\n\n levels++;\n } // end while (levels < maxit)\n} // end CR_iteration()\n\n\n/*************************************************************************\n * Implementing the CR algorithm\n ************************************************************************/\n// ----------------------------\n// specialization for device\n// ----------------------------\ntemplate \nvoid CR_Selector >\n::markCoarseFinePoints_1x1( Matrix_d &A,\n IVector &cf_map)\n{\n if (A.hasProps(DIAG)) { FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED); }\n\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // Dimension of A.\n const int AnumRows = (int) A.get_num_rows();\n\n if (AnumRows == 0) { return; }\n\n // Choose blocksize. Using 1 thread / row for now\n const int blocksize = 256;\n const int numBlocks = std::min(AMGX_GRID_MAX_SIZE, (int) ((AnumRows + blocksize - 1) / blocksize));\n const int maxit = 5; // Max num of CR iterations.\n ValueType norm0; // Stores the energy norm of the smoothed error in CR iteration.\n int *cf_map_ptr = cf_map.raw();\n // Initially all points are set to FINE - this is done in the level\n // Initial number of fine points (all points are fine).\n int numFine = AnumRows;\n // Randomly initialize error vector v_err: AnumRows x 1 vector.\n Vector_d v_err(AnumRows);\n initRandom(v_err);\n // Perform Compatible Relaxation (CR) iteration.\n CR_iteration(A, cf_map, numFine, v_err, norm0, maxit);\n int numCoarse = AnumRows - numFine;\n printf(\"CR: numrows = %d, numCoarse = %d\\n\", AnumRows, numCoarse);\n} // end markCoarseFinePoints_1x1 (device specialization)\n\n\ntemplate \nvoid CR_SelectorBase::markCoarseFinePoints( Matrix &A,\n FVector &weights,\n const BVector &s_con,\n IVector &cf_map,\n IVector &scratch,\n int cf_map_init)\n{\n ViewType oldView = A.currentView();\n A.setView(OWNED);\n\n if (A.get_block_size() == 1)\n {\n markCoarseFinePoints_1x1(A, cf_map);\n }\n else\n {\n FatalError(\"Unsupported block size CR selector\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n}\n\n\n\n/****************************************\n * Explicit instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class CR_SelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class CR_Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace classical\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \"solvers/solver.h\"\n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n\nnamespace amgx\n{\n\nnamespace classical\n{\n\n///////////////////////////////////////////////////////////////////////////////////////////////////\n\n/*************************************************************************\n * Generate uniform random vector\n ************************************************************************/\nstruct prg\n{\n float a, b;\n\n __host__ __device__\n prg(float _a = 0.f, float _b = 1.f) : a(_a), b(_b) {};\n\n __host__ __device__\n float operator()(const unsigned int n) const\n {\n amgx::thrust::default_random_engine rng;\n amgx::thrust::uniform_real_distribution dist(a, b);\n rng.discard(n);\n return dist(rng);\n }\n};\n\n\n// Initialization of test error vector with random entries\ntemplate \nvoid initRandom(Vector &vec)\n{\n const unsigned int size = vec.size();\n amgx::thrust::counting_iterator index_sequence_begin(0);\n thrust_wrapper::transform(index_sequence_begin, index_sequence_begin + size,\n vec.begin(), prg(0.f, 1.f));\n}\n\n// Zero out the entries of a vector corresponding to the coarse points\ntemplate \n__global__\nvoid zero_coarse_points_kernel(const int *cf_map, const int numRows, ValueType *vec)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n // if point is coarse, zero out the corresponding entry of vec\n if (cf_map[tIdx] == COARSE)\n {\n vec[tIdx] = 0.;\n }\n }\n}\n\n// Create a vector which extracts fine points from the error vector and is\n// zero at the coarse points\ntemplate \n__global__\nvoid get_fine_error_kernel( const int *cf_map, const int numRows,\n const ValueType *v_err, ValueType *vec )\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n // if point is coarse, zero out the corresponding entry of vec\n if (cf_map[tIdx] == COARSE)\n {\n vec[tIdx] = 0.;\n }\n else\n {\n vec[tIdx] = v_err[tIdx];\n }\n }\n}\n\n\n// Update only the fine points of the error vector (v_err) based on input vec.\ntemplate \n__global__\nvoid update_fine_zero_coarse_error_kernel(const int *cf_map, const int numRows,\n ValueType *v_err, const ValueType *vec)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n // if point is coarse, zero out the corresponding entry of vec\n if (cf_map[tIdx] == FINE)\n {\n v_err[tIdx] = vec[tIdx];\n }\n else\n {\n v_err[tIdx] = (ValueType) 0;\n }\n }\n}\n\n\n// Update coarse and fine points from the matrix coloring taking the first\n// levels colors to form the independent set.\ntemplate \n__global__\nvoid update_cf_map_kernel(IndexType *cf_map, const int numRows, const int levels, const IndexType *A_row_colors)\n{\n //coarse = [coarse; find(independent)];\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n // go over each row\n if (A_row_colors[tIdx] <= levels) // double check that this is what we want\n {\n cf_map[tIdx] = COARSE;\n }\n }\n}\n\n\n\n#define EXPERIMENTAL_CR\n\n#ifdef EXPERIMENTAL_CR\n\n\n\n// Update only the fine points of the error vector (v_err) based on input vec.\ntemplate \n__global__\nvoid compute_Asc_nnzPerRow_kernel2(const IndexType *ArowOffsets, const IndexType *AcolInd,\n const ValueType *Avalues, const ValueType *AdiagValues,\n const IndexType AnumRows, const IndexType *cf_map,\n IndexType *Asc_nnzPerRow, ValueType *row_thresh, const ValueType thresh)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AnumRows; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on one \"fine\" row of A at a time\n const int ArowBegin = ArowOffsets[tIdx];\n const int ArowEnd = ArowOffsets[tIdx + 1];\n\n if (cf_map[tIdx] == FINE)\n {\n // work with fine pts only\n // Matlab pseudo-code:\n // Asc(:,i) = A(:,i)/A(i,i)<-thresh;\n // but this thresh is dynamic\n ValueType fine_avg_thresh = 0.;\n int fine_count = 0;\n\n for (int j = ArowBegin; j < ArowEnd; j++)\n {\n int AjColIndex = AcolInd[j];\n\n if ((AjColIndex != tIdx) && (cf_map[AjColIndex] == FINE))\n {\n fine_avg_thresh += Avalues[j];\n fine_count ++;\n }\n }\n\n fine_avg_thresh *= thresh / fine_count;\n row_thresh[tIdx] = fine_avg_thresh;\n\n for (int j = ArowBegin; j < ArowEnd; j++)\n {\n int AjColIndex = AcolInd[j];\n\n if (AjColIndex == tIdx)\n {\n /*Asc_nnzPerRow[tIdx]++;*/\n continue;\n }\n\n if (cf_map[AjColIndex] == FINE)\n {\n if (Avalues[j] / AdiagValues[tIdx] < fine_avg_thresh / AdiagValues[tIdx] || Avalues[j] / AdiagValues[AjColIndex] < fine_avg_thresh / AdiagValues[AjColIndex])\n {\n Asc_nnzPerRow[tIdx]++;\n }\n }\n }\n }\n }\n}\n\n\n\n// Update only the fine points of the error vector (v_err) based on input vec.\ntemplate \n__global__\nvoid compute_AscColInd_kernel(const IndexType *ArowOffsets, const IndexType *AcolInd,\n const ValueType *Avalues, const ValueType *AdiagValues,\n const IndexType AnumRows, const IndexType *cf_map,\n const IndexType *AscRowOffsets, IndexType *AscColInd,\n const ValueType *row_thresh)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AnumRows; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on one \"coarse\" row of A at a time\n const int ArowBegin = ArowOffsets[tIdx];\n const int ArowEnd = ArowOffsets[tIdx + 1];\n const int AscRowEnd = AscRowOffsets[tIdx + 1];\n ValueType rowMax(0);\n\n if (cf_map[tIdx] == FINE) // tIdx = j (in matlab) i.e. row of A\n {\n // work with fine pts only\n // Matlab pseudo-code:\n // Asc(:,i) = A(:,i)/A(i,i)<-thresh;\n int AscEntryNum = AscRowOffsets[tIdx]; // starting location of the current column of P\n\n // calculate max(-rowentries)\n for (int j = ArowBegin; j < ArowEnd && AscEntryNum < AscRowEnd; j++)\n {\n rowMax = max( rowMax, -Avalues[j]);\n }\n\n for (int j = ArowBegin; j < ArowEnd && AscEntryNum < AscRowEnd; j++)\n {\n // go thru each \"fine\" column of A\n int AjColIndex = AcolInd[j];\n\n if (AjColIndex == tIdx) { /*AscColInd[AscEntryNum] = AjColIndex; AscEntryNum++;*/ continue; }\n\n if (cf_map[AjColIndex] == FINE)\n {\n if (-Avalues[j] >= row_thresh[tIdx]*rowMax)\n {\n AscColInd[AscEntryNum] = AjColIndex;\n AscEntryNum++;\n }\n }\n }\n }\n }\n}\n\n#else\n\n\n// Update only the fine points of the error vector (v_err) based on input vec.\ntemplate \n__global__\nvoid compute_Asc_nnzPerRow_kernel(const IndexType *ArowOffsets, const IndexType *AcolInd,\n const ValueType *Avalues, const ValueType *AdiagValues,\n const IndexType AnumRows, const IndexType *cf_map,\n IndexType *Asc_nnzPerRow, const ValueType thresh)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AnumRows; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on one \"fine\" row of A at a time\n const int ArowBegin = ArowOffsets[tIdx];\n const int ArowEnd = ArowOffsets[tIdx + 1];\n\n if (cf_map[tIdx] == FINE)\n {\n // work with fine pts only\n // Matlab pseudo-code:\n // Asc(:,i) = A(:,i)/A(i,i)<-thresh;\n for (int j = ArowBegin; j < ArowEnd; j++)\n {\n // go thru each \"fine\" column of A\n int AjColIndex = AcolInd[j];\n\n if (AjColIndex == tIdx)\n {\n /*Asc_nnzPerRow[tIdx]++;*/\n continue;\n }\n\n if (cf_map[AjColIndex] == FINE)\n {\n if (Avalues[j] / AdiagValues[tIdx] < -thresh || Avalues[j] / AdiagValues[AjColIndex] < -thresh)\n {\n Asc_nnzPerRow[tIdx]++;\n }\n }\n }\n }\n }\n}\n\n\n\n// Update only the fine points of the error vector (v_err) based on input vec.\ntemplate \n__global__\nvoid compute_AscColInd_kernel(const IndexType *ArowOffsets, const IndexType *AcolInd,\n const ValueType *Avalues, const ValueType *AdiagValues,\n const IndexType AnumRows, const IndexType *cf_map,\n const IndexType *AscRowOffsets, IndexType *AscColInd,\n const ValueType thresh)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < AnumRows; tIdx += blockDim.x * gridDim.x)\n {\n // each thread works on one \"coarse\" row of A at a time\n const int ArowBegin = ArowOffsets[tIdx];\n const int ArowEnd = ArowOffsets[tIdx + 1];\n const int AscRowEnd = AscRowOffsets[tIdx + 1];\n\n if (cf_map[tIdx] == FINE) // tIdx = j (in matlab) i.e. row of A\n {\n // work with fine pts only\n // Matlab pseudo-code:\n // Asc(:,i) = A(:,i)/A(i,i)<-thresh;\n int AscEntryNum = AscRowOffsets[tIdx]; // starting location of the current column of P\n\n for (int j = ArowBegin; j < ArowEnd && AscEntryNum < AscRowEnd; j++)\n {\n // go thru each \"fine\" column of A\n int AjColIndex = AcolInd[j];\n\n if (AjColIndex == tIdx) { /*AscColInd[AscEntryNum] = AjColIndex; AscEntryNum++;*/ continue; }\n\n if (cf_map[AjColIndex] == FINE)\n {\n if ((Avalues[j] / AdiagValues[tIdx] < -thresh) || (Avalues[j] / AdiagValues[AjColIndex] < -thresh))\n {\n AscColInd[AscEntryNum] = AjColIndex;\n AscEntryNum++;\n }\n }\n }\n }\n }\n}\n\n#endif\n///////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nCR_Selector >::\nCR_Selector(AMG_Config &cfg, const std::string &cfg_scope) : \n CR_SelectorBase(cfg, cfg_scope),\n m_smoother(0) {}\n\n\ntemplate \nCR_Selector >::\n~CR_Selector()\n{\n if (m_smoother) { delete m_smoother; }\n}\n\n\n// Private functions for the implementation on device\n\n// Presmooth error\ntemplate \nvoid CR_Selector >\n::presmoothFineError( Matrix_d &A, const VVector &AdiagValues, const IVector &cf_map,\n Vector_d &v_u, Vector_d &v_tmp, Vector_d &v_z,\n Solver *smoother, ValueType &norm0,\n const ValueType rho_thresh, const int pre )\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Vector_d::value_type ValueTypeB;\n // Matrix A\n const IndexType AnumRows = (IndexType) A.get_num_rows();\n // Coarse/Fine map\n const int *cf_map_ptr = cf_map.raw();\n // Vector with the initial 'fine' error\n ValueTypeB *v_u_ptr = v_u.raw();\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (AnumRows / blocksize + 1));\n // Compute 'energy norm' of v_u (normalize error before smoothing)\n // v_tmp = diag(Aff).*v_u\n thrust_wrapper::transform(AdiagValues.begin(), AdiagValues.end(),\n v_u.begin(), v_tmp.begin(),\n amgx::thrust::multiplies());\n cudaCheckError();\n // norm0 = sqrt(v_tmp' * v_u)\n norm0 = sqrt( amgx::thrust::inner_product(v_u.begin(), v_u.end(),\n v_tmp.begin(), ValueTypeB(0.)) );\n cudaCheckError();\n // normalize: v_u = v_u / (norm0 + 1e-12);\n thrust_wrapper::transform(v_u.begin(), v_u.end(),\n amgx::thrust::make_constant_iterator(norm0 + 1.0e-12), v_u.begin(),\n amgx::thrust::divides());\n cudaCheckError();\n\n // ---------------------------- begin error presmoothing --------------------------\n for (int k = 1; k <= 5; k++)\n {\n // Presmooth with Symmetric Gauss-Seidel (or ilu) for pre number of sm(ooth) s(teps)\n // WARNING: zeroing out technique used here is not equivalent to the fine matrix pre-smoothing!\n for (int sms = 1; sms <= pre; sms++)\n {\n smoother->solve_iteration( v_z, v_u, false );\n zero_coarse_points_kernel\n <<< numBlocks, blocksize>>>(cf_map_ptr, AnumRows, v_u_ptr);\n cudaCheckError();\n }\n\n // -------------------- compute 'energy norm' of v_u ------------------------------\n // v_tmp = diag(Aff).*v_u\n thrust_wrapper::transform(AdiagValues.begin(), AdiagValues.end(),\n v_u.begin(), v_tmp.begin(),\n amgx::thrust::multiplies());\n cudaCheckError();\n // norm0 = sqrt(v_tmp' * v_u)\n norm0 = sqrt( amgx::thrust::inner_product(v_u.begin(), v_u.end(),\n v_tmp.begin(), ValueTypeB(0.)) );\n cudaCheckError();\n // -------------------- done with 'energy norm' of v_u ----------------------------\n // normalize: v_u = v_u / (norm0 + 1e-12);\n thrust_wrapper::transform(v_u.begin(), v_u.end(),\n amgx::thrust::make_constant_iterator(norm0 + 1.0e-12), v_u.begin(),\n amgx::thrust::divides());\n cudaCheckError();\n\n // norm0 := rho\n if (norm0 > 5 || norm0 < rho_thresh) { break; }\n }\n\n // ---------------------------- end error presmoothing ------------------------------\n}\n\n\n// A single iteration of CR while loop.\ntemplate \nvoid CR_Selector >\n::CR_iteration( Matrix_d &A, IVector &cf_map, int &numFine,\n Vector_d &v_err, ValueType &norm0,\n const ValueType maxit )\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n typedef typename Vector_d::value_type ValueTypeB;\n // Matrix A\n const IndexType AnumRows = (IndexType) A.get_num_rows();\n // Coarse/Fine map\n IndexType *cf_map_ptr = cf_map.raw();\n // Raw ptr for the row colors from the multi-color smoother\n const IndexType *A_row_colors_ptr;\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (AnumRows / blocksize + 1));\n const int pre = 10;\n const ValueType alpha = 0.7;\n const ValueType rho_thresh = 1.0e-2;\n const IndexType *ArowOffsets_ptr = A.row_offsets.raw();\n const IndexType *AcolInd_ptr = A.col_indices.raw();\n const ValueType *Avalues_ptr = A.values.raw();\n // temporary vectors - diagonal and non-zero offsets\n VVector AdiagValues(AnumRows, 0.0);\n // Extract the diagonal of A\n find_diag_kernel_indexed_dia <<< numBlocks, blocksize>>>(AnumRows, A.diag.raw(), Avalues_ptr, AdiagValues.raw());\n cudaCheckError();\n // ----------------------------------- define smoother ------------------------------------------\n // Set the smoother to Gauss-Seidel (json config format)\n AMG_Config cfg;\n std::string cfg_string = \"{ \\\"config_version\\\": 2\";\n cfg_string += \", \\\"solver\\\": { \\\"scope\\\": \\\"presmoothError\\\", \\\"solver\\\": \\\"MULTICOLOR_GS\\\"\";\n cfg_string += \" ,\\\"matrix_coloring_scheme\\\": \\\"PARALLEL_GREEDY\\\"\";\n cfg_string += \" } }\";\n cfg.parseParameterString(cfg_string.c_str());\n // Allocate smoother\n m_smoother = SolverFactory::allocate(cfg, \"default\", \"solver\");\n // --------------------------------- end define smoother ----------------------------------------\n Vector_d v_z(AnumRows, 0.0);\n Vector_d v_u(AnumRows);\n Vector_d v_tmp(v_u.size());\n ValueTypeB *v_err_ptr = v_err.raw();\n ValueTypeB *v_u_ptr = v_u.raw();\n // Store the number of non-zeros in each column of the interpolation matrix P\n IVector Asc_nnzPerRow(AnumRows, 0);\n#ifdef EXPERIMENTAL_CR\n VVector Asc_thresh_per_row(AnumRows, 0);\n ValueTypeB CR_P_trunc = 0.9; // truncation coefficient for CR P matrix, [0.0 ... 1.0]\n ValueTypeB Asc_thresh = (CR_P_trunc < epsilon(CR_P_trunc)) ? (-1. / epsilon(CR_P_trunc)) : (-1. / CR_P_trunc + 1);\n compute_Asc_nnzPerRow_kernel2\n <<< numBlocks, blocksize>>>(ArowOffsets_ptr, AcolInd_ptr,\n Avalues_ptr, AdiagValues.raw(),\n AnumRows, cf_map_ptr,\n Asc_nnzPerRow.raw(), Asc_thresh_per_row.raw(), Asc_thresh);\n cudaCheckError();\n#else\n compute_Asc_nnzPerRow_kernel\n <<< numBlocks, blocksize>>>(ArowOffsets_ptr, AcolInd_ptr,\n Avalues_ptr, AdiagValues.raw(),\n AnumRows, cf_map_ptr,\n Asc_nnzPerRow.raw(), 0.1);\n cudaCheckError();\n#endif\n cudaCheckError();\n // get the offsets in Asc with an inclusive scan\n thrust_wrapper::inclusive_scan(Asc_nnzPerRow.begin(), Asc_nnzPerRow.end(), Asc_nnzPerRow.begin());\n cudaCheckError();\n // get total num of non-zeros in P\n const int Asc_nnz = Asc_nnzPerRow[AnumRows - 1];\n // Declare Asc (strong connections of A)\n Matrix_d Asc;\n // resize Asc\n Asc.resize(0, 0, 0, 1);\n Asc.addProps(CSR);\n Asc.resize(AnumRows, AnumRows, Asc_nnz, 1);\n Asc.row_offsets[0] = 0;\n // set P offsets (P column offsets or P^T row offsets)\n amgx::thrust::copy(Asc_nnzPerRow.begin(), Asc_nnzPerRow.end(), Asc.row_offsets.begin() + 1);\n cudaCheckError();\n#ifdef EXPERIMENTAL_CR\n compute_AscColInd_kernel\n <<< numBlocks, blocksize>>>(ArowOffsets_ptr, AcolInd_ptr,\n Avalues_ptr, AdiagValues.raw(),\n AnumRows, cf_map_ptr,\n Asc.row_offsets.raw(), Asc.col_indices.raw(), Asc_thresh_per_row.raw());\n cudaCheckError();\n#else\n compute_AscColInd_kernel\n <<< numBlocks, blocksize>>>(ArowOffsets_ptr, AcolInd_ptr,\n Avalues_ptr, AdiagValues.raw(),\n AnumRows, cf_map_ptr,\n Asc.row_offsets.raw(), Asc.col_indices.raw(), 0.25);\n cudaCheckError();\n#endif\n cudaCheckError();\n Asc.set_initialized(0);\n Asc.colorMatrix(cfg, \"presmoothError\");\n Asc.set_initialized(1);\n // get the coloring of matrix\n A_row_colors_ptr = Asc.getMatrixColoring().getRowColors().raw();\n m_smoother->setup(A, false);\n int levels = 1; // number of CR iterations\n\n while (levels < 5)\n {\n // Here, we normally (i.e. as is done in Brannick's matlab code) we need to smooth using Aff,\n // however, since we are using Gauss-Seidel (symmetric or not) we can instead use A and\n // simply zero out v_err at the coarse points of each iteration.\n // v_u = v_err(fine)\n get_fine_error_kernel\n <<< numBlocks, blocksize>>>(cf_map_ptr, AnumRows, v_err_ptr, v_u_ptr);\n cudaCheckError();\n // Presmooth error at the current FINE points.\n presmoothFineError( A, AdiagValues, cf_map, v_u, v_tmp, v_z,\n m_smoother, norm0, rho_thresh, pre );\n\n // check for convergence and break if it's fast enough\n if (norm0 <= alpha && levels > 1)\n {\n break;\n }\n\n // v_err(fine) = v_u;\n update_fine_zero_coarse_error_kernel\n <<< numBlocks, blocksize>>>(cf_map_ptr, AnumRows, v_err_ptr, v_u_ptr);\n cudaCheckError();\n // TODO: add aggressive coarsening here.\n // if (AggressiveCoarsening) { A2 = Al*Al; }\n // else { A2 = Al; }\n // Add nodes with lowest colors (<=levels) to the coarse set.\n update_cf_map_kernel\n <<< numBlocks, blocksize>>>(cf_map_ptr, AnumRows, levels - 1, A_row_colors_ptr);\n cudaCheckError();\n // Get the new FINE points count.\n numFine = (int) amgx::thrust::count(cf_map.begin(), cf_map.end(), (int)FINE);\n cudaCheckError();\n\n if (numFine == 0) { break; } // no fine points left, just exit\n\n levels++;\n } // end while (levels < maxit)\n} // end CR_iteration()\n\n\n/*************************************************************************\n * Implementing the CR algorithm\n ************************************************************************/\n// ----------------------------\n// specialization for device\n// ----------------------------\ntemplate \nvoid CR_Selector >\n::markCoarseFinePoints_1x1( Matrix_d &A,\n IVector &cf_map)\n{\n if (A.hasProps(DIAG)) { FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED); }\n\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // Dimension of A.\n const int AnumRows = (int) A.get_num_rows();\n\n if (AnumRows == 0) { return; }\n\n // Choose blocksize. Using 1 thread / row for now\n const int blocksize = 256;\n const int numBlocks = std::min(AMGX_GRID_MAX_SIZE, (int) ((AnumRows + blocksize - 1) / blocksize));\n const int maxit = 5; // Max num of CR iterations.\n ValueType norm0; // Stores the energy norm of the smoothed error in CR iteration.\n int *cf_map_ptr = cf_map.raw();\n // Initially all points are set to FINE - this is done in the level\n // Initial number of fine points (all points are fine).\n int numFine = AnumRows;\n // Randomly initialize error vector v_err: AnumRows x 1 vector.\n Vector_d v_err(AnumRows);\n initRandom(v_err);\n // Perform Compatible Relaxation (CR) iteration.\n CR_iteration(A, cf_map, numFine, v_err, norm0, maxit);\n int numCoarse = AnumRows - numFine;\n printf(\"CR: numrows = %d, numCoarse = %d\\n\", AnumRows, numCoarse);\n} // end markCoarseFinePoints_1x1 (device specialization)\n\n\ntemplate \nvoid CR_SelectorBase::markCoarseFinePoints( Matrix &A,\n FVector &weights,\n const BVector &s_con,\n IVector &cf_map,\n IVector &scratch,\n int cf_map_init)\n{\n ViewType oldView = A.currentView();\n A.setView(OWNED);\n\n if (A.get_block_size() == 1)\n {\n markCoarseFinePoints_1x1(A, cf_map);\n }\n else\n {\n FatalError(\"Unsupported block size CR selector\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n}\n\n\n\n/****************************************\n * Explicit instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class CR_SelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class CR_Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace classical\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/classical/selectors/cr.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n\n#include\n\nnamespace amgx\n{\n\nnamespace classical\n{\n/*************************************************************************\n * Perform the initial marking for each row\n ************************************************************************/\ntemplate \n__global__\nvoid initialMarkingKernel(IndexType *cf_map, const IndexType numRows)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n cf_map[tIdx] = (tIdx % 2) ? COARSE : FINE;\n }\n}\n\n\n\n/*************************************************************************\n * mark coarse and fine points\n ************************************************************************/\ntemplate \nvoid Dummy_Selector >\n::markCoarseFinePoints_1x1(Matrix_h &A,\n FVector &weights,\n const BVector &s_con,\n IVector &cf_map,\n IVector &scratch,\n int cf_map_init)\n{\n FatalError(\"DummySelector not implemented on host\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid Dummy_Selector >::markCoarseFinePoints_1x1(Matrix_d &A,\n FVector &weights,\n const BVector &s_con,\n IVector &cf_map,\n IVector &scratch,\n int cf_map_init)\n{\n if (A.hasProps(DIAG))\n {\n FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // Choose blocksize. Using 1 thread / row for now\n const int blockSize = 256;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) ((A.get_num_rows() + blockSize - 1) / blockSize));\n const int numRows = (int) A.get_num_rows();\n // raw pointers from the cusp arrays\n const IndexType *offsets_ptr = A.row_offsets.raw();\n const IndexType *column_indices_ptr = A.col_indices.raw();\n const ValueType *values_ptr = A.values.raw();\n float *weights_ptr = weights.raw();\n const bool *s_con_ptr = s_con.raw();\n int *scratch_ptr = scratch.raw();\n // do the initial marking of points\n initialMarkingKernel <<< numBlocks, blockSize>>>(cf_map.raw(), A.get_num_rows());\n cudaCheckError();\n cf_map.dirtybit = 1;\n\n // Do a two ring exchange before\n // TODO: This should query the interpolator to figure out what is needed in more general case\n // By default we do 2 rings in classical\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_2ring(cf_map, cf_map.tag);\n }\n}\n\ntemplate \nvoid Dummy_SelectorBase< T_Config>::markCoarseFinePoints(Matrix< T_Config> &A,\n FVector &weights,\n const BVector &s_con,\n IVector &cf_map,\n IVector &scratch,\n int cf_map_init)\n{\n ViewType oldView = A.currentView();\n A.setView(OWNED);\n\n if (A.get_block_size() == 1)\n {\n markCoarseFinePoints_1x1(A, weights, s_con, cf_map, scratch);\n }\n else\n {\n FatalError(\"Unsupported block size Dummy selector\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n // Mark cf_map as dirty\n cf_map.dirtybit = 1;\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class Dummy_SelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Dummy_Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace classical\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n#include\n\nnamespace amgx\n{\n\nnamespace classical\n{\n/*************************************************************************\n * Perform the initial marking for each row\n ************************************************************************/\ntemplate \n__global__\nvoid initialMarkingKernel(IndexType *cf_map, const IndexType numRows)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n cf_map[tIdx] = (tIdx % 2) ? COARSE : FINE;\n }\n}\n\n\n\n/*************************************************************************\n * mark coarse and fine points\n ************************************************************************/\ntemplate \nvoid Dummy_Selector >\n::markCoarseFinePoints_1x1(Matrix_h &A,\n FVector &weights,\n const BVector &s_con,\n IVector &cf_map,\n IVector &scratch,\n int cf_map_init)\n{\n FatalError(\"DummySelector not implemented on host\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \nvoid Dummy_Selector >::markCoarseFinePoints_1x1(Matrix_d &A,\n FVector &weights,\n const BVector &s_con,\n IVector &cf_map,\n IVector &scratch,\n int cf_map_init)\n{\n if (A.hasProps(DIAG))\n {\n FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // Choose blocksize. Using 1 thread / row for now\n const int blockSize = 256;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) ((A.get_num_rows() + blockSize - 1) / blockSize));\n const int numRows = (int) A.get_num_rows();\n // raw pointers from the cusp arrays\n const IndexType *offsets_ptr = A.row_offsets.raw();\n const IndexType *column_indices_ptr = A.col_indices.raw();\n const ValueType *values_ptr = A.values.raw();\n float *weights_ptr = weights.raw();\n const bool *s_con_ptr = s_con.raw();\n int *scratch_ptr = scratch.raw();\n // do the initial marking of points\n initialMarkingKernel <<< numBlocks, blockSize>>>(cf_map.raw(), A.get_num_rows());\n cudaCheckError();\n cf_map.dirtybit = 1;\n\n // Do a two ring exchange before\n // TODO: This should query the interpolator to figure out what is needed in more general case\n // By default we do 2 rings in classical\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo_2ring(cf_map, cf_map.tag);\n }\n}\n\ntemplate \nvoid Dummy_SelectorBase< T_Config>::markCoarseFinePoints(Matrix< T_Config> &A,\n FVector &weights,\n const BVector &s_con,\n IVector &cf_map,\n IVector &scratch,\n int cf_map_init)\n{\n ViewType oldView = A.currentView();\n A.setView(OWNED);\n\n if (A.get_block_size() == 1)\n {\n markCoarseFinePoints_1x1(A, weights, s_con, cf_map, scratch);\n }\n else\n {\n FatalError(\"Unsupported block size Dummy selector\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n // Mark cf_map as dirty\n cf_map.dirtybit = 1;\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class Dummy_SelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Dummy_Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace classical\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/classical/selectors/dummy_selector.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \n\nnamespace amgx\n{\nnamespace classical\n{\nstruct is_non_neg\n{\n __host__ __device__\n bool operator()(const int &x)\n {\n return x >= 0;\n }\n};\n\n__global__\nvoid coarseMarkKernel(int *cf_map, int *mark, const int numEntries)\n{\n for (int tIdx = threadIdx.x + blockIdx.x * blockDim.x; tIdx < numEntries; tIdx += gridDim.x * blockDim.x)\n {\n if (cf_map[tIdx] == COARSE)\n {\n mark[tIdx] = 1;\n }\n }\n}\n\n__global__\nvoid modifyCoarseMapKernel(int *cf_map, int *mark, const int numEntries)\n{\n for (int tIdx = threadIdx.x + blockIdx.x * blockDim.x; tIdx < numEntries; tIdx += gridDim.x * blockDim.x)\n {\n // entry mark[i] is 0 if fine, desired #+1 if coarse\n if (cf_map[tIdx] == COARSE)\n {\n cf_map[tIdx] = mark[tIdx] - 1;\n }\n }\n}\n\n/*\n * Count the # of non-zeros per row\n */\n__global__\nvoid nonZerosPerCoarseRowKernel(const int num_rows, const int *cf_map, const int *C_hat_start,\n const int *C_hat_end, int *nonZerosPerCoarseRow)\n{\n for (int tIdx = threadIdx.x + blockIdx.x * blockDim.x; tIdx < num_rows; tIdx += gridDim.x * blockDim.x)\n {\n // If coarse point\n if (cf_map[tIdx] >= 0)\n {\n nonZerosPerCoarseRow[cf_map[tIdx]] = C_hat_end[tIdx] - C_hat_start[tIdx];\n }\n }\n}\n\n\ntemplate \n__global__\nvoid resolve_boundary(const IndexType *offsets, const IndexType *column_indices, bool *s_con,\n const int *cf_map_in, const int base_offset, const int numRows)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n if (cf_map_in[base_offset + tIdx] == COARSE)\n {\n // for each column\n for (int j = offset; j < offset + numj; j++)\n {\n if (!s_con[j]) { continue; } // skip weakly connected points\n\n int jcol = column_indices[j];\n\n if (cf_map_in[jcol] == COARSE) // coarse neighbour\n {\n s_con[j] = false; //make it weak\n }\n }\n }\n }\n}\n\nnamespace selector\n{\n#include \n#include // Included inside the namespace to solve name colisions.\n\n__device__ __forceinline__ int get_work( int *queue, int warp_id )\n{\n int offset = -1;\n\n if ( utils::lane_id() == 0 )\n {\n offset = atomicAdd( queue, 1 );\n }\n\n return utils::shfl( offset, 0 );\n}\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nestimate_c_hat_size_kernel( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const int *cf_map,\n const bool *s_con,\n int *C_hat_offsets )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // A shared location where threads propose a row of B to load.\n __shared__ int s_b_row_ids[CTA_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += gridDim.x * NUM_WARPS )\n {\n // Skip fine rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id < 0)\n {\n continue;\n }\n\n // Load A row IDs.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // The number of elements.\n int count(0);\n\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n __syncthreads();\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0, num_rows = __popc(vote) ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [b_col_it, b_col_end).\n int b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n // Load the range of B.\n int b_col_it = A_rows[b_row_id + 0];\n int b_col_end = A_rows[b_row_id + 1];\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any(b_col_it < b_col_end) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != b_row_id;\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_off_diagonal && s_con[b_col_it] && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n }\n }\n }\n\n // Store the number of columns in each row.\n if ( lane_id == 0 )\n {\n C_hat_offsets[a_row_id] = count;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nestimate_c_hat_size_kernel( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const int *cf_map,\n const bool *s_con,\n int *C_hat_offsets )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // A shared location where threads propose a row of B to load.\n __shared__ int s_b_row_ids[CTA_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += gridDim.x * NUM_WARPS )\n {\n // Skip fine rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id < 0)\n {\n continue;\n }\n\n // Load A row IDs.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // The number of elements.\n int count(0);\n\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n utils::syncwarp();\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0, num_rows = __popc(vote) ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [b_col_it, b_col_end).\n int b_row_id = -1;\n\n if ( is_active_k )\n {\n b_row_id = s_b_row_ids[warp_id * WARP_SIZE + local_k];\n }\n\n // Load the range of B.\n int b_col_it = 0, b_col_end = 0;\n\n if ( is_active_k )\n {\n b_col_it = A_rows[b_row_id + 0];\n b_col_end = A_rows[b_row_id + 1];\n }\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != b_row_id;\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_off_diagonal && s_con[b_col_it] && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n }\n }\n }\n\n // Store the number of columns in each row.\n if ( lane_id == 0 )\n {\n C_hat_offsets[a_row_id] = count;\n }\n }\n}\n\n\ntemplate< int CTA_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nestimate_c_hat_size_kernel_opt( const int nunfine,\n const int *A_rows,\n const int *A_cols,\n const int *cf_map,\n const bool *s_con,\n int *C_hat_offsets,\n int* unfine_set )\n{\n // First threads load the row IDs of A needed by the CTA...\n int gid = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Loop over rows of A.\n if (gid >= nunfine) return;\n\n int a_row_id = unfine_set[gid];\n\n int count = 0;\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = A_rows[a_row_id ]; a_col_it < A_rows[a_row_id + 1]; ++a_col_it )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = A_cols[a_col_it];\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n\n if(is_coarse_strongly_connected) \n {\n count++;\n }\n\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n\n if (!is_fine_strongly_connected ) continue;\n\n int b_col_beg = A_rows[a_col_id];\n\n // _iterate over the range of columns of B.\n for (int b_col_it = b_col_beg; b_col_it < A_rows[a_col_id+1]; ++b_col_it)\n {\n // The ID of the column.\n int b_col_id = A_cols[b_col_it];\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_id != a_col_id;\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_off_diagonal && s_con[b_col_it] && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n if(is_coarse_strongly_connected)\n {\n count++;\n }\n }\n }\n\n C_hat_offsets[a_row_id] = count;\n}\n\ntemplate< int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_c_hat_kernel( int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_end,\n int *__restrict C_hat,\n int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Shared memory to vote.\n __shared__ int s_b_row_ids[CTA_SIZE];\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Skip fine rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id < 0)\n {\n continue;\n }\n\n // Clear the set.\n set.clear();\n // Load the range of the row.\n __syncthreads();\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_begin = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n __syncthreads();\n\n // _iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = is_active && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n int item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n item = a_col_id;\n }\n\n set.insert( item, wk_status );\n //if (a_row_id == 7) printf(\"inserting %d\\n\",item);\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n utils::syncwarp();\n\n int num_rows = __popc( vote );\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int uniform_b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n // Load the range of the row of B.\n int b_col_it = A_rows[uniform_b_row_id + 0];\n int b_col_end = A_rows[uniform_b_row_id + 1];\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any(b_col_it < b_col_end) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != uniform_b_row_id;\n // Is it a strongly connected node.\n is_strongly_connected = is_off_diagonal && s_con[b_col_it];\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_strongly_connected && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n int b_item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n b_item = b_col_id;\n }\n\n set.insert( b_item, wk_status );\n //if (a_row_id == 7) printf(\"inserting %d\\n\",b_item);\n }\n }\n }\n\n int c_col_it = C_hat_start[a_row_id];\n int count = set.store(&C_hat[c_col_it]);\n\n if ( lane_id == 0 )\n {\n C_hat_end[a_row_id] = c_col_it + count;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_c_hat_kernel( int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_end,\n int *__restrict C_hat,\n int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // Shared memory to vote.\n __shared__ int s_b_row_ids[CTA_SIZE];\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Skip fine rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id < 0)\n {\n continue;\n }\n\n // Clear the set.\n set.clear();\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_begin = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // _iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = is_active && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n int item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n item = a_col_id;\n }\n\n set.insert( item, wk_status );\n //if (a_row_id == 7) printf(\"inserting item %d, is_off_diagonal = %d, \\n\",item,is_off_diagonal);\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n utils::syncwarp();\n\n int num_rows = __popc( vote );\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int uniform_b_row_id = -1;\n\n if ( is_active_k )\n {\n uniform_b_row_id = s_b_row_ids[warp_id * WARP_SIZE + local_k];\n }\n\n // Load the range of the row of B.\n int b_col_it = 0, b_col_end = 0;\n\n if ( is_active_k )\n {\n b_col_it = A_rows[uniform_b_row_id + 0];\n b_col_end = A_rows[uniform_b_row_id + 1];\n }\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != uniform_b_row_id && b_col_id != a_row_id;\n // Is it a strongly connected node.\n is_strongly_connected = is_off_diagonal && s_con[b_col_it];\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_strongly_connected && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n int b_item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n b_item = b_col_id;\n }\n\n set.insert( b_item, wk_status );\n //if (a_row_id == 7) printf(\"inserting b_item %d, is_off_diagonal = %d, \\n\",b_item,is_off_diagonal);\n }\n }\n }\n\n int c_col_it = C_hat_start[a_row_id];\n int count = set.store(&C_hat[c_col_it]) ;\n\n if ( lane_id == 0 )\n {\n C_hat_end[a_row_id] = c_col_it + count;\n }\n }\n}\n\ntemplate \n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_c_hat_kernel_opt( int unfine_set_size,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const bool *__restrict s_con,\n const int *__restrict C_hat_start,\n int *__restrict C_hat,\n int *__restrict C_hat_end,\n int *__restrict cf_map,\n int *__restrict unfine_set)\n{\n constexpr int SLOT_VACANT = -1;\n\n // Group indices\n constexpr int ngroups = CTA_SIZE / GROUP_SIZE;\n const int group_id = threadIdx.x / GROUP_SIZE;\n const int lane_id = threadIdx.x % GROUP_SIZE;\n\n const int gid = blockIdx.x * ngroups + group_id;\n\n // Dynamic sized shared memory\n extern __shared__ int s[];\n\n // Hash container storage\n int* key_s = (int*)s;\n int* key_group_s = &key_s[group_id*HASH_SIZE];\n\n // Initialise the keys and values.\n#pragma unroll\n for(int i = threadIdx.x; i < ngroups*HASH_SIZE; i += CTA_SIZE)\n {\n key_s[i] = SLOT_VACANT; // Inserted keys will be in range [0,N]\n }\n\n int* count_s = (int*)&key_s[ngroups*HASH_SIZE];\n if(lane_id == 0)\n {\n count_s[group_id] = 0;\n }\n\n __syncthreads();\n\n int a_row_id = -1;\n\n if(gid < unfine_set_size)\n {\n // One row of A per group\n a_row_id = unfine_set[gid];\n\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = A_rows[a_row_id] + lane_id ; a_col_it < A_rows[a_row_id + 1]; a_col_it += GROUP_SIZE )\n {\n int a_col_id = A_cols[a_col_it];\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_id != a_row_id;\n\n // Check if we have a strong connection to the point\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n\n if(is_coarse_strongly_connected) \n {\n int hash = a_col_id % HASH_SIZE;\n\n while(true)\n {\n // XXX You can just avoid this and directly atomicCAS, which is faster?\n int curr_key = key_group_s[hash];\n if(curr_key == a_col_id) break;\n\n if(curr_key == SLOT_VACANT)\n {\n int old_key = atomicCAS(&key_group_s[hash], SLOT_VACANT, a_col_id);\n if(old_key == SLOT_VACANT || old_key == a_col_id)\n {\n break;\n }\n }\n\n // We did not secure a slot, so linear probe to next slot\n hash = (hash + 1) % HASH_SIZE;\n }\n }\n\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n\n if(!is_fine_strongly_connected) continue;\n\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int b_col_beg = A_rows[a_col_id];\n\n // Iterate over the range of columns of B.\n for (int b_col_it = b_col_beg; b_col_it < A_rows[a_col_id+1]; ++b_col_it)\n {\n int b_col_id = A_cols[b_col_it];\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_id != a_col_id && b_col_id != a_row_id;\n // Is it a strongly connected node.\n is_strongly_connected = is_off_diagonal && s_con[b_col_it];\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_strongly_connected && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n\n if(!is_coarse_strongly_connected) continue;\n\n int hash = b_col_id % HASH_SIZE;\n\n while(true)\n {\n // XXX You can just avoid this and directly atomicCAS, which is faster?\n int curr_key = key_group_s[hash];\n if(curr_key == b_col_id) break;\n \n if(curr_key == SLOT_VACANT)\n {\n int old_key = atomicCAS(&key_group_s[hash], SLOT_VACANT, b_col_id);\n if(old_key == SLOT_VACANT || old_key == b_col_id)\n {\n break;\n }\n }\n\n // We did not secure a slot, so linear probe to next slot\n hash = (hash + 1) % HASH_SIZE;\n }\n }\n }\n }\n\n __syncthreads();\n\n // XXX Could change to poll by ballot style compress\n if(gid < unfine_set_size)\n {\n // Store the results.\n#pragma unroll\n for(int i = lane_id; i < HASH_SIZE; i += GROUP_SIZE)\n {\n int key = key_group_s[i];\n if(key != SLOT_VACANT)\n {\n int ind = atomicAdd(&count_s[group_id], 1);\n C_hat[C_hat_start[a_row_id] + ind] = key;\n }\n }\n }\n\n __syncthreads();\n\n if(gid < unfine_set_size)\n {\n // Store the results.\n if(lane_id == 0)\n {\n C_hat_end[a_row_id] = C_hat_start[a_row_id] + count_s[group_id];\n }\n }\n}\n} // namespace selector\n\ntemplate \nvoid __global__ createCfMapGlobal(const IndexType *cf_map, int64_t *cf_map_global, const int64_t my_part_offset, const IndexType num_owned_fine_pts)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_owned_fine_pts ; tidx += blockDim.x * gridDim.x)\n {\n // Only renumber the interior points\n if (cf_map[tidx] >= 0 && tidx < num_owned_fine_pts)\n {\n // Simply shift\n cf_map_global[tidx] = (int64_t) cf_map[tidx] /* my_local_id */ + my_part_offset;\n }\n else\n {\n cf_map_global[tidx] = -1;\n }\n }\n}\n\n#include \n\nenum { WARP_SIZE = 32, GRID_SIZE = 1024, SMEM_SIZE = 128 };\n\ntemplate< typename Value_type, int CTA_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nfillS2ColIndices( const int A_num_rows,\n const int *__restrict cf_map,\n const int *__restrict C_hat,\n const int *__restrict C_hat_start,\n const int *__restrict C_hat_end,\n const int *__restrict S2_rows,\n int *S2_cols,\n int64_t *S2_global_cols,\n int64_t *cf_map_global)\n{\n const int NUM_WARPS = CTA_SIZE / 32;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows; a_row_id += NUM_WARPS_PER_GRID )\n {\n int coarse_fine_id = cf_map[a_row_id];\n\n // Skip fine rows.\n if ( coarse_fine_id < 0 )\n {\n continue;\n }\n\n // Rebuild C_hat.\n int c_hat_it = C_hat_start[a_row_id];\n int c_hat_end = C_hat_end [a_row_id];\n int s2_col_it = S2_rows[coarse_fine_id];\n\n // Iterate over the columns of A.\n for ( c_hat_it += lane_id ; utils::any( c_hat_it < c_hat_end ) ; c_hat_it += WARP_SIZE, s2_col_it += WARP_SIZE)\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int s2_col_id(-1);\n\n if ( c_hat_it < c_hat_end )\n {\n s2_col_id = C_hat[c_hat_it];\n }\n\n if (s2_col_id != -1)\n {\n S2_cols[s2_col_it + lane_id] = cf_map[s2_col_id];\n\n if (S2_global_cols != NULL)\n {\n S2_global_cols[s2_col_it + lane_id] = cf_map_global[s2_col_id];\n }\n }\n }\n }\n}\n\n\n\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid Selector >::renumberAndCountCoarsePoints( IVector &cf_map,\n int &num_coarse_points,\n int num_rows)\n{\n IVector mark(cf_map.size(), 0);\n int blockSize = 128;\n int num_blocks = std::min( 4096, (int) (cf_map.size() + blockSize - 1) / blockSize );\n\n if (cf_map.size() > 0)\n {\n coarseMarkKernel <<< num_blocks, blockSize>>>(cf_map.raw(), mark.raw(), cf_map.size());\n cudaCheckError();\n }\n\n // get the sequence of values\n thrust_wrapper::inclusive_scan(mark.begin(), mark.end(), mark.begin());\n cudaCheckError();\n\n // assign to cf_map\n if (cf_map.size() > 0)\n {\n modifyCoarseMapKernel <<< num_blocks, blockSize>>>(cf_map.raw(), mark.raw(), cf_map.size());\n cudaCheckError();\n }\n\n\n num_coarse_points = (int) thrust_wrapper::count_if(cf_map.begin(), cf_map.begin() + num_rows, is_non_neg());\n cudaCheckError();\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid Selector >::renumberAndCountCoarsePoints( IVector &cf_map,\n int &num_coarse_points,\n int num_rows)\n{\n num_coarse_points = 0;\n\n for (int i = 0; i < num_rows; i++)\n {\n // if coarse\n if (cf_map[i] == COARSE)\n {\n cf_map[i] = num_coarse_points++;\n }\n }\n}\n\n__global__\nvoid correctCfMapKernel(int *cf_map, const int *cf_map_scanned, const int *cf_map_s2, int size)\n{\n for (int tIdx = threadIdx.x + blockIdx.x * blockDim.x; tIdx < size; tIdx += gridDim.x * blockDim.x)\n {\n if (cf_map[tIdx] == COARSE) // Coarse point\n {\n int coarse_id = cf_map_scanned[tIdx];\n int coarse_id_s2 = cf_map_s2[coarse_id];\n // if cf_map_s2 is STRONG_FINE, mark as COARSE\n cf_map[tIdx] = (coarse_id_s2 == STRONG_FINE) ? COARSE : coarse_id_s2;\n }\n }\n}\n\ntemplate \nvoid Selector >\n::correctCfMap(IVector &cf_map, IVector &cf_map_scanned, IVector &cf_map_S2)\n{\n FatalError(\"CorrectCFMap not implemented on host\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n\ntemplate \nvoid Selector >::correctCfMap(IVector &cf_map, IVector &cf_map_scanned, IVector &cf_map_S2)\n{\n // One thread per row of A\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // get raw pointers to data\n const int cta_size = 128;\n const int grid_size = std::min( 4096, ((int) cf_map.size() + cta_size - 1) / cta_size);\n correctCfMapKernel <<< grid_size, cta_size>>>(cf_map.raw(), cf_map_scanned.raw(), cf_map_S2.raw(), cf_map.size());\n cudaCheckError();\n}\n\ntemplate \nvoid Selector >\n::createS2(Matrix_h &A,\n Matrix_h &S2,\n const BVector &s_con,\n IVector &cf_map)\n{\n FatalError(\"Aggressive PMIS selector not implemented on host\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \ntemplate \nvoid Selector>::compute_c_hat_opt_dispatch(\n const Matrix_d &A,\n const bool *s_con,\n const int *C_hat_start,\n int *C_hat,\n int *C_hat_end,\n int *cf_map,\n IntVector &unfine_set)\n{\n typedef typename IndPrecisionMap::Type Index_type;\n\n constexpr int cta_size = 128;\n constexpr int ngroups = cta_size / group_size;\n\n const int grid_size = unfine_set.size() / ngroups + 1;\n\n cudaDeviceProp deviceProps = getDeviceProperties();\n size_t max_shmem_size = deviceProps.sharedMemPerMultiprocessor;\n\n constexpr int shmem_size =\n (sizeof(int)+sizeof(Index_type))*ngroups*hash_size + group_size; \n\n if(shmem_size > max_shmem_size) \n { \n FatalError(\"In compute_values_opt the requested hash size is larger than max.\\n\", \n AMGX_ERR_NOT_IMPLEMENTED); \n } \n\n cudaFuncSetAttribute(amgx::classical::selector::compute_c_hat_kernel_opt \n , \n cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size); \n\n amgx::classical::selector::compute_c_hat_kernel_opt\n \n <<>>(\n unfine_set.size(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n s_con,\n C_hat_start,\n C_hat,\n C_hat_end,\n cf_map,\n unfine_set.raw());\n cudaCheckError();\n}\n\n// Determines which nodes are not fine, hence excluded\n__global__ void set_unfine(int n, int* cf_map, int* unfine, int* unfine_offs)\n{\n int i = blockIdx.x*blockDim.x + threadIdx.x;\n if(i >= n) return;\n\n unfine[i] = (cf_map[i] >= 0) ? 1 : 0;\n\n if(i == 0) \n {\n unfine_offs[0] = 0;\n }\n}\n\n// Compresses indirection\n__global__ void unfine_set_fill(int n, int* unfine, int* unfine_offs, int* unfine_set)\n{\n int i = blockIdx.x*blockDim.x + threadIdx.x;\n if(i >= n) return;\n\n if(unfine[i]) \n { \n unfine_set[unfine_offs[i]] = i; \n } \n}\n\n\ntemplate \nvoid Selector >::createS2(Matrix_d &A,\n Matrix_d &S2,\n const BVector &s_con,\n IVector &cf_map)\n{\n const int blockSize = 256;\n // One thread per row of A\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // get raw pointers to data\n const IndexType *Aoffsets = A.row_offsets.raw();\n const IndexType *Acolumn_indices = A.col_indices.raw();\n const ValueType *Avalues = A.values.raw();\n const IndexType Anum_rows = (int) A.get_num_rows();\n int *cf_map_ptr = cf_map.raw();\n Hash_Workspace exp_wk(true, GRID_SIZE);\n IntVector C_hat_start( A.get_num_rows() + 1, 0 ), C_hat_end( A.get_num_rows() + 1, 0 );\n\n int max_c_hat_size = 0;\n int nunfine = 0;\n IntVector unfine_set;\n\n if(this->m_use_opt_kernels)\n {\n IntVector unfine(A.get_num_rows());\n IntVector unfine_offs(A.get_num_rows()+1);\n constexpr int nthreads = 128;\n const int nblocks = A.get_num_rows() / nthreads + 1;\n set_unfine<<>>(A.get_num_rows(), cf_map.raw(), unfine.raw(), unfine_offs.raw());\n cudaCheckError();\n thrust_wrapper::inclusive_scan(unfine.begin(), unfine.end(), unfine_offs.begin()+1);\n\n nunfine = unfine_offs[unfine_offs.size()-1];\n if(nunfine > 0)\n {\n unfine_set.resize(nunfine);\n unfine_set_fill<<>>(A.get_num_rows(), unfine.raw(), unfine_offs.raw(), unfine_set.raw());\n cudaCheckError();\n }\n\n const int CTA_SIZE = 128;\n int grid_size = nunfine/CTA_SIZE + 1;\n\n amgx::classical::selector::estimate_c_hat_size_kernel_opt< CTA_SIZE > <<< grid_size, CTA_SIZE>>>(\n nunfine,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw(),\n unfine_set.raw());\n cudaCheckError();\n\n max_c_hat_size = thrust::reduce(C_hat_start.begin(), C_hat_start.end(), -1, thrust::maximum());\n }\n else\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int avg_nz_per_row = A.get_num_nz() / A.get_num_rows();\n int grid_size = A.get_num_rows()/NUM_WARPS + 1;\n\n if ( avg_nz_per_row < 16 )\n {\n amgx::classical::selector::estimate_c_hat_size_kernel< 8, CTA_SIZE, WARP_SIZE> <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw());\n cudaCheckError();\n }\n else\n {\n amgx::classical::selector::estimate_c_hat_size_kernel <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw());\n }\n\n cudaCheckError();\n }\n\n // Compute row offsets.\n thrust_wrapper::exclusive_scan(C_hat_start.begin(), C_hat_start.end(), C_hat_start.begin());\n cudaCheckError();\n int nVals = C_hat_start[C_hat_start.size() - 1];\n IntVector C_hat( nVals );\n\n if(this->m_use_opt_kernels)\n {\n int hash_size = pow(2, ceil(log2(max_c_hat_size)));\n\n if(nunfine > 0)\n {\n switch(hash_size)\n {\n case 1:\n case 2:\n case 4:\n case 8:\n case 16:\n case 32:\n compute_c_hat_opt_dispatch<32, 8>(A, s_con.raw(), C_hat_start.raw(), C_hat.raw(), C_hat_end.raw(), cf_map.raw(), unfine_set);\n break;\n case 64:\n compute_c_hat_opt_dispatch<64, 8>(A, s_con.raw(), C_hat_start.raw(), C_hat.raw(), C_hat_end.raw(), cf_map.raw(), unfine_set);\n break;\n case 128:\n compute_c_hat_opt_dispatch<128, 8>(A, s_con.raw(), C_hat_start.raw(), C_hat.raw(), C_hat_end.raw(), cf_map.raw(), unfine_set);\n break;\n case 256:\n compute_c_hat_opt_dispatch<256, 16>(A, s_con.raw(), C_hat_start.raw(), C_hat.raw(), C_hat_end.raw(), cf_map.raw(), unfine_set);\n break;\n case 512:\n compute_c_hat_opt_dispatch<512, 32>(A, s_con.raw(), C_hat_start.raw(), C_hat.raw(), C_hat_end.raw(), cf_map.raw(), unfine_set);\n break;\n default:\n FatalError(\"Size of C_hat in a single row too large, solve with original solvers.\", AMGX_ERR_INTERNAL);\n }\n }\n }\n else\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int avg_nz_per_row = A.get_num_nz() / A.get_num_rows();\n int attempt = 0;\n\n for ( bool done = false ; !done && attempt < 10 ; ++attempt )\n {\n // Double the amount of GMEM (if needed).\n if ( attempt > 0 )\n {\n exp_wk.expand();\n }\n\n // Reset the status. \n // TODO: Optimize with async copies.\n int status = 0;\n cudaMemcpy( exp_wk.get_status(), &status, sizeof(int), cudaMemcpyHostToDevice );\n cudaCheckError();\n // Compute the set C_hat.\n int work_offset = GRID_SIZE * NUM_WARPS;\n cudaMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), cudaMemcpyHostToDevice );\n cudaCheckError();\n\n // Run the computation.\n if ( avg_nz_per_row < 16 )\n {\n amgx::classical::selector::compute_c_hat_kernel< 8, CTA_SIZE, SMEM_SIZE, WARP_SIZE> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n C_hat.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_work_queue(),\n exp_wk.get_status() );\n cudaCheckError();\n }\n else\n {\n amgx::classical::selector::compute_c_hat_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n C_hat.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_work_queue(),\n exp_wk.get_status() );\n }\n\n cudaCheckError();\n // Read the result from count_non_zeroes.\n cudaMemcpy( &status, exp_wk.get_status(), sizeof(int), cudaMemcpyDeviceToHost );\n cudaCheckError();\n done = status == 0;\n }\n }\n // get pointers to data\n int *C_hat_ptr = C_hat.raw();\n int *C_hat_start_ptr = C_hat_start.raw();\n int *C_hat_end_ptr = C_hat_end.raw();\n int S2_num_rows;;\n DistributedArranger *prep = NULL;\n\n if (A.is_matrix_distributed())\n {\n prep = new DistributedArranger;\n int num_owned_fine_pts = A.get_num_rows();\n int num_owned_coarse_pts = thrust_wrapper::count_if(cf_map.begin(), cf_map.begin() + num_owned_fine_pts, is_non_neg());\n int num_halo_coarse_pts = thrust_wrapper::count_if(cf_map.begin() + num_owned_fine_pts, cf_map.end(), is_non_neg());\n cudaCheckError();\n S2_num_rows = num_owned_coarse_pts;\n prep->initialize_manager(A, S2, num_owned_coarse_pts);\n }\n else\n {\n S2_num_rows = (int) thrust_wrapper::count_if(cf_map.begin(), cf_map.end(), is_non_neg());\n cudaCheckError();\n }\n\n // count the number of non-zeros in the interpolation matrix\n int numBlocks = (int) (A.get_num_rows() + blockSize - 1) / blockSize;\n IntVector nonZeroOffsets(S2_num_rows + 1);\n IntVector nonZerosPerRow(S2_num_rows);\n // Updating the number of nonZeros for your own coarse rows\n int *nonZerosPerRow_ptr = nonZerosPerRow.raw();\n nonZerosPerCoarseRowKernel <<< numBlocks, blockSize>>>(Anum_rows, cf_map.raw(), C_hat_start_ptr, C_hat_end_ptr,\n nonZerosPerRow_ptr);\n cudaCheckError();\n // get total with a reduction\n int nonZeros = thrust_wrapper::reduce(nonZerosPerRow.begin(), nonZerosPerRow.end());\n cudaCheckError();\n // get the offsets with an exclusive scan\n thrust_wrapper::exclusive_scan(nonZerosPerRow.begin(), nonZerosPerRow.end(), nonZeroOffsets.begin());\n cudaCheckError();\n nonZeroOffsets[S2_num_rows] = nonZeros;\n // resize S2\n S2.resize(0, 0, 0, 1);\n S2.addProps(CSR);\n S2.set_num_rows(S2_num_rows);\n S2.set_num_cols(S2_num_rows);\n S2.set_num_nz(nonZeros);\n S2.row_offsets.resize(S2_num_rows + 1);\n S2.col_indices.resize(nonZeros);\n S2.diag.resize(S2_num_rows);\n S2.set_block_dimx(A.get_block_dimx());\n S2.set_block_dimy(A.get_block_dimy());\n cudaCheckError();\n // copy nonzero offsets to the P matrix\n amgx::thrust::copy(nonZeroOffsets.begin(), nonZeroOffsets.end(), S2.row_offsets.begin());\n cudaCheckError();\n {\n const int CTA_SIZE = 256;\n // Run the computation.\n typedef typename MatPrecisionMap::Type Value_type;\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n if (!A.is_matrix_distributed())\n {\n int grid_size = (A.get_num_rows() / NUM_WARPS) + 1;\n\n fillS2ColIndices <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n cf_map.raw(),\n C_hat.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n S2.row_offsets.raw(),\n S2.col_indices.raw(),\n NULL,\n NULL);\n cudaCheckError();\n }\n else\n {\n I64Vector_d cf_map_global(cf_map.size());\n int num_owned_fine_pts = A.get_num_rows();\n int my_rank = A.manager->global_id();\n const int cta_size = 128;\n int grid_size = (num_owned_fine_pts + cta_size - 1) / cta_size;\n createCfMapGlobal <<< grid_size, cta_size>>>(cf_map.raw(), cf_map_global.raw(), S2.manager->part_offsets_h[my_rank], num_owned_fine_pts);\n cudaCheckError();\n\n // Exchange the cf_map_global so that we know the coarse global id of halo nodes\n cf_map_global.dirtybit = 1;\n A.manager->exchange_halo_2ring(cf_map_global, cf_map_global.tag);\n I64Vector_d S2_col_indices_global(S2.col_indices.size());\n\n grid_size = (A.get_num_rows() / NUM_WARPS) + 1;\n fillS2ColIndices <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n cf_map.raw(),\n C_hat.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n S2.row_offsets.raw(),\n S2.col_indices.raw(),\n S2_col_indices_global.raw(),\n cf_map_global.raw());\n cudaCheckError();\n prep->initialize_manager_from_global_col_indices(S2, S2_col_indices_global);\n }\n\n if (prep != NULL) { delete prep; }\n\n cudaCheckError();\n }\n}\n\n\n\ntemplate \nvoid Selector >::demoteStrongEdges(const Matrix_h &A,\n const FVector &weights,\n BVector &s_con,\n const IVector &cf_map,\n const IndexType base_offset)\n{\n for ( int tIdx = 0; tIdx < A.get_num_rows() - base_offset; tIdx++)\n {\n const int offset = A.row_offsets[tIdx];\n const int numj = A.row_offsets[tIdx + 1] - offset;\n\n if (cf_map[base_offset + tIdx] == COARSE)\n {\n // for each column\n for (int j = offset; j < offset + numj; j++)\n {\n if (!s_con[j]) { continue; } // skip weakly connected points\n\n int jcol = A.col_indices[j];\n\n if (cf_map[jcol] == COARSE) // coarse neighbour\n {\n s_con[j] = false; //make it weak\n }\n }\n }\n }\n}\n\ntemplate \nvoid Selector >::demoteStrongEdges(const Matrix_d &A,\n const FVector &weights,\n BVector &s_con,\n const IVector &cf_map,\n const IndexType offset)\n{\n const int blockSize = 256;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (((A.get_num_rows() - offset) + blockSize - 1) / blockSize));\n const int numRows = (int)(A.get_num_rows() - offset);\n resolve_boundary <<< blockSize, numBlocks>>>(A.row_offsets.raw() + offset, A.col_indices.raw(), s_con.raw(), cf_map.raw(), offset, A.get_num_rows() - offset);\n cudaCheckError();\n}\n\n\n\ntemplate\nstd::map*> &\nSelectorFactory::getFactories( )\n{\n static std::map*> s_factories;\n return s_factories;\n}\n\ntemplate\nvoid SelectorFactory::registerFactory(std::string name, SelectorFactory *f)\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::const_iterator it = factories.find(name);\n\n if (it != factories.end())\n {\n std::string error = \"SelectorFactory '\" + name + \"' has already been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n factories[name] = f;\n}\n\ntemplate\nvoid SelectorFactory::unregisterFactory(std::string name)\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.find(name);\n\n if (it == factories.end())\n {\n std::string error = \"SelectorFactory '\" + name + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n SelectorFactory *factory = it->second;\n assert( factory != NULL );\n delete factory;\n factories.erase(it);\n}\n\ntemplate\nvoid SelectorFactory::unregisterFactories( )\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.begin( );\n\n for ( ; it != factories.end( ) ; )\n {\n SelectorFactory *factory = it->second;\n assert( factory != NULL );\n it++;\n delete factory;\n }\n\n factories.clear( );\n}\n\ntemplate\nSelector *SelectorFactory::allocate(AMG_Config &cfg, const std::string &cfg_scope)\n{\n std::map*> &factories = getFactories( );\n std::string selector = cfg.getParameter(\"selector\", cfg_scope);\n typename std::map *>::const_iterator it = factories.find(selector);\n\n if (it == factories.end())\n {\n std::string error = \"SelectorFactory '\" + selector + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n return it->second->create(cfg, cfg_scope);\n};\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class SelectorFactory::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n} // namespace classical\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \n\nnamespace amgx\n{\nnamespace classical\n{\nstruct is_non_neg\n{\n __host__ __device__\n bool operator()(const int &x)\n {\n return x >= 0;\n }\n};\n\n__global__\nvoid coarseMarkKernel(int *cf_map, int *mark, const int numEntries)\n{\n for (int tIdx = threadIdx.x + blockIdx.x * blockDim.x; tIdx < numEntries; tIdx += gridDim.x * blockDim.x)\n {\n if (cf_map[tIdx] == COARSE)\n {\n mark[tIdx] = 1;\n }\n }\n}\n\n__global__\nvoid modifyCoarseMapKernel(int *cf_map, int *mark, const int numEntries)\n{\n for (int tIdx = threadIdx.x + blockIdx.x * blockDim.x; tIdx < numEntries; tIdx += gridDim.x * blockDim.x)\n {\n // entry mark[i] is 0 if fine, desired #+1 if coarse\n if (cf_map[tIdx] == COARSE)\n {\n cf_map[tIdx] = mark[tIdx] - 1;\n }\n }\n}\n\n/*\n * Count the # of non-zeros per row\n */\n__global__\nvoid nonZerosPerCoarseRowKernel(const int num_rows, const int *cf_map, const int *C_hat_start,\n const int *C_hat_end, int *nonZerosPerCoarseRow)\n{\n for (int tIdx = threadIdx.x + blockIdx.x * blockDim.x; tIdx < num_rows; tIdx += gridDim.x * blockDim.x)\n {\n // If coarse point\n if (cf_map[tIdx] >= 0)\n {\n nonZerosPerCoarseRow[cf_map[tIdx]] = C_hat_end[tIdx] - C_hat_start[tIdx];\n }\n }\n}\n\n\ntemplate \n__global__\nvoid resolve_boundary(const IndexType *offsets, const IndexType *column_indices, bool *s_con,\n const int *cf_map_in, const int base_offset, const int numRows)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n if (cf_map_in[base_offset + tIdx] == COARSE)\n {\n // for each column\n for (int j = offset; j < offset + numj; j++)\n {\n if (!s_con[j]) { continue; } // skip weakly connected points\n\n int jcol = column_indices[j];\n\n if (cf_map_in[jcol] == COARSE) // coarse neighbour\n {\n s_con[j] = false; //make it weak\n }\n }\n }\n }\n}\n\nnamespace selector\n{\n#include \n#include // Included inside the namespace to solve name colisions.\n\n__device__ __forceinline__ int get_work( int *queue, int warp_id )\n{\n int offset = -1;\n\n if ( utils::lane_id() == 0 )\n {\n offset = atomicAdd( queue, 1 );\n }\n\n return utils::shfl( offset, 0 );\n}\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nestimate_c_hat_size_kernel( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const int *cf_map,\n const bool *s_con,\n int *C_hat_offsets )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // A shared location where threads propose a row of B to load.\n __shared__ int s_b_row_ids[CTA_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += gridDim.x * NUM_WARPS )\n {\n // Skip fine rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id < 0)\n {\n continue;\n }\n\n // Load A row IDs.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // The number of elements.\n int count(0);\n\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n __syncthreads();\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0, num_rows = __popc(vote) ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [b_col_it, b_col_end).\n int b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n // Load the range of B.\n int b_col_it = A_rows[b_row_id + 0];\n int b_col_end = A_rows[b_row_id + 1];\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any(b_col_it < b_col_end) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != b_row_id;\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_off_diagonal && s_con[b_col_it] && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n }\n }\n }\n\n // Store the number of columns in each row.\n if ( lane_id == 0 )\n {\n C_hat_offsets[a_row_id] = count;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nestimate_c_hat_size_kernel( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const int *cf_map,\n const bool *s_con,\n int *C_hat_offsets )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // A shared location where threads propose a row of B to load.\n __shared__ int s_b_row_ids[CTA_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += gridDim.x * NUM_WARPS )\n {\n // Skip fine rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id < 0)\n {\n continue;\n }\n\n // Load A row IDs.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // The number of elements.\n int count(0);\n\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n utils::syncwarp();\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0, num_rows = __popc(vote) ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [b_col_it, b_col_end).\n int b_row_id = -1;\n\n if ( is_active_k )\n {\n b_row_id = s_b_row_ids[warp_id * WARP_SIZE + local_k];\n }\n\n // Load the range of B.\n int b_col_it = 0, b_col_end = 0;\n\n if ( is_active_k )\n {\n b_col_it = A_rows[b_row_id + 0];\n b_col_end = A_rows[b_row_id + 1];\n }\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != b_row_id;\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_off_diagonal && s_con[b_col_it] && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n }\n }\n }\n\n // Store the number of columns in each row.\n if ( lane_id == 0 )\n {\n C_hat_offsets[a_row_id] = count;\n }\n }\n}\n\n\ntemplate< int CTA_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nestimate_c_hat_size_kernel_opt( const int nunfine,\n const int *A_rows,\n const int *A_cols,\n const int *cf_map,\n const bool *s_con,\n int *C_hat_offsets,\n int* unfine_set )\n{\n // First threads load the row IDs of A needed by the CTA...\n int gid = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Loop over rows of A.\n if (gid >= nunfine) return;\n\n int a_row_id = unfine_set[gid];\n\n int count = 0;\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = A_rows[a_row_id ]; a_col_it < A_rows[a_row_id + 1]; ++a_col_it )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = A_cols[a_col_it];\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n\n if(is_coarse_strongly_connected) \n {\n count++;\n }\n\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n\n if (!is_fine_strongly_connected ) continue;\n\n int b_col_beg = A_rows[a_col_id];\n\n // _iterate over the range of columns of B.\n for (int b_col_it = b_col_beg; b_col_it < A_rows[a_col_id+1]; ++b_col_it)\n {\n // The ID of the column.\n int b_col_id = A_cols[b_col_it];\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_id != a_col_id;\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_off_diagonal && s_con[b_col_it] && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n if(is_coarse_strongly_connected)\n {\n count++;\n }\n }\n }\n\n C_hat_offsets[a_row_id] = count;\n}\n\ntemplate< int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_c_hat_kernel( int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_end,\n int *__restrict C_hat,\n int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Shared memory to vote.\n __shared__ int s_b_row_ids[CTA_SIZE];\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Skip fine rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id < 0)\n {\n continue;\n }\n\n // Clear the set.\n set.clear();\n // Load the range of the row.\n __syncthreads();\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_begin = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n __syncthreads();\n\n // _iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = is_active && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n int item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n item = a_col_id;\n }\n\n set.insert( item, wk_status );\n //if (a_row_id == 7) printf(\"inserting %d\\n\",item);\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n utils::syncwarp();\n\n int num_rows = __popc( vote );\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int uniform_b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n // Load the range of the row of B.\n int b_col_it = A_rows[uniform_b_row_id + 0];\n int b_col_end = A_rows[uniform_b_row_id + 1];\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any(b_col_it < b_col_end) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != uniform_b_row_id;\n // Is it a strongly connected node.\n is_strongly_connected = is_off_diagonal && s_con[b_col_it];\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_strongly_connected && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n int b_item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n b_item = b_col_id;\n }\n\n set.insert( b_item, wk_status );\n //if (a_row_id == 7) printf(\"inserting %d\\n\",b_item);\n }\n }\n }\n\n int c_col_it = C_hat_start[a_row_id];\n int count = set.store(&C_hat[c_col_it]);\n\n if ( lane_id == 0 )\n {\n C_hat_end[a_row_id] = c_col_it + count;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_c_hat_kernel( int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_end,\n int *__restrict C_hat,\n int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // Shared memory to vote.\n __shared__ int s_b_row_ids[CTA_SIZE];\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Skip fine rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id < 0)\n {\n continue;\n }\n\n // Clear the set.\n set.clear();\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_begin = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // _iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = is_active && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n int item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n item = a_col_id;\n }\n\n set.insert( item, wk_status );\n //if (a_row_id == 7) printf(\"inserting item %d, is_off_diagonal = %d, \\n\",item,is_off_diagonal);\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n utils::syncwarp();\n\n int num_rows = __popc( vote );\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int uniform_b_row_id = -1;\n\n if ( is_active_k )\n {\n uniform_b_row_id = s_b_row_ids[warp_id * WARP_SIZE + local_k];\n }\n\n // Load the range of the row of B.\n int b_col_it = 0, b_col_end = 0;\n\n if ( is_active_k )\n {\n b_col_it = A_rows[uniform_b_row_id + 0];\n b_col_end = A_rows[uniform_b_row_id + 1];\n }\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != uniform_b_row_id && b_col_id != a_row_id;\n // Is it a strongly connected node.\n is_strongly_connected = is_off_diagonal && s_con[b_col_it];\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_strongly_connected && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n int b_item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n b_item = b_col_id;\n }\n\n set.insert( b_item, wk_status );\n //if (a_row_id == 7) printf(\"inserting b_item %d, is_off_diagonal = %d, \\n\",b_item,is_off_diagonal);\n }\n }\n }\n\n int c_col_it = C_hat_start[a_row_id];\n int count = set.store(&C_hat[c_col_it]) ;\n\n if ( lane_id == 0 )\n {\n C_hat_end[a_row_id] = c_col_it + count;\n }\n }\n}\n\ntemplate \n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_c_hat_kernel_opt( int unfine_set_size,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const bool *__restrict s_con,\n const int *__restrict C_hat_start,\n int *__restrict C_hat,\n int *__restrict C_hat_end,\n int *__restrict cf_map,\n int *__restrict unfine_set)\n{\n constexpr int SLOT_VACANT = -1;\n\n // Group indices\n constexpr int ngroups = CTA_SIZE / GROUP_SIZE;\n const int group_id = threadIdx.x / GROUP_SIZE;\n const int lane_id = threadIdx.x % GROUP_SIZE;\n\n const int gid = blockIdx.x * ngroups + group_id;\n\n // Dynamic sized shared memory\n extern __shared__ int s[];\n\n // Hash container storage\n int* key_s = (int*)s;\n int* key_group_s = &key_s[group_id*HASH_SIZE];\n\n // Initialise the keys and values.\n#pragma unroll\n for(int i = threadIdx.x; i < ngroups*HASH_SIZE; i += CTA_SIZE)\n {\n key_s[i] = SLOT_VACANT; // Inserted keys will be in range [0,N]\n }\n\n int* count_s = (int*)&key_s[ngroups*HASH_SIZE];\n if(lane_id == 0)\n {\n count_s[group_id] = 0;\n }\n\n __syncthreads();\n\n int a_row_id = -1;\n\n if(gid < unfine_set_size)\n {\n // One row of A per group\n a_row_id = unfine_set[gid];\n\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = A_rows[a_row_id] + lane_id ; a_col_it < A_rows[a_row_id + 1]; a_col_it += GROUP_SIZE )\n {\n int a_col_id = A_cols[a_col_it];\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_id != a_row_id;\n\n // Check if we have a strong connection to the point\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n\n if(is_coarse_strongly_connected) \n {\n int hash = a_col_id % HASH_SIZE;\n\n while(true)\n {\n // XXX You can just avoid this and directly atomicCAS, which is faster?\n int curr_key = key_group_s[hash];\n if(curr_key == a_col_id) break;\n\n if(curr_key == SLOT_VACANT)\n {\n int old_key = atomicCAS(&key_group_s[hash], SLOT_VACANT, a_col_id);\n if(old_key == SLOT_VACANT || old_key == a_col_id)\n {\n break;\n }\n }\n\n // We did not secure a slot, so linear probe to next slot\n hash = (hash + 1) % HASH_SIZE;\n }\n }\n\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n\n if(!is_fine_strongly_connected) continue;\n\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int b_col_beg = A_rows[a_col_id];\n\n // Iterate over the range of columns of B.\n for (int b_col_it = b_col_beg; b_col_it < A_rows[a_col_id+1]; ++b_col_it)\n {\n int b_col_id = A_cols[b_col_it];\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_id != a_col_id && b_col_id != a_row_id;\n // Is it a strongly connected node.\n is_strongly_connected = is_off_diagonal && s_con[b_col_it];\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_strongly_connected && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n\n if(!is_coarse_strongly_connected) continue;\n\n int hash = b_col_id % HASH_SIZE;\n\n while(true)\n {\n // XXX You can just avoid this and directly atomicCAS, which is faster?\n int curr_key = key_group_s[hash];\n if(curr_key == b_col_id) break;\n \n if(curr_key == SLOT_VACANT)\n {\n int old_key = atomicCAS(&key_group_s[hash], SLOT_VACANT, b_col_id);\n if(old_key == SLOT_VACANT || old_key == b_col_id)\n {\n break;\n }\n }\n\n // We did not secure a slot, so linear probe to next slot\n hash = (hash + 1) % HASH_SIZE;\n }\n }\n }\n }\n\n __syncthreads();\n\n // XXX Could change to poll by ballot style compress\n if(gid < unfine_set_size)\n {\n // Store the results.\n#pragma unroll\n for(int i = lane_id; i < HASH_SIZE; i += GROUP_SIZE)\n {\n int key = key_group_s[i];\n if(key != SLOT_VACANT)\n {\n int ind = atomicAdd(&count_s[group_id], 1);\n C_hat[C_hat_start[a_row_id] + ind] = key;\n }\n }\n }\n\n __syncthreads();\n\n if(gid < unfine_set_size)\n {\n // Store the results.\n if(lane_id == 0)\n {\n C_hat_end[a_row_id] = C_hat_start[a_row_id] + count_s[group_id];\n }\n }\n}\n} // namespace selector\n\ntemplate \nvoid __global__ createCfMapGlobal(const IndexType *cf_map, int64_t *cf_map_global, const int64_t my_part_offset, const IndexType num_owned_fine_pts)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_owned_fine_pts ; tidx += blockDim.x * gridDim.x)\n {\n // Only renumber the interior points\n if (cf_map[tidx] >= 0 && tidx < num_owned_fine_pts)\n {\n // Simply shift\n cf_map_global[tidx] = (int64_t) cf_map[tidx] /* my_local_id */ + my_part_offset;\n }\n else\n {\n cf_map_global[tidx] = -1;\n }\n }\n}\n\n#include \n\nenum { WARP_SIZE = 32, GRID_SIZE = 1024, SMEM_SIZE = 128 };\n\ntemplate< typename Value_type, int CTA_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nfillS2ColIndices( const int A_num_rows,\n const int *__restrict cf_map,\n const int *__restrict C_hat,\n const int *__restrict C_hat_start,\n const int *__restrict C_hat_end,\n const int *__restrict S2_rows,\n int *S2_cols,\n int64_t *S2_global_cols,\n int64_t *cf_map_global)\n{\n const int NUM_WARPS = CTA_SIZE / 32;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows; a_row_id += NUM_WARPS_PER_GRID )\n {\n int coarse_fine_id = cf_map[a_row_id];\n\n // Skip fine rows.\n if ( coarse_fine_id < 0 )\n {\n continue;\n }\n\n // Rebuild C_hat.\n int c_hat_it = C_hat_start[a_row_id];\n int c_hat_end = C_hat_end [a_row_id];\n int s2_col_it = S2_rows[coarse_fine_id];\n\n // Iterate over the columns of A.\n for ( c_hat_it += lane_id ; utils::any( c_hat_it < c_hat_end ) ; c_hat_it += WARP_SIZE, s2_col_it += WARP_SIZE)\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int s2_col_id(-1);\n\n if ( c_hat_it < c_hat_end )\n {\n s2_col_id = C_hat[c_hat_it];\n }\n\n if (s2_col_id != -1)\n {\n S2_cols[s2_col_it + lane_id] = cf_map[s2_col_id];\n\n if (S2_global_cols != NULL)\n {\n S2_global_cols[s2_col_it + lane_id] = cf_map_global[s2_col_id];\n }\n }\n }\n }\n}\n\n\n\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid Selector >::renumberAndCountCoarsePoints( IVector &cf_map,\n int &num_coarse_points,\n int num_rows)\n{\n IVector mark(cf_map.size(), 0);\n int blockSize = 128;\n int num_blocks = std::min( 4096, (int) (cf_map.size() + blockSize - 1) / blockSize );\n\n if (cf_map.size() > 0)\n {\n coarseMarkKernel <<< num_blocks, blockSize>>>(cf_map.raw(), mark.raw(), cf_map.size());\n cudaCheckError();\n }\n\n // get the sequence of values\n thrust_wrapper::inclusive_scan(mark.begin(), mark.end(), mark.begin());\n cudaCheckError();\n\n // assign to cf_map\n if (cf_map.size() > 0)\n {\n modifyCoarseMapKernel <<< num_blocks, blockSize>>>(cf_map.raw(), mark.raw(), cf_map.size());\n cudaCheckError();\n }\n\n\n num_coarse_points = (int) thrust_wrapper::count_if(cf_map.begin(), cf_map.begin() + num_rows, is_non_neg());\n cudaCheckError();\n}\n\n\ntemplate< AMGX_VecPrecision V, AMGX_MatPrecision M, AMGX_IndPrecision I >\nvoid Selector >::renumberAndCountCoarsePoints( IVector &cf_map,\n int &num_coarse_points,\n int num_rows)\n{\n num_coarse_points = 0;\n\n for (int i = 0; i < num_rows; i++)\n {\n // if coarse\n if (cf_map[i] == COARSE)\n {\n cf_map[i] = num_coarse_points++;\n }\n }\n}\n\n__global__\nvoid correctCfMapKernel(int *cf_map, const int *cf_map_scanned, const int *cf_map_s2, int size)\n{\n for (int tIdx = threadIdx.x + blockIdx.x * blockDim.x; tIdx < size; tIdx += gridDim.x * blockDim.x)\n {\n if (cf_map[tIdx] == COARSE) // Coarse point\n {\n int coarse_id = cf_map_scanned[tIdx];\n int coarse_id_s2 = cf_map_s2[coarse_id];\n // if cf_map_s2 is STRONG_FINE, mark as COARSE\n cf_map[tIdx] = (coarse_id_s2 == STRONG_FINE) ? COARSE : coarse_id_s2;\n }\n }\n}\n\ntemplate \nvoid Selector >\n::correctCfMap(IVector &cf_map, IVector &cf_map_scanned, IVector &cf_map_S2)\n{\n FatalError(\"CorrectCFMap not implemented on host\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\n\ntemplate \nvoid Selector >::correctCfMap(IVector &cf_map, IVector &cf_map_scanned, IVector &cf_map_S2)\n{\n // One thread per row of A\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // get raw pointers to data\n const int cta_size = 128;\n const int grid_size = std::min( 4096, ((int) cf_map.size() + cta_size - 1) / cta_size);\n correctCfMapKernel <<< grid_size, cta_size>>>(cf_map.raw(), cf_map_scanned.raw(), cf_map_S2.raw(), cf_map.size());\n cudaCheckError();\n}\n\ntemplate \nvoid Selector >\n::createS2(Matrix_h &A,\n Matrix_h &S2,\n const BVector &s_con,\n IVector &cf_map)\n{\n FatalError(\"Aggressive PMIS selector not implemented on host\", AMGX_ERR_NOT_IMPLEMENTED);\n}\n\ntemplate \ntemplate \nvoid Selector>::compute_c_hat_opt_dispatch(\n const Matrix_d &A,\n const bool *s_con,\n const int *C_hat_start,\n int *C_hat,\n int *C_hat_end,\n int *cf_map,\n IntVector &unfine_set)\n{\n typedef typename IndPrecisionMap::Type Index_type;\n\n constexpr int cta_size = 128;\n constexpr int ngroups = cta_size / group_size;\n\n const int grid_size = unfine_set.size() / ngroups + 1;\n\n hipDeviceProp_t deviceProps = getDeviceProperties();\n size_t max_shmem_size = deviceProps.sharedMemPerMultiprocessor;\n\n constexpr int shmem_size =\n (sizeof(int)+sizeof(Index_type))*ngroups*hash_size + group_size; \n\n if(shmem_size > max_shmem_size) \n { \n FatalError(\"In compute_values_opt the requested hash size is larger than max.\\n\", \n AMGX_ERR_NOT_IMPLEMENTED); \n } \n\n hipFuncSetAttribute(reinterpret_cast(amgx::classical::selector::compute_c_hat_kernel_opt \n ), \n hipFuncAttributeMaxDynamicSharedMemorySize, shmem_size); \n\n amgx::classical::selector::compute_c_hat_kernel_opt\n \n <<>>(\n unfine_set.size(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n s_con,\n C_hat_start,\n C_hat,\n C_hat_end,\n cf_map,\n unfine_set.raw());\n cudaCheckError();\n}\n\n// Determines which nodes are not fine, hence excluded\n__global__ void set_unfine(int n, int* cf_map, int* unfine, int* unfine_offs)\n{\n int i = blockIdx.x*blockDim.x + threadIdx.x;\n if(i >= n) return;\n\n unfine[i] = (cf_map[i] >= 0) ? 1 : 0;\n\n if(i == 0) \n {\n unfine_offs[0] = 0;\n }\n}\n\n// Compresses indirection\n__global__ void unfine_set_fill(int n, int* unfine, int* unfine_offs, int* unfine_set)\n{\n int i = blockIdx.x*blockDim.x + threadIdx.x;\n if(i >= n) return;\n\n if(unfine[i]) \n { \n unfine_set[unfine_offs[i]] = i; \n } \n}\n\n\ntemplate \nvoid Selector >::createS2(Matrix_d &A,\n Matrix_d &S2,\n const BVector &s_con,\n IVector &cf_map)\n{\n const int blockSize = 256;\n // One thread per row of A\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // get raw pointers to data\n const IndexType *Aoffsets = A.row_offsets.raw();\n const IndexType *Acolumn_indices = A.col_indices.raw();\n const ValueType *Avalues = A.values.raw();\n const IndexType Anum_rows = (int) A.get_num_rows();\n int *cf_map_ptr = cf_map.raw();\n Hash_Workspace exp_wk(true, GRID_SIZE);\n IntVector C_hat_start( A.get_num_rows() + 1, 0 ), C_hat_end( A.get_num_rows() + 1, 0 );\n\n int max_c_hat_size = 0;\n int nunfine = 0;\n IntVector unfine_set;\n\n if(this->m_use_opt_kernels)\n {\n IntVector unfine(A.get_num_rows());\n IntVector unfine_offs(A.get_num_rows()+1);\n constexpr int nthreads = 128;\n const int nblocks = A.get_num_rows() / nthreads + 1;\n set_unfine<<>>(A.get_num_rows(), cf_map.raw(), unfine.raw(), unfine_offs.raw());\n cudaCheckError();\n thrust_wrapper::inclusive_scan(unfine.begin(), unfine.end(), unfine_offs.begin()+1);\n\n nunfine = unfine_offs[unfine_offs.size()-1];\n if(nunfine > 0)\n {\n unfine_set.resize(nunfine);\n unfine_set_fill<<>>(A.get_num_rows(), unfine.raw(), unfine_offs.raw(), unfine_set.raw());\n cudaCheckError();\n }\n\n const int CTA_SIZE = 128;\n int grid_size = nunfine/CTA_SIZE + 1;\n\n amgx::classical::selector::estimate_c_hat_size_kernel_opt< CTA_SIZE > <<< grid_size, CTA_SIZE>>>(\n nunfine,\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw(),\n unfine_set.raw());\n cudaCheckError();\n\n max_c_hat_size = thrust::reduce(C_hat_start.begin(), C_hat_start.end(), -1, thrust::maximum());\n }\n else\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int avg_nz_per_row = A.get_num_nz() / A.get_num_rows();\n int grid_size = A.get_num_rows()/NUM_WARPS + 1;\n\n if ( avg_nz_per_row < 16 )\n {\n amgx::classical::selector::estimate_c_hat_size_kernel< 8, CTA_SIZE, WARP_SIZE> <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw());\n cudaCheckError();\n }\n else\n {\n amgx::classical::selector::estimate_c_hat_size_kernel <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw());\n }\n\n cudaCheckError();\n }\n\n // Compute row offsets.\n thrust_wrapper::exclusive_scan(C_hat_start.begin(), C_hat_start.end(), C_hat_start.begin());\n cudaCheckError();\n int nVals = C_hat_start[C_hat_start.size() - 1];\n IntVector C_hat( nVals );\n\n if(this->m_use_opt_kernels)\n {\n int hash_size = pow(2, ceil(log2(max_c_hat_size)));\n\n if(nunfine > 0)\n {\n switch(hash_size)\n {\n case 1:\n case 2:\n case 4:\n case 8:\n case 16:\n case 32:\n compute_c_hat_opt_dispatch<32, 8>(A, s_con.raw(), C_hat_start.raw(), C_hat.raw(), C_hat_end.raw(), cf_map.raw(), unfine_set);\n break;\n case 64:\n compute_c_hat_opt_dispatch<64, 8>(A, s_con.raw(), C_hat_start.raw(), C_hat.raw(), C_hat_end.raw(), cf_map.raw(), unfine_set);\n break;\n case 128:\n compute_c_hat_opt_dispatch<128, 8>(A, s_con.raw(), C_hat_start.raw(), C_hat.raw(), C_hat_end.raw(), cf_map.raw(), unfine_set);\n break;\n case 256:\n compute_c_hat_opt_dispatch<256, 16>(A, s_con.raw(), C_hat_start.raw(), C_hat.raw(), C_hat_end.raw(), cf_map.raw(), unfine_set);\n break;\n case 512:\n compute_c_hat_opt_dispatch<512, 32>(A, s_con.raw(), C_hat_start.raw(), C_hat.raw(), C_hat_end.raw(), cf_map.raw(), unfine_set);\n break;\n default:\n FatalError(\"Size of C_hat in a single row too large, solve with original solvers.\", AMGX_ERR_INTERNAL);\n }\n }\n }\n else\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int avg_nz_per_row = A.get_num_nz() / A.get_num_rows();\n int attempt = 0;\n\n for ( bool done = false ; !done && attempt < 10 ; ++attempt )\n {\n // Double the amount of GMEM (if needed).\n if ( attempt > 0 )\n {\n exp_wk.expand();\n }\n\n // Reset the status. \n // TODO: Optimize with async copies.\n int status = 0;\n hipMemcpy( exp_wk.get_status(), &status, sizeof(int), hipMemcpyHostToDevice );\n cudaCheckError();\n // Compute the set C_hat.\n int work_offset = GRID_SIZE * NUM_WARPS;\n hipMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), hipMemcpyHostToDevice );\n cudaCheckError();\n\n // Run the computation.\n if ( avg_nz_per_row < 16 )\n {\n amgx::classical::selector::compute_c_hat_kernel< 8, CTA_SIZE, SMEM_SIZE, WARP_SIZE> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n C_hat.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_work_queue(),\n exp_wk.get_status() );\n cudaCheckError();\n }\n else\n {\n amgx::classical::selector::compute_c_hat_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n C_hat.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_work_queue(),\n exp_wk.get_status() );\n }\n\n cudaCheckError();\n // Read the result from count_non_zeroes.\n hipMemcpy( &status, exp_wk.get_status(), sizeof(int), hipMemcpyDeviceToHost );\n cudaCheckError();\n done = status == 0;\n }\n }\n // get pointers to data\n int *C_hat_ptr = C_hat.raw();\n int *C_hat_start_ptr = C_hat_start.raw();\n int *C_hat_end_ptr = C_hat_end.raw();\n int S2_num_rows;;\n DistributedArranger *prep = NULL;\n\n if (A.is_matrix_distributed())\n {\n prep = new DistributedArranger;\n int num_owned_fine_pts = A.get_num_rows();\n int num_owned_coarse_pts = thrust_wrapper::count_if(cf_map.begin(), cf_map.begin() + num_owned_fine_pts, is_non_neg());\n int num_halo_coarse_pts = thrust_wrapper::count_if(cf_map.begin() + num_owned_fine_pts, cf_map.end(), is_non_neg());\n cudaCheckError();\n S2_num_rows = num_owned_coarse_pts;\n prep->initialize_manager(A, S2, num_owned_coarse_pts);\n }\n else\n {\n S2_num_rows = (int) thrust_wrapper::count_if(cf_map.begin(), cf_map.end(), is_non_neg());\n cudaCheckError();\n }\n\n // count the number of non-zeros in the interpolation matrix\n int numBlocks = (int) (A.get_num_rows() + blockSize - 1) / blockSize;\n IntVector nonZeroOffsets(S2_num_rows + 1);\n IntVector nonZerosPerRow(S2_num_rows);\n // Updating the number of nonZeros for your own coarse rows\n int *nonZerosPerRow_ptr = nonZerosPerRow.raw();\n nonZerosPerCoarseRowKernel <<< numBlocks, blockSize>>>(Anum_rows, cf_map.raw(), C_hat_start_ptr, C_hat_end_ptr,\n nonZerosPerRow_ptr);\n cudaCheckError();\n // get total with a reduction\n int nonZeros = thrust_wrapper::reduce(nonZerosPerRow.begin(), nonZerosPerRow.end());\n cudaCheckError();\n // get the offsets with an exclusive scan\n thrust_wrapper::exclusive_scan(nonZerosPerRow.begin(), nonZerosPerRow.end(), nonZeroOffsets.begin());\n cudaCheckError();\n nonZeroOffsets[S2_num_rows] = nonZeros;\n // resize S2\n S2.resize(0, 0, 0, 1);\n S2.addProps(CSR);\n S2.set_num_rows(S2_num_rows);\n S2.set_num_cols(S2_num_rows);\n S2.set_num_nz(nonZeros);\n S2.row_offsets.resize(S2_num_rows + 1);\n S2.col_indices.resize(nonZeros);\n S2.diag.resize(S2_num_rows);\n S2.set_block_dimx(A.get_block_dimx());\n S2.set_block_dimy(A.get_block_dimy());\n cudaCheckError();\n // copy nonzero offsets to the P matrix\n amgx::thrust::copy(nonZeroOffsets.begin(), nonZeroOffsets.end(), S2.row_offsets.begin());\n cudaCheckError();\n {\n const int CTA_SIZE = 256;\n // Run the computation.\n typedef typename MatPrecisionMap::Type Value_type;\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n if (!A.is_matrix_distributed())\n {\n int grid_size = (A.get_num_rows() / NUM_WARPS) + 1;\n\n fillS2ColIndices <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n cf_map.raw(),\n C_hat.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n S2.row_offsets.raw(),\n S2.col_indices.raw(),\n NULL,\n NULL);\n cudaCheckError();\n }\n else\n {\n I64Vector_d cf_map_global(cf_map.size());\n int num_owned_fine_pts = A.get_num_rows();\n int my_rank = A.manager->global_id();\n const int cta_size = 128;\n int grid_size = (num_owned_fine_pts + cta_size - 1) / cta_size;\n createCfMapGlobal <<< grid_size, cta_size>>>(cf_map.raw(), cf_map_global.raw(), S2.manager->part_offsets_h[my_rank], num_owned_fine_pts);\n cudaCheckError();\n\n // Exchange the cf_map_global so that we know the coarse global id of halo nodes\n cf_map_global.dirtybit = 1;\n A.manager->exchange_halo_2ring(cf_map_global, cf_map_global.tag);\n I64Vector_d S2_col_indices_global(S2.col_indices.size());\n\n grid_size = (A.get_num_rows() / NUM_WARPS) + 1;\n fillS2ColIndices <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n cf_map.raw(),\n C_hat.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n S2.row_offsets.raw(),\n S2.col_indices.raw(),\n S2_col_indices_global.raw(),\n cf_map_global.raw());\n cudaCheckError();\n prep->initialize_manager_from_global_col_indices(S2, S2_col_indices_global);\n }\n\n if (prep != NULL) { delete prep; }\n\n cudaCheckError();\n }\n}\n\n\n\ntemplate \nvoid Selector >::demoteStrongEdges(const Matrix_h &A,\n const FVector &weights,\n BVector &s_con,\n const IVector &cf_map,\n const IndexType base_offset)\n{\n for ( int tIdx = 0; tIdx < A.get_num_rows() - base_offset; tIdx++)\n {\n const int offset = A.row_offsets[tIdx];\n const int numj = A.row_offsets[tIdx + 1] - offset;\n\n if (cf_map[base_offset + tIdx] == COARSE)\n {\n // for each column\n for (int j = offset; j < offset + numj; j++)\n {\n if (!s_con[j]) { continue; } // skip weakly connected points\n\n int jcol = A.col_indices[j];\n\n if (cf_map[jcol] == COARSE) // coarse neighbour\n {\n s_con[j] = false; //make it weak\n }\n }\n }\n }\n}\n\ntemplate \nvoid Selector >::demoteStrongEdges(const Matrix_d &A,\n const FVector &weights,\n BVector &s_con,\n const IVector &cf_map,\n const IndexType offset)\n{\n const int blockSize = 256;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (((A.get_num_rows() - offset) + blockSize - 1) / blockSize));\n const int numRows = (int)(A.get_num_rows() - offset);\n resolve_boundary <<< blockSize, numBlocks>>>(A.row_offsets.raw() + offset, A.col_indices.raw(), s_con.raw(), cf_map.raw(), offset, A.get_num_rows() - offset);\n cudaCheckError();\n}\n\n\n\ntemplate\nstd::map*> &\nSelectorFactory::getFactories( )\n{\n static std::map*> s_factories;\n return s_factories;\n}\n\ntemplate\nvoid SelectorFactory::registerFactory(std::string name, SelectorFactory *f)\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::const_iterator it = factories.find(name);\n\n if (it != factories.end())\n {\n std::string error = \"SelectorFactory '\" + name + \"' has already been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n factories[name] = f;\n}\n\ntemplate\nvoid SelectorFactory::unregisterFactory(std::string name)\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.find(name);\n\n if (it == factories.end())\n {\n std::string error = \"SelectorFactory '\" + name + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n SelectorFactory *factory = it->second;\n assert( factory != NULL );\n delete factory;\n factories.erase(it);\n}\n\ntemplate\nvoid SelectorFactory::unregisterFactories( )\n{\n std::map*> &factories = getFactories( );\n typename std::map *>::iterator it = factories.begin( );\n\n for ( ; it != factories.end( ) ; )\n {\n SelectorFactory *factory = it->second;\n assert( factory != NULL );\n it++;\n delete factory;\n }\n\n factories.clear( );\n}\n\ntemplate\nSelector *SelectorFactory::allocate(AMG_Config &cfg, const std::string &cfg_scope)\n{\n std::map*> &factories = getFactories( );\n std::string selector = cfg.getParameter(\"selector\", cfg_scope);\n typename std::map *>::const_iterator it = factories.find(selector);\n\n if (it == factories.end())\n {\n std::string error = \"SelectorFactory '\" + selector + \"' has not been registered\\n\";\n FatalError(error.c_str(), AMGX_ERR_CORE);\n }\n\n return it->second->create(cfg, cfg_scope);\n};\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class SelectorFactory::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\nAMGX_FORCOMPLEX_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n} // namespace classical\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/classical/selectors/selector.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\nnamespace amgx\n{\n\nnamespace classical\n{\n/*************************************************************************\n * marks the strongest connected (and indepent) points as coarse\n ************************************************************************/\n\ntemplate< AMGX_VecPrecision t_vecPrec, AMGX_MatPrecision t_matPrec, AMGX_IndPrecision t_indPrec >\nstruct MarkAdditionalPoints\n{\n static const AMGX_VecPrecision vecPrec = (AMGX_VecPrecision)t_indPrec;\n typedef Vector > IndVector;\n static int markAdditionalCoarsePoints(const Matrix > &A,\n const Vector > &weights,\n const Vector > &s_con,\n IndVector &cf_map_in,\n IndVector &cf_map_out,\n IndVector &mark)\n {\n int num_coarse = 0;\n\n //deep copy cf_map_in into cf_map_out\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n cf_map_out[i] = cf_map_in[i];\n }\n\n // set all unassigned points to be coarse\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n mark[i] = 0;\n\n if (cf_map_in[i] == UNASSIGNED)\n {\n cf_map_out[i] = COARSE;\n mark[i] = 1;\n }\n }\n\n // now run through newly marked coarse points and make them unassigned again if necessary\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n if (mark[i] > 0) // newly marked C point\n {\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n if (!s_con[j]) { continue; } // skip weak connections\n\n int jcol = A.col_indices[j];\n\n if (mark[jcol] && weights[jcol] > 1) // other point is also newly coarse\n {\n // if my weight is larger, they can't be a coarse point\n if (weights[i] > weights[jcol])\n {\n cf_map_out[jcol] = UNASSIGNED;\n }\n // if their weight is large, I can't be a coarse point\n else if (weights[jcol] > weights[i])\n {\n cf_map_out[i] = UNASSIGNED;\n }\n }\n }\n }\n }\n\n // finally count the number of coarse points\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n num_coarse += static_cast(cf_map_out[i] == COARSE);\n }\n\n return num_coarse;\n }\n\n /*************************************************************************\n * marks all neighbors of coarse points as fine\n ************************************************************************/\n static int markAdditionalFinePoints(const Matrix > &A,\n const Vector > &s_con,\n IndVector &cf_map_in,\n IndVector &cf_map_out)\n {\n int num_fine = 0;\n\n //deep copy cf_map_in into cf_map_out\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n cf_map_out[i] = cf_map_in[i];\n }\n\n //for each row of the matrix\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n //if already fine add it to the count\n num_fine += static_cast(cf_map_in[i] == FINE);\n\n if (cf_map_in[i] == UNASSIGNED)\n {\n // for each off diagonal, if I have a new strongly connected coarse neighbour\n // set to fine\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n if (!s_con[j]) { continue; } //skip weakly connected points\n\n int jcol = A.col_indices[j];\n\n if (cf_map_in[jcol] == COARSE)\n {\n cf_map_out[i] = FINE;\n num_fine++;\n break;\n }\n }\n }\n }\n\n return num_fine;\n }\n};\n\n/*************************************************************************\n * mark coarse and fine points\n ************************************************************************/\ntemplate \nvoid PMIS_Selector >\n::markCoarseFinePoints_1x1(Matrix_h &A,\n FVector &weights,\n const BVector &s_con,\n IVector &cf_map,\n IVector &scratch,\n int cf_map_init)\n{\n int num_coarse = 0, num_fine = 0;\n int num_strong_fine = 0;\n\n //mark all unconnected points as fine\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n int numj = A.row_offsets[i + 1] - A.row_offsets[i];\n\n if (numj == 0) //if there are no columns\n {\n cf_map[i] = FINE;\n }\n else if (numj == 1 && A.col_indices[A.row_offsets[i]] == i) //if there is only 1 column and it is the diagonal\n {\n cf_map[i] = FINE;\n }\n else if (weights[i] < 1) // no point strongly depends on point\n {\n cf_map[i] = FINE;\n }\n else\n {\n cf_map[i] = UNASSIGNED;\n }\n\n // check to see if a point has no strong connections to anyone\n // needs more justification / knowledge of everywhere this condition is applied\n bool isolated = true;\n\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n if (s_con[j])\n {\n isolated = false;\n break;\n }\n }\n\n if (isolated)\n {\n cf_map[i] = STRONG_FINE;\n num_strong_fine++;\n weights[i] = 0.F;\n }\n }\n\n // temp storage for marking coarse points\n IntVector mark(A.get_num_rows(), 0);\n int iterations = 0;\n\n do\n {\n ++iterations;\n num_coarse = MarkAdditionalPoints::markAdditionalCoarsePoints(A, weights, s_con, cf_map, scratch, mark);\n num_fine = MarkAdditionalPoints::markAdditionalFinePoints(A, s_con, scratch, cf_map);\n }\n while (num_coarse + num_fine + num_strong_fine < A.get_num_rows());\n\n // now we have \\hat{C} and the diagonals, we can start calculating the weights\n // remap cf_map (to be like in distance1)\n}\n\n/*************************************************************************\n * device kernels\n ************************************************************************/\n\n/*************************************************************************\n * Perform the initial marking for each row\n ************************************************************************/\ntemplate \n__global__\nvoid initialMarkingKernel(const IndexType *offsets, const IndexType *column_indices,\n const int numRows, int *cf_map, float *weights, const bool *s_con)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n // set initial C/F/U\n if (numj == 0) // no points\n {\n cf_map[tIdx] = FINE;\n }\n else if (numj == 1 && column_indices[offset] == tIdx) // only the diagonal\n {\n cf_map[tIdx] = FINE;\n }\n else if (weights[tIdx] < 1) // no strong dependancies\n {\n cf_map[tIdx] = FINE;\n }\n else\n {\n cf_map[tIdx] = UNASSIGNED;\n }\n\n bool isolated = true;\n\n for (int j = offsets[tIdx]; j < offsets[tIdx + 1]; j++)\n {\n if (s_con[j])\n {\n isolated = false;\n break;\n }\n }\n\n if (isolated)\n {\n cf_map[tIdx] = STRONG_FINE;\n weights[tIdx] = 0.F;\n }\n }\n}\n\ntemplate \n__global__\nvoid initialMarkingCFInit3Kernel(const IndexType *offsets, const IndexType *column_indices,\n const int numRows, int *cf_map, float *weights, const bool *s_con)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n // set initial C/F/U\n if (numj == 0) // no points\n {\n cf_map[tIdx] = FINE;\n }\n else if (numj == 1 && column_indices[offset] == tIdx) // only the diagonal\n {\n cf_map[tIdx] = FINE;\n }\n else if (weights[tIdx] < 1) // no strong dependancies\n {\n cf_map[tIdx] = FINE;\n }\n else\n {\n cf_map[tIdx] = UNASSIGNED;\n }\n\n bool isolated = true;\n\n for (int j = offsets[tIdx]; j < offsets[tIdx + 1]; j++)\n {\n if (s_con[j])\n {\n isolated = false;\n break;\n }\n }\n\n if (isolated)\n {\n cf_map[tIdx] = COARSE;\n weights[tIdx] = 0.F;\n }\n }\n}\n\n/*************************************************************************\n * Perform the initial marking for each row\n ************************************************************************/\ntemplate \n__global__\nvoid initialMarkingCfInitKernel(const IndexType *offsets, const IndexType *column_indices,\n const int numRows, int *cf_map, float *weights, const bool *s_con, IndexType *mark)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n // Check if on boundary, this is unmarking even coarse points\n for (int j = offsets[tIdx]; j < offsets[tIdx + 1]; j++)\n {\n if (column_indices[j] > numRows)\n {\n cf_map[tIdx] = UNASSIGNED;\n mark[tIdx] = 1;\n break;\n }\n }\n\n // set initial C/F/U\n if (numj == 0) // no points\n {\n cf_map[tIdx] = FINE;\n }\n else if (numj == 1 && column_indices[offset] == tIdx) // only the diagonal\n {\n cf_map[tIdx] = FINE;\n }\n else if (weights[tIdx] < 1) // no strong dependancies\n {\n cf_map[tIdx] = FINE;\n }\n else if (cf_map[tIdx] == STRONG_FINE)\n {\n // Do nothing, leave as strong fine\n weights[tIdx] = 0.F;\n }\n else if (cf_map[tIdx] == FINE)\n {\n // Fine point, mark as unassigned\n cf_map[tIdx] = UNASSIGNED;\n mark[tIdx] = 1;\n }\n }\n}\n\n\n/*************************************************************************\n * Mark every currently unassigned point as coarse\n ************************************************************************/\ntemplate \n__global__\nvoid markUnassignedAsCoarseKernel(const int numRows,\n const float *weights,\n IndexType *cf_map_in,\n IndexType *cf_map_out, IndexType *mark)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n const IndexType map_in = cf_map_in[tIdx];\n const int unassigned = map_in == UNASSIGNED ? 1 : 0;\n mark[tIdx] = unassigned;\n\n // Mark unassigned piont as coarse\n if (weights[tIdx] > 1.)\n {\n cf_map_out[tIdx] = unassigned ? COARSE : map_in;\n }\n else\n {\n cf_map_out[tIdx] = map_in;\n }\n }\n}\n\n/*************************************************************************\n * Unset all newly set coarse points who shouldn't be\n ************************************************************************/\ntemplate \n__global__\nvoid markAdditionalCoarsePointsKernel(const IndexType *offsets,\n const IndexType *column_indices,\n const int numRows, const float *weights,\n const bool *s_con,\n IndexType *cf_map_out, IndexType *mark)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n // run through newly marked coarse points & set to unassigned again if necessary\n if (mark[tIdx] > 0)\n {\n for ( int j = offsets[tIdx], j_end = offsets[tIdx + 1] ; j < j_end; j++ )\n {\n if (!s_con[j]) { continue; } //these are always false for edges pointing outside of the 1-ring halo\n\n int jcol = column_indices[j];\n float w_col = weights[jcol];\n float w_row = weights[tIdx];\n\n if ( mark[jcol] && w_col > 1.0f )\n {\n // If my weight is bigger, they can't be a coarse point\n if ( w_row > w_col )\n {\n cf_map_out[jcol] = UNASSIGNED;\n }\n // If their weight is bigger, I can't be a coarse point\n else if ( w_col > w_row )\n {\n cf_map_out[tIdx] = UNASSIGNED;\n }\n }\n }\n }\n }\n}\n\n/*************************************************************************\n * mark every neighbour of a newly assigned coarse point to be fine\n ************************************************************************/\ntemplate \n__global__\nvoid markAdditionalFinePointsKernel(const IndexType *offsets, const IndexType *column_indices,\n const int numRows, const bool *s_con,\n int *cf_map_in, int *cf_map_out)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x ; tIdx < numRows ; tIdx += blockDim.x * gridDim.x )\n {\n const int map_in = cf_map_in[tIdx];\n // go over each row\n bool fine = false;\n\n if ( map_in == UNASSIGNED )\n {\n // for each column\n for ( int j = offsets[tIdx], j_end = offsets[tIdx + 1] ; !fine && j < j_end ; ++j )\n {\n if ( !s_con[j] )\n {\n continue; // skip weakly connected points\n }\n\n int jcol = column_indices[j];\n fine = cf_map_in[jcol] == COARSE;\n }\n }\n\n cf_map_out[tIdx] = fine ? FINE : map_in;\n }\n}\n\ntemplate \nvoid PMIS_Selector >::markCoarseFinePoints_1x1(Matrix_d &A,\n FVector &weights,\n const BVector &s_con,\n IVector &cf_map,\n IVector &scratch,\n int cf_map_init)\n{\n if (A.hasProps(DIAG))\n {\n FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // Choose blocksize. Using 1 thread / row for now\n const int blockSize = 256;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) ((A.get_num_rows() + blockSize - 1) / blockSize));\n const int numRows = (int) A.get_num_rows();\n // raw pointers from the cusp arrays\n const IndexType *offsets_ptr = A.row_offsets.raw();\n const IndexType *column_indices_ptr = A.col_indices.raw();\n float *weights_ptr = weights.raw();\n const bool *s_con_ptr = s_con.raw();\n int *cf_map_ptr = cf_map.raw();\n int *scratch_ptr = scratch.raw();\n IVector mark(cf_map.size());\n // do the initial marking of points\n thrust_wrapper::fill(mark.begin(), mark.end(), 0);\n cudaCheckError();\n\n if (numRows > 0)\n {\n if (cf_map_init == 0)\n {\n initialMarkingKernel <<< numBlocks, blockSize>>>(offsets_ptr, column_indices_ptr, numRows, cf_map_ptr, weights_ptr, s_con_ptr);\n cudaCheckError();\n }\n else if (cf_map_init == 1)\n {\n initialMarkingCfInitKernel <<< numBlocks, blockSize>>>(offsets_ptr, column_indices_ptr, numRows, cf_map_ptr, weights_ptr, s_con_ptr, mark.raw());\n cudaCheckError();\n }\n else if (cf_map_init == 3)\n {\n initialMarkingCFInit3Kernel <<< numBlocks, blockSize>>>(offsets_ptr, column_indices_ptr, numRows, cf_map_ptr, weights_ptr, s_con_ptr);\n cudaCheckError();\n }\n else\n {\n FatalError(\"invalid value of cf_map_init in pmis selector, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n\n cudaCheckError();\n // exchange initial cf_map for the 1st ring halos\n cf_map.dirtybit = 1;\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo(cf_map, cf_map.tag);\n }\n\n // count the # of coarse and fine points we have now\n int numUnassigned = 0;\n // array needed for marking of additional coarse points\n int one_ring_size, offset;\n\n if (!A.is_matrix_singleGPU())\n {\n A.getOffsetAndSizeForView(FULL, &offset, &one_ring_size);\n }\n else\n {\n one_ring_size = numRows;\n }\n\n const int numBlocks1Ring = std::min (AMGX_GRID_MAX_SIZE, (int) ((one_ring_size + blockSize - 1) / blockSize));\n // iterate until all points have been classified\n int numUnassignedMax = numRows;\n //int numUnassignedMaxPrevious;\n int iter = 0;\n\n do\n {\n //numUnassignedMaxPrevious = numUnassignedMax;\n if (numRows > 0)\n {\n if (iter || !cf_map_init)\n {\n // Laucnhed for 1-ring\n markUnassignedAsCoarseKernel <<< numBlocks1Ring, blockSize>>>(one_ring_size, weights_ptr, cf_map_ptr,\n scratch_ptr, mark.raw());\n cudaCheckError();\n markAdditionalCoarsePointsKernel <<< numBlocks, blockSize>>>(offsets_ptr, column_indices_ptr,\n numRows, weights_ptr, s_con_ptr,\n scratch_ptr, mark.raw());\n cudaCheckError();\n }\n else\n {\n amgx::thrust::copy(cf_map.begin(), cf_map.begin() + one_ring_size, scratch.begin());\n cudaCheckError();\n }\n\n // If a neighbor marked one of my node as unassigned, need to mark it\n // Taking the min (unassigned == -4, coarse = -1)\n scratch.dirtybit = 1;\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->min_from_halo(scratch, scratch.tag);\n }\n\n cudaCheckError();\n // Now exchange with neighbors scratch for my own nodes\n scratch.dirtybit = 1;\n\n if (!A.is_matrix_singleGPU())\n {\n // exchange updated CF map for halos (stored in scratch)\n A.manager->exchange_halo(scratch, scratch.tag);\n }\n\n cudaCheckError();\n markAdditionalFinePointsKernel <<< numBlocks, blockSize>>>(offsets_ptr, column_indices_ptr,\n numRows, s_con_ptr, scratch_ptr,\n cf_map_ptr);\n cudaCheckError();\n } // num rows > 0\n\n // count # of points still unassigned\n numUnassigned = (int) thrust_wrapper::count(cf_map.begin(), cf_map.begin() + numRows, (int)UNASSIGNED);\n cudaCheckError();\n numUnassignedMax = numUnassigned;\n cudaCheckError();\n cf_map.dirtybit = 1;\n\n if (A.is_matrix_distributed())\n {\n // exchange CF map for halos\n A.manager->exchange_halo(cf_map, cf_map.tag);\n // get max across all processes\n A.manager->getComms()->all_reduce_max(numUnassigned, numUnassignedMax);\n }\n\n iter++;\n }\n while ( numUnassignedMax != 0 );\n}\n\ntemplate \nvoid PMIS_SelectorBase< T_Config>::markCoarseFinePoints(Matrix< T_Config> &A,\n FVector &weights,\n const BVector &s_con,\n IVector &cf_map,\n IVector &scratch,\n int cf_map_init)\n{\n ViewType oldView = A.currentView();\n A.setView(OWNED);\n\n if (A.get_block_size() == 1)\n {\n markCoarseFinePoints_1x1(A, weights, s_con, cf_map, scratch, cf_map_init);\n }\n else\n {\n FatalError(\"Unsupported block size PMIS selector\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class PMIS_SelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class PMIS_Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace classical\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\nnamespace amgx\n{\n\nnamespace classical\n{\n/*************************************************************************\n * marks the strongest connected (and indepent) points as coarse\n ************************************************************************/\n\ntemplate< AMGX_VecPrecision t_vecPrec, AMGX_MatPrecision t_matPrec, AMGX_IndPrecision t_indPrec >\nstruct MarkAdditionalPoints\n{\n static const AMGX_VecPrecision vecPrec = (AMGX_VecPrecision)t_indPrec;\n typedef Vector > IndVector;\n static int markAdditionalCoarsePoints(const Matrix > &A,\n const Vector > &weights,\n const Vector > &s_con,\n IndVector &cf_map_in,\n IndVector &cf_map_out,\n IndVector &mark)\n {\n int num_coarse = 0;\n\n //deep copy cf_map_in into cf_map_out\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n cf_map_out[i] = cf_map_in[i];\n }\n\n // set all unassigned points to be coarse\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n mark[i] = 0;\n\n if (cf_map_in[i] == UNASSIGNED)\n {\n cf_map_out[i] = COARSE;\n mark[i] = 1;\n }\n }\n\n // now run through newly marked coarse points and make them unassigned again if necessary\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n if (mark[i] > 0) // newly marked C point\n {\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n if (!s_con[j]) { continue; } // skip weak connections\n\n int jcol = A.col_indices[j];\n\n if (mark[jcol] && weights[jcol] > 1) // other point is also newly coarse\n {\n // if my weight is larger, they can't be a coarse point\n if (weights[i] > weights[jcol])\n {\n cf_map_out[jcol] = UNASSIGNED;\n }\n // if their weight is large, I can't be a coarse point\n else if (weights[jcol] > weights[i])\n {\n cf_map_out[i] = UNASSIGNED;\n }\n }\n }\n }\n }\n\n // finally count the number of coarse points\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n num_coarse += static_cast(cf_map_out[i] == COARSE);\n }\n\n return num_coarse;\n }\n\n /*************************************************************************\n * marks all neighbors of coarse points as fine\n ************************************************************************/\n static int markAdditionalFinePoints(const Matrix > &A,\n const Vector > &s_con,\n IndVector &cf_map_in,\n IndVector &cf_map_out)\n {\n int num_fine = 0;\n\n //deep copy cf_map_in into cf_map_out\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n cf_map_out[i] = cf_map_in[i];\n }\n\n //for each row of the matrix\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n //if already fine add it to the count\n num_fine += static_cast(cf_map_in[i] == FINE);\n\n if (cf_map_in[i] == UNASSIGNED)\n {\n // for each off diagonal, if I have a new strongly connected coarse neighbour\n // set to fine\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n if (!s_con[j]) { continue; } //skip weakly connected points\n\n int jcol = A.col_indices[j];\n\n if (cf_map_in[jcol] == COARSE)\n {\n cf_map_out[i] = FINE;\n num_fine++;\n break;\n }\n }\n }\n }\n\n return num_fine;\n }\n};\n\n/*************************************************************************\n * mark coarse and fine points\n ************************************************************************/\ntemplate \nvoid PMIS_Selector >\n::markCoarseFinePoints_1x1(Matrix_h &A,\n FVector &weights,\n const BVector &s_con,\n IVector &cf_map,\n IVector &scratch,\n int cf_map_init)\n{\n int num_coarse = 0, num_fine = 0;\n int num_strong_fine = 0;\n\n //mark all unconnected points as fine\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n int numj = A.row_offsets[i + 1] - A.row_offsets[i];\n\n if (numj == 0) //if there are no columns\n {\n cf_map[i] = FINE;\n }\n else if (numj == 1 && A.col_indices[A.row_offsets[i]] == i) //if there is only 1 column and it is the diagonal\n {\n cf_map[i] = FINE;\n }\n else if (weights[i] < 1) // no point strongly depends on point\n {\n cf_map[i] = FINE;\n }\n else\n {\n cf_map[i] = UNASSIGNED;\n }\n\n // check to see if a point has no strong connections to anyone\n // needs more justification / knowledge of everywhere this condition is applied\n bool isolated = true;\n\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n if (s_con[j])\n {\n isolated = false;\n break;\n }\n }\n\n if (isolated)\n {\n cf_map[i] = STRONG_FINE;\n num_strong_fine++;\n weights[i] = 0.F;\n }\n }\n\n // temp storage for marking coarse points\n IntVector mark(A.get_num_rows(), 0);\n int iterations = 0;\n\n do\n {\n ++iterations;\n num_coarse = MarkAdditionalPoints::markAdditionalCoarsePoints(A, weights, s_con, cf_map, scratch, mark);\n num_fine = MarkAdditionalPoints::markAdditionalFinePoints(A, s_con, scratch, cf_map);\n }\n while (num_coarse + num_fine + num_strong_fine < A.get_num_rows());\n\n // now we have \\hat{C} and the diagonals, we can start calculating the weights\n // remap cf_map (to be like in distance1)\n}\n\n/*************************************************************************\n * device kernels\n ************************************************************************/\n\n/*************************************************************************\n * Perform the initial marking for each row\n ************************************************************************/\ntemplate \n__global__\nvoid initialMarkingKernel(const IndexType *offsets, const IndexType *column_indices,\n const int numRows, int *cf_map, float *weights, const bool *s_con)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n // set initial C/F/U\n if (numj == 0) // no points\n {\n cf_map[tIdx] = FINE;\n }\n else if (numj == 1 && column_indices[offset] == tIdx) // only the diagonal\n {\n cf_map[tIdx] = FINE;\n }\n else if (weights[tIdx] < 1) // no strong dependancies\n {\n cf_map[tIdx] = FINE;\n }\n else\n {\n cf_map[tIdx] = UNASSIGNED;\n }\n\n bool isolated = true;\n\n for (int j = offsets[tIdx]; j < offsets[tIdx + 1]; j++)\n {\n if (s_con[j])\n {\n isolated = false;\n break;\n }\n }\n\n if (isolated)\n {\n cf_map[tIdx] = STRONG_FINE;\n weights[tIdx] = 0.F;\n }\n }\n}\n\ntemplate \n__global__\nvoid initialMarkingCFInit3Kernel(const IndexType *offsets, const IndexType *column_indices,\n const int numRows, int *cf_map, float *weights, const bool *s_con)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n // set initial C/F/U\n if (numj == 0) // no points\n {\n cf_map[tIdx] = FINE;\n }\n else if (numj == 1 && column_indices[offset] == tIdx) // only the diagonal\n {\n cf_map[tIdx] = FINE;\n }\n else if (weights[tIdx] < 1) // no strong dependancies\n {\n cf_map[tIdx] = FINE;\n }\n else\n {\n cf_map[tIdx] = UNASSIGNED;\n }\n\n bool isolated = true;\n\n for (int j = offsets[tIdx]; j < offsets[tIdx + 1]; j++)\n {\n if (s_con[j])\n {\n isolated = false;\n break;\n }\n }\n\n if (isolated)\n {\n cf_map[tIdx] = COARSE;\n weights[tIdx] = 0.F;\n }\n }\n}\n\n/*************************************************************************\n * Perform the initial marking for each row\n ************************************************************************/\ntemplate \n__global__\nvoid initialMarkingCfInitKernel(const IndexType *offsets, const IndexType *column_indices,\n const int numRows, int *cf_map, float *weights, const bool *s_con, IndexType *mark)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n // Check if on boundary, this is unmarking even coarse points\n for (int j = offsets[tIdx]; j < offsets[tIdx + 1]; j++)\n {\n if (column_indices[j] > numRows)\n {\n cf_map[tIdx] = UNASSIGNED;\n mark[tIdx] = 1;\n break;\n }\n }\n\n // set initial C/F/U\n if (numj == 0) // no points\n {\n cf_map[tIdx] = FINE;\n }\n else if (numj == 1 && column_indices[offset] == tIdx) // only the diagonal\n {\n cf_map[tIdx] = FINE;\n }\n else if (weights[tIdx] < 1) // no strong dependancies\n {\n cf_map[tIdx] = FINE;\n }\n else if (cf_map[tIdx] == STRONG_FINE)\n {\n // Do nothing, leave as strong fine\n weights[tIdx] = 0.F;\n }\n else if (cf_map[tIdx] == FINE)\n {\n // Fine point, mark as unassigned\n cf_map[tIdx] = UNASSIGNED;\n mark[tIdx] = 1;\n }\n }\n}\n\n\n/*************************************************************************\n * Mark every currently unassigned point as coarse\n ************************************************************************/\ntemplate \n__global__\nvoid markUnassignedAsCoarseKernel(const int numRows,\n const float *weights,\n IndexType *cf_map_in,\n IndexType *cf_map_out, IndexType *mark)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n const IndexType map_in = cf_map_in[tIdx];\n const int unassigned = map_in == UNASSIGNED ? 1 : 0;\n mark[tIdx] = unassigned;\n\n // Mark unassigned piont as coarse\n if (weights[tIdx] > 1.)\n {\n cf_map_out[tIdx] = unassigned ? COARSE : map_in;\n }\n else\n {\n cf_map_out[tIdx] = map_in;\n }\n }\n}\n\n/*************************************************************************\n * Unset all newly set coarse points who shouldn't be\n ************************************************************************/\ntemplate \n__global__\nvoid markAdditionalCoarsePointsKernel(const IndexType *offsets,\n const IndexType *column_indices,\n const int numRows, const float *weights,\n const bool *s_con,\n IndexType *cf_map_out, IndexType *mark)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x )\n {\n // run through newly marked coarse points & set to unassigned again if necessary\n if (mark[tIdx] > 0)\n {\n for ( int j = offsets[tIdx], j_end = offsets[tIdx + 1] ; j < j_end; j++ )\n {\n if (!s_con[j]) { continue; } //these are always false for edges pointing outside of the 1-ring halo\n\n int jcol = column_indices[j];\n float w_col = weights[jcol];\n float w_row = weights[tIdx];\n\n if ( mark[jcol] && w_col > 1.0f )\n {\n // If my weight is bigger, they can't be a coarse point\n if ( w_row > w_col )\n {\n cf_map_out[jcol] = UNASSIGNED;\n }\n // If their weight is bigger, I can't be a coarse point\n else if ( w_col > w_row )\n {\n cf_map_out[tIdx] = UNASSIGNED;\n }\n }\n }\n }\n }\n}\n\n/*************************************************************************\n * mark every neighbour of a newly assigned coarse point to be fine\n ************************************************************************/\ntemplate \n__global__\nvoid markAdditionalFinePointsKernel(const IndexType *offsets, const IndexType *column_indices,\n const int numRows, const bool *s_con,\n int *cf_map_in, int *cf_map_out)\n{\n for ( int tIdx = threadIdx.x + blockDim.x * blockIdx.x ; tIdx < numRows ; tIdx += blockDim.x * gridDim.x )\n {\n const int map_in = cf_map_in[tIdx];\n // go over each row\n bool fine = false;\n\n if ( map_in == UNASSIGNED )\n {\n // for each column\n for ( int j = offsets[tIdx], j_end = offsets[tIdx + 1] ; !fine && j < j_end ; ++j )\n {\n if ( !s_con[j] )\n {\n continue; // skip weakly connected points\n }\n\n int jcol = column_indices[j];\n fine = cf_map_in[jcol] == COARSE;\n }\n }\n\n cf_map_out[tIdx] = fine ? FINE : map_in;\n }\n}\n\ntemplate \nvoid PMIS_Selector >::markCoarseFinePoints_1x1(Matrix_d &A,\n FVector &weights,\n const BVector &s_con,\n IVector &cf_map,\n IVector &scratch,\n int cf_map_init)\n{\n if (A.hasProps(DIAG))\n {\n FatalError(\"Unsupported separate diag\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // Choose blocksize. Using 1 thread / row for now\n const int blockSize = 256;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) ((A.get_num_rows() + blockSize - 1) / blockSize));\n const int numRows = (int) A.get_num_rows();\n // raw pointers from the cusp arrays\n const IndexType *offsets_ptr = A.row_offsets.raw();\n const IndexType *column_indices_ptr = A.col_indices.raw();\n float *weights_ptr = weights.raw();\n const bool *s_con_ptr = s_con.raw();\n int *cf_map_ptr = cf_map.raw();\n int *scratch_ptr = scratch.raw();\n IVector mark(cf_map.size());\n // do the initial marking of points\n thrust_wrapper::fill(mark.begin(), mark.end(), 0);\n cudaCheckError();\n\n if (numRows > 0)\n {\n if (cf_map_init == 0)\n {\n initialMarkingKernel <<< numBlocks, blockSize>>>(offsets_ptr, column_indices_ptr, numRows, cf_map_ptr, weights_ptr, s_con_ptr);\n cudaCheckError();\n }\n else if (cf_map_init == 1)\n {\n initialMarkingCfInitKernel <<< numBlocks, blockSize>>>(offsets_ptr, column_indices_ptr, numRows, cf_map_ptr, weights_ptr, s_con_ptr, mark.raw());\n cudaCheckError();\n }\n else if (cf_map_init == 3)\n {\n initialMarkingCFInit3Kernel <<< numBlocks, blockSize>>>(offsets_ptr, column_indices_ptr, numRows, cf_map_ptr, weights_ptr, s_con_ptr);\n cudaCheckError();\n }\n else\n {\n FatalError(\"invalid value of cf_map_init in pmis selector, exiting\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n }\n\n cudaCheckError();\n // exchange initial cf_map for the 1st ring halos\n cf_map.dirtybit = 1;\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->exchange_halo(cf_map, cf_map.tag);\n }\n\n // count the # of coarse and fine points we have now\n int numUnassigned = 0;\n // array needed for marking of additional coarse points\n int one_ring_size, offset;\n\n if (!A.is_matrix_singleGPU())\n {\n A.getOffsetAndSizeForView(FULL, &offset, &one_ring_size);\n }\n else\n {\n one_ring_size = numRows;\n }\n\n const int numBlocks1Ring = std::min (AMGX_GRID_MAX_SIZE, (int) ((one_ring_size + blockSize - 1) / blockSize));\n // iterate until all points have been classified\n int numUnassignedMax = numRows;\n //int numUnassignedMaxPrevious;\n int iter = 0;\n\n do\n {\n //numUnassignedMaxPrevious = numUnassignedMax;\n if (numRows > 0)\n {\n if (iter || !cf_map_init)\n {\n // Laucnhed for 1-ring\n markUnassignedAsCoarseKernel <<< numBlocks1Ring, blockSize>>>(one_ring_size, weights_ptr, cf_map_ptr,\n scratch_ptr, mark.raw());\n cudaCheckError();\n markAdditionalCoarsePointsKernel <<< numBlocks, blockSize>>>(offsets_ptr, column_indices_ptr,\n numRows, weights_ptr, s_con_ptr,\n scratch_ptr, mark.raw());\n cudaCheckError();\n }\n else\n {\n amgx::thrust::copy(cf_map.begin(), cf_map.begin() + one_ring_size, scratch.begin());\n cudaCheckError();\n }\n\n // If a neighbor marked one of my node as unassigned, need to mark it\n // Taking the min (unassigned == -4, coarse = -1)\n scratch.dirtybit = 1;\n\n if (!A.is_matrix_singleGPU())\n {\n A.manager->min_from_halo(scratch, scratch.tag);\n }\n\n cudaCheckError();\n // Now exchange with neighbors scratch for my own nodes\n scratch.dirtybit = 1;\n\n if (!A.is_matrix_singleGPU())\n {\n // exchange updated CF map for halos (stored in scratch)\n A.manager->exchange_halo(scratch, scratch.tag);\n }\n\n cudaCheckError();\n markAdditionalFinePointsKernel <<< numBlocks, blockSize>>>(offsets_ptr, column_indices_ptr,\n numRows, s_con_ptr, scratch_ptr,\n cf_map_ptr);\n cudaCheckError();\n } // num rows > 0\n\n // count # of points still unassigned\n numUnassigned = (int) thrust_wrapper::count(cf_map.begin(), cf_map.begin() + numRows, (int)UNASSIGNED);\n cudaCheckError();\n numUnassignedMax = numUnassigned;\n cudaCheckError();\n cf_map.dirtybit = 1;\n\n if (A.is_matrix_distributed())\n {\n // exchange CF map for halos\n A.manager->exchange_halo(cf_map, cf_map.tag);\n // get max across all processes\n A.manager->getComms()->all_reduce_max(numUnassigned, numUnassignedMax);\n }\n\n iter++;\n }\n while ( numUnassignedMax != 0 );\n}\n\ntemplate \nvoid PMIS_SelectorBase< T_Config>::markCoarseFinePoints(Matrix< T_Config> &A,\n FVector &weights,\n const BVector &s_con,\n IVector &cf_map,\n IVector &scratch,\n int cf_map_init)\n{\n ViewType oldView = A.currentView();\n A.setView(OWNED);\n\n if (A.get_block_size() == 1)\n {\n markCoarseFinePoints_1x1(A, weights, s_con, cf_map, scratch, cf_map_init);\n }\n else\n {\n FatalError(\"Unsupported block size PMIS selector\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n}\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class PMIS_SelectorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class PMIS_Selector::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace classical\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/classical/selectors/pmis.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n\nnamespace amgx\n{\n\n/*\n * Count the # of non-zeros per row\n */\n__global__\nvoid nonZerosPerRowKernel(const int num_rows, const int *cf_map, const int *C_hat_start,\n const int *C_hat_end, int *nonZerosPerRow)\n{\n for (int tIdx = threadIdx.x + blockIdx.x * blockDim.x; tIdx < num_rows; tIdx += gridDim.x * blockDim.x)\n {\n int nonZeros = 0;\n\n if (cf_map[tIdx] >= 0)\n {\n nonZeros = 1;\n }\n else if (cf_map[tIdx] == STRONG_FINE)\n {\n nonZeros = 0;\n }\n else\n {\n nonZeros = C_hat_end[tIdx] - C_hat_start[tIdx];\n }\n\n nonZerosPerRow[tIdx] = nonZeros;\n }\n}\n\n__global__\nvoid nonZerosPerRowSizeKernel(const int num_rows, const int *cf_map,\n const int *C_hat_size, int *nonZerosPerRow)\n{\n for (int tIdx = threadIdx.x + blockIdx.x * blockDim.x; tIdx < num_rows; tIdx += gridDim.x * blockDim.x)\n {\n int nonZeros = 0;\n\n if (cf_map[tIdx] >= 0)\n {\n nonZeros = 1;\n }\n else if (cf_map[tIdx] == STRONG_FINE)\n {\n nonZeros = 0;\n }\n else\n {\n nonZeros = C_hat_size[tIdx];\n }\n\n nonZerosPerRow[tIdx] = nonZeros;\n }\n}\n\n\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\nnamespace amgx\n{\n\n/*\n * Count the # of non-zeros per row\n */\n__global__\nvoid nonZerosPerRowKernel(const int num_rows, const int *cf_map, const int *C_hat_start,\n const int *C_hat_end, int *nonZerosPerRow)\n{\n for (int tIdx = threadIdx.x + blockIdx.x * blockDim.x; tIdx < num_rows; tIdx += gridDim.x * blockDim.x)\n {\n int nonZeros = 0;\n\n if (cf_map[tIdx] >= 0)\n {\n nonZeros = 1;\n }\n else if (cf_map[tIdx] == STRONG_FINE)\n {\n nonZeros = 0;\n }\n else\n {\n nonZeros = C_hat_end[tIdx] - C_hat_start[tIdx];\n }\n\n nonZerosPerRow[tIdx] = nonZeros;\n }\n}\n\n__global__\nvoid nonZerosPerRowSizeKernel(const int num_rows, const int *cf_map,\n const int *C_hat_size, int *nonZerosPerRow)\n{\n for (int tIdx = threadIdx.x + blockIdx.x * blockDim.x; tIdx < num_rows; tIdx += gridDim.x * blockDim.x)\n {\n int nonZeros = 0;\n\n if (cf_map[tIdx] >= 0)\n {\n nonZeros = 1;\n }\n else if (cf_map[tIdx] == STRONG_FINE)\n {\n nonZeros = 0;\n }\n else\n {\n nonZeros = C_hat_size[tIdx];\n }\n\n nonZerosPerRow[tIdx] = nonZeros;\n }\n}\n\n\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/classical/interpolators/common.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\nusing std::set;\nusing std::vector;\n\ntemplate < typename IndexType, typename ValueType >\n__global__\nvoid check_P_values(IndexType *offsets, ValueType *val, int size, int flag)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n // Simply cast\n for (int i = offsets[tidx]; i < offsets[tidx + 1]; i++)\n {\n if (isnan(val[i])) { printf(\"%d, row = %d, i=%d\\n\", flag, tidx, i); }\n }\n }\n}\n\ntemplate \n__global__\nvoid fill_cf_map_global(const IndexType *cf_map, KeyType *cf_map_global, IndexType num_owned_fine_pts)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_owned_fine_pts ; tidx += blockDim.x * gridDim.x)\n {\n // Simply cast\n cf_map_global[tidx] = (KeyType) cf_map[tidx];\n }\n}\n\ntemplate \n__global__\nvoid fill_P_col_indices(const KeyType *global_cols, IndexType *cols, IndexType num_owned_fine_pts)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_owned_fine_pts ; tidx += blockDim.x * gridDim.x)\n {\n // Simply cast\n cols[tidx] = (IndexType) global_cols[tidx];\n }\n}\n\n\n\ntemplate \n__global__\nvoid createCfMapGlobal(const IndexType *cf_map, KeyType *cf_map_global, const KeyType my_part_offset, const IndexType num_owned_fine_pts)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_owned_fine_pts ; tidx += blockDim.x * gridDim.x)\n {\n // Only renumber the interior points\n if (cf_map[tidx] >= 0 && tidx < num_owned_fine_pts)\n {\n // Simply shift\n cf_map_global[tidx] = (KeyType) cf_map[tidx] /* my_local_id */ + my_part_offset;\n }\n else\n {\n cf_map_global[tidx] = -1;\n }\n }\n}\n\n\ntemplate\n__global__\nvoid initializeAssignedArray(const IndexType *cf_map, IndexType *assigned, const IndexType *A_rows, const IndexType *A_cols, const bool *s_con, IndexType *C_hat_offsets, int A_num_rows)\n{\n const int nWarps = cta_size / 32;\n // The coordinates of the thread inside the CTA/warp.\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n // Loop over rows of A.\n for ( int aRowId = blockIdx.x * nWarps + warpId ; aRowId < A_num_rows ; aRowId += gridDim.x * nWarps )\n {\n int coarse_fine_id = cf_map[aRowId];\n\n if ( coarse_fine_id >= 0)\n {\n assigned[aRowId] = 0;\n }\n else if (coarse_fine_id == FINE)\n {\n // Load A row IDs.\n int aColIt = A_rows[aRowId ];\n int aColEnd = A_rows[aRowId + 1];\n int coarse_count = 0;\n\n // Iterate over the columns of A.\n for ( aColIt += laneId ; utils::any( aColIt < aColEnd ) ; aColIt += 32 )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int aColId = aColIt < aColEnd ? A_cols[aColIt] : -1;\n // Is it an off-diagonal element.\n bool is_off_diagonal = aColIt < aColEnd && aColId != aRowId;\n // Is it fine.\n bool is_coarse = is_off_diagonal && cf_map[aColId] >= 0;\n // Is it strongly connected ?\n bool is_coarse_strongly_connected = aColIt < aColEnd && s_con[aColIt] && is_coarse;\n // Push coarse and strongly connected nodes in the set.\n coarse_count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n }\n\n // If fine point connected to coarse point, assigned = 1,\n if (coarse_count)\n {\n assigned[aRowId] = 1;\n C_hat_offsets[aRowId] = coarse_count;\n }\n }\n else // Strong fine\n {\n assigned[aRowId] = -1;\n }\n }\n}\n\ntemplate \n__global__\nvoid initializeAssignedArray_opt(const IndexType *cf_map, IndexType *assigned, const IndexType *A_rows, const IndexType *A_cols, const bool *s_con, IndexType *C_hat_offsets, int A_num_rows)\n{\n const int ngroups = cta_size / group_size;\n\n // The coordinates of the thread inside the CTA/group.\n const int group_id = threadIdx.x / group_size;\n const int lane_id = threadIdx.x % group_size;\n\n int a_row_id = blockIdx.x * ngroups + group_id;\n\n bool active = true;\n if(a_row_id > A_num_rows)\n {\n active = false;\n }\n\n int coarse_fine_id = UNASSIGNED;\n\n if(active)\n {\n coarse_fine_id = cf_map[a_row_id];\n\n if(coarse_fine_id >= 0)\n {\n if(lane_id == 0)\n {\n assigned[a_row_id] = 0;\n }\n\n active = false;\n }\n else if(coarse_fine_id != FINE) // Strong fine\n {\n if(lane_id == 0)\n {\n assigned[a_row_id] = -1;\n }\n\n active = false;\n }\n }\n\n active = (active && coarse_fine_id == FINE);\n\n constexpr int warp_size = 32;\n unsigned int base_mask = 0xFFFFFFFF >> (warp_size - group_size);\n int ngroups_in_warp = warp_size / group_size;\n int group_id_in_warp = group_id % ngroups_in_warp;\n unsigned int group_mask = base_mask << (group_id_in_warp*group_size);\n\n // Load A row IDs.\n int aColIt = A_rows[a_row_id ];\n int aColEnd = A_rows[a_row_id + 1];\n\n // Iterate over the columns of A.\n int coarse_count = 0;\n for ( aColIt += lane_id ; utils::any( active && aColIt < aColEnd ) ; aColIt += group_size )\n {\n bool is_coarse_strongly_connected = false;\n\n if(active)\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int aColId = aColIt < aColEnd ? A_cols[aColIt] : -1;\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = aColIt < aColEnd && aColId != a_row_id;\n\n // Is it fine.\n bool is_coarse = is_off_diagonal && cf_map[aColId] >= 0;\n\n // Is it strongly connected ?\n is_coarse_strongly_connected = s_con[aColIt] && is_coarse;\n }\n\n // Count any of the threads in the group that had strong connections\n int coarse_exist_flags = utils::ballot(is_coarse_strongly_connected);\n if(lane_id == 0)\n {\n coarse_count += __popc(coarse_exist_flags & group_mask);\n }\n }\n\n // If fine point connected to coarse point, assigned = 1,\n if (active && coarse_count > 0)\n {\n if(lane_id == 0)\n {\n assigned[a_row_id] = 1;\n C_hat_offsets[a_row_id] = coarse_count;\n }\n }\n}\n\ntemplate\n__global__\nvoid fillAssignedArray(IndexType *assigned, const IndexType *A_rows, const IndexType *A_cols, const bool *s_con, int A_num_rows, int pass)\n{\n const int nWarps = cta_size / 32;\n // The coordinates of the thread inside the CTA/warp.\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n // Loop over rows of A.\n for ( int aRowId = blockIdx.x * nWarps + warpId ; aRowId < A_num_rows ; aRowId += gridDim.x * nWarps )\n {\n if ( assigned[aRowId] == -1)\n {\n // Load A row IDs.\n int aColIt = A_rows[aRowId ];\n int aColEnd = A_rows[aRowId + 1];\n int assigned_count = 0;\n\n // Iterate over the columns of A.\n for ( aColIt += laneId ; utils::any( aColIt < aColEnd ) ; aColIt += 32 )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int aColId = aColIt < aColEnd ? A_cols[aColIt] : -1;\n // Is it an off-diagonal element.\n bool is_off_diagonal = aColIt < aColEnd && aColId != aRowId;\n // Is it assigned\n bool is_assigned = is_off_diagonal && assigned[aColId] == pass - 1;\n // Is it strongly connected ?\n bool is_assigned_strongly_connected = aColIt < aColEnd && s_con[aColIt] && is_assigned;\n // Push coarse and strongly connected nodes in the set.\n assigned_count += __popc( utils::ballot( is_assigned_strongly_connected ) );\n\n if (assigned_count) { break; }\n }\n\n // If one of the strong neighbor has been assigned, set assigned to current pass number\n if (assigned_count)\n {\n assigned[aRowId] = pass;\n }\n }\n }\n}\n\ntemplate < typename IndexType, int cta_size, int group_size >\n__global__\nvoid fillAssignedArray_opt(IndexType *assigned, const IndexType *A_rows, const IndexType *A_cols, const bool *s_con, int A_num_rows, int pass)\n{\n // The coordinates of the thread inside the CTA/groups\n const int group_id = threadIdx.x / group_size;\n const int lane_id = threadIdx.x % group_size;\n const int ngroups_per_block = cta_size / group_size;\n\n int a_row_id = blockIdx.x * ngroups_per_block + group_id;\n\n bool active = true;\n if (a_row_id >= A_num_rows || assigned[a_row_id] != -1)\n {\n active = false;\n }\n\n // Mask for the subgroup\n constexpr int warp_size = 32;\n unsigned int base_mask = 0xFFFFFFFF >> (warp_size - group_size);\n int ngroups_in_warp = warp_size / group_size;\n int group_id_in_warp = group_id % ngroups_in_warp;\n unsigned int group_mask = base_mask << (group_id_in_warp*group_size);\n\n // Load A row IDs.\n int a_col_begin = 0;\n int a_col_end = 0;\n\n if(active)\n {\n a_col_begin = A_rows[a_row_id ];\n a_col_end = A_rows[a_row_id + 1];\n }\n\n // Iterate over the columns of A.\n for (int a_col_it = a_col_begin + lane_id ; utils::any(active && a_col_it < a_col_end); a_col_it += group_size)\n {\n bool is_assigned_strongly_connected = false;\n\n if(active)\n {\n // Columns of A maps to rows of B. Each thread of the group loads its A-col/B-row ID.\n int a_col_id = a_col_it < a_col_end ? A_cols[a_col_it] : -1;\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n\n // Is it assigned\n bool is_assigned = is_off_diagonal && assigned[a_col_id] == pass - 1;\n\n // Is it strongly connected ?\n is_assigned_strongly_connected = s_con[a_col_it] && is_assigned;\n }\n\n int group_strongly_connected = utils::ballot(is_assigned_strongly_connected);\n\n // Push coarse and strongly connected nodes in the set.\n if(active && (group_strongly_connected & group_mask))\n {\n if(lane_id == 0)\n {\n assigned[a_row_id] = pass;\n }\n\n active = false;\n }\n }\n}\n\n\n/*\n * remove all duplicates from a vector\n */\ntemplate \nvoid removeDuplicates(vector &vec)\n{\n std::sort(vec.begin(), vec.end());\n vec.erase(std::unique(vec.begin(), vec.end()), vec.end());\n}\n\n/*\n * find the value at a particular column of a row if it exists\n */\ntemplate \ntypename VVector::value_type findValue(const IVector &columns,\n const VVector &values,\n const int row_begin, const int row_end, const int column_needed)\n{\n typename VVector::value_type v = 0;\n\n // NAIVE\n for (int k = row_begin; k < row_end; k++)\n {\n int kcol = columns[k];\n\n if (kcol == column_needed)\n {\n v = values[k];\n }\n }\n\n return v;\n}\n\n\n\n/*************************************************************************\n* create the interpolation matrix P\n************************************************************************/\ntemplate \nvoid Multipass_Interpolator >::generateInterpolationMatrix_1x1(Matrix_h &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix_h &P)\n{\n FatalError(\"multipass interpolation not implemented on host\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n} // end multipass interpolator\n\n\n///////////////////////////////////////////////////////////////////////////////\n\n// DEVICE CODE\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#include \n#include // Included inside the namespace to solve name colisions.\n\n__device__ __forceinline__ int get_work( int *queue, int warp_id )\n{\n int offset = -1;\n\n if ( utils::lane_id() == 0 )\n {\n offset = atomicAdd( queue, 1 );\n }\n\n return utils::shfl( offset, 0 );\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nnamespace multipass\n{\n\n#include \n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nestimate_c_hat_size_kernel( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const bool *s_con,\n const int *assigned,\n int *C_hat_offsets,\n int pass )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Number of items per grid.\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n volatile int a_row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += NUM_WARPS_PER_GRID )\n {\n // Skip coarse rows and fine rows touching a coarse point\n if (assigned[a_row_id] != pass)\n {\n continue;\n }\n\n // Load A row IDs.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // The number of elements.\n int my_count = 0;\n\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // assigned == pass-1\n bool is_assigned = is_off_diagonal && (assigned[a_col_id] == pass - 1);\n // Is it a coarse, strongly-connected column and already assigned?\n bool is_assigned_strongly_connected = is_strongly_connected && is_assigned;\n // If so, add their C_hat_offset value to my_count\n my_count += (is_assigned_strongly_connected ? C_hat_offsets[a_col_id] : 0);\n }\n\n // Do reduction\n#pragma unroll\n\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n my_count += utils::shfl_xor( my_count, mask );\n }\n\n // Write result -- RACE CONDITION!\n if (lane_id == 0)\n {\n C_hat_offsets[a_row_id] = my_count;\n }\n }\n}\n\ntemplate< int CTA_SIZE, int group_size >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nestimate_c_hat_size_kernel_opt( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const bool *s_con,\n const int *assigned,\n int *C_hat_offsets,\n int pass )\n{\n const int group_id = threadIdx.x / group_size;\n const int lane_id = threadIdx.x % group_size;\n\n int ngroups = CTA_SIZE / group_size;\n int a_row_id = blockIdx.x * ngroups + group_id;\n\n // Skip coarse rows and fine rows touching a coarse point\n bool active = true;\n if (a_row_id >= A_num_rows || assigned[a_row_id] != pass)\n {\n active = false;\n }\n\n int my_count = 0;\n\n // Load A row IDs.\n if(active)\n {\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Iterate over the columns of A to build C_hat.\n for (int a_col_it = a_col_begin + lane_id; a_col_it < a_col_end; a_col_it += group_size )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = A_cols[a_col_it];\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_id != a_row_id;\n\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n\n // assigned == pass-1\n bool is_assigned = is_off_diagonal && (assigned[a_col_id] == pass - 1);\n\n // Is it a coarse, strongly-connected column and already assigned?\n bool is_assigned_strongly_connected = is_strongly_connected && is_assigned;\n\n // If so, add their C_hat_offset value to my_count\n my_count += (is_assigned_strongly_connected ? C_hat_offsets[a_col_id] : 0);\n }\n }\n\n#pragma unroll\n for(int i = group_size / 2; i > 0; i /= 2)\n {\n my_count += utils::shfl_down(my_count, i);\n }\n\n if (active && lane_id == 0)\n {\n C_hat_offsets[a_row_id] = my_count;\n }\n}\n\n template< int CTA_SIZE, int WARP_SIZE, typename KeyType >\n__global__ __launch_bounds__( CTA_SIZE )\n void\ncompute_c_hat_first_pass_kernel( int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const bool *__restrict s_con,\n const KeyType *__restrict cf_map_global,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_size,\n KeyType *__restrict C_hat,\n int *__restrict assigned)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += NUM_WARPS_PER_GRID)\n {\n // Skip if not in current pass\n if (assigned[a_row_id] != 1)\n {\n continue;\n }\n\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_begin = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n int count = 0;\n int c_col_it = C_hat_start[a_row_id];\n\n // _iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = is_active && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // assigned == pass -1\n bool is_assigned = is_off_diagonal && (assigned[a_col_id] == 0);\n // Is it an assigned==pass -1 and strongly-connected column.\n bool is_assigned_strongly_connected = is_strongly_connected && is_assigned;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_assigned_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_assigned_strongly_connected )\n {\n //if (cf_map[a_col_id] < 0)\n // printf(\"error in multipass 1\\n\");\n C_hat[c_col_it + count + dest] = cf_map_global[a_col_id];\n }\n\n count += __popc( vote );\n }\n\n if ( lane_id == 0 )\n {\n C_hat_size[a_row_id] = count;\n }\n }\n}\n\n template< int CTA_SIZE, int WARP_SIZE, int group_size, typename KeyType>\n__global__ __launch_bounds__( CTA_SIZE )\n void\ncompute_c_hat_first_pass_kernel_opt( int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const bool *__restrict s_con,\n const KeyType *__restrict cf_map_global,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_size,\n KeyType *__restrict C_hat,\n int *__restrict assigned)\n{\n const int ngroups = CTA_SIZE / group_size;\n\n // The coordinates of the thread inside the CTA/warp.\n const int group_id = threadIdx.x / group_size;\n const int lane_id = threadIdx.x % group_size;\n\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * ngroups + group_id;\n\n bool active = true;\n\n // Skip if not in current pass\n if (a_row_id >= A_num_rows || assigned[a_row_id] != 1)\n {\n active = false;\n }\n\n int a_col_begin = 0;\n int a_col_end = 0;\n int c_col_it = 0;\n\n if(active)\n {\n a_col_begin = A_rows[a_row_id];\n a_col_end = A_rows[a_row_id + 1];\n c_col_it = C_hat_start[a_row_id];\n }\n\n unsigned int base_mask = 0xFFFFFFFF >> (WARP_SIZE - group_size);\n int ngroups_in_warp = WARP_SIZE / group_size;\n int group_id_in_warp = group_id % ngroups_in_warp;\n unsigned int group_mask = base_mask << (group_id_in_warp*group_size);\n\n // _iterate over the columns of A to build C_hat.\n int count = 0;\n for ( int a_col_it = a_col_begin + lane_id ; utils::any(active && a_col_it < a_col_end) ; a_col_it += group_size )\n {\n bool is_assigned_strongly_connected = false;\n\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if(active)\n {\n if (a_col_it < a_col_end)\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = (a_col_it < a_col_end) && (a_col_id != a_row_id);\n\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n\n // assigned == pass -1\n bool is_assigned = is_off_diagonal && (assigned[a_col_id] == 0);\n\n // Is it an assigned==pass -1 and strongly-connected column.\n is_assigned_strongly_connected = is_strongly_connected && is_assigned;\n }\n\n // We collect fine and strongly-collected columns.\n int all_vote = utils::ballot( is_assigned_strongly_connected );\n int group_vote = all_vote & group_mask;\n int dest = __popc( group_vote & utils::lane_mask_lt() );\n\n if ( active && is_assigned_strongly_connected )\n {\n C_hat[c_col_it + count + dest] = cf_map_global[a_col_id];\n }\n\n count += __popc( group_vote );\n }\n\n if ( active && lane_id == 0 )\n {\n C_hat_size[a_row_id] = count;\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, typename KeyType >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_c_hat_kernel( int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const bool *__restrict s_con,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_size,\n KeyType *__restrict C_hat,\n int *__restrict C_hat_pos,\n int *__restrict assigned,\n int gmem_size,\n KeyType *g_keys,\n int *wk_work_queue,\n int *wk_status,\n int pass )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // Shared memory to vote.\n __shared__ int s_b_row_ids[CTA_SIZE];\n // The hash keys stored in shared memory.\n __shared__ KeyType s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Skip if not in current pass\n if (assigned[a_row_id] != pass)\n {\n continue;\n }\n\n // Clear the set.\n set.clear();\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_begin = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // _iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = is_active && a_col_id != a_row_id;\n // Check if assigned == pass -1\n bool is_assigned = is_off_diagonal && (assigned[a_col_id] == pass - 1);\n // Check if we have a strong connection to the point\n bool is_strongly_connected = is_active && s_con[a_col_it];\n // Check if strong connection to assigned == pass-1\n bool is_assigned_strongly_connected = is_assigned && is_strongly_connected;\n // We collect columns for which assigned == pass - 1\n int vote = utils::ballot( is_assigned_strongly_connected);\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_assigned_strongly_connected)\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n __syncwarp();\n\n int num_rows = __popc( vote );\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int uniform_b_row_id = -1;\n\n if ( is_active_k )\n {\n uniform_b_row_id = s_b_row_ids[warp_id * WARP_SIZE + local_k];\n }\n\n // Load the range of the row of B.\n int b_col_it = 0, b_col_end = 0;\n\n if ( is_active_k )\n {\n // this loops over elements of C_hat for row uniform_b_row_id\n b_col_it = C_hat_start[uniform_b_row_id];\n b_col_end = b_col_it + C_hat_size[uniform_b_row_id];\n }\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n // The ID of the interpolatory point\n KeyType b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = C_hat[b_col_it];\n }\n\n // Is it valid\n is_off_diagonal = b_col_it < b_col_end; // && b_col_id != local_k;\n // Push node to set\n KeyType b_item = -1;\n\n if ( is_off_diagonal )\n {\n b_item = b_col_id;\n }\n\n // Insert the interpolatory point to the set\n set.insert( b_item, wk_status );\n }\n }\n }\n\n int c_col_it = C_hat_start[a_row_id];\n int count = set.store_with_positions( &C_hat[c_col_it], &C_hat_pos[c_col_it] );\n\n if ( lane_id == 0 )\n {\n C_hat_size[a_row_id] = count;\n }\n }\n}\n\ntemplate \n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_c_hat_kernel_opt( int assigned_set_size,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const bool *__restrict s_con,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_size,\n KeyType *__restrict C_hat,\n int *__restrict C_hat_pos,\n int *__restrict assigned,\n int *__restrict assigned_set,\n int pass )\n{\n constexpr KeyType SLOT_VACANT = -1;\n\n // Group indices\n constexpr int ngroups = CTA_SIZE / GROUP_SIZE;\n const int group_id = threadIdx.x / GROUP_SIZE;\n const int lane_id = threadIdx.x % GROUP_SIZE;\n\n const int gid = blockIdx.x * ngroups + group_id;\n\n // Dynamic sized shared memory\n extern __shared__ int s[];\n\n // Block-level hash container storage\n KeyType* key_s = (KeyType*)s;\n\n //KeyType* data_s = (KeyType*)&key_s[ngroups*HASH_SIZE];\n int* col_ind_s = (int*)&key_s[ngroups*HASH_SIZE];\n\n // Group-level hash containers\n KeyType* key_group_s = &key_s[group_id*HASH_SIZE];\n //KeyType* data_group_s = &data_s[group_id*HASH_SIZE];\n\n // Initialise the keys and values.\n#pragma unroll\n for(int i = threadIdx.x; i < ngroups*HASH_SIZE; i += CTA_SIZE)\n {\n key_s[i] = SLOT_VACANT; // Inserted keys will be in range [0,N]\n //data_s[i] = 0; // We will sum into values\n }\n\n if(lane_id == 0)\n {\n col_ind_s[group_id] = 0;\n }\n\n __syncthreads();\n\n int a_row_id = -1;\n\n if(gid < assigned_set_size)\n {\n // One row of A per group\n a_row_id = assigned_set[gid];\n\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = A_rows[a_row_id] + lane_id ; a_col_it < A_rows[a_row_id + 1]; a_col_it += GROUP_SIZE )\n {\n int a_col_id = A_cols[a_col_it];\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_id != a_row_id;\n\n // Check if assigned == pass -1\n bool is_assigned = is_off_diagonal && (assigned[a_col_id] == pass - 1);\n\n // Check if we have a strong connection to the point\n bool is_strongly_connected = s_con[a_col_it];\n\n // Check if strong connection to assigned == pass-1\n bool is_assigned_strongly_connected = is_assigned && is_strongly_connected;\n\n if(!is_assigned_strongly_connected) continue;\n\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int b_col_beg = C_hat_start[a_col_id];\n\n // Iterate over the range of columns of B.\n for (int b_col_it = b_col_beg; b_col_it < b_col_beg + C_hat_size[a_col_id]; ++b_col_it)\n {\n KeyType key = C_hat[b_col_it];\n\n int hash = key % HASH_SIZE;\n\n while(true)\n {\n // XXX You can just avoid this and directly atomicCAS, which is faster?\n int curr_key = key_group_s[hash];\n if(curr_key == key) break;\n \n if(curr_key == SLOT_VACANT)\n {\n KeyType old_key = (KeyType)atomicCAS((unsigned long long*)&key_group_s[hash], (unsigned long long)SLOT_VACANT, (unsigned long long)key);\n if(old_key == SLOT_VACANT || old_key == key)\n {\n break;\n }\n }\n\n // We did not secure a slot, so linear probe to next slot\n hash = (hash + 1) % HASH_SIZE;\n }\n }\n }\n }\n\n __syncthreads();\n\n // XXX Could change to poll by ballot style compress\n if(gid < assigned_set_size)\n {\n // Store the results.\n int c_col_it = C_hat_start[a_row_id];\n\n#pragma unroll\n for(int i = lane_id; i < HASH_SIZE; i += GROUP_SIZE)\n {\n KeyType key = key_group_s[i];\n if(key != SLOT_VACANT)\n {\n int ind = atomicAdd(&col_ind_s[group_id], 1);\n C_hat[c_col_it + ind] = key;\n //C_hat_pos[c_col_it + ind] = i;\n }\n }\n }\n\n __syncthreads();\n\n if(gid < assigned_set_size)\n {\n if(lane_id == 0)\n {\n C_hat_size[a_row_id] = col_ind_s[group_id];\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, typename KeyType >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_interp_weight_first_pass_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const KeyType *__restrict cf_map_global,\n const bool *__restrict s_con,\n const Value_type *__restrict diag,\n const int *__restrict P_rows,\n KeyType *P_cols,\n Value_type *P_vals,\n const int *__restrict assigned)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / 32;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += NUM_WARPS_PER_GRID)\n {\n // Only do work if assigned[row_id] == pass\n KeyType coarse_fine_id = cf_map_global[a_row_id];\n\n // If coarse rows\n if ( assigned[a_row_id] == 0 )\n {\n if ( lane_id == 0 )\n {\n int p_row_it = P_rows[a_row_id];\n P_cols[p_row_it] = coarse_fine_id;\n P_vals[p_row_it] = Value_type( 1 );\n }\n }\n // If strong fine\n else if ( assigned[a_row_id] == -1 )\n {\n // Do nothing\n }\n // strongly connected to a coarse point\n else if (assigned[a_row_id] == 1)\n {\n // Load A row IDs.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n int p_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n p_col_tmp = P_rows[a_row_id + lane_id];\n }\n\n int p_col_it = utils::shfl( p_col_tmp, 0 );\n int p_col_end = utils::shfl( p_col_tmp, 1 );\n\n // Weak value.\n Value_type sum_N(0), sum_C(0), alfa(0);\n int count = 0;\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id(-1);\n Value_type a_value(0);\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n a_value = A_vals[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it a strongly-connected column.\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it a weakly connected node.\n //bool is_weakly_connected = is_off_diagonal && !is_strongly_connected;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map_global[a_col_id] == STRONG_FINE;\n // Is it a fine and strongly-connected column.\n bool is_assigned_strongly_connected = is_strongly_connected && assigned[a_col_id] == 0;\n\n // Update the weak value.\n if ( is_off_diagonal && !is_strong_fine)\n {\n sum_N += a_value;\n }\n\n if ( is_assigned_strongly_connected )\n {\n sum_C += a_value;\n }\n\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_assigned_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_assigned_strongly_connected )\n {\n P_cols[p_col_it + count + dest] = cf_map_global[a_col_id];\n P_vals[p_col_it + count + dest] = a_value;\n }\n\n count += __popc( vote );\n }\n\n // We're done with that row of A, now reduce sum_N and sum_C\n#pragma unroll\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n sum_C += utils::shfl_xor( sum_C, mask );\n sum_N += utils::shfl_xor( sum_N, mask );\n }\n\n if ( lane_id == 0 )\n {\n // NOTE this matches the check for zero in HYPRE\n double div = (fabs(sum_C * diag[a_row_id]) == 0.0) ? 1. : sum_C * diag[a_row_id];\n alfa = -sum_N / div;\n }\n\n alfa = utils::shfl( alfa, 0 );\n\n // Scale the value of P\n for ( p_col_it += lane_id ; utils::any( p_col_it < p_col_end ) ; p_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n if ( p_col_it < p_col_end )\n {\n P_vals[p_col_it] *= alfa;\n }\n }\n } // if (assigned[a_row_id] == 1)\n } // Loop over rows\n}\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, typename KeyType >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_interp_weight_first_pass_kernel_opt( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const KeyType *__restrict cf_map_global,\n const bool *__restrict s_con,\n const Value_type *__restrict diag,\n const int *__restrict P_rows,\n KeyType *P_cols,\n Value_type *P_vals,\n const int *__restrict assigned)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / 32;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += NUM_WARPS_PER_GRID)\n {\n // Only do work if assigned[row_id] == pass\n KeyType coarse_fine_id = cf_map_global[a_row_id];\n\n // If coarse rows\n if ( assigned[a_row_id] == 0 )\n {\n if ( lane_id == 0 )\n {\n int p_row_it = P_rows[a_row_id];\n P_cols[p_row_it] = coarse_fine_id;\n P_vals[p_row_it] = Value_type( 1 );\n }\n }\n // If strong fine\n else if ( assigned[a_row_id] == -1 )\n {\n // Do nothing\n }\n // strongly connected to a coarse point\n else if (assigned[a_row_id] == 1)\n {\n // Load A row IDs.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n int p_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n p_col_tmp = P_rows[a_row_id + lane_id];\n }\n\n int p_col_it = utils::shfl( p_col_tmp, 0 );\n int p_col_end = utils::shfl( p_col_tmp, 1 );\n\n // Weak value.\n Value_type sum_N(0), sum_C(0), alfa(0);\n int count = 0;\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id(-1);\n Value_type a_value(0);\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n a_value = A_vals[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it a strongly-connected column.\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it a weakly connected node.\n //bool is_weakly_connected = is_off_diagonal && !is_strongly_connected;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map_global[a_col_id] == STRONG_FINE;\n // Is it a fine and strongly-connected column.\n bool is_assigned_strongly_connected = is_strongly_connected && assigned[a_col_id] == 0;\n\n // Update the weak value.\n if ( is_off_diagonal && !is_strong_fine)\n {\n sum_N += a_value;\n }\n\n if ( is_assigned_strongly_connected )\n {\n sum_C += a_value;\n }\n\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_assigned_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_assigned_strongly_connected )\n {\n P_cols[p_col_it + count + dest] = cf_map_global[a_col_id];\n P_vals[p_col_it + count + dest] = a_value;\n }\n\n count += __popc( vote );\n }\n\n // We're done with that row of A, now reduce sum_N and sum_C\n#pragma unroll\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n sum_C += utils::shfl_xor( sum_C, mask );\n sum_N += utils::shfl_xor( sum_N, mask );\n }\n\n if ( lane_id == 0 )\n {\n // NOTE this matches the check for zero in HYPRE\n double div = (fabs(sum_C * diag[a_row_id]) == 0.0) ? 1. : sum_C * diag[a_row_id];\n alfa = -sum_N / div;\n }\n\n alfa = utils::shfl( alfa, 0 );\n\n // Scale the value of P\n for ( p_col_it += lane_id ; utils::any( p_col_it < p_col_end ) ; p_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n if ( p_col_it < p_col_end )\n {\n P_vals[p_col_it] *= alfa;\n }\n }\n } // if (assigned[a_row_id] == 1)\n } // Loop over rows\n}\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, typename KeyType >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_interp_weight_kernel_opt( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const KeyType *__restrict C_hat,\n const int *__restrict C_hat_pos,\n const int *__restrict C_hat_start,\n const int *__restrict C_hat_size,\n const Value_type *__restrict diag,\n const int *__restrict P_rows,\n KeyType *P_cols,\n Value_type *P_vals,\n const int gmem_size,\n KeyType *g_keys,\n Value_type *g_vals,\n int *wk_work_queue,\n const int *__restrict assigned,\n const int *__restrict assigned_set,\n int pass)\n{\n const int NUM_WARPS = CTA_SIZE / 32;\n // The hash keys stored in shared memory.\n __shared__ KeyType s_keys[NUM_WARPS * SMEM_SIZE];\n // A shared location where threads propose a row of B to load.\n __shared__ int s_b_row_ids[CTA_SIZE];\n // A shared location where threads store a value of B to load.\n __shared__ Value_type s_b_values[CTA_SIZE];\n // The hash values stored in shared memory.\n __shared__ Value_type s_vals[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_map map( &s_keys[warp_id * SMEM_SIZE],\n &g_keys[a * gmem_size],\n &s_vals[warp_id * SMEM_SIZE],\n &g_vals[a * gmem_size],\n gmem_size );\n // Loop over rows of A.\n for ( ; a < A_num_rows ; a += gridDim.x*NUM_WARPS )\n {\n int a_row_id = assigned_set[a];\n\n // Clear the table.\n map.clear();\n // Rebuild C_hat.\n int c_hat_it = C_hat_start[a_row_id];\n int c_hat_end = c_hat_it + C_hat_size[a_row_id];\n map.load( c_hat_end - c_hat_it, &C_hat[c_hat_it], &C_hat_pos[c_hat_it] );\n // Load A row IDs.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // sums\n Value_type sum_N(0), sum_C(0), alfa(0);\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id(-1);\n Value_type a_value(0);\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n a_value = A_vals[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it a strongly-connected column.\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a fine and strongly-connected column.\n bool is_assigned_strongly_connected = is_strongly_connected && assigned[a_col_id] == pass - 1;\n // Is it a weakly connected node.\n bool is_weakly_connected = is_off_diagonal && !is_assigned_strongly_connected;\n\n // Update the weak value.\n if ( is_weakly_connected && !is_strong_fine)\n {\n sum_N += a_value;\n }\n\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_assigned_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_assigned_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n s_b_values[warp_id * WARP_SIZE + dest] = a_value;\n }\n\n int num_rows = __popc( vote );\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [b_col_it, bCol_end).\n int b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n Value_type uniform_val = s_b_values[warp_id * WARP_SIZE + k];\n // TODO: make sure we have better memory accesses. b_colBegin rather than bCol_it because we iterate twice.\n int b_col_begin = P_rows[b_row_id ];\n int b_col_end = P_rows[b_row_id + 1];\n\n // _iterate over the range of columns of B.\n for ( int b_col_it = b_col_begin + lane_id ; utils::any( b_col_it < b_col_end ) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n KeyType b_col_id(-1);\n Value_type tmp(0);\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = P_cols[b_col_it];\n tmp = P_vals[b_col_it] * uniform_val;\n sum_C += tmp;\n sum_N += tmp;\n }\n\n // Update C_hat values. If the value is not in C_hat it will be skipped (true parameter).\n // Need to get fine level id for b_col_id\n // if (a_row_id == 8 && b_col_id == 3) printf(\"updating 8,3: %lg\\n\",tmp);\n bool found = map.update( b_col_id, tmp );\n bool any_found = (__popc(utils::ballot(found)));\n }\n }\n }\n\n // We're done with that row of A, now reduce sum_N and sum_C\n#pragma unroll\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n sum_C += utils::shfl_xor( sum_C, mask );\n sum_N += utils::shfl_xor( sum_N, mask );\n }\n\n if ( lane_id == 0 )\n {\n // NOTE this matches the check for zero in HYPRE\n double div = (fabs(sum_C * diag[a_row_id]) == 0.0) ? 1. : sum_C * diag[a_row_id];\n alfa = -sum_N / div;\n // alfa = -sum_N/(sum_C*diag[a_row_id]);\n }\n\n alfa = utils::shfl( alfa, 0 );\n\n int p_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n p_col_tmp = P_rows[a_row_id + lane_id];\n }\n\n int p_col_it = utils::shfl( p_col_tmp, 0 );\n int p_col_end = utils::shfl( p_col_tmp, 1 );\n\n map.store_keys_scale_values( p_col_end - p_col_it, &P_cols[p_col_it], alfa, &P_vals[p_col_it] );\n }\n}\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, typename KeyType >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_interp_weight_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const KeyType *__restrict C_hat,\n const int *__restrict C_hat_pos,\n const int *__restrict C_hat_start,\n const int *__restrict C_hat_size,\n const Value_type *__restrict diag,\n const int *__restrict P_rows,\n KeyType *P_cols,\n Value_type *P_vals,\n const int gmem_size,\n KeyType *g_keys,\n Value_type *g_vals,\n int *wk_work_queue,\n const int *__restrict assigned,\n int pass)\n{\n const int NUM_WARPS = CTA_SIZE / 32;\n // The hash keys stored in shared memory.\n __shared__ KeyType s_keys[NUM_WARPS * SMEM_SIZE];\n // A shared location where threads propose a row of B to load.\n __shared__ int s_b_row_ids[CTA_SIZE];\n // A shared location where threads store a value of B to load.\n __shared__ Value_type s_b_values[CTA_SIZE];\n // The hash values stored in shared memory.\n __shared__ Value_type s_vals[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_map map( &s_keys[warp_id * SMEM_SIZE],\n &g_keys[a_row_id * gmem_size],\n &s_vals[warp_id * SMEM_SIZE],\n &g_vals[a_row_id * gmem_size],\n gmem_size );\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Only do work if assigned[row_id] == pass\n if (assigned[a_row_id] != pass)\n {\n continue;\n }\n\n // Clear the table.\n map.clear();\n // Rebuild C_hat.\n int c_hat_it = C_hat_start[a_row_id];\n int c_hat_end = c_hat_it + C_hat_size[a_row_id];\n map.load( c_hat_end - c_hat_it, &C_hat[c_hat_it], &C_hat_pos[c_hat_it] );\n // Load A row IDs.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // sums\n Value_type sum_N(0), sum_C(0), alfa(0);\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id(-1);\n Value_type a_value(0);\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n a_value = A_vals[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it a strongly-connected column.\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a fine and strongly-connected column.\n bool is_assigned_strongly_connected = is_strongly_connected && assigned[a_col_id] == pass - 1;\n // Is it a weakly connected node.\n bool is_weakly_connected = is_off_diagonal && !is_assigned_strongly_connected;\n\n // Update the weak value.\n if ( is_weakly_connected && !is_strong_fine)\n {\n sum_N += a_value;\n }\n\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_assigned_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_assigned_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n s_b_values[warp_id * WARP_SIZE + dest] = a_value;\n }\n\n __syncwarp();\n\n int num_rows = __popc( vote );\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [b_col_it, bCol_end).\n int b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n Value_type uniform_val = s_b_values[warp_id * WARP_SIZE + k];\n // TODO: make sure we have better memory accesses. b_colBegin rather than bCol_it because we iterate twice.\n int b_col_begin = P_rows[b_row_id ];\n int b_col_end = P_rows[b_row_id + 1];\n\n // _iterate over the range of columns of B.\n for ( int b_col_it = b_col_begin + lane_id ; utils::any( b_col_it < b_col_end ) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n KeyType b_col_id(-1);\n Value_type tmp(0);\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = P_cols[b_col_it];\n tmp = P_vals[b_col_it] * uniform_val;\n sum_C += tmp;\n sum_N += tmp;\n }\n\n // Update C_hat values. If the value is not in C_hat it will be skipped (true parameter).\n // Need to get fine level id for b_col_id\n // if (a_row_id == 8 && b_col_id == 3) printf(\"updating 8,3: %lg\\n\",tmp);\n bool found = map.update( b_col_id, tmp );\n bool any_found = (__popc(utils::ballot(found)));\n }\n }\n }\n\n // We're done with that row of A, now reduce sum_N and sum_C\n#pragma unroll\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n sum_C += utils::shfl_xor( sum_C, mask );\n sum_N += utils::shfl_xor( sum_N, mask );\n }\n\n if ( lane_id == 0 )\n {\n // NOTE this matches the check for zero in HYPRE\n double div = (fabs(sum_C * diag[a_row_id]) == 0.0) ? 1. : sum_C * diag[a_row_id];\n alfa = -sum_N / div;\n // alfa = -sum_N/(sum_C*diag[a_row_id]);\n }\n\n alfa = utils::shfl( alfa, 0 );\n\n int p_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n p_col_tmp = P_rows[a_row_id + lane_id];\n }\n\n int p_col_it = utils::shfl( p_col_tmp, 0 );\n int p_col_end = utils::shfl( p_col_tmp, 1 );\n\n map.store_keys_scale_values( p_col_end - p_col_it, &P_cols[p_col_it], alfa, &P_vals[p_col_it] );\n }\n}\n\n\n} // namespace multipass\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n/*\n * grab the diagonal of the matrix\n */\ntemplate \n__global__\nvoid getDiagonalKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, const IndexType numRows, ValueType *diagonal)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n for (int j = offset; j < offset + numj; j++)\n {\n int jcol = column_indices[j];\n\n if (tIdx == jcol)\n {\n diagonal[tIdx] = values[j];\n }\n }\n }\n}\n/*************************************************************************\n * Implementing Extended+i algorithm from \\S 4.5 of:\n * \"Distance-two interpolation for parallel algebraic multigrid\"\n * Reference [4] on wiki\n ************************************************************************/\ntemplate \nMultipass_Interpolator >::Multipass_Interpolator(AMG_Config &cfg, const std::string &cfg_scope) : Base(cfg, cfg_scope)\n{\n this->m_use_opt_kernels = cfg.getParameter(\"use_opt_kernels\", \"default\");\n}\n\ntemplate \nMultipass_Interpolator >::~Multipass_Interpolator()\n{}\n\nenum { WARP_SIZE = 32, GRID_SIZE = 1024, SMEM_SIZE = 128 };\n\n\nstruct is_less_than_zero\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x < 0;\n }\n};\n\nstruct is_strong_fine\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x == -3;\n }\n};\n\n__global__ void set_assigned_in_pass(int n, int pass, int* assigned, int* assigned_in_pass, int* assigned_offs)\n{\n int i = blockIdx.x*blockDim.x + threadIdx.x;\n if(i >= n) return;\n\n assigned_in_pass[i] = (assigned[i] == pass) ? 1 : 0;\n\n if(i == 0) \n {\n assigned_offs[0] = 0;\n }\n}\n\n__global__ void assigned_set_fill(int n, int* assigned_in_pass, int* assigned_offs, int* assigned_set)\n{\n int i = blockIdx.x*blockDim.x + threadIdx.x;\n if(i >= n) return;\n\n if(assigned_in_pass[i]) \n { \n assigned_set[assigned_offs[i]] = i; \n } \n}\n\ntemplate \ntemplate \nvoid Multipass_Interpolator>::compute_c_hat_opt_dispatch(\n const Matrix_d &A,\n const bool *s_con,\n const int *C_hat_start,\n int *C_hat_size,\n KeyType *C_hat,\n int *C_hat_pos,\n int *assigned,\n IntVector &assigned_set,\n int pass )\n{\n typedef typename IndPrecisionMap::Type Index_type;\n\n constexpr int cta_size = 128;\n constexpr int ngroups = cta_size / group_size;\n\n const int grid_size = assigned_set.size() / ngroups + 1;\n\n cudaDeviceProp deviceProps = getDeviceProperties();\n size_t max_shmem_size = deviceProps.sharedMemPerMultiprocessor;\n\n constexpr int shmem_size =\n (sizeof(KeyType)+sizeof(Index_type))*ngroups*hash_size + group_size; \n\n if(shmem_size > max_shmem_size) \n { \n FatalError(\"In compute_values_opt the requested hash size is larger than max.\\n\", \n AMGX_ERR_NOT_IMPLEMENTED); \n } \n\n cudaFuncSetAttribute(multipass::compute_c_hat_kernel_opt \n , \n cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size); \n\n multipass::compute_c_hat_kernel_opt\n \n <<>>(\n assigned_set.size(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n s_con,\n C_hat_start,\n C_hat_size,\n C_hat,\n C_hat_pos,\n assigned,\n assigned_set.raw(),\n pass);\n cudaCheckError();\n}\n\ntemplate \nvoid Multipass_Interpolator >::generateInterpolationMatrix_1x1(Matrix_d &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix_d &P)\n{\n nvtxRange nvr_gim(__func__);\n\n // Implementation based on paper \"On long range interpolation operators for aggressive coarsening\" by Ulrike Meier Yang, section 4.3\n const int blockSize = 256;\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // get raw pointers to data\n const IndexType *Aoffsets = A.row_offsets.raw();\n const IndexType *Acolumn_indices = A.col_indices.raw();\n const ValueType *Avalues = A.values.raw();\n const IndexType Anum_rows = (int) A.get_num_rows();\n int *cf_map_ptr = cf_map.raw();\n bool *s_con_ptr = s_con.raw();\n int size_two_ring;\n int size_one_ring;\n\n if (A.is_matrix_distributed())\n {\n int offset;\n A.getOffsetAndSizeForView(ALL, &offset, &size_two_ring);\n A.getOffsetAndSizeForView(FULL, &offset, &size_one_ring);\n }\n else\n {\n size_one_ring = A.get_num_rows();\n size_two_ring = A.get_num_rows();\n }\n\n int numBlocksOneRing = std::min( 4096, (int) (size_one_ring + blockSize - 1) / blockSize );\n int numBlocksTwoRing = std::min( 4096, (int) (size_two_ring + blockSize - 1) / blockSize );\n int numBlocks = std::min( 4096, (int) (A.get_num_rows() + blockSize - 1) / blockSize );\n // ----------------------------------------------------------\n // First fill out the assigned array and count # of passes\n // ----------------------------------------------------------\n IntVector assigned( size_one_ring, -1);\n IntVector C_hat_start( size_one_ring + 1, 0 );\n // assigned[i] = -1: unassigned point\n // assigned[i] = 0: coarse point\n // assigned[i] = 1: fine point directly connected to coarse point\n // assigned[i] = 2: fine point distance 2 away from strongly connected coarse point\n // assigned[i] = ...\n // Initialize assigned array by marking the coarse and fine point directly connected to coarse points\n const int cta_size = 256;\n const int nWarps = cta_size / 32;\n int grid_size = std::min( 4096, (int) (A.get_num_rows() + nWarps - 1) / nWarps);\n if(this->m_use_opt_kernels)\n {\n constexpr int group_size = 8;\n int ngroups_per_block = cta_size / group_size;\n int grid_size = A.get_num_rows() / ngroups_per_block + 1;\n initializeAssignedArray_opt <<< grid_size, cta_size>>>(cf_map.raw(), assigned.raw(), A.row_offsets.raw(), A.col_indices.raw(), s_con.raw(), C_hat_start.raw(), Anum_rows);\n cudaCheckError();\n }\n else\n {\n initializeAssignedArray <<< grid_size, cta_size>>>(cf_map.raw(), assigned.raw(), A.row_offsets.raw(), A.col_indices.raw(), s_con.raw(), C_hat_start.raw(), Anum_rows);\n }\n cudaCheckError();\n // Count the number of passes and fill assigned for pass > 1\n int pass = 2;\n int max_num_passes = 10;\n // Count the number of unassigned nodes by checking which nodes have assigned[i] = -1\n IndexType num_unassigned = thrust_wrapper::count_if(assigned.begin(), assigned.end(), is_less_than_zero());\n cudaCheckError();\n IndexType num_strong_fine = thrust_wrapper::count_if(cf_map.begin(), cf_map.end(), is_strong_fine());\n cudaCheckError();\n IndexType num_unassigned_max = num_unassigned - num_strong_fine;\n\n // Exchange 1 ring of assigned array\n if (A.is_matrix_distributed())\n {\n assigned.dirtybit = 1;\n A.manager->exchange_halo(assigned, assigned.tag);\n // get max across all processes\n A.manager->getComms()->all_reduce_max(num_unassigned, num_unassigned_max);\n }\n\n // Fill the assigned array for the other nodes, going one pass at the time\n while (num_unassigned_max && pass < max_num_passes)\n {\n if(this->m_use_opt_kernels)\n {\n constexpr int group_size = 8;\n int ngroups_per_block = cta_size / group_size;\n int nblocks = A.get_num_rows() / ngroups_per_block + 1;\n fillAssignedArray_opt <<< nblocks, cta_size>>>(assigned.raw(), A.row_offsets.raw(), A.col_indices.raw(), s_con.raw(), Anum_rows, pass);\n cudaCheckError();\n }\n else\n {\n fillAssignedArray <<< grid_size, cta_size>>>(assigned.raw(), A.row_offsets.raw(), A.col_indices.raw(), s_con.raw(), Anum_rows, pass);\n cudaCheckError();\n }\n num_unassigned = thrust_wrapper::count_if(assigned.begin(), assigned.begin() + A.get_num_rows(), is_less_than_zero());\n cudaCheckError();\n num_unassigned_max = num_unassigned - num_strong_fine;\n\n if (A.is_matrix_distributed())\n {\n assigned.dirtybit = 1;\n A.manager->exchange_halo(assigned, assigned.tag);\n // get max across all processes\n A.manager->getComms()->all_reduce_max(num_unassigned, num_unassigned_max);\n }\n\n // printf(\"after pass %d, %d unassigned\\n\",pass, num_unassigned_max);\n pass++;\n }\n\n cudaCheckError();\n // Check that we haven't reached maximum number of passes\n //if (pass == max_num_passes)\n //FatalError(\"Multipass intepolator reached maximum number of passes, exiting\\n\",AMGX_ERR_NOT_IMPLEMENTED);\n int num_passes = pass;\n // -------------------------------------------------------------------------------------\n // Create array cf_map_global, which assigns a global index to all owned coarse points\n // -------------------------------------------------------------------------------------\n I64Vector_d cf_map_global(size_one_ring);\n DistributedArranger *prep = NULL;\n I64Vector_d P_col_indices_global;\n int coarsePoints;\n\n if (A.is_matrix_distributed())\n {\n prep = new DistributedArranger;\n int num_owned_fine_pts = A.get_num_rows();\n int num_owned_coarse_pts = thrust_wrapper::count_if(cf_map.begin(), cf_map.begin() + num_owned_fine_pts, is_non_neg());\n cudaCheckError();\n int num_halo_coarse_pts = thrust_wrapper::count_if(cf_map.begin() + num_owned_fine_pts, cf_map.end(), is_non_neg());\n cudaCheckError();\n coarsePoints = num_owned_coarse_pts + num_halo_coarse_pts;\n // Partially initialize the distributed manager of matrix P, using num_owned_coarse_pts\n prep->initialize_manager(A, P, num_owned_coarse_pts);\n int my_rank = A.manager->global_id();\n const int cta_size = 128;\n const int grid_size = std::min( 4096, (num_owned_fine_pts + cta_size - 1) / cta_size);\n // Fill the cf_map_global array, which will assign for each coarse pt, the global index of that coarse pt\n createCfMapGlobal <<< grid_size, cta_size>>>(cf_map.raw(), cf_map_global.raw(), P.manager->part_offsets_h[my_rank], num_owned_fine_pts);\n cudaCheckError();\n // Exchange the cf_map_global so that we know the coarse global id of halo nodes in one ring\n cf_map_global.dirtybit = 1;\n A.manager->exchange_halo(cf_map_global, cf_map_global.tag);\n }\n else\n {\n coarsePoints = (int) thrust_wrapper::count_if(cf_map.begin(), cf_map.end(), is_non_neg());\n cudaCheckError();\n const int cta_size = 128;\n const int grid_size = std::min( 4096, (A.get_num_rows() + cta_size - 1) / cta_size);\n fill_cf_map_global <<< grid_size, cta_size>>>(cf_map.raw(), cf_map_global.raw(), A.get_num_rows());\n cudaCheckError();\n }\n\n // ----------------------------------------------------------\n // Create an upper bound for the length of each row P\n // ----------------------------------------------------------\n Hash_Workspace exp_wk(true, GRID_SIZE);\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int work_offset = GRID_SIZE * NUM_WARPS;\n cudaMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), cudaMemcpyHostToDevice );\n cudaCheckError();\n\n // Do one ring exchange of C_hat_start, which have been filled for nodes with assigned=1\n if (A.is_matrix_distributed())\n {\n C_hat_start.dirtybit = 1;\n A.manager->exchange_halo(C_hat_start, C_hat_start.tag);\n }\n\n for (int i = 2; i < num_passes; i++)\n {\n if(this->m_use_opt_kernels)\n {\n constexpr int group_size = 8;\n int ngroups_per_block = CTA_SIZE / group_size;\n int nblocks = A.get_num_rows() / ngroups_per_block + 1;\n multipass::estimate_c_hat_size_kernel_opt< CTA_SIZE, group_size> <<< nblocks, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n s_con.raw(),\n assigned.raw(),\n C_hat_start.raw(),\n i);\n cudaCheckError();\n }\n else\n {\n multipass::estimate_c_hat_size_kernel< CTA_SIZE, WARP_SIZE> <<< 2048, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n s_con.raw(),\n assigned.raw(),\n C_hat_start.raw(),\n i);\n cudaCheckError();\n }\n\n if (A.is_matrix_distributed())\n {\n // TODO: This could be improved by only doing the exchange for points in current pass\n C_hat_start.dirtybit = 1;\n A.manager->exchange_halo(C_hat_start, C_hat_start.tag);\n }\n }\n\n cudaCheckError();\n }\n\n // Do an exclusive scan on C_hat_start to compute compute upper bound on row offsets\n thrust_wrapper::exclusive_scan( C_hat_start.begin( ), C_hat_start.end( ), C_hat_start.begin( ) );\n cudaCheckError();\n // Now C_hat_start contains the offsets\n // Create a temporary manager that can be used to exchange halo information on values and/or column indices, where there's more than 1 value associated with each node\n Matrix C_hat_matrix(1, 1, 1, 0);\n\n // Create manager to exchange C_hat data\n if (A.is_matrix_distributed())\n {\n prep->createTempManager(C_hat_matrix, A, C_hat_start);\n }\n\n // ----------------------------------------------------------\n // Count the actual number of nonzeroes in each row of P\n // ----------------------------------------------------------\n // Allocate memory to store columns/values.\n int nVals = C_hat_start[C_hat_start.size() - 1];\n // C_hat stores the global indices of the coarse nodes that each fine point depends on\n I64Vector_d C_hat( nVals, -1 );\n // C_hat_pos is an array used by the hash set structure (check with Julien for more details)\n IntVector C_hat_pos( nVals );\n // C_hat_size stores the number of coarse points on which each fine point depends on\n IntVector C_hat_size( size_one_ring + 1, 0 );\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int avg_nz_per_row = A.get_num_nz() / A.get_num_rows();\n if(this->m_use_opt_kernels)\n {\n constexpr int group_size = 8;\n int ngroups_per_block = CTA_SIZE / group_size;\n int nblocks = A.get_num_rows() / ngroups_per_block + 1;\n\n // This will fill C_hat for nodes with assigned = 1 (for fine points strongly connected to coarse points)\n multipass::compute_c_hat_first_pass_kernel_opt< CTA_SIZE, WARP_SIZE, group_size, int64_t> <<< nblocks, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n s_con.raw(),\n cf_map_global.raw(),\n C_hat_start.raw(),\n C_hat_size.raw(),\n C_hat.raw(),\n assigned.raw());\n\n cudaCheckError();\n }\n else\n {\n // This will fill C_hat for nodes with assigned = 1 (for fine points strongly connected to coarse points)\n multipass::compute_c_hat_first_pass_kernel< CTA_SIZE, WARP_SIZE, int64_t> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n s_con.raw(),\n cf_map_global.raw(),\n C_hat_start.raw(),\n C_hat_size.raw(),\n C_hat.raw(),\n assigned.raw()\n );\n cudaCheckError();\n }\n\n // Here exchange C_hat_size and C_hat for one ring neighbors\n if (A.is_matrix_distributed())\n {\n //TODO: We could probably reduce the amount of data exchanged here, since we only need to exchange for recently updated nodes\n C_hat.dirtybit = 1;\n C_hat_matrix.manager->exchange_halo(C_hat, C_hat.tag);\n C_hat_size.dirtybit = 1;\n A.manager->exchange_halo(C_hat_size, C_hat_size.tag);\n }\n\n for (int i = 2; i < num_passes; i++)\n {\n if(false && this->m_use_opt_kernels)\n {\n I64Vector_d C_hat_(C_hat);\n IntVector C_hat_pos_(C_hat_pos);\n IntVector C_hat_size_(C_hat_size);\n\n IntVector assigned_in_pass(assigned.size());\n IntVector assigned_offs(assigned.size()+1);\n\n int nthreads = 128;\n int nblocks = assigned.size() / nthreads + 1;\n set_assigned_in_pass<<>>(assigned.size(), i, assigned.raw(), assigned_in_pass.raw(), assigned_offs.raw());\n cudaCheckError();\n\n thrust_wrapper::inclusive_scan(assigned_in_pass.begin(), assigned_in_pass.end(), assigned_offs.begin()+1);\n\n int nassigned_in_pass = assigned_offs[assigned_offs.size()-1] ;\n if(nassigned_in_pass > 0)\n {\n IntVector assigned_set(nassigned_in_pass);\n assigned_set_fill<<>>(assigned.size(), assigned_in_pass.raw(), assigned_offs.raw(), assigned_set.raw());\n cudaCheckError();\n\n int max_c_hat_size = thrust::reduce(C_hat_size_.begin(), C_hat_size_.end(), -1, thrust::maximum());\n\n int hash_size = pow(2, ceil(log2(max_c_hat_size)));\n\n switch(hash_size)\n {\n case 1:\n case 2:\n case 4:\n case 8:\n case 16:\n case 32:\n compute_c_hat_opt_dispatch<32, 32>(\n A, s_con.raw(), C_hat_start.raw(),\n C_hat_size_.raw(), C_hat_.raw(), C_hat_pos_.raw(), assigned.raw(),\n assigned_set, i);\n break;\n case 64:\n compute_c_hat_opt_dispatch<64, 32>(\n A, s_con.raw(), C_hat_start.raw(),\n C_hat_size_.raw(), C_hat_.raw(), C_hat_pos_.raw(), assigned.raw(),\n assigned_set, i);\n break;\n case 128:\n compute_c_hat_opt_dispatch<128, 32>(\n A, s_con.raw(), C_hat_start.raw(),\n C_hat_size_.raw(), C_hat_.raw(), C_hat_pos_.raw(), assigned.raw(),\n assigned_set, i);\n break;\n case 256:\n compute_c_hat_opt_dispatch<256, 32>(\n A, s_con.raw(), C_hat_start.raw(),\n C_hat_size_.raw(), C_hat_.raw(), C_hat_pos_.raw(), assigned.raw(),\n assigned_set, i);\n break;\n case 512:\n compute_c_hat_opt_dispatch<512, 32>(\n A, s_con.raw(), C_hat_start.raw(),\n C_hat_size_.raw(), C_hat_.raw(), C_hat_pos_.raw(), assigned.raw(),\n assigned_set, i);\n break;\n default:\n FatalError(\"Size of C_hat in a single row too large, solve with original solvers.\", AMGX_ERR_INTERNAL);\n }\n }\n\n cudaDeviceSynchronize();\n cudaCheckError();\n\n#if 0\n if (A.is_matrix_distributed())\n {\n //TODO: We could probably reduce the amount of data exchanged here, since we only need to exchange for recently updated nodes\n C_hat_.dirtybit = 1;\n C_hat_matrix.manager->exchange_halo(C_hat_, C_hat_.tag);\n C_hat_size_.dirtybit = 1;\n A.manager->exchange_halo(C_hat_size_, C_hat_size_.tag);\n }\n#endif\n }\n else\n {\n }\n\n // In case the hash table is not big enough, run multiple attemps\n int attempt = 0;\n\n for ( bool done = false ; !done && attempt < 10 ; ++attempt )\n {\n // Double the amount of GMEM (if needed).\n if ( attempt > 0 )\n {\n exp_wk.expand();\n }\n\n // Reset the status. TODO: Launch async copies.\n int status = 0;\n cudaMemcpy( exp_wk.get_status(), &status, sizeof(int), cudaMemcpyHostToDevice );\n cudaCheckError();\n // Compute the set C_hat.\n int work_offset = GRID_SIZE * NUM_WARPS;\n cudaMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), cudaMemcpyHostToDevice );\n cudaCheckError();\n // Run the computation.\n multipass::compute_c_hat_kernel< 8, CTA_SIZE, SMEM_SIZE, WARP_SIZE, int64_t> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n s_con.raw(),\n C_hat_start.raw(),\n C_hat_size.raw(),\n C_hat.raw(),\n C_hat_pos.raw(),\n assigned.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_work_queue(),\n exp_wk.get_status(),\n i);\n cudaCheckError();\n\n if (A.is_matrix_distributed())\n {\n //TODO: We could probably reduce the amount of data exchanged here, since we only need to exchange for recently updated nodes\n C_hat.dirtybit = 1;\n C_hat_matrix.manager->exchange_halo(C_hat, C_hat.tag);\n C_hat_size.dirtybit = 1;\n A.manager->exchange_halo(C_hat_size, C_hat_size.tag);\n }\n\n cudaCheckError();\n // Read the result from count_non_zeroes.\n cudaMemcpy( &status, exp_wk.get_status(), sizeof(int), cudaMemcpyDeviceToHost );\n cudaCheckError();\n done = status == 0;\n }\n }\n }\n // At this stage, for each fine node (even in 1 ring halo), we have the global indices\n // of the interpolatory coarse points\n // count the number of non-zeros in the interpolation matrix,\n // For coarse point, nonzeros = 1\n // for strong fine, nonzeros = 0\n // for other points, nonzeros[i] = C_hat_size[i]\n IntVector nonZeroOffsets(size_one_ring + 1);\n IntVector nonZerosPerRow(size_one_ring);\n nonZerosPerRowSizeKernel <<< numBlocks, blockSize>>>(Anum_rows, cf_map_ptr, C_hat_size.raw(),\n nonZerosPerRow.raw());\n cudaCheckError();\n\n // Do a one ring halo exchange to know the lengths of the neighbors row\n if (A.is_matrix_distributed())\n {\n nonZerosPerRow.dirtybit = 1;\n A.manager->exchange_halo(nonZerosPerRow, nonZerosPerRow.tag);\n }\n\n // get total number of non-zeros with reduction\n int nonZeros = thrust_wrapper::reduce(nonZerosPerRow.begin(), nonZerosPerRow.end());\n cudaCheckError();\n // get the offsets with an exclusive scan\n thrust_wrapper::exclusive_scan(nonZerosPerRow.begin(), nonZerosPerRow.end(), nonZeroOffsets.begin());\n cudaCheckError();\n nonZeroOffsets[nonZeroOffsets.size() - 1] = nonZeros;\n // resize P\n P.resize(0, 0, 0, 1);\n P.addProps(CSR);\n // TODO: is the number of columns used anywhere?\n P.resize(size_one_ring, coarsePoints, nonZeros, 1);\n cudaCheckError();\n P_col_indices_global.resize(P.col_indices.size());\n Matrix P_temp(1, 1, 1, 0);\n\n if (A.is_matrix_distributed())\n {\n prep->createTempManager(P_temp, A, nonZeroOffsets);\n }\n\n IndexType *Poffsets = P.row_offsets.raw();\n IndexType *Pcolumn_indices = P.col_indices.raw();\n ValueType *Pvalues = P.values.raw();\n // copy nonzero offsets to the P matrix\n amgx::thrust::copy(nonZeroOffsets.begin(), nonZeroOffsets.end(), P.row_offsets.begin());\n cudaCheckError();\n // grab the diagonal terms\n VVector diag(size_one_ring);\n ValueType *diag_ptr = diag.raw();\n {\n find_diag_kernel_indexed_dia <<< numBlocksOneRing, blockSize>>>(\n size_one_ring,\n A.diag.raw(),\n A.values.raw(),\n diag.raw());\n }\n cudaCheckError();\n\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int work_offset = GRID_SIZE * NUM_WARPS;\n typedef typename MatPrecisionMap::Type Value_type;\n\n // Fill P.col_indices_global, P.values\n if(this->m_use_opt_kernels)\n {\n int grid_size = A.get_num_rows() / NUM_WARPS + 1;\n\n multipass::compute_interp_weight_first_pass_kernel_opt\n \n <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n cf_map_global.raw(),\n s_con.raw(),\n diag.raw(),\n P.row_offsets.raw(),\n P_col_indices_global.raw(),\n P.values.raw(),\n assigned.raw());\n cudaCheckError();\n }\n else\n {\n cudaMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), cudaMemcpyHostToDevice );\n cudaCheckError();\n\n // Compute the set C_hat.\n multipass::compute_interp_weight_first_pass_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n cf_map_global.raw(),\n s_con.raw(),\n diag.raw(),\n P.row_offsets.raw(),\n P_col_indices_global.raw(),\n P.values.raw(),\n assigned.raw());\n cudaCheckError();\n }\n\n if (A.is_matrix_distributed())\n {\n //TODO: We could probably reduce the amount of data exchanged here, since we only need to exchange for recently updated nodes\n P_col_indices_global.dirtybit = 1;\n P_temp.manager->exchange_halo(P_col_indices_global, P_col_indices_global.tag);\n P.values.dirtybit = 1;\n P_temp.manager->exchange_halo(P.values, P.values.tag);\n }\n\n for (int i = 2; i < num_passes; i++)\n {\n if(this->m_use_opt_kernels)\n {\n IntVector assigned_in_pass(assigned.size());\n IntVector assigned_offs(assigned.size() + 1);\n\n int nthreads = 128;\n int nblocks = assigned.size() / nthreads + 1;\n set_assigned_in_pass<<>>(assigned.size(), i, assigned.raw(), assigned_in_pass.raw(), assigned_offs.raw());\n cudaCheckError();\n thrust_wrapper::inclusive_scan(assigned_in_pass.begin(), assigned_in_pass.end(), assigned_offs.begin()+1);\n\n IntVector assigned_set(assigned_offs[assigned_offs.size()-1]);\n assigned_set_fill<<>>(assigned.size(), assigned_in_pass.raw(), assigned_offs.raw(), assigned_set.raw());\n cudaCheckError();\n\n Hash_Workspace exp_wk2(true, 4096);\n multipass::compute_interp_weight_kernel_opt <<< 4096, CTA_SIZE>>>(\n assigned_set.size(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat.raw(),\n C_hat_pos.raw(),\n C_hat_start.raw(),\n C_hat_size.raw(),\n diag.raw(),\n P.row_offsets.raw(),\n P_col_indices_global.raw(),\n P.values.raw(),\n exp_wk2.get_gmem_size(),\n exp_wk2.get_keys(),\n exp_wk2.get_vals(),\n exp_wk2.get_work_queue(),\n assigned.raw(),\n assigned_set.raw(),\n i);\n cudaCheckError();\n }\n else\n {\n cudaMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), cudaMemcpyHostToDevice );\n cudaCheckError();\n multipass::compute_interp_weight_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat.raw(),\n C_hat_pos.raw(),\n C_hat_start.raw(),\n C_hat_size.raw(),\n diag.raw(),\n P.row_offsets.raw(),\n P_col_indices_global.raw(),\n P.values.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_vals(),\n exp_wk.get_work_queue(),\n assigned.raw(),\n i);\n cudaCheckError();\n }\n\n if (A.is_matrix_distributed())\n {\n //TODO: This is exchanging too much data\n P_col_indices_global.dirtybit = 1;\n P_temp.manager->exchange_halo(P_col_indices_global, P_col_indices_global.tag);\n P.values.dirtybit = 1;\n P_temp.manager->exchange_halo(P.values, P.values.tag);\n }\n\n cudaCheckError();\n }\n\n if (!A.is_matrix_distributed())\n {\n const int cta_size = 128;\n const int grid_size = std::min( 4096, ( (int) P.col_indices.size() + cta_size - 1) / cta_size);\n // Simply cast from int64 to int\n fill_P_col_indices <<< grid_size, cta_size>>>(P_col_indices_global.raw(), P.col_indices.raw(), (int) P.col_indices.size());\n cudaCheckError();\n }\n else\n {\n // Resize P to only include the owned rows and columns\n int P_num_nnz = P.row_offsets[P.get_num_rows()];\n P.resize(A.get_num_rows(), P.manager->halo_offsets[0], P_num_nnz, 1);\n // Using the global col indices, identify neighbors for P.manager, fill B2L_maps, etc...\n prep->initialize_manager_from_global_col_indices(P, P_col_indices_global);\n // Create the list of interior and boundary nodes (this is necessary if doing masked SpMV operations)\n prep->createRowsLists(P, true);\n }\n\n // Delete the temporary distributed arranger\n if (prep != NULL) { delete prep; }\n}\n\ntemplate< class T_Config>\nvoid Multipass_InterpolatorBase::generateInterpolationMatrix(Matrix &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix &P)\n{\n P.set_initialized(0);\n ViewType oldView = A.currentView();\n A.setView(OWNED);\n\n if (A.get_block_size() == 1)\n {\n generateInterpolationMatrix_1x1(A, cf_map, s_con, scratch, P);\n }\n else\n {\n FatalError(\"Unsupported dimensions for multipass interpolator\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n //P.computeDiagonal();\n P.set_initialized(1);\n#if 0\n P.printToFile(\"P.out\", \"\", -1, -1);\n#endif\n}\n\n#define AMGX_CASE_LINE(CASE) template class Multipass_InterpolatorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Multipass_Interpolator::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\nusing std::set;\nusing std::vector;\n\ntemplate < typename IndexType, typename ValueType >\n__global__\nvoid check_P_values(IndexType *offsets, ValueType *val, int size, int flag)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < size; tidx += blockDim.x * gridDim.x)\n {\n // Simply cast\n for (int i = offsets[tidx]; i < offsets[tidx + 1]; i++)\n {\n if (isnan(val[i])) { printf(\"%d, row = %d, i=%d\\n\", flag, tidx, i); }\n }\n }\n}\n\ntemplate \n__global__\nvoid fill_cf_map_global(const IndexType *cf_map, KeyType *cf_map_global, IndexType num_owned_fine_pts)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_owned_fine_pts ; tidx += blockDim.x * gridDim.x)\n {\n // Simply cast\n cf_map_global[tidx] = (KeyType) cf_map[tidx];\n }\n}\n\ntemplate \n__global__\nvoid fill_P_col_indices(const KeyType *global_cols, IndexType *cols, IndexType num_owned_fine_pts)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_owned_fine_pts ; tidx += blockDim.x * gridDim.x)\n {\n // Simply cast\n cols[tidx] = (IndexType) global_cols[tidx];\n }\n}\n\n\n\ntemplate \n__global__\nvoid createCfMapGlobal(const IndexType *cf_map, KeyType *cf_map_global, const KeyType my_part_offset, const IndexType num_owned_fine_pts)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_owned_fine_pts ; tidx += blockDim.x * gridDim.x)\n {\n // Only renumber the interior points\n if (cf_map[tidx] >= 0 && tidx < num_owned_fine_pts)\n {\n // Simply shift\n cf_map_global[tidx] = (KeyType) cf_map[tidx] /* my_local_id */ + my_part_offset;\n }\n else\n {\n cf_map_global[tidx] = -1;\n }\n }\n}\n\n\ntemplate\n__global__\nvoid initializeAssignedArray(const IndexType *cf_map, IndexType *assigned, const IndexType *A_rows, const IndexType *A_cols, const bool *s_con, IndexType *C_hat_offsets, int A_num_rows)\n{\n const int nWarps = cta_size / 32;\n // The coordinates of the thread inside the CTA/warp.\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n // Loop over rows of A.\n for ( int aRowId = blockIdx.x * nWarps + warpId ; aRowId < A_num_rows ; aRowId += gridDim.x * nWarps )\n {\n int coarse_fine_id = cf_map[aRowId];\n\n if ( coarse_fine_id >= 0)\n {\n assigned[aRowId] = 0;\n }\n else if (coarse_fine_id == FINE)\n {\n // Load A row IDs.\n int aColIt = A_rows[aRowId ];\n int aColEnd = A_rows[aRowId + 1];\n int coarse_count = 0;\n\n // Iterate over the columns of A.\n for ( aColIt += laneId ; utils::any( aColIt < aColEnd ) ; aColIt += 32 )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int aColId = aColIt < aColEnd ? A_cols[aColIt] : -1;\n // Is it an off-diagonal element.\n bool is_off_diagonal = aColIt < aColEnd && aColId != aRowId;\n // Is it fine.\n bool is_coarse = is_off_diagonal && cf_map[aColId] >= 0;\n // Is it strongly connected ?\n bool is_coarse_strongly_connected = aColIt < aColEnd && s_con[aColIt] && is_coarse;\n // Push coarse and strongly connected nodes in the set.\n coarse_count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n }\n\n // If fine point connected to coarse point, assigned = 1,\n if (coarse_count)\n {\n assigned[aRowId] = 1;\n C_hat_offsets[aRowId] = coarse_count;\n }\n }\n else // Strong fine\n {\n assigned[aRowId] = -1;\n }\n }\n}\n\ntemplate \n__global__\nvoid initializeAssignedArray_opt(const IndexType *cf_map, IndexType *assigned, const IndexType *A_rows, const IndexType *A_cols, const bool *s_con, IndexType *C_hat_offsets, int A_num_rows)\n{\n const int ngroups = cta_size / group_size;\n\n // The coordinates of the thread inside the CTA/group.\n const int group_id = threadIdx.x / group_size;\n const int lane_id = threadIdx.x % group_size;\n\n int a_row_id = blockIdx.x * ngroups + group_id;\n\n bool active = true;\n if(a_row_id > A_num_rows)\n {\n active = false;\n }\n\n int coarse_fine_id = UNASSIGNED;\n\n if(active)\n {\n coarse_fine_id = cf_map[a_row_id];\n\n if(coarse_fine_id >= 0)\n {\n if(lane_id == 0)\n {\n assigned[a_row_id] = 0;\n }\n\n active = false;\n }\n else if(coarse_fine_id != FINE) // Strong fine\n {\n if(lane_id == 0)\n {\n assigned[a_row_id] = -1;\n }\n\n active = false;\n }\n }\n\n active = (active && coarse_fine_id == FINE);\n\n constexpr int warp_size = 32;\n unsigned int base_mask = 0xFFFFFFFF >> (warp_size - group_size);\n int ngroups_in_warp = warp_size / group_size;\n int group_id_in_warp = group_id % ngroups_in_warp;\n unsigned int group_mask = base_mask << (group_id_in_warp*group_size);\n\n // Load A row IDs.\n int aColIt = A_rows[a_row_id ];\n int aColEnd = A_rows[a_row_id + 1];\n\n // Iterate over the columns of A.\n int coarse_count = 0;\n for ( aColIt += lane_id ; utils::any( active && aColIt < aColEnd ) ; aColIt += group_size )\n {\n bool is_coarse_strongly_connected = false;\n\n if(active)\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int aColId = aColIt < aColEnd ? A_cols[aColIt] : -1;\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = aColIt < aColEnd && aColId != a_row_id;\n\n // Is it fine.\n bool is_coarse = is_off_diagonal && cf_map[aColId] >= 0;\n\n // Is it strongly connected ?\n is_coarse_strongly_connected = s_con[aColIt] && is_coarse;\n }\n\n // Count any of the threads in the group that had strong connections\n int coarse_exist_flags = utils::ballot(is_coarse_strongly_connected);\n if(lane_id == 0)\n {\n coarse_count += __popc(coarse_exist_flags & group_mask);\n }\n }\n\n // If fine point connected to coarse point, assigned = 1,\n if (active && coarse_count > 0)\n {\n if(lane_id == 0)\n {\n assigned[a_row_id] = 1;\n C_hat_offsets[a_row_id] = coarse_count;\n }\n }\n}\n\ntemplate\n__global__\nvoid fillAssignedArray(IndexType *assigned, const IndexType *A_rows, const IndexType *A_cols, const bool *s_con, int A_num_rows, int pass)\n{\n const int nWarps = cta_size / 32;\n // The coordinates of the thread inside the CTA/warp.\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n // Loop over rows of A.\n for ( int aRowId = blockIdx.x * nWarps + warpId ; aRowId < A_num_rows ; aRowId += gridDim.x * nWarps )\n {\n if ( assigned[aRowId] == -1)\n {\n // Load A row IDs.\n int aColIt = A_rows[aRowId ];\n int aColEnd = A_rows[aRowId + 1];\n int assigned_count = 0;\n\n // Iterate over the columns of A.\n for ( aColIt += laneId ; utils::any( aColIt < aColEnd ) ; aColIt += 32 )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int aColId = aColIt < aColEnd ? A_cols[aColIt] : -1;\n // Is it an off-diagonal element.\n bool is_off_diagonal = aColIt < aColEnd && aColId != aRowId;\n // Is it assigned\n bool is_assigned = is_off_diagonal && assigned[aColId] == pass - 1;\n // Is it strongly connected ?\n bool is_assigned_strongly_connected = aColIt < aColEnd && s_con[aColIt] && is_assigned;\n // Push coarse and strongly connected nodes in the set.\n assigned_count += __popc( utils::ballot( is_assigned_strongly_connected ) );\n\n if (assigned_count) { break; }\n }\n\n // If one of the strong neighbor has been assigned, set assigned to current pass number\n if (assigned_count)\n {\n assigned[aRowId] = pass;\n }\n }\n }\n}\n\ntemplate < typename IndexType, int cta_size, int group_size >\n__global__\nvoid fillAssignedArray_opt(IndexType *assigned, const IndexType *A_rows, const IndexType *A_cols, const bool *s_con, int A_num_rows, int pass)\n{\n // The coordinates of the thread inside the CTA/groups\n const int group_id = threadIdx.x / group_size;\n const int lane_id = threadIdx.x % group_size;\n const int ngroups_per_block = cta_size / group_size;\n\n int a_row_id = blockIdx.x * ngroups_per_block + group_id;\n\n bool active = true;\n if (a_row_id >= A_num_rows || assigned[a_row_id] != -1)\n {\n active = false;\n }\n\n // Mask for the subgroup\n constexpr int warp_size = 32;\n unsigned int base_mask = 0xFFFFFFFF >> (warp_size - group_size);\n int ngroups_in_warp = warp_size / group_size;\n int group_id_in_warp = group_id % ngroups_in_warp;\n unsigned int group_mask = base_mask << (group_id_in_warp*group_size);\n\n // Load A row IDs.\n int a_col_begin = 0;\n int a_col_end = 0;\n\n if(active)\n {\n a_col_begin = A_rows[a_row_id ];\n a_col_end = A_rows[a_row_id + 1];\n }\n\n // Iterate over the columns of A.\n for (int a_col_it = a_col_begin + lane_id ; utils::any(active && a_col_it < a_col_end); a_col_it += group_size)\n {\n bool is_assigned_strongly_connected = false;\n\n if(active)\n {\n // Columns of A maps to rows of B. Each thread of the group loads its A-col/B-row ID.\n int a_col_id = a_col_it < a_col_end ? A_cols[a_col_it] : -1;\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n\n // Is it assigned\n bool is_assigned = is_off_diagonal && assigned[a_col_id] == pass - 1;\n\n // Is it strongly connected ?\n is_assigned_strongly_connected = s_con[a_col_it] && is_assigned;\n }\n\n int group_strongly_connected = utils::ballot(is_assigned_strongly_connected);\n\n // Push coarse and strongly connected nodes in the set.\n if(active && (group_strongly_connected & group_mask))\n {\n if(lane_id == 0)\n {\n assigned[a_row_id] = pass;\n }\n\n active = false;\n }\n }\n}\n\n\n/*\n * remove all duplicates from a vector\n */\ntemplate \nvoid removeDuplicates(vector &vec)\n{\n std::sort(vec.begin(), vec.end());\n vec.erase(std::unique(vec.begin(), vec.end()), vec.end());\n}\n\n/*\n * find the value at a particular column of a row if it exists\n */\ntemplate \ntypename VVector::value_type findValue(const IVector &columns,\n const VVector &values,\n const int row_begin, const int row_end, const int column_needed)\n{\n typename VVector::value_type v = 0;\n\n // NAIVE\n for (int k = row_begin; k < row_end; k++)\n {\n int kcol = columns[k];\n\n if (kcol == column_needed)\n {\n v = values[k];\n }\n }\n\n return v;\n}\n\n\n\n/*************************************************************************\n* create the interpolation matrix P\n************************************************************************/\ntemplate \nvoid Multipass_Interpolator >::generateInterpolationMatrix_1x1(Matrix_h &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix_h &P)\n{\n FatalError(\"multipass interpolation not implemented on host\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n} // end multipass interpolator\n\n\n///////////////////////////////////////////////////////////////////////////////\n\n// DEVICE CODE\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#include \n#include // Included inside the namespace to solve name colisions.\n\n__device__ __forceinline__ int get_work( int *queue, int warp_id )\n{\n int offset = -1;\n\n if ( utils::lane_id() == 0 )\n {\n offset = atomicAdd( queue, 1 );\n }\n\n return utils::shfl( offset, 0 );\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nnamespace multipass\n{\n\n#include \n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nestimate_c_hat_size_kernel( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const bool *s_con,\n const int *assigned,\n int *C_hat_offsets,\n int pass )\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n // Number of items per grid.\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n volatile int a_row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += NUM_WARPS_PER_GRID )\n {\n // Skip coarse rows and fine rows touching a coarse point\n if (assigned[a_row_id] != pass)\n {\n continue;\n }\n\n // Load A row IDs.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // The number of elements.\n int my_count = 0;\n\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // assigned == pass-1\n bool is_assigned = is_off_diagonal && (assigned[a_col_id] == pass - 1);\n // Is it a coarse, strongly-connected column and already assigned?\n bool is_assigned_strongly_connected = is_strongly_connected && is_assigned;\n // If so, add their C_hat_offset value to my_count\n my_count += (is_assigned_strongly_connected ? C_hat_offsets[a_col_id] : 0);\n }\n\n // Do reduction\n#pragma unroll\n\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n my_count += utils::shfl_xor( my_count, mask );\n }\n\n // Write result -- RACE CONDITION!\n if (lane_id == 0)\n {\n C_hat_offsets[a_row_id] = my_count;\n }\n }\n}\n\ntemplate< int CTA_SIZE, int group_size >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nestimate_c_hat_size_kernel_opt( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const bool *s_con,\n const int *assigned,\n int *C_hat_offsets,\n int pass )\n{\n const int group_id = threadIdx.x / group_size;\n const int lane_id = threadIdx.x % group_size;\n\n int ngroups = CTA_SIZE / group_size;\n int a_row_id = blockIdx.x * ngroups + group_id;\n\n // Skip coarse rows and fine rows touching a coarse point\n bool active = true;\n if (a_row_id >= A_num_rows || assigned[a_row_id] != pass)\n {\n active = false;\n }\n\n int my_count = 0;\n\n // Load A row IDs.\n if(active)\n {\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n\n // Iterate over the columns of A to build C_hat.\n for (int a_col_it = a_col_begin + lane_id; a_col_it < a_col_end; a_col_it += group_size )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = A_cols[a_col_it];\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_id != a_row_id;\n\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n\n // assigned == pass-1\n bool is_assigned = is_off_diagonal && (assigned[a_col_id] == pass - 1);\n\n // Is it a coarse, strongly-connected column and already assigned?\n bool is_assigned_strongly_connected = is_strongly_connected && is_assigned;\n\n // If so, add their C_hat_offset value to my_count\n my_count += (is_assigned_strongly_connected ? C_hat_offsets[a_col_id] : 0);\n }\n }\n\n#pragma unroll\n for(int i = group_size / 2; i > 0; i /= 2)\n {\n my_count += utils::shfl_down(my_count, i);\n }\n\n if (active && lane_id == 0)\n {\n C_hat_offsets[a_row_id] = my_count;\n }\n}\n\n template< int CTA_SIZE, int WARP_SIZE, typename KeyType >\n__global__ __launch_bounds__( CTA_SIZE )\n void\ncompute_c_hat_first_pass_kernel( int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const bool *__restrict s_con,\n const KeyType *__restrict cf_map_global,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_size,\n KeyType *__restrict C_hat,\n int *__restrict assigned)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / WARP_SIZE;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += NUM_WARPS_PER_GRID)\n {\n // Skip if not in current pass\n if (assigned[a_row_id] != 1)\n {\n continue;\n }\n\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_begin = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n int count = 0;\n int c_col_it = C_hat_start[a_row_id];\n\n // _iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = is_active && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // assigned == pass -1\n bool is_assigned = is_off_diagonal && (assigned[a_col_id] == 0);\n // Is it an assigned==pass -1 and strongly-connected column.\n bool is_assigned_strongly_connected = is_strongly_connected && is_assigned;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_assigned_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_assigned_strongly_connected )\n {\n //if (cf_map[a_col_id] < 0)\n // printf(\"error in multipass 1\\n\");\n C_hat[c_col_it + count + dest] = cf_map_global[a_col_id];\n }\n\n count += __popc( vote );\n }\n\n if ( lane_id == 0 )\n {\n C_hat_size[a_row_id] = count;\n }\n }\n}\n\n template< int CTA_SIZE, int WARP_SIZE, int group_size, typename KeyType>\n__global__ __launch_bounds__( CTA_SIZE )\n void\ncompute_c_hat_first_pass_kernel_opt( int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const bool *__restrict s_con,\n const KeyType *__restrict cf_map_global,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_size,\n KeyType *__restrict C_hat,\n int *__restrict assigned)\n{\n const int ngroups = CTA_SIZE / group_size;\n\n // The coordinates of the thread inside the CTA/warp.\n const int group_id = threadIdx.x / group_size;\n const int lane_id = threadIdx.x % group_size;\n\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * ngroups + group_id;\n\n bool active = true;\n\n // Skip if not in current pass\n if (a_row_id >= A_num_rows || assigned[a_row_id] != 1)\n {\n active = false;\n }\n\n int a_col_begin = 0;\n int a_col_end = 0;\n int c_col_it = 0;\n\n if(active)\n {\n a_col_begin = A_rows[a_row_id];\n a_col_end = A_rows[a_row_id + 1];\n c_col_it = C_hat_start[a_row_id];\n }\n\n unsigned int base_mask = 0xFFFFFFFF >> (WARP_SIZE - group_size);\n int ngroups_in_warp = WARP_SIZE / group_size;\n int group_id_in_warp = group_id % ngroups_in_warp;\n unsigned int group_mask = base_mask << (group_id_in_warp*group_size);\n\n // _iterate over the columns of A to build C_hat.\n int count = 0;\n for ( int a_col_it = a_col_begin + lane_id ; utils::any(active && a_col_it < a_col_end) ; a_col_it += group_size )\n {\n bool is_assigned_strongly_connected = false;\n\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if(active)\n {\n if (a_col_it < a_col_end)\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = (a_col_it < a_col_end) && (a_col_id != a_row_id);\n\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n\n // assigned == pass -1\n bool is_assigned = is_off_diagonal && (assigned[a_col_id] == 0);\n\n // Is it an assigned==pass -1 and strongly-connected column.\n is_assigned_strongly_connected = is_strongly_connected && is_assigned;\n }\n\n // We collect fine and strongly-collected columns.\n int all_vote = utils::ballot( is_assigned_strongly_connected );\n int group_vote = all_vote & group_mask;\n int dest = __popc( group_vote & utils::lane_mask_lt() );\n\n if ( active && is_assigned_strongly_connected )\n {\n C_hat[c_col_it + count + dest] = cf_map_global[a_col_id];\n }\n\n count += __popc( group_vote );\n }\n\n if ( active && lane_id == 0 )\n {\n C_hat_size[a_row_id] = count;\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, typename KeyType >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_c_hat_kernel( int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const bool *__restrict s_con,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_size,\n KeyType *__restrict C_hat,\n int *__restrict C_hat_pos,\n int *__restrict assigned,\n int gmem_size,\n KeyType *g_keys,\n int *wk_work_queue,\n int *wk_status,\n int pass )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // Shared memory to vote.\n __shared__ int s_b_row_ids[CTA_SIZE];\n // The hash keys stored in shared memory.\n __shared__ KeyType s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Skip if not in current pass\n if (assigned[a_row_id] != pass)\n {\n continue;\n }\n\n // Clear the set.\n set.clear();\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_begin = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // _iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = is_active && a_col_id != a_row_id;\n // Check if assigned == pass -1\n bool is_assigned = is_off_diagonal && (assigned[a_col_id] == pass - 1);\n // Check if we have a strong connection to the point\n bool is_strongly_connected = is_active && s_con[a_col_it];\n // Check if strong connection to assigned == pass-1\n bool is_assigned_strongly_connected = is_assigned && is_strongly_connected;\n // We collect columns for which assigned == pass - 1\n int vote = utils::ballot( is_assigned_strongly_connected);\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_assigned_strongly_connected)\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n __syncwarp();\n\n int num_rows = __popc( vote );\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int uniform_b_row_id = -1;\n\n if ( is_active_k )\n {\n uniform_b_row_id = s_b_row_ids[warp_id * WARP_SIZE + local_k];\n }\n\n // Load the range of the row of B.\n int b_col_it = 0, b_col_end = 0;\n\n if ( is_active_k )\n {\n // this loops over elements of C_hat for row uniform_b_row_id\n b_col_it = C_hat_start[uniform_b_row_id];\n b_col_end = b_col_it + C_hat_size[uniform_b_row_id];\n }\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n // The ID of the interpolatory point\n KeyType b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = C_hat[b_col_it];\n }\n\n // Is it valid\n is_off_diagonal = b_col_it < b_col_end; // && b_col_id != local_k;\n // Push node to set\n KeyType b_item = -1;\n\n if ( is_off_diagonal )\n {\n b_item = b_col_id;\n }\n\n // Insert the interpolatory point to the set\n set.insert( b_item, wk_status );\n }\n }\n }\n\n int c_col_it = C_hat_start[a_row_id];\n int count = set.store_with_positions( &C_hat[c_col_it], &C_hat_pos[c_col_it] );\n\n if ( lane_id == 0 )\n {\n C_hat_size[a_row_id] = count;\n }\n }\n}\n\ntemplate \n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_c_hat_kernel_opt( int assigned_set_size,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const bool *__restrict s_con,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_size,\n KeyType *__restrict C_hat,\n int *__restrict C_hat_pos,\n int *__restrict assigned,\n int *__restrict assigned_set,\n int pass )\n{\n constexpr KeyType SLOT_VACANT = -1;\n\n // Group indices\n constexpr int ngroups = CTA_SIZE / GROUP_SIZE;\n const int group_id = threadIdx.x / GROUP_SIZE;\n const int lane_id = threadIdx.x % GROUP_SIZE;\n\n const int gid = blockIdx.x * ngroups + group_id;\n\n // Dynamic sized shared memory\n extern __shared__ int s[];\n\n // Block-level hash container storage\n KeyType* key_s = (KeyType*)s;\n\n //KeyType* data_s = (KeyType*)&key_s[ngroups*HASH_SIZE];\n int* col_ind_s = (int*)&key_s[ngroups*HASH_SIZE];\n\n // Group-level hash containers\n KeyType* key_group_s = &key_s[group_id*HASH_SIZE];\n //KeyType* data_group_s = &data_s[group_id*HASH_SIZE];\n\n // Initialise the keys and values.\n#pragma unroll\n for(int i = threadIdx.x; i < ngroups*HASH_SIZE; i += CTA_SIZE)\n {\n key_s[i] = SLOT_VACANT; // Inserted keys will be in range [0,N]\n //data_s[i] = 0; // We will sum into values\n }\n\n if(lane_id == 0)\n {\n col_ind_s[group_id] = 0;\n }\n\n __syncthreads();\n\n int a_row_id = -1;\n\n if(gid < assigned_set_size)\n {\n // One row of A per group\n a_row_id = assigned_set[gid];\n\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = A_rows[a_row_id] + lane_id ; a_col_it < A_rows[a_row_id + 1]; a_col_it += GROUP_SIZE )\n {\n int a_col_id = A_cols[a_col_it];\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_id != a_row_id;\n\n // Check if assigned == pass -1\n bool is_assigned = is_off_diagonal && (assigned[a_col_id] == pass - 1);\n\n // Check if we have a strong connection to the point\n bool is_strongly_connected = s_con[a_col_it];\n\n // Check if strong connection to assigned == pass-1\n bool is_assigned_strongly_connected = is_assigned && is_strongly_connected;\n\n if(!is_assigned_strongly_connected) continue;\n\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int b_col_beg = C_hat_start[a_col_id];\n\n // Iterate over the range of columns of B.\n for (int b_col_it = b_col_beg; b_col_it < b_col_beg + C_hat_size[a_col_id]; ++b_col_it)\n {\n KeyType key = C_hat[b_col_it];\n\n int hash = key % HASH_SIZE;\n\n while(true)\n {\n // XXX You can just avoid this and directly atomicCAS, which is faster?\n int curr_key = key_group_s[hash];\n if(curr_key == key) break;\n \n if(curr_key == SLOT_VACANT)\n {\n KeyType old_key = (KeyType)atomicCAS((unsigned long long*)&key_group_s[hash], (unsigned long long)SLOT_VACANT, (unsigned long long)key);\n if(old_key == SLOT_VACANT || old_key == key)\n {\n break;\n }\n }\n\n // We did not secure a slot, so linear probe to next slot\n hash = (hash + 1) % HASH_SIZE;\n }\n }\n }\n }\n\n __syncthreads();\n\n // XXX Could change to poll by ballot style compress\n if(gid < assigned_set_size)\n {\n // Store the results.\n int c_col_it = C_hat_start[a_row_id];\n\n#pragma unroll\n for(int i = lane_id; i < HASH_SIZE; i += GROUP_SIZE)\n {\n KeyType key = key_group_s[i];\n if(key != SLOT_VACANT)\n {\n int ind = atomicAdd(&col_ind_s[group_id], 1);\n C_hat[c_col_it + ind] = key;\n //C_hat_pos[c_col_it + ind] = i;\n }\n }\n }\n\n __syncthreads();\n\n if(gid < assigned_set_size)\n {\n if(lane_id == 0)\n {\n C_hat_size[a_row_id] = col_ind_s[group_id];\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, typename KeyType >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_interp_weight_first_pass_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const KeyType *__restrict cf_map_global,\n const bool *__restrict s_con,\n const Value_type *__restrict diag,\n const int *__restrict P_rows,\n KeyType *P_cols,\n Value_type *P_vals,\n const int *__restrict assigned)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / 32;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += NUM_WARPS_PER_GRID)\n {\n // Only do work if assigned[row_id] == pass\n KeyType coarse_fine_id = cf_map_global[a_row_id];\n\n // If coarse rows\n if ( assigned[a_row_id] == 0 )\n {\n if ( lane_id == 0 )\n {\n int p_row_it = P_rows[a_row_id];\n P_cols[p_row_it] = coarse_fine_id;\n P_vals[p_row_it] = Value_type( 1 );\n }\n }\n // If strong fine\n else if ( assigned[a_row_id] == -1 )\n {\n // Do nothing\n }\n // strongly connected to a coarse point\n else if (assigned[a_row_id] == 1)\n {\n // Load A row IDs.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n int p_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n p_col_tmp = P_rows[a_row_id + lane_id];\n }\n\n int p_col_it = utils::shfl( p_col_tmp, 0 );\n int p_col_end = utils::shfl( p_col_tmp, 1 );\n\n // Weak value.\n Value_type sum_N(0), sum_C(0), alfa(0);\n int count = 0;\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id(-1);\n Value_type a_value(0);\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n a_value = A_vals[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it a strongly-connected column.\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it a weakly connected node.\n //bool is_weakly_connected = is_off_diagonal && !is_strongly_connected;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map_global[a_col_id] == STRONG_FINE;\n // Is it a fine and strongly-connected column.\n bool is_assigned_strongly_connected = is_strongly_connected && assigned[a_col_id] == 0;\n\n // Update the weak value.\n if ( is_off_diagonal && !is_strong_fine)\n {\n sum_N += a_value;\n }\n\n if ( is_assigned_strongly_connected )\n {\n sum_C += a_value;\n }\n\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_assigned_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_assigned_strongly_connected )\n {\n P_cols[p_col_it + count + dest] = cf_map_global[a_col_id];\n P_vals[p_col_it + count + dest] = a_value;\n }\n\n count += __popc( vote );\n }\n\n // We're done with that row of A, now reduce sum_N and sum_C\n#pragma unroll\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n sum_C += utils::shfl_xor( sum_C, mask );\n sum_N += utils::shfl_xor( sum_N, mask );\n }\n\n if ( lane_id == 0 )\n {\n // NOTE this matches the check for zero in HYPRE\n double div = (fabs(sum_C * diag[a_row_id]) == 0.0) ? 1. : sum_C * diag[a_row_id];\n alfa = -sum_N / div;\n }\n\n alfa = utils::shfl( alfa, 0 );\n\n // Scale the value of P\n for ( p_col_it += lane_id ; utils::any( p_col_it < p_col_end ) ; p_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n if ( p_col_it < p_col_end )\n {\n P_vals[p_col_it] *= alfa;\n }\n }\n } // if (assigned[a_row_id] == 1)\n } // Loop over rows\n}\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, typename KeyType >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_interp_weight_first_pass_kernel_opt( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const KeyType *__restrict cf_map_global,\n const bool *__restrict s_con,\n const Value_type *__restrict diag,\n const int *__restrict P_rows,\n KeyType *P_cols,\n Value_type *P_vals,\n const int *__restrict assigned)\n{\n const int NUM_WARPS_PER_CTA = CTA_SIZE / 32;\n const int NUM_WARPS_PER_GRID = gridDim.x * NUM_WARPS_PER_CTA;\n\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS_PER_CTA + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += NUM_WARPS_PER_GRID)\n {\n // Only do work if assigned[row_id] == pass\n KeyType coarse_fine_id = cf_map_global[a_row_id];\n\n // If coarse rows\n if ( assigned[a_row_id] == 0 )\n {\n if ( lane_id == 0 )\n {\n int p_row_it = P_rows[a_row_id];\n P_cols[p_row_it] = coarse_fine_id;\n P_vals[p_row_it] = Value_type( 1 );\n }\n }\n // If strong fine\n else if ( assigned[a_row_id] == -1 )\n {\n // Do nothing\n }\n // strongly connected to a coarse point\n else if (assigned[a_row_id] == 1)\n {\n // Load A row IDs.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n int p_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n p_col_tmp = P_rows[a_row_id + lane_id];\n }\n\n int p_col_it = utils::shfl( p_col_tmp, 0 );\n int p_col_end = utils::shfl( p_col_tmp, 1 );\n\n // Weak value.\n Value_type sum_N(0), sum_C(0), alfa(0);\n int count = 0;\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id(-1);\n Value_type a_value(0);\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n a_value = A_vals[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it a strongly-connected column.\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it a weakly connected node.\n //bool is_weakly_connected = is_off_diagonal && !is_strongly_connected;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map_global[a_col_id] == STRONG_FINE;\n // Is it a fine and strongly-connected column.\n bool is_assigned_strongly_connected = is_strongly_connected && assigned[a_col_id] == 0;\n\n // Update the weak value.\n if ( is_off_diagonal && !is_strong_fine)\n {\n sum_N += a_value;\n }\n\n if ( is_assigned_strongly_connected )\n {\n sum_C += a_value;\n }\n\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_assigned_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_assigned_strongly_connected )\n {\n P_cols[p_col_it + count + dest] = cf_map_global[a_col_id];\n P_vals[p_col_it + count + dest] = a_value;\n }\n\n count += __popc( vote );\n }\n\n // We're done with that row of A, now reduce sum_N and sum_C\n#pragma unroll\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n sum_C += utils::shfl_xor( sum_C, mask );\n sum_N += utils::shfl_xor( sum_N, mask );\n }\n\n if ( lane_id == 0 )\n {\n // NOTE this matches the check for zero in HYPRE\n double div = (fabs(sum_C * diag[a_row_id]) == 0.0) ? 1. : sum_C * diag[a_row_id];\n alfa = -sum_N / div;\n }\n\n alfa = utils::shfl( alfa, 0 );\n\n // Scale the value of P\n for ( p_col_it += lane_id ; utils::any( p_col_it < p_col_end ) ; p_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n if ( p_col_it < p_col_end )\n {\n P_vals[p_col_it] *= alfa;\n }\n }\n } // if (assigned[a_row_id] == 1)\n } // Loop over rows\n}\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, typename KeyType >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_interp_weight_kernel_opt( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const KeyType *__restrict C_hat,\n const int *__restrict C_hat_pos,\n const int *__restrict C_hat_start,\n const int *__restrict C_hat_size,\n const Value_type *__restrict diag,\n const int *__restrict P_rows,\n KeyType *P_cols,\n Value_type *P_vals,\n const int gmem_size,\n KeyType *g_keys,\n Value_type *g_vals,\n int *wk_work_queue,\n const int *__restrict assigned,\n const int *__restrict assigned_set,\n int pass)\n{\n const int NUM_WARPS = CTA_SIZE / 32;\n // The hash keys stored in shared memory.\n __shared__ KeyType s_keys[NUM_WARPS * SMEM_SIZE];\n // A shared location where threads propose a row of B to load.\n __shared__ int s_b_row_ids[CTA_SIZE];\n // A shared location where threads store a value of B to load.\n __shared__ Value_type s_b_values[CTA_SIZE];\n // The hash values stored in shared memory.\n __shared__ Value_type s_vals[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_map map( &s_keys[warp_id * SMEM_SIZE],\n &g_keys[a * gmem_size],\n &s_vals[warp_id * SMEM_SIZE],\n &g_vals[a * gmem_size],\n gmem_size );\n // Loop over rows of A.\n for ( ; a < A_num_rows ; a += gridDim.x*NUM_WARPS )\n {\n int a_row_id = assigned_set[a];\n\n // Clear the table.\n map.clear();\n // Rebuild C_hat.\n int c_hat_it = C_hat_start[a_row_id];\n int c_hat_end = c_hat_it + C_hat_size[a_row_id];\n map.load( c_hat_end - c_hat_it, &C_hat[c_hat_it], &C_hat_pos[c_hat_it] );\n // Load A row IDs.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // sums\n Value_type sum_N(0), sum_C(0), alfa(0);\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id(-1);\n Value_type a_value(0);\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n a_value = A_vals[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it a strongly-connected column.\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a fine and strongly-connected column.\n bool is_assigned_strongly_connected = is_strongly_connected && assigned[a_col_id] == pass - 1;\n // Is it a weakly connected node.\n bool is_weakly_connected = is_off_diagonal && !is_assigned_strongly_connected;\n\n // Update the weak value.\n if ( is_weakly_connected && !is_strong_fine)\n {\n sum_N += a_value;\n }\n\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_assigned_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_assigned_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n s_b_values[warp_id * WARP_SIZE + dest] = a_value;\n }\n\n int num_rows = __popc( vote );\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [b_col_it, bCol_end).\n int b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n Value_type uniform_val = s_b_values[warp_id * WARP_SIZE + k];\n // TODO: make sure we have better memory accesses. b_colBegin rather than bCol_it because we iterate twice.\n int b_col_begin = P_rows[b_row_id ];\n int b_col_end = P_rows[b_row_id + 1];\n\n // _iterate over the range of columns of B.\n for ( int b_col_it = b_col_begin + lane_id ; utils::any( b_col_it < b_col_end ) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n KeyType b_col_id(-1);\n Value_type tmp(0);\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = P_cols[b_col_it];\n tmp = P_vals[b_col_it] * uniform_val;\n sum_C += tmp;\n sum_N += tmp;\n }\n\n // Update C_hat values. If the value is not in C_hat it will be skipped (true parameter).\n // Need to get fine level id for b_col_id\n // if (a_row_id == 8 && b_col_id == 3) printf(\"updating 8,3: %lg\\n\",tmp);\n bool found = map.update( b_col_id, tmp );\n bool any_found = (__popc(utils::ballot(found)));\n }\n }\n }\n\n // We're done with that row of A, now reduce sum_N and sum_C\n#pragma unroll\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n sum_C += utils::shfl_xor( sum_C, mask );\n sum_N += utils::shfl_xor( sum_N, mask );\n }\n\n if ( lane_id == 0 )\n {\n // NOTE this matches the check for zero in HYPRE\n double div = (fabs(sum_C * diag[a_row_id]) == 0.0) ? 1. : sum_C * diag[a_row_id];\n alfa = -sum_N / div;\n // alfa = -sum_N/(sum_C*diag[a_row_id]);\n }\n\n alfa = utils::shfl( alfa, 0 );\n\n int p_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n p_col_tmp = P_rows[a_row_id + lane_id];\n }\n\n int p_col_it = utils::shfl( p_col_tmp, 0 );\n int p_col_end = utils::shfl( p_col_tmp, 1 );\n\n map.store_keys_scale_values( p_col_end - p_col_it, &P_cols[p_col_it], alfa, &P_vals[p_col_it] );\n }\n}\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE, typename KeyType >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_interp_weight_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const KeyType *__restrict C_hat,\n const int *__restrict C_hat_pos,\n const int *__restrict C_hat_start,\n const int *__restrict C_hat_size,\n const Value_type *__restrict diag,\n const int *__restrict P_rows,\n KeyType *P_cols,\n Value_type *P_vals,\n const int gmem_size,\n KeyType *g_keys,\n Value_type *g_vals,\n int *wk_work_queue,\n const int *__restrict assigned,\n int pass)\n{\n const int NUM_WARPS = CTA_SIZE / 32;\n // The hash keys stored in shared memory.\n __shared__ KeyType s_keys[NUM_WARPS * SMEM_SIZE];\n // A shared location where threads propose a row of B to load.\n __shared__ int s_b_row_ids[CTA_SIZE];\n // A shared location where threads store a value of B to load.\n __shared__ Value_type s_b_values[CTA_SIZE];\n // The hash values stored in shared memory.\n __shared__ Value_type s_vals[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_map map( &s_keys[warp_id * SMEM_SIZE],\n &g_keys[a_row_id * gmem_size],\n &s_vals[warp_id * SMEM_SIZE],\n &g_vals[a_row_id * gmem_size],\n gmem_size );\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Only do work if assigned[row_id] == pass\n if (assigned[a_row_id] != pass)\n {\n continue;\n }\n\n // Clear the table.\n map.clear();\n // Rebuild C_hat.\n int c_hat_it = C_hat_start[a_row_id];\n int c_hat_end = c_hat_it + C_hat_size[a_row_id];\n map.load( c_hat_end - c_hat_it, &C_hat[c_hat_it], &C_hat_pos[c_hat_it] );\n // Load A row IDs.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // sums\n Value_type sum_N(0), sum_C(0), alfa(0);\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id(-1);\n Value_type a_value(0);\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n a_value = A_vals[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it a strongly-connected column.\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a fine and strongly-connected column.\n bool is_assigned_strongly_connected = is_strongly_connected && assigned[a_col_id] == pass - 1;\n // Is it a weakly connected node.\n bool is_weakly_connected = is_off_diagonal && !is_assigned_strongly_connected;\n\n // Update the weak value.\n if ( is_weakly_connected && !is_strong_fine)\n {\n sum_N += a_value;\n }\n\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_assigned_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_assigned_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n s_b_values[warp_id * WARP_SIZE + dest] = a_value;\n }\n\n __syncwarp();\n\n int num_rows = __popc( vote );\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [b_col_it, bCol_end).\n int b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n Value_type uniform_val = s_b_values[warp_id * WARP_SIZE + k];\n // TODO: make sure we have better memory accesses. b_colBegin rather than bCol_it because we iterate twice.\n int b_col_begin = P_rows[b_row_id ];\n int b_col_end = P_rows[b_row_id + 1];\n\n // _iterate over the range of columns of B.\n for ( int b_col_it = b_col_begin + lane_id ; utils::any( b_col_it < b_col_end ) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n KeyType b_col_id(-1);\n Value_type tmp(0);\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = P_cols[b_col_it];\n tmp = P_vals[b_col_it] * uniform_val;\n sum_C += tmp;\n sum_N += tmp;\n }\n\n // Update C_hat values. If the value is not in C_hat it will be skipped (true parameter).\n // Need to get fine level id for b_col_id\n // if (a_row_id == 8 && b_col_id == 3) printf(\"updating 8,3: %lg\\n\",tmp);\n bool found = map.update( b_col_id, tmp );\n bool any_found = (__popc(utils::ballot(found)));\n }\n }\n }\n\n // We're done with that row of A, now reduce sum_N and sum_C\n#pragma unroll\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n sum_C += utils::shfl_xor( sum_C, mask );\n sum_N += utils::shfl_xor( sum_N, mask );\n }\n\n if ( lane_id == 0 )\n {\n // NOTE this matches the check for zero in HYPRE\n double div = (fabs(sum_C * diag[a_row_id]) == 0.0) ? 1. : sum_C * diag[a_row_id];\n alfa = -sum_N / div;\n // alfa = -sum_N/(sum_C*diag[a_row_id]);\n }\n\n alfa = utils::shfl( alfa, 0 );\n\n int p_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n p_col_tmp = P_rows[a_row_id + lane_id];\n }\n\n int p_col_it = utils::shfl( p_col_tmp, 0 );\n int p_col_end = utils::shfl( p_col_tmp, 1 );\n\n map.store_keys_scale_values( p_col_end - p_col_it, &P_cols[p_col_it], alfa, &P_vals[p_col_it] );\n }\n}\n\n\n} // namespace multipass\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n/*\n * grab the diagonal of the matrix\n */\ntemplate \n__global__\nvoid getDiagonalKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, const IndexType numRows, ValueType *diagonal)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n for (int j = offset; j < offset + numj; j++)\n {\n int jcol = column_indices[j];\n\n if (tIdx == jcol)\n {\n diagonal[tIdx] = values[j];\n }\n }\n }\n}\n/*************************************************************************\n * Implementing Extended+i algorithm from \\S 4.5 of:\n * \"Distance-two interpolation for parallel algebraic multigrid\"\n * Reference [4] on wiki\n ************************************************************************/\ntemplate \nMultipass_Interpolator >::Multipass_Interpolator(AMG_Config &cfg, const std::string &cfg_scope) : Base(cfg, cfg_scope)\n{\n this->m_use_opt_kernels = cfg.getParameter(\"use_opt_kernels\", \"default\");\n}\n\ntemplate \nMultipass_Interpolator >::~Multipass_Interpolator()\n{}\n\nenum { WARP_SIZE = 32, GRID_SIZE = 1024, SMEM_SIZE = 128 };\n\n\nstruct is_less_than_zero\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x < 0;\n }\n};\n\nstruct is_strong_fine\n{\n __host__ __device__\n bool operator()(int x)\n {\n return x == -3;\n }\n};\n\n__global__ void set_assigned_in_pass(int n, int pass, int* assigned, int* assigned_in_pass, int* assigned_offs)\n{\n int i = blockIdx.x*blockDim.x + threadIdx.x;\n if(i >= n) return;\n\n assigned_in_pass[i] = (assigned[i] == pass) ? 1 : 0;\n\n if(i == 0) \n {\n assigned_offs[0] = 0;\n }\n}\n\n__global__ void assigned_set_fill(int n, int* assigned_in_pass, int* assigned_offs, int* assigned_set)\n{\n int i = blockIdx.x*blockDim.x + threadIdx.x;\n if(i >= n) return;\n\n if(assigned_in_pass[i]) \n { \n assigned_set[assigned_offs[i]] = i; \n } \n}\n\ntemplate \ntemplate \nvoid Multipass_Interpolator>::compute_c_hat_opt_dispatch(\n const Matrix_d &A,\n const bool *s_con,\n const int *C_hat_start,\n int *C_hat_size,\n KeyType *C_hat,\n int *C_hat_pos,\n int *assigned,\n IntVector &assigned_set,\n int pass )\n{\n typedef typename IndPrecisionMap::Type Index_type;\n\n constexpr int cta_size = 128;\n constexpr int ngroups = cta_size / group_size;\n\n const int grid_size = assigned_set.size() / ngroups + 1;\n\n hipDeviceProp_t deviceProps = getDeviceProperties();\n size_t max_shmem_size = deviceProps.sharedMemPerMultiprocessor;\n\n constexpr int shmem_size =\n (sizeof(KeyType)+sizeof(Index_type))*ngroups*hash_size + group_size; \n\n if(shmem_size > max_shmem_size) \n { \n FatalError(\"In compute_values_opt the requested hash size is larger than max.\\n\", \n AMGX_ERR_NOT_IMPLEMENTED); \n } \n\n hipFuncSetAttribute(reinterpret_cast(multipass::compute_c_hat_kernel_opt \n ), \n hipFuncAttributeMaxDynamicSharedMemorySize, shmem_size); \n\n multipass::compute_c_hat_kernel_opt\n \n <<>>(\n assigned_set.size(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n s_con,\n C_hat_start,\n C_hat_size,\n C_hat,\n C_hat_pos,\n assigned,\n assigned_set.raw(),\n pass);\n cudaCheckError();\n}\n\ntemplate \nvoid Multipass_Interpolator >::generateInterpolationMatrix_1x1(Matrix_d &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix_d &P)\n{\n nvtxRange nvr_gim(__func__);\n\n // Implementation based on paper \"On long range interpolation operators for aggressive coarsening\" by Ulrike Meier Yang, section 4.3\n const int blockSize = 256;\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // get raw pointers to data\n const IndexType *Aoffsets = A.row_offsets.raw();\n const IndexType *Acolumn_indices = A.col_indices.raw();\n const ValueType *Avalues = A.values.raw();\n const IndexType Anum_rows = (int) A.get_num_rows();\n int *cf_map_ptr = cf_map.raw();\n bool *s_con_ptr = s_con.raw();\n int size_two_ring;\n int size_one_ring;\n\n if (A.is_matrix_distributed())\n {\n int offset;\n A.getOffsetAndSizeForView(ALL, &offset, &size_two_ring);\n A.getOffsetAndSizeForView(FULL, &offset, &size_one_ring);\n }\n else\n {\n size_one_ring = A.get_num_rows();\n size_two_ring = A.get_num_rows();\n }\n\n int numBlocksOneRing = std::min( 4096, (int) (size_one_ring + blockSize - 1) / blockSize );\n int numBlocksTwoRing = std::min( 4096, (int) (size_two_ring + blockSize - 1) / blockSize );\n int numBlocks = std::min( 4096, (int) (A.get_num_rows() + blockSize - 1) / blockSize );\n // ----------------------------------------------------------\n // First fill out the assigned array and count # of passes\n // ----------------------------------------------------------\n IntVector assigned( size_one_ring, -1);\n IntVector C_hat_start( size_one_ring + 1, 0 );\n // assigned[i] = -1: unassigned point\n // assigned[i] = 0: coarse point\n // assigned[i] = 1: fine point directly connected to coarse point\n // assigned[i] = 2: fine point distance 2 away from strongly connected coarse point\n // assigned[i] = ...\n // Initialize assigned array by marking the coarse and fine point directly connected to coarse points\n const int cta_size = 256;\n const int nWarps = cta_size / 32;\n int grid_size = std::min( 4096, (int) (A.get_num_rows() + nWarps - 1) / nWarps);\n if(this->m_use_opt_kernels)\n {\n constexpr int group_size = 8;\n int ngroups_per_block = cta_size / group_size;\n int grid_size = A.get_num_rows() / ngroups_per_block + 1;\n initializeAssignedArray_opt <<< grid_size, cta_size>>>(cf_map.raw(), assigned.raw(), A.row_offsets.raw(), A.col_indices.raw(), s_con.raw(), C_hat_start.raw(), Anum_rows);\n cudaCheckError();\n }\n else\n {\n initializeAssignedArray <<< grid_size, cta_size>>>(cf_map.raw(), assigned.raw(), A.row_offsets.raw(), A.col_indices.raw(), s_con.raw(), C_hat_start.raw(), Anum_rows);\n }\n cudaCheckError();\n // Count the number of passes and fill assigned for pass > 1\n int pass = 2;\n int max_num_passes = 10;\n // Count the number of unassigned nodes by checking which nodes have assigned[i] = -1\n IndexType num_unassigned = thrust_wrapper::count_if(assigned.begin(), assigned.end(), is_less_than_zero());\n cudaCheckError();\n IndexType num_strong_fine = thrust_wrapper::count_if(cf_map.begin(), cf_map.end(), is_strong_fine());\n cudaCheckError();\n IndexType num_unassigned_max = num_unassigned - num_strong_fine;\n\n // Exchange 1 ring of assigned array\n if (A.is_matrix_distributed())\n {\n assigned.dirtybit = 1;\n A.manager->exchange_halo(assigned, assigned.tag);\n // get max across all processes\n A.manager->getComms()->all_reduce_max(num_unassigned, num_unassigned_max);\n }\n\n // Fill the assigned array for the other nodes, going one pass at the time\n while (num_unassigned_max && pass < max_num_passes)\n {\n if(this->m_use_opt_kernels)\n {\n constexpr int group_size = 8;\n int ngroups_per_block = cta_size / group_size;\n int nblocks = A.get_num_rows() / ngroups_per_block + 1;\n fillAssignedArray_opt <<< nblocks, cta_size>>>(assigned.raw(), A.row_offsets.raw(), A.col_indices.raw(), s_con.raw(), Anum_rows, pass);\n cudaCheckError();\n }\n else\n {\n fillAssignedArray <<< grid_size, cta_size>>>(assigned.raw(), A.row_offsets.raw(), A.col_indices.raw(), s_con.raw(), Anum_rows, pass);\n cudaCheckError();\n }\n num_unassigned = thrust_wrapper::count_if(assigned.begin(), assigned.begin() + A.get_num_rows(), is_less_than_zero());\n cudaCheckError();\n num_unassigned_max = num_unassigned - num_strong_fine;\n\n if (A.is_matrix_distributed())\n {\n assigned.dirtybit = 1;\n A.manager->exchange_halo(assigned, assigned.tag);\n // get max across all processes\n A.manager->getComms()->all_reduce_max(num_unassigned, num_unassigned_max);\n }\n\n // printf(\"after pass %d, %d unassigned\\n\",pass, num_unassigned_max);\n pass++;\n }\n\n cudaCheckError();\n // Check that we haven't reached maximum number of passes\n //if (pass == max_num_passes)\n //FatalError(\"Multipass intepolator reached maximum number of passes, exiting\\n\",AMGX_ERR_NOT_IMPLEMENTED);\n int num_passes = pass;\n // -------------------------------------------------------------------------------------\n // Create array cf_map_global, which assigns a global index to all owned coarse points\n // -------------------------------------------------------------------------------------\n I64Vector_d cf_map_global(size_one_ring);\n DistributedArranger *prep = NULL;\n I64Vector_d P_col_indices_global;\n int coarsePoints;\n\n if (A.is_matrix_distributed())\n {\n prep = new DistributedArranger;\n int num_owned_fine_pts = A.get_num_rows();\n int num_owned_coarse_pts = thrust_wrapper::count_if(cf_map.begin(), cf_map.begin() + num_owned_fine_pts, is_non_neg());\n cudaCheckError();\n int num_halo_coarse_pts = thrust_wrapper::count_if(cf_map.begin() + num_owned_fine_pts, cf_map.end(), is_non_neg());\n cudaCheckError();\n coarsePoints = num_owned_coarse_pts + num_halo_coarse_pts;\n // Partially initialize the distributed manager of matrix P, using num_owned_coarse_pts\n prep->initialize_manager(A, P, num_owned_coarse_pts);\n int my_rank = A.manager->global_id();\n const int cta_size = 128;\n const int grid_size = std::min( 4096, (num_owned_fine_pts + cta_size - 1) / cta_size);\n // Fill the cf_map_global array, which will assign for each coarse pt, the global index of that coarse pt\n createCfMapGlobal <<< grid_size, cta_size>>>(cf_map.raw(), cf_map_global.raw(), P.manager->part_offsets_h[my_rank], num_owned_fine_pts);\n cudaCheckError();\n // Exchange the cf_map_global so that we know the coarse global id of halo nodes in one ring\n cf_map_global.dirtybit = 1;\n A.manager->exchange_halo(cf_map_global, cf_map_global.tag);\n }\n else\n {\n coarsePoints = (int) thrust_wrapper::count_if(cf_map.begin(), cf_map.end(), is_non_neg());\n cudaCheckError();\n const int cta_size = 128;\n const int grid_size = std::min( 4096, (A.get_num_rows() + cta_size - 1) / cta_size);\n fill_cf_map_global <<< grid_size, cta_size>>>(cf_map.raw(), cf_map_global.raw(), A.get_num_rows());\n cudaCheckError();\n }\n\n // ----------------------------------------------------------\n // Create an upper bound for the length of each row P\n // ----------------------------------------------------------\n Hash_Workspace exp_wk(true, GRID_SIZE);\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int work_offset = GRID_SIZE * NUM_WARPS;\n hipMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), hipMemcpyHostToDevice );\n cudaCheckError();\n\n // Do one ring exchange of C_hat_start, which have been filled for nodes with assigned=1\n if (A.is_matrix_distributed())\n {\n C_hat_start.dirtybit = 1;\n A.manager->exchange_halo(C_hat_start, C_hat_start.tag);\n }\n\n for (int i = 2; i < num_passes; i++)\n {\n if(this->m_use_opt_kernels)\n {\n constexpr int group_size = 8;\n int ngroups_per_block = CTA_SIZE / group_size;\n int nblocks = A.get_num_rows() / ngroups_per_block + 1;\n multipass::estimate_c_hat_size_kernel_opt< CTA_SIZE, group_size> <<< nblocks, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n s_con.raw(),\n assigned.raw(),\n C_hat_start.raw(),\n i);\n cudaCheckError();\n }\n else\n {\n multipass::estimate_c_hat_size_kernel< CTA_SIZE, WARP_SIZE> <<< 2048, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n s_con.raw(),\n assigned.raw(),\n C_hat_start.raw(),\n i);\n cudaCheckError();\n }\n\n if (A.is_matrix_distributed())\n {\n // TODO: This could be improved by only doing the exchange for points in current pass\n C_hat_start.dirtybit = 1;\n A.manager->exchange_halo(C_hat_start, C_hat_start.tag);\n }\n }\n\n cudaCheckError();\n }\n\n // Do an exclusive scan on C_hat_start to compute compute upper bound on row offsets\n thrust_wrapper::exclusive_scan( C_hat_start.begin( ), C_hat_start.end( ), C_hat_start.begin( ) );\n cudaCheckError();\n // Now C_hat_start contains the offsets\n // Create a temporary manager that can be used to exchange halo information on values and/or column indices, where there's more than 1 value associated with each node\n Matrix C_hat_matrix(1, 1, 1, 0);\n\n // Create manager to exchange C_hat data\n if (A.is_matrix_distributed())\n {\n prep->createTempManager(C_hat_matrix, A, C_hat_start);\n }\n\n // ----------------------------------------------------------\n // Count the actual number of nonzeroes in each row of P\n // ----------------------------------------------------------\n // Allocate memory to store columns/values.\n int nVals = C_hat_start[C_hat_start.size() - 1];\n // C_hat stores the global indices of the coarse nodes that each fine point depends on\n I64Vector_d C_hat( nVals, -1 );\n // C_hat_pos is an array used by the hash set structure (check with Julien for more details)\n IntVector C_hat_pos( nVals );\n // C_hat_size stores the number of coarse points on which each fine point depends on\n IntVector C_hat_size( size_one_ring + 1, 0 );\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int avg_nz_per_row = A.get_num_nz() / A.get_num_rows();\n if(this->m_use_opt_kernels)\n {\n constexpr int group_size = 8;\n int ngroups_per_block = CTA_SIZE / group_size;\n int nblocks = A.get_num_rows() / ngroups_per_block + 1;\n\n // This will fill C_hat for nodes with assigned = 1 (for fine points strongly connected to coarse points)\n multipass::compute_c_hat_first_pass_kernel_opt< CTA_SIZE, WARP_SIZE, group_size, int64_t> <<< nblocks, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n s_con.raw(),\n cf_map_global.raw(),\n C_hat_start.raw(),\n C_hat_size.raw(),\n C_hat.raw(),\n assigned.raw());\n\n cudaCheckError();\n }\n else\n {\n // This will fill C_hat for nodes with assigned = 1 (for fine points strongly connected to coarse points)\n multipass::compute_c_hat_first_pass_kernel< CTA_SIZE, WARP_SIZE, int64_t> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n s_con.raw(),\n cf_map_global.raw(),\n C_hat_start.raw(),\n C_hat_size.raw(),\n C_hat.raw(),\n assigned.raw()\n );\n cudaCheckError();\n }\n\n // Here exchange C_hat_size and C_hat for one ring neighbors\n if (A.is_matrix_distributed())\n {\n //TODO: We could probably reduce the amount of data exchanged here, since we only need to exchange for recently updated nodes\n C_hat.dirtybit = 1;\n C_hat_matrix.manager->exchange_halo(C_hat, C_hat.tag);\n C_hat_size.dirtybit = 1;\n A.manager->exchange_halo(C_hat_size, C_hat_size.tag);\n }\n\n for (int i = 2; i < num_passes; i++)\n {\n if(false && this->m_use_opt_kernels)\n {\n I64Vector_d C_hat_(C_hat);\n IntVector C_hat_pos_(C_hat_pos);\n IntVector C_hat_size_(C_hat_size);\n\n IntVector assigned_in_pass(assigned.size());\n IntVector assigned_offs(assigned.size()+1);\n\n int nthreads = 128;\n int nblocks = assigned.size() / nthreads + 1;\n set_assigned_in_pass<<>>(assigned.size(), i, assigned.raw(), assigned_in_pass.raw(), assigned_offs.raw());\n cudaCheckError();\n\n thrust_wrapper::inclusive_scan(assigned_in_pass.begin(), assigned_in_pass.end(), assigned_offs.begin()+1);\n\n int nassigned_in_pass = assigned_offs[assigned_offs.size()-1] ;\n if(nassigned_in_pass > 0)\n {\n IntVector assigned_set(nassigned_in_pass);\n assigned_set_fill<<>>(assigned.size(), assigned_in_pass.raw(), assigned_offs.raw(), assigned_set.raw());\n cudaCheckError();\n\n int max_c_hat_size = thrust::reduce(C_hat_size_.begin(), C_hat_size_.end(), -1, thrust::maximum());\n\n int hash_size = pow(2, ceil(log2(max_c_hat_size)));\n\n switch(hash_size)\n {\n case 1:\n case 2:\n case 4:\n case 8:\n case 16:\n case 32:\n compute_c_hat_opt_dispatch<32, 32>(\n A, s_con.raw(), C_hat_start.raw(),\n C_hat_size_.raw(), C_hat_.raw(), C_hat_pos_.raw(), assigned.raw(),\n assigned_set, i);\n break;\n case 64:\n compute_c_hat_opt_dispatch<64, 32>(\n A, s_con.raw(), C_hat_start.raw(),\n C_hat_size_.raw(), C_hat_.raw(), C_hat_pos_.raw(), assigned.raw(),\n assigned_set, i);\n break;\n case 128:\n compute_c_hat_opt_dispatch<128, 32>(\n A, s_con.raw(), C_hat_start.raw(),\n C_hat_size_.raw(), C_hat_.raw(), C_hat_pos_.raw(), assigned.raw(),\n assigned_set, i);\n break;\n case 256:\n compute_c_hat_opt_dispatch<256, 32>(\n A, s_con.raw(), C_hat_start.raw(),\n C_hat_size_.raw(), C_hat_.raw(), C_hat_pos_.raw(), assigned.raw(),\n assigned_set, i);\n break;\n case 512:\n compute_c_hat_opt_dispatch<512, 32>(\n A, s_con.raw(), C_hat_start.raw(),\n C_hat_size_.raw(), C_hat_.raw(), C_hat_pos_.raw(), assigned.raw(),\n assigned_set, i);\n break;\n default:\n FatalError(\"Size of C_hat in a single row too large, solve with original solvers.\", AMGX_ERR_INTERNAL);\n }\n }\n\n hipDeviceSynchronize();\n cudaCheckError();\n\n#if 0\n if (A.is_matrix_distributed())\n {\n //TODO: We could probably reduce the amount of data exchanged here, since we only need to exchange for recently updated nodes\n C_hat_.dirtybit = 1;\n C_hat_matrix.manager->exchange_halo(C_hat_, C_hat_.tag);\n C_hat_size_.dirtybit = 1;\n A.manager->exchange_halo(C_hat_size_, C_hat_size_.tag);\n }\n#endif\n }\n else\n {\n }\n\n // In case the hash table is not big enough, run multiple attemps\n int attempt = 0;\n\n for ( bool done = false ; !done && attempt < 10 ; ++attempt )\n {\n // Double the amount of GMEM (if needed).\n if ( attempt > 0 )\n {\n exp_wk.expand();\n }\n\n // Reset the status. TODO: Launch async copies.\n int status = 0;\n hipMemcpy( exp_wk.get_status(), &status, sizeof(int), hipMemcpyHostToDevice );\n cudaCheckError();\n // Compute the set C_hat.\n int work_offset = GRID_SIZE * NUM_WARPS;\n hipMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), hipMemcpyHostToDevice );\n cudaCheckError();\n // Run the computation.\n multipass::compute_c_hat_kernel< 8, CTA_SIZE, SMEM_SIZE, WARP_SIZE, int64_t> <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n s_con.raw(),\n C_hat_start.raw(),\n C_hat_size.raw(),\n C_hat.raw(),\n C_hat_pos.raw(),\n assigned.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_work_queue(),\n exp_wk.get_status(),\n i);\n cudaCheckError();\n\n if (A.is_matrix_distributed())\n {\n //TODO: We could probably reduce the amount of data exchanged here, since we only need to exchange for recently updated nodes\n C_hat.dirtybit = 1;\n C_hat_matrix.manager->exchange_halo(C_hat, C_hat.tag);\n C_hat_size.dirtybit = 1;\n A.manager->exchange_halo(C_hat_size, C_hat_size.tag);\n }\n\n cudaCheckError();\n // Read the result from count_non_zeroes.\n hipMemcpy( &status, exp_wk.get_status(), sizeof(int), hipMemcpyDeviceToHost );\n cudaCheckError();\n done = status == 0;\n }\n }\n }\n // At this stage, for each fine node (even in 1 ring halo), we have the global indices\n // of the interpolatory coarse points\n // count the number of non-zeros in the interpolation matrix,\n // For coarse point, nonzeros = 1\n // for strong fine, nonzeros = 0\n // for other points, nonzeros[i] = C_hat_size[i]\n IntVector nonZeroOffsets(size_one_ring + 1);\n IntVector nonZerosPerRow(size_one_ring);\n nonZerosPerRowSizeKernel <<< numBlocks, blockSize>>>(Anum_rows, cf_map_ptr, C_hat_size.raw(),\n nonZerosPerRow.raw());\n cudaCheckError();\n\n // Do a one ring halo exchange to know the lengths of the neighbors row\n if (A.is_matrix_distributed())\n {\n nonZerosPerRow.dirtybit = 1;\n A.manager->exchange_halo(nonZerosPerRow, nonZerosPerRow.tag);\n }\n\n // get total number of non-zeros with reduction\n int nonZeros = thrust_wrapper::reduce(nonZerosPerRow.begin(), nonZerosPerRow.end());\n cudaCheckError();\n // get the offsets with an exclusive scan\n thrust_wrapper::exclusive_scan(nonZerosPerRow.begin(), nonZerosPerRow.end(), nonZeroOffsets.begin());\n cudaCheckError();\n nonZeroOffsets[nonZeroOffsets.size() - 1] = nonZeros;\n // resize P\n P.resize(0, 0, 0, 1);\n P.addProps(CSR);\n // TODO: is the number of columns used anywhere?\n P.resize(size_one_ring, coarsePoints, nonZeros, 1);\n cudaCheckError();\n P_col_indices_global.resize(P.col_indices.size());\n Matrix P_temp(1, 1, 1, 0);\n\n if (A.is_matrix_distributed())\n {\n prep->createTempManager(P_temp, A, nonZeroOffsets);\n }\n\n IndexType *Poffsets = P.row_offsets.raw();\n IndexType *Pcolumn_indices = P.col_indices.raw();\n ValueType *Pvalues = P.values.raw();\n // copy nonzero offsets to the P matrix\n amgx::thrust::copy(nonZeroOffsets.begin(), nonZeroOffsets.end(), P.row_offsets.begin());\n cudaCheckError();\n // grab the diagonal terms\n VVector diag(size_one_ring);\n ValueType *diag_ptr = diag.raw();\n {\n find_diag_kernel_indexed_dia <<< numBlocksOneRing, blockSize>>>(\n size_one_ring,\n A.diag.raw(),\n A.values.raw(),\n diag.raw());\n }\n cudaCheckError();\n\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int work_offset = GRID_SIZE * NUM_WARPS;\n typedef typename MatPrecisionMap::Type Value_type;\n\n // Fill P.col_indices_global, P.values\n if(this->m_use_opt_kernels)\n {\n int grid_size = A.get_num_rows() / NUM_WARPS + 1;\n\n multipass::compute_interp_weight_first_pass_kernel_opt\n \n <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n cf_map_global.raw(),\n s_con.raw(),\n diag.raw(),\n P.row_offsets.raw(),\n P_col_indices_global.raw(),\n P.values.raw(),\n assigned.raw());\n cudaCheckError();\n }\n else\n {\n hipMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), hipMemcpyHostToDevice );\n cudaCheckError();\n\n // Compute the set C_hat.\n multipass::compute_interp_weight_first_pass_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n cf_map_global.raw(),\n s_con.raw(),\n diag.raw(),\n P.row_offsets.raw(),\n P_col_indices_global.raw(),\n P.values.raw(),\n assigned.raw());\n cudaCheckError();\n }\n\n if (A.is_matrix_distributed())\n {\n //TODO: We could probably reduce the amount of data exchanged here, since we only need to exchange for recently updated nodes\n P_col_indices_global.dirtybit = 1;\n P_temp.manager->exchange_halo(P_col_indices_global, P_col_indices_global.tag);\n P.values.dirtybit = 1;\n P_temp.manager->exchange_halo(P.values, P.values.tag);\n }\n\n for (int i = 2; i < num_passes; i++)\n {\n if(this->m_use_opt_kernels)\n {\n IntVector assigned_in_pass(assigned.size());\n IntVector assigned_offs(assigned.size() + 1);\n\n int nthreads = 128;\n int nblocks = assigned.size() / nthreads + 1;\n set_assigned_in_pass<<>>(assigned.size(), i, assigned.raw(), assigned_in_pass.raw(), assigned_offs.raw());\n cudaCheckError();\n thrust_wrapper::inclusive_scan(assigned_in_pass.begin(), assigned_in_pass.end(), assigned_offs.begin()+1);\n\n IntVector assigned_set(assigned_offs[assigned_offs.size()-1]);\n assigned_set_fill<<>>(assigned.size(), assigned_in_pass.raw(), assigned_offs.raw(), assigned_set.raw());\n cudaCheckError();\n\n Hash_Workspace exp_wk2(true, 4096);\n multipass::compute_interp_weight_kernel_opt <<< 4096, CTA_SIZE>>>(\n assigned_set.size(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat.raw(),\n C_hat_pos.raw(),\n C_hat_start.raw(),\n C_hat_size.raw(),\n diag.raw(),\n P.row_offsets.raw(),\n P_col_indices_global.raw(),\n P.values.raw(),\n exp_wk2.get_gmem_size(),\n exp_wk2.get_keys(),\n exp_wk2.get_vals(),\n exp_wk2.get_work_queue(),\n assigned.raw(),\n assigned_set.raw(),\n i);\n cudaCheckError();\n }\n else\n {\n hipMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), hipMemcpyHostToDevice );\n cudaCheckError();\n multipass::compute_interp_weight_kernel <<< GRID_SIZE, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat.raw(),\n C_hat_pos.raw(),\n C_hat_start.raw(),\n C_hat_size.raw(),\n diag.raw(),\n P.row_offsets.raw(),\n P_col_indices_global.raw(),\n P.values.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_vals(),\n exp_wk.get_work_queue(),\n assigned.raw(),\n i);\n cudaCheckError();\n }\n\n if (A.is_matrix_distributed())\n {\n //TODO: This is exchanging too much data\n P_col_indices_global.dirtybit = 1;\n P_temp.manager->exchange_halo(P_col_indices_global, P_col_indices_global.tag);\n P.values.dirtybit = 1;\n P_temp.manager->exchange_halo(P.values, P.values.tag);\n }\n\n cudaCheckError();\n }\n\n if (!A.is_matrix_distributed())\n {\n const int cta_size = 128;\n const int grid_size = std::min( 4096, ( (int) P.col_indices.size() + cta_size - 1) / cta_size);\n // Simply cast from int64 to int\n fill_P_col_indices <<< grid_size, cta_size>>>(P_col_indices_global.raw(), P.col_indices.raw(), (int) P.col_indices.size());\n cudaCheckError();\n }\n else\n {\n // Resize P to only include the owned rows and columns\n int P_num_nnz = P.row_offsets[P.get_num_rows()];\n P.resize(A.get_num_rows(), P.manager->halo_offsets[0], P_num_nnz, 1);\n // Using the global col indices, identify neighbors for P.manager, fill B2L_maps, etc...\n prep->initialize_manager_from_global_col_indices(P, P_col_indices_global);\n // Create the list of interior and boundary nodes (this is necessary if doing masked SpMV operations)\n prep->createRowsLists(P, true);\n }\n\n // Delete the temporary distributed arranger\n if (prep != NULL) { delete prep; }\n}\n\ntemplate< class T_Config>\nvoid Multipass_InterpolatorBase::generateInterpolationMatrix(Matrix &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix &P)\n{\n P.set_initialized(0);\n ViewType oldView = A.currentView();\n A.setView(OWNED);\n\n if (A.get_block_size() == 1)\n {\n generateInterpolationMatrix_1x1(A, cf_map, s_con, scratch, P);\n }\n else\n {\n FatalError(\"Unsupported dimensions for multipass interpolator\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n //P.computeDiagonal();\n P.set_initialized(1);\n#if 0\n P.printToFile(\"P.out\", \"\", -1, -1);\n#endif\n}\n\n#define AMGX_CASE_LINE(CASE) template class Multipass_InterpolatorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Multipass_Interpolator::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/classical/interpolators/multipass.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace amgx\n{\n\nnamespace detail\n{\n\ntemplate< class CsrMatrixType >\nstruct find_row_diagonal\n{\n// typedef typename CsrMatrixType::values_array_type::value_type value_type;\n typedef typename CsrMatrixType::value_type value_type;\n const CsrMatrixType &my_A;\n find_row_diagonal( const CsrMatrixType &A ) : my_A( A ) {}\n\n // Operator to extract the diagonal value.\n value_type operator()( int irow ) const\n {\n for ( int i = my_A.row_offsets[irow], end = my_A.row_offsets[irow + 1] ; i < end ; ++i )\n if ( my_A.col_indices[i] == irow )\n {\n return my_A.values[i];\n }\n\n // assert( false );\n return value_type( 0 );\n }\n};\n\ntemplate< class CsrMatrixType >\nstruct count_nnz_per_row\n{\n const CsrMatrixType &my_A;\n static const AMGX_VecPrecision vecPrec = CsrMatrixType::TConfig::vecPrec;\n static const AMGX_MatPrecision matPrec = CsrMatrixType::TConfig::matPrec;\n static const AMGX_IndPrecision indPrec = CsrMatrixType::TConfig::indPrec;\n typedef Vector > BVector;\n typedef Vector > IntVector;\n\n const BVector &my_are_sc;\n const IntVector &my_cf_map;\n count_nnz_per_row( const CsrMatrixType &A,\n const BVector &are_sc,\n const IntVector &cf_map ) :\n my_A( A ),\n my_are_sc( are_sc ),\n my_cf_map( cf_map )\n {}\n\n // Operator to extract the diagonal value.\n int operator()( int i ) const\n {\n if ( my_cf_map[i] >= 0 ) // If it's a coarse node, only one NZ element.\n {\n return 1;\n }\n\n // Count the number of strongly connected coarse nodes.\n int count = 0;\n\n for ( int itj = my_A.row_offsets[i] ; itj < my_A.row_offsets[i + 1] ; ++itj )\n {\n int j = my_A.col_indices[itj];\n\n if ( my_cf_map[j] >= 0 && my_are_sc[itj] )\n {\n count++;\n }\n }\n\n return count;\n }\n};\n\ntemplate< class CsrMatrixType >\nstruct compute_weights\n{\n typedef typename CsrMatrixType::value_type value_type;\n typedef typename CsrMatrixType::MVector VVector;\n static const AMGX_VecPrecision vecPrec = CsrMatrixType::TConfig::vecPrec;\n static const AMGX_MatPrecision matPrec = CsrMatrixType::TConfig::matPrec;\n static const AMGX_IndPrecision indPrec = CsrMatrixType::TConfig::indPrec;\n typedef Vector > BVector;\n typedef Vector > IntVector;\n const CsrMatrixType &my_A;\n const VVector &my_diag;\n const BVector &my_are_sc;\n const IntVector &my_cf_map;\n IntVector *my_edges_markers;\n CsrMatrixType &my_P;\n\n // Constructor.\n compute_weights( const CsrMatrixType &A,\n const VVector &diag,\n const BVector &are_sc,\n const IntVector &cf_map,\n IntVector *edges_markers,\n CsrMatrixType &P ) :\n my_A( A ),\n my_diag( diag ),\n my_are_sc( are_sc ),\n my_cf_map( cf_map ),\n my_edges_markers( edges_markers ),\n my_P( P )\n { }\n\n // Get the thread index. Always 0 if OpenMP is off.\n inline\n int get_thread_id( ) const\n {\n return 0;\n }\n\n // Operator to find the weights.\n void operator()( int i )\n {\n // Skip coarse nodes.\n int nnz = my_P.row_offsets[i];\n int id = my_cf_map[i];\n\n if ( id >= 0 )\n {\n my_P.col_indices[nnz] = id;\n my_P.values[nnz] = value_type( 1 );\n return;\n }\n\n assert( my_A.row_offsets[i] < my_A.row_offsets[i + 1] );\n const int strong_marker = -2;\n\n // Initialize the F-C connections to 0 and mark F-F connections.\n for ( int itj = my_A.row_offsets[i] ; itj < my_A.row_offsets[i + 1] ; ++itj )\n {\n if ( !my_are_sc[itj] )\n {\n continue;\n }\n\n int j = my_A.col_indices[itj];\n\n if ( my_cf_map[j] == FINE ) // It's a F-F connection, mark it.\n {\n assert( my_cf_map[j] < 0 );\n my_edges_markers[get_thread_id( )][j] = strong_marker;\n continue;\n }\n\n assert( nnz < my_P.row_offsets[i + 1] );\n my_edges_markers[get_thread_id( )][j] = nnz; // It's a F-C connection, mark the coarse end point.\n my_P.col_indices[nnz] = my_cf_map[j];\n my_P.values[nnz] = value_type( 0 );\n nnz++;\n }\n\n // Make sure we exhausted all the indices.\n assert( nnz == my_P.row_offsets[i + 1] );\n // Reset NNZ variable.\n nnz = my_P.row_offsets[i];\n // The diagonal element a_ii.\n value_type diag = my_diag[i];\n\n // Loop over the neighbours of the node n_i to add contributions.\n for ( int itj = my_A.row_offsets[i] ; itj < my_A.row_offsets[i + 1] ; ++itj )\n {\n int j = my_A.col_indices[itj];\n\n if ( i == j ) // Skip the diagonal.\n {\n continue;\n }\n\n // 1st case: j is a coarse node and it strongly influences i, so we accumulate a_ij.\n if ( my_edges_markers[get_thread_id( )][j] >= nnz )\n {\n int itk = my_edges_markers[get_thread_id( )][j];\n my_P.values[itk] += my_A.values[itj];\n continue;\n }\n // 2nd case: j is a fine node and strongly influences i. We distribute a_ij to the coarse\n // nodes that strongly influence i (not the diagonal).\n else if ( my_edges_markers[get_thread_id( )][j] == strong_marker )\n {\n value_type sum( 0 );\n int diag_sign = std::signbit( my_diag[i] );\n\n for ( int itk = my_A.row_offsets[j] ; itk < my_A.row_offsets[j + 1] ; ++itk )\n {\n int k = my_A.col_indices[itk];\n\n if ( my_edges_markers[get_thread_id( )][k] >= nnz && std::signbit( my_A.values[itk] ) != diag_sign )\n {\n sum += my_A.values[itk];\n }\n }\n\n if ( sum == value_type( 0 ) )\n {\n diag += my_A.values[itj];\n continue;\n }\n\n value_type d = my_A.values[itj] / sum;\n\n for ( int itk = my_A.row_offsets[j] ; itk < my_A.row_offsets[j + 1] ; ++itk )\n {\n int it = my_edges_markers[get_thread_id( )][my_A.col_indices[itk]];\n\n if ( it >= nnz && std::signbit( my_A.values[itk] ) != diag_sign )\n {\n my_P.values[it] += d * my_A.values[itk];\n }\n }\n }\n else if ( my_cf_map[j] == FINE )\n {\n diag += my_A.values[itj];\n }\n }\n\n if ( diag == value_type( 0 ) )\n {\n diag = my_diag[i];\n\n if ( diag == value_type( 0 ) ) // REALLY ?\n {\n diag = value_type( 1 );\n }\n }\n\n assert( diag != value_type( 0 ) );\n diag = value_type( -1.0 ) / diag;\n\n for ( ; nnz < my_P.row_offsets[i + 1] ; ++nnz )\n {\n my_P.values[nnz] *= diag;\n }\n\n // strong_marker--;\n }\n};\n\n}\n\n/*************************************************************************\n* create the interpolation matrix P\n************************************************************************/\ntemplate \nvoid Distance1_Interpolator >::generateInterpolationMatrix_1x1(\n Matrix_h &A,\n IntVector &cf_map,\n BVector &are_sc,\n IntVector &wk,\n Matrix_h &P)\n{\n // The diagonal of A.\n VVector diag( A.get_num_rows() );\n thrust_wrapper::transform( amgx::thrust::make_counting_iterator( 0 ),\n amgx::thrust::make_counting_iterator( A.get_num_rows() ),\n diag.begin( ),\n detail::find_row_diagonal( A ) );\n cudaCheckError();\n // Count the number of coarse points.\n int n_coarse = 0;\n\n for ( int i = 0 ; i < A.get_num_rows() ; ++i )\n if ( cf_map[i] >= 0) { n_coarse++; }\n\n // Count the number of non-zero elements per row of P.\n IntVector nnz_per_row( A.get_num_rows() + 1 );\n amgx::thrust::transform( amgx::thrust::make_counting_iterator( 0 ),\n amgx::thrust::make_counting_iterator( A.get_num_rows() ),\n nnz_per_row.begin( ),\n detail::count_nnz_per_row( A, are_sc, cf_map ) );\n cudaCheckError();\n nnz_per_row[A.get_num_rows()] = 0;\n // Find the total number of non-zeroes. \n // TODO: merge with the exclusive scan.\n int nnz = amgx::thrust::reduce( nnz_per_row.begin( ), nnz_per_row.end( ) );\n cudaCheckError();\n // Resize P so that there is enough space to store the non-zero elements.\n P.addProps(CSR);\n P.resize( A.get_num_rows(), n_coarse, nnz );\n // Compute row offsets of P.\n thrust_wrapper::exclusive_scan( nnz_per_row.begin( ),\n nnz_per_row.end( ),\n P.row_offsets.begin( ) );\n cudaCheckError();\n // For each row we compute the weights.\n IntVector edges_markers_on_stack( A.get_num_rows(), -1 );\n IntVector *edges_markers = &edges_markers_on_stack;\n detail::compute_weights compute_fct( A, diag, are_sc, cf_map, edges_markers, P );\n\n amgx::thrust::for_each( amgx::thrust::host,\n amgx::thrust::make_counting_iterator( 0 ),\n amgx::thrust::make_counting_iterator( A.get_num_rows() ),\n compute_fct );\n cudaCheckError();\n}\n\n/*************************************************************************\n* create the interpolation matrix P on GPU\n************************************************************************/\n\nenum SETS {SET_STRONG_COARSE = 1,\n SET_WEAK_COARSE = 2,\n SET_STRONG_FINE = 4,\n SET_STRONG_FINE_NO_COMMON = 8,\n SET_WEAK_FINE = 16,\n SET_DIAGONAL = 32\n };\n\n/*\n * grab the diagonal of the matrix\n */\n/*template \n__global__\nvoid getDiagonalKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, const IndexType numRows, ValueType *diagonal)\n{\n const int tIdx = threadIdx.x + blockDim.x*blockIdx.x;\n\n if (tIdx >= numRows) return;\n\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx+1]-offset;\n\n for (int j=offset; j < offset+numj; j++)\n {\n int jcol = column_indices[j];\n if (tIdx == jcol)\n {\n diagonal[tIdx] = values[j];\n }\n }\n}*/\n\n/*\n * Count the non-zeros in final output\n */\ntemplate \n__global__\nvoid numNonZerosVecKernel(const IndexType *offsets, const IndexType *column_indices,\n const IndexType numRows, int *cf_map, const bool *s_con,\n int *nonZerosVec)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n int numOnRow = 0;\n\n // count strong coarse influences\n if (cf_map[tIdx] == FINE)\n {\n for (int j = offset; j < offset + numj; j++)\n {\n if (!s_con[j]) { continue; } // if weakly connected, ignore\n\n int jcol = column_indices[j];\n\n if (cf_map[jcol] >= 0) // if coarse\n {\n numOnRow++;\n }\n }\n }\n else\n {\n numOnRow++;\n }\n\n nonZerosVec[tIdx] = numOnRow;\n }\n}\n\n/*\n * Calculate D_i vector\n * - Weakly influencing and strong fine with no common C-points to be added to diagonal term\n */\ntemplate \n__global__\nvoid calculateDKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, int num_rows, ValueType *diagonal,\n int *set_fields, ValueType *D)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < num_rows; tIdx += blockDim.x * gridDim.x)\n {\n const int offset = offsets[tIdx];\n const int numk = offsets[tIdx + 1] - offset;\n ValueType sum = 0;\n int mask = SET_WEAK_FINE;\n\n for (int k = offset; k < offset + numk; k++)\n {\n if (set_fields[k] & mask)\n {\n sum += values[k];\n }\n }\n\n D[tIdx] += sum;\n }\n}\n\n/*\n * Calculate the k dependent term in w_i,j\n */\ntemplate \n__global__\nvoid calculateBKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, int num_rows, ValueType *diagonal, ValueType *diag_inc,\n int *set_fields, int *cf_map, int *nonZerosVec, ValueType *B)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < num_rows; tIdx += blockDim.x * gridDim.x)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n int localNonZeros = 0;\n const int nonZerosOffset = nonZerosVec[tIdx];\n ValueType sum = 0.0, top = 0.0, bottom_sum = 0.0, a_ik;\n\n if (cf_map[tIdx] >= 0) // If I am a coarse row\n {\n B[nonZerosOffset] = 1;\n return;\n }\n\n diag_inc[tIdx] = 0;\n int flag = 0;\n int first_j_loop = 0;\n ValueType tol = 1e-10;\n\n // otherwise\n for (int j = offset; j < offset + numj; j++)\n {\n sum = 0.0;\n bottom_sum = 0.0;\n\n if (set_fields[j] & SET_STRONG_COARSE) // we actually care about this j\n {\n int jcol = column_indices[j]; // this is the only part of j I need / care about.\n\n if (flag == 0)\n {\n first_j_loop = 1;\n flag = 1;\n }\n else\n {\n first_j_loop = 0;\n }\n\n // run through the columns of row i, summing over all k\n for (int k = offset; k < offset + numj; k++)\n {\n top = 0.0;\n bottom_sum = 0.0;\n\n if ( (set_fields[k] & SET_STRONG_FINE) || (set_fields[k] & SET_STRONG_FINE_NO_COMMON)) // k we're interested in\n {\n // part of the desired set - get column, diagonal and a_ik\n int kcol = column_indices[k];\n a_ik = values[k]; // get a_i,k term\n // now find a_k,j\n ValueType diag = diagonal[kcol];\n int sgn = diag < 0.0 ? -1 : 1;\n\n // now I need to find the j-th column on the k-th row. another loop (sigh)\n for (int n = offsets[kcol]; n < offsets[kcol + 1]; n++)\n {\n int ncol = column_indices[n];\n\n if (ncol == jcol) // this is the entry I need\n {\n if (sgn * values[n] < 0)\n {\n // add to the sum\n top = a_ik * values[n]; // discard diagonal\n }\n }\n }\n\n // calculate the bottom sum - loop through columns of i\n for (int m = offset; m < offset + numj; m++)\n {\n if (set_fields[m] & SET_STRONG_COARSE)\n {\n // this is a point I care about\n int mcol = column_indices[m];\n\n // now find the m-th column of the k-th row and add to sum\n for (int n = offsets[kcol]; n < offsets[kcol + 1]; n++)\n {\n int ncol = column_indices[n];\n\n if (ncol == mcol)\n {\n if (sgn * values[n] < 0)\n {\n bottom_sum += values[n];\n }\n }\n }\n }\n }\n\n if (abs(bottom_sum) < tol)\n {\n if (first_j_loop == 1)\n {\n //add it to the diagonal....\n diag_inc[tIdx] += values[k];\n }\n }\n else\n {\n sum += top / bottom_sum;\n }\n }\n }\n\n B[nonZerosOffset + localNonZeros] = sum;\n localNonZeros++;\n }\n }\n }\n}\n\n/*\n * count the number of strong coarse points per row\n * - was used for degugging, might be useful in the future\n */\ntemplate \n__global__\nvoid coarsePerRow(const IndexType *offsets, int num_rows, int *set_fields, int *output)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < num_rows; tIdx += blockDim.x * gridDim.x)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n int localCoarse = 0;\n\n for (int j = offset; j < offset + numj; j++)\n {\n if (set_fields[j] & SET_STRONG_COARSE)\n {\n localCoarse++;\n }\n }\n\n output[tIdx] = localCoarse;\n }\n}\n\n/*\n * Assemble final w_ij terms\n */\ntemplate \n__global__\nvoid calculateWKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, int num_rows, ValueType *diagonal,\n int *set_fields, int *nonZerosVec, ValueType *D, ValueType *B,\n ValueType *w, IndexType *Pcolumn_indices, int *cf_map)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < num_rows; tIdx += blockDim.x * gridDim.x)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n int localNonZeros = 0;\n const int nonZerosOffset = nonZerosVec[tIdx];\n ValueType diag = diagonal[tIdx];\n\n if (cf_map[tIdx] >= 0)\n {\n w[nonZerosOffset] = 1.0;\n Pcolumn_indices[nonZerosOffset] = cf_map[tIdx];\n continue;\n }\n\n for (int j = offset; j < offset + numj; j++)\n {\n if (set_fields[j] & SET_STRONG_COARSE)\n {\n ValueType aij = values[j];\n ValueType Bij = B[nonZerosOffset + localNonZeros];\n ValueType Di = D[tIdx];\n Pcolumn_indices[nonZerosOffset + localNonZeros] = cf_map[column_indices[j]];\n double bottom_sum = (fabs(diag + Di) < 1e-10) ? 1. : diag + Di;\n w[nonZerosOffset + localNonZeros] = -1.0 / bottom_sum * (aij + Bij);\n localNonZeros++;\n }\n }\n }\n}\n\ntemplate \n__global__\nvoid markCoarseEdgesKernel(const IndexType *offsets, const IndexType *column_indices,\n const size_t numRows, const int *cf_map, const bool *s_con,\n bool *coarseMark)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x)\n {\n const int rowBegin = offsets[tIdx];\n const int rowEnd = offsets[tIdx + 1];\n\n // loop over columns\n for (int j = rowBegin; j < rowEnd; j++)\n {\n int jcol = column_indices[j];\n\n if (tIdx == jcol) { continue; } // skip the diagonal\n\n if (cf_map[jcol] >= 0)\n {\n // mark this point\n coarseMark[j] = true;\n }\n }\n }\n}\n\ntemplate \n__device__\nbool sparseIntersectDevice(const IndexType *column_indices, ValueType *values,\n int begin1, int end1, int begin2, int end2)\n{\n int i1 = begin1;\n int i2 = begin2;\n\n // check arrays non zero length\n if (begin1 >= end1 || begin2 >= end2) { return false; }\n\n bool r = false;\n\n while ( true )\n {\n int col1 = column_indices[i1];\n int col2 = column_indices[i2];\n\n if (col1 == col2) // this is a match\n {\n if (values[i1] && values[i2])\n {\n r = true;\n break;\n }\n\n i1++;\n i2++;\n\n if ( i1 >= end1 || i2 >= end2 ) { break; }\n }\n else if (col1 > col2)\n {\n i2++;\n\n if (i2 >= end2) { break; }\n }\n else\n {\n i1++;\n\n if (i1 >= end1) { break; }\n }\n }\n\n return r;\n}\n\ntemplate \n__global__\nvoid categoriseEdgesKernel(const IndexType *offsets, const IndexType *column_indices,\n const size_t numRows, const int *cf_map, const bool *s_con,\n bool *coarse_mark, int *set_fields)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x)\n {\n // set_fields[tIdx] = 1;\n if (cf_map[tIdx] >= 0) { continue; } // we don't need to categorise C-? edges\n\n const int rowBegin = offsets[tIdx];\n const int rowEnd = offsets[tIdx + 1];\n\n // we know I'm a fine row now - loop over non-zeros\n for (int j = rowBegin; j < rowEnd; j++)\n {\n int jcol = column_indices[j];\n\n if (jcol == tIdx) { continue; }\n\n // check F-C edges\n if (cf_map[jcol] >= 0)\n {\n if (s_con[j])\n {\n set_fields[j] |= SET_STRONG_COARSE;\n }\n else\n {\n set_fields[j] |= SET_WEAK_COARSE;\n }\n }\n else // F-F edges\n {\n if (!s_con[j])\n {\n set_fields[j] |= SET_WEAK_FINE;\n }\n else // check if there exists a common C-point\n {\n // intersect i and jcol\n bool common_C = true;\n common_C = sparseIntersectDevice(column_indices, coarse_mark, rowBegin, rowEnd,\n offsets[jcol], offsets[jcol + 1]);\n\n if (common_C)\n {\n set_fields[j] |= SET_STRONG_FINE;\n }\n else\n {\n set_fields[j] |= SET_STRONG_FINE_NO_COMMON;\n }\n }\n }\n }\n }\n}\n\ntemplate \nvoid Distance1_Interpolator >::setSetsGPU(const Matrix_d &A, const BVector &s_con,\n const IntVector &cf_map,\n IntVector &set_fields)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n const int blocksize = 64;\n const int numBlocks = std::min( AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() / blocksize + 1) );\n // raw pointers for matrix structure\n const IndexType *offsets = A.row_offsets.raw();\n const IndexType *column_indices = A.col_indices.raw();\n // raw pointers for coarse/fine map and strength of connection array\n const bool *s_con_ptr = s_con.raw();\n const int *cf_map_ptr = cf_map.raw();\n int *set_fields_ptr = set_fields.raw();\n // memory for CoarseMark array\n BVector coarseMark(A.get_num_nz(), false);\n bool *coarseMark_ptr = coarseMark.raw();\n markCoarseEdgesKernel <<< numBlocks, blocksize>>>(offsets, column_indices, A.get_num_rows(),\n cf_map_ptr, s_con_ptr, coarseMark_ptr);\n cudaCheckError();\n // now we have the strong coarse edges marked, categorise all edges\n categoriseEdgesKernel <<< numBlocks, blocksize>>>(offsets, column_indices, A.get_num_rows(),\n cf_map_ptr, s_con_ptr, coarseMark_ptr,\n set_fields_ptr);\n cudaCheckError();\n}\n\ntemplate \nvoid Distance1_Interpolator >::generateInterpolationMatrix_1x1(Matrix_d &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix_d &P)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // Assemble all the raw pointers needed\n // Matrix A\n const IndexType *Aoffsets_ptr = A.row_offsets.raw();\n const IndexType *Acolumn_indices_ptr = A.col_indices.raw();\n const ValueType *Avalues_ptr = A.values.raw();\n // Coarse/Fine map, strength of connection, scratch\n int *cf_map_ptr = cf_map.raw();\n const bool *s_con_ptr = s_con.raw();\n int *scratch_ptr = scratch.raw();\n // temporary vectors - diagonal and non-zero offsets\n VVector diag(A.get_num_rows(), 0.0);\n IntVector nonZerosVec(A.get_num_rows() + 1, 0);\n ValueType *diag_ptr = diag.raw();\n int *nonZerosVec_ptr = nonZerosVec.raw();\n // NB - will get pointers for P *after* it has been initialised\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() / blocksize + 1));\n const IndexType numRows = (IndexType) A.get_num_rows();\n // extract the diagonal of A\n find_diag_kernel_indexed_dia <<< numBlocks, blocksize>>>(A.get_num_rows(),\n A.diag.raw(),\n A.values.raw(),\n diag_ptr);\n cudaCheckError();\n /* getDiagonalKernel<<>>(Aoffsets_ptr,Acolumn_indices_ptr,Avalues_ptr,\n numRows,diag_ptr);*/\n cudaCheckError();\n // count the number of non-zeros in final output per row\n numNonZerosVecKernel <<< numBlocks, blocksize>>>(Aoffsets_ptr, Acolumn_indices_ptr, numRows,\n cf_map_ptr, s_con_ptr, nonZerosVec_ptr);\n cudaCheckError();\n // now I have the non-zeros per row for matrix P, count the non-zeros on each row\n // to get total NNZ\n // sum non-zeros per row\n int NNZidx = thrust_wrapper::reduce(nonZerosVec.begin(), nonZerosVec.end());\n cudaCheckError();\n thrust_wrapper::exclusive_scan(nonZerosVec.begin(), nonZerosVec.end(), nonZerosVec.begin());\n cudaCheckError();\n nonZerosVec[A.get_num_rows()] = NNZidx;\n // generate sets on the device\n IntVector set_fields(A.get_num_nz(), 0);\n setSetsGPU(A, s_con, cf_map, set_fields);\n int *set_field_ptr = set_fields.raw();\n // count the coarse points first\n // (no need to worry about >4B coarse points)\n int coarsePoints = (int) amgx::thrust::count_if(cf_map.begin(), cf_map.end(), is_non_neg());\n cudaCheckError();\n // I have all the information needed - create P and get the pointers to its data\n P.addProps(CSR);\n P.resize(A.get_num_rows(), coarsePoints, NNZidx, 1);\n IndexType *Pcolumn_indices_ptr = P.col_indices.raw();\n ValueType *Pvalues_ptr = P.values.raw();\n // now P is initialised, get row offsets from nonZerosVec with an exclusive prefix sum\n // set final value\n amgx::thrust::copy(nonZerosVec.begin(), nonZerosVec.end(), P.row_offsets.begin());\n cudaCheckError();\n IndexType *Poffsets_ptr = P.row_offsets.raw();\n // assign temp memory\n VVector B(NNZidx, 0.0);\n VVector D(A.get_num_rows(), 0.0);\n VVector W(NNZidx, 0.0);\n // get the raw pointers for temp memory\n ValueType *D_ptr = D.raw();\n ValueType *B_ptr = B.raw();\n ValueType *W_ptr = W.raw();\n // call each kernel in turn\n calculateBKernel <<< numBlocks, blocksize>>>(Aoffsets_ptr, Acolumn_indices_ptr, Avalues_ptr,\n (int) A.get_num_rows(), diag_ptr, D_ptr, set_field_ptr, cf_map_ptr,\n nonZerosVec_ptr, B_ptr);\n cudaCheckError();\n calculateDKernel <<< numBlocks, blocksize>>>(Aoffsets_ptr, Acolumn_indices_ptr, Avalues_ptr,\n (int) A.get_num_rows(), diag_ptr, set_field_ptr, D_ptr);\n cudaCheckError();\n calculateWKernel <<< numBlocks, blocksize>>>(Aoffsets_ptr, Acolumn_indices_ptr, Avalues_ptr,\n (int) A.get_num_rows(), diag_ptr, set_field_ptr, nonZerosVec_ptr,\n D_ptr, B_ptr, Pvalues_ptr, Pcolumn_indices_ptr,\n cf_map_ptr);\n cudaCheckError();\n}\n\ntemplate< class T_Config>\nvoid Distance1_InterpolatorBase::generateInterpolationMatrix(Matrix &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix &P)\n{\n P.set_initialized(0);\n\n if (A.get_block_size() == 1)\n {\n generateInterpolationMatrix_1x1(A, cf_map, s_con, scratch, P);\n }\n else\n {\n FatalError(\"Unsupported block size for distance1 interpolator\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n P.computeDiagonal();\n P.set_initialized(1);\n}\n\n#define AMGX_CASE_LINE(CASE) template class Distance1_InterpolatorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Distance1_Interpolator::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace amgx\n{\n\nnamespace detail\n{\n\ntemplate< class CsrMatrixType >\nstruct find_row_diagonal\n{\n// typedef typename CsrMatrixType::values_array_type::value_type value_type;\n typedef typename CsrMatrixType::value_type value_type;\n const CsrMatrixType &my_A;\n find_row_diagonal( const CsrMatrixType &A ) : my_A( A ) {}\n\n // Operator to extract the diagonal value.\n value_type operator()( int irow ) const\n {\n for ( int i = my_A.row_offsets[irow], end = my_A.row_offsets[irow + 1] ; i < end ; ++i )\n if ( my_A.col_indices[i] == irow )\n {\n return my_A.values[i];\n }\n\n // assert( false );\n return value_type( 0 );\n }\n};\n\ntemplate< class CsrMatrixType >\nstruct count_nnz_per_row\n{\n const CsrMatrixType &my_A;\n static const AMGX_VecPrecision vecPrec = CsrMatrixType::TConfig::vecPrec;\n static const AMGX_MatPrecision matPrec = CsrMatrixType::TConfig::matPrec;\n static const AMGX_IndPrecision indPrec = CsrMatrixType::TConfig::indPrec;\n typedef Vector > BVector;\n typedef Vector > IntVector;\n\n const BVector &my_are_sc;\n const IntVector &my_cf_map;\n count_nnz_per_row( const CsrMatrixType &A,\n const BVector &are_sc,\n const IntVector &cf_map ) :\n my_A( A ),\n my_are_sc( are_sc ),\n my_cf_map( cf_map )\n {}\n\n // Operator to extract the diagonal value.\n int operator()( int i ) const\n {\n if ( my_cf_map[i] >= 0 ) // If it's a coarse node, only one NZ element.\n {\n return 1;\n }\n\n // Count the number of strongly connected coarse nodes.\n int count = 0;\n\n for ( int itj = my_A.row_offsets[i] ; itj < my_A.row_offsets[i + 1] ; ++itj )\n {\n int j = my_A.col_indices[itj];\n\n if ( my_cf_map[j] >= 0 && my_are_sc[itj] )\n {\n count++;\n }\n }\n\n return count;\n }\n};\n\ntemplate< class CsrMatrixType >\nstruct compute_weights\n{\n typedef typename CsrMatrixType::value_type value_type;\n typedef typename CsrMatrixType::MVector VVector;\n static const AMGX_VecPrecision vecPrec = CsrMatrixType::TConfig::vecPrec;\n static const AMGX_MatPrecision matPrec = CsrMatrixType::TConfig::matPrec;\n static const AMGX_IndPrecision indPrec = CsrMatrixType::TConfig::indPrec;\n typedef Vector > BVector;\n typedef Vector > IntVector;\n const CsrMatrixType &my_A;\n const VVector &my_diag;\n const BVector &my_are_sc;\n const IntVector &my_cf_map;\n IntVector *my_edges_markers;\n CsrMatrixType &my_P;\n\n // Constructor.\n compute_weights( const CsrMatrixType &A,\n const VVector &diag,\n const BVector &are_sc,\n const IntVector &cf_map,\n IntVector *edges_markers,\n CsrMatrixType &P ) :\n my_A( A ),\n my_diag( diag ),\n my_are_sc( are_sc ),\n my_cf_map( cf_map ),\n my_edges_markers( edges_markers ),\n my_P( P )\n { }\n\n // Get the thread index. Always 0 if OpenMP is off.\n inline\n int get_thread_id( ) const\n {\n return 0;\n }\n\n // Operator to find the weights.\n void operator()( int i )\n {\n // Skip coarse nodes.\n int nnz = my_P.row_offsets[i];\n int id = my_cf_map[i];\n\n if ( id >= 0 )\n {\n my_P.col_indices[nnz] = id;\n my_P.values[nnz] = value_type( 1 );\n return;\n }\n\n assert( my_A.row_offsets[i] < my_A.row_offsets[i + 1] );\n const int strong_marker = -2;\n\n // Initialize the F-C connections to 0 and mark F-F connections.\n for ( int itj = my_A.row_offsets[i] ; itj < my_A.row_offsets[i + 1] ; ++itj )\n {\n if ( !my_are_sc[itj] )\n {\n continue;\n }\n\n int j = my_A.col_indices[itj];\n\n if ( my_cf_map[j] == FINE ) // It's a F-F connection, mark it.\n {\n assert( my_cf_map[j] < 0 );\n my_edges_markers[get_thread_id( )][j] = strong_marker;\n continue;\n }\n\n assert( nnz < my_P.row_offsets[i + 1] );\n my_edges_markers[get_thread_id( )][j] = nnz; // It's a F-C connection, mark the coarse end point.\n my_P.col_indices[nnz] = my_cf_map[j];\n my_P.values[nnz] = value_type( 0 );\n nnz++;\n }\n\n // Make sure we exhausted all the indices.\n assert( nnz == my_P.row_offsets[i + 1] );\n // Reset NNZ variable.\n nnz = my_P.row_offsets[i];\n // The diagonal element a_ii.\n value_type diag = my_diag[i];\n\n // Loop over the neighbours of the node n_i to add contributions.\n for ( int itj = my_A.row_offsets[i] ; itj < my_A.row_offsets[i + 1] ; ++itj )\n {\n int j = my_A.col_indices[itj];\n\n if ( i == j ) // Skip the diagonal.\n {\n continue;\n }\n\n // 1st case: j is a coarse node and it strongly influences i, so we accumulate a_ij.\n if ( my_edges_markers[get_thread_id( )][j] >= nnz )\n {\n int itk = my_edges_markers[get_thread_id( )][j];\n my_P.values[itk] += my_A.values[itj];\n continue;\n }\n // 2nd case: j is a fine node and strongly influences i. We distribute a_ij to the coarse\n // nodes that strongly influence i (not the diagonal).\n else if ( my_edges_markers[get_thread_id( )][j] == strong_marker )\n {\n value_type sum( 0 );\n int diag_sign = std::signbit( my_diag[i] );\n\n for ( int itk = my_A.row_offsets[j] ; itk < my_A.row_offsets[j + 1] ; ++itk )\n {\n int k = my_A.col_indices[itk];\n\n if ( my_edges_markers[get_thread_id( )][k] >= nnz && std::signbit( my_A.values[itk] ) != diag_sign )\n {\n sum += my_A.values[itk];\n }\n }\n\n if ( sum == value_type( 0 ) )\n {\n diag += my_A.values[itj];\n continue;\n }\n\n value_type d = my_A.values[itj] / sum;\n\n for ( int itk = my_A.row_offsets[j] ; itk < my_A.row_offsets[j + 1] ; ++itk )\n {\n int it = my_edges_markers[get_thread_id( )][my_A.col_indices[itk]];\n\n if ( it >= nnz && std::signbit( my_A.values[itk] ) != diag_sign )\n {\n my_P.values[it] += d * my_A.values[itk];\n }\n }\n }\n else if ( my_cf_map[j] == FINE )\n {\n diag += my_A.values[itj];\n }\n }\n\n if ( diag == value_type( 0 ) )\n {\n diag = my_diag[i];\n\n if ( diag == value_type( 0 ) ) // REALLY ?\n {\n diag = value_type( 1 );\n }\n }\n\n assert( diag != value_type( 0 ) );\n diag = value_type( -1.0 ) / diag;\n\n for ( ; nnz < my_P.row_offsets[i + 1] ; ++nnz )\n {\n my_P.values[nnz] *= diag;\n }\n\n // strong_marker--;\n }\n};\n\n}\n\n/*************************************************************************\n* create the interpolation matrix P\n************************************************************************/\ntemplate \nvoid Distance1_Interpolator >::generateInterpolationMatrix_1x1(\n Matrix_h &A,\n IntVector &cf_map,\n BVector &are_sc,\n IntVector &wk,\n Matrix_h &P)\n{\n // The diagonal of A.\n VVector diag( A.get_num_rows() );\n thrust_wrapper::transform( amgx::thrust::make_counting_iterator( 0 ),\n amgx::thrust::make_counting_iterator( A.get_num_rows() ),\n diag.begin( ),\n detail::find_row_diagonal( A ) );\n cudaCheckError();\n // Count the number of coarse points.\n int n_coarse = 0;\n\n for ( int i = 0 ; i < A.get_num_rows() ; ++i )\n if ( cf_map[i] >= 0) { n_coarse++; }\n\n // Count the number of non-zero elements per row of P.\n IntVector nnz_per_row( A.get_num_rows() + 1 );\n amgx::thrust::transform( amgx::thrust::make_counting_iterator( 0 ),\n amgx::thrust::make_counting_iterator( A.get_num_rows() ),\n nnz_per_row.begin( ),\n detail::count_nnz_per_row( A, are_sc, cf_map ) );\n cudaCheckError();\n nnz_per_row[A.get_num_rows()] = 0;\n // Find the total number of non-zeroes. \n // TODO: merge with the exclusive scan.\n int nnz = amgx::thrust::reduce( nnz_per_row.begin( ), nnz_per_row.end( ) );\n cudaCheckError();\n // Resize P so that there is enough space to store the non-zero elements.\n P.addProps(CSR);\n P.resize( A.get_num_rows(), n_coarse, nnz );\n // Compute row offsets of P.\n thrust_wrapper::exclusive_scan( nnz_per_row.begin( ),\n nnz_per_row.end( ),\n P.row_offsets.begin( ) );\n cudaCheckError();\n // For each row we compute the weights.\n IntVector edges_markers_on_stack( A.get_num_rows(), -1 );\n IntVector *edges_markers = &edges_markers_on_stack;\n detail::compute_weights compute_fct( A, diag, are_sc, cf_map, edges_markers, P );\n\n amgx::thrust::for_each( amgx::thrust::host,\n amgx::thrust::make_counting_iterator( 0 ),\n amgx::thrust::make_counting_iterator( A.get_num_rows() ),\n compute_fct );\n cudaCheckError();\n}\n\n/*************************************************************************\n* create the interpolation matrix P on GPU\n************************************************************************/\n\nenum SETS {SET_STRONG_COARSE = 1,\n SET_WEAK_COARSE = 2,\n SET_STRONG_FINE = 4,\n SET_STRONG_FINE_NO_COMMON = 8,\n SET_WEAK_FINE = 16,\n SET_DIAGONAL = 32\n };\n\n/*\n * grab the diagonal of the matrix\n */\n/*template \n__global__\nvoid getDiagonalKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, const IndexType numRows, ValueType *diagonal)\n{\n const int tIdx = threadIdx.x + blockDim.x*blockIdx.x;\n\n if (tIdx >= numRows) return;\n\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx+1]-offset;\n\n for (int j=offset; j < offset+numj; j++)\n {\n int jcol = column_indices[j];\n if (tIdx == jcol)\n {\n diagonal[tIdx] = values[j];\n }\n }\n}*/\n\n/*\n * Count the non-zeros in final output\n */\ntemplate \n__global__\nvoid numNonZerosVecKernel(const IndexType *offsets, const IndexType *column_indices,\n const IndexType numRows, int *cf_map, const bool *s_con,\n int *nonZerosVec)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n int numOnRow = 0;\n\n // count strong coarse influences\n if (cf_map[tIdx] == FINE)\n {\n for (int j = offset; j < offset + numj; j++)\n {\n if (!s_con[j]) { continue; } // if weakly connected, ignore\n\n int jcol = column_indices[j];\n\n if (cf_map[jcol] >= 0) // if coarse\n {\n numOnRow++;\n }\n }\n }\n else\n {\n numOnRow++;\n }\n\n nonZerosVec[tIdx] = numOnRow;\n }\n}\n\n/*\n * Calculate D_i vector\n * - Weakly influencing and strong fine with no common C-points to be added to diagonal term\n */\ntemplate \n__global__\nvoid calculateDKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, int num_rows, ValueType *diagonal,\n int *set_fields, ValueType *D)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < num_rows; tIdx += blockDim.x * gridDim.x)\n {\n const int offset = offsets[tIdx];\n const int numk = offsets[tIdx + 1] - offset;\n ValueType sum = 0;\n int mask = SET_WEAK_FINE;\n\n for (int k = offset; k < offset + numk; k++)\n {\n if (set_fields[k] & mask)\n {\n sum += values[k];\n }\n }\n\n D[tIdx] += sum;\n }\n}\n\n/*\n * Calculate the k dependent term in w_i,j\n */\ntemplate \n__global__\nvoid calculateBKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, int num_rows, ValueType *diagonal, ValueType *diag_inc,\n int *set_fields, int *cf_map, int *nonZerosVec, ValueType *B)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < num_rows; tIdx += blockDim.x * gridDim.x)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n int localNonZeros = 0;\n const int nonZerosOffset = nonZerosVec[tIdx];\n ValueType sum = 0.0, top = 0.0, bottom_sum = 0.0, a_ik;\n\n if (cf_map[tIdx] >= 0) // If I am a coarse row\n {\n B[nonZerosOffset] = 1;\n return;\n }\n\n diag_inc[tIdx] = 0;\n int flag = 0;\n int first_j_loop = 0;\n ValueType tol = 1e-10;\n\n // otherwise\n for (int j = offset; j < offset + numj; j++)\n {\n sum = 0.0;\n bottom_sum = 0.0;\n\n if (set_fields[j] & SET_STRONG_COARSE) // we actually care about this j\n {\n int jcol = column_indices[j]; // this is the only part of j I need / care about.\n\n if (flag == 0)\n {\n first_j_loop = 1;\n flag = 1;\n }\n else\n {\n first_j_loop = 0;\n }\n\n // run through the columns of row i, summing over all k\n for (int k = offset; k < offset + numj; k++)\n {\n top = 0.0;\n bottom_sum = 0.0;\n\n if ( (set_fields[k] & SET_STRONG_FINE) || (set_fields[k] & SET_STRONG_FINE_NO_COMMON)) // k we're interested in\n {\n // part of the desired set - get column, diagonal and a_ik\n int kcol = column_indices[k];\n a_ik = values[k]; // get a_i,k term\n // now find a_k,j\n ValueType diag = diagonal[kcol];\n int sgn = diag < 0.0 ? -1 : 1;\n\n // now I need to find the j-th column on the k-th row. another loop (sigh)\n for (int n = offsets[kcol]; n < offsets[kcol + 1]; n++)\n {\n int ncol = column_indices[n];\n\n if (ncol == jcol) // this is the entry I need\n {\n if (sgn * values[n] < 0)\n {\n // add to the sum\n top = a_ik * values[n]; // discard diagonal\n }\n }\n }\n\n // calculate the bottom sum - loop through columns of i\n for (int m = offset; m < offset + numj; m++)\n {\n if (set_fields[m] & SET_STRONG_COARSE)\n {\n // this is a point I care about\n int mcol = column_indices[m];\n\n // now find the m-th column of the k-th row and add to sum\n for (int n = offsets[kcol]; n < offsets[kcol + 1]; n++)\n {\n int ncol = column_indices[n];\n\n if (ncol == mcol)\n {\n if (sgn * values[n] < 0)\n {\n bottom_sum += values[n];\n }\n }\n }\n }\n }\n\n if (abs(bottom_sum) < tol)\n {\n if (first_j_loop == 1)\n {\n //add it to the diagonal....\n diag_inc[tIdx] += values[k];\n }\n }\n else\n {\n sum += top / bottom_sum;\n }\n }\n }\n\n B[nonZerosOffset + localNonZeros] = sum;\n localNonZeros++;\n }\n }\n }\n}\n\n/*\n * count the number of strong coarse points per row\n * - was used for degugging, might be useful in the future\n */\ntemplate \n__global__\nvoid coarsePerRow(const IndexType *offsets, int num_rows, int *set_fields, int *output)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < num_rows; tIdx += blockDim.x * gridDim.x)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n int localCoarse = 0;\n\n for (int j = offset; j < offset + numj; j++)\n {\n if (set_fields[j] & SET_STRONG_COARSE)\n {\n localCoarse++;\n }\n }\n\n output[tIdx] = localCoarse;\n }\n}\n\n/*\n * Assemble final w_ij terms\n */\ntemplate \n__global__\nvoid calculateWKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, int num_rows, ValueType *diagonal,\n int *set_fields, int *nonZerosVec, ValueType *D, ValueType *B,\n ValueType *w, IndexType *Pcolumn_indices, int *cf_map)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < num_rows; tIdx += blockDim.x * gridDim.x)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n int localNonZeros = 0;\n const int nonZerosOffset = nonZerosVec[tIdx];\n ValueType diag = diagonal[tIdx];\n\n if (cf_map[tIdx] >= 0)\n {\n w[nonZerosOffset] = 1.0;\n Pcolumn_indices[nonZerosOffset] = cf_map[tIdx];\n continue;\n }\n\n for (int j = offset; j < offset + numj; j++)\n {\n if (set_fields[j] & SET_STRONG_COARSE)\n {\n ValueType aij = values[j];\n ValueType Bij = B[nonZerosOffset + localNonZeros];\n ValueType Di = D[tIdx];\n Pcolumn_indices[nonZerosOffset + localNonZeros] = cf_map[column_indices[j]];\n double bottom_sum = (fabs(diag + Di) < 1e-10) ? 1. : diag + Di;\n w[nonZerosOffset + localNonZeros] = -1.0 / bottom_sum * (aij + Bij);\n localNonZeros++;\n }\n }\n }\n}\n\ntemplate \n__global__\nvoid markCoarseEdgesKernel(const IndexType *offsets, const IndexType *column_indices,\n const size_t numRows, const int *cf_map, const bool *s_con,\n bool *coarseMark)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x)\n {\n const int rowBegin = offsets[tIdx];\n const int rowEnd = offsets[tIdx + 1];\n\n // loop over columns\n for (int j = rowBegin; j < rowEnd; j++)\n {\n int jcol = column_indices[j];\n\n if (tIdx == jcol) { continue; } // skip the diagonal\n\n if (cf_map[jcol] >= 0)\n {\n // mark this point\n coarseMark[j] = true;\n }\n }\n }\n}\n\ntemplate \n__device__\nbool sparseIntersectDevice(const IndexType *column_indices, ValueType *values,\n int begin1, int end1, int begin2, int end2)\n{\n int i1 = begin1;\n int i2 = begin2;\n\n // check arrays non zero length\n if (begin1 >= end1 || begin2 >= end2) { return false; }\n\n bool r = false;\n\n while ( true )\n {\n int col1 = column_indices[i1];\n int col2 = column_indices[i2];\n\n if (col1 == col2) // this is a match\n {\n if (values[i1] && values[i2])\n {\n r = true;\n break;\n }\n\n i1++;\n i2++;\n\n if ( i1 >= end1 || i2 >= end2 ) { break; }\n }\n else if (col1 > col2)\n {\n i2++;\n\n if (i2 >= end2) { break; }\n }\n else\n {\n i1++;\n\n if (i1 >= end1) { break; }\n }\n }\n\n return r;\n}\n\ntemplate \n__global__\nvoid categoriseEdgesKernel(const IndexType *offsets, const IndexType *column_indices,\n const size_t numRows, const int *cf_map, const bool *s_con,\n bool *coarse_mark, int *set_fields)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x)\n {\n // set_fields[tIdx] = 1;\n if (cf_map[tIdx] >= 0) { continue; } // we don't need to categorise C-? edges\n\n const int rowBegin = offsets[tIdx];\n const int rowEnd = offsets[tIdx + 1];\n\n // we know I'm a fine row now - loop over non-zeros\n for (int j = rowBegin; j < rowEnd; j++)\n {\n int jcol = column_indices[j];\n\n if (jcol == tIdx) { continue; }\n\n // check F-C edges\n if (cf_map[jcol] >= 0)\n {\n if (s_con[j])\n {\n set_fields[j] |= SET_STRONG_COARSE;\n }\n else\n {\n set_fields[j] |= SET_WEAK_COARSE;\n }\n }\n else // F-F edges\n {\n if (!s_con[j])\n {\n set_fields[j] |= SET_WEAK_FINE;\n }\n else // check if there exists a common C-point\n {\n // intersect i and jcol\n bool common_C = true;\n common_C = sparseIntersectDevice(column_indices, coarse_mark, rowBegin, rowEnd,\n offsets[jcol], offsets[jcol + 1]);\n\n if (common_C)\n {\n set_fields[j] |= SET_STRONG_FINE;\n }\n else\n {\n set_fields[j] |= SET_STRONG_FINE_NO_COMMON;\n }\n }\n }\n }\n }\n}\n\ntemplate \nvoid Distance1_Interpolator >::setSetsGPU(const Matrix_d &A, const BVector &s_con,\n const IntVector &cf_map,\n IntVector &set_fields)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n const int blocksize = 64;\n const int numBlocks = std::min( AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() / blocksize + 1) );\n // raw pointers for matrix structure\n const IndexType *offsets = A.row_offsets.raw();\n const IndexType *column_indices = A.col_indices.raw();\n // raw pointers for coarse/fine map and strength of connection array\n const bool *s_con_ptr = s_con.raw();\n const int *cf_map_ptr = cf_map.raw();\n int *set_fields_ptr = set_fields.raw();\n // memory for CoarseMark array\n BVector coarseMark(A.get_num_nz(), false);\n bool *coarseMark_ptr = coarseMark.raw();\n markCoarseEdgesKernel <<< numBlocks, blocksize>>>(offsets, column_indices, A.get_num_rows(),\n cf_map_ptr, s_con_ptr, coarseMark_ptr);\n cudaCheckError();\n // now we have the strong coarse edges marked, categorise all edges\n categoriseEdgesKernel <<< numBlocks, blocksize>>>(offsets, column_indices, A.get_num_rows(),\n cf_map_ptr, s_con_ptr, coarseMark_ptr,\n set_fields_ptr);\n cudaCheckError();\n}\n\ntemplate \nvoid Distance1_Interpolator >::generateInterpolationMatrix_1x1(Matrix_d &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix_d &P)\n{\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // Assemble all the raw pointers needed\n // Matrix A\n const IndexType *Aoffsets_ptr = A.row_offsets.raw();\n const IndexType *Acolumn_indices_ptr = A.col_indices.raw();\n const ValueType *Avalues_ptr = A.values.raw();\n // Coarse/Fine map, strength of connection, scratch\n int *cf_map_ptr = cf_map.raw();\n const bool *s_con_ptr = s_con.raw();\n int *scratch_ptr = scratch.raw();\n // temporary vectors - diagonal and non-zero offsets\n VVector diag(A.get_num_rows(), 0.0);\n IntVector nonZerosVec(A.get_num_rows() + 1, 0);\n ValueType *diag_ptr = diag.raw();\n int *nonZerosVec_ptr = nonZerosVec.raw();\n // NB - will get pointers for P *after* it has been initialised\n // choose blocksize. Using 1 thread / row for now\n const int blocksize = 64;\n const int numBlocks = std::min (AMGX_GRID_MAX_SIZE, (int) (A.get_num_rows() / blocksize + 1));\n const IndexType numRows = (IndexType) A.get_num_rows();\n // extract the diagonal of A\n find_diag_kernel_indexed_dia <<< numBlocks, blocksize>>>(A.get_num_rows(),\n A.diag.raw(),\n A.values.raw(),\n diag_ptr);\n cudaCheckError();\n /* getDiagonalKernel<<>>(Aoffsets_ptr,Acolumn_indices_ptr,Avalues_ptr,\n numRows,diag_ptr);*/\n cudaCheckError();\n // count the number of non-zeros in final output per row\n numNonZerosVecKernel <<< numBlocks, blocksize>>>(Aoffsets_ptr, Acolumn_indices_ptr, numRows,\n cf_map_ptr, s_con_ptr, nonZerosVec_ptr);\n cudaCheckError();\n // now I have the non-zeros per row for matrix P, count the non-zeros on each row\n // to get total NNZ\n // sum non-zeros per row\n int NNZidx = thrust_wrapper::reduce(nonZerosVec.begin(), nonZerosVec.end());\n cudaCheckError();\n thrust_wrapper::exclusive_scan(nonZerosVec.begin(), nonZerosVec.end(), nonZerosVec.begin());\n cudaCheckError();\n nonZerosVec[A.get_num_rows()] = NNZidx;\n // generate sets on the device\n IntVector set_fields(A.get_num_nz(), 0);\n setSetsGPU(A, s_con, cf_map, set_fields);\n int *set_field_ptr = set_fields.raw();\n // count the coarse points first\n // (no need to worry about >4B coarse points)\n int coarsePoints = (int) amgx::thrust::count_if(cf_map.begin(), cf_map.end(), is_non_neg());\n cudaCheckError();\n // I have all the information needed - create P and get the pointers to its data\n P.addProps(CSR);\n P.resize(A.get_num_rows(), coarsePoints, NNZidx, 1);\n IndexType *Pcolumn_indices_ptr = P.col_indices.raw();\n ValueType *Pvalues_ptr = P.values.raw();\n // now P is initialised, get row offsets from nonZerosVec with an exclusive prefix sum\n // set final value\n amgx::thrust::copy(nonZerosVec.begin(), nonZerosVec.end(), P.row_offsets.begin());\n cudaCheckError();\n IndexType *Poffsets_ptr = P.row_offsets.raw();\n // assign temp memory\n VVector B(NNZidx, 0.0);\n VVector D(A.get_num_rows(), 0.0);\n VVector W(NNZidx, 0.0);\n // get the raw pointers for temp memory\n ValueType *D_ptr = D.raw();\n ValueType *B_ptr = B.raw();\n ValueType *W_ptr = W.raw();\n // call each kernel in turn\n calculateBKernel <<< numBlocks, blocksize>>>(Aoffsets_ptr, Acolumn_indices_ptr, Avalues_ptr,\n (int) A.get_num_rows(), diag_ptr, D_ptr, set_field_ptr, cf_map_ptr,\n nonZerosVec_ptr, B_ptr);\n cudaCheckError();\n calculateDKernel <<< numBlocks, blocksize>>>(Aoffsets_ptr, Acolumn_indices_ptr, Avalues_ptr,\n (int) A.get_num_rows(), diag_ptr, set_field_ptr, D_ptr);\n cudaCheckError();\n calculateWKernel <<< numBlocks, blocksize>>>(Aoffsets_ptr, Acolumn_indices_ptr, Avalues_ptr,\n (int) A.get_num_rows(), diag_ptr, set_field_ptr, nonZerosVec_ptr,\n D_ptr, B_ptr, Pvalues_ptr, Pcolumn_indices_ptr,\n cf_map_ptr);\n cudaCheckError();\n}\n\ntemplate< class T_Config>\nvoid Distance1_InterpolatorBase::generateInterpolationMatrix(Matrix &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix &P)\n{\n P.set_initialized(0);\n\n if (A.get_block_size() == 1)\n {\n generateInterpolationMatrix_1x1(A, cf_map, s_con, scratch, P);\n }\n else\n {\n FatalError(\"Unsupported block size for distance1 interpolator\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n P.computeDiagonal();\n P.set_initialized(1);\n}\n\n#define AMGX_CASE_LINE(CASE) template class Distance1_InterpolatorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Distance1_Interpolator::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n\n", "source": "/kaggle/working/amgx/src/classical/interpolators/distance1.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\nusing std::set;\nusing std::vector;\n\n/*\n * remove all duplicates from a vector\n */\ntemplate \nvoid removeDuplicates(vector &vec)\n{\n std::sort(vec.begin(), vec.end());\n vec.erase(std::unique(vec.begin(), vec.end()), vec.end());\n}\n\n/*\n * find the value at a particular column of a row if it exists\n */\ntemplate \ntypename VVector::value_type findValue(const IVector &columns,\n const VVector &values,\n const int row_begin, const int row_end, const int column_needed)\n{\n typename VVector::value_type v = 0;\n\n // NAIVE\n for (int k = row_begin; k < row_end; k++)\n {\n int kcol = columns[k];\n\n if (kcol == column_needed)\n {\n v = values[k];\n }\n }\n\n return v;\n}\n\n\ntemplate \nvoid __global__ createCfMapGlobal(const IndexType *cf_map, int64_t *cf_map_global, const int64_t my_part_offset, const IndexType num_owned_fine_pts)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_owned_fine_pts ; tidx += blockDim.x * gridDim.x)\n {\n // Only renumber the interior points\n if (cf_map[tidx] >= 0 && tidx < num_owned_fine_pts)\n {\n // Simply shift\n cf_map_global[tidx] = (int64_t) cf_map[tidx] /* my_local_id */ + my_part_offset;\n }\n else\n {\n cf_map_global[tidx] = -1;\n }\n }\n}\n\ntemplate \n__global__ void fill_P_global_col_indices_kernel(index_type *row_offsets, index_type *C_hat_start, index_type *C_hat_end, index_type *C_hat, int64_t *cf_map_global, int64_t *global_col_indices, index_type num_rows)\n{\n int coopIdx = threadIdx.x % coop;\n\n for (int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop; row < num_rows ; row += blockDim.x * gridDim.x / coop)\n {\n int64_t coarse_fine_id = cf_map_global[row];\n int j = row_offsets[row] + coopIdx;\n\n // coarse row\n if (coarse_fine_id >= 0)\n {\n if (coopIdx == 0)\n {\n global_col_indices[j] = coarse_fine_id;\n }\n\n continue;\n }\n\n // strong fine row\n if (coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // fine row\n for (int i = C_hat_start[row] + coopIdx; i < C_hat_end[row]; i += coop, j += coop)\n {\n index_type fine_col = C_hat[i];\n int64_t global_col = cf_map_global[fine_col];\n global_col_indices[j] = global_col;\n }\n }\n}\n\n\n/*\n * calculate \\tilde{a}_{ii} from eqn. 4.11 in [4]\n */\ntemplate \nvoid Distance2_Interpolator >::calculateD( const Matrix_h &A, IVector &cf_map, BVector &s_con,\n vector *C_hat, VVector &diag, VVector &D,\n set *weak_lists, VVector &innerSum,\n IVector &innerSumOffsets)\n{\n // only do the non-weak part for now\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n int outerOffset = innerSumOffsets[i], innerOffset = 0;\n ValueType sum = 0;\n\n // loop over all non-zeros\n for (int k = A.row_offsets[i]; k < A.row_offsets[i + 1]; k++)\n {\n int kcol = A.col_indices[k];\n\n if (cf_map[kcol] == FINE && s_con[k]) // strongly connected fine, ignore diagonal\n {\n int sgn = diag[kcol] < 0 ? -1 : 1;\n // find a_ki\n ValueType temp = findValue(A.col_indices, A.values, A.row_offsets[kcol], A.row_offsets[kcol + 1], i);\n ValueType a_ki = (sgn * temp < 0) * temp;\n sum += a_ki * innerSum[outerOffset + innerOffset];\n innerOffset++;\n }\n }\n\n // now calculate the weak part\n set difference;\n\n for (set::iterator it = weak_lists[i].begin(); it != weak_lists[i].end(); ++it)\n {\n difference.insert(*it);\n }\n\n for (int l = 0; l < C_hat[i].size(); l++)\n {\n difference.erase(C_hat[i][l]);\n }\n\n ValueType weak = 0;\n\n for (set::iterator it = difference.begin(); it != difference.end(); ++it)\n {\n weak += findValue(A.col_indices, A.values, A.row_offsets[i], A.row_offsets[i + 1], *it);\n }\n\n D[i] = diag[i] + weak + sum;\n }\n}\n\n/*\n * Generate the lower sum from eqns. 4.10 and 4.11 of [4]\n * strongly resembles a sparse dot product\n */\ntemplate \ntypename VVector::value_type sparseSum(const IVector &columns,\n const VVector &values,\n const int row_begin, const int row_end, const std::vector &C_hat_row, const int sgn)\n{\n int i1 = row_begin, i2 = 0;\n\n // check lengths are non-zero\n if (row_begin >= row_end || C_hat_row.size() == 0) { return 0; }\n\n typename VVector::value_type sum = 0;\n\n // if column exists in both C_hat_row and row of matrix, add matrix value\n while (true)\n {\n int col1 = columns[i1];\n int col2 = C_hat_row[i2];\n\n if (col1 == col2) // match\n {\n if (sgn * values[i1] < 0)\n {\n sum += values[i1];\n }\n\n i1++;\n i2++;\n\n if (i1 >= row_end || i2 >= C_hat_row.size()) { break; }\n }\n else if (col1 > col2)\n {\n // instead of linear increase, binary search\n /* int min = i2++, max = C_hat_row.size();\n while(min <= max)\n {\n int mid = (min+max)/2;\n col2 = C_hat_row[mid];\n if (col1 > col2)\n min = mid+1;\n else if (col1 < col2)\n max = mid-1;\n else\n {\n i2 = mid;\n break;\n }\n }*/\n i2++;\n\n if (i2 >= C_hat_row.size()) { break; }\n }\n else\n {\n i1++;\n\n if (i1 >= row_end) { break; }\n }\n }\n\n return sum;\n}\n\n/*\n * Generate the repeatedly reused inner sum from eqns 4.10, 4.11 in [4]\n * \\sum_{k\\in F^s_i} a_{ik}\\frac{1}{\\sum_{l\\in\\hat{C}_i\\cup\\{i\\}} \\bar{a}_{kl}}\n */\ntemplate \nvoid Distance2_Interpolator >::generateInnerSum(Matrix_h &A, const IntVector &cf_map,\n const BVector &s_con,\n vector *C_hat, VVector &diag,\n VVector &innerSum,\n IntVector &innerSumOffsets)\n{\n // calculate necessary storage for innerSum\n int connections = 0;\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n int jcol = A.col_indices[j];\n\n if (cf_map[jcol] == FINE && s_con[j] && jcol != i)\n {\n connections++;\n }\n }\n }\n\n innerSum.resize(connections);\n innerSumOffsets.resize(A.get_num_rows() + 1);\n innerSumOffsets[0] = 0;\n int innerOffset = 0;\n\n // generate the values\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n for (int k = A.row_offsets[i]; k < A.row_offsets[i + 1]; k++)\n {\n int kcol = A.col_indices[k];\n\n if (cf_map[kcol] == FINE && s_con[k] && kcol != i)\n {\n ValueType bottom_sum = 0;\n ValueType a_ik = A.values[k];\n int sgn = diag[kcol] < 0 ? -1 : 1;\n // generate the bottom sum - looks similar to sparse dot product\n bottom_sum = sparseSum(A.col_indices, A.values, A.row_offsets[kcol], A.row_offsets[kcol + 1],\n C_hat[i], sgn);\n // add \\union i\n ValueType iTerm = findValue(A.col_indices, A.values, A.row_offsets[kcol], A.row_offsets[kcol + 1], i);\n bottom_sum += (sgn * iTerm < 0) * iTerm;\n\n if (bottom_sum == 0 && cf_map[i] == FINE)\n {\n char buf[255];\n sprintf(buf, \"DIVISION BY ZERO: row: %d, col: %d\\n\", i, kcol);\n amgx_output(buf, strlen(buf));\n }\n\n innerSum[innerOffset] = a_ik / bottom_sum;\n innerOffset++;\n }\n }\n\n innerSumOffsets[i + 1] = innerOffset;\n }\n}\n\n/*************************************************************************\n* create the interpolation matrix P\n************************************************************************/\ntemplate \nvoid Distance2_Interpolator >::generateInterpolationMatrix_1x1(Matrix_h &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix_h &P)\n{\n if (!A.is_matrix_singleGPU())\n {\n FatalError(\"Distributed classical AMG not implemented on host\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename Matrix_h::value_type ValueType;\n // implementation begins here\n set *coarse_lists = new set[A.get_num_rows()];\n set *weak_lists = new set[A.get_num_rows()];\n vector *C_hat = new vector[A.get_num_rows()]; // hold full \\hat{C} set for each row\n set *C_hat_set = new set[A.get_num_rows()];\n set::iterator set_it;\n VVector diag(A.get_num_rows(), 0);\n\n // grab the diagonal\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n if (A.col_indices[j] == i)\n {\n diag[i] = A.values[j];\n }\n }\n }\n\n // generate coarse sets\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n // loop over non-zero elements\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n int jcol = A.col_indices[j];\n\n if (jcol == i) { continue; } // skip the diagonal\n\n // Coarse and strongly connected\n if (cf_map[jcol] >= 0 && s_con[j])\n {\n coarse_lists[i].insert(jcol);\n }\n else if (!s_con[j] && jcol != i && cf_map[jcol] != STRONG_FINE)\n {\n weak_lists[i].insert(jcol);\n }\n }\n }\n\n int numWTotal = 0;\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n numWTotal += (int) weak_lists[i].size();\n }\n\n typedef typename set::iterator iSetIter;\n\n // generate \\hat{C} for each row, i\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n C_hat[i].clear();\n\n if (cf_map[i] == FINE)\n {\n // copy coarse_lists[i] to C_hat_set[i]\n for (iSetIter it = coarse_lists[i].begin(); it != coarse_lists[i].end(); ++it)\n {\n C_hat_set[i].insert(*it);\n }\n\n // loop over non-zeros\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n int jcol = A.col_indices[j];\n\n if (jcol == i) { continue; } // don't want to take into account myself\n\n // if this is a strong-fine connection, union of the sets\n if (cf_map[jcol] == FINE && s_con[j])\n {\n for (iSetIter it = coarse_lists[jcol].begin(); it != coarse_lists[jcol].end(); ++it)\n {\n C_hat_set[i].insert(*it);\n }\n }\n }\n\n removeDuplicates(C_hat[i]);\n }\n\n // copy C_hat_set[i] into C_hat[i]\n for (iSetIter it = C_hat_set[i].begin(); it != C_hat_set[i].end(); ++it)\n {\n C_hat[i].push_back(*it);\n }\n }\n\n // check that \\hat{C_i} for fine points does not contain i\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n if (cf_map[i] == FINE)\n {\n vector::iterator it = std::find(C_hat[i].begin(), C_hat[i].end(), i);\n {\n assert(it == C_hat[i].end());\n }\n }\n }\n\n // generate the inner sum for all i, k\n VVector innerSum;\n IntVector innerSumOffsets;\n generateInnerSum(A, cf_map, s_con, C_hat, diag, innerSum, innerSumOffsets);\n // calculate the ~a_ii term\n VVector D(A.get_num_rows(), 0);\n calculateD(A, cf_map, s_con, C_hat, diag, D, weak_lists, innerSum, innerSumOffsets);\n // delete all storage we have no more need for\n delete [] coarse_lists;\n delete [] weak_lists;\n delete [] C_hat_set;\n // now we have \\hat{C} and the diagonals, we can start calculating the weights\n // remap cf_map (to be like in distance1)\n int coarsePoints = 0;\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n // if coarse\n if (cf_map[i] >= 0) { coarsePoints++; }\n }\n\n // count the non-zeros - 1 per row if coarse, len(C_hat[i]) otherwise\n IntVector nonZeroOffsets(A.get_num_rows() + 1);\n IntVector nonZerosPerRow(A.get_num_rows());\n int nonZeros = 0;\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n if (cf_map[i] >= 0)\n {\n nonZerosPerRow[i] = 1;\n nonZeros++;\n }\n else if (cf_map[i] == STRONG_FINE)\n {\n nonZerosPerRow[i] = 0;\n }\n else\n {\n nonZerosPerRow[i] = (int) C_hat[i].size();\n nonZeros += (int) C_hat[i].size();\n }\n }\n\n // get the offsets with an exclusive scan\n thrust_wrapper::exclusive_scan(nonZerosPerRow.begin(), nonZerosPerRow.end(), nonZeroOffsets.begin());\n cudaCheckError();\n nonZeroOffsets[A.get_num_rows()] = nonZeros;\n // resize P\n P.addProps(CSR);\n P.resize(A.get_num_rows(), coarsePoints, nonZeros, 1);\n // copy nonzero offsets to the P matrix\n amgx::thrust::copy(nonZeroOffsets.begin(), nonZeroOffsets.end(), P.row_offsets.begin());\n cudaCheckError();\n // main loop\n int outerOffset = 0, innerOffset;\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n int nonZeroOffset = nonZeroOffsets[i];\n int localNonZeros = 0;\n outerOffset = innerSumOffsets[i];\n\n if (cf_map[i] >= 0)\n {\n // will be setting here & continuing\n P.values[nonZeroOffset + localNonZeros] = 1.0;\n P.col_indices[nonZeroOffset + localNonZeros] = cf_map[i];\n continue;\n }\n else if (cf_map[i] != STRONG_FINE) // fine row\n {\n //\n for (int j = 0; j < C_hat[i].size(); j++)\n {\n innerOffset = 0; // set to 'row' offset\n int jcol = C_hat[i][j]; // j can potentially not exist on row i\n\n if (jcol == i) // skip diagonal\n {\n FatalError(\"Error - Diagonal should never be in C_hat[i]\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n ValueType value = 0;\n // loop over i and find a_ij - if not found, doesn't exist, keep as 0\n ValueType a_ij = 0;\n a_ij = findValue(A.col_indices, A.values, A.row_offsets[i], A.row_offsets[i + 1], jcol);\n\n // now loop over and find strong-fine connections\n for (int k = A.row_offsets[i]; k < A.row_offsets[i + 1]; k++)\n {\n int kcol = A.col_indices[k];\n\n if (cf_map[kcol] == FINE && s_con[k]) // Strong-Fine edge, ignore diagonal\n {\n if (kcol == i)\n {\n FatalError(\"Error - point should never to strongly connected to itself\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n // point we care about\n int sgn = diag[kcol] < 0 ? -1 : 1;\n ValueType a_kj = 0;\n // get \\bar{a_kj}\n ValueType temp = findValue(A.col_indices, A.values, A.row_offsets[kcol], A.row_offsets[kcol + 1], jcol);\n a_kj = (sgn * temp < 0) * temp;\n // see if we need to add \\bar{a_ki} in\n vector::iterator it = std::find(C_hat[i].begin(), C_hat[i].end(), i);\n\n if (it != C_hat[i].end())\n {\n FatalError(\"Error - point i should never be in\\\\hat{C_i}\", AMGX_ERR_BAD_PARAMETERS);\n }\n else\n {\n // we now have everything we need - add to sum\n // not valid in non-symmetric matrices\n value += innerSum[outerOffset + innerOffset] * a_kj;\n innerOffset++;\n }\n } // strong fine connection end\n } // k loop (find strong fine connections) end\n\n ValueType diagonal = (D[i]);\n // now set column & value for P\n P.col_indices[nonZeroOffset + localNonZeros] = cf_map[jcol];\n P.values[nonZeroOffset + localNonZeros] = -1.0 / (diagonal) * (a_ij + value);\n localNonZeros++;\n } // end j loop\n } // fine connection loop\n\n outerOffset += innerOffset;\n } // end i loop\n\n delete [] C_hat;\n} // end distance2 interpolator\n\n\n///////////////////////////////////////////////////////////////////////////////\n\n// DEVICE CODE\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#include \n#include // Included inside the namespace to solve name colisions.\n\n__device__ __forceinline__ int get_work( int *queue, int warp_id )\n{\n int offset = -1;\n\n if ( utils::lane_id() == 0 )\n {\n offset = atomicAdd( queue, 1 );\n }\n\n return utils::shfl( offset, 0 );\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nnamespace distance2\n{\n\n#include \n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nestimate_c_hat_size_kernel( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const int *cf_map,\n const bool *s_con,\n int *C_hat_offsets )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // A shared location where threads propose a row of B to load.\n __shared__ volatile int s_b_row_ids[CTA_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += gridDim.x * NUM_WARPS )\n {\n // Skip coarse and strong fine rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id >= 0 || coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // Load A row IDs.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // The number of elements.\n int count(0);\n\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n __syncthreads();\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0, num_rows = __popc(vote) ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [b_col_it, b_col_end).\n int b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n // Load the range of B.\n int b_col_it = A_rows[b_row_id + 0];\n int b_col_end = A_rows[b_row_id + 1];\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any(b_col_it < b_col_end) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != b_row_id;\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_off_diagonal && s_con[b_col_it] && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n }\n }\n }\n\n // Store the number of columns in each row.\n if ( lane_id == 0 )\n {\n C_hat_offsets[a_row_id] = count;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nestimate_c_hat_size_kernel( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const int *cf_map,\n const bool *s_con,\n int *C_hat_offsets )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // A shared location where threads propose a row of B to load.\n __shared__ volatile int s_b_row_ids[CTA_SIZE];\n s_b_row_ids[threadIdx.x] = 0;\n __syncthreads();\n\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += gridDim.x * NUM_WARPS )\n {\n // Skip coarse rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id >= 0 || coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // Load A row IDs.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // The number of elements.\n int count(0);\n\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n utils::syncwarp();\n\n // For each warp, we have up to 32 rows of B to proceed.\n\n int num_rows = __popc(vote);\n for ( int k = 0; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [b_col_it, b_col_end).\n int b_row_id = -1;\n\n if ( is_active_k )\n {\n b_row_id = s_b_row_ids[warp_id * WARP_SIZE + local_k];\n }\n\n // Load the range of B.\n int b_col_it = 0, b_col_end = 0;\n\n if ( is_active_k )\n {\n b_col_it = A_rows[b_row_id + 0];\n b_col_end = A_rows[b_row_id + 1];\n }\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != b_row_id;\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_off_diagonal && s_con[b_col_it] && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n }\n }\n }\n\n // Store the number of columns in each row.\n if ( lane_id == 0 )\n {\n C_hat_offsets[a_row_id] = count;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_c_hat_kernel( int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_end,\n int *__restrict C_hat,\n int *__restrict C_hat_pos,\n int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Shared memory to vote.\n __shared__ volatile int s_b_row_ids[CTA_SIZE];\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Skip coarse rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id >= 0 || coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // Clear the set.\n set.clear();\n // Load the range of the row.\n __syncthreads();\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_begin = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n __syncthreads();\n\n // _iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = is_active && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n int item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n item = a_col_id;\n }\n\n set.insert( item, wk_status );\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n utils::syncwarp();\n\n int num_rows = __popc( vote );\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int uniform_b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n // Load the range of the row of B.\n int b_col_it = A_rows[uniform_b_row_id + 0];\n int b_col_end = A_rows[uniform_b_row_id + 1];\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any(b_col_it < b_col_end) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != uniform_b_row_id;\n // Is it a strongly connected node.\n is_strongly_connected = is_off_diagonal && s_con[b_col_it];\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_strongly_connected && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n int b_item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n b_item = b_col_id;\n }\n\n set.insert( b_item, wk_status );\n }\n }\n }\n\n int c_col_it = C_hat_start[a_row_id];\n int count = set.store_with_positions( &C_hat[c_col_it], &C_hat_pos[c_col_it] );\n\n if ( lane_id == 0 )\n {\n C_hat_end[a_row_id] = c_col_it + count;\n }\n }\n\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_c_hat_kernel( int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_end,\n int *__restrict C_hat,\n int *__restrict C_hat_pos,\n int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // Shared memory to vote.\n __shared__ volatile int s_b_row_ids[CTA_SIZE];\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Skip coarse rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id >= 0 || coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // Clear the set.\n set.clear();\n\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_begin = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // _iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = is_active && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n int item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n item = a_col_id;\n }\n\n set.insert( item, wk_status );\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n utils::syncwarp();\n\n int num_rows = __popc( vote );\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int uniform_b_row_id = -1;\n\n if ( is_active_k )\n {\n uniform_b_row_id = s_b_row_ids[warp_id * WARP_SIZE + local_k];\n }\n\n // Load the range of the row of B.\n int b_col_it = 0, b_col_end = 0;\n\n if ( is_active_k )\n {\n b_col_it = A_rows[uniform_b_row_id + 0];\n b_col_end = A_rows[uniform_b_row_id + 1];\n }\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != uniform_b_row_id;\n // Is it a strongly connected node.\n is_strongly_connected = is_off_diagonal && s_con[b_col_it];\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_strongly_connected && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n int b_item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n b_item = b_col_id;\n }\n\n set.insert( b_item, wk_status );\n }\n }\n }\n\n int c_col_it = C_hat_start[a_row_id];\n int count = set.store_with_positions( &C_hat[c_col_it], &C_hat_pos[c_col_it] );\n\n if ( lane_id == 0 )\n {\n C_hat_end[a_row_id] = c_col_it + count;\n }\n }\n\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_inner_sum_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const int *__restrict C_hat,\n const int *__restrict C_hat_pos,\n const int *__restrict C_hat_start,\n const int *__restrict C_hat_end,\n const Value_type *__restrict diag,\n const int *__restrict inner_sum_offsets,\n Value_type *inner_sum,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue )\n{\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // A shared location where threads propose a row of B to load.\n __shared__ volatile int s_b_row_ids[CTA_SIZE];\n // A shared location where threads propose a value.\n __shared__ volatile Value_type s_a_values[CTA_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Skip coarse rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id >= 0 || coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // Clear the set.\n set.clear();\n // Rebuild C_hat.\n int c_hat_it = C_hat_start[a_row_id];\n int c_hat_end = C_hat_end [a_row_id];\n set.load( c_hat_end - c_hat_it, &C_hat[c_hat_it], &C_hat_pos[c_hat_it] );\n // The offset in the inner sum table.\n int inner_sum_offset = inner_sum_offsets[a_row_id];\n // And share the value of the diagonal.\n bool sign_diag = false;\n\n if ( lane_id == 0 )\n {\n sign_diag = sign( diag[a_row_id] );\n }\n\n sign_diag = utils::shfl( sign_diag, 0 );\n\n // Load A row IDs.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // _columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_off_diagonal && s_con[a_col_it] && cf_map[a_col_id] == FINE;\n // Read the associated value from A.\n Value_type a_value(0);\n\n if ( is_fine_strongly_connected )\n {\n a_value = A_vals[a_col_it];\n }\n\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n s_a_values [warp_id * WARP_SIZE + dest] = a_value;\n }\n\n utils::syncwarp();\n\n int num_rows = __popc( vote );\n // First n_rows threads reload the correct value.\n a_value = s_a_values[warp_id * WARP_SIZE + lane_id];\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Each thread keeps its own sum.\n Value_type bottom_sum(0);\n // Threads in the warp proceeds columns of B in the range [b_col_it, b_col_end).\n int b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n // TODO: make sure we have better memory accesses.\n int b_col_it = A_rows[b_row_id ];\n int b_col_end = A_rows[b_row_id + 1];\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any( b_col_it < b_col_end ) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it needed to compute the bottom sum.\n bool is_needed = set.contains( b_col_id ) || b_col_id == a_row_id;\n // Read the associated value from A.\n Value_type b_value(0);\n\n if ( is_needed )\n {\n b_value = A_vals[b_col_it];\n }\n\n // Update the sum if needed.\n if ( sign_diag != sign( b_value ) )\n {\n bottom_sum += b_value;\n }\n }\n\n // Reduce the row to a single value.\n#pragma unroll\n for ( int offset = WARP_SIZE / 2 ; offset > 0 ; offset >>= 1 )\n {\n int hi = __double2hiint(bottom_sum);\n hi = utils::shfl_xor( hi, offset );\n int lo = __double2loint(bottom_sum);\n lo = utils::shfl_xor( lo, offset );\n bottom_sum += __hiloint2double(hi, lo);\n }\n\n bottom_sum = utils::shfl( bottom_sum, 0 );\n\n if ( lane_id == k && bottom_sum != Value_type(0) )\n {\n a_value /= bottom_sum;\n }\n }\n\n // Write the results.\n if ( lane_id < num_rows )\n {\n inner_sum[inner_sum_offset + lane_id] = a_value;\n }\n\n inner_sum_offset += num_rows;\n }\n }\n\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_inner_sum_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const int *__restrict C_hat,\n const int *__restrict C_hat_pos,\n const int *__restrict C_hat_start,\n const int *__restrict C_hat_end,\n const Value_type *__restrict diag,\n const int *__restrict inner_sum_offsets,\n Value_type *inner_sum,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue )\n{\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // A shared location where threads propose a row of B to load.\n __shared__ volatile int s_b_row_ids[CTA_SIZE];\n // A shared location where threads propose a value.\n __shared__ volatile Value_type s_a_values[CTA_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Skip coarse rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id >= 0 || coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // Clear the set.\n set.clear();\n // Rebuild C_hat.\n int c_hat_it = C_hat_start[a_row_id];\n int c_hat_end = C_hat_end [a_row_id];\n set.load( c_hat_end - c_hat_it, &C_hat[c_hat_it], &C_hat_pos[c_hat_it] );\n // The offset in the inner sum table.\n int inner_sum_offset = inner_sum_offsets[a_row_id];\n // And share the value of the diagonal.\n bool sign_diag = false;\n\n if ( lane_id == 0 )\n {\n sign_diag = sign( diag[a_row_id] );\n }\n\n sign_diag = utils::shfl( sign_diag, 0 );\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // _columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_off_diagonal && s_con[a_col_it] && cf_map[a_col_id] == FINE;\n // Read the associated value from A.\n Value_type a_value(0);\n\n if ( is_fine_strongly_connected )\n {\n a_value = A_vals[a_col_it];\n }\n\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n s_a_values [warp_id * WARP_SIZE + dest] = a_value;\n }\n\n int num_rows = __popc( vote );\n // First n_rows threads reload the correct value.\n a_value = s_a_values[warp_id * WARP_SIZE + lane_id];\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n // Each thread keeps its own sum.\n Value_type bottom_sum(0);\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [b_col_it, b_col_end).\n int b_row_id = -1;\n\n if ( is_active_k )\n {\n b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k + lane_id_div_num_threads];\n }\n\n // TODO: make sure we have better memory accesses.\n int b_col_it = 0, b_col_end = 0;\n\n if ( is_active_k )\n {\n b_col_it = A_rows[b_row_id ];\n b_col_end = A_rows[b_row_id + 1];\n }\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any( b_col_it < b_col_end ) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it needed to compute the bottom sum.\n bool is_needed = set.contains( b_col_id ) || b_col_id == a_row_id;\n // Read the associated value from A.\n Value_type b_value(0);\n\n if ( is_needed )\n {\n b_value = A_vals[b_col_it];\n }\n\n // Update the sum if needed.\n if ( sign_diag != sign( b_value ) )\n {\n bottom_sum += b_value;\n }\n }\n\n // Reduce the row to a single value.\n#pragma unroll\n for ( int offset = NUM_THREADS_PER_ROW / 2 ; offset > 0 ; offset >>= 1 )\n {\n int hi = __double2hiint(bottom_sum);\n hi = utils::shfl_down( hi, offset, NUM_THREADS_PER_ROW );\n int lo = __double2loint(bottom_sum);\n lo = utils::shfl_down( lo, offset, NUM_THREADS_PER_ROW );\n bottom_sum += __hiloint2double(hi, lo);\n }\n\n bottom_sum = utils::shfl( bottom_sum, lane_id_mod_num_threads * NUM_THREADS_PER_ROW );\n\n if ( lane_id >= k && lane_id < k + NUM_THREADS_PER_ROW && bottom_sum != Value_type(0) )\n {\n a_value /= bottom_sum;\n }\n }\n\n // Write the results.\n if ( lane_id < num_rows )\n {\n inner_sum[inner_sum_offset + lane_id] = a_value;\n }\n\n inner_sum_offset += num_rows;\n }\n }\n\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_interp_weight_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const int *__restrict C_hat,\n const int *__restrict C_hat_pos,\n const int *__restrict C_hat_start,\n const int *__restrict C_hat_end,\n const Value_type *__restrict diag,\n const int *__restrict inner_sum_offsets,\n const Value_type *__restrict inner_sum,\n const int *__restrict P_rows,\n int *P_cols,\n Value_type *P_vals,\n const int gmem_size,\n int *g_keys,\n Value_type *g_vals,\n int *wk_work_queue )\n{\n \n const int NUM_WARPS = CTA_SIZE / 32;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // A shared location where threads propose a row of B to load.\n __shared__ volatile int s_b_row_ids[CTA_SIZE];\n // A shared location where threads propose a value.\n __shared__ volatile Value_type s_aki[NUM_WARPS];\n // The hash values stored in shared memory.\n __shared__ Value_type s_vals[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_map map( &s_keys[warp_id * SMEM_SIZE],\n &g_keys[a_row_id * gmem_size],\n &s_vals[warp_id * SMEM_SIZE],\n &g_vals[a_row_id * gmem_size],\n gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n int coarse_fine_id = cf_map[a_row_id];\n\n // Skip coarse rows.\n if ( coarse_fine_id >= 0 )\n {\n if ( lane_id == 0 )\n {\n int p_row_it = P_rows[a_row_id];\n P_cols[p_row_it] = coarse_fine_id;\n P_vals[p_row_it] = Value_type( 1 );\n }\n\n continue;\n }\n\n if ( coarse_fine_id == STRONG_FINE )\n {\n continue;\n }\n\n // Clear the table.\n map.clear();\n // Rebuild C_hat.\n int c_hat_it = C_hat_start[a_row_id];\n int c_hat_end = C_hat_end [a_row_id];\n map.load( c_hat_end - c_hat_it, &C_hat[c_hat_it], &C_hat_pos[c_hat_it] );\n // Load A row IDs.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // The offset in the inner sum table.\n int inner_sum_offset = inner_sum_offsets[a_row_id];\n // Weak value.\n Value_type weak(0), sum(0);\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id(-1);\n Value_type a_value(0);\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n a_value = A_vals[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it a strongly-connected column.\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it a weakly connected node.\n bool is_weakly_connected = is_off_diagonal && !is_strongly_connected;\n // Is isolated\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Update C_hat values. If the value is not in C_hat it will be skipped (true parameter).\n bool is_in_c_hat = map.update( a_col_id, a_value );\n\n // Update the weak value.\n if ( is_weakly_connected && !is_in_c_hat && !is_strong_fine)\n {\n weak += a_value;\n }\n\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && cf_map[a_col_id] == FINE;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n utils::syncwarp();\n\n int num_rows = __popc( vote );\n // We pre-load inner sums.\n sum = Value_type(0);\n\n if ( lane_id < num_rows )\n {\n sum = inner_sum[inner_sum_offset + lane_id];\n }\n\n inner_sum_offset += num_rows;\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [b_col_it, bCol_end).\n int b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n // TODO: make sure we have better memory accesses. b_colBegin rather than bCol_it because we iterate twice.\n int b_col_begin = A_rows[b_row_id ];\n int b_col_end = A_rows[b_row_id + 1];\n // Diagonal element.\n Value_type b_diag = diag[b_row_id];\n\n // Set s_aki to 0 because a_row_id may not appear in a row.\n if ( lane_id == 0 )\n {\n s_aki[warp_id] = Value_type(0);\n }\n\n utils::syncwarp();\n\n // Load the kth inner sum.\n Value_type uniform_val = utils::shfl( sum, k );\n\n // _iterate over the range of columns of B.\n for ( int b_col_it = b_col_begin + lane_id ; utils::any( b_col_it < b_col_end ) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n int b_col_id(-1);\n Value_type b_value(0);\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n b_value = A_vals[b_col_it];\n }\n\n // Update bValue if needed.\n if ( sign( b_diag ) == sign( b_value ) )\n {\n b_value = Value_type(0);\n }\n\n // See if we have found a_row_id in the row. Only one thread can evaluate to true.\n if ( a_row_id == b_col_id )\n {\n s_aki[warp_id] = b_value;\n }\n\n // Update C_hat values. If the value is not in C_hat it will be skipped (true parameter).\n map.update( b_col_id, b_value * uniform_val );\n }\n\n utils::syncwarp();\n\n // The thread k updates the sum.\n if ( lane_id == 0 )\n {\n weak += s_aki[warp_id] * uniform_val;\n }\n }\n }\n\n // We're done with that row of A. We compute D.\n#pragma unroll\n\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n weak += utils::shfl_xor( weak, mask );\n }\n\n if ( lane_id == 0 )\n {\n weak += diag[a_row_id];\n sum = Value_type(-1) / weak;\n }\n\n sum = utils::shfl( sum, 0 );\n\n int p_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n p_col_tmp = P_rows[a_row_id + lane_id];\n }\n\n int p_col_it = utils::shfl( p_col_tmp, 0 );\n int p_col_end = utils::shfl( p_col_tmp, 1 );\n\n map.store_map_keys_scale_values( p_col_end - p_col_it, cf_map, &P_cols[p_col_it], sum, &P_vals[p_col_it] );\n }\n}\n\n\n} // namespace distance2\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename IndexType, int kCtaSize >\n__global__\nvoid\ncalculateInnerSumStorageKernel( const IndexType *A_rows,\n const IndexType *A_cols,\n const int A_num_rows,\n const int *cf_map,\n const bool *s_con,\n int *innerSum_offsets )\n{\n const int nWarps = kCtaSize / 32;\n // The coordinates of the thread inside the CTA/warp.\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n // Loop over rows of A.\n for ( int aRowId = blockIdx.x * nWarps + warpId ; aRowId < A_num_rows ; aRowId += gridDim.x * nWarps )\n {\n // Skip coarse rows.\n int coarse_fine_id = cf_map[aRowId];\n\n if ( coarse_fine_id >= 0 || coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // The sum for each thread.\n int sum = 0;\n // Load A row IDs.\n int aColIt = A_rows[aRowId ];\n int aColEnd = A_rows[aRowId + 1];\n\n // Iterate over the columns of A.\n for ( aColIt += laneId ; utils::any( aColIt < aColEnd ) ; aColIt += 32 )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int aColId = aColIt < aColEnd ? A_cols[aColIt] : -1;\n // Is it an off-diagonal element.\n bool is_off_diagonal = aColIt < aColEnd && aColId != aRowId;\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_off_diagonal && s_con[aColIt] && cf_map[aColId] == FINE;\n // Update the sum.\n sum += __popc( utils::ballot( is_fine_strongly_connected ) );\n }\n\n // Update the value.\n if ( laneId == 0 )\n {\n innerSum_offsets[aRowId] = sum;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n/*\n * grab the diagonal of the matrix\n */\ntemplate \n__global__\nvoid getDiagonalKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, const IndexType numRows, ValueType *diagonal)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n for (int j = offset; j < offset + numj; j++)\n {\n int jcol = column_indices[j];\n\n if (tIdx == jcol)\n {\n diagonal[tIdx] = values[j];\n }\n }\n }\n}\n\n/*************************************************************************\n * Implementing Extended+i algorithm from \\S 4.5 of:\n * \"Distance-two interpolation for parallel algebraic multigrid\"\n * Reference [4] on wiki\n ************************************************************************/\ntemplate \nDistance2_Interpolator >::Distance2_Interpolator(AMG_Config &cfg, const std::string &cfg_scope) : Base(cfg, cfg_scope)\n{}\n\ntemplate \nDistance2_Interpolator >::~Distance2_Interpolator()\n{}\n\nenum { WARP_SIZE = 32, SMEM_SIZE = 128 };\n\ntemplate \nvoid Distance2_Interpolator >::generateInterpolationMatrix_1x1(Matrix_d &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix_d &P)\n{\n nvtxRange gim_nr(__func__);\n\n const int blockSize = 256;\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // get raw pointers to data\n const IndexType *Aoffsets = A.row_offsets.raw();\n const IndexType *Acolumn_indices = A.col_indices.raw();\n const ValueType *Avalues = A.values.raw();\n const IndexType Anum_rows = (int) A.get_num_rows();\n\n cudaDeviceProp props = getDeviceProperties();\n int grid_size = (props.major >= 7) ? 1024 : 128;\n int gmem_size = (props.major >= 7) ? 512 : 2048;\n\n Hash_Workspace exp_wk(true, grid_size, 8, gmem_size);\n\n IntVector C_hat_start( A.get_num_rows() + 1, 0 ), C_hat_end( A.get_num_rows() + 1, 0 );\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int work_offset = grid_size * NUM_WARPS;\n cudaMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), cudaMemcpyHostToDevice );\n cudaCheckError();\n int avg_nz_per_row = (A.get_num_rows() == 0) ? 0 : A.get_num_nz() / A.get_num_rows();\n\n if ( avg_nz_per_row < 16 )\n {\n distance2::estimate_c_hat_size_kernel< 8, CTA_SIZE, WARP_SIZE> <<< 2048, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw());\n cudaCheckError();\n }\n else\n {\n distance2::estimate_c_hat_size_kernel <<< 2048, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw());\n }\n\n cudaCheckError();\n }\n // Compute row offsets.\n thrust_wrapper::exclusive_scan( C_hat_start.begin( ), C_hat_start.end( ), C_hat_start.begin( ) );\n cudaCheckError();\n // Allocate memory to store columns/values.\n int nVals = C_hat_start[C_hat_start.size() - 1];\n IntVector C_hat( nVals ), C_hat_pos( nVals );\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int avg_nz_per_row = (A.get_num_rows() == 0) ? 0 : A.get_num_nz() / A.get_num_rows();\n int attempt = 0;\n\n for ( bool done = false ; !done && attempt < 10 ; ++attempt )\n {\n // Double the amount of GMEM (if needed).\n if ( attempt > 0 )\n {\n exp_wk.expand();\n }\n\n // Reset the status. TODO: Launch async copies.\n int status = 0;\n cudaMemcpy( exp_wk.get_status(), &status, sizeof(int), cudaMemcpyHostToDevice );\n cudaCheckError();\n // Compute the set C_hat.\n int work_offset = grid_size * NUM_WARPS;\n cudaMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), cudaMemcpyHostToDevice );\n cudaCheckError();\n\n // Run the computation.\n if ( avg_nz_per_row < 16 )\n {\n distance2::compute_c_hat_kernel< 8, CTA_SIZE, SMEM_SIZE, WARP_SIZE> <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n C_hat.raw(),\n C_hat_pos.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_work_queue(),\n exp_wk.get_status() );\n cudaCheckError();\n }\n else\n {\n distance2::compute_c_hat_kernel <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n C_hat.raw(),\n C_hat_pos.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_work_queue(),\n exp_wk.get_status() );\n }\n\n cudaCheckError();\n // Read the result from count_non_zeroes.\n cudaMemcpy( &status, exp_wk.get_status(), sizeof(int), cudaMemcpyDeviceToHost );\n cudaCheckError();\n done = status == 0;\n }\n }\n // get pointers to data\n int *C_hat_ptr = C_hat.raw();\n int *C_hat_start_ptr = C_hat_start.raw();\n int *C_hat_end_ptr = C_hat_end.raw();\n // modify Coarse/Fine map array to have the column index of the new interpolation matrix\n // for evry coarse point\n int coarsePoints;\n DistributedArranger *prep;\n\n if (A.is_matrix_distributed())\n {\n prep = new DistributedArranger;\n int num_owned_fine_pts = A.get_num_rows();\n int num_owned_coarse_pts = thrust_wrapper::count_if(cf_map.begin(), cf_map.begin() + num_owned_fine_pts, is_non_neg());\n cudaCheckError();\n int num_halo_coarse_pts = thrust_wrapper::count_if(cf_map.begin() + num_owned_fine_pts, cf_map.end(), is_non_neg());\n cudaCheckError();\n coarsePoints = num_owned_coarse_pts + num_halo_coarse_pts;\n // Using the number of owned_coarse_pts (number of rows of Ac), initialize P manager\n // This will set the part_offsets array, index_range, halo_offsets[0] value\n prep->initialize_manager(A, P, num_owned_coarse_pts);\n }\n else\n {\n coarsePoints = (int) thrust_wrapper::count_if(cf_map.begin(), cf_map.end(), is_non_neg());\n cudaCheckError();\n }\n\n // mark as 1 if a coarse point\n int size_one_ring;\n\n if (!A.is_matrix_singleGPU())\n {\n int offset;\n A.getOffsetAndSizeForView(FULL, &offset, &size_one_ring);\n }\n else\n {\n size_one_ring = A.get_num_rows();\n }\n\n int numBlocksOneRing = std::min( 4096, (int) (size_one_ring + blockSize - 1) / blockSize );\n int numBlocks = std::min( 4096, (int) (A.get_num_rows() + blockSize - 1) / blockSize );\n // count the number of non-zeros in the interpolation matrix\n IntVector nonZeroOffsets(A.get_num_rows() + 1);\n IntVector nonZerosPerRow(A.get_num_rows());\n int *nonZerosPerRow_ptr = nonZerosPerRow.raw();\n\n if (A.get_num_rows() > 0)\n {\n nonZerosPerRowKernel <<< numBlocks, blockSize>>>(Anum_rows, cf_map.raw(), C_hat_start_ptr, C_hat_end_ptr,\n nonZerosPerRow_ptr);\n cudaCheckError();\n }\n\n // get total with a reduction\n int nonZeros = thrust_wrapper::reduce(nonZerosPerRow.begin(), nonZerosPerRow.end());\n cudaCheckError();\n // get the offsets with an exclusive scan\n thrust_wrapper::exclusive_scan(nonZerosPerRow.begin(), nonZerosPerRow.end(), nonZeroOffsets.begin());\n cudaCheckError();\n nonZeroOffsets[A.get_num_rows()] = nonZeros;\n // resize P\n P.resize(0, 0, 0, 1);\n P.addProps(CSR);\n P.resize(A.get_num_rows(), coarsePoints, nonZeros, 1);\n IndexType *Poffsets = P.row_offsets.raw();\n IndexType *Pcolumn_indices = P.col_indices.raw();\n ValueType *Pvalues = P.values.raw();\n // copy nonzero offsets to the P matrix\n amgx::thrust::copy(nonZeroOffsets.begin(), nonZeroOffsets.end(), P.row_offsets.begin());\n cudaCheckError();\n // grab the diagonal terms\n VVector diag(size_one_ring);\n ValueType *diag_ptr = diag.raw();\n\n if (A.get_num_rows() > 0)\n {\n find_diag_kernel_indexed_dia <<< numBlocksOneRing, blockSize>>>(\n size_one_ring,\n A.diag.raw(),\n A.values.raw(),\n diag.raw());\n }\n\n cudaCheckError();\n IntVector innerSumOffset( A.get_num_rows() + 1, 0 );\n\n if (A.get_num_rows() > 0)\n {\n const int kCtaSize = 256;\n const int nWarps = kCtaSize / 32;\n const int grid_size = std::min( 4096, ( Anum_rows + nWarps - 1 ) / nWarps );\n calculateInnerSumStorageKernel <<< grid_size, kCtaSize>>>(\n Aoffsets,\n Acolumn_indices,\n Anum_rows,\n cf_map.raw(),\n s_con.raw(),\n innerSumOffset.raw() );\n cudaCheckError();\n }\n\n // get the offsets\n thrust_wrapper::exclusive_scan( innerSumOffset.begin(), innerSumOffset.end(), innerSumOffset.begin() );\n cudaCheckError();\n // assign memory & get pointer\n int numInnerSum = innerSumOffset[A.get_num_rows()];\n VVector innerSum( numInnerSum, 0 );\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int avg_nz_per_row = (A.get_num_rows() == 0) ? 0 : A.get_num_nz() / A.get_num_rows();\n // Compute the set C_hat.\n int work_offset = grid_size * NUM_WARPS;\n cudaMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), cudaMemcpyHostToDevice );\n cudaCheckError();\n // Run the computation.\n typedef typename MatPrecisionMap::Type Value_type;\n {\n distance2::compute_inner_sum_kernel <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat.raw(),\n C_hat_pos.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n diag.raw(),\n innerSumOffset.raw(),\n innerSum.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_work_queue());\n }\n cudaCheckError();\n }\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int avg_nz_per_row = (A.get_num_rows() == 0) ? 0 : A.get_num_nz() / A.get_num_rows();\n // Compute the set C_hat.\n int work_offset = grid_size * NUM_WARPS;\n cudaMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), cudaMemcpyHostToDevice );\n cudaCheckError();\n // Run the computation.\n typedef typename MatPrecisionMap::Type Value_type;\n distance2::compute_interp_weight_kernel <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat.raw(),\n C_hat_pos.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n diag.raw(),\n innerSumOffset.raw(),\n innerSum.raw(),\n P.row_offsets.raw(),\n P.col_indices.raw(),\n P.values.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_vals(),\n exp_wk.get_work_queue());\n cudaCheckError();\n }\n\n if (A.is_matrix_distributed())\n {\n //TODO: cf_map_global could be smaller (only contain halo nodes)\n I64Vector_d cf_map_global(cf_map.size());\n int num_owned_fine_pts = P.get_num_rows();\n int my_rank = A.manager->global_id();\n const int cta_size = 128;\n const int grid_size = std::min( 4096, (num_owned_fine_pts + cta_size - 1) / cta_size);\n\n // Write the global coarse index of all the owned fine points (-1 if not a coarse point)\n if (num_owned_fine_pts > 0)\n {\n createCfMapGlobal <<< grid_size, cta_size>>>(cf_map.raw(), cf_map_global.raw(), P.manager->part_offsets_h[my_rank], num_owned_fine_pts);\n cudaCheckError();\n }\n\n // Exchange the cf_map_global so that we know the coarse global id of halo nodes\n cf_map_global.dirtybit = 1;\n A.manager->exchange_halo_2ring(cf_map_global, cf_map_global.tag);\n I64Vector_d P_col_indices_global(P.col_indices.size());\n\n if (num_owned_fine_pts > 0)\n {\n fill_P_global_col_indices_kernel<16, int> <<< grid_size, cta_size>>>(P.row_offsets.raw(), C_hat_start.raw(), C_hat_end.raw(), C_hat.raw(), cf_map_global.raw(), P_col_indices_global.raw(), num_owned_fine_pts);\n cudaCheckError();\n }\n\n // Now that we know the global col indices of P, we can find:\n // the list of neighbors, create the B2L_maps for P, create local numbering for halo nodes, create local_to_global_map array, etc...\n prep->initialize_manager_from_global_col_indices(P, P_col_indices_global);\n prep->createRowsLists(P, true);\n delete prep;\n }\n}\n\ntemplate< class T_Config>\nvoid Distance2_InterpolatorBase::generateInterpolationMatrix(Matrix &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix &P)\n{\n P.set_initialized(0);\n ViewType oldView = A.currentView();\n A.setView(OWNED);\n\n // If multi-GPU do a 2-ring exchange of cf_map since this is a distance two interpolator\n if (A.get_block_size() == 1)\n {\n generateInterpolationMatrix_1x1(A, cf_map, s_con, scratch, P);\n }\n else\n {\n FatalError(\"Unsupported dimensions for distance2 interpolator\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n P.set_initialized(1);\n}\n\n#define AMGX_CASE_LINE(CASE) template class Distance2_InterpolatorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Distance2_Interpolator::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace amgx\n{\n\nusing std::set;\nusing std::vector;\n\n/*\n * remove all duplicates from a vector\n */\ntemplate \nvoid removeDuplicates(vector &vec)\n{\n std::sort(vec.begin(), vec.end());\n vec.erase(std::unique(vec.begin(), vec.end()), vec.end());\n}\n\n/*\n * find the value at a particular column of a row if it exists\n */\ntemplate \ntypename VVector::value_type findValue(const IVector &columns,\n const VVector &values,\n const int row_begin, const int row_end, const int column_needed)\n{\n typename VVector::value_type v = 0;\n\n // NAIVE\n for (int k = row_begin; k < row_end; k++)\n {\n int kcol = columns[k];\n\n if (kcol == column_needed)\n {\n v = values[k];\n }\n }\n\n return v;\n}\n\n\ntemplate \nvoid __global__ createCfMapGlobal(const IndexType *cf_map, int64_t *cf_map_global, const int64_t my_part_offset, const IndexType num_owned_fine_pts)\n{\n for (int tidx = threadIdx.x + blockIdx.x * blockDim.x; tidx < num_owned_fine_pts ; tidx += blockDim.x * gridDim.x)\n {\n // Only renumber the interior points\n if (cf_map[tidx] >= 0 && tidx < num_owned_fine_pts)\n {\n // Simply shift\n cf_map_global[tidx] = (int64_t) cf_map[tidx] /* my_local_id */ + my_part_offset;\n }\n else\n {\n cf_map_global[tidx] = -1;\n }\n }\n}\n\ntemplate \n__global__ void fill_P_global_col_indices_kernel(index_type *row_offsets, index_type *C_hat_start, index_type *C_hat_end, index_type *C_hat, int64_t *cf_map_global, int64_t *global_col_indices, index_type num_rows)\n{\n int coopIdx = threadIdx.x % coop;\n\n for (int row = blockIdx.x * blockDim.x / coop + threadIdx.x / coop; row < num_rows ; row += blockDim.x * gridDim.x / coop)\n {\n int64_t coarse_fine_id = cf_map_global[row];\n int j = row_offsets[row] + coopIdx;\n\n // coarse row\n if (coarse_fine_id >= 0)\n {\n if (coopIdx == 0)\n {\n global_col_indices[j] = coarse_fine_id;\n }\n\n continue;\n }\n\n // strong fine row\n if (coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // fine row\n for (int i = C_hat_start[row] + coopIdx; i < C_hat_end[row]; i += coop, j += coop)\n {\n index_type fine_col = C_hat[i];\n int64_t global_col = cf_map_global[fine_col];\n global_col_indices[j] = global_col;\n }\n }\n}\n\n\n/*\n * calculate \\tilde{a}_{ii} from eqn. 4.11 in [4]\n */\ntemplate \nvoid Distance2_Interpolator >::calculateD( const Matrix_h &A, IVector &cf_map, BVector &s_con,\n vector *C_hat, VVector &diag, VVector &D,\n set *weak_lists, VVector &innerSum,\n IVector &innerSumOffsets)\n{\n // only do the non-weak part for now\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n int outerOffset = innerSumOffsets[i], innerOffset = 0;\n ValueType sum = 0;\n\n // loop over all non-zeros\n for (int k = A.row_offsets[i]; k < A.row_offsets[i + 1]; k++)\n {\n int kcol = A.col_indices[k];\n\n if (cf_map[kcol] == FINE && s_con[k]) // strongly connected fine, ignore diagonal\n {\n int sgn = diag[kcol] < 0 ? -1 : 1;\n // find a_ki\n ValueType temp = findValue(A.col_indices, A.values, A.row_offsets[kcol], A.row_offsets[kcol + 1], i);\n ValueType a_ki = (sgn * temp < 0) * temp;\n sum += a_ki * innerSum[outerOffset + innerOffset];\n innerOffset++;\n }\n }\n\n // now calculate the weak part\n set difference;\n\n for (set::iterator it = weak_lists[i].begin(); it != weak_lists[i].end(); ++it)\n {\n difference.insert(*it);\n }\n\n for (int l = 0; l < C_hat[i].size(); l++)\n {\n difference.erase(C_hat[i][l]);\n }\n\n ValueType weak = 0;\n\n for (set::iterator it = difference.begin(); it != difference.end(); ++it)\n {\n weak += findValue(A.col_indices, A.values, A.row_offsets[i], A.row_offsets[i + 1], *it);\n }\n\n D[i] = diag[i] + weak + sum;\n }\n}\n\n/*\n * Generate the lower sum from eqns. 4.10 and 4.11 of [4]\n * strongly resembles a sparse dot product\n */\ntemplate \ntypename VVector::value_type sparseSum(const IVector &columns,\n const VVector &values,\n const int row_begin, const int row_end, const std::vector &C_hat_row, const int sgn)\n{\n int i1 = row_begin, i2 = 0;\n\n // check lengths are non-zero\n if (row_begin >= row_end || C_hat_row.size() == 0) { return 0; }\n\n typename VVector::value_type sum = 0;\n\n // if column exists in both C_hat_row and row of matrix, add matrix value\n while (true)\n {\n int col1 = columns[i1];\n int col2 = C_hat_row[i2];\n\n if (col1 == col2) // match\n {\n if (sgn * values[i1] < 0)\n {\n sum += values[i1];\n }\n\n i1++;\n i2++;\n\n if (i1 >= row_end || i2 >= C_hat_row.size()) { break; }\n }\n else if (col1 > col2)\n {\n // instead of linear increase, binary search\n /* int min = i2++, max = C_hat_row.size();\n while(min <= max)\n {\n int mid = (min+max)/2;\n col2 = C_hat_row[mid];\n if (col1 > col2)\n min = mid+1;\n else if (col1 < col2)\n max = mid-1;\n else\n {\n i2 = mid;\n break;\n }\n }*/\n i2++;\n\n if (i2 >= C_hat_row.size()) { break; }\n }\n else\n {\n i1++;\n\n if (i1 >= row_end) { break; }\n }\n }\n\n return sum;\n}\n\n/*\n * Generate the repeatedly reused inner sum from eqns 4.10, 4.11 in [4]\n * \\sum_{k\\in F^s_i} a_{ik}\\frac{1}{\\sum_{l\\in\\hat{C}_i\\cup\\{i\\}} \\bar{a}_{kl}}\n */\ntemplate \nvoid Distance2_Interpolator >::generateInnerSum(Matrix_h &A, const IntVector &cf_map,\n const BVector &s_con,\n vector *C_hat, VVector &diag,\n VVector &innerSum,\n IntVector &innerSumOffsets)\n{\n // calculate necessary storage for innerSum\n int connections = 0;\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n int jcol = A.col_indices[j];\n\n if (cf_map[jcol] == FINE && s_con[j] && jcol != i)\n {\n connections++;\n }\n }\n }\n\n innerSum.resize(connections);\n innerSumOffsets.resize(A.get_num_rows() + 1);\n innerSumOffsets[0] = 0;\n int innerOffset = 0;\n\n // generate the values\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n for (int k = A.row_offsets[i]; k < A.row_offsets[i + 1]; k++)\n {\n int kcol = A.col_indices[k];\n\n if (cf_map[kcol] == FINE && s_con[k] && kcol != i)\n {\n ValueType bottom_sum = 0;\n ValueType a_ik = A.values[k];\n int sgn = diag[kcol] < 0 ? -1 : 1;\n // generate the bottom sum - looks similar to sparse dot product\n bottom_sum = sparseSum(A.col_indices, A.values, A.row_offsets[kcol], A.row_offsets[kcol + 1],\n C_hat[i], sgn);\n // add \\union i\n ValueType iTerm = findValue(A.col_indices, A.values, A.row_offsets[kcol], A.row_offsets[kcol + 1], i);\n bottom_sum += (sgn * iTerm < 0) * iTerm;\n\n if (bottom_sum == 0 && cf_map[i] == FINE)\n {\n char buf[255];\n sprintf(buf, \"DIVISION BY ZERO: row: %d, col: %d\\n\", i, kcol);\n amgx_output(buf, strlen(buf));\n }\n\n innerSum[innerOffset] = a_ik / bottom_sum;\n innerOffset++;\n }\n }\n\n innerSumOffsets[i + 1] = innerOffset;\n }\n}\n\n/*************************************************************************\n* create the interpolation matrix P\n************************************************************************/\ntemplate \nvoid Distance2_Interpolator >::generateInterpolationMatrix_1x1(Matrix_h &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix_h &P)\n{\n if (!A.is_matrix_singleGPU())\n {\n FatalError(\"Distributed classical AMG not implemented on host\\n\", AMGX_ERR_NOT_IMPLEMENTED);\n }\n\n typedef typename Matrix_h::value_type ValueType;\n // implementation begins here\n set *coarse_lists = new set[A.get_num_rows()];\n set *weak_lists = new set[A.get_num_rows()];\n vector *C_hat = new vector[A.get_num_rows()]; // hold full \\hat{C} set for each row\n set *C_hat_set = new set[A.get_num_rows()];\n set::iterator set_it;\n VVector diag(A.get_num_rows(), 0);\n\n // grab the diagonal\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n if (A.col_indices[j] == i)\n {\n diag[i] = A.values[j];\n }\n }\n }\n\n // generate coarse sets\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n // loop over non-zero elements\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n int jcol = A.col_indices[j];\n\n if (jcol == i) { continue; } // skip the diagonal\n\n // Coarse and strongly connected\n if (cf_map[jcol] >= 0 && s_con[j])\n {\n coarse_lists[i].insert(jcol);\n }\n else if (!s_con[j] && jcol != i && cf_map[jcol] != STRONG_FINE)\n {\n weak_lists[i].insert(jcol);\n }\n }\n }\n\n int numWTotal = 0;\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n numWTotal += (int) weak_lists[i].size();\n }\n\n typedef typename set::iterator iSetIter;\n\n // generate \\hat{C} for each row, i\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n C_hat[i].clear();\n\n if (cf_map[i] == FINE)\n {\n // copy coarse_lists[i] to C_hat_set[i]\n for (iSetIter it = coarse_lists[i].begin(); it != coarse_lists[i].end(); ++it)\n {\n C_hat_set[i].insert(*it);\n }\n\n // loop over non-zeros\n for (int j = A.row_offsets[i]; j < A.row_offsets[i + 1]; j++)\n {\n int jcol = A.col_indices[j];\n\n if (jcol == i) { continue; } // don't want to take into account myself\n\n // if this is a strong-fine connection, union of the sets\n if (cf_map[jcol] == FINE && s_con[j])\n {\n for (iSetIter it = coarse_lists[jcol].begin(); it != coarse_lists[jcol].end(); ++it)\n {\n C_hat_set[i].insert(*it);\n }\n }\n }\n\n removeDuplicates(C_hat[i]);\n }\n\n // copy C_hat_set[i] into C_hat[i]\n for (iSetIter it = C_hat_set[i].begin(); it != C_hat_set[i].end(); ++it)\n {\n C_hat[i].push_back(*it);\n }\n }\n\n // check that \\hat{C_i} for fine points does not contain i\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n if (cf_map[i] == FINE)\n {\n vector::iterator it = std::find(C_hat[i].begin(), C_hat[i].end(), i);\n {\n assert(it == C_hat[i].end());\n }\n }\n }\n\n // generate the inner sum for all i, k\n VVector innerSum;\n IntVector innerSumOffsets;\n generateInnerSum(A, cf_map, s_con, C_hat, diag, innerSum, innerSumOffsets);\n // calculate the ~a_ii term\n VVector D(A.get_num_rows(), 0);\n calculateD(A, cf_map, s_con, C_hat, diag, D, weak_lists, innerSum, innerSumOffsets);\n // delete all storage we have no more need for\n delete [] coarse_lists;\n delete [] weak_lists;\n delete [] C_hat_set;\n // now we have \\hat{C} and the diagonals, we can start calculating the weights\n // remap cf_map (to be like in distance1)\n int coarsePoints = 0;\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n // if coarse\n if (cf_map[i] >= 0) { coarsePoints++; }\n }\n\n // count the non-zeros - 1 per row if coarse, len(C_hat[i]) otherwise\n IntVector nonZeroOffsets(A.get_num_rows() + 1);\n IntVector nonZerosPerRow(A.get_num_rows());\n int nonZeros = 0;\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n if (cf_map[i] >= 0)\n {\n nonZerosPerRow[i] = 1;\n nonZeros++;\n }\n else if (cf_map[i] == STRONG_FINE)\n {\n nonZerosPerRow[i] = 0;\n }\n else\n {\n nonZerosPerRow[i] = (int) C_hat[i].size();\n nonZeros += (int) C_hat[i].size();\n }\n }\n\n // get the offsets with an exclusive scan\n thrust_wrapper::exclusive_scan(nonZerosPerRow.begin(), nonZerosPerRow.end(), nonZeroOffsets.begin());\n cudaCheckError();\n nonZeroOffsets[A.get_num_rows()] = nonZeros;\n // resize P\n P.addProps(CSR);\n P.resize(A.get_num_rows(), coarsePoints, nonZeros, 1);\n // copy nonzero offsets to the P matrix\n amgx::thrust::copy(nonZeroOffsets.begin(), nonZeroOffsets.end(), P.row_offsets.begin());\n cudaCheckError();\n // main loop\n int outerOffset = 0, innerOffset;\n\n for (int i = 0; i < A.get_num_rows(); i++)\n {\n int nonZeroOffset = nonZeroOffsets[i];\n int localNonZeros = 0;\n outerOffset = innerSumOffsets[i];\n\n if (cf_map[i] >= 0)\n {\n // will be setting here & continuing\n P.values[nonZeroOffset + localNonZeros] = 1.0;\n P.col_indices[nonZeroOffset + localNonZeros] = cf_map[i];\n continue;\n }\n else if (cf_map[i] != STRONG_FINE) // fine row\n {\n //\n for (int j = 0; j < C_hat[i].size(); j++)\n {\n innerOffset = 0; // set to 'row' offset\n int jcol = C_hat[i][j]; // j can potentially not exist on row i\n\n if (jcol == i) // skip diagonal\n {\n FatalError(\"Error - Diagonal should never be in C_hat[i]\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n ValueType value = 0;\n // loop over i and find a_ij - if not found, doesn't exist, keep as 0\n ValueType a_ij = 0;\n a_ij = findValue(A.col_indices, A.values, A.row_offsets[i], A.row_offsets[i + 1], jcol);\n\n // now loop over and find strong-fine connections\n for (int k = A.row_offsets[i]; k < A.row_offsets[i + 1]; k++)\n {\n int kcol = A.col_indices[k];\n\n if (cf_map[kcol] == FINE && s_con[k]) // Strong-Fine edge, ignore diagonal\n {\n if (kcol == i)\n {\n FatalError(\"Error - point should never to strongly connected to itself\", AMGX_ERR_BAD_PARAMETERS);\n }\n\n // point we care about\n int sgn = diag[kcol] < 0 ? -1 : 1;\n ValueType a_kj = 0;\n // get \\bar{a_kj}\n ValueType temp = findValue(A.col_indices, A.values, A.row_offsets[kcol], A.row_offsets[kcol + 1], jcol);\n a_kj = (sgn * temp < 0) * temp;\n // see if we need to add \\bar{a_ki} in\n vector::iterator it = std::find(C_hat[i].begin(), C_hat[i].end(), i);\n\n if (it != C_hat[i].end())\n {\n FatalError(\"Error - point i should never be in\\\\hat{C_i}\", AMGX_ERR_BAD_PARAMETERS);\n }\n else\n {\n // we now have everything we need - add to sum\n // not valid in non-symmetric matrices\n value += innerSum[outerOffset + innerOffset] * a_kj;\n innerOffset++;\n }\n } // strong fine connection end\n } // k loop (find strong fine connections) end\n\n ValueType diagonal = (D[i]);\n // now set column & value for P\n P.col_indices[nonZeroOffset + localNonZeros] = cf_map[jcol];\n P.values[nonZeroOffset + localNonZeros] = -1.0 / (diagonal) * (a_ij + value);\n localNonZeros++;\n } // end j loop\n } // fine connection loop\n\n outerOffset += innerOffset;\n } // end i loop\n\n delete [] C_hat;\n} // end distance2 interpolator\n\n\n///////////////////////////////////////////////////////////////////////////////\n\n// DEVICE CODE\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#include \n#include // Included inside the namespace to solve name colisions.\n\n__device__ __forceinline__ int get_work( int *queue, int warp_id )\n{\n int offset = -1;\n\n if ( utils::lane_id() == 0 )\n {\n offset = atomicAdd( queue, 1 );\n }\n\n return utils::shfl( offset, 0 );\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\nnamespace distance2\n{\n\n#include \n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int CTA_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nestimate_c_hat_size_kernel( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const int *cf_map,\n const bool *s_con,\n int *C_hat_offsets )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // A shared location where threads propose a row of B to load.\n __shared__ volatile int s_b_row_ids[CTA_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += gridDim.x * NUM_WARPS )\n {\n // Skip coarse and strong fine rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id >= 0 || coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // Load A row IDs.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // The number of elements.\n int count(0);\n\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n __syncthreads();\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0, num_rows = __popc(vote) ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [b_col_it, b_col_end).\n int b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n // Load the range of B.\n int b_col_it = A_rows[b_row_id + 0];\n int b_col_end = A_rows[b_row_id + 1];\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any(b_col_it < b_col_end) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != b_row_id;\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_off_diagonal && s_con[b_col_it] && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n }\n }\n }\n\n // Store the number of columns in each row.\n if ( lane_id == 0 )\n {\n C_hat_offsets[a_row_id] = count;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\nestimate_c_hat_size_kernel( const int A_num_rows,\n const int *A_rows,\n const int *A_cols,\n const int *cf_map,\n const bool *s_con,\n int *C_hat_offsets )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // A shared location where threads propose a row of B to load.\n __shared__ volatile int s_b_row_ids[CTA_SIZE];\n s_b_row_ids[threadIdx.x] = 0;\n __syncthreads();\n\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id += gridDim.x * NUM_WARPS )\n {\n // Skip coarse rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id >= 0 || coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // Load A row IDs.\n int a_col_begin = A_rows[a_row_id ];\n int a_col_end = A_rows[a_row_id + 1];\n // The number of elements.\n int count(0);\n\n // Iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n utils::syncwarp();\n\n // For each warp, we have up to 32 rows of B to proceed.\n\n int num_rows = __popc(vote);\n for ( int k = 0; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [b_col_it, b_col_end).\n int b_row_id = -1;\n\n if ( is_active_k )\n {\n b_row_id = s_b_row_ids[warp_id * WARP_SIZE + local_k];\n }\n\n // Load the range of B.\n int b_col_it = 0, b_col_end = 0;\n\n if ( is_active_k )\n {\n b_col_it = A_rows[b_row_id + 0];\n b_col_end = A_rows[b_row_id + 1];\n }\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != b_row_id;\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_off_diagonal && s_con[b_col_it] && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n count += __popc( utils::ballot( is_coarse_strongly_connected ) );\n }\n }\n }\n\n // Store the number of columns in each row.\n if ( lane_id == 0 )\n {\n C_hat_offsets[a_row_id] = count;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_c_hat_kernel( int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_end,\n int *__restrict C_hat,\n int *__restrict C_hat_pos,\n int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // Shared memory to vote.\n __shared__ volatile int s_b_row_ids[CTA_SIZE];\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Skip coarse rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id >= 0 || coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // Clear the set.\n set.clear();\n // Load the range of the row.\n __syncthreads();\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_begin = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n __syncthreads();\n\n // _iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = is_active && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n int item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n item = a_col_id;\n }\n\n set.insert( item, wk_status );\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n utils::syncwarp();\n\n int num_rows = __popc( vote );\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int uniform_b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n // Load the range of the row of B.\n int b_col_it = A_rows[uniform_b_row_id + 0];\n int b_col_end = A_rows[uniform_b_row_id + 1];\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any(b_col_it < b_col_end) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != uniform_b_row_id;\n // Is it a strongly connected node.\n is_strongly_connected = is_off_diagonal && s_con[b_col_it];\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_strongly_connected && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n int b_item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n b_item = b_col_id;\n }\n\n set.insert( b_item, wk_status );\n }\n }\n }\n\n int c_col_it = C_hat_start[a_row_id];\n int count = set.store_with_positions( &C_hat[c_col_it], &C_hat_pos[c_col_it] );\n\n if ( lane_id == 0 )\n {\n C_hat_end[a_row_id] = c_col_it + count;\n }\n }\n\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_c_hat_kernel( int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const int *__restrict C_hat_start,\n int *__restrict C_hat_end,\n int *__restrict C_hat,\n int *__restrict C_hat_pos,\n int gmem_size,\n int *g_keys,\n int *wk_work_queue,\n int *wk_status )\n{\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // Shared memory to vote.\n __shared__ volatile int s_b_row_ids[CTA_SIZE];\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id( );\n const int lane_id = utils::lane_id( );\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Skip coarse rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id >= 0 || coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // Clear the set.\n set.clear();\n\n // Load the range of the row.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_begin = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // _iterate over the columns of A to build C_hat.\n for ( int a_col_it = a_col_begin + lane_id ; utils::any(a_col_it < a_col_end) ; a_col_it += WARP_SIZE )\n {\n // Is it an active thread.\n const bool is_active = a_col_it < a_col_end;\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( is_active )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = is_active && a_col_id != a_row_id;\n // Is it strongly connected ?\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it fine.\n bool is_fine = is_off_diagonal && cf_map[a_col_id] == FINE;\n // Is it isolated.\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Is it a coarse and strongly-connected column.\n bool is_coarse_strongly_connected = is_strongly_connected && !is_fine && !is_strong_fine;\n // Push coarse and strongly connected nodes in the set.\n int item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n item = a_col_id;\n }\n\n set.insert( item, wk_status );\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && is_fine;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n utils::syncwarp();\n\n int num_rows = __popc( vote );\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [bColIt, bColEnd).\n int uniform_b_row_id = -1;\n\n if ( is_active_k )\n {\n uniform_b_row_id = s_b_row_ids[warp_id * WARP_SIZE + local_k];\n }\n\n // Load the range of the row of B.\n int b_col_it = 0, b_col_end = 0;\n\n if ( is_active_k )\n {\n b_col_it = A_rows[uniform_b_row_id + 0];\n b_col_end = A_rows[uniform_b_row_id + 1];\n }\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any(b_col_it < b_col_end) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it an off-diagonal element.\n is_off_diagonal = b_col_it < b_col_end && b_col_id != uniform_b_row_id;\n // Is it a strongly connected node.\n is_strongly_connected = is_off_diagonal && s_con[b_col_it];\n // Is it a coarse and strongly-connected column.\n is_coarse_strongly_connected = is_strongly_connected && (cf_map[b_col_id] != FINE && cf_map[b_col_id] != STRONG_FINE);\n // Push coarse and strongly connected nodes in the set.\n int b_item = -1;\n\n if ( is_coarse_strongly_connected )\n {\n b_item = b_col_id;\n }\n\n set.insert( b_item, wk_status );\n }\n }\n }\n\n int c_col_it = C_hat_start[a_row_id];\n int count = set.store_with_positions( &C_hat[c_col_it], &C_hat_pos[c_col_it] );\n\n if ( lane_id == 0 )\n {\n C_hat_end[a_row_id] = c_col_it + count;\n }\n }\n\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_inner_sum_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const int *__restrict C_hat,\n const int *__restrict C_hat_pos,\n const int *__restrict C_hat_start,\n const int *__restrict C_hat_end,\n const Value_type *__restrict diag,\n const int *__restrict inner_sum_offsets,\n Value_type *inner_sum,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue )\n{\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // A shared location where threads propose a row of B to load.\n __shared__ volatile int s_b_row_ids[CTA_SIZE];\n // A shared location where threads propose a value.\n __shared__ volatile Value_type s_a_values[CTA_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Skip coarse rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id >= 0 || coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // Clear the set.\n set.clear();\n // Rebuild C_hat.\n int c_hat_it = C_hat_start[a_row_id];\n int c_hat_end = C_hat_end [a_row_id];\n set.load( c_hat_end - c_hat_it, &C_hat[c_hat_it], &C_hat_pos[c_hat_it] );\n // The offset in the inner sum table.\n int inner_sum_offset = inner_sum_offsets[a_row_id];\n // And share the value of the diagonal.\n bool sign_diag = false;\n\n if ( lane_id == 0 )\n {\n sign_diag = sign( diag[a_row_id] );\n }\n\n sign_diag = utils::shfl( sign_diag, 0 );\n\n // Load A row IDs.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // _columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_off_diagonal && s_con[a_col_it] && cf_map[a_col_id] == FINE;\n // Read the associated value from A.\n Value_type a_value(0);\n\n if ( is_fine_strongly_connected )\n {\n a_value = A_vals[a_col_it];\n }\n\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n s_a_values [warp_id * WARP_SIZE + dest] = a_value;\n }\n\n utils::syncwarp();\n\n int num_rows = __popc( vote );\n // First n_rows threads reload the correct value.\n a_value = s_a_values[warp_id * WARP_SIZE + lane_id];\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Each thread keeps its own sum.\n Value_type bottom_sum(0);\n // Threads in the warp proceeds columns of B in the range [b_col_it, b_col_end).\n int b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n // TODO: make sure we have better memory accesses.\n int b_col_it = A_rows[b_row_id ];\n int b_col_end = A_rows[b_row_id + 1];\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id ; utils::any( b_col_it < b_col_end ) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it needed to compute the bottom sum.\n bool is_needed = set.contains( b_col_id ) || b_col_id == a_row_id;\n // Read the associated value from A.\n Value_type b_value(0);\n\n if ( is_needed )\n {\n b_value = A_vals[b_col_it];\n }\n\n // Update the sum if needed.\n if ( sign_diag != sign( b_value ) )\n {\n bottom_sum += b_value;\n }\n }\n\n // Reduce the row to a single value.\n#pragma unroll\n for ( int offset = WARP_SIZE / 2 ; offset > 0 ; offset >>= 1 )\n {\n int hi = __double2hiint(bottom_sum);\n hi = utils::shfl_xor( hi, offset );\n int lo = __double2loint(bottom_sum);\n lo = utils::shfl_xor( lo, offset );\n bottom_sum += __hiloint2double(hi, lo);\n }\n\n bottom_sum = utils::shfl( bottom_sum, 0 );\n\n if ( lane_id == k && bottom_sum != Value_type(0) )\n {\n a_value /= bottom_sum;\n }\n }\n\n // Write the results.\n if ( lane_id < num_rows )\n {\n inner_sum[inner_sum_offset + lane_id] = a_value;\n }\n\n inner_sum_offset += num_rows;\n }\n }\n\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< int NUM_THREADS_PER_ROW, typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_inner_sum_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const int *__restrict C_hat,\n const int *__restrict C_hat_pos,\n const int *__restrict C_hat_start,\n const int *__restrict C_hat_end,\n const Value_type *__restrict diag,\n const int *__restrict inner_sum_offsets,\n Value_type *inner_sum,\n const int gmem_size,\n int *g_keys,\n int *wk_work_queue )\n{\n\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n const int NUM_LOADED_ROWS = WARP_SIZE / NUM_THREADS_PER_ROW;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // A shared location where threads propose a row of B to load.\n __shared__ volatile int s_b_row_ids[CTA_SIZE];\n // A shared location where threads propose a value.\n __shared__ volatile Value_type s_a_values[CTA_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // Constants.\n const int lane_id_div_num_threads = lane_id / NUM_THREADS_PER_ROW;\n const int lane_id_mod_num_threads = lane_id % NUM_THREADS_PER_ROW;\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_set set( &s_keys[warp_id * SMEM_SIZE], &g_keys[a_row_id * gmem_size], gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n // Skip coarse rows.\n int coarse_fine_id = cf_map[a_row_id];\n\n if ( coarse_fine_id >= 0 || coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // Clear the set.\n set.clear();\n // Rebuild C_hat.\n int c_hat_it = C_hat_start[a_row_id];\n int c_hat_end = C_hat_end [a_row_id];\n set.load( c_hat_end - c_hat_it, &C_hat[c_hat_it], &C_hat_pos[c_hat_it] );\n // The offset in the inner sum table.\n int inner_sum_offset = inner_sum_offsets[a_row_id];\n // And share the value of the diagonal.\n bool sign_diag = false;\n\n if ( lane_id == 0 )\n {\n sign_diag = sign( diag[a_row_id] );\n }\n\n sign_diag = utils::shfl( sign_diag, 0 );\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // _columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id = -1;\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_off_diagonal && s_con[a_col_it] && cf_map[a_col_id] == FINE;\n // Read the associated value from A.\n Value_type a_value(0);\n\n if ( is_fine_strongly_connected )\n {\n a_value = A_vals[a_col_it];\n }\n\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n s_a_values [warp_id * WARP_SIZE + dest] = a_value;\n }\n\n int num_rows = __popc( vote );\n // First n_rows threads reload the correct value.\n a_value = s_a_values[warp_id * WARP_SIZE + lane_id];\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; k += NUM_LOADED_ROWS )\n {\n // Each thread keeps its own sum.\n Value_type bottom_sum(0);\n int local_k = k + lane_id_div_num_threads;\n // Is it an active thread.\n bool is_active_k = local_k < num_rows;\n // Threads in the warp proceeds columns of B in the range [b_col_it, b_col_end).\n int b_row_id = -1;\n\n if ( is_active_k )\n {\n b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k + lane_id_div_num_threads];\n }\n\n // TODO: make sure we have better memory accesses.\n int b_col_it = 0, b_col_end = 0;\n\n if ( is_active_k )\n {\n b_col_it = A_rows[b_row_id ];\n b_col_end = A_rows[b_row_id + 1];\n }\n\n // _iterate over the range of columns of B.\n for ( b_col_it += lane_id_mod_num_threads ; utils::any( b_col_it < b_col_end ) ; b_col_it += NUM_THREADS_PER_ROW )\n {\n // The ID of the column.\n int b_col_id = -1;\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n }\n\n // Is it needed to compute the bottom sum.\n bool is_needed = set.contains( b_col_id ) || b_col_id == a_row_id;\n // Read the associated value from A.\n Value_type b_value(0);\n\n if ( is_needed )\n {\n b_value = A_vals[b_col_it];\n }\n\n // Update the sum if needed.\n if ( sign_diag != sign( b_value ) )\n {\n bottom_sum += b_value;\n }\n }\n\n // Reduce the row to a single value.\n#pragma unroll\n for ( int offset = NUM_THREADS_PER_ROW / 2 ; offset > 0 ; offset >>= 1 )\n {\n int hi = __double2hiint(bottom_sum);\n hi = utils::shfl_down( hi, offset, NUM_THREADS_PER_ROW );\n int lo = __double2loint(bottom_sum);\n lo = utils::shfl_down( lo, offset, NUM_THREADS_PER_ROW );\n bottom_sum += __hiloint2double(hi, lo);\n }\n\n bottom_sum = utils::shfl( bottom_sum, lane_id_mod_num_threads * NUM_THREADS_PER_ROW );\n\n if ( lane_id >= k && lane_id < k + NUM_THREADS_PER_ROW && bottom_sum != Value_type(0) )\n {\n a_value /= bottom_sum;\n }\n }\n\n // Write the results.\n if ( lane_id < num_rows )\n {\n inner_sum[inner_sum_offset + lane_id] = a_value;\n }\n\n inner_sum_offset += num_rows;\n }\n }\n\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename Value_type, int CTA_SIZE, int SMEM_SIZE, int WARP_SIZE >\n__global__ __launch_bounds__( CTA_SIZE )\nvoid\ncompute_interp_weight_kernel( const int A_num_rows,\n const int *__restrict A_rows,\n const int *__restrict A_cols,\n const Value_type *__restrict A_vals,\n const int *__restrict cf_map,\n const bool *__restrict s_con,\n const int *__restrict C_hat,\n const int *__restrict C_hat_pos,\n const int *__restrict C_hat_start,\n const int *__restrict C_hat_end,\n const Value_type *__restrict diag,\n const int *__restrict inner_sum_offsets,\n const Value_type *__restrict inner_sum,\n const int *__restrict P_rows,\n int *P_cols,\n Value_type *P_vals,\n const int gmem_size,\n int *g_keys,\n Value_type *g_vals,\n int *wk_work_queue )\n{\n \n const int NUM_WARPS = CTA_SIZE / 32;\n // The hash keys stored in shared memory.\n __shared__ int s_keys[NUM_WARPS * SMEM_SIZE];\n // A shared location where threads propose a row of B to load.\n __shared__ volatile int s_b_row_ids[CTA_SIZE];\n // A shared location where threads propose a value.\n __shared__ volatile Value_type s_aki[NUM_WARPS];\n // The hash values stored in shared memory.\n __shared__ Value_type s_vals[NUM_WARPS * SMEM_SIZE];\n // The coordinates of the thread inside the CTA/warp.\n const int warp_id = utils::warp_id();\n const int lane_id = utils::lane_id();\n // First threads load the row IDs of A needed by the CTA...\n int a_row_id = blockIdx.x * NUM_WARPS + warp_id;\n // Create local storage for the set.\n Hash_map map( &s_keys[warp_id * SMEM_SIZE],\n &g_keys[a_row_id * gmem_size],\n &s_vals[warp_id * SMEM_SIZE],\n &g_vals[a_row_id * gmem_size],\n gmem_size );\n // Loop over rows of A.\n for ( ; a_row_id < A_num_rows ; a_row_id = get_work( wk_work_queue, warp_id ) )\n {\n int coarse_fine_id = cf_map[a_row_id];\n\n // Skip coarse rows.\n if ( coarse_fine_id >= 0 )\n {\n if ( lane_id == 0 )\n {\n int p_row_it = P_rows[a_row_id];\n P_cols[p_row_it] = coarse_fine_id;\n P_vals[p_row_it] = Value_type( 1 );\n }\n\n continue;\n }\n\n if ( coarse_fine_id == STRONG_FINE )\n {\n continue;\n }\n\n // Clear the table.\n map.clear();\n // Rebuild C_hat.\n int c_hat_it = C_hat_start[a_row_id];\n int c_hat_end = C_hat_end [a_row_id];\n map.load( c_hat_end - c_hat_it, &C_hat[c_hat_it], &C_hat_pos[c_hat_it] );\n // Load A row IDs.\n int a_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n a_col_tmp = A_rows[a_row_id + lane_id];\n }\n\n int a_col_it = utils::shfl( a_col_tmp, 0 );\n int a_col_end = utils::shfl( a_col_tmp, 1 );\n\n // The offset in the inner sum table.\n int inner_sum_offset = inner_sum_offsets[a_row_id];\n // Weak value.\n Value_type weak(0), sum(0);\n\n // Iterate over the columns of A.\n for ( a_col_it += lane_id ; utils::any( a_col_it < a_col_end ) ; a_col_it += WARP_SIZE )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int a_col_id(-1);\n Value_type a_value(0);\n\n if ( a_col_it < a_col_end )\n {\n a_col_id = A_cols[a_col_it];\n a_value = A_vals[a_col_it];\n }\n\n // Is it an off-diagonal element.\n bool is_off_diagonal = a_col_it < a_col_end && a_col_id != a_row_id;\n // Is it a strongly-connected column.\n bool is_strongly_connected = is_off_diagonal && s_con[a_col_it];\n // Is it a weakly connected node.\n bool is_weakly_connected = is_off_diagonal && !is_strongly_connected;\n // Is isolated\n bool is_strong_fine = is_off_diagonal && cf_map[a_col_id] == STRONG_FINE;\n // Update C_hat values. If the value is not in C_hat it will be skipped (true parameter).\n bool is_in_c_hat = map.update( a_col_id, a_value );\n\n // Update the weak value.\n if ( is_weakly_connected && !is_in_c_hat && !is_strong_fine)\n {\n weak += a_value;\n }\n\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_strongly_connected && cf_map[a_col_id] == FINE;\n // We collect fine and strongly-collected columns.\n int vote = utils::ballot( is_fine_strongly_connected );\n int dest = __popc( vote & utils::lane_mask_lt() );\n\n if ( is_fine_strongly_connected )\n {\n s_b_row_ids[warp_id * WARP_SIZE + dest] = a_col_id;\n }\n\n utils::syncwarp();\n\n int num_rows = __popc( vote );\n // We pre-load inner sums.\n sum = Value_type(0);\n\n if ( lane_id < num_rows )\n {\n sum = inner_sum[inner_sum_offset + lane_id];\n }\n\n inner_sum_offset += num_rows;\n\n // For each warp, we have up to 32 rows of B to proceed.\n for ( int k = 0 ; k < num_rows ; ++k )\n {\n // Threads in the warp proceeds columns of B in the range [b_col_it, bCol_end).\n int b_row_id = s_b_row_ids[warp_id * WARP_SIZE + k];\n // TODO: make sure we have better memory accesses. b_colBegin rather than bCol_it because we iterate twice.\n int b_col_begin = A_rows[b_row_id ];\n int b_col_end = A_rows[b_row_id + 1];\n // Diagonal element.\n Value_type b_diag = diag[b_row_id];\n\n // Set s_aki to 0 because a_row_id may not appear in a row.\n if ( lane_id == 0 )\n {\n s_aki[warp_id] = Value_type(0);\n }\n\n utils::syncwarp();\n\n // Load the kth inner sum.\n Value_type uniform_val = utils::shfl( sum, k );\n\n // _iterate over the range of columns of B.\n for ( int b_col_it = b_col_begin + lane_id ; utils::any( b_col_it < b_col_end ) ; b_col_it += WARP_SIZE )\n {\n // The ID of the column.\n int b_col_id(-1);\n Value_type b_value(0);\n\n if ( b_col_it < b_col_end )\n {\n b_col_id = A_cols[b_col_it];\n b_value = A_vals[b_col_it];\n }\n\n // Update bValue if needed.\n if ( sign( b_diag ) == sign( b_value ) )\n {\n b_value = Value_type(0);\n }\n\n // See if we have found a_row_id in the row. Only one thread can evaluate to true.\n if ( a_row_id == b_col_id )\n {\n s_aki[warp_id] = b_value;\n }\n\n // Update C_hat values. If the value is not in C_hat it will be skipped (true parameter).\n map.update( b_col_id, b_value * uniform_val );\n }\n\n utils::syncwarp();\n\n // The thread k updates the sum.\n if ( lane_id == 0 )\n {\n weak += s_aki[warp_id] * uniform_val;\n }\n }\n }\n\n // We're done with that row of A. We compute D.\n#pragma unroll\n\n for ( int mask = WARP_SIZE / 2 ; mask > 0 ; mask >>= 1 )\n {\n weak += utils::shfl_xor( weak, mask );\n }\n\n if ( lane_id == 0 )\n {\n weak += diag[a_row_id];\n sum = Value_type(-1) / weak;\n }\n\n sum = utils::shfl( sum, 0 );\n\n int p_col_tmp = -1;\n\n if ( lane_id < 2 )\n {\n p_col_tmp = P_rows[a_row_id + lane_id];\n }\n\n int p_col_it = utils::shfl( p_col_tmp, 0 );\n int p_col_end = utils::shfl( p_col_tmp, 1 );\n\n map.store_map_keys_scale_values( p_col_end - p_col_it, cf_map, &P_cols[p_col_it], sum, &P_vals[p_col_it] );\n }\n}\n\n\n} // namespace distance2\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate< typename IndexType, int kCtaSize >\n__global__\nvoid\ncalculateInnerSumStorageKernel( const IndexType *A_rows,\n const IndexType *A_cols,\n const int A_num_rows,\n const int *cf_map,\n const bool *s_con,\n int *innerSum_offsets )\n{\n const int nWarps = kCtaSize / 32;\n // The coordinates of the thread inside the CTA/warp.\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n // Loop over rows of A.\n for ( int aRowId = blockIdx.x * nWarps + warpId ; aRowId < A_num_rows ; aRowId += gridDim.x * nWarps )\n {\n // Skip coarse rows.\n int coarse_fine_id = cf_map[aRowId];\n\n if ( coarse_fine_id >= 0 || coarse_fine_id == STRONG_FINE)\n {\n continue;\n }\n\n // The sum for each thread.\n int sum = 0;\n // Load A row IDs.\n int aColIt = A_rows[aRowId ];\n int aColEnd = A_rows[aRowId + 1];\n\n // Iterate over the columns of A.\n for ( aColIt += laneId ; utils::any( aColIt < aColEnd ) ; aColIt += 32 )\n {\n // Columns of A maps to rows of B. Each thread of the warp loads its A-col/B-row ID.\n int aColId = aColIt < aColEnd ? A_cols[aColIt] : -1;\n // Is it an off-diagonal element.\n bool is_off_diagonal = aColIt < aColEnd && aColId != aRowId;\n // Is it a fine and strongly-connected column.\n bool is_fine_strongly_connected = is_off_diagonal && s_con[aColIt] && cf_map[aColId] == FINE;\n // Update the sum.\n sum += __popc( utils::ballot( is_fine_strongly_connected ) );\n }\n\n // Update the value.\n if ( laneId == 0 )\n {\n innerSum_offsets[aRowId] = sum;\n }\n }\n}\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n/*\n * grab the diagonal of the matrix\n */\ntemplate \n__global__\nvoid getDiagonalKernel(const IndexType *offsets, const IndexType *column_indices,\n const ValueType *values, const IndexType numRows, ValueType *diagonal)\n{\n for (int tIdx = threadIdx.x + blockDim.x * blockIdx.x; tIdx < numRows; tIdx += blockDim.x * gridDim.x)\n {\n const int offset = offsets[tIdx];\n const int numj = offsets[tIdx + 1] - offset;\n\n for (int j = offset; j < offset + numj; j++)\n {\n int jcol = column_indices[j];\n\n if (tIdx == jcol)\n {\n diagonal[tIdx] = values[j];\n }\n }\n }\n}\n\n/*************************************************************************\n * Implementing Extended+i algorithm from \\S 4.5 of:\n * \"Distance-two interpolation for parallel algebraic multigrid\"\n * Reference [4] on wiki\n ************************************************************************/\ntemplate \nDistance2_Interpolator >::Distance2_Interpolator(AMG_Config &cfg, const std::string &cfg_scope) : Base(cfg, cfg_scope)\n{}\n\ntemplate \nDistance2_Interpolator >::~Distance2_Interpolator()\n{}\n\nenum { WARP_SIZE = 32, SMEM_SIZE = 128 };\n\ntemplate \nvoid Distance2_Interpolator >::generateInterpolationMatrix_1x1(Matrix_d &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix_d &P)\n{\n nvtxRange gim_nr(__func__);\n\n const int blockSize = 256;\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // get raw pointers to data\n const IndexType *Aoffsets = A.row_offsets.raw();\n const IndexType *Acolumn_indices = A.col_indices.raw();\n const ValueType *Avalues = A.values.raw();\n const IndexType Anum_rows = (int) A.get_num_rows();\n\n hipDeviceProp_t props = getDeviceProperties();\n int grid_size = (props.major >= 7) ? 1024 : 128;\n int gmem_size = (props.major >= 7) ? 512 : 2048;\n\n Hash_Workspace exp_wk(true, grid_size, 8, gmem_size);\n\n IntVector C_hat_start( A.get_num_rows() + 1, 0 ), C_hat_end( A.get_num_rows() + 1, 0 );\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int work_offset = grid_size * NUM_WARPS;\n hipMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), hipMemcpyHostToDevice );\n cudaCheckError();\n int avg_nz_per_row = (A.get_num_rows() == 0) ? 0 : A.get_num_nz() / A.get_num_rows();\n\n if ( avg_nz_per_row < 16 )\n {\n distance2::estimate_c_hat_size_kernel< 8, CTA_SIZE, WARP_SIZE> <<< 2048, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw());\n cudaCheckError();\n }\n else\n {\n distance2::estimate_c_hat_size_kernel <<< 2048, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw());\n }\n\n cudaCheckError();\n }\n // Compute row offsets.\n thrust_wrapper::exclusive_scan( C_hat_start.begin( ), C_hat_start.end( ), C_hat_start.begin( ) );\n cudaCheckError();\n // Allocate memory to store columns/values.\n int nVals = C_hat_start[C_hat_start.size() - 1];\n IntVector C_hat( nVals ), C_hat_pos( nVals );\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int avg_nz_per_row = (A.get_num_rows() == 0) ? 0 : A.get_num_nz() / A.get_num_rows();\n int attempt = 0;\n\n for ( bool done = false ; !done && attempt < 10 ; ++attempt )\n {\n // Double the amount of GMEM (if needed).\n if ( attempt > 0 )\n {\n exp_wk.expand();\n }\n\n // Reset the status. TODO: Launch async copies.\n int status = 0;\n hipMemcpy( exp_wk.get_status(), &status, sizeof(int), hipMemcpyHostToDevice );\n cudaCheckError();\n // Compute the set C_hat.\n int work_offset = grid_size * NUM_WARPS;\n hipMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), hipMemcpyHostToDevice );\n cudaCheckError();\n\n // Run the computation.\n if ( avg_nz_per_row < 16 )\n {\n distance2::compute_c_hat_kernel< 8, CTA_SIZE, SMEM_SIZE, WARP_SIZE> <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n C_hat.raw(),\n C_hat_pos.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_work_queue(),\n exp_wk.get_status() );\n cudaCheckError();\n }\n else\n {\n distance2::compute_c_hat_kernel <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n C_hat.raw(),\n C_hat_pos.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_work_queue(),\n exp_wk.get_status() );\n }\n\n cudaCheckError();\n // Read the result from count_non_zeroes.\n hipMemcpy( &status, exp_wk.get_status(), sizeof(int), hipMemcpyDeviceToHost );\n cudaCheckError();\n done = status == 0;\n }\n }\n // get pointers to data\n int *C_hat_ptr = C_hat.raw();\n int *C_hat_start_ptr = C_hat_start.raw();\n int *C_hat_end_ptr = C_hat_end.raw();\n // modify Coarse/Fine map array to have the column index of the new interpolation matrix\n // for evry coarse point\n int coarsePoints;\n DistributedArranger *prep;\n\n if (A.is_matrix_distributed())\n {\n prep = new DistributedArranger;\n int num_owned_fine_pts = A.get_num_rows();\n int num_owned_coarse_pts = thrust_wrapper::count_if(cf_map.begin(), cf_map.begin() + num_owned_fine_pts, is_non_neg());\n cudaCheckError();\n int num_halo_coarse_pts = thrust_wrapper::count_if(cf_map.begin() + num_owned_fine_pts, cf_map.end(), is_non_neg());\n cudaCheckError();\n coarsePoints = num_owned_coarse_pts + num_halo_coarse_pts;\n // Using the number of owned_coarse_pts (number of rows of Ac), initialize P manager\n // This will set the part_offsets array, index_range, halo_offsets[0] value\n prep->initialize_manager(A, P, num_owned_coarse_pts);\n }\n else\n {\n coarsePoints = (int) thrust_wrapper::count_if(cf_map.begin(), cf_map.end(), is_non_neg());\n cudaCheckError();\n }\n\n // mark as 1 if a coarse point\n int size_one_ring;\n\n if (!A.is_matrix_singleGPU())\n {\n int offset;\n A.getOffsetAndSizeForView(FULL, &offset, &size_one_ring);\n }\n else\n {\n size_one_ring = A.get_num_rows();\n }\n\n int numBlocksOneRing = std::min( 4096, (int) (size_one_ring + blockSize - 1) / blockSize );\n int numBlocks = std::min( 4096, (int) (A.get_num_rows() + blockSize - 1) / blockSize );\n // count the number of non-zeros in the interpolation matrix\n IntVector nonZeroOffsets(A.get_num_rows() + 1);\n IntVector nonZerosPerRow(A.get_num_rows());\n int *nonZerosPerRow_ptr = nonZerosPerRow.raw();\n\n if (A.get_num_rows() > 0)\n {\n nonZerosPerRowKernel <<< numBlocks, blockSize>>>(Anum_rows, cf_map.raw(), C_hat_start_ptr, C_hat_end_ptr,\n nonZerosPerRow_ptr);\n cudaCheckError();\n }\n\n // get total with a reduction\n int nonZeros = thrust_wrapper::reduce(nonZerosPerRow.begin(), nonZerosPerRow.end());\n cudaCheckError();\n // get the offsets with an exclusive scan\n thrust_wrapper::exclusive_scan(nonZerosPerRow.begin(), nonZerosPerRow.end(), nonZeroOffsets.begin());\n cudaCheckError();\n nonZeroOffsets[A.get_num_rows()] = nonZeros;\n // resize P\n P.resize(0, 0, 0, 1);\n P.addProps(CSR);\n P.resize(A.get_num_rows(), coarsePoints, nonZeros, 1);\n IndexType *Poffsets = P.row_offsets.raw();\n IndexType *Pcolumn_indices = P.col_indices.raw();\n ValueType *Pvalues = P.values.raw();\n // copy nonzero offsets to the P matrix\n amgx::thrust::copy(nonZeroOffsets.begin(), nonZeroOffsets.end(), P.row_offsets.begin());\n cudaCheckError();\n // grab the diagonal terms\n VVector diag(size_one_ring);\n ValueType *diag_ptr = diag.raw();\n\n if (A.get_num_rows() > 0)\n {\n find_diag_kernel_indexed_dia <<< numBlocksOneRing, blockSize>>>(\n size_one_ring,\n A.diag.raw(),\n A.values.raw(),\n diag.raw());\n }\n\n cudaCheckError();\n IntVector innerSumOffset( A.get_num_rows() + 1, 0 );\n\n if (A.get_num_rows() > 0)\n {\n const int kCtaSize = 256;\n const int nWarps = kCtaSize / 32;\n const int grid_size = std::min( 4096, ( Anum_rows + nWarps - 1 ) / nWarps );\n calculateInnerSumStorageKernel <<< grid_size, kCtaSize>>>(\n Aoffsets,\n Acolumn_indices,\n Anum_rows,\n cf_map.raw(),\n s_con.raw(),\n innerSumOffset.raw() );\n cudaCheckError();\n }\n\n // get the offsets\n thrust_wrapper::exclusive_scan( innerSumOffset.begin(), innerSumOffset.end(), innerSumOffset.begin() );\n cudaCheckError();\n // assign memory & get pointer\n int numInnerSum = innerSumOffset[A.get_num_rows()];\n VVector innerSum( numInnerSum, 0 );\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int avg_nz_per_row = (A.get_num_rows() == 0) ? 0 : A.get_num_nz() / A.get_num_rows();\n // Compute the set C_hat.\n int work_offset = grid_size * NUM_WARPS;\n hipMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), hipMemcpyHostToDevice );\n cudaCheckError();\n // Run the computation.\n typedef typename MatPrecisionMap::Type Value_type;\n {\n distance2::compute_inner_sum_kernel <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat.raw(),\n C_hat_pos.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n diag.raw(),\n innerSumOffset.raw(),\n innerSum.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_work_queue());\n }\n cudaCheckError();\n }\n {\n const int CTA_SIZE = 256;\n const int NUM_WARPS = CTA_SIZE / WARP_SIZE;\n int avg_nz_per_row = (A.get_num_rows() == 0) ? 0 : A.get_num_nz() / A.get_num_rows();\n // Compute the set C_hat.\n int work_offset = grid_size * NUM_WARPS;\n hipMemcpy( exp_wk.get_work_queue(), &work_offset, sizeof(int), hipMemcpyHostToDevice );\n cudaCheckError();\n // Run the computation.\n typedef typename MatPrecisionMap::Type Value_type;\n distance2::compute_interp_weight_kernel <<< grid_size, CTA_SIZE>>>(\n A.get_num_rows(),\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n cf_map.raw(),\n s_con.raw(),\n C_hat.raw(),\n C_hat_pos.raw(),\n C_hat_start.raw(),\n C_hat_end.raw(),\n diag.raw(),\n innerSumOffset.raw(),\n innerSum.raw(),\n P.row_offsets.raw(),\n P.col_indices.raw(),\n P.values.raw(),\n exp_wk.get_gmem_size(),\n exp_wk.get_keys(),\n exp_wk.get_vals(),\n exp_wk.get_work_queue());\n cudaCheckError();\n }\n\n if (A.is_matrix_distributed())\n {\n //TODO: cf_map_global could be smaller (only contain halo nodes)\n I64Vector_d cf_map_global(cf_map.size());\n int num_owned_fine_pts = P.get_num_rows();\n int my_rank = A.manager->global_id();\n const int cta_size = 128;\n const int grid_size = std::min( 4096, (num_owned_fine_pts + cta_size - 1) / cta_size);\n\n // Write the global coarse index of all the owned fine points (-1 if not a coarse point)\n if (num_owned_fine_pts > 0)\n {\n createCfMapGlobal <<< grid_size, cta_size>>>(cf_map.raw(), cf_map_global.raw(), P.manager->part_offsets_h[my_rank], num_owned_fine_pts);\n cudaCheckError();\n }\n\n // Exchange the cf_map_global so that we know the coarse global id of halo nodes\n cf_map_global.dirtybit = 1;\n A.manager->exchange_halo_2ring(cf_map_global, cf_map_global.tag);\n I64Vector_d P_col_indices_global(P.col_indices.size());\n\n if (num_owned_fine_pts > 0)\n {\n fill_P_global_col_indices_kernel<16, int> <<< grid_size, cta_size>>>(P.row_offsets.raw(), C_hat_start.raw(), C_hat_end.raw(), C_hat.raw(), cf_map_global.raw(), P_col_indices_global.raw(), num_owned_fine_pts);\n cudaCheckError();\n }\n\n // Now that we know the global col indices of P, we can find:\n // the list of neighbors, create the B2L_maps for P, create local numbering for halo nodes, create local_to_global_map array, etc...\n prep->initialize_manager_from_global_col_indices(P, P_col_indices_global);\n prep->createRowsLists(P, true);\n delete prep;\n }\n}\n\ntemplate< class T_Config>\nvoid Distance2_InterpolatorBase::generateInterpolationMatrix(Matrix &A,\n IntVector &cf_map,\n BVector &s_con,\n IntVector &scratch,\n Matrix &P)\n{\n P.set_initialized(0);\n ViewType oldView = A.currentView();\n A.setView(OWNED);\n\n // If multi-GPU do a 2-ring exchange of cf_map since this is a distance two interpolator\n if (A.get_block_size() == 1)\n {\n generateInterpolationMatrix_1x1(A, cf_map, s_con, scratch, P);\n }\n else\n {\n FatalError(\"Unsupported dimensions for distance2 interpolator\", AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n }\n\n A.setView(oldView);\n P.set_initialized(1);\n}\n\n#define AMGX_CASE_LINE(CASE) template class Distance2_InterpolatorBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Distance2_Interpolator::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n", "source": "/kaggle/working/amgx/src/classical/interpolators/distance2.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n\n/*\n* Note:\n* This implementation assumes that off-diag entries all have the opposite sign\n* comparing to the diag entry. This is true for most practical cases.\n* It would even work if the offending off-diag entries are just a few.\n* But if there are many off-diag entries violate this assumption,\n* the interpolation based on this strength would be inaccurate.\n* This is explained in \"Intro to Algebraic multigrid\" by K. Stuben.\n*/\n\n\nnamespace amgx\n{\n\ntemplate< typename T_Config >\nStrength_BaseBase::Strength_BaseBase(AMG_Config &cfg,\n const std::string &cfg_scope)\n{\n alpha = cfg.AMG_Config::template getParameter(\"strength_threshold\", cfg_scope);\n m_use_opt_kernels = cfg.AMG_Config::template getParameter(\"use_opt_kernels\", \"default\");\n}\n\n/*************************************************************************\n* \"random\" hash function for both device and host\n************************************************************************/\n__host__ __device__ __forceinline__\nstatic float ourHash(const int i)\n{\n unsigned int a = i;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return (a ^ 0x4a51e590) / (float)UINT_MAX;\n}\n\n/*************************************************************************\n* Computes the strength matrix and the connection weights\n* Described in \\S 4.1 of:\n* \"Reducing complexity in parallel algebraic multigrid preconditioners\"\n*\n************************************************************************/\ntemplate \nvoid\nStrength_Base >::\ncomputeStrongConnectionsAndWeights_1x1(Matrix_h &A,\n BVector &s_con,\n FVector &weights,\n const double max_row_sum)\n{\n bool compute_row_sum = (max_row_sum < 1.0);\n VVector sums_ptr;\n\n // get the (normalised) row sums\n if (compute_row_sum)\n {\n sums_ptr.resize(A.get_num_rows());\n weightedRowSum(A, sums_ptr);\n cudaCheckError();\n }\n\n // get min/max off-diag, depending on sign of diagonal\n for (int row = 0; row < A.get_num_rows(); row++)\n {\n ValueType diag(0), minVal(0), maxVal(0);\n ValueType row_sum = compute_row_sum ? sums_ptr[row] : -1;\n int rowEnd = A.row_offsets[row + 1];\n\n for (int j = A.row_offsets[row]; j < rowEnd; j++)\n {\n int col = A.col_indices[j];\n ValueType val = A.values[j];\n\n if (col == row)\n {\n diag = val;\n }\n else\n {\n minVal = std::min(minVal, val);\n maxVal = std::max(maxVal, val);\n }\n }\n\n //set the threshold for being strongly connected\n ValueType threshold = (diag < 0) ? maxVal : minVal;\n threshold *= this->alpha;\n\n //initialize the weight to a small random number\n if (A.is_matrix_singleGPU())\n {\n weights[row] += ourHash(row);\n }\n else\n {\n weights[row] += ourHash((int)A.manager->base_index() + row);\n }\n\n // sum the column of S - # of points that strongly connect to me\n for (int j = A.row_offsets[row]; j < rowEnd; j++)\n {\n bool is_strongly_connected = false;\n\n if (compute_row_sum && row_sum > max_row_sum)\n {\n is_strongly_connected = false;\n }\n else\n is_strongly_connected =\n this->strongly_connected(A.values[j], threshold, diag);\n\n int col = A.col_indices[j];\n s_con[j] = (col != row) && is_strongly_connected;\n weights[A.col_indices[j]] += s_con[j] ? 1. : 0.;\n }\n }\n}\n\n\ntemplate \nvoid\nStrength_Base >::\ncomputeWeights_1x1(Matrix_h &S,\n FVector &weights)\n{\n for (int row = 0; row < S.get_num_rows(); row++)\n {\n //initialize the weight to a small random number\n if (S.is_matrix_singleGPU())\n {\n weights[row] += ourHash(row);\n }\n else\n {\n weights[row] += ourHash(S.manager->base_index() + row);\n }\n\n int rowEnd = S.row_offsets[row + 1];\n\n for (int j = S.row_offsets[row]; j < rowEnd; j++)\n {\n int col = S.col_indices[j];\n\n if (col != row) { weights[col] += 1.; }\n }\n }\n}\n\n\ntemplate \n__device__\nbool stronglyConnectedAHat(ValueType value, ValueType threshold, ValueType diag)\n{\n return (diag < ValueType(0)) ? value > threshold : value < threshold;\n}\n\n/*************************************************************************\n* Computes the strength matrix and the connection weights (device)\n************************************************************************/\ntemplate< typename IndexType, typename ValueType, int kCtaSize, bool singleGPU >\n__global__\nvoid computeStrongConnectionsAndWeightsKernel( const IndexType *A_rows,\n const IndexType *A_cols,\n const ValueType *A_vals,\n int A_num_rows,\n bool *s_con,\n float *weights,\n ValueType alpha,\n ValueType *row_sum,\n const double max_row_sum,\n int64_t base_index)\n{\n // One warp works on each row and hence one iteration handles\n // num_warps*numBlock rows. This means atomicAdd() is inevitable.\n const int num_warps = kCtaSize / 32;\n const int num_rows_per_iter = num_warps * gridDim.x;\n __shared__ volatile ValueType smem[kCtaSize];\n __shared__ volatile ValueType s_diag[num_warps];\n __shared__ volatile ValueType s_threshold[num_warps];\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n for ( int aRowId = blockIdx.x * num_warps + warpId ; aRowId < A_num_rows ;\n aRowId += num_rows_per_iter )\n {\n ValueType minVal(0), maxVal(0);\n\n if ( laneId == 0 ) // Reset the diagonal\n {\n s_diag[warpId] = ValueType(0);\n }\n\n utils::syncwarp();\n\n // Row sum\n ValueType rowSum = -1.0;\n\n if (max_row_sum < 1.0) { rowSum = row_sum[aRowId]; }\n\n // get diagonal, min/max off-diagonals\n const int aRowBegin = A_rows[aRowId ];\n const int aRowEnd = A_rows[aRowId + 1];\n\n for ( IndexType aRowIt = aRowBegin + laneId ; utils::any( aRowIt < aRowEnd ) ;\n aRowIt += 32 )\n {\n IndexType aColId = aRowIt < aRowEnd ? A_cols[aRowIt] : -1;\n ValueType aValue = aRowIt < aRowEnd ? A_vals[aRowIt] : ValueType(0);\n\n if ( aColId == aRowId ) // only one thread evaluates to true.\n {\n s_diag[warpId] = aValue;\n }\n\n bool is_off_diagonal = aRowIt < aRowEnd && aColId != aRowId;\n\n if ( is_off_diagonal )\n {\n minVal = min( minVal, aValue );\n maxVal = max( maxVal, aValue );\n }\n }\n\n // init weights[] with a random number\n if ( laneId == 0 )\n {\n if ( singleGPU )\n {\n atomicAdd( &weights[aRowId], ourHash(aRowId) );\n }\n else\n {\n atomicAdd( &weights[aRowId], ourHash( (int) base_index + aRowId) );\n }\n }\n\n utils::syncwarp();\n\n // Big assumption: diag and off-diag always have the opposite sign.\n // If diag entry is negative, then all off-diag entries must be positive.\n // This means max off-diag is to be used to compute the threshold.\n // If diag entry is positve, the min off-diag is used instead.\n if ( s_diag[warpId] < ValueType(0) )\n {\n smem[threadIdx.x] = maxVal;\n\n utils::syncwarp();\n\n#pragma unroll\n for ( int offset = 16 ; offset > 0 ; offset /= 2 )\n {\n if ( laneId < offset )\n {\n smem[threadIdx.x] = maxVal = max( maxVal, smem[threadIdx.x + offset] );\n }\n utils::syncwarp();\n }\n }\n else\n {\n smem[threadIdx.x] = minVal;\n\n utils::syncwarp();\n\n#pragma unroll\n for ( int offset = 16 ; offset > 0 ; offset /= 2 )\n {\n if ( laneId < offset )\n {\n smem[threadIdx.x] = minVal = min( minVal, smem[threadIdx.x + offset] );\n }\n utils::syncwarp();\n }\n }\n\n if ( laneId == 0 )\n {\n // If laneId=0, then maxVal or minVal is in smem[threadIdx.x].\n s_threshold[warpId] = smem[threadIdx.x] * alpha;\n }\n\n utils::syncwarp();\n\n // sum of the column of S\n for ( IndexType aRowIt = aRowBegin + laneId ; utils::any( aRowIt < aRowEnd ) ;\n aRowIt += 32 )\n {\n IndexType aColId = aRowIt < aRowEnd ? A_cols[aRowIt] : -1;\n ValueType aValue = aRowIt < aRowEnd ? A_vals[aRowIt] : ValueType(0);\n bool is_strongly_connected = false;\n\n if (max_row_sum < 1.0 && rowSum > max_row_sum)\n {\n is_strongly_connected = false;\n }\n else\n {\n bool is_off_diagonal = aRowIt < aRowEnd && aColId != aRowId;\n is_strongly_connected = is_off_diagonal &&\n stronglyConnectedAHat( aValue, s_threshold[warpId], s_diag[warpId] );\n }\n\n if ( is_strongly_connected && aRowIt < aRowEnd && aColId < A_num_rows)\n {\n atomicAdd( &weights[aColId], 1.0f );\n }\n\n if ( aRowIt < aRowEnd )\n {\n s_con[aRowIt] = is_strongly_connected;\n }\n }\n }\n}\n\n/*************************************************************************\n* Computes the strength matrix and the connection weights (device)\n************************************************************************/\ntemplate< typename IndexType, typename ValueType, int kCtaSize, bool singleGPU >\n__global__\nvoid computeStrongConnectionsAndWeightsKernel_opt( const IndexType *A_rows,\n const IndexType *A_cols,\n const ValueType *A_vals,\n int A_num_rows,\n bool *s_con,\n float *weights,\n ValueType alpha,\n ValueType *row_sum,\n const double max_row_sum,\n int64_t base_index)\n{\n // One warp works on each row and hence one iteration handles\n // num_warps*numBlock rows. This means atomicAdd() is inevitable.\n int aRowId = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (aRowId < A_num_rows)\n {\n ValueType minVal(0), maxVal(0);\n\n ValueType diag = ValueType(0);\n\n utils::syncwarp();\n\n // Row sum\n ValueType rowSum = -1.0;\n\n if (max_row_sum < 1.0) { rowSum = row_sum[aRowId]; }\n\n // get diagonal, min/max off-diagonals\n const int aRowBegin = A_rows[aRowId ];\n const int aRowEnd = A_rows[aRowId + 1];\n\n for ( IndexType aRowIt = aRowBegin; aRowIt < aRowEnd ; ++aRowIt)\n {\n IndexType aColId = A_cols[aRowIt];\n ValueType aValue = A_vals[aRowIt];\n\n if ( aColId == aRowId ) // only one thread evaluates to true.\n {\n diag = aValue;\n }\n\n if ( aColId != aRowId )\n {\n minVal = min( minVal, aValue );\n maxVal = max( maxVal, aValue );\n }\n }\n\n if ( singleGPU )\n {\n atomicAdd(&weights[aRowId], ourHash(aRowId));\n }\n else\n {\n atomicAdd(&weights[aRowId], ourHash( (int) base_index + aRowId));\n }\n\n utils::syncwarp();\n\n // Big assumption: diag and off-diag always have the opposite sign.\n // If diag entry is negative, then all off-diag entries must be positive.\n // This means max off-diag is to be used to compute the threshold.\n // If diag entry is positve, the min off-diag is used instead.\n ValueType threshold;\n if ( diag < ValueType(0) )\n {\n threshold = maxVal*alpha;\n }\n else\n {\n threshold = minVal*alpha;\n }\n\n utils::syncwarp();\n\n // sum of the column of S\n for ( IndexType aRowIt = aRowBegin; aRowIt < aRowEnd ; ++aRowIt)\n {\n IndexType aColId = A_cols[aRowIt];\n ValueType aValue = A_vals[aRowIt];\n bool is_strongly_connected = false;\n\n if (max_row_sum < 1.0 && rowSum > max_row_sum)\n {\n is_strongly_connected = false;\n }\n else\n {\n bool is_off_diagonal = aColId != aRowId;\n is_strongly_connected = is_off_diagonal &&\n stronglyConnectedAHat( aValue, threshold, diag );\n }\n\n if ( is_strongly_connected && aColId < A_num_rows)\n {\n atomicAdd( &weights[aColId], 1.0f );\n }\n\n s_con[aRowIt] = is_strongly_connected;\n }\n }\n}\n\n\ntemplate< typename IndexType, typename ValueType, int kCtaSize, bool singleGPU >\n__global__\nvoid computeWeightsKernel( const IndexType *A_rows,\n const IndexType *A_cols,\n int A_num_rows,\n float *weights,\n int64_t base_index)\n{\n const int num_warps = kCtaSize / 32;\n const int num_rows_per_iter = num_warps * gridDim.x;\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n for ( int aRowId = blockIdx.x * num_warps + warpId ; aRowId < A_num_rows ;\n aRowId += num_rows_per_iter )\n {\n if ( laneId == 0 )\n {\n if ( singleGPU )\n {\n atomicAdd( &weights[aRowId], ourHash(aRowId) );\n }\n else\n {\n atomicAdd( &weights[aRowId], ourHash( (int) base_index + aRowId) );\n }\n }\n\n const int aRowBegin = A_rows[aRowId ];\n const int aRowEnd = A_rows[aRowId + 1];\n\n for ( IndexType aRowIt = aRowBegin + laneId ; utils::any( aRowIt < aRowEnd ) ;\n aRowIt += 32 )\n {\n IndexType aColId = aRowIt < aRowEnd ? A_cols[aRowIt] : -1;\n bool is_off_diagonal = aRowIt < aRowEnd && aColId != aRowId;\n\n if (is_off_diagonal)\n {\n atomicAdd( &weights[aColId], 1.0f );\n }\n }\n }\n}\n\n\n/*************************************************************************\n* Computes the strength matrix and the connection weights (device)\n************************************************************************/\ntemplate \nvoid Strength_Base >::\ncomputeStrongConnectionsAndWeights_1x1(Matrix_d &A,\n BVector &s_con,\n FVector &weights,\n const double max_row_sum)\n{\n typedef TemplateConfig TConfig_d;\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // get the raw pointers for everything I need\n const IndexType *offsets_ptr = A.row_offsets.raw();\n const IndexType *column_indices_ptr = A.col_indices.raw();\n const ValueType *values_ptr = A.values.raw();\n bool *s_con_ptr = s_con.raw();\n float *weights_ptr = weights.raw();\n bool compute_row_sum = (max_row_sum < 1.0);\n\n if (A.get_num_rows() == 0) { compute_row_sum = false; }\n\n VVector sums_ptr;\n\n if (compute_row_sum)\n {\n sums_ptr.resize(A.get_num_rows());\n weightedRowSum(A, sums_ptr);\n cudaCheckError();\n }\n\n // choose a blocksize. Use 1 warp per row\n const int blockSize = 256;\n const int numWarps = blockSize / 32;\n const int numBlocks = std::min( 4096, (int) (A.get_num_rows() + numWarps - 1) / numWarps );\n\n if (A.get_num_rows() > 0)\n {\n if(this->m_use_opt_kernels)\n {\n // One row per thread\n const int blockSize = 128;\n const int numBlocks = A.get_num_rows() / blockSize + 1;\n\n if (A.is_matrix_singleGPU()) {\n computeStrongConnectionsAndWeightsKernel_opt\n <<< numBlocks, blockSize>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n A.get_num_rows(),\n s_con.raw(),\n weights.raw(),\n this->alpha,\n compute_row_sum ? sums_ptr.raw() : NULL,\n max_row_sum,\n 0);\n cudaCheckError();\n }\n else {\n computeStrongConnectionsAndWeightsKernel_opt\n <<< numBlocks, blockSize>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n A.get_num_rows(),\n s_con.raw(),\n weights.raw(),\n this->alpha,\n compute_row_sum ? sums_ptr.raw() : NULL,\n max_row_sum,\n A.manager->base_index());\n cudaCheckError();\n }\n }\n else\n {\n const int blockSize = 256;\n const int numWarps = blockSize / 32;\n const int numBlocks = min( 4096, (int) (A.get_num_rows() + numWarps - 1) / numWarps );\n\n if (A.is_matrix_singleGPU()) {\n computeStrongConnectionsAndWeightsKernel\n <<< numBlocks, blockSize>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n A.get_num_rows(),\n s_con.raw(),\n weights.raw(),\n this->alpha,\n compute_row_sum ? sums_ptr.raw() : NULL,\n max_row_sum,\n 0);\n cudaCheckError();\n }\n else {\n computeStrongConnectionsAndWeightsKernel\n <<< numBlocks, blockSize>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n A.get_num_rows(),\n s_con.raw(),\n weights.raw(),\n this->alpha,\n compute_row_sum ? sums_ptr.raw() : NULL,\n max_row_sum,\n A.manager->base_index());\n cudaCheckError();\n }\n }\n }\n\n if (!A.is_matrix_singleGPU() && A.currentView() == OWNED)\n {\n // Need to add neighbors contribution to my weights\n weights.dirtybit = 1;\n A.manager->add_from_halo(weights, weights.tag);\n }\n\n cudaCheckError();\n}\n\n\ntemplate \nvoid Strength_Base >::\ncomputeWeights_1x1(Matrix_d &S,\n FVector &weights)\n{\n typedef TemplateConfig TConfig_d;\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // get the raw pointers for everything I need\n const IndexType *offsets_ptr = S.row_offsets.raw();\n const IndexType *column_indices_ptr = S.col_indices.raw();\n float *weights_ptr = weights.raw();\n // choose a blocksize. Use 1 thread per row\n const int blockSize = 256;\n const int numWarps = blockSize / 32;\n const int numBlocks = std::min( 4096, (int) (S.get_num_rows() + numWarps - 1) / numWarps );\n cudaDeviceSynchronize();\n cudaCheckError();\n\n // call the CUDA kernel\n if (S.is_matrix_singleGPU())\n computeWeightsKernel\n <<< numBlocks, blockSize>>>(\n S.row_offsets.raw(),\n S.col_indices.raw(),\n S.get_num_rows(),\n weights.raw(),\n 0);\n else\n computeWeightsKernel\n <<< numBlocks, blockSize>>>(\n S.row_offsets.raw(),\n S.col_indices.raw(),\n S.get_num_rows(),\n weights.raw(),\n S.manager->base_index());\n\n cudaCheckError();\n\n if (!S.is_matrix_singleGPU() && S.currentView() == OWNED)\n {\n // Need to add neighbors contribution to my weights\n weights.dirtybit = 1;\n S.manager->add_from_halo(weights, weights.tag);\n }\n\n cudaCheckError();\n}\n\ntemplate\nvoid Strength_BaseBase::\ncomputeStrongConnectionsAndWeights(Matrix &A,\n BVector &s_con,\n FVector &weights,\n const double max_row_sum\n )\n{\n if (A.get_block_size() == 1)\n {\n computeStrongConnectionsAndWeights_1x1(A, s_con, weights, max_row_sum);\n }\n else\n FatalError(\"Unsupported block size for strong connections\",\n AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate\nvoid Strength_BaseBase::computeWeights(Matrix &S,\n FVector &weights\n )\n{\n if (S.get_block_size() == 1)\n {\n computeWeights_1x1(S, weights);\n }\n else\n FatalError(\"Unsupported block size for strong connections\",\n AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class Strength_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Strength_BaseBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n\n/*\n* Note:\n* This implementation assumes that off-diag entries all have the opposite sign\n* comparing to the diag entry. This is true for most practical cases.\n* It would even work if the offending off-diag entries are just a few.\n* But if there are many off-diag entries violate this assumption,\n* the interpolation based on this strength would be inaccurate.\n* This is explained in \"Intro to Algebraic multigrid\" by K. Stuben.\n*/\n\n\nnamespace amgx\n{\n\ntemplate< typename T_Config >\nStrength_BaseBase::Strength_BaseBase(AMG_Config &cfg,\n const std::string &cfg_scope)\n{\n alpha = cfg.AMG_Config::template getParameter(\"strength_threshold\", cfg_scope);\n m_use_opt_kernels = cfg.AMG_Config::template getParameter(\"use_opt_kernels\", \"default\");\n}\n\n/*************************************************************************\n* \"random\" hash function for both device and host\n************************************************************************/\n__host__ __device__ __forceinline__\nstatic float ourHash(const int i)\n{\n unsigned int a = i;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return (a ^ 0x4a51e590) / (float)UINT_MAX;\n}\n\n/*************************************************************************\n* Computes the strength matrix and the connection weights\n* Described in \\S 4.1 of:\n* \"Reducing complexity in parallel algebraic multigrid preconditioners\"\n*\n************************************************************************/\ntemplate \nvoid\nStrength_Base >::\ncomputeStrongConnectionsAndWeights_1x1(Matrix_h &A,\n BVector &s_con,\n FVector &weights,\n const double max_row_sum)\n{\n bool compute_row_sum = (max_row_sum < 1.0);\n VVector sums_ptr;\n\n // get the (normalised) row sums\n if (compute_row_sum)\n {\n sums_ptr.resize(A.get_num_rows());\n weightedRowSum(A, sums_ptr);\n cudaCheckError();\n }\n\n // get min/max off-diag, depending on sign of diagonal\n for (int row = 0; row < A.get_num_rows(); row++)\n {\n ValueType diag(0), minVal(0), maxVal(0);\n ValueType row_sum = compute_row_sum ? sums_ptr[row] : -1;\n int rowEnd = A.row_offsets[row + 1];\n\n for (int j = A.row_offsets[row]; j < rowEnd; j++)\n {\n int col = A.col_indices[j];\n ValueType val = A.values[j];\n\n if (col == row)\n {\n diag = val;\n }\n else\n {\n minVal = std::min(minVal, val);\n maxVal = std::max(maxVal, val);\n }\n }\n\n //set the threshold for being strongly connected\n ValueType threshold = (diag < 0) ? maxVal : minVal;\n threshold *= this->alpha;\n\n //initialize the weight to a small random number\n if (A.is_matrix_singleGPU())\n {\n weights[row] += ourHash(row);\n }\n else\n {\n weights[row] += ourHash((int)A.manager->base_index() + row);\n }\n\n // sum the column of S - # of points that strongly connect to me\n for (int j = A.row_offsets[row]; j < rowEnd; j++)\n {\n bool is_strongly_connected = false;\n\n if (compute_row_sum && row_sum > max_row_sum)\n {\n is_strongly_connected = false;\n }\n else\n is_strongly_connected =\n this->strongly_connected(A.values[j], threshold, diag);\n\n int col = A.col_indices[j];\n s_con[j] = (col != row) && is_strongly_connected;\n weights[A.col_indices[j]] += s_con[j] ? 1. : 0.;\n }\n }\n}\n\n\ntemplate \nvoid\nStrength_Base >::\ncomputeWeights_1x1(Matrix_h &S,\n FVector &weights)\n{\n for (int row = 0; row < S.get_num_rows(); row++)\n {\n //initialize the weight to a small random number\n if (S.is_matrix_singleGPU())\n {\n weights[row] += ourHash(row);\n }\n else\n {\n weights[row] += ourHash(S.manager->base_index() + row);\n }\n\n int rowEnd = S.row_offsets[row + 1];\n\n for (int j = S.row_offsets[row]; j < rowEnd; j++)\n {\n int col = S.col_indices[j];\n\n if (col != row) { weights[col] += 1.; }\n }\n }\n}\n\n\ntemplate \n__device__\nbool stronglyConnectedAHat(ValueType value, ValueType threshold, ValueType diag)\n{\n return (diag < ValueType(0)) ? value > threshold : value < threshold;\n}\n\n/*************************************************************************\n* Computes the strength matrix and the connection weights (device)\n************************************************************************/\ntemplate< typename IndexType, typename ValueType, int kCtaSize, bool singleGPU >\n__global__\nvoid computeStrongConnectionsAndWeightsKernel( const IndexType *A_rows,\n const IndexType *A_cols,\n const ValueType *A_vals,\n int A_num_rows,\n bool *s_con,\n float *weights,\n ValueType alpha,\n ValueType *row_sum,\n const double max_row_sum,\n int64_t base_index)\n{\n // One warp works on each row and hence one iteration handles\n // num_warps*numBlock rows. This means atomicAdd() is inevitable.\n const int num_warps = kCtaSize / 32;\n const int num_rows_per_iter = num_warps * gridDim.x;\n __shared__ volatile ValueType smem[kCtaSize];\n __shared__ volatile ValueType s_diag[num_warps];\n __shared__ volatile ValueType s_threshold[num_warps];\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n for ( int aRowId = blockIdx.x * num_warps + warpId ; aRowId < A_num_rows ;\n aRowId += num_rows_per_iter )\n {\n ValueType minVal(0), maxVal(0);\n\n if ( laneId == 0 ) // Reset the diagonal\n {\n s_diag[warpId] = ValueType(0);\n }\n\n utils::syncwarp();\n\n // Row sum\n ValueType rowSum = -1.0;\n\n if (max_row_sum < 1.0) { rowSum = row_sum[aRowId]; }\n\n // get diagonal, min/max off-diagonals\n const int aRowBegin = A_rows[aRowId ];\n const int aRowEnd = A_rows[aRowId + 1];\n\n for ( IndexType aRowIt = aRowBegin + laneId ; utils::any( aRowIt < aRowEnd ) ;\n aRowIt += 32 )\n {\n IndexType aColId = aRowIt < aRowEnd ? A_cols[aRowIt] : -1;\n ValueType aValue = aRowIt < aRowEnd ? A_vals[aRowIt] : ValueType(0);\n\n if ( aColId == aRowId ) // only one thread evaluates to true.\n {\n s_diag[warpId] = aValue;\n }\n\n bool is_off_diagonal = aRowIt < aRowEnd && aColId != aRowId;\n\n if ( is_off_diagonal )\n {\n minVal = min( minVal, aValue );\n maxVal = max( maxVal, aValue );\n }\n }\n\n // init weights[] with a random number\n if ( laneId == 0 )\n {\n if ( singleGPU )\n {\n atomicAdd( &weights[aRowId], ourHash(aRowId) );\n }\n else\n {\n atomicAdd( &weights[aRowId], ourHash( (int) base_index + aRowId) );\n }\n }\n\n utils::syncwarp();\n\n // Big assumption: diag and off-diag always have the opposite sign.\n // If diag entry is negative, then all off-diag entries must be positive.\n // This means max off-diag is to be used to compute the threshold.\n // If diag entry is positve, the min off-diag is used instead.\n if ( s_diag[warpId] < ValueType(0) )\n {\n smem[threadIdx.x] = maxVal;\n\n utils::syncwarp();\n\n#pragma unroll\n for ( int offset = 16 ; offset > 0 ; offset /= 2 )\n {\n if ( laneId < offset )\n {\n smem[threadIdx.x] = maxVal = max( maxVal, smem[threadIdx.x + offset] );\n }\n utils::syncwarp();\n }\n }\n else\n {\n smem[threadIdx.x] = minVal;\n\n utils::syncwarp();\n\n#pragma unroll\n for ( int offset = 16 ; offset > 0 ; offset /= 2 )\n {\n if ( laneId < offset )\n {\n smem[threadIdx.x] = minVal = min( minVal, smem[threadIdx.x + offset] );\n }\n utils::syncwarp();\n }\n }\n\n if ( laneId == 0 )\n {\n // If laneId=0, then maxVal or minVal is in smem[threadIdx.x].\n s_threshold[warpId] = smem[threadIdx.x] * alpha;\n }\n\n utils::syncwarp();\n\n // sum of the column of S\n for ( IndexType aRowIt = aRowBegin + laneId ; utils::any( aRowIt < aRowEnd ) ;\n aRowIt += 32 )\n {\n IndexType aColId = aRowIt < aRowEnd ? A_cols[aRowIt] : -1;\n ValueType aValue = aRowIt < aRowEnd ? A_vals[aRowIt] : ValueType(0);\n bool is_strongly_connected = false;\n\n if (max_row_sum < 1.0 && rowSum > max_row_sum)\n {\n is_strongly_connected = false;\n }\n else\n {\n bool is_off_diagonal = aRowIt < aRowEnd && aColId != aRowId;\n is_strongly_connected = is_off_diagonal &&\n stronglyConnectedAHat( aValue, s_threshold[warpId], s_diag[warpId] );\n }\n\n if ( is_strongly_connected && aRowIt < aRowEnd && aColId < A_num_rows)\n {\n atomicAdd( &weights[aColId], 1.0f );\n }\n\n if ( aRowIt < aRowEnd )\n {\n s_con[aRowIt] = is_strongly_connected;\n }\n }\n }\n}\n\n/*************************************************************************\n* Computes the strength matrix and the connection weights (device)\n************************************************************************/\ntemplate< typename IndexType, typename ValueType, int kCtaSize, bool singleGPU >\n__global__\nvoid computeStrongConnectionsAndWeightsKernel_opt( const IndexType *A_rows,\n const IndexType *A_cols,\n const ValueType *A_vals,\n int A_num_rows,\n bool *s_con,\n float *weights,\n ValueType alpha,\n ValueType *row_sum,\n const double max_row_sum,\n int64_t base_index)\n{\n // One warp works on each row and hence one iteration handles\n // num_warps*numBlock rows. This means atomicAdd() is inevitable.\n int aRowId = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (aRowId < A_num_rows)\n {\n ValueType minVal(0), maxVal(0);\n\n ValueType diag = ValueType(0);\n\n utils::syncwarp();\n\n // Row sum\n ValueType rowSum = -1.0;\n\n if (max_row_sum < 1.0) { rowSum = row_sum[aRowId]; }\n\n // get diagonal, min/max off-diagonals\n const int aRowBegin = A_rows[aRowId ];\n const int aRowEnd = A_rows[aRowId + 1];\n\n for ( IndexType aRowIt = aRowBegin; aRowIt < aRowEnd ; ++aRowIt)\n {\n IndexType aColId = A_cols[aRowIt];\n ValueType aValue = A_vals[aRowIt];\n\n if ( aColId == aRowId ) // only one thread evaluates to true.\n {\n diag = aValue;\n }\n\n if ( aColId != aRowId )\n {\n minVal = min( minVal, aValue );\n maxVal = max( maxVal, aValue );\n }\n }\n\n if ( singleGPU )\n {\n atomicAdd(&weights[aRowId], ourHash(aRowId));\n }\n else\n {\n atomicAdd(&weights[aRowId], ourHash( (int) base_index + aRowId));\n }\n\n utils::syncwarp();\n\n // Big assumption: diag and off-diag always have the opposite sign.\n // If diag entry is negative, then all off-diag entries must be positive.\n // This means max off-diag is to be used to compute the threshold.\n // If diag entry is positve, the min off-diag is used instead.\n ValueType threshold;\n if ( diag < ValueType(0) )\n {\n threshold = maxVal*alpha;\n }\n else\n {\n threshold = minVal*alpha;\n }\n\n utils::syncwarp();\n\n // sum of the column of S\n for ( IndexType aRowIt = aRowBegin; aRowIt < aRowEnd ; ++aRowIt)\n {\n IndexType aColId = A_cols[aRowIt];\n ValueType aValue = A_vals[aRowIt];\n bool is_strongly_connected = false;\n\n if (max_row_sum < 1.0 && rowSum > max_row_sum)\n {\n is_strongly_connected = false;\n }\n else\n {\n bool is_off_diagonal = aColId != aRowId;\n is_strongly_connected = is_off_diagonal &&\n stronglyConnectedAHat( aValue, threshold, diag );\n }\n\n if ( is_strongly_connected && aColId < A_num_rows)\n {\n atomicAdd( &weights[aColId], 1.0f );\n }\n\n s_con[aRowIt] = is_strongly_connected;\n }\n }\n}\n\n\ntemplate< typename IndexType, typename ValueType, int kCtaSize, bool singleGPU >\n__global__\nvoid computeWeightsKernel( const IndexType *A_rows,\n const IndexType *A_cols,\n int A_num_rows,\n float *weights,\n int64_t base_index)\n{\n const int num_warps = kCtaSize / 32;\n const int num_rows_per_iter = num_warps * gridDim.x;\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n for ( int aRowId = blockIdx.x * num_warps + warpId ; aRowId < A_num_rows ;\n aRowId += num_rows_per_iter )\n {\n if ( laneId == 0 )\n {\n if ( singleGPU )\n {\n atomicAdd( &weights[aRowId], ourHash(aRowId) );\n }\n else\n {\n atomicAdd( &weights[aRowId], ourHash( (int) base_index + aRowId) );\n }\n }\n\n const int aRowBegin = A_rows[aRowId ];\n const int aRowEnd = A_rows[aRowId + 1];\n\n for ( IndexType aRowIt = aRowBegin + laneId ; utils::any( aRowIt < aRowEnd ) ;\n aRowIt += 32 )\n {\n IndexType aColId = aRowIt < aRowEnd ? A_cols[aRowIt] : -1;\n bool is_off_diagonal = aRowIt < aRowEnd && aColId != aRowId;\n\n if (is_off_diagonal)\n {\n atomicAdd( &weights[aColId], 1.0f );\n }\n }\n }\n}\n\n\n/*************************************************************************\n* Computes the strength matrix and the connection weights (device)\n************************************************************************/\ntemplate \nvoid Strength_Base >::\ncomputeStrongConnectionsAndWeights_1x1(Matrix_d &A,\n BVector &s_con,\n FVector &weights,\n const double max_row_sum)\n{\n typedef TemplateConfig TConfig_d;\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // get the raw pointers for everything I need\n const IndexType *offsets_ptr = A.row_offsets.raw();\n const IndexType *column_indices_ptr = A.col_indices.raw();\n const ValueType *values_ptr = A.values.raw();\n bool *s_con_ptr = s_con.raw();\n float *weights_ptr = weights.raw();\n bool compute_row_sum = (max_row_sum < 1.0);\n\n if (A.get_num_rows() == 0) { compute_row_sum = false; }\n\n VVector sums_ptr;\n\n if (compute_row_sum)\n {\n sums_ptr.resize(A.get_num_rows());\n weightedRowSum(A, sums_ptr);\n cudaCheckError();\n }\n\n // choose a blocksize. Use 1 warp per row\n const int blockSize = 256;\n const int numWarps = blockSize / 32;\n const int numBlocks = std::min( 4096, (int) (A.get_num_rows() + numWarps - 1) / numWarps );\n\n if (A.get_num_rows() > 0)\n {\n if(this->m_use_opt_kernels)\n {\n // One row per thread\n const int blockSize = 128;\n const int numBlocks = A.get_num_rows() / blockSize + 1;\n\n if (A.is_matrix_singleGPU()) {\n computeStrongConnectionsAndWeightsKernel_opt\n <<< numBlocks, blockSize>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n A.get_num_rows(),\n s_con.raw(),\n weights.raw(),\n this->alpha,\n compute_row_sum ? sums_ptr.raw() : NULL,\n max_row_sum,\n 0);\n cudaCheckError();\n }\n else {\n computeStrongConnectionsAndWeightsKernel_opt\n <<< numBlocks, blockSize>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n A.get_num_rows(),\n s_con.raw(),\n weights.raw(),\n this->alpha,\n compute_row_sum ? sums_ptr.raw() : NULL,\n max_row_sum,\n A.manager->base_index());\n cudaCheckError();\n }\n }\n else\n {\n const int blockSize = 256;\n const int numWarps = blockSize / 32;\n const int numBlocks = min( 4096, (int) (A.get_num_rows() + numWarps - 1) / numWarps );\n\n if (A.is_matrix_singleGPU()) {\n computeStrongConnectionsAndWeightsKernel\n <<< numBlocks, blockSize>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n A.get_num_rows(),\n s_con.raw(),\n weights.raw(),\n this->alpha,\n compute_row_sum ? sums_ptr.raw() : NULL,\n max_row_sum,\n 0);\n cudaCheckError();\n }\n else {\n computeStrongConnectionsAndWeightsKernel\n <<< numBlocks, blockSize>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n A.values.raw(),\n A.get_num_rows(),\n s_con.raw(),\n weights.raw(),\n this->alpha,\n compute_row_sum ? sums_ptr.raw() : NULL,\n max_row_sum,\n A.manager->base_index());\n cudaCheckError();\n }\n }\n }\n\n if (!A.is_matrix_singleGPU() && A.currentView() == OWNED)\n {\n // Need to add neighbors contribution to my weights\n weights.dirtybit = 1;\n A.manager->add_from_halo(weights, weights.tag);\n }\n\n cudaCheckError();\n}\n\n\ntemplate \nvoid Strength_Base >::\ncomputeWeights_1x1(Matrix_d &S,\n FVector &weights)\n{\n typedef TemplateConfig TConfig_d;\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueType;\n // get the raw pointers for everything I need\n const IndexType *offsets_ptr = S.row_offsets.raw();\n const IndexType *column_indices_ptr = S.col_indices.raw();\n float *weights_ptr = weights.raw();\n // choose a blocksize. Use 1 thread per row\n const int blockSize = 256;\n const int numWarps = blockSize / 32;\n const int numBlocks = std::min( 4096, (int) (S.get_num_rows() + numWarps - 1) / numWarps );\n hipDeviceSynchronize();\n cudaCheckError();\n\n // call the CUDA kernel\n if (S.is_matrix_singleGPU())\n computeWeightsKernel\n <<< numBlocks, blockSize>>>(\n S.row_offsets.raw(),\n S.col_indices.raw(),\n S.get_num_rows(),\n weights.raw(),\n 0);\n else\n computeWeightsKernel\n <<< numBlocks, blockSize>>>(\n S.row_offsets.raw(),\n S.col_indices.raw(),\n S.get_num_rows(),\n weights.raw(),\n S.manager->base_index());\n\n cudaCheckError();\n\n if (!S.is_matrix_singleGPU() && S.currentView() == OWNED)\n {\n // Need to add neighbors contribution to my weights\n weights.dirtybit = 1;\n S.manager->add_from_halo(weights, weights.tag);\n }\n\n cudaCheckError();\n}\n\ntemplate\nvoid Strength_BaseBase::\ncomputeStrongConnectionsAndWeights(Matrix &A,\n BVector &s_con,\n FVector &weights,\n const double max_row_sum\n )\n{\n if (A.get_block_size() == 1)\n {\n computeStrongConnectionsAndWeights_1x1(A, s_con, weights, max_row_sum);\n }\n else\n FatalError(\"Unsupported block size for strong connections\",\n AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\ntemplate\nvoid Strength_BaseBase::computeWeights(Matrix &S,\n FVector &weights\n )\n{\n if (S.get_block_size() == 1)\n {\n computeWeights_1x1(S, weights);\n }\n else\n FatalError(\"Unsupported block size for strong connections\",\n AMGX_ERR_NOT_SUPPORTED_BLOCKSIZE);\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n#define AMGX_CASE_LINE(CASE) template class Strength_Base::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Strength_BaseBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n\n", "source": "/kaggle/working/amgx/src/classical/strength/strength_base.cu"} {"input": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \n#include \n#include \n//#include \n//#include \n#include \n#include \n#include \n\n#include \n\n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \n/*\n* Note:\n* This implementation assumes that off-diag entries all have the opposite sign\n* comparing to the diag entry. This is true for most practical cases.\n* It would even work if the offending off-diag entries are just a few.\n* But if there are many off-diag entries violate this assumption,\n* the interpolation based on this strength would be inaccurate.\n* This is explained in \"Intro to Algebraic multigrid\" by K. Stuben.\n*/\n\n\nnamespace amgx\n{\n\nnamespace\n{\nint level_counter = 0;\nbool is_random_initialized = false;\n}\n\ntemplate \nStrength_Affinity >::Strength_Affinity(AMG_Config &cfg,\n const std::string &cfg_scope) : Strength_AffinityBase(cfg, cfg_scope)\n{\n AMG_Config default_cfg;\n std::string default_cfg_scope = \"default\";\n affinity_iters = cfg.AMG_Config::template getParameter(\"affinity_iterations\", cfg_scope);\n n_TV = cfg.AMG_Config::template getParameter(\"affinity_vectors\", cfg_scope);\n\n if (n_TV > 32)\n {\n FatalError(\"Strength affinity: Test vectors > 32 is not supported, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n }\n\n solver = new MulticolorGaussSeidelSolver(cfg, cfg_scope);\n this->solver->set_max_iters(affinity_iters);\n}\n\n/*************************************************************************\n* \"random\" hash function for both device and host\n************************************************************************/\n__host__ __device__ __forceinline__\nstatic float ourHash(const int i)\n{\n unsigned int a = i;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return (a ^ 0x4a51e590) / (float)UINT_MAX;\n}\n\nstruct prg\n{\n float a, b;\n\n __host__ __device__\n prg(float _a = -1.f, float _b = 1.f) : a(_a), b(_b) {};\n\n __host__ __device__\n float operator()(const unsigned int n) const\n {\n#if 1\n amgx::thrust::default_random_engine rng;\n amgx::thrust::uniform_real_distribution dist(a, b);\n rng.discard(n);\n return dist(rng);\n#else\n float ru = ourHash(n);\n return (ru - 0.5f) * 2.0f;\n#endif\n }\n};\n\n\ntemplate \nvoid initRandom(Vector &vec, int size)\n{\n vec.resize(size);\n amgx::thrust::counting_iterator index_sequence_begin(0);\n thrust_wrapper::transform(index_sequence_begin,\n index_sequence_begin + size,\n vec.begin(),\n prg(-1.f, 1.f));\n}\n\n// Device utils\ntemplate \nstatic __device__ __inline__ T cahedRead (const T *ar)\n{\n return utils::Ld::load(ar);\n}\n\n// ! shfl is defined for int/float only in arch 600\ntemplate< typename ValueType >\n__forceinline__ __device__ void reduce_inwarp_mul(\n const ValueType vValueA,\n const ValueType vValueB,\n const int vecId,\n const int n_per_warp,\n volatile ValueType *smem,\n double *s_xy)\n{\n ValueType sum = vValueA * vValueB;\n#pragma unroll\n\n for ( int offset = 16 / n_per_warp ; offset > 0 ; offset /= 2 )\n {\n sum += utils::shfl_xor(sum, offset, 16);\n }\n\n *s_xy = sum;\n}\n\ntemplate< typename ValueType >\n__forceinline__ __device__ void reduce_inblock_mul(\n const int blockSize,\n const ValueType vValueA,\n const ValueType vValueB,\n const int vecId,\n const int n_per_block,\n volatile ValueType *smem,\n double *s_xy)\n{\n ValueType sum = vValueA * vValueB;\n smem[threadIdx.x] = sum;\n __syncthreads();\n\n if (blockSize >= 512)\n {\n if (vecId < 256 / n_per_block)\n {\n smem[threadIdx.x] = sum = sum + smem[threadIdx.x + 256 / n_per_block];\n }\n\n __syncthreads();\n }\n\n if (blockSize >= 256)\n {\n if (vecId < 128 / n_per_block)\n {\n smem[threadIdx.x] = sum = sum + smem[threadIdx.x + 128 / n_per_block];\n }\n\n __syncthreads();\n }\n\n if (blockSize >= 128)\n {\n if (vecId < 64 / n_per_block)\n {\n smem[threadIdx.x] = sum = sum + smem[threadIdx.x + 64 / n_per_block];\n }\n\n __syncthreads();\n }\n\n#pragma unroll\n\n for ( int offset = 16 ; offset > 0 ; offset /= 2 )\n if ( vecId < offset )\n {\n smem[threadIdx.x] = sum = sum + smem[threadIdx.x + offset];\n }\n\n if ( vecId == 0 )\n {\n // If laneId=0, then sum is in smem[threadIdx.x].\n *s_xy = sum;\n }\n}\n\n/*************************************************************************\n* Computes affinity matrix (device)\n************************************************************************/\n\ntemplate< typename IndexType, typename ValueTypeA, typename ValueTypeB, int kCtaSize >\n__global__\nvoid computeAffinity_1x1_Kernel(const int *A_rows,\n const int *A_cols,\n const ValueTypeB *X,\n const int nTV,\n const int A_num_rows,\n ValueTypeA *affinity\n )\n{\n const ValueTypeB epsilon = 1.e-12;\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n const int vid = tid % nTV;\n const int num_vecs_per_warp = 32 / nTV;\n const int num_rows_per_iter = gridDim.x * blockDim.x / nTV;\n //const int num_vecs_per_block = blockDim.x / nTV;\n ValueTypeB *smem = NULL;\n double s_xx, s_xy, s_yy;\n\n for ( int aRowId = tid / nTV ; aRowId < A_num_rows ;\n aRowId += num_rows_per_iter )\n {\n const int aRowBegin = A_rows[aRowId];\n const int aRowEnd = A_rows[aRowId + 1];\n ValueTypeB vValueA = X[aRowId * nTV + vid];\n reduce_inwarp_mul(vValueA, vValueA, vid, num_vecs_per_warp, smem, &s_xx);\n#if 0\n\n if (num_vecs_per_warp > 0)\n {\n reduce_inwarp_mul(vValueA, vValueA, vid, num_vecs_per_warp, smem, &s_xx);\n }\n else\n {\n reduce_inblock_mul(kCtaSize, vValueA, vValueA, vid, num_vecs_per_block, smem, &s_xx);\n }\n\n#endif\n\n for ( IndexType aRowIt = aRowBegin ; aRowIt < aRowEnd;\n aRowIt++ )\n {\n IndexType aColId = A_cols[aRowIt];\n\n if (aColId != aRowId)\n {\n ValueTypeB vValueB = cahedRead(X + aColId * nTV + vid);\n reduce_inwarp_mul(vValueA, vValueB, vid, num_vecs_per_warp, smem, &s_xy);\n reduce_inwarp_mul(vValueB, vValueB, vid, num_vecs_per_warp, smem, &s_yy);\n\n if (vid == 0)\n {\n affinity[aRowIt] = (s_xy * s_xy / (s_xx * s_yy + epsilon));\n }\n }\n }\n }\n}\n\n\n/*************************************************************************\n* Computes the strength matrix and the connection weights (device)\n************************************************************************/\n\ntemplate< typename IndexType, typename ValueType, int kCtaSize, bool singleGPU >\n__global__\nvoid computeStrongConnectionsAndWeightsFromAffinityKernel(\n const int *A_rows,\n const int *A_cols,\n const ValueType *A_vals,\n int A_num_rows,\n bool *s_con,\n float *weights,\n ValueType alpha,\n int64_t base_index)\n{\n // One warp works on each row and hence one iteration handles\n // num_warps*numBlock rows. This means atomicAdd() is inevitable.\n const int num_warps = kCtaSize / 32;\n const int num_rows_per_iter = num_warps * gridDim.x;\n __shared__ volatile ValueType smem[kCtaSize];\n __shared__ volatile ValueType s_threshold[num_warps];\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n for ( int aRowId = blockIdx.x * num_warps + warpId ; aRowId < A_num_rows ;\n aRowId += num_rows_per_iter )\n {\n ValueType maxVal(0);\n // get diagonal, min/max off-idagonals\n const int aRowBegin = A_rows[aRowId ];\n const int aRowEnd = A_rows[aRowId + 1];\n\n for ( IndexType aRowIt = aRowBegin + laneId ; utils::any( aRowIt < aRowEnd ) ;\n aRowIt += 32 )\n {\n IndexType aColId = aRowIt < aRowEnd ? A_cols[aRowIt] : -1;\n ValueType aValue = aRowIt < aRowEnd ? A_vals[aRowIt] : ValueType(0);\n bool is_off_diagonal = aRowIt < aRowEnd && aColId != aRowId;\n\n if ( is_off_diagonal )\n {\n maxVal = max( maxVal, aValue );\n }\n }\n\n // init weights[] with a random number\n if ( laneId == 0 )\n {\n if ( singleGPU )\n {\n atomicAdd( &weights[aRowId], ourHash(aRowId) );\n }\n else\n {\n atomicAdd( &weights[aRowId], ourHash( (int) base_index + aRowId) );\n }\n }\n\n smem[threadIdx.x] = maxVal;\n#pragma unroll\n\n for ( int offset = 16 ; offset > 0 ; offset /= 2 )\n if ( laneId < offset )\n {\n smem[threadIdx.x] = maxVal = max( maxVal, smem[threadIdx.x + offset] );\n }\n\n if ( laneId == 0 )\n {\n // If laneId=0, then maxVal or minVal is in smem[threadIdx.x].\n s_threshold[warpId] = smem[threadIdx.x] * alpha;\n }\n\n // sum of the column of S\n for ( IndexType aRowIt = aRowBegin + laneId ; utils::any( aRowIt < aRowEnd ) ;\n aRowIt += 32 )\n {\n IndexType aColId = aRowIt < aRowEnd ? A_cols[aRowIt] : -1;\n ValueType aValue = aRowIt < aRowEnd ? A_vals[aRowIt] : ValueType(0);\n bool is_off_diagonal = aRowIt < aRowEnd && aColId != aRowId;\n bool is_strongly_connected = is_off_diagonal && (aValue > s_threshold[warpId]);\n\n if ( is_strongly_connected && aRowIt < aRowEnd && aColId < A_num_rows)\n {\n atomicAdd( &weights[aColId], 1.0f );\n }\n\n if ( aRowIt < aRowEnd )\n {\n s_con[aRowIt] = is_strongly_connected;\n }\n }\n }\n}\n\n// Save to file\ntemplate \nvoid dump_matrix_vector(const char *fname, const Matrix &A, const Vector &solution)\n{\n unsigned int num_nz = A.get_num_nz();\n unsigned int num_rows = A.get_num_rows();\n unsigned int vec_sz = solution.size();\n std::ofstream fout;\n fout.open(fname, std::ofstream::out | std::ofstream::app);\n fout << num_rows << \" \" << num_nz << \" \" << vec_sz << std::endl;\n typename Vector::value_type val;\n\n for (int i = 0; i < num_rows + 1; i++)\n {\n fout << A.row_offsets[i] << std::endl;\n }\n\n for (int i = 0; i < num_nz; i++)\n {\n fout << A.col_indices[i] << std::endl;\n }\n\n fout.precision(14);\n fout.width(16);\n\n for (int i = 0; i < num_nz; i++)\n {\n val = A.values[i];\n fout << val << std::endl;\n }\n\n for (int i = 0; i < vec_sz; i++)\n {\n val = solution[i];\n fout << val << std::endl;\n }\n\n fout << std::endl;\n fout.close();\n}\n\n/*************************************************************************\n* Computes the strength matrix and the connection weights (device)\n************************************************************************/\ntemplate \nvoid Strength_Affinity >::\ncomputeStrongConnectionsAndWeights_1x1(Matrix_d &A,\n BVector &s_con,\n FVector &weights,\n const double max_row_sum)\n{\n typedef TemplateConfig TConfig_d;\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n typedef typename Vector::value_type ValueTypeB;\n bool level_opt_flag = ((int) max_row_sum != -1);\n m_aff.copy(A);\n\n if (!is_random_initialized)\n {\n unsigned int nTV_num_rows = this->n_TV * A.get_num_rows();\n m_rhs.resize(nTV_num_rows, 0.0);\n initRandom(m_x, nTV_num_rows);\n\n if (level_opt_flag) { is_random_initialized = true; }\n\n level_counter = n_TV;\n }\n else\n {\n this->n_TV = min(level_counter, 32);\n unsigned int nTV_num_rows = this->n_TV * A.get_num_rows();\n m_rhs.resize(nTV_num_rows, 0.0);\n\n if (m_x.size() < nTV_num_rows)\n {\n initRandom(m_x, nTV_num_rows);\n }\n }\n\n level_counter *= 2;\n this->solver->setup_no_throw(m_aff, false);\n TimerCPU timer(true);\n timer.start();\n\n for (int iter = 0; iter < affinity_iters; iter ++)\n {\n this->solver->batch_smooth_1x1_fast(m_aff, this->n_TV, m_rhs, m_x);\n }\n\n cudaDeviceSynchronize();\n cudaCheckError();\n double elapsed = timer.elapsed();\n// End of GS check\n // get the raw pointers for everything I need\n const IndexType *offsets_ptr = m_aff.row_offsets.raw();\n const IndexType *column_indices_ptr = m_aff.col_indices.raw();\n bool *s_con_ptr = s_con.raw();\n float *weights_ptr = weights.raw();\n ValueTypeA *m_aff_values_ptr = m_aff.values.raw();\n //m_aff_values.raw();\n ValueTypeB *m_x_ptr = m_x.raw();\n // choose a blocksize. Use 1 warp per row\n const int blockSize = 256;\n const int numWarps = blockSize / 32;\n const int numBlocks = (int) (A.get_num_rows() + numWarps - 1) / numWarps;\n\n if (A.get_num_rows() > 0)\n {\n computeAffinity_1x1_Kernel\n <<< numBlocks, blockSize>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n m_x_ptr,\n this->n_TV,\n A.get_num_rows(),\n m_aff_values_ptr\n );\n cudaCheckError();\n\n if (A.is_matrix_singleGPU())\n computeStrongConnectionsAndWeightsFromAffinityKernel\n <<< numBlocks, blockSize>>>(\n m_aff.row_offsets.raw(),\n m_aff.col_indices.raw(),\n m_aff_values_ptr,\n m_aff.get_num_rows(),\n s_con.raw(),\n weights.raw(),\n this->alpha,\n int64_t(0));\n else\n computeStrongConnectionsAndWeightsFromAffinityKernel\n <<< numBlocks, blockSize>>>(\n m_aff.row_offsets.raw(),\n m_aff.col_indices.raw(),\n m_aff_values_ptr,\n m_aff.get_num_rows(),\n s_con.raw(),\n weights.raw(),\n this->alpha,\n A.manager->base_index());\n \n cudaCheckError();\n }\n\n if (!A.is_matrix_singleGPU() && A.currentView() == OWNED)\n {\n // Need to add neighbors contribution to my weights\n weights.dirtybit = 1;\n A.manager->add_from_halo(weights, weights.tag);\n }\n\n cudaCheckError();\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n\n#define AMGX_CASE_LINE(CASE) template class Strength_AffinityBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Strength_Affinity::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n\n\n", "output": "// SPDX-FileCopyrightText: 2011 - 2025 NVIDIA CORPORATION. All Rights Reserved.\n//\n// SPDX-License-Identifier: BSD-3-Clause\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n//#include \n//#include \n#include \n#include \n#include \n\n#include \n\n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \n/*\n* Note:\n* This implementation assumes that off-diag entries all have the opposite sign\n* comparing to the diag entry. This is true for most practical cases.\n* It would even work if the offending off-diag entries are just a few.\n* But if there are many off-diag entries violate this assumption,\n* the interpolation based on this strength would be inaccurate.\n* This is explained in \"Intro to Algebraic multigrid\" by K. Stuben.\n*/\n\n\nnamespace amgx\n{\n\nnamespace\n{\nint level_counter = 0;\nbool is_random_initialized = false;\n}\n\ntemplate \nStrength_Affinity >::Strength_Affinity(AMG_Config &cfg,\n const std::string &cfg_scope) : Strength_AffinityBase(cfg, cfg_scope)\n{\n AMG_Config default_cfg;\n std::string default_cfg_scope = \"default\";\n affinity_iters = cfg.AMG_Config::template getParameter(\"affinity_iterations\", cfg_scope);\n n_TV = cfg.AMG_Config::template getParameter(\"affinity_vectors\", cfg_scope);\n\n if (n_TV > 32)\n {\n FatalError(\"Strength affinity: Test vectors > 32 is not supported, exiting\", AMGX_ERR_NOT_SUPPORTED_TARGET);\n }\n\n solver = new MulticolorGaussSeidelSolver(cfg, cfg_scope);\n this->solver->set_max_iters(affinity_iters);\n}\n\n/*************************************************************************\n* \"random\" hash function for both device and host\n************************************************************************/\n__host__ __device__ __forceinline__\nstatic float ourHash(const int i)\n{\n unsigned int a = i;\n a = (a + 0x7ed55d16) + (a << 12);\n a = (a ^ 0xc761c23c) + (a >> 19);\n a = (a + 0x165667b1) + (a << 5);\n a = (a ^ 0xd3a2646c) + (a << 9);\n a = (a + 0xfd7046c5) + (a << 3);\n a = (a ^ 0xb55a4f09) + (a >> 16);\n return (a ^ 0x4a51e590) / (float)UINT_MAX;\n}\n\nstruct prg\n{\n float a, b;\n\n __host__ __device__\n prg(float _a = -1.f, float _b = 1.f) : a(_a), b(_b) {};\n\n __host__ __device__\n float operator()(const unsigned int n) const\n {\n#if 1\n amgx::thrust::default_random_engine rng;\n amgx::thrust::uniform_real_distribution dist(a, b);\n rng.discard(n);\n return dist(rng);\n#else\n float ru = ourHash(n);\n return (ru - 0.5f) * 2.0f;\n#endif\n }\n};\n\n\ntemplate \nvoid initRandom(Vector &vec, int size)\n{\n vec.resize(size);\n amgx::thrust::counting_iterator index_sequence_begin(0);\n thrust_wrapper::transform(index_sequence_begin,\n index_sequence_begin + size,\n vec.begin(),\n prg(-1.f, 1.f));\n}\n\n// Device utils\ntemplate \nstatic __device__ __inline__ T cahedRead (const T *ar)\n{\n return utils::Ld::load(ar);\n}\n\n// ! shfl is defined for int/float only in arch 600\ntemplate< typename ValueType >\n__forceinline__ __device__ void reduce_inwarp_mul(\n const ValueType vValueA,\n const ValueType vValueB,\n const int vecId,\n const int n_per_warp,\n volatile ValueType *smem,\n double *s_xy)\n{\n ValueType sum = vValueA * vValueB;\n#pragma unroll\n\n for ( int offset = 16 / n_per_warp ; offset > 0 ; offset /= 2 )\n {\n sum += utils::shfl_xor(sum, offset, 16);\n }\n\n *s_xy = sum;\n}\n\ntemplate< typename ValueType >\n__forceinline__ __device__ void reduce_inblock_mul(\n const int blockSize,\n const ValueType vValueA,\n const ValueType vValueB,\n const int vecId,\n const int n_per_block,\n volatile ValueType *smem,\n double *s_xy)\n{\n ValueType sum = vValueA * vValueB;\n smem[threadIdx.x] = sum;\n __syncthreads();\n\n if (blockSize >= 512)\n {\n if (vecId < 256 / n_per_block)\n {\n smem[threadIdx.x] = sum = sum + smem[threadIdx.x + 256 / n_per_block];\n }\n\n __syncthreads();\n }\n\n if (blockSize >= 256)\n {\n if (vecId < 128 / n_per_block)\n {\n smem[threadIdx.x] = sum = sum + smem[threadIdx.x + 128 / n_per_block];\n }\n\n __syncthreads();\n }\n\n if (blockSize >= 128)\n {\n if (vecId < 64 / n_per_block)\n {\n smem[threadIdx.x] = sum = sum + smem[threadIdx.x + 64 / n_per_block];\n }\n\n __syncthreads();\n }\n\n#pragma unroll\n\n for ( int offset = 16 ; offset > 0 ; offset /= 2 )\n if ( vecId < offset )\n {\n smem[threadIdx.x] = sum = sum + smem[threadIdx.x + offset];\n }\n\n if ( vecId == 0 )\n {\n // If laneId=0, then sum is in smem[threadIdx.x].\n *s_xy = sum;\n }\n}\n\n/*************************************************************************\n* Computes affinity matrix (device)\n************************************************************************/\n\ntemplate< typename IndexType, typename ValueTypeA, typename ValueTypeB, int kCtaSize >\n__global__\nvoid computeAffinity_1x1_Kernel(const int *A_rows,\n const int *A_cols,\n const ValueTypeB *X,\n const int nTV,\n const int A_num_rows,\n ValueTypeA *affinity\n )\n{\n const ValueTypeB epsilon = 1.e-12;\n const int tid = blockDim.x * blockIdx.x + threadIdx.x;\n const int vid = tid % nTV;\n const int num_vecs_per_warp = 32 / nTV;\n const int num_rows_per_iter = gridDim.x * blockDim.x / nTV;\n //const int num_vecs_per_block = blockDim.x / nTV;\n ValueTypeB *smem = NULL;\n double s_xx, s_xy, s_yy;\n\n for ( int aRowId = tid / nTV ; aRowId < A_num_rows ;\n aRowId += num_rows_per_iter )\n {\n const int aRowBegin = A_rows[aRowId];\n const int aRowEnd = A_rows[aRowId + 1];\n ValueTypeB vValueA = X[aRowId * nTV + vid];\n reduce_inwarp_mul(vValueA, vValueA, vid, num_vecs_per_warp, smem, &s_xx);\n#if 0\n\n if (num_vecs_per_warp > 0)\n {\n reduce_inwarp_mul(vValueA, vValueA, vid, num_vecs_per_warp, smem, &s_xx);\n }\n else\n {\n reduce_inblock_mul(kCtaSize, vValueA, vValueA, vid, num_vecs_per_block, smem, &s_xx);\n }\n\n#endif\n\n for ( IndexType aRowIt = aRowBegin ; aRowIt < aRowEnd;\n aRowIt++ )\n {\n IndexType aColId = A_cols[aRowIt];\n\n if (aColId != aRowId)\n {\n ValueTypeB vValueB = cahedRead(X + aColId * nTV + vid);\n reduce_inwarp_mul(vValueA, vValueB, vid, num_vecs_per_warp, smem, &s_xy);\n reduce_inwarp_mul(vValueB, vValueB, vid, num_vecs_per_warp, smem, &s_yy);\n\n if (vid == 0)\n {\n affinity[aRowIt] = (s_xy * s_xy / (s_xx * s_yy + epsilon));\n }\n }\n }\n }\n}\n\n\n/*************************************************************************\n* Computes the strength matrix and the connection weights (device)\n************************************************************************/\n\ntemplate< typename IndexType, typename ValueType, int kCtaSize, bool singleGPU >\n__global__\nvoid computeStrongConnectionsAndWeightsFromAffinityKernel(\n const int *A_rows,\n const int *A_cols,\n const ValueType *A_vals,\n int A_num_rows,\n bool *s_con,\n float *weights,\n ValueType alpha,\n int64_t base_index)\n{\n // One warp works on each row and hence one iteration handles\n // num_warps*numBlock rows. This means atomicAdd() is inevitable.\n const int num_warps = kCtaSize / 32;\n const int num_rows_per_iter = num_warps * gridDim.x;\n __shared__ volatile ValueType smem[kCtaSize];\n __shared__ volatile ValueType s_threshold[num_warps];\n const int warpId = threadIdx.x / 32;\n const int laneId = threadIdx.x % 32;\n\n for ( int aRowId = blockIdx.x * num_warps + warpId ; aRowId < A_num_rows ;\n aRowId += num_rows_per_iter )\n {\n ValueType maxVal(0);\n // get diagonal, min/max off-idagonals\n const int aRowBegin = A_rows[aRowId ];\n const int aRowEnd = A_rows[aRowId + 1];\n\n for ( IndexType aRowIt = aRowBegin + laneId ; utils::any( aRowIt < aRowEnd ) ;\n aRowIt += 32 )\n {\n IndexType aColId = aRowIt < aRowEnd ? A_cols[aRowIt] : -1;\n ValueType aValue = aRowIt < aRowEnd ? A_vals[aRowIt] : ValueType(0);\n bool is_off_diagonal = aRowIt < aRowEnd && aColId != aRowId;\n\n if ( is_off_diagonal )\n {\n maxVal = max( maxVal, aValue );\n }\n }\n\n // init weights[] with a random number\n if ( laneId == 0 )\n {\n if ( singleGPU )\n {\n atomicAdd( &weights[aRowId], ourHash(aRowId) );\n }\n else\n {\n atomicAdd( &weights[aRowId], ourHash( (int) base_index + aRowId) );\n }\n }\n\n smem[threadIdx.x] = maxVal;\n#pragma unroll\n\n for ( int offset = 16 ; offset > 0 ; offset /= 2 )\n if ( laneId < offset )\n {\n smem[threadIdx.x] = maxVal = max( maxVal, smem[threadIdx.x + offset] );\n }\n\n if ( laneId == 0 )\n {\n // If laneId=0, then maxVal or minVal is in smem[threadIdx.x].\n s_threshold[warpId] = smem[threadIdx.x] * alpha;\n }\n\n // sum of the column of S\n for ( IndexType aRowIt = aRowBegin + laneId ; utils::any( aRowIt < aRowEnd ) ;\n aRowIt += 32 )\n {\n IndexType aColId = aRowIt < aRowEnd ? A_cols[aRowIt] : -1;\n ValueType aValue = aRowIt < aRowEnd ? A_vals[aRowIt] : ValueType(0);\n bool is_off_diagonal = aRowIt < aRowEnd && aColId != aRowId;\n bool is_strongly_connected = is_off_diagonal && (aValue > s_threshold[warpId]);\n\n if ( is_strongly_connected && aRowIt < aRowEnd && aColId < A_num_rows)\n {\n atomicAdd( &weights[aColId], 1.0f );\n }\n\n if ( aRowIt < aRowEnd )\n {\n s_con[aRowIt] = is_strongly_connected;\n }\n }\n }\n}\n\n// Save to file\ntemplate \nvoid dump_matrix_vector(const char *fname, const Matrix &A, const Vector &solution)\n{\n unsigned int num_nz = A.get_num_nz();\n unsigned int num_rows = A.get_num_rows();\n unsigned int vec_sz = solution.size();\n std::ofstream fout;\n fout.open(fname, std::ofstream::out | std::ofstream::app);\n fout << num_rows << \" \" << num_nz << \" \" << vec_sz << std::endl;\n typename Vector::value_type val;\n\n for (int i = 0; i < num_rows + 1; i++)\n {\n fout << A.row_offsets[i] << std::endl;\n }\n\n for (int i = 0; i < num_nz; i++)\n {\n fout << A.col_indices[i] << std::endl;\n }\n\n fout.precision(14);\n fout.width(16);\n\n for (int i = 0; i < num_nz; i++)\n {\n val = A.values[i];\n fout << val << std::endl;\n }\n\n for (int i = 0; i < vec_sz; i++)\n {\n val = solution[i];\n fout << val << std::endl;\n }\n\n fout << std::endl;\n fout.close();\n}\n\n/*************************************************************************\n* Computes the strength matrix and the connection weights (device)\n************************************************************************/\ntemplate \nvoid Strength_Affinity >::\ncomputeStrongConnectionsAndWeights_1x1(Matrix_d &A,\n BVector &s_con,\n FVector &weights,\n const double max_row_sum)\n{\n typedef TemplateConfig TConfig_d;\n typedef typename Matrix_d::index_type IndexType;\n typedef typename Matrix_d::value_type ValueTypeA;\n typedef typename Vector::value_type ValueTypeB;\n bool level_opt_flag = ((int) max_row_sum != -1);\n m_aff.copy(A);\n\n if (!is_random_initialized)\n {\n unsigned int nTV_num_rows = this->n_TV * A.get_num_rows();\n m_rhs.resize(nTV_num_rows, 0.0);\n initRandom(m_x, nTV_num_rows);\n\n if (level_opt_flag) { is_random_initialized = true; }\n\n level_counter = n_TV;\n }\n else\n {\n this->n_TV = min(level_counter, 32);\n unsigned int nTV_num_rows = this->n_TV * A.get_num_rows();\n m_rhs.resize(nTV_num_rows, 0.0);\n\n if (m_x.size() < nTV_num_rows)\n {\n initRandom(m_x, nTV_num_rows);\n }\n }\n\n level_counter *= 2;\n this->solver->setup_no_throw(m_aff, false);\n TimerCPU timer(true);\n timer.start();\n\n for (int iter = 0; iter < affinity_iters; iter ++)\n {\n this->solver->batch_smooth_1x1_fast(m_aff, this->n_TV, m_rhs, m_x);\n }\n\n hipDeviceSynchronize();\n cudaCheckError();\n double elapsed = timer.elapsed();\n// End of GS check\n // get the raw pointers for everything I need\n const IndexType *offsets_ptr = m_aff.row_offsets.raw();\n const IndexType *column_indices_ptr = m_aff.col_indices.raw();\n bool *s_con_ptr = s_con.raw();\n float *weights_ptr = weights.raw();\n ValueTypeA *m_aff_values_ptr = m_aff.values.raw();\n //m_aff_values.raw();\n ValueTypeB *m_x_ptr = m_x.raw();\n // choose a blocksize. Use 1 warp per row\n const int blockSize = 256;\n const int numWarps = blockSize / 32;\n const int numBlocks = (int) (A.get_num_rows() + numWarps - 1) / numWarps;\n\n if (A.get_num_rows() > 0)\n {\n computeAffinity_1x1_Kernel\n <<< numBlocks, blockSize>>>(\n A.row_offsets.raw(),\n A.col_indices.raw(),\n m_x_ptr,\n this->n_TV,\n A.get_num_rows(),\n m_aff_values_ptr\n );\n cudaCheckError();\n\n if (A.is_matrix_singleGPU())\n computeStrongConnectionsAndWeightsFromAffinityKernel\n <<< numBlocks, blockSize>>>(\n m_aff.row_offsets.raw(),\n m_aff.col_indices.raw(),\n m_aff_values_ptr,\n m_aff.get_num_rows(),\n s_con.raw(),\n weights.raw(),\n this->alpha,\n int64_t(0));\n else\n computeStrongConnectionsAndWeightsFromAffinityKernel\n <<< numBlocks, blockSize>>>(\n m_aff.row_offsets.raw(),\n m_aff.col_indices.raw(),\n m_aff_values_ptr,\n m_aff.get_num_rows(),\n s_con.raw(),\n weights.raw(),\n this->alpha,\n A.manager->base_index());\n \n cudaCheckError();\n }\n\n if (!A.is_matrix_singleGPU() && A.currentView() == OWNED)\n {\n // Need to add neighbors contribution to my weights\n weights.dirtybit = 1;\n A.manager->add_from_halo(weights, weights.tag);\n }\n\n cudaCheckError();\n}\n\n\n/****************************************\n * Explict instantiations\n ***************************************/\n\n#define AMGX_CASE_LINE(CASE) template class Strength_AffinityBase::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n#define AMGX_CASE_LINE(CASE) template class Strength_Affinity::Type>;\nAMGX_FORALL_BUILDS(AMGX_CASE_LINE)\n#undef AMGX_CASE_LINE\n\n} // namespace amgx\n\n\n", "source": "/kaggle/working/amgx/src/classical/strength/affinity.cu"} {"input": "\n#include \n\n// CUDA-compatible types (no C++ standard library headers for NVRTC)\ntypedef unsigned char uint8_t;\ntypedef unsigned long long uint64_t;\ntypedef uint64_t uintptr_t;\n\n// ============================================================================\n// Kernel Configuration Constants (from ChunkedKernel.h)\n// ============================================================================\n\nconstexpr int CHUNK_SIZE = 128 * 1024; // 128KB per chunk\nconstexpr int NUM_BLOCKS = 75; // Fixed grid size\nconstexpr int THREADS_PER_BLOCK = 1024; // Fixed block size\nconstexpr int FLOAT4_SIZE = 16; // 16 bytes per float4\nconstexpr int MAX_CHUNKS_PER_BLOCK = 512; // Max chunks per block for shared memory\n\nextern \"C\" {\n\n/**\n * Chunked batched copy kernel implementation\n *\n * This kernel performs efficient batched memory copies using:\n * 1. Contiguous block assignment for better load balancing\n * 2. Shared memory prefetching of chunk metadata\n * 3. Vectorized float4 (16-byte) copies for aligned data\n * 4. Byte-by-byte fallback for unaligned or small data\n */\n__global__ void chunked_batched_copy_kernel(\n uint8_t** src_addrs,\n uint8_t** dst_addrs,\n size_t* sizes,\n int total_chunks\n) {\n // Shared memory for metadata prefetching\n __shared__ uint8_t* s_src_addrs[MAX_CHUNKS_PER_BLOCK];\n __shared__ uint8_t* s_dst_addrs[MAX_CHUNKS_PER_BLOCK];\n __shared__ size_t s_sizes[MAX_CHUNKS_PER_BLOCK];\n\n // Contiguous block assignment: block i processes chunks [start_chunk, end_chunk)\n int chunks_per_block = (total_chunks + gridDim.x - 1) / gridDim.x; // Ceiling division\n int start_chunk = blockIdx.x * chunks_per_block;\n int end_chunk = start_chunk + chunks_per_block;\n if (end_chunk > total_chunks) {\n end_chunk = total_chunks;\n }\n int num_chunks_this_block = end_chunk - start_chunk;\n\n // Phase 1: Cooperative loading of metadata to shared memory\n // All 1024 threads cooperate to load metadata from global memory\n for (int i = threadIdx.x; i < num_chunks_this_block; i += blockDim.x) {\n int global_chunk_id = start_chunk + i;\n s_src_addrs[i] = src_addrs[global_chunk_id];\n s_dst_addrs[i] = dst_addrs[global_chunk_id];\n s_sizes[i] = sizes[global_chunk_id];\n }\n __syncthreads();\n\n // Phase 2: Process each chunk using metadata from shared memory\n for (int chunk_id = 0; chunk_id < num_chunks_this_block; chunk_id++) {\n uint8_t* src = s_src_addrs[chunk_id];\n uint8_t* dst = s_dst_addrs[chunk_id];\n size_t size = s_sizes[chunk_id];\n\n // Check if both src and dst are aligned to 16 bytes for float4 access\n uintptr_t src_addr = (uintptr_t)src;\n uintptr_t dst_addr = (uintptr_t)dst;\n bool is_aligned = ((src_addr % FLOAT4_SIZE) == 0) && ((dst_addr % FLOAT4_SIZE) == 0);\n\n if (is_aligned && size >= FLOAT4_SIZE) {\n // Fast path: vectorized float4 copies\n size_t aligned_size = (size / FLOAT4_SIZE) * FLOAT4_SIZE;\n\n // All 1024 threads cooperate on float4 copies\n #pragma unroll 4\n for (size_t offset = threadIdx.x * FLOAT4_SIZE;\n offset < aligned_size;\n offset += blockDim.x * FLOAT4_SIZE) {\n // Vectorized 16-byte load and store\n float4 data = *((float4*)(src + offset));\n *((float4*)(dst + offset)) = data;\n }\n\n // Handle remaining bytes (< 16 bytes) with byte-by-byte copy\n for (size_t offset = aligned_size + threadIdx.x;\n offset < size;\n offset += blockDim.x) {\n dst[offset] = src[offset];\n }\n } else {\n // Fallback path: byte-by-byte copy for unaligned addresses\n // Still use all threads for parallelism\n for (size_t offset = threadIdx.x; offset < size; offset += blockDim.x) {\n dst[offset] = src[offset];\n }\n }\n }\n\n // System-wide memory fence: ensures all writes from this thread are visible\n // to system agents (NIC DMA engines, other GPUs via PCIe/NVLink).\n // Required so that NVSHMEM put() on another CUDA stream can read correct\n // pack data from the send buffer. Without this, writes may remain in L2\n // cache, invisible to the NIC's DMA engine.\n __threadfence_system();\n}\n\n}\n\n\n", "output": "\n#include \n\n// CUDA-compatible types (no C++ standard library headers for NVRTC)\ntypedef unsigned char uint8_t;\ntypedef unsigned long long uint64_t;\ntypedef uint64_t uintptr_t;\n\n// ============================================================================\n// Kernel Configuration Constants (from ChunkedKernel.h)\n// ============================================================================\n\nconstexpr int CHUNK_SIZE = 128 * 1024; // 128KB per chunk\nconstexpr int NUM_BLOCKS = 75; // Fixed grid size\nconstexpr int THREADS_PER_BLOCK = 1024; // Fixed block size\nconstexpr int FLOAT4_SIZE = 16; // 16 bytes per float4\nconstexpr int MAX_CHUNKS_PER_BLOCK = 512; // Max chunks per block for shared memory\n\nextern \"C\" {\n\n/**\n * Chunked batched copy kernel implementation\n *\n * This kernel performs efficient batched memory copies using:\n * 1. Contiguous block assignment for better load balancing\n * 2. Shared memory prefetching of chunk metadata\n * 3. Vectorized float4 (16-byte) copies for aligned data\n * 4. Byte-by-byte fallback for unaligned or small data\n */\n__global__ void chunked_batched_copy_kernel(\n uint8_t** src_addrs,\n uint8_t** dst_addrs,\n size_t* sizes,\n int total_chunks\n) {\n // Shared memory for metadata prefetching\n __shared__ uint8_t* s_src_addrs[MAX_CHUNKS_PER_BLOCK];\n __shared__ uint8_t* s_dst_addrs[MAX_CHUNKS_PER_BLOCK];\n __shared__ size_t s_sizes[MAX_CHUNKS_PER_BLOCK];\n\n // Contiguous block assignment: block i processes chunks [start_chunk, end_chunk)\n int chunks_per_block = (total_chunks + gridDim.x - 1) / gridDim.x; // Ceiling division\n int start_chunk = blockIdx.x * chunks_per_block;\n int end_chunk = start_chunk + chunks_per_block;\n if (end_chunk > total_chunks) {\n end_chunk = total_chunks;\n }\n int num_chunks_this_block = end_chunk - start_chunk;\n\n // Phase 1: Cooperative loading of metadata to shared memory\n // All 1024 threads cooperate to load metadata from global memory\n for (int i = threadIdx.x; i < num_chunks_this_block; i += blockDim.x) {\n int global_chunk_id = start_chunk + i;\n s_src_addrs[i] = src_addrs[global_chunk_id];\n s_dst_addrs[i] = dst_addrs[global_chunk_id];\n s_sizes[i] = sizes[global_chunk_id];\n }\n __syncthreads();\n\n // Phase 2: Process each chunk using metadata from shared memory\n for (int chunk_id = 0; chunk_id < num_chunks_this_block; chunk_id++) {\n uint8_t* src = s_src_addrs[chunk_id];\n uint8_t* dst = s_dst_addrs[chunk_id];\n size_t size = s_sizes[chunk_id];\n\n // Check if both src and dst are aligned to 16 bytes for float4 access\n uintptr_t src_addr = (uintptr_t)src;\n uintptr_t dst_addr = (uintptr_t)dst;\n bool is_aligned = ((src_addr % FLOAT4_SIZE) == 0) && ((dst_addr % FLOAT4_SIZE) == 0);\n\n if (is_aligned && size >= FLOAT4_SIZE) {\n // Fast path: vectorized float4 copies\n size_t aligned_size = (size / FLOAT4_SIZE) * FLOAT4_SIZE;\n\n // All 1024 threads cooperate on float4 copies\n #pragma unroll 4\n for (size_t offset = threadIdx.x * FLOAT4_SIZE;\n offset < aligned_size;\n offset += blockDim.x * FLOAT4_SIZE) {\n // Vectorized 16-byte load and store\n float4 data = *((float4*)(src + offset));\n *((float4*)(dst + offset)) = data;\n }\n\n // Handle remaining bytes (< 16 bytes) with byte-by-byte copy\n for (size_t offset = aligned_size + threadIdx.x;\n offset < size;\n offset += blockDim.x) {\n dst[offset] = src[offset];\n }\n } else {\n // Fallback path: byte-by-byte copy for unaligned addresses\n // Still use all threads for parallelism\n for (size_t offset = threadIdx.x; offset < size; offset += blockDim.x) {\n dst[offset] = src[offset];\n }\n }\n }\n\n // System-wide memory fence: ensures all writes from this thread are visible\n // to system agents (NIC DMA engines, other GPUs via PCIe/NVLink).\n // Required so that NVSHMEM put() on another CUDA stream can read correct\n // pack data from the send buffer. Without this, writes may remain in L2\n // cache, invisible to the NIC's DMA engine.\n __threadfence_system();\n}\n\n}\n\n\n", "source": "/kaggle/working/Megatron-LM/megatron/core/resharding/nvshmem_copy_service/kernels/chunked_kernel.cu"} {"input": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.\n * SPDX-License-Identifier: Apache-2.0\n */\n\n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \n\nnamespace MLCommon {\nnamespace Stats {\nnamespace opg {\n\ntemplate \nstatic __global__ void varsPartitionKernelColMajor(\n Type* var, const Type* data, const Type* mu, IdxType D, IdxType N, IdxType K)\n{\n typedef cub::BlockReduce BlockReduce;\n __shared__ typename BlockReduce::TempStorage temp_storage;\n Type thread_data = Type(0);\n IdxType colStart = N * blockIdx.x;\n Type m = mu[blockIdx.x];\n for (IdxType i = threadIdx.x; i < N; i += TPB) {\n IdxType idx = colStart + i;\n Type diff = data[idx] - m;\n thread_data += diff * diff;\n }\n Type acc = BlockReduce(temp_storage).Sum(thread_data);\n if (threadIdx.x == 0) { var[blockIdx.x] = acc / K; }\n}\n\ntemplate \nvoid varPartition(\n Type* std, const Type* data, const Type* mu, IdxType D, IdxType N, IdxType K, cudaStream_t stream)\n{\n static const int TPB = 256;\n varsPartitionKernelColMajor<<>>(std, data, mu, D, N, K);\n\n RAFT_CUDA_TRY(cudaPeekAtLastError());\n}\n\ntemplate \nvoid var_impl(const raft::handle_t& handle,\n Matrix::Data& out,\n const std::vector*>& in,\n const Matrix::PartDescriptor& inDesc,\n const T* mu,\n cudaStream_t* streams,\n int n_streams)\n{\n auto& comm = handle.get_comms();\n\n rmm::device_uvector local_vars_tmp(in.size() * inDesc.N, streams[0]);\n rmm::device_uvector local_vars_tmp_t(in.size() * inDesc.N, streams[0]);\n\n std::vector localBlocks = inDesc.blocksOwnedBy(comm.get_rank());\n\n for (size_t i = 0; i < localBlocks.size(); i++) {\n T* loc = local_vars_tmp.data() + (i * inDesc.N);\n Stats::opg::varPartition(\n loc, in[i]->ptr, mu, inDesc.N, localBlocks[i]->size, inDesc.M, streams[i % n_streams]);\n }\n\n for (int i = 0; i < n_streams; i++) {\n RAFT_CUDA_TRY(cudaStreamSynchronize(streams[i]));\n }\n\n raft::linalg::transpose(handle,\n local_vars_tmp.data(),\n local_vars_tmp_t.data(),\n inDesc.N,\n localBlocks.size(),\n streams[0]);\n\n raft::stats::sum(\n local_vars_tmp.data(), local_vars_tmp_t.data(), inDesc.N, localBlocks.size(), streams[0]);\n\n comm.allreduce(local_vars_tmp.data(), out.ptr, inDesc.N, raft::comms::op_t::SUM, streams[0]);\n\n comm.sync_stream(streams[0]);\n}\n\nvoid var(const raft::handle_t& handle,\n Matrix::Data& out,\n const std::vector*>& in,\n const Matrix::PartDescriptor& inDesc,\n const double* mu,\n cudaStream_t* streams,\n int n_streams)\n{\n var_impl(handle, out, in, inDesc, mu, streams, n_streams);\n}\n\nvoid var(const raft::handle_t& handle,\n Matrix::Data& out,\n const std::vector*>& in,\n const Matrix::PartDescriptor& inDesc,\n const float* mu,\n cudaStream_t* streams,\n int n_streams)\n{\n var_impl(handle, out, in, inDesc, mu, streams, n_streams);\n}\n\n}; // namespace opg\n}; // namespace Stats\n}; // namespace MLCommon\n", "output": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.\n * SPDX-License-Identifier: Apache-2.0\n */\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \n\nnamespace MLCommon {\nnamespace Stats {\nnamespace opg {\n\ntemplate \nstatic __global__ void varsPartitionKernelColMajor(\n Type* var, const Type* data, const Type* mu, IdxType D, IdxType N, IdxType K)\n{\n typedef hipcub::BlockReduce BlockReduce;\n __shared__ typename BlockReduce::TempStorage temp_storage;\n Type thread_data = Type(0);\n IdxType colStart = N * blockIdx.x;\n Type m = mu[blockIdx.x];\n for (IdxType i = threadIdx.x; i < N; i += TPB) {\n IdxType idx = colStart + i;\n Type diff = data[idx] - m;\n thread_data += diff * diff;\n }\n Type acc = BlockReduce(temp_storage).Sum(thread_data);\n if (threadIdx.x == 0) { var[blockIdx.x] = acc / K; }\n}\n\ntemplate \nvoid varPartition(\n Type* std, const Type* data, const Type* mu, IdxType D, IdxType N, IdxType K, hipStream_t stream)\n{\n static const int TPB = 256;\n varsPartitionKernelColMajor<<>>(std, data, mu, D, N, K);\n\n RAFT_CUDA_TRY(hipPeekAtLastError());\n}\n\ntemplate \nvoid var_impl(const raft::handle_t& handle,\n Matrix::Data& out,\n const std::vector*>& in,\n const Matrix::PartDescriptor& inDesc,\n const T* mu,\n hipStream_t* streams,\n int n_streams)\n{\n auto& comm = handle.get_comms();\n\n rmm::device_uvector local_vars_tmp(in.size() * inDesc.N, streams[0]);\n rmm::device_uvector local_vars_tmp_t(in.size() * inDesc.N, streams[0]);\n\n std::vector localBlocks = inDesc.blocksOwnedBy(comm.get_rank());\n\n for (size_t i = 0; i < localBlocks.size(); i++) {\n T* loc = local_vars_tmp.data() + (i * inDesc.N);\n Stats::opg::varPartition(\n loc, in[i]->ptr, mu, inDesc.N, localBlocks[i]->size, inDesc.M, streams[i % n_streams]);\n }\n\n for (int i = 0; i < n_streams; i++) {\n RAFT_CUDA_TRY(hipStreamSynchronize(streams[i]));\n }\n\n raft::linalg::transpose(handle,\n local_vars_tmp.data(),\n local_vars_tmp_t.data(),\n inDesc.N,\n localBlocks.size(),\n streams[0]);\n\n raft::stats::sum(\n local_vars_tmp.data(), local_vars_tmp_t.data(), inDesc.N, localBlocks.size(), streams[0]);\n\n comm.allreduce(local_vars_tmp.data(), out.ptr, inDesc.N, raft::comms::op_t::SUM, streams[0]);\n\n comm.sync_stream(streams[0]);\n}\n\nvoid var(const raft::handle_t& handle,\n Matrix::Data& out,\n const std::vector*>& in,\n const Matrix::PartDescriptor& inDesc,\n const double* mu,\n hipStream_t* streams,\n int n_streams)\n{\n var_impl(handle, out, in, inDesc, mu, streams, n_streams);\n}\n\nvoid var(const raft::handle_t& handle,\n Matrix::Data& out,\n const std::vector*>& in,\n const Matrix::PartDescriptor& inDesc,\n const float* mu,\n hipStream_t* streams,\n int n_streams)\n{\n var_impl(handle, out, in, inDesc, mu, streams, n_streams);\n}\n\n}; // namespace opg\n}; // namespace Stats\n}; // namespace MLCommon\n", "source": "/kaggle/working/cuml/cpp/src_prims/opg/stats/stddev.cu"} {"input": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.\n * SPDX-License-Identifier: Apache-2.0\n */\n#include \n\n#include \n#include \n\n#include \n#include \n#include \n\nnamespace MLCommon {\nnamespace Matrix {\nnamespace opg {\n\ntemplate \nvoid gatherPart(const raft::handle_t& h,\n T* gatheredPart,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int partIndex,\n int rootRank,\n int myRank,\n cudaStream_t stream)\n{\n ASSERT(partIndex < desc.partsToRanks.size(),\n \"MLCommon::Matrix::opg::gatherPart: Part index is out of range\");\n const auto& comm = h.get_comms();\n std::vector requests;\n int items = desc.partsToRanks[partIndex]->size * desc.N;\n comm.sync_stream(stream);\n if (desc.partsToRanks[partIndex]->rank == myRank) {\n requests.resize(requests.size() + 1);\n int localIndex = 0;\n for (int i = 0; i < partIndex; i++) {\n if (desc.partsToRanks[i]->rank == myRank) { localIndex++; }\n }\n\n comm.isend(parts[localIndex]->ptr, items, rootRank, 0, &requests.back());\n }\n if (myRank == rootRank) {\n requests.resize(requests.size() + 1);\n comm.irecv(gatheredPart, items, desc.partsToRanks[partIndex]->rank, 0, &requests.back());\n }\n comm.waitall(requests.size(), requests.data());\n}\n\ntemplate \nvoid allGatherPart(const raft::handle_t& h,\n T* gatheredPart,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int partIndex,\n int myRank,\n cudaStream_t stream)\n{\n ASSERT(partIndex < desc.partsToRanks.size(),\n \"MLCommon::Matrix::opg::allGatherPart: Part index is out of range\");\n\n const auto& comm = h.get_comms();\n\n int partLength = desc.partsToRanks[partIndex]->size * desc.N;\n\n if (desc.partsToRanks[partIndex]->rank == myRank) {\n int localIndex = 0;\n for (int i = 0; i < partIndex; i++) {\n if (desc.partsToRanks[i]->rank == myRank) { localIndex++; }\n }\n raft::copy(gatheredPart, parts[localIndex]->ptr, partLength, stream);\n comm.sync_stream(stream);\n }\n comm.bcast(gatheredPart, partLength, desc.partsToRanks[partIndex]->rank, stream);\n}\n\ntemplate \nvoid gather(const raft::handle_t& h,\n T* gatheredMatrix,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int rootRank,\n int myRank,\n cudaStream_t stream)\n{\n size_t maxPartSize = 0;\n for (size_t i = 0; i < desc.partsToRanks.size(); i++) {\n if (maxPartSize < desc.partsToRanks[i]->size) { maxPartSize = desc.partsToRanks[i]->size; }\n }\n\n rmm::device_uvector partBuffer(0, stream);\n if (myRank == rootRank) { partBuffer.resize(maxPartSize * desc.N, stream); }\n size_t offset = 0;\n for (size_t i = 0; i < desc.partsToRanks.size(); i++) {\n gatherPart(h, partBuffer.data(), parts, desc, i, rootRank, myRank, stream);\n\n if (myRank == rootRank) {\n if (desc.layout == Matrix::LayoutRowMajor) {\n raft::copy(\n gatheredMatrix + offset, partBuffer.data(), desc.partsToRanks[i]->size * desc.N, stream);\n } else {\n raft::linalg::transpose(h,\n partBuffer.data(),\n gatheredMatrix + offset,\n (int)desc.partsToRanks[i]->size,\n (int)desc.N,\n stream);\n }\n }\n offset += desc.partsToRanks[i]->size * desc.N;\n }\n RAFT_CUDA_TRY(cudaStreamSynchronize(stream));\n}\n\ntemplate \nvoid allGather(const raft::handle_t& h,\n T* gatheredMatrix,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream)\n{\n size_t maxPartSize = 0;\n for (size_t i = 0; i < desc.partsToRanks.size(); i++) {\n if (maxPartSize < desc.partsToRanks[i]->size) { maxPartSize = desc.partsToRanks[i]->size; }\n }\n\n rmm::device_uvector partBuffer(maxPartSize * desc.N, stream);\n size_t offset = 0;\n for (size_t i = 0; i < desc.partsToRanks.size(); i++) {\n allGatherPart(h, partBuffer.data(), parts, desc, i, myRank, stream);\n if (desc.layout == Matrix::LayoutRowMajor) {\n raft::copy(\n gatheredMatrix + offset, partBuffer.data(), desc.partsToRanks[i]->size * desc.N, stream);\n } else {\n raft::linalg::transpose(h,\n partBuffer.data(),\n gatheredMatrix + offset,\n (int)desc.partsToRanks[i]->size,\n (int)desc.N,\n stream);\n }\n offset += desc.partsToRanks[i]->size * desc.N;\n }\n}\n\ntemplate \nvoid allocate(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream)\n{\n auto allocator = rmm::mr::get_current_device_resource_ref();\n for (size_t i = 0; i < desc.partsToRanks.size(); i++) {\n if (myRank == desc.partsToRanks[i]->rank) {\n int partSize = desc.partsToRanks[i]->size * desc.N;\n T* partMem = (T*)allocator.allocate(stream, partSize * sizeof(T));\n Matrix::Data* part = new Matrix::Data(partMem, partSize);\n parts.push_back(part);\n }\n }\n return;\n}\n\ntemplate \nvoid deallocate(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream)\n{\n auto allocator = rmm::mr::get_current_device_resource_ref();\n for (size_t i = 0, localIndex = 0; i < desc.partsToRanks.size(); i++) {\n if (myRank == desc.partsToRanks[i]->rank) {\n int partSize = desc.partsToRanks[i]->size * desc.N;\n allocator.deallocate(stream, parts[localIndex]->ptr, partSize * sizeof(T));\n localIndex++;\n }\n }\n for (size_t i = 0, localIndex = 0; i < desc.partsToRanks.size(); i++) {\n if (myRank == desc.partsToRanks[i]->rank) {\n delete parts[localIndex];\n localIndex++;\n }\n }\n return;\n}\n\ntemplate \nvoid randomize(const raft::handle_t& h,\n raft::random::Rng& r,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream,\n T low = T(-1.0),\n T high = T(1.0))\n{\n for (size_t i = 0, localIndex = 0; i < desc.partsToRanks.size(); i++) {\n if (myRank == desc.partsToRanks[i]->rank) {\n int partSize = desc.partsToRanks[i]->size * desc.N;\n r.uniform(parts[localIndex]->ptr, partSize, low, high, stream);\n localIndex++;\n }\n }\n}\n\ntemplate \nvoid reset(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream)\n{\n for (size_t i = 0, localIndex = 0; i < desc.partsToRanks.size(); i++) {\n if (myRank == desc.partsToRanks[i]->rank) {\n size_t partSize = desc.partsToRanks[i]->size * desc.N * sizeof(T);\n RAFT_CUDA_TRY(cudaMemsetAsync(parts[localIndex]->ptr, 0, partSize, stream));\n localIndex++;\n }\n }\n}\n\ntemplate \nstatic __global__ void printRaw2DKernel(T* buffer, int rows, int cols, bool isColMajor)\n{\n for (int r = 0; r < rows; r++) {\n for (int c = 0; c < cols; c++) {\n if (isColMajor) {\n printf(\"%8.6f \", buffer[c * rows + r]);\n } else {\n printf(\"%8.6f \", buffer[r * cols + c]);\n }\n }\n printf(\"\\n\");\n }\n return;\n}\n\ntemplate \nvoid printRaw2D(T* buffer, int rows, int cols, bool isColMajor, cudaStream_t stream)\n{\n printRaw2DKernel<<<1, 1, 0, stream>>>(buffer, rows, cols, isColMajor);\n}\n\ntemplate \nvoid print(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n const char* matrixName,\n int myRank,\n cudaStream_t stream)\n{\n rmm::device_uvector buffer(desc.M * desc.N, stream);\n if (myRank == 0) { buffer.resize(desc.M * desc.N, stream); }\n Matrix::opg::gather(h, buffer.data(), parts, desc, 0, myRank, stream);\n if (myRank == 0) {\n printf(\"%s = [\", matrixName);\n printRaw2D(buffer.data(), desc.M, desc.N, false, stream);\n RAFT_CUDA_TRY(cudaStreamSynchronize(stream));\n printf(\"];\\n#-------------------------------\\n\");\n }\n}\n\n// Instantiations\n\nvoid gatherPart(const raft::handle_t& h,\n float* gatheredPart,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int partIndex,\n int rootRank,\n int myRank,\n cudaStream_t stream)\n{\n gatherPart(h, gatheredPart, parts, desc, partIndex, rootRank, myRank, stream);\n}\n\nvoid allGatherPart(const raft::handle_t& h,\n float* gatheredPart,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int partIndex,\n int myRank,\n cudaStream_t stream)\n{\n allGatherPart(h, gatheredPart, parts, desc, partIndex, myRank, stream);\n}\n\nvoid gather(const raft::handle_t& h,\n float* gatheredMatrix,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int rootRank,\n int myRank,\n cudaStream_t stream)\n{\n gather(h, gatheredMatrix, parts, desc, rootRank, myRank, stream);\n}\n\nvoid allGather(const raft::handle_t& h,\n float* gatheredMatrix,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream)\n{\n allGather(h, gatheredMatrix, parts, desc, myRank, stream);\n}\n\nvoid allocate(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream)\n{\n allocate(h, parts, desc, myRank, stream);\n return;\n}\n\nvoid deallocate(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream)\n{\n deallocate(h, parts, desc, myRank, stream);\n}\n\nvoid randomize(const raft::handle_t& h,\n raft::random::Rng& r,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream,\n float low,\n float high)\n{\n randomize(h, r, parts, desc, myRank, stream, low, high);\n}\n\nvoid reset(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream)\n{\n reset(h, parts, desc, myRank, stream);\n}\n\nvoid printRaw2D(float* buffer, int rows, int cols, bool isColMajor, cudaStream_t stream)\n{\n printRaw2D(buffer, rows, cols, isColMajor, stream);\n}\n\nvoid print(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n const char* matrixName,\n int myRank,\n cudaStream_t stream)\n{\n print(h, parts, desc, matrixName, myRank, stream);\n}\n\nvoid gatherPart(const raft::handle_t& h,\n double* gatheredPart,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int partIndex,\n int rootRank,\n int myRank,\n cudaStream_t stream)\n{\n gatherPart(h, gatheredPart, parts, desc, partIndex, rootRank, myRank, stream);\n}\n\nvoid allGatherPart(const raft::handle_t& h,\n double* gatheredPart,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int partIndex,\n int myRank,\n cudaStream_t stream)\n{\n allGatherPart(h, gatheredPart, parts, desc, partIndex, myRank, stream);\n}\n\nvoid gather(const raft::handle_t& h,\n double* gatheredMatrix,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int rootRank,\n int myRank,\n cudaStream_t stream)\n{\n gather(h, gatheredMatrix, parts, desc, rootRank, myRank, stream);\n}\n\nvoid allGather(const raft::handle_t& h,\n double* gatheredMatrix,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream)\n{\n allGather(h, gatheredMatrix, parts, desc, myRank, stream);\n}\n\nvoid allocate(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream)\n{\n allocate(h, parts, desc, myRank, stream);\n}\n\nvoid deallocate(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream)\n{\n deallocate(h, parts, desc, myRank, stream);\n}\n\nvoid randomize(const raft::handle_t& h,\n raft::random::Rng& r,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream,\n double low,\n double high)\n{\n randomize(h, r, parts, desc, myRank, stream, low, high);\n}\n\nvoid reset(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n cudaStream_t stream)\n{\n reset(h, parts, desc, myRank, stream);\n}\n\nvoid printRaw2D(double* buffer, int rows, int cols, bool isColMajor, cudaStream_t stream)\n{\n printRaw2D(buffer, rows, cols, isColMajor, stream);\n}\n\nvoid print(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n const char* matrixName,\n int myRank,\n cudaStream_t stream)\n{\n print(h, parts, desc, matrixName, myRank, stream);\n}\n\n} // end namespace opg\n} // end namespace Matrix\n} // end namespace MLCommon\n", "output": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.\n * SPDX-License-Identifier: Apache-2.0\n */\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \n#include \n\n#include \n#include \n#include \n\nnamespace MLCommon {\nnamespace Matrix {\nnamespace opg {\n\ntemplate \nvoid gatherPart(const raft::handle_t& h,\n T* gatheredPart,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int partIndex,\n int rootRank,\n int myRank,\n hipStream_t stream)\n{\n ASSERT(partIndex < desc.partsToRanks.size(),\n \"MLCommon::Matrix::opg::gatherPart: Part index is out of range\");\n const auto& comm = h.get_comms();\n std::vector requests;\n int items = desc.partsToRanks[partIndex]->size * desc.N;\n comm.sync_stream(stream);\n if (desc.partsToRanks[partIndex]->rank == myRank) {\n requests.resize(requests.size() + 1);\n int localIndex = 0;\n for (int i = 0; i < partIndex; i++) {\n if (desc.partsToRanks[i]->rank == myRank) { localIndex++; }\n }\n\n comm.isend(parts[localIndex]->ptr, items, rootRank, 0, &requests.back());\n }\n if (myRank == rootRank) {\n requests.resize(requests.size() + 1);\n comm.irecv(gatheredPart, items, desc.partsToRanks[partIndex]->rank, 0, &requests.back());\n }\n comm.waitall(requests.size(), requests.data());\n}\n\ntemplate \nvoid allGatherPart(const raft::handle_t& h,\n T* gatheredPart,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int partIndex,\n int myRank,\n hipStream_t stream)\n{\n ASSERT(partIndex < desc.partsToRanks.size(),\n \"MLCommon::Matrix::opg::allGatherPart: Part index is out of range\");\n\n const auto& comm = h.get_comms();\n\n int partLength = desc.partsToRanks[partIndex]->size * desc.N;\n\n if (desc.partsToRanks[partIndex]->rank == myRank) {\n int localIndex = 0;\n for (int i = 0; i < partIndex; i++) {\n if (desc.partsToRanks[i]->rank == myRank) { localIndex++; }\n }\n raft::copy(gatheredPart, parts[localIndex]->ptr, partLength, stream);\n comm.sync_stream(stream);\n }\n comm.bcast(gatheredPart, partLength, desc.partsToRanks[partIndex]->rank, stream);\n}\n\ntemplate \nvoid gather(const raft::handle_t& h,\n T* gatheredMatrix,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int rootRank,\n int myRank,\n hipStream_t stream)\n{\n size_t maxPartSize = 0;\n for (size_t i = 0; i < desc.partsToRanks.size(); i++) {\n if (maxPartSize < desc.partsToRanks[i]->size) { maxPartSize = desc.partsToRanks[i]->size; }\n }\n\n rmm::device_uvector partBuffer(0, stream);\n if (myRank == rootRank) { partBuffer.resize(maxPartSize * desc.N, stream); }\n size_t offset = 0;\n for (size_t i = 0; i < desc.partsToRanks.size(); i++) {\n gatherPart(h, partBuffer.data(), parts, desc, i, rootRank, myRank, stream);\n\n if (myRank == rootRank) {\n if (desc.layout == Matrix::LayoutRowMajor) {\n raft::copy(\n gatheredMatrix + offset, partBuffer.data(), desc.partsToRanks[i]->size * desc.N, stream);\n } else {\n raft::linalg::transpose(h,\n partBuffer.data(),\n gatheredMatrix + offset,\n (int)desc.partsToRanks[i]->size,\n (int)desc.N,\n stream);\n }\n }\n offset += desc.partsToRanks[i]->size * desc.N;\n }\n RAFT_CUDA_TRY(hipStreamSynchronize(stream));\n}\n\ntemplate \nvoid allGather(const raft::handle_t& h,\n T* gatheredMatrix,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream)\n{\n size_t maxPartSize = 0;\n for (size_t i = 0; i < desc.partsToRanks.size(); i++) {\n if (maxPartSize < desc.partsToRanks[i]->size) { maxPartSize = desc.partsToRanks[i]->size; }\n }\n\n rmm::device_uvector partBuffer(maxPartSize * desc.N, stream);\n size_t offset = 0;\n for (size_t i = 0; i < desc.partsToRanks.size(); i++) {\n allGatherPart(h, partBuffer.data(), parts, desc, i, myRank, stream);\n if (desc.layout == Matrix::LayoutRowMajor) {\n raft::copy(\n gatheredMatrix + offset, partBuffer.data(), desc.partsToRanks[i]->size * desc.N, stream);\n } else {\n raft::linalg::transpose(h,\n partBuffer.data(),\n gatheredMatrix + offset,\n (int)desc.partsToRanks[i]->size,\n (int)desc.N,\n stream);\n }\n offset += desc.partsToRanks[i]->size * desc.N;\n }\n}\n\ntemplate \nvoid allocate(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream)\n{\n auto allocator = rmm::mr::get_current_device_resource_ref();\n for (size_t i = 0; i < desc.partsToRanks.size(); i++) {\n if (myRank == desc.partsToRanks[i]->rank) {\n int partSize = desc.partsToRanks[i]->size * desc.N;\n T* partMem = (T*)allocator.allocate(stream, partSize * sizeof(T));\n Matrix::Data* part = new Matrix::Data(partMem, partSize);\n parts.push_back(part);\n }\n }\n return;\n}\n\ntemplate \nvoid deallocate(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream)\n{\n auto allocator = rmm::mr::get_current_device_resource_ref();\n for (size_t i = 0, localIndex = 0; i < desc.partsToRanks.size(); i++) {\n if (myRank == desc.partsToRanks[i]->rank) {\n int partSize = desc.partsToRanks[i]->size * desc.N;\n allocator.deallocate(stream, parts[localIndex]->ptr, partSize * sizeof(T));\n localIndex++;\n }\n }\n for (size_t i = 0, localIndex = 0; i < desc.partsToRanks.size(); i++) {\n if (myRank == desc.partsToRanks[i]->rank) {\n delete parts[localIndex];\n localIndex++;\n }\n }\n return;\n}\n\ntemplate \nvoid randomize(const raft::handle_t& h,\n raft::random::Rng& r,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream,\n T low = T(-1.0),\n T high = T(1.0))\n{\n for (size_t i = 0, localIndex = 0; i < desc.partsToRanks.size(); i++) {\n if (myRank == desc.partsToRanks[i]->rank) {\n int partSize = desc.partsToRanks[i]->size * desc.N;\n r.uniform(parts[localIndex]->ptr, partSize, low, high, stream);\n localIndex++;\n }\n }\n}\n\ntemplate \nvoid reset(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream)\n{\n for (size_t i = 0, localIndex = 0; i < desc.partsToRanks.size(); i++) {\n if (myRank == desc.partsToRanks[i]->rank) {\n size_t partSize = desc.partsToRanks[i]->size * desc.N * sizeof(T);\n RAFT_CUDA_TRY(hipMemsetAsync(parts[localIndex]->ptr, 0, partSize, stream));\n localIndex++;\n }\n }\n}\n\ntemplate \nstatic __global__ void printRaw2DKernel(T* buffer, int rows, int cols, bool isColMajor)\n{\n for (int r = 0; r < rows; r++) {\n for (int c = 0; c < cols; c++) {\n if (isColMajor) {\n printf(\"%8.6f \", buffer[c * rows + r]);\n } else {\n printf(\"%8.6f \", buffer[r * cols + c]);\n }\n }\n printf(\"\\n\");\n }\n return;\n}\n\ntemplate \nvoid printRaw2D(T* buffer, int rows, int cols, bool isColMajor, hipStream_t stream)\n{\n printRaw2DKernel<<<1, 1, 0, stream>>>(buffer, rows, cols, isColMajor);\n}\n\ntemplate \nvoid print(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n const char* matrixName,\n int myRank,\n hipStream_t stream)\n{\n rmm::device_uvector buffer(desc.M * desc.N, stream);\n if (myRank == 0) { buffer.resize(desc.M * desc.N, stream); }\n Matrix::opg::gather(h, buffer.data(), parts, desc, 0, myRank, stream);\n if (myRank == 0) {\n printf(\"%s = [\", matrixName);\n printRaw2D(buffer.data(), desc.M, desc.N, false, stream);\n RAFT_CUDA_TRY(hipStreamSynchronize(stream));\n printf(\"];\\n#-------------------------------\\n\");\n }\n}\n\n// Instantiations\n\nvoid gatherPart(const raft::handle_t& h,\n float* gatheredPart,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int partIndex,\n int rootRank,\n int myRank,\n hipStream_t stream)\n{\n gatherPart(h, gatheredPart, parts, desc, partIndex, rootRank, myRank, stream);\n}\n\nvoid allGatherPart(const raft::handle_t& h,\n float* gatheredPart,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int partIndex,\n int myRank,\n hipStream_t stream)\n{\n allGatherPart(h, gatheredPart, parts, desc, partIndex, myRank, stream);\n}\n\nvoid gather(const raft::handle_t& h,\n float* gatheredMatrix,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int rootRank,\n int myRank,\n hipStream_t stream)\n{\n gather(h, gatheredMatrix, parts, desc, rootRank, myRank, stream);\n}\n\nvoid allGather(const raft::handle_t& h,\n float* gatheredMatrix,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream)\n{\n allGather(h, gatheredMatrix, parts, desc, myRank, stream);\n}\n\nvoid allocate(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream)\n{\n allocate(h, parts, desc, myRank, stream);\n return;\n}\n\nvoid deallocate(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream)\n{\n deallocate(h, parts, desc, myRank, stream);\n}\n\nvoid randomize(const raft::handle_t& h,\n raft::random::Rng& r,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream,\n float low,\n float high)\n{\n randomize(h, r, parts, desc, myRank, stream, low, high);\n}\n\nvoid reset(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream)\n{\n reset(h, parts, desc, myRank, stream);\n}\n\nvoid printRaw2D(float* buffer, int rows, int cols, bool isColMajor, hipStream_t stream)\n{\n printRaw2D(buffer, rows, cols, isColMajor, stream);\n}\n\nvoid print(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n const char* matrixName,\n int myRank,\n hipStream_t stream)\n{\n print(h, parts, desc, matrixName, myRank, stream);\n}\n\nvoid gatherPart(const raft::handle_t& h,\n double* gatheredPart,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int partIndex,\n int rootRank,\n int myRank,\n hipStream_t stream)\n{\n gatherPart(h, gatheredPart, parts, desc, partIndex, rootRank, myRank, stream);\n}\n\nvoid allGatherPart(const raft::handle_t& h,\n double* gatheredPart,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int partIndex,\n int myRank,\n hipStream_t stream)\n{\n allGatherPart(h, gatheredPart, parts, desc, partIndex, myRank, stream);\n}\n\nvoid gather(const raft::handle_t& h,\n double* gatheredMatrix,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int rootRank,\n int myRank,\n hipStream_t stream)\n{\n gather(h, gatheredMatrix, parts, desc, rootRank, myRank, stream);\n}\n\nvoid allGather(const raft::handle_t& h,\n double* gatheredMatrix,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream)\n{\n allGather(h, gatheredMatrix, parts, desc, myRank, stream);\n}\n\nvoid allocate(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream)\n{\n allocate(h, parts, desc, myRank, stream);\n}\n\nvoid deallocate(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream)\n{\n deallocate(h, parts, desc, myRank, stream);\n}\n\nvoid randomize(const raft::handle_t& h,\n raft::random::Rng& r,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream,\n double low,\n double high)\n{\n randomize(h, r, parts, desc, myRank, stream, low, high);\n}\n\nvoid reset(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n int myRank,\n hipStream_t stream)\n{\n reset(h, parts, desc, myRank, stream);\n}\n\nvoid printRaw2D(double* buffer, int rows, int cols, bool isColMajor, hipStream_t stream)\n{\n printRaw2D(buffer, rows, cols, isColMajor, stream);\n}\n\nvoid print(const raft::handle_t& h,\n std::vector*>& parts,\n Matrix::PartDescriptor& desc,\n const char* matrixName,\n int myRank,\n hipStream_t stream)\n{\n print(h, parts, desc, matrixName, myRank, stream);\n}\n\n} // end namespace opg\n} // end namespace Matrix\n} // end namespace MLCommon\n", "source": "/kaggle/working/cuml/cpp/src_prims/opg/matrix/matrix_utils.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n\n#if THRUST_CPP_DIALECT >= 2011\n#include \n#endif\n\n// This example demonstrates two ways to achieve algorithm invocations that are asynchronous with\n// the calling thread.\n//\n// The first method wraps a call to thrust::reduce inside a __global__ function. Since __global__ function\n// launches are asynchronous with the launching thread, this achieves asynchrony. The result of the reduction\n// is stored to a pointer to CUDA global memory. The calling thread waits for the result of the reduction to \n// be ready by synchronizing with the CUDA stream on which the __global__ function is launched.\n//\n// The second method uses the C++11 library function, std::async, to create concurrency. The lambda function\n// given to std::async returns the result of thrust::reduce to a std::future. The calling thread can use the\n// std::future to wait for the result of the reduction. This method requires a compiler which supports\n// C++11-capable language and library constructs.\n\n#ifdef THRUST_EXAMPLE_DEVICE_SIDE\ntemplate\n__global__ void reduce_kernel(Iterator first, Iterator last, T init, BinaryOperation binary_op, Pointer result)\n{\n *result = thrust::reduce(thrust::cuda::par, first, last, init, binary_op);\n}\n#endif\n\nint main()\n{\n size_t n = 1 << 20;\n thrust::device_vector data(n, 1);\n thrust::device_vector result(1, 0);\n\n // method 1: call thrust::reduce from an asynchronous CUDA kernel launch\n\n // create a CUDA stream \n cudaStream_t s;\n cudaStreamCreate(&s);\n\n // launch a CUDA kernel with only 1 thread on our stream\n#ifdef THRUST_EXAMPLE_DEVICE_SIDE\n reduce_kernel<<<1,1,0,s>>>(data.begin(), data.end(), 0, thrust::plus(), result.data());\n#else\n result[0] = thrust::reduce(thrust::cuda::par, data.begin(), data.end(), 0, thrust::plus());\n#endif\n\n // wait for the stream to finish\n cudaStreamSynchronize(s);\n\n // our result should be ready\n assert(result[0] == n);\n\n cudaStreamDestroy(s);\n\n // reset the result\n result[0] = 0;\n\n#if THRUST_CPP_DIALECT >= 2011\n // method 2: use std::async to create asynchrony\n\n // copy all the algorithm parameters\n auto begin = data.begin();\n auto end = data.end();\n unsigned int init = 0;\n auto binary_op = thrust::plus();\n\n // std::async captures the algorithm parameters by value\n // use std::launch::async to ensure the creation of a new thread\n std::future future_result = std::async(std::launch::async, [=]\n {\n return thrust::reduce(begin, end, init, binary_op);\n });\n\n // wait on the result and check that it is correct\n assert(future_result.get() == n);\n#endif\n\n return 0;\n}\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n\n#if THRUST_CPP_DIALECT >= 2011\n#include \n#endif\n\n// This example demonstrates two ways to achieve algorithm invocations that are asynchronous with\n// the calling thread.\n//\n// The first method wraps a call to thrust::reduce inside a __global__ function. Since __global__ function\n// launches are asynchronous with the launching thread, this achieves asynchrony. The result of the reduction\n// is stored to a pointer to CUDA global memory. The calling thread waits for the result of the reduction to \n// be ready by synchronizing with the CUDA stream on which the __global__ function is launched.\n//\n// The second method uses the C++11 library function, std::async, to create concurrency. The lambda function\n// given to std::async returns the result of thrust::reduce to a std::future. The calling thread can use the\n// std::future to wait for the result of the reduction. This method requires a compiler which supports\n// C++11-capable language and library constructs.\n\n#ifdef THRUST_EXAMPLE_DEVICE_SIDE\ntemplate\n__global__ void reduce_kernel(Iterator first, Iterator last, T init, BinaryOperation binary_op, Pointer result)\n{\n *result = thrust::reduce(thrust::cuda::par, first, last, init, binary_op);\n}\n#endif\n\nint main()\n{\n size_t n = 1 << 20;\n thrust::device_vector data(n, 1);\n thrust::device_vector result(1, 0);\n\n // method 1: call thrust::reduce from an asynchronous CUDA kernel launch\n\n // create a CUDA stream \n hipStream_t s;\n hipStreamCreate(&s);\n\n // launch a CUDA kernel with only 1 thread on our stream\n#ifdef THRUST_EXAMPLE_DEVICE_SIDE\n reduce_kernel<<<1,1,0,s>>>(data.begin(), data.end(), 0, thrust::plus(), result.data());\n#else\n result[0] = thrust::reduce(thrust::cuda::par, data.begin(), data.end(), 0, thrust::plus());\n#endif\n\n // wait for the stream to finish\n hipStreamSynchronize(s);\n\n // our result should be ready\n assert(result[0] == n);\n\n hipStreamDestroy(s);\n\n // reset the result\n result[0] = 0;\n\n#if THRUST_CPP_DIALECT >= 2011\n // method 2: use std::async to create asynchrony\n\n // copy all the algorithm parameters\n auto begin = data.begin();\n auto end = data.end();\n unsigned int init = 0;\n auto binary_op = thrust::plus();\n\n // std::async captures the algorithm parameters by value\n // use std::launch::async to ensure the creation of a new thread\n std::future future_result = std::async(std::launch::async, [=]\n {\n return thrust::reduce(begin, end, init, binary_op);\n });\n\n // wait on the result and check that it is correct\n assert(future_result.get() == n);\n#endif\n\n return 0;\n}\n\n", "source": "/kaggle/working/thrust/examples/cuda/async_reduce.cu"} {"input": "#include \n\n#include \n#include \n#include \n#include \n#include \n\n\nTHRUST_DISABLE_MSVC_POSSIBLE_LOSS_OF_DATA_WARNING_BEGIN\n\n\ntemplate \n#if THRUST_DEVICE_SYSTEM == THRUST_DEVICE_SYSTEM_CUDA\n__global__\n#endif \nvoid simple_copy_on_device(Iterator1 first1, Iterator1 last1, Iterator2 first2)\n{\n while(first1 != last1)\n *(first2++) = *(first1++);\n}\n\ntemplate \nvoid simple_copy(Iterator1 first1, Iterator1 last1, Iterator2 first2)\n{\n#if THRUST_DEVICE_SYSTEM == THRUST_DEVICE_SYSTEM_CUDA\n simple_copy_on_device<<<1,1>>>(first1, last1, first2);\n#else\n simple_copy_on_device(first1, last1, first2);\n#endif\n}\n\n\nvoid TestDeviceDereferenceDeviceVectorIterator(void)\n{\n thrust::device_vector input = unittest::random_integers(100); \n thrust::device_vector output(input.size(), 0);\n\n simple_copy(input.begin(), input.end(), output.begin());\n\n ASSERT_EQUAL(input, output);\n}\nDECLARE_UNITTEST(TestDeviceDereferenceDeviceVectorIterator);\n\nvoid TestDeviceDereferenceDevicePtr(void)\n{\n thrust::device_vector input = unittest::random_integers(100); \n thrust::device_vector output(input.size(), 0);\n\n thrust::device_ptr _first1 = &input[0];\n thrust::device_ptr _last1 = _first1 + input.size();\n thrust::device_ptr _first2 = &output[0];\n\n simple_copy(_first1, _last1, _first2);\n\n ASSERT_EQUAL(input, output);\n}\nDECLARE_UNITTEST(TestDeviceDereferenceDevicePtr);\n\nvoid TestDeviceDereferenceTransformIterator(void)\n{\n thrust::device_vector input = unittest::random_integers(100); \n thrust::device_vector output(input.size(), 0);\n\n simple_copy(thrust::make_transform_iterator(input.begin(), thrust::identity()),\n thrust::make_transform_iterator(input.end (), thrust::identity()),\n output.begin());\n\n ASSERT_EQUAL(input, output);\n}\nDECLARE_UNITTEST(TestDeviceDereferenceTransformIterator);\n\nvoid TestDeviceDereferenceCountingIterator(void)\n{\n thrust::counting_iterator first(1);\n thrust::counting_iterator last(6);\n\n thrust::device_vector output(5);\n\n simple_copy(first, last, output.begin());\n\n ASSERT_EQUAL(output[0], 1);\n ASSERT_EQUAL(output[1], 2);\n ASSERT_EQUAL(output[2], 3);\n ASSERT_EQUAL(output[3], 4);\n ASSERT_EQUAL(output[4], 5);\n}\nDECLARE_UNITTEST(TestDeviceDereferenceCountingIterator);\n\nvoid TestDeviceDereferenceTransformedCountingIterator(void)\n{\n thrust::counting_iterator first(1);\n thrust::counting_iterator last(6);\n\n thrust::device_vector output(5);\n\n simple_copy(thrust::make_transform_iterator(first, thrust::negate()),\n thrust::make_transform_iterator(last, thrust::negate()),\n output.begin());\n\n ASSERT_EQUAL(output[0], -1);\n ASSERT_EQUAL(output[1], -2);\n ASSERT_EQUAL(output[2], -3);\n ASSERT_EQUAL(output[3], -4);\n ASSERT_EQUAL(output[4], -5);\n}\nDECLARE_UNITTEST(TestDeviceDereferenceTransformedCountingIterator);\n\nTHRUST_DISABLE_MSVC_POSSIBLE_LOSS_OF_DATA_WARNING_END\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n\n#include \n#include \n#include \n#include \n#include \n\n\nTHRUST_DISABLE_MSVC_POSSIBLE_LOSS_OF_DATA_WARNING_BEGIN\n\n\ntemplate \n#if THRUST_DEVICE_SYSTEM == THRUST_DEVICE_SYSTEM_CUDA\n__global__\n#endif \nvoid simple_copy_on_device(Iterator1 first1, Iterator1 last1, Iterator2 first2)\n{\n while(first1 != last1)\n *(first2++) = *(first1++);\n}\n\ntemplate \nvoid simple_copy(Iterator1 first1, Iterator1 last1, Iterator2 first2)\n{\n#if THRUST_DEVICE_SYSTEM == THRUST_DEVICE_SYSTEM_CUDA\n simple_copy_on_device<<<1,1>>>(first1, last1, first2);\n#else\n simple_copy_on_device(first1, last1, first2);\n#endif\n}\n\n\nvoid TestDeviceDereferenceDeviceVectorIterator(void)\n{\n thrust::device_vector input = unittest::random_integers(100); \n thrust::device_vector output(input.size(), 0);\n\n simple_copy(input.begin(), input.end(), output.begin());\n\n ASSERT_EQUAL(input, output);\n}\nDECLARE_UNITTEST(TestDeviceDereferenceDeviceVectorIterator);\n\nvoid TestDeviceDereferenceDevicePtr(void)\n{\n thrust::device_vector input = unittest::random_integers(100); \n thrust::device_vector output(input.size(), 0);\n\n thrust::device_ptr _first1 = &input[0];\n thrust::device_ptr _last1 = _first1 + input.size();\n thrust::device_ptr _first2 = &output[0];\n\n simple_copy(_first1, _last1, _first2);\n\n ASSERT_EQUAL(input, output);\n}\nDECLARE_UNITTEST(TestDeviceDereferenceDevicePtr);\n\nvoid TestDeviceDereferenceTransformIterator(void)\n{\n thrust::device_vector input = unittest::random_integers(100); \n thrust::device_vector output(input.size(), 0);\n\n simple_copy(thrust::make_transform_iterator(input.begin(), thrust::identity()),\n thrust::make_transform_iterator(input.end (), thrust::identity()),\n output.begin());\n\n ASSERT_EQUAL(input, output);\n}\nDECLARE_UNITTEST(TestDeviceDereferenceTransformIterator);\n\nvoid TestDeviceDereferenceCountingIterator(void)\n{\n thrust::counting_iterator first(1);\n thrust::counting_iterator last(6);\n\n thrust::device_vector output(5);\n\n simple_copy(first, last, output.begin());\n\n ASSERT_EQUAL(output[0], 1);\n ASSERT_EQUAL(output[1], 2);\n ASSERT_EQUAL(output[2], 3);\n ASSERT_EQUAL(output[3], 4);\n ASSERT_EQUAL(output[4], 5);\n}\nDECLARE_UNITTEST(TestDeviceDereferenceCountingIterator);\n\nvoid TestDeviceDereferenceTransformedCountingIterator(void)\n{\n thrust::counting_iterator first(1);\n thrust::counting_iterator last(6);\n\n thrust::device_vector output(5);\n\n simple_copy(thrust::make_transform_iterator(first, thrust::negate()),\n thrust::make_transform_iterator(last, thrust::negate()),\n output.begin());\n\n ASSERT_EQUAL(output[0], -1);\n ASSERT_EQUAL(output[1], -2);\n ASSERT_EQUAL(output[2], -3);\n ASSERT_EQUAL(output[3], -4);\n ASSERT_EQUAL(output[4], -5);\n}\nDECLARE_UNITTEST(TestDeviceDereferenceTransformedCountingIterator);\n\nTHRUST_DISABLE_MSVC_POSSIBLE_LOSS_OF_DATA_WARNING_END\n", "source": "/kaggle/working/thrust/testing/dereference.cu"} {"input": "#include \n#include \n#include \n#include \n\n\ntemplate\nstruct is_equal_div_10_unique\n{\n __host__ __device__\n bool operator()(const T x, const T& y) const { return ((int) x / 10) == ((int) y / 10); }\n};\n\n\ntemplate\nvoid initialize_keys(Vector& keys)\n{\n keys.resize(9);\n keys[0] = 11;\n keys[1] = 11;\n keys[2] = 21;\n keys[3] = 20;\n keys[4] = 21;\n keys[5] = 21;\n keys[6] = 21;\n keys[7] = 37;\n keys[8] = 37;\n}\n\n\ntemplate\nvoid initialize_values(Vector& values)\n{\n values.resize(9);\n values[0] = 0; \n values[1] = 1;\n values[2] = 2;\n values[3] = 3;\n values[4] = 4;\n values[5] = 5;\n values[6] = 6;\n values[7] = 7;\n values[8] = 8;\n}\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid unique_by_key_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, Iterator3 result)\n{\n *result = thrust::unique_by_key(exec, keys_first, keys_last, values_first);\n}\n\n\ntemplate\n__global__\nvoid unique_by_key_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, BinaryPredicate pred, Iterator3 result)\n{\n *result = thrust::unique_by_key(exec, keys_first, keys_last, values_first, pred);\n}\n\n\ntemplate\nvoid TestUniqueByKeyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector keys;\n Vector values;\n \n typedef thrust::pair iter_pair;\n thrust::device_vector new_last_vec(1);\n iter_pair new_last;\n \n // basic test\n initialize_keys(keys); initialize_values(values);\n \n unique_by_key_kernel<<<1,1>>>(exec, keys.begin(), keys.end(), values.begin(), new_last_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n new_last = new_last_vec[0];\n \n ASSERT_EQUAL(new_last.first - keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - values.begin(), 5);\n ASSERT_EQUAL(keys[0], 11);\n ASSERT_EQUAL(keys[1], 21);\n ASSERT_EQUAL(keys[2], 20);\n ASSERT_EQUAL(keys[3], 21);\n ASSERT_EQUAL(keys[4], 37);\n \n ASSERT_EQUAL(values[0], 0);\n ASSERT_EQUAL(values[1], 2);\n ASSERT_EQUAL(values[2], 3);\n ASSERT_EQUAL(values[3], 4);\n ASSERT_EQUAL(values[4], 7);\n \n // test BinaryPredicate\n initialize_keys(keys); initialize_values(values);\n \n unique_by_key_kernel<<<1,1>>>(exec, keys.begin(), keys.end(), values.begin(), is_equal_div_10_unique(), new_last_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n new_last = new_last_vec[0];\n \n ASSERT_EQUAL(new_last.first - keys.begin(), 3);\n ASSERT_EQUAL(new_last.second - values.begin(), 3);\n ASSERT_EQUAL(keys[0], 11);\n ASSERT_EQUAL(keys[1], 21);\n ASSERT_EQUAL(keys[2], 37);\n \n ASSERT_EQUAL(values[0], 0);\n ASSERT_EQUAL(values[1], 2);\n ASSERT_EQUAL(values[2], 7);\n}\n\nvoid TestUniqueByKeyDeviceSeq()\n{\n TestUniqueByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUniqueByKeyDeviceSeq);\n\n\nvoid TestUniqueByKeyDeviceDevice()\n{\n TestUniqueByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUniqueByKeyDeviceDevice);\n\n\nvoid TestUniqueByKeyDeviceNoSync()\n{\n TestUniqueByKeyDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueByKeyDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestUniqueByKeyCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector keys;\n Vector values;\n \n typedef thrust::pair iter_pair;\n iter_pair new_last;\n \n // basic test\n initialize_keys(keys); initialize_values(values);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n \n new_last = thrust::unique_by_key(streampolicy, keys.begin(), keys.end(), values.begin());\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(new_last.first - keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - values.begin(), 5);\n ASSERT_EQUAL(keys[0], 11);\n ASSERT_EQUAL(keys[1], 21);\n ASSERT_EQUAL(keys[2], 20);\n ASSERT_EQUAL(keys[3], 21);\n ASSERT_EQUAL(keys[4], 37);\n \n ASSERT_EQUAL(values[0], 0);\n ASSERT_EQUAL(values[1], 2);\n ASSERT_EQUAL(values[2], 3);\n ASSERT_EQUAL(values[3], 4);\n ASSERT_EQUAL(values[4], 7);\n \n // test BinaryPredicate\n initialize_keys(keys); initialize_values(values);\n \n new_last = thrust::unique_by_key(streampolicy, keys.begin(), keys.end(), values.begin(), is_equal_div_10_unique());\n \n ASSERT_EQUAL(new_last.first - keys.begin(), 3);\n ASSERT_EQUAL(new_last.second - values.begin(), 3);\n ASSERT_EQUAL(keys[0], 11);\n ASSERT_EQUAL(keys[1], 21);\n ASSERT_EQUAL(keys[2], 37);\n \n ASSERT_EQUAL(values[0], 0);\n ASSERT_EQUAL(values[1], 2);\n ASSERT_EQUAL(values[2], 7);\n\n cudaStreamDestroy(s);\n}\n\nvoid TestUniqueByKeyCudaStreamsSync()\n{\n TestUniqueByKeyCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestUniqueByKeyCudaStreamsSync);\n\n\nvoid TestUniqueByKeyCudaStreamsNoSync()\n{\n TestUniqueByKeyCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueByKeyCudaStreamsNoSync);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid unique_by_key_copy_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, Iterator3 keys_result, Iterator4 values_result, Iterator5 result)\n{\n *result = thrust::unique_by_key_copy(exec, keys_first, keys_last, values_first, keys_result, values_result);\n}\n\n\ntemplate\n__global__\nvoid unique_by_key_copy_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, Iterator3 keys_result, Iterator4 values_result, BinaryPredicate pred, Iterator5 result)\n{\n *result = thrust::unique_by_key_copy(exec, keys_first, keys_last, values_first, keys_result, values_result, pred);\n}\n\n\ntemplate\nvoid TestUniqueCopyByKeyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector keys;\n Vector values;\n\n typedef thrust::pair iter_pair;\n thrust::device_vector new_last_vec(1);\n iter_pair new_last;\n\n // basic test\n initialize_keys(keys); initialize_values(values);\n\n Vector output_keys(keys.size());\n Vector output_values(values.size());\n\n unique_by_key_copy_kernel<<<1,1>>>(exec, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), new_last_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n new_last = new_last_vec[0];\n\n ASSERT_EQUAL(new_last.first - output_keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 5);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 20);\n ASSERT_EQUAL(output_keys[3], 21);\n ASSERT_EQUAL(output_keys[4], 37);\n \n ASSERT_EQUAL(output_values[0], 0);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 3);\n ASSERT_EQUAL(output_values[3], 4);\n ASSERT_EQUAL(output_values[4], 7);\n\n // test BinaryPredicate\n initialize_keys(keys); initialize_values(values);\n \n unique_by_key_copy_kernel<<<1,1>>>(exec, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), is_equal_div_10_unique(), new_last_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n new_last = new_last_vec[0];\n\n ASSERT_EQUAL(new_last.first - output_keys.begin(), 3);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 3);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 37);\n \n ASSERT_EQUAL(output_values[0], 0);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 7);\n}\n\n\nvoid TestUniqueCopyByKeyDeviceSeq()\n{\n TestUniqueCopyByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUniqueCopyByKeyDeviceSeq);\n\n\nvoid TestUniqueCopyByKeyDeviceDevice()\n{\n TestUniqueCopyByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUniqueCopyByKeyDeviceDevice);\n\n\nvoid TestUniqueCopyByKeyDeviceNoSync()\n{\n TestUniqueCopyByKeyDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueCopyByKeyDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestUniqueCopyByKeyCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector keys;\n Vector values;\n\n typedef thrust::pair iter_pair;\n iter_pair new_last;\n\n // basic test\n initialize_keys(keys); initialize_values(values);\n\n Vector output_keys(keys.size());\n Vector output_values(values.size());\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n\n new_last = thrust::unique_by_key_copy(streampolicy, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(new_last.first - output_keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 5);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 20);\n ASSERT_EQUAL(output_keys[3], 21);\n ASSERT_EQUAL(output_keys[4], 37);\n \n ASSERT_EQUAL(output_values[0], 0);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 3);\n ASSERT_EQUAL(output_values[3], 4);\n ASSERT_EQUAL(output_values[4], 7);\n\n // test BinaryPredicate\n initialize_keys(keys); initialize_values(values);\n \n new_last = thrust::unique_by_key_copy(streampolicy, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), is_equal_div_10_unique());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(new_last.first - output_keys.begin(), 3);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 3);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 37);\n \n ASSERT_EQUAL(output_values[0], 0);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 7);\n\n cudaStreamDestroy(s);\n}\n\nvoid TestUniqueCopyByKeyCudaStreamsSync()\n{\n TestUniqueCopyByKeyCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestUniqueCopyByKeyCudaStreamsSync);\n\n\nvoid TestUniqueCopyByKeyCudaStreamsNoSync()\n{\n TestUniqueCopyByKeyCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueCopyByKeyCudaStreamsNoSync);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n\ntemplate\nstruct is_equal_div_10_unique\n{\n __host__ __device__\n bool operator()(const T x, const T& y) const { return ((int) x / 10) == ((int) y / 10); }\n};\n\n\ntemplate\nvoid initialize_keys(Vector& keys)\n{\n keys.resize(9);\n keys[0] = 11;\n keys[1] = 11;\n keys[2] = 21;\n keys[3] = 20;\n keys[4] = 21;\n keys[5] = 21;\n keys[6] = 21;\n keys[7] = 37;\n keys[8] = 37;\n}\n\n\ntemplate\nvoid initialize_values(Vector& values)\n{\n values.resize(9);\n values[0] = 0; \n values[1] = 1;\n values[2] = 2;\n values[3] = 3;\n values[4] = 4;\n values[5] = 5;\n values[6] = 6;\n values[7] = 7;\n values[8] = 8;\n}\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid unique_by_key_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, Iterator3 result)\n{\n *result = thrust::unique_by_key(exec, keys_first, keys_last, values_first);\n}\n\n\ntemplate\n__global__\nvoid unique_by_key_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, BinaryPredicate pred, Iterator3 result)\n{\n *result = thrust::unique_by_key(exec, keys_first, keys_last, values_first, pred);\n}\n\n\ntemplate\nvoid TestUniqueByKeyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector keys;\n Vector values;\n \n typedef thrust::pair iter_pair;\n thrust::device_vector new_last_vec(1);\n iter_pair new_last;\n \n // basic test\n initialize_keys(keys); initialize_values(values);\n \n unique_by_key_kernel<<<1,1>>>(exec, keys.begin(), keys.end(), values.begin(), new_last_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n new_last = new_last_vec[0];\n \n ASSERT_EQUAL(new_last.first - keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - values.begin(), 5);\n ASSERT_EQUAL(keys[0], 11);\n ASSERT_EQUAL(keys[1], 21);\n ASSERT_EQUAL(keys[2], 20);\n ASSERT_EQUAL(keys[3], 21);\n ASSERT_EQUAL(keys[4], 37);\n \n ASSERT_EQUAL(values[0], 0);\n ASSERT_EQUAL(values[1], 2);\n ASSERT_EQUAL(values[2], 3);\n ASSERT_EQUAL(values[3], 4);\n ASSERT_EQUAL(values[4], 7);\n \n // test BinaryPredicate\n initialize_keys(keys); initialize_values(values);\n \n unique_by_key_kernel<<<1,1>>>(exec, keys.begin(), keys.end(), values.begin(), is_equal_div_10_unique(), new_last_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n new_last = new_last_vec[0];\n \n ASSERT_EQUAL(new_last.first - keys.begin(), 3);\n ASSERT_EQUAL(new_last.second - values.begin(), 3);\n ASSERT_EQUAL(keys[0], 11);\n ASSERT_EQUAL(keys[1], 21);\n ASSERT_EQUAL(keys[2], 37);\n \n ASSERT_EQUAL(values[0], 0);\n ASSERT_EQUAL(values[1], 2);\n ASSERT_EQUAL(values[2], 7);\n}\n\nvoid TestUniqueByKeyDeviceSeq()\n{\n TestUniqueByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUniqueByKeyDeviceSeq);\n\n\nvoid TestUniqueByKeyDeviceDevice()\n{\n TestUniqueByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUniqueByKeyDeviceDevice);\n\n\nvoid TestUniqueByKeyDeviceNoSync()\n{\n TestUniqueByKeyDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueByKeyDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestUniqueByKeyCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector keys;\n Vector values;\n \n typedef thrust::pair iter_pair;\n iter_pair new_last;\n \n // basic test\n initialize_keys(keys); initialize_values(values);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n \n new_last = thrust::unique_by_key(streampolicy, keys.begin(), keys.end(), values.begin());\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(new_last.first - keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - values.begin(), 5);\n ASSERT_EQUAL(keys[0], 11);\n ASSERT_EQUAL(keys[1], 21);\n ASSERT_EQUAL(keys[2], 20);\n ASSERT_EQUAL(keys[3], 21);\n ASSERT_EQUAL(keys[4], 37);\n \n ASSERT_EQUAL(values[0], 0);\n ASSERT_EQUAL(values[1], 2);\n ASSERT_EQUAL(values[2], 3);\n ASSERT_EQUAL(values[3], 4);\n ASSERT_EQUAL(values[4], 7);\n \n // test BinaryPredicate\n initialize_keys(keys); initialize_values(values);\n \n new_last = thrust::unique_by_key(streampolicy, keys.begin(), keys.end(), values.begin(), is_equal_div_10_unique());\n \n ASSERT_EQUAL(new_last.first - keys.begin(), 3);\n ASSERT_EQUAL(new_last.second - values.begin(), 3);\n ASSERT_EQUAL(keys[0], 11);\n ASSERT_EQUAL(keys[1], 21);\n ASSERT_EQUAL(keys[2], 37);\n \n ASSERT_EQUAL(values[0], 0);\n ASSERT_EQUAL(values[1], 2);\n ASSERT_EQUAL(values[2], 7);\n\n hipStreamDestroy(s);\n}\n\nvoid TestUniqueByKeyCudaStreamsSync()\n{\n TestUniqueByKeyCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestUniqueByKeyCudaStreamsSync);\n\n\nvoid TestUniqueByKeyCudaStreamsNoSync()\n{\n TestUniqueByKeyCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueByKeyCudaStreamsNoSync);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid unique_by_key_copy_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, Iterator3 keys_result, Iterator4 values_result, Iterator5 result)\n{\n *result = thrust::unique_by_key_copy(exec, keys_first, keys_last, values_first, keys_result, values_result);\n}\n\n\ntemplate\n__global__\nvoid unique_by_key_copy_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, Iterator3 keys_result, Iterator4 values_result, BinaryPredicate pred, Iterator5 result)\n{\n *result = thrust::unique_by_key_copy(exec, keys_first, keys_last, values_first, keys_result, values_result, pred);\n}\n\n\ntemplate\nvoid TestUniqueCopyByKeyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector keys;\n Vector values;\n\n typedef thrust::pair iter_pair;\n thrust::device_vector new_last_vec(1);\n iter_pair new_last;\n\n // basic test\n initialize_keys(keys); initialize_values(values);\n\n Vector output_keys(keys.size());\n Vector output_values(values.size());\n\n unique_by_key_copy_kernel<<<1,1>>>(exec, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), new_last_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n new_last = new_last_vec[0];\n\n ASSERT_EQUAL(new_last.first - output_keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 5);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 20);\n ASSERT_EQUAL(output_keys[3], 21);\n ASSERT_EQUAL(output_keys[4], 37);\n \n ASSERT_EQUAL(output_values[0], 0);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 3);\n ASSERT_EQUAL(output_values[3], 4);\n ASSERT_EQUAL(output_values[4], 7);\n\n // test BinaryPredicate\n initialize_keys(keys); initialize_values(values);\n \n unique_by_key_copy_kernel<<<1,1>>>(exec, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), is_equal_div_10_unique(), new_last_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n new_last = new_last_vec[0];\n\n ASSERT_EQUAL(new_last.first - output_keys.begin(), 3);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 3);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 37);\n \n ASSERT_EQUAL(output_values[0], 0);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 7);\n}\n\n\nvoid TestUniqueCopyByKeyDeviceSeq()\n{\n TestUniqueCopyByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUniqueCopyByKeyDeviceSeq);\n\n\nvoid TestUniqueCopyByKeyDeviceDevice()\n{\n TestUniqueCopyByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUniqueCopyByKeyDeviceDevice);\n\n\nvoid TestUniqueCopyByKeyDeviceNoSync()\n{\n TestUniqueCopyByKeyDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueCopyByKeyDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestUniqueCopyByKeyCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector keys;\n Vector values;\n\n typedef thrust::pair iter_pair;\n iter_pair new_last;\n\n // basic test\n initialize_keys(keys); initialize_values(values);\n\n Vector output_keys(keys.size());\n Vector output_values(values.size());\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n\n new_last = thrust::unique_by_key_copy(streampolicy, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(new_last.first - output_keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 5);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 20);\n ASSERT_EQUAL(output_keys[3], 21);\n ASSERT_EQUAL(output_keys[4], 37);\n \n ASSERT_EQUAL(output_values[0], 0);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 3);\n ASSERT_EQUAL(output_values[3], 4);\n ASSERT_EQUAL(output_values[4], 7);\n\n // test BinaryPredicate\n initialize_keys(keys); initialize_values(values);\n \n new_last = thrust::unique_by_key_copy(streampolicy, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), is_equal_div_10_unique());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(new_last.first - output_keys.begin(), 3);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 3);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 37);\n \n ASSERT_EQUAL(output_values[0], 0);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 7);\n\n hipStreamDestroy(s);\n}\n\nvoid TestUniqueCopyByKeyCudaStreamsSync()\n{\n TestUniqueCopyByKeyCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestUniqueCopyByKeyCudaStreamsSync);\n\n\nvoid TestUniqueCopyByKeyCudaStreamsNoSync()\n{\n TestUniqueCopyByKeyCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueCopyByKeyCudaStreamsNoSync);\n\n", "source": "/kaggle/working/thrust/testing/cuda/unique_by_key.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__ void mismatch_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2, Iterator3 result)\n{\n *result = thrust::mismatch(exec, first1, last1, first2);\n}\n\n\ntemplate\nvoid TestMismatchDevice(ExecutionPolicy exec)\n{\n thrust::device_vector a(4);\n thrust::device_vector b(4);\n a[0] = 1; b[0] = 1;\n a[1] = 2; b[1] = 2;\n a[2] = 3; b[2] = 4;\n a[3] = 4; b[3] = 3;\n\n typedef thrust::pair<\n typename thrust::device_vector::iterator,\n typename thrust::device_vector::iterator\n > pair_type;\n\n thrust::device_vector d_result(1);\n \n mismatch_kernel<<<1,1>>>(exec, a.begin(), a.end(), b.begin(), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(2, ((pair_type)d_result[0]).first - a.begin());\n ASSERT_EQUAL(2, ((pair_type)d_result[0]).second - b.begin());\n \n b[2] = 3;\n \n mismatch_kernel<<<1,1>>>(exec, a.begin(), a.end(), b.begin(), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(3, ((pair_type)d_result[0]).first - a.begin());\n ASSERT_EQUAL(3, ((pair_type)d_result[0]).second - b.begin());\n \n b[3] = 4;\n \n mismatch_kernel<<<1,1>>>(exec, a.begin(), a.end(), b.begin(), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(4, ((pair_type)d_result[0]).first - a.begin());\n ASSERT_EQUAL(4, ((pair_type)d_result[0]).second - b.begin());\n}\n\n\nvoid TestMismatchDeviceSeq()\n{\n TestMismatchDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestMismatchDeviceSeq);\n\n\nvoid TestMismatchDeviceDevice()\n{\n TestMismatchDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestMismatchDeviceDevice);\n#endif\n\n\nvoid TestMismatchCudaStreams()\n{\n typedef thrust::device_vector Vector;\n\n Vector a(4); Vector b(4);\n a[0] = 1; b[0] = 1;\n a[1] = 2; b[1] = 2;\n a[2] = 3; b[2] = 4;\n a[3] = 4; b[3] = 3;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n ASSERT_EQUAL(thrust::mismatch(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin()).first - a.begin(), 2);\n ASSERT_EQUAL(thrust::mismatch(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin()).second - b.begin(), 2);\n\n b[2] = 3;\n \n ASSERT_EQUAL(thrust::mismatch(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin()).first - a.begin(), 3);\n ASSERT_EQUAL(thrust::mismatch(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin()).second - b.begin(), 3);\n \n b[3] = 4;\n \n ASSERT_EQUAL(thrust::mismatch(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin()).first - a.begin(), 4);\n ASSERT_EQUAL(thrust::mismatch(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin()).second - b.begin(), 4);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestMismatchCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__ void mismatch_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2, Iterator3 result)\n{\n *result = thrust::mismatch(exec, first1, last1, first2);\n}\n\n\ntemplate\nvoid TestMismatchDevice(ExecutionPolicy exec)\n{\n thrust::device_vector a(4);\n thrust::device_vector b(4);\n a[0] = 1; b[0] = 1;\n a[1] = 2; b[1] = 2;\n a[2] = 3; b[2] = 4;\n a[3] = 4; b[3] = 3;\n\n typedef thrust::pair<\n typename thrust::device_vector::iterator,\n typename thrust::device_vector::iterator\n > pair_type;\n\n thrust::device_vector d_result(1);\n \n mismatch_kernel<<<1,1>>>(exec, a.begin(), a.end(), b.begin(), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(2, ((pair_type)d_result[0]).first - a.begin());\n ASSERT_EQUAL(2, ((pair_type)d_result[0]).second - b.begin());\n \n b[2] = 3;\n \n mismatch_kernel<<<1,1>>>(exec, a.begin(), a.end(), b.begin(), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(3, ((pair_type)d_result[0]).first - a.begin());\n ASSERT_EQUAL(3, ((pair_type)d_result[0]).second - b.begin());\n \n b[3] = 4;\n \n mismatch_kernel<<<1,1>>>(exec, a.begin(), a.end(), b.begin(), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(4, ((pair_type)d_result[0]).first - a.begin());\n ASSERT_EQUAL(4, ((pair_type)d_result[0]).second - b.begin());\n}\n\n\nvoid TestMismatchDeviceSeq()\n{\n TestMismatchDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestMismatchDeviceSeq);\n\n\nvoid TestMismatchDeviceDevice()\n{\n TestMismatchDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestMismatchDeviceDevice);\n#endif\n\n\nvoid TestMismatchCudaStreams()\n{\n typedef thrust::device_vector Vector;\n\n Vector a(4); Vector b(4);\n a[0] = 1; b[0] = 1;\n a[1] = 2; b[1] = 2;\n a[2] = 3; b[2] = 4;\n a[3] = 4; b[3] = 3;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n ASSERT_EQUAL(thrust::mismatch(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin()).first - a.begin(), 2);\n ASSERT_EQUAL(thrust::mismatch(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin()).second - b.begin(), 2);\n\n b[2] = 3;\n \n ASSERT_EQUAL(thrust::mismatch(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin()).first - a.begin(), 3);\n ASSERT_EQUAL(thrust::mismatch(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin()).second - b.begin(), 3);\n \n b[3] = 4;\n \n ASSERT_EQUAL(thrust::mismatch(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin()).first - a.begin(), 4);\n ASSERT_EQUAL(thrust::mismatch(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin()).second - b.begin(), 4);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestMismatchCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/mismatch.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid inner_product_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2, T init, Iterator3 result)\n{\n *result = thrust::inner_product(exec, first1, last1, first2, init);\n}\n\n\ntemplate\nvoid TestInnerProductDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n\n thrust::host_vector h_v1 = unittest::random_integers(n);\n thrust::host_vector h_v2 = unittest::random_integers(n);\n \n thrust::device_vector d_v1 = h_v1;\n thrust::device_vector d_v2 = h_v2;\n \n thrust::device_vector result(1);\n \n int init = 13;\n \n int expected = thrust::inner_product(h_v1.begin(), h_v1.end(), h_v2.begin(), init);\n\n inner_product_kernel<<<1,1>>>(exec, d_v1.begin(), d_v1.end(), d_v2.begin(), init, result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(expected, result[0]);\n}\n\n\nvoid TestInnerProductDeviceSeq()\n{\n TestInnerProductDevice(thrust::seq);\n};\nDECLARE_UNITTEST(TestInnerProductDeviceSeq);\n\n\nvoid TestInnerProductDeviceDevice()\n{\n TestInnerProductDevice(thrust::device);\n};\nDECLARE_UNITTEST(TestInnerProductDeviceDevice);\n#endif\n\n\nvoid TestInnerProductCudaStreams()\n{\n thrust::device_vector v1(3);\n thrust::device_vector v2(3);\n v1[0] = 1; v1[1] = -2; v1[2] = 3;\n v2[0] = -4; v2[1] = 5; v2[2] = 6;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n int init = 3;\n int result = thrust::inner_product(thrust::cuda::par.on(s), v1.begin(), v1.end(), v2.begin(), init);\n ASSERT_EQUAL(result, 7);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestInnerProductCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid inner_product_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2, T init, Iterator3 result)\n{\n *result = thrust::inner_product(exec, first1, last1, first2, init);\n}\n\n\ntemplate\nvoid TestInnerProductDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n\n thrust::host_vector h_v1 = unittest::random_integers(n);\n thrust::host_vector h_v2 = unittest::random_integers(n);\n \n thrust::device_vector d_v1 = h_v1;\n thrust::device_vector d_v2 = h_v2;\n \n thrust::device_vector result(1);\n \n int init = 13;\n \n int expected = thrust::inner_product(h_v1.begin(), h_v1.end(), h_v2.begin(), init);\n\n inner_product_kernel<<<1,1>>>(exec, d_v1.begin(), d_v1.end(), d_v2.begin(), init, result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(expected, result[0]);\n}\n\n\nvoid TestInnerProductDeviceSeq()\n{\n TestInnerProductDevice(thrust::seq);\n};\nDECLARE_UNITTEST(TestInnerProductDeviceSeq);\n\n\nvoid TestInnerProductDeviceDevice()\n{\n TestInnerProductDevice(thrust::device);\n};\nDECLARE_UNITTEST(TestInnerProductDeviceDevice);\n#endif\n\n\nvoid TestInnerProductCudaStreams()\n{\n thrust::device_vector v1(3);\n thrust::device_vector v2(3);\n v1[0] = 1; v1[1] = -2; v1[2] = 3;\n v2[0] = -4; v2[1] = 5; v2[2] = 6;\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n int init = 3;\n int result = thrust::inner_product(thrust::cuda::par.on(s), v1.begin(), v1.end(), v2.begin(), init);\n ASSERT_EQUAL(result, 7);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestInnerProductCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/inner_product.cu"} {"input": "#include \n#include \n#include \n\n\ntemplate\nstruct is_equal_div_10_unique\n{\n __host__ __device__\n bool operator()(const T x, const T& y) const { return ((int) x / 10) == ((int) y / 10); }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid unique_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n *result = thrust::unique(exec, first, last);\n}\n\n\ntemplate\n__global__\nvoid unique_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, BinaryPredicate pred, Iterator2 result)\n{\n *result = thrust::unique(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestUniqueDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(10);\n data[0] = 11; \n data[1] = 11; \n data[2] = 12;\n data[3] = 20; \n data[4] = 29; \n data[5] = 21; \n data[6] = 21; \n data[7] = 31; \n data[8] = 31; \n data[9] = 37; \n\n thrust::device_vector new_last_vec(1);\n Vector::iterator new_last;\n \n unique_kernel<<<1,1>>>(exec, data.begin(), data.end(), new_last_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n new_last = new_last_vec[0];\n\n ASSERT_EQUAL(new_last - data.begin(), 7);\n ASSERT_EQUAL(data[0], 11);\n ASSERT_EQUAL(data[1], 12);\n ASSERT_EQUAL(data[2], 20);\n ASSERT_EQUAL(data[3], 29);\n ASSERT_EQUAL(data[4], 21);\n ASSERT_EQUAL(data[5], 31);\n ASSERT_EQUAL(data[6], 37);\n\n unique_kernel<<<1,1>>>(exec, data.begin(), new_last, is_equal_div_10_unique(), new_last_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n new_last = new_last_vec[0];\n\n ASSERT_EQUAL(new_last - data.begin(), 3);\n ASSERT_EQUAL(data[0], 11);\n ASSERT_EQUAL(data[1], 20);\n ASSERT_EQUAL(data[2], 31);\n}\n\n\nvoid TestUniqueDeviceSeq()\n{\n TestUniqueDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUniqueDeviceSeq);\n\n\nvoid TestUniqueDeviceDevice()\n{\n TestUniqueDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUniqueDeviceDevice);\n\n\nvoid TestUniqueDeviceNoSync()\n{\n TestUniqueDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestUniqueCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(10);\n data[0] = 11; \n data[1] = 11; \n data[2] = 12;\n data[3] = 20; \n data[4] = 29; \n data[5] = 21; \n data[6] = 21; \n data[7] = 31; \n data[8] = 31; \n data[9] = 37; \n\n thrust::device_vector new_last_vec(1);\n Vector::iterator new_last;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n \n new_last = thrust::unique(streampolicy, data.begin(), data.end());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(new_last - data.begin(), 7);\n ASSERT_EQUAL(data[0], 11);\n ASSERT_EQUAL(data[1], 12);\n ASSERT_EQUAL(data[2], 20);\n ASSERT_EQUAL(data[3], 29);\n ASSERT_EQUAL(data[4], 21);\n ASSERT_EQUAL(data[5], 31);\n ASSERT_EQUAL(data[6], 37);\n\n new_last = thrust::unique(streampolicy, data.begin(), new_last, is_equal_div_10_unique());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(new_last - data.begin(), 3);\n ASSERT_EQUAL(data[0], 11);\n ASSERT_EQUAL(data[1], 20);\n ASSERT_EQUAL(data[2], 31);\n\n cudaStreamDestroy(s);\n}\n\nvoid TestUniqueCudaStreamsSync()\n{\n TestUniqueCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestUniqueCudaStreamsSync);\n\n\nvoid TestUniqueCudaStreamsNoSync()\n{\n TestUniqueCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueCudaStreamsNoSync);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid unique_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, Iterator3 result2)\n{\n *result2 = thrust::unique_copy(exec, first, last, result1);\n}\n\n\ntemplate\n__global__\nvoid unique_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, BinaryPredicate pred, Iterator3 result2)\n{\n *result2 = thrust::unique_copy(exec, first, last, result1, pred);\n}\n\n\ntemplate\nvoid TestUniqueCopyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(10);\n data[0] = 11; \n data[1] = 11; \n data[2] = 12;\n data[3] = 20; \n data[4] = 29; \n data[5] = 21; \n data[6] = 21; \n data[7] = 31; \n data[8] = 31; \n data[9] = 37; \n \n Vector output(10, -1);\n\n thrust::device_vector new_last_vec(1);\n Vector::iterator new_last;\n \n unique_copy_kernel<<<1,1>>>(exec, data.begin(), data.end(), output.begin(), new_last_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n new_last = new_last_vec[0];\n\n ASSERT_EQUAL(new_last - output.begin(), 7);\n ASSERT_EQUAL(output[0], 11);\n ASSERT_EQUAL(output[1], 12);\n ASSERT_EQUAL(output[2], 20);\n ASSERT_EQUAL(output[3], 29);\n ASSERT_EQUAL(output[4], 21);\n ASSERT_EQUAL(output[5], 31);\n ASSERT_EQUAL(output[6], 37);\n\n unique_copy_kernel<<<1,1>>>(exec, output.begin(), new_last, data.begin(), is_equal_div_10_unique(), new_last_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n new_last = new_last_vec[0];\n\n ASSERT_EQUAL(new_last - data.begin(), 3);\n ASSERT_EQUAL(data[0], 11);\n ASSERT_EQUAL(data[1], 20);\n ASSERT_EQUAL(data[2], 31);\n}\n\n\nvoid TestUniqueCopyDeviceSeq()\n{\n TestUniqueCopyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUniqueCopyDeviceSeq);\n\n\nvoid TestUniqueCopyDeviceDevice()\n{\n TestUniqueCopyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUniqueCopyDeviceDevice);\n\n\nvoid TestUniqueCopyDeviceNoSync()\n{\n TestUniqueCopyDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueCopyDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestUniqueCopyCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(10);\n data[0] = 11; \n data[1] = 11; \n data[2] = 12;\n data[3] = 20; \n data[4] = 29; \n data[5] = 21; \n data[6] = 21; \n data[7] = 31; \n data[8] = 31; \n data[9] = 37; \n \n Vector output(10, -1);\n\n thrust::device_vector new_last_vec(1);\n Vector::iterator new_last;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n \n new_last = thrust::unique_copy(streampolicy, data.begin(), data.end(), output.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(new_last - output.begin(), 7);\n ASSERT_EQUAL(output[0], 11);\n ASSERT_EQUAL(output[1], 12);\n ASSERT_EQUAL(output[2], 20);\n ASSERT_EQUAL(output[3], 29);\n ASSERT_EQUAL(output[4], 21);\n ASSERT_EQUAL(output[5], 31);\n ASSERT_EQUAL(output[6], 37);\n\n new_last = thrust::unique_copy(streampolicy, output.begin(), new_last, data.begin(), is_equal_div_10_unique());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(new_last - data.begin(), 3);\n ASSERT_EQUAL(data[0], 11);\n ASSERT_EQUAL(data[1], 20);\n ASSERT_EQUAL(data[2], 31);\n\n cudaStreamDestroy(s);\n}\n\nvoid TestUniqueCopyCudaStreamsSync()\n{\n TestUniqueCopyCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestUniqueCopyCudaStreamsSync);\n\n\nvoid TestUniqueCopyCudaStreamsNoSync()\n{\n TestUniqueCopyCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueCopyCudaStreamsNoSync);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid unique_count_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n *result = thrust::unique_count(exec, first, last);\n}\n\n\ntemplate\n__global__\nvoid unique_count_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, BinaryPredicate pred, Iterator2 result)\n{\n *result = thrust::unique_count(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestUniqueCountDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(10);\n data[0] = 11; \n data[1] = 11; \n data[2] = 12;\n data[3] = 20; \n data[4] = 29; \n data[5] = 21; \n data[6] = 21; \n data[7] = 31; \n data[8] = 31; \n data[9] = 37; \n \n Vector output(1, -1);\n \n unique_count_kernel<<<1,1>>>(exec, data.begin(), data.end(), output.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(output[0], 7);\n\n unique_count_kernel<<<1,1>>>(exec, data.begin(), data.end(), is_equal_div_10_unique(), output.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(output[0], 3);\n}\n\n\nvoid TestUniqueCountDeviceSeq()\n{\n TestUniqueCountDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUniqueCountDeviceSeq);\n\n\nvoid TestUniqueCountDeviceDevice()\n{\n TestUniqueCountDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUniqueCountDeviceDevice);\n\n\nvoid TestUniqueCountDeviceNoSync()\n{\n TestUniqueCountDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueCountDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestUniqueCountCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(10);\n data[0] = 11; \n data[1] = 11; \n data[2] = 12;\n data[3] = 20; \n data[4] = 29; \n data[5] = 21; \n data[6] = 21; \n data[7] = 31; \n data[8] = 31; \n data[9] = 37;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n \n int result = thrust::unique_count(streampolicy, data.begin(), data.end());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(result, 7);\n\n result = thrust::unique_count(streampolicy, data.begin(), data.end(), is_equal_div_10_unique());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(result, 3);\n\n cudaStreamDestroy(s);\n}\n\nvoid TestUniqueCountCudaStreamsSync()\n{\n TestUniqueCountCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestUniqueCountCudaStreamsSync);\n\n\nvoid TestUniqueCountCudaStreamsNoSync()\n{\n TestUniqueCountCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueCountCudaStreamsNoSync);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\ntemplate\nstruct is_equal_div_10_unique\n{\n __host__ __device__\n bool operator()(const T x, const T& y) const { return ((int) x / 10) == ((int) y / 10); }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid unique_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n *result = thrust::unique(exec, first, last);\n}\n\n\ntemplate\n__global__\nvoid unique_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, BinaryPredicate pred, Iterator2 result)\n{\n *result = thrust::unique(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestUniqueDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(10);\n data[0] = 11; \n data[1] = 11; \n data[2] = 12;\n data[3] = 20; \n data[4] = 29; \n data[5] = 21; \n data[6] = 21; \n data[7] = 31; \n data[8] = 31; \n data[9] = 37; \n\n thrust::device_vector new_last_vec(1);\n Vector::iterator new_last;\n \n unique_kernel<<<1,1>>>(exec, data.begin(), data.end(), new_last_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n new_last = new_last_vec[0];\n\n ASSERT_EQUAL(new_last - data.begin(), 7);\n ASSERT_EQUAL(data[0], 11);\n ASSERT_EQUAL(data[1], 12);\n ASSERT_EQUAL(data[2], 20);\n ASSERT_EQUAL(data[3], 29);\n ASSERT_EQUAL(data[4], 21);\n ASSERT_EQUAL(data[5], 31);\n ASSERT_EQUAL(data[6], 37);\n\n unique_kernel<<<1,1>>>(exec, data.begin(), new_last, is_equal_div_10_unique(), new_last_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n new_last = new_last_vec[0];\n\n ASSERT_EQUAL(new_last - data.begin(), 3);\n ASSERT_EQUAL(data[0], 11);\n ASSERT_EQUAL(data[1], 20);\n ASSERT_EQUAL(data[2], 31);\n}\n\n\nvoid TestUniqueDeviceSeq()\n{\n TestUniqueDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUniqueDeviceSeq);\n\n\nvoid TestUniqueDeviceDevice()\n{\n TestUniqueDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUniqueDeviceDevice);\n\n\nvoid TestUniqueDeviceNoSync()\n{\n TestUniqueDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestUniqueCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(10);\n data[0] = 11; \n data[1] = 11; \n data[2] = 12;\n data[3] = 20; \n data[4] = 29; \n data[5] = 21; \n data[6] = 21; \n data[7] = 31; \n data[8] = 31; \n data[9] = 37; \n\n thrust::device_vector new_last_vec(1);\n Vector::iterator new_last;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n \n new_last = thrust::unique(streampolicy, data.begin(), data.end());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(new_last - data.begin(), 7);\n ASSERT_EQUAL(data[0], 11);\n ASSERT_EQUAL(data[1], 12);\n ASSERT_EQUAL(data[2], 20);\n ASSERT_EQUAL(data[3], 29);\n ASSERT_EQUAL(data[4], 21);\n ASSERT_EQUAL(data[5], 31);\n ASSERT_EQUAL(data[6], 37);\n\n new_last = thrust::unique(streampolicy, data.begin(), new_last, is_equal_div_10_unique());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(new_last - data.begin(), 3);\n ASSERT_EQUAL(data[0], 11);\n ASSERT_EQUAL(data[1], 20);\n ASSERT_EQUAL(data[2], 31);\n\n hipStreamDestroy(s);\n}\n\nvoid TestUniqueCudaStreamsSync()\n{\n TestUniqueCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestUniqueCudaStreamsSync);\n\n\nvoid TestUniqueCudaStreamsNoSync()\n{\n TestUniqueCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueCudaStreamsNoSync);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid unique_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, Iterator3 result2)\n{\n *result2 = thrust::unique_copy(exec, first, last, result1);\n}\n\n\ntemplate\n__global__\nvoid unique_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, BinaryPredicate pred, Iterator3 result2)\n{\n *result2 = thrust::unique_copy(exec, first, last, result1, pred);\n}\n\n\ntemplate\nvoid TestUniqueCopyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(10);\n data[0] = 11; \n data[1] = 11; \n data[2] = 12;\n data[3] = 20; \n data[4] = 29; \n data[5] = 21; \n data[6] = 21; \n data[7] = 31; \n data[8] = 31; \n data[9] = 37; \n \n Vector output(10, -1);\n\n thrust::device_vector new_last_vec(1);\n Vector::iterator new_last;\n \n unique_copy_kernel<<<1,1>>>(exec, data.begin(), data.end(), output.begin(), new_last_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n new_last = new_last_vec[0];\n\n ASSERT_EQUAL(new_last - output.begin(), 7);\n ASSERT_EQUAL(output[0], 11);\n ASSERT_EQUAL(output[1], 12);\n ASSERT_EQUAL(output[2], 20);\n ASSERT_EQUAL(output[3], 29);\n ASSERT_EQUAL(output[4], 21);\n ASSERT_EQUAL(output[5], 31);\n ASSERT_EQUAL(output[6], 37);\n\n unique_copy_kernel<<<1,1>>>(exec, output.begin(), new_last, data.begin(), is_equal_div_10_unique(), new_last_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n new_last = new_last_vec[0];\n\n ASSERT_EQUAL(new_last - data.begin(), 3);\n ASSERT_EQUAL(data[0], 11);\n ASSERT_EQUAL(data[1], 20);\n ASSERT_EQUAL(data[2], 31);\n}\n\n\nvoid TestUniqueCopyDeviceSeq()\n{\n TestUniqueCopyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUniqueCopyDeviceSeq);\n\n\nvoid TestUniqueCopyDeviceDevice()\n{\n TestUniqueCopyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUniqueCopyDeviceDevice);\n\n\nvoid TestUniqueCopyDeviceNoSync()\n{\n TestUniqueCopyDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueCopyDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestUniqueCopyCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(10);\n data[0] = 11; \n data[1] = 11; \n data[2] = 12;\n data[3] = 20; \n data[4] = 29; \n data[5] = 21; \n data[6] = 21; \n data[7] = 31; \n data[8] = 31; \n data[9] = 37; \n \n Vector output(10, -1);\n\n thrust::device_vector new_last_vec(1);\n Vector::iterator new_last;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n \n new_last = thrust::unique_copy(streampolicy, data.begin(), data.end(), output.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(new_last - output.begin(), 7);\n ASSERT_EQUAL(output[0], 11);\n ASSERT_EQUAL(output[1], 12);\n ASSERT_EQUAL(output[2], 20);\n ASSERT_EQUAL(output[3], 29);\n ASSERT_EQUAL(output[4], 21);\n ASSERT_EQUAL(output[5], 31);\n ASSERT_EQUAL(output[6], 37);\n\n new_last = thrust::unique_copy(streampolicy, output.begin(), new_last, data.begin(), is_equal_div_10_unique());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(new_last - data.begin(), 3);\n ASSERT_EQUAL(data[0], 11);\n ASSERT_EQUAL(data[1], 20);\n ASSERT_EQUAL(data[2], 31);\n\n hipStreamDestroy(s);\n}\n\nvoid TestUniqueCopyCudaStreamsSync()\n{\n TestUniqueCopyCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestUniqueCopyCudaStreamsSync);\n\n\nvoid TestUniqueCopyCudaStreamsNoSync()\n{\n TestUniqueCopyCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueCopyCudaStreamsNoSync);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid unique_count_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n *result = thrust::unique_count(exec, first, last);\n}\n\n\ntemplate\n__global__\nvoid unique_count_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, BinaryPredicate pred, Iterator2 result)\n{\n *result = thrust::unique_count(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestUniqueCountDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(10);\n data[0] = 11; \n data[1] = 11; \n data[2] = 12;\n data[3] = 20; \n data[4] = 29; \n data[5] = 21; \n data[6] = 21; \n data[7] = 31; \n data[8] = 31; \n data[9] = 37; \n \n Vector output(1, -1);\n \n unique_count_kernel<<<1,1>>>(exec, data.begin(), data.end(), output.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(output[0], 7);\n\n unique_count_kernel<<<1,1>>>(exec, data.begin(), data.end(), is_equal_div_10_unique(), output.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(output[0], 3);\n}\n\n\nvoid TestUniqueCountDeviceSeq()\n{\n TestUniqueCountDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUniqueCountDeviceSeq);\n\n\nvoid TestUniqueCountDeviceDevice()\n{\n TestUniqueCountDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUniqueCountDeviceDevice);\n\n\nvoid TestUniqueCountDeviceNoSync()\n{\n TestUniqueCountDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueCountDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestUniqueCountCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(10);\n data[0] = 11; \n data[1] = 11; \n data[2] = 12;\n data[3] = 20; \n data[4] = 29; \n data[5] = 21; \n data[6] = 21; \n data[7] = 31; \n data[8] = 31; \n data[9] = 37;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n \n int result = thrust::unique_count(streampolicy, data.begin(), data.end());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(result, 7);\n\n result = thrust::unique_count(streampolicy, data.begin(), data.end(), is_equal_div_10_unique());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(result, 3);\n\n hipStreamDestroy(s);\n}\n\nvoid TestUniqueCountCudaStreamsSync()\n{\n TestUniqueCountCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestUniqueCountCudaStreamsSync);\n\n\nvoid TestUniqueCountCudaStreamsNoSync()\n{\n TestUniqueCountCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestUniqueCountCudaStreamsNoSync);\n\n", "source": "/kaggle/working/thrust/testing/cuda/unique.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_symmetric_difference_kernel(ExecutionPolicy exec,\n Iterator1 first1, Iterator1 last1,\n Iterator2 first2, Iterator2 last2,\n Iterator3 result1,\n Iterator4 result2)\n{\n *result2 = thrust::set_symmetric_difference(exec, first1, last1, first2, last2, result1);\n}\n\n\ntemplate\nvoid TestSetSymmetricDifferenceDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::iterator Iterator;\n\n Vector a(4), b(5);\n\n a[0] = 0; a[1] = 2; a[2] = 4; a[3] = 6;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4; b[4] = 7;\n\n Vector ref(5);\n ref[0] = 2; ref[1] = 3; ref[2] = 3; ref[3] = 6; ref[4] = 7;\n\n Vector result(5);\n thrust::device_vector end_vec(1);\n\n set_symmetric_difference_kernel<<<1,1>>>(exec,\n a.begin(), a.end(),\n b.begin(), b.end(),\n result.begin(),\n end_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n Iterator end = end_vec[0];\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n}\n\n\nvoid TestSetSymmetricDifferenceDeviceSeq()\n{\n TestSetSymmetricDifferenceDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetSymmetricDifferenceDeviceSeq);\n\n\nvoid TestSetSymmetricDifferenceDeviceDevice()\n{\n TestSetSymmetricDifferenceDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetSymmetricDifferenceDeviceDevice);\n#endif\n\n\nvoid TestSetSymmetricDifferenceCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a(4), b(5);\n\n a[0] = 0; a[1] = 2; a[2] = 4; a[3] = 6;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4; b[4] = 7;\n\n Vector ref(5);\n ref[0] = 2; ref[1] = 3; ref[2] = 3; ref[3] = 6; ref[4] = 7;\n\n Vector result(5);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n Iterator end = thrust::set_symmetric_difference(thrust::cuda::par.on(s),\n a.begin(), a.end(),\n b.begin(), b.end(),\n result.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSetSymmetricDifferenceCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_symmetric_difference_kernel(ExecutionPolicy exec,\n Iterator1 first1, Iterator1 last1,\n Iterator2 first2, Iterator2 last2,\n Iterator3 result1,\n Iterator4 result2)\n{\n *result2 = thrust::set_symmetric_difference(exec, first1, last1, first2, last2, result1);\n}\n\n\ntemplate\nvoid TestSetSymmetricDifferenceDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::iterator Iterator;\n\n Vector a(4), b(5);\n\n a[0] = 0; a[1] = 2; a[2] = 4; a[3] = 6;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4; b[4] = 7;\n\n Vector ref(5);\n ref[0] = 2; ref[1] = 3; ref[2] = 3; ref[3] = 6; ref[4] = 7;\n\n Vector result(5);\n thrust::device_vector end_vec(1);\n\n set_symmetric_difference_kernel<<<1,1>>>(exec,\n a.begin(), a.end(),\n b.begin(), b.end(),\n result.begin(),\n end_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n Iterator end = end_vec[0];\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n}\n\n\nvoid TestSetSymmetricDifferenceDeviceSeq()\n{\n TestSetSymmetricDifferenceDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetSymmetricDifferenceDeviceSeq);\n\n\nvoid TestSetSymmetricDifferenceDeviceDevice()\n{\n TestSetSymmetricDifferenceDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetSymmetricDifferenceDeviceDevice);\n#endif\n\n\nvoid TestSetSymmetricDifferenceCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a(4), b(5);\n\n a[0] = 0; a[1] = 2; a[2] = 4; a[3] = 6;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4; b[4] = 7;\n\n Vector ref(5);\n ref[0] = 2; ref[1] = 3; ref[2] = 3; ref[3] = 6; ref[4] = 7;\n\n Vector result(5);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n Iterator end = thrust::set_symmetric_difference(thrust::cuda::par.on(s),\n a.begin(), a.end(),\n b.begin(), b.end(),\n result.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSetSymmetricDifferenceCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/set_symmetric_difference.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n thrust::copy(exec, first, last, result);\n}\n\n\ntemplate\nvoid TestCopyDevice(ExecutionPolicy exec, size_t n)\n{\n thrust::host_vector h_src = unittest::random_integers(n);\n thrust::host_vector h_dst(n);\n\n thrust::device_vector d_src = h_src;\n thrust::device_vector d_dst(n);\n \n thrust::copy(h_src.begin(), h_src.end(), h_dst.begin());\n copy_kernel<<<1,1>>>(exec, d_src.begin(), d_src.end(), d_dst.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_dst, d_dst);\n}\n\n\ntemplate\nvoid TestCopyDeviceSeq(size_t n)\n{\n TestCopyDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCopyDeviceSeq);\n\n\ntemplate\nvoid TestCopyDeviceDevice(size_t n)\n{\n TestCopyDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCopyDeviceDevice);\n\n\ntemplate\n__global__\nvoid copy_n_kernel(ExecutionPolicy exec, Iterator1 first, Size n, Iterator2 result)\n{\n thrust::copy_n(exec, first, n, result);\n}\n\n\ntemplate\nvoid TestCopyNDevice(ExecutionPolicy exec, size_t n)\n{\n thrust::host_vector h_src = unittest::random_integers(n);\n thrust::host_vector h_dst(n);\n\n thrust::device_vector d_src = h_src;\n thrust::device_vector d_dst(n);\n \n thrust::copy_n(h_src.begin(), h_src.size(), h_dst.begin());\n copy_n_kernel<<<1,1>>>(exec, d_src.begin(), d_src.size(), d_dst.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_dst, d_dst);\n}\n\n\ntemplate\nvoid TestCopyNDeviceSeq(size_t n)\n{\n TestCopyNDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCopyNDeviceSeq);\n\n\ntemplate\nvoid TestCopyNDeviceDevice(size_t n)\n{\n TestCopyNDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCopyNDeviceDevice);\n#endif\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n thrust::copy(exec, first, last, result);\n}\n\n\ntemplate\nvoid TestCopyDevice(ExecutionPolicy exec, size_t n)\n{\n thrust::host_vector h_src = unittest::random_integers(n);\n thrust::host_vector h_dst(n);\n\n thrust::device_vector d_src = h_src;\n thrust::device_vector d_dst(n);\n \n thrust::copy(h_src.begin(), h_src.end(), h_dst.begin());\n copy_kernel<<<1,1>>>(exec, d_src.begin(), d_src.end(), d_dst.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_dst, d_dst);\n}\n\n\ntemplate\nvoid TestCopyDeviceSeq(size_t n)\n{\n TestCopyDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCopyDeviceSeq);\n\n\ntemplate\nvoid TestCopyDeviceDevice(size_t n)\n{\n TestCopyDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCopyDeviceDevice);\n\n\ntemplate\n__global__\nvoid copy_n_kernel(ExecutionPolicy exec, Iterator1 first, Size n, Iterator2 result)\n{\n thrust::copy_n(exec, first, n, result);\n}\n\n\ntemplate\nvoid TestCopyNDevice(ExecutionPolicy exec, size_t n)\n{\n thrust::host_vector h_src = unittest::random_integers(n);\n thrust::host_vector h_dst(n);\n\n thrust::device_vector d_src = h_src;\n thrust::device_vector d_dst(n);\n \n thrust::copy_n(h_src.begin(), h_src.size(), h_dst.begin());\n copy_n_kernel<<<1,1>>>(exec, d_src.begin(), d_src.size(), d_dst.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_dst, d_dst);\n}\n\n\ntemplate\nvoid TestCopyNDeviceSeq(size_t n)\n{\n TestCopyNDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCopyNDeviceSeq);\n\n\ntemplate\nvoid TestCopyNDeviceDevice(size_t n)\n{\n TestCopyNDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCopyNDeviceDevice);\n#endif\n\n", "source": "/kaggle/working/thrust/testing/cuda/copy.cu"} {"input": "#include \n#include \n#include \n#include \n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid scatter_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 map_first, Iterator3 result)\n{\n thrust::scatter(exec, first, last, map_first, result);\n}\n\n\ntemplate\nvoid TestScatterDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n const size_t output_size = std::min((size_t) 10, 2 * n);\n \n thrust::host_vector h_input(n, 1);\n thrust::device_vector d_input(n, 1);\n \n thrust::host_vector h_map = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n {\n h_map[i] = h_map[i] % output_size;\n }\n \n thrust::device_vector d_map = h_map;\n \n thrust::host_vector h_output(output_size, 0);\n thrust::device_vector d_output(output_size, 0);\n \n thrust::scatter(h_input.begin(), h_input.end(), h_map.begin(), h_output.begin());\n\n scatter_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_map.begin(), d_output.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_EQUAL(h_output, d_output);\n}\n\nvoid TestScatterDeviceSeq()\n{\n TestScatterDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestScatterDeviceSeq);\n\nvoid TestScatterDeviceDevice()\n{\n TestScatterDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestScatterDeviceDevice);\n\n\ntemplate\n__global__\nvoid scatter_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 map_first, Iterator3 stencil_first, Iterator4 result, Function f)\n{\n thrust::scatter_if(exec, first, last, map_first, stencil_first, result, f);\n}\n\n\ntemplate\nstruct is_even_scatter_if\n{\n __host__ __device__ bool operator()(const T i) const { return (i % 2) == 0; }\n};\n\n\ntemplate\nvoid TestScatterIfDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n const size_t output_size = std::min((size_t) 10, 2 * n);\n \n thrust::host_vector h_input(n, 1);\n thrust::device_vector d_input(n, 1);\n \n thrust::host_vector h_map = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n {\n h_map[i] = h_map[i] % output_size;\n }\n \n thrust::device_vector d_map = h_map;\n \n thrust::host_vector h_output(output_size, 0);\n thrust::device_vector d_output(output_size, 0);\n \n thrust::scatter_if(h_input.begin(), h_input.end(), h_map.begin(), h_map.begin(), h_output.begin(), is_even_scatter_if());\n\n scatter_if_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_map.begin(), d_map.begin(), d_output.begin(), is_even_scatter_if());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_EQUAL(h_output, d_output);\n}\n\n\nvoid TestScatterIfDeviceSeq()\n{\n TestScatterIfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestScatterIfDeviceSeq);\n\n\nvoid TestScatterIfDeviceDevice()\n{\n TestScatterIfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestScatterIfDeviceDevice);\n#endif\n\n\nvoid TestScatterCudaStreams()\n{\n typedef thrust::device_vector Vector;\n\n Vector map(5); // scatter indices\n Vector src(5); // source vector\n Vector dst(8); // destination vector\n\n map[0] = 6; map[1] = 3; map[2] = 1; map[3] = 7; map[4] = 2;\n src[0] = 0; src[1] = 1; src[2] = 2; src[3] = 3; src[4] = 4;\n dst[0] = 0; dst[1] = 0; dst[2] = 0; dst[3] = 0; dst[4] = 0; dst[5] = 0; dst[6] = 0; dst[7] = 0;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::scatter(thrust::cuda::par.on(s), src.begin(), src.end(), map.begin(), dst.begin());\n\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(dst[0], 0);\n ASSERT_EQUAL(dst[1], 2);\n ASSERT_EQUAL(dst[2], 4);\n ASSERT_EQUAL(dst[3], 1);\n ASSERT_EQUAL(dst[4], 0);\n ASSERT_EQUAL(dst[5], 0);\n ASSERT_EQUAL(dst[6], 0);\n ASSERT_EQUAL(dst[7], 3);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestScatterCudaStreams);\n\n\nvoid TestScatterIfCudaStreams()\n{\n typedef thrust::device_vector Vector;\n \n Vector flg(5); // predicate array\n Vector map(5); // scatter indices\n Vector src(5); // source vector\n Vector dst(8); // destination vector\n \n flg[0] = 0; flg[1] = 1; flg[2] = 0; flg[3] = 1; flg[4] = 0;\n map[0] = 6; map[1] = 3; map[2] = 1; map[3] = 7; map[4] = 2;\n src[0] = 0; src[1] = 1; src[2] = 2; src[3] = 3; src[4] = 4;\n dst[0] = 0; dst[1] = 0; dst[2] = 0; dst[3] = 0; dst[4] = 0; dst[5] = 0; dst[6] = 0; dst[7] = 0;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n thrust::scatter_if(thrust::cuda::par.on(s), src.begin(), src.end(), map.begin(), flg.begin(), dst.begin());\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(dst[0], 0);\n ASSERT_EQUAL(dst[1], 0);\n ASSERT_EQUAL(dst[2], 0);\n ASSERT_EQUAL(dst[3], 1);\n ASSERT_EQUAL(dst[4], 0);\n ASSERT_EQUAL(dst[5], 0);\n ASSERT_EQUAL(dst[6], 0);\n ASSERT_EQUAL(dst[7], 3);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestScatterIfCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid scatter_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 map_first, Iterator3 result)\n{\n thrust::scatter(exec, first, last, map_first, result);\n}\n\n\ntemplate\nvoid TestScatterDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n const size_t output_size = std::min((size_t) 10, 2 * n);\n \n thrust::host_vector h_input(n, 1);\n thrust::device_vector d_input(n, 1);\n \n thrust::host_vector h_map = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n {\n h_map[i] = h_map[i] % output_size;\n }\n \n thrust::device_vector d_map = h_map;\n \n thrust::host_vector h_output(output_size, 0);\n thrust::device_vector d_output(output_size, 0);\n \n thrust::scatter(h_input.begin(), h_input.end(), h_map.begin(), h_output.begin());\n\n scatter_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_map.begin(), d_output.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_EQUAL(h_output, d_output);\n}\n\nvoid TestScatterDeviceSeq()\n{\n TestScatterDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestScatterDeviceSeq);\n\nvoid TestScatterDeviceDevice()\n{\n TestScatterDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestScatterDeviceDevice);\n\n\ntemplate\n__global__\nvoid scatter_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 map_first, Iterator3 stencil_first, Iterator4 result, Function f)\n{\n thrust::scatter_if(exec, first, last, map_first, stencil_first, result, f);\n}\n\n\ntemplate\nstruct is_even_scatter_if\n{\n __host__ __device__ bool operator()(const T i) const { return (i % 2) == 0; }\n};\n\n\ntemplate\nvoid TestScatterIfDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n const size_t output_size = std::min((size_t) 10, 2 * n);\n \n thrust::host_vector h_input(n, 1);\n thrust::device_vector d_input(n, 1);\n \n thrust::host_vector h_map = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n {\n h_map[i] = h_map[i] % output_size;\n }\n \n thrust::device_vector d_map = h_map;\n \n thrust::host_vector h_output(output_size, 0);\n thrust::device_vector d_output(output_size, 0);\n \n thrust::scatter_if(h_input.begin(), h_input.end(), h_map.begin(), h_map.begin(), h_output.begin(), is_even_scatter_if());\n\n scatter_if_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_map.begin(), d_map.begin(), d_output.begin(), is_even_scatter_if());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_EQUAL(h_output, d_output);\n}\n\n\nvoid TestScatterIfDeviceSeq()\n{\n TestScatterIfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestScatterIfDeviceSeq);\n\n\nvoid TestScatterIfDeviceDevice()\n{\n TestScatterIfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestScatterIfDeviceDevice);\n#endif\n\n\nvoid TestScatterCudaStreams()\n{\n typedef thrust::device_vector Vector;\n\n Vector map(5); // scatter indices\n Vector src(5); // source vector\n Vector dst(8); // destination vector\n\n map[0] = 6; map[1] = 3; map[2] = 1; map[3] = 7; map[4] = 2;\n src[0] = 0; src[1] = 1; src[2] = 2; src[3] = 3; src[4] = 4;\n dst[0] = 0; dst[1] = 0; dst[2] = 0; dst[3] = 0; dst[4] = 0; dst[5] = 0; dst[6] = 0; dst[7] = 0;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::scatter(thrust::cuda::par.on(s), src.begin(), src.end(), map.begin(), dst.begin());\n\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(dst[0], 0);\n ASSERT_EQUAL(dst[1], 2);\n ASSERT_EQUAL(dst[2], 4);\n ASSERT_EQUAL(dst[3], 1);\n ASSERT_EQUAL(dst[4], 0);\n ASSERT_EQUAL(dst[5], 0);\n ASSERT_EQUAL(dst[6], 0);\n ASSERT_EQUAL(dst[7], 3);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestScatterCudaStreams);\n\n\nvoid TestScatterIfCudaStreams()\n{\n typedef thrust::device_vector Vector;\n \n Vector flg(5); // predicate array\n Vector map(5); // scatter indices\n Vector src(5); // source vector\n Vector dst(8); // destination vector\n \n flg[0] = 0; flg[1] = 1; flg[2] = 0; flg[3] = 1; flg[4] = 0;\n map[0] = 6; map[1] = 3; map[2] = 1; map[3] = 7; map[4] = 2;\n src[0] = 0; src[1] = 1; src[2] = 2; src[3] = 3; src[4] = 4;\n dst[0] = 0; dst[1] = 0; dst[2] = 0; dst[3] = 0; dst[4] = 0; dst[5] = 0; dst[6] = 0; dst[7] = 0;\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n thrust::scatter_if(thrust::cuda::par.on(s), src.begin(), src.end(), map.begin(), flg.begin(), dst.begin());\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(dst[0], 0);\n ASSERT_EQUAL(dst[1], 0);\n ASSERT_EQUAL(dst[2], 0);\n ASSERT_EQUAL(dst[3], 1);\n ASSERT_EQUAL(dst[4], 0);\n ASSERT_EQUAL(dst[5], 0);\n ASSERT_EQUAL(dst[6], 0);\n ASSERT_EQUAL(dst[7], 3);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestScatterIfCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/scatter.cu"} {"input": "#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_difference_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first1, Iterator1 keys_last1,\n Iterator2 keys_first2, Iterator2 keys_last2,\n Iterator3 values_first1,\n Iterator4 values_first2,\n Iterator5 keys_result,\n Iterator6 values_result,\n Iterator7 result)\n{\n *result = thrust::set_difference_by_key(exec,\n keys_first1, keys_last1,\n keys_first2, keys_last2,\n values_first1,\n values_first2,\n keys_result,\n values_result);\n}\n\n\ntemplate\nvoid TestSetDifferenceByKeyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::iterator Iterator;\n\n Vector a_key(4), b_key(5);\n Vector a_val(4), b_val(5);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4; a_key[3] = 5;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0; a_val[3] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4; b_key[4] = 6;\n b_val[0] = 1; b_val[1] = 1; b_val[2] = 1; b_val[3] = 1; b_val[4] = 1;\n\n Vector ref_key(2), ref_val(2);\n ref_key[0] = 2; ref_key[1] = 5;\n ref_val[0] = 0; ref_val[1] = 0;\n\n Vector result_key(2), result_val(2);\n\n typedef thrust::pair iter_pair;\n\n thrust::device_vector end_vec(1);\n\n set_difference_by_key_kernel<<<1,1>>>(exec,\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n b_val.begin(),\n result_key.begin(),\n result_val.begin(),\n end_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n iter_pair end = end_vec.front();\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n}\n\n\nvoid TestSetDifferenceByKeyDeviceSeq()\n{\n TestSetDifferenceByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetDifferenceByKeyDeviceSeq);\n\n\nvoid TestSetDifferenceByKeyDeviceDevice()\n{\n TestSetDifferenceByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetDifferenceByKeyDeviceDevice);\n#endif\n\n\nvoid TestSetDifferenceByKeyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a_key(4), b_key(5);\n Vector a_val(4), b_val(5);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4; a_key[3] = 5;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0; a_val[3] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4; b_key[4] = 6;\n b_val[0] = 1; b_val[1] = 1; b_val[2] = 1; b_val[3] = 1; b_val[4] = 1;\n\n Vector ref_key(2), ref_val(2);\n ref_key[0] = 2; ref_key[1] = 5;\n ref_val[0] = 0; ref_val[1] = 0;\n\n Vector result_key(2), result_val(2);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::pair end =\n thrust::set_difference_by_key(thrust::cuda::par.on(s),\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n b_val.begin(),\n result_key.begin(),\n result_val.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSetDifferenceByKeyCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_difference_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first1, Iterator1 keys_last1,\n Iterator2 keys_first2, Iterator2 keys_last2,\n Iterator3 values_first1,\n Iterator4 values_first2,\n Iterator5 keys_result,\n Iterator6 values_result,\n Iterator7 result)\n{\n *result = thrust::set_difference_by_key(exec,\n keys_first1, keys_last1,\n keys_first2, keys_last2,\n values_first1,\n values_first2,\n keys_result,\n values_result);\n}\n\n\ntemplate\nvoid TestSetDifferenceByKeyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::iterator Iterator;\n\n Vector a_key(4), b_key(5);\n Vector a_val(4), b_val(5);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4; a_key[3] = 5;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0; a_val[3] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4; b_key[4] = 6;\n b_val[0] = 1; b_val[1] = 1; b_val[2] = 1; b_val[3] = 1; b_val[4] = 1;\n\n Vector ref_key(2), ref_val(2);\n ref_key[0] = 2; ref_key[1] = 5;\n ref_val[0] = 0; ref_val[1] = 0;\n\n Vector result_key(2), result_val(2);\n\n typedef thrust::pair iter_pair;\n\n thrust::device_vector end_vec(1);\n\n set_difference_by_key_kernel<<<1,1>>>(exec,\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n b_val.begin(),\n result_key.begin(),\n result_val.begin(),\n end_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n iter_pair end = end_vec.front();\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n}\n\n\nvoid TestSetDifferenceByKeyDeviceSeq()\n{\n TestSetDifferenceByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetDifferenceByKeyDeviceSeq);\n\n\nvoid TestSetDifferenceByKeyDeviceDevice()\n{\n TestSetDifferenceByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetDifferenceByKeyDeviceDevice);\n#endif\n\n\nvoid TestSetDifferenceByKeyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a_key(4), b_key(5);\n Vector a_val(4), b_val(5);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4; a_key[3] = 5;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0; a_val[3] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4; b_key[4] = 6;\n b_val[0] = 1; b_val[1] = 1; b_val[2] = 1; b_val[3] = 1; b_val[4] = 1;\n\n Vector ref_key(2), ref_val(2);\n ref_key[0] = 2; ref_key[1] = 5;\n ref_val[0] = 0; ref_val[1] = 0;\n\n Vector result_key(2), result_val(2);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::pair end =\n thrust::set_difference_by_key(thrust::cuda::par.on(s),\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n b_val.begin(),\n result_key.begin(),\n result_val.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSetDifferenceByKeyCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/set_difference_by_key.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid transform_inclusive_scan_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, Function1 f1, Function2 f2, Iterator3 result2)\n{\n *result2 = thrust::transform_inclusive_scan(exec, first, last, result1, f1, f2);\n}\n\n\ntemplate\n__global__\nvoid transform_exclusive_scan_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result, Function1 f1, T init, Function2 f2, Iterator3 result2)\n{\n *result2 = thrust::transform_exclusive_scan(exec, first, last, result, f1, init, f2);\n}\n\n\ntemplate\nvoid TestTransformScanDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::value_type T;\n \n typename Vector::iterator iter;\n \n Vector input(5);\n Vector ref(5);\n Vector output(5);\n \n input[0] = 1; input[1] = 3; input[2] = -2; input[3] = 4; input[4] = -5;\n \n Vector input_copy(input);\n\n thrust::device_vector iter_vec(1);\n \n // inclusive scan\n transform_inclusive_scan_kernel<<<1,1>>>(exec, input.begin(), input.end(), output.begin(), thrust::negate(), thrust::plus(), iter_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n iter = iter_vec[0];\n ref[0] = -1; ref[1] = -4; ref[2] = -2; ref[3] = -6; ref[4] = -1;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(ref, output);\n \n // exclusive scan with 0 init\n transform_exclusive_scan_kernel<<<1,1>>>(exec, input.begin(), input.end(), output.begin(), thrust::negate(), 0, thrust::plus(), iter_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ref[0] = 0; ref[1] = -1; ref[2] = -4; ref[3] = -2; ref[4] = -6;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(ref, output);\n \n // exclusive scan with nonzero init\n transform_exclusive_scan_kernel<<<1,1>>>(exec, input.begin(), input.end(), output.begin(), thrust::negate(), 3, thrust::plus(), iter_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n iter = iter_vec[0];\n ref[0] = 3; ref[1] = 2; ref[2] = -1; ref[3] = 1; ref[4] = -3;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(ref, output);\n \n // inplace inclusive scan\n input = input_copy;\n transform_inclusive_scan_kernel<<<1,1>>>(exec, input.begin(), input.end(), input.begin(), thrust::negate(), thrust::plus(), iter_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n iter = iter_vec[0];\n ref[0] = -1; ref[1] = -4; ref[2] = -2; ref[3] = -6; ref[4] = -1;\n ASSERT_EQUAL(std::size_t(iter - input.begin()), input.size());\n ASSERT_EQUAL(ref, input);\n \n // inplace exclusive scan with init\n input = input_copy;\n transform_exclusive_scan_kernel<<<1,1>>>(exec, input.begin(), input.end(), input.begin(), thrust::negate(), 3, thrust::plus(), iter_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n iter = iter_vec[0];\n ref[0] = 3; ref[1] = 2; ref[2] = -1; ref[3] = 1; ref[4] = -3;\n ASSERT_EQUAL(std::size_t(iter - input.begin()), input.size());\n ASSERT_EQUAL(ref, input);\n}\n\n\nvoid TestTransformScanDeviceSeq()\n{\n TestTransformScanDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTransformScanDeviceSeq);\n\n\nvoid TestTransformScanDeviceDevice()\n{\n TestTransformScanDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTransformScanDeviceDevice);\n#endif\n\n\nvoid TestTransformScanCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector::iterator iter;\n\n Vector input(5);\n Vector result(5);\n Vector output(5);\n\n input[0] = 1; input[1] = 3; input[2] = -2; input[3] = 4; input[4] = -5;\n\n Vector input_copy(input);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n // inclusive scan\n iter = thrust::transform_inclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), thrust::negate(), thrust::plus());\n cudaStreamSynchronize(s);\n\n result[0] = -1; result[1] = -4; result[2] = -2; result[3] = -6; result[4] = -1;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n \n // exclusive scan with 0 init\n iter = thrust::transform_exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), thrust::negate(), 0, thrust::plus());\n cudaStreamSynchronize(s);\n\n result[0] = 0; result[1] = -1; result[2] = -4; result[3] = -2; result[4] = -6;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n \n // exclusive scan with nonzero init\n iter = thrust::transform_exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), thrust::negate(), 3, thrust::plus());\n cudaStreamSynchronize(s);\n\n result[0] = 3; result[1] = 2; result[2] = -1; result[3] = 1; result[4] = -3;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n \n // inplace inclusive scan\n input = input_copy;\n iter = thrust::transform_inclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), input.begin(), thrust::negate(), thrust::plus());\n cudaStreamSynchronize(s);\n\n result[0] = -1; result[1] = -4; result[2] = -2; result[3] = -6; result[4] = -1;\n ASSERT_EQUAL(std::size_t(iter - input.begin()), input.size());\n ASSERT_EQUAL(input, result);\n\n // inplace exclusive scan with init\n input = input_copy;\n iter = thrust::transform_exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), input.begin(), thrust::negate(), 3, thrust::plus());\n cudaStreamSynchronize(s);\n\n result[0] = 3; result[1] = 2; result[2] = -1; result[3] = 1; result[4] = -3;\n ASSERT_EQUAL(std::size_t(iter - input.begin()), input.size());\n ASSERT_EQUAL(input, result);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestTransformScanCudaStreams);\n\nvoid TestTransformScanConstAccumulator()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector::iterator iter;\n\n Vector input(5);\n Vector reference(5);\n Vector output(5);\n\n input[0] = 1;\n input[1] = 3;\n input[2] = -2;\n input[3] = 4;\n input[4] = -5;\n\n thrust::transform_inclusive_scan(input.begin(),\n input.end(),\n output.begin(),\n thrust::identity(),\n thrust::plus());\n thrust::inclusive_scan(input.begin(), input.end(), reference.begin(), thrust::plus());\n\n ASSERT_EQUAL(output, reference);\n}\nDECLARE_UNITTEST(TestTransformScanConstAccumulator);\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid transform_inclusive_scan_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, Function1 f1, Function2 f2, Iterator3 result2)\n{\n *result2 = thrust::transform_inclusive_scan(exec, first, last, result1, f1, f2);\n}\n\n\ntemplate\n__global__\nvoid transform_exclusive_scan_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result, Function1 f1, T init, Function2 f2, Iterator3 result2)\n{\n *result2 = thrust::transform_exclusive_scan(exec, first, last, result, f1, init, f2);\n}\n\n\ntemplate\nvoid TestTransformScanDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::value_type T;\n \n typename Vector::iterator iter;\n \n Vector input(5);\n Vector ref(5);\n Vector output(5);\n \n input[0] = 1; input[1] = 3; input[2] = -2; input[3] = 4; input[4] = -5;\n \n Vector input_copy(input);\n\n thrust::device_vector iter_vec(1);\n \n // inclusive scan\n transform_inclusive_scan_kernel<<<1,1>>>(exec, input.begin(), input.end(), output.begin(), thrust::negate(), thrust::plus(), iter_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n iter = iter_vec[0];\n ref[0] = -1; ref[1] = -4; ref[2] = -2; ref[3] = -6; ref[4] = -1;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(ref, output);\n \n // exclusive scan with 0 init\n transform_exclusive_scan_kernel<<<1,1>>>(exec, input.begin(), input.end(), output.begin(), thrust::negate(), 0, thrust::plus(), iter_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ref[0] = 0; ref[1] = -1; ref[2] = -4; ref[3] = -2; ref[4] = -6;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(ref, output);\n \n // exclusive scan with nonzero init\n transform_exclusive_scan_kernel<<<1,1>>>(exec, input.begin(), input.end(), output.begin(), thrust::negate(), 3, thrust::plus(), iter_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n iter = iter_vec[0];\n ref[0] = 3; ref[1] = 2; ref[2] = -1; ref[3] = 1; ref[4] = -3;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(ref, output);\n \n // inplace inclusive scan\n input = input_copy;\n transform_inclusive_scan_kernel<<<1,1>>>(exec, input.begin(), input.end(), input.begin(), thrust::negate(), thrust::plus(), iter_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n iter = iter_vec[0];\n ref[0] = -1; ref[1] = -4; ref[2] = -2; ref[3] = -6; ref[4] = -1;\n ASSERT_EQUAL(std::size_t(iter - input.begin()), input.size());\n ASSERT_EQUAL(ref, input);\n \n // inplace exclusive scan with init\n input = input_copy;\n transform_exclusive_scan_kernel<<<1,1>>>(exec, input.begin(), input.end(), input.begin(), thrust::negate(), 3, thrust::plus(), iter_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n iter = iter_vec[0];\n ref[0] = 3; ref[1] = 2; ref[2] = -1; ref[3] = 1; ref[4] = -3;\n ASSERT_EQUAL(std::size_t(iter - input.begin()), input.size());\n ASSERT_EQUAL(ref, input);\n}\n\n\nvoid TestTransformScanDeviceSeq()\n{\n TestTransformScanDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTransformScanDeviceSeq);\n\n\nvoid TestTransformScanDeviceDevice()\n{\n TestTransformScanDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTransformScanDeviceDevice);\n#endif\n\n\nvoid TestTransformScanCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector::iterator iter;\n\n Vector input(5);\n Vector result(5);\n Vector output(5);\n\n input[0] = 1; input[1] = 3; input[2] = -2; input[3] = 4; input[4] = -5;\n\n Vector input_copy(input);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n // inclusive scan\n iter = thrust::transform_inclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), thrust::negate(), thrust::plus());\n hipStreamSynchronize(s);\n\n result[0] = -1; result[1] = -4; result[2] = -2; result[3] = -6; result[4] = -1;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n \n // exclusive scan with 0 init\n iter = thrust::transform_exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), thrust::negate(), 0, thrust::plus());\n hipStreamSynchronize(s);\n\n result[0] = 0; result[1] = -1; result[2] = -4; result[3] = -2; result[4] = -6;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n \n // exclusive scan with nonzero init\n iter = thrust::transform_exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), thrust::negate(), 3, thrust::plus());\n hipStreamSynchronize(s);\n\n result[0] = 3; result[1] = 2; result[2] = -1; result[3] = 1; result[4] = -3;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n \n // inplace inclusive scan\n input = input_copy;\n iter = thrust::transform_inclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), input.begin(), thrust::negate(), thrust::plus());\n hipStreamSynchronize(s);\n\n result[0] = -1; result[1] = -4; result[2] = -2; result[3] = -6; result[4] = -1;\n ASSERT_EQUAL(std::size_t(iter - input.begin()), input.size());\n ASSERT_EQUAL(input, result);\n\n // inplace exclusive scan with init\n input = input_copy;\n iter = thrust::transform_exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), input.begin(), thrust::negate(), 3, thrust::plus());\n hipStreamSynchronize(s);\n\n result[0] = 3; result[1] = 2; result[2] = -1; result[3] = 1; result[4] = -3;\n ASSERT_EQUAL(std::size_t(iter - input.begin()), input.size());\n ASSERT_EQUAL(input, result);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestTransformScanCudaStreams);\n\nvoid TestTransformScanConstAccumulator()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector::iterator iter;\n\n Vector input(5);\n Vector reference(5);\n Vector output(5);\n\n input[0] = 1;\n input[1] = 3;\n input[2] = -2;\n input[3] = 4;\n input[4] = -5;\n\n thrust::transform_inclusive_scan(input.begin(),\n input.end(),\n output.begin(),\n thrust::identity(),\n thrust::plus());\n thrust::inclusive_scan(input.begin(), input.end(), reference.begin(), thrust::plus());\n\n ASSERT_EQUAL(output, reference);\n}\nDECLARE_UNITTEST(TestTransformScanConstAccumulator);\n", "source": "/kaggle/working/thrust/testing/cuda/transform_scan.cu"} {"input": "#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid partition_point_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Predicate pred, Iterator2 result)\n{\n *result = thrust::partition_point(exec, first, last, pred);\n}\n\n\ntemplate\nstruct is_even\n{\n __host__ __device__\n bool operator()(T x) const { return ((int) x % 2) == 0; }\n};\n\n\ntemplate\nvoid TestPartitionPointDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::device_vector v = unittest::random_integers(n);\n typedef typename thrust::device_vector::iterator iterator;\n\n iterator ref = thrust::stable_partition(v.begin(), v.end(), is_even());\n\n thrust::device_vector result(1);\n partition_point_kernel<<<1,1>>>(exec, v.begin(), v.end(), is_even(), result.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n ASSERT_EQUAL(ref - v.begin(), (iterator)result[0] - v.begin());\n}\n\n\nvoid TestPartitionPointDeviceSeq()\n{\n TestPartitionPointDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestPartitionPointDeviceSeq);\n\n\nvoid TestPartitionPointDeviceDevice()\n{\n TestPartitionPointDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestPartitionPointDeviceDevice);\n#endif\n\n\nvoid TestPartitionPointCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n typedef Vector::iterator Iterator;\n\n Vector v(4);\n v[0] = 1; v[1] = 1; v[2] = 1; v[3] = 0;\n\n Iterator first = v.begin();\n\n Iterator last = v.begin() + 4;\n Iterator ref = first + 3;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n ASSERT_EQUAL_QUIET(ref, thrust::partition_point(thrust::cuda::par.on(s), first, last, thrust::identity()));\n\n last = v.begin() + 3;\n ref = last;\n ASSERT_EQUAL_QUIET(ref, thrust::partition_point(thrust::cuda::par.on(s), first, last, thrust::identity()));\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestPartitionPointCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid partition_point_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Predicate pred, Iterator2 result)\n{\n *result = thrust::partition_point(exec, first, last, pred);\n}\n\n\ntemplate\nstruct is_even\n{\n __host__ __device__\n bool operator()(T x) const { return ((int) x % 2) == 0; }\n};\n\n\ntemplate\nvoid TestPartitionPointDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::device_vector v = unittest::random_integers(n);\n typedef typename thrust::device_vector::iterator iterator;\n\n iterator ref = thrust::stable_partition(v.begin(), v.end(), is_even());\n\n thrust::device_vector result(1);\n partition_point_kernel<<<1,1>>>(exec, v.begin(), v.end(), is_even(), result.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n ASSERT_EQUAL(ref - v.begin(), (iterator)result[0] - v.begin());\n}\n\n\nvoid TestPartitionPointDeviceSeq()\n{\n TestPartitionPointDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestPartitionPointDeviceSeq);\n\n\nvoid TestPartitionPointDeviceDevice()\n{\n TestPartitionPointDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestPartitionPointDeviceDevice);\n#endif\n\n\nvoid TestPartitionPointCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n typedef Vector::iterator Iterator;\n\n Vector v(4);\n v[0] = 1; v[1] = 1; v[2] = 1; v[3] = 0;\n\n Iterator first = v.begin();\n\n Iterator last = v.begin() + 4;\n Iterator ref = first + 3;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n ASSERT_EQUAL_QUIET(ref, thrust::partition_point(thrust::cuda::par.on(s), first, last, thrust::identity()));\n\n last = v.begin() + 3;\n ref = last;\n ASSERT_EQUAL_QUIET(ref, thrust::partition_point(thrust::cuda::par.on(s), first, last, thrust::identity()));\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestPartitionPointCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/partition_point.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid merge_kernel(ExecutionPolicy exec,\n Iterator1 first1, Iterator1 last1,\n Iterator2 first2, Iterator2 last2,\n Iterator3 result1,\n Iterator4 result2)\n{\n *result2 = thrust::merge(exec, first1, last1, first2, last2, result1);\n}\n\n\ntemplate\nvoid TestMergeDevice(ExecutionPolicy exec)\n{\n size_t n = 10000;\n size_t sizes[] = {0, 1, n / 2, n, n + 1, 2 * n};\n size_t num_sizes = sizeof(sizes) / sizeof(size_t);\n\n thrust::host_vector random = unittest::random_integers(n + *thrust::max_element(sizes, sizes + num_sizes));\n\n thrust::host_vector h_a(random.begin(), random.begin() + n);\n thrust::host_vector h_b(random.begin() + n, random.end());\n\n thrust::stable_sort(h_a.begin(), h_a.end());\n thrust::stable_sort(h_b.begin(), h_b.end());\n \n thrust::device_vector d_a = h_a;\n thrust::device_vector d_b = h_b;\n\n for(size_t i = 0; i < num_sizes; i++)\n {\n size_t size = sizes[i];\n \n thrust::host_vector h_result(n + size);\n thrust::device_vector d_result(n + size);\n\n typename thrust::host_vector::iterator h_end;\n\n typedef typename thrust::device_vector::iterator iter_type;\n thrust::device_vector d_end(1);\n \n h_end = thrust::merge(h_a.begin(), h_a.end(),\n h_b.begin(), h_b.begin() + size,\n h_result.begin());\n h_result.resize(h_end - h_result.begin());\n\n merge_kernel<<<1,1>>>(exec,\n d_a.begin(), d_a.end(),\n d_b.begin(), d_b.begin() + size,\n d_result.begin(),\n d_end.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n d_result.resize((iter_type)d_end[0] - d_result.begin());\n\n ASSERT_EQUAL(h_result, d_result);\n }\n}\n\n\nvoid TestMergeDeviceSeq()\n{\n TestMergeDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestMergeDeviceSeq);\n\n\nvoid TestMergeDeviceDevice()\n{\n TestMergeDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestMergeDeviceDevice);\n#endif\n\n\nvoid TestMergeCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a(3), b(4);\n\n a[0] = 0; a[1] = 2; a[2] = 4;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4;\n\n Vector ref(7);\n ref[0] = 0;\n ref[1] = 0;\n ref[2] = 2;\n ref[3] = 3;\n ref[4] = 3;\n ref[5] = 4;\n ref[6] = 4;\n\n Vector result(7);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n Iterator end = thrust::merge(thrust::cuda::par.on(s),\n a.begin(), a.end(),\n b.begin(), b.end(),\n result.begin());\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestMergeCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid merge_kernel(ExecutionPolicy exec,\n Iterator1 first1, Iterator1 last1,\n Iterator2 first2, Iterator2 last2,\n Iterator3 result1,\n Iterator4 result2)\n{\n *result2 = thrust::merge(exec, first1, last1, first2, last2, result1);\n}\n\n\ntemplate\nvoid TestMergeDevice(ExecutionPolicy exec)\n{\n size_t n = 10000;\n size_t sizes[] = {0, 1, n / 2, n, n + 1, 2 * n};\n size_t num_sizes = sizeof(sizes) / sizeof(size_t);\n\n thrust::host_vector random = unittest::random_integers(n + *thrust::max_element(sizes, sizes + num_sizes));\n\n thrust::host_vector h_a(random.begin(), random.begin() + n);\n thrust::host_vector h_b(random.begin() + n, random.end());\n\n thrust::stable_sort(h_a.begin(), h_a.end());\n thrust::stable_sort(h_b.begin(), h_b.end());\n \n thrust::device_vector d_a = h_a;\n thrust::device_vector d_b = h_b;\n\n for(size_t i = 0; i < num_sizes; i++)\n {\n size_t size = sizes[i];\n \n thrust::host_vector h_result(n + size);\n thrust::device_vector d_result(n + size);\n\n typename thrust::host_vector::iterator h_end;\n\n typedef typename thrust::device_vector::iterator iter_type;\n thrust::device_vector d_end(1);\n \n h_end = thrust::merge(h_a.begin(), h_a.end(),\n h_b.begin(), h_b.begin() + size,\n h_result.begin());\n h_result.resize(h_end - h_result.begin());\n\n merge_kernel<<<1,1>>>(exec,\n d_a.begin(), d_a.end(),\n d_b.begin(), d_b.begin() + size,\n d_result.begin(),\n d_end.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n d_result.resize((iter_type)d_end[0] - d_result.begin());\n\n ASSERT_EQUAL(h_result, d_result);\n }\n}\n\n\nvoid TestMergeDeviceSeq()\n{\n TestMergeDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestMergeDeviceSeq);\n\n\nvoid TestMergeDeviceDevice()\n{\n TestMergeDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestMergeDeviceDevice);\n#endif\n\n\nvoid TestMergeCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a(3), b(4);\n\n a[0] = 0; a[1] = 2; a[2] = 4;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4;\n\n Vector ref(7);\n ref[0] = 0;\n ref[1] = 0;\n ref[2] = 2;\n ref[3] = 3;\n ref[4] = 3;\n ref[5] = 4;\n ref[6] = 4;\n\n Vector result(7);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n Iterator end = thrust::merge(thrust::cuda::par.on(s),\n a.begin(), a.end(),\n b.begin(), b.end(),\n result.begin());\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestMergeCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/merge.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid swap_ranges_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2)\n{\n thrust::swap_ranges(exec, first1, last1, first2);\n}\n\n\ntemplate\nvoid TestSwapRangesDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n\n Vector v1(5);\n v1[0] = 0; v1[1] = 1; v1[2] = 2; v1[3] = 3; v1[4] = 4;\n\n Vector v2(5);\n v2[0] = 5; v2[1] = 6; v2[2] = 7; v2[3] = 8; v2[4] = 9;\n\n swap_ranges_kernel<<<1,1>>>(exec, v1.begin(), v1.end(), v2.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n ASSERT_EQUAL(v1[0], 5);\n ASSERT_EQUAL(v1[1], 6);\n ASSERT_EQUAL(v1[2], 7);\n ASSERT_EQUAL(v1[3], 8);\n ASSERT_EQUAL(v1[4], 9);\n \n ASSERT_EQUAL(v2[0], 0);\n ASSERT_EQUAL(v2[1], 1);\n ASSERT_EQUAL(v2[2], 2);\n ASSERT_EQUAL(v2[3], 3);\n ASSERT_EQUAL(v2[4], 4);\n}\n\nvoid TestSwapRangesDeviceSeq()\n{\n TestSwapRangesDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSwapRangesDeviceSeq);\n\nvoid TestSwapRangesDeviceDevice()\n{\n TestSwapRangesDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSwapRangesDeviceDevice);\n#endif\n\nvoid TestSwapRangesCudaStreams()\n{\n typedef thrust::device_vector Vector;\n\n Vector v1(5);\n v1[0] = 0; v1[1] = 1; v1[2] = 2; v1[3] = 3; v1[4] = 4;\n\n Vector v2(5);\n v2[0] = 5; v2[1] = 6; v2[2] = 7; v2[3] = 8; v2[4] = 9;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::swap_ranges(thrust::cuda::par.on(s), v1.begin(), v1.end(), v2.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(v1[0], 5);\n ASSERT_EQUAL(v1[1], 6);\n ASSERT_EQUAL(v1[2], 7);\n ASSERT_EQUAL(v1[3], 8);\n ASSERT_EQUAL(v1[4], 9);\n \n ASSERT_EQUAL(v2[0], 0);\n ASSERT_EQUAL(v2[1], 1);\n ASSERT_EQUAL(v2[2], 2);\n ASSERT_EQUAL(v2[3], 3);\n ASSERT_EQUAL(v2[4], 4);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSwapRangesCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid swap_ranges_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2)\n{\n thrust::swap_ranges(exec, first1, last1, first2);\n}\n\n\ntemplate\nvoid TestSwapRangesDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n\n Vector v1(5);\n v1[0] = 0; v1[1] = 1; v1[2] = 2; v1[3] = 3; v1[4] = 4;\n\n Vector v2(5);\n v2[0] = 5; v2[1] = 6; v2[2] = 7; v2[3] = 8; v2[4] = 9;\n\n swap_ranges_kernel<<<1,1>>>(exec, v1.begin(), v1.end(), v2.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n ASSERT_EQUAL(v1[0], 5);\n ASSERT_EQUAL(v1[1], 6);\n ASSERT_EQUAL(v1[2], 7);\n ASSERT_EQUAL(v1[3], 8);\n ASSERT_EQUAL(v1[4], 9);\n \n ASSERT_EQUAL(v2[0], 0);\n ASSERT_EQUAL(v2[1], 1);\n ASSERT_EQUAL(v2[2], 2);\n ASSERT_EQUAL(v2[3], 3);\n ASSERT_EQUAL(v2[4], 4);\n}\n\nvoid TestSwapRangesDeviceSeq()\n{\n TestSwapRangesDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSwapRangesDeviceSeq);\n\nvoid TestSwapRangesDeviceDevice()\n{\n TestSwapRangesDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSwapRangesDeviceDevice);\n#endif\n\nvoid TestSwapRangesCudaStreams()\n{\n typedef thrust::device_vector Vector;\n\n Vector v1(5);\n v1[0] = 0; v1[1] = 1; v1[2] = 2; v1[3] = 3; v1[4] = 4;\n\n Vector v2(5);\n v2[0] = 5; v2[1] = 6; v2[2] = 7; v2[3] = 8; v2[4] = 9;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::swap_ranges(thrust::cuda::par.on(s), v1.begin(), v1.end(), v2.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(v1[0], 5);\n ASSERT_EQUAL(v1[1], 6);\n ASSERT_EQUAL(v1[2], 7);\n ASSERT_EQUAL(v1[3], 8);\n ASSERT_EQUAL(v1[4], 9);\n \n ASSERT_EQUAL(v2[0], 0);\n ASSERT_EQUAL(v2[1], 1);\n ASSERT_EQUAL(v2[2], 2);\n ASSERT_EQUAL(v2[3], 3);\n ASSERT_EQUAL(v2[4], 4);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSwapRangesCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/swap_ranges.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid transform_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, Function f, Iterator3 result2)\n{\n *result2 = thrust::transform(exec, first, last, result1, f);\n}\n\n\ntemplate\nvoid TestTransformUnaryDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::value_type T;\n \n typename Vector::iterator iter;\n \n Vector input(3);\n Vector output(3);\n Vector result(3);\n input[0] = 1; input[1] = -2; input[2] = 3;\n result[0] = -1; result[1] = 2; result[2] = -3;\n\n thrust::device_vector iter_vec(1);\n \n transform_kernel<<<1,1>>>(exec, input.begin(), input.end(), output.begin(), thrust::negate(), iter_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(output, result);\n}\n\nvoid TestTransformUnaryDeviceSeq()\n{\n TestTransformUnaryDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTransformUnaryDeviceSeq);\n\nvoid TestTransformUnaryDeviceDevice()\n{\n TestTransformUnaryDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTransformUnaryDeviceDevice);\n\n\ntemplate\n__global__\nvoid transform_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, Function f, Predicate pred, Iterator3 result2)\n{\n *result2 = thrust::transform_if(exec, first, last, result1, f, pred);\n}\n\n\ntemplate\nvoid TestTransformIfUnaryNoStencilDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::value_type T;\n \n typename Vector::iterator iter;\n \n Vector input(3);\n Vector output(3);\n Vector result(3);\n \n input[0] = 0; input[1] = -2; input[2] = 0;\n output[0] = -1; output[1] = -2; output[2] = -3; \n result[0] = -1; result[1] = 2; result[2] = -3;\n\n thrust::device_vector iter_vec(1);\n \n transform_if_kernel<<<1,1>>>(exec,\n input.begin(), input.end(),\n output.begin(),\n thrust::negate(),\n thrust::identity(),\n iter_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(output, result);\n}\n\nvoid TestTransformIfUnaryNoStencilDeviceSeq()\n{\n TestTransformIfUnaryNoStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTransformIfUnaryNoStencilDeviceSeq);\n\nvoid TestTransformIfUnaryNoStencilDeviceDevice()\n{\n TestTransformIfUnaryNoStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTransformIfUnaryNoStencilDeviceDevice);\n\n\ntemplate\n__global__\nvoid transform_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Iterator3 result1, Function f, Predicate pred, Iterator4 result2)\n{\n *result2 = thrust::transform_if(exec, first, last, stencil_first, result1, f, pred);\n}\n\n\ntemplate\nvoid TestTransformIfUnaryDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::value_type T;\n \n typename Vector::iterator iter;\n \n Vector input(3);\n Vector stencil(3);\n Vector output(3);\n Vector result(3);\n \n input[0] = 1; input[1] = -2; input[2] = 3;\n output[0] = 1; output[1] = 2; output[2] = 3; \n stencil[0] = 1; stencil[1] = 0; stencil[2] = 1;\n result[0] = -1; result[1] = 2; result[2] = -3;\n\n thrust::device_vector iter_vec(1);\n \n transform_if_kernel<<<1,1>>>(exec,\n input.begin(), input.end(),\n stencil.begin(),\n output.begin(),\n thrust::negate(),\n thrust::identity(),\n iter_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(output, result);\n}\n\nvoid TestTransformIfUnaryDeviceSeq()\n{\n TestTransformIfUnaryDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTransformIfUnaryDeviceSeq);\n\nvoid TestTransformIfUnaryDeviceDevice()\n{\n TestTransformIfUnaryDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTransformIfUnaryDeviceDevice);\n\n\ntemplate\n__global__\nvoid transform_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2, Iterator3 result1, Function f, Iterator4 result2)\n{\n *result2 = thrust::transform(exec, first1, last1, first2, result1, f);\n}\n\n\ntemplate\nvoid TestTransformBinaryDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::value_type T;\n \n typename Vector::iterator iter;\n \n Vector input1(3);\n Vector input2(3);\n Vector output(3);\n Vector result(3);\n input1[0] = 1; input1[1] = -2; input1[2] = 3;\n input2[0] = -4; input2[1] = 5; input2[2] = 6;\n result[0] = 5; result[1] = -7; result[2] = -3;\n\n thrust::device_vector iter_vec(1);\n \n transform_kernel<<<1,1>>>(exec, input1.begin(), input1.end(), input2.begin(), output.begin(), thrust::minus(), iter_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(std::size_t(iter - output.begin()), input1.size());\n ASSERT_EQUAL(output, result);\n}\n\nvoid TestTransformBinaryDeviceSeq()\n{\n TestTransformBinaryDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTransformBinaryDeviceSeq);\n\nvoid TestTransformBinaryDeviceDevice()\n{\n TestTransformBinaryDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTransformBinaryDeviceDevice);\n\n\ntemplate\n__global__\nvoid transform_if_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2, Iterator3 stencil_first, Iterator4 result1, Function f, Predicate pred, Iterator5 result2)\n{\n *result2 = thrust::transform_if(exec, first1, last1, first2, stencil_first, result1, f, pred);\n}\n\n\ntemplate\nvoid TestTransformIfBinaryDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::value_type T;\n \n typename Vector::iterator iter;\n \n Vector input1(3);\n Vector input2(3);\n Vector stencil(3);\n Vector output(3);\n Vector result(3);\n \n input1[0] = 1; input1[1] = -2; input1[2] = 3;\n input2[0] = -4; input2[1] = 5; input2[2] = 6;\n stencil[0] = 0; stencil[1] = 1; stencil[2] = 0;\n output[0] = 1; output[1] = 2; output[2] = 3;\n result[0] = 5; result[1] = 2; result[2] = -3;\n \n thrust::identity identity;\n\n thrust::device_vector iter_vec(1);\n \n transform_if_kernel<<<1,1>>>(exec,\n input1.begin(), input1.end(),\n input2.begin(),\n stencil.begin(),\n output.begin(),\n thrust::minus(),\n thrust::not1(identity),\n iter_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(std::size_t(iter - output.begin()), input1.size());\n ASSERT_EQUAL(output, result);\n}\n\nvoid TestTransformIfBinaryDeviceSeq()\n{\n TestTransformIfBinaryDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTransformIfBinaryDeviceSeq);\n\nvoid TestTransformIfBinaryDeviceDevice()\n{\n TestTransformIfBinaryDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTransformIfBinaryDeviceDevice);\n#endif\n\nvoid TestTransformUnaryCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector::iterator iter;\n\n Vector input(3);\n Vector output(3);\n Vector result(3);\n input[0] = 1; input[1] = -2; input[2] = 3;\n result[0] = -1; result[1] = 2; result[2] = -3;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n iter = thrust::transform(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), thrust::negate());\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(output, result);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestTransformUnaryCudaStreams);\n\n\nvoid TestTransformBinaryCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector::iterator iter;\n\n Vector input1(3);\n Vector input2(3);\n Vector output(3);\n Vector result(3);\n input1[0] = 1; input1[1] = -2; input1[2] = 3;\n input2[0] = -4; input2[1] = 5; input2[2] = 6;\n result[0] = 5; result[1] = -7; result[2] = -3;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n iter = thrust::transform(thrust::cuda::par.on(s), input1.begin(), input1.end(), input2.begin(), output.begin(), thrust::minus());\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(std::size_t(iter - output.begin()), input1.size());\n ASSERT_EQUAL(output, result);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestTransformBinaryCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid transform_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, Function f, Iterator3 result2)\n{\n *result2 = thrust::transform(exec, first, last, result1, f);\n}\n\n\ntemplate\nvoid TestTransformUnaryDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::value_type T;\n \n typename Vector::iterator iter;\n \n Vector input(3);\n Vector output(3);\n Vector result(3);\n input[0] = 1; input[1] = -2; input[2] = 3;\n result[0] = -1; result[1] = 2; result[2] = -3;\n\n thrust::device_vector iter_vec(1);\n \n transform_kernel<<<1,1>>>(exec, input.begin(), input.end(), output.begin(), thrust::negate(), iter_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(output, result);\n}\n\nvoid TestTransformUnaryDeviceSeq()\n{\n TestTransformUnaryDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTransformUnaryDeviceSeq);\n\nvoid TestTransformUnaryDeviceDevice()\n{\n TestTransformUnaryDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTransformUnaryDeviceDevice);\n\n\ntemplate\n__global__\nvoid transform_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, Function f, Predicate pred, Iterator3 result2)\n{\n *result2 = thrust::transform_if(exec, first, last, result1, f, pred);\n}\n\n\ntemplate\nvoid TestTransformIfUnaryNoStencilDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::value_type T;\n \n typename Vector::iterator iter;\n \n Vector input(3);\n Vector output(3);\n Vector result(3);\n \n input[0] = 0; input[1] = -2; input[2] = 0;\n output[0] = -1; output[1] = -2; output[2] = -3; \n result[0] = -1; result[1] = 2; result[2] = -3;\n\n thrust::device_vector iter_vec(1);\n \n transform_if_kernel<<<1,1>>>(exec,\n input.begin(), input.end(),\n output.begin(),\n thrust::negate(),\n thrust::identity(),\n iter_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(output, result);\n}\n\nvoid TestTransformIfUnaryNoStencilDeviceSeq()\n{\n TestTransformIfUnaryNoStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTransformIfUnaryNoStencilDeviceSeq);\n\nvoid TestTransformIfUnaryNoStencilDeviceDevice()\n{\n TestTransformIfUnaryNoStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTransformIfUnaryNoStencilDeviceDevice);\n\n\ntemplate\n__global__\nvoid transform_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Iterator3 result1, Function f, Predicate pred, Iterator4 result2)\n{\n *result2 = thrust::transform_if(exec, first, last, stencil_first, result1, f, pred);\n}\n\n\ntemplate\nvoid TestTransformIfUnaryDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::value_type T;\n \n typename Vector::iterator iter;\n \n Vector input(3);\n Vector stencil(3);\n Vector output(3);\n Vector result(3);\n \n input[0] = 1; input[1] = -2; input[2] = 3;\n output[0] = 1; output[1] = 2; output[2] = 3; \n stencil[0] = 1; stencil[1] = 0; stencil[2] = 1;\n result[0] = -1; result[1] = 2; result[2] = -3;\n\n thrust::device_vector iter_vec(1);\n \n transform_if_kernel<<<1,1>>>(exec,\n input.begin(), input.end(),\n stencil.begin(),\n output.begin(),\n thrust::negate(),\n thrust::identity(),\n iter_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(output, result);\n}\n\nvoid TestTransformIfUnaryDeviceSeq()\n{\n TestTransformIfUnaryDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTransformIfUnaryDeviceSeq);\n\nvoid TestTransformIfUnaryDeviceDevice()\n{\n TestTransformIfUnaryDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTransformIfUnaryDeviceDevice);\n\n\ntemplate\n__global__\nvoid transform_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2, Iterator3 result1, Function f, Iterator4 result2)\n{\n *result2 = thrust::transform(exec, first1, last1, first2, result1, f);\n}\n\n\ntemplate\nvoid TestTransformBinaryDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::value_type T;\n \n typename Vector::iterator iter;\n \n Vector input1(3);\n Vector input2(3);\n Vector output(3);\n Vector result(3);\n input1[0] = 1; input1[1] = -2; input1[2] = 3;\n input2[0] = -4; input2[1] = 5; input2[2] = 6;\n result[0] = 5; result[1] = -7; result[2] = -3;\n\n thrust::device_vector iter_vec(1);\n \n transform_kernel<<<1,1>>>(exec, input1.begin(), input1.end(), input2.begin(), output.begin(), thrust::minus(), iter_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(std::size_t(iter - output.begin()), input1.size());\n ASSERT_EQUAL(output, result);\n}\n\nvoid TestTransformBinaryDeviceSeq()\n{\n TestTransformBinaryDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTransformBinaryDeviceSeq);\n\nvoid TestTransformBinaryDeviceDevice()\n{\n TestTransformBinaryDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTransformBinaryDeviceDevice);\n\n\ntemplate\n__global__\nvoid transform_if_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2, Iterator3 stencil_first, Iterator4 result1, Function f, Predicate pred, Iterator5 result2)\n{\n *result2 = thrust::transform_if(exec, first1, last1, first2, stencil_first, result1, f, pred);\n}\n\n\ntemplate\nvoid TestTransformIfBinaryDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::value_type T;\n \n typename Vector::iterator iter;\n \n Vector input1(3);\n Vector input2(3);\n Vector stencil(3);\n Vector output(3);\n Vector result(3);\n \n input1[0] = 1; input1[1] = -2; input1[2] = 3;\n input2[0] = -4; input2[1] = 5; input2[2] = 6;\n stencil[0] = 0; stencil[1] = 1; stencil[2] = 0;\n output[0] = 1; output[1] = 2; output[2] = 3;\n result[0] = 5; result[1] = 2; result[2] = -3;\n \n thrust::identity identity;\n\n thrust::device_vector iter_vec(1);\n \n transform_if_kernel<<<1,1>>>(exec,\n input1.begin(), input1.end(),\n input2.begin(),\n stencil.begin(),\n output.begin(),\n thrust::minus(),\n thrust::not1(identity),\n iter_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(std::size_t(iter - output.begin()), input1.size());\n ASSERT_EQUAL(output, result);\n}\n\nvoid TestTransformIfBinaryDeviceSeq()\n{\n TestTransformIfBinaryDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTransformIfBinaryDeviceSeq);\n\nvoid TestTransformIfBinaryDeviceDevice()\n{\n TestTransformIfBinaryDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTransformIfBinaryDeviceDevice);\n#endif\n\nvoid TestTransformUnaryCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector::iterator iter;\n\n Vector input(3);\n Vector output(3);\n Vector result(3);\n input[0] = 1; input[1] = -2; input[2] = 3;\n result[0] = -1; result[1] = 2; result[2] = -3;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n iter = thrust::transform(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), thrust::negate());\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(output, result);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestTransformUnaryCudaStreams);\n\n\nvoid TestTransformBinaryCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector::iterator iter;\n\n Vector input1(3);\n Vector input2(3);\n Vector output(3);\n Vector result(3);\n input1[0] = 1; input1[1] = -2; input1[2] = 3;\n input2[0] = -4; input2[1] = 5; input2[2] = 6;\n result[0] = 5; result[1] = -7; result[2] = -3;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n iter = thrust::transform(thrust::cuda::par.on(s), input1.begin(), input1.end(), input2.begin(), output.begin(), thrust::minus());\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(std::size_t(iter - output.begin()), input1.size());\n ASSERT_EQUAL(output, result);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestTransformBinaryCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/transform.cu"} {"input": "#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid inclusive_scan_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n thrust::inclusive_scan(exec, first, last, result);\n}\n\n\ntemplate\n__global__\nvoid exclusive_scan_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n thrust::exclusive_scan(exec, first, last, result);\n}\n\n\ntemplate\n__global__\nvoid exclusive_scan_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result, T init)\n{\n thrust::exclusive_scan(exec, first, last, result, init);\n}\n\n\ntemplate\nvoid TestScanDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_input = unittest::random_integers(n);\n thrust::device_vector d_input = h_input;\n \n thrust::host_vector h_output(n);\n thrust::device_vector d_output(n);\n \n thrust::inclusive_scan(h_input.begin(), h_input.end(), h_output.begin());\n\n inclusive_scan_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_output.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(d_output, h_output);\n \n thrust::exclusive_scan(h_input.begin(), h_input.end(), h_output.begin());\n\n exclusive_scan_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_output.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(d_output, h_output);\n \n thrust::exclusive_scan(h_input.begin(), h_input.end(), h_output.begin(), (T) 11);\n\n exclusive_scan_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_output.begin(), (T) 11);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(d_output, h_output);\n \n // in-place scans\n h_output = h_input;\n d_output = d_input;\n\n thrust::inclusive_scan(h_output.begin(), h_output.end(), h_output.begin());\n\n inclusive_scan_kernel<<<1,1>>>(exec, d_output.begin(), d_output.end(), d_output.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(d_output, h_output);\n \n h_output = h_input;\n d_output = d_input;\n \n thrust::exclusive_scan(h_output.begin(), h_output.end(), h_output.begin());\n\n exclusive_scan_kernel<<<1,1>>>(exec, d_output.begin(), d_output.end(), d_output.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(d_output, h_output);\n}\n\n\ntemplate\nstruct TestScanDeviceSeq\n{\n void operator()(const size_t n)\n {\n TestScanDevice(thrust::seq, n);\n }\n};\nVariableUnitTest TestScanDeviceSeqInstance;\n\n\ntemplate\nstruct TestScanDeviceDevice\n{\n void operator()(const size_t n)\n {\n TestScanDevice(thrust::device, n);\n }\n};\nVariableUnitTest TestScanDeviceDeviceInstance;\n#endif\n\n\nvoid TestScanCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector::iterator iter;\n\n Vector input(5);\n Vector result(5);\n Vector output(5);\n\n input[0] = 1; input[1] = 3; input[2] = -2; input[3] = 4; input[4] = -5;\n\n Vector input_copy(input);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n // inclusive scan\n iter = thrust::inclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin());\n cudaStreamSynchronize(s);\n\n result[0] = 1; result[1] = 4; result[2] = 2; result[3] = 6; result[4] = 1;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n \n // exclusive scan\n iter = thrust::exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), 0);\n cudaStreamSynchronize(s);\n\n result[0] = 0; result[1] = 1; result[2] = 4; result[3] = 2; result[4] = 6;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n \n // exclusive scan with init\n iter = thrust::exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), 3);\n cudaStreamSynchronize(s);\n\n result[0] = 3; result[1] = 4; result[2] = 7; result[3] = 5; result[4] = 9;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n \n // inclusive scan with op\n iter = thrust::inclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), thrust::plus());\n cudaStreamSynchronize(s);\n\n result[0] = 1; result[1] = 4; result[2] = 2; result[3] = 6; result[4] = 1;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n\n // exclusive scan with init and op\n iter = thrust::exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), 3, thrust::plus());\n cudaStreamSynchronize(s);\n\n result[0] = 3; result[1] = 4; result[2] = 7; result[3] = 5; result[4] = 9;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n\n // inplace inclusive scan\n input = input_copy;\n iter = thrust::inclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), input.begin());\n cudaStreamSynchronize(s);\n\n result[0] = 1; result[1] = 4; result[2] = 2; result[3] = 6; result[4] = 1;\n ASSERT_EQUAL(std::size_t(iter - input.begin()), input.size());\n ASSERT_EQUAL(input, result);\n\n // inplace exclusive scan with init\n input = input_copy;\n iter = thrust::exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), input.begin(), 3);\n cudaStreamSynchronize(s);\n\n result[0] = 3; result[1] = 4; result[2] = 7; result[3] = 5; result[4] = 9;\n ASSERT_EQUAL(std::size_t(iter - input.begin()), input.size());\n ASSERT_EQUAL(input, result);\n\n // inplace exclusive scan with implicit init=0\n input = input_copy;\n iter = thrust::exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), input.begin());\n cudaStreamSynchronize(s);\n\n result[0] = 0; result[1] = 1; result[2] = 4; result[3] = 2; result[4] = 6;\n ASSERT_EQUAL(std::size_t(iter - input.begin()), input.size());\n ASSERT_EQUAL(input, result);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestScanCudaStreams);\n\ntemplate \nstruct const_ref_plus_mod3\n{\n T * table;\n\n const_ref_plus_mod3(T * table) : table(table) {}\n\n __host__ __device__\n const T& operator()(T a, T b)\n {\n return table[(int) (a + b)];\n }\n};\n\nstatic void TestInclusiveScanWithConstAccumulator(void)\n{\n // add numbers modulo 3 with external lookup table\n thrust::device_vector data(7);\n data[0] = 0;\n data[1] = 1;\n data[2] = 2;\n data[3] = 1;\n data[4] = 2;\n data[5] = 0;\n data[6] = 1;\n\n thrust::device_vector table(6);\n table[0] = 0;\n table[1] = 1;\n table[2] = 2;\n table[3] = 0;\n table[4] = 1;\n table[5] = 2;\n\n thrust::inclusive_scan(data.begin(), data.end(), data.begin(), const_ref_plus_mod3(thrust::raw_pointer_cast(&table[0])));\n \n ASSERT_EQUAL(data[0], 0);\n ASSERT_EQUAL(data[1], 1);\n ASSERT_EQUAL(data[2], 0);\n ASSERT_EQUAL(data[3], 1);\n ASSERT_EQUAL(data[4], 0);\n ASSERT_EQUAL(data[5], 0);\n ASSERT_EQUAL(data[6], 1);\n}\nDECLARE_UNITTEST(TestInclusiveScanWithConstAccumulator);\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid inclusive_scan_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n thrust::inclusive_scan(exec, first, last, result);\n}\n\n\ntemplate\n__global__\nvoid exclusive_scan_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n thrust::exclusive_scan(exec, first, last, result);\n}\n\n\ntemplate\n__global__\nvoid exclusive_scan_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result, T init)\n{\n thrust::exclusive_scan(exec, first, last, result, init);\n}\n\n\ntemplate\nvoid TestScanDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_input = unittest::random_integers(n);\n thrust::device_vector d_input = h_input;\n \n thrust::host_vector h_output(n);\n thrust::device_vector d_output(n);\n \n thrust::inclusive_scan(h_input.begin(), h_input.end(), h_output.begin());\n\n inclusive_scan_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_output.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(d_output, h_output);\n \n thrust::exclusive_scan(h_input.begin(), h_input.end(), h_output.begin());\n\n exclusive_scan_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_output.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(d_output, h_output);\n \n thrust::exclusive_scan(h_input.begin(), h_input.end(), h_output.begin(), (T) 11);\n\n exclusive_scan_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_output.begin(), (T) 11);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(d_output, h_output);\n \n // in-place scans\n h_output = h_input;\n d_output = d_input;\n\n thrust::inclusive_scan(h_output.begin(), h_output.end(), h_output.begin());\n\n inclusive_scan_kernel<<<1,1>>>(exec, d_output.begin(), d_output.end(), d_output.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(d_output, h_output);\n \n h_output = h_input;\n d_output = d_input;\n \n thrust::exclusive_scan(h_output.begin(), h_output.end(), h_output.begin());\n\n exclusive_scan_kernel<<<1,1>>>(exec, d_output.begin(), d_output.end(), d_output.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(d_output, h_output);\n}\n\n\ntemplate\nstruct TestScanDeviceSeq\n{\n void operator()(const size_t n)\n {\n TestScanDevice(thrust::seq, n);\n }\n};\nVariableUnitTest TestScanDeviceSeqInstance;\n\n\ntemplate\nstruct TestScanDeviceDevice\n{\n void operator()(const size_t n)\n {\n TestScanDevice(thrust::device, n);\n }\n};\nVariableUnitTest TestScanDeviceDeviceInstance;\n#endif\n\n\nvoid TestScanCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector::iterator iter;\n\n Vector input(5);\n Vector result(5);\n Vector output(5);\n\n input[0] = 1; input[1] = 3; input[2] = -2; input[3] = 4; input[4] = -5;\n\n Vector input_copy(input);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n // inclusive scan\n iter = thrust::inclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin());\n hipStreamSynchronize(s);\n\n result[0] = 1; result[1] = 4; result[2] = 2; result[3] = 6; result[4] = 1;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n \n // exclusive scan\n iter = thrust::exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), 0);\n hipStreamSynchronize(s);\n\n result[0] = 0; result[1] = 1; result[2] = 4; result[3] = 2; result[4] = 6;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n \n // exclusive scan with init\n iter = thrust::exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), 3);\n hipStreamSynchronize(s);\n\n result[0] = 3; result[1] = 4; result[2] = 7; result[3] = 5; result[4] = 9;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n \n // inclusive scan with op\n iter = thrust::inclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), thrust::plus());\n hipStreamSynchronize(s);\n\n result[0] = 1; result[1] = 4; result[2] = 2; result[3] = 6; result[4] = 1;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n\n // exclusive scan with init and op\n iter = thrust::exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), 3, thrust::plus());\n hipStreamSynchronize(s);\n\n result[0] = 3; result[1] = 4; result[2] = 7; result[3] = 5; result[4] = 9;\n ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());\n ASSERT_EQUAL(input, input_copy);\n ASSERT_EQUAL(output, result);\n\n // inplace inclusive scan\n input = input_copy;\n iter = thrust::inclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), input.begin());\n hipStreamSynchronize(s);\n\n result[0] = 1; result[1] = 4; result[2] = 2; result[3] = 6; result[4] = 1;\n ASSERT_EQUAL(std::size_t(iter - input.begin()), input.size());\n ASSERT_EQUAL(input, result);\n\n // inplace exclusive scan with init\n input = input_copy;\n iter = thrust::exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), input.begin(), 3);\n hipStreamSynchronize(s);\n\n result[0] = 3; result[1] = 4; result[2] = 7; result[3] = 5; result[4] = 9;\n ASSERT_EQUAL(std::size_t(iter - input.begin()), input.size());\n ASSERT_EQUAL(input, result);\n\n // inplace exclusive scan with implicit init=0\n input = input_copy;\n iter = thrust::exclusive_scan(thrust::cuda::par.on(s), input.begin(), input.end(), input.begin());\n hipStreamSynchronize(s);\n\n result[0] = 0; result[1] = 1; result[2] = 4; result[3] = 2; result[4] = 6;\n ASSERT_EQUAL(std::size_t(iter - input.begin()), input.size());\n ASSERT_EQUAL(input, result);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestScanCudaStreams);\n\ntemplate \nstruct const_ref_plus_mod3\n{\n T * table;\n\n const_ref_plus_mod3(T * table) : table(table) {}\n\n __host__ __device__\n const T& operator()(T a, T b)\n {\n return table[(int) (a + b)];\n }\n};\n\nstatic void TestInclusiveScanWithConstAccumulator(void)\n{\n // add numbers modulo 3 with external lookup table\n thrust::device_vector data(7);\n data[0] = 0;\n data[1] = 1;\n data[2] = 2;\n data[3] = 1;\n data[4] = 2;\n data[5] = 0;\n data[6] = 1;\n\n thrust::device_vector table(6);\n table[0] = 0;\n table[1] = 1;\n table[2] = 2;\n table[3] = 0;\n table[4] = 1;\n table[5] = 2;\n\n thrust::inclusive_scan(data.begin(), data.end(), data.begin(), const_ref_plus_mod3(thrust::raw_pointer_cast(&table[0])));\n \n ASSERT_EQUAL(data[0], 0);\n ASSERT_EQUAL(data[1], 1);\n ASSERT_EQUAL(data[2], 0);\n ASSERT_EQUAL(data[3], 1);\n ASSERT_EQUAL(data[4], 0);\n ASSERT_EQUAL(data[5], 0);\n ASSERT_EQUAL(data[6], 1);\n}\nDECLARE_UNITTEST(TestInclusiveScanWithConstAccumulator);\n", "source": "/kaggle/working/thrust/testing/cuda/scan.cu"} {"input": "#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid stable_sort_kernel(ExecutionPolicy exec, Iterator first, Iterator last)\n{\n thrust::stable_sort(exec, first, last);\n}\n\n\nstruct make_pair_functor\n{\n template\n __host__ __device__\n thrust::pair operator()(const T1 &x, const T2 &y)\n {\n return thrust::make_pair(x,y);\n } // end operator()()\n}; // end make_pair_functor\n\n\ntemplate\nvoid TestPairStableSortDevice(ExecutionPolicy exec)\n{\n size_t n = 10000;\n typedef thrust::pair P;\n\n thrust::host_vector h_p1 = unittest::random_integers(n);\n thrust::host_vector h_p2 = unittest::random_integers(n);\n thrust::host_vector

h_pairs(n);\n\n // zip up pairs on the host\n thrust::transform(h_p1.begin(), h_p1.end(), h_p2.begin(), h_pairs.begin(), make_pair_functor());\n\n thrust::device_vector

d_pairs = h_pairs;\n\n stable_sort_kernel<<<1,1>>>(exec, d_pairs.begin(), d_pairs.end());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n // sort on the host\n thrust::stable_sort(h_pairs.begin(), h_pairs.end());\n\n ASSERT_EQUAL_QUIET(h_pairs, d_pairs);\n};\n\n\nvoid TestPairStableSortDeviceSeq()\n{\n TestPairStableSortDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestPairStableSortDeviceSeq);\n\n\nvoid TestPairStableSortDeviceDevice()\n{\n TestPairStableSortDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestPairStableSortDeviceDevice);\n#endif\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid stable_sort_kernel(ExecutionPolicy exec, Iterator first, Iterator last)\n{\n thrust::stable_sort(exec, first, last);\n}\n\n\nstruct make_pair_functor\n{\n template\n __host__ __device__\n thrust::pair operator()(const T1 &x, const T2 &y)\n {\n return thrust::make_pair(x,y);\n } // end operator()()\n}; // end make_pair_functor\n\n\ntemplate\nvoid TestPairStableSortDevice(ExecutionPolicy exec)\n{\n size_t n = 10000;\n typedef thrust::pair P;\n\n thrust::host_vector h_p1 = unittest::random_integers(n);\n thrust::host_vector h_p2 = unittest::random_integers(n);\n thrust::host_vector

h_pairs(n);\n\n // zip up pairs on the host\n thrust::transform(h_p1.begin(), h_p1.end(), h_p2.begin(), h_pairs.begin(), make_pair_functor());\n\n thrust::device_vector

d_pairs = h_pairs;\n\n stable_sort_kernel<<<1,1>>>(exec, d_pairs.begin(), d_pairs.end());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n // sort on the host\n thrust::stable_sort(h_pairs.begin(), h_pairs.end());\n\n ASSERT_EQUAL_QUIET(h_pairs, d_pairs);\n};\n\n\nvoid TestPairStableSortDeviceSeq()\n{\n TestPairStableSortDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestPairStableSortDeviceSeq);\n\n\nvoid TestPairStableSortDeviceDevice()\n{\n TestPairStableSortDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestPairStableSortDeviceDevice);\n#endif\n\n", "source": "/kaggle/working/thrust/testing/cuda/pair_sort.cu"} {"input": "#include \n#include \n#include \n#include \n\n\ntemplate\nstruct my_less\n{\n __host__ __device__\n bool operator()(const T& lhs, const T& rhs) const\n {\n return lhs < rhs;\n }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid sort_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Compare comp)\n{\n thrust::sort(exec, first, last, comp);\n}\n\n\ntemplate\nvoid TestComparisonSortDevice(ExecutionPolicy exec, const size_t n, Compare comp)\n{\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n sort_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), comp);\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n\n thrust::sort(h_data.begin(), h_data.end(), comp);\n\n ASSERT_EQUAL(h_data, d_data);\n};\n\n\ntemplate\n struct TestComparisonSortDeviceSeq\n{\n void operator()(const size_t n)\n {\n TestComparisonSortDevice(thrust::seq, n, my_less());\n }\n};\nVariableUnitTest<\n TestComparisonSortDeviceSeq,\n unittest::type_list\n> TestComparisonSortDeviceSeqInstance;\n\n\ntemplate\n struct TestComparisonSortDeviceDevice\n{\n void operator()(const size_t n)\n {\n TestComparisonSortDevice(thrust::device, n, my_less());\n }\n};\nVariableUnitTest<\n TestComparisonSortDeviceDevice,\n unittest::type_list\n> TestComparisonSortDeviceDeviceDeviceInstance;\n\n\ntemplate\nvoid TestSortDevice(ExecutionPolicy exec, const size_t n)\n{\n TestComparisonSortDevice(exec, n, thrust::less());\n};\n\n\ntemplate\n struct TestSortDeviceSeq\n{\n void operator()(const size_t n)\n {\n TestSortDevice(thrust::seq, n);\n }\n};\nVariableUnitTest<\n TestSortDeviceSeq,\n unittest::type_list\n> TestSortDeviceSeqInstance;\n\n\ntemplate\n struct TestSortDeviceDevice\n{\n void operator()(const size_t n)\n {\n TestSortDevice(thrust::device, n);\n }\n};\nVariableUnitTest<\n TestSortDeviceDevice,\n unittest::type_list\n> TestSortDeviceDeviceInstance;\n#endif\n\n\nvoid TestSortCudaStreams()\n{\n thrust::device_vector keys(10);\n\n keys[0] = 9;\n keys[1] = 3;\n keys[2] = 2;\n keys[3] = 0;\n keys[4] = 4;\n keys[5] = 7;\n keys[6] = 8;\n keys[7] = 1;\n keys[8] = 5;\n keys[9] = 6;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::sort(thrust::cuda::par.on(s), keys.begin(), keys.end());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(true, thrust::is_sorted(keys.begin(), keys.end()));\n \n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSortCudaStreams);\n\n\nvoid TestComparisonSortCudaStreams()\n{\n thrust::device_vector keys(10);\n\n keys[0] = 9;\n keys[1] = 3;\n keys[2] = 2;\n keys[3] = 0;\n keys[4] = 4;\n keys[5] = 7;\n keys[6] = 8;\n keys[7] = 1;\n keys[8] = 5;\n keys[9] = 6;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::sort(thrust::cuda::par.on(s), keys.begin(), keys.end(), my_less());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(true, thrust::is_sorted(keys.begin(), keys.end(), my_less()));\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestComparisonSortCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n\ntemplate\nstruct my_less\n{\n __host__ __device__\n bool operator()(const T& lhs, const T& rhs) const\n {\n return lhs < rhs;\n }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid sort_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Compare comp)\n{\n thrust::sort(exec, first, last, comp);\n}\n\n\ntemplate\nvoid TestComparisonSortDevice(ExecutionPolicy exec, const size_t n, Compare comp)\n{\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n sort_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), comp);\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n\n thrust::sort(h_data.begin(), h_data.end(), comp);\n\n ASSERT_EQUAL(h_data, d_data);\n};\n\n\ntemplate\n struct TestComparisonSortDeviceSeq\n{\n void operator()(const size_t n)\n {\n TestComparisonSortDevice(thrust::seq, n, my_less());\n }\n};\nVariableUnitTest<\n TestComparisonSortDeviceSeq,\n unittest::type_list\n> TestComparisonSortDeviceSeqInstance;\n\n\ntemplate\n struct TestComparisonSortDeviceDevice\n{\n void operator()(const size_t n)\n {\n TestComparisonSortDevice(thrust::device, n, my_less());\n }\n};\nVariableUnitTest<\n TestComparisonSortDeviceDevice,\n unittest::type_list\n> TestComparisonSortDeviceDeviceDeviceInstance;\n\n\ntemplate\nvoid TestSortDevice(ExecutionPolicy exec, const size_t n)\n{\n TestComparisonSortDevice(exec, n, thrust::less());\n};\n\n\ntemplate\n struct TestSortDeviceSeq\n{\n void operator()(const size_t n)\n {\n TestSortDevice(thrust::seq, n);\n }\n};\nVariableUnitTest<\n TestSortDeviceSeq,\n unittest::type_list\n> TestSortDeviceSeqInstance;\n\n\ntemplate\n struct TestSortDeviceDevice\n{\n void operator()(const size_t n)\n {\n TestSortDevice(thrust::device, n);\n }\n};\nVariableUnitTest<\n TestSortDeviceDevice,\n unittest::type_list\n> TestSortDeviceDeviceInstance;\n#endif\n\n\nvoid TestSortCudaStreams()\n{\n thrust::device_vector keys(10);\n\n keys[0] = 9;\n keys[1] = 3;\n keys[2] = 2;\n keys[3] = 0;\n keys[4] = 4;\n keys[5] = 7;\n keys[6] = 8;\n keys[7] = 1;\n keys[8] = 5;\n keys[9] = 6;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::sort(thrust::cuda::par.on(s), keys.begin(), keys.end());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(true, thrust::is_sorted(keys.begin(), keys.end()));\n \n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSortCudaStreams);\n\n\nvoid TestComparisonSortCudaStreams()\n{\n thrust::device_vector keys(10);\n\n keys[0] = 9;\n keys[1] = 3;\n keys[2] = 2;\n keys[3] = 0;\n keys[4] = 4;\n keys[5] = 7;\n keys[6] = 8;\n keys[7] = 1;\n keys[8] = 5;\n keys[9] = 6;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::sort(thrust::cuda::par.on(s), keys.begin(), keys.end(), my_less());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(true, thrust::is_sorted(keys.begin(), keys.end(), my_less()));\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestComparisonSortCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/sort.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid reduce_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first, Iterator1 keys_last,\n Iterator2 values_first,\n Iterator3 keys_result,\n Iterator4 values_result,\n Iterator5 result)\n{\n *result = thrust::reduce_by_key(exec, keys_first, keys_last, values_first, keys_result, values_result);\n}\n\n\ntemplate\n__global__\nvoid reduce_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first, Iterator1 keys_last,\n Iterator2 values_first,\n Iterator3 keys_result,\n Iterator4 values_result,\n BinaryPredicate pred,\n Iterator5 result)\n{\n *result = thrust::reduce_by_key(exec, keys_first, keys_last, values_first, keys_result, values_result, pred);\n}\n\n\ntemplate\n__global__\nvoid reduce_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first, Iterator1 keys_last,\n Iterator2 values_first,\n Iterator3 keys_result,\n Iterator4 values_result,\n BinaryPredicate pred,\n BinaryFunction binary_op,\n Iterator5 result)\n{\n *result = thrust::reduce_by_key(exec, keys_first, keys_last, values_first, keys_result, values_result, pred, binary_op);\n}\n#endif\n\n\ntemplate\nstruct is_equal_div_10_reduce\n{\n __host__ __device__\n bool operator()(const T x, const T& y) const { return ((int) x / 10) == ((int) y / 10); }\n};\n\n\ntemplate\nvoid initialize_keys(Vector& keys)\n{\n keys.resize(9);\n keys[0] = 11;\n keys[1] = 11;\n keys[2] = 21;\n keys[3] = 20;\n keys[4] = 21;\n keys[5] = 21;\n keys[6] = 21;\n keys[7] = 37;\n keys[8] = 37;\n}\n\n\ntemplate\nvoid initialize_values(Vector& values)\n{\n values.resize(9);\n values[0] = 0; \n values[1] = 1;\n values[2] = 2;\n values[3] = 3;\n values[4] = 4;\n values[5] = 5;\n values[6] = 6;\n values[7] = 7;\n values[8] = 8;\n}\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\nvoid TestReduceByKeyDevice(ExecutionPolicy exec)\n{\n typedef int T;\n \n thrust::device_vector keys;\n thrust::device_vector values;\n\n typedef typename thrust::pair<\n typename thrust::device_vector::iterator,\n typename thrust::device_vector::iterator\n > iterator_pair;\n\n thrust::device_vector new_last_vec(1);\n iterator_pair new_last;\n \n // basic test\n initialize_keys(keys); initialize_values(values);\n \n thrust::device_vector output_keys(keys.size());\n thrust::device_vector output_values(values.size());\n \n reduce_by_key_kernel<<<1,1>>>(exec, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), new_last_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n new_last = new_last_vec[0];\n \n ASSERT_EQUAL(new_last.first - output_keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 5);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 20);\n ASSERT_EQUAL(output_keys[3], 21);\n ASSERT_EQUAL(output_keys[4], 37);\n \n ASSERT_EQUAL(output_values[0], 1);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 3);\n ASSERT_EQUAL(output_values[3], 15);\n ASSERT_EQUAL(output_values[4], 15);\n \n // test BinaryPredicate\n initialize_keys(keys); initialize_values(values);\n \n reduce_by_key_kernel<<<1,1>>>(exec, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), is_equal_div_10_reduce(), new_last_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n new_last = new_last_vec[0];\n \n ASSERT_EQUAL(new_last.first - output_keys.begin(), 3);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 3);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 37);\n \n ASSERT_EQUAL(output_values[0], 1);\n ASSERT_EQUAL(output_values[1], 20);\n ASSERT_EQUAL(output_values[2], 15);\n \n // test BinaryFunction\n initialize_keys(keys); initialize_values(values);\n \n reduce_by_key_kernel<<<1,1>>>(exec, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), thrust::equal_to(), thrust::plus(), new_last_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n new_last = new_last_vec[0];\n \n ASSERT_EQUAL(new_last.first - output_keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 5);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 20);\n ASSERT_EQUAL(output_keys[3], 21);\n ASSERT_EQUAL(output_keys[4], 37);\n \n ASSERT_EQUAL(output_values[0], 1);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 3);\n ASSERT_EQUAL(output_values[3], 15);\n ASSERT_EQUAL(output_values[4], 15);\n}\n\n\nvoid TestReduceByKeyDeviceSeq()\n{\n TestReduceByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReduceByKeyDeviceSeq);\n\n\nvoid TestReduceByKeyDeviceDevice()\n{\n TestReduceByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReduceByKeyDeviceDevice);\n\n\nvoid TestReduceByKeyDeviceNoSync()\n{\n TestReduceByKeyDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestReduceByKeyDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestReduceByKeyCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector keys;\n Vector values;\n\n thrust::pair new_last;\n\n // basic test\n initialize_keys(keys); initialize_values(values);\n\n Vector output_keys(keys.size());\n Vector output_values(values.size());\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n\n new_last = thrust::reduce_by_key(streampolicy, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin());\n\n ASSERT_EQUAL(new_last.first - output_keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 5);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 20);\n ASSERT_EQUAL(output_keys[3], 21);\n ASSERT_EQUAL(output_keys[4], 37);\n \n ASSERT_EQUAL(output_values[0], 1);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 3);\n ASSERT_EQUAL(output_values[3], 15);\n ASSERT_EQUAL(output_values[4], 15);\n\n // test BinaryPredicate\n initialize_keys(keys); initialize_values(values);\n \n new_last = thrust::reduce_by_key(streampolicy, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), is_equal_div_10_reduce());\n\n ASSERT_EQUAL(new_last.first - output_keys.begin(), 3);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 3);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 37);\n \n ASSERT_EQUAL(output_values[0], 1);\n ASSERT_EQUAL(output_values[1], 20);\n ASSERT_EQUAL(output_values[2], 15);\n\n // test BinaryFunction\n initialize_keys(keys); initialize_values(values);\n\n new_last = thrust::reduce_by_key(streampolicy, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), thrust::equal_to(), thrust::plus());\n\n ASSERT_EQUAL(new_last.first - output_keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 5);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 20);\n ASSERT_EQUAL(output_keys[3], 21);\n ASSERT_EQUAL(output_keys[4], 37);\n \n ASSERT_EQUAL(output_values[0], 1);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 3);\n ASSERT_EQUAL(output_values[3], 15);\n ASSERT_EQUAL(output_values[4], 15);\n\n cudaStreamDestroy(s);\n}\n\nvoid TestReduceByKeyCudaStreamsSync()\n{\n TestReduceByKeyCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestReduceByKeyCudaStreamsSync);\n\n\nvoid TestReduceByKeyCudaStreamsNoSync()\n{\n TestReduceByKeyCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestReduceByKeyCudaStreamsNoSync);\n\n\n// Maps indices to key ids\nclass div_op : public thrust::unary_function\n{\n std::int64_t m_divisor;\n\npublic:\n __host__ div_op(std::int64_t divisor)\n : m_divisor(divisor)\n {}\n\n __host__ __device__\n std::int64_t operator()(std::int64_t x) const\n {\n return x / m_divisor;\n }\n};\n\n// Produces unique sequence for key\nclass mod_op : public thrust::unary_function\n{\n std::int64_t m_divisor;\n\npublic:\n __host__ mod_op(std::int64_t divisor)\n : m_divisor(divisor)\n {}\n\n __host__ __device__\n std::int64_t operator()(std::int64_t x) const\n {\n // div: 2 \n // idx: 0 1 2 3 4 5 \n // key: 0 0 | 1 1 | 2 2 \n // mod: 0 1 | 0 1 | 0 1\n // ret: 0 1 1 2 2 3\n return (x % m_divisor) + (x / m_divisor);\n }\n};\n\n\nvoid TestReduceByKeyWithBigIndexesHelper(int magnitude)\n{\n const std::int64_t key_size_magnitude = 8;\n ASSERT_EQUAL(true, key_size_magnitude < magnitude);\n\n const std::int64_t num_items = 1ll << magnitude;\n const std::int64_t num_unique_keys = 1ll << key_size_magnitude;\n\n // Size of each key group\n const std::int64_t key_size = num_items / num_unique_keys;\n\n using counting_it = thrust::counting_iterator;\n using transform_key_it = thrust::transform_iterator;\n using transform_val_it = thrust::transform_iterator;\n\n counting_it count_begin(0ll);\n counting_it count_end = count_begin + num_items;\n ASSERT_EQUAL(static_cast(thrust::distance(count_begin, count_end)),\n num_items);\n\n transform_key_it keys_begin(count_begin, div_op{key_size});\n transform_key_it keys_end(count_end, div_op{key_size});\n\n transform_val_it values_begin(count_begin, mod_op{key_size});\n\n thrust::device_vector output_keys(num_unique_keys);\n thrust::device_vector output_values(num_unique_keys);\n\n // example:\n // items: 6\n // unique_keys: 2\n // key_size: 3\n // keys: 0 0 0 | 1 1 1 \n // values: 0 1 2 | 1 2 3\n // result: 3 6 = sum(range(key_size)) + key_size * key_id\n thrust::reduce_by_key(keys_begin,\n keys_end,\n values_begin,\n output_keys.begin(),\n output_values.begin());\n\n ASSERT_EQUAL(\n true,\n thrust::equal(output_keys.begin(), output_keys.end(), count_begin));\n\n thrust::host_vector result = output_values;\n\n const std::int64_t sum = (key_size - 1) * key_size / 2;\n for (std::int64_t key_id = 0; key_id < num_unique_keys; key_id++)\n {\n ASSERT_EQUAL(result[key_id], sum + key_id * key_size);\n }\n}\n\nvoid TestReduceByKeyWithBigIndexes()\n{\n TestReduceByKeyWithBigIndexesHelper(30);\n TestReduceByKeyWithBigIndexesHelper(31);\n TestReduceByKeyWithBigIndexesHelper(32);\n TestReduceByKeyWithBigIndexesHelper(33);\n}\nDECLARE_UNITTEST(TestReduceByKeyWithBigIndexes);\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid reduce_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first, Iterator1 keys_last,\n Iterator2 values_first,\n Iterator3 keys_result,\n Iterator4 values_result,\n Iterator5 result)\n{\n *result = thrust::reduce_by_key(exec, keys_first, keys_last, values_first, keys_result, values_result);\n}\n\n\ntemplate\n__global__\nvoid reduce_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first, Iterator1 keys_last,\n Iterator2 values_first,\n Iterator3 keys_result,\n Iterator4 values_result,\n BinaryPredicate pred,\n Iterator5 result)\n{\n *result = thrust::reduce_by_key(exec, keys_first, keys_last, values_first, keys_result, values_result, pred);\n}\n\n\ntemplate\n__global__\nvoid reduce_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first, Iterator1 keys_last,\n Iterator2 values_first,\n Iterator3 keys_result,\n Iterator4 values_result,\n BinaryPredicate pred,\n BinaryFunction binary_op,\n Iterator5 result)\n{\n *result = thrust::reduce_by_key(exec, keys_first, keys_last, values_first, keys_result, values_result, pred, binary_op);\n}\n#endif\n\n\ntemplate\nstruct is_equal_div_10_reduce\n{\n __host__ __device__\n bool operator()(const T x, const T& y) const { return ((int) x / 10) == ((int) y / 10); }\n};\n\n\ntemplate\nvoid initialize_keys(Vector& keys)\n{\n keys.resize(9);\n keys[0] = 11;\n keys[1] = 11;\n keys[2] = 21;\n keys[3] = 20;\n keys[4] = 21;\n keys[5] = 21;\n keys[6] = 21;\n keys[7] = 37;\n keys[8] = 37;\n}\n\n\ntemplate\nvoid initialize_values(Vector& values)\n{\n values.resize(9);\n values[0] = 0; \n values[1] = 1;\n values[2] = 2;\n values[3] = 3;\n values[4] = 4;\n values[5] = 5;\n values[6] = 6;\n values[7] = 7;\n values[8] = 8;\n}\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\nvoid TestReduceByKeyDevice(ExecutionPolicy exec)\n{\n typedef int T;\n \n thrust::device_vector keys;\n thrust::device_vector values;\n\n typedef typename thrust::pair<\n typename thrust::device_vector::iterator,\n typename thrust::device_vector::iterator\n > iterator_pair;\n\n thrust::device_vector new_last_vec(1);\n iterator_pair new_last;\n \n // basic test\n initialize_keys(keys); initialize_values(values);\n \n thrust::device_vector output_keys(keys.size());\n thrust::device_vector output_values(values.size());\n \n reduce_by_key_kernel<<<1,1>>>(exec, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), new_last_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n new_last = new_last_vec[0];\n \n ASSERT_EQUAL(new_last.first - output_keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 5);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 20);\n ASSERT_EQUAL(output_keys[3], 21);\n ASSERT_EQUAL(output_keys[4], 37);\n \n ASSERT_EQUAL(output_values[0], 1);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 3);\n ASSERT_EQUAL(output_values[3], 15);\n ASSERT_EQUAL(output_values[4], 15);\n \n // test BinaryPredicate\n initialize_keys(keys); initialize_values(values);\n \n reduce_by_key_kernel<<<1,1>>>(exec, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), is_equal_div_10_reduce(), new_last_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n new_last = new_last_vec[0];\n \n ASSERT_EQUAL(new_last.first - output_keys.begin(), 3);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 3);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 37);\n \n ASSERT_EQUAL(output_values[0], 1);\n ASSERT_EQUAL(output_values[1], 20);\n ASSERT_EQUAL(output_values[2], 15);\n \n // test BinaryFunction\n initialize_keys(keys); initialize_values(values);\n \n reduce_by_key_kernel<<<1,1>>>(exec, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), thrust::equal_to(), thrust::plus(), new_last_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n new_last = new_last_vec[0];\n \n ASSERT_EQUAL(new_last.first - output_keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 5);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 20);\n ASSERT_EQUAL(output_keys[3], 21);\n ASSERT_EQUAL(output_keys[4], 37);\n \n ASSERT_EQUAL(output_values[0], 1);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 3);\n ASSERT_EQUAL(output_values[3], 15);\n ASSERT_EQUAL(output_values[4], 15);\n}\n\n\nvoid TestReduceByKeyDeviceSeq()\n{\n TestReduceByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReduceByKeyDeviceSeq);\n\n\nvoid TestReduceByKeyDeviceDevice()\n{\n TestReduceByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReduceByKeyDeviceDevice);\n\n\nvoid TestReduceByKeyDeviceNoSync()\n{\n TestReduceByKeyDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestReduceByKeyDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestReduceByKeyCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector keys;\n Vector values;\n\n thrust::pair new_last;\n\n // basic test\n initialize_keys(keys); initialize_values(values);\n\n Vector output_keys(keys.size());\n Vector output_values(values.size());\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n\n new_last = thrust::reduce_by_key(streampolicy, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin());\n\n ASSERT_EQUAL(new_last.first - output_keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 5);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 20);\n ASSERT_EQUAL(output_keys[3], 21);\n ASSERT_EQUAL(output_keys[4], 37);\n \n ASSERT_EQUAL(output_values[0], 1);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 3);\n ASSERT_EQUAL(output_values[3], 15);\n ASSERT_EQUAL(output_values[4], 15);\n\n // test BinaryPredicate\n initialize_keys(keys); initialize_values(values);\n \n new_last = thrust::reduce_by_key(streampolicy, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), is_equal_div_10_reduce());\n\n ASSERT_EQUAL(new_last.first - output_keys.begin(), 3);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 3);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 37);\n \n ASSERT_EQUAL(output_values[0], 1);\n ASSERT_EQUAL(output_values[1], 20);\n ASSERT_EQUAL(output_values[2], 15);\n\n // test BinaryFunction\n initialize_keys(keys); initialize_values(values);\n\n new_last = thrust::reduce_by_key(streampolicy, keys.begin(), keys.end(), values.begin(), output_keys.begin(), output_values.begin(), thrust::equal_to(), thrust::plus());\n\n ASSERT_EQUAL(new_last.first - output_keys.begin(), 5);\n ASSERT_EQUAL(new_last.second - output_values.begin(), 5);\n ASSERT_EQUAL(output_keys[0], 11);\n ASSERT_EQUAL(output_keys[1], 21);\n ASSERT_EQUAL(output_keys[2], 20);\n ASSERT_EQUAL(output_keys[3], 21);\n ASSERT_EQUAL(output_keys[4], 37);\n \n ASSERT_EQUAL(output_values[0], 1);\n ASSERT_EQUAL(output_values[1], 2);\n ASSERT_EQUAL(output_values[2], 3);\n ASSERT_EQUAL(output_values[3], 15);\n ASSERT_EQUAL(output_values[4], 15);\n\n hipStreamDestroy(s);\n}\n\nvoid TestReduceByKeyCudaStreamsSync()\n{\n TestReduceByKeyCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestReduceByKeyCudaStreamsSync);\n\n\nvoid TestReduceByKeyCudaStreamsNoSync()\n{\n TestReduceByKeyCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestReduceByKeyCudaStreamsNoSync);\n\n\n// Maps indices to key ids\nclass div_op : public thrust::unary_function\n{\n std::int64_t m_divisor;\n\npublic:\n __host__ div_op(std::int64_t divisor)\n : m_divisor(divisor)\n {}\n\n __host__ __device__\n std::int64_t operator()(std::int64_t x) const\n {\n return x / m_divisor;\n }\n};\n\n// Produces unique sequence for key\nclass mod_op : public thrust::unary_function\n{\n std::int64_t m_divisor;\n\npublic:\n __host__ mod_op(std::int64_t divisor)\n : m_divisor(divisor)\n {}\n\n __host__ __device__\n std::int64_t operator()(std::int64_t x) const\n {\n // div: 2 \n // idx: 0 1 2 3 4 5 \n // key: 0 0 | 1 1 | 2 2 \n // mod: 0 1 | 0 1 | 0 1\n // ret: 0 1 1 2 2 3\n return (x % m_divisor) + (x / m_divisor);\n }\n};\n\n\nvoid TestReduceByKeyWithBigIndexesHelper(int magnitude)\n{\n const std::int64_t key_size_magnitude = 8;\n ASSERT_EQUAL(true, key_size_magnitude < magnitude);\n\n const std::int64_t num_items = 1ll << magnitude;\n const std::int64_t num_unique_keys = 1ll << key_size_magnitude;\n\n // Size of each key group\n const std::int64_t key_size = num_items / num_unique_keys;\n\n using counting_it = thrust::counting_iterator;\n using transform_key_it = thrust::transform_iterator;\n using transform_val_it = thrust::transform_iterator;\n\n counting_it count_begin(0ll);\n counting_it count_end = count_begin + num_items;\n ASSERT_EQUAL(static_cast(thrust::distance(count_begin, count_end)),\n num_items);\n\n transform_key_it keys_begin(count_begin, div_op{key_size});\n transform_key_it keys_end(count_end, div_op{key_size});\n\n transform_val_it values_begin(count_begin, mod_op{key_size});\n\n thrust::device_vector output_keys(num_unique_keys);\n thrust::device_vector output_values(num_unique_keys);\n\n // example:\n // items: 6\n // unique_keys: 2\n // key_size: 3\n // keys: 0 0 0 | 1 1 1 \n // values: 0 1 2 | 1 2 3\n // result: 3 6 = sum(range(key_size)) + key_size * key_id\n thrust::reduce_by_key(keys_begin,\n keys_end,\n values_begin,\n output_keys.begin(),\n output_values.begin());\n\n ASSERT_EQUAL(\n true,\n thrust::equal(output_keys.begin(), output_keys.end(), count_begin));\n\n thrust::host_vector result = output_values;\n\n const std::int64_t sum = (key_size - 1) * key_size / 2;\n for (std::int64_t key_id = 0; key_id < num_unique_keys; key_id++)\n {\n ASSERT_EQUAL(result[key_id], sum + key_id * key_size);\n }\n}\n\nvoid TestReduceByKeyWithBigIndexes()\n{\n TestReduceByKeyWithBigIndexesHelper(30);\n TestReduceByKeyWithBigIndexesHelper(31);\n TestReduceByKeyWithBigIndexesHelper(32);\n TestReduceByKeyWithBigIndexesHelper(33);\n}\nDECLARE_UNITTEST(TestReduceByKeyWithBigIndexes);\n", "source": "/kaggle/working/thrust/testing/cuda/reduce_by_key.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid remove_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T val, Iterator2 result)\n{\n *result = thrust::remove(exec, first, last, val);\n}\n\n\ntemplate\n__global__\nvoid remove_if_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Predicate pred, Iterator2 result)\n{\n *result = thrust::remove_if(exec, first, last, pred);\n}\n\n\ntemplate\n__global__\nvoid remove_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Predicate pred, Iterator3 result)\n{\n *result = thrust::remove_if(exec, first, last, stencil_first, pred);\n}\n\n\ntemplate\n__global__\nvoid remove_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, T val, Iterator3 result2)\n{\n *result2 = thrust::remove_copy(exec, first, last, result1, val);\n}\n\n\ntemplate\n__global__\nvoid remove_copy_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result, Predicate pred, Iterator3 result_end)\n{\n *result_end = thrust::remove_copy_if(exec, first, last, result, pred);\n}\n\n\ntemplate\n__global__\nvoid remove_copy_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Iterator3 result, Predicate pred, Iterator4 result_end)\n{\n *result_end = thrust::remove_copy_if(exec, first, last, stencil_first, result, pred);\n}\n#endif\n\n\ntemplate\nstruct is_even\n : thrust::unary_function\n{\n __host__ __device__\n bool operator()(T x) { return (static_cast(x) & 1) == 0; }\n};\n\n\ntemplate\nstruct is_true\n : thrust::unary_function\n{\n __host__ __device__\n bool operator()(T x) { return x ? true : false; }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\nvoid TestRemoveDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n\n typedef typename thrust::device_vector::iterator iterator;\n thrust::device_vector d_result(1);\n \n size_t h_size = thrust::remove(h_data.begin(), h_data.end(), 0) - h_data.begin();\n\n remove_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), 0, d_result.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n size_t d_size = (iterator)d_result[0] - d_data.begin();\n \n ASSERT_EQUAL(h_size, d_size);\n \n h_data.resize(h_size);\n d_data.resize(d_size);\n \n ASSERT_EQUAL(h_data, d_data);\n}\n\n\nvoid TestRemoveDeviceSeq()\n{\n TestRemoveDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestRemoveDeviceSeq);\n\n\nvoid TestRemoveDeviceDevice()\n{\n TestRemoveDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestRemoveDeviceDevice);\n\n\ntemplate\nvoid TestRemoveIfDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n\n typedef typename thrust::device_vector::iterator iterator;\n thrust::device_vector d_result(1);\n \n size_t h_size = thrust::remove_if(h_data.begin(), h_data.end(), is_true()) - h_data.begin();\n\n remove_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), is_true(), d_result.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n size_t d_size = (iterator)d_result[0] - d_data.begin();\n \n ASSERT_EQUAL(h_size, d_size);\n \n h_data.resize(h_size);\n d_data.resize(d_size);\n \n ASSERT_EQUAL(h_data, d_data);\n}\n\n\nvoid TestRemoveIfDeviceSeq()\n{\n TestRemoveIfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestRemoveIfDeviceSeq);\n\n\nvoid TestRemoveIfDeviceDevice()\n{\n TestRemoveIfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestRemoveIfDeviceDevice);\n\n\ntemplate\nvoid TestRemoveIfStencilDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n\n typedef typename thrust::device_vector::iterator iterator;\n thrust::device_vector d_result(1);\n \n thrust::host_vector h_stencil = unittest::random_integers(n);\n thrust::device_vector d_stencil = h_stencil;\n \n size_t h_size = thrust::remove_if(h_data.begin(), h_data.end(), h_stencil.begin(), is_true()) - h_data.begin();\n\n remove_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_stencil.begin(), is_true(), d_result.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n size_t d_size = (iterator)d_result[0] - d_data.begin();\n \n ASSERT_EQUAL(h_size, d_size);\n \n h_data.resize(h_size);\n d_data.resize(d_size);\n \n ASSERT_EQUAL(h_data, d_data);\n}\n\n\nvoid TestRemoveIfStencilDeviceSeq()\n{\n TestRemoveIfStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestRemoveIfStencilDeviceSeq);\n\n\nvoid TestRemoveIfStencilDeviceDevice()\n{\n TestRemoveIfStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestRemoveIfStencilDeviceDevice);\n\n\ntemplate\nvoid TestRemoveCopyDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n\n typedef typename thrust::device_vector::iterator iterator;\n thrust::device_vector d_new_end(1);\n \n size_t h_size = thrust::remove_copy(h_data.begin(), h_data.end(), h_result.begin(), 0) - h_result.begin();\n\n remove_copy_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin(), 0, d_new_end.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n size_t d_size = (iterator)d_new_end[0] - d_result.begin();\n \n ASSERT_EQUAL(h_size, d_size);\n \n h_result.resize(h_size);\n d_result.resize(d_size);\n \n ASSERT_EQUAL(h_result, d_result);\n}\n\n\nvoid TestRemoveCopyDeviceSeq()\n{\n TestRemoveCopyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestRemoveCopyDeviceSeq);\n\n\nvoid TestRemoveCopyDeviceDevice()\n{\n TestRemoveCopyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestRemoveCopyDeviceDevice);\n\n\ntemplate\nvoid TestRemoveCopyIfDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n\n typedef typename thrust::device_vector::iterator iterator;\n thrust::device_vector d_new_end(1);\n \n size_t h_size = thrust::remove_copy_if(h_data.begin(), h_data.end(), h_result.begin(), is_true()) - h_result.begin();\n\n remove_copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin(), is_true(), d_new_end.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n size_t d_size = (iterator)d_new_end[0] - d_result.begin();\n \n ASSERT_EQUAL(h_size, d_size);\n \n h_result.resize(h_size);\n d_result.resize(d_size);\n \n ASSERT_EQUAL(h_result, d_result);\n}\n\n\nvoid TestRemoveCopyIfDeviceSeq()\n{\n TestRemoveCopyIfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestRemoveCopyIfDeviceSeq);\n\n\nvoid TestRemoveCopyIfDeviceDevice()\n{\n TestRemoveCopyIfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestRemoveCopyIfDeviceDevice);\n\n\ntemplate\nvoid TestRemoveCopyIfStencilDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n\n typedef typename thrust::device_vector::iterator iterator;\n thrust::device_vector d_new_end(1);\n\n thrust::host_vector h_stencil = unittest::random_integers(n);\n thrust::device_vector d_stencil = h_stencil;\n \n size_t h_size = thrust::remove_copy_if(h_data.begin(), h_data.end(), h_stencil.begin(), h_result.begin(), is_true()) - h_result.begin();\n\n remove_copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_stencil.begin(), d_result.begin(), is_true(), d_new_end.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n size_t d_size = (iterator)d_new_end[0] - d_result.begin();\n \n ASSERT_EQUAL(h_size, d_size);\n \n h_result.resize(h_size);\n d_result.resize(d_size);\n \n ASSERT_EQUAL(h_result, d_result);\n}\n\n\nvoid TestRemoveCopyIfStencilDeviceSeq()\n{\n TestRemoveCopyIfStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestRemoveCopyIfStencilDeviceSeq);\n\n\nvoid TestRemoveCopyIfStencilDeviceDevice()\n{\n TestRemoveCopyIfStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestRemoveCopyIfStencilDeviceDevice);\n#endif\n\n\nvoid TestRemoveCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n Vector::iterator end = thrust::remove(thrust::cuda::par.on(s),\n data.begin(), \n data.end(), \n (T) 2);\n\n ASSERT_EQUAL(end - data.begin(), 3);\n\n ASSERT_EQUAL(data[0], 1);\n ASSERT_EQUAL(data[1], 1);\n ASSERT_EQUAL(data[2], 3);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestRemoveCudaStreams);\n\n\nvoid TestRemoveCopyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n Vector result(5);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n Vector::iterator end = thrust::remove_copy(thrust::cuda::par.on(s),\n data.begin(), \n data.end(), \n result.begin(), \n (T) 2);\n\n ASSERT_EQUAL(end - result.begin(), 3);\n\n ASSERT_EQUAL(result[0], 1);\n ASSERT_EQUAL(result[1], 1);\n ASSERT_EQUAL(result[2], 3);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestRemoveCopyCudaStreams);\n\n\nvoid TestRemoveIfCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n Vector::iterator end = thrust::remove_if(thrust::cuda::par.on(s),\n data.begin(), \n data.end(), \n is_even());\n\n ASSERT_EQUAL(end - data.begin(), 3);\n\n ASSERT_EQUAL(data[0], 1);\n ASSERT_EQUAL(data[1], 1);\n ASSERT_EQUAL(data[2], 3);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestRemoveIfCudaStreams);\n\n\nvoid TestRemoveIfStencilCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n Vector stencil(5);\n stencil[0] = 0;\n stencil[1] = 1;\n stencil[2] = 0;\n stencil[3] = 0;\n stencil[4] = 1;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n Vector::iterator end = thrust::remove_if(thrust::cuda::par.on(s),\n data.begin(), \n data.end(),\n stencil.begin(),\n thrust::identity());\n\n ASSERT_EQUAL(end - data.begin(), 3);\n\n ASSERT_EQUAL(data[0], 1);\n ASSERT_EQUAL(data[1], 1);\n ASSERT_EQUAL(data[2], 3);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestRemoveIfStencilCudaStreams);\n\n\nvoid TestRemoveCopyIfCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n Vector result(5);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n Vector::iterator end = thrust::remove_copy_if(thrust::cuda::par.on(s),\n data.begin(), \n data.end(), \n result.begin(), \n is_even());\n\n ASSERT_EQUAL(end - result.begin(), 3);\n\n ASSERT_EQUAL(result[0], 1);\n ASSERT_EQUAL(result[1], 1);\n ASSERT_EQUAL(result[2], 3);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestRemoveCopyIfCudaStreams);\n\n\nvoid TestRemoveCopyIfStencilCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n Vector stencil(5);\n stencil[0] = 0;\n stencil[1] = 1;\n stencil[2] = 0;\n stencil[3] = 0;\n stencil[4] = 1;\n\n Vector result(5);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n Vector::iterator end = thrust::remove_copy_if(thrust::cuda::par.on(s),\n data.begin(), \n data.end(), \n stencil.begin(),\n result.begin(), \n thrust::identity());\n\n ASSERT_EQUAL(end - result.begin(), 3);\n\n ASSERT_EQUAL(result[0], 1);\n ASSERT_EQUAL(result[1], 1);\n ASSERT_EQUAL(result[2], 3);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestRemoveCopyIfStencilCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid remove_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T val, Iterator2 result)\n{\n *result = thrust::remove(exec, first, last, val);\n}\n\n\ntemplate\n__global__\nvoid remove_if_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Predicate pred, Iterator2 result)\n{\n *result = thrust::remove_if(exec, first, last, pred);\n}\n\n\ntemplate\n__global__\nvoid remove_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Predicate pred, Iterator3 result)\n{\n *result = thrust::remove_if(exec, first, last, stencil_first, pred);\n}\n\n\ntemplate\n__global__\nvoid remove_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, T val, Iterator3 result2)\n{\n *result2 = thrust::remove_copy(exec, first, last, result1, val);\n}\n\n\ntemplate\n__global__\nvoid remove_copy_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result, Predicate pred, Iterator3 result_end)\n{\n *result_end = thrust::remove_copy_if(exec, first, last, result, pred);\n}\n\n\ntemplate\n__global__\nvoid remove_copy_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Iterator3 result, Predicate pred, Iterator4 result_end)\n{\n *result_end = thrust::remove_copy_if(exec, first, last, stencil_first, result, pred);\n}\n#endif\n\n\ntemplate\nstruct is_even\n : thrust::unary_function\n{\n __host__ __device__\n bool operator()(T x) { return (static_cast(x) & 1) == 0; }\n};\n\n\ntemplate\nstruct is_true\n : thrust::unary_function\n{\n __host__ __device__\n bool operator()(T x) { return x ? true : false; }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\nvoid TestRemoveDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n\n typedef typename thrust::device_vector::iterator iterator;\n thrust::device_vector d_result(1);\n \n size_t h_size = thrust::remove(h_data.begin(), h_data.end(), 0) - h_data.begin();\n\n remove_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), 0, d_result.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n size_t d_size = (iterator)d_result[0] - d_data.begin();\n \n ASSERT_EQUAL(h_size, d_size);\n \n h_data.resize(h_size);\n d_data.resize(d_size);\n \n ASSERT_EQUAL(h_data, d_data);\n}\n\n\nvoid TestRemoveDeviceSeq()\n{\n TestRemoveDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestRemoveDeviceSeq);\n\n\nvoid TestRemoveDeviceDevice()\n{\n TestRemoveDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestRemoveDeviceDevice);\n\n\ntemplate\nvoid TestRemoveIfDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n\n typedef typename thrust::device_vector::iterator iterator;\n thrust::device_vector d_result(1);\n \n size_t h_size = thrust::remove_if(h_data.begin(), h_data.end(), is_true()) - h_data.begin();\n\n remove_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), is_true(), d_result.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n size_t d_size = (iterator)d_result[0] - d_data.begin();\n \n ASSERT_EQUAL(h_size, d_size);\n \n h_data.resize(h_size);\n d_data.resize(d_size);\n \n ASSERT_EQUAL(h_data, d_data);\n}\n\n\nvoid TestRemoveIfDeviceSeq()\n{\n TestRemoveIfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestRemoveIfDeviceSeq);\n\n\nvoid TestRemoveIfDeviceDevice()\n{\n TestRemoveIfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestRemoveIfDeviceDevice);\n\n\ntemplate\nvoid TestRemoveIfStencilDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n\n typedef typename thrust::device_vector::iterator iterator;\n thrust::device_vector d_result(1);\n \n thrust::host_vector h_stencil = unittest::random_integers(n);\n thrust::device_vector d_stencil = h_stencil;\n \n size_t h_size = thrust::remove_if(h_data.begin(), h_data.end(), h_stencil.begin(), is_true()) - h_data.begin();\n\n remove_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_stencil.begin(), is_true(), d_result.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n size_t d_size = (iterator)d_result[0] - d_data.begin();\n \n ASSERT_EQUAL(h_size, d_size);\n \n h_data.resize(h_size);\n d_data.resize(d_size);\n \n ASSERT_EQUAL(h_data, d_data);\n}\n\n\nvoid TestRemoveIfStencilDeviceSeq()\n{\n TestRemoveIfStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestRemoveIfStencilDeviceSeq);\n\n\nvoid TestRemoveIfStencilDeviceDevice()\n{\n TestRemoveIfStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestRemoveIfStencilDeviceDevice);\n\n\ntemplate\nvoid TestRemoveCopyDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n\n typedef typename thrust::device_vector::iterator iterator;\n thrust::device_vector d_new_end(1);\n \n size_t h_size = thrust::remove_copy(h_data.begin(), h_data.end(), h_result.begin(), 0) - h_result.begin();\n\n remove_copy_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin(), 0, d_new_end.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n size_t d_size = (iterator)d_new_end[0] - d_result.begin();\n \n ASSERT_EQUAL(h_size, d_size);\n \n h_result.resize(h_size);\n d_result.resize(d_size);\n \n ASSERT_EQUAL(h_result, d_result);\n}\n\n\nvoid TestRemoveCopyDeviceSeq()\n{\n TestRemoveCopyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestRemoveCopyDeviceSeq);\n\n\nvoid TestRemoveCopyDeviceDevice()\n{\n TestRemoveCopyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestRemoveCopyDeviceDevice);\n\n\ntemplate\nvoid TestRemoveCopyIfDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n\n typedef typename thrust::device_vector::iterator iterator;\n thrust::device_vector d_new_end(1);\n \n size_t h_size = thrust::remove_copy_if(h_data.begin(), h_data.end(), h_result.begin(), is_true()) - h_result.begin();\n\n remove_copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin(), is_true(), d_new_end.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n size_t d_size = (iterator)d_new_end[0] - d_result.begin();\n \n ASSERT_EQUAL(h_size, d_size);\n \n h_result.resize(h_size);\n d_result.resize(d_size);\n \n ASSERT_EQUAL(h_result, d_result);\n}\n\n\nvoid TestRemoveCopyIfDeviceSeq()\n{\n TestRemoveCopyIfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestRemoveCopyIfDeviceSeq);\n\n\nvoid TestRemoveCopyIfDeviceDevice()\n{\n TestRemoveCopyIfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestRemoveCopyIfDeviceDevice);\n\n\ntemplate\nvoid TestRemoveCopyIfStencilDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n\n typedef typename thrust::device_vector::iterator iterator;\n thrust::device_vector d_new_end(1);\n\n thrust::host_vector h_stencil = unittest::random_integers(n);\n thrust::device_vector d_stencil = h_stencil;\n \n size_t h_size = thrust::remove_copy_if(h_data.begin(), h_data.end(), h_stencil.begin(), h_result.begin(), is_true()) - h_result.begin();\n\n remove_copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_stencil.begin(), d_result.begin(), is_true(), d_new_end.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n size_t d_size = (iterator)d_new_end[0] - d_result.begin();\n \n ASSERT_EQUAL(h_size, d_size);\n \n h_result.resize(h_size);\n d_result.resize(d_size);\n \n ASSERT_EQUAL(h_result, d_result);\n}\n\n\nvoid TestRemoveCopyIfStencilDeviceSeq()\n{\n TestRemoveCopyIfStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestRemoveCopyIfStencilDeviceSeq);\n\n\nvoid TestRemoveCopyIfStencilDeviceDevice()\n{\n TestRemoveCopyIfStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestRemoveCopyIfStencilDeviceDevice);\n#endif\n\n\nvoid TestRemoveCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n hipStream_t s;\n hipStreamCreate(&s);\n\n Vector::iterator end = thrust::remove(thrust::cuda::par.on(s),\n data.begin(), \n data.end(), \n (T) 2);\n\n ASSERT_EQUAL(end - data.begin(), 3);\n\n ASSERT_EQUAL(data[0], 1);\n ASSERT_EQUAL(data[1], 1);\n ASSERT_EQUAL(data[2], 3);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestRemoveCudaStreams);\n\n\nvoid TestRemoveCopyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n Vector result(5);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n Vector::iterator end = thrust::remove_copy(thrust::cuda::par.on(s),\n data.begin(), \n data.end(), \n result.begin(), \n (T) 2);\n\n ASSERT_EQUAL(end - result.begin(), 3);\n\n ASSERT_EQUAL(result[0], 1);\n ASSERT_EQUAL(result[1], 1);\n ASSERT_EQUAL(result[2], 3);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestRemoveCopyCudaStreams);\n\n\nvoid TestRemoveIfCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n hipStream_t s;\n hipStreamCreate(&s);\n\n Vector::iterator end = thrust::remove_if(thrust::cuda::par.on(s),\n data.begin(), \n data.end(), \n is_even());\n\n ASSERT_EQUAL(end - data.begin(), 3);\n\n ASSERT_EQUAL(data[0], 1);\n ASSERT_EQUAL(data[1], 1);\n ASSERT_EQUAL(data[2], 3);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestRemoveIfCudaStreams);\n\n\nvoid TestRemoveIfStencilCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n Vector stencil(5);\n stencil[0] = 0;\n stencil[1] = 1;\n stencil[2] = 0;\n stencil[3] = 0;\n stencil[4] = 1;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n Vector::iterator end = thrust::remove_if(thrust::cuda::par.on(s),\n data.begin(), \n data.end(),\n stencil.begin(),\n thrust::identity());\n\n ASSERT_EQUAL(end - data.begin(), 3);\n\n ASSERT_EQUAL(data[0], 1);\n ASSERT_EQUAL(data[1], 1);\n ASSERT_EQUAL(data[2], 3);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestRemoveIfStencilCudaStreams);\n\n\nvoid TestRemoveCopyIfCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n Vector result(5);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n Vector::iterator end = thrust::remove_copy_if(thrust::cuda::par.on(s),\n data.begin(), \n data.end(), \n result.begin(), \n is_even());\n\n ASSERT_EQUAL(end - result.begin(), 3);\n\n ASSERT_EQUAL(result[0], 1);\n ASSERT_EQUAL(result[1], 1);\n ASSERT_EQUAL(result[2], 3);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestRemoveCopyIfCudaStreams);\n\n\nvoid TestRemoveCopyIfStencilCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n Vector stencil(5);\n stencil[0] = 0;\n stencil[1] = 1;\n stencil[2] = 0;\n stencil[3] = 0;\n stencil[4] = 1;\n\n Vector result(5);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n Vector::iterator end = thrust::remove_copy_if(thrust::cuda::par.on(s),\n data.begin(), \n data.end(), \n stencil.begin(),\n result.begin(), \n thrust::identity());\n\n ASSERT_EQUAL(end - result.begin(), 3);\n\n ASSERT_EQUAL(result[0], 1);\n ASSERT_EQUAL(result[1], 1);\n ASSERT_EQUAL(result[2], 3);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestRemoveCopyIfStencilCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/remove.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid max_element_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Iterator2 result)\n{\n *result = thrust::max_element(exec, first, last);\n}\n\n\ntemplate\n__global__\nvoid max_element_kernel(ExecutionPolicy exec, Iterator first, Iterator last, BinaryPredicate pred, Iterator2 result)\n{\n *result = thrust::max_element(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestMaxElementDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n\n typedef typename thrust::device_vector::iterator iter_type;\n\n thrust::device_vector d_result(1);\n \n typename thrust::host_vector::iterator h_max = thrust::max_element(h_data.begin(), h_data.end());\n\n max_element_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(h_max - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n\n \n typename thrust::host_vector::iterator h_min = thrust::max_element(h_data.begin(), h_data.end(), thrust::greater());\n\n max_element_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), thrust::greater(), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(h_min - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n}\n\n\nvoid TestMaxElementDeviceSeq()\n{\n TestMaxElementDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestMaxElementDeviceSeq);\n\n\nvoid TestMaxElementDeviceDevice()\n{\n TestMaxElementDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestMaxElementDeviceDevice);\n\n\nvoid TestMaxElementDeviceNoSync()\n{\n TestMaxElementDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestMaxElementDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestMaxElementCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(6);\n data[0] = 3;\n data[1] = 5;\n data[2] = 1;\n data[3] = 2;\n data[4] = 5;\n data[5] = 1;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n\n ASSERT_EQUAL( *thrust::max_element(streampolicy, data.begin(), data.end()), 5);\n ASSERT_EQUAL( thrust::max_element(streampolicy, data.begin(), data.end()) - data.begin(), 1);\n \n ASSERT_EQUAL( *thrust::max_element(streampolicy, data.begin(), data.end(), thrust::greater()), 1);\n ASSERT_EQUAL( thrust::max_element(streampolicy, data.begin(), data.end(), thrust::greater()) - data.begin(), 2);\n\n cudaStreamDestroy(s);\n}\n\nvoid TestMaxElementCudaStreamsSync(){\n TestMaxElementCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestMaxElementCudaStreamsSync);\n\n\nvoid TestMaxElementCudaStreamsNoSync(){\n TestMaxElementCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestMaxElementCudaStreamsNoSync);\n\n\nvoid TestMaxElementDevicePointer()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(6);\n data[0] = 3;\n data[1] = 5;\n data[2] = 1;\n data[3] = 2;\n data[4] = 5;\n data[5] = 1;\n\n T* raw_ptr = thrust::raw_pointer_cast(data.data());\n size_t n = data.size();\n ASSERT_EQUAL( thrust::max_element(thrust::device, raw_ptr, raw_ptr+n) - raw_ptr, 1);\n ASSERT_EQUAL( thrust::max_element(thrust::device, raw_ptr, raw_ptr+n, thrust::greater()) - raw_ptr, 2);\n}\nDECLARE_UNITTEST(TestMaxElementDevicePointer);\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid max_element_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Iterator2 result)\n{\n *result = thrust::max_element(exec, first, last);\n}\n\n\ntemplate\n__global__\nvoid max_element_kernel(ExecutionPolicy exec, Iterator first, Iterator last, BinaryPredicate pred, Iterator2 result)\n{\n *result = thrust::max_element(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestMaxElementDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n\n typedef typename thrust::device_vector::iterator iter_type;\n\n thrust::device_vector d_result(1);\n \n typename thrust::host_vector::iterator h_max = thrust::max_element(h_data.begin(), h_data.end());\n\n max_element_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(h_max - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n\n \n typename thrust::host_vector::iterator h_min = thrust::max_element(h_data.begin(), h_data.end(), thrust::greater());\n\n max_element_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), thrust::greater(), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(h_min - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n}\n\n\nvoid TestMaxElementDeviceSeq()\n{\n TestMaxElementDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestMaxElementDeviceSeq);\n\n\nvoid TestMaxElementDeviceDevice()\n{\n TestMaxElementDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestMaxElementDeviceDevice);\n\n\nvoid TestMaxElementDeviceNoSync()\n{\n TestMaxElementDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestMaxElementDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestMaxElementCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(6);\n data[0] = 3;\n data[1] = 5;\n data[2] = 1;\n data[3] = 2;\n data[4] = 5;\n data[5] = 1;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n\n ASSERT_EQUAL( *thrust::max_element(streampolicy, data.begin(), data.end()), 5);\n ASSERT_EQUAL( thrust::max_element(streampolicy, data.begin(), data.end()) - data.begin(), 1);\n \n ASSERT_EQUAL( *thrust::max_element(streampolicy, data.begin(), data.end(), thrust::greater()), 1);\n ASSERT_EQUAL( thrust::max_element(streampolicy, data.begin(), data.end(), thrust::greater()) - data.begin(), 2);\n\n hipStreamDestroy(s);\n}\n\nvoid TestMaxElementCudaStreamsSync(){\n TestMaxElementCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestMaxElementCudaStreamsSync);\n\n\nvoid TestMaxElementCudaStreamsNoSync(){\n TestMaxElementCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestMaxElementCudaStreamsNoSync);\n\n\nvoid TestMaxElementDevicePointer()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(6);\n data[0] = 3;\n data[1] = 5;\n data[2] = 1;\n data[3] = 2;\n data[4] = 5;\n data[5] = 1;\n\n T* raw_ptr = thrust::raw_pointer_cast(data.data());\n size_t n = data.size();\n ASSERT_EQUAL( thrust::max_element(thrust::device, raw_ptr, raw_ptr+n) - raw_ptr, 1);\n ASSERT_EQUAL( thrust::max_element(thrust::device, raw_ptr, raw_ptr+n, thrust::greater()) - raw_ptr, 2);\n}\nDECLARE_UNITTEST(TestMaxElementDevicePointer);\n", "source": "/kaggle/working/thrust/testing/cuda/max_element.cu"} {"input": "#include \n#include \n#include \n#include \n\n\ntemplate\nstruct my_less\n{\n __host__ __device__\n bool operator()(const T& lhs, const T& rhs) const\n {\n return lhs < rhs;\n }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid sort_by_key_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, Compare comp)\n{\n thrust::sort_by_key(exec, keys_first, keys_last, values_first, comp);\n}\n\n\ntemplate\nvoid TestComparisonSortByKeyDevice(ExecutionPolicy exec, const size_t n, Compare comp)\n{\n thrust::host_vector h_keys = unittest::random_integers(n);\n thrust::device_vector d_keys = h_keys;\n\n thrust::host_vector h_values = h_keys;\n thrust::device_vector d_values = d_keys;\n\n sort_by_key_kernel<<<1,1>>>(exec, d_keys.begin(), d_keys.end(), d_values.begin(), comp);\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n thrust::sort_by_key(h_keys.begin(), h_keys.end(), h_values.begin(), comp);\n\n ASSERT_EQUAL(h_keys, d_keys);\n ASSERT_EQUAL(h_values, d_values);\n};\n\n\ntemplate\n struct TestComparisonSortByKeyDeviceSeq\n{\n void operator()(const size_t n)\n {\n TestComparisonSortByKeyDevice(thrust::seq, n, my_less());\n }\n};\nVariableUnitTest<\n TestComparisonSortByKeyDeviceSeq,\n unittest::type_list\n> TestComparisonSortByKeyDeviceSeqInstance;\n\n\ntemplate\n struct TestComparisonSortByKeyDeviceDevice\n{\n void operator()(const size_t n)\n {\n TestComparisonSortByKeyDevice(thrust::device, n, my_less());\n }\n};\nVariableUnitTest<\n TestComparisonSortByKeyDeviceDevice,\n unittest::type_list\n> TestComparisonSortByKeyDeviceDeviceDeviceInstance;\n\n\ntemplate\nvoid TestSortByKeyDevice(ExecutionPolicy exec, const size_t n)\n{\n TestComparisonSortByKeyDevice(exec, n, thrust::less());\n};\n\n\ntemplate\n struct TestSortByKeyDeviceSeq\n{\n void operator()(const size_t n)\n {\n TestSortByKeyDevice(thrust::seq, n);\n }\n};\nVariableUnitTest<\n TestSortByKeyDeviceSeq,\n unittest::type_list\n> TestSortByKeyDeviceSeqInstance;\n\n\ntemplate\n struct TestSortByKeyDeviceDevice\n{\n void operator()(const size_t n)\n {\n TestSortByKeyDevice(thrust::device, n);\n }\n};\nVariableUnitTest<\n TestSortByKeyDeviceDevice,\n unittest::type_list\n> TestSortByKeyDeviceDeviceInstance;\n#endif\n\n\nvoid TestComparisonSortByKeyCudaStreams()\n{\n thrust::device_vector keys(10);\n thrust::device_vector vals(10);\n\n keys[0] = 9; vals[0] = 9;\n keys[1] = 3; vals[1] = 3;\n keys[2] = 2; vals[2] = 2;\n keys[3] = 0; vals[3] = 0;\n keys[4] = 4; vals[4] = 4;\n keys[5] = 7; vals[5] = 7;\n keys[6] = 8; vals[6] = 8;\n keys[7] = 1; vals[7] = 1;\n keys[8] = 5; vals[8] = 5;\n keys[9] = 6; vals[9] = 6;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::sort_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), my_less());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(true, thrust::is_sorted(keys.begin(), keys.end()));\n ASSERT_EQUAL(true, thrust::is_sorted(vals.begin(), vals.end()));\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestComparisonSortByKeyCudaStreams);\n\n\nvoid TestSortByKeyCudaStreams()\n{\n thrust::device_vector keys(10);\n thrust::device_vector vals(10);\n\n keys[0] = 9; vals[0] = 9;\n keys[1] = 3; vals[1] = 3;\n keys[2] = 2; vals[2] = 2;\n keys[3] = 0; vals[3] = 0;\n keys[4] = 4; vals[4] = 4;\n keys[5] = 7; vals[5] = 7;\n keys[6] = 8; vals[6] = 8;\n keys[7] = 1; vals[7] = 1;\n keys[8] = 5; vals[8] = 5;\n keys[9] = 6; vals[9] = 6;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::sort_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(true, thrust::is_sorted(keys.begin(), keys.end()));\n ASSERT_EQUAL(true, thrust::is_sorted(vals.begin(), vals.end()));\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSortByKeyCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n\ntemplate\nstruct my_less\n{\n __host__ __device__\n bool operator()(const T& lhs, const T& rhs) const\n {\n return lhs < rhs;\n }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid sort_by_key_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, Compare comp)\n{\n thrust::sort_by_key(exec, keys_first, keys_last, values_first, comp);\n}\n\n\ntemplate\nvoid TestComparisonSortByKeyDevice(ExecutionPolicy exec, const size_t n, Compare comp)\n{\n thrust::host_vector h_keys = unittest::random_integers(n);\n thrust::device_vector d_keys = h_keys;\n\n thrust::host_vector h_values = h_keys;\n thrust::device_vector d_values = d_keys;\n\n sort_by_key_kernel<<<1,1>>>(exec, d_keys.begin(), d_keys.end(), d_values.begin(), comp);\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n thrust::sort_by_key(h_keys.begin(), h_keys.end(), h_values.begin(), comp);\n\n ASSERT_EQUAL(h_keys, d_keys);\n ASSERT_EQUAL(h_values, d_values);\n};\n\n\ntemplate\n struct TestComparisonSortByKeyDeviceSeq\n{\n void operator()(const size_t n)\n {\n TestComparisonSortByKeyDevice(thrust::seq, n, my_less());\n }\n};\nVariableUnitTest<\n TestComparisonSortByKeyDeviceSeq,\n unittest::type_list\n> TestComparisonSortByKeyDeviceSeqInstance;\n\n\ntemplate\n struct TestComparisonSortByKeyDeviceDevice\n{\n void operator()(const size_t n)\n {\n TestComparisonSortByKeyDevice(thrust::device, n, my_less());\n }\n};\nVariableUnitTest<\n TestComparisonSortByKeyDeviceDevice,\n unittest::type_list\n> TestComparisonSortByKeyDeviceDeviceDeviceInstance;\n\n\ntemplate\nvoid TestSortByKeyDevice(ExecutionPolicy exec, const size_t n)\n{\n TestComparisonSortByKeyDevice(exec, n, thrust::less());\n};\n\n\ntemplate\n struct TestSortByKeyDeviceSeq\n{\n void operator()(const size_t n)\n {\n TestSortByKeyDevice(thrust::seq, n);\n }\n};\nVariableUnitTest<\n TestSortByKeyDeviceSeq,\n unittest::type_list\n> TestSortByKeyDeviceSeqInstance;\n\n\ntemplate\n struct TestSortByKeyDeviceDevice\n{\n void operator()(const size_t n)\n {\n TestSortByKeyDevice(thrust::device, n);\n }\n};\nVariableUnitTest<\n TestSortByKeyDeviceDevice,\n unittest::type_list\n> TestSortByKeyDeviceDeviceInstance;\n#endif\n\n\nvoid TestComparisonSortByKeyCudaStreams()\n{\n thrust::device_vector keys(10);\n thrust::device_vector vals(10);\n\n keys[0] = 9; vals[0] = 9;\n keys[1] = 3; vals[1] = 3;\n keys[2] = 2; vals[2] = 2;\n keys[3] = 0; vals[3] = 0;\n keys[4] = 4; vals[4] = 4;\n keys[5] = 7; vals[5] = 7;\n keys[6] = 8; vals[6] = 8;\n keys[7] = 1; vals[7] = 1;\n keys[8] = 5; vals[8] = 5;\n keys[9] = 6; vals[9] = 6;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::sort_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), my_less());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(true, thrust::is_sorted(keys.begin(), keys.end()));\n ASSERT_EQUAL(true, thrust::is_sorted(vals.begin(), vals.end()));\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestComparisonSortByKeyCudaStreams);\n\n\nvoid TestSortByKeyCudaStreams()\n{\n thrust::device_vector keys(10);\n thrust::device_vector vals(10);\n\n keys[0] = 9; vals[0] = 9;\n keys[1] = 3; vals[1] = 3;\n keys[2] = 2; vals[2] = 2;\n keys[3] = 0; vals[3] = 0;\n keys[4] = 4; vals[4] = 4;\n keys[5] = 7; vals[5] = 7;\n keys[6] = 8; vals[6] = 8;\n keys[7] = 1; vals[7] = 1;\n keys[8] = 5; vals[8] = 5;\n keys[9] = 6; vals[9] = 6;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::sort_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(true, thrust::is_sorted(keys.begin(), keys.end()));\n ASSERT_EQUAL(true, thrust::is_sorted(vals.begin(), vals.end()));\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSortByKeyCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/sort_by_key.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid count_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T value, Iterator2 result)\n{\n *result = thrust::count(exec, first, last, value);\n}\n\n\ntemplate\nvoid TestCountDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n\n thrust::device_vector d_result(1);\n \n size_t h_result = thrust::count(h_data.begin(), h_data.end(), T(5));\n\n count_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), T(5), d_result.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_EQUAL(h_result, d_result[0]);\n}\n\n\ntemplate\nvoid TestCountDeviceSeq(const size_t n)\n{\n TestCountDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCountDeviceSeq);\n\n\ntemplate\nvoid TestCountDeviceDevice(const size_t n)\n{\n TestCountDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCountDeviceDevice);\n\n\ntemplate\n__global__\nvoid count_if_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Predicate pred, Iterator2 result)\n{\n *result = thrust::count_if(exec, first, last, pred);\n}\n\n\ntemplate\nstruct greater_than_five\n{\n __host__ __device__ bool operator()(const T &x) const {return x > 5;}\n};\n\n\ntemplate\nvoid TestCountIfDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n\n thrust::device_vector d_result(1);\n \n size_t h_result = thrust::count_if(h_data.begin(), h_data.end(), greater_than_five());\n count_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), greater_than_five(), d_result.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_EQUAL(h_result, d_result[0]);\n}\n\n\ntemplate\nvoid TestCountIfDeviceSeq(const size_t n)\n{\n TestCountIfDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCountIfDeviceSeq);\n\n\ntemplate\nvoid TestCountIfDeviceDevice(const size_t n)\n{\n TestCountIfDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCountIfDeviceDevice);\n#endif\n\n\nvoid TestCountCudaStreams()\n{\n thrust::device_vector data(5);\n data[0] = 1; data[1] = 1; data[2] = 0; data[3] = 0; data[4] = 1;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n ASSERT_EQUAL(thrust::count(thrust::cuda::par.on(s), data.begin(), data.end(), 0), 2);\n ASSERT_EQUAL(thrust::count(thrust::cuda::par.on(s), data.begin(), data.end(), 1), 3);\n ASSERT_EQUAL(thrust::count(thrust::cuda::par.on(s), data.begin(), data.end(), 2), 0);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestCountCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid count_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T value, Iterator2 result)\n{\n *result = thrust::count(exec, first, last, value);\n}\n\n\ntemplate\nvoid TestCountDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n\n thrust::device_vector d_result(1);\n \n size_t h_result = thrust::count(h_data.begin(), h_data.end(), T(5));\n\n count_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), T(5), d_result.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_EQUAL(h_result, d_result[0]);\n}\n\n\ntemplate\nvoid TestCountDeviceSeq(const size_t n)\n{\n TestCountDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCountDeviceSeq);\n\n\ntemplate\nvoid TestCountDeviceDevice(const size_t n)\n{\n TestCountDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCountDeviceDevice);\n\n\ntemplate\n__global__\nvoid count_if_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Predicate pred, Iterator2 result)\n{\n *result = thrust::count_if(exec, first, last, pred);\n}\n\n\ntemplate\nstruct greater_than_five\n{\n __host__ __device__ bool operator()(const T &x) const {return x > 5;}\n};\n\n\ntemplate\nvoid TestCountIfDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n\n thrust::device_vector d_result(1);\n \n size_t h_result = thrust::count_if(h_data.begin(), h_data.end(), greater_than_five());\n count_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), greater_than_five(), d_result.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_EQUAL(h_result, d_result[0]);\n}\n\n\ntemplate\nvoid TestCountIfDeviceSeq(const size_t n)\n{\n TestCountIfDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCountIfDeviceSeq);\n\n\ntemplate\nvoid TestCountIfDeviceDevice(const size_t n)\n{\n TestCountIfDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestCountIfDeviceDevice);\n#endif\n\n\nvoid TestCountCudaStreams()\n{\n thrust::device_vector data(5);\n data[0] = 1; data[1] = 1; data[2] = 0; data[3] = 0; data[4] = 1;\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n ASSERT_EQUAL(thrust::count(thrust::cuda::par.on(s), data.begin(), data.end(), 0), 2);\n ASSERT_EQUAL(thrust::count(thrust::cuda::par.on(s), data.begin(), data.end(), 1), 3);\n ASSERT_EQUAL(thrust::count(thrust::cuda::par.on(s), data.begin(), data.end(), 2), 0);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestCountCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/count.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid is_sorted_until_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n *result = thrust::is_sorted_until(exec, first, last);\n}\n\n\ntemplate\nvoid TestIsSortedUntilDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n\n thrust::device_vector v = unittest::random_integers(n);\n\n typedef typename thrust::device_vector::iterator iter_type;\n\n thrust::device_vector result(1);\n\n v[0] = 1;\n v[1] = 0;\n \n is_sorted_until_kernel<<<1,1>>>(exec, v.begin(), v.end(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL_QUIET(v.begin() + 1, (iter_type)result[0]);\n \n thrust::sort(v.begin(), v.end());\n \n is_sorted_until_kernel<<<1,1>>>(exec, v.begin(), v.end(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL_QUIET(v.end(), (iter_type)result[0]);\n}\n\n\nvoid TestIsSortedUntilDeviceSeq()\n{\n TestIsSortedUntilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestIsSortedUntilDeviceSeq);\n\n\nvoid TestIsSortedUntilDeviceDevice()\n{\n TestIsSortedUntilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestIsSortedUntilDeviceDevice);\n#endif\n\n\nvoid TestIsSortedUntilCudaStreams()\n{\n typedef thrust::device_vector Vector;\n\n typedef Vector::value_type T;\n typedef Vector::iterator Iterator;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n Vector v(4);\n v[0] = 0; v[1] = 5; v[2] = 8; v[3] = 0;\n\n Iterator first = v.begin();\n\n Iterator last = v.begin() + 0;\n Iterator ref = last;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last));\n\n last = v.begin() + 1;\n ref = last;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last));\n\n last = v.begin() + 2;\n ref = last;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last));\n\n last = v.begin() + 3;\n ref = v.begin() + 3;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last));\n\n last = v.begin() + 4;\n ref = v.begin() + 3;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last));\n\n last = v.begin() + 3;\n ref = v.begin() + 3;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last, thrust::less()));\n\n last = v.begin() + 4;\n ref = v.begin() + 3;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last, thrust::less()));\n\n last = v.begin() + 1;\n ref = v.begin() + 1;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last, thrust::greater()));\n\n last = v.begin() + 4;\n ref = v.begin() + 1;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last, thrust::greater()));\n\n first = v.begin() + 2;\n last = v.begin() + 4;\n ref = v.begin() + 4;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last, thrust::greater()));\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestIsSortedUntilCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid is_sorted_until_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n *result = thrust::is_sorted_until(exec, first, last);\n}\n\n\ntemplate\nvoid TestIsSortedUntilDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n\n thrust::device_vector v = unittest::random_integers(n);\n\n typedef typename thrust::device_vector::iterator iter_type;\n\n thrust::device_vector result(1);\n\n v[0] = 1;\n v[1] = 0;\n \n is_sorted_until_kernel<<<1,1>>>(exec, v.begin(), v.end(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL_QUIET(v.begin() + 1, (iter_type)result[0]);\n \n thrust::sort(v.begin(), v.end());\n \n is_sorted_until_kernel<<<1,1>>>(exec, v.begin(), v.end(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL_QUIET(v.end(), (iter_type)result[0]);\n}\n\n\nvoid TestIsSortedUntilDeviceSeq()\n{\n TestIsSortedUntilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestIsSortedUntilDeviceSeq);\n\n\nvoid TestIsSortedUntilDeviceDevice()\n{\n TestIsSortedUntilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestIsSortedUntilDeviceDevice);\n#endif\n\n\nvoid TestIsSortedUntilCudaStreams()\n{\n typedef thrust::device_vector Vector;\n\n typedef Vector::value_type T;\n typedef Vector::iterator Iterator;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n Vector v(4);\n v[0] = 0; v[1] = 5; v[2] = 8; v[3] = 0;\n\n Iterator first = v.begin();\n\n Iterator last = v.begin() + 0;\n Iterator ref = last;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last));\n\n last = v.begin() + 1;\n ref = last;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last));\n\n last = v.begin() + 2;\n ref = last;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last));\n\n last = v.begin() + 3;\n ref = v.begin() + 3;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last));\n\n last = v.begin() + 4;\n ref = v.begin() + 3;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last));\n\n last = v.begin() + 3;\n ref = v.begin() + 3;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last, thrust::less()));\n\n last = v.begin() + 4;\n ref = v.begin() + 3;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last, thrust::less()));\n\n last = v.begin() + 1;\n ref = v.begin() + 1;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last, thrust::greater()));\n\n last = v.begin() + 4;\n ref = v.begin() + 1;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last, thrust::greater()));\n\n first = v.begin() + 2;\n last = v.begin() + 4;\n ref = v.begin() + 4;\n ASSERT_EQUAL_QUIET(ref, thrust::is_sorted_until(thrust::cuda::par.on(s), first, last, thrust::greater()));\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestIsSortedUntilCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/is_sorted_until.cu"} {"input": "#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid minmax_element_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n *result = thrust::minmax_element(exec, first, last);\n}\n\n\ntemplate\n__global__\nvoid minmax_element_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, BinaryPredicate pred, Iterator2 result)\n{\n *result = thrust::minmax_element(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestMinMaxElementDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n typename thrust::host_vector::iterator h_min;\n typename thrust::host_vector::iterator h_max;\n typename thrust::device_vector::iterator d_min;\n typename thrust::device_vector::iterator d_max;\n\n typedef thrust::pair<\n typename thrust::device_vector::iterator,\n typename thrust::device_vector::iterator\n > pair_type;\n\n thrust::device_vector d_result(1);\n \n h_min = thrust::minmax_element(h_data.begin(), h_data.end()).first;\n h_max = thrust::minmax_element(h_data.begin(), h_data.end()).second;\n\n d_min = thrust::minmax_element(d_data.begin(), d_data.end()).first;\n d_max = thrust::minmax_element(d_data.begin(), d_data.end()).second;\n\n minmax_element_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n d_min = ((pair_type)d_result[0]).first;\n d_max = ((pair_type)d_result[0]).second;\n \n ASSERT_EQUAL(h_min - h_data.begin(), d_min - d_data.begin());\n ASSERT_EQUAL(h_max - h_data.begin(), d_max - d_data.begin());\n \n h_max = thrust::minmax_element(h_data.begin(), h_data.end(), thrust::greater()).first;\n h_min = thrust::minmax_element(h_data.begin(), h_data.end(), thrust::greater()).second;\n\n minmax_element_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), thrust::greater(), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n d_max = ((pair_type)d_result[0]).first;\n d_min = ((pair_type)d_result[0]).second;\n \n ASSERT_EQUAL(h_min - h_data.begin(), d_min - d_data.begin());\n ASSERT_EQUAL(h_max - h_data.begin(), d_max - d_data.begin());\n}\n\n\nvoid TestMinMaxElementDeviceSeq()\n{\n TestMinMaxElementDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestMinMaxElementDeviceSeq);\n\n\nvoid TestMinMaxElementDeviceDevice()\n{\n TestMinMaxElementDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestMinMaxElementDeviceDevice);\n#endif\n\n\nvoid TestMinMaxElementCudaStreams()\n{\n typedef thrust::device_vector Vector;\n\n Vector data(6);\n data[0] = 3;\n data[1] = 5;\n data[2] = 1;\n data[3] = 2;\n data[4] = 5;\n data[5] = 1;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n ASSERT_EQUAL( *thrust::minmax_element(thrust::cuda::par.on(s), data.begin(), data.end()).first, 1);\n ASSERT_EQUAL( *thrust::minmax_element(thrust::cuda::par.on(s), data.begin(), data.end()).second, 5);\n ASSERT_EQUAL( thrust::minmax_element(thrust::cuda::par.on(s), data.begin(), data.end()).first - data.begin(), 2);\n ASSERT_EQUAL( thrust::minmax_element(thrust::cuda::par.on(s), data.begin(), data.end()).second - data.begin(), 1);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestMinMaxElementCudaStreams);\n\nvoid TestMinMaxElementDevicePointer()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(6);\n data[0] = 3;\n data[1] = 5;\n data[2] = 1;\n data[3] = 2;\n data[4] = 5;\n data[5] = 1;\n\n T* raw_ptr = thrust::raw_pointer_cast(data.data());\n size_t n = data.size();\n ASSERT_EQUAL( thrust::minmax_element(thrust::device, raw_ptr, raw_ptr+n).first - raw_ptr, 2);\n ASSERT_EQUAL( thrust::minmax_element(thrust::device, raw_ptr, raw_ptr+n).second - raw_ptr, 1);\n}\nDECLARE_UNITTEST(TestMinMaxElementDevicePointer);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid minmax_element_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n *result = thrust::minmax_element(exec, first, last);\n}\n\n\ntemplate\n__global__\nvoid minmax_element_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, BinaryPredicate pred, Iterator2 result)\n{\n *result = thrust::minmax_element(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestMinMaxElementDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n typename thrust::host_vector::iterator h_min;\n typename thrust::host_vector::iterator h_max;\n typename thrust::device_vector::iterator d_min;\n typename thrust::device_vector::iterator d_max;\n\n typedef thrust::pair<\n typename thrust::device_vector::iterator,\n typename thrust::device_vector::iterator\n > pair_type;\n\n thrust::device_vector d_result(1);\n \n h_min = thrust::minmax_element(h_data.begin(), h_data.end()).first;\n h_max = thrust::minmax_element(h_data.begin(), h_data.end()).second;\n\n d_min = thrust::minmax_element(d_data.begin(), d_data.end()).first;\n d_max = thrust::minmax_element(d_data.begin(), d_data.end()).second;\n\n minmax_element_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n d_min = ((pair_type)d_result[0]).first;\n d_max = ((pair_type)d_result[0]).second;\n \n ASSERT_EQUAL(h_min - h_data.begin(), d_min - d_data.begin());\n ASSERT_EQUAL(h_max - h_data.begin(), d_max - d_data.begin());\n \n h_max = thrust::minmax_element(h_data.begin(), h_data.end(), thrust::greater()).first;\n h_min = thrust::minmax_element(h_data.begin(), h_data.end(), thrust::greater()).second;\n\n minmax_element_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), thrust::greater(), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n d_max = ((pair_type)d_result[0]).first;\n d_min = ((pair_type)d_result[0]).second;\n \n ASSERT_EQUAL(h_min - h_data.begin(), d_min - d_data.begin());\n ASSERT_EQUAL(h_max - h_data.begin(), d_max - d_data.begin());\n}\n\n\nvoid TestMinMaxElementDeviceSeq()\n{\n TestMinMaxElementDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestMinMaxElementDeviceSeq);\n\n\nvoid TestMinMaxElementDeviceDevice()\n{\n TestMinMaxElementDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestMinMaxElementDeviceDevice);\n#endif\n\n\nvoid TestMinMaxElementCudaStreams()\n{\n typedef thrust::device_vector Vector;\n\n Vector data(6);\n data[0] = 3;\n data[1] = 5;\n data[2] = 1;\n data[3] = 2;\n data[4] = 5;\n data[5] = 1;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n ASSERT_EQUAL( *thrust::minmax_element(thrust::cuda::par.on(s), data.begin(), data.end()).first, 1);\n ASSERT_EQUAL( *thrust::minmax_element(thrust::cuda::par.on(s), data.begin(), data.end()).second, 5);\n ASSERT_EQUAL( thrust::minmax_element(thrust::cuda::par.on(s), data.begin(), data.end()).first - data.begin(), 2);\n ASSERT_EQUAL( thrust::minmax_element(thrust::cuda::par.on(s), data.begin(), data.end()).second - data.begin(), 1);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestMinMaxElementCudaStreams);\n\nvoid TestMinMaxElementDevicePointer()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(6);\n data[0] = 3;\n data[1] = 5;\n data[2] = 1;\n data[3] = 2;\n data[4] = 5;\n data[5] = 1;\n\n T* raw_ptr = thrust::raw_pointer_cast(data.data());\n size_t n = data.size();\n ASSERT_EQUAL( thrust::minmax_element(thrust::device, raw_ptr, raw_ptr+n).first - raw_ptr, 2);\n ASSERT_EQUAL( thrust::minmax_element(thrust::device, raw_ptr, raw_ptr+n).second - raw_ptr, 1);\n}\nDECLARE_UNITTEST(TestMinMaxElementDevicePointer);\n\n", "source": "/kaggle/working/thrust/testing/cuda/minmax_element.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_symmetric_difference_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first1, Iterator1 keys_last1,\n Iterator2 keys_first2, Iterator2 keys_last2,\n Iterator3 values_first1,\n Iterator4 values_first2,\n Iterator5 keys_result,\n Iterator6 values_result,\n Iterator7 result)\n{\n *result = thrust::set_symmetric_difference_by_key(exec, keys_first1, keys_last1, keys_first2, keys_last2, values_first1, values_first2, keys_result, values_result);\n}\n\n\ntemplate\nvoid TestSetSymmetricDifferenceByKeyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::iterator Iterator;\n\n Vector a_key(4), b_key(5);\n Vector a_val(4), b_val(5);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4; a_key[3] = 6;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0; a_val[3] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4; b_key[4] = 7;\n b_val[0] = 1; b_val[1] = 1; b_val[2] = 1; b_val[3] = 1; b_val[4] = 1;\n\n Vector ref_key(5), ref_val(5);\n ref_key[0] = 2; ref_key[1] = 3; ref_key[2] = 3; ref_key[3] = 6; ref_key[4] = 7;\n ref_val[0] = 0; ref_val[1] = 1; ref_val[2] = 1; ref_val[3] = 0; ref_val[4] = 1;\n\n Vector result_key(5), result_val(5);\n\n typedef thrust::pair iter_pair;\n thrust::device_vector end_vec(1);\n\n set_symmetric_difference_by_key_kernel<<<1,1>>>(exec,\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n b_val.begin(),\n result_key.begin(),\n result_val.begin(),\n end_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n iter_pair end = end_vec[0];\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n}\n\n\nvoid TestSetSymmetricDifferenceByKeyDeviceSeq()\n{\n TestSetSymmetricDifferenceByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetSymmetricDifferenceByKeyDeviceSeq);\n\n\nvoid TestSetSymmetricDifferenceByKeyDeviceDevice()\n{\n TestSetSymmetricDifferenceByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetSymmetricDifferenceByKeyDeviceDevice);\n#endif\n\n\nvoid TestSetSymmetricDifferenceByKeyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a_key(4), b_key(5);\n Vector a_val(4), b_val(5);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4; a_key[3] = 6;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0; a_val[3] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4; b_key[4] = 7;\n b_val[0] = 1; b_val[1] = 1; b_val[2] = 1; b_val[3] = 1; b_val[4] = 1;\n\n Vector ref_key(5), ref_val(5);\n ref_key[0] = 2; ref_key[1] = 3; ref_key[2] = 3; ref_key[3] = 6; ref_key[4] = 7;\n ref_val[0] = 0; ref_val[1] = 1; ref_val[2] = 1; ref_val[3] = 0; ref_val[4] = 1;\n\n Vector result_key(5), result_val(5);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::pair end =\n thrust::set_symmetric_difference_by_key(thrust::cuda::par.on(s),\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n b_val.begin(),\n result_key.begin(),\n result_val.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSetSymmetricDifferenceByKeyCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_symmetric_difference_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first1, Iterator1 keys_last1,\n Iterator2 keys_first2, Iterator2 keys_last2,\n Iterator3 values_first1,\n Iterator4 values_first2,\n Iterator5 keys_result,\n Iterator6 values_result,\n Iterator7 result)\n{\n *result = thrust::set_symmetric_difference_by_key(exec, keys_first1, keys_last1, keys_first2, keys_last2, values_first1, values_first2, keys_result, values_result);\n}\n\n\ntemplate\nvoid TestSetSymmetricDifferenceByKeyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::iterator Iterator;\n\n Vector a_key(4), b_key(5);\n Vector a_val(4), b_val(5);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4; a_key[3] = 6;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0; a_val[3] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4; b_key[4] = 7;\n b_val[0] = 1; b_val[1] = 1; b_val[2] = 1; b_val[3] = 1; b_val[4] = 1;\n\n Vector ref_key(5), ref_val(5);\n ref_key[0] = 2; ref_key[1] = 3; ref_key[2] = 3; ref_key[3] = 6; ref_key[4] = 7;\n ref_val[0] = 0; ref_val[1] = 1; ref_val[2] = 1; ref_val[3] = 0; ref_val[4] = 1;\n\n Vector result_key(5), result_val(5);\n\n typedef thrust::pair iter_pair;\n thrust::device_vector end_vec(1);\n\n set_symmetric_difference_by_key_kernel<<<1,1>>>(exec,\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n b_val.begin(),\n result_key.begin(),\n result_val.begin(),\n end_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n iter_pair end = end_vec[0];\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n}\n\n\nvoid TestSetSymmetricDifferenceByKeyDeviceSeq()\n{\n TestSetSymmetricDifferenceByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetSymmetricDifferenceByKeyDeviceSeq);\n\n\nvoid TestSetSymmetricDifferenceByKeyDeviceDevice()\n{\n TestSetSymmetricDifferenceByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetSymmetricDifferenceByKeyDeviceDevice);\n#endif\n\n\nvoid TestSetSymmetricDifferenceByKeyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a_key(4), b_key(5);\n Vector a_val(4), b_val(5);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4; a_key[3] = 6;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0; a_val[3] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4; b_key[4] = 7;\n b_val[0] = 1; b_val[1] = 1; b_val[2] = 1; b_val[3] = 1; b_val[4] = 1;\n\n Vector ref_key(5), ref_val(5);\n ref_key[0] = 2; ref_key[1] = 3; ref_key[2] = 3; ref_key[3] = 6; ref_key[4] = 7;\n ref_val[0] = 0; ref_val[1] = 1; ref_val[2] = 1; ref_val[3] = 0; ref_val[4] = 1;\n\n Vector result_key(5), result_val(5);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::pair end =\n thrust::set_symmetric_difference_by_key(thrust::cuda::par.on(s),\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n b_val.begin(),\n result_key.begin(),\n result_val.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSetSymmetricDifferenceByKeyCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/set_symmetric_difference_by_key.cu"} {"input": "#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid inclusive_scan_by_key_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, Iterator3 result)\n{\n thrust::inclusive_scan_by_key(exec, keys_first, keys_last, values_first, result);\n}\n\n\ntemplate\n__global__\nvoid exclusive_scan_by_key_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, Iterator3 result)\n{\n thrust::exclusive_scan_by_key(exec, keys_first, keys_last, values_first, result);\n}\n\n\ntemplate\n__global__\nvoid exclusive_scan_by_key_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, Iterator3 result, T init)\n{\n thrust::exclusive_scan_by_key(exec, keys_first, keys_last, values_first, result, init);\n}\n\n\ntemplate\nvoid TestScanByKeyDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n\n thrust::host_vector h_keys(n);\n for(size_t i = 0, k = 0; i < n; i++)\n {\n h_keys[i] = static_cast(k);\n if(rand() % 10 == 0)\n {\n k++;\n }\n }\n thrust::device_vector d_keys = h_keys;\n \n thrust::host_vector h_vals = unittest::random_integers(n);\n for(size_t i = 0; i < n; i++)\n {\n h_vals[i] = i % 10;\n }\n thrust::device_vector d_vals = h_vals;\n \n thrust::host_vector h_output(n);\n thrust::device_vector d_output(n);\n \n thrust::inclusive_scan_by_key(h_keys.begin(), h_keys.end(), h_vals.begin(), h_output.begin());\n inclusive_scan_by_key_kernel<<<1,1>>>(exec, d_keys.begin(), d_keys.end(), d_vals.begin(), d_output.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n ASSERT_EQUAL(d_output, h_output);\n \n thrust::exclusive_scan_by_key(h_keys.begin(), h_keys.end(), h_vals.begin(), h_output.begin());\n exclusive_scan_by_key_kernel<<<1,1>>>(exec, d_keys.begin(), d_keys.end(), d_vals.begin(), d_output.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n ASSERT_EQUAL(d_output, h_output);\n \n thrust::exclusive_scan_by_key(h_keys.begin(), h_keys.end(), h_vals.begin(), h_output.begin(), 11);\n exclusive_scan_by_key_kernel<<<1,1>>>(exec, d_keys.begin(), d_keys.end(), d_vals.begin(), d_output.begin(), 11);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n ASSERT_EQUAL(d_output, h_output);\n \n // in-place scans: in/out values aliasing\n h_output = h_vals;\n d_output = d_vals;\n thrust::inclusive_scan_by_key(h_keys.begin(), h_keys.end(), h_output.begin(), h_output.begin());\n inclusive_scan_by_key_kernel<<<1,1>>>(exec,d_keys.begin(), d_keys.end(), d_output.begin(), d_output.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n ASSERT_EQUAL(d_output, h_output);\n \n h_output = h_vals;\n d_output = d_vals;\n thrust::exclusive_scan_by_key(h_keys.begin(), h_keys.end(), h_output.begin(), h_output.begin(), 11);\n exclusive_scan_by_key_kernel<<<1,1>>>(exec, d_keys.begin(), d_keys.end(), d_output.begin(), d_output.begin(), 11);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n ASSERT_EQUAL(d_output, h_output);\n\n // in-place scans: keys/values aliasing\n thrust::inclusive_scan_by_key(h_keys.begin(), h_keys.end(), h_vals.begin(), h_output.begin());\n inclusive_scan_by_key_kernel<<<1,1>>>(exec, d_keys.begin(), d_keys.end(), d_vals.begin(), d_keys.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n ASSERT_EQUAL(d_keys, h_output);\n\n d_keys = h_keys;\n thrust::exclusive_scan_by_key(h_keys.begin(), h_keys.end(), h_vals.begin(), h_output.begin(), 11);\n exclusive_scan_by_key_kernel<<<1,1>>>(exec, d_keys.begin(), d_keys.end(), d_vals.begin(), d_keys.begin(), 11);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n ASSERT_EQUAL(d_keys, h_output);\n}\n\n\nvoid TestScanByKeyDeviceSeq()\n{\n TestScanByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestScanByKeyDeviceSeq);\n\n\nvoid TestScanByKeyDeviceDevice()\n{\n TestScanByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestScanByKeyDeviceDevice);\n#endif\n\n\nvoid TestInclusiveScanByKeyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n typedef Vector::iterator Iterator;\n\n Vector keys(7);\n Vector vals(7);\n\n Vector output(7, 0);\n\n keys[0] = 0; vals[0] = 1;\n keys[1] = 1; vals[1] = 2;\n keys[2] = 1; vals[2] = 3;\n keys[3] = 1; vals[3] = 4;\n keys[4] = 2; vals[4] = 5;\n keys[5] = 3; vals[5] = 6;\n keys[6] = 3; vals[6] = 7;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n Iterator iter = thrust::inclusive_scan_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), output.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(iter, output.end());\n\n ASSERT_EQUAL(output[0], 1);\n ASSERT_EQUAL(output[1], 2);\n ASSERT_EQUAL(output[2], 5);\n ASSERT_EQUAL(output[3], 9);\n ASSERT_EQUAL(output[4], 5);\n ASSERT_EQUAL(output[5], 6);\n ASSERT_EQUAL(output[6], 13);\n \n thrust::inclusive_scan_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), output.begin(), thrust::equal_to(), thrust::multiplies());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(output[0], 1);\n ASSERT_EQUAL(output[1], 2);\n ASSERT_EQUAL(output[2], 6);\n ASSERT_EQUAL(output[3], 24);\n ASSERT_EQUAL(output[4], 5);\n ASSERT_EQUAL(output[5], 6);\n ASSERT_EQUAL(output[6], 42);\n \n thrust::inclusive_scan_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), output.begin(), thrust::equal_to());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(output[0], 1);\n ASSERT_EQUAL(output[1], 2);\n ASSERT_EQUAL(output[2], 5);\n ASSERT_EQUAL(output[3], 9);\n ASSERT_EQUAL(output[4], 5);\n ASSERT_EQUAL(output[5], 6);\n ASSERT_EQUAL(output[6], 13);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestInclusiveScanByKeyCudaStreams);\n\n\nvoid TestExclusiveScanByKeyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n typedef Vector::iterator Iterator;\n\n Vector keys(7);\n Vector vals(7);\n\n Vector output(7, 0);\n\n keys[0] = 0; vals[0] = 1;\n keys[1] = 1; vals[1] = 2;\n keys[2] = 1; vals[2] = 3;\n keys[3] = 1; vals[3] = 4;\n keys[4] = 2; vals[4] = 5;\n keys[5] = 3; vals[5] = 6;\n keys[6] = 3; vals[6] = 7;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n Iterator iter = thrust::exclusive_scan_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), output.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(iter, output.end());\n\n ASSERT_EQUAL(output[0], 0);\n ASSERT_EQUAL(output[1], 0);\n ASSERT_EQUAL(output[2], 2);\n ASSERT_EQUAL(output[3], 5);\n ASSERT_EQUAL(output[4], 0);\n ASSERT_EQUAL(output[5], 0);\n ASSERT_EQUAL(output[6], 6);\n\n thrust::exclusive_scan_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), output.begin(), T(10));\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(output[0], 10);\n ASSERT_EQUAL(output[1], 10);\n ASSERT_EQUAL(output[2], 12);\n ASSERT_EQUAL(output[3], 15);\n ASSERT_EQUAL(output[4], 10);\n ASSERT_EQUAL(output[5], 10);\n ASSERT_EQUAL(output[6], 16);\n \n thrust::exclusive_scan_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), output.begin(), T(10), thrust::equal_to(), thrust::multiplies());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(output[0], 10);\n ASSERT_EQUAL(output[1], 10);\n ASSERT_EQUAL(output[2], 20);\n ASSERT_EQUAL(output[3], 60);\n ASSERT_EQUAL(output[4], 10);\n ASSERT_EQUAL(output[5], 10);\n ASSERT_EQUAL(output[6], 60);\n \n thrust::exclusive_scan_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), output.begin(), T(10), thrust::equal_to());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(output[0], 10);\n ASSERT_EQUAL(output[1], 10);\n ASSERT_EQUAL(output[2], 12);\n ASSERT_EQUAL(output[3], 15);\n ASSERT_EQUAL(output[4], 10);\n ASSERT_EQUAL(output[5], 10);\n ASSERT_EQUAL(output[6], 16);\n}\nDECLARE_UNITTEST(TestExclusiveScanByKeyCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid inclusive_scan_by_key_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, Iterator3 result)\n{\n thrust::inclusive_scan_by_key(exec, keys_first, keys_last, values_first, result);\n}\n\n\ntemplate\n__global__\nvoid exclusive_scan_by_key_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, Iterator3 result)\n{\n thrust::exclusive_scan_by_key(exec, keys_first, keys_last, values_first, result);\n}\n\n\ntemplate\n__global__\nvoid exclusive_scan_by_key_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first, Iterator3 result, T init)\n{\n thrust::exclusive_scan_by_key(exec, keys_first, keys_last, values_first, result, init);\n}\n\n\ntemplate\nvoid TestScanByKeyDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n\n thrust::host_vector h_keys(n);\n for(size_t i = 0, k = 0; i < n; i++)\n {\n h_keys[i] = static_cast(k);\n if(rand() % 10 == 0)\n {\n k++;\n }\n }\n thrust::device_vector d_keys = h_keys;\n \n thrust::host_vector h_vals = unittest::random_integers(n);\n for(size_t i = 0; i < n; i++)\n {\n h_vals[i] = i % 10;\n }\n thrust::device_vector d_vals = h_vals;\n \n thrust::host_vector h_output(n);\n thrust::device_vector d_output(n);\n \n thrust::inclusive_scan_by_key(h_keys.begin(), h_keys.end(), h_vals.begin(), h_output.begin());\n inclusive_scan_by_key_kernel<<<1,1>>>(exec, d_keys.begin(), d_keys.end(), d_vals.begin(), d_output.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n ASSERT_EQUAL(d_output, h_output);\n \n thrust::exclusive_scan_by_key(h_keys.begin(), h_keys.end(), h_vals.begin(), h_output.begin());\n exclusive_scan_by_key_kernel<<<1,1>>>(exec, d_keys.begin(), d_keys.end(), d_vals.begin(), d_output.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n ASSERT_EQUAL(d_output, h_output);\n \n thrust::exclusive_scan_by_key(h_keys.begin(), h_keys.end(), h_vals.begin(), h_output.begin(), 11);\n exclusive_scan_by_key_kernel<<<1,1>>>(exec, d_keys.begin(), d_keys.end(), d_vals.begin(), d_output.begin(), 11);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n ASSERT_EQUAL(d_output, h_output);\n \n // in-place scans: in/out values aliasing\n h_output = h_vals;\n d_output = d_vals;\n thrust::inclusive_scan_by_key(h_keys.begin(), h_keys.end(), h_output.begin(), h_output.begin());\n inclusive_scan_by_key_kernel<<<1,1>>>(exec,d_keys.begin(), d_keys.end(), d_output.begin(), d_output.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n ASSERT_EQUAL(d_output, h_output);\n \n h_output = h_vals;\n d_output = d_vals;\n thrust::exclusive_scan_by_key(h_keys.begin(), h_keys.end(), h_output.begin(), h_output.begin(), 11);\n exclusive_scan_by_key_kernel<<<1,1>>>(exec, d_keys.begin(), d_keys.end(), d_output.begin(), d_output.begin(), 11);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n ASSERT_EQUAL(d_output, h_output);\n\n // in-place scans: keys/values aliasing\n thrust::inclusive_scan_by_key(h_keys.begin(), h_keys.end(), h_vals.begin(), h_output.begin());\n inclusive_scan_by_key_kernel<<<1,1>>>(exec, d_keys.begin(), d_keys.end(), d_vals.begin(), d_keys.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n ASSERT_EQUAL(d_keys, h_output);\n\n d_keys = h_keys;\n thrust::exclusive_scan_by_key(h_keys.begin(), h_keys.end(), h_vals.begin(), h_output.begin(), 11);\n exclusive_scan_by_key_kernel<<<1,1>>>(exec, d_keys.begin(), d_keys.end(), d_vals.begin(), d_keys.begin(), 11);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n ASSERT_EQUAL(d_keys, h_output);\n}\n\n\nvoid TestScanByKeyDeviceSeq()\n{\n TestScanByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestScanByKeyDeviceSeq);\n\n\nvoid TestScanByKeyDeviceDevice()\n{\n TestScanByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestScanByKeyDeviceDevice);\n#endif\n\n\nvoid TestInclusiveScanByKeyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n typedef Vector::iterator Iterator;\n\n Vector keys(7);\n Vector vals(7);\n\n Vector output(7, 0);\n\n keys[0] = 0; vals[0] = 1;\n keys[1] = 1; vals[1] = 2;\n keys[2] = 1; vals[2] = 3;\n keys[3] = 1; vals[3] = 4;\n keys[4] = 2; vals[4] = 5;\n keys[5] = 3; vals[5] = 6;\n keys[6] = 3; vals[6] = 7;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n Iterator iter = thrust::inclusive_scan_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), output.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(iter, output.end());\n\n ASSERT_EQUAL(output[0], 1);\n ASSERT_EQUAL(output[1], 2);\n ASSERT_EQUAL(output[2], 5);\n ASSERT_EQUAL(output[3], 9);\n ASSERT_EQUAL(output[4], 5);\n ASSERT_EQUAL(output[5], 6);\n ASSERT_EQUAL(output[6], 13);\n \n thrust::inclusive_scan_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), output.begin(), thrust::equal_to(), thrust::multiplies());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(output[0], 1);\n ASSERT_EQUAL(output[1], 2);\n ASSERT_EQUAL(output[2], 6);\n ASSERT_EQUAL(output[3], 24);\n ASSERT_EQUAL(output[4], 5);\n ASSERT_EQUAL(output[5], 6);\n ASSERT_EQUAL(output[6], 42);\n \n thrust::inclusive_scan_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), output.begin(), thrust::equal_to());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(output[0], 1);\n ASSERT_EQUAL(output[1], 2);\n ASSERT_EQUAL(output[2], 5);\n ASSERT_EQUAL(output[3], 9);\n ASSERT_EQUAL(output[4], 5);\n ASSERT_EQUAL(output[5], 6);\n ASSERT_EQUAL(output[6], 13);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestInclusiveScanByKeyCudaStreams);\n\n\nvoid TestExclusiveScanByKeyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n typedef Vector::iterator Iterator;\n\n Vector keys(7);\n Vector vals(7);\n\n Vector output(7, 0);\n\n keys[0] = 0; vals[0] = 1;\n keys[1] = 1; vals[1] = 2;\n keys[2] = 1; vals[2] = 3;\n keys[3] = 1; vals[3] = 4;\n keys[4] = 2; vals[4] = 5;\n keys[5] = 3; vals[5] = 6;\n keys[6] = 3; vals[6] = 7;\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n Iterator iter = thrust::exclusive_scan_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), output.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(iter, output.end());\n\n ASSERT_EQUAL(output[0], 0);\n ASSERT_EQUAL(output[1], 0);\n ASSERT_EQUAL(output[2], 2);\n ASSERT_EQUAL(output[3], 5);\n ASSERT_EQUAL(output[4], 0);\n ASSERT_EQUAL(output[5], 0);\n ASSERT_EQUAL(output[6], 6);\n\n thrust::exclusive_scan_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), output.begin(), T(10));\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(output[0], 10);\n ASSERT_EQUAL(output[1], 10);\n ASSERT_EQUAL(output[2], 12);\n ASSERT_EQUAL(output[3], 15);\n ASSERT_EQUAL(output[4], 10);\n ASSERT_EQUAL(output[5], 10);\n ASSERT_EQUAL(output[6], 16);\n \n thrust::exclusive_scan_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), output.begin(), T(10), thrust::equal_to(), thrust::multiplies());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(output[0], 10);\n ASSERT_EQUAL(output[1], 10);\n ASSERT_EQUAL(output[2], 20);\n ASSERT_EQUAL(output[3], 60);\n ASSERT_EQUAL(output[4], 10);\n ASSERT_EQUAL(output[5], 10);\n ASSERT_EQUAL(output[6], 60);\n \n thrust::exclusive_scan_by_key(thrust::cuda::par.on(s), keys.begin(), keys.end(), vals.begin(), output.begin(), T(10), thrust::equal_to());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(output[0], 10);\n ASSERT_EQUAL(output[1], 10);\n ASSERT_EQUAL(output[2], 12);\n ASSERT_EQUAL(output[3], 15);\n ASSERT_EQUAL(output[4], 10);\n ASSERT_EQUAL(output[5], 10);\n ASSERT_EQUAL(output[6], 16);\n}\nDECLARE_UNITTEST(TestExclusiveScanByKeyCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/scan_by_key.cu"} {"input": "#include \n\n#include \n\ntemplate \n__host__ __device__ void universal_vector_access(VecT &in, thrust::universal_vector &out)\n{\n const int expected_front = 4;\n const int expected_back = 2;\n\n out[0] = in.size() == 2 && //\n in[0] == expected_front && //\n in.front() == expected_front && //\n *in.data() == expected_front && //\n in[1] == expected_back && //\n in.back() == expected_back;\n}\n\n#if defined(THRUST_TEST_DEVICE_SIDE)\ntemplate \n__global__ void universal_vector_device_access_kernel(VecT &vec,\n thrust::universal_vector &out)\n{\n universal_vector_access(vec, out);\n}\n\ntemplate \nvoid test_universal_vector_access(VecT &vec, thrust::universal_vector &out)\n{\n universal_vector_device_access_kernel<<<1, 1>>>(vec, out);\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n ASSERT_EQUAL(out[0], true);\n}\n#else\ntemplate \nvoid test_universal_vector_access(VecT &vec, thrust::universal_vector &out)\n{\n universal_vector_access(vec, out);\n ASSERT_EQUAL(out[0], true);\n}\n#endif\n\nvoid TestUniversalVectorDeviceAccess()\n{\n thrust::universal_vector> in_storage(1);\n thrust::universal_vector &in = *thrust::raw_pointer_cast(in_storage.data());\n\n in.resize(2);\n in[0] = 4;\n in[1] = 2;\n\n thrust::universal_vector> out_storage(1);\n thrust::universal_vector &out = *thrust::raw_pointer_cast(out_storage.data());\n out.resize(1);\n out[0] = false;\n\n test_universal_vector_access(in, out);\n}\nDECLARE_UNITTEST(TestUniversalVectorDeviceAccess);\n\nvoid TestConstUniversalVectorDeviceAccess()\n{\n thrust::universal_vector> in_storage(1);\n\n {\n thrust::universal_vector &in = *thrust::raw_pointer_cast(in_storage.data());\n\n in.resize(2);\n in[0] = 4;\n in[1] = 2;\n }\n\n const thrust::universal_vector &in = *thrust::raw_pointer_cast(in_storage.data());\n\n thrust::universal_vector> out_storage(1);\n thrust::universal_vector &out = *thrust::raw_pointer_cast(out_storage.data());\n\n out.resize(1);\n out[0] = false;\n\n test_universal_vector_access(in, out);\n}\nDECLARE_UNITTEST(TestConstUniversalVectorDeviceAccess);\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n\n#include \n\ntemplate \n__host__ __device__ void universal_vector_access(VecT &in, thrust::universal_vector &out)\n{\n const int expected_front = 4;\n const int expected_back = 2;\n\n out[0] = in.size() == 2 && //\n in[0] == expected_front && //\n in.front() == expected_front && //\n *in.data() == expected_front && //\n in[1] == expected_back && //\n in.back() == expected_back;\n}\n\n#if defined(THRUST_TEST_DEVICE_SIDE)\ntemplate \n__global__ void universal_vector_device_access_kernel(VecT &vec,\n thrust::universal_vector &out)\n{\n universal_vector_access(vec, out);\n}\n\ntemplate \nvoid test_universal_vector_access(VecT &vec, thrust::universal_vector &out)\n{\n universal_vector_device_access_kernel<<<1, 1>>>(vec, out);\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n ASSERT_EQUAL(out[0], true);\n}\n#else\ntemplate \nvoid test_universal_vector_access(VecT &vec, thrust::universal_vector &out)\n{\n universal_vector_access(vec, out);\n ASSERT_EQUAL(out[0], true);\n}\n#endif\n\nvoid TestUniversalVectorDeviceAccess()\n{\n thrust::universal_vector> in_storage(1);\n thrust::universal_vector &in = *thrust::raw_pointer_cast(in_storage.data());\n\n in.resize(2);\n in[0] = 4;\n in[1] = 2;\n\n thrust::universal_vector> out_storage(1);\n thrust::universal_vector &out = *thrust::raw_pointer_cast(out_storage.data());\n out.resize(1);\n out[0] = false;\n\n test_universal_vector_access(in, out);\n}\nDECLARE_UNITTEST(TestUniversalVectorDeviceAccess);\n\nvoid TestConstUniversalVectorDeviceAccess()\n{\n thrust::universal_vector> in_storage(1);\n\n {\n thrust::universal_vector &in = *thrust::raw_pointer_cast(in_storage.data());\n\n in.resize(2);\n in[0] = 4;\n in[1] = 2;\n }\n\n const thrust::universal_vector &in = *thrust::raw_pointer_cast(in_storage.data());\n\n thrust::universal_vector> out_storage(1);\n thrust::universal_vector &out = *thrust::raw_pointer_cast(out_storage.data());\n\n out.resize(1);\n out[0] = false;\n\n test_universal_vector_access(in, out);\n}\nDECLARE_UNITTEST(TestConstUniversalVectorDeviceAccess);\n", "source": "/kaggle/working/thrust/testing/cuda/device_side_universal_vector.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_union_kernel(ExecutionPolicy exec,\n Iterator1 first1, Iterator1 last1,\n Iterator2 first2, Iterator2 last2,\n Iterator3 result1,\n Iterator4 result2)\n{\n *result2 = thrust::set_union(exec, first1, last1, first2, last2, result1);\n}\n\n\ntemplate\nvoid TestSetUnionDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::iterator Iterator;\n\n Vector a(3), b(4);\n\n a[0] = 0; a[1] = 2; a[2] = 4;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4;\n\n Vector ref(5);\n ref[0] = 0; ref[1] = 2; ref[2] = 3; ref[3] = 3; ref[4] = 4;\n\n Vector result(5);\n thrust::device_vector end_vec(1);\n\n set_union_kernel<<<1,1>>>(exec,\n a.begin(), a.end(),\n b.begin(), b.end(),\n result.begin(),\n end_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n Iterator end = end_vec[0];\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n}\n\n\nvoid TestSetUnionDeviceSeq()\n{\n TestSetUnionDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetUnionDeviceSeq);\n\n\nvoid TestSetUnionDeviceDevice()\n{\n TestSetUnionDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetUnionDeviceDevice);\n#endif\n\n\nvoid TestSetUnionCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a(3), b(4);\n\n a[0] = 0; a[1] = 2; a[2] = 4;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4;\n\n Vector ref(5);\n ref[0] = 0; ref[1] = 2; ref[2] = 3; ref[3] = 3; ref[4] = 4;\n\n Vector result(5);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n Iterator end = thrust::set_union(thrust::cuda::par.on(s),\n a.begin(), a.end(),\n b.begin(), b.end(),\n result.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSetUnionCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_union_kernel(ExecutionPolicy exec,\n Iterator1 first1, Iterator1 last1,\n Iterator2 first2, Iterator2 last2,\n Iterator3 result1,\n Iterator4 result2)\n{\n *result2 = thrust::set_union(exec, first1, last1, first2, last2, result1);\n}\n\n\ntemplate\nvoid TestSetUnionDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::iterator Iterator;\n\n Vector a(3), b(4);\n\n a[0] = 0; a[1] = 2; a[2] = 4;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4;\n\n Vector ref(5);\n ref[0] = 0; ref[1] = 2; ref[2] = 3; ref[3] = 3; ref[4] = 4;\n\n Vector result(5);\n thrust::device_vector end_vec(1);\n\n set_union_kernel<<<1,1>>>(exec,\n a.begin(), a.end(),\n b.begin(), b.end(),\n result.begin(),\n end_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n Iterator end = end_vec[0];\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n}\n\n\nvoid TestSetUnionDeviceSeq()\n{\n TestSetUnionDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetUnionDeviceSeq);\n\n\nvoid TestSetUnionDeviceDevice()\n{\n TestSetUnionDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetUnionDeviceDevice);\n#endif\n\n\nvoid TestSetUnionCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a(3), b(4);\n\n a[0] = 0; a[1] = 2; a[2] = 4;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4;\n\n Vector ref(5);\n ref[0] = 0; ref[1] = 2; ref[2] = 3; ref[3] = 3; ref[4] = 4;\n\n Vector result(5);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n Iterator end = thrust::set_union(thrust::cuda::par.on(s),\n a.begin(), a.end(),\n b.begin(), b.end(),\n result.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSetUnionCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/set_union.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid reverse_kernel(ExecutionPolicy exec, Iterator first, Iterator last)\n{\n thrust::reverse(exec, first, last);\n}\n\n\ntemplate\nvoid TestReverseDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n thrust::reverse(h_data.begin(), h_data.end());\n\n reverse_kernel<<<1,1>>>(exec, raw_pointer_cast(d_data.data()), raw_pointer_cast(d_data.data() + d_data.size()));\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_EQUAL(h_data, d_data);\n};\n\n\nvoid TestReverseDeviceSeq()\n{\n TestReverseDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReverseDeviceSeq);\n\n\nvoid TestReverseDeviceDevice()\n{\n TestReverseDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReverseDeviceDevice);\n\n\ntemplate\n__global__\nvoid reverse_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n thrust::reverse_copy(exec, first, last, result);\n}\n\n\ntemplate\nvoid TestReverseCopyDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n\n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n\n thrust::reverse_copy(h_data.begin(), h_data.end(), h_result.begin());\n\n reverse_copy_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n ASSERT_EQUAL(h_result, d_result);\n};\n\n\nvoid TestReverseCopyDeviceSeq()\n{\n TestReverseCopyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReverseCopyDeviceSeq);\n\n\nvoid TestReverseCopyDeviceDevice()\n{\n TestReverseCopyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReverseCopyDeviceDevice);\n#endif\n\n\nvoid TestReverseCudaStreams()\n{\n typedef thrust::device_vector Vector;\n Vector data(5);\n data[0] = 1;\n data[1] = 2;\n data[2] = 3;\n data[3] = 4;\n data[4] = 5;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::reverse(thrust::cuda::par.on(s), data.begin(), data.end());\n\n cudaStreamSynchronize(s);\n\n Vector ref(5);\n ref[0] = 5;\n ref[1] = 4;\n ref[2] = 3;\n ref[3] = 2;\n ref[4] = 1;\n\n ASSERT_EQUAL(ref, data);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestReverseCudaStreams);\n\n\nvoid TestReverseCopyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n Vector data(5);\n data[0] = 1;\n data[1] = 2;\n data[2] = 3;\n data[3] = 4;\n data[4] = 5;\n\n Vector result(5);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::reverse_copy(thrust::cuda::par.on(s), data.begin(), data.end(), result.begin());\n\n cudaStreamSynchronize(s);\n\n Vector ref(5);\n ref[0] = 5;\n ref[1] = 4;\n ref[2] = 3;\n ref[3] = 2;\n ref[4] = 1;\n\n ASSERT_EQUAL(ref, result);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestReverseCopyCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid reverse_kernel(ExecutionPolicy exec, Iterator first, Iterator last)\n{\n thrust::reverse(exec, first, last);\n}\n\n\ntemplate\nvoid TestReverseDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n thrust::reverse(h_data.begin(), h_data.end());\n\n reverse_kernel<<<1,1>>>(exec, raw_pointer_cast(d_data.data()), raw_pointer_cast(d_data.data() + d_data.size()));\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_EQUAL(h_data, d_data);\n};\n\n\nvoid TestReverseDeviceSeq()\n{\n TestReverseDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReverseDeviceSeq);\n\n\nvoid TestReverseDeviceDevice()\n{\n TestReverseDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReverseDeviceDevice);\n\n\ntemplate\n__global__\nvoid reverse_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n thrust::reverse_copy(exec, first, last, result);\n}\n\n\ntemplate\nvoid TestReverseCopyDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n\n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n\n thrust::reverse_copy(h_data.begin(), h_data.end(), h_result.begin());\n\n reverse_copy_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n ASSERT_EQUAL(h_result, d_result);\n};\n\n\nvoid TestReverseCopyDeviceSeq()\n{\n TestReverseCopyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReverseCopyDeviceSeq);\n\n\nvoid TestReverseCopyDeviceDevice()\n{\n TestReverseCopyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReverseCopyDeviceDevice);\n#endif\n\n\nvoid TestReverseCudaStreams()\n{\n typedef thrust::device_vector Vector;\n Vector data(5);\n data[0] = 1;\n data[1] = 2;\n data[2] = 3;\n data[3] = 4;\n data[4] = 5;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::reverse(thrust::cuda::par.on(s), data.begin(), data.end());\n\n hipStreamSynchronize(s);\n\n Vector ref(5);\n ref[0] = 5;\n ref[1] = 4;\n ref[2] = 3;\n ref[3] = 2;\n ref[4] = 1;\n\n ASSERT_EQUAL(ref, data);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestReverseCudaStreams);\n\n\nvoid TestReverseCopyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n Vector data(5);\n data[0] = 1;\n data[1] = 2;\n data[2] = 3;\n data[3] = 4;\n data[4] = 5;\n\n Vector result(5);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::reverse_copy(thrust::cuda::par.on(s), data.begin(), data.end(), result.begin());\n\n hipStreamSynchronize(s);\n\n Vector ref(5);\n ref[0] = 5;\n ref[1] = 4;\n ref[2] = 3;\n ref[3] = 2;\n ref[4] = 1;\n\n ASSERT_EQUAL(ref, result);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestReverseCopyCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/reverse.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid uninitialized_fill_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T val)\n{\n thrust::uninitialized_fill(exec, first, last, val);\n}\n\n\ntemplate\nvoid TestUninitializedFillDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector v(5);\n v[0] = 0; v[1] = 1; v[2] = 2; v[3] = 3; v[4] = 4;\n \n T exemplar(7);\n \n uninitialized_fill_kernel<<<1,1>>>(exec, v.begin() + 1, v.begin() + 4, exemplar);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], 4);\n \n exemplar = 8;\n \n uninitialized_fill_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 3, exemplar);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(v[0], exemplar);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], 7);\n ASSERT_EQUAL(v[4], 4);\n \n exemplar = 9;\n \n uninitialized_fill_kernel<<<1,1>>>(exec, v.begin() + 2, v.end(), exemplar);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(v[0], 8);\n ASSERT_EQUAL(v[1], 8);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], 9);\n \n exemplar = 1;\n \n uninitialized_fill_kernel<<<1,1>>>(exec, v.begin(), v.end(), exemplar);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(v[0], exemplar);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], exemplar);\n}\n\n\nvoid TestUninitializedFillDeviceSeq()\n{\n TestUninitializedFillDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUninitializedFillDeviceSeq);\n\n\nvoid TestUninitializedFillDeviceDevice()\n{\n TestUninitializedFillDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUninitializedFillDeviceDevice);\n#endif\n\n\nvoid TestUninitializedFillCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector v(5);\n v[0] = 0; v[1] = 1; v[2] = 2; v[3] = 3; v[4] = 4;\n \n T exemplar(7);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n thrust::uninitialized_fill(thrust::cuda::par.on(s), v.begin(), v.end(), exemplar);\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(v[0], exemplar);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], exemplar);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestUninitializedFillCudaStreams);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid uninitialized_fill_n_kernel(ExecutionPolicy exec, Iterator1 first, Size n, T val, Iterator2 result)\n{\n *result = thrust::uninitialized_fill_n(exec, first, n, val);\n}\n\n\ntemplate\nvoid TestUninitializedFillNDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector v(5);\n v[0] = 0; v[1] = 1; v[2] = 2; v[3] = 3; v[4] = 4;\n \n T exemplar(7);\n\n thrust::device_vector iter_vec(1);\n \n uninitialized_fill_n_kernel<<<1,1>>>(exec, v.begin() + 1, 3, exemplar, iter_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n Vector::iterator iter = iter_vec[0];\n \n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], 4);\n ASSERT_EQUAL_QUIET(v.begin() + 4, iter);\n \n exemplar = 8;\n \n uninitialized_fill_n_kernel<<<1,1>>>(exec, v.begin() + 0, 3, exemplar, iter_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(v[0], exemplar);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], 7);\n ASSERT_EQUAL(v[4], 4);\n ASSERT_EQUAL_QUIET(v.begin() + 3, iter);\n \n exemplar = 9;\n \n uninitialized_fill_n_kernel<<<1,1>>>(exec, v.begin() + 2, 3, exemplar, iter_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(v[0], 8);\n ASSERT_EQUAL(v[1], 8);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], 9);\n ASSERT_EQUAL_QUIET(v.end(), iter);\n \n exemplar = 1;\n \n uninitialized_fill_n_kernel<<<1,1>>>(exec, v.begin(), v.size(), exemplar, iter_vec.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(v[0], exemplar);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], exemplar);\n ASSERT_EQUAL_QUIET(v.end(), iter);\n}\n\n\nvoid TestUninitializedFillNDeviceSeq()\n{\n TestUninitializedFillNDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUninitializedFillNDeviceSeq);\n\n\nvoid TestUninitializedFillNDeviceDevice()\n{\n TestUninitializedFillNDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUninitializedFillNDeviceDevice);\n#endif\n\n\nvoid TestUninitializedFillNCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector v(5);\n v[0] = 0; v[1] = 1; v[2] = 2; v[3] = 3; v[4] = 4;\n \n T exemplar(7);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n thrust::uninitialized_fill_n(thrust::cuda::par.on(s), v.begin(), v.size(), exemplar);\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(v[0], exemplar);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], exemplar);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestUninitializedFillNCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid uninitialized_fill_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T val)\n{\n thrust::uninitialized_fill(exec, first, last, val);\n}\n\n\ntemplate\nvoid TestUninitializedFillDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector v(5);\n v[0] = 0; v[1] = 1; v[2] = 2; v[3] = 3; v[4] = 4;\n \n T exemplar(7);\n \n uninitialized_fill_kernel<<<1,1>>>(exec, v.begin() + 1, v.begin() + 4, exemplar);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], 4);\n \n exemplar = 8;\n \n uninitialized_fill_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 3, exemplar);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(v[0], exemplar);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], 7);\n ASSERT_EQUAL(v[4], 4);\n \n exemplar = 9;\n \n uninitialized_fill_kernel<<<1,1>>>(exec, v.begin() + 2, v.end(), exemplar);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(v[0], 8);\n ASSERT_EQUAL(v[1], 8);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], 9);\n \n exemplar = 1;\n \n uninitialized_fill_kernel<<<1,1>>>(exec, v.begin(), v.end(), exemplar);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(v[0], exemplar);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], exemplar);\n}\n\n\nvoid TestUninitializedFillDeviceSeq()\n{\n TestUninitializedFillDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUninitializedFillDeviceSeq);\n\n\nvoid TestUninitializedFillDeviceDevice()\n{\n TestUninitializedFillDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUninitializedFillDeviceDevice);\n#endif\n\n\nvoid TestUninitializedFillCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector v(5);\n v[0] = 0; v[1] = 1; v[2] = 2; v[3] = 3; v[4] = 4;\n \n T exemplar(7);\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n thrust::uninitialized_fill(thrust::cuda::par.on(s), v.begin(), v.end(), exemplar);\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(v[0], exemplar);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], exemplar);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestUninitializedFillCudaStreams);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid uninitialized_fill_n_kernel(ExecutionPolicy exec, Iterator1 first, Size n, T val, Iterator2 result)\n{\n *result = thrust::uninitialized_fill_n(exec, first, n, val);\n}\n\n\ntemplate\nvoid TestUninitializedFillNDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector v(5);\n v[0] = 0; v[1] = 1; v[2] = 2; v[3] = 3; v[4] = 4;\n \n T exemplar(7);\n\n thrust::device_vector iter_vec(1);\n \n uninitialized_fill_n_kernel<<<1,1>>>(exec, v.begin() + 1, 3, exemplar, iter_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n Vector::iterator iter = iter_vec[0];\n \n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], 4);\n ASSERT_EQUAL_QUIET(v.begin() + 4, iter);\n \n exemplar = 8;\n \n uninitialized_fill_n_kernel<<<1,1>>>(exec, v.begin() + 0, 3, exemplar, iter_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(v[0], exemplar);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], 7);\n ASSERT_EQUAL(v[4], 4);\n ASSERT_EQUAL_QUIET(v.begin() + 3, iter);\n \n exemplar = 9;\n \n uninitialized_fill_n_kernel<<<1,1>>>(exec, v.begin() + 2, 3, exemplar, iter_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(v[0], 8);\n ASSERT_EQUAL(v[1], 8);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], 9);\n ASSERT_EQUAL_QUIET(v.end(), iter);\n \n exemplar = 1;\n \n uninitialized_fill_n_kernel<<<1,1>>>(exec, v.begin(), v.size(), exemplar, iter_vec.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n iter = iter_vec[0];\n \n ASSERT_EQUAL(v[0], exemplar);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], exemplar);\n ASSERT_EQUAL_QUIET(v.end(), iter);\n}\n\n\nvoid TestUninitializedFillNDeviceSeq()\n{\n TestUninitializedFillNDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUninitializedFillNDeviceSeq);\n\n\nvoid TestUninitializedFillNDeviceDevice()\n{\n TestUninitializedFillNDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUninitializedFillNDeviceDevice);\n#endif\n\n\nvoid TestUninitializedFillNCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector v(5);\n v[0] = 0; v[1] = 1; v[2] = 2; v[3] = 3; v[4] = 4;\n \n T exemplar(7);\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n thrust::uninitialized_fill_n(thrust::cuda::par.on(s), v.begin(), v.size(), exemplar);\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(v[0], exemplar);\n ASSERT_EQUAL(v[1], exemplar);\n ASSERT_EQUAL(v[2], exemplar);\n ASSERT_EQUAL(v[3], exemplar);\n ASSERT_EQUAL(v[4], exemplar);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestUninitializedFillNCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/uninitialized_fill.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid merge_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first1, Iterator1 keys_last1,\n Iterator2 keys_first2, Iterator2 keys_last2,\n Iterator3 values_first1,\n Iterator4 values_first2,\n Iterator5 keys_result,\n Iterator6 values_result,\n Iterator7 result)\n{\n *result = thrust::merge_by_key(exec, keys_first1, keys_last1, keys_first2, keys_last2, values_first1, values_first2, keys_result, values_result);\n}\n\n\ntemplate\nvoid TestMergeByKeyDevice(ExecutionPolicy exec)\n{\n thrust::device_vector a_key(3), a_val(3), b_key(4), b_val(4);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4;\n a_val[0] = 13; a_val[1] = 7; a_val[2] = 42;\n\n b_key[0] = 0 ; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4;\n b_val[0] = 42; b_val[1] = 42; b_val[2] = 7; b_val[3] = 13;\n\n thrust::device_vector ref_key(7), ref_val(7);\n ref_key[0] = 0; ref_val[0] = 13;\n ref_key[1] = 0; ref_val[1] = 42;\n ref_key[2] = 2; ref_val[2] = 7;\n ref_key[3] = 3; ref_val[3] = 42;\n ref_key[4] = 3; ref_val[4] = 7;\n ref_key[5] = 4; ref_val[5] = 42;\n ref_key[6] = 4; ref_val[6] = 13;\n\n thrust::device_vector result_key(7), result_val(7);\n\n typedef typename thrust::device_vector::iterator Iterator;\n\n thrust::device_vector > result_ends(1);\n\n merge_by_key_kernel<<<1,1>>>(exec,\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(), b_val.begin(),\n result_key.begin(),\n result_val.begin(),\n result_ends.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n thrust::pair ends = result_ends[0];\n\n ASSERT_EQUAL_QUIET(result_key.end(), ends.first);\n ASSERT_EQUAL_QUIET(result_val.end(), ends.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n}\n\n\nvoid TestMergeByKeyDeviceSeq()\n{\n TestMergeByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestMergeByKeyDeviceSeq);\n\n\nvoid TestMergeByKeyDeviceDevice()\n{\n TestMergeByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestMergeByKeyDeviceDevice);\n#endif\n\n\nvoid TestMergeByKeyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a_key(3), a_val(3), b_key(4), b_val(4);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4;\n a_val[0] = 13; a_val[1] = 7; a_val[2] = 42;\n\n b_key[0] = 0 ; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4;\n b_val[0] = 42; b_val[1] = 42; b_val[2] = 7; b_val[3] = 13;\n\n Vector ref_key(7), ref_val(7);\n ref_key[0] = 0; ref_val[0] = 13;\n ref_key[1] = 0; ref_val[1] = 42;\n ref_key[2] = 2; ref_val[2] = 7;\n ref_key[3] = 3; ref_val[3] = 42;\n ref_key[4] = 3; ref_val[4] = 7;\n ref_key[5] = 4; ref_val[5] = 42;\n ref_key[6] = 4; ref_val[6] = 13;\n\n Vector result_key(7), result_val(7);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::pair ends =\n thrust::merge_by_key(thrust::cuda::par.on(s),\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(), b_val.begin(),\n result_key.begin(),\n result_val.begin());\n\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result_key.end(), ends.first);\n ASSERT_EQUAL_QUIET(result_val.end(), ends.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestMergeByKeyCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid merge_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first1, Iterator1 keys_last1,\n Iterator2 keys_first2, Iterator2 keys_last2,\n Iterator3 values_first1,\n Iterator4 values_first2,\n Iterator5 keys_result,\n Iterator6 values_result,\n Iterator7 result)\n{\n *result = thrust::merge_by_key(exec, keys_first1, keys_last1, keys_first2, keys_last2, values_first1, values_first2, keys_result, values_result);\n}\n\n\ntemplate\nvoid TestMergeByKeyDevice(ExecutionPolicy exec)\n{\n thrust::device_vector a_key(3), a_val(3), b_key(4), b_val(4);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4;\n a_val[0] = 13; a_val[1] = 7; a_val[2] = 42;\n\n b_key[0] = 0 ; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4;\n b_val[0] = 42; b_val[1] = 42; b_val[2] = 7; b_val[3] = 13;\n\n thrust::device_vector ref_key(7), ref_val(7);\n ref_key[0] = 0; ref_val[0] = 13;\n ref_key[1] = 0; ref_val[1] = 42;\n ref_key[2] = 2; ref_val[2] = 7;\n ref_key[3] = 3; ref_val[3] = 42;\n ref_key[4] = 3; ref_val[4] = 7;\n ref_key[5] = 4; ref_val[5] = 42;\n ref_key[6] = 4; ref_val[6] = 13;\n\n thrust::device_vector result_key(7), result_val(7);\n\n typedef typename thrust::device_vector::iterator Iterator;\n\n thrust::device_vector > result_ends(1);\n\n merge_by_key_kernel<<<1,1>>>(exec,\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(), b_val.begin(),\n result_key.begin(),\n result_val.begin(),\n result_ends.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n thrust::pair ends = result_ends[0];\n\n ASSERT_EQUAL_QUIET(result_key.end(), ends.first);\n ASSERT_EQUAL_QUIET(result_val.end(), ends.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n}\n\n\nvoid TestMergeByKeyDeviceSeq()\n{\n TestMergeByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestMergeByKeyDeviceSeq);\n\n\nvoid TestMergeByKeyDeviceDevice()\n{\n TestMergeByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestMergeByKeyDeviceDevice);\n#endif\n\n\nvoid TestMergeByKeyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a_key(3), a_val(3), b_key(4), b_val(4);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4;\n a_val[0] = 13; a_val[1] = 7; a_val[2] = 42;\n\n b_key[0] = 0 ; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4;\n b_val[0] = 42; b_val[1] = 42; b_val[2] = 7; b_val[3] = 13;\n\n Vector ref_key(7), ref_val(7);\n ref_key[0] = 0; ref_val[0] = 13;\n ref_key[1] = 0; ref_val[1] = 42;\n ref_key[2] = 2; ref_val[2] = 7;\n ref_key[3] = 3; ref_val[3] = 42;\n ref_key[4] = 3; ref_val[4] = 7;\n ref_key[5] = 4; ref_val[5] = 42;\n ref_key[6] = 4; ref_val[6] = 13;\n\n Vector result_key(7), result_val(7);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::pair ends =\n thrust::merge_by_key(thrust::cuda::par.on(s),\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(), b_val.begin(),\n result_key.begin(),\n result_val.begin());\n\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result_key.end(), ends.first);\n ASSERT_EQUAL_QUIET(result_val.end(), ends.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestMergeByKeyCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/merge_by_key.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_union_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first1, Iterator1 keys_last1,\n Iterator2 keys_first2, Iterator2 keys_last2,\n Iterator3 values_first1,\n Iterator4 values_first2,\n Iterator5 keys_result,\n Iterator6 values_result,\n Iterator7 result)\n{\n *result = thrust::set_union_by_key(exec, keys_first1, keys_last1, keys_first2, keys_last2, values_first1, values_first2, keys_result, values_result);\n}\n\n\ntemplate\nvoid TestSetUnionByKeyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::iterator Iterator;\n\n Vector a_key(3), b_key(4);\n Vector a_val(3), b_val(4);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4;\n b_val[0] = 1; b_val[1] = 1; b_val[2] = 1; b_val[3] = 1;\n\n Vector ref_key(5), ref_val(5);\n ref_key[0] = 0; ref_key[1] = 2; ref_key[2] = 3; ref_key[3] = 3; ref_key[4] = 4;\n ref_val[0] = 0; ref_val[1] = 0; ref_val[2] = 1; ref_val[3] = 1; ref_val[4] = 0;\n\n Vector result_key(5), result_val(5);\n\n thrust::device_vector > end_vec(1);\n\n set_union_by_key_kernel<<<1,1>>>(exec,\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n b_val.begin(),\n result_key.begin(),\n result_val.begin(),\n end_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n thrust::pair end = end_vec[0];\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n}\n\n\nvoid TestSetUnionByKeyDeviceSeq()\n{\n TestSetUnionByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetUnionByKeyDeviceSeq);\n\n\nvoid TestSetUnionByKeyDeviceDevice()\n{\n TestSetUnionByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetUnionByKeyDeviceDevice);\n#endif\n\n\nvoid TestSetUnionByKeyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a_key(3), b_key(4);\n Vector a_val(3), b_val(4);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4;\n b_val[0] = 1; b_val[1] = 1; b_val[2] = 1; b_val[3] = 1;\n\n Vector ref_key(5), ref_val(5);\n ref_key[0] = 0; ref_key[1] = 2; ref_key[2] = 3; ref_key[3] = 3; ref_key[4] = 4;\n ref_val[0] = 0; ref_val[1] = 0; ref_val[2] = 1; ref_val[3] = 1; ref_val[4] = 0;\n\n Vector result_key(5), result_val(5);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::pair end =\n thrust::set_union_by_key(thrust::cuda::par.on(s),\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n b_val.begin(),\n result_key.begin(),\n result_val.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSetUnionByKeyCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_union_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first1, Iterator1 keys_last1,\n Iterator2 keys_first2, Iterator2 keys_last2,\n Iterator3 values_first1,\n Iterator4 values_first2,\n Iterator5 keys_result,\n Iterator6 values_result,\n Iterator7 result)\n{\n *result = thrust::set_union_by_key(exec, keys_first1, keys_last1, keys_first2, keys_last2, values_first1, values_first2, keys_result, values_result);\n}\n\n\ntemplate\nvoid TestSetUnionByKeyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::iterator Iterator;\n\n Vector a_key(3), b_key(4);\n Vector a_val(3), b_val(4);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4;\n b_val[0] = 1; b_val[1] = 1; b_val[2] = 1; b_val[3] = 1;\n\n Vector ref_key(5), ref_val(5);\n ref_key[0] = 0; ref_key[1] = 2; ref_key[2] = 3; ref_key[3] = 3; ref_key[4] = 4;\n ref_val[0] = 0; ref_val[1] = 0; ref_val[2] = 1; ref_val[3] = 1; ref_val[4] = 0;\n\n Vector result_key(5), result_val(5);\n\n thrust::device_vector > end_vec(1);\n\n set_union_by_key_kernel<<<1,1>>>(exec,\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n b_val.begin(),\n result_key.begin(),\n result_val.begin(),\n end_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n thrust::pair end = end_vec[0];\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n}\n\n\nvoid TestSetUnionByKeyDeviceSeq()\n{\n TestSetUnionByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetUnionByKeyDeviceSeq);\n\n\nvoid TestSetUnionByKeyDeviceDevice()\n{\n TestSetUnionByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetUnionByKeyDeviceDevice);\n#endif\n\n\nvoid TestSetUnionByKeyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a_key(3), b_key(4);\n Vector a_val(3), b_val(4);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4;\n b_val[0] = 1; b_val[1] = 1; b_val[2] = 1; b_val[3] = 1;\n\n Vector ref_key(5), ref_val(5);\n ref_key[0] = 0; ref_key[1] = 2; ref_key[2] = 3; ref_key[3] = 3; ref_key[4] = 4;\n ref_val[0] = 0; ref_val[1] = 0; ref_val[2] = 1; ref_val[3] = 1; ref_val[4] = 0;\n\n Vector result_key(5), result_val(5);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::pair end =\n thrust::set_union_by_key(thrust::cuda::par.on(s),\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n b_val.begin(),\n result_key.begin(),\n result_val.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSetUnionByKeyCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/set_union_by_key.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid min_element_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Iterator2 result)\n{\n *result = thrust::min_element(exec, first, last);\n}\n\n\ntemplate\n__global__\nvoid min_element_kernel(ExecutionPolicy exec, Iterator first, Iterator last, BinaryPredicate pred, Iterator2 result)\n{\n *result = thrust::min_element(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestMinElementDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n\n typedef typename thrust::device_vector::iterator iter_type;\n\n thrust::device_vector d_result(1);\n \n typename thrust::host_vector::iterator h_min = thrust::min_element(h_data.begin(), h_data.end());\n\n min_element_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(h_min - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n\n typename thrust::host_vector::iterator h_max = thrust::min_element(h_data.begin(), h_data.end(), thrust::greater());\n\n min_element_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), thrust::greater(), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(h_max - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n}\n\n\nvoid TestMinElementDeviceSeq()\n{\n TestMinElementDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestMinElementDeviceSeq);\n\n\nvoid TestMinElementDeviceDevice()\n{\n TestMinElementDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestMinElementDeviceDevice);\n#endif\n\n\nvoid TestMinElementCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(6);\n data[0] = 3;\n data[1] = 5;\n data[2] = 1;\n data[3] = 2;\n data[4] = 5;\n data[5] = 1;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n ASSERT_EQUAL( *thrust::min_element(thrust::cuda::par.on(s), data.begin(), data.end()), 1);\n ASSERT_EQUAL( thrust::min_element(thrust::cuda::par.on(s), data.begin(), data.end()) - data.begin(), 2);\n \n ASSERT_EQUAL( *thrust::min_element(thrust::cuda::par.on(s), data.begin(), data.end(), thrust::greater()), 5);\n ASSERT_EQUAL( thrust::min_element(thrust::cuda::par.on(s), data.begin(), data.end(), thrust::greater()) - data.begin(), 1);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestMinElementCudaStreams);\n\nvoid TestMinElementDevicePointer()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(6);\n data[0] = 3;\n data[1] = 5;\n data[2] = 1;\n data[3] = 2;\n data[4] = 5;\n data[5] = 1;\n\n T* raw_ptr = thrust::raw_pointer_cast(data.data());\n size_t n = data.size();\n ASSERT_EQUAL( thrust::min_element(thrust::device, raw_ptr, raw_ptr+n) - raw_ptr, 2);\n ASSERT_EQUAL( thrust::min_element(thrust::device, raw_ptr, raw_ptr+n, thrust::greater()) - raw_ptr, 1);\n}\nDECLARE_UNITTEST(TestMinElementDevicePointer);\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid min_element_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Iterator2 result)\n{\n *result = thrust::min_element(exec, first, last);\n}\n\n\ntemplate\n__global__\nvoid min_element_kernel(ExecutionPolicy exec, Iterator first, Iterator last, BinaryPredicate pred, Iterator2 result)\n{\n *result = thrust::min_element(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestMinElementDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n\n typedef typename thrust::device_vector::iterator iter_type;\n\n thrust::device_vector d_result(1);\n \n typename thrust::host_vector::iterator h_min = thrust::min_element(h_data.begin(), h_data.end());\n\n min_element_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(h_min - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n\n typename thrust::host_vector::iterator h_max = thrust::min_element(h_data.begin(), h_data.end(), thrust::greater());\n\n min_element_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), thrust::greater(), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(h_max - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n}\n\n\nvoid TestMinElementDeviceSeq()\n{\n TestMinElementDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestMinElementDeviceSeq);\n\n\nvoid TestMinElementDeviceDevice()\n{\n TestMinElementDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestMinElementDeviceDevice);\n#endif\n\n\nvoid TestMinElementCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(6);\n data[0] = 3;\n data[1] = 5;\n data[2] = 1;\n data[3] = 2;\n data[4] = 5;\n data[5] = 1;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n ASSERT_EQUAL( *thrust::min_element(thrust::cuda::par.on(s), data.begin(), data.end()), 1);\n ASSERT_EQUAL( thrust::min_element(thrust::cuda::par.on(s), data.begin(), data.end()) - data.begin(), 2);\n \n ASSERT_EQUAL( *thrust::min_element(thrust::cuda::par.on(s), data.begin(), data.end(), thrust::greater()), 5);\n ASSERT_EQUAL( thrust::min_element(thrust::cuda::par.on(s), data.begin(), data.end(), thrust::greater()) - data.begin(), 1);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestMinElementCudaStreams);\n\nvoid TestMinElementDevicePointer()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(6);\n data[0] = 3;\n data[1] = 5;\n data[2] = 1;\n data[3] = 2;\n data[4] = 5;\n data[5] = 1;\n\n T* raw_ptr = thrust::raw_pointer_cast(data.data());\n size_t n = data.size();\n ASSERT_EQUAL( thrust::min_element(thrust::device, raw_ptr, raw_ptr+n) - raw_ptr, 2);\n ASSERT_EQUAL( thrust::min_element(thrust::device, raw_ptr, raw_ptr+n, thrust::greater()) - raw_ptr, 1);\n}\nDECLARE_UNITTEST(TestMinElementDevicePointer);\n", "source": "/kaggle/working/thrust/testing/cuda/min_element.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_intersection_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1,\n Iterator2 first2, Iterator2 last2,\n Iterator3 result1,\n Iterator4 result2)\n{\n *result2 = thrust::set_intersection(exec, first1, last1, first2, last2, result1);\n}\n\n\ntemplate\nvoid TestSetIntersectionDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a(3), b(4);\n\n a[0] = 0; a[1] = 2; a[2] = 4;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4;\n\n Vector ref(2);\n ref[0] = 0; ref[1] = 4;\n\n Vector result(2);\n thrust::device_vector end_vec(1);\n\n set_intersection_kernel<<<1,1>>>(exec, a.begin(), a.end(), b.begin(), b.end(), result.begin(), end_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n Iterator end = end_vec.front();\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n}\n\n\nvoid TestSetIntersectionDeviceSeq()\n{\n TestSetIntersectionDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetIntersectionDeviceSeq);\n\n\nvoid TestSetIntersectionDeviceDevice()\n{\n TestSetIntersectionDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetIntersectionDeviceDevice);\n\n\nvoid TestSetIntersectionDeviceNoSync()\n{\n TestSetIntersectionDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestSetIntersectionDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestSetIntersectionCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a(3), b(4);\n\n a[0] = 0; a[1] = 2; a[2] = 4;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4;\n\n Vector ref(2);\n ref[0] = 0; ref[1] = 4;\n\n Vector result(2);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n\n Iterator end = thrust::set_intersection(streampolicy,\n a.begin(), a.end(),\n b.begin(), b.end(),\n result.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n\n cudaStreamDestroy(s);\n}\n\nvoid TestSetIntersectionCudaStreamsSync()\n{\n TestSetIntersectionCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestSetIntersectionCudaStreamsSync);\n\n\nvoid TestSetIntersectionCudaStreamsNoSync()\n{\n TestSetIntersectionCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestSetIntersectionCudaStreamsNoSync);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_intersection_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1,\n Iterator2 first2, Iterator2 last2,\n Iterator3 result1,\n Iterator4 result2)\n{\n *result2 = thrust::set_intersection(exec, first1, last1, first2, last2, result1);\n}\n\n\ntemplate\nvoid TestSetIntersectionDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a(3), b(4);\n\n a[0] = 0; a[1] = 2; a[2] = 4;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4;\n\n Vector ref(2);\n ref[0] = 0; ref[1] = 4;\n\n Vector result(2);\n thrust::device_vector end_vec(1);\n\n set_intersection_kernel<<<1,1>>>(exec, a.begin(), a.end(), b.begin(), b.end(), result.begin(), end_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n Iterator end = end_vec.front();\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n}\n\n\nvoid TestSetIntersectionDeviceSeq()\n{\n TestSetIntersectionDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetIntersectionDeviceSeq);\n\n\nvoid TestSetIntersectionDeviceDevice()\n{\n TestSetIntersectionDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetIntersectionDeviceDevice);\n\n\nvoid TestSetIntersectionDeviceNoSync()\n{\n TestSetIntersectionDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestSetIntersectionDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestSetIntersectionCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a(3), b(4);\n\n a[0] = 0; a[1] = 2; a[2] = 4;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4;\n\n Vector ref(2);\n ref[0] = 0; ref[1] = 4;\n\n Vector result(2);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n\n Iterator end = thrust::set_intersection(streampolicy,\n a.begin(), a.end(),\n b.begin(), b.end(),\n result.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n\n hipStreamDestroy(s);\n}\n\nvoid TestSetIntersectionCudaStreamsSync()\n{\n TestSetIntersectionCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestSetIntersectionCudaStreamsSync);\n\n\nvoid TestSetIntersectionCudaStreamsNoSync()\n{\n TestSetIntersectionCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestSetIntersectionCudaStreamsNoSync);\n\n", "source": "/kaggle/working/thrust/testing/cuda/set_intersection.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_intersection_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first1, Iterator1 keys_last1,\n Iterator2 keys_first2, Iterator2 keys_last2,\n Iterator3 values_first1,\n Iterator4 keys_result,\n Iterator5 values_result,\n Iterator6 result)\n{\n *result = thrust::set_intersection_by_key(exec, keys_first1, keys_last1, keys_first2, keys_last2, values_first1, keys_result, values_result);\n}\n\n\ntemplate\nvoid TestSetIntersectionByKeyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::iterator Iterator;\n\n Vector a_key(3), b_key(4);\n Vector a_val(3);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4;\n\n Vector ref_key(2), ref_val(2);\n ref_key[0] = 0; ref_key[1] = 4;\n ref_val[0] = 0; ref_val[1] = 0;\n\n Vector result_key(2), result_val(2);\n\n typedef thrust::pair iter_pair;\n thrust::device_vector end_vec(1);\n\n set_intersection_by_key_kernel<<<1,1>>>(exec,\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n result_key.begin(),\n result_val.begin(),\n end_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n thrust::pair end = end_vec.front();\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n}\n\n\nvoid TestSetIntersectionByKeyDeviceSeq()\n{\n TestSetIntersectionByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetIntersectionByKeyDeviceSeq);\n\n\nvoid TestSetIntersectionByKeyDeviceDevice()\n{\n TestSetIntersectionByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetIntersectionByKeyDeviceDevice);\n\n\nvoid TestSetIntersectionByKeyDeviceNoSync()\n{\n TestSetIntersectionByKeyDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestSetIntersectionByKeyDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestSetIntersectionByKeyCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a_key(3), b_key(4);\n Vector a_val(3);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4;\n\n Vector ref_key(2), ref_val(2);\n ref_key[0] = 0; ref_key[1] = 4;\n ref_val[0] = 0; ref_val[1] = 0;\n\n Vector result_key(2), result_val(2);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n\n thrust::pair end =\n thrust::set_intersection_by_key(streampolicy,\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n result_key.begin(),\n result_val.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n\n cudaStreamDestroy(s);\n}\n\nvoid TestSetIntersectionByKeyCudaStreamsSync()\n{\n TestSetIntersectionByKeyCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestSetIntersectionByKeyCudaStreamsSync);\n\n\nvoid TestSetIntersectionByKeyCudaStreamsNoSync()\n{\n TestSetIntersectionByKeyCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestSetIntersectionByKeyCudaStreamsNoSync);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_intersection_by_key_kernel(ExecutionPolicy exec,\n Iterator1 keys_first1, Iterator1 keys_last1,\n Iterator2 keys_first2, Iterator2 keys_last2,\n Iterator3 values_first1,\n Iterator4 keys_result,\n Iterator5 values_result,\n Iterator6 result)\n{\n *result = thrust::set_intersection_by_key(exec, keys_first1, keys_last1, keys_first2, keys_last2, values_first1, keys_result, values_result);\n}\n\n\ntemplate\nvoid TestSetIntersectionByKeyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::iterator Iterator;\n\n Vector a_key(3), b_key(4);\n Vector a_val(3);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4;\n\n Vector ref_key(2), ref_val(2);\n ref_key[0] = 0; ref_key[1] = 4;\n ref_val[0] = 0; ref_val[1] = 0;\n\n Vector result_key(2), result_val(2);\n\n typedef thrust::pair iter_pair;\n thrust::device_vector end_vec(1);\n\n set_intersection_by_key_kernel<<<1,1>>>(exec,\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n result_key.begin(),\n result_val.begin(),\n end_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n thrust::pair end = end_vec.front();\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n}\n\n\nvoid TestSetIntersectionByKeyDeviceSeq()\n{\n TestSetIntersectionByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetIntersectionByKeyDeviceSeq);\n\n\nvoid TestSetIntersectionByKeyDeviceDevice()\n{\n TestSetIntersectionByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetIntersectionByKeyDeviceDevice);\n\n\nvoid TestSetIntersectionByKeyDeviceNoSync()\n{\n TestSetIntersectionByKeyDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestSetIntersectionByKeyDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestSetIntersectionByKeyCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a_key(3), b_key(4);\n Vector a_val(3);\n\n a_key[0] = 0; a_key[1] = 2; a_key[2] = 4;\n a_val[0] = 0; a_val[1] = 0; a_val[2] = 0;\n\n b_key[0] = 0; b_key[1] = 3; b_key[2] = 3; b_key[3] = 4;\n\n Vector ref_key(2), ref_val(2);\n ref_key[0] = 0; ref_key[1] = 4;\n ref_val[0] = 0; ref_val[1] = 0;\n\n Vector result_key(2), result_val(2);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n\n thrust::pair end =\n thrust::set_intersection_by_key(streampolicy,\n a_key.begin(), a_key.end(),\n b_key.begin(), b_key.end(),\n a_val.begin(),\n result_key.begin(),\n result_val.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result_key.end(), end.first);\n ASSERT_EQUAL_QUIET(result_val.end(), end.second);\n ASSERT_EQUAL(ref_key, result_key);\n ASSERT_EQUAL(ref_val, result_val);\n\n hipStreamDestroy(s);\n}\n\nvoid TestSetIntersectionByKeyCudaStreamsSync()\n{\n TestSetIntersectionByKeyCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestSetIntersectionByKeyCudaStreamsSync);\n\n\nvoid TestSetIntersectionByKeyCudaStreamsNoSync()\n{\n TestSetIntersectionByKeyCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestSetIntersectionByKeyCudaStreamsNoSync);\n\n", "source": "/kaggle/working/thrust/testing/cuda/set_intersection_by_key.cu"} {"input": "#include \n#include \n#include \n#include \n\n\ntemplate\n__global__\nvoid reduce_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T init, Iterator2 result)\n{\n *result = thrust::reduce(exec, first, last, init);\n}\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\nvoid TestReduceDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n thrust::device_vector d_result(1);\n \n T init = 13;\n \n T h_result = thrust::reduce(h_data.begin(), h_data.end(), init);\n \n reduce_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), init, d_result.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_EQUAL(h_result, d_result[0]);\n}\n\n\ntemplate\nstruct TestReduceDeviceSeq\n{\n void operator()(const size_t n)\n {\n TestReduceDevice(thrust::seq, n);\n }\n};\nVariableUnitTest TestReduceDeviceSeqInstance;\n\n\ntemplate\nstruct TestReduceDeviceDevice\n{\n void operator()(const size_t n)\n {\n TestReduceDevice(thrust::device, n);\n }\n};\nVariableUnitTest TestReduceDeviceDeviceInstance;\n\n\ntemplate\nstruct TestReduceDeviceNoSync\n{\n void operator()(const size_t n)\n {\n TestReduceDevice(thrust::cuda::par_nosync, n);\n }\n};\nVariableUnitTest TestReduceDeviceNoSyncInstance;\n#endif\n\n\ntemplate\nvoid TestReduceCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n\n Vector v(3);\n v[0] = 1; v[1] = -2; v[2] = 3;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n\n // no initializer\n ASSERT_EQUAL(thrust::reduce(streampolicy, v.begin(), v.end()), 2);\n\n // with initializer\n ASSERT_EQUAL(thrust::reduce(streampolicy, v.begin(), v.end(), 10), 12);\n\n cudaStreamDestroy(s);\n}\n\nvoid TestReduceCudaStreamsSync()\n{\n TestReduceCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestReduceCudaStreamsSync);\n\n\nvoid TestReduceCudaStreamsNoSync()\n{\n TestReduceCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestReduceCudaStreamsNoSync);\n\n#if defined(THRUST_RDC_ENABLED)\nvoid TestReduceLargeInput()\n{\n using T = unsigned long long;\n using OffsetT = std::size_t;\n const OffsetT num_items = 1ull << 32;\n\n thrust::constant_iterator d_data(T{1});\n thrust::device_vector d_result(1);\n\n reduce_kernel<<<1,1>>>(thrust::device, d_data, d_data + num_items, T{}, d_result.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_EQUAL(num_items, d_result[0]);\n}\nDECLARE_UNITTEST(TestReduceLargeInput);\n#endif\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n\ntemplate\n__global__\nvoid reduce_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T init, Iterator2 result)\n{\n *result = thrust::reduce(exec, first, last, init);\n}\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\nvoid TestReduceDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n thrust::device_vector d_result(1);\n \n T init = 13;\n \n T h_result = thrust::reduce(h_data.begin(), h_data.end(), init);\n \n reduce_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), init, d_result.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_EQUAL(h_result, d_result[0]);\n}\n\n\ntemplate\nstruct TestReduceDeviceSeq\n{\n void operator()(const size_t n)\n {\n TestReduceDevice(thrust::seq, n);\n }\n};\nVariableUnitTest TestReduceDeviceSeqInstance;\n\n\ntemplate\nstruct TestReduceDeviceDevice\n{\n void operator()(const size_t n)\n {\n TestReduceDevice(thrust::device, n);\n }\n};\nVariableUnitTest TestReduceDeviceDeviceInstance;\n\n\ntemplate\nstruct TestReduceDeviceNoSync\n{\n void operator()(const size_t n)\n {\n TestReduceDevice(thrust::cuda::par_nosync, n);\n }\n};\nVariableUnitTest TestReduceDeviceNoSyncInstance;\n#endif\n\n\ntemplate\nvoid TestReduceCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n\n Vector v(3);\n v[0] = 1; v[1] = -2; v[2] = 3;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n\n // no initializer\n ASSERT_EQUAL(thrust::reduce(streampolicy, v.begin(), v.end()), 2);\n\n // with initializer\n ASSERT_EQUAL(thrust::reduce(streampolicy, v.begin(), v.end(), 10), 12);\n\n hipStreamDestroy(s);\n}\n\nvoid TestReduceCudaStreamsSync()\n{\n TestReduceCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestReduceCudaStreamsSync);\n\n\nvoid TestReduceCudaStreamsNoSync()\n{\n TestReduceCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestReduceCudaStreamsNoSync);\n\n#if defined(THRUST_RDC_ENABLED)\nvoid TestReduceLargeInput()\n{\n using T = unsigned long long;\n using OffsetT = std::size_t;\n const OffsetT num_items = 1ull << 32;\n\n thrust::constant_iterator d_data(T{1});\n thrust::device_vector d_result(1);\n\n reduce_kernel<<<1,1>>>(thrust::device, d_data, d_data + num_items, T{}, d_result.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_EQUAL(num_items, d_result[0]);\n}\nDECLARE_UNITTEST(TestReduceLargeInput);\n#endif\n\n", "source": "/kaggle/working/thrust/testing/cuda/reduce.cu"} {"input": "#include \n#include \n#include \n\n\ntemplate\nstruct return_value\n{\n T val;\n \n return_value(void){}\n return_value(T v):val(v){}\n \n __host__ __device__\n T operator()(void){ return val; }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid generate_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Function f)\n{\n thrust::generate(exec, first, last, f);\n}\n\n\ntemplate\nvoid TestGenerateDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n \n T value = 13;\n return_value f(value);\n \n thrust::generate(h_result.begin(), h_result.end(), f);\n\n generate_kernel<<<1,1>>>(exec, d_result.begin(), d_result.end(), f);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_result, d_result);\n}\n\n\ntemplate\nvoid TestGenerateDeviceSeq(const size_t n)\n{\n TestGenerateDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGenerateDeviceSeq);\n\n\ntemplate\nvoid TestGenerateDeviceDevice(const size_t n)\n{\n TestGenerateDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGenerateDeviceDevice);\n#endif\n\n\nvoid TestGenerateCudaStreams()\n{\n thrust::device_vector result(5);\n \n int value = 13;\n \n return_value f(value);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n thrust::generate(thrust::cuda::par.on(s), result.begin(), result.end(), f);\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(result[0], value);\n ASSERT_EQUAL(result[1], value);\n ASSERT_EQUAL(result[2], value);\n ASSERT_EQUAL(result[3], value);\n ASSERT_EQUAL(result[4], value);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestGenerateCudaStreams);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid generate_n_kernel(ExecutionPolicy exec, Iterator first, Size n, Function f)\n{\n thrust::generate_n(exec, first, n, f);\n}\n\n\ntemplate\nvoid TestGenerateNDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n \n T value = 13;\n return_value f(value);\n \n thrust::generate_n(h_result.begin(), h_result.size(), f);\n\n generate_n_kernel<<<1,1>>>(exec, d_result.begin(), d_result.size(), f);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_result, d_result);\n}\n\n\ntemplate\nvoid TestGenerateNDeviceSeq(const size_t n)\n{\n TestGenerateNDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGenerateNDeviceSeq);\n\n\ntemplate\nvoid TestGenerateNDeviceDevice(const size_t n)\n{\n TestGenerateNDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGenerateNDeviceDevice);\n#endif\n\n\nvoid TestGenerateNCudaStreams()\n{\n thrust::device_vector result(5);\n \n int value = 13;\n \n return_value f(value);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n thrust::generate_n(thrust::cuda::par.on(s), result.begin(), result.size(), f);\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(result[0], value);\n ASSERT_EQUAL(result[1], value);\n ASSERT_EQUAL(result[2], value);\n ASSERT_EQUAL(result[3], value);\n ASSERT_EQUAL(result[4], value);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestGenerateNCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\ntemplate\nstruct return_value\n{\n T val;\n \n return_value(void){}\n return_value(T v):val(v){}\n \n __host__ __device__\n T operator()(void){ return val; }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid generate_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Function f)\n{\n thrust::generate(exec, first, last, f);\n}\n\n\ntemplate\nvoid TestGenerateDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n \n T value = 13;\n return_value f(value);\n \n thrust::generate(h_result.begin(), h_result.end(), f);\n\n generate_kernel<<<1,1>>>(exec, d_result.begin(), d_result.end(), f);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_result, d_result);\n}\n\n\ntemplate\nvoid TestGenerateDeviceSeq(const size_t n)\n{\n TestGenerateDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGenerateDeviceSeq);\n\n\ntemplate\nvoid TestGenerateDeviceDevice(const size_t n)\n{\n TestGenerateDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGenerateDeviceDevice);\n#endif\n\n\nvoid TestGenerateCudaStreams()\n{\n thrust::device_vector result(5);\n \n int value = 13;\n \n return_value f(value);\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n thrust::generate(thrust::cuda::par.on(s), result.begin(), result.end(), f);\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(result[0], value);\n ASSERT_EQUAL(result[1], value);\n ASSERT_EQUAL(result[2], value);\n ASSERT_EQUAL(result[3], value);\n ASSERT_EQUAL(result[4], value);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestGenerateCudaStreams);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid generate_n_kernel(ExecutionPolicy exec, Iterator first, Size n, Function f)\n{\n thrust::generate_n(exec, first, n, f);\n}\n\n\ntemplate\nvoid TestGenerateNDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n \n T value = 13;\n return_value f(value);\n \n thrust::generate_n(h_result.begin(), h_result.size(), f);\n\n generate_n_kernel<<<1,1>>>(exec, d_result.begin(), d_result.size(), f);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_result, d_result);\n}\n\n\ntemplate\nvoid TestGenerateNDeviceSeq(const size_t n)\n{\n TestGenerateNDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGenerateNDeviceSeq);\n\n\ntemplate\nvoid TestGenerateNDeviceDevice(const size_t n)\n{\n TestGenerateNDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGenerateNDeviceDevice);\n#endif\n\n\nvoid TestGenerateNCudaStreams()\n{\n thrust::device_vector result(5);\n \n int value = 13;\n \n return_value f(value);\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n thrust::generate_n(thrust::cuda::par.on(s), result.begin(), result.size(), f);\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(result[0], value);\n ASSERT_EQUAL(result[1], value);\n ASSERT_EQUAL(result[2], value);\n ASSERT_EQUAL(result[3], value);\n ASSERT_EQUAL(result[4], value);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestGenerateNCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/generate.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid uninitialized_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n thrust::uninitialized_copy(exec, first, last, result);\n}\n\n\ntemplate\nvoid TestUninitializedCopyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n \n Vector v1(5);\n v1[0] = 0; v1[1] = 1; v1[2] = 2; v1[3] = 3; v1[4] = 4;\n \n // copy to Vector\n Vector v2(5);\n uninitialized_copy_kernel<<<1,1>>>(exec, v1.begin(), v1.end(), v2.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n ASSERT_EQUAL(v2[0], 0);\n ASSERT_EQUAL(v2[1], 1);\n ASSERT_EQUAL(v2[2], 2);\n ASSERT_EQUAL(v2[3], 3);\n ASSERT_EQUAL(v2[4], 4);\n}\n\n\nvoid TestUninitializedCopyDeviceSeq()\n{\n TestUninitializedCopyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUninitializedCopyDeviceSeq);\n\n\nvoid TestUninitializedCopyDeviceDevice()\n{\n TestUninitializedCopyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUninitializedCopyDeviceDevice);\n#endif\n\n\nvoid TestUninitializedCopyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n \n Vector v1(5);\n v1[0] = 0; v1[1] = 1; v1[2] = 2; v1[3] = 3; v1[4] = 4;\n \n // copy to Vector\n Vector v2(5);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::uninitialized_copy(thrust::cuda::par.on(s), v1.begin(), v1.end(), v2.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(v2[0], 0);\n ASSERT_EQUAL(v2[1], 1);\n ASSERT_EQUAL(v2[2], 2);\n ASSERT_EQUAL(v2[3], 3);\n ASSERT_EQUAL(v2[4], 4);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestUninitializedCopyCudaStreams);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid uninitialized_copy_n_kernel(ExecutionPolicy exec, Iterator1 first, Size n, Iterator2 result)\n{\n thrust::uninitialized_copy_n(exec, first, n, result);\n}\n\n\ntemplate\nvoid TestUninitializedCopyNDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n \n Vector v1(5);\n v1[0] = 0; v1[1] = 1; v1[2] = 2; v1[3] = 3; v1[4] = 4;\n \n // copy to Vector\n Vector v2(5);\n uninitialized_copy_n_kernel<<<1,1>>>(exec, v1.begin(), v1.size(), v2.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n ASSERT_EQUAL(v2[0], 0);\n ASSERT_EQUAL(v2[1], 1);\n ASSERT_EQUAL(v2[2], 2);\n ASSERT_EQUAL(v2[3], 3);\n ASSERT_EQUAL(v2[4], 4);\n}\n\n\nvoid TestUninitializedCopyNDeviceSeq()\n{\n TestUninitializedCopyNDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUninitializedCopyNDeviceSeq);\n\n\nvoid TestUninitializedCopyNDeviceDevice()\n{\n TestUninitializedCopyNDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUninitializedCopyNDeviceDevice);\n#endif\n\n\nvoid TestUninitializedCopyNCudaStreams()\n{\n typedef thrust::device_vector Vector;\n \n Vector v1(5);\n v1[0] = 0; v1[1] = 1; v1[2] = 2; v1[3] = 3; v1[4] = 4;\n \n // copy to Vector\n Vector v2(5);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::uninitialized_copy_n(thrust::cuda::par.on(s), v1.begin(), v1.size(), v2.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(v2[0], 0);\n ASSERT_EQUAL(v2[1], 1);\n ASSERT_EQUAL(v2[2], 2);\n ASSERT_EQUAL(v2[3], 3);\n ASSERT_EQUAL(v2[4], 4);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestUninitializedCopyNCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid uninitialized_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n thrust::uninitialized_copy(exec, first, last, result);\n}\n\n\ntemplate\nvoid TestUninitializedCopyDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n \n Vector v1(5);\n v1[0] = 0; v1[1] = 1; v1[2] = 2; v1[3] = 3; v1[4] = 4;\n \n // copy to Vector\n Vector v2(5);\n uninitialized_copy_kernel<<<1,1>>>(exec, v1.begin(), v1.end(), v2.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n ASSERT_EQUAL(v2[0], 0);\n ASSERT_EQUAL(v2[1], 1);\n ASSERT_EQUAL(v2[2], 2);\n ASSERT_EQUAL(v2[3], 3);\n ASSERT_EQUAL(v2[4], 4);\n}\n\n\nvoid TestUninitializedCopyDeviceSeq()\n{\n TestUninitializedCopyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUninitializedCopyDeviceSeq);\n\n\nvoid TestUninitializedCopyDeviceDevice()\n{\n TestUninitializedCopyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUninitializedCopyDeviceDevice);\n#endif\n\n\nvoid TestUninitializedCopyCudaStreams()\n{\n typedef thrust::device_vector Vector;\n \n Vector v1(5);\n v1[0] = 0; v1[1] = 1; v1[2] = 2; v1[3] = 3; v1[4] = 4;\n \n // copy to Vector\n Vector v2(5);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::uninitialized_copy(thrust::cuda::par.on(s), v1.begin(), v1.end(), v2.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(v2[0], 0);\n ASSERT_EQUAL(v2[1], 1);\n ASSERT_EQUAL(v2[2], 2);\n ASSERT_EQUAL(v2[3], 3);\n ASSERT_EQUAL(v2[4], 4);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestUninitializedCopyCudaStreams);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid uninitialized_copy_n_kernel(ExecutionPolicy exec, Iterator1 first, Size n, Iterator2 result)\n{\n thrust::uninitialized_copy_n(exec, first, n, result);\n}\n\n\ntemplate\nvoid TestUninitializedCopyNDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n \n Vector v1(5);\n v1[0] = 0; v1[1] = 1; v1[2] = 2; v1[3] = 3; v1[4] = 4;\n \n // copy to Vector\n Vector v2(5);\n uninitialized_copy_n_kernel<<<1,1>>>(exec, v1.begin(), v1.size(), v2.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n ASSERT_EQUAL(v2[0], 0);\n ASSERT_EQUAL(v2[1], 1);\n ASSERT_EQUAL(v2[2], 2);\n ASSERT_EQUAL(v2[3], 3);\n ASSERT_EQUAL(v2[4], 4);\n}\n\n\nvoid TestUninitializedCopyNDeviceSeq()\n{\n TestUninitializedCopyNDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestUninitializedCopyNDeviceSeq);\n\n\nvoid TestUninitializedCopyNDeviceDevice()\n{\n TestUninitializedCopyNDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestUninitializedCopyNDeviceDevice);\n#endif\n\n\nvoid TestUninitializedCopyNCudaStreams()\n{\n typedef thrust::device_vector Vector;\n \n Vector v1(5);\n v1[0] = 0; v1[1] = 1; v1[2] = 2; v1[3] = 3; v1[4] = 4;\n \n // copy to Vector\n Vector v2(5);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::uninitialized_copy_n(thrust::cuda::par.on(s), v1.begin(), v1.size(), v2.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(v2[0], 0);\n ASSERT_EQUAL(v2[1], 1);\n ASSERT_EQUAL(v2[2], 2);\n ASSERT_EQUAL(v2[3], 3);\n ASSERT_EQUAL(v2[4], 4);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestUninitializedCopyNCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/uninitialized_copy.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid stable_sort_by_key_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first)\n{\n thrust::stable_sort_by_key(exec, keys_first, keys_last, values_first);\n}\n\n\nstruct make_pair_functor\n{\n template\n __host__ __device__\n thrust::pair operator()(const T1 &x, const T2 &y)\n {\n return thrust::make_pair(x,y);\n } // end operator()()\n}; // end make_pair_functor\n\n\ntemplate\nvoid TestPairStableSortByKeyDevice(ExecutionPolicy exec)\n{\n size_t n = 10000;\n typedef thrust::pair P;\n\n // host arrays\n thrust::host_vector h_p1 = unittest::random_integers(n);\n thrust::host_vector h_p2 = unittest::random_integers(n);\n thrust::host_vector

h_pairs(n);\n\n thrust::host_vector h_values(n);\n thrust::sequence(h_values.begin(), h_values.end());\n\n // zip up pairs on the host\n thrust::transform(h_p1.begin(), h_p1.end(), h_p2.begin(), h_pairs.begin(), make_pair_functor());\n\n // device arrays\n thrust::device_vector

d_pairs = h_pairs;\n thrust::device_vector d_values = h_values;\n\n // sort on the device\n stable_sort_by_key_kernel<<<1,1>>>(exec, d_pairs.begin(), d_pairs.end(), d_values.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n // sort on the host\n thrust::stable_sort_by_key(h_pairs.begin(), h_pairs.end(), h_values.begin());\n\n ASSERT_EQUAL_QUIET(h_pairs, d_pairs);\n ASSERT_EQUAL(h_values, d_values);\n};\n\n\nvoid TestPairStableSortByKeyDeviceSeq()\n{\n TestPairStableSortByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestPairStableSortByKeyDeviceSeq);\n\n\nvoid TestPairStableSortByKeyDeviceDevice()\n{\n TestPairStableSortByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestPairStableSortByKeyDeviceDevice);\n#endif\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid stable_sort_by_key_kernel(ExecutionPolicy exec, Iterator1 keys_first, Iterator1 keys_last, Iterator2 values_first)\n{\n thrust::stable_sort_by_key(exec, keys_first, keys_last, values_first);\n}\n\n\nstruct make_pair_functor\n{\n template\n __host__ __device__\n thrust::pair operator()(const T1 &x, const T2 &y)\n {\n return thrust::make_pair(x,y);\n } // end operator()()\n}; // end make_pair_functor\n\n\ntemplate\nvoid TestPairStableSortByKeyDevice(ExecutionPolicy exec)\n{\n size_t n = 10000;\n typedef thrust::pair P;\n\n // host arrays\n thrust::host_vector h_p1 = unittest::random_integers(n);\n thrust::host_vector h_p2 = unittest::random_integers(n);\n thrust::host_vector

h_pairs(n);\n\n thrust::host_vector h_values(n);\n thrust::sequence(h_values.begin(), h_values.end());\n\n // zip up pairs on the host\n thrust::transform(h_p1.begin(), h_p1.end(), h_p2.begin(), h_pairs.begin(), make_pair_functor());\n\n // device arrays\n thrust::device_vector

d_pairs = h_pairs;\n thrust::device_vector d_values = h_values;\n\n // sort on the device\n stable_sort_by_key_kernel<<<1,1>>>(exec, d_pairs.begin(), d_pairs.end(), d_values.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n // sort on the host\n thrust::stable_sort_by_key(h_pairs.begin(), h_pairs.end(), h_values.begin());\n\n ASSERT_EQUAL_QUIET(h_pairs, d_pairs);\n ASSERT_EQUAL(h_values, d_values);\n};\n\n\nvoid TestPairStableSortByKeyDeviceSeq()\n{\n TestPairStableSortByKeyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestPairStableSortByKeyDeviceSeq);\n\n\nvoid TestPairStableSortByKeyDeviceDevice()\n{\n TestPairStableSortByKeyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestPairStableSortByKeyDeviceDevice);\n#endif\n\n", "source": "/kaggle/working/thrust/testing/cuda/pair_sort_by_key.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid transform_reduce_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Function1 f1, T init, Function2 f2, Iterator2 result)\n{\n *result = thrust::transform_reduce(exec, first, last, f1, init, f2);\n}\n\n\ntemplate\nvoid TestTransformReduceDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::value_type T;\n \n Vector data(3);\n data[0] = 1; data[1] = -2; data[2] = 3;\n \n T init = 10;\n\n thrust::device_vector result(1);\n\n transform_reduce_kernel<<<1,1>>>(exec, data.begin(), data.end(), thrust::negate(), init, thrust::plus(), result.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_EQUAL(8, (T)result[0]);\n}\n\n\nvoid TestTransformReduceDeviceSeq()\n{\n TestTransformReduceDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTransformReduceDeviceSeq);\n\n\nvoid TestTransformReduceDeviceDevice()\n{\n TestTransformReduceDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTransformReduceDeviceDevice);\n#endif\n\n\nvoid TestTransformReduceCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector data(3);\n data[0] = 1; data[1] = -2; data[2] = 3;\n \n T init = 10;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n T result = thrust::transform_reduce(thrust::cuda::par.on(s), data.begin(), data.end(), thrust::negate(), init, thrust::plus());\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(8, result);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestTransformReduceCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid transform_reduce_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Function1 f1, T init, Function2 f2, Iterator2 result)\n{\n *result = thrust::transform_reduce(exec, first, last, f1, init, f2);\n}\n\n\ntemplate\nvoid TestTransformReduceDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::value_type T;\n \n Vector data(3);\n data[0] = 1; data[1] = -2; data[2] = 3;\n \n T init = 10;\n\n thrust::device_vector result(1);\n\n transform_reduce_kernel<<<1,1>>>(exec, data.begin(), data.end(), thrust::negate(), init, thrust::plus(), result.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_EQUAL(8, (T)result[0]);\n}\n\n\nvoid TestTransformReduceDeviceSeq()\n{\n TestTransformReduceDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTransformReduceDeviceSeq);\n\n\nvoid TestTransformReduceDeviceDevice()\n{\n TestTransformReduceDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTransformReduceDeviceDevice);\n#endif\n\n\nvoid TestTransformReduceCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector data(3);\n data[0] = 1; data[1] = -2; data[2] = 3;\n \n T init = 10;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n T result = thrust::transform_reduce(thrust::cuda::par.on(s), data.begin(), data.end(), thrust::negate(), init, thrust::plus());\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(8, result);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestTransformReduceCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/transform_reduce.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_difference_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2, Iterator2 last2, Iterator3 result1, Iterator4 result2)\n{\n *result2 = thrust::set_difference(exec, first1, last1, first2, last2, result1);\n}\n\n\ntemplate\nvoid TestSetDifferenceDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::iterator Iterator;\n\n Vector a(4), b(5);\n\n a[0] = 0; a[1] = 2; a[2] = 4; a[3] = 5;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4; b[4] = 6;\n\n Vector ref(2);\n ref[0] = 2; ref[1] = 5;\n\n Vector result(2);\n\n thrust::device_vector end_vec(1);\n\n set_difference_kernel<<<1,1>>>(exec, a.begin(), a.end(), b.begin(), b.end(), result.begin(), end_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n Iterator end = end_vec.front();\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n}\n\n\nvoid TestSetDifferenceDeviceSeq()\n{\n TestSetDifferenceDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetDifferenceDeviceSeq);\n\n\nvoid TestSetDifferenceDeviceDevice()\n{\n TestSetDifferenceDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetDifferenceDeviceDevice);\n#endif\n\n\nvoid TestSetDifferenceCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a(4), b(5);\n\n a[0] = 0; a[1] = 2; a[2] = 4; a[3] = 5;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4; b[4] = 6;\n\n Vector ref(2);\n ref[0] = 2; ref[1] = 5;\n\n Vector result(2);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n Iterator end = thrust::set_difference(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin(), b.end(), result.begin());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSetDifferenceCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid set_difference_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2, Iterator2 last2, Iterator3 result1, Iterator4 result2)\n{\n *result2 = thrust::set_difference(exec, first1, last1, first2, last2, result1);\n}\n\n\ntemplate\nvoid TestSetDifferenceDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n typedef typename Vector::iterator Iterator;\n\n Vector a(4), b(5);\n\n a[0] = 0; a[1] = 2; a[2] = 4; a[3] = 5;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4; b[4] = 6;\n\n Vector ref(2);\n ref[0] = 2; ref[1] = 5;\n\n Vector result(2);\n\n thrust::device_vector end_vec(1);\n\n set_difference_kernel<<<1,1>>>(exec, a.begin(), a.end(), b.begin(), b.end(), result.begin(), end_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n Iterator end = end_vec.front();\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n}\n\n\nvoid TestSetDifferenceDeviceSeq()\n{\n TestSetDifferenceDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSetDifferenceDeviceSeq);\n\n\nvoid TestSetDifferenceDeviceDevice()\n{\n TestSetDifferenceDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSetDifferenceDeviceDevice);\n#endif\n\n\nvoid TestSetDifferenceCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::iterator Iterator;\n\n Vector a(4), b(5);\n\n a[0] = 0; a[1] = 2; a[2] = 4; a[3] = 5;\n b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4; b[4] = 6;\n\n Vector ref(2);\n ref[0] = 2; ref[1] = 5;\n\n Vector result(2);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n Iterator end = thrust::set_difference(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin(), b.end(), result.begin());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL_QUIET(result.end(), end);\n ASSERT_EQUAL(ref, result);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSetDifferenceCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/set_difference.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__ void adjacent_difference_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n thrust::adjacent_difference(exec, first, last, result);\n}\n\n\ntemplate\n__global__ void adjacent_difference_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result, BinaryFunction f)\n{\n thrust::adjacent_difference(exec, first, last, result, f);\n}\n\n\ntemplate\nvoid TestAdjacentDifferenceDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_input = unittest::random_samples(n);\n thrust::device_vector d_input = h_input;\n\n thrust::host_vector h_output(n);\n thrust::device_vector d_output(n);\n\n thrust::adjacent_difference(h_input.begin(), h_input.end(), h_output.begin());\n adjacent_difference_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_output.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(h_output, d_output);\n\n thrust::adjacent_difference(h_input.begin(), h_input.end(), h_output.begin(), thrust::plus());\n adjacent_difference_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_output.begin(), thrust::plus());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(h_output, d_output);\n\n // in-place operation\n thrust::adjacent_difference(h_input.begin(), h_input.end(), h_input.begin(), thrust::plus());\n adjacent_difference_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_input.begin(), thrust::plus());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(h_input, h_output); //computed previously\n ASSERT_EQUAL(d_input, d_output); //computed previously\n}\n\n\ntemplate\nvoid TestAdjacentDifferenceDeviceSeq(const size_t n)\n{\n TestAdjacentDifferenceDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestAdjacentDifferenceDeviceSeq);\n\n\ntemplate\nvoid TestAdjacentDifferenceDeviceDevice(const size_t n)\n{\n TestAdjacentDifferenceDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestAdjacentDifferenceDeviceDevice);\n#endif\n\n\nvoid TestAdjacentDifferenceCudaStreams()\n{\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::device_vector input(3);\n thrust::device_vector output(3);\n input[0] = 1; input[1] = 4; input[2] = 6;\n\n thrust::adjacent_difference(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin());\n\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(output[0], 1);\n ASSERT_EQUAL(output[1], 3);\n ASSERT_EQUAL(output[2], 2);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestAdjacentDifferenceCudaStreams);\n\nstruct detect_wrong_difference\n{\n using difference_type = void;\n using value_type = void;\n using pointer = void;\n using reference = void;\n using iterator_category = std::output_iterator_tag;\n\n bool * flag;\n\n __host__ __device__ detect_wrong_difference operator++() const { return *this; }\n __host__ __device__ detect_wrong_difference operator*() const { return *this; }\n template\n __host__ __device__ detect_wrong_difference operator+(Difference) const { return *this; }\n template\n __host__ __device__ detect_wrong_difference operator[](Index) const { return *this; }\n\n __device__\n void operator=(long long difference) const\n {\n if (difference != 1)\n {\n *flag = false;\n }\n }\n};\n\nvoid TestAdjacentDifferenceWithBigIndexesHelper(int magnitude)\n{\n thrust::counting_iterator begin(1);\n thrust::counting_iterator end = begin + (1ll << magnitude);\n ASSERT_EQUAL(thrust::distance(begin, end), 1ll << magnitude);\n\n thrust::device_ptr all_differences_correct = thrust::device_malloc(1);\n *all_differences_correct = true;\n\n detect_wrong_difference out = { thrust::raw_pointer_cast(all_differences_correct) };\n\n thrust::adjacent_difference(thrust::device, begin, end, out);\n\n bool all_differences_correct_h = *all_differences_correct;\n thrust::device_free(all_differences_correct);\n\n ASSERT_EQUAL(all_differences_correct_h, true);\n}\n\nvoid TestAdjacentDifferenceWithBigIndexes()\n{\n TestAdjacentDifferenceWithBigIndexesHelper(30);\n TestAdjacentDifferenceWithBigIndexesHelper(31);\n TestAdjacentDifferenceWithBigIndexesHelper(32);\n TestAdjacentDifferenceWithBigIndexesHelper(33);\n}\nDECLARE_UNITTEST(TestAdjacentDifferenceWithBigIndexes);\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__ void adjacent_difference_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result)\n{\n thrust::adjacent_difference(exec, first, last, result);\n}\n\n\ntemplate\n__global__ void adjacent_difference_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result, BinaryFunction f)\n{\n thrust::adjacent_difference(exec, first, last, result, f);\n}\n\n\ntemplate\nvoid TestAdjacentDifferenceDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_input = unittest::random_samples(n);\n thrust::device_vector d_input = h_input;\n\n thrust::host_vector h_output(n);\n thrust::device_vector d_output(n);\n\n thrust::adjacent_difference(h_input.begin(), h_input.end(), h_output.begin());\n adjacent_difference_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_output.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(h_output, d_output);\n\n thrust::adjacent_difference(h_input.begin(), h_input.end(), h_output.begin(), thrust::plus());\n adjacent_difference_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_output.begin(), thrust::plus());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(h_output, d_output);\n\n // in-place operation\n thrust::adjacent_difference(h_input.begin(), h_input.end(), h_input.begin(), thrust::plus());\n adjacent_difference_kernel<<<1,1>>>(exec, d_input.begin(), d_input.end(), d_input.begin(), thrust::plus());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(h_input, h_output); //computed previously\n ASSERT_EQUAL(d_input, d_output); //computed previously\n}\n\n\ntemplate\nvoid TestAdjacentDifferenceDeviceSeq(const size_t n)\n{\n TestAdjacentDifferenceDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestAdjacentDifferenceDeviceSeq);\n\n\ntemplate\nvoid TestAdjacentDifferenceDeviceDevice(const size_t n)\n{\n TestAdjacentDifferenceDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestAdjacentDifferenceDeviceDevice);\n#endif\n\n\nvoid TestAdjacentDifferenceCudaStreams()\n{\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::device_vector input(3);\n thrust::device_vector output(3);\n input[0] = 1; input[1] = 4; input[2] = 6;\n\n thrust::adjacent_difference(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin());\n\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(output[0], 1);\n ASSERT_EQUAL(output[1], 3);\n ASSERT_EQUAL(output[2], 2);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestAdjacentDifferenceCudaStreams);\n\nstruct detect_wrong_difference\n{\n using difference_type = void;\n using value_type = void;\n using pointer = void;\n using reference = void;\n using iterator_category = std::output_iterator_tag;\n\n bool * flag;\n\n __host__ __device__ detect_wrong_difference operator++() const { return *this; }\n __host__ __device__ detect_wrong_difference operator*() const { return *this; }\n template\n __host__ __device__ detect_wrong_difference operator+(Difference) const { return *this; }\n template\n __host__ __device__ detect_wrong_difference operator[](Index) const { return *this; }\n\n __device__\n void operator=(long long difference) const\n {\n if (difference != 1)\n {\n *flag = false;\n }\n }\n};\n\nvoid TestAdjacentDifferenceWithBigIndexesHelper(int magnitude)\n{\n thrust::counting_iterator begin(1);\n thrust::counting_iterator end = begin + (1ll << magnitude);\n ASSERT_EQUAL(thrust::distance(begin, end), 1ll << magnitude);\n\n thrust::device_ptr all_differences_correct = thrust::device_malloc(1);\n *all_differences_correct = true;\n\n detect_wrong_difference out = { thrust::raw_pointer_cast(all_differences_correct) };\n\n thrust::adjacent_difference(thrust::device, begin, end, out);\n\n bool all_differences_correct_h = *all_differences_correct;\n thrust::device_free(all_differences_correct);\n\n ASSERT_EQUAL(all_differences_correct_h, true);\n}\n\nvoid TestAdjacentDifferenceWithBigIndexes()\n{\n TestAdjacentDifferenceWithBigIndexesHelper(30);\n TestAdjacentDifferenceWithBigIndexesHelper(31);\n TestAdjacentDifferenceWithBigIndexesHelper(32);\n TestAdjacentDifferenceWithBigIndexesHelper(33);\n}\nDECLARE_UNITTEST(TestAdjacentDifferenceWithBigIndexes);\n", "source": "/kaggle/working/thrust/testing/cuda/adjacent_difference.cu"} {"input": "#include \n#include \n#include \n\n\ntemplate \nstruct less_than_five\n{\n __host__ __device__ bool operator()(const T &val) const {return val < 5;}\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid replace_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T1 old_value, T2 new_value)\n{\n thrust::replace(exec, first, last, old_value, new_value);\n}\n\n\ntemplate\nvoid TestReplaceDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n T old_value = 0;\n T new_value = 1;\n \n thrust::replace(h_data.begin(), h_data.end(), old_value, new_value);\n\n replace_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), old_value, new_value);\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_ALMOST_EQUAL(h_data, d_data);\n}\n\n\ntemplate\nvoid TestReplaceDeviceSeq(const size_t n)\n{\n TestReplaceDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestReplaceDeviceSeq);\n\ntemplate\nvoid TestReplaceDeviceDevice(const size_t n)\n{\n TestReplaceDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestReplaceDeviceDevice);\n\n\ntemplate\n__global__\nvoid replace_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result, T1 old_value, T2 new_value)\n{\n thrust::replace_copy(exec, first, last, result, old_value, new_value);\n}\n\n\ntemplate\nvoid TestReplaceCopyDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n int old_value = 0;\n int new_value = 1;\n \n thrust::host_vector h_dest(n);\n thrust::device_vector d_dest(n);\n \n thrust::replace_copy(h_data.begin(), h_data.end(), h_dest.begin(), old_value, new_value);\n\n replace_copy_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_dest.begin(), old_value, new_value);\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_ALMOST_EQUAL(h_data, d_data);\n ASSERT_ALMOST_EQUAL(h_dest, d_dest);\n}\n\nvoid TestReplaceCopyDeviceSeq()\n{\n TestReplaceCopyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReplaceCopyDeviceSeq);\n\nvoid TestReplaceCopyDeviceDevice()\n{\n TestReplaceCopyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReplaceCopyDeviceDevice);\n\n\ntemplate\n__global__\nvoid replace_if_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Predicate pred, T new_value)\n{\n thrust::replace_if(exec, first, last, pred, new_value);\n}\n\n\ntemplate\nvoid TestReplaceIfDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n thrust::replace_if(h_data.begin(), h_data.end(), less_than_five(), 0);\n\n replace_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), less_than_five(), 0);\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_ALMOST_EQUAL(h_data, d_data);\n}\n\nvoid TestReplaceIfDeviceSeq()\n{\n TestReplaceIfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReplaceIfDeviceSeq);\n\nvoid TestReplaceIfDeviceDevice()\n{\n TestReplaceIfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReplaceIfDeviceDevice);\n\n\ntemplate\n__global__\nvoid replace_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Predicate pred, T new_value)\n{\n thrust::replace_if(exec, first, last, stencil_first, pred, new_value);\n}\n\n\ntemplate\nvoid TestReplaceIfStencilDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n thrust::host_vector h_stencil = unittest::random_samples(n);\n thrust::device_vector d_stencil = h_stencil;\n \n thrust::replace_if(h_data.begin(), h_data.end(), h_stencil.begin(), less_than_five(), 0);\n\n replace_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_stencil.begin(), less_than_five(), 0);\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_ALMOST_EQUAL(h_data, d_data);\n}\n\nvoid TestReplaceIfStencilDeviceSeq()\n{\n TestReplaceIfStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReplaceIfStencilDeviceSeq);\n\nvoid TestReplaceIfStencilDeviceDevice()\n{\n TestReplaceIfStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReplaceIfStencilDeviceDevice);\n\n\ntemplate\n__global__\nvoid replace_copy_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result, Predicate pred, T new_value)\n{\n thrust::replace_copy_if(exec, first, last, result, pred, new_value);\n}\n\n\ntemplate\nvoid TestReplaceCopyIfDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n thrust::host_vector h_dest(n);\n thrust::device_vector d_dest(n);\n \n thrust::replace_copy_if(h_data.begin(), h_data.end(), h_dest.begin(), less_than_five(), 0);\n\n replace_copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_dest.begin(), less_than_five(), 0);\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_ALMOST_EQUAL(h_data, d_data);\n ASSERT_ALMOST_EQUAL(h_dest, d_dest);\n}\n\nvoid TestReplaceCopyIfDeviceSeq()\n{\n TestReplaceCopyIfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReplaceCopyIfDeviceSeq);\n\nvoid TestReplaceCopyIfDeviceDevice()\n{\n TestReplaceCopyIfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReplaceCopyIfDeviceDevice);\n\n\ntemplate\n__global__\nvoid replace_copy_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Iterator3 result, Predicate pred, T new_value)\n{\n thrust::replace_copy_if(exec, first, last, stencil_first, result, pred, new_value);\n}\n\n\ntemplate\nvoid TestReplaceCopyIfStencilDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n thrust::host_vector h_stencil = unittest::random_samples(n);\n thrust::device_vector d_stencil = h_stencil;\n \n thrust::host_vector h_dest(n);\n thrust::device_vector d_dest(n);\n \n thrust::replace_copy_if(h_data.begin(), h_data.end(), h_stencil.begin(), h_dest.begin(), less_than_five(), 0);\n\n replace_copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_stencil.begin(), d_dest.begin(), less_than_five(), 0);\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_ALMOST_EQUAL(h_data, d_data);\n ASSERT_ALMOST_EQUAL(h_dest, d_dest);\n}\n\n\nvoid TestReplaceCopyIfStencilDeviceSeq()\n{\n TestReplaceCopyIfStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReplaceCopyIfStencilDeviceSeq);\n\n\nvoid TestReplaceCopyIfStencilDeviceDevice()\n{\n TestReplaceCopyIfStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReplaceCopyIfStencilDeviceDevice);\n#endif\n\n\nvoid TestReplaceCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::replace(thrust::cuda::par.on(s), data.begin(), data.end(), (T) 1, (T) 4);\n thrust::replace(thrust::cuda::par.on(s), data.begin(), data.end(), (T) 2, (T) 5);\n\n cudaStreamSynchronize(s);\n\n Vector result(5);\n result[0] = 4; \n result[1] = 5; \n result[2] = 4;\n result[3] = 3; \n result[4] = 5; \n\n ASSERT_EQUAL(data, result);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestReplaceCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\ntemplate \nstruct less_than_five\n{\n __host__ __device__ bool operator()(const T &val) const {return val < 5;}\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid replace_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T1 old_value, T2 new_value)\n{\n thrust::replace(exec, first, last, old_value, new_value);\n}\n\n\ntemplate\nvoid TestReplaceDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n T old_value = 0;\n T new_value = 1;\n \n thrust::replace(h_data.begin(), h_data.end(), old_value, new_value);\n\n replace_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), old_value, new_value);\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_ALMOST_EQUAL(h_data, d_data);\n}\n\n\ntemplate\nvoid TestReplaceDeviceSeq(const size_t n)\n{\n TestReplaceDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestReplaceDeviceSeq);\n\ntemplate\nvoid TestReplaceDeviceDevice(const size_t n)\n{\n TestReplaceDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestReplaceDeviceDevice);\n\n\ntemplate\n__global__\nvoid replace_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result, T1 old_value, T2 new_value)\n{\n thrust::replace_copy(exec, first, last, result, old_value, new_value);\n}\n\n\ntemplate\nvoid TestReplaceCopyDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n int old_value = 0;\n int new_value = 1;\n \n thrust::host_vector h_dest(n);\n thrust::device_vector d_dest(n);\n \n thrust::replace_copy(h_data.begin(), h_data.end(), h_dest.begin(), old_value, new_value);\n\n replace_copy_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_dest.begin(), old_value, new_value);\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_ALMOST_EQUAL(h_data, d_data);\n ASSERT_ALMOST_EQUAL(h_dest, d_dest);\n}\n\nvoid TestReplaceCopyDeviceSeq()\n{\n TestReplaceCopyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReplaceCopyDeviceSeq);\n\nvoid TestReplaceCopyDeviceDevice()\n{\n TestReplaceCopyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReplaceCopyDeviceDevice);\n\n\ntemplate\n__global__\nvoid replace_if_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Predicate pred, T new_value)\n{\n thrust::replace_if(exec, first, last, pred, new_value);\n}\n\n\ntemplate\nvoid TestReplaceIfDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n thrust::replace_if(h_data.begin(), h_data.end(), less_than_five(), 0);\n\n replace_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), less_than_five(), 0);\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_ALMOST_EQUAL(h_data, d_data);\n}\n\nvoid TestReplaceIfDeviceSeq()\n{\n TestReplaceIfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReplaceIfDeviceSeq);\n\nvoid TestReplaceIfDeviceDevice()\n{\n TestReplaceIfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReplaceIfDeviceDevice);\n\n\ntemplate\n__global__\nvoid replace_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Predicate pred, T new_value)\n{\n thrust::replace_if(exec, first, last, stencil_first, pred, new_value);\n}\n\n\ntemplate\nvoid TestReplaceIfStencilDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n thrust::host_vector h_stencil = unittest::random_samples(n);\n thrust::device_vector d_stencil = h_stencil;\n \n thrust::replace_if(h_data.begin(), h_data.end(), h_stencil.begin(), less_than_five(), 0);\n\n replace_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_stencil.begin(), less_than_five(), 0);\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_ALMOST_EQUAL(h_data, d_data);\n}\n\nvoid TestReplaceIfStencilDeviceSeq()\n{\n TestReplaceIfStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReplaceIfStencilDeviceSeq);\n\nvoid TestReplaceIfStencilDeviceDevice()\n{\n TestReplaceIfStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReplaceIfStencilDeviceDevice);\n\n\ntemplate\n__global__\nvoid replace_copy_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result, Predicate pred, T new_value)\n{\n thrust::replace_copy_if(exec, first, last, result, pred, new_value);\n}\n\n\ntemplate\nvoid TestReplaceCopyIfDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n thrust::host_vector h_dest(n);\n thrust::device_vector d_dest(n);\n \n thrust::replace_copy_if(h_data.begin(), h_data.end(), h_dest.begin(), less_than_five(), 0);\n\n replace_copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_dest.begin(), less_than_five(), 0);\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_ALMOST_EQUAL(h_data, d_data);\n ASSERT_ALMOST_EQUAL(h_dest, d_dest);\n}\n\nvoid TestReplaceCopyIfDeviceSeq()\n{\n TestReplaceCopyIfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReplaceCopyIfDeviceSeq);\n\nvoid TestReplaceCopyIfDeviceDevice()\n{\n TestReplaceCopyIfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReplaceCopyIfDeviceDevice);\n\n\ntemplate\n__global__\nvoid replace_copy_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Iterator3 result, Predicate pred, T new_value)\n{\n thrust::replace_copy_if(exec, first, last, stencil_first, result, pred, new_value);\n}\n\n\ntemplate\nvoid TestReplaceCopyIfStencilDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_samples(n);\n thrust::device_vector d_data = h_data;\n \n thrust::host_vector h_stencil = unittest::random_samples(n);\n thrust::device_vector d_stencil = h_stencil;\n \n thrust::host_vector h_dest(n);\n thrust::device_vector d_dest(n);\n \n thrust::replace_copy_if(h_data.begin(), h_data.end(), h_stencil.begin(), h_dest.begin(), less_than_five(), 0);\n\n replace_copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_stencil.begin(), d_dest.begin(), less_than_five(), 0);\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_ALMOST_EQUAL(h_data, d_data);\n ASSERT_ALMOST_EQUAL(h_dest, d_dest);\n}\n\n\nvoid TestReplaceCopyIfStencilDeviceSeq()\n{\n TestReplaceCopyIfStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestReplaceCopyIfStencilDeviceSeq);\n\n\nvoid TestReplaceCopyIfStencilDeviceDevice()\n{\n TestReplaceCopyIfStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestReplaceCopyIfStencilDeviceDevice);\n#endif\n\n\nvoid TestReplaceCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::replace(thrust::cuda::par.on(s), data.begin(), data.end(), (T) 1, (T) 4);\n thrust::replace(thrust::cuda::par.on(s), data.begin(), data.end(), (T) 2, (T) 5);\n\n hipStreamSynchronize(s);\n\n Vector result(5);\n result[0] = 4; \n result[1] = 5; \n result[2] = 4;\n result[3] = 3; \n result[4] = 5; \n\n ASSERT_EQUAL(data, result);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestReplaceCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/replace.cu"} {"input": "#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid tabulate_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Function f)\n{\n thrust::tabulate(exec, first, last, f);\n}\n\n\ntemplate\nvoid TestTabulateDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n using namespace thrust::placeholders;\n typedef typename Vector::value_type T;\n \n Vector v(5);\n\n tabulate_kernel<<<1,1>>>(exec, v.begin(), v.end(), thrust::identity());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], 1);\n ASSERT_EQUAL(v[2], 2);\n ASSERT_EQUAL(v[3], 3);\n ASSERT_EQUAL(v[4], 4);\n\n tabulate_kernel<<<1,1>>>(exec, v.begin(), v.end(), -_1);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], -1);\n ASSERT_EQUAL(v[2], -2);\n ASSERT_EQUAL(v[3], -3);\n ASSERT_EQUAL(v[4], -4);\n \n tabulate_kernel<<<1,1>>>(exec, v.begin(), v.end(), _1 * _1 * _1);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], 1);\n ASSERT_EQUAL(v[2], 8);\n ASSERT_EQUAL(v[3], 27);\n ASSERT_EQUAL(v[4], 64);\n}\n\nvoid TestTabulateDeviceSeq()\n{\n TestTabulateDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTabulateDeviceSeq);\n\nvoid TestTabulateDeviceDevice()\n{\n TestTabulateDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTabulateDeviceDevice);\n#endif\n\nvoid TestTabulateCudaStreams()\n{\n using namespace thrust::placeholders;\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector v(5);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::tabulate(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], 1);\n ASSERT_EQUAL(v[2], 2);\n ASSERT_EQUAL(v[3], 3);\n ASSERT_EQUAL(v[4], 4);\n\n thrust::tabulate(thrust::cuda::par.on(s), v.begin(), v.end(), -_1);\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], -1);\n ASSERT_EQUAL(v[2], -2);\n ASSERT_EQUAL(v[3], -3);\n ASSERT_EQUAL(v[4], -4);\n \n thrust::tabulate(thrust::cuda::par.on(s), v.begin(), v.end(), _1 * _1 * _1);\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], 1);\n ASSERT_EQUAL(v[2], 8);\n ASSERT_EQUAL(v[3], 27);\n ASSERT_EQUAL(v[4], 64);\n\n cudaStreamSynchronize(s);\n}\nDECLARE_UNITTEST(TestTabulateCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid tabulate_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Function f)\n{\n thrust::tabulate(exec, first, last, f);\n}\n\n\ntemplate\nvoid TestTabulateDevice(ExecutionPolicy exec)\n{\n typedef thrust::device_vector Vector;\n using namespace thrust::placeholders;\n typedef typename Vector::value_type T;\n \n Vector v(5);\n\n tabulate_kernel<<<1,1>>>(exec, v.begin(), v.end(), thrust::identity());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], 1);\n ASSERT_EQUAL(v[2], 2);\n ASSERT_EQUAL(v[3], 3);\n ASSERT_EQUAL(v[4], 4);\n\n tabulate_kernel<<<1,1>>>(exec, v.begin(), v.end(), -_1);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], -1);\n ASSERT_EQUAL(v[2], -2);\n ASSERT_EQUAL(v[3], -3);\n ASSERT_EQUAL(v[4], -4);\n \n tabulate_kernel<<<1,1>>>(exec, v.begin(), v.end(), _1 * _1 * _1);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], 1);\n ASSERT_EQUAL(v[2], 8);\n ASSERT_EQUAL(v[3], 27);\n ASSERT_EQUAL(v[4], 64);\n}\n\nvoid TestTabulateDeviceSeq()\n{\n TestTabulateDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestTabulateDeviceSeq);\n\nvoid TestTabulateDeviceDevice()\n{\n TestTabulateDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestTabulateDeviceDevice);\n#endif\n\nvoid TestTabulateCudaStreams()\n{\n using namespace thrust::placeholders;\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector v(5);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::tabulate(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], 1);\n ASSERT_EQUAL(v[2], 2);\n ASSERT_EQUAL(v[3], 3);\n ASSERT_EQUAL(v[4], 4);\n\n thrust::tabulate(thrust::cuda::par.on(s), v.begin(), v.end(), -_1);\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], -1);\n ASSERT_EQUAL(v[2], -2);\n ASSERT_EQUAL(v[3], -3);\n ASSERT_EQUAL(v[4], -4);\n \n thrust::tabulate(thrust::cuda::par.on(s), v.begin(), v.end(), _1 * _1 * _1);\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], 1);\n ASSERT_EQUAL(v[2], 8);\n ASSERT_EQUAL(v[3], 27);\n ASSERT_EQUAL(v[4], 64);\n\n hipStreamSynchronize(s);\n}\nDECLARE_UNITTEST(TestTabulateCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/tabulate.cu"} {"input": "#include \n#include \n#include \n\n\ntemplate\nstruct equal_to_value_pred\n{\n T value;\n\n equal_to_value_pred(T value) : value(value) {}\n\n __host__ __device__\n bool operator()(T v) const { return v == value; }\n};\n\n\ntemplate\nstruct not_equal_to_value_pred\n{\n T value;\n\n not_equal_to_value_pred(T value) : value(value) {}\n\n __host__ __device__\n bool operator()(T v) const { return v != value; }\n};\n\n\ntemplate\nstruct less_than_value_pred\n{\n T value;\n\n less_than_value_pred(T value) : value(value) {}\n\n __host__ __device__\n bool operator()(T v) const { return v < value; }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__ void find_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T value, Iterator2 result)\n{\n *result = thrust::find(exec, first, last, value);\n}\n\n\ntemplate\nvoid TestFindDevice(ExecutionPolicy exec)\n{\n size_t n = 100;\n\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n typename thrust::host_vector::iterator h_iter;\n \n typedef typename thrust::device_vector::iterator iter_type;\n thrust::device_vector d_result(1);\n \n h_iter = thrust::find(h_data.begin(), h_data.end(), int(0));\n\n find_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), int(0), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_iter - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n \n for(size_t i = 1; i < n; i *= 2)\n {\n int sample = h_data[i];\n\n h_iter = thrust::find(h_data.begin(), h_data.end(), sample);\n\n find_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), sample, d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(h_iter - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n }\n}\n\n\nvoid TestFindDeviceSeq()\n{\n TestFindDevice(thrust::seq);\n};\nDECLARE_UNITTEST(TestFindDeviceSeq);\n\n\nvoid TestFindDeviceDevice()\n{\n TestFindDevice(thrust::device);\n};\nDECLARE_UNITTEST(TestFindDeviceDevice);\n\n\ntemplate\n__global__ void find_if_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Predicate pred, Iterator2 result)\n{\n *result = thrust::find_if(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestFindIfDevice(ExecutionPolicy exec)\n{\n size_t n = 100;\n\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n typename thrust::host_vector::iterator h_iter;\n \n typedef typename thrust::device_vector::iterator iter_type;\n thrust::device_vector d_result(1);\n \n h_iter = thrust::find_if(h_data.begin(), h_data.end(), equal_to_value_pred(0));\n\n find_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), equal_to_value_pred(0), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(h_iter - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n \n for (size_t i = 1; i < n; i *= 2)\n {\n int sample = h_data[i];\n\n h_iter = thrust::find_if(h_data.begin(), h_data.end(), equal_to_value_pred(sample));\n\n find_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), equal_to_value_pred(sample), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(h_iter - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n }\n}\n\n\nvoid TestFindIfDeviceSeq()\n{\n TestFindIfDevice(thrust::seq);\n};\nDECLARE_UNITTEST(TestFindIfDeviceSeq);\n\n\nvoid TestFindIfDeviceDevice()\n{\n TestFindIfDevice(thrust::device);\n};\nDECLARE_UNITTEST(TestFindIfDeviceDevice);\n\n\ntemplate\n__global__ void find_if_not_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Predicate pred, Iterator2 result)\n{\n *result = thrust::find_if_not(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestFindIfNotDevice(ExecutionPolicy exec)\n{\n size_t n = 100;\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n typename thrust::host_vector::iterator h_iter;\n \n typedef typename thrust::device_vector::iterator iter_type;\n thrust::device_vector d_result(1);\n \n h_iter = thrust::find_if_not(h_data.begin(), h_data.end(), not_equal_to_value_pred(0));\n\n find_if_not_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), not_equal_to_value_pred(0), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(h_iter - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n \n for(size_t i = 1; i < n; i *= 2)\n {\n int sample = h_data[i];\n\n h_iter = thrust::find_if_not(h_data.begin(), h_data.end(), not_equal_to_value_pred(sample));\n\n find_if_not_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), not_equal_to_value_pred(sample), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(h_iter - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n }\n}\n\n\nvoid TestFindIfNotDeviceSeq()\n{\n TestFindIfNotDevice(thrust::seq);\n};\nDECLARE_UNITTEST(TestFindIfNotDeviceSeq);\n\n\nvoid TestFindIfNotDeviceDevice()\n{\n TestFindIfNotDevice(thrust::device);\n};\nDECLARE_UNITTEST(TestFindIfNotDeviceDevice);\n#endif\n\n\nvoid TestFindCudaStreams()\n{\n thrust::device_vector vec(5);\n vec[0] = 1;\n vec[1] = 2;\n vec[2] = 3;\n vec[3] = 3;\n vec[4] = 5;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n ASSERT_EQUAL(thrust::find(thrust::cuda::par.on(s), vec.begin(), vec.end(), 0) - vec.begin(), 5);\n ASSERT_EQUAL(thrust::find(thrust::cuda::par.on(s), vec.begin(), vec.end(), 1) - vec.begin(), 0);\n ASSERT_EQUAL(thrust::find(thrust::cuda::par.on(s), vec.begin(), vec.end(), 2) - vec.begin(), 1);\n ASSERT_EQUAL(thrust::find(thrust::cuda::par.on(s), vec.begin(), vec.end(), 3) - vec.begin(), 2);\n ASSERT_EQUAL(thrust::find(thrust::cuda::par.on(s), vec.begin(), vec.end(), 4) - vec.begin(), 5);\n ASSERT_EQUAL(thrust::find(thrust::cuda::par.on(s), vec.begin(), vec.end(), 5) - vec.begin(), 4);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestFindCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\ntemplate\nstruct equal_to_value_pred\n{\n T value;\n\n equal_to_value_pred(T value) : value(value) {}\n\n __host__ __device__\n bool operator()(T v) const { return v == value; }\n};\n\n\ntemplate\nstruct not_equal_to_value_pred\n{\n T value;\n\n not_equal_to_value_pred(T value) : value(value) {}\n\n __host__ __device__\n bool operator()(T v) const { return v != value; }\n};\n\n\ntemplate\nstruct less_than_value_pred\n{\n T value;\n\n less_than_value_pred(T value) : value(value) {}\n\n __host__ __device__\n bool operator()(T v) const { return v < value; }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__ void find_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T value, Iterator2 result)\n{\n *result = thrust::find(exec, first, last, value);\n}\n\n\ntemplate\nvoid TestFindDevice(ExecutionPolicy exec)\n{\n size_t n = 100;\n\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n typename thrust::host_vector::iterator h_iter;\n \n typedef typename thrust::device_vector::iterator iter_type;\n thrust::device_vector d_result(1);\n \n h_iter = thrust::find(h_data.begin(), h_data.end(), int(0));\n\n find_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), int(0), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_iter - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n \n for(size_t i = 1; i < n; i *= 2)\n {\n int sample = h_data[i];\n\n h_iter = thrust::find(h_data.begin(), h_data.end(), sample);\n\n find_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), sample, d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(h_iter - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n }\n}\n\n\nvoid TestFindDeviceSeq()\n{\n TestFindDevice(thrust::seq);\n};\nDECLARE_UNITTEST(TestFindDeviceSeq);\n\n\nvoid TestFindDeviceDevice()\n{\n TestFindDevice(thrust::device);\n};\nDECLARE_UNITTEST(TestFindDeviceDevice);\n\n\ntemplate\n__global__ void find_if_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Predicate pred, Iterator2 result)\n{\n *result = thrust::find_if(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestFindIfDevice(ExecutionPolicy exec)\n{\n size_t n = 100;\n\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n typename thrust::host_vector::iterator h_iter;\n \n typedef typename thrust::device_vector::iterator iter_type;\n thrust::device_vector d_result(1);\n \n h_iter = thrust::find_if(h_data.begin(), h_data.end(), equal_to_value_pred(0));\n\n find_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), equal_to_value_pred(0), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(h_iter - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n \n for (size_t i = 1; i < n; i *= 2)\n {\n int sample = h_data[i];\n\n h_iter = thrust::find_if(h_data.begin(), h_data.end(), equal_to_value_pred(sample));\n\n find_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), equal_to_value_pred(sample), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(h_iter - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n }\n}\n\n\nvoid TestFindIfDeviceSeq()\n{\n TestFindIfDevice(thrust::seq);\n};\nDECLARE_UNITTEST(TestFindIfDeviceSeq);\n\n\nvoid TestFindIfDeviceDevice()\n{\n TestFindIfDevice(thrust::device);\n};\nDECLARE_UNITTEST(TestFindIfDeviceDevice);\n\n\ntemplate\n__global__ void find_if_not_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Predicate pred, Iterator2 result)\n{\n *result = thrust::find_if_not(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestFindIfNotDevice(ExecutionPolicy exec)\n{\n size_t n = 100;\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n typename thrust::host_vector::iterator h_iter;\n \n typedef typename thrust::device_vector::iterator iter_type;\n thrust::device_vector d_result(1);\n \n h_iter = thrust::find_if_not(h_data.begin(), h_data.end(), not_equal_to_value_pred(0));\n\n find_if_not_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), not_equal_to_value_pred(0), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(h_iter - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n \n for(size_t i = 1; i < n; i *= 2)\n {\n int sample = h_data[i];\n\n h_iter = thrust::find_if_not(h_data.begin(), h_data.end(), not_equal_to_value_pred(sample));\n\n find_if_not_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), not_equal_to_value_pred(sample), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(h_iter - h_data.begin(), (iter_type)d_result[0] - d_data.begin());\n }\n}\n\n\nvoid TestFindIfNotDeviceSeq()\n{\n TestFindIfNotDevice(thrust::seq);\n};\nDECLARE_UNITTEST(TestFindIfNotDeviceSeq);\n\n\nvoid TestFindIfNotDeviceDevice()\n{\n TestFindIfNotDevice(thrust::device);\n};\nDECLARE_UNITTEST(TestFindIfNotDeviceDevice);\n#endif\n\n\nvoid TestFindCudaStreams()\n{\n thrust::device_vector vec(5);\n vec[0] = 1;\n vec[1] = 2;\n vec[2] = 3;\n vec[3] = 3;\n vec[4] = 5;\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n ASSERT_EQUAL(thrust::find(thrust::cuda::par.on(s), vec.begin(), vec.end(), 0) - vec.begin(), 5);\n ASSERT_EQUAL(thrust::find(thrust::cuda::par.on(s), vec.begin(), vec.end(), 1) - vec.begin(), 0);\n ASSERT_EQUAL(thrust::find(thrust::cuda::par.on(s), vec.begin(), vec.end(), 2) - vec.begin(), 1);\n ASSERT_EQUAL(thrust::find(thrust::cuda::par.on(s), vec.begin(), vec.end(), 3) - vec.begin(), 2);\n ASSERT_EQUAL(thrust::find(thrust::cuda::par.on(s), vec.begin(), vec.end(), 4) - vec.begin(), 5);\n ASSERT_EQUAL(thrust::find(thrust::cuda::par.on(s), vec.begin(), vec.end(), 5) - vec.begin(), 4);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestFindCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/find.cu"} {"input": "#include \n#include \n#include \n#include \n\nstatic const size_t NUM_REGISTERS = 64;\n\ntemplate __host__ __device__ void f (int * x) { int temp = *x; f(x + 1); *x = temp;};\ntemplate <> __host__ __device__ void f<0>(int * /*x*/) { }\ntemplate \nstruct CopyFunctorWithManyRegisters\n{\n __host__ __device__\n void operator()(int * ptr)\n {\n f(ptr);\n }\n};\n\n\nvoid TestForEachLargeRegisterFootprint()\n{\n int current_device = -1;\n cudaGetDevice(¤t_device);\n cudaDeviceProp prop;\n cudaGetDeviceProperties(&prop, current_device);\n\n thrust::device_vector data(NUM_REGISTERS, 12345);\n\n thrust::device_vector input(1, thrust::raw_pointer_cast(&data[0])); // length is irrelevant\n \n thrust::for_each(input.begin(), input.end(), CopyFunctorWithManyRegisters());\n}\nDECLARE_UNITTEST(TestForEachLargeRegisterFootprint);\n\n\nvoid TestForEachNLargeRegisterFootprint()\n{\n int current_device = -1;\n cudaGetDevice(¤t_device);\n cudaDeviceProp prop;\n cudaGetDeviceProperties(&prop, current_device);\n\n thrust::device_vector data(NUM_REGISTERS, 12345);\n\n thrust::device_vector input(1, thrust::raw_pointer_cast(&data[0])); // length is irrelevant\n \n thrust::for_each_n(input.begin(), input.size(), CopyFunctorWithManyRegisters());\n}\nDECLARE_UNITTEST(TestForEachNLargeRegisterFootprint);\n\n\ntemplate \nstruct mark_present_for_each\n{\n T * ptr;\n __host__ __device__ void\n operator()(T x){ ptr[(int) x] = 1; }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__ void for_each_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Function f)\n{\n thrust::for_each(exec, first, last, f);\n}\n\n\ntemplate\nvoid TestForEachDeviceSeq(const size_t n)\n{\n const size_t output_size = std::min((size_t) 10, 2 * n);\n \n thrust::host_vector h_input = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n h_input[i] = ((size_t) h_input[i]) % output_size;\n \n thrust::device_vector d_input = h_input;\n \n thrust::host_vector h_output(output_size, (T) 0);\n thrust::device_vector d_output(output_size, (T) 0);\n \n mark_present_for_each h_f;\n mark_present_for_each d_f;\n h_f.ptr = &h_output[0];\n d_f.ptr = (&d_output[0]).get();\n \n thrust::for_each(h_input.begin(), h_input.end(), h_f);\n \n for_each_kernel<<<1,1>>>(thrust::seq, d_input.begin(), d_input.end(), d_f);\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_EQUAL(h_output, d_output);\n}\nDECLARE_VARIABLE_UNITTEST(TestForEachDeviceSeq);\n\n\ntemplate\nvoid TestForEachDeviceDevice(const size_t n)\n{\n const size_t output_size = std::min((size_t) 10, 2 * n);\n \n thrust::host_vector h_input = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n h_input[i] = ((size_t) h_input[i]) % output_size;\n \n thrust::device_vector d_input = h_input;\n \n thrust::host_vector h_output(output_size, (T) 0);\n thrust::device_vector d_output(output_size, (T) 0);\n \n mark_present_for_each h_f;\n mark_present_for_each d_f;\n h_f.ptr = &h_output[0];\n d_f.ptr = (&d_output[0]).get();\n \n thrust::for_each(h_input.begin(), h_input.end(), h_f);\n \n for_each_kernel<<<1,1>>>(thrust::device, d_input.begin(), d_input.end(), d_f);\n {\n cudaError_t const err = cudaGetLastError();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_output, d_output);\n}\nDECLARE_VARIABLE_UNITTEST(TestForEachDeviceDevice);\n\n\ntemplate\n__global__\nvoid for_each_n_kernel(ExecutionPolicy exec, Iterator first, Size n, Function f)\n{\n thrust::for_each_n(exec, first, n, f);\n}\n\n\ntemplate\nvoid TestForEachNDeviceSeq(const size_t n)\n{\n const size_t output_size = std::min((size_t) 10, 2 * n);\n \n thrust::host_vector h_input = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n h_input[i] = static_cast(((size_t) h_input[i]) % output_size);\n \n thrust::device_vector d_input = h_input;\n \n thrust::host_vector h_output(output_size, (T) 0);\n thrust::device_vector d_output(output_size, (T) 0);\n \n mark_present_for_each h_f;\n mark_present_for_each d_f;\n h_f.ptr = &h_output[0];\n d_f.ptr = (&d_output[0]).get();\n \n thrust::for_each_n(h_input.begin(), h_input.size(), h_f);\n \n for_each_n_kernel<<<1,1>>>(thrust::seq, d_input.begin(), d_input.size(), d_f);\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_EQUAL(h_output, d_output);\n}\nDECLARE_VARIABLE_UNITTEST(TestForEachNDeviceSeq);\n\n\ntemplate\nvoid TestForEachNDeviceDevice(const size_t n)\n{\n const size_t output_size = std::min((size_t) 10, 2 * n);\n \n thrust::host_vector h_input = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n h_input[i] = static_cast(((size_t) h_input[i]) % output_size);\n \n thrust::device_vector d_input = h_input;\n \n thrust::host_vector h_output(output_size, (T) 0);\n thrust::device_vector d_output(output_size, (T) 0);\n \n mark_present_for_each h_f;\n mark_present_for_each d_f;\n h_f.ptr = &h_output[0];\n d_f.ptr = (&d_output[0]).get();\n \n thrust::for_each_n(h_input.begin(), h_input.size(), h_f);\n \n for_each_n_kernel<<<1,1>>>(thrust::device, d_input.begin(), d_input.size(), d_f);\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n ASSERT_EQUAL(h_output, d_output);\n}\nDECLARE_VARIABLE_UNITTEST(TestForEachNDeviceDevice);\n#endif\n\n\nvoid TestForEachCudaStreams()\n{\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n thrust::device_vector input(5);\n thrust::device_vector output(7, 0);\n \n input[0] = 3; input[1] = 2; input[2] = 3; input[3] = 4; input[4] = 6;\n \n mark_present_for_each f;\n f.ptr = thrust::raw_pointer_cast(output.data());\n \n thrust::for_each(thrust::cuda::par.on(s), input.begin(), input.end(), f);\n\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(output[0], 0);\n ASSERT_EQUAL(output[1], 0);\n ASSERT_EQUAL(output[2], 1);\n ASSERT_EQUAL(output[3], 1);\n ASSERT_EQUAL(output[4], 1);\n ASSERT_EQUAL(output[5], 0);\n ASSERT_EQUAL(output[6], 1);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestForEachCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\nstatic const size_t NUM_REGISTERS = 64;\n\ntemplate __host__ __device__ void f (int * x) { int temp = *x; f(x + 1); *x = temp;};\ntemplate <> __host__ __device__ void f<0>(int * /*x*/) { }\ntemplate \nstruct CopyFunctorWithManyRegisters\n{\n __host__ __device__\n void operator()(int * ptr)\n {\n f(ptr);\n }\n};\n\n\nvoid TestForEachLargeRegisterFootprint()\n{\n int current_device = -1;\n hipGetDevice(¤t_device);\n hipDeviceProp_t prop;\n hipGetDeviceProperties(&prop, current_device);\n\n thrust::device_vector data(NUM_REGISTERS, 12345);\n\n thrust::device_vector input(1, thrust::raw_pointer_cast(&data[0])); // length is irrelevant\n \n thrust::for_each(input.begin(), input.end(), CopyFunctorWithManyRegisters());\n}\nDECLARE_UNITTEST(TestForEachLargeRegisterFootprint);\n\n\nvoid TestForEachNLargeRegisterFootprint()\n{\n int current_device = -1;\n hipGetDevice(¤t_device);\n hipDeviceProp_t prop;\n hipGetDeviceProperties(&prop, current_device);\n\n thrust::device_vector data(NUM_REGISTERS, 12345);\n\n thrust::device_vector input(1, thrust::raw_pointer_cast(&data[0])); // length is irrelevant\n \n thrust::for_each_n(input.begin(), input.size(), CopyFunctorWithManyRegisters());\n}\nDECLARE_UNITTEST(TestForEachNLargeRegisterFootprint);\n\n\ntemplate \nstruct mark_present_for_each\n{\n T * ptr;\n __host__ __device__ void\n operator()(T x){ ptr[(int) x] = 1; }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__ void for_each_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Function f)\n{\n thrust::for_each(exec, first, last, f);\n}\n\n\ntemplate\nvoid TestForEachDeviceSeq(const size_t n)\n{\n const size_t output_size = std::min((size_t) 10, 2 * n);\n \n thrust::host_vector h_input = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n h_input[i] = ((size_t) h_input[i]) % output_size;\n \n thrust::device_vector d_input = h_input;\n \n thrust::host_vector h_output(output_size, (T) 0);\n thrust::device_vector d_output(output_size, (T) 0);\n \n mark_present_for_each h_f;\n mark_present_for_each d_f;\n h_f.ptr = &h_output[0];\n d_f.ptr = (&d_output[0]).get();\n \n thrust::for_each(h_input.begin(), h_input.end(), h_f);\n \n for_each_kernel<<<1,1>>>(thrust::seq, d_input.begin(), d_input.end(), d_f);\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_EQUAL(h_output, d_output);\n}\nDECLARE_VARIABLE_UNITTEST(TestForEachDeviceSeq);\n\n\ntemplate\nvoid TestForEachDeviceDevice(const size_t n)\n{\n const size_t output_size = std::min((size_t) 10, 2 * n);\n \n thrust::host_vector h_input = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n h_input[i] = ((size_t) h_input[i]) % output_size;\n \n thrust::device_vector d_input = h_input;\n \n thrust::host_vector h_output(output_size, (T) 0);\n thrust::device_vector d_output(output_size, (T) 0);\n \n mark_present_for_each h_f;\n mark_present_for_each d_f;\n h_f.ptr = &h_output[0];\n d_f.ptr = (&d_output[0]).get();\n \n thrust::for_each(h_input.begin(), h_input.end(), h_f);\n \n for_each_kernel<<<1,1>>>(thrust::device, d_input.begin(), d_input.end(), d_f);\n {\n hipError_t const err = hipGetLastError();\n ASSERT_EQUAL(hipSuccess, err);\n }\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_output, d_output);\n}\nDECLARE_VARIABLE_UNITTEST(TestForEachDeviceDevice);\n\n\ntemplate\n__global__\nvoid for_each_n_kernel(ExecutionPolicy exec, Iterator first, Size n, Function f)\n{\n thrust::for_each_n(exec, first, n, f);\n}\n\n\ntemplate\nvoid TestForEachNDeviceSeq(const size_t n)\n{\n const size_t output_size = std::min((size_t) 10, 2 * n);\n \n thrust::host_vector h_input = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n h_input[i] = static_cast(((size_t) h_input[i]) % output_size);\n \n thrust::device_vector d_input = h_input;\n \n thrust::host_vector h_output(output_size, (T) 0);\n thrust::device_vector d_output(output_size, (T) 0);\n \n mark_present_for_each h_f;\n mark_present_for_each d_f;\n h_f.ptr = &h_output[0];\n d_f.ptr = (&d_output[0]).get();\n \n thrust::for_each_n(h_input.begin(), h_input.size(), h_f);\n \n for_each_n_kernel<<<1,1>>>(thrust::seq, d_input.begin(), d_input.size(), d_f);\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_EQUAL(h_output, d_output);\n}\nDECLARE_VARIABLE_UNITTEST(TestForEachNDeviceSeq);\n\n\ntemplate\nvoid TestForEachNDeviceDevice(const size_t n)\n{\n const size_t output_size = std::min((size_t) 10, 2 * n);\n \n thrust::host_vector h_input = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n h_input[i] = static_cast(((size_t) h_input[i]) % output_size);\n \n thrust::device_vector d_input = h_input;\n \n thrust::host_vector h_output(output_size, (T) 0);\n thrust::device_vector d_output(output_size, (T) 0);\n \n mark_present_for_each h_f;\n mark_present_for_each d_f;\n h_f.ptr = &h_output[0];\n d_f.ptr = (&d_output[0]).get();\n \n thrust::for_each_n(h_input.begin(), h_input.size(), h_f);\n \n for_each_n_kernel<<<1,1>>>(thrust::device, d_input.begin(), d_input.size(), d_f);\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n ASSERT_EQUAL(h_output, d_output);\n}\nDECLARE_VARIABLE_UNITTEST(TestForEachNDeviceDevice);\n#endif\n\n\nvoid TestForEachCudaStreams()\n{\n hipStream_t s;\n hipStreamCreate(&s);\n \n thrust::device_vector input(5);\n thrust::device_vector output(7, 0);\n \n input[0] = 3; input[1] = 2; input[2] = 3; input[3] = 4; input[4] = 6;\n \n mark_present_for_each f;\n f.ptr = thrust::raw_pointer_cast(output.data());\n \n thrust::for_each(thrust::cuda::par.on(s), input.begin(), input.end(), f);\n\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(output[0], 0);\n ASSERT_EQUAL(output[1], 0);\n ASSERT_EQUAL(output[2], 1);\n ASSERT_EQUAL(output[3], 1);\n ASSERT_EQUAL(output[4], 1);\n ASSERT_EQUAL(output[5], 0);\n ASSERT_EQUAL(output[6], 1);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestForEachCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/for_each.cu"} {"input": "#include \n#include \n#include \n#include \n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid gather_kernel(ExecutionPolicy exec, Iterator1 map_first, Iterator1 map_last, Iterator2 elements_first, Iterator3 result)\n{\n thrust::gather(exec, map_first, map_last, elements_first, result);\n}\n\n\ntemplate\nvoid TestGatherDevice(ExecutionPolicy exec, const size_t n)\n{\n const size_t source_size = std::min((size_t) 10, 2 * n);\n \n // source vectors to gather from\n thrust::host_vector h_source = unittest::random_samples(source_size);\n thrust::device_vector d_source = h_source;\n \n // gather indices\n thrust::host_vector h_map = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n h_map[i] = h_map[i] % source_size;\n \n thrust::device_vector d_map = h_map;\n \n // gather destination\n thrust::host_vector h_output(n);\n thrust::device_vector d_output(n);\n \n thrust::gather(h_map.begin(), h_map.end(), h_source.begin(), h_output.begin());\n\n gather_kernel<<<1,1>>>(exec, d_map.begin(), d_map.end(), d_source.begin(), d_output.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_output, d_output);\n}\n\ntemplate\nvoid TestGatherDeviceSeq(const size_t n)\n{\n TestGatherDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGatherDeviceSeq);\n\ntemplate\nvoid TestGatherDeviceDevice(const size_t n)\n{\n TestGatherDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGatherDeviceDevice);\n#endif\n\n\nvoid TestGatherCudaStreams()\n{\n thrust::device_vector map(5); // gather indices\n thrust::device_vector src(8); // source vector\n thrust::device_vector dst(5); // destination vector\n \n map[0] = 6; map[1] = 2; map[2] = 1; map[3] = 7; map[4] = 2;\n src[0] = 0; src[1] = 1; src[2] = 2; src[3] = 3; src[4] = 4; src[5] = 5; src[6] = 6; src[7] = 7;\n dst[0] = 0; dst[1] = 0; dst[2] = 0; dst[3] = 0; dst[4] = 0;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n thrust::gather(thrust::cuda::par.on(s), map.begin(), map.end(), src.begin(), dst.begin());\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(dst[0], 6);\n ASSERT_EQUAL(dst[1], 2);\n ASSERT_EQUAL(dst[2], 1);\n ASSERT_EQUAL(dst[3], 7);\n ASSERT_EQUAL(dst[4], 2);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestGatherCudaStreams);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid gather_if_kernel(ExecutionPolicy exec, Iterator1 map_first, Iterator1 map_last, Iterator2 stencil_first, Iterator3 elements_first, Iterator4 result, Predicate pred)\n{\n thrust::gather_if(exec, map_first, map_last, stencil_first, elements_first, result, pred);\n}\n\n\ntemplate\nstruct is_even_gather_if\n{\n __host__ __device__\n bool operator()(const T i) const\n { \n return (i % 2) == 0;\n }\n};\n\n\ntemplate\nvoid TestGatherIfDevice(ExecutionPolicy exec, const size_t n)\n{\n const size_t source_size = std::min((size_t) 10, 2 * n);\n \n // source vectors to gather from\n thrust::host_vector h_source = unittest::random_samples(source_size);\n thrust::device_vector d_source = h_source;\n \n // gather indices\n thrust::host_vector h_map = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n h_map[i] = h_map[i] % source_size;\n \n thrust::device_vector d_map = h_map;\n \n // gather stencil\n thrust::host_vector h_stencil = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n h_stencil[i] = h_stencil[i] % 2;\n \n thrust::device_vector d_stencil = h_stencil;\n \n // gather destination\n thrust::host_vector h_output(n);\n thrust::device_vector d_output(n);\n \n thrust::gather_if(h_map.begin(), h_map.end(), h_stencil.begin(), h_source.begin(), h_output.begin(), is_even_gather_if());\n\n gather_if_kernel<<<1,1>>>(exec, d_map.begin(), d_map.end(), d_stencil.begin(), d_source.begin(), d_output.begin(), is_even_gather_if());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_output, d_output);\n}\n\ntemplate\nvoid TestGatherIfDeviceSeq(const size_t n)\n{\n TestGatherIfDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGatherIfDeviceSeq);\n\ntemplate\nvoid TestGatherIfDeviceDevice(const size_t n)\n{\n TestGatherIfDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGatherIfDeviceDevice);\n#endif\n\nvoid TestGatherIfCudaStreams(void)\n{\n thrust::device_vector flg(5); // predicate array\n thrust::device_vector map(5); // gather indices\n thrust::device_vector src(8); // source vector\n thrust::device_vector dst(5); // destination vector\n \n flg[0] = 0; flg[1] = 1; flg[2] = 0; flg[3] = 1; flg[4] = 0;\n map[0] = 6; map[1] = 2; map[2] = 1; map[3] = 7; map[4] = 2;\n src[0] = 0; src[1] = 1; src[2] = 2; src[3] = 3; src[4] = 4; src[5] = 5; src[6] = 6; src[7] = 7;\n dst[0] = 0; dst[1] = 0; dst[2] = 0; dst[3] = 0; dst[4] = 0;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n thrust::gather_if(thrust::cuda::par.on(s), map.begin(), map.end(), flg.begin(), src.begin(), dst.begin());\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(dst[0], 0);\n ASSERT_EQUAL(dst[1], 2);\n ASSERT_EQUAL(dst[2], 0);\n ASSERT_EQUAL(dst[3], 7);\n ASSERT_EQUAL(dst[4], 0);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestGatherIfCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid gather_kernel(ExecutionPolicy exec, Iterator1 map_first, Iterator1 map_last, Iterator2 elements_first, Iterator3 result)\n{\n thrust::gather(exec, map_first, map_last, elements_first, result);\n}\n\n\ntemplate\nvoid TestGatherDevice(ExecutionPolicy exec, const size_t n)\n{\n const size_t source_size = std::min((size_t) 10, 2 * n);\n \n // source vectors to gather from\n thrust::host_vector h_source = unittest::random_samples(source_size);\n thrust::device_vector d_source = h_source;\n \n // gather indices\n thrust::host_vector h_map = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n h_map[i] = h_map[i] % source_size;\n \n thrust::device_vector d_map = h_map;\n \n // gather destination\n thrust::host_vector h_output(n);\n thrust::device_vector d_output(n);\n \n thrust::gather(h_map.begin(), h_map.end(), h_source.begin(), h_output.begin());\n\n gather_kernel<<<1,1>>>(exec, d_map.begin(), d_map.end(), d_source.begin(), d_output.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_output, d_output);\n}\n\ntemplate\nvoid TestGatherDeviceSeq(const size_t n)\n{\n TestGatherDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGatherDeviceSeq);\n\ntemplate\nvoid TestGatherDeviceDevice(const size_t n)\n{\n TestGatherDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGatherDeviceDevice);\n#endif\n\n\nvoid TestGatherCudaStreams()\n{\n thrust::device_vector map(5); // gather indices\n thrust::device_vector src(8); // source vector\n thrust::device_vector dst(5); // destination vector\n \n map[0] = 6; map[1] = 2; map[2] = 1; map[3] = 7; map[4] = 2;\n src[0] = 0; src[1] = 1; src[2] = 2; src[3] = 3; src[4] = 4; src[5] = 5; src[6] = 6; src[7] = 7;\n dst[0] = 0; dst[1] = 0; dst[2] = 0; dst[3] = 0; dst[4] = 0;\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n thrust::gather(thrust::cuda::par.on(s), map.begin(), map.end(), src.begin(), dst.begin());\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(dst[0], 6);\n ASSERT_EQUAL(dst[1], 2);\n ASSERT_EQUAL(dst[2], 1);\n ASSERT_EQUAL(dst[3], 7);\n ASSERT_EQUAL(dst[4], 2);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestGatherCudaStreams);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid gather_if_kernel(ExecutionPolicy exec, Iterator1 map_first, Iterator1 map_last, Iterator2 stencil_first, Iterator3 elements_first, Iterator4 result, Predicate pred)\n{\n thrust::gather_if(exec, map_first, map_last, stencil_first, elements_first, result, pred);\n}\n\n\ntemplate\nstruct is_even_gather_if\n{\n __host__ __device__\n bool operator()(const T i) const\n { \n return (i % 2) == 0;\n }\n};\n\n\ntemplate\nvoid TestGatherIfDevice(ExecutionPolicy exec, const size_t n)\n{\n const size_t source_size = std::min((size_t) 10, 2 * n);\n \n // source vectors to gather from\n thrust::host_vector h_source = unittest::random_samples(source_size);\n thrust::device_vector d_source = h_source;\n \n // gather indices\n thrust::host_vector h_map = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n h_map[i] = h_map[i] % source_size;\n \n thrust::device_vector d_map = h_map;\n \n // gather stencil\n thrust::host_vector h_stencil = unittest::random_integers(n);\n \n for(size_t i = 0; i < n; i++)\n h_stencil[i] = h_stencil[i] % 2;\n \n thrust::device_vector d_stencil = h_stencil;\n \n // gather destination\n thrust::host_vector h_output(n);\n thrust::device_vector d_output(n);\n \n thrust::gather_if(h_map.begin(), h_map.end(), h_stencil.begin(), h_source.begin(), h_output.begin(), is_even_gather_if());\n\n gather_if_kernel<<<1,1>>>(exec, d_map.begin(), d_map.end(), d_stencil.begin(), d_source.begin(), d_output.begin(), is_even_gather_if());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_output, d_output);\n}\n\ntemplate\nvoid TestGatherIfDeviceSeq(const size_t n)\n{\n TestGatherIfDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGatherIfDeviceSeq);\n\ntemplate\nvoid TestGatherIfDeviceDevice(const size_t n)\n{\n TestGatherIfDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestGatherIfDeviceDevice);\n#endif\n\nvoid TestGatherIfCudaStreams(void)\n{\n thrust::device_vector flg(5); // predicate array\n thrust::device_vector map(5); // gather indices\n thrust::device_vector src(8); // source vector\n thrust::device_vector dst(5); // destination vector\n \n flg[0] = 0; flg[1] = 1; flg[2] = 0; flg[3] = 1; flg[4] = 0;\n map[0] = 6; map[1] = 2; map[2] = 1; map[3] = 7; map[4] = 2;\n src[0] = 0; src[1] = 1; src[2] = 2; src[3] = 3; src[4] = 4; src[5] = 5; src[6] = 6; src[7] = 7;\n dst[0] = 0; dst[1] = 0; dst[2] = 0; dst[3] = 0; dst[4] = 0;\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n thrust::gather_if(thrust::cuda::par.on(s), map.begin(), map.end(), flg.begin(), src.begin(), dst.begin());\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(dst[0], 0);\n ASSERT_EQUAL(dst[1], 2);\n ASSERT_EQUAL(dst[2], 0);\n ASSERT_EQUAL(dst[3], 7);\n ASSERT_EQUAL(dst[4], 0);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestGatherIfCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/gather.cu"} {"input": "#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid is_partitioned_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Predicate pred, Iterator2 result)\n{\n *result = thrust::is_partitioned(exec, first, last, pred);\n}\n\n\ntemplate\nstruct is_even\n{\n __host__ __device__\n bool operator()(T x) const { return ((int) x % 2) == 0; }\n};\n\n\ntemplate\nvoid TestIsPartitionedDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n\n n = thrust::max(n, 2);\n\n thrust::device_vector v = unittest::random_integers(n);\n\n thrust::device_vector result(1);\n\n v[0] = 1;\n v[1] = 0;\n\n is_partitioned_kernel<<<1,1>>>(exec, v.begin(), v.end(), is_even(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n\n thrust::partition(v.begin(), v.end(), is_even());\n\n is_partitioned_kernel<<<1,1>>>(exec, v.begin(), v.end(), is_even(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n}\n\n\nvoid TestIsPartitionedDeviceSeq()\n{\n TestIsPartitionedDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestIsPartitionedDeviceSeq);\n\n\nvoid TestIsPartitionedDeviceDevice()\n{\n TestIsPartitionedDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestIsPartitionedDeviceDevice);\n#endif\n\n\nvoid TestIsPartitionedCudaStreams()\n{\n thrust::device_vector v(4);\n v[0] = 1; v[1] = 1; v[2] = 1; v[3] = 0;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n // empty partition\n ASSERT_EQUAL_QUIET(true, thrust::is_partitioned(thrust::cuda::par.on(s), v.begin(), v.begin(), thrust::identity()));\n\n // one element true partition\n ASSERT_EQUAL_QUIET(true, thrust::is_partitioned(thrust::cuda::par.on(s), v.begin(), v.begin() + 1, thrust::identity()));\n\n // just true partition\n ASSERT_EQUAL_QUIET(true, thrust::is_partitioned(thrust::cuda::par.on(s), v.begin(), v.begin() + 2, thrust::identity()));\n\n // both true & false partitions\n ASSERT_EQUAL_QUIET(true, thrust::is_partitioned(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()));\n\n // one element false partition\n ASSERT_EQUAL_QUIET(true, thrust::is_partitioned(thrust::cuda::par.on(s), v.begin() + 3, v.end(), thrust::identity()));\n\n v[0] = 1; v[1] = 0; v[2] = 1; v[3] = 1;\n\n // not partitioned\n ASSERT_EQUAL_QUIET(false, thrust::is_partitioned(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()));\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestIsPartitionedCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid is_partitioned_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Predicate pred, Iterator2 result)\n{\n *result = thrust::is_partitioned(exec, first, last, pred);\n}\n\n\ntemplate\nstruct is_even\n{\n __host__ __device__\n bool operator()(T x) const { return ((int) x % 2) == 0; }\n};\n\n\ntemplate\nvoid TestIsPartitionedDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n\n n = thrust::max(n, 2);\n\n thrust::device_vector v = unittest::random_integers(n);\n\n thrust::device_vector result(1);\n\n v[0] = 1;\n v[1] = 0;\n\n is_partitioned_kernel<<<1,1>>>(exec, v.begin(), v.end(), is_even(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n\n thrust::partition(v.begin(), v.end(), is_even());\n\n is_partitioned_kernel<<<1,1>>>(exec, v.begin(), v.end(), is_even(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n}\n\n\nvoid TestIsPartitionedDeviceSeq()\n{\n TestIsPartitionedDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestIsPartitionedDeviceSeq);\n\n\nvoid TestIsPartitionedDeviceDevice()\n{\n TestIsPartitionedDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestIsPartitionedDeviceDevice);\n#endif\n\n\nvoid TestIsPartitionedCudaStreams()\n{\n thrust::device_vector v(4);\n v[0] = 1; v[1] = 1; v[2] = 1; v[3] = 0;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n // empty partition\n ASSERT_EQUAL_QUIET(true, thrust::is_partitioned(thrust::cuda::par.on(s), v.begin(), v.begin(), thrust::identity()));\n\n // one element true partition\n ASSERT_EQUAL_QUIET(true, thrust::is_partitioned(thrust::cuda::par.on(s), v.begin(), v.begin() + 1, thrust::identity()));\n\n // just true partition\n ASSERT_EQUAL_QUIET(true, thrust::is_partitioned(thrust::cuda::par.on(s), v.begin(), v.begin() + 2, thrust::identity()));\n\n // both true & false partitions\n ASSERT_EQUAL_QUIET(true, thrust::is_partitioned(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()));\n\n // one element false partition\n ASSERT_EQUAL_QUIET(true, thrust::is_partitioned(thrust::cuda::par.on(s), v.begin() + 3, v.end(), thrust::identity()));\n\n v[0] = 1; v[1] = 0; v[2] = 1; v[3] = 1;\n\n // not partitioned\n ASSERT_EQUAL_QUIET(false, thrust::is_partitioned(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()));\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestIsPartitionedCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/is_partitioned.cu"} {"input": "#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid equal_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2, Iterator3 result)\n{\n *result = thrust::equal(exec, first1, last1, first2);\n}\n\n\ntemplate\n__global__\nvoid equal_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2, BinaryPredicate pred, Iterator3 result)\n{\n *result = thrust::equal(exec, first1, last1, first2, pred);\n}\n\n\ntemplate\nvoid TestEqualDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::device_vector d_data1 = unittest::random_samples(n);\n thrust::device_vector d_data2 = unittest::random_samples(n);\n thrust::device_vector d_result(1, false);\n \n //empty ranges\n equal_kernel<<<1,1>>>(exec, d_data1.begin(), d_data1.begin(), d_data1.begin(), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(d_result[0], true);\n \n //symmetric cases\n equal_kernel<<<1,1>>>(exec, d_data1.begin(), d_data1.end(), d_data1.begin(), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(d_result[0], true);\n \n if(n > 0)\n {\n d_data1[0] = 0; d_data2[0] = 1;\n \n //different vectors\n equal_kernel<<<1,1>>>(exec, d_data1.begin(), d_data1.end(), d_data2.begin(), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(d_result[0], false);\n \n //different predicates\n equal_kernel<<<1,1>>>(exec, d_data1.begin(), d_data1.begin() + 1, d_data2.begin(), thrust::less(), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(d_result[0], true);\n\n equal_kernel<<<1,1>>>(exec, d_data1.begin(), d_data1.begin() + 1, d_data2.begin(), thrust::greater(), d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(d_result[0], false);\n }\n}\n\n\ntemplate\nvoid TestEqualDeviceSeq(const size_t n)\n{\n TestEqualDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestEqualDeviceSeq);\n\n\ntemplate\nvoid TestEqualDeviceDevice(const size_t n)\n{\n TestEqualDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestEqualDeviceDevice);\n#endif\n\n\nvoid TestEqualCudaStreams()\n{\n thrust::device_vector v1(5);\n thrust::device_vector v2(5);\n v1[0] = 5; v1[1] = 2; v1[2] = 0; v1[3] = 0; v1[4] = 0;\n v2[0] = 5; v2[1] = 2; v2[2] = 0; v2[3] = 6; v2[4] = 1;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.end(), v1.begin()), true);\n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.end(), v2.begin()), false);\n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v2.begin(), v2.end(), v2.begin()), true);\n \n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.begin() + 0, v1.begin()), true);\n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.begin() + 1, v1.begin()), true);\n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.begin() + 3, v2.begin()), true);\n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.begin() + 4, v2.begin()), false);\n \n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.end(), v2.begin(), thrust::less_equal()), true);\n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.end(), v2.begin(), thrust::greater()), false);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestEqualCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid equal_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2, Iterator3 result)\n{\n *result = thrust::equal(exec, first1, last1, first2);\n}\n\n\ntemplate\n__global__\nvoid equal_kernel(ExecutionPolicy exec, Iterator1 first1, Iterator1 last1, Iterator2 first2, BinaryPredicate pred, Iterator3 result)\n{\n *result = thrust::equal(exec, first1, last1, first2, pred);\n}\n\n\ntemplate\nvoid TestEqualDevice(ExecutionPolicy exec, const size_t n)\n{\n thrust::device_vector d_data1 = unittest::random_samples(n);\n thrust::device_vector d_data2 = unittest::random_samples(n);\n thrust::device_vector d_result(1, false);\n \n //empty ranges\n equal_kernel<<<1,1>>>(exec, d_data1.begin(), d_data1.begin(), d_data1.begin(), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(d_result[0], true);\n \n //symmetric cases\n equal_kernel<<<1,1>>>(exec, d_data1.begin(), d_data1.end(), d_data1.begin(), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(d_result[0], true);\n \n if(n > 0)\n {\n d_data1[0] = 0; d_data2[0] = 1;\n \n //different vectors\n equal_kernel<<<1,1>>>(exec, d_data1.begin(), d_data1.end(), d_data2.begin(), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(d_result[0], false);\n \n //different predicates\n equal_kernel<<<1,1>>>(exec, d_data1.begin(), d_data1.begin() + 1, d_data2.begin(), thrust::less(), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(d_result[0], true);\n\n equal_kernel<<<1,1>>>(exec, d_data1.begin(), d_data1.begin() + 1, d_data2.begin(), thrust::greater(), d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(d_result[0], false);\n }\n}\n\n\ntemplate\nvoid TestEqualDeviceSeq(const size_t n)\n{\n TestEqualDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestEqualDeviceSeq);\n\n\ntemplate\nvoid TestEqualDeviceDevice(const size_t n)\n{\n TestEqualDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestEqualDeviceDevice);\n#endif\n\n\nvoid TestEqualCudaStreams()\n{\n thrust::device_vector v1(5);\n thrust::device_vector v2(5);\n v1[0] = 5; v1[1] = 2; v1[2] = 0; v1[3] = 0; v1[4] = 0;\n v2[0] = 5; v2[1] = 2; v2[2] = 0; v2[3] = 6; v2[4] = 1;\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.end(), v1.begin()), true);\n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.end(), v2.begin()), false);\n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v2.begin(), v2.end(), v2.begin()), true);\n \n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.begin() + 0, v1.begin()), true);\n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.begin() + 1, v1.begin()), true);\n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.begin() + 3, v2.begin()), true);\n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.begin() + 4, v2.begin()), false);\n \n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.end(), v2.begin(), thrust::less_equal()), true);\n ASSERT_EQUAL(thrust::equal(thrust::cuda::par.on(s), v1.begin(), v1.end(), v2.begin(), thrust::greater()), false);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestEqualCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/equal.cu"} {"input": "#include \n#include \n#include \n#include \n\ntemplate\nstruct is_even\n{\n __host__ __device__\n bool operator()(T x) { return (static_cast(x) & 1) == 0; }\n};\n\n\ntemplate\nstruct mod_3\n{\n __host__ __device__\n unsigned int operator()(T x) { return static_cast(x) % 3; }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__ void copy_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, Predicate pred, Iterator3 result2)\n{\n *result2 = thrust::copy_if(exec, first, last, result1, pred);\n}\n\n\ntemplate\nvoid TestCopyIfDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n typename thrust::host_vector::iterator h_new_end;\n typename thrust::device_vector::iterator d_new_end;\n\n thrust::device_vector<\n typename thrust::device_vector::iterator\n > d_new_end_vec(1);\n \n // test with Predicate that returns a bool\n {\n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n \n h_new_end = thrust::copy_if(h_data.begin(), h_data.end(), h_result.begin(), is_even());\n\n copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin(), is_even(), d_new_end_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n d_new_end = d_new_end_vec[0];\n \n h_result.resize(h_new_end - h_result.begin());\n d_result.resize(d_new_end - d_result.begin());\n \n ASSERT_EQUAL(h_result, d_result);\n }\n \n // test with Predicate that returns a non-bool\n {\n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n \n h_new_end = thrust::copy_if(h_data.begin(), h_data.end(), h_result.begin(), mod_3());\n\n copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin(), mod_3(), d_new_end_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n d_new_end = d_new_end_vec[0];\n \n h_result.resize(h_new_end - h_result.begin());\n d_result.resize(d_new_end - d_result.begin());\n \n ASSERT_EQUAL(h_result, d_result);\n }\n}\n\n\nvoid TestCopyIfDeviceSeq()\n{\n TestCopyIfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestCopyIfDeviceSeq);\n\n\nvoid TestCopyIfDeviceDevice()\n{\n TestCopyIfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestCopyIfDeviceDevice);\n\n\nvoid TestCopyIfDeviceNoSync()\n{\n TestCopyIfDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestCopyIfDeviceNoSync);\n#endif\n\ntemplate\nvoid TestCopyIfCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n Vector result(5);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n Vector::iterator end = thrust::copy_if(policy.on(s),\n data.begin(), \n data.end(), \n result.begin(),\n is_even());\n\n ASSERT_EQUAL(end - result.begin(), 2);\n\n ASSERT_EQUAL(result[0], 2);\n ASSERT_EQUAL(result[1], 2);\n\n cudaStreamDestroy(s);\n}\n\nvoid TestCopyIfCudaStreamsSync(){\n TestCopyIfCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestCopyIfCudaStreamsSync);\n\nvoid TestCopyIfCudaStreamsNoSync(){\n TestCopyIfCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestCopyIfCudaStreamsNoSync);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__ void copy_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Iterator3 result1, Predicate pred, Iterator4 result2)\n{\n *result2 = thrust::copy_if(exec, first, last, stencil_first, result1, pred);\n}\n\n\ntemplate\nvoid TestCopyIfStencilDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data(n); thrust::sequence(h_data.begin(), h_data.end());\n thrust::device_vector d_data(n); thrust::sequence(d_data.begin(), d_data.end()); \n \n thrust::host_vector h_stencil = unittest::random_integers(n);\n thrust::device_vector d_stencil = unittest::random_integers(n);\n \n typename thrust::host_vector::iterator h_new_end;\n typename thrust::device_vector::iterator d_new_end;\n\n thrust::device_vector<\n typename thrust::device_vector::iterator\n > d_new_end_vec(1);\n \n // test with Predicate that returns a bool\n {\n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n \n h_new_end = thrust::copy_if(h_data.begin(), h_data.end(), h_result.begin(), is_even());\n\n copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin(), is_even(), d_new_end_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n d_new_end = d_new_end_vec[0];\n \n h_result.resize(h_new_end - h_result.begin());\n d_result.resize(d_new_end - d_result.begin());\n \n ASSERT_EQUAL(h_result, d_result);\n }\n \n // test with Predicate that returns a non-bool\n {\n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n \n h_new_end = thrust::copy_if(h_data.begin(), h_data.end(), h_result.begin(), mod_3());\n\n copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin(), mod_3(), d_new_end_vec.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n d_new_end = d_new_end_vec[0];\n \n h_result.resize(h_new_end - h_result.begin());\n d_result.resize(d_new_end - d_result.begin());\n \n ASSERT_EQUAL(h_result, d_result);\n }\n}\n\n\nvoid TestCopyIfStencilDeviceSeq()\n{\n TestCopyIfStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestCopyIfStencilDeviceSeq);\n\n\nvoid TestCopyIfStencilDeviceDevice()\n{\n TestCopyIfStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestCopyIfStencilDeviceDevice);\n\n\nvoid TestCopyIfStencilDeviceNoSync()\n{\n TestCopyIfStencilDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestCopyIfStencilDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestCopyIfStencilCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n Vector result(5);\n\n Vector stencil(5);\n stencil[0] = 0;\n stencil[1] = 1;\n stencil[2] = 0;\n stencil[3] = 0;\n stencil[4] = 1;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n Vector::iterator end = thrust::copy_if(policy.on(s),\n data.begin(), \n data.end(),\n stencil.begin(),\n result.begin(),\n thrust::identity());\n\n ASSERT_EQUAL(end - result.begin(), 2);\n\n ASSERT_EQUAL(result[0], 2);\n ASSERT_EQUAL(result[1], 2);\n\n cudaStreamDestroy(s);\n}\n\nvoid TestCopyIfStencilCudaStreamsSync()\n{\n TestCopyIfStencilCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestCopyIfStencilCudaStreamsSync);\n\n\nvoid TestCopyIfStencilCudaStreamsNoSync()\n{\n TestCopyIfStencilCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestCopyIfStencilCudaStreamsNoSync);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\ntemplate\nstruct is_even\n{\n __host__ __device__\n bool operator()(T x) { return (static_cast(x) & 1) == 0; }\n};\n\n\ntemplate\nstruct mod_3\n{\n __host__ __device__\n unsigned int operator()(T x) { return static_cast(x) % 3; }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__ void copy_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, Predicate pred, Iterator3 result2)\n{\n *result2 = thrust::copy_if(exec, first, last, result1, pred);\n}\n\n\ntemplate\nvoid TestCopyIfDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n typename thrust::host_vector::iterator h_new_end;\n typename thrust::device_vector::iterator d_new_end;\n\n thrust::device_vector<\n typename thrust::device_vector::iterator\n > d_new_end_vec(1);\n \n // test with Predicate that returns a bool\n {\n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n \n h_new_end = thrust::copy_if(h_data.begin(), h_data.end(), h_result.begin(), is_even());\n\n copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin(), is_even(), d_new_end_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n d_new_end = d_new_end_vec[0];\n \n h_result.resize(h_new_end - h_result.begin());\n d_result.resize(d_new_end - d_result.begin());\n \n ASSERT_EQUAL(h_result, d_result);\n }\n \n // test with Predicate that returns a non-bool\n {\n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n \n h_new_end = thrust::copy_if(h_data.begin(), h_data.end(), h_result.begin(), mod_3());\n\n copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin(), mod_3(), d_new_end_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n d_new_end = d_new_end_vec[0];\n \n h_result.resize(h_new_end - h_result.begin());\n d_result.resize(d_new_end - d_result.begin());\n \n ASSERT_EQUAL(h_result, d_result);\n }\n}\n\n\nvoid TestCopyIfDeviceSeq()\n{\n TestCopyIfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestCopyIfDeviceSeq);\n\n\nvoid TestCopyIfDeviceDevice()\n{\n TestCopyIfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestCopyIfDeviceDevice);\n\n\nvoid TestCopyIfDeviceNoSync()\n{\n TestCopyIfDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestCopyIfDeviceNoSync);\n#endif\n\ntemplate\nvoid TestCopyIfCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n Vector result(5);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n Vector::iterator end = thrust::copy_if(policy.on(s),\n data.begin(), \n data.end(), \n result.begin(),\n is_even());\n\n ASSERT_EQUAL(end - result.begin(), 2);\n\n ASSERT_EQUAL(result[0], 2);\n ASSERT_EQUAL(result[1], 2);\n\n hipStreamDestroy(s);\n}\n\nvoid TestCopyIfCudaStreamsSync(){\n TestCopyIfCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestCopyIfCudaStreamsSync);\n\nvoid TestCopyIfCudaStreamsNoSync(){\n TestCopyIfCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestCopyIfCudaStreamsNoSync);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__ void copy_if_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Iterator3 result1, Predicate pred, Iterator4 result2)\n{\n *result2 = thrust::copy_if(exec, first, last, stencil_first, result1, pred);\n}\n\n\ntemplate\nvoid TestCopyIfStencilDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n thrust::host_vector h_data(n); thrust::sequence(h_data.begin(), h_data.end());\n thrust::device_vector d_data(n); thrust::sequence(d_data.begin(), d_data.end()); \n \n thrust::host_vector h_stencil = unittest::random_integers(n);\n thrust::device_vector d_stencil = unittest::random_integers(n);\n \n typename thrust::host_vector::iterator h_new_end;\n typename thrust::device_vector::iterator d_new_end;\n\n thrust::device_vector<\n typename thrust::device_vector::iterator\n > d_new_end_vec(1);\n \n // test with Predicate that returns a bool\n {\n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n \n h_new_end = thrust::copy_if(h_data.begin(), h_data.end(), h_result.begin(), is_even());\n\n copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin(), is_even(), d_new_end_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n d_new_end = d_new_end_vec[0];\n \n h_result.resize(h_new_end - h_result.begin());\n d_result.resize(d_new_end - d_result.begin());\n \n ASSERT_EQUAL(h_result, d_result);\n }\n \n // test with Predicate that returns a non-bool\n {\n thrust::host_vector h_result(n);\n thrust::device_vector d_result(n);\n \n h_new_end = thrust::copy_if(h_data.begin(), h_data.end(), h_result.begin(), mod_3());\n\n copy_if_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), d_result.begin(), mod_3(), d_new_end_vec.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n d_new_end = d_new_end_vec[0];\n \n h_result.resize(h_new_end - h_result.begin());\n d_result.resize(d_new_end - d_result.begin());\n \n ASSERT_EQUAL(h_result, d_result);\n }\n}\n\n\nvoid TestCopyIfStencilDeviceSeq()\n{\n TestCopyIfStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestCopyIfStencilDeviceSeq);\n\n\nvoid TestCopyIfStencilDeviceDevice()\n{\n TestCopyIfStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestCopyIfStencilDeviceDevice);\n\n\nvoid TestCopyIfStencilDeviceNoSync()\n{\n TestCopyIfStencilDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestCopyIfStencilDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestCopyIfStencilCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 3; \n data[4] = 2; \n\n Vector result(5);\n\n Vector stencil(5);\n stencil[0] = 0;\n stencil[1] = 1;\n stencil[2] = 0;\n stencil[3] = 0;\n stencil[4] = 1;\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n Vector::iterator end = thrust::copy_if(policy.on(s),\n data.begin(), \n data.end(),\n stencil.begin(),\n result.begin(),\n thrust::identity());\n\n ASSERT_EQUAL(end - result.begin(), 2);\n\n ASSERT_EQUAL(result[0], 2);\n ASSERT_EQUAL(result[1], 2);\n\n hipStreamDestroy(s);\n}\n\nvoid TestCopyIfStencilCudaStreamsSync()\n{\n TestCopyIfStencilCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestCopyIfStencilCudaStreamsSync);\n\n\nvoid TestCopyIfStencilCudaStreamsNoSync()\n{\n TestCopyIfStencilCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestCopyIfStencilCudaStreamsNoSync);\n\n", "source": "/kaggle/working/thrust/testing/cuda/copy_if.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid is_sorted_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Iterator2 result)\n{\n *result = thrust::is_sorted(exec, first, last);\n}\n\n\ntemplate\nvoid TestIsSortedDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n\n thrust::device_vector v = unittest::random_integers(n);\n\n thrust::device_vector result(1);\n\n v[0] = 1;\n v[1] = 0;\n\n is_sorted_kernel<<<1,1>>>(exec, v.begin(), v.end(), result.begin());\n\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n\n thrust::sort(v.begin(), v.end());\n\n is_sorted_kernel<<<1,1>>>(exec, v.begin(), v.end(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n}\n\nvoid TestIsSortedDeviceSeq()\n{\n TestIsSortedDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestIsSortedDeviceSeq);\n\n\nvoid TestIsSortedDeviceDevice()\n{\n TestIsSortedDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestIsSortedDeviceDevice);\n#endif\n\n\nvoid TestIsSortedCudaStreams()\n{\n thrust::device_vector v(4);\n v[0] = 0; v[1] = 5; v[2] = 8; v[3] = 0;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 0), true);\n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 1), true);\n \n // the following line crashes gcc 4.3\n#if (__GNUC__ == 4) && (__GNUC_MINOR__ == 3)\n // do nothing\n#else\n // compile this line on other compilers\n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 2), true);\n#endif // GCC\n\n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 3), true);\n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 4), false);\n \n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 3, thrust::less()), true);\n \n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 1, thrust::greater()), true);\n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 4, thrust::greater()), false);\n \n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.end()), false);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestIsSortedCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid is_sorted_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Iterator2 result)\n{\n *result = thrust::is_sorted(exec, first, last);\n}\n\n\ntemplate\nvoid TestIsSortedDevice(ExecutionPolicy exec)\n{\n size_t n = 1000;\n\n thrust::device_vector v = unittest::random_integers(n);\n\n thrust::device_vector result(1);\n\n v[0] = 1;\n v[1] = 0;\n\n is_sorted_kernel<<<1,1>>>(exec, v.begin(), v.end(), result.begin());\n\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n\n thrust::sort(v.begin(), v.end());\n\n is_sorted_kernel<<<1,1>>>(exec, v.begin(), v.end(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n}\n\nvoid TestIsSortedDeviceSeq()\n{\n TestIsSortedDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestIsSortedDeviceSeq);\n\n\nvoid TestIsSortedDeviceDevice()\n{\n TestIsSortedDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestIsSortedDeviceDevice);\n#endif\n\n\nvoid TestIsSortedCudaStreams()\n{\n thrust::device_vector v(4);\n v[0] = 0; v[1] = 5; v[2] = 8; v[3] = 0;\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 0), true);\n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 1), true);\n \n // the following line crashes gcc 4.3\n#if (__GNUC__ == 4) && (__GNUC_MINOR__ == 3)\n // do nothing\n#else\n // compile this line on other compilers\n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 2), true);\n#endif // GCC\n\n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 3), true);\n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 4), false);\n \n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 3, thrust::less()), true);\n \n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 1, thrust::greater()), true);\n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.begin() + 4, thrust::greater()), false);\n \n ASSERT_EQUAL(thrust::is_sorted(thrust::cuda::par.on(s), v.begin(), v.end()), false);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestIsSortedCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/is_sorted.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n\n\ntemplate\nbool are_same_type(const T1 &, const T2 &)\n{\n return false;\n}\n\n\ntemplate\nbool are_same_type(const T &, const T &)\n{\n return true;\n}\n\n\nvoid TestSelectSystemCudaToCpp()\n{\n using thrust::system::detail::generic::select_system;\n\n thrust::cuda::tag cuda_tag;\n thrust::cpp::tag cpp_tag;\n thrust::cuda_cub::cross_system cuda_to_cpp(cuda_tag, cpp_tag);\n\n // select_system(cuda::tag, thrust::host_system_tag) should return cuda_to_cpp\n bool is_cuda_to_cpp = are_same_type(cuda_to_cpp, select_system(cuda_tag, cpp_tag));\n ASSERT_EQUAL(true, is_cuda_to_cpp);\n}\nDECLARE_UNITTEST(TestSelectSystemCudaToCpp);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__ void get_temporary_buffer_kernel(size_t n, Iterator result)\n{\n *result = thrust::get_temporary_buffer(thrust::seq, n);\n}\n\n\ntemplate\n__global__ void return_temporary_buffer_kernel(Pointer ptr, std::ptrdiff_t n)\n{\n thrust::return_temporary_buffer(thrust::seq, ptr, n);\n}\n\n\nvoid TestGetTemporaryBufferDeviceSeq()\n{\n const std::ptrdiff_t n = 9001;\n\n typedef thrust::pointer pointer;\n typedef thrust::pair ptr_and_sz_type;\n thrust::device_vector d_result(1);\n \n get_temporary_buffer_kernel<<<1,1>>>(n, d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ptr_and_sz_type ptr_and_sz = d_result[0];\n\n if(ptr_and_sz.second > 0)\n {\n ASSERT_EQUAL(ptr_and_sz.second, n);\n\n const int ref_val = 13;\n thrust::device_vector ref(n, ref_val);\n\n thrust::fill_n(thrust::device, ptr_and_sz.first, n, ref_val);\n\n ASSERT_EQUAL(true, thrust::all_of(thrust::device, ptr_and_sz.first, ptr_and_sz.first + n, thrust::placeholders::_1 == ref_val));\n\n return_temporary_buffer_kernel<<<1,1>>>(ptr_and_sz.first, ptr_and_sz.second);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n }\n}\nDECLARE_UNITTEST(TestGetTemporaryBufferDeviceSeq);\n\n\ntemplate\n__global__ void malloc_kernel(size_t n, Iterator result)\n{\n *result = static_cast(thrust::malloc(thrust::seq, sizeof(int) * n).get());\n}\n\n\ntemplate\n__global__ void free_kernel(Pointer ptr)\n{\n thrust::free(thrust::seq, ptr);\n}\n\n\nvoid TestMallocDeviceSeq()\n{\n const std::ptrdiff_t n = 9001;\n\n typedef thrust::pointer pointer;\n thrust::device_vector d_result(1);\n \n malloc_kernel<<<1,1>>>(n, d_result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n pointer ptr = d_result[0];\n\n if(ptr.get() != 0)\n {\n const int ref_val = 13;\n thrust::device_vector ref(n, ref_val);\n\n thrust::fill_n(thrust::device, ptr, n, ref_val);\n\n ASSERT_EQUAL(true, thrust::all_of(thrust::device, ptr, ptr + n, thrust::placeholders::_1 == ref_val));\n\n free_kernel<<<1,1>>>(ptr);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n }\n}\nDECLARE_UNITTEST(TestMallocDeviceSeq);\n#endif\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\n\ntemplate\nbool are_same_type(const T1 &, const T2 &)\n{\n return false;\n}\n\n\ntemplate\nbool are_same_type(const T &, const T &)\n{\n return true;\n}\n\n\nvoid TestSelectSystemCudaToCpp()\n{\n using thrust::system::detail::generic::select_system;\n\n thrust::cuda::tag cuda_tag;\n thrust::cpp::tag cpp_tag;\n thrust::cuda_cub::cross_system cuda_to_cpp(cuda_tag, cpp_tag);\n\n // select_system(cuda::tag, thrust::host_system_tag) should return cuda_to_cpp\n bool is_cuda_to_cpp = are_same_type(cuda_to_cpp, select_system(cuda_tag, cpp_tag));\n ASSERT_EQUAL(true, is_cuda_to_cpp);\n}\nDECLARE_UNITTEST(TestSelectSystemCudaToCpp);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__ void get_temporary_buffer_kernel(size_t n, Iterator result)\n{\n *result = thrust::get_temporary_buffer(thrust::seq, n);\n}\n\n\ntemplate\n__global__ void return_temporary_buffer_kernel(Pointer ptr, std::ptrdiff_t n)\n{\n thrust::return_temporary_buffer(thrust::seq, ptr, n);\n}\n\n\nvoid TestGetTemporaryBufferDeviceSeq()\n{\n const std::ptrdiff_t n = 9001;\n\n typedef thrust::pointer pointer;\n typedef thrust::pair ptr_and_sz_type;\n thrust::device_vector d_result(1);\n \n get_temporary_buffer_kernel<<<1,1>>>(n, d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ptr_and_sz_type ptr_and_sz = d_result[0];\n\n if(ptr_and_sz.second > 0)\n {\n ASSERT_EQUAL(ptr_and_sz.second, n);\n\n const int ref_val = 13;\n thrust::device_vector ref(n, ref_val);\n\n thrust::fill_n(thrust::device, ptr_and_sz.first, n, ref_val);\n\n ASSERT_EQUAL(true, thrust::all_of(thrust::device, ptr_and_sz.first, ptr_and_sz.first + n, thrust::placeholders::_1 == ref_val));\n\n return_temporary_buffer_kernel<<<1,1>>>(ptr_and_sz.first, ptr_and_sz.second);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n }\n}\nDECLARE_UNITTEST(TestGetTemporaryBufferDeviceSeq);\n\n\ntemplate\n__global__ void malloc_kernel(size_t n, Iterator result)\n{\n *result = static_cast(thrust::malloc(thrust::seq, sizeof(int) * n).get());\n}\n\n\ntemplate\n__global__ void free_kernel(Pointer ptr)\n{\n thrust::free(thrust::seq, ptr);\n}\n\n\nvoid TestMallocDeviceSeq()\n{\n const std::ptrdiff_t n = 9001;\n\n typedef thrust::pointer pointer;\n thrust::device_vector d_result(1);\n \n malloc_kernel<<<1,1>>>(n, d_result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n pointer ptr = d_result[0];\n\n if(ptr.get() != 0)\n {\n const int ref_val = 13;\n thrust::device_vector ref(n, ref_val);\n\n thrust::fill_n(thrust::device, ptr, n, ref_val);\n\n ASSERT_EQUAL(true, thrust::all_of(thrust::device, ptr, ptr + n, thrust::placeholders::_1 == ref_val));\n\n free_kernel<<<1,1>>>(ptr);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n }\n}\nDECLARE_UNITTEST(TestMallocDeviceSeq);\n#endif\n\n", "source": "/kaggle/working/thrust/testing/cuda/memory.cu"} {"input": "#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid sequence_kernel(ExecutionPolicy exec, Iterator first, Iterator last)\n{\n thrust::sequence(exec, first, last);\n}\n\n\ntemplate\n__global__\nvoid sequence_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T init)\n{\n thrust::sequence(exec, first, last, init);\n}\n\n\ntemplate\n__global__\nvoid sequence_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T init, T step)\n{\n thrust::sequence(exec, first, last, init, step);\n}\n\n\ntemplate\nvoid TestSequenceDevice(ExecutionPolicy exec)\n{\n thrust::device_vector v(5);\n \n sequence_kernel<<<1,1>>>(exec, v.begin(), v.end());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], 1);\n ASSERT_EQUAL(v[2], 2);\n ASSERT_EQUAL(v[3], 3);\n ASSERT_EQUAL(v[4], 4);\n \n sequence_kernel<<<1,1>>>(exec, v.begin(), v.end(), 10);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(v[0], 10);\n ASSERT_EQUAL(v[1], 11);\n ASSERT_EQUAL(v[2], 12);\n ASSERT_EQUAL(v[3], 13);\n ASSERT_EQUAL(v[4], 14);\n \n sequence_kernel<<<1,1>>>(exec, v.begin(), v.end(), 10, 2);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(v[0], 10);\n ASSERT_EQUAL(v[1], 12);\n ASSERT_EQUAL(v[2], 14);\n ASSERT_EQUAL(v[3], 16);\n ASSERT_EQUAL(v[4], 18);\n}\n\nvoid TestSequenceDeviceSeq()\n{\n TestSequenceDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSequenceDeviceSeq);\n\nvoid TestSequenceDeviceDevice()\n{\n TestSequenceDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSequenceDeviceDevice);\n#endif\n\nvoid TestSequenceCudaStreams()\n{\n typedef thrust::device_vector Vector;\n \n Vector v(5);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n thrust::sequence(thrust::cuda::par.on(s), v.begin(), v.end());\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], 1);\n ASSERT_EQUAL(v[2], 2);\n ASSERT_EQUAL(v[3], 3);\n ASSERT_EQUAL(v[4], 4);\n\n thrust::sequence(thrust::cuda::par.on(s), v.begin(), v.end(), 10);\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(v[0], 10);\n ASSERT_EQUAL(v[1], 11);\n ASSERT_EQUAL(v[2], 12);\n ASSERT_EQUAL(v[3], 13);\n ASSERT_EQUAL(v[4], 14);\n \n thrust::sequence(thrust::cuda::par.on(s), v.begin(), v.end(), 10, 2);\n cudaStreamSynchronize(s);\n\n ASSERT_EQUAL(v[0], 10);\n ASSERT_EQUAL(v[1], 12);\n ASSERT_EQUAL(v[2], 14);\n ASSERT_EQUAL(v[3], 16);\n ASSERT_EQUAL(v[4], 18);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSequenceCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid sequence_kernel(ExecutionPolicy exec, Iterator first, Iterator last)\n{\n thrust::sequence(exec, first, last);\n}\n\n\ntemplate\n__global__\nvoid sequence_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T init)\n{\n thrust::sequence(exec, first, last, init);\n}\n\n\ntemplate\n__global__\nvoid sequence_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T init, T step)\n{\n thrust::sequence(exec, first, last, init, step);\n}\n\n\ntemplate\nvoid TestSequenceDevice(ExecutionPolicy exec)\n{\n thrust::device_vector v(5);\n \n sequence_kernel<<<1,1>>>(exec, v.begin(), v.end());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], 1);\n ASSERT_EQUAL(v[2], 2);\n ASSERT_EQUAL(v[3], 3);\n ASSERT_EQUAL(v[4], 4);\n \n sequence_kernel<<<1,1>>>(exec, v.begin(), v.end(), 10);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(v[0], 10);\n ASSERT_EQUAL(v[1], 11);\n ASSERT_EQUAL(v[2], 12);\n ASSERT_EQUAL(v[3], 13);\n ASSERT_EQUAL(v[4], 14);\n \n sequence_kernel<<<1,1>>>(exec, v.begin(), v.end(), 10, 2);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(v[0], 10);\n ASSERT_EQUAL(v[1], 12);\n ASSERT_EQUAL(v[2], 14);\n ASSERT_EQUAL(v[3], 16);\n ASSERT_EQUAL(v[4], 18);\n}\n\nvoid TestSequenceDeviceSeq()\n{\n TestSequenceDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestSequenceDeviceSeq);\n\nvoid TestSequenceDeviceDevice()\n{\n TestSequenceDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestSequenceDeviceDevice);\n#endif\n\nvoid TestSequenceCudaStreams()\n{\n typedef thrust::device_vector Vector;\n \n Vector v(5);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n thrust::sequence(thrust::cuda::par.on(s), v.begin(), v.end());\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], 1);\n ASSERT_EQUAL(v[2], 2);\n ASSERT_EQUAL(v[3], 3);\n ASSERT_EQUAL(v[4], 4);\n\n thrust::sequence(thrust::cuda::par.on(s), v.begin(), v.end(), 10);\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(v[0], 10);\n ASSERT_EQUAL(v[1], 11);\n ASSERT_EQUAL(v[2], 12);\n ASSERT_EQUAL(v[3], 13);\n ASSERT_EQUAL(v[4], 14);\n \n thrust::sequence(thrust::cuda::par.on(s), v.begin(), v.end(), 10, 2);\n hipStreamSynchronize(s);\n\n ASSERT_EQUAL(v[0], 10);\n ASSERT_EQUAL(v[1], 12);\n ASSERT_EQUAL(v[2], 14);\n ASSERT_EQUAL(v[3], 16);\n ASSERT_EQUAL(v[4], 18);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestSequenceCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/sequence.cu"} {"input": "#include \n#include \n#include \n#include \n\n\ntemplate\nstruct is_even\n{\n __host__ __device__\n bool operator()(T x) const { return ((int) x % 2) == 0; }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid partition_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Predicate pred, Iterator2 result)\n{\n *result = thrust::partition(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestPartitionDevice(ExecutionPolicy exec)\n{\n typedef int T;\n typedef typename thrust::device_vector::iterator iterator;\n \n thrust::device_vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 1; \n data[4] = 2; \n\n thrust::device_vector result(1);\n \n partition_kernel<<<1,1>>>(exec, data.begin(), data.end(), is_even(), result.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n thrust::device_vector ref(5);\n ref[0] = 2;\n ref[1] = 2;\n ref[2] = 1;\n ref[3] = 1;\n ref[4] = 1;\n \n ASSERT_EQUAL(2, (iterator)result[0] - data.begin());\n ASSERT_EQUAL(ref, data);\n}\n\n\nvoid TestPartitionDeviceSeq()\n{\n TestPartitionDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestPartitionDeviceSeq);\n\n\nvoid TestPartitionDeviceDevice()\n{\n TestPartitionDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestPartitionDeviceDevice);\n\n\nvoid TestPartitionDeviceNoSync()\n{\n TestPartitionDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestPartitionDeviceNoSync);\n\n\ntemplate\n__global__\nvoid partition_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Predicate pred, Iterator3 result)\n{\n *result = thrust::partition(exec, first, last, stencil_first, pred);\n}\n\n\ntemplate\nvoid TestPartitionStencilDevice(ExecutionPolicy exec)\n{\n typedef int T;\n typedef typename thrust::device_vector::iterator iterator;\n \n thrust::device_vector data(5);\n data[0] = 0;\n data[1] = 1;\n data[2] = 0;\n data[3] = 0;\n data[4] = 1;\n \n thrust::device_vector stencil(5);\n stencil[0] = 1; \n stencil[1] = 2; \n stencil[2] = 1;\n stencil[3] = 1; \n stencil[4] = 2; \n\n thrust::device_vector result(1);\n \n partition_kernel<<<1,1>>>(exec, data.begin(), data.end(), stencil.begin(), is_even(), result.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n thrust::device_vector ref(5);\n ref[0] = 1;\n ref[1] = 1;\n ref[2] = 0;\n ref[3] = 0;\n ref[4] = 0;\n \n ASSERT_EQUAL(2, (iterator)result[0] - data.begin());\n ASSERT_EQUAL(ref, data);\n}\n\n\nvoid TestPartitionStencilDeviceSeq()\n{\n TestPartitionStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestPartitionStencilDeviceSeq);\n\n\nvoid TestPartitionStencilDeviceDevice()\n{\n TestPartitionStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestPartitionStencilDeviceDevice);\n\n\nvoid TestPartitionStencilDeviceNoSync()\n{\n TestPartitionStencilDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestPartitionStencilDeviceNoSync);\n\n\ntemplate\n__global__\nvoid partition_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 true_result, Iterator3 false_result, Predicate pred, Iterator4 result)\n{\n *result = thrust::partition_copy(exec, first, last, true_result, false_result, pred);\n}\n\n\ntemplate\nvoid TestPartitionCopyDevice(ExecutionPolicy exec)\n{\n typedef int T;\n typedef thrust::device_vector::iterator iterator;\n \n thrust::device_vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 1; \n data[4] = 2; \n \n thrust::device_vector true_results(2);\n thrust::device_vector false_results(3);\n\n typedef thrust::pair pair_type;\n thrust::device_vector iterators(1);\n \n partition_copy_kernel<<<1,1>>>(exec, data.begin(), data.end(), true_results.begin(), false_results.begin(), is_even(), iterators.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n thrust::device_vector true_ref(2);\n true_ref[0] = 2;\n true_ref[1] = 2;\n \n thrust::device_vector false_ref(3);\n false_ref[0] = 1;\n false_ref[1] = 1;\n false_ref[2] = 1;\n\n pair_type ends = iterators[0];\n \n ASSERT_EQUAL(2, ends.first - true_results.begin());\n ASSERT_EQUAL(3, ends.second - false_results.begin());\n ASSERT_EQUAL(true_ref, true_results);\n ASSERT_EQUAL(false_ref, false_results);\n}\n\n\nvoid TestPartitionCopyDeviceSeq()\n{\n TestPartitionCopyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestPartitionCopyDeviceSeq);\n\n\nvoid TestPartitionCopyDeviceDevice()\n{\n TestPartitionCopyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestPartitionCopyDeviceDevice);\n\n\nvoid TestPartitionCopyDeviceNoSync()\n{\n TestPartitionCopyDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestPartitionCopyDeviceNoSync);\n\n\ntemplate\n__global__\nvoid partition_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Iterator3 true_result, Iterator4 false_result, Predicate pred, Iterator5 result)\n{\n *result = thrust::partition_copy(exec, first, last, stencil_first, true_result, false_result, pred);\n}\n\n\ntemplate\nvoid TestPartitionCopyStencilDevice(ExecutionPolicy exec)\n{\n typedef int T;\n \n thrust::device_vector data(5);\n data[0] = 0; \n data[1] = 1; \n data[2] = 0;\n data[3] = 0; \n data[4] = 1; \n \n thrust::device_vector stencil(5);\n stencil[0] = 1; \n stencil[1] = 2; \n stencil[2] = 1;\n stencil[3] = 1; \n stencil[4] = 2; \n \n thrust::device_vector true_results(2);\n thrust::device_vector false_results(3);\n\n typedef typename thrust::device_vector::iterator iterator;\n typedef thrust::pair pair_type;\n thrust::device_vector iterators(1);\n\n partition_copy_kernel<<<1,1>>>(exec, data.begin(), data.end(), stencil.begin(), true_results.begin(), false_results.begin(), is_even(), iterators.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n pair_type ends = iterators[0];\n \n thrust::device_vector true_ref(2);\n true_ref[0] = 1;\n true_ref[1] = 1;\n \n thrust::device_vector false_ref(3);\n false_ref[0] = 0;\n false_ref[1] = 0;\n false_ref[2] = 0;\n \n ASSERT_EQUAL(2, ends.first - true_results.begin());\n ASSERT_EQUAL(3, ends.second - false_results.begin());\n ASSERT_EQUAL(true_ref, true_results);\n ASSERT_EQUAL(false_ref, false_results);\n}\n\n\nvoid TestPartitionCopyStencilDeviceSeq()\n{\n TestPartitionCopyStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestPartitionCopyStencilDeviceSeq);\n\n\nvoid TestPartitionCopyStencilDeviceDevice()\n{\n TestPartitionCopyStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestPartitionCopyStencilDeviceDevice);\n\n\nvoid TestPartitionCopyStencilDeviceNoSync()\n{\n TestPartitionCopyStencilDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestPartitionCopyStencilDeviceNoSync);\n\n\ntemplate\n__global__\nvoid stable_partition_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Predicate pred, Iterator2 result)\n{\n *result = thrust::stable_partition(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestStablePartitionDevice(ExecutionPolicy exec)\n{\n typedef int T;\n typedef typename thrust::device_vector::iterator iterator;\n \n thrust::device_vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 1; \n data[4] = 2; \n\n thrust::device_vector result(1);\n\n stable_partition_kernel<<<1,1>>>(exec, data.begin(), data.end(), is_even(), result.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n thrust::device_vector ref(5);\n ref[0] = 2;\n ref[1] = 2;\n ref[2] = 1;\n ref[3] = 1;\n ref[4] = 1;\n \n ASSERT_EQUAL(2, (iterator)result[0] - data.begin());\n ASSERT_EQUAL(ref, data);\n}\n\n\nvoid TestStablePartitionDeviceSeq()\n{\n TestStablePartitionDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestStablePartitionDeviceSeq);\n\n\nvoid TestStablePartitionDeviceDevice()\n{\n TestStablePartitionDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestStablePartitionDeviceDevice);\n\n\nvoid TestStablePartitionDeviceNoSync()\n{\n TestStablePartitionDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestStablePartitionDeviceNoSync);\n\n\ntemplate\n__global__\nvoid stable_partition_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Predicate pred, Iterator3 result)\n{\n *result = thrust::stable_partition(exec, first, last, stencil_first, pred);\n}\n\n\ntemplate\nvoid TestStablePartitionStencilDevice(ExecutionPolicy exec)\n{\n typedef int T;\n typedef typename thrust::device_vector::iterator iterator;\n \n thrust::device_vector data(5);\n data[0] = 0;\n data[1] = 1;\n data[2] = 0;\n data[3] = 0;\n data[4] = 1;\n \n thrust::device_vector stencil(5);\n stencil[0] = 1; \n stencil[1] = 2; \n stencil[2] = 1;\n stencil[3] = 1; \n stencil[4] = 2; \n\n thrust::device_vector result(1);\n\n stable_partition_kernel<<<1,1>>>(exec, data.begin(), data.end(), stencil.begin(), is_even(), result.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n thrust::device_vector ref(5);\n ref[0] = 1;\n ref[1] = 1;\n ref[2] = 0;\n ref[3] = 0;\n ref[4] = 0;\n \n ASSERT_EQUAL(2, (iterator)result[0] - data.begin());\n ASSERT_EQUAL(ref, data);\n}\n\n\nvoid TestStablePartitionStencilDeviceSeq()\n{\n TestStablePartitionStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestStablePartitionStencilDeviceSeq);\n\n\nvoid TestStablePartitionStencilDeviceDevice()\n{\n TestStablePartitionStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestStablePartitionStencilDeviceDevice);\n\n\nvoid TestStablePartitionStencilDeviceNoSync()\n{\n TestStablePartitionStencilDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestStablePartitionStencilDeviceNoSync);\n\n\ntemplate\n__global__\nvoid stable_partition_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 true_result, Iterator3 false_result, Predicate pred, Iterator4 result)\n{\n *result = thrust::stable_partition_copy(exec, first, last, true_result, false_result, pred);\n}\n\n\ntemplate\nvoid TestStablePartitionCopyDevice(ExecutionPolicy exec)\n{\n typedef int T;\n typedef thrust::device_vector::iterator iterator;\n \n thrust::device_vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 1; \n data[4] = 2; \n \n thrust::device_vector true_results(2);\n thrust::device_vector false_results(3);\n\n typedef thrust::pair pair_type;\n thrust::device_vector iterators(1);\n \n stable_partition_copy_kernel<<<1,1>>>(exec, data.begin(), data.end(), true_results.begin(), false_results.begin(), is_even(), iterators.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n \n thrust::device_vector true_ref(2);\n true_ref[0] = 2;\n true_ref[1] = 2;\n \n thrust::device_vector false_ref(3);\n false_ref[0] = 1;\n false_ref[1] = 1;\n false_ref[2] = 1;\n\n pair_type ends = iterators[0];\n \n ASSERT_EQUAL(2, ends.first - true_results.begin());\n ASSERT_EQUAL(3, ends.second - false_results.begin());\n ASSERT_EQUAL(true_ref, true_results);\n ASSERT_EQUAL(false_ref, false_results);\n}\n\n\nvoid TestStablePartitionCopyDeviceSeq()\n{\n TestStablePartitionCopyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestStablePartitionCopyDeviceSeq);\n\n\nvoid TestStablePartitionCopyDeviceDevice()\n{\n TestStablePartitionCopyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestStablePartitionCopyDeviceDevice);\n\n\nvoid TestStablePartitionCopyDeviceNoSync()\n{\n TestStablePartitionCopyDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestStablePartitionCopyDeviceNoSync);\n\n\ntemplate\n__global__\nvoid stable_partition_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Iterator3 true_result, Iterator4 false_result, Predicate pred, Iterator5 result)\n{\n *result = thrust::stable_partition_copy(exec, first, last, stencil_first, true_result, false_result, pred);\n}\n\n\ntemplate\nvoid TestStablePartitionCopyStencilDevice(ExecutionPolicy exec)\n{\n typedef int T;\n \n thrust::device_vector data(5);\n data[0] = 0; \n data[1] = 1; \n data[2] = 0;\n data[3] = 0; \n data[4] = 1; \n \n thrust::device_vector stencil(5);\n stencil[0] = 1; \n stencil[1] = 2; \n stencil[2] = 1;\n stencil[3] = 1; \n stencil[4] = 2; \n \n thrust::device_vector true_results(2);\n thrust::device_vector false_results(3);\n\n typedef typename thrust::device_vector::iterator iterator;\n typedef thrust::pair pair_type;\n thrust::device_vector iterators(1);\n\n stable_partition_copy_kernel<<<1,1>>>(exec, data.begin(), data.end(), stencil.begin(), true_results.begin(), false_results.begin(), is_even(), iterators.begin());\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n\n pair_type ends = iterators[0];\n \n thrust::device_vector true_ref(2);\n true_ref[0] = 1;\n true_ref[1] = 1;\n \n thrust::device_vector false_ref(3);\n false_ref[0] = 0;\n false_ref[1] = 0;\n false_ref[2] = 0;\n \n ASSERT_EQUAL(2, ends.first - true_results.begin());\n ASSERT_EQUAL(3, ends.second - false_results.begin());\n ASSERT_EQUAL(true_ref, true_results);\n ASSERT_EQUAL(false_ref, false_results);\n}\n\n\nvoid TestStablePartitionCopyStencilDeviceSeq()\n{\n TestStablePartitionCopyStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestStablePartitionCopyStencilDeviceSeq);\n\n\nvoid TestStablePartitionCopyStencilDeviceDevice()\n{\n TestStablePartitionCopyStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestStablePartitionCopyStencilDeviceDevice);\n\n\nvoid TestStablePartitionCopyStencilDeviceNoSync()\n{\n TestStablePartitionCopyStencilDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestStablePartitionCopyStencilDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestPartitionCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n typedef Vector::iterator Iterator;\n \n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 1; \n data[4] = 2; \n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n \n Iterator iter = thrust::partition(streampolicy, data.begin(), data.end(), is_even());\n \n Vector ref(5);\n ref[0] = 2;\n ref[1] = 2;\n ref[2] = 1;\n ref[3] = 1;\n ref[4] = 1;\n \n ASSERT_EQUAL(iter - data.begin(), 2);\n ASSERT_EQUAL(data, ref);\n\n cudaStreamDestroy(s);\n}\n\nvoid TestPartitionCudaStreamsSync()\n{\n TestPartitionCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestPartitionCudaStreamsSync);\n\n\nvoid TestPartitionCudaStreamsNoSync()\n{\n TestPartitionCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestPartitionCudaStreamsNoSync);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n\ntemplate\nstruct is_even\n{\n __host__ __device__\n bool operator()(T x) const { return ((int) x % 2) == 0; }\n};\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid partition_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Predicate pred, Iterator2 result)\n{\n *result = thrust::partition(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestPartitionDevice(ExecutionPolicy exec)\n{\n typedef int T;\n typedef typename thrust::device_vector::iterator iterator;\n \n thrust::device_vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 1; \n data[4] = 2; \n\n thrust::device_vector result(1);\n \n partition_kernel<<<1,1>>>(exec, data.begin(), data.end(), is_even(), result.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n thrust::device_vector ref(5);\n ref[0] = 2;\n ref[1] = 2;\n ref[2] = 1;\n ref[3] = 1;\n ref[4] = 1;\n \n ASSERT_EQUAL(2, (iterator)result[0] - data.begin());\n ASSERT_EQUAL(ref, data);\n}\n\n\nvoid TestPartitionDeviceSeq()\n{\n TestPartitionDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestPartitionDeviceSeq);\n\n\nvoid TestPartitionDeviceDevice()\n{\n TestPartitionDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestPartitionDeviceDevice);\n\n\nvoid TestPartitionDeviceNoSync()\n{\n TestPartitionDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestPartitionDeviceNoSync);\n\n\ntemplate\n__global__\nvoid partition_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Predicate pred, Iterator3 result)\n{\n *result = thrust::partition(exec, first, last, stencil_first, pred);\n}\n\n\ntemplate\nvoid TestPartitionStencilDevice(ExecutionPolicy exec)\n{\n typedef int T;\n typedef typename thrust::device_vector::iterator iterator;\n \n thrust::device_vector data(5);\n data[0] = 0;\n data[1] = 1;\n data[2] = 0;\n data[3] = 0;\n data[4] = 1;\n \n thrust::device_vector stencil(5);\n stencil[0] = 1; \n stencil[1] = 2; \n stencil[2] = 1;\n stencil[3] = 1; \n stencil[4] = 2; \n\n thrust::device_vector result(1);\n \n partition_kernel<<<1,1>>>(exec, data.begin(), data.end(), stencil.begin(), is_even(), result.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n thrust::device_vector ref(5);\n ref[0] = 1;\n ref[1] = 1;\n ref[2] = 0;\n ref[3] = 0;\n ref[4] = 0;\n \n ASSERT_EQUAL(2, (iterator)result[0] - data.begin());\n ASSERT_EQUAL(ref, data);\n}\n\n\nvoid TestPartitionStencilDeviceSeq()\n{\n TestPartitionStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestPartitionStencilDeviceSeq);\n\n\nvoid TestPartitionStencilDeviceDevice()\n{\n TestPartitionStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestPartitionStencilDeviceDevice);\n\n\nvoid TestPartitionStencilDeviceNoSync()\n{\n TestPartitionStencilDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestPartitionStencilDeviceNoSync);\n\n\ntemplate\n__global__\nvoid partition_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 true_result, Iterator3 false_result, Predicate pred, Iterator4 result)\n{\n *result = thrust::partition_copy(exec, first, last, true_result, false_result, pred);\n}\n\n\ntemplate\nvoid TestPartitionCopyDevice(ExecutionPolicy exec)\n{\n typedef int T;\n typedef thrust::device_vector::iterator iterator;\n \n thrust::device_vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 1; \n data[4] = 2; \n \n thrust::device_vector true_results(2);\n thrust::device_vector false_results(3);\n\n typedef thrust::pair pair_type;\n thrust::device_vector iterators(1);\n \n partition_copy_kernel<<<1,1>>>(exec, data.begin(), data.end(), true_results.begin(), false_results.begin(), is_even(), iterators.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n thrust::device_vector true_ref(2);\n true_ref[0] = 2;\n true_ref[1] = 2;\n \n thrust::device_vector false_ref(3);\n false_ref[0] = 1;\n false_ref[1] = 1;\n false_ref[2] = 1;\n\n pair_type ends = iterators[0];\n \n ASSERT_EQUAL(2, ends.first - true_results.begin());\n ASSERT_EQUAL(3, ends.second - false_results.begin());\n ASSERT_EQUAL(true_ref, true_results);\n ASSERT_EQUAL(false_ref, false_results);\n}\n\n\nvoid TestPartitionCopyDeviceSeq()\n{\n TestPartitionCopyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestPartitionCopyDeviceSeq);\n\n\nvoid TestPartitionCopyDeviceDevice()\n{\n TestPartitionCopyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestPartitionCopyDeviceDevice);\n\n\nvoid TestPartitionCopyDeviceNoSync()\n{\n TestPartitionCopyDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestPartitionCopyDeviceNoSync);\n\n\ntemplate\n__global__\nvoid partition_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Iterator3 true_result, Iterator4 false_result, Predicate pred, Iterator5 result)\n{\n *result = thrust::partition_copy(exec, first, last, stencil_first, true_result, false_result, pred);\n}\n\n\ntemplate\nvoid TestPartitionCopyStencilDevice(ExecutionPolicy exec)\n{\n typedef int T;\n \n thrust::device_vector data(5);\n data[0] = 0; \n data[1] = 1; \n data[2] = 0;\n data[3] = 0; \n data[4] = 1; \n \n thrust::device_vector stencil(5);\n stencil[0] = 1; \n stencil[1] = 2; \n stencil[2] = 1;\n stencil[3] = 1; \n stencil[4] = 2; \n \n thrust::device_vector true_results(2);\n thrust::device_vector false_results(3);\n\n typedef typename thrust::device_vector::iterator iterator;\n typedef thrust::pair pair_type;\n thrust::device_vector iterators(1);\n\n partition_copy_kernel<<<1,1>>>(exec, data.begin(), data.end(), stencil.begin(), true_results.begin(), false_results.begin(), is_even(), iterators.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n pair_type ends = iterators[0];\n \n thrust::device_vector true_ref(2);\n true_ref[0] = 1;\n true_ref[1] = 1;\n \n thrust::device_vector false_ref(3);\n false_ref[0] = 0;\n false_ref[1] = 0;\n false_ref[2] = 0;\n \n ASSERT_EQUAL(2, ends.first - true_results.begin());\n ASSERT_EQUAL(3, ends.second - false_results.begin());\n ASSERT_EQUAL(true_ref, true_results);\n ASSERT_EQUAL(false_ref, false_results);\n}\n\n\nvoid TestPartitionCopyStencilDeviceSeq()\n{\n TestPartitionCopyStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestPartitionCopyStencilDeviceSeq);\n\n\nvoid TestPartitionCopyStencilDeviceDevice()\n{\n TestPartitionCopyStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestPartitionCopyStencilDeviceDevice);\n\n\nvoid TestPartitionCopyStencilDeviceNoSync()\n{\n TestPartitionCopyStencilDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestPartitionCopyStencilDeviceNoSync);\n\n\ntemplate\n__global__\nvoid stable_partition_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Predicate pred, Iterator2 result)\n{\n *result = thrust::stable_partition(exec, first, last, pred);\n}\n\n\ntemplate\nvoid TestStablePartitionDevice(ExecutionPolicy exec)\n{\n typedef int T;\n typedef typename thrust::device_vector::iterator iterator;\n \n thrust::device_vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 1; \n data[4] = 2; \n\n thrust::device_vector result(1);\n\n stable_partition_kernel<<<1,1>>>(exec, data.begin(), data.end(), is_even(), result.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n thrust::device_vector ref(5);\n ref[0] = 2;\n ref[1] = 2;\n ref[2] = 1;\n ref[3] = 1;\n ref[4] = 1;\n \n ASSERT_EQUAL(2, (iterator)result[0] - data.begin());\n ASSERT_EQUAL(ref, data);\n}\n\n\nvoid TestStablePartitionDeviceSeq()\n{\n TestStablePartitionDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestStablePartitionDeviceSeq);\n\n\nvoid TestStablePartitionDeviceDevice()\n{\n TestStablePartitionDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestStablePartitionDeviceDevice);\n\n\nvoid TestStablePartitionDeviceNoSync()\n{\n TestStablePartitionDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestStablePartitionDeviceNoSync);\n\n\ntemplate\n__global__\nvoid stable_partition_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Predicate pred, Iterator3 result)\n{\n *result = thrust::stable_partition(exec, first, last, stencil_first, pred);\n}\n\n\ntemplate\nvoid TestStablePartitionStencilDevice(ExecutionPolicy exec)\n{\n typedef int T;\n typedef typename thrust::device_vector::iterator iterator;\n \n thrust::device_vector data(5);\n data[0] = 0;\n data[1] = 1;\n data[2] = 0;\n data[3] = 0;\n data[4] = 1;\n \n thrust::device_vector stencil(5);\n stencil[0] = 1; \n stencil[1] = 2; \n stencil[2] = 1;\n stencil[3] = 1; \n stencil[4] = 2; \n\n thrust::device_vector result(1);\n\n stable_partition_kernel<<<1,1>>>(exec, data.begin(), data.end(), stencil.begin(), is_even(), result.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n thrust::device_vector ref(5);\n ref[0] = 1;\n ref[1] = 1;\n ref[2] = 0;\n ref[3] = 0;\n ref[4] = 0;\n \n ASSERT_EQUAL(2, (iterator)result[0] - data.begin());\n ASSERT_EQUAL(ref, data);\n}\n\n\nvoid TestStablePartitionStencilDeviceSeq()\n{\n TestStablePartitionStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestStablePartitionStencilDeviceSeq);\n\n\nvoid TestStablePartitionStencilDeviceDevice()\n{\n TestStablePartitionStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestStablePartitionStencilDeviceDevice);\n\n\nvoid TestStablePartitionStencilDeviceNoSync()\n{\n TestStablePartitionStencilDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestStablePartitionStencilDeviceNoSync);\n\n\ntemplate\n__global__\nvoid stable_partition_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 true_result, Iterator3 false_result, Predicate pred, Iterator4 result)\n{\n *result = thrust::stable_partition_copy(exec, first, last, true_result, false_result, pred);\n}\n\n\ntemplate\nvoid TestStablePartitionCopyDevice(ExecutionPolicy exec)\n{\n typedef int T;\n typedef thrust::device_vector::iterator iterator;\n \n thrust::device_vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 1; \n data[4] = 2; \n \n thrust::device_vector true_results(2);\n thrust::device_vector false_results(3);\n\n typedef thrust::pair pair_type;\n thrust::device_vector iterators(1);\n \n stable_partition_copy_kernel<<<1,1>>>(exec, data.begin(), data.end(), true_results.begin(), false_results.begin(), is_even(), iterators.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n \n thrust::device_vector true_ref(2);\n true_ref[0] = 2;\n true_ref[1] = 2;\n \n thrust::device_vector false_ref(3);\n false_ref[0] = 1;\n false_ref[1] = 1;\n false_ref[2] = 1;\n\n pair_type ends = iterators[0];\n \n ASSERT_EQUAL(2, ends.first - true_results.begin());\n ASSERT_EQUAL(3, ends.second - false_results.begin());\n ASSERT_EQUAL(true_ref, true_results);\n ASSERT_EQUAL(false_ref, false_results);\n}\n\n\nvoid TestStablePartitionCopyDeviceSeq()\n{\n TestStablePartitionCopyDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestStablePartitionCopyDeviceSeq);\n\n\nvoid TestStablePartitionCopyDeviceDevice()\n{\n TestStablePartitionCopyDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestStablePartitionCopyDeviceDevice);\n\n\nvoid TestStablePartitionCopyDeviceNoSync()\n{\n TestStablePartitionCopyDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestStablePartitionCopyDeviceNoSync);\n\n\ntemplate\n__global__\nvoid stable_partition_copy_kernel(ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 stencil_first, Iterator3 true_result, Iterator4 false_result, Predicate pred, Iterator5 result)\n{\n *result = thrust::stable_partition_copy(exec, first, last, stencil_first, true_result, false_result, pred);\n}\n\n\ntemplate\nvoid TestStablePartitionCopyStencilDevice(ExecutionPolicy exec)\n{\n typedef int T;\n \n thrust::device_vector data(5);\n data[0] = 0; \n data[1] = 1; \n data[2] = 0;\n data[3] = 0; \n data[4] = 1; \n \n thrust::device_vector stencil(5);\n stencil[0] = 1; \n stencil[1] = 2; \n stencil[2] = 1;\n stencil[3] = 1; \n stencil[4] = 2; \n \n thrust::device_vector true_results(2);\n thrust::device_vector false_results(3);\n\n typedef typename thrust::device_vector::iterator iterator;\n typedef thrust::pair pair_type;\n thrust::device_vector iterators(1);\n\n stable_partition_copy_kernel<<<1,1>>>(exec, data.begin(), data.end(), stencil.begin(), true_results.begin(), false_results.begin(), is_even(), iterators.begin());\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n\n pair_type ends = iterators[0];\n \n thrust::device_vector true_ref(2);\n true_ref[0] = 1;\n true_ref[1] = 1;\n \n thrust::device_vector false_ref(3);\n false_ref[0] = 0;\n false_ref[1] = 0;\n false_ref[2] = 0;\n \n ASSERT_EQUAL(2, ends.first - true_results.begin());\n ASSERT_EQUAL(3, ends.second - false_results.begin());\n ASSERT_EQUAL(true_ref, true_results);\n ASSERT_EQUAL(false_ref, false_results);\n}\n\n\nvoid TestStablePartitionCopyStencilDeviceSeq()\n{\n TestStablePartitionCopyStencilDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestStablePartitionCopyStencilDeviceSeq);\n\n\nvoid TestStablePartitionCopyStencilDeviceDevice()\n{\n TestStablePartitionCopyStencilDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestStablePartitionCopyStencilDeviceDevice);\n\n\nvoid TestStablePartitionCopyStencilDeviceNoSync()\n{\n TestStablePartitionCopyStencilDevice(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestStablePartitionCopyStencilDeviceNoSync);\n#endif\n\n\ntemplate\nvoid TestPartitionCudaStreams(ExecutionPolicy policy)\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n typedef Vector::iterator Iterator;\n \n Vector data(5);\n data[0] = 1; \n data[1] = 2; \n data[2] = 1;\n data[3] = 1; \n data[4] = 2; \n\n hipStream_t s;\n hipStreamCreate(&s);\n\n auto streampolicy = policy.on(s);\n \n Iterator iter = thrust::partition(streampolicy, data.begin(), data.end(), is_even());\n \n Vector ref(5);\n ref[0] = 2;\n ref[1] = 2;\n ref[2] = 1;\n ref[3] = 1;\n ref[4] = 1;\n \n ASSERT_EQUAL(iter - data.begin(), 2);\n ASSERT_EQUAL(data, ref);\n\n hipStreamDestroy(s);\n}\n\nvoid TestPartitionCudaStreamsSync()\n{\n TestPartitionCudaStreams(thrust::cuda::par);\n}\nDECLARE_UNITTEST(TestPartitionCudaStreamsSync);\n\n\nvoid TestPartitionCudaStreamsNoSync()\n{\n TestPartitionCudaStreams(thrust::cuda::par_nosync);\n}\nDECLARE_UNITTEST(TestPartitionCudaStreamsNoSync);\n\n", "source": "/kaggle/working/thrust/testing/cuda/partition.cu"} {"input": "#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid all_of_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Function f, Iterator2 result)\n{\n *result = thrust::all_of(exec, first, last, f);\n}\n\n\ntemplate\nvoid TestAllOfDevice(ExecutionPolicy exec)\n{\n typedef int T;\n thrust::device_vector v(3, 1);\n thrust::device_vector result(1);\n \n all_of_kernel<<<1,1>>>(exec, v.begin(), v.end(), thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n \n v[1] = 0;\n \n all_of_kernel<<<1,1>>>(exec, v.begin(), v.end(), thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n \n all_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 0, thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n\n all_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 1, thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n\n all_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 2, thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n\n all_of_kernel<<<1,1>>>(exec, v.begin() + 1, v.begin() + 2, thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n}\n\n\nvoid TestAllOfDeviceSeq()\n{\n TestAllOfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestAllOfDeviceSeq);\n\n\nvoid TestAllOfDeviceDevice()\n{\n TestAllOfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestAllOfDeviceDevice);\n#endif\n\n\nvoid TestAllOfCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector v(3, 1);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n ASSERT_EQUAL(thrust::all_of(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()), true);\n \n v[1] = 0;\n \n ASSERT_EQUAL(thrust::all_of(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()), false);\n \n ASSERT_EQUAL(thrust::all_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 0, thrust::identity()), true);\n ASSERT_EQUAL(thrust::all_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 1, thrust::identity()), true);\n ASSERT_EQUAL(thrust::all_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 2, thrust::identity()), false);\n ASSERT_EQUAL(thrust::all_of(thrust::cuda::par.on(s), v.begin() + 1, v.begin() + 2, thrust::identity()), false);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestAllOfCudaStreams);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid any_of_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Function f, Iterator2 result)\n{\n *result = thrust::any_of(exec, first, last, f);\n}\n\n\ntemplate\nvoid TestAnyOfDevice(ExecutionPolicy exec)\n{\n typedef int T;\n \n thrust::device_vector v(3, 1);\n thrust::device_vector result(1);\n \n any_of_kernel<<<1,1>>>(exec, v.begin(), v.end(), thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n \n v[1] = 0;\n \n any_of_kernel<<<1,1>>>(exec, v.begin(), v.end(), thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n \n any_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 0, thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n\n any_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 1, thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n\n any_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 2, thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n\n any_of_kernel<<<1,1>>>(exec, v.begin() + 1, v.begin() + 2, thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n}\n\n\nvoid TestAnyOfDeviceSeq()\n{\n TestAnyOfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestAnyOfDeviceSeq);\n\n\nvoid TestAnyOfDeviceDevice()\n{\n TestAnyOfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestAnyOfDeviceDevice);\n#endif\n\n\nvoid TestAnyOfCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector v(3, 1);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n ASSERT_EQUAL(thrust::any_of(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()), true);\n\n v[1] = 0;\n \n ASSERT_EQUAL(thrust::any_of(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()), true);\n\n ASSERT_EQUAL(thrust::any_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 0, thrust::identity()), false);\n ASSERT_EQUAL(thrust::any_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 1, thrust::identity()), true);\n ASSERT_EQUAL(thrust::any_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 2, thrust::identity()), true);\n ASSERT_EQUAL(thrust::any_of(thrust::cuda::par.on(s), v.begin() + 1, v.begin() + 2, thrust::identity()), false);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestAnyOfCudaStreams);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid none_of_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Function f, Iterator2 result)\n{\n *result = thrust::none_of(exec, first, last, f);\n}\n\n\ntemplate\nvoid TestNoneOfDevice(ExecutionPolicy exec)\n{\n typedef int T;\n \n thrust::device_vector v(3, 1);\n thrust::device_vector result(1);\n \n none_of_kernel<<<1,1>>>(exec, v.begin(), v.end(), thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n \n v[1] = 0;\n \n none_of_kernel<<<1,1>>>(exec, v.begin(), v.end(), thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n \n none_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 0, thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n\n none_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 1, thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n\n none_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 2, thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n\n none_of_kernel<<<1,1>>>(exec, v.begin() + 1, v.begin() + 2, thrust::identity(), result.begin());\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n}\n\n\nvoid TestNoneOfDeviceSeq()\n{\n TestNoneOfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestNoneOfDeviceSeq);\n\n\nvoid TestNoneOfDeviceDevice()\n{\n TestNoneOfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestNoneOfDeviceDevice);\n#endif\n\n\nvoid TestNoneOfCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector v(3, 1);\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n\n ASSERT_EQUAL(thrust::none_of(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()), false);\n\n v[1] = 0;\n \n ASSERT_EQUAL(thrust::none_of(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()), false);\n\n ASSERT_EQUAL(thrust::none_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 0, thrust::identity()), true);\n ASSERT_EQUAL(thrust::none_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 1, thrust::identity()), false);\n ASSERT_EQUAL(thrust::none_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 2, thrust::identity()), false);\n ASSERT_EQUAL(thrust::none_of(thrust::cuda::par.on(s), v.begin() + 1, v.begin() + 2, thrust::identity()), true);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestNoneOfCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid all_of_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Function f, Iterator2 result)\n{\n *result = thrust::all_of(exec, first, last, f);\n}\n\n\ntemplate\nvoid TestAllOfDevice(ExecutionPolicy exec)\n{\n typedef int T;\n thrust::device_vector v(3, 1);\n thrust::device_vector result(1);\n \n all_of_kernel<<<1,1>>>(exec, v.begin(), v.end(), thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n \n v[1] = 0;\n \n all_of_kernel<<<1,1>>>(exec, v.begin(), v.end(), thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n \n all_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 0, thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n\n all_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 1, thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n\n all_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 2, thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n\n all_of_kernel<<<1,1>>>(exec, v.begin() + 1, v.begin() + 2, thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n}\n\n\nvoid TestAllOfDeviceSeq()\n{\n TestAllOfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestAllOfDeviceSeq);\n\n\nvoid TestAllOfDeviceDevice()\n{\n TestAllOfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestAllOfDeviceDevice);\n#endif\n\n\nvoid TestAllOfCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n \n Vector v(3, 1);\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n ASSERT_EQUAL(thrust::all_of(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()), true);\n \n v[1] = 0;\n \n ASSERT_EQUAL(thrust::all_of(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()), false);\n \n ASSERT_EQUAL(thrust::all_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 0, thrust::identity()), true);\n ASSERT_EQUAL(thrust::all_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 1, thrust::identity()), true);\n ASSERT_EQUAL(thrust::all_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 2, thrust::identity()), false);\n ASSERT_EQUAL(thrust::all_of(thrust::cuda::par.on(s), v.begin() + 1, v.begin() + 2, thrust::identity()), false);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestAllOfCudaStreams);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid any_of_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Function f, Iterator2 result)\n{\n *result = thrust::any_of(exec, first, last, f);\n}\n\n\ntemplate\nvoid TestAnyOfDevice(ExecutionPolicy exec)\n{\n typedef int T;\n \n thrust::device_vector v(3, 1);\n thrust::device_vector result(1);\n \n any_of_kernel<<<1,1>>>(exec, v.begin(), v.end(), thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n \n v[1] = 0;\n \n any_of_kernel<<<1,1>>>(exec, v.begin(), v.end(), thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n \n any_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 0, thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n\n any_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 1, thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n\n any_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 2, thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n\n any_of_kernel<<<1,1>>>(exec, v.begin() + 1, v.begin() + 2, thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n}\n\n\nvoid TestAnyOfDeviceSeq()\n{\n TestAnyOfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestAnyOfDeviceSeq);\n\n\nvoid TestAnyOfDeviceDevice()\n{\n TestAnyOfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestAnyOfDeviceDevice);\n#endif\n\n\nvoid TestAnyOfCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector v(3, 1);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n ASSERT_EQUAL(thrust::any_of(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()), true);\n\n v[1] = 0;\n \n ASSERT_EQUAL(thrust::any_of(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()), true);\n\n ASSERT_EQUAL(thrust::any_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 0, thrust::identity()), false);\n ASSERT_EQUAL(thrust::any_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 1, thrust::identity()), true);\n ASSERT_EQUAL(thrust::any_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 2, thrust::identity()), true);\n ASSERT_EQUAL(thrust::any_of(thrust::cuda::par.on(s), v.begin() + 1, v.begin() + 2, thrust::identity()), false);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestAnyOfCudaStreams);\n\n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid none_of_kernel(ExecutionPolicy exec, Iterator first, Iterator last, Function f, Iterator2 result)\n{\n *result = thrust::none_of(exec, first, last, f);\n}\n\n\ntemplate\nvoid TestNoneOfDevice(ExecutionPolicy exec)\n{\n typedef int T;\n \n thrust::device_vector v(3, 1);\n thrust::device_vector result(1);\n \n none_of_kernel<<<1,1>>>(exec, v.begin(), v.end(), thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n \n v[1] = 0;\n \n none_of_kernel<<<1,1>>>(exec, v.begin(), v.end(), thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n \n none_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 0, thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n\n none_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 1, thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n\n none_of_kernel<<<1,1>>>(exec, v.begin() + 0, v.begin() + 2, thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(false, result[0]);\n\n none_of_kernel<<<1,1>>>(exec, v.begin() + 1, v.begin() + 2, thrust::identity(), result.begin());\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n\n ASSERT_EQUAL(true, result[0]);\n}\n\n\nvoid TestNoneOfDeviceSeq()\n{\n TestNoneOfDevice(thrust::seq);\n}\nDECLARE_UNITTEST(TestNoneOfDeviceSeq);\n\n\nvoid TestNoneOfDeviceDevice()\n{\n TestNoneOfDevice(thrust::device);\n}\nDECLARE_UNITTEST(TestNoneOfDeviceDevice);\n#endif\n\n\nvoid TestNoneOfCudaStreams()\n{\n typedef thrust::device_vector Vector;\n typedef Vector::value_type T;\n\n Vector v(3, 1);\n\n hipStream_t s;\n hipStreamCreate(&s);\n\n ASSERT_EQUAL(thrust::none_of(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()), false);\n\n v[1] = 0;\n \n ASSERT_EQUAL(thrust::none_of(thrust::cuda::par.on(s), v.begin(), v.end(), thrust::identity()), false);\n\n ASSERT_EQUAL(thrust::none_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 0, thrust::identity()), true);\n ASSERT_EQUAL(thrust::none_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 1, thrust::identity()), false);\n ASSERT_EQUAL(thrust::none_of(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 2, thrust::identity()), false);\n ASSERT_EQUAL(thrust::none_of(thrust::cuda::par.on(s), v.begin() + 1, v.begin() + 2, thrust::identity()), true);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestNoneOfCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/logical.cu"} {"input": "#include \n#include \n#include \n#include \n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid fill_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T value)\n{\n thrust::fill(exec, first, last, value);\n}\n\n\ntemplate\nvoid TestFillDevice(ExecutionPolicy exec, size_t n)\n{\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n thrust::fill(h_data.begin() + std::min((size_t)1, n), h_data.begin() + std::min((size_t)3, n), (T) 0);\n\n fill_kernel<<<1,1>>>(exec, d_data.begin() + std::min((size_t)1, n), d_data.begin() + std::min((size_t)3, n), (T) 0);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n thrust::fill(h_data.begin() + std::min((size_t)117, n), h_data.begin() + std::min((size_t)367, n), (T) 1);\n\n fill_kernel<<<1,1>>>(exec, d_data.begin() + std::min((size_t)117, n), d_data.begin() + std::min((size_t)367, n), (T) 1);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n thrust::fill(h_data.begin() + std::min((size_t)8, n), h_data.begin() + std::min((size_t)259, n), (T) 2);\n\n fill_kernel<<<1,1>>>(exec, d_data.begin() + std::min((size_t)8, n), d_data.begin() + std::min((size_t)259, n), (T) 2);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n thrust::fill(h_data.begin() + std::min((size_t)3, n), h_data.end(), (T) 3);\n\n fill_kernel<<<1,1>>>(exec, d_data.begin() + std::min((size_t)3, n), d_data.end(), (T) 3);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n thrust::fill(h_data.begin(), h_data.end(), (T) 4);\n\n fill_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), (T) 4);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n}\n\ntemplate\nvoid TestFillDeviceSeq(size_t n)\n{\n TestFillDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestFillDeviceSeq);\n\ntemplate\nvoid TestFillDeviceDevice(size_t n)\n{\n TestFillDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestFillDeviceDevice);\n\n\ntemplate\n__global__\nvoid fill_n_kernel(ExecutionPolicy exec, Iterator first, Size n, T value)\n{\n thrust::fill_n(exec, first, n, value);\n}\n\n\ntemplate\nvoid TestFillNDevice(ExecutionPolicy exec, size_t n)\n{\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n size_t begin_offset = std::min(1,n);\n\n thrust::fill_n(h_data.begin() + begin_offset, std::min((size_t)3, n) - begin_offset, (T) 0);\n\n fill_n_kernel<<<1,1>>>(exec, d_data.begin() + begin_offset, std::min((size_t)3, n) - begin_offset, (T) 0);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n begin_offset = std::min(117, n);\n\n thrust::fill_n(h_data.begin() + begin_offset, std::min((size_t)367, n) - begin_offset, (T) 1);\n\n fill_n_kernel<<<1,1>>>(exec, d_data.begin() + begin_offset, std::min((size_t)367, n) - begin_offset, (T) 1);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n begin_offset = std::min(8, n);\n\n thrust::fill_n(h_data.begin() + begin_offset, std::min((size_t)259, n) - begin_offset, (T) 2);\n\n fill_n_kernel<<<1,1>>>(exec, d_data.begin() + begin_offset, std::min((size_t)259, n) - begin_offset, (T) 2);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n begin_offset = std::min(3, n);\n\n thrust::fill_n(h_data.begin() + begin_offset, h_data.size() - begin_offset, (T) 3);\n\n fill_n_kernel<<<1,1>>>(exec, d_data.begin() + begin_offset, d_data.size() - begin_offset, (T) 3);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n thrust::fill_n(h_data.begin(), h_data.size(), (T) 4);\n\n fill_n_kernel<<<1,1>>>(exec, d_data.begin(), d_data.size(), (T) 4);\n {\n cudaError_t const err = cudaDeviceSynchronize();\n ASSERT_EQUAL(cudaSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n}\n\ntemplate\nvoid TestFillNDeviceSeq(size_t n)\n{\n TestFillNDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestFillNDeviceSeq);\n\ntemplate\nvoid TestFillNDeviceDevice(size_t n)\n{\n TestFillNDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestFillNDeviceDevice);\n#endif\n\nvoid TestFillCudaStreams()\n{\n thrust::device_vector v(5);\n v[0] = 0; v[1] = 1; v[2] = 2; v[3] = 3; v[4] = 4;\n\n cudaStream_t s;\n cudaStreamCreate(&s);\n \n thrust::fill(thrust::cuda::par.on(s), v.begin() + 1, v.begin() + 4, 7);\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], 7);\n ASSERT_EQUAL(v[2], 7);\n ASSERT_EQUAL(v[3], 7);\n ASSERT_EQUAL(v[4], 4);\n \n thrust::fill(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 3, 8);\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(v[0], 8);\n ASSERT_EQUAL(v[1], 8);\n ASSERT_EQUAL(v[2], 8);\n ASSERT_EQUAL(v[3], 7);\n ASSERT_EQUAL(v[4], 4);\n \n thrust::fill(thrust::cuda::par.on(s), v.begin() + 2, v.end(), 9);\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(v[0], 8);\n ASSERT_EQUAL(v[1], 8);\n ASSERT_EQUAL(v[2], 9);\n ASSERT_EQUAL(v[3], 9);\n ASSERT_EQUAL(v[4], 9);\n \n thrust::fill(thrust::cuda::par.on(s), v.begin(), v.end(), 1);\n cudaStreamSynchronize(s);\n \n ASSERT_EQUAL(v[0], 1);\n ASSERT_EQUAL(v[1], 1);\n ASSERT_EQUAL(v[2], 1);\n ASSERT_EQUAL(v[3], 1);\n ASSERT_EQUAL(v[4], 1);\n\n cudaStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestFillCudaStreams);\n\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n#ifdef THRUST_TEST_DEVICE_SIDE\ntemplate\n__global__\nvoid fill_kernel(ExecutionPolicy exec, Iterator first, Iterator last, T value)\n{\n thrust::fill(exec, first, last, value);\n}\n\n\ntemplate\nvoid TestFillDevice(ExecutionPolicy exec, size_t n)\n{\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n thrust::fill(h_data.begin() + std::min((size_t)1, n), h_data.begin() + std::min((size_t)3, n), (T) 0);\n\n fill_kernel<<<1,1>>>(exec, d_data.begin() + std::min((size_t)1, n), d_data.begin() + std::min((size_t)3, n), (T) 0);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n thrust::fill(h_data.begin() + std::min((size_t)117, n), h_data.begin() + std::min((size_t)367, n), (T) 1);\n\n fill_kernel<<<1,1>>>(exec, d_data.begin() + std::min((size_t)117, n), d_data.begin() + std::min((size_t)367, n), (T) 1);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n thrust::fill(h_data.begin() + std::min((size_t)8, n), h_data.begin() + std::min((size_t)259, n), (T) 2);\n\n fill_kernel<<<1,1>>>(exec, d_data.begin() + std::min((size_t)8, n), d_data.begin() + std::min((size_t)259, n), (T) 2);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n thrust::fill(h_data.begin() + std::min((size_t)3, n), h_data.end(), (T) 3);\n\n fill_kernel<<<1,1>>>(exec, d_data.begin() + std::min((size_t)3, n), d_data.end(), (T) 3);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n thrust::fill(h_data.begin(), h_data.end(), (T) 4);\n\n fill_kernel<<<1,1>>>(exec, d_data.begin(), d_data.end(), (T) 4);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n}\n\ntemplate\nvoid TestFillDeviceSeq(size_t n)\n{\n TestFillDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestFillDeviceSeq);\n\ntemplate\nvoid TestFillDeviceDevice(size_t n)\n{\n TestFillDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestFillDeviceDevice);\n\n\ntemplate\n__global__\nvoid fill_n_kernel(ExecutionPolicy exec, Iterator first, Size n, T value)\n{\n thrust::fill_n(exec, first, n, value);\n}\n\n\ntemplate\nvoid TestFillNDevice(ExecutionPolicy exec, size_t n)\n{\n thrust::host_vector h_data = unittest::random_integers(n);\n thrust::device_vector d_data = h_data;\n \n size_t begin_offset = std::min(1,n);\n\n thrust::fill_n(h_data.begin() + begin_offset, std::min((size_t)3, n) - begin_offset, (T) 0);\n\n fill_n_kernel<<<1,1>>>(exec, d_data.begin() + begin_offset, std::min((size_t)3, n) - begin_offset, (T) 0);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n begin_offset = std::min(117, n);\n\n thrust::fill_n(h_data.begin() + begin_offset, std::min((size_t)367, n) - begin_offset, (T) 1);\n\n fill_n_kernel<<<1,1>>>(exec, d_data.begin() + begin_offset, std::min((size_t)367, n) - begin_offset, (T) 1);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n begin_offset = std::min(8, n);\n\n thrust::fill_n(h_data.begin() + begin_offset, std::min((size_t)259, n) - begin_offset, (T) 2);\n\n fill_n_kernel<<<1,1>>>(exec, d_data.begin() + begin_offset, std::min((size_t)259, n) - begin_offset, (T) 2);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n begin_offset = std::min(3, n);\n\n thrust::fill_n(h_data.begin() + begin_offset, h_data.size() - begin_offset, (T) 3);\n\n fill_n_kernel<<<1,1>>>(exec, d_data.begin() + begin_offset, d_data.size() - begin_offset, (T) 3);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n \n thrust::fill_n(h_data.begin(), h_data.size(), (T) 4);\n\n fill_n_kernel<<<1,1>>>(exec, d_data.begin(), d_data.size(), (T) 4);\n {\n hipError_t const err = hipDeviceSynchronize();\n ASSERT_EQUAL(hipSuccess, err);\n }\n \n ASSERT_EQUAL(h_data, d_data);\n}\n\ntemplate\nvoid TestFillNDeviceSeq(size_t n)\n{\n TestFillNDevice(thrust::seq, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestFillNDeviceSeq);\n\ntemplate\nvoid TestFillNDeviceDevice(size_t n)\n{\n TestFillNDevice(thrust::device, n);\n}\nDECLARE_VARIABLE_UNITTEST(TestFillNDeviceDevice);\n#endif\n\nvoid TestFillCudaStreams()\n{\n thrust::device_vector v(5);\n v[0] = 0; v[1] = 1; v[2] = 2; v[3] = 3; v[4] = 4;\n\n hipStream_t s;\n hipStreamCreate(&s);\n \n thrust::fill(thrust::cuda::par.on(s), v.begin() + 1, v.begin() + 4, 7);\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(v[0], 0);\n ASSERT_EQUAL(v[1], 7);\n ASSERT_EQUAL(v[2], 7);\n ASSERT_EQUAL(v[3], 7);\n ASSERT_EQUAL(v[4], 4);\n \n thrust::fill(thrust::cuda::par.on(s), v.begin() + 0, v.begin() + 3, 8);\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(v[0], 8);\n ASSERT_EQUAL(v[1], 8);\n ASSERT_EQUAL(v[2], 8);\n ASSERT_EQUAL(v[3], 7);\n ASSERT_EQUAL(v[4], 4);\n \n thrust::fill(thrust::cuda::par.on(s), v.begin() + 2, v.end(), 9);\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(v[0], 8);\n ASSERT_EQUAL(v[1], 8);\n ASSERT_EQUAL(v[2], 9);\n ASSERT_EQUAL(v[3], 9);\n ASSERT_EQUAL(v[4], 9);\n \n thrust::fill(thrust::cuda::par.on(s), v.begin(), v.end(), 1);\n hipStreamSynchronize(s);\n \n ASSERT_EQUAL(v[0], 1);\n ASSERT_EQUAL(v[1], 1);\n ASSERT_EQUAL(v[2], 1);\n ASSERT_EQUAL(v[3], 1);\n ASSERT_EQUAL(v[4], 1);\n\n hipStreamDestroy(s);\n}\nDECLARE_UNITTEST(TestFillCudaStreams);\n\n", "source": "/kaggle/working/thrust/testing/cuda/fill.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n\n__global__ void dummy_kernel() {}\n\nbool binary_exists_for_current_device()\n{\n // check against the dummy_kernel\n // if we're unable to get the attributes, then\n // we didn't compile a binary compatible with the current device\n cudaFuncAttributes attr;\n cudaError_t error = cudaFuncGetAttributes(&attr, dummy_kernel);\n\n // clear the CUDA global error state if we just set it, so that\n // check_cuda_error doesn't complain\n if (cudaSuccess != error) (void)cudaGetLastError();\n\n return cudaSuccess == error;\n}\n\nvoid list_devices(void)\n{\n int deviceCount;\n cudaGetDeviceCount(&deviceCount);\n if(deviceCount == 0)\n {\n std::cout << \"There is no device supporting CUDA\" << std::endl;\n }\n\n int selected_device;\n cudaGetDevice(&selected_device);\n\n for (int dev = 0; dev < deviceCount; ++dev)\n {\n cudaDeviceProp deviceProp;\n cudaGetDeviceProperties(&deviceProp, dev);\n\n if(dev == 0)\n {\n if(deviceProp.major == 9999 && deviceProp.minor == 9999)\n std::cout << \"There is no device supporting CUDA.\" << std::endl;\n else if(deviceCount == 1)\n std::cout << \"There is 1 device supporting CUDA\" << std:: endl;\n else\n std::cout << \"There are \" << deviceCount << \" devices supporting CUDA\" << std:: endl;\n }\n\n std::cout << \"\\nDevice \" << dev << \": \\\"\" << deviceProp.name << \"\\\"\";\n if(dev == selected_device)\n std::cout << \" [SELECTED]\";\n std::cout << std::endl;\n\n std::cout << \" Major revision number: \" << deviceProp.major << std::endl;\n std::cout << \" Minor revision number: \" << deviceProp.minor << std::endl;\n std::cout << \" Total amount of global memory: \" << deviceProp.totalGlobalMem << \" bytes\" << std::endl;\n }\n std::cout << std::endl;\n}\n\n// provide next, which c++03 doesn't have\ntemplate Iterator my_next(Iterator iter)\n{\n return ++iter;\n}\n\n\nstd::vector CUDATestDriver::target_devices(const ArgumentMap &kwargs)\n{\n std::vector result;\n\n // by default, test all devices in the system (device id -1)\n int device_id = kwargs.count(\"device\") ? atoi(kwargs.find(\"device\")->second.c_str()) : -1;\n\n if(device_id < 0)\n {\n // target all devices in the system\n int count = 0;\n cudaGetDeviceCount(&count);\n\n result.resize(count);\n std::iota(result.begin(), result.end(), 0);\n }\n else\n {\n // target the specified device\n result = std::vector(1,device_id);\n }\n\n return result;\n}\n\nbool CUDATestDriver::check_cuda_error(bool concise)\n{\n cudaError_t const error = cudaGetLastError();\n if(cudaSuccess != error)\n {\n if(!concise)\n {\n std::cout << \"[ERROR] CUDA error detected before running tests: [\"\n << std::string(cudaGetErrorName(error))\n << \": \"\n << std::string(cudaGetErrorString(error))\n << \"]\" << std::endl;\n }\n }\n\n return cudaSuccess != error;\n}\n\nbool CUDATestDriver::post_test_smoke_check(const UnitTest &test, bool concise)\n{\n cudaError_t const error = cudaDeviceSynchronize();\n if(cudaSuccess != error)\n {\n if(!concise)\n {\n std::cout << \"\\t[ERROR] CUDA error detected after running \" << test.name << \": [\"\n << std::string(cudaGetErrorName(error))\n << \": \"\n << std::string(cudaGetErrorString(error))\n << \"]\" << std::endl;\n }\n }\n\n return cudaSuccess == error;\n}\n\nbool CUDATestDriver::run_tests(const ArgumentSet &args, const ArgumentMap &kwargs)\n{\n bool verbose = kwargs.count(\"verbose\");\n bool concise = kwargs.count(\"concise\");\n\n if(verbose && concise)\n {\n std::cout << \"--verbose and --concise cannot be used together\" << std::endl;\n exit(EXIT_FAILURE);\n }\n\n // check error status before doing anything\n if(check_cuda_error(concise)) return false;\n\n bool result = true;\n\n if(kwargs.count(\"verbose\"))\n {\n list_devices();\n }\n\n // figure out which devices to target\n std::vector devices = target_devices(kwargs);\n\n // target each device\n for(std::vector::iterator device = devices.begin();\n device != devices.end();\n ++device)\n {\n cudaDeviceSynchronize();\n\n // set the device\n cudaSetDevice(*device);\n\n // check if a binary exists for this device\n // if none exists, skip the device silently unless this is the only one we're targeting\n if(devices.size() > 1 && !binary_exists_for_current_device())\n {\n // note which device we're skipping\n cudaDeviceProp deviceProp;\n cudaGetDeviceProperties(&deviceProp, *device);\n\n std::cout << \"Skipping Device \" << *device << \": \\\"\" << deviceProp.name << \"\\\"\" << std::endl;\n\n continue;\n }\n\n if(!concise)\n {\n // note which device we're testing\n cudaDeviceProp deviceProp;\n cudaGetDeviceProperties(&deviceProp, *device);\n\n std::cout << \"Testing Device \" << *device << \": \\\"\" << deviceProp.name << \"\\\"\" << std::endl;\n }\n\n // check error status before running any tests\n if(check_cuda_error(concise)) return false;\n\n // run tests\n result &= UnitTestDriver::run_tests(args, kwargs);\n\n if(!concise && my_next(device) != devices.end())\n {\n // provide some separation between the output of separate tests\n std::cout << std::endl;\n }\n }\n\n return result;\n}\n\nint CUDATestDriver::current_device_architecture() const\n{\n int current = -1;\n cudaGetDevice(¤t);\n cudaDeviceProp deviceProp;\n cudaGetDeviceProperties(&deviceProp, current);\n\n return 100 * deviceProp.major + 10 * deviceProp.minor;\n}\n\nUnitTestDriver &driver_instance(thrust::system::cuda::tag)\n{\n static CUDATestDriver s_instance;\n return s_instance;\n}\n\n", "output": "#include \n#include \n#include \n#include \n#include \n\n__global__ void dummy_kernel() {}\n\nbool binary_exists_for_current_device()\n{\n // check against the dummy_kernel\n // if we're unable to get the attributes, then\n // we didn't compile a binary compatible with the current device\n hipFuncAttributes attr;\n hipError_t error = hipFuncGetAttributes(&attr,reinterpret_cast(dummy_kernel));\n\n // clear the CUDA global error state if we just set it, so that\n // check_cuda_error doesn't complain\n if (hipSuccess != error) (void)hipGetLastError();\n\n return hipSuccess == error;\n}\n\nvoid list_devices(void)\n{\n int deviceCount;\n hipGetDeviceCount(&deviceCount);\n if(deviceCount == 0)\n {\n std::cout << \"There is no device supporting CUDA\" << std::endl;\n }\n\n int selected_device;\n hipGetDevice(&selected_device);\n\n for (int dev = 0; dev < deviceCount; ++dev)\n {\n hipDeviceProp_t deviceProp;\n hipGetDeviceProperties(&deviceProp, dev);\n\n if(dev == 0)\n {\n if(deviceProp.major == 9999 && deviceProp.minor == 9999)\n std::cout << \"There is no device supporting CUDA.\" << std::endl;\n else if(deviceCount == 1)\n std::cout << \"There is 1 device supporting CUDA\" << std:: endl;\n else\n std::cout << \"There are \" << deviceCount << \" devices supporting CUDA\" << std:: endl;\n }\n\n std::cout << \"\\nDevice \" << dev << \": \\\"\" << deviceProp.name << \"\\\"\";\n if(dev == selected_device)\n std::cout << \" [SELECTED]\";\n std::cout << std::endl;\n\n std::cout << \" Major revision number: \" << deviceProp.major << std::endl;\n std::cout << \" Minor revision number: \" << deviceProp.minor << std::endl;\n std::cout << \" Total amount of global memory: \" << deviceProp.totalGlobalMem << \" bytes\" << std::endl;\n }\n std::cout << std::endl;\n}\n\n// provide next, which c++03 doesn't have\ntemplate Iterator my_next(Iterator iter)\n{\n return ++iter;\n}\n\n\nstd::vector CUDATestDriver::target_devices(const ArgumentMap &kwargs)\n{\n std::vector result;\n\n // by default, test all devices in the system (device id -1)\n int device_id = kwargs.count(\"device\") ? atoi(kwargs.find(\"device\")->second.c_str()) : -1;\n\n if(device_id < 0)\n {\n // target all devices in the system\n int count = 0;\n hipGetDeviceCount(&count);\n\n result.resize(count);\n std::iota(result.begin(), result.end(), 0);\n }\n else\n {\n // target the specified device\n result = std::vector(1,device_id);\n }\n\n return result;\n}\n\nbool CUDATestDriver::check_cuda_error(bool concise)\n{\n hipError_t const error = hipGetLastError();\n if(hipSuccess != error)\n {\n if(!concise)\n {\n std::cout << \"[ERROR] CUDA error detected before running tests: [\"\n << std::string(hipGetErrorName(error))\n << \": \"\n << std::string(hipGetErrorString(error))\n << \"]\" << std::endl;\n }\n }\n\n return hipSuccess != error;\n}\n\nbool CUDATestDriver::post_test_smoke_check(const UnitTest &test, bool concise)\n{\n hipError_t const error = hipDeviceSynchronize();\n if(hipSuccess != error)\n {\n if(!concise)\n {\n std::cout << \"\\t[ERROR] CUDA error detected after running \" << test.name << \": [\"\n << std::string(hipGetErrorName(error))\n << \": \"\n << std::string(hipGetErrorString(error))\n << \"]\" << std::endl;\n }\n }\n\n return hipSuccess == error;\n}\n\nbool CUDATestDriver::run_tests(const ArgumentSet &args, const ArgumentMap &kwargs)\n{\n bool verbose = kwargs.count(\"verbose\");\n bool concise = kwargs.count(\"concise\");\n\n if(verbose && concise)\n {\n std::cout << \"--verbose and --concise cannot be used together\" << std::endl;\n exit(EXIT_FAILURE);\n }\n\n // check error status before doing anything\n if(check_cuda_error(concise)) return false;\n\n bool result = true;\n\n if(kwargs.count(\"verbose\"))\n {\n list_devices();\n }\n\n // figure out which devices to target\n std::vector devices = target_devices(kwargs);\n\n // target each device\n for(std::vector::iterator device = devices.begin();\n device != devices.end();\n ++device)\n {\n hipDeviceSynchronize();\n\n // set the device\n hipSetDevice(*device);\n\n // check if a binary exists for this device\n // if none exists, skip the device silently unless this is the only one we're targeting\n if(devices.size() > 1 && !binary_exists_for_current_device())\n {\n // note which device we're skipping\n hipDeviceProp_t deviceProp;\n hipGetDeviceProperties(&deviceProp, *device);\n\n std::cout << \"Skipping Device \" << *device << \": \\\"\" << deviceProp.name << \"\\\"\" << std::endl;\n\n continue;\n }\n\n if(!concise)\n {\n // note which device we're testing\n hipDeviceProp_t deviceProp;\n hipGetDeviceProperties(&deviceProp, *device);\n\n std::cout << \"Testing Device \" << *device << \": \\\"\" << deviceProp.name << \"\\\"\" << std::endl;\n }\n\n // check error status before running any tests\n if(check_cuda_error(concise)) return false;\n\n // run tests\n result &= UnitTestDriver::run_tests(args, kwargs);\n\n if(!concise && my_next(device) != devices.end())\n {\n // provide some separation between the output of separate tests\n std::cout << std::endl;\n }\n }\n\n return result;\n}\n\nint CUDATestDriver::current_device_architecture() const\n{\n int current = -1;\n hipGetDevice(¤t);\n hipDeviceProp_t deviceProp;\n hipGetDeviceProperties(&deviceProp, current);\n\n return 100 * deviceProp.major + 10 * deviceProp.minor;\n}\n\nUnitTestDriver &driver_instance(thrust::system::cuda::tag)\n{\n static CUDATestDriver s_instance;\n return s_instance;\n}\n\n", "source": "/kaggle/working/thrust/testing/unittest/cuda/testframework.cu"} {"input": "#include \n\nvoid helloCPU()\n{\n printf(\"Hello from the CPU.\\n\");\n}\n\n/*\n * The addition of `__global__` signifies that this function\n * should be launced on the GPU.\n */\n\n__global__ void helloGPU()\n{\n printf(\"Hello from the GPU.\\n\");\n}\n\nint main()\n{\n helloCPU();\n\n\n /*\n * Add an execution configuration with the <<<...>>> syntax\n * will launch this function as a kernel on the GPU.\n */\n\n helloGPU<<<1, 1>>>();\n\n /*\n * `cudaDeviceSynchronize` will block the CPU stream until\n * all GPU kernels have completed.\n */\n\n cudaDeviceSynchronize();\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n\nvoid helloCPU()\n{\n printf(\"Hello from the CPU.\\n\");\n}\n\n/*\n * The addition of `__global__` signifies that this function\n * should be launced on the GPU.\n */\n\n__global__ void helloGPU()\n{\n printf(\"Hello from the GPU.\\n\");\n}\n\nint main()\n{\n helloCPU();\n\n\n /*\n * Add an execution configuration with the <<<...>>> syntax\n * will launch this function as a kernel on the GPU.\n */\n\n helloGPU<<<1, 1>>>();\n\n /*\n * `cudaDeviceSynchronize` will block the CPU stream until\n * all GPU kernels have completed.\n */\n\n hipDeviceSynchronize();\n}\n", "source": "/kaggle/working/GPU_Acceleration_Using_CUDA_C_CPP/01-hello-gpu-solution.cu"} {"input": "#include \n#include \n\n// Simple define to index into a 1D array from 2D space\n#define I2D(num, c, r) ((r)*(num)+(c))\n\n__global__\nvoid step_kernel_mod(int ni, int nj, float fact, float* temp_in, float* temp_out)\n{\n int i00, im10, ip10, i0m1, i0p1;\n float d2tdx2, d2tdy2;\n\n int j = blockIdx.x * blockDim.x + threadIdx.x;\n int i = blockIdx.y * blockDim.y + threadIdx.y;\n\n // loop over all points in domain (except boundary)\n if (j > 0 && i > 0 && j < nj-1 && i < ni-1) {\n // find indices into linear memory\n // for central point and neighbours\n i00 = I2D(ni, i, j);\n im10 = I2D(ni, i-1, j);\n ip10 = I2D(ni, i+1, j);\n i0m1 = I2D(ni, i, j-1);\n i0p1 = I2D(ni, i, j+1);\n\n // evaluate derivatives\n d2tdx2 = temp_in[im10]-2*temp_in[i00]+temp_in[ip10];\n d2tdy2 = temp_in[i0m1]-2*temp_in[i00]+temp_in[i0p1];\n\n // update temperatures\n temp_out[i00] = temp_in[i00]+fact*(d2tdx2 + d2tdy2);\n }\n}\n\nvoid step_kernel_ref(int ni, int nj, float fact, float* temp_in, float* temp_out)\n{\n int i00, im10, ip10, i0m1, i0p1;\n float d2tdx2, d2tdy2;\n\n\n // loop over all points in domain (except boundary)\n for ( int j=1; j < nj-1; j++ ) {\n for ( int i=1; i < ni-1; i++ ) {\n // find indices into linear memory\n // for central point and neighbours\n i00 = I2D(ni, i, j);\n im10 = I2D(ni, i-1, j);\n ip10 = I2D(ni, i+1, j);\n i0m1 = I2D(ni, i, j-1);\n i0p1 = I2D(ni, i, j+1);\n\n // evaluate derivatives\n d2tdx2 = temp_in[im10]-2*temp_in[i00]+temp_in[ip10];\n d2tdy2 = temp_in[i0m1]-2*temp_in[i00]+temp_in[i0p1];\n\n // update temperatures\n temp_out[i00] = temp_in[i00]+fact*(d2tdx2 + d2tdy2);\n }\n }\n}\n\nint main()\n{\n int istep;\n int nstep = 200; // number of time steps\n\n // Specify our 2D dimensions\n const int ni = 200;\n const int nj = 100;\n float tfac = 8.418e-5; // thermal diffusivity of silver\n\n float *temp1_ref, *temp2_ref, *temp1, *temp2, *temp_tmp;\n\n const int size = ni * nj * sizeof(float);\n\n temp1_ref = (float*)malloc(size);\n temp2_ref = (float*)malloc(size);\n cudaMallocManaged(&temp1, size);\n cudaMallocManaged(&temp2, size);\n\n // Initialize with random data\n for( int i = 0; i < ni*nj; ++i) {\n temp1_ref[i] = temp2_ref[i] = temp1[i] = temp2[i] = (float)rand()/(float)(RAND_MAX/100.0f);\n }\n\n // Execute the CPU-only reference version\n for (istep=0; istep < nstep; istep++) {\n step_kernel_ref(ni, nj, tfac, temp1_ref, temp2_ref);\n\n // swap the temperature pointers\n temp_tmp = temp1_ref;\n temp1_ref = temp2_ref;\n temp2_ref= temp_tmp;\n }\n\n dim3 tblocks(32, 16, 1);\n dim3 grid((nj/tblocks.x)+1, (ni/tblocks.y)+1, 1);\n cudaError_t ierrSync, ierrAsync;\n\n // Execute the modified version using same data\n for (istep=0; istep < nstep; istep++) {\n step_kernel_mod<<< grid, tblocks >>>(ni, nj, tfac, temp1, temp2);\n\n ierrSync = cudaGetLastError();\n ierrAsync = cudaDeviceSynchronize(); // Wait for the GPU to finish\n if (ierrSync != cudaSuccess) { printf(\"Sync error: %s\\n\", cudaGetErrorString(ierrSync)); }\n if (ierrAsync != cudaSuccess) { printf(\"Async error: %s\\n\", cudaGetErrorString(ierrAsync)); }\n\n // swap the temperature pointers\n temp_tmp = temp1;\n temp1 = temp2;\n temp2= temp_tmp;\n }\n\n float maxError = 0;\n // Output should always be stored in the temp1 and temp1_ref at this point\n for( int i = 0; i < ni*nj; ++i ) {\n if (abs(temp1[i]-temp1_ref[i]) > maxError) { maxError = abs(temp1[i]-temp1_ref[i]); }\n }\n\n // Check and see if our maxError is greater than an error bound\n if (maxError > 0.0005f)\n printf(\"Problem! The Max Error of %.5f is NOT within acceptable bounds.\\n\", maxError);\n else\n printf(\"The Max Error of %.5f is within acceptable bounds.\\n\", maxError);\n\n free( temp1_ref );\n free( temp2_ref );\n cudaFree( temp1 );\n cudaFree( temp2 );\n\n return 0;\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n\n// Simple define to index into a 1D array from 2D space\n#define I2D(num, c, r) ((r)*(num)+(c))\n\n__global__\nvoid step_kernel_mod(int ni, int nj, float fact, float* temp_in, float* temp_out)\n{\n int i00, im10, ip10, i0m1, i0p1;\n float d2tdx2, d2tdy2;\n\n int j = blockIdx.x * blockDim.x + threadIdx.x;\n int i = blockIdx.y * blockDim.y + threadIdx.y;\n\n // loop over all points in domain (except boundary)\n if (j > 0 && i > 0 && j < nj-1 && i < ni-1) {\n // find indices into linear memory\n // for central point and neighbours\n i00 = I2D(ni, i, j);\n im10 = I2D(ni, i-1, j);\n ip10 = I2D(ni, i+1, j);\n i0m1 = I2D(ni, i, j-1);\n i0p1 = I2D(ni, i, j+1);\n\n // evaluate derivatives\n d2tdx2 = temp_in[im10]-2*temp_in[i00]+temp_in[ip10];\n d2tdy2 = temp_in[i0m1]-2*temp_in[i00]+temp_in[i0p1];\n\n // update temperatures\n temp_out[i00] = temp_in[i00]+fact*(d2tdx2 + d2tdy2);\n }\n}\n\nvoid step_kernel_ref(int ni, int nj, float fact, float* temp_in, float* temp_out)\n{\n int i00, im10, ip10, i0m1, i0p1;\n float d2tdx2, d2tdy2;\n\n\n // loop over all points in domain (except boundary)\n for ( int j=1; j < nj-1; j++ ) {\n for ( int i=1; i < ni-1; i++ ) {\n // find indices into linear memory\n // for central point and neighbours\n i00 = I2D(ni, i, j);\n im10 = I2D(ni, i-1, j);\n ip10 = I2D(ni, i+1, j);\n i0m1 = I2D(ni, i, j-1);\n i0p1 = I2D(ni, i, j+1);\n\n // evaluate derivatives\n d2tdx2 = temp_in[im10]-2*temp_in[i00]+temp_in[ip10];\n d2tdy2 = temp_in[i0m1]-2*temp_in[i00]+temp_in[i0p1];\n\n // update temperatures\n temp_out[i00] = temp_in[i00]+fact*(d2tdx2 + d2tdy2);\n }\n }\n}\n\nint main()\n{\n int istep;\n int nstep = 200; // number of time steps\n\n // Specify our 2D dimensions\n const int ni = 200;\n const int nj = 100;\n float tfac = 8.418e-5; // thermal diffusivity of silver\n\n float *temp1_ref, *temp2_ref, *temp1, *temp2, *temp_tmp;\n\n const int size = ni * nj * sizeof(float);\n\n temp1_ref = (float*)malloc(size);\n temp2_ref = (float*)malloc(size);\n hipMallocManaged(&temp1, size);\n hipMallocManaged(&temp2, size);\n\n // Initialize with random data\n for( int i = 0; i < ni*nj; ++i) {\n temp1_ref[i] = temp2_ref[i] = temp1[i] = temp2[i] = (float)rand()/(float)(RAND_MAX/100.0f);\n }\n\n // Execute the CPU-only reference version\n for (istep=0; istep < nstep; istep++) {\n step_kernel_ref(ni, nj, tfac, temp1_ref, temp2_ref);\n\n // swap the temperature pointers\n temp_tmp = temp1_ref;\n temp1_ref = temp2_ref;\n temp2_ref= temp_tmp;\n }\n\n dim3 tblocks(32, 16, 1);\n dim3 grid((nj/tblocks.x)+1, (ni/tblocks.y)+1, 1);\n hipError_t ierrSync, ierrAsync;\n\n // Execute the modified version using same data\n for (istep=0; istep < nstep; istep++) {\n step_kernel_mod<<< grid, tblocks >>>(ni, nj, tfac, temp1, temp2);\n\n ierrSync = hipGetLastError();\n ierrAsync = hipDeviceSynchronize(); // Wait for the GPU to finish\n if (ierrSync != hipSuccess) { printf(\"Sync error: %s\\n\", hipGetErrorString(ierrSync)); }\n if (ierrAsync != hipSuccess) { printf(\"Async error: %s\\n\", hipGetErrorString(ierrAsync)); }\n\n // swap the temperature pointers\n temp_tmp = temp1;\n temp1 = temp2;\n temp2= temp_tmp;\n }\n\n float maxError = 0;\n // Output should always be stored in the temp1 and temp1_ref at this point\n for( int i = 0; i < ni*nj; ++i ) {\n if (abs(temp1[i]-temp1_ref[i]) > maxError) { maxError = abs(temp1[i]-temp1_ref[i]); }\n }\n\n // Check and see if our maxError is greater than an error bound\n if (maxError > 0.0005f)\n printf(\"Problem! The Max Error of %.5f is NOT within acceptable bounds.\\n\", maxError);\n else\n printf(\"The Max Error of %.5f is within acceptable bounds.\\n\", maxError);\n\n free( temp1_ref );\n free( temp2_ref );\n hipFree( temp1 );\n hipFree( temp2 );\n\n return 0;\n}\n", "source": "/kaggle/working/GPU_Acceleration_Using_CUDA_C_CPP/01-heat-conduction-solution.cu"} {"input": "#include \n\n#define N 64\n\n__global__ void matrixMulGPU( int * a, int * b, int * c )\n{\n /*\n * Build out this kernel.\n */\n}\n\n/*\n * This CPU function already works, and will run to create a solution matrix\n * against which to verify your work building out the matrixMulGPU kernel.\n */\n\nvoid matrixMulCPU( int * a, int * b, int * c )\n{\n int val = 0;\n\n for( int row = 0; row < N; ++row )\n for( int col = 0; col < N; ++col )\n {\n val = 0;\n for ( int k = 0; k < N; ++k )\n val += a[row * N + k] * b[k * N + col];\n c[row * N + col] = val;\n }\n}\n\nint main()\n{\n int *a, *b, *c_cpu, *c_gpu; // Allocate a solution matrix for both the CPU and the GPU operations\n\n int size = N * N * sizeof (int); // Number of bytes of an N x N matrix\n\n // Allocate memory\n cudaMallocManaged (&a, size);\n cudaMallocManaged (&b, size);\n cudaMallocManaged (&c_cpu, size);\n cudaMallocManaged (&c_gpu, size);\n\n // Initialize memory; create 2D matrices\n for( int row = 0; row < N; ++row )\n for( int col = 0; col < N; ++col )\n {\n a[row*N + col] = row;\n b[row*N + col] = col+2;\n c_cpu[row*N + col] = 0;\n c_gpu[row*N + col] = 0;\n }\n\n /*\n * Assign `threads_per_block` and `number_of_blocks` 2D values\n * that can be used in matrixMulGPU above.\n */\n\n dim3 threads_per_block;\n dim3 number_of_blocks;\n\n matrixMulGPU <<< number_of_blocks, threads_per_block >>> ( a, b, c_gpu );\n\n cudaDeviceSynchronize();\n\n // Call the CPU version to check our work\n matrixMulCPU( a, b, c_cpu );\n\n // Compare the two answers to make sure they are equal\n bool error = false;\n for( int row = 0; row < N && !error; ++row )\n for( int col = 0; col < N && !error; ++col )\n if (c_cpu[row * N + col] != c_gpu[row * N + col])\n {\n printf(\"FOUND ERROR at c[%d][%d]\\n\", row, col);\n error = true;\n break;\n }\n if (!error)\n printf(\"Success!\\n\");\n\n // Free all our allocated memory\n cudaFree(a); cudaFree(b);\n cudaFree( c_cpu ); cudaFree( c_gpu );\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n\n#define N 64\n\n__global__ void matrixMulGPU( int * a, int * b, int * c )\n{\n /*\n * Build out this kernel.\n */\n}\n\n/*\n * This CPU function already works, and will run to create a solution matrix\n * against which to verify your work building out the matrixMulGPU kernel.\n */\n\nvoid matrixMulCPU( int * a, int * b, int * c )\n{\n int val = 0;\n\n for( int row = 0; row < N; ++row )\n for( int col = 0; col < N; ++col )\n {\n val = 0;\n for ( int k = 0; k < N; ++k )\n val += a[row * N + k] * b[k * N + col];\n c[row * N + col] = val;\n }\n}\n\nint main()\n{\n int *a, *b, *c_cpu, *c_gpu; // Allocate a solution matrix for both the CPU and the GPU operations\n\n int size = N * N * sizeof (int); // Number of bytes of an N x N matrix\n\n // Allocate memory\n hipMallocManaged (&a, size);\n hipMallocManaged (&b, size);\n hipMallocManaged (&c_cpu, size);\n hipMallocManaged (&c_gpu, size);\n\n // Initialize memory; create 2D matrices\n for( int row = 0; row < N; ++row )\n for( int col = 0; col < N; ++col )\n {\n a[row*N + col] = row;\n b[row*N + col] = col+2;\n c_cpu[row*N + col] = 0;\n c_gpu[row*N + col] = 0;\n }\n\n /*\n * Assign `threads_per_block` and `number_of_blocks` 2D values\n * that can be used in matrixMulGPU above.\n */\n\n dim3 threads_per_block;\n dim3 number_of_blocks;\n\n matrixMulGPU <<< number_of_blocks, threads_per_block >>> ( a, b, c_gpu );\n\n hipDeviceSynchronize();\n\n // Call the CPU version to check our work\n matrixMulCPU( a, b, c_cpu );\n\n // Compare the two answers to make sure they are equal\n bool error = false;\n for( int row = 0; row < N && !error; ++row )\n for( int col = 0; col < N && !error; ++col )\n if (c_cpu[row * N + col] != c_gpu[row * N + col])\n {\n printf(\"FOUND ERROR at c[%d][%d]\\n\", row, col);\n error = true;\n break;\n }\n if (!error)\n printf(\"Success!\\n\");\n\n // Free all our allocated memory\n hipFree(a); hipFree(b);\n hipFree( c_cpu ); hipFree( c_gpu );\n}\n", "source": "/kaggle/working/GPU_Acceleration_Using_CUDA_C_CPP/01-matrix-multiply-2d.cu"} {"input": "#include \n\nvoid init(int *a, int N)\n{\n int i;\n for (i = 0; i < N; ++i)\n {\n a[i] = i;\n }\n}\n\n__global__\nvoid doubleElements(int *a, int N)\n{\n\n /*\n * Use a grid-stride loop so each thread does work\n * on more than one element in the array.\n */\n\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = gridDim.x * blockDim.x;\n\n for (int i = idx; i < N; i += stride)\n {\n a[i] *= 2;\n }\n}\n\nbool checkElementsAreDoubled(int *a, int N)\n{\n int i;\n for (i = 0; i < N; ++i)\n {\n if (a[i] != i*2) return false;\n }\n return true;\n}\n\nint main()\n{\n int N = 10000;\n int *a;\n\n size_t size = N * sizeof(int);\n cudaMallocManaged(&a, size);\n\n init(a, N);\n\n size_t threads_per_block = 256;\n size_t number_of_blocks = 32;\n\n doubleElements<<>>(a, N);\n cudaDeviceSynchronize();\n\n bool areDoubled = checkElementsAreDoubled(a, N);\n printf(\"All elements were doubled? %s\\n\", areDoubled ? \"TRUE\" : \"FALSE\");\n\n cudaFree(a);\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n\nvoid init(int *a, int N)\n{\n int i;\n for (i = 0; i < N; ++i)\n {\n a[i] = i;\n }\n}\n\n__global__\nvoid doubleElements(int *a, int N)\n{\n\n /*\n * Use a grid-stride loop so each thread does work\n * on more than one element in the array.\n */\n\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = gridDim.x * blockDim.x;\n\n for (int i = idx; i < N; i += stride)\n {\n a[i] *= 2;\n }\n}\n\nbool checkElementsAreDoubled(int *a, int N)\n{\n int i;\n for (i = 0; i < N; ++i)\n {\n if (a[i] != i*2) return false;\n }\n return true;\n}\n\nint main()\n{\n int N = 10000;\n int *a;\n\n size_t size = N * sizeof(int);\n hipMallocManaged(&a, size);\n\n init(a, N);\n\n size_t threads_per_block = 256;\n size_t number_of_blocks = 32;\n\n doubleElements<<>>(a, N);\n hipDeviceSynchronize();\n\n bool areDoubled = checkElementsAreDoubled(a, N);\n printf(\"All elements were doubled? %s\\n\", areDoubled ? \"TRUE\" : \"FALSE\");\n\n hipFree(a);\n}\n", "source": "/kaggle/working/GPU_Acceleration_Using_CUDA_C_CPP/03-grid-stride-double-solution.cu"} {"input": "#include \n\n\n/*\n * Notice the absence of the previously expected argument `N`.\n */\n\n__global__ void loop()\n{\n /*\n * This kernel does the work of only 1 iteration\n * of the original for loop. Indication of which\n * \"iteration\" is being executed by this kernel is\n * still available via `threadIdx.x`.\n */\n\n printf(\"This is iteration number %d\\n\", threadIdx.x);\n}\n\nint main()\n{\n /*\n * It is the execution context that sets how many \"iterations\"\n * of the \"loop\" will be done.\n */\n\n loop<<<1, 10>>>();\n cudaDeviceSynchronize();\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n\n\n/*\n * Notice the absence of the previously expected argument `N`.\n */\n\n__global__ void loop()\n{\n /*\n * This kernel does the work of only 1 iteration\n * of the original for loop. Indication of which\n * \"iteration\" is being executed by this kernel is\n * still available via `threadIdx.x`.\n */\n\n printf(\"This is iteration number %d\\n\", threadIdx.x);\n}\n\nint main()\n{\n /*\n * It is the execution context that sets how many \"iterations\"\n * of the \"loop\" will be done.\n */\n\n loop<<<1, 10>>>();\n hipDeviceSynchronize();\n}\n", "source": "/kaggle/working/GPU_Acceleration_Using_CUDA_C_CPP/01-single-block-loop-solution.cu"} {"input": "#include \n\n__global__ void printSuccessForCorrectExecutionConfiguration()\n{\n\n if(threadIdx.x == 1023 && blockIdx.x == 255)\n {\n printf(\"Success!\\n\");\n }\n}\n\nint main()\n{\n /*\n * This is one possible execution context that will make\n * the kernel launch print its success message.\n */\n\n printSuccessForCorrectExecutionConfiguration<<<256, 1024>>>();\n\n /*\n * Don't forget kernel execution is asynchronous and you must\n * sync on its completion.\n */\n\n cudaDeviceSynchronize();\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n\n__global__ void printSuccessForCorrectExecutionConfiguration()\n{\n\n if(threadIdx.x == 1023 && blockIdx.x == 255)\n {\n printf(\"Success!\\n\");\n }\n}\n\nint main()\n{\n /*\n * This is one possible execution context that will make\n * the kernel launch print its success message.\n */\n\n printSuccessForCorrectExecutionConfiguration<<<256, 1024>>>();\n\n /*\n * Don't forget kernel execution is asynchronous and you must\n * sync on its completion.\n */\n\n hipDeviceSynchronize();\n}\n", "source": "/kaggle/working/GPU_Acceleration_Using_CUDA_C_CPP/01-thread-and-block-idx-solution.cu"} {"input": "#include \n\nvoid init(int *a, int N)\n{\n int i;\n for (i = 0; i < N; ++i)\n {\n a[i] = i;\n }\n}\n\n__global__\nvoid doubleElements(int *a, int N)\n{\n int i;\n i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i < N)\n {\n a[i] *= 2;\n }\n}\n\nbool checkElementsAreDoubled(int *a, int N)\n{\n int i;\n for (i = 0; i < N; ++i)\n {\n if (a[i] != i*2) return false;\n }\n return true;\n}\n\nint main()\n{\n int N = 1000;\n int *a;\n\n size_t size = N * sizeof(int);\n\n /*\n * Use `cudaMallocManaged` to allocate pointer `a` available\n * on both the host and the device.\n */\n\n cudaMallocManaged(&a, size);\n\n init(a, N);\n\n size_t threads_per_block = 256;\n size_t number_of_blocks = (N + threads_per_block - 1) / threads_per_block;\n\n doubleElements<<>>(a, N);\n cudaDeviceSynchronize();\n\n bool areDoubled = checkElementsAreDoubled(a, N);\n printf(\"All elements were doubled? %s\\n\", areDoubled ? \"TRUE\" : \"FALSE\");\n\n /*\n * Use `cudaFree` to free memory allocated\n * with `cudaMallocManaged`.\n */\n\n cudaFree(a);\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n\nvoid init(int *a, int N)\n{\n int i;\n for (i = 0; i < N; ++i)\n {\n a[i] = i;\n }\n}\n\n__global__\nvoid doubleElements(int *a, int N)\n{\n int i;\n i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i < N)\n {\n a[i] *= 2;\n }\n}\n\nbool checkElementsAreDoubled(int *a, int N)\n{\n int i;\n for (i = 0; i < N; ++i)\n {\n if (a[i] != i*2) return false;\n }\n return true;\n}\n\nint main()\n{\n int N = 1000;\n int *a;\n\n size_t size = N * sizeof(int);\n\n /*\n * Use `cudaMallocManaged` to allocate pointer `a` available\n * on both the host and the device.\n */\n\n hipMallocManaged(&a, size);\n\n init(a, N);\n\n size_t threads_per_block = 256;\n size_t number_of_blocks = (N + threads_per_block - 1) / threads_per_block;\n\n doubleElements<<>>(a, N);\n hipDeviceSynchronize();\n\n bool areDoubled = checkElementsAreDoubled(a, N);\n printf(\"All elements were doubled? %s\\n\", areDoubled ? \"TRUE\" : \"FALSE\");\n\n /*\n * Use `cudaFree` to free memory allocated\n * with `cudaMallocManaged`.\n */\n\n hipFree(a);\n}\n", "source": "/kaggle/working/GPU_Acceleration_Using_CUDA_C_CPP/01-double-elements-solution.cu"} {"input": "#include \n\n#define N 64\n\n__global__ void matrixMulGPU( int * a, int * b, int * c )\n{\n int val = 0;\n\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n int col = blockIdx.y * blockDim.y + threadIdx.y;\n\n if (row < N && col < N)\n {\n for ( int k = 0; k < N; ++k )\n val += a[row * N + k] * b[k * N + col];\n c[row * N + col] = val;\n }\n}\n\nvoid matrixMulCPU( int * a, int * b, int * c )\n{\n int val = 0;\n\n for( int row = 0; row < N; ++row )\n for( int col = 0; col < N; ++col )\n {\n val = 0;\n for ( int k = 0; k < N; ++k )\n val += a[row * N + k] * b[k * N + col];\n c[row * N + col] = val;\n }\n}\n\nint main()\n{\n int *a, *b, *c_cpu, *c_gpu;\n\n int size = N * N * sizeof (int); // Number of bytes of an N x N matrix\n\n // Allocate memory\n cudaMallocManaged (&a, size);\n cudaMallocManaged (&b, size);\n cudaMallocManaged (&c_cpu, size);\n cudaMallocManaged (&c_gpu, size);\n\n // Initialize memory\n for( int row = 0; row < N; ++row )\n for( int col = 0; col < N; ++col )\n {\n a[row*N + col] = row;\n b[row*N + col] = col+2;\n c_cpu[row*N + col] = 0;\n c_gpu[row*N + col] = 0;\n }\n\n dim3 threads_per_block (16, 16, 1); // A 16 x 16 block threads\n dim3 number_of_blocks ((N / threads_per_block.x) + 1, (N / threads_per_block.y) + 1, 1);\n\n matrixMulGPU <<< number_of_blocks, threads_per_block >>> ( a, b, c_gpu );\n\n cudaDeviceSynchronize(); // Wait for the GPU to finish before proceeding\n\n // Call the CPU version to check our work\n matrixMulCPU( a, b, c_cpu );\n\n // Compare the two answers to make sure they are equal\n bool error = false;\n for( int row = 0; row < N && !error; ++row )\n for( int col = 0; col < N && !error; ++col )\n if (c_cpu[row * N + col] != c_gpu[row * N + col])\n {\n printf(\"FOUND ERROR at c[%d][%d]\\n\", row, col);\n error = true;\n break;\n }\n if (!error)\n printf(\"Success!\\n\");\n\n // Free all our allocated memory\n cudaFree(a); cudaFree(b);\n cudaFree( c_cpu ); cudaFree( c_gpu );\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n\n#define N 64\n\n__global__ void matrixMulGPU( int * a, int * b, int * c )\n{\n int val = 0;\n\n int row = blockIdx.x * blockDim.x + threadIdx.x;\n int col = blockIdx.y * blockDim.y + threadIdx.y;\n\n if (row < N && col < N)\n {\n for ( int k = 0; k < N; ++k )\n val += a[row * N + k] * b[k * N + col];\n c[row * N + col] = val;\n }\n}\n\nvoid matrixMulCPU( int * a, int * b, int * c )\n{\n int val = 0;\n\n for( int row = 0; row < N; ++row )\n for( int col = 0; col < N; ++col )\n {\n val = 0;\n for ( int k = 0; k < N; ++k )\n val += a[row * N + k] * b[k * N + col];\n c[row * N + col] = val;\n }\n}\n\nint main()\n{\n int *a, *b, *c_cpu, *c_gpu;\n\n int size = N * N * sizeof (int); // Number of bytes of an N x N matrix\n\n // Allocate memory\n hipMallocManaged (&a, size);\n hipMallocManaged (&b, size);\n hipMallocManaged (&c_cpu, size);\n hipMallocManaged (&c_gpu, size);\n\n // Initialize memory\n for( int row = 0; row < N; ++row )\n for( int col = 0; col < N; ++col )\n {\n a[row*N + col] = row;\n b[row*N + col] = col+2;\n c_cpu[row*N + col] = 0;\n c_gpu[row*N + col] = 0;\n }\n\n dim3 threads_per_block (16, 16, 1); // A 16 x 16 block threads\n dim3 number_of_blocks ((N / threads_per_block.x) + 1, (N / threads_per_block.y) + 1, 1);\n\n matrixMulGPU <<< number_of_blocks, threads_per_block >>> ( a, b, c_gpu );\n\n hipDeviceSynchronize(); // Wait for the GPU to finish before proceeding\n\n // Call the CPU version to check our work\n matrixMulCPU( a, b, c_cpu );\n\n // Compare the two answers to make sure they are equal\n bool error = false;\n for( int row = 0; row < N && !error; ++row )\n for( int col = 0; col < N && !error; ++col )\n if (c_cpu[row * N + col] != c_gpu[row * N + col])\n {\n printf(\"FOUND ERROR at c[%d][%d]\\n\", row, col);\n error = true;\n break;\n }\n if (!error)\n printf(\"Success!\\n\");\n\n // Free all our allocated memory\n hipFree(a); hipFree(b);\n hipFree( c_cpu ); hipFree( c_gpu );\n}\n", "source": "/kaggle/working/GPU_Acceleration_Using_CUDA_C_CPP/01-matrix-multiply-2d-solution.cu"} {"input": "#include \n\n__global__ void loop()\n{\n /*\n * This idiomatic expression gives each thread\n * a unique index within the entire grid.\n */\n\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n printf(\"%d\\n\", i);\n}\n\nint main()\n{\n /*\n * Additional execution configurations that would\n * work and meet the exercises contraints are:\n *\n * <<<5, 2>>>\n * <<<10, 1>>>\n */\n\n loop<<<2, 5>>>();\n cudaDeviceSynchronize();\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n\n__global__ void loop()\n{\n /*\n * This idiomatic expression gives each thread\n * a unique index within the entire grid.\n */\n\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n printf(\"%d\\n\", i);\n}\n\nint main()\n{\n /*\n * Additional execution configurations that would\n * work and meet the exercises contraints are:\n *\n * <<<5, 2>>>\n * <<<10, 1>>>\n */\n\n loop<<<2, 5>>>();\n hipDeviceSynchronize();\n}\n", "source": "/kaggle/working/GPU_Acceleration_Using_CUDA_C_CPP/02-multi-block-loop-solution.cu"} {"input": "#include \n#include \n\ninline cudaError_t checkCuda(cudaError_t result)\n{\n if (result != cudaSuccess) {\n fprintf(stderr, \"CUDA Runtime Error: %s\\n\", cudaGetErrorString(result));\n assert(result == cudaSuccess);\n }\n return result;\n}\n\nvoid initWith(float num, float *a, int N)\n{\n for(int i = 0; i < N; ++i)\n {\n a[i] = num;\n }\n}\n\n__global__\nvoid addVectorsInto(float *result, float *a, float *b, int N)\n{\n int index = threadIdx.x + blockIdx.x * blockDim.x;\n int stride = blockDim.x * gridDim.x;\n\n for(int i = index; i < N; i += stride)\n {\n result[i] = a[i] + b[i];\n }\n}\n\nvoid checkElementsAre(float target, float *array, int N)\n{\n for(int i = 0; i < N; i++)\n {\n if(array[i] != target)\n {\n printf(\"FAIL: array[%d] - %0.0f does not equal %0.0f\\n\", i, array[i], target);\n exit(1);\n }\n }\n}\n\nint main()\n{\n const int N = 2<<20;\n size_t size = N * sizeof(float);\n\n float *a;\n float *b;\n float *c;\n\n checkCuda( cudaMallocManaged(&a, size) );\n checkCuda( cudaMallocManaged(&b, size) );\n checkCuda( cudaMallocManaged(&c, size) );\n\n initWith(3, a, N);\n initWith(4, b, N);\n initWith(0, c, N);\n\n size_t threadsPerBlock;\n size_t numberOfBlocks;\n\n threadsPerBlock = 256;\n numberOfBlocks = (N + threadsPerBlock - 1) / threadsPerBlock;\n\n addVectorsInto<<>>(c, a, b, N);\n\n checkCuda( cudaGetLastError() );\n checkCuda( cudaDeviceSynchronize() );\n\n checkElementsAre(7, c, N);\n\n checkCuda( cudaFree(a) );\n checkCuda( cudaFree(b) );\n checkCuda( cudaFree(c) );\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n\ninline hipError_t checkCuda(hipError_t result)\n{\n if (result != hipSuccess) {\n fprintf(stderr, \"CUDA Runtime Error: %s\\n\", hipGetErrorString(result));\n assert(result == hipSuccess);\n }\n return result;\n}\n\nvoid initWith(float num, float *a, int N)\n{\n for(int i = 0; i < N; ++i)\n {\n a[i] = num;\n }\n}\n\n__global__\nvoid addVectorsInto(float *result, float *a, float *b, int N)\n{\n int index = threadIdx.x + blockIdx.x * blockDim.x;\n int stride = blockDim.x * gridDim.x;\n\n for(int i = index; i < N; i += stride)\n {\n result[i] = a[i] + b[i];\n }\n}\n\nvoid checkElementsAre(float target, float *array, int N)\n{\n for(int i = 0; i < N; i++)\n {\n if(array[i] != target)\n {\n printf(\"FAIL: array[%d] - %0.0f does not equal %0.0f\\n\", i, array[i], target);\n exit(1);\n }\n }\n}\n\nint main()\n{\n const int N = 2<<20;\n size_t size = N * sizeof(float);\n\n float *a;\n float *b;\n float *c;\n\n checkCuda( hipMallocManaged(&a, size) );\n checkCuda( hipMallocManaged(&b, size) );\n checkCuda( hipMallocManaged(&c, size) );\n\n initWith(3, a, N);\n initWith(4, b, N);\n initWith(0, c, N);\n\n size_t threadsPerBlock;\n size_t numberOfBlocks;\n\n threadsPerBlock = 256;\n numberOfBlocks = (N + threadsPerBlock - 1) / threadsPerBlock;\n\n addVectorsInto<<>>(c, a, b, N);\n\n checkCuda( hipGetLastError() );\n checkCuda( hipDeviceSynchronize() );\n\n checkElementsAre(7, c, N);\n\n checkCuda( hipFree(a) );\n checkCuda( hipFree(b) );\n checkCuda( hipFree(c) );\n}\n", "source": "/kaggle/working/GPU_Acceleration_Using_CUDA_C_CPP/01-add-error-handling-solution (1).cu"} {"input": "#include \n\nvoid init(int *a, int N)\n{\n int i;\n for (i = 0; i < N; ++i)\n {\n a[i] = i;\n }\n}\n\n__global__\nvoid doubleElements(int *a, int N)\n{\n\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = gridDim.x * blockDim.x;\n\n /*\n * The previous code (now commented out) attempted\n * to access an element outside the range of `a`.\n */\n\n // for (int i = idx; i < N + stride; i += stride)\n for (int i = idx; i < N; i += stride)\n {\n a[i] *= 2;\n }\n}\n\nbool checkElementsAreDoubled(int *a, int N)\n{\n int i;\n for (i = 0; i < N; ++i)\n {\n if (a[i] != i*2) return false;\n }\n return true;\n}\n\nint main()\n{\n int N = 10000;\n int *a;\n\n size_t size = N * sizeof(int);\n cudaMallocManaged(&a, size);\n\n init(a, N);\n\n /*\n * The previous code (now commented out) attempted to launch\n * the kernel with more than the maximum number of threads per\n * block, which is 1024.\n */\n\n size_t threads_per_block = 1024;\n /* size_t threads_per_block = 2048; */\n size_t number_of_blocks = 32;\n\n cudaError_t syncErr, asyncErr;\n\n doubleElements<<>>(a, N);\n\n /*\n * Catch errors for both the kernel launch above and any\n * errors that occur during the asynchronous `doubleElements`\n * kernel execution.\n */\n\n syncErr = cudaGetLastError();\n asyncErr = cudaDeviceSynchronize();\n\n /*\n * Print errors should they exist.\n */\n\n if (syncErr != cudaSuccess) printf(\"Error: %s\\n\", cudaGetErrorString(syncErr));\n if (asyncErr != cudaSuccess) printf(\"Error: %s\\n\", cudaGetErrorString(asyncErr));\n\n bool areDoubled = checkElementsAreDoubled(a, N);\n printf(\"All elements were doubled? %s\\n\", areDoubled ? \"TRUE\" : \"FALSE\");\n\n cudaFree(a);\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n\nvoid init(int *a, int N)\n{\n int i;\n for (i = 0; i < N; ++i)\n {\n a[i] = i;\n }\n}\n\n__global__\nvoid doubleElements(int *a, int N)\n{\n\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = gridDim.x * blockDim.x;\n\n /*\n * The previous code (now commented out) attempted\n * to access an element outside the range of `a`.\n */\n\n // for (int i = idx; i < N + stride; i += stride)\n for (int i = idx; i < N; i += stride)\n {\n a[i] *= 2;\n }\n}\n\nbool checkElementsAreDoubled(int *a, int N)\n{\n int i;\n for (i = 0; i < N; ++i)\n {\n if (a[i] != i*2) return false;\n }\n return true;\n}\n\nint main()\n{\n int N = 10000;\n int *a;\n\n size_t size = N * sizeof(int);\n hipMallocManaged(&a, size);\n\n init(a, N);\n\n /*\n * The previous code (now commented out) attempted to launch\n * the kernel with more than the maximum number of threads per\n * block, which is 1024.\n */\n\n size_t threads_per_block = 1024;\n /* size_t threads_per_block = 2048; */\n size_t number_of_blocks = 32;\n\n hipError_t syncErr, asyncErr;\n\n doubleElements<<>>(a, N);\n\n /*\n * Catch errors for both the kernel launch above and any\n * errors that occur during the asynchronous `doubleElements`\n * kernel execution.\n */\n\n syncErr = hipGetLastError();\n asyncErr = hipDeviceSynchronize();\n\n /*\n * Print errors should they exist.\n */\n\n if (syncErr != hipSuccess) printf(\"Error: %s\\n\", hipGetErrorString(syncErr));\n if (asyncErr != hipSuccess) printf(\"Error: %s\\n\", hipGetErrorString(asyncErr));\n\n bool areDoubled = checkElementsAreDoubled(a, N);\n printf(\"All elements were doubled? %s\\n\", areDoubled ? \"TRUE\" : \"FALSE\");\n\n hipFree(a);\n}\n", "source": "/kaggle/working/GPU_Acceleration_Using_CUDA_C_CPP/01-add-error-handling-solution.cu"} {"input": "#include \n\n/*\n * Refactor firstParallel so that it can run on the GPU.\n */\n\n__global__ void firstParallel()\n{\n printf(\"This should be running in parallel.\\n\");\n}\n\nint main()\n{\n /*\n * Refactor this call to firstParallel to execute in parallel\n * on the GPU.\n */\n\n firstParallel<<<5,5>>>();\n\n /*\n * Some code is needed below so that the CPU will wait\n * for the GPU kernels to complete before proceeding.\n */\n cudaDeviceSynchronize();\n\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n\n/*\n * Refactor firstParallel so that it can run on the GPU.\n */\n\n__global__ void firstParallel()\n{\n printf(\"This should be running in parallel.\\n\");\n}\n\nint main()\n{\n /*\n * Refactor this call to firstParallel to execute in parallel\n * on the GPU.\n */\n\n firstParallel<<<5,5>>>();\n\n /*\n * Some code is needed below so that the CPU will wait\n * for the GPU kernels to complete before proceeding.\n */\n hipDeviceSynchronize();\n\n}\n", "source": "/kaggle/working/GPU_Acceleration_Using_CUDA_C_CPP/01-basic-parallel.cu"} {"input": "#include \n#include \n\ninline cudaError_t checkCuda(cudaError_t result)\n{\n if (result != cudaSuccess) {\n fprintf(stderr, \"CUDA Runtime Error: %s\\n\", cudaGetErrorString(result));\n assert(result == cudaSuccess);\n }\n return result;\n}\n\nvoid initWith(float num, float *a, int N)\n{\n for(int i = 0; i < N; ++i)\n {\n a[i] = num;\n }\n}\n\n__global__\nvoid addVectorsInto(float *result, float *a, float *b, int N)\n{\n int index = threadIdx.x + blockIdx.x * blockDim.x;\n int stride = blockDim.x * gridDim.x;\n\n for(int i = index; i < N; i += stride)\n {\n result[i] = a[i] + b[i];\n }\n}\n\nvoid checkElementsAre(float target, float *array, int N)\n{\n for(int i = 0; i < N; i++)\n {\n if(array[i] != target)\n {\n printf(\"FAIL: array[%d] - %0.0f does not equal %0.0f\\n\", i, array[i], target);\n exit(1);\n }\n }\n}\n\nint main()\n{\n const int N = 2<<20;\n size_t size = N * sizeof(float);\n\n float *a;\n float *b;\n float *c;\n\n checkCuda( cudaMallocManaged(&a, size) );\n checkCuda( cudaMallocManaged(&b, size) );\n checkCuda( cudaMallocManaged(&c, size) );\n\n initWith(3, a, N);\n initWith(4, b, N);\n initWith(0, c, N);\n\n size_t threadsPerBlock;\n size_t numberOfBlocks;\n\n threadsPerBlock = 256;\n numberOfBlocks = (N + threadsPerBlock - 1) / threadsPerBlock;\n\n addVectorsInto<<>>(c, a, b, N);\n\n checkCuda( cudaGetLastError() );\n checkCuda( cudaDeviceSynchronize() );\n\n checkElementsAre(7, c, N);\n\n checkCuda( cudaFree(a) );\n checkCuda( cudaFree(b) );\n checkCuda( cudaFree(c) );\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n\ninline hipError_t checkCuda(hipError_t result)\n{\n if (result != hipSuccess) {\n fprintf(stderr, \"CUDA Runtime Error: %s\\n\", hipGetErrorString(result));\n assert(result == hipSuccess);\n }\n return result;\n}\n\nvoid initWith(float num, float *a, int N)\n{\n for(int i = 0; i < N; ++i)\n {\n a[i] = num;\n }\n}\n\n__global__\nvoid addVectorsInto(float *result, float *a, float *b, int N)\n{\n int index = threadIdx.x + blockIdx.x * blockDim.x;\n int stride = blockDim.x * gridDim.x;\n\n for(int i = index; i < N; i += stride)\n {\n result[i] = a[i] + b[i];\n }\n}\n\nvoid checkElementsAre(float target, float *array, int N)\n{\n for(int i = 0; i < N; i++)\n {\n if(array[i] != target)\n {\n printf(\"FAIL: array[%d] - %0.0f does not equal %0.0f\\n\", i, array[i], target);\n exit(1);\n }\n }\n}\n\nint main()\n{\n const int N = 2<<20;\n size_t size = N * sizeof(float);\n\n float *a;\n float *b;\n float *c;\n\n checkCuda( hipMallocManaged(&a, size) );\n checkCuda( hipMallocManaged(&b, size) );\n checkCuda( hipMallocManaged(&c, size) );\n\n initWith(3, a, N);\n initWith(4, b, N);\n initWith(0, c, N);\n\n size_t threadsPerBlock;\n size_t numberOfBlocks;\n\n threadsPerBlock = 256;\n numberOfBlocks = (N + threadsPerBlock - 1) / threadsPerBlock;\n\n addVectorsInto<<>>(c, a, b, N);\n\n checkCuda( hipGetLastError() );\n checkCuda( hipDeviceSynchronize() );\n\n checkElementsAre(7, c, N);\n\n checkCuda( hipFree(a) );\n checkCuda( hipFree(b) );\n checkCuda( hipFree(c) );\n}\n", "source": "/kaggle/working/GPU_Acceleration_Using_CUDA_C_CPP/01-vector-add-solution.cu"} {"input": "#include \n\n\n__global__ void initializeElementsTo(int initialValue, int *a, int N)\n{\n int i = threadIdx.x + blockIdx.x * blockDim.x;\n if (i < N)\n {\n a[i] = initialValue;\n }\n}\n\nint main()\n{\n /*\n * Do not modify `N`.\n */\n\n int N = 1000;\n\n int *a;\n size_t size = N * sizeof(int);\n\n cudaMallocManaged(&a, size);\n\n /*\n * Assume we have reason to want the number of threads\n * fixed at `256`: do not modify `threads_per_block`.\n */\n\n size_t threads_per_block = 256;\n\n /*\n * The following is idiomatic CUDA to make sure there are at\n * least as many threads in the grid as there are `N` elements.\n */\n\n size_t number_of_blocks = (N + threads_per_block - 1) / threads_per_block;\n\n int initialValue = 6;\n\n initializeElementsTo<<>>(initialValue, a, N);\n cudaDeviceSynchronize();\n\n /*\n * Check to make sure all values in `a`, were initialized.\n */\n\n for (int i = 0; i < N; ++i)\n {\n if(a[i] != initialValue)\n {\n printf(\"FAILURE: target value: %d\\t a[%d]: %d\\n\", initialValue, i, a[i]);\n exit(1);\n }\n }\n printf(\"SUCCESS!\\n\");\n\n cudaFree(a);\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n\n\n__global__ void initializeElementsTo(int initialValue, int *a, int N)\n{\n int i = threadIdx.x + blockIdx.x * blockDim.x;\n if (i < N)\n {\n a[i] = initialValue;\n }\n}\n\nint main()\n{\n /*\n * Do not modify `N`.\n */\n\n int N = 1000;\n\n int *a;\n size_t size = N * sizeof(int);\n\n hipMallocManaged(&a, size);\n\n /*\n * Assume we have reason to want the number of threads\n * fixed at `256`: do not modify `threads_per_block`.\n */\n\n size_t threads_per_block = 256;\n\n /*\n * The following is idiomatic CUDA to make sure there are at\n * least as many threads in the grid as there are `N` elements.\n */\n\n size_t number_of_blocks = (N + threads_per_block - 1) / threads_per_block;\n\n int initialValue = 6;\n\n initializeElementsTo<<>>(initialValue, a, N);\n hipDeviceSynchronize();\n\n /*\n * Check to make sure all values in `a`, were initialized.\n */\n\n for (int i = 0; i < N; ++i)\n {\n if(a[i] != initialValue)\n {\n printf(\"FAILURE: target value: %d\\t a[%d]: %d\\n\", initialValue, i, a[i]);\n exit(1);\n }\n }\n printf(\"SUCCESS!\\n\");\n\n hipFree(a);\n}\n", "source": "/kaggle/working/GPU_Acceleration_Using_CUDA_C_CPP/02-mismatched-config-loop-solution.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n\n#include \n\n#include \"./common.h\"\n#include \"./utils.cuh\"\n#include \"common/util/cuda_runtime.h\"\n#include \"common/util/logging.h\"\n\nnamespace transformer_engine {\n\nnamespace {\n\n__global__ void __launch_bounds__(1)\n update_tensor_scale_inv_kernel(const float *__restrict__ scale_ptr,\n float *__restrict__ scale_inv_ptr) {\n const float scale = scale_ptr == nullptr ? 1 : *scale_ptr;\n reciprocal(scale_inv_ptr, scale);\n}\n\n} // namespace\n\ncudaDataType_t get_cuda_dtype(const transformer_engine::DType t) {\n using namespace transformer_engine;\n switch (t) {\n case DType::kFloat16:\n return CUDA_R_16F;\n case DType::kFloat32:\n return CUDA_R_32F;\n case DType::kBFloat16:\n return CUDA_R_16BF;\n case DType::kFloat8E4M3:\n return CUDA_R_8F_E4M3;\n case DType::kFloat8E5M2:\n return CUDA_R_8F_E5M2;\n#if CUDA_VERSION >= 12080\n case DType::kFloat4E2M1:\n return CUDA_R_4F_E2M1;\n#endif\n default:\n NVTE_ERROR(\"Invalid type\");\n }\n}\n\nvoid update_tensor_scale_inv(Tensor *t, cudaStream_t stream) {\n if (is_fp8_dtype(t->data.dtype) && is_tensor_scaling(t->scaling_mode)) {\n NVTE_CHECK(t->scale_inv.dptr != nullptr, \"Tensor should have allocated scale_inv.\");\n update_tensor_scale_inv_kernel<<<1, 1, 0, stream>>>(\n reinterpret_cast(t->scale.dptr),\n reinterpret_cast(t->scale_inv.dptr));\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\nnamespace {\n\nconstexpr size_t kThreadsPerBlock = 256;\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock)\n memset_kernel(void *__restrict__ ptr, int value, size_t size_in_bytes) {\n size_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx * sizeof(TVectorized) >= size_in_bytes) {\n return; // Out of bounds\n }\n\n if ((idx + 1) * sizeof(TVectorized) > size_in_bytes) {\n // If the buffer size is not an even multiple of the vectorization, manually set the remaining bytes unvectorized.\n size_t remaining_bytes = size_in_bytes - idx * sizeof(TVectorized);\n memset(reinterpret_cast(ptr) + idx * sizeof(TVectorized), value, remaining_bytes);\n return;\n }\n\n union {\n TVectorized value;\n uint8_t data[sizeof(TVectorized)];\n } data;\n for (size_t i = 0; i < sizeof(TVectorized); ++i) {\n data.data[i] = static_cast(value);\n }\n reinterpret_cast(ptr)[idx] = data.value;\n}\n\n__global__ void __launch_bounds__(kThreadsPerBlock)\n splits_to_offsets_kernel(const int64_t *__restrict__ first_dims, int64_t *__restrict__ output,\n size_t num_tensors, int64_t logical_last_dim) {\n __shared__ int64_t block_scan[kThreadsPerBlock];\n __shared__ int64_t chunk_prefix;\n\n const size_t tid = threadIdx.x;\n if (tid == 0) {\n output[0] = 0;\n chunk_prefix = 0;\n }\n __syncthreads();\n\n for (size_t chunk_start = 0; chunk_start < num_tensors; chunk_start += kThreadsPerBlock) {\n const size_t idx = chunk_start + tid;\n int64_t value = 0;\n if (idx < num_tensors) {\n value = first_dims[idx] * logical_last_dim;\n }\n block_scan[tid] = value;\n __syncthreads();\n\n // Inclusive scan in shared memory.\n for (size_t offset = 1; offset < kThreadsPerBlock; offset <<= 1) {\n const int64_t addend = (tid >= offset) ? block_scan[tid - offset] : 0;\n __syncthreads();\n block_scan[tid] += addend;\n __syncthreads();\n }\n\n if (idx < num_tensors) {\n output[idx + 1] = chunk_prefix + block_scan[tid];\n }\n __syncthreads();\n\n if (tid == kThreadsPerBlock - 1) {\n chunk_prefix += block_scan[tid];\n }\n __syncthreads();\n }\n}\n\n} // namespace\n\n#define MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, vectorizedType, stream) \\\n if (size_in_bytes >= sizeof(vectorizedType) && \\\n reinterpret_cast(ptr) % sizeof(vectorizedType) == 0) { \\\n size_t numBlocks = DIVUP(size_in_bytes, kThreadsPerBlock * sizeof(vectorizedType)); \\\n dim3 grid(numBlocks, 1, 1); \\\n memset_kernel \\\n <<>>(ptr, value, size_in_bytes); \\\n NVTE_CHECK_CUDA(cudaGetLastError()); \\\n return; \\\n }\n\nextern \"C\" {\nvoid nvte_memset(void *ptr, int value, size_t size_in_bytes, cudaStream_t stream) {\n NVTE_API_CALL(nvte_memset);\n NVTE_CHECK(ptr != nullptr, \"Pointer for memset must be allocated.\");\n\n if (size_in_bytes > 4096) {\n // Use cudaMemsetAsync for larger sizes.\n NVTE_CHECK_CUDA(cudaMemsetAsync(ptr, value, size_in_bytes, stream));\n return;\n }\n\n MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, float4, stream);\n MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, float2, stream);\n MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, float, stream);\n MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, uint8_t, stream);\n}\n\nvoid nvte_splits_to_offsets(const int64_t *first_dims, int64_t *output, size_t num_tensors,\n int64_t logical_last_dim, cudaStream_t stream) {\n NVTE_API_CALL(nvte_splits_to_offsets);\n NVTE_CHECK(output != nullptr, \"Output pointer must be allocated.\");\n NVTE_CHECK(num_tensors > 0, \"num_tensors must be greater than 0.\");\n NVTE_CHECK(first_dims != nullptr, \"first_dims pointer must be allocated.\");\n NVTE_CHECK(logical_last_dim > 0, \"logical_last_dim must be greater than 0.\");\n\n splits_to_offsets_kernel<<<1, kThreadsPerBlock, 0, stream>>>(first_dims, output, num_tensors,\n logical_last_dim);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n} // extern \"C\"\n\nvoid checkCuDriverContext(CUstream stream) {\n // Ensure the thread's \"current\" CUDA context is set.\n cuda_driver::ensure_context_exists();\n\n CUcontext ctx;\n const CUresult driver_status = cuda_driver::call(\"cuStreamGetCtx\", stream, &ctx);\n switch (driver_status) {\n case CUDA_SUCCESS:\n break;\n\n case CUDA_ERROR_INVALID_CONTEXT:\n int current_device;\n NVTE_CHECK_CUDA(cudaGetDevice(¤t_device));\n NVTE_CALL_CHECK_CUDA_DRIVER(cuDevicePrimaryCtxRetain, &ctx, current_device);\n NVTE_CALL_CHECK_CUDA_DRIVER(cuCtxSetCurrent, ctx);\n break;\n\n default:\n const char *desc_NVTE_CHECK_CUDA_DRIVER;\n cuda_driver::call(\"cuGetErrorString\", driver_status, &desc_NVTE_CHECK_CUDA_DRIVER);\n NVTE_ERROR(\"CUDA Error: \", desc_NVTE_CHECK_CUDA_DRIVER);\n }\n}\n\nCUtensorMapDataType get_CUtensorMapDataType(DType dtype) {\n static const std::unordered_map dtypeMapping = []() {\n std::unordered_map typeMapping = {\n {DType::kByte, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT8},\n {DType::kFloat32, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_FLOAT32},\n {DType::kFloat16, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_FLOAT16},\n {DType::kBFloat16, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16},\n {DType::kFloat8E4M3, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT8},\n {DType::kFloat8E5M2, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT8}};\n#if FP4_TYPE_SUPPORTED\n typeMapping.insert(\n {DType::kFloat4E2M1, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B});\n#endif\n return typeMapping;\n }();\n return dtypeMapping.at(dtype);\n}\n\n// Set up parameters to create TMA descriptor.\nvoid create_2D_tensor_map(CUtensorMap &tensorMap, const SimpleTensor &tensor,\n const uint64_t globalY, const uint64_t globalX, const uint32_t shmemY,\n const uint32_t shmemX, const uint32_t stride_elems,\n const uint32_t offset_elems, const size_t type_num_bits,\n const CUtensorMapSwizzle swizzle) {\n cuda_driver::ensure_context_exists();\n // Get a function pointer to the cuTensorMapEncodeTiled driver API\n // Note: PFN_cuTensorMapEncodeTiled is not defined in cuda13\n static PFN_cuTensorMapEncodeTiled_v12000 cuDriverTensorMapEncodeTiled = []() {\n void *driver_ptr = cuda_driver::get_symbol(\"cuTensorMapEncodeTiled\");\n return reinterpret_cast(driver_ptr);\n }();\n // rank is the number of dimensions of the array\n constexpr uint32_t rank = 2;\n\n // Dimension for the packed data types must reflect the number of individual U# values.\n uint64_t size[rank] = {globalX, globalY};\n\n // The stride is the number of bytes to traverse from the first element of one row to the next\n uint64_t stride[rank - 1] = {(stride_elems * type_num_bits) / 8};\n\n // The boxSize is the size of the shared memory buffer that is used as the\n // source/destination of a TMA transfer\n uint32_t boxSize[rank] = {shmemX, shmemY};\n\n // The distance between elements in units of sizeof(element)\n uint32_t elemStride[rank] = {1, 1};\n\n const CUtensorMapDataType tensorDataType = get_CUtensorMapDataType(tensor.dtype);\n void *dataPtr = reinterpret_cast(reinterpret_cast(tensor.dptr) +\n (offset_elems * type_num_bits) / 8);\n\n NVTE_CHECK(is_aligned_ptr(dataPtr, TMA_GMEM_ALIGNMENT),\n \"Tensor data pointer must be 16B aligned\");\n\n const int TMA_needed_size = (TMA_GMEM_ALIGNMENT * 8) / type_num_bits;\n NVTE_CHECK(globalX % TMA_needed_size == 0, \"Shape not supported. For \", type_num_bits,\n \"-bit data type, expected multiple of \", TMA_needed_size, \", got \", globalX);\n\n // Create the tensor descriptor.\n NVTE_CHECK_CUDA_DRIVER(cuDriverTensorMapEncodeTiled(\n &tensorMap, // CUtensorMap *tensorMap,\n tensorDataType,\n rank, // cuuint32_t tensorRank,\n dataPtr, // void *globalAddress,\n size, // const cuuint64_t *globalDim,\n stride, // const cuuint64_t *globalStrides,\n boxSize, // const cuuint32_t *boxDim,\n elemStride, // const cuuint32_t *elementStrides,\n // Interleave patterns can be used to accelerate loading of values that\n // are less than 4 bytes long.\n CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,\n\n // Swizzling can be used to avoid shared memory bank conflicts.\n swizzle,\n\n // L2 Promotion can be used to widen the effect of a cache-policy to a wider\n // set of L2 cache lines.\n CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,\n // CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_256B,\n\n // Any element that is outside of bounds will be set to zero by the TMA transfer.\n CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));\n}\n\nbool is_supported_by_CC_100() {\n int deviceComputeCapability = cuda::sm_arch(cuda::current_device());\n\n return deviceComputeCapability >= 100;\n}\n\nstd::vector> convert_tensor_array(NVTETensor **nvte_tensors,\n size_t outer_size, size_t inner_size) {\n std::vector> ret;\n for (size_t i = 0; i < outer_size; ++i) {\n ret.emplace_back();\n for (size_t j = 0; j < inner_size; ++j) {\n ret.back().push_back(convertNVTETensor(nvte_tensors[i][j]));\n }\n }\n return ret;\n}\n\nsize_t get_buffer_size_bytes(const size_t elements_num, const DType buffer_dtype) {\n return (elements_num * typeToNumBits(buffer_dtype)) / 8;\n}\n\nsize_t get_buffer_size_bytes(const size_t dim_first, const size_t dim_last,\n const DType buffer_dtype) {\n if (buffer_dtype == DType::kFloat4E2M1) {\n NVTE_CHECK(dim_last % 2 == 0,\n \"Last dimension of a tensor with FP4 type of data must be an even number!\");\n }\n const size_t elements_num = dim_first * dim_last;\n return get_buffer_size_bytes(elements_num, buffer_dtype);\n}\n\n} // namespace transformer_engine\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \n\n#include \"./common.h\"\n#include \"./utils.cuh\"\n#include \"common/util/hip/hip_runtime.h\"\n#include \"common/util/logging.h\"\n\nnamespace transformer_engine {\n\nnamespace {\n\n__global__ void __launch_bounds__(1)\n update_tensor_scale_inv_kernel(const float *__restrict__ scale_ptr,\n float *__restrict__ scale_inv_ptr) {\n const float scale = scale_ptr == nullptr ? 1 : *scale_ptr;\n reciprocal(scale_inv_ptr, scale);\n}\n\n} // namespace\n\nhipDataType get_cuda_dtype(const transformer_engine::DType t) {\n using namespace transformer_engine;\n switch (t) {\n case DType::kFloat16:\n return HIP_R_16F;\n case DType::kFloat32:\n return HIP_R_32F;\n case DType::kBFloat16:\n return HIP_R_16BF;\n case DType::kFloat8E4M3:\n return HIP_R_8F_E4M3;\n case DType::kFloat8E5M2:\n return HIP_R_8F_E5M2;\n#if CUDA_VERSION >= 12080\n case DType::kFloat4E2M1:\n return HIP_R_4F_E2M1;\n#endif\n default:\n NVTE_ERROR(\"Invalid type\");\n }\n}\n\nvoid update_tensor_scale_inv(Tensor *t, hipStream_t stream) {\n if (is_fp8_dtype(t->data.dtype) && is_tensor_scaling(t->scaling_mode)) {\n NVTE_CHECK(t->scale_inv.dptr != nullptr, \"Tensor should have allocated scale_inv.\");\n update_tensor_scale_inv_kernel<<<1, 1, 0, stream>>>(\n reinterpret_cast(t->scale.dptr),\n reinterpret_cast(t->scale_inv.dptr));\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\nnamespace {\n\nconstexpr size_t kThreadsPerBlock = 256;\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock)\n memset_kernel(void *__restrict__ ptr, int value, size_t size_in_bytes) {\n size_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (idx * sizeof(TVectorized) >= size_in_bytes) {\n return; // Out of bounds\n }\n\n if ((idx + 1) * sizeof(TVectorized) > size_in_bytes) {\n // If the buffer size is not an even multiple of the vectorization, manually set the remaining bytes unvectorized.\n size_t remaining_bytes = size_in_bytes - idx * sizeof(TVectorized);\n memset(reinterpret_cast(ptr) + idx * sizeof(TVectorized), value, remaining_bytes);\n return;\n }\n\n union {\n TVectorized value;\n uint8_t data[sizeof(TVectorized)];\n } data;\n for (size_t i = 0; i < sizeof(TVectorized); ++i) {\n data.data[i] = static_cast(value);\n }\n reinterpret_cast(ptr)[idx] = data.value;\n}\n\n__global__ void __launch_bounds__(kThreadsPerBlock)\n splits_to_offsets_kernel(const int64_t *__restrict__ first_dims, int64_t *__restrict__ output,\n size_t num_tensors, int64_t logical_last_dim) {\n __shared__ int64_t block_scan[kThreadsPerBlock];\n __shared__ int64_t chunk_prefix;\n\n const size_t tid = threadIdx.x;\n if (tid == 0) {\n output[0] = 0;\n chunk_prefix = 0;\n }\n __syncthreads();\n\n for (size_t chunk_start = 0; chunk_start < num_tensors; chunk_start += kThreadsPerBlock) {\n const size_t idx = chunk_start + tid;\n int64_t value = 0;\n if (idx < num_tensors) {\n value = first_dims[idx] * logical_last_dim;\n }\n block_scan[tid] = value;\n __syncthreads();\n\n // Inclusive scan in shared memory.\n for (size_t offset = 1; offset < kThreadsPerBlock; offset <<= 1) {\n const int64_t addend = (tid >= offset) ? block_scan[tid - offset] : 0;\n __syncthreads();\n block_scan[tid] += addend;\n __syncthreads();\n }\n\n if (idx < num_tensors) {\n output[idx + 1] = chunk_prefix + block_scan[tid];\n }\n __syncthreads();\n\n if (tid == kThreadsPerBlock - 1) {\n chunk_prefix += block_scan[tid];\n }\n __syncthreads();\n }\n}\n\n} // namespace\n\n#define MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, vectorizedType, stream) \\\n if (size_in_bytes >= sizeof(vectorizedType) && \\\n reinterpret_cast(ptr) % sizeof(vectorizedType) == 0) { \\\n size_t numBlocks = DIVUP(size_in_bytes, kThreadsPerBlock * sizeof(vectorizedType)); \\\n dim3 grid(numBlocks, 1, 1); \\\n memset_kernel \\\n <<>>(ptr, value, size_in_bytes); \\\n NVTE_CHECK_CUDA(hipGetLastError()); \\\n return; \\\n }\n\nextern \"C\" {\nvoid nvte_memset(void *ptr, int value, size_t size_in_bytes, hipStream_t stream) {\n NVTE_API_CALL(nvte_memset);\n NVTE_CHECK(ptr != nullptr, \"Pointer for memset must be allocated.\");\n\n if (size_in_bytes > 4096) {\n // Use cudaMemsetAsync for larger sizes.\n NVTE_CHECK_CUDA(hipMemsetAsync(ptr, value, size_in_bytes, stream));\n return;\n }\n\n MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, float4, stream);\n MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, float2, stream);\n MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, float, stream);\n MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, uint8_t, stream);\n}\n\nvoid nvte_splits_to_offsets(const int64_t *first_dims, int64_t *output, size_t num_tensors,\n int64_t logical_last_dim, hipStream_t stream) {\n NVTE_API_CALL(nvte_splits_to_offsets);\n NVTE_CHECK(output != nullptr, \"Output pointer must be allocated.\");\n NVTE_CHECK(num_tensors > 0, \"num_tensors must be greater than 0.\");\n NVTE_CHECK(first_dims != nullptr, \"first_dims pointer must be allocated.\");\n NVTE_CHECK(logical_last_dim > 0, \"logical_last_dim must be greater than 0.\");\n\n splits_to_offsets_kernel<<<1, kThreadsPerBlock, 0, stream>>>(first_dims, output, num_tensors,\n logical_last_dim);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n} // extern \"C\"\n\nvoid checkCuDriverContext(hipStream_t stream) {\n // Ensure the thread's \"current\" CUDA context is set.\n cuda_driver::ensure_context_exists();\n\n hipCtx_t ctx;\n const hipError_t driver_status = cuda_driver::call(\"cuStreamGetCtx\", stream, &ctx);\n switch (driver_status) {\n case hipSuccess:\n break;\n\n case hipErrorInvalidContext:\n int current_device;\n NVTE_CHECK_CUDA(hipGetDevice(¤t_device));\n NVTE_CALL_CHECK_CUDA_DRIVER(hipDevicePrimaryCtxRetain, &ctx, current_device);\n NVTE_CALL_CHECK_CUDA_DRIVER(hipCtxSetCurrent, ctx);\n break;\n\n default:\n const char *desc_NVTE_CHECK_CUDA_DRIVER;\n cuda_driver::call(\"hipDrvGetErrorString\", driver_status, &desc_NVTE_CHECK_CUDA_DRIVER);\n NVTE_ERROR(\"CUDA Error: \", desc_NVTE_CHECK_CUDA_DRIVER);\n }\n}\n\nCUtensorMapDataType get_CUtensorMapDataType(DType dtype) {\n static const std::unordered_map dtypeMapping = []() {\n std::unordered_map typeMapping = {\n {DType::kByte, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT8},\n {DType::kFloat32, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_FLOAT32},\n {DType::kFloat16, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_FLOAT16},\n {DType::kBFloat16, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16},\n {DType::kFloat8E4M3, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT8},\n {DType::kFloat8E5M2, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT8}};\n#if FP4_TYPE_SUPPORTED\n typeMapping.insert(\n {DType::kFloat4E2M1, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B});\n#endif\n return typeMapping;\n }();\n return dtypeMapping.at(dtype);\n}\n\n// Set up parameters to create TMA descriptor.\nvoid create_2D_tensor_map(CUtensorMap &tensorMap, const SimpleTensor &tensor,\n const uint64_t globalY, const uint64_t globalX, const uint32_t shmemY,\n const uint32_t shmemX, const uint32_t stride_elems,\n const uint32_t offset_elems, const size_t type_num_bits,\n const CUtensorMapSwizzle swizzle) {\n cuda_driver::ensure_context_exists();\n // Get a function pointer to the cuTensorMapEncodeTiled driver API\n // Note: PFN_cuTensorMapEncodeTiled is not defined in cuda13\n static PFN_cuTensorMapEncodeTiled_v12000 cuDriverTensorMapEncodeTiled = []() {\n void *driver_ptr = cuda_driver::get_symbol(\"cuTensorMapEncodeTiled\");\n return reinterpret_cast(driver_ptr);\n }();\n // rank is the number of dimensions of the array\n constexpr uint32_t rank = 2;\n\n // Dimension for the packed data types must reflect the number of individual U# values.\n uint64_t size[rank] = {globalX, globalY};\n\n // The stride is the number of bytes to traverse from the first element of one row to the next\n uint64_t stride[rank - 1] = {(stride_elems * type_num_bits) / 8};\n\n // The boxSize is the size of the shared memory buffer that is used as the\n // source/destination of a TMA transfer\n uint32_t boxSize[rank] = {shmemX, shmemY};\n\n // The distance between elements in units of sizeof(element)\n uint32_t elemStride[rank] = {1, 1};\n\n const CUtensorMapDataType tensorDataType = get_CUtensorMapDataType(tensor.dtype);\n void *dataPtr = reinterpret_cast(reinterpret_cast(tensor.dptr) +\n (offset_elems * type_num_bits) / 8);\n\n NVTE_CHECK(is_aligned_ptr(dataPtr, TMA_GMEM_ALIGNMENT),\n \"Tensor data pointer must be 16B aligned\");\n\n const int TMA_needed_size = (TMA_GMEM_ALIGNMENT * 8) / type_num_bits;\n NVTE_CHECK(globalX % TMA_needed_size == 0, \"Shape not supported. For \", type_num_bits,\n \"-bit data type, expected multiple of \", TMA_needed_size, \", got \", globalX);\n\n // Create the tensor descriptor.\n NVTE_CHECK_CUDA_DRIVER(cuDriverTensorMapEncodeTiled(\n &tensorMap, // CUtensorMap *tensorMap,\n tensorDataType,\n rank, // cuuint32_t tensorRank,\n dataPtr, // void *globalAddress,\n size, // const cuuint64_t *globalDim,\n stride, // const cuuint64_t *globalStrides,\n boxSize, // const cuuint32_t *boxDim,\n elemStride, // const cuuint32_t *elementStrides,\n // Interleave patterns can be used to accelerate loading of values that\n // are less than 4 bytes long.\n CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,\n\n // Swizzling can be used to avoid shared memory bank conflicts.\n swizzle,\n\n // L2 Promotion can be used to widen the effect of a cache-policy to a wider\n // set of L2 cache lines.\n CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,\n // CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_256B,\n\n // Any element that is outside of bounds will be set to zero by the TMA transfer.\n CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));\n}\n\nbool is_supported_by_CC_100() {\n int deviceComputeCapability = cuda::sm_arch(cuda::current_device());\n\n return deviceComputeCapability >= 100;\n}\n\nstd::vector> convert_tensor_array(NVTETensor **nvte_tensors,\n size_t outer_size, size_t inner_size) {\n std::vector> ret;\n for (size_t i = 0; i < outer_size; ++i) {\n ret.emplace_back();\n for (size_t j = 0; j < inner_size; ++j) {\n ret.back().push_back(convertNVTETensor(nvte_tensors[i][j]));\n }\n }\n return ret;\n}\n\nsize_t get_buffer_size_bytes(const size_t elements_num, const DType buffer_dtype) {\n return (elements_num * typeToNumBits(buffer_dtype)) / 8;\n}\n\nsize_t get_buffer_size_bytes(const size_t dim_first, const size_t dim_last,\n const DType buffer_dtype) {\n if (buffer_dtype == DType::kFloat4E2M1) {\n NVTE_CHECK(dim_last % 2 == 0,\n \"Last dimension of a tensor with FP4 type of data must be an even number!\");\n }\n const size_t elements_num = dim_first * dim_last;\n return get_buffer_size_bytes(elements_num, buffer_dtype);\n}\n\n} // namespace transformer_engine\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/common.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"../common.h\"\n#include \"transformer_engine/fused_attn.h\"\n\nnamespace transformer_engine {\nnamespace kv_cache {\n\nconstexpr int block_size = 1024;\n\ntemplate \n__global__ void reindex_kv_cache_kernel(dtype *k_cache, dtype *v_cache, int *batch_indices,\n int *cu_new_lens, int *cu_cached_lens, int h_kv, int d_k,\n int d_v, int b, int max_seq_len) {\n // k_cache, v_cache: bshd\n // batch_indices: [b]; cu_new_lens, cu_cached_lens: [b + 1]\n int actual_b = b;\n for (int i = 0; i < b - 1; i++) {\n if (batch_indices[i + 1] < batch_indices[i]) {\n actual_b = i + 1;\n }\n }\n bool flag = (batch_indices[0] != 0);\n for (int batch_idx = 0; batch_idx < actual_b; batch_idx++) {\n if (flag || ((batch_indices[batch_idx] - batch_indices[0]) != batch_idx)) {\n int num_tokens = (cu_cached_lens[batch_idx + 1] - cu_cached_lens[batch_idx]) -\n (cu_new_lens[batch_idx + 1] - cu_new_lens[batch_idx]);\n int num_elts_k = h_kv * d_k;\n int num_elts_v = h_kv * d_v;\n int num_elts = max(num_elts_k, num_elts_v);\n for (int token_idx = blockIdx.x; token_idx < num_tokens; token_idx += gridDim.x) {\n int src_offset = batch_indices[batch_idx] * max_seq_len + token_idx;\n int des_offset = batch_idx * max_seq_len + token_idx;\n dtype *k_cache_src_offset = k_cache + src_offset * num_elts_k;\n dtype *k_cache_des_offset = k_cache + des_offset * num_elts_k;\n dtype *v_cache_src_offset = v_cache + src_offset * num_elts_v;\n dtype *v_cache_des_offset = v_cache + des_offset * num_elts_v;\n for (int i = threadIdx.x; i < num_elts; i += blockDim.x) {\n if (i < num_elts_k) {\n *(k_cache_des_offset + i) = *(k_cache_src_offset + i);\n }\n if (i < num_elts_v) {\n *(v_cache_des_offset + i) = *(v_cache_src_offset + i);\n }\n }\n }\n }\n }\n}\n\ntemplate \n__global__ void copy_to_kv_cache_kernel(dtype *new_k, dtype *new_v, dtype *k_cache, dtype *v_cache,\n int *page_table, int *cu_new_lens, int *cu_cached_lens,\n NVTE_QKV_Format qkv_format, int h_kv, int d_k, int d_v,\n int b, int max_ctx_len, int max_seq_len,\n int max_pages_per_seq, bool is_non_paged) {\n // new_k, new_v: qkv_format; k_cache, v_cache: bshd\n // cu_new_lens, cu_cached_lens: [b + 1]\n // page_table: [b, max_pages_per_seq]\n int page_size = max_seq_len / max_pages_per_seq;\n if (qkv_format == NVTE_QKV_Format::NVTE_BSHD) {\n for (int batch_idx = blockIdx.x; batch_idx < b; batch_idx += gridDim.x) {\n int *page_list = is_non_paged ? nullptr : page_table + batch_idx * max_pages_per_seq;\n int cached_len = cu_cached_lens[batch_idx + 1] - cu_cached_lens[batch_idx];\n int new_len = cu_new_lens[batch_idx + 1] - cu_new_lens[batch_idx];\n int num_elts_k = h_kv * d_k;\n int num_elts_v = h_kv * d_v;\n int hd = h_kv * max(d_k, d_v);\n for (int i = blockIdx.y; i < new_len; i += gridDim.y) {\n int page_idx = is_non_paged ? batch_idx : page_list[(cached_len - new_len + i) / page_size];\n dtype *new_token_id_k = new_k + (batch_idx * max_ctx_len + i) * num_elts_k;\n dtype *new_token_id_v = new_v + (batch_idx * max_ctx_len + i) * num_elts_v;\n dtype *token_id_k =\n k_cache + (page_idx * page_size + (cached_len - new_len + i) % page_size) * num_elts_k;\n dtype *token_id_v =\n v_cache + (page_idx * page_size + (cached_len - new_len + i) % page_size) * num_elts_v;\n for (int j = threadIdx.x; j < hd; j += blockDim.x) {\n if (j < num_elts_k) {\n *(token_id_k + j) = *(new_token_id_k + j);\n }\n if (j < num_elts_v) {\n *(token_id_v + j) = *(new_token_id_v + j);\n }\n }\n }\n }\n } else if (qkv_format == NVTE_QKV_Format::NVTE_SBHD) {\n for (int batch_idx = blockIdx.x; batch_idx < b; batch_idx += gridDim.x) {\n int *page_list = is_non_paged ? nullptr : page_table + batch_idx * max_pages_per_seq;\n int cached_len = cu_cached_lens[batch_idx + 1] - cu_cached_lens[batch_idx];\n int new_len = cu_new_lens[batch_idx + 1] - cu_new_lens[batch_idx];\n int num_elts_k = h_kv * d_k;\n int num_elts_v = h_kv * d_v;\n int hd = h_kv * max(d_k, d_v);\n for (int i = blockIdx.y; i < new_len; i += gridDim.y) {\n int page_idx = is_non_paged ? batch_idx : page_list[(cached_len - new_len + i) / page_size];\n dtype *new_token_id_k = new_k + (i * b + batch_idx) * num_elts_k;\n dtype *new_token_id_v = new_v + (i * b + batch_idx) * num_elts_v;\n dtype *token_id_k =\n k_cache + (page_idx * page_size + (cached_len - new_len + i) % page_size) * num_elts_k;\n dtype *token_id_v =\n v_cache + (page_idx * page_size + (cached_len - new_len + i) % page_size) * num_elts_v;\n for (int j = threadIdx.x; j < hd; j += blockDim.x) {\n if (j < num_elts_k) {\n *(token_id_k + j) = *(new_token_id_k + j);\n }\n if (j < num_elts_v) {\n *(token_id_v + j) = *(new_token_id_v + j);\n }\n }\n }\n }\n } else if (qkv_format == NVTE_QKV_Format::NVTE_THD) {\n for (int batch_idx = blockIdx.x; batch_idx < b; batch_idx += gridDim.x) {\n int *page_list = is_non_paged ? nullptr : page_table + batch_idx * max_pages_per_seq;\n int cached_len = cu_cached_lens[batch_idx + 1] - cu_cached_lens[batch_idx];\n int new_len = cu_new_lens[batch_idx + 1] - cu_new_lens[batch_idx];\n int num_elts_k = h_kv * d_k;\n int num_elts_v = h_kv * d_v;\n int hd = h_kv * max(d_k, d_v);\n for (int i = blockIdx.y; i < new_len; i += gridDim.y) {\n int page_idx = is_non_paged ? batch_idx : page_list[(cached_len - new_len + i) / page_size];\n dtype *new_token_id_k = new_k + (cu_new_lens[batch_idx] + i) * num_elts_k;\n dtype *new_token_id_v = new_v + (cu_new_lens[batch_idx] + i) * num_elts_v;\n dtype *token_id_k =\n k_cache + (page_idx * page_size + (cached_len - new_len + i) % page_size) * num_elts_k;\n dtype *token_id_v =\n v_cache + (page_idx * page_size + (cached_len - new_len + i) % page_size) * num_elts_v;\n for (int j = threadIdx.x; j < hd; j += blockDim.x) {\n if (j < num_elts_k) {\n *(token_id_k + j) = *(new_token_id_k + j);\n }\n if (j < num_elts_v) {\n *(token_id_v + j) = *(new_token_id_v + j);\n }\n }\n }\n }\n }\n}\n\ntemplate \nvoid copy_to_kv_cache_launcher(Tensor new_k, Tensor new_v, Tensor k_cache, Tensor v_cache,\n Tensor page_table, Tensor cu_new_lens, Tensor cu_cached_lens,\n NVTE_QKV_Format qkv_format, int h_kv, int d_k, int d_v, int b,\n int max_ctx_len, int max_seq_len, int max_pages_per_seq,\n bool is_non_paged, cudaStream_t stream) {\n if (new_k.has_data() && new_v.has_data() && k_cache.has_data() && v_cache.has_data()) {\n if (is_non_paged) {\n reindex_kv_cache_kernel<<>>(\n reinterpret_cast(k_cache.data.dptr),\n reinterpret_cast(v_cache.data.dptr),\n reinterpret_cast(page_table.data.dptr),\n reinterpret_cast(cu_new_lens.data.dptr),\n reinterpret_cast(cu_cached_lens.data.dptr), h_kv, d_k, d_v, b, max_seq_len);\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n dim3 grid_size(b, max_ctx_len);\n copy_to_kv_cache_kernel<<>>(\n reinterpret_cast(new_k.data.dptr), reinterpret_cast(new_v.data.dptr),\n reinterpret_cast(k_cache.data.dptr), reinterpret_cast(v_cache.data.dptr),\n reinterpret_cast(page_table.data.dptr),\n reinterpret_cast(cu_new_lens.data.dptr),\n reinterpret_cast(cu_cached_lens.data.dptr), qkv_format, h_kv, d_k, d_v, b,\n max_ctx_len, max_seq_len, max_pages_per_seq, is_non_paged);\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\nvoid copy_to_kv_cache(Tensor new_k, Tensor new_v, Tensor k_cache, Tensor v_cache, Tensor page_table,\n Tensor cu_new_lens, Tensor cu_cached_lens, NVTE_QKV_Format qkv_format, int b,\n int max_ctx_len, int max_seq_len, int max_pages_per_seq, bool is_non_paged,\n cudaStream_t stream) {\n int h_kv = new_k.shape()[new_k.dim() - 2];\n int d_k = new_k.shape()[new_k.dim() - 1];\n int d_v = new_v.shape()[new_v.dim() - 1];\n NVTE_CHECK(k_cache.dtype() == v_cache.dtype() && new_k.dtype() == new_v.dtype() &&\n new_k.dtype() == k_cache.dtype(),\n \"new_k, new_v, k_cache and v_cache must be of the same data type.\");\n NVTE_CHECK(qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD ||\n qkv_format == NVTE_QKV_Format::NVTE_THD,\n \"qkv_format must be {BSHD, SBHD, THD}.\");\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_FLOAT(\n k_cache.dtype(), dtype,\n copy_to_kv_cache_launcher(new_k, new_v, k_cache, v_cache, page_table, cu_new_lens,\n cu_cached_lens, qkv_format, h_kv, d_k, d_v, b, max_ctx_len,\n max_seq_len, max_pages_per_seq, is_non_paged, stream););\n}\n\ntemplate \n__global__ void convert_thd_to_bshd_kernel(scalar_t *tensor, scalar_t *new_tensor, int *cu_seqlens,\n int b, int max_seq_len, int h, int d) {\n // tensor: thd; new_tensor: bshd\n // cu_seqlens: [b + 1]\n for (int batch_idx = blockIdx.x; batch_idx < b; batch_idx += gridDim.x) {\n int num_elts = (cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx]) * h * d;\n int thd_offset = cu_seqlens[batch_idx] * h * d;\n int bshd_offset = batch_idx * max_seq_len * h * d;\n scalar_t *thd_token = tensor + thd_offset;\n scalar_t *bshd_token = new_tensor + bshd_offset;\n for (int i = threadIdx.x; i < num_elts; i += blockDim.x) {\n *(bshd_token + i) = *(thd_token + i);\n }\n }\n}\n\ntemplate \nvoid convert_thd_to_bshd_launcher(Tensor tensor, Tensor new_tensor, Tensor cu_seqlens, int b,\n int max_seq_len, int h, int d, cudaStream_t stream) {\n using namespace transformer_engine;\n convert_thd_to_bshd_kernel<<<16, 256, 0, stream>>>(\n reinterpret_cast(tensor.data.dptr),\n reinterpret_cast(new_tensor.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), b, max_seq_len, h, d);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid convert_thd_to_bshd(Tensor tensor, Tensor cu_seqlens, Tensor new_tensor, int b,\n int max_seq_len, cudaStream_t stream) {\n using namespace transformer_engine;\n\n auto tensor_shape = tensor.shape();\n TRANSFORMER_ENGINE_TYPE_SWITCH_FLOAT(\n new_tensor.dtype(), dtype,\n convert_thd_to_bshd_launcher(tensor, new_tensor, cu_seqlens, b, max_seq_len,\n tensor_shape[1], tensor_shape[2], stream););\n}\n\ntemplate \n__global__ void convert_bshd_to_thd_kernel(scalar_t *tensor, scalar_t *new_tensor, int *cu_seqlens,\n int b, int max_seq_len, int h, int d) {\n // tensor: bshd; new_tensor: thd\n // cu_seqlens: [b + 1]\n for (int batch_idx = blockIdx.x; batch_idx < b; batch_idx += gridDim.x) {\n int seqlen = cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx];\n int num_elts = seqlen * h * d;\n int bshd_offset = batch_idx * max_seq_len * h * d;\n int thd_offset = cu_seqlens[batch_idx] * h * d;\n scalar_t *bshd_token = tensor + bshd_offset;\n scalar_t *thd_token = new_tensor + thd_offset;\n for (int i = threadIdx.x; i < num_elts; i += blockDim.x) {\n *(thd_token + i) = *(bshd_token + i);\n }\n }\n}\n\ntemplate \nvoid convert_bshd_to_thd_launcher(Tensor tensor, Tensor new_tensor, Tensor cu_seqlens, int b,\n int max_seq_len, int h, int d, cudaStream_t stream) {\n using namespace transformer_engine;\n convert_bshd_to_thd_kernel<<<16, 256, 0, stream>>>(\n reinterpret_cast(tensor.data.dptr),\n reinterpret_cast(new_tensor.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), b, max_seq_len, h, d);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid convert_bshd_to_thd(Tensor tensor, Tensor cu_seqlens, Tensor new_tensor, int t,\n cudaStream_t stream) {\n using namespace transformer_engine;\n\n auto tensor_shape = tensor.shape();\n TRANSFORMER_ENGINE_TYPE_SWITCH_FLOAT(\n tensor.dtype(), dtype,\n convert_bshd_to_thd_launcher(tensor, new_tensor, cu_seqlens, tensor_shape[0],\n tensor_shape[1], tensor_shape[2], tensor_shape[3],\n stream););\n}\n\n} // namespace kv_cache\n} // namespace transformer_engine\n\n/***************************************************************************************************\n * KV Cache: Copy new KV tokens to the KV cache\n * 1. new_k and new_v are in qkv_format; k_cache and v_cache are in 'bshd' format\n * 2. cu_new_lens and cu_cached_lens are of shape [b + 1]; cu_cached_lens include the added lens\n * in current step\n * 3. Non-paged KV cache is a special case of paged KV cache, with page_table = [b, 1] and\n * max_pages_per_seq = 1. We use the same underlying kernel for both non-paged and paged.\n * Set is_non_paged = True/False to indicate as such.\n * 4. is_non_paged = True also re-indexes the KV cache, e.g. the initial batch indices [0, 3, 1, 2]\n * becomes [0, 1, 1, 2]. The page_table = batch_indices.unsqueeze(1) is however unchanged.\n * batch_indices_post can be used for monotonical indexing, i.e. [0, 1, 2, 3]. batch_indices is\n * preserved for the next layer in the same iteration.\n * 5. Only supports same page_table for k_cache and v_cache\n * 6. Only pad_between_seqs = False when qkv_format = thd, i.e. there should be no pad tokens\n * between sequences in new_k and new_v such as [a a a 0..0 b b 0..0 c 0..0].\n **************************************************************************************************/\n\nvoid nvte_copy_to_kv_cache(NVTETensor new_k, NVTETensor new_v, NVTETensor k_cache,\n NVTETensor v_cache, NVTETensor page_table, NVTETensor cu_new_lens,\n NVTETensor cu_cached_lens, NVTE_QKV_Format qkv_format, int b,\n int max_ctx_len, int max_seq_len, int max_pages_per_seq,\n int is_non_paged, cudaStream_t stream) {\n NVTE_API_CALL(nvte_copy_to_kv_cache);\n using namespace transformer_engine;\n\n kv_cache::copy_to_kv_cache(*convertNVTETensorCheck(new_k), *convertNVTETensorCheck(new_v),\n *convertNVTETensorCheck(k_cache), *convertNVTETensorCheck(v_cache),\n *convertNVTETensorCheck(page_table),\n *convertNVTETensorCheck(cu_new_lens),\n *convertNVTETensorCheck(cu_cached_lens), qkv_format, b, max_ctx_len,\n max_seq_len, max_pages_per_seq, is_non_paged, stream);\n}\n\n/***************************************************************************************************\n * KV Cache: Convert a tensor from qkv_format = thd to qkv_format = bshd\n **************************************************************************************************/\n\nvoid nvte_convert_thd_to_bshd(NVTETensor tensor, NVTETensor cu_seqlens, NVTETensor new_tensor,\n int b, int max_seq_len, cudaStream_t stream) {\n NVTE_API_CALL(nvte_convert_thd_to_bshd);\n using namespace transformer_engine;\n\n kv_cache::convert_thd_to_bshd(*convertNVTETensorCheck(tensor),\n *convertNVTETensorCheck(cu_seqlens),\n *convertNVTETensorCheck(new_tensor), b, max_seq_len, stream);\n}\n\n/***************************************************************************************************\n * KV Cache: Convert a tensor from qkv_format = bshd to qkv_format = thd\n **************************************************************************************************/\n\nvoid nvte_convert_bshd_to_thd(NVTETensor tensor, NVTETensor cu_seqlens, NVTETensor new_tensor,\n int t, cudaStream_t stream) {\n NVTE_API_CALL(nvte_convert_bshd_to_thd);\n using namespace transformer_engine;\n\n kv_cache::convert_bshd_to_thd(*convertNVTETensorCheck(tensor),\n *convertNVTETensorCheck(cu_seqlens),\n *convertNVTETensorCheck(new_tensor), t, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \"../common.h\"\n#include \"transformer_engine/fused_attn.h\"\n\nnamespace transformer_engine {\nnamespace kv_cache {\n\nconstexpr int block_size = 1024;\n\ntemplate \n__global__ void reindex_kv_cache_kernel(dtype *k_cache, dtype *v_cache, int *batch_indices,\n int *cu_new_lens, int *cu_cached_lens, int h_kv, int d_k,\n int d_v, int b, int max_seq_len) {\n // k_cache, v_cache: bshd\n // batch_indices: [b]; cu_new_lens, cu_cached_lens: [b + 1]\n int actual_b = b;\n for (int i = 0; i < b - 1; i++) {\n if (batch_indices[i + 1] < batch_indices[i]) {\n actual_b = i + 1;\n }\n }\n bool flag = (batch_indices[0] != 0);\n for (int batch_idx = 0; batch_idx < actual_b; batch_idx++) {\n if (flag || ((batch_indices[batch_idx] - batch_indices[0]) != batch_idx)) {\n int num_tokens = (cu_cached_lens[batch_idx + 1] - cu_cached_lens[batch_idx]) -\n (cu_new_lens[batch_idx + 1] - cu_new_lens[batch_idx]);\n int num_elts_k = h_kv * d_k;\n int num_elts_v = h_kv * d_v;\n int num_elts = max(num_elts_k, num_elts_v);\n for (int token_idx = blockIdx.x; token_idx < num_tokens; token_idx += gridDim.x) {\n int src_offset = batch_indices[batch_idx] * max_seq_len + token_idx;\n int des_offset = batch_idx * max_seq_len + token_idx;\n dtype *k_cache_src_offset = k_cache + src_offset * num_elts_k;\n dtype *k_cache_des_offset = k_cache + des_offset * num_elts_k;\n dtype *v_cache_src_offset = v_cache + src_offset * num_elts_v;\n dtype *v_cache_des_offset = v_cache + des_offset * num_elts_v;\n for (int i = threadIdx.x; i < num_elts; i += blockDim.x) {\n if (i < num_elts_k) {\n *(k_cache_des_offset + i) = *(k_cache_src_offset + i);\n }\n if (i < num_elts_v) {\n *(v_cache_des_offset + i) = *(v_cache_src_offset + i);\n }\n }\n }\n }\n }\n}\n\ntemplate \n__global__ void copy_to_kv_cache_kernel(dtype *new_k, dtype *new_v, dtype *k_cache, dtype *v_cache,\n int *page_table, int *cu_new_lens, int *cu_cached_lens,\n NVTE_QKV_Format qkv_format, int h_kv, int d_k, int d_v,\n int b, int max_ctx_len, int max_seq_len,\n int max_pages_per_seq, bool is_non_paged) {\n // new_k, new_v: qkv_format; k_cache, v_cache: bshd\n // cu_new_lens, cu_cached_lens: [b + 1]\n // page_table: [b, max_pages_per_seq]\n int page_size = max_seq_len / max_pages_per_seq;\n if (qkv_format == NVTE_QKV_Format::NVTE_BSHD) {\n for (int batch_idx = blockIdx.x; batch_idx < b; batch_idx += gridDim.x) {\n int *page_list = is_non_paged ? nullptr : page_table + batch_idx * max_pages_per_seq;\n int cached_len = cu_cached_lens[batch_idx + 1] - cu_cached_lens[batch_idx];\n int new_len = cu_new_lens[batch_idx + 1] - cu_new_lens[batch_idx];\n int num_elts_k = h_kv * d_k;\n int num_elts_v = h_kv * d_v;\n int hd = h_kv * max(d_k, d_v);\n for (int i = blockIdx.y; i < new_len; i += gridDim.y) {\n int page_idx = is_non_paged ? batch_idx : page_list[(cached_len - new_len + i) / page_size];\n dtype *new_token_id_k = new_k + (batch_idx * max_ctx_len + i) * num_elts_k;\n dtype *new_token_id_v = new_v + (batch_idx * max_ctx_len + i) * num_elts_v;\n dtype *token_id_k =\n k_cache + (page_idx * page_size + (cached_len - new_len + i) % page_size) * num_elts_k;\n dtype *token_id_v =\n v_cache + (page_idx * page_size + (cached_len - new_len + i) % page_size) * num_elts_v;\n for (int j = threadIdx.x; j < hd; j += blockDim.x) {\n if (j < num_elts_k) {\n *(token_id_k + j) = *(new_token_id_k + j);\n }\n if (j < num_elts_v) {\n *(token_id_v + j) = *(new_token_id_v + j);\n }\n }\n }\n }\n } else if (qkv_format == NVTE_QKV_Format::NVTE_SBHD) {\n for (int batch_idx = blockIdx.x; batch_idx < b; batch_idx += gridDim.x) {\n int *page_list = is_non_paged ? nullptr : page_table + batch_idx * max_pages_per_seq;\n int cached_len = cu_cached_lens[batch_idx + 1] - cu_cached_lens[batch_idx];\n int new_len = cu_new_lens[batch_idx + 1] - cu_new_lens[batch_idx];\n int num_elts_k = h_kv * d_k;\n int num_elts_v = h_kv * d_v;\n int hd = h_kv * max(d_k, d_v);\n for (int i = blockIdx.y; i < new_len; i += gridDim.y) {\n int page_idx = is_non_paged ? batch_idx : page_list[(cached_len - new_len + i) / page_size];\n dtype *new_token_id_k = new_k + (i * b + batch_idx) * num_elts_k;\n dtype *new_token_id_v = new_v + (i * b + batch_idx) * num_elts_v;\n dtype *token_id_k =\n k_cache + (page_idx * page_size + (cached_len - new_len + i) % page_size) * num_elts_k;\n dtype *token_id_v =\n v_cache + (page_idx * page_size + (cached_len - new_len + i) % page_size) * num_elts_v;\n for (int j = threadIdx.x; j < hd; j += blockDim.x) {\n if (j < num_elts_k) {\n *(token_id_k + j) = *(new_token_id_k + j);\n }\n if (j < num_elts_v) {\n *(token_id_v + j) = *(new_token_id_v + j);\n }\n }\n }\n }\n } else if (qkv_format == NVTE_QKV_Format::NVTE_THD) {\n for (int batch_idx = blockIdx.x; batch_idx < b; batch_idx += gridDim.x) {\n int *page_list = is_non_paged ? nullptr : page_table + batch_idx * max_pages_per_seq;\n int cached_len = cu_cached_lens[batch_idx + 1] - cu_cached_lens[batch_idx];\n int new_len = cu_new_lens[batch_idx + 1] - cu_new_lens[batch_idx];\n int num_elts_k = h_kv * d_k;\n int num_elts_v = h_kv * d_v;\n int hd = h_kv * max(d_k, d_v);\n for (int i = blockIdx.y; i < new_len; i += gridDim.y) {\n int page_idx = is_non_paged ? batch_idx : page_list[(cached_len - new_len + i) / page_size];\n dtype *new_token_id_k = new_k + (cu_new_lens[batch_idx] + i) * num_elts_k;\n dtype *new_token_id_v = new_v + (cu_new_lens[batch_idx] + i) * num_elts_v;\n dtype *token_id_k =\n k_cache + (page_idx * page_size + (cached_len - new_len + i) % page_size) * num_elts_k;\n dtype *token_id_v =\n v_cache + (page_idx * page_size + (cached_len - new_len + i) % page_size) * num_elts_v;\n for (int j = threadIdx.x; j < hd; j += blockDim.x) {\n if (j < num_elts_k) {\n *(token_id_k + j) = *(new_token_id_k + j);\n }\n if (j < num_elts_v) {\n *(token_id_v + j) = *(new_token_id_v + j);\n }\n }\n }\n }\n }\n}\n\ntemplate \nvoid copy_to_kv_cache_launcher(Tensor new_k, Tensor new_v, Tensor k_cache, Tensor v_cache,\n Tensor page_table, Tensor cu_new_lens, Tensor cu_cached_lens,\n NVTE_QKV_Format qkv_format, int h_kv, int d_k, int d_v, int b,\n int max_ctx_len, int max_seq_len, int max_pages_per_seq,\n bool is_non_paged, hipStream_t stream) {\n if (new_k.has_data() && new_v.has_data() && k_cache.has_data() && v_cache.has_data()) {\n if (is_non_paged) {\n reindex_kv_cache_kernel<<>>(\n reinterpret_cast(k_cache.data.dptr),\n reinterpret_cast(v_cache.data.dptr),\n reinterpret_cast(page_table.data.dptr),\n reinterpret_cast(cu_new_lens.data.dptr),\n reinterpret_cast(cu_cached_lens.data.dptr), h_kv, d_k, d_v, b, max_seq_len);\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n dim3 grid_size(b, max_ctx_len);\n copy_to_kv_cache_kernel<<>>(\n reinterpret_cast(new_k.data.dptr), reinterpret_cast(new_v.data.dptr),\n reinterpret_cast(k_cache.data.dptr), reinterpret_cast(v_cache.data.dptr),\n reinterpret_cast(page_table.data.dptr),\n reinterpret_cast(cu_new_lens.data.dptr),\n reinterpret_cast(cu_cached_lens.data.dptr), qkv_format, h_kv, d_k, d_v, b,\n max_ctx_len, max_seq_len, max_pages_per_seq, is_non_paged);\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\nvoid copy_to_kv_cache(Tensor new_k, Tensor new_v, Tensor k_cache, Tensor v_cache, Tensor page_table,\n Tensor cu_new_lens, Tensor cu_cached_lens, NVTE_QKV_Format qkv_format, int b,\n int max_ctx_len, int max_seq_len, int max_pages_per_seq, bool is_non_paged,\n hipStream_t stream) {\n int h_kv = new_k.shape()[new_k.dim() - 2];\n int d_k = new_k.shape()[new_k.dim() - 1];\n int d_v = new_v.shape()[new_v.dim() - 1];\n NVTE_CHECK(k_cache.dtype() == v_cache.dtype() && new_k.dtype() == new_v.dtype() &&\n new_k.dtype() == k_cache.dtype(),\n \"new_k, new_v, k_cache and v_cache must be of the same data type.\");\n NVTE_CHECK(qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD ||\n qkv_format == NVTE_QKV_Format::NVTE_THD,\n \"qkv_format must be {BSHD, SBHD, THD}.\");\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_FLOAT(\n k_cache.dtype(), dtype,\n copy_to_kv_cache_launcher(new_k, new_v, k_cache, v_cache, page_table, cu_new_lens,\n cu_cached_lens, qkv_format, h_kv, d_k, d_v, b, max_ctx_len,\n max_seq_len, max_pages_per_seq, is_non_paged, stream););\n}\n\ntemplate \n__global__ void convert_thd_to_bshd_kernel(scalar_t *tensor, scalar_t *new_tensor, int *cu_seqlens,\n int b, int max_seq_len, int h, int d) {\n // tensor: thd; new_tensor: bshd\n // cu_seqlens: [b + 1]\n for (int batch_idx = blockIdx.x; batch_idx < b; batch_idx += gridDim.x) {\n int num_elts = (cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx]) * h * d;\n int thd_offset = cu_seqlens[batch_idx] * h * d;\n int bshd_offset = batch_idx * max_seq_len * h * d;\n scalar_t *thd_token = tensor + thd_offset;\n scalar_t *bshd_token = new_tensor + bshd_offset;\n for (int i = threadIdx.x; i < num_elts; i += blockDim.x) {\n *(bshd_token + i) = *(thd_token + i);\n }\n }\n}\n\ntemplate \nvoid convert_thd_to_bshd_launcher(Tensor tensor, Tensor new_tensor, Tensor cu_seqlens, int b,\n int max_seq_len, int h, int d, hipStream_t stream) {\n using namespace transformer_engine;\n convert_thd_to_bshd_kernel<<<16, 256, 0, stream>>>(\n reinterpret_cast(tensor.data.dptr),\n reinterpret_cast(new_tensor.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), b, max_seq_len, h, d);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid convert_thd_to_bshd(Tensor tensor, Tensor cu_seqlens, Tensor new_tensor, int b,\n int max_seq_len, hipStream_t stream) {\n using namespace transformer_engine;\n\n auto tensor_shape = tensor.shape();\n TRANSFORMER_ENGINE_TYPE_SWITCH_FLOAT(\n new_tensor.dtype(), dtype,\n convert_thd_to_bshd_launcher(tensor, new_tensor, cu_seqlens, b, max_seq_len,\n tensor_shape[1], tensor_shape[2], stream););\n}\n\ntemplate \n__global__ void convert_bshd_to_thd_kernel(scalar_t *tensor, scalar_t *new_tensor, int *cu_seqlens,\n int b, int max_seq_len, int h, int d) {\n // tensor: bshd; new_tensor: thd\n // cu_seqlens: [b + 1]\n for (int batch_idx = blockIdx.x; batch_idx < b; batch_idx += gridDim.x) {\n int seqlen = cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx];\n int num_elts = seqlen * h * d;\n int bshd_offset = batch_idx * max_seq_len * h * d;\n int thd_offset = cu_seqlens[batch_idx] * h * d;\n scalar_t *bshd_token = tensor + bshd_offset;\n scalar_t *thd_token = new_tensor + thd_offset;\n for (int i = threadIdx.x; i < num_elts; i += blockDim.x) {\n *(thd_token + i) = *(bshd_token + i);\n }\n }\n}\n\ntemplate \nvoid convert_bshd_to_thd_launcher(Tensor tensor, Tensor new_tensor, Tensor cu_seqlens, int b,\n int max_seq_len, int h, int d, hipStream_t stream) {\n using namespace transformer_engine;\n convert_bshd_to_thd_kernel<<<16, 256, 0, stream>>>(\n reinterpret_cast(tensor.data.dptr),\n reinterpret_cast(new_tensor.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), b, max_seq_len, h, d);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid convert_bshd_to_thd(Tensor tensor, Tensor cu_seqlens, Tensor new_tensor, int t,\n hipStream_t stream) {\n using namespace transformer_engine;\n\n auto tensor_shape = tensor.shape();\n TRANSFORMER_ENGINE_TYPE_SWITCH_FLOAT(\n tensor.dtype(), dtype,\n convert_bshd_to_thd_launcher(tensor, new_tensor, cu_seqlens, tensor_shape[0],\n tensor_shape[1], tensor_shape[2], tensor_shape[3],\n stream););\n}\n\n} // namespace kv_cache\n} // namespace transformer_engine\n\n/***************************************************************************************************\n * KV Cache: Copy new KV tokens to the KV cache\n * 1. new_k and new_v are in qkv_format; k_cache and v_cache are in 'bshd' format\n * 2. cu_new_lens and cu_cached_lens are of shape [b + 1]; cu_cached_lens include the added lens\n * in current step\n * 3. Non-paged KV cache is a special case of paged KV cache, with page_table = [b, 1] and\n * max_pages_per_seq = 1. We use the same underlying kernel for both non-paged and paged.\n * Set is_non_paged = True/False to indicate as such.\n * 4. is_non_paged = True also re-indexes the KV cache, e.g. the initial batch indices [0, 3, 1, 2]\n * becomes [0, 1, 1, 2]. The page_table = batch_indices.unsqueeze(1) is however unchanged.\n * batch_indices_post can be used for monotonical indexing, i.e. [0, 1, 2, 3]. batch_indices is\n * preserved for the next layer in the same iteration.\n * 5. Only supports same page_table for k_cache and v_cache\n * 6. Only pad_between_seqs = False when qkv_format = thd, i.e. there should be no pad tokens\n * between sequences in new_k and new_v such as [a a a 0..0 b b 0..0 c 0..0].\n **************************************************************************************************/\n\nvoid nvte_copy_to_kv_cache(NVTETensor new_k, NVTETensor new_v, NVTETensor k_cache,\n NVTETensor v_cache, NVTETensor page_table, NVTETensor cu_new_lens,\n NVTETensor cu_cached_lens, NVTE_QKV_Format qkv_format, int b,\n int max_ctx_len, int max_seq_len, int max_pages_per_seq,\n int is_non_paged, hipStream_t stream) {\n NVTE_API_CALL(nvte_copy_to_kv_cache);\n using namespace transformer_engine;\n\n kv_cache::copy_to_kv_cache(*convertNVTETensorCheck(new_k), *convertNVTETensorCheck(new_v),\n *convertNVTETensorCheck(k_cache), *convertNVTETensorCheck(v_cache),\n *convertNVTETensorCheck(page_table),\n *convertNVTETensorCheck(cu_new_lens),\n *convertNVTETensorCheck(cu_cached_lens), qkv_format, b, max_ctx_len,\n max_seq_len, max_pages_per_seq, is_non_paged, stream);\n}\n\n/***************************************************************************************************\n * KV Cache: Convert a tensor from qkv_format = thd to qkv_format = bshd\n **************************************************************************************************/\n\nvoid nvte_convert_thd_to_bshd(NVTETensor tensor, NVTETensor cu_seqlens, NVTETensor new_tensor,\n int b, int max_seq_len, hipStream_t stream) {\n NVTE_API_CALL(nvte_convert_thd_to_bshd);\n using namespace transformer_engine;\n\n kv_cache::convert_thd_to_bshd(*convertNVTETensorCheck(tensor),\n *convertNVTETensorCheck(cu_seqlens),\n *convertNVTETensorCheck(new_tensor), b, max_seq_len, stream);\n}\n\n/***************************************************************************************************\n * KV Cache: Convert a tensor from qkv_format = bshd to qkv_format = thd\n **************************************************************************************************/\n\nvoid nvte_convert_bshd_to_thd(NVTETensor tensor, NVTETensor cu_seqlens, NVTETensor new_tensor,\n int t, hipStream_t stream) {\n NVTE_API_CALL(nvte_convert_bshd_to_thd);\n using namespace transformer_engine;\n\n kv_cache::convert_bshd_to_thd(*convertNVTETensorCheck(tensor),\n *convertNVTETensorCheck(cu_seqlens),\n *convertNVTETensorCheck(new_tensor), t, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/fused_attn/kv_cache.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/cuda_driver.h\"\n#include \"../util/cuda_runtime.h\"\n#include \"../util/ptx.cuh\"\n#include \"../utils.cuh\"\n#include \"transformer_engine/fused_attn.h\"\n\nnamespace transformer_engine {\n\n// ============================================================================\n// prepare_flash_attn: SBH3D <-> BSHD_BSHD_BSHD for the FlashAttention backend\n// ============================================================================\n\nnamespace flash_attention {\n\n/// Packed vector of N elements of T; alignment matches a single wide load/store of N * sizeof(T) bytes.\ntemplate \nstruct alignas(sizeof(T) * N) Vec {\n T data[N];\n};\n\nconstexpr int warp_size = 32;\nconstexpr int type_size = 2; // FP16 or BF16\nconstexpr int nvec = sizeof(uint64_t) / type_size;\nconstexpr int load_size = warp_size * nvec;\nconstexpr int block_size = 512;\n\ntemplate \n__launch_bounds__(block_size) __global__\n void prepare_kernel_fwd(const T *qkvi, T *qkv, const size_t B, const size_t S, const size_t Z,\n const size_t W) {\n const int warpid = (blockDim.x * blockIdx.x + threadIdx.x) / warp_size;\n const int id_in_warp = threadIdx.x % warp_size;\n const size_t offset_input = blockIdx.y * W + warpid * 3 * W * Z + id_in_warp * nvec;\n const T *my_input = qkvi + offset_input;\n\n const size_t s = warpid / B;\n if (s >= S) return;\n\n const size_t b = warpid % B;\n\n const size_t offset_output = blockIdx.y * B * S * Z * W + (s + b * S) * W * Z + id_in_warp * nvec;\n\n T *my_output = qkv + offset_output;\n\n for (int i = 0; i < Z; ++i) {\n Vec *const out = reinterpret_cast *>(my_output + i * load_size);\n *out = *reinterpret_cast *>(my_input + i * load_size * 3);\n }\n}\n\ntemplate \n__launch_bounds__(block_size) __global__\n void prepare_kernel_bwd(const T *q, const T *k, const T *v, T *qkv, const size_t B,\n const size_t S, const size_t Z, const size_t W) {\n const T *input = blockIdx.y == 0 ? q : (blockIdx.y == 1 ? k : v);\n\n const int warpid = (blockDim.x * blockIdx.x + threadIdx.x) / warp_size;\n const int id_in_warp = threadIdx.x % warp_size;\n const size_t offset_input = warpid * W * Z + id_in_warp * nvec;\n const T *my_input = input + offset_input;\n\n const size_t b = warpid / S;\n if (b >= B) return;\n\n const size_t s = warpid % S;\n\n const size_t offset_output = (b + s * B) * 3 * W * Z + id_in_warp * nvec + blockIdx.y * W;\n\n T *my_output = qkv + offset_output;\n\n for (int i = 0; i < Z; ++i) {\n Vec *const out = reinterpret_cast *>(my_output + i * load_size * 3);\n *out = *reinterpret_cast *>(my_input + i * load_size);\n }\n}\n\nvoid prepare_flash_attn_fwd(Tensor qkvi, Tensor qkv, cudaStream_t stream) {\n using namespace transformer_engine;\n\n NVTE_CHECK(qkvi.dim() == 4, \"Expected 4-dim tensor.\");\n NVTE_CHECK(qkvi.dtype() == DType::kFloat16 || qkvi.dtype() == DType::kBFloat16);\n\n auto qkvi_shape = qkvi.shape();\n\n NVTE_CHECK(qkvi_shape[3] % load_size == 0);\n NVTE_CHECK(qkvi_shape[3] == load_size);\n\n // [s, b, n, h * 3] -> [3, b, s, n, h]\n std::vector shape = {3, qkvi_shape[1], qkvi_shape[0], qkvi_shape[2], qkvi_shape[3]};\n\n size_t warps = qkvi_shape[0] * qkvi_shape[1];\n size_t warps_per_block = block_size / warp_size;\n size_t blocks = (warps + warps_per_block - 1) / warps_per_block;\n dim3 grid(blocks, 3);\n int threads = block_size;\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n qkvi.dtype(), dtype,\n prepare_kernel_fwd<<>>(\n reinterpret_cast(qkvi.data.dptr), reinterpret_cast(qkv.data.dptr),\n shape[1], shape[2], shape[3], shape[4]););\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid prepare_flash_attn_bwd(Tensor q, Tensor k, Tensor v, Tensor qkv, cudaStream_t stream) {\n using namespace transformer_engine;\n\n NVTE_CHECK(q.dim() == 4, \"Expected 4-dim tensor.\");\n NVTE_CHECK(k.dim() == 4, \"Expected 4-dim tensor.\");\n NVTE_CHECK(v.dim() == 4, \"Expected 4-dim tensor.\");\n NVTE_CHECK(q.dtype() == DType::kFloat16 || q.dtype() == DType::kBFloat16);\n NVTE_CHECK(k.dtype() == q.dtype());\n NVTE_CHECK(v.dtype() == q.dtype());\n\n auto q_shape = q.shape();\n auto k_shape = k.shape();\n auto v_shape = v.shape();\n\n NVTE_CHECK(q_shape[3] % load_size == 0);\n NVTE_CHECK(q_shape[3] == load_size);\n NVTE_CHECK(k_shape[3] % load_size == 0);\n NVTE_CHECK(k_shape[3] == load_size);\n NVTE_CHECK(v_shape[3] % load_size == 0);\n NVTE_CHECK(v_shape[3] == load_size);\n\n // 3 x [s, b, n, h] -> [b, s, n, 3 * h]\n std::vector shape = {q_shape[1], q_shape[0], q_shape[2], 3 * q_shape[3]};\n\n size_t warps = q_shape[0] * q_shape[1];\n size_t warps_per_block = block_size / warp_size;\n size_t blocks = (warps + warps_per_block - 1) / warps_per_block;\n dim3 grid(blocks, 3);\n int threads = block_size;\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n q.dtype(), dtype,\n prepare_kernel_bwd<<>>(\n reinterpret_cast(q.data.dptr), reinterpret_cast(k.data.dptr),\n reinterpret_cast(v.data.dptr), reinterpret_cast(qkv.data.dptr),\n q_shape[0], q_shape[1], q_shape[2], q_shape[3]););\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n} // namespace flash_attention\n\n// ============================================================================\n// multi_tensor_transpose_to_bhsd: BSHD/SBHD -> BHSD\n// ============================================================================\n\nnamespace multi_tensor_transpose_to_bhsd {\n\nusing flash_attention::Vec;\n\nconstexpr int kMaxPermuteTensors = 16;\n\nstruct PermuteSlot {\n const void *input;\n void *output;\n size_t S, H, D_in, D_out;\n};\n\nstruct PermuteParams {\n PermuteSlot slots[kMaxPermuteTensors];\n};\n\nstruct TmaMapParams {\n CUtensorMap maps[kMaxPermuteTensors];\n};\n\n// ---------- path 3: fallback_not_vec_aligned ----------\n\n__device__ __forceinline__ void copy_row_bytes(const char *__restrict__ src, char *__restrict__ dst,\n size_t D_bytes) {\n size_t off = 0;\n for (; off + 16 <= D_bytes; off += 16) {\n uint4 tmp;\n memcpy(&tmp, src + off, 16);\n memcpy(dst + off, &tmp, 16);\n }\n for (; off + 8 <= D_bytes; off += 8) {\n uint2 tmp;\n memcpy(&tmp, src + off, 8);\n memcpy(dst + off, &tmp, 8);\n }\n for (; off + 4 <= D_bytes; off += 4) {\n unsigned int tmp;\n memcpy(&tmp, src + off, 4);\n memcpy(dst + off, &tmp, 4);\n }\n for (; off + 2 <= D_bytes; off += 2) {\n uint16_t tmp;\n memcpy(&tmp, src + off, 2);\n memcpy(dst + off, &tmp, 2);\n }\n for (; off < D_bytes; ++off) dst[off] = src[off];\n}\n\n__device__ __forceinline__ void copy_and_pad_row_bytes(const char *__restrict__ src,\n char *__restrict__ dst, size_t D_bytes,\n size_t D_out_bytes) {\n copy_row_bytes(src, dst, D_bytes);\n for (size_t off = D_bytes; off < D_out_bytes; ++off) dst[off] = 0;\n}\n\nconstexpr int TRANSPOSE_TILE = 32;\nconstexpr int TRANSPOSE_BLOCK = 256;\nconstexpr int TRANSPOSE_WARPS = TRANSPOSE_BLOCK / 32; // 8\n\ntemplate \n__launch_bounds__(TRANSPOSE_BLOCK) __global__\n void transpose_to_bhsd_fallback_not_vec_aligned_kernel(PermuteParams params, size_t b,\n unsigned int s_tiles) {\n const auto &slot = params.slots[blockIdx.z];\n const T *__restrict__ in = reinterpret_cast(slot.input);\n T *__restrict__ out = reinterpret_cast(slot.output);\n const size_t S = slot.S;\n const size_t H = slot.H;\n const size_t D = slot.D_in;\n const size_t D_out = slot.D_out;\n const size_t D_bytes = D * sizeof(T);\n const size_t D_out_bytes = D_out * sizeof(T);\n const size_t D_smem_pad = (D_bytes + 3u) & ~size_t(3);\n\n const size_t tile_s = static_cast(blockIdx.x) % static_cast(s_tiles);\n const size_t b_i = static_cast(blockIdx.x) / static_cast(s_tiles);\n if (b_i >= b) return;\n const size_t tile_h = static_cast(blockIdx.y);\n\n const size_t s_base = tile_s * TRANSPOSE_TILE;\n const size_t h_base = tile_h * TRANSPOSE_TILE;\n\n extern __shared__ char smem[];\n const size_t smem_row = static_cast(TRANSPOSE_TILE) * D_smem_pad + 4;\n\n // ---- Phase 1: global \u2192 smem (sweep consecutive H \u2192 coalesced reads) ----\n for (unsigned int warp_off = threadIdx.x >> 5; warp_off < TRANSPOSE_TILE;\n warp_off += TRANSPOSE_WARPS) {\n const size_t local_s = warp_off;\n const size_t local_h = threadIdx.x & 31u;\n const size_t s_i = s_base + local_s;\n const size_t h_i = h_base + local_h;\n if (s_i < S && h_i < H) {\n const char *__restrict__ src;\n if constexpr (kIsBshd)\n src = reinterpret_cast(in + b_i * S * H * D + s_i * H * D + h_i * D);\n else\n src = reinterpret_cast(in + s_i * b * H * D + b_i * H * D + h_i * D);\n copy_row_bytes(src, smem + local_s * smem_row + local_h * D_smem_pad, D_bytes);\n }\n }\n\n __syncthreads();\n\n // ---- Phase 2: smem \u2192 global (sweep consecutive S \u2192 coalesced writes, with padding) ----\n for (unsigned int warp_off = threadIdx.x >> 5; warp_off < TRANSPOSE_TILE;\n warp_off += TRANSPOSE_WARPS) {\n const size_t local_h = warp_off;\n const size_t local_s = threadIdx.x & 31u;\n const size_t s_i = s_base + local_s;\n const size_t h_i = h_base + local_h;\n if (s_i < S && h_i < H) {\n copy_and_pad_row_bytes(\n smem + local_s * smem_row + local_h * D_smem_pad,\n reinterpret_cast(out + b_i * H * S * D_out + h_i * S * D_out + s_i * D_out),\n D_bytes, D_out_bytes);\n }\n }\n}\n\n// ---------- path 2: fallback_vec_aligned ----------\n\nconstexpr int fallback_permute_threads = 1024;\n\ntemplate \n__device__ __forceinline__ void permute_vec_loop(const T *__restrict__ in, T *__restrict__ out,\n size_t b, size_t S, size_t H, size_t D,\n size_t D_out, size_t b_i, size_t h_i,\n size_t s_begin, size_t S_chunk) {\n const size_t out_base = b_i * H * S * D_out + h_i * S * D_out;\n const size_t d_vec = D / static_cast(N);\n const size_t total_work = S_chunk * d_vec;\n for (size_t w = static_cast(threadIdx.x); w < total_work;\n w += static_cast(blockDim.x)) {\n const size_t s_local = w / d_vec;\n const size_t s_i = s_begin + s_local;\n const size_t d_off = (w % d_vec) * static_cast(N);\n const T *__restrict__ in_ptr;\n if constexpr (kIsBshd) {\n in_ptr = in + b_i * (S * H * D) + s_i * (H * D) + h_i * D + d_off;\n } else {\n in_ptr = in + s_i * (b * H * D) + b_i * (H * D) + h_i * D + d_off;\n }\n T *__restrict__ out_ptr = out + out_base + s_i * D_out + d_off;\n *reinterpret_cast *>(out_ptr) = *reinterpret_cast *>(in_ptr);\n }\n if (D_out > D) {\n const size_t pad_elems = D_out - D;\n const size_t total_pad = S_chunk * pad_elems;\n for (size_t w = static_cast(threadIdx.x); w < total_pad;\n w += static_cast(blockDim.x)) {\n const size_t s_local = w / pad_elems;\n const size_t s_i = s_begin + s_local;\n const size_t d_off = D + (w % pad_elems);\n out[out_base + s_i * D_out + d_off] = static_cast(0.f);\n }\n }\n}\n\ntemplate \n__launch_bounds__(fallback_permute_threads) __global__\n void transpose_to_bhsd_fallback_vec_aligned_kernel(PermuteParams params, size_t b,\n unsigned int permute_s_splits,\n size_t h_grid) {\n const auto &slot = params.slots[blockIdx.z];\n const T *__restrict__ in = reinterpret_cast(slot.input);\n T *__restrict__ out = reinterpret_cast(slot.output);\n const size_t S = slot.S;\n const size_t H = slot.H;\n const size_t D = slot.D_in;\n const size_t D_out = slot.D_out;\n\n const size_t b_i = static_cast(blockIdx.x) / h_grid;\n const size_t h_i = static_cast(blockIdx.x) % h_grid;\n if (b_i >= b) return;\n if (h_i >= H) return;\n\n const unsigned int s_part = blockIdx.y;\n const size_t s_begin = (S * static_cast(s_part)) / static_cast(permute_s_splits);\n const size_t s_end =\n (S * static_cast(s_part + 1)) / static_cast(permute_s_splits);\n if (s_begin >= s_end) return;\n const size_t S_chunk = s_end - s_begin;\n\n const size_t D_bytes = D * sizeof(T);\n\n if (D_bytes % 16 == 0) {\n constexpr size_t N = 16 / sizeof(T);\n permute_vec_loop(in, out, b, S, H, D, D_out, b_i, h_i, s_begin, S_chunk);\n return;\n }\n if (D_bytes % 8 == 0) {\n constexpr size_t N = 8 / sizeof(T);\n permute_vec_loop(in, out, b, S, H, D, D_out, b_i, h_i, s_begin, S_chunk);\n return;\n }\n if constexpr (sizeof(T) <= 4) {\n if (D_bytes % 4 == 0) {\n constexpr size_t N = 4 / sizeof(T);\n permute_vec_loop(in, out, b, S, H, D, D_out, b_i, h_i, s_begin, S_chunk);\n return;\n }\n }\n}\n\n// ---------- path 1: TMA ----------\n\nconstexpr int tma_permute_threads = 128;\nconstexpr int tma_permute_s_tile_default = 32;\n\n__device__ __forceinline__ void cp_async_bulk_tensor_4d_global_to_shared(\n void *dst_shmem, const CUtensorMap *tensor_map, uint32_t c0, uint32_t c1, uint32_t c2,\n uint32_t c3, uint64_t *mbar) {\n#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n uint32_t dst = __cvta_generic_to_shared(dst_shmem);\n uint32_t bar = __cvta_generic_to_shared(mbar);\n asm volatile(\n \"cp.async.bulk.tensor.4d.shared::cluster.global.tile\"\n \".mbarrier::complete_tx::bytes [%0], [%1, {%2, %3, %4, %5}], [%6];\" ::\"r\"(dst),\n \"l\"(tensor_map), \"r\"(c0), \"r\"(c1), \"r\"(c2), \"r\"(c3), \"r\"(bar)\n : \"memory\");\n#else\n NVTE_DEVICE_ERROR(\"cp_async_bulk_tensor_4d_global_to_shared requires SM 10.0+.\");\n#endif\n}\n\n__device__ __forceinline__ void cp_async_bulk_tensor_4d_shared_to_global(\n const CUtensorMap *tensor_map, uint32_t c0, uint32_t c1, uint32_t c2, uint32_t c3,\n void *src_shmem) {\n#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n uint32_t src = __cvta_generic_to_shared(src_shmem);\n asm volatile(\n \"cp.async.bulk.tensor.4d.global.shared::cta.bulk_group\"\n \" [%0, {%1, %2, %3, %4}], [%5];\" ::\"l\"(tensor_map),\n \"r\"(c0), \"r\"(c1), \"r\"(c2), \"r\"(c3), \"r\"(src)\n : \"memory\");\n#else\n NVTE_DEVICE_ERROR(\"cp_async_bulk_tensor_4d_shared_to_global requires SM 10.0+.\");\n#endif\n}\n\nstatic void create_4D_tensor_map(CUtensorMap &tensorMap, void *dataPtr, DType dtype, uint64_t dim0,\n uint64_t dim1, uint64_t dim2, uint64_t dim3, uint32_t box0,\n uint32_t box1, uint32_t box2, uint32_t box3) {\n cuda_driver::ensure_context_exists();\n static PFN_cuTensorMapEncodeTiled_v12000 cuDriverTensorMapEncodeTiled = []() {\n void *ptr = cuda_driver::get_symbol(\"cuTensorMapEncodeTiled\");\n return reinterpret_cast(ptr);\n }();\n\n CUtensorMapDataType tma_dtype;\n size_t elem_bytes;\n switch (dtype) {\n case DType::kFloat16:\n tma_dtype = CU_TENSOR_MAP_DATA_TYPE_FLOAT16;\n elem_bytes = 2;\n break;\n case DType::kBFloat16:\n tma_dtype = CU_TENSOR_MAP_DATA_TYPE_BFLOAT16;\n elem_bytes = 2;\n break;\n case DType::kFloat8E4M3:\n case DType::kFloat8E5M2:\n case DType::kFloat8E8M0:\n case DType::kByte:\n tma_dtype = CU_TENSOR_MAP_DATA_TYPE_UINT8;\n elem_bytes = 1;\n break;\n default:\n NVTE_ERROR(\"create_4D_tensor_map: unsupported dtype \", to_string(static_cast(dtype)));\n }\n\n constexpr uint32_t rank = 4;\n uint64_t size[rank] = {dim0, dim1, dim2, dim3};\n uint64_t stride[rank - 1] = {\n dim0 * elem_bytes,\n dim0 * dim1 * elem_bytes,\n dim0 * dim1 * dim2 * elem_bytes,\n };\n uint32_t boxSize[rank] = {box0, box1, box2, box3};\n uint32_t elemStride[rank] = {1, 1, 1, 1};\n\n const auto oob_fill = (tma_dtype == CU_TENSOR_MAP_DATA_TYPE_UINT8)\n ? CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE\n : CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA;\n\n NVTE_CHECK_CUDA_DRIVER(cuDriverTensorMapEncodeTiled(\n &tensorMap, tma_dtype, rank, dataPtr, size, stride, boxSize, elemStride,\n CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE,\n oob_fill));\n}\n\ntemplate \n__device__ __forceinline__ void issue_tma_load_strided(T *smem_buf, const CUtensorMap *tma,\n size_t h_i, size_t s_tile, size_t b_i,\n uint64_t *mbar, size_t tile_bytes) {\n ptx::mbarrier_arrive_expect_tx(mbar, static_cast(tile_bytes));\n if constexpr (kIsBshd) {\n cp_async_bulk_tensor_4d_global_to_shared(smem_buf, tma, 0, static_cast(h_i),\n static_cast(s_tile),\n static_cast(b_i), mbar);\n } else {\n cp_async_bulk_tensor_4d_global_to_shared(smem_buf, tma, 0, static_cast(h_i),\n static_cast(b_i),\n static_cast(s_tile), mbar);\n }\n}\n\n__device__ __forceinline__ void st_global_cs_uint4(uint4 *ptr, uint4 val) {\n asm volatile(\"st.global.cs.v4.b32 [%0], {%1, %2, %3, %4};\" ::\"l\"(ptr), \"r\"(val.x), \"r\"(val.y),\n \"r\"(val.z), \"r\"(val.w)\n : \"memory\");\n}\n// TMA loads from strided input to smem + non-temporal stores to contiguous output in gmem\n\ntemplate \n__launch_bounds__(tma_permute_threads) __global__\n void transpose_to_bhsd_kernel(const __grid_constant__ TmaMapParams tma_maps,\n PermuteParams params, size_t b, size_t h_grid,\n unsigned int permute_s_splits, size_t s_tile_size) {\n#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n const auto &slot = params.slots[blockIdx.z];\n const CUtensorMap *tma_in = &tma_maps.maps[blockIdx.z];\n T *__restrict__ tensor_out = reinterpret_cast(slot.output);\n const size_t Sdim = slot.S;\n const size_t Hdim = slot.H;\n const size_t Ddim = slot.D_in;\n const size_t Ddim_out = slot.D_out;\n\n const size_t b_i = static_cast(blockIdx.x) / h_grid;\n const size_t h_i = static_cast(blockIdx.x) % h_grid;\n\n if (b_i >= b) return;\n if (h_i >= Hdim) return;\n\n const unsigned int s_part = blockIdx.y;\n const size_t s_begin =\n (Sdim * static_cast(s_part)) / static_cast(permute_s_splits);\n const size_t s_end =\n (Sdim * static_cast(s_part + 1)) / static_cast(permute_s_splits);\n if (s_begin >= s_end) return;\n\n const size_t out_base = b_i * Hdim * Sdim * Ddim_out + h_i * Sdim * Ddim_out;\n\n extern __shared__ __align__(128) char smem_raw[];\n T *smem = reinterpret_cast(smem_raw);\n\n __shared__ __align__(8) uint64_t mbar;\n const bool is_leader = (threadIdx.x == 0);\n\n if (is_leader) {\n ptx::mbarrier_init(&mbar, static_cast(blockDim.x));\n ptx::fence_proxy_async_shared_cta();\n }\n __syncthreads();\n\n const size_t S_TILE = s_tile_size;\n const uint32_t tile_bytes = static_cast(S_TILE * Ddim * sizeof(T));\n int parity = 0;\n\n for (size_t s_tile = s_begin; s_tile < s_end; s_tile += S_TILE) {\n const size_t tile_rows = min(S_TILE, s_end - s_tile);\n\n if (is_leader) {\n issue_tma_load_strided(smem, tma_in, h_i, s_tile, b_i, &mbar, tile_bytes);\n } else {\n ptx::mbarrier_arrive(&mbar);\n }\n\n ptx::mbarrier_wait_parity(&mbar, parity);\n parity ^= 1;\n\n T *__restrict__ out_ptr = tensor_out + out_base + s_tile * Ddim_out;\n constexpr size_t vec_elems = sizeof(uint4) / sizeof(T);\n\n if (Ddim_out == Ddim) {\n const size_t total_elems = tile_rows * Ddim;\n for (size_t i = threadIdx.x * vec_elems; i < total_elems;\n i += static_cast(blockDim.x) * vec_elems) {\n uint4 v = *reinterpret_cast(smem + i);\n st_global_cs_uint4(reinterpret_cast(out_ptr + i), v);\n }\n } else {\n const size_t total_out_elems = tile_rows * Ddim_out;\n for (size_t i = threadIdx.x * vec_elems; i < total_out_elems;\n i += static_cast(blockDim.x) * vec_elems) {\n const size_t row = i / Ddim_out;\n const size_t col = i % Ddim_out;\n uint4 v;\n if (col + vec_elems <= Ddim) {\n v = *reinterpret_cast(smem + row * Ddim + col);\n } else {\n memset(&v, 0, sizeof(v));\n const size_t smem_off = row * Ddim + col;\n size_t copy_elems = (col < Ddim) ? (Ddim - col) : 0;\n if (copy_elems > 0) memcpy(&v, smem + smem_off, copy_elems * sizeof(T));\n }\n st_global_cs_uint4(reinterpret_cast(out_ptr + i), v);\n }\n }\n\n __syncthreads();\n }\n\n if (is_leader) {\n ptx::mbarrier_invalid(&mbar);\n }\n#endif\n}\n\n// 4D TMA descriptor:\n// [B, S, H, D]: TMA dims [D, H, S, B], box [D, 1, S_TILE, 1]\n// [S, B, H, D]: TMA dims [D, H, B, S], box [D, 1, 1, S_TILE]\n\nstatic void create_strided_tensor_map(CUtensorMap &map, void *ptr, DType dtype, size_t b, size_t s,\n size_t h, size_t d, size_t s_tile, bool is_bshd) {\n if (is_bshd) {\n create_4D_tensor_map(map, ptr, dtype, static_cast(d), static_cast(h),\n static_cast(s), static_cast(b),\n static_cast(d), 1, static_cast(s_tile), 1);\n } else {\n create_4D_tensor_map(map, ptr, dtype, static_cast(d), static_cast(h),\n static_cast(b), static_cast(s),\n static_cast(d), 1, 1, static_cast(s_tile));\n }\n}\n\nvoid multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_tensors,\n NVTE_QKV_Format original_format, cudaStream_t stream) {\n using namespace transformer_engine;\n if (num_tensors == 0) return;\n NVTE_CHECK(num_tensors <= static_cast(kMaxPermuteTensors), \"num_tensors must be in [1, \",\n kMaxPermuteTensors, \"], got \", num_tensors, \".\");\n\n const bool is_bshd = (original_format == NVTE_QKV_Format::NVTE_BSHD);\n const DType dtype = inputs[0].dtype();\n const size_t elem_size = typeToSize(dtype);\n const size_t b = outputs[0].shape()[0];\n\n PermuteParams params{};\n size_t s_max = 0, h_max = 0, s_min = SIZE_MAX;\n size_t d_in_max = 0, d_out_max = 0;\n bool any_not_vec_aligned = false;\n bool all_tma_ok = true;\n\n for (size_t i = 0; i < num_tensors; ++i) {\n const size_t H = outputs[i].shape()[1];\n const size_t S = outputs[i].shape()[2];\n const size_t D_in = inputs[i].shape()[inputs[i].shape().size() - 1];\n const size_t D_out = outputs[i].shape()[3];\n params.slots[i] = {inputs[i].data.dptr, outputs[i].data.dptr, S, H, D_in, D_out};\n s_max = std::max(s_max, S);\n h_max = std::max(h_max, H);\n s_min = std::min(s_min, S);\n d_in_max = std::max(d_in_max, D_in);\n d_out_max = std::max(d_out_max, D_out);\n if ((D_in * elem_size) % 4 != 0) any_not_vec_aligned = true;\n const size_t inner = D_in * elem_size;\n if (inner < 32 || inner % 16 != 0) all_tma_ok = false;\n }\n\n if (all_tma_ok) {\n const int sm = cuda::sm_arch(cuda::current_device());\n if (sm < 100) {\n all_tma_ok = false;\n } else {\n switch (dtype) {\n case DType::kFloat16:\n case DType::kBFloat16:\n case DType::kFloat8E4M3:\n case DType::kFloat8E5M2:\n case DType::kFloat8E8M0:\n case DType::kByte:\n break;\n default:\n all_tma_ok = false;\n }\n }\n }\n\n // Dispatch order:\n // 1. TMA path: SM 10.0+, D_in*elem >= 32 && 16-aligned, supported dtype,\n // and s_tile*D_in*elem is uint4-aligned.\n // 2. Fallback path (vec-aligned): vectorized loads/stores when D_in*elem % 4 == 0.\n // 3. Fallback path (not-vec-aligned): shared-memory transpose when D_in*elem % 4 != 0.\n if (all_tma_ok) {\n const size_t s_tile = std::min(static_cast(tma_permute_s_tile_default), s_min);\n bool tma_aligned = true;\n for (size_t i = 0; i < num_tensors && tma_aligned; ++i) {\n if ((s_tile * params.slots[i].D_in * elem_size) % sizeof(uint4) != 0) tma_aligned = false;\n }\n\n if (tma_aligned) {\n TmaMapParams tma_maps{};\n for (size_t i = 0; i < num_tensors; ++i) {\n const auto &slot = params.slots[i];\n create_strided_tensor_map(tma_maps.maps[i], const_cast(slot.input), dtype, b,\n slot.S, slot.H, slot.D_in, s_tile, is_bshd);\n }\n\n const unsigned int permute_s_splits = std::max(1u, static_cast(s_min / s_tile));\n dim3 grid(static_cast(b * h_max), permute_s_splits,\n static_cast(num_tensors));\n const size_t smem_bytes = s_tile * d_in_max * elem_size;\n\n if (is_bshd) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n dtype, dtype_t, auto kernel = transpose_to_bhsd_kernel;\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes));\n kernel<<>>(tma_maps, params, b, h_max,\n permute_s_splits, s_tile););\n } else {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n dtype, dtype_t, auto kernel = transpose_to_bhsd_kernel;\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes));\n kernel<<>>(tma_maps, params, b, h_max,\n permute_s_splits, s_tile););\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n return;\n }\n }\n\n if (!any_not_vec_aligned) {\n const unsigned int permute_s_splits = std::max(\n 1u, static_cast(s_min / static_cast(fallback_permute_threads)));\n dim3 grid(static_cast(b * h_max), permute_s_splits,\n static_cast(num_tensors));\n\n if (is_bshd) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n dtype, dtype_t,\n transpose_to_bhsd_fallback_vec_aligned_kernel\n <<>>(params, b, permute_s_splits, h_max););\n } else {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n dtype, dtype_t,\n transpose_to_bhsd_fallback_vec_aligned_kernel\n <<>>(params, b, permute_s_splits, h_max););\n }\n } else {\n const unsigned int st =\n static_cast((s_max + TRANSPOSE_TILE - 1) / TRANSPOSE_TILE);\n const unsigned int ht =\n static_cast((h_max + TRANSPOSE_TILE - 1) / TRANSPOSE_TILE);\n dim3 grid(static_cast(b) * st, ht, static_cast(num_tensors));\n const size_t D_pad = (d_in_max * elem_size + 3u) & ~size_t(3);\n const size_t smem_bytes =\n static_cast(TRANSPOSE_TILE) * (static_cast(TRANSPOSE_TILE) * D_pad + 4);\n\n if (is_bshd) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n dtype, dtype_t,\n transpose_to_bhsd_fallback_not_vec_aligned_kernel\n <<>>(params, b, st););\n } else {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n dtype, dtype_t,\n transpose_to_bhsd_fallback_not_vec_aligned_kernel\n <<>>(params, b, st););\n }\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n} // namespace multi_tensor_transpose_to_bhsd\n\n// ===================================================================================\n// multi_tensor_pad_last_dim: pad the last dim of multiple tensors to certain alignment\n// ===================================================================================\n\nnamespace multi_tensor_pad_last_dim {\n\nconstexpr int pad_threads_per_block = 256;\nconstexpr int kMaxPadTensors = 16;\n\nstruct PadLastDimArgs {\n const uint8_t *input;\n uint32_t *output;\n size_t n_uint32;\n uint32_t in_row_bytes;\n uint32_t out_row_uint32;\n};\n\nstruct MultiPadParams {\n PadLastDimArgs tensors[kMaxPadTensors];\n};\n\n__launch_bounds__(pad_threads_per_block) __global__\n void multi_tensor_pad_last_dim_kernel(MultiPadParams params) {\n const auto &a = params.tensors[blockIdx.y];\n\n for (size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < a.n_uint32;\n idx += static_cast(gridDim.x) * blockDim.x) {\n const uint32_t col_byte = (idx % a.out_row_uint32) * 4;\n const size_t row = idx / a.out_row_uint32;\n const uint8_t *__restrict__ src = a.input + row * static_cast(a.in_row_bytes);\n\n uint32_t val;\n if (col_byte + 4 <= a.in_row_bytes) {\n memcpy(&val, src + col_byte, 4);\n } else if (col_byte >= a.in_row_bytes) {\n val = 0;\n } else {\n val = 0;\n memcpy(&val, src + col_byte, a.in_row_bytes - col_byte);\n }\n a.output[idx] = val;\n }\n}\n\nvoid launch_pad_batch(MultiPadParams ¶ms, int kernel_count, size_t max_n_uint32,\n cudaStream_t stream) {\n if (kernel_count == 0) return;\n constexpr int threads = pad_threads_per_block;\n const int blocks_x = static_cast(\n std::min(DIVUP(max_n_uint32, static_cast(threads)), static_cast(65535)));\n dim3 grid(blocks_x, kernel_count);\n multi_tensor_pad_last_dim_kernel<<>>(params);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid multi_tensor_pad_last_dim(Tensor *inputs, Tensor *outputs, size_t num_tensors,\n cudaStream_t stream) {\n using namespace transformer_engine;\n\n if (num_tensors == 0) return;\n\n MultiPadParams params{};\n size_t max_n_uint32 = 0;\n int kernel_count = 0;\n\n for (size_t i = 0; i < num_tensors; ++i) {\n auto &inp = inputs[i];\n auto &out = outputs[i];\n\n NVTE_CHECK(inp.data.shape.size() == 2, \"Expected 2D input tensor at index \", i, \".\");\n NVTE_CHECK(out.data.shape.size() == 2, \"Expected 2D output tensor at index \", i, \".\");\n NVTE_CHECK(inp.data.dtype == out.data.dtype, \"Dtype mismatch at index \", i, \".\");\n\n const size_t rows = inp.data.shape[0];\n const size_t in_cols = inp.data.shape[1];\n const size_t out_cols = out.data.shape[1];\n\n NVTE_CHECK(out.data.shape[0] == rows, \"Row count mismatch at index \", i, \".\");\n NVTE_CHECK(out_cols >= in_cols, \"out_cols < in_cols at index \", i, \".\");\n\n if (rows == 0) continue;\n\n if (in_cols == out_cols) {\n const size_t total_bytes = rows * in_cols * typeToSize(inp.data.dtype);\n NVTE_CHECK_CUDA(cudaMemcpyAsync(out.data.dptr, inp.data.dptr, total_bytes,\n cudaMemcpyDeviceToDevice, stream));\n continue;\n }\n\n if (kernel_count == kMaxPadTensors) {\n launch_pad_batch(params, kernel_count, max_n_uint32, stream);\n params = MultiPadParams{};\n kernel_count = 0;\n max_n_uint32 = 0;\n }\n\n const size_t elem_size = typeToSize(inp.data.dtype);\n const auto in_row_bytes = static_cast(in_cols * elem_size);\n const auto out_row_bytes = static_cast(out_cols * elem_size);\n NVTE_CHECK(out_row_bytes % 4 == 0, \"Padded row size in bytes (\", out_row_bytes,\n \") must be a multiple of 4.\");\n\n const uint32_t out_row_uint32 = out_row_bytes / 4;\n const size_t n_uint32 = rows * out_row_uint32;\n\n params.tensors[kernel_count] = {reinterpret_cast(inp.data.dptr),\n reinterpret_cast(out.data.dptr), n_uint32,\n in_row_bytes, out_row_uint32};\n max_n_uint32 = std::max(max_n_uint32, n_uint32);\n ++kernel_count;\n }\n\n launch_pad_batch(params, kernel_count, max_n_uint32, stream);\n}\n\n} // namespace multi_tensor_pad_last_dim\n} // namespace transformer_engine\n\nvoid nvte_prepare_flash_attn_fwd(NVTETensor qkvi, NVTETensor qkv, cudaStream_t stream) {\n NVTE_API_CALL(nvte_prepare_flash_attn_fwd);\n using namespace transformer_engine;\n\n flash_attention::prepare_flash_attn_fwd(*convertNVTETensorCheck(qkvi),\n *convertNVTETensorCheck(qkv), stream);\n}\n\nvoid nvte_prepare_flash_attn_bwd(NVTETensor q, NVTETensor k, NVTETensor v, NVTETensor qkv,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_prepare_flash_attn_bwd);\n using namespace transformer_engine;\n\n flash_attention::prepare_flash_attn_bwd(*convertNVTETensorCheck(q), *convertNVTETensorCheck(k),\n *convertNVTETensorCheck(v), *convertNVTETensorCheck(qkv),\n stream);\n}\n\nvoid nvte_multi_tensor_transpose_to_bhsd(NVTETensor *inputs, NVTETensor *outputs,\n size_t num_tensors, NVTE_QKV_Format original_format,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_transpose_to_bhsd);\n NVTE_CHECK(original_format == NVTE_QKV_Format::NVTE_BSHD ||\n original_format == NVTE_QKV_Format::NVTE_SBHD,\n \"nvte_multi_tensor_transpose_to_bhsd: only BSHD/SBHD -> BHSD is currently \"\n \"supported.\");\n using namespace transformer_engine;\n\n std::vector in_vec(num_tensors), out_vec(num_tensors);\n for (size_t i = 0; i < num_tensors; ++i) {\n in_vec[i] = *convertNVTETensorCheck(inputs[i]);\n out_vec[i] = *convertNVTETensorCheck(outputs[i]);\n }\n constexpr size_t kBatch = multi_tensor_transpose_to_bhsd::kMaxPermuteTensors;\n for (size_t offset = 0; offset < num_tensors; offset += kBatch) {\n const size_t batch = std::min(num_tensors - offset, kBatch);\n multi_tensor_transpose_to_bhsd::multi_tensor_transpose_to_bhsd(\n in_vec.data() + offset, out_vec.data() + offset, batch, original_format, stream);\n }\n}\n\nvoid nvte_multi_tensor_pad_last_dim(NVTETensor *inputs, NVTETensor *outputs, size_t num_tensors,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_pad_last_dim);\n using namespace transformer_engine;\n\n std::vector in_vec(num_tensors), out_vec(num_tensors);\n for (size_t i = 0; i < num_tensors; ++i) {\n in_vec[i] = *convertNVTETensorCheck(inputs[i]);\n out_vec[i] = *convertNVTETensorCheck(outputs[i]);\n }\n multi_tensor_pad_last_dim::multi_tensor_pad_last_dim(in_vec.data(), out_vec.data(), num_tensors,\n stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/cuda_driver.h\"\n#include \"../util/hip/hip_runtime.h\"\n#include \"../util/ptx.cuh\"\n#include \"../utils.cuh\"\n#include \"transformer_engine/fused_attn.h\"\n\nnamespace transformer_engine {\n\n// ============================================================================\n// prepare_flash_attn: SBH3D <-> BSHD_BSHD_BSHD for the FlashAttention backend\n// ============================================================================\n\nnamespace flash_attention {\n\n/// Packed vector of N elements of T; alignment matches a single wide load/store of N * sizeof(T) bytes.\ntemplate \nstruct alignas(sizeof(T) * N) Vec {\n T data[N];\n};\n\nconstexpr int warp_size = 32;\nconstexpr int type_size = 2; // FP16 or BF16\nconstexpr int nvec = sizeof(uint64_t) / type_size;\nconstexpr int load_size = warp_size * nvec;\nconstexpr int block_size = 512;\n\ntemplate \n__launch_bounds__(block_size) __global__\n void prepare_kernel_fwd(const T *qkvi, T *qkv, const size_t B, const size_t S, const size_t Z,\n const size_t W) {\n const int warpid = (blockDim.x * blockIdx.x + threadIdx.x) / warp_size;\n const int id_in_warp = threadIdx.x % warp_size;\n const size_t offset_input = blockIdx.y * W + warpid * 3 * W * Z + id_in_warp * nvec;\n const T *my_input = qkvi + offset_input;\n\n const size_t s = warpid / B;\n if (s >= S) return;\n\n const size_t b = warpid % B;\n\n const size_t offset_output = blockIdx.y * B * S * Z * W + (s + b * S) * W * Z + id_in_warp * nvec;\n\n T *my_output = qkv + offset_output;\n\n for (int i = 0; i < Z; ++i) {\n Vec *const out = reinterpret_cast *>(my_output + i * load_size);\n *out = *reinterpret_cast *>(my_input + i * load_size * 3);\n }\n}\n\ntemplate \n__launch_bounds__(block_size) __global__\n void prepare_kernel_bwd(const T *q, const T *k, const T *v, T *qkv, const size_t B,\n const size_t S, const size_t Z, const size_t W) {\n const T *input = blockIdx.y == 0 ? q : (blockIdx.y == 1 ? k : v);\n\n const int warpid = (blockDim.x * blockIdx.x + threadIdx.x) / warp_size;\n const int id_in_warp = threadIdx.x % warp_size;\n const size_t offset_input = warpid * W * Z + id_in_warp * nvec;\n const T *my_input = input + offset_input;\n\n const size_t b = warpid / S;\n if (b >= B) return;\n\n const size_t s = warpid % S;\n\n const size_t offset_output = (b + s * B) * 3 * W * Z + id_in_warp * nvec + blockIdx.y * W;\n\n T *my_output = qkv + offset_output;\n\n for (int i = 0; i < Z; ++i) {\n Vec *const out = reinterpret_cast *>(my_output + i * load_size * 3);\n *out = *reinterpret_cast *>(my_input + i * load_size);\n }\n}\n\nvoid prepare_flash_attn_fwd(Tensor qkvi, Tensor qkv, hipStream_t stream) {\n using namespace transformer_engine;\n\n NVTE_CHECK(qkvi.dim() == 4, \"Expected 4-dim tensor.\");\n NVTE_CHECK(qkvi.dtype() == DType::kFloat16 || qkvi.dtype() == DType::kBFloat16);\n\n auto qkvi_shape = qkvi.shape();\n\n NVTE_CHECK(qkvi_shape[3] % load_size == 0);\n NVTE_CHECK(qkvi_shape[3] == load_size);\n\n // [s, b, n, h * 3] -> [3, b, s, n, h]\n std::vector shape = {3, qkvi_shape[1], qkvi_shape[0], qkvi_shape[2], qkvi_shape[3]};\n\n size_t warps = qkvi_shape[0] * qkvi_shape[1];\n size_t warps_per_block = block_size / warp_size;\n size_t blocks = (warps + warps_per_block - 1) / warps_per_block;\n dim3 grid(blocks, 3);\n int threads = block_size;\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n qkvi.dtype(), dtype,\n prepare_kernel_fwd<<>>(\n reinterpret_cast(qkvi.data.dptr), reinterpret_cast(qkv.data.dptr),\n shape[1], shape[2], shape[3], shape[4]););\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid prepare_flash_attn_bwd(Tensor q, Tensor k, Tensor v, Tensor qkv, hipStream_t stream) {\n using namespace transformer_engine;\n\n NVTE_CHECK(q.dim() == 4, \"Expected 4-dim tensor.\");\n NVTE_CHECK(k.dim() == 4, \"Expected 4-dim tensor.\");\n NVTE_CHECK(v.dim() == 4, \"Expected 4-dim tensor.\");\n NVTE_CHECK(q.dtype() == DType::kFloat16 || q.dtype() == DType::kBFloat16);\n NVTE_CHECK(k.dtype() == q.dtype());\n NVTE_CHECK(v.dtype() == q.dtype());\n\n auto q_shape = q.shape();\n auto k_shape = k.shape();\n auto v_shape = v.shape();\n\n NVTE_CHECK(q_shape[3] % load_size == 0);\n NVTE_CHECK(q_shape[3] == load_size);\n NVTE_CHECK(k_shape[3] % load_size == 0);\n NVTE_CHECK(k_shape[3] == load_size);\n NVTE_CHECK(v_shape[3] % load_size == 0);\n NVTE_CHECK(v_shape[3] == load_size);\n\n // 3 x [s, b, n, h] -> [b, s, n, 3 * h]\n std::vector shape = {q_shape[1], q_shape[0], q_shape[2], 3 * q_shape[3]};\n\n size_t warps = q_shape[0] * q_shape[1];\n size_t warps_per_block = block_size / warp_size;\n size_t blocks = (warps + warps_per_block - 1) / warps_per_block;\n dim3 grid(blocks, 3);\n int threads = block_size;\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n q.dtype(), dtype,\n prepare_kernel_bwd<<>>(\n reinterpret_cast(q.data.dptr), reinterpret_cast(k.data.dptr),\n reinterpret_cast(v.data.dptr), reinterpret_cast(qkv.data.dptr),\n q_shape[0], q_shape[1], q_shape[2], q_shape[3]););\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n} // namespace flash_attention\n\n// ============================================================================\n// multi_tensor_transpose_to_bhsd: BSHD/SBHD -> BHSD\n// ============================================================================\n\nnamespace multi_tensor_transpose_to_bhsd {\n\nusing flash_attention::Vec;\n\nconstexpr int kMaxPermuteTensors = 16;\n\nstruct PermuteSlot {\n const void *input;\n void *output;\n size_t S, H, D_in, D_out;\n};\n\nstruct PermuteParams {\n PermuteSlot slots[kMaxPermuteTensors];\n};\n\nstruct TmaMapParams {\n CUtensorMap maps[kMaxPermuteTensors];\n};\n\n// ---------- path 3: fallback_not_vec_aligned ----------\n\n__device__ __forceinline__ void copy_row_bytes(const char *__restrict__ src, char *__restrict__ dst,\n size_t D_bytes) {\n size_t off = 0;\n for (; off + 16 <= D_bytes; off += 16) {\n uint4 tmp;\n memcpy(&tmp, src + off, 16);\n memcpy(dst + off, &tmp, 16);\n }\n for (; off + 8 <= D_bytes; off += 8) {\n uint2 tmp;\n memcpy(&tmp, src + off, 8);\n memcpy(dst + off, &tmp, 8);\n }\n for (; off + 4 <= D_bytes; off += 4) {\n unsigned int tmp;\n memcpy(&tmp, src + off, 4);\n memcpy(dst + off, &tmp, 4);\n }\n for (; off + 2 <= D_bytes; off += 2) {\n uint16_t tmp;\n memcpy(&tmp, src + off, 2);\n memcpy(dst + off, &tmp, 2);\n }\n for (; off < D_bytes; ++off) dst[off] = src[off];\n}\n\n__device__ __forceinline__ void copy_and_pad_row_bytes(const char *__restrict__ src,\n char *__restrict__ dst, size_t D_bytes,\n size_t D_out_bytes) {\n copy_row_bytes(src, dst, D_bytes);\n for (size_t off = D_bytes; off < D_out_bytes; ++off) dst[off] = 0;\n}\n\nconstexpr int TRANSPOSE_TILE = 32;\nconstexpr int TRANSPOSE_BLOCK = 256;\nconstexpr int TRANSPOSE_WARPS = TRANSPOSE_BLOCK / 32; // 8\n\ntemplate \n__launch_bounds__(TRANSPOSE_BLOCK) __global__\n void transpose_to_bhsd_fallback_not_vec_aligned_kernel(PermuteParams params, size_t b,\n unsigned int s_tiles) {\n const auto &slot = params.slots[blockIdx.z];\n const T *__restrict__ in = reinterpret_cast(slot.input);\n T *__restrict__ out = reinterpret_cast(slot.output);\n const size_t S = slot.S;\n const size_t H = slot.H;\n const size_t D = slot.D_in;\n const size_t D_out = slot.D_out;\n const size_t D_bytes = D * sizeof(T);\n const size_t D_out_bytes = D_out * sizeof(T);\n const size_t D_smem_pad = (D_bytes + 3u) & ~size_t(3);\n\n const size_t tile_s = static_cast(blockIdx.x) % static_cast(s_tiles);\n const size_t b_i = static_cast(blockIdx.x) / static_cast(s_tiles);\n if (b_i >= b) return;\n const size_t tile_h = static_cast(blockIdx.y);\n\n const size_t s_base = tile_s * TRANSPOSE_TILE;\n const size_t h_base = tile_h * TRANSPOSE_TILE;\n\n extern __shared__ char smem[];\n const size_t smem_row = static_cast(TRANSPOSE_TILE) * D_smem_pad + 4;\n\n // ---- Phase 1: global \u2192 smem (sweep consecutive H \u2192 coalesced reads) ----\n for (unsigned int warp_off = threadIdx.x >> 5; warp_off < TRANSPOSE_TILE;\n warp_off += TRANSPOSE_WARPS) {\n const size_t local_s = warp_off;\n const size_t local_h = threadIdx.x & 31u;\n const size_t s_i = s_base + local_s;\n const size_t h_i = h_base + local_h;\n if (s_i < S && h_i < H) {\n const char *__restrict__ src;\n if constexpr (kIsBshd)\n src = reinterpret_cast(in + b_i * S * H * D + s_i * H * D + h_i * D);\n else\n src = reinterpret_cast(in + s_i * b * H * D + b_i * H * D + h_i * D);\n copy_row_bytes(src, smem + local_s * smem_row + local_h * D_smem_pad, D_bytes);\n }\n }\n\n __syncthreads();\n\n // ---- Phase 2: smem \u2192 global (sweep consecutive S \u2192 coalesced writes, with padding) ----\n for (unsigned int warp_off = threadIdx.x >> 5; warp_off < TRANSPOSE_TILE;\n warp_off += TRANSPOSE_WARPS) {\n const size_t local_h = warp_off;\n const size_t local_s = threadIdx.x & 31u;\n const size_t s_i = s_base + local_s;\n const size_t h_i = h_base + local_h;\n if (s_i < S && h_i < H) {\n copy_and_pad_row_bytes(\n smem + local_s * smem_row + local_h * D_smem_pad,\n reinterpret_cast(out + b_i * H * S * D_out + h_i * S * D_out + s_i * D_out),\n D_bytes, D_out_bytes);\n }\n }\n}\n\n// ---------- path 2: fallback_vec_aligned ----------\n\nconstexpr int fallback_permute_threads = 1024;\n\ntemplate \n__device__ __forceinline__ void permute_vec_loop(const T *__restrict__ in, T *__restrict__ out,\n size_t b, size_t S, size_t H, size_t D,\n size_t D_out, size_t b_i, size_t h_i,\n size_t s_begin, size_t S_chunk) {\n const size_t out_base = b_i * H * S * D_out + h_i * S * D_out;\n const size_t d_vec = D / static_cast(N);\n const size_t total_work = S_chunk * d_vec;\n for (size_t w = static_cast(threadIdx.x); w < total_work;\n w += static_cast(blockDim.x)) {\n const size_t s_local = w / d_vec;\n const size_t s_i = s_begin + s_local;\n const size_t d_off = (w % d_vec) * static_cast(N);\n const T *__restrict__ in_ptr;\n if constexpr (kIsBshd) {\n in_ptr = in + b_i * (S * H * D) + s_i * (H * D) + h_i * D + d_off;\n } else {\n in_ptr = in + s_i * (b * H * D) + b_i * (H * D) + h_i * D + d_off;\n }\n T *__restrict__ out_ptr = out + out_base + s_i * D_out + d_off;\n *reinterpret_cast *>(out_ptr) = *reinterpret_cast *>(in_ptr);\n }\n if (D_out > D) {\n const size_t pad_elems = D_out - D;\n const size_t total_pad = S_chunk * pad_elems;\n for (size_t w = static_cast(threadIdx.x); w < total_pad;\n w += static_cast(blockDim.x)) {\n const size_t s_local = w / pad_elems;\n const size_t s_i = s_begin + s_local;\n const size_t d_off = D + (w % pad_elems);\n out[out_base + s_i * D_out + d_off] = static_cast(0.f);\n }\n }\n}\n\ntemplate \n__launch_bounds__(fallback_permute_threads) __global__\n void transpose_to_bhsd_fallback_vec_aligned_kernel(PermuteParams params, size_t b,\n unsigned int permute_s_splits,\n size_t h_grid) {\n const auto &slot = params.slots[blockIdx.z];\n const T *__restrict__ in = reinterpret_cast(slot.input);\n T *__restrict__ out = reinterpret_cast(slot.output);\n const size_t S = slot.S;\n const size_t H = slot.H;\n const size_t D = slot.D_in;\n const size_t D_out = slot.D_out;\n\n const size_t b_i = static_cast(blockIdx.x) / h_grid;\n const size_t h_i = static_cast(blockIdx.x) % h_grid;\n if (b_i >= b) return;\n if (h_i >= H) return;\n\n const unsigned int s_part = blockIdx.y;\n const size_t s_begin = (S * static_cast(s_part)) / static_cast(permute_s_splits);\n const size_t s_end =\n (S * static_cast(s_part + 1)) / static_cast(permute_s_splits);\n if (s_begin >= s_end) return;\n const size_t S_chunk = s_end - s_begin;\n\n const size_t D_bytes = D * sizeof(T);\n\n if (D_bytes % 16 == 0) {\n constexpr size_t N = 16 / sizeof(T);\n permute_vec_loop(in, out, b, S, H, D, D_out, b_i, h_i, s_begin, S_chunk);\n return;\n }\n if (D_bytes % 8 == 0) {\n constexpr size_t N = 8 / sizeof(T);\n permute_vec_loop(in, out, b, S, H, D, D_out, b_i, h_i, s_begin, S_chunk);\n return;\n }\n if constexpr (sizeof(T) <= 4) {\n if (D_bytes % 4 == 0) {\n constexpr size_t N = 4 / sizeof(T);\n permute_vec_loop(in, out, b, S, H, D, D_out, b_i, h_i, s_begin, S_chunk);\n return;\n }\n }\n}\n\n// ---------- path 1: TMA ----------\n\nconstexpr int tma_permute_threads = 128;\nconstexpr int tma_permute_s_tile_default = 32;\n\n__device__ __forceinline__ void cp_async_bulk_tensor_4d_global_to_shared(\n void *dst_shmem, const CUtensorMap *tensor_map, uint32_t c0, uint32_t c1, uint32_t c2,\n uint32_t c3, uint64_t *mbar) {\n#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n uint32_t dst = __cvta_generic_to_shared(dst_shmem);\n uint32_t bar = __cvta_generic_to_shared(mbar);\n asm volatile(\n \"cp.async.bulk.tensor.4d.shared::cluster.global.tile\"\n \".mbarrier::complete_tx::bytes [%0], [%1, {%2, %3, %4, %5}], [%6];\" ::\"r\"(dst),\n \"l\"(tensor_map), \"r\"(c0), \"r\"(c1), \"r\"(c2), \"r\"(c3), \"r\"(bar)\n : \"memory\");\n#else\n NVTE_DEVICE_ERROR(\"cp_async_bulk_tensor_4d_global_to_shared requires SM 10.0+.\");\n#endif\n}\n\n__device__ __forceinline__ void cp_async_bulk_tensor_4d_shared_to_global(\n const CUtensorMap *tensor_map, uint32_t c0, uint32_t c1, uint32_t c2, uint32_t c3,\n void *src_shmem) {\n#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n uint32_t src = __cvta_generic_to_shared(src_shmem);\n asm volatile(\n \"cp.async.bulk.tensor.4d.global.shared::cta.bulk_group\"\n \" [%0, {%1, %2, %3, %4}], [%5];\" ::\"l\"(tensor_map),\n \"r\"(c0), \"r\"(c1), \"r\"(c2), \"r\"(c3), \"r\"(src)\n : \"memory\");\n#else\n NVTE_DEVICE_ERROR(\"cp_async_bulk_tensor_4d_shared_to_global requires SM 10.0+.\");\n#endif\n}\n\nstatic void create_4D_tensor_map(CUtensorMap &tensorMap, void *dataPtr, DType dtype, uint64_t dim0,\n uint64_t dim1, uint64_t dim2, uint64_t dim3, uint32_t box0,\n uint32_t box1, uint32_t box2, uint32_t box3) {\n cuda_driver::ensure_context_exists();\n static PFN_cuTensorMapEncodeTiled_v12000 cuDriverTensorMapEncodeTiled = []() {\n void *ptr = cuda_driver::get_symbol(\"cuTensorMapEncodeTiled\");\n return reinterpret_cast(ptr);\n }();\n\n CUtensorMapDataType tma_dtype;\n size_t elem_bytes;\n switch (dtype) {\n case DType::kFloat16:\n tma_dtype = CU_TENSOR_MAP_DATA_TYPE_FLOAT16;\n elem_bytes = 2;\n break;\n case DType::kBFloat16:\n tma_dtype = CU_TENSOR_MAP_DATA_TYPE_BFLOAT16;\n elem_bytes = 2;\n break;\n case DType::kFloat8E4M3:\n case DType::kFloat8E5M2:\n case DType::kFloat8E8M0:\n case DType::kByte:\n tma_dtype = CU_TENSOR_MAP_DATA_TYPE_UINT8;\n elem_bytes = 1;\n break;\n default:\n NVTE_ERROR(\"create_4D_tensor_map: unsupported dtype \", to_string(static_cast(dtype)));\n }\n\n constexpr uint32_t rank = 4;\n uint64_t size[rank] = {dim0, dim1, dim2, dim3};\n uint64_t stride[rank - 1] = {\n dim0 * elem_bytes,\n dim0 * dim1 * elem_bytes,\n dim0 * dim1 * dim2 * elem_bytes,\n };\n uint32_t boxSize[rank] = {box0, box1, box2, box3};\n uint32_t elemStride[rank] = {1, 1, 1, 1};\n\n const auto oob_fill = (tma_dtype == CU_TENSOR_MAP_DATA_TYPE_UINT8)\n ? CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE\n : CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA;\n\n NVTE_CHECK_CUDA_DRIVER(cuDriverTensorMapEncodeTiled(\n &tensorMap, tma_dtype, rank, dataPtr, size, stride, boxSize, elemStride,\n CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE,\n oob_fill));\n}\n\ntemplate \n__device__ __forceinline__ void issue_tma_load_strided(T *smem_buf, const CUtensorMap *tma,\n size_t h_i, size_t s_tile, size_t b_i,\n uint64_t *mbar, size_t tile_bytes) {\n ptx::mbarrier_arrive_expect_tx(mbar, static_cast(tile_bytes));\n if constexpr (kIsBshd) {\n cp_async_bulk_tensor_4d_global_to_shared(smem_buf, tma, 0, static_cast(h_i),\n static_cast(s_tile),\n static_cast(b_i), mbar);\n } else {\n cp_async_bulk_tensor_4d_global_to_shared(smem_buf, tma, 0, static_cast(h_i),\n static_cast(b_i),\n static_cast(s_tile), mbar);\n }\n}\n\n__device__ __forceinline__ void st_global_cs_uint4(uint4 *ptr, uint4 val) {\n asm volatile(\"st.global.cs.v4.b32 [%0], {%1, %2, %3, %4};\" ::\"l\"(ptr), \"r\"(val.x), \"r\"(val.y),\n \"r\"(val.z), \"r\"(val.w)\n : \"memory\");\n}\n// TMA loads from strided input to smem + non-temporal stores to contiguous output in gmem\n\ntemplate \n__launch_bounds__(tma_permute_threads) __global__\n void transpose_to_bhsd_kernel(const __grid_constant__ TmaMapParams tma_maps,\n PermuteParams params, size_t b, size_t h_grid,\n unsigned int permute_s_splits, size_t s_tile_size) {\n#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n const auto &slot = params.slots[blockIdx.z];\n const CUtensorMap *tma_in = &tma_maps.maps[blockIdx.z];\n T *__restrict__ tensor_out = reinterpret_cast(slot.output);\n const size_t Sdim = slot.S;\n const size_t Hdim = slot.H;\n const size_t Ddim = slot.D_in;\n const size_t Ddim_out = slot.D_out;\n\n const size_t b_i = static_cast(blockIdx.x) / h_grid;\n const size_t h_i = static_cast(blockIdx.x) % h_grid;\n\n if (b_i >= b) return;\n if (h_i >= Hdim) return;\n\n const unsigned int s_part = blockIdx.y;\n const size_t s_begin =\n (Sdim * static_cast(s_part)) / static_cast(permute_s_splits);\n const size_t s_end =\n (Sdim * static_cast(s_part + 1)) / static_cast(permute_s_splits);\n if (s_begin >= s_end) return;\n\n const size_t out_base = b_i * Hdim * Sdim * Ddim_out + h_i * Sdim * Ddim_out;\n\n extern __shared__ __align__(128) char smem_raw[];\n T *smem = reinterpret_cast(smem_raw);\n\n __shared__ __align__(8) uint64_t mbar;\n const bool is_leader = (threadIdx.x == 0);\n\n if (is_leader) {\n ptx::mbarrier_init(&mbar, static_cast(blockDim.x));\n ptx::fence_proxy_async_shared_cta();\n }\n __syncthreads();\n\n const size_t S_TILE = s_tile_size;\n const uint32_t tile_bytes = static_cast(S_TILE * Ddim * sizeof(T));\n int parity = 0;\n\n for (size_t s_tile = s_begin; s_tile < s_end; s_tile += S_TILE) {\n const size_t tile_rows = min(S_TILE, s_end - s_tile);\n\n if (is_leader) {\n issue_tma_load_strided(smem, tma_in, h_i, s_tile, b_i, &mbar, tile_bytes);\n } else {\n ptx::mbarrier_arrive(&mbar);\n }\n\n ptx::mbarrier_wait_parity(&mbar, parity);\n parity ^= 1;\n\n T *__restrict__ out_ptr = tensor_out + out_base + s_tile * Ddim_out;\n constexpr size_t vec_elems = sizeof(uint4) / sizeof(T);\n\n if (Ddim_out == Ddim) {\n const size_t total_elems = tile_rows * Ddim;\n for (size_t i = threadIdx.x * vec_elems; i < total_elems;\n i += static_cast(blockDim.x) * vec_elems) {\n uint4 v = *reinterpret_cast(smem + i);\n st_global_cs_uint4(reinterpret_cast(out_ptr + i), v);\n }\n } else {\n const size_t total_out_elems = tile_rows * Ddim_out;\n for (size_t i = threadIdx.x * vec_elems; i < total_out_elems;\n i += static_cast(blockDim.x) * vec_elems) {\n const size_t row = i / Ddim_out;\n const size_t col = i % Ddim_out;\n uint4 v;\n if (col + vec_elems <= Ddim) {\n v = *reinterpret_cast(smem + row * Ddim + col);\n } else {\n memset(&v, 0, sizeof(v));\n const size_t smem_off = row * Ddim + col;\n size_t copy_elems = (col < Ddim) ? (Ddim - col) : 0;\n if (copy_elems > 0) memcpy(&v, smem + smem_off, copy_elems * sizeof(T));\n }\n st_global_cs_uint4(reinterpret_cast(out_ptr + i), v);\n }\n }\n\n __syncthreads();\n }\n\n if (is_leader) {\n ptx::mbarrier_invalid(&mbar);\n }\n#endif\n}\n\n// 4D TMA descriptor:\n// [B, S, H, D]: TMA dims [D, H, S, B], box [D, 1, S_TILE, 1]\n// [S, B, H, D]: TMA dims [D, H, B, S], box [D, 1, 1, S_TILE]\n\nstatic void create_strided_tensor_map(CUtensorMap &map, void *ptr, DType dtype, size_t b, size_t s,\n size_t h, size_t d, size_t s_tile, bool is_bshd) {\n if (is_bshd) {\n create_4D_tensor_map(map, ptr, dtype, static_cast(d), static_cast(h),\n static_cast(s), static_cast(b),\n static_cast(d), 1, static_cast(s_tile), 1);\n } else {\n create_4D_tensor_map(map, ptr, dtype, static_cast(d), static_cast(h),\n static_cast(b), static_cast(s),\n static_cast(d), 1, 1, static_cast(s_tile));\n }\n}\n\nvoid multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_tensors,\n NVTE_QKV_Format original_format, hipStream_t stream) {\n using namespace transformer_engine;\n if (num_tensors == 0) return;\n NVTE_CHECK(num_tensors <= static_cast(kMaxPermuteTensors), \"num_tensors must be in [1, \",\n kMaxPermuteTensors, \"], got \", num_tensors, \".\");\n\n const bool is_bshd = (original_format == NVTE_QKV_Format::NVTE_BSHD);\n const DType dtype = inputs[0].dtype();\n const size_t elem_size = typeToSize(dtype);\n const size_t b = outputs[0].shape()[0];\n\n PermuteParams params{};\n size_t s_max = 0, h_max = 0, s_min = SIZE_MAX;\n size_t d_in_max = 0, d_out_max = 0;\n bool any_not_vec_aligned = false;\n bool all_tma_ok = true;\n\n for (size_t i = 0; i < num_tensors; ++i) {\n const size_t H = outputs[i].shape()[1];\n const size_t S = outputs[i].shape()[2];\n const size_t D_in = inputs[i].shape()[inputs[i].shape().size() - 1];\n const size_t D_out = outputs[i].shape()[3];\n params.slots[i] = {inputs[i].data.dptr, outputs[i].data.dptr, S, H, D_in, D_out};\n s_max = std::max(s_max, S);\n h_max = std::max(h_max, H);\n s_min = std::min(s_min, S);\n d_in_max = std::max(d_in_max, D_in);\n d_out_max = std::max(d_out_max, D_out);\n if ((D_in * elem_size) % 4 != 0) any_not_vec_aligned = true;\n const size_t inner = D_in * elem_size;\n if (inner < 32 || inner % 16 != 0) all_tma_ok = false;\n }\n\n if (all_tma_ok) {\n const int sm = cuda::sm_arch(cuda::current_device());\n if (sm < 100) {\n all_tma_ok = false;\n } else {\n switch (dtype) {\n case DType::kFloat16:\n case DType::kBFloat16:\n case DType::kFloat8E4M3:\n case DType::kFloat8E5M2:\n case DType::kFloat8E8M0:\n case DType::kByte:\n break;\n default:\n all_tma_ok = false;\n }\n }\n }\n\n // Dispatch order:\n // 1. TMA path: SM 10.0+, D_in*elem >= 32 && 16-aligned, supported dtype,\n // and s_tile*D_in*elem is uint4-aligned.\n // 2. Fallback path (vec-aligned): vectorized loads/stores when D_in*elem % 4 == 0.\n // 3. Fallback path (not-vec-aligned): shared-memory transpose when D_in*elem % 4 != 0.\n if (all_tma_ok) {\n const size_t s_tile = std::min(static_cast(tma_permute_s_tile_default), s_min);\n bool tma_aligned = true;\n for (size_t i = 0; i < num_tensors && tma_aligned; ++i) {\n if ((s_tile * params.slots[i].D_in * elem_size) % sizeof(uint4) != 0) tma_aligned = false;\n }\n\n if (tma_aligned) {\n TmaMapParams tma_maps{};\n for (size_t i = 0; i < num_tensors; ++i) {\n const auto &slot = params.slots[i];\n create_strided_tensor_map(tma_maps.maps[i], const_cast(slot.input), dtype, b,\n slot.S, slot.H, slot.D_in, s_tile, is_bshd);\n }\n\n const unsigned int permute_s_splits = std::max(1u, static_cast(s_min / s_tile));\n dim3 grid(static_cast(b * h_max), permute_s_splits,\n static_cast(num_tensors));\n const size_t smem_bytes = s_tile * d_in_max * elem_size;\n\n if (is_bshd) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n dtype, dtype_t, auto kernel = transpose_to_bhsd_kernel;\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(kernel), hipFuncAttributeMaxDynamicSharedMemorySize, smem_bytes));\n kernel<<>>(tma_maps, params, b, h_max,\n permute_s_splits, s_tile););\n } else {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n dtype, dtype_t, auto kernel = transpose_to_bhsd_kernel;\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(kernel), hipFuncAttributeMaxDynamicSharedMemorySize, smem_bytes));\n kernel<<>>(tma_maps, params, b, h_max,\n permute_s_splits, s_tile););\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n return;\n }\n }\n\n if (!any_not_vec_aligned) {\n const unsigned int permute_s_splits = std::max(\n 1u, static_cast(s_min / static_cast(fallback_permute_threads)));\n dim3 grid(static_cast(b * h_max), permute_s_splits,\n static_cast(num_tensors));\n\n if (is_bshd) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n dtype, dtype_t,\n transpose_to_bhsd_fallback_vec_aligned_kernel\n <<>>(params, b, permute_s_splits, h_max););\n } else {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n dtype, dtype_t,\n transpose_to_bhsd_fallback_vec_aligned_kernel\n <<>>(params, b, permute_s_splits, h_max););\n }\n } else {\n const unsigned int st =\n static_cast((s_max + TRANSPOSE_TILE - 1) / TRANSPOSE_TILE);\n const unsigned int ht =\n static_cast((h_max + TRANSPOSE_TILE - 1) / TRANSPOSE_TILE);\n dim3 grid(static_cast(b) * st, ht, static_cast(num_tensors));\n const size_t D_pad = (d_in_max * elem_size + 3u) & ~size_t(3);\n const size_t smem_bytes =\n static_cast(TRANSPOSE_TILE) * (static_cast(TRANSPOSE_TILE) * D_pad + 4);\n\n if (is_bshd) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n dtype, dtype_t,\n transpose_to_bhsd_fallback_not_vec_aligned_kernel\n <<>>(params, b, st););\n } else {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n dtype, dtype_t,\n transpose_to_bhsd_fallback_not_vec_aligned_kernel\n <<>>(params, b, st););\n }\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n} // namespace multi_tensor_transpose_to_bhsd\n\n// ===================================================================================\n// multi_tensor_pad_last_dim: pad the last dim of multiple tensors to certain alignment\n// ===================================================================================\n\nnamespace multi_tensor_pad_last_dim {\n\nconstexpr int pad_threads_per_block = 256;\nconstexpr int kMaxPadTensors = 16;\n\nstruct PadLastDimArgs {\n const uint8_t *input;\n uint32_t *output;\n size_t n_uint32;\n uint32_t in_row_bytes;\n uint32_t out_row_uint32;\n};\n\nstruct MultiPadParams {\n PadLastDimArgs tensors[kMaxPadTensors];\n};\n\n__launch_bounds__(pad_threads_per_block) __global__\n void multi_tensor_pad_last_dim_kernel(MultiPadParams params) {\n const auto &a = params.tensors[blockIdx.y];\n\n for (size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < a.n_uint32;\n idx += static_cast(gridDim.x) * blockDim.x) {\n const uint32_t col_byte = (idx % a.out_row_uint32) * 4;\n const size_t row = idx / a.out_row_uint32;\n const uint8_t *__restrict__ src = a.input + row * static_cast(a.in_row_bytes);\n\n uint32_t val;\n if (col_byte + 4 <= a.in_row_bytes) {\n memcpy(&val, src + col_byte, 4);\n } else if (col_byte >= a.in_row_bytes) {\n val = 0;\n } else {\n val = 0;\n memcpy(&val, src + col_byte, a.in_row_bytes - col_byte);\n }\n a.output[idx] = val;\n }\n}\n\nvoid launch_pad_batch(MultiPadParams ¶ms, int kernel_count, size_t max_n_uint32,\n hipStream_t stream) {\n if (kernel_count == 0) return;\n constexpr int threads = pad_threads_per_block;\n const int blocks_x = static_cast(\n std::min(DIVUP(max_n_uint32, static_cast(threads)), static_cast(65535)));\n dim3 grid(blocks_x, kernel_count);\n multi_tensor_pad_last_dim_kernel<<>>(params);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid multi_tensor_pad_last_dim(Tensor *inputs, Tensor *outputs, size_t num_tensors,\n hipStream_t stream) {\n using namespace transformer_engine;\n\n if (num_tensors == 0) return;\n\n MultiPadParams params{};\n size_t max_n_uint32 = 0;\n int kernel_count = 0;\n\n for (size_t i = 0; i < num_tensors; ++i) {\n auto &inp = inputs[i];\n auto &out = outputs[i];\n\n NVTE_CHECK(inp.data.shape.size() == 2, \"Expected 2D input tensor at index \", i, \".\");\n NVTE_CHECK(out.data.shape.size() == 2, \"Expected 2D output tensor at index \", i, \".\");\n NVTE_CHECK(inp.data.dtype == out.data.dtype, \"Dtype mismatch at index \", i, \".\");\n\n const size_t rows = inp.data.shape[0];\n const size_t in_cols = inp.data.shape[1];\n const size_t out_cols = out.data.shape[1];\n\n NVTE_CHECK(out.data.shape[0] == rows, \"Row count mismatch at index \", i, \".\");\n NVTE_CHECK(out_cols >= in_cols, \"out_cols < in_cols at index \", i, \".\");\n\n if (rows == 0) continue;\n\n if (in_cols == out_cols) {\n const size_t total_bytes = rows * in_cols * typeToSize(inp.data.dtype);\n NVTE_CHECK_CUDA(hipMemcpyAsync(out.data.dptr, inp.data.dptr, total_bytes,\n hipMemcpyDeviceToDevice, stream));\n continue;\n }\n\n if (kernel_count == kMaxPadTensors) {\n launch_pad_batch(params, kernel_count, max_n_uint32, stream);\n params = MultiPadParams{};\n kernel_count = 0;\n max_n_uint32 = 0;\n }\n\n const size_t elem_size = typeToSize(inp.data.dtype);\n const auto in_row_bytes = static_cast(in_cols * elem_size);\n const auto out_row_bytes = static_cast(out_cols * elem_size);\n NVTE_CHECK(out_row_bytes % 4 == 0, \"Padded row size in bytes (\", out_row_bytes,\n \") must be a multiple of 4.\");\n\n const uint32_t out_row_uint32 = out_row_bytes / 4;\n const size_t n_uint32 = rows * out_row_uint32;\n\n params.tensors[kernel_count] = {reinterpret_cast(inp.data.dptr),\n reinterpret_cast(out.data.dptr), n_uint32,\n in_row_bytes, out_row_uint32};\n max_n_uint32 = std::max(max_n_uint32, n_uint32);\n ++kernel_count;\n }\n\n launch_pad_batch(params, kernel_count, max_n_uint32, stream);\n}\n\n} // namespace multi_tensor_pad_last_dim\n} // namespace transformer_engine\n\nvoid nvte_prepare_flash_attn_fwd(NVTETensor qkvi, NVTETensor qkv, hipStream_t stream) {\n NVTE_API_CALL(nvte_prepare_flash_attn_fwd);\n using namespace transformer_engine;\n\n flash_attention::prepare_flash_attn_fwd(*convertNVTETensorCheck(qkvi),\n *convertNVTETensorCheck(qkv), stream);\n}\n\nvoid nvte_prepare_flash_attn_bwd(NVTETensor q, NVTETensor k, NVTETensor v, NVTETensor qkv,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_prepare_flash_attn_bwd);\n using namespace transformer_engine;\n\n flash_attention::prepare_flash_attn_bwd(*convertNVTETensorCheck(q), *convertNVTETensorCheck(k),\n *convertNVTETensorCheck(v), *convertNVTETensorCheck(qkv),\n stream);\n}\n\nvoid nvte_multi_tensor_transpose_to_bhsd(NVTETensor *inputs, NVTETensor *outputs,\n size_t num_tensors, NVTE_QKV_Format original_format,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_transpose_to_bhsd);\n NVTE_CHECK(original_format == NVTE_QKV_Format::NVTE_BSHD ||\n original_format == NVTE_QKV_Format::NVTE_SBHD,\n \"nvte_multi_tensor_transpose_to_bhsd: only BSHD/SBHD -> BHSD is currently \"\n \"supported.\");\n using namespace transformer_engine;\n\n std::vector in_vec(num_tensors), out_vec(num_tensors);\n for (size_t i = 0; i < num_tensors; ++i) {\n in_vec[i] = *convertNVTETensorCheck(inputs[i]);\n out_vec[i] = *convertNVTETensorCheck(outputs[i]);\n }\n constexpr size_t kBatch = multi_tensor_transpose_to_bhsd::kMaxPermuteTensors;\n for (size_t offset = 0; offset < num_tensors; offset += kBatch) {\n const size_t batch = std::min(num_tensors - offset, kBatch);\n multi_tensor_transpose_to_bhsd::multi_tensor_transpose_to_bhsd(\n in_vec.data() + offset, out_vec.data() + offset, batch, original_format, stream);\n }\n}\n\nvoid nvte_multi_tensor_pad_last_dim(NVTETensor *inputs, NVTETensor *outputs, size_t num_tensors,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_pad_last_dim);\n using namespace transformer_engine;\n\n std::vector in_vec(num_tensors), out_vec(num_tensors);\n for (size_t i = 0; i < num_tensors; ++i) {\n in_vec[i] = *convertNVTETensorCheck(inputs[i]);\n out_vec[i] = *convertNVTETensorCheck(outputs[i]);\n }\n multi_tensor_pad_last_dim::multi_tensor_pad_last_dim(in_vec.data(), out_vec.data(), num_tensors,\n stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/fused_attn/flash_attn.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"transformer_engine/fused_attn.h\"\n\nnamespace transformer_engine {\nnamespace context_parallel {\n\nstruct LseCorrectionFunctor {\n __forceinline__ __device__ static void run(float *lse, float *half_lse, size_t idx,\n size_t half_idx) {\n float val = lse[idx];\n float val_per_step = half_lse[half_idx];\n float max_scale = max(val, val_per_step);\n float min_scale = min(val, val_per_step);\n lse[idx] = max_scale + log1pf(expf(min_scale - max_scale));\n }\n};\n\nstruct ReadLseFunctor {\n __forceinline__ __device__ static void run(float *lse, float *half_lse, size_t idx,\n size_t half_idx) {\n half_lse[half_idx] = lse[idx];\n }\n};\n\nstruct EmptyFunctor {\n __forceinline__ __device__ static void run(void *token, void *token_per_step, int idx) {}\n};\n\nstruct CopyFunctor {\n __forceinline__ __device__ static void run(void *token, void *token_per_step, int idx) {\n reinterpret_cast(token)[idx] = reinterpret_cast(token_per_step)[idx];\n }\n};\n\ntemplate \nstruct AddFunctor {\n __forceinline__ __device__ static void run(dtype *token, dtype *token_per_step, int idx) {\n float4 d_ = reinterpret_cast(token)[idx];\n dtype *p_ = reinterpret_cast(&d_);\n\n float4 d = reinterpret_cast(token_per_step)[idx];\n dtype *p = reinterpret_cast(&d);\n\n#pragma unroll\n for (int i = 0; i < sizeof(float4) / sizeof(dtype); i++) {\n p_[i] = p_[i] + p[i];\n }\n\n reinterpret_cast(token)[idx] = d_;\n }\n};\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Binary search an array for a target value\n **************************************************************************************************/\n\n__forceinline__ __device__ int binary_search(int target, int *array, int len) {\n int left = 1, right = len - 1;\n while (left < right) {\n int mid = (left + right) / 2;\n if (array[mid] <= target) {\n left = mid + 1;\n } else {\n right = mid;\n }\n }\n return left - 1;\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Generate partitioned indices for input tokens\n **************************************************************************************************/\n__global__ void thd_partition_indices_kernel(int *output, int *cu_seqlens, int batch,\n int total_tokens, int world_size, int rank) {\n extern __shared__ int cu_seqlens_s[];\n for (int i = threadIdx.x; i <= batch; i += blockDim.x) {\n int seqlen = cu_seqlens[i];\n // Currently we assume that each sequence length is divisible by (world_size*2) since we have\n // to distribute each sequence evenly to different GPUs.\n assert(seqlen % (world_size * 2) == 0);\n cu_seqlens_s[i] = seqlen / world_size;\n }\n __syncthreads();\n\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n int num_threads = blockDim.x * gridDim.x;\n\n for (int token_id = tid; token_id < total_tokens / world_size; token_id += num_threads) {\n int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1);\n int seq_len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id];\n int index = token_id - cu_seqlens_s[seq_id];\n int offset = index < seq_len / 2 ? rank : (world_size - 1) * 2 - rank;\n index += cu_seqlens_s[seq_id] * world_size + seq_len / 2 * offset;\n output[token_id] = index;\n }\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Read the half of a THD tensor\n **************************************************************************************************/\n\n__global__ void thd_read_half_tensor_kernel(void *half, void *tensor, int *cu_seqlens, int batch,\n int hidden_size_in_bytes, int half_idx,\n int dim_size_of_token) {\n extern __shared__ int cu_seqlens_s[];\n for (int i = threadIdx.x; i <= batch; i += blockDim.x) {\n cu_seqlens_s[i] = cu_seqlens[i] / 2;\n }\n __syncthreads();\n\n int warpid = (blockIdx.x * blockDim.x + threadIdx.x) / 32;\n int laneid = threadIdx.x % 32;\n int num_warps = (blockDim.x * gridDim.x) / 32;\n int num_total_tokens = cu_seqlens_s[batch];\n int num_float4s_per_token = hidden_size_in_bytes / sizeof(float4);\n\n size_t offset = static_cast(dim_size_of_token) * hidden_size_in_bytes;\n half = reinterpret_cast(reinterpret_cast(half) + offset / 2 * blockIdx.y);\n tensor = reinterpret_cast(reinterpret_cast(tensor) + offset * blockIdx.y);\n\n for (int token_id = warpid; token_id < num_total_tokens; token_id += num_warps) {\n int seqid = binary_search(token_id, cu_seqlens_s, batch + 1);\n\n size_t offset_in_bytes = static_cast(token_id) * hidden_size_in_bytes;\n float4 *cur_half_token =\n reinterpret_cast(reinterpret_cast(half) + offset_in_bytes);\n\n offset_in_bytes =\n (static_cast(token_id) + cu_seqlens_s[seqid + half_idx]) * hidden_size_in_bytes;\n float4 *cur_token =\n reinterpret_cast(reinterpret_cast(tensor) + offset_in_bytes);\n\n for (int idx = laneid; idx < num_float4s_per_token; idx += 32) {\n cur_half_token[idx] = cur_token[idx];\n }\n }\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: softmax_lse related operations\n **************************************************************************************************/\n\ntemplate \n__global__ void thd_lse_kernel(float *lse, float *half_lse, int *cu_seqlens, int batch,\n int num_heads, int lse_seqlen, int second_half_lse_seqlen) {\n extern __shared__ int cu_seqlens_s[];\n for (int i = threadIdx.x; i <= batch; i += blockDim.x) {\n cu_seqlens_s[i] = cu_seqlens[i] / 2;\n }\n __syncthreads();\n\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n int num_threads = blockDim.x * gridDim.x;\n int num_total_tokens = cu_seqlens_s[batch];\n\n for (int token_id = tid; token_id < num_total_tokens; token_id += num_threads) {\n int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1);\n for (int head_id = blockIdx.y; head_id < num_heads; head_id += gridDim.y) {\n size_t idx, half_idx;\n if constexpr (lse_packed) {\n idx = head_id * lse_seqlen + token_id + cu_seqlens_s[seq_id + 1];\n half_idx = head_id * second_half_lse_seqlen + token_id;\n } else {\n size_t row = static_cast(seq_id) * num_heads + head_id;\n int col = token_id - cu_seqlens_s[seq_id];\n int seq_len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id];\n\n idx = row * lse_seqlen + col + seq_len;\n half_idx = row * second_half_lse_seqlen + col;\n }\n\n Functor::run(lse, half_lse, idx, half_idx);\n }\n }\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Out correction in forward\n **************************************************************************************************/\n\ntemplate \n__global__ void thd_out_correction_kernel(dtype *out, dtype *out_per_step, float *lse,\n float *lse_per_step, int *cu_seqlens, int batch,\n int num_heads, int dim_per_head, int lse_seqlen,\n int lse_per_step_seqlen) {\n extern __shared__ int cu_seqlens_s[];\n for (int i = threadIdx.x; i <= batch; i += blockDim.x) {\n cu_seqlens_s[i] = cu_seqlens[i] / (only_second_half + 1);\n }\n __syncthreads();\n\n int tile_id = (blockIdx.x * blockDim.x + threadIdx.x) / tile_size;\n int lane_id = threadIdx.x % tile_size;\n int num_tiles = (blockDim.x * gridDim.x) / tile_size;\n int num_total_tokens = cu_seqlens_s[batch];\n int num_loops_per_head = dim_per_head * sizeof(dtype) / sizeof(float4);\n\n for (int token_id = tile_id; token_id < num_total_tokens; token_id += num_tiles) {\n int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1);\n for (int head_id = blockIdx.y; head_id < num_heads; head_id += gridDim.y) {\n size_t idx, idx_per_step;\n\n if constexpr (lse_packed) {\n idx = head_id * lse_seqlen + token_id + cu_seqlens_s[seq_id + 1] * only_second_half;\n idx_per_step = head_id * lse_per_step_seqlen + token_id;\n } else {\n size_t row = static_cast(seq_id) * num_heads + head_id;\n int col = token_id - cu_seqlens_s[seq_id];\n int seq_len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id];\n idx = row * lse_seqlen + col + seq_len * only_second_half;\n idx_per_step = row * lse_per_step_seqlen + col;\n }\n float lse_corrected_exp = expf(lse_per_step[idx_per_step] - lse[idx]);\n\n idx = token_id + cu_seqlens_s[seq_id + 1] * only_second_half;\n idx = (idx * num_heads + head_id) * dim_per_head;\n idx_per_step = (static_cast(token_id) * num_heads + head_id) * dim_per_head;\n dtype *cur_out = out + idx;\n dtype *cur_out_per_step = out_per_step + idx_per_step;\n\n for (int j = lane_id; j < num_loops_per_head; j += tile_size) {\n float4 data_per_step = reinterpret_cast(cur_out_per_step)[j];\n float4 data = reinterpret_cast(cur_out)[j];\n dtype *p_per_step = reinterpret_cast(&data_per_step);\n dtype *p = reinterpret_cast(&data);\n for (int k = 0; k < sizeof(float4) / sizeof(dtype); k++) {\n p[k] = p[k] +\n (p_per_step[k] == static_cast(0.f)\n ? static_cast(0.f)\n : static_cast(static_cast(p_per_step[k]) * lse_corrected_exp));\n }\n reinterpret_cast(cur_out)[j] = data;\n }\n }\n }\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Gradients correction in backward\n **************************************************************************************************/\n\ntemplate \n__global__ void thd_grad_correction_kernel(dtype *grad, dtype *grad_per_step, int *cu_seqlens,\n int batch, int hidden_size, int dim_size_of_token) {\n extern __shared__ int cu_seqlens_s[];\n for (int i = threadIdx.x; i <= batch; i += blockDim.x) {\n if constexpr (functor_idx < 2) {\n cu_seqlens_s[i] = cu_seqlens[i] / 2;\n } else {\n cu_seqlens_s[i] = cu_seqlens[i];\n }\n }\n __syncthreads();\n\n int group_id = (blockIdx.x * blockDim.x + threadIdx.x) / group_size;\n int lane_id = threadIdx.x % group_size;\n int num_groups = (blockDim.x * gridDim.x) / group_size;\n int num_total_tokens = cu_seqlens_s[batch];\n int num_inner_loops = hidden_size * sizeof(dtype) / sizeof(float4);\n\n size_t offset = static_cast(dim_size_of_token) * hidden_size;\n if constexpr (functor_idx < 2) {\n grad_per_step = grad_per_step + offset / 2 * blockIdx.y;\n } else {\n grad_per_step = grad_per_step + offset * blockIdx.y;\n }\n grad = grad + offset * blockIdx.y;\n\n for (int token_id = group_id; token_id < num_total_tokens; token_id += num_groups) {\n int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1);\n\n int token_offset;\n bool is_first_half;\n if constexpr (functor_idx < 2) {\n token_offset = cu_seqlens_s[seq_id + functor_idx];\n is_first_half = (functor_idx == 0);\n } else {\n token_offset = 0;\n int len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id];\n is_first_half = (token_id - cu_seqlens_s[seq_id]) < (len / 2);\n }\n\n dtype *token = &grad[(token_id + token_offset) * static_cast(hidden_size)];\n dtype *token_per_step = &grad_per_step[token_id * static_cast(hidden_size)];\n for (int idx = lane_id; idx < num_inner_loops; idx += group_size) {\n if (is_first_half) {\n Functor_0::run(token, token_per_step, idx);\n } else {\n Functor_1::run(token, token_per_step, idx);\n }\n }\n }\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Read the half of a THD tensor\n **************************************************************************************************/\n\nvoid thd_read_half_tensor(const Tensor &tensor, const Tensor &cu_seqlens, Tensor &half,\n int half_idx, cudaStream_t stream) {\n using namespace transformer_engine;\n NVTE_CHECK(tensor.dim() == 3 || tensor.dim() == 4);\n NVTE_CHECK(cu_seqlens.dtype() == DType::kInt32);\n\n auto cu_seqlens_shape = cu_seqlens.shape();\n auto tensor_shape = tensor.shape();\n\n NVTE_CHECK(cu_seqlens.dim() == 1);\n NVTE_CHECK(cu_seqlens_shape[0] >= 2);\n\n // Shapes of q and dq are [t, h, d], so the dimension of \"t\" is 0\n // Shapes of kv and dkv are [2, t, h, d], so the dimension of \"t\" is 1\n int seq_dim = tensor.dim() == 3 ? 0 : 1;\n\n int batch = cu_seqlens_shape[0] - 1;\n int num_heads = tensor_shape[seq_dim + 1];\n int dim_per_head = tensor_shape[seq_dim + 2];\n int hidden_size_in_bytes = (num_heads * dim_per_head * typeToNumBits(tensor.dtype())) / 8;\n\n // For 128-bits load/store\n NVTE_CHECK(hidden_size_in_bytes % 16 == 0);\n\n // Launch Kernel\n constexpr unsigned int block = 256;\n unsigned int grid_x = (tensor_shape[seq_dim] / 2 * 32 + block - 1) / block;\n unsigned int grid_y = 1;\n for (int i = 0; i < seq_dim; i++) {\n grid_y *= tensor_shape[i];\n }\n dim3 grid = {grid_x, grid_y};\n thd_read_half_tensor_kernel<<>>(\n half.data.dptr, tensor.data.dptr, reinterpret_cast(cu_seqlens.data.dptr), batch,\n hidden_size_in_bytes, half_idx, tensor_shape[seq_dim]);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: softmax_lse related operations\n **************************************************************************************************/\n\nvoid thd_second_half_lse_correction(Tensor lse, const Tensor &lse_per_step,\n const Tensor &cu_seqlens, bool lse_packed,\n cudaStream_t stream) {\n using namespace transformer_engine;\n NVTE_CHECK(lse.dtype() == DType::kFloat32);\n NVTE_CHECK(lse_per_step.dtype() == DType::kFloat32);\n NVTE_CHECK(cu_seqlens.dtype() == DType::kInt32);\n NVTE_CHECK(cu_seqlens.dim() == 1);\n\n int batch, num_heads, lse_seqlen, second_half_lse_seqlen;\n auto cu_seqlens_shape = cu_seqlens.shape();\n auto lse_shape = lse.shape();\n auto lse_per_step_shape = lse_per_step.shape();\n\n if (lse_packed) {\n NVTE_CHECK(lse.dim() == 2);\n NVTE_CHECK(lse_per_step.dim() == 2);\n\n batch = cu_seqlens_shape[0] - 1;\n num_heads = lse_shape[0];\n lse_seqlen = lse_shape[1];\n second_half_lse_seqlen = lse_per_step_shape[1];\n\n NVTE_CHECK(lse_per_step_shape[0] == num_heads);\n NVTE_CHECK(second_half_lse_seqlen >= lse_seqlen / 2);\n } else {\n NVTE_CHECK(lse.dim() == 3);\n NVTE_CHECK(lse_per_step.dim() == 3);\n\n batch = lse_shape[0];\n num_heads = lse_shape[1];\n lse_seqlen = lse_shape[2];\n second_half_lse_seqlen = lse_per_step_shape[2];\n\n NVTE_CHECK(lse_per_step_shape[0] == batch);\n NVTE_CHECK(lse_per_step_shape[1] == num_heads);\n NVTE_CHECK(second_half_lse_seqlen == lse_seqlen / 2);\n NVTE_CHECK(cu_seqlens_shape[0] == batch + 1);\n }\n\n constexpr unsigned int block = 256;\n unsigned int grid_x = (lse_seqlen / 2 + block - 1) / block;\n unsigned int grid_y = num_heads;\n dim3 grid = {grid_x, grid_y};\n\n if (lse_packed) {\n thd_lse_kernel<<>>(\n reinterpret_cast(lse.data.dptr), reinterpret_cast(lse_per_step.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, lse_seqlen,\n second_half_lse_seqlen);\n NVTE_CHECK_CUDA(cudaGetLastError());\n } else {\n thd_lse_kernel<<>>(\n reinterpret_cast(lse.data.dptr), reinterpret_cast(lse_per_step.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, lse_seqlen,\n second_half_lse_seqlen);\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\nvoid thd_read_second_half_lse(const Tensor &lse, const Tensor &cu_seqlens, Tensor &half_lse,\n bool lse_packed, int second_half_lse_seqlen, cudaStream_t stream) {\n using namespace transformer_engine;\n NVTE_CHECK(lse.dtype() == DType::kFloat32);\n NVTE_CHECK(cu_seqlens.dtype() == DType::kInt32);\n NVTE_CHECK(cu_seqlens.dim() == 1);\n\n int batch, num_heads, lse_seqlen;\n\n auto cu_seqlens_shape = cu_seqlens.shape();\n auto lse_shape = lse.shape();\n\n if (lse_packed) {\n NVTE_CHECK(lse.dim() == 2);\n\n batch = cu_seqlens_shape[0] - 1;\n num_heads = lse_shape[0];\n lse_seqlen = lse_shape[1];\n\n NVTE_CHECK(second_half_lse_seqlen >= lse_seqlen / 2);\n } else {\n NVTE_CHECK(lse.dim() == 3);\n\n batch = lse_shape[0];\n num_heads = lse_shape[1];\n lse_seqlen = lse_shape[2];\n\n NVTE_CHECK(cu_seqlens_shape[0] == batch + 1);\n NVTE_CHECK(second_half_lse_seqlen == lse_seqlen / 2);\n }\n\n constexpr unsigned int block = 256;\n unsigned int grid_x = (lse_seqlen / 2 + block - 1) / block;\n unsigned int grid_y = num_heads;\n dim3 grid = {grid_x, grid_y};\n\n if (lse_packed) {\n thd_lse_kernel<<>>(\n reinterpret_cast(lse.data.dptr), reinterpret_cast(half_lse.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, lse_seqlen,\n second_half_lse_seqlen);\n NVTE_CHECK_CUDA(cudaGetLastError());\n } else {\n thd_lse_kernel<<>>(\n reinterpret_cast(lse.data.dptr), reinterpret_cast(half_lse.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, lse_seqlen,\n second_half_lse_seqlen);\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Out correction in forward\n **************************************************************************************************/\n\ntemplate \nstatic void thd_out_correction_helper(Tensor out, const Tensor &out_per_step, const Tensor &lse,\n const Tensor &lse_per_step, const Tensor &cu_seqlens,\n bool lse_packed, cudaStream_t stream) {\n using namespace transformer_engine;\n NVTE_CHECK(out.dtype() == out_per_step.dtype());\n NVTE_CHECK(lse.dtype() == DType::kFloat32);\n NVTE_CHECK(lse_per_step.dtype() == DType::kFloat32);\n NVTE_CHECK(cu_seqlens.dtype() == DType::kInt32);\n\n auto out_shape = out.shape();\n auto lse_shape = lse.shape();\n auto out_per_step_shape = out_per_step.shape();\n auto lse_per_step_shape = lse_per_step.shape();\n auto cu_seqlens_shape = cu_seqlens.shape();\n\n int total_tokens = out_shape[0];\n int num_heads = out_shape[1];\n int dim_per_head = out_shape[2];\n\n NVTE_CHECK(out_per_step_shape[0] == total_tokens / (only_second_half + 1));\n NVTE_CHECK(out_per_step_shape[1] == num_heads);\n NVTE_CHECK(out_per_step_shape[2] == dim_per_head);\n\n int batch, lse_seqlen, lse_per_step_seqlen;\n if (lse_packed) {\n batch = cu_seqlens_shape[0] - 1;\n lse_seqlen = lse_shape[1];\n lse_per_step_seqlen = lse_per_step_shape[1];\n\n NVTE_CHECK(lse_shape[0] == num_heads);\n NVTE_CHECK(lse_seqlen >= total_tokens);\n NVTE_CHECK(lse_per_step_shape[0] == num_heads);\n NVTE_CHECK(lse_per_step_seqlen >= lse_seqlen / (only_second_half + 1));\n } else {\n batch = lse_shape[0];\n lse_seqlen = lse_shape[2];\n lse_per_step_seqlen = lse_per_step_shape[2];\n\n NVTE_CHECK(lse_shape[1] == num_heads);\n NVTE_CHECK(lse_per_step_shape[0] == batch);\n NVTE_CHECK(lse_per_step_shape[1] == num_heads);\n NVTE_CHECK(lse_per_step_seqlen == lse_seqlen / (only_second_half + 1));\n NVTE_CHECK(cu_seqlens_shape[0] == batch + 1);\n }\n\n constexpr int tile = 16;\n constexpr int block = 512;\n unsigned int grid_x =\n (static_cast(total_tokens) / (only_second_half + 1) * tile + block - 1) / block;\n dim3 grid = {grid_x, (unsigned int)num_heads};\n\n if (lse_packed) {\n thd_out_correction_kernel\n <<>>(\n reinterpret_cast(out.data.dptr),\n reinterpret_cast(out_per_step.data.dptr),\n reinterpret_cast(lse.data.dptr),\n reinterpret_cast(lse_per_step.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, dim_per_head,\n lse_seqlen, lse_per_step_seqlen);\n NVTE_CHECK_CUDA(cudaGetLastError());\n } else {\n thd_out_correction_kernel\n <<>>(\n reinterpret_cast(out.data.dptr),\n reinterpret_cast(out_per_step.data.dptr),\n reinterpret_cast(lse.data.dptr),\n reinterpret_cast(lse_per_step.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, dim_per_head,\n lse_seqlen, lse_per_step_seqlen);\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\nvoid thd_out_correction(Tensor out, const Tensor &out_per_step, const Tensor &lse,\n const Tensor &lse_per_step, const Tensor &cu_seqlens, bool only_second_half,\n bool lse_packed, cudaStream_t stream) {\n using namespace transformer_engine;\n if (only_second_half) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n out.dtype(), dtype,\n thd_out_correction_helper(out, out_per_step, lse, lse_per_step, cu_seqlens,\n lse_packed, stream););\n } else {\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n out.dtype(), dtype,\n thd_out_correction_helper(out, out_per_step, lse, lse_per_step, cu_seqlens,\n lse_packed, stream););\n }\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Gradients correction in backward\n **************************************************************************************************/\n\ntemplate \nstatic void thd_grad_correction_helper(Tensor grad, const Tensor &grad_per_step,\n const Tensor &cu_seqlens, cudaStream_t stream) {\n using namespace transformer_engine;\n NVTE_CHECK(grad.dim() == 3 || grad.dim() == 4);\n NVTE_CHECK(cu_seqlens.dtype() == DType::kInt32);\n NVTE_CHECK(cu_seqlens.dim() == 1);\n\n auto grad_shape = grad.shape();\n auto cu_seqlens_shape = cu_seqlens.shape();\n auto grad_per_step_shape = grad_per_step.shape();\n\n // Shape of dq is [t, h, d], so the dimension of \"t\" is 0\n // Shape of dkv is [2, t, h, d], so the dimension of \"t\" is 1\n int seq_dim = grad.dim() == 3 ? 0 : 1;\n\n int total_tokens = grad_shape[seq_dim];\n int num_heads = grad_shape[seq_dim + 1];\n int dim_per_head = grad_shape[seq_dim + 2];\n int batch = cu_seqlens_shape[0] - 1;\n\n if constexpr (functor_idx < 2) {\n NVTE_CHECK(grad_per_step_shape[seq_dim] == total_tokens / 2);\n } else {\n NVTE_CHECK(grad_per_step_shape[seq_dim] == total_tokens);\n }\n NVTE_CHECK(grad_per_step_shape[seq_dim + 1] == num_heads);\n NVTE_CHECK(grad_per_step_shape[seq_dim + 2] == dim_per_head);\n\n size_t hidden_size = num_heads * dim_per_head;\n NVTE_CHECK(((hidden_size * typeToNumBits(grad.dtype())) / 8) % 16 == 0);\n\n constexpr unsigned int block = 256;\n unsigned int grid_x;\n if constexpr (functor_idx < 2) {\n grid_x = (total_tokens / 2 * 32 + block - 1) / block;\n } else {\n grid_x = (total_tokens * 32 + block - 1) / block;\n }\n unsigned int grid_y = 1;\n for (int i = 0; i < seq_dim; i++) {\n grid_y *= grad_shape[i];\n }\n dim3 grid = {grid_x, grid_y};\n\n thd_grad_correction_kernel\n <<>>(\n reinterpret_cast(grad.data.dptr),\n reinterpret_cast(grad_per_step.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), batch, hidden_size, total_tokens);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\ntemplate \nstatic void thd_grad_dispatcher(Tensor grad, const Tensor &grad_per_step, const Tensor &cu_seqlens,\n const std::string &first_half, const std::string &second_half,\n cudaStream_t stream) {\n using namespace transformer_engine;\n if (first_half == \"add\" && second_half == \"none\") {\n thd_grad_correction_helper, EmptyFunctor, 0>(grad, grad_per_step,\n cu_seqlens, stream);\n } else if (first_half == \"copy\" && second_half == \"none\") {\n thd_grad_correction_helper(grad, grad_per_step, cu_seqlens,\n stream);\n } else if (first_half == \"none\" && second_half == \"add\") {\n thd_grad_correction_helper, 1>(grad, grad_per_step,\n cu_seqlens, stream);\n } else if (first_half == \"none\" && second_half == \"copy\") {\n thd_grad_correction_helper(grad, grad_per_step, cu_seqlens,\n stream);\n } else if (first_half == \"add\" && second_half == \"copy\") {\n thd_grad_correction_helper, CopyFunctor, 2>(grad, grad_per_step,\n cu_seqlens, stream);\n } else if (first_half == \"copy\" && second_half == \"add\") {\n thd_grad_correction_helper, 2>(grad, grad_per_step,\n cu_seqlens, stream);\n } else {\n NVTE_ERROR(\"Unsupported Functor of first half and second_half\\n\");\n }\n}\n\nvoid thd_grad_correction(Tensor grad, const Tensor &grad_per_step, const Tensor &cu_seqlens,\n const std::string &first_half, const std::string &second_half,\n cudaStream_t stream) {\n using namespace transformer_engine;\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n grad.dtype(), dtype,\n thd_grad_dispatcher(grad, grad_per_step, cu_seqlens, first_half, second_half,\n stream););\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Generate partitioned indices for input tokens\n **************************************************************************************************/\n\nvoid thd_get_partitioned_indices(const Tensor &cu_seqlens, Tensor output, int total_tokens,\n int world_size, int rank, cudaStream_t stream) {\n using namespace transformer_engine;\n NVTE_CHECK(cu_seqlens.dtype() == DType::kInt32);\n NVTE_CHECK(cu_seqlens.dim() == 1);\n\n auto cu_seqlens_shape = cu_seqlens.shape();\n auto output_shape = output.shape();\n\n NVTE_CHECK(cu_seqlens_shape[0] >= 2);\n NVTE_CHECK(rank >= 0 && rank < world_size);\n NVTE_CHECK(world_size > 0);\n NVTE_CHECK(total_tokens > 0 && total_tokens % (world_size * 2) == 0);\n\n int batch = cu_seqlens_shape[0] - 1;\n\n constexpr unsigned int block = 256;\n unsigned int grid = (output_shape[0] + block - 1) / block;\n thd_partition_indices_kernel<<>>(\n reinterpret_cast(output.data.dptr), reinterpret_cast(cu_seqlens.data.dptr),\n batch, total_tokens, world_size, rank);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n} // namespace context_parallel\n} // namespace transformer_engine\n\nvoid nvte_cp_thd_read_half_tensor(const NVTETensor &tensor, const NVTETensor &cu_seqlens,\n NVTETensor half, int half_idx, cudaStream_t stream) {\n NVTE_API_CALL(nvte_thd_read_half_tensor);\n using namespace transformer_engine;\n\n context_parallel::thd_read_half_tensor(*convertNVTETensorCheck(tensor),\n *convertNVTETensorCheck(cu_seqlens),\n *convertNVTETensorCheck(half), half_idx, stream);\n}\n\nvoid nvte_cp_thd_second_half_lse_correction(NVTETensor lse, const NVTETensor &lse_per_step,\n const NVTETensor &cu_seqlens, int lse_packed,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_thd_second_half_lse_correction);\n using namespace transformer_engine;\n\n context_parallel::thd_second_half_lse_correction(\n *convertNVTETensorCheck(lse), *convertNVTETensorCheck(lse_per_step),\n *convertNVTETensorCheck(cu_seqlens), lse_packed, stream);\n}\n\nvoid nvte_cp_thd_read_second_half_lse(const NVTETensor &lse, const NVTETensor &cu_seqlens,\n NVTETensor half_lse, int lse_packed,\n int second_half_lse_seqlen, cudaStream_t stream) {\n NVTE_API_CALL(nvte_thd_read_second_half_lse);\n using namespace transformer_engine;\n\n context_parallel::thd_read_second_half_lse(\n *convertNVTETensorCheck(lse), *convertNVTETensorCheck(cu_seqlens),\n *convertNVTETensorCheck(half_lse), lse_packed, second_half_lse_seqlen, stream);\n}\n\nvoid nvte_cp_thd_out_correction(NVTETensor out, const NVTETensor &out_per_step,\n const NVTETensor &lse, const NVTETensor &lse_per_step,\n const NVTETensor &cu_seqlens, int only_second_half, int lse_packed,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_thd_out_correction);\n using namespace transformer_engine;\n\n context_parallel::thd_out_correction(\n *convertNVTETensorCheck(out), *convertNVTETensorCheck(out_per_step),\n *convertNVTETensorCheck(lse), *convertNVTETensorCheck(lse_per_step),\n *convertNVTETensorCheck(cu_seqlens), only_second_half, lse_packed, stream);\n}\n\nvoid nvte_cp_thd_grad_correction(NVTETensor grad, const NVTETensor &grad_per_step,\n const NVTETensor &cu_seqlens, const char *first_half,\n const char *second_half, cudaStream_t stream) {\n NVTE_API_CALL(nvte_thd_grad_correction);\n using namespace transformer_engine;\n\n std::string first_half_str(first_half);\n std::string second_half_str(second_half);\n\n context_parallel::thd_grad_correction(\n *convertNVTETensorCheck(grad), *convertNVTETensorCheck(grad_per_step),\n *convertNVTETensorCheck(cu_seqlens), first_half_str, second_half_str, stream);\n}\n\nvoid nvte_cp_thd_get_partitioned_indices(const NVTETensor &cu_seqlens, NVTETensor output,\n int total_tokens, int world_size, int rank,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_thd_get_partitioned_indices);\n using namespace transformer_engine;\n\n context_parallel::thd_get_partitioned_indices(*convertNVTETensorCheck(cu_seqlens),\n *convertNVTETensorCheck(output), total_tokens,\n world_size, rank, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"transformer_engine/fused_attn.h\"\n\nnamespace transformer_engine {\nnamespace context_parallel {\n\nstruct LseCorrectionFunctor {\n __forceinline__ __device__ static void run(float *lse, float *half_lse, size_t idx,\n size_t half_idx) {\n float val = lse[idx];\n float val_per_step = half_lse[half_idx];\n float max_scale = max(val, val_per_step);\n float min_scale = min(val, val_per_step);\n lse[idx] = max_scale + log1pf(expf(min_scale - max_scale));\n }\n};\n\nstruct ReadLseFunctor {\n __forceinline__ __device__ static void run(float *lse, float *half_lse, size_t idx,\n size_t half_idx) {\n half_lse[half_idx] = lse[idx];\n }\n};\n\nstruct EmptyFunctor {\n __forceinline__ __device__ static void run(void *token, void *token_per_step, int idx) {}\n};\n\nstruct CopyFunctor {\n __forceinline__ __device__ static void run(void *token, void *token_per_step, int idx) {\n reinterpret_cast(token)[idx] = reinterpret_cast(token_per_step)[idx];\n }\n};\n\ntemplate \nstruct AddFunctor {\n __forceinline__ __device__ static void run(dtype *token, dtype *token_per_step, int idx) {\n float4 d_ = reinterpret_cast(token)[idx];\n dtype *p_ = reinterpret_cast(&d_);\n\n float4 d = reinterpret_cast(token_per_step)[idx];\n dtype *p = reinterpret_cast(&d);\n\n#pragma unroll\n for (int i = 0; i < sizeof(float4) / sizeof(dtype); i++) {\n p_[i] = p_[i] + p[i];\n }\n\n reinterpret_cast(token)[idx] = d_;\n }\n};\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Binary search an array for a target value\n **************************************************************************************************/\n\n__forceinline__ __device__ int binary_search(int target, int *array, int len) {\n int left = 1, right = len - 1;\n while (left < right) {\n int mid = (left + right) / 2;\n if (array[mid] <= target) {\n left = mid + 1;\n } else {\n right = mid;\n }\n }\n return left - 1;\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Generate partitioned indices for input tokens\n **************************************************************************************************/\n__global__ void thd_partition_indices_kernel(int *output, int *cu_seqlens, int batch,\n int total_tokens, int world_size, int rank) {\n extern __shared__ int cu_seqlens_s[];\n for (int i = threadIdx.x; i <= batch; i += blockDim.x) {\n int seqlen = cu_seqlens[i];\n // Currently we assume that each sequence length is divisible by (world_size*2) since we have\n // to distribute each sequence evenly to different GPUs.\n assert(seqlen % (world_size * 2) == 0);\n cu_seqlens_s[i] = seqlen / world_size;\n }\n __syncthreads();\n\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n int num_threads = blockDim.x * gridDim.x;\n\n for (int token_id = tid; token_id < total_tokens / world_size; token_id += num_threads) {\n int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1);\n int seq_len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id];\n int index = token_id - cu_seqlens_s[seq_id];\n int offset = index < seq_len / 2 ? rank : (world_size - 1) * 2 - rank;\n index += cu_seqlens_s[seq_id] * world_size + seq_len / 2 * offset;\n output[token_id] = index;\n }\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Read the half of a THD tensor\n **************************************************************************************************/\n\n__global__ void thd_read_half_tensor_kernel(void *half, void *tensor, int *cu_seqlens, int batch,\n int hidden_size_in_bytes, int half_idx,\n int dim_size_of_token) {\n extern __shared__ int cu_seqlens_s[];\n for (int i = threadIdx.x; i <= batch; i += blockDim.x) {\n cu_seqlens_s[i] = cu_seqlens[i] / 2;\n }\n __syncthreads();\n\n int warpid = (blockIdx.x * blockDim.x + threadIdx.x) / 32;\n int laneid = threadIdx.x % 32;\n int num_warps = (blockDim.x * gridDim.x) / 32;\n int num_total_tokens = cu_seqlens_s[batch];\n int num_float4s_per_token = hidden_size_in_bytes / sizeof(float4);\n\n size_t offset = static_cast(dim_size_of_token) * hidden_size_in_bytes;\n half = reinterpret_cast(reinterpret_cast(half) + offset / 2 * blockIdx.y);\n tensor = reinterpret_cast(reinterpret_cast(tensor) + offset * blockIdx.y);\n\n for (int token_id = warpid; token_id < num_total_tokens; token_id += num_warps) {\n int seqid = binary_search(token_id, cu_seqlens_s, batch + 1);\n\n size_t offset_in_bytes = static_cast(token_id) * hidden_size_in_bytes;\n float4 *cur_half_token =\n reinterpret_cast(reinterpret_cast(half) + offset_in_bytes);\n\n offset_in_bytes =\n (static_cast(token_id) + cu_seqlens_s[seqid + half_idx]) * hidden_size_in_bytes;\n float4 *cur_token =\n reinterpret_cast(reinterpret_cast(tensor) + offset_in_bytes);\n\n for (int idx = laneid; idx < num_float4s_per_token; idx += 32) {\n cur_half_token[idx] = cur_token[idx];\n }\n }\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: softmax_lse related operations\n **************************************************************************************************/\n\ntemplate \n__global__ void thd_lse_kernel(float *lse, float *half_lse, int *cu_seqlens, int batch,\n int num_heads, int lse_seqlen, int second_half_lse_seqlen) {\n extern __shared__ int cu_seqlens_s[];\n for (int i = threadIdx.x; i <= batch; i += blockDim.x) {\n cu_seqlens_s[i] = cu_seqlens[i] / 2;\n }\n __syncthreads();\n\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n int num_threads = blockDim.x * gridDim.x;\n int num_total_tokens = cu_seqlens_s[batch];\n\n for (int token_id = tid; token_id < num_total_tokens; token_id += num_threads) {\n int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1);\n for (int head_id = blockIdx.y; head_id < num_heads; head_id += gridDim.y) {\n size_t idx, half_idx;\n if constexpr (lse_packed) {\n idx = head_id * lse_seqlen + token_id + cu_seqlens_s[seq_id + 1];\n half_idx = head_id * second_half_lse_seqlen + token_id;\n } else {\n size_t row = static_cast(seq_id) * num_heads + head_id;\n int col = token_id - cu_seqlens_s[seq_id];\n int seq_len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id];\n\n idx = row * lse_seqlen + col + seq_len;\n half_idx = row * second_half_lse_seqlen + col;\n }\n\n Functor::run(lse, half_lse, idx, half_idx);\n }\n }\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Out correction in forward\n **************************************************************************************************/\n\ntemplate \n__global__ void thd_out_correction_kernel(dtype *out, dtype *out_per_step, float *lse,\n float *lse_per_step, int *cu_seqlens, int batch,\n int num_heads, int dim_per_head, int lse_seqlen,\n int lse_per_step_seqlen) {\n extern __shared__ int cu_seqlens_s[];\n for (int i = threadIdx.x; i <= batch; i += blockDim.x) {\n cu_seqlens_s[i] = cu_seqlens[i] / (only_second_half + 1);\n }\n __syncthreads();\n\n int tile_id = (blockIdx.x * blockDim.x + threadIdx.x) / tile_size;\n int lane_id = threadIdx.x % tile_size;\n int num_tiles = (blockDim.x * gridDim.x) / tile_size;\n int num_total_tokens = cu_seqlens_s[batch];\n int num_loops_per_head = dim_per_head * sizeof(dtype) / sizeof(float4);\n\n for (int token_id = tile_id; token_id < num_total_tokens; token_id += num_tiles) {\n int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1);\n for (int head_id = blockIdx.y; head_id < num_heads; head_id += gridDim.y) {\n size_t idx, idx_per_step;\n\n if constexpr (lse_packed) {\n idx = head_id * lse_seqlen + token_id + cu_seqlens_s[seq_id + 1] * only_second_half;\n idx_per_step = head_id * lse_per_step_seqlen + token_id;\n } else {\n size_t row = static_cast(seq_id) * num_heads + head_id;\n int col = token_id - cu_seqlens_s[seq_id];\n int seq_len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id];\n idx = row * lse_seqlen + col + seq_len * only_second_half;\n idx_per_step = row * lse_per_step_seqlen + col;\n }\n float lse_corrected_exp = expf(lse_per_step[idx_per_step] - lse[idx]);\n\n idx = token_id + cu_seqlens_s[seq_id + 1] * only_second_half;\n idx = (idx * num_heads + head_id) * dim_per_head;\n idx_per_step = (static_cast(token_id) * num_heads + head_id) * dim_per_head;\n dtype *cur_out = out + idx;\n dtype *cur_out_per_step = out_per_step + idx_per_step;\n\n for (int j = lane_id; j < num_loops_per_head; j += tile_size) {\n float4 data_per_step = reinterpret_cast(cur_out_per_step)[j];\n float4 data = reinterpret_cast(cur_out)[j];\n dtype *p_per_step = reinterpret_cast(&data_per_step);\n dtype *p = reinterpret_cast(&data);\n for (int k = 0; k < sizeof(float4) / sizeof(dtype); k++) {\n p[k] = p[k] +\n (p_per_step[k] == static_cast(0.f)\n ? static_cast(0.f)\n : static_cast(static_cast(p_per_step[k]) * lse_corrected_exp));\n }\n reinterpret_cast(cur_out)[j] = data;\n }\n }\n }\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Gradients correction in backward\n **************************************************************************************************/\n\ntemplate \n__global__ void thd_grad_correction_kernel(dtype *grad, dtype *grad_per_step, int *cu_seqlens,\n int batch, int hidden_size, int dim_size_of_token) {\n extern __shared__ int cu_seqlens_s[];\n for (int i = threadIdx.x; i <= batch; i += blockDim.x) {\n if constexpr (functor_idx < 2) {\n cu_seqlens_s[i] = cu_seqlens[i] / 2;\n } else {\n cu_seqlens_s[i] = cu_seqlens[i];\n }\n }\n __syncthreads();\n\n int group_id = (blockIdx.x * blockDim.x + threadIdx.x) / group_size;\n int lane_id = threadIdx.x % group_size;\n int num_groups = (blockDim.x * gridDim.x) / group_size;\n int num_total_tokens = cu_seqlens_s[batch];\n int num_inner_loops = hidden_size * sizeof(dtype) / sizeof(float4);\n\n size_t offset = static_cast(dim_size_of_token) * hidden_size;\n if constexpr (functor_idx < 2) {\n grad_per_step = grad_per_step + offset / 2 * blockIdx.y;\n } else {\n grad_per_step = grad_per_step + offset * blockIdx.y;\n }\n grad = grad + offset * blockIdx.y;\n\n for (int token_id = group_id; token_id < num_total_tokens; token_id += num_groups) {\n int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1);\n\n int token_offset;\n bool is_first_half;\n if constexpr (functor_idx < 2) {\n token_offset = cu_seqlens_s[seq_id + functor_idx];\n is_first_half = (functor_idx == 0);\n } else {\n token_offset = 0;\n int len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id];\n is_first_half = (token_id - cu_seqlens_s[seq_id]) < (len / 2);\n }\n\n dtype *token = &grad[(token_id + token_offset) * static_cast(hidden_size)];\n dtype *token_per_step = &grad_per_step[token_id * static_cast(hidden_size)];\n for (int idx = lane_id; idx < num_inner_loops; idx += group_size) {\n if (is_first_half) {\n Functor_0::run(token, token_per_step, idx);\n } else {\n Functor_1::run(token, token_per_step, idx);\n }\n }\n }\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Read the half of a THD tensor\n **************************************************************************************************/\n\nvoid thd_read_half_tensor(const Tensor &tensor, const Tensor &cu_seqlens, Tensor &half,\n int half_idx, hipStream_t stream) {\n using namespace transformer_engine;\n NVTE_CHECK(tensor.dim() == 3 || tensor.dim() == 4);\n NVTE_CHECK(cu_seqlens.dtype() == DType::kInt32);\n\n auto cu_seqlens_shape = cu_seqlens.shape();\n auto tensor_shape = tensor.shape();\n\n NVTE_CHECK(cu_seqlens.dim() == 1);\n NVTE_CHECK(cu_seqlens_shape[0] >= 2);\n\n // Shapes of q and dq are [t, h, d], so the dimension of \"t\" is 0\n // Shapes of kv and dkv are [2, t, h, d], so the dimension of \"t\" is 1\n int seq_dim = tensor.dim() == 3 ? 0 : 1;\n\n int batch = cu_seqlens_shape[0] - 1;\n int num_heads = tensor_shape[seq_dim + 1];\n int dim_per_head = tensor_shape[seq_dim + 2];\n int hidden_size_in_bytes = (num_heads * dim_per_head * typeToNumBits(tensor.dtype())) / 8;\n\n // For 128-bits load/store\n NVTE_CHECK(hidden_size_in_bytes % 16 == 0);\n\n // Launch Kernel\n constexpr unsigned int block = 256;\n unsigned int grid_x = (tensor_shape[seq_dim] / 2 * 32 + block - 1) / block;\n unsigned int grid_y = 1;\n for (int i = 0; i < seq_dim; i++) {\n grid_y *= tensor_shape[i];\n }\n dim3 grid = {grid_x, grid_y};\n thd_read_half_tensor_kernel<<>>(\n half.data.dptr, tensor.data.dptr, reinterpret_cast(cu_seqlens.data.dptr), batch,\n hidden_size_in_bytes, half_idx, tensor_shape[seq_dim]);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: softmax_lse related operations\n **************************************************************************************************/\n\nvoid thd_second_half_lse_correction(Tensor lse, const Tensor &lse_per_step,\n const Tensor &cu_seqlens, bool lse_packed,\n hipStream_t stream) {\n using namespace transformer_engine;\n NVTE_CHECK(lse.dtype() == DType::kFloat32);\n NVTE_CHECK(lse_per_step.dtype() == DType::kFloat32);\n NVTE_CHECK(cu_seqlens.dtype() == DType::kInt32);\n NVTE_CHECK(cu_seqlens.dim() == 1);\n\n int batch, num_heads, lse_seqlen, second_half_lse_seqlen;\n auto cu_seqlens_shape = cu_seqlens.shape();\n auto lse_shape = lse.shape();\n auto lse_per_step_shape = lse_per_step.shape();\n\n if (lse_packed) {\n NVTE_CHECK(lse.dim() == 2);\n NVTE_CHECK(lse_per_step.dim() == 2);\n\n batch = cu_seqlens_shape[0] - 1;\n num_heads = lse_shape[0];\n lse_seqlen = lse_shape[1];\n second_half_lse_seqlen = lse_per_step_shape[1];\n\n NVTE_CHECK(lse_per_step_shape[0] == num_heads);\n NVTE_CHECK(second_half_lse_seqlen >= lse_seqlen / 2);\n } else {\n NVTE_CHECK(lse.dim() == 3);\n NVTE_CHECK(lse_per_step.dim() == 3);\n\n batch = lse_shape[0];\n num_heads = lse_shape[1];\n lse_seqlen = lse_shape[2];\n second_half_lse_seqlen = lse_per_step_shape[2];\n\n NVTE_CHECK(lse_per_step_shape[0] == batch);\n NVTE_CHECK(lse_per_step_shape[1] == num_heads);\n NVTE_CHECK(second_half_lse_seqlen == lse_seqlen / 2);\n NVTE_CHECK(cu_seqlens_shape[0] == batch + 1);\n }\n\n constexpr unsigned int block = 256;\n unsigned int grid_x = (lse_seqlen / 2 + block - 1) / block;\n unsigned int grid_y = num_heads;\n dim3 grid = {grid_x, grid_y};\n\n if (lse_packed) {\n thd_lse_kernel<<>>(\n reinterpret_cast(lse.data.dptr), reinterpret_cast(lse_per_step.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, lse_seqlen,\n second_half_lse_seqlen);\n NVTE_CHECK_CUDA(hipGetLastError());\n } else {\n thd_lse_kernel<<>>(\n reinterpret_cast(lse.data.dptr), reinterpret_cast(lse_per_step.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, lse_seqlen,\n second_half_lse_seqlen);\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\nvoid thd_read_second_half_lse(const Tensor &lse, const Tensor &cu_seqlens, Tensor &half_lse,\n bool lse_packed, int second_half_lse_seqlen, hipStream_t stream) {\n using namespace transformer_engine;\n NVTE_CHECK(lse.dtype() == DType::kFloat32);\n NVTE_CHECK(cu_seqlens.dtype() == DType::kInt32);\n NVTE_CHECK(cu_seqlens.dim() == 1);\n\n int batch, num_heads, lse_seqlen;\n\n auto cu_seqlens_shape = cu_seqlens.shape();\n auto lse_shape = lse.shape();\n\n if (lse_packed) {\n NVTE_CHECK(lse.dim() == 2);\n\n batch = cu_seqlens_shape[0] - 1;\n num_heads = lse_shape[0];\n lse_seqlen = lse_shape[1];\n\n NVTE_CHECK(second_half_lse_seqlen >= lse_seqlen / 2);\n } else {\n NVTE_CHECK(lse.dim() == 3);\n\n batch = lse_shape[0];\n num_heads = lse_shape[1];\n lse_seqlen = lse_shape[2];\n\n NVTE_CHECK(cu_seqlens_shape[0] == batch + 1);\n NVTE_CHECK(second_half_lse_seqlen == lse_seqlen / 2);\n }\n\n constexpr unsigned int block = 256;\n unsigned int grid_x = (lse_seqlen / 2 + block - 1) / block;\n unsigned int grid_y = num_heads;\n dim3 grid = {grid_x, grid_y};\n\n if (lse_packed) {\n thd_lse_kernel<<>>(\n reinterpret_cast(lse.data.dptr), reinterpret_cast(half_lse.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, lse_seqlen,\n second_half_lse_seqlen);\n NVTE_CHECK_CUDA(hipGetLastError());\n } else {\n thd_lse_kernel<<>>(\n reinterpret_cast(lse.data.dptr), reinterpret_cast(half_lse.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, lse_seqlen,\n second_half_lse_seqlen);\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Out correction in forward\n **************************************************************************************************/\n\ntemplate \nstatic void thd_out_correction_helper(Tensor out, const Tensor &out_per_step, const Tensor &lse,\n const Tensor &lse_per_step, const Tensor &cu_seqlens,\n bool lse_packed, hipStream_t stream) {\n using namespace transformer_engine;\n NVTE_CHECK(out.dtype() == out_per_step.dtype());\n NVTE_CHECK(lse.dtype() == DType::kFloat32);\n NVTE_CHECK(lse_per_step.dtype() == DType::kFloat32);\n NVTE_CHECK(cu_seqlens.dtype() == DType::kInt32);\n\n auto out_shape = out.shape();\n auto lse_shape = lse.shape();\n auto out_per_step_shape = out_per_step.shape();\n auto lse_per_step_shape = lse_per_step.shape();\n auto cu_seqlens_shape = cu_seqlens.shape();\n\n int total_tokens = out_shape[0];\n int num_heads = out_shape[1];\n int dim_per_head = out_shape[2];\n\n NVTE_CHECK(out_per_step_shape[0] == total_tokens / (only_second_half + 1));\n NVTE_CHECK(out_per_step_shape[1] == num_heads);\n NVTE_CHECK(out_per_step_shape[2] == dim_per_head);\n\n int batch, lse_seqlen, lse_per_step_seqlen;\n if (lse_packed) {\n batch = cu_seqlens_shape[0] - 1;\n lse_seqlen = lse_shape[1];\n lse_per_step_seqlen = lse_per_step_shape[1];\n\n NVTE_CHECK(lse_shape[0] == num_heads);\n NVTE_CHECK(lse_seqlen >= total_tokens);\n NVTE_CHECK(lse_per_step_shape[0] == num_heads);\n NVTE_CHECK(lse_per_step_seqlen >= lse_seqlen / (only_second_half + 1));\n } else {\n batch = lse_shape[0];\n lse_seqlen = lse_shape[2];\n lse_per_step_seqlen = lse_per_step_shape[2];\n\n NVTE_CHECK(lse_shape[1] == num_heads);\n NVTE_CHECK(lse_per_step_shape[0] == batch);\n NVTE_CHECK(lse_per_step_shape[1] == num_heads);\n NVTE_CHECK(lse_per_step_seqlen == lse_seqlen / (only_second_half + 1));\n NVTE_CHECK(cu_seqlens_shape[0] == batch + 1);\n }\n\n constexpr int tile = 16;\n constexpr int block = 512;\n unsigned int grid_x =\n (static_cast(total_tokens) / (only_second_half + 1) * tile + block - 1) / block;\n dim3 grid = {grid_x, (unsigned int)num_heads};\n\n if (lse_packed) {\n thd_out_correction_kernel\n <<>>(\n reinterpret_cast(out.data.dptr),\n reinterpret_cast(out_per_step.data.dptr),\n reinterpret_cast(lse.data.dptr),\n reinterpret_cast(lse_per_step.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, dim_per_head,\n lse_seqlen, lse_per_step_seqlen);\n NVTE_CHECK_CUDA(hipGetLastError());\n } else {\n thd_out_correction_kernel\n <<>>(\n reinterpret_cast(out.data.dptr),\n reinterpret_cast(out_per_step.data.dptr),\n reinterpret_cast(lse.data.dptr),\n reinterpret_cast(lse_per_step.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), batch, num_heads, dim_per_head,\n lse_seqlen, lse_per_step_seqlen);\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\nvoid thd_out_correction(Tensor out, const Tensor &out_per_step, const Tensor &lse,\n const Tensor &lse_per_step, const Tensor &cu_seqlens, bool only_second_half,\n bool lse_packed, hipStream_t stream) {\n using namespace transformer_engine;\n if (only_second_half) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n out.dtype(), dtype,\n thd_out_correction_helper(out, out_per_step, lse, lse_per_step, cu_seqlens,\n lse_packed, stream););\n } else {\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n out.dtype(), dtype,\n thd_out_correction_helper(out, out_per_step, lse, lse_per_step, cu_seqlens,\n lse_packed, stream););\n }\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Gradients correction in backward\n **************************************************************************************************/\n\ntemplate \nstatic void thd_grad_correction_helper(Tensor grad, const Tensor &grad_per_step,\n const Tensor &cu_seqlens, hipStream_t stream) {\n using namespace transformer_engine;\n NVTE_CHECK(grad.dim() == 3 || grad.dim() == 4);\n NVTE_CHECK(cu_seqlens.dtype() == DType::kInt32);\n NVTE_CHECK(cu_seqlens.dim() == 1);\n\n auto grad_shape = grad.shape();\n auto cu_seqlens_shape = cu_seqlens.shape();\n auto grad_per_step_shape = grad_per_step.shape();\n\n // Shape of dq is [t, h, d], so the dimension of \"t\" is 0\n // Shape of dkv is [2, t, h, d], so the dimension of \"t\" is 1\n int seq_dim = grad.dim() == 3 ? 0 : 1;\n\n int total_tokens = grad_shape[seq_dim];\n int num_heads = grad_shape[seq_dim + 1];\n int dim_per_head = grad_shape[seq_dim + 2];\n int batch = cu_seqlens_shape[0] - 1;\n\n if constexpr (functor_idx < 2) {\n NVTE_CHECK(grad_per_step_shape[seq_dim] == total_tokens / 2);\n } else {\n NVTE_CHECK(grad_per_step_shape[seq_dim] == total_tokens);\n }\n NVTE_CHECK(grad_per_step_shape[seq_dim + 1] == num_heads);\n NVTE_CHECK(grad_per_step_shape[seq_dim + 2] == dim_per_head);\n\n size_t hidden_size = num_heads * dim_per_head;\n NVTE_CHECK(((hidden_size * typeToNumBits(grad.dtype())) / 8) % 16 == 0);\n\n constexpr unsigned int block = 256;\n unsigned int grid_x;\n if constexpr (functor_idx < 2) {\n grid_x = (total_tokens / 2 * 32 + block - 1) / block;\n } else {\n grid_x = (total_tokens * 32 + block - 1) / block;\n }\n unsigned int grid_y = 1;\n for (int i = 0; i < seq_dim; i++) {\n grid_y *= grad_shape[i];\n }\n dim3 grid = {grid_x, grid_y};\n\n thd_grad_correction_kernel\n <<>>(\n reinterpret_cast(grad.data.dptr),\n reinterpret_cast(grad_per_step.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr), batch, hidden_size, total_tokens);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\ntemplate \nstatic void thd_grad_dispatcher(Tensor grad, const Tensor &grad_per_step, const Tensor &cu_seqlens,\n const std::string &first_half, const std::string &second_half,\n hipStream_t stream) {\n using namespace transformer_engine;\n if (first_half == \"add\" && second_half == \"none\") {\n thd_grad_correction_helper, EmptyFunctor, 0>(grad, grad_per_step,\n cu_seqlens, stream);\n } else if (first_half == \"copy\" && second_half == \"none\") {\n thd_grad_correction_helper(grad, grad_per_step, cu_seqlens,\n stream);\n } else if (first_half == \"none\" && second_half == \"add\") {\n thd_grad_correction_helper, 1>(grad, grad_per_step,\n cu_seqlens, stream);\n } else if (first_half == \"none\" && second_half == \"copy\") {\n thd_grad_correction_helper(grad, grad_per_step, cu_seqlens,\n stream);\n } else if (first_half == \"add\" && second_half == \"copy\") {\n thd_grad_correction_helper, CopyFunctor, 2>(grad, grad_per_step,\n cu_seqlens, stream);\n } else if (first_half == \"copy\" && second_half == \"add\") {\n thd_grad_correction_helper, 2>(grad, grad_per_step,\n cu_seqlens, stream);\n } else {\n NVTE_ERROR(\"Unsupported Functor of first half and second_half\\n\");\n }\n}\n\nvoid thd_grad_correction(Tensor grad, const Tensor &grad_per_step, const Tensor &cu_seqlens,\n const std::string &first_half, const std::string &second_half,\n hipStream_t stream) {\n using namespace transformer_engine;\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n grad.dtype(), dtype,\n thd_grad_dispatcher(grad, grad_per_step, cu_seqlens, first_half, second_half,\n stream););\n}\n\n/***************************************************************************************************\n * Support THD format for Context Parallel: Generate partitioned indices for input tokens\n **************************************************************************************************/\n\nvoid thd_get_partitioned_indices(const Tensor &cu_seqlens, Tensor output, int total_tokens,\n int world_size, int rank, hipStream_t stream) {\n using namespace transformer_engine;\n NVTE_CHECK(cu_seqlens.dtype() == DType::kInt32);\n NVTE_CHECK(cu_seqlens.dim() == 1);\n\n auto cu_seqlens_shape = cu_seqlens.shape();\n auto output_shape = output.shape();\n\n NVTE_CHECK(cu_seqlens_shape[0] >= 2);\n NVTE_CHECK(rank >= 0 && rank < world_size);\n NVTE_CHECK(world_size > 0);\n NVTE_CHECK(total_tokens > 0 && total_tokens % (world_size * 2) == 0);\n\n int batch = cu_seqlens_shape[0] - 1;\n\n constexpr unsigned int block = 256;\n unsigned int grid = (output_shape[0] + block - 1) / block;\n thd_partition_indices_kernel<<>>(\n reinterpret_cast(output.data.dptr), reinterpret_cast(cu_seqlens.data.dptr),\n batch, total_tokens, world_size, rank);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n} // namespace context_parallel\n} // namespace transformer_engine\n\nvoid nvte_cp_thd_read_half_tensor(const NVTETensor &tensor, const NVTETensor &cu_seqlens,\n NVTETensor half, int half_idx, hipStream_t stream) {\n NVTE_API_CALL(nvte_thd_read_half_tensor);\n using namespace transformer_engine;\n\n context_parallel::thd_read_half_tensor(*convertNVTETensorCheck(tensor),\n *convertNVTETensorCheck(cu_seqlens),\n *convertNVTETensorCheck(half), half_idx, stream);\n}\n\nvoid nvte_cp_thd_second_half_lse_correction(NVTETensor lse, const NVTETensor &lse_per_step,\n const NVTETensor &cu_seqlens, int lse_packed,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_thd_second_half_lse_correction);\n using namespace transformer_engine;\n\n context_parallel::thd_second_half_lse_correction(\n *convertNVTETensorCheck(lse), *convertNVTETensorCheck(lse_per_step),\n *convertNVTETensorCheck(cu_seqlens), lse_packed, stream);\n}\n\nvoid nvte_cp_thd_read_second_half_lse(const NVTETensor &lse, const NVTETensor &cu_seqlens,\n NVTETensor half_lse, int lse_packed,\n int second_half_lse_seqlen, hipStream_t stream) {\n NVTE_API_CALL(nvte_thd_read_second_half_lse);\n using namespace transformer_engine;\n\n context_parallel::thd_read_second_half_lse(\n *convertNVTETensorCheck(lse), *convertNVTETensorCheck(cu_seqlens),\n *convertNVTETensorCheck(half_lse), lse_packed, second_half_lse_seqlen, stream);\n}\n\nvoid nvte_cp_thd_out_correction(NVTETensor out, const NVTETensor &out_per_step,\n const NVTETensor &lse, const NVTETensor &lse_per_step,\n const NVTETensor &cu_seqlens, int only_second_half, int lse_packed,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_thd_out_correction);\n using namespace transformer_engine;\n\n context_parallel::thd_out_correction(\n *convertNVTETensorCheck(out), *convertNVTETensorCheck(out_per_step),\n *convertNVTETensorCheck(lse), *convertNVTETensorCheck(lse_per_step),\n *convertNVTETensorCheck(cu_seqlens), only_second_half, lse_packed, stream);\n}\n\nvoid nvte_cp_thd_grad_correction(NVTETensor grad, const NVTETensor &grad_per_step,\n const NVTETensor &cu_seqlens, const char *first_half,\n const char *second_half, hipStream_t stream) {\n NVTE_API_CALL(nvte_thd_grad_correction);\n using namespace transformer_engine;\n\n std::string first_half_str(first_half);\n std::string second_half_str(second_half);\n\n context_parallel::thd_grad_correction(\n *convertNVTETensorCheck(grad), *convertNVTETensorCheck(grad_per_step),\n *convertNVTETensorCheck(cu_seqlens), first_half_str, second_half_str, stream);\n}\n\nvoid nvte_cp_thd_get_partitioned_indices(const NVTETensor &cu_seqlens, NVTETensor output,\n int total_tokens, int world_size, int rank,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_thd_get_partitioned_indices);\n using namespace transformer_engine;\n\n context_parallel::thd_get_partitioned_indices(*convertNVTETensorCheck(cu_seqlens),\n *convertNVTETensorCheck(output), total_tokens,\n world_size, rank, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/fused_attn/context_parallel.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \"../common.h\"\n#include \"../cudnn_utils.h\"\n#include \"transformer_engine/fused_attn.h\"\n#include \"utils.h\"\n\nnamespace transformer_engine {\nnamespace fused_attn {\n\nusing namespace transformer_engine;\n\n// get matrix strides based on matrix type\nvoid generateMatrixStrides(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d,\n int64_t *strideA, NVTE_QKV_Layout layout, NVTE_QKV_Matrix matrix) {\n constexpr int batch_dim_idx = 0;\n constexpr int head_dim_idx = 1;\n constexpr int seqlen_dim_idx = 2;\n constexpr int hidden_dim_idx = 3;\n\n constexpr int seqlen_transpose_dim_idx = 3;\n constexpr int hidden_transpose_dim_idx = 2;\n\n constexpr int seqlen_q_dim_idx = 2;\n constexpr int seqlen_kv_dim_idx = 3;\n\n switch (layout) {\n case NVTE_QKV_Layout::NVTE_SB3HD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = 3 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * 3 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = 3 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = b * 3 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_SBH3D:\n if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = 3 * h * d;\n strideA[head_dim_idx] = 3 * d;\n strideA[seqlen_dim_idx] = b * 3 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = 3 * h * d;\n strideA[head_dim_idx] = 3 * d;\n strideA[seqlen_transpose_dim_idx] = b * 3 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_SBHD_SB2HD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = 2 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * 2 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = 2 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = b * 2 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_SBHD_SBH2D:\n if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = 2 * h * d;\n strideA[head_dim_idx] = 2 * d;\n strideA[seqlen_dim_idx] = b * 2 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = 2 * h * d;\n strideA[head_dim_idx] = 2 * d;\n strideA[seqlen_transpose_dim_idx] = b * 2 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_SBHD_SBHD_SBHD:\n case NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_SBHD_SBHD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = b * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_BS3HD:\n case NVTE_QKV_Layout::NVTE_T3HD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = s_q * 3 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = 3 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = s_q * 3 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = 3 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix) {\n strideA[batch_dim_idx] = s_q * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_BSH3D:\n case NVTE_QKV_Layout::NVTE_TH3D:\n if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = s_q * 3 * h * d;\n strideA[head_dim_idx] = 3 * d;\n strideA[seqlen_dim_idx] = 3 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = s_q * 3 * h * d;\n strideA[head_dim_idx] = 3 * d;\n strideA[seqlen_transpose_dim_idx] = 3 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix) {\n strideA[batch_dim_idx] = s_q * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_BSHD_BS2HD:\n case NVTE_QKV_Layout::NVTE_THD_T2HD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = s_kv * 2 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = 2 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = s_kv * 2 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = 2 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = s_q * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_BSHD_BSH2D:\n case NVTE_QKV_Layout::NVTE_THD_TH2D:\n if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = s_kv * 2 * h * d;\n strideA[head_dim_idx] = 2 * d;\n strideA[seqlen_dim_idx] = 2 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = s_kv * 2 * h * d;\n strideA[head_dim_idx] = 2 * d;\n strideA[seqlen_transpose_dim_idx] = 2 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = s_q * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_BSHD_BSHD_BSHD:\n case NVTE_QKV_Layout::NVTE_THD_THD_THD:\n case NVTE_QKV_Layout::NVTE_THD_BSHD_BSHD:\n case NVTE_QKV_Layout::NVTE_Paged_KV_BSHD_BSHD_BSHD:\n case NVTE_QKV_Layout::NVTE_Paged_KV_THD_BSHD_BSHD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = s_q * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = s_kv * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = s_kv * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_SBHD_BSHD_BSHD:\n case NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_BSHD_BSHD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = s_kv * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = s_kv * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_BSHD_SBHD_SBHD:\n case NVTE_QKV_Layout::NVTE_THD_SBHD_SBHD:\n case NVTE_QKV_Layout::NVTE_Paged_KV_BSHD_SBHD_SBHD:\n case NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = b * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = s_q * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = h * s_q * d;\n strideA[head_dim_idx] = s_q * d;\n strideA[seqlen_dim_idx] = d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = h * s_kv * d;\n strideA[head_dim_idx] = s_kv * d;\n strideA[seqlen_dim_idx] = d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = h * s_kv * d;\n strideA[head_dim_idx] = s_kv * d;\n strideA[seqlen_transpose_dim_idx] = d;\n strideA[hidden_transpose_dim_idx] = 1;\n }\n break;\n }\n\n if (matrix == NVTE_QKV_Matrix::NVTE_S_Matrix) {\n strideA[seqlen_kv_dim_idx] = 1;\n strideA[seqlen_q_dim_idx] = s_kv;\n strideA[head_dim_idx] = s_q * s_kv;\n strideA[batch_dim_idx] = h * s_q * s_kv;\n }\n}\n\nbool allowAllConfig(cudnnBackendDescriptor_t engine_config) {\n (void)engine_config;\n return false;\n}\n\ncudnn_frontend::Tensor tensor_create(cudnnDataType_t type, int64_t id, int64_t const *dim,\n int64_t const *stride, bool is_virtual, bool is_value) {\n int nbDims = 4;\n auto tensor_created =\n cudnn_frontend::TensorBuilder()\n .setDim(nbDims, dim)\n .setStride(nbDims, stride)\n .setId(id)\n .setAlignment(16) // 16B alignment is needed to run a tensor core engine\n .setDataType(type)\n .setVirtual(is_virtual)\n .setByValue(is_value)\n .build();\n return tensor_created;\n}\n\ncudnn_frontend::Tensor tensor_create_with_offset(\n cudnnDataType_t type, int64_t id, int64_t const *dim, int64_t const *stride, bool is_virtual,\n bool is_value, std::shared_ptr raggedOffset) {\n int nbDims = 4;\n auto tensor_created =\n cudnn_frontend::TensorBuilder()\n .setDim(nbDims, dim)\n .setStride(nbDims, stride)\n .setId(id)\n .setAlignment(16) // 16B alignment is needed to run a tensor core engine\n .setDataType(type)\n .setVirtual(is_virtual)\n .setByValue(is_value)\n .setRaggedOffset(raggedOffset)\n .build();\n return tensor_created;\n}\n\ncudnn_frontend::PointWiseDesc pw_desc_create(cudnnDataType_t type, cudnnPointwiseMode_t mode) {\n auto pw_desc_created =\n cudnn_frontend::PointWiseDescBuilder().setMode(mode).setComputeType(type).build();\n return pw_desc_created;\n}\n\ncudnn_frontend::Operation unary_pw_op_create(cudnn_frontend::Tensor const &xDesc,\n cudnn_frontend::Tensor const &yDesc,\n cudnn_frontend::PointWiseDesc const &pwDesc) {\n auto pw_op_created =\n cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(xDesc)\n .setyDesc(yDesc)\n .setpwDesc(pwDesc)\n .build();\n return pw_op_created;\n}\n\ncudnn_frontend::Operation binary_pw_op_create(cudnn_frontend::Tensor const &xDesc,\n cudnn_frontend::Tensor const &bDesc,\n cudnn_frontend::Tensor const &yDesc,\n cudnn_frontend::PointWiseDesc const &pwDesc) {\n auto pw_op_created =\n cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(xDesc)\n .setbDesc(bDesc)\n .setyDesc(yDesc)\n .setpwDesc(pwDesc)\n .build();\n return pw_op_created;\n}\n\ncudnn_frontend::Operation ternary_pw_op_create(cudnn_frontend::Tensor const &xDesc,\n cudnn_frontend::Tensor const &bDesc,\n cudnn_frontend::Tensor const &tDesc,\n cudnn_frontend::Tensor const &yDesc,\n cudnn_frontend::PointWiseDesc const &pwDesc) {\n auto pw_op_created =\n cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(xDesc)\n .setbDesc(bDesc)\n .settDesc(tDesc)\n .setyDesc(yDesc)\n .setpwDesc(pwDesc)\n .build();\n return pw_op_created;\n}\n\n// convert cu_seqlens to actual_seqlens\n__global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b,\n int32_t const *const q_cu_seqlens,\n int32_t const *const kv_cu_seqlens, int32_t *q_seqlens,\n int32_t *kv_seqlens) {\n size_t tid = blockIdx.x * blockDim.x + threadIdx.x;\n if (tid < actual_b) {\n q_seqlens[tid] = q_cu_seqlens[tid + 1] - q_cu_seqlens[tid];\n kv_seqlens[tid] = kv_cu_seqlens[tid + 1] - kv_cu_seqlens[tid];\n } else if (tid < max_b) {\n q_seqlens[tid] = 0;\n kv_seqlens[tid] = 0;\n }\n}\n\n// convert cu_seqlens_padded to offsets\ntemplate \n__device__ void cu_seqlens_padded_to_offsets_impl(\n NVTE_QKV_Layout_Group layout_group, int64_t actual_b, int64_t max_b, int64_t h, int64_t hg,\n int64_t d_qk, int64_t d_v, const int32_t *cu_seqlens_q_padded,\n const int32_t *cu_seqlens_kv_padded, OFFSETS_T *offsets_q, OFFSETS_T *offsets_k,\n OFFSETS_T *offsets_v, OFFSETS_T *offsets_o, OFFSETS_T *offsets_s) {\n size_t tid = blockIdx.x * blockDim.x + threadIdx.x;\n auto cu_seqlens_id = min(tid, actual_b);\n if (tid <= max_b) {\n if (offsets_s != nullptr) {\n offsets_s[tid] = h * cu_seqlens_q_padded[cu_seqlens_id];\n }\n if (offsets_q != nullptr && offsets_o != nullptr) {\n offsets_o[tid] = h * d_v * cu_seqlens_q_padded[cu_seqlens_id];\n switch (layout_group) {\n case NVTE_QKV_Layout_Group::NVTE_HD_HD_HD:\n case NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD:\n offsets_q[tid] = h * d_qk * cu_seqlens_q_padded[cu_seqlens_id];\n break;\n case NVTE_QKV_Layout_Group::NVTE_3HD:\n case NVTE_QKV_Layout_Group::NVTE_H3D:\n offsets_q[tid] = 3 * h * d_qk * cu_seqlens_q_padded[cu_seqlens_id];\n break;\n case NVTE_QKV_Layout_Group::NVTE_HD_2HD:\n case NVTE_QKV_Layout_Group::NVTE_HD_H2D:\n offsets_q[tid] = h * d_qk * cu_seqlens_q_padded[cu_seqlens_id];\n break;\n }\n }\n if (offsets_k != nullptr && offsets_v != nullptr) {\n switch (layout_group) {\n case NVTE_QKV_Layout_Group::NVTE_HD_HD_HD:\n case NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD:\n offsets_k[tid] = hg * d_qk * cu_seqlens_kv_padded[cu_seqlens_id];\n offsets_v[tid] = hg * d_v * cu_seqlens_kv_padded[cu_seqlens_id];\n break;\n case NVTE_QKV_Layout_Group::NVTE_3HD:\n case NVTE_QKV_Layout_Group::NVTE_H3D:\n offsets_k[tid] = 3 * h * d_qk * cu_seqlens_q_padded[cu_seqlens_id];\n offsets_v[tid] = offsets_k[cu_seqlens_id];\n break;\n case NVTE_QKV_Layout_Group::NVTE_HD_2HD:\n case NVTE_QKV_Layout_Group::NVTE_HD_H2D:\n offsets_k[tid] = 2 * hg * d_qk * cu_seqlens_kv_padded[cu_seqlens_id];\n offsets_v[tid] = offsets_k[cu_seqlens_id];\n break;\n }\n }\n }\n}\n\n__global__ void cu_seqlens_padded_to_offsets(NVTE_QKV_Layout_Group layout_group, int64_t actual_b,\n int64_t max_b, int64_t h, int64_t hg, int64_t d_qk,\n int64_t d_v, const int32_t *cu_seqlens_q_padded,\n const int32_t *cu_seqlens_kv_padded,\n DType offset_dtype, void *offsets_q, void *offsets_k,\n void *offsets_v, void *offsets_o, void *offsets_s) {\n if (offset_dtype == DType::kInt32) {\n cu_seqlens_padded_to_offsets_impl(\n layout_group, actual_b, max_b, h, hg, d_qk, d_v, cu_seqlens_q_padded, cu_seqlens_kv_padded,\n reinterpret_cast(offsets_q), reinterpret_cast(offsets_k),\n reinterpret_cast(offsets_v), reinterpret_cast(offsets_o),\n reinterpret_cast(offsets_s));\n } else {\n assert(offset_dtype == DType::kInt64 && \"expect int64\");\n cu_seqlens_padded_to_offsets_impl(\n layout_group, actual_b, max_b, h, hg, d_qk, d_v, cu_seqlens_q_padded, cu_seqlens_kv_padded,\n reinterpret_cast(offsets_q), reinterpret_cast(offsets_k),\n reinterpret_cast(offsets_v), reinterpret_cast(offsets_o),\n reinterpret_cast(offsets_s));\n }\n}\n\nDType get_ragged_offset_dtype(NVTE_QKV_Layout_Group layout_group, int64_t num_attn_heads,\n int64_t num_gqa_groups, int64_t max_seqlen_q, int64_t max_seqlen_kv,\n int64_t head_dim_qk, int64_t head_dim_v) {\n std::array offsets_qkvo{};\n switch (layout_group) {\n case NVTE_QKV_Layout_Group::NVTE_HD_HD_HD:\n case NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD:\n offsets_qkvo[0] = num_attn_heads * head_dim_qk * max_seqlen_q;\n offsets_qkvo[1] = num_gqa_groups * head_dim_qk * max_seqlen_kv;\n offsets_qkvo[2] = num_gqa_groups * head_dim_v * max_seqlen_kv;\n break;\n case NVTE_QKV_Layout_Group::NVTE_3HD:\n case NVTE_QKV_Layout_Group::NVTE_H3D:\n offsets_qkvo[0] = 3 * num_attn_heads * head_dim_qk * max_seqlen_q;\n offsets_qkvo[1] = offsets_qkvo[0];\n offsets_qkvo[2] = offsets_qkvo[0];\n break;\n case NVTE_QKV_Layout_Group::NVTE_HD_2HD:\n case NVTE_QKV_Layout_Group::NVTE_HD_H2D:\n offsets_qkvo[0] = num_attn_heads * head_dim_qk * max_seqlen_q;\n offsets_qkvo[1] = 2 * num_gqa_groups * head_dim_qk * max_seqlen_kv;\n offsets_qkvo[2] = offsets_qkvo[1];\n break;\n }\n\n offsets_qkvo[3] = num_attn_heads * head_dim_qk * max_seqlen_q;\n\n size_t max_offset = *std::max_element(offsets_qkvo.begin(), offsets_qkvo.end());\n if (max_offset > std::numeric_limits::max()) {\n return DType::kInt64;\n }\n\n return DType::kInt32;\n}\n\n// quantize batch size\nsize_t get_max_batch_size(size_t batch_size) {\n size_t max_b = batch_size;\n size_t log2_b = ceil(log2(batch_size));\n // batch size is expected to be 10s-100s\n // b = 1, ..., 32 -> max_b = 32\n // b = 33, ..., 512 -> max_b = next power of 2\n // b = 513, ... -> max_b = increment by 512\n if (log2_b <= 5) {\n max_b = 32;\n } else if (log2_b <= 9) {\n max_b = pow(2, log2_b);\n } else {\n max_b = (batch_size + 511) / 512 * 512;\n }\n return max_b;\n}\n\n// quantize token count\nsize_t get_max_tokens(size_t num_tokens) {\n // token count is expected to be 1k's-100k's\n // t = 0, ..., 1024 -> max_t = 1024\n // t = 1025, ..., 32k -> max_t = next power of 2\n // t = 32k+1, ... -> max_t = increment by 32k\n size_t log2_t = ceil(log2(num_tokens));\n size_t max_t = 0;\n if (log2_t <= 10) {\n max_t = 1024;\n } else if (log2_t <= 15) {\n max_t = pow(2, log2_t);\n } else {\n max_t = (num_tokens + 32767) / 32768 * 32768;\n }\n return max_t;\n}\n\n__global__ void populate_rng_state_kernel(int64_t *rng_state_dst, const int64_t *const seed,\n int64_t offset) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n if (tid > 0) return;\n rng_state_dst[0] = seed[0];\n rng_state_dst[1] = offset;\n}\n\n__global__ void get_runtime_num_segments_kernel(int32_t *cu_seqlen, size_t len, uint32_t *out) {\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n if (tid >= len) return;\n\n if (cu_seqlen[tid] > 0) {\n // atomicAdd only support 32 bits dtype\n atomicAdd(out, 1);\n }\n}\n\nvoid PopulateRngStateAsync(void *rng_state_dst, const void *seed, size_t q_max_seqlen,\n size_t kv_max_seqlen, NVTE_Fused_Attn_Backend backend,\n cudaStream_t stream) {\n size_t increment = 0;\n if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) {\n increment = 16;\n } else {\n constexpr int threads_per_cta = 128;\n increment = (q_max_seqlen * kv_max_seqlen + threads_per_cta - 1) / threads_per_cta;\n }\n auto offset = FusedAttnOffsetManager::Instance().GetAndUpdateOffset(increment);\n populate_rng_state_kernel<<<1, 1, 0, stream>>>(reinterpret_cast(rng_state_dst),\n reinterpret_cast(seed), offset);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nuint32_t GetRuntimeNumSegments(void *cu_seqlen, void *workspace, size_t len, cudaStream_t stream) {\n // workspace size requires 4 bytes\n uint32_t *dout = static_cast(workspace);\n uint32_t hout{};\n NVTE_CHECK_CUDA(cudaMemsetAsync(dout, 0, sizeof(uint32_t), stream));\n constexpr int threads = 128;\n const int blocks = (len - 1) / threads + 1;\n get_runtime_num_segments_kernel<<>>(static_cast(cu_seqlen),\n len, dout);\n NVTE_CHECK_CUDA(cudaGetLastError());\n NVTE_CHECK_CUDA(cudaMemcpyAsync(&hout, dout, sizeof(uint32_t), cudaMemcpyDeviceToHost, stream));\n NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));\n return hout;\n}\n\n__global__ void extract_seed_and_offset(int64_t *rng_state_ptr, bool captured, int64_t *seed_ptr,\n uint64_t seed_val, int64_t *offset_ptr, uint64_t offset_val,\n uint32_t offset_intragraph) {\n if (captured) {\n rng_state_ptr[0] = *seed_ptr;\n rng_state_ptr[1] = static_cast(*offset_ptr + static_cast(offset_intragraph));\n } else {\n rng_state_ptr[0] = static_cast(seed_val);\n rng_state_ptr[1] = static_cast(offset_val);\n }\n}\n\n} // namespace fused_attn\n} // namespace transformer_engine\n\nvoid nvte_extract_seed_and_offset(int64_t *rng_state_ptr, int captured, int64_t *seed_ptr,\n uint64_t seed_val, int64_t *offset_ptr, uint64_t offset_val,\n uint32_t offset_intragraph, cudaStream_t stream) {\n NVTE_API_CALL(nvte_extract_seed_and_offset);\n using namespace transformer_engine;\n\n fused_attn::extract_seed_and_offset<<<1, 1, 0, stream>>>(\n rng_state_ptr, captured, seed_ptr, seed_val, offset_ptr, offset_val, offset_intragraph);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n#include \"../common.h\"\n#include \"../cudnn_utils.h\"\n#include \"transformer_engine/fused_attn.h\"\n#include \"utils.h\"\n\nnamespace transformer_engine {\nnamespace fused_attn {\n\nusing namespace transformer_engine;\n\n// get matrix strides based on matrix type\nvoid generateMatrixStrides(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d,\n int64_t *strideA, NVTE_QKV_Layout layout, NVTE_QKV_Matrix matrix) {\n constexpr int batch_dim_idx = 0;\n constexpr int head_dim_idx = 1;\n constexpr int seqlen_dim_idx = 2;\n constexpr int hidden_dim_idx = 3;\n\n constexpr int seqlen_transpose_dim_idx = 3;\n constexpr int hidden_transpose_dim_idx = 2;\n\n constexpr int seqlen_q_dim_idx = 2;\n constexpr int seqlen_kv_dim_idx = 3;\n\n switch (layout) {\n case NVTE_QKV_Layout::NVTE_SB3HD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = 3 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * 3 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = 3 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = b * 3 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_SBH3D:\n if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = 3 * h * d;\n strideA[head_dim_idx] = 3 * d;\n strideA[seqlen_dim_idx] = b * 3 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = 3 * h * d;\n strideA[head_dim_idx] = 3 * d;\n strideA[seqlen_transpose_dim_idx] = b * 3 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_SBHD_SB2HD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = 2 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * 2 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = 2 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = b * 2 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_SBHD_SBH2D:\n if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = 2 * h * d;\n strideA[head_dim_idx] = 2 * d;\n strideA[seqlen_dim_idx] = b * 2 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = 2 * h * d;\n strideA[head_dim_idx] = 2 * d;\n strideA[seqlen_transpose_dim_idx] = b * 2 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_SBHD_SBHD_SBHD:\n case NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_SBHD_SBHD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = b * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_BS3HD:\n case NVTE_QKV_Layout::NVTE_T3HD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = s_q * 3 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = 3 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = s_q * 3 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = 3 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix) {\n strideA[batch_dim_idx] = s_q * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_BSH3D:\n case NVTE_QKV_Layout::NVTE_TH3D:\n if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = s_q * 3 * h * d;\n strideA[head_dim_idx] = 3 * d;\n strideA[seqlen_dim_idx] = 3 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = s_q * 3 * h * d;\n strideA[head_dim_idx] = 3 * d;\n strideA[seqlen_transpose_dim_idx] = 3 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix) {\n strideA[batch_dim_idx] = s_q * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_BSHD_BS2HD:\n case NVTE_QKV_Layout::NVTE_THD_T2HD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = s_kv * 2 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = 2 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = s_kv * 2 * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = 2 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = s_q * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_BSHD_BSH2D:\n case NVTE_QKV_Layout::NVTE_THD_TH2D:\n if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = s_kv * 2 * h * d;\n strideA[head_dim_idx] = 2 * d;\n strideA[seqlen_dim_idx] = 2 * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = s_kv * 2 * h * d;\n strideA[head_dim_idx] = 2 * d;\n strideA[seqlen_transpose_dim_idx] = 2 * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = s_q * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_BSHD_BSHD_BSHD:\n case NVTE_QKV_Layout::NVTE_THD_THD_THD:\n case NVTE_QKV_Layout::NVTE_THD_BSHD_BSHD:\n case NVTE_QKV_Layout::NVTE_Paged_KV_BSHD_BSHD_BSHD:\n case NVTE_QKV_Layout::NVTE_Paged_KV_THD_BSHD_BSHD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = s_q * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = s_kv * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = s_kv * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_SBHD_BSHD_BSHD:\n case NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_BSHD_BSHD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = s_kv * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = s_kv * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_BSHD_SBHD_SBHD:\n case NVTE_QKV_Layout::NVTE_THD_SBHD_SBHD:\n case NVTE_QKV_Layout::NVTE_Paged_KV_BSHD_SBHD_SBHD:\n case NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = b * h * d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_transpose_dim_idx] = b * h * d;\n strideA[hidden_transpose_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = s_q * h * d;\n strideA[head_dim_idx] = d;\n strideA[seqlen_dim_idx] = h * d;\n strideA[hidden_dim_idx] = 1;\n }\n break;\n case NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD:\n if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) {\n strideA[batch_dim_idx] = h * s_q * d;\n strideA[head_dim_idx] = s_q * d;\n strideA[seqlen_dim_idx] = d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) {\n strideA[batch_dim_idx] = h * s_kv * d;\n strideA[head_dim_idx] = s_kv * d;\n strideA[seqlen_dim_idx] = d;\n strideA[hidden_dim_idx] = 1;\n } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) ||\n (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) {\n strideA[batch_dim_idx] = h * s_kv * d;\n strideA[head_dim_idx] = s_kv * d;\n strideA[seqlen_transpose_dim_idx] = d;\n strideA[hidden_transpose_dim_idx] = 1;\n }\n break;\n }\n\n if (matrix == NVTE_QKV_Matrix::NVTE_S_Matrix) {\n strideA[seqlen_kv_dim_idx] = 1;\n strideA[seqlen_q_dim_idx] = s_kv;\n strideA[head_dim_idx] = s_q * s_kv;\n strideA[batch_dim_idx] = h * s_q * s_kv;\n }\n}\n\nbool allowAllConfig(cudnnBackendDescriptor_t engine_config) {\n (void)engine_config;\n return false;\n}\n\ncudnn_frontend::Tensor tensor_create(hipdnnDataType_t type, int64_t id, int64_t const *dim,\n int64_t const *stride, bool is_virtual, bool is_value) {\n int nbDims = 4;\n auto tensor_created =\n cudnn_frontend::TensorBuilder()\n .setDim(nbDims, dim)\n .setStride(nbDims, stride)\n .setId(id)\n .setAlignment(16) // 16B alignment is needed to run a tensor core engine\n .setDataType(type)\n .setVirtual(is_virtual)\n .setByValue(is_value)\n .build();\n return tensor_created;\n}\n\ncudnn_frontend::Tensor tensor_create_with_offset(\n hipdnnDataType_t type, int64_t id, int64_t const *dim, int64_t const *stride, bool is_virtual,\n bool is_value, std::shared_ptr raggedOffset) {\n int nbDims = 4;\n auto tensor_created =\n cudnn_frontend::TensorBuilder()\n .setDim(nbDims, dim)\n .setStride(nbDims, stride)\n .setId(id)\n .setAlignment(16) // 16B alignment is needed to run a tensor core engine\n .setDataType(type)\n .setVirtual(is_virtual)\n .setByValue(is_value)\n .setRaggedOffset(raggedOffset)\n .build();\n return tensor_created;\n}\n\ncudnn_frontend::PointWiseDesc pw_desc_create(hipdnnDataType_t type, cudnnPointwiseMode_t mode) {\n auto pw_desc_created =\n cudnn_frontend::PointWiseDescBuilder().setMode(mode).setComputeType(type).build();\n return pw_desc_created;\n}\n\ncudnn_frontend::Operation unary_pw_op_create(cudnn_frontend::Tensor const &xDesc,\n cudnn_frontend::Tensor const &yDesc,\n cudnn_frontend::PointWiseDesc const &pwDesc) {\n auto pw_op_created =\n cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(xDesc)\n .setyDesc(yDesc)\n .setpwDesc(pwDesc)\n .build();\n return pw_op_created;\n}\n\ncudnn_frontend::Operation binary_pw_op_create(cudnn_frontend::Tensor const &xDesc,\n cudnn_frontend::Tensor const &bDesc,\n cudnn_frontend::Tensor const &yDesc,\n cudnn_frontend::PointWiseDesc const &pwDesc) {\n auto pw_op_created =\n cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(xDesc)\n .setbDesc(bDesc)\n .setyDesc(yDesc)\n .setpwDesc(pwDesc)\n .build();\n return pw_op_created;\n}\n\ncudnn_frontend::Operation ternary_pw_op_create(cudnn_frontend::Tensor const &xDesc,\n cudnn_frontend::Tensor const &bDesc,\n cudnn_frontend::Tensor const &tDesc,\n cudnn_frontend::Tensor const &yDesc,\n cudnn_frontend::PointWiseDesc const &pwDesc) {\n auto pw_op_created =\n cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(xDesc)\n .setbDesc(bDesc)\n .settDesc(tDesc)\n .setyDesc(yDesc)\n .setpwDesc(pwDesc)\n .build();\n return pw_op_created;\n}\n\n// convert cu_seqlens to actual_seqlens\n__global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b,\n int32_t const *const q_cu_seqlens,\n int32_t const *const kv_cu_seqlens, int32_t *q_seqlens,\n int32_t *kv_seqlens) {\n size_t tid = blockIdx.x * blockDim.x + threadIdx.x;\n if (tid < actual_b) {\n q_seqlens[tid] = q_cu_seqlens[tid + 1] - q_cu_seqlens[tid];\n kv_seqlens[tid] = kv_cu_seqlens[tid + 1] - kv_cu_seqlens[tid];\n } else if (tid < max_b) {\n q_seqlens[tid] = 0;\n kv_seqlens[tid] = 0;\n }\n}\n\n// convert cu_seqlens_padded to offsets\ntemplate \n__device__ void cu_seqlens_padded_to_offsets_impl(\n NVTE_QKV_Layout_Group layout_group, int64_t actual_b, int64_t max_b, int64_t h, int64_t hg,\n int64_t d_qk, int64_t d_v, const int32_t *cu_seqlens_q_padded,\n const int32_t *cu_seqlens_kv_padded, OFFSETS_T *offsets_q, OFFSETS_T *offsets_k,\n OFFSETS_T *offsets_v, OFFSETS_T *offsets_o, OFFSETS_T *offsets_s) {\n size_t tid = blockIdx.x * blockDim.x + threadIdx.x;\n auto cu_seqlens_id = min(tid, actual_b);\n if (tid <= max_b) {\n if (offsets_s != nullptr) {\n offsets_s[tid] = h * cu_seqlens_q_padded[cu_seqlens_id];\n }\n if (offsets_q != nullptr && offsets_o != nullptr) {\n offsets_o[tid] = h * d_v * cu_seqlens_q_padded[cu_seqlens_id];\n switch (layout_group) {\n case NVTE_QKV_Layout_Group::NVTE_HD_HD_HD:\n case NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD:\n offsets_q[tid] = h * d_qk * cu_seqlens_q_padded[cu_seqlens_id];\n break;\n case NVTE_QKV_Layout_Group::NVTE_3HD:\n case NVTE_QKV_Layout_Group::NVTE_H3D:\n offsets_q[tid] = 3 * h * d_qk * cu_seqlens_q_padded[cu_seqlens_id];\n break;\n case NVTE_QKV_Layout_Group::NVTE_HD_2HD:\n case NVTE_QKV_Layout_Group::NVTE_HD_H2D:\n offsets_q[tid] = h * d_qk * cu_seqlens_q_padded[cu_seqlens_id];\n break;\n }\n }\n if (offsets_k != nullptr && offsets_v != nullptr) {\n switch (layout_group) {\n case NVTE_QKV_Layout_Group::NVTE_HD_HD_HD:\n case NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD:\n offsets_k[tid] = hg * d_qk * cu_seqlens_kv_padded[cu_seqlens_id];\n offsets_v[tid] = hg * d_v * cu_seqlens_kv_padded[cu_seqlens_id];\n break;\n case NVTE_QKV_Layout_Group::NVTE_3HD:\n case NVTE_QKV_Layout_Group::NVTE_H3D:\n offsets_k[tid] = 3 * h * d_qk * cu_seqlens_q_padded[cu_seqlens_id];\n offsets_v[tid] = offsets_k[cu_seqlens_id];\n break;\n case NVTE_QKV_Layout_Group::NVTE_HD_2HD:\n case NVTE_QKV_Layout_Group::NVTE_HD_H2D:\n offsets_k[tid] = 2 * hg * d_qk * cu_seqlens_kv_padded[cu_seqlens_id];\n offsets_v[tid] = offsets_k[cu_seqlens_id];\n break;\n }\n }\n }\n}\n\n__global__ void cu_seqlens_padded_to_offsets(NVTE_QKV_Layout_Group layout_group, int64_t actual_b,\n int64_t max_b, int64_t h, int64_t hg, int64_t d_qk,\n int64_t d_v, const int32_t *cu_seqlens_q_padded,\n const int32_t *cu_seqlens_kv_padded,\n DType offset_dtype, void *offsets_q, void *offsets_k,\n void *offsets_v, void *offsets_o, void *offsets_s) {\n if (offset_dtype == DType::kInt32) {\n cu_seqlens_padded_to_offsets_impl(\n layout_group, actual_b, max_b, h, hg, d_qk, d_v, cu_seqlens_q_padded, cu_seqlens_kv_padded,\n reinterpret_cast(offsets_q), reinterpret_cast(offsets_k),\n reinterpret_cast(offsets_v), reinterpret_cast(offsets_o),\n reinterpret_cast(offsets_s));\n } else {\n assert(offset_dtype == DType::kInt64 && \"expect int64\");\n cu_seqlens_padded_to_offsets_impl(\n layout_group, actual_b, max_b, h, hg, d_qk, d_v, cu_seqlens_q_padded, cu_seqlens_kv_padded,\n reinterpret_cast(offsets_q), reinterpret_cast(offsets_k),\n reinterpret_cast(offsets_v), reinterpret_cast(offsets_o),\n reinterpret_cast(offsets_s));\n }\n}\n\nDType get_ragged_offset_dtype(NVTE_QKV_Layout_Group layout_group, int64_t num_attn_heads,\n int64_t num_gqa_groups, int64_t max_seqlen_q, int64_t max_seqlen_kv,\n int64_t head_dim_qk, int64_t head_dim_v) {\n std::array offsets_qkvo{};\n switch (layout_group) {\n case NVTE_QKV_Layout_Group::NVTE_HD_HD_HD:\n case NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD:\n offsets_qkvo[0] = num_attn_heads * head_dim_qk * max_seqlen_q;\n offsets_qkvo[1] = num_gqa_groups * head_dim_qk * max_seqlen_kv;\n offsets_qkvo[2] = num_gqa_groups * head_dim_v * max_seqlen_kv;\n break;\n case NVTE_QKV_Layout_Group::NVTE_3HD:\n case NVTE_QKV_Layout_Group::NVTE_H3D:\n offsets_qkvo[0] = 3 * num_attn_heads * head_dim_qk * max_seqlen_q;\n offsets_qkvo[1] = offsets_qkvo[0];\n offsets_qkvo[2] = offsets_qkvo[0];\n break;\n case NVTE_QKV_Layout_Group::NVTE_HD_2HD:\n case NVTE_QKV_Layout_Group::NVTE_HD_H2D:\n offsets_qkvo[0] = num_attn_heads * head_dim_qk * max_seqlen_q;\n offsets_qkvo[1] = 2 * num_gqa_groups * head_dim_qk * max_seqlen_kv;\n offsets_qkvo[2] = offsets_qkvo[1];\n break;\n }\n\n offsets_qkvo[3] = num_attn_heads * head_dim_qk * max_seqlen_q;\n\n size_t max_offset = *std::max_element(offsets_qkvo.begin(), offsets_qkvo.end());\n if (max_offset > std::numeric_limits::max()) {\n return DType::kInt64;\n }\n\n return DType::kInt32;\n}\n\n// quantize batch size\nsize_t get_max_batch_size(size_t batch_size) {\n size_t max_b = batch_size;\n size_t log2_b = ceil(log2(batch_size));\n // batch size is expected to be 10s-100s\n // b = 1, ..., 32 -> max_b = 32\n // b = 33, ..., 512 -> max_b = next power of 2\n // b = 513, ... -> max_b = increment by 512\n if (log2_b <= 5) {\n max_b = 32;\n } else if (log2_b <= 9) {\n max_b = pow(2, log2_b);\n } else {\n max_b = (batch_size + 511) / 512 * 512;\n }\n return max_b;\n}\n\n// quantize token count\nsize_t get_max_tokens(size_t num_tokens) {\n // token count is expected to be 1k's-100k's\n // t = 0, ..., 1024 -> max_t = 1024\n // t = 1025, ..., 32k -> max_t = next power of 2\n // t = 32k+1, ... -> max_t = increment by 32k\n size_t log2_t = ceil(log2(num_tokens));\n size_t max_t = 0;\n if (log2_t <= 10) {\n max_t = 1024;\n } else if (log2_t <= 15) {\n max_t = pow(2, log2_t);\n } else {\n max_t = (num_tokens + 32767) / 32768 * 32768;\n }\n return max_t;\n}\n\n__global__ void populate_rng_state_kernel(int64_t *rng_state_dst, const int64_t *const seed,\n int64_t offset) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n if (tid > 0) return;\n rng_state_dst[0] = seed[0];\n rng_state_dst[1] = offset;\n}\n\n__global__ void get_runtime_num_segments_kernel(int32_t *cu_seqlen, size_t len, uint32_t *out) {\n int tid = blockDim.x * blockIdx.x + threadIdx.x;\n if (tid >= len) return;\n\n if (cu_seqlen[tid] > 0) {\n // atomicAdd only support 32 bits dtype\n atomicAdd(out, 1);\n }\n}\n\nvoid PopulateRngStateAsync(void *rng_state_dst, const void *seed, size_t q_max_seqlen,\n size_t kv_max_seqlen, NVTE_Fused_Attn_Backend backend,\n hipStream_t stream) {\n size_t increment = 0;\n if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) {\n increment = 16;\n } else {\n constexpr int threads_per_cta = 128;\n increment = (q_max_seqlen * kv_max_seqlen + threads_per_cta - 1) / threads_per_cta;\n }\n auto offset = FusedAttnOffsetManager::Instance().GetAndUpdateOffset(increment);\n populate_rng_state_kernel<<<1, 1, 0, stream>>>(reinterpret_cast(rng_state_dst),\n reinterpret_cast(seed), offset);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nuint32_t GetRuntimeNumSegments(void *cu_seqlen, void *workspace, size_t len, hipStream_t stream) {\n // workspace size requires 4 bytes\n uint32_t *dout = static_cast(workspace);\n uint32_t hout{};\n NVTE_CHECK_CUDA(hipMemsetAsync(dout, 0, sizeof(uint32_t), stream));\n constexpr int threads = 128;\n const int blocks = (len - 1) / threads + 1;\n get_runtime_num_segments_kernel<<>>(static_cast(cu_seqlen),\n len, dout);\n NVTE_CHECK_CUDA(hipGetLastError());\n NVTE_CHECK_CUDA(hipMemcpyAsync(&hout, dout, sizeof(uint32_t), hipMemcpyDeviceToHost, stream));\n NVTE_CHECK_CUDA(hipStreamSynchronize(stream));\n return hout;\n}\n\n__global__ void extract_seed_and_offset(int64_t *rng_state_ptr, bool captured, int64_t *seed_ptr,\n uint64_t seed_val, int64_t *offset_ptr, uint64_t offset_val,\n uint32_t offset_intragraph) {\n if (captured) {\n rng_state_ptr[0] = *seed_ptr;\n rng_state_ptr[1] = static_cast(*offset_ptr + static_cast(offset_intragraph));\n } else {\n rng_state_ptr[0] = static_cast(seed_val);\n rng_state_ptr[1] = static_cast(offset_val);\n }\n}\n\n} // namespace fused_attn\n} // namespace transformer_engine\n\nvoid nvte_extract_seed_and_offset(int64_t *rng_state_ptr, int captured, int64_t *seed_ptr,\n uint64_t seed_val, int64_t *offset_ptr, uint64_t offset_val,\n uint32_t offset_intragraph, hipStream_t stream) {\n NVTE_API_CALL(nvte_extract_seed_and_offset);\n using namespace transformer_engine;\n\n fused_attn::extract_seed_and_offset<<<1, 1, 0, stream>>>(\n rng_state_ptr, captured, seed_ptr, seed_val, offset_ptr, offset_val, offset_intragraph);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/fused_attn/utils.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n\nnamespace {\n\nconstexpr int64_t kMaxKernelAddresses = 256;\n\nstruct HostPointersArgs {\n uint64_t ptrs[kMaxKernelAddresses];\n};\n\n__global__ void write_pointers_kernel(HostPointersArgs args, uint64_t *out, int64_t count,\n int64_t offset) {\n const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x;\n if (idx < count) {\n out[offset + idx] = args.ptrs[idx];\n }\n}\n\n} // namespace\n\nvoid nvte_convert_pointers_to_tensor(const uint64_t *host_ptrs, NVTETensor output, int64_t count,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_convert_pointers_to_tensor);\n using namespace transformer_engine;\n Tensor *out_tensor = convertNVTETensorCheck(output);\n uint64_t *out_ptr = static_cast(out_tensor->data.dptr);\n NVTE_CHECK(out_ptr != nullptr, \"Output tensor data pointer is null.\");\n\n int64_t offset = 0;\n while (offset < count) {\n const int64_t chunk = std::min(kMaxKernelAddresses, count - offset);\n HostPointersArgs args{};\n for (int64_t i = 0; i < chunk; ++i) {\n args.ptrs[i] = host_ptrs[offset + i];\n }\n constexpr int threads = kMaxKernelAddresses;\n write_pointers_kernel<<<1, threads, 0, stream>>>(args, out_ptr, chunk, offset);\n NVTE_CHECK_CUDA(cudaGetLastError());\n offset += chunk;\n }\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n\nnamespace {\n\nconstexpr int64_t kMaxKernelAddresses = 256;\n\nstruct HostPointersArgs {\n uint64_t ptrs[kMaxKernelAddresses];\n};\n\n__global__ void write_pointers_kernel(HostPointersArgs args, uint64_t *out, int64_t count,\n int64_t offset) {\n const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x;\n if (idx < count) {\n out[offset + idx] = args.ptrs[idx];\n }\n}\n\n} // namespace\n\nvoid nvte_convert_pointers_to_tensor(const uint64_t *host_ptrs, NVTETensor output, int64_t count,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_convert_pointers_to_tensor);\n using namespace transformer_engine;\n Tensor *out_tensor = convertNVTETensorCheck(output);\n uint64_t *out_ptr = static_cast(out_tensor->data.dptr);\n NVTE_CHECK(out_ptr != nullptr, \"Output tensor data pointer is null.\");\n\n int64_t offset = 0;\n while (offset < count) {\n const int64_t chunk = std::min(kMaxKernelAddresses, count - offset);\n HostPointersArgs args{};\n for (int64_t i = 0; i < chunk; ++i) {\n args.ptrs[i] = host_ptrs[offset + i];\n }\n constexpr int threads = kMaxKernelAddresses;\n write_pointers_kernel<<<1, threads, 0, stream>>>(args, out_ptr, chunk, offset);\n NVTE_CHECK_CUDA(hipGetLastError());\n offset += chunk;\n }\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/util/utils.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\n\nnamespace {\n\n// Parameters to tune\nconstexpr int n_warps_per_tile = 4;\nconstexpr int threads_per_block = THREADS_PER_WARP * n_warps_per_tile;\nconstexpr int desired_load_store_size = 8;\nconstexpr int kMaxTensorsPerKernel = 64; // Args must be <4 KB\n\nstruct MultiPaddingArgs {\n // (input) Data buffers for input tensors\n void* input_list[kMaxTensorsPerKernel];\n // (output) Data buffers for cast output tensors\n void* output_list[kMaxTensorsPerKernel];\n // Input matrix heights\n int num_rows_list[kMaxTensorsPerKernel];\n // Input matrix heights (padded)\n int padded_num_rows_list[kMaxTensorsPerKernel];\n // Input matrix widths\n int row_length_list[kMaxTensorsPerKernel];\n // Prefix sum (with leading zero) of CUDA blocks needed for each\n // tensor\n int block_range[kMaxTensorsPerKernel + 1];\n // Number of tensors being processed by kernel\n int num_tensors;\n};\n\ntemplate \n__global__ void __launch_bounds__(threads_per_block) multi_padding_kernel(MultiPaddingArgs args) {\n using Vec = Vec;\n\n // Thread indices\n // Note: Block is interpreted as a warp_size x num_warps grid\n constexpr int bdimx = THREADS_PER_WARP;\n constexpr int bdimy = n_warps_per_tile;\n const int tid = threadIdx.x;\n const int tidx = tid % bdimx;\n const int tidy = tid / bdimx;\n const int bid = blockIdx.x;\n\n // Input tensors are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec x nvec subtiles\n constexpr int tile_dim_m = THREADS_PER_WARP * nvec;\n constexpr int tile_dim_n = THREADS_PER_WARP * nvec;\n\n // Number of nvec x nvec subtiles for each thread to\n // load/store\n constexpr int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n\n // Find tensor corresponding to block\n int tensor_id = 0;\n while (args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n const Type* input = reinterpret_cast(args.input_list[tensor_id]);\n Type* output = reinterpret_cast(args.output_list[tensor_id]);\n const int num_rows = args.num_rows_list[tensor_id];\n const int padded_num_rows = args.padded_num_rows_list[tensor_id];\n const int row_length = args.row_length_list[tensor_id];\n\n // Find position of tile within tensor\n const int num_tiles_n = (row_length + tile_dim_n - 1) / tile_dim_n;\n const int tile_id = bid - args.block_range[tensor_id];\n const int tile_id_m = tile_id / num_tiles_n;\n const int tile_id_n = tile_id % num_tiles_n;\n const int tile_row = tile_id_m * tile_dim_m;\n const int tile_col = tile_id_n * tile_dim_n;\n\n // Load input and store to registers\n // Note: Each thread loads n_iterations subtiles, casts to output\n // type, and transposes in registers.\n Type local_zero = static_cast(0.f);\n for (int iter = 0; iter < n_iterations; ++iter) {\n const int i1 = tidy + iter * bdimy;\n const int j1 = tidx;\n#pragma unroll\n for (int i2 = 0; i2 < nvec; ++i2) {\n const int row = tile_row + i1 * nvec + i2;\n const int col = tile_col + j1 * nvec;\n Vec local_input;\n Vec local_output;\n local_input.clear();\n if (row < num_rows) {\n for (int j2 = 0; j2 < nvec; ++j2) {\n if (col + j2 < row_length) {\n local_input.data.elt[j2] = input[static_cast(row) * row_length + col + j2];\n }\n }\n }\n#pragma unroll\n for (int j2 = 0; j2 < nvec; ++j2) {\n local_output.data.elt[j2] = local_input.data.elt[j2];\n }\n if (row < num_rows) {\n for (int j2 = 0; j2 < nvec; ++j2) {\n if (col + j2 < row_length) {\n output[static_cast(row) * row_length + col + j2] = local_output.data.elt[j2];\n }\n }\n } else if (row < padded_num_rows) {\n // padding\n for (int j2 = 0; j2 < nvec; ++j2) {\n if (col + j2 < row_length) {\n output[static_cast(row) * row_length + col + j2] = local_zero;\n }\n }\n }\n }\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(threads_per_block) multi_unpadding_kernel(MultiPaddingArgs args) {\n using Vec = Vec;\n\n // Thread indices\n // Note: Block is interpreted as a warp_size x num_warps grid\n constexpr int bdimx = THREADS_PER_WARP;\n constexpr int bdimy = n_warps_per_tile;\n const int tid = threadIdx.x;\n const int tidx = tid % bdimx;\n const int tidy = tid / bdimx;\n const int bid = blockIdx.x;\n\n // Input tensors are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec x nvec subtiles\n constexpr int tile_dim_m = THREADS_PER_WARP * nvec;\n constexpr int tile_dim_n = THREADS_PER_WARP * nvec;\n\n // Number of nvec x nvec subtiles for each thread to\n // load/store\n constexpr int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n\n // Find tensor corresponding to block\n int tensor_id = 0;\n while (args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n const Type* input = reinterpret_cast(args.input_list[tensor_id]);\n Type* output = reinterpret_cast(args.output_list[tensor_id]);\n const int num_rows = args.num_rows_list[tensor_id];\n const int row_length = args.row_length_list[tensor_id];\n\n // Find position of tile within tensor\n const int num_tiles_n = (row_length + tile_dim_n - 1) / tile_dim_n;\n const int tile_id = bid - args.block_range[tensor_id];\n const int tile_id_m = tile_id / num_tiles_n;\n const int tile_id_n = tile_id % num_tiles_n;\n const int tile_row = tile_id_m * tile_dim_m;\n const int tile_col = tile_id_n * tile_dim_n;\n\n // Load input and store to registers\n // Note: Each thread loads n_iterations subtiles, casts to output\n // type, and transposes in registers.\n Type local_zero = static_cast(0.f);\n for (int iter = 0; iter < n_iterations; ++iter) {\n const int i1 = tidy + iter * bdimy;\n const int j1 = tidx;\n#pragma unroll\n for (int i2 = 0; i2 < nvec; ++i2) {\n const int row = tile_row + i1 * nvec + i2;\n const int col = tile_col + j1 * nvec;\n Vec local_input;\n Vec local_output;\n local_input.clear();\n if (row < num_rows) {\n for (int j2 = 0; j2 < nvec; ++j2) {\n if (col + j2 < row_length) {\n local_input.data.elt[j2] = input[static_cast(row) * row_length + col + j2];\n }\n }\n }\n#pragma unroll\n for (int j2 = 0; j2 < nvec; ++j2) {\n local_output.data.elt[j2] = local_input.data.elt[j2];\n }\n if (row < num_rows) {\n for (int j2 = 0; j2 < nvec; ++j2) {\n if (col + j2 < row_length) {\n output[static_cast(row) * row_length + col + j2] = local_output.data.elt[j2];\n }\n }\n }\n }\n }\n}\n\n} // namespace\n\nvoid multi_padding(const std::vector input_list, std::vector output_list,\n const std::vector padded_num_rows_list, cudaStream_t stream) {\n // Check that number of tensors is valid\n NVTE_CHECK(output_list.size() == input_list.size(),\n \"Number of input and output tensors must match\");\n if (input_list.empty()) {\n return;\n }\n\n // Check that tensor properties are valid\n DType type = input_list[0]->data.dtype;\n for (size_t tensor_id = 0; tensor_id < input_list.size(); ++tensor_id) {\n const auto& input = *input_list[tensor_id];\n const auto& output = *output_list[tensor_id];\n CheckInputTensor(input, \"multi_padding_input_\" + std::to_string(tensor_id));\n CheckInputTensor(output, \"multi_padding_output_\" + std::to_string(tensor_id));\n\n NVTE_CHECK(input.data.dtype == type, \"Input tensor types do not match.\");\n NVTE_CHECK(output.data.dtype == type, \"Output tensor types do not match.\");\n\n NVTE_CHECK(input.data.shape.size() == 2, \"Input tensor must have 2 dimensions.\");\n NVTE_CHECK(output.data.shape[0] == padded_num_rows_list[tensor_id],\n \"output tensor shape does not match padded input shape.\");\n }\n\n // Input matrices are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec x nvec subtiles\n const int tile_dim_m = THREADS_PER_WARP * desired_load_store_size * 8 / typeToNumBits(type);\n const int tile_dim_n = THREADS_PER_WARP * desired_load_store_size * 8 / typeToNumBits(type);\n\n // Add tensors to kernel argument struct\n MultiPaddingArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.block_range[0] = 0;\n for (size_t tensor_id = 0; tensor_id < input_list.size(); ++tensor_id) {\n // Launch kernel if argument struct is full\n if (kernel_args.num_tensors == kMaxTensorsPerKernel) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n type, Type, constexpr int nvec = desired_load_store_size / sizeof(Type);\n const int n_blocks = kernel_args.block_range[kernel_args.num_tensors];\n multi_padding_kernel\n <<>>(kernel_args);); // NOLINT(*)\n NVTE_CHECK_CUDA(cudaGetLastError());\n kernel_args.num_tensors = 0;\n }\n\n // Calculate number of thread blocks needed for tensor\n const int num_rows = input_list[tensor_id]->data.shape[0];\n const int padded_num_rows = padded_num_rows_list[tensor_id];\n const int row_length = input_list[tensor_id]->data.shape[1];\n const int num_tiles_m = (padded_num_rows + tile_dim_m - 1) / tile_dim_m;\n const int num_tiles_n = (row_length + tile_dim_n - 1) / tile_dim_n;\n const int num_tiles = num_tiles_m * num_tiles_n;\n\n // Add tensor to kernel argument struct\n const int pos = kernel_args.num_tensors;\n kernel_args.input_list[pos] = const_cast(input_list[tensor_id]->data.dptr);\n kernel_args.output_list[pos] = output_list[tensor_id]->data.dptr;\n kernel_args.num_rows_list[pos] = num_rows;\n kernel_args.padded_num_rows_list[pos] = padded_num_rows;\n kernel_args.row_length_list[pos] = row_length;\n kernel_args.block_range[pos + 1] = kernel_args.block_range[pos] + num_tiles;\n kernel_args.num_tensors++;\n }\n\n // Launch kernel\n if (kernel_args.num_tensors > 0) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n type, Type, constexpr int nvec = desired_load_store_size / sizeof(Type);\n const int n_blocks = kernel_args.block_range[kernel_args.num_tensors];\n multi_padding_kernel\n <<>>(kernel_args);); // NOLINT(*)\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\nvoid multi_unpadding(const std::vector input_list, std::vector output_list,\n const std::vector unpadded_num_rows_list, cudaStream_t stream) {\n // Check that number of tensors is valid\n NVTE_CHECK(output_list.size() == input_list.size(),\n \"Number of input and output tensors must match\");\n if (input_list.empty()) {\n return;\n }\n\n // Check that tensor properties are valid\n DType type = input_list[0]->data.dtype;\n for (size_t tensor_id = 0; tensor_id < input_list.size(); ++tensor_id) {\n const auto& input = *input_list[tensor_id];\n const auto& output = *output_list[tensor_id];\n CheckInputTensor(input, \"multi_unpadding_input_\" + std::to_string(tensor_id));\n CheckInputTensor(output, \"multi_unpadding_output_\" + std::to_string(tensor_id));\n\n NVTE_CHECK(input.data.dtype == type, \"Input tensor types do not match.\");\n NVTE_CHECK(output.data.dtype == type, \"Output tensor types do not match.\");\n\n NVTE_CHECK(input.data.shape.size() == 2, \"Input tensor must have 2 dimensions.\");\n NVTE_CHECK(output.data.shape[0] == unpadded_num_rows_list[tensor_id],\n \"output tensor shape does not match padded input shape.\");\n }\n\n // Input matrices are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec x nvec subtiles\n const int tile_dim_m = THREADS_PER_WARP * desired_load_store_size / typeToSize(type);\n const int tile_dim_n = THREADS_PER_WARP * desired_load_store_size / typeToSize(type);\n\n // Add tensors to kernel argument struct\n MultiPaddingArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.block_range[0] = 0;\n for (size_t tensor_id = 0; tensor_id < input_list.size(); ++tensor_id) {\n // Launch kernel if argument struct is full\n if (kernel_args.num_tensors == kMaxTensorsPerKernel) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n type, Type, constexpr int nvec = desired_load_store_size / sizeof(Type);\n const int n_blocks = kernel_args.block_range[kernel_args.num_tensors];\n multi_unpadding_kernel\n <<>>(kernel_args);); // NOLINT(*)\n NVTE_CHECK_CUDA(cudaGetLastError());\n kernel_args.num_tensors = 0;\n }\n\n // Calculate number of thread blocks needed for tensor\n const int num_rows = unpadded_num_rows_list[tensor_id];\n const int row_length = input_list[tensor_id]->data.shape[1];\n const int num_tiles_m = (num_rows + tile_dim_m - 1) / tile_dim_m;\n const int num_tiles_n = (row_length + tile_dim_n - 1) / tile_dim_n;\n const int num_tiles = num_tiles_m * num_tiles_n;\n\n // Add tensor to kernel argument struct\n const int pos = kernel_args.num_tensors;\n kernel_args.input_list[pos] = const_cast(input_list[tensor_id]->data.dptr);\n kernel_args.output_list[pos] = output_list[tensor_id]->data.dptr;\n kernel_args.num_rows_list[pos] = num_rows;\n kernel_args.row_length_list[pos] = row_length;\n kernel_args.block_range[pos + 1] = kernel_args.block_range[pos] + num_tiles;\n kernel_args.num_tensors++;\n }\n\n // Launch kernel\n if (kernel_args.num_tensors > 0) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n type, Type, constexpr int nvec = desired_load_store_size / sizeof(Type);\n const int n_blocks = kernel_args.block_range[kernel_args.num_tensors];\n multi_unpadding_kernel\n <<>>(kernel_args);); // NOLINT(*)\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\n} // namespace transformer_engine\n\nvoid nvte_multi_padding(size_t num_tensors, const NVTETensor* input_list, NVTETensor* output_list,\n const int* padded_num_rows_list, cudaStream_t stream) {\n NVTE_API_CALL(nvte_multi_padding);\n using namespace transformer_engine;\n std::vector input_list_, output_list_;\n std::vector padded_num_rows_list_;\n for (size_t i = 0; i < num_tensors; ++i) {\n input_list_.push_back(convertNVTETensorCheck(input_list[i]));\n output_list_.push_back(convertNVTETensorCheck(output_list[i]));\n padded_num_rows_list_.push_back(padded_num_rows_list[i]);\n }\n multi_padding(input_list_, output_list_, padded_num_rows_list_, stream);\n}\n\nvoid nvte_multi_unpadding(size_t num_tensors, const NVTETensor* input_list, NVTETensor* output_list,\n const int* unpadded_num_rows_list, cudaStream_t stream) {\n NVTE_API_CALL(nvte_multi_unpadding);\n using namespace transformer_engine;\n std::vector input_list_, output_list_;\n std::vector unpadded_num_rows_list_;\n for (size_t i = 0; i < num_tensors; ++i) {\n input_list_.push_back(convertNVTETensorCheck(input_list[i]));\n output_list_.push_back(convertNVTETensorCheck(output_list[i]));\n unpadded_num_rows_list_.push_back(unpadded_num_rows_list[i]);\n }\n multi_unpadding(input_list_, output_list_, unpadded_num_rows_list_, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\n\nnamespace {\n\n// Parameters to tune\nconstexpr int n_warps_per_tile = 4;\nconstexpr int threads_per_block = THREADS_PER_WARP * n_warps_per_tile;\nconstexpr int desired_load_store_size = 8;\nconstexpr int kMaxTensorsPerKernel = 64; // Args must be <4 KB\n\nstruct MultiPaddingArgs {\n // (input) Data buffers for input tensors\n void* input_list[kMaxTensorsPerKernel];\n // (output) Data buffers for cast output tensors\n void* output_list[kMaxTensorsPerKernel];\n // Input matrix heights\n int num_rows_list[kMaxTensorsPerKernel];\n // Input matrix heights (padded)\n int padded_num_rows_list[kMaxTensorsPerKernel];\n // Input matrix widths\n int row_length_list[kMaxTensorsPerKernel];\n // Prefix sum (with leading zero) of CUDA blocks needed for each\n // tensor\n int block_range[kMaxTensorsPerKernel + 1];\n // Number of tensors being processed by kernel\n int num_tensors;\n};\n\ntemplate \n__global__ void __launch_bounds__(threads_per_block) multi_padding_kernel(MultiPaddingArgs args) {\n using Vec = Vec;\n\n // Thread indices\n // Note: Block is interpreted as a warp_size x num_warps grid\n constexpr int bdimx = THREADS_PER_WARP;\n constexpr int bdimy = n_warps_per_tile;\n const int tid = threadIdx.x;\n const int tidx = tid % bdimx;\n const int tidy = tid / bdimx;\n const int bid = blockIdx.x;\n\n // Input tensors are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec x nvec subtiles\n constexpr int tile_dim_m = THREADS_PER_WARP * nvec;\n constexpr int tile_dim_n = THREADS_PER_WARP * nvec;\n\n // Number of nvec x nvec subtiles for each thread to\n // load/store\n constexpr int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n\n // Find tensor corresponding to block\n int tensor_id = 0;\n while (args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n const Type* input = reinterpret_cast(args.input_list[tensor_id]);\n Type* output = reinterpret_cast(args.output_list[tensor_id]);\n const int num_rows = args.num_rows_list[tensor_id];\n const int padded_num_rows = args.padded_num_rows_list[tensor_id];\n const int row_length = args.row_length_list[tensor_id];\n\n // Find position of tile within tensor\n const int num_tiles_n = (row_length + tile_dim_n - 1) / tile_dim_n;\n const int tile_id = bid - args.block_range[tensor_id];\n const int tile_id_m = tile_id / num_tiles_n;\n const int tile_id_n = tile_id % num_tiles_n;\n const int tile_row = tile_id_m * tile_dim_m;\n const int tile_col = tile_id_n * tile_dim_n;\n\n // Load input and store to registers\n // Note: Each thread loads n_iterations subtiles, casts to output\n // type, and transposes in registers.\n Type local_zero = static_cast(0.f);\n for (int iter = 0; iter < n_iterations; ++iter) {\n const int i1 = tidy + iter * bdimy;\n const int j1 = tidx;\n#pragma unroll\n for (int i2 = 0; i2 < nvec; ++i2) {\n const int row = tile_row + i1 * nvec + i2;\n const int col = tile_col + j1 * nvec;\n Vec local_input;\n Vec local_output;\n local_input.clear();\n if (row < num_rows) {\n for (int j2 = 0; j2 < nvec; ++j2) {\n if (col + j2 < row_length) {\n local_input.data.elt[j2] = input[static_cast(row) * row_length + col + j2];\n }\n }\n }\n#pragma unroll\n for (int j2 = 0; j2 < nvec; ++j2) {\n local_output.data.elt[j2] = local_input.data.elt[j2];\n }\n if (row < num_rows) {\n for (int j2 = 0; j2 < nvec; ++j2) {\n if (col + j2 < row_length) {\n output[static_cast(row) * row_length + col + j2] = local_output.data.elt[j2];\n }\n }\n } else if (row < padded_num_rows) {\n // padding\n for (int j2 = 0; j2 < nvec; ++j2) {\n if (col + j2 < row_length) {\n output[static_cast(row) * row_length + col + j2] = local_zero;\n }\n }\n }\n }\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(threads_per_block) multi_unpadding_kernel(MultiPaddingArgs args) {\n using Vec = Vec;\n\n // Thread indices\n // Note: Block is interpreted as a warp_size x num_warps grid\n constexpr int bdimx = THREADS_PER_WARP;\n constexpr int bdimy = n_warps_per_tile;\n const int tid = threadIdx.x;\n const int tidx = tid % bdimx;\n const int tidy = tid / bdimx;\n const int bid = blockIdx.x;\n\n // Input tensors are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec x nvec subtiles\n constexpr int tile_dim_m = THREADS_PER_WARP * nvec;\n constexpr int tile_dim_n = THREADS_PER_WARP * nvec;\n\n // Number of nvec x nvec subtiles for each thread to\n // load/store\n constexpr int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n\n // Find tensor corresponding to block\n int tensor_id = 0;\n while (args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n const Type* input = reinterpret_cast(args.input_list[tensor_id]);\n Type* output = reinterpret_cast(args.output_list[tensor_id]);\n const int num_rows = args.num_rows_list[tensor_id];\n const int row_length = args.row_length_list[tensor_id];\n\n // Find position of tile within tensor\n const int num_tiles_n = (row_length + tile_dim_n - 1) / tile_dim_n;\n const int tile_id = bid - args.block_range[tensor_id];\n const int tile_id_m = tile_id / num_tiles_n;\n const int tile_id_n = tile_id % num_tiles_n;\n const int tile_row = tile_id_m * tile_dim_m;\n const int tile_col = tile_id_n * tile_dim_n;\n\n // Load input and store to registers\n // Note: Each thread loads n_iterations subtiles, casts to output\n // type, and transposes in registers.\n Type local_zero = static_cast(0.f);\n for (int iter = 0; iter < n_iterations; ++iter) {\n const int i1 = tidy + iter * bdimy;\n const int j1 = tidx;\n#pragma unroll\n for (int i2 = 0; i2 < nvec; ++i2) {\n const int row = tile_row + i1 * nvec + i2;\n const int col = tile_col + j1 * nvec;\n Vec local_input;\n Vec local_output;\n local_input.clear();\n if (row < num_rows) {\n for (int j2 = 0; j2 < nvec; ++j2) {\n if (col + j2 < row_length) {\n local_input.data.elt[j2] = input[static_cast(row) * row_length + col + j2];\n }\n }\n }\n#pragma unroll\n for (int j2 = 0; j2 < nvec; ++j2) {\n local_output.data.elt[j2] = local_input.data.elt[j2];\n }\n if (row < num_rows) {\n for (int j2 = 0; j2 < nvec; ++j2) {\n if (col + j2 < row_length) {\n output[static_cast(row) * row_length + col + j2] = local_output.data.elt[j2];\n }\n }\n }\n }\n }\n}\n\n} // namespace\n\nvoid multi_padding(const std::vector input_list, std::vector output_list,\n const std::vector padded_num_rows_list, hipStream_t stream) {\n // Check that number of tensors is valid\n NVTE_CHECK(output_list.size() == input_list.size(),\n \"Number of input and output tensors must match\");\n if (input_list.empty()) {\n return;\n }\n\n // Check that tensor properties are valid\n DType type = input_list[0]->data.dtype;\n for (size_t tensor_id = 0; tensor_id < input_list.size(); ++tensor_id) {\n const auto& input = *input_list[tensor_id];\n const auto& output = *output_list[tensor_id];\n CheckInputTensor(input, \"multi_padding_input_\" + std::to_string(tensor_id));\n CheckInputTensor(output, \"multi_padding_output_\" + std::to_string(tensor_id));\n\n NVTE_CHECK(input.data.dtype == type, \"Input tensor types do not match.\");\n NVTE_CHECK(output.data.dtype == type, \"Output tensor types do not match.\");\n\n NVTE_CHECK(input.data.shape.size() == 2, \"Input tensor must have 2 dimensions.\");\n NVTE_CHECK(output.data.shape[0] == padded_num_rows_list[tensor_id],\n \"output tensor shape does not match padded input shape.\");\n }\n\n // Input matrices are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec x nvec subtiles\n const int tile_dim_m = THREADS_PER_WARP * desired_load_store_size * 8 / typeToNumBits(type);\n const int tile_dim_n = THREADS_PER_WARP * desired_load_store_size * 8 / typeToNumBits(type);\n\n // Add tensors to kernel argument struct\n MultiPaddingArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.block_range[0] = 0;\n for (size_t tensor_id = 0; tensor_id < input_list.size(); ++tensor_id) {\n // Launch kernel if argument struct is full\n if (kernel_args.num_tensors == kMaxTensorsPerKernel) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n type, Type, constexpr int nvec = desired_load_store_size / sizeof(Type);\n const int n_blocks = kernel_args.block_range[kernel_args.num_tensors];\n multi_padding_kernel\n <<>>(kernel_args);); // NOLINT(*)\n NVTE_CHECK_CUDA(hipGetLastError());\n kernel_args.num_tensors = 0;\n }\n\n // Calculate number of thread blocks needed for tensor\n const int num_rows = input_list[tensor_id]->data.shape[0];\n const int padded_num_rows = padded_num_rows_list[tensor_id];\n const int row_length = input_list[tensor_id]->data.shape[1];\n const int num_tiles_m = (padded_num_rows + tile_dim_m - 1) / tile_dim_m;\n const int num_tiles_n = (row_length + tile_dim_n - 1) / tile_dim_n;\n const int num_tiles = num_tiles_m * num_tiles_n;\n\n // Add tensor to kernel argument struct\n const int pos = kernel_args.num_tensors;\n kernel_args.input_list[pos] = const_cast(input_list[tensor_id]->data.dptr);\n kernel_args.output_list[pos] = output_list[tensor_id]->data.dptr;\n kernel_args.num_rows_list[pos] = num_rows;\n kernel_args.padded_num_rows_list[pos] = padded_num_rows;\n kernel_args.row_length_list[pos] = row_length;\n kernel_args.block_range[pos + 1] = kernel_args.block_range[pos] + num_tiles;\n kernel_args.num_tensors++;\n }\n\n // Launch kernel\n if (kernel_args.num_tensors > 0) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n type, Type, constexpr int nvec = desired_load_store_size / sizeof(Type);\n const int n_blocks = kernel_args.block_range[kernel_args.num_tensors];\n multi_padding_kernel\n <<>>(kernel_args);); // NOLINT(*)\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\nvoid multi_unpadding(const std::vector input_list, std::vector output_list,\n const std::vector unpadded_num_rows_list, hipStream_t stream) {\n // Check that number of tensors is valid\n NVTE_CHECK(output_list.size() == input_list.size(),\n \"Number of input and output tensors must match\");\n if (input_list.empty()) {\n return;\n }\n\n // Check that tensor properties are valid\n DType type = input_list[0]->data.dtype;\n for (size_t tensor_id = 0; tensor_id < input_list.size(); ++tensor_id) {\n const auto& input = *input_list[tensor_id];\n const auto& output = *output_list[tensor_id];\n CheckInputTensor(input, \"multi_unpadding_input_\" + std::to_string(tensor_id));\n CheckInputTensor(output, \"multi_unpadding_output_\" + std::to_string(tensor_id));\n\n NVTE_CHECK(input.data.dtype == type, \"Input tensor types do not match.\");\n NVTE_CHECK(output.data.dtype == type, \"Output tensor types do not match.\");\n\n NVTE_CHECK(input.data.shape.size() == 2, \"Input tensor must have 2 dimensions.\");\n NVTE_CHECK(output.data.shape[0] == unpadded_num_rows_list[tensor_id],\n \"output tensor shape does not match padded input shape.\");\n }\n\n // Input matrices are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec x nvec subtiles\n const int tile_dim_m = THREADS_PER_WARP * desired_load_store_size / typeToSize(type);\n const int tile_dim_n = THREADS_PER_WARP * desired_load_store_size / typeToSize(type);\n\n // Add tensors to kernel argument struct\n MultiPaddingArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.block_range[0] = 0;\n for (size_t tensor_id = 0; tensor_id < input_list.size(); ++tensor_id) {\n // Launch kernel if argument struct is full\n if (kernel_args.num_tensors == kMaxTensorsPerKernel) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n type, Type, constexpr int nvec = desired_load_store_size / sizeof(Type);\n const int n_blocks = kernel_args.block_range[kernel_args.num_tensors];\n multi_unpadding_kernel\n <<>>(kernel_args);); // NOLINT(*)\n NVTE_CHECK_CUDA(hipGetLastError());\n kernel_args.num_tensors = 0;\n }\n\n // Calculate number of thread blocks needed for tensor\n const int num_rows = unpadded_num_rows_list[tensor_id];\n const int row_length = input_list[tensor_id]->data.shape[1];\n const int num_tiles_m = (num_rows + tile_dim_m - 1) / tile_dim_m;\n const int num_tiles_n = (row_length + tile_dim_n - 1) / tile_dim_n;\n const int num_tiles = num_tiles_m * num_tiles_n;\n\n // Add tensor to kernel argument struct\n const int pos = kernel_args.num_tensors;\n kernel_args.input_list[pos] = const_cast(input_list[tensor_id]->data.dptr);\n kernel_args.output_list[pos] = output_list[tensor_id]->data.dptr;\n kernel_args.num_rows_list[pos] = num_rows;\n kernel_args.row_length_list[pos] = row_length;\n kernel_args.block_range[pos + 1] = kernel_args.block_range[pos] + num_tiles;\n kernel_args.num_tensors++;\n }\n\n // Launch kernel\n if (kernel_args.num_tensors > 0) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(\n type, Type, constexpr int nvec = desired_load_store_size / sizeof(Type);\n const int n_blocks = kernel_args.block_range[kernel_args.num_tensors];\n multi_unpadding_kernel\n <<>>(kernel_args);); // NOLINT(*)\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\n} // namespace transformer_engine\n\nvoid nvte_multi_padding(size_t num_tensors, const NVTETensor* input_list, NVTETensor* output_list,\n const int* padded_num_rows_list, hipStream_t stream) {\n NVTE_API_CALL(nvte_multi_padding);\n using namespace transformer_engine;\n std::vector input_list_, output_list_;\n std::vector padded_num_rows_list_;\n for (size_t i = 0; i < num_tensors; ++i) {\n input_list_.push_back(convertNVTETensorCheck(input_list[i]));\n output_list_.push_back(convertNVTETensorCheck(output_list[i]));\n padded_num_rows_list_.push_back(padded_num_rows_list[i]);\n }\n multi_padding(input_list_, output_list_, padded_num_rows_list_, stream);\n}\n\nvoid nvte_multi_unpadding(size_t num_tensors, const NVTETensor* input_list, NVTETensor* output_list,\n const int* unpadded_num_rows_list, hipStream_t stream) {\n NVTE_API_CALL(nvte_multi_unpadding);\n using namespace transformer_engine;\n std::vector input_list_, output_list_;\n std::vector unpadded_num_rows_list_;\n for (size_t i = 0; i < num_tensors; ++i) {\n input_list_.push_back(convertNVTETensorCheck(input_list[i]));\n output_list_.push_back(convertNVTETensorCheck(output_list[i]));\n unpadded_num_rows_list_.push_back(unpadded_num_rows_list[i]);\n }\n multi_unpadding(input_list_, output_list_, unpadded_num_rows_list_, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/util/padding.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\n\ntemplate \n__device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);\n\ntemplate <>\n__device__ __inline__ void copy_vector(bf16 *dst, const bf16 *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(bf16 *dst, const bf16 *src) {\n *((float2 *)dst) = *((float2 *)src); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(half *dst, const half *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(half *dst, const half *src) {\n *((float2 *)dst) = *((float2 *)src); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(uint8_t *dst, const uint8_t *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(uint8_t *dst, const uint8_t *src) {\n *((half2 *)dst) = *((half2 *)src); // NOLINT(*)\n}\n\ntemplate \nstruct Add {\n __device__ __forceinline__ T operator()(T a, T b) const { return a + b; }\n};\n\ntemplate \nstruct Max {\n __device__ __forceinline__ T operator()(T a, T b) const { return a < b ? b : a; }\n};\n\ntemplate \n__device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize,\n unsigned int mask = 0xffffffff) {\n#if CUDA_VERSION >= 9000\n return __shfl_xor_sync(mask, value, laneMask, width);\n#else\n return __shfl_xor(value, laneMask, width);\n#endif\n}\n\ntemplate class ReduceOp>\n__device__ __forceinline__ void warp_reduce(acc_t *sum) {\n ReduceOp r;\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n acc_t b = WARP_SHFL_XOR_NATIVE(sum[i], offset, WARP_SIZE);\n sum[i] = r(sum[i], b);\n }\n }\n}\n\n/*\n * Extended softmax (from native aten pytorch) with following additional features\n * 1) input scaling\n */\ntemplate \n__global__ void scaled_softmax_warp_forward(output_t *dst, const input_t *src, const acc_t scale,\n int micro_batch_size, int element_count) {\n // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and\n // warp_size of method warp_softmax_forward_kernel.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_SIZE =\n (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;\n constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;\n constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;\n\n // blockDim/threadIdx = (WARP_SIZE, WARPS_PER_BLOCK, )\n // gridDim/blockIdx = (seq_len, attn_heads, batches)\n size_t first_batch =\n (blockDim.y * (blockIdx.x + gridDim.x * (blockIdx.y + gridDim.y * blockIdx.z)) +\n threadIdx.y) *\n WARP_BATCH;\n\n // micro_batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = micro_batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the batch\n int local_idx = threadIdx.x;\n\n size_t thread_offset = first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;\n src += thread_offset;\n dst += thread_offset;\n\n // load data from global memory\n acc_t elements[WARP_BATCH][WARP_ITERATIONS];\n input_t temp_data[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n\n if (element_index < batch_element_count) {\n int itr_idx = i * element_count + it * WARP_SIZE;\n copy_vector(temp_data, src + itr_idx);\n\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n elements[i][it + element] = (acc_t)temp_data[element] * scale;\n }\n } else {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n elements[i][it + element] = -std::numeric_limits::infinity();\n }\n }\n }\n }\n\n // compute max_value\n acc_t max_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = elements[i][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];\n }\n }\n warp_reduce(max_value);\n\n acc_t sum[WARP_BATCH]{0.0f};\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[i][it] = std::exp((elements[i][it] - max_value[i]));\n sum[i] += elements[i][it];\n }\n }\n warp_reduce(sum);\n\n // store result\n output_t out[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = elements[i][it + element] / sum[i];\n }\n copy_vector(dst + i * element_count + it * WARP_SIZE, out);\n } else {\n break;\n }\n }\n }\n}\n\n/*\n * Extended softmax (from native aten pytorch) with following additional features\n * 1) input scaling\n * 2) Explicit masking\n */\ntemplate \n__global__ void scaled_masked_softmax_warp_forward(output_t *dst, const input_t *src,\n const uint8_t *mask, const acc_t scale,\n int micro_batch_size, int element_count,\n int pad_batches) {\n // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and\n // warp_size of method warp_softmax_forward_kernel.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_SIZE =\n (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;\n constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;\n constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;\n\n // blockDim/threadIdx = (WARP_SIZE, WARPS_PER_BLOCK, )\n // gridDim/blockIdx = (seq_len, attn_heads, batches)\n size_t first_batch =\n (blockDim.y * (blockIdx.x + gridDim.x * (blockIdx.y + gridDim.y * blockIdx.z)) +\n threadIdx.y) *\n WARP_BATCH;\n size_t pad_first_batch = 0;\n if (pad_batches != 1) { // bert style\n pad_first_batch =\n (blockDim.y * (blockIdx.x + gridDim.x * blockIdx.z) + threadIdx.y) * WARP_BATCH;\n } else { // gpt2 style\n pad_first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n }\n\n // micro_batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = micro_batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the batch\n int local_idx = threadIdx.x;\n\n size_t thread_offset_src_dst = first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;\n size_t thread_offset_mask = pad_first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;\n src += thread_offset_src_dst;\n dst += thread_offset_src_dst;\n mask += thread_offset_mask;\n\n // load data from global memory\n acc_t elements[WARP_BATCH][WARP_ITERATIONS];\n input_t temp_data[ELEMENTS_PER_LDG_STG];\n uint8_t temp_mask[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n\n if (element_index < batch_element_count) {\n int itr_idx = i * element_count + it * WARP_SIZE;\n copy_vector(temp_data, src + itr_idx);\n copy_vector(temp_mask, mask + itr_idx);\n\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if (temp_mask[element] != 1) {\n elements[i][it + element] = (acc_t)temp_data[element] * scale;\n } else {\n elements[i][it + element] = -10000.0;\n }\n }\n } else {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n elements[i][it + element] = -std::numeric_limits::infinity();\n }\n }\n }\n }\n\n // compute max_value\n acc_t max_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = elements[i][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];\n }\n }\n warp_reduce(max_value);\n\n // compute scale value to account for full mask\n acc_t scale_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n scale_value[i] = (max_value[i] == -10000.0) ? 0.0 : 1.0;\n }\n\n acc_t sum[WARP_BATCH]{0.0f};\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[i][it] = std::exp((elements[i][it] - max_value[i]));\n sum[i] += elements[i][it];\n }\n }\n warp_reduce(sum);\n\n // store result\n output_t out[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = elements[i][it + element] * scale_value[i] / sum[i];\n }\n copy_vector(dst + i * element_count + it * WARP_SIZE, out);\n } else {\n break;\n }\n }\n }\n}\n\ntemplate \n__global__ void scaled_masked_softmax_warp_backward(output_t *gradInput, const input_t *grad,\n const input_t *output, acc_t scale,\n int micro_batch_size, int element_count) {\n // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and\n // warp_size of method warp_softmax_backward_kernel.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_SIZE =\n (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;\n constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;\n constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;\n\n // blockDim/threadIdx = (WARP_SIZE, WARPS_PER_BLOCK, )\n // gridDim/blockIdx = (seq_len, attn_heads, batches)\n size_t first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n\n // micro_batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = micro_batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the batch\n int local_idx = threadIdx.x;\n\n // the first element to process by the current thread\n size_t thread_offset = first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;\n grad += thread_offset;\n output += thread_offset;\n gradInput += thread_offset;\n\n // load data from global memory\n acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS]{0.0f};\n acc_t output_reg[WARP_BATCH][WARP_ITERATIONS]{0.0f};\n input_t temp_grad[ELEMENTS_PER_LDG_STG];\n input_t temp_output[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < batch_element_count) {\n copy_vector(temp_grad,\n grad + i * element_count + it * WARP_SIZE);\n copy_vector(temp_output,\n output + i * element_count + it * WARP_SIZE);\n\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n output_reg[i][it + element] = (acc_t)temp_output[element];\n }\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n grad_reg[i][it + element] = (acc_t)temp_grad[element] * output_reg[i][it + element];\n }\n }\n }\n }\n\n acc_t sum[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] = grad_reg[i][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n sum[i] += grad_reg[i][it];\n }\n }\n warp_reduce(sum);\n\n // store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n // compute gradients\n output_t out[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = (output_t)(scale * (grad_reg[i][it + element] -\n output_reg[i][it + element] * sum[i]));\n }\n copy_vector(gradInput + i * element_count + it * WARP_SIZE,\n out);\n }\n }\n }\n}\n\ntemplate \nvoid dispatch_scaled_softmax_forward(output_t *dst, const input_t *src, const input_t scale,\n int query_seq_len, int key_seq_len, int batches,\n int attn_heads, cudaStream_t stream) {\n NVTE_CHECK(key_seq_len >= 0 && key_seq_len <= 16384, \"Unsupported shape.\");\n if (key_seq_len == 0) {\n return;\n } else {\n int log2_elements = log2_ceil(key_seq_len);\n const int next_power_of_two = 1 << log2_elements;\n int batch_count = batches * attn_heads * query_seq_len;\n\n // This value must match the WARP_SIZE constexpr\n // value computed inside softmax_warp_forward.\n int warp_size = (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n\n // This value must match the WARP_BATCH constexpr\n // value computed inside softmax_warp_forward.\n int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximimize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n NVTE_CHECK(query_seq_len % batches_per_block == 0, \"Unsupported shape.\");\n dim3 blocks(query_seq_len / batches_per_block, attn_heads, batches);\n dim3 threads(warp_size, warps_per_block, 1);\n // Launch code would be more elegant if C++ supported FOR CONSTEXPR\n switch (log2_elements) {\n case 0: // 1\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 1: // 2\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 2: // 4\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 3: // 8\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 4: // 16\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 5: // 32\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 6: // 64\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 7: // 128\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 8: // 256\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 9: // 512\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 10: // 1024\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 11: // 2048\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 12: // 4096\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 13: // 8192\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 14: // 16384\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n default:\n break;\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\ntemplate \nvoid dispatch_scaled_masked_softmax_forward(output_t *dst, const input_t *src, const uint8_t *mask,\n const input_t scale, int query_seq_len, int key_seq_len,\n int batches, int attn_heads, int pad_batches,\n cudaStream_t stream) {\n NVTE_CHECK(key_seq_len >= 0 && key_seq_len <= 16384, \"Unsupported shape.\");\n if (key_seq_len == 0) {\n return;\n } else {\n int log2_elements = log2_ceil(key_seq_len);\n const int next_power_of_two = 1 << log2_elements;\n int batch_count = batches * attn_heads * query_seq_len;\n\n // This value must match the WARP_SIZE constexpr\n // value computed inside softmax_warp_forward.\n int warp_size = (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n\n // This value must match the WARP_BATCH constexpr\n // value computed inside softmax_warp_forward.\n int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximimize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n NVTE_CHECK(query_seq_len % batches_per_block == 0, \"Unsupported shape.\");\n dim3 blocks(query_seq_len / batches_per_block, attn_heads, batches);\n dim3 threads(warp_size, warps_per_block, 1);\n // Launch code would be more elegant if C++ supported FOR CONSTEXPR\n switch (log2_elements) {\n case 0: // 1\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 1: // 2\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 2: // 4\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 3: // 8\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 4: // 16\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 5: // 32\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 6: // 64\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 7: // 128\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 8: // 256\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 9: // 512\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 10: // 1024\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 11: // 2048\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 12: // 4096\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 13: // 8192\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 14: // 16384\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n default:\n break;\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\ntemplate \nvoid dispatch_scaled_masked_softmax_backward(output_t *grad_input, const input_t *grad,\n const input_t *output, const acc_t scale,\n int query_seq_len, int key_seq_len, int batches,\n int attn_heads, cudaStream_t stream) {\n NVTE_CHECK(key_seq_len >= 0 && key_seq_len <= 16384, \"Unsupported shape.\");\n if (key_seq_len == 0) {\n return;\n } else {\n int log2_elements = log2_ceil(key_seq_len);\n const int next_power_of_two = 1 << log2_elements;\n int batch_count = batches * attn_heads * query_seq_len;\n\n // This value must match the WARP_SIZE constexpr\n // value computed inside softmax_warp_backward.\n int warp_size = (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n\n // This value must match the WARP_BATCH constexpr\n // value computed inside softmax_warp_backward.\n int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximimize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = batch_count / batches_per_block;\n dim3 threads(warp_size, warps_per_block, 1);\n // Launch code would be more elegant if C++ supported FOR CONSTEXPR\n switch (log2_elements) {\n case 0: // 1\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 1: // 2\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 2: // 4\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 3: // 8\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 4: // 16\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 5: // 32\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 6: // 64\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 7: // 128\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 8: // 256\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 9: // 512\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 10: // 1024\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 11: // 2048\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 12: // 4096\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 13: // 8192\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 14: // 16384\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n default:\n break;\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\nvoid scaled_softmax_forward(const Tensor &input, Tensor *softmax_results, float scale_factor,\n cudaStream_t stream) {\n const int batches = input.data.shape[0];\n const int attn_heads = input.data.shape[1];\n const int query_seq_len = input.data.shape[2];\n const int key_seq_len = input.data.shape[3];\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n input.data.dtype, softmax_type,\n dispatch_scaled_softmax_forward(\n reinterpret_cast(softmax_results->data.dptr),\n reinterpret_cast(input.data.dptr), scale_factor, query_seq_len,\n key_seq_len, batches, attn_heads, stream););\n}\n\nvoid scaled_softmax_backward(Tensor output_grads, const Tensor incoming_grads,\n const Tensor softmax_results, float scale_factor,\n cudaStream_t stream) {\n // output grads is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len]\n const int batches = output_grads.data.shape[0];\n const int attn_heads = output_grads.data.shape[1];\n const int query_seq_len = output_grads.data.shape[2];\n const int key_seq_len = output_grads.data.shape[3];\n\n // Softmax Grad\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n output_grads.data.dtype, softmax_type,\n dispatch_scaled_masked_softmax_backward(\n reinterpret_cast(output_grads.data.dptr),\n reinterpret_cast(incoming_grads.data.dptr),\n reinterpret_cast(softmax_results.data.dptr), scale_factor,\n query_seq_len, key_seq_len, batches, attn_heads, stream););\n}\n\nvoid scaled_masked_softmax_forward(const Tensor input, const Tensor mask, Tensor *softmax_results,\n float scale_factor, cudaStream_t stream) {\n const int batches = input.data.shape[0];\n const int pad_batches = mask.data.shape[0];\n const int attn_heads = input.data.shape[1];\n const int query_seq_len = input.data.shape[2];\n const int key_seq_len = input.data.shape[3];\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n input.data.dtype, softmax_type,\n dispatch_scaled_masked_softmax_forward(\n reinterpret_cast(softmax_results->data.dptr),\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(mask.data.dptr), scale_factor, query_seq_len,\n key_seq_len, batches, attn_heads, pad_batches, stream););\n}\n\nvoid scaled_masked_softmax_backward(Tensor output_grads, const Tensor incoming_grads,\n const Tensor softmax_results, float scale_factor,\n cudaStream_t stream) {\n // output grads is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len]\n const int batches = output_grads.data.shape[0];\n const int attn_heads = output_grads.data.shape[1];\n const int query_seq_len = output_grads.data.shape[2];\n const int key_seq_len = output_grads.data.shape[3];\n\n // Softmax Grad\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n output_grads.data.dtype, softmax_type,\n dispatch_scaled_masked_softmax_backward(\n reinterpret_cast(output_grads.data.dptr),\n reinterpret_cast(incoming_grads.data.dptr),\n reinterpret_cast(softmax_results.data.dptr), scale_factor,\n query_seq_len, key_seq_len, batches, attn_heads, stream););\n}\n\n} // end namespace transformer_engine\n\nvoid nvte_scaled_softmax_forward(const NVTETensor input, NVTETensor softmax_results,\n float scale_factor, cudaStream_t stream) {\n NVTE_API_CALL(nvte_scaled_softmax_forward);\n using namespace transformer_engine;\n scaled_softmax_forward(*convertNVTETensorCheck(input), convertNVTETensorCheck(softmax_results),\n scale_factor, stream);\n}\n\nvoid nvte_scaled_softmax_backward(const NVTETensor incoming_grads, const NVTETensor softmax_results,\n NVTETensor output_grads, float scale_factor,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_scaled_softmax_backward);\n using namespace transformer_engine;\n scaled_softmax_backward(*convertNVTETensorCheck(output_grads),\n *convertNVTETensorCheck(incoming_grads),\n *convertNVTETensorCheck(softmax_results), scale_factor, stream);\n}\n\nvoid nvte_scaled_masked_softmax_forward(const NVTETensor input, const NVTETensor mask,\n NVTETensor softmax_results, float scale_factor,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_scaled_masked_softmax_forward);\n using namespace transformer_engine;\n scaled_masked_softmax_forward(*convertNVTETensorCheck(input), *convertNVTETensorCheck(mask),\n convertNVTETensorCheck(softmax_results), scale_factor, stream);\n}\n\nvoid nvte_scaled_masked_softmax_backward(const NVTETensor incoming_grads,\n const NVTETensor softmax_results, NVTETensor output_grads,\n float scale_factor, cudaStream_t stream) {\n NVTE_API_CALL(nvte_scaled_masked_softmax_backward);\n using namespace transformer_engine;\n scaled_masked_softmax_backward(*convertNVTETensorCheck(output_grads),\n *convertNVTETensorCheck(incoming_grads),\n *convertNVTETensorCheck(softmax_results), scale_factor, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\n\ntemplate \n__device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);\n\ntemplate <>\n__device__ __inline__ void copy_vector(bf16 *dst, const bf16 *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(bf16 *dst, const bf16 *src) {\n *((float2 *)dst) = *((float2 *)src); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(half *dst, const half *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(half *dst, const half *src) {\n *((float2 *)dst) = *((float2 *)src); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(uint8_t *dst, const uint8_t *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(uint8_t *dst, const uint8_t *src) {\n *((half2 *)dst) = *((half2 *)src); // NOLINT(*)\n}\n\ntemplate \nstruct Add {\n __device__ __forceinline__ T operator()(T a, T b) const { return a + b; }\n};\n\ntemplate \nstruct Max {\n __device__ __forceinline__ T operator()(T a, T b) const { return a < b ? b : a; }\n};\n\ntemplate \n__device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize,\n unsigned int mask = 0xffffffff) {\n#if CUDA_VERSION >= 9000\n return __shfl_xor_sync(mask, value, laneMask, width);\n#else\n return __shfl_xor(value, laneMask, width);\n#endif\n}\n\ntemplate class ReduceOp>\n__device__ __forceinline__ void warp_reduce(acc_t *sum) {\n ReduceOp r;\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n acc_t b = WARP_SHFL_XOR_NATIVE(sum[i], offset, WARP_SIZE);\n sum[i] = r(sum[i], b);\n }\n }\n}\n\n/*\n * Extended softmax (from native aten pytorch) with following additional features\n * 1) input scaling\n */\ntemplate \n__global__ void scaled_softmax_warp_forward(output_t *dst, const input_t *src, const acc_t scale,\n int micro_batch_size, int element_count) {\n // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and\n // warp_size of method warp_softmax_forward_kernel.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_SIZE =\n (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;\n constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;\n constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;\n\n // blockDim/threadIdx = (WARP_SIZE, WARPS_PER_BLOCK, )\n // gridDim/blockIdx = (seq_len, attn_heads, batches)\n size_t first_batch =\n (blockDim.y * (blockIdx.x + gridDim.x * (blockIdx.y + gridDim.y * blockIdx.z)) +\n threadIdx.y) *\n WARP_BATCH;\n\n // micro_batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = micro_batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the batch\n int local_idx = threadIdx.x;\n\n size_t thread_offset = first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;\n src += thread_offset;\n dst += thread_offset;\n\n // load data from global memory\n acc_t elements[WARP_BATCH][WARP_ITERATIONS];\n input_t temp_data[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n\n if (element_index < batch_element_count) {\n int itr_idx = i * element_count + it * WARP_SIZE;\n copy_vector(temp_data, src + itr_idx);\n\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n elements[i][it + element] = (acc_t)temp_data[element] * scale;\n }\n } else {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n elements[i][it + element] = -std::numeric_limits::infinity();\n }\n }\n }\n }\n\n // compute max_value\n acc_t max_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = elements[i][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];\n }\n }\n warp_reduce(max_value);\n\n acc_t sum[WARP_BATCH]{0.0f};\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[i][it] = std::exp((elements[i][it] - max_value[i]));\n sum[i] += elements[i][it];\n }\n }\n warp_reduce(sum);\n\n // store result\n output_t out[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = elements[i][it + element] / sum[i];\n }\n copy_vector(dst + i * element_count + it * WARP_SIZE, out);\n } else {\n break;\n }\n }\n }\n}\n\n/*\n * Extended softmax (from native aten pytorch) with following additional features\n * 1) input scaling\n * 2) Explicit masking\n */\ntemplate \n__global__ void scaled_masked_softmax_warp_forward(output_t *dst, const input_t *src,\n const uint8_t *mask, const acc_t scale,\n int micro_batch_size, int element_count,\n int pad_batches) {\n // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and\n // warp_size of method warp_softmax_forward_kernel.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_SIZE =\n (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;\n constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;\n constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;\n\n // blockDim/threadIdx = (WARP_SIZE, WARPS_PER_BLOCK, )\n // gridDim/blockIdx = (seq_len, attn_heads, batches)\n size_t first_batch =\n (blockDim.y * (blockIdx.x + gridDim.x * (blockIdx.y + gridDim.y * blockIdx.z)) +\n threadIdx.y) *\n WARP_BATCH;\n size_t pad_first_batch = 0;\n if (pad_batches != 1) { // bert style\n pad_first_batch =\n (blockDim.y * (blockIdx.x + gridDim.x * blockIdx.z) + threadIdx.y) * WARP_BATCH;\n } else { // gpt2 style\n pad_first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n }\n\n // micro_batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = micro_batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the batch\n int local_idx = threadIdx.x;\n\n size_t thread_offset_src_dst = first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;\n size_t thread_offset_mask = pad_first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;\n src += thread_offset_src_dst;\n dst += thread_offset_src_dst;\n mask += thread_offset_mask;\n\n // load data from global memory\n acc_t elements[WARP_BATCH][WARP_ITERATIONS];\n input_t temp_data[ELEMENTS_PER_LDG_STG];\n uint8_t temp_mask[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n\n if (element_index < batch_element_count) {\n int itr_idx = i * element_count + it * WARP_SIZE;\n copy_vector(temp_data, src + itr_idx);\n copy_vector(temp_mask, mask + itr_idx);\n\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if (temp_mask[element] != 1) {\n elements[i][it + element] = (acc_t)temp_data[element] * scale;\n } else {\n elements[i][it + element] = -10000.0;\n }\n }\n } else {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n elements[i][it + element] = -std::numeric_limits::infinity();\n }\n }\n }\n }\n\n // compute max_value\n acc_t max_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = elements[i][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];\n }\n }\n warp_reduce(max_value);\n\n // compute scale value to account for full mask\n acc_t scale_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n scale_value[i] = (max_value[i] == -10000.0) ? 0.0 : 1.0;\n }\n\n acc_t sum[WARP_BATCH]{0.0f};\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[i][it] = std::exp((elements[i][it] - max_value[i]));\n sum[i] += elements[i][it];\n }\n }\n warp_reduce(sum);\n\n // store result\n output_t out[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = elements[i][it + element] * scale_value[i] / sum[i];\n }\n copy_vector(dst + i * element_count + it * WARP_SIZE, out);\n } else {\n break;\n }\n }\n }\n}\n\ntemplate \n__global__ void scaled_masked_softmax_warp_backward(output_t *gradInput, const input_t *grad,\n const input_t *output, acc_t scale,\n int micro_batch_size, int element_count) {\n // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and\n // warp_size of method warp_softmax_backward_kernel.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_SIZE =\n (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;\n constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;\n constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;\n\n // blockDim/threadIdx = (WARP_SIZE, WARPS_PER_BLOCK, )\n // gridDim/blockIdx = (seq_len, attn_heads, batches)\n size_t first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n\n // micro_batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = micro_batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the batch\n int local_idx = threadIdx.x;\n\n // the first element to process by the current thread\n size_t thread_offset = first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;\n grad += thread_offset;\n output += thread_offset;\n gradInput += thread_offset;\n\n // load data from global memory\n acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS]{0.0f};\n acc_t output_reg[WARP_BATCH][WARP_ITERATIONS]{0.0f};\n input_t temp_grad[ELEMENTS_PER_LDG_STG];\n input_t temp_output[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < batch_element_count) {\n copy_vector(temp_grad,\n grad + i * element_count + it * WARP_SIZE);\n copy_vector(temp_output,\n output + i * element_count + it * WARP_SIZE);\n\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n output_reg[i][it + element] = (acc_t)temp_output[element];\n }\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n grad_reg[i][it + element] = (acc_t)temp_grad[element] * output_reg[i][it + element];\n }\n }\n }\n }\n\n acc_t sum[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] = grad_reg[i][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n sum[i] += grad_reg[i][it];\n }\n }\n warp_reduce(sum);\n\n // store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n // compute gradients\n output_t out[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = (output_t)(scale * (grad_reg[i][it + element] -\n output_reg[i][it + element] * sum[i]));\n }\n copy_vector(gradInput + i * element_count + it * WARP_SIZE,\n out);\n }\n }\n }\n}\n\ntemplate \nvoid dispatch_scaled_softmax_forward(output_t *dst, const input_t *src, const input_t scale,\n int query_seq_len, int key_seq_len, int batches,\n int attn_heads, hipStream_t stream) {\n NVTE_CHECK(key_seq_len >= 0 && key_seq_len <= 16384, \"Unsupported shape.\");\n if (key_seq_len == 0) {\n return;\n } else {\n int log2_elements = log2_ceil(key_seq_len);\n const int next_power_of_two = 1 << log2_elements;\n int batch_count = batches * attn_heads * query_seq_len;\n\n // This value must match the WARP_SIZE constexpr\n // value computed inside softmax_warp_forward.\n int warp_size = (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n\n // This value must match the WARP_BATCH constexpr\n // value computed inside softmax_warp_forward.\n int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximimize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n NVTE_CHECK(query_seq_len % batches_per_block == 0, \"Unsupported shape.\");\n dim3 blocks(query_seq_len / batches_per_block, attn_heads, batches);\n dim3 threads(warp_size, warps_per_block, 1);\n // Launch code would be more elegant if C++ supported FOR CONSTEXPR\n switch (log2_elements) {\n case 0: // 1\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 1: // 2\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 2: // 4\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 3: // 8\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 4: // 16\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 5: // 32\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 6: // 64\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 7: // 128\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 8: // 256\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 9: // 512\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 10: // 1024\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 11: // 2048\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 12: // 4096\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 13: // 8192\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n case 14: // 16384\n scaled_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, key_seq_len);\n break;\n default:\n break;\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\ntemplate \nvoid dispatch_scaled_masked_softmax_forward(output_t *dst, const input_t *src, const uint8_t *mask,\n const input_t scale, int query_seq_len, int key_seq_len,\n int batches, int attn_heads, int pad_batches,\n hipStream_t stream) {\n NVTE_CHECK(key_seq_len >= 0 && key_seq_len <= 16384, \"Unsupported shape.\");\n if (key_seq_len == 0) {\n return;\n } else {\n int log2_elements = log2_ceil(key_seq_len);\n const int next_power_of_two = 1 << log2_elements;\n int batch_count = batches * attn_heads * query_seq_len;\n\n // This value must match the WARP_SIZE constexpr\n // value computed inside softmax_warp_forward.\n int warp_size = (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n\n // This value must match the WARP_BATCH constexpr\n // value computed inside softmax_warp_forward.\n int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximimize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n NVTE_CHECK(query_seq_len % batches_per_block == 0, \"Unsupported shape.\");\n dim3 blocks(query_seq_len / batches_per_block, attn_heads, batches);\n dim3 threads(warp_size, warps_per_block, 1);\n // Launch code would be more elegant if C++ supported FOR CONSTEXPR\n switch (log2_elements) {\n case 0: // 1\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 1: // 2\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 2: // 4\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 3: // 8\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 4: // 16\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 5: // 32\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 6: // 64\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 7: // 128\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 8: // 256\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 9: // 512\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 10: // 1024\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 11: // 2048\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 12: // 4096\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 13: // 8192\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n case 14: // 16384\n scaled_masked_softmax_warp_forward\n <<>>(dst, src, mask, scale, batch_count, key_seq_len,\n pad_batches);\n break;\n default:\n break;\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\ntemplate \nvoid dispatch_scaled_masked_softmax_backward(output_t *grad_input, const input_t *grad,\n const input_t *output, const acc_t scale,\n int query_seq_len, int key_seq_len, int batches,\n int attn_heads, hipStream_t stream) {\n NVTE_CHECK(key_seq_len >= 0 && key_seq_len <= 16384, \"Unsupported shape.\");\n if (key_seq_len == 0) {\n return;\n } else {\n int log2_elements = log2_ceil(key_seq_len);\n const int next_power_of_two = 1 << log2_elements;\n int batch_count = batches * attn_heads * query_seq_len;\n\n // This value must match the WARP_SIZE constexpr\n // value computed inside softmax_warp_backward.\n int warp_size = (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n\n // This value must match the WARP_BATCH constexpr\n // value computed inside softmax_warp_backward.\n int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximimize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = batch_count / batches_per_block;\n dim3 threads(warp_size, warps_per_block, 1);\n // Launch code would be more elegant if C++ supported FOR CONSTEXPR\n switch (log2_elements) {\n case 0: // 1\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 1: // 2\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 2: // 4\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 3: // 8\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 4: // 16\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 5: // 32\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 6: // 64\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 7: // 128\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 8: // 256\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 9: // 512\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 10: // 1024\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 11: // 2048\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 12: // 4096\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 13: // 8192\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n case 14: // 16384\n scaled_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n key_seq_len);\n break;\n default:\n break;\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\nvoid scaled_softmax_forward(const Tensor &input, Tensor *softmax_results, float scale_factor,\n hipStream_t stream) {\n const int batches = input.data.shape[0];\n const int attn_heads = input.data.shape[1];\n const int query_seq_len = input.data.shape[2];\n const int key_seq_len = input.data.shape[3];\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n input.data.dtype, softmax_type,\n dispatch_scaled_softmax_forward(\n reinterpret_cast(softmax_results->data.dptr),\n reinterpret_cast(input.data.dptr), scale_factor, query_seq_len,\n key_seq_len, batches, attn_heads, stream););\n}\n\nvoid scaled_softmax_backward(Tensor output_grads, const Tensor incoming_grads,\n const Tensor softmax_results, float scale_factor,\n hipStream_t stream) {\n // output grads is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len]\n const int batches = output_grads.data.shape[0];\n const int attn_heads = output_grads.data.shape[1];\n const int query_seq_len = output_grads.data.shape[2];\n const int key_seq_len = output_grads.data.shape[3];\n\n // Softmax Grad\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n output_grads.data.dtype, softmax_type,\n dispatch_scaled_masked_softmax_backward(\n reinterpret_cast(output_grads.data.dptr),\n reinterpret_cast(incoming_grads.data.dptr),\n reinterpret_cast(softmax_results.data.dptr), scale_factor,\n query_seq_len, key_seq_len, batches, attn_heads, stream););\n}\n\nvoid scaled_masked_softmax_forward(const Tensor input, const Tensor mask, Tensor *softmax_results,\n float scale_factor, hipStream_t stream) {\n const int batches = input.data.shape[0];\n const int pad_batches = mask.data.shape[0];\n const int attn_heads = input.data.shape[1];\n const int query_seq_len = input.data.shape[2];\n const int key_seq_len = input.data.shape[3];\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n input.data.dtype, softmax_type,\n dispatch_scaled_masked_softmax_forward(\n reinterpret_cast(softmax_results->data.dptr),\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(mask.data.dptr), scale_factor, query_seq_len,\n key_seq_len, batches, attn_heads, pad_batches, stream););\n}\n\nvoid scaled_masked_softmax_backward(Tensor output_grads, const Tensor incoming_grads,\n const Tensor softmax_results, float scale_factor,\n hipStream_t stream) {\n // output grads is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len]\n const int batches = output_grads.data.shape[0];\n const int attn_heads = output_grads.data.shape[1];\n const int query_seq_len = output_grads.data.shape[2];\n const int key_seq_len = output_grads.data.shape[3];\n\n // Softmax Grad\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n output_grads.data.dtype, softmax_type,\n dispatch_scaled_masked_softmax_backward(\n reinterpret_cast(output_grads.data.dptr),\n reinterpret_cast(incoming_grads.data.dptr),\n reinterpret_cast(softmax_results.data.dptr), scale_factor,\n query_seq_len, key_seq_len, batches, attn_heads, stream););\n}\n\n} // end namespace transformer_engine\n\nvoid nvte_scaled_softmax_forward(const NVTETensor input, NVTETensor softmax_results,\n float scale_factor, hipStream_t stream) {\n NVTE_API_CALL(nvte_scaled_softmax_forward);\n using namespace transformer_engine;\n scaled_softmax_forward(*convertNVTETensorCheck(input), convertNVTETensorCheck(softmax_results),\n scale_factor, stream);\n}\n\nvoid nvte_scaled_softmax_backward(const NVTETensor incoming_grads, const NVTETensor softmax_results,\n NVTETensor output_grads, float scale_factor,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_scaled_softmax_backward);\n using namespace transformer_engine;\n scaled_softmax_backward(*convertNVTETensorCheck(output_grads),\n *convertNVTETensorCheck(incoming_grads),\n *convertNVTETensorCheck(softmax_results), scale_factor, stream);\n}\n\nvoid nvte_scaled_masked_softmax_forward(const NVTETensor input, const NVTETensor mask,\n NVTETensor softmax_results, float scale_factor,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_scaled_masked_softmax_forward);\n using namespace transformer_engine;\n scaled_masked_softmax_forward(*convertNVTETensorCheck(input), *convertNVTETensorCheck(mask),\n convertNVTETensorCheck(softmax_results), scale_factor, stream);\n}\n\nvoid nvte_scaled_masked_softmax_backward(const NVTETensor incoming_grads,\n const NVTETensor softmax_results, NVTETensor output_grads,\n float scale_factor, hipStream_t stream) {\n NVTE_API_CALL(nvte_scaled_masked_softmax_backward);\n using namespace transformer_engine;\n scaled_masked_softmax_backward(*convertNVTETensorCheck(output_grads),\n *convertNVTETensorCheck(incoming_grads),\n *convertNVTETensorCheck(softmax_results), scale_factor, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\n\ntemplate \n__device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);\n\ntemplate <>\n__device__ __inline__ void copy_vector(bf16 *dst, const bf16 *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(bf16 *dst, const bf16 *src) {\n *((uint64_t *)dst) = *((uint64_t *)src); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(fp16 *dst, const fp16 *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(fp16 *dst, const fp16 *src) {\n *((uint64_t *)dst) = *((uint64_t *)src); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(uint8_t *dst, const uint8_t *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(uint8_t *dst, const uint8_t *src) {\n *((uint32_t *)dst) = *((uint32_t *)src); // NOLINT(*)\n}\n\ntemplate \n__device__ __inline__ void copy_zero_vector(Datatype *dst);\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(bf16 *dst) {\n *dst = 0.0f;\n}\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(bf16 *dst) {\n *((float2 *)dst) = make_float2(0.0f, 0.0f); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(fp16 *dst) {\n *dst = 0.0f;\n}\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(fp16 *dst) {\n *((float2 *)dst) = make_float2(0.0f, 0.0f); // NOLINT(*)\n}\n\ntemplate \nstruct Add {\n __device__ __forceinline__ T operator()(T a, T b) const { return a + b; }\n};\n\ntemplate \nstruct Max {\n __device__ __forceinline__ T operator()(T a, T b) const { return a < b ? b : a; }\n};\n\ntemplate \n__device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize,\n unsigned int mask = 0xffffffff) {\n#if CUDA_VERSION >= 9000\n return __shfl_xor_sync(mask, value, laneMask, width);\n#else\n return __shfl_xor(value, laneMask, width);\n#endif\n}\n\ntemplate class ReduceOp>\n__device__ __forceinline__ void warp_reduce(acc_t *sum) {\n ReduceOp r;\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_ROWS; ++i) {\n acc_t b = WARP_SHFL_XOR_NATIVE(sum[i], offset, WARP_SIZE);\n sum[i] = r(sum[i], b);\n }\n }\n}\n\n/*\n * Extended softmax (from native aten pytorch) with the following additional features\n * 1) input scaling\n * 2) implicit causal masking\n *\n * works for all cases:\n * k > q\n * k < q\n * k = q\n *\n * where:\n * microbatches = batches * attn_heads * query_seq_len\n * rows = query_seq_len\n * cols = key_seq_len\n */\ntemplate \n__global__ void scaled_aligned_causal_masked_softmax_warp_forward(output_t *dst, const input_t *src,\n const acc_t scale,\n const int microbatches,\n const int rows, const int cols) {\n // 1) WARP_WIDTH must match the value of warp_size\n // 2) WARP_ROWS must match the value of rows_per_warp\n // of the dispatch_scaled_aligned_causal_masked_softmax_forward method.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_WIDTH =\n (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_WIDTH;\n constexpr int WARP_ROWS = (next_power_of_two <= 128) ? 2 : 1;\n constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;\n\n const int global_row_idx = (blockIdx.x * blockDim.y + threadIdx.y) * WARP_ROWS;\n const int col = threadIdx.x * ELEMENTS_PER_LDG_STG;\n\n const size_t thread_offset = global_row_idx * cols + col;\n\n src += thread_offset;\n dst += thread_offset;\n\n // load data from global memory into registers WITH scaling\n acc_t elements[WARP_ROWS][WARP_ITERATIONS];\n input_t temp_data[ELEMENTS_PER_LDG_STG];\n\n#pragma unroll\n for (int w = 0; w < WARP_ROWS; ++w) {\n const int microbatch = global_row_idx + w;\n const int i = microbatch % rows; // local row index of attention matrix\n const int masked_elements = i + cols - rows + 1;\n\n if (microbatch >= microbatches) {\n break;\n }\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n const int j = col + it * WARP_WIDTH;\n const int itr_idx = w * cols + it * WARP_WIDTH;\n\n if (j < masked_elements) {\n copy_vector(temp_data, src + itr_idx);\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if (j + element < masked_elements) {\n elements[w][it + element] = (acc_t)temp_data[element] * scale;\n } else {\n elements[w][it + element] = (acc_t)(-10'000);\n }\n }\n } else {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n elements[w][it + element] = (acc_t)(-10'000);\n }\n }\n }\n }\n\n // compute max_value\n acc_t max_value[WARP_ROWS];\n#pragma unroll\n for (int w = 0; w < WARP_ROWS; ++w) {\n max_value[w] = elements[w][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n max_value[w] = (max_value[w] > elements[w][it]) ? max_value[w] : elements[w][it];\n }\n }\n warp_reduce(max_value);\n\n acc_t sum[WARP_ROWS]{0.0f};\n#pragma unroll\n for (int w = 0; w < WARP_ROWS; ++w) {\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[w][it] = expf((elements[w][it] - max_value[w]));\n sum[w] += elements[w][it];\n }\n }\n warp_reduce(sum);\n\n output_t out[ELEMENTS_PER_LDG_STG]{0.0f};\n// store result\n#pragma unroll\n for (int w = 0; w < WARP_ROWS; ++w) {\n const int microbatch = global_row_idx + w;\n const int i = microbatch % rows;\n const int masked_elements = i + cols - rows + 1;\n\n // out of Attention matrix bounds (rows)\n if (microbatch >= microbatches) {\n break;\n }\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n const int j = col + it * WARP_WIDTH; // index of the first column\n const int itr_idx = w * cols + it * WARP_WIDTH;\n\n if (j < masked_elements) {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if (j + element < masked_elements) {\n out[element] = elements[w][it + element] / sum[w];\n } else {\n out[element] = (output_t)(0.0f);\n }\n }\n copy_vector(dst + itr_idx, out);\n } else if (j < cols) {\n copy_zero_vector(dst + itr_idx);\n } else {\n break;\n }\n }\n }\n}\n\ntemplate \n__global__ void scaled_aligned_causal_masked_softmax_warp_backward(\n output_t *gradInput, const input_t *grad, const input_t *softmax_output, const acc_t scale,\n const int microbatches, const int rows, const int cols) {\n // 1) WARP_WIDTH must match the value of warp_size\n // 2) WARP_ROWS must match the value of rows_per_warp\n // of the dispatch_scaled_aligned_causal_masked_softmax_forward method.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_WIDTH =\n (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_WIDTH;\n constexpr int WARP_ROWS = (next_power_of_two <= 128) ? 2 : 1;\n constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;\n\n const int global_row_idx = (blockIdx.x * blockDim.y + threadIdx.y) * WARP_ROWS;\n const int col = threadIdx.x * ELEMENTS_PER_LDG_STG;\n\n const size_t thread_offset = global_row_idx * cols + col;\n\n grad += thread_offset;\n softmax_output += thread_offset;\n gradInput += thread_offset;\n\n // load data from global memory into registers\n acc_t grad_reg[WARP_ROWS][WARP_ITERATIONS]{0.0f};\n acc_t softmax_output_reg[WARP_ROWS][WARP_ITERATIONS]{0.0f};\n input_t temp_grad[ELEMENTS_PER_LDG_STG];\n input_t temp_output[ELEMENTS_PER_LDG_STG];\n\n#pragma unroll\n for (int w = 0; w < WARP_ROWS; ++w) {\n const int microbatch = global_row_idx + w;\n const int i = microbatch % rows; // local row index of attention matrix\n const int masked_elements = i + cols - rows + 1;\n\n if (microbatch >= microbatches) {\n break;\n }\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n const int j = col + it * WARP_WIDTH; // index of the first column\n const int itr_idx = w * cols + it * WARP_WIDTH;\n\n if (j < masked_elements) {\n copy_vector(temp_grad, grad + itr_idx);\n copy_vector(temp_output, softmax_output + itr_idx);\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if (j + element < masked_elements) {\n softmax_output_reg[w][it + element] = (acc_t)temp_output[element];\n grad_reg[w][it + element] =\n (acc_t)temp_grad[element] * softmax_output_reg[w][it + element];\n }\n }\n }\n }\n }\n\n acc_t sum[WARP_ROWS];\n#pragma unroll\n for (int w = 0; w < WARP_ROWS; ++w) {\n sum[w] = grad_reg[w][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n sum[w] += grad_reg[w][it];\n }\n }\n\n warp_reduce(sum);\n\n// store result\n#pragma unroll\n for (int w = 0; w < WARP_ROWS; ++w) {\n const int microbatch = global_row_idx + w;\n if (microbatch >= microbatches) {\n break;\n }\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n const int j = col + it * WARP_WIDTH; // index of the first column\n const int itr_idx = w * cols + it * WARP_WIDTH;\n\n if (j < cols) {\n output_t out[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = (output_t)(scale * (grad_reg[w][it + element] -\n softmax_output_reg[w][it + element] * sum[w]));\n }\n copy_vector(gradInput + itr_idx, out);\n }\n }\n }\n}\n\ntemplate \nvoid call_kernel_scaled_aligned_causal_masked_softmax_forward(\n dim3 grid_size, dim3 block_size, const int shmem_size, cudaStream_t stream, output_t *dst,\n const input_t *src, const acc_t scale, const int microbatches, const int query_seq_len,\n const int key_seq_len) {\n scaled_aligned_causal_masked_softmax_warp_forward\n <<>>(dst, src, scale, microbatches, query_seq_len,\n key_seq_len);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\ntemplate \nvoid call_kernel_scaled_aligned_causal_masked_softmax_backward(\n dim3 grid_size, dim3 block_size, const int shmem_size, cudaStream_t stream, output_t *gradInput,\n const input_t *grad, const input_t *output, const acc_t scale, const int microbatches,\n const int query_seq_len, const int key_seq_len) {\n scaled_aligned_causal_masked_softmax_warp_backward\n <<>>(gradInput, grad, output, scale, microbatches,\n query_seq_len, key_seq_len);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\ntemplate \nstruct FunctionWrapper {\n using ForwardType =\n std::function;\n using BackwardType = std::function;\n};\n\nconstexpr int MIN_SUPPORTED_POWER = 4;\nconstexpr int MAX_SUPPORTED_POWER = 14;\nconstexpr int MIN_POWER = MIN_SUPPORTED_POWER - 1;\nconstexpr int MAX_POWER = MAX_SUPPORTED_POWER + 1;\n\n// Recursively instantiate the function for the limit of \"log2_elements\",\n// i.e. \"MAX_POWER\" defined above.\ntemplate \nstruct CompileTimeLoopForward {\n using ForwardFuncType = typename FunctionWrapper::ForwardType;\n static void populate(std::array *arr) {\n CompileTimeLoopForward::populate(arr);\n (*arr)[log2_elements] =\n &call_kernel_scaled_aligned_causal_masked_softmax_forward;\n }\n};\n\ntemplate \nstruct CompileTimeLoopForward {\n using ForwardFuncType = typename FunctionWrapper::ForwardType;\n static void populate(std::array *arr) { (*arr)[MIN_POWER] = nullptr; }\n};\n\ntemplate \nstruct CompileTimeLoopBackward {\n using BackwardFuncType = typename FunctionWrapper::BackwardType;\n static void populate(std::array *arr) {\n CompileTimeLoopBackward::populate(arr);\n (*arr)[log2_elements] =\n &call_kernel_scaled_aligned_causal_masked_softmax_backward;\n }\n};\n\ntemplate \nstruct CompileTimeLoopBackward {\n using BackwardFuncType = typename FunctionWrapper::BackwardType;\n static void populate(std::array *arr) {\n (*arr)[MIN_POWER] = nullptr;\n }\n};\n\ntemplate \nvoid dispatch_scaled_aligned_causal_masked_softmax_forward(output_t *dst, const input_t *src,\n const input_t scale, int query_seq_len,\n int key_seq_len, int batches,\n int attn_heads, cudaStream_t stream) {\n NVTE_CHECK(key_seq_len >= 0 && key_seq_len <= 16384, \"Unsupported shape.\");\n\n if (key_seq_len == 0) {\n return;\n }\n int log2_elements = log2_ceil(key_seq_len);\n const int next_power_of_two = 1 << log2_elements;\n\n // This value must match the WARP_WIDTH constexpr\n // value computed inside scaled_aligned_causal_masked_softmax_warp_forward.\n int warp_width = (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n\n // This value must match the WARP_ROWS constexpr\n // value computed inside scaled_aligned_causal_masked_softmax_warp_forward.\n int microbatches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximimize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = threads_per_block / warp_width;\n int microbatches_per_block = warps_per_block * microbatches_per_warp;\n int microbatches = batches * attn_heads * query_seq_len;\n int blocks = DIVUP(microbatches, microbatches_per_block);\n\n dim3 block_size(warp_width, warps_per_block);\n dim3 grid_size(blocks);\n\n // create an array of pointers to functions\n using ForwardFuncType = typename FunctionWrapper::ForwardType;\n static std::array forwardFunctionArray;\n static bool is_initialized = false;\n if (!is_initialized) {\n CompileTimeLoopForward::populate(\n &forwardFunctionArray);\n is_initialized = true;\n }\n // Call the corresponding kernel\n forwardFunctionArray[log2_elements](grid_size, block_size, 0, stream, dst, src, scale,\n microbatches, query_seq_len, key_seq_len);\n}\n\ntemplate \nvoid dispatch_scaled_aligned_causal_masked_softmax_backward(\n output_t *grad_input, const input_t *grad, const input_t *output, const acc_t scale,\n int query_seq_len, int key_seq_len, int batches, int attn_heads, cudaStream_t stream) {\n NVTE_CHECK(key_seq_len >= 0 && key_seq_len <= 16384, \"Unsupported shape.\");\n\n if (key_seq_len == 0) {\n return;\n }\n int log2_elements = log2_ceil(key_seq_len);\n const int next_power_of_two = 1 << log2_elements;\n\n // This value must match the WARP_WIDTH constexpr\n // value computed inside scaled_aligned_causal_masked_softmax_warp_forward.\n int warp_width = (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n\n // This value must match the WARP_ROWS constexpr\n // value computed inside scaled_aligned_causal_masked_softmax_warp_forward.\n int microbatches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximimize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = threads_per_block / warp_width;\n int microbatches_per_block = warps_per_block * microbatches_per_warp;\n int microbatches = batches * attn_heads * query_seq_len;\n int blocks = DIVUP(microbatches, microbatches_per_block);\n\n dim3 block_size(warp_width, warps_per_block);\n dim3 grid_size(blocks);\n\n // create an array of pointers to functions\n using BackwardFuncType = typename FunctionWrapper::BackwardType;\n static std::array backwardFunctionArray;\n static bool is_initialized = false;\n if (!is_initialized) {\n CompileTimeLoopBackward::populate(\n &backwardFunctionArray);\n is_initialized = true;\n }\n // Call the corresponding kernel\n backwardFunctionArray[log2_elements](grid_size, block_size, 0, stream, grad_input, grad, output,\n scale, microbatches, query_seq_len, key_seq_len);\n}\n\nvoid scaled_aligned_causal_masked_softmax_forward(const Tensor &input, Tensor *softmax_results,\n float scale_factor, cudaStream_t stream) {\n const int batches = input.data.shape[0];\n const int attn_heads = input.data.shape[1];\n const int query_seq_len = input.data.shape[2];\n const int key_seq_len = input.data.shape[3];\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n input.data.dtype, softmax_type,\n dispatch_scaled_aligned_causal_masked_softmax_forward(\n reinterpret_cast(softmax_results->data.dptr),\n reinterpret_cast(input.data.dptr), scale_factor, query_seq_len,\n key_seq_len, batches, attn_heads, stream););\n}\n\nvoid scaled_aligned_causal_masked_softmax_backward(Tensor output_grads, const Tensor incoming_grads,\n const Tensor softmax_results, float scale_factor,\n cudaStream_t stream) {\n // output grads is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len]\n const int batches = output_grads.data.shape[0];\n const int attn_heads = output_grads.data.shape[1];\n const int query_seq_len = output_grads.data.shape[2];\n const int key_seq_len = output_grads.data.shape[3];\n\n // Softmax Grad\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n output_grads.data.dtype, softmax_type,\n dispatch_scaled_aligned_causal_masked_softmax_backward(\n reinterpret_cast(output_grads.data.dptr),\n reinterpret_cast(incoming_grads.data.dptr),\n reinterpret_cast(softmax_results.data.dptr), scale_factor,\n query_seq_len, key_seq_len, batches, attn_heads, stream););\n}\n} // end namespace transformer_engine\n\nvoid nvte_scaled_aligned_causal_masked_softmax_forward(const NVTETensor input,\n NVTETensor softmax_results,\n float scale_factor, cudaStream_t stream) {\n NVTE_API_CALL(nvte_scaled_aligned_causal_masked_softmax_forward);\n using namespace transformer_engine;\n scaled_aligned_causal_masked_softmax_forward(*convertNVTETensorCheck(input),\n convertNVTETensorCheck(softmax_results),\n scale_factor, stream);\n}\n\nvoid nvte_scaled_aligned_causal_masked_softmax_backward(const NVTETensor incoming_grads,\n const NVTETensor softmax_results,\n NVTETensor output_grads, float scale_factor,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_scaled_aligned_causal_masked_softmax_backward);\n using namespace transformer_engine;\n scaled_aligned_causal_masked_softmax_backward(\n *convertNVTETensorCheck(output_grads), *convertNVTETensorCheck(incoming_grads),\n *convertNVTETensorCheck(softmax_results), scale_factor, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\n\ntemplate \n__device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);\n\ntemplate <>\n__device__ __inline__ void copy_vector(bf16 *dst, const bf16 *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(bf16 *dst, const bf16 *src) {\n *((uint64_t *)dst) = *((uint64_t *)src); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(fp16 *dst, const fp16 *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(fp16 *dst, const fp16 *src) {\n *((uint64_t *)dst) = *((uint64_t *)src); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(uint8_t *dst, const uint8_t *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(uint8_t *dst, const uint8_t *src) {\n *((uint32_t *)dst) = *((uint32_t *)src); // NOLINT(*)\n}\n\ntemplate \n__device__ __inline__ void copy_zero_vector(Datatype *dst);\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(bf16 *dst) {\n *dst = 0.0f;\n}\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(bf16 *dst) {\n *((float2 *)dst) = make_float2(0.0f, 0.0f); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(fp16 *dst) {\n *dst = 0.0f;\n}\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(fp16 *dst) {\n *((float2 *)dst) = make_float2(0.0f, 0.0f); // NOLINT(*)\n}\n\ntemplate \nstruct Add {\n __device__ __forceinline__ T operator()(T a, T b) const { return a + b; }\n};\n\ntemplate \nstruct Max {\n __device__ __forceinline__ T operator()(T a, T b) const { return a < b ? b : a; }\n};\n\ntemplate \n__device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize,\n unsigned int mask = 0xffffffff) {\n#if CUDA_VERSION >= 9000\n return __shfl_xor_sync(mask, value, laneMask, width);\n#else\n return __shfl_xor(value, laneMask, width);\n#endif\n}\n\ntemplate class ReduceOp>\n__device__ __forceinline__ void warp_reduce(acc_t *sum) {\n ReduceOp r;\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_ROWS; ++i) {\n acc_t b = WARP_SHFL_XOR_NATIVE(sum[i], offset, WARP_SIZE);\n sum[i] = r(sum[i], b);\n }\n }\n}\n\n/*\n * Extended softmax (from native aten pytorch) with the following additional features\n * 1) input scaling\n * 2) implicit causal masking\n *\n * works for all cases:\n * k > q\n * k < q\n * k = q\n *\n * where:\n * microbatches = batches * attn_heads * query_seq_len\n * rows = query_seq_len\n * cols = key_seq_len\n */\ntemplate \n__global__ void scaled_aligned_causal_masked_softmax_warp_forward(output_t *dst, const input_t *src,\n const acc_t scale,\n const int microbatches,\n const int rows, const int cols) {\n // 1) WARP_WIDTH must match the value of warp_size\n // 2) WARP_ROWS must match the value of rows_per_warp\n // of the dispatch_scaled_aligned_causal_masked_softmax_forward method.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_WIDTH =\n (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_WIDTH;\n constexpr int WARP_ROWS = (next_power_of_two <= 128) ? 2 : 1;\n constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;\n\n const int global_row_idx = (blockIdx.x * blockDim.y + threadIdx.y) * WARP_ROWS;\n const int col = threadIdx.x * ELEMENTS_PER_LDG_STG;\n\n const size_t thread_offset = global_row_idx * cols + col;\n\n src += thread_offset;\n dst += thread_offset;\n\n // load data from global memory into registers WITH scaling\n acc_t elements[WARP_ROWS][WARP_ITERATIONS];\n input_t temp_data[ELEMENTS_PER_LDG_STG];\n\n#pragma unroll\n for (int w = 0; w < WARP_ROWS; ++w) {\n const int microbatch = global_row_idx + w;\n const int i = microbatch % rows; // local row index of attention matrix\n const int masked_elements = i + cols - rows + 1;\n\n if (microbatch >= microbatches) {\n break;\n }\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n const int j = col + it * WARP_WIDTH;\n const int itr_idx = w * cols + it * WARP_WIDTH;\n\n if (j < masked_elements) {\n copy_vector(temp_data, src + itr_idx);\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if (j + element < masked_elements) {\n elements[w][it + element] = (acc_t)temp_data[element] * scale;\n } else {\n elements[w][it + element] = (acc_t)(-10'000);\n }\n }\n } else {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n elements[w][it + element] = (acc_t)(-10'000);\n }\n }\n }\n }\n\n // compute max_value\n acc_t max_value[WARP_ROWS];\n#pragma unroll\n for (int w = 0; w < WARP_ROWS; ++w) {\n max_value[w] = elements[w][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n max_value[w] = (max_value[w] > elements[w][it]) ? max_value[w] : elements[w][it];\n }\n }\n warp_reduce(max_value);\n\n acc_t sum[WARP_ROWS]{0.0f};\n#pragma unroll\n for (int w = 0; w < WARP_ROWS; ++w) {\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[w][it] = expf((elements[w][it] - max_value[w]));\n sum[w] += elements[w][it];\n }\n }\n warp_reduce(sum);\n\n output_t out[ELEMENTS_PER_LDG_STG]{0.0f};\n// store result\n#pragma unroll\n for (int w = 0; w < WARP_ROWS; ++w) {\n const int microbatch = global_row_idx + w;\n const int i = microbatch % rows;\n const int masked_elements = i + cols - rows + 1;\n\n // out of Attention matrix bounds (rows)\n if (microbatch >= microbatches) {\n break;\n }\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n const int j = col + it * WARP_WIDTH; // index of the first column\n const int itr_idx = w * cols + it * WARP_WIDTH;\n\n if (j < masked_elements) {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if (j + element < masked_elements) {\n out[element] = elements[w][it + element] / sum[w];\n } else {\n out[element] = (output_t)(0.0f);\n }\n }\n copy_vector(dst + itr_idx, out);\n } else if (j < cols) {\n copy_zero_vector(dst + itr_idx);\n } else {\n break;\n }\n }\n }\n}\n\ntemplate \n__global__ void scaled_aligned_causal_masked_softmax_warp_backward(\n output_t *gradInput, const input_t *grad, const input_t *softmax_output, const acc_t scale,\n const int microbatches, const int rows, const int cols) {\n // 1) WARP_WIDTH must match the value of warp_size\n // 2) WARP_ROWS must match the value of rows_per_warp\n // of the dispatch_scaled_aligned_causal_masked_softmax_forward method.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_WIDTH =\n (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_WIDTH;\n constexpr int WARP_ROWS = (next_power_of_two <= 128) ? 2 : 1;\n constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;\n\n const int global_row_idx = (blockIdx.x * blockDim.y + threadIdx.y) * WARP_ROWS;\n const int col = threadIdx.x * ELEMENTS_PER_LDG_STG;\n\n const size_t thread_offset = global_row_idx * cols + col;\n\n grad += thread_offset;\n softmax_output += thread_offset;\n gradInput += thread_offset;\n\n // load data from global memory into registers\n acc_t grad_reg[WARP_ROWS][WARP_ITERATIONS]{0.0f};\n acc_t softmax_output_reg[WARP_ROWS][WARP_ITERATIONS]{0.0f};\n input_t temp_grad[ELEMENTS_PER_LDG_STG];\n input_t temp_output[ELEMENTS_PER_LDG_STG];\n\n#pragma unroll\n for (int w = 0; w < WARP_ROWS; ++w) {\n const int microbatch = global_row_idx + w;\n const int i = microbatch % rows; // local row index of attention matrix\n const int masked_elements = i + cols - rows + 1;\n\n if (microbatch >= microbatches) {\n break;\n }\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n const int j = col + it * WARP_WIDTH; // index of the first column\n const int itr_idx = w * cols + it * WARP_WIDTH;\n\n if (j < masked_elements) {\n copy_vector(temp_grad, grad + itr_idx);\n copy_vector(temp_output, softmax_output + itr_idx);\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if (j + element < masked_elements) {\n softmax_output_reg[w][it + element] = (acc_t)temp_output[element];\n grad_reg[w][it + element] =\n (acc_t)temp_grad[element] * softmax_output_reg[w][it + element];\n }\n }\n }\n }\n }\n\n acc_t sum[WARP_ROWS];\n#pragma unroll\n for (int w = 0; w < WARP_ROWS; ++w) {\n sum[w] = grad_reg[w][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n sum[w] += grad_reg[w][it];\n }\n }\n\n warp_reduce(sum);\n\n// store result\n#pragma unroll\n for (int w = 0; w < WARP_ROWS; ++w) {\n const int microbatch = global_row_idx + w;\n if (microbatch >= microbatches) {\n break;\n }\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n const int j = col + it * WARP_WIDTH; // index of the first column\n const int itr_idx = w * cols + it * WARP_WIDTH;\n\n if (j < cols) {\n output_t out[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = (output_t)(scale * (grad_reg[w][it + element] -\n softmax_output_reg[w][it + element] * sum[w]));\n }\n copy_vector(gradInput + itr_idx, out);\n }\n }\n }\n}\n\ntemplate \nvoid call_kernel_scaled_aligned_causal_masked_softmax_forward(\n dim3 grid_size, dim3 block_size, const int shmem_size, hipStream_t stream, output_t *dst,\n const input_t *src, const acc_t scale, const int microbatches, const int query_seq_len,\n const int key_seq_len) {\n scaled_aligned_causal_masked_softmax_warp_forward\n <<>>(dst, src, scale, microbatches, query_seq_len,\n key_seq_len);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\ntemplate \nvoid call_kernel_scaled_aligned_causal_masked_softmax_backward(\n dim3 grid_size, dim3 block_size, const int shmem_size, hipStream_t stream, output_t *gradInput,\n const input_t *grad, const input_t *output, const acc_t scale, const int microbatches,\n const int query_seq_len, const int key_seq_len) {\n scaled_aligned_causal_masked_softmax_warp_backward\n <<>>(gradInput, grad, output, scale, microbatches,\n query_seq_len, key_seq_len);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\ntemplate \nstruct FunctionWrapper {\n using ForwardType =\n std::function;\n using BackwardType = std::function;\n};\n\nconstexpr int MIN_SUPPORTED_POWER = 4;\nconstexpr int MAX_SUPPORTED_POWER = 14;\nconstexpr int MIN_POWER = MIN_SUPPORTED_POWER - 1;\nconstexpr int MAX_POWER = MAX_SUPPORTED_POWER + 1;\n\n// Recursively instantiate the function for the limit of \"log2_elements\",\n// i.e. \"MAX_POWER\" defined above.\ntemplate \nstruct CompileTimeLoopForward {\n using ForwardFuncType = typename FunctionWrapper::ForwardType;\n static void populate(std::array *arr) {\n CompileTimeLoopForward::populate(arr);\n (*arr)[log2_elements] =\n &call_kernel_scaled_aligned_causal_masked_softmax_forward;\n }\n};\n\ntemplate \nstruct CompileTimeLoopForward {\n using ForwardFuncType = typename FunctionWrapper::ForwardType;\n static void populate(std::array *arr) { (*arr)[MIN_POWER] = nullptr; }\n};\n\ntemplate \nstruct CompileTimeLoopBackward {\n using BackwardFuncType = typename FunctionWrapper::BackwardType;\n static void populate(std::array *arr) {\n CompileTimeLoopBackward::populate(arr);\n (*arr)[log2_elements] =\n &call_kernel_scaled_aligned_causal_masked_softmax_backward;\n }\n};\n\ntemplate \nstruct CompileTimeLoopBackward {\n using BackwardFuncType = typename FunctionWrapper::BackwardType;\n static void populate(std::array *arr) {\n (*arr)[MIN_POWER] = nullptr;\n }\n};\n\ntemplate \nvoid dispatch_scaled_aligned_causal_masked_softmax_forward(output_t *dst, const input_t *src,\n const input_t scale, int query_seq_len,\n int key_seq_len, int batches,\n int attn_heads, hipStream_t stream) {\n NVTE_CHECK(key_seq_len >= 0 && key_seq_len <= 16384, \"Unsupported shape.\");\n\n if (key_seq_len == 0) {\n return;\n }\n int log2_elements = log2_ceil(key_seq_len);\n const int next_power_of_two = 1 << log2_elements;\n\n // This value must match the WARP_WIDTH constexpr\n // value computed inside scaled_aligned_causal_masked_softmax_warp_forward.\n int warp_width = (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n\n // This value must match the WARP_ROWS constexpr\n // value computed inside scaled_aligned_causal_masked_softmax_warp_forward.\n int microbatches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximimize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = threads_per_block / warp_width;\n int microbatches_per_block = warps_per_block * microbatches_per_warp;\n int microbatches = batches * attn_heads * query_seq_len;\n int blocks = DIVUP(microbatches, microbatches_per_block);\n\n dim3 block_size(warp_width, warps_per_block);\n dim3 grid_size(blocks);\n\n // create an array of pointers to functions\n using ForwardFuncType = typename FunctionWrapper::ForwardType;\n static std::array forwardFunctionArray;\n static bool is_initialized = false;\n if (!is_initialized) {\n CompileTimeLoopForward::populate(\n &forwardFunctionArray);\n is_initialized = true;\n }\n // Call the corresponding kernel\n forwardFunctionArray[log2_elements](grid_size, block_size, 0, stream, dst, src, scale,\n microbatches, query_seq_len, key_seq_len);\n}\n\ntemplate \nvoid dispatch_scaled_aligned_causal_masked_softmax_backward(\n output_t *grad_input, const input_t *grad, const input_t *output, const acc_t scale,\n int query_seq_len, int key_seq_len, int batches, int attn_heads, hipStream_t stream) {\n NVTE_CHECK(key_seq_len >= 0 && key_seq_len <= 16384, \"Unsupported shape.\");\n\n if (key_seq_len == 0) {\n return;\n }\n int log2_elements = log2_ceil(key_seq_len);\n const int next_power_of_two = 1 << log2_elements;\n\n // This value must match the WARP_WIDTH constexpr\n // value computed inside scaled_aligned_causal_masked_softmax_warp_forward.\n int warp_width = (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n\n // This value must match the WARP_ROWS constexpr\n // value computed inside scaled_aligned_causal_masked_softmax_warp_forward.\n int microbatches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximimize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = threads_per_block / warp_width;\n int microbatches_per_block = warps_per_block * microbatches_per_warp;\n int microbatches = batches * attn_heads * query_seq_len;\n int blocks = DIVUP(microbatches, microbatches_per_block);\n\n dim3 block_size(warp_width, warps_per_block);\n dim3 grid_size(blocks);\n\n // create an array of pointers to functions\n using BackwardFuncType = typename FunctionWrapper::BackwardType;\n static std::array backwardFunctionArray;\n static bool is_initialized = false;\n if (!is_initialized) {\n CompileTimeLoopBackward::populate(\n &backwardFunctionArray);\n is_initialized = true;\n }\n // Call the corresponding kernel\n backwardFunctionArray[log2_elements](grid_size, block_size, 0, stream, grad_input, grad, output,\n scale, microbatches, query_seq_len, key_seq_len);\n}\n\nvoid scaled_aligned_causal_masked_softmax_forward(const Tensor &input, Tensor *softmax_results,\n float scale_factor, hipStream_t stream) {\n const int batches = input.data.shape[0];\n const int attn_heads = input.data.shape[1];\n const int query_seq_len = input.data.shape[2];\n const int key_seq_len = input.data.shape[3];\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n input.data.dtype, softmax_type,\n dispatch_scaled_aligned_causal_masked_softmax_forward(\n reinterpret_cast(softmax_results->data.dptr),\n reinterpret_cast(input.data.dptr), scale_factor, query_seq_len,\n key_seq_len, batches, attn_heads, stream););\n}\n\nvoid scaled_aligned_causal_masked_softmax_backward(Tensor output_grads, const Tensor incoming_grads,\n const Tensor softmax_results, float scale_factor,\n hipStream_t stream) {\n // output grads is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len]\n const int batches = output_grads.data.shape[0];\n const int attn_heads = output_grads.data.shape[1];\n const int query_seq_len = output_grads.data.shape[2];\n const int key_seq_len = output_grads.data.shape[3];\n\n // Softmax Grad\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n output_grads.data.dtype, softmax_type,\n dispatch_scaled_aligned_causal_masked_softmax_backward(\n reinterpret_cast(output_grads.data.dptr),\n reinterpret_cast(incoming_grads.data.dptr),\n reinterpret_cast(softmax_results.data.dptr), scale_factor,\n query_seq_len, key_seq_len, batches, attn_heads, stream););\n}\n} // end namespace transformer_engine\n\nvoid nvte_scaled_aligned_causal_masked_softmax_forward(const NVTETensor input,\n NVTETensor softmax_results,\n float scale_factor, hipStream_t stream) {\n NVTE_API_CALL(nvte_scaled_aligned_causal_masked_softmax_forward);\n using namespace transformer_engine;\n scaled_aligned_causal_masked_softmax_forward(*convertNVTETensorCheck(input),\n convertNVTETensorCheck(softmax_results),\n scale_factor, stream);\n}\n\nvoid nvte_scaled_aligned_causal_masked_softmax_backward(const NVTETensor incoming_grads,\n const NVTETensor softmax_results,\n NVTETensor output_grads, float scale_factor,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_scaled_aligned_causal_masked_softmax_backward);\n using namespace transformer_engine;\n scaled_aligned_causal_masked_softmax_backward(\n *convertNVTETensorCheck(output_grads), *convertNVTETensorCheck(incoming_grads),\n *convertNVTETensorCheck(softmax_results), scale_factor, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\n\ntemplate \n__device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);\n\ntemplate <>\n__device__ __inline__ void copy_vector(bf16 *dst, const bf16 *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(bf16 *dst, const bf16 *src) {\n *((float2 *)dst) = *((float2 *)src); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(fp16 *dst, const fp16 *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(fp16 *dst, const fp16 *src) {\n *((float2 *)dst) = *((float2 *)src); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(uint8_t *dst, const uint8_t *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(uint8_t *dst, const uint8_t *src) {\n *((half2 *)dst) = *((half2 *)src); // NOLINT(*)\n}\n\ntemplate \n__device__ __inline__ void copy_zero_vector(Datatype *dst);\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(bf16 *dst) {\n *dst = 0.0f;\n}\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(bf16 *dst) {\n *((float2 *)dst) = make_float2(0.0f, 0.0f); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(fp16 *dst) {\n *dst = 0.0f;\n}\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(fp16 *dst) {\n *((float2 *)dst) = make_float2(0.0f, 0.0f); // NOLINT(*)\n}\n\ntemplate \nstruct Add {\n __device__ __forceinline__ T operator()(T a, T b) const { return a + b; }\n};\n\ntemplate \nstruct Max {\n __device__ __forceinline__ T operator()(T a, T b) const { return a < b ? b : a; }\n};\n\ntemplate \n__device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize,\n unsigned int mask = 0xffffffff) {\n#if CUDA_VERSION >= 9000\n return __shfl_xor_sync(mask, value, laneMask, width);\n#else\n return __shfl_xor(value, laneMask, width);\n#endif\n}\n\ntemplate class ReduceOp>\n__device__ __forceinline__ void warp_reduce(acc_t *sum) {\n ReduceOp r;\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n acc_t b = WARP_SHFL_XOR_NATIVE(sum[i], offset, WARP_SIZE);\n sum[i] = r(sum[i], b);\n }\n }\n}\n\n/*\n * Extended softmax (from native aten pytorch) with following additional features\n * 1) input scaling\n * 2) Implicit time (diagonal masking)\n */\ntemplate \n__global__ void scaled_upper_triang_masked_softmax_warp_forward(output_t *dst, const input_t *src,\n const acc_t scale,\n int micro_batch_size, int stride,\n int element_count) {\n // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and\n // warp_size of method warp_softmax_forward_kernel.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_SIZE =\n (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;\n constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;\n constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;\n\n size_t first_batch =\n (blockDim.y * blockIdx.y + threadIdx.y) * gridDim.x * WARP_BATCH + blockIdx.x;\n int local_seq = blockIdx.x + 1;\n int warp_iteration_limit = (local_seq + ELEMENTS_PER_LDG_STG * WARP_SIZE - 1) / WARP_SIZE;\n\n // micro_batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = micro_batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the batch\n int local_idx = threadIdx.x;\n\n size_t thread_offset = first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;\n src += thread_offset;\n dst += thread_offset;\n\n // load data from global memory\n acc_t elements[WARP_BATCH][WARP_ITERATIONS];\n input_t temp_data[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : local_seq;\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n\n if (element_index < batch_element_count) {\n copy_vector(\n temp_data, src + i * element_count * stride + it * WARP_SIZE);\n\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if ((element_index + element) < batch_element_count) {\n elements[i][it + element] = (acc_t)temp_data[element] * scale;\n } else {\n elements[i][it + element] = -std::numeric_limits::infinity();\n }\n }\n } else {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n elements[i][it + element] = -std::numeric_limits::infinity();\n }\n }\n }\n }\n\n // compute max_value\n acc_t max_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = elements[i][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];\n }\n }\n warp_reduce(max_value);\n\n acc_t sum[WARP_BATCH]{0.0f};\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n if (it < warp_iteration_limit) {\n elements[i][it] = std::exp((elements[i][it] - max_value[i]));\n sum[i] += elements[i][it];\n }\n }\n }\n warp_reduce(sum);\n\n // store result\n output_t out[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n\n if (element_index < local_seq) {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if (element_index + element < local_seq) {\n out[element] = elements[i][it + element] / sum[i];\n } else {\n out[element] = 0.0f;\n }\n }\n copy_vector(\n dst + i * element_count * stride + it * WARP_SIZE, out);\n } else if (element_index < element_count) {\n copy_zero_vector(dst + i * element_count * stride +\n it * WARP_SIZE);\n } else {\n break;\n }\n }\n }\n}\n\ntemplate \n__global__ void scaled_upper_triang_masked_softmax_warp_backward(output_t *gradInput,\n const input_t *grad,\n const input_t *output, acc_t scale,\n int micro_batch_size, int stride,\n int element_count) {\n // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and\n // warp_size of method warp_softmax_backward_kernel.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_SIZE =\n (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;\n constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;\n constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;\n\n size_t first_batch =\n (blockDim.y * blockIdx.y + threadIdx.y) * gridDim.x * WARP_BATCH + blockIdx.x;\n int local_seq = blockIdx.x + 1;\n\n // micro_batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = micro_batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the batch\n int local_idx = threadIdx.x;\n\n // the first element to process by the current thread\n size_t thread_offset = first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;\n grad += thread_offset;\n output += thread_offset;\n gradInput += thread_offset;\n\n // load data from global memory\n acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS]{0.0f};\n acc_t output_reg[WARP_BATCH][WARP_ITERATIONS]{0.0f};\n input_t temp_grad[ELEMENTS_PER_LDG_STG];\n input_t temp_output[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : local_seq;\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < batch_element_count) {\n copy_vector(\n temp_grad, grad + i * element_count * stride + it * WARP_SIZE);\n copy_vector(\n temp_output, output + i * element_count * stride + it * WARP_SIZE);\n\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if (element_index + element < batch_element_count) {\n output_reg[i][it + element] = (acc_t)temp_output[element];\n }\n }\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if (element_index + element < batch_element_count) {\n grad_reg[i][it + element] = (acc_t)temp_grad[element] * output_reg[i][it + element];\n }\n }\n }\n }\n }\n\n acc_t sum[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] = grad_reg[i][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n sum[i] += grad_reg[i][it];\n }\n }\n warp_reduce(sum);\n\n // store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n // compute gradients\n output_t out[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = (output_t)(scale * (grad_reg[i][it + element] -\n output_reg[i][it + element] * sum[i]));\n }\n copy_vector(\n gradInput + i * element_count * stride + it * WARP_SIZE, out);\n }\n }\n }\n}\n\ntemplate \nvoid dispatch_scaled_upper_triang_masked_softmax_forward(output_t *dst, const input_t *src,\n const input_t scale, int softmax_elements,\n int softmax_elements_stride,\n int attn_batches, cudaStream_t stream) {\n NVTE_CHECK(softmax_elements >= 0 && softmax_elements <= 16384, \"Unsupported shape.\");\n if (softmax_elements == 0) {\n return;\n } else {\n int log2_elements = log2_ceil(softmax_elements);\n const int next_power_of_two = 1 << log2_elements;\n int seq_len = softmax_elements;\n int batch_count = attn_batches * seq_len;\n\n // This value must match the WARP_SIZE constexpr\n // value computed inside softmax_warp_forward.\n int warp_size = (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n\n // This value must match the WARP_BATCH constexpr\n // value computed inside softmax_warp_forward.\n int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximimize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n NVTE_CHECK(attn_batches % batches_per_block == 0, \"Unsupported shape.\");\n\n int blocks_per_seq = attn_batches / batches_per_block;\n dim3 blocks(seq_len, blocks_per_seq, 1);\n dim3 threads(warp_size, warps_per_block, 1);\n // Launch code would be more elegant if C++ supported FOR CONSTEXPR\n switch (log2_elements) {\n case 0: // 1\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 1: // 2\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 2: // 4\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 3: // 8\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 4: // 16\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 5: // 32\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 6: // 64\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 7: // 128\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 8: // 256\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 9: // 512\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 10: // 1024\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 11: // 2048\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 12: // 4096\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 13: // 8192\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 14: // 16384\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n default:\n break;\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\ntemplate \nvoid dispatch_scaled_upper_triang_masked_softmax_backward(output_t *grad_input, const input_t *grad,\n const input_t *output, const acc_t scale,\n int softmax_elements,\n int softmax_elements_stride,\n int attn_batches, cudaStream_t stream) {\n NVTE_CHECK(softmax_elements >= 0 && softmax_elements <= 16384, \"Unsupported shape.\");\n if (softmax_elements == 0) {\n return;\n } else {\n int log2_elements = log2_ceil(softmax_elements);\n const int next_power_of_two = 1 << log2_elements;\n int seq_len = softmax_elements;\n int batch_count = attn_batches * seq_len;\n\n // This value must match the WARP_SIZE constexpr\n // value computed inside softmax_warp_backward.\n int warp_size = (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n\n // This value must match the WARP_BATCH constexpr\n // value computed inside softmax_warp_backward.\n int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximimize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n NVTE_CHECK(attn_batches % batches_per_block == 0, \"Unsupported shape.\");\n\n int blocks_per_seq = attn_batches / batches_per_block;\n dim3 blocks(seq_len, blocks_per_seq, 1);\n dim3 threads(warp_size, warps_per_block, 1);\n // Launch code would be more elegant if C++ supported FOR CONSTEXPR\n switch (log2_elements) {\n case 0: // 1\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 1: // 2\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 2: // 4\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 3: // 8\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 4: // 16\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 5: // 32\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 6: // 64\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 7: // 128\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 8: // 256\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 9: // 512\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 10: // 1024\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 11: // 2048\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 12: // 4096\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 13: // 8192\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 14: // 16384\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n default:\n break;\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\nvoid scaled_upper_triang_masked_softmax_forward(const Tensor input, Tensor *softmax_results,\n float scale_factor, cudaStream_t stream) {\n const int attn_batches = input.data.shape[0];\n const int seq_len = input.data.shape[1];\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n input.data.dtype, softmax_type,\n dispatch_scaled_upper_triang_masked_softmax_forward(\n reinterpret_cast(softmax_results->data.dptr),\n reinterpret_cast(input.data.dptr), scale_factor, seq_len, seq_len,\n attn_batches, stream););\n}\n\nvoid scaled_upper_triang_masked_softmax_backward(Tensor output_grads, const Tensor incoming_grads,\n const Tensor softmax_results, float scale_factor,\n cudaStream_t stream) {\n const int attn_batches = output_grads.data.shape[0];\n const int seq_len = output_grads.data.shape[1];\n\n // Softmax Grad\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n output_grads.data.dtype, softmax_type,\n dispatch_scaled_upper_triang_masked_softmax_backward(\n reinterpret_cast(output_grads.data.dptr),\n reinterpret_cast(incoming_grads.data.dptr),\n reinterpret_cast(softmax_results.data.dptr), scale_factor, seq_len,\n seq_len, attn_batches, stream););\n}\n\n} // end namespace transformer_engine\n\nvoid nvte_scaled_upper_triang_masked_softmax_forward(const NVTETensor input,\n NVTETensor softmax_results, float scale_factor,\n cudaStream_t stream) {\n using namespace transformer_engine;\n scaled_upper_triang_masked_softmax_forward(*convertNVTETensorCheck(input),\n convertNVTETensorCheck(softmax_results), scale_factor,\n stream);\n}\n\nvoid nvte_scaled_upper_triang_masked_softmax_backward(const NVTETensor incoming_grads,\n const NVTETensor softmax_results,\n NVTETensor output_grads, float scale_factor,\n cudaStream_t stream) {\n using namespace transformer_engine;\n scaled_upper_triang_masked_softmax_backward(\n *convertNVTETensorCheck(output_grads), *convertNVTETensorCheck(incoming_grads),\n *convertNVTETensorCheck(softmax_results), scale_factor, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\n\ntemplate \n__device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);\n\ntemplate <>\n__device__ __inline__ void copy_vector(bf16 *dst, const bf16 *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(bf16 *dst, const bf16 *src) {\n *((float2 *)dst) = *((float2 *)src); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(fp16 *dst, const fp16 *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(fp16 *dst, const fp16 *src) {\n *((float2 *)dst) = *((float2 *)src); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(uint8_t *dst, const uint8_t *src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(uint8_t *dst, const uint8_t *src) {\n *((half2 *)dst) = *((half2 *)src); // NOLINT(*)\n}\n\ntemplate \n__device__ __inline__ void copy_zero_vector(Datatype *dst);\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(bf16 *dst) {\n *dst = 0.0f;\n}\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(bf16 *dst) {\n *((float2 *)dst) = make_float2(0.0f, 0.0f); // NOLINT(*)\n}\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(fp16 *dst) {\n *dst = 0.0f;\n}\n\ntemplate <>\n__device__ __inline__ void copy_zero_vector(fp16 *dst) {\n *((float2 *)dst) = make_float2(0.0f, 0.0f); // NOLINT(*)\n}\n\ntemplate \nstruct Add {\n __device__ __forceinline__ T operator()(T a, T b) const { return a + b; }\n};\n\ntemplate \nstruct Max {\n __device__ __forceinline__ T operator()(T a, T b) const { return a < b ? b : a; }\n};\n\ntemplate \n__device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize,\n unsigned int mask = 0xffffffff) {\n#if CUDA_VERSION >= 9000\n return __shfl_xor_sync(mask, value, laneMask, width);\n#else\n return __shfl_xor(value, laneMask, width);\n#endif\n}\n\ntemplate class ReduceOp>\n__device__ __forceinline__ void warp_reduce(acc_t *sum) {\n ReduceOp r;\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n acc_t b = WARP_SHFL_XOR_NATIVE(sum[i], offset, WARP_SIZE);\n sum[i] = r(sum[i], b);\n }\n }\n}\n\n/*\n * Extended softmax (from native aten pytorch) with following additional features\n * 1) input scaling\n * 2) Implicit time (diagonal masking)\n */\ntemplate \n__global__ void scaled_upper_triang_masked_softmax_warp_forward(output_t *dst, const input_t *src,\n const acc_t scale,\n int micro_batch_size, int stride,\n int element_count) {\n // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and\n // warp_size of method warp_softmax_forward_kernel.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_SIZE =\n (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;\n constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;\n constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;\n\n size_t first_batch =\n (blockDim.y * blockIdx.y + threadIdx.y) * gridDim.x * WARP_BATCH + blockIdx.x;\n int local_seq = blockIdx.x + 1;\n int warp_iteration_limit = (local_seq + ELEMENTS_PER_LDG_STG * WARP_SIZE - 1) / WARP_SIZE;\n\n // micro_batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = micro_batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the batch\n int local_idx = threadIdx.x;\n\n size_t thread_offset = first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;\n src += thread_offset;\n dst += thread_offset;\n\n // load data from global memory\n acc_t elements[WARP_BATCH][WARP_ITERATIONS];\n input_t temp_data[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : local_seq;\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n\n if (element_index < batch_element_count) {\n copy_vector(\n temp_data, src + i * element_count * stride + it * WARP_SIZE);\n\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if ((element_index + element) < batch_element_count) {\n elements[i][it + element] = (acc_t)temp_data[element] * scale;\n } else {\n elements[i][it + element] = -std::numeric_limits::infinity();\n }\n }\n } else {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n elements[i][it + element] = -std::numeric_limits::infinity();\n }\n }\n }\n }\n\n // compute max_value\n acc_t max_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = elements[i][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];\n }\n }\n warp_reduce(max_value);\n\n acc_t sum[WARP_BATCH]{0.0f};\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n if (it < warp_iteration_limit) {\n elements[i][it] = std::exp((elements[i][it] - max_value[i]));\n sum[i] += elements[i][it];\n }\n }\n }\n warp_reduce(sum);\n\n // store result\n output_t out[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n\n if (element_index < local_seq) {\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if (element_index + element < local_seq) {\n out[element] = elements[i][it + element] / sum[i];\n } else {\n out[element] = 0.0f;\n }\n }\n copy_vector(\n dst + i * element_count * stride + it * WARP_SIZE, out);\n } else if (element_index < element_count) {\n copy_zero_vector(dst + i * element_count * stride +\n it * WARP_SIZE);\n } else {\n break;\n }\n }\n }\n}\n\ntemplate \n__global__ void scaled_upper_triang_masked_softmax_warp_backward(output_t *gradInput,\n const input_t *grad,\n const input_t *output, acc_t scale,\n int micro_batch_size, int stride,\n int element_count) {\n // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and\n // warp_size of method warp_softmax_backward_kernel.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_SIZE =\n (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;\n constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;\n constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;\n\n size_t first_batch =\n (blockDim.y * blockIdx.y + threadIdx.y) * gridDim.x * WARP_BATCH + blockIdx.x;\n int local_seq = blockIdx.x + 1;\n\n // micro_batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = micro_batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the batch\n int local_idx = threadIdx.x;\n\n // the first element to process by the current thread\n size_t thread_offset = first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;\n grad += thread_offset;\n output += thread_offset;\n gradInput += thread_offset;\n\n // load data from global memory\n acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS]{0.0f};\n acc_t output_reg[WARP_BATCH][WARP_ITERATIONS]{0.0f};\n input_t temp_grad[ELEMENTS_PER_LDG_STG];\n input_t temp_output[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : local_seq;\n\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < batch_element_count) {\n copy_vector(\n temp_grad, grad + i * element_count * stride + it * WARP_SIZE);\n copy_vector(\n temp_output, output + i * element_count * stride + it * WARP_SIZE);\n\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if (element_index + element < batch_element_count) {\n output_reg[i][it + element] = (acc_t)temp_output[element];\n }\n }\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n if (element_index + element < batch_element_count) {\n grad_reg[i][it + element] = (acc_t)temp_grad[element] * output_reg[i][it + element];\n }\n }\n }\n }\n }\n\n acc_t sum[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] = grad_reg[i][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n sum[i] += grad_reg[i][it];\n }\n }\n warp_reduce(sum);\n\n // store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n // compute gradients\n output_t out[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = (output_t)(scale * (grad_reg[i][it + element] -\n output_reg[i][it + element] * sum[i]));\n }\n copy_vector(\n gradInput + i * element_count * stride + it * WARP_SIZE, out);\n }\n }\n }\n}\n\ntemplate \nvoid dispatch_scaled_upper_triang_masked_softmax_forward(output_t *dst, const input_t *src,\n const input_t scale, int softmax_elements,\n int softmax_elements_stride,\n int attn_batches, hipStream_t stream) {\n NVTE_CHECK(softmax_elements >= 0 && softmax_elements <= 16384, \"Unsupported shape.\");\n if (softmax_elements == 0) {\n return;\n } else {\n int log2_elements = log2_ceil(softmax_elements);\n const int next_power_of_two = 1 << log2_elements;\n int seq_len = softmax_elements;\n int batch_count = attn_batches * seq_len;\n\n // This value must match the WARP_SIZE constexpr\n // value computed inside softmax_warp_forward.\n int warp_size = (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n\n // This value must match the WARP_BATCH constexpr\n // value computed inside softmax_warp_forward.\n int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximimize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n NVTE_CHECK(attn_batches % batches_per_block == 0, \"Unsupported shape.\");\n\n int blocks_per_seq = attn_batches / batches_per_block;\n dim3 blocks(seq_len, blocks_per_seq, 1);\n dim3 threads(warp_size, warps_per_block, 1);\n // Launch code would be more elegant if C++ supported FOR CONSTEXPR\n switch (log2_elements) {\n case 0: // 1\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 1: // 2\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 2: // 4\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 3: // 8\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 4: // 16\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 5: // 32\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 6: // 64\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 7: // 128\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 8: // 256\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 9: // 512\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 10: // 1024\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 11: // 2048\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 12: // 4096\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 13: // 8192\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n case 14: // 16384\n scaled_upper_triang_masked_softmax_warp_forward\n <<>>(dst, src, scale, batch_count, softmax_elements_stride,\n softmax_elements);\n break;\n default:\n break;\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\ntemplate \nvoid dispatch_scaled_upper_triang_masked_softmax_backward(output_t *grad_input, const input_t *grad,\n const input_t *output, const acc_t scale,\n int softmax_elements,\n int softmax_elements_stride,\n int attn_batches, hipStream_t stream) {\n NVTE_CHECK(softmax_elements >= 0 && softmax_elements <= 16384, \"Unsupported shape.\");\n if (softmax_elements == 0) {\n return;\n } else {\n int log2_elements = log2_ceil(softmax_elements);\n const int next_power_of_two = 1 << log2_elements;\n int seq_len = softmax_elements;\n int batch_count = attn_batches * seq_len;\n\n // This value must match the WARP_SIZE constexpr\n // value computed inside softmax_warp_backward.\n int warp_size = (next_power_of_two < THREADS_PER_WARP) ? next_power_of_two : THREADS_PER_WARP;\n\n // This value must match the WARP_BATCH constexpr\n // value computed inside softmax_warp_backward.\n int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximimize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n NVTE_CHECK(attn_batches % batches_per_block == 0, \"Unsupported shape.\");\n\n int blocks_per_seq = attn_batches / batches_per_block;\n dim3 blocks(seq_len, blocks_per_seq, 1);\n dim3 threads(warp_size, warps_per_block, 1);\n // Launch code would be more elegant if C++ supported FOR CONSTEXPR\n switch (log2_elements) {\n case 0: // 1\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 1: // 2\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 2: // 4\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 3: // 8\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 4: // 16\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 5: // 32\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 6: // 64\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 7: // 128\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 8: // 256\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 9: // 512\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 10: // 1024\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 11: // 2048\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 12: // 4096\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 13: // 8192\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 14: // 16384\n scaled_upper_triang_masked_softmax_warp_backward\n <<>>(grad_input, grad, output, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n default:\n break;\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\nvoid scaled_upper_triang_masked_softmax_forward(const Tensor input, Tensor *softmax_results,\n float scale_factor, hipStream_t stream) {\n const int attn_batches = input.data.shape[0];\n const int seq_len = input.data.shape[1];\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n input.data.dtype, softmax_type,\n dispatch_scaled_upper_triang_masked_softmax_forward(\n reinterpret_cast(softmax_results->data.dptr),\n reinterpret_cast(input.data.dptr), scale_factor, seq_len, seq_len,\n attn_batches, stream););\n}\n\nvoid scaled_upper_triang_masked_softmax_backward(Tensor output_grads, const Tensor incoming_grads,\n const Tensor softmax_results, float scale_factor,\n hipStream_t stream) {\n const int attn_batches = output_grads.data.shape[0];\n const int seq_len = output_grads.data.shape[1];\n\n // Softmax Grad\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n output_grads.data.dtype, softmax_type,\n dispatch_scaled_upper_triang_masked_softmax_backward(\n reinterpret_cast(output_grads.data.dptr),\n reinterpret_cast(incoming_grads.data.dptr),\n reinterpret_cast(softmax_results.data.dptr), scale_factor, seq_len,\n seq_len, attn_batches, stream););\n}\n\n} // end namespace transformer_engine\n\nvoid nvte_scaled_upper_triang_masked_softmax_forward(const NVTETensor input,\n NVTETensor softmax_results, float scale_factor,\n hipStream_t stream) {\n using namespace transformer_engine;\n scaled_upper_triang_masked_softmax_forward(*convertNVTETensorCheck(input),\n convertNVTETensorCheck(softmax_results), scale_factor,\n stream);\n}\n\nvoid nvte_scaled_upper_triang_masked_softmax_backward(const NVTETensor incoming_grads,\n const NVTETensor softmax_results,\n NVTETensor output_grads, float scale_factor,\n hipStream_t stream) {\n using namespace transformer_engine;\n scaled_upper_triang_masked_softmax_backward(\n *convertNVTETensorCheck(output_grads), *convertNVTETensorCheck(incoming_grads),\n *convertNVTETensorCheck(softmax_results), scale_factor, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"transformer_engine/transformer_engine.h\"\n\nnamespace transformer_engine {\nnamespace {\nconstexpr uint32_t WARP_SIZE = 32;\n} // namespace\nnamespace swizzle_kernel_1d {\nconstexpr uint32_t WARPS_X_PER_TB = 2; // configurable\nconstexpr uint32_t WARPS_Y_PER_TB = 2; // configurable\n\n// Transposes a 4x4 matrix of bytes stored across four threads with consecutive thread ids where\n// each thread stores a single row (of four bytes).\n// Example:\n// lane0.row = 0x00010203\n// lane1.row = 0x04050607\n// lane2.row = 0x08090a0b\n// lane3.row = 0x0c0d0e0f\n// Becomes:\n// lane0.row = 0x0004080c\n// lane1.row = 0x0105090d\n// lane2.row = 0x02060a0e\n// lane3.row = 0x03070b0f\nuint32_t __device__ __forceinline__ transpose_4x4_byte_matrix(const uint32_t row,\n const uint32_t lane,\n const uint32_t active_mask) {\n using cu = const uint32_t;\n\n // Threads operate in groups of 4, and each thread stores 4 bytes at a time.\n // The bytes in this 4x4 matrix are labeled in hex. We shuffle around bytes\n // until we have transposed the 4x4 matrix.\n cu m_0123_4567_89ab_cdef = row;\n cu m_4567_0123_cdef_89ab = __shfl_xor_sync(active_mask, m_0123_4567_89ab_cdef, 1, 4);\n cu m_0426_4062_8cae_c8ea = __byte_perm(m_0123_4567_89ab_cdef, m_4567_0123_cdef_89ab, 0x6240);\n cu m_5173_1537_d9fb_9dbf = __byte_perm(m_0123_4567_89ab_cdef, m_4567_0123_cdef_89ab, 0x3715);\n cu m_0426_1537_8cae_9dbf = (lane & 1) ? m_5173_1537_d9fb_9dbf : m_0426_4062_8cae_c8ea;\n cu m_8cae_9dbf_0426_1537 = __shfl_xor_sync(active_mask, m_0426_1537_8cae_9dbf, 2, 4);\n cu m_048c_159d_8c04_9d15 = __byte_perm(m_0426_1537_8cae_9dbf, m_8cae_9dbf_0426_1537, 0x5410);\n cu m_ae26_bf37_26ae_37bf = __byte_perm(m_0426_1537_8cae_9dbf, m_8cae_9dbf_0426_1537, 0x3276);\n cu m_048c_159d_26ae_37bf = (lane & 2) ? m_ae26_bf37_26ae_37bf : m_048c_159d_8c04_9d15;\n\n return m_048c_159d_26ae_37bf;\n}\n\n// Expands a uint32_t to a uint4 by duplicating each byte four times.\n// Example: 0x01020304u becomes uint4{0x01010101, 0x02020202, 0x03030303, 0x04040404}\nuint4 __device__ __forceinline__ broadcast_uint32_t_to_uint4(uint32_t x) {\n return {__byte_perm(x, 0, 0x0000), __byte_perm(x, 0, 0x1111), __byte_perm(x, 0, 0x2222),\n __byte_perm(x, 0, 0x3333)};\n}\n\n// Tag struct denoting whether the number of rows of the input fp8 block scaling tensor's data\n// matrix is divisible by 128. If it is not, some threads could read out of bounds scaling factors.\nstruct no_oob_tag_t {};\nconstexpr no_oob_tag_t NO_OOB_TAG;\n\ntemplate \nvoid __global__ __launch_bounds__(WARPS_X_PER_TB* WARPS_Y_PER_TB* WARP_SIZE)\n swizzle_block_scaling_1d_to_mxfp8_scaling_factors_kernel(\n const void* __restrict__ const in, void* __restrict__ const out, const uint32_t tiles_x,\n const uint32_t tiles_y, const uint32_t in_y_stride, const uint32_t out_y_stride,\n OOBT first_oob) {\n // resolve kernel variant\n constexpr bool no_oob = std::is_same_v;\n static_assert(no_oob || std::is_same_v);\n\n // load thread indices\n const uint32_t lane = threadIdx.x;\n __builtin_assume(lane < WARP_SIZE);\n const uint32_t warp_x = threadIdx.z;\n __builtin_assume(warp_x < WARPS_X_PER_TB);\n const uint32_t warp_y = threadIdx.y;\n __builtin_assume(warp_y < WARPS_Y_PER_TB);\n\n // compute tile indices\n const uint32_t out_tile_y = blockIdx.y * WARPS_Y_PER_TB + warp_y;\n const uint32_t out_tile_x = blockIdx.x * WARPS_X_PER_TB + warp_x;\n const uint32_t in_tile_y = out_tile_x;\n const uint32_t in_tile_x = out_tile_y;\n\n // bounds check; uniform branch\n if (out_tile_y >= tiles_y || out_tile_x >= tiles_x) {\n return;\n }\n\n // calculate this warp's input base pointer\n constexpr uint32_t in_x_stride = WARP_SIZE * sizeof(uint4);\n const void* const warp_src =\n (reinterpret_cast(in) + in_tile_y * in_y_stride + in_tile_x * in_x_stride);\n\n // load scaling factors for this lane's initial four 1x128 tiles\n uint4 sf;\n if constexpr (no_oob) {\n sf = reinterpret_cast(warp_src)[lane];\n } else {\n if ((out_tile_y < tiles_y - 1) || lane < first_oob) {\n sf = reinterpret_cast(warp_src)[lane];\n } else {\n sf = uint4{0, 0, 0, 0};\n }\n }\n\n // pack the exponent bits of the scaling factors\n uint32_t packed_exponents = ((sf.x >> 23) & 0xFF) | (((sf.y >> 23) & 0xFF) << 8) |\n (((sf.z >> 23) & 0xFF) << 16) | (((sf.w >> 23) & 0xFF) << 24);\n\n // partially swizzle the scaling factors\n constexpr uint32_t ACTIVE_MASK = 0xFFFFFFFF; // no divergent branches\n const uint32_t lane_load_idx = (lane % 4) * 8 + (lane / 4);\n packed_exponents = __shfl_sync(ACTIVE_MASK, packed_exponents, lane_load_idx);\n\n // transpose 4x4 matrices of scaling factors\n packed_exponents = transpose_4x4_byte_matrix(packed_exponents, lane % 4, ACTIVE_MASK);\n\n // broadcast the scaling factors for sixteen 1x32 tiles\n sf = broadcast_uint32_t_to_uint4(packed_exponents);\n\n // store them cooperatively for 512 1x32 tiles in a 128x128 tile\n constexpr uint32_t out_x_stride = 512;\n void* const warp_dst =\n (reinterpret_cast(out) + out_tile_y * out_y_stride + out_tile_x * out_x_stride);\n reinterpret_cast(warp_dst)[lane] = sf;\n}\n\nvoid launch_kernel(const void* const in, void* const out, uint32_t data_rows, uint32_t data_cols,\n cudaStream_t stream) {\n NVTE_CHECK(is_aligned_ptr(in, alignof(uint4)), \"Input scaling factor pointer must be aligned to \",\n alignof(uint4), \" bytes\");\n NVTE_CHECK(is_aligned_ptr(out, alignof(uint4)),\n \"Output scaling factor pointer must be aligned to \", alignof(uint4), \" bytes\");\n NVTE_CHECK(data_rows % 4 == 0, \"Input tensor must not have any padding scaling factors\");\n\n const uint32_t tiles_x = DIVUP(data_cols, 128u);\n const uint32_t tiles_y = DIVUP(data_rows, 128u);\n const dim3 grid_dim{DIVUP(tiles_x, WARPS_X_PER_TB), DIVUP(tiles_y, WARPS_Y_PER_TB), 1};\n const dim3 block_dim{WARP_SIZE, WARPS_Y_PER_TB, WARPS_X_PER_TB};\n\n // Each 128x128 tile in the data corresponds to a 128x1 tile in the input scales\n // and a 128x4 tile in the output scales. The input scales are in transposed order.\n const uint32_t input_scale_inv_cols = DIVUP(data_rows, 4u) * 4;\n const uint32_t output_scale_inv_cols = tiles_x * 128 * 4;\n const uint32_t in_y_stride = input_scale_inv_cols * sizeof(float);\n const uint32_t out_y_stride = output_scale_inv_cols * sizeof(uint8_t);\n\n const uint32_t first_oob = (input_scale_inv_cols % 128) / 4;\n\n if (first_oob == 0) {\n swizzle_block_scaling_1d_to_mxfp8_scaling_factors_kernel<<>>(\n in, out, tiles_x, tiles_y, in_y_stride, out_y_stride, NO_OOB_TAG);\n } else {\n swizzle_block_scaling_1d_to_mxfp8_scaling_factors_kernel<<>>(\n in, out, tiles_x, tiles_y, in_y_stride, out_y_stride, first_oob);\n }\n}\n} // namespace swizzle_kernel_1d\nnamespace swizzle_kernel_2d {\nconstexpr uint32_t WARPS_X_PER_TB = 2; // configurable\nconstexpr uint32_t WARPS_Y_PER_TB = 2; // configurable\n\nvoid __global__ __launch_bounds__(WARPS_X_PER_TB* WARPS_Y_PER_TB* WARP_SIZE)\n swizzle_block_scaling_2d_to_mxfp8_scaling_factors_kernel(\n const void* __restrict__ const in, void* __restrict__ const out, const uint32_t tiles_x,\n const uint32_t tiles_y, const uint32_t in_y_stride, const uint32_t out_y_stride) {\n // load thread indices\n const uint32_t lane = threadIdx.x;\n __builtin_assume(lane < WARP_SIZE);\n const uint32_t warp_x = threadIdx.z;\n __builtin_assume(warp_x < WARPS_X_PER_TB);\n const uint32_t warp_y = threadIdx.y;\n __builtin_assume(warp_y < WARPS_Y_PER_TB);\n\n // compute tile indices\n const uint32_t out_tile_y = blockIdx.y * WARPS_Y_PER_TB + warp_y;\n const uint32_t out_tile_x = blockIdx.x * WARPS_X_PER_TB + warp_x;\n const uint32_t in_tile_y = out_tile_y;\n const uint32_t in_tile_x = out_tile_x;\n\n // bounds check; uniform branch\n if (out_tile_y >= tiles_y || out_tile_x >= tiles_x) {\n return;\n }\n\n // calculate this warp's input base pointer\n constexpr uint32_t in_x_stride = sizeof(float);\n const void* const warp_src =\n (reinterpret_cast(in) + in_tile_y * in_y_stride + in_tile_x * in_x_stride);\n\n // load scaling factor for this warp's 128x128 tile\n uint32_t sf = *reinterpret_cast(warp_src);\n\n // broadcast it to four scaling factors for 1x32 tiles\n // extract and broadcast the exponent byte to four bytes for E8M0 format\n uint32_t exp_byte = (sf >> 23) & 0xFF;\n sf = exp_byte | (exp_byte << 8) | (exp_byte << 16) | (exp_byte << 24);\n\n // broadcast it to sixteen scaling factors for 1x32 tiles\n const uint4 sf4{sf, sf, sf, sf};\n\n // store it cooperatively for 512 1x32 tiles in a 128x128 tile\n constexpr uint32_t out_x_stride = 512;\n void* const warp_dst =\n (reinterpret_cast(out) + out_tile_y * out_y_stride + out_tile_x * out_x_stride);\n reinterpret_cast(warp_dst)[lane] = sf4;\n}\n\nvoid launch_kernel(const void* const in, void* const out, uint32_t data_rows, uint32_t data_cols,\n cudaStream_t stream) {\n NVTE_CHECK(is_aligned_ptr(in, alignof(float)), \"Input scaling factor pointer must be aligned to \",\n alignof(float), \" bytes\");\n NVTE_CHECK(is_aligned_ptr(out, alignof(uint4)),\n \"Output scaling factor pointer must be aligned to \", alignof(uint4), \" bytes\");\n\n const uint32_t tiles_x = DIVUP(data_cols, 128u);\n const uint32_t tiles_y = DIVUP(data_rows, 128u);\n const dim3 grid_dim{DIVUP(tiles_x, WARPS_X_PER_TB), DIVUP(tiles_y, WARPS_Y_PER_TB), 1};\n const dim3 block_dim{WARP_SIZE, WARPS_Y_PER_TB, WARPS_X_PER_TB};\n\n // Each 128x128 tile in the data corresponds to a 1x1 tile in the input scales\n // and a 128x4 tile in the output scales.\n const uint32_t input_scale_inv_cols = DIVUP(data_cols, 512u) * 4;\n const uint32_t output_scale_inv_cols = tiles_x * 128 * 4;\n const uint32_t in_y_stride = input_scale_inv_cols * sizeof(float);\n const uint32_t out_y_stride = output_scale_inv_cols * sizeof(uint8_t);\n\n swizzle_block_scaling_2d_to_mxfp8_scaling_factors_kernel<<>>(\n in, out, tiles_x, tiles_y, in_y_stride, out_y_stride);\n}\n} // namespace swizzle_kernel_2d\n\nvoid swizzle_block_scaling_to_mxfp8_scaling_factors(const Tensor* input, Tensor* output,\n cudaStream_t stream) {\n // Do nothing if tensor is empty\n if (input->data.numel() == 0) {\n return;\n }\n\n CheckInputTensor(*input, \"block_scaling_scaling_factor_input\");\n CheckInputTensor(*output, \"mxfp8_scaling_factor_output\");\n\n const NVTEScalingMode scaling_mode = input->scaling_mode;\n NVTE_CHECK(scaling_mode == NVTE_BLOCK_SCALING_1D || scaling_mode == NVTE_BLOCK_SCALING_2D,\n \"Input tensor must be a block scaling tensor\");\n NVTE_CHECK(output->scaling_mode == NVTE_MXFP8_1D_SCALING,\n \"Output tensor must be an mxfp8 tensor\");\n\n NVTE_CHECK(input->data.dtype == transformer_engine::DType::kFloat8E4M3 ||\n input->data.dtype == transformer_engine::DType::kFloat8E5M2,\n \"Input data must have FP8E4M3 or FP8E5M2 dtype to be compatible with MXFP8\");\n NVTE_CHECK(output->data.dtype == input->data.dtype,\n \"Output data must have the same dtype as input data\");\n NVTE_CHECK(input->scale_inv.dtype == DType::kFloat32, \"Input must have FP32 scaling factors\");\n NVTE_CHECK(output->scale_inv.dtype == DType::kFloat8E8M0,\n \"Output must have E8M0 scaling factors\");\n\n NVTE_CHECK(output->with_gemm_swizzled_scales,\n \"Expected output tensor with scales in GEMM swizzled format.\");\n\n NVTE_CHECK(input->data.dptr != nullptr, \"Input must have rowwise data\");\n NVTE_CHECK(output->data.dptr == input->data.dptr, \"Output must share data with input\");\n NVTE_CHECK(input->scale_inv.dptr != nullptr, \"Input must have rowwise scaling factors\");\n NVTE_CHECK(output->scale_inv.dptr != nullptr, \"Output must have rowwise scaling factors\");\n\n NVTE_CHECK(input->data.shape.size() == 2, \"Input data must be a matrix\");\n NVTE_CHECK(output->data.shape == input->data.shape,\n \"Output data must have the same shape as input data\");\n NVTE_CHECK(input->scale_inv.shape.size() == 2, \"Input scaling factors must be a matrix\");\n NVTE_CHECK(output->scale_inv.shape.size() == 2, \"Output scaling factors must be a matrix\");\n\n const size_t data_rows = input->data.shape[0];\n const size_t data_cols = input->data.shape[1];\n const size_t input_scale_inv_rows = input->scale_inv.shape[0];\n const size_t input_scale_inv_cols = input->scale_inv.shape[1];\n const size_t output_scale_inv_rows = output->scale_inv.shape[0];\n const size_t output_scale_inv_cols = output->scale_inv.shape[1];\n\n NVTE_CHECK(output_scale_inv_rows == DIVUP(data_rows, 128) * 128,\n \"Expected the output scaling factor matrix to have \",\n DIVUP(data_rows, 128) * 128, \" rows, but it has \", output_scale_inv_rows,\n \" rows instead.\");\n NVTE_CHECK(output_scale_inv_cols == DIVUP(data_cols, 128) * 4,\n \"Expected the output scaling factor matrix to have \",\n DIVUP(data_cols, 128) * 4, \" columns, but it has \", output_scale_inv_cols,\n \" columns instead.\");\n\n if (scaling_mode == NVTE_BLOCK_SCALING_1D) {\n NVTE_CHECK(input_scale_inv_rows == DIVUP(data_cols, 128),\n \"Expected the input scaling factor matrix to have \", DIVUP(data_cols, 128),\n \" rows, but it has \", input_scale_inv_rows, \" rows instead.\");\n NVTE_CHECK(input_scale_inv_cols == DIVUP(data_rows, 4) * 4,\n \"Expected the input scaling factor matrix to have \", DIVUP(data_rows, 4) * 4,\n \" columns, but it has \", input_scale_inv_cols, \" columns instead.\");\n\n swizzle_kernel_1d::launch_kernel(input->scale_inv.dptr, output->scale_inv.dptr, data_rows,\n data_cols, stream);\n } else { // scaling_mode == NVTE_BLOCK_SCALING_2D\n NVTE_CHECK(input_scale_inv_rows == DIVUP(data_rows, 128),\n \"Expected the input scaling factor matrix to have \", DIVUP(data_rows, 128),\n \" rows, but it has \", input_scale_inv_rows, \" rows instead.\");\n NVTE_CHECK(input_scale_inv_cols == DIVUP(data_cols, 512) * 4,\n \"Expected the input scaling factor matrix to have \",\n DIVUP(data_cols, 512) * 4, \" columns, but it has \", input_scale_inv_cols,\n \" columns instead.\");\n\n swizzle_kernel_2d::launch_kernel(input->scale_inv.dptr, output->scale_inv.dptr, data_rows,\n data_cols, stream);\n }\n}\n\n} // namespace transformer_engine\n\nvoid nvte_swizzle_block_scaling_to_mxfp8_scaling_factors(const NVTETensor input, NVTETensor output,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_swizzle_block_scaling_to_mxfp8_scaling_factors);\n using namespace transformer_engine;\n swizzle_block_scaling_to_mxfp8_scaling_factors(convertNVTETensorCheck(input),\n convertNVTETensorCheck(output), stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"transformer_engine/transformer_engine.h\"\n\nnamespace transformer_engine {\nnamespace {\nconstexpr uint32_t WARP_SIZE = 32;\n} // namespace\nnamespace swizzle_kernel_1d {\nconstexpr uint32_t WARPS_X_PER_TB = 2; // configurable\nconstexpr uint32_t WARPS_Y_PER_TB = 2; // configurable\n\n// Transposes a 4x4 matrix of bytes stored across four threads with consecutive thread ids where\n// each thread stores a single row (of four bytes).\n// Example:\n// lane0.row = 0x00010203\n// lane1.row = 0x04050607\n// lane2.row = 0x08090a0b\n// lane3.row = 0x0c0d0e0f\n// Becomes:\n// lane0.row = 0x0004080c\n// lane1.row = 0x0105090d\n// lane2.row = 0x02060a0e\n// lane3.row = 0x03070b0f\nuint32_t __device__ __forceinline__ transpose_4x4_byte_matrix(const uint32_t row,\n const uint32_t lane,\n const uint32_t active_mask) {\n using cu = const uint32_t;\n\n // Threads operate in groups of 4, and each thread stores 4 bytes at a time.\n // The bytes in this 4x4 matrix are labeled in hex. We shuffle around bytes\n // until we have transposed the 4x4 matrix.\n cu m_0123_4567_89ab_cdef = row;\n cu m_4567_0123_cdef_89ab = __shfl_xor_sync(active_mask, m_0123_4567_89ab_cdef, 1, 4);\n cu m_0426_4062_8cae_c8ea = __byte_perm(m_0123_4567_89ab_cdef, m_4567_0123_cdef_89ab, 0x6240);\n cu m_5173_1537_d9fb_9dbf = __byte_perm(m_0123_4567_89ab_cdef, m_4567_0123_cdef_89ab, 0x3715);\n cu m_0426_1537_8cae_9dbf = (lane & 1) ? m_5173_1537_d9fb_9dbf : m_0426_4062_8cae_c8ea;\n cu m_8cae_9dbf_0426_1537 = __shfl_xor_sync(active_mask, m_0426_1537_8cae_9dbf, 2, 4);\n cu m_048c_159d_8c04_9d15 = __byte_perm(m_0426_1537_8cae_9dbf, m_8cae_9dbf_0426_1537, 0x5410);\n cu m_ae26_bf37_26ae_37bf = __byte_perm(m_0426_1537_8cae_9dbf, m_8cae_9dbf_0426_1537, 0x3276);\n cu m_048c_159d_26ae_37bf = (lane & 2) ? m_ae26_bf37_26ae_37bf : m_048c_159d_8c04_9d15;\n\n return m_048c_159d_26ae_37bf;\n}\n\n// Expands a uint32_t to a uint4 by duplicating each byte four times.\n// Example: 0x01020304u becomes uint4{0x01010101, 0x02020202, 0x03030303, 0x04040404}\nuint4 __device__ __forceinline__ broadcast_uint32_t_to_uint4(uint32_t x) {\n return {__byte_perm(x, 0, 0x0000), __byte_perm(x, 0, 0x1111), __byte_perm(x, 0, 0x2222),\n __byte_perm(x, 0, 0x3333)};\n}\n\n// Tag struct denoting whether the number of rows of the input fp8 block scaling tensor's data\n// matrix is divisible by 128. If it is not, some threads could read out of bounds scaling factors.\nstruct no_oob_tag_t {};\nconstexpr no_oob_tag_t NO_OOB_TAG;\n\ntemplate \nvoid __global__ __launch_bounds__(WARPS_X_PER_TB* WARPS_Y_PER_TB* WARP_SIZE)\n swizzle_block_scaling_1d_to_mxfp8_scaling_factors_kernel(\n const void* __restrict__ const in, void* __restrict__ const out, const uint32_t tiles_x,\n const uint32_t tiles_y, const uint32_t in_y_stride, const uint32_t out_y_stride,\n OOBT first_oob) {\n // resolve kernel variant\n constexpr bool no_oob = std::is_same_v;\n static_assert(no_oob || std::is_same_v);\n\n // load thread indices\n const uint32_t lane = threadIdx.x;\n __builtin_assume(lane < WARP_SIZE);\n const uint32_t warp_x = threadIdx.z;\n __builtin_assume(warp_x < WARPS_X_PER_TB);\n const uint32_t warp_y = threadIdx.y;\n __builtin_assume(warp_y < WARPS_Y_PER_TB);\n\n // compute tile indices\n const uint32_t out_tile_y = blockIdx.y * WARPS_Y_PER_TB + warp_y;\n const uint32_t out_tile_x = blockIdx.x * WARPS_X_PER_TB + warp_x;\n const uint32_t in_tile_y = out_tile_x;\n const uint32_t in_tile_x = out_tile_y;\n\n // bounds check; uniform branch\n if (out_tile_y >= tiles_y || out_tile_x >= tiles_x) {\n return;\n }\n\n // calculate this warp's input base pointer\n constexpr uint32_t in_x_stride = WARP_SIZE * sizeof(uint4);\n const void* const warp_src =\n (reinterpret_cast(in) + in_tile_y * in_y_stride + in_tile_x * in_x_stride);\n\n // load scaling factors for this lane's initial four 1x128 tiles\n uint4 sf;\n if constexpr (no_oob) {\n sf = reinterpret_cast(warp_src)[lane];\n } else {\n if ((out_tile_y < tiles_y - 1) || lane < first_oob) {\n sf = reinterpret_cast(warp_src)[lane];\n } else {\n sf = uint4{0, 0, 0, 0};\n }\n }\n\n // pack the exponent bits of the scaling factors\n uint32_t packed_exponents = ((sf.x >> 23) & 0xFF) | (((sf.y >> 23) & 0xFF) << 8) |\n (((sf.z >> 23) & 0xFF) << 16) | (((sf.w >> 23) & 0xFF) << 24);\n\n // partially swizzle the scaling factors\n constexpr uint32_t ACTIVE_MASK = 0xFFFFFFFF; // no divergent branches\n const uint32_t lane_load_idx = (lane % 4) * 8 + (lane / 4);\n packed_exponents = __shfl_sync(ACTIVE_MASK, packed_exponents, lane_load_idx);\n\n // transpose 4x4 matrices of scaling factors\n packed_exponents = transpose_4x4_byte_matrix(packed_exponents, lane % 4, ACTIVE_MASK);\n\n // broadcast the scaling factors for sixteen 1x32 tiles\n sf = broadcast_uint32_t_to_uint4(packed_exponents);\n\n // store them cooperatively for 512 1x32 tiles in a 128x128 tile\n constexpr uint32_t out_x_stride = 512;\n void* const warp_dst =\n (reinterpret_cast(out) + out_tile_y * out_y_stride + out_tile_x * out_x_stride);\n reinterpret_cast(warp_dst)[lane] = sf;\n}\n\nvoid launch_kernel(const void* const in, void* const out, uint32_t data_rows, uint32_t data_cols,\n hipStream_t stream) {\n NVTE_CHECK(is_aligned_ptr(in, alignof(uint4)), \"Input scaling factor pointer must be aligned to \",\n alignof(uint4), \" bytes\");\n NVTE_CHECK(is_aligned_ptr(out, alignof(uint4)),\n \"Output scaling factor pointer must be aligned to \", alignof(uint4), \" bytes\");\n NVTE_CHECK(data_rows % 4 == 0, \"Input tensor must not have any padding scaling factors\");\n\n const uint32_t tiles_x = DIVUP(data_cols, 128u);\n const uint32_t tiles_y = DIVUP(data_rows, 128u);\n const dim3 grid_dim{DIVUP(tiles_x, WARPS_X_PER_TB), DIVUP(tiles_y, WARPS_Y_PER_TB), 1};\n const dim3 block_dim{WARP_SIZE, WARPS_Y_PER_TB, WARPS_X_PER_TB};\n\n // Each 128x128 tile in the data corresponds to a 128x1 tile in the input scales\n // and a 128x4 tile in the output scales. The input scales are in transposed order.\n const uint32_t input_scale_inv_cols = DIVUP(data_rows, 4u) * 4;\n const uint32_t output_scale_inv_cols = tiles_x * 128 * 4;\n const uint32_t in_y_stride = input_scale_inv_cols * sizeof(float);\n const uint32_t out_y_stride = output_scale_inv_cols * sizeof(uint8_t);\n\n const uint32_t first_oob = (input_scale_inv_cols % 128) / 4;\n\n if (first_oob == 0) {\n swizzle_block_scaling_1d_to_mxfp8_scaling_factors_kernel<<>>(\n in, out, tiles_x, tiles_y, in_y_stride, out_y_stride, NO_OOB_TAG);\n } else {\n swizzle_block_scaling_1d_to_mxfp8_scaling_factors_kernel<<>>(\n in, out, tiles_x, tiles_y, in_y_stride, out_y_stride, first_oob);\n }\n}\n} // namespace swizzle_kernel_1d\nnamespace swizzle_kernel_2d {\nconstexpr uint32_t WARPS_X_PER_TB = 2; // configurable\nconstexpr uint32_t WARPS_Y_PER_TB = 2; // configurable\n\nvoid __global__ __launch_bounds__(WARPS_X_PER_TB* WARPS_Y_PER_TB* WARP_SIZE)\n swizzle_block_scaling_2d_to_mxfp8_scaling_factors_kernel(\n const void* __restrict__ const in, void* __restrict__ const out, const uint32_t tiles_x,\n const uint32_t tiles_y, const uint32_t in_y_stride, const uint32_t out_y_stride) {\n // load thread indices\n const uint32_t lane = threadIdx.x;\n __builtin_assume(lane < WARP_SIZE);\n const uint32_t warp_x = threadIdx.z;\n __builtin_assume(warp_x < WARPS_X_PER_TB);\n const uint32_t warp_y = threadIdx.y;\n __builtin_assume(warp_y < WARPS_Y_PER_TB);\n\n // compute tile indices\n const uint32_t out_tile_y = blockIdx.y * WARPS_Y_PER_TB + warp_y;\n const uint32_t out_tile_x = blockIdx.x * WARPS_X_PER_TB + warp_x;\n const uint32_t in_tile_y = out_tile_y;\n const uint32_t in_tile_x = out_tile_x;\n\n // bounds check; uniform branch\n if (out_tile_y >= tiles_y || out_tile_x >= tiles_x) {\n return;\n }\n\n // calculate this warp's input base pointer\n constexpr uint32_t in_x_stride = sizeof(float);\n const void* const warp_src =\n (reinterpret_cast(in) + in_tile_y * in_y_stride + in_tile_x * in_x_stride);\n\n // load scaling factor for this warp's 128x128 tile\n uint32_t sf = *reinterpret_cast(warp_src);\n\n // broadcast it to four scaling factors for 1x32 tiles\n // extract and broadcast the exponent byte to four bytes for E8M0 format\n uint32_t exp_byte = (sf >> 23) & 0xFF;\n sf = exp_byte | (exp_byte << 8) | (exp_byte << 16) | (exp_byte << 24);\n\n // broadcast it to sixteen scaling factors for 1x32 tiles\n const uint4 sf4{sf, sf, sf, sf};\n\n // store it cooperatively for 512 1x32 tiles in a 128x128 tile\n constexpr uint32_t out_x_stride = 512;\n void* const warp_dst =\n (reinterpret_cast(out) + out_tile_y * out_y_stride + out_tile_x * out_x_stride);\n reinterpret_cast(warp_dst)[lane] = sf4;\n}\n\nvoid launch_kernel(const void* const in, void* const out, uint32_t data_rows, uint32_t data_cols,\n hipStream_t stream) {\n NVTE_CHECK(is_aligned_ptr(in, alignof(float)), \"Input scaling factor pointer must be aligned to \",\n alignof(float), \" bytes\");\n NVTE_CHECK(is_aligned_ptr(out, alignof(uint4)),\n \"Output scaling factor pointer must be aligned to \", alignof(uint4), \" bytes\");\n\n const uint32_t tiles_x = DIVUP(data_cols, 128u);\n const uint32_t tiles_y = DIVUP(data_rows, 128u);\n const dim3 grid_dim{DIVUP(tiles_x, WARPS_X_PER_TB), DIVUP(tiles_y, WARPS_Y_PER_TB), 1};\n const dim3 block_dim{WARP_SIZE, WARPS_Y_PER_TB, WARPS_X_PER_TB};\n\n // Each 128x128 tile in the data corresponds to a 1x1 tile in the input scales\n // and a 128x4 tile in the output scales.\n const uint32_t input_scale_inv_cols = DIVUP(data_cols, 512u) * 4;\n const uint32_t output_scale_inv_cols = tiles_x * 128 * 4;\n const uint32_t in_y_stride = input_scale_inv_cols * sizeof(float);\n const uint32_t out_y_stride = output_scale_inv_cols * sizeof(uint8_t);\n\n swizzle_block_scaling_2d_to_mxfp8_scaling_factors_kernel<<>>(\n in, out, tiles_x, tiles_y, in_y_stride, out_y_stride);\n}\n} // namespace swizzle_kernel_2d\n\nvoid swizzle_block_scaling_to_mxfp8_scaling_factors(const Tensor* input, Tensor* output,\n hipStream_t stream) {\n // Do nothing if tensor is empty\n if (input->data.numel() == 0) {\n return;\n }\n\n CheckInputTensor(*input, \"block_scaling_scaling_factor_input\");\n CheckInputTensor(*output, \"mxfp8_scaling_factor_output\");\n\n const NVTEScalingMode scaling_mode = input->scaling_mode;\n NVTE_CHECK(scaling_mode == NVTE_BLOCK_SCALING_1D || scaling_mode == NVTE_BLOCK_SCALING_2D,\n \"Input tensor must be a block scaling tensor\");\n NVTE_CHECK(output->scaling_mode == NVTE_MXFP8_1D_SCALING,\n \"Output tensor must be an mxfp8 tensor\");\n\n NVTE_CHECK(input->data.dtype == transformer_engine::DType::kFloat8E4M3 ||\n input->data.dtype == transformer_engine::DType::kFloat8E5M2,\n \"Input data must have FP8E4M3 or FP8E5M2 dtype to be compatible with MXFP8\");\n NVTE_CHECK(output->data.dtype == input->data.dtype,\n \"Output data must have the same dtype as input data\");\n NVTE_CHECK(input->scale_inv.dtype == DType::kFloat32, \"Input must have FP32 scaling factors\");\n NVTE_CHECK(output->scale_inv.dtype == DType::kFloat8E8M0,\n \"Output must have E8M0 scaling factors\");\n\n NVTE_CHECK(output->with_gemm_swizzled_scales,\n \"Expected output tensor with scales in GEMM swizzled format.\");\n\n NVTE_CHECK(input->data.dptr != nullptr, \"Input must have rowwise data\");\n NVTE_CHECK(output->data.dptr == input->data.dptr, \"Output must share data with input\");\n NVTE_CHECK(input->scale_inv.dptr != nullptr, \"Input must have rowwise scaling factors\");\n NVTE_CHECK(output->scale_inv.dptr != nullptr, \"Output must have rowwise scaling factors\");\n\n NVTE_CHECK(input->data.shape.size() == 2, \"Input data must be a matrix\");\n NVTE_CHECK(output->data.shape == input->data.shape,\n \"Output data must have the same shape as input data\");\n NVTE_CHECK(input->scale_inv.shape.size() == 2, \"Input scaling factors must be a matrix\");\n NVTE_CHECK(output->scale_inv.shape.size() == 2, \"Output scaling factors must be a matrix\");\n\n const size_t data_rows = input->data.shape[0];\n const size_t data_cols = input->data.shape[1];\n const size_t input_scale_inv_rows = input->scale_inv.shape[0];\n const size_t input_scale_inv_cols = input->scale_inv.shape[1];\n const size_t output_scale_inv_rows = output->scale_inv.shape[0];\n const size_t output_scale_inv_cols = output->scale_inv.shape[1];\n\n NVTE_CHECK(output_scale_inv_rows == DIVUP(data_rows, 128) * 128,\n \"Expected the output scaling factor matrix to have \",\n DIVUP(data_rows, 128) * 128, \" rows, but it has \", output_scale_inv_rows,\n \" rows instead.\");\n NVTE_CHECK(output_scale_inv_cols == DIVUP(data_cols, 128) * 4,\n \"Expected the output scaling factor matrix to have \",\n DIVUP(data_cols, 128) * 4, \" columns, but it has \", output_scale_inv_cols,\n \" columns instead.\");\n\n if (scaling_mode == NVTE_BLOCK_SCALING_1D) {\n NVTE_CHECK(input_scale_inv_rows == DIVUP(data_cols, 128),\n \"Expected the input scaling factor matrix to have \", DIVUP(data_cols, 128),\n \" rows, but it has \", input_scale_inv_rows, \" rows instead.\");\n NVTE_CHECK(input_scale_inv_cols == DIVUP(data_rows, 4) * 4,\n \"Expected the input scaling factor matrix to have \", DIVUP(data_rows, 4) * 4,\n \" columns, but it has \", input_scale_inv_cols, \" columns instead.\");\n\n swizzle_kernel_1d::launch_kernel(input->scale_inv.dptr, output->scale_inv.dptr, data_rows,\n data_cols, stream);\n } else { // scaling_mode == NVTE_BLOCK_SCALING_2D\n NVTE_CHECK(input_scale_inv_rows == DIVUP(data_rows, 128),\n \"Expected the input scaling factor matrix to have \", DIVUP(data_rows, 128),\n \" rows, but it has \", input_scale_inv_rows, \" rows instead.\");\n NVTE_CHECK(input_scale_inv_cols == DIVUP(data_cols, 512) * 4,\n \"Expected the input scaling factor matrix to have \",\n DIVUP(data_cols, 512) * 4, \" columns, but it has \", input_scale_inv_cols,\n \" columns instead.\");\n\n swizzle_kernel_2d::launch_kernel(input->scale_inv.dptr, output->scale_inv.dptr, data_rows,\n data_cols, stream);\n }\n}\n\n} // namespace transformer_engine\n\nvoid nvte_swizzle_block_scaling_to_mxfp8_scaling_factors(const NVTETensor input, NVTETensor output,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_swizzle_block_scaling_to_mxfp8_scaling_factors);\n using namespace transformer_engine;\n swizzle_block_scaling_to_mxfp8_scaling_factors(convertNVTETensorCheck(input),\n convertNVTETensorCheck(output), stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/swizzle/swizzle_block_scaling.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/cuda_runtime.h\"\n#include \"../util/logging.h\"\n#include \"transformer_engine/transformer_engine.h\"\n\nnamespace transformer_engine {\nnamespace {\n\nconstexpr int MXFP8_BLOCK_SIZE = 32;\nconstexpr int NVFP4_BLOCK_SIZE = 16;\n\nint get_max_dynamic_smem() {\n static int max_smem = -1;\n if (max_smem < 0) {\n int device;\n NVTE_CHECK_CUDA(cudaGetDevice(&device));\n NVTE_CHECK_CUDA(\n cudaDeviceGetAttribute(&max_smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device));\n }\n return max_smem;\n}\n\nconstexpr __device__ __host__ int TB_DIM = 32;\nconstexpr __device__ __host__ int NEW_SF_TILE_DIM_K = 16;\nconstexpr __device__ __host__ int N_SF_PER_TD_PER_TILE = 4;\n\n// output is in ~K-major interleaved blocks\nconstexpr __device__ __host__ int NEW_SF_TILE_DIM_K_I32 = NEW_SF_TILE_DIM_K / 4;\nconstexpr __device__ __host__ int NEW_SF_TILE_DIM_M_I32 = 32;\n\ntemplate \n__device__ inline void regs_shuffle_with_bit_shifts(LType* regs_vec) {\n // inp, 4-byte chunks [0,1,2,3, 4,5,6,7, 8,9,10,11, 12,13,14,15]\n // out, swapping byte to form new 4-byte chunks [0,4,8,12, 1,5,9,13, 2,6,10,14, 3,7,11,15]\n\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int kVectorSize = N_SF_PER_TD_PER_TILE * N_TILE_PER_TD;\n int32_t new_regs[kVectorSize];\n int32_t* regs = reinterpret_cast(regs_vec);\n\n#pragma unroll\n for (int i = 0; i < N_TILE_PER_TD; i++) {\n#pragma unroll\n for (int j = 0; j < N_SF_PER_TD_PER_TILE; j++) {\n new_regs[i * N_SF_PER_TD_PER_TILE + j] =\n (((regs[i + 0 * N_TILE_PER_TD] >> 8 * j) & 0xFF)) |\n (((regs[i + 1 * N_TILE_PER_TD] >> 8 * j) & 0xFF) << 8) |\n (((regs[i + 2 * N_TILE_PER_TD] >> 8 * j) & 0xFF) << 16) |\n (((regs[i + 3 * N_TILE_PER_TD] >> 8 * j) & 0xFF) << 24);\n }\n }\n#pragma unroll\n for (int i = 0; i < kVectorSize; i++) regs[i] = new_regs[i];\n}\n\ntemplate \n__device__ inline void regs_unshuffle_with_bit_shifts(LType* regs_vec) {\n // Inverse of regs_shuffle_with_bit_shifts\n // inp, 4-byte chunks [0,4,8,12, 1,5,9,13, 2,6,10,14, 3,7,11,15]\n // out, swapping byte to form new 4-byte chunks [0,1,2,3, 4,5,6,7, 8,9,10,11, 12,13,14,15]\n\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int kVectorSize = N_SF_PER_TD_PER_TILE * N_TILE_PER_TD;\n int32_t new_regs[kVectorSize];\n int32_t* regs = reinterpret_cast(regs_vec);\n\n#pragma unroll\n for (int i = 0; i < N_TILE_PER_TD; i++) {\n#pragma unroll\n for (int j = 0; j < N_SF_PER_TD_PER_TILE; j++) {\n new_regs[i + j * N_TILE_PER_TD] =\n ((regs[i * N_SF_PER_TD_PER_TILE + 0] >> 8 * j) & 0xFF) |\n (((regs[i * N_SF_PER_TD_PER_TILE + 1] >> 8 * j) & 0xFF) << 8) |\n (((regs[i * N_SF_PER_TD_PER_TILE + 2] >> 8 * j) & 0xFF) << 16) |\n (((regs[i * N_SF_PER_TD_PER_TILE + 3] >> 8 * j) & 0xFF) << 24);\n }\n }\n#pragma unroll\n for (int i = 0; i < kVectorSize; i++) regs[i] = new_regs[i];\n}\n\n// IS_PADDED_K / IS_PADDED_M select the boundary-block specialization at compile\n// time so the inner load loop avoids the per-iteration runtime checks. The\n// caller computes the runtime predicates from blockIdx/gridDim once per block\n// (uniform across the block) and dispatches to the right specialization.\ntemplate \n__device__ void swizzle_col_scaling_kernel_impl(const void* input, void* output, const int M,\n const int K, const int original_M,\n const int original_K, const int bid_x,\n const int bid_y, const int grid_dim_x,\n const int grid_dim_y) {\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int N_SF_PER_TD = N_TILE_PER_TD * N_SF_PER_TD_PER_TILE;\n constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4;\n\n // input is in M-major\n constexpr int SF_TILE_DIM_M_I32 = SF_TILE_DIM_M / 4;\n constexpr int SF_TILE_DIM_K_I32 = SF_TILE_DIM_K;\n\n const int M_i32 = M / 4;\n const int K_i32 = K;\n\n int m_tiles_in_tb = N_TILE_PER_TD;\n int k_tiles_in_tb = TB_DIM;\n if (bid_x == grid_dim_x - 1) {\n k_tiles_in_tb = (K_i32 / SF_TILE_DIM_K_I32 - 1) % k_tiles_in_tb + 1;\n }\n if (bid_y == grid_dim_y - 1) {\n m_tiles_in_tb = (M_i32 / SF_TILE_DIM_M_I32 - 1) % m_tiles_in_tb + 1;\n }\n\n const int input_offset =\n bid_x * TB_DIM * SF_TILE_DIM_K_I32 * M_i32 + bid_y * N_TILE_PER_TD * SF_TILE_DIM_M_I32;\n const int32_t* input_i32 = reinterpret_cast(input) + input_offset;\n int32_t* output_i32[N_TILE_PER_TD];\n#pragma unroll\n for (int i = 0; i < m_tiles_in_tb; i++) {\n output_i32[i] = reinterpret_cast(output) + bid_x * TB_DIM * SF_TILE_SIZE_I32 +\n (bid_y * N_TILE_PER_TD + i) * SF_TILE_DIM_M_I32 * K_i32;\n }\n extern __shared__ int slm[];\n\n // load, global -> regs\n // Each register read for a given i is along the M direction at K-coord\n // (bid_x * TB_DIM * SF_TILE_DIM_K + threadIdx.y * SF_TILE_DIM_K + i). When that\n // K-coord is past original_K, the entire register is out of the per-tensor data\n // region (which may be the unpadded compact extent), so we must NOT issue the\n // __ldg there -- it could read past the per-tensor buffer (and, for the last\n // tensor in a grouped allocation, past the end of the allocation entirely).\n LType regs_vec[N_SF_PER_TD_PER_TILE];\n if (threadIdx.x * N_TILE_PER_TD < m_tiles_in_tb * SF_TILE_DIM_M_I32 &&\n threadIdx.y < k_tiles_in_tb) {\n const int k_base = bid_x * TB_DIM * SF_TILE_DIM_K + threadIdx.y * SF_TILE_DIM_K;\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) {\n const int thread_offset =\n (threadIdx.y * SF_TILE_DIM_K_I32 + i) * M_i32 + threadIdx.x * N_TILE_PER_TD;\n const int k_coord = k_base + i;\n if constexpr (IS_PADDED_K) {\n if (k_coord >= original_K) {\n // Entire register is past original_K: zero directly without loading.\n uint8_t* zero_bytes = reinterpret_cast(regs_vec + i);\n#pragma unroll\n for (int j = 0; j < static_cast(sizeof(LType)); j++) zero_bytes[j] = 0;\n continue;\n }\n }\n regs_vec[i] = __ldg(reinterpret_cast(input_i32 + thread_offset));\n // Per-byte M masking is still needed when only part of the register is past\n // original_M (i.e. K-coord is in range but the M position spans the boundary).\n if constexpr (IS_PADDED_M) {\n for (int j = 0; j < N_TILE_PER_TD * sizeof(int); j++) {\n const int index = (input_offset + thread_offset) * sizeof(int) + j;\n if (index % M >= original_M) {\n reinterpret_cast(regs_vec + i)[j] = 0;\n }\n }\n }\n }\n\n // local shuffle\n regs_shuffle_with_bit_shifts(regs_vec);\n\n // store, regs -> shared\n int tM = threadIdx.x * N_SF_PER_TD;\n int* slm_tile = slm + (threadIdx.y * SF_TILE_SIZE_I32 +\n tM / SF_TILE_DIM_M * k_tiles_in_tb * SF_TILE_SIZE_I32);\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD; i++) {\n /* TODO rotate_i */\n slm_tile[(tM % SF_TILE_DIM_M) / NEW_SF_TILE_DIM_M_I32 +\n ((tM + i) % NEW_SF_TILE_DIM_M_I32) * NEW_SF_TILE_DIM_K_I32] =\n reinterpret_cast(regs_vec)[i];\n }\n }\n __syncthreads();\n\n // store, shared -> global\n int linear_id = threadIdx.y * blockDim.x + threadIdx.x;\n#pragma unroll\n for (int i = 0; i < m_tiles_in_tb; i++) {\n __align__(16) int4* output_v4i = reinterpret_cast(output_i32[i]);\n __align__(16) int4* slm_v4i =\n reinterpret_cast(slm + i * k_tiles_in_tb * SF_TILE_SIZE_I32);\n#pragma unroll\n for (int j = linear_id; j < SF_TILE_SIZE_I32 * k_tiles_in_tb / 4;\n j += blockDim.x * blockDim.y) {\n output_v4i[j] = slm_v4i[j];\n }\n }\n}\n\n// Dispatch helper: pick the right (IS_PADDED_K, IS_PADDED_M) col-scaling impl\n// specialization at runtime based on the per-block padding predicates. The\n// branching here is uniform across all threads in the block, so the indirect\n// path each block takes still inlines cleanly.\ntemplate \n__device__ __forceinline__ void dispatch_swizzle_col_scaling_kernel_impl(\n const void* input, void* output, const int M, const int K, const int original_M,\n const int original_K, const int bid_x, const int bid_y, const int grid_dim_x,\n const int grid_dim_y, const bool padding_k, const bool padding_m) {\n if (padding_k && padding_m) {\n swizzle_col_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n } else if (padding_k) {\n swizzle_col_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n } else if (padding_m) {\n swizzle_col_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n } else {\n swizzle_col_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n swizzle_col_scaling_kernel(const void* input, void* output, const int M, const int K,\n const int original_M, const int original_K) {\n const bool padding_m = (blockIdx.y == gridDim.y - 1) && (original_M < M);\n const bool padding_k = (blockIdx.x == gridDim.x - 1) && (original_K < K);\n dispatch_swizzle_col_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y,\n padding_k, padding_m);\n}\n\ntemplate \n__device__ inline void regs_shuffle(LType* regs_vec) {\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n if constexpr (N_TILE_PER_TD == 1) return;\n\n constexpr int kVectorSize = N_SF_PER_TD_PER_TILE * N_TILE_PER_TD;\n int32_t tmp[kVectorSize];\n int32_t* ptr = reinterpret_cast(regs_vec);\n#pragma unroll\n for (int i = 0; i < kVectorSize; i++)\n tmp[i % N_TILE_PER_TD * N_SF_PER_TD_PER_TILE + i / N_TILE_PER_TD] = ptr[i];\n\n#pragma unroll\n for (int i = 0; i < kVectorSize; i++) ptr[i] = tmp[i];\n}\n\n// Inverse of regs_shuffle.\ntemplate \n__device__ inline void regs_unshuffle(LType* regs_vec) {\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n if constexpr (N_TILE_PER_TD == 1) return;\n\n constexpr int kVectorSize = N_SF_PER_TD_PER_TILE * N_TILE_PER_TD;\n int32_t tmp[kVectorSize];\n int32_t* ptr = reinterpret_cast(regs_vec);\n#pragma unroll\n for (int i = 0; i < kVectorSize; i++)\n tmp[i % N_SF_PER_TD_PER_TILE * N_TILE_PER_TD + i / N_SF_PER_TD_PER_TILE] = ptr[i];\n\n#pragma unroll\n for (int i = 0; i < kVectorSize; i++) ptr[i] = tmp[i];\n}\n\n// IS_PADDED_K / IS_PADDED_M select the boundary-block specialization at compile\n// time so the inner load loop avoids the per-iteration runtime checks. The\n// caller computes the runtime predicates from blockIdx/gridDim once per block\n// (uniform across the block) and dispatches to the right specialization.\ntemplate \n__device__ void swizzle_row_scaling_kernel_impl(const void* input, void* output, const int M,\n const int K, const int original_M,\n const int original_K, const int bid_x,\n const int bid_y, const int grid_dim_x,\n const int grid_dim_y) {\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int N_TILES_IN_TB = TB_DIM * N_TILE_PER_TD;\n\n // input is in K-major\n constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4;\n constexpr int SF_TILE_DIM_M_I32 = SF_TILE_DIM_M;\n\n int n_tiles_in_tb = N_TILES_IN_TB;\n const int K_i32 = K / 4;\n if (bid_x == grid_dim_x - 1) {\n n_tiles_in_tb = (K_i32 - 1) % N_TILES_IN_TB + 1;\n }\n\n const int input_offset = bid_y * SF_TILE_DIM_M_I32 * K_i32 + bid_x * N_TILES_IN_TB;\n const int* input_i32 = reinterpret_cast(input) + input_offset;\n int* output_i32 = reinterpret_cast(output) + bid_y * SF_TILE_DIM_M_I32 * K_i32 +\n bid_x * N_TILES_IN_TB * SF_TILE_SIZE_I32;\n\n extern __shared__ int4 slm_v4i[];\n\n // load, global -> regs\n // Each register read for a given i is along the K direction at row\n // (bid_y * SF_TILE_DIM_M + i * TB_DIM + threadIdx.y). When that row is past\n // original_M, the entire register is out of the per-tensor data region (which\n // may be the unpadded compact extent), so we must NOT issue the __ldg there --\n // it could read past the per-tensor buffer (and, for the last tensor in a\n // grouped allocation, past the end of the allocation entirely).\n LType regs_vec[N_SF_PER_TD_PER_TILE];\n if (threadIdx.x * N_TILE_PER_TD < n_tiles_in_tb) {\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) {\n const int row = bid_y * SF_TILE_DIM_M + i * TB_DIM + threadIdx.y;\n const int thread_offset = (i * TB_DIM + threadIdx.y) * K_i32 + threadIdx.x * N_TILE_PER_TD;\n if constexpr (IS_PADDED_M) {\n if (row >= original_M) {\n // Entire register is past original_M: zero directly without loading.\n uint8_t* zero_bytes = reinterpret_cast(regs_vec + i);\n#pragma unroll\n for (int j = 0; j < static_cast(sizeof(LType)); j++) zero_bytes[j] = 0;\n continue;\n }\n }\n regs_vec[i] = __ldg(reinterpret_cast(input_i32 + thread_offset));\n // Per-byte K masking is still needed when only part of the register is past\n // original_K (i.e. row is in range but the K position spans the boundary).\n if constexpr (IS_PADDED_K) {\n#pragma unroll\n for (int j = 0; j < N_TILE_PER_TD * sizeof(int); j++) {\n const int index = (input_offset + thread_offset) * sizeof(int) + j;\n if (index % K >= original_K) {\n reinterpret_cast(regs_vec + i)[j] = 0;\n }\n }\n }\n }\n\n // shuffle regs\n regs_shuffle(regs_vec);\n\n// store, regs -> shared\n#pragma unroll\n for (int i = 0; i < N_TILE_PER_TD; i++) {\n /* TODO rotate i */\n slm_v4i[(threadIdx.x * N_TILE_PER_TD + i) * SF_TILE_SIZE_I32 / 4 + threadIdx.y] =\n reinterpret_cast(regs_vec)[i];\n }\n }\n __syncthreads();\n\n // store, shared -> global\n int linear_id = threadIdx.y * blockDim.x + threadIdx.x;\n __align__(16) int4* output_v4i = reinterpret_cast(output_i32);\n#pragma unroll\n for (int i = linear_id; i < SF_TILE_SIZE_I32 * n_tiles_in_tb / 4; i += blockDim.x * blockDim.y) {\n output_v4i[i] = slm_v4i[i];\n }\n}\n\n// Dispatch helper: pick the right (IS_PADDED_K, IS_PADDED_M) row-scaling impl\n// specialization at runtime based on the per-block padding predicates. The\n// branching here is uniform across all threads in the block, so the indirect\n// path each block takes still inlines cleanly.\ntemplate \n__device__ __forceinline__ void dispatch_swizzle_row_scaling_kernel_impl(\n const void* input, void* output, const int M, const int K, const int original_M,\n const int original_K, const int bid_x, const int bid_y, const int grid_dim_x,\n const int grid_dim_y, const bool padding_k, const bool padding_m) {\n if (padding_k && padding_m) {\n swizzle_row_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n } else if (padding_k) {\n swizzle_row_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n } else if (padding_m) {\n swizzle_row_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n } else {\n swizzle_row_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n swizzle_row_scaling_kernel(const void* input, void* output, const int M, const int K,\n const int original_M, const int original_K) {\n const bool padding_m = (blockIdx.y == gridDim.y - 1) && (original_M < M);\n const bool padding_k = (blockIdx.x == gridDim.x - 1) && (original_K < K);\n dispatch_swizzle_row_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y,\n padding_k, padding_m);\n}\n\n// Narrow-K specialization for row scaling swizzle.\n// When K is small (num_tiles_k < TB_DIM), the standard kernel wastes threadIdx.x\n// because there aren't enough K-tiles to distribute across threads.\n// This kernel repurposes the thread dimensions: threadIdx.x iterates rows within\n// an M-tile, threadIdx.y indexes M-tiles within the block, processing TB_DIM\n// M-tiles per block with full thread utilization.\ntemplate \n__device__ void swizzle_row_scaling_narrow_k_kernel_impl(const void* input, void* output,\n const int M, const int K,\n const int original_M, const int original_K,\n const int bid, const int grid_dim) {\n constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4;\n const int K_i32 = K / 4;\n const int num_tiles_m = M / SF_TILE_DIM_M;\n\n const int m_tile = bid * blockDim.y + threadIdx.y;\n const bool active = (m_tile < num_tiles_m);\n\n extern __shared__ int4 slm_v4i[];\n const int slm_tile_v4i = K_i32 * (SF_TILE_SIZE_I32 / 4);\n\n if (active) {\n const bool padding_m = (m_tile == num_tiles_m - 1) && (original_M < M);\n const bool padding_k = (original_K < K);\n\n int4* my_slm = slm_v4i + threadIdx.y * slm_tile_v4i;\n\n for (int k = 0; k < K_i32; k++) {\n const int input_base = m_tile * SF_TILE_DIM_M * K_i32 + k;\n const int* input_i32 = reinterpret_cast(input) + input_base;\n\n int regs[N_SF_PER_TD_PER_TILE];\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) {\n const int row = i * TB_DIM + threadIdx.x;\n regs[i] = __ldg(input_i32 + row * K_i32);\n if (padding_m || padding_k) {\n for (int j = 0; j < 4; j++) {\n const int byte_row = m_tile * SF_TILE_DIM_M + row;\n const int byte_col = k * 4 + j;\n if (byte_row >= original_M || byte_col >= original_K) {\n reinterpret_cast(®s[i])[j] = 0;\n }\n }\n }\n }\n\n my_slm[k * (SF_TILE_SIZE_I32 / 4) + threadIdx.x] = *reinterpret_cast(regs);\n }\n }\n\n __syncthreads();\n\n if (active) {\n int4* my_slm = slm_v4i + threadIdx.y * slm_tile_v4i;\n int4* out_v4i =\n reinterpret_cast(reinterpret_cast(output) + m_tile * SF_TILE_DIM_M * K_i32);\n\n for (int i = threadIdx.x; i < slm_tile_v4i; i += blockDim.x) {\n out_v4i[i] = my_slm[i];\n }\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n swizzle_row_scaling_narrow_k_kernel(const void* input, void* output, const int M, const int K,\n const int original_M, const int original_K) {\n swizzle_row_scaling_narrow_k_kernel_impl(\n input, output, M, K, original_M, original_K, blockIdx.x, gridDim.x);\n}\n\n// Narrow-M variant of the column scaling swizzle kernel, for when num_tiles_m < TB_DIM.\n// Analogous to the narrow-K row kernel: when the M dimension is small, the normal\n// col kernel underutilizes threads in the load phase because threadIdx.x covers M\n// positions with vectorized loads, leaving many threads idle. This kernel repurposes\n// thread dimensions: threadIdx.y indexes K-tiles within the block, threadIdx.x covers\n// one int32 column of an M-tile, and M-tiles are iterated serially.\ntemplate \n__device__ void swizzle_col_scaling_narrow_m_kernel_impl(const void* input, void* output,\n const int M, const int K,\n const int original_M, const int original_K,\n const int bid, const int grid_dim) {\n constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4;\n constexpr int SF_TILE_DIM_M_I32 = SF_TILE_DIM_M / 4;\n constexpr int SF_TILE_DIM_K_I32 = SF_TILE_DIM_K;\n\n const int M_i32 = M / 4;\n const int K_i32 = K;\n const int num_tiles_m = M / SF_TILE_DIM_M;\n const int num_tiles_k = K / SF_TILE_DIM_K;\n\n const int k_tile = bid * blockDim.y + threadIdx.y;\n const bool active = (k_tile < num_tiles_k);\n const int remaining = num_tiles_k - bid * static_cast(blockDim.y);\n const int k_tiles_in_block = remaining <= 0 ? 0 : (remaining < TB_DIM ? remaining : TB_DIM);\n\n extern __shared__ int slm_narrow_m[];\n\n if (active) {\n const bool padding_k = (k_tile == num_tiles_k - 1) && (original_K < K);\n const int32_t* input_i32 = reinterpret_cast(input);\n\n for (int m_tile = 0; m_tile < num_tiles_m; m_tile++) {\n const bool padding_m = (m_tile == num_tiles_m - 1) && (original_M < M);\n\n int regs[N_SF_PER_TD_PER_TILE];\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) {\n const int k_row = k_tile * SF_TILE_DIM_K_I32 + i;\n const int m_col = m_tile * SF_TILE_DIM_M_I32 + threadIdx.x;\n regs[i] = __ldg(input_i32 + k_row * M_i32 + m_col);\n if (padding_m || padding_k) {\n for (int j = 0; j < 4; j++) {\n if (m_col * 4 + j >= original_M || k_row >= original_K) {\n reinterpret_cast(®s[i])[j] = 0;\n }\n }\n }\n }\n\n regs_shuffle_with_bit_shifts(regs);\n\n int tM = threadIdx.x * N_SF_PER_TD_PER_TILE;\n int* slm_tile =\n slm_narrow_m + m_tile * TB_DIM * SF_TILE_SIZE_I32 + threadIdx.y * SF_TILE_SIZE_I32;\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) {\n slm_tile[(tM % SF_TILE_DIM_M) / NEW_SF_TILE_DIM_M_I32 +\n ((tM + i) % NEW_SF_TILE_DIM_M_I32) * NEW_SF_TILE_DIM_K_I32] = regs[i];\n }\n }\n }\n\n __syncthreads();\n\n const int linear_id = threadIdx.y * blockDim.x + threadIdx.x;\n for (int m_tile = 0; m_tile < num_tiles_m; m_tile++) {\n int4* out_v4i = reinterpret_cast(reinterpret_cast(output) +\n m_tile * SF_TILE_DIM_M_I32 * K_i32 +\n bid * TB_DIM * SF_TILE_SIZE_I32);\n int4* slm_v4i = reinterpret_cast(slm_narrow_m + m_tile * TB_DIM * SF_TILE_SIZE_I32);\n const int n_v4i = k_tiles_in_block * SF_TILE_SIZE_I32 / 4;\n for (int j = linear_id; j < n_v4i; j += blockDim.x * blockDim.y) {\n out_v4i[j] = slm_v4i[j];\n }\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n swizzle_col_scaling_narrow_m_kernel(const void* input, void* output, const int M, const int K,\n const int original_M, const int original_K) {\n swizzle_col_scaling_narrow_m_kernel_impl(\n input, output, M, K, original_M, original_K, blockIdx.x, gridDim.x);\n}\n\ntemplate \n__device__ void unswizzle_row_scaling_kernel_impl(const void* input, void* output, const int M,\n const int K, const int bid_x, const int bid_y,\n const int grid_dim_x, const int grid_dim_y) {\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int N_TILES_IN_TB = TB_DIM * N_TILE_PER_TD;\n\n constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4;\n constexpr int SF_TILE_DIM_M_I32 = SF_TILE_DIM_M;\n\n int n_tiles_in_tb = N_TILES_IN_TB;\n const int K_i32 = K / 4;\n if (bid_x == grid_dim_x - 1) {\n n_tiles_in_tb = (K_i32 - 1) % N_TILES_IN_TB + 1;\n }\n\n const int input_offset =\n bid_y * SF_TILE_DIM_M_I32 * K_i32 + bid_x * N_TILES_IN_TB * SF_TILE_SIZE_I32;\n const int* input_i32 = reinterpret_cast(input) + input_offset;\n const int output_offset = bid_y * SF_TILE_DIM_M_I32 * K_i32 + bid_x * N_TILES_IN_TB;\n int* output_i32 = reinterpret_cast(output) + output_offset;\n\n extern __shared__ int4 slm_v4i[];\n\n int linear_id = threadIdx.y * blockDim.x + threadIdx.x;\n const int4* input_v4i = reinterpret_cast(input_i32);\n#pragma unroll\n for (int i = linear_id; i < SF_TILE_SIZE_I32 * n_tiles_in_tb / 4; i += blockDim.x * blockDim.y) {\n slm_v4i[i] = input_v4i[i];\n }\n __syncthreads();\n\n LType regs_vec[N_SF_PER_TD_PER_TILE];\n if (threadIdx.x * N_TILE_PER_TD < n_tiles_in_tb) {\n#pragma unroll\n for (int i = 0; i < N_TILE_PER_TD; i++) {\n reinterpret_cast(regs_vec)[i] =\n slm_v4i[(threadIdx.x * N_TILE_PER_TD + i) * SF_TILE_SIZE_I32 / 4 + threadIdx.y];\n }\n\n regs_unshuffle(regs_vec);\n\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) {\n const int thread_offset = (i * TB_DIM + threadIdx.y) * K_i32 + threadIdx.x * N_TILE_PER_TD;\n reinterpret_cast(output_i32 + thread_offset)[0] = regs_vec[i];\n }\n }\n}\n\ntemplate \n__device__ void unswizzle_col_scaling_kernel_impl(const void* input, void* output, const int M,\n const int K, const int bid_x, const int bid_y,\n const int grid_dim_x, const int grid_dim_y) {\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int N_SF_PER_TD = N_TILE_PER_TD * N_SF_PER_TD_PER_TILE;\n constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4;\n\n constexpr int SF_TILE_DIM_M_I32 = SF_TILE_DIM_M / 4;\n constexpr int SF_TILE_DIM_K_I32 = SF_TILE_DIM_K;\n\n const int M_i32 = M / 4;\n const int K_i32 = K;\n\n int m_tiles_in_tb = N_TILE_PER_TD;\n int k_tiles_in_tb = TB_DIM;\n if (bid_x == grid_dim_x - 1) {\n k_tiles_in_tb = (K_i32 / SF_TILE_DIM_K_I32 - 1) % k_tiles_in_tb + 1;\n }\n if (bid_y == grid_dim_y - 1) {\n m_tiles_in_tb = (M_i32 / SF_TILE_DIM_M_I32 - 1) % m_tiles_in_tb + 1;\n }\n\n const int32_t* input_i32[N_TILE_PER_TD];\n#pragma unroll\n for (int i = 0; i < m_tiles_in_tb; i++) {\n input_i32[i] = reinterpret_cast(input) + bid_x * TB_DIM * SF_TILE_SIZE_I32 +\n (bid_y * N_TILE_PER_TD + i) * SF_TILE_DIM_M_I32 * K_i32;\n }\n const int output_offset =\n bid_x * TB_DIM * SF_TILE_DIM_K_I32 * M_i32 + bid_y * N_TILE_PER_TD * SF_TILE_DIM_M_I32;\n int* output_i32 = reinterpret_cast(output) + output_offset;\n\n extern __shared__ int slm[];\n\n int linear_id = threadIdx.y * blockDim.x + threadIdx.x;\n#pragma unroll\n for (int i = 0; i < m_tiles_in_tb; i++) {\n __align__(16) const int4* input_v4i = reinterpret_cast(input_i32[i]);\n __align__(16) int4* slm_v4i =\n reinterpret_cast(slm + i * k_tiles_in_tb * SF_TILE_SIZE_I32);\n#pragma unroll\n for (int j = linear_id; j < SF_TILE_SIZE_I32 * k_tiles_in_tb / 4;\n j += blockDim.x * blockDim.y) {\n slm_v4i[j] = input_v4i[j];\n }\n }\n __syncthreads();\n\n LType regs_vec[N_SF_PER_TD_PER_TILE];\n if (threadIdx.x * N_TILE_PER_TD < m_tiles_in_tb * SF_TILE_DIM_M_I32 &&\n threadIdx.y < k_tiles_in_tb) {\n int tM = threadIdx.x * N_SF_PER_TD;\n int* slm_tile = slm + (threadIdx.y * SF_TILE_SIZE_I32 +\n tM / SF_TILE_DIM_M * k_tiles_in_tb * SF_TILE_SIZE_I32);\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD; i++) {\n reinterpret_cast(regs_vec)[i] =\n slm_tile[(tM % SF_TILE_DIM_M) / NEW_SF_TILE_DIM_M_I32 +\n ((tM + i) % NEW_SF_TILE_DIM_M_I32) * NEW_SF_TILE_DIM_K_I32];\n }\n\n regs_unshuffle_with_bit_shifts(regs_vec);\n\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) {\n const int thread_offset =\n (threadIdx.y * SF_TILE_DIM_K_I32 + i) * M_i32 + threadIdx.x * N_TILE_PER_TD;\n reinterpret_cast(output_i32 + thread_offset)[0] = regs_vec[i];\n }\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n unswizzle_scaling_kernel(const void* input, void* output, const int M, const int K,\n const bool row_scaling) {\n const int bid_x = blockIdx.x;\n const int bid_y = blockIdx.y;\n const int grid_dim_x = gridDim.x;\n const int grid_dim_y = gridDim.y;\n if (row_scaling) {\n unswizzle_row_scaling_kernel_impl(\n input, output, M, K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n } else {\n unswizzle_col_scaling_kernel_impl(\n input, output, M, K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n }\n}\n\nconstexpr int kMaxTensorsPerKernel = 64; // Args must be <4 KB\nstruct MultiSwizzleArgs {\n // (input) Data buffers for input scaling factors\n void* input_list[kMaxTensorsPerKernel];\n // (output) Data buffers for swizzled scaling factors\n void* output_list[kMaxTensorsPerKernel];\n // Input scaling factor m\n int m_list[kMaxTensorsPerKernel];\n // Input scaling factor k\n int k_list[kMaxTensorsPerKernel];\n // Input scaling factor m before padding\n int original_m_list[kMaxTensorsPerKernel];\n // Input scaling factor k before padding\n int original_k_list[kMaxTensorsPerKernel];\n // Prefix sum (with leading zero) of CUDA blocks needed for each\n // tensor\n int block_range[kMaxTensorsPerKernel + 1];\n // Number of tensors being processed by kernel\n int num_tensors;\n};\n\nconstexpr size_t round_up_to_multiple(size_t value, size_t multiple) {\n return DIVUP(value, multiple) * multiple;\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n grouped_swizzle_row_scaling_uniform_shape_kernel(const void* input, void* output, const int M,\n const int K, const int original_M,\n const int original_K,\n const size_t input_stride_bytes,\n const size_t output_stride_bytes) {\n const int tensor_id = blockIdx.z;\n // Input and output strides may differ: input is in the kernel-produced \"compact\"\n // layout (per-tensor stride = original_M * padded_k * elem_size) when callers\n // pass the unswizzled grouped scale buffer as-is, while the output is always in\n // the per-tensor padded (\"swizzle-ready\") layout (padded_m * padded_k * elem_size).\n const uint8_t* input_base =\n reinterpret_cast(input) + tensor_id * input_stride_bytes;\n uint8_t* output_base = reinterpret_cast(output) + tensor_id * output_stride_bytes;\n const bool padding_m = (blockIdx.y == gridDim.y - 1) && (original_M < M);\n const bool padding_k = (blockIdx.x == gridDim.x - 1) && (original_K < K);\n dispatch_swizzle_row_scaling_kernel_impl(\n input_base, output_base, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x,\n gridDim.y, padding_k, padding_m);\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n grouped_swizzle_col_scaling_uniform_shape_kernel(const void* input, void* output, const int M,\n const int K, const int original_M,\n const int original_K,\n const size_t input_stride_bytes,\n const size_t output_stride_bytes) {\n const int tensor_id = blockIdx.z;\n // See the rowwise kernel for stride semantics. For columnwise the per-tensor\n // compact stride is DIVUP(original_K, 1) * padded_m * elem_size (i.e. the\n // unpadded scale-row count in the K direction times the padded M extent).\n const uint8_t* input_base =\n reinterpret_cast(input) + tensor_id * input_stride_bytes;\n uint8_t* output_base = reinterpret_cast(output) + tensor_id * output_stride_bytes;\n const bool padding_m = (blockIdx.y == gridDim.y - 1) && (original_M < M);\n const bool padding_k = (blockIdx.x == gridDim.x - 1) && (original_K < K);\n dispatch_swizzle_col_scaling_kernel_impl(\n input_base, output_base, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x,\n gridDim.y, padding_k, padding_m);\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n grouped_unswizzle_scaling_uniform_shape_kernel(const void* input, void* output, const int M,\n const int K, const size_t scale_stride_bytes,\n const bool row_scaling) {\n const int tensor_id = blockIdx.z;\n const uint8_t* input_base =\n reinterpret_cast(input) + tensor_id * scale_stride_bytes;\n uint8_t* output_base = reinterpret_cast(output) + tensor_id * scale_stride_bytes;\n if (row_scaling) {\n unswizzle_row_scaling_kernel_impl(\n input_base, output_base, M, K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y);\n } else {\n unswizzle_col_scaling_kernel_impl(\n input_base, output_base, M, K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y);\n }\n}\n\ntemplate \n__global__ void multi_tensor_unswizzle_row_scaling_kernel(MultiSwizzleArgs kernel_args) {\n const int bid = blockIdx.x;\n int tensor_id = 0;\n while (kernel_args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n const void* input = kernel_args.input_list[tensor_id];\n void* output = kernel_args.output_list[tensor_id];\n const int M = kernel_args.m_list[tensor_id];\n const int K = kernel_args.k_list[tensor_id];\n\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int N_TILES_IN_TB = TB_DIM * N_TILE_PER_TD;\n\n const int num_tiles_k = K / SF_TILE_DIM_K;\n const int num_tiles_m = M / SF_TILE_DIM_M;\n const int flat_offset = bid - kernel_args.block_range[tensor_id];\n const int grid_dim_x = DIVUP(num_tiles_k, N_TILES_IN_TB);\n const int grid_dim_y = num_tiles_m;\n const int bid_x = flat_offset / grid_dim_y;\n const int bid_y = flat_offset % grid_dim_y;\n\n unswizzle_row_scaling_kernel_impl(\n input, output, M, K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n}\n\ntemplate \n__global__ void multi_tensor_unswizzle_col_scaling_kernel(MultiSwizzleArgs kernel_args) {\n const int bid = blockIdx.x;\n int tensor_id = 0;\n while (kernel_args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n const void* input = kernel_args.input_list[tensor_id];\n void* output = kernel_args.output_list[tensor_id];\n const int M = kernel_args.m_list[tensor_id];\n const int K = kernel_args.k_list[tensor_id];\n\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n\n const int num_tiles_k = K / SF_TILE_DIM_K;\n const int num_tiles_m = M / SF_TILE_DIM_M;\n const int flat_offset = bid - kernel_args.block_range[tensor_id];\n const int grid_dim_x = DIVUP(num_tiles_k, TB_DIM);\n const int grid_dim_y = DIVUP(num_tiles_m, N_TILE_PER_TD);\n const int bid_x = flat_offset / grid_dim_y;\n const int bid_y = flat_offset % grid_dim_y;\n\n unswizzle_col_scaling_kernel_impl(\n input, output, M, K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n}\n\ntemplate \n__global__ void multi_tensor_swizzle_row_scaling_kernel(MultiSwizzleArgs kernel_args) {\n // Find tensor corresponding to block\n const int bid = blockIdx.x;\n int tensor_id = 0;\n while (kernel_args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n // Get args corresponding to block\n const void* input = kernel_args.input_list[tensor_id];\n void* output = kernel_args.output_list[tensor_id];\n const int M = kernel_args.m_list[tensor_id];\n const int K = kernel_args.k_list[tensor_id];\n const int original_M = kernel_args.original_m_list[tensor_id];\n const int original_K = kernel_args.original_k_list[tensor_id];\n\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int N_TILES_IN_TB = TB_DIM * N_TILE_PER_TD;\n\n // Get block index in grid. Emulate 2D grid.\n const int num_tiles_k = K / SF_TILE_DIM_K;\n const int num_tiles_m = M / SF_TILE_DIM_M;\n const int grid_dim_x = DIVUP(num_tiles_k, N_TILES_IN_TB);\n const int grid_dim_y = num_tiles_m;\n const int bid_x = (bid - kernel_args.block_range[tensor_id]) / grid_dim_y;\n const int bid_y = (bid - kernel_args.block_range[tensor_id]) % grid_dim_y;\n\n const bool padding_m = (bid_y == grid_dim_y - 1) && (original_M < M);\n const bool padding_k = (bid_x == grid_dim_x - 1) && (original_K < K);\n dispatch_swizzle_row_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y, padding_k,\n padding_m);\n}\n\ntemplate \n__global__ void multi_tensor_swizzle_col_scaling_kernel(MultiSwizzleArgs kernel_args) {\n // Find tensor corresponding to block\n const int bid = blockIdx.x;\n int tensor_id = 0;\n while (kernel_args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n // Get args corresponding to block\n const void* input = kernel_args.input_list[tensor_id];\n void* output = kernel_args.output_list[tensor_id];\n const int M = kernel_args.m_list[tensor_id];\n const int K = kernel_args.k_list[tensor_id];\n const int original_M = kernel_args.original_m_list[tensor_id];\n const int original_K = kernel_args.original_k_list[tensor_id];\n\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n\n // Get block index in grid. Emulate 2D grid.\n const int num_tiles_k = K / SF_TILE_DIM_K;\n const int num_tiles_m = M / SF_TILE_DIM_M;\n const int grid_dim_x = DIVUP(num_tiles_k, TB_DIM);\n const int grid_dim_y = DIVUP(num_tiles_m, N_TILE_PER_TD);\n const int bid_x = (bid - kernel_args.block_range[tensor_id]) / grid_dim_y;\n const int bid_y = (bid - kernel_args.block_range[tensor_id]) % grid_dim_y;\n\n const bool padding_m = (bid_y == grid_dim_y - 1) && (original_M < M);\n const bool padding_k = (bid_x == grid_dim_x - 1) && (original_K < K);\n dispatch_swizzle_col_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y, padding_k,\n padding_m);\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n multi_tensor_swizzle_row_scaling_narrow_k_kernel(MultiSwizzleArgs kernel_args) {\n const int bid = blockIdx.x;\n int tensor_id = 0;\n while (kernel_args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n const void* input = kernel_args.input_list[tensor_id];\n void* output = kernel_args.output_list[tensor_id];\n const int M = kernel_args.m_list[tensor_id];\n const int K = kernel_args.k_list[tensor_id];\n const int original_M = kernel_args.original_m_list[tensor_id];\n const int original_K = kernel_args.original_k_list[tensor_id];\n const int flat_bid = bid - kernel_args.block_range[tensor_id];\n const int num_tiles_m = M / SF_TILE_DIM_M;\n const int grid_dim = DIVUP(num_tiles_m, TB_DIM);\n\n swizzle_row_scaling_narrow_k_kernel_impl(\n input, output, M, K, original_M, original_K, flat_bid, grid_dim);\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n multi_tensor_swizzle_col_scaling_narrow_m_kernel(MultiSwizzleArgs kernel_args) {\n const int bid = blockIdx.x;\n int tensor_id = 0;\n while (kernel_args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n const void* input = kernel_args.input_list[tensor_id];\n void* output = kernel_args.output_list[tensor_id];\n const int M = kernel_args.m_list[tensor_id];\n const int K = kernel_args.k_list[tensor_id];\n const int original_M = kernel_args.original_m_list[tensor_id];\n const int original_K = kernel_args.original_k_list[tensor_id];\n const int flat_bid = bid - kernel_args.block_range[tensor_id];\n const int num_tiles_k = K / SF_TILE_DIM_K;\n const int grid_dim = DIVUP(num_tiles_k, TB_DIM);\n\n swizzle_col_scaling_narrow_m_kernel_impl(\n input, output, M, K, original_M, original_K, flat_bid, grid_dim);\n}\n\n} // namespace\n\nvoid swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t stream) {\n // Check scaling mode\n const auto& scaling_mode = input->scaling_mode;\n NVTE_CHECK(scaling_mode == NVTE_MXFP8_1D_SCALING || scaling_mode == NVTE_NVFP4_1D_SCALING,\n \"Input tensor has invalid scaling mode (\", to_string(input->scaling_mode), \").\");\n\n // Check tensors\n CheckInputTensor(*input, \"scaling_factor_input\");\n CheckInputTensor(*output, \"scaling_factor_output\");\n NVTE_CHECK(!input->with_gemm_swizzled_scales,\n \"Expected input tensor with scales in compact format.\");\n NVTE_CHECK(output->with_gemm_swizzled_scales,\n \"Expected output tensor with scales in GEMM swizzled format.\");\n switch (scaling_mode) {\n case NVTE_MXFP8_1D_SCALING:\n NVTE_CHECK(is_fp8_dtype(input->dtype()), \"Input tensor has invalid dtype (expected FP8, got \",\n to_string(input->dtype()), \").\");\n break;\n case NVTE_NVFP4_1D_SCALING:\n NVTE_CHECK(is_fp4_dtype(input->dtype()), \"Input tensor has invalid dtype (expected FP4, got \",\n to_string(input->dtype()), \").\");\n break;\n default:\n NVTE_ERROR(\"Invalid scaling mode\");\n }\n\n // Check if scaling factors are non-trivial\n const bool has_rowwise_scale_inv = input->scale_inv.has_data();\n const bool has_columnwise_scale_inv = input->columnwise_scale_inv.has_data();\n NVTE_CHECK(!has_rowwise_scale_inv || !has_columnwise_scale_inv,\n \"Input tensor has both row-wise and column-wise scaling factors\");\n if (!has_rowwise_scale_inv && !has_columnwise_scale_inv) {\n return;\n }\n\n // Deduce tensor dims\n int m{0}, k{0};\n switch (scaling_mode) {\n case NVTE_MXFP8_1D_SCALING: {\n if (has_rowwise_scale_inv) {\n NVTE_CHECK(input->scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", input->scale_inv.shape, \".\");\n m = input->scale_inv.shape[0];\n k = input->scale_inv.shape[1];\n } else if (has_columnwise_scale_inv) {\n NVTE_CHECK(input->columnwise_scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", input->columnwise_scale_inv.shape,\n \".\");\n m = input->columnwise_scale_inv.shape[1];\n k = input->columnwise_scale_inv.shape[0];\n }\n break;\n }\n case NVTE_NVFP4_1D_SCALING: {\n if (has_rowwise_scale_inv) {\n NVTE_CHECK(input->scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", input->scale_inv.shape, \".\");\n m = input->scale_inv.shape[0];\n k = input->scale_inv.shape[1];\n } else if (has_columnwise_scale_inv) {\n NVTE_CHECK(input->columnwise_scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", input->columnwise_scale_inv.shape,\n \".\");\n m = input->columnwise_scale_inv.shape[0];\n k = input->columnwise_scale_inv.shape[1];\n }\n break;\n }\n default:\n NVTE_ERROR(\"Invalid scaling mode\");\n }\n\n // Check dims\n constexpr int SF_TILE_DIM_M = 128;\n constexpr int SF_TILE_DIM_K = 4;\n NVTE_CHECK(m % SF_TILE_DIM_M == 0, \"Input should be padded in M/N dimension!\");\n NVTE_CHECK(k % SF_TILE_DIM_K == 0, \"Input should be padded in K dimension!\");\n\n // Check that output tensor matches input tensor\n if (has_rowwise_scale_inv) {\n NVTE_CHECK(output->scale_inv.has_data(),\n \"Output tensor does not have row-wise scaling factors.\");\n NVTE_CHECK(m * k == output->scale_inv.numel(), \"Expected output tensor to have \", m * k,\n \" row-wise scaling factors, but got shape=\", output->scale_inv.shape, \".\");\n }\n if (has_columnwise_scale_inv) {\n NVTE_CHECK(output->columnwise_scale_inv.has_data(),\n \"Output tensor does not have column-wise scaling factors.\");\n NVTE_CHECK(\n m * k == output->columnwise_scale_inv.numel(), \"Expected output tensor to have \", m * k,\n \" column-wise scaling factors, but got shape=\", output->columnwise_scale_inv.shape, \".\");\n }\n\n // Choose swizzle implementation\n bool rowwise_swizzle{false}, columnwise_swizzle{false};\n switch (scaling_mode) {\n case NVTE_MXFP8_1D_SCALING: {\n rowwise_swizzle = has_rowwise_scale_inv;\n columnwise_swizzle = has_columnwise_scale_inv;\n break;\n }\n case NVTE_NVFP4_1D_SCALING: {\n // NVFP4 column-wise data is transposed, so row-wise and\n // column-wise scales have same swizzling format\n rowwise_swizzle = true;\n columnwise_swizzle = false;\n break;\n }\n default:\n NVTE_ERROR(\"Invalid scaling mode\");\n }\n\n const dim3 block_size(TB_DIM, TB_DIM);\n const int num_tiles_m = m / SF_TILE_DIM_M;\n const int num_tiles_k = k / SF_TILE_DIM_K;\n\n // Perform row-wise swizzle\n if (rowwise_swizzle) {\n int original_M{0}, original_K{0};\n void *input_scale_inv_ptr{nullptr}, *output_scale_inv_ptr{nullptr};\n switch (scaling_mode) {\n case NVTE_MXFP8_1D_SCALING: {\n original_M = input->flat_first_dim();\n original_K = input->flat_last_dim() / MXFP8_BLOCK_SIZE;\n input_scale_inv_ptr = input->scale_inv.dptr;\n output_scale_inv_ptr = output->scale_inv.dptr;\n break;\n }\n case NVTE_NVFP4_1D_SCALING: {\n if (has_rowwise_scale_inv) {\n original_M = input->flat_first_dim();\n original_K = input->flat_last_dim() / NVFP4_BLOCK_SIZE;\n input_scale_inv_ptr = input->scale_inv.dptr;\n output_scale_inv_ptr = output->scale_inv.dptr;\n } else if (has_columnwise_scale_inv) {\n original_M = input->flat_last_dim();\n original_K = input->flat_first_dim() / NVFP4_BLOCK_SIZE;\n input_scale_inv_ptr = input->columnwise_scale_inv.dptr;\n output_scale_inv_ptr = output->columnwise_scale_inv.dptr;\n }\n break;\n }\n default:\n NVTE_ERROR(\"Invalid scaling mode\");\n }\n\n const int narrow_k_slm_size =\n TB_DIM * num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t));\n if (num_tiles_k < TB_DIM && narrow_k_slm_size <= get_max_dynamic_smem()) {\n // Narrow-K: batch TB_DIM M-tiles per block, fully utilizing all threads.\n dim3 num_blocks_narrow(DIVUP(num_tiles_m, TB_DIM));\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(swizzle_row_scaling_narrow_k_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, narrow_k_slm_size));\n swizzle_row_scaling_narrow_k_kernel\n <<>>(\n input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K);\n } else {\n int vec_load_size = (num_tiles_k - 1) % 4 + 1;\n /* there is no int3 and misaligned if using int4/int2 */\n if (vec_load_size == 3) vec_load_size = 1;\n int n_tiles_in_tb = TB_DIM * vec_load_size;\n dim3 num_blocks(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m);\n int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n\n switch (vec_load_size) {\n case 4:\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(swizzle_row_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n swizzle_row_scaling_kernel\n <<>>(\n input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K);\n break;\n case 2:\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(swizzle_row_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n swizzle_row_scaling_kernel\n <<>>(\n input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K);\n break;\n case 1:\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(swizzle_row_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n swizzle_row_scaling_kernel\n <<>>(\n input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K);\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n break;\n }\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n\n // Perform column-wise swizzle\n if (columnwise_swizzle) {\n const int original_M = input->flat_last_dim();\n const int original_K = input->flat_first_dim() / MXFP8_BLOCK_SIZE;\n\n const int narrow_m_slm_size =\n TB_DIM * num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t));\n if (num_tiles_m < TB_DIM && narrow_m_slm_size <= get_max_dynamic_smem()) {\n // Narrow-M: batch TB_DIM K-tiles per block, fully utilizing all threads.\n dim3 num_blocks_narrow(DIVUP(num_tiles_k, TB_DIM));\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(swizzle_col_scaling_narrow_m_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, narrow_m_slm_size));\n swizzle_col_scaling_narrow_m_kernel\n <<>>(\n input->columnwise_scale_inv.dptr, output->columnwise_scale_inv.dptr, m, k, original_M,\n original_K);\n } else {\n int vec_load_size = (num_tiles_m - 1) % 4 + 1;\n if (vec_load_size == 3) vec_load_size = 1; /* no int3 and misaligned if using int4/int2 */\n int n_tiles_in_tb = TB_DIM * vec_load_size;\n dim3 num_blocks(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size));\n int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n\n switch (vec_load_size) {\n case 4:\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(swizzle_col_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n swizzle_col_scaling_kernel\n <<>>(input->columnwise_scale_inv.dptr,\n output->columnwise_scale_inv.dptr, m,\n k, original_M, original_K);\n break;\n case 2:\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(swizzle_col_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n swizzle_col_scaling_kernel\n <<>>(input->columnwise_scale_inv.dptr,\n output->columnwise_scale_inv.dptr, m,\n k, original_M, original_K);\n break;\n case 1:\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(swizzle_col_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n swizzle_col_scaling_kernel\n <<>>(input->columnwise_scale_inv.dptr,\n output->columnwise_scale_inv.dptr, m,\n k, original_M, original_K);\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n break;\n }\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\ntemplate \nvoid launch_multi_tensor_swizzle_scaling_factors(MultiSwizzleArgs& kernel_args,\n const int vec_load_size, const bool is_rowwise,\n const bool use_narrow_k, const bool use_narrow_m,\n cudaStream_t stream) {\n // cudaFuncSetAttribute is a host-synchronous driver call; cache the max shared memory\n // setting per kernel variant so we only pay the cost when slm_size actually increases.\n auto set_smem_if_needed = [](auto kernel_fn, int slm, int& cached) {\n if (cached < slm) {\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, slm));\n cached = slm;\n }\n };\n\n dim3 block_size(TB_DIM, TB_DIM);\n\n if (is_rowwise && use_narrow_k) {\n // Narrow-K path: each block handles TB_DIM M-tiles with full thread utilization.\n // slm_size depends on num_tiles_k, which can vary per tensor \u2014 use the max.\n int max_num_tiles_k = 0;\n for (size_t j = 0; j < kernel_args.num_tensors; j++) {\n const int num_tiles_m = kernel_args.m_list[j] / SF_TILE_DIM_M;\n const int num_tiles_k = kernel_args.k_list[j] / SF_TILE_DIM_K;\n max_num_tiles_k = std::max(max_num_tiles_k, num_tiles_k);\n kernel_args.block_range[j + 1] = kernel_args.block_range[j] + DIVUP(num_tiles_m, TB_DIM);\n }\n int slm_size = TB_DIM * max_num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n const int num_blocks = kernel_args.block_range[kernel_args.num_tensors];\n\n static int cached_narrow_k = -1;\n set_smem_if_needed(\n multi_tensor_swizzle_row_scaling_narrow_k_kernel, slm_size,\n cached_narrow_k);\n multi_tensor_swizzle_row_scaling_narrow_k_kernel\n <<>>(kernel_args);\n } else if (!is_rowwise && use_narrow_m) {\n // Narrow-M path: each block handles TB_DIM K-tiles with full thread utilization.\n // slm_size depends on num_tiles_m, which can vary per tensor \u2014 use the max.\n int max_num_tiles_m = 0;\n for (size_t j = 0; j < kernel_args.num_tensors; j++) {\n const int num_tiles_m = kernel_args.m_list[j] / SF_TILE_DIM_M;\n const int num_tiles_k = kernel_args.k_list[j] / SF_TILE_DIM_K;\n max_num_tiles_m = std::max(max_num_tiles_m, num_tiles_m);\n kernel_args.block_range[j + 1] = kernel_args.block_range[j] + DIVUP(num_tiles_k, TB_DIM);\n }\n int slm_size = TB_DIM * max_num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n const int num_blocks = kernel_args.block_range[kernel_args.num_tensors];\n\n static int cached_narrow_m = -1;\n set_smem_if_needed(\n multi_tensor_swizzle_col_scaling_narrow_m_kernel, slm_size,\n cached_narrow_m);\n multi_tensor_swizzle_col_scaling_narrow_m_kernel\n <<>>(kernel_args);\n } else {\n int n_tiles_in_tb = TB_DIM * vec_load_size;\n int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n /* Calculate number of CUDA blocks needed for each tensor.\n * We have to do it here because we have to iterate over all tensors in this batch to\n * get the minimum vec_load_size.\n */\n for (size_t j = 0; j < kernel_args.num_tensors; j++) {\n const int m = kernel_args.m_list[j];\n const int k = kernel_args.k_list[j];\n int num_tiles_m = m / SF_TILE_DIM_M;\n int num_tiles_k = k / SF_TILE_DIM_K;\n if (is_rowwise) {\n kernel_args.block_range[j + 1] =\n kernel_args.block_range[j] + DIVUP(num_tiles_k, n_tiles_in_tb) * num_tiles_m;\n } else {\n kernel_args.block_range[j + 1] =\n kernel_args.block_range[j] +\n DIVUP(num_tiles_k, TB_DIM) * DIVUP(num_tiles_m, vec_load_size);\n }\n }\n const int num_blocks = kernel_args.block_range[kernel_args.num_tensors];\n\n static int cached_row_int4 = -1, cached_row_int2 = -1, cached_row_int1 = -1;\n static int cached_col_int4 = -1, cached_col_int2 = -1, cached_col_int1 = -1;\n\n if (is_rowwise) {\n switch (vec_load_size) {\n case 4:\n set_smem_if_needed(\n multi_tensor_swizzle_row_scaling_kernel, slm_size,\n cached_row_int4);\n multi_tensor_swizzle_row_scaling_kernel\n <<>>(kernel_args);\n break;\n case 2:\n set_smem_if_needed(\n multi_tensor_swizzle_row_scaling_kernel, slm_size,\n cached_row_int2);\n multi_tensor_swizzle_row_scaling_kernel\n <<>>(kernel_args);\n break;\n case 1:\n set_smem_if_needed(\n multi_tensor_swizzle_row_scaling_kernel, slm_size,\n cached_row_int1);\n multi_tensor_swizzle_row_scaling_kernel\n <<>>(kernel_args);\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n break;\n }\n } else {\n switch (vec_load_size) {\n case 4:\n set_smem_if_needed(\n multi_tensor_swizzle_col_scaling_kernel, slm_size,\n cached_col_int4);\n multi_tensor_swizzle_col_scaling_kernel\n <<>>(kernel_args);\n break;\n case 2:\n set_smem_if_needed(\n multi_tensor_swizzle_col_scaling_kernel, slm_size,\n cached_col_int2);\n multi_tensor_swizzle_col_scaling_kernel\n <<>>(kernel_args);\n break;\n case 1:\n set_smem_if_needed(\n multi_tensor_swizzle_col_scaling_kernel, slm_size,\n cached_col_int1);\n multi_tensor_swizzle_col_scaling_kernel\n <<>>(kernel_args);\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n break;\n }\n }\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\ntemplate \nvoid launch_multi_tensor_unswizzle_scaling_factors(MultiSwizzleArgs& kernel_args,\n const int vec_load_size, const bool is_rowwise,\n cudaStream_t stream) {\n int n_tiles_in_tb = TB_DIM * vec_load_size;\n int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n for (size_t j = 0; j < kernel_args.num_tensors; j++) {\n const int m = kernel_args.m_list[j];\n const int k = kernel_args.k_list[j];\n int num_tiles_m = m / SF_TILE_DIM_M;\n int num_tiles_k = k / SF_TILE_DIM_K;\n if (is_rowwise) {\n kernel_args.block_range[j + 1] =\n kernel_args.block_range[j] + DIVUP(num_tiles_k, n_tiles_in_tb) * num_tiles_m;\n } else {\n kernel_args.block_range[j + 1] =\n kernel_args.block_range[j] +\n DIVUP(num_tiles_k, TB_DIM) * DIVUP(num_tiles_m, vec_load_size);\n }\n }\n\n int num_blocks = kernel_args.block_range[kernel_args.num_tensors];\n if (num_blocks > 0) {\n dim3 block_size(TB_DIM, TB_DIM);\n if (is_rowwise) {\n switch (vec_load_size) {\n case 4:\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n multi_tensor_unswizzle_row_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n multi_tensor_unswizzle_row_scaling_kernel\n <<>>(kernel_args);\n break;\n case 2:\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n multi_tensor_unswizzle_row_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n multi_tensor_unswizzle_row_scaling_kernel\n <<>>(kernel_args);\n break;\n case 1:\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n multi_tensor_unswizzle_row_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n multi_tensor_unswizzle_row_scaling_kernel\n <<>>(kernel_args);\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n break;\n }\n } else {\n switch (vec_load_size) {\n case 4:\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n multi_tensor_unswizzle_col_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n multi_tensor_unswizzle_col_scaling_kernel\n <<>>(kernel_args);\n break;\n case 2:\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n multi_tensor_unswizzle_col_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n multi_tensor_unswizzle_col_scaling_kernel\n <<>>(kernel_args);\n break;\n case 1:\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n multi_tensor_unswizzle_col_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n multi_tensor_unswizzle_col_scaling_kernel\n <<>>(kernel_args);\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n break;\n }\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\nvoid multi_tensor_swizzle_scaling_factors(const std::vector& input,\n std::vector& output, cudaStream_t stream,\n bool check_scale_inv_shapes) {\n auto num_tensors = input.size();\n bool all_has_data = true;\n bool all_has_columnwise_data = true;\n bool all_nvfp4 = true;\n for (size_t i = 0; i < num_tensors; i++) {\n auto scaling_mode = input[i]->scaling_mode;\n auto is_fp8 = is_fp8_dtype(input[i]->dtype());\n auto is_fp4 = is_fp4_dtype(input[i]->dtype());\n NVTE_CHECK(\n (is_fp8 && is_mxfp8_scaling(scaling_mode)) || (is_fp4 && is_nvfp4_scaling(scaling_mode)),\n \"Not implemented scaling mode \" + to_string(scaling_mode) + \".\");\n NVTE_CHECK(!input[i]->with_gemm_swizzled_scales,\n \"Expected input tensors with scales in compact format.\");\n NVTE_CHECK(output[i]->with_gemm_swizzled_scales,\n \"Expected output tensors with scales in GEMM swizzled format.\");\n\n // We don't allow empty tensors. They should be filtered out before calling this function.\n NVTE_CHECK(input[i]->numel() != 0, \"Tensor input[\", i, \"] is empty.\");\n CheckInputTensor(*input[i], \"scaling_factor_input[\" + std::to_string(i) + \"]\",\n check_scale_inv_shapes);\n CheckInputTensor(*output[i], \"scaling_factor_output[\" + std::to_string(i) + \"]\",\n check_scale_inv_shapes);\n all_has_data = all_has_data && input[i]->scale_inv.has_data();\n all_has_columnwise_data =\n (all_has_columnwise_data && input[i]->columnwise_scale_inv.has_data());\n all_nvfp4 = all_nvfp4 && is_nvfp4_scaling(scaling_mode);\n }\n NVTE_CHECK(all_has_data || all_has_columnwise_data,\n \"All tensors should have data or columnwise data.\");\n NVTE_CHECK(!all_has_data || !all_has_columnwise_data,\n \"All tensors have both data and columnwise data.\");\n\n const bool rowwise_swizzle = all_has_data || all_nvfp4;\n const bool columnwise_swizzle = all_has_columnwise_data && !all_nvfp4;\n\n constexpr int SF_TILE_DIM_M = 128;\n constexpr int SF_TILE_DIM_K = 4;\n if (rowwise_swizzle) {\n MultiSwizzleArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.block_range[0] = 0;\n int vec_load_size = 4;\n bool all_narrow_k = true;\n for (size_t i = 0; i < num_tensors; i++) {\n //Launch kernel if argument struct is full\n if (kernel_args.num_tensors == kMaxTensorsPerKernel) {\n // There is no int3 and misaligned if using int4/int2.\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_swizzle_scaling_factors(\n kernel_args, vec_load_size, true, all_narrow_k, false, stream);\n // Reset the argument struct and vec_load_size\n kernel_args.num_tensors = 0;\n vec_load_size = 4;\n all_narrow_k = true;\n }\n\n int m, k;\n\n if (all_has_data) {\n m = input[i]->scale_inv.shape[0];\n k = input[i]->scale_inv.shape[1];\n } else {\n NVTE_CHECK(all_nvfp4, \"When doing rowwise swizzle with rowwise data, it has to be NVFP4\");\n m = input[i]->columnwise_scale_inv.shape[0];\n k = input[i]->columnwise_scale_inv.shape[1];\n }\n\n NVTE_CHECK(m % SF_TILE_DIM_M == 0, \"Input should be padded in M/N dimension!\");\n NVTE_CHECK(k % SF_TILE_DIM_K == 0, \"Input should be padded in K dimension!\");\n NVTE_CHECK(k > 0, \"Input scale inverse should be 2D!\");\n\n if (all_has_data) {\n NVTE_CHECK(output[i]->scale_inv.has_data(), \"Output tensor \", i,\n \" does not have row-wise scaling factors.\");\n NVTE_CHECK(m * k == output[i]->scale_inv.numel(), \"Expected output tensor \", i, \" to have \",\n m * k, \" row-wise scaling factors, but got shape=\", output[i]->scale_inv.shape,\n \".\");\n }\n if (all_has_columnwise_data) {\n NVTE_CHECK(output[i]->columnwise_scale_inv.has_data(), \"Output tensor \", i,\n \" does not have column-wise scaling factors.\");\n NVTE_CHECK(m * k == output[i]->columnwise_scale_inv.numel(), \"Expected output tensor \", i,\n \" to have \", m * k, \" column-wise scaling factors, but got shape=\",\n output[i]->columnwise_scale_inv.shape, \".\");\n }\n\n int num_tiles_k = k / SF_TILE_DIM_K;\n const int narrow_k_slm =\n TB_DIM * num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t));\n all_narrow_k =\n all_narrow_k && (num_tiles_k < TB_DIM) && (narrow_k_slm <= get_max_dynamic_smem());\n int vec_load_size_i = (num_tiles_k - 1) % 4 + 1;\n // We use the minimum vec_load_size across all tensors.\n // TODO(zhongbo): fix vec_load_size for NVFP4\n // Current unit test won't capture this issue, but in E2E\n // using vec_load_size = 1 other than 1 will lead to mis-aligned\n // address error in MOE training\n vec_load_size = all_nvfp4 ? 1 : std::min(vec_load_size, vec_load_size_i);\n\n const int pos = kernel_args.num_tensors;\n kernel_args.m_list[pos] = m;\n kernel_args.k_list[pos] = k;\n if (!all_nvfp4 || all_has_data) {\n int block_scale_size = all_nvfp4 ? NVFP4_BLOCK_SIZE : MXFP8_BLOCK_SIZE;\n kernel_args.input_list[pos] = const_cast(input[i]->scale_inv.dptr);\n kernel_args.output_list[pos] = output[i]->scale_inv.dptr;\n kernel_args.original_m_list[pos] = input[i]->flat_first_dim();\n kernel_args.original_k_list[pos] = input[i]->flat_last_dim() / block_scale_size;\n } else {\n kernel_args.input_list[pos] = const_cast(input[i]->columnwise_scale_inv.dptr);\n kernel_args.output_list[pos] = output[i]->columnwise_scale_inv.dptr;\n kernel_args.original_m_list[pos] = input[i]->flat_last_dim();\n kernel_args.original_k_list[pos] = input[i]->flat_first_dim() / NVFP4_BLOCK_SIZE;\n }\n kernel_args.num_tensors++;\n }\n // Launch the remaining tensors\n // There is no int3 and misaligned if using int4/int2.\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_swizzle_scaling_factors(\n kernel_args, vec_load_size, true, all_narrow_k, false, stream);\n }\n\n if (columnwise_swizzle) {\n // NVFP4 shouldn't end up here because it only needs rowwise swizzle\n NVTE_CHECK(!all_nvfp4, \"NVFP4 shouldn't end up here because it only needs rowwise swizzle\");\n\n MultiSwizzleArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.block_range[0] = 0;\n int vec_load_size = 4;\n bool all_narrow_m = true;\n for (size_t i = 0; i < num_tensors; i++) {\n //Launch kernel if argument struct is full\n if (kernel_args.num_tensors == kMaxTensorsPerKernel) {\n // There is no int3 and misaligned if using int4/int2.\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_swizzle_scaling_factors(\n kernel_args, vec_load_size, false, false, all_narrow_m, stream);\n // Reset the argument struct and vec_load_size\n kernel_args.num_tensors = 0;\n vec_load_size = 4;\n all_narrow_m = true;\n }\n const int m = input[i]->columnwise_scale_inv.shape[1];\n const int k = input[i]->columnwise_scale_inv.shape[0];\n\n NVTE_CHECK(m % SF_TILE_DIM_M == 0, \"Input should be padded in M/N dimension!\");\n NVTE_CHECK(k % SF_TILE_DIM_K == 0, \"Input should be padded in K dimension!\");\n NVTE_CHECK(k > 0, \"Input scale inverse should be 2D!\");\n NVTE_CHECK(m * k == std::accumulate(output[i]->columnwise_scale_inv.shape.begin(),\n output[i]->columnwise_scale_inv.shape.end(), 1,\n std::multiplies()),\n \"Input.columnwise_scale_inv size is not equal to \"\n \"Output.columnwise_scale_inv size!\");\n\n int num_tiles_m = m / SF_TILE_DIM_M;\n int num_tiles_k = k / SF_TILE_DIM_K;\n const int narrow_m_slm =\n TB_DIM * num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t));\n all_narrow_m =\n all_narrow_m && (num_tiles_m < TB_DIM) && (narrow_m_slm <= get_max_dynamic_smem());\n int vec_load_size_i = (num_tiles_k - 1) % 4 + 1;\n // We use the minimum vec_load_size across all tensors.\n vec_load_size = std::min(vec_load_size, vec_load_size_i);\n\n const int pos = kernel_args.num_tensors;\n kernel_args.input_list[pos] = const_cast(input[i]->columnwise_scale_inv.dptr);\n kernel_args.output_list[pos] = output[i]->columnwise_scale_inv.dptr;\n kernel_args.m_list[pos] = m;\n kernel_args.k_list[pos] = k;\n kernel_args.original_m_list[pos] = input[i]->flat_last_dim();\n kernel_args.original_k_list[pos] = input[i]->flat_first_dim() / MXFP8_BLOCK_SIZE;\n kernel_args.num_tensors++;\n }\n // Launch the remaining tensors\n // There is no int3 and misaligned if using int4/int2.\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_swizzle_scaling_factors(\n kernel_args, vec_load_size, false, false, all_narrow_m, stream);\n }\n}\n\nvoid unswizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t stream) {\n const auto& scaling_mode = output->scaling_mode;\n NVTE_CHECK(scaling_mode == NVTE_MXFP8_1D_SCALING || scaling_mode == NVTE_NVFP4_1D_SCALING,\n \"Output tensor has invalid scaling mode (\", to_string(output->scaling_mode), \").\");\n\n CheckInputTensor(*input, \"scaling_factor_input\");\n CheckInputTensor(*output, \"scaling_factor_output\");\n NVTE_CHECK(input->with_gemm_swizzled_scales, \"Expected input tensor with swizzled scales.\");\n NVTE_CHECK(!output->with_gemm_swizzled_scales,\n \"Expected output tensor in row-major compact format.\");\n NVTE_CHECK(input->scaling_mode == scaling_mode,\n \"Input and output tensors must have matching scaling modes, but got \",\n to_string(input->scaling_mode), \" and \", to_string(output->scaling_mode), \".\");\n\n const bool has_rowwise_scale_inv = output->scale_inv.has_data();\n const bool has_columnwise_scale_inv = output->columnwise_scale_inv.has_data();\n NVTE_CHECK(!has_rowwise_scale_inv || !has_columnwise_scale_inv,\n \"Output tensor has both row-wise and column-wise scaling factors\");\n if (!has_rowwise_scale_inv && !has_columnwise_scale_inv) {\n return;\n }\n if (has_rowwise_scale_inv) {\n NVTE_CHECK(input->scale_inv.has_data(),\n \"Output tensor requests row-wise scaling factors, but input tensor does not \"\n \"provide them.\");\n } else if (has_columnwise_scale_inv) {\n NVTE_CHECK(input->columnwise_scale_inv.has_data(),\n \"Output tensor requests column-wise scaling factors, but input tensor does not \"\n \"provide them.\");\n }\n\n constexpr int SF_TILE_DIM_M = 128;\n constexpr int SF_TILE_DIM_K = 4;\n const dim3 block_size(TB_DIM, TB_DIM);\n\n int m{0}, k{0};\n void* input_ptr{nullptr};\n void* output_ptr{nullptr};\n bool rowwise{false};\n\n switch (scaling_mode) {\n case NVTE_MXFP8_1D_SCALING: {\n NVTE_CHECK(is_fp8_dtype(input->dtype()), \"Input tensor has invalid dtype (expected FP8, got \",\n to_string(input->dtype()), \").\");\n if (has_rowwise_scale_inv) {\n NVTE_CHECK(output->scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", output->scale_inv.shape, \".\");\n m = output->scale_inv.shape[0];\n k = output->scale_inv.shape[1];\n NVTE_CHECK(static_cast(m) * k == input->scale_inv.numel(),\n \"Expected input tensor to have \", static_cast(m) * k,\n \" row-wise scaling factors, but got shape=\", input->scale_inv.shape, \".\");\n NVTE_CHECK(static_cast(m) * k == output->scale_inv.numel(),\n \"Expected output tensor to have \", static_cast(m) * k,\n \" row-wise scaling factors, but got shape=\", output->scale_inv.shape, \".\");\n input_ptr = input->scale_inv.dptr;\n output_ptr = output->scale_inv.dptr;\n rowwise = true;\n } else if (has_columnwise_scale_inv) {\n NVTE_CHECK(output->columnwise_scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", output->columnwise_scale_inv.shape,\n \".\");\n m = output->columnwise_scale_inv.shape[1];\n k = output->columnwise_scale_inv.shape[0];\n NVTE_CHECK(\n static_cast(m) * k == input->columnwise_scale_inv.numel(),\n \"Expected input tensor to have \", static_cast(m) * k,\n \" column-wise scaling factors, but got shape=\", input->columnwise_scale_inv.shape, \".\");\n NVTE_CHECK(static_cast(m) * k == output->columnwise_scale_inv.numel(),\n \"Expected output tensor to have \", static_cast(m) * k,\n \" column-wise scaling factors, but got shape=\",\n output->columnwise_scale_inv.shape, \".\");\n input_ptr = input->columnwise_scale_inv.dptr;\n output_ptr = output->columnwise_scale_inv.dptr;\n rowwise = false;\n }\n break;\n }\n case NVTE_NVFP4_1D_SCALING: {\n NVTE_CHECK(is_fp4_dtype(input->dtype()), \"Input tensor has invalid dtype (expected FP4, got \",\n to_string(input->dtype()), \").\");\n // NVFP4: always unswizzle rowwise regardless of which scale buffer holds the data\n if (has_rowwise_scale_inv) {\n NVTE_CHECK(output->scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", output->scale_inv.shape, \".\");\n m = output->scale_inv.shape[0];\n k = output->scale_inv.shape[1];\n // Example for NVFP4 rowwise path:\n NVTE_CHECK(static_cast(m) * k == input->scale_inv.numel(),\n \"Expected input tensor to have \", static_cast(m) * k,\n \" row-wise scaling factors, but got shape=\", input->scale_inv.shape, \".\");\n NVTE_CHECK(static_cast(m) * k == output->scale_inv.numel(),\n \"Expected output tensor to have \", static_cast(m) * k,\n \" row-wise scaling factors, but got shape=\", output->scale_inv.shape, \".\");\n input_ptr = input->scale_inv.dptr;\n output_ptr = output->scale_inv.dptr;\n } else if (has_columnwise_scale_inv) {\n NVTE_CHECK(output->columnwise_scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", output->columnwise_scale_inv.shape,\n \".\");\n m = output->columnwise_scale_inv.shape[0];\n k = output->columnwise_scale_inv.shape[1];\n NVTE_CHECK(\n static_cast(m) * k == input->columnwise_scale_inv.numel(),\n \"Expected input tensor to have \", static_cast(m) * k,\n \" column-wise scaling factors, but got shape=\", input->columnwise_scale_inv.shape, \".\");\n NVTE_CHECK(static_cast(m) * k == output->columnwise_scale_inv.numel(),\n \"Expected output tensor to have \", static_cast(m) * k,\n \" column-wise scaling factors, but got shape=\",\n output->columnwise_scale_inv.shape, \".\");\n input_ptr = input->columnwise_scale_inv.dptr;\n output_ptr = output->columnwise_scale_inv.dptr;\n }\n rowwise = true;\n break;\n }\n default:\n NVTE_ERROR(\"Invalid scaling mode\");\n }\n\n NVTE_CHECK(m % SF_TILE_DIM_M == 0, \"Output should be padded in M/N dimension!\");\n NVTE_CHECK(k % SF_TILE_DIM_K == 0, \"Output should be padded in K dimension!\");\n\n const int num_tiles_m = m / SF_TILE_DIM_M;\n const int num_tiles_k = k / SF_TILE_DIM_K;\n\n auto launch_unswizzle = [&](int vec_load_size, const dim3& num_blocks, int slm_size) {\n switch (vec_load_size) {\n case 4:\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(unswizzle_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n unswizzle_scaling_kernel\n <<>>(input_ptr, output_ptr, m, k, rowwise);\n break;\n case 2:\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(unswizzle_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n unswizzle_scaling_kernel\n <<>>(input_ptr, output_ptr, m, k, rowwise);\n break;\n case 1:\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(unswizzle_scaling_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n unswizzle_scaling_kernel\n <<>>(input_ptr, output_ptr, m, k, rowwise);\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n };\n\n int vec_load_size = rowwise ? (num_tiles_k - 1) % 4 + 1 : (num_tiles_m - 1) % 4 + 1;\n if (vec_load_size == 3) vec_load_size = 1;\n int n_tiles_in_tb = TB_DIM * vec_load_size;\n dim3 num_blocks = rowwise ? dim3(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m)\n : dim3(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size));\n int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n launch_unswizzle(vec_load_size, num_blocks, slm_size);\n}\n\nvoid multi_tensor_unswizzle_scaling_factors(const std::vector& input,\n std::vector& output, cudaStream_t stream) {\n size_t num_tensors = output.size();\n const auto& first_scaling_mode = output[0]->scaling_mode;\n\n constexpr int SF_TILE_DIM_M = 128;\n constexpr int SF_TILE_DIM_K = 4;\n\n bool all_has_data = true;\n bool all_has_columnwise_data = true;\n bool all_nvfp4 = true;\n for (size_t i = 0; i < num_tensors; i++) {\n const auto scaling_mode = output[i]->scaling_mode;\n const auto is_fp8 = is_fp8_dtype(input[i]->dtype());\n const auto is_fp4 = is_fp4_dtype(input[i]->dtype());\n\n NVTE_CHECK(scaling_mode == first_scaling_mode,\n \"All tensors should have the same scaling mode in multi-tensor unswizzle.\");\n NVTE_CHECK(\n (is_fp8 && is_mxfp8_scaling(scaling_mode)) || (is_fp4 && is_nvfp4_scaling(scaling_mode)),\n \"Not implemented scaling mode \" + to_string(scaling_mode) + \".\");\n NVTE_CHECK(input[i]->with_gemm_swizzled_scales,\n \"Expected input tensors with scales in GEMM swizzled format.\");\n NVTE_CHECK(!output[i]->with_gemm_swizzled_scales,\n \"Expected output tensors with scales in compact format.\");\n NVTE_CHECK(input[i]->numel() != 0, \"Tensor input[\", i, \"] is empty.\");\n CheckInputTensor(*input[i], \"scaling_factor_input[\" + std::to_string(i) + \"]\");\n CheckInputTensor(*output[i], \"scaling_factor_output[\" + std::to_string(i) + \"]\");\n\n all_has_data = all_has_data && output[i]->scale_inv.has_data();\n all_has_columnwise_data =\n (all_has_columnwise_data && output[i]->columnwise_scale_inv.has_data());\n all_nvfp4 = all_nvfp4 && is_nvfp4_scaling(scaling_mode);\n }\n NVTE_CHECK(all_has_data || all_has_columnwise_data,\n \"All tensors should have data or columnwise data.\");\n NVTE_CHECK(!all_has_data || !all_has_columnwise_data,\n \"All tensors have both data and columnwise data.\");\n\n const bool rowwise_unswizzle = all_has_data || all_nvfp4;\n const bool columnwise_unswizzle = all_has_columnwise_data && !all_nvfp4;\n\n if (rowwise_unswizzle) {\n MultiSwizzleArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.block_range[0] = 0;\n int vec_load_size = 4;\n for (size_t i = 0; i < num_tensors; i++) {\n if (kernel_args.num_tensors == kMaxTensorsPerKernel) {\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_unswizzle_scaling_factors(\n kernel_args, vec_load_size, true, stream);\n kernel_args.num_tensors = 0;\n vec_load_size = 4;\n }\n int m, k;\n if (all_has_data) {\n NVTE_CHECK(input[i]->scale_inv.has_data(), \"Input tensor \", i,\n \" does not have row-wise scaling factors.\");\n NVTE_CHECK(output[i]->scale_inv.shape.size() == 2, \"Expected output tensor \", i,\n \" to have \", \"2D scaling factors, got shape=\", output[i]->scale_inv.shape, \".\");\n m = output[i]->scale_inv.shape[0];\n k = output[i]->scale_inv.shape[1];\n NVTE_CHECK(m * k == input[i]->scale_inv.numel(), \"Expected input tensor \", i, \" to have \",\n m * k, \" row-wise scaling factors, but got shape=\", input[i]->scale_inv.shape,\n \".\");\n }\n\n if (all_has_columnwise_data) {\n NVTE_CHECK(all_nvfp4,\n \"When doing rowwise unswizzle with columnwise data, it has to be NVFP4\");\n NVTE_CHECK(input[i]->columnwise_scale_inv.has_data(), \"Input tensor \", i,\n \" does not have column-wise scaling factors.\");\n NVTE_CHECK(output[i]->columnwise_scale_inv.shape.size() == 2, \"Expected output tensor \", i,\n \" to have \",\n \"2D scaling factors, got shape=\", output[i]->columnwise_scale_inv.shape, \".\");\n m = output[i]->columnwise_scale_inv.shape[0];\n k = output[i]->columnwise_scale_inv.shape[1];\n NVTE_CHECK(m * k == input[i]->columnwise_scale_inv.numel(), \"Expected input tensor \", i,\n \" to have \", m * k, \" column-wise scaling factors, but got shape=\",\n input[i]->columnwise_scale_inv.shape, \".\");\n }\n\n NVTE_CHECK(m % SF_TILE_DIM_M == 0, \"Output should be padded in M/N dimension!\");\n NVTE_CHECK(k % SF_TILE_DIM_K == 0, \"Output should be padded in K dimension!\");\n NVTE_CHECK(k > 0, \"Output scale inverse should be 2D!\");\n\n int num_tiles_k = k / SF_TILE_DIM_K;\n int vec_load_size_i = (num_tiles_k - 1) % 4 + 1;\n vec_load_size = all_nvfp4 ? 1 : std::min(vec_load_size, vec_load_size_i);\n\n const int pos = kernel_args.num_tensors;\n kernel_args.m_list[pos] = m;\n kernel_args.k_list[pos] = k;\n if (!all_nvfp4 || all_has_data) {\n kernel_args.input_list[pos] = const_cast(input[i]->scale_inv.dptr);\n kernel_args.output_list[pos] = output[i]->scale_inv.dptr;\n } else {\n kernel_args.input_list[pos] = const_cast(input[i]->columnwise_scale_inv.dptr);\n kernel_args.output_list[pos] = output[i]->columnwise_scale_inv.dptr;\n }\n kernel_args.num_tensors++;\n }\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_unswizzle_scaling_factors(\n kernel_args, vec_load_size, true, stream);\n }\n\n if (columnwise_unswizzle) {\n NVTE_CHECK(!all_nvfp4, \"NVFP4 shouldn't end up here because it only needs rowwise unswizzle\");\n\n MultiSwizzleArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.block_range[0] = 0;\n int vec_load_size = 4;\n for (size_t i = 0; i < num_tensors; i++) {\n if (kernel_args.num_tensors == kMaxTensorsPerKernel) {\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_unswizzle_scaling_factors(\n kernel_args, vec_load_size, false, stream);\n kernel_args.num_tensors = 0;\n vec_load_size = 4;\n }\n NVTE_CHECK(output[i]->columnwise_scale_inv.shape.size() == 2, \"Expected output tensor \", i,\n \" to have \",\n \"2D scaling factors, got shape=\", output[i]->columnwise_scale_inv.shape, \".\");\n const int m = output[i]->columnwise_scale_inv.shape[1];\n const int k = output[i]->columnwise_scale_inv.shape[0];\n\n NVTE_CHECK(m % SF_TILE_DIM_M == 0, \"Output should be padded in M/N dimension!\");\n NVTE_CHECK(k % SF_TILE_DIM_K == 0, \"Output should be padded in K dimension!\");\n NVTE_CHECK(k > 0, \"Output scale inverse should be 2D!\");\n NVTE_CHECK(m * k == std::accumulate(input[i]->columnwise_scale_inv.shape.begin(),\n input[i]->columnwise_scale_inv.shape.end(), 1,\n std::multiplies()),\n \"Input.columnwise_scale_inv size is not equal to \"\n \"Output.columnwise_scale_inv size!\");\n\n int num_tiles_k = k / SF_TILE_DIM_K;\n int vec_load_size_i = (num_tiles_k - 1) % 4 + 1;\n vec_load_size = std::min(vec_load_size, vec_load_size_i);\n\n const int pos = kernel_args.num_tensors;\n kernel_args.input_list[pos] = const_cast(input[i]->columnwise_scale_inv.dptr);\n kernel_args.output_list[pos] = output[i]->columnwise_scale_inv.dptr;\n kernel_args.m_list[pos] = m;\n kernel_args.k_list[pos] = k;\n kernel_args.num_tensors++;\n }\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_unswizzle_scaling_factors(\n kernel_args, vec_load_size, false, stream);\n }\n}\n} // namespace transformer_engine\n\n/*\n* WIP (Phuong):\n* - Opt for bank conflicts\n* - Adding swizzle for 2d-block scaling.\n*/\nvoid nvte_swizzle_scaling_factors(const NVTETensor input, NVTETensor output, cudaStream_t stream) {\n NVTE_API_CALL(nvte_swizzle_scaling_factors);\n using namespace transformer_engine;\n swizzle_scaling_factors(convertNVTETensorCheck(input), convertNVTETensorCheck(output), stream);\n}\n\nvoid nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETensor* outputs,\n const size_t num_tensors, cudaStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_swizzle_scaling_factors);\n using namespace transformer_engine;\n NVTE_CHECK(num_tensors > 0, \"Number of tensors should be greater than 0.\");\n std::vector input_list, output_list;\n for (size_t i = 0; i < num_tensors; i++) {\n input_list.push_back(convertNVTETensorCheck(inputs[i]));\n output_list.push_back(convertNVTETensorCheck(outputs[i]));\n }\n multi_tensor_swizzle_scaling_factors(input_list, output_list, stream,\n /*check_scale_inv_shapes=*/true);\n}\n\nvoid nvte_multi_tensor_swizzle_scaling_factors_unchecked(const NVTETensor* inputs,\n NVTETensor* outputs,\n const size_t num_tensors,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_swizzle_scaling_factors_unchecked);\n using namespace transformer_engine;\n NVTE_CHECK(num_tensors > 0, \"Number of tensors should be greater than 0.\");\n std::vector input_list, output_list;\n for (size_t i = 0; i < num_tensors; i++) {\n input_list.push_back(convertNVTETensorCheck(inputs[i]));\n output_list.push_back(convertNVTETensorCheck(outputs[i]));\n }\n multi_tensor_swizzle_scaling_factors(input_list, output_list, stream,\n /*check_scale_inv_shapes=*/false);\n}\n\nvoid nvte_unswizzle_scaling_factors(const NVTETensor input, NVTETensor output,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_unswizzle_scaling_factors);\n using namespace transformer_engine;\n unswizzle_scaling_factors(convertNVTETensorCheck(input), convertNVTETensorCheck(output), stream);\n}\n\nvoid nvte_multi_tensor_unswizzle_scaling_factors(const NVTETensor* inputs, NVTETensor* outputs,\n const size_t num_tensors, cudaStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_unswizzle_scaling_factors);\n using namespace transformer_engine;\n NVTE_CHECK(num_tensors > 0, \"Number of tensors should be greater than 0.\");\n std::vector input_list, output_list;\n for (size_t i = 0; i < num_tensors; i++) {\n input_list.push_back(convertNVTETensorCheck(inputs[i]));\n output_list.push_back(convertNVTETensorCheck(outputs[i]));\n }\n multi_tensor_unswizzle_scaling_factors(input_list, output_list, stream);\n}\n\nnamespace transformer_engine {\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n grouped_swizzle_scaling_variable_shape_kernel(const void* input, void* output,\n const int64_t* m_array, const int64_t* k_array,\n int num_tensors, bool rowwise,\n size_t scale_elem_size, size_t common_m,\n size_t common_k) {\n extern __shared__ int s_metadata[];\n int* s_total_blocks = &s_metadata[0];\n\n // Warp reduction to compute total workload\n if (threadIdx.x < 32 && threadIdx.y == 0) {\n int local_blocks = 0;\n for (int i = threadIdx.x; i < num_tensors; i += 32) {\n size_t m = rowwise ? (m_array ? m_array[i] : common_m) : (k_array ? k_array[i] : common_k);\n size_t k = rowwise ? (k_array ? k_array[i] : common_k) : (m_array ? m_array[i] : common_m);\n\n size_t padded_m = round_up_to_multiple(m, 128);\n size_t padded_k = round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4);\n\n int num_tiles_m = padded_m / SF_TILE_DIM_M;\n int num_tiles_k = padded_k / SF_TILE_DIM_K;\n\n int vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1));\n if (vec_load_size == 3) vec_load_size = 1;\n int n_tiles_in_tb = TB_DIM * vec_load_size;\n\n int grid_dim_x = rowwise ? DIVUP(num_tiles_k, n_tiles_in_tb) : DIVUP(num_tiles_k, TB_DIM);\n int grid_dim_y = rowwise ? num_tiles_m : DIVUP(num_tiles_m, vec_load_size);\n local_blocks += grid_dim_x * grid_dim_y;\n }\n\n for (int offset = 16; offset > 0; offset /= 2) {\n local_blocks += __shfl_down_sync(0xffffffff, local_blocks, offset);\n }\n if (threadIdx.x == 0) *s_total_blocks = local_blocks;\n }\n __syncthreads();\n\n const int total_blocks = *s_total_blocks;\n\n // Persistent-grid loop\n for (int linear_block_id = blockIdx.x; linear_block_id < total_blocks;\n linear_block_id += gridDim.x) {\n // Discover tensor_id and local_block_id via linear scan\n int tensor_id = 0;\n int current_block_base = 0;\n size_t current_scale_base = 0;\n int grid_dim_x = 0;\n int grid_dim_y = 0;\n size_t M = 0, K = 0;\n int vec_load_size = 0;\n\n for (int i = 0; i < num_tensors; ++i) {\n M = rowwise ? (m_array ? m_array[i] : common_m) : (k_array ? k_array[i] : common_k);\n K = rowwise ? (k_array ? k_array[i] : common_k) : (m_array ? m_array[i] : common_m);\n\n size_t padded_m = round_up_to_multiple(M, 128);\n size_t padded_k = round_up_to_multiple(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4);\n\n int num_tiles_m = padded_m / SF_TILE_DIM_M;\n int num_tiles_k = padded_k / SF_TILE_DIM_K;\n\n vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1));\n if (vec_load_size == 3) vec_load_size = 1;\n int n_tiles_in_tb = TB_DIM * vec_load_size;\n\n grid_dim_x = rowwise ? DIVUP(num_tiles_k, n_tiles_in_tb) : DIVUP(num_tiles_k, TB_DIM);\n grid_dim_y = rowwise ? num_tiles_m : DIVUP(num_tiles_m, vec_load_size);\n int blocks_i = grid_dim_x * grid_dim_y;\n\n if (linear_block_id < current_block_base + blocks_i) {\n tensor_id = i;\n break;\n }\n current_block_base += blocks_i;\n current_scale_base += padded_m * padded_k * scale_elem_size;\n }\n\n int local_block_id = linear_block_id - current_block_base;\n int block_x = local_block_id % grid_dim_x;\n int block_y = local_block_id / grid_dim_x;\n\n const uint8_t* input_base = reinterpret_cast(input) + current_scale_base;\n uint8_t* output_base = reinterpret_cast(output) + current_scale_base;\n\n const int padded_m = static_cast(round_up_to_multiple(M, 128));\n const int padded_k =\n static_cast(round_up_to_multiple(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4));\n const int original_M = static_cast(M);\n const int original_K = static_cast(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)));\n const bool padding_m = (block_y == grid_dim_y - 1) && (original_M < padded_m);\n const bool padding_k = (block_x == grid_dim_x - 1) && (original_K < padded_k);\n\n if (rowwise) {\n if (vec_load_size == 4) {\n dispatch_swizzle_row_scaling_kernel_impl(\n input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y,\n grid_dim_x, grid_dim_y, padding_k, padding_m);\n } else if (vec_load_size == 2) {\n dispatch_swizzle_row_scaling_kernel_impl(\n input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y,\n grid_dim_x, grid_dim_y, padding_k, padding_m);\n } else {\n dispatch_swizzle_row_scaling_kernel_impl(\n input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y,\n grid_dim_x, grid_dim_y, padding_k, padding_m);\n }\n } else {\n if (vec_load_size == 4) {\n dispatch_swizzle_col_scaling_kernel_impl(\n input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y,\n grid_dim_x, grid_dim_y, padding_k, padding_m);\n } else if (vec_load_size == 2) {\n dispatch_swizzle_col_scaling_kernel_impl(\n input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y,\n grid_dim_x, grid_dim_y, padding_k, padding_m);\n } else {\n dispatch_swizzle_col_scaling_kernel_impl(\n input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y,\n grid_dim_x, grid_dim_y, padding_k, padding_m);\n }\n }\n }\n}\n\ntemplate \nint grouped_swizzle_variable_max_active_blocks_per_sm(int device_id) {\n static std::vector cache(cuda::num_devices(), -1);\n static std::vector flags(cuda::num_devices());\n NVTE_CHECK(0 <= device_id && device_id < cuda::num_devices(), \"invalid CUDA device ID\");\n\n auto init = [&]() {\n constexpr int metadata_shmem = sizeof(int); // s_total_blocks\n constexpr int dynamic_smem_size =\n TB_DIM * 4 * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t) + metadata_shmem;\n int max_active_blocks_per_sm;\n NVTE_CHECK_CUDA(cudaOccupancyMaxActiveBlocksPerMultiprocessor(\n &max_active_blocks_per_sm,\n grouped_swizzle_scaling_variable_shape_kernel,\n TB_DIM * TB_DIM, dynamic_smem_size));\n NVTE_CHECK(max_active_blocks_per_sm > 0, \"Occupancy query returned 0 blocks per SM.\");\n cache[device_id] = max_active_blocks_per_sm;\n };\n std::call_once(flags[device_id], init);\n return cache[device_id];\n}\n\nvoid swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* output,\n cudaStream_t stream) {\n // Check scaling mode\n NVTE_CHECK(input->scaling_mode == NVTE_MXFP8_1D_SCALING,\n \"Grouped swizzle supports only MXFP8 scaling.\");\n\n // Check tensors\n CheckInputGroupedTensor(*input, \"input\");\n CheckOutputGroupedTensor(*output, \"output\", false);\n NVTE_CHECK(!input->with_gemm_swizzled_scales,\n \"Expected input grouped tensor with scales in compact format.\");\n NVTE_CHECK(output->with_gemm_swizzled_scales,\n \"Expected output grouped tensor with scales in GEMM swizzled format.\");\n\n // Check scaling factors availability\n const bool has_rowwise_scale_inv = input->scale_inv.has_data();\n const bool has_columnwise_scale_inv = input->columnwise_scale_inv.has_data();\n if (!has_rowwise_scale_inv && !has_columnwise_scale_inv) {\n return;\n }\n\n const int64_t* m_array = reinterpret_cast(input->first_dims.dptr);\n const int64_t* k_array = reinterpret_cast(input->last_dims.dptr);\n const bool is_variable_shape = !input->all_same_shape();\n\n if (!is_variable_shape) {\n // Fallback to uniform shape implementation\n // Assumption is that all the tensors share the same shapes and are contgiuous.\n // And so we dont need to pass array of input/output pointers(due to conttiguity)\n // as well as array of shapes(due to uniform shapes).\n const size_t first_dim = input->get_common_first_dim();\n const size_t last_dim = input->get_common_last_dim();\n\n constexpr int SF_TILE_DIM_M = 128;\n constexpr int SF_TILE_DIM_K = 4;\n const dim3 block_size(TB_DIM, TB_DIM);\n\n auto launch_grouped_swizzle = [&](bool rowwise) {\n const size_t m = rowwise ? first_dim : last_dim;\n const size_t k = rowwise ? last_dim : first_dim;\n const size_t padded_m = round_up_to_multiple(m, 128);\n const size_t padded_k =\n round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4);\n // Per-tensor scale-element counts:\n // - \"padded\" layout: each tensor occupies padded_m * padded_k elements\n // (total buffer = num_tensors * padded_m * padded_k).\n // - \"compact\" layout (what the grouped MXFP8 quantize kernel actually writes):\n // per-tensor stride is m * padded_k (rowwise) or DIVUP(k,32) * padded_m\n // (columnwise) and the total buffer the C++ allocator hands out has its\n // grouped first dim padded up to a multiple of 128 (rowwise) or 4\n // (columnwise) -- so the buffer may be slightly larger than\n // num_tensors * compact_scale_elems, with trailing alignment slack at\n // the very end (never read because of the per-tensor row/k guard in the\n // kernel impl).\n // The output is always written in the padded layout. The input may be in\n // either layout; the kernel handles the compact case safely by using\n // different per-tensor strides for input vs output and skipping loads past\n // the per-tensor extent.\n const size_t padded_scale_elems = padded_m * padded_k;\n const size_t compact_scale_elems =\n rowwise ? m * padded_k : DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)) * padded_m;\n const size_t compact_total_scale_elems =\n rowwise ? round_up_to_multiple(input->num_tensors * m, 128) * padded_k\n : round_up_to_multiple(\n input->num_tensors * DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4) *\n padded_m;\n\n const size_t scale_elem_size = rowwise ? typeToSize(input->scale_inv.dtype)\n : typeToSize(input->columnwise_scale_inv.dtype);\n\n const size_t input_scale_numel =\n rowwise ? input->scale_inv.numel() : input->columnwise_scale_inv.numel();\n const size_t output_scale_numel =\n rowwise ? output->scale_inv.numel() : output->columnwise_scale_inv.numel();\n\n bool input_is_compact;\n if (input_scale_numel == input->num_tensors * padded_scale_elems) {\n input_is_compact = false;\n } else if (input_scale_numel == compact_total_scale_elems) {\n input_is_compact = true;\n } else {\n NVTE_ERROR(\"Grouped input \", (rowwise ? \"scale_inv\" : \"columnwise_scale_inv\"),\n \" size does not match expected packed size (got \", input_scale_numel,\n \", expected either \", input->num_tensors * padded_scale_elems,\n \" (per-tensor padded) or \", compact_total_scale_elems, \" (compact)).\");\n }\n NVTE_CHECK(output_scale_numel == input->num_tensors * padded_scale_elems, \"Grouped output \",\n (rowwise ? \"scale_inv\" : \"columnwise_scale_inv\"),\n \" size does not match expected per-tensor padded size.\");\n\n const size_t input_stride_bytes =\n (input_is_compact ? compact_scale_elems : padded_scale_elems) * scale_elem_size;\n const size_t output_stride_bytes = padded_scale_elems * scale_elem_size;\n\n const int num_tiles_m = padded_m / SF_TILE_DIM_M;\n const int num_tiles_k = padded_k / SF_TILE_DIM_K;\n int vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1));\n if (vec_load_size == 3) vec_load_size = 1;\n const int n_tiles_in_tb = TB_DIM * vec_load_size;\n\n dim3 num_blocks;\n if (rowwise) {\n num_blocks = dim3(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m, input->num_tensors);\n } else {\n num_blocks =\n dim3(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size), input->num_tensors);\n }\n const int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n\n const int original_M = static_cast(rowwise ? first_dim : last_dim);\n const int original_K = static_cast(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)));\n const void* input_ptr = rowwise ? input->scale_inv.dptr : input->columnwise_scale_inv.dptr;\n void* output_ptr = rowwise ? output->scale_inv.dptr : output->columnwise_scale_inv.dptr;\n\n if (rowwise) {\n TRANSFORMER_ENGINE_VECTORIZED_LOAD_INTEGER_TYPE_SWITCH(vec_load_size, LType, {\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n grouped_swizzle_row_scaling_uniform_shape_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n grouped_swizzle_row_scaling_uniform_shape_kernel\n <<>>(\n input_ptr, output_ptr, padded_m, padded_k, original_M, original_K,\n input_stride_bytes, output_stride_bytes);\n });\n } else {\n TRANSFORMER_ENGINE_VECTORIZED_LOAD_INTEGER_TYPE_SWITCH(vec_load_size, LType, {\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n grouped_swizzle_col_scaling_uniform_shape_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n grouped_swizzle_col_scaling_uniform_shape_kernel\n <<>>(\n input_ptr, output_ptr, padded_m, padded_k, original_M, original_K,\n input_stride_bytes, output_stride_bytes);\n });\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n };\n\n if (has_rowwise_scale_inv) {\n launch_grouped_swizzle(true);\n }\n if (has_columnwise_scale_inv) {\n launch_grouped_swizzle(false);\n }\n } else {\n // Variable shape implementation using Device-Side Block Scheduler\n size_t num_tensors = input->num_tensors;\n\n constexpr int SF_TILE_DIM_M = 128;\n constexpr int SF_TILE_DIM_K = 4;\n const dim3 block_size(TB_DIM, TB_DIM);\n const int max_slm_size = TB_DIM * 4 * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n const int metadata_shmem = sizeof(int); // s_total_blocks\n const int dynamic_smem_size = max_slm_size + metadata_shmem;\n\n size_t common_m = input->all_same_first_dim() ? input->get_common_first_dim() : 0;\n size_t common_k = input->all_same_last_dim() ? input->get_common_last_dim() : 0;\n\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n grouped_swizzle_scaling_variable_shape_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, dynamic_smem_size));\n\n const int device_id = cuda::current_device();\n const int num_SMs = cuda::sm_count(device_id);\n const int max_active_blocks_per_sm =\n grouped_swizzle_variable_max_active_blocks_per_sm(device_id);\n const int persistent_blocks = num_SMs * max_active_blocks_per_sm;\n const dim3 num_blocks(persistent_blocks);\n\n auto launch_grouped_swizzle_variable = [&](bool rowwise) {\n const size_t scale_elem_size = rowwise ? typeToSize(input->scale_inv.dtype)\n : typeToSize(input->columnwise_scale_inv.dtype);\n\n const void* input_ptr = rowwise ? input->scale_inv.dptr : input->columnwise_scale_inv.dptr;\n void* output_ptr = rowwise ? output->scale_inv.dptr : output->columnwise_scale_inv.dptr;\n\n grouped_swizzle_scaling_variable_shape_kernel\n <<>>(\n input_ptr, output_ptr, m_array, k_array, num_tensors, rowwise, scale_elem_size,\n common_m, common_k);\n\n NVTE_CHECK_CUDA(cudaGetLastError());\n };\n\n if (has_rowwise_scale_inv) {\n launch_grouped_swizzle_variable(true);\n }\n if (has_columnwise_scale_inv) {\n launch_grouped_swizzle_variable(false);\n }\n }\n}\n\nvoid unswizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* output,\n cudaStream_t stream) {\n NVTE_CHECK(output->scaling_mode == NVTE_MXFP8_1D_SCALING,\n \"Grouped unswizzle supports only MXFP8 scaling.\");\n\n CheckInputGroupedTensor(*input, \"input\");\n CheckOutputGroupedTensor(*output, \"output\", false);\n NVTE_CHECK(input->with_gemm_swizzled_scales,\n \"Expected input grouped tensor with scales in GEMM swizzled format.\");\n NVTE_CHECK(!output->with_gemm_swizzled_scales,\n \"Expected output grouped tensor with scales in compact format.\");\n NVTE_CHECK(input->scaling_mode == output->scaling_mode,\n \"Input and output grouped tensors must have matching scaling modes.\");\n NVTE_CHECK(input->num_tensors == output->num_tensors,\n \"Input and output grouped tensors must have the same number of tensors.\");\n\n const bool has_rowwise_scale_inv = output->scale_inv.has_data();\n const bool has_columnwise_scale_inv = output->columnwise_scale_inv.has_data();\n if (!has_rowwise_scale_inv && !has_columnwise_scale_inv) {\n return;\n }\n\n NVTE_CHECK(input->all_same_shape() && output->all_same_shape(),\n \"Grouped unswizzle requires uniform tensor shapes.\");\n\n const size_t first_dim = output->get_common_first_dim();\n const size_t last_dim = output->get_common_last_dim();\n\n constexpr int SF_TILE_DIM_M = 128;\n constexpr int SF_TILE_DIM_K = 4;\n const dim3 block_size(TB_DIM, TB_DIM);\n\n auto launch_grouped_unswizzle = [&](bool rowwise) {\n const size_t m = rowwise ? first_dim : last_dim;\n const size_t k = rowwise ? last_dim : first_dim;\n const size_t padded_m = round_up_to_multiple(m, 128);\n const size_t padded_k =\n round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4);\n const size_t scale_elems = padded_m * padded_k;\n\n const size_t scale_elem_size = rowwise ? typeToSize(output->scale_inv.dtype)\n : typeToSize(output->columnwise_scale_inv.dtype);\n const size_t scale_stride_bytes = scale_elems * scale_elem_size;\n\n if (rowwise) {\n NVTE_CHECK(input->scale_inv.numel() == input->num_tensors * scale_elems,\n \"Grouped input scale_inv size does not match expected packed size.\");\n NVTE_CHECK(output->scale_inv.numel() == output->num_tensors * scale_elems,\n \"Grouped output scale_inv size does not match expected packed size.\");\n } else {\n NVTE_CHECK(input->columnwise_scale_inv.numel() == input->num_tensors * scale_elems,\n \"Grouped input columnwise_scale_inv size does not match expected packed size.\");\n NVTE_CHECK(output->columnwise_scale_inv.numel() == output->num_tensors * scale_elems,\n \"Grouped output columnwise_scale_inv size does not match expected packed size.\");\n }\n\n const int num_tiles_m = padded_m / SF_TILE_DIM_M;\n const int num_tiles_k = padded_k / SF_TILE_DIM_K;\n int vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1));\n if (vec_load_size == 3) vec_load_size = 1;\n const int n_tiles_in_tb = TB_DIM * vec_load_size;\n\n dim3 num_blocks;\n if (rowwise) {\n num_blocks = dim3(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m, output->num_tensors);\n } else {\n num_blocks =\n dim3(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size), output->num_tensors);\n }\n const int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n\n const void* input_ptr = rowwise ? input->scale_inv.dptr : input->columnwise_scale_inv.dptr;\n void* output_ptr = rowwise ? output->scale_inv.dptr : output->columnwise_scale_inv.dptr;\n\n using kernel_t = void (*)(const void*, void*, const int, const int, const size_t, const bool);\n kernel_t kernel_fn = nullptr;\n switch (vec_load_size) {\n case 4:\n kernel_fn =\n grouped_unswizzle_scaling_uniform_shape_kernel;\n break;\n case 2:\n kernel_fn =\n grouped_unswizzle_scaling_uniform_shape_kernel;\n break;\n case 1:\n kernel_fn =\n grouped_unswizzle_scaling_uniform_shape_kernel;\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n }\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n kernel_fn<<>>(input_ptr, output_ptr, padded_m,\n padded_k, scale_stride_bytes, rowwise);\n NVTE_CHECK_CUDA(cudaGetLastError());\n };\n\n if (has_rowwise_scale_inv) {\n launch_grouped_unswizzle(true);\n }\n if (has_columnwise_scale_inv) {\n launch_grouped_unswizzle(false);\n }\n}\n\n} // namespace transformer_engine\n\nvoid nvte_swizzle_grouped_scaling_factors(const NVTEGroupedTensor input, NVTEGroupedTensor output,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_swizzle_grouped_scaling_factors);\n using namespace transformer_engine;\n swizzle_grouped_scaling_factors(convertNVTEGroupedTensorCheck(input),\n convertNVTEGroupedTensorCheck(output), stream);\n}\n\nvoid nvte_unswizzle_grouped_scaling_factors(const NVTEGroupedTensor input, NVTEGroupedTensor output,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_unswizzle_grouped_scaling_factors);\n using namespace transformer_engine;\n unswizzle_grouped_scaling_factors(convertNVTEGroupedTensorCheck(input),\n convertNVTEGroupedTensorCheck(output), stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/hip/hip_runtime.h\"\n#include \"../util/logging.h\"\n#include \"transformer_engine/transformer_engine.h\"\n\nnamespace transformer_engine {\nnamespace {\n\nconstexpr int MXFP8_BLOCK_SIZE = 32;\nconstexpr int NVFP4_BLOCK_SIZE = 16;\n\nint get_max_dynamic_smem() {\n static int max_smem = -1;\n if (max_smem < 0) {\n int device;\n NVTE_CHECK_CUDA(hipGetDevice(&device));\n NVTE_CHECK_CUDA(\n hipDeviceGetAttribute(&max_smem, hipDeviceAttributeSharedMemPerBlockOptin, device));\n }\n return max_smem;\n}\n\nconstexpr __device__ __host__ int TB_DIM = 32;\nconstexpr __device__ __host__ int NEW_SF_TILE_DIM_K = 16;\nconstexpr __device__ __host__ int N_SF_PER_TD_PER_TILE = 4;\n\n// output is in ~K-major interleaved blocks\nconstexpr __device__ __host__ int NEW_SF_TILE_DIM_K_I32 = NEW_SF_TILE_DIM_K / 4;\nconstexpr __device__ __host__ int NEW_SF_TILE_DIM_M_I32 = 32;\n\ntemplate \n__device__ inline void regs_shuffle_with_bit_shifts(LType* regs_vec) {\n // inp, 4-byte chunks [0,1,2,3, 4,5,6,7, 8,9,10,11, 12,13,14,15]\n // out, swapping byte to form new 4-byte chunks [0,4,8,12, 1,5,9,13, 2,6,10,14, 3,7,11,15]\n\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int kVectorSize = N_SF_PER_TD_PER_TILE * N_TILE_PER_TD;\n int32_t new_regs[kVectorSize];\n int32_t* regs = reinterpret_cast(regs_vec);\n\n#pragma unroll\n for (int i = 0; i < N_TILE_PER_TD; i++) {\n#pragma unroll\n for (int j = 0; j < N_SF_PER_TD_PER_TILE; j++) {\n new_regs[i * N_SF_PER_TD_PER_TILE + j] =\n (((regs[i + 0 * N_TILE_PER_TD] >> 8 * j) & 0xFF)) |\n (((regs[i + 1 * N_TILE_PER_TD] >> 8 * j) & 0xFF) << 8) |\n (((regs[i + 2 * N_TILE_PER_TD] >> 8 * j) & 0xFF) << 16) |\n (((regs[i + 3 * N_TILE_PER_TD] >> 8 * j) & 0xFF) << 24);\n }\n }\n#pragma unroll\n for (int i = 0; i < kVectorSize; i++) regs[i] = new_regs[i];\n}\n\ntemplate \n__device__ inline void regs_unshuffle_with_bit_shifts(LType* regs_vec) {\n // Inverse of regs_shuffle_with_bit_shifts\n // inp, 4-byte chunks [0,4,8,12, 1,5,9,13, 2,6,10,14, 3,7,11,15]\n // out, swapping byte to form new 4-byte chunks [0,1,2,3, 4,5,6,7, 8,9,10,11, 12,13,14,15]\n\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int kVectorSize = N_SF_PER_TD_PER_TILE * N_TILE_PER_TD;\n int32_t new_regs[kVectorSize];\n int32_t* regs = reinterpret_cast(regs_vec);\n\n#pragma unroll\n for (int i = 0; i < N_TILE_PER_TD; i++) {\n#pragma unroll\n for (int j = 0; j < N_SF_PER_TD_PER_TILE; j++) {\n new_regs[i + j * N_TILE_PER_TD] =\n ((regs[i * N_SF_PER_TD_PER_TILE + 0] >> 8 * j) & 0xFF) |\n (((regs[i * N_SF_PER_TD_PER_TILE + 1] >> 8 * j) & 0xFF) << 8) |\n (((regs[i * N_SF_PER_TD_PER_TILE + 2] >> 8 * j) & 0xFF) << 16) |\n (((regs[i * N_SF_PER_TD_PER_TILE + 3] >> 8 * j) & 0xFF) << 24);\n }\n }\n#pragma unroll\n for (int i = 0; i < kVectorSize; i++) regs[i] = new_regs[i];\n}\n\n// IS_PADDED_K / IS_PADDED_M select the boundary-block specialization at compile\n// time so the inner load loop avoids the per-iteration runtime checks. The\n// caller computes the runtime predicates from blockIdx/gridDim once per block\n// (uniform across the block) and dispatches to the right specialization.\ntemplate \n__device__ void swizzle_col_scaling_kernel_impl(const void* input, void* output, const int M,\n const int K, const int original_M,\n const int original_K, const int bid_x,\n const int bid_y, const int grid_dim_x,\n const int grid_dim_y) {\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int N_SF_PER_TD = N_TILE_PER_TD * N_SF_PER_TD_PER_TILE;\n constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4;\n\n // input is in M-major\n constexpr int SF_TILE_DIM_M_I32 = SF_TILE_DIM_M / 4;\n constexpr int SF_TILE_DIM_K_I32 = SF_TILE_DIM_K;\n\n const int M_i32 = M / 4;\n const int K_i32 = K;\n\n int m_tiles_in_tb = N_TILE_PER_TD;\n int k_tiles_in_tb = TB_DIM;\n if (bid_x == grid_dim_x - 1) {\n k_tiles_in_tb = (K_i32 / SF_TILE_DIM_K_I32 - 1) % k_tiles_in_tb + 1;\n }\n if (bid_y == grid_dim_y - 1) {\n m_tiles_in_tb = (M_i32 / SF_TILE_DIM_M_I32 - 1) % m_tiles_in_tb + 1;\n }\n\n const int input_offset =\n bid_x * TB_DIM * SF_TILE_DIM_K_I32 * M_i32 + bid_y * N_TILE_PER_TD * SF_TILE_DIM_M_I32;\n const int32_t* input_i32 = reinterpret_cast(input) + input_offset;\n int32_t* output_i32[N_TILE_PER_TD];\n#pragma unroll\n for (int i = 0; i < m_tiles_in_tb; i++) {\n output_i32[i] = reinterpret_cast(output) + bid_x * TB_DIM * SF_TILE_SIZE_I32 +\n (bid_y * N_TILE_PER_TD + i) * SF_TILE_DIM_M_I32 * K_i32;\n }\n extern __shared__ int slm[];\n\n // load, global -> regs\n // Each register read for a given i is along the M direction at K-coord\n // (bid_x * TB_DIM * SF_TILE_DIM_K + threadIdx.y * SF_TILE_DIM_K + i). When that\n // K-coord is past original_K, the entire register is out of the per-tensor data\n // region (which may be the unpadded compact extent), so we must NOT issue the\n // __ldg there -- it could read past the per-tensor buffer (and, for the last\n // tensor in a grouped allocation, past the end of the allocation entirely).\n LType regs_vec[N_SF_PER_TD_PER_TILE];\n if (threadIdx.x * N_TILE_PER_TD < m_tiles_in_tb * SF_TILE_DIM_M_I32 &&\n threadIdx.y < k_tiles_in_tb) {\n const int k_base = bid_x * TB_DIM * SF_TILE_DIM_K + threadIdx.y * SF_TILE_DIM_K;\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) {\n const int thread_offset =\n (threadIdx.y * SF_TILE_DIM_K_I32 + i) * M_i32 + threadIdx.x * N_TILE_PER_TD;\n const int k_coord = k_base + i;\n if constexpr (IS_PADDED_K) {\n if (k_coord >= original_K) {\n // Entire register is past original_K: zero directly without loading.\n uint8_t* zero_bytes = reinterpret_cast(regs_vec + i);\n#pragma unroll\n for (int j = 0; j < static_cast(sizeof(LType)); j++) zero_bytes[j] = 0;\n continue;\n }\n }\n regs_vec[i] = __ldg(reinterpret_cast(input_i32 + thread_offset));\n // Per-byte M masking is still needed when only part of the register is past\n // original_M (i.e. K-coord is in range but the M position spans the boundary).\n if constexpr (IS_PADDED_M) {\n for (int j = 0; j < N_TILE_PER_TD * sizeof(int); j++) {\n const int index = (input_offset + thread_offset) * sizeof(int) + j;\n if (index % M >= original_M) {\n reinterpret_cast(regs_vec + i)[j] = 0;\n }\n }\n }\n }\n\n // local shuffle\n regs_shuffle_with_bit_shifts(regs_vec);\n\n // store, regs -> shared\n int tM = threadIdx.x * N_SF_PER_TD;\n int* slm_tile = slm + (threadIdx.y * SF_TILE_SIZE_I32 +\n tM / SF_TILE_DIM_M * k_tiles_in_tb * SF_TILE_SIZE_I32);\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD; i++) {\n /* TODO rotate_i */\n slm_tile[(tM % SF_TILE_DIM_M) / NEW_SF_TILE_DIM_M_I32 +\n ((tM + i) % NEW_SF_TILE_DIM_M_I32) * NEW_SF_TILE_DIM_K_I32] =\n reinterpret_cast(regs_vec)[i];\n }\n }\n __syncthreads();\n\n // store, shared -> global\n int linear_id = threadIdx.y * blockDim.x + threadIdx.x;\n#pragma unroll\n for (int i = 0; i < m_tiles_in_tb; i++) {\n __align__(16) int4* output_v4i = reinterpret_cast(output_i32[i]);\n __align__(16) int4* slm_v4i =\n reinterpret_cast(slm + i * k_tiles_in_tb * SF_TILE_SIZE_I32);\n#pragma unroll\n for (int j = linear_id; j < SF_TILE_SIZE_I32 * k_tiles_in_tb / 4;\n j += blockDim.x * blockDim.y) {\n output_v4i[j] = slm_v4i[j];\n }\n }\n}\n\n// Dispatch helper: pick the right (IS_PADDED_K, IS_PADDED_M) col-scaling impl\n// specialization at runtime based on the per-block padding predicates. The\n// branching here is uniform across all threads in the block, so the indirect\n// path each block takes still inlines cleanly.\ntemplate \n__device__ __forceinline__ void dispatch_swizzle_col_scaling_kernel_impl(\n const void* input, void* output, const int M, const int K, const int original_M,\n const int original_K, const int bid_x, const int bid_y, const int grid_dim_x,\n const int grid_dim_y, const bool padding_k, const bool padding_m) {\n if (padding_k && padding_m) {\n swizzle_col_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n } else if (padding_k) {\n swizzle_col_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n } else if (padding_m) {\n swizzle_col_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n } else {\n swizzle_col_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n swizzle_col_scaling_kernel(const void* input, void* output, const int M, const int K,\n const int original_M, const int original_K) {\n const bool padding_m = (blockIdx.y == gridDim.y - 1) && (original_M < M);\n const bool padding_k = (blockIdx.x == gridDim.x - 1) && (original_K < K);\n dispatch_swizzle_col_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y,\n padding_k, padding_m);\n}\n\ntemplate \n__device__ inline void regs_shuffle(LType* regs_vec) {\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n if constexpr (N_TILE_PER_TD == 1) return;\n\n constexpr int kVectorSize = N_SF_PER_TD_PER_TILE * N_TILE_PER_TD;\n int32_t tmp[kVectorSize];\n int32_t* ptr = reinterpret_cast(regs_vec);\n#pragma unroll\n for (int i = 0; i < kVectorSize; i++)\n tmp[i % N_TILE_PER_TD * N_SF_PER_TD_PER_TILE + i / N_TILE_PER_TD] = ptr[i];\n\n#pragma unroll\n for (int i = 0; i < kVectorSize; i++) ptr[i] = tmp[i];\n}\n\n// Inverse of regs_shuffle.\ntemplate \n__device__ inline void regs_unshuffle(LType* regs_vec) {\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n if constexpr (N_TILE_PER_TD == 1) return;\n\n constexpr int kVectorSize = N_SF_PER_TD_PER_TILE * N_TILE_PER_TD;\n int32_t tmp[kVectorSize];\n int32_t* ptr = reinterpret_cast(regs_vec);\n#pragma unroll\n for (int i = 0; i < kVectorSize; i++)\n tmp[i % N_SF_PER_TD_PER_TILE * N_TILE_PER_TD + i / N_SF_PER_TD_PER_TILE] = ptr[i];\n\n#pragma unroll\n for (int i = 0; i < kVectorSize; i++) ptr[i] = tmp[i];\n}\n\n// IS_PADDED_K / IS_PADDED_M select the boundary-block specialization at compile\n// time so the inner load loop avoids the per-iteration runtime checks. The\n// caller computes the runtime predicates from blockIdx/gridDim once per block\n// (uniform across the block) and dispatches to the right specialization.\ntemplate \n__device__ void swizzle_row_scaling_kernel_impl(const void* input, void* output, const int M,\n const int K, const int original_M,\n const int original_K, const int bid_x,\n const int bid_y, const int grid_dim_x,\n const int grid_dim_y) {\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int N_TILES_IN_TB = TB_DIM * N_TILE_PER_TD;\n\n // input is in K-major\n constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4;\n constexpr int SF_TILE_DIM_M_I32 = SF_TILE_DIM_M;\n\n int n_tiles_in_tb = N_TILES_IN_TB;\n const int K_i32 = K / 4;\n if (bid_x == grid_dim_x - 1) {\n n_tiles_in_tb = (K_i32 - 1) % N_TILES_IN_TB + 1;\n }\n\n const int input_offset = bid_y * SF_TILE_DIM_M_I32 * K_i32 + bid_x * N_TILES_IN_TB;\n const int* input_i32 = reinterpret_cast(input) + input_offset;\n int* output_i32 = reinterpret_cast(output) + bid_y * SF_TILE_DIM_M_I32 * K_i32 +\n bid_x * N_TILES_IN_TB * SF_TILE_SIZE_I32;\n\n extern __shared__ int4 slm_v4i[];\n\n // load, global -> regs\n // Each register read for a given i is along the K direction at row\n // (bid_y * SF_TILE_DIM_M + i * TB_DIM + threadIdx.y). When that row is past\n // original_M, the entire register is out of the per-tensor data region (which\n // may be the unpadded compact extent), so we must NOT issue the __ldg there --\n // it could read past the per-tensor buffer (and, for the last tensor in a\n // grouped allocation, past the end of the allocation entirely).\n LType regs_vec[N_SF_PER_TD_PER_TILE];\n if (threadIdx.x * N_TILE_PER_TD < n_tiles_in_tb) {\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) {\n const int row = bid_y * SF_TILE_DIM_M + i * TB_DIM + threadIdx.y;\n const int thread_offset = (i * TB_DIM + threadIdx.y) * K_i32 + threadIdx.x * N_TILE_PER_TD;\n if constexpr (IS_PADDED_M) {\n if (row >= original_M) {\n // Entire register is past original_M: zero directly without loading.\n uint8_t* zero_bytes = reinterpret_cast(regs_vec + i);\n#pragma unroll\n for (int j = 0; j < static_cast(sizeof(LType)); j++) zero_bytes[j] = 0;\n continue;\n }\n }\n regs_vec[i] = __ldg(reinterpret_cast(input_i32 + thread_offset));\n // Per-byte K masking is still needed when only part of the register is past\n // original_K (i.e. row is in range but the K position spans the boundary).\n if constexpr (IS_PADDED_K) {\n#pragma unroll\n for (int j = 0; j < N_TILE_PER_TD * sizeof(int); j++) {\n const int index = (input_offset + thread_offset) * sizeof(int) + j;\n if (index % K >= original_K) {\n reinterpret_cast(regs_vec + i)[j] = 0;\n }\n }\n }\n }\n\n // shuffle regs\n regs_shuffle(regs_vec);\n\n// store, regs -> shared\n#pragma unroll\n for (int i = 0; i < N_TILE_PER_TD; i++) {\n /* TODO rotate i */\n slm_v4i[(threadIdx.x * N_TILE_PER_TD + i) * SF_TILE_SIZE_I32 / 4 + threadIdx.y] =\n reinterpret_cast(regs_vec)[i];\n }\n }\n __syncthreads();\n\n // store, shared -> global\n int linear_id = threadIdx.y * blockDim.x + threadIdx.x;\n __align__(16) int4* output_v4i = reinterpret_cast(output_i32);\n#pragma unroll\n for (int i = linear_id; i < SF_TILE_SIZE_I32 * n_tiles_in_tb / 4; i += blockDim.x * blockDim.y) {\n output_v4i[i] = slm_v4i[i];\n }\n}\n\n// Dispatch helper: pick the right (IS_PADDED_K, IS_PADDED_M) row-scaling impl\n// specialization at runtime based on the per-block padding predicates. The\n// branching here is uniform across all threads in the block, so the indirect\n// path each block takes still inlines cleanly.\ntemplate \n__device__ __forceinline__ void dispatch_swizzle_row_scaling_kernel_impl(\n const void* input, void* output, const int M, const int K, const int original_M,\n const int original_K, const int bid_x, const int bid_y, const int grid_dim_x,\n const int grid_dim_y, const bool padding_k, const bool padding_m) {\n if (padding_k && padding_m) {\n swizzle_row_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n } else if (padding_k) {\n swizzle_row_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n } else if (padding_m) {\n swizzle_row_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n } else {\n swizzle_row_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n swizzle_row_scaling_kernel(const void* input, void* output, const int M, const int K,\n const int original_M, const int original_K) {\n const bool padding_m = (blockIdx.y == gridDim.y - 1) && (original_M < M);\n const bool padding_k = (blockIdx.x == gridDim.x - 1) && (original_K < K);\n dispatch_swizzle_row_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y,\n padding_k, padding_m);\n}\n\n// Narrow-K specialization for row scaling swizzle.\n// When K is small (num_tiles_k < TB_DIM), the standard kernel wastes threadIdx.x\n// because there aren't enough K-tiles to distribute across threads.\n// This kernel repurposes the thread dimensions: threadIdx.x iterates rows within\n// an M-tile, threadIdx.y indexes M-tiles within the block, processing TB_DIM\n// M-tiles per block with full thread utilization.\ntemplate \n__device__ void swizzle_row_scaling_narrow_k_kernel_impl(const void* input, void* output,\n const int M, const int K,\n const int original_M, const int original_K,\n const int bid, const int grid_dim) {\n constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4;\n const int K_i32 = K / 4;\n const int num_tiles_m = M / SF_TILE_DIM_M;\n\n const int m_tile = bid * blockDim.y + threadIdx.y;\n const bool active = (m_tile < num_tiles_m);\n\n extern __shared__ int4 slm_v4i[];\n const int slm_tile_v4i = K_i32 * (SF_TILE_SIZE_I32 / 4);\n\n if (active) {\n const bool padding_m = (m_tile == num_tiles_m - 1) && (original_M < M);\n const bool padding_k = (original_K < K);\n\n int4* my_slm = slm_v4i + threadIdx.y * slm_tile_v4i;\n\n for (int k = 0; k < K_i32; k++) {\n const int input_base = m_tile * SF_TILE_DIM_M * K_i32 + k;\n const int* input_i32 = reinterpret_cast(input) + input_base;\n\n int regs[N_SF_PER_TD_PER_TILE];\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) {\n const int row = i * TB_DIM + threadIdx.x;\n regs[i] = __ldg(input_i32 + row * K_i32);\n if (padding_m || padding_k) {\n for (int j = 0; j < 4; j++) {\n const int byte_row = m_tile * SF_TILE_DIM_M + row;\n const int byte_col = k * 4 + j;\n if (byte_row >= original_M || byte_col >= original_K) {\n reinterpret_cast(®s[i])[j] = 0;\n }\n }\n }\n }\n\n my_slm[k * (SF_TILE_SIZE_I32 / 4) + threadIdx.x] = *reinterpret_cast(regs);\n }\n }\n\n __syncthreads();\n\n if (active) {\n int4* my_slm = slm_v4i + threadIdx.y * slm_tile_v4i;\n int4* out_v4i =\n reinterpret_cast(reinterpret_cast(output) + m_tile * SF_TILE_DIM_M * K_i32);\n\n for (int i = threadIdx.x; i < slm_tile_v4i; i += blockDim.x) {\n out_v4i[i] = my_slm[i];\n }\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n swizzle_row_scaling_narrow_k_kernel(const void* input, void* output, const int M, const int K,\n const int original_M, const int original_K) {\n swizzle_row_scaling_narrow_k_kernel_impl(\n input, output, M, K, original_M, original_K, blockIdx.x, gridDim.x);\n}\n\n// Narrow-M variant of the column scaling swizzle kernel, for when num_tiles_m < TB_DIM.\n// Analogous to the narrow-K row kernel: when the M dimension is small, the normal\n// col kernel underutilizes threads in the load phase because threadIdx.x covers M\n// positions with vectorized loads, leaving many threads idle. This kernel repurposes\n// thread dimensions: threadIdx.y indexes K-tiles within the block, threadIdx.x covers\n// one int32 column of an M-tile, and M-tiles are iterated serially.\ntemplate \n__device__ void swizzle_col_scaling_narrow_m_kernel_impl(const void* input, void* output,\n const int M, const int K,\n const int original_M, const int original_K,\n const int bid, const int grid_dim) {\n constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4;\n constexpr int SF_TILE_DIM_M_I32 = SF_TILE_DIM_M / 4;\n constexpr int SF_TILE_DIM_K_I32 = SF_TILE_DIM_K;\n\n const int M_i32 = M / 4;\n const int K_i32 = K;\n const int num_tiles_m = M / SF_TILE_DIM_M;\n const int num_tiles_k = K / SF_TILE_DIM_K;\n\n const int k_tile = bid * blockDim.y + threadIdx.y;\n const bool active = (k_tile < num_tiles_k);\n const int remaining = num_tiles_k - bid * static_cast(blockDim.y);\n const int k_tiles_in_block = remaining <= 0 ? 0 : (remaining < TB_DIM ? remaining : TB_DIM);\n\n extern __shared__ int slm_narrow_m[];\n\n if (active) {\n const bool padding_k = (k_tile == num_tiles_k - 1) && (original_K < K);\n const int32_t* input_i32 = reinterpret_cast(input);\n\n for (int m_tile = 0; m_tile < num_tiles_m; m_tile++) {\n const bool padding_m = (m_tile == num_tiles_m - 1) && (original_M < M);\n\n int regs[N_SF_PER_TD_PER_TILE];\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) {\n const int k_row = k_tile * SF_TILE_DIM_K_I32 + i;\n const int m_col = m_tile * SF_TILE_DIM_M_I32 + threadIdx.x;\n regs[i] = __ldg(input_i32 + k_row * M_i32 + m_col);\n if (padding_m || padding_k) {\n for (int j = 0; j < 4; j++) {\n if (m_col * 4 + j >= original_M || k_row >= original_K) {\n reinterpret_cast(®s[i])[j] = 0;\n }\n }\n }\n }\n\n regs_shuffle_with_bit_shifts(regs);\n\n int tM = threadIdx.x * N_SF_PER_TD_PER_TILE;\n int* slm_tile =\n slm_narrow_m + m_tile * TB_DIM * SF_TILE_SIZE_I32 + threadIdx.y * SF_TILE_SIZE_I32;\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) {\n slm_tile[(tM % SF_TILE_DIM_M) / NEW_SF_TILE_DIM_M_I32 +\n ((tM + i) % NEW_SF_TILE_DIM_M_I32) * NEW_SF_TILE_DIM_K_I32] = regs[i];\n }\n }\n }\n\n __syncthreads();\n\n const int linear_id = threadIdx.y * blockDim.x + threadIdx.x;\n for (int m_tile = 0; m_tile < num_tiles_m; m_tile++) {\n int4* out_v4i = reinterpret_cast(reinterpret_cast(output) +\n m_tile * SF_TILE_DIM_M_I32 * K_i32 +\n bid * TB_DIM * SF_TILE_SIZE_I32);\n int4* slm_v4i = reinterpret_cast(slm_narrow_m + m_tile * TB_DIM * SF_TILE_SIZE_I32);\n const int n_v4i = k_tiles_in_block * SF_TILE_SIZE_I32 / 4;\n for (int j = linear_id; j < n_v4i; j += blockDim.x * blockDim.y) {\n out_v4i[j] = slm_v4i[j];\n }\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n swizzle_col_scaling_narrow_m_kernel(const void* input, void* output, const int M, const int K,\n const int original_M, const int original_K) {\n swizzle_col_scaling_narrow_m_kernel_impl(\n input, output, M, K, original_M, original_K, blockIdx.x, gridDim.x);\n}\n\ntemplate \n__device__ void unswizzle_row_scaling_kernel_impl(const void* input, void* output, const int M,\n const int K, const int bid_x, const int bid_y,\n const int grid_dim_x, const int grid_dim_y) {\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int N_TILES_IN_TB = TB_DIM * N_TILE_PER_TD;\n\n constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4;\n constexpr int SF_TILE_DIM_M_I32 = SF_TILE_DIM_M;\n\n int n_tiles_in_tb = N_TILES_IN_TB;\n const int K_i32 = K / 4;\n if (bid_x == grid_dim_x - 1) {\n n_tiles_in_tb = (K_i32 - 1) % N_TILES_IN_TB + 1;\n }\n\n const int input_offset =\n bid_y * SF_TILE_DIM_M_I32 * K_i32 + bid_x * N_TILES_IN_TB * SF_TILE_SIZE_I32;\n const int* input_i32 = reinterpret_cast(input) + input_offset;\n const int output_offset = bid_y * SF_TILE_DIM_M_I32 * K_i32 + bid_x * N_TILES_IN_TB;\n int* output_i32 = reinterpret_cast(output) + output_offset;\n\n extern __shared__ int4 slm_v4i[];\n\n int linear_id = threadIdx.y * blockDim.x + threadIdx.x;\n const int4* input_v4i = reinterpret_cast(input_i32);\n#pragma unroll\n for (int i = linear_id; i < SF_TILE_SIZE_I32 * n_tiles_in_tb / 4; i += blockDim.x * blockDim.y) {\n slm_v4i[i] = input_v4i[i];\n }\n __syncthreads();\n\n LType regs_vec[N_SF_PER_TD_PER_TILE];\n if (threadIdx.x * N_TILE_PER_TD < n_tiles_in_tb) {\n#pragma unroll\n for (int i = 0; i < N_TILE_PER_TD; i++) {\n reinterpret_cast(regs_vec)[i] =\n slm_v4i[(threadIdx.x * N_TILE_PER_TD + i) * SF_TILE_SIZE_I32 / 4 + threadIdx.y];\n }\n\n regs_unshuffle(regs_vec);\n\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) {\n const int thread_offset = (i * TB_DIM + threadIdx.y) * K_i32 + threadIdx.x * N_TILE_PER_TD;\n reinterpret_cast(output_i32 + thread_offset)[0] = regs_vec[i];\n }\n }\n}\n\ntemplate \n__device__ void unswizzle_col_scaling_kernel_impl(const void* input, void* output, const int M,\n const int K, const int bid_x, const int bid_y,\n const int grid_dim_x, const int grid_dim_y) {\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int N_SF_PER_TD = N_TILE_PER_TD * N_SF_PER_TD_PER_TILE;\n constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4;\n\n constexpr int SF_TILE_DIM_M_I32 = SF_TILE_DIM_M / 4;\n constexpr int SF_TILE_DIM_K_I32 = SF_TILE_DIM_K;\n\n const int M_i32 = M / 4;\n const int K_i32 = K;\n\n int m_tiles_in_tb = N_TILE_PER_TD;\n int k_tiles_in_tb = TB_DIM;\n if (bid_x == grid_dim_x - 1) {\n k_tiles_in_tb = (K_i32 / SF_TILE_DIM_K_I32 - 1) % k_tiles_in_tb + 1;\n }\n if (bid_y == grid_dim_y - 1) {\n m_tiles_in_tb = (M_i32 / SF_TILE_DIM_M_I32 - 1) % m_tiles_in_tb + 1;\n }\n\n const int32_t* input_i32[N_TILE_PER_TD];\n#pragma unroll\n for (int i = 0; i < m_tiles_in_tb; i++) {\n input_i32[i] = reinterpret_cast(input) + bid_x * TB_DIM * SF_TILE_SIZE_I32 +\n (bid_y * N_TILE_PER_TD + i) * SF_TILE_DIM_M_I32 * K_i32;\n }\n const int output_offset =\n bid_x * TB_DIM * SF_TILE_DIM_K_I32 * M_i32 + bid_y * N_TILE_PER_TD * SF_TILE_DIM_M_I32;\n int* output_i32 = reinterpret_cast(output) + output_offset;\n\n extern __shared__ int slm[];\n\n int linear_id = threadIdx.y * blockDim.x + threadIdx.x;\n#pragma unroll\n for (int i = 0; i < m_tiles_in_tb; i++) {\n __align__(16) const int4* input_v4i = reinterpret_cast(input_i32[i]);\n __align__(16) int4* slm_v4i =\n reinterpret_cast(slm + i * k_tiles_in_tb * SF_TILE_SIZE_I32);\n#pragma unroll\n for (int j = linear_id; j < SF_TILE_SIZE_I32 * k_tiles_in_tb / 4;\n j += blockDim.x * blockDim.y) {\n slm_v4i[j] = input_v4i[j];\n }\n }\n __syncthreads();\n\n LType regs_vec[N_SF_PER_TD_PER_TILE];\n if (threadIdx.x * N_TILE_PER_TD < m_tiles_in_tb * SF_TILE_DIM_M_I32 &&\n threadIdx.y < k_tiles_in_tb) {\n int tM = threadIdx.x * N_SF_PER_TD;\n int* slm_tile = slm + (threadIdx.y * SF_TILE_SIZE_I32 +\n tM / SF_TILE_DIM_M * k_tiles_in_tb * SF_TILE_SIZE_I32);\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD; i++) {\n reinterpret_cast(regs_vec)[i] =\n slm_tile[(tM % SF_TILE_DIM_M) / NEW_SF_TILE_DIM_M_I32 +\n ((tM + i) % NEW_SF_TILE_DIM_M_I32) * NEW_SF_TILE_DIM_K_I32];\n }\n\n regs_unshuffle_with_bit_shifts(regs_vec);\n\n#pragma unroll\n for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) {\n const int thread_offset =\n (threadIdx.y * SF_TILE_DIM_K_I32 + i) * M_i32 + threadIdx.x * N_TILE_PER_TD;\n reinterpret_cast(output_i32 + thread_offset)[0] = regs_vec[i];\n }\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n unswizzle_scaling_kernel(const void* input, void* output, const int M, const int K,\n const bool row_scaling) {\n const int bid_x = blockIdx.x;\n const int bid_y = blockIdx.y;\n const int grid_dim_x = gridDim.x;\n const int grid_dim_y = gridDim.y;\n if (row_scaling) {\n unswizzle_row_scaling_kernel_impl(\n input, output, M, K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n } else {\n unswizzle_col_scaling_kernel_impl(\n input, output, M, K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n }\n}\n\nconstexpr int kMaxTensorsPerKernel = 64; // Args must be <4 KB\nstruct MultiSwizzleArgs {\n // (input) Data buffers for input scaling factors\n void* input_list[kMaxTensorsPerKernel];\n // (output) Data buffers for swizzled scaling factors\n void* output_list[kMaxTensorsPerKernel];\n // Input scaling factor m\n int m_list[kMaxTensorsPerKernel];\n // Input scaling factor k\n int k_list[kMaxTensorsPerKernel];\n // Input scaling factor m before padding\n int original_m_list[kMaxTensorsPerKernel];\n // Input scaling factor k before padding\n int original_k_list[kMaxTensorsPerKernel];\n // Prefix sum (with leading zero) of CUDA blocks needed for each\n // tensor\n int block_range[kMaxTensorsPerKernel + 1];\n // Number of tensors being processed by kernel\n int num_tensors;\n};\n\nconstexpr size_t round_up_to_multiple(size_t value, size_t multiple) {\n return DIVUP(value, multiple) * multiple;\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n grouped_swizzle_row_scaling_uniform_shape_kernel(const void* input, void* output, const int M,\n const int K, const int original_M,\n const int original_K,\n const size_t input_stride_bytes,\n const size_t output_stride_bytes) {\n const int tensor_id = blockIdx.z;\n // Input and output strides may differ: input is in the kernel-produced \"compact\"\n // layout (per-tensor stride = original_M * padded_k * elem_size) when callers\n // pass the unswizzled grouped scale buffer as-is, while the output is always in\n // the per-tensor padded (\"swizzle-ready\") layout (padded_m * padded_k * elem_size).\n const uint8_t* input_base =\n reinterpret_cast(input) + tensor_id * input_stride_bytes;\n uint8_t* output_base = reinterpret_cast(output) + tensor_id * output_stride_bytes;\n const bool padding_m = (blockIdx.y == gridDim.y - 1) && (original_M < M);\n const bool padding_k = (blockIdx.x == gridDim.x - 1) && (original_K < K);\n dispatch_swizzle_row_scaling_kernel_impl(\n input_base, output_base, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x,\n gridDim.y, padding_k, padding_m);\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n grouped_swizzle_col_scaling_uniform_shape_kernel(const void* input, void* output, const int M,\n const int K, const int original_M,\n const int original_K,\n const size_t input_stride_bytes,\n const size_t output_stride_bytes) {\n const int tensor_id = blockIdx.z;\n // See the rowwise kernel for stride semantics. For columnwise the per-tensor\n // compact stride is DIVUP(original_K, 1) * padded_m * elem_size (i.e. the\n // unpadded scale-row count in the K direction times the padded M extent).\n const uint8_t* input_base =\n reinterpret_cast(input) + tensor_id * input_stride_bytes;\n uint8_t* output_base = reinterpret_cast(output) + tensor_id * output_stride_bytes;\n const bool padding_m = (blockIdx.y == gridDim.y - 1) && (original_M < M);\n const bool padding_k = (blockIdx.x == gridDim.x - 1) && (original_K < K);\n dispatch_swizzle_col_scaling_kernel_impl(\n input_base, output_base, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x,\n gridDim.y, padding_k, padding_m);\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n grouped_unswizzle_scaling_uniform_shape_kernel(const void* input, void* output, const int M,\n const int K, const size_t scale_stride_bytes,\n const bool row_scaling) {\n const int tensor_id = blockIdx.z;\n const uint8_t* input_base =\n reinterpret_cast(input) + tensor_id * scale_stride_bytes;\n uint8_t* output_base = reinterpret_cast(output) + tensor_id * scale_stride_bytes;\n if (row_scaling) {\n unswizzle_row_scaling_kernel_impl(\n input_base, output_base, M, K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y);\n } else {\n unswizzle_col_scaling_kernel_impl(\n input_base, output_base, M, K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y);\n }\n}\n\ntemplate \n__global__ void multi_tensor_unswizzle_row_scaling_kernel(MultiSwizzleArgs kernel_args) {\n const int bid = blockIdx.x;\n int tensor_id = 0;\n while (kernel_args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n const void* input = kernel_args.input_list[tensor_id];\n void* output = kernel_args.output_list[tensor_id];\n const int M = kernel_args.m_list[tensor_id];\n const int K = kernel_args.k_list[tensor_id];\n\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int N_TILES_IN_TB = TB_DIM * N_TILE_PER_TD;\n\n const int num_tiles_k = K / SF_TILE_DIM_K;\n const int num_tiles_m = M / SF_TILE_DIM_M;\n const int flat_offset = bid - kernel_args.block_range[tensor_id];\n const int grid_dim_x = DIVUP(num_tiles_k, N_TILES_IN_TB);\n const int grid_dim_y = num_tiles_m;\n const int bid_x = flat_offset / grid_dim_y;\n const int bid_y = flat_offset % grid_dim_y;\n\n unswizzle_row_scaling_kernel_impl(\n input, output, M, K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n}\n\ntemplate \n__global__ void multi_tensor_unswizzle_col_scaling_kernel(MultiSwizzleArgs kernel_args) {\n const int bid = blockIdx.x;\n int tensor_id = 0;\n while (kernel_args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n const void* input = kernel_args.input_list[tensor_id];\n void* output = kernel_args.output_list[tensor_id];\n const int M = kernel_args.m_list[tensor_id];\n const int K = kernel_args.k_list[tensor_id];\n\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n\n const int num_tiles_k = K / SF_TILE_DIM_K;\n const int num_tiles_m = M / SF_TILE_DIM_M;\n const int flat_offset = bid - kernel_args.block_range[tensor_id];\n const int grid_dim_x = DIVUP(num_tiles_k, TB_DIM);\n const int grid_dim_y = DIVUP(num_tiles_m, N_TILE_PER_TD);\n const int bid_x = flat_offset / grid_dim_y;\n const int bid_y = flat_offset % grid_dim_y;\n\n unswizzle_col_scaling_kernel_impl(\n input, output, M, K, bid_x, bid_y, grid_dim_x, grid_dim_y);\n}\n\ntemplate \n__global__ void multi_tensor_swizzle_row_scaling_kernel(MultiSwizzleArgs kernel_args) {\n // Find tensor corresponding to block\n const int bid = blockIdx.x;\n int tensor_id = 0;\n while (kernel_args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n // Get args corresponding to block\n const void* input = kernel_args.input_list[tensor_id];\n void* output = kernel_args.output_list[tensor_id];\n const int M = kernel_args.m_list[tensor_id];\n const int K = kernel_args.k_list[tensor_id];\n const int original_M = kernel_args.original_m_list[tensor_id];\n const int original_K = kernel_args.original_k_list[tensor_id];\n\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n constexpr int N_TILES_IN_TB = TB_DIM * N_TILE_PER_TD;\n\n // Get block index in grid. Emulate 2D grid.\n const int num_tiles_k = K / SF_TILE_DIM_K;\n const int num_tiles_m = M / SF_TILE_DIM_M;\n const int grid_dim_x = DIVUP(num_tiles_k, N_TILES_IN_TB);\n const int grid_dim_y = num_tiles_m;\n const int bid_x = (bid - kernel_args.block_range[tensor_id]) / grid_dim_y;\n const int bid_y = (bid - kernel_args.block_range[tensor_id]) % grid_dim_y;\n\n const bool padding_m = (bid_y == grid_dim_y - 1) && (original_M < M);\n const bool padding_k = (bid_x == grid_dim_x - 1) && (original_K < K);\n dispatch_swizzle_row_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y, padding_k,\n padding_m);\n}\n\ntemplate \n__global__ void multi_tensor_swizzle_col_scaling_kernel(MultiSwizzleArgs kernel_args) {\n // Find tensor corresponding to block\n const int bid = blockIdx.x;\n int tensor_id = 0;\n while (kernel_args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n // Get args corresponding to block\n const void* input = kernel_args.input_list[tensor_id];\n void* output = kernel_args.output_list[tensor_id];\n const int M = kernel_args.m_list[tensor_id];\n const int K = kernel_args.k_list[tensor_id];\n const int original_M = kernel_args.original_m_list[tensor_id];\n const int original_K = kernel_args.original_k_list[tensor_id];\n\n constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int);\n\n // Get block index in grid. Emulate 2D grid.\n const int num_tiles_k = K / SF_TILE_DIM_K;\n const int num_tiles_m = M / SF_TILE_DIM_M;\n const int grid_dim_x = DIVUP(num_tiles_k, TB_DIM);\n const int grid_dim_y = DIVUP(num_tiles_m, N_TILE_PER_TD);\n const int bid_x = (bid - kernel_args.block_range[tensor_id]) / grid_dim_y;\n const int bid_y = (bid - kernel_args.block_range[tensor_id]) % grid_dim_y;\n\n const bool padding_m = (bid_y == grid_dim_y - 1) && (original_M < M);\n const bool padding_k = (bid_x == grid_dim_x - 1) && (original_K < K);\n dispatch_swizzle_col_scaling_kernel_impl(\n input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y, padding_k,\n padding_m);\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n multi_tensor_swizzle_row_scaling_narrow_k_kernel(MultiSwizzleArgs kernel_args) {\n const int bid = blockIdx.x;\n int tensor_id = 0;\n while (kernel_args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n const void* input = kernel_args.input_list[tensor_id];\n void* output = kernel_args.output_list[tensor_id];\n const int M = kernel_args.m_list[tensor_id];\n const int K = kernel_args.k_list[tensor_id];\n const int original_M = kernel_args.original_m_list[tensor_id];\n const int original_K = kernel_args.original_k_list[tensor_id];\n const int flat_bid = bid - kernel_args.block_range[tensor_id];\n const int num_tiles_m = M / SF_TILE_DIM_M;\n const int grid_dim = DIVUP(num_tiles_m, TB_DIM);\n\n swizzle_row_scaling_narrow_k_kernel_impl(\n input, output, M, K, original_M, original_K, flat_bid, grid_dim);\n}\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n multi_tensor_swizzle_col_scaling_narrow_m_kernel(MultiSwizzleArgs kernel_args) {\n const int bid = blockIdx.x;\n int tensor_id = 0;\n while (kernel_args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n const void* input = kernel_args.input_list[tensor_id];\n void* output = kernel_args.output_list[tensor_id];\n const int M = kernel_args.m_list[tensor_id];\n const int K = kernel_args.k_list[tensor_id];\n const int original_M = kernel_args.original_m_list[tensor_id];\n const int original_K = kernel_args.original_k_list[tensor_id];\n const int flat_bid = bid - kernel_args.block_range[tensor_id];\n const int num_tiles_k = K / SF_TILE_DIM_K;\n const int grid_dim = DIVUP(num_tiles_k, TB_DIM);\n\n swizzle_col_scaling_narrow_m_kernel_impl(\n input, output, M, K, original_M, original_K, flat_bid, grid_dim);\n}\n\n} // namespace\n\nvoid swizzle_scaling_factors(const Tensor* input, Tensor* output, hipStream_t stream) {\n // Check scaling mode\n const auto& scaling_mode = input->scaling_mode;\n NVTE_CHECK(scaling_mode == NVTE_MXFP8_1D_SCALING || scaling_mode == NVTE_NVFP4_1D_SCALING,\n \"Input tensor has invalid scaling mode (\", to_string(input->scaling_mode), \").\");\n\n // Check tensors\n CheckInputTensor(*input, \"scaling_factor_input\");\n CheckInputTensor(*output, \"scaling_factor_output\");\n NVTE_CHECK(!input->with_gemm_swizzled_scales,\n \"Expected input tensor with scales in compact format.\");\n NVTE_CHECK(output->with_gemm_swizzled_scales,\n \"Expected output tensor with scales in GEMM swizzled format.\");\n switch (scaling_mode) {\n case NVTE_MXFP8_1D_SCALING:\n NVTE_CHECK(is_fp8_dtype(input->dtype()), \"Input tensor has invalid dtype (expected FP8, got \",\n to_string(input->dtype()), \").\");\n break;\n case NVTE_NVFP4_1D_SCALING:\n NVTE_CHECK(is_fp4_dtype(input->dtype()), \"Input tensor has invalid dtype (expected FP4, got \",\n to_string(input->dtype()), \").\");\n break;\n default:\n NVTE_ERROR(\"Invalid scaling mode\");\n }\n\n // Check if scaling factors are non-trivial\n const bool has_rowwise_scale_inv = input->scale_inv.has_data();\n const bool has_columnwise_scale_inv = input->columnwise_scale_inv.has_data();\n NVTE_CHECK(!has_rowwise_scale_inv || !has_columnwise_scale_inv,\n \"Input tensor has both row-wise and column-wise scaling factors\");\n if (!has_rowwise_scale_inv && !has_columnwise_scale_inv) {\n return;\n }\n\n // Deduce tensor dims\n int m{0}, k{0};\n switch (scaling_mode) {\n case NVTE_MXFP8_1D_SCALING: {\n if (has_rowwise_scale_inv) {\n NVTE_CHECK(input->scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", input->scale_inv.shape, \".\");\n m = input->scale_inv.shape[0];\n k = input->scale_inv.shape[1];\n } else if (has_columnwise_scale_inv) {\n NVTE_CHECK(input->columnwise_scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", input->columnwise_scale_inv.shape,\n \".\");\n m = input->columnwise_scale_inv.shape[1];\n k = input->columnwise_scale_inv.shape[0];\n }\n break;\n }\n case NVTE_NVFP4_1D_SCALING: {\n if (has_rowwise_scale_inv) {\n NVTE_CHECK(input->scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", input->scale_inv.shape, \".\");\n m = input->scale_inv.shape[0];\n k = input->scale_inv.shape[1];\n } else if (has_columnwise_scale_inv) {\n NVTE_CHECK(input->columnwise_scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", input->columnwise_scale_inv.shape,\n \".\");\n m = input->columnwise_scale_inv.shape[0];\n k = input->columnwise_scale_inv.shape[1];\n }\n break;\n }\n default:\n NVTE_ERROR(\"Invalid scaling mode\");\n }\n\n // Check dims\n constexpr int SF_TILE_DIM_M = 128;\n constexpr int SF_TILE_DIM_K = 4;\n NVTE_CHECK(m % SF_TILE_DIM_M == 0, \"Input should be padded in M/N dimension!\");\n NVTE_CHECK(k % SF_TILE_DIM_K == 0, \"Input should be padded in K dimension!\");\n\n // Check that output tensor matches input tensor\n if (has_rowwise_scale_inv) {\n NVTE_CHECK(output->scale_inv.has_data(),\n \"Output tensor does not have row-wise scaling factors.\");\n NVTE_CHECK(m * k == output->scale_inv.numel(), \"Expected output tensor to have \", m * k,\n \" row-wise scaling factors, but got shape=\", output->scale_inv.shape, \".\");\n }\n if (has_columnwise_scale_inv) {\n NVTE_CHECK(output->columnwise_scale_inv.has_data(),\n \"Output tensor does not have column-wise scaling factors.\");\n NVTE_CHECK(\n m * k == output->columnwise_scale_inv.numel(), \"Expected output tensor to have \", m * k,\n \" column-wise scaling factors, but got shape=\", output->columnwise_scale_inv.shape, \".\");\n }\n\n // Choose swizzle implementation\n bool rowwise_swizzle{false}, columnwise_swizzle{false};\n switch (scaling_mode) {\n case NVTE_MXFP8_1D_SCALING: {\n rowwise_swizzle = has_rowwise_scale_inv;\n columnwise_swizzle = has_columnwise_scale_inv;\n break;\n }\n case NVTE_NVFP4_1D_SCALING: {\n // NVFP4 column-wise data is transposed, so row-wise and\n // column-wise scales have same swizzling format\n rowwise_swizzle = true;\n columnwise_swizzle = false;\n break;\n }\n default:\n NVTE_ERROR(\"Invalid scaling mode\");\n }\n\n const dim3 block_size(TB_DIM, TB_DIM);\n const int num_tiles_m = m / SF_TILE_DIM_M;\n const int num_tiles_k = k / SF_TILE_DIM_K;\n\n // Perform row-wise swizzle\n if (rowwise_swizzle) {\n int original_M{0}, original_K{0};\n void *input_scale_inv_ptr{nullptr}, *output_scale_inv_ptr{nullptr};\n switch (scaling_mode) {\n case NVTE_MXFP8_1D_SCALING: {\n original_M = input->flat_first_dim();\n original_K = input->flat_last_dim() / MXFP8_BLOCK_SIZE;\n input_scale_inv_ptr = input->scale_inv.dptr;\n output_scale_inv_ptr = output->scale_inv.dptr;\n break;\n }\n case NVTE_NVFP4_1D_SCALING: {\n if (has_rowwise_scale_inv) {\n original_M = input->flat_first_dim();\n original_K = input->flat_last_dim() / NVFP4_BLOCK_SIZE;\n input_scale_inv_ptr = input->scale_inv.dptr;\n output_scale_inv_ptr = output->scale_inv.dptr;\n } else if (has_columnwise_scale_inv) {\n original_M = input->flat_last_dim();\n original_K = input->flat_first_dim() / NVFP4_BLOCK_SIZE;\n input_scale_inv_ptr = input->columnwise_scale_inv.dptr;\n output_scale_inv_ptr = output->columnwise_scale_inv.dptr;\n }\n break;\n }\n default:\n NVTE_ERROR(\"Invalid scaling mode\");\n }\n\n const int narrow_k_slm_size =\n TB_DIM * num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t));\n if (num_tiles_k < TB_DIM && narrow_k_slm_size <= get_max_dynamic_smem()) {\n // Narrow-K: batch TB_DIM M-tiles per block, fully utilizing all threads.\n dim3 num_blocks_narrow(DIVUP(num_tiles_m, TB_DIM));\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(swizzle_row_scaling_narrow_k_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, narrow_k_slm_size));\n swizzle_row_scaling_narrow_k_kernel\n <<>>(\n input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K);\n } else {\n int vec_load_size = (num_tiles_k - 1) % 4 + 1;\n /* there is no int3 and misaligned if using int4/int2 */\n if (vec_load_size == 3) vec_load_size = 1;\n int n_tiles_in_tb = TB_DIM * vec_load_size;\n dim3 num_blocks(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m);\n int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n\n switch (vec_load_size) {\n case 4:\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(swizzle_row_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n swizzle_row_scaling_kernel\n <<>>(\n input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K);\n break;\n case 2:\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(swizzle_row_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n swizzle_row_scaling_kernel\n <<>>(\n input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K);\n break;\n case 1:\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(swizzle_row_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n swizzle_row_scaling_kernel\n <<>>(\n input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K);\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n break;\n }\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n\n // Perform column-wise swizzle\n if (columnwise_swizzle) {\n const int original_M = input->flat_last_dim();\n const int original_K = input->flat_first_dim() / MXFP8_BLOCK_SIZE;\n\n const int narrow_m_slm_size =\n TB_DIM * num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t));\n if (num_tiles_m < TB_DIM && narrow_m_slm_size <= get_max_dynamic_smem()) {\n // Narrow-M: batch TB_DIM K-tiles per block, fully utilizing all threads.\n dim3 num_blocks_narrow(DIVUP(num_tiles_k, TB_DIM));\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(swizzle_col_scaling_narrow_m_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, narrow_m_slm_size));\n swizzle_col_scaling_narrow_m_kernel\n <<>>(\n input->columnwise_scale_inv.dptr, output->columnwise_scale_inv.dptr, m, k, original_M,\n original_K);\n } else {\n int vec_load_size = (num_tiles_m - 1) % 4 + 1;\n if (vec_load_size == 3) vec_load_size = 1; /* no int3 and misaligned if using int4/int2 */\n int n_tiles_in_tb = TB_DIM * vec_load_size;\n dim3 num_blocks(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size));\n int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n\n switch (vec_load_size) {\n case 4:\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(swizzle_col_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n swizzle_col_scaling_kernel\n <<>>(input->columnwise_scale_inv.dptr,\n output->columnwise_scale_inv.dptr, m,\n k, original_M, original_K);\n break;\n case 2:\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(swizzle_col_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n swizzle_col_scaling_kernel\n <<>>(input->columnwise_scale_inv.dptr,\n output->columnwise_scale_inv.dptr, m,\n k, original_M, original_K);\n break;\n case 1:\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(swizzle_col_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n swizzle_col_scaling_kernel\n <<>>(input->columnwise_scale_inv.dptr,\n output->columnwise_scale_inv.dptr, m,\n k, original_M, original_K);\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n break;\n }\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\ntemplate \nvoid launch_multi_tensor_swizzle_scaling_factors(MultiSwizzleArgs& kernel_args,\n const int vec_load_size, const bool is_rowwise,\n const bool use_narrow_k, const bool use_narrow_m,\n hipStream_t stream) {\n // cudaFuncSetAttribute is a host-synchronous driver call; cache the max shared memory\n // setting per kernel variant so we only pay the cost when slm_size actually increases.\n auto set_smem_if_needed = [](auto kernel_fn, int slm, int& cached) {\n if (cached < slm) {\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(kernel_fn), hipFuncAttributeMaxDynamicSharedMemorySize, slm));\n cached = slm;\n }\n };\n\n dim3 block_size(TB_DIM, TB_DIM);\n\n if (is_rowwise && use_narrow_k) {\n // Narrow-K path: each block handles TB_DIM M-tiles with full thread utilization.\n // slm_size depends on num_tiles_k, which can vary per tensor \u2014 use the max.\n int max_num_tiles_k = 0;\n for (size_t j = 0; j < kernel_args.num_tensors; j++) {\n const int num_tiles_m = kernel_args.m_list[j] / SF_TILE_DIM_M;\n const int num_tiles_k = kernel_args.k_list[j] / SF_TILE_DIM_K;\n max_num_tiles_k = std::max(max_num_tiles_k, num_tiles_k);\n kernel_args.block_range[j + 1] = kernel_args.block_range[j] + DIVUP(num_tiles_m, TB_DIM);\n }\n int slm_size = TB_DIM * max_num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n const int num_blocks = kernel_args.block_range[kernel_args.num_tensors];\n\n static int cached_narrow_k = -1;\n set_smem_if_needed(\n multi_tensor_swizzle_row_scaling_narrow_k_kernel, slm_size,\n cached_narrow_k);\n multi_tensor_swizzle_row_scaling_narrow_k_kernel\n <<>>(kernel_args);\n } else if (!is_rowwise && use_narrow_m) {\n // Narrow-M path: each block handles TB_DIM K-tiles with full thread utilization.\n // slm_size depends on num_tiles_m, which can vary per tensor \u2014 use the max.\n int max_num_tiles_m = 0;\n for (size_t j = 0; j < kernel_args.num_tensors; j++) {\n const int num_tiles_m = kernel_args.m_list[j] / SF_TILE_DIM_M;\n const int num_tiles_k = kernel_args.k_list[j] / SF_TILE_DIM_K;\n max_num_tiles_m = std::max(max_num_tiles_m, num_tiles_m);\n kernel_args.block_range[j + 1] = kernel_args.block_range[j] + DIVUP(num_tiles_k, TB_DIM);\n }\n int slm_size = TB_DIM * max_num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n const int num_blocks = kernel_args.block_range[kernel_args.num_tensors];\n\n static int cached_narrow_m = -1;\n set_smem_if_needed(\n multi_tensor_swizzle_col_scaling_narrow_m_kernel, slm_size,\n cached_narrow_m);\n multi_tensor_swizzle_col_scaling_narrow_m_kernel\n <<>>(kernel_args);\n } else {\n int n_tiles_in_tb = TB_DIM * vec_load_size;\n int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n /* Calculate number of CUDA blocks needed for each tensor.\n * We have to do it here because we have to iterate over all tensors in this batch to\n * get the minimum vec_load_size.\n */\n for (size_t j = 0; j < kernel_args.num_tensors; j++) {\n const int m = kernel_args.m_list[j];\n const int k = kernel_args.k_list[j];\n int num_tiles_m = m / SF_TILE_DIM_M;\n int num_tiles_k = k / SF_TILE_DIM_K;\n if (is_rowwise) {\n kernel_args.block_range[j + 1] =\n kernel_args.block_range[j] + DIVUP(num_tiles_k, n_tiles_in_tb) * num_tiles_m;\n } else {\n kernel_args.block_range[j + 1] =\n kernel_args.block_range[j] +\n DIVUP(num_tiles_k, TB_DIM) * DIVUP(num_tiles_m, vec_load_size);\n }\n }\n const int num_blocks = kernel_args.block_range[kernel_args.num_tensors];\n\n static int cached_row_int4 = -1, cached_row_int2 = -1, cached_row_int1 = -1;\n static int cached_col_int4 = -1, cached_col_int2 = -1, cached_col_int1 = -1;\n\n if (is_rowwise) {\n switch (vec_load_size) {\n case 4:\n set_smem_if_needed(\n multi_tensor_swizzle_row_scaling_kernel, slm_size,\n cached_row_int4);\n multi_tensor_swizzle_row_scaling_kernel\n <<>>(kernel_args);\n break;\n case 2:\n set_smem_if_needed(\n multi_tensor_swizzle_row_scaling_kernel, slm_size,\n cached_row_int2);\n multi_tensor_swizzle_row_scaling_kernel\n <<>>(kernel_args);\n break;\n case 1:\n set_smem_if_needed(\n multi_tensor_swizzle_row_scaling_kernel, slm_size,\n cached_row_int1);\n multi_tensor_swizzle_row_scaling_kernel\n <<>>(kernel_args);\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n break;\n }\n } else {\n switch (vec_load_size) {\n case 4:\n set_smem_if_needed(\n multi_tensor_swizzle_col_scaling_kernel, slm_size,\n cached_col_int4);\n multi_tensor_swizzle_col_scaling_kernel\n <<>>(kernel_args);\n break;\n case 2:\n set_smem_if_needed(\n multi_tensor_swizzle_col_scaling_kernel, slm_size,\n cached_col_int2);\n multi_tensor_swizzle_col_scaling_kernel\n <<>>(kernel_args);\n break;\n case 1:\n set_smem_if_needed(\n multi_tensor_swizzle_col_scaling_kernel, slm_size,\n cached_col_int1);\n multi_tensor_swizzle_col_scaling_kernel\n <<>>(kernel_args);\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n break;\n }\n }\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\ntemplate \nvoid launch_multi_tensor_unswizzle_scaling_factors(MultiSwizzleArgs& kernel_args,\n const int vec_load_size, const bool is_rowwise,\n hipStream_t stream) {\n int n_tiles_in_tb = TB_DIM * vec_load_size;\n int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n for (size_t j = 0; j < kernel_args.num_tensors; j++) {\n const int m = kernel_args.m_list[j];\n const int k = kernel_args.k_list[j];\n int num_tiles_m = m / SF_TILE_DIM_M;\n int num_tiles_k = k / SF_TILE_DIM_K;\n if (is_rowwise) {\n kernel_args.block_range[j + 1] =\n kernel_args.block_range[j] + DIVUP(num_tiles_k, n_tiles_in_tb) * num_tiles_m;\n } else {\n kernel_args.block_range[j + 1] =\n kernel_args.block_range[j] +\n DIVUP(num_tiles_k, TB_DIM) * DIVUP(num_tiles_m, vec_load_size);\n }\n }\n\n int num_blocks = kernel_args.block_range[kernel_args.num_tensors];\n if (num_blocks > 0) {\n dim3 block_size(TB_DIM, TB_DIM);\n if (is_rowwise) {\n switch (vec_load_size) {\n case 4:\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(multi_tensor_unswizzle_row_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n multi_tensor_unswizzle_row_scaling_kernel\n <<>>(kernel_args);\n break;\n case 2:\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(multi_tensor_unswizzle_row_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n multi_tensor_unswizzle_row_scaling_kernel\n <<>>(kernel_args);\n break;\n case 1:\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(multi_tensor_unswizzle_row_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n multi_tensor_unswizzle_row_scaling_kernel\n <<>>(kernel_args);\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n break;\n }\n } else {\n switch (vec_load_size) {\n case 4:\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(multi_tensor_unswizzle_col_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n multi_tensor_unswizzle_col_scaling_kernel\n <<>>(kernel_args);\n break;\n case 2:\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(multi_tensor_unswizzle_col_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n multi_tensor_unswizzle_col_scaling_kernel\n <<>>(kernel_args);\n break;\n case 1:\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(multi_tensor_unswizzle_col_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n multi_tensor_unswizzle_col_scaling_kernel\n <<>>(kernel_args);\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n break;\n }\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\nvoid multi_tensor_swizzle_scaling_factors(const std::vector& input,\n std::vector& output, hipStream_t stream,\n bool check_scale_inv_shapes) {\n auto num_tensors = input.size();\n bool all_has_data = true;\n bool all_has_columnwise_data = true;\n bool all_nvfp4 = true;\n for (size_t i = 0; i < num_tensors; i++) {\n auto scaling_mode = input[i]->scaling_mode;\n auto is_fp8 = is_fp8_dtype(input[i]->dtype());\n auto is_fp4 = is_fp4_dtype(input[i]->dtype());\n NVTE_CHECK(\n (is_fp8 && is_mxfp8_scaling(scaling_mode)) || (is_fp4 && is_nvfp4_scaling(scaling_mode)),\n \"Not implemented scaling mode \" + to_string(scaling_mode) + \".\");\n NVTE_CHECK(!input[i]->with_gemm_swizzled_scales,\n \"Expected input tensors with scales in compact format.\");\n NVTE_CHECK(output[i]->with_gemm_swizzled_scales,\n \"Expected output tensors with scales in GEMM swizzled format.\");\n\n // We don't allow empty tensors. They should be filtered out before calling this function.\n NVTE_CHECK(input[i]->numel() != 0, \"Tensor input[\", i, \"] is empty.\");\n CheckInputTensor(*input[i], \"scaling_factor_input[\" + std::to_string(i) + \"]\",\n check_scale_inv_shapes);\n CheckInputTensor(*output[i], \"scaling_factor_output[\" + std::to_string(i) + \"]\",\n check_scale_inv_shapes);\n all_has_data = all_has_data && input[i]->scale_inv.has_data();\n all_has_columnwise_data =\n (all_has_columnwise_data && input[i]->columnwise_scale_inv.has_data());\n all_nvfp4 = all_nvfp4 && is_nvfp4_scaling(scaling_mode);\n }\n NVTE_CHECK(all_has_data || all_has_columnwise_data,\n \"All tensors should have data or columnwise data.\");\n NVTE_CHECK(!all_has_data || !all_has_columnwise_data,\n \"All tensors have both data and columnwise data.\");\n\n const bool rowwise_swizzle = all_has_data || all_nvfp4;\n const bool columnwise_swizzle = all_has_columnwise_data && !all_nvfp4;\n\n constexpr int SF_TILE_DIM_M = 128;\n constexpr int SF_TILE_DIM_K = 4;\n if (rowwise_swizzle) {\n MultiSwizzleArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.block_range[0] = 0;\n int vec_load_size = 4;\n bool all_narrow_k = true;\n for (size_t i = 0; i < num_tensors; i++) {\n //Launch kernel if argument struct is full\n if (kernel_args.num_tensors == kMaxTensorsPerKernel) {\n // There is no int3 and misaligned if using int4/int2.\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_swizzle_scaling_factors(\n kernel_args, vec_load_size, true, all_narrow_k, false, stream);\n // Reset the argument struct and vec_load_size\n kernel_args.num_tensors = 0;\n vec_load_size = 4;\n all_narrow_k = true;\n }\n\n int m, k;\n\n if (all_has_data) {\n m = input[i]->scale_inv.shape[0];\n k = input[i]->scale_inv.shape[1];\n } else {\n NVTE_CHECK(all_nvfp4, \"When doing rowwise swizzle with rowwise data, it has to be NVFP4\");\n m = input[i]->columnwise_scale_inv.shape[0];\n k = input[i]->columnwise_scale_inv.shape[1];\n }\n\n NVTE_CHECK(m % SF_TILE_DIM_M == 0, \"Input should be padded in M/N dimension!\");\n NVTE_CHECK(k % SF_TILE_DIM_K == 0, \"Input should be padded in K dimension!\");\n NVTE_CHECK(k > 0, \"Input scale inverse should be 2D!\");\n\n if (all_has_data) {\n NVTE_CHECK(output[i]->scale_inv.has_data(), \"Output tensor \", i,\n \" does not have row-wise scaling factors.\");\n NVTE_CHECK(m * k == output[i]->scale_inv.numel(), \"Expected output tensor \", i, \" to have \",\n m * k, \" row-wise scaling factors, but got shape=\", output[i]->scale_inv.shape,\n \".\");\n }\n if (all_has_columnwise_data) {\n NVTE_CHECK(output[i]->columnwise_scale_inv.has_data(), \"Output tensor \", i,\n \" does not have column-wise scaling factors.\");\n NVTE_CHECK(m * k == output[i]->columnwise_scale_inv.numel(), \"Expected output tensor \", i,\n \" to have \", m * k, \" column-wise scaling factors, but got shape=\",\n output[i]->columnwise_scale_inv.shape, \".\");\n }\n\n int num_tiles_k = k / SF_TILE_DIM_K;\n const int narrow_k_slm =\n TB_DIM * num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t));\n all_narrow_k =\n all_narrow_k && (num_tiles_k < TB_DIM) && (narrow_k_slm <= get_max_dynamic_smem());\n int vec_load_size_i = (num_tiles_k - 1) % 4 + 1;\n // We use the minimum vec_load_size across all tensors.\n // TODO(zhongbo): fix vec_load_size for NVFP4\n // Current unit test won't capture this issue, but in E2E\n // using vec_load_size = 1 other than 1 will lead to mis-aligned\n // address error in MOE training\n vec_load_size = all_nvfp4 ? 1 : std::min(vec_load_size, vec_load_size_i);\n\n const int pos = kernel_args.num_tensors;\n kernel_args.m_list[pos] = m;\n kernel_args.k_list[pos] = k;\n if (!all_nvfp4 || all_has_data) {\n int block_scale_size = all_nvfp4 ? NVFP4_BLOCK_SIZE : MXFP8_BLOCK_SIZE;\n kernel_args.input_list[pos] = const_cast(input[i]->scale_inv.dptr);\n kernel_args.output_list[pos] = output[i]->scale_inv.dptr;\n kernel_args.original_m_list[pos] = input[i]->flat_first_dim();\n kernel_args.original_k_list[pos] = input[i]->flat_last_dim() / block_scale_size;\n } else {\n kernel_args.input_list[pos] = const_cast(input[i]->columnwise_scale_inv.dptr);\n kernel_args.output_list[pos] = output[i]->columnwise_scale_inv.dptr;\n kernel_args.original_m_list[pos] = input[i]->flat_last_dim();\n kernel_args.original_k_list[pos] = input[i]->flat_first_dim() / NVFP4_BLOCK_SIZE;\n }\n kernel_args.num_tensors++;\n }\n // Launch the remaining tensors\n // There is no int3 and misaligned if using int4/int2.\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_swizzle_scaling_factors(\n kernel_args, vec_load_size, true, all_narrow_k, false, stream);\n }\n\n if (columnwise_swizzle) {\n // NVFP4 shouldn't end up here because it only needs rowwise swizzle\n NVTE_CHECK(!all_nvfp4, \"NVFP4 shouldn't end up here because it only needs rowwise swizzle\");\n\n MultiSwizzleArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.block_range[0] = 0;\n int vec_load_size = 4;\n bool all_narrow_m = true;\n for (size_t i = 0; i < num_tensors; i++) {\n //Launch kernel if argument struct is full\n if (kernel_args.num_tensors == kMaxTensorsPerKernel) {\n // There is no int3 and misaligned if using int4/int2.\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_swizzle_scaling_factors(\n kernel_args, vec_load_size, false, false, all_narrow_m, stream);\n // Reset the argument struct and vec_load_size\n kernel_args.num_tensors = 0;\n vec_load_size = 4;\n all_narrow_m = true;\n }\n const int m = input[i]->columnwise_scale_inv.shape[1];\n const int k = input[i]->columnwise_scale_inv.shape[0];\n\n NVTE_CHECK(m % SF_TILE_DIM_M == 0, \"Input should be padded in M/N dimension!\");\n NVTE_CHECK(k % SF_TILE_DIM_K == 0, \"Input should be padded in K dimension!\");\n NVTE_CHECK(k > 0, \"Input scale inverse should be 2D!\");\n NVTE_CHECK(m * k == std::accumulate(output[i]->columnwise_scale_inv.shape.begin(),\n output[i]->columnwise_scale_inv.shape.end(), 1,\n std::multiplies()),\n \"Input.columnwise_scale_inv size is not equal to \"\n \"Output.columnwise_scale_inv size!\");\n\n int num_tiles_m = m / SF_TILE_DIM_M;\n int num_tiles_k = k / SF_TILE_DIM_K;\n const int narrow_m_slm =\n TB_DIM * num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t));\n all_narrow_m =\n all_narrow_m && (num_tiles_m < TB_DIM) && (narrow_m_slm <= get_max_dynamic_smem());\n int vec_load_size_i = (num_tiles_k - 1) % 4 + 1;\n // We use the minimum vec_load_size across all tensors.\n vec_load_size = std::min(vec_load_size, vec_load_size_i);\n\n const int pos = kernel_args.num_tensors;\n kernel_args.input_list[pos] = const_cast(input[i]->columnwise_scale_inv.dptr);\n kernel_args.output_list[pos] = output[i]->columnwise_scale_inv.dptr;\n kernel_args.m_list[pos] = m;\n kernel_args.k_list[pos] = k;\n kernel_args.original_m_list[pos] = input[i]->flat_last_dim();\n kernel_args.original_k_list[pos] = input[i]->flat_first_dim() / MXFP8_BLOCK_SIZE;\n kernel_args.num_tensors++;\n }\n // Launch the remaining tensors\n // There is no int3 and misaligned if using int4/int2.\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_swizzle_scaling_factors(\n kernel_args, vec_load_size, false, false, all_narrow_m, stream);\n }\n}\n\nvoid unswizzle_scaling_factors(const Tensor* input, Tensor* output, hipStream_t stream) {\n const auto& scaling_mode = output->scaling_mode;\n NVTE_CHECK(scaling_mode == NVTE_MXFP8_1D_SCALING || scaling_mode == NVTE_NVFP4_1D_SCALING,\n \"Output tensor has invalid scaling mode (\", to_string(output->scaling_mode), \").\");\n\n CheckInputTensor(*input, \"scaling_factor_input\");\n CheckInputTensor(*output, \"scaling_factor_output\");\n NVTE_CHECK(input->with_gemm_swizzled_scales, \"Expected input tensor with swizzled scales.\");\n NVTE_CHECK(!output->with_gemm_swizzled_scales,\n \"Expected output tensor in row-major compact format.\");\n NVTE_CHECK(input->scaling_mode == scaling_mode,\n \"Input and output tensors must have matching scaling modes, but got \",\n to_string(input->scaling_mode), \" and \", to_string(output->scaling_mode), \".\");\n\n const bool has_rowwise_scale_inv = output->scale_inv.has_data();\n const bool has_columnwise_scale_inv = output->columnwise_scale_inv.has_data();\n NVTE_CHECK(!has_rowwise_scale_inv || !has_columnwise_scale_inv,\n \"Output tensor has both row-wise and column-wise scaling factors\");\n if (!has_rowwise_scale_inv && !has_columnwise_scale_inv) {\n return;\n }\n if (has_rowwise_scale_inv) {\n NVTE_CHECK(input->scale_inv.has_data(),\n \"Output tensor requests row-wise scaling factors, but input tensor does not \"\n \"provide them.\");\n } else if (has_columnwise_scale_inv) {\n NVTE_CHECK(input->columnwise_scale_inv.has_data(),\n \"Output tensor requests column-wise scaling factors, but input tensor does not \"\n \"provide them.\");\n }\n\n constexpr int SF_TILE_DIM_M = 128;\n constexpr int SF_TILE_DIM_K = 4;\n const dim3 block_size(TB_DIM, TB_DIM);\n\n int m{0}, k{0};\n void* input_ptr{nullptr};\n void* output_ptr{nullptr};\n bool rowwise{false};\n\n switch (scaling_mode) {\n case NVTE_MXFP8_1D_SCALING: {\n NVTE_CHECK(is_fp8_dtype(input->dtype()), \"Input tensor has invalid dtype (expected FP8, got \",\n to_string(input->dtype()), \").\");\n if (has_rowwise_scale_inv) {\n NVTE_CHECK(output->scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", output->scale_inv.shape, \".\");\n m = output->scale_inv.shape[0];\n k = output->scale_inv.shape[1];\n NVTE_CHECK(static_cast(m) * k == input->scale_inv.numel(),\n \"Expected input tensor to have \", static_cast(m) * k,\n \" row-wise scaling factors, but got shape=\", input->scale_inv.shape, \".\");\n NVTE_CHECK(static_cast(m) * k == output->scale_inv.numel(),\n \"Expected output tensor to have \", static_cast(m) * k,\n \" row-wise scaling factors, but got shape=\", output->scale_inv.shape, \".\");\n input_ptr = input->scale_inv.dptr;\n output_ptr = output->scale_inv.dptr;\n rowwise = true;\n } else if (has_columnwise_scale_inv) {\n NVTE_CHECK(output->columnwise_scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", output->columnwise_scale_inv.shape,\n \".\");\n m = output->columnwise_scale_inv.shape[1];\n k = output->columnwise_scale_inv.shape[0];\n NVTE_CHECK(\n static_cast(m) * k == input->columnwise_scale_inv.numel(),\n \"Expected input tensor to have \", static_cast(m) * k,\n \" column-wise scaling factors, but got shape=\", input->columnwise_scale_inv.shape, \".\");\n NVTE_CHECK(static_cast(m) * k == output->columnwise_scale_inv.numel(),\n \"Expected output tensor to have \", static_cast(m) * k,\n \" column-wise scaling factors, but got shape=\",\n output->columnwise_scale_inv.shape, \".\");\n input_ptr = input->columnwise_scale_inv.dptr;\n output_ptr = output->columnwise_scale_inv.dptr;\n rowwise = false;\n }\n break;\n }\n case NVTE_NVFP4_1D_SCALING: {\n NVTE_CHECK(is_fp4_dtype(input->dtype()), \"Input tensor has invalid dtype (expected FP4, got \",\n to_string(input->dtype()), \").\");\n // NVFP4: always unswizzle rowwise regardless of which scale buffer holds the data\n if (has_rowwise_scale_inv) {\n NVTE_CHECK(output->scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", output->scale_inv.shape, \".\");\n m = output->scale_inv.shape[0];\n k = output->scale_inv.shape[1];\n // Example for NVFP4 rowwise path:\n NVTE_CHECK(static_cast(m) * k == input->scale_inv.numel(),\n \"Expected input tensor to have \", static_cast(m) * k,\n \" row-wise scaling factors, but got shape=\", input->scale_inv.shape, \".\");\n NVTE_CHECK(static_cast(m) * k == output->scale_inv.numel(),\n \"Expected output tensor to have \", static_cast(m) * k,\n \" row-wise scaling factors, but got shape=\", output->scale_inv.shape, \".\");\n input_ptr = input->scale_inv.dptr;\n output_ptr = output->scale_inv.dptr;\n } else if (has_columnwise_scale_inv) {\n NVTE_CHECK(output->columnwise_scale_inv.shape.size() == 2,\n \"Expected 2D scaling factors, got shape=\", output->columnwise_scale_inv.shape,\n \".\");\n m = output->columnwise_scale_inv.shape[0];\n k = output->columnwise_scale_inv.shape[1];\n NVTE_CHECK(\n static_cast(m) * k == input->columnwise_scale_inv.numel(),\n \"Expected input tensor to have \", static_cast(m) * k,\n \" column-wise scaling factors, but got shape=\", input->columnwise_scale_inv.shape, \".\");\n NVTE_CHECK(static_cast(m) * k == output->columnwise_scale_inv.numel(),\n \"Expected output tensor to have \", static_cast(m) * k,\n \" column-wise scaling factors, but got shape=\",\n output->columnwise_scale_inv.shape, \".\");\n input_ptr = input->columnwise_scale_inv.dptr;\n output_ptr = output->columnwise_scale_inv.dptr;\n }\n rowwise = true;\n break;\n }\n default:\n NVTE_ERROR(\"Invalid scaling mode\");\n }\n\n NVTE_CHECK(m % SF_TILE_DIM_M == 0, \"Output should be padded in M/N dimension!\");\n NVTE_CHECK(k % SF_TILE_DIM_K == 0, \"Output should be padded in K dimension!\");\n\n const int num_tiles_m = m / SF_TILE_DIM_M;\n const int num_tiles_k = k / SF_TILE_DIM_K;\n\n auto launch_unswizzle = [&](int vec_load_size, const dim3& num_blocks, int slm_size) {\n switch (vec_load_size) {\n case 4:\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(unswizzle_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n unswizzle_scaling_kernel\n <<>>(input_ptr, output_ptr, m, k, rowwise);\n break;\n case 2:\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(unswizzle_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n unswizzle_scaling_kernel\n <<>>(input_ptr, output_ptr, m, k, rowwise);\n break;\n case 1:\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(unswizzle_scaling_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n unswizzle_scaling_kernel\n <<>>(input_ptr, output_ptr, m, k, rowwise);\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n };\n\n int vec_load_size = rowwise ? (num_tiles_k - 1) % 4 + 1 : (num_tiles_m - 1) % 4 + 1;\n if (vec_load_size == 3) vec_load_size = 1;\n int n_tiles_in_tb = TB_DIM * vec_load_size;\n dim3 num_blocks = rowwise ? dim3(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m)\n : dim3(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size));\n int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n launch_unswizzle(vec_load_size, num_blocks, slm_size);\n}\n\nvoid multi_tensor_unswizzle_scaling_factors(const std::vector& input,\n std::vector& output, hipStream_t stream) {\n size_t num_tensors = output.size();\n const auto& first_scaling_mode = output[0]->scaling_mode;\n\n constexpr int SF_TILE_DIM_M = 128;\n constexpr int SF_TILE_DIM_K = 4;\n\n bool all_has_data = true;\n bool all_has_columnwise_data = true;\n bool all_nvfp4 = true;\n for (size_t i = 0; i < num_tensors; i++) {\n const auto scaling_mode = output[i]->scaling_mode;\n const auto is_fp8 = is_fp8_dtype(input[i]->dtype());\n const auto is_fp4 = is_fp4_dtype(input[i]->dtype());\n\n NVTE_CHECK(scaling_mode == first_scaling_mode,\n \"All tensors should have the same scaling mode in multi-tensor unswizzle.\");\n NVTE_CHECK(\n (is_fp8 && is_mxfp8_scaling(scaling_mode)) || (is_fp4 && is_nvfp4_scaling(scaling_mode)),\n \"Not implemented scaling mode \" + to_string(scaling_mode) + \".\");\n NVTE_CHECK(input[i]->with_gemm_swizzled_scales,\n \"Expected input tensors with scales in GEMM swizzled format.\");\n NVTE_CHECK(!output[i]->with_gemm_swizzled_scales,\n \"Expected output tensors with scales in compact format.\");\n NVTE_CHECK(input[i]->numel() != 0, \"Tensor input[\", i, \"] is empty.\");\n CheckInputTensor(*input[i], \"scaling_factor_input[\" + std::to_string(i) + \"]\");\n CheckInputTensor(*output[i], \"scaling_factor_output[\" + std::to_string(i) + \"]\");\n\n all_has_data = all_has_data && output[i]->scale_inv.has_data();\n all_has_columnwise_data =\n (all_has_columnwise_data && output[i]->columnwise_scale_inv.has_data());\n all_nvfp4 = all_nvfp4 && is_nvfp4_scaling(scaling_mode);\n }\n NVTE_CHECK(all_has_data || all_has_columnwise_data,\n \"All tensors should have data or columnwise data.\");\n NVTE_CHECK(!all_has_data || !all_has_columnwise_data,\n \"All tensors have both data and columnwise data.\");\n\n const bool rowwise_unswizzle = all_has_data || all_nvfp4;\n const bool columnwise_unswizzle = all_has_columnwise_data && !all_nvfp4;\n\n if (rowwise_unswizzle) {\n MultiSwizzleArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.block_range[0] = 0;\n int vec_load_size = 4;\n for (size_t i = 0; i < num_tensors; i++) {\n if (kernel_args.num_tensors == kMaxTensorsPerKernel) {\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_unswizzle_scaling_factors(\n kernel_args, vec_load_size, true, stream);\n kernel_args.num_tensors = 0;\n vec_load_size = 4;\n }\n int m, k;\n if (all_has_data) {\n NVTE_CHECK(input[i]->scale_inv.has_data(), \"Input tensor \", i,\n \" does not have row-wise scaling factors.\");\n NVTE_CHECK(output[i]->scale_inv.shape.size() == 2, \"Expected output tensor \", i,\n \" to have \", \"2D scaling factors, got shape=\", output[i]->scale_inv.shape, \".\");\n m = output[i]->scale_inv.shape[0];\n k = output[i]->scale_inv.shape[1];\n NVTE_CHECK(m * k == input[i]->scale_inv.numel(), \"Expected input tensor \", i, \" to have \",\n m * k, \" row-wise scaling factors, but got shape=\", input[i]->scale_inv.shape,\n \".\");\n }\n\n if (all_has_columnwise_data) {\n NVTE_CHECK(all_nvfp4,\n \"When doing rowwise unswizzle with columnwise data, it has to be NVFP4\");\n NVTE_CHECK(input[i]->columnwise_scale_inv.has_data(), \"Input tensor \", i,\n \" does not have column-wise scaling factors.\");\n NVTE_CHECK(output[i]->columnwise_scale_inv.shape.size() == 2, \"Expected output tensor \", i,\n \" to have \",\n \"2D scaling factors, got shape=\", output[i]->columnwise_scale_inv.shape, \".\");\n m = output[i]->columnwise_scale_inv.shape[0];\n k = output[i]->columnwise_scale_inv.shape[1];\n NVTE_CHECK(m * k == input[i]->columnwise_scale_inv.numel(), \"Expected input tensor \", i,\n \" to have \", m * k, \" column-wise scaling factors, but got shape=\",\n input[i]->columnwise_scale_inv.shape, \".\");\n }\n\n NVTE_CHECK(m % SF_TILE_DIM_M == 0, \"Output should be padded in M/N dimension!\");\n NVTE_CHECK(k % SF_TILE_DIM_K == 0, \"Output should be padded in K dimension!\");\n NVTE_CHECK(k > 0, \"Output scale inverse should be 2D!\");\n\n int num_tiles_k = k / SF_TILE_DIM_K;\n int vec_load_size_i = (num_tiles_k - 1) % 4 + 1;\n vec_load_size = all_nvfp4 ? 1 : std::min(vec_load_size, vec_load_size_i);\n\n const int pos = kernel_args.num_tensors;\n kernel_args.m_list[pos] = m;\n kernel_args.k_list[pos] = k;\n if (!all_nvfp4 || all_has_data) {\n kernel_args.input_list[pos] = const_cast(input[i]->scale_inv.dptr);\n kernel_args.output_list[pos] = output[i]->scale_inv.dptr;\n } else {\n kernel_args.input_list[pos] = const_cast(input[i]->columnwise_scale_inv.dptr);\n kernel_args.output_list[pos] = output[i]->columnwise_scale_inv.dptr;\n }\n kernel_args.num_tensors++;\n }\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_unswizzle_scaling_factors(\n kernel_args, vec_load_size, true, stream);\n }\n\n if (columnwise_unswizzle) {\n NVTE_CHECK(!all_nvfp4, \"NVFP4 shouldn't end up here because it only needs rowwise unswizzle\");\n\n MultiSwizzleArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.block_range[0] = 0;\n int vec_load_size = 4;\n for (size_t i = 0; i < num_tensors; i++) {\n if (kernel_args.num_tensors == kMaxTensorsPerKernel) {\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_unswizzle_scaling_factors(\n kernel_args, vec_load_size, false, stream);\n kernel_args.num_tensors = 0;\n vec_load_size = 4;\n }\n NVTE_CHECK(output[i]->columnwise_scale_inv.shape.size() == 2, \"Expected output tensor \", i,\n \" to have \",\n \"2D scaling factors, got shape=\", output[i]->columnwise_scale_inv.shape, \".\");\n const int m = output[i]->columnwise_scale_inv.shape[1];\n const int k = output[i]->columnwise_scale_inv.shape[0];\n\n NVTE_CHECK(m % SF_TILE_DIM_M == 0, \"Output should be padded in M/N dimension!\");\n NVTE_CHECK(k % SF_TILE_DIM_K == 0, \"Output should be padded in K dimension!\");\n NVTE_CHECK(k > 0, \"Output scale inverse should be 2D!\");\n NVTE_CHECK(m * k == std::accumulate(input[i]->columnwise_scale_inv.shape.begin(),\n input[i]->columnwise_scale_inv.shape.end(), 1,\n std::multiplies()),\n \"Input.columnwise_scale_inv size is not equal to \"\n \"Output.columnwise_scale_inv size!\");\n\n int num_tiles_k = k / SF_TILE_DIM_K;\n int vec_load_size_i = (num_tiles_k - 1) % 4 + 1;\n vec_load_size = std::min(vec_load_size, vec_load_size_i);\n\n const int pos = kernel_args.num_tensors;\n kernel_args.input_list[pos] = const_cast(input[i]->columnwise_scale_inv.dptr);\n kernel_args.output_list[pos] = output[i]->columnwise_scale_inv.dptr;\n kernel_args.m_list[pos] = m;\n kernel_args.k_list[pos] = k;\n kernel_args.num_tensors++;\n }\n if (vec_load_size == 3) vec_load_size = 1;\n launch_multi_tensor_unswizzle_scaling_factors(\n kernel_args, vec_load_size, false, stream);\n }\n}\n} // namespace transformer_engine\n\n/*\n* WIP (Phuong):\n* - Opt for bank conflicts\n* - Adding swizzle for 2d-block scaling.\n*/\nvoid nvte_swizzle_scaling_factors(const NVTETensor input, NVTETensor output, hipStream_t stream) {\n NVTE_API_CALL(nvte_swizzle_scaling_factors);\n using namespace transformer_engine;\n swizzle_scaling_factors(convertNVTETensorCheck(input), convertNVTETensorCheck(output), stream);\n}\n\nvoid nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETensor* outputs,\n const size_t num_tensors, hipStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_swizzle_scaling_factors);\n using namespace transformer_engine;\n NVTE_CHECK(num_tensors > 0, \"Number of tensors should be greater than 0.\");\n std::vector input_list, output_list;\n for (size_t i = 0; i < num_tensors; i++) {\n input_list.push_back(convertNVTETensorCheck(inputs[i]));\n output_list.push_back(convertNVTETensorCheck(outputs[i]));\n }\n multi_tensor_swizzle_scaling_factors(input_list, output_list, stream,\n /*check_scale_inv_shapes=*/true);\n}\n\nvoid nvte_multi_tensor_swizzle_scaling_factors_unchecked(const NVTETensor* inputs,\n NVTETensor* outputs,\n const size_t num_tensors,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_swizzle_scaling_factors_unchecked);\n using namespace transformer_engine;\n NVTE_CHECK(num_tensors > 0, \"Number of tensors should be greater than 0.\");\n std::vector input_list, output_list;\n for (size_t i = 0; i < num_tensors; i++) {\n input_list.push_back(convertNVTETensorCheck(inputs[i]));\n output_list.push_back(convertNVTETensorCheck(outputs[i]));\n }\n multi_tensor_swizzle_scaling_factors(input_list, output_list, stream,\n /*check_scale_inv_shapes=*/false);\n}\n\nvoid nvte_unswizzle_scaling_factors(const NVTETensor input, NVTETensor output,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_unswizzle_scaling_factors);\n using namespace transformer_engine;\n unswizzle_scaling_factors(convertNVTETensorCheck(input), convertNVTETensorCheck(output), stream);\n}\n\nvoid nvte_multi_tensor_unswizzle_scaling_factors(const NVTETensor* inputs, NVTETensor* outputs,\n const size_t num_tensors, hipStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_unswizzle_scaling_factors);\n using namespace transformer_engine;\n NVTE_CHECK(num_tensors > 0, \"Number of tensors should be greater than 0.\");\n std::vector input_list, output_list;\n for (size_t i = 0; i < num_tensors; i++) {\n input_list.push_back(convertNVTETensorCheck(inputs[i]));\n output_list.push_back(convertNVTETensorCheck(outputs[i]));\n }\n multi_tensor_unswizzle_scaling_factors(input_list, output_list, stream);\n}\n\nnamespace transformer_engine {\n\ntemplate \n__global__ void __launch_bounds__(TB_DIM* TB_DIM)\n grouped_swizzle_scaling_variable_shape_kernel(const void* input, void* output,\n const int64_t* m_array, const int64_t* k_array,\n int num_tensors, bool rowwise,\n size_t scale_elem_size, size_t common_m,\n size_t common_k) {\n extern __shared__ int s_metadata[];\n int* s_total_blocks = &s_metadata[0];\n\n // Warp reduction to compute total workload\n if (threadIdx.x < 32 && threadIdx.y == 0) {\n int local_blocks = 0;\n for (int i = threadIdx.x; i < num_tensors; i += 32) {\n size_t m = rowwise ? (m_array ? m_array[i] : common_m) : (k_array ? k_array[i] : common_k);\n size_t k = rowwise ? (k_array ? k_array[i] : common_k) : (m_array ? m_array[i] : common_m);\n\n size_t padded_m = round_up_to_multiple(m, 128);\n size_t padded_k = round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4);\n\n int num_tiles_m = padded_m / SF_TILE_DIM_M;\n int num_tiles_k = padded_k / SF_TILE_DIM_K;\n\n int vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1));\n if (vec_load_size == 3) vec_load_size = 1;\n int n_tiles_in_tb = TB_DIM * vec_load_size;\n\n int grid_dim_x = rowwise ? DIVUP(num_tiles_k, n_tiles_in_tb) : DIVUP(num_tiles_k, TB_DIM);\n int grid_dim_y = rowwise ? num_tiles_m : DIVUP(num_tiles_m, vec_load_size);\n local_blocks += grid_dim_x * grid_dim_y;\n }\n\n for (int offset = 16; offset > 0; offset /= 2) {\n local_blocks += __shfl_down_sync(0xffffffff, local_blocks, offset);\n }\n if (threadIdx.x == 0) *s_total_blocks = local_blocks;\n }\n __syncthreads();\n\n const int total_blocks = *s_total_blocks;\n\n // Persistent-grid loop\n for (int linear_block_id = blockIdx.x; linear_block_id < total_blocks;\n linear_block_id += gridDim.x) {\n // Discover tensor_id and local_block_id via linear scan\n int tensor_id = 0;\n int current_block_base = 0;\n size_t current_scale_base = 0;\n int grid_dim_x = 0;\n int grid_dim_y = 0;\n size_t M = 0, K = 0;\n int vec_load_size = 0;\n\n for (int i = 0; i < num_tensors; ++i) {\n M = rowwise ? (m_array ? m_array[i] : common_m) : (k_array ? k_array[i] : common_k);\n K = rowwise ? (k_array ? k_array[i] : common_k) : (m_array ? m_array[i] : common_m);\n\n size_t padded_m = round_up_to_multiple(M, 128);\n size_t padded_k = round_up_to_multiple(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4);\n\n int num_tiles_m = padded_m / SF_TILE_DIM_M;\n int num_tiles_k = padded_k / SF_TILE_DIM_K;\n\n vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1));\n if (vec_load_size == 3) vec_load_size = 1;\n int n_tiles_in_tb = TB_DIM * vec_load_size;\n\n grid_dim_x = rowwise ? DIVUP(num_tiles_k, n_tiles_in_tb) : DIVUP(num_tiles_k, TB_DIM);\n grid_dim_y = rowwise ? num_tiles_m : DIVUP(num_tiles_m, vec_load_size);\n int blocks_i = grid_dim_x * grid_dim_y;\n\n if (linear_block_id < current_block_base + blocks_i) {\n tensor_id = i;\n break;\n }\n current_block_base += blocks_i;\n current_scale_base += padded_m * padded_k * scale_elem_size;\n }\n\n int local_block_id = linear_block_id - current_block_base;\n int block_x = local_block_id % grid_dim_x;\n int block_y = local_block_id / grid_dim_x;\n\n const uint8_t* input_base = reinterpret_cast(input) + current_scale_base;\n uint8_t* output_base = reinterpret_cast(output) + current_scale_base;\n\n const int padded_m = static_cast(round_up_to_multiple(M, 128));\n const int padded_k =\n static_cast(round_up_to_multiple(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4));\n const int original_M = static_cast(M);\n const int original_K = static_cast(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)));\n const bool padding_m = (block_y == grid_dim_y - 1) && (original_M < padded_m);\n const bool padding_k = (block_x == grid_dim_x - 1) && (original_K < padded_k);\n\n if (rowwise) {\n if (vec_load_size == 4) {\n dispatch_swizzle_row_scaling_kernel_impl(\n input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y,\n grid_dim_x, grid_dim_y, padding_k, padding_m);\n } else if (vec_load_size == 2) {\n dispatch_swizzle_row_scaling_kernel_impl(\n input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y,\n grid_dim_x, grid_dim_y, padding_k, padding_m);\n } else {\n dispatch_swizzle_row_scaling_kernel_impl(\n input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y,\n grid_dim_x, grid_dim_y, padding_k, padding_m);\n }\n } else {\n if (vec_load_size == 4) {\n dispatch_swizzle_col_scaling_kernel_impl(\n input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y,\n grid_dim_x, grid_dim_y, padding_k, padding_m);\n } else if (vec_load_size == 2) {\n dispatch_swizzle_col_scaling_kernel_impl(\n input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y,\n grid_dim_x, grid_dim_y, padding_k, padding_m);\n } else {\n dispatch_swizzle_col_scaling_kernel_impl(\n input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y,\n grid_dim_x, grid_dim_y, padding_k, padding_m);\n }\n }\n }\n}\n\ntemplate \nint grouped_swizzle_variable_max_active_blocks_per_sm(int device_id) {\n static std::vector cache(cuda::num_devices(), -1);\n static std::vector flags(cuda::num_devices());\n NVTE_CHECK(0 <= device_id && device_id < cuda::num_devices(), \"invalid CUDA device ID\");\n\n auto init = [&]() {\n constexpr int metadata_shmem = sizeof(int); // s_total_blocks\n constexpr int dynamic_smem_size =\n TB_DIM * 4 * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t) + metadata_shmem;\n int max_active_blocks_per_sm;\n NVTE_CHECK_CUDA(hipOccupancyMaxActiveBlocksPerMultiprocessor(\n &max_active_blocks_per_sm,\n grouped_swizzle_scaling_variable_shape_kernel,\n TB_DIM * TB_DIM, dynamic_smem_size));\n NVTE_CHECK(max_active_blocks_per_sm > 0, \"Occupancy query returned 0 blocks per SM.\");\n cache[device_id] = max_active_blocks_per_sm;\n };\n std::call_once(flags[device_id], init);\n return cache[device_id];\n}\n\nvoid swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* output,\n hipStream_t stream) {\n // Check scaling mode\n NVTE_CHECK(input->scaling_mode == NVTE_MXFP8_1D_SCALING,\n \"Grouped swizzle supports only MXFP8 scaling.\");\n\n // Check tensors\n CheckInputGroupedTensor(*input, \"input\");\n CheckOutputGroupedTensor(*output, \"output\", false);\n NVTE_CHECK(!input->with_gemm_swizzled_scales,\n \"Expected input grouped tensor with scales in compact format.\");\n NVTE_CHECK(output->with_gemm_swizzled_scales,\n \"Expected output grouped tensor with scales in GEMM swizzled format.\");\n\n // Check scaling factors availability\n const bool has_rowwise_scale_inv = input->scale_inv.has_data();\n const bool has_columnwise_scale_inv = input->columnwise_scale_inv.has_data();\n if (!has_rowwise_scale_inv && !has_columnwise_scale_inv) {\n return;\n }\n\n const int64_t* m_array = reinterpret_cast(input->first_dims.dptr);\n const int64_t* k_array = reinterpret_cast(input->last_dims.dptr);\n const bool is_variable_shape = !input->all_same_shape();\n\n if (!is_variable_shape) {\n // Fallback to uniform shape implementation\n // Assumption is that all the tensors share the same shapes and are contgiuous.\n // And so we dont need to pass array of input/output pointers(due to conttiguity)\n // as well as array of shapes(due to uniform shapes).\n const size_t first_dim = input->get_common_first_dim();\n const size_t last_dim = input->get_common_last_dim();\n\n constexpr int SF_TILE_DIM_M = 128;\n constexpr int SF_TILE_DIM_K = 4;\n const dim3 block_size(TB_DIM, TB_DIM);\n\n auto launch_grouped_swizzle = [&](bool rowwise) {\n const size_t m = rowwise ? first_dim : last_dim;\n const size_t k = rowwise ? last_dim : first_dim;\n const size_t padded_m = round_up_to_multiple(m, 128);\n const size_t padded_k =\n round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4);\n // Per-tensor scale-element counts:\n // - \"padded\" layout: each tensor occupies padded_m * padded_k elements\n // (total buffer = num_tensors * padded_m * padded_k).\n // - \"compact\" layout (what the grouped MXFP8 quantize kernel actually writes):\n // per-tensor stride is m * padded_k (rowwise) or DIVUP(k,32) * padded_m\n // (columnwise) and the total buffer the C++ allocator hands out has its\n // grouped first dim padded up to a multiple of 128 (rowwise) or 4\n // (columnwise) -- so the buffer may be slightly larger than\n // num_tensors * compact_scale_elems, with trailing alignment slack at\n // the very end (never read because of the per-tensor row/k guard in the\n // kernel impl).\n // The output is always written in the padded layout. The input may be in\n // either layout; the kernel handles the compact case safely by using\n // different per-tensor strides for input vs output and skipping loads past\n // the per-tensor extent.\n const size_t padded_scale_elems = padded_m * padded_k;\n const size_t compact_scale_elems =\n rowwise ? m * padded_k : DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)) * padded_m;\n const size_t compact_total_scale_elems =\n rowwise ? round_up_to_multiple(input->num_tensors * m, 128) * padded_k\n : round_up_to_multiple(\n input->num_tensors * DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4) *\n padded_m;\n\n const size_t scale_elem_size = rowwise ? typeToSize(input->scale_inv.dtype)\n : typeToSize(input->columnwise_scale_inv.dtype);\n\n const size_t input_scale_numel =\n rowwise ? input->scale_inv.numel() : input->columnwise_scale_inv.numel();\n const size_t output_scale_numel =\n rowwise ? output->scale_inv.numel() : output->columnwise_scale_inv.numel();\n\n bool input_is_compact;\n if (input_scale_numel == input->num_tensors * padded_scale_elems) {\n input_is_compact = false;\n } else if (input_scale_numel == compact_total_scale_elems) {\n input_is_compact = true;\n } else {\n NVTE_ERROR(\"Grouped input \", (rowwise ? \"scale_inv\" : \"columnwise_scale_inv\"),\n \" size does not match expected packed size (got \", input_scale_numel,\n \", expected either \", input->num_tensors * padded_scale_elems,\n \" (per-tensor padded) or \", compact_total_scale_elems, \" (compact)).\");\n }\n NVTE_CHECK(output_scale_numel == input->num_tensors * padded_scale_elems, \"Grouped output \",\n (rowwise ? \"scale_inv\" : \"columnwise_scale_inv\"),\n \" size does not match expected per-tensor padded size.\");\n\n const size_t input_stride_bytes =\n (input_is_compact ? compact_scale_elems : padded_scale_elems) * scale_elem_size;\n const size_t output_stride_bytes = padded_scale_elems * scale_elem_size;\n\n const int num_tiles_m = padded_m / SF_TILE_DIM_M;\n const int num_tiles_k = padded_k / SF_TILE_DIM_K;\n int vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1));\n if (vec_load_size == 3) vec_load_size = 1;\n const int n_tiles_in_tb = TB_DIM * vec_load_size;\n\n dim3 num_blocks;\n if (rowwise) {\n num_blocks = dim3(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m, input->num_tensors);\n } else {\n num_blocks =\n dim3(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size), input->num_tensors);\n }\n const int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n\n const int original_M = static_cast(rowwise ? first_dim : last_dim);\n const int original_K = static_cast(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)));\n const void* input_ptr = rowwise ? input->scale_inv.dptr : input->columnwise_scale_inv.dptr;\n void* output_ptr = rowwise ? output->scale_inv.dptr : output->columnwise_scale_inv.dptr;\n\n if (rowwise) {\n TRANSFORMER_ENGINE_VECTORIZED_LOAD_INTEGER_TYPE_SWITCH(vec_load_size, LType, {\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(grouped_swizzle_row_scaling_uniform_shape_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n grouped_swizzle_row_scaling_uniform_shape_kernel\n <<>>(\n input_ptr, output_ptr, padded_m, padded_k, original_M, original_K,\n input_stride_bytes, output_stride_bytes);\n });\n } else {\n TRANSFORMER_ENGINE_VECTORIZED_LOAD_INTEGER_TYPE_SWITCH(vec_load_size, LType, {\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(grouped_swizzle_col_scaling_uniform_shape_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n grouped_swizzle_col_scaling_uniform_shape_kernel\n <<>>(\n input_ptr, output_ptr, padded_m, padded_k, original_M, original_K,\n input_stride_bytes, output_stride_bytes);\n });\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n };\n\n if (has_rowwise_scale_inv) {\n launch_grouped_swizzle(true);\n }\n if (has_columnwise_scale_inv) {\n launch_grouped_swizzle(false);\n }\n } else {\n // Variable shape implementation using Device-Side Block Scheduler\n size_t num_tensors = input->num_tensors;\n\n constexpr int SF_TILE_DIM_M = 128;\n constexpr int SF_TILE_DIM_K = 4;\n const dim3 block_size(TB_DIM, TB_DIM);\n const int max_slm_size = TB_DIM * 4 * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n const int metadata_shmem = sizeof(int); // s_total_blocks\n const int dynamic_smem_size = max_slm_size + metadata_shmem;\n\n size_t common_m = input->all_same_first_dim() ? input->get_common_first_dim() : 0;\n size_t common_k = input->all_same_last_dim() ? input->get_common_last_dim() : 0;\n\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(grouped_swizzle_scaling_variable_shape_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, dynamic_smem_size));\n\n const int device_id = cuda::current_device();\n const int num_SMs = cuda::sm_count(device_id);\n const int max_active_blocks_per_sm =\n grouped_swizzle_variable_max_active_blocks_per_sm(device_id);\n const int persistent_blocks = num_SMs * max_active_blocks_per_sm;\n const dim3 num_blocks(persistent_blocks);\n\n auto launch_grouped_swizzle_variable = [&](bool rowwise) {\n const size_t scale_elem_size = rowwise ? typeToSize(input->scale_inv.dtype)\n : typeToSize(input->columnwise_scale_inv.dtype);\n\n const void* input_ptr = rowwise ? input->scale_inv.dptr : input->columnwise_scale_inv.dptr;\n void* output_ptr = rowwise ? output->scale_inv.dptr : output->columnwise_scale_inv.dptr;\n\n grouped_swizzle_scaling_variable_shape_kernel\n <<>>(\n input_ptr, output_ptr, m_array, k_array, num_tensors, rowwise, scale_elem_size,\n common_m, common_k);\n\n NVTE_CHECK_CUDA(hipGetLastError());\n };\n\n if (has_rowwise_scale_inv) {\n launch_grouped_swizzle_variable(true);\n }\n if (has_columnwise_scale_inv) {\n launch_grouped_swizzle_variable(false);\n }\n }\n}\n\nvoid unswizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* output,\n hipStream_t stream) {\n NVTE_CHECK(output->scaling_mode == NVTE_MXFP8_1D_SCALING,\n \"Grouped unswizzle supports only MXFP8 scaling.\");\n\n CheckInputGroupedTensor(*input, \"input\");\n CheckOutputGroupedTensor(*output, \"output\", false);\n NVTE_CHECK(input->with_gemm_swizzled_scales,\n \"Expected input grouped tensor with scales in GEMM swizzled format.\");\n NVTE_CHECK(!output->with_gemm_swizzled_scales,\n \"Expected output grouped tensor with scales in compact format.\");\n NVTE_CHECK(input->scaling_mode == output->scaling_mode,\n \"Input and output grouped tensors must have matching scaling modes.\");\n NVTE_CHECK(input->num_tensors == output->num_tensors,\n \"Input and output grouped tensors must have the same number of tensors.\");\n\n const bool has_rowwise_scale_inv = output->scale_inv.has_data();\n const bool has_columnwise_scale_inv = output->columnwise_scale_inv.has_data();\n if (!has_rowwise_scale_inv && !has_columnwise_scale_inv) {\n return;\n }\n\n NVTE_CHECK(input->all_same_shape() && output->all_same_shape(),\n \"Grouped unswizzle requires uniform tensor shapes.\");\n\n const size_t first_dim = output->get_common_first_dim();\n const size_t last_dim = output->get_common_last_dim();\n\n constexpr int SF_TILE_DIM_M = 128;\n constexpr int SF_TILE_DIM_K = 4;\n const dim3 block_size(TB_DIM, TB_DIM);\n\n auto launch_grouped_unswizzle = [&](bool rowwise) {\n const size_t m = rowwise ? first_dim : last_dim;\n const size_t k = rowwise ? last_dim : first_dim;\n const size_t padded_m = round_up_to_multiple(m, 128);\n const size_t padded_k =\n round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4);\n const size_t scale_elems = padded_m * padded_k;\n\n const size_t scale_elem_size = rowwise ? typeToSize(output->scale_inv.dtype)\n : typeToSize(output->columnwise_scale_inv.dtype);\n const size_t scale_stride_bytes = scale_elems * scale_elem_size;\n\n if (rowwise) {\n NVTE_CHECK(input->scale_inv.numel() == input->num_tensors * scale_elems,\n \"Grouped input scale_inv size does not match expected packed size.\");\n NVTE_CHECK(output->scale_inv.numel() == output->num_tensors * scale_elems,\n \"Grouped output scale_inv size does not match expected packed size.\");\n } else {\n NVTE_CHECK(input->columnwise_scale_inv.numel() == input->num_tensors * scale_elems,\n \"Grouped input columnwise_scale_inv size does not match expected packed size.\");\n NVTE_CHECK(output->columnwise_scale_inv.numel() == output->num_tensors * scale_elems,\n \"Grouped output columnwise_scale_inv size does not match expected packed size.\");\n }\n\n const int num_tiles_m = padded_m / SF_TILE_DIM_M;\n const int num_tiles_k = padded_k / SF_TILE_DIM_K;\n int vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1));\n if (vec_load_size == 3) vec_load_size = 1;\n const int n_tiles_in_tb = TB_DIM * vec_load_size;\n\n dim3 num_blocks;\n if (rowwise) {\n num_blocks = dim3(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m, output->num_tensors);\n } else {\n num_blocks =\n dim3(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size), output->num_tensors);\n }\n const int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t);\n\n const void* input_ptr = rowwise ? input->scale_inv.dptr : input->columnwise_scale_inv.dptr;\n void* output_ptr = rowwise ? output->scale_inv.dptr : output->columnwise_scale_inv.dptr;\n\n using kernel_t = void (*)(const void*, void*, const int, const int, const size_t, const bool);\n kernel_t kernel_fn = nullptr;\n switch (vec_load_size) {\n case 4:\n kernel_fn =\n grouped_unswizzle_scaling_uniform_shape_kernel;\n break;\n case 2:\n kernel_fn =\n grouped_unswizzle_scaling_uniform_shape_kernel;\n break;\n case 1:\n kernel_fn =\n grouped_unswizzle_scaling_uniform_shape_kernel;\n break;\n default:\n NVTE_ERROR(\"Not valid vec_load_size.\");\n }\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(kernel_fn), hipFuncAttributeMaxDynamicSharedMemorySize, slm_size));\n kernel_fn<<>>(input_ptr, output_ptr, padded_m,\n padded_k, scale_stride_bytes, rowwise);\n NVTE_CHECK_CUDA(hipGetLastError());\n };\n\n if (has_rowwise_scale_inv) {\n launch_grouped_unswizzle(true);\n }\n if (has_columnwise_scale_inv) {\n launch_grouped_unswizzle(false);\n }\n}\n\n} // namespace transformer_engine\n\nvoid nvte_swizzle_grouped_scaling_factors(const NVTEGroupedTensor input, NVTEGroupedTensor output,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_swizzle_grouped_scaling_factors);\n using namespace transformer_engine;\n swizzle_grouped_scaling_factors(convertNVTEGroupedTensorCheck(input),\n convertNVTEGroupedTensorCheck(output), stream);\n}\n\nvoid nvte_unswizzle_grouped_scaling_factors(const NVTEGroupedTensor input, NVTEGroupedTensor output,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_unswizzle_grouped_scaling_factors);\n using namespace transformer_engine;\n unswizzle_grouped_scaling_factors(convertNVTEGroupedTensorCheck(input),\n convertNVTEGroupedTensorCheck(output), stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/swizzle/swizzle.cu"} {"input": "/*************************************************************************\n* Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n*\n* See LICENSE for license information.\n************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"../cast/mxfp8/swizzle.cuh\"\n#include \"../common.h\"\n#include \"../util/cuda_runtime.h\"\n#include \"../util/handle_manager.h\"\n#include \"../util/logging.h\"\n#include \"../util/vectorized_pointwise.h\"\n#include \"./config.h\"\n\nnamespace {\n\ninline void CreateCublasHandle(cublasLtHandle_t *handle) {\n NVTE_CHECK_CUBLAS(cublasLtCreate(handle));\n}\n\n} // namespace\n\n// MXFP8 support for grouped GEMM requires cuBLAS 13.3+\n#define CUBLAS_MXFP8_GROUPED_GEMM_VERSION 130300\n// BF16 support for grouped GEMM requires cuBLAS 13.3+\n#define CUBLAS_GROUPED_GEMM_VERSION 130300\n\n#if CUBLAS_VERSION >= CUBLAS_GROUPED_GEMM_VERSION\n\nnamespace {\n\n// Helper struct to pass per-tensor shape/offset info (pointer or uniform value)\nstruct TensorShapeInfo {\n const int64_t *first_dims; // nullptr if uniform\n const int64_t *last_dims; // nullptr if uniform\n const int64_t *offsets; // nullptr if need to compute\n int64_t uniform_first; // used if first_dims == nullptr\n int64_t uniform_last; // used if last_dims == nullptr\n\n // Create from GroupedTensor\n static TensorShapeInfo from_tensor(const transformer_engine::GroupedTensor *t) {\n const bool has_first = t->first_dims.has_data();\n const bool has_last = t->last_dims.has_data();\n // When per-tensor dims are not provided, we must be in the uniform-shape case.\n NVTE_CHECK(has_first || t->all_same_first_dim(),\n \"GroupedTensor is missing first_dims for varying shapes\");\n NVTE_CHECK(has_last || t->all_same_last_dim(),\n \"GroupedTensor is missing last_dims for varying shapes\");\n\n const int64_t *first_ptr =\n has_first ? static_cast(t->first_dims.dptr) : nullptr;\n const int64_t *last_ptr = has_last ? static_cast(t->last_dims.dptr) : nullptr;\n\n const int64_t uniform_first = has_first ? 0 : static_cast(t->get_common_first_dim());\n const int64_t uniform_last = has_last ? 0 : static_cast(t->get_common_last_dim());\n\n return {first_ptr, last_ptr,\n t->tensor_offsets.has_data() ? static_cast(t->tensor_offsets.dptr)\n : nullptr,\n uniform_first, uniform_last};\n }\n\n // Create for C tensor (uses D's dimensions, only has offsets)\n static TensorShapeInfo create_shape_info_for_C(const transformer_engine::GroupedTensor *C,\n const transformer_engine::GroupedTensor *D) {\n const bool has_first = D->first_dims.has_data();\n const bool has_last = D->last_dims.has_data();\n NVTE_CHECK(has_first || D->all_same_first_dim(),\n \"GroupedTensor D is missing first_dims for varying shapes\");\n NVTE_CHECK(has_last || D->all_same_last_dim(),\n \"GroupedTensor D is missing last_dims for varying shapes\");\n\n const int64_t *first_ptr =\n has_first ? static_cast(D->first_dims.dptr) : nullptr;\n const int64_t *last_ptr = has_last ? static_cast(D->last_dims.dptr) : nullptr;\n const int64_t uniform_first = has_first ? 0 : static_cast(D->get_common_first_dim());\n const int64_t uniform_last = has_last ? 0 : static_cast(D->get_common_last_dim());\n\n return {first_ptr, last_ptr,\n C->tensor_offsets.has_data() ? static_cast(C->tensor_offsets.dptr)\n : nullptr,\n uniform_first, uniform_last};\n }\n};\n\n// Helper functions to compute average dimensions for cuBLASLt algorithm-selection heuristics.\n//\n// logical_shape encoding (from build_grouped_tensor):\n// all_same: {num_tensors * M, N}\n// varying_first: {sum_of_first_dims, common_last}\n// varying_last: {common_first, sum_of_last_dims}\n// varying_both: {1, total_elements} <-- lossy, can't recover per-dim averages\n//\n// We use all_same_first/last_dim() + get_common_first/last_dim() to get exact\n// answers whenever possible, falling back to logical_shape division otherwise.\n// For varying_both, per-dim averages are unrecoverable without a D2H copy,\n// so we return 1 \u2014 a valid non-zero hint that won't skip work.\ninline int64_t compute_avg_first_dim(const transformer_engine::GroupedTensor *t) {\n if (t->all_same_first_dim()) {\n return static_cast(t->get_common_first_dim());\n }\n const int64_t n = static_cast(t->num_tensors);\n if (t->all_same_last_dim()) {\n // varying_first only: logical_shape = {sum_of_first_dims, common_last}\n return static_cast(t->logical_shape.data[0]) / n;\n }\n // varying_both: logical_shape = {1, total_elements}, no way to recover avg first dim\n return 1;\n}\n\ninline int64_t compute_avg_last_dim(const transformer_engine::GroupedTensor *t) {\n if (t->all_same_last_dim()) {\n // logical_shape[1] is the common N\n return static_cast(t->logical_shape.data[1]);\n }\n // When varying, logical_shape[1] should be sum of last dims if provided; otherwise fallback to avg via division.\n return static_cast(t->logical_shape.data[1]) / static_cast(t->num_tensors);\n}\n\n// Constants for grouped GEMM workspace (declared early for use in helpers)\nstatic constexpr size_t kGroupedGemmAlignment = 256;\nstatic constexpr size_t kGroupedGemmCublasWorkspaceSize = 32ull * 1024 * 1024; // 32 MiB\n\n// Workspace layout for grouped GEMM\nstruct GroupedGemmSetupWorkspace {\n void **A_ptrs;\n void **B_ptrs;\n void **C_ptrs;\n void **D_ptrs;\n float **alpha_ptrs;\n float **beta_ptrs;\n void **\n a_scale_inv_ptrs; // Per-tensor FP8 scale pointers for A (float* for tensor scaling, E8M0* for MXFP8)\n void **\n b_scale_inv_ptrs; // Per-tensor FP8 scale pointers for B (float* for tensor scaling, E8M0* for MXFP8)\n // Storage dimensions for cuBLAS matrix layouts\n int *a_rows;\n int *a_cols;\n int *b_rows;\n int *b_cols;\n int *d_rows; // M (first dim) - also used for C\n int *d_cols; // N (last dim) - also used for C\n\n // Initialize from workspace buffer\n // Layout: all pointer arrays first (16-byte aligned for cuBLAS), then int arrays\n static GroupedGemmSetupWorkspace from_buffers(char *setup_ws_ptr, size_t num_tensors) {\n GroupedGemmSetupWorkspace ws;\n size_t offset = 0;\n const size_t ptr_size = num_tensors * sizeof(void *);\n const size_t int_size = num_tensors * sizeof(int);\n constexpr size_t kPtrAlignment = 16; // cuBLAS requires 16-byte alignment for pointer arrays\n\n // Helper to align offset to kPtrAlignment\n auto align_offset = [&]() {\n offset = (offset + kPtrAlignment - 1) / kPtrAlignment * kPtrAlignment;\n };\n\n // Pointer arrays first (all 16-byte aligned for cuBLAS grouped GEMM)\n align_offset();\n ws.A_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n align_offset();\n ws.B_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n align_offset();\n ws.C_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n align_offset();\n ws.D_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n align_offset();\n ws.alpha_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n align_offset();\n ws.beta_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n align_offset();\n ws.a_scale_inv_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n align_offset();\n ws.b_scale_inv_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n\n // Int arrays for storage dimensions (4-byte aligned is fine)\n align_offset();\n ws.a_rows = reinterpret_cast(setup_ws_ptr + offset);\n offset += int_size;\n ws.a_cols = reinterpret_cast(setup_ws_ptr + offset);\n offset += int_size;\n ws.b_rows = reinterpret_cast(setup_ws_ptr + offset);\n offset += int_size;\n ws.b_cols = reinterpret_cast(setup_ws_ptr + offset);\n offset += int_size;\n ws.d_rows = reinterpret_cast(setup_ws_ptr + offset);\n offset += int_size;\n ws.d_cols = reinterpret_cast(setup_ws_ptr + offset);\n\n return ws;\n }\n\n // Calculate required size for setup workspace\n static size_t required_setup_size(size_t num_tensors, size_t alignment) {\n const size_t ptr_size = num_tensors * sizeof(void *);\n const size_t int_size = num_tensors * sizeof(int);\n constexpr size_t kPtrAlignment = 16; // Must match from_buffers\n\n // Layout: 8 ptr arrays (each 16-byte aligned), then 6 int arrays\n // Each ptr array takes ptr_size bytes but needs to start at 16-byte boundary\n auto aligned_ptr_size = ((ptr_size + kPtrAlignment - 1) / kPtrAlignment) * kPtrAlignment;\n size_t size = 8 * aligned_ptr_size + 6 * int_size;\n size = ((size + alignment - 1) / alignment) * alignment;\n return size;\n }\n};\n\ninline size_t validate_grouped_gemm_inputs(\n size_t num_tensors, std::initializer_list inputs,\n const transformer_engine::Tensor *alpha_tensor, const transformer_engine::Tensor *beta_tensor) {\n NVTE_CHECK(num_tensors >= 1, \"Grouped GEMM: number of tensors must be at least 1\");\n for (const auto *tensor : inputs) {\n NVTE_CHECK(tensor->num_tensors == num_tensors,\n \"Grouped GEMM: inputs must have the same number of tensors\");\n }\n\n const size_t alpha_numel = alpha_tensor->data.numel();\n const size_t beta_numel = beta_tensor->data.numel();\n NVTE_CHECK(alpha_numel == num_tensors, \"Grouped GEMM: alpha must have num_tensors (\", num_tensors,\n \") elements, got \", alpha_numel);\n NVTE_CHECK(beta_numel == num_tensors, \"Grouped GEMM: beta must have num_tensors (\", num_tensors,\n \") elements, got \", beta_numel);\n\n auto is_supported_input_dtype = [](transformer_engine::DType dtype) {\n return dtype == transformer_engine::DType::kFloat8E4M3 ||\n dtype == transformer_engine::DType::kFloat8E5M2 ||\n dtype == transformer_engine::DType::kBFloat16 ||\n dtype == transformer_engine::DType::kFloat16;\n };\n for (const auto *tensor : inputs) {\n if (tensor->has_data() || tensor->has_columnwise_data()) {\n NVTE_CHECK(is_supported_input_dtype(tensor->dtype()),\n \"Grouped GEMM inputs must be FP8, BF16, or FP16, got \",\n transformer_engine::to_string(tensor->dtype()), \".\");\n }\n }\n // Cross-operand consistency across all inputs (skip tensors without data).\n const transformer_engine::GroupedTensor *ref = nullptr;\n for (const auto *tensor : inputs) {\n if (tensor->has_data() || tensor->has_columnwise_data()) {\n ref = tensor;\n break;\n }\n }\n if (ref != nullptr) {\n const bool ref_is_fp8 = is_fp8_dtype(ref->dtype());\n const bool ref_is_mxfp8 = transformer_engine::is_mxfp_scaling(ref->scaling_mode);\n for (const auto *tensor : inputs) {\n if (!(tensor->has_data() || tensor->has_columnwise_data())) continue;\n NVTE_CHECK(is_fp8_dtype(tensor->dtype()) == ref_is_fp8,\n \"Grouped GEMM: A and B must both be FP8 or both be non-FP8.\");\n NVTE_CHECK(transformer_engine::is_mxfp_scaling(tensor->scaling_mode) == ref_is_mxfp8,\n \"Grouped GEMM: A and B must both use MXFP8 scaling or both use tensor scaling.\");\n if (ref_is_mxfp8) {\n NVTE_CHECK(tensor->with_gemm_swizzled_scales,\n \"MXFP8 grouped GEMM: scales must be swizzled for GEMM.\");\n }\n }\n }\n return num_tensors;\n}\n\ninline void validate_grouped_gemm_outputs(\n size_t num_tensors, std::initializer_list outputs) {\n auto is_output_dtype = [](transformer_engine::DType dtype) {\n return dtype == transformer_engine::DType::kBFloat16 ||\n dtype == transformer_engine::DType::kFloat16 ||\n dtype == transformer_engine::DType::kFloat32;\n };\n for (const auto *tensor : outputs) {\n if (tensor == nullptr) {\n continue;\n }\n NVTE_CHECK(tensor->num_tensors == num_tensors,\n \"Grouped GEMM: outputs must have the same number of tensors as inputs\");\n NVTE_CHECK(is_output_dtype(tensor->dtype()),\n \"Grouped GEMM: outputs must be BF16, FP16, or FP32.\");\n }\n}\n\ninline size_t grouped_gemm_setup_workspace_size(size_t num_tensors) {\n return GroupedGemmSetupWorkspace::required_setup_size(num_tensors, kGroupedGemmAlignment);\n}\n\ninline void check_grouped_gemm_requirements(const char *api_name) {\n const int current_device = transformer_engine::cuda::current_device();\n NVTE_CHECK(transformer_engine::cuda::sm_arch(current_device) >= 100, api_name,\n \" requires Blackwell (SM100) or newer architecture.\");\n NVTE_CHECK(transformer_engine::cuda::cublas_version() >= CUBLAS_GROUPED_GEMM_VERSION, api_name,\n \" requires cuBLAS 13.3+, but run-time cuBLAS version is \",\n transformer_engine::cuda::cublas_version());\n}\n\ninline transformer_engine::GroupedMatmulConfig parse_grouped_gemm_config(\n NVTEGroupedMatmulConfig config) {\n transformer_engine::GroupedMatmulConfig config_;\n if (config != nullptr) {\n config_ = *reinterpret_cast(config);\n }\n return config_;\n}\n\n// Select row-wise vs column-wise storage and adjust transpose flag for grouped GEMM.\n// Mirrors the non-grouped GEMM logic for FP8 layout handling (TN-only on Hopper) and\n// fallback to column-wise data when row-wise is absent.\n// Contains all information needed for GEMM setup - shape already accounts for storage layout.\nstruct GroupedOperandSelection {\n TensorShapeInfo shape; // Shape info with dims already swapped for columnwise if needed\n char *dptr = nullptr;\n void *scale_inv = nullptr; // Contiguous array of scales (input)\n transformer_engine::DType dtype = transformer_engine::DType::kNumTypes;\n NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING;\n bool with_gemm_swizzled_scales = false;\n bool trans = false;\n bool rowwise = true;\n};\n\nconstexpr int kMaxGroups = 64;\n// Arguments for the grouped GEMM kernel that operates on multiple output tensors.\nstruct MultiTensorGroupGemmOutputArgs {\n void *data_ptrs[kMaxGroups];\n int rows[kMaxGroups];\n int cols[kMaxGroups];\n};\n\n// Arguments for the grouped GEMM kernel that operates on multiple inputA tensors.\nstruct MultiTensorGroupGemmInputArgs {\n void *data_ptrs[kMaxGroups];\n void *scale_inv_ptrs[kMaxGroups];\n int rows[kMaxGroups];\n int cols[kMaxGroups];\n};\nstruct MultiTensorListInfo {\n bool all_row = true;\n bool all_col = true;\n transformer_engine::DType row_dtype = transformer_engine::DType::kNumTypes;\n transformer_engine::DType col_dtype = transformer_engine::DType::kNumTypes;\n NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING;\n bool with_gemm_swizzled_scales = false;\n};\n\nstruct OperandStorageChoice {\n bool use_rowwise = true;\n bool swap_dims = true;\n bool trans = false;\n};\n\ninline OperandStorageChoice choose_grouped_operand_storage(bool trans, bool is_A, bool is_mxfp8,\n bool is_fp8, bool non_tn_fp8_ok,\n bool has_row, bool has_col,\n const char *name) {\n NVTE_CHECK(has_row || has_col, \"Grouped GEMM: \", name,\n \" is missing both row-wise and column-wise data\");\n if (is_mxfp8) {\n if (is_A) {\n if (trans) {\n NVTE_CHECK(has_row, \"Grouped GEMM: MXFP8 transposed \", name, \" is missing row-wise data\");\n return {true, true, trans};\n }\n NVTE_CHECK(has_col, \"Grouped GEMM: MXFP8 non-transposed \", name,\n \" is missing column-wise data\");\n return {false, false, trans};\n }\n if (trans) {\n NVTE_CHECK(has_col, \"Grouped GEMM: MXFP8 transposed \", name, \" is missing column-wise data\");\n return {false, false, trans};\n }\n NVTE_CHECK(has_row, \"Grouped GEMM: MXFP8 non-transposed \", name, \" is missing row-wise data\");\n return {true, true, trans};\n }\n\n // Hopper-style TN-only FP8: force TN by switching layout and flipping transpose when needed.\n if (is_fp8 && !non_tn_fp8_ok) {\n if (is_A && !trans) {\n NVTE_CHECK(has_col, \"Grouped GEMM: \", name,\n \" is missing column-wise data needed for FP8 TN layout\");\n return {false, true, true};\n }\n if (!is_A && trans) {\n NVTE_CHECK(has_col, \"Grouped GEMM: \", name,\n \" is missing column-wise data needed for FP8 TN layout\");\n return {false, true, false};\n }\n }\n\n // If only column-wise data is available, mirror the transpose flag (pre-transposed storage).\n if (!has_row && has_col) {\n NVTE_CHECK(!is_fp8 || non_tn_fp8_ok,\n \"Grouped GEMM: FP8 on Hopper requires row-wise data for this transpose config.\");\n return {false, true, !trans};\n }\n\n NVTE_CHECK(has_row, \"Grouped GEMM: \", name, \" is missing row-wise data\");\n return {true, true, trans};\n}\n\n// Build Kernel Arguments detailing out addresses and other metadata for list of C/D tensors\n// passed to the grouped GEMM kernel. Use-case: C/D --> List of wgrads for experts in MOE\ninline MultiTensorGroupGemmOutputArgs build_grouped_gemm_multi_out_args(\n const NVTETensor *tensor_list, size_t list_size, size_t expected_num_tensors,\n transformer_engine::DType expected_dtype, const char *name) {\n MultiTensorGroupGemmOutputArgs args{};\n if (list_size == 0) {\n NVTE_CHECK(tensor_list == nullptr, \"Grouped GEMM: \", name, \"_list provided with num_\", name,\n \"_tensors=0\");\n return args;\n }\n NVTE_CHECK(tensor_list != nullptr, \"Grouped GEMM: \", name, \"_list is null but num_\", name,\n \"_tensors=\", list_size);\n NVTE_CHECK(list_size == expected_num_tensors, \"Grouped GEMM: \", name,\n \"_list must have num_tensors (\", expected_num_tensors, \") entries, got \", list_size);\n NVTE_CHECK(list_size <= static_cast(kMaxGroups), \"Grouped GEMM: \", name,\n \"_list supports up to \", kMaxGroups, \" tensors per kernel, got \", list_size);\n\n for (size_t i = 0; i < list_size; ++i) {\n const transformer_engine::Tensor *t =\n transformer_engine::convertNVTETensorCheck(tensor_list[i]);\n NVTE_CHECK(t->has_data(), \"Grouped GEMM: \", name, \"_list tensor \", i, \" has no data\");\n NVTE_CHECK(t->dtype() == expected_dtype, \"Grouped GEMM: \", name, \"_list tensor \", i,\n \" dtype mismatch. Expected \", transformer_engine::to_string(expected_dtype), \" got \",\n transformer_engine::to_string(t->dtype()));\n const auto &shape = t->shape();\n NVTE_CHECK(shape.size() == 2, \"Grouped GEMM: \", name, \"_list tensor \", i, \" must be 2D.\");\n args.data_ptrs[i] = t->data.dptr;\n args.rows[i] = static_cast(shape[1]);\n args.cols[i] = static_cast(shape[0]);\n }\n return args;\n}\n\n// Build Kernel Arguments detailing out addresses and other metadata for list of A tensors\n// passed to the grouped GEMM kernel. Use-case: A --> List of Expert weights\ninline MultiTensorGroupGemmInputArgs build_grouped_gemm_multi_inputA_args(\n const NVTETensor *tensor_list, size_t list_size, bool use_rowwise, bool is_fp8,\n int64_t *avg_first_dim, int64_t *avg_last_dim, const char *name) {\n using namespace transformer_engine;\n MultiTensorGroupGemmInputArgs args{};\n *avg_first_dim = 0;\n *avg_last_dim = 0;\n if (list_size == 0) {\n return args;\n }\n for (size_t i = 0; i < list_size; ++i) {\n const transformer_engine::Tensor *t =\n transformer_engine::convertNVTETensorCheck(tensor_list[i]);\n const transformer_engine::SimpleTensor &data = use_rowwise ? t->data : t->columnwise_data;\n const transformer_engine::SimpleTensor &scale_inv =\n use_rowwise ? t->scale_inv : t->columnwise_scale_inv;\n NVTE_CHECK(data.has_data(), \"Grouped GEMM: \", name, \"_list tensor \", i,\n \" is missing required data.\");\n NVTE_CHECK(data.shape.size() == 2, \"Grouped GEMM: \", name, \"_list tensor \", i, \" must be 2D.\");\n args.data_ptrs[i] = data.dptr;\n args.rows[i] = static_cast(data.shape[1]);\n args.cols[i] = static_cast(data.shape[0]);\n *avg_first_dim += static_cast(data.shape[0]);\n *avg_last_dim += static_cast(data.shape[1]);\n\n if (is_fp8) {\n NVTE_CHECK(scale_inv.has_data(), \"Grouped GEMM: \", name, \"_list tensor \", i,\n \" requires scale_inv for FP8.\");\n args.scale_inv_ptrs[i] = scale_inv.dptr;\n } else {\n args.scale_inv_ptrs[i] = nullptr;\n }\n }\n *avg_first_dim /= static_cast(list_size);\n *avg_last_dim /= static_cast(list_size);\n return args;\n}\n\ninline MultiTensorListInfo validate_grouped_gemm_multi_inputA_list(const NVTETensor *tensor_list,\n size_t list_size,\n size_t expected_num_tensors,\n const char *name) {\n using namespace transformer_engine;\n MultiTensorListInfo info{};\n if (list_size == 0) {\n NVTE_CHECK(tensor_list == nullptr, \"Grouped GEMM: \", name, \"_list provided with num_\", name,\n \"_tensors=0\");\n return info;\n }\n NVTE_CHECK(tensor_list != nullptr, \"Grouped GEMM: \", name, \"_list is null but num_\", name,\n \"_tensors=\", list_size);\n NVTE_CHECK(list_size == expected_num_tensors, \"Grouped GEMM: \", name,\n \"_list must have num_tensors (\", expected_num_tensors, \") entries, got \", list_size);\n NVTE_CHECK(list_size <= static_cast(kMaxGroups), \"Grouped GEMM: \", name,\n \"_list supports up to \", kMaxGroups, \" tensors per kernel, got \", list_size);\n\n const transformer_engine::Tensor *t0 = transformer_engine::convertNVTETensorCheck(tensor_list[0]);\n info.scaling_mode = t0->scaling_mode;\n info.with_gemm_swizzled_scales = t0->with_gemm_swizzled_scales;\n const bool mxfp8 = transformer_engine::is_mxfp_scaling(info.scaling_mode);\n NVTE_CHECK(info.scaling_mode == NVTE_DELAYED_TENSOR_SCALING || mxfp8,\n \"Grouped GEMM: input list only supports tensor scaling or MXFP8.\");\n\n for (size_t i = 0; i < list_size; ++i) {\n const transformer_engine::Tensor *t =\n transformer_engine::convertNVTETensorCheck(tensor_list[i]);\n NVTE_CHECK(t->scaling_mode == info.scaling_mode, \"Grouped GEMM: \", name,\n \"_list tensors must share the same scaling mode.\");\n NVTE_CHECK(t->with_gemm_swizzled_scales == info.with_gemm_swizzled_scales,\n \"Grouped GEMM: \", name, \"_list tensors must share GEMM swizzled scale state.\");\n\n if (t->has_data()) {\n if (info.row_dtype == DType::kNumTypes) {\n info.row_dtype = t->data.dtype;\n }\n // Check all tensors have the same dtype\n NVTE_CHECK(t->data.dtype == info.row_dtype, \"Grouped GEMM: \", name,\n \"_list rowwise dtypes must match.\");\n } else {\n // All tensors must have either data or columnwise data\n info.all_row = false;\n }\n\n if (t->has_columnwise_data()) {\n if (info.col_dtype == DType::kNumTypes) {\n info.col_dtype = t->columnwise_data.dtype;\n }\n NVTE_CHECK(t->columnwise_data.dtype == info.col_dtype, \"Grouped GEMM: \", name,\n \"_list columnwise dtypes must match.\");\n } else {\n // All tensors must have either data or columnwise data\n info.all_col = false;\n }\n }\n\n return info;\n}\n\n// Helper to create TensorShapeInfo from a GroupedTensor, optionally swapping first/last dims.\n// When swap_dims=true, first_dims and last_dims are swapped to account for columnwise storage.\n// Note: tensor_offsets are the same for rowwise and columnwise data (same element count per tensor).\ninline TensorShapeInfo create_shape_info(const transformer_engine::GroupedTensor *t,\n bool swap_dims) {\n const bool has_first = t->first_dims.has_data();\n const bool has_last = t->last_dims.has_data();\n NVTE_CHECK(has_first || t->all_same_first_dim(),\n \"GroupedTensor is missing first_dims for varying shapes\");\n NVTE_CHECK(has_last || t->all_same_last_dim(),\n \"GroupedTensor is missing last_dims for varying shapes\");\n\n const int64_t *first_ptr = has_first ? static_cast(t->first_dims.dptr) : nullptr;\n const int64_t *last_ptr = has_last ? static_cast(t->last_dims.dptr) : nullptr;\n const int64_t uniform_first = has_first ? 0 : static_cast(t->get_common_first_dim());\n const int64_t uniform_last = has_last ? 0 : static_cast(t->get_common_last_dim());\n\n const int64_t *offsets_ptr =\n t->tensor_offsets.has_data() ? static_cast(t->tensor_offsets.dptr) : nullptr;\n\n if (swap_dims) {\n // Swap first/last to account for columnwise (transposed) storage\n return {last_ptr, first_ptr, offsets_ptr, uniform_last, uniform_first};\n }\n return {first_ptr, last_ptr, offsets_ptr, uniform_first, uniform_last};\n}\n\ninline GroupedOperandSelection select_grouped_operand(const transformer_engine::GroupedTensor *t,\n bool trans, bool is_A) {\n using namespace transformer_engine;\n const bool has_row = t->has_data();\n const bool has_col = t->has_columnwise_data();\n\n if (!has_row && !has_col) {\n GroupedOperandSelection sel{};\n sel.trans = trans;\n sel.scaling_mode = t->scaling_mode;\n sel.dtype = t->dtype();\n sel.shape = create_shape_info(t, /*swap_dims=*/false);\n return sel;\n }\n\n const auto sm = t->scaling_mode;\n const bool mxfp8 = is_mxfp_scaling(sm);\n\n // Validate scaling mode\n NVTE_CHECK(sm == NVTE_DELAYED_TENSOR_SCALING || mxfp8,\n \"Grouped GEMM is only supported with bf16, fp8 tensor scaling and MXFP8\");\n\n const DType row_dtype = t->data.dtype;\n const DType col_dtype = t->columnwise_data.dtype;\n GroupedOperandSelection sel{};\n sel.trans = trans;\n sel.scaling_mode = sm;\n sel.with_gemm_swizzled_scales = t->with_gemm_swizzled_scales;\n\n const DType rep_dtype = has_row ? row_dtype : col_dtype;\n const bool is_fp8 = is_fp8_dtype(rep_dtype);\n const bool non_tn_fp8_ok = nvte_is_non_tn_fp8_gemm_supported();\n\n // Helper to select columnwise storage.\n // swap_dims=true (default): swap first/last dims in shape info (used when columnwise == transposed).\n // swap_dims=false: keep original dims (MXFP8: columnwise data has different scale direction,\n // but the logical matrix shape and transpose flag remain unchanged).\n auto use_columnwise = [&](bool swap_dims = true) {\n sel.dptr = static_cast(t->columnwise_data.dptr);\n sel.scale_inv = t->columnwise_scale_inv.dptr;\n sel.dtype = col_dtype;\n sel.rowwise = false;\n sel.shape = create_shape_info(t, swap_dims);\n };\n\n // Helper to select row-wise storage\n auto use_rowwise = [&]() {\n sel.dptr = static_cast(t->data.dptr);\n sel.scale_inv = t->scale_inv.dptr;\n sel.dtype = row_dtype;\n sel.rowwise = true;\n sel.shape = create_shape_info(t, /*swap_dims=*/false);\n };\n\n const auto choice = choose_grouped_operand_storage(trans, is_A, mxfp8, is_fp8, non_tn_fp8_ok,\n has_row, has_col, is_A ? \"A\" : \"B\");\n sel.trans = choice.trans;\n if (choice.use_rowwise) {\n use_rowwise();\n } else {\n use_columnwise(choice.swap_dims);\n }\n return sel;\n}\n\ninline void *validate_and_get_workspace_ptr(transformer_engine::Tensor *ws, size_t required_size,\n const char *workspace_name) {\n NVTE_CHECK(ws != nullptr, workspace_name, \" tensor is null.\");\n const size_t provided_size = get_buffer_size_bytes(ws->data.numel(), ws->data.dtype);\n NVTE_CHECK(provided_size >= required_size, \"Grouped GEMM: Insufficient \", workspace_name,\n \". Required: \", required_size, \" bytes, Available: \", provided_size, \" bytes.\");\n return ws->data.dptr;\n}\n\ninline void init_matrix_layouts(\n cublasLtMatrixLayoutOpaque_t &descA, cublasLtMatrixLayoutOpaque_t &descB,\n cublasLtMatrixLayoutOpaque_t &descC, cublasLtMatrixLayoutOpaque_t &descD,\n const GroupedGemmSetupWorkspace &ws, const GroupedOperandSelection &A_sel,\n const GroupedOperandSelection &B_sel, transformer_engine::DType d_dtype, size_t num_tensors) {\n const cudaDataType_t A_type = get_cuda_dtype(A_sel.dtype);\n const cudaDataType_t B_type = get_cuda_dtype(B_sel.dtype);\n const cudaDataType_t D_type = get_cuda_dtype(d_dtype);\n\n // Storage dimensions computed by kernel, leading dimension = rows\n NVTE_CHECK_CUBLAS(cublasLtGroupedMatrixLayoutInit(&descA, A_type, num_tensors, ws.a_rows,\n ws.a_cols, ws.a_rows));\n NVTE_CHECK_CUBLAS(cublasLtGroupedMatrixLayoutInit(&descB, B_type, num_tensors, ws.b_rows,\n ws.b_cols, ws.b_rows));\n NVTE_CHECK_CUBLAS(cublasLtGroupedMatrixLayoutInit(&descC, D_type, num_tensors, ws.d_rows,\n ws.d_cols, ws.d_rows));\n NVTE_CHECK_CUBLAS(cublasLtGroupedMatrixLayoutInit(&descD, D_type, num_tensors, ws.d_rows,\n ws.d_cols, ws.d_rows));\n}\n\ninline void init_matmul_desc(cublasLtMatmulDescOpaque_t &matmulDesc, cublasOperation_t op_A,\n cublasOperation_t op_B, bool use_fp8, bool use_split_accumulator) {\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescInit(&matmulDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F));\n\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_TRANSA, &op_A,\n sizeof(op_A)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_TRANSB, &op_B,\n sizeof(op_B)));\n\n cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_DEVICE;\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_POINTER_MODE,\n &pointer_mode, sizeof(pointer_mode)));\n\n int64_t alphabeta_batch_stride = 1;\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc,\n CUBLASLT_MATMUL_DESC_ALPHA_BATCH_STRIDE,\n &alphabeta_batch_stride, sizeof(int64_t)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc,\n CUBLASLT_MATMUL_DESC_BETA_BATCH_STRIDE,\n &alphabeta_batch_stride, sizeof(int64_t)));\n\n // Fast accumulation is only supported for FP8 (mirrors non-grouped GEMM logic).\n int8_t fastAccuMode = use_split_accumulator ? 0 : static_cast(use_fp8);\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_FAST_ACCUM,\n &fastAccuMode, sizeof(fastAccuMode)));\n}\n\n// Configures cuBLAS for MXFP8 grouped GEMM: sets VEC32_UE8M0 scale mode and scale pointers\n// for both A and B.\ninline void set_mxfp8_scale_pointers(cublasLtMatmulDescOpaque_t &matmulDesc,\n void **a_scale_inv_ptrs, void **b_scale_inv_ptrs) {\n#if CUBLAS_VERSION >= CUBLAS_MXFP8_GROUPED_GEMM_VERSION\n NVTE_CHECK(transformer_engine::cuda::cublas_version() >= CUBLAS_MXFP8_GROUPED_GEMM_VERSION,\n \"MXFP8 grouped GEMM requires cuBLAS \", CUBLAS_MXFP8_GROUPED_GEMM_VERSION,\n \"+, but run-time cuBLAS version is \", transformer_engine::cuda::cublas_version());\n const cublasLtMatmulMatrixScale_t scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0;\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_A_SCALE_MODE,\n &scale_mode, sizeof(scale_mode)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_B_SCALE_MODE,\n &scale_mode, sizeof(scale_mode)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc,\n CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,\n &a_scale_inv_ptrs, sizeof(a_scale_inv_ptrs)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc,\n CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,\n &b_scale_inv_ptrs, sizeof(b_scale_inv_ptrs)));\n#else\n NVTE_CHECK(false, \"MXFP8 grouped GEMM requires cuBLAS \", CUBLAS_MXFP8_GROUPED_GEMM_VERSION,\n \"+, but compile-time cuBLAS version is \", CUBLAS_VERSION);\n#endif // CUBLAS_VERSION >= CUBLAS_MXFP8_GROUPED_GEMM_VERSION\n}\n\n// Configures cuBLAS for tensor-scaling FP8 grouped GEMM: sets PER_BATCH_SCALAR_32F scale mode\n// and scale pointers for A and B. Both operands are guaranteed FP8 by the caller.\ninline void set_fp8_scale_pointers(cublasLtMatmulDescOpaque_t &matmulDesc, void **a_scale_inv_ptrs,\n void **b_scale_inv_ptrs) {\n const cublasLtMatmulMatrixScale_t scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_PER_BATCH_SCALAR_32F;\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_A_SCALE_MODE,\n &scale_mode, sizeof(scale_mode)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc,\n CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,\n &a_scale_inv_ptrs, sizeof(a_scale_inv_ptrs)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_B_SCALE_MODE,\n &scale_mode, sizeof(scale_mode)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc,\n CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,\n &b_scale_inv_ptrs, sizeof(b_scale_inv_ptrs)));\n}\ninline cublasLtMatmulAlgo_t select_grouped_gemm_algo(cublasLtHandle_t handle,\n cublasLtMatmulDescOpaque_t &matmulDesc,\n cublasLtMatrixLayoutOpaque_t &descA,\n cublasLtMatrixLayoutOpaque_t &descB,\n cublasLtMatrixLayoutOpaque_t &descC,\n cublasLtMatrixLayoutOpaque_t &descD,\n int64_t avg_m, int64_t avg_n, int64_t avg_k) {\n cublasLtMatmulPreferenceOpaque_t preference;\n NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceInit(&preference));\n NVTE_CHECK_CUBLAS(\n cublasLtMatmulPreferenceSetAttribute(&preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,\n &kGroupedGemmCublasWorkspaceSize, sizeof(size_t)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute(\n &preference, CUBLASLT_MATMUL_PREF_GROUPED_DESC_D_AVERAGE_ROWS, &avg_m, sizeof(int64_t)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute(\n &preference, CUBLASLT_MATMUL_PREF_GROUPED_DESC_D_AVERAGE_COLS, &avg_n, sizeof(int64_t)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute(\n &preference, CUBLASLT_MATMUL_PREF_GROUPED_AVERAGE_REDUCTION_DIM, &avg_k, sizeof(int64_t)));\n\n cublasLtMatmulHeuristicResult_t heuristicResult;\n int returnedResults = 0;\n auto status = cublasLtMatmulAlgoGetHeuristic(handle, &matmulDesc, &descA, &descB, &descC, &descD,\n &preference, 1, &heuristicResult, &returnedResults);\n NVTE_CHECK(status != CUBLAS_STATUS_NOT_SUPPORTED,\n \"Unable to find suitable cuBLAS grouped GEMM algorithm\");\n NVTE_CHECK_CUBLAS(status);\n NVTE_CHECK(returnedResults > 0, \"No suitable algorithm found for grouped GEMM\");\n return heuristicResult.algo;\n}\n\nstruct GroupedGemmWorkspace {\n GroupedGemmSetupWorkspace setup_workspace;\n void *cublas_workspace_ptr = nullptr;\n size_t num_tensors = 0;\n};\n\ninline GroupedGemmWorkspace setup_grouped_gemm_workspace(transformer_engine::Tensor *wspace_setup,\n transformer_engine::Tensor *wspace_cublas,\n size_t num_tensors) {\n const size_t setup_workspace_size = grouped_gemm_setup_workspace_size(num_tensors);\n const size_t cublas_workspace_size = kGroupedGemmCublasWorkspaceSize;\n void *setup_workspace_ptr = validate_and_get_workspace_ptr(wspace_setup, setup_workspace_size,\n \"Grouped GEMM setup workspace\");\n void *cublas_workspace_ptr = validate_and_get_workspace_ptr(wspace_cublas, cublas_workspace_size,\n \"Grouped GEMM cuBLAS workspace\");\n auto setup_workspace = GroupedGemmSetupWorkspace::from_buffers(\n static_cast(setup_workspace_ptr), num_tensors);\n return {std::move(setup_workspace), cublas_workspace_ptr, num_tensors};\n}\n\ninline void execute_grouped_gemm(const GroupedGemmSetupWorkspace &setup_workspace,\n const GroupedOperandSelection &A_sel,\n const GroupedOperandSelection &B_sel,\n transformer_engine::DType d_dtype, size_t num_tensors,\n bool use_split_accumulator, bool use_fp8, int64_t avg_m_val,\n int64_t avg_n_val, int64_t avg_k_val, void *cublas_workspace_ptr,\n cudaStream_t stream, int math_sm_count = 0) {\n using cublasHandleManager =\n transformer_engine::detail::HandleManager;\n cublasLtHandle_t handle = cublasHandleManager::Instance().GetHandle();\n\n cublasOperation_t op_A = A_sel.trans ? CUBLAS_OP_T : CUBLAS_OP_N;\n cublasOperation_t op_B = B_sel.trans ? CUBLAS_OP_T : CUBLAS_OP_N;\n\n cublasLtMatrixLayoutOpaque_t descA, descB, descC, descD;\n init_matrix_layouts(descA, descB, descC, descD, setup_workspace, A_sel, B_sel, d_dtype,\n num_tensors);\n\n cublasLtMatmulDescOpaque_t matmulDesc;\n init_matmul_desc(matmulDesc, op_A, op_B, use_fp8, use_split_accumulator);\n if (transformer_engine::is_mxfp_scaling(A_sel.scaling_mode)) {\n set_mxfp8_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs,\n setup_workspace.b_scale_inv_ptrs);\n } else if (use_fp8) {\n set_fp8_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs,\n setup_workspace.b_scale_inv_ptrs);\n }\n if (math_sm_count != 0) {\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n &matmulDesc, CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET, &math_sm_count, sizeof(math_sm_count)));\n }\n cublasLtMatmulAlgo_t algo = select_grouped_gemm_algo(handle, matmulDesc, descA, descB, descC,\n descD, avg_m_val, avg_n_val, avg_k_val);\n\n NVTE_CHECK_CUBLAS(cublasLtMatmul(handle, &matmulDesc, setup_workspace.alpha_ptrs,\n setup_workspace.A_ptrs, &descA, setup_workspace.B_ptrs, &descB,\n setup_workspace.beta_ptrs, setup_workspace.C_ptrs, &descC,\n setup_workspace.D_ptrs, &descD, &algo, cublas_workspace_ptr,\n kGroupedGemmCublasWorkspaceSize, stream));\n}\n\n// Device helper: compute the element offset for tensor `idx` given shape metadata.\n// Three cases:\n// 1. Explicit per-tensor offset array provided \u2192 use it directly.\n// 2. Per-tensor first/last dims provided but no offsets \u2192 cumulative sum of (first*last) products.\n// 3. Fully uniform shapes \u2192 idx * uniform_first * uniform_last.\n__forceinline__ __device__ int64_t compute_grouped_tensor_offset(const TensorShapeInfo &meta,\n size_t idx) {\n if (meta.offsets) {\n return meta.offsets[idx];\n } else if (meta.first_dims != nullptr || meta.last_dims != nullptr) {\n // offset[i] = sum_{j < i} (first_dims[j] * last_dims[j])\n int64_t cumsum = 0;\n for (size_t i = 0; i < idx; i++) {\n int64_t f = meta.first_dims ? meta.first_dims[i] : meta.uniform_first;\n int64_t l = meta.last_dims ? meta.last_dims[i] : meta.uniform_last;\n cumsum += f * l;\n }\n return cumsum;\n } else {\n return static_cast(idx) * meta.uniform_first * meta.uniform_last;\n }\n}\n\n__forceinline__ __device__ int64_t padded_mxfp8_scale_inv_bytes(int64_t first, int64_t last,\n bool rowwise) {\n namespace mxfp8_swizzle = transformer_engine::dispatch::mxfp8::swizzle;\n constexpr int64_t kMxfp8BlockSize = 32;\n // x is the dimension along which quantization is applied, y is other dimension\n const int64_t scale_tile_y = static_cast(mxfp8_swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_Y);\n const int64_t scale_tile_x = static_cast(mxfp8_swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_X);\n // Padded byte size of the swizzled MXFP8 scale_inv for a single tensor with data\n // shape (first, last). Rowwise scales use rows=first, cols=last; columnwise\n // scales swap the orientation since they are stored in column-major order.\n const int64_t scale_dim_y = rowwise ? first : last;\n const int64_t padded_scale_dim_y =\n ((scale_dim_y + scale_tile_y - 1) / scale_tile_y) * scale_tile_y;\n const int64_t data_dim_x = rowwise ? last : first;\n const int64_t scale_dim_x = (data_dim_x + kMxfp8BlockSize - 1) / kMxfp8BlockSize;\n const int64_t padded_scale_dim_x =\n ((scale_dim_x + scale_tile_x - 1) / scale_tile_x) * scale_tile_x;\n // MXFP8 scales are E8M0 (1 byte per element), so element count == byte count.\n return padded_scale_dim_y * padded_scale_dim_x;\n}\n\n// Device helper: byte offset into a contiguous grouped MXFP8 scale_inv buffer for\n// tensor `idx`. Each expert's scale_inv is expected to be padded\n// to the 128x4 swizzled layout.\n__forceinline__ __device__ int64_t compute_grouped_tensor_mxfp8_scale_inv_offset(\n const TensorShapeInfo &meta, size_t idx, bool rowwise) {\n if (meta.first_dims != nullptr || meta.last_dims != nullptr) {\n int64_t cumsum = 0;\n for (size_t i = 0; i < idx; i++) {\n const int64_t f = meta.first_dims ? meta.first_dims[i] : meta.uniform_first;\n const int64_t l = meta.last_dims ? meta.last_dims[i] : meta.uniform_last;\n cumsum += padded_mxfp8_scale_inv_bytes(f, l, rowwise);\n }\n return cumsum;\n }\n return static_cast(idx) *\n padded_mxfp8_scale_inv_bytes(meta.uniform_first, meta.uniform_last, rowwise);\n}\n\n// Linear scan to find which tensor contains the given row.\n// Returns the tensor index and writes the exclusive end-row of that tensor to *out_tensor_row_end.\n__forceinline__ __device__ int find_tensor_for_row(const int64_t *first_dims, int64_t uniform_first,\n int row, int num_tensors,\n int *out_tensor_row_end) {\n int offset = 0;\n for (int i = 0; i < num_tensors; i++) {\n int dim = first_dims ? static_cast(first_dims[i]) : static_cast(uniform_first);\n offset += dim;\n if (row < offset) {\n *out_tensor_row_end = offset;\n return i;\n }\n }\n *out_tensor_row_end = offset;\n return num_tensors - 1;\n}\n\n// Kernel that performs (optionally scaled) bias addition to Grouped GEMM output tensors.\n// SM-filling grid with grid-stride over row chunks.\n// 2D grid: blockIdx.x = SM-filling row blocks, blockIdx.y = column chunk.\n// Each block grid-strides over kRowsPerBlock-sized row chunks, processing\n// all chunks that map to it. Safe when sum(first_dims) <= total_rows.\ntemplate \n__global__ void grouped_bias_add_kernel(char *__restrict__ d_base,\n const char *__restrict__ bias_base,\n const float *__restrict__ scale_base,\n TensorShapeInfo d_meta, int n, int total_rows,\n int num_tensors) {\n using VecStorage = transformer_engine::VectorizedStorage;\n using VecType = typename VecStorage::LType;\n\n const int tid = static_cast(threadIdx.x);\n const int row_bid = static_cast(blockIdx.x);\n const int col_bid = static_cast(blockIdx.y);\n const int row_grid_stride = static_cast(gridDim.x);\n\n // Single-warp reduction to compute valid_rows = sum(first_dims).\n // kMaxGroups <= 64 so warp 0 (32 lanes) covers it with <=2 loads each.\n __shared__ int s_valid_rows;\n if (tid < 32) {\n int local_sum = 0;\n for (int i = tid; i < num_tensors; i += 32) {\n local_sum += d_meta.first_dims ? static_cast(d_meta.first_dims[i])\n : static_cast(d_meta.uniform_first);\n }\n for (int offset = 16; offset > 0; offset >>= 1) {\n local_sum += __shfl_down_sync(0xffffffff, local_sum, offset);\n }\n if (tid == 0) s_valid_rows = local_sum;\n }\n __syncthreads();\n const int valid_rows = s_valid_rows;\n\n const int block_cols = kBlockDim * kVec;\n const int col = col_bid * block_cols + tid * kVec;\n if (col >= n) return;\n\n T *__restrict__ d = reinterpret_cast(d_base);\n const T *__restrict__ bias = reinterpret_cast(bias_base);\n\n // Grid-stride loop over row chunks.\n for (int chunk_start = row_bid * kRowsPerBlock; chunk_start < valid_rows;\n chunk_start += row_grid_stride * kRowsPerBlock) {\n const int row_start = chunk_start;\n const int row_end = min(row_start + kRowsPerBlock, valid_rows);\n\n // Linear scan to find the starting row's tensor and its boundary.\n int tensor_row_end;\n int tensor_idx = find_tensor_for_row(d_meta.first_dims, d_meta.uniform_first, row_start,\n num_tensors, &tensor_row_end);\n int bias_idx = tensor_idx * n;\n\n VecStorage b_in;\n\n // Walk tensor segments within this chunk's row range.\n int seg_start = row_start;\n while (seg_start < row_end) {\n while (tensor_idx < num_tensors - 1 && tensor_row_end <= seg_start) {\n tensor_idx++;\n bias_idx += n;\n int dim = d_meta.first_dims ? static_cast(d_meta.first_dims[tensor_idx])\n : static_cast(d_meta.uniform_first);\n tensor_row_end += dim;\n }\n b_in.scratch_.aligned = *reinterpret_cast(bias + bias_idx + col);\n const int seg_end = min(tensor_row_end, row_end);\n\n for (int row = seg_start; row < seg_end; row++) {\n T *d_ptr = d + row * n + col;\n VecStorage d_in;\n d_in.scratch_.aligned = *reinterpret_cast(d_ptr);\n\n [[maybe_unused]] float s_val;\n if constexpr (UseScale) s_val = scale_base[row];\n\n#pragma unroll\n for (int i = 0; i < kVec; ++i) {\n if constexpr (UseScale) {\n d_in.scratch_.separate[i] =\n static_cast(fmaf(static_cast(b_in.scratch_.separate[i]), s_val,\n static_cast(d_in.scratch_.separate[i])));\n } else {\n d_in.scratch_.separate[i] =\n static_cast(static_cast(d_in.scratch_.separate[i]) +\n static_cast(b_in.scratch_.separate[i]));\n }\n }\n *reinterpret_cast(d_ptr) = d_in.scratch_.aligned;\n }\n\n seg_start = seg_end;\n }\n }\n}\n\n// Single kernel that sets up all GEMM parameters.\n// Rationale: cuBLASLt grouped matmul API needs flat arrays of pointers and per-matrix dimensions,\n// but NVTEGroupedTensor stores a single contiguous buffer + optional per-tensor offsets/shapes.\n// We bridge the mismatch on GPU by computing per-group pointers and storage dims in one kernel.\n__global__ void setup_grouped_gemm_kernel(\n // Output arrays\n void **A_ptrs, void **B_ptrs, void **C_ptrs, void **D_ptrs, int *a_rows, int *a_cols,\n int *b_rows, int *b_cols, int *d_rows, int *d_cols, float **alpha_ptrs, float **beta_ptrs,\n void **a_scale_inv_ptrs, void **b_scale_inv_ptrs,\n // Inputs\n char *a_base, char *b_base, char *c_base, char *d_base, TensorShapeInfo A_meta,\n TensorShapeInfo B_meta, TensorShapeInfo C_meta, TensorShapeInfo D_meta, size_t a_elem_size,\n size_t b_elem_size, size_t c_elem_size, size_t d_elem_size, float *alpha_ptr, float *beta_ptr,\n // Scale inputs: for tensor scaling, pass float* and set mxfp8_base to nullptr\n // For MXFP8, pass nullptr for tensor_scale and set mxfp8_base\n float *a_scale_base, float *b_scale_base, bool a_rowwise, bool b_rowwise,\n NVTEScalingMode scaling_mode, size_t num_tensors,\n MultiTensorGroupGemmInputArgs a_multi_tensor_args,\n MultiTensorGroupGemmOutputArgs c_multi_tensor_args,\n MultiTensorGroupGemmOutputArgs d_multi_tensor_args) {\n size_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= num_tensors) return;\n\n // Get dimensions for this tensor (from array or uniform value)\n const bool has_a_multi_tensor = (a_base == nullptr);\n const bool has_c_multi_tensor = (c_base == nullptr);\n const bool has_d_multi_tensor = (d_base == nullptr);\n int64_t a_first = 0;\n int64_t a_last = 0;\n if (has_a_multi_tensor) {\n a_first = static_cast(a_multi_tensor_args.cols[idx]);\n a_last = static_cast(a_multi_tensor_args.rows[idx]);\n } else {\n a_first = A_meta.first_dims ? A_meta.first_dims[idx] : A_meta.uniform_first;\n a_last = A_meta.last_dims ? A_meta.last_dims[idx] : A_meta.uniform_last;\n }\n int64_t b_first = B_meta.first_dims ? B_meta.first_dims[idx] : B_meta.uniform_first;\n int64_t b_last = B_meta.last_dims ? B_meta.last_dims[idx] : B_meta.uniform_last;\n int64_t d_first = D_meta.first_dims ? D_meta.first_dims[idx] : D_meta.uniform_first;\n int64_t d_last = D_meta.last_dims ? D_meta.last_dims[idx] : D_meta.uniform_last;\n\n // Compute offsets (from explicit array, cumulative from per-tensor dims, or uniform)\n int64_t a_offset = has_a_multi_tensor ? 0 : compute_grouped_tensor_offset(A_meta, idx);\n int64_t b_offset = compute_grouped_tensor_offset(B_meta, idx);\n int64_t c_offset = compute_grouped_tensor_offset(C_meta, idx);\n int64_t d_offset = compute_grouped_tensor_offset(D_meta, idx);\n\n // Compute data pointers\n A_ptrs[idx] =\n has_a_multi_tensor ? a_multi_tensor_args.data_ptrs[idx] : (a_base + a_offset * a_elem_size);\n B_ptrs[idx] = b_base + b_offset * b_elem_size;\n C_ptrs[idx] =\n has_c_multi_tensor ? c_multi_tensor_args.data_ptrs[idx] : (c_base + c_offset * c_elem_size);\n D_ptrs[idx] =\n has_d_multi_tensor ? d_multi_tensor_args.data_ptrs[idx] : (d_base + d_offset * d_elem_size);\n\n // Compute storage dimensions for cuBLAS matrix layouts.\n // For INPUTS (A, B): Row-wise storage is seen as transposed column-major by cuBLAS,\n // so rows=last, cols=first. For columnwise, dims are already swapped.\n a_rows[idx] = static_cast(a_last);\n a_cols[idx] = static_cast(a_first);\n b_rows[idx] = static_cast(b_last);\n b_cols[idx] = static_cast(b_first);\n if (has_d_multi_tensor) {\n d_rows[idx] = d_multi_tensor_args.rows[idx];\n d_cols[idx] = d_multi_tensor_args.cols[idx];\n } else {\n d_rows[idx] = static_cast(d_last);\n d_cols[idx] = static_cast(d_first);\n }\n\n // Fill alpha/beta pointers (per-matrix)\n alpha_ptrs[idx] = alpha_ptr + idx;\n beta_ptrs[idx] = beta_ptr + idx;\n\n // Fill scale pointers (per-matrix).\n // The interpretation of the scale buffers depends on the shared scaling recipe:\n // otherwise : one float per tensor, indexed by tensor index\n if (a_scale_base) {\n if (scaling_mode == NVTE_MXFP8_1D_SCALING) {\n const int64_t a_scale_offset =\n compute_grouped_tensor_mxfp8_scale_inv_offset(A_meta, idx, a_rowwise);\n a_scale_inv_ptrs[idx] = reinterpret_cast(\n static_cast(static_cast(a_scale_base)) + a_scale_offset);\n } else {\n a_scale_inv_ptrs[idx] = static_cast(a_scale_base) + idx;\n }\n } else {\n a_scale_inv_ptrs[idx] = a_multi_tensor_args.scale_inv_ptrs[idx];\n }\n if (b_scale_base) {\n if (scaling_mode == NVTE_MXFP8_1D_SCALING) {\n const int64_t b_scale_offset =\n compute_grouped_tensor_mxfp8_scale_inv_offset(B_meta, idx, b_rowwise);\n b_scale_inv_ptrs[idx] = reinterpret_cast(\n static_cast(static_cast(b_scale_base)) + b_scale_offset);\n } else {\n b_scale_inv_ptrs[idx] = static_cast(b_scale_base) + idx;\n }\n }\n}\n\n// Launch the setup kernel to populate workspace arrays\ninline void launch_grouped_gemm_setup(\n const GroupedGemmSetupWorkspace &ws, const GroupedOperandSelection &A_sel,\n const GroupedOperandSelection &B_sel, const transformer_engine::GroupedTensor *C,\n const transformer_engine::GroupedTensor *D, const transformer_engine::Tensor *alpha_tensor,\n const transformer_engine::Tensor *beta_tensor, size_t num_tensors, cudaStream_t stream,\n const MultiTensorGroupGemmInputArgs &a_multi_tensor_args, const NVTETensor *C_list,\n const NVTETensor *D_list, char *a_base, transformer_engine::DType c_dtype,\n transformer_engine::DType d_dtype) {\n // Use shape info from selection (already accounts for columnwise dimension swap)\n TensorShapeInfo A_meta = A_sel.shape;\n TensorShapeInfo B_meta = B_sel.shape;\n TensorShapeInfo C_meta{};\n TensorShapeInfo D_meta{};\n\n const bool has_d_multi_tensor = (D_list != nullptr);\n const bool has_c_multi_tensor = (C_list != nullptr) || has_d_multi_tensor;\n MultiTensorGroupGemmOutputArgs c_multi_tensor_args{};\n MultiTensorGroupGemmOutputArgs d_multi_tensor_args{};\n if (has_d_multi_tensor) {\n d_multi_tensor_args =\n build_grouped_gemm_multi_out_args(D_list, num_tensors, num_tensors, d_dtype, \"D\");\n }\n if (C_list != nullptr) {\n c_multi_tensor_args =\n build_grouped_gemm_multi_out_args(C_list, num_tensors, num_tensors, d_dtype, \"C\");\n } else if (has_d_multi_tensor) {\n c_multi_tensor_args = d_multi_tensor_args;\n }\n\n char *c_base = nullptr;\n char *d_base = nullptr;\n\n if (!has_c_multi_tensor) {\n NVTE_CHECK(C != nullptr && D != nullptr,\n \"Grouped GEMM: C/D grouped tensors are required when no C list is provided\");\n C_meta = TensorShapeInfo::create_shape_info_for_C(C, D);\n c_base = static_cast(C->data.dptr);\n }\n if (!has_d_multi_tensor) {\n NVTE_CHECK(D != nullptr,\n \"Grouped GEMM: D grouped tensor is required when no D list is provided\");\n D_meta = TensorShapeInfo::from_tensor(D);\n d_base = static_cast(D->data.dptr);\n }\n\n const size_t a_elem_size = transformer_engine::typeToSize(A_sel.dtype);\n const size_t b_elem_size = transformer_engine::typeToSize(B_sel.dtype);\n const size_t c_elem_size = transformer_engine::typeToSize(c_dtype);\n const size_t d_elem_size = transformer_engine::typeToSize(d_dtype);\n\n const int threads_per_block = 256;\n const int num_blocks = (num_tensors + threads_per_block - 1) / threads_per_block;\n\n // A and B share the same scaling recipe (validated in validate_grouped_gemm_inputs).\n // Pass scale buffers as void* and let the kernel interpret them via scaling_mode.\n\n // Scale rowwise flag for MXFP8/NVFP4: to calculate scale_inv padding based offsets\n // within kernel. Ignored for tensor scaling.\n const bool a_rowwise = A_sel.rowwise;\n const bool b_rowwise = B_sel.rowwise;\n setup_grouped_gemm_kernel<<>>(\n ws.A_ptrs, ws.B_ptrs, ws.C_ptrs, ws.D_ptrs, ws.a_rows, ws.a_cols, ws.b_rows, ws.b_cols,\n ws.d_rows, ws.d_cols, ws.alpha_ptrs, ws.beta_ptrs, ws.a_scale_inv_ptrs, ws.b_scale_inv_ptrs,\n A_sel.dptr, B_sel.dptr, c_base, d_base, A_meta, B_meta, C_meta, D_meta, a_elem_size,\n b_elem_size, c_elem_size, d_elem_size, static_cast(alpha_tensor->data.dptr),\n static_cast(beta_tensor->data.dptr), reinterpret_cast(A_sel.scale_inv),\n reinterpret_cast(B_sel.scale_inv), a_rowwise, b_rowwise, A_sel.scaling_mode,\n num_tensors, a_multi_tensor_args, c_multi_tensor_args, d_multi_tensor_args);\n\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n} // namespace\n\nsize_t nvte_get_grouped_gemm_setup_workspace_size(size_t num_tensors) {\n NVTE_API_CALL(nvte_get_grouped_gemm_setup_workspace_size);\n return grouped_gemm_setup_workspace_size(num_tensors);\n}\n\nvoid nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedTensor B, int transb,\n const NVTEGroupedTensor C, NVTEGroupedTensor D, const NVTETensor alpha,\n const NVTETensor beta, NVTETensor workspace_setup,\n NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_grouped_gemm);\n using namespace transformer_engine;\n\n // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.3+\n check_grouped_gemm_requirements(\"nvte_grouped_gemm\");\n\n // Convert to internal types\n const GroupedTensor *inputA = convertNVTEGroupedTensorCheck(A);\n const GroupedTensor *inputB = convertNVTEGroupedTensorCheck(B);\n const GroupedTensor *inputC_raw = convertNVTEGroupedTensor(C); // Can be NULL\n GroupedTensor *outputD = convertNVTEGroupedTensorCheck(D);\n const Tensor *alpha_tensor = convertNVTETensorCheck(alpha);\n const Tensor *beta_tensor = convertNVTETensorCheck(beta);\n Tensor *wspace_setup = convertNVTETensor(workspace_setup);\n Tensor *wspace_cublas = convertNVTETensor(workspace_cublas);\n\n // Parse config (if provided)\n GroupedMatmulConfig config_ = parse_grouped_gemm_config(config);\n\n // Validate inputs and outputs.\n const size_t num_tensors = validate_grouped_gemm_inputs(inputA->num_tensors, {inputA, inputB},\n alpha_tensor, beta_tensor);\n validate_grouped_gemm_outputs(num_tensors, {inputC_raw, outputD});\n\n // If C is NULL, use D as C (valid when beta=0, cuBLAS won't read C data)\n const GroupedTensor *inputC = (inputC_raw != nullptr) ? inputC_raw : outputD;\n // num_tensors validated above.\n // Select operand storage (row-wise vs column-wise) and adjust transpose flags to\n // mirror the non-grouped GEMM logic for FP8 layout constraints.\n auto A_sel = select_grouped_operand(inputA, static_cast(transa), /*is_A=*/true);\n auto B_sel = select_grouped_operand(inputB, static_cast(transb), /*is_A=*/false);\n\n // Workspaces: setup (pointer arrays) and cuBLAS\n auto workspace = setup_grouped_gemm_workspace(wspace_setup, wspace_cublas, num_tensors);\n\n MultiTensorGroupGemmInputArgs a_multi_tensor_args{};\n launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor,\n beta_tensor, num_tensors, stream, a_multi_tensor_args,\n /*C_list=*/nullptr, /*D_list=*/nullptr, A_sel.dptr, inputC->dtype(),\n outputD->dtype());\n\n // Compute average dimensions for heuristics\n // K dimension: if transa, K is A's first dim; if not, K is A's last dim\n // Use original inputA and transa for heuristics (not modified A_sel.trans)\n int64_t avg_m_val = config_.avg_m.value_or(compute_avg_first_dim(outputD));\n int64_t avg_n_val = config_.avg_n.value_or(compute_avg_last_dim(outputD));\n int64_t avg_k_val =\n config_.avg_k.value_or(transa ? compute_avg_first_dim(inputA) : compute_avg_last_dim(inputA));\n const bool use_fp8 = is_fp8_dtype(A_sel.dtype) || is_fp8_dtype(B_sel.dtype);\n execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, outputD->dtype(), num_tensors,\n config_.use_split_accumulator, use_fp8, avg_m_val, avg_n_val, avg_k_val,\n workspace.cublas_workspace_ptr, stream, config_.sm_count);\n}\n\nvoid nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors,\n int transa, const NVTEGroupedTensor B, int transb,\n const NVTEGroupedTensor C, NVTEGroupedTensor D,\n const NVTETensor alpha, const NVTETensor beta,\n NVTETensor workspace_setup, NVTETensor workspace_cublas,\n NVTEGroupedMatmulConfig config, cudaStream_t stream) {\n NVTE_API_CALL(nvte_grouped_gemm_with_discrete_inputA);\n using namespace transformer_engine;\n\n // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.3+\n check_grouped_gemm_requirements(\"nvte_grouped_gemm_with_discrete_inputA\");\n\n NVTE_CHECK(A_list != nullptr, \"Grouped GEMM: A_list is null.\");\n NVTE_CHECK(num_a_tensors > 0, \"Grouped GEMM: num_a_tensors must be > 0.\");\n\n const GroupedTensor *inputB = convertNVTEGroupedTensorCheck(B);\n const GroupedTensor *inputC_raw = convertNVTEGroupedTensor(C); // Can be NULL\n GroupedTensor *outputD = convertNVTEGroupedTensorCheck(D);\n const Tensor *alpha_tensor = convertNVTETensorCheck(alpha);\n const Tensor *beta_tensor = convertNVTETensorCheck(beta);\n Tensor *wspace_setup = convertNVTETensor(workspace_setup);\n Tensor *wspace_cublas = convertNVTETensor(workspace_cublas);\n\n // Parse config (if provided)\n GroupedMatmulConfig config_ = parse_grouped_gemm_config(config);\n\n // Validate inputs and outputs.\n const size_t num_tensors =\n validate_grouped_gemm_inputs(num_a_tensors, {inputB}, alpha_tensor, beta_tensor);\n\n validate_grouped_gemm_outputs(num_tensors, {inputC_raw, outputD});\n\n // If C is NULL, use D as C (valid when beta=0, cuBLAS won't read C data)\n const GroupedTensor *inputC = (inputC_raw != nullptr) ? inputC_raw : outputD;\n\n // Validate A list and selection\n auto A_list_info =\n validate_grouped_gemm_multi_inputA_list(A_list, num_a_tensors, num_tensors, \"A\");\n auto is_fp8_or_16bit = [](transformer_engine::DType dtype) {\n return dtype == transformer_engine::DType::kFloat8E4M3 ||\n dtype == transformer_engine::DType::kFloat8E5M2 ||\n dtype == transformer_engine::DType::kBFloat16 ||\n dtype == transformer_engine::DType::kFloat16;\n };\n NVTE_CHECK(is_fp8_or_16bit(A_list_info.all_row ? A_list_info.row_dtype : A_list_info.col_dtype),\n \"Grouped GEMM: A_list tensors must be FP8, BF16, or FP16.\");\n\n // Cross-operand consistency (mirrors validate_grouped_gemm_inputs).\n const DType a_rep_dtype = A_list_info.all_row ? A_list_info.row_dtype : A_list_info.col_dtype;\n NVTE_CHECK(is_fp8_dtype(a_rep_dtype) == is_fp8_dtype(inputB->dtype()),\n \"Grouped GEMM: A and B must both be FP8 or both be non-FP8.\");\n NVTE_CHECK(transformer_engine::is_mxfp_scaling(A_list_info.scaling_mode) ==\n transformer_engine::is_mxfp_scaling(inputB->scaling_mode),\n \"Grouped GEMM: A and B must both use MXFP8 scaling or both use tensor scaling.\");\n if (transformer_engine::is_mxfp_scaling(A_list_info.scaling_mode)) {\n NVTE_CHECK(A_list_info.with_gemm_swizzled_scales,\n \"MXFP8 grouped GEMM: A scales must be swizzled for GEMM.\");\n NVTE_CHECK(inputB->with_gemm_swizzled_scales,\n \"MXFP8 grouped GEMM: B scales must be swizzled for GEMM.\");\n }\n\n // Select operand storage for B (row-wise vs column-wise)\n auto B_sel = select_grouped_operand(inputB, static_cast(transb), /*is_A=*/false);\n\n GroupedOperandSelection A_sel{};\n A_sel.scaling_mode = A_list_info.scaling_mode;\n A_sel.with_gemm_swizzled_scales = A_list_info.with_gemm_swizzled_scales;\n A_sel.trans = static_cast(transa);\n\n const DType rep_dtype = A_list_info.all_row ? A_list_info.row_dtype : A_list_info.col_dtype;\n const bool is_fp8 = is_fp8_dtype(rep_dtype);\n const bool non_tn_fp8_ok = nvte_is_non_tn_fp8_gemm_supported();\n const bool mxfp8 = transformer_engine::is_mxfp_scaling(A_list_info.scaling_mode);\n\n int64_t avg_first_dim = 0;\n int64_t avg_last_dim = 0;\n MultiTensorGroupGemmInputArgs a_multi_tensor_args{};\n\n const auto choice =\n choose_grouped_operand_storage(static_cast(transa), /*is_A=*/true, mxfp8, is_fp8,\n non_tn_fp8_ok, A_list_info.all_row, A_list_info.all_col, \"A\");\n A_sel.trans = choice.trans;\n A_sel.rowwise = choice.use_rowwise;\n if (choice.use_rowwise) {\n NVTE_CHECK(A_list_info.all_row, \"Grouped GEMM: A_list is missing row-wise data\");\n A_sel.dtype = A_list_info.row_dtype;\n a_multi_tensor_args = build_grouped_gemm_multi_inputA_args(\n A_list, num_a_tensors, /*use_rowwise=*/true, is_fp8, &avg_first_dim, &avg_last_dim, \"A\");\n } else {\n NVTE_CHECK(A_list_info.all_col, \"Grouped GEMM: A_list is missing column-wise data\");\n A_sel.dtype = A_list_info.col_dtype;\n a_multi_tensor_args = build_grouped_gemm_multi_inputA_args(\n A_list, num_a_tensors, /*use_rowwise=*/false, is_fp8, &avg_first_dim, &avg_last_dim, \"A\");\n }\n\n // For discrete A_list, scale pointers are per-tensor; use multi-tensor args.\n // Base pointer is unused when providing per-tensor pointers.\n A_sel.scale_inv = nullptr;\n A_sel.dptr = nullptr;\n\n // Workspaces: setup (pointer arrays) and cuBLAS\n auto workspace = setup_grouped_gemm_workspace(wspace_setup, wspace_cublas, num_tensors);\n\n launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor,\n beta_tensor, num_tensors, stream, a_multi_tensor_args,\n /*C_list=*/nullptr, /*D_list=*/nullptr, nullptr, inputC->dtype(),\n outputD->dtype());\n\n // Compute average dimensions for heuristics\n int64_t avg_m_val = config_.avg_m.value_or(compute_avg_first_dim(outputD));\n int64_t avg_n_val =\n config_.avg_n.value_or(transb ? compute_avg_first_dim(inputB) : compute_avg_last_dim(inputB));\n int64_t avg_k_val =\n config_.avg_k.value_or(static_cast(transa) ? avg_last_dim : avg_first_dim);\n const bool use_fp8 = is_fp8_dtype(A_sel.dtype) || is_fp8_dtype(B_sel.dtype);\n execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, outputD->dtype(), num_tensors,\n config_.use_split_accumulator, use_fp8, avg_m_val, avg_n_val, avg_k_val,\n workspace.cublas_workspace_ptr, stream, config_.sm_count);\n}\n\nvoid nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa,\n const NVTEGroupedTensor B, int transb,\n const NVTETensor *C_list, size_t num_c_tensors,\n NVTETensor *D_list, size_t num_d_tensors,\n const NVTETensor alpha, const NVTETensor beta,\n NVTETensor workspace_setup, NVTETensor workspace_cublas,\n NVTEGroupedMatmulConfig config, cudaStream_t stream) {\n NVTE_API_CALL(nvte_grouped_gemm_with_discrete_out);\n using namespace transformer_engine;\n\n // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.3+\n check_grouped_gemm_requirements(\"nvte_grouped_gemm_with_discrete_out\");\n\n NVTE_CHECK(D_list != nullptr, \"Grouped GEMM: D_list is null.\");\n NVTE_CHECK(num_d_tensors > 0, \"Grouped GEMM: num_d_tensors must be > 0.\");\n if (num_c_tensors > 0) {\n NVTE_CHECK(C_list != nullptr, \"Grouped GEMM: C_list is null but num_c_tensors > 0.\");\n }\n\n const GroupedTensor *inputA = convertNVTEGroupedTensorCheck(A);\n const GroupedTensor *inputB = convertNVTEGroupedTensorCheck(B);\n const Tensor *alpha_tensor = convertNVTETensorCheck(alpha);\n const Tensor *beta_tensor = convertNVTETensorCheck(beta);\n Tensor *wspace_setup = convertNVTETensor(workspace_setup);\n Tensor *wspace_cublas = convertNVTETensor(workspace_cublas);\n\n const Tensor *d0 = convertNVTETensorCheck(D_list[0]);\n const DType d_dtype = d0->dtype();\n\n const size_t num_tensors = validate_grouped_gemm_inputs(inputA->num_tensors, {inputA, inputB},\n alpha_tensor, beta_tensor);\n NVTE_CHECK(num_d_tensors == num_tensors, \"Grouped GEMM: D_list must have num_tensors (\",\n num_tensors, \") entries, got \", num_d_tensors);\n if (num_c_tensors > 0) {\n NVTE_CHECK(num_c_tensors == num_tensors, \"Grouped GEMM: C_list must have num_tensors (\",\n num_tensors, \") entries, got \", num_c_tensors);\n }\n auto is_output_dtype = [](transformer_engine::DType dtype) {\n return dtype == transformer_engine::DType::kBFloat16 ||\n dtype == transformer_engine::DType::kFloat16 ||\n dtype == transformer_engine::DType::kFloat32;\n };\n NVTE_CHECK(is_output_dtype(d_dtype), \"Grouped GEMM: D must be BF16, FP16, or FP32.\");\n\n // Parse config (if provided)\n GroupedMatmulConfig config_ = parse_grouped_gemm_config(config);\n\n // Select operand storage (row-wise vs column-wise) and adjust transpose flags to\n // mirror the non-grouped GEMM logic for FP8 layout constraints.\n auto A_sel = select_grouped_operand(inputA, static_cast(transa), /*is_A=*/true);\n auto B_sel = select_grouped_operand(inputB, static_cast(transb), /*is_A=*/false);\n // Workspaces: setup (pointer arrays) and cuBLAS\n auto workspace = setup_grouped_gemm_workspace(wspace_setup, wspace_cublas, num_tensors);\n\n MultiTensorGroupGemmInputArgs a_multi_tensor_args{};\n launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, /*C=*/nullptr, /*D=*/nullptr,\n alpha_tensor, beta_tensor, num_tensors, stream, a_multi_tensor_args,\n C_list, D_list, A_sel.dptr, d_dtype, d_dtype);\n\n // Compute average dimensions for heuristics\n int64_t avg_m_val =\n config_.avg_m.value_or(transa ? compute_avg_last_dim(inputA) : compute_avg_first_dim(inputA));\n int64_t avg_n_val =\n config_.avg_n.value_or(transb ? compute_avg_first_dim(inputB) : compute_avg_last_dim(inputB));\n int64_t avg_k_val =\n config_.avg_k.value_or(transa ? compute_avg_first_dim(inputA) : compute_avg_last_dim(inputA));\n const bool use_fp8 = is_fp8_dtype(A_sel.dtype) || is_fp8_dtype(B_sel.dtype);\n execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, d_dtype, num_tensors,\n config_.use_split_accumulator, use_fp8, avg_m_val, avg_n_val, avg_k_val,\n workspace.cublas_workspace_ptr, stream, config_.sm_count);\n}\n\nnamespace {\n\nvoid launch_grouped_bias_add(const transformer_engine::GroupedTensor *outputD,\n const transformer_engine::GroupedTensor *bias_tensor,\n const float *scale_ptr, bool use_scale, cudaStream_t stream) {\n using namespace transformer_engine;\n\n const char *api_name = use_scale ? \"Grouped scaled bias add\" : \"Grouped bias add\";\n\n NVTE_CHECK(outputD->num_tensors >= 1, api_name, \": number of tensors must be at least 1\");\n NVTE_CHECK(outputD->num_tensors == bias_tensor->num_tensors, api_name,\n \": output and bias must have the same number of tensors\");\n NVTE_CHECK(outputD->has_data(), api_name, \": output is missing row-wise data\");\n NVTE_CHECK(bias_tensor->has_data(), api_name, \": bias is missing row-wise data\");\n NVTE_CHECK(outputD->dtype() == bias_tensor->dtype(), api_name,\n \": output and bias must have matching dtypes\");\n NVTE_CHECK(bias_tensor->all_same_first_dim(), api_name,\n \": bias must have uniform first dim (expected 1)\");\n NVTE_CHECK(bias_tensor->get_common_first_dim() == 1, api_name, \": bias first dim must be 1\");\n NVTE_CHECK(outputD->all_same_last_dim() && bias_tensor->all_same_last_dim(), api_name,\n \": requires uniform last dim for output and bias\");\n NVTE_CHECK(outputD->get_common_last_dim() == bias_tensor->get_common_last_dim(), api_name,\n \": output and bias last dims must match\");\n\n const TensorShapeInfo d_meta = TensorShapeInfo::from_tensor(outputD);\n\n const DType dtype = outputD->dtype();\n constexpr int kThreads = 128;\n\n const int num_tensors = static_cast(outputD->num_tensors);\n NVTE_CHECK(num_tensors <= kMaxGroups, api_name, \" supports at most \", kMaxGroups,\n \" tensors, got \", num_tensors);\n const int total_rows = static_cast(outputD->logical_shape.data[0]);\n const int n = static_cast(outputD->get_common_last_dim());\n\n const size_t elem_size = typeToSize(dtype);\n const int kVec = (elem_size <= 2) ? 8 : 4;\n NVTE_CHECK(n % kVec == 0, api_name, \": requires last dim divisible by \", kVec);\n\n constexpr int kRowsPerBlock = 16;\n constexpr int kBlocksPerSM = 32;\n\n const int num_sms = transformer_engine::cuda::sm_count();\n\n const int block_cols = kThreads * kVec;\n const int col_blocks = (n + block_cols - 1) / block_cols;\n const int max_row_chunks = (total_rows + kRowsPerBlock - 1) / kRowsPerBlock;\n const int row_blocks = std::min(max_row_chunks, num_sms * kBlocksPerSM / col_blocks);\n const dim3 grid(std::max(1, row_blocks), col_blocks);\n const dim3 block(kThreads);\n\n auto launch = [&](auto use_scale_tag) {\n constexpr bool kUseScale = decltype(use_scale_tag)::value;\n if (elem_size <= 2) {\n constexpr int kV = 8;\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(dtype, T, {\n grouped_bias_add_kernel\n <<>>(static_cast(outputD->data.dptr),\n static_cast(bias_tensor->data.dptr),\n scale_ptr, d_meta, n, total_rows, num_tensors);\n });\n } else {\n constexpr int kV = 4;\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(dtype, T, {\n grouped_bias_add_kernel\n <<>>(static_cast(outputD->data.dptr),\n static_cast(bias_tensor->data.dptr),\n scale_ptr, d_meta, n, total_rows, num_tensors);\n });\n }\n };\n\n if (use_scale) {\n launch(std::true_type{});\n } else {\n launch(std::false_type{});\n }\n\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n} // namespace\n\nvoid nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_grouped_bias_add);\n using namespace transformer_engine;\n const GroupedTensor *outputD = convertNVTEGroupedTensorCheck(output);\n const GroupedTensor *bias_tensor = convertNVTEGroupedTensorCheck(bias);\n launch_grouped_bias_add(outputD, bias_tensor, nullptr, false, stream);\n}\n\nvoid nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias,\n const NVTETensor scale, cudaStream_t stream) {\n NVTE_API_CALL(nvte_grouped_scaled_bias_add);\n using namespace transformer_engine;\n const GroupedTensor *outputD = convertNVTEGroupedTensorCheck(output);\n const GroupedTensor *bias_tensor = convertNVTEGroupedTensorCheck(bias);\n const Tensor *scale_tensor = convertNVTETensorCheck(scale);\n\n NVTE_CHECK(scale_tensor->data.dptr != nullptr,\n \"Grouped scaled bias add: scale tensor must not be null\");\n NVTE_CHECK(scale_tensor->dtype() == DType::kFloat32,\n \"Grouped scaled bias add: scale must be float32\");\n NVTE_CHECK(scale_tensor->data.shape.size() == 1,\n \"Grouped scaled bias add: scale must be 1D, got \", scale_tensor->data.shape.size(),\n \"D\");\n const size_t total_rows = static_cast(outputD->logical_shape.data[0]);\n NVTE_CHECK(scale_tensor->data.shape[0] == total_rows, \"Grouped scaled bias add: scale size (\",\n scale_tensor->data.shape[0], \") must equal total rows (\", total_rows, \")\");\n\n const float *scale_ptr = static_cast(scale_tensor->data.dptr);\n launch_grouped_bias_add(outputD, bias_tensor, scale_ptr, true, stream);\n}\n\n#else // CUBLAS_VERSION < CUBLAS_GROUPED_GEMM_VERSION\n\nvoid nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedTensor B, int transb,\n const NVTEGroupedTensor C, NVTEGroupedTensor D, const NVTETensor alpha,\n const NVTETensor beta, NVTETensor workspace_setup,\n NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config,\n cudaStream_t stream) {\n NVTE_ERROR(\"nvte_grouped_gemm requires cuBLAS 13.3+, but compile-time cuBLAS version is \",\n CUBLAS_VERSION, \". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer.\");\n}\n\nvoid nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors,\n int transa, const NVTEGroupedTensor B, int transb,\n const NVTEGroupedTensor C, NVTEGroupedTensor D,\n const NVTETensor alpha, const NVTETensor beta,\n NVTETensor workspace_setup, NVTETensor workspace_cublas,\n NVTEGroupedMatmulConfig config, cudaStream_t stream) {\n NVTE_ERROR(\n \"nvte_grouped_gemm_with_discrete_inputA requires cuBLAS 13.3+, but compile-time \"\n \"cuBLAS version is \",\n CUBLAS_VERSION, \". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer.\");\n}\n\nvoid nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa,\n const NVTEGroupedTensor B, int transb,\n const NVTETensor *C_list, size_t num_c_tensors,\n NVTETensor *D_list, size_t num_d_tensors,\n const NVTETensor alpha, const NVTETensor beta,\n NVTETensor workspace_setup, NVTETensor workspace_cublas,\n NVTEGroupedMatmulConfig config, cudaStream_t stream) {\n NVTE_ERROR(\n \"nvte_grouped_gemm_with_discrete_out requires cuBLAS 13.3+, but compile-time \"\n \"cuBLAS version is \",\n CUBLAS_VERSION, \". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer.\");\n}\n\nvoid nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias,\n cudaStream_t stream) {\n NVTE_ERROR(\"nvte_grouped_bias_add requires cuBLAS 13.3+, but compile-time cuBLAS version is \",\n CUBLAS_VERSION, \". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer.\");\n}\n\nvoid nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias,\n const NVTETensor scale, cudaStream_t stream) {\n NVTE_ERROR(\n \"nvte_grouped_scaled_bias_add requires cuBLAS 13.3+, but compile-time cuBLAS version \"\n \"is \",\n CUBLAS_VERSION, \". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer.\");\n}\n\nsize_t nvte_get_grouped_gemm_setup_workspace_size(size_t num_tensors) {\n NVTE_ERROR(\n \"nvte_get_grouped_gemm_setup_workspace_size requires cuBLAS 13.3+, but compile-time cuBLAS \"\n \"version is \",\n CUBLAS_VERSION, \". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer.\");\n return 0;\n}\n\n#endif // CUBLAS_VERSION >= CUBLAS_GROUPED_GEMM_VERSION\n\nnamespace {\n\n__global__ void convert_int32_to_int64_kernel(const int32_t *src, int64_t *dst, size_t n) {\n size_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < n) dst[idx] = static_cast(src[idx]);\n}\n\n// Like convert_int32_to_int64_kernel but scales each element by multiplier.\n// Used to convert per-expert slice counts to per-expert row counts for multi-dim tensors.\n__global__ void convert_int32_to_int64_with_multiplier_kernel(const int32_t *src, int64_t *dst,\n size_t n, int64_t multiplier) {\n size_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < n) dst[idx] = static_cast(src[idx]) * multiplier;\n}\n\n// Computes exclusive prefix sums: offsets[0]=0, offsets[i]=sum(first_dims[0..i-1]*last_dim).\n// Produces n_groups+1 values. Single-threaded sequential scan; n_groups is typically small.\n__global__ void compute_grouped_tensor_offsets_kernel(const int64_t *first_dims, int64_t *offsets,\n size_t n_groups, int64_t last_dim) {\n offsets[0] = 0;\n for (size_t i = 0; i < n_groups; i++) {\n offsets[i + 1] = offsets[i] + first_dims[i] * last_dim;\n }\n}\n\n} // namespace\n\nvoid nvte_convert_int32_to_int64(const int32_t *src, int64_t *dst, size_t n, cudaStream_t stream) {\n NVTE_API_CALL(nvte_convert_int32_to_int64);\n if (n == 0) return;\n const int threads = 256;\n const int blocks = static_cast((n + threads - 1) / threads);\n convert_int32_to_int64_kernel<<>>(src, dst, n);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid nvte_convert_int32_to_int64_with_multiplier(const int32_t *src, int64_t *dst, size_t n,\n int64_t multiplier, cudaStream_t stream) {\n NVTE_API_CALL(nvte_convert_int32_to_int64_with_multiplier);\n if (n == 0) return;\n const int threads = 256;\n const int blocks = static_cast((n + threads - 1) / threads);\n convert_int32_to_int64_with_multiplier_kernel<<>>(src, dst, n,\n multiplier);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid nvte_compute_grouped_tensor_offsets(const int64_t *first_dims, int64_t *offsets,\n size_t n_groups, int64_t last_dim, cudaStream_t stream) {\n NVTE_API_CALL(nvte_compute_grouped_tensor_offsets);\n // Always write at least offsets[0]=0 (needed even for n_groups==0).\n compute_grouped_tensor_offsets_kernel<<<1, 1, 0, stream>>>(first_dims, offsets, n_groups,\n last_dim);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n", "output": "/*************************************************************************\n* Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n*\n* See LICENSE for license information.\n************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"../cast/mxfp8/swizzle.cuh\"\n#include \"../common.h\"\n#include \"../util/hip/hip_runtime.h\"\n#include \"../util/handle_manager.h\"\n#include \"../util/logging.h\"\n#include \"../util/vectorized_pointwise.h\"\n#include \"./config.h\"\n\nnamespace {\n\ninline void CreateCublasHandle(hipblasLtHandle_t *handle) {\n NVTE_CHECK_CUBLAS(hipblasLtCreate(handle));\n}\n\n} // namespace\n\n// MXFP8 support for grouped GEMM requires cuBLAS 13.3+\n#define CUBLAS_MXFP8_GROUPED_GEMM_VERSION 130300\n// BF16 support for grouped GEMM requires cuBLAS 13.3+\n#define CUBLAS_GROUPED_GEMM_VERSION 130300\n\n#if CUBLAS_VERSION >= CUBLAS_GROUPED_GEMM_VERSION\n\nnamespace {\n\n// Helper struct to pass per-tensor shape/offset info (pointer or uniform value)\nstruct TensorShapeInfo {\n const int64_t *first_dims; // nullptr if uniform\n const int64_t *last_dims; // nullptr if uniform\n const int64_t *offsets; // nullptr if need to compute\n int64_t uniform_first; // used if first_dims == nullptr\n int64_t uniform_last; // used if last_dims == nullptr\n\n // Create from GroupedTensor\n static TensorShapeInfo from_tensor(const transformer_engine::GroupedTensor *t) {\n const bool has_first = t->first_dims.has_data();\n const bool has_last = t->last_dims.has_data();\n // When per-tensor dims are not provided, we must be in the uniform-shape case.\n NVTE_CHECK(has_first || t->all_same_first_dim(),\n \"GroupedTensor is missing first_dims for varying shapes\");\n NVTE_CHECK(has_last || t->all_same_last_dim(),\n \"GroupedTensor is missing last_dims for varying shapes\");\n\n const int64_t *first_ptr =\n has_first ? static_cast(t->first_dims.dptr) : nullptr;\n const int64_t *last_ptr = has_last ? static_cast(t->last_dims.dptr) : nullptr;\n\n const int64_t uniform_first = has_first ? 0 : static_cast(t->get_common_first_dim());\n const int64_t uniform_last = has_last ? 0 : static_cast(t->get_common_last_dim());\n\n return {first_ptr, last_ptr,\n t->tensor_offsets.has_data() ? static_cast(t->tensor_offsets.dptr)\n : nullptr,\n uniform_first, uniform_last};\n }\n\n // Create for C tensor (uses D's dimensions, only has offsets)\n static TensorShapeInfo create_shape_info_for_C(const transformer_engine::GroupedTensor *C,\n const transformer_engine::GroupedTensor *D) {\n const bool has_first = D->first_dims.has_data();\n const bool has_last = D->last_dims.has_data();\n NVTE_CHECK(has_first || D->all_same_first_dim(),\n \"GroupedTensor D is missing first_dims for varying shapes\");\n NVTE_CHECK(has_last || D->all_same_last_dim(),\n \"GroupedTensor D is missing last_dims for varying shapes\");\n\n const int64_t *first_ptr =\n has_first ? static_cast(D->first_dims.dptr) : nullptr;\n const int64_t *last_ptr = has_last ? static_cast(D->last_dims.dptr) : nullptr;\n const int64_t uniform_first = has_first ? 0 : static_cast(D->get_common_first_dim());\n const int64_t uniform_last = has_last ? 0 : static_cast(D->get_common_last_dim());\n\n return {first_ptr, last_ptr,\n C->tensor_offsets.has_data() ? static_cast(C->tensor_offsets.dptr)\n : nullptr,\n uniform_first, uniform_last};\n }\n};\n\n// Helper functions to compute average dimensions for cuBLASLt algorithm-selection heuristics.\n//\n// logical_shape encoding (from build_grouped_tensor):\n// all_same: {num_tensors * M, N}\n// varying_first: {sum_of_first_dims, common_last}\n// varying_last: {common_first, sum_of_last_dims}\n// varying_both: {1, total_elements} <-- lossy, can't recover per-dim averages\n//\n// We use all_same_first/last_dim() + get_common_first/last_dim() to get exact\n// answers whenever possible, falling back to logical_shape division otherwise.\n// For varying_both, per-dim averages are unrecoverable without a D2H copy,\n// so we return 1 \u2014 a valid non-zero hint that won't skip work.\ninline int64_t compute_avg_first_dim(const transformer_engine::GroupedTensor *t) {\n if (t->all_same_first_dim()) {\n return static_cast(t->get_common_first_dim());\n }\n const int64_t n = static_cast(t->num_tensors);\n if (t->all_same_last_dim()) {\n // varying_first only: logical_shape = {sum_of_first_dims, common_last}\n return static_cast(t->logical_shape.data[0]) / n;\n }\n // varying_both: logical_shape = {1, total_elements}, no way to recover avg first dim\n return 1;\n}\n\ninline int64_t compute_avg_last_dim(const transformer_engine::GroupedTensor *t) {\n if (t->all_same_last_dim()) {\n // logical_shape[1] is the common N\n return static_cast(t->logical_shape.data[1]);\n }\n // When varying, logical_shape[1] should be sum of last dims if provided; otherwise fallback to avg via division.\n return static_cast(t->logical_shape.data[1]) / static_cast(t->num_tensors);\n}\n\n// Constants for grouped GEMM workspace (declared early for use in helpers)\nstatic constexpr size_t kGroupedGemmAlignment = 256;\nstatic constexpr size_t kGroupedGemmCublasWorkspaceSize = 32ull * 1024 * 1024; // 32 MiB\n\n// Workspace layout for grouped GEMM\nstruct GroupedGemmSetupWorkspace {\n void **A_ptrs;\n void **B_ptrs;\n void **C_ptrs;\n void **D_ptrs;\n float **alpha_ptrs;\n float **beta_ptrs;\n void **\n a_scale_inv_ptrs; // Per-tensor FP8 scale pointers for A (float* for tensor scaling, E8M0* for MXFP8)\n void **\n b_scale_inv_ptrs; // Per-tensor FP8 scale pointers for B (float* for tensor scaling, E8M0* for MXFP8)\n // Storage dimensions for cuBLAS matrix layouts\n int *a_rows;\n int *a_cols;\n int *b_rows;\n int *b_cols;\n int *d_rows; // M (first dim) - also used for C\n int *d_cols; // N (last dim) - also used for C\n\n // Initialize from workspace buffer\n // Layout: all pointer arrays first (16-byte aligned for cuBLAS), then int arrays\n static GroupedGemmSetupWorkspace from_buffers(char *setup_ws_ptr, size_t num_tensors) {\n GroupedGemmSetupWorkspace ws;\n size_t offset = 0;\n const size_t ptr_size = num_tensors * sizeof(void *);\n const size_t int_size = num_tensors * sizeof(int);\n constexpr size_t kPtrAlignment = 16; // cuBLAS requires 16-byte alignment for pointer arrays\n\n // Helper to align offset to kPtrAlignment\n auto align_offset = [&]() {\n offset = (offset + kPtrAlignment - 1) / kPtrAlignment * kPtrAlignment;\n };\n\n // Pointer arrays first (all 16-byte aligned for cuBLAS grouped GEMM)\n align_offset();\n ws.A_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n align_offset();\n ws.B_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n align_offset();\n ws.C_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n align_offset();\n ws.D_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n align_offset();\n ws.alpha_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n align_offset();\n ws.beta_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n align_offset();\n ws.a_scale_inv_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n align_offset();\n ws.b_scale_inv_ptrs = reinterpret_cast(setup_ws_ptr + offset);\n offset += ptr_size;\n\n // Int arrays for storage dimensions (4-byte aligned is fine)\n align_offset();\n ws.a_rows = reinterpret_cast(setup_ws_ptr + offset);\n offset += int_size;\n ws.a_cols = reinterpret_cast(setup_ws_ptr + offset);\n offset += int_size;\n ws.b_rows = reinterpret_cast(setup_ws_ptr + offset);\n offset += int_size;\n ws.b_cols = reinterpret_cast(setup_ws_ptr + offset);\n offset += int_size;\n ws.d_rows = reinterpret_cast(setup_ws_ptr + offset);\n offset += int_size;\n ws.d_cols = reinterpret_cast(setup_ws_ptr + offset);\n\n return ws;\n }\n\n // Calculate required size for setup workspace\n static size_t required_setup_size(size_t num_tensors, size_t alignment) {\n const size_t ptr_size = num_tensors * sizeof(void *);\n const size_t int_size = num_tensors * sizeof(int);\n constexpr size_t kPtrAlignment = 16; // Must match from_buffers\n\n // Layout: 8 ptr arrays (each 16-byte aligned), then 6 int arrays\n // Each ptr array takes ptr_size bytes but needs to start at 16-byte boundary\n auto aligned_ptr_size = ((ptr_size + kPtrAlignment - 1) / kPtrAlignment) * kPtrAlignment;\n size_t size = 8 * aligned_ptr_size + 6 * int_size;\n size = ((size + alignment - 1) / alignment) * alignment;\n return size;\n }\n};\n\ninline size_t validate_grouped_gemm_inputs(\n size_t num_tensors, std::initializer_list inputs,\n const transformer_engine::Tensor *alpha_tensor, const transformer_engine::Tensor *beta_tensor) {\n NVTE_CHECK(num_tensors >= 1, \"Grouped GEMM: number of tensors must be at least 1\");\n for (const auto *tensor : inputs) {\n NVTE_CHECK(tensor->num_tensors == num_tensors,\n \"Grouped GEMM: inputs must have the same number of tensors\");\n }\n\n const size_t alpha_numel = alpha_tensor->data.numel();\n const size_t beta_numel = beta_tensor->data.numel();\n NVTE_CHECK(alpha_numel == num_tensors, \"Grouped GEMM: alpha must have num_tensors (\", num_tensors,\n \") elements, got \", alpha_numel);\n NVTE_CHECK(beta_numel == num_tensors, \"Grouped GEMM: beta must have num_tensors (\", num_tensors,\n \") elements, got \", beta_numel);\n\n auto is_supported_input_dtype = [](transformer_engine::DType dtype) {\n return dtype == transformer_engine::DType::kFloat8E4M3 ||\n dtype == transformer_engine::DType::kFloat8E5M2 ||\n dtype == transformer_engine::DType::kBFloat16 ||\n dtype == transformer_engine::DType::kFloat16;\n };\n for (const auto *tensor : inputs) {\n if (tensor->has_data() || tensor->has_columnwise_data()) {\n NVTE_CHECK(is_supported_input_dtype(tensor->dtype()),\n \"Grouped GEMM inputs must be FP8, BF16, or FP16, got \",\n transformer_engine::to_string(tensor->dtype()), \".\");\n }\n }\n // Cross-operand consistency across all inputs (skip tensors without data).\n const transformer_engine::GroupedTensor *ref = nullptr;\n for (const auto *tensor : inputs) {\n if (tensor->has_data() || tensor->has_columnwise_data()) {\n ref = tensor;\n break;\n }\n }\n if (ref != nullptr) {\n const bool ref_is_fp8 = is_fp8_dtype(ref->dtype());\n const bool ref_is_mxfp8 = transformer_engine::is_mxfp_scaling(ref->scaling_mode);\n for (const auto *tensor : inputs) {\n if (!(tensor->has_data() || tensor->has_columnwise_data())) continue;\n NVTE_CHECK(is_fp8_dtype(tensor->dtype()) == ref_is_fp8,\n \"Grouped GEMM: A and B must both be FP8 or both be non-FP8.\");\n NVTE_CHECK(transformer_engine::is_mxfp_scaling(tensor->scaling_mode) == ref_is_mxfp8,\n \"Grouped GEMM: A and B must both use MXFP8 scaling or both use tensor scaling.\");\n if (ref_is_mxfp8) {\n NVTE_CHECK(tensor->with_gemm_swizzled_scales,\n \"MXFP8 grouped GEMM: scales must be swizzled for GEMM.\");\n }\n }\n }\n return num_tensors;\n}\n\ninline void validate_grouped_gemm_outputs(\n size_t num_tensors, std::initializer_list outputs) {\n auto is_output_dtype = [](transformer_engine::DType dtype) {\n return dtype == transformer_engine::DType::kBFloat16 ||\n dtype == transformer_engine::DType::kFloat16 ||\n dtype == transformer_engine::DType::kFloat32;\n };\n for (const auto *tensor : outputs) {\n if (tensor == nullptr) {\n continue;\n }\n NVTE_CHECK(tensor->num_tensors == num_tensors,\n \"Grouped GEMM: outputs must have the same number of tensors as inputs\");\n NVTE_CHECK(is_output_dtype(tensor->dtype()),\n \"Grouped GEMM: outputs must be BF16, FP16, or FP32.\");\n }\n}\n\ninline size_t grouped_gemm_setup_workspace_size(size_t num_tensors) {\n return GroupedGemmSetupWorkspace::required_setup_size(num_tensors, kGroupedGemmAlignment);\n}\n\ninline void check_grouped_gemm_requirements(const char *api_name) {\n const int current_device = transformer_engine::cuda::current_device();\n NVTE_CHECK(transformer_engine::cuda::sm_arch(current_device) >= 100, api_name,\n \" requires Blackwell (SM100) or newer architecture.\");\n NVTE_CHECK(transformer_engine::cuda::cublas_version() >= CUBLAS_GROUPED_GEMM_VERSION, api_name,\n \" requires cuBLAS 13.3+, but run-time cuBLAS version is \",\n transformer_engine::cuda::cublas_version());\n}\n\ninline transformer_engine::GroupedMatmulConfig parse_grouped_gemm_config(\n NVTEGroupedMatmulConfig config) {\n transformer_engine::GroupedMatmulConfig config_;\n if (config != nullptr) {\n config_ = *reinterpret_cast(config);\n }\n return config_;\n}\n\n// Select row-wise vs column-wise storage and adjust transpose flag for grouped GEMM.\n// Mirrors the non-grouped GEMM logic for FP8 layout handling (TN-only on Hopper) and\n// fallback to column-wise data when row-wise is absent.\n// Contains all information needed for GEMM setup - shape already accounts for storage layout.\nstruct GroupedOperandSelection {\n TensorShapeInfo shape; // Shape info with dims already swapped for columnwise if needed\n char *dptr = nullptr;\n void *scale_inv = nullptr; // Contiguous array of scales (input)\n transformer_engine::DType dtype = transformer_engine::DType::kNumTypes;\n NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING;\n bool with_gemm_swizzled_scales = false;\n bool trans = false;\n bool rowwise = true;\n};\n\nconstexpr int kMaxGroups = 64;\n// Arguments for the grouped GEMM kernel that operates on multiple output tensors.\nstruct MultiTensorGroupGemmOutputArgs {\n void *data_ptrs[kMaxGroups];\n int rows[kMaxGroups];\n int cols[kMaxGroups];\n};\n\n// Arguments for the grouped GEMM kernel that operates on multiple inputA tensors.\nstruct MultiTensorGroupGemmInputArgs {\n void *data_ptrs[kMaxGroups];\n void *scale_inv_ptrs[kMaxGroups];\n int rows[kMaxGroups];\n int cols[kMaxGroups];\n};\nstruct MultiTensorListInfo {\n bool all_row = true;\n bool all_col = true;\n transformer_engine::DType row_dtype = transformer_engine::DType::kNumTypes;\n transformer_engine::DType col_dtype = transformer_engine::DType::kNumTypes;\n NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING;\n bool with_gemm_swizzled_scales = false;\n};\n\nstruct OperandStorageChoice {\n bool use_rowwise = true;\n bool swap_dims = true;\n bool trans = false;\n};\n\ninline OperandStorageChoice choose_grouped_operand_storage(bool trans, bool is_A, bool is_mxfp8,\n bool is_fp8, bool non_tn_fp8_ok,\n bool has_row, bool has_col,\n const char *name) {\n NVTE_CHECK(has_row || has_col, \"Grouped GEMM: \", name,\n \" is missing both row-wise and column-wise data\");\n if (is_mxfp8) {\n if (is_A) {\n if (trans) {\n NVTE_CHECK(has_row, \"Grouped GEMM: MXFP8 transposed \", name, \" is missing row-wise data\");\n return {true, true, trans};\n }\n NVTE_CHECK(has_col, \"Grouped GEMM: MXFP8 non-transposed \", name,\n \" is missing column-wise data\");\n return {false, false, trans};\n }\n if (trans) {\n NVTE_CHECK(has_col, \"Grouped GEMM: MXFP8 transposed \", name, \" is missing column-wise data\");\n return {false, false, trans};\n }\n NVTE_CHECK(has_row, \"Grouped GEMM: MXFP8 non-transposed \", name, \" is missing row-wise data\");\n return {true, true, trans};\n }\n\n // Hopper-style TN-only FP8: force TN by switching layout and flipping transpose when needed.\n if (is_fp8 && !non_tn_fp8_ok) {\n if (is_A && !trans) {\n NVTE_CHECK(has_col, \"Grouped GEMM: \", name,\n \" is missing column-wise data needed for FP8 TN layout\");\n return {false, true, true};\n }\n if (!is_A && trans) {\n NVTE_CHECK(has_col, \"Grouped GEMM: \", name,\n \" is missing column-wise data needed for FP8 TN layout\");\n return {false, true, false};\n }\n }\n\n // If only column-wise data is available, mirror the transpose flag (pre-transposed storage).\n if (!has_row && has_col) {\n NVTE_CHECK(!is_fp8 || non_tn_fp8_ok,\n \"Grouped GEMM: FP8 on Hopper requires row-wise data for this transpose config.\");\n return {false, true, !trans};\n }\n\n NVTE_CHECK(has_row, \"Grouped GEMM: \", name, \" is missing row-wise data\");\n return {true, true, trans};\n}\n\n// Build Kernel Arguments detailing out addresses and other metadata for list of C/D tensors\n// passed to the grouped GEMM kernel. Use-case: C/D --> List of wgrads for experts in MOE\ninline MultiTensorGroupGemmOutputArgs build_grouped_gemm_multi_out_args(\n const NVTETensor *tensor_list, size_t list_size, size_t expected_num_tensors,\n transformer_engine::DType expected_dtype, const char *name) {\n MultiTensorGroupGemmOutputArgs args{};\n if (list_size == 0) {\n NVTE_CHECK(tensor_list == nullptr, \"Grouped GEMM: \", name, \"_list provided with num_\", name,\n \"_tensors=0\");\n return args;\n }\n NVTE_CHECK(tensor_list != nullptr, \"Grouped GEMM: \", name, \"_list is null but num_\", name,\n \"_tensors=\", list_size);\n NVTE_CHECK(list_size == expected_num_tensors, \"Grouped GEMM: \", name,\n \"_list must have num_tensors (\", expected_num_tensors, \") entries, got \", list_size);\n NVTE_CHECK(list_size <= static_cast(kMaxGroups), \"Grouped GEMM: \", name,\n \"_list supports up to \", kMaxGroups, \" tensors per kernel, got \", list_size);\n\n for (size_t i = 0; i < list_size; ++i) {\n const transformer_engine::Tensor *t =\n transformer_engine::convertNVTETensorCheck(tensor_list[i]);\n NVTE_CHECK(t->has_data(), \"Grouped GEMM: \", name, \"_list tensor \", i, \" has no data\");\n NVTE_CHECK(t->dtype() == expected_dtype, \"Grouped GEMM: \", name, \"_list tensor \", i,\n \" dtype mismatch. Expected \", transformer_engine::to_string(expected_dtype), \" got \",\n transformer_engine::to_string(t->dtype()));\n const auto &shape = t->shape();\n NVTE_CHECK(shape.size() == 2, \"Grouped GEMM: \", name, \"_list tensor \", i, \" must be 2D.\");\n args.data_ptrs[i] = t->data.dptr;\n args.rows[i] = static_cast(shape[1]);\n args.cols[i] = static_cast(shape[0]);\n }\n return args;\n}\n\n// Build Kernel Arguments detailing out addresses and other metadata for list of A tensors\n// passed to the grouped GEMM kernel. Use-case: A --> List of Expert weights\ninline MultiTensorGroupGemmInputArgs build_grouped_gemm_multi_inputA_args(\n const NVTETensor *tensor_list, size_t list_size, bool use_rowwise, bool is_fp8,\n int64_t *avg_first_dim, int64_t *avg_last_dim, const char *name) {\n using namespace transformer_engine;\n MultiTensorGroupGemmInputArgs args{};\n *avg_first_dim = 0;\n *avg_last_dim = 0;\n if (list_size == 0) {\n return args;\n }\n for (size_t i = 0; i < list_size; ++i) {\n const transformer_engine::Tensor *t =\n transformer_engine::convertNVTETensorCheck(tensor_list[i]);\n const transformer_engine::SimpleTensor &data = use_rowwise ? t->data : t->columnwise_data;\n const transformer_engine::SimpleTensor &scale_inv =\n use_rowwise ? t->scale_inv : t->columnwise_scale_inv;\n NVTE_CHECK(data.has_data(), \"Grouped GEMM: \", name, \"_list tensor \", i,\n \" is missing required data.\");\n NVTE_CHECK(data.shape.size() == 2, \"Grouped GEMM: \", name, \"_list tensor \", i, \" must be 2D.\");\n args.data_ptrs[i] = data.dptr;\n args.rows[i] = static_cast(data.shape[1]);\n args.cols[i] = static_cast(data.shape[0]);\n *avg_first_dim += static_cast(data.shape[0]);\n *avg_last_dim += static_cast(data.shape[1]);\n\n if (is_fp8) {\n NVTE_CHECK(scale_inv.has_data(), \"Grouped GEMM: \", name, \"_list tensor \", i,\n \" requires scale_inv for FP8.\");\n args.scale_inv_ptrs[i] = scale_inv.dptr;\n } else {\n args.scale_inv_ptrs[i] = nullptr;\n }\n }\n *avg_first_dim /= static_cast(list_size);\n *avg_last_dim /= static_cast(list_size);\n return args;\n}\n\ninline MultiTensorListInfo validate_grouped_gemm_multi_inputA_list(const NVTETensor *tensor_list,\n size_t list_size,\n size_t expected_num_tensors,\n const char *name) {\n using namespace transformer_engine;\n MultiTensorListInfo info{};\n if (list_size == 0) {\n NVTE_CHECK(tensor_list == nullptr, \"Grouped GEMM: \", name, \"_list provided with num_\", name,\n \"_tensors=0\");\n return info;\n }\n NVTE_CHECK(tensor_list != nullptr, \"Grouped GEMM: \", name, \"_list is null but num_\", name,\n \"_tensors=\", list_size);\n NVTE_CHECK(list_size == expected_num_tensors, \"Grouped GEMM: \", name,\n \"_list must have num_tensors (\", expected_num_tensors, \") entries, got \", list_size);\n NVTE_CHECK(list_size <= static_cast(kMaxGroups), \"Grouped GEMM: \", name,\n \"_list supports up to \", kMaxGroups, \" tensors per kernel, got \", list_size);\n\n const transformer_engine::Tensor *t0 = transformer_engine::convertNVTETensorCheck(tensor_list[0]);\n info.scaling_mode = t0->scaling_mode;\n info.with_gemm_swizzled_scales = t0->with_gemm_swizzled_scales;\n const bool mxfp8 = transformer_engine::is_mxfp_scaling(info.scaling_mode);\n NVTE_CHECK(info.scaling_mode == NVTE_DELAYED_TENSOR_SCALING || mxfp8,\n \"Grouped GEMM: input list only supports tensor scaling or MXFP8.\");\n\n for (size_t i = 0; i < list_size; ++i) {\n const transformer_engine::Tensor *t =\n transformer_engine::convertNVTETensorCheck(tensor_list[i]);\n NVTE_CHECK(t->scaling_mode == info.scaling_mode, \"Grouped GEMM: \", name,\n \"_list tensors must share the same scaling mode.\");\n NVTE_CHECK(t->with_gemm_swizzled_scales == info.with_gemm_swizzled_scales,\n \"Grouped GEMM: \", name, \"_list tensors must share GEMM swizzled scale state.\");\n\n if (t->has_data()) {\n if (info.row_dtype == DType::kNumTypes) {\n info.row_dtype = t->data.dtype;\n }\n // Check all tensors have the same dtype\n NVTE_CHECK(t->data.dtype == info.row_dtype, \"Grouped GEMM: \", name,\n \"_list rowwise dtypes must match.\");\n } else {\n // All tensors must have either data or columnwise data\n info.all_row = false;\n }\n\n if (t->has_columnwise_data()) {\n if (info.col_dtype == DType::kNumTypes) {\n info.col_dtype = t->columnwise_data.dtype;\n }\n NVTE_CHECK(t->columnwise_data.dtype == info.col_dtype, \"Grouped GEMM: \", name,\n \"_list columnwise dtypes must match.\");\n } else {\n // All tensors must have either data or columnwise data\n info.all_col = false;\n }\n }\n\n return info;\n}\n\n// Helper to create TensorShapeInfo from a GroupedTensor, optionally swapping first/last dims.\n// When swap_dims=true, first_dims and last_dims are swapped to account for columnwise storage.\n// Note: tensor_offsets are the same for rowwise and columnwise data (same element count per tensor).\ninline TensorShapeInfo create_shape_info(const transformer_engine::GroupedTensor *t,\n bool swap_dims) {\n const bool has_first = t->first_dims.has_data();\n const bool has_last = t->last_dims.has_data();\n NVTE_CHECK(has_first || t->all_same_first_dim(),\n \"GroupedTensor is missing first_dims for varying shapes\");\n NVTE_CHECK(has_last || t->all_same_last_dim(),\n \"GroupedTensor is missing last_dims for varying shapes\");\n\n const int64_t *first_ptr = has_first ? static_cast(t->first_dims.dptr) : nullptr;\n const int64_t *last_ptr = has_last ? static_cast(t->last_dims.dptr) : nullptr;\n const int64_t uniform_first = has_first ? 0 : static_cast(t->get_common_first_dim());\n const int64_t uniform_last = has_last ? 0 : static_cast(t->get_common_last_dim());\n\n const int64_t *offsets_ptr =\n t->tensor_offsets.has_data() ? static_cast(t->tensor_offsets.dptr) : nullptr;\n\n if (swap_dims) {\n // Swap first/last to account for columnwise (transposed) storage\n return {last_ptr, first_ptr, offsets_ptr, uniform_last, uniform_first};\n }\n return {first_ptr, last_ptr, offsets_ptr, uniform_first, uniform_last};\n}\n\ninline GroupedOperandSelection select_grouped_operand(const transformer_engine::GroupedTensor *t,\n bool trans, bool is_A) {\n using namespace transformer_engine;\n const bool has_row = t->has_data();\n const bool has_col = t->has_columnwise_data();\n\n if (!has_row && !has_col) {\n GroupedOperandSelection sel{};\n sel.trans = trans;\n sel.scaling_mode = t->scaling_mode;\n sel.dtype = t->dtype();\n sel.shape = create_shape_info(t, /*swap_dims=*/false);\n return sel;\n }\n\n const auto sm = t->scaling_mode;\n const bool mxfp8 = is_mxfp_scaling(sm);\n\n // Validate scaling mode\n NVTE_CHECK(sm == NVTE_DELAYED_TENSOR_SCALING || mxfp8,\n \"Grouped GEMM is only supported with bf16, fp8 tensor scaling and MXFP8\");\n\n const DType row_dtype = t->data.dtype;\n const DType col_dtype = t->columnwise_data.dtype;\n GroupedOperandSelection sel{};\n sel.trans = trans;\n sel.scaling_mode = sm;\n sel.with_gemm_swizzled_scales = t->with_gemm_swizzled_scales;\n\n const DType rep_dtype = has_row ? row_dtype : col_dtype;\n const bool is_fp8 = is_fp8_dtype(rep_dtype);\n const bool non_tn_fp8_ok = nvte_is_non_tn_fp8_gemm_supported();\n\n // Helper to select columnwise storage.\n // swap_dims=true (default): swap first/last dims in shape info (used when columnwise == transposed).\n // swap_dims=false: keep original dims (MXFP8: columnwise data has different scale direction,\n // but the logical matrix shape and transpose flag remain unchanged).\n auto use_columnwise = [&](bool swap_dims = true) {\n sel.dptr = static_cast(t->columnwise_data.dptr);\n sel.scale_inv = t->columnwise_scale_inv.dptr;\n sel.dtype = col_dtype;\n sel.rowwise = false;\n sel.shape = create_shape_info(t, swap_dims);\n };\n\n // Helper to select row-wise storage\n auto use_rowwise = [&]() {\n sel.dptr = static_cast(t->data.dptr);\n sel.scale_inv = t->scale_inv.dptr;\n sel.dtype = row_dtype;\n sel.rowwise = true;\n sel.shape = create_shape_info(t, /*swap_dims=*/false);\n };\n\n const auto choice = choose_grouped_operand_storage(trans, is_A, mxfp8, is_fp8, non_tn_fp8_ok,\n has_row, has_col, is_A ? \"A\" : \"B\");\n sel.trans = choice.trans;\n if (choice.use_rowwise) {\n use_rowwise();\n } else {\n use_columnwise(choice.swap_dims);\n }\n return sel;\n}\n\ninline void *validate_and_get_workspace_ptr(transformer_engine::Tensor *ws, size_t required_size,\n const char *workspace_name) {\n NVTE_CHECK(ws != nullptr, workspace_name, \" tensor is null.\");\n const size_t provided_size = get_buffer_size_bytes(ws->data.numel(), ws->data.dtype);\n NVTE_CHECK(provided_size >= required_size, \"Grouped GEMM: Insufficient \", workspace_name,\n \". Required: \", required_size, \" bytes, Available: \", provided_size, \" bytes.\");\n return ws->data.dptr;\n}\n\ninline void init_matrix_layouts(\n hipblasLtMatrixLayoutOpaque_t &descA, hipblasLtMatrixLayoutOpaque_t &descB,\n hipblasLtMatrixLayoutOpaque_t &descC, hipblasLtMatrixLayoutOpaque_t &descD,\n const GroupedGemmSetupWorkspace &ws, const GroupedOperandSelection &A_sel,\n const GroupedOperandSelection &B_sel, transformer_engine::DType d_dtype, size_t num_tensors) {\n const hipDataType A_type = get_cuda_dtype(A_sel.dtype);\n const hipDataType B_type = get_cuda_dtype(B_sel.dtype);\n const hipDataType D_type = get_cuda_dtype(d_dtype);\n\n // Storage dimensions computed by kernel, leading dimension = rows\n NVTE_CHECK_CUBLAS(cublasLtGroupedMatrixLayoutInit(&descA, A_type, num_tensors, ws.a_rows,\n ws.a_cols, ws.a_rows));\n NVTE_CHECK_CUBLAS(cublasLtGroupedMatrixLayoutInit(&descB, B_type, num_tensors, ws.b_rows,\n ws.b_cols, ws.b_rows));\n NVTE_CHECK_CUBLAS(cublasLtGroupedMatrixLayoutInit(&descC, D_type, num_tensors, ws.d_rows,\n ws.d_cols, ws.d_rows));\n NVTE_CHECK_CUBLAS(cublasLtGroupedMatrixLayoutInit(&descD, D_type, num_tensors, ws.d_rows,\n ws.d_cols, ws.d_rows));\n}\n\ninline void init_matmul_desc(hipblasLtMatmulDescOpaque_t &matmulDesc, hipblasOperation_t op_A,\n hipblasOperation_t op_B, bool use_fp8, bool use_split_accumulator) {\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescInit(&matmulDesc, HIPBLAS_COMPUTE_32F, HIP_R_32F));\n\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(&matmulDesc, HIPBLASLT_MATMUL_DESC_TRANSA, &op_A,\n sizeof(op_A)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(&matmulDesc, HIPBLASLT_MATMUL_DESC_TRANSB, &op_B,\n sizeof(op_B)));\n\n hipblasLtPointerMode_t pointer_mode = HIPBLASLT_POINTER_MODE_DEVICE;\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(&matmulDesc, HIPBLASLT_MATMUL_DESC_POINTER_MODE,\n &pointer_mode, sizeof(pointer_mode)));\n\n int64_t alphabeta_batch_stride = 1;\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(&matmulDesc,\n CUBLASLT_MATMUL_DESC_ALPHA_BATCH_STRIDE,\n &alphabeta_batch_stride, sizeof(int64_t)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(&matmulDesc,\n CUBLASLT_MATMUL_DESC_BETA_BATCH_STRIDE,\n &alphabeta_batch_stride, sizeof(int64_t)));\n\n // Fast accumulation is only supported for FP8 (mirrors non-grouped GEMM logic).\n int8_t fastAccuMode = use_split_accumulator ? 0 : static_cast(use_fp8);\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_FAST_ACCUM,\n &fastAccuMode, sizeof(fastAccuMode)));\n}\n\n// Configures cuBLAS for MXFP8 grouped GEMM: sets VEC32_UE8M0 scale mode and scale pointers\n// for both A and B.\ninline void set_mxfp8_scale_pointers(hipblasLtMatmulDescOpaque_t &matmulDesc,\n void **a_scale_inv_ptrs, void **b_scale_inv_ptrs) {\n#if CUBLAS_VERSION >= CUBLAS_MXFP8_GROUPED_GEMM_VERSION\n NVTE_CHECK(transformer_engine::cuda::cublas_version() >= CUBLAS_MXFP8_GROUPED_GEMM_VERSION,\n \"MXFP8 grouped GEMM requires cuBLAS \", CUBLAS_MXFP8_GROUPED_GEMM_VERSION,\n \"+, but run-time cuBLAS version is \", transformer_engine::cuda::cublas_version());\n const hipblasLtMatmulMatrixScale_t scale_mode = HIPBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0;\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(&matmulDesc, HIPBLASLT_MATMUL_DESC_A_SCALE_MODE,\n &scale_mode, sizeof(scale_mode)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(&matmulDesc, HIPBLASLT_MATMUL_DESC_B_SCALE_MODE,\n &scale_mode, sizeof(scale_mode)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(&matmulDesc,\n HIPBLASLT_MATMUL_DESC_A_SCALE_POINTER,\n &a_scale_inv_ptrs, sizeof(a_scale_inv_ptrs)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(&matmulDesc,\n HIPBLASLT_MATMUL_DESC_B_SCALE_POINTER,\n &b_scale_inv_ptrs, sizeof(b_scale_inv_ptrs)));\n#else\n NVTE_CHECK(false, \"MXFP8 grouped GEMM requires cuBLAS \", CUBLAS_MXFP8_GROUPED_GEMM_VERSION,\n \"+, but compile-time cuBLAS version is \", CUBLAS_VERSION);\n#endif // CUBLAS_VERSION >= CUBLAS_MXFP8_GROUPED_GEMM_VERSION\n}\n\n// Configures cuBLAS for tensor-scaling FP8 grouped GEMM: sets PER_BATCH_SCALAR_32F scale mode\n// and scale pointers for A and B. Both operands are guaranteed FP8 by the caller.\ninline void set_fp8_scale_pointers(hipblasLtMatmulDescOpaque_t &matmulDesc, void **a_scale_inv_ptrs,\n void **b_scale_inv_ptrs) {\n const hipblasLtMatmulMatrixScale_t scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_PER_BATCH_SCALAR_32F;\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(&matmulDesc, HIPBLASLT_MATMUL_DESC_A_SCALE_MODE,\n &scale_mode, sizeof(scale_mode)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(&matmulDesc,\n HIPBLASLT_MATMUL_DESC_A_SCALE_POINTER,\n &a_scale_inv_ptrs, sizeof(a_scale_inv_ptrs)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(&matmulDesc, HIPBLASLT_MATMUL_DESC_B_SCALE_MODE,\n &scale_mode, sizeof(scale_mode)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(&matmulDesc,\n HIPBLASLT_MATMUL_DESC_B_SCALE_POINTER,\n &b_scale_inv_ptrs, sizeof(b_scale_inv_ptrs)));\n}\ninline hipblasLtMatmulAlgo_t select_grouped_gemm_algo(hipblasLtHandle_t handle,\n hipblasLtMatmulDescOpaque_t &matmulDesc,\n hipblasLtMatrixLayoutOpaque_t &descA,\n hipblasLtMatrixLayoutOpaque_t &descB,\n hipblasLtMatrixLayoutOpaque_t &descC,\n hipblasLtMatrixLayoutOpaque_t &descD,\n int64_t avg_m, int64_t avg_n, int64_t avg_k) {\n hipblasLtMatmulPreferenceOpaque_t preference;\n NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceInit(&preference));\n NVTE_CHECK_CUBLAS(\n hipblasLtMatmulPreferenceSetAttribute(&preference, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,\n &kGroupedGemmCublasWorkspaceSize, sizeof(size_t)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulPreferenceSetAttribute(\n &preference, CUBLASLT_MATMUL_PREF_GROUPED_DESC_D_AVERAGE_ROWS, &avg_m, sizeof(int64_t)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulPreferenceSetAttribute(\n &preference, CUBLASLT_MATMUL_PREF_GROUPED_DESC_D_AVERAGE_COLS, &avg_n, sizeof(int64_t)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulPreferenceSetAttribute(\n &preference, CUBLASLT_MATMUL_PREF_GROUPED_AVERAGE_REDUCTION_DIM, &avg_k, sizeof(int64_t)));\n\n hipblasLtMatmulHeuristicResult_t heuristicResult;\n int returnedResults = 0;\n auto status = hipblasLtMatmulAlgoGetHeuristic(handle, &matmulDesc, &descA, &descB, &descC, &descD,\n &preference, 1, &heuristicResult, &returnedResults);\n NVTE_CHECK(status != HIPBLAS_STATUS_NOT_SUPPORTED,\n \"Unable to find suitable cuBLAS grouped GEMM algorithm\");\n NVTE_CHECK_CUBLAS(status);\n NVTE_CHECK(returnedResults > 0, \"No suitable algorithm found for grouped GEMM\");\n return heuristicResult.algo;\n}\n\nstruct GroupedGemmWorkspace {\n GroupedGemmSetupWorkspace setup_workspace;\n void *cublas_workspace_ptr = nullptr;\n size_t num_tensors = 0;\n};\n\ninline GroupedGemmWorkspace setup_grouped_gemm_workspace(transformer_engine::Tensor *wspace_setup,\n transformer_engine::Tensor *wspace_cublas,\n size_t num_tensors) {\n const size_t setup_workspace_size = grouped_gemm_setup_workspace_size(num_tensors);\n const size_t cublas_workspace_size = kGroupedGemmCublasWorkspaceSize;\n void *setup_workspace_ptr = validate_and_get_workspace_ptr(wspace_setup, setup_workspace_size,\n \"Grouped GEMM setup workspace\");\n void *cublas_workspace_ptr = validate_and_get_workspace_ptr(wspace_cublas, cublas_workspace_size,\n \"Grouped GEMM cuBLAS workspace\");\n auto setup_workspace = GroupedGemmSetupWorkspace::from_buffers(\n static_cast(setup_workspace_ptr), num_tensors);\n return {std::move(setup_workspace), cublas_workspace_ptr, num_tensors};\n}\n\ninline void execute_grouped_gemm(const GroupedGemmSetupWorkspace &setup_workspace,\n const GroupedOperandSelection &A_sel,\n const GroupedOperandSelection &B_sel,\n transformer_engine::DType d_dtype, size_t num_tensors,\n bool use_split_accumulator, bool use_fp8, int64_t avg_m_val,\n int64_t avg_n_val, int64_t avg_k_val, void *cublas_workspace_ptr,\n hipStream_t stream, int math_sm_count = 0) {\n using cublasHandleManager =\n transformer_engine::detail::HandleManager;\n hipblasLtHandle_t handle = cublasHandleManager::Instance().GetHandle();\n\n hipblasOperation_t op_A = A_sel.trans ? HIPBLAS_OP_T : HIPBLAS_OP_N;\n hipblasOperation_t op_B = B_sel.trans ? HIPBLAS_OP_T : HIPBLAS_OP_N;\n\n hipblasLtMatrixLayoutOpaque_t descA, descB, descC, descD;\n init_matrix_layouts(descA, descB, descC, descD, setup_workspace, A_sel, B_sel, d_dtype,\n num_tensors);\n\n hipblasLtMatmulDescOpaque_t matmulDesc;\n init_matmul_desc(matmulDesc, op_A, op_B, use_fp8, use_split_accumulator);\n if (transformer_engine::is_mxfp_scaling(A_sel.scaling_mode)) {\n set_mxfp8_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs,\n setup_workspace.b_scale_inv_ptrs);\n } else if (use_fp8) {\n set_fp8_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs,\n setup_workspace.b_scale_inv_ptrs);\n }\n if (math_sm_count != 0) {\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n &matmulDesc, CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET, &math_sm_count, sizeof(math_sm_count)));\n }\n hipblasLtMatmulAlgo_t algo = select_grouped_gemm_algo(handle, matmulDesc, descA, descB, descC,\n descD, avg_m_val, avg_n_val, avg_k_val);\n\n NVTE_CHECK_CUBLAS(hipblasLtMatmul(handle, &matmulDesc, setup_workspace.alpha_ptrs,\n setup_workspace.A_ptrs, &descA, setup_workspace.B_ptrs, &descB,\n setup_workspace.beta_ptrs, setup_workspace.C_ptrs, &descC,\n setup_workspace.D_ptrs, &descD, &algo, cublas_workspace_ptr,\n kGroupedGemmCublasWorkspaceSize, stream));\n}\n\n// Device helper: compute the element offset for tensor `idx` given shape metadata.\n// Three cases:\n// 1. Explicit per-tensor offset array provided \u2192 use it directly.\n// 2. Per-tensor first/last dims provided but no offsets \u2192 cumulative sum of (first*last) products.\n// 3. Fully uniform shapes \u2192 idx * uniform_first * uniform_last.\n__forceinline__ __device__ int64_t compute_grouped_tensor_offset(const TensorShapeInfo &meta,\n size_t idx) {\n if (meta.offsets) {\n return meta.offsets[idx];\n } else if (meta.first_dims != nullptr || meta.last_dims != nullptr) {\n // offset[i] = sum_{j < i} (first_dims[j] * last_dims[j])\n int64_t cumsum = 0;\n for (size_t i = 0; i < idx; i++) {\n int64_t f = meta.first_dims ? meta.first_dims[i] : meta.uniform_first;\n int64_t l = meta.last_dims ? meta.last_dims[i] : meta.uniform_last;\n cumsum += f * l;\n }\n return cumsum;\n } else {\n return static_cast(idx) * meta.uniform_first * meta.uniform_last;\n }\n}\n\n__forceinline__ __device__ int64_t padded_mxfp8_scale_inv_bytes(int64_t first, int64_t last,\n bool rowwise) {\n namespace mxfp8_swizzle = transformer_engine::dispatch::mxfp8::swizzle;\n constexpr int64_t kMxfp8BlockSize = 32;\n // x is the dimension along which quantization is applied, y is other dimension\n const int64_t scale_tile_y = static_cast(mxfp8_swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_Y);\n const int64_t scale_tile_x = static_cast(mxfp8_swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_X);\n // Padded byte size of the swizzled MXFP8 scale_inv for a single tensor with data\n // shape (first, last). Rowwise scales use rows=first, cols=last; columnwise\n // scales swap the orientation since they are stored in column-major order.\n const int64_t scale_dim_y = rowwise ? first : last;\n const int64_t padded_scale_dim_y =\n ((scale_dim_y + scale_tile_y - 1) / scale_tile_y) * scale_tile_y;\n const int64_t data_dim_x = rowwise ? last : first;\n const int64_t scale_dim_x = (data_dim_x + kMxfp8BlockSize - 1) / kMxfp8BlockSize;\n const int64_t padded_scale_dim_x =\n ((scale_dim_x + scale_tile_x - 1) / scale_tile_x) * scale_tile_x;\n // MXFP8 scales are E8M0 (1 byte per element), so element count == byte count.\n return padded_scale_dim_y * padded_scale_dim_x;\n}\n\n// Device helper: byte offset into a contiguous grouped MXFP8 scale_inv buffer for\n// tensor `idx`. Each expert's scale_inv is expected to be padded\n// to the 128x4 swizzled layout.\n__forceinline__ __device__ int64_t compute_grouped_tensor_mxfp8_scale_inv_offset(\n const TensorShapeInfo &meta, size_t idx, bool rowwise) {\n if (meta.first_dims != nullptr || meta.last_dims != nullptr) {\n int64_t cumsum = 0;\n for (size_t i = 0; i < idx; i++) {\n const int64_t f = meta.first_dims ? meta.first_dims[i] : meta.uniform_first;\n const int64_t l = meta.last_dims ? meta.last_dims[i] : meta.uniform_last;\n cumsum += padded_mxfp8_scale_inv_bytes(f, l, rowwise);\n }\n return cumsum;\n }\n return static_cast(idx) *\n padded_mxfp8_scale_inv_bytes(meta.uniform_first, meta.uniform_last, rowwise);\n}\n\n// Linear scan to find which tensor contains the given row.\n// Returns the tensor index and writes the exclusive end-row of that tensor to *out_tensor_row_end.\n__forceinline__ __device__ int find_tensor_for_row(const int64_t *first_dims, int64_t uniform_first,\n int row, int num_tensors,\n int *out_tensor_row_end) {\n int offset = 0;\n for (int i = 0; i < num_tensors; i++) {\n int dim = first_dims ? static_cast(first_dims[i]) : static_cast(uniform_first);\n offset += dim;\n if (row < offset) {\n *out_tensor_row_end = offset;\n return i;\n }\n }\n *out_tensor_row_end = offset;\n return num_tensors - 1;\n}\n\n// Kernel that performs (optionally scaled) bias addition to Grouped GEMM output tensors.\n// SM-filling grid with grid-stride over row chunks.\n// 2D grid: blockIdx.x = SM-filling row blocks, blockIdx.y = column chunk.\n// Each block grid-strides over kRowsPerBlock-sized row chunks, processing\n// all chunks that map to it. Safe when sum(first_dims) <= total_rows.\ntemplate \n__global__ void grouped_bias_add_kernel(char *__restrict__ d_base,\n const char *__restrict__ bias_base,\n const float *__restrict__ scale_base,\n TensorShapeInfo d_meta, int n, int total_rows,\n int num_tensors) {\n using VecStorage = transformer_engine::VectorizedStorage;\n using VecType = typename VecStorage::LType;\n\n const int tid = static_cast(threadIdx.x);\n const int row_bid = static_cast(blockIdx.x);\n const int col_bid = static_cast(blockIdx.y);\n const int row_grid_stride = static_cast(gridDim.x);\n\n // Single-warp reduction to compute valid_rows = sum(first_dims).\n // kMaxGroups <= 64 so warp 0 (32 lanes) covers it with <=2 loads each.\n __shared__ int s_valid_rows;\n if (tid < 32) {\n int local_sum = 0;\n for (int i = tid; i < num_tensors; i += 32) {\n local_sum += d_meta.first_dims ? static_cast(d_meta.first_dims[i])\n : static_cast(d_meta.uniform_first);\n }\n for (int offset = 16; offset > 0; offset >>= 1) {\n local_sum += __shfl_down_sync(0xffffffff, local_sum, offset);\n }\n if (tid == 0) s_valid_rows = local_sum;\n }\n __syncthreads();\n const int valid_rows = s_valid_rows;\n\n const int block_cols = kBlockDim * kVec;\n const int col = col_bid * block_cols + tid * kVec;\n if (col >= n) return;\n\n T *__restrict__ d = reinterpret_cast(d_base);\n const T *__restrict__ bias = reinterpret_cast(bias_base);\n\n // Grid-stride loop over row chunks.\n for (int chunk_start = row_bid * kRowsPerBlock; chunk_start < valid_rows;\n chunk_start += row_grid_stride * kRowsPerBlock) {\n const int row_start = chunk_start;\n const int row_end = min(row_start + kRowsPerBlock, valid_rows);\n\n // Linear scan to find the starting row's tensor and its boundary.\n int tensor_row_end;\n int tensor_idx = find_tensor_for_row(d_meta.first_dims, d_meta.uniform_first, row_start,\n num_tensors, &tensor_row_end);\n int bias_idx = tensor_idx * n;\n\n VecStorage b_in;\n\n // Walk tensor segments within this chunk's row range.\n int seg_start = row_start;\n while (seg_start < row_end) {\n while (tensor_idx < num_tensors - 1 && tensor_row_end <= seg_start) {\n tensor_idx++;\n bias_idx += n;\n int dim = d_meta.first_dims ? static_cast(d_meta.first_dims[tensor_idx])\n : static_cast(d_meta.uniform_first);\n tensor_row_end += dim;\n }\n b_in.scratch_.aligned = *reinterpret_cast(bias + bias_idx + col);\n const int seg_end = min(tensor_row_end, row_end);\n\n for (int row = seg_start; row < seg_end; row++) {\n T *d_ptr = d + row * n + col;\n VecStorage d_in;\n d_in.scratch_.aligned = *reinterpret_cast(d_ptr);\n\n [[maybe_unused]] float s_val;\n if constexpr (UseScale) s_val = scale_base[row];\n\n#pragma unroll\n for (int i = 0; i < kVec; ++i) {\n if constexpr (UseScale) {\n d_in.scratch_.separate[i] =\n static_cast(fmaf(static_cast(b_in.scratch_.separate[i]), s_val,\n static_cast(d_in.scratch_.separate[i])));\n } else {\n d_in.scratch_.separate[i] =\n static_cast(static_cast(d_in.scratch_.separate[i]) +\n static_cast(b_in.scratch_.separate[i]));\n }\n }\n *reinterpret_cast(d_ptr) = d_in.scratch_.aligned;\n }\n\n seg_start = seg_end;\n }\n }\n}\n\n// Single kernel that sets up all GEMM parameters.\n// Rationale: cuBLASLt grouped matmul API needs flat arrays of pointers and per-matrix dimensions,\n// but NVTEGroupedTensor stores a single contiguous buffer + optional per-tensor offsets/shapes.\n// We bridge the mismatch on GPU by computing per-group pointers and storage dims in one kernel.\n__global__ void setup_grouped_gemm_kernel(\n // Output arrays\n void **A_ptrs, void **B_ptrs, void **C_ptrs, void **D_ptrs, int *a_rows, int *a_cols,\n int *b_rows, int *b_cols, int *d_rows, int *d_cols, float **alpha_ptrs, float **beta_ptrs,\n void **a_scale_inv_ptrs, void **b_scale_inv_ptrs,\n // Inputs\n char *a_base, char *b_base, char *c_base, char *d_base, TensorShapeInfo A_meta,\n TensorShapeInfo B_meta, TensorShapeInfo C_meta, TensorShapeInfo D_meta, size_t a_elem_size,\n size_t b_elem_size, size_t c_elem_size, size_t d_elem_size, float *alpha_ptr, float *beta_ptr,\n // Scale inputs: for tensor scaling, pass float* and set mxfp8_base to nullptr\n // For MXFP8, pass nullptr for tensor_scale and set mxfp8_base\n float *a_scale_base, float *b_scale_base, bool a_rowwise, bool b_rowwise,\n NVTEScalingMode scaling_mode, size_t num_tensors,\n MultiTensorGroupGemmInputArgs a_multi_tensor_args,\n MultiTensorGroupGemmOutputArgs c_multi_tensor_args,\n MultiTensorGroupGemmOutputArgs d_multi_tensor_args) {\n size_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= num_tensors) return;\n\n // Get dimensions for this tensor (from array or uniform value)\n const bool has_a_multi_tensor = (a_base == nullptr);\n const bool has_c_multi_tensor = (c_base == nullptr);\n const bool has_d_multi_tensor = (d_base == nullptr);\n int64_t a_first = 0;\n int64_t a_last = 0;\n if (has_a_multi_tensor) {\n a_first = static_cast(a_multi_tensor_args.cols[idx]);\n a_last = static_cast(a_multi_tensor_args.rows[idx]);\n } else {\n a_first = A_meta.first_dims ? A_meta.first_dims[idx] : A_meta.uniform_first;\n a_last = A_meta.last_dims ? A_meta.last_dims[idx] : A_meta.uniform_last;\n }\n int64_t b_first = B_meta.first_dims ? B_meta.first_dims[idx] : B_meta.uniform_first;\n int64_t b_last = B_meta.last_dims ? B_meta.last_dims[idx] : B_meta.uniform_last;\n int64_t d_first = D_meta.first_dims ? D_meta.first_dims[idx] : D_meta.uniform_first;\n int64_t d_last = D_meta.last_dims ? D_meta.last_dims[idx] : D_meta.uniform_last;\n\n // Compute offsets (from explicit array, cumulative from per-tensor dims, or uniform)\n int64_t a_offset = has_a_multi_tensor ? 0 : compute_grouped_tensor_offset(A_meta, idx);\n int64_t b_offset = compute_grouped_tensor_offset(B_meta, idx);\n int64_t c_offset = compute_grouped_tensor_offset(C_meta, idx);\n int64_t d_offset = compute_grouped_tensor_offset(D_meta, idx);\n\n // Compute data pointers\n A_ptrs[idx] =\n has_a_multi_tensor ? a_multi_tensor_args.data_ptrs[idx] : (a_base + a_offset * a_elem_size);\n B_ptrs[idx] = b_base + b_offset * b_elem_size;\n C_ptrs[idx] =\n has_c_multi_tensor ? c_multi_tensor_args.data_ptrs[idx] : (c_base + c_offset * c_elem_size);\n D_ptrs[idx] =\n has_d_multi_tensor ? d_multi_tensor_args.data_ptrs[idx] : (d_base + d_offset * d_elem_size);\n\n // Compute storage dimensions for cuBLAS matrix layouts.\n // For INPUTS (A, B): Row-wise storage is seen as transposed column-major by cuBLAS,\n // so rows=last, cols=first. For columnwise, dims are already swapped.\n a_rows[idx] = static_cast(a_last);\n a_cols[idx] = static_cast(a_first);\n b_rows[idx] = static_cast(b_last);\n b_cols[idx] = static_cast(b_first);\n if (has_d_multi_tensor) {\n d_rows[idx] = d_multi_tensor_args.rows[idx];\n d_cols[idx] = d_multi_tensor_args.cols[idx];\n } else {\n d_rows[idx] = static_cast(d_last);\n d_cols[idx] = static_cast(d_first);\n }\n\n // Fill alpha/beta pointers (per-matrix)\n alpha_ptrs[idx] = alpha_ptr + idx;\n beta_ptrs[idx] = beta_ptr + idx;\n\n // Fill scale pointers (per-matrix).\n // The interpretation of the scale buffers depends on the shared scaling recipe:\n // otherwise : one float per tensor, indexed by tensor index\n if (a_scale_base) {\n if (scaling_mode == NVTE_MXFP8_1D_SCALING) {\n const int64_t a_scale_offset =\n compute_grouped_tensor_mxfp8_scale_inv_offset(A_meta, idx, a_rowwise);\n a_scale_inv_ptrs[idx] = reinterpret_cast(\n static_cast(static_cast(a_scale_base)) + a_scale_offset);\n } else {\n a_scale_inv_ptrs[idx] = static_cast(a_scale_base) + idx;\n }\n } else {\n a_scale_inv_ptrs[idx] = a_multi_tensor_args.scale_inv_ptrs[idx];\n }\n if (b_scale_base) {\n if (scaling_mode == NVTE_MXFP8_1D_SCALING) {\n const int64_t b_scale_offset =\n compute_grouped_tensor_mxfp8_scale_inv_offset(B_meta, idx, b_rowwise);\n b_scale_inv_ptrs[idx] = reinterpret_cast(\n static_cast(static_cast(b_scale_base)) + b_scale_offset);\n } else {\n b_scale_inv_ptrs[idx] = static_cast(b_scale_base) + idx;\n }\n }\n}\n\n// Launch the setup kernel to populate workspace arrays\ninline void launch_grouped_gemm_setup(\n const GroupedGemmSetupWorkspace &ws, const GroupedOperandSelection &A_sel,\n const GroupedOperandSelection &B_sel, const transformer_engine::GroupedTensor *C,\n const transformer_engine::GroupedTensor *D, const transformer_engine::Tensor *alpha_tensor,\n const transformer_engine::Tensor *beta_tensor, size_t num_tensors, hipStream_t stream,\n const MultiTensorGroupGemmInputArgs &a_multi_tensor_args, const NVTETensor *C_list,\n const NVTETensor *D_list, char *a_base, transformer_engine::DType c_dtype,\n transformer_engine::DType d_dtype) {\n // Use shape info from selection (already accounts for columnwise dimension swap)\n TensorShapeInfo A_meta = A_sel.shape;\n TensorShapeInfo B_meta = B_sel.shape;\n TensorShapeInfo C_meta{};\n TensorShapeInfo D_meta{};\n\n const bool has_d_multi_tensor = (D_list != nullptr);\n const bool has_c_multi_tensor = (C_list != nullptr) || has_d_multi_tensor;\n MultiTensorGroupGemmOutputArgs c_multi_tensor_args{};\n MultiTensorGroupGemmOutputArgs d_multi_tensor_args{};\n if (has_d_multi_tensor) {\n d_multi_tensor_args =\n build_grouped_gemm_multi_out_args(D_list, num_tensors, num_tensors, d_dtype, \"D\");\n }\n if (C_list != nullptr) {\n c_multi_tensor_args =\n build_grouped_gemm_multi_out_args(C_list, num_tensors, num_tensors, d_dtype, \"C\");\n } else if (has_d_multi_tensor) {\n c_multi_tensor_args = d_multi_tensor_args;\n }\n\n char *c_base = nullptr;\n char *d_base = nullptr;\n\n if (!has_c_multi_tensor) {\n NVTE_CHECK(C != nullptr && D != nullptr,\n \"Grouped GEMM: C/D grouped tensors are required when no C list is provided\");\n C_meta = TensorShapeInfo::create_shape_info_for_C(C, D);\n c_base = static_cast(C->data.dptr);\n }\n if (!has_d_multi_tensor) {\n NVTE_CHECK(D != nullptr,\n \"Grouped GEMM: D grouped tensor is required when no D list is provided\");\n D_meta = TensorShapeInfo::from_tensor(D);\n d_base = static_cast(D->data.dptr);\n }\n\n const size_t a_elem_size = transformer_engine::typeToSize(A_sel.dtype);\n const size_t b_elem_size = transformer_engine::typeToSize(B_sel.dtype);\n const size_t c_elem_size = transformer_engine::typeToSize(c_dtype);\n const size_t d_elem_size = transformer_engine::typeToSize(d_dtype);\n\n const int threads_per_block = 256;\n const int num_blocks = (num_tensors + threads_per_block - 1) / threads_per_block;\n\n // A and B share the same scaling recipe (validated in validate_grouped_gemm_inputs).\n // Pass scale buffers as void* and let the kernel interpret them via scaling_mode.\n\n // Scale rowwise flag for MXFP8/NVFP4: to calculate scale_inv padding based offsets\n // within kernel. Ignored for tensor scaling.\n const bool a_rowwise = A_sel.rowwise;\n const bool b_rowwise = B_sel.rowwise;\n setup_grouped_gemm_kernel<<>>(\n ws.A_ptrs, ws.B_ptrs, ws.C_ptrs, ws.D_ptrs, ws.a_rows, ws.a_cols, ws.b_rows, ws.b_cols,\n ws.d_rows, ws.d_cols, ws.alpha_ptrs, ws.beta_ptrs, ws.a_scale_inv_ptrs, ws.b_scale_inv_ptrs,\n A_sel.dptr, B_sel.dptr, c_base, d_base, A_meta, B_meta, C_meta, D_meta, a_elem_size,\n b_elem_size, c_elem_size, d_elem_size, static_cast(alpha_tensor->data.dptr),\n static_cast(beta_tensor->data.dptr), reinterpret_cast(A_sel.scale_inv),\n reinterpret_cast(B_sel.scale_inv), a_rowwise, b_rowwise, A_sel.scaling_mode,\n num_tensors, a_multi_tensor_args, c_multi_tensor_args, d_multi_tensor_args);\n\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n} // namespace\n\nsize_t nvte_get_grouped_gemm_setup_workspace_size(size_t num_tensors) {\n NVTE_API_CALL(nvte_get_grouped_gemm_setup_workspace_size);\n return grouped_gemm_setup_workspace_size(num_tensors);\n}\n\nvoid nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedTensor B, int transb,\n const NVTEGroupedTensor C, NVTEGroupedTensor D, const NVTETensor alpha,\n const NVTETensor beta, NVTETensor workspace_setup,\n NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_grouped_gemm);\n using namespace transformer_engine;\n\n // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.3+\n check_grouped_gemm_requirements(\"nvte_grouped_gemm\");\n\n // Convert to internal types\n const GroupedTensor *inputA = convertNVTEGroupedTensorCheck(A);\n const GroupedTensor *inputB = convertNVTEGroupedTensorCheck(B);\n const GroupedTensor *inputC_raw = convertNVTEGroupedTensor(C); // Can be NULL\n GroupedTensor *outputD = convertNVTEGroupedTensorCheck(D);\n const Tensor *alpha_tensor = convertNVTETensorCheck(alpha);\n const Tensor *beta_tensor = convertNVTETensorCheck(beta);\n Tensor *wspace_setup = convertNVTETensor(workspace_setup);\n Tensor *wspace_cublas = convertNVTETensor(workspace_cublas);\n\n // Parse config (if provided)\n GroupedMatmulConfig config_ = parse_grouped_gemm_config(config);\n\n // Validate inputs and outputs.\n const size_t num_tensors = validate_grouped_gemm_inputs(inputA->num_tensors, {inputA, inputB},\n alpha_tensor, beta_tensor);\n validate_grouped_gemm_outputs(num_tensors, {inputC_raw, outputD});\n\n // If C is NULL, use D as C (valid when beta=0, cuBLAS won't read C data)\n const GroupedTensor *inputC = (inputC_raw != nullptr) ? inputC_raw : outputD;\n // num_tensors validated above.\n // Select operand storage (row-wise vs column-wise) and adjust transpose flags to\n // mirror the non-grouped GEMM logic for FP8 layout constraints.\n auto A_sel = select_grouped_operand(inputA, static_cast(transa), /*is_A=*/true);\n auto B_sel = select_grouped_operand(inputB, static_cast(transb), /*is_A=*/false);\n\n // Workspaces: setup (pointer arrays) and cuBLAS\n auto workspace = setup_grouped_gemm_workspace(wspace_setup, wspace_cublas, num_tensors);\n\n MultiTensorGroupGemmInputArgs a_multi_tensor_args{};\n launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor,\n beta_tensor, num_tensors, stream, a_multi_tensor_args,\n /*C_list=*/nullptr, /*D_list=*/nullptr, A_sel.dptr, inputC->dtype(),\n outputD->dtype());\n\n // Compute average dimensions for heuristics\n // K dimension: if transa, K is A's first dim; if not, K is A's last dim\n // Use original inputA and transa for heuristics (not modified A_sel.trans)\n int64_t avg_m_val = config_.avg_m.value_or(compute_avg_first_dim(outputD));\n int64_t avg_n_val = config_.avg_n.value_or(compute_avg_last_dim(outputD));\n int64_t avg_k_val =\n config_.avg_k.value_or(transa ? compute_avg_first_dim(inputA) : compute_avg_last_dim(inputA));\n const bool use_fp8 = is_fp8_dtype(A_sel.dtype) || is_fp8_dtype(B_sel.dtype);\n execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, outputD->dtype(), num_tensors,\n config_.use_split_accumulator, use_fp8, avg_m_val, avg_n_val, avg_k_val,\n workspace.cublas_workspace_ptr, stream, config_.sm_count);\n}\n\nvoid nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors,\n int transa, const NVTEGroupedTensor B, int transb,\n const NVTEGroupedTensor C, NVTEGroupedTensor D,\n const NVTETensor alpha, const NVTETensor beta,\n NVTETensor workspace_setup, NVTETensor workspace_cublas,\n NVTEGroupedMatmulConfig config, hipStream_t stream) {\n NVTE_API_CALL(nvte_grouped_gemm_with_discrete_inputA);\n using namespace transformer_engine;\n\n // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.3+\n check_grouped_gemm_requirements(\"nvte_grouped_gemm_with_discrete_inputA\");\n\n NVTE_CHECK(A_list != nullptr, \"Grouped GEMM: A_list is null.\");\n NVTE_CHECK(num_a_tensors > 0, \"Grouped GEMM: num_a_tensors must be > 0.\");\n\n const GroupedTensor *inputB = convertNVTEGroupedTensorCheck(B);\n const GroupedTensor *inputC_raw = convertNVTEGroupedTensor(C); // Can be NULL\n GroupedTensor *outputD = convertNVTEGroupedTensorCheck(D);\n const Tensor *alpha_tensor = convertNVTETensorCheck(alpha);\n const Tensor *beta_tensor = convertNVTETensorCheck(beta);\n Tensor *wspace_setup = convertNVTETensor(workspace_setup);\n Tensor *wspace_cublas = convertNVTETensor(workspace_cublas);\n\n // Parse config (if provided)\n GroupedMatmulConfig config_ = parse_grouped_gemm_config(config);\n\n // Validate inputs and outputs.\n const size_t num_tensors =\n validate_grouped_gemm_inputs(num_a_tensors, {inputB}, alpha_tensor, beta_tensor);\n\n validate_grouped_gemm_outputs(num_tensors, {inputC_raw, outputD});\n\n // If C is NULL, use D as C (valid when beta=0, cuBLAS won't read C data)\n const GroupedTensor *inputC = (inputC_raw != nullptr) ? inputC_raw : outputD;\n\n // Validate A list and selection\n auto A_list_info =\n validate_grouped_gemm_multi_inputA_list(A_list, num_a_tensors, num_tensors, \"A\");\n auto is_fp8_or_16bit = [](transformer_engine::DType dtype) {\n return dtype == transformer_engine::DType::kFloat8E4M3 ||\n dtype == transformer_engine::DType::kFloat8E5M2 ||\n dtype == transformer_engine::DType::kBFloat16 ||\n dtype == transformer_engine::DType::kFloat16;\n };\n NVTE_CHECK(is_fp8_or_16bit(A_list_info.all_row ? A_list_info.row_dtype : A_list_info.col_dtype),\n \"Grouped GEMM: A_list tensors must be FP8, BF16, or FP16.\");\n\n // Cross-operand consistency (mirrors validate_grouped_gemm_inputs).\n const DType a_rep_dtype = A_list_info.all_row ? A_list_info.row_dtype : A_list_info.col_dtype;\n NVTE_CHECK(is_fp8_dtype(a_rep_dtype) == is_fp8_dtype(inputB->dtype()),\n \"Grouped GEMM: A and B must both be FP8 or both be non-FP8.\");\n NVTE_CHECK(transformer_engine::is_mxfp_scaling(A_list_info.scaling_mode) ==\n transformer_engine::is_mxfp_scaling(inputB->scaling_mode),\n \"Grouped GEMM: A and B must both use MXFP8 scaling or both use tensor scaling.\");\n if (transformer_engine::is_mxfp_scaling(A_list_info.scaling_mode)) {\n NVTE_CHECK(A_list_info.with_gemm_swizzled_scales,\n \"MXFP8 grouped GEMM: A scales must be swizzled for GEMM.\");\n NVTE_CHECK(inputB->with_gemm_swizzled_scales,\n \"MXFP8 grouped GEMM: B scales must be swizzled for GEMM.\");\n }\n\n // Select operand storage for B (row-wise vs column-wise)\n auto B_sel = select_grouped_operand(inputB, static_cast(transb), /*is_A=*/false);\n\n GroupedOperandSelection A_sel{};\n A_sel.scaling_mode = A_list_info.scaling_mode;\n A_sel.with_gemm_swizzled_scales = A_list_info.with_gemm_swizzled_scales;\n A_sel.trans = static_cast(transa);\n\n const DType rep_dtype = A_list_info.all_row ? A_list_info.row_dtype : A_list_info.col_dtype;\n const bool is_fp8 = is_fp8_dtype(rep_dtype);\n const bool non_tn_fp8_ok = nvte_is_non_tn_fp8_gemm_supported();\n const bool mxfp8 = transformer_engine::is_mxfp_scaling(A_list_info.scaling_mode);\n\n int64_t avg_first_dim = 0;\n int64_t avg_last_dim = 0;\n MultiTensorGroupGemmInputArgs a_multi_tensor_args{};\n\n const auto choice =\n choose_grouped_operand_storage(static_cast(transa), /*is_A=*/true, mxfp8, is_fp8,\n non_tn_fp8_ok, A_list_info.all_row, A_list_info.all_col, \"A\");\n A_sel.trans = choice.trans;\n A_sel.rowwise = choice.use_rowwise;\n if (choice.use_rowwise) {\n NVTE_CHECK(A_list_info.all_row, \"Grouped GEMM: A_list is missing row-wise data\");\n A_sel.dtype = A_list_info.row_dtype;\n a_multi_tensor_args = build_grouped_gemm_multi_inputA_args(\n A_list, num_a_tensors, /*use_rowwise=*/true, is_fp8, &avg_first_dim, &avg_last_dim, \"A\");\n } else {\n NVTE_CHECK(A_list_info.all_col, \"Grouped GEMM: A_list is missing column-wise data\");\n A_sel.dtype = A_list_info.col_dtype;\n a_multi_tensor_args = build_grouped_gemm_multi_inputA_args(\n A_list, num_a_tensors, /*use_rowwise=*/false, is_fp8, &avg_first_dim, &avg_last_dim, \"A\");\n }\n\n // For discrete A_list, scale pointers are per-tensor; use multi-tensor args.\n // Base pointer is unused when providing per-tensor pointers.\n A_sel.scale_inv = nullptr;\n A_sel.dptr = nullptr;\n\n // Workspaces: setup (pointer arrays) and cuBLAS\n auto workspace = setup_grouped_gemm_workspace(wspace_setup, wspace_cublas, num_tensors);\n\n launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor,\n beta_tensor, num_tensors, stream, a_multi_tensor_args,\n /*C_list=*/nullptr, /*D_list=*/nullptr, nullptr, inputC->dtype(),\n outputD->dtype());\n\n // Compute average dimensions for heuristics\n int64_t avg_m_val = config_.avg_m.value_or(compute_avg_first_dim(outputD));\n int64_t avg_n_val =\n config_.avg_n.value_or(transb ? compute_avg_first_dim(inputB) : compute_avg_last_dim(inputB));\n int64_t avg_k_val =\n config_.avg_k.value_or(static_cast(transa) ? avg_last_dim : avg_first_dim);\n const bool use_fp8 = is_fp8_dtype(A_sel.dtype) || is_fp8_dtype(B_sel.dtype);\n execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, outputD->dtype(), num_tensors,\n config_.use_split_accumulator, use_fp8, avg_m_val, avg_n_val, avg_k_val,\n workspace.cublas_workspace_ptr, stream, config_.sm_count);\n}\n\nvoid nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa,\n const NVTEGroupedTensor B, int transb,\n const NVTETensor *C_list, size_t num_c_tensors,\n NVTETensor *D_list, size_t num_d_tensors,\n const NVTETensor alpha, const NVTETensor beta,\n NVTETensor workspace_setup, NVTETensor workspace_cublas,\n NVTEGroupedMatmulConfig config, hipStream_t stream) {\n NVTE_API_CALL(nvte_grouped_gemm_with_discrete_out);\n using namespace transformer_engine;\n\n // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.3+\n check_grouped_gemm_requirements(\"nvte_grouped_gemm_with_discrete_out\");\n\n NVTE_CHECK(D_list != nullptr, \"Grouped GEMM: D_list is null.\");\n NVTE_CHECK(num_d_tensors > 0, \"Grouped GEMM: num_d_tensors must be > 0.\");\n if (num_c_tensors > 0) {\n NVTE_CHECK(C_list != nullptr, \"Grouped GEMM: C_list is null but num_c_tensors > 0.\");\n }\n\n const GroupedTensor *inputA = convertNVTEGroupedTensorCheck(A);\n const GroupedTensor *inputB = convertNVTEGroupedTensorCheck(B);\n const Tensor *alpha_tensor = convertNVTETensorCheck(alpha);\n const Tensor *beta_tensor = convertNVTETensorCheck(beta);\n Tensor *wspace_setup = convertNVTETensor(workspace_setup);\n Tensor *wspace_cublas = convertNVTETensor(workspace_cublas);\n\n const Tensor *d0 = convertNVTETensorCheck(D_list[0]);\n const DType d_dtype = d0->dtype();\n\n const size_t num_tensors = validate_grouped_gemm_inputs(inputA->num_tensors, {inputA, inputB},\n alpha_tensor, beta_tensor);\n NVTE_CHECK(num_d_tensors == num_tensors, \"Grouped GEMM: D_list must have num_tensors (\",\n num_tensors, \") entries, got \", num_d_tensors);\n if (num_c_tensors > 0) {\n NVTE_CHECK(num_c_tensors == num_tensors, \"Grouped GEMM: C_list must have num_tensors (\",\n num_tensors, \") entries, got \", num_c_tensors);\n }\n auto is_output_dtype = [](transformer_engine::DType dtype) {\n return dtype == transformer_engine::DType::kBFloat16 ||\n dtype == transformer_engine::DType::kFloat16 ||\n dtype == transformer_engine::DType::kFloat32;\n };\n NVTE_CHECK(is_output_dtype(d_dtype), \"Grouped GEMM: D must be BF16, FP16, or FP32.\");\n\n // Parse config (if provided)\n GroupedMatmulConfig config_ = parse_grouped_gemm_config(config);\n\n // Select operand storage (row-wise vs column-wise) and adjust transpose flags to\n // mirror the non-grouped GEMM logic for FP8 layout constraints.\n auto A_sel = select_grouped_operand(inputA, static_cast(transa), /*is_A=*/true);\n auto B_sel = select_grouped_operand(inputB, static_cast(transb), /*is_A=*/false);\n // Workspaces: setup (pointer arrays) and cuBLAS\n auto workspace = setup_grouped_gemm_workspace(wspace_setup, wspace_cublas, num_tensors);\n\n MultiTensorGroupGemmInputArgs a_multi_tensor_args{};\n launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, /*C=*/nullptr, /*D=*/nullptr,\n alpha_tensor, beta_tensor, num_tensors, stream, a_multi_tensor_args,\n C_list, D_list, A_sel.dptr, d_dtype, d_dtype);\n\n // Compute average dimensions for heuristics\n int64_t avg_m_val =\n config_.avg_m.value_or(transa ? compute_avg_last_dim(inputA) : compute_avg_first_dim(inputA));\n int64_t avg_n_val =\n config_.avg_n.value_or(transb ? compute_avg_first_dim(inputB) : compute_avg_last_dim(inputB));\n int64_t avg_k_val =\n config_.avg_k.value_or(transa ? compute_avg_first_dim(inputA) : compute_avg_last_dim(inputA));\n const bool use_fp8 = is_fp8_dtype(A_sel.dtype) || is_fp8_dtype(B_sel.dtype);\n execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, d_dtype, num_tensors,\n config_.use_split_accumulator, use_fp8, avg_m_val, avg_n_val, avg_k_val,\n workspace.cublas_workspace_ptr, stream, config_.sm_count);\n}\n\nnamespace {\n\nvoid launch_grouped_bias_add(const transformer_engine::GroupedTensor *outputD,\n const transformer_engine::GroupedTensor *bias_tensor,\n const float *scale_ptr, bool use_scale, hipStream_t stream) {\n using namespace transformer_engine;\n\n const char *api_name = use_scale ? \"Grouped scaled bias add\" : \"Grouped bias add\";\n\n NVTE_CHECK(outputD->num_tensors >= 1, api_name, \": number of tensors must be at least 1\");\n NVTE_CHECK(outputD->num_tensors == bias_tensor->num_tensors, api_name,\n \": output and bias must have the same number of tensors\");\n NVTE_CHECK(outputD->has_data(), api_name, \": output is missing row-wise data\");\n NVTE_CHECK(bias_tensor->has_data(), api_name, \": bias is missing row-wise data\");\n NVTE_CHECK(outputD->dtype() == bias_tensor->dtype(), api_name,\n \": output and bias must have matching dtypes\");\n NVTE_CHECK(bias_tensor->all_same_first_dim(), api_name,\n \": bias must have uniform first dim (expected 1)\");\n NVTE_CHECK(bias_tensor->get_common_first_dim() == 1, api_name, \": bias first dim must be 1\");\n NVTE_CHECK(outputD->all_same_last_dim() && bias_tensor->all_same_last_dim(), api_name,\n \": requires uniform last dim for output and bias\");\n NVTE_CHECK(outputD->get_common_last_dim() == bias_tensor->get_common_last_dim(), api_name,\n \": output and bias last dims must match\");\n\n const TensorShapeInfo d_meta = TensorShapeInfo::from_tensor(outputD);\n\n const DType dtype = outputD->dtype();\n constexpr int kThreads = 128;\n\n const int num_tensors = static_cast(outputD->num_tensors);\n NVTE_CHECK(num_tensors <= kMaxGroups, api_name, \" supports at most \", kMaxGroups,\n \" tensors, got \", num_tensors);\n const int total_rows = static_cast(outputD->logical_shape.data[0]);\n const int n = static_cast(outputD->get_common_last_dim());\n\n const size_t elem_size = typeToSize(dtype);\n const int kVec = (elem_size <= 2) ? 8 : 4;\n NVTE_CHECK(n % kVec == 0, api_name, \": requires last dim divisible by \", kVec);\n\n constexpr int kRowsPerBlock = 16;\n constexpr int kBlocksPerSM = 32;\n\n const int num_sms = transformer_engine::cuda::sm_count();\n\n const int block_cols = kThreads * kVec;\n const int col_blocks = (n + block_cols - 1) / block_cols;\n const int max_row_chunks = (total_rows + kRowsPerBlock - 1) / kRowsPerBlock;\n const int row_blocks = std::min(max_row_chunks, num_sms * kBlocksPerSM / col_blocks);\n const dim3 grid(std::max(1, row_blocks), col_blocks);\n const dim3 block(kThreads);\n\n auto launch = [&](auto use_scale_tag) {\n constexpr bool kUseScale = decltype(use_scale_tag)::value;\n if (elem_size <= 2) {\n constexpr int kV = 8;\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(dtype, T, {\n grouped_bias_add_kernel\n <<>>(static_cast(outputD->data.dptr),\n static_cast(bias_tensor->data.dptr),\n scale_ptr, d_meta, n, total_rows, num_tensors);\n });\n } else {\n constexpr int kV = 4;\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(dtype, T, {\n grouped_bias_add_kernel\n <<>>(static_cast(outputD->data.dptr),\n static_cast(bias_tensor->data.dptr),\n scale_ptr, d_meta, n, total_rows, num_tensors);\n });\n }\n };\n\n if (use_scale) {\n launch(std::true_type{});\n } else {\n launch(std::false_type{});\n }\n\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n} // namespace\n\nvoid nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_grouped_bias_add);\n using namespace transformer_engine;\n const GroupedTensor *outputD = convertNVTEGroupedTensorCheck(output);\n const GroupedTensor *bias_tensor = convertNVTEGroupedTensorCheck(bias);\n launch_grouped_bias_add(outputD, bias_tensor, nullptr, false, stream);\n}\n\nvoid nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias,\n const NVTETensor scale, hipStream_t stream) {\n NVTE_API_CALL(nvte_grouped_scaled_bias_add);\n using namespace transformer_engine;\n const GroupedTensor *outputD = convertNVTEGroupedTensorCheck(output);\n const GroupedTensor *bias_tensor = convertNVTEGroupedTensorCheck(bias);\n const Tensor *scale_tensor = convertNVTETensorCheck(scale);\n\n NVTE_CHECK(scale_tensor->data.dptr != nullptr,\n \"Grouped scaled bias add: scale tensor must not be null\");\n NVTE_CHECK(scale_tensor->dtype() == DType::kFloat32,\n \"Grouped scaled bias add: scale must be float32\");\n NVTE_CHECK(scale_tensor->data.shape.size() == 1,\n \"Grouped scaled bias add: scale must be 1D, got \", scale_tensor->data.shape.size(),\n \"D\");\n const size_t total_rows = static_cast(outputD->logical_shape.data[0]);\n NVTE_CHECK(scale_tensor->data.shape[0] == total_rows, \"Grouped scaled bias add: scale size (\",\n scale_tensor->data.shape[0], \") must equal total rows (\", total_rows, \")\");\n\n const float *scale_ptr = static_cast(scale_tensor->data.dptr);\n launch_grouped_bias_add(outputD, bias_tensor, scale_ptr, true, stream);\n}\n\n#else // CUBLAS_VERSION < CUBLAS_GROUPED_GEMM_VERSION\n\nvoid nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedTensor B, int transb,\n const NVTEGroupedTensor C, NVTEGroupedTensor D, const NVTETensor alpha,\n const NVTETensor beta, NVTETensor workspace_setup,\n NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config,\n hipStream_t stream) {\n NVTE_ERROR(\"nvte_grouped_gemm requires cuBLAS 13.3+, but compile-time cuBLAS version is \",\n CUBLAS_VERSION, \". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer.\");\n}\n\nvoid nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors,\n int transa, const NVTEGroupedTensor B, int transb,\n const NVTEGroupedTensor C, NVTEGroupedTensor D,\n const NVTETensor alpha, const NVTETensor beta,\n NVTETensor workspace_setup, NVTETensor workspace_cublas,\n NVTEGroupedMatmulConfig config, hipStream_t stream) {\n NVTE_ERROR(\n \"nvte_grouped_gemm_with_discrete_inputA requires cuBLAS 13.3+, but compile-time \"\n \"cuBLAS version is \",\n CUBLAS_VERSION, \". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer.\");\n}\n\nvoid nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa,\n const NVTEGroupedTensor B, int transb,\n const NVTETensor *C_list, size_t num_c_tensors,\n NVTETensor *D_list, size_t num_d_tensors,\n const NVTETensor alpha, const NVTETensor beta,\n NVTETensor workspace_setup, NVTETensor workspace_cublas,\n NVTEGroupedMatmulConfig config, hipStream_t stream) {\n NVTE_ERROR(\n \"nvte_grouped_gemm_with_discrete_out requires cuBLAS 13.3+, but compile-time \"\n \"cuBLAS version is \",\n CUBLAS_VERSION, \". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer.\");\n}\n\nvoid nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias,\n hipStream_t stream) {\n NVTE_ERROR(\"nvte_grouped_bias_add requires cuBLAS 13.3+, but compile-time cuBLAS version is \",\n CUBLAS_VERSION, \". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer.\");\n}\n\nvoid nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias,\n const NVTETensor scale, hipStream_t stream) {\n NVTE_ERROR(\n \"nvte_grouped_scaled_bias_add requires cuBLAS 13.3+, but compile-time cuBLAS version \"\n \"is \",\n CUBLAS_VERSION, \". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer.\");\n}\n\nsize_t nvte_get_grouped_gemm_setup_workspace_size(size_t num_tensors) {\n NVTE_ERROR(\n \"nvte_get_grouped_gemm_setup_workspace_size requires cuBLAS 13.3+, but compile-time cuBLAS \"\n \"version is \",\n CUBLAS_VERSION, \". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer.\");\n return 0;\n}\n\n#endif // CUBLAS_VERSION >= CUBLAS_GROUPED_GEMM_VERSION\n\nnamespace {\n\n__global__ void convert_int32_to_int64_kernel(const int32_t *src, int64_t *dst, size_t n) {\n size_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < n) dst[idx] = static_cast(src[idx]);\n}\n\n// Like convert_int32_to_int64_kernel but scales each element by multiplier.\n// Used to convert per-expert slice counts to per-expert row counts for multi-dim tensors.\n__global__ void convert_int32_to_int64_with_multiplier_kernel(const int32_t *src, int64_t *dst,\n size_t n, int64_t multiplier) {\n size_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < n) dst[idx] = static_cast(src[idx]) * multiplier;\n}\n\n// Computes exclusive prefix sums: offsets[0]=0, offsets[i]=sum(first_dims[0..i-1]*last_dim).\n// Produces n_groups+1 values. Single-threaded sequential scan; n_groups is typically small.\n__global__ void compute_grouped_tensor_offsets_kernel(const int64_t *first_dims, int64_t *offsets,\n size_t n_groups, int64_t last_dim) {\n offsets[0] = 0;\n for (size_t i = 0; i < n_groups; i++) {\n offsets[i + 1] = offsets[i] + first_dims[i] * last_dim;\n }\n}\n\n} // namespace\n\nvoid nvte_convert_int32_to_int64(const int32_t *src, int64_t *dst, size_t n, hipStream_t stream) {\n NVTE_API_CALL(nvte_convert_int32_to_int64);\n if (n == 0) return;\n const int threads = 256;\n const int blocks = static_cast((n + threads - 1) / threads);\n convert_int32_to_int64_kernel<<>>(src, dst, n);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid nvte_convert_int32_to_int64_with_multiplier(const int32_t *src, int64_t *dst, size_t n,\n int64_t multiplier, hipStream_t stream) {\n NVTE_API_CALL(nvte_convert_int32_to_int64_with_multiplier);\n if (n == 0) return;\n const int threads = 256;\n const int blocks = static_cast((n + threads - 1) / threads);\n convert_int32_to_int64_with_multiplier_kernel<<>>(src, dst, n,\n multiplier);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid nvte_compute_grouped_tensor_offsets(const int64_t *first_dims, int64_t *offsets,\n size_t n_groups, int64_t last_dim, hipStream_t stream) {\n NVTE_API_CALL(nvte_compute_grouped_tensor_offsets);\n // Always write at least offsets[0]=0 (needed even for n_groups==0).\n compute_grouped_tensor_offsets_kernel<<<1, 1, 0, stream>>>(first_dims, offsets, n_groups,\n last_dim);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/cuda_runtime.h\"\n#include \"../util/handle_manager.h\"\n#include \"../util/logging.h\"\n#include \"../util/multi_stream.h\"\n#include \"./config.h\"\n#include \"./cutlass_grouped_gemm.cuh\"\n\nnamespace {\n\n/* Use CUDA const memory to store scalar 1 and 0 for cublas usage\n*/\n__device__ __constant__ float one_device;\n__device__ __constant__ float zero_device;\n\ninline float *GetScalarOne() {\n static std::once_flag init_flag;\n std::call_once(init_flag, []() {\n float one = 1.0f;\n NVTE_CHECK_CUDA(cudaMemcpyToSymbol(one_device, &one, sizeof(float)));\n });\n // return address by cudaGetSymbolAddress\n float *dev_ptr;\n NVTE_CHECK_CUDA(cudaGetSymbolAddress(reinterpret_cast(&dev_ptr), one_device));\n return dev_ptr;\n}\n\ninline float *GetScalarZero() {\n static std::once_flag init_flag;\n std::call_once(init_flag, []() {\n float zero = 0.0f;\n NVTE_CHECK_CUDA(cudaMemcpyToSymbol(zero_device, &zero, sizeof(float)));\n });\n // return address by cudaGetSymbolAddress\n float *dev_ptr;\n NVTE_CHECK_CUDA(cudaGetSymbolAddress(reinterpret_cast(&dev_ptr), zero_device));\n return dev_ptr;\n}\n\n__global__ __launch_bounds__(1) void set_float_kernel(float *ptr, float val) { *ptr = val; }\n\nuint32_t _getAlignment(uintptr_t address) {\n // alignment are in bytes\n uint32_t alignment = 256;\n for (;; alignment /= 2) {\n if (address % alignment == 0) {\n return alignment;\n }\n }\n}\n\ninline void CreateCublasHandle(cublasLtHandle_t *handle) {\n NVTE_CHECK_CUBLAS(cublasLtCreate(handle));\n}\n\n/* Parameters for cuBLAS GEMM\n *\n * cuBLAS follows the BLAS convention of column-major ordering. This\n * is different than the row-major that is typically used in\n * Transformer Engine.\n *\n */\nstruct GemmParam {\n void *A = nullptr;\n void *B = nullptr;\n cublasOperation_t transA = CUBLAS_OP_N;\n cublasOperation_t transB = CUBLAS_OP_N;\n transformer_engine::DType Atype = transformer_engine::DType::kNumTypes;\n transformer_engine::DType Btype = transformer_engine::DType::kNumTypes;\n void *A_scale_inv = nullptr;\n void *B_scale_inv = nullptr;\n int lda = 0; // A column strides\n int ldb = 0; // B column strides\n};\n\n/* Populate parameters for cuBLAS GEMM\n *\n * cuBLAS follows the BLAS convention of column-major ordering. This\n * is different than the row-major that is typically used in\n * Transformer Engine.\n *\n */\nGemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cublasOperation_t transA,\n const transformer_engine::Tensor &B, const cublasOperation_t transB,\n int m, int n, int k) {\n using namespace transformer_engine;\n NVTE_CHECK(\n A.scaling_mode == B.scaling_mode ||\n (A.scaling_mode == NVTE_BLOCK_SCALING_1D && B.scaling_mode == NVTE_BLOCK_SCALING_2D) ||\n (A.scaling_mode == NVTE_BLOCK_SCALING_2D && B.scaling_mode == NVTE_BLOCK_SCALING_1D),\n \"Inputs A and B to GEMM need to have compatible scaling modes, but got A.scaling_mode = \" +\n to_string(A.scaling_mode) + \", B.scaling_mode = \" + to_string(B.scaling_mode));\n NVTE_CHECK(A.has_data() || A.has_columnwise_data(), \"Input A does not hold any data!\");\n NVTE_CHECK(B.has_data() || B.has_columnwise_data(), \"Input B does not hold any data!\");\n GemmParam ret;\n\n // Transpose mode with column-major ordering\n bool is_A_transposed = transA == CUBLAS_OP_T;\n bool is_B_transposed = transB == CUBLAS_OP_T;\n\n // Set conditions for MXFP8 and NVFP4 gemm execution.\n const auto nvfp4 = is_nvfp_scaling(A.scaling_mode) && is_nvfp_scaling(B.scaling_mode);\n const auto mxfp8 = !nvfp4 && is_mxfp_scaling(A.scaling_mode) && is_mxfp_scaling(B.scaling_mode);\n int is_nvte_non_tn_fp8_gemm_supported = 0; // needed only for per tensor scaling\n if (is_tensor_scaling(A.scaling_mode) || is_tensor_scaling(B.scaling_mode)) {\n is_nvte_non_tn_fp8_gemm_supported = nvte_is_non_tn_fp8_gemm_supported();\n }\n\n // Configure A matrix\n if (is_tensor_scaling(A.scaling_mode)) {\n // Unscaled or FP8 tensor scaling\n ret.A = A.data.dptr;\n ret.transA = transA;\n ret.Atype = A.data.dtype;\n ret.A_scale_inv = A.scale_inv.dptr;\n ret.lda = is_A_transposed ? k : m;\n if (!is_nvte_non_tn_fp8_gemm_supported && !is_A_transposed) {\n // Hopper only supports TN GEMMs for FP8. \"Column-wise data\" is transpose of data.\n if (A.has_columnwise_data() && is_fp8_dtype(A.columnwise_data.dtype)) {\n ret.A = A.columnwise_data.dptr;\n ret.transA = CUBLAS_OP_T;\n ret.Atype = A.columnwise_data.dtype;\n ret.A_scale_inv = A.columnwise_scale_inv.dptr;\n ret.lda = k;\n } else {\n NVTE_CHECK(!is_fp8_dtype(ret.Atype), \"Input A is missing column-wise usage\");\n }\n } else if (is_nvte_non_tn_fp8_gemm_supported && !A.has_data()) {\n // Blackwell supports any GEMM layout for FP8, so we can use column-wise/transposed\n // data with the mirrored transpose-flag if we don't have row-wise data.\n NVTE_CHECK(A.has_columnwise_data() && is_fp8_dtype(A.columnwise_data.dtype),\n \"Input A is missing column-wise usage\");\n ret.A = A.columnwise_data.dptr;\n ret.transA = is_A_transposed ? CUBLAS_OP_N : CUBLAS_OP_T;\n ret.Atype = A.columnwise_data.dtype;\n ret.A_scale_inv = A.columnwise_scale_inv.dptr;\n ret.lda = is_A_transposed ? m : k;\n }\n\n if (is_fp8_dtype(ret.Atype)) {\n // Requirements from https://docs.nvidia.com/cuda/cublas/#tensor-core-usage\n NVTE_CHECK(ret.lda % 16 == 0,\n \"Leading dimension requirement on A for FP8 GEMM. Caller must pad.\");\n }\n } else if (nvfp4) {\n // NVFP4 GEMM. Either the pure NVFP4 recipe or the FWD pass of the Hybrid NVFP4/MXFP8 recipe.\n\n if (is_A_transposed) {\n NVTE_CHECK(A.has_data(), \"Input A is missing row-wise usage\");\n } else {\n NVTE_CHECK(is_nvfp4_scaling(A.scaling_mode),\n \"Input A has unsupported combination of recipe and layout\");\n NVTE_CHECK(A.has_columnwise_data(), \"Input A is missing column-wise usage\");\n }\n ret.A = is_A_transposed ? A.data.dptr : A.columnwise_data.dptr;\n ret.transA = CUBLAS_OP_T; // NVFP4 gemm is only supported in TN layout.\n ret.Atype = is_A_transposed ? A.data.dtype : A.columnwise_data.dtype;\n ret.A_scale_inv = is_A_transposed ? A.scale_inv.dptr : A.columnwise_scale_inv.dptr;\n ret.lda = k;\n } else if (mxfp8) {\n // MXFP8 GEMM. Either for pure MXFP8 recipe or backward of Hybrid NVFP4 recipe.\n // Note: Row-wise and column-wise data are scaled along different\n // dimensions (with matrix interpreted in row-major order).\n\n if (is_A_transposed) {\n NVTE_CHECK(A.has_data(), \"Input A is missing row-wise usage\");\n } else {\n NVTE_CHECK(A.has_columnwise_data(), \"Input A is missing column-wise usage\");\n }\n ret.A = is_A_transposed ? A.data.dptr : A.columnwise_data.dptr;\n ret.transA = transA;\n ret.Atype = is_A_transposed ? A.data.dtype : A.columnwise_data.dtype;\n ret.A_scale_inv = is_A_transposed ? A.scale_inv.dptr : A.columnwise_scale_inv.dptr;\n ret.lda = is_A_transposed ? k : m;\n } else if (A.scaling_mode == NVTE_BLOCK_SCALING_1D || A.scaling_mode == NVTE_BLOCK_SCALING_2D) {\n // FP8 block scaling\n // Note: Hopper only supports TN GEMMs for FP8. \"Column-wise data\" is transpose of data.\n if (is_A_transposed) {\n NVTE_CHECK(A.has_data(), \"Input A is missing row-wise usage\");\n } else {\n NVTE_CHECK(A.has_columnwise_data(), \"Input A is missing column-wise usage\");\n }\n ret.A = is_A_transposed ? A.data.dptr : A.columnwise_data.dptr;\n ret.transA = CUBLAS_OP_T;\n ret.Atype = is_A_transposed ? A.data.dtype : A.columnwise_data.dtype;\n ret.A_scale_inv = is_A_transposed ? A.scale_inv.dptr : A.columnwise_scale_inv.dptr;\n ret.lda = k;\n\n // Requirements from https://docs.nvidia.com/cuda/cublas/#tensor-core-usage\n NVTE_CHECK((ret.lda % 16) == 0,\n \"Leading dimension requirement on NVTE_BLOCK_SCALING GEMM. Caller must pad.\");\n // Divisibility of 8 derived from FP8 (m * CTypeSize) % 16 == 0 requirement.\n // Smallest supported CType is 2 bytes in this scaling mode.\n NVTE_CHECK((m % 8) == 0,\n \"Outer dimension requirement on A for NVTE_BLOCK_SCALING GEMM. Caller must pad.\");\n } else {\n NVTE_ERROR(\"A has unsupported scaling mode\");\n }\n\n // Configure B matrix\n if (is_tensor_scaling(B.scaling_mode)) {\n // Unscaled or FP8 tensor scaling\n ret.B = B.data.dptr;\n ret.transB = transB;\n ret.Btype = B.data.dtype;\n ret.B_scale_inv = B.scale_inv.dptr;\n ret.ldb = is_B_transposed ? n : k;\n if (!is_nvte_non_tn_fp8_gemm_supported && is_B_transposed) {\n // Hopper only supports TN GEMMs for FP8. \"Column-wise data\" is transpose of data.\n if (B.has_columnwise_data() && is_fp8_dtype(B.columnwise_data.dtype)) {\n ret.B = B.columnwise_data.dptr;\n ret.transB = CUBLAS_OP_N;\n ret.Btype = B.columnwise_data.dtype;\n ret.B_scale_inv = B.columnwise_scale_inv.dptr;\n ret.ldb = k;\n } else {\n NVTE_CHECK(!is_fp8_dtype(ret.Btype), \"Input B is missing column-wise usage\");\n }\n } else if (is_nvte_non_tn_fp8_gemm_supported && !B.has_data()) {\n // Blackwell supports any GEMM layout for FP8, so we can use column-wise/transposed\n // data with the mirrored transpose-flag if we don't have row-wise data.\n NVTE_CHECK(B.has_columnwise_data() && is_fp8_dtype(B.columnwise_data.dtype),\n \"Input B is missing column-wise usage\");\n ret.B = B.columnwise_data.dptr;\n ret.transB = is_B_transposed ? CUBLAS_OP_N : CUBLAS_OP_T;\n ret.Btype = B.columnwise_data.dtype;\n ret.B_scale_inv = B.columnwise_scale_inv.dptr;\n ret.ldb = is_B_transposed ? k : n;\n }\n\n if (is_fp8_dtype(ret.Atype)) {\n // Requirements from https://docs.nvidia.com/cuda/cublas/#tensor-core-usage\n NVTE_CHECK(ret.ldb % 16 == 0,\n \"Leading dimension requirement on B for FP8 GEMM. Caller must pad.\");\n }\n } else if (nvfp4) {\n if (is_B_transposed) {\n NVTE_CHECK(is_nvfp4_scaling(B.scaling_mode),\n \"Input B has unsupported combination of recipe and layout\");\n NVTE_CHECK(B.has_columnwise_data(), \"Input B is missing column-wise usage\");\n } else {\n NVTE_CHECK(B.has_data(), \"Input B is missing row-wise usage\");\n }\n ret.B = is_B_transposed ? B.columnwise_data.dptr : B.data.dptr;\n ret.transB = CUBLAS_OP_N; // NVFP4 gemm is only supported in TN layout.\n ret.Btype = is_B_transposed ? B.columnwise_data.dtype : B.data.dtype;\n ret.B_scale_inv = is_B_transposed ? B.columnwise_scale_inv.dptr : B.scale_inv.dptr;\n ret.ldb = k;\n } else if (mxfp8) {\n if (is_B_transposed) {\n NVTE_CHECK(B.has_columnwise_data(), \"Input B is missing column-wise usage\");\n } else {\n NVTE_CHECK(B.has_data(), \"Input B is missing row-wise usage\");\n }\n ret.B = is_B_transposed ? B.columnwise_data.dptr : B.data.dptr;\n ret.transB = transB;\n ret.Btype = is_B_transposed ? B.columnwise_data.dtype : B.data.dtype;\n ret.B_scale_inv = is_B_transposed ? B.columnwise_scale_inv.dptr : B.scale_inv.dptr;\n ret.ldb = is_B_transposed ? n : k;\n } else if (B.scaling_mode == NVTE_BLOCK_SCALING_1D || B.scaling_mode == NVTE_BLOCK_SCALING_2D) {\n // FP8 block scaling\n // Note: Hopper only supports TN GEMMs for FP8. \"Column-wise data\" is transpose of data.\n if (is_B_transposed) {\n NVTE_CHECK(B.has_columnwise_data(), \"Input B is missing column-wise usage\");\n } else {\n NVTE_CHECK(B.has_data(), \"Input B is missing row-wise usage\");\n }\n ret.B = is_B_transposed ? B.columnwise_data.dptr : B.data.dptr;\n ret.transB = CUBLAS_OP_N;\n ret.Btype = is_B_transposed ? B.columnwise_data.dtype : B.data.dtype;\n ret.B_scale_inv = is_B_transposed ? B.columnwise_scale_inv.dptr : B.scale_inv.dptr;\n ret.ldb = k;\n\n // Requirements from\n // https://docs.nvidia.com/cuda/cublas/#tensor-core-usage\n NVTE_CHECK((ret.ldb % 16) == 0,\n \"B tensor stride requirement on NVTE_BLOCK_SCALING GEMM. Caller must pad.\");\n if (B.scaling_mode == NVTE_BLOCK_SCALING_1D) {\n // Observed this requirement only present for B tensor is 1D quantized.\n NVTE_CHECK((n % 8) == 0,\n \"Outer dimension requirement on B for NVTE_BLOCK_SCALING GEMM. Caller must pad.\");\n }\n } else {\n NVTE_ERROR(\"B has unsupported scaling mode\");\n }\n\n return ret;\n}\n\n} // namespace\n\nnamespace transformer_engine {\n\nusing cublasHandleManager = detail::HandleManager;\n\nvoid cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD,\n const Tensor *inputBias, Tensor *outputPreGelu, cublasOperation_t transa,\n cublasOperation_t transb, bool grad, void *workspace, size_t workspaceSize,\n const void *alpha, const void *beta, bool use_split_accumulator, int math_sm_count,\n int m_split, int n_split, bool gemm_producer, const Tensor *inputCounter,\n cudaStream_t stream) {\n NVTE_CHECK(!inputA->row_scaled_nvfp4 && !inputB->row_scaled_nvfp4,\n \"cuBLAS GEMM does not support row-scaled NVFP4 inputs.\");\n\n // Tensor dims in row-major order\n const int A0 = inputA->flat_first_dim();\n const int A1 = inputA->flat_last_dim();\n const int B0 = inputB->flat_first_dim();\n const int B1 = inputB->flat_last_dim();\n\n // GEMM dims in column-major order\n const int m = transa == CUBLAS_OP_T ? A0 : A1;\n const int n = transb == CUBLAS_OP_T ? B1 : B0;\n const int k = transa == CUBLAS_OP_T ? A1 : A0;\n NVTE_CHECK((transb == CUBLAS_OP_T ? B0 : B1) == k,\n \"GEMM inputs have incompatible dimensions (A is \", A0, \"x\", A1, \", B is \", B0, \"x\", B1,\n \")\");\n const int ldd = m;\n\n // Return immediately if GEMM is trivial\n if (m <= 0 || n <= 0) {\n return;\n }\n NVTE_CHECK(k > 0);\n\n const GemmParam param = CanonicalizeGemmInput(*inputA, transa, *inputB, transb, m, n, k);\n\n void *C = outputD->data.dptr;\n void *D = outputD->data.dptr;\n void *D_scale = outputD->scale.dptr;\n void *D_amax = outputD->amax.dptr;\n void *bias_ptr = inputBias->data.dptr;\n const bool bias = bias_ptr != nullptr;\n void *pre_gelu_out = outputPreGelu->data.dptr;\n void *counter = nullptr;\n if (inputCounter != nullptr) {\n counter = inputCounter->data.dptr;\n }\n const bool gelu = pre_gelu_out != nullptr;\n const bool use_fp8 = is_fp8_dtype(param.Atype) || is_fp8_dtype(param.Btype);\n const bool use_fp4 = is_fp4_dtype(param.Atype) || is_fp4_dtype(param.Btype);\n\n // Update scaling factors with NVFP4 tensor scales\n // TODO: Check whether scales are on CPU/GPU or add API to control.\n // Currently scales are assumed to be on CPU when amax is provided\n // and on GPU when not provided, but this is brittle.\n if (use_fp4 &&\n ((transa == CUBLAS_OP_T ? inputA->amax.dptr : inputA->columnwise_amax.dptr) != nullptr ||\n (transb == CUBLAS_OP_T ? inputB->columnwise_amax.dptr : inputB->amax.dptr) != nullptr)) {\n // Reserve some workspace for alpha scale\n NVTE_CHECK(workspaceSize >= 4,\n \"NVFP4 GEMM requires at least 4 byte workspace for alpha scale, but only has \",\n workspaceSize, \" bytes remaining.\");\n workspaceSize = (workspaceSize / 4) * 4 - 4; // Remove last 4 aligned bytes\n uint8_t *workspace_ptr = reinterpret_cast(workspace);\n float *new_alpha_ptr = reinterpret_cast(&workspace_ptr[workspaceSize]);\n\n // Update alpha scale on device\n // Note: Compute NVFP4 tensor scales based on amaxes and then\n // divide from alpha scale. This way we only need to apply NVFP4\n // tensor scales in matmul output, instead of in matmul inputs.\n float old_alpha = *reinterpret_cast(alpha); // Assumed to be on CPU\n TensorWrapper new_alpha_tensor(new_alpha_ptr, std::vector{1}, DType::kFloat32);\n bool a_rowwise_amax = transa == CUBLAS_OP_T;\n bool b_rowwise_amax = transb != CUBLAS_OP_T;\n nvte_nvfp4_compute_per_tensor_scale(inputA->nvte_tensor, a_rowwise_amax, inputB->nvte_tensor,\n b_rowwise_amax, old_alpha, new_alpha_tensor.data(), stream);\n alpha = new_alpha_ptr;\n\n // Make sure beta scale is on device\n float old_beta = *reinterpret_cast(beta); // Assumed to be on CPU\n if (old_beta == 0) {\n beta = GetScalarZero(); // Device constant memory\n } else if (old_beta == 1) {\n beta = GetScalarOne(); // Device constant memory\n } else {\n // Move beta to workspace\n NVTE_CHECK(workspaceSize >= 4,\n \"NVFP4 GEMM requires at least 4 byte workspace for beta scale, but only has \",\n workspaceSize, \" bytes remaining.\");\n workspaceSize = (workspaceSize / 4) * 4 - 4; // Remove last 4 aligned bytes\n float *new_beta_ptr = reinterpret_cast(&workspace_ptr[workspaceSize]);\n set_float_kernel<<<1, 1, 0, stream>>>(new_beta_ptr, old_beta);\n NVTE_CHECK_CUDA(cudaGetLastError());\n beta = new_beta_ptr;\n }\n }\n\n const cudaDataType_t A_type = get_cuda_dtype(param.Atype);\n const cudaDataType_t B_type = get_cuda_dtype(param.Btype);\n const cudaDataType_t D_type = get_cuda_dtype(outputD->data.dtype);\n const cudaDataType_t bias_type = get_cuda_dtype(inputBias->data.dtype);\n\n NVTE_CHECK(!is_fp8_dtype(param.Atype) || param.A_scale_inv != nullptr,\n \"FP8 input to GEMM requires inverse of scale!\");\n NVTE_CHECK(!is_fp8_dtype(param.Btype) || param.B_scale_inv != nullptr,\n \"FP8 input to GEMM requires inverse of scale!\");\n NVTE_CHECK(!is_fp4_dtype(param.Atype) || param.A_scale_inv != nullptr,\n \"FP4 input to GEMM requires inverse of scale!\");\n NVTE_CHECK(!is_fp4_dtype(param.Btype) || param.B_scale_inv != nullptr,\n \"FP4 input to GEMM requires inverse of scale!\");\n\n // check consistency of arguments:\n // if fp8 is desired, context cannot be null\n // fp8 + gelu fusion + fp8 aux is unavailable right now.\n if ((use_fp8 || use_fp4) && gelu) {\n NVTE_CHECK(!is_fp8_dtype(outputPreGelu->data.dtype),\n \"fp8 Aux output for gemm + gelu fusion not supported!\");\n }\n if (is_fp4_dtype(outputD->data.dtype)) {\n NVTE_ERROR(\"FP4 GEMM output is not supported!\");\n }\n if (use_fp4 && (D_type == CUDA_R_16F)) {\n NVTE_ERROR(\"FP4 GEMM does not support FP16 output!\");\n }\n\n cublasLtHandle_t handle = cublasHandleManager::Instance().GetHandle();\n\n cublasLtMatmulDesc_t operationDesc = nullptr;\n cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr, Ddesc = nullptr;\n cublasLtMatmulPreference_t preference = nullptr;\n int returnedResults = 0;\n cublasLtMatmulHeuristicResult_t heuristicResult = {};\n cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DEFAULT;\n\n int64_t ld_gelumat = (int64_t)ldd;\n\n // Use TF32 only for pure FP32 GEMM.\n cublasComputeType_t gemm_compute_type = CUBLAS_COMPUTE_32F;\n if (A_type == CUDA_R_32F && B_type == CUDA_R_32F && D_type == CUDA_R_32F) {\n gemm_compute_type = CUBLAS_COMPUTE_32F_FAST_TF32;\n }\n\n // Create matrix descriptors. Not setting any extra attributes.\n NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Adesc, A_type, param.transA == CUBLAS_OP_N ? m : k,\n param.transA == CUBLAS_OP_N ? k : m, param.lda));\n NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Bdesc, B_type, param.transB == CUBLAS_OP_N ? k : n,\n param.transB == CUBLAS_OP_N ? n : k, param.ldb));\n\n NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Ddesc, D_type, m, n, ldd));\n\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescCreate(&operationDesc, gemm_compute_type, CUDA_R_32F));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSA,\n ¶m.transA, sizeof(param.transA)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSB,\n ¶m.transB, sizeof(param.transB)));\n // Set math SM count\n if (math_sm_count != 0) {\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET,\n &math_sm_count, sizeof(math_sm_count)));\n }\n\n // set fp8/fp4 attributes -- input and output types should already be set to fp8/fp4\n // as appropriate. Note: gelu fusion isn't available right now, and we don't need\n // amax(D) either (next op is high precision).\n const bool mxfp8_gemm = !use_fp4 && is_mxfp8_scaling(inputA->scaling_mode);\n\n if (use_fp8 || use_fp4) {\n // Fast accumulation is only supported for FP8.\n const int8_t fastAccuMode = (use_split_accumulator) ? 0 : use_fp8;\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_FAST_ACCUM,\n &fastAccuMode, sizeof(fastAccuMode)));\n\n // Scaling factors.\n#if CUBLAS_VERSION >= 120800\n cublasLtMatmulMatrixScale_t scaling_mode_a;\n cublasLtMatmulMatrixScale_t scaling_mode_b;\n#endif // CUBLAS_VERSION >= 120800\n if (is_tensor_scaling(inputA->scaling_mode) && is_tensor_scaling(inputB->scaling_mode)) {\n void *A_scale_inverse = param.A_scale_inv;\n void *B_scale_inverse = param.B_scale_inv;\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,\n &A_scale_inverse, sizeof(A_scale_inverse)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,\n &B_scale_inverse, sizeof(B_scale_inverse)));\n#if CUBLAS_VERSION >= 120800\n scaling_mode_a = CUBLASLT_MATMUL_MATRIX_SCALE_SCALAR_32F;\n scaling_mode_b = CUBLASLT_MATMUL_MATRIX_SCALE_SCALAR_32F;\n#endif // CUBLAS_VERSION >= 120800\n } else if (mxfp8_gemm) {\n#if CUBLAS_VERSION >= 120800\n NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120800,\n \"MXFP8 requires cuBLAS 12.8+, but run-time cuBLAS version is \",\n transformer_engine::cuda::cublas_version());\n\n // Check that scales are in expected format\n NVTE_CHECK(inputA->with_gemm_swizzled_scales,\n \"MXFP8 scales are not in format expected by GEMM\");\n NVTE_CHECK(inputB->with_gemm_swizzled_scales,\n \"MXFP8 scales are not in format expected by GEMM\");\n\n // Configure cuBLAS scales\n fp8e8m0 *A_scale_inverse = reinterpret_cast(param.A_scale_inv);\n fp8e8m0 *B_scale_inverse = reinterpret_cast(param.B_scale_inv);\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,\n &A_scale_inverse, sizeof(A_scale_inverse)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,\n &B_scale_inverse, sizeof(B_scale_inverse)));\n scaling_mode_a = CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0;\n scaling_mode_b = CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0;\n\n // Workaround for heuristic cache bug in cublasLt. This separates the MXFP8 cache key from non-block scaling.\n // CUBLASLT_MATMUL_DESC_ALPHA_VECTOR_BATCH_STRIDE is unused for block scaling so it's safe to set.\n if (transformer_engine::cuda::cublas_version() <= 120803) {\n const int64_t dummy_a_vec_stride = 1;\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_ALPHA_VECTOR_BATCH_STRIDE, &dummy_a_vec_stride,\n sizeof(dummy_a_vec_stride)));\n }\n#else\n NVTE_ERROR(\"MXFP8 requires cuBLAS 12.8+, but compile-time cuBLAS version is \",\n CUBLAS_VERSION);\n#endif // CUBLAS_VERSION >= 120800\n } else if (use_fp4) { // NVFP4 GEMM\n#if CUBLAS_VERSION >= 120800\n NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120800,\n \"FP4 requires cuBLAS 12.8+, but run-time cuBLAS version is \",\n transformer_engine::cuda::cublas_version());\n\n // Check that scales are in expected format\n NVTE_CHECK(inputA->with_gemm_swizzled_scales,\n \"NVFP4 block scales are not in format expected by GEMM\");\n NVTE_CHECK(inputB->with_gemm_swizzled_scales,\n \"NVFP4 block scales are not in format expected by GEMM\");\n\n // alpha and beta are device pointers to FP32\n const cublasDataType_t scale_type = CUDA_R_32F;\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_SCALE_TYPE, &scale_type, sizeof(scale_type)));\n const cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_DEVICE;\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_POINTER_MODE, &pointer_mode, sizeof(pointer_mode)));\n\n // Configure cuBLAS scales\n fp8e4m3 *A_scale_inverse = reinterpret_cast(param.A_scale_inv);\n fp8e4m3 *B_scale_inverse = reinterpret_cast(param.B_scale_inv);\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,\n &A_scale_inverse, sizeof(A_scale_inverse)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,\n &B_scale_inverse, sizeof(B_scale_inverse)));\n scaling_mode_a = CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3;\n scaling_mode_b = CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3;\n#else\n NVTE_ERROR(\"FP4 requires cuBLAS 12.8+, but compile-time cuBLAS version is \", CUBLAS_VERSION);\n#endif // CUBLAS_VERSION >= 120800\n } else if ((inputA->scaling_mode == NVTE_BLOCK_SCALING_1D ||\n inputA->scaling_mode == NVTE_BLOCK_SCALING_2D) &&\n (inputB->scaling_mode == NVTE_BLOCK_SCALING_1D ||\n inputB->scaling_mode == NVTE_BLOCK_SCALING_2D)) {\n#if CUBLAS_VERSION >= 120900\n NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120900,\n \"FP8 block scaling requires cuBLAS 12.9+, but run-time cuBLAS version is \",\n transformer_engine::cuda::cublas_version());\n\n // Check that matrix formats are valid\n NVTE_CHECK((!(inputA->scaling_mode == NVTE_BLOCK_SCALING_2D &&\n inputB->scaling_mode == NVTE_BLOCK_SCALING_2D)),\n \"Only 1D by 1D, 1D by 2D, and 2D by 1D block scaling GEMM is supported, \"\n \"but got 2D by 2D\");\n\n // Configure cuBLAS scales\n float *A_scale_inverse = reinterpret_cast(param.A_scale_inv);\n float *B_scale_inverse = reinterpret_cast(param.B_scale_inv);\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,\n &A_scale_inverse, sizeof(A_scale_inverse)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,\n &B_scale_inverse, sizeof(B_scale_inverse)));\n scaling_mode_a = inputA->scaling_mode == NVTE_BLOCK_SCALING_1D\n ? CUBLASLT_MATMUL_MATRIX_SCALE_VEC128_32F\n : CUBLASLT_MATMUL_MATRIX_SCALE_BLK128x128_32F;\n scaling_mode_b = inputB->scaling_mode == NVTE_BLOCK_SCALING_1D\n ? CUBLASLT_MATMUL_MATRIX_SCALE_VEC128_32F\n : CUBLASLT_MATMUL_MATRIX_SCALE_BLK128x128_32F;\n#else\n NVTE_ERROR(\"FP8 block scaling requires cuBLAS 12.9+, but compile-time cuBLAS version is \",\n CUBLAS_VERSION);\n#endif // CUBLAS_VERSION >= 120900\n } else {\n NVTE_ERROR(\"Not implemented scaling modes: \" + to_string(inputA->scaling_mode) + \" and \" +\n to_string(inputB->scaling_mode) + \".\");\n }\n\n#if CUBLAS_VERSION >= 120800\n if (transformer_engine::cuda::cublas_version() >= 120800) {\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_A_SCALE_MODE,\n &scaling_mode_a, sizeof(scaling_mode_a)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_B_SCALE_MODE,\n &scaling_mode_b, sizeof(scaling_mode_b)));\n }\n#endif // CUBLAS_VERSION >= 120800\n if (is_fp8_dtype(outputD->data.dtype)) {\n // Accumulation mode not supported for FP8 output\n C = nullptr;\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_D_SCALE_POINTER, &D_scale, sizeof(D_scale)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_AMAX_D_POINTER, &D_amax, sizeof(D_amax)));\n#if CUBLAS_VERSION >= 120800\n if (transformer_engine::cuda::cublas_version() >= 120800) {\n // NOTE: In all current cases where FP8 output is supported, the input is\n // scaled identically to the output.\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_D_SCALE_MODE,\n &scaling_mode_a, sizeof(scaling_mode_a)));\n }\n#endif // CUBLAS_VERSION >= 120800\n // For FP8 output, cuBLAS requires C_type to match bias_type and\n // be FP16/BF16\n const cudaDataType_t C_type = bias ? bias_type : CUDA_R_16BF;\n NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, C_type, m, n, ldd));\n } else {\n NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));\n }\n if (bias) {\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE, &bias_type, sizeof(bias_type)));\n }\n } else {\n NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));\n }\n\n if (bias && gelu) {\n if (grad) {\n epilogue = CUBLASLT_EPILOGUE_DGELU_BGRAD;\n } else {\n epilogue = CUBLASLT_EPILOGUE_GELU_AUX_BIAS;\n }\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, sizeof(bias_ptr)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER,\n &pre_gelu_out, sizeof(pre_gelu_out)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_LD, &ld_gelumat, sizeof(ld_gelumat)));\n const cudaDataType_t aux_type = get_cuda_dtype(outputPreGelu->data.dtype);\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_DATA_TYPE, &aux_type, sizeof(aux_type)));\n } else if (bias) {\n if (grad) {\n // grad output is always input B\n epilogue = CUBLASLT_EPILOGUE_BGRADB;\n } else {\n epilogue = CUBLASLT_EPILOGUE_BIAS;\n }\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, sizeof(bias_ptr)));\n } else if (gelu) {\n if (grad) {\n epilogue = CUBLASLT_EPILOGUE_DGELU;\n } else {\n epilogue = CUBLASLT_EPILOGUE_GELU_AUX;\n }\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER,\n &pre_gelu_out, sizeof(pre_gelu_out)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_LD, &ld_gelumat, sizeof(ld_gelumat)));\n const cudaDataType_t aux_type = get_cuda_dtype(outputPreGelu->data.dtype);\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_DATA_TYPE, &aux_type, sizeof(aux_type)));\n }\n\n if ((inputA->scaling_mode == NVTE_BLOCK_SCALING_1D) ||\n (inputA->scaling_mode == NVTE_BLOCK_SCALING_2D)) {\n NVTE_CHECK((epilogue == CUBLASLT_EPILOGUE_DEFAULT || epilogue == CUBLASLT_EPILOGUE_BIAS ||\n epilogue == CUBLASLT_EPILOGUE_DGELU),\n \"Epilogue requested outside of the available and tested cuBLAS functionality for \"\n \"float8 block scaled GEMM\");\n }\n\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE,\n &epilogue, sizeof(epilogue)));\n\n if (counter != nullptr) {\n#if !(CUDA_VERSION >= 12020 && CUDA_VERSION < 13000)\n NVTE_ERROR(\"Atomic GEMM requires CUDA >=12.2.0 and <13.0.0, but compile-time CUDA version is \",\n CUDA_VERSION);\n#elif !(CUBLAS_VERSION >= 120205 && CUBLAS_VERSION < 130000)\n NVTE_ERROR(\n \"Atomic GEMM requires cuBLAS >=12.2.5 and <13.0.0, but compile-time cuBLAS version is \",\n CUBLAS_VERSION);\n#else\n NVTE_CHECK(transformer_engine::cuda::cudart_version() >= 12020 &&\n transformer_engine::cuda::cudart_version() < 13000,\n \"Atomic GEMM requires CUDA >=12.2.0 and <13.0.0, but run-time CUDA version is \",\n transformer_engine::cuda::cudart_version());\n NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120205 &&\n transformer_engine::cuda::cublas_version() < 130000,\n \"Atomic GEMM requires cuBLAS >=12.2.5 and <13.0.0, but run-time cuBLAS version is \",\n transformer_engine::cuda::cublas_version());\n if (m_split == 0) m_split = 1;\n if (n_split == 0) n_split = 1;\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_ATOMIC_SYNC_NUM_CHUNKS_D_ROWS, &m_split,\n sizeof(m_split)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_ATOMIC_SYNC_NUM_CHUNKS_D_COLS, &n_split,\n sizeof(n_split)));\n if (gemm_producer) {\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_ATOMIC_SYNC_OUT_COUNTERS_POINTER, &counter,\n sizeof(counter)));\n } else {\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_ATOMIC_SYNC_IN_COUNTERS_POINTER, &counter,\n sizeof(counter)));\n }\n#endif\n }\n\n // align the workspace to 256 B\n const int required_alignment = 256;\n const auto original_workspace_alignment = _getAlignment(reinterpret_cast(workspace));\n uint8_t *aligned_workspace_ptr =\n reinterpret_cast(workspace) + required_alignment - original_workspace_alignment;\n workspaceSize = workspaceSize - required_alignment + original_workspace_alignment;\n const auto new_workspace_alignment =\n _getAlignment(reinterpret_cast(aligned_workspace_ptr));\n NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceCreate(&preference));\n NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute(\n preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize)));\n const auto A_alignment = _getAlignment(reinterpret_cast(param.A));\n const auto B_alignment = _getAlignment(reinterpret_cast(param.B));\n const auto C_alignment = _getAlignment(reinterpret_cast(C));\n const auto D_alignment = _getAlignment(reinterpret_cast(D));\n NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute(\n preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_A_BYTES, &A_alignment, sizeof(A_alignment)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute(\n preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_B_BYTES, &B_alignment, sizeof(B_alignment)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute(\n preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_C_BYTES, &C_alignment, sizeof(C_alignment)));\n NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute(\n preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_D_BYTES, &D_alignment, sizeof(D_alignment)));\n NVTE_CHECK(new_workspace_alignment % 256 == 0,\n \"cuBLAS workspace pointer must be aligned to 256 bytes, got \",\n new_workspace_alignment);\n\n const auto status =\n cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Cdesc, Ddesc, preference,\n 1, &heuristicResult, &returnedResults);\n NVTE_CHECK(status != CUBLAS_STATUS_NOT_SUPPORTED,\n \"Unable to find suitable cuBLAS GEMM algorithm\");\n NVTE_CHECK_CUBLAS(status);\n if (returnedResults == 0) NVTE_ERROR(\"Unable to find any suitable algorithms\");\n\n // D = alpha * (A * B) + beta * C\n NVTE_CHECK_CUBLAS(cublasLtMatmul(handle, operationDesc, alpha, /* alpha */\n param.A, /* A */\n Adesc, param.B, /* B */\n Bdesc, beta, /* beta */\n C, /* C */\n Cdesc, D, /* D */\n Ddesc, &heuristicResult.algo, /* algo */\n aligned_workspace_ptr, /* workspace */\n workspaceSize, stream)); /* stream */\n\n // Update FP8 scale-inv in output tensor\n // Note: This is a WAR for the case when we have fp8 output but D->scale_inv is not allocated.\n // TODO: Changing gemm interface so that D->scale_inv is allocated and the scale_inv can be\n // calculated here.\n if (is_fp8_dtype(outputD->data.dtype) && outputD->scale_inv.dptr) {\n update_tensor_scale_inv(outputD, stream);\n }\n\n NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceDestroy(preference));\n NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Ddesc));\n NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Cdesc));\n NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Bdesc));\n NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Adesc));\n NVTE_CHECK_CUBLAS(cublasLtMatmulDescDestroy(operationDesc));\n}\n\n} // namespace transformer_engine\n\nvoid nvte_cublas_gemm(const NVTETensor A, const NVTETensor B, NVTETensor D, const NVTETensor bias,\n NVTETensor pre_gelu_out, bool transa, bool transb, bool grad,\n NVTETensor workspace, bool accumulate, bool use_split_accumulator,\n int math_sm_count, cudaStream_t stream) {\n NVTE_API_CALL(nvte_cublas_gemm);\n using namespace transformer_engine;\n\n // Tensors\n const Tensor *inputA = convertNVTETensorCheck(A);\n const Tensor *inputB = convertNVTETensorCheck(B);\n Tensor *outputD = convertNVTETensorCheck(D);\n const Tensor *biasTensor = convertNVTETensor(bias);\n Tensor *outputGelu = convertNVTETensor(pre_gelu_out);\n Tensor *wspace = convertNVTETensor(workspace);\n\n // Scales\n const float alpha = 1;\n const float beta = accumulate ? 1 : 0;\n\n // Check for NVFP4\n // TODO Remove once alpha scale logic is moved into cublas_gemm function\n if (is_nvfp_scaling(inputA->scaling_mode) || is_nvfp_scaling(inputB->scaling_mode)) {\n NVTE_ERROR(\"nvte_cublas_gemm does not support NVFP4 data. Use nvte_cublas_gemm_v2 instead.\");\n }\n\n // Launch GEMM\n cublas_gemm(inputA, inputB, outputD, biasTensor, outputGelu, (transa) ? CUBLAS_OP_T : CUBLAS_OP_N,\n (transb) ? CUBLAS_OP_T : CUBLAS_OP_N, grad, wspace->data.dptr, wspace->data.shape[0],\n &alpha, &beta, use_split_accumulator, math_sm_count, 0, 0, false, nullptr, stream);\n}\n\nvoid nvte_cublas_gemm_v2(int transa, int transb, const float *alpha, const NVTETensor A,\n const NVTETensor B, const float *beta, const NVTETensor C, NVTETensor D,\n NVTETensor workspace, NVTEMatmulConfig config, cudaStream_t stream) {\n NVTE_API_CALL(nvte_cublas_gemm_v2);\n using namespace transformer_engine;\n\n // Data tensors\n const Tensor *A_tensor = convertNVTETensorCheck(A);\n const Tensor *B_tensor = convertNVTETensorCheck(B);\n const Tensor *C_tensor = convertNVTETensorCheck(C);\n Tensor *D_tensor = convertNVTETensorCheck(D);\n NVTE_CHECK(C_tensor == D_tensor,\n \"Currently nvte_cublas_gemm_v2 does not support different C and D tensors.\");\n\n // Workspace\n void *workspace_ptr = nullptr;\n size_t workspace_size = 0;\n Tensor *workspace_tensor = convertNVTETensor(workspace);\n if (workspace_tensor != nullptr) {\n workspace_ptr = workspace_tensor->data.dptr;\n workspace_size =\n get_buffer_size_bytes(workspace_tensor->data.numel(), workspace_tensor->data.dtype);\n }\n\n // Additional config\n MatmulConfig config_;\n if (config != nullptr) {\n config_ = *reinterpret_cast(config);\n }\n\n // Configure GEMM epilogue\n const bool with_grad_epilogue = (config_.dbias_tensor != nullptr || config_.with_dgelu_epilogue);\n if (with_grad_epilogue) {\n NVTE_CHECK(config_.bias_tensor == nullptr && !config_.with_gelu_epilogue,\n \"Invalid epilogue (bias=\", config_.bias_tensor != nullptr,\n \", dbias=\", config_.dbias_tensor != nullptr, \", gelu=\", config_.with_gelu_epilogue,\n \", dgelu=\", config_.with_dgelu_epilogue, \").\");\n }\n Tensor dummy_tensor;\n Tensor *epilogue_bias_tensor = &dummy_tensor;\n if (!with_grad_epilogue && config_.bias_tensor != nullptr) {\n epilogue_bias_tensor = convertNVTETensorCheck(config_.bias_tensor);\n } else if (with_grad_epilogue && config_.dbias_tensor != nullptr) {\n epilogue_bias_tensor = convertNVTETensorCheck(config_.dbias_tensor);\n }\n Tensor *epilogue_aux_tensor = &dummy_tensor;\n if (config_.with_gelu_epilogue || config_.with_dgelu_epilogue) {\n NVTE_CHECK(config_.epilogue_aux_tensor != nullptr,\n \"Requested epilogue (bias=\", config_.bias_tensor != nullptr,\n \", dbias=\", config_.dbias_tensor != nullptr, \", gelu=\", config_.with_gelu_epilogue,\n \", dgelu=\", config_.with_dgelu_epilogue, \") without providing aux tensor.\");\n epilogue_aux_tensor = convertNVTETensor(config_.epilogue_aux_tensor);\n }\n\n // Launch GEMM\n cublas_gemm(A_tensor, B_tensor, D_tensor, epilogue_bias_tensor, epilogue_aux_tensor,\n transa ? CUBLAS_OP_T : CUBLAS_OP_N, transb ? CUBLAS_OP_T : CUBLAS_OP_N,\n with_grad_epilogue, workspace_ptr, workspace_size, alpha, beta,\n config_.use_split_accumulator, config_.sm_count, 0, 0, false, nullptr, stream);\n}\n\nvoid nvte_cublas_gemm_scaled(const NVTETensor A, const NVTETensor B, NVTETensor D,\n const NVTETensor bias, NVTETensor pre_gelu_out, bool transa,\n bool transb, bool grad, NVTETensor workspace, float alpha, float beta,\n bool use_split_accumulator, int math_sm_count, cudaStream_t stream) {\n NVTE_API_CALL(nvte_cublas_gemm);\n using namespace transformer_engine;\n\n // Tensors\n const Tensor *inputA = convertNVTETensorCheck(A);\n const Tensor *inputB = convertNVTETensorCheck(B);\n Tensor *outputD = convertNVTETensorCheck(D);\n const Tensor *biasTensor = convertNVTETensor(bias);\n Tensor *outputGelu = convertNVTETensor(pre_gelu_out);\n Tensor *wspace = convertNVTETensor(workspace);\n\n // Check for NVFP4\n // TODO Remove once alpha scale logic is moved into cublas_gemm function\n if (is_nvfp_scaling(inputA->scaling_mode) || is_nvfp_scaling(inputB->scaling_mode)) {\n NVTE_ERROR(\"nvte_cublas_gemm does not support NVFP4 data. Use nvte_cublas_gemm_v2 instead.\");\n }\n\n // Launch GEMM\n cublas_gemm(inputA, inputB, outputD, biasTensor, outputGelu, (transa) ? CUBLAS_OP_T : CUBLAS_OP_N,\n (transb) ? CUBLAS_OP_T : CUBLAS_OP_N, grad, wspace->data.dptr, wspace->data.shape[0],\n &alpha, &beta, use_split_accumulator, math_sm_count, 0, 0, false, nullptr, stream);\n}\n\nvoid nvte_cublas_atomic_gemm(const NVTETensor A, const NVTETensor B, NVTETensor D,\n const NVTETensor bias, NVTETensor pre_gelu_out, bool transa,\n bool transb, bool grad, NVTETensor workspace, bool accumulate,\n bool use_split_accumulator, int math_sm_count, int m_split,\n int n_split, bool gemm_producer, const NVTETensor counter,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_cublas_atomic_gemm);\n using namespace transformer_engine;\n#if !(CUDA_VERSION >= 12020 && CUDA_VERSION < 13000)\n NVTE_ERROR(\"Atomic GEMM requires CUDA >=12.2.0 and <13.0.0, but compile-time CUDA version is \",\n CUDA_VERSION);\n#elif !(CUBLAS_VERSION >= 120205 && CUBLAS_VERSION < 130000)\n NVTE_ERROR(\n \"Atomic GEMM requires cuBLAS >=12.2.5 and <13.0.0, but compile-time cuBLAS version is \",\n CUBLAS_VERSION);\n#else\n NVTE_CHECK(\n transformer_engine::cuda::cudart_version() >= 12020 &&\n transformer_engine::cuda::cudart_version() < 13000,\n \"Atomic GEMM requires CUDA version >=12.2.0 and <13.0.0, but run-time CUDA version is \",\n transformer_engine::cuda::cudart_version());\n NVTE_CHECK(\n transformer_engine::cuda::cublas_version() >= 120205 &&\n transformer_engine::cuda::cublas_version() < 130000,\n \"Atomic GEMM requires cuBLAS version >=12.2.5 and <13.0.0, but run-time cuBLAS version is \",\n transformer_engine::cuda::cublas_version());\n\n const Tensor *inputA = convertNVTETensorCheck(A);\n const Tensor *inputB = convertNVTETensorCheck(B);\n Tensor *outputD = convertNVTETensor(D);\n const Tensor *biasTensor = convertNVTETensor(bias);\n Tensor *outputGelu = convertNVTETensor(pre_gelu_out);\n const Tensor *inputCounter = convertNVTETensor(counter);\n Tensor *wspace = convertNVTETensor(workspace);\n\n const void *alpha_ptr = GetScalarOne();\n const void *beta_ptr = accumulate ? GetScalarOne() : GetScalarZero();\n\n NVTE_CHECK(is_delayed_tensor_scaling(inputA->scaling_mode) &&\n is_delayed_tensor_scaling(inputB->scaling_mode),\n \"Atomic GEMM only supports delayed scaling.\");\n cublas_gemm(inputA, inputB, outputD, biasTensor, outputGelu, (transa) ? CUBLAS_OP_T : CUBLAS_OP_N,\n (transb) ? CUBLAS_OP_T : CUBLAS_OP_N, grad, wspace->data.dptr, wspace->data.shape[0],\n alpha_ptr, beta_ptr, use_split_accumulator, math_sm_count, m_split, n_split,\n gemm_producer, inputCounter, stream);\n#endif\n}\n\nvoid multi_stream_cublas_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor *D,\n const NVTETensor *bias, NVTETensor *pre_gelu_out, const int num_gemms,\n bool transa, bool transb, bool grad, NVTETensor *workspace,\n bool accumulate, bool use_split_accumulator, int math_sm_count,\n cudaStream_t stream) {\n using namespace transformer_engine;\n\n int num_streams = nvte_get_num_compute_streams();\n\n int num_stream_used = std::min(num_streams, num_gemms);\n // wait for current stream to finish\n NVTE_CHECK_CUDA(cudaEventRecord(detail::get_compute_stream_event(0), stream));\n for (int s = 0; s < num_stream_used; s++) {\n NVTE_CHECK_CUDA(\n cudaStreamWaitEvent(detail::get_compute_stream(s), detail::get_compute_stream_event(0)));\n }\n\n for (int i = 0; i < num_gemms; i++) {\n // Check whether GELU or dGELU epilogue is requested\n Tensor *pre_gelu_tensor = convertNVTETensor(pre_gelu_out[i]);\n bool with_gelu_dgelu_epilogue =\n (pre_gelu_tensor != nullptr && pre_gelu_tensor->data.dptr != nullptr);\n\n // Construct config\n MatmulConfig config;\n if (grad) {\n config.dbias_tensor = bias[i];\n config.with_dgelu_epilogue = with_gelu_dgelu_epilogue;\n } else {\n config.bias_tensor = bias[i];\n config.with_gelu_epilogue = with_gelu_dgelu_epilogue;\n }\n config.epilogue_aux_tensor = pre_gelu_out[i];\n config.use_split_accumulator = use_split_accumulator;\n config.sm_count = math_sm_count;\n\n // Launch GEMM\n const float alpha = 1.f;\n const float beta = accumulate ? 1.f : 0.f;\n nvte_cublas_gemm_v2(transa, transb, &alpha, A[i], B[i], &beta, D[i], D[i],\n workspace[i % num_streams], &config,\n detail::get_compute_stream(i % num_streams));\n }\n\n // record events on compute streams\n for (int s = 0; s < num_stream_used; s++) {\n NVTE_CHECK_CUDA(\n cudaEventRecord(detail::get_compute_stream_event(s), detail::get_compute_stream(s)));\n }\n // wait for all compute streams to finish\n for (int s = 0; s < num_stream_used; s++) {\n NVTE_CHECK_CUDA(cudaStreamWaitEvent(stream, detail::get_compute_stream_event(s)));\n }\n}\n\nvoid nvte_multi_stream_cublas_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor *D,\n const NVTETensor *bias, NVTETensor *pre_gelu_out,\n const int num_gemms, bool transa, bool transb, bool grad,\n NVTETensor *workspace, bool accumulate,\n bool use_split_accumulator, int math_sm_count,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_multi_stream_cublas_gemm);\n using namespace transformer_engine;\n\n // Deprecation warning\n NVTE_WARN(\n \"nvte_multi_stream_cublas_gemm is deprecated and will be removed in a future release. \"\n \"Please migrate to nvte_multi_tensor_gemm (with CUTLASS Grouped GEMM support when \"\n \"applicable).\");\n\n multi_stream_cublas_gemm(A, B, D, bias, pre_gelu_out, num_gemms, transa, transb, grad, workspace,\n accumulate, use_split_accumulator, math_sm_count, stream);\n}\n\nnamespace transformer_engine {\n\nusing cublasHandleManager = detail::HandleManager;\n\nvoid nvte_cublas_handle_init() { auto _ = cublasHandleManager::Instance().GetHandle(); }\n\n} // namespace transformer_engine\n\nvoid nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor *D,\n const NVTETensor *bias, NVTETensor *pre_gelu_out, const int num_gemms,\n bool transa, bool transb, bool grad, NVTETensor *workspace,\n bool accumulate, bool use_split_accumulator, int math_sm_count,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_gemm);\n\n const int current_device = transformer_engine::cuda::current_device();\n const bool is_hopper = (transformer_engine::cuda::sm_arch(current_device) == 90);\n const bool use_cutlass = transformer_engine::getenv(\"NVTE_USE_CUTLASS_GROUPED_GEMM\", false);\n const bool warn_fallback =\n transformer_engine::getenv(\"NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK\", false);\n\n auto cublas_path = [&]() {\n multi_stream_cublas_gemm(A, B, D, bias, pre_gelu_out, num_gemms, transa, transb, grad,\n workspace, accumulate, use_split_accumulator, math_sm_count, stream);\n };\n\n // Currently only support cutlass group gemm on Hopper Arch\n if (!(is_hopper && use_cutlass)) {\n cublas_path();\n return;\n }\n\n auto is_empty_arr = [&](const NVTETensor *p) -> bool {\n if (p == nullptr) return true;\n for (int i = 0; i < num_gemms; ++i) {\n if (transformer_engine::convertNVTETensor(p[i])->has_data()) return false;\n }\n return true;\n };\n\n auto all_groups_uniform_k128 = [&](const NVTETensor *p, bool trans) -> bool {\n int64_t ref_k = -1;\n for (size_t i = 0; i < num_gemms; i++) {\n const auto tensor = transformer_engine::convertNVTETensorCheck(p[i]);\n const int k = trans ? tensor->data.shape[0] : tensor->data.shape[1];\n\n if ((k & 127) != 0) return false;\n\n if (ref_k < 0)\n ref_k = k;\n else if (k != ref_k)\n return false;\n }\n\n return true;\n };\n\n auto is_supported_dtype = [&]() -> bool {\n auto *inputA = transformer_engine::convertNVTETensorCheck(A[0]);\n auto *inputB = transformer_engine::convertNVTETensorCheck(B[0]);\n auto *OutputD = transformer_engine::convertNVTETensorCheck(D[0]);\n auto A_type = get_cuda_dtype(inputA->data.dtype);\n auto B_type = get_cuda_dtype(inputB->data.dtype);\n auto D_type = get_cuda_dtype(OutputD->data.dtype);\n\n return (A_type == B_type) && (A_type == D_type) &&\n ((A_type == CUDA_R_16BF) || (A_type == CUDA_R_16F));\n };\n\n // CUTLASS Grouped GEMM fast path (SM90/TMA)\n // Conditions:\n // - No fused epilogue: both bias and pre_gelu_out are empty.\n // - Supported dtypes only: FP16/BF16 (FP32 accumulate).\n // - Uniform K across groups and K % 128 == 0.\n // - use_split_accumulator is ignored for FP16/BF16.\n // - grad is irrelevant when bias/pre_gelu_out are empty.\n //\n // Otherwise, fall back to cuBLAS.\n if (is_empty_arr(bias) && is_empty_arr(pre_gelu_out) && is_supported_dtype() &&\n all_groups_uniform_k128(B, transb)) {\n cutlass_grouped_gemm(A, B, D, num_gemms, transa, transb, grad, workspace, accumulate,\n current_device, math_sm_count, stream);\n } else {\n if (warn_fallback) {\n NVTE_WARN(\"Fallback to cuBLAS grouped GEMM.\");\n }\n cublas_path();\n }\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/hip/hip_runtime.h\"\n#include \"../util/handle_manager.h\"\n#include \"../util/logging.h\"\n#include \"../util/multi_stream.h\"\n#include \"./config.h\"\n#include \"./cutlass_grouped_gemm.cuh\"\n\nnamespace {\n\n/* Use CUDA const memory to store scalar 1 and 0 for cublas usage\n*/\n__device__ __constant__ float one_device;\n__device__ __constant__ float zero_device;\n\ninline float *GetScalarOne() {\n static std::once_flag init_flag;\n std::call_once(init_flag, []() {\n float one = 1.0f;\n NVTE_CHECK_CUDA(hipMemcpyToSymbol(HIP_SYMBOL(one_device), &one, sizeof(float)));\n });\n // return address by cudaGetSymbolAddress\n float *dev_ptr;\n NVTE_CHECK_CUDA(hipGetSymbolAddress(reinterpret_cast(&dev_ptr),HIP_SYMBOL(one_device)));\n return dev_ptr;\n}\n\ninline float *GetScalarZero() {\n static std::once_flag init_flag;\n std::call_once(init_flag, []() {\n float zero = 0.0f;\n NVTE_CHECK_CUDA(hipMemcpyToSymbol(HIP_SYMBOL(zero_device), &zero, sizeof(float)));\n });\n // return address by cudaGetSymbolAddress\n float *dev_ptr;\n NVTE_CHECK_CUDA(hipGetSymbolAddress(reinterpret_cast(&dev_ptr),HIP_SYMBOL(zero_device)));\n return dev_ptr;\n}\n\n__global__ __launch_bounds__(1) void set_float_kernel(float *ptr, float val) { *ptr = val; }\n\nuint32_t _getAlignment(uintptr_t address) {\n // alignment are in bytes\n uint32_t alignment = 256;\n for (;; alignment /= 2) {\n if (address % alignment == 0) {\n return alignment;\n }\n }\n}\n\ninline void CreateCublasHandle(hipblasLtHandle_t *handle) {\n NVTE_CHECK_CUBLAS(hipblasLtCreate(handle));\n}\n\n/* Parameters for cuBLAS GEMM\n *\n * cuBLAS follows the BLAS convention of column-major ordering. This\n * is different than the row-major that is typically used in\n * Transformer Engine.\n *\n */\nstruct GemmParam {\n void *A = nullptr;\n void *B = nullptr;\n hipblasOperation_t transA = HIPBLAS_OP_N;\n hipblasOperation_t transB = HIPBLAS_OP_N;\n transformer_engine::DType Atype = transformer_engine::DType::kNumTypes;\n transformer_engine::DType Btype = transformer_engine::DType::kNumTypes;\n void *A_scale_inv = nullptr;\n void *B_scale_inv = nullptr;\n int lda = 0; // A column strides\n int ldb = 0; // B column strides\n};\n\n/* Populate parameters for cuBLAS GEMM\n *\n * cuBLAS follows the BLAS convention of column-major ordering. This\n * is different than the row-major that is typically used in\n * Transformer Engine.\n *\n */\nGemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const hipblasOperation_t transA,\n const transformer_engine::Tensor &B, const hipblasOperation_t transB,\n int m, int n, int k) {\n using namespace transformer_engine;\n NVTE_CHECK(\n A.scaling_mode == B.scaling_mode ||\n (A.scaling_mode == NVTE_BLOCK_SCALING_1D && B.scaling_mode == NVTE_BLOCK_SCALING_2D) ||\n (A.scaling_mode == NVTE_BLOCK_SCALING_2D && B.scaling_mode == NVTE_BLOCK_SCALING_1D),\n \"Inputs A and B to GEMM need to have compatible scaling modes, but got A.scaling_mode = \" +\n to_string(A.scaling_mode) + \", B.scaling_mode = \" + to_string(B.scaling_mode));\n NVTE_CHECK(A.has_data() || A.has_columnwise_data(), \"Input A does not hold any data!\");\n NVTE_CHECK(B.has_data() || B.has_columnwise_data(), \"Input B does not hold any data!\");\n GemmParam ret;\n\n // Transpose mode with column-major ordering\n bool is_A_transposed = transA == HIPBLAS_OP_T;\n bool is_B_transposed = transB == HIPBLAS_OP_T;\n\n // Set conditions for MXFP8 and NVFP4 gemm execution.\n const auto nvfp4 = is_nvfp_scaling(A.scaling_mode) && is_nvfp_scaling(B.scaling_mode);\n const auto mxfp8 = !nvfp4 && is_mxfp_scaling(A.scaling_mode) && is_mxfp_scaling(B.scaling_mode);\n int is_nvte_non_tn_fp8_gemm_supported = 0; // needed only for per tensor scaling\n if (is_tensor_scaling(A.scaling_mode) || is_tensor_scaling(B.scaling_mode)) {\n is_nvte_non_tn_fp8_gemm_supported = nvte_is_non_tn_fp8_gemm_supported();\n }\n\n // Configure A matrix\n if (is_tensor_scaling(A.scaling_mode)) {\n // Unscaled or FP8 tensor scaling\n ret.A = A.data.dptr;\n ret.transA = transA;\n ret.Atype = A.data.dtype;\n ret.A_scale_inv = A.scale_inv.dptr;\n ret.lda = is_A_transposed ? k : m;\n if (!is_nvte_non_tn_fp8_gemm_supported && !is_A_transposed) {\n // Hopper only supports TN GEMMs for FP8. \"Column-wise data\" is transpose of data.\n if (A.has_columnwise_data() && is_fp8_dtype(A.columnwise_data.dtype)) {\n ret.A = A.columnwise_data.dptr;\n ret.transA = HIPBLAS_OP_T;\n ret.Atype = A.columnwise_data.dtype;\n ret.A_scale_inv = A.columnwise_scale_inv.dptr;\n ret.lda = k;\n } else {\n NVTE_CHECK(!is_fp8_dtype(ret.Atype), \"Input A is missing column-wise usage\");\n }\n } else if (is_nvte_non_tn_fp8_gemm_supported && !A.has_data()) {\n // Blackwell supports any GEMM layout for FP8, so we can use column-wise/transposed\n // data with the mirrored transpose-flag if we don't have row-wise data.\n NVTE_CHECK(A.has_columnwise_data() && is_fp8_dtype(A.columnwise_data.dtype),\n \"Input A is missing column-wise usage\");\n ret.A = A.columnwise_data.dptr;\n ret.transA = is_A_transposed ? HIPBLAS_OP_N : HIPBLAS_OP_T;\n ret.Atype = A.columnwise_data.dtype;\n ret.A_scale_inv = A.columnwise_scale_inv.dptr;\n ret.lda = is_A_transposed ? m : k;\n }\n\n if (is_fp8_dtype(ret.Atype)) {\n // Requirements from https://docs.nvidia.com/cuda/cublas/#tensor-core-usage\n NVTE_CHECK(ret.lda % 16 == 0,\n \"Leading dimension requirement on A for FP8 GEMM. Caller must pad.\");\n }\n } else if (nvfp4) {\n // NVFP4 GEMM. Either the pure NVFP4 recipe or the FWD pass of the Hybrid NVFP4/MXFP8 recipe.\n\n if (is_A_transposed) {\n NVTE_CHECK(A.has_data(), \"Input A is missing row-wise usage\");\n } else {\n NVTE_CHECK(is_nvfp4_scaling(A.scaling_mode),\n \"Input A has unsupported combination of recipe and layout\");\n NVTE_CHECK(A.has_columnwise_data(), \"Input A is missing column-wise usage\");\n }\n ret.A = is_A_transposed ? A.data.dptr : A.columnwise_data.dptr;\n ret.transA = HIPBLAS_OP_T; // NVFP4 gemm is only supported in TN layout.\n ret.Atype = is_A_transposed ? A.data.dtype : A.columnwise_data.dtype;\n ret.A_scale_inv = is_A_transposed ? A.scale_inv.dptr : A.columnwise_scale_inv.dptr;\n ret.lda = k;\n } else if (mxfp8) {\n // MXFP8 GEMM. Either for pure MXFP8 recipe or backward of Hybrid NVFP4 recipe.\n // Note: Row-wise and column-wise data are scaled along different\n // dimensions (with matrix interpreted in row-major order).\n\n if (is_A_transposed) {\n NVTE_CHECK(A.has_data(), \"Input A is missing row-wise usage\");\n } else {\n NVTE_CHECK(A.has_columnwise_data(), \"Input A is missing column-wise usage\");\n }\n ret.A = is_A_transposed ? A.data.dptr : A.columnwise_data.dptr;\n ret.transA = transA;\n ret.Atype = is_A_transposed ? A.data.dtype : A.columnwise_data.dtype;\n ret.A_scale_inv = is_A_transposed ? A.scale_inv.dptr : A.columnwise_scale_inv.dptr;\n ret.lda = is_A_transposed ? k : m;\n } else if (A.scaling_mode == NVTE_BLOCK_SCALING_1D || A.scaling_mode == NVTE_BLOCK_SCALING_2D) {\n // FP8 block scaling\n // Note: Hopper only supports TN GEMMs for FP8. \"Column-wise data\" is transpose of data.\n if (is_A_transposed) {\n NVTE_CHECK(A.has_data(), \"Input A is missing row-wise usage\");\n } else {\n NVTE_CHECK(A.has_columnwise_data(), \"Input A is missing column-wise usage\");\n }\n ret.A = is_A_transposed ? A.data.dptr : A.columnwise_data.dptr;\n ret.transA = HIPBLAS_OP_T;\n ret.Atype = is_A_transposed ? A.data.dtype : A.columnwise_data.dtype;\n ret.A_scale_inv = is_A_transposed ? A.scale_inv.dptr : A.columnwise_scale_inv.dptr;\n ret.lda = k;\n\n // Requirements from https://docs.nvidia.com/cuda/cublas/#tensor-core-usage\n NVTE_CHECK((ret.lda % 16) == 0,\n \"Leading dimension requirement on NVTE_BLOCK_SCALING GEMM. Caller must pad.\");\n // Divisibility of 8 derived from FP8 (m * CTypeSize) % 16 == 0 requirement.\n // Smallest supported CType is 2 bytes in this scaling mode.\n NVTE_CHECK((m % 8) == 0,\n \"Outer dimension requirement on A for NVTE_BLOCK_SCALING GEMM. Caller must pad.\");\n } else {\n NVTE_ERROR(\"A has unsupported scaling mode\");\n }\n\n // Configure B matrix\n if (is_tensor_scaling(B.scaling_mode)) {\n // Unscaled or FP8 tensor scaling\n ret.B = B.data.dptr;\n ret.transB = transB;\n ret.Btype = B.data.dtype;\n ret.B_scale_inv = B.scale_inv.dptr;\n ret.ldb = is_B_transposed ? n : k;\n if (!is_nvte_non_tn_fp8_gemm_supported && is_B_transposed) {\n // Hopper only supports TN GEMMs for FP8. \"Column-wise data\" is transpose of data.\n if (B.has_columnwise_data() && is_fp8_dtype(B.columnwise_data.dtype)) {\n ret.B = B.columnwise_data.dptr;\n ret.transB = HIPBLAS_OP_N;\n ret.Btype = B.columnwise_data.dtype;\n ret.B_scale_inv = B.columnwise_scale_inv.dptr;\n ret.ldb = k;\n } else {\n NVTE_CHECK(!is_fp8_dtype(ret.Btype), \"Input B is missing column-wise usage\");\n }\n } else if (is_nvte_non_tn_fp8_gemm_supported && !B.has_data()) {\n // Blackwell supports any GEMM layout for FP8, so we can use column-wise/transposed\n // data with the mirrored transpose-flag if we don't have row-wise data.\n NVTE_CHECK(B.has_columnwise_data() && is_fp8_dtype(B.columnwise_data.dtype),\n \"Input B is missing column-wise usage\");\n ret.B = B.columnwise_data.dptr;\n ret.transB = is_B_transposed ? HIPBLAS_OP_N : HIPBLAS_OP_T;\n ret.Btype = B.columnwise_data.dtype;\n ret.B_scale_inv = B.columnwise_scale_inv.dptr;\n ret.ldb = is_B_transposed ? k : n;\n }\n\n if (is_fp8_dtype(ret.Atype)) {\n // Requirements from https://docs.nvidia.com/cuda/cublas/#tensor-core-usage\n NVTE_CHECK(ret.ldb % 16 == 0,\n \"Leading dimension requirement on B for FP8 GEMM. Caller must pad.\");\n }\n } else if (nvfp4) {\n if (is_B_transposed) {\n NVTE_CHECK(is_nvfp4_scaling(B.scaling_mode),\n \"Input B has unsupported combination of recipe and layout\");\n NVTE_CHECK(B.has_columnwise_data(), \"Input B is missing column-wise usage\");\n } else {\n NVTE_CHECK(B.has_data(), \"Input B is missing row-wise usage\");\n }\n ret.B = is_B_transposed ? B.columnwise_data.dptr : B.data.dptr;\n ret.transB = HIPBLAS_OP_N; // NVFP4 gemm is only supported in TN layout.\n ret.Btype = is_B_transposed ? B.columnwise_data.dtype : B.data.dtype;\n ret.B_scale_inv = is_B_transposed ? B.columnwise_scale_inv.dptr : B.scale_inv.dptr;\n ret.ldb = k;\n } else if (mxfp8) {\n if (is_B_transposed) {\n NVTE_CHECK(B.has_columnwise_data(), \"Input B is missing column-wise usage\");\n } else {\n NVTE_CHECK(B.has_data(), \"Input B is missing row-wise usage\");\n }\n ret.B = is_B_transposed ? B.columnwise_data.dptr : B.data.dptr;\n ret.transB = transB;\n ret.Btype = is_B_transposed ? B.columnwise_data.dtype : B.data.dtype;\n ret.B_scale_inv = is_B_transposed ? B.columnwise_scale_inv.dptr : B.scale_inv.dptr;\n ret.ldb = is_B_transposed ? n : k;\n } else if (B.scaling_mode == NVTE_BLOCK_SCALING_1D || B.scaling_mode == NVTE_BLOCK_SCALING_2D) {\n // FP8 block scaling\n // Note: Hopper only supports TN GEMMs for FP8. \"Column-wise data\" is transpose of data.\n if (is_B_transposed) {\n NVTE_CHECK(B.has_columnwise_data(), \"Input B is missing column-wise usage\");\n } else {\n NVTE_CHECK(B.has_data(), \"Input B is missing row-wise usage\");\n }\n ret.B = is_B_transposed ? B.columnwise_data.dptr : B.data.dptr;\n ret.transB = HIPBLAS_OP_N;\n ret.Btype = is_B_transposed ? B.columnwise_data.dtype : B.data.dtype;\n ret.B_scale_inv = is_B_transposed ? B.columnwise_scale_inv.dptr : B.scale_inv.dptr;\n ret.ldb = k;\n\n // Requirements from\n // https://docs.nvidia.com/cuda/cublas/#tensor-core-usage\n NVTE_CHECK((ret.ldb % 16) == 0,\n \"B tensor stride requirement on NVTE_BLOCK_SCALING GEMM. Caller must pad.\");\n if (B.scaling_mode == NVTE_BLOCK_SCALING_1D) {\n // Observed this requirement only present for B tensor is 1D quantized.\n NVTE_CHECK((n % 8) == 0,\n \"Outer dimension requirement on B for NVTE_BLOCK_SCALING GEMM. Caller must pad.\");\n }\n } else {\n NVTE_ERROR(\"B has unsupported scaling mode\");\n }\n\n return ret;\n}\n\n} // namespace\n\nnamespace transformer_engine {\n\nusing cublasHandleManager = detail::HandleManager;\n\nvoid cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD,\n const Tensor *inputBias, Tensor *outputPreGelu, hipblasOperation_t transa,\n hipblasOperation_t transb, bool grad, void *workspace, size_t workspaceSize,\n const void *alpha, const void *beta, bool use_split_accumulator, int math_sm_count,\n int m_split, int n_split, bool gemm_producer, const Tensor *inputCounter,\n hipStream_t stream) {\n NVTE_CHECK(!inputA->row_scaled_nvfp4 && !inputB->row_scaled_nvfp4,\n \"cuBLAS GEMM does not support row-scaled NVFP4 inputs.\");\n\n // Tensor dims in row-major order\n const int A0 = inputA->flat_first_dim();\n const int A1 = inputA->flat_last_dim();\n const int B0 = inputB->flat_first_dim();\n const int B1 = inputB->flat_last_dim();\n\n // GEMM dims in column-major order\n const int m = transa == HIPBLAS_OP_T ? A0 : A1;\n const int n = transb == HIPBLAS_OP_T ? B1 : B0;\n const int k = transa == HIPBLAS_OP_T ? A1 : A0;\n NVTE_CHECK((transb == HIPBLAS_OP_T ? B0 : B1) == k,\n \"GEMM inputs have incompatible dimensions (A is \", A0, \"x\", A1, \", B is \", B0, \"x\", B1,\n \")\");\n const int ldd = m;\n\n // Return immediately if GEMM is trivial\n if (m <= 0 || n <= 0) {\n return;\n }\n NVTE_CHECK(k > 0);\n\n const GemmParam param = CanonicalizeGemmInput(*inputA, transa, *inputB, transb, m, n, k);\n\n void *C = outputD->data.dptr;\n void *D = outputD->data.dptr;\n void *D_scale = outputD->scale.dptr;\n void *D_amax = outputD->amax.dptr;\n void *bias_ptr = inputBias->data.dptr;\n const bool bias = bias_ptr != nullptr;\n void *pre_gelu_out = outputPreGelu->data.dptr;\n void *counter = nullptr;\n if (inputCounter != nullptr) {\n counter = inputCounter->data.dptr;\n }\n const bool gelu = pre_gelu_out != nullptr;\n const bool use_fp8 = is_fp8_dtype(param.Atype) || is_fp8_dtype(param.Btype);\n const bool use_fp4 = is_fp4_dtype(param.Atype) || is_fp4_dtype(param.Btype);\n\n // Update scaling factors with NVFP4 tensor scales\n // TODO: Check whether scales are on CPU/GPU or add API to control.\n // Currently scales are assumed to be on CPU when amax is provided\n // and on GPU when not provided, but this is brittle.\n if (use_fp4 &&\n ((transa == HIPBLAS_OP_T ? inputA->amax.dptr : inputA->columnwise_amax.dptr) != nullptr ||\n (transb == HIPBLAS_OP_T ? inputB->columnwise_amax.dptr : inputB->amax.dptr) != nullptr)) {\n // Reserve some workspace for alpha scale\n NVTE_CHECK(workspaceSize >= 4,\n \"NVFP4 GEMM requires at least 4 byte workspace for alpha scale, but only has \",\n workspaceSize, \" bytes remaining.\");\n workspaceSize = (workspaceSize / 4) * 4 - 4; // Remove last 4 aligned bytes\n uint8_t *workspace_ptr = reinterpret_cast(workspace);\n float *new_alpha_ptr = reinterpret_cast(&workspace_ptr[workspaceSize]);\n\n // Update alpha scale on device\n // Note: Compute NVFP4 tensor scales based on amaxes and then\n // divide from alpha scale. This way we only need to apply NVFP4\n // tensor scales in matmul output, instead of in matmul inputs.\n float old_alpha = *reinterpret_cast(alpha); // Assumed to be on CPU\n TensorWrapper new_alpha_tensor(new_alpha_ptr, std::vector{1}, DType::kFloat32);\n bool a_rowwise_amax = transa == HIPBLAS_OP_T;\n bool b_rowwise_amax = transb != HIPBLAS_OP_T;\n nvte_nvfp4_compute_per_tensor_scale(inputA->nvte_tensor, a_rowwise_amax, inputB->nvte_tensor,\n b_rowwise_amax, old_alpha, new_alpha_tensor.data(), stream);\n alpha = new_alpha_ptr;\n\n // Make sure beta scale is on device\n float old_beta = *reinterpret_cast(beta); // Assumed to be on CPU\n if (old_beta == 0) {\n beta = GetScalarZero(); // Device constant memory\n } else if (old_beta == 1) {\n beta = GetScalarOne(); // Device constant memory\n } else {\n // Move beta to workspace\n NVTE_CHECK(workspaceSize >= 4,\n \"NVFP4 GEMM requires at least 4 byte workspace for beta scale, but only has \",\n workspaceSize, \" bytes remaining.\");\n workspaceSize = (workspaceSize / 4) * 4 - 4; // Remove last 4 aligned bytes\n float *new_beta_ptr = reinterpret_cast(&workspace_ptr[workspaceSize]);\n set_float_kernel<<<1, 1, 0, stream>>>(new_beta_ptr, old_beta);\n NVTE_CHECK_CUDA(hipGetLastError());\n beta = new_beta_ptr;\n }\n }\n\n const hipDataType A_type = get_cuda_dtype(param.Atype);\n const hipDataType B_type = get_cuda_dtype(param.Btype);\n const hipDataType D_type = get_cuda_dtype(outputD->data.dtype);\n const hipDataType bias_type = get_cuda_dtype(inputBias->data.dtype);\n\n NVTE_CHECK(!is_fp8_dtype(param.Atype) || param.A_scale_inv != nullptr,\n \"FP8 input to GEMM requires inverse of scale!\");\n NVTE_CHECK(!is_fp8_dtype(param.Btype) || param.B_scale_inv != nullptr,\n \"FP8 input to GEMM requires inverse of scale!\");\n NVTE_CHECK(!is_fp4_dtype(param.Atype) || param.A_scale_inv != nullptr,\n \"FP4 input to GEMM requires inverse of scale!\");\n NVTE_CHECK(!is_fp4_dtype(param.Btype) || param.B_scale_inv != nullptr,\n \"FP4 input to GEMM requires inverse of scale!\");\n\n // check consistency of arguments:\n // if fp8 is desired, context cannot be null\n // fp8 + gelu fusion + fp8 aux is unavailable right now.\n if ((use_fp8 || use_fp4) && gelu) {\n NVTE_CHECK(!is_fp8_dtype(outputPreGelu->data.dtype),\n \"fp8 Aux output for gemm + gelu fusion not supported!\");\n }\n if (is_fp4_dtype(outputD->data.dtype)) {\n NVTE_ERROR(\"FP4 GEMM output is not supported!\");\n }\n if (use_fp4 && (D_type == HIP_R_16F)) {\n NVTE_ERROR(\"FP4 GEMM does not support FP16 output!\");\n }\n\n hipblasLtHandle_t handle = cublasHandleManager::Instance().GetHandle();\n\n hipblasLtMatmulDesc_t operationDesc = nullptr;\n hipblasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr, Ddesc = nullptr;\n hipblasLtMatmulPreference_t preference = nullptr;\n int returnedResults = 0;\n hipblasLtMatmulHeuristicResult_t heuristicResult = {};\n hipblasLtEpilogue_t epilogue = HIPBLASLT_EPILOGUE_DEFAULT;\n\n int64_t ld_gelumat = (int64_t)ldd;\n\n // Use TF32 only for pure FP32 GEMM.\n hipblasComputeType_t gemm_compute_type = HIPBLAS_COMPUTE_32F;\n if (A_type == HIP_R_32F && B_type == HIP_R_32F && D_type == HIP_R_32F) {\n gemm_compute_type = HIPBLAS_COMPUTE_32F_FAST_TF32;\n }\n\n // Create matrix descriptors. Not setting any extra attributes.\n NVTE_CHECK_CUBLAS(hipblasLtMatrixLayoutCreate(&Adesc, A_type, param.transA == HIPBLAS_OP_N ? m : k,\n param.transA == HIPBLAS_OP_N ? k : m, param.lda));\n NVTE_CHECK_CUBLAS(hipblasLtMatrixLayoutCreate(&Bdesc, B_type, param.transB == HIPBLAS_OP_N ? k : n,\n param.transB == HIPBLAS_OP_N ? n : k, param.ldb));\n\n NVTE_CHECK_CUBLAS(hipblasLtMatrixLayoutCreate(&Ddesc, D_type, m, n, ldd));\n\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescCreate(&operationDesc, gemm_compute_type, HIP_R_32F));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc, HIPBLASLT_MATMUL_DESC_TRANSA,\n ¶m.transA, sizeof(param.transA)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc, HIPBLASLT_MATMUL_DESC_TRANSB,\n ¶m.transB, sizeof(param.transB)));\n // Set math SM count\n if (math_sm_count != 0) {\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET,\n &math_sm_count, sizeof(math_sm_count)));\n }\n\n // set fp8/fp4 attributes -- input and output types should already be set to fp8/fp4\n // as appropriate. Note: gelu fusion isn't available right now, and we don't need\n // amax(D) either (next op is high precision).\n const bool mxfp8_gemm = !use_fp4 && is_mxfp8_scaling(inputA->scaling_mode);\n\n if (use_fp8 || use_fp4) {\n // Fast accumulation is only supported for FP8.\n const int8_t fastAccuMode = (use_split_accumulator) ? 0 : use_fp8;\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_FAST_ACCUM,\n &fastAccuMode, sizeof(fastAccuMode)));\n\n // Scaling factors.\n#if CUBLAS_VERSION >= 120800\n hipblasLtMatmulMatrixScale_t scaling_mode_a;\n hipblasLtMatmulMatrixScale_t scaling_mode_b;\n#endif // CUBLAS_VERSION >= 120800\n if (is_tensor_scaling(inputA->scaling_mode) && is_tensor_scaling(inputB->scaling_mode)) {\n void *A_scale_inverse = param.A_scale_inv;\n void *B_scale_inverse = param.B_scale_inv;\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc,\n HIPBLASLT_MATMUL_DESC_A_SCALE_POINTER,\n &A_scale_inverse, sizeof(A_scale_inverse)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc,\n HIPBLASLT_MATMUL_DESC_B_SCALE_POINTER,\n &B_scale_inverse, sizeof(B_scale_inverse)));\n#if CUBLAS_VERSION >= 120800\n scaling_mode_a = HIPBLASLT_MATMUL_MATRIX_SCALE_SCALAR_32F;\n scaling_mode_b = HIPBLASLT_MATMUL_MATRIX_SCALE_SCALAR_32F;\n#endif // CUBLAS_VERSION >= 120800\n } else if (mxfp8_gemm) {\n#if CUBLAS_VERSION >= 120800\n NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120800,\n \"MXFP8 requires cuBLAS 12.8+, but run-time cuBLAS version is \",\n transformer_engine::cuda::cublas_version());\n\n // Check that scales are in expected format\n NVTE_CHECK(inputA->with_gemm_swizzled_scales,\n \"MXFP8 scales are not in format expected by GEMM\");\n NVTE_CHECK(inputB->with_gemm_swizzled_scales,\n \"MXFP8 scales are not in format expected by GEMM\");\n\n // Configure cuBLAS scales\n fp8e8m0 *A_scale_inverse = reinterpret_cast(param.A_scale_inv);\n fp8e8m0 *B_scale_inverse = reinterpret_cast(param.B_scale_inv);\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc,\n HIPBLASLT_MATMUL_DESC_A_SCALE_POINTER,\n &A_scale_inverse, sizeof(A_scale_inverse)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc,\n HIPBLASLT_MATMUL_DESC_B_SCALE_POINTER,\n &B_scale_inverse, sizeof(B_scale_inverse)));\n scaling_mode_a = HIPBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0;\n scaling_mode_b = HIPBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0;\n\n // Workaround for heuristic cache bug in cublasLt. This separates the MXFP8 cache key from non-block scaling.\n // CUBLASLT_MATMUL_DESC_ALPHA_VECTOR_BATCH_STRIDE is unused for block scaling so it's safe to set.\n if (transformer_engine::cuda::cublas_version() <= 120803) {\n const int64_t dummy_a_vec_stride = 1;\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_ALPHA_VECTOR_BATCH_STRIDE, &dummy_a_vec_stride,\n sizeof(dummy_a_vec_stride)));\n }\n#else\n NVTE_ERROR(\"MXFP8 requires cuBLAS 12.8+, but compile-time cuBLAS version is \",\n CUBLAS_VERSION);\n#endif // CUBLAS_VERSION >= 120800\n } else if (use_fp4) { // NVFP4 GEMM\n#if CUBLAS_VERSION >= 120800\n NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120800,\n \"FP4 requires cuBLAS 12.8+, but run-time cuBLAS version is \",\n transformer_engine::cuda::cublas_version());\n\n // Check that scales are in expected format\n NVTE_CHECK(inputA->with_gemm_swizzled_scales,\n \"NVFP4 block scales are not in format expected by GEMM\");\n NVTE_CHECK(inputB->with_gemm_swizzled_scales,\n \"NVFP4 block scales are not in format expected by GEMM\");\n\n // alpha and beta are device pointers to FP32\n const hipDataType scale_type = HIP_R_32F;\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_SCALE_TYPE, &scale_type, sizeof(scale_type)));\n const hipblasLtPointerMode_t pointer_mode = HIPBLASLT_POINTER_MODE_DEVICE;\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, HIPBLASLT_MATMUL_DESC_POINTER_MODE, &pointer_mode, sizeof(pointer_mode)));\n\n // Configure cuBLAS scales\n fp8e4m3 *A_scale_inverse = reinterpret_cast(param.A_scale_inv);\n fp8e4m3 *B_scale_inverse = reinterpret_cast(param.B_scale_inv);\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc,\n HIPBLASLT_MATMUL_DESC_A_SCALE_POINTER,\n &A_scale_inverse, sizeof(A_scale_inverse)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc,\n HIPBLASLT_MATMUL_DESC_B_SCALE_POINTER,\n &B_scale_inverse, sizeof(B_scale_inverse)));\n scaling_mode_a = HIPBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3;\n scaling_mode_b = HIPBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3;\n#else\n NVTE_ERROR(\"FP4 requires cuBLAS 12.8+, but compile-time cuBLAS version is \", CUBLAS_VERSION);\n#endif // CUBLAS_VERSION >= 120800\n } else if ((inputA->scaling_mode == NVTE_BLOCK_SCALING_1D ||\n inputA->scaling_mode == NVTE_BLOCK_SCALING_2D) &&\n (inputB->scaling_mode == NVTE_BLOCK_SCALING_1D ||\n inputB->scaling_mode == NVTE_BLOCK_SCALING_2D)) {\n#if CUBLAS_VERSION >= 120900\n NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120900,\n \"FP8 block scaling requires cuBLAS 12.9+, but run-time cuBLAS version is \",\n transformer_engine::cuda::cublas_version());\n\n // Check that matrix formats are valid\n NVTE_CHECK((!(inputA->scaling_mode == NVTE_BLOCK_SCALING_2D &&\n inputB->scaling_mode == NVTE_BLOCK_SCALING_2D)),\n \"Only 1D by 1D, 1D by 2D, and 2D by 1D block scaling GEMM is supported, \"\n \"but got 2D by 2D\");\n\n // Configure cuBLAS scales\n float *A_scale_inverse = reinterpret_cast(param.A_scale_inv);\n float *B_scale_inverse = reinterpret_cast(param.B_scale_inv);\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc,\n HIPBLASLT_MATMUL_DESC_A_SCALE_POINTER,\n &A_scale_inverse, sizeof(A_scale_inverse)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc,\n HIPBLASLT_MATMUL_DESC_B_SCALE_POINTER,\n &B_scale_inverse, sizeof(B_scale_inverse)));\n scaling_mode_a = inputA->scaling_mode == NVTE_BLOCK_SCALING_1D\n ? HIPBLASLT_MATMUL_MATRIX_SCALE_VEC128_32F\n : HIPBLASLT_MATMUL_MATRIX_SCALE_BLK128x128_32F;\n scaling_mode_b = inputB->scaling_mode == NVTE_BLOCK_SCALING_1D\n ? HIPBLASLT_MATMUL_MATRIX_SCALE_VEC128_32F\n : HIPBLASLT_MATMUL_MATRIX_SCALE_BLK128x128_32F;\n#else\n NVTE_ERROR(\"FP8 block scaling requires cuBLAS 12.9+, but compile-time cuBLAS version is \",\n CUBLAS_VERSION);\n#endif // CUBLAS_VERSION >= 120900\n } else {\n NVTE_ERROR(\"Not implemented scaling modes: \" + to_string(inputA->scaling_mode) + \" and \" +\n to_string(inputB->scaling_mode) + \".\");\n }\n\n#if CUBLAS_VERSION >= 120800\n if (transformer_engine::cuda::cublas_version() >= 120800) {\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc,\n HIPBLASLT_MATMUL_DESC_A_SCALE_MODE,\n &scaling_mode_a, sizeof(scaling_mode_a)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc,\n HIPBLASLT_MATMUL_DESC_B_SCALE_MODE,\n &scaling_mode_b, sizeof(scaling_mode_b)));\n }\n#endif // CUBLAS_VERSION >= 120800\n if (is_fp8_dtype(outputD->data.dtype)) {\n // Accumulation mode not supported for FP8 output\n C = nullptr;\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, HIPBLASLT_MATMUL_DESC_D_SCALE_POINTER, &D_scale, sizeof(D_scale)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, HIPBLASLT_MATMUL_DESC_AMAX_D_POINTER, &D_amax, sizeof(D_amax)));\n#if CUBLAS_VERSION >= 120800\n if (transformer_engine::cuda::cublas_version() >= 120800) {\n // NOTE: In all current cases where FP8 output is supported, the input is\n // scaled identically to the output.\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc,\n CUBLASLT_MATMUL_DESC_D_SCALE_MODE,\n &scaling_mode_a, sizeof(scaling_mode_a)));\n }\n#endif // CUBLAS_VERSION >= 120800\n // For FP8 output, cuBLAS requires C_type to match bias_type and\n // be FP16/BF16\n const hipDataType C_type = bias ? bias_type : HIP_R_16BF;\n NVTE_CHECK_CUBLAS(hipblasLtMatrixLayoutCreate(&Cdesc, C_type, m, n, ldd));\n } else {\n NVTE_CHECK_CUBLAS(hipblasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));\n }\n if (bias) {\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, HIPBLASLT_MATMUL_DESC_BIAS_DATA_TYPE, &bias_type, sizeof(bias_type)));\n }\n } else {\n NVTE_CHECK_CUBLAS(hipblasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));\n }\n\n if (bias && gelu) {\n if (grad) {\n epilogue = HIPBLASLT_EPILOGUE_DGELU_BGRAD;\n } else {\n epilogue = HIPBLASLT_EPILOGUE_GELU_AUX_BIAS;\n }\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, HIPBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, sizeof(bias_ptr)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc,\n HIPBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER,\n &pre_gelu_out, sizeof(pre_gelu_out)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, HIPBLASLT_MATMUL_DESC_EPILOGUE_AUX_LD, &ld_gelumat, sizeof(ld_gelumat)));\n const hipDataType aux_type = get_cuda_dtype(outputPreGelu->data.dtype);\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, HIPBLASLT_MATMUL_DESC_EPILOGUE_AUX_DATA_TYPE, &aux_type, sizeof(aux_type)));\n } else if (bias) {\n if (grad) {\n // grad output is always input B\n epilogue = HIPBLASLT_EPILOGUE_BGRADB;\n } else {\n epilogue = HIPBLASLT_EPILOGUE_BIAS;\n }\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, HIPBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, sizeof(bias_ptr)));\n } else if (gelu) {\n if (grad) {\n epilogue = HIPBLASLT_EPILOGUE_DGELU;\n } else {\n epilogue = HIPBLASLT_EPILOGUE_GELU_AUX;\n }\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc,\n HIPBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER,\n &pre_gelu_out, sizeof(pre_gelu_out)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, HIPBLASLT_MATMUL_DESC_EPILOGUE_AUX_LD, &ld_gelumat, sizeof(ld_gelumat)));\n const hipDataType aux_type = get_cuda_dtype(outputPreGelu->data.dtype);\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, HIPBLASLT_MATMUL_DESC_EPILOGUE_AUX_DATA_TYPE, &aux_type, sizeof(aux_type)));\n }\n\n if ((inputA->scaling_mode == NVTE_BLOCK_SCALING_1D) ||\n (inputA->scaling_mode == NVTE_BLOCK_SCALING_2D)) {\n NVTE_CHECK((epilogue == HIPBLASLT_EPILOGUE_DEFAULT || epilogue == HIPBLASLT_EPILOGUE_BIAS ||\n epilogue == HIPBLASLT_EPILOGUE_DGELU),\n \"Epilogue requested outside of the available and tested cuBLAS functionality for \"\n \"float8 block scaled GEMM\");\n }\n\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(operationDesc, HIPBLASLT_MATMUL_DESC_EPILOGUE,\n &epilogue, sizeof(epilogue)));\n\n if (counter != nullptr) {\n#if !(CUDA_VERSION >= 12020 && CUDA_VERSION < 13000)\n NVTE_ERROR(\"Atomic GEMM requires CUDA >=12.2.0 and <13.0.0, but compile-time CUDA version is \",\n CUDA_VERSION);\n#elif !(CUBLAS_VERSION >= 120205 && CUBLAS_VERSION < 130000)\n NVTE_ERROR(\n \"Atomic GEMM requires cuBLAS >=12.2.5 and <13.0.0, but compile-time cuBLAS version is \",\n CUBLAS_VERSION);\n#else\n NVTE_CHECK(transformer_engine::cuda::cudart_version() >= 12020 &&\n transformer_engine::cuda::cudart_version() < 13000,\n \"Atomic GEMM requires CUDA >=12.2.0 and <13.0.0, but run-time CUDA version is \",\n transformer_engine::cuda::cudart_version());\n NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120205 &&\n transformer_engine::cuda::cublas_version() < 130000,\n \"Atomic GEMM requires cuBLAS >=12.2.5 and <13.0.0, but run-time cuBLAS version is \",\n transformer_engine::cuda::cublas_version());\n if (m_split == 0) m_split = 1;\n if (n_split == 0) n_split = 1;\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_ATOMIC_SYNC_NUM_CHUNKS_D_ROWS, &m_split,\n sizeof(m_split)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_ATOMIC_SYNC_NUM_CHUNKS_D_COLS, &n_split,\n sizeof(n_split)));\n if (gemm_producer) {\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_ATOMIC_SYNC_OUT_COUNTERS_POINTER, &counter,\n sizeof(counter)));\n } else {\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescSetAttribute(\n operationDesc, CUBLASLT_MATMUL_DESC_ATOMIC_SYNC_IN_COUNTERS_POINTER, &counter,\n sizeof(counter)));\n }\n#endif\n }\n\n // align the workspace to 256 B\n const int required_alignment = 256;\n const auto original_workspace_alignment = _getAlignment(reinterpret_cast(workspace));\n uint8_t *aligned_workspace_ptr =\n reinterpret_cast(workspace) + required_alignment - original_workspace_alignment;\n workspaceSize = workspaceSize - required_alignment + original_workspace_alignment;\n const auto new_workspace_alignment =\n _getAlignment(reinterpret_cast(aligned_workspace_ptr));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulPreferenceCreate(&preference));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulPreferenceSetAttribute(\n preference, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize)));\n const auto A_alignment = _getAlignment(reinterpret_cast(param.A));\n const auto B_alignment = _getAlignment(reinterpret_cast(param.B));\n const auto C_alignment = _getAlignment(reinterpret_cast(C));\n const auto D_alignment = _getAlignment(reinterpret_cast(D));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulPreferenceSetAttribute(\n preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_A_BYTES, &A_alignment, sizeof(A_alignment)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulPreferenceSetAttribute(\n preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_B_BYTES, &B_alignment, sizeof(B_alignment)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulPreferenceSetAttribute(\n preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_C_BYTES, &C_alignment, sizeof(C_alignment)));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulPreferenceSetAttribute(\n preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_D_BYTES, &D_alignment, sizeof(D_alignment)));\n NVTE_CHECK(new_workspace_alignment % 256 == 0,\n \"cuBLAS workspace pointer must be aligned to 256 bytes, got \",\n new_workspace_alignment);\n\n const auto status =\n hipblasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Cdesc, Ddesc, preference,\n 1, &heuristicResult, &returnedResults);\n NVTE_CHECK(status != HIPBLAS_STATUS_NOT_SUPPORTED,\n \"Unable to find suitable cuBLAS GEMM algorithm\");\n NVTE_CHECK_CUBLAS(status);\n if (returnedResults == 0) NVTE_ERROR(\"Unable to find any suitable algorithms\");\n\n // D = alpha * (A * B) + beta * C\n NVTE_CHECK_CUBLAS(hipblasLtMatmul(handle, operationDesc, alpha, /* alpha */\n param.A, /* A */\n Adesc, param.B, /* B */\n Bdesc, beta, /* beta */\n C, /* C */\n Cdesc, D, /* D */\n Ddesc, &heuristicResult.algo, /* algo */\n aligned_workspace_ptr, /* workspace */\n workspaceSize, stream)); /* stream */\n\n // Update FP8 scale-inv in output tensor\n // Note: This is a WAR for the case when we have fp8 output but D->scale_inv is not allocated.\n // TODO: Changing gemm interface so that D->scale_inv is allocated and the scale_inv can be\n // calculated here.\n if (is_fp8_dtype(outputD->data.dtype) && outputD->scale_inv.dptr) {\n update_tensor_scale_inv(outputD, stream);\n }\n\n NVTE_CHECK_CUBLAS(hipblasLtMatmulPreferenceDestroy(preference));\n NVTE_CHECK_CUBLAS(hipblasLtMatrixLayoutDestroy(Ddesc));\n NVTE_CHECK_CUBLAS(hipblasLtMatrixLayoutDestroy(Cdesc));\n NVTE_CHECK_CUBLAS(hipblasLtMatrixLayoutDestroy(Bdesc));\n NVTE_CHECK_CUBLAS(hipblasLtMatrixLayoutDestroy(Adesc));\n NVTE_CHECK_CUBLAS(hipblasLtMatmulDescDestroy(operationDesc));\n}\n\n} // namespace transformer_engine\n\nvoid nvte_cublas_gemm(const NVTETensor A, const NVTETensor B, NVTETensor D, const NVTETensor bias,\n NVTETensor pre_gelu_out, bool transa, bool transb, bool grad,\n NVTETensor workspace, bool accumulate, bool use_split_accumulator,\n int math_sm_count, hipStream_t stream) {\n NVTE_API_CALL(nvte_cublas_gemm);\n using namespace transformer_engine;\n\n // Tensors\n const Tensor *inputA = convertNVTETensorCheck(A);\n const Tensor *inputB = convertNVTETensorCheck(B);\n Tensor *outputD = convertNVTETensorCheck(D);\n const Tensor *biasTensor = convertNVTETensor(bias);\n Tensor *outputGelu = convertNVTETensor(pre_gelu_out);\n Tensor *wspace = convertNVTETensor(workspace);\n\n // Scales\n const float alpha = 1;\n const float beta = accumulate ? 1 : 0;\n\n // Check for NVFP4\n // TODO Remove once alpha scale logic is moved into cublas_gemm function\n if (is_nvfp_scaling(inputA->scaling_mode) || is_nvfp_scaling(inputB->scaling_mode)) {\n NVTE_ERROR(\"nvte_cublas_gemm does not support NVFP4 data. Use nvte_cublas_gemm_v2 instead.\");\n }\n\n // Launch GEMM\n cublas_gemm(inputA, inputB, outputD, biasTensor, outputGelu, (transa) ? HIPBLAS_OP_T : HIPBLAS_OP_N,\n (transb) ? HIPBLAS_OP_T : HIPBLAS_OP_N, grad, wspace->data.dptr, wspace->data.shape[0],\n &alpha, &beta, use_split_accumulator, math_sm_count, 0, 0, false, nullptr, stream);\n}\n\nvoid nvte_cublas_gemm_v2(int transa, int transb, const float *alpha, const NVTETensor A,\n const NVTETensor B, const float *beta, const NVTETensor C, NVTETensor D,\n NVTETensor workspace, NVTEMatmulConfig config, hipStream_t stream) {\n NVTE_API_CALL(nvte_cublas_gemm_v2);\n using namespace transformer_engine;\n\n // Data tensors\n const Tensor *A_tensor = convertNVTETensorCheck(A);\n const Tensor *B_tensor = convertNVTETensorCheck(B);\n const Tensor *C_tensor = convertNVTETensorCheck(C);\n Tensor *D_tensor = convertNVTETensorCheck(D);\n NVTE_CHECK(C_tensor == D_tensor,\n \"Currently nvte_cublas_gemm_v2 does not support different C and D tensors.\");\n\n // Workspace\n void *workspace_ptr = nullptr;\n size_t workspace_size = 0;\n Tensor *workspace_tensor = convertNVTETensor(workspace);\n if (workspace_tensor != nullptr) {\n workspace_ptr = workspace_tensor->data.dptr;\n workspace_size =\n get_buffer_size_bytes(workspace_tensor->data.numel(), workspace_tensor->data.dtype);\n }\n\n // Additional config\n MatmulConfig config_;\n if (config != nullptr) {\n config_ = *reinterpret_cast(config);\n }\n\n // Configure GEMM epilogue\n const bool with_grad_epilogue = (config_.dbias_tensor != nullptr || config_.with_dgelu_epilogue);\n if (with_grad_epilogue) {\n NVTE_CHECK(config_.bias_tensor == nullptr && !config_.with_gelu_epilogue,\n \"Invalid epilogue (bias=\", config_.bias_tensor != nullptr,\n \", dbias=\", config_.dbias_tensor != nullptr, \", gelu=\", config_.with_gelu_epilogue,\n \", dgelu=\", config_.with_dgelu_epilogue, \").\");\n }\n Tensor dummy_tensor;\n Tensor *epilogue_bias_tensor = &dummy_tensor;\n if (!with_grad_epilogue && config_.bias_tensor != nullptr) {\n epilogue_bias_tensor = convertNVTETensorCheck(config_.bias_tensor);\n } else if (with_grad_epilogue && config_.dbias_tensor != nullptr) {\n epilogue_bias_tensor = convertNVTETensorCheck(config_.dbias_tensor);\n }\n Tensor *epilogue_aux_tensor = &dummy_tensor;\n if (config_.with_gelu_epilogue || config_.with_dgelu_epilogue) {\n NVTE_CHECK(config_.epilogue_aux_tensor != nullptr,\n \"Requested epilogue (bias=\", config_.bias_tensor != nullptr,\n \", dbias=\", config_.dbias_tensor != nullptr, \", gelu=\", config_.with_gelu_epilogue,\n \", dgelu=\", config_.with_dgelu_epilogue, \") without providing aux tensor.\");\n epilogue_aux_tensor = convertNVTETensor(config_.epilogue_aux_tensor);\n }\n\n // Launch GEMM\n cublas_gemm(A_tensor, B_tensor, D_tensor, epilogue_bias_tensor, epilogue_aux_tensor,\n transa ? HIPBLAS_OP_T : HIPBLAS_OP_N, transb ? HIPBLAS_OP_T : HIPBLAS_OP_N,\n with_grad_epilogue, workspace_ptr, workspace_size, alpha, beta,\n config_.use_split_accumulator, config_.sm_count, 0, 0, false, nullptr, stream);\n}\n\nvoid nvte_cublas_gemm_scaled(const NVTETensor A, const NVTETensor B, NVTETensor D,\n const NVTETensor bias, NVTETensor pre_gelu_out, bool transa,\n bool transb, bool grad, NVTETensor workspace, float alpha, float beta,\n bool use_split_accumulator, int math_sm_count, hipStream_t stream) {\n NVTE_API_CALL(nvte_cublas_gemm);\n using namespace transformer_engine;\n\n // Tensors\n const Tensor *inputA = convertNVTETensorCheck(A);\n const Tensor *inputB = convertNVTETensorCheck(B);\n Tensor *outputD = convertNVTETensorCheck(D);\n const Tensor *biasTensor = convertNVTETensor(bias);\n Tensor *outputGelu = convertNVTETensor(pre_gelu_out);\n Tensor *wspace = convertNVTETensor(workspace);\n\n // Check for NVFP4\n // TODO Remove once alpha scale logic is moved into cublas_gemm function\n if (is_nvfp_scaling(inputA->scaling_mode) || is_nvfp_scaling(inputB->scaling_mode)) {\n NVTE_ERROR(\"nvte_cublas_gemm does not support NVFP4 data. Use nvte_cublas_gemm_v2 instead.\");\n }\n\n // Launch GEMM\n cublas_gemm(inputA, inputB, outputD, biasTensor, outputGelu, (transa) ? HIPBLAS_OP_T : HIPBLAS_OP_N,\n (transb) ? HIPBLAS_OP_T : HIPBLAS_OP_N, grad, wspace->data.dptr, wspace->data.shape[0],\n &alpha, &beta, use_split_accumulator, math_sm_count, 0, 0, false, nullptr, stream);\n}\n\nvoid nvte_cublas_atomic_gemm(const NVTETensor A, const NVTETensor B, NVTETensor D,\n const NVTETensor bias, NVTETensor pre_gelu_out, bool transa,\n bool transb, bool grad, NVTETensor workspace, bool accumulate,\n bool use_split_accumulator, int math_sm_count, int m_split,\n int n_split, bool gemm_producer, const NVTETensor counter,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_cublas_atomic_gemm);\n using namespace transformer_engine;\n#if !(CUDA_VERSION >= 12020 && CUDA_VERSION < 13000)\n NVTE_ERROR(\"Atomic GEMM requires CUDA >=12.2.0 and <13.0.0, but compile-time CUDA version is \",\n CUDA_VERSION);\n#elif !(CUBLAS_VERSION >= 120205 && CUBLAS_VERSION < 130000)\n NVTE_ERROR(\n \"Atomic GEMM requires cuBLAS >=12.2.5 and <13.0.0, but compile-time cuBLAS version is \",\n CUBLAS_VERSION);\n#else\n NVTE_CHECK(\n transformer_engine::cuda::cudart_version() >= 12020 &&\n transformer_engine::cuda::cudart_version() < 13000,\n \"Atomic GEMM requires CUDA version >=12.2.0 and <13.0.0, but run-time CUDA version is \",\n transformer_engine::cuda::cudart_version());\n NVTE_CHECK(\n transformer_engine::cuda::cublas_version() >= 120205 &&\n transformer_engine::cuda::cublas_version() < 130000,\n \"Atomic GEMM requires cuBLAS version >=12.2.5 and <13.0.0, but run-time cuBLAS version is \",\n transformer_engine::cuda::cublas_version());\n\n const Tensor *inputA = convertNVTETensorCheck(A);\n const Tensor *inputB = convertNVTETensorCheck(B);\n Tensor *outputD = convertNVTETensor(D);\n const Tensor *biasTensor = convertNVTETensor(bias);\n Tensor *outputGelu = convertNVTETensor(pre_gelu_out);\n const Tensor *inputCounter = convertNVTETensor(counter);\n Tensor *wspace = convertNVTETensor(workspace);\n\n const void *alpha_ptr = GetScalarOne();\n const void *beta_ptr = accumulate ? GetScalarOne() : GetScalarZero();\n\n NVTE_CHECK(is_delayed_tensor_scaling(inputA->scaling_mode) &&\n is_delayed_tensor_scaling(inputB->scaling_mode),\n \"Atomic GEMM only supports delayed scaling.\");\n cublas_gemm(inputA, inputB, outputD, biasTensor, outputGelu, (transa) ? HIPBLAS_OP_T : HIPBLAS_OP_N,\n (transb) ? HIPBLAS_OP_T : HIPBLAS_OP_N, grad, wspace->data.dptr, wspace->data.shape[0],\n alpha_ptr, beta_ptr, use_split_accumulator, math_sm_count, m_split, n_split,\n gemm_producer, inputCounter, stream);\n#endif\n}\n\nvoid multi_stream_cublas_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor *D,\n const NVTETensor *bias, NVTETensor *pre_gelu_out, const int num_gemms,\n bool transa, bool transb, bool grad, NVTETensor *workspace,\n bool accumulate, bool use_split_accumulator, int math_sm_count,\n hipStream_t stream) {\n using namespace transformer_engine;\n\n int num_streams = nvte_get_num_compute_streams();\n\n int num_stream_used = std::min(num_streams, num_gemms);\n // wait for current stream to finish\n NVTE_CHECK_CUDA(hipEventRecord(detail::get_compute_stream_event(0), stream));\n for (int s = 0; s < num_stream_used; s++) {\n NVTE_CHECK_CUDA(\n hipStreamWaitEvent(detail::get_compute_stream(s), detail::get_compute_stream_event(0)));\n }\n\n for (int i = 0; i < num_gemms; i++) {\n // Check whether GELU or dGELU epilogue is requested\n Tensor *pre_gelu_tensor = convertNVTETensor(pre_gelu_out[i]);\n bool with_gelu_dgelu_epilogue =\n (pre_gelu_tensor != nullptr && pre_gelu_tensor->data.dptr != nullptr);\n\n // Construct config\n MatmulConfig config;\n if (grad) {\n config.dbias_tensor = bias[i];\n config.with_dgelu_epilogue = with_gelu_dgelu_epilogue;\n } else {\n config.bias_tensor = bias[i];\n config.with_gelu_epilogue = with_gelu_dgelu_epilogue;\n }\n config.epilogue_aux_tensor = pre_gelu_out[i];\n config.use_split_accumulator = use_split_accumulator;\n config.sm_count = math_sm_count;\n\n // Launch GEMM\n const float alpha = 1.f;\n const float beta = accumulate ? 1.f : 0.f;\n nvte_cublas_gemm_v2(transa, transb, &alpha, A[i], B[i], &beta, D[i], D[i],\n workspace[i % num_streams], &config,\n detail::get_compute_stream(i % num_streams));\n }\n\n // record events on compute streams\n for (int s = 0; s < num_stream_used; s++) {\n NVTE_CHECK_CUDA(\n hipEventRecord(detail::get_compute_stream_event(s), detail::get_compute_stream(s)));\n }\n // wait for all compute streams to finish\n for (int s = 0; s < num_stream_used; s++) {\n NVTE_CHECK_CUDA(hipStreamWaitEvent(stream, detail::get_compute_stream_event(s)));\n }\n}\n\nvoid nvte_multi_stream_cublas_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor *D,\n const NVTETensor *bias, NVTETensor *pre_gelu_out,\n const int num_gemms, bool transa, bool transb, bool grad,\n NVTETensor *workspace, bool accumulate,\n bool use_split_accumulator, int math_sm_count,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_multi_stream_cublas_gemm);\n using namespace transformer_engine;\n\n // Deprecation warning\n NVTE_WARN(\n \"nvte_multi_stream_cublas_gemm is deprecated and will be removed in a future release. \"\n \"Please migrate to nvte_multi_tensor_gemm (with CUTLASS Grouped GEMM support when \"\n \"applicable).\");\n\n multi_stream_cublas_gemm(A, B, D, bias, pre_gelu_out, num_gemms, transa, transb, grad, workspace,\n accumulate, use_split_accumulator, math_sm_count, stream);\n}\n\nnamespace transformer_engine {\n\nusing cublasHandleManager = detail::HandleManager;\n\nvoid nvte_cublas_handle_init() { auto _ = cublasHandleManager::Instance().GetHandle(); }\n\n} // namespace transformer_engine\n\nvoid nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor *D,\n const NVTETensor *bias, NVTETensor *pre_gelu_out, const int num_gemms,\n bool transa, bool transb, bool grad, NVTETensor *workspace,\n bool accumulate, bool use_split_accumulator, int math_sm_count,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_gemm);\n\n const int current_device = transformer_engine::cuda::current_device();\n const bool is_hopper = (transformer_engine::cuda::sm_arch(current_device) == 90);\n const bool use_cutlass = transformer_engine::getenv(\"NVTE_USE_CUTLASS_GROUPED_GEMM\", false);\n const bool warn_fallback =\n transformer_engine::getenv(\"NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK\", false);\n\n auto cublas_path = [&]() {\n multi_stream_cublas_gemm(A, B, D, bias, pre_gelu_out, num_gemms, transa, transb, grad,\n workspace, accumulate, use_split_accumulator, math_sm_count, stream);\n };\n\n // Currently only support cutlass group gemm on Hopper Arch\n if (!(is_hopper && use_cutlass)) {\n cublas_path();\n return;\n }\n\n auto is_empty_arr = [&](const NVTETensor *p) -> bool {\n if (p == nullptr) return true;\n for (int i = 0; i < num_gemms; ++i) {\n if (transformer_engine::convertNVTETensor(p[i])->has_data()) return false;\n }\n return true;\n };\n\n auto all_groups_uniform_k128 = [&](const NVTETensor *p, bool trans) -> bool {\n int64_t ref_k = -1;\n for (size_t i = 0; i < num_gemms; i++) {\n const auto tensor = transformer_engine::convertNVTETensorCheck(p[i]);\n const int k = trans ? tensor->data.shape[0] : tensor->data.shape[1];\n\n if ((k & 127) != 0) return false;\n\n if (ref_k < 0)\n ref_k = k;\n else if (k != ref_k)\n return false;\n }\n\n return true;\n };\n\n auto is_supported_dtype = [&]() -> bool {\n auto *inputA = transformer_engine::convertNVTETensorCheck(A[0]);\n auto *inputB = transformer_engine::convertNVTETensorCheck(B[0]);\n auto *OutputD = transformer_engine::convertNVTETensorCheck(D[0]);\n auto A_type = get_cuda_dtype(inputA->data.dtype);\n auto B_type = get_cuda_dtype(inputB->data.dtype);\n auto D_type = get_cuda_dtype(OutputD->data.dtype);\n\n return (A_type == B_type) && (A_type == D_type) &&\n ((A_type == HIP_R_16BF) || (A_type == HIP_R_16F));\n };\n\n // CUTLASS Grouped GEMM fast path (SM90/TMA)\n // Conditions:\n // - No fused epilogue: both bias and pre_gelu_out are empty.\n // - Supported dtypes only: FP16/BF16 (FP32 accumulate).\n // - Uniform K across groups and K % 128 == 0.\n // - use_split_accumulator is ignored for FP16/BF16.\n // - grad is irrelevant when bias/pre_gelu_out are empty.\n //\n // Otherwise, fall back to cuBLAS.\n if (is_empty_arr(bias) && is_empty_arr(pre_gelu_out) && is_supported_dtype() &&\n all_groups_uniform_k128(B, transb)) {\n cutlass_grouped_gemm(A, B, D, num_gemms, transa, transb, grad, workspace, accumulate,\n current_device, math_sm_count, stream);\n } else {\n if (warn_fallback) {\n NVTE_WARN(\"Fallback to cuBLAS grouped GEMM.\");\n }\n cublas_path();\n }\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/gemm/cublaslt_gemm.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/cuda_runtime.h\"\n#include \"../util/logging.h\"\n#include \"../util/rtc.h\"\n#include \"../util/string.h\"\n\nnamespace transformer_engine {\n\nnamespace {\n\n// String with RTC kernel implementation\n#include \"string_code_transpose_rtc_swap_first_dims_cu.h\"\n\n// Hard-coded kernel parameters\nconstexpr size_t block_size = 128;\n\n/* Performance heuristics for optimized kernel parameters */\nstruct KernelConfig {\n /* Vector load/store size */\n size_t vector_size;\n\n /* Whether config is valid */\n bool valid = false;\n /* Number of CUDA blocks */\n size_t num_blocks = 0;\n /* Whether each thread needs to make exactly one load/store */\n bool single_load_store = true;\n\n /* Number of active SMs */\n size_t active_sm_count = 0;\n /* Used bytes per L1 cache load */\n size_t bytes_per_load = 0;\n /* Used bytes per L1 cache store */\n size_t bytes_per_store = 0;\n\n KernelConfig(size_t dim0, size_t dim1, size_t dim2, size_t sm_count, size_t vector_size_)\n : vector_size{vector_size_} {\n // Check that tiles are correctly aligned\n if (dim2 % vector_size_ != 0) {\n return;\n }\n valid = true;\n\n // Number of CUDA blocks\n num_blocks = DIVUP(dim0 * dim1 * dim2 / vector_size, block_size);\n if (num_blocks > 2147483647ull) {\n // Maximum number of CUDA blocks\n single_load_store = false;\n num_blocks = 2147483647ull;\n } else if (num_blocks * block_size != dim0 * dim1 * dim2 / vector_size) {\n single_load_store = false;\n }\n\n // SM occupancy\n constexpr size_t warp_size = 32;\n constexpr size_t warps_per_sm = 16; // Rough estimate for saturated SMs\n active_sm_count = std::min(DIVUP(num_blocks * block_size / warp_size, warps_per_sm), sm_count);\n\n // L1 cache efficiency\n constexpr size_t cache_line_size = 128;\n bytes_per_store = std::min(cache_line_size, warp_size * vector_size); // Contiguous writes\n bytes_per_load = bytes_per_store;\n if (dim2 % (vector_size * warp_size) != 0) {\n // Some warps are reading from two non-contiguous regions\n bytes_per_load /= 2;\n }\n }\n\n /* Compare by estimated cost */\n bool operator<(const KernelConfig &other) const {\n if (this->valid && other.valid) {\n // cost ~ (1/bytes_per_load + 1/bytes_per_store) / active_sms\n // Note: Integer arithmetic ensures stable ordering\n const auto &l1 = this->bytes_per_load;\n const auto &s1 = this->bytes_per_store;\n const auto &p1 = this->active_sm_count;\n const auto &l2 = other.bytes_per_load;\n const auto &s2 = other.bytes_per_store;\n const auto &p2 = other.active_sm_count;\n const auto scale = l1 * s1 * p1 * l2 * s2 * p2;\n const auto cost1 = (scale / l1 + scale / s1) / p1;\n const auto cost2 = (scale / l2 + scale / s2) / p2;\n return cost1 < cost2;\n } else {\n return this->valid && !other.valid;\n }\n }\n};\n\ntemplate \n__global__ void __launch_bounds__(block_size)\n swap_first_dims_untuned_kernel(const Type *__restrict__ input, Type *__restrict__ output,\n const size_t dim0, const size_t dim1, const size_t dim2) {\n const size_t gid = threadIdx.x + blockIdx.x * block_size;\n const size_t nthreads = gridDim.x * block_size;\n for (size_t idx = gid; idx < dim0 * dim1 * dim2; idx += nthreads) {\n const auto idx2 = idx % dim2;\n const auto idx1 = (idx / dim2) % dim1;\n const auto idx0 = (idx / dim2) / dim1;\n const auto in_offset = idx1 * dim0 * dim2 + idx0 * dim2 + idx2;\n output[idx] = input[in_offset];\n }\n}\n\n} // namespace\n\nvoid swap_first_dims(const Tensor &input, Tensor &output, cudaStream_t stream) {\n // Check tensors\n NVTE_CHECK(input.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor must be simple tensor, but scaling mode is \",\n to_string(input.scaling_mode), \".\");\n NVTE_CHECK(output.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Output tensor must be simple tensor, but scaling mode is \",\n to_string(output.scaling_mode), \".\");\n NVTE_CHECK(input.dtype() == output.dtype(), \"Input tensor (dtype=\", to_string(input.dtype()),\n \") and output tensor (dtype=\", to_string(output.dtype()), \") do not match.\");\n NVTE_CHECK(input.data.dptr != nullptr, \"Input is not allocated.\");\n NVTE_CHECK(output.data.dptr != nullptr, \"Output is not allocated.\");\n\n // Check tensor dimensions\n const auto input_shape = input.shape();\n const auto output_shape = output.shape();\n NVTE_CHECK(input_shape.size() >= 2, \"Invalid input tensor dimensions (shape=\", input_shape, \").\");\n NVTE_CHECK(output_shape.size() == input_shape.size(), \"Input tensor (shape=\", input_shape,\n \") and output tensor (shape=\", output_shape, \") do not match.\");\n NVTE_CHECK(input_shape[0] == output_shape[1], \"Input tensor (shape=\", input_shape,\n \") and output tensor (shape=\", output_shape, \") do not match.\");\n NVTE_CHECK(input_shape[1] == output_shape[0], \"Input tensor (shape=\", input_shape,\n \") and output tensor (shape=\", output_shape, \") do not match.\");\n for (size_t i = 2; i < input_shape.size(); ++i) {\n NVTE_CHECK(input_shape[i] == output_shape[i], \"Input tensor (shape=\", input_shape,\n \") and output tensor (shape=\", output_shape, \") do not match.\");\n }\n\n // Reinterpret tensors as 3D tensors of bytes\n const size_t dim0 = output_shape[0];\n const size_t dim1 = output_shape[1];\n size_t dim2 = 1;\n for (size_t i = 2; i < output_shape.size(); ++i) {\n dim2 *= output_shape[i];\n }\n dim2 = get_buffer_size_bytes(dim2, output.dtype());\n\n // Choose kernel config with performance heuristics\n const size_t sm_count = static_cast(cuda::sm_count());\n KernelConfig config(dim0, dim1, dim2, sm_count, 1);\n if (rtc::is_enabled()) {\n auto try_config = [&](size_t vector_size) {\n KernelConfig new_config(dim0, dim1, dim2, sm_count, vector_size);\n if (new_config < config) {\n config = new_config;\n }\n };\n try_config(16);\n try_config(8);\n try_config(4);\n try_config(2);\n }\n const size_t vector_size = config.vector_size;\n\n // Launch kernel\n if (vector_size == 1) {\n // General kernel\n swap_first_dims_untuned_kernel<<>>(\n static_cast(input.data.dptr), static_cast(output.data.dptr),\n dim0, dim1, dim2);\n NVTE_CHECK_CUDA(cudaGetLastError());\n } else {\n // Compile NVRTC kernel if needed\n auto &rtc_manager = rtc::KernelManager::instance();\n const std::string kernel_label =\n concat_strings(\"swap_first_dims,vector_size=\", vector_size,\n \",single_load_store=\", config.single_load_store);\n if (!rtc_manager.is_compiled(kernel_label)) {\n std::string code = string_code_transpose_rtc_swap_first_dims_cu;\n code = regex_replace(code, \"__VECTOR_SIZE__\", vector_size);\n code = regex_replace(code, \"__BLOCK_SIZE__\", block_size);\n code =\n regex_replace(code, \"__SINGLE_LOAD_STORE__\", static_cast(config.single_load_store));\n rtc_manager.compile(kernel_label, \"swap_first_dims_kernel\", code,\n \"transformer_engine/common/transpose/rtc/swap_first_dims.cu\");\n }\n\n // Launch NVRTC kernel\n rtc_manager.launch(kernel_label, config.num_blocks, block_size, 0, stream, input.data.dptr,\n output.data.dptr, dim0, dim1, dim2 / vector_size);\n }\n}\n\n} // namespace transformer_engine\n\nvoid nvte_swap_first_dims(const NVTETensor input, NVTETensor output, cudaStream_t stream) {\n NVTE_API_CALL(nvte_swap_first_dims);\n using namespace transformer_engine;\n swap_first_dims(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output), stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/hip/hip_runtime.h\"\n#include \"../util/logging.h\"\n#include \"../util/rtc.h\"\n#include \"../util/string.h\"\n\nnamespace transformer_engine {\n\nnamespace {\n\n// String with RTC kernel implementation\n#include \"string_code_transpose_rtc_swap_first_dims_cu.h\"\n\n// Hard-coded kernel parameters\nconstexpr size_t block_size = 128;\n\n/* Performance heuristics for optimized kernel parameters */\nstruct KernelConfig {\n /* Vector load/store size */\n size_t vector_size;\n\n /* Whether config is valid */\n bool valid = false;\n /* Number of CUDA blocks */\n size_t num_blocks = 0;\n /* Whether each thread needs to make exactly one load/store */\n bool single_load_store = true;\n\n /* Number of active SMs */\n size_t active_sm_count = 0;\n /* Used bytes per L1 cache load */\n size_t bytes_per_load = 0;\n /* Used bytes per L1 cache store */\n size_t bytes_per_store = 0;\n\n KernelConfig(size_t dim0, size_t dim1, size_t dim2, size_t sm_count, size_t vector_size_)\n : vector_size{vector_size_} {\n // Check that tiles are correctly aligned\n if (dim2 % vector_size_ != 0) {\n return;\n }\n valid = true;\n\n // Number of CUDA blocks\n num_blocks = DIVUP(dim0 * dim1 * dim2 / vector_size, block_size);\n if (num_blocks > 2147483647ull) {\n // Maximum number of CUDA blocks\n single_load_store = false;\n num_blocks = 2147483647ull;\n } else if (num_blocks * block_size != dim0 * dim1 * dim2 / vector_size) {\n single_load_store = false;\n }\n\n // SM occupancy\n constexpr size_t warp_size = 32;\n constexpr size_t warps_per_sm = 16; // Rough estimate for saturated SMs\n active_sm_count = std::min(DIVUP(num_blocks * block_size / warp_size, warps_per_sm), sm_count);\n\n // L1 cache efficiency\n constexpr size_t cache_line_size = 128;\n bytes_per_store = std::min(cache_line_size, warp_size * vector_size); // Contiguous writes\n bytes_per_load = bytes_per_store;\n if (dim2 % (vector_size * warp_size) != 0) {\n // Some warps are reading from two non-contiguous regions\n bytes_per_load /= 2;\n }\n }\n\n /* Compare by estimated cost */\n bool operator<(const KernelConfig &other) const {\n if (this->valid && other.valid) {\n // cost ~ (1/bytes_per_load + 1/bytes_per_store) / active_sms\n // Note: Integer arithmetic ensures stable ordering\n const auto &l1 = this->bytes_per_load;\n const auto &s1 = this->bytes_per_store;\n const auto &p1 = this->active_sm_count;\n const auto &l2 = other.bytes_per_load;\n const auto &s2 = other.bytes_per_store;\n const auto &p2 = other.active_sm_count;\n const auto scale = l1 * s1 * p1 * l2 * s2 * p2;\n const auto cost1 = (scale / l1 + scale / s1) / p1;\n const auto cost2 = (scale / l2 + scale / s2) / p2;\n return cost1 < cost2;\n } else {\n return this->valid && !other.valid;\n }\n }\n};\n\ntemplate \n__global__ void __launch_bounds__(block_size)\n swap_first_dims_untuned_kernel(const Type *__restrict__ input, Type *__restrict__ output,\n const size_t dim0, const size_t dim1, const size_t dim2) {\n const size_t gid = threadIdx.x + blockIdx.x * block_size;\n const size_t nthreads = gridDim.x * block_size;\n for (size_t idx = gid; idx < dim0 * dim1 * dim2; idx += nthreads) {\n const auto idx2 = idx % dim2;\n const auto idx1 = (idx / dim2) % dim1;\n const auto idx0 = (idx / dim2) / dim1;\n const auto in_offset = idx1 * dim0 * dim2 + idx0 * dim2 + idx2;\n output[idx] = input[in_offset];\n }\n}\n\n} // namespace\n\nvoid swap_first_dims(const Tensor &input, Tensor &output, hipStream_t stream) {\n // Check tensors\n NVTE_CHECK(input.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor must be simple tensor, but scaling mode is \",\n to_string(input.scaling_mode), \".\");\n NVTE_CHECK(output.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Output tensor must be simple tensor, but scaling mode is \",\n to_string(output.scaling_mode), \".\");\n NVTE_CHECK(input.dtype() == output.dtype(), \"Input tensor (dtype=\", to_string(input.dtype()),\n \") and output tensor (dtype=\", to_string(output.dtype()), \") do not match.\");\n NVTE_CHECK(input.data.dptr != nullptr, \"Input is not allocated.\");\n NVTE_CHECK(output.data.dptr != nullptr, \"Output is not allocated.\");\n\n // Check tensor dimensions\n const auto input_shape = input.shape();\n const auto output_shape = output.shape();\n NVTE_CHECK(input_shape.size() >= 2, \"Invalid input tensor dimensions (shape=\", input_shape, \").\");\n NVTE_CHECK(output_shape.size() == input_shape.size(), \"Input tensor (shape=\", input_shape,\n \") and output tensor (shape=\", output_shape, \") do not match.\");\n NVTE_CHECK(input_shape[0] == output_shape[1], \"Input tensor (shape=\", input_shape,\n \") and output tensor (shape=\", output_shape, \") do not match.\");\n NVTE_CHECK(input_shape[1] == output_shape[0], \"Input tensor (shape=\", input_shape,\n \") and output tensor (shape=\", output_shape, \") do not match.\");\n for (size_t i = 2; i < input_shape.size(); ++i) {\n NVTE_CHECK(input_shape[i] == output_shape[i], \"Input tensor (shape=\", input_shape,\n \") and output tensor (shape=\", output_shape, \") do not match.\");\n }\n\n // Reinterpret tensors as 3D tensors of bytes\n const size_t dim0 = output_shape[0];\n const size_t dim1 = output_shape[1];\n size_t dim2 = 1;\n for (size_t i = 2; i < output_shape.size(); ++i) {\n dim2 *= output_shape[i];\n }\n dim2 = get_buffer_size_bytes(dim2, output.dtype());\n\n // Choose kernel config with performance heuristics\n const size_t sm_count = static_cast(cuda::sm_count());\n KernelConfig config(dim0, dim1, dim2, sm_count, 1);\n if (rtc::is_enabled()) {\n auto try_config = [&](size_t vector_size) {\n KernelConfig new_config(dim0, dim1, dim2, sm_count, vector_size);\n if (new_config < config) {\n config = new_config;\n }\n };\n try_config(16);\n try_config(8);\n try_config(4);\n try_config(2);\n }\n const size_t vector_size = config.vector_size;\n\n // Launch kernel\n if (vector_size == 1) {\n // General kernel\n swap_first_dims_untuned_kernel<<>>(\n static_cast(input.data.dptr), static_cast(output.data.dptr),\n dim0, dim1, dim2);\n NVTE_CHECK_CUDA(hipGetLastError());\n } else {\n // Compile NVRTC kernel if needed\n auto &rtc_manager = rtc::KernelManager::instance();\n const std::string kernel_label =\n concat_strings(\"swap_first_dims,vector_size=\", vector_size,\n \",single_load_store=\", config.single_load_store);\n if (!rtc_manager.is_compiled(kernel_label)) {\n std::string code = string_code_transpose_rtc_swap_first_dims_cu;\n code = regex_replace(code, \"__VECTOR_SIZE__\", vector_size);\n code = regex_replace(code, \"__BLOCK_SIZE__\", block_size);\n code =\n regex_replace(code, \"__SINGLE_LOAD_STORE__\", static_cast(config.single_load_store));\n rtc_manager.compile(kernel_label, \"swap_first_dims_kernel\", code,\n \"transformer_engine/common/transpose/rtc/swap_first_dims.cu\");\n }\n\n // Launch NVRTC kernel\n rtc_manager.launch(kernel_label, config.num_blocks, block_size, 0, stream, input.data.dptr,\n output.data.dptr, dim0, dim1, dim2 / vector_size);\n }\n}\n\n} // namespace transformer_engine\n\nvoid nvte_swap_first_dims(const NVTETensor input, NVTETensor output, hipStream_t stream) {\n NVTE_API_CALL(nvte_swap_first_dims);\n using namespace transformer_engine;\n swap_first_dims(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output), stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/transpose/swap_first_dims.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \n\n#include \"../common.h\"\n#include \"../util/rtc.h\"\n#include \"../util/string.h\"\n#include \"../utils.cuh\"\n#include \"./transpose.h\"\n\nnamespace transformer_engine {\nnamespace detail {\n\nnamespace {\n\n// String with RTC kernel implementation\n#include \"string_code_transpose_rtc_transpose_cu.h\"\n\n// Hard-coded kernel parameters\nconstexpr size_t warps_per_tile = 4;\nconstexpr size_t block_size = THREADS_PER_WARP * warps_per_tile;\n\n/* Performance heuristics for optimized kernel parameters */\nstruct KernelConfig {\n /** Vector load size */\n size_t load_size;\n /** Vector store size */\n size_t store_size;\n\n /* Whether config is valid */\n bool valid = false;\n /* Number of CUDA blocks */\n size_t num_blocks = 0;\n\n /* Number of active SMs */\n size_t active_sm_count = 0;\n /* Elements per L1 cache load */\n size_t elements_per_load = 0;\n /* Elements per L1 cache store */\n size_t elements_per_store = 0;\n\n KernelConfig(size_t row_length, size_t num_rows, size_t type_size, size_t load_size_,\n size_t store_size_, size_t sm_count)\n : load_size{load_size_}, store_size{store_size_} {\n // Check that tiles are correctly aligned\n constexpr size_t cache_line_size = 128;\n if (load_size % type_size != 0 || store_size % type_size != 0 ||\n cache_line_size % type_size != 0) {\n return;\n }\n const size_t row_tile_elements = load_size * THREADS_PER_WARP / type_size;\n const size_t col_tile_elements = store_size * THREADS_PER_WARP / type_size;\n valid = (row_length % row_tile_elements == 0 && num_rows % col_tile_elements == 0);\n if (!valid) {\n return;\n }\n\n // Number of CUDA blocks\n num_blocks = (row_length / row_tile_elements) * (num_rows / col_tile_elements);\n\n // Parameters for performance model\n constexpr size_t warps_per_sm = 16; // Rough estimate for saturated SMs\n active_sm_count = std::min(DIVUP(num_blocks * warps_per_tile, warps_per_sm), sm_count);\n elements_per_load = (std::min(cache_line_size, row_tile_elements * type_size) / type_size);\n elements_per_store = (std::min(cache_line_size, col_tile_elements * type_size) / type_size);\n }\n\n /* Compare by estimated cost */\n bool operator<(const KernelConfig &other) const {\n if (this->valid && other.valid) {\n // cost ~ (1/elements_per_load + 1/elements_per_store) / active_sms\n // Note: Integer arithmetic ensures stable ordering\n const auto &l1 = this->elements_per_load;\n const auto &s1 = this->elements_per_store;\n const auto &p1 = this->active_sm_count;\n const auto &l2 = other.elements_per_load;\n const auto &s2 = other.elements_per_store;\n const auto &p2 = other.active_sm_count;\n const auto scale = l1 * s1 * p1 * l2 * s2 * p2;\n const auto cost1 = (scale / l1 + scale / s1) / p1;\n const auto cost2 = (scale / l2 + scale / s2) / p2;\n return cost1 < cost2;\n } else {\n return this->valid && !other.valid;\n }\n }\n};\n\ntemplate \n__global__ void __launch_bounds__(block_size)\n transpose_general_kernel(const Type *__restrict__ const input, const fp32 *const noop,\n Type *__restrict__ const output, const size_t row_length,\n const size_t num_rows) {\n if (noop != nullptr && noop[0] == 1.0f) return;\n\n // Vectorized load/store sizes\n constexpr size_t nvec_in = load_size / sizeof(Type);\n constexpr size_t nvec_out = store_size / sizeof(Type);\n using IVec = Vec;\n using OVec = Vec;\n\n // Thread indices\n // Note: Block is interpreted as a warp_size x num_warps grid\n constexpr size_t bdimx = THREADS_PER_WARP;\n constexpr size_t bdimy = warps_per_tile;\n const size_t tid = threadIdx.x;\n const size_t tidx = tid % bdimx;\n const size_t tidy = tid / bdimx;\n const size_t bid = blockIdx.x;\n\n // Input tensors are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec_out x nvec_in subtiles\n constexpr size_t tile_dim_m = THREADS_PER_WARP * nvec_out;\n constexpr size_t tile_dim_n = THREADS_PER_WARP * nvec_in;\n\n // Position of tile within tensor\n const size_t num_tiles_m = (num_rows + tile_dim_m - 1) / tile_dim_m;\n const size_t tile_id_m = bid % num_tiles_m;\n const size_t tile_id_n = bid / num_tiles_m;\n const size_t tile_row = tile_id_m * tile_dim_m;\n const size_t tile_col = tile_id_n * tile_dim_n;\n\n // Number of nvec_out x nvec_in subtiles for each thread to\n // load/store\n constexpr size_t num_iterations = THREADS_PER_WARP / warps_per_tile;\n\n // Load input and store to registers\n // Note: Each thread loads num_iterations subtiles and transposes in\n // registers.\n OVec local_output[nvec_in][num_iterations];\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n#pragma unroll\n for (size_t i2 = 0; i2 < nvec_out; ++i2) {\n const size_t row = tile_row + i1 * nvec_out + i2;\n const size_t col = tile_col + j1 * nvec_in;\n IVec local_input;\n local_input.clear();\n if (row < num_rows) {\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n if (col + j2 < row_length) {\n local_input.data.elt[j2] = input[row * row_length + col + j2];\n }\n }\n }\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n local_output[j2][iter].data.elt[i2] = local_input.data.elt[j2];\n }\n }\n }\n\n // Copy transposed output from registers to global memory\n __shared__ OVec shared_output[THREADS_PER_WARP][THREADS_PER_WARP + 1];\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n shared_output[j1][i1] = local_output[j2][iter];\n }\n __syncthreads();\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidx;\n const size_t j1 = tidy + iter * bdimy;\n const size_t row = tile_row + i1 * nvec_out;\n const size_t col = tile_col + j1 * nvec_in + j2;\n if (col < row_length) {\n#pragma unroll\n for (size_t i2 = 0; i2 < nvec_out; ++i2) {\n if (row + i2 < num_rows) {\n output[col * num_rows + row + i2] = shared_output[j1][i1].data.elt[i2];\n }\n }\n }\n }\n __syncthreads();\n }\n}\n\n} // namespace\n\nvoid transpose(const Tensor &input, const Tensor &noop, Tensor *output_, cudaStream_t stream) {\n Tensor &output = *output_;\n NVTE_CHECK(input.data.shape.size() == 2, \"Input must have 2 dimensions.\");\n NVTE_CHECK(output.data.shape.size() == 2, \"Output must have 2 dimensions.\");\n const size_t row_length = input.data.shape[1];\n const size_t num_rows = input.data.shape[0];\n\n NVTE_CHECK(output.data.shape[0] == row_length, \"Wrong dimension of output.\");\n NVTE_CHECK(output.data.shape[1] == num_rows, \"Wrong dimension of output.\");\n\n NVTE_CHECK(input.data.dptr != nullptr, \"Input is not allocated.\");\n NVTE_CHECK(output.data.dptr != nullptr, \"Output is not allocated.\");\n NVTE_CHECK(input.data.dtype == output.data.dtype, \"Input (dtype=\", to_string(input.data.dtype),\n \") and output (dtype=\", to_string(output.data.dtype), \") do not match.\");\n\n if (noop.data.dptr != nullptr) {\n NVTE_CHECK(noop.numel() == 1, \"Expected 1 element, \", \"but found \", noop.numel(), \".\");\n NVTE_CHECK(noop.data.dtype == DType::kFloat32);\n NVTE_CHECK(noop.data.dptr != nullptr);\n }\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n input.data.dtype, Type, constexpr const char *type_name = TypeInfo::name;\n constexpr size_t type_size = sizeof(Type);\n\n // Choose between runtime-compiled or statically-compiled kernel\n const bool aligned = (row_length % THREADS_PER_WARP == 0 && num_rows % THREADS_PER_WARP == 0);\n if (aligned && rtc::is_enabled()) { // Runtime-compiled tuned kernel\n // Pick kernel config\n std::vector kernel_configs;\n kernel_configs.reserve(16);\n const size_t sm_count = static_cast(cuda::sm_count());\n auto add_config = [&](size_t load_size, size_t store_size) {\n kernel_configs.emplace_back(row_length, num_rows, type_size, load_size, store_size,\n sm_count);\n };\n add_config(8, 8);\n add_config(4, 8);\n add_config(8, 4);\n add_config(4, 4);\n add_config(2, 8);\n add_config(8, 2);\n add_config(2, 4);\n add_config(4, 2);\n add_config(2, 2);\n add_config(1, 8);\n add_config(8, 1);\n add_config(1, 4);\n add_config(4, 1);\n add_config(1, 2);\n add_config(2, 1);\n add_config(1, 1);\n const auto &kernel_config = *std::min_element(kernel_configs.begin(), kernel_configs.end());\n NVTE_CHECK(kernel_config.valid, \"invalid kernel config\");\n const size_t load_size = kernel_config.load_size;\n const size_t store_size = kernel_config.store_size;\n const size_t num_blocks = kernel_config.num_blocks;\n\n // Compile NVRTC kernel if needed and launch\n auto &rtc_manager = rtc::KernelManager::instance();\n const std::string kernel_label = concat_strings(\n \"transpose\"\n \",type=\",\n type_name, \",load_size=\", load_size, \",store_size=\", store_size);\n if (!rtc_manager.is_compiled(kernel_label)) {\n std::string code = string_code_transpose_rtc_transpose_cu;\n code = regex_replace(code, \"__TYPE__\", type_name);\n code = regex_replace(code, \"__LOAD_SIZE__\", load_size);\n code = regex_replace(code, \"__STORE_SIZE__\", store_size);\n code = regex_replace(code, \"__WARPS_PER_TILE__\", warps_per_tile);\n code = regex_replace(code, \"__BLOCK_SIZE__\", block_size);\n rtc_manager.compile(kernel_label, \"transpose_optimized_kernel\", code,\n \"transformer_engine/common/transpose/rtc/transpose.cu\");\n }\n rtc_manager.launch(kernel_label, num_blocks, block_size, 0, stream,\n static_cast(input.data.dptr),\n static_cast(noop.data.dptr),\n static_cast(output.data.dptr), row_length, num_rows);\n } else { // Statically-compiled general kernel\n constexpr size_t load_size = 4;\n constexpr size_t store_size = 4;\n constexpr size_t row_tile_size = load_size / type_size * THREADS_PER_WARP;\n constexpr size_t col_tile_size = store_size / type_size * THREADS_PER_WARP;\n const int num_blocks = (DIVUP(row_length, row_tile_size) * DIVUP(num_rows, col_tile_size));\n transpose_general_kernel\n <<>>(static_cast(input.data.dptr),\n static_cast(noop.data.dptr),\n static_cast(output.data.dptr),\n row_length, num_rows);\n NVTE_CHECK_CUDA(cudaGetLastError());\n }); // NOLINT(*)\n}\n\n} // namespace detail\n} // namespace transformer_engine\n\nvoid nvte_transpose(const NVTETensor input, NVTETensor output, cudaStream_t stream) {\n NVTE_API_CALL(nvte_transpose);\n using namespace transformer_engine;\n auto noop = Tensor();\n detail::transpose(*convertNVTETensorCheck(input), noop, convertNVTETensor(output), stream);\n}\n\nvoid nvte_transpose_with_noop(const NVTETensor input, const NVTETensor noop, NVTETensor output,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_transpose_with_noop);\n using namespace transformer_engine;\n detail::transpose(*convertNVTETensorCheck(input), *convertNVTETensorCheck(noop),\n convertNVTETensor(output), stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \n\n#include \"../common.h\"\n#include \"../util/rtc.h\"\n#include \"../util/string.h\"\n#include \"../utils.cuh\"\n#include \"./transpose.h\"\n\nnamespace transformer_engine {\nnamespace detail {\n\nnamespace {\n\n// String with RTC kernel implementation\n#include \"string_code_transpose_rtc_transpose_cu.h\"\n\n// Hard-coded kernel parameters\nconstexpr size_t warps_per_tile = 4;\nconstexpr size_t block_size = THREADS_PER_WARP * warps_per_tile;\n\n/* Performance heuristics for optimized kernel parameters */\nstruct KernelConfig {\n /** Vector load size */\n size_t load_size;\n /** Vector store size */\n size_t store_size;\n\n /* Whether config is valid */\n bool valid = false;\n /* Number of CUDA blocks */\n size_t num_blocks = 0;\n\n /* Number of active SMs */\n size_t active_sm_count = 0;\n /* Elements per L1 cache load */\n size_t elements_per_load = 0;\n /* Elements per L1 cache store */\n size_t elements_per_store = 0;\n\n KernelConfig(size_t row_length, size_t num_rows, size_t type_size, size_t load_size_,\n size_t store_size_, size_t sm_count)\n : load_size{load_size_}, store_size{store_size_} {\n // Check that tiles are correctly aligned\n constexpr size_t cache_line_size = 128;\n if (load_size % type_size != 0 || store_size % type_size != 0 ||\n cache_line_size % type_size != 0) {\n return;\n }\n const size_t row_tile_elements = load_size * THREADS_PER_WARP / type_size;\n const size_t col_tile_elements = store_size * THREADS_PER_WARP / type_size;\n valid = (row_length % row_tile_elements == 0 && num_rows % col_tile_elements == 0);\n if (!valid) {\n return;\n }\n\n // Number of CUDA blocks\n num_blocks = (row_length / row_tile_elements) * (num_rows / col_tile_elements);\n\n // Parameters for performance model\n constexpr size_t warps_per_sm = 16; // Rough estimate for saturated SMs\n active_sm_count = std::min(DIVUP(num_blocks * warps_per_tile, warps_per_sm), sm_count);\n elements_per_load = (std::min(cache_line_size, row_tile_elements * type_size) / type_size);\n elements_per_store = (std::min(cache_line_size, col_tile_elements * type_size) / type_size);\n }\n\n /* Compare by estimated cost */\n bool operator<(const KernelConfig &other) const {\n if (this->valid && other.valid) {\n // cost ~ (1/elements_per_load + 1/elements_per_store) / active_sms\n // Note: Integer arithmetic ensures stable ordering\n const auto &l1 = this->elements_per_load;\n const auto &s1 = this->elements_per_store;\n const auto &p1 = this->active_sm_count;\n const auto &l2 = other.elements_per_load;\n const auto &s2 = other.elements_per_store;\n const auto &p2 = other.active_sm_count;\n const auto scale = l1 * s1 * p1 * l2 * s2 * p2;\n const auto cost1 = (scale / l1 + scale / s1) / p1;\n const auto cost2 = (scale / l2 + scale / s2) / p2;\n return cost1 < cost2;\n } else {\n return this->valid && !other.valid;\n }\n }\n};\n\ntemplate \n__global__ void __launch_bounds__(block_size)\n transpose_general_kernel(const Type *__restrict__ const input, const fp32 *const noop,\n Type *__restrict__ const output, const size_t row_length,\n const size_t num_rows) {\n if (noop != nullptr && noop[0] == 1.0f) return;\n\n // Vectorized load/store sizes\n constexpr size_t nvec_in = load_size / sizeof(Type);\n constexpr size_t nvec_out = store_size / sizeof(Type);\n using IVec = Vec;\n using OVec = Vec;\n\n // Thread indices\n // Note: Block is interpreted as a warp_size x num_warps grid\n constexpr size_t bdimx = THREADS_PER_WARP;\n constexpr size_t bdimy = warps_per_tile;\n const size_t tid = threadIdx.x;\n const size_t tidx = tid % bdimx;\n const size_t tidy = tid / bdimx;\n const size_t bid = blockIdx.x;\n\n // Input tensors are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec_out x nvec_in subtiles\n constexpr size_t tile_dim_m = THREADS_PER_WARP * nvec_out;\n constexpr size_t tile_dim_n = THREADS_PER_WARP * nvec_in;\n\n // Position of tile within tensor\n const size_t num_tiles_m = (num_rows + tile_dim_m - 1) / tile_dim_m;\n const size_t tile_id_m = bid % num_tiles_m;\n const size_t tile_id_n = bid / num_tiles_m;\n const size_t tile_row = tile_id_m * tile_dim_m;\n const size_t tile_col = tile_id_n * tile_dim_n;\n\n // Number of nvec_out x nvec_in subtiles for each thread to\n // load/store\n constexpr size_t num_iterations = THREADS_PER_WARP / warps_per_tile;\n\n // Load input and store to registers\n // Note: Each thread loads num_iterations subtiles and transposes in\n // registers.\n OVec local_output[nvec_in][num_iterations];\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n#pragma unroll\n for (size_t i2 = 0; i2 < nvec_out; ++i2) {\n const size_t row = tile_row + i1 * nvec_out + i2;\n const size_t col = tile_col + j1 * nvec_in;\n IVec local_input;\n local_input.clear();\n if (row < num_rows) {\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n if (col + j2 < row_length) {\n local_input.data.elt[j2] = input[row * row_length + col + j2];\n }\n }\n }\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n local_output[j2][iter].data.elt[i2] = local_input.data.elt[j2];\n }\n }\n }\n\n // Copy transposed output from registers to global memory\n __shared__ OVec shared_output[THREADS_PER_WARP][THREADS_PER_WARP + 1];\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n shared_output[j1][i1] = local_output[j2][iter];\n }\n __syncthreads();\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidx;\n const size_t j1 = tidy + iter * bdimy;\n const size_t row = tile_row + i1 * nvec_out;\n const size_t col = tile_col + j1 * nvec_in + j2;\n if (col < row_length) {\n#pragma unroll\n for (size_t i2 = 0; i2 < nvec_out; ++i2) {\n if (row + i2 < num_rows) {\n output[col * num_rows + row + i2] = shared_output[j1][i1].data.elt[i2];\n }\n }\n }\n }\n __syncthreads();\n }\n}\n\n} // namespace\n\nvoid transpose(const Tensor &input, const Tensor &noop, Tensor *output_, hipStream_t stream) {\n Tensor &output = *output_;\n NVTE_CHECK(input.data.shape.size() == 2, \"Input must have 2 dimensions.\");\n NVTE_CHECK(output.data.shape.size() == 2, \"Output must have 2 dimensions.\");\n const size_t row_length = input.data.shape[1];\n const size_t num_rows = input.data.shape[0];\n\n NVTE_CHECK(output.data.shape[0] == row_length, \"Wrong dimension of output.\");\n NVTE_CHECK(output.data.shape[1] == num_rows, \"Wrong dimension of output.\");\n\n NVTE_CHECK(input.data.dptr != nullptr, \"Input is not allocated.\");\n NVTE_CHECK(output.data.dptr != nullptr, \"Output is not allocated.\");\n NVTE_CHECK(input.data.dtype == output.data.dtype, \"Input (dtype=\", to_string(input.data.dtype),\n \") and output (dtype=\", to_string(output.data.dtype), \") do not match.\");\n\n if (noop.data.dptr != nullptr) {\n NVTE_CHECK(noop.numel() == 1, \"Expected 1 element, \", \"but found \", noop.numel(), \".\");\n NVTE_CHECK(noop.data.dtype == DType::kFloat32);\n NVTE_CHECK(noop.data.dptr != nullptr);\n }\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n input.data.dtype, Type, constexpr const char *type_name = TypeInfo::name;\n constexpr size_t type_size = sizeof(Type);\n\n // Choose between runtime-compiled or statically-compiled kernel\n const bool aligned = (row_length % THREADS_PER_WARP == 0 && num_rows % THREADS_PER_WARP == 0);\n if (aligned && rtc::is_enabled()) { // Runtime-compiled tuned kernel\n // Pick kernel config\n std::vector kernel_configs;\n kernel_configs.reserve(16);\n const size_t sm_count = static_cast(cuda::sm_count());\n auto add_config = [&](size_t load_size, size_t store_size) {\n kernel_configs.emplace_back(row_length, num_rows, type_size, load_size, store_size,\n sm_count);\n };\n add_config(8, 8);\n add_config(4, 8);\n add_config(8, 4);\n add_config(4, 4);\n add_config(2, 8);\n add_config(8, 2);\n add_config(2, 4);\n add_config(4, 2);\n add_config(2, 2);\n add_config(1, 8);\n add_config(8, 1);\n add_config(1, 4);\n add_config(4, 1);\n add_config(1, 2);\n add_config(2, 1);\n add_config(1, 1);\n const auto &kernel_config = *std::min_element(kernel_configs.begin(), kernel_configs.end());\n NVTE_CHECK(kernel_config.valid, \"invalid kernel config\");\n const size_t load_size = kernel_config.load_size;\n const size_t store_size = kernel_config.store_size;\n const size_t num_blocks = kernel_config.num_blocks;\n\n // Compile NVRTC kernel if needed and launch\n auto &rtc_manager = rtc::KernelManager::instance();\n const std::string kernel_label = concat_strings(\n \"transpose\"\n \",type=\",\n type_name, \",load_size=\", load_size, \",store_size=\", store_size);\n if (!rtc_manager.is_compiled(kernel_label)) {\n std::string code = string_code_transpose_rtc_transpose_cu;\n code = regex_replace(code, \"__TYPE__\", type_name);\n code = regex_replace(code, \"__LOAD_SIZE__\", load_size);\n code = regex_replace(code, \"__STORE_SIZE__\", store_size);\n code = regex_replace(code, \"__WARPS_PER_TILE__\", warps_per_tile);\n code = regex_replace(code, \"__BLOCK_SIZE__\", block_size);\n rtc_manager.compile(kernel_label, \"transpose_optimized_kernel\", code,\n \"transformer_engine/common/transpose/rtc/transpose.cu\");\n }\n rtc_manager.launch(kernel_label, num_blocks, block_size, 0, stream,\n static_cast(input.data.dptr),\n static_cast(noop.data.dptr),\n static_cast(output.data.dptr), row_length, num_rows);\n } else { // Statically-compiled general kernel\n constexpr size_t load_size = 4;\n constexpr size_t store_size = 4;\n constexpr size_t row_tile_size = load_size / type_size * THREADS_PER_WARP;\n constexpr size_t col_tile_size = store_size / type_size * THREADS_PER_WARP;\n const int num_blocks = (DIVUP(row_length, row_tile_size) * DIVUP(num_rows, col_tile_size));\n transpose_general_kernel\n <<>>(static_cast(input.data.dptr),\n static_cast(noop.data.dptr),\n static_cast(output.data.dptr),\n row_length, num_rows);\n NVTE_CHECK_CUDA(hipGetLastError());\n }); // NOLINT(*)\n}\n\n} // namespace detail\n} // namespace transformer_engine\n\nvoid nvte_transpose(const NVTETensor input, NVTETensor output, hipStream_t stream) {\n NVTE_API_CALL(nvte_transpose);\n using namespace transformer_engine;\n auto noop = Tensor();\n detail::transpose(*convertNVTETensorCheck(input), noop, convertNVTETensor(output), stream);\n}\n\nvoid nvte_transpose_with_noop(const NVTETensor input, const NVTETensor noop, NVTETensor output,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_transpose_with_noop);\n using namespace transformer_engine;\n detail::transpose(*convertNVTETensorCheck(input), *convertNVTETensorCheck(noop),\n convertNVTETensor(output), stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/transpose/transpose.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"../util/math.h\"\n#include \"../util/rtc.h\"\n#include \"../util/string.h\"\n#include \"../utils.cuh\"\n#include \"cast_transpose.h\"\n#include \"common/common.h\"\n\nnamespace transformer_engine {\n\nnamespace detail {\n\n// String with RTC kernel implementation\n#include \"string_code_transpose_rtc_cast_transpose_fusion_cu.h\"\n\n// STUFF TO TUNE\nconstexpr size_t n_warps_per_tile = 8;\nconstexpr size_t desired_load_size = 8;\nconstexpr size_t desired_store_size = 8;\nconstexpr size_t desired_load_size_dact = 4; // dAct fusion kernels use more registers\nconstexpr size_t desired_store_size_dact = 4;\n\nconstexpr size_t threads_per_warp = static_cast(THREADS_PER_WARP);\nconstexpr size_t max_threads_per_block = 256;\nconstexpr size_t reduce_dbias_num_threads = 256;\nconstexpr size_t cast_transpose_num_threads = n_warps_per_tile * threads_per_warp;\nconstexpr size_t n_warps_per_block = cast_transpose_num_threads / threads_per_warp;\nstatic_assert(cast_transpose_num_threads <= max_threads_per_block);\n\n/* Performance heuristics for optimized kernel parameters */\nstruct KernelConfig {\n size_t load_size = 0; // Vector load size\n size_t store_size = 0; // Vector store size to transposed output\n\n bool valid = false; // Whether config is valid\n bool is_dact = false; // Whether dact is used\n size_t num_blocks = 0; // Number of CUDA blocks\n\n size_t active_sm_count = 0; // Number of active SMs\n size_t elements_per_load = 0; // Elements per L1 cache load\n size_t elements_per_load_dact = 0; // Elements per L1 cache load dact\n size_t elements_per_store_c = 0; // Elements per L1 cache store to cast output\n size_t elements_per_store_t = 0; // Elements per L1 cache store to transposed output\n\n KernelConfig(size_t row_length, size_t num_rows, size_t itype_size, size_t itype2_size,\n size_t otype_size, size_t load_size_, size_t store_size_, size_t sm_count,\n bool is_dact_)\n : load_size{load_size_}, store_size{store_size_}, is_dact{is_dact_} {\n if (is_dact) {\n if (load_size > desired_load_size_dact || store_size > desired_store_size_dact) {\n return;\n }\n }\n\n // Check that tiles are correctly aligned\n constexpr size_t cache_line_size = 128;\n if (load_size % itype_size != 0 || store_size % otype_size != 0 ||\n cache_line_size % itype_size != 0 || cache_line_size % otype_size != 0) {\n return;\n }\n /* row_tile_elements */\n const size_t tile_size_x = (load_size * THREADS_PER_WARP) / itype_size;\n /* col_tile_elements */\n const size_t tile_size_y = (store_size * THREADS_PER_WARP) / otype_size;\n const size_t num_tiles_x = row_length / tile_size_x;\n const size_t num_tiles_y = num_rows / tile_size_y;\n\n valid = (row_length % tile_size_x == 0 && num_rows % tile_size_y == 0);\n if (!valid) {\n return;\n }\n\n // Number of CUDA blocks\n num_blocks = num_tiles_x * num_tiles_y;\n\n // Parameters for performance model\n constexpr size_t warps_per_sm = 16; // Rough estimate for saturated SMs\n active_sm_count = std::min(DIVUP(num_blocks * n_warps_per_tile, warps_per_sm), sm_count);\n elements_per_load = (std::min(cache_line_size, tile_size_x * itype_size) / itype_size);\n elements_per_load_dact = (std::min(cache_line_size, tile_size_x * itype2_size) / itype2_size);\n elements_per_store_c = (std::min(cache_line_size, tile_size_x * otype_size) / otype_size);\n elements_per_store_t = (std::min(cache_line_size, tile_size_y * otype_size) / otype_size);\n }\n\n /* Compare by estimated cost */\n bool operator<(const KernelConfig &other) const {\n if (this->valid && other.valid) {\n // cost ~ (1/elements_per_load\n // + 1/elements_per_load_dact\n // + 1/elements_per_store_c\n // + 1/elements_per_store_t) / active_sms\n // Note: Integer arithmetic ensures stable ordering\n const auto &l1 = this->elements_per_load;\n const auto &la1 = this->elements_per_load_dact;\n const auto &sc1 = this->elements_per_store_c;\n const auto &st1 = this->elements_per_store_t;\n const auto &p1 = this->active_sm_count;\n const auto &l2 = other.elements_per_load;\n const auto &la2 = other.elements_per_load_dact;\n const auto &sc2 = other.elements_per_store_c;\n const auto &st2 = other.elements_per_store_t;\n const auto &p2 = other.active_sm_count;\n const auto scale1 = l1 * sc1 * st1 * p1 * (is_dact ? la1 : 1);\n const auto scale2 = l2 * sc2 * st2 * p2 * (is_dact ? la2 : 1);\n const auto scale = scale1 * scale2;\n const auto cost1 =\n (scale / l1 + scale / sc1 + scale / st1 + (is_dact ? (scale / la1) : 0)) / p1;\n const auto cost2 =\n (scale / l2 + scale / sc2 + scale / st2 + (is_dact ? (scale / la2) : 0)) / p2;\n\n return cost1 < cost2;\n } else {\n return this->valid && !other.valid;\n }\n }\n};\n\ntemplate \ninline __device__ void cast_and_transpose_regs(const CVec (&in)[nvec_out],\n OVec (&out_trans)[nvec_in],\n CVec &out_dbias, // NOLINT(*)\n typename OVec::type *output_cast_tile,\n const size_t current_place, const size_t stride,\n const CType scale,\n CType &amax, // NOLINT(*)\n const int dbias_shfl_src_lane,\n const bool valid_store) {\n using OType = typename OVec::type;\n using OVecC = Vec;\n\n CVec step_dbias;\n if constexpr (IS_DBIAS) {\n step_dbias.clear();\n }\n\n#pragma unroll\n for (unsigned int i = 0; i < nvec_out; ++i) {\n OVecC out_cast;\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n const CType tmp = in[i].data.elt[j];\n if constexpr (IS_DBIAS) {\n step_dbias.data.elt[j] += tmp; // dbias: thread tile local accumulation\n }\n out_cast.data.elt[j] = static_cast(tmp * scale);\n out_trans[j].data.elt[i] = static_cast(tmp * scale); // thread tile transpose\n\n __builtin_assume(amax >= 0);\n amax = fmaxf(fabsf(tmp), amax);\n }\n if (IS_FULL_TILE || valid_store) {\n out_cast.store_to(output_cast_tile, current_place + stride * i);\n }\n }\n\n if constexpr (IS_DBIAS) {\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n CType elt = step_dbias.data.elt[j];\n elt = __shfl_sync(0xffffffff, elt, dbias_shfl_src_lane); // shuffle data in a warp\n out_dbias.data.elt[j] += elt;\n }\n }\n}\n\nvoid populate_cast_transpose_dbias_workspace_config(const Tensor &cast_output, /*cast*/\n Tensor *workspace, const int nvec_out) {\n const size_t row_length = cast_output.flat_last_dim();\n const size_t num_rows = cast_output.flat_first_dim();\n\n const size_t tile_size_y = (nvec_out * THREADS_PER_WARP);\n NVTE_CHECK(num_rows % nvec_out == 0, \"Unsupported shape.\");\n\n const size_t num_rows_partial_dbias = DIVUP(num_rows, tile_size_y);\n\n if (workspace->data.dptr == nullptr) {\n workspace->data.shape = {num_rows_partial_dbias, row_length};\n workspace->data.dtype = DType::kFloat32;\n } else {\n // Check that workspace matches expected size\n const size_t workspace_size = get_buffer_size_bytes(\n std::accumulate(workspace->data.shape.begin(), workspace->data.shape.end(), 1,\n std::multiplies()),\n workspace->data.dtype);\n const size_t required_size =\n get_buffer_size_bytes(num_rows_partial_dbias, row_length, DType::kFloat32);\n NVTE_CHECK(workspace_size >= required_size, \"Invalid workspace (expected dims=(\",\n num_rows_partial_dbias, \",\", row_length, \"), dtype=\", to_string(DType::kFloat32),\n \"; found dims=\", workspace->data.shape,\n \", dtype=\", typeToNumBits(workspace->data.dtype), \" bits)\");\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(reduce_dbias_num_threads)\n reduce_dbias_kernel(OutputType *const dbias_output, const ComputeType *const dbias_partial,\n const int row_length, const int num_rows) {\n using ComputeVec = Vec;\n using OutputVec = Vec;\n\n const int thread_id = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (thread_id * nvec >= row_length) {\n return;\n }\n\n const ComputeType *const thread_in_base = dbias_partial + thread_id * nvec;\n OutputType *const thread_out_base = dbias_output + thread_id * nvec;\n\n const int stride_in_vec = row_length / nvec;\n\n ComputeVec ldg_vec;\n ComputeVec acc_vec;\n acc_vec.clear();\n for (int i = 0; i < num_rows; ++i) {\n ldg_vec.load_from(thread_in_base, i * stride_in_vec);\n#pragma unroll\n for (int e = 0; e < nvec; ++e) {\n acc_vec.data.elt[e] += ldg_vec.data.elt[e];\n }\n }\n\n OutputVec stg_vec;\n#pragma unroll\n for (int e = 0; e < nvec; ++e) {\n stg_vec.data.elt[e] = OutputType(acc_vec.data.elt[e]);\n }\n stg_vec.store_to(thread_out_base, 0);\n}\n\ntemplate \nvoid reduce_dbias(const Tensor &workspace, Tensor *dbias, const size_t row_length,\n const size_t num_rows, const int nvec_out, cudaStream_t stream) {\n constexpr int reduce_dbias_store_bytes = 8; // stg.64\n constexpr int reduce_dbias_nvec = reduce_dbias_store_bytes / sizeof(InputType);\n\n NVTE_CHECK(row_length % reduce_dbias_nvec == 0, \"Unsupported shape.\");\n\n const size_t reduce_dbias_row_length = row_length;\n const size_t reduce_dbias_num_rows =\n DIVUP(num_rows, static_cast(nvec_out * THREADS_PER_WARP));\n const size_t reduce_dbias_num_blocks =\n DIVUP(row_length, reduce_dbias_num_threads * reduce_dbias_nvec);\n\n using DbiasOutputType = fp32;\n reduce_dbias_kernel\n <<>>(\n reinterpret_cast(dbias->data.dptr),\n reinterpret_cast(workspace.data.dptr), reduce_dbias_row_length,\n reduce_dbias_num_rows);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\ntemplate \n__global__ void __launch_bounds__(cast_transpose_num_threads)\n cast_transpose_fused_kernel_notaligned(const Param param, const size_t row_length,\n const size_t num_rows, const size_t num_tiles) {\n static_assert(!(IS_DACT && IS_ACT), \"forward and backward activation are mutually exclusive\");\n using IType = typename Param::InputType;\n using IType2 = typename Param::InputType2;\n using OType = typename Param::OutputType;\n using CType = typename Param::ComputeType;\n using IVec = Vec;\n using IVec2 = Vec;\n using OVec = Vec;\n using CVec = Vec;\n\n extern __shared__ char scratch[];\n\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const unsigned int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;\n const size_t num_tiles_x =\n (row_length + nvec_in * THREADS_PER_WARP - 1) / (nvec_in * THREADS_PER_WARP);\n const size_t tile_id =\n blockIdx.x * blockDim.x / (THREADS_PER_WARP * n_warps_per_tile) + warp_id / n_warps_per_tile;\n if (tile_id >= num_tiles) {\n return;\n }\n\n const size_t tile_id_x = tile_id % num_tiles_x;\n const size_t tile_id_y = tile_id / num_tiles_x;\n\n const size_t tile_offset =\n (tile_id_x * nvec_in + tile_id_y * row_length * nvec_out) * THREADS_PER_WARP;\n const size_t tile_offset_transp =\n (tile_id_y * nvec_out + tile_id_x * num_rows * nvec_in) * THREADS_PER_WARP;\n\n const IType *const my_input_tile = param.input + tile_offset;\n const IType2 *const my_act_input_tile = param.act_input + tile_offset;\n OType *const my_output_c_tile = param.output_c + tile_offset;\n OType *const my_output_t_tile = param.output_t + tile_offset_transp;\n CType *const my_partial_dbias_tile =\n param.workspace + (tile_id_x * (nvec_in * THREADS_PER_WARP) + tile_id_y * row_length);\n\n const size_t stride = row_length / nvec_in;\n const size_t output_stride = num_rows / nvec_out;\n const size_t row_length_rest = stride - tile_id_x * THREADS_PER_WARP;\n const size_t row_height_rest = output_stride - tile_id_y * THREADS_PER_WARP;\n const unsigned int tile_length =\n row_length_rest > THREADS_PER_WARP ? THREADS_PER_WARP : row_length_rest;\n const unsigned int tile_height =\n row_height_rest > THREADS_PER_WARP ? THREADS_PER_WARP : row_height_rest;\n\n OVec *const my_scratch =\n reinterpret_cast(scratch) +\n (my_id_in_warp + warp_id / n_warps_per_tile * THREADS_PER_WARP) * (THREADS_PER_WARP + 1);\n\n CVec *const my_dbias_scratch = reinterpret_cast(scratch);\n\n IVec in[2][nvec_out];\n IVec2 act_in[2][nvec_out];\n const unsigned int warp_id_in_tile = warp_id % n_warps_per_tile;\n constexpr unsigned int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n OVec out_space[n_iterations][nvec_in];\n\n size_t current_stride = warp_id_in_tile * n_iterations * nvec_out * stride;\n size_t current_row = (tile_id_y * THREADS_PER_WARP + warp_id_in_tile * n_iterations) * nvec_out;\n unsigned int my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n CType amax = 0;\n const CType scale = param.scale_ptr != nullptr ? *param.scale_ptr : 1;\n\n CVec partial_dbias;\n if constexpr (IS_DBIAS) {\n partial_dbias.clear();\n }\n\n {\n const bool valid_load = my_place < tile_length && warp_id_in_tile * n_iterations < tile_height;\n#pragma unroll\n for (unsigned int i = 0; i < nvec_out; ++i) {\n if (valid_load) {\n const size_t ld_offset = current_stride + my_place + stride * i;\n in[0][i].load_from(my_input_tile, ld_offset);\n if constexpr (IS_DACT) {\n act_in[0][i].load_from(my_act_input_tile, ld_offset);\n }\n } else {\n in[0][i].clear();\n if constexpr (IS_DACT) {\n act_in[0][i].clear();\n }\n }\n }\n }\n\n#pragma unroll\n for (unsigned int i = 0; i < n_iterations; ++i) {\n const size_t current_place = current_stride + my_place;\n const unsigned int my_place_in = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n const unsigned int current_in = (i + 1) % 2;\n if (i < n_iterations - 1) {\n const bool valid_load =\n my_place_in < tile_length && warp_id_in_tile * n_iterations + i + 1 < tile_height;\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n if (valid_load) {\n const size_t ld_offset = current_stride + my_place_in + stride * (nvec_out + j);\n in[current_in][j].load_from(my_input_tile, ld_offset);\n if constexpr (IS_DACT) {\n act_in[current_in][j].load_from(my_act_input_tile, ld_offset);\n }\n } else {\n in[current_in][j].clear();\n if constexpr (IS_DACT) {\n act_in[current_in][j].clear();\n }\n }\n }\n }\n CVec after_dact[nvec_out]; // NOLINT(*)\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n#pragma unroll\n for (unsigned int k = 0; k < nvec_in; ++k) {\n if constexpr (IS_DACT) {\n after_dact[j].data.elt[k] = CType(in[current_in ^ 1][j].data.elt[k]) *\n OP(act_in[current_in ^ 1][j].data.elt[k], {});\n } else if constexpr (IS_ACT) {\n after_dact[j].data.elt[k] = OP(in[current_in ^ 1][j].data.elt[k], {});\n } else {\n after_dact[j].data.elt[k] = CType(in[current_in ^ 1][j].data.elt[k]);\n }\n }\n }\n const int dbias_shfl_src_lane =\n (my_id_in_warp + i + warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n constexpr bool IS_FULL_TILE = false;\n const bool valid_store =\n (my_place < tile_length) && (warp_id_in_tile * n_iterations + i < tile_height);\n\n cast_and_transpose_regs(after_dact, out_space[i], partial_dbias,\n my_output_c_tile, current_place, stride, scale,\n amax, dbias_shfl_src_lane, valid_store);\n\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += nvec_out * stride;\n current_row += nvec_out;\n }\n\n for (unsigned int i = 0; i < nvec_in; ++i) {\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * nvec_in;\n for (unsigned int j = 0; warp_id_in_tile * n_iterations + j < tile_length; ++j) {\n const bool valid_store = my_place < tile_height;\n if (valid_store) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile,\n current_stride + my_place);\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * nvec_in;\n }\n __syncthreads();\n }\n\n if constexpr (IS_DBIAS) {\n my_dbias_scratch[threadIdx.x] = partial_dbias;\n __syncthreads();\n if (warp_id_in_tile == 0) {\n#pragma unroll\n for (unsigned int i = 1; i < n_warps_per_tile; ++i) {\n CVec tmp = my_dbias_scratch[threadIdx.x + i * THREADS_PER_WARP];\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n partial_dbias.data.elt[j] += tmp.data.elt[j];\n }\n }\n if (my_id_in_warp < tile_length) {\n partial_dbias.store_to(my_partial_dbias_tile, my_id_in_warp);\n }\n }\n }\n\n // Reduce amax over block\n if (param.amax != nullptr) {\n amax = reduce_max(amax, warp_id);\n if (threadIdx.x == 0) {\n static_assert(std::is_same::value);\n atomicMaxFloat(param.amax, amax);\n }\n }\n\n // Update scale-inverse\n if (blockIdx.x == 0 && threadIdx.x == 0 && param.scale_inv != nullptr) {\n reciprocal(param.scale_inv, scale);\n }\n}\n\nstatic const char *ActTypeToString[] = {\n \"none\", // 0\n \"sigmoid\", // 1\n \"dsigmoid\", // 2\n \"gelu\", // 3\n \"dgelu\", // 4\n \"qgelu\", // 5\n \"dqgelu\", // 6\n \"silu\", // 7\n \"dsilu\", // 8\n \"relu\", // 9\n \"drelu\", // 10\n \"srelu\", // 11\n \"dsrelu\" // 12\n};\n\ntemplate \nconstexpr int get_activation_type() {\n constexpr decltype(OP) ActivationList[] = {\n nullptr, // 0\n &sigmoid, // 1\n &dsigmoid, // 2\n &gelu, // 3\n &dgelu, // 4\n &qgelu, // 5\n &dqgelu, // 6\n &silu, // 7\n &dsilu, // 8\n &relu, // 9\n &drelu, // 10\n &srelu, // 11\n &dsrelu // 12\n };\n#pragma unroll\n for (int i = 0; i < sizeof(ActivationList) / sizeof(ActivationList[0]); ++i) {\n if (OP == ActivationList[i]) {\n return i;\n }\n }\n return 0;\n}\n\ntemplate \nvoid cast_transpose_fused(const Tensor &input, const Tensor *act_input, Tensor *output,\n Tensor *dbias, Tensor *workspace, cudaStream_t stream) {\n // Check tensors, unless querying dbias workspace\n if (!IS_DBIAS || workspace->data.dptr != nullptr) {\n CheckInputTensor(input, \"cast_transpose_fused_input\");\n CheckOutputTensor(*output, \"output\");\n if constexpr (IS_DBIAS) {\n NVTE_CHECK(dbias != nullptr && dbias->has_data());\n CheckOutputTensor(*dbias, \"dbias\");\n }\n if constexpr (IS_DACT) {\n NVTE_CHECK(act_input != nullptr && act_input->has_data());\n CheckInputTensor(*act_input, \"act_input\");\n }\n }\n\n // Check that inputs and outputs are available\n NVTE_CHECK(input.has_data(), \"Input is not allocated\");\n NVTE_CHECK(output->has_data(), \"Output rowwise data is not allocated\");\n NVTE_CHECK(output->has_columnwise_data(), \"Output columnwise data is not allocated\");\n\n // Flatten tensor to 2D\n NVTE_CHECK(input.data.shape == output->data.shape,\n \"Input and output shapes do not match (input=\", input.data.shape,\n \", output=\", output->data.shape);\n const size_t row_length = input.flat_last_dim();\n const size_t num_rows = input.flat_first_dim();\n\n // Check that cast and transposed output data matches\n NVTE_CHECK(output->data.dtype == output->columnwise_data.dtype,\n \"Cast and transposed output types must match.\");\n NVTE_CHECK(output->scale_inv.dptr == output->columnwise_scale_inv.dptr,\n \"Cast and transposed outputs need to share scale-inverse tensor.\");\n\n if constexpr (IS_DBIAS) {\n NVTE_CHECK(dbias->data.dtype == input.data.dtype, \"DBias must have the same type as input.\");\n NVTE_CHECK(dbias->data.shape == std::vector{row_length}, \"Wrong shape of DBias.\");\n }\n if constexpr (IS_DACT) {\n NVTE_CHECK(input.dtype() == act_input->dtype(), \"Types of both inputs must match.\");\n NVTE_CHECK(input.data.shape == act_input->data.shape, \"Shapes of both inputs must match.\");\n }\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.dtype(), InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n output->dtype(), OutputType, using InputType2 = InputType;\n using Param = CTDBiasDActParam;\n\n constexpr int itype_size = sizeof(InputType);\n constexpr int itype2_size = sizeof(InputType2);\n constexpr int otype_size = sizeof(OutputType);\n\n const bool aligned =\n (row_length % THREADS_PER_WARP == 0) && (num_rows % THREADS_PER_WARP == 0);\n const bool jit_compiled = aligned && rtc::is_enabled();\n\n size_t load_size = (IS_DACT ? desired_load_size_dact : desired_load_size);\n size_t store_size = (IS_DACT ? desired_store_size_dact : desired_store_size);\n size_t num_blocks;\n\n if (jit_compiled) {\n // Pick kernel config\n std::vector kernel_configs;\n kernel_configs.reserve(16);\n const size_t sm_count = static_cast(cuda::sm_count());\n auto add_config = [&](size_t load_size_config, size_t store_size_config) {\n kernel_configs.emplace_back(row_length, num_rows, itype_size, itype2_size, otype_size,\n load_size_config, store_size_config, sm_count, IS_DACT);\n };\n add_config(8, 8);\n add_config(4, 8);\n add_config(8, 4);\n add_config(4, 4);\n add_config(2, 8);\n add_config(8, 2);\n add_config(2, 4);\n add_config(4, 2);\n add_config(2, 2);\n add_config(1, 8);\n add_config(8, 1);\n add_config(1, 4);\n add_config(4, 1);\n add_config(1, 2);\n add_config(2, 1);\n add_config(1, 1);\n\n // Select the kernel configuration with the lowest cost\n const auto &kernel_config =\n *std::min_element(kernel_configs.begin(), kernel_configs.end());\n NVTE_CHECK(kernel_config.valid, \"invalid kernel config\");\n load_size = kernel_config.load_size;\n store_size = kernel_config.store_size;\n num_blocks = kernel_config.num_blocks;\n }\n\n const size_t nvec_in = load_size / itype_size;\n const size_t nvec_out = store_size / otype_size;\n const size_t tile_size_x = nvec_in * threads_per_warp;\n const size_t tile_size_y = nvec_out * threads_per_warp;\n const size_t num_tiles_x = DIVUP(row_length, tile_size_x);\n const size_t num_tiles_y = DIVUP(num_rows, tile_size_y);\n const size_t num_tiles = num_tiles_x * num_tiles_y;\n\n NVTE_CHECK(row_length % nvec_in == 0, \"Unsupported shape.\");\n NVTE_CHECK(num_rows % nvec_out == 0, \"Unsupported shape.\");\n\n if (!jit_compiled) {\n num_blocks = DIVUP(num_tiles * n_warps_per_tile, n_warps_per_block);\n } if constexpr (IS_DBIAS) {\n // Check workspace size\n populate_cast_transpose_dbias_workspace_config(*output, workspace, nvec_out);\n if (workspace->data.dptr == nullptr) {\n return;\n }\n }\n\n size_t VecOutputTypeSize;\n switch (nvec_out) {\n case 1:\n VecOutputTypeSize = sizeof(Vec);\n break;\n case 2:\n VecOutputTypeSize = sizeof(Vec);\n break;\n case 4:\n VecOutputTypeSize = sizeof(Vec);\n break;\n case 8:\n VecOutputTypeSize = sizeof(Vec);\n break;\n } size_t shared_size_transpose = cast_transpose_num_threads / n_warps_per_tile *\n (threads_per_warp + 1) * VecOutputTypeSize;\n\n if constexpr (IS_DBIAS) {\n size_t VecComputeTypeSize;\n switch (nvec_in) {\n case 1:\n VecComputeTypeSize = sizeof(Vec);\n break;\n case 2:\n VecComputeTypeSize = sizeof(Vec);\n break;\n case 4:\n VecComputeTypeSize = sizeof(Vec);\n break;\n case 8:\n VecComputeTypeSize = sizeof(Vec);\n break;\n }\n const size_t shared_size_dbias = cast_transpose_num_threads * VecComputeTypeSize;\n if (shared_size_transpose < shared_size_dbias) {\n shared_size_transpose = shared_size_dbias;\n }\n }\n\n Param param;\n param.input = reinterpret_cast(input.data.dptr);\n param.output_c = reinterpret_cast(output->data.dptr);\n param.output_t = reinterpret_cast(output->columnwise_data.dptr);\n param.scale_ptr = reinterpret_cast(output->scale.dptr);\n param.amax = reinterpret_cast(output->amax.dptr);\n param.scale_inv = reinterpret_cast(output->scale_inv.dptr);\n if constexpr (IS_DBIAS) {\n param.workspace = reinterpret_cast(workspace->data.dptr);\n } if constexpr (IS_DACT) {\n param.act_input = reinterpret_cast(act_input->data.dptr);\n }\n\n // Runtime-compiled tuned kernel\n if (jit_compiled) {\n constexpr const char *itype_name = TypeInfo::name;\n constexpr const char *itype2_name = TypeInfo::name;\n constexpr const char *otype_name = TypeInfo::name;\n\n int actType = 0;\n if constexpr (IS_DACT || IS_ACT) {\n actType = get_activation_type();\n }\n\n // Compile NVRTC kernel if needed and launch\n auto &rtc_manager = rtc::KernelManager::instance();\n const std::string kernel_label = concat_strings(\n \"cast_transpose_fusion\"\n \",itype=\",\n itype_name, \",itype2=\", itype2_name, \",otype=\", otype_name,\n \",load_size=\", load_size, \",store_size=\", store_size, \",IS_DBIAS=\", IS_DBIAS,\n \",IS_DACT=\", IS_DACT, \",IS_ACT=\", IS_ACT,\n \",activationType=\", ActTypeToString[actType]);\n\n if (!rtc_manager.is_compiled(kernel_label)) {\n std::string code = string_code_transpose_rtc_cast_transpose_fusion_cu;\n code = regex_replace(code, \"__ITYPE__\", itype_name);\n code = regex_replace(code, \"__ITYPE2__\", itype2_name);\n code = regex_replace(code, \"__OTYPE__\", otype_name);\n code = regex_replace(code, \"__LOAD_SIZE__\", load_size);\n code = regex_replace(code, \"__STORE_SIZE__\", store_size);\n code = regex_replace(code, \"__WARPS_PER_TILE__\", n_warps_per_tile);\n code = regex_replace(code, \"__BLOCK_SIZE__\", cast_transpose_num_threads);\n code = regex_replace(code, \"__IS_DBIAS__\", IS_DBIAS);\n code = regex_replace(code, \"__IS_DACT__\", IS_DACT);\n code = regex_replace(code, \"__IS_ACT__\", IS_ACT);\n code = regex_replace(code, \"__ACTIVATION_TYPE__\", actType);\n\n rtc_manager.compile(\n kernel_label, \"cast_transpose_fusion_kernel_optimized\", code,\n \"transformer_engine/common/transpose/rtc/cast_transpose_fusion.cu\");\n }\n\n rtc_manager.set_cache_config(kernel_label, CU_FUNC_CACHE_PREFER_SHARED);\n\n rtc_manager.launch(kernel_label, num_blocks, cast_transpose_num_threads,\n shared_size_transpose, stream, param, row_length, num_rows,\n num_tiles);\n } else { // Statically-compiled general kernel\n constexpr size_t load_size = IS_DACT ? desired_load_size_dact : desired_load_size;\n constexpr size_t store_size = IS_DACT ? desired_store_size_dact : desired_store_size;\n constexpr size_t nvec_in = load_size / itype_size;\n constexpr size_t nvec_out = store_size / otype_size;\n\n NVTE_CHECK(row_length % nvec_in == 0, \"Unsupported shape.\");\n NVTE_CHECK(num_rows % nvec_out == 0, \"Unsupported shape.\");\n\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n cast_transpose_fused_kernel_notaligned,\n cudaFuncAttributePreferredSharedMemoryCarveout, 100));\n cast_transpose_fused_kernel_notaligned\n <<>>(\n param, row_length, num_rows, num_tiles);\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n\n if constexpr (IS_DBIAS) {\n reduce_dbias(*workspace, dbias, row_length, num_rows, nvec_out, stream);\n }); // NOLINT(*)\n ); // NOLINT(*)\n}\n\ntemplate \n__global__ void __launch_bounds__(cast_transpose_num_threads)\n dgated_act_cast_transpose_kernel(const IType *const input, const IType *const act_input,\n OType *const output_c, OType *const output_t,\n const CType *const scale_ptr, CType *const amax,\n CType *const scale_inv, const size_t row_length,\n const size_t num_rows, const size_t num_tiles) {\n using IVec = Vec;\n using OVec = Vec;\n using CVec = Vec;\n\n extern __shared__ char scratch[];\n\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;\n const size_t num_tiles_x = row_length / (nvec_in * THREADS_PER_WARP);\n const size_t tile_id =\n blockIdx.x * blockDim.x / (THREADS_PER_WARP * n_warps_per_tile) + warp_id / n_warps_per_tile;\n if (tile_id >= num_tiles) {\n return;\n }\n\n const size_t tile_id_x = tile_id % num_tiles_x;\n const size_t tile_id_y = tile_id / num_tiles_x;\n\n const IType *const my_input_tile =\n input + (tile_id_x * nvec_in + tile_id_y * row_length * nvec_out) * THREADS_PER_WARP;\n const IType *const my_act_input_tile =\n act_input + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP;\n const IType *const my_gate_input_tile =\n act_input + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP +\n row_length;\n OType *const my_output_c_tile_0 =\n output_c + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP;\n OType *const my_output_c_tile_1 =\n output_c + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP +\n row_length;\n OType *const my_output_t_tile_0 =\n output_t + (tile_id_y * nvec_out + tile_id_x * num_rows * nvec_in) * THREADS_PER_WARP;\n OType *const my_output_t_tile_1 =\n output_t + (tile_id_y * nvec_out + tile_id_x * num_rows * nvec_in) * THREADS_PER_WARP +\n row_length * num_rows;\n OVec *const my_scratch =\n reinterpret_cast(scratch) +\n (my_id_in_warp + warp_id / n_warps_per_tile * THREADS_PER_WARP) * (THREADS_PER_WARP + 1);\n\n IVec in[2][nvec_out];\n IVec act_in[2][nvec_out];\n IVec gate_in[2][nvec_out];\n const unsigned int warp_id_in_tile = warp_id % n_warps_per_tile;\n constexpr unsigned int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n OVec out_space_0[n_iterations][nvec_in];\n OVec out_space_1[n_iterations][nvec_in];\n\n const size_t stride = row_length / nvec_in;\n const size_t output_stride = num_rows / nvec_out;\n size_t current_stride = warp_id_in_tile * n_iterations * nvec_out * stride;\n unsigned int my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n const size_t stride2 = 2 * row_length / nvec_in;\n size_t current_stride2 = warp_id_in_tile * n_iterations * nvec_out * stride2;\n CType max = 0;\n const CType scale = scale_ptr != nullptr ? *scale_ptr : 1;\n\n CVec partial_dbias;\n\n#pragma unroll\n for (unsigned int i = 0; i < nvec_out; ++i) {\n in[0][i].load_from(my_input_tile, current_stride + my_place + stride * i);\n act_in[0][i].load_from(my_act_input_tile, current_stride2 + my_place + stride2 * i);\n gate_in[0][i].load_from(my_gate_input_tile, current_stride2 + my_place + stride2 * i);\n }\n#pragma unroll\n for (unsigned int i = 0; i < n_iterations; ++i) {\n const size_t current_place = current_stride2 + my_place;\n const unsigned int my_place_in = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n const unsigned int current_in = (i + 1) % 2;\n if (i < n_iterations - 1) {\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n in[current_in][j].load_from(my_input_tile,\n current_stride + my_place_in + stride * (nvec_out + j));\n act_in[current_in][j].load_from(my_act_input_tile,\n current_stride2 + my_place_in + stride2 * (nvec_out + j));\n gate_in[current_in][j].load_from(my_gate_input_tile,\n current_stride2 + my_place_in + stride2 * (nvec_out + j));\n }\n }\n CVec after_dact[nvec_out]; // NOLINT(*)\n CVec after_dgate[nvec_out]; // NOLINT(*)\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n#pragma unroll\n for (unsigned int k = 0; k < nvec_in; ++k) {\n after_dact[j].data.elt[k] = OP1(act_in[current_in ^ 1][j].data.elt[k], {}) *\n CType(in[current_in ^ 1][j].data.elt[k]) *\n CType(gate_in[current_in ^ 1][j].data.elt[k]);\n after_dgate[j].data.elt[k] = CType(in[current_in ^ 1][j].data.elt[k]) *\n OP2(act_in[current_in ^ 1][j].data.elt[k], {});\n }\n }\n OVec out_trans_0[nvec_in]; // NOLINT(*)\n OVec out_trans_1[nvec_in]; // NOLINT(*)\n\n constexpr bool IS_DBIAS = false;\n constexpr bool IS_FULL_TILE = true;\n constexpr bool valid_store = true;\n constexpr int dbias_shfl_src_lane = 0;\n\n cast_and_transpose_regs(after_dact, out_trans_0, partial_dbias,\n my_output_c_tile_0, current_place, stride2,\n scale, max, dbias_shfl_src_lane, valid_store);\n\n cast_and_transpose_regs(after_dgate, out_trans_1, partial_dbias,\n my_output_c_tile_1, current_place, stride2,\n scale, max, dbias_shfl_src_lane, valid_store);\n\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n out_space_0[i][j].data.vec = out_trans_0[j].data.vec;\n out_space_1[i][j].data.vec = out_trans_1[j].data.vec;\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += nvec_out * stride;\n current_stride2 += nvec_out * stride2;\n }\n\n for (unsigned int i = 0; i < nvec_in; ++i) {\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space_0[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * nvec_in;\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile_0,\n current_stride + my_place);\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * nvec_in;\n }\n __syncthreads();\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space_1[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * nvec_in;\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile_1,\n current_stride + my_place);\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * nvec_in;\n }\n __syncthreads();\n }\n\n // Reduce amax over block\n if (amax != nullptr) {\n max = reduce_max(max, warp_id);\n if (threadIdx.x == 0) {\n static_assert(std::is_same::value);\n atomicMaxFloat(amax, max);\n }\n }\n\n // Update scale-inverse\n if (blockIdx.x == 0 && threadIdx.x == 0 && scale_inv != nullptr) {\n reciprocal(scale_inv, scale);\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(cast_transpose_num_threads)\n dgated_act_cast_transpose_kernel_notaligned(const IType *const input,\n const IType *const act_input, OType *const output_c,\n OType *const output_t, const CType *const scale_ptr,\n CType *const amax, CType *const scale_inv,\n const size_t row_length, const size_t num_rows,\n const size_t num_tiles) {\n using IVec = Vec;\n using OVec = Vec;\n using CVec = Vec;\n\n extern __shared__ char scratch[];\n\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;\n const size_t num_tiles_x =\n (row_length + nvec_in * THREADS_PER_WARP - 1) / (nvec_in * THREADS_PER_WARP);\n const size_t tile_id =\n blockIdx.x * blockDim.x / (THREADS_PER_WARP * n_warps_per_tile) + warp_id / n_warps_per_tile;\n if (tile_id >= num_tiles) return;\n const size_t tile_id_x = tile_id % num_tiles_x;\n const size_t tile_id_y = tile_id / num_tiles_x;\n\n const IType *const my_input_tile =\n input + (tile_id_x * nvec_in + tile_id_y * row_length * nvec_out) * THREADS_PER_WARP;\n const IType *const my_act_input_tile =\n act_input + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP;\n const IType *const my_gate_input_tile =\n act_input + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP +\n row_length;\n OType *const my_output_c_tile_0 =\n output_c + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP;\n OType *const my_output_c_tile_1 =\n output_c + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP +\n row_length;\n OType *const my_output_t_tile_0 =\n output_t + (tile_id_y * nvec_out + tile_id_x * num_rows * nvec_in) * THREADS_PER_WARP;\n OType *const my_output_t_tile_1 =\n output_t + (tile_id_y * nvec_out + tile_id_x * num_rows * nvec_in) * THREADS_PER_WARP +\n row_length * num_rows;\n const size_t stride = row_length / nvec_in;\n const size_t stride2 = 2 * row_length / nvec_in;\n const size_t output_stride = num_rows / nvec_out;\n const size_t row_length_rest = stride - tile_id_x * THREADS_PER_WARP;\n const size_t row_height_rest = output_stride - tile_id_y * THREADS_PER_WARP;\n const unsigned int tile_length =\n row_length_rest > THREADS_PER_WARP ? THREADS_PER_WARP : row_length_rest;\n const unsigned int tile_height =\n row_height_rest > THREADS_PER_WARP ? THREADS_PER_WARP : row_height_rest;\n\n OVec *const my_scratch =\n reinterpret_cast(scratch) +\n (my_id_in_warp + warp_id / n_warps_per_tile * THREADS_PER_WARP) * (THREADS_PER_WARP + 1);\n\n IVec in[2][nvec_out];\n IVec act_in[2][nvec_out];\n IVec gate_in[2][nvec_out];\n const unsigned int warp_id_in_tile = warp_id % n_warps_per_tile;\n constexpr unsigned int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n OVec out_space_0[n_iterations][nvec_in];\n OVec out_space_1[n_iterations][nvec_in];\n\n size_t current_stride = warp_id_in_tile * n_iterations * nvec_out * stride;\n size_t current_stride2 = warp_id_in_tile * n_iterations * nvec_out * stride2;\n unsigned int my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n CType max = 0;\n const CType scale = scale_ptr != nullptr ? *scale_ptr : 1;\n\n CVec partial_dbias;\n\n {\n const bool valid_load = my_place < tile_length && warp_id_in_tile * n_iterations < tile_height;\n#pragma unroll\n for (unsigned int i = 0; i < nvec_out; ++i) {\n if (valid_load) {\n in[0][i].load_from(my_input_tile, current_stride + my_place + stride * i);\n act_in[0][i].load_from(my_act_input_tile, current_stride2 + my_place + stride2 * i);\n gate_in[0][i].load_from(my_gate_input_tile, current_stride2 + my_place + stride2 * i);\n } else {\n in[0][i].clear();\n act_in[0][i].clear();\n gate_in[0][i].clear();\n }\n }\n }\n#pragma unroll\n for (unsigned int i = 0; i < n_iterations; ++i) {\n const size_t current_place = current_stride2 + my_place;\n const unsigned int my_place_in = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n const unsigned int current_in = (i + 1) % 2;\n if (i < n_iterations - 1) {\n {\n const bool valid_load =\n my_place_in < tile_length && warp_id_in_tile * n_iterations + i + 1 < tile_height;\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n if (valid_load) {\n in[current_in][j].load_from(my_input_tile,\n current_stride + my_place_in + stride * (nvec_out + j));\n act_in[current_in][j].load_from(\n my_act_input_tile, current_stride2 + my_place_in + stride2 * (nvec_out + j));\n gate_in[current_in][j].load_from(\n my_gate_input_tile, current_stride2 + my_place_in + stride2 * (nvec_out + j));\n } else {\n in[current_in][j].clear();\n act_in[current_in][j].clear();\n gate_in[current_in][j].clear();\n }\n }\n }\n }\n CVec after_dact[nvec_out]; // NOLINT(*)\n CVec after_dgate[nvec_out]; // NOLINT(*)\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n#pragma unroll\n for (unsigned int k = 0; k < nvec_in; ++k) {\n after_dact[j].data.elt[k] = OP1(act_in[current_in ^ 1][j].data.elt[k], {}) *\n CType(in[current_in ^ 1][j].data.elt[k]) *\n CType(gate_in[current_in ^ 1][j].data.elt[k]);\n after_dgate[j].data.elt[k] = CType(in[current_in ^ 1][j].data.elt[k]) *\n OP2(act_in[current_in ^ 1][j].data.elt[k], {});\n }\n }\n OVec out_trans_0[nvec_in]; // NOLINT(*)\n OVec out_trans_1[nvec_in]; // NOLINT(*)\n\n constexpr bool IS_DBIAS = false;\n constexpr bool IS_FULL_TILE = false;\n constexpr int dbias_shfl_src_lane = 0;\n const bool valid_store =\n (my_place < tile_length) && (warp_id_in_tile * n_iterations + i < tile_height);\n\n cast_and_transpose_regs(after_dact, out_trans_0, partial_dbias,\n my_output_c_tile_0, current_place, stride2,\n scale, max, dbias_shfl_src_lane, valid_store);\n cast_and_transpose_regs(after_dgate, out_trans_1, partial_dbias,\n my_output_c_tile_1, current_place, stride2,\n scale, max, dbias_shfl_src_lane, valid_store);\n\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n out_space_0[i][j].data.vec = out_trans_0[j].data.vec;\n out_space_1[i][j].data.vec = out_trans_1[j].data.vec;\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += nvec_out * stride;\n current_stride2 += nvec_out * stride2;\n }\n\n for (unsigned int i = 0; i < nvec_in; ++i) {\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space_0[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * nvec_in;\n for (unsigned int j = 0; warp_id_in_tile * n_iterations + j < tile_length; ++j) {\n const bool valid_store = my_place < tile_height;\n if (valid_store) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile_0,\n current_stride + my_place);\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * nvec_in;\n }\n __syncthreads();\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space_1[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * nvec_in;\n for (unsigned int j = 0; warp_id_in_tile * n_iterations + j < tile_length; ++j) {\n const bool valid_store = my_place < tile_height;\n if (valid_store) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile_1,\n current_stride + my_place);\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * nvec_in;\n }\n __syncthreads();\n }\n\n // Reduce amax over block\n if (amax != nullptr) {\n max = reduce_max(max, warp_id);\n if (threadIdx.x == 0) {\n static_assert(std::is_same::value);\n atomicMaxFloat(amax, max);\n }\n }\n\n // Update scale-inverse\n if (blockIdx.x == 0 && threadIdx.x == 0 && scale_inv != nullptr) {\n reciprocal(scale_inv, scale);\n }\n}\n\ntemplate \nvoid dgated_act_cast_transpose(const Tensor &input, const Tensor &gated_act_input, Tensor *output,\n cudaStream_t stream) {\n CheckInputTensor(input, \"dgated_act_cast_transpose_input\");\n CheckInputTensor(gated_act_input, \"dgated_act_cast_transpose_gated_act_input\");\n CheckOutputTensor(*output, \"dgated_act_cast_transpose_output\");\n\n NVTE_CHECK(input.data.shape.size() == 2, \"Input must have 2 dimensions.\");\n NVTE_CHECK(gated_act_input.data.shape.size() == 2, \"Input must have 2 dimensions.\");\n NVTE_CHECK(output->has_data() && output->has_columnwise_data(),\n \"Both rowwise and columnwise data need to be allocated.\");\n NVTE_CHECK(output->data.shape.size() == 2, \"C output must have 2 dimensions.\");\n NVTE_CHECK(output->columnwise_data.shape.size() == 2, \"T output must have 2 dimensions.\");\n const size_t row_length = input.data.shape[1];\n const size_t num_rows = input.data.shape[0];\n\n NVTE_CHECK(gated_act_input.data.shape[0] == num_rows, \"Wrong dimension of output.\");\n NVTE_CHECK(gated_act_input.data.shape[1] == row_length * 2, \"Wrong dimension of output.\");\n NVTE_CHECK(output->data.shape[0] == num_rows, \"Wrong dimension of output.\");\n NVTE_CHECK(output->data.shape[1] == row_length * 2, \"Wrong dimension of output.\");\n NVTE_CHECK(output->columnwise_data.shape[0] == row_length * 2, \"Wrong dimension of T output.\");\n NVTE_CHECK(output->columnwise_data.shape[1] == num_rows, \"Wrong dimension of T output.\");\n\n NVTE_CHECK(input.data.dtype == gated_act_input.data.dtype, \"Types of both inputs must match.\");\n\n NVTE_CHECK(output->data.dtype == output->columnwise_data.dtype,\n \"C and T outputs need to have the same type.\");\n NVTE_CHECK(output->scale_inv.dptr == output->columnwise_scale_inv.dptr,\n \"C and T outputs need to share scale inverse tensor.\");\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.dtype(), InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n output->dtype(), OutputType, using InputType2 = InputType;\n /* dact fusion kernel uses more registers */\n constexpr int desired_load_size_dact = 4;\n constexpr int desired_store_size_dact = 4; constexpr int itype_size = sizeof(InputType);\n constexpr int otype_size = sizeof(OutputType);\n constexpr int nvec_in = desired_load_size_dact / itype_size;\n constexpr int nvec_out = desired_store_size_dact / otype_size;\n\n NVTE_CHECK(row_length % nvec_in == 0, \"Unsupported shape.\");\n NVTE_CHECK(num_rows % nvec_out == 0, \"Unsupported shape.\");\n const size_t n_tiles =\n DIVUP(row_length, static_cast(nvec_in * THREADS_PER_WARP)) *\n DIVUP(num_rows, static_cast(nvec_out * THREADS_PER_WARP));\n const size_t n_warps_per_block = cast_transpose_num_threads / THREADS_PER_WARP;\n const size_t n_blocks = DIVUP(n_tiles * n_warps_per_tile, n_warps_per_block);\n\n const bool full_tile = row_length % (nvec_in * THREADS_PER_WARP) == 0 &&\n num_rows % (nvec_out * THREADS_PER_WARP) == 0;\n const size_t shmem_size = cast_transpose_num_threads / n_warps_per_tile *\n (THREADS_PER_WARP + 1) * sizeof(Vec);\n if (full_tile) {\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n dgated_act_cast_transpose_kernel,\n cudaFuncAttributePreferredSharedMemoryCarveout, 100));\n\n dgated_act_cast_transpose_kernel\n <<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(gated_act_input.data.dptr),\n reinterpret_cast(output->data.dptr),\n reinterpret_cast(output->columnwise_data.dptr),\n reinterpret_cast(output->scale.dptr),\n reinterpret_cast(output->amax.dptr),\n reinterpret_cast(output->scale_inv.dptr), row_length, num_rows,\n n_tiles);\n NVTE_CHECK_CUDA(cudaGetLastError());\n } else {\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n dgated_act_cast_transpose_kernel_notaligned,\n cudaFuncAttributePreferredSharedMemoryCarveout, 100));\n dgated_act_cast_transpose_kernel_notaligned\n <<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(gated_act_input.data.dptr),\n reinterpret_cast(output->data.dptr),\n reinterpret_cast(output->columnwise_data.dptr),\n reinterpret_cast(output->scale.dptr),\n reinterpret_cast(output->amax.dptr),\n reinterpret_cast(output->scale_inv.dptr), row_length, num_rows,\n n_tiles);\n NVTE_CHECK_CUDA(cudaGetLastError());\n }); // NOLINT(*)\n ); // NOLINT(*)\n}\n\n// Explicit template instantiation\ntemplate void cast_transpose_fused(\n const Tensor &, const Tensor *, Tensor *, Tensor *, Tensor *, cudaStream_t);\n#define NVTE_INSTANTIATE_ACTIVATION(op) \\\n template void cast_transpose_fused>( \\\n const Tensor &, const Tensor *, Tensor *, Tensor *, Tensor *, cudaStream_t); \\\n template void cast_transpose_fused>( \\\n const Tensor &, const Tensor *, Tensor *, Tensor *, Tensor *, cudaStream_t);\nNVTE_INSTANTIATE_ACTIVATION(relu);\nNVTE_INSTANTIATE_ACTIVATION(srelu);\nNVTE_INSTANTIATE_ACTIVATION(gelu);\nNVTE_INSTANTIATE_ACTIVATION(qgelu);\nNVTE_INSTANTIATE_ACTIVATION(silu);\n#undef NVTE_INSTANTIATE_ACTIVATION\n\n} // namespace detail\n\n} // namespace transformer_engine\n\nusing ComputeType = typename transformer_engine::fp32;\n\nvoid nvte_cast_transpose_dbias(const NVTETensor input, NVTETensor output, NVTETensor dbias,\n NVTETensor workspace, cudaStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose_dbias);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n constexpr bool IS_DBIAS = true;\n constexpr bool IS_DACT = false;\n constexpr bool IS_ACT = false;\n\n constexpr const NVTETensor activation_input = nullptr;\n\n cast_transpose_fused(\n *convertNVTETensorCheck(input), convertNVTETensor(activation_input),\n convertNVTETensor(output), convertNVTETensor(dbias), convertNVTETensor(workspace), stream);\n}\n\nvoid nvte_cast_transpose_dbias_dgelu(const NVTETensor input, const NVTETensor act_input,\n NVTETensor output, NVTETensor dbias, NVTETensor workspace,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose_dbias_dgelu);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n constexpr bool IS_DBIAS = true;\n constexpr bool IS_DACT = true;\n constexpr bool IS_ACT = false;\n\n cast_transpose_fused>(\n *convertNVTETensorCheck(input), convertNVTETensorCheck(act_input),\n convertNVTETensorCheck(output), convertNVTETensorCheck(dbias), convertNVTETensor(workspace),\n stream);\n}\n\nvoid nvte_cast_transpose_dbias_dsilu(const NVTETensor input, const NVTETensor silu_input,\n NVTETensor output, NVTETensor dbias, NVTETensor workspace,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose_dbias_dsilu);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n constexpr bool IS_DBIAS = true;\n constexpr bool IS_DACT = true;\n constexpr bool IS_ACT = false;\n\n cast_transpose_fused>(\n *convertNVTETensorCheck(input), convertNVTETensorCheck(silu_input),\n convertNVTETensorCheck(output), convertNVTETensorCheck(dbias), convertNVTETensor(workspace),\n stream);\n}\n\nvoid nvte_cast_transpose_dbias_drelu(const NVTETensor input, const NVTETensor relu_input,\n NVTETensor output, NVTETensor dbias, NVTETensor workspace,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose_dbias_drelu);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n constexpr bool IS_DBIAS = true;\n constexpr bool IS_DACT = true;\n constexpr bool IS_ACT = false;\n\n cast_transpose_fused>(\n *convertNVTETensorCheck(input), convertNVTETensorCheck(relu_input),\n convertNVTETensorCheck(output), convertNVTETensorCheck(dbias), convertNVTETensor(workspace),\n stream);\n}\n\nvoid nvte_cast_transpose_dbias_dsrelu(const NVTETensor input, const NVTETensor srelu_input,\n NVTETensor output, NVTETensor dbias, NVTETensor workspace,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose_dbias_dsrelu);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n constexpr bool IS_DBIAS = true;\n constexpr bool IS_DACT = true;\n constexpr bool IS_ACT = false;\n\n cast_transpose_fused>(\n *convertNVTETensorCheck(input), convertNVTETensorCheck(srelu_input),\n convertNVTETensorCheck(output), convertNVTETensorCheck(dbias), convertNVTETensor(workspace),\n stream);\n}\n\nvoid nvte_cast_transpose_dbias_dqgelu(const NVTETensor input, const NVTETensor qgelu_input,\n NVTETensor output, NVTETensor dbias, NVTETensor workspace,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose_dbias_dqgelu);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n constexpr bool IS_DBIAS = true;\n constexpr bool IS_DACT = true;\n constexpr bool IS_ACT = false;\n\n cast_transpose_fused>(\n *convertNVTETensorCheck(input), convertNVTETensorCheck(qgelu_input),\n convertNVTETensorCheck(output), convertNVTETensorCheck(dbias), convertNVTETensor(workspace),\n stream);\n}\n\nvoid nvte_dgeglu_cast_transpose(const NVTETensor input, const NVTETensor gated_act_input,\n NVTETensor output, cudaStream_t stream) {\n NVTE_API_CALL(nvte_dgeglu_cast_transpose);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n dgated_act_cast_transpose, gelu>(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(gated_act_input),\n convertNVTETensorCheck(output), stream);\n}\n\nvoid nvte_dswiglu_cast_transpose(const NVTETensor input, const NVTETensor swiglu_input,\n NVTETensor output, cudaStream_t stream) {\n NVTE_API_CALL(nvte_dswiglu_cast_transpose);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n dgated_act_cast_transpose, silu>(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(swiglu_input),\n convertNVTETensorCheck(output), stream);\n}\n\nvoid nvte_dreglu_cast_transpose(const NVTETensor input, const NVTETensor gated_act_input,\n NVTETensor output, cudaStream_t stream) {\n NVTE_API_CALL(nvte_dreglu_cast_transpose);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n dgated_act_cast_transpose, relu>(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(gated_act_input),\n convertNVTETensorCheck(output), stream);\n}\n\nvoid nvte_dsreglu_cast_transpose(const NVTETensor input, const NVTETensor gated_act_input,\n NVTETensor output, cudaStream_t stream) {\n NVTE_API_CALL(nvte_dsreglu_cast_transpose);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n dgated_act_cast_transpose, srelu>(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(gated_act_input),\n convertNVTETensorCheck(output), stream);\n}\n\nvoid nvte_dqgeglu_cast_transpose(const NVTETensor input, const NVTETensor gated_act_input,\n NVTETensor output, cudaStream_t stream) {\n NVTE_API_CALL(nvte_dqgeglu_cast_transpose);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n dgated_act_cast_transpose, qgelu>(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(gated_act_input),\n convertNVTETensorCheck(output), stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"../util/math.h\"\n#include \"../util/rtc.h\"\n#include \"../util/string.h\"\n#include \"../utils.cuh\"\n#include \"cast_transpose.h\"\n#include \"common/common.h\"\n\nnamespace transformer_engine {\n\nnamespace detail {\n\n// String with RTC kernel implementation\n#include \"string_code_transpose_rtc_cast_transpose_fusion_cu.h\"\n\n// STUFF TO TUNE\nconstexpr size_t n_warps_per_tile = 8;\nconstexpr size_t desired_load_size = 8;\nconstexpr size_t desired_store_size = 8;\nconstexpr size_t desired_load_size_dact = 4; // dAct fusion kernels use more registers\nconstexpr size_t desired_store_size_dact = 4;\n\nconstexpr size_t threads_per_warp = static_cast(THREADS_PER_WARP);\nconstexpr size_t max_threads_per_block = 256;\nconstexpr size_t reduce_dbias_num_threads = 256;\nconstexpr size_t cast_transpose_num_threads = n_warps_per_tile * threads_per_warp;\nconstexpr size_t n_warps_per_block = cast_transpose_num_threads / threads_per_warp;\nstatic_assert(cast_transpose_num_threads <= max_threads_per_block);\n\n/* Performance heuristics for optimized kernel parameters */\nstruct KernelConfig {\n size_t load_size = 0; // Vector load size\n size_t store_size = 0; // Vector store size to transposed output\n\n bool valid = false; // Whether config is valid\n bool is_dact = false; // Whether dact is used\n size_t num_blocks = 0; // Number of CUDA blocks\n\n size_t active_sm_count = 0; // Number of active SMs\n size_t elements_per_load = 0; // Elements per L1 cache load\n size_t elements_per_load_dact = 0; // Elements per L1 cache load dact\n size_t elements_per_store_c = 0; // Elements per L1 cache store to cast output\n size_t elements_per_store_t = 0; // Elements per L1 cache store to transposed output\n\n KernelConfig(size_t row_length, size_t num_rows, size_t itype_size, size_t itype2_size,\n size_t otype_size, size_t load_size_, size_t store_size_, size_t sm_count,\n bool is_dact_)\n : load_size{load_size_}, store_size{store_size_}, is_dact{is_dact_} {\n if (is_dact) {\n if (load_size > desired_load_size_dact || store_size > desired_store_size_dact) {\n return;\n }\n }\n\n // Check that tiles are correctly aligned\n constexpr size_t cache_line_size = 128;\n if (load_size % itype_size != 0 || store_size % otype_size != 0 ||\n cache_line_size % itype_size != 0 || cache_line_size % otype_size != 0) {\n return;\n }\n /* row_tile_elements */\n const size_t tile_size_x = (load_size * THREADS_PER_WARP) / itype_size;\n /* col_tile_elements */\n const size_t tile_size_y = (store_size * THREADS_PER_WARP) / otype_size;\n const size_t num_tiles_x = row_length / tile_size_x;\n const size_t num_tiles_y = num_rows / tile_size_y;\n\n valid = (row_length % tile_size_x == 0 && num_rows % tile_size_y == 0);\n if (!valid) {\n return;\n }\n\n // Number of CUDA blocks\n num_blocks = num_tiles_x * num_tiles_y;\n\n // Parameters for performance model\n constexpr size_t warps_per_sm = 16; // Rough estimate for saturated SMs\n active_sm_count = std::min(DIVUP(num_blocks * n_warps_per_tile, warps_per_sm), sm_count);\n elements_per_load = (std::min(cache_line_size, tile_size_x * itype_size) / itype_size);\n elements_per_load_dact = (std::min(cache_line_size, tile_size_x * itype2_size) / itype2_size);\n elements_per_store_c = (std::min(cache_line_size, tile_size_x * otype_size) / otype_size);\n elements_per_store_t = (std::min(cache_line_size, tile_size_y * otype_size) / otype_size);\n }\n\n /* Compare by estimated cost */\n bool operator<(const KernelConfig &other) const {\n if (this->valid && other.valid) {\n // cost ~ (1/elements_per_load\n // + 1/elements_per_load_dact\n // + 1/elements_per_store_c\n // + 1/elements_per_store_t) / active_sms\n // Note: Integer arithmetic ensures stable ordering\n const auto &l1 = this->elements_per_load;\n const auto &la1 = this->elements_per_load_dact;\n const auto &sc1 = this->elements_per_store_c;\n const auto &st1 = this->elements_per_store_t;\n const auto &p1 = this->active_sm_count;\n const auto &l2 = other.elements_per_load;\n const auto &la2 = other.elements_per_load_dact;\n const auto &sc2 = other.elements_per_store_c;\n const auto &st2 = other.elements_per_store_t;\n const auto &p2 = other.active_sm_count;\n const auto scale1 = l1 * sc1 * st1 * p1 * (is_dact ? la1 : 1);\n const auto scale2 = l2 * sc2 * st2 * p2 * (is_dact ? la2 : 1);\n const auto scale = scale1 * scale2;\n const auto cost1 =\n (scale / l1 + scale / sc1 + scale / st1 + (is_dact ? (scale / la1) : 0)) / p1;\n const auto cost2 =\n (scale / l2 + scale / sc2 + scale / st2 + (is_dact ? (scale / la2) : 0)) / p2;\n\n return cost1 < cost2;\n } else {\n return this->valid && !other.valid;\n }\n }\n};\n\ntemplate \ninline __device__ void cast_and_transpose_regs(const CVec (&in)[nvec_out],\n OVec (&out_trans)[nvec_in],\n CVec &out_dbias, // NOLINT(*)\n typename OVec::type *output_cast_tile,\n const size_t current_place, const size_t stride,\n const CType scale,\n CType &amax, // NOLINT(*)\n const int dbias_shfl_src_lane,\n const bool valid_store) {\n using OType = typename OVec::type;\n using OVecC = Vec;\n\n CVec step_dbias;\n if constexpr (IS_DBIAS) {\n step_dbias.clear();\n }\n\n#pragma unroll\n for (unsigned int i = 0; i < nvec_out; ++i) {\n OVecC out_cast;\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n const CType tmp = in[i].data.elt[j];\n if constexpr (IS_DBIAS) {\n step_dbias.data.elt[j] += tmp; // dbias: thread tile local accumulation\n }\n out_cast.data.elt[j] = static_cast(tmp * scale);\n out_trans[j].data.elt[i] = static_cast(tmp * scale); // thread tile transpose\n\n __builtin_assume(amax >= 0);\n amax = fmaxf(fabsf(tmp), amax);\n }\n if (IS_FULL_TILE || valid_store) {\n out_cast.store_to(output_cast_tile, current_place + stride * i);\n }\n }\n\n if constexpr (IS_DBIAS) {\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n CType elt = step_dbias.data.elt[j];\n elt = __shfl_sync(0xffffffff, elt, dbias_shfl_src_lane); // shuffle data in a warp\n out_dbias.data.elt[j] += elt;\n }\n }\n}\n\nvoid populate_cast_transpose_dbias_workspace_config(const Tensor &cast_output, /*cast*/\n Tensor *workspace, const int nvec_out) {\n const size_t row_length = cast_output.flat_last_dim();\n const size_t num_rows = cast_output.flat_first_dim();\n\n const size_t tile_size_y = (nvec_out * THREADS_PER_WARP);\n NVTE_CHECK(num_rows % nvec_out == 0, \"Unsupported shape.\");\n\n const size_t num_rows_partial_dbias = DIVUP(num_rows, tile_size_y);\n\n if (workspace->data.dptr == nullptr) {\n workspace->data.shape = {num_rows_partial_dbias, row_length};\n workspace->data.dtype = DType::kFloat32;\n } else {\n // Check that workspace matches expected size\n const size_t workspace_size = get_buffer_size_bytes(\n std::accumulate(workspace->data.shape.begin(), workspace->data.shape.end(), 1,\n std::multiplies()),\n workspace->data.dtype);\n const size_t required_size =\n get_buffer_size_bytes(num_rows_partial_dbias, row_length, DType::kFloat32);\n NVTE_CHECK(workspace_size >= required_size, \"Invalid workspace (expected dims=(\",\n num_rows_partial_dbias, \",\", row_length, \"), dtype=\", to_string(DType::kFloat32),\n \"; found dims=\", workspace->data.shape,\n \", dtype=\", typeToNumBits(workspace->data.dtype), \" bits)\");\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(reduce_dbias_num_threads)\n reduce_dbias_kernel(OutputType *const dbias_output, const ComputeType *const dbias_partial,\n const int row_length, const int num_rows) {\n using ComputeVec = Vec;\n using OutputVec = Vec;\n\n const int thread_id = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (thread_id * nvec >= row_length) {\n return;\n }\n\n const ComputeType *const thread_in_base = dbias_partial + thread_id * nvec;\n OutputType *const thread_out_base = dbias_output + thread_id * nvec;\n\n const int stride_in_vec = row_length / nvec;\n\n ComputeVec ldg_vec;\n ComputeVec acc_vec;\n acc_vec.clear();\n for (int i = 0; i < num_rows; ++i) {\n ldg_vec.load_from(thread_in_base, i * stride_in_vec);\n#pragma unroll\n for (int e = 0; e < nvec; ++e) {\n acc_vec.data.elt[e] += ldg_vec.data.elt[e];\n }\n }\n\n OutputVec stg_vec;\n#pragma unroll\n for (int e = 0; e < nvec; ++e) {\n stg_vec.data.elt[e] = OutputType(acc_vec.data.elt[e]);\n }\n stg_vec.store_to(thread_out_base, 0);\n}\n\ntemplate \nvoid reduce_dbias(const Tensor &workspace, Tensor *dbias, const size_t row_length,\n const size_t num_rows, const int nvec_out, hipStream_t stream) {\n constexpr int reduce_dbias_store_bytes = 8; // stg.64\n constexpr int reduce_dbias_nvec = reduce_dbias_store_bytes / sizeof(InputType);\n\n NVTE_CHECK(row_length % reduce_dbias_nvec == 0, \"Unsupported shape.\");\n\n const size_t reduce_dbias_row_length = row_length;\n const size_t reduce_dbias_num_rows =\n DIVUP(num_rows, static_cast(nvec_out * THREADS_PER_WARP));\n const size_t reduce_dbias_num_blocks =\n DIVUP(row_length, reduce_dbias_num_threads * reduce_dbias_nvec);\n\n using DbiasOutputType = fp32;\n reduce_dbias_kernel\n <<>>(\n reinterpret_cast(dbias->data.dptr),\n reinterpret_cast(workspace.data.dptr), reduce_dbias_row_length,\n reduce_dbias_num_rows);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\ntemplate \n__global__ void __launch_bounds__(cast_transpose_num_threads)\n cast_transpose_fused_kernel_notaligned(const Param param, const size_t row_length,\n const size_t num_rows, const size_t num_tiles) {\n static_assert(!(IS_DACT && IS_ACT), \"forward and backward activation are mutually exclusive\");\n using IType = typename Param::InputType;\n using IType2 = typename Param::InputType2;\n using OType = typename Param::OutputType;\n using CType = typename Param::ComputeType;\n using IVec = Vec;\n using IVec2 = Vec;\n using OVec = Vec;\n using CVec = Vec;\n\n extern __shared__ char scratch[];\n\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const unsigned int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;\n const size_t num_tiles_x =\n (row_length + nvec_in * THREADS_PER_WARP - 1) / (nvec_in * THREADS_PER_WARP);\n const size_t tile_id =\n blockIdx.x * blockDim.x / (THREADS_PER_WARP * n_warps_per_tile) + warp_id / n_warps_per_tile;\n if (tile_id >= num_tiles) {\n return;\n }\n\n const size_t tile_id_x = tile_id % num_tiles_x;\n const size_t tile_id_y = tile_id / num_tiles_x;\n\n const size_t tile_offset =\n (tile_id_x * nvec_in + tile_id_y * row_length * nvec_out) * THREADS_PER_WARP;\n const size_t tile_offset_transp =\n (tile_id_y * nvec_out + tile_id_x * num_rows * nvec_in) * THREADS_PER_WARP;\n\n const IType *const my_input_tile = param.input + tile_offset;\n const IType2 *const my_act_input_tile = param.act_input + tile_offset;\n OType *const my_output_c_tile = param.output_c + tile_offset;\n OType *const my_output_t_tile = param.output_t + tile_offset_transp;\n CType *const my_partial_dbias_tile =\n param.workspace + (tile_id_x * (nvec_in * THREADS_PER_WARP) + tile_id_y * row_length);\n\n const size_t stride = row_length / nvec_in;\n const size_t output_stride = num_rows / nvec_out;\n const size_t row_length_rest = stride - tile_id_x * THREADS_PER_WARP;\n const size_t row_height_rest = output_stride - tile_id_y * THREADS_PER_WARP;\n const unsigned int tile_length =\n row_length_rest > THREADS_PER_WARP ? THREADS_PER_WARP : row_length_rest;\n const unsigned int tile_height =\n row_height_rest > THREADS_PER_WARP ? THREADS_PER_WARP : row_height_rest;\n\n OVec *const my_scratch =\n reinterpret_cast(scratch) +\n (my_id_in_warp + warp_id / n_warps_per_tile * THREADS_PER_WARP) * (THREADS_PER_WARP + 1);\n\n CVec *const my_dbias_scratch = reinterpret_cast(scratch);\n\n IVec in[2][nvec_out];\n IVec2 act_in[2][nvec_out];\n const unsigned int warp_id_in_tile = warp_id % n_warps_per_tile;\n constexpr unsigned int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n OVec out_space[n_iterations][nvec_in];\n\n size_t current_stride = warp_id_in_tile * n_iterations * nvec_out * stride;\n size_t current_row = (tile_id_y * THREADS_PER_WARP + warp_id_in_tile * n_iterations) * nvec_out;\n unsigned int my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n CType amax = 0;\n const CType scale = param.scale_ptr != nullptr ? *param.scale_ptr : 1;\n\n CVec partial_dbias;\n if constexpr (IS_DBIAS) {\n partial_dbias.clear();\n }\n\n {\n const bool valid_load = my_place < tile_length && warp_id_in_tile * n_iterations < tile_height;\n#pragma unroll\n for (unsigned int i = 0; i < nvec_out; ++i) {\n if (valid_load) {\n const size_t ld_offset = current_stride + my_place + stride * i;\n in[0][i].load_from(my_input_tile, ld_offset);\n if constexpr (IS_DACT) {\n act_in[0][i].load_from(my_act_input_tile, ld_offset);\n }\n } else {\n in[0][i].clear();\n if constexpr (IS_DACT) {\n act_in[0][i].clear();\n }\n }\n }\n }\n\n#pragma unroll\n for (unsigned int i = 0; i < n_iterations; ++i) {\n const size_t current_place = current_stride + my_place;\n const unsigned int my_place_in = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n const unsigned int current_in = (i + 1) % 2;\n if (i < n_iterations - 1) {\n const bool valid_load =\n my_place_in < tile_length && warp_id_in_tile * n_iterations + i + 1 < tile_height;\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n if (valid_load) {\n const size_t ld_offset = current_stride + my_place_in + stride * (nvec_out + j);\n in[current_in][j].load_from(my_input_tile, ld_offset);\n if constexpr (IS_DACT) {\n act_in[current_in][j].load_from(my_act_input_tile, ld_offset);\n }\n } else {\n in[current_in][j].clear();\n if constexpr (IS_DACT) {\n act_in[current_in][j].clear();\n }\n }\n }\n }\n CVec after_dact[nvec_out]; // NOLINT(*)\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n#pragma unroll\n for (unsigned int k = 0; k < nvec_in; ++k) {\n if constexpr (IS_DACT) {\n after_dact[j].data.elt[k] = CType(in[current_in ^ 1][j].data.elt[k]) *\n OP(act_in[current_in ^ 1][j].data.elt[k], {});\n } else if constexpr (IS_ACT) {\n after_dact[j].data.elt[k] = OP(in[current_in ^ 1][j].data.elt[k], {});\n } else {\n after_dact[j].data.elt[k] = CType(in[current_in ^ 1][j].data.elt[k]);\n }\n }\n }\n const int dbias_shfl_src_lane =\n (my_id_in_warp + i + warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n constexpr bool IS_FULL_TILE = false;\n const bool valid_store =\n (my_place < tile_length) && (warp_id_in_tile * n_iterations + i < tile_height);\n\n cast_and_transpose_regs(after_dact, out_space[i], partial_dbias,\n my_output_c_tile, current_place, stride, scale,\n amax, dbias_shfl_src_lane, valid_store);\n\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += nvec_out * stride;\n current_row += nvec_out;\n }\n\n for (unsigned int i = 0; i < nvec_in; ++i) {\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * nvec_in;\n for (unsigned int j = 0; warp_id_in_tile * n_iterations + j < tile_length; ++j) {\n const bool valid_store = my_place < tile_height;\n if (valid_store) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile,\n current_stride + my_place);\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * nvec_in;\n }\n __syncthreads();\n }\n\n if constexpr (IS_DBIAS) {\n my_dbias_scratch[threadIdx.x] = partial_dbias;\n __syncthreads();\n if (warp_id_in_tile == 0) {\n#pragma unroll\n for (unsigned int i = 1; i < n_warps_per_tile; ++i) {\n CVec tmp = my_dbias_scratch[threadIdx.x + i * THREADS_PER_WARP];\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n partial_dbias.data.elt[j] += tmp.data.elt[j];\n }\n }\n if (my_id_in_warp < tile_length) {\n partial_dbias.store_to(my_partial_dbias_tile, my_id_in_warp);\n }\n }\n }\n\n // Reduce amax over block\n if (param.amax != nullptr) {\n amax = reduce_max(amax, warp_id);\n if (threadIdx.x == 0) {\n static_assert(std::is_same::value);\n atomicMaxFloat(param.amax, amax);\n }\n }\n\n // Update scale-inverse\n if (blockIdx.x == 0 && threadIdx.x == 0 && param.scale_inv != nullptr) {\n reciprocal(param.scale_inv, scale);\n }\n}\n\nstatic const char *ActTypeToString[] = {\n \"none\", // 0\n \"sigmoid\", // 1\n \"dsigmoid\", // 2\n \"gelu\", // 3\n \"dgelu\", // 4\n \"qgelu\", // 5\n \"dqgelu\", // 6\n \"silu\", // 7\n \"dsilu\", // 8\n \"relu\", // 9\n \"drelu\", // 10\n \"srelu\", // 11\n \"dsrelu\" // 12\n};\n\ntemplate \nconstexpr int get_activation_type() {\n constexpr decltype(OP) ActivationList[] = {\n nullptr, // 0\n &sigmoid, // 1\n &dsigmoid, // 2\n &gelu, // 3\n &dgelu, // 4\n &qgelu, // 5\n &dqgelu, // 6\n &silu, // 7\n &dsilu, // 8\n &relu, // 9\n &drelu, // 10\n &srelu, // 11\n &dsrelu // 12\n };\n#pragma unroll\n for (int i = 0; i < sizeof(ActivationList) / sizeof(ActivationList[0]); ++i) {\n if (OP == ActivationList[i]) {\n return i;\n }\n }\n return 0;\n}\n\ntemplate \nvoid cast_transpose_fused(const Tensor &input, const Tensor *act_input, Tensor *output,\n Tensor *dbias, Tensor *workspace, hipStream_t stream) {\n // Check tensors, unless querying dbias workspace\n if (!IS_DBIAS || workspace->data.dptr != nullptr) {\n CheckInputTensor(input, \"cast_transpose_fused_input\");\n CheckOutputTensor(*output, \"output\");\n if constexpr (IS_DBIAS) {\n NVTE_CHECK(dbias != nullptr && dbias->has_data());\n CheckOutputTensor(*dbias, \"dbias\");\n }\n if constexpr (IS_DACT) {\n NVTE_CHECK(act_input != nullptr && act_input->has_data());\n CheckInputTensor(*act_input, \"act_input\");\n }\n }\n\n // Check that inputs and outputs are available\n NVTE_CHECK(input.has_data(), \"Input is not allocated\");\n NVTE_CHECK(output->has_data(), \"Output rowwise data is not allocated\");\n NVTE_CHECK(output->has_columnwise_data(), \"Output columnwise data is not allocated\");\n\n // Flatten tensor to 2D\n NVTE_CHECK(input.data.shape == output->data.shape,\n \"Input and output shapes do not match (input=\", input.data.shape,\n \", output=\", output->data.shape);\n const size_t row_length = input.flat_last_dim();\n const size_t num_rows = input.flat_first_dim();\n\n // Check that cast and transposed output data matches\n NVTE_CHECK(output->data.dtype == output->columnwise_data.dtype,\n \"Cast and transposed output types must match.\");\n NVTE_CHECK(output->scale_inv.dptr == output->columnwise_scale_inv.dptr,\n \"Cast and transposed outputs need to share scale-inverse tensor.\");\n\n if constexpr (IS_DBIAS) {\n NVTE_CHECK(dbias->data.dtype == input.data.dtype, \"DBias must have the same type as input.\");\n NVTE_CHECK(dbias->data.shape == std::vector{row_length}, \"Wrong shape of DBias.\");\n }\n if constexpr (IS_DACT) {\n NVTE_CHECK(input.dtype() == act_input->dtype(), \"Types of both inputs must match.\");\n NVTE_CHECK(input.data.shape == act_input->data.shape, \"Shapes of both inputs must match.\");\n }\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.dtype(), InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n output->dtype(), OutputType, using InputType2 = InputType;\n using Param = CTDBiasDActParam;\n\n constexpr int itype_size = sizeof(InputType);\n constexpr int itype2_size = sizeof(InputType2);\n constexpr int otype_size = sizeof(OutputType);\n\n const bool aligned =\n (row_length % THREADS_PER_WARP == 0) && (num_rows % THREADS_PER_WARP == 0);\n const bool jit_compiled = aligned && rtc::is_enabled();\n\n size_t load_size = (IS_DACT ? desired_load_size_dact : desired_load_size);\n size_t store_size = (IS_DACT ? desired_store_size_dact : desired_store_size);\n size_t num_blocks;\n\n if (jit_compiled) {\n // Pick kernel config\n std::vector kernel_configs;\n kernel_configs.reserve(16);\n const size_t sm_count = static_cast(cuda::sm_count());\n auto add_config = [&](size_t load_size_config, size_t store_size_config) {\n kernel_configs.emplace_back(row_length, num_rows, itype_size, itype2_size, otype_size,\n load_size_config, store_size_config, sm_count, IS_DACT);\n };\n add_config(8, 8);\n add_config(4, 8);\n add_config(8, 4);\n add_config(4, 4);\n add_config(2, 8);\n add_config(8, 2);\n add_config(2, 4);\n add_config(4, 2);\n add_config(2, 2);\n add_config(1, 8);\n add_config(8, 1);\n add_config(1, 4);\n add_config(4, 1);\n add_config(1, 2);\n add_config(2, 1);\n add_config(1, 1);\n\n // Select the kernel configuration with the lowest cost\n const auto &kernel_config =\n *std::min_element(kernel_configs.begin(), kernel_configs.end());\n NVTE_CHECK(kernel_config.valid, \"invalid kernel config\");\n load_size = kernel_config.load_size;\n store_size = kernel_config.store_size;\n num_blocks = kernel_config.num_blocks;\n }\n\n const size_t nvec_in = load_size / itype_size;\n const size_t nvec_out = store_size / otype_size;\n const size_t tile_size_x = nvec_in * threads_per_warp;\n const size_t tile_size_y = nvec_out * threads_per_warp;\n const size_t num_tiles_x = DIVUP(row_length, tile_size_x);\n const size_t num_tiles_y = DIVUP(num_rows, tile_size_y);\n const size_t num_tiles = num_tiles_x * num_tiles_y;\n\n NVTE_CHECK(row_length % nvec_in == 0, \"Unsupported shape.\");\n NVTE_CHECK(num_rows % nvec_out == 0, \"Unsupported shape.\");\n\n if (!jit_compiled) {\n num_blocks = DIVUP(num_tiles * n_warps_per_tile, n_warps_per_block);\n } if constexpr (IS_DBIAS) {\n // Check workspace size\n populate_cast_transpose_dbias_workspace_config(*output, workspace, nvec_out);\n if (workspace->data.dptr == nullptr) {\n return;\n }\n }\n\n size_t VecOutputTypeSize;\n switch (nvec_out) {\n case 1:\n VecOutputTypeSize = sizeof(Vec);\n break;\n case 2:\n VecOutputTypeSize = sizeof(Vec);\n break;\n case 4:\n VecOutputTypeSize = sizeof(Vec);\n break;\n case 8:\n VecOutputTypeSize = sizeof(Vec);\n break;\n } size_t shared_size_transpose = cast_transpose_num_threads / n_warps_per_tile *\n (threads_per_warp + 1) * VecOutputTypeSize;\n\n if constexpr (IS_DBIAS) {\n size_t VecComputeTypeSize;\n switch (nvec_in) {\n case 1:\n VecComputeTypeSize = sizeof(Vec);\n break;\n case 2:\n VecComputeTypeSize = sizeof(Vec);\n break;\n case 4:\n VecComputeTypeSize = sizeof(Vec);\n break;\n case 8:\n VecComputeTypeSize = sizeof(Vec);\n break;\n }\n const size_t shared_size_dbias = cast_transpose_num_threads * VecComputeTypeSize;\n if (shared_size_transpose < shared_size_dbias) {\n shared_size_transpose = shared_size_dbias;\n }\n }\n\n Param param;\n param.input = reinterpret_cast(input.data.dptr);\n param.output_c = reinterpret_cast(output->data.dptr);\n param.output_t = reinterpret_cast(output->columnwise_data.dptr);\n param.scale_ptr = reinterpret_cast(output->scale.dptr);\n param.amax = reinterpret_cast(output->amax.dptr);\n param.scale_inv = reinterpret_cast(output->scale_inv.dptr);\n if constexpr (IS_DBIAS) {\n param.workspace = reinterpret_cast(workspace->data.dptr);\n } if constexpr (IS_DACT) {\n param.act_input = reinterpret_cast(act_input->data.dptr);\n }\n\n // Runtime-compiled tuned kernel\n if (jit_compiled) {\n constexpr const char *itype_name = TypeInfo::name;\n constexpr const char *itype2_name = TypeInfo::name;\n constexpr const char *otype_name = TypeInfo::name;\n\n int actType = 0;\n if constexpr (IS_DACT || IS_ACT) {\n actType = get_activation_type();\n }\n\n // Compile NVRTC kernel if needed and launch\n auto &rtc_manager = rtc::KernelManager::instance();\n const std::string kernel_label = concat_strings(\n \"cast_transpose_fusion\"\n \",itype=\",\n itype_name, \",itype2=\", itype2_name, \",otype=\", otype_name,\n \",load_size=\", load_size, \",store_size=\", store_size, \",IS_DBIAS=\", IS_DBIAS,\n \",IS_DACT=\", IS_DACT, \",IS_ACT=\", IS_ACT,\n \",activationType=\", ActTypeToString[actType]);\n\n if (!rtc_manager.is_compiled(kernel_label)) {\n std::string code = string_code_transpose_rtc_cast_transpose_fusion_cu;\n code = regex_replace(code, \"__ITYPE__\", itype_name);\n code = regex_replace(code, \"__ITYPE2__\", itype2_name);\n code = regex_replace(code, \"__OTYPE__\", otype_name);\n code = regex_replace(code, \"__LOAD_SIZE__\", load_size);\n code = regex_replace(code, \"__STORE_SIZE__\", store_size);\n code = regex_replace(code, \"__WARPS_PER_TILE__\", n_warps_per_tile);\n code = regex_replace(code, \"__BLOCK_SIZE__\", cast_transpose_num_threads);\n code = regex_replace(code, \"__IS_DBIAS__\", IS_DBIAS);\n code = regex_replace(code, \"__IS_DACT__\", IS_DACT);\n code = regex_replace(code, \"__IS_ACT__\", IS_ACT);\n code = regex_replace(code, \"__ACTIVATION_TYPE__\", actType);\n\n rtc_manager.compile(\n kernel_label, \"cast_transpose_fusion_kernel_optimized\", code,\n \"transformer_engine/common/transpose/rtc/cast_transpose_fusion.cu\");\n }\n\n rtc_manager.set_cache_config(kernel_label, hipFuncCachePreferShared);\n\n rtc_manager.launch(kernel_label, num_blocks, cast_transpose_num_threads,\n shared_size_transpose, stream, param, row_length, num_rows,\n num_tiles);\n } else { // Statically-compiled general kernel\n constexpr size_t load_size = IS_DACT ? desired_load_size_dact : desired_load_size;\n constexpr size_t store_size = IS_DACT ? desired_store_size_dact : desired_store_size;\n constexpr size_t nvec_in = load_size / itype_size;\n constexpr size_t nvec_out = store_size / otype_size;\n\n NVTE_CHECK(row_length % nvec_in == 0, \"Unsupported shape.\");\n NVTE_CHECK(num_rows % nvec_out == 0, \"Unsupported shape.\");\n\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(cast_transpose_fused_kernel_notaligned),\n hipFuncAttributePreferredSharedMemoryCarveout, 100));\n cast_transpose_fused_kernel_notaligned\n <<>>(\n param, row_length, num_rows, num_tiles);\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n\n if constexpr (IS_DBIAS) {\n reduce_dbias(*workspace, dbias, row_length, num_rows, nvec_out, stream);\n }); // NOLINT(*)\n ); // NOLINT(*)\n}\n\ntemplate \n__global__ void __launch_bounds__(cast_transpose_num_threads)\n dgated_act_cast_transpose_kernel(const IType *const input, const IType *const act_input,\n OType *const output_c, OType *const output_t,\n const CType *const scale_ptr, CType *const amax,\n CType *const scale_inv, const size_t row_length,\n const size_t num_rows, const size_t num_tiles) {\n using IVec = Vec;\n using OVec = Vec;\n using CVec = Vec;\n\n extern __shared__ char scratch[];\n\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;\n const size_t num_tiles_x = row_length / (nvec_in * THREADS_PER_WARP);\n const size_t tile_id =\n blockIdx.x * blockDim.x / (THREADS_PER_WARP * n_warps_per_tile) + warp_id / n_warps_per_tile;\n if (tile_id >= num_tiles) {\n return;\n }\n\n const size_t tile_id_x = tile_id % num_tiles_x;\n const size_t tile_id_y = tile_id / num_tiles_x;\n\n const IType *const my_input_tile =\n input + (tile_id_x * nvec_in + tile_id_y * row_length * nvec_out) * THREADS_PER_WARP;\n const IType *const my_act_input_tile =\n act_input + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP;\n const IType *const my_gate_input_tile =\n act_input + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP +\n row_length;\n OType *const my_output_c_tile_0 =\n output_c + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP;\n OType *const my_output_c_tile_1 =\n output_c + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP +\n row_length;\n OType *const my_output_t_tile_0 =\n output_t + (tile_id_y * nvec_out + tile_id_x * num_rows * nvec_in) * THREADS_PER_WARP;\n OType *const my_output_t_tile_1 =\n output_t + (tile_id_y * nvec_out + tile_id_x * num_rows * nvec_in) * THREADS_PER_WARP +\n row_length * num_rows;\n OVec *const my_scratch =\n reinterpret_cast(scratch) +\n (my_id_in_warp + warp_id / n_warps_per_tile * THREADS_PER_WARP) * (THREADS_PER_WARP + 1);\n\n IVec in[2][nvec_out];\n IVec act_in[2][nvec_out];\n IVec gate_in[2][nvec_out];\n const unsigned int warp_id_in_tile = warp_id % n_warps_per_tile;\n constexpr unsigned int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n OVec out_space_0[n_iterations][nvec_in];\n OVec out_space_1[n_iterations][nvec_in];\n\n const size_t stride = row_length / nvec_in;\n const size_t output_stride = num_rows / nvec_out;\n size_t current_stride = warp_id_in_tile * n_iterations * nvec_out * stride;\n unsigned int my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n const size_t stride2 = 2 * row_length / nvec_in;\n size_t current_stride2 = warp_id_in_tile * n_iterations * nvec_out * stride2;\n CType max = 0;\n const CType scale = scale_ptr != nullptr ? *scale_ptr : 1;\n\n CVec partial_dbias;\n\n#pragma unroll\n for (unsigned int i = 0; i < nvec_out; ++i) {\n in[0][i].load_from(my_input_tile, current_stride + my_place + stride * i);\n act_in[0][i].load_from(my_act_input_tile, current_stride2 + my_place + stride2 * i);\n gate_in[0][i].load_from(my_gate_input_tile, current_stride2 + my_place + stride2 * i);\n }\n#pragma unroll\n for (unsigned int i = 0; i < n_iterations; ++i) {\n const size_t current_place = current_stride2 + my_place;\n const unsigned int my_place_in = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n const unsigned int current_in = (i + 1) % 2;\n if (i < n_iterations - 1) {\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n in[current_in][j].load_from(my_input_tile,\n current_stride + my_place_in + stride * (nvec_out + j));\n act_in[current_in][j].load_from(my_act_input_tile,\n current_stride2 + my_place_in + stride2 * (nvec_out + j));\n gate_in[current_in][j].load_from(my_gate_input_tile,\n current_stride2 + my_place_in + stride2 * (nvec_out + j));\n }\n }\n CVec after_dact[nvec_out]; // NOLINT(*)\n CVec after_dgate[nvec_out]; // NOLINT(*)\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n#pragma unroll\n for (unsigned int k = 0; k < nvec_in; ++k) {\n after_dact[j].data.elt[k] = OP1(act_in[current_in ^ 1][j].data.elt[k], {}) *\n CType(in[current_in ^ 1][j].data.elt[k]) *\n CType(gate_in[current_in ^ 1][j].data.elt[k]);\n after_dgate[j].data.elt[k] = CType(in[current_in ^ 1][j].data.elt[k]) *\n OP2(act_in[current_in ^ 1][j].data.elt[k], {});\n }\n }\n OVec out_trans_0[nvec_in]; // NOLINT(*)\n OVec out_trans_1[nvec_in]; // NOLINT(*)\n\n constexpr bool IS_DBIAS = false;\n constexpr bool IS_FULL_TILE = true;\n constexpr bool valid_store = true;\n constexpr int dbias_shfl_src_lane = 0;\n\n cast_and_transpose_regs(after_dact, out_trans_0, partial_dbias,\n my_output_c_tile_0, current_place, stride2,\n scale, max, dbias_shfl_src_lane, valid_store);\n\n cast_and_transpose_regs(after_dgate, out_trans_1, partial_dbias,\n my_output_c_tile_1, current_place, stride2,\n scale, max, dbias_shfl_src_lane, valid_store);\n\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n out_space_0[i][j].data.vec = out_trans_0[j].data.vec;\n out_space_1[i][j].data.vec = out_trans_1[j].data.vec;\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += nvec_out * stride;\n current_stride2 += nvec_out * stride2;\n }\n\n for (unsigned int i = 0; i < nvec_in; ++i) {\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space_0[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * nvec_in;\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile_0,\n current_stride + my_place);\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * nvec_in;\n }\n __syncthreads();\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space_1[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * nvec_in;\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile_1,\n current_stride + my_place);\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * nvec_in;\n }\n __syncthreads();\n }\n\n // Reduce amax over block\n if (amax != nullptr) {\n max = reduce_max(max, warp_id);\n if (threadIdx.x == 0) {\n static_assert(std::is_same::value);\n atomicMaxFloat(amax, max);\n }\n }\n\n // Update scale-inverse\n if (blockIdx.x == 0 && threadIdx.x == 0 && scale_inv != nullptr) {\n reciprocal(scale_inv, scale);\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(cast_transpose_num_threads)\n dgated_act_cast_transpose_kernel_notaligned(const IType *const input,\n const IType *const act_input, OType *const output_c,\n OType *const output_t, const CType *const scale_ptr,\n CType *const amax, CType *const scale_inv,\n const size_t row_length, const size_t num_rows,\n const size_t num_tiles) {\n using IVec = Vec;\n using OVec = Vec;\n using CVec = Vec;\n\n extern __shared__ char scratch[];\n\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;\n const size_t num_tiles_x =\n (row_length + nvec_in * THREADS_PER_WARP - 1) / (nvec_in * THREADS_PER_WARP);\n const size_t tile_id =\n blockIdx.x * blockDim.x / (THREADS_PER_WARP * n_warps_per_tile) + warp_id / n_warps_per_tile;\n if (tile_id >= num_tiles) return;\n const size_t tile_id_x = tile_id % num_tiles_x;\n const size_t tile_id_y = tile_id / num_tiles_x;\n\n const IType *const my_input_tile =\n input + (tile_id_x * nvec_in + tile_id_y * row_length * nvec_out) * THREADS_PER_WARP;\n const IType *const my_act_input_tile =\n act_input + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP;\n const IType *const my_gate_input_tile =\n act_input + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP +\n row_length;\n OType *const my_output_c_tile_0 =\n output_c + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP;\n OType *const my_output_c_tile_1 =\n output_c + (tile_id_x * nvec_in + tile_id_y * row_length * 2 * nvec_out) * THREADS_PER_WARP +\n row_length;\n OType *const my_output_t_tile_0 =\n output_t + (tile_id_y * nvec_out + tile_id_x * num_rows * nvec_in) * THREADS_PER_WARP;\n OType *const my_output_t_tile_1 =\n output_t + (tile_id_y * nvec_out + tile_id_x * num_rows * nvec_in) * THREADS_PER_WARP +\n row_length * num_rows;\n const size_t stride = row_length / nvec_in;\n const size_t stride2 = 2 * row_length / nvec_in;\n const size_t output_stride = num_rows / nvec_out;\n const size_t row_length_rest = stride - tile_id_x * THREADS_PER_WARP;\n const size_t row_height_rest = output_stride - tile_id_y * THREADS_PER_WARP;\n const unsigned int tile_length =\n row_length_rest > THREADS_PER_WARP ? THREADS_PER_WARP : row_length_rest;\n const unsigned int tile_height =\n row_height_rest > THREADS_PER_WARP ? THREADS_PER_WARP : row_height_rest;\n\n OVec *const my_scratch =\n reinterpret_cast(scratch) +\n (my_id_in_warp + warp_id / n_warps_per_tile * THREADS_PER_WARP) * (THREADS_PER_WARP + 1);\n\n IVec in[2][nvec_out];\n IVec act_in[2][nvec_out];\n IVec gate_in[2][nvec_out];\n const unsigned int warp_id_in_tile = warp_id % n_warps_per_tile;\n constexpr unsigned int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n OVec out_space_0[n_iterations][nvec_in];\n OVec out_space_1[n_iterations][nvec_in];\n\n size_t current_stride = warp_id_in_tile * n_iterations * nvec_out * stride;\n size_t current_stride2 = warp_id_in_tile * n_iterations * nvec_out * stride2;\n unsigned int my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n CType max = 0;\n const CType scale = scale_ptr != nullptr ? *scale_ptr : 1;\n\n CVec partial_dbias;\n\n {\n const bool valid_load = my_place < tile_length && warp_id_in_tile * n_iterations < tile_height;\n#pragma unroll\n for (unsigned int i = 0; i < nvec_out; ++i) {\n if (valid_load) {\n in[0][i].load_from(my_input_tile, current_stride + my_place + stride * i);\n act_in[0][i].load_from(my_act_input_tile, current_stride2 + my_place + stride2 * i);\n gate_in[0][i].load_from(my_gate_input_tile, current_stride2 + my_place + stride2 * i);\n } else {\n in[0][i].clear();\n act_in[0][i].clear();\n gate_in[0][i].clear();\n }\n }\n }\n#pragma unroll\n for (unsigned int i = 0; i < n_iterations; ++i) {\n const size_t current_place = current_stride2 + my_place;\n const unsigned int my_place_in = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n const unsigned int current_in = (i + 1) % 2;\n if (i < n_iterations - 1) {\n {\n const bool valid_load =\n my_place_in < tile_length && warp_id_in_tile * n_iterations + i + 1 < tile_height;\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n if (valid_load) {\n in[current_in][j].load_from(my_input_tile,\n current_stride + my_place_in + stride * (nvec_out + j));\n act_in[current_in][j].load_from(\n my_act_input_tile, current_stride2 + my_place_in + stride2 * (nvec_out + j));\n gate_in[current_in][j].load_from(\n my_gate_input_tile, current_stride2 + my_place_in + stride2 * (nvec_out + j));\n } else {\n in[current_in][j].clear();\n act_in[current_in][j].clear();\n gate_in[current_in][j].clear();\n }\n }\n }\n }\n CVec after_dact[nvec_out]; // NOLINT(*)\n CVec after_dgate[nvec_out]; // NOLINT(*)\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n#pragma unroll\n for (unsigned int k = 0; k < nvec_in; ++k) {\n after_dact[j].data.elt[k] = OP1(act_in[current_in ^ 1][j].data.elt[k], {}) *\n CType(in[current_in ^ 1][j].data.elt[k]) *\n CType(gate_in[current_in ^ 1][j].data.elt[k]);\n after_dgate[j].data.elt[k] = CType(in[current_in ^ 1][j].data.elt[k]) *\n OP2(act_in[current_in ^ 1][j].data.elt[k], {});\n }\n }\n OVec out_trans_0[nvec_in]; // NOLINT(*)\n OVec out_trans_1[nvec_in]; // NOLINT(*)\n\n constexpr bool IS_DBIAS = false;\n constexpr bool IS_FULL_TILE = false;\n constexpr int dbias_shfl_src_lane = 0;\n const bool valid_store =\n (my_place < tile_length) && (warp_id_in_tile * n_iterations + i < tile_height);\n\n cast_and_transpose_regs(after_dact, out_trans_0, partial_dbias,\n my_output_c_tile_0, current_place, stride2,\n scale, max, dbias_shfl_src_lane, valid_store);\n cast_and_transpose_regs(after_dgate, out_trans_1, partial_dbias,\n my_output_c_tile_1, current_place, stride2,\n scale, max, dbias_shfl_src_lane, valid_store);\n\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n out_space_0[i][j].data.vec = out_trans_0[j].data.vec;\n out_space_1[i][j].data.vec = out_trans_1[j].data.vec;\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += nvec_out * stride;\n current_stride2 += nvec_out * stride2;\n }\n\n for (unsigned int i = 0; i < nvec_in; ++i) {\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space_0[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * nvec_in;\n for (unsigned int j = 0; warp_id_in_tile * n_iterations + j < tile_length; ++j) {\n const bool valid_store = my_place < tile_height;\n if (valid_store) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile_0,\n current_stride + my_place);\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * nvec_in;\n }\n __syncthreads();\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space_1[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * nvec_in;\n for (unsigned int j = 0; warp_id_in_tile * n_iterations + j < tile_length; ++j) {\n const bool valid_store = my_place < tile_height;\n if (valid_store) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile_1,\n current_stride + my_place);\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * nvec_in;\n }\n __syncthreads();\n }\n\n // Reduce amax over block\n if (amax != nullptr) {\n max = reduce_max(max, warp_id);\n if (threadIdx.x == 0) {\n static_assert(std::is_same::value);\n atomicMaxFloat(amax, max);\n }\n }\n\n // Update scale-inverse\n if (blockIdx.x == 0 && threadIdx.x == 0 && scale_inv != nullptr) {\n reciprocal(scale_inv, scale);\n }\n}\n\ntemplate \nvoid dgated_act_cast_transpose(const Tensor &input, const Tensor &gated_act_input, Tensor *output,\n hipStream_t stream) {\n CheckInputTensor(input, \"dgated_act_cast_transpose_input\");\n CheckInputTensor(gated_act_input, \"dgated_act_cast_transpose_gated_act_input\");\n CheckOutputTensor(*output, \"dgated_act_cast_transpose_output\");\n\n NVTE_CHECK(input.data.shape.size() == 2, \"Input must have 2 dimensions.\");\n NVTE_CHECK(gated_act_input.data.shape.size() == 2, \"Input must have 2 dimensions.\");\n NVTE_CHECK(output->has_data() && output->has_columnwise_data(),\n \"Both rowwise and columnwise data need to be allocated.\");\n NVTE_CHECK(output->data.shape.size() == 2, \"C output must have 2 dimensions.\");\n NVTE_CHECK(output->columnwise_data.shape.size() == 2, \"T output must have 2 dimensions.\");\n const size_t row_length = input.data.shape[1];\n const size_t num_rows = input.data.shape[0];\n\n NVTE_CHECK(gated_act_input.data.shape[0] == num_rows, \"Wrong dimension of output.\");\n NVTE_CHECK(gated_act_input.data.shape[1] == row_length * 2, \"Wrong dimension of output.\");\n NVTE_CHECK(output->data.shape[0] == num_rows, \"Wrong dimension of output.\");\n NVTE_CHECK(output->data.shape[1] == row_length * 2, \"Wrong dimension of output.\");\n NVTE_CHECK(output->columnwise_data.shape[0] == row_length * 2, \"Wrong dimension of T output.\");\n NVTE_CHECK(output->columnwise_data.shape[1] == num_rows, \"Wrong dimension of T output.\");\n\n NVTE_CHECK(input.data.dtype == gated_act_input.data.dtype, \"Types of both inputs must match.\");\n\n NVTE_CHECK(output->data.dtype == output->columnwise_data.dtype,\n \"C and T outputs need to have the same type.\");\n NVTE_CHECK(output->scale_inv.dptr == output->columnwise_scale_inv.dptr,\n \"C and T outputs need to share scale inverse tensor.\");\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.dtype(), InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n output->dtype(), OutputType, using InputType2 = InputType;\n /* dact fusion kernel uses more registers */\n constexpr int desired_load_size_dact = 4;\n constexpr int desired_store_size_dact = 4; constexpr int itype_size = sizeof(InputType);\n constexpr int otype_size = sizeof(OutputType);\n constexpr int nvec_in = desired_load_size_dact / itype_size;\n constexpr int nvec_out = desired_store_size_dact / otype_size;\n\n NVTE_CHECK(row_length % nvec_in == 0, \"Unsupported shape.\");\n NVTE_CHECK(num_rows % nvec_out == 0, \"Unsupported shape.\");\n const size_t n_tiles =\n DIVUP(row_length, static_cast(nvec_in * THREADS_PER_WARP)) *\n DIVUP(num_rows, static_cast(nvec_out * THREADS_PER_WARP));\n const size_t n_warps_per_block = cast_transpose_num_threads / THREADS_PER_WARP;\n const size_t n_blocks = DIVUP(n_tiles * n_warps_per_tile, n_warps_per_block);\n\n const bool full_tile = row_length % (nvec_in * THREADS_PER_WARP) == 0 &&\n num_rows % (nvec_out * THREADS_PER_WARP) == 0;\n const size_t shmem_size = cast_transpose_num_threads / n_warps_per_tile *\n (THREADS_PER_WARP + 1) * sizeof(Vec);\n if (full_tile) {\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(dgated_act_cast_transpose_kernel),\n hipFuncAttributePreferredSharedMemoryCarveout, 100));\n\n dgated_act_cast_transpose_kernel\n <<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(gated_act_input.data.dptr),\n reinterpret_cast(output->data.dptr),\n reinterpret_cast(output->columnwise_data.dptr),\n reinterpret_cast(output->scale.dptr),\n reinterpret_cast(output->amax.dptr),\n reinterpret_cast(output->scale_inv.dptr), row_length, num_rows,\n n_tiles);\n NVTE_CHECK_CUDA(hipGetLastError());\n } else {\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(dgated_act_cast_transpose_kernel_notaligned),\n hipFuncAttributePreferredSharedMemoryCarveout, 100));\n dgated_act_cast_transpose_kernel_notaligned\n <<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(gated_act_input.data.dptr),\n reinterpret_cast(output->data.dptr),\n reinterpret_cast(output->columnwise_data.dptr),\n reinterpret_cast(output->scale.dptr),\n reinterpret_cast(output->amax.dptr),\n reinterpret_cast(output->scale_inv.dptr), row_length, num_rows,\n n_tiles);\n NVTE_CHECK_CUDA(hipGetLastError());\n }); // NOLINT(*)\n ); // NOLINT(*)\n}\n\n// Explicit template instantiation\ntemplate void cast_transpose_fused(\n const Tensor &, const Tensor *, Tensor *, Tensor *, Tensor *, hipStream_t);\n#define NVTE_INSTANTIATE_ACTIVATION(op) \\\n template void cast_transpose_fused>( \\\n const Tensor &, const Tensor *, Tensor *, Tensor *, Tensor *, hipStream_t); \\\n template void cast_transpose_fused>( \\\n const Tensor &, const Tensor *, Tensor *, Tensor *, Tensor *, hipStream_t);\nNVTE_INSTANTIATE_ACTIVATION(relu);\nNVTE_INSTANTIATE_ACTIVATION(srelu);\nNVTE_INSTANTIATE_ACTIVATION(gelu);\nNVTE_INSTANTIATE_ACTIVATION(qgelu);\nNVTE_INSTANTIATE_ACTIVATION(silu);\n#undef NVTE_INSTANTIATE_ACTIVATION\n\n} // namespace detail\n\n} // namespace transformer_engine\n\nusing ComputeType = typename transformer_engine::fp32;\n\nvoid nvte_cast_transpose_dbias(const NVTETensor input, NVTETensor output, NVTETensor dbias,\n NVTETensor workspace, hipStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose_dbias);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n constexpr bool IS_DBIAS = true;\n constexpr bool IS_DACT = false;\n constexpr bool IS_ACT = false;\n\n constexpr const NVTETensor activation_input = nullptr;\n\n cast_transpose_fused(\n *convertNVTETensorCheck(input), convertNVTETensor(activation_input),\n convertNVTETensor(output), convertNVTETensor(dbias), convertNVTETensor(workspace), stream);\n}\n\nvoid nvte_cast_transpose_dbias_dgelu(const NVTETensor input, const NVTETensor act_input,\n NVTETensor output, NVTETensor dbias, NVTETensor workspace,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose_dbias_dgelu);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n constexpr bool IS_DBIAS = true;\n constexpr bool IS_DACT = true;\n constexpr bool IS_ACT = false;\n\n cast_transpose_fused>(\n *convertNVTETensorCheck(input), convertNVTETensorCheck(act_input),\n convertNVTETensorCheck(output), convertNVTETensorCheck(dbias), convertNVTETensor(workspace),\n stream);\n}\n\nvoid nvte_cast_transpose_dbias_dsilu(const NVTETensor input, const NVTETensor silu_input,\n NVTETensor output, NVTETensor dbias, NVTETensor workspace,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose_dbias_dsilu);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n constexpr bool IS_DBIAS = true;\n constexpr bool IS_DACT = true;\n constexpr bool IS_ACT = false;\n\n cast_transpose_fused>(\n *convertNVTETensorCheck(input), convertNVTETensorCheck(silu_input),\n convertNVTETensorCheck(output), convertNVTETensorCheck(dbias), convertNVTETensor(workspace),\n stream);\n}\n\nvoid nvte_cast_transpose_dbias_drelu(const NVTETensor input, const NVTETensor relu_input,\n NVTETensor output, NVTETensor dbias, NVTETensor workspace,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose_dbias_drelu);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n constexpr bool IS_DBIAS = true;\n constexpr bool IS_DACT = true;\n constexpr bool IS_ACT = false;\n\n cast_transpose_fused>(\n *convertNVTETensorCheck(input), convertNVTETensorCheck(relu_input),\n convertNVTETensorCheck(output), convertNVTETensorCheck(dbias), convertNVTETensor(workspace),\n stream);\n}\n\nvoid nvte_cast_transpose_dbias_dsrelu(const NVTETensor input, const NVTETensor srelu_input,\n NVTETensor output, NVTETensor dbias, NVTETensor workspace,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose_dbias_dsrelu);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n constexpr bool IS_DBIAS = true;\n constexpr bool IS_DACT = true;\n constexpr bool IS_ACT = false;\n\n cast_transpose_fused>(\n *convertNVTETensorCheck(input), convertNVTETensorCheck(srelu_input),\n convertNVTETensorCheck(output), convertNVTETensorCheck(dbias), convertNVTETensor(workspace),\n stream);\n}\n\nvoid nvte_cast_transpose_dbias_dqgelu(const NVTETensor input, const NVTETensor qgelu_input,\n NVTETensor output, NVTETensor dbias, NVTETensor workspace,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose_dbias_dqgelu);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n constexpr bool IS_DBIAS = true;\n constexpr bool IS_DACT = true;\n constexpr bool IS_ACT = false;\n\n cast_transpose_fused>(\n *convertNVTETensorCheck(input), convertNVTETensorCheck(qgelu_input),\n convertNVTETensorCheck(output), convertNVTETensorCheck(dbias), convertNVTETensor(workspace),\n stream);\n}\n\nvoid nvte_dgeglu_cast_transpose(const NVTETensor input, const NVTETensor gated_act_input,\n NVTETensor output, hipStream_t stream) {\n NVTE_API_CALL(nvte_dgeglu_cast_transpose);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n dgated_act_cast_transpose, gelu>(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(gated_act_input),\n convertNVTETensorCheck(output), stream);\n}\n\nvoid nvte_dswiglu_cast_transpose(const NVTETensor input, const NVTETensor swiglu_input,\n NVTETensor output, hipStream_t stream) {\n NVTE_API_CALL(nvte_dswiglu_cast_transpose);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n dgated_act_cast_transpose, silu>(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(swiglu_input),\n convertNVTETensorCheck(output), stream);\n}\n\nvoid nvte_dreglu_cast_transpose(const NVTETensor input, const NVTETensor gated_act_input,\n NVTETensor output, hipStream_t stream) {\n NVTE_API_CALL(nvte_dreglu_cast_transpose);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n dgated_act_cast_transpose, relu>(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(gated_act_input),\n convertNVTETensorCheck(output), stream);\n}\n\nvoid nvte_dsreglu_cast_transpose(const NVTETensor input, const NVTETensor gated_act_input,\n NVTETensor output, hipStream_t stream) {\n NVTE_API_CALL(nvte_dsreglu_cast_transpose);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n dgated_act_cast_transpose, srelu>(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(gated_act_input),\n convertNVTETensorCheck(output), stream);\n}\n\nvoid nvte_dqgeglu_cast_transpose(const NVTETensor input, const NVTETensor gated_act_input,\n NVTETensor output, hipStream_t stream) {\n NVTE_API_CALL(nvte_dqgeglu_cast_transpose);\n using namespace transformer_engine;\n using namespace transformer_engine::detail;\n\n dgated_act_cast_transpose, qgelu>(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(gated_act_input),\n convertNVTETensorCheck(output), stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/transpose/cast_transpose_fusion.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\n\ntemplate \ninline __device__ void transpose_regs_partial_dbias(const IVec (&in)[nvec_out],\n OVec (&out_trans)[nvec_in],\n CVec &out_dbias, // NOLINT(*)\n const CType scale_inv,\n const int dbias_shfl_src_lane) {\n using T = typename OVec::type;\n using OVecC = Vec;\n\n CVec step_dbias;\n step_dbias.clear();\n\n#pragma unroll\n for (unsigned int i = 0; i < nvec_out; ++i) {\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n const CType tmp = static_cast(in[i].data.elt[j]) * scale_inv;\n const T elt_o = in[i].data.elt[j];\n\n /* dbias: thread tile local accumulation */\n step_dbias.data.elt[j] += tmp;\n\n out_trans[j].data.elt[i] = elt_o; // thread tile transpose\n }\n }\n\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n CType elt = step_dbias.data.elt[j];\n elt = __shfl_sync(0xffffffff, elt, dbias_shfl_src_lane); // shuffle data in warp\n out_dbias.data.elt[j] += elt;\n }\n}\n\n// STUFF TO TUNE\nconstexpr unsigned int n_warps_per_tile = 4;\nconstexpr int desired_load_size = 8;\nconstexpr int desired_store_size = 8;\n\nconstexpr unsigned int max_threads_per_block = 256;\nstatic_assert(n_warps_per_tile * THREADS_PER_WARP <= max_threads_per_block);\nconstexpr unsigned int cast_transpose_num_threads = n_warps_per_tile * THREADS_PER_WARP;\n\nnamespace {\n\ntemplate \nstruct TDBiasParam {\n using InputType = IType;\n using OutputType = OType;\n using ComputeType = CType;\n const IType *input;\n OType *output_t;\n const CType *scale_inv;\n CType *workspace;\n};\n\n} // namespace\n\ntemplate \n__global__ void __launch_bounds__(cast_transpose_num_threads)\n transpose_dbias_kernel(const Param param, const size_t row_length, const size_t num_rows,\n const size_t num_tiles) {\n using IType = typename Param::InputType;\n using OType = typename Param::OutputType;\n using CType = typename Param::ComputeType;\n using IVec = Vec;\n using OVec = Vec;\n using CVec = Vec;\n\n extern __shared__ char scratch[];\n\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const unsigned int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;\n const size_t num_tiles_x = row_length / (nvec_in * THREADS_PER_WARP);\n // const size_t num_tiles_y = num_rows / (nvec * THREADS_PER_WARP);\n const size_t tile_id =\n blockIdx.x * blockDim.x / (THREADS_PER_WARP * n_warps_per_tile) + warp_id / n_warps_per_tile;\n if (tile_id >= num_tiles) return;\n const size_t tile_id_x = tile_id % num_tiles_x;\n const size_t tile_id_y = tile_id / num_tiles_x;\n\n const IType *const my_input_tile =\n param.input + (tile_id_x * nvec_in + tile_id_y * row_length * nvec_out) * THREADS_PER_WARP;\n OType *const my_output_t_tile =\n param.output_t + (tile_id_y * nvec_out + tile_id_x * num_rows * nvec_in) * THREADS_PER_WARP;\n CType *const my_partial_dbias_tile =\n param.workspace + (tile_id_x * (nvec_in * THREADS_PER_WARP) + tile_id_y * row_length);\n\n OVec *const my_scratch =\n reinterpret_cast(scratch) +\n (my_id_in_warp + warp_id / n_warps_per_tile * THREADS_PER_WARP) * (THREADS_PER_WARP + 1);\n\n CVec *const my_dbias_scratch = reinterpret_cast(scratch);\n\n IVec in[2][nvec_out];\n const unsigned int warp_id_in_tile = warp_id % n_warps_per_tile;\n constexpr unsigned int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n OVec out_space[n_iterations][nvec_in];\n CVec partial_dbias;\n\n const size_t stride = row_length / nvec_in;\n const size_t output_stride = num_rows / nvec_out;\n size_t current_stride = warp_id_in_tile * n_iterations * nvec_out * stride;\n unsigned int my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n const CType scale_inv = param.scale_inv != nullptr ? *param.scale_inv : 1;\n\n partial_dbias.clear();\n\n#pragma unroll\n for (unsigned int i = 0; i < nvec_out; ++i) {\n in[0][i].load_from(my_input_tile, current_stride + my_place + stride * i);\n }\n#pragma unroll\n for (unsigned int i = 0; i < n_iterations; ++i) {\n const unsigned int my_place_in = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n const unsigned int current_in = (i + 1) % 2;\n if (i < n_iterations - 1) {\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n in[current_in][j].load_from(my_input_tile,\n current_stride + my_place_in + stride * (nvec_out + j));\n }\n }\n OVec out_trans[nvec_in]; // NOLINT(*)\n transpose_regs_partial_dbias(\n in[current_in ^ 1], out_trans, partial_dbias, scale_inv,\n (my_id_in_warp + i + warp_id_in_tile * n_iterations) % THREADS_PER_WARP);\n\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n out_space[i][j].data.vec = out_trans[j].data.vec;\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += nvec_out * stride;\n }\n\n for (unsigned int i = 0; i < nvec_in; ++i) {\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * nvec_in;\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile,\n current_stride + my_place);\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * nvec_in;\n }\n __syncthreads();\n }\n\n my_dbias_scratch[threadIdx.x] = partial_dbias;\n __syncthreads();\n // TODO(ptredak): check if the regular reduction is better\n if (warp_id_in_tile == 0) {\n#pragma unroll\n for (unsigned int i = 1; i < n_warps_per_tile; ++i) {\n CVec tmp = my_dbias_scratch[threadIdx.x + i * THREADS_PER_WARP];\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n partial_dbias.data.elt[j] += tmp.data.elt[j];\n }\n }\n\n partial_dbias.store_to(my_partial_dbias_tile, my_id_in_warp);\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(cast_transpose_num_threads)\n transpose_dbias_kernel_notaligned(const Param param, const size_t row_length,\n const size_t num_rows, const size_t num_tiles) {\n using IType = typename Param::InputType;\n using OType = typename Param::OutputType;\n using CType = typename Param::ComputeType;\n using IVec = Vec;\n using OVec = Vec;\n using CVec = Vec;\n\n extern __shared__ char scratch[];\n\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const unsigned int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;\n const size_t num_tiles_x =\n (row_length + nvec_in * THREADS_PER_WARP - 1) / (nvec_in * THREADS_PER_WARP);\n const size_t tile_id =\n blockIdx.x * blockDim.x / (THREADS_PER_WARP * n_warps_per_tile) + warp_id / n_warps_per_tile;\n if (tile_id >= num_tiles) return;\n const size_t tile_id_x = tile_id % num_tiles_x;\n const size_t tile_id_y = tile_id / num_tiles_x;\n\n const IType *const my_input_tile =\n param.input + (tile_id_x * nvec_in + tile_id_y * row_length * nvec_out) * THREADS_PER_WARP;\n OType *const my_output_t_tile =\n param.output_t + (tile_id_y * nvec_out + tile_id_x * num_rows * nvec_in) * THREADS_PER_WARP;\n CType *const my_partial_dbias_tile =\n param.workspace + (tile_id_x * (nvec_in * THREADS_PER_WARP) + tile_id_y * row_length);\n\n const size_t stride = row_length / nvec_in;\n const size_t output_stride = num_rows / nvec_out;\n const size_t row_length_rest = stride - tile_id_x * THREADS_PER_WARP;\n const size_t row_height_rest = output_stride - tile_id_y * THREADS_PER_WARP;\n const unsigned int tile_length =\n row_length_rest > THREADS_PER_WARP ? THREADS_PER_WARP : row_length_rest;\n const unsigned int tile_height =\n row_height_rest > THREADS_PER_WARP ? THREADS_PER_WARP : row_height_rest;\n\n OVec *const my_scratch =\n reinterpret_cast(scratch) +\n (my_id_in_warp + warp_id / n_warps_per_tile * THREADS_PER_WARP) * (THREADS_PER_WARP + 1);\n\n CVec *const my_dbias_scratch = reinterpret_cast(scratch);\n\n IVec in[2][nvec_out];\n const unsigned int warp_id_in_tile = warp_id % n_warps_per_tile;\n constexpr unsigned int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n OVec out_space[n_iterations][nvec_in];\n CVec partial_dbias;\n\n size_t current_stride = warp_id_in_tile * n_iterations * nvec_out * stride;\n unsigned int my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n const CType scale_inv = param.scale_inv != nullptr ? *param.scale_inv : 1;\n\n partial_dbias.clear();\n\n {\n const bool valid_load = my_place < tile_length && warp_id_in_tile * n_iterations < tile_height;\n#pragma unroll\n for (unsigned int i = 0; i < nvec_out; ++i) {\n if (valid_load) {\n in[0][i].load_from(my_input_tile, current_stride + my_place + stride * i);\n } else {\n in[0][i].clear();\n }\n }\n }\n#pragma unroll\n for (unsigned int i = 0; i < n_iterations; ++i) {\n const unsigned int my_place_in = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n const unsigned int current_in = (i + 1) % 2;\n if (i < n_iterations - 1) {\n const bool valid_load =\n my_place_in < tile_length && warp_id_in_tile * n_iterations + i + 1 < tile_height;\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n if (valid_load) {\n in[current_in][j].load_from(my_input_tile,\n current_stride + my_place_in + stride * (nvec_out + j));\n } else {\n in[current_in][j].clear();\n }\n }\n }\n OVec out_trans[nvec_in]; // NOLINT(*)\n transpose_regs_partial_dbias(\n in[current_in ^ 1], out_trans, partial_dbias, scale_inv,\n (my_id_in_warp + i + warp_id_in_tile * n_iterations) % THREADS_PER_WARP);\n\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n out_space[i][j].data.vec = out_trans[j].data.vec;\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += nvec_out * stride;\n }\n\n for (unsigned int i = 0; i < nvec_in; ++i) {\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * nvec_in;\n for (unsigned int j = 0; warp_id_in_tile * n_iterations + j < tile_length; ++j) {\n const bool valid_store = my_place < tile_height;\n if (valid_store) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile,\n current_stride + my_place);\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * nvec_in;\n }\n __syncthreads();\n }\n\n my_dbias_scratch[threadIdx.x] = partial_dbias;\n __syncthreads();\n // TODO(ptredak): check if the regular reduction is better\n if (warp_id_in_tile == 0) {\n#pragma unroll\n for (unsigned int i = 1; i < n_warps_per_tile; ++i) {\n CVec tmp = my_dbias_scratch[threadIdx.x + i * THREADS_PER_WARP];\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n partial_dbias.data.elt[j] += tmp.data.elt[j];\n }\n }\n\n if (my_id_in_warp < tile_length) {\n partial_dbias.store_to(my_partial_dbias_tile, my_id_in_warp);\n }\n }\n}\n\nconstexpr size_t reduce_dbias_num_threads = 256;\n\ntemplate \n__global__ void __launch_bounds__(reduce_dbias_num_threads)\n reduce_dbias_kernel(OutputType *const dbias_output, const ComputeType *const dbias_partial,\n const int row_length, const int num_rows) {\n using ComputeVec = Vec;\n using OutputVec = Vec;\n\n const int thread_id = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (thread_id * nvec >= row_length) return;\n\n const ComputeType *const thread_in_base = dbias_partial + thread_id * nvec;\n OutputType *const thread_out_base = dbias_output + thread_id * nvec;\n\n const int stride_in_vec = row_length / nvec;\n\n ComputeVec ldg_vec;\n ComputeVec acc_vec;\n acc_vec.clear();\n for (int i = 0; i < num_rows; ++i) {\n ldg_vec.load_from(thread_in_base, i * stride_in_vec);\n#pragma unroll\n for (int e = 0; e < nvec; ++e) {\n acc_vec.data.elt[e] += ldg_vec.data.elt[e];\n }\n }\n\n OutputVec stg_vec;\n#pragma unroll\n for (int e = 0; e < nvec; ++e) {\n stg_vec.data.elt[e] = OutputType(acc_vec.data.elt[e]);\n }\n stg_vec.store_to(thread_out_base, 0);\n}\n\nvoid populate_transpose_dbias_workspace_config(const Tensor &input, /*cast*/\n Tensor *workspace, const int nvec_out) {\n const size_t row_length = input.data.shape[1];\n const size_t num_rows = input.data.shape[0];\n\n const size_t tile_size_y = (nvec_out * THREADS_PER_WARP);\n NVTE_CHECK(num_rows % nvec_out == 0, \"Unsupported shape.\");\n\n const size_t num_rows_partial_dbias = DIVUP(num_rows, tile_size_y);\n\n if (workspace->data.dptr == nullptr) {\n // Set workspace size\n workspace->data.shape = {num_rows_partial_dbias, row_length};\n workspace->data.dtype = DType::kFloat32;\n } else {\n // Check that workspace matches expected size\n const size_t workspace_size = get_buffer_size_bytes(\n std::accumulate(workspace->data.shape.begin(), workspace->data.shape.end(), 1,\n std::multiplies()),\n workspace->data.dtype);\n const size_t required_size =\n get_buffer_size_bytes(num_rows_partial_dbias, row_length, DType::kFloat32);\n NVTE_CHECK(workspace_size >= required_size, \"Invalid workspace (expected dims=(\",\n num_rows_partial_dbias, \",\", row_length, \"), dtype=\", to_string(DType::kFloat32),\n \"; found dims=\", workspace->data.shape,\n \", dtype=\", typeToNumBits(workspace->data.dtype), \" bits)\");\n }\n}\n\ntemplate \nvoid reduce_dbias(const Tensor &workspace, Tensor *dbias, const size_t row_length,\n const size_t num_rows, const int nvec_out, cudaStream_t stream) {\n constexpr int reduce_dbias_store_bytes = 8; // stg.64\n constexpr int reduce_dbias_nvec = reduce_dbias_store_bytes / sizeof(BiasType);\n\n NVTE_CHECK(row_length % reduce_dbias_nvec == 0, \"Unsupported shape.\");\n\n const size_t reduce_dbias_row_length = row_length;\n const size_t reduce_dbias_num_rows =\n DIVUP(num_rows, static_cast(nvec_out * THREADS_PER_WARP));\n const size_t reduce_dbias_num_blocks =\n DIVUP(row_length, reduce_dbias_num_threads * reduce_dbias_nvec);\n\n reduce_dbias_kernel\n <<>>(\n reinterpret_cast(dbias->data.dptr),\n reinterpret_cast(workspace.data.dptr), reduce_dbias_row_length,\n reduce_dbias_num_rows);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid fp8_transpose_dbias(const Tensor &input, Tensor *transposed_output, Tensor *dbias,\n Tensor *workspace, cudaStream_t stream) {\n CheckInputTensor(input, \"fp8_transpose_dbias_input\");\n CheckOutputTensor(*transposed_output, \"transposed_output\");\n CheckOutputTensor(*dbias, \"dbias\");\n\n NVTE_CHECK(input.data.shape.size() == 2, \"Input must have 2 dimensions.\");\n NVTE_CHECK(transposed_output->data.shape.size() == 2, \"T output must have 2 dimensions.\");\n const size_t row_length = input.data.shape[1];\n const size_t num_rows = input.data.shape[0];\n\n NVTE_CHECK(transposed_output->data.shape[0] == row_length, \"Wrong dimension of T output.\");\n NVTE_CHECK(transposed_output->data.shape[1] == num_rows, \"Wrong dimension of T output.\");\n\n NVTE_CHECK(transposed_output->data.dtype == input.data.dtype,\n \"T output must have the same type as input.\");\n NVTE_CHECK(dbias->data.shape == std::vector{row_length}, \"Wrong shape of DBias.\");\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n dbias->data.dtype, BiasType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(\n input.data.dtype, Type, constexpr int type_size = sizeof(Type);\n constexpr int nvec_in = desired_load_size / type_size;\n constexpr int nvec_out = desired_store_size / type_size;\n\n // Check workspace size\n populate_transpose_dbias_workspace_config(input, workspace, nvec_out);\n if (workspace->data.dptr == nullptr) { return; }\n\n NVTE_CHECK(row_length % nvec_in == 0, \"Unsupported shape.\");\n NVTE_CHECK(num_rows % nvec_out == 0, \"Unsupported shape.\");\n const size_t n_tiles =\n DIVUP(row_length, static_cast(nvec_in * THREADS_PER_WARP)) *\n DIVUP(num_rows, static_cast(nvec_out * THREADS_PER_WARP));\n const size_t n_warps_per_block = cast_transpose_num_threads / THREADS_PER_WARP;\n const size_t n_blocks = DIVUP(n_tiles * n_warps_per_tile, n_warps_per_block);\n\n const bool full_tile = row_length % (nvec_in * THREADS_PER_WARP) == 0 &&\n num_rows % (nvec_out * THREADS_PER_WARP) == 0;\n\n using ComputeType = fp32; constexpr size_t shared_size_transpose =\n cast_transpose_num_threads / n_warps_per_tile *\n (THREADS_PER_WARP + 1) * sizeof(Vec);\n constexpr size_t shared_size_dbias =\n cast_transpose_num_threads * sizeof(Vec);\n static_assert(shared_size_transpose >= shared_size_dbias);\n using Param = TDBiasParam; Param param;\n param.input = reinterpret_cast(input.data.dptr);\n param.output_t = reinterpret_cast(transposed_output->data.dptr);\n param.scale_inv =\n reinterpret_cast(transposed_output->scale_inv.dptr);\n param.workspace = reinterpret_cast(workspace->data.dptr);\n\n if (full_tile) {\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(transpose_dbias_kernel,\n cudaFuncAttributePreferredSharedMemoryCarveout,\n 100));\n transpose_dbias_kernel\n <<>>(\n param, row_length, num_rows, n_tiles);\n NVTE_CHECK_CUDA(cudaGetLastError());\n } else {\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(transpose_dbias_kernel_notaligned,\n cudaFuncAttributePreferredSharedMemoryCarveout, 100));\n transpose_dbias_kernel_notaligned\n <<>>(\n param, row_length, num_rows, n_tiles);\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n\n reduce_dbias(*workspace, dbias, row_length, num_rows, nvec_out,\n stream);); // NOLINT(*)\n ); // NOLINT(*)\n}\n\n} // namespace transformer_engine\n\nvoid nvte_fp8_transpose_dbias(const NVTETensor input, NVTETensor transposed_output,\n NVTETensor dbias, NVTETensor workspace, cudaStream_t stream) {\n NVTE_API_CALL(nvte_fp8_transpose_dbias);\n using namespace transformer_engine;\n fp8_transpose_dbias(*convertNVTETensorCheck(input), convertNVTETensor(transposed_output),\n convertNVTETensor(dbias), convertNVTETensor(workspace), stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\n\ntemplate \ninline __device__ void transpose_regs_partial_dbias(const IVec (&in)[nvec_out],\n OVec (&out_trans)[nvec_in],\n CVec &out_dbias, // NOLINT(*)\n const CType scale_inv,\n const int dbias_shfl_src_lane) {\n using T = typename OVec::type;\n using OVecC = Vec;\n\n CVec step_dbias;\n step_dbias.clear();\n\n#pragma unroll\n for (unsigned int i = 0; i < nvec_out; ++i) {\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n const CType tmp = static_cast(in[i].data.elt[j]) * scale_inv;\n const T elt_o = in[i].data.elt[j];\n\n /* dbias: thread tile local accumulation */\n step_dbias.data.elt[j] += tmp;\n\n out_trans[j].data.elt[i] = elt_o; // thread tile transpose\n }\n }\n\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n CType elt = step_dbias.data.elt[j];\n elt = __shfl_sync(0xffffffff, elt, dbias_shfl_src_lane); // shuffle data in warp\n out_dbias.data.elt[j] += elt;\n }\n}\n\n// STUFF TO TUNE\nconstexpr unsigned int n_warps_per_tile = 4;\nconstexpr int desired_load_size = 8;\nconstexpr int desired_store_size = 8;\n\nconstexpr unsigned int max_threads_per_block = 256;\nstatic_assert(n_warps_per_tile * THREADS_PER_WARP <= max_threads_per_block);\nconstexpr unsigned int cast_transpose_num_threads = n_warps_per_tile * THREADS_PER_WARP;\n\nnamespace {\n\ntemplate \nstruct TDBiasParam {\n using InputType = IType;\n using OutputType = OType;\n using ComputeType = CType;\n const IType *input;\n OType *output_t;\n const CType *scale_inv;\n CType *workspace;\n};\n\n} // namespace\n\ntemplate \n__global__ void __launch_bounds__(cast_transpose_num_threads)\n transpose_dbias_kernel(const Param param, const size_t row_length, const size_t num_rows,\n const size_t num_tiles) {\n using IType = typename Param::InputType;\n using OType = typename Param::OutputType;\n using CType = typename Param::ComputeType;\n using IVec = Vec;\n using OVec = Vec;\n using CVec = Vec;\n\n extern __shared__ char scratch[];\n\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const unsigned int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;\n const size_t num_tiles_x = row_length / (nvec_in * THREADS_PER_WARP);\n // const size_t num_tiles_y = num_rows / (nvec * THREADS_PER_WARP);\n const size_t tile_id =\n blockIdx.x * blockDim.x / (THREADS_PER_WARP * n_warps_per_tile) + warp_id / n_warps_per_tile;\n if (tile_id >= num_tiles) return;\n const size_t tile_id_x = tile_id % num_tiles_x;\n const size_t tile_id_y = tile_id / num_tiles_x;\n\n const IType *const my_input_tile =\n param.input + (tile_id_x * nvec_in + tile_id_y * row_length * nvec_out) * THREADS_PER_WARP;\n OType *const my_output_t_tile =\n param.output_t + (tile_id_y * nvec_out + tile_id_x * num_rows * nvec_in) * THREADS_PER_WARP;\n CType *const my_partial_dbias_tile =\n param.workspace + (tile_id_x * (nvec_in * THREADS_PER_WARP) + tile_id_y * row_length);\n\n OVec *const my_scratch =\n reinterpret_cast(scratch) +\n (my_id_in_warp + warp_id / n_warps_per_tile * THREADS_PER_WARP) * (THREADS_PER_WARP + 1);\n\n CVec *const my_dbias_scratch = reinterpret_cast(scratch);\n\n IVec in[2][nvec_out];\n const unsigned int warp_id_in_tile = warp_id % n_warps_per_tile;\n constexpr unsigned int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n OVec out_space[n_iterations][nvec_in];\n CVec partial_dbias;\n\n const size_t stride = row_length / nvec_in;\n const size_t output_stride = num_rows / nvec_out;\n size_t current_stride = warp_id_in_tile * n_iterations * nvec_out * stride;\n unsigned int my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n const CType scale_inv = param.scale_inv != nullptr ? *param.scale_inv : 1;\n\n partial_dbias.clear();\n\n#pragma unroll\n for (unsigned int i = 0; i < nvec_out; ++i) {\n in[0][i].load_from(my_input_tile, current_stride + my_place + stride * i);\n }\n#pragma unroll\n for (unsigned int i = 0; i < n_iterations; ++i) {\n const unsigned int my_place_in = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n const unsigned int current_in = (i + 1) % 2;\n if (i < n_iterations - 1) {\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n in[current_in][j].load_from(my_input_tile,\n current_stride + my_place_in + stride * (nvec_out + j));\n }\n }\n OVec out_trans[nvec_in]; // NOLINT(*)\n transpose_regs_partial_dbias(\n in[current_in ^ 1], out_trans, partial_dbias, scale_inv,\n (my_id_in_warp + i + warp_id_in_tile * n_iterations) % THREADS_PER_WARP);\n\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n out_space[i][j].data.vec = out_trans[j].data.vec;\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += nvec_out * stride;\n }\n\n for (unsigned int i = 0; i < nvec_in; ++i) {\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * nvec_in;\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile,\n current_stride + my_place);\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * nvec_in;\n }\n __syncthreads();\n }\n\n my_dbias_scratch[threadIdx.x] = partial_dbias;\n __syncthreads();\n // TODO(ptredak): check if the regular reduction is better\n if (warp_id_in_tile == 0) {\n#pragma unroll\n for (unsigned int i = 1; i < n_warps_per_tile; ++i) {\n CVec tmp = my_dbias_scratch[threadIdx.x + i * THREADS_PER_WARP];\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n partial_dbias.data.elt[j] += tmp.data.elt[j];\n }\n }\n\n partial_dbias.store_to(my_partial_dbias_tile, my_id_in_warp);\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(cast_transpose_num_threads)\n transpose_dbias_kernel_notaligned(const Param param, const size_t row_length,\n const size_t num_rows, const size_t num_tiles) {\n using IType = typename Param::InputType;\n using OType = typename Param::OutputType;\n using CType = typename Param::ComputeType;\n using IVec = Vec;\n using OVec = Vec;\n using CVec = Vec;\n\n extern __shared__ char scratch[];\n\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const unsigned int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;\n const size_t num_tiles_x =\n (row_length + nvec_in * THREADS_PER_WARP - 1) / (nvec_in * THREADS_PER_WARP);\n const size_t tile_id =\n blockIdx.x * blockDim.x / (THREADS_PER_WARP * n_warps_per_tile) + warp_id / n_warps_per_tile;\n if (tile_id >= num_tiles) return;\n const size_t tile_id_x = tile_id % num_tiles_x;\n const size_t tile_id_y = tile_id / num_tiles_x;\n\n const IType *const my_input_tile =\n param.input + (tile_id_x * nvec_in + tile_id_y * row_length * nvec_out) * THREADS_PER_WARP;\n OType *const my_output_t_tile =\n param.output_t + (tile_id_y * nvec_out + tile_id_x * num_rows * nvec_in) * THREADS_PER_WARP;\n CType *const my_partial_dbias_tile =\n param.workspace + (tile_id_x * (nvec_in * THREADS_PER_WARP) + tile_id_y * row_length);\n\n const size_t stride = row_length / nvec_in;\n const size_t output_stride = num_rows / nvec_out;\n const size_t row_length_rest = stride - tile_id_x * THREADS_PER_WARP;\n const size_t row_height_rest = output_stride - tile_id_y * THREADS_PER_WARP;\n const unsigned int tile_length =\n row_length_rest > THREADS_PER_WARP ? THREADS_PER_WARP : row_length_rest;\n const unsigned int tile_height =\n row_height_rest > THREADS_PER_WARP ? THREADS_PER_WARP : row_height_rest;\n\n OVec *const my_scratch =\n reinterpret_cast(scratch) +\n (my_id_in_warp + warp_id / n_warps_per_tile * THREADS_PER_WARP) * (THREADS_PER_WARP + 1);\n\n CVec *const my_dbias_scratch = reinterpret_cast(scratch);\n\n IVec in[2][nvec_out];\n const unsigned int warp_id_in_tile = warp_id % n_warps_per_tile;\n constexpr unsigned int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n OVec out_space[n_iterations][nvec_in];\n CVec partial_dbias;\n\n size_t current_stride = warp_id_in_tile * n_iterations * nvec_out * stride;\n unsigned int my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n const CType scale_inv = param.scale_inv != nullptr ? *param.scale_inv : 1;\n\n partial_dbias.clear();\n\n {\n const bool valid_load = my_place < tile_length && warp_id_in_tile * n_iterations < tile_height;\n#pragma unroll\n for (unsigned int i = 0; i < nvec_out; ++i) {\n if (valid_load) {\n in[0][i].load_from(my_input_tile, current_stride + my_place + stride * i);\n } else {\n in[0][i].clear();\n }\n }\n }\n#pragma unroll\n for (unsigned int i = 0; i < n_iterations; ++i) {\n const unsigned int my_place_in = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n const unsigned int current_in = (i + 1) % 2;\n if (i < n_iterations - 1) {\n const bool valid_load =\n my_place_in < tile_length && warp_id_in_tile * n_iterations + i + 1 < tile_height;\n#pragma unroll\n for (unsigned int j = 0; j < nvec_out; ++j) {\n if (valid_load) {\n in[current_in][j].load_from(my_input_tile,\n current_stride + my_place_in + stride * (nvec_out + j));\n } else {\n in[current_in][j].clear();\n }\n }\n }\n OVec out_trans[nvec_in]; // NOLINT(*)\n transpose_regs_partial_dbias(\n in[current_in ^ 1], out_trans, partial_dbias, scale_inv,\n (my_id_in_warp + i + warp_id_in_tile * n_iterations) % THREADS_PER_WARP);\n\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n out_space[i][j].data.vec = out_trans[j].data.vec;\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += nvec_out * stride;\n }\n\n for (unsigned int i = 0; i < nvec_in; ++i) {\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * nvec_in;\n for (unsigned int j = 0; warp_id_in_tile * n_iterations + j < tile_length; ++j) {\n const bool valid_store = my_place < tile_height;\n if (valid_store) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile,\n current_stride + my_place);\n }\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * nvec_in;\n }\n __syncthreads();\n }\n\n my_dbias_scratch[threadIdx.x] = partial_dbias;\n __syncthreads();\n // TODO(ptredak): check if the regular reduction is better\n if (warp_id_in_tile == 0) {\n#pragma unroll\n for (unsigned int i = 1; i < n_warps_per_tile; ++i) {\n CVec tmp = my_dbias_scratch[threadIdx.x + i * THREADS_PER_WARP];\n#pragma unroll\n for (unsigned int j = 0; j < nvec_in; ++j) {\n partial_dbias.data.elt[j] += tmp.data.elt[j];\n }\n }\n\n if (my_id_in_warp < tile_length) {\n partial_dbias.store_to(my_partial_dbias_tile, my_id_in_warp);\n }\n }\n}\n\nconstexpr size_t reduce_dbias_num_threads = 256;\n\ntemplate \n__global__ void __launch_bounds__(reduce_dbias_num_threads)\n reduce_dbias_kernel(OutputType *const dbias_output, const ComputeType *const dbias_partial,\n const int row_length, const int num_rows) {\n using ComputeVec = Vec;\n using OutputVec = Vec;\n\n const int thread_id = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (thread_id * nvec >= row_length) return;\n\n const ComputeType *const thread_in_base = dbias_partial + thread_id * nvec;\n OutputType *const thread_out_base = dbias_output + thread_id * nvec;\n\n const int stride_in_vec = row_length / nvec;\n\n ComputeVec ldg_vec;\n ComputeVec acc_vec;\n acc_vec.clear();\n for (int i = 0; i < num_rows; ++i) {\n ldg_vec.load_from(thread_in_base, i * stride_in_vec);\n#pragma unroll\n for (int e = 0; e < nvec; ++e) {\n acc_vec.data.elt[e] += ldg_vec.data.elt[e];\n }\n }\n\n OutputVec stg_vec;\n#pragma unroll\n for (int e = 0; e < nvec; ++e) {\n stg_vec.data.elt[e] = OutputType(acc_vec.data.elt[e]);\n }\n stg_vec.store_to(thread_out_base, 0);\n}\n\nvoid populate_transpose_dbias_workspace_config(const Tensor &input, /*cast*/\n Tensor *workspace, const int nvec_out) {\n const size_t row_length = input.data.shape[1];\n const size_t num_rows = input.data.shape[0];\n\n const size_t tile_size_y = (nvec_out * THREADS_PER_WARP);\n NVTE_CHECK(num_rows % nvec_out == 0, \"Unsupported shape.\");\n\n const size_t num_rows_partial_dbias = DIVUP(num_rows, tile_size_y);\n\n if (workspace->data.dptr == nullptr) {\n // Set workspace size\n workspace->data.shape = {num_rows_partial_dbias, row_length};\n workspace->data.dtype = DType::kFloat32;\n } else {\n // Check that workspace matches expected size\n const size_t workspace_size = get_buffer_size_bytes(\n std::accumulate(workspace->data.shape.begin(), workspace->data.shape.end(), 1,\n std::multiplies()),\n workspace->data.dtype);\n const size_t required_size =\n get_buffer_size_bytes(num_rows_partial_dbias, row_length, DType::kFloat32);\n NVTE_CHECK(workspace_size >= required_size, \"Invalid workspace (expected dims=(\",\n num_rows_partial_dbias, \",\", row_length, \"), dtype=\", to_string(DType::kFloat32),\n \"; found dims=\", workspace->data.shape,\n \", dtype=\", typeToNumBits(workspace->data.dtype), \" bits)\");\n }\n}\n\ntemplate \nvoid reduce_dbias(const Tensor &workspace, Tensor *dbias, const size_t row_length,\n const size_t num_rows, const int nvec_out, hipStream_t stream) {\n constexpr int reduce_dbias_store_bytes = 8; // stg.64\n constexpr int reduce_dbias_nvec = reduce_dbias_store_bytes / sizeof(BiasType);\n\n NVTE_CHECK(row_length % reduce_dbias_nvec == 0, \"Unsupported shape.\");\n\n const size_t reduce_dbias_row_length = row_length;\n const size_t reduce_dbias_num_rows =\n DIVUP(num_rows, static_cast(nvec_out * THREADS_PER_WARP));\n const size_t reduce_dbias_num_blocks =\n DIVUP(row_length, reduce_dbias_num_threads * reduce_dbias_nvec);\n\n reduce_dbias_kernel\n <<>>(\n reinterpret_cast(dbias->data.dptr),\n reinterpret_cast(workspace.data.dptr), reduce_dbias_row_length,\n reduce_dbias_num_rows);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid fp8_transpose_dbias(const Tensor &input, Tensor *transposed_output, Tensor *dbias,\n Tensor *workspace, hipStream_t stream) {\n CheckInputTensor(input, \"fp8_transpose_dbias_input\");\n CheckOutputTensor(*transposed_output, \"transposed_output\");\n CheckOutputTensor(*dbias, \"dbias\");\n\n NVTE_CHECK(input.data.shape.size() == 2, \"Input must have 2 dimensions.\");\n NVTE_CHECK(transposed_output->data.shape.size() == 2, \"T output must have 2 dimensions.\");\n const size_t row_length = input.data.shape[1];\n const size_t num_rows = input.data.shape[0];\n\n NVTE_CHECK(transposed_output->data.shape[0] == row_length, \"Wrong dimension of T output.\");\n NVTE_CHECK(transposed_output->data.shape[1] == num_rows, \"Wrong dimension of T output.\");\n\n NVTE_CHECK(transposed_output->data.dtype == input.data.dtype,\n \"T output must have the same type as input.\");\n NVTE_CHECK(dbias->data.shape == std::vector{row_length}, \"Wrong shape of DBias.\");\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n dbias->data.dtype, BiasType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(\n input.data.dtype, Type, constexpr int type_size = sizeof(Type);\n constexpr int nvec_in = desired_load_size / type_size;\n constexpr int nvec_out = desired_store_size / type_size;\n\n // Check workspace size\n populate_transpose_dbias_workspace_config(input, workspace, nvec_out);\n if (workspace->data.dptr == nullptr) { return; }\n\n NVTE_CHECK(row_length % nvec_in == 0, \"Unsupported shape.\");\n NVTE_CHECK(num_rows % nvec_out == 0, \"Unsupported shape.\");\n const size_t n_tiles =\n DIVUP(row_length, static_cast(nvec_in * THREADS_PER_WARP)) *\n DIVUP(num_rows, static_cast(nvec_out * THREADS_PER_WARP));\n const size_t n_warps_per_block = cast_transpose_num_threads / THREADS_PER_WARP;\n const size_t n_blocks = DIVUP(n_tiles * n_warps_per_tile, n_warps_per_block);\n\n const bool full_tile = row_length % (nvec_in * THREADS_PER_WARP) == 0 &&\n num_rows % (nvec_out * THREADS_PER_WARP) == 0;\n\n using ComputeType = fp32; constexpr size_t shared_size_transpose =\n cast_transpose_num_threads / n_warps_per_tile *\n (THREADS_PER_WARP + 1) * sizeof(Vec);\n constexpr size_t shared_size_dbias =\n cast_transpose_num_threads * sizeof(Vec);\n static_assert(shared_size_transpose >= shared_size_dbias);\n using Param = TDBiasParam; Param param;\n param.input = reinterpret_cast(input.data.dptr);\n param.output_t = reinterpret_cast(transposed_output->data.dptr);\n param.scale_inv =\n reinterpret_cast(transposed_output->scale_inv.dptr);\n param.workspace = reinterpret_cast(workspace->data.dptr);\n\n if (full_tile) {\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(transpose_dbias_kernel),\n hipFuncAttributePreferredSharedMemoryCarveout,\n 100));\n transpose_dbias_kernel\n <<>>(\n param, row_length, num_rows, n_tiles);\n NVTE_CHECK_CUDA(hipGetLastError());\n } else {\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(transpose_dbias_kernel_notaligned),\n hipFuncAttributePreferredSharedMemoryCarveout, 100));\n transpose_dbias_kernel_notaligned\n <<>>(\n param, row_length, num_rows, n_tiles);\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n\n reduce_dbias(*workspace, dbias, row_length, num_rows, nvec_out,\n stream);); // NOLINT(*)\n ); // NOLINT(*)\n}\n\n} // namespace transformer_engine\n\nvoid nvte_fp8_transpose_dbias(const NVTETensor input, NVTETensor transposed_output,\n NVTETensor dbias, NVTETensor workspace, hipStream_t stream) {\n NVTE_API_CALL(nvte_fp8_transpose_dbias);\n using namespace transformer_engine;\n fp8_transpose_dbias(*convertNVTETensorCheck(input), convertNVTETensor(transposed_output),\n convertNVTETensor(dbias), convertNVTETensor(workspace), stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/transpose/transpose_fusion.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/recipe/recipe_common.cuh\"\n#include \"common/transpose/cast_transpose.h\"\n#include \"common/util/cuda_runtime.h\"\n#include \"common/utils.cuh\"\n\nnamespace transformer_engine {\nnamespace {\n\nusing transformer_engine::detail::FP8BlockwiseColumnwiseOption;\nusing transformer_engine::detail::FP8BlockwiseRowwiseOption;\n\n// clang-format off\n/*\n\nStep 1: Load input to shared memory\n* shard memory: 128x128 elements with type=InputType (below graph doesn't consider padding)\n* 8 warps\n* Loop 8 times\n* What each thread does in each loop:\n * 8 elements are read from the input at a time\n * 2 elements are written to the shared memory at a time, for a total of 4 times\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| T0 | T1 | T2 | T3 | T4 | T5 | T6 | T7 | T8 | T9 | T10 | T11 | T12 | T13 | T14 | T15 |\n| T16 | T17 | T18 | T19 | T20 | T21 | T22 | T23 | T24 | T25 | T26 | T27 | T28 | T29 | T30 | T31 |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| Warp 1 |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| Warp 7 |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n| ... |\n| Loop 8 times |\n| ... |\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n\nStep 2: Cast and store to output_c\n* shard memory: 128x128 elements with type=InputType (below graph doesn't consider padding)\n* 8 warps\n* Loop 4 times\n* What each thread does in each loop:\n * 2 elements are read from the shared memory at a time, for a total of 8 times\n * Every 8 consecutive threads do reduction and calculate the amax of each row\n * 16 elements are quantized and write to output_c at a time\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| T0 | T1 | T2 | T3 | T4 | T5 | T6 | T7 |\n| T8 | T9 | T10 | T11 | T12 | T13 | T14 | T15 |\n| T16 | T17 | T18 | T19 | T20 | T21 | T22 | T23 |\n| T24 | T25 | T26 | T27 | T28 | T29 | T30 | T31 |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| |\n| Warp 1 |\n| |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| |\n| Warp 7 |\n| |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n| ... |\n| Loop 4 times |\n| ... |\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n\nStep 3 (if columnwise transpose is True, GEMM_READY): Transpose, cast and store to output_t\n* shard memory: 128x128 elements with type=InputType (below graph doesn't consider padding)\n* 8 warps\n* Loop 2 times\n* What each thread does in each loop:\n * 2 elements (in a row) are read from the shared memory at a time, for a total of 16 times\n * Every 8 consecutive threads do reduction and calculate the amax of each column\n * 16 elements are quantized and write to output_t at a time, for a total of 2 times\n+------8 elements-------+------8 elements-------+-----40 elements-------+------8 elements-------+------8 elements-------+------8 elements-------+-----40 elements-------+------8 elements-------+\n| T0 | T8 | T16 | T24 | | | | T0 | T8 | T16 | T24 | | | |\n| T1 | T9 | T17 | T25 | | | | T1 | T9 | T17 | T25 | | | |\n| T2 | T10 | T18 | T26 | | | | T2 | T10 | T18 | T26 | | | |\n| T3 | T11 | T19 | T27 | Warp 1 | ... | Warp 7 | T3 | T11 | T19 | T27 | Warp 1 | ... | Warp 7 |\n| T4 | T12 | T20 | T28 | | | | T4 | T12 | T20 | T28 | | | |\n| T5 | T13 | T21 | T29 | | | | T5 | T13 | T21 | T29 | | | |\n| T6 | T14 | T22 | T30 | | | | T6 | T14 | T22 | T30 | | | |\n| T7 | T15 | T23 | T31 | | | | T7 | T15 | T23 | T31 | | | |\n+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+\n\nStep 3 (if columnwise transpose is False, COMPACT format): Skip Transpose, cast and store to output_t\n* shard memory: 128x128 elements with type=InputType (below graph doesn't consider padding)\n* 8 warps\n* Loop 1 times\n* What each thread does in each loop:\n * 16 elements (in a row) are read from the shared memory, for a total of 4 rows,\n * it needs 8 reads in smem to get 16 elements in a row, thread tile shape is 16x4\n * Every 32 consecutive threads in a warp do reduction and calculate the amax of each column,\n * so each thread will do warp shuffle 16 times to get the amax of each column\n * 16 elements are quantized and write to output_t at a time, for a total of 4 times\n+------16 elements-------+------16 elements-------+-----80 elements-----+------16 elements------+\n| T0 | | | |\n| T1 | | | |\n| T2 | | | |\n| T3 | | | |\n| T4 | | | |\n| T5 | | | |\n| T6 | | | |\n| T7 | | | |\n| ... | | | |\n| T31 | | | |\n+-----------------------+-----------------------+-----------------------+-----------------------+\n\n*/\n// clang-format on\n\nconstexpr size_t kThreadsPerWarp = 32;\n\n// Hyperparameters for performance tuning\nconstexpr int kTileDim = 128; // Fixed to 128 beacause we are using 1x128 and 128x1 quantization\nconstexpr int kNVecIn = 8; // The number of elements each LDG touches\nconstexpr int kNVecOut = 16; // The number of elements each STG touches\nconstexpr int kNVecSMem = 2; // The number of elements each LDS/STS touches\nconstexpr int kThreadsPerBlock = 256; // Thread block size, 8 warps in total\n\n// Auto-calculated constants, do not modify directly)\nstatic_assert(kNVecIn % kNVecSMem == 0, \"kNVecIn must be divisible by kNVecSMem\");\nstatic_assert(kNVecOut % kNVecSMem == 0, \"kNVecOut must be divisible by kNVecSMem\");\nconstexpr int kSMemRow = kTileDim;\nconstexpr int kSMemCol = (kTileDim / kNVecSMem) + 1;\nconstexpr int kSMemSize = kSMemRow * kSMemCol * kNVecSMem;\nconstexpr int kNumThreadsLoad = kTileDim / kNVecIn;\nconstexpr int kNumThreadsStore = kTileDim / kNVecOut;\nstatic_assert(kNumThreadsLoad <= kThreadsPerWarp, \"kNumThreadsLoad must be <= kThreadsPerWarp\");\nstatic_assert(kNumThreadsStore <= kThreadsPerWarp, \"kNumThreadsStore must be <= kThreadsPerWarp\");\nconstexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp;\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpose_kernel(\n const IType* const input, OType* const output_c, OType* const output_t,\n CType* const tile_scales_inv_c, CType* const tile_scales_inv_t, const size_t row_length,\n const size_t num_rows, const size_t scale_stride_x, const size_t scale_stride_y,\n const size_t scale_t_stride_x, const size_t scale_t_stride_y, const float epsilon,\n FP8BlockwiseRowwiseOption rowwise_option, FP8BlockwiseColumnwiseOption columnwise_option,\n const bool pow_2_scaling, const float* noop_ptr) {\n // skip execution if noop\n if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {\n return;\n }\n\n bool return_rowwise = rowwise_option != FP8BlockwiseRowwiseOption::NONE;\n bool return_columnwise_gemm_ready =\n columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY;\n bool return_columnwise_compact =\n columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT;\n\n using SMemVec = Vec;\n using OVec = Vec;\n union IVec {\n Vec input_type;\n Vec smem_type;\n };\n\n extern __shared__ char smem_base[];\n SMemVec* smem = reinterpret_cast(&smem_base[0]);\n\n // Step 1: Load input to shared memory\n {\n constexpr int r_stride = kThreadsPerBlock / kNumThreadsLoad; // stride in rows of shared memory\n constexpr int num_iterations = kTileDim / r_stride;\n const int c_s =\n (threadIdx.x % kNumThreadsLoad) * (kNVecIn / kNVecSMem); // Column in shared memory\n int r_s = threadIdx.x / kNumThreadsLoad; // Row in shared memory\n const size_t c_g =\n static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem; // Column in global memory\n size_t r_g = static_cast(blockIdx.y) * kTileDim + r_s; // Row in global memory\n const size_t stride_g = static_cast(r_stride) * row_length; // Stride in global memory\n const size_t num_ele = c_g < row_length ? min(static_cast(kNVecIn), row_length - c_g)\n : 0; // For not aligned case\n const IType* input_g = &input[r_g * row_length + c_g]; // Input address in global memory\n#pragma unroll\n for (int iter = 0; iter < num_iterations; ++iter) {\n IVec input_vec;\n // Step 1.1: Load from global memory (input) to registers\n if constexpr (kAligned) {\n input_vec.input_type.load_from(input_g);\n } else {\n if (r_g < num_rows) {\n input_vec.input_type.load_from_elts(input_g, 0, num_ele);\n } else {\n input_vec.input_type.clear();\n }\n }\n // Step 1.2: Write to shared memory\n#pragma unroll\n for (int i = 0; i < kNVecIn / kNVecSMem; ++i) {\n int c = c_s + i;\n int r = r_s;\n smem[r * kSMemCol + c] = input_vec.smem_type.data.elt[i];\n }\n // Step 1.3: Update input address, row index of shared memory, (and row index of global memory for not aligned case)\n input_g += stride_g;\n r_s += r_stride;\n if constexpr (!kAligned) {\n r_g += r_stride;\n }\n }\n }\n\n __syncthreads();\n\n // Step 2: Cast and store to output_c\n if (return_rowwise) {\n constexpr int r_stride =\n kThreadsPerBlock / kNumThreadsStore; // stride in rows of shared memory\n constexpr int num_iterations = kTileDim / r_stride;\n const int c_s =\n (threadIdx.x % kNumThreadsStore) * (kNVecOut / kNVecSMem); // Column in shared memory\n int r_s = threadIdx.x / kNumThreadsStore; // Row in shared memory\n const size_t c_g =\n static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem; // Column in global memory\n size_t r_g = static_cast(blockIdx.y) * kTileDim + r_s; // Row in global memory\n const size_t stride_g = static_cast(r_stride) * row_length; // Stride in global memory\n const size_t num_ele = c_g < row_length ? min(static_cast(kNVecOut), row_length - c_g)\n : 0; // For not aligned case\n OType* output_g = &output_c[r_g * row_length + c_g]; // Output address in global memory\n // Each kNumThreadsStore threads form a warp process one row, we need to find the lane id of\n // the first thread to do the reduction.\n const unsigned src_lane = (threadIdx.x % kThreadsPerWarp) / kNumThreadsStore * kNumThreadsStore;\n // This mask represents which threads should do the reduction together.\n const unsigned mask = ((1 << kNumThreadsStore) - 1) << src_lane;\n const bool is_src_lane = (threadIdx.x % kNumThreadsStore) == 0;\n#pragma unroll\n for (int iter = 0; iter < num_iterations; ++iter) {\n SMemVec smem_vec[kNVecOut / kNVecSMem];\n // Step 2.1: Load from shared memory to registers\n#pragma unroll\n for (int i = 0; i < kNVecOut / kNVecSMem; ++i) {\n int c = c_s + i;\n int r = r_s;\n smem_vec[i] = smem[r * kSMemCol + c];\n }\n // Step 2.2: Compute local amax\n CType amax = 0;\n#pragma unroll\n for (int i = 0; i < kNVecOut / kNVecSMem; ++i) {\n#pragma unroll\n for (int j = 0; j < kNVecSMem; ++j) {\n __builtin_assume(amax >= 0);\n amax = fmaxf(amax, fabsf(smem_vec[i].data.elt[j]));\n }\n }\n // Step 2.3: Reduce amax\n#pragma unroll\n for (int delta = kNumThreadsStore / 2; delta > 0; delta /= 2) {\n const float other_amax = __shfl_down_sync(mask, amax, delta);\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n amax = __shfl_sync(mask, amax, src_lane);\n CType scale;\n // Step 2.4: Compute scale\n scale = compute_scale_from_types(amax, epsilon, pow_2_scaling);\n // Step 2.5: Write scale_inv\n bool write_scale_inv = is_src_lane;\n if constexpr (!kAligned) {\n write_scale_inv &= (r_g < num_rows);\n }\n if (write_scale_inv) {\n CType scale_inv = 1.0 / scale;\n size_t row_idx = static_cast(blockIdx.y) * kTileDim + r_s;\n size_t col_idx = static_cast(blockIdx.x);\n tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv;\n }\n // Step 2.6: Quantize\n OVec output_vec;\n#pragma unroll\n for (int i = 0; i < kNVecOut / kNVecSMem; ++i) {\n#pragma unroll\n for (int j = 0; j < kNVecSMem; ++j) {\n output_vec.data.elt[i * kNVecSMem + j] =\n static_cast(static_cast(smem_vec[i].data.elt[j]) * scale);\n }\n }\n // Step 2.7: Store output_c\n if constexpr (kAligned) {\n output_vec.store_to(output_g);\n } else {\n if (r_g < num_rows) {\n output_vec.store_to_elts(output_g, 0, num_ele);\n }\n }\n // Step 2.8: Update output address, row index of shared memory (and row index of global memory for not aligned case)\n output_g += stride_g;\n r_s += r_stride;\n if constexpr (!kAligned) {\n r_g += r_stride;\n }\n }\n }\n\n // Step 3 (return_columnwise_gemm_ready): Transpose, cast and store to output_t\n if (return_columnwise_gemm_ready) {\n constexpr int c_stride =\n kThreadsPerBlock / kNumThreadsStore; // Stride in columns of shared memory\n constexpr int num_iterations = kTileDim / (c_stride * kNVecSMem);\n const int r_s = (threadIdx.x % kNumThreadsStore) * kNVecOut; // Row in shared memory\n int c_s = threadIdx.x / kNumThreadsStore; // Column in shared memory\n size_t r_g =\n static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem; // Row in global memory\n const size_t c_g = static_cast(blockIdx.y) * kTileDim + r_s; // Column in global memory\n const size_t stride_g =\n static_cast(c_stride) * kNVecSMem * num_rows; // Stride in global memory\n const size_t num_ele = c_g < num_rows ? min(static_cast(kNVecOut), num_rows - c_g)\n : 0; // For not aligned case\n OType* output_g = &output_t[r_g * num_rows + c_g]; // Output address in global memory\n // Each kNumThreadsStore threads form a warp process one row, we need to find the lane id of\n // the first thread to do the reduction.\n const unsigned src_lane = (threadIdx.x % kThreadsPerWarp) / kNumThreadsStore * kNumThreadsStore;\n // This mask represents which threads should do the reduction together.\n const unsigned mask = ((1 << kNumThreadsStore) - 1) << src_lane;\n const bool is_src_lane = (threadIdx.x % kNumThreadsStore) == 0;\n#pragma unroll\n for (int iter = 0; iter < num_iterations; ++iter) {\n SMemVec smem_vec[kNVecOut];\n // Step 3.1: Load from shared memory to registers\n#pragma unroll\n for (int i = 0; i < kNVecOut; ++i) {\n int r = r_s + i;\n int c = c_s;\n smem_vec[i] = smem[r * kSMemCol + c];\n }\n#pragma unroll\n for (int smem_idx = 0; smem_idx < kNVecSMem; ++smem_idx) {\n // Step 3.2: Compute local amax\n CType amax = 0;\n#pragma unroll\n for (int i = 0; i < kNVecOut; ++i) {\n amax = fmaxf(amax, fabsf(smem_vec[i].data.elt[smem_idx]));\n }\n // Step 3.3: Reduce amax\n#pragma unroll\n for (int delta = kNumThreadsStore / 2; delta > 0; delta /= 2) {\n const float other_amax = __shfl_down_sync(mask, amax, delta);\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n amax = __shfl_sync(mask, amax, src_lane);\n // Step 3.4: Compute scale\n CType scale;\n scale = compute_scale_from_types(amax, epsilon, pow_2_scaling);\n // Step 3.5: Write scale_inv_t\n bool write_scale_inv = is_src_lane;\n if constexpr (!kAligned) {\n write_scale_inv &= (r_g + smem_idx < row_length);\n }\n if (write_scale_inv) {\n CType scale_inv = 1.0 / scale;\n size_t row_idx = static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem + smem_idx;\n size_t col_idx = static_cast(blockIdx.y);\n tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv;\n }\n // Step 3.6: Quantize\n OVec output_vec;\n#pragma unroll\n for (int i = 0; i < kNVecOut; ++i) {\n output_vec.data.elt[i] =\n static_cast(static_cast(smem_vec[i].data.elt[smem_idx]) * scale);\n }\n // Step 3.7: Store output_t\n if constexpr (kAligned) {\n output_vec.store_to(output_g + smem_idx * num_rows);\n } else {\n if (r_g + smem_idx < row_length) {\n output_vec.store_to_elts(output_g + smem_idx * num_rows, 0, num_ele);\n }\n }\n }\n // Step 3.8: Update output address, column index of shared memory (and row index of global memory for not aligned case)\n output_g += stride_g;\n c_s += c_stride;\n if constexpr (!kAligned) {\n r_g += c_stride * kNVecSMem;\n }\n }\n }\n\n // Step 4 (return_columnwise_compact): cast in 128x1 style and store to output, skip transpose\n if (return_columnwise_compact) {\n // thread tile should be 4x16, 16 means 8 smem reads\n constexpr int kThreadTileRow = kTileDim / kThreadsPerWarp;\n constexpr int kThreadTileCol = kNVecOut;\n using RegVec = Vec;\n using RegScaleVec = Vec;\n constexpr int num_smem_reads = kNVecOut / kNVecSMem;\n // c_stride will not be used here because we only have one iteration\n // constexpr int c_stride = kThreadTileCol * kNumWarps / kNVecSMem;\n constexpr int num_iterations =\n kTileDim / (kNumWarps * kThreadTileCol); // should be only one iteration\n static_assert(num_iterations == 1,\n \"num_iterations should be 1 for columnwise non-transpose case\");\n const int thr_idx_in_warp = threadIdx.x % kThreadsPerWarp;\n const int warp_idx = threadIdx.x / kThreadsPerWarp;\n const int r_s = thr_idx_in_warp * kThreadTileRow; // Row in shared memory\n int c_s = warp_idx * num_smem_reads; // Column in shared memory\n size_t r_g = static_cast(blockIdx.y) * kTileDim + r_s; // Row in global memory\n const size_t c_g =\n static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem; // Column in global memory\n const size_t num_ele = c_g < row_length\n ? min(static_cast(kThreadTileCol), row_length - c_g)\n : 0; // For not aligned case\n#pragma unroll\n for (int iter = 0; iter < num_iterations; ++iter) {\n RegVec reg_vec[kThreadTileRow];\n RegScaleVec thr_scale;\n\n // Step 3.1: Load from shared memory to registers\n#pragma unroll\n for (int i = 0; i < kThreadTileRow; ++i) {\n int r = r_s + i;\n#pragma unroll\n for (int j = 0; j < num_smem_reads; ++j) {\n int c = c_s + j;\n SMemVec smem_vec = smem[r * kSMemCol + c];\n // copy smem_vec to reg vec with its elements\n#pragma unroll\n for (int k = 0; k < kNVecSMem; ++k) {\n reg_vec[i].data.elt[j * kNVecSMem + k] = smem_vec.data.elt[k];\n }\n }\n }\n#pragma unroll\n for (int reg_idx = 0; reg_idx < kThreadTileCol; ++reg_idx) {\n // Step 3.2: Compute local amax\n CType amax = 0;\n#pragma unroll\n for (int i = 0; i < kThreadTileRow; ++i) {\n amax = fmaxf(amax, fabsf(reg_vec[i].data.elt[reg_idx]));\n }\n // Step 3.3: Reduce amax\n const bool is_src_lane = thr_idx_in_warp == 0;\n amax = warp_reduce_max(amax);\n constexpr int lane_zero = 0;\n amax = __shfl_sync(0xFFFFFFFF, amax, lane_zero);\n // Step 3.4: Compute scale\n CType scale;\n scale = compute_scale_from_types(amax, epsilon, pow_2_scaling);\n thr_scale.data.elt[reg_idx] = scale;\n // Step 3.5: Write scale_inv_t\n bool write_scale_inv = is_src_lane;\n if constexpr (!kAligned) {\n write_scale_inv &= (c_g + reg_idx < row_length);\n }\n if (write_scale_inv) {\n CType scale_inv = 1.0 / scale;\n size_t row_idx = static_cast(blockIdx.y);\n size_t col_idx = static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem + reg_idx;\n tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv;\n }\n }\n // Step 3.6: Quantize\n for (int row_idx = 0; row_idx < kThreadTileRow; ++row_idx) {\n OType* output_g =\n &output_t[(r_g + row_idx) * row_length + c_g]; // Output address in global memory\n OVec output_vec;\n#pragma unroll\n for (int i = 0; i < kThreadTileCol; ++i) {\n output_vec.data.elt[i] = static_cast(\n static_cast(reg_vec[row_idx].data.elt[i]) * thr_scale.data.elt[i]);\n }\n // Step 3.7: Store output_t\n if constexpr (kAligned) {\n output_vec.store_to(output_g);\n } else {\n if (r_g + row_idx < num_rows) {\n output_vec.store_to_elts(output_g, 0, num_ele);\n }\n }\n }\n // Step 3.8: Update output address, column index of shared memory\n // this section shouldn't matter since we only have one iteration\n }\n }\n}\n\n} // namespace\n} // namespace transformer_engine\n\nnamespace transformer_engine::detail {\n\nvoid quantize_transpose_vector_blockwise(const SimpleTensor& input, SimpleTensor& scale_inv,\n SimpleTensor& scale_inv_t, SimpleTensor& output,\n SimpleTensor& output_t, const float epsilon,\n FP8BlockwiseRowwiseOption rowwise_option,\n FP8BlockwiseColumnwiseOption columnwise_option,\n const bool pow2_scale, const SimpleTensor& noop_tensor,\n cudaStream_t stream) {\n NVTE_API_CALL(quantize_transpose_vector_blockwise);\n\n if (transformer_engine::cuda::sm_arch() >= 100) {\n NVTE_CHECK(pow2_scale, \"On Blackwell and newer, the FP8 block scaling recipe is emulated \",\n \"with MXFP8, which requires using power of two scaling factors.\");\n }\n\n const size_t row_length = input.shape.size() > 0 ? input.shape.at(input.shape.size() - 1) : 1u;\n size_t num_elements = row_length;\n size_t num_rows = 1;\n for (size_t i = 0; (i < input.shape.size() - 1) && (input.shape.size() > 0); ++i) {\n num_rows *= input.shape.at(i);\n num_elements *= input.shape.at(i);\n }\n\n // Early return if the input tensor is empty\n if (num_elements == 0) {\n return;\n }\n\n // Options for scale layout of cuBLAS GEMM kernel.\n size_t scale_stride_x = 0;\n size_t scale_stride_y = 0;\n size_t scale_t_stride_x = 0;\n size_t scale_t_stride_y = 0;\n\n if (rowwise_option != FP8BlockwiseRowwiseOption::NONE) {\n NVTE_CHECK(rowwise_option == FP8BlockwiseRowwiseOption::ROWWISE_GEMM_READY ||\n rowwise_option == FP8BlockwiseRowwiseOption::ROWWISE_COMPACT,\n \"Unexpected rowwise enum value\");\n NVTE_CHECK(input.shape == output.shape, \"Input and output must have the same shape.\");\n NVTE_CHECK(scale_inv.shape.size() == 2, \"Scale dimension must be 2.\");\n size_t scale_k = scale_inv.shape[1];\n bool rowwise_compact = rowwise_option == FP8BlockwiseRowwiseOption::ROWWISE_COMPACT;\n scale_stride_x = rowwise_compact ? 1 : scale_k;\n scale_stride_y = rowwise_compact ? scale_k : 1;\n }\n\n if (columnwise_option != FP8BlockwiseColumnwiseOption::NONE) {\n NVTE_CHECK(output_t.shape.size() == input.shape.size(),\n \"output_t must have same number of dimensions as input.\");\n\n if (output_t.shape.size() > 0) {\n if (columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY) {\n NVTE_CHECK(output_t.shape[0] == row_length, \"Wrong dimension 0 of output_t.\");\n for (size_t i = 1; i < output_t.shape.size(); ++i) {\n NVTE_CHECK(output_t.shape.at(i) == input.shape.at(i - 1), \"Wrong dimension in output_t\");\n }\n } else {\n NVTE_CHECK(columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT,\n \"Unexpected columnwise option enum value\");\n NVTE_CHECK(output_t.shape[0] == input.shape[0], \"Wrong dimension 0 of output_t.\");\n NVTE_CHECK(\n input.shape == output_t.shape,\n \"Input and output_t must have the same shape for columnwise non-transpose case.\");\n }\n }\n if (rowwise_option != FP8BlockwiseRowwiseOption::NONE) {\n // output may not be defined if rowwise quantization is not needed.\n NVTE_CHECK(output.dtype == output_t.dtype,\n \"output and output_t need to have the same dtype.\");\n }\n NVTE_CHECK(scale_inv_t.shape.size() == 2, \"Scale_t dimension must be 2.\");\n bool columnwise_compact = columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT;\n size_t scale_t_k = scale_inv_t.shape[1];\n scale_t_stride_x = columnwise_compact ? 1 : scale_t_k;\n scale_t_stride_y = columnwise_compact ? scale_t_k : 1;\n }\n auto output_dtype =\n rowwise_option != FP8BlockwiseRowwiseOption::NONE ? output.dtype : output_t.dtype;\n\n const size_t num_blocks_x = DIVUP(row_length, (size_t)kTileDim);\n const size_t num_blocks_y = DIVUP(num_rows, (size_t)kTileDim);\n\n const float* noop_ptr = reinterpret_cast(noop_tensor.dptr);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.dtype, InputType,\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(\n output_dtype, OutputType,\n\n dim3 grid(num_blocks_x, num_blocks_y, 1);\n\n const bool full_tile = row_length % kTileDim == 0 && num_rows % kTileDim == 0;\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n full_tile, kAligned,\n\n size_t smem_bytes = kSMemSize * sizeof(InputType);\n // shared memory must be requested up\n if (smem_bytes >= 48 * 1024) {\n cudaError_t err = cudaFuncSetAttribute(\n &block_scaled_1d_cast_transpose_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);\n NVTE_CHECK(err == cudaSuccess, \"Failed to set dynamic shared memory size.\");\n } block_scaled_1d_cast_transpose_kernel\n <<>>(\n reinterpret_cast(input.dptr),\n reinterpret_cast(output.dptr),\n reinterpret_cast(output_t.dptr),\n reinterpret_cast(scale_inv.dptr),\n reinterpret_cast(scale_inv_t.dptr), row_length, num_rows, scale_stride_x,\n scale_stride_y, scale_t_stride_x, scale_t_stride_y, epsilon, rowwise_option,\n columnwise_option, pow2_scale, noop_ptr);) // kAligned\n ) // OutputType\n ) // InputType\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n} // namespace transformer_engine::detail\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/recipe/recipe_common.cuh\"\n#include \"common/transpose/cast_transpose.h\"\n#include \"common/util/hip/hip_runtime.h\"\n#include \"common/utils.cuh\"\n\nnamespace transformer_engine {\nnamespace {\n\nusing transformer_engine::detail::FP8BlockwiseColumnwiseOption;\nusing transformer_engine::detail::FP8BlockwiseRowwiseOption;\n\n// clang-format off\n/*\n\nStep 1: Load input to shared memory\n* shard memory: 128x128 elements with type=InputType (below graph doesn't consider padding)\n* 8 warps\n* Loop 8 times\n* What each thread does in each loop:\n * 8 elements are read from the input at a time\n * 2 elements are written to the shared memory at a time, for a total of 4 times\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| T0 | T1 | T2 | T3 | T4 | T5 | T6 | T7 | T8 | T9 | T10 | T11 | T12 | T13 | T14 | T15 |\n| T16 | T17 | T18 | T19 | T20 | T21 | T22 | T23 | T24 | T25 | T26 | T27 | T28 | T29 | T30 | T31 |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| Warp 1 |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| Warp 7 |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n| ... |\n| Loop 8 times |\n| ... |\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n\nStep 2: Cast and store to output_c\n* shard memory: 128x128 elements with type=InputType (below graph doesn't consider padding)\n* 8 warps\n* Loop 4 times\n* What each thread does in each loop:\n * 2 elements are read from the shared memory at a time, for a total of 8 times\n * Every 8 consecutive threads do reduction and calculate the amax of each row\n * 16 elements are quantized and write to output_c at a time\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| T0 | T1 | T2 | T3 | T4 | T5 | T6 | T7 |\n| T8 | T9 | T10 | T11 | T12 | T13 | T14 | T15 |\n| T16 | T17 | T18 | T19 | T20 | T21 | T22 | T23 |\n| T24 | T25 | T26 | T27 | T28 | T29 | T30 | T31 |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| |\n| Warp 1 |\n| |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| |\n| Warp 7 |\n| |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n| ... |\n| Loop 4 times |\n| ... |\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n\nStep 3 (if columnwise transpose is True, GEMM_READY): Transpose, cast and store to output_t\n* shard memory: 128x128 elements with type=InputType (below graph doesn't consider padding)\n* 8 warps\n* Loop 2 times\n* What each thread does in each loop:\n * 2 elements (in a row) are read from the shared memory at a time, for a total of 16 times\n * Every 8 consecutive threads do reduction and calculate the amax of each column\n * 16 elements are quantized and write to output_t at a time, for a total of 2 times\n+------8 elements-------+------8 elements-------+-----40 elements-------+------8 elements-------+------8 elements-------+------8 elements-------+-----40 elements-------+------8 elements-------+\n| T0 | T8 | T16 | T24 | | | | T0 | T8 | T16 | T24 | | | |\n| T1 | T9 | T17 | T25 | | | | T1 | T9 | T17 | T25 | | | |\n| T2 | T10 | T18 | T26 | | | | T2 | T10 | T18 | T26 | | | |\n| T3 | T11 | T19 | T27 | Warp 1 | ... | Warp 7 | T3 | T11 | T19 | T27 | Warp 1 | ... | Warp 7 |\n| T4 | T12 | T20 | T28 | | | | T4 | T12 | T20 | T28 | | | |\n| T5 | T13 | T21 | T29 | | | | T5 | T13 | T21 | T29 | | | |\n| T6 | T14 | T22 | T30 | | | | T6 | T14 | T22 | T30 | | | |\n| T7 | T15 | T23 | T31 | | | | T7 | T15 | T23 | T31 | | | |\n+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+\n\nStep 3 (if columnwise transpose is False, COMPACT format): Skip Transpose, cast and store to output_t\n* shard memory: 128x128 elements with type=InputType (below graph doesn't consider padding)\n* 8 warps\n* Loop 1 times\n* What each thread does in each loop:\n * 16 elements (in a row) are read from the shared memory, for a total of 4 rows,\n * it needs 8 reads in smem to get 16 elements in a row, thread tile shape is 16x4\n * Every 32 consecutive threads in a warp do reduction and calculate the amax of each column,\n * so each thread will do warp shuffle 16 times to get the amax of each column\n * 16 elements are quantized and write to output_t at a time, for a total of 4 times\n+------16 elements-------+------16 elements-------+-----80 elements-----+------16 elements------+\n| T0 | | | |\n| T1 | | | |\n| T2 | | | |\n| T3 | | | |\n| T4 | | | |\n| T5 | | | |\n| T6 | | | |\n| T7 | | | |\n| ... | | | |\n| T31 | | | |\n+-----------------------+-----------------------+-----------------------+-----------------------+\n\n*/\n// clang-format on\n\nconstexpr size_t kThreadsPerWarp = 32;\n\n// Hyperparameters for performance tuning\nconstexpr int kTileDim = 128; // Fixed to 128 beacause we are using 1x128 and 128x1 quantization\nconstexpr int kNVecIn = 8; // The number of elements each LDG touches\nconstexpr int kNVecOut = 16; // The number of elements each STG touches\nconstexpr int kNVecSMem = 2; // The number of elements each LDS/STS touches\nconstexpr int kThreadsPerBlock = 256; // Thread block size, 8 warps in total\n\n// Auto-calculated constants, do not modify directly)\nstatic_assert(kNVecIn % kNVecSMem == 0, \"kNVecIn must be divisible by kNVecSMem\");\nstatic_assert(kNVecOut % kNVecSMem == 0, \"kNVecOut must be divisible by kNVecSMem\");\nconstexpr int kSMemRow = kTileDim;\nconstexpr int kSMemCol = (kTileDim / kNVecSMem) + 1;\nconstexpr int kSMemSize = kSMemRow * kSMemCol * kNVecSMem;\nconstexpr int kNumThreadsLoad = kTileDim / kNVecIn;\nconstexpr int kNumThreadsStore = kTileDim / kNVecOut;\nstatic_assert(kNumThreadsLoad <= kThreadsPerWarp, \"kNumThreadsLoad must be <= kThreadsPerWarp\");\nstatic_assert(kNumThreadsStore <= kThreadsPerWarp, \"kNumThreadsStore must be <= kThreadsPerWarp\");\nconstexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp;\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpose_kernel(\n const IType* const input, OType* const output_c, OType* const output_t,\n CType* const tile_scales_inv_c, CType* const tile_scales_inv_t, const size_t row_length,\n const size_t num_rows, const size_t scale_stride_x, const size_t scale_stride_y,\n const size_t scale_t_stride_x, const size_t scale_t_stride_y, const float epsilon,\n FP8BlockwiseRowwiseOption rowwise_option, FP8BlockwiseColumnwiseOption columnwise_option,\n const bool pow_2_scaling, const float* noop_ptr) {\n // skip execution if noop\n if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {\n return;\n }\n\n bool return_rowwise = rowwise_option != FP8BlockwiseRowwiseOption::NONE;\n bool return_columnwise_gemm_ready =\n columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY;\n bool return_columnwise_compact =\n columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT;\n\n using SMemVec = Vec;\n using OVec = Vec;\n union IVec {\n Vec input_type;\n Vec smem_type;\n };\n\n extern __shared__ char smem_base[];\n SMemVec* smem = reinterpret_cast(&smem_base[0]);\n\n // Step 1: Load input to shared memory\n {\n constexpr int r_stride = kThreadsPerBlock / kNumThreadsLoad; // stride in rows of shared memory\n constexpr int num_iterations = kTileDim / r_stride;\n const int c_s =\n (threadIdx.x % kNumThreadsLoad) * (kNVecIn / kNVecSMem); // Column in shared memory\n int r_s = threadIdx.x / kNumThreadsLoad; // Row in shared memory\n const size_t c_g =\n static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem; // Column in global memory\n size_t r_g = static_cast(blockIdx.y) * kTileDim + r_s; // Row in global memory\n const size_t stride_g = static_cast(r_stride) * row_length; // Stride in global memory\n const size_t num_ele = c_g < row_length ? min(static_cast(kNVecIn), row_length - c_g)\n : 0; // For not aligned case\n const IType* input_g = &input[r_g * row_length + c_g]; // Input address in global memory\n#pragma unroll\n for (int iter = 0; iter < num_iterations; ++iter) {\n IVec input_vec;\n // Step 1.1: Load from global memory (input) to registers\n if constexpr (kAligned) {\n input_vec.input_type.load_from(input_g);\n } else {\n if (r_g < num_rows) {\n input_vec.input_type.load_from_elts(input_g, 0, num_ele);\n } else {\n input_vec.input_type.clear();\n }\n }\n // Step 1.2: Write to shared memory\n#pragma unroll\n for (int i = 0; i < kNVecIn / kNVecSMem; ++i) {\n int c = c_s + i;\n int r = r_s;\n smem[r * kSMemCol + c] = input_vec.smem_type.data.elt[i];\n }\n // Step 1.3: Update input address, row index of shared memory, (and row index of global memory for not aligned case)\n input_g += stride_g;\n r_s += r_stride;\n if constexpr (!kAligned) {\n r_g += r_stride;\n }\n }\n }\n\n __syncthreads();\n\n // Step 2: Cast and store to output_c\n if (return_rowwise) {\n constexpr int r_stride =\n kThreadsPerBlock / kNumThreadsStore; // stride in rows of shared memory\n constexpr int num_iterations = kTileDim / r_stride;\n const int c_s =\n (threadIdx.x % kNumThreadsStore) * (kNVecOut / kNVecSMem); // Column in shared memory\n int r_s = threadIdx.x / kNumThreadsStore; // Row in shared memory\n const size_t c_g =\n static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem; // Column in global memory\n size_t r_g = static_cast(blockIdx.y) * kTileDim + r_s; // Row in global memory\n const size_t stride_g = static_cast(r_stride) * row_length; // Stride in global memory\n const size_t num_ele = c_g < row_length ? min(static_cast(kNVecOut), row_length - c_g)\n : 0; // For not aligned case\n OType* output_g = &output_c[r_g * row_length + c_g]; // Output address in global memory\n // Each kNumThreadsStore threads form a warp process one row, we need to find the lane id of\n // the first thread to do the reduction.\n const unsigned src_lane = (threadIdx.x % kThreadsPerWarp) / kNumThreadsStore * kNumThreadsStore;\n // This mask represents which threads should do the reduction together.\n const unsigned mask = ((1 << kNumThreadsStore) - 1) << src_lane;\n const bool is_src_lane = (threadIdx.x % kNumThreadsStore) == 0;\n#pragma unroll\n for (int iter = 0; iter < num_iterations; ++iter) {\n SMemVec smem_vec[kNVecOut / kNVecSMem];\n // Step 2.1: Load from shared memory to registers\n#pragma unroll\n for (int i = 0; i < kNVecOut / kNVecSMem; ++i) {\n int c = c_s + i;\n int r = r_s;\n smem_vec[i] = smem[r * kSMemCol + c];\n }\n // Step 2.2: Compute local amax\n CType amax = 0;\n#pragma unroll\n for (int i = 0; i < kNVecOut / kNVecSMem; ++i) {\n#pragma unroll\n for (int j = 0; j < kNVecSMem; ++j) {\n __builtin_assume(amax >= 0);\n amax = fmaxf(amax, fabsf(smem_vec[i].data.elt[j]));\n }\n }\n // Step 2.3: Reduce amax\n#pragma unroll\n for (int delta = kNumThreadsStore / 2; delta > 0; delta /= 2) {\n const float other_amax = __shfl_down_sync(mask, amax, delta);\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n amax = __shfl_sync(mask, amax, src_lane);\n CType scale;\n // Step 2.4: Compute scale\n scale = compute_scale_from_types(amax, epsilon, pow_2_scaling);\n // Step 2.5: Write scale_inv\n bool write_scale_inv = is_src_lane;\n if constexpr (!kAligned) {\n write_scale_inv &= (r_g < num_rows);\n }\n if (write_scale_inv) {\n CType scale_inv = 1.0 / scale;\n size_t row_idx = static_cast(blockIdx.y) * kTileDim + r_s;\n size_t col_idx = static_cast(blockIdx.x);\n tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv;\n }\n // Step 2.6: Quantize\n OVec output_vec;\n#pragma unroll\n for (int i = 0; i < kNVecOut / kNVecSMem; ++i) {\n#pragma unroll\n for (int j = 0; j < kNVecSMem; ++j) {\n output_vec.data.elt[i * kNVecSMem + j] =\n static_cast(static_cast(smem_vec[i].data.elt[j]) * scale);\n }\n }\n // Step 2.7: Store output_c\n if constexpr (kAligned) {\n output_vec.store_to(output_g);\n } else {\n if (r_g < num_rows) {\n output_vec.store_to_elts(output_g, 0, num_ele);\n }\n }\n // Step 2.8: Update output address, row index of shared memory (and row index of global memory for not aligned case)\n output_g += stride_g;\n r_s += r_stride;\n if constexpr (!kAligned) {\n r_g += r_stride;\n }\n }\n }\n\n // Step 3 (return_columnwise_gemm_ready): Transpose, cast and store to output_t\n if (return_columnwise_gemm_ready) {\n constexpr int c_stride =\n kThreadsPerBlock / kNumThreadsStore; // Stride in columns of shared memory\n constexpr int num_iterations = kTileDim / (c_stride * kNVecSMem);\n const int r_s = (threadIdx.x % kNumThreadsStore) * kNVecOut; // Row in shared memory\n int c_s = threadIdx.x / kNumThreadsStore; // Column in shared memory\n size_t r_g =\n static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem; // Row in global memory\n const size_t c_g = static_cast(blockIdx.y) * kTileDim + r_s; // Column in global memory\n const size_t stride_g =\n static_cast(c_stride) * kNVecSMem * num_rows; // Stride in global memory\n const size_t num_ele = c_g < num_rows ? min(static_cast(kNVecOut), num_rows - c_g)\n : 0; // For not aligned case\n OType* output_g = &output_t[r_g * num_rows + c_g]; // Output address in global memory\n // Each kNumThreadsStore threads form a warp process one row, we need to find the lane id of\n // the first thread to do the reduction.\n const unsigned src_lane = (threadIdx.x % kThreadsPerWarp) / kNumThreadsStore * kNumThreadsStore;\n // This mask represents which threads should do the reduction together.\n const unsigned mask = ((1 << kNumThreadsStore) - 1) << src_lane;\n const bool is_src_lane = (threadIdx.x % kNumThreadsStore) == 0;\n#pragma unroll\n for (int iter = 0; iter < num_iterations; ++iter) {\n SMemVec smem_vec[kNVecOut];\n // Step 3.1: Load from shared memory to registers\n#pragma unroll\n for (int i = 0; i < kNVecOut; ++i) {\n int r = r_s + i;\n int c = c_s;\n smem_vec[i] = smem[r * kSMemCol + c];\n }\n#pragma unroll\n for (int smem_idx = 0; smem_idx < kNVecSMem; ++smem_idx) {\n // Step 3.2: Compute local amax\n CType amax = 0;\n#pragma unroll\n for (int i = 0; i < kNVecOut; ++i) {\n amax = fmaxf(amax, fabsf(smem_vec[i].data.elt[smem_idx]));\n }\n // Step 3.3: Reduce amax\n#pragma unroll\n for (int delta = kNumThreadsStore / 2; delta > 0; delta /= 2) {\n const float other_amax = __shfl_down_sync(mask, amax, delta);\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n amax = __shfl_sync(mask, amax, src_lane);\n // Step 3.4: Compute scale\n CType scale;\n scale = compute_scale_from_types(amax, epsilon, pow_2_scaling);\n // Step 3.5: Write scale_inv_t\n bool write_scale_inv = is_src_lane;\n if constexpr (!kAligned) {\n write_scale_inv &= (r_g + smem_idx < row_length);\n }\n if (write_scale_inv) {\n CType scale_inv = 1.0 / scale;\n size_t row_idx = static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem + smem_idx;\n size_t col_idx = static_cast(blockIdx.y);\n tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv;\n }\n // Step 3.6: Quantize\n OVec output_vec;\n#pragma unroll\n for (int i = 0; i < kNVecOut; ++i) {\n output_vec.data.elt[i] =\n static_cast(static_cast(smem_vec[i].data.elt[smem_idx]) * scale);\n }\n // Step 3.7: Store output_t\n if constexpr (kAligned) {\n output_vec.store_to(output_g + smem_idx * num_rows);\n } else {\n if (r_g + smem_idx < row_length) {\n output_vec.store_to_elts(output_g + smem_idx * num_rows, 0, num_ele);\n }\n }\n }\n // Step 3.8: Update output address, column index of shared memory (and row index of global memory for not aligned case)\n output_g += stride_g;\n c_s += c_stride;\n if constexpr (!kAligned) {\n r_g += c_stride * kNVecSMem;\n }\n }\n }\n\n // Step 4 (return_columnwise_compact): cast in 128x1 style and store to output, skip transpose\n if (return_columnwise_compact) {\n // thread tile should be 4x16, 16 means 8 smem reads\n constexpr int kThreadTileRow = kTileDim / kThreadsPerWarp;\n constexpr int kThreadTileCol = kNVecOut;\n using RegVec = Vec;\n using RegScaleVec = Vec;\n constexpr int num_smem_reads = kNVecOut / kNVecSMem;\n // c_stride will not be used here because we only have one iteration\n // constexpr int c_stride = kThreadTileCol * kNumWarps / kNVecSMem;\n constexpr int num_iterations =\n kTileDim / (kNumWarps * kThreadTileCol); // should be only one iteration\n static_assert(num_iterations == 1,\n \"num_iterations should be 1 for columnwise non-transpose case\");\n const int thr_idx_in_warp = threadIdx.x % kThreadsPerWarp;\n const int warp_idx = threadIdx.x / kThreadsPerWarp;\n const int r_s = thr_idx_in_warp * kThreadTileRow; // Row in shared memory\n int c_s = warp_idx * num_smem_reads; // Column in shared memory\n size_t r_g = static_cast(blockIdx.y) * kTileDim + r_s; // Row in global memory\n const size_t c_g =\n static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem; // Column in global memory\n const size_t num_ele = c_g < row_length\n ? min(static_cast(kThreadTileCol), row_length - c_g)\n : 0; // For not aligned case\n#pragma unroll\n for (int iter = 0; iter < num_iterations; ++iter) {\n RegVec reg_vec[kThreadTileRow];\n RegScaleVec thr_scale;\n\n // Step 3.1: Load from shared memory to registers\n#pragma unroll\n for (int i = 0; i < kThreadTileRow; ++i) {\n int r = r_s + i;\n#pragma unroll\n for (int j = 0; j < num_smem_reads; ++j) {\n int c = c_s + j;\n SMemVec smem_vec = smem[r * kSMemCol + c];\n // copy smem_vec to reg vec with its elements\n#pragma unroll\n for (int k = 0; k < kNVecSMem; ++k) {\n reg_vec[i].data.elt[j * kNVecSMem + k] = smem_vec.data.elt[k];\n }\n }\n }\n#pragma unroll\n for (int reg_idx = 0; reg_idx < kThreadTileCol; ++reg_idx) {\n // Step 3.2: Compute local amax\n CType amax = 0;\n#pragma unroll\n for (int i = 0; i < kThreadTileRow; ++i) {\n amax = fmaxf(amax, fabsf(reg_vec[i].data.elt[reg_idx]));\n }\n // Step 3.3: Reduce amax\n const bool is_src_lane = thr_idx_in_warp == 0;\n amax = warp_reduce_max(amax);\n constexpr int lane_zero = 0;\n amax = __shfl_sync(0xFFFFFFFF, amax, lane_zero);\n // Step 3.4: Compute scale\n CType scale;\n scale = compute_scale_from_types(amax, epsilon, pow_2_scaling);\n thr_scale.data.elt[reg_idx] = scale;\n // Step 3.5: Write scale_inv_t\n bool write_scale_inv = is_src_lane;\n if constexpr (!kAligned) {\n write_scale_inv &= (c_g + reg_idx < row_length);\n }\n if (write_scale_inv) {\n CType scale_inv = 1.0 / scale;\n size_t row_idx = static_cast(blockIdx.y);\n size_t col_idx = static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem + reg_idx;\n tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv;\n }\n }\n // Step 3.6: Quantize\n for (int row_idx = 0; row_idx < kThreadTileRow; ++row_idx) {\n OType* output_g =\n &output_t[(r_g + row_idx) * row_length + c_g]; // Output address in global memory\n OVec output_vec;\n#pragma unroll\n for (int i = 0; i < kThreadTileCol; ++i) {\n output_vec.data.elt[i] = static_cast(\n static_cast(reg_vec[row_idx].data.elt[i]) * thr_scale.data.elt[i]);\n }\n // Step 3.7: Store output_t\n if constexpr (kAligned) {\n output_vec.store_to(output_g);\n } else {\n if (r_g + row_idx < num_rows) {\n output_vec.store_to_elts(output_g, 0, num_ele);\n }\n }\n }\n // Step 3.8: Update output address, column index of shared memory\n // this section shouldn't matter since we only have one iteration\n }\n }\n}\n\n} // namespace\n} // namespace transformer_engine\n\nnamespace transformer_engine::detail {\n\nvoid quantize_transpose_vector_blockwise(const SimpleTensor& input, SimpleTensor& scale_inv,\n SimpleTensor& scale_inv_t, SimpleTensor& output,\n SimpleTensor& output_t, const float epsilon,\n FP8BlockwiseRowwiseOption rowwise_option,\n FP8BlockwiseColumnwiseOption columnwise_option,\n const bool pow2_scale, const SimpleTensor& noop_tensor,\n hipStream_t stream) {\n NVTE_API_CALL(quantize_transpose_vector_blockwise);\n\n if (transformer_engine::cuda::sm_arch() >= 100) {\n NVTE_CHECK(pow2_scale, \"On Blackwell and newer, the FP8 block scaling recipe is emulated \",\n \"with MXFP8, which requires using power of two scaling factors.\");\n }\n\n const size_t row_length = input.shape.size() > 0 ? input.shape.at(input.shape.size() - 1) : 1u;\n size_t num_elements = row_length;\n size_t num_rows = 1;\n for (size_t i = 0; (i < input.shape.size() - 1) && (input.shape.size() > 0); ++i) {\n num_rows *= input.shape.at(i);\n num_elements *= input.shape.at(i);\n }\n\n // Early return if the input tensor is empty\n if (num_elements == 0) {\n return;\n }\n\n // Options for scale layout of cuBLAS GEMM kernel.\n size_t scale_stride_x = 0;\n size_t scale_stride_y = 0;\n size_t scale_t_stride_x = 0;\n size_t scale_t_stride_y = 0;\n\n if (rowwise_option != FP8BlockwiseRowwiseOption::NONE) {\n NVTE_CHECK(rowwise_option == FP8BlockwiseRowwiseOption::ROWWISE_GEMM_READY ||\n rowwise_option == FP8BlockwiseRowwiseOption::ROWWISE_COMPACT,\n \"Unexpected rowwise enum value\");\n NVTE_CHECK(input.shape == output.shape, \"Input and output must have the same shape.\");\n NVTE_CHECK(scale_inv.shape.size() == 2, \"Scale dimension must be 2.\");\n size_t scale_k = scale_inv.shape[1];\n bool rowwise_compact = rowwise_option == FP8BlockwiseRowwiseOption::ROWWISE_COMPACT;\n scale_stride_x = rowwise_compact ? 1 : scale_k;\n scale_stride_y = rowwise_compact ? scale_k : 1;\n }\n\n if (columnwise_option != FP8BlockwiseColumnwiseOption::NONE) {\n NVTE_CHECK(output_t.shape.size() == input.shape.size(),\n \"output_t must have same number of dimensions as input.\");\n\n if (output_t.shape.size() > 0) {\n if (columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY) {\n NVTE_CHECK(output_t.shape[0] == row_length, \"Wrong dimension 0 of output_t.\");\n for (size_t i = 1; i < output_t.shape.size(); ++i) {\n NVTE_CHECK(output_t.shape.at(i) == input.shape.at(i - 1), \"Wrong dimension in output_t\");\n }\n } else {\n NVTE_CHECK(columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT,\n \"Unexpected columnwise option enum value\");\n NVTE_CHECK(output_t.shape[0] == input.shape[0], \"Wrong dimension 0 of output_t.\");\n NVTE_CHECK(\n input.shape == output_t.shape,\n \"Input and output_t must have the same shape for columnwise non-transpose case.\");\n }\n }\n if (rowwise_option != FP8BlockwiseRowwiseOption::NONE) {\n // output may not be defined if rowwise quantization is not needed.\n NVTE_CHECK(output.dtype == output_t.dtype,\n \"output and output_t need to have the same dtype.\");\n }\n NVTE_CHECK(scale_inv_t.shape.size() == 2, \"Scale_t dimension must be 2.\");\n bool columnwise_compact = columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT;\n size_t scale_t_k = scale_inv_t.shape[1];\n scale_t_stride_x = columnwise_compact ? 1 : scale_t_k;\n scale_t_stride_y = columnwise_compact ? scale_t_k : 1;\n }\n auto output_dtype =\n rowwise_option != FP8BlockwiseRowwiseOption::NONE ? output.dtype : output_t.dtype;\n\n const size_t num_blocks_x = DIVUP(row_length, (size_t)kTileDim);\n const size_t num_blocks_y = DIVUP(num_rows, (size_t)kTileDim);\n\n const float* noop_ptr = reinterpret_cast(noop_tensor.dptr);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.dtype, InputType,\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(\n output_dtype, OutputType,\n\n dim3 grid(num_blocks_x, num_blocks_y, 1);\n\n const bool full_tile = row_length % kTileDim == 0 && num_rows % kTileDim == 0;\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n full_tile, kAligned,\n\n size_t smem_bytes = kSMemSize * sizeof(InputType);\n // shared memory must be requested up\n if (smem_bytes >= 48 * 1024) {\n hipError_t err = hipFuncSetAttribute(reinterpret_cast(&block_scaled_1d_cast_transpose_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);\n NVTE_CHECK(err == hipSuccess, \"Failed to set dynamic shared memory size.\");\n } block_scaled_1d_cast_transpose_kernel\n <<>>(\n reinterpret_cast(input.dptr),\n reinterpret_cast(output.dptr),\n reinterpret_cast(output_t.dptr),\n reinterpret_cast(scale_inv.dptr),\n reinterpret_cast(scale_inv_t.dptr), row_length, num_rows, scale_stride_x,\n scale_stride_y, scale_t_stride_x, scale_t_stride_y, epsilon, rowwise_option,\n columnwise_option, pow2_scale, noop_ptr);) // kAligned\n ) // OutputType\n ) // InputType\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n} // namespace transformer_engine::detail\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \n\n#include \"../util/rtc.h\"\n#include \"../util/string.h\"\n#include \"../utils.cuh\"\n#include \"cast_transpose.h\"\n\nnamespace transformer_engine::detail {\n\nnamespace {\n\n// String with RTC kernel implementation\n#include \"string_code_transpose_rtc_cast_transpose_cu.h\"\n\n// Hard-coded kernel parameters\nusing CType = float;\nconstexpr size_t warps_per_tile = 4;\nconstexpr size_t block_size = THREADS_PER_WARP * warps_per_tile;\n\n/* Performance heuristics for optimized kernel parameters */\nstruct KernelConfig {\n /** Vector load size */\n size_t load_size = 0;\n /** Vector store size to transposed output */\n size_t store_size = 0;\n\n /* Whether config is valid */\n bool valid = false;\n /* Number of CUDA blocks */\n size_t num_blocks = 0;\n\n /* Number of active SMs */\n size_t active_sm_count = 0;\n /* Elements per L1 cache load */\n size_t elements_per_load = 0;\n /* Elements per L1 cache store to cast output*/\n size_t elements_per_store_c = 0;\n /* Elements per L1 cache store to transposed output */\n size_t elements_per_store_t = 0;\n\n KernelConfig(size_t row_length, size_t num_rows, size_t itype_size, size_t otype_size,\n size_t load_size_, size_t store_size_, size_t sm_count)\n : load_size{load_size_}, store_size{store_size_} {\n // Check that tiles are correctly aligned\n constexpr size_t cache_line_size = 128;\n if (load_size % itype_size != 0 || store_size % otype_size != 0 ||\n cache_line_size % itype_size != 0 || cache_line_size % otype_size != 0) {\n return;\n }\n const size_t row_tile_elements = load_size * THREADS_PER_WARP / itype_size;\n const size_t col_tile_elements = store_size * THREADS_PER_WARP / otype_size;\n valid = (row_length % row_tile_elements == 0 && num_rows % col_tile_elements == 0);\n if (!valid) {\n return;\n }\n\n // Number of CUDA blocks\n num_blocks = (row_length / row_tile_elements) * (num_rows / col_tile_elements);\n\n // Parameters for performance model\n constexpr size_t warps_per_sm = 16; // Rough estimate for saturated SMs\n active_sm_count = std::min(DIVUP(num_blocks * warps_per_tile, warps_per_sm), sm_count);\n elements_per_load = (std::min(cache_line_size, row_tile_elements * itype_size) / itype_size);\n elements_per_store_c = (std::min(cache_line_size, row_tile_elements * otype_size) / otype_size);\n elements_per_store_t = (std::min(cache_line_size, col_tile_elements * otype_size) / otype_size);\n }\n\n /* Compare by estimated cost */\n bool operator<(const KernelConfig &other) const {\n if (this->valid && other.valid) {\n // cost ~ (1/elements_per_load\n // + 1/elements_per_store_c\n // + 1/elements_per_store_t) / active_sms\n // Note: Integer arithmetic ensures stable ordering\n const auto &l1 = this->elements_per_load;\n const auto &sc1 = this->elements_per_store_c;\n const auto &st1 = this->elements_per_store_t;\n const auto &p1 = this->active_sm_count;\n const auto &l2 = other.elements_per_load;\n const auto &sc2 = other.elements_per_store_c;\n const auto &st2 = other.elements_per_store_t;\n const auto &p2 = other.active_sm_count;\n const auto scale = l1 * sc1 * st1 * p1 * l2 * sc2 * st2 * p2;\n const auto cost1 = (scale / l1 + scale / sc1 + scale / st1) / p1;\n const auto cost2 = (scale / l2 + scale / sc2 + scale / st2) / p2;\n return cost1 < cost2;\n } else {\n return this->valid && !other.valid;\n }\n }\n};\n\ntemplate \n__global__ void __launch_bounds__(block_size) cast_transpose_general_kernel(\n const IType *__restrict__ const input, const CType *__restrict__ const noop,\n OType *__restrict__ const output_c, OType *__restrict__ const output_t,\n const CType *__restrict__ const scale_ptr, CType *__restrict__ const amax_ptr,\n CType *__restrict__ const scale_inv_ptr, const size_t row_length, const size_t num_rows) {\n if (noop != nullptr && noop[0] == 1.0f) return;\n\n // Vectorized load/store sizes\n constexpr size_t nvec_in = load_size / sizeof(IType);\n constexpr size_t nvec_out = store_size / sizeof(OType);\n using IVec = Vec;\n using OVecT = Vec;\n\n // Thread indices\n // Note: Block is interpreted as a warp_size x num_warps grid\n constexpr size_t bdimx = THREADS_PER_WARP;\n constexpr size_t bdimy = warps_per_tile;\n const size_t tid = threadIdx.x;\n const size_t tidx = tid % bdimx;\n const size_t tidy = tid / bdimx;\n const size_t bid = blockIdx.x;\n\n // Input tensors are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec_out x nvec_in subtiles\n constexpr size_t tile_dim_m = THREADS_PER_WARP * nvec_out;\n constexpr size_t tile_dim_n = THREADS_PER_WARP * nvec_in;\n\n // Position of tile within tensor\n const size_t num_tiles_m = (num_rows + tile_dim_m - 1) / tile_dim_m;\n const size_t tile_id_m = bid % num_tiles_m;\n const size_t tile_id_n = bid / num_tiles_m;\n const size_t tile_row = tile_id_m * tile_dim_m;\n const size_t tile_col = tile_id_n * tile_dim_n;\n\n // Number of nvec_out x nvec_in subtiles for each thread to\n // load/store\n constexpr size_t num_iterations = THREADS_PER_WARP / warps_per_tile;\n\n // FP8 factors\n const CType scale = scale_ptr == nullptr ? 1 : *scale_ptr;\n CType amax = 0;\n\n // Load input and store to registers\n // Note: Each thread loads num_iterations subtiles, computes amax,\n // casts type, and transposes in registers.\n OVecT local_output_t[nvec_in][num_iterations];\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n#pragma unroll\n for (size_t i2 = 0; i2 < nvec_out; ++i2) {\n const size_t row = tile_row + i1 * nvec_out + i2;\n const size_t col = tile_col + j1 * nvec_in;\n if (row < num_rows) {\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n if (col + j2 < row_length) {\n const CType in = input[row * row_length + col + j2];\n const OType out = OType(in * scale);\n __builtin_assume(amax >= 0);\n amax = fmaxf(fabsf(in), amax);\n output_c[row * row_length + col + j2] = out;\n local_output_t[j2][iter].data.elt[i2] = out;\n }\n }\n }\n }\n }\n\n // Copy transposed output from registers to global memory\n __shared__ OVecT shared_output_t[THREADS_PER_WARP][THREADS_PER_WARP + 1];\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n shared_output_t[j1][i1] = local_output_t[j2][iter];\n }\n __syncthreads();\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidx;\n const size_t j1 = tidy + iter * bdimy;\n const size_t row = tile_row + i1 * nvec_out;\n const size_t col = tile_col + j1 * nvec_in + j2;\n if (col < row_length) {\n#pragma unroll\n for (size_t i2 = 0; i2 < nvec_out; ++i2) {\n if (row + i2 < num_rows) {\n output_t[col * num_rows + row + i2] = shared_output_t[j1][i1].data.elt[i2];\n }\n }\n }\n }\n __syncthreads();\n }\n\n // Reduce amax over block\n if (amax_ptr != nullptr) {\n amax = reduce_max(amax, tidy);\n if (threadIdx.x == 0) {\n static_assert(std::is_same::value);\n atomicMaxFloat(amax_ptr, amax);\n }\n }\n\n // Update scale-inverse\n if (blockIdx.x == 0 && threadIdx.x == 0 && scale_inv_ptr != nullptr) {\n reciprocal(scale_inv_ptr, scale);\n }\n}\n\n} // namespace\n\nvoid cast_transpose(const Tensor &input, const Tensor &noop, Tensor *output_, cudaStream_t stream) {\n Tensor &output = *output_;\n\n CheckNoopTensor(noop, \"cast_transpose_noop\");\n CheckInputTensor(input, \"cast_transpose_input\");\n CheckOutputTensor(output, \"cast_transpose_output\");\n\n // Check that inputs and outputs are available\n NVTE_CHECK(input.has_data(), \"Input is not allocated\");\n NVTE_CHECK(output.has_data(), \"Output rowwise data is not allocated\");\n NVTE_CHECK(output.has_columnwise_data(), \"Output columnwise is not allocated\");\n\n // Flatten tensor to 2D\n NVTE_CHECK(input.data.shape == output.data.shape,\n \"Input and output shapes do not match (input=\", input.data.shape,\n \", output=\", output.data.shape);\n const size_t row_length = input.flat_last_dim();\n const size_t num_rows = input.flat_first_dim();\n NVTE_CHECK(output.flat_first_dim() == num_rows && output.flat_last_dim() == row_length,\n \"Invalid output dimensions (expected \", std::vector{num_rows, row_length},\n \", got \", std::vector{output.flat_first_dim(), output.flat_last_dim()}, \")\");\n\n // Check that cast and transposed output data matches\n NVTE_CHECK(output.data.dtype == output.columnwise_data.dtype,\n \"Cast and transposed output types must match.\");\n NVTE_CHECK(output.scale_inv.dptr == output.columnwise_scale_inv.dptr,\n \"Cast and transposed outputs need to share scale-inverse tensor.\");\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.dtype(), InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n output.dtype(), OutputType,\n if (is_tensor_scaling(output.scaling_mode)) {\n // delayed scaling and current scaling are two variants of per-tensor scaling\n\n constexpr const char *itype_name = TypeInfo::name;\n constexpr const char *otype_name = TypeInfo::name;\n constexpr size_t itype_size = sizeof(InputType);\n constexpr size_t otype_size = sizeof(OutputType);\n\n // Choose between runtime-compiled or statically-compiled kernel\n const bool aligned =\n (row_length % THREADS_PER_WARP == 0 && num_rows % THREADS_PER_WARP == 0);\n if (aligned && rtc::is_enabled()) { // Runtime-compiled tuned kernel\n // Pick kernel config\n std::vector kernel_configs;\n kernel_configs.reserve(16);\n const size_t sm_count = static_cast(cuda::sm_count());\n auto add_config = [&](size_t load_size, size_t store_size) {\n kernel_configs.emplace_back(row_length, num_rows, itype_size, otype_size, load_size,\n store_size, sm_count);\n };\n add_config(8, 8);\n add_config(4, 8);\n add_config(8, 4);\n add_config(4, 4);\n add_config(2, 8);\n add_config(8, 2);\n add_config(2, 4);\n add_config(4, 2);\n add_config(2, 2);\n add_config(1, 8);\n add_config(8, 1);\n add_config(1, 4);\n add_config(4, 1);\n add_config(1, 2);\n add_config(2, 1);\n add_config(1, 1);\n const auto &kernel_config =\n *std::min_element(kernel_configs.begin(), kernel_configs.end());\n NVTE_CHECK(kernel_config.valid, \"invalid kernel config\");\n const size_t load_size = kernel_config.load_size;\n const size_t store_size = kernel_config.store_size;\n const size_t num_blocks = kernel_config.num_blocks;\n\n // Compile NVRTC kernel if needed and launch\n auto &rtc_manager = rtc::KernelManager::instance();\n const std::string kernel_label = concat_strings(\n \"cast_transpose\"\n \",itype=\",\n itype_name, \",otype=\", otype_name, \",load_size=\", load_size,\n \",store_size=\", store_size);\n if (!rtc_manager.is_compiled(kernel_label)) {\n std::string code = string_code_transpose_rtc_cast_transpose_cu;\n code = regex_replace(code, \"__ITYPE__\", itype_name);\n code = regex_replace(code, \"__OTYPE__\", otype_name);\n code = regex_replace(code, \"__LOAD_SIZE__\", load_size);\n code = regex_replace(code, \"__STORE_SIZE__\", store_size);\n code = regex_replace(code, \"__WARPS_PER_TILE__\", warps_per_tile);\n code = regex_replace(code, \"__BLOCK_SIZE__\", block_size);\n rtc_manager.compile(kernel_label, \"cast_transpose_optimized_kernel\", code,\n \"transformer_engine/common/transpose/rtc/cast_transpose.cu\");\n }\n rtc_manager.launch(kernel_label, num_blocks, block_size, 0, stream,\n static_cast(input.data.dptr),\n reinterpret_cast(noop.data.dptr),\n static_cast(output.data.dptr),\n static_cast(output.columnwise_data.dptr),\n static_cast(output.scale.dptr),\n static_cast(output.amax.dptr),\n static_cast(output.scale_inv.dptr), row_length, num_rows);\n } else { // Statically-compiled general kernel\n constexpr size_t load_size = 4;\n constexpr size_t store_size = 4;\n constexpr size_t row_tile_size = load_size / itype_size * THREADS_PER_WARP;\n constexpr size_t col_tile_size = store_size / otype_size * THREADS_PER_WARP;\n const int num_blocks =\n (DIVUP(row_length, row_tile_size) * DIVUP(num_rows, col_tile_size));\n\n cast_transpose_general_kernel\n <<>>(\n static_cast(input.data.dptr),\n reinterpret_cast(noop.data.dptr),\n static_cast(output.data.dptr),\n static_cast(output.columnwise_data.dptr),\n static_cast(output.scale.dptr),\n static_cast(output.amax.dptr),\n static_cast(output.scale_inv.dptr), row_length, num_rows);\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n } else {\n NVTE_ERROR(\"Not implemented scaling mode: \", to_string(output.scaling_mode));\n }); // NOLINT(*)\n ); // NOLINT(*)\n}\n\n} // namespace transformer_engine::detail\n\nvoid nvte_cast_transpose(const NVTETensor input, NVTETensor output, cudaStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose);\n using namespace transformer_engine;\n auto noop = Tensor();\n transformer_engine::detail::cast_transpose(*convertNVTETensorCheck(input), noop,\n convertNVTETensor(output), stream);\n}\n\nvoid nvte_cast_transpose_with_noop(const NVTETensor input, const NVTETensor noop, NVTETensor output,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose_with_noop);\n using namespace transformer_engine;\n transformer_engine::detail::cast_transpose(*convertNVTETensorCheck(input),\n *convertNVTETensorCheck(noop),\n convertNVTETensor(output), stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \n\n#include \"../util/rtc.h\"\n#include \"../util/string.h\"\n#include \"../utils.cuh\"\n#include \"cast_transpose.h\"\n\nnamespace transformer_engine::detail {\n\nnamespace {\n\n// String with RTC kernel implementation\n#include \"string_code_transpose_rtc_cast_transpose_cu.h\"\n\n// Hard-coded kernel parameters\nusing CType = float;\nconstexpr size_t warps_per_tile = 4;\nconstexpr size_t block_size = THREADS_PER_WARP * warps_per_tile;\n\n/* Performance heuristics for optimized kernel parameters */\nstruct KernelConfig {\n /** Vector load size */\n size_t load_size = 0;\n /** Vector store size to transposed output */\n size_t store_size = 0;\n\n /* Whether config is valid */\n bool valid = false;\n /* Number of CUDA blocks */\n size_t num_blocks = 0;\n\n /* Number of active SMs */\n size_t active_sm_count = 0;\n /* Elements per L1 cache load */\n size_t elements_per_load = 0;\n /* Elements per L1 cache store to cast output*/\n size_t elements_per_store_c = 0;\n /* Elements per L1 cache store to transposed output */\n size_t elements_per_store_t = 0;\n\n KernelConfig(size_t row_length, size_t num_rows, size_t itype_size, size_t otype_size,\n size_t load_size_, size_t store_size_, size_t sm_count)\n : load_size{load_size_}, store_size{store_size_} {\n // Check that tiles are correctly aligned\n constexpr size_t cache_line_size = 128;\n if (load_size % itype_size != 0 || store_size % otype_size != 0 ||\n cache_line_size % itype_size != 0 || cache_line_size % otype_size != 0) {\n return;\n }\n const size_t row_tile_elements = load_size * THREADS_PER_WARP / itype_size;\n const size_t col_tile_elements = store_size * THREADS_PER_WARP / otype_size;\n valid = (row_length % row_tile_elements == 0 && num_rows % col_tile_elements == 0);\n if (!valid) {\n return;\n }\n\n // Number of CUDA blocks\n num_blocks = (row_length / row_tile_elements) * (num_rows / col_tile_elements);\n\n // Parameters for performance model\n constexpr size_t warps_per_sm = 16; // Rough estimate for saturated SMs\n active_sm_count = std::min(DIVUP(num_blocks * warps_per_tile, warps_per_sm), sm_count);\n elements_per_load = (std::min(cache_line_size, row_tile_elements * itype_size) / itype_size);\n elements_per_store_c = (std::min(cache_line_size, row_tile_elements * otype_size) / otype_size);\n elements_per_store_t = (std::min(cache_line_size, col_tile_elements * otype_size) / otype_size);\n }\n\n /* Compare by estimated cost */\n bool operator<(const KernelConfig &other) const {\n if (this->valid && other.valid) {\n // cost ~ (1/elements_per_load\n // + 1/elements_per_store_c\n // + 1/elements_per_store_t) / active_sms\n // Note: Integer arithmetic ensures stable ordering\n const auto &l1 = this->elements_per_load;\n const auto &sc1 = this->elements_per_store_c;\n const auto &st1 = this->elements_per_store_t;\n const auto &p1 = this->active_sm_count;\n const auto &l2 = other.elements_per_load;\n const auto &sc2 = other.elements_per_store_c;\n const auto &st2 = other.elements_per_store_t;\n const auto &p2 = other.active_sm_count;\n const auto scale = l1 * sc1 * st1 * p1 * l2 * sc2 * st2 * p2;\n const auto cost1 = (scale / l1 + scale / sc1 + scale / st1) / p1;\n const auto cost2 = (scale / l2 + scale / sc2 + scale / st2) / p2;\n return cost1 < cost2;\n } else {\n return this->valid && !other.valid;\n }\n }\n};\n\ntemplate \n__global__ void __launch_bounds__(block_size) cast_transpose_general_kernel(\n const IType *__restrict__ const input, const CType *__restrict__ const noop,\n OType *__restrict__ const output_c, OType *__restrict__ const output_t,\n const CType *__restrict__ const scale_ptr, CType *__restrict__ const amax_ptr,\n CType *__restrict__ const scale_inv_ptr, const size_t row_length, const size_t num_rows) {\n if (noop != nullptr && noop[0] == 1.0f) return;\n\n // Vectorized load/store sizes\n constexpr size_t nvec_in = load_size / sizeof(IType);\n constexpr size_t nvec_out = store_size / sizeof(OType);\n using IVec = Vec;\n using OVecT = Vec;\n\n // Thread indices\n // Note: Block is interpreted as a warp_size x num_warps grid\n constexpr size_t bdimx = THREADS_PER_WARP;\n constexpr size_t bdimy = warps_per_tile;\n const size_t tid = threadIdx.x;\n const size_t tidx = tid % bdimx;\n const size_t tidy = tid / bdimx;\n const size_t bid = blockIdx.x;\n\n // Input tensors are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec_out x nvec_in subtiles\n constexpr size_t tile_dim_m = THREADS_PER_WARP * nvec_out;\n constexpr size_t tile_dim_n = THREADS_PER_WARP * nvec_in;\n\n // Position of tile within tensor\n const size_t num_tiles_m = (num_rows + tile_dim_m - 1) / tile_dim_m;\n const size_t tile_id_m = bid % num_tiles_m;\n const size_t tile_id_n = bid / num_tiles_m;\n const size_t tile_row = tile_id_m * tile_dim_m;\n const size_t tile_col = tile_id_n * tile_dim_n;\n\n // Number of nvec_out x nvec_in subtiles for each thread to\n // load/store\n constexpr size_t num_iterations = THREADS_PER_WARP / warps_per_tile;\n\n // FP8 factors\n const CType scale = scale_ptr == nullptr ? 1 : *scale_ptr;\n CType amax = 0;\n\n // Load input and store to registers\n // Note: Each thread loads num_iterations subtiles, computes amax,\n // casts type, and transposes in registers.\n OVecT local_output_t[nvec_in][num_iterations];\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n#pragma unroll\n for (size_t i2 = 0; i2 < nvec_out; ++i2) {\n const size_t row = tile_row + i1 * nvec_out + i2;\n const size_t col = tile_col + j1 * nvec_in;\n if (row < num_rows) {\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n if (col + j2 < row_length) {\n const CType in = input[row * row_length + col + j2];\n const OType out = OType(in * scale);\n __builtin_assume(amax >= 0);\n amax = fmaxf(fabsf(in), amax);\n output_c[row * row_length + col + j2] = out;\n local_output_t[j2][iter].data.elt[i2] = out;\n }\n }\n }\n }\n }\n\n // Copy transposed output from registers to global memory\n __shared__ OVecT shared_output_t[THREADS_PER_WARP][THREADS_PER_WARP + 1];\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n shared_output_t[j1][i1] = local_output_t[j2][iter];\n }\n __syncthreads();\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidx;\n const size_t j1 = tidy + iter * bdimy;\n const size_t row = tile_row + i1 * nvec_out;\n const size_t col = tile_col + j1 * nvec_in + j2;\n if (col < row_length) {\n#pragma unroll\n for (size_t i2 = 0; i2 < nvec_out; ++i2) {\n if (row + i2 < num_rows) {\n output_t[col * num_rows + row + i2] = shared_output_t[j1][i1].data.elt[i2];\n }\n }\n }\n }\n __syncthreads();\n }\n\n // Reduce amax over block\n if (amax_ptr != nullptr) {\n amax = reduce_max(amax, tidy);\n if (threadIdx.x == 0) {\n static_assert(std::is_same::value);\n atomicMaxFloat(amax_ptr, amax);\n }\n }\n\n // Update scale-inverse\n if (blockIdx.x == 0 && threadIdx.x == 0 && scale_inv_ptr != nullptr) {\n reciprocal(scale_inv_ptr, scale);\n }\n}\n\n} // namespace\n\nvoid cast_transpose(const Tensor &input, const Tensor &noop, Tensor *output_, hipStream_t stream) {\n Tensor &output = *output_;\n\n CheckNoopTensor(noop, \"cast_transpose_noop\");\n CheckInputTensor(input, \"cast_transpose_input\");\n CheckOutputTensor(output, \"cast_transpose_output\");\n\n // Check that inputs and outputs are available\n NVTE_CHECK(input.has_data(), \"Input is not allocated\");\n NVTE_CHECK(output.has_data(), \"Output rowwise data is not allocated\");\n NVTE_CHECK(output.has_columnwise_data(), \"Output columnwise is not allocated\");\n\n // Flatten tensor to 2D\n NVTE_CHECK(input.data.shape == output.data.shape,\n \"Input and output shapes do not match (input=\", input.data.shape,\n \", output=\", output.data.shape);\n const size_t row_length = input.flat_last_dim();\n const size_t num_rows = input.flat_first_dim();\n NVTE_CHECK(output.flat_first_dim() == num_rows && output.flat_last_dim() == row_length,\n \"Invalid output dimensions (expected \", std::vector{num_rows, row_length},\n \", got \", std::vector{output.flat_first_dim(), output.flat_last_dim()}, \")\");\n\n // Check that cast and transposed output data matches\n NVTE_CHECK(output.data.dtype == output.columnwise_data.dtype,\n \"Cast and transposed output types must match.\");\n NVTE_CHECK(output.scale_inv.dptr == output.columnwise_scale_inv.dptr,\n \"Cast and transposed outputs need to share scale-inverse tensor.\");\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.dtype(), InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n output.dtype(), OutputType,\n if (is_tensor_scaling(output.scaling_mode)) {\n // delayed scaling and current scaling are two variants of per-tensor scaling\n\n constexpr const char *itype_name = TypeInfo::name;\n constexpr const char *otype_name = TypeInfo::name;\n constexpr size_t itype_size = sizeof(InputType);\n constexpr size_t otype_size = sizeof(OutputType);\n\n // Choose between runtime-compiled or statically-compiled kernel\n const bool aligned =\n (row_length % THREADS_PER_WARP == 0 && num_rows % THREADS_PER_WARP == 0);\n if (aligned && rtc::is_enabled()) { // Runtime-compiled tuned kernel\n // Pick kernel config\n std::vector kernel_configs;\n kernel_configs.reserve(16);\n const size_t sm_count = static_cast(cuda::sm_count());\n auto add_config = [&](size_t load_size, size_t store_size) {\n kernel_configs.emplace_back(row_length, num_rows, itype_size, otype_size, load_size,\n store_size, sm_count);\n };\n add_config(8, 8);\n add_config(4, 8);\n add_config(8, 4);\n add_config(4, 4);\n add_config(2, 8);\n add_config(8, 2);\n add_config(2, 4);\n add_config(4, 2);\n add_config(2, 2);\n add_config(1, 8);\n add_config(8, 1);\n add_config(1, 4);\n add_config(4, 1);\n add_config(1, 2);\n add_config(2, 1);\n add_config(1, 1);\n const auto &kernel_config =\n *std::min_element(kernel_configs.begin(), kernel_configs.end());\n NVTE_CHECK(kernel_config.valid, \"invalid kernel config\");\n const size_t load_size = kernel_config.load_size;\n const size_t store_size = kernel_config.store_size;\n const size_t num_blocks = kernel_config.num_blocks;\n\n // Compile NVRTC kernel if needed and launch\n auto &rtc_manager = rtc::KernelManager::instance();\n const std::string kernel_label = concat_strings(\n \"cast_transpose\"\n \",itype=\",\n itype_name, \",otype=\", otype_name, \",load_size=\", load_size,\n \",store_size=\", store_size);\n if (!rtc_manager.is_compiled(kernel_label)) {\n std::string code = string_code_transpose_rtc_cast_transpose_cu;\n code = regex_replace(code, \"__ITYPE__\", itype_name);\n code = regex_replace(code, \"__OTYPE__\", otype_name);\n code = regex_replace(code, \"__LOAD_SIZE__\", load_size);\n code = regex_replace(code, \"__STORE_SIZE__\", store_size);\n code = regex_replace(code, \"__WARPS_PER_TILE__\", warps_per_tile);\n code = regex_replace(code, \"__BLOCK_SIZE__\", block_size);\n rtc_manager.compile(kernel_label, \"cast_transpose_optimized_kernel\", code,\n \"transformer_engine/common/transpose/rtc/cast_transpose.cu\");\n }\n rtc_manager.launch(kernel_label, num_blocks, block_size, 0, stream,\n static_cast(input.data.dptr),\n reinterpret_cast(noop.data.dptr),\n static_cast(output.data.dptr),\n static_cast(output.columnwise_data.dptr),\n static_cast(output.scale.dptr),\n static_cast(output.amax.dptr),\n static_cast(output.scale_inv.dptr), row_length, num_rows);\n } else { // Statically-compiled general kernel\n constexpr size_t load_size = 4;\n constexpr size_t store_size = 4;\n constexpr size_t row_tile_size = load_size / itype_size * THREADS_PER_WARP;\n constexpr size_t col_tile_size = store_size / otype_size * THREADS_PER_WARP;\n const int num_blocks =\n (DIVUP(row_length, row_tile_size) * DIVUP(num_rows, col_tile_size));\n\n cast_transpose_general_kernel\n <<>>(\n static_cast(input.data.dptr),\n reinterpret_cast(noop.data.dptr),\n static_cast(output.data.dptr),\n static_cast(output.columnwise_data.dptr),\n static_cast(output.scale.dptr),\n static_cast(output.amax.dptr),\n static_cast(output.scale_inv.dptr), row_length, num_rows);\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n } else {\n NVTE_ERROR(\"Not implemented scaling mode: \", to_string(output.scaling_mode));\n }); // NOLINT(*)\n ); // NOLINT(*)\n}\n\n} // namespace transformer_engine::detail\n\nvoid nvte_cast_transpose(const NVTETensor input, NVTETensor output, hipStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose);\n using namespace transformer_engine;\n auto noop = Tensor();\n transformer_engine::detail::cast_transpose(*convertNVTETensorCheck(input), noop,\n convertNVTETensor(output), stream);\n}\n\nvoid nvte_cast_transpose_with_noop(const NVTETensor input, const NVTETensor noop, NVTETensor output,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_cast_transpose_with_noop);\n using namespace transformer_engine;\n transformer_engine::detail::cast_transpose(*convertNVTETensorCheck(input),\n *convertNVTETensorCheck(noop),\n convertNVTETensor(output), stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/transpose/cast_transpose.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\n\nnamespace {\n\n// Parameters to tune\nconstexpr int n_warps_per_tile = 4;\nconstexpr int threads_per_block = THREADS_PER_WARP * n_warps_per_tile;\nconstexpr int desired_load_size = 8;\nconstexpr int desired_store_size = 8;\nconstexpr int kMaxTensorsPerKernel = 64; // Args must be <4 KB\n\nstruct MultiCastTransposeArgs {\n // (input) Data buffers for input tensors\n void* input_list[kMaxTensorsPerKernel];\n // (output) Data buffers for cast output tensors\n void* output_c_list[kMaxTensorsPerKernel];\n // (output) Data buffers for transpose output tensors\n void* output_t_list[kMaxTensorsPerKernel];\n // (input) Scaling factor for output tensors\n void* scale_list[kMaxTensorsPerKernel];\n // (output) AMAX's of input tensors\n void* amax_list[kMaxTensorsPerKernel];\n // (output) Inverse of scaling factor for output tensors\n void* scale_inv_list[kMaxTensorsPerKernel];\n // Input matrix heights\n int num_rows_list[kMaxTensorsPerKernel];\n // Input matrix widths\n int row_length_list[kMaxTensorsPerKernel];\n // Prefix sum (with leading zero) of CUDA blocks needed for each\n // tensor\n int block_range[kMaxTensorsPerKernel + 1];\n // Number of tensors being processed by kernel\n int num_tensors;\n};\n\ntemplate \n__global__ void __launch_bounds__(threads_per_block)\n multi_cast_transpose_kernel(MultiCastTransposeArgs args) {\n using IVec = Vec;\n using OVecC = Vec;\n using OVecT = Vec;\n\n // Thread indices\n // Note: Block is interpreted as a warp_size x num_warps grid\n constexpr int bdimx = THREADS_PER_WARP;\n constexpr int bdimy = n_warps_per_tile;\n const int tid = threadIdx.x;\n const int tidx = tid % bdimx;\n const int tidy = tid / bdimx;\n const int bid = blockIdx.x;\n\n // Input tensors are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec_out x nvec_in subtiles\n constexpr int tile_dim_m = THREADS_PER_WARP * nvec_out;\n constexpr int tile_dim_n = THREADS_PER_WARP * nvec_in;\n\n // Number of nvec_out x nvec_in subtiles for each thread to\n // load/store\n constexpr int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n\n // Find tensor corresponding to block\n int tensor_id = 0;\n while (args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n const IType* input = reinterpret_cast(args.input_list[tensor_id]);\n OType* output_c = reinterpret_cast(args.output_c_list[tensor_id]);\n OType* output_t = reinterpret_cast(args.output_t_list[tensor_id]);\n const CType* scale_ptr = reinterpret_cast(args.scale_list[tensor_id]);\n const CType scale = scale_ptr == nullptr ? 1 : *scale_ptr;\n CType* amax_ptr = reinterpret_cast(args.amax_list[tensor_id]);\n CType* scale_inv_ptr = reinterpret_cast(args.scale_inv_list[tensor_id]);\n const int num_rows = args.num_rows_list[tensor_id];\n const int row_length = args.row_length_list[tensor_id];\n\n // Find position of tile within tensor\n const int num_tiles_n = (row_length + tile_dim_n - 1) / tile_dim_n;\n const int tile_id = bid - args.block_range[tensor_id];\n const int tile_id_m = tile_id / num_tiles_n;\n const int tile_id_n = tile_id % num_tiles_n;\n const int tile_row = tile_id_m * tile_dim_m;\n const int tile_col = tile_id_n * tile_dim_n;\n\n // Load input and store to registers\n // Note: Each thread loads n_iterations subtiles, casts to output\n // type, and transposes in registers.\n OVecT local_output_t[nvec_in][n_iterations];\n CType local_amax = 0;\n#pragma unroll\n for (int iter = 0; iter < n_iterations; ++iter) {\n const int i1 = tidy + iter * bdimy;\n const int j1 = tidx;\n#pragma unroll\n for (int i2 = 0; i2 < nvec_out; ++i2) {\n const int row = tile_row + i1 * nvec_out + i2;\n const int col = tile_col + j1 * nvec_in;\n IVec local_input;\n OVecC local_output_c;\n if constexpr (aligned) {\n local_input.load_from(&input[row * row_length + col]);\n } else {\n local_input.clear();\n if (row < num_rows) {\n#pragma unroll\n for (int j2 = 0; j2 < nvec_in; ++j2) {\n if (col + j2 < row_length) {\n local_input.data.elt[j2] = input[row * row_length + col + j2];\n }\n }\n }\n }\n#pragma unroll\n for (int j2 = 0; j2 < nvec_in; ++j2) {\n const CType x = CType(local_input.data.elt[j2]);\n const OType y = OType(scale * x);\n local_output_c.data.elt[j2] = y;\n local_output_t[j2][iter].data.elt[i2] = y;\n __builtin_assume(local_amax >= 0);\n local_amax = fmaxf(fabsf(x), local_amax);\n }\n if constexpr (aligned) {\n local_output_c.store_to(&output_c[row * row_length + col]);\n } else {\n if (row < num_rows) {\n#pragma unroll\n for (int j2 = 0; j2 < nvec_in; ++j2) {\n if (col + j2 < row_length) {\n output_c[row * row_length + col + j2] = local_output_c.data.elt[j2];\n }\n }\n }\n }\n }\n }\n\n // Copy transposed output from registers to global memory\n __shared__ OVecT shared_output_t[THREADS_PER_WARP][THREADS_PER_WARP + 1];\n#pragma unroll\n for (int j2 = 0; j2 < nvec_in; ++j2) {\n#pragma unroll\n for (int iter = 0; iter < n_iterations; ++iter) {\n const int i1 = tidy + iter * bdimy;\n const int j1 = tidx;\n shared_output_t[j1][i1] = local_output_t[j2][iter];\n }\n __syncthreads();\n#pragma unroll\n for (int iter = 0; iter < n_iterations; ++iter) {\n const int i1 = tidx;\n const int j1 = tidy + iter * bdimy;\n const int row = tile_row + i1 * nvec_out;\n const int col = tile_col + j1 * nvec_in + j2;\n if constexpr (aligned) {\n shared_output_t[j1][i1].store_to(&output_t[col * num_rows + row]);\n } else {\n if (col < row_length) {\n#pragma unroll\n for (int i2 = 0; i2 < nvec_out; ++i2) {\n if (row + i2 < num_rows) {\n output_t[col * num_rows + row + i2] = shared_output_t[j1][i1].data.elt[i2];\n }\n }\n }\n }\n }\n __syncthreads();\n }\n\n // Finalize fp8 factors\n local_amax = reduce_max(local_amax, tidy);\n if (tid == 0) {\n static_assert(std::is_same::value);\n if (amax_ptr != nullptr) atomicMaxFloat(amax_ptr, local_amax);\n }\n if (tile_id == 0 && tid == 0 && scale_inv_ptr != nullptr) {\n reciprocal(scale_inv_ptr, scale);\n }\n}\n\n} // namespace\n\nvoid multi_cast_transpose(const std::vector input_list, std::vector output_list,\n cudaStream_t stream) {\n // Check that number of tensors is valid\n NVTE_CHECK(output_list.size() == input_list.size(),\n \"Number of input and output tensors must match\");\n if (input_list.empty()) {\n return;\n }\n\n // Check that tensor properties are valid\n DType itype = input_list[0]->data.dtype;\n DType otype = output_list[0]->dtype();\n for (size_t tensor_id = 0; tensor_id < input_list.size(); ++tensor_id) {\n const auto& input = *input_list[tensor_id];\n const auto& output = *output_list[tensor_id];\n CheckInputTensor(input, \"multi_cast_transpose_input_\" + std::to_string(tensor_id));\n CheckInputTensor(output, \"multi_cast_transpose_output_\" + std::to_string(tensor_id));\n //std::cout << *static_cast(output.data.dptr) << std::endl;\n NVTE_CHECK(output.has_data() && output.has_columnwise_data(),\n \"Both rowwise and columnwise output data needs to be allocated.\");\n\n NVTE_CHECK(input.data.dtype == itype, \"Input tensor types do not match.\");\n NVTE_CHECK(output.data.dtype == otype, \"C output tensor types do not match.\");\n NVTE_CHECK(output.data.dtype == otype, \"T output tensor types do not match.\");\n\n NVTE_CHECK(input.data.shape.size() == 2, \"Input tensor must have 2 dimensions, but shape is \",\n input.data.shape);\n NVTE_CHECK(output.data.shape == input.data.shape, \"C output tensor shape \", output.data.shape,\n \"does not match input tensor shape \", input.data.shape);\n NVTE_CHECK(output.columnwise_data.shape.size() == 2, \"T output tensor shape \",\n output.columnwise_data.shape, \"does not match input tensor shape \",\n input.data.shape);\n NVTE_CHECK(output.columnwise_data.shape[0] == input.data.shape[1], \"T output tensor shape \",\n output.columnwise_data.shape, \"does not match input tensor shape \",\n input.data.shape);\n NVTE_CHECK(output.columnwise_data.shape[1] == input.data.shape[0], \"T output tensor shape \",\n output.columnwise_data.shape, \"does not match input tensor shape \",\n input.data.shape);\n }\n\n // Input matrices are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec_out x nvec_in subtiles\n const int tile_dim_m = THREADS_PER_WARP * desired_store_size * 8 / typeToNumBits(otype);\n const int tile_dim_n = THREADS_PER_WARP * desired_load_size * 8 / typeToNumBits(itype);\n\n // Add tensors to kernel argument struct\n MultiCastTransposeArgs kernel_args_aligned, kernel_args_unaligned;\n kernel_args_aligned.num_tensors = 0;\n kernel_args_aligned.block_range[0] = 0;\n kernel_args_unaligned.num_tensors = 0;\n kernel_args_unaligned.block_range[0] = 0;\n for (size_t tensor_id = 0; tensor_id < input_list.size(); ++tensor_id) {\n // Launch kernel if argument struct is full\n if (kernel_args_aligned.num_tensors == kMaxTensorsPerKernel) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n itype, InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n otype, OutputType, constexpr int nvec_in = desired_load_size / sizeof(InputType);\n constexpr int nvec_out = desired_store_size / sizeof(OutputType);\n const int n_blocks = kernel_args_aligned.block_range[kernel_args_aligned.num_tensors];\n multi_cast_transpose_kernel\n <<>>(kernel_args_aligned);); // NOLINT(*)\n ); // NOLINT(*)\n NVTE_CHECK_CUDA(cudaGetLastError());\n kernel_args_aligned.num_tensors = 0;\n }\n if (kernel_args_unaligned.num_tensors == kMaxTensorsPerKernel) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n itype, InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n otype, OutputType, constexpr int nvec_in = desired_load_size / sizeof(InputType);\n constexpr int nvec_out = desired_store_size / sizeof(OutputType);\n const int n_blocks =\n kernel_args_unaligned.block_range[kernel_args_unaligned.num_tensors];\n multi_cast_transpose_kernel\n <<>>(kernel_args_unaligned);); // NOLINT(*)\n ); // NOLINT(*)\n NVTE_CHECK_CUDA(cudaGetLastError());\n kernel_args_unaligned.num_tensors = 0;\n }\n\n // Calculate number of thread blocks needed for tensor\n const int num_rows = input_list[tensor_id]->data.shape[0];\n const int row_length = input_list[tensor_id]->data.shape[1];\n const int num_tiles_m = (num_rows + tile_dim_m - 1) / tile_dim_m;\n const int num_tiles_n = (row_length + tile_dim_n - 1) / tile_dim_n;\n const int num_tiles = num_tiles_m * num_tiles_n;\n\n // Figure out whether to use aligned or unaligned kernel\n const bool aligned =\n ((num_tiles_m * tile_dim_m == num_rows) && (num_tiles_n * tile_dim_n == row_length));\n auto& kernel_args = aligned ? kernel_args_aligned : kernel_args_unaligned;\n\n // Add tensor to kernel argument struct\n const int pos = kernel_args.num_tensors;\n kernel_args.input_list[pos] = const_cast(input_list[tensor_id]->data.dptr);\n kernel_args.output_c_list[pos] = output_list[tensor_id]->data.dptr;\n kernel_args.output_t_list[pos] = output_list[tensor_id]->columnwise_data.dptr;\n kernel_args.scale_list[pos] = output_list[tensor_id]->scale.dptr;\n kernel_args.amax_list[pos] = output_list[tensor_id]->amax.dptr;\n kernel_args.scale_inv_list[pos] = output_list[tensor_id]->scale_inv.dptr;\n kernel_args.num_rows_list[pos] = num_rows;\n kernel_args.row_length_list[pos] = row_length;\n kernel_args.block_range[pos + 1] = kernel_args.block_range[pos] + num_tiles;\n kernel_args.num_tensors++;\n }\n\n // Launch kernel\n if (kernel_args_aligned.num_tensors > 0) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n itype, InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n otype, OutputType, constexpr int nvec_in = desired_load_size / sizeof(InputType);\n constexpr int nvec_out = desired_store_size / sizeof(OutputType);\n const int n_blocks = kernel_args_aligned.block_range[kernel_args_aligned.num_tensors];\n multi_cast_transpose_kernel\n <<>>(kernel_args_aligned);); // NOLINT(*)\n ); // NOLINT(*)\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n if (kernel_args_unaligned.num_tensors > 0) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n itype, InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n otype, OutputType, constexpr int nvec_in = desired_load_size / sizeof(InputType);\n constexpr int nvec_out = desired_store_size / sizeof(OutputType);\n const int n_blocks =\n kernel_args_unaligned.block_range[kernel_args_unaligned.num_tensors];\n multi_cast_transpose_kernel\n <<>>(kernel_args_unaligned);); // NOLINT(*)\n ); // NOLINT(*)\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\n} // namespace transformer_engine\n\nvoid nvte_multi_cast_transpose(size_t num_tensors, const NVTETensor* input_list,\n NVTETensor* output_list, cudaStream_t stream) {\n NVTE_API_CALL(nvte_multi_cast_transpose);\n using namespace transformer_engine;\n std::vector input_list_, output_list_;\n for (size_t i = 0; i < num_tensors; ++i) {\n input_list_.push_back(convertNVTETensorCheck(input_list[i]));\n output_list_.push_back(convertNVTETensorCheck(output_list[i]));\n }\n multi_cast_transpose(input_list_, output_list_, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\n\nnamespace {\n\n// Parameters to tune\nconstexpr int n_warps_per_tile = 4;\nconstexpr int threads_per_block = THREADS_PER_WARP * n_warps_per_tile;\nconstexpr int desired_load_size = 8;\nconstexpr int desired_store_size = 8;\nconstexpr int kMaxTensorsPerKernel = 64; // Args must be <4 KB\n\nstruct MultiCastTransposeArgs {\n // (input) Data buffers for input tensors\n void* input_list[kMaxTensorsPerKernel];\n // (output) Data buffers for cast output tensors\n void* output_c_list[kMaxTensorsPerKernel];\n // (output) Data buffers for transpose output tensors\n void* output_t_list[kMaxTensorsPerKernel];\n // (input) Scaling factor for output tensors\n void* scale_list[kMaxTensorsPerKernel];\n // (output) AMAX's of input tensors\n void* amax_list[kMaxTensorsPerKernel];\n // (output) Inverse of scaling factor for output tensors\n void* scale_inv_list[kMaxTensorsPerKernel];\n // Input matrix heights\n int num_rows_list[kMaxTensorsPerKernel];\n // Input matrix widths\n int row_length_list[kMaxTensorsPerKernel];\n // Prefix sum (with leading zero) of CUDA blocks needed for each\n // tensor\n int block_range[kMaxTensorsPerKernel + 1];\n // Number of tensors being processed by kernel\n int num_tensors;\n};\n\ntemplate \n__global__ void __launch_bounds__(threads_per_block)\n multi_cast_transpose_kernel(MultiCastTransposeArgs args) {\n using IVec = Vec;\n using OVecC = Vec;\n using OVecT = Vec;\n\n // Thread indices\n // Note: Block is interpreted as a warp_size x num_warps grid\n constexpr int bdimx = THREADS_PER_WARP;\n constexpr int bdimy = n_warps_per_tile;\n const int tid = threadIdx.x;\n const int tidx = tid % bdimx;\n const int tidy = tid / bdimx;\n const int bid = blockIdx.x;\n\n // Input tensors are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec_out x nvec_in subtiles\n constexpr int tile_dim_m = THREADS_PER_WARP * nvec_out;\n constexpr int tile_dim_n = THREADS_PER_WARP * nvec_in;\n\n // Number of nvec_out x nvec_in subtiles for each thread to\n // load/store\n constexpr int n_iterations = THREADS_PER_WARP / n_warps_per_tile;\n\n // Find tensor corresponding to block\n int tensor_id = 0;\n while (args.block_range[tensor_id + 1] <= bid) {\n ++tensor_id;\n }\n const IType* input = reinterpret_cast(args.input_list[tensor_id]);\n OType* output_c = reinterpret_cast(args.output_c_list[tensor_id]);\n OType* output_t = reinterpret_cast(args.output_t_list[tensor_id]);\n const CType* scale_ptr = reinterpret_cast(args.scale_list[tensor_id]);\n const CType scale = scale_ptr == nullptr ? 1 : *scale_ptr;\n CType* amax_ptr = reinterpret_cast(args.amax_list[tensor_id]);\n CType* scale_inv_ptr = reinterpret_cast(args.scale_inv_list[tensor_id]);\n const int num_rows = args.num_rows_list[tensor_id];\n const int row_length = args.row_length_list[tensor_id];\n\n // Find position of tile within tensor\n const int num_tiles_n = (row_length + tile_dim_n - 1) / tile_dim_n;\n const int tile_id = bid - args.block_range[tensor_id];\n const int tile_id_m = tile_id / num_tiles_n;\n const int tile_id_n = tile_id % num_tiles_n;\n const int tile_row = tile_id_m * tile_dim_m;\n const int tile_col = tile_id_n * tile_dim_n;\n\n // Load input and store to registers\n // Note: Each thread loads n_iterations subtiles, casts to output\n // type, and transposes in registers.\n OVecT local_output_t[nvec_in][n_iterations];\n CType local_amax = 0;\n#pragma unroll\n for (int iter = 0; iter < n_iterations; ++iter) {\n const int i1 = tidy + iter * bdimy;\n const int j1 = tidx;\n#pragma unroll\n for (int i2 = 0; i2 < nvec_out; ++i2) {\n const int row = tile_row + i1 * nvec_out + i2;\n const int col = tile_col + j1 * nvec_in;\n IVec local_input;\n OVecC local_output_c;\n if constexpr (aligned) {\n local_input.load_from(&input[row * row_length + col]);\n } else {\n local_input.clear();\n if (row < num_rows) {\n#pragma unroll\n for (int j2 = 0; j2 < nvec_in; ++j2) {\n if (col + j2 < row_length) {\n local_input.data.elt[j2] = input[row * row_length + col + j2];\n }\n }\n }\n }\n#pragma unroll\n for (int j2 = 0; j2 < nvec_in; ++j2) {\n const CType x = CType(local_input.data.elt[j2]);\n const OType y = OType(scale * x);\n local_output_c.data.elt[j2] = y;\n local_output_t[j2][iter].data.elt[i2] = y;\n __builtin_assume(local_amax >= 0);\n local_amax = fmaxf(fabsf(x), local_amax);\n }\n if constexpr (aligned) {\n local_output_c.store_to(&output_c[row * row_length + col]);\n } else {\n if (row < num_rows) {\n#pragma unroll\n for (int j2 = 0; j2 < nvec_in; ++j2) {\n if (col + j2 < row_length) {\n output_c[row * row_length + col + j2] = local_output_c.data.elt[j2];\n }\n }\n }\n }\n }\n }\n\n // Copy transposed output from registers to global memory\n __shared__ OVecT shared_output_t[THREADS_PER_WARP][THREADS_PER_WARP + 1];\n#pragma unroll\n for (int j2 = 0; j2 < nvec_in; ++j2) {\n#pragma unroll\n for (int iter = 0; iter < n_iterations; ++iter) {\n const int i1 = tidy + iter * bdimy;\n const int j1 = tidx;\n shared_output_t[j1][i1] = local_output_t[j2][iter];\n }\n __syncthreads();\n#pragma unroll\n for (int iter = 0; iter < n_iterations; ++iter) {\n const int i1 = tidx;\n const int j1 = tidy + iter * bdimy;\n const int row = tile_row + i1 * nvec_out;\n const int col = tile_col + j1 * nvec_in + j2;\n if constexpr (aligned) {\n shared_output_t[j1][i1].store_to(&output_t[col * num_rows + row]);\n } else {\n if (col < row_length) {\n#pragma unroll\n for (int i2 = 0; i2 < nvec_out; ++i2) {\n if (row + i2 < num_rows) {\n output_t[col * num_rows + row + i2] = shared_output_t[j1][i1].data.elt[i2];\n }\n }\n }\n }\n }\n __syncthreads();\n }\n\n // Finalize fp8 factors\n local_amax = reduce_max(local_amax, tidy);\n if (tid == 0) {\n static_assert(std::is_same::value);\n if (amax_ptr != nullptr) atomicMaxFloat(amax_ptr, local_amax);\n }\n if (tile_id == 0 && tid == 0 && scale_inv_ptr != nullptr) {\n reciprocal(scale_inv_ptr, scale);\n }\n}\n\n} // namespace\n\nvoid multi_cast_transpose(const std::vector input_list, std::vector output_list,\n hipStream_t stream) {\n // Check that number of tensors is valid\n NVTE_CHECK(output_list.size() == input_list.size(),\n \"Number of input and output tensors must match\");\n if (input_list.empty()) {\n return;\n }\n\n // Check that tensor properties are valid\n DType itype = input_list[0]->data.dtype;\n DType otype = output_list[0]->dtype();\n for (size_t tensor_id = 0; tensor_id < input_list.size(); ++tensor_id) {\n const auto& input = *input_list[tensor_id];\n const auto& output = *output_list[tensor_id];\n CheckInputTensor(input, \"multi_cast_transpose_input_\" + std::to_string(tensor_id));\n CheckInputTensor(output, \"multi_cast_transpose_output_\" + std::to_string(tensor_id));\n //std::cout << *static_cast(output.data.dptr) << std::endl;\n NVTE_CHECK(output.has_data() && output.has_columnwise_data(),\n \"Both rowwise and columnwise output data needs to be allocated.\");\n\n NVTE_CHECK(input.data.dtype == itype, \"Input tensor types do not match.\");\n NVTE_CHECK(output.data.dtype == otype, \"C output tensor types do not match.\");\n NVTE_CHECK(output.data.dtype == otype, \"T output tensor types do not match.\");\n\n NVTE_CHECK(input.data.shape.size() == 2, \"Input tensor must have 2 dimensions, but shape is \",\n input.data.shape);\n NVTE_CHECK(output.data.shape == input.data.shape, \"C output tensor shape \", output.data.shape,\n \"does not match input tensor shape \", input.data.shape);\n NVTE_CHECK(output.columnwise_data.shape.size() == 2, \"T output tensor shape \",\n output.columnwise_data.shape, \"does not match input tensor shape \",\n input.data.shape);\n NVTE_CHECK(output.columnwise_data.shape[0] == input.data.shape[1], \"T output tensor shape \",\n output.columnwise_data.shape, \"does not match input tensor shape \",\n input.data.shape);\n NVTE_CHECK(output.columnwise_data.shape[1] == input.data.shape[0], \"T output tensor shape \",\n output.columnwise_data.shape, \"does not match input tensor shape \",\n input.data.shape);\n }\n\n // Input matrices are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec_out x nvec_in subtiles\n const int tile_dim_m = THREADS_PER_WARP * desired_store_size * 8 / typeToNumBits(otype);\n const int tile_dim_n = THREADS_PER_WARP * desired_load_size * 8 / typeToNumBits(itype);\n\n // Add tensors to kernel argument struct\n MultiCastTransposeArgs kernel_args_aligned, kernel_args_unaligned;\n kernel_args_aligned.num_tensors = 0;\n kernel_args_aligned.block_range[0] = 0;\n kernel_args_unaligned.num_tensors = 0;\n kernel_args_unaligned.block_range[0] = 0;\n for (size_t tensor_id = 0; tensor_id < input_list.size(); ++tensor_id) {\n // Launch kernel if argument struct is full\n if (kernel_args_aligned.num_tensors == kMaxTensorsPerKernel) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n itype, InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n otype, OutputType, constexpr int nvec_in = desired_load_size / sizeof(InputType);\n constexpr int nvec_out = desired_store_size / sizeof(OutputType);\n const int n_blocks = kernel_args_aligned.block_range[kernel_args_aligned.num_tensors];\n multi_cast_transpose_kernel\n <<>>(kernel_args_aligned);); // NOLINT(*)\n ); // NOLINT(*)\n NVTE_CHECK_CUDA(hipGetLastError());\n kernel_args_aligned.num_tensors = 0;\n }\n if (kernel_args_unaligned.num_tensors == kMaxTensorsPerKernel) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n itype, InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n otype, OutputType, constexpr int nvec_in = desired_load_size / sizeof(InputType);\n constexpr int nvec_out = desired_store_size / sizeof(OutputType);\n const int n_blocks =\n kernel_args_unaligned.block_range[kernel_args_unaligned.num_tensors];\n multi_cast_transpose_kernel\n <<>>(kernel_args_unaligned);); // NOLINT(*)\n ); // NOLINT(*)\n NVTE_CHECK_CUDA(hipGetLastError());\n kernel_args_unaligned.num_tensors = 0;\n }\n\n // Calculate number of thread blocks needed for tensor\n const int num_rows = input_list[tensor_id]->data.shape[0];\n const int row_length = input_list[tensor_id]->data.shape[1];\n const int num_tiles_m = (num_rows + tile_dim_m - 1) / tile_dim_m;\n const int num_tiles_n = (row_length + tile_dim_n - 1) / tile_dim_n;\n const int num_tiles = num_tiles_m * num_tiles_n;\n\n // Figure out whether to use aligned or unaligned kernel\n const bool aligned =\n ((num_tiles_m * tile_dim_m == num_rows) && (num_tiles_n * tile_dim_n == row_length));\n auto& kernel_args = aligned ? kernel_args_aligned : kernel_args_unaligned;\n\n // Add tensor to kernel argument struct\n const int pos = kernel_args.num_tensors;\n kernel_args.input_list[pos] = const_cast(input_list[tensor_id]->data.dptr);\n kernel_args.output_c_list[pos] = output_list[tensor_id]->data.dptr;\n kernel_args.output_t_list[pos] = output_list[tensor_id]->columnwise_data.dptr;\n kernel_args.scale_list[pos] = output_list[tensor_id]->scale.dptr;\n kernel_args.amax_list[pos] = output_list[tensor_id]->amax.dptr;\n kernel_args.scale_inv_list[pos] = output_list[tensor_id]->scale_inv.dptr;\n kernel_args.num_rows_list[pos] = num_rows;\n kernel_args.row_length_list[pos] = row_length;\n kernel_args.block_range[pos + 1] = kernel_args.block_range[pos] + num_tiles;\n kernel_args.num_tensors++;\n }\n\n // Launch kernel\n if (kernel_args_aligned.num_tensors > 0) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n itype, InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n otype, OutputType, constexpr int nvec_in = desired_load_size / sizeof(InputType);\n constexpr int nvec_out = desired_store_size / sizeof(OutputType);\n const int n_blocks = kernel_args_aligned.block_range[kernel_args_aligned.num_tensors];\n multi_cast_transpose_kernel\n <<>>(kernel_args_aligned);); // NOLINT(*)\n ); // NOLINT(*)\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n if (kernel_args_unaligned.num_tensors > 0) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n itype, InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(\n otype, OutputType, constexpr int nvec_in = desired_load_size / sizeof(InputType);\n constexpr int nvec_out = desired_store_size / sizeof(OutputType);\n const int n_blocks =\n kernel_args_unaligned.block_range[kernel_args_unaligned.num_tensors];\n multi_cast_transpose_kernel\n <<>>(kernel_args_unaligned);); // NOLINT(*)\n ); // NOLINT(*)\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\n} // namespace transformer_engine\n\nvoid nvte_multi_cast_transpose(size_t num_tensors, const NVTETensor* input_list,\n NVTETensor* output_list, hipStream_t stream) {\n NVTE_API_CALL(nvte_multi_cast_transpose);\n using namespace transformer_engine;\n std::vector input_list_, output_list_;\n for (size_t i = 0; i < num_tensors; ++i) {\n input_list_.push_back(convertNVTETensorCheck(input_list[i]));\n output_list_.push_back(convertNVTETensorCheck(output_list[i]));\n }\n multi_cast_transpose(input_list_, output_list_, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/transpose/multi_cast_transpose.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/recipe/recipe_common.cuh\"\n#include \"common/transpose/cast_transpose.h\"\n#include \"common/util/curanddx.hpp\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n\nnamespace transformer_engine {\n\n#if CUDA_VERSION >= 12080\nnamespace quantize_transpose_nvfp4 {\nnamespace {\n\nusing std::int32_t;\nusing std::uint32_t;\nusing std::uint8_t;\n\nusing transformer_engine::detail::TypeExtrema;\n\n// clang-format off\n/*\n\nStep 1: Load input to shared memory\n* shard memory: 128x128 elements with type=InputType (below graph doesn't consider padding)\n* 8 warps\n* Loop 8 times\n* What each thread does in each loop:\n * 8 elements are read from the input at a time\n * 2 elements are written to the shared memory at a time, for a total of 4 times\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| T0 | T1 | T2 | T3 | T4 | T5 | T6 | T7 | T8 | T9 | T10 | T11 | T12 | T13 | T14 | T15 |\n| T16 | T17 | T18 | T19 | T20 | T21 | T22 | T23 | T24 | T25 | T26 | T27 | T28 | T29 | T30 | T31 |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| Warp 1 |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| Warp 7 |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n| ... |\n| Loop 8 times |\n| ... |\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n\nStep 2: Cast and store to output_c\n* shard memory: 128x128 elements with type=InputType (below graph doesn't consider padding)\n* 8 warps\n* Loop 4 times\n* What each thread does in each loop:\n * 2 elements are read from the shared memory at a time, for a total of 8 times\n * Every 8 consecutive threads do reduction and calculate the amax of each row\n * 16 elements are quantized and write to output_c at a time\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| T0 | T1 | T2 | T3 | T4 | T5 | T6 | T7 |\n| T8 | T9 | T10 | T11 | T12 | T13 | T14 | T15 |\n| T16 | T17 | T18 | T19 | T20 | T21 | T22 | T23 |\n| T24 | T25 | T26 | T27 | T28 | T29 | T30 | T31 |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| |\n| Warp 1 |\n| |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| |\n| Warp 7 |\n| |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n| ... |\n| Loop 4 times |\n| ... |\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n\nStep 3: Transpose, cast and store to output_t\n* shard memory: 128x128 elements with type=InputType (below graph doesn't consider padding)\n* 8 warps\n* Loop 2 times\n* What each thread does in each loop:\n * 2 elements (in a row) are read from the shared memory at a time, for a total of 16 times\n * Every 8 consecutive threads do reduction and calculate the amax of each column\n * 16 elements are quantized and write to output_c at a time, for a total of 2 times\n+------8 elements-------+------8 elements-------+-----40 elements-------+------8 elements-------+------8 elements-------+------8 elements-------+-----40 elements-------+------8 elements-------+\n| T0 | T8 | T16 | T24 | | | | T0 | T8 | T16 | T24 | | | |\n| T1 | T9 | T17 | T25 | | | | T1 | T9 | T17 | T25 | | | |\n| T2 | T10 | T18 | T26 | | | | T2 | T10 | T18 | T26 | | | |\n| T3 | T11 | T19 | T27 | Warp 1 | ... | Warp 7 | T3 | T11 | T19 | T27 | Warp 1 | ... | Warp 7 |\n| T4 | T12 | T20 | T28 | | | | T4 | T12 | T20 | T28 | | | |\n| T5 | T13 | T21 | T29 | | | | T5 | T13 | T21 | T29 | | | |\n| T6 | T14 | T22 | T30 | | | | T6 | T14 | T22 | T30 | | | |\n| T7 | T15 | T23 | T31 | | | | T7 | T15 | T23 | T31 | | | |\n+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+\n\n*/\n// clang-format on\n\nconstexpr int kThreadsPerWarp = 32;\n\n// for fp4, we use uint8_t to store 2 fp4 numbers\nconstexpr int kNFP4PerContainer = 2;\n\n// Hyperparameters for performance tuning\nconstexpr int kTileDim = 128;\n// constexpr int kScaleDim = 32;\nconstexpr int kNVecIn = 8; // The number of elements each LDG touches\nconstexpr int kNVecOut = 16; // The number of elements each STG touches\nconstexpr int kNVecSMem = 2; // The number of elements each LDS/STS touches\nconstexpr int kThreadsPerBlock = 256; // Thread block size, 8 warps in total\n\n// Auto-calculated constants, do not modify directly)\nstatic_assert(kNVecIn % kNVecSMem == 0, \"kNVecIn must be divisible by kNVecSMem\");\nstatic_assert(kNVecOut % kNVecSMem == 0, \"kNVecOut must be divisible by kNVecSMem\");\nconstexpr int kSMemRow = kTileDim;\nconstexpr int kSMemCol = (kTileDim / kNVecSMem) + 1;\nconstexpr int kSMemSize = kSMemRow * kSMemCol * kNVecSMem;\nconstexpr int kNumThreadsLoad = kTileDim / kNVecIn; // 16\nconstexpr int kNumThreadsStore = kTileDim / kNVecOut; // 8\n// constexpr int kNumThreadsReduce = kScaleDim / kNVecOut;\nstatic_assert(kNumThreadsLoad <= kThreadsPerWarp, \"kNumThreadsLoad must be <= kThreadsPerWarp\");\nstatic_assert(kNumThreadsStore <= kThreadsPerWarp, \"kNumThreadsStore must be <= kThreadsPerWarp\");\n\n// for 2D block scaling, we need to reduce amax in warp\nstatic __device__ constexpr unsigned int WARP_REDUCE_AMAX_GROUP_MASKS[8] = {\n 0x01010101, 0x02020202, 0x04040404, 0x08080808, 0x10101010, 0x20202020, 0x40404040, 0x80808080};\n\n// max for every group_size elements in warp\ntemplate \n__device__ __forceinline__ float groupMax(float val, unsigned int groupMask) {\n for (int offset = group_size / 2; offset > 0; offset /= 2) {\n val = max(val, __shfl_down_sync(groupMask, val, offset * shfl_down_stride));\n }\n return val;\n}\n\ntemplate \n__device__ __forceinline__ ScaleType\nComputeDecodeScaleFP4(const float amax, const float global_encode_scale_multiplier) {\n float decode_scale = amax * global_encode_scale_multiplier;\n decode_scale = fminf(decode_scale, TypeExtrema::max);\n return static_cast(decode_scale);\n}\n\ntemplate \n__device__ __forceinline__ float ComputeEncodeScaleFP4(ScaleType decode_scale,\n const float global_decode_scale) {\n return fminf(1.0f / (static_cast(decode_scale) * global_decode_scale),\n TypeExtrema::max);\n}\n\ntemplate \n__device__ __forceinline__ float ComputeOutputFP4(IType input, float encode_scale) {\n return static_cast(input) * encode_scale;\n}\n\n__device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax) {\n constexpr float fp8_max = TypeExtrema::max;\n constexpr float fp4_max = TypeExtrema::max;\n float global_encode_scale = fp8_max * fp4_max / global_amax;\n // If scale is infinity, return max value of float32\n global_encode_scale = fminf(global_encode_scale, TypeExtrema::max);\n // If global amax is 0 or infinity, return 1\n if (global_amax == 0.f || global_encode_scale == 0.f) {\n return 1.f;\n }\n return global_encode_scale;\n}\n\n__device__ __forceinline__ uint32_t get_rbits(\n transformer_engine::curanddx::detail::philox4x32_native_state&\n rng, // NVTE_BUILD_NUM_PHILOX_ROUNDS rounds of philox4x32\n uint4& random_uint4, int& rnd_idx) {\n if (rnd_idx == 4) {\n rnd_idx = 0;\n random_uint4 = rng.generate4();\n }\n\n // Treat uint4 as an array of 4x uint32_t elements for indexing\n const uint32_t* const rbits_arr = reinterpret_cast(&random_uint4);\n const uint32_t rbits = rbits_arr[rnd_idx++];\n return rbits;\n}\n\ntemplate \n__device__ __forceinline__ size_t scale_factor_swizzled_offset(size_t row_idx, size_t col_idx,\n uint32_t col_length) {\n // This function takes in indices from the scale factor matrix and returns an offset in the\n // swizzled format. row_idx, col_idx are original indices from the scale factor matrix (unswizzled\n // index). col_length is the column length of the scale factor matrix. tile_scales_inv is the\n // pointer to the scale factor matrix.\n\n // https://github.com/NVIDIA/cutlass/blob/main/media/docs/cpp/blackwell_functionality.md#scale-factor-layouts\n // For any scale factor matrix, it's 512B base block. Each base block consists of 128 rows and 4\n // columns. Base block is divided into 4 column blocks, each column block has 32 rows and 4\n // columns.\n\n // NOTE: There are not a lot of good illustrations about the swizzled scale factor matrix.\n // To think in high level, the swizzled scale factor matrix could be composed as:\n // unswizzled_scale_factor_matrix = torch.empty((M, N // 16), dtype=torch.uint8)\n // cbg_cnt = N // 16 // 4 # Assuming N is divisible by 64\n // rb_cnt = M // 128 # Assuming M is divisible by 128\n // tmp = unswizzled_scale_factor_matrix.reshape(rb_cnt, 4, 32, cbg_cnt, 4)\n // tmp = torch.permute(tmp, (0, 3, 2, 1, 4))\n // swizzled_scale_factor_matrix = tmp.reshape((-1, 128, 4))\n\n constexpr uint32_t kTotalRowsPerBaseBlock = 128;\n constexpr uint32_t kRowsPerBaseBlockCol = 32;\n constexpr uint32_t kColsPerBaseBlockCol = 4;\n\n const size_t rb = row_idx / kTotalRowsPerBaseBlock;\n const size_t rem = row_idx % kTotalRowsPerBaseBlock;\n const size_t d4 = rem / kRowsPerBaseBlockCol;\n const size_t d3 = rem % kRowsPerBaseBlockCol;\n const size_t cbg = col_idx / kColsPerBaseBlockCol;\n const size_t d5 = col_idx % kColsPerBaseBlockCol;\n\n const size_t cbg_cnt = DIVUP(col_length, kColsPerBaseBlockCol);\n // row-major offset in the logical shape\n // (rb_cnt , cbg_cnt , 32 , 4 , 4)\n // Magic number 16 below comes from the fact we have kColsPerBaseBlockCol = 4, and d4 ([0-128] /\n // 32 = [0-4])\n return ((rb * cbg_cnt + cbg) * kRowsPerBaseBlockCol + d3) * 16 + d4 * kColsPerBaseBlockCol + d5;\n}\n\n__device__ __forceinline__ __nv_fp4x4_e2m1 cvt_fp32_to_fp4_4x_with_stochastic_rounding(\n const float2 in01, const float2 in23, const uint32_t rbits) {\n constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING;\n if constexpr (has_rs) {\n uint16_t out_4x;\n asm volatile(\n \"{\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %0, {%3, %4, %1, %2}, %5; \\n\\t\"\n \"}\"\n : \"=h\"(out_4x)\n : \"f\"(in01.y), \"f\"(in01.x), \"f\"(in23.y), \"f\"(in23.x), \"r\"(rbits));\n return *reinterpret_cast<__nv_fp4x4_e2m1*>(&out_4x);\n } else {\n NVTE_DEVICE_ERROR(\n \"FP4 cvt.rs PTX instructions are architecture-specific. \"\n \"Try recompiling with sm_XXXa instead of sm_XXX.\");\n uint16_t dummy = 0;\n return *reinterpret_cast<__nv_fp4x4_e2m1*>(&dummy);\n }\n}\n\n__device__ __forceinline__ __nv_fp4x4_e2m1 cvt_fp32_to_fp4_4x_with_rn(const float2 in01,\n const float2 in23,\n const uint32_t rbits) {\n constexpr bool has_fp4 = ARCH_BLACKWELL_FAMILY;\n if constexpr (has_fp4) {\n // NOTE: rbits unused for rn.\n uint32_t out_4x; // Only need 16 bit. Using 32 bit container for packing.\n asm volatile(\n \"{\\n\"\n \".reg.b8 f0; \\n\\t\"\n \".reg.b8 f1; \\n\\t\"\n \"cvt.rn.satfinite.e2m1x2.f32 f0, %1, %2;\\n\\t\"\n \"cvt.rn.satfinite.e2m1x2.f32 f1, %3, %4;\\n\\t\"\n \"mov.b32 %0, {f0, f1, f0, f1};\\n\\t\"\n \"}\"\n : \"=r\"(out_4x)\n : \"f\"(in01.y), \"f\"(in01.x), \"f\"(in23.y), \"f\"(in23.x));\n return reinterpret_cast<__nv_fp4x4_e2m1*>(&out_4x)[0];\n } else {\n NVTE_DEVICE_ERROR(\n \"FP4 cvt PTX instructions are architecture-specific. \"\n \"Try recompiling with sm_XXXa instead of sm_XXX.\");\n uint16_t dummy = 0;\n return *reinterpret_cast<__nv_fp4x4_e2m1*>(&dummy);\n }\n}\n\ntemplate \n__device__ __forceinline__ __nv_fp4x4_e2m1 cvt_fp32_to_fp4_4x(const float2 in01, const float2 in23,\n const uint32_t rbits) {\n if constexpr (kApplyStochasticRounding) {\n return cvt_fp32_to_fp4_4x_with_stochastic_rounding(in01, in23, rbits);\n } else {\n return cvt_fp32_to_fp4_4x_with_rn(in01, in23, rbits);\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpose_kernel(\n const IType* const input, const float* global_amax, OType* const output_c,\n OType* const output_t, ScaleType* const tile_scales_inv_c, ScaleType* const tile_scales_inv_t,\n const size_t row_length, const size_t num_rows, const size_t scale_stride_x,\n const size_t scale_stride_y, const size_t scale_t_stride_x, const size_t scale_t_stride_y,\n const size_t kScaleBlockDim, const float epsilon, const size_t* rng_state,\n const float* noop_ptr) {\n constexpr int kNVecContainer = kNVecOut / kNFP4PerContainer;\n using SMemVec = Vec;\n using OVec = Vec;\n union IVec {\n Vec input_type;\n Vec smem_type;\n };\n\n if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {\n return;\n }\n\n const size_t block_idx_x = blockIdx.x;\n const size_t block_idx_y = blockIdx.y;\n const size_t rng_sequence =\n threadIdx.x + block_idx_x * kThreadsPerBlock + block_idx_y * gridDim.x * kThreadsPerBlock;\n const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0;\n const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0;\n\n transformer_engine::curanddx::detail::philox4x32_native_state rng;\n rng.init(rng_seed, rng_sequence, rng_offset);\n uint4 random_uint4 = kApplyStochasticRounding ? rng.generate4() : uint4{0, 0, 0, 0};\n\n int rnd_idx =\n 0; // Index of the random number. It increments each time when used and resets to 0 if reaches 4x\n\n extern __shared__ char smem_base[];\n SMemVec* smem = reinterpret_cast(&smem_base[0]);\n\n // 2D block scaling is not supported for E8 scaling MXFP4 or for colwise only mode.\n // Instead of static_assert, return early if these invalid modes are detected.\n if constexpr (kIs2DBlockScaling && kIsE8Scaling) {\n return;\n }\n if constexpr (kIs2DBlockScaling && !kReturnIdentity) {\n return;\n }\n // for 128x128 block, 2D block scaling means there will be 8x8 amax values for nvfp4, 4x4 for 2D mxfp4\n // use constexpr to define the size, when not using 2D, use minimal size 1x1\n constexpr int kFP4BlockScalingSize = 16;\n constexpr int k2DBlockAmaxDim = kIs2DBlockScaling ? (kTileDim / kFP4BlockScalingSize) : 1;\n constexpr int kNumRowsPerWarp = kThreadsPerWarp / kNumThreadsStore; // 4\n constexpr int k2DBlockAmaxReduceDim =\n kIs2DBlockScaling ? (kFP4BlockScalingSize / kNumRowsPerWarp) : 1;\n __shared__ CType amax_smem_red[k2DBlockAmaxDim][k2DBlockAmaxDim][k2DBlockAmaxReduceDim];\n __shared__ CType amax_smem[k2DBlockAmaxDim][k2DBlockAmaxDim];\n\n // Step 1: Load input to shared memory\n {\n constexpr int r_stride = kThreadsPerBlock / kNumThreadsLoad; // stride in rows of shared memory\n constexpr int num_iterations = kTileDim / r_stride;\n const int c_s =\n (threadIdx.x % kNumThreadsLoad) * (kNVecIn / kNVecSMem); // Column in shared memory\n int r_s = threadIdx.x / kNumThreadsLoad; // Row in shared memory\n const size_t c_g = block_idx_x * kTileDim + c_s * kNVecSMem; // Column in global memory\n size_t r_g = block_idx_y * kTileDim + r_s; // Row in global memory\n const size_t stride_g = static_cast(r_stride) * row_length; // Stride in global memory\n const size_t num_ele = (c_g < row_length ? min(static_cast(kNVecIn), row_length - c_g)\n : 0); // For not aligned case\n const IType* input_g = &input[r_g * row_length + c_g]; // Input address in global memory\n#pragma unroll\n for (int iter = 0; iter < num_iterations; ++iter) {\n IVec input_vec;\n // Step 1.1: Load from global memory (input) to registers\n if constexpr (kAligned) {\n input_vec.input_type.load_from(input_g);\n } else {\n if (r_g < num_rows) {\n input_vec.input_type.load_from_elts(input_g, 0, num_ele);\n } else {\n input_vec.input_type.clear();\n }\n }\n // Step 1.2: Write to shared memory\n#pragma unroll\n for (int i = 0; i < kNVecIn / kNVecSMem; ++i) {\n int c = c_s + i;\n int r = r_s;\n smem[r * kSMemCol + c] = input_vec.smem_type.data.elt[i];\n }\n // Step 1.3: Update input address, row index of shared memory, (and row index of global memory\n // for not aligned case)\n input_g += stride_g;\n r_s += r_stride;\n if constexpr (!kAligned) {\n r_g += r_stride;\n }\n }\n }\n\n __syncthreads();\n\n const int kNumThreadsReduce = kScaleBlockDim / kNVecOut;\n const float global_encode_scale =\n kIsE8Scaling ? 1.0f : ComputeGlobalEncodeScaleFP4(global_amax[0]);\n constexpr float fp4_max_inv = 1.0f / TypeExtrema::max;\n const float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n const float global_decode_scale = 1.0 / global_encode_scale;\n\n // Step 2: Cast and store to output_c\n if constexpr (kReturnIdentity) {\n constexpr int r_stride =\n kThreadsPerBlock / kNumThreadsStore; // stride in rows of shared memory\n constexpr int num_iterations = kTileDim / r_stride;\n const int c_s =\n (threadIdx.x % kNumThreadsStore) * (kNVecOut / kNVecSMem); // Column in shared memory\n int r_s = threadIdx.x / kNumThreadsStore; // Row in shared memory\n const size_t c_g = block_idx_x * kTileDim + c_s * kNVecSMem; // Column in global memory\n size_t r_g = block_idx_y * kTileDim + r_s; // Row in global memory\n const size_t stride_g = static_cast(r_stride) * row_length; // Stride in global memory\n const size_t num_ele =\n (c_g < row_length ? min(static_cast(kNVecOut / kNFP4PerContainer),\n (row_length - c_g) / kNFP4PerContainer)\n : 0); // For not aligned case\n OType* output_g =\n &output_c[(r_g * row_length + c_g) / kNFP4PerContainer]; // Output address in global memory\n // Each kNumThreadsStore threads form a warp process one row, we need to find the lane id of\n // the first thread to do the reduction.\n const unsigned src_lane =\n (threadIdx.x % kThreadsPerWarp) / kNumThreadsReduce * kNumThreadsReduce;\n // This mask represents which threads should do the reduction together.\n const unsigned mask = ((1 << kNumThreadsReduce) - 1) << src_lane;\n const bool is_src_lane = (threadIdx.x % kNumThreadsReduce) == 0;\n#pragma unroll\n for (int iter = 0; iter < num_iterations; ++iter) {\n SMemVec smem_vec[kNVecOut / kNVecSMem];\n // Step 2.1: Load from shared memory to registers\n#pragma unroll\n for (int i = 0; i < kNVecOut / kNVecSMem; ++i) {\n int c = c_s + i;\n int r = r_s;\n smem_vec[i] = smem[r * kSMemCol + c];\n }\n // Step 2.2: Compute local amax\n CType amax = 0;\n#pragma unroll\n for (int i = 0; i < kNVecOut / kNVecSMem; ++i) {\n#pragma unroll\n for (int j = 0; j < kNVecSMem; ++j) {\n __builtin_assume(amax >= 0);\n amax = fmaxf(amax, fabsf(smem_vec[i].data.elt[j]));\n }\n }\n // Step 2.3: Reduce amax\n if constexpr (kIsE8Scaling) {\n#pragma unroll\n for (int delta = kNumThreadsReduce / 2; delta > 0; delta /= 2) {\n const float other_amax = __shfl_down_sync(mask, amax, delta);\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n amax = __shfl_sync(mask, amax, src_lane);\n }\n // doing shuffle sync for 2D block scaling (not applicable for E8 scaling)\n if constexpr (kIs2DBlockScaling) {\n // first amax shuffle sync in warp, then reduce in smem\n // T0 T8 T16 T24 should do amax reduction together\n constexpr int kNumRowsPerIter = kThreadsPerBlock / kNumThreadsStore; // 32\n int warp_idx = threadIdx.x / kThreadsPerWarp; // 0 ~ 7\n int tid_in_warp_x = threadIdx.x % kNumThreadsStore;\n int tid_in_warp_y = (threadIdx.x / kNumThreadsStore) % kNumRowsPerWarp;\n CType amax_warp_reduced = groupMax(\n amax, WARP_REDUCE_AMAX_GROUP_MASKS[tid_in_warp_x]);\n // now T0 ~ T8 in each warp has the reduced amax values\n int data_row_idx = iter * kNumRowsPerIter + warp_idx * kNumRowsPerWarp + tid_in_warp_y;\n if (tid_in_warp_y == 0) {\n amax_smem_red[data_row_idx / kFP4BlockScalingSize][tid_in_warp_x]\n [warp_idx % k2DBlockAmaxReduceDim] = amax_warp_reduced;\n }\n __syncthreads();\n\n if (data_row_idx % kFP4BlockScalingSize == 0) {\n CType amax_2d = 0.0;\n for (int i = 0; i < k2DBlockAmaxReduceDim; i++) {\n amax_2d = fmaxf(amax_2d,\n amax_smem_red[data_row_idx / kFP4BlockScalingSize][tid_in_warp_x][i]);\n }\n amax_smem[data_row_idx / kFP4BlockScalingSize][tid_in_warp_x] = amax_2d;\n }\n __syncthreads();\n // every thread now knows 2D amax\n amax = amax_smem[data_row_idx / kFP4BlockScalingSize][tid_in_warp_x];\n }\n // Step 2.4: Compute scale\n const size_t row_idx = block_idx_y * kTileDim + r_s;\n float row_global_encode_scale = global_encode_scale;\n if constexpr (kRowScaledNVFP4) {\n row_global_encode_scale =\n row_idx < num_rows ? ComputeGlobalEncodeScaleFP4(global_amax[row_idx]) : 1.0f;\n }\n const float row_global_encode_scale_multiplier =\n kRowScaledNVFP4 ? row_global_encode_scale * fp4_max_inv : global_encode_scale_multiplier;\n const float row_global_decode_scale =\n kRowScaledNVFP4 ? 1.0f / row_global_encode_scale : global_decode_scale;\n ScaleType scale_inv =\n ComputeDecodeScaleFP4(amax, row_global_encode_scale_multiplier);\n float encode_scale = ComputeEncodeScaleFP4(scale_inv, row_global_decode_scale);\n // Step 2.5: Write scale_inv\n bool write_scale_inv = is_src_lane;\n if constexpr (!kAligned) {\n write_scale_inv &= (r_g < num_rows);\n write_scale_inv &= (c_g < row_length);\n }\n if (write_scale_inv) {\n size_t row_idx = block_idx_y * kTileDim + r_s;\n size_t col_idx = block_idx_x * (kNumThreadsStore / kNumThreadsReduce) +\n (threadIdx.x % kNumThreadsStore) / kNumThreadsReduce;\n if constexpr (kSwizzledScale) {\n size_t offset = scale_factor_swizzled_offset(\n row_idx, col_idx, DIVUP(row_length, kScaleBlockDim));\n tile_scales_inv_c[offset] = scale_inv;\n } else {\n tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv;\n }\n }\n // Step 2.6: Quantize\n OVec output_vec;\n#pragma unroll\n for (int i = 0; i < kNVecOut / kNVecSMem; i += 2) {\n // Pack two elements into __nv_bfloat162\n float2 f2_a;\n float2 f2_b;\n f2_a.x = ComputeOutputFP4(smem_vec[i].data.elt[0], encode_scale);\n f2_a.y = ComputeOutputFP4(smem_vec[i].data.elt[1], encode_scale);\n f2_b.x = ComputeOutputFP4(smem_vec[i + 1].data.elt[0], encode_scale);\n f2_b.y = ComputeOutputFP4(smem_vec[i + 1].data.elt[1], encode_scale);\n const uint32_t rbits = kApplyStochasticRounding ? get_rbits(rng, random_uint4, rnd_idx) : 0;\n // Convert to __nv_fp4x4_e2m1\n __nv_fp4x4_e2m1 out_4x = cvt_fp32_to_fp4_4x(f2_a, f2_b, rbits);\n\n output_vec.data.elt[i] = reinterpret_cast<__nv_fp4x2_storage_t*>(&out_4x)[0];\n output_vec.data.elt[i + 1] = reinterpret_cast<__nv_fp4x2_storage_t*>(&out_4x)[1];\n }\n // Step 2.7: Store output_c\n if constexpr (kAligned) {\n output_vec.store_to(output_g);\n } else {\n if (r_g < num_rows) {\n output_vec.store_to_elts(output_g, 0, num_ele);\n }\n }\n // Step 2.8: Update output address, row index of shared memory (and row index of global memory\n // for not aligned case)\n output_g += stride_g / kNFP4PerContainer;\n r_s += r_stride;\n if constexpr (!kAligned) {\n r_g += r_stride;\n }\n }\n }\n\n // Step 3: Transpose, cast and store to output_t\n if constexpr (kReturnTranspose) {\n constexpr int c_stride =\n kThreadsPerBlock / kNumThreadsStore; // Stride in columns of shared memory\n constexpr int num_iterations = kTileDim / (c_stride * kNVecSMem);\n const int r_s = (threadIdx.x % kNumThreadsStore) * kNVecOut; // Row in shared memory\n int c_s = threadIdx.x / kNumThreadsStore; // Column in shared memory\n size_t r_g = block_idx_x * kTileDim + c_s * kNVecSMem; // Row in global memory\n const size_t c_g = block_idx_y * kTileDim + r_s; // Column in global memory\n const size_t stride_g =\n static_cast(c_stride) * kNVecSMem * num_rows; // Stride in global memory\n const size_t num_ele = (c_g < num_rows ? min(static_cast(kNVecOut / kNFP4PerContainer),\n (num_rows - c_g) / kNFP4PerContainer)\n : 0); // For not aligned case\n OType* output_g =\n &output_t[(r_g * num_rows + c_g) / kNFP4PerContainer]; // Output address in global memory\n // Each kNumThreadsStore threads form a warp process one row, we need to find the lane id of\n // the first thread to do the reduction.\n const unsigned src_lane =\n (threadIdx.x % kThreadsPerWarp) / kNumThreadsReduce * kNumThreadsReduce;\n // This mask represents which threads should do the reduction together.\n const unsigned mask = ((1 << kNumThreadsReduce) - 1) << src_lane;\n const bool is_src_lane = (threadIdx.x % kNumThreadsReduce) == 0;\n#pragma unroll\n for (int iter = 0; iter < num_iterations; ++iter) {\n SMemVec smem_vec[kNVecOut];\n // Step 3.1: Load from shared memory to registers\n#pragma unroll\n for (int i = 0; i < kNVecOut; ++i) {\n int r = r_s + i;\n int c = c_s;\n smem_vec[i] = smem[r * kSMemCol + c];\n }\n#pragma unroll\n for (int smem_idx = 0; smem_idx < kNVecSMem; ++smem_idx) {\n // Step 3.2: Compute local amax\n CType amax = 0;\n if constexpr (kIs2DBlockScaling) {\n // TODO(zhongbo): 2D block scaling, directly read from amax_smem\n int warp_idx = threadIdx.x / kThreadsPerWarp; // 0 ~ 7\n constexpr int kNumColsPerWarp =\n kThreadsPerWarp / kNumThreadsStore * kNVecSMem; // 8 elements\n constexpr int kNumWarpsPerBlock =\n kThreadsPerBlock / kThreadsPerWarp; // 8 warps per block\n constexpr int kNumColsPerIter = kNumColsPerWarp * kNumWarpsPerBlock;\n int tid_in_warp_x = (threadIdx.x / kNumThreadsStore) % kNumColsPerWarp;\n int tid_in_warp_y = (threadIdx.x % kThreadsPerWarp) % kNumThreadsStore;\n int data_col_idx = iter * kNumColsPerIter + warp_idx * kNumColsPerWarp + tid_in_warp_x;\n amax = amax_smem[tid_in_warp_y][data_col_idx / kFP4BlockScalingSize];\n } else {\n#pragma unroll\n for (int i = 0; i < kNVecOut; ++i) {\n amax = fmaxf(amax, fabsf(smem_vec[i].data.elt[smem_idx]));\n }\n }\n // Step 3.3: Reduce amax\n if constexpr (kIsE8Scaling) {\n#pragma unroll\n for (int delta = kNumThreadsReduce / 2; delta > 0; delta /= 2) {\n const float other_amax = __shfl_down_sync(mask, amax, delta);\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n amax = __shfl_sync(mask, amax, src_lane);\n }\n // Step 3.4: Compute scale\n ScaleType scale_inv =\n ComputeDecodeScaleFP4(amax, global_encode_scale_multiplier);\n float encode_scale = ComputeEncodeScaleFP4(scale_inv, global_decode_scale);\n // Step 3.5: Write scale_inv_t\n bool write_scale_inv = is_src_lane;\n if constexpr (!kAligned) {\n write_scale_inv &= (r_g + smem_idx < row_length);\n write_scale_inv &= (c_g < num_rows);\n }\n if (write_scale_inv) {\n size_t row_idx = block_idx_x * kTileDim + c_s * kNVecSMem + smem_idx;\n size_t col_idx = (block_idx_y * (kNumThreadsStore / kNumThreadsReduce) +\n (threadIdx.x % kNumThreadsStore) / kNumThreadsReduce);\n if constexpr (kSwizzledScale) {\n size_t offset = scale_factor_swizzled_offset(\n row_idx, col_idx, DIVUP(num_rows, kScaleBlockDim));\n tile_scales_inv_t[offset] = scale_inv;\n } else {\n tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv;\n }\n }\n // Step 3.6: Quantize\n OVec output_vec;\n#pragma unroll\n for (int i = 0; i < kNVecOut / kNFP4PerContainer; i += 2) {\n // Pack two elements into __nv_bfloat162\n float2 f2_a;\n float2 f2_b;\n f2_a.x =\n ComputeOutputFP4(smem_vec[2 * i].data.elt[smem_idx], encode_scale);\n f2_a.y = ComputeOutputFP4(smem_vec[2 * i + 1].data.elt[smem_idx],\n encode_scale);\n f2_b.x = ComputeOutputFP4(smem_vec[2 * (i + 1)].data.elt[smem_idx],\n encode_scale);\n f2_b.y = ComputeOutputFP4(smem_vec[2 * (i + 1) + 1].data.elt[smem_idx],\n encode_scale);\n const uint32_t rbits =\n kApplyStochasticRounding ? get_rbits(rng, random_uint4, rnd_idx) : 0;\n // Convert to __nv_fp4x4_e2m1\n __nv_fp4x4_e2m1 out_4x = cvt_fp32_to_fp4_4x(f2_a, f2_b, rbits);\n\n output_vec.data.elt[i] = reinterpret_cast<__nv_fp4x2_storage_t*>(&out_4x)[0];\n output_vec.data.elt[i + 1] = reinterpret_cast<__nv_fp4x2_storage_t*>(&out_4x)[1];\n }\n // Step 3.7: Store output_t\n if constexpr (kAligned) {\n output_vec.store_to(output_g + smem_idx * num_rows / kNFP4PerContainer);\n } else {\n if (r_g + smem_idx < row_length) {\n output_vec.store_to_elts(output_g + smem_idx * num_rows / kNFP4PerContainer, 0,\n num_ele);\n }\n }\n }\n // Step 3.8: Update output address, column index of shared memory (and row index of global\n // memory for not aligned case)\n output_g += stride_g / kNFP4PerContainer;\n c_s += c_stride;\n if constexpr (!kAligned) {\n r_g += c_stride * kNVecSMem;\n }\n }\n }\n}\n\n} // namespace\n} // namespace quantize_transpose_nvfp4\n#endif // CUDA_VERSION >= 12080\n\nnamespace detail {\n\nvoid quantize_transpose_vector_blockwise_fp4(\n const SimpleTensor& input, const SimpleTensor& global_amax, SimpleTensor& scale_inv,\n SimpleTensor& scale_inv_t, SimpleTensor& output, SimpleTensor& output_t, const float epsilon,\n const bool return_identity, const bool return_transpose, const bool pow2_scale,\n const bool swizzled_scale, const bool use_stochastic_rounding,\n const NVTETensor rng_state_tensor, const bool use_2d_quantization, const bool row_scaled_nvfp4,\n const SimpleTensor& noop_tensor, cudaStream_t stream) {\n NVTE_API_CALL(quantize_transpose_vector_blockwise_fp4);\n#if CUDA_VERSION >= 12080\n\n // pow 2 scale is for MXFP4 since it's using E8M0 scaling\n // raise error if pow2_scale is true\n NVTE_CHECK(!pow2_scale, \"No support for pow2_scale for MXFP4 for now\");\n\n NVTE_CHECK(return_identity || return_transpose,\n \"At least one of return_identity or return_transpose must be true.\");\n\n NVTE_CHECK(return_identity || !use_2d_quantization,\n \"2D block quantization is only supported when return_identity is true.\");\n NVTE_CHECK(!row_scaled_nvfp4 || (return_identity && !return_transpose),\n \"Row-scaled NVFP4 quantization only supports rowwise quantization.\");\n NVTE_CHECK(!row_scaled_nvfp4 || !use_2d_quantization,\n \"Row-scaled NVFP4 quantization does not support 2D quantization.\");\n\n const size_t row_length = input.shape.size() > 0 ? input.shape.at(input.shape.size() - 1) : 1u;\n size_t num_elements = row_length;\n size_t num_rows = 1;\n for (size_t i = 0; (i < input.shape.size() - 1) && (input.shape.size() > 0); ++i) {\n num_rows *= input.shape.at(i);\n num_elements *= input.shape.at(i);\n }\n\n // Early return if the input tensor is empty\n if (num_elements == 0) {\n return;\n }\n\n size_t scale_stride_x = 0;\n size_t scale_stride_y = 0;\n\n if (return_identity) {\n scale_stride_x = 1;\n scale_stride_y = scale_inv.shape[1];\n }\n\n size_t scale_t_stride_x = 0;\n size_t scale_t_stride_y = 0;\n\n if (return_transpose) {\n scale_t_stride_x = 1;\n scale_t_stride_y = scale_inv_t.shape[1];\n }\n\n using namespace transformer_engine::quantize_transpose_nvfp4;\n\n const size_t num_blocks_x = DIVUP(row_length, static_cast(kTileDim));\n const size_t num_blocks_y = DIVUP(num_rows, static_cast(kTileDim));\n\n // noop tensor for cuda graph\n const float* noop_ptr = reinterpret_cast(noop_tensor.dptr);\n\n const size_t* rng_state = nullptr;\n if (rng_state_tensor != nullptr) {\n Tensor& rng_state_te_tensor = *convertNVTETensor(rng_state_tensor);\n NVTE_CHECK(rng_state_te_tensor.dtype() == DType::kInt64,\n \"RNG state should contain 2 64-bit values.\");\n NVTE_CHECK(rng_state_te_tensor.data.shape == std::vector{2},\n \"Shape of the RNG state should be [2], but got \", rng_state_te_tensor.data.shape);\n rng_state = reinterpret_cast(rng_state_te_tensor.data.dptr);\n }\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.dtype, InputType,\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_FP4x2_ONLY(\n return_identity ? output.dtype : output_t.dtype, 2, OutputType,\n\n dim3 grid(num_blocks_x, num_blocks_y, 1);\n\n using ScaleType = fp8e4m3; constexpr int kScaleBlockDim = 16;\n constexpr bool kPow2Scale = false;\n\n const bool full_tile = row_length % kTileDim == 0 && num_rows % kTileDim == 0;\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_identity, kReturnIdentity,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_transpose, kReturnTranspose,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n full_tile, kAligned,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n swizzled_scale, kSwizzledScale,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_stochastic_rounding, kApplyStochasticRounding,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_2d_quantization, kIs2DBlockScaling,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n row_scaled_nvfp4, kRowScaledNVFP4,\n\n size_t smem_bytes = kSMemSize * sizeof(InputType);\n auto kernel = block_scaled_1d_cast_transpose_kernel<\n kReturnIdentity, kReturnTranspose, kPow2Scale, kAligned,\n float, InputType, OutputType, ScaleType, kSwizzledScale,\n kApplyStochasticRounding, kIs2DBlockScaling,\n kRowScaledNVFP4>;\n if (smem_bytes >= 48 * 1024) {\n cudaError_t err = cudaFuncSetAttribute(\n kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,\n smem_bytes);\n NVTE_CHECK(err == cudaSuccess,\n \"Failed to set dynamic shared memory size.\");\n } kernel<<>>(\n reinterpret_cast(input.dptr),\n reinterpret_cast(global_amax.dptr),\n reinterpret_cast(output.dptr),\n reinterpret_cast(output_t.dptr),\n reinterpret_cast(scale_inv.dptr),\n reinterpret_cast(scale_inv_t.dptr),\n row_length, num_rows, scale_stride_x, scale_stride_y,\n scale_t_stride_x, scale_t_stride_y, kScaleBlockDim,\n epsilon, rng_state,\n noop_ptr);) // kRowScaledNVFP4\n ) // kIs2DBlockScaling\n ) // kApplyStochasticRounding\n ) // kSwizzledScale\n ) // kAligned\n ) // kReturnTranspose\n ) // kReturnIdentity\n ) // OutputType\n ) // InputType\n\n NVTE_CHECK_CUDA(cudaGetLastError());\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // CUDA_VERSION >= 12080\n}\n\n} // namespace detail\n} // namespace transformer_engine\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/recipe/recipe_common.cuh\"\n#include \"common/transpose/cast_transpose.h\"\n#include \"common/util/curanddx.hpp\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n\nnamespace transformer_engine {\n\n#if CUDA_VERSION >= 12080\nnamespace quantize_transpose_nvfp4 {\nnamespace {\n\nusing std::int32_t;\nusing std::uint32_t;\nusing std::uint8_t;\n\nusing transformer_engine::detail::TypeExtrema;\n\n// clang-format off\n/*\n\nStep 1: Load input to shared memory\n* shard memory: 128x128 elements with type=InputType (below graph doesn't consider padding)\n* 8 warps\n* Loop 8 times\n* What each thread does in each loop:\n * 8 elements are read from the input at a time\n * 2 elements are written to the shared memory at a time, for a total of 4 times\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| T0 | T1 | T2 | T3 | T4 | T5 | T6 | T7 | T8 | T9 | T10 | T11 | T12 | T13 | T14 | T15 |\n| T16 | T17 | T18 | T19 | T20 | T21 | T22 | T23 | T24 | T25 | T26 | T27 | T28 | T29 | T30 | T31 |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| Warp 1 |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| Warp 7 |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n| ... |\n| Loop 8 times |\n| ... |\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n\nStep 2: Cast and store to output_c\n* shard memory: 128x128 elements with type=InputType (below graph doesn't consider padding)\n* 8 warps\n* Loop 4 times\n* What each thread does in each loop:\n * 2 elements are read from the shared memory at a time, for a total of 8 times\n * Every 8 consecutive threads do reduction and calculate the amax of each row\n * 16 elements are quantized and write to output_c at a time\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| T0 | T1 | T2 | T3 | T4 | T5 | T6 | T7 |\n| T8 | T9 | T10 | T11 | T12 | T13 | T14 | T15 |\n| T16 | T17 | T18 | T19 | T20 | T21 | T22 | T23 |\n| T24 | T25 | T26 | T27 | T28 | T29 | T30 | T31 |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| |\n| Warp 1 |\n| |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| |\n| Warp 7 |\n| |\n| |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n| ... |\n| ... |\n| ... |\n| ... |\n| Loop 4 times |\n| ... |\n| ... |\n| ... |\n| ... |\n+-------------------------------+-------------------------------+-------------------------------+-------------------------------+\n\nStep 3: Transpose, cast and store to output_t\n* shard memory: 128x128 elements with type=InputType (below graph doesn't consider padding)\n* 8 warps\n* Loop 2 times\n* What each thread does in each loop:\n * 2 elements (in a row) are read from the shared memory at a time, for a total of 16 times\n * Every 8 consecutive threads do reduction and calculate the amax of each column\n * 16 elements are quantized and write to output_c at a time, for a total of 2 times\n+------8 elements-------+------8 elements-------+-----40 elements-------+------8 elements-------+------8 elements-------+------8 elements-------+-----40 elements-------+------8 elements-------+\n| T0 | T8 | T16 | T24 | | | | T0 | T8 | T16 | T24 | | | |\n| T1 | T9 | T17 | T25 | | | | T1 | T9 | T17 | T25 | | | |\n| T2 | T10 | T18 | T26 | | | | T2 | T10 | T18 | T26 | | | |\n| T3 | T11 | T19 | T27 | Warp 1 | ... | Warp 7 | T3 | T11 | T19 | T27 | Warp 1 | ... | Warp 7 |\n| T4 | T12 | T20 | T28 | | | | T4 | T12 | T20 | T28 | | | |\n| T5 | T13 | T21 | T29 | | | | T5 | T13 | T21 | T29 | | | |\n| T6 | T14 | T22 | T30 | | | | T6 | T14 | T22 | T30 | | | |\n| T7 | T15 | T23 | T31 | | | | T7 | T15 | T23 | T31 | | | |\n+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+\n\n*/\n// clang-format on\n\nconstexpr int kThreadsPerWarp = 32;\n\n// for fp4, we use uint8_t to store 2 fp4 numbers\nconstexpr int kNFP4PerContainer = 2;\n\n// Hyperparameters for performance tuning\nconstexpr int kTileDim = 128;\n// constexpr int kScaleDim = 32;\nconstexpr int kNVecIn = 8; // The number of elements each LDG touches\nconstexpr int kNVecOut = 16; // The number of elements each STG touches\nconstexpr int kNVecSMem = 2; // The number of elements each LDS/STS touches\nconstexpr int kThreadsPerBlock = 256; // Thread block size, 8 warps in total\n\n// Auto-calculated constants, do not modify directly)\nstatic_assert(kNVecIn % kNVecSMem == 0, \"kNVecIn must be divisible by kNVecSMem\");\nstatic_assert(kNVecOut % kNVecSMem == 0, \"kNVecOut must be divisible by kNVecSMem\");\nconstexpr int kSMemRow = kTileDim;\nconstexpr int kSMemCol = (kTileDim / kNVecSMem) + 1;\nconstexpr int kSMemSize = kSMemRow * kSMemCol * kNVecSMem;\nconstexpr int kNumThreadsLoad = kTileDim / kNVecIn; // 16\nconstexpr int kNumThreadsStore = kTileDim / kNVecOut; // 8\n// constexpr int kNumThreadsReduce = kScaleDim / kNVecOut;\nstatic_assert(kNumThreadsLoad <= kThreadsPerWarp, \"kNumThreadsLoad must be <= kThreadsPerWarp\");\nstatic_assert(kNumThreadsStore <= kThreadsPerWarp, \"kNumThreadsStore must be <= kThreadsPerWarp\");\n\n// for 2D block scaling, we need to reduce amax in warp\nstatic __device__ constexpr unsigned int WARP_REDUCE_AMAX_GROUP_MASKS[8] = {\n 0x01010101, 0x02020202, 0x04040404, 0x08080808, 0x10101010, 0x20202020, 0x40404040, 0x80808080};\n\n// max for every group_size elements in warp\ntemplate \n__device__ __forceinline__ float groupMax(float val, unsigned int groupMask) {\n for (int offset = group_size / 2; offset > 0; offset /= 2) {\n val = max(val, __shfl_down_sync(groupMask, val, offset * shfl_down_stride));\n }\n return val;\n}\n\ntemplate \n__device__ __forceinline__ ScaleType\nComputeDecodeScaleFP4(const float amax, const float global_encode_scale_multiplier) {\n float decode_scale = amax * global_encode_scale_multiplier;\n decode_scale = fminf(decode_scale, TypeExtrema::max);\n return static_cast(decode_scale);\n}\n\ntemplate \n__device__ __forceinline__ float ComputeEncodeScaleFP4(ScaleType decode_scale,\n const float global_decode_scale) {\n return fminf(1.0f / (static_cast(decode_scale) * global_decode_scale),\n TypeExtrema::max);\n}\n\ntemplate \n__device__ __forceinline__ float ComputeOutputFP4(IType input, float encode_scale) {\n return static_cast(input) * encode_scale;\n}\n\n__device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax) {\n constexpr float fp8_max = TypeExtrema::max;\n constexpr float fp4_max = TypeExtrema::max;\n float global_encode_scale = fp8_max * fp4_max / global_amax;\n // If scale is infinity, return max value of float32\n global_encode_scale = fminf(global_encode_scale, TypeExtrema::max);\n // If global amax is 0 or infinity, return 1\n if (global_amax == 0.f || global_encode_scale == 0.f) {\n return 1.f;\n }\n return global_encode_scale;\n}\n\n__device__ __forceinline__ uint32_t get_rbits(\n transformer_engine::curanddx::detail::philox4x32_native_state&\n rng, // NVTE_BUILD_NUM_PHILOX_ROUNDS rounds of philox4x32\n uint4& random_uint4, int& rnd_idx) {\n if (rnd_idx == 4) {\n rnd_idx = 0;\n random_uint4 = rng.generate4();\n }\n\n // Treat uint4 as an array of 4x uint32_t elements for indexing\n const uint32_t* const rbits_arr = reinterpret_cast(&random_uint4);\n const uint32_t rbits = rbits_arr[rnd_idx++];\n return rbits;\n}\n\ntemplate \n__device__ __forceinline__ size_t scale_factor_swizzled_offset(size_t row_idx, size_t col_idx,\n uint32_t col_length) {\n // This function takes in indices from the scale factor matrix and returns an offset in the\n // swizzled format. row_idx, col_idx are original indices from the scale factor matrix (unswizzled\n // index). col_length is the column length of the scale factor matrix. tile_scales_inv is the\n // pointer to the scale factor matrix.\n\n // https://github.com/NVIDIA/cutlass/blob/main/media/docs/cpp/blackwell_functionality.md#scale-factor-layouts\n // For any scale factor matrix, it's 512B base block. Each base block consists of 128 rows and 4\n // columns. Base block is divided into 4 column blocks, each column block has 32 rows and 4\n // columns.\n\n // NOTE: There are not a lot of good illustrations about the swizzled scale factor matrix.\n // To think in high level, the swizzled scale factor matrix could be composed as:\n // unswizzled_scale_factor_matrix = torch.empty((M, N // 16), dtype=torch.uint8)\n // cbg_cnt = N // 16 // 4 # Assuming N is divisible by 64\n // rb_cnt = M // 128 # Assuming M is divisible by 128\n // tmp = unswizzled_scale_factor_matrix.reshape(rb_cnt, 4, 32, cbg_cnt, 4)\n // tmp = torch.permute(tmp, (0, 3, 2, 1, 4))\n // swizzled_scale_factor_matrix = tmp.reshape((-1, 128, 4))\n\n constexpr uint32_t kTotalRowsPerBaseBlock = 128;\n constexpr uint32_t kRowsPerBaseBlockCol = 32;\n constexpr uint32_t kColsPerBaseBlockCol = 4;\n\n const size_t rb = row_idx / kTotalRowsPerBaseBlock;\n const size_t rem = row_idx % kTotalRowsPerBaseBlock;\n const size_t d4 = rem / kRowsPerBaseBlockCol;\n const size_t d3 = rem % kRowsPerBaseBlockCol;\n const size_t cbg = col_idx / kColsPerBaseBlockCol;\n const size_t d5 = col_idx % kColsPerBaseBlockCol;\n\n const size_t cbg_cnt = DIVUP(col_length, kColsPerBaseBlockCol);\n // row-major offset in the logical shape\n // (rb_cnt , cbg_cnt , 32 , 4 , 4)\n // Magic number 16 below comes from the fact we have kColsPerBaseBlockCol = 4, and d4 ([0-128] /\n // 32 = [0-4])\n return ((rb * cbg_cnt + cbg) * kRowsPerBaseBlockCol + d3) * 16 + d4 * kColsPerBaseBlockCol + d5;\n}\n\n__device__ __forceinline__ __hip_fp4x4_e2m1 cvt_fp32_to_fp4_4x_with_stochastic_rounding(\n const float2 in01, const float2 in23, const uint32_t rbits) {\n constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING;\n if constexpr (has_rs) {\n uint16_t out_4x;\n asm volatile(\n \"{\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %0, {%3, %4, %1, %2}, %5; \\n\\t\"\n \"}\"\n : \"=h\"(out_4x)\n : \"f\"(in01.y), \"f\"(in01.x), \"f\"(in23.y), \"f\"(in23.x), \"r\"(rbits));\n return *reinterpret_cast<__hip_fp4x4_e2m1*>(&out_4x);\n } else {\n NVTE_DEVICE_ERROR(\n \"FP4 cvt.rs PTX instructions are architecture-specific. \"\n \"Try recompiling with sm_XXXa instead of sm_XXX.\");\n uint16_t dummy = 0;\n return *reinterpret_cast<__hip_fp4x4_e2m1*>(&dummy);\n }\n}\n\n__device__ __forceinline__ __hip_fp4x4_e2m1 cvt_fp32_to_fp4_4x_with_rn(const float2 in01,\n const float2 in23,\n const uint32_t rbits) {\n constexpr bool has_fp4 = ARCH_BLACKWELL_FAMILY;\n if constexpr (has_fp4) {\n // NOTE: rbits unused for rn.\n uint32_t out_4x; // Only need 16 bit. Using 32 bit container for packing.\n asm volatile(\n \"{\\n\"\n \".reg.b8 f0; \\n\\t\"\n \".reg.b8 f1; \\n\\t\"\n \"cvt.rn.satfinite.e2m1x2.f32 f0, %1, %2;\\n\\t\"\n \"cvt.rn.satfinite.e2m1x2.f32 f1, %3, %4;\\n\\t\"\n \"mov.b32 %0, {f0, f1, f0, f1};\\n\\t\"\n \"}\"\n : \"=r\"(out_4x)\n : \"f\"(in01.y), \"f\"(in01.x), \"f\"(in23.y), \"f\"(in23.x));\n return reinterpret_cast<__hip_fp4x4_e2m1*>(&out_4x)[0];\n } else {\n NVTE_DEVICE_ERROR(\n \"FP4 cvt PTX instructions are architecture-specific. \"\n \"Try recompiling with sm_XXXa instead of sm_XXX.\");\n uint16_t dummy = 0;\n return *reinterpret_cast<__hip_fp4x4_e2m1*>(&dummy);\n }\n}\n\ntemplate \n__device__ __forceinline__ __hip_fp4x4_e2m1 cvt_fp32_to_fp4_4x(const float2 in01, const float2 in23,\n const uint32_t rbits) {\n if constexpr (kApplyStochasticRounding) {\n return cvt_fp32_to_fp4_4x_with_stochastic_rounding(in01, in23, rbits);\n } else {\n return cvt_fp32_to_fp4_4x_with_rn(in01, in23, rbits);\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpose_kernel(\n const IType* const input, const float* global_amax, OType* const output_c,\n OType* const output_t, ScaleType* const tile_scales_inv_c, ScaleType* const tile_scales_inv_t,\n const size_t row_length, const size_t num_rows, const size_t scale_stride_x,\n const size_t scale_stride_y, const size_t scale_t_stride_x, const size_t scale_t_stride_y,\n const size_t kScaleBlockDim, const float epsilon, const size_t* rng_state,\n const float* noop_ptr) {\n constexpr int kNVecContainer = kNVecOut / kNFP4PerContainer;\n using SMemVec = Vec;\n using OVec = Vec;\n union IVec {\n Vec input_type;\n Vec smem_type;\n };\n\n if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {\n return;\n }\n\n const size_t block_idx_x = blockIdx.x;\n const size_t block_idx_y = blockIdx.y;\n const size_t rng_sequence =\n threadIdx.x + block_idx_x * kThreadsPerBlock + block_idx_y * gridDim.x * kThreadsPerBlock;\n const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0;\n const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0;\n\n transformer_engine::curanddx::detail::philox4x32_native_state rng;\n rng.init(rng_seed, rng_sequence, rng_offset);\n uint4 random_uint4 = kApplyStochasticRounding ? rng.generate4() : uint4{0, 0, 0, 0};\n\n int rnd_idx =\n 0; // Index of the random number. It increments each time when used and resets to 0 if reaches 4x\n\n extern __shared__ char smem_base[];\n SMemVec* smem = reinterpret_cast(&smem_base[0]);\n\n // 2D block scaling is not supported for E8 scaling MXFP4 or for colwise only mode.\n // Instead of static_assert, return early if these invalid modes are detected.\n if constexpr (kIs2DBlockScaling && kIsE8Scaling) {\n return;\n }\n if constexpr (kIs2DBlockScaling && !kReturnIdentity) {\n return;\n }\n // for 128x128 block, 2D block scaling means there will be 8x8 amax values for nvfp4, 4x4 for 2D mxfp4\n // use constexpr to define the size, when not using 2D, use minimal size 1x1\n constexpr int kFP4BlockScalingSize = 16;\n constexpr int k2DBlockAmaxDim = kIs2DBlockScaling ? (kTileDim / kFP4BlockScalingSize) : 1;\n constexpr int kNumRowsPerWarp = kThreadsPerWarp / kNumThreadsStore; // 4\n constexpr int k2DBlockAmaxReduceDim =\n kIs2DBlockScaling ? (kFP4BlockScalingSize / kNumRowsPerWarp) : 1;\n __shared__ CType amax_smem_red[k2DBlockAmaxDim][k2DBlockAmaxDim][k2DBlockAmaxReduceDim];\n __shared__ CType amax_smem[k2DBlockAmaxDim][k2DBlockAmaxDim];\n\n // Step 1: Load input to shared memory\n {\n constexpr int r_stride = kThreadsPerBlock / kNumThreadsLoad; // stride in rows of shared memory\n constexpr int num_iterations = kTileDim / r_stride;\n const int c_s =\n (threadIdx.x % kNumThreadsLoad) * (kNVecIn / kNVecSMem); // Column in shared memory\n int r_s = threadIdx.x / kNumThreadsLoad; // Row in shared memory\n const size_t c_g = block_idx_x * kTileDim + c_s * kNVecSMem; // Column in global memory\n size_t r_g = block_idx_y * kTileDim + r_s; // Row in global memory\n const size_t stride_g = static_cast(r_stride) * row_length; // Stride in global memory\n const size_t num_ele = (c_g < row_length ? min(static_cast(kNVecIn), row_length - c_g)\n : 0); // For not aligned case\n const IType* input_g = &input[r_g * row_length + c_g]; // Input address in global memory\n#pragma unroll\n for (int iter = 0; iter < num_iterations; ++iter) {\n IVec input_vec;\n // Step 1.1: Load from global memory (input) to registers\n if constexpr (kAligned) {\n input_vec.input_type.load_from(input_g);\n } else {\n if (r_g < num_rows) {\n input_vec.input_type.load_from_elts(input_g, 0, num_ele);\n } else {\n input_vec.input_type.clear();\n }\n }\n // Step 1.2: Write to shared memory\n#pragma unroll\n for (int i = 0; i < kNVecIn / kNVecSMem; ++i) {\n int c = c_s + i;\n int r = r_s;\n smem[r * kSMemCol + c] = input_vec.smem_type.data.elt[i];\n }\n // Step 1.3: Update input address, row index of shared memory, (and row index of global memory\n // for not aligned case)\n input_g += stride_g;\n r_s += r_stride;\n if constexpr (!kAligned) {\n r_g += r_stride;\n }\n }\n }\n\n __syncthreads();\n\n const int kNumThreadsReduce = kScaleBlockDim / kNVecOut;\n const float global_encode_scale =\n kIsE8Scaling ? 1.0f : ComputeGlobalEncodeScaleFP4(global_amax[0]);\n constexpr float fp4_max_inv = 1.0f / TypeExtrema::max;\n const float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n const float global_decode_scale = 1.0 / global_encode_scale;\n\n // Step 2: Cast and store to output_c\n if constexpr (kReturnIdentity) {\n constexpr int r_stride =\n kThreadsPerBlock / kNumThreadsStore; // stride in rows of shared memory\n constexpr int num_iterations = kTileDim / r_stride;\n const int c_s =\n (threadIdx.x % kNumThreadsStore) * (kNVecOut / kNVecSMem); // Column in shared memory\n int r_s = threadIdx.x / kNumThreadsStore; // Row in shared memory\n const size_t c_g = block_idx_x * kTileDim + c_s * kNVecSMem; // Column in global memory\n size_t r_g = block_idx_y * kTileDim + r_s; // Row in global memory\n const size_t stride_g = static_cast(r_stride) * row_length; // Stride in global memory\n const size_t num_ele =\n (c_g < row_length ? min(static_cast(kNVecOut / kNFP4PerContainer),\n (row_length - c_g) / kNFP4PerContainer)\n : 0); // For not aligned case\n OType* output_g =\n &output_c[(r_g * row_length + c_g) / kNFP4PerContainer]; // Output address in global memory\n // Each kNumThreadsStore threads form a warp process one row, we need to find the lane id of\n // the first thread to do the reduction.\n const unsigned src_lane =\n (threadIdx.x % kThreadsPerWarp) / kNumThreadsReduce * kNumThreadsReduce;\n // This mask represents which threads should do the reduction together.\n const unsigned mask = ((1 << kNumThreadsReduce) - 1) << src_lane;\n const bool is_src_lane = (threadIdx.x % kNumThreadsReduce) == 0;\n#pragma unroll\n for (int iter = 0; iter < num_iterations; ++iter) {\n SMemVec smem_vec[kNVecOut / kNVecSMem];\n // Step 2.1: Load from shared memory to registers\n#pragma unroll\n for (int i = 0; i < kNVecOut / kNVecSMem; ++i) {\n int c = c_s + i;\n int r = r_s;\n smem_vec[i] = smem[r * kSMemCol + c];\n }\n // Step 2.2: Compute local amax\n CType amax = 0;\n#pragma unroll\n for (int i = 0; i < kNVecOut / kNVecSMem; ++i) {\n#pragma unroll\n for (int j = 0; j < kNVecSMem; ++j) {\n __builtin_assume(amax >= 0);\n amax = fmaxf(amax, fabsf(smem_vec[i].data.elt[j]));\n }\n }\n // Step 2.3: Reduce amax\n if constexpr (kIsE8Scaling) {\n#pragma unroll\n for (int delta = kNumThreadsReduce / 2; delta > 0; delta /= 2) {\n const float other_amax = __shfl_down_sync(mask, amax, delta);\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n amax = __shfl_sync(mask, amax, src_lane);\n }\n // doing shuffle sync for 2D block scaling (not applicable for E8 scaling)\n if constexpr (kIs2DBlockScaling) {\n // first amax shuffle sync in warp, then reduce in smem\n // T0 T8 T16 T24 should do amax reduction together\n constexpr int kNumRowsPerIter = kThreadsPerBlock / kNumThreadsStore; // 32\n int warp_idx = threadIdx.x / kThreadsPerWarp; // 0 ~ 7\n int tid_in_warp_x = threadIdx.x % kNumThreadsStore;\n int tid_in_warp_y = (threadIdx.x / kNumThreadsStore) % kNumRowsPerWarp;\n CType amax_warp_reduced = groupMax(\n amax, WARP_REDUCE_AMAX_GROUP_MASKS[tid_in_warp_x]);\n // now T0 ~ T8 in each warp has the reduced amax values\n int data_row_idx = iter * kNumRowsPerIter + warp_idx * kNumRowsPerWarp + tid_in_warp_y;\n if (tid_in_warp_y == 0) {\n amax_smem_red[data_row_idx / kFP4BlockScalingSize][tid_in_warp_x]\n [warp_idx % k2DBlockAmaxReduceDim] = amax_warp_reduced;\n }\n __syncthreads();\n\n if (data_row_idx % kFP4BlockScalingSize == 0) {\n CType amax_2d = 0.0;\n for (int i = 0; i < k2DBlockAmaxReduceDim; i++) {\n amax_2d = fmaxf(amax_2d,\n amax_smem_red[data_row_idx / kFP4BlockScalingSize][tid_in_warp_x][i]);\n }\n amax_smem[data_row_idx / kFP4BlockScalingSize][tid_in_warp_x] = amax_2d;\n }\n __syncthreads();\n // every thread now knows 2D amax\n amax = amax_smem[data_row_idx / kFP4BlockScalingSize][tid_in_warp_x];\n }\n // Step 2.4: Compute scale\n const size_t row_idx = block_idx_y * kTileDim + r_s;\n float row_global_encode_scale = global_encode_scale;\n if constexpr (kRowScaledNVFP4) {\n row_global_encode_scale =\n row_idx < num_rows ? ComputeGlobalEncodeScaleFP4(global_amax[row_idx]) : 1.0f;\n }\n const float row_global_encode_scale_multiplier =\n kRowScaledNVFP4 ? row_global_encode_scale * fp4_max_inv : global_encode_scale_multiplier;\n const float row_global_decode_scale =\n kRowScaledNVFP4 ? 1.0f / row_global_encode_scale : global_decode_scale;\n ScaleType scale_inv =\n ComputeDecodeScaleFP4(amax, row_global_encode_scale_multiplier);\n float encode_scale = ComputeEncodeScaleFP4(scale_inv, row_global_decode_scale);\n // Step 2.5: Write scale_inv\n bool write_scale_inv = is_src_lane;\n if constexpr (!kAligned) {\n write_scale_inv &= (r_g < num_rows);\n write_scale_inv &= (c_g < row_length);\n }\n if (write_scale_inv) {\n size_t row_idx = block_idx_y * kTileDim + r_s;\n size_t col_idx = block_idx_x * (kNumThreadsStore / kNumThreadsReduce) +\n (threadIdx.x % kNumThreadsStore) / kNumThreadsReduce;\n if constexpr (kSwizzledScale) {\n size_t offset = scale_factor_swizzled_offset(\n row_idx, col_idx, DIVUP(row_length, kScaleBlockDim));\n tile_scales_inv_c[offset] = scale_inv;\n } else {\n tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv;\n }\n }\n // Step 2.6: Quantize\n OVec output_vec;\n#pragma unroll\n for (int i = 0; i < kNVecOut / kNVecSMem; i += 2) {\n // Pack two elements into __nv_bfloat162\n float2 f2_a;\n float2 f2_b;\n f2_a.x = ComputeOutputFP4(smem_vec[i].data.elt[0], encode_scale);\n f2_a.y = ComputeOutputFP4(smem_vec[i].data.elt[1], encode_scale);\n f2_b.x = ComputeOutputFP4(smem_vec[i + 1].data.elt[0], encode_scale);\n f2_b.y = ComputeOutputFP4(smem_vec[i + 1].data.elt[1], encode_scale);\n const uint32_t rbits = kApplyStochasticRounding ? get_rbits(rng, random_uint4, rnd_idx) : 0;\n // Convert to __nv_fp4x4_e2m1\n __hip_fp4x4_e2m1 out_4x = cvt_fp32_to_fp4_4x(f2_a, f2_b, rbits);\n\n output_vec.data.elt[i] = reinterpret_cast<__hip_fp4x2_storage_t*>(&out_4x)[0];\n output_vec.data.elt[i + 1] = reinterpret_cast<__hip_fp4x2_storage_t*>(&out_4x)[1];\n }\n // Step 2.7: Store output_c\n if constexpr (kAligned) {\n output_vec.store_to(output_g);\n } else {\n if (r_g < num_rows) {\n output_vec.store_to_elts(output_g, 0, num_ele);\n }\n }\n // Step 2.8: Update output address, row index of shared memory (and row index of global memory\n // for not aligned case)\n output_g += stride_g / kNFP4PerContainer;\n r_s += r_stride;\n if constexpr (!kAligned) {\n r_g += r_stride;\n }\n }\n }\n\n // Step 3: Transpose, cast and store to output_t\n if constexpr (kReturnTranspose) {\n constexpr int c_stride =\n kThreadsPerBlock / kNumThreadsStore; // Stride in columns of shared memory\n constexpr int num_iterations = kTileDim / (c_stride * kNVecSMem);\n const int r_s = (threadIdx.x % kNumThreadsStore) * kNVecOut; // Row in shared memory\n int c_s = threadIdx.x / kNumThreadsStore; // Column in shared memory\n size_t r_g = block_idx_x * kTileDim + c_s * kNVecSMem; // Row in global memory\n const size_t c_g = block_idx_y * kTileDim + r_s; // Column in global memory\n const size_t stride_g =\n static_cast(c_stride) * kNVecSMem * num_rows; // Stride in global memory\n const size_t num_ele = (c_g < num_rows ? min(static_cast(kNVecOut / kNFP4PerContainer),\n (num_rows - c_g) / kNFP4PerContainer)\n : 0); // For not aligned case\n OType* output_g =\n &output_t[(r_g * num_rows + c_g) / kNFP4PerContainer]; // Output address in global memory\n // Each kNumThreadsStore threads form a warp process one row, we need to find the lane id of\n // the first thread to do the reduction.\n const unsigned src_lane =\n (threadIdx.x % kThreadsPerWarp) / kNumThreadsReduce * kNumThreadsReduce;\n // This mask represents which threads should do the reduction together.\n const unsigned mask = ((1 << kNumThreadsReduce) - 1) << src_lane;\n const bool is_src_lane = (threadIdx.x % kNumThreadsReduce) == 0;\n#pragma unroll\n for (int iter = 0; iter < num_iterations; ++iter) {\n SMemVec smem_vec[kNVecOut];\n // Step 3.1: Load from shared memory to registers\n#pragma unroll\n for (int i = 0; i < kNVecOut; ++i) {\n int r = r_s + i;\n int c = c_s;\n smem_vec[i] = smem[r * kSMemCol + c];\n }\n#pragma unroll\n for (int smem_idx = 0; smem_idx < kNVecSMem; ++smem_idx) {\n // Step 3.2: Compute local amax\n CType amax = 0;\n if constexpr (kIs2DBlockScaling) {\n // TODO(zhongbo): 2D block scaling, directly read from amax_smem\n int warp_idx = threadIdx.x / kThreadsPerWarp; // 0 ~ 7\n constexpr int kNumColsPerWarp =\n kThreadsPerWarp / kNumThreadsStore * kNVecSMem; // 8 elements\n constexpr int kNumWarpsPerBlock =\n kThreadsPerBlock / kThreadsPerWarp; // 8 warps per block\n constexpr int kNumColsPerIter = kNumColsPerWarp * kNumWarpsPerBlock;\n int tid_in_warp_x = (threadIdx.x / kNumThreadsStore) % kNumColsPerWarp;\n int tid_in_warp_y = (threadIdx.x % kThreadsPerWarp) % kNumThreadsStore;\n int data_col_idx = iter * kNumColsPerIter + warp_idx * kNumColsPerWarp + tid_in_warp_x;\n amax = amax_smem[tid_in_warp_y][data_col_idx / kFP4BlockScalingSize];\n } else {\n#pragma unroll\n for (int i = 0; i < kNVecOut; ++i) {\n amax = fmaxf(amax, fabsf(smem_vec[i].data.elt[smem_idx]));\n }\n }\n // Step 3.3: Reduce amax\n if constexpr (kIsE8Scaling) {\n#pragma unroll\n for (int delta = kNumThreadsReduce / 2; delta > 0; delta /= 2) {\n const float other_amax = __shfl_down_sync(mask, amax, delta);\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n amax = __shfl_sync(mask, amax, src_lane);\n }\n // Step 3.4: Compute scale\n ScaleType scale_inv =\n ComputeDecodeScaleFP4(amax, global_encode_scale_multiplier);\n float encode_scale = ComputeEncodeScaleFP4(scale_inv, global_decode_scale);\n // Step 3.5: Write scale_inv_t\n bool write_scale_inv = is_src_lane;\n if constexpr (!kAligned) {\n write_scale_inv &= (r_g + smem_idx < row_length);\n write_scale_inv &= (c_g < num_rows);\n }\n if (write_scale_inv) {\n size_t row_idx = block_idx_x * kTileDim + c_s * kNVecSMem + smem_idx;\n size_t col_idx = (block_idx_y * (kNumThreadsStore / kNumThreadsReduce) +\n (threadIdx.x % kNumThreadsStore) / kNumThreadsReduce);\n if constexpr (kSwizzledScale) {\n size_t offset = scale_factor_swizzled_offset(\n row_idx, col_idx, DIVUP(num_rows, kScaleBlockDim));\n tile_scales_inv_t[offset] = scale_inv;\n } else {\n tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv;\n }\n }\n // Step 3.6: Quantize\n OVec output_vec;\n#pragma unroll\n for (int i = 0; i < kNVecOut / kNFP4PerContainer; i += 2) {\n // Pack two elements into __nv_bfloat162\n float2 f2_a;\n float2 f2_b;\n f2_a.x =\n ComputeOutputFP4(smem_vec[2 * i].data.elt[smem_idx], encode_scale);\n f2_a.y = ComputeOutputFP4(smem_vec[2 * i + 1].data.elt[smem_idx],\n encode_scale);\n f2_b.x = ComputeOutputFP4(smem_vec[2 * (i + 1)].data.elt[smem_idx],\n encode_scale);\n f2_b.y = ComputeOutputFP4(smem_vec[2 * (i + 1) + 1].data.elt[smem_idx],\n encode_scale);\n const uint32_t rbits =\n kApplyStochasticRounding ? get_rbits(rng, random_uint4, rnd_idx) : 0;\n // Convert to __nv_fp4x4_e2m1\n __hip_fp4x4_e2m1 out_4x = cvt_fp32_to_fp4_4x(f2_a, f2_b, rbits);\n\n output_vec.data.elt[i] = reinterpret_cast<__hip_fp4x2_storage_t*>(&out_4x)[0];\n output_vec.data.elt[i + 1] = reinterpret_cast<__hip_fp4x2_storage_t*>(&out_4x)[1];\n }\n // Step 3.7: Store output_t\n if constexpr (kAligned) {\n output_vec.store_to(output_g + smem_idx * num_rows / kNFP4PerContainer);\n } else {\n if (r_g + smem_idx < row_length) {\n output_vec.store_to_elts(output_g + smem_idx * num_rows / kNFP4PerContainer, 0,\n num_ele);\n }\n }\n }\n // Step 3.8: Update output address, column index of shared memory (and row index of global\n // memory for not aligned case)\n output_g += stride_g / kNFP4PerContainer;\n c_s += c_stride;\n if constexpr (!kAligned) {\n r_g += c_stride * kNVecSMem;\n }\n }\n }\n}\n\n} // namespace\n} // namespace quantize_transpose_nvfp4\n#endif // CUDA_VERSION >= 12080\n\nnamespace detail {\n\nvoid quantize_transpose_vector_blockwise_fp4(\n const SimpleTensor& input, const SimpleTensor& global_amax, SimpleTensor& scale_inv,\n SimpleTensor& scale_inv_t, SimpleTensor& output, SimpleTensor& output_t, const float epsilon,\n const bool return_identity, const bool return_transpose, const bool pow2_scale,\n const bool swizzled_scale, const bool use_stochastic_rounding,\n const NVTETensor rng_state_tensor, const bool use_2d_quantization, const bool row_scaled_nvfp4,\n const SimpleTensor& noop_tensor, hipStream_t stream) {\n NVTE_API_CALL(quantize_transpose_vector_blockwise_fp4);\n#if CUDA_VERSION >= 12080\n\n // pow 2 scale is for MXFP4 since it's using E8M0 scaling\n // raise error if pow2_scale is true\n NVTE_CHECK(!pow2_scale, \"No support for pow2_scale for MXFP4 for now\");\n\n NVTE_CHECK(return_identity || return_transpose,\n \"At least one of return_identity or return_transpose must be true.\");\n\n NVTE_CHECK(return_identity || !use_2d_quantization,\n \"2D block quantization is only supported when return_identity is true.\");\n NVTE_CHECK(!row_scaled_nvfp4 || (return_identity && !return_transpose),\n \"Row-scaled NVFP4 quantization only supports rowwise quantization.\");\n NVTE_CHECK(!row_scaled_nvfp4 || !use_2d_quantization,\n \"Row-scaled NVFP4 quantization does not support 2D quantization.\");\n\n const size_t row_length = input.shape.size() > 0 ? input.shape.at(input.shape.size() - 1) : 1u;\n size_t num_elements = row_length;\n size_t num_rows = 1;\n for (size_t i = 0; (i < input.shape.size() - 1) && (input.shape.size() > 0); ++i) {\n num_rows *= input.shape.at(i);\n num_elements *= input.shape.at(i);\n }\n\n // Early return if the input tensor is empty\n if (num_elements == 0) {\n return;\n }\n\n size_t scale_stride_x = 0;\n size_t scale_stride_y = 0;\n\n if (return_identity) {\n scale_stride_x = 1;\n scale_stride_y = scale_inv.shape[1];\n }\n\n size_t scale_t_stride_x = 0;\n size_t scale_t_stride_y = 0;\n\n if (return_transpose) {\n scale_t_stride_x = 1;\n scale_t_stride_y = scale_inv_t.shape[1];\n }\n\n using namespace transformer_engine::quantize_transpose_nvfp4;\n\n const size_t num_blocks_x = DIVUP(row_length, static_cast(kTileDim));\n const size_t num_blocks_y = DIVUP(num_rows, static_cast(kTileDim));\n\n // noop tensor for cuda graph\n const float* noop_ptr = reinterpret_cast(noop_tensor.dptr);\n\n const size_t* rng_state = nullptr;\n if (rng_state_tensor != nullptr) {\n Tensor& rng_state_te_tensor = *convertNVTETensor(rng_state_tensor);\n NVTE_CHECK(rng_state_te_tensor.dtype() == DType::kInt64,\n \"RNG state should contain 2 64-bit values.\");\n NVTE_CHECK(rng_state_te_tensor.data.shape == std::vector{2},\n \"Shape of the RNG state should be [2], but got \", rng_state_te_tensor.data.shape);\n rng_state = reinterpret_cast(rng_state_te_tensor.data.dptr);\n }\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.dtype, InputType,\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_FP4x2_ONLY(\n return_identity ? output.dtype : output_t.dtype, 2, OutputType,\n\n dim3 grid(num_blocks_x, num_blocks_y, 1);\n\n using ScaleType = fp8e4m3; constexpr int kScaleBlockDim = 16;\n constexpr bool kPow2Scale = false;\n\n const bool full_tile = row_length % kTileDim == 0 && num_rows % kTileDim == 0;\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_identity, kReturnIdentity,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_transpose, kReturnTranspose,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n full_tile, kAligned,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n swizzled_scale, kSwizzledScale,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_stochastic_rounding, kApplyStochasticRounding,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_2d_quantization, kIs2DBlockScaling,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n row_scaled_nvfp4, kRowScaledNVFP4,\n\n size_t smem_bytes = kSMemSize * sizeof(InputType);\n auto kernel = block_scaled_1d_cast_transpose_kernel<\n kReturnIdentity, kReturnTranspose, kPow2Scale, kAligned,\n float, InputType, OutputType, ScaleType, kSwizzledScale,\n kApplyStochasticRounding, kIs2DBlockScaling,\n kRowScaledNVFP4>;\n if (smem_bytes >= 48 * 1024) {\n hipError_t err = hipFuncSetAttribute(reinterpret_cast(kernel), hipFuncAttributeMaxDynamicSharedMemorySize,\n smem_bytes);\n NVTE_CHECK(err == hipSuccess,\n \"Failed to set dynamic shared memory size.\");\n } kernel<<>>(\n reinterpret_cast(input.dptr),\n reinterpret_cast(global_amax.dptr),\n reinterpret_cast(output.dptr),\n reinterpret_cast(output_t.dptr),\n reinterpret_cast(scale_inv.dptr),\n reinterpret_cast(scale_inv_t.dptr),\n row_length, num_rows, scale_stride_x, scale_stride_y,\n scale_t_stride_x, scale_t_stride_y, kScaleBlockDim,\n epsilon, rng_state,\n noop_ptr);) // kRowScaledNVFP4\n ) // kIs2DBlockScaling\n ) // kApplyStochasticRounding\n ) // kSwizzledScale\n ) // kAligned\n ) // kReturnTranspose\n ) // kReturnIdentity\n ) // OutputType\n ) // InputType\n\n NVTE_CHECK_CUDA(hipGetLastError());\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // CUDA_VERSION >= 12080\n}\n\n} // namespace detail\n} // namespace transformer_engine\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/recipe/recipe_common.cuh\"\n#include \"common/util/cuda_runtime.h\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n\n#if (!defined(__CUDA_MINIMUM_ARCH__) && __CUDA_ARCH__ >= 900) || \\\n (defined(__CUDA_MINIMUM_ARCH__) && __CUDA_MINIMUM_ARCH__ >= 900)\n#define TMA_HW_SUPPORTED\n#endif\n\nnamespace transformer_engine {\nnamespace {\n\n// const values configuration\n\nconstexpr size_t kThreadsPerWarp = 32;\n#ifdef TMA_HW_SUPPORTED\nconstexpr size_t BLOCK_TILE_DIM = 128;\nconstexpr size_t WARP_TILE_DIM_X = 32;\nconstexpr size_t WARP_TILE_DIM_Y = 64;\nconstexpr size_t THREAD_TILE_DIM_X = 16;\nconstexpr size_t THREAD_TILE_DIM_Y = 4;\n#else\nconstexpr size_t BLOCK_TILE_DIM = 128;\nconstexpr size_t WARP_TILE_DIM_X = 64;\nconstexpr size_t WARP_TILE_DIM_Y = 32;\nconstexpr size_t THREAD_TILE_DIM_X = 8;\nconstexpr size_t THREAD_TILE_DIM_Y = 8;\n#endif\n\n#ifdef TMA_HW_SUPPORTED\nconstexpr size_t NUM_BYTES_PER_BANK = 4;\nconstexpr size_t NUM_BANKS_PER_SHARED_ELEM = THREAD_TILE_DIM_Y / NUM_BYTES_PER_BANK;\nconstexpr size_t SHARED_BLOCK_TILE_DIM_Y = BLOCK_TILE_DIM;\nconstexpr size_t SHARED_BLOCK_TILE_DIM_X_BANKS =\n BLOCK_TILE_DIM / (NUM_BYTES_PER_BANK * NUM_BANKS_PER_SHARED_ELEM);\nconstexpr size_t NUM_BANKS_Y_IN_WARP = WARP_TILE_DIM_Y / NUM_BYTES_PER_BANK;\n#endif\nconstexpr size_t ELE_PER_THREAD = THREAD_TILE_DIM_X * THREAD_TILE_DIM_Y;\nconstexpr size_t THREADS_PER_BLOCK = BLOCK_TILE_DIM * BLOCK_TILE_DIM / ELE_PER_THREAD;\nconstexpr size_t NUM_WARPS_X_IN_BLOCK = BLOCK_TILE_DIM / WARP_TILE_DIM_X;\nconstexpr size_t NUM_WARPS_Y_IN_BLOCK = BLOCK_TILE_DIM / WARP_TILE_DIM_Y;\nconstexpr size_t NUM_WARPS_IN_BLOCK = NUM_WARPS_X_IN_BLOCK * NUM_WARPS_Y_IN_BLOCK;\n\nconstexpr size_t NUM_THREADS_X_IN_WARP = WARP_TILE_DIM_X / THREAD_TILE_DIM_X;\nconstexpr size_t NUM_THREADS_Y_IN_WARP = kThreadsPerWarp / NUM_THREADS_X_IN_WARP;\n\n#define MIN(a, b) (a < b ? a : b)\n\ntemplate \n__global__ void __launch_bounds__(THREADS_PER_BLOCK)\n block_scaled_cast_transpose_kernel(const IType* const input, OType* const output_c,\n OType* const output_t, CType* const tile_scales_inv_c,\n CType* const tile_scales_inv_t, const size_t row_length,\n const size_t num_rows, const size_t scale_stride_x,\n const size_t scale_stride_y, const size_t scale_t_stride_x,\n const size_t scale_t_stride_y, const float epsilon,\n const __grid_constant__ CUtensorMap tensor_map_output_t,\n bool pow_2_scaling, const float* noop_ptr) {\n using IVec = Vec;\n using OVecCast = Vec;\n using OVecTrans = Vec;\n\n if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {\n return;\n }\n\n // shared mem for amax reduction in entire block, each warp produces one amax, there are\n // NUM_WARPS_IN_BLOCK amax to reduce\n __shared__ CType block_tile_amax_shared[NUM_WARPS_IN_BLOCK];\n\n IVec thrd_tile_input[THREAD_TILE_DIM_Y];\n constexpr int THREAD_TILE_DIM_X_ = kReturnTranspose ? THREAD_TILE_DIM_X : 1;\n OVecTrans thrd_tile_out_trans[THREAD_TILE_DIM_X_];\n\n const int tid_in_warp = threadIdx.x % kThreadsPerWarp;\n const int tid_in_warp_x = tid_in_warp % NUM_THREADS_X_IN_WARP;\n const int tid_in_warp_y = tid_in_warp / NUM_THREADS_X_IN_WARP;\n const int warp_id_in_block = threadIdx.x / kThreadsPerWarp;\n const int warp_id_in_block_x = warp_id_in_block % NUM_WARPS_X_IN_BLOCK;\n const int warp_id_in_block_y = warp_id_in_block / NUM_WARPS_X_IN_BLOCK;\n\n // This is ONLY true if the input is a full tile\n const int tile_id_x = blockIdx.x;\n const int tile_id_y = blockIdx.y;\n\n const size_t block_tile_start_idx =\n tile_id_y * BLOCK_TILE_DIM * row_length + tile_id_x * BLOCK_TILE_DIM;\n const size_t warp_tile_start_idx =\n block_tile_start_idx +\n warp_id_in_block_y * THREAD_TILE_DIM_Y * NUM_THREADS_Y_IN_WARP * row_length +\n warp_id_in_block_x * THREAD_TILE_DIM_X * NUM_THREADS_X_IN_WARP;\n const size_t thread_tile_start_idx = warp_tile_start_idx +\n tid_in_warp_y * THREAD_TILE_DIM_Y * row_length +\n tid_in_warp_x * THREAD_TILE_DIM_X;\n\n CType warp_tile_amax;\n CType block_tile_amax;\n CType block_tile_scale;\n CType amax = 0;\n\n// Step 1: Load a block tile of input data into thread tiles on registers\n#pragma unroll\n for (int i = 0; i < THREAD_TILE_DIM_Y; i++) {\n thrd_tile_input[i].load_from(input + thread_tile_start_idx + i * row_length);\n }\n\n // Step 2: calculate block tile amax and scale\n // Calculate thread_tile amax\n for (int i = 0; i < THREAD_TILE_DIM_Y; i++) {\n#pragma unroll\n for (int j = 0; j < THREAD_TILE_DIM_X; j++) {\n __builtin_assume(amax >= 0);\n amax = fmaxf(amax, fabsf(static_cast(thrd_tile_input[i].data.elt[j])));\n }\n }\n // Reduce amax in the warp (32x32 tile)\n warp_tile_amax = warp_reduce_max(amax);\n // broadcast the amax to all threads in a warp from the lane 0\n constexpr int lane_zero = 0;\n warp_tile_amax = __shfl_sync(0xFFFFFFFF, warp_tile_amax, lane_zero);\n\n // reduce warp_tile_amax across multiple warps in a thread block using shared mem\n if (tid_in_warp == 0) {\n block_tile_amax_shared[warp_id_in_block_y * NUM_WARPS_X_IN_BLOCK + warp_id_in_block_x] =\n warp_tile_amax;\n }\n __syncthreads();\n // only 8 elements needs reduction, if using reduction tree, multiple _syncthreads will be needed,\n // instead we just let thread 0 do the job\n if (threadIdx.x == 0) {\n CType blk_amax = block_tile_amax_shared[0];\n#pragma unroll\n for (int idx = 1; idx < NUM_WARPS_IN_BLOCK; idx++) {\n blk_amax = fmaxf(blk_amax, block_tile_amax_shared[idx]);\n }\n block_tile_amax_shared[0] = blk_amax;\n }\n __syncthreads();\n block_tile_amax = block_tile_amax_shared[0];\n\n block_tile_scale =\n compute_scale_from_types(block_tile_amax, epsilon, pow_2_scaling);\n\n if (threadIdx.x == 0) {\n static_assert(std::is_same::value);\n const CType scale_inv = 1.0f / block_tile_scale;\n\n size_t row_idx = tile_id_y;\n size_t col_idx = tile_id_x;\n tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv;\n\n if constexpr (kReturnTranspose) {\n row_idx = tile_id_x;\n col_idx = tile_id_y;\n tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv;\n }\n }\n\n // Step 3: Store cast output, Step 4: do transpose within thread tile\n OVecCast tmp_output_c;\n\n for (int i = 0; i < THREAD_TILE_DIM_Y; i++) {\n#pragma unroll\n for (int j = 0; j < THREAD_TILE_DIM_X; j++) {\n // Step 3: Store cast output\n CType scale_data = block_tile_scale;\n\n OType scaled_elt =\n static_cast(static_cast(thrd_tile_input[i].data.elt[j]) * scale_data);\n tmp_output_c.data.elt[j] = scaled_elt;\n // Step 4: do transpose within thread tile\n if constexpr (kReturnTranspose) {\n thrd_tile_out_trans[j].data.elt[i] = scaled_elt;\n }\n }\n tmp_output_c.store_to(output_c + thread_tile_start_idx + i * row_length);\n }\n\n // Step 4: store transpose into shared memory\n if constexpr (kReturnTranspose) {\n#ifdef TMA_HW_SUPPORTED\n __shared__ alignas(128)\n OVecTrans block_tile_trans_shared[SHARED_BLOCK_TILE_DIM_Y][SHARED_BLOCK_TILE_DIM_X_BANKS];\n OType(*block_tile_trans_shared_otype_ptr)[BLOCK_TILE_DIM] =\n reinterpret_cast(block_tile_trans_shared);\n\n#pragma unroll\n for (int i = 0; i < THREAD_TILE_DIM_X; i++) {\n auto warp_id_in_block_x_ = warp_id_in_block_y;\n auto warp_id_in_block_y_ = warp_id_in_block_x;\n int row_idx = warp_id_in_block_y_ * THREAD_TILE_DIM_X * NUM_THREADS_X_IN_WARP +\n tid_in_warp_x * THREAD_TILE_DIM_X + i;\n int col_idx =\n warp_id_in_block_x_ * (NUM_BANKS_Y_IN_WARP / NUM_BANKS_PER_SHARED_ELEM) + tid_in_warp_y;\n block_tile_trans_shared[row_idx][col_idx] = thrd_tile_out_trans[i];\n }\n\n // Wait for shared memory writes to be visible to TMA engine.\n ptx::fence_proxy_async_shared_cta();\n __syncthreads();\n // After syncthreads, writes by all threads are visible to TMA engine.\n\n // Step 5: store transpose output\n // Initiate TMA transfer to copy shared memory to global memory\n if (threadIdx.x == 0) {\n ptx::cp_async_bulk_tensor_2d_shared_to_global(\n reinterpret_cast(&tensor_map_output_t), tile_id_y * BLOCK_TILE_DIM,\n tile_id_x * BLOCK_TILE_DIM,\n reinterpret_cast(block_tile_trans_shared_otype_ptr));\n // Wait for TMA transfer to have finished reading shared memory.\n // Create a \"bulk async-group\" out of the previous bulk copy operation.\n ptx::cp_async_bulk_commit_group();\n // Wait for the group to have completed reading from shared memory.\n ptx::cp_async_bulk_wait_group_read<0>();\n }\n#else\n // Step 4 Alternative (when TMA is not available, skip writing to shared memory)\n const size_t block_tile_t_start_idx =\n tile_id_x * BLOCK_TILE_DIM * num_rows + tile_id_y * BLOCK_TILE_DIM;\n const size_t warp_tile_t_start_idx =\n block_tile_t_start_idx +\n warp_id_in_block_x * THREAD_TILE_DIM_X * NUM_THREADS_X_IN_WARP * num_rows +\n warp_id_in_block_y * THREAD_TILE_DIM_Y * NUM_THREADS_Y_IN_WARP;\n const size_t thread_tile_t_start_idx = warp_tile_t_start_idx +\n tid_in_warp_x * THREAD_TILE_DIM_X * num_rows +\n tid_in_warp_y * THREAD_TILE_DIM_Y;\n#pragma unroll\n for (int i = 0; i < THREAD_TILE_DIM_X; i++) {\n thrd_tile_out_trans[i].store_to(output_t + thread_tile_t_start_idx + i * num_rows);\n }\n#endif\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(THREADS_PER_BLOCK) block_scaled_cast_transpose_kernel_notaligned(\n const IType* const input, OType* const output_c, OType* const output_t,\n CType* const tile_scales_inv_c, CType* const tile_scales_inv_t, const size_t row_length,\n const size_t num_rows, const size_t scale_stride_x, const size_t scale_stride_y,\n const size_t scale_t_stride_x, const size_t scale_t_stride_y, const float epsilon,\n bool pow_2_scaling, const float* noop_ptr) {\n using IVec = Vec;\n using OVecCast = Vec;\n using OVecTrans = Vec;\n\n if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {\n return;\n }\n\n // shared mem for amax reduction in entire block, each warp produces one amax, there are\n // NUM_WARPS_IN_BLOCK amax to reduce\n __shared__ CType block_tile_amax_shared[NUM_WARPS_IN_BLOCK];\n\n IVec thrd_tile_input[THREAD_TILE_DIM_Y];\n constexpr int THREAD_TILE_DIM_X_ = kReturnTranspose ? THREAD_TILE_DIM_X : 1;\n OVecTrans thrd_tile_out_trans[THREAD_TILE_DIM_X_];\n\n const int tid_in_warp = threadIdx.x % kThreadsPerWarp;\n const int tid_in_warp_x = tid_in_warp % NUM_THREADS_X_IN_WARP;\n const int tid_in_warp_y = tid_in_warp / NUM_THREADS_X_IN_WARP;\n const int warp_id_in_block = threadIdx.x / kThreadsPerWarp;\n const int warp_id_in_block_x = warp_id_in_block % NUM_WARPS_X_IN_BLOCK;\n const int warp_id_in_block_y = warp_id_in_block / NUM_WARPS_X_IN_BLOCK;\n\n const int tile_id_x = blockIdx.x;\n const int tile_id_y = blockIdx.y;\n\n const size_t block_tile_start_row_idx = tile_id_y * BLOCK_TILE_DIM;\n const size_t block_tile_start_col_idx = tile_id_x * BLOCK_TILE_DIM;\n const size_t block_tile_start_idx =\n block_tile_start_row_idx * row_length + block_tile_start_col_idx;\n const size_t warp_tile_start_idx =\n block_tile_start_idx +\n warp_id_in_block_y * THREAD_TILE_DIM_Y * NUM_THREADS_Y_IN_WARP * row_length +\n warp_id_in_block_x * THREAD_TILE_DIM_X * NUM_THREADS_X_IN_WARP;\n const size_t thread_tile_start_idx = warp_tile_start_idx +\n tid_in_warp_y * THREAD_TILE_DIM_Y * row_length +\n tid_in_warp_x * THREAD_TILE_DIM_X;\n\n // handle non-full tile\n // check for three cases: full thread tile, nonfull thread tile, empty thread tile\n // for empty thread tile, directly write zero to the transposed shared mem buffer\n // for nonfull thread tile, fill zero to thread tile and act as if it's full\n const size_t thread_tile_start_row_idx =\n tile_id_y * BLOCK_TILE_DIM + warp_id_in_block_y * THREAD_TILE_DIM_Y * NUM_THREADS_Y_IN_WARP +\n tid_in_warp_y * THREAD_TILE_DIM_Y;\n const size_t thread_tile_start_col_idx =\n tile_id_x * BLOCK_TILE_DIM + warp_id_in_block_x * THREAD_TILE_DIM_X * NUM_THREADS_X_IN_WARP +\n tid_in_warp_x * THREAD_TILE_DIM_X;\n\n const size_t thread_tile_end_row_idx = thread_tile_start_row_idx + THREAD_TILE_DIM_Y - 1;\n const size_t thread_tile_end_col_idx = thread_tile_start_col_idx + THREAD_TILE_DIM_X - 1;\n\n bool full_thrd_tile =\n (thread_tile_end_row_idx < num_rows) && (thread_tile_end_col_idx < row_length);\n bool empty_thrd_tile =\n (thread_tile_start_row_idx >= num_rows) || (thread_tile_start_col_idx >= row_length);\n bool nonfull_thrd_tile = (!full_thrd_tile) && (!empty_thrd_tile);\n\n const size_t thread_tile_ncols =\n MIN(THREAD_TILE_DIM_X,\n (MIN(thread_tile_end_col_idx, row_length - 1) - thread_tile_start_col_idx + 1));\n const size_t thread_tile_nrows =\n MIN(THREAD_TILE_DIM_Y,\n (MIN(thread_tile_end_row_idx, num_rows - 1) - thread_tile_start_row_idx + 1));\n\n CType warp_tile_amax;\n CType block_tile_amax;\n CType block_tile_scale;\n CType amax = 0;\n\n if (!empty_thrd_tile) {\n // Step 1: Load a block tile of input data into thread tiles on registers\n // Edge case: nonfull thread tile case, will use the partial load function here\n if (nonfull_thrd_tile) {\n#pragma unroll\n for (int i = 0; i < THREAD_TILE_DIM_Y; i++) {\n if (i >= thread_tile_nrows) {\n thrd_tile_input[i].clear();\n } else {\n thrd_tile_input[i].load_from_elts(input + thread_tile_start_idx + i * row_length, 0,\n thread_tile_ncols);\n }\n }\n } else {\n#pragma unroll\n for (int i = 0; i < THREAD_TILE_DIM_Y; i++) {\n thrd_tile_input[i].load_from_elts(input + thread_tile_start_idx + i * row_length, 0,\n THREAD_TILE_DIM_X);\n }\n }\n\n // Step 2: calculate block tile amax and scale\n // Calculate thread_tile amax\n for (int i = 0; i < THREAD_TILE_DIM_Y; i++) {\n#pragma unroll\n for (int j = 0; j < THREAD_TILE_DIM_X; j++) {\n __builtin_assume(amax >= 0);\n amax = fmaxf(amax, fabsf(static_cast(thrd_tile_input[i].data.elt[j])));\n }\n }\n }\n // Reduce amax in the warp (32x32 tile)\n warp_tile_amax = warp_reduce_max(amax);\n // broadcast the amax to all threads in a warp from the lane 0\n constexpr int lane_zero = 0;\n warp_tile_amax = __shfl_sync(0xFFFFFFFF, warp_tile_amax, lane_zero);\n\n // reduce warp_tile_amax across multiple warps in a thread block using shared mem\n if (tid_in_warp == 0) {\n block_tile_amax_shared[warp_id_in_block_y * NUM_WARPS_X_IN_BLOCK + warp_id_in_block_x] =\n warp_tile_amax;\n }\n __syncthreads();\n // only 8 elements needs reduction, if using reduction tree, multiple _syncthreads will be needed,\n // instead we just let thread 0 do the job\n if (threadIdx.x == 0) {\n CType blk_amax = block_tile_amax_shared[0];\n#pragma unroll\n for (int idx = 1; idx < NUM_WARPS_IN_BLOCK; idx++) {\n blk_amax = fmaxf(blk_amax, block_tile_amax_shared[idx]);\n }\n block_tile_amax_shared[0] = blk_amax;\n }\n __syncthreads();\n block_tile_amax = block_tile_amax_shared[0];\n\n block_tile_scale =\n compute_scale_from_types(block_tile_amax, epsilon, pow_2_scaling);\n\n if (threadIdx.x == 0) {\n static_assert(std::is_same::value);\n const CType scale_inv = 1.0f / block_tile_scale;\n\n size_t row_idx = tile_id_y;\n size_t col_idx = tile_id_x;\n tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv;\n\n if constexpr (kReturnTranspose) {\n row_idx = tile_id_x;\n col_idx = tile_id_y;\n tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv;\n }\n }\n\n // Step 3: Store cast output, Step 4: do transpose within thread tile\n // Edge case: in the non-full tile case, there are three subcases\n // for full thread tile, it's the same thing here\n // for nonfull thread tile, pay attention when saving tmp_output_c to global\n // memory, cannot vec store_to, but need to elt store to for empty tile,\n // it should not enter this step, skip to Step 4\n\n // set thrd_tile_out_trans to all zero\n if constexpr (kReturnTranspose) {\n#pragma unroll\n for (int j = 0; j < THREAD_TILE_DIM_X; j++) {\n thrd_tile_out_trans[j].clear();\n }\n }\n\n if (!empty_thrd_tile) {\n OVecCast tmp_output_c;\n for (int i = 0; i < THREAD_TILE_DIM_Y; i++) {\n if (i >= thread_tile_nrows) {\n continue;\n }\n#pragma unroll\n for (int j = 0; j < THREAD_TILE_DIM_X; j++) {\n // Step 3: Store cast output\n CType scale_data = block_tile_scale;\n\n OType scaled_elt =\n static_cast(static_cast(thrd_tile_input[i].data.elt[j]) * scale_data);\n tmp_output_c.data.elt[j] = scaled_elt;\n // Step 4: do transpose within thread tile\n if constexpr (kReturnTranspose) {\n thrd_tile_out_trans[j].data.elt[i] = scaled_elt;\n }\n }\n tmp_output_c.store_to_elts(output_c + thread_tile_start_idx + i * row_length, 0,\n thread_tile_ncols);\n }\n\n if constexpr (kReturnTranspose) {\n const size_t block_tile_t_start_idx =\n tile_id_x * BLOCK_TILE_DIM * num_rows + tile_id_y * BLOCK_TILE_DIM;\n const size_t warp_tile_t_start_idx =\n block_tile_t_start_idx +\n warp_id_in_block_x * THREAD_TILE_DIM_X * NUM_THREADS_X_IN_WARP * num_rows +\n warp_id_in_block_y * THREAD_TILE_DIM_Y * NUM_THREADS_Y_IN_WARP;\n const size_t thread_tile_t_start_idx = warp_tile_t_start_idx +\n tid_in_warp_x * THREAD_TILE_DIM_X * num_rows +\n tid_in_warp_y * THREAD_TILE_DIM_Y;\n#pragma unroll\n for (int i = 0; i < thread_tile_ncols; i++) {\n thrd_tile_out_trans[i].store_to_elts(output_t + thread_tile_t_start_idx + i * num_rows, 0,\n thread_tile_nrows);\n }\n }\n }\n}\n\ntemplate \nCUtensorMap get_tensor_map(const SimpleTensor& tensor, size_t global_dim_x, size_t global_dim_y) {\n CUtensorMapDataType dataType;\n if constexpr (std::is_same_v ||\n std::is_same_v) {\n dataType = CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT8;\n } else {\n NVTE_ERROR(\n \"Invalid output type for blockwise transpose (must be FP8: Float8E4M3 or Float8E5M2).\");\n }\n\n CUtensorMap tensor_map_output_trans{};\n create_2D_tensor_map(tensor_map_output_trans, tensor, global_dim_y, global_dim_x,\n /*shmemY=*/BLOCK_TILE_DIM, /*shmemX=*/BLOCK_TILE_DIM,\n /*stride_elems=*/global_dim_x, /*offset_elems=*/0, sizeof(OutputType) * 8);\n return tensor_map_output_trans;\n}\n\n} // namespace\n} // namespace transformer_engine\n\nnamespace transformer_engine::detail {\n\nvoid quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor& scale_inv,\n SimpleTensor& scale_inv_t, SimpleTensor& output,\n SimpleTensor& output_t, const float epsilon,\n const bool return_transpose, const bool pow_2_scale,\n const SimpleTensor& noop_tensor, cudaStream_t stream) {\n NVTE_API_CALL(quantize_transpose_square_blockwise);\n checkCuDriverContext(stream);\n\n if (transformer_engine::cuda::sm_arch() >= 100) {\n NVTE_CHECK(pow_2_scale, \"On Blackwell and newer, the FP8 block scaling recipe is emulated \",\n \"with MXFP8, which requires using power of two scaling factors.\");\n }\n\n NVTE_CHECK(input.shape == output.shape, \"Input and output must have the same shape.\");\n const size_t row_length = input.shape.size() > 0 ? input.shape.back() : 1;\n size_t num_rows = 1;\n for (size_t i = 0; (i < input.shape.size() - 1) && (input.shape.size() > 0); ++i) {\n num_rows *= input.shape.at(i);\n }\n\n NVTE_CHECK(scale_inv.shape.size() == 2, \"scale_inv must have 2 dimensions.\");\n\n size_t scale_k = scale_inv.shape[1];\n\n const size_t scale_stride_x = 1;\n const size_t scale_stride_y = scale_k;\n\n size_t scale_t_stride_x = 0;\n size_t scale_t_stride_y = 0;\n\n const float* noop_ptr = reinterpret_cast(noop_tensor.dptr);\n\n if (return_transpose) {\n NVTE_CHECK(output_t.shape.size() == input.shape.size(), \"input (shape=\", input.shape,\n \") and output_t (shape=\", output_t.shape, \") have incompatible dims.\");\n if (output_t.shape.size() > 0) {\n NVTE_CHECK(output_t.shape.front() == input.shape.back(), \"input (shape=\", input.shape,\n \") and output_t (shape=\", output_t.shape, \") have incompatible dims.\");\n for (size_t i = 1; i < output_t.shape.size(); ++i) {\n NVTE_CHECK(output_t.shape[i] == input.shape[i - 1], \"input (shape=\", input.shape,\n \") and output_t (shape=\", output_t.shape, \") have incompatible dims.\");\n }\n }\n NVTE_CHECK(output.dtype == output_t.dtype, \"output and output_t need to have the same type.\");\n\n NVTE_CHECK(scale_inv_t.shape.size() == 2, \"scale_inv_t must have 2 dimensions.\");\n\n scale_t_stride_x = 1;\n scale_t_stride_y = scale_inv_t.shape[1];\n }\n\n const size_t num_blocks_x = DIVUP(row_length, BLOCK_TILE_DIM);\n const size_t num_blocks_y = DIVUP(num_rows, BLOCK_TILE_DIM);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.dtype, InputType,\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(\n output.dtype, OutputType,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_transpose, kReturnTranspose,\n\n dim3 grid(num_blocks_x, num_blocks_y, 1);\n const bool full_tile =\n row_length % BLOCK_TILE_DIM == 0 && num_rows % BLOCK_TILE_DIM == 0;\n\n if (full_tile) {\n CUtensorMap tensor_map_output_trans;\n if (return_transpose) {\n tensor_map_output_trans =\n get_tensor_map(output_t, num_rows, row_length);\n }\n block_scaled_cast_transpose_kernel\n <<>>(\n reinterpret_cast(input.dptr),\n reinterpret_cast(output.dptr),\n reinterpret_cast(output_t.dptr),\n reinterpret_cast(scale_inv.dptr),\n reinterpret_cast(scale_inv_t.dptr), row_length, num_rows,\n scale_stride_x, scale_stride_y, scale_t_stride_x, scale_t_stride_y, epsilon,\n tensor_map_output_trans, pow_2_scale, noop_ptr);\n } else {\n block_scaled_cast_transpose_kernel_notaligned\n <<>>(\n reinterpret_cast(input.dptr),\n reinterpret_cast(output.dptr),\n reinterpret_cast(output_t.dptr),\n reinterpret_cast(scale_inv.dptr),\n reinterpret_cast(scale_inv_t.dptr), row_length, num_rows,\n scale_stride_x, scale_stride_y, scale_t_stride_x, scale_t_stride_y, epsilon,\n pow_2_scale, noop_ptr);\n } // full-tile\n ) // return_transpose\n ) // OutputType\n ) // InputType\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n} // namespace transformer_engine::detail\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/recipe/recipe_common.cuh\"\n#include \"common/util/hip/hip_runtime.h\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n\n#if (!defined(__CUDA_MINIMUM_ARCH__) && __CUDA_ARCH__ >= 900) || \\\n (defined(__CUDA_MINIMUM_ARCH__) && __CUDA_MINIMUM_ARCH__ >= 900)\n#define TMA_HW_SUPPORTED\n#endif\n\nnamespace transformer_engine {\nnamespace {\n\n// const values configuration\n\nconstexpr size_t kThreadsPerWarp = 32;\n#ifdef TMA_HW_SUPPORTED\nconstexpr size_t BLOCK_TILE_DIM = 128;\nconstexpr size_t WARP_TILE_DIM_X = 32;\nconstexpr size_t WARP_TILE_DIM_Y = 64;\nconstexpr size_t THREAD_TILE_DIM_X = 16;\nconstexpr size_t THREAD_TILE_DIM_Y = 4;\n#else\nconstexpr size_t BLOCK_TILE_DIM = 128;\nconstexpr size_t WARP_TILE_DIM_X = 64;\nconstexpr size_t WARP_TILE_DIM_Y = 32;\nconstexpr size_t THREAD_TILE_DIM_X = 8;\nconstexpr size_t THREAD_TILE_DIM_Y = 8;\n#endif\n\n#ifdef TMA_HW_SUPPORTED\nconstexpr size_t NUM_BYTES_PER_BANK = 4;\nconstexpr size_t NUM_BANKS_PER_SHARED_ELEM = THREAD_TILE_DIM_Y / NUM_BYTES_PER_BANK;\nconstexpr size_t SHARED_BLOCK_TILE_DIM_Y = BLOCK_TILE_DIM;\nconstexpr size_t SHARED_BLOCK_TILE_DIM_X_BANKS =\n BLOCK_TILE_DIM / (NUM_BYTES_PER_BANK * NUM_BANKS_PER_SHARED_ELEM);\nconstexpr size_t NUM_BANKS_Y_IN_WARP = WARP_TILE_DIM_Y / NUM_BYTES_PER_BANK;\n#endif\nconstexpr size_t ELE_PER_THREAD = THREAD_TILE_DIM_X * THREAD_TILE_DIM_Y;\nconstexpr size_t THREADS_PER_BLOCK = BLOCK_TILE_DIM * BLOCK_TILE_DIM / ELE_PER_THREAD;\nconstexpr size_t NUM_WARPS_X_IN_BLOCK = BLOCK_TILE_DIM / WARP_TILE_DIM_X;\nconstexpr size_t NUM_WARPS_Y_IN_BLOCK = BLOCK_TILE_DIM / WARP_TILE_DIM_Y;\nconstexpr size_t NUM_WARPS_IN_BLOCK = NUM_WARPS_X_IN_BLOCK * NUM_WARPS_Y_IN_BLOCK;\n\nconstexpr size_t NUM_THREADS_X_IN_WARP = WARP_TILE_DIM_X / THREAD_TILE_DIM_X;\nconstexpr size_t NUM_THREADS_Y_IN_WARP = kThreadsPerWarp / NUM_THREADS_X_IN_WARP;\n\n#define MIN(a, b) (a < b ? a : b)\n\ntemplate \n__global__ void __launch_bounds__(THREADS_PER_BLOCK)\n block_scaled_cast_transpose_kernel(const IType* const input, OType* const output_c,\n OType* const output_t, CType* const tile_scales_inv_c,\n CType* const tile_scales_inv_t, const size_t row_length,\n const size_t num_rows, const size_t scale_stride_x,\n const size_t scale_stride_y, const size_t scale_t_stride_x,\n const size_t scale_t_stride_y, const float epsilon,\n const __grid_constant__ CUtensorMap tensor_map_output_t,\n bool pow_2_scaling, const float* noop_ptr) {\n using IVec = Vec;\n using OVecCast = Vec;\n using OVecTrans = Vec;\n\n if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {\n return;\n }\n\n // shared mem for amax reduction in entire block, each warp produces one amax, there are\n // NUM_WARPS_IN_BLOCK amax to reduce\n __shared__ CType block_tile_amax_shared[NUM_WARPS_IN_BLOCK];\n\n IVec thrd_tile_input[THREAD_TILE_DIM_Y];\n constexpr int THREAD_TILE_DIM_X_ = kReturnTranspose ? THREAD_TILE_DIM_X : 1;\n OVecTrans thrd_tile_out_trans[THREAD_TILE_DIM_X_];\n\n const int tid_in_warp = threadIdx.x % kThreadsPerWarp;\n const int tid_in_warp_x = tid_in_warp % NUM_THREADS_X_IN_WARP;\n const int tid_in_warp_y = tid_in_warp / NUM_THREADS_X_IN_WARP;\n const int warp_id_in_block = threadIdx.x / kThreadsPerWarp;\n const int warp_id_in_block_x = warp_id_in_block % NUM_WARPS_X_IN_BLOCK;\n const int warp_id_in_block_y = warp_id_in_block / NUM_WARPS_X_IN_BLOCK;\n\n // This is ONLY true if the input is a full tile\n const int tile_id_x = blockIdx.x;\n const int tile_id_y = blockIdx.y;\n\n const size_t block_tile_start_idx =\n tile_id_y * BLOCK_TILE_DIM * row_length + tile_id_x * BLOCK_TILE_DIM;\n const size_t warp_tile_start_idx =\n block_tile_start_idx +\n warp_id_in_block_y * THREAD_TILE_DIM_Y * NUM_THREADS_Y_IN_WARP * row_length +\n warp_id_in_block_x * THREAD_TILE_DIM_X * NUM_THREADS_X_IN_WARP;\n const size_t thread_tile_start_idx = warp_tile_start_idx +\n tid_in_warp_y * THREAD_TILE_DIM_Y * row_length +\n tid_in_warp_x * THREAD_TILE_DIM_X;\n\n CType warp_tile_amax;\n CType block_tile_amax;\n CType block_tile_scale;\n CType amax = 0;\n\n// Step 1: Load a block tile of input data into thread tiles on registers\n#pragma unroll\n for (int i = 0; i < THREAD_TILE_DIM_Y; i++) {\n thrd_tile_input[i].load_from(input + thread_tile_start_idx + i * row_length);\n }\n\n // Step 2: calculate block tile amax and scale\n // Calculate thread_tile amax\n for (int i = 0; i < THREAD_TILE_DIM_Y; i++) {\n#pragma unroll\n for (int j = 0; j < THREAD_TILE_DIM_X; j++) {\n __builtin_assume(amax >= 0);\n amax = fmaxf(amax, fabsf(static_cast(thrd_tile_input[i].data.elt[j])));\n }\n }\n // Reduce amax in the warp (32x32 tile)\n warp_tile_amax = warp_reduce_max(amax);\n // broadcast the amax to all threads in a warp from the lane 0\n constexpr int lane_zero = 0;\n warp_tile_amax = __shfl_sync(0xFFFFFFFF, warp_tile_amax, lane_zero);\n\n // reduce warp_tile_amax across multiple warps in a thread block using shared mem\n if (tid_in_warp == 0) {\n block_tile_amax_shared[warp_id_in_block_y * NUM_WARPS_X_IN_BLOCK + warp_id_in_block_x] =\n warp_tile_amax;\n }\n __syncthreads();\n // only 8 elements needs reduction, if using reduction tree, multiple _syncthreads will be needed,\n // instead we just let thread 0 do the job\n if (threadIdx.x == 0) {\n CType blk_amax = block_tile_amax_shared[0];\n#pragma unroll\n for (int idx = 1; idx < NUM_WARPS_IN_BLOCK; idx++) {\n blk_amax = fmaxf(blk_amax, block_tile_amax_shared[idx]);\n }\n block_tile_amax_shared[0] = blk_amax;\n }\n __syncthreads();\n block_tile_amax = block_tile_amax_shared[0];\n\n block_tile_scale =\n compute_scale_from_types(block_tile_amax, epsilon, pow_2_scaling);\n\n if (threadIdx.x == 0) {\n static_assert(std::is_same::value);\n const CType scale_inv = 1.0f / block_tile_scale;\n\n size_t row_idx = tile_id_y;\n size_t col_idx = tile_id_x;\n tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv;\n\n if constexpr (kReturnTranspose) {\n row_idx = tile_id_x;\n col_idx = tile_id_y;\n tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv;\n }\n }\n\n // Step 3: Store cast output, Step 4: do transpose within thread tile\n OVecCast tmp_output_c;\n\n for (int i = 0; i < THREAD_TILE_DIM_Y; i++) {\n#pragma unroll\n for (int j = 0; j < THREAD_TILE_DIM_X; j++) {\n // Step 3: Store cast output\n CType scale_data = block_tile_scale;\n\n OType scaled_elt =\n static_cast(static_cast(thrd_tile_input[i].data.elt[j]) * scale_data);\n tmp_output_c.data.elt[j] = scaled_elt;\n // Step 4: do transpose within thread tile\n if constexpr (kReturnTranspose) {\n thrd_tile_out_trans[j].data.elt[i] = scaled_elt;\n }\n }\n tmp_output_c.store_to(output_c + thread_tile_start_idx + i * row_length);\n }\n\n // Step 4: store transpose into shared memory\n if constexpr (kReturnTranspose) {\n#ifdef TMA_HW_SUPPORTED\n __shared__ alignas(128)\n OVecTrans block_tile_trans_shared[SHARED_BLOCK_TILE_DIM_Y][SHARED_BLOCK_TILE_DIM_X_BANKS];\n OType(*block_tile_trans_shared_otype_ptr)[BLOCK_TILE_DIM] =\n reinterpret_cast(block_tile_trans_shared);\n\n#pragma unroll\n for (int i = 0; i < THREAD_TILE_DIM_X; i++) {\n auto warp_id_in_block_x_ = warp_id_in_block_y;\n auto warp_id_in_block_y_ = warp_id_in_block_x;\n int row_idx = warp_id_in_block_y_ * THREAD_TILE_DIM_X * NUM_THREADS_X_IN_WARP +\n tid_in_warp_x * THREAD_TILE_DIM_X + i;\n int col_idx =\n warp_id_in_block_x_ * (NUM_BANKS_Y_IN_WARP / NUM_BANKS_PER_SHARED_ELEM) + tid_in_warp_y;\n block_tile_trans_shared[row_idx][col_idx] = thrd_tile_out_trans[i];\n }\n\n // Wait for shared memory writes to be visible to TMA engine.\n ptx::fence_proxy_async_shared_cta();\n __syncthreads();\n // After syncthreads, writes by all threads are visible to TMA engine.\n\n // Step 5: store transpose output\n // Initiate TMA transfer to copy shared memory to global memory\n if (threadIdx.x == 0) {\n ptx::cp_async_bulk_tensor_2d_shared_to_global(\n reinterpret_cast(&tensor_map_output_t), tile_id_y * BLOCK_TILE_DIM,\n tile_id_x * BLOCK_TILE_DIM,\n reinterpret_cast(block_tile_trans_shared_otype_ptr));\n // Wait for TMA transfer to have finished reading shared memory.\n // Create a \"bulk async-group\" out of the previous bulk copy operation.\n ptx::cp_async_bulk_commit_group();\n // Wait for the group to have completed reading from shared memory.\n ptx::cp_async_bulk_wait_group_read<0>();\n }\n#else\n // Step 4 Alternative (when TMA is not available, skip writing to shared memory)\n const size_t block_tile_t_start_idx =\n tile_id_x * BLOCK_TILE_DIM * num_rows + tile_id_y * BLOCK_TILE_DIM;\n const size_t warp_tile_t_start_idx =\n block_tile_t_start_idx +\n warp_id_in_block_x * THREAD_TILE_DIM_X * NUM_THREADS_X_IN_WARP * num_rows +\n warp_id_in_block_y * THREAD_TILE_DIM_Y * NUM_THREADS_Y_IN_WARP;\n const size_t thread_tile_t_start_idx = warp_tile_t_start_idx +\n tid_in_warp_x * THREAD_TILE_DIM_X * num_rows +\n tid_in_warp_y * THREAD_TILE_DIM_Y;\n#pragma unroll\n for (int i = 0; i < THREAD_TILE_DIM_X; i++) {\n thrd_tile_out_trans[i].store_to(output_t + thread_tile_t_start_idx + i * num_rows);\n }\n#endif\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(THREADS_PER_BLOCK) block_scaled_cast_transpose_kernel_notaligned(\n const IType* const input, OType* const output_c, OType* const output_t,\n CType* const tile_scales_inv_c, CType* const tile_scales_inv_t, const size_t row_length,\n const size_t num_rows, const size_t scale_stride_x, const size_t scale_stride_y,\n const size_t scale_t_stride_x, const size_t scale_t_stride_y, const float epsilon,\n bool pow_2_scaling, const float* noop_ptr) {\n using IVec = Vec;\n using OVecCast = Vec;\n using OVecTrans = Vec;\n\n if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {\n return;\n }\n\n // shared mem for amax reduction in entire block, each warp produces one amax, there are\n // NUM_WARPS_IN_BLOCK amax to reduce\n __shared__ CType block_tile_amax_shared[NUM_WARPS_IN_BLOCK];\n\n IVec thrd_tile_input[THREAD_TILE_DIM_Y];\n constexpr int THREAD_TILE_DIM_X_ = kReturnTranspose ? THREAD_TILE_DIM_X : 1;\n OVecTrans thrd_tile_out_trans[THREAD_TILE_DIM_X_];\n\n const int tid_in_warp = threadIdx.x % kThreadsPerWarp;\n const int tid_in_warp_x = tid_in_warp % NUM_THREADS_X_IN_WARP;\n const int tid_in_warp_y = tid_in_warp / NUM_THREADS_X_IN_WARP;\n const int warp_id_in_block = threadIdx.x / kThreadsPerWarp;\n const int warp_id_in_block_x = warp_id_in_block % NUM_WARPS_X_IN_BLOCK;\n const int warp_id_in_block_y = warp_id_in_block / NUM_WARPS_X_IN_BLOCK;\n\n const int tile_id_x = blockIdx.x;\n const int tile_id_y = blockIdx.y;\n\n const size_t block_tile_start_row_idx = tile_id_y * BLOCK_TILE_DIM;\n const size_t block_tile_start_col_idx = tile_id_x * BLOCK_TILE_DIM;\n const size_t block_tile_start_idx =\n block_tile_start_row_idx * row_length + block_tile_start_col_idx;\n const size_t warp_tile_start_idx =\n block_tile_start_idx +\n warp_id_in_block_y * THREAD_TILE_DIM_Y * NUM_THREADS_Y_IN_WARP * row_length +\n warp_id_in_block_x * THREAD_TILE_DIM_X * NUM_THREADS_X_IN_WARP;\n const size_t thread_tile_start_idx = warp_tile_start_idx +\n tid_in_warp_y * THREAD_TILE_DIM_Y * row_length +\n tid_in_warp_x * THREAD_TILE_DIM_X;\n\n // handle non-full tile\n // check for three cases: full thread tile, nonfull thread tile, empty thread tile\n // for empty thread tile, directly write zero to the transposed shared mem buffer\n // for nonfull thread tile, fill zero to thread tile and act as if it's full\n const size_t thread_tile_start_row_idx =\n tile_id_y * BLOCK_TILE_DIM + warp_id_in_block_y * THREAD_TILE_DIM_Y * NUM_THREADS_Y_IN_WARP +\n tid_in_warp_y * THREAD_TILE_DIM_Y;\n const size_t thread_tile_start_col_idx =\n tile_id_x * BLOCK_TILE_DIM + warp_id_in_block_x * THREAD_TILE_DIM_X * NUM_THREADS_X_IN_WARP +\n tid_in_warp_x * THREAD_TILE_DIM_X;\n\n const size_t thread_tile_end_row_idx = thread_tile_start_row_idx + THREAD_TILE_DIM_Y - 1;\n const size_t thread_tile_end_col_idx = thread_tile_start_col_idx + THREAD_TILE_DIM_X - 1;\n\n bool full_thrd_tile =\n (thread_tile_end_row_idx < num_rows) && (thread_tile_end_col_idx < row_length);\n bool empty_thrd_tile =\n (thread_tile_start_row_idx >= num_rows) || (thread_tile_start_col_idx >= row_length);\n bool nonfull_thrd_tile = (!full_thrd_tile) && (!empty_thrd_tile);\n\n const size_t thread_tile_ncols =\n MIN(THREAD_TILE_DIM_X,\n (MIN(thread_tile_end_col_idx, row_length - 1) - thread_tile_start_col_idx + 1));\n const size_t thread_tile_nrows =\n MIN(THREAD_TILE_DIM_Y,\n (MIN(thread_tile_end_row_idx, num_rows - 1) - thread_tile_start_row_idx + 1));\n\n CType warp_tile_amax;\n CType block_tile_amax;\n CType block_tile_scale;\n CType amax = 0;\n\n if (!empty_thrd_tile) {\n // Step 1: Load a block tile of input data into thread tiles on registers\n // Edge case: nonfull thread tile case, will use the partial load function here\n if (nonfull_thrd_tile) {\n#pragma unroll\n for (int i = 0; i < THREAD_TILE_DIM_Y; i++) {\n if (i >= thread_tile_nrows) {\n thrd_tile_input[i].clear();\n } else {\n thrd_tile_input[i].load_from_elts(input + thread_tile_start_idx + i * row_length, 0,\n thread_tile_ncols);\n }\n }\n } else {\n#pragma unroll\n for (int i = 0; i < THREAD_TILE_DIM_Y; i++) {\n thrd_tile_input[i].load_from_elts(input + thread_tile_start_idx + i * row_length, 0,\n THREAD_TILE_DIM_X);\n }\n }\n\n // Step 2: calculate block tile amax and scale\n // Calculate thread_tile amax\n for (int i = 0; i < THREAD_TILE_DIM_Y; i++) {\n#pragma unroll\n for (int j = 0; j < THREAD_TILE_DIM_X; j++) {\n __builtin_assume(amax >= 0);\n amax = fmaxf(amax, fabsf(static_cast(thrd_tile_input[i].data.elt[j])));\n }\n }\n }\n // Reduce amax in the warp (32x32 tile)\n warp_tile_amax = warp_reduce_max(amax);\n // broadcast the amax to all threads in a warp from the lane 0\n constexpr int lane_zero = 0;\n warp_tile_amax = __shfl_sync(0xFFFFFFFF, warp_tile_amax, lane_zero);\n\n // reduce warp_tile_amax across multiple warps in a thread block using shared mem\n if (tid_in_warp == 0) {\n block_tile_amax_shared[warp_id_in_block_y * NUM_WARPS_X_IN_BLOCK + warp_id_in_block_x] =\n warp_tile_amax;\n }\n __syncthreads();\n // only 8 elements needs reduction, if using reduction tree, multiple _syncthreads will be needed,\n // instead we just let thread 0 do the job\n if (threadIdx.x == 0) {\n CType blk_amax = block_tile_amax_shared[0];\n#pragma unroll\n for (int idx = 1; idx < NUM_WARPS_IN_BLOCK; idx++) {\n blk_amax = fmaxf(blk_amax, block_tile_amax_shared[idx]);\n }\n block_tile_amax_shared[0] = blk_amax;\n }\n __syncthreads();\n block_tile_amax = block_tile_amax_shared[0];\n\n block_tile_scale =\n compute_scale_from_types(block_tile_amax, epsilon, pow_2_scaling);\n\n if (threadIdx.x == 0) {\n static_assert(std::is_same::value);\n const CType scale_inv = 1.0f / block_tile_scale;\n\n size_t row_idx = tile_id_y;\n size_t col_idx = tile_id_x;\n tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv;\n\n if constexpr (kReturnTranspose) {\n row_idx = tile_id_x;\n col_idx = tile_id_y;\n tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv;\n }\n }\n\n // Step 3: Store cast output, Step 4: do transpose within thread tile\n // Edge case: in the non-full tile case, there are three subcases\n // for full thread tile, it's the same thing here\n // for nonfull thread tile, pay attention when saving tmp_output_c to global\n // memory, cannot vec store_to, but need to elt store to for empty tile,\n // it should not enter this step, skip to Step 4\n\n // set thrd_tile_out_trans to all zero\n if constexpr (kReturnTranspose) {\n#pragma unroll\n for (int j = 0; j < THREAD_TILE_DIM_X; j++) {\n thrd_tile_out_trans[j].clear();\n }\n }\n\n if (!empty_thrd_tile) {\n OVecCast tmp_output_c;\n for (int i = 0; i < THREAD_TILE_DIM_Y; i++) {\n if (i >= thread_tile_nrows) {\n continue;\n }\n#pragma unroll\n for (int j = 0; j < THREAD_TILE_DIM_X; j++) {\n // Step 3: Store cast output\n CType scale_data = block_tile_scale;\n\n OType scaled_elt =\n static_cast(static_cast(thrd_tile_input[i].data.elt[j]) * scale_data);\n tmp_output_c.data.elt[j] = scaled_elt;\n // Step 4: do transpose within thread tile\n if constexpr (kReturnTranspose) {\n thrd_tile_out_trans[j].data.elt[i] = scaled_elt;\n }\n }\n tmp_output_c.store_to_elts(output_c + thread_tile_start_idx + i * row_length, 0,\n thread_tile_ncols);\n }\n\n if constexpr (kReturnTranspose) {\n const size_t block_tile_t_start_idx =\n tile_id_x * BLOCK_TILE_DIM * num_rows + tile_id_y * BLOCK_TILE_DIM;\n const size_t warp_tile_t_start_idx =\n block_tile_t_start_idx +\n warp_id_in_block_x * THREAD_TILE_DIM_X * NUM_THREADS_X_IN_WARP * num_rows +\n warp_id_in_block_y * THREAD_TILE_DIM_Y * NUM_THREADS_Y_IN_WARP;\n const size_t thread_tile_t_start_idx = warp_tile_t_start_idx +\n tid_in_warp_x * THREAD_TILE_DIM_X * num_rows +\n tid_in_warp_y * THREAD_TILE_DIM_Y;\n#pragma unroll\n for (int i = 0; i < thread_tile_ncols; i++) {\n thrd_tile_out_trans[i].store_to_elts(output_t + thread_tile_t_start_idx + i * num_rows, 0,\n thread_tile_nrows);\n }\n }\n }\n}\n\ntemplate \nCUtensorMap get_tensor_map(const SimpleTensor& tensor, size_t global_dim_x, size_t global_dim_y) {\n CUtensorMapDataType dataType;\n if constexpr (std::is_same_v ||\n std::is_same_v) {\n dataType = CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT8;\n } else {\n NVTE_ERROR(\n \"Invalid output type for blockwise transpose (must be FP8: Float8E4M3 or Float8E5M2).\");\n }\n\n CUtensorMap tensor_map_output_trans{};\n create_2D_tensor_map(tensor_map_output_trans, tensor, global_dim_y, global_dim_x,\n /*shmemY=*/BLOCK_TILE_DIM, /*shmemX=*/BLOCK_TILE_DIM,\n /*stride_elems=*/global_dim_x, /*offset_elems=*/0, sizeof(OutputType) * 8);\n return tensor_map_output_trans;\n}\n\n} // namespace\n} // namespace transformer_engine\n\nnamespace transformer_engine::detail {\n\nvoid quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor& scale_inv,\n SimpleTensor& scale_inv_t, SimpleTensor& output,\n SimpleTensor& output_t, const float epsilon,\n const bool return_transpose, const bool pow_2_scale,\n const SimpleTensor& noop_tensor, hipStream_t stream) {\n NVTE_API_CALL(quantize_transpose_square_blockwise);\n checkCuDriverContext(stream);\n\n if (transformer_engine::cuda::sm_arch() >= 100) {\n NVTE_CHECK(pow_2_scale, \"On Blackwell and newer, the FP8 block scaling recipe is emulated \",\n \"with MXFP8, which requires using power of two scaling factors.\");\n }\n\n NVTE_CHECK(input.shape == output.shape, \"Input and output must have the same shape.\");\n const size_t row_length = input.shape.size() > 0 ? input.shape.back() : 1;\n size_t num_rows = 1;\n for (size_t i = 0; (i < input.shape.size() - 1) && (input.shape.size() > 0); ++i) {\n num_rows *= input.shape.at(i);\n }\n\n NVTE_CHECK(scale_inv.shape.size() == 2, \"scale_inv must have 2 dimensions.\");\n\n size_t scale_k = scale_inv.shape[1];\n\n const size_t scale_stride_x = 1;\n const size_t scale_stride_y = scale_k;\n\n size_t scale_t_stride_x = 0;\n size_t scale_t_stride_y = 0;\n\n const float* noop_ptr = reinterpret_cast(noop_tensor.dptr);\n\n if (return_transpose) {\n NVTE_CHECK(output_t.shape.size() == input.shape.size(), \"input (shape=\", input.shape,\n \") and output_t (shape=\", output_t.shape, \") have incompatible dims.\");\n if (output_t.shape.size() > 0) {\n NVTE_CHECK(output_t.shape.front() == input.shape.back(), \"input (shape=\", input.shape,\n \") and output_t (shape=\", output_t.shape, \") have incompatible dims.\");\n for (size_t i = 1; i < output_t.shape.size(); ++i) {\n NVTE_CHECK(output_t.shape[i] == input.shape[i - 1], \"input (shape=\", input.shape,\n \") and output_t (shape=\", output_t.shape, \") have incompatible dims.\");\n }\n }\n NVTE_CHECK(output.dtype == output_t.dtype, \"output and output_t need to have the same type.\");\n\n NVTE_CHECK(scale_inv_t.shape.size() == 2, \"scale_inv_t must have 2 dimensions.\");\n\n scale_t_stride_x = 1;\n scale_t_stride_y = scale_inv_t.shape[1];\n }\n\n const size_t num_blocks_x = DIVUP(row_length, BLOCK_TILE_DIM);\n const size_t num_blocks_y = DIVUP(num_rows, BLOCK_TILE_DIM);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.dtype, InputType,\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(\n output.dtype, OutputType,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_transpose, kReturnTranspose,\n\n dim3 grid(num_blocks_x, num_blocks_y, 1);\n const bool full_tile =\n row_length % BLOCK_TILE_DIM == 0 && num_rows % BLOCK_TILE_DIM == 0;\n\n if (full_tile) {\n CUtensorMap tensor_map_output_trans;\n if (return_transpose) {\n tensor_map_output_trans =\n get_tensor_map(output_t, num_rows, row_length);\n }\n block_scaled_cast_transpose_kernel\n <<>>(\n reinterpret_cast(input.dptr),\n reinterpret_cast(output.dptr),\n reinterpret_cast(output_t.dptr),\n reinterpret_cast(scale_inv.dptr),\n reinterpret_cast(scale_inv_t.dptr), row_length, num_rows,\n scale_stride_x, scale_stride_y, scale_t_stride_x, scale_t_stride_y, epsilon,\n tensor_map_output_trans, pow_2_scale, noop_ptr);\n } else {\n block_scaled_cast_transpose_kernel_notaligned\n <<>>(\n reinterpret_cast(input.dptr),\n reinterpret_cast(output.dptr),\n reinterpret_cast(output_t.dptr),\n reinterpret_cast(scale_inv.dptr),\n reinterpret_cast(scale_inv_t.dptr), row_length, num_rows,\n scale_stride_x, scale_stride_y, scale_t_stride_x, scale_t_stride_y, epsilon,\n pow_2_scale, noop_ptr);\n } // full-tile\n ) // return_transpose\n ) // OutputType\n ) // InputType\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n} // namespace transformer_engine::detail\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"utils.cuh\"\n\nusing namespace transformer_engine;\n\nnamespace {\n\n// Parameters\nusing VectorType = BytesToType<__VECTOR_SIZE__>::Type;\nconstexpr size_t block_size = __BLOCK_SIZE__;\n\n} // namespace\n\n__global__ void __launch_bounds__(block_size)\n swap_first_dims_kernel(const VectorType* __restrict__ const input,\n VectorType* __restrict__ const output, const size_t dim0,\n const size_t dim1, const size_t dim2) {\n const size_t gid = threadIdx.x + blockIdx.x * block_size;\n#if __SINGLE_LOAD_STORE__\n const auto idx = gid;\n#else\n const size_t nthreads = gridDim.x * block_size;\n for (size_t idx = gid; idx < dim0 * dim1 * dim2; idx += nthreads)\n#endif // __SINGLE_LOAD_STORE__\n {\n const auto idx2 = idx % dim2;\n const auto idx1 = (idx / dim2) % dim1;\n const auto idx0 = (idx / dim2) / dim1;\n const auto in_offset = idx1 * dim0 * dim2 + idx0 * dim2 + idx2;\n output[idx] = input[in_offset];\n }\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \"utils.cuh\"\n\nusing namespace transformer_engine;\n\nnamespace {\n\n// Parameters\nusing VectorType = BytesToType<__VECTOR_SIZE__>::Type;\nconstexpr size_t block_size = __BLOCK_SIZE__;\n\n} // namespace\n\n__global__ void __launch_bounds__(block_size)\n swap_first_dims_kernel(const VectorType* __restrict__ const input,\n VectorType* __restrict__ const output, const size_t dim0,\n const size_t dim1, const size_t dim2) {\n const size_t gid = threadIdx.x + blockIdx.x * block_size;\n#if __SINGLE_LOAD_STORE__\n const auto idx = gid;\n#else\n const size_t nthreads = gridDim.x * block_size;\n for (size_t idx = gid; idx < dim0 * dim1 * dim2; idx += nthreads)\n#endif // __SINGLE_LOAD_STORE__\n {\n const auto idx2 = idx % dim2;\n const auto idx1 = (idx / dim2) % dim1;\n const auto idx0 = (idx / dim2) / dim1;\n const auto in_offset = idx1 * dim0 * dim2 + idx0 * dim2 + idx2;\n output[idx] = input[in_offset];\n }\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/transpose/rtc/swap_first_dims.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"utils.cuh\"\n\nusing namespace transformer_engine;\n\nnamespace {\n\n// Parameters\nusing Type = __TYPE__;\nconstexpr size_t load_size = __LOAD_SIZE__;\nconstexpr size_t store_size = __STORE_SIZE__;\nconstexpr size_t warps_per_tile = __WARPS_PER_TILE__;\nconstexpr size_t block_size = __BLOCK_SIZE__;\n\n} // namespace\n\n__global__ void __launch_bounds__(block_size)\n transpose_optimized_kernel(const Type* __restrict__ const input, const float* const noop,\n Type* __restrict__ const output, const size_t row_length,\n const size_t num_rows) {\n if (noop != nullptr && noop[0] == 1.0f) return;\n\n // Vectorized load/store sizes\n constexpr size_t nvec_in = load_size / sizeof(Type);\n constexpr size_t nvec_out = store_size / sizeof(Type);\n using IVec = Vec;\n using OVec = Vec;\n\n // Thread indices\n // Note: Block is interpreted as a warp_size x num_warps grid\n constexpr size_t bdimx = THREADS_PER_WARP;\n constexpr size_t bdimy = warps_per_tile;\n const size_t tid = threadIdx.x;\n const size_t tidx = tid % bdimx;\n const size_t tidy = tid / bdimx;\n const size_t bid = blockIdx.x;\n\n // Input tensors are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec_out x nvec_in subtiles\n constexpr size_t tile_dim_m = THREADS_PER_WARP * nvec_out;\n constexpr size_t tile_dim_n = THREADS_PER_WARP * nvec_in;\n\n // Position of tile within tensor\n const size_t num_tiles_m = num_rows / tile_dim_m;\n const size_t tile_id_m = bid % num_tiles_m;\n const size_t tile_id_n = bid / num_tiles_m;\n const size_t tile_row = tile_id_m * tile_dim_m;\n const size_t tile_col = tile_id_n * tile_dim_n;\n\n // Number of nvec_out x nvec_in subtiles for each thread to\n // load/store\n constexpr size_t num_iterations = THREADS_PER_WARP / warps_per_tile;\n\n // Load input to registers and transpose\n // Note: Each thread loads num_iterations subtiles and transposes in\n // registers.\n OVec local_output[nvec_in][num_iterations];\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n#pragma unroll\n for (size_t i2 = 0; i2 < nvec_out; ++i2) {\n const size_t row = tile_row + i1 * nvec_out + i2;\n const size_t col = tile_col + j1 * nvec_in;\n IVec local_input;\n local_input.load_from(&input[row * row_length + col]);\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n local_output[j2][iter].data.elt[i2] = local_input.data.elt[j2];\n }\n }\n }\n\n // Copy from registers to shared memory to global memory\n __shared__ OVec shared_output[THREADS_PER_WARP][THREADS_PER_WARP + 1];\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n shared_output[j1][i1] = local_output[j2][iter];\n }\n __syncthreads();\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidx;\n const size_t j1 = tidy + iter * bdimy;\n const size_t row = tile_row + i1 * nvec_out;\n const size_t col = tile_col + j1 * nvec_in + j2;\n shared_output[j1][i1].store_to(&output[col * num_rows + row]);\n }\n __syncthreads();\n }\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \"utils.cuh\"\n\nusing namespace transformer_engine;\n\nnamespace {\n\n// Parameters\nusing Type = __TYPE__;\nconstexpr size_t load_size = __LOAD_SIZE__;\nconstexpr size_t store_size = __STORE_SIZE__;\nconstexpr size_t warps_per_tile = __WARPS_PER_TILE__;\nconstexpr size_t block_size = __BLOCK_SIZE__;\n\n} // namespace\n\n__global__ void __launch_bounds__(block_size)\n transpose_optimized_kernel(const Type* __restrict__ const input, const float* const noop,\n Type* __restrict__ const output, const size_t row_length,\n const size_t num_rows) {\n if (noop != nullptr && noop[0] == 1.0f) return;\n\n // Vectorized load/store sizes\n constexpr size_t nvec_in = load_size / sizeof(Type);\n constexpr size_t nvec_out = store_size / sizeof(Type);\n using IVec = Vec;\n using OVec = Vec;\n\n // Thread indices\n // Note: Block is interpreted as a warp_size x num_warps grid\n constexpr size_t bdimx = THREADS_PER_WARP;\n constexpr size_t bdimy = warps_per_tile;\n const size_t tid = threadIdx.x;\n const size_t tidx = tid % bdimx;\n const size_t tidy = tid / bdimx;\n const size_t bid = blockIdx.x;\n\n // Input tensors are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec_out x nvec_in subtiles\n constexpr size_t tile_dim_m = THREADS_PER_WARP * nvec_out;\n constexpr size_t tile_dim_n = THREADS_PER_WARP * nvec_in;\n\n // Position of tile within tensor\n const size_t num_tiles_m = num_rows / tile_dim_m;\n const size_t tile_id_m = bid % num_tiles_m;\n const size_t tile_id_n = bid / num_tiles_m;\n const size_t tile_row = tile_id_m * tile_dim_m;\n const size_t tile_col = tile_id_n * tile_dim_n;\n\n // Number of nvec_out x nvec_in subtiles for each thread to\n // load/store\n constexpr size_t num_iterations = THREADS_PER_WARP / warps_per_tile;\n\n // Load input to registers and transpose\n // Note: Each thread loads num_iterations subtiles and transposes in\n // registers.\n OVec local_output[nvec_in][num_iterations];\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n#pragma unroll\n for (size_t i2 = 0; i2 < nvec_out; ++i2) {\n const size_t row = tile_row + i1 * nvec_out + i2;\n const size_t col = tile_col + j1 * nvec_in;\n IVec local_input;\n local_input.load_from(&input[row * row_length + col]);\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n local_output[j2][iter].data.elt[i2] = local_input.data.elt[j2];\n }\n }\n }\n\n // Copy from registers to shared memory to global memory\n __shared__ OVec shared_output[THREADS_PER_WARP][THREADS_PER_WARP + 1];\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n shared_output[j1][i1] = local_output[j2][iter];\n }\n __syncthreads();\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidx;\n const size_t j1 = tidy + iter * bdimy;\n const size_t row = tile_row + i1 * nvec_out;\n const size_t col = tile_col + j1 * nvec_in + j2;\n shared_output[j1][i1].store_to(&output[col * num_rows + row]);\n }\n __syncthreads();\n }\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/transpose/rtc/transpose.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"util/math.h\"\n#include \"utils.cuh\"\n\nusing namespace transformer_engine;\n\nnamespace {\n\n// Parameters\nusing CType = float;\nusing IType = __ITYPE__;\nusing IType2 = __ITYPE2__;\nusing OType = __OTYPE__;\nconstexpr size_t LOAD_SIZE = __LOAD_SIZE__;\nconstexpr size_t STORE_SIZE = __STORE_SIZE__;\nconstexpr size_t WARPS_PER_TILE = __WARPS_PER_TILE__;\nconstexpr size_t BLOCK_SIZE = __BLOCK_SIZE__;\nconstexpr bool IS_DBIAS = __IS_DBIAS__;\nconstexpr bool IS_DACT = __IS_DACT__;\nconstexpr bool IS_ACT = __IS_ACT__;\nstatic_assert(!(IS_DACT && IS_ACT), \"forward and backward activation are mutually exclusive\");\nconstexpr size_t ACT_TYPE = __ACTIVATION_TYPE__;\n\nconstexpr size_t NVEC_IN = LOAD_SIZE / sizeof(IType);\nconstexpr size_t NVEC_OUT = STORE_SIZE / sizeof(OType);\nusing CVec = Vec;\nusing IVec = Vec;\nusing IVec2 = Vec;\nusing OVec = Vec;\nusing Param = CTDBiasDActParam;\n\nusing OP = CType (*)(const CType, const Empty &);\nconstexpr OP ActivationList[] = {\n nullptr, // 0\n &sigmoid, // 1\n &dsigmoid, // 2\n &gelu, // 3\n &dgelu, // 4\n &qgelu, // 5\n &dqgelu, // 6\n &silu, // 7\n &dsilu, // 8\n &relu, // 9\n &drelu, // 10\n &srelu, // 11\n &dsrelu // 12\n};\n\n} // namespace\n\ninline __device__ void cast_and_transpose_regs_optimized(const CVec (&in)[NVEC_OUT],\n OVec (&out_trans)[NVEC_IN],\n CVec &out_dbias, // NOLINT(*)\n typename OVec::type *output_cast_tile,\n const size_t current_place,\n const size_t stride, const CType scale,\n CType &amax, // NOLINT(*)\n const int dbias_shfl_src_lane) {\n using OVecC = Vec;\n\n CVec step_dbias;\n if constexpr (IS_DBIAS) {\n step_dbias.clear();\n }\n\n#pragma unroll\n for (unsigned int i = 0; i < NVEC_OUT; ++i) {\n OVecC out_cast;\n#pragma unroll\n for (unsigned int j = 0; j < NVEC_IN; ++j) {\n const CType tmp = in[i].data.elt[j];\n if constexpr (IS_DBIAS) {\n step_dbias.data.elt[j] += tmp; // dbias: thread tile local accumulation\n }\n out_cast.data.elt[j] = static_cast(tmp * scale);\n out_trans[j].data.elt[i] = static_cast(tmp * scale); // thread tile transpose\n\n __builtin_assume(amax >= 0);\n amax = fmaxf(fabsf(tmp), amax);\n }\n out_cast.store_to(output_cast_tile, current_place + stride * i);\n }\n\n if constexpr (IS_DBIAS) {\n#pragma unroll\n for (unsigned int j = 0; j < NVEC_IN; ++j) {\n CType elt = step_dbias.data.elt[j];\n elt = __shfl_sync(0xffffffff, elt, dbias_shfl_src_lane); // shuffle data in a warp\n out_dbias.data.elt[j] += elt;\n }\n }\n}\n\n__global__ void __launch_bounds__(BLOCK_SIZE)\n cast_transpose_fusion_kernel_optimized(const Param param, const size_t row_length,\n const size_t num_rows, const size_t num_tiles) {\n extern __shared__ char scratch[];\n\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const unsigned int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;\n const size_t num_tiles_x = row_length / (NVEC_IN * THREADS_PER_WARP);\n const size_t tile_id =\n blockIdx.x * blockDim.x / (THREADS_PER_WARP * WARPS_PER_TILE) + warp_id / WARPS_PER_TILE;\n if (tile_id >= num_tiles) {\n return;\n }\n\n const size_t tile_id_x = tile_id % num_tiles_x;\n const size_t tile_id_y = tile_id / num_tiles_x;\n\n const size_t tile_offset =\n (tile_id_x * NVEC_IN + tile_id_y * row_length * NVEC_OUT) * THREADS_PER_WARP;\n const size_t tile_offset_transp =\n (tile_id_y * NVEC_OUT + tile_id_x * num_rows * NVEC_IN) * THREADS_PER_WARP;\n\n const IType *const my_input_tile = param.input + tile_offset;\n const IType2 *const my_act_input_tile = param.act_input + tile_offset;\n OType *const my_output_c_tile = param.output_c + tile_offset;\n OType *const my_output_t_tile = param.output_t + tile_offset_transp;\n CType *const my_partial_dbias_tile =\n param.workspace + (tile_id_x * (NVEC_IN * THREADS_PER_WARP) + tile_id_y * row_length);\n\n OVec *const my_scratch =\n reinterpret_cast(scratch) +\n (my_id_in_warp + warp_id / WARPS_PER_TILE * THREADS_PER_WARP) * (THREADS_PER_WARP + 1);\n\n CVec *const my_dbias_scratch = reinterpret_cast(scratch);\n\n IVec in[2][NVEC_OUT];\n IVec2 act_in[2][NVEC_OUT];\n\n const unsigned int warp_id_in_tile = warp_id % WARPS_PER_TILE;\n constexpr unsigned int n_iterations = THREADS_PER_WARP / WARPS_PER_TILE;\n OVec out_space[n_iterations][NVEC_IN];\n\n const size_t stride = row_length / NVEC_IN;\n const size_t output_stride = num_rows / NVEC_OUT;\n size_t current_stride = warp_id_in_tile * n_iterations * NVEC_OUT * stride;\n size_t current_row = (tile_id_y * THREADS_PER_WARP + warp_id_in_tile * n_iterations) * NVEC_OUT;\n unsigned int my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n\n CType amax = 0.0f;\n const CType scale = param.scale_ptr != nullptr ? *param.scale_ptr : 1;\n\n CVec partial_dbias;\n if constexpr (IS_DBIAS) {\n partial_dbias.clear();\n }\n\n#pragma unroll\n for (unsigned int i = 0; i < NVEC_OUT; ++i) {\n in[0][i].load_from(my_input_tile, current_stride + my_place + stride * i);\n if constexpr (IS_DACT) {\n act_in[0][i].load_from(my_act_input_tile, current_stride + my_place + stride * i);\n }\n }\n#pragma unroll\n for (unsigned int i = 0; i < n_iterations; ++i) {\n const size_t current_place = current_stride + my_place;\n const unsigned int my_place_in = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n const unsigned int current_in = (i + 1) % 2;\n if (i < n_iterations - 1) {\n#pragma unroll\n for (unsigned int j = 0; j < NVEC_OUT; ++j) {\n const size_t ld_offset = current_stride + my_place_in + stride * (NVEC_OUT + j);\n in[current_in][j].load_from(my_input_tile, ld_offset);\n if constexpr (IS_DACT) {\n act_in[current_in][j].load_from(my_act_input_tile, ld_offset);\n }\n }\n }\n CVec in_cast_fp32[NVEC_OUT]; // NOLINT(*)\n#pragma unroll\n for (unsigned int j = 0; j < NVEC_OUT; ++j) {\n#pragma unroll\n for (unsigned int k = 0; k < NVEC_IN; ++k) {\n if constexpr (IS_DACT) {\n in_cast_fp32[j].data.elt[k] =\n static_cast(in[current_in ^ 1][j].data.elt[k]) *\n ActivationList[ACT_TYPE](act_in[current_in ^ 1][j].data.elt[k], {});\n } else if constexpr (IS_ACT) {\n in_cast_fp32[j].data.elt[k] =\n ActivationList[ACT_TYPE](in[current_in ^ 1][j].data.elt[k], {});\n } else {\n in_cast_fp32[j].data.elt[k] = static_cast(in[current_in ^ 1][j].data.elt[k]);\n }\n }\n }\n\n const int dbias_shfl_src_lane =\n (my_id_in_warp + i + warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n\n cast_and_transpose_regs_optimized(in_cast_fp32, out_space[i], partial_dbias, my_output_c_tile,\n current_place, stride, scale, amax, dbias_shfl_src_lane);\n\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += NVEC_OUT * stride;\n current_row += NVEC_OUT;\n }\n\n#pragma unroll\n for (unsigned int i = 0; i < NVEC_IN; ++i) {\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * NVEC_IN;\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile,\n current_stride + my_place);\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * NVEC_IN;\n }\n __syncthreads();\n }\n\n if constexpr (IS_DBIAS) {\n my_dbias_scratch[threadIdx.x] = partial_dbias;\n __syncthreads();\n if (warp_id_in_tile == 0) {\n#pragma unroll\n for (unsigned int i = 1; i < WARPS_PER_TILE; ++i) {\n CVec tmp = my_dbias_scratch[threadIdx.x + i * THREADS_PER_WARP];\n#pragma unroll\n for (unsigned int j = 0; j < NVEC_IN; ++j) {\n partial_dbias.data.elt[j] += tmp.data.elt[j];\n }\n }\n partial_dbias.store_to(my_partial_dbias_tile, my_id_in_warp);\n }\n }\n\n // Reduce amax over block\n if (param.amax != nullptr) {\n const CType max_block = reduce_max(amax, warp_id);\n if (threadIdx.x == 0) {\n atomicMaxFloat(param.amax, max_block);\n }\n }\n\n // Update scale-inverse\n if (blockIdx.x == 0 && threadIdx.x == 0 && param.scale_inv != nullptr) {\n reciprocal(param.scale_inv, scale);\n }\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \"util/math.h\"\n#include \"utils.cuh\"\n\nusing namespace transformer_engine;\n\nnamespace {\n\n// Parameters\nusing CType = float;\nusing IType = __ITYPE__;\nusing IType2 = __ITYPE2__;\nusing OType = __OTYPE__;\nconstexpr size_t LOAD_SIZE = __LOAD_SIZE__;\nconstexpr size_t STORE_SIZE = __STORE_SIZE__;\nconstexpr size_t WARPS_PER_TILE = __WARPS_PER_TILE__;\nconstexpr size_t BLOCK_SIZE = __BLOCK_SIZE__;\nconstexpr bool IS_DBIAS = __IS_DBIAS__;\nconstexpr bool IS_DACT = __IS_DACT__;\nconstexpr bool IS_ACT = __IS_ACT__;\nstatic_assert(!(IS_DACT && IS_ACT), \"forward and backward activation are mutually exclusive\");\nconstexpr size_t ACT_TYPE = __ACTIVATION_TYPE__;\n\nconstexpr size_t NVEC_IN = LOAD_SIZE / sizeof(IType);\nconstexpr size_t NVEC_OUT = STORE_SIZE / sizeof(OType);\nusing CVec = Vec;\nusing IVec = Vec;\nusing IVec2 = Vec;\nusing OVec = Vec;\nusing Param = CTDBiasDActParam;\n\nusing OP = CType (*)(const CType, const Empty &);\nconstexpr OP ActivationList[] = {\n nullptr, // 0\n &sigmoid, // 1\n &dsigmoid, // 2\n &gelu, // 3\n &dgelu, // 4\n &qgelu, // 5\n &dqgelu, // 6\n &silu, // 7\n &dsilu, // 8\n &relu, // 9\n &drelu, // 10\n &srelu, // 11\n &dsrelu // 12\n};\n\n} // namespace\n\ninline __device__ void cast_and_transpose_regs_optimized(const CVec (&in)[NVEC_OUT],\n OVec (&out_trans)[NVEC_IN],\n CVec &out_dbias, // NOLINT(*)\n typename OVec::type *output_cast_tile,\n const size_t current_place,\n const size_t stride, const CType scale,\n CType &amax, // NOLINT(*)\n const int dbias_shfl_src_lane) {\n using OVecC = Vec;\n\n CVec step_dbias;\n if constexpr (IS_DBIAS) {\n step_dbias.clear();\n }\n\n#pragma unroll\n for (unsigned int i = 0; i < NVEC_OUT; ++i) {\n OVecC out_cast;\n#pragma unroll\n for (unsigned int j = 0; j < NVEC_IN; ++j) {\n const CType tmp = in[i].data.elt[j];\n if constexpr (IS_DBIAS) {\n step_dbias.data.elt[j] += tmp; // dbias: thread tile local accumulation\n }\n out_cast.data.elt[j] = static_cast(tmp * scale);\n out_trans[j].data.elt[i] = static_cast(tmp * scale); // thread tile transpose\n\n __builtin_assume(amax >= 0);\n amax = fmaxf(fabsf(tmp), amax);\n }\n out_cast.store_to(output_cast_tile, current_place + stride * i);\n }\n\n if constexpr (IS_DBIAS) {\n#pragma unroll\n for (unsigned int j = 0; j < NVEC_IN; ++j) {\n CType elt = step_dbias.data.elt[j];\n elt = __shfl_sync(0xffffffff, elt, dbias_shfl_src_lane); // shuffle data in a warp\n out_dbias.data.elt[j] += elt;\n }\n }\n}\n\n__global__ void __launch_bounds__(BLOCK_SIZE)\n cast_transpose_fusion_kernel_optimized(const Param param, const size_t row_length,\n const size_t num_rows, const size_t num_tiles) {\n extern __shared__ char scratch[];\n\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const unsigned int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;\n const size_t num_tiles_x = row_length / (NVEC_IN * THREADS_PER_WARP);\n const size_t tile_id =\n blockIdx.x * blockDim.x / (THREADS_PER_WARP * WARPS_PER_TILE) + warp_id / WARPS_PER_TILE;\n if (tile_id >= num_tiles) {\n return;\n }\n\n const size_t tile_id_x = tile_id % num_tiles_x;\n const size_t tile_id_y = tile_id / num_tiles_x;\n\n const size_t tile_offset =\n (tile_id_x * NVEC_IN + tile_id_y * row_length * NVEC_OUT) * THREADS_PER_WARP;\n const size_t tile_offset_transp =\n (tile_id_y * NVEC_OUT + tile_id_x * num_rows * NVEC_IN) * THREADS_PER_WARP;\n\n const IType *const my_input_tile = param.input + tile_offset;\n const IType2 *const my_act_input_tile = param.act_input + tile_offset;\n OType *const my_output_c_tile = param.output_c + tile_offset;\n OType *const my_output_t_tile = param.output_t + tile_offset_transp;\n CType *const my_partial_dbias_tile =\n param.workspace + (tile_id_x * (NVEC_IN * THREADS_PER_WARP) + tile_id_y * row_length);\n\n OVec *const my_scratch =\n reinterpret_cast(scratch) +\n (my_id_in_warp + warp_id / WARPS_PER_TILE * THREADS_PER_WARP) * (THREADS_PER_WARP + 1);\n\n CVec *const my_dbias_scratch = reinterpret_cast(scratch);\n\n IVec in[2][NVEC_OUT];\n IVec2 act_in[2][NVEC_OUT];\n\n const unsigned int warp_id_in_tile = warp_id % WARPS_PER_TILE;\n constexpr unsigned int n_iterations = THREADS_PER_WARP / WARPS_PER_TILE;\n OVec out_space[n_iterations][NVEC_IN];\n\n const size_t stride = row_length / NVEC_IN;\n const size_t output_stride = num_rows / NVEC_OUT;\n size_t current_stride = warp_id_in_tile * n_iterations * NVEC_OUT * stride;\n size_t current_row = (tile_id_y * THREADS_PER_WARP + warp_id_in_tile * n_iterations) * NVEC_OUT;\n unsigned int my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n\n CType amax = 0.0f;\n const CType scale = param.scale_ptr != nullptr ? *param.scale_ptr : 1;\n\n CVec partial_dbias;\n if constexpr (IS_DBIAS) {\n partial_dbias.clear();\n }\n\n#pragma unroll\n for (unsigned int i = 0; i < NVEC_OUT; ++i) {\n in[0][i].load_from(my_input_tile, current_stride + my_place + stride * i);\n if constexpr (IS_DACT) {\n act_in[0][i].load_from(my_act_input_tile, current_stride + my_place + stride * i);\n }\n }\n#pragma unroll\n for (unsigned int i = 0; i < n_iterations; ++i) {\n const size_t current_place = current_stride + my_place;\n const unsigned int my_place_in = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n const unsigned int current_in = (i + 1) % 2;\n if (i < n_iterations - 1) {\n#pragma unroll\n for (unsigned int j = 0; j < NVEC_OUT; ++j) {\n const size_t ld_offset = current_stride + my_place_in + stride * (NVEC_OUT + j);\n in[current_in][j].load_from(my_input_tile, ld_offset);\n if constexpr (IS_DACT) {\n act_in[current_in][j].load_from(my_act_input_tile, ld_offset);\n }\n }\n }\n CVec in_cast_fp32[NVEC_OUT]; // NOLINT(*)\n#pragma unroll\n for (unsigned int j = 0; j < NVEC_OUT; ++j) {\n#pragma unroll\n for (unsigned int k = 0; k < NVEC_IN; ++k) {\n if constexpr (IS_DACT) {\n in_cast_fp32[j].data.elt[k] =\n static_cast(in[current_in ^ 1][j].data.elt[k]) *\n ActivationList[ACT_TYPE](act_in[current_in ^ 1][j].data.elt[k], {});\n } else if constexpr (IS_ACT) {\n in_cast_fp32[j].data.elt[k] =\n ActivationList[ACT_TYPE](in[current_in ^ 1][j].data.elt[k], {});\n } else {\n in_cast_fp32[j].data.elt[k] = static_cast(in[current_in ^ 1][j].data.elt[k]);\n }\n }\n }\n\n const int dbias_shfl_src_lane =\n (my_id_in_warp + i + warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n\n cast_and_transpose_regs_optimized(in_cast_fp32, out_space[i], partial_dbias, my_output_c_tile,\n current_place, stride, scale, amax, dbias_shfl_src_lane);\n\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += NVEC_OUT * stride;\n current_row += NVEC_OUT;\n }\n\n#pragma unroll\n for (unsigned int i = 0; i < NVEC_IN; ++i) {\n#pragma unroll\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[(my_id_in_warp + THREADS_PER_WARP - j - warp_id_in_tile * n_iterations) %\n THREADS_PER_WARP] = out_space[j][i];\n }\n __syncthreads();\n my_place =\n (my_id_in_warp + THREADS_PER_WARP - warp_id_in_tile * n_iterations) % THREADS_PER_WARP;\n current_stride = i * output_stride + warp_id_in_tile * n_iterations * output_stride * NVEC_IN;\n for (unsigned int j = 0; j < n_iterations; ++j) {\n my_scratch[j + warp_id_in_tile * n_iterations].store_to(my_output_t_tile,\n current_stride + my_place);\n my_place = (my_place + THREADS_PER_WARP - 1) % THREADS_PER_WARP;\n current_stride += output_stride * NVEC_IN;\n }\n __syncthreads();\n }\n\n if constexpr (IS_DBIAS) {\n my_dbias_scratch[threadIdx.x] = partial_dbias;\n __syncthreads();\n if (warp_id_in_tile == 0) {\n#pragma unroll\n for (unsigned int i = 1; i < WARPS_PER_TILE; ++i) {\n CVec tmp = my_dbias_scratch[threadIdx.x + i * THREADS_PER_WARP];\n#pragma unroll\n for (unsigned int j = 0; j < NVEC_IN; ++j) {\n partial_dbias.data.elt[j] += tmp.data.elt[j];\n }\n }\n partial_dbias.store_to(my_partial_dbias_tile, my_id_in_warp);\n }\n }\n\n // Reduce amax over block\n if (param.amax != nullptr) {\n const CType max_block = reduce_max(amax, warp_id);\n if (threadIdx.x == 0) {\n atomicMaxFloat(param.amax, max_block);\n }\n }\n\n // Update scale-inverse\n if (blockIdx.x == 0 && threadIdx.x == 0 && param.scale_inv != nullptr) {\n reciprocal(param.scale_inv, scale);\n }\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/transpose/rtc/cast_transpose_fusion.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"utils.cuh\"\n\nusing namespace transformer_engine;\n\nnamespace {\n\n// Parameters\nusing CType = float;\nusing IType = __ITYPE__;\nusing OType = __OTYPE__;\nconstexpr size_t load_size = __LOAD_SIZE__;\nconstexpr size_t store_size = __STORE_SIZE__;\nconstexpr size_t warps_per_tile = __WARPS_PER_TILE__;\nconstexpr size_t block_size = __BLOCK_SIZE__;\n\n} // namespace\n\n__global__ void __launch_bounds__(block_size) cast_transpose_optimized_kernel(\n const IType* __restrict__ const input, const CType* __restrict__ const noop,\n OType* __restrict__ const output_c, OType* __restrict__ const output_t,\n const CType* __restrict__ const scale_ptr, CType* __restrict__ const amax_ptr,\n CType* __restrict__ const scale_inv_ptr, const size_t row_length, const size_t num_rows) {\n if (noop != nullptr && noop[0] == 1.0f) return;\n\n // Vectorized load/store sizes\n constexpr size_t nvec_in = load_size / sizeof(IType);\n constexpr size_t nvec_out = store_size / sizeof(OType);\n using IVec = Vec;\n using OVecC = Vec;\n using OVecT = Vec;\n\n // Thread indices\n // Note: Block is interpreted as a warp_size x num_warps grid\n constexpr size_t bdimx = THREADS_PER_WARP;\n constexpr size_t bdimy = warps_per_tile;\n const size_t tid = threadIdx.x;\n const size_t tidx = tid % bdimx;\n const size_t tidy = tid / bdimx;\n const size_t bid = blockIdx.x;\n\n // Input tensors are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec_out x nvec_in subtiles\n constexpr size_t tile_dim_m = THREADS_PER_WARP * nvec_out;\n constexpr size_t tile_dim_n = THREADS_PER_WARP * nvec_in;\n\n // Position of tile within tensor\n const size_t num_tiles_m = num_rows / tile_dim_m;\n const size_t tile_id_m = bid % num_tiles_m;\n const size_t tile_id_n = bid / num_tiles_m;\n const size_t tile_row = tile_id_m * tile_dim_m;\n const size_t tile_col = tile_id_n * tile_dim_n;\n\n // Number of nvec_out x nvec_in subtiles for each thread to\n // load/store\n constexpr size_t num_iterations = THREADS_PER_WARP / warps_per_tile;\n\n // FP8 factors\n const CType scale = scale_ptr == nullptr ? 1 : *scale_ptr;\n CType amax = 0;\n\n // Load input to registers and transpose\n // Note: Each thread loads num_iterations subtiles, computes amax,\n // casts type, and transposes in registers.\n OVecT local_output_t[nvec_in][num_iterations];\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n#pragma unroll\n for (size_t i2 = 0; i2 < nvec_out; ++i2) {\n const size_t row = tile_row + i1 * nvec_out + i2;\n const size_t col = tile_col + j1 * nvec_in;\n IVec local_input;\n OVecC local_output_c;\n local_input.load_from(&input[row * row_length + col]);\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n const CType in = static_cast(local_input.data.elt[j2]);\n const OType out = OType(in * scale);\n __builtin_assume(amax >= 0);\n amax = fmaxf(fabsf(in), amax);\n local_output_c.data.elt[j2] = out;\n local_output_t[j2][iter].data.elt[i2] = out;\n }\n local_output_c.store_to(&output_c[row * row_length + col]);\n }\n }\n\n // Copy from registers to shared memory to global memory\n __shared__ OVecT shared_output_t[THREADS_PER_WARP][THREADS_PER_WARP + 1];\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n shared_output_t[j1][i1] = local_output_t[j2][iter];\n }\n __syncthreads();\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidx;\n const size_t j1 = tidy + iter * bdimy;\n const size_t row = tile_row + i1 * nvec_out;\n const size_t col = tile_col + j1 * nvec_in + j2;\n shared_output_t[j1][i1].store_to(&output_t[col * num_rows + row]);\n }\n __syncthreads();\n }\n\n // Reduce amax over block\n if (amax_ptr != nullptr) {\n amax = reduce_max(amax, tidy);\n if (threadIdx.x == 0) {\n atomicMaxFloat(amax_ptr, amax);\n }\n }\n\n // Update scale-inverse\n if (blockIdx.x == 0 && threadIdx.x == 0 && scale_inv_ptr != nullptr) {\n reciprocal(scale_inv_ptr, scale);\n }\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \"utils.cuh\"\n\nusing namespace transformer_engine;\n\nnamespace {\n\n// Parameters\nusing CType = float;\nusing IType = __ITYPE__;\nusing OType = __OTYPE__;\nconstexpr size_t load_size = __LOAD_SIZE__;\nconstexpr size_t store_size = __STORE_SIZE__;\nconstexpr size_t warps_per_tile = __WARPS_PER_TILE__;\nconstexpr size_t block_size = __BLOCK_SIZE__;\n\n} // namespace\n\n__global__ void __launch_bounds__(block_size) cast_transpose_optimized_kernel(\n const IType* __restrict__ const input, const CType* __restrict__ const noop,\n OType* __restrict__ const output_c, OType* __restrict__ const output_t,\n const CType* __restrict__ const scale_ptr, CType* __restrict__ const amax_ptr,\n CType* __restrict__ const scale_inv_ptr, const size_t row_length, const size_t num_rows) {\n if (noop != nullptr && noop[0] == 1.0f) return;\n\n // Vectorized load/store sizes\n constexpr size_t nvec_in = load_size / sizeof(IType);\n constexpr size_t nvec_out = store_size / sizeof(OType);\n using IVec = Vec;\n using OVecC = Vec;\n using OVecT = Vec;\n\n // Thread indices\n // Note: Block is interpreted as a warp_size x num_warps grid\n constexpr size_t bdimx = THREADS_PER_WARP;\n constexpr size_t bdimy = warps_per_tile;\n const size_t tid = threadIdx.x;\n const size_t tidx = tid % bdimx;\n const size_t tidy = tid / bdimx;\n const size_t bid = blockIdx.x;\n\n // Input tensors are divided into tiles\n // Note: Each tile is a warp_size x warp_size grid of nvec_out x nvec_in subtiles\n constexpr size_t tile_dim_m = THREADS_PER_WARP * nvec_out;\n constexpr size_t tile_dim_n = THREADS_PER_WARP * nvec_in;\n\n // Position of tile within tensor\n const size_t num_tiles_m = num_rows / tile_dim_m;\n const size_t tile_id_m = bid % num_tiles_m;\n const size_t tile_id_n = bid / num_tiles_m;\n const size_t tile_row = tile_id_m * tile_dim_m;\n const size_t tile_col = tile_id_n * tile_dim_n;\n\n // Number of nvec_out x nvec_in subtiles for each thread to\n // load/store\n constexpr size_t num_iterations = THREADS_PER_WARP / warps_per_tile;\n\n // FP8 factors\n const CType scale = scale_ptr == nullptr ? 1 : *scale_ptr;\n CType amax = 0;\n\n // Load input to registers and transpose\n // Note: Each thread loads num_iterations subtiles, computes amax,\n // casts type, and transposes in registers.\n OVecT local_output_t[nvec_in][num_iterations];\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n#pragma unroll\n for (size_t i2 = 0; i2 < nvec_out; ++i2) {\n const size_t row = tile_row + i1 * nvec_out + i2;\n const size_t col = tile_col + j1 * nvec_in;\n IVec local_input;\n OVecC local_output_c;\n local_input.load_from(&input[row * row_length + col]);\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n const CType in = static_cast(local_input.data.elt[j2]);\n const OType out = OType(in * scale);\n __builtin_assume(amax >= 0);\n amax = fmaxf(fabsf(in), amax);\n local_output_c.data.elt[j2] = out;\n local_output_t[j2][iter].data.elt[i2] = out;\n }\n local_output_c.store_to(&output_c[row * row_length + col]);\n }\n }\n\n // Copy from registers to shared memory to global memory\n __shared__ OVecT shared_output_t[THREADS_PER_WARP][THREADS_PER_WARP + 1];\n#pragma unroll\n for (size_t j2 = 0; j2 < nvec_in; ++j2) {\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidy + iter * bdimy;\n const size_t j1 = tidx;\n shared_output_t[j1][i1] = local_output_t[j2][iter];\n }\n __syncthreads();\n#pragma unroll\n for (size_t iter = 0; iter < num_iterations; ++iter) {\n const size_t i1 = tidx;\n const size_t j1 = tidy + iter * bdimy;\n const size_t row = tile_row + i1 * nvec_out;\n const size_t col = tile_col + j1 * nvec_in + j2;\n shared_output_t[j1][i1].store_to(&output_t[col * num_rows + row]);\n }\n __syncthreads();\n }\n\n // Reduce amax over block\n if (amax_ptr != nullptr) {\n amax = reduce_max(amax, tidy);\n if (threadIdx.x == 0) {\n atomicMaxFloat(amax_ptr, amax);\n }\n }\n\n // Update scale-inverse\n if (blockIdx.x == 0 && threadIdx.x == 0 && scale_inv_ptr != nullptr) {\n reciprocal(scale_inv_ptr, scale);\n }\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/transpose/rtc/cast_transpose.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \n\n#include \"../common.h\"\n#include \"../utils.cuh\"\n#include \"transformer_engine/dropout.h\"\n\nnamespace transformer_engine {\nnamespace {\n\n// RNG kernels process chunks of 16 entries\nconstexpr size_t rng_chunk_size = 16;\n\n// CUDA block size\nconstexpr size_t block_size = 128;\n\n// Vector class to help with vectorized memory accesses\ntemplate \nunion Vector {\n using StorageType = typename BytesToType::Type;\n StorageType storage;\n T entries[kSize];\n};\n\n/* Byte-wise less-than comparison\n *\n * Results are stored in each byte's most-significant bit (MSB). All\n * other bits are zero.\n */\n__device__ __forceinline__ uint32_t bytewise_less_than(uint32_t a, uint32_t b) {\n // Compare low bits by masking MSBs and subtracting. The resulting\n // MSBs are 0 if the low bits of a are less than the low bits of b.\n uint32_t result = (a | 0x80808080) - (b & 0x7F7F7F7F);\n\n // Bitwise logical op to get answer in MSBs\n // Equivalent logic: result = (a == b) ? !result : b\n asm(\"lop3.b32 %0, %1, %2, %3, 0x4D;\\n\\t\" : \"=r\"(result) : \"r\"(a), \"r\"(b), \"r\"(result));\n\n // Mask out everything except MSBs and return\n result &= 0x80808080;\n return result;\n}\n\n/* Generate dropout mask with 16 bits.\n *\n * 1 corresponds to keep and 0 to drop.\n *\n * Consumes 4 values from cuRAND Philox generator.\n */\n__device__ __forceinline__ uint16_t make_16bit_mask(uint64_t chunk_idx, uint64_t rng_seed,\n uint64_t rng_offset,\n uint32_t bytewise_drop_prob) {\n // Generate random bits\n curandStatePhilox4_32_10_t state;\n curand_init(rng_seed, chunk_idx, rng_offset, &state);\n const uint4 rand_bits = curand4(&state);\n\n // Compute mask\n // Note: bytewise_less_than fills MSBs (bits 7, 15, 23, 31). By\n // shifting 2 bits after every call, every other bit will be filled.\n uint32_t result = bytewise_less_than(rand_bits.x, bytewise_drop_prob);\n result = (result >> 2) | bytewise_less_than(rand_bits.y, bytewise_drop_prob);\n result = (result >> 2) | bytewise_less_than(rand_bits.z, bytewise_drop_prob);\n result = (result >> 2) | bytewise_less_than(rand_bits.w, bytewise_drop_prob);\n\n // Consolidate mask in lowest 16 bits\n result |= result >> 17;\n\n // Flip bits so 0 corresponds to drop\n result = ~result;\n\n return result;\n}\n\n// Dropout forward with FP16/BF16 input and output.\ntemplate \n__global__ void __launch_bounds__(block_size)\n dropout_kernel_fwd_f16(const T *__restrict__ input_ptr, T *__restrict__ output_ptr,\n uint8_t *__restrict__ mask_ptr,\n const uint64_t *__restrict__ rng_state_ptr, size_t num_chunks,\n uint32_t bytewise_drop_prob, float scale) {\n static_assert(sizeof(T) == 2);\n\n // Each thread processes a chunk of 16 entries\n const size_t gid = threadIdx.x + blockIdx.x * block_size;\n const size_t nthreads = gridDim.x * block_size;\n for (size_t chunk_idx = gid; chunk_idx < num_chunks; chunk_idx += nthreads) {\n // Generate dropout mask\n auto local_mask =\n make_16bit_mask(chunk_idx, rng_state_ptr[0], rng_state_ptr[1], bytewise_drop_prob);\n reinterpret_cast(mask_ptr)[chunk_idx] = local_mask;\n\n // Read input data\n using VectorType = Vector;\n VectorType local_data;\n local_data = reinterpret_cast(input_ptr)[chunk_idx];\n\n // Apply dropout based on mask\n#pragma unroll\n for (size_t i = 0; i < rng_chunk_size; i++) {\n float val = static_cast(local_data.entries[i]);\n if ((local_mask & 0x1) == 0) {\n val = 0;\n }\n val *= scale;\n local_data.entries[i] = static_cast(val);\n local_mask >>= 1;\n }\n\n // Write output data\n reinterpret_cast(output_ptr)[chunk_idx] = local_data;\n }\n}\n\n// Dropout forward with FP8 input and FP16/BF16 output.\ntemplate \n__global__ void __launch_bounds__(block_size)\n dropout_kernel_fwd_fp8(const InputType *__restrict__ input_ptr,\n const float *__restrict__ input_scale_inv_ptr,\n OutputType *__restrict__ output_ptr, uint8_t *__restrict__ mask_ptr,\n const uint64_t *__restrict__ rng_state_ptr, size_t num_chunks,\n uint32_t bytewise_drop_prob, float scale) {\n static_assert(sizeof(InputType) == 1);\n static_assert(sizeof(OutputType) == 2);\n const float input_scale_inv = *input_scale_inv_ptr;\n\n // Each thread processes a chunk of 16 entries\n const size_t gid = threadIdx.x + blockIdx.x * block_size;\n const size_t nthreads = gridDim.x * block_size;\n for (size_t chunk_idx = gid; chunk_idx < num_chunks; chunk_idx += nthreads) {\n // Generate dropout mask\n auto local_mask =\n make_16bit_mask(chunk_idx, rng_state_ptr[0], rng_state_ptr[1], bytewise_drop_prob);\n reinterpret_cast(mask_ptr)[chunk_idx] = local_mask;\n\n // Read input data\n using InputVectorType = Vector;\n InputVectorType local_input;\n local_input = reinterpret_cast(input_ptr)[chunk_idx];\n\n // Apply dropout based on mask\n using OutputVectorType = Vector;\n OutputVectorType local_output;\n#pragma unroll\n for (size_t i = 0; i < rng_chunk_size; i++) {\n float val = static_cast(local_input.entries[i]);\n val *= input_scale_inv;\n if ((local_mask & 0x1) == 0) {\n val = 0;\n }\n val *= scale;\n local_output.entries[i] = static_cast(val);\n local_mask >>= 1;\n }\n\n // Write output data\n reinterpret_cast(output_ptr)[chunk_idx] = local_output;\n }\n}\n\n// Apply dropout mask and scale.\ntemplate \n__global__ void __launch_bounds__(block_size)\n apply_dropout_mask(const T *__restrict__ input_ptr, const uint8_t *__restrict__ mask_ptr,\n T *__restrict__ output_ptr, size_t num_chunks, float scale) {\n // Each thread processes a chunk of 8 entries.\n const size_t gid = threadIdx.x + blockIdx.x * block_size;\n const size_t nthreads = gridDim.x * block_size;\n constexpr size_t chunk_size = 8;\n for (size_t chunk_idx = gid; chunk_idx < num_chunks; chunk_idx += nthreads) {\n // Read dropout mask\n uint8_t local_mask = mask_ptr[chunk_idx];\n\n // Read input data\n using VectorType = Vector;\n VectorType local_data;\n local_data = reinterpret_cast(input_ptr)[chunk_idx];\n\n // Apply dropout based on mask\n#pragma unroll\n for (size_t i = 0; i < chunk_size; i++) {\n float val = static_cast(local_data.entries[i]);\n if ((local_mask & 0x1) == 0) {\n val = 0;\n }\n val *= scale;\n local_data.entries[i] = static_cast(val);\n local_mask >>= 1;\n }\n\n // Write output data\n reinterpret_cast(output_ptr)[chunk_idx] = local_data;\n }\n}\n\n} // namespace\n\nvoid dropout_fwd(const Tensor &input, Tensor &output, Tensor &mask, Tensor &rng_state,\n float dropout_probability, cudaStream_t stream) {\n // Check tensors\n const size_t numel = input.numel();\n NVTE_CHECK(input.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor must be FP16/BF16 tensor or tensor-scaled FP8 tensor, \",\n \"but scaling mode is \", to_string(input.scaling_mode), \".\");\n NVTE_CHECK(output.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Output tensor must be FP16/BF16 tensor, \", \"but scaling mode is \",\n to_string(output.scaling_mode), \".\");\n NVTE_CHECK(mask.scaling_mode == NVTE_DELAYED_TENSOR_SCALING, \"Mask tensor must be plain tensor, \",\n \"but scaling mode is \", to_string(mask.scaling_mode), \".\");\n NVTE_CHECK(rng_state.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"RNG state tensor must be INT64 tensor with two entries, \", \"but scaling mode is \",\n to_string(rng_state.scaling_mode), \".\");\n NVTE_CHECK(output.dtype() == DType::kFloat16 || output.dtype() == DType::kBFloat16,\n \"Output tensor must be FP16/BF16 tensor, but dtype is \", to_string(output.dtype()),\n \".\");\n NVTE_CHECK(rng_state.dtype() == DType::kInt64,\n \"RNG state tensor must be INT64 tensor with two entries, but dtype is \",\n to_string(rng_state.dtype()), \".\");\n NVTE_CHECK(numel % 16 == 0,\n \"Input tensor number of elements must be divisible by 16, but shape is \",\n input.shape(), \".\");\n NVTE_CHECK(numel == output.numel(), \"Input tensor (shape=\", input.shape(),\n \") and output tensor (shape=\", output.shape(), \") do not match.\");\n NVTE_CHECK(typeToNumBits(mask.dtype()) * mask.numel() == numel, \"Mask tensor must have \", numel,\n \" bits, but found dtype=\", to_string(mask.dtype()), \" and shape=\", mask.shape(), \".\");\n NVTE_CHECK(rng_state.numel() == 2, \"RNG state tensor must be INT64 tensor with two entries, \",\n \"but shape is \", rng_state.shape(), \".\");\n NVTE_CHECK(input.data.dptr != nullptr, \"Input tensor is missing data.\");\n NVTE_CHECK(output.data.dptr != nullptr, \"Output tensor is missing data.\");\n NVTE_CHECK(mask.data.dptr != nullptr, \"Mask tensor is missing data.\");\n NVTE_CHECK(rng_state.data.dptr != nullptr, \"RNG state tensor is missing data.\");\n\n // Convert dropout probablity to scale and 8-bit representation\n NVTE_CHECK(dropout_probability >= 0 && dropout_probability < 1, \"Invalid dropout probability (\",\n dropout_probability, \").\");\n const float scale = 1 / (1 - dropout_probability);\n uint32_t bytewise_drop_prob = static_cast(std::floor(dropout_probability * 256));\n bytewise_drop_prob |= bytewise_drop_prob << 8;\n bytewise_drop_prob |= bytewise_drop_prob << 16;\n\n // CUDA config\n const size_t num_chunks = numel / rng_chunk_size;\n const size_t num_blocks = DIVUP(num_chunks, block_size);\n\n // Launch kernel depending on input dtype\n if (input.dtype() == DType::kFloat16 || input.dtype() == DType::kBFloat16) {\n NVTE_CHECK(input.dtype() == output.dtype(), \"Input tensor (dtype=\", to_string(input.dtype()),\n \") and output tensor (dtype=\", to_string(output.dtype()), \") do not match.\");\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n input.dtype(), DType,\n dropout_kernel_fwd_f16<<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(output.data.dptr),\n reinterpret_cast(mask.data.dptr),\n reinterpret_cast(rng_state.data.dptr), num_chunks, bytewise_drop_prob,\n scale););\n NVTE_CHECK_CUDA(cudaGetLastError());\n } else if (input.dtype() == DType::kFloat8E4M3 || input.dtype() == DType::kFloat8E5M2) {\n NVTE_CHECK(input.scale_inv.dptr != nullptr, \"Input tensor scale-inverse is not allocated.\");\n TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(\n input.dtype(), InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n output.dtype(), OutputType,\n dropout_kernel_fwd_fp8<<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(input.scale_inv.dptr),\n reinterpret_cast(output.data.dptr),\n reinterpret_cast(mask.data.dptr),\n reinterpret_cast(rng_state.data.dptr), num_chunks,\n bytewise_drop_prob, scale);\n\n ););\n NVTE_CHECK_CUDA(cudaGetLastError());\n } else {\n NVTE_ERROR(\"Input tensor must be FP16/BF16 tensor or tensor-scaled FP8 tensor, \",\n \"but dtype is \", to_string(input.dtype()), \".\");\n }\n}\n\nvoid dropout_bwd(const Tensor &grad_output, const Tensor &mask, Tensor &grad_input,\n float dropout_probability, cudaStream_t stream) {\n // Check tensors\n const size_t numel = grad_output.numel();\n NVTE_CHECK(grad_output.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Grad output tensor must be FP16/BF16 tensor, \", \"but scaling mode is \",\n to_string(grad_output.scaling_mode), \".\");\n NVTE_CHECK(grad_input.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Grad input tensor must be FP16/BF16 tensor, \", \"but scaling mode is \",\n to_string(grad_input.scaling_mode), \".\");\n NVTE_CHECK(mask.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Mask tensor must be a plain tensor, but scaling mode is \",\n to_string(mask.scaling_mode), \".\");\n NVTE_CHECK(grad_output.dtype() == DType::kFloat16 || grad_output.dtype() == DType::kBFloat16,\n \"Grad output tensor must be FP16/BF16 tensor, but dtype is \",\n to_string(grad_output.dtype()), \".\");\n NVTE_CHECK(grad_output.dtype() == grad_input.dtype(),\n \"Grad output tensor (dtype=\", to_string(grad_output.dtype()),\n \") and grad input tensor (dtype=\", to_string(grad_input.dtype()), \") do not match.\");\n NVTE_CHECK(numel % 16 == 0,\n \"Grad output tensor number of elements must be divisible by 16, but shape is \",\n grad_output.shape(), \".\");\n NVTE_CHECK(numel == grad_input.numel(), \"Grad output tensor (shape=\", grad_output.shape(),\n \") and grad input tensor (shape=\", grad_input.shape(), \") do not match.\");\n NVTE_CHECK(typeToNumBits(mask.dtype()) * mask.numel() == numel, \"Mask tensor must have \", numel,\n \" bits, but found dtype=\", to_string(mask.dtype()), \" and shape=\", mask.shape(), \".\");\n NVTE_CHECK(grad_output.data.dptr != nullptr, \"Grad output tensor is missing data.\");\n NVTE_CHECK(grad_input.data.dptr != nullptr, \"Grad input tensor is missing data.\");\n NVTE_CHECK(mask.data.dptr != nullptr, \"Mask tensor is missing data.\");\n\n // Convert dropout probablity to scale\n NVTE_CHECK(dropout_probability >= 0 && dropout_probability < 1, \"Invalid dropout probability (\",\n dropout_probability, \").\");\n const float scale = 1 / (1 - dropout_probability);\n\n // CUDA config\n const size_t num_chunks = numel / 8;\n const size_t num_blocks = DIVUP(num_chunks, block_size);\n\n // Launch kernel\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n grad_output.dtype(), DType,\n apply_dropout_mask<<>>(\n reinterpret_cast(grad_output.data.dptr),\n reinterpret_cast(mask.data.dptr),\n reinterpret_cast(grad_input.data.dptr), num_chunks, scale););\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n} // namespace transformer_engine\n\nvoid nvte_dropout_fwd(const NVTETensor input, NVTETensor output, NVTETensor mask,\n NVTETensor rng_state, float dropout_probability, cudaStream_t stream) {\n NVTE_API_CALL(nvte_dropout_fwd);\n using namespace transformer_engine;\n dropout_fwd(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output),\n *convertNVTETensorCheck(mask), *convertNVTETensorCheck(rng_state),\n dropout_probability, stream);\n}\n\nvoid nvte_dropout_bwd(const NVTETensor grad_output, const NVTETensor mask, NVTETensor grad_input,\n float dropout_probability, cudaStream_t stream) {\n NVTE_API_CALL(nvte_dropout_bwd);\n using namespace transformer_engine;\n dropout_bwd(*convertNVTETensorCheck(grad_output), *convertNVTETensorCheck(mask),\n *convertNVTETensorCheck(grad_input), dropout_probability, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n#include \n\n#include \"../common.h\"\n#include \"../utils.cuh\"\n#include \"transformer_engine/dropout.h\"\n\nnamespace transformer_engine {\nnamespace {\n\n// RNG kernels process chunks of 16 entries\nconstexpr size_t rng_chunk_size = 16;\n\n// CUDA block size\nconstexpr size_t block_size = 128;\n\n// Vector class to help with vectorized memory accesses\ntemplate \nunion Vector {\n using StorageType = typename BytesToType::Type;\n StorageType storage;\n T entries[kSize];\n};\n\n/* Byte-wise less-than comparison\n *\n * Results are stored in each byte's most-significant bit (MSB). All\n * other bits are zero.\n */\n__device__ __forceinline__ uint32_t bytewise_less_than(uint32_t a, uint32_t b) {\n // Compare low bits by masking MSBs and subtracting. The resulting\n // MSBs are 0 if the low bits of a are less than the low bits of b.\n uint32_t result = (a | 0x80808080) - (b & 0x7F7F7F7F);\n\n // Bitwise logical op to get answer in MSBs\n // Equivalent logic: result = (a == b) ? !result : b\n asm(\"lop3.b32 %0, %1, %2, %3, 0x4D;\\n\\t\" : \"=r\"(result) : \"r\"(a), \"r\"(b), \"r\"(result));\n\n // Mask out everything except MSBs and return\n result &= 0x80808080;\n return result;\n}\n\n/* Generate dropout mask with 16 bits.\n *\n * 1 corresponds to keep and 0 to drop.\n *\n * Consumes 4 values from cuRAND Philox generator.\n */\n__device__ __forceinline__ uint16_t make_16bit_mask(uint64_t chunk_idx, uint64_t rng_seed,\n uint64_t rng_offset,\n uint32_t bytewise_drop_prob) {\n // Generate random bits\n hiprandStatePhilox4_32_10_t state;\n hiprand_init(rng_seed, chunk_idx, rng_offset, &state);\n const uint4 rand_bits = curand4(&state);\n\n // Compute mask\n // Note: bytewise_less_than fills MSBs (bits 7, 15, 23, 31). By\n // shifting 2 bits after every call, every other bit will be filled.\n uint32_t result = bytewise_less_than(rand_bits.x, bytewise_drop_prob);\n result = (result >> 2) | bytewise_less_than(rand_bits.y, bytewise_drop_prob);\n result = (result >> 2) | bytewise_less_than(rand_bits.z, bytewise_drop_prob);\n result = (result >> 2) | bytewise_less_than(rand_bits.w, bytewise_drop_prob);\n\n // Consolidate mask in lowest 16 bits\n result |= result >> 17;\n\n // Flip bits so 0 corresponds to drop\n result = ~result;\n\n return result;\n}\n\n// Dropout forward with FP16/BF16 input and output.\ntemplate \n__global__ void __launch_bounds__(block_size)\n dropout_kernel_fwd_f16(const T *__restrict__ input_ptr, T *__restrict__ output_ptr,\n uint8_t *__restrict__ mask_ptr,\n const uint64_t *__restrict__ rng_state_ptr, size_t num_chunks,\n uint32_t bytewise_drop_prob, float scale) {\n static_assert(sizeof(T) == 2);\n\n // Each thread processes a chunk of 16 entries\n const size_t gid = threadIdx.x + blockIdx.x * block_size;\n const size_t nthreads = gridDim.x * block_size;\n for (size_t chunk_idx = gid; chunk_idx < num_chunks; chunk_idx += nthreads) {\n // Generate dropout mask\n auto local_mask =\n make_16bit_mask(chunk_idx, rng_state_ptr[0], rng_state_ptr[1], bytewise_drop_prob);\n reinterpret_cast(mask_ptr)[chunk_idx] = local_mask;\n\n // Read input data\n using VectorType = Vector;\n VectorType local_data;\n local_data = reinterpret_cast(input_ptr)[chunk_idx];\n\n // Apply dropout based on mask\n#pragma unroll\n for (size_t i = 0; i < rng_chunk_size; i++) {\n float val = static_cast(local_data.entries[i]);\n if ((local_mask & 0x1) == 0) {\n val = 0;\n }\n val *= scale;\n local_data.entries[i] = static_cast(val);\n local_mask >>= 1;\n }\n\n // Write output data\n reinterpret_cast(output_ptr)[chunk_idx] = local_data;\n }\n}\n\n// Dropout forward with FP8 input and FP16/BF16 output.\ntemplate \n__global__ void __launch_bounds__(block_size)\n dropout_kernel_fwd_fp8(const InputType *__restrict__ input_ptr,\n const float *__restrict__ input_scale_inv_ptr,\n OutputType *__restrict__ output_ptr, uint8_t *__restrict__ mask_ptr,\n const uint64_t *__restrict__ rng_state_ptr, size_t num_chunks,\n uint32_t bytewise_drop_prob, float scale) {\n static_assert(sizeof(InputType) == 1);\n static_assert(sizeof(OutputType) == 2);\n const float input_scale_inv = *input_scale_inv_ptr;\n\n // Each thread processes a chunk of 16 entries\n const size_t gid = threadIdx.x + blockIdx.x * block_size;\n const size_t nthreads = gridDim.x * block_size;\n for (size_t chunk_idx = gid; chunk_idx < num_chunks; chunk_idx += nthreads) {\n // Generate dropout mask\n auto local_mask =\n make_16bit_mask(chunk_idx, rng_state_ptr[0], rng_state_ptr[1], bytewise_drop_prob);\n reinterpret_cast(mask_ptr)[chunk_idx] = local_mask;\n\n // Read input data\n using InputVectorType = Vector;\n InputVectorType local_input;\n local_input = reinterpret_cast(input_ptr)[chunk_idx];\n\n // Apply dropout based on mask\n using OutputVectorType = Vector;\n OutputVectorType local_output;\n#pragma unroll\n for (size_t i = 0; i < rng_chunk_size; i++) {\n float val = static_cast(local_input.entries[i]);\n val *= input_scale_inv;\n if ((local_mask & 0x1) == 0) {\n val = 0;\n }\n val *= scale;\n local_output.entries[i] = static_cast(val);\n local_mask >>= 1;\n }\n\n // Write output data\n reinterpret_cast(output_ptr)[chunk_idx] = local_output;\n }\n}\n\n// Apply dropout mask and scale.\ntemplate \n__global__ void __launch_bounds__(block_size)\n apply_dropout_mask(const T *__restrict__ input_ptr, const uint8_t *__restrict__ mask_ptr,\n T *__restrict__ output_ptr, size_t num_chunks, float scale) {\n // Each thread processes a chunk of 8 entries.\n const size_t gid = threadIdx.x + blockIdx.x * block_size;\n const size_t nthreads = gridDim.x * block_size;\n constexpr size_t chunk_size = 8;\n for (size_t chunk_idx = gid; chunk_idx < num_chunks; chunk_idx += nthreads) {\n // Read dropout mask\n uint8_t local_mask = mask_ptr[chunk_idx];\n\n // Read input data\n using VectorType = Vector;\n VectorType local_data;\n local_data = reinterpret_cast(input_ptr)[chunk_idx];\n\n // Apply dropout based on mask\n#pragma unroll\n for (size_t i = 0; i < chunk_size; i++) {\n float val = static_cast(local_data.entries[i]);\n if ((local_mask & 0x1) == 0) {\n val = 0;\n }\n val *= scale;\n local_data.entries[i] = static_cast(val);\n local_mask >>= 1;\n }\n\n // Write output data\n reinterpret_cast(output_ptr)[chunk_idx] = local_data;\n }\n}\n\n} // namespace\n\nvoid dropout_fwd(const Tensor &input, Tensor &output, Tensor &mask, Tensor &rng_state,\n float dropout_probability, hipStream_t stream) {\n // Check tensors\n const size_t numel = input.numel();\n NVTE_CHECK(input.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor must be FP16/BF16 tensor or tensor-scaled FP8 tensor, \",\n \"but scaling mode is \", to_string(input.scaling_mode), \".\");\n NVTE_CHECK(output.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Output tensor must be FP16/BF16 tensor, \", \"but scaling mode is \",\n to_string(output.scaling_mode), \".\");\n NVTE_CHECK(mask.scaling_mode == NVTE_DELAYED_TENSOR_SCALING, \"Mask tensor must be plain tensor, \",\n \"but scaling mode is \", to_string(mask.scaling_mode), \".\");\n NVTE_CHECK(rng_state.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"RNG state tensor must be INT64 tensor with two entries, \", \"but scaling mode is \",\n to_string(rng_state.scaling_mode), \".\");\n NVTE_CHECK(output.dtype() == DType::kFloat16 || output.dtype() == DType::kBFloat16,\n \"Output tensor must be FP16/BF16 tensor, but dtype is \", to_string(output.dtype()),\n \".\");\n NVTE_CHECK(rng_state.dtype() == DType::kInt64,\n \"RNG state tensor must be INT64 tensor with two entries, but dtype is \",\n to_string(rng_state.dtype()), \".\");\n NVTE_CHECK(numel % 16 == 0,\n \"Input tensor number of elements must be divisible by 16, but shape is \",\n input.shape(), \".\");\n NVTE_CHECK(numel == output.numel(), \"Input tensor (shape=\", input.shape(),\n \") and output tensor (shape=\", output.shape(), \") do not match.\");\n NVTE_CHECK(typeToNumBits(mask.dtype()) * mask.numel() == numel, \"Mask tensor must have \", numel,\n \" bits, but found dtype=\", to_string(mask.dtype()), \" and shape=\", mask.shape(), \".\");\n NVTE_CHECK(rng_state.numel() == 2, \"RNG state tensor must be INT64 tensor with two entries, \",\n \"but shape is \", rng_state.shape(), \".\");\n NVTE_CHECK(input.data.dptr != nullptr, \"Input tensor is missing data.\");\n NVTE_CHECK(output.data.dptr != nullptr, \"Output tensor is missing data.\");\n NVTE_CHECK(mask.data.dptr != nullptr, \"Mask tensor is missing data.\");\n NVTE_CHECK(rng_state.data.dptr != nullptr, \"RNG state tensor is missing data.\");\n\n // Convert dropout probablity to scale and 8-bit representation\n NVTE_CHECK(dropout_probability >= 0 && dropout_probability < 1, \"Invalid dropout probability (\",\n dropout_probability, \").\");\n const float scale = 1 / (1 - dropout_probability);\n uint32_t bytewise_drop_prob = static_cast(std::floor(dropout_probability * 256));\n bytewise_drop_prob |= bytewise_drop_prob << 8;\n bytewise_drop_prob |= bytewise_drop_prob << 16;\n\n // CUDA config\n const size_t num_chunks = numel / rng_chunk_size;\n const size_t num_blocks = DIVUP(num_chunks, block_size);\n\n // Launch kernel depending on input dtype\n if (input.dtype() == DType::kFloat16 || input.dtype() == DType::kBFloat16) {\n NVTE_CHECK(input.dtype() == output.dtype(), \"Input tensor (dtype=\", to_string(input.dtype()),\n \") and output tensor (dtype=\", to_string(output.dtype()), \") do not match.\");\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n input.dtype(), DType,\n dropout_kernel_fwd_f16<<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(output.data.dptr),\n reinterpret_cast(mask.data.dptr),\n reinterpret_cast(rng_state.data.dptr), num_chunks, bytewise_drop_prob,\n scale););\n NVTE_CHECK_CUDA(hipGetLastError());\n } else if (input.dtype() == DType::kFloat8E4M3 || input.dtype() == DType::kFloat8E5M2) {\n NVTE_CHECK(input.scale_inv.dptr != nullptr, \"Input tensor scale-inverse is not allocated.\");\n TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(\n input.dtype(), InputType,\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n output.dtype(), OutputType,\n dropout_kernel_fwd_fp8<<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(input.scale_inv.dptr),\n reinterpret_cast(output.data.dptr),\n reinterpret_cast(mask.data.dptr),\n reinterpret_cast(rng_state.data.dptr), num_chunks,\n bytewise_drop_prob, scale);\n\n ););\n NVTE_CHECK_CUDA(hipGetLastError());\n } else {\n NVTE_ERROR(\"Input tensor must be FP16/BF16 tensor or tensor-scaled FP8 tensor, \",\n \"but dtype is \", to_string(input.dtype()), \".\");\n }\n}\n\nvoid dropout_bwd(const Tensor &grad_output, const Tensor &mask, Tensor &grad_input,\n float dropout_probability, hipStream_t stream) {\n // Check tensors\n const size_t numel = grad_output.numel();\n NVTE_CHECK(grad_output.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Grad output tensor must be FP16/BF16 tensor, \", \"but scaling mode is \",\n to_string(grad_output.scaling_mode), \".\");\n NVTE_CHECK(grad_input.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Grad input tensor must be FP16/BF16 tensor, \", \"but scaling mode is \",\n to_string(grad_input.scaling_mode), \".\");\n NVTE_CHECK(mask.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Mask tensor must be a plain tensor, but scaling mode is \",\n to_string(mask.scaling_mode), \".\");\n NVTE_CHECK(grad_output.dtype() == DType::kFloat16 || grad_output.dtype() == DType::kBFloat16,\n \"Grad output tensor must be FP16/BF16 tensor, but dtype is \",\n to_string(grad_output.dtype()), \".\");\n NVTE_CHECK(grad_output.dtype() == grad_input.dtype(),\n \"Grad output tensor (dtype=\", to_string(grad_output.dtype()),\n \") and grad input tensor (dtype=\", to_string(grad_input.dtype()), \") do not match.\");\n NVTE_CHECK(numel % 16 == 0,\n \"Grad output tensor number of elements must be divisible by 16, but shape is \",\n grad_output.shape(), \".\");\n NVTE_CHECK(numel == grad_input.numel(), \"Grad output tensor (shape=\", grad_output.shape(),\n \") and grad input tensor (shape=\", grad_input.shape(), \") do not match.\");\n NVTE_CHECK(typeToNumBits(mask.dtype()) * mask.numel() == numel, \"Mask tensor must have \", numel,\n \" bits, but found dtype=\", to_string(mask.dtype()), \" and shape=\", mask.shape(), \".\");\n NVTE_CHECK(grad_output.data.dptr != nullptr, \"Grad output tensor is missing data.\");\n NVTE_CHECK(grad_input.data.dptr != nullptr, \"Grad input tensor is missing data.\");\n NVTE_CHECK(mask.data.dptr != nullptr, \"Mask tensor is missing data.\");\n\n // Convert dropout probablity to scale\n NVTE_CHECK(dropout_probability >= 0 && dropout_probability < 1, \"Invalid dropout probability (\",\n dropout_probability, \").\");\n const float scale = 1 / (1 - dropout_probability);\n\n // CUDA config\n const size_t num_chunks = numel / 8;\n const size_t num_blocks = DIVUP(num_chunks, block_size);\n\n // Launch kernel\n TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(\n grad_output.dtype(), DType,\n apply_dropout_mask<<>>(\n reinterpret_cast(grad_output.data.dptr),\n reinterpret_cast(mask.data.dptr),\n reinterpret_cast(grad_input.data.dptr), num_chunks, scale););\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n} // namespace transformer_engine\n\nvoid nvte_dropout_fwd(const NVTETensor input, NVTETensor output, NVTETensor mask,\n NVTETensor rng_state, float dropout_probability, hipStream_t stream) {\n NVTE_API_CALL(nvte_dropout_fwd);\n using namespace transformer_engine;\n dropout_fwd(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output),\n *convertNVTETensorCheck(mask), *convertNVTETensorCheck(rng_state),\n dropout_probability, stream);\n}\n\nvoid nvte_dropout_bwd(const NVTETensor grad_output, const NVTETensor mask, NVTETensor grad_input,\n float dropout_probability, hipStream_t stream) {\n NVTE_API_CALL(nvte_dropout_bwd);\n using namespace transformer_engine;\n dropout_bwd(*convertNVTETensorCheck(grad_output), *convertNVTETensorCheck(mask),\n *convertNVTETensorCheck(grad_input), dropout_probability, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/dropout/dropout.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"../utils.cuh\"\n#include \"common/util/cuda_runtime.h\"\n#include \"utils.h\"\n\nnamespace transformer_engine {\nnamespace fused_router {\n\ntemplate \n__global__ void fused_moe_aux_loss_forward_kernel(const DataType* probs,\n const IndexType* tokens_per_expert,\n int total_num_tokens, int num_rows, int num_cols,\n int topk, float coeff, float* Coeff_buf) {\n // -----------------------------------------------------------------------\n // 1) Write the CPU-computed coefficient into a device buffer to re-use in BWD\n // -----------------------------------------------------------------------\n if (threadIdx.x == 0 && blockIdx.x == 0) {\n Coeff_buf[0] = coeff;\n }\n\n // -----------------------------------------------------------------------\n // 2) Each CTA computes a partial dot-product:\n // Sigma_col ( Sigma_row probs[row, col] ) * tokens_per_expert[col]\n // -----------------------------------------------------------------------\n CompType thread_sum = CompType(0);\n\n // Grid-stride over rows so that every row is processed exactly once.\n // Each thread processes a subset of columns.\n for (int col = threadIdx.x; col < num_cols; col += blockDim.x) {\n CompType col_sum = CompType(0);\n\n // Accumulate probs over the rows assigned to this CTA (grid-stride).\n for (int row = blockIdx.x; row < num_rows; row += gridDim.x) {\n col_sum += CompType(probs[row * num_cols + col]);\n }\n\n // Multiply by the token count for this expert.\n col_sum *= CompType(tokens_per_expert[col]);\n\n // Accumulate the per-column contribution into the thread-local sum.\n thread_sum += col_sum;\n }\n\n // -----------------------------------------------------------------------\n // 3) Block-level reduction of thread_sum using warp_reduce_on_shmem\n // -----------------------------------------------------------------------\n extern __shared__ float shmem[];\n CompType* shmem_block = reinterpret_cast(shmem);\n shmem_block[threadIdx.x] = thread_sum;\n __syncthreads();\n\n const int warp_id = threadIdx.x / kThreadsPerWarp;\n const int lane_id = threadIdx.x % kThreadsPerWarp;\n if (warp_id == 0) {\n CompType block_sum = warp_reduce_on_shmem(shmem_block, static_cast(blockDim.x),\n ReduceFuncType::SUM, lane_id);\n if (lane_id == 0) {\n atomicAdd(&Coeff_buf[1], static_cast(block_sum * coeff));\n }\n }\n}\n\n// Small kernel to convert the float accumulator to the output DataType.\ntemplate \n__global__ void convert_accum_to_output(const float* Coeff_buf, DataType* aux_loss) {\n aux_loss[0] = static_cast(Coeff_buf[1]);\n}\n\n/* -------------------------------------------------------------------------\n * Kernel launcher -- simplified (no cluster launch).\n * ------------------------------------------------------------------------- */\ntemplate \nvoid fused_moe_aux_loss_forward_kernel_launcher(const DataType* probs,\n const IndexType* tokens_per_expert,\n int total_num_tokens, int num_experts, int num_rows,\n int num_cols, int topk, float coeff,\n DataType* aux_loss, float* Coeff_buf,\n cudaStream_t stream) {\n NVTE_CHECK(num_experts == num_cols, \"Number of experts (\", num_experts,\n \") must be equal to number of input columns (\", num_cols, \").\");\n\n // Round up to a multiple of warp size for correct warp shuffles.\n const int block_size = ((std::min(1024, num_cols) + static_cast(kThreadsPerWarp) - 1) /\n static_cast(kThreadsPerWarp)) *\n static_cast(kThreadsPerWarp);\n const int grid_size = cuda::sm_count() * 2;\n\n // One CompType per thread in shared memory.\n const size_t smem_size = block_size * sizeof(CompType);\n check_shared_memory_capacity_num_experts(smem_size, num_experts);\n\n // Compute final coefficient and zero the float accumulator (Coeff_buf[1]) before launch.\n const float C_coeff = (num_experts * coeff) / topk / total_num_tokens / total_num_tokens;\n NVTE_CHECK_CUDA(cudaMemsetAsync(Coeff_buf + 1, 0, sizeof(float), stream));\n fused_moe_aux_loss_forward_kernel\n <<>>(probs, tokens_per_expert, total_num_tokens,\n num_rows, num_cols, topk, C_coeff, Coeff_buf);\n NVTE_CHECK_CUDA(cudaGetLastError());\n\n // Convert the float accumulator to the output DataType.\n convert_accum_to_output<<<1, 1, 0, stream>>>(Coeff_buf, aux_loss);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid fused_moe_aux_loss_forward(const Tensor& probs, const Tensor& tokens_per_expert,\n int total_num_tokens, int num_experts, int num_rows, int num_cols,\n int topk, float coeff, Tensor& aux_loss, Tensor& Coeff_buf,\n cudaStream_t stream) {\n TE_ROUTER_PROBS_TYPE_SWITCH_ALL(\n probs.data.dtype, DataType,\n TE_ROUTER_INDEX_TYPE_SWITCH_ALL(\n tokens_per_expert.data.dtype, IndexType,\n fused_moe_aux_loss_forward_kernel_launcher(\n reinterpret_cast(probs.data.dptr),\n reinterpret_cast(tokens_per_expert.data.dptr), total_num_tokens,\n num_experts, num_rows, num_cols, topk, coeff,\n reinterpret_cast(aux_loss.data.dptr),\n reinterpret_cast(Coeff_buf.data.dptr), stream);););\n}\n\ntemplate \n__global__ void fused_moe_aux_loss_backward_kernel(const float* Const_buf,\n const IndexType* tokens_per_expert, int num_rows,\n int num_cols, DataType* grad_aux_loss,\n DataType* grad_probs) {\n int global_warp_num = gridDim.x * blockDim.x / kThreadsPerWarp;\n int global_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / kThreadsPerWarp;\n int lane_id = threadIdx.x % kThreadsPerWarp;\n\n // Loop: for all positions in each row\n for (int i = lane_id; i < num_cols; i += kThreadsPerWarp) {\n float C_coeff = Const_buf[0];\n CompType tokens_per_expert_i = static_cast(tokens_per_expert[i]);\n CompType grad_aux_loss_value = static_cast(grad_aux_loss[0]);\n // Loop: for all rows\n for (int j = global_warp_id; j < num_rows; j += global_warp_num) {\n grad_probs[j * num_cols + i] = C_coeff * tokens_per_expert_i * grad_aux_loss_value;\n }\n }\n}\n\ntemplate \nvoid fused_moe_aux_loss_backward_kernel_launcher(const float* Const_buf,\n const IndexType* tokens_per_expert, int num_rows,\n int num_cols, DataType* grad_aux_loss,\n DataType* grad_probs, cudaStream_t stream) {\n // Meta data for the kernel\n int block_size = 256;\n int grid_size = (num_rows + block_size - 1) / block_size;\n fused_moe_aux_loss_backward_kernel<<>>(\n Const_buf, tokens_per_expert, num_rows, num_cols, grad_aux_loss, grad_probs);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid fused_moe_aux_loss_backward(const Tensor& Const_buf, const Tensor& tokens_per_expert,\n int num_rows, int num_cols, Tensor& grad_aux_loss,\n Tensor& grad_probs, cudaStream_t stream) {\n TE_ROUTER_PROBS_TYPE_SWITCH_ALL(\n grad_aux_loss.data.dtype, DataType,\n TE_ROUTER_INDEX_TYPE_SWITCH_ALL(\n tokens_per_expert.data.dtype, IndexType,\n fused_moe_aux_loss_backward_kernel_launcher(\n reinterpret_cast(Const_buf.data.dptr),\n reinterpret_cast(tokens_per_expert.data.dptr), num_rows, num_cols,\n reinterpret_cast(grad_aux_loss.data.dptr),\n reinterpret_cast(grad_probs.data.dptr), stream);););\n}\n\n} // namespace fused_router\n} // namespace transformer_engine\n\nvoid nvte_fused_moe_aux_loss_forward(const NVTETensor probs, const NVTETensor tokens_per_expert,\n int total_num_tokens, int num_experts, int num_rows,\n int num_cols, int topk, float coeff, NVTETensor aux_loss,\n NVTETensor Coeff_buf, cudaStream_t stream) {\n NVTE_API_CALL(nvte_fused_moe_aux_loss_forward);\n using namespace transformer_engine;\n fused_router::fused_moe_aux_loss_forward(\n *convertNVTETensorCheck(probs), *convertNVTETensorCheck(tokens_per_expert), total_num_tokens,\n num_experts, num_rows, num_cols, topk, coeff, *convertNVTETensorCheck(aux_loss),\n *convertNVTETensorCheck(Coeff_buf), stream);\n}\n\nvoid nvte_fused_moe_aux_loss_backward(const NVTETensor Const_buf,\n const NVTETensor tokens_per_expert, int num_rows,\n int num_cols, NVTETensor grad_aux_loss, NVTETensor grad_probs,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_fused_moe_aux_loss_backward);\n using namespace transformer_engine;\n fused_router::fused_moe_aux_loss_backward(*convertNVTETensorCheck(Const_buf),\n *convertNVTETensorCheck(tokens_per_expert), num_rows,\n num_cols, *convertNVTETensorCheck(grad_aux_loss),\n *convertNVTETensorCheck(grad_probs), stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"../utils.cuh\"\n#include \"common/util/hip/hip_runtime.h\"\n#include \"utils.h\"\n\nnamespace transformer_engine {\nnamespace fused_router {\n\ntemplate \n__global__ void fused_moe_aux_loss_forward_kernel(const DataType* probs,\n const IndexType* tokens_per_expert,\n int total_num_tokens, int num_rows, int num_cols,\n int topk, float coeff, float* Coeff_buf) {\n // -----------------------------------------------------------------------\n // 1) Write the CPU-computed coefficient into a device buffer to re-use in BWD\n // -----------------------------------------------------------------------\n if (threadIdx.x == 0 && blockIdx.x == 0) {\n Coeff_buf[0] = coeff;\n }\n\n // -----------------------------------------------------------------------\n // 2) Each CTA computes a partial dot-product:\n // Sigma_col ( Sigma_row probs[row, col] ) * tokens_per_expert[col]\n // -----------------------------------------------------------------------\n CompType thread_sum = CompType(0);\n\n // Grid-stride over rows so that every row is processed exactly once.\n // Each thread processes a subset of columns.\n for (int col = threadIdx.x; col < num_cols; col += blockDim.x) {\n CompType col_sum = CompType(0);\n\n // Accumulate probs over the rows assigned to this CTA (grid-stride).\n for (int row = blockIdx.x; row < num_rows; row += gridDim.x) {\n col_sum += CompType(probs[row * num_cols + col]);\n }\n\n // Multiply by the token count for this expert.\n col_sum *= CompType(tokens_per_expert[col]);\n\n // Accumulate the per-column contribution into the thread-local sum.\n thread_sum += col_sum;\n }\n\n // -----------------------------------------------------------------------\n // 3) Block-level reduction of thread_sum using warp_reduce_on_shmem\n // -----------------------------------------------------------------------\n extern __shared__ float shmem[];\n CompType* shmem_block = reinterpret_cast(shmem);\n shmem_block[threadIdx.x] = thread_sum;\n __syncthreads();\n\n const int warp_id = threadIdx.x / kThreadsPerWarp;\n const int lane_id = threadIdx.x % kThreadsPerWarp;\n if (warp_id == 0) {\n CompType block_sum = warp_reduce_on_shmem(shmem_block, static_cast(blockDim.x),\n ReduceFuncType::SUM, lane_id);\n if (lane_id == 0) {\n atomicAdd(&Coeff_buf[1], static_cast(block_sum * coeff));\n }\n }\n}\n\n// Small kernel to convert the float accumulator to the output DataType.\ntemplate \n__global__ void convert_accum_to_output(const float* Coeff_buf, DataType* aux_loss) {\n aux_loss[0] = static_cast(Coeff_buf[1]);\n}\n\n/* -------------------------------------------------------------------------\n * Kernel launcher -- simplified (no cluster launch).\n * ------------------------------------------------------------------------- */\ntemplate \nvoid fused_moe_aux_loss_forward_kernel_launcher(const DataType* probs,\n const IndexType* tokens_per_expert,\n int total_num_tokens, int num_experts, int num_rows,\n int num_cols, int topk, float coeff,\n DataType* aux_loss, float* Coeff_buf,\n hipStream_t stream) {\n NVTE_CHECK(num_experts == num_cols, \"Number of experts (\", num_experts,\n \") must be equal to number of input columns (\", num_cols, \").\");\n\n // Round up to a multiple of warp size for correct warp shuffles.\n const int block_size = ((std::min(1024, num_cols) + static_cast(kThreadsPerWarp) - 1) /\n static_cast(kThreadsPerWarp)) *\n static_cast(kThreadsPerWarp);\n const int grid_size = cuda::sm_count() * 2;\n\n // One CompType per thread in shared memory.\n const size_t smem_size = block_size * sizeof(CompType);\n check_shared_memory_capacity_num_experts(smem_size, num_experts);\n\n // Compute final coefficient and zero the float accumulator (Coeff_buf[1]) before launch.\n const float C_coeff = (num_experts * coeff) / topk / total_num_tokens / total_num_tokens;\n NVTE_CHECK_CUDA(hipMemsetAsync(Coeff_buf + 1, 0, sizeof(float), stream));\n fused_moe_aux_loss_forward_kernel\n <<>>(probs, tokens_per_expert, total_num_tokens,\n num_rows, num_cols, topk, C_coeff, Coeff_buf);\n NVTE_CHECK_CUDA(hipGetLastError());\n\n // Convert the float accumulator to the output DataType.\n convert_accum_to_output<<<1, 1, 0, stream>>>(Coeff_buf, aux_loss);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid fused_moe_aux_loss_forward(const Tensor& probs, const Tensor& tokens_per_expert,\n int total_num_tokens, int num_experts, int num_rows, int num_cols,\n int topk, float coeff, Tensor& aux_loss, Tensor& Coeff_buf,\n hipStream_t stream) {\n TE_ROUTER_PROBS_TYPE_SWITCH_ALL(\n probs.data.dtype, DataType,\n TE_ROUTER_INDEX_TYPE_SWITCH_ALL(\n tokens_per_expert.data.dtype, IndexType,\n fused_moe_aux_loss_forward_kernel_launcher(\n reinterpret_cast(probs.data.dptr),\n reinterpret_cast(tokens_per_expert.data.dptr), total_num_tokens,\n num_experts, num_rows, num_cols, topk, coeff,\n reinterpret_cast(aux_loss.data.dptr),\n reinterpret_cast(Coeff_buf.data.dptr), stream);););\n}\n\ntemplate \n__global__ void fused_moe_aux_loss_backward_kernel(const float* Const_buf,\n const IndexType* tokens_per_expert, int num_rows,\n int num_cols, DataType* grad_aux_loss,\n DataType* grad_probs) {\n int global_warp_num = gridDim.x * blockDim.x / kThreadsPerWarp;\n int global_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / kThreadsPerWarp;\n int lane_id = threadIdx.x % kThreadsPerWarp;\n\n // Loop: for all positions in each row\n for (int i = lane_id; i < num_cols; i += kThreadsPerWarp) {\n float C_coeff = Const_buf[0];\n CompType tokens_per_expert_i = static_cast(tokens_per_expert[i]);\n CompType grad_aux_loss_value = static_cast(grad_aux_loss[0]);\n // Loop: for all rows\n for (int j = global_warp_id; j < num_rows; j += global_warp_num) {\n grad_probs[j * num_cols + i] = C_coeff * tokens_per_expert_i * grad_aux_loss_value;\n }\n }\n}\n\ntemplate \nvoid fused_moe_aux_loss_backward_kernel_launcher(const float* Const_buf,\n const IndexType* tokens_per_expert, int num_rows,\n int num_cols, DataType* grad_aux_loss,\n DataType* grad_probs, hipStream_t stream) {\n // Meta data for the kernel\n int block_size = 256;\n int grid_size = (num_rows + block_size - 1) / block_size;\n fused_moe_aux_loss_backward_kernel<<>>(\n Const_buf, tokens_per_expert, num_rows, num_cols, grad_aux_loss, grad_probs);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid fused_moe_aux_loss_backward(const Tensor& Const_buf, const Tensor& tokens_per_expert,\n int num_rows, int num_cols, Tensor& grad_aux_loss,\n Tensor& grad_probs, hipStream_t stream) {\n TE_ROUTER_PROBS_TYPE_SWITCH_ALL(\n grad_aux_loss.data.dtype, DataType,\n TE_ROUTER_INDEX_TYPE_SWITCH_ALL(\n tokens_per_expert.data.dtype, IndexType,\n fused_moe_aux_loss_backward_kernel_launcher(\n reinterpret_cast(Const_buf.data.dptr),\n reinterpret_cast(tokens_per_expert.data.dptr), num_rows, num_cols,\n reinterpret_cast(grad_aux_loss.data.dptr),\n reinterpret_cast(grad_probs.data.dptr), stream);););\n}\n\n} // namespace fused_router\n} // namespace transformer_engine\n\nvoid nvte_fused_moe_aux_loss_forward(const NVTETensor probs, const NVTETensor tokens_per_expert,\n int total_num_tokens, int num_experts, int num_rows,\n int num_cols, int topk, float coeff, NVTETensor aux_loss,\n NVTETensor Coeff_buf, hipStream_t stream) {\n NVTE_API_CALL(nvte_fused_moe_aux_loss_forward);\n using namespace transformer_engine;\n fused_router::fused_moe_aux_loss_forward(\n *convertNVTETensorCheck(probs), *convertNVTETensorCheck(tokens_per_expert), total_num_tokens,\n num_experts, num_rows, num_cols, topk, coeff, *convertNVTETensorCheck(aux_loss),\n *convertNVTETensorCheck(Coeff_buf), stream);\n}\n\nvoid nvte_fused_moe_aux_loss_backward(const NVTETensor Const_buf,\n const NVTETensor tokens_per_expert, int num_rows,\n int num_cols, NVTETensor grad_aux_loss, NVTETensor grad_probs,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_fused_moe_aux_loss_backward);\n using namespace transformer_engine;\n fused_router::fused_moe_aux_loss_backward(*convertNVTETensorCheck(Const_buf),\n *convertNVTETensorCheck(tokens_per_expert), num_rows,\n num_cols, *convertNVTETensorCheck(grad_aux_loss),\n *convertNVTETensorCheck(grad_probs), stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/fused_router/fused_moe_aux_loss.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"utils.h\"\n\nnamespace transformer_engine {\nnamespace fused_router {\n\ntemplate \n__global__ void fused_topk_with_score_function_forward_kernel(\n const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax,\n int num_groups, int group_topk, float scaling_factor, int score_function,\n const BiasType *expert_bias, DataType *probs, bool *routing_map,\n CompType *intermediate_output) {\n /***\n * Section: Global Variables/Addresses init\n * - Each warp is responsible for one token, and has own shared memory buffer.\n * Then __syncwarp() is used instead of __syncthreads()\n */\n // Used variables/addresses init\n int num_token_per_block = blockDim.x / kThreadsPerWarp;\n int warp_id = threadIdx.x / kThreadsPerWarp;\n int lane_id = threadIdx.x % kThreadsPerWarp;\n extern __shared__ float shmem[];\n CompType *scores_buf = reinterpret_cast(shmem);\n CompType *topk_scores_buf = scores_buf + num_experts * num_token_per_block;\n CompType *group_scores_buf = nullptr, *masked_scores_buf = nullptr;\n int *topk_indices_buf = nullptr;\n if (group_topk > 0) {\n masked_scores_buf = topk_scores_buf + topk * num_token_per_block;\n group_scores_buf = masked_scores_buf + num_experts * num_token_per_block;\n topk_indices_buf = reinterpret_cast(group_scores_buf + num_groups * num_token_per_block);\n } else {\n topk_indices_buf = reinterpret_cast(topk_scores_buf + topk * num_token_per_block);\n }\n // The address of buffers on the current warp\n CompType *scores = scores_buf + warp_id * num_experts;\n CompType *topk_scores = topk_scores_buf + warp_id * topk;\n CompType *masked_scores = masked_scores_buf + warp_id * num_experts;\n CompType *group_scores = group_scores_buf + warp_id * num_groups;\n int *topk_indices = topk_indices_buf + warp_id * topk;\n\n /***\n * Section: Main Loop\n * - Each warp is responsible for one token\n */\n int total_round = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n for (int round = blockIdx.x; round < total_round; round += gridDim.x) {\n int token_offset_cur_warp = round * num_token_per_block + warp_id;\n // Each warp is responsible for one token\n if (token_offset_cur_warp >= num_tokens) break;\n\n /***\n * Section: Init buffer\n * - Clear the global buffer which will accept the result of this round\n * - Clear/Init the shmem buffer used by current warp this round\n * - Load the logits to shmem\n */\n int pos_offset = token_offset_cur_warp * num_experts;\n // Clear the probs/routing_map (num_experts)\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n probs[pos_offset + i] = 0.0;\n routing_map[pos_offset + i] = false;\n if (score_function == 1) {\n intermediate_output[pos_offset + i] = -std::numeric_limits::infinity();\n }\n }\n // Load the logits to shmem\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n scores[i] = logits[pos_offset + i];\n }\n // If group_topk > 0, init the masked_scores to -inf\n if (group_topk > 0) {\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n masked_scores[i] = -std::numeric_limits::infinity();\n }\n }\n __threadfence_block();\n __syncwarp();\n\n /***\n * Section: Preprocess\n * Possible preprocess the scores before the topk operation\n * - Pre-softmax\n * - Sigmoid\n * - Sqrtsoftplus\n * - Expert bias\n * This is in-place scores update\n */\n if (use_pre_softmax && score_function == 1) { // score_function == 1 means softmax\n // Apply softmax to the logits before the topk\n apply_softmax_on_float(scores, num_experts, lane_id);\n __syncwarp();\n // Save the softmax output for backward\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n intermediate_output[pos_offset + i] = scores[i];\n }\n } else if (score_function == 0) { // score_function == 0 means sigmoid\n // Apply sigmoid to the logits\n apply_sigmoid_on_float(scores, num_experts, lane_id);\n __syncwarp();\n // Save the sigmoid output for backward\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n intermediate_output[pos_offset + i] = scores[i];\n }\n } else if (score_function == 2) { // score_function == 2 means sqrtsoftplus\n // First save the original logits for backward (needed for sqrtsoftplus gradient computation)\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n intermediate_output[pos_offset + i] = scores[i]; // Save original logits\n }\n __syncwarp();\n // Apply sqrtsoftplus to the logits\n apply_sqrtsoftplus_on_float(scores, num_experts, lane_id);\n }\n\n __syncwarp(); //Confirm the scores is written to the output\n\n // Expert bias is only used at the sigmoid/sqrtsoftplus case\n if (expert_bias && (score_function == 0 || score_function == 2)) {\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n scores[i] += static_cast(expert_bias[i]);\n }\n __syncwarp();\n }\n\n /***\n * Section: Topk\n * Get the topk indices\n * - group_topk\n * - naive topk\n * - topk with expert bias\n */\n // Topk on the scores\n // The bias being not empty happens at the sigmoid/sqrtsoftplus case\n if (group_topk > 0) {\n int group_size = num_experts / num_groups;\n // Top2\n for (int i = 0; i < num_groups; i++) {\n topk_and_mask(\n /*scores ptr = */ scores + i * group_size,\n /*data size = */ group_size,\n /*topk = */ topk / group_topk,\n /*topk indices ptr = */ topk_indices,\n /*topk scores ptr = */ topk_scores,\n /*lane id = */ lane_id);\n __syncwarp();\n // Compute the group score\n if (lane_id == 0) {\n CompType tmp = 0.0;\n for (int j = 0; j < topk / group_topk; j++) {\n tmp = tmp + topk_scores[j];\n }\n group_scores[i] = tmp;\n }\n __syncwarp();\n }\n\n // select the topk groups\n topk_and_mask(\n /*scores ptr = */ group_scores,\n /*data size = */ num_groups,\n /*topk = */ group_topk,\n /*topk indices ptr = */ topk_indices,\n /*topk scores ptr = */ topk_scores,\n /*lane id = */ lane_id);\n __syncwarp();\n // Copy the unmasked scores to the buffer\n for (int i = 0; i < group_topk; i++) {\n int st = topk_indices[i] * group_size;\n int ed = st + group_size;\n for (int j = st + lane_id; j < ed; j += kThreadsPerWarp) {\n masked_scores[j] = scores[j];\n }\n }\n __syncwarp();\n topk_and_mask(masked_scores, num_experts, topk, topk_indices, topk_scores, lane_id);\n\n } else {\n topk_and_mask(scores, num_experts, topk, topk_indices, topk_scores, lane_id);\n }\n __syncwarp();\n\n /***\n * Section: Postprocess\n * Possible postprocess the scores after the topk operation\n * - Revert Expert bias\n * - Softmax\n * - Sigmoid/Sqrtsoftplus post-processing when topk > 1\n * - Write the result with scaling_factor\n */\n // Revert Expert bias from the topk scores\n if (expert_bias && (score_function == 0 || score_function == 2)) {\n for (int i = lane_id; i < topk; i += kThreadsPerWarp) {\n topk_scores[i] = topk_scores[i] - static_cast(expert_bias[topk_indices[i]]);\n }\n __syncwarp();\n }\n\n // score_function == 1 means softmax\n if (!use_pre_softmax && score_function == 1) {\n // Apply softmax to the topk logits\n apply_softmax_on_float(topk_scores, topk, lane_id);\n __syncwarp();\n // Save the softmax output for backward\n for (int i = lane_id; i < topk; i += kThreadsPerWarp) {\n intermediate_output[pos_offset + topk_indices[i]] = topk_scores[i];\n }\n __syncwarp();\n }\n\n // Sigmoid/Sqrtsoftplus post-processing when topk > 1\n if (score_function == 0 || score_function == 2) {\n if (topk > 1) {\n CompType sum_scores = warp_reduce_on_shmem(topk_scores, topk, ReduceFuncType::SUM, lane_id);\n for (int i = lane_id; i < topk; i += kThreadsPerWarp) {\n topk_scores[i] = topk_scores[i] / (sum_scores + epsilon);\n }\n }\n __syncwarp();\n }\n\n // Write the probs/routing_map to the output tensor\n for (int i = lane_id; i < topk; i += kThreadsPerWarp) {\n routing_map[pos_offset + topk_indices[i]] = true;\n probs[pos_offset + topk_indices[i]] = scaling_factor * topk_scores[i];\n }\n __threadfence_block();\n __syncwarp();\n }\n}\n\ntemplate \nvoid fused_topk_with_score_function_forward_kernel_launcher(\n const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax,\n int num_groups, int group_topk, float scaling_factor, int score_function,\n const BiasType *expert_bias, DataType *probs, bool *routing_map, CompType *intermediate_output,\n cudaStream_t stream) {\n size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp;\n size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // scores\n + topk * num_token_per_block * sizeof(CompType) // topk_scores\n + topk * num_token_per_block * sizeof(int); // topk_indices\n if (group_topk > 0) {\n shared_memory_size += num_groups * num_token_per_block * sizeof(CompType); // group_scores\n shared_memory_size += num_experts * num_token_per_block * sizeof(CompType); // maksed_scores\n }\n check_shared_memory_capacity_num_experts(shared_memory_size, num_experts);\n // Radix selection is O(E), independent of K, but it needs 4 passes for 32-bit float;\n // switch at K=16 where naive O(K^2*E) starts to dominate\n if (topk < 16) {\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n fused_topk_with_score_function_forward_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size));\n fused_topk_with_score_function_forward_kernel\n <<>>(\n logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk,\n scaling_factor, score_function, expert_bias, probs, routing_map, intermediate_output);\n } else {\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n fused_topk_with_score_function_forward_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size));\n fused_topk_with_score_function_forward_kernel\n <<>>(\n logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk,\n scaling_factor, score_function, expert_bias, probs, routing_map, intermediate_output);\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid fused_topk_with_score_function_forward(const Tensor logits, int num_tokens, int num_experts,\n int topk, bool use_pre_softmax, int num_groups,\n int group_topk, float scaling_factor,\n int score_function, const Tensor expert_bias,\n Tensor probs, Tensor routing_map,\n Tensor intermediate_output, cudaStream_t stream) {\n TE_ROUTER_PROBS_TYPE_SWITCH_ALL(\n logits.data.dtype, DataType,\n TE_ROUTER_PROBS_TYPE_SWITCH_ALL(\n expert_bias.data.dtype, BiasType,\n fused_topk_with_score_function_forward_kernel_launcher(\n reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk,\n use_pre_softmax, num_groups, group_topk, scaling_factor, score_function,\n reinterpret_cast(expert_bias.data.dptr),\n reinterpret_cast(probs.data.dptr),\n reinterpret_cast(routing_map.data.dptr),\n reinterpret_cast(intermediate_output.data.dptr), stream);););\n}\n\ntemplate \n__global__ void fused_topk_with_score_function_backward_kernel(\n // Inputs tensor\n const bool *routing_map, const CompType *intermediate_output, const DataType *grad_probs,\n // Other parameters\n int num_tokens, int num_experts, int topk, bool use_pre_softmax, float scaling_factor,\n int score_function,\n // Output tensor\n DataType *grad_logits) {\n /***\n * Section: Global Variables/Addresses init\n * - Each warp is responsible for one token, and has own shared memory buffer.\n * Then __syncwarp() is used instead of __syncthreads()\n */\n // Used variables/addresses init\n int num_token_per_block = blockDim.x / kThreadsPerWarp;\n int warp_id = threadIdx.x / kThreadsPerWarp;\n int lane_id = threadIdx.x % kThreadsPerWarp;\n extern __shared__ float shmem[];\n CompType *grad_probs_buf = reinterpret_cast(shmem);\n // To store the output of softmax/sigmoid from fwd, or original logits for sqrtsoftplus\n CompType *act_from_fwd_buf = grad_probs_buf + num_experts * num_token_per_block;\n CompType *comp_buf = act_from_fwd_buf + num_experts * num_token_per_block;\n // To store the routing_map from the fwd\n bool *routing_map_buf = reinterpret_cast(comp_buf + num_experts * num_token_per_block);\n // The address of buffers on the current warp\n CompType *local_grad = grad_probs_buf + warp_id * num_experts;\n CompType *local_act_from_fwd = act_from_fwd_buf + warp_id * num_experts;\n CompType *local_comp_buf = comp_buf + warp_id * num_experts;\n bool *local_routing_map = routing_map_buf + warp_id * num_experts;\n\n /***\n * Section: Main Loop\n * - Each warp is responsible for one token\n */\n int total_round = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n for (int round = blockIdx.x; round < total_round; round += gridDim.x) {\n int token_offset_cur_warp = round * num_token_per_block + warp_id;\n // Each warp is responsible for one token\n if (token_offset_cur_warp >= num_tokens) break;\n\n /***\n * Section: Init buffer\n * - Clear the global buffer which will accept the result of this round\n * - Clear/Init the shmem buffer used by current warp this round\n * - Load the dgrad/output_from_fwd to shmem\n */\n int pos_offset = token_offset_cur_warp * num_experts;\n // Load the dgrad/output_from_fwd to shmem\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_grad[i] = grad_probs[pos_offset + i];\n local_act_from_fwd[i] = intermediate_output[pos_offset + i];\n local_routing_map[i] = routing_map[pos_offset + i];\n }\n __threadfence_block();\n __syncwarp();\n\n /***\n * Section: Backward of ops after the topk\n * - Backward of the used scaling_factor\n * - Sigmoid/Sqrtsoftplus Post-processing bwd when topk > 1\n * - Softmax bwd if use_pre_softmax is false\n */\n // Backward of the used scaling_factor\n // In-place update\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n if (local_routing_map[i]) {\n local_grad[i] = local_grad[i] * scaling_factor;\n }\n }\n __syncwarp();\n\n // Sqrtsoftplus: First compute sqrtsoftplus output from original logits\n // (needed for both post-processing bwd and activation bwd, compute once here)\n // For sqrtsoftplus, intermediate_output stores original logits\n if (score_function == 2) {\n // Copy original logits to local_comp_buf and apply sqrtsoftplus in-place\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_comp_buf[i] = local_act_from_fwd[i];\n }\n __syncwarp();\n apply_sqrtsoftplus_on_float(local_comp_buf, num_experts, lane_id);\n __syncwarp();\n }\n\n // Sigmoid/Sqrtsoftplus Post-processing bwd when topk > 1 (normalization backward)\n if (topk > 1 && (score_function == 0 || score_function == 2)) {\n // Select the correct activation output buffer:\n // - Sigmoid: local_act_from_fwd already contains sigmoid output\n // - Sqrtsoftplus: local_comp_buf contains sqrtsoftplus output computed above\n CompType *act_output = (score_function == 0) ? local_act_from_fwd : local_comp_buf;\n\n CompType sum_fwd_input = masked_warp_reduce_on_shmem(\n /*data ptr = */ act_output,\n /*mask ptr = */ local_routing_map,\n /*data size = */ num_experts,\n /*reduce func = */ ReduceFuncType::SUM, lane_id);\n // Compute sum of output * grad using registers\n CompType local_sum_Output_x_Grad = 0.0;\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n if (local_routing_map[i]) {\n local_sum_Output_x_Grad += local_grad[i] * act_output[i];\n }\n }\n // Warp reduce the sum\n for (int s = 16; s > 0; s /= 2) {\n local_sum_Output_x_Grad += __shfl_xor_sync(0xffffffff, local_sum_Output_x_Grad, s);\n }\n CompType sum_Output_x_Grad = local_sum_Output_x_Grad;\n // In-place update\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n if (local_routing_map[i]) {\n local_grad[i] =\n local_grad[i] / (sum_fwd_input + epsilon) -\n sum_Output_x_Grad / ((sum_fwd_input + epsilon) * (sum_fwd_input + epsilon));\n } else {\n local_grad[i] = 0.0;\n }\n }\n __syncwarp();\n }\n\n // Softmax bwd if use_pre_softmax is false\n if (!use_pre_softmax && score_function == 1) {\n apply_softmax_bwd_on_float(local_grad, local_act_from_fwd, local_comp_buf, local_routing_map,\n num_experts, lane_id);\n __syncwarp();\n }\n\n /***\n * Section: Backward of topk\n * mask the unselected position in the grad\n */\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n if (!local_routing_map[i]) {\n local_grad[i] = 0.0;\n }\n }\n __syncwarp();\n\n /***\n * Section: Backward of ops before the topk\n * - Pre-softmax bwd\n * - Sigmoid bwd\n * - Sqrtsoftplus bwd\n * - Write the grad_logits to the global mem\n */\n // Pre-softmax bwd\n if (score_function == 1 && use_pre_softmax) {\n apply_softmax_bwd_on_float(local_grad, local_act_from_fwd, local_comp_buf, nullptr,\n num_experts, lane_id);\n __syncwarp();\n }\n // Sigmoid bwd\n if (score_function == 0) {\n apply_sigmoid_bwd_on_float(local_grad, local_act_from_fwd, num_experts, lane_id);\n __syncwarp();\n }\n // Sqrtsoftplus bwd\n // For sqrtsoftplus, local_comp_buf already contains sqrtsoftplus output computed earlier\n // Now compute gradient: dy/dx = sigmoid(x) / (2 * y)\n if (score_function == 2) {\n apply_sqrtsoftplus_bwd_on_float(local_grad, local_comp_buf, local_act_from_fwd, num_experts,\n lane_id);\n __syncwarp();\n }\n // Write the grad_logits to the global mem\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n grad_logits[pos_offset + i] = local_grad[i];\n }\n __syncwarp();\n }\n}\n\ntemplate \nvoid fused_topk_with_score_function_backward_kernel_launcher(\n const bool *routing_map, const CompType *intermediate_output, const DataType *grad_probs,\n int num_tokens, int num_experts, int topk, bool use_pre_softmax, float scaling_factor,\n int score_function, DataType *grad_logits, cudaStream_t stream) {\n // Meta data for the kernel\n size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp;\n size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // grad_probs\n +\n num_experts * num_token_per_block * sizeof(CompType) // act_from_fwd\n + num_experts * num_token_per_block * sizeof(CompType) // comp_buf\n + num_experts * num_token_per_block * sizeof(bool); // routing_map\n check_shared_memory_capacity_num_experts(shared_memory_size, num_experts);\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(fused_topk_with_score_function_backward_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize,\n shared_memory_size));\n fused_topk_with_score_function_backward_kernel\n <<>>(\n routing_map, intermediate_output, grad_probs, num_tokens, num_experts, topk,\n use_pre_softmax, scaling_factor, score_function, grad_logits);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid fused_topk_with_score_function_backward(const Tensor &routing_map,\n const Tensor &intermediate_output,\n const Tensor &grad_probs, int num_tokens,\n int num_experts, int topk, bool use_pre_softmax,\n float scaling_factor, int score_function,\n Tensor &grad_logits, cudaStream_t stream) {\n TE_ROUTER_PROBS_TYPE_SWITCH_ALL(\n grad_logits.data.dtype, DataType,\n fused_topk_with_score_function_backward_kernel_launcher(\n reinterpret_cast(routing_map.data.dptr),\n reinterpret_cast(intermediate_output.data.dptr),\n reinterpret_cast(grad_probs.data.dptr), num_tokens, num_experts, topk,\n use_pre_softmax, scaling_factor, score_function,\n reinterpret_cast(grad_logits.data.dptr), stream););\n}\n\n} // namespace fused_router\n} // namespace transformer_engine\n\nvoid nvte_fused_topk_with_score_function_forward(\n const NVTETensor logits, int num_tokens, int num_experts, int topk, int use_pre_softmax,\n int num_groups, int group_topk, float scaling_factor, int score_function,\n const NVTETensor expert_bias, NVTETensor probs, NVTETensor routing_map,\n NVTETensor intermediate_output, cudaStream_t stream) {\n NVTE_API_CALL(nvte_fused_topk_with_score_function_forward);\n using namespace transformer_engine;\n fused_router::fused_topk_with_score_function_forward(\n *convertNVTETensorCheck(logits), num_tokens, num_experts, topk,\n static_cast(use_pre_softmax), num_groups, group_topk, scaling_factor, score_function,\n *convertNVTETensorCheck(expert_bias), *convertNVTETensorCheck(probs),\n *convertNVTETensorCheck(routing_map), *convertNVTETensorCheck(intermediate_output), stream);\n}\n\nvoid nvte_fused_topk_with_score_function_backward(const NVTETensor routing_map,\n const NVTETensor intermediate_output,\n const NVTETensor grad_probs, int num_tokens,\n int num_experts, int topk, int use_pre_softmax,\n float scaling_factor, int score_function,\n NVTETensor grad_logits, cudaStream_t stream) {\n NVTE_API_CALL(nvte_fused_topk_with_score_function_backward);\n using namespace transformer_engine;\n fused_router::fused_topk_with_score_function_backward(\n *convertNVTETensorCheck(routing_map), *convertNVTETensorCheck(intermediate_output),\n *convertNVTETensorCheck(grad_probs), num_tokens, num_experts, topk,\n static_cast(use_pre_softmax), scaling_factor, score_function,\n *convertNVTETensorCheck(grad_logits), stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"utils.h\"\n\nnamespace transformer_engine {\nnamespace fused_router {\n\ntemplate \n__global__ void fused_topk_with_score_function_forward_kernel(\n const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax,\n int num_groups, int group_topk, float scaling_factor, int score_function,\n const BiasType *expert_bias, DataType *probs, bool *routing_map,\n CompType *intermediate_output) {\n /***\n * Section: Global Variables/Addresses init\n * - Each warp is responsible for one token, and has own shared memory buffer.\n * Then __syncwarp() is used instead of __syncthreads()\n */\n // Used variables/addresses init\n int num_token_per_block = blockDim.x / kThreadsPerWarp;\n int warp_id = threadIdx.x / kThreadsPerWarp;\n int lane_id = threadIdx.x % kThreadsPerWarp;\n extern __shared__ float shmem[];\n CompType *scores_buf = reinterpret_cast(shmem);\n CompType *topk_scores_buf = scores_buf + num_experts * num_token_per_block;\n CompType *group_scores_buf = nullptr, *masked_scores_buf = nullptr;\n int *topk_indices_buf = nullptr;\n if (group_topk > 0) {\n masked_scores_buf = topk_scores_buf + topk * num_token_per_block;\n group_scores_buf = masked_scores_buf + num_experts * num_token_per_block;\n topk_indices_buf = reinterpret_cast(group_scores_buf + num_groups * num_token_per_block);\n } else {\n topk_indices_buf = reinterpret_cast(topk_scores_buf + topk * num_token_per_block);\n }\n // The address of buffers on the current warp\n CompType *scores = scores_buf + warp_id * num_experts;\n CompType *topk_scores = topk_scores_buf + warp_id * topk;\n CompType *masked_scores = masked_scores_buf + warp_id * num_experts;\n CompType *group_scores = group_scores_buf + warp_id * num_groups;\n int *topk_indices = topk_indices_buf + warp_id * topk;\n\n /***\n * Section: Main Loop\n * - Each warp is responsible for one token\n */\n int total_round = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n for (int round = blockIdx.x; round < total_round; round += gridDim.x) {\n int token_offset_cur_warp = round * num_token_per_block + warp_id;\n // Each warp is responsible for one token\n if (token_offset_cur_warp >= num_tokens) break;\n\n /***\n * Section: Init buffer\n * - Clear the global buffer which will accept the result of this round\n * - Clear/Init the shmem buffer used by current warp this round\n * - Load the logits to shmem\n */\n int pos_offset = token_offset_cur_warp * num_experts;\n // Clear the probs/routing_map (num_experts)\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n probs[pos_offset + i] = 0.0;\n routing_map[pos_offset + i] = false;\n if (score_function == 1) {\n intermediate_output[pos_offset + i] = -std::numeric_limits::infinity();\n }\n }\n // Load the logits to shmem\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n scores[i] = logits[pos_offset + i];\n }\n // If group_topk > 0, init the masked_scores to -inf\n if (group_topk > 0) {\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n masked_scores[i] = -std::numeric_limits::infinity();\n }\n }\n __threadfence_block();\n __syncwarp();\n\n /***\n * Section: Preprocess\n * Possible preprocess the scores before the topk operation\n * - Pre-softmax\n * - Sigmoid\n * - Sqrtsoftplus\n * - Expert bias\n * This is in-place scores update\n */\n if (use_pre_softmax && score_function == 1) { // score_function == 1 means softmax\n // Apply softmax to the logits before the topk\n apply_softmax_on_float(scores, num_experts, lane_id);\n __syncwarp();\n // Save the softmax output for backward\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n intermediate_output[pos_offset + i] = scores[i];\n }\n } else if (score_function == 0) { // score_function == 0 means sigmoid\n // Apply sigmoid to the logits\n apply_sigmoid_on_float(scores, num_experts, lane_id);\n __syncwarp();\n // Save the sigmoid output for backward\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n intermediate_output[pos_offset + i] = scores[i];\n }\n } else if (score_function == 2) { // score_function == 2 means sqrtsoftplus\n // First save the original logits for backward (needed for sqrtsoftplus gradient computation)\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n intermediate_output[pos_offset + i] = scores[i]; // Save original logits\n }\n __syncwarp();\n // Apply sqrtsoftplus to the logits\n apply_sqrtsoftplus_on_float(scores, num_experts, lane_id);\n }\n\n __syncwarp(); //Confirm the scores is written to the output\n\n // Expert bias is only used at the sigmoid/sqrtsoftplus case\n if (expert_bias && (score_function == 0 || score_function == 2)) {\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n scores[i] += static_cast(expert_bias[i]);\n }\n __syncwarp();\n }\n\n /***\n * Section: Topk\n * Get the topk indices\n * - group_topk\n * - naive topk\n * - topk with expert bias\n */\n // Topk on the scores\n // The bias being not empty happens at the sigmoid/sqrtsoftplus case\n if (group_topk > 0) {\n int group_size = num_experts / num_groups;\n // Top2\n for (int i = 0; i < num_groups; i++) {\n topk_and_mask(\n /*scores ptr = */ scores + i * group_size,\n /*data size = */ group_size,\n /*topk = */ topk / group_topk,\n /*topk indices ptr = */ topk_indices,\n /*topk scores ptr = */ topk_scores,\n /*lane id = */ lane_id);\n __syncwarp();\n // Compute the group score\n if (lane_id == 0) {\n CompType tmp = 0.0;\n for (int j = 0; j < topk / group_topk; j++) {\n tmp = tmp + topk_scores[j];\n }\n group_scores[i] = tmp;\n }\n __syncwarp();\n }\n\n // select the topk groups\n topk_and_mask(\n /*scores ptr = */ group_scores,\n /*data size = */ num_groups,\n /*topk = */ group_topk,\n /*topk indices ptr = */ topk_indices,\n /*topk scores ptr = */ topk_scores,\n /*lane id = */ lane_id);\n __syncwarp();\n // Copy the unmasked scores to the buffer\n for (int i = 0; i < group_topk; i++) {\n int st = topk_indices[i] * group_size;\n int ed = st + group_size;\n for (int j = st + lane_id; j < ed; j += kThreadsPerWarp) {\n masked_scores[j] = scores[j];\n }\n }\n __syncwarp();\n topk_and_mask(masked_scores, num_experts, topk, topk_indices, topk_scores, lane_id);\n\n } else {\n topk_and_mask(scores, num_experts, topk, topk_indices, topk_scores, lane_id);\n }\n __syncwarp();\n\n /***\n * Section: Postprocess\n * Possible postprocess the scores after the topk operation\n * - Revert Expert bias\n * - Softmax\n * - Sigmoid/Sqrtsoftplus post-processing when topk > 1\n * - Write the result with scaling_factor\n */\n // Revert Expert bias from the topk scores\n if (expert_bias && (score_function == 0 || score_function == 2)) {\n for (int i = lane_id; i < topk; i += kThreadsPerWarp) {\n topk_scores[i] = topk_scores[i] - static_cast(expert_bias[topk_indices[i]]);\n }\n __syncwarp();\n }\n\n // score_function == 1 means softmax\n if (!use_pre_softmax && score_function == 1) {\n // Apply softmax to the topk logits\n apply_softmax_on_float(topk_scores, topk, lane_id);\n __syncwarp();\n // Save the softmax output for backward\n for (int i = lane_id; i < topk; i += kThreadsPerWarp) {\n intermediate_output[pos_offset + topk_indices[i]] = topk_scores[i];\n }\n __syncwarp();\n }\n\n // Sigmoid/Sqrtsoftplus post-processing when topk > 1\n if (score_function == 0 || score_function == 2) {\n if (topk > 1) {\n CompType sum_scores = warp_reduce_on_shmem(topk_scores, topk, ReduceFuncType::SUM, lane_id);\n for (int i = lane_id; i < topk; i += kThreadsPerWarp) {\n topk_scores[i] = topk_scores[i] / (sum_scores + epsilon);\n }\n }\n __syncwarp();\n }\n\n // Write the probs/routing_map to the output tensor\n for (int i = lane_id; i < topk; i += kThreadsPerWarp) {\n routing_map[pos_offset + topk_indices[i]] = true;\n probs[pos_offset + topk_indices[i]] = scaling_factor * topk_scores[i];\n }\n __threadfence_block();\n __syncwarp();\n }\n}\n\ntemplate \nvoid fused_topk_with_score_function_forward_kernel_launcher(\n const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax,\n int num_groups, int group_topk, float scaling_factor, int score_function,\n const BiasType *expert_bias, DataType *probs, bool *routing_map, CompType *intermediate_output,\n hipStream_t stream) {\n size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp;\n size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // scores\n + topk * num_token_per_block * sizeof(CompType) // topk_scores\n + topk * num_token_per_block * sizeof(int); // topk_indices\n if (group_topk > 0) {\n shared_memory_size += num_groups * num_token_per_block * sizeof(CompType); // group_scores\n shared_memory_size += num_experts * num_token_per_block * sizeof(CompType); // maksed_scores\n }\n check_shared_memory_capacity_num_experts(shared_memory_size, num_experts);\n // Radix selection is O(E), independent of K, but it needs 4 passes for 32-bit float;\n // switch at K=16 where naive O(K^2*E) starts to dominate\n if (topk < 16) {\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(fused_topk_with_score_function_forward_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size));\n fused_topk_with_score_function_forward_kernel\n <<>>(\n logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk,\n scaling_factor, score_function, expert_bias, probs, routing_map, intermediate_output);\n } else {\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(fused_topk_with_score_function_forward_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size));\n fused_topk_with_score_function_forward_kernel\n <<>>(\n logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk,\n scaling_factor, score_function, expert_bias, probs, routing_map, intermediate_output);\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid fused_topk_with_score_function_forward(const Tensor logits, int num_tokens, int num_experts,\n int topk, bool use_pre_softmax, int num_groups,\n int group_topk, float scaling_factor,\n int score_function, const Tensor expert_bias,\n Tensor probs, Tensor routing_map,\n Tensor intermediate_output, hipStream_t stream) {\n TE_ROUTER_PROBS_TYPE_SWITCH_ALL(\n logits.data.dtype, DataType,\n TE_ROUTER_PROBS_TYPE_SWITCH_ALL(\n expert_bias.data.dtype, BiasType,\n fused_topk_with_score_function_forward_kernel_launcher(\n reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk,\n use_pre_softmax, num_groups, group_topk, scaling_factor, score_function,\n reinterpret_cast(expert_bias.data.dptr),\n reinterpret_cast(probs.data.dptr),\n reinterpret_cast(routing_map.data.dptr),\n reinterpret_cast(intermediate_output.data.dptr), stream);););\n}\n\ntemplate \n__global__ void fused_topk_with_score_function_backward_kernel(\n // Inputs tensor\n const bool *routing_map, const CompType *intermediate_output, const DataType *grad_probs,\n // Other parameters\n int num_tokens, int num_experts, int topk, bool use_pre_softmax, float scaling_factor,\n int score_function,\n // Output tensor\n DataType *grad_logits) {\n /***\n * Section: Global Variables/Addresses init\n * - Each warp is responsible for one token, and has own shared memory buffer.\n * Then __syncwarp() is used instead of __syncthreads()\n */\n // Used variables/addresses init\n int num_token_per_block = blockDim.x / kThreadsPerWarp;\n int warp_id = threadIdx.x / kThreadsPerWarp;\n int lane_id = threadIdx.x % kThreadsPerWarp;\n extern __shared__ float shmem[];\n CompType *grad_probs_buf = reinterpret_cast(shmem);\n // To store the output of softmax/sigmoid from fwd, or original logits for sqrtsoftplus\n CompType *act_from_fwd_buf = grad_probs_buf + num_experts * num_token_per_block;\n CompType *comp_buf = act_from_fwd_buf + num_experts * num_token_per_block;\n // To store the routing_map from the fwd\n bool *routing_map_buf = reinterpret_cast(comp_buf + num_experts * num_token_per_block);\n // The address of buffers on the current warp\n CompType *local_grad = grad_probs_buf + warp_id * num_experts;\n CompType *local_act_from_fwd = act_from_fwd_buf + warp_id * num_experts;\n CompType *local_comp_buf = comp_buf + warp_id * num_experts;\n bool *local_routing_map = routing_map_buf + warp_id * num_experts;\n\n /***\n * Section: Main Loop\n * - Each warp is responsible for one token\n */\n int total_round = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n for (int round = blockIdx.x; round < total_round; round += gridDim.x) {\n int token_offset_cur_warp = round * num_token_per_block + warp_id;\n // Each warp is responsible for one token\n if (token_offset_cur_warp >= num_tokens) break;\n\n /***\n * Section: Init buffer\n * - Clear the global buffer which will accept the result of this round\n * - Clear/Init the shmem buffer used by current warp this round\n * - Load the dgrad/output_from_fwd to shmem\n */\n int pos_offset = token_offset_cur_warp * num_experts;\n // Load the dgrad/output_from_fwd to shmem\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_grad[i] = grad_probs[pos_offset + i];\n local_act_from_fwd[i] = intermediate_output[pos_offset + i];\n local_routing_map[i] = routing_map[pos_offset + i];\n }\n __threadfence_block();\n __syncwarp();\n\n /***\n * Section: Backward of ops after the topk\n * - Backward of the used scaling_factor\n * - Sigmoid/Sqrtsoftplus Post-processing bwd when topk > 1\n * - Softmax bwd if use_pre_softmax is false\n */\n // Backward of the used scaling_factor\n // In-place update\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n if (local_routing_map[i]) {\n local_grad[i] = local_grad[i] * scaling_factor;\n }\n }\n __syncwarp();\n\n // Sqrtsoftplus: First compute sqrtsoftplus output from original logits\n // (needed for both post-processing bwd and activation bwd, compute once here)\n // For sqrtsoftplus, intermediate_output stores original logits\n if (score_function == 2) {\n // Copy original logits to local_comp_buf and apply sqrtsoftplus in-place\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_comp_buf[i] = local_act_from_fwd[i];\n }\n __syncwarp();\n apply_sqrtsoftplus_on_float(local_comp_buf, num_experts, lane_id);\n __syncwarp();\n }\n\n // Sigmoid/Sqrtsoftplus Post-processing bwd when topk > 1 (normalization backward)\n if (topk > 1 && (score_function == 0 || score_function == 2)) {\n // Select the correct activation output buffer:\n // - Sigmoid: local_act_from_fwd already contains sigmoid output\n // - Sqrtsoftplus: local_comp_buf contains sqrtsoftplus output computed above\n CompType *act_output = (score_function == 0) ? local_act_from_fwd : local_comp_buf;\n\n CompType sum_fwd_input = masked_warp_reduce_on_shmem(\n /*data ptr = */ act_output,\n /*mask ptr = */ local_routing_map,\n /*data size = */ num_experts,\n /*reduce func = */ ReduceFuncType::SUM, lane_id);\n // Compute sum of output * grad using registers\n CompType local_sum_Output_x_Grad = 0.0;\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n if (local_routing_map[i]) {\n local_sum_Output_x_Grad += local_grad[i] * act_output[i];\n }\n }\n // Warp reduce the sum\n for (int s = 16; s > 0; s /= 2) {\n local_sum_Output_x_Grad += __shfl_xor_sync(0xffffffff, local_sum_Output_x_Grad, s);\n }\n CompType sum_Output_x_Grad = local_sum_Output_x_Grad;\n // In-place update\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n if (local_routing_map[i]) {\n local_grad[i] =\n local_grad[i] / (sum_fwd_input + epsilon) -\n sum_Output_x_Grad / ((sum_fwd_input + epsilon) * (sum_fwd_input + epsilon));\n } else {\n local_grad[i] = 0.0;\n }\n }\n __syncwarp();\n }\n\n // Softmax bwd if use_pre_softmax is false\n if (!use_pre_softmax && score_function == 1) {\n apply_softmax_bwd_on_float(local_grad, local_act_from_fwd, local_comp_buf, local_routing_map,\n num_experts, lane_id);\n __syncwarp();\n }\n\n /***\n * Section: Backward of topk\n * mask the unselected position in the grad\n */\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n if (!local_routing_map[i]) {\n local_grad[i] = 0.0;\n }\n }\n __syncwarp();\n\n /***\n * Section: Backward of ops before the topk\n * - Pre-softmax bwd\n * - Sigmoid bwd\n * - Sqrtsoftplus bwd\n * - Write the grad_logits to the global mem\n */\n // Pre-softmax bwd\n if (score_function == 1 && use_pre_softmax) {\n apply_softmax_bwd_on_float(local_grad, local_act_from_fwd, local_comp_buf, nullptr,\n num_experts, lane_id);\n __syncwarp();\n }\n // Sigmoid bwd\n if (score_function == 0) {\n apply_sigmoid_bwd_on_float(local_grad, local_act_from_fwd, num_experts, lane_id);\n __syncwarp();\n }\n // Sqrtsoftplus bwd\n // For sqrtsoftplus, local_comp_buf already contains sqrtsoftplus output computed earlier\n // Now compute gradient: dy/dx = sigmoid(x) / (2 * y)\n if (score_function == 2) {\n apply_sqrtsoftplus_bwd_on_float(local_grad, local_comp_buf, local_act_from_fwd, num_experts,\n lane_id);\n __syncwarp();\n }\n // Write the grad_logits to the global mem\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n grad_logits[pos_offset + i] = local_grad[i];\n }\n __syncwarp();\n }\n}\n\ntemplate \nvoid fused_topk_with_score_function_backward_kernel_launcher(\n const bool *routing_map, const CompType *intermediate_output, const DataType *grad_probs,\n int num_tokens, int num_experts, int topk, bool use_pre_softmax, float scaling_factor,\n int score_function, DataType *grad_logits, hipStream_t stream) {\n // Meta data for the kernel\n size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp;\n size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // grad_probs\n +\n num_experts * num_token_per_block * sizeof(CompType) // act_from_fwd\n + num_experts * num_token_per_block * sizeof(CompType) // comp_buf\n + num_experts * num_token_per_block * sizeof(bool); // routing_map\n check_shared_memory_capacity_num_experts(shared_memory_size, num_experts);\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(fused_topk_with_score_function_backward_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize,\n shared_memory_size));\n fused_topk_with_score_function_backward_kernel\n <<>>(\n routing_map, intermediate_output, grad_probs, num_tokens, num_experts, topk,\n use_pre_softmax, scaling_factor, score_function, grad_logits);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid fused_topk_with_score_function_backward(const Tensor &routing_map,\n const Tensor &intermediate_output,\n const Tensor &grad_probs, int num_tokens,\n int num_experts, int topk, bool use_pre_softmax,\n float scaling_factor, int score_function,\n Tensor &grad_logits, hipStream_t stream) {\n TE_ROUTER_PROBS_TYPE_SWITCH_ALL(\n grad_logits.data.dtype, DataType,\n fused_topk_with_score_function_backward_kernel_launcher(\n reinterpret_cast(routing_map.data.dptr),\n reinterpret_cast(intermediate_output.data.dptr),\n reinterpret_cast(grad_probs.data.dptr), num_tokens, num_experts, topk,\n use_pre_softmax, scaling_factor, score_function,\n reinterpret_cast(grad_logits.data.dptr), stream););\n}\n\n} // namespace fused_router\n} // namespace transformer_engine\n\nvoid nvte_fused_topk_with_score_function_forward(\n const NVTETensor logits, int num_tokens, int num_experts, int topk, int use_pre_softmax,\n int num_groups, int group_topk, float scaling_factor, int score_function,\n const NVTETensor expert_bias, NVTETensor probs, NVTETensor routing_map,\n NVTETensor intermediate_output, hipStream_t stream) {\n NVTE_API_CALL(nvte_fused_topk_with_score_function_forward);\n using namespace transformer_engine;\n fused_router::fused_topk_with_score_function_forward(\n *convertNVTETensorCheck(logits), num_tokens, num_experts, topk,\n static_cast(use_pre_softmax), num_groups, group_topk, scaling_factor, score_function,\n *convertNVTETensorCheck(expert_bias), *convertNVTETensorCheck(probs),\n *convertNVTETensorCheck(routing_map), *convertNVTETensorCheck(intermediate_output), stream);\n}\n\nvoid nvte_fused_topk_with_score_function_backward(const NVTETensor routing_map,\n const NVTETensor intermediate_output,\n const NVTETensor grad_probs, int num_tokens,\n int num_experts, int topk, int use_pre_softmax,\n float scaling_factor, int score_function,\n NVTETensor grad_logits, hipStream_t stream) {\n NVTE_API_CALL(nvte_fused_topk_with_score_function_backward);\n using namespace transformer_engine;\n fused_router::fused_topk_with_score_function_backward(\n *convertNVTETensorCheck(routing_map), *convertNVTETensorCheck(intermediate_output),\n *convertNVTETensorCheck(grad_probs), num_tokens, num_experts, topk,\n static_cast(use_pre_softmax), scaling_factor, score_function,\n *convertNVTETensorCheck(grad_logits), stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/fused_router/fused_topk_with_score_function.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"../utils.cuh\"\n#include \"utils.h\"\n\nnamespace transformer_engine {\nnamespace fused_router {\n\ntemplate \n__global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logits, int num_tokens,\n int num_experts, int topk,\n int score_function, float *scores,\n bool *routing_map,\n CompType *intermediate_output) {\n /***\n * Section: Global Variables/Addresses init\n * - Each warp is responsible for one token, and has own shared memory buffer.\n * Then __syncwarp() is used instead of __syncthreads()\n */\n // Used variables/addresses init\n int num_token_per_block = blockDim.x / kThreadsPerWarp;\n int warp_id = threadIdx.x / kThreadsPerWarp;\n int lane_id = threadIdx.x % kThreadsPerWarp;\n extern __shared__ float shmem_scores_for_aux_loss[];\n CompType *logits_buf = reinterpret_cast(shmem_scores_for_aux_loss);\n CompType *topk_logits_buf =\n reinterpret_cast(logits_buf + num_experts * num_token_per_block);\n int *topk_indices_buf = reinterpret_cast(topk_logits_buf + topk * num_token_per_block);\n // The address of buffers on the current warp\n CompType *local_logits = logits_buf + warp_id * num_experts;\n CompType *topk_logits = topk_logits_buf + warp_id * topk;\n int *topk_indices = topk_indices_buf + warp_id * topk;\n\n /***\n * Section: Main Loop\n * - Each warp is responsible for one token\n */\n int total_round = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n for (int round = blockIdx.x; round < total_round; round += gridDim.x) {\n int token_offset_cur_warp = round * num_token_per_block + warp_id;\n // Each warp is responsible for one token\n if (token_offset_cur_warp >= num_tokens) break;\n\n /***\n * Section: Init buffer\n * - Clear the global buffer which will accept the result of this round\n * - Clear/Init the shmem buffer used by current warp this round\n * - Load the logits to shmem\n */\n int pos_offset = token_offset_cur_warp * num_experts;\n // Clear the routing_map (num_experts)\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n routing_map[pos_offset + i] = false;\n if (score_function == 1) {\n intermediate_output[pos_offset + i] = -std::numeric_limits::infinity();\n }\n }\n // Load the logits to shmem\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_logits[i] = static_cast(logits[pos_offset + i]);\n }\n __threadfence_block();\n __syncwarp();\n\n /***\n * Section: Preprocess\n * Possible preprocess the scores before the topk operation\n * - Pre-softmax\n * - Sigmoid\n * - Sqrtsoftplus\n * - Sigmoid/Sqrtsoftplus post-processing when topk > 1\n * This is in-place scores update\n */\n if (score_function == 1) { // score_function == 1 means softmax\n // Apply softmax to the logits before the topk\n apply_softmax_on_float(local_logits, num_experts, lane_id);\n __syncwarp();\n // Save the softmax output for backward\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n intermediate_output[pos_offset + i] = local_logits[i];\n }\n } else if (score_function == 0) { // score_function == 0 means sigmoid\n // Apply sigmoid to the logits\n apply_sigmoid_on_float(local_logits, num_experts, lane_id);\n __syncwarp();\n // Save the sigmoid output for backward\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n intermediate_output[pos_offset + i] = local_logits[i];\n }\n } else if (score_function == 2) { // score_function == 2 means sqrtsoftplus\n // First save the original logits for backward (needed for gradient computation)\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n intermediate_output[pos_offset + i] = local_logits[i]; // Save original logits\n }\n __syncwarp();\n // Apply sqrtsoftplus to the logits\n apply_sqrtsoftplus_on_float(local_logits, num_experts, lane_id);\n }\n\n __syncwarp(); //Confirm the scores is written to the output\n\n // Sigmoid/Sqrtsoftplus post-processing\n if (score_function == 0 || score_function == 2) {\n auto sum_logits =\n warp_reduce_on_shmem(local_logits, num_experts, ReduceFuncType::SUM, lane_id);\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_logits[i] /= (sum_logits + epsilon);\n }\n __syncwarp();\n }\n\n /***\n * Section: Topk\n * Get the topk indices\n */\n topk_and_mask(local_logits, num_experts, topk, topk_indices, topk_logits, lane_id);\n __syncwarp();\n\n // Write the routing_map to the output tensor\n for (int i = lane_id; i < topk; i += kThreadsPerWarp) {\n routing_map[pos_offset + topk_indices[i]] = true;\n }\n // Write the scores to the output tensor\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n scores[pos_offset + i] = local_logits[i];\n }\n __threadfence_block();\n __syncwarp();\n }\n}\n\ntemplate \nvoid fused_score_for_moe_aux_loss_forward_kernel_launcher(\n const DataType *logits, int num_tokens, int num_experts, int topk, int score_function,\n float *scores, bool *routing_map, CompType *intermediate_output, cudaStream_t stream) {\n // Meta data for the kernel\n size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp;\n size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // logits\n + topk * num_token_per_block * sizeof(CompType) // topk_logits\n + topk * num_token_per_block * sizeof(int); // topk_indices\n check_shared_memory_capacity_num_experts(shared_memory_size, num_experts);\n // Radix selection is O(E), independent of K, but it needs 4 passes for 32-bit float;\n // switch at K=16 where naive O(K^2*E) starts to dominate\n if (topk < 16) {\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n fused_score_for_moe_aux_loss_forward_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size));\n fused_score_for_moe_aux_loss_forward_kernel\n <<>>(\n logits, num_tokens, num_experts, topk, score_function, scores, routing_map,\n intermediate_output);\n } else {\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(\n fused_score_for_moe_aux_loss_forward_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size));\n fused_score_for_moe_aux_loss_forward_kernel\n <<>>(\n logits, num_tokens, num_experts, topk, score_function, scores, routing_map,\n intermediate_output);\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid fused_score_for_moe_aux_loss_forward(const Tensor &logits, int num_tokens, int num_experts,\n int topk, int score_function, Tensor &scores,\n Tensor &routing_map, Tensor &intermediate_output,\n cudaStream_t stream) {\n TE_ROUTER_PROBS_TYPE_SWITCH_ALL(\n logits.data.dtype, DataType,\n fused_score_for_moe_aux_loss_forward_kernel_launcher(\n reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk,\n score_function, reinterpret_cast(scores.data.dptr),\n reinterpret_cast(routing_map.data.dptr),\n reinterpret_cast(intermediate_output.data.dptr), stream););\n}\n\ntemplate \n__global__ void fused_score_for_moe_aux_loss_backward_kernel(const CompType *intermediate_output,\n const float *grad_scores,\n int num_tokens, int num_experts,\n int topk, int score_function,\n DataType *grad_logits) {\n /***\n * Section: Global Variables/Addresses init\n * - Each warp is responsible for one token, and has own shared memory buffer.\n * Then __syncwarp() is used instead of __syncthreads()\n */\n // Used variables/addresses init\n int num_token_per_block = blockDim.x / kThreadsPerWarp;\n int warp_id = threadIdx.x / kThreadsPerWarp;\n int lane_id = threadIdx.x % kThreadsPerWarp;\n extern __shared__ float shmem[];\n CompType *grad_scores_buf = reinterpret_cast(shmem);\n // To store the output of softmax/sigmoid from fwd, or original logits for sqrtsoftplus\n CompType *act_from_fwd_buf = grad_scores_buf + num_experts * num_token_per_block;\n CompType *comp_buf = act_from_fwd_buf + num_experts * num_token_per_block;\n // The address of buffers on the current warp\n CompType *local_grad = grad_scores_buf + warp_id * num_experts;\n CompType *local_act_from_fwd = act_from_fwd_buf + warp_id * num_experts;\n CompType *local_comp_buf = comp_buf + warp_id * num_experts;\n\n /***\n * Section: Main Loop\n * - Each warp is responsible for one token\n */\n int total_round = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n for (int round = blockIdx.x; round < total_round; round += gridDim.x) {\n int token_offset_cur_warp = round * num_token_per_block + warp_id;\n // Each warp is responsible for one token\n if (token_offset_cur_warp >= num_tokens) break;\n\n /***\n * Section: Init buffer\n * - Clear the global buffer which will accept the result of this round\n * - Clear/Init the shmem buffer used by current warp this round\n * - Load the dgrad/output_from_fwd to shmem\n */\n int pos_offset = token_offset_cur_warp * num_experts;\n // Load the dgrad/output_from_fwd to shmem\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_grad[i] = grad_scores[pos_offset + i];\n local_act_from_fwd[i] = intermediate_output[pos_offset + i];\n }\n __threadfence_block();\n __syncwarp();\n\n /***\n * Section: Backward of ops before the topk\n * - Pre-softmax bwd\n * - Sigmoid/Sqrtsoftplus Post-processing bwd when topk > 1\n * - Sigmoid bwd\n * - Sqrtsoftplus bwd\n * - Write the grad_logits to the global mem\n */\n // Sqrtsoftplus: First compute sqrtsoftplus output from original logits\n // (needed for both post-processing bwd and activation bwd, compute once here)\n // For sqrtsoftplus, intermediate_output stores original logits\n if (score_function == 2) {\n // Copy original logits to local_comp_buf and apply sqrtsoftplus in-place\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_comp_buf[i] = local_act_from_fwd[i];\n }\n __syncwarp();\n apply_sqrtsoftplus_on_float(local_comp_buf, num_experts, lane_id);\n __syncwarp();\n }\n\n // Sigmoid/Sqrtsoftplus Post-processing bwd (normalization backward)\n if (score_function == 0 || score_function == 2) {\n // Select the correct activation output buffer:\n // - Sigmoid: local_act_from_fwd already contains sigmoid output\n // - Sqrtsoftplus: local_comp_buf contains sqrtsoftplus output computed above\n CompType *act_output = (score_function == 0) ? local_act_from_fwd : local_comp_buf;\n\n auto sum_fwd_input =\n warp_reduce_on_shmem(act_output, num_experts, ReduceFuncType::SUM, lane_id);\n // Compute sum of output * grad using registers\n CompType local_sum_Output_x_Grad = 0.0;\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_sum_Output_x_Grad += local_grad[i] * act_output[i];\n }\n // Warp reduce the sum\n for (int s = 16; s > 0; s /= 2) {\n local_sum_Output_x_Grad += __shfl_xor_sync(0xffffffff, local_sum_Output_x_Grad, s);\n }\n CompType sum_Output_x_Grad = local_sum_Output_x_Grad;\n // In-place update\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_grad[i] = local_grad[i] / (sum_fwd_input + epsilon) -\n sum_Output_x_Grad / ((sum_fwd_input + epsilon) * (sum_fwd_input + epsilon));\n }\n __syncwarp();\n }\n\n // Pre-softmax bwd\n if (score_function == 1) {\n apply_softmax_bwd_on_float(local_grad, local_act_from_fwd, local_comp_buf, nullptr,\n num_experts, lane_id);\n __syncwarp();\n }\n // Sigmoid bwd\n if (score_function == 0) {\n apply_sigmoid_bwd_on_float(local_grad, local_act_from_fwd, num_experts, lane_id);\n __syncwarp();\n }\n // Sqrtsoftplus bwd\n // For sqrtsoftplus, local_comp_buf already contains sqrtsoftplus output computed earlier\n // Now compute gradient: dy/dx = sigmoid(x) / (2 * y)\n if (score_function == 2) {\n apply_sqrtsoftplus_bwd_on_float(local_grad, local_comp_buf, local_act_from_fwd, num_experts,\n lane_id);\n __syncwarp();\n }\n // Write the grad_logits to the global mem\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n grad_logits[pos_offset + i] = static_cast(local_grad[i]);\n }\n __syncwarp();\n }\n}\n\ntemplate \nvoid fused_score_for_moe_aux_loss_backward_kernel_launcher(\n const CompType *intermediate_output, const float *grad_scores, int num_tokens, int num_experts,\n int topk, int score_function, DataType *grad_logits, cudaStream_t stream) {\n // Meta data for the kernel\n size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp;\n size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // grad_scores\n +\n num_experts * num_token_per_block * sizeof(CompType) // act_from_fwd\n + num_experts * num_token_per_block * sizeof(CompType); // comp_buf\n check_shared_memory_capacity_num_experts(shared_memory_size, num_experts);\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(fused_score_for_moe_aux_loss_backward_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize,\n shared_memory_size));\n fused_score_for_moe_aux_loss_backward_kernel\n <<>>(\n intermediate_output, grad_scores, num_tokens, num_experts, topk, score_function,\n grad_logits);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid fused_score_for_moe_aux_loss_backward(const Tensor &intermediate_output,\n const Tensor &grad_scores, int num_tokens,\n int num_experts, int topk, int score_function,\n Tensor &grad_logits, cudaStream_t stream) {\n TE_ROUTER_PROBS_TYPE_SWITCH_ALL(\n grad_logits.data.dtype, DataType,\n fused_score_for_moe_aux_loss_backward_kernel_launcher(\n reinterpret_cast(intermediate_output.data.dptr),\n reinterpret_cast(grad_scores.data.dptr), num_tokens, num_experts, topk,\n score_function, reinterpret_cast(grad_logits.data.dptr), stream););\n}\n\n} // namespace fused_router\n} // namespace transformer_engine\n\nvoid nvte_fused_score_for_moe_aux_loss_forward(const NVTETensor logits, int num_tokens,\n int num_experts, int topk, int score_function,\n NVTETensor scores, const NVTETensor routing_map,\n const NVTETensor intermediate_output,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_fused_score_for_moe_aux_loss_forward);\n using namespace transformer_engine;\n fused_router::fused_score_for_moe_aux_loss_forward(\n *convertNVTETensorCheck(logits), num_tokens, num_experts, topk, score_function,\n *convertNVTETensorCheck(scores), *convertNVTETensorCheck(routing_map),\n *convertNVTETensorCheck(intermediate_output), stream);\n}\n\nvoid nvte_fused_score_for_moe_aux_loss_backward(const NVTETensor intermediate_output,\n const NVTETensor grad_scores, int num_tokens,\n int num_experts, int topk, int score_function,\n NVTETensor grad_logits, cudaStream_t stream) {\n NVTE_API_CALL(nvte_fused_score_for_moe_aux_loss_backward);\n using namespace transformer_engine;\n fused_router::fused_score_for_moe_aux_loss_backward(\n *convertNVTETensorCheck(intermediate_output), *convertNVTETensorCheck(grad_scores),\n num_tokens, num_experts, topk, score_function, *convertNVTETensorCheck(grad_logits), stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"../utils.cuh\"\n#include \"utils.h\"\n\nnamespace transformer_engine {\nnamespace fused_router {\n\ntemplate \n__global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logits, int num_tokens,\n int num_experts, int topk,\n int score_function, float *scores,\n bool *routing_map,\n CompType *intermediate_output) {\n /***\n * Section: Global Variables/Addresses init\n * - Each warp is responsible for one token, and has own shared memory buffer.\n * Then __syncwarp() is used instead of __syncthreads()\n */\n // Used variables/addresses init\n int num_token_per_block = blockDim.x / kThreadsPerWarp;\n int warp_id = threadIdx.x / kThreadsPerWarp;\n int lane_id = threadIdx.x % kThreadsPerWarp;\n extern __shared__ float shmem_scores_for_aux_loss[];\n CompType *logits_buf = reinterpret_cast(shmem_scores_for_aux_loss);\n CompType *topk_logits_buf =\n reinterpret_cast(logits_buf + num_experts * num_token_per_block);\n int *topk_indices_buf = reinterpret_cast(topk_logits_buf + topk * num_token_per_block);\n // The address of buffers on the current warp\n CompType *local_logits = logits_buf + warp_id * num_experts;\n CompType *topk_logits = topk_logits_buf + warp_id * topk;\n int *topk_indices = topk_indices_buf + warp_id * topk;\n\n /***\n * Section: Main Loop\n * - Each warp is responsible for one token\n */\n int total_round = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n for (int round = blockIdx.x; round < total_round; round += gridDim.x) {\n int token_offset_cur_warp = round * num_token_per_block + warp_id;\n // Each warp is responsible for one token\n if (token_offset_cur_warp >= num_tokens) break;\n\n /***\n * Section: Init buffer\n * - Clear the global buffer which will accept the result of this round\n * - Clear/Init the shmem buffer used by current warp this round\n * - Load the logits to shmem\n */\n int pos_offset = token_offset_cur_warp * num_experts;\n // Clear the routing_map (num_experts)\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n routing_map[pos_offset + i] = false;\n if (score_function == 1) {\n intermediate_output[pos_offset + i] = -std::numeric_limits::infinity();\n }\n }\n // Load the logits to shmem\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_logits[i] = static_cast(logits[pos_offset + i]);\n }\n __threadfence_block();\n __syncwarp();\n\n /***\n * Section: Preprocess\n * Possible preprocess the scores before the topk operation\n * - Pre-softmax\n * - Sigmoid\n * - Sqrtsoftplus\n * - Sigmoid/Sqrtsoftplus post-processing when topk > 1\n * This is in-place scores update\n */\n if (score_function == 1) { // score_function == 1 means softmax\n // Apply softmax to the logits before the topk\n apply_softmax_on_float(local_logits, num_experts, lane_id);\n __syncwarp();\n // Save the softmax output for backward\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n intermediate_output[pos_offset + i] = local_logits[i];\n }\n } else if (score_function == 0) { // score_function == 0 means sigmoid\n // Apply sigmoid to the logits\n apply_sigmoid_on_float(local_logits, num_experts, lane_id);\n __syncwarp();\n // Save the sigmoid output for backward\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n intermediate_output[pos_offset + i] = local_logits[i];\n }\n } else if (score_function == 2) { // score_function == 2 means sqrtsoftplus\n // First save the original logits for backward (needed for gradient computation)\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n intermediate_output[pos_offset + i] = local_logits[i]; // Save original logits\n }\n __syncwarp();\n // Apply sqrtsoftplus to the logits\n apply_sqrtsoftplus_on_float(local_logits, num_experts, lane_id);\n }\n\n __syncwarp(); //Confirm the scores is written to the output\n\n // Sigmoid/Sqrtsoftplus post-processing\n if (score_function == 0 || score_function == 2) {\n auto sum_logits =\n warp_reduce_on_shmem(local_logits, num_experts, ReduceFuncType::SUM, lane_id);\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_logits[i] /= (sum_logits + epsilon);\n }\n __syncwarp();\n }\n\n /***\n * Section: Topk\n * Get the topk indices\n */\n topk_and_mask(local_logits, num_experts, topk, topk_indices, topk_logits, lane_id);\n __syncwarp();\n\n // Write the routing_map to the output tensor\n for (int i = lane_id; i < topk; i += kThreadsPerWarp) {\n routing_map[pos_offset + topk_indices[i]] = true;\n }\n // Write the scores to the output tensor\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n scores[pos_offset + i] = local_logits[i];\n }\n __threadfence_block();\n __syncwarp();\n }\n}\n\ntemplate \nvoid fused_score_for_moe_aux_loss_forward_kernel_launcher(\n const DataType *logits, int num_tokens, int num_experts, int topk, int score_function,\n float *scores, bool *routing_map, CompType *intermediate_output, hipStream_t stream) {\n // Meta data for the kernel\n size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp;\n size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // logits\n + topk * num_token_per_block * sizeof(CompType) // topk_logits\n + topk * num_token_per_block * sizeof(int); // topk_indices\n check_shared_memory_capacity_num_experts(shared_memory_size, num_experts);\n // Radix selection is O(E), independent of K, but it needs 4 passes for 32-bit float;\n // switch at K=16 where naive O(K^2*E) starts to dominate\n if (topk < 16) {\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(fused_score_for_moe_aux_loss_forward_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size));\n fused_score_for_moe_aux_loss_forward_kernel\n <<>>(\n logits, num_tokens, num_experts, topk, score_function, scores, routing_map,\n intermediate_output);\n } else {\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(fused_score_for_moe_aux_loss_forward_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size));\n fused_score_for_moe_aux_loss_forward_kernel\n <<>>(\n logits, num_tokens, num_experts, topk, score_function, scores, routing_map,\n intermediate_output);\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid fused_score_for_moe_aux_loss_forward(const Tensor &logits, int num_tokens, int num_experts,\n int topk, int score_function, Tensor &scores,\n Tensor &routing_map, Tensor &intermediate_output,\n hipStream_t stream) {\n TE_ROUTER_PROBS_TYPE_SWITCH_ALL(\n logits.data.dtype, DataType,\n fused_score_for_moe_aux_loss_forward_kernel_launcher(\n reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk,\n score_function, reinterpret_cast(scores.data.dptr),\n reinterpret_cast(routing_map.data.dptr),\n reinterpret_cast(intermediate_output.data.dptr), stream););\n}\n\ntemplate \n__global__ void fused_score_for_moe_aux_loss_backward_kernel(const CompType *intermediate_output,\n const float *grad_scores,\n int num_tokens, int num_experts,\n int topk, int score_function,\n DataType *grad_logits) {\n /***\n * Section: Global Variables/Addresses init\n * - Each warp is responsible for one token, and has own shared memory buffer.\n * Then __syncwarp() is used instead of __syncthreads()\n */\n // Used variables/addresses init\n int num_token_per_block = blockDim.x / kThreadsPerWarp;\n int warp_id = threadIdx.x / kThreadsPerWarp;\n int lane_id = threadIdx.x % kThreadsPerWarp;\n extern __shared__ float shmem[];\n CompType *grad_scores_buf = reinterpret_cast(shmem);\n // To store the output of softmax/sigmoid from fwd, or original logits for sqrtsoftplus\n CompType *act_from_fwd_buf = grad_scores_buf + num_experts * num_token_per_block;\n CompType *comp_buf = act_from_fwd_buf + num_experts * num_token_per_block;\n // The address of buffers on the current warp\n CompType *local_grad = grad_scores_buf + warp_id * num_experts;\n CompType *local_act_from_fwd = act_from_fwd_buf + warp_id * num_experts;\n CompType *local_comp_buf = comp_buf + warp_id * num_experts;\n\n /***\n * Section: Main Loop\n * - Each warp is responsible for one token\n */\n int total_round = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n for (int round = blockIdx.x; round < total_round; round += gridDim.x) {\n int token_offset_cur_warp = round * num_token_per_block + warp_id;\n // Each warp is responsible for one token\n if (token_offset_cur_warp >= num_tokens) break;\n\n /***\n * Section: Init buffer\n * - Clear the global buffer which will accept the result of this round\n * - Clear/Init the shmem buffer used by current warp this round\n * - Load the dgrad/output_from_fwd to shmem\n */\n int pos_offset = token_offset_cur_warp * num_experts;\n // Load the dgrad/output_from_fwd to shmem\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_grad[i] = grad_scores[pos_offset + i];\n local_act_from_fwd[i] = intermediate_output[pos_offset + i];\n }\n __threadfence_block();\n __syncwarp();\n\n /***\n * Section: Backward of ops before the topk\n * - Pre-softmax bwd\n * - Sigmoid/Sqrtsoftplus Post-processing bwd when topk > 1\n * - Sigmoid bwd\n * - Sqrtsoftplus bwd\n * - Write the grad_logits to the global mem\n */\n // Sqrtsoftplus: First compute sqrtsoftplus output from original logits\n // (needed for both post-processing bwd and activation bwd, compute once here)\n // For sqrtsoftplus, intermediate_output stores original logits\n if (score_function == 2) {\n // Copy original logits to local_comp_buf and apply sqrtsoftplus in-place\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_comp_buf[i] = local_act_from_fwd[i];\n }\n __syncwarp();\n apply_sqrtsoftplus_on_float(local_comp_buf, num_experts, lane_id);\n __syncwarp();\n }\n\n // Sigmoid/Sqrtsoftplus Post-processing bwd (normalization backward)\n if (score_function == 0 || score_function == 2) {\n // Select the correct activation output buffer:\n // - Sigmoid: local_act_from_fwd already contains sigmoid output\n // - Sqrtsoftplus: local_comp_buf contains sqrtsoftplus output computed above\n CompType *act_output = (score_function == 0) ? local_act_from_fwd : local_comp_buf;\n\n auto sum_fwd_input =\n warp_reduce_on_shmem(act_output, num_experts, ReduceFuncType::SUM, lane_id);\n // Compute sum of output * grad using registers\n CompType local_sum_Output_x_Grad = 0.0;\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_sum_Output_x_Grad += local_grad[i] * act_output[i];\n }\n // Warp reduce the sum\n for (int s = 16; s > 0; s /= 2) {\n local_sum_Output_x_Grad += __shfl_xor_sync(0xffffffff, local_sum_Output_x_Grad, s);\n }\n CompType sum_Output_x_Grad = local_sum_Output_x_Grad;\n // In-place update\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n local_grad[i] = local_grad[i] / (sum_fwd_input + epsilon) -\n sum_Output_x_Grad / ((sum_fwd_input + epsilon) * (sum_fwd_input + epsilon));\n }\n __syncwarp();\n }\n\n // Pre-softmax bwd\n if (score_function == 1) {\n apply_softmax_bwd_on_float(local_grad, local_act_from_fwd, local_comp_buf, nullptr,\n num_experts, lane_id);\n __syncwarp();\n }\n // Sigmoid bwd\n if (score_function == 0) {\n apply_sigmoid_bwd_on_float(local_grad, local_act_from_fwd, num_experts, lane_id);\n __syncwarp();\n }\n // Sqrtsoftplus bwd\n // For sqrtsoftplus, local_comp_buf already contains sqrtsoftplus output computed earlier\n // Now compute gradient: dy/dx = sigmoid(x) / (2 * y)\n if (score_function == 2) {\n apply_sqrtsoftplus_bwd_on_float(local_grad, local_comp_buf, local_act_from_fwd, num_experts,\n lane_id);\n __syncwarp();\n }\n // Write the grad_logits to the global mem\n for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) {\n grad_logits[pos_offset + i] = static_cast(local_grad[i]);\n }\n __syncwarp();\n }\n}\n\ntemplate \nvoid fused_score_for_moe_aux_loss_backward_kernel_launcher(\n const CompType *intermediate_output, const float *grad_scores, int num_tokens, int num_experts,\n int topk, int score_function, DataType *grad_logits, hipStream_t stream) {\n // Meta data for the kernel\n size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp;\n size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block;\n size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // grad_scores\n +\n num_experts * num_token_per_block * sizeof(CompType) // act_from_fwd\n + num_experts * num_token_per_block * sizeof(CompType); // comp_buf\n check_shared_memory_capacity_num_experts(shared_memory_size, num_experts);\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(fused_score_for_moe_aux_loss_backward_kernel),\n hipFuncAttributeMaxDynamicSharedMemorySize,\n shared_memory_size));\n fused_score_for_moe_aux_loss_backward_kernel\n <<>>(\n intermediate_output, grad_scores, num_tokens, num_experts, topk, score_function,\n grad_logits);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid fused_score_for_moe_aux_loss_backward(const Tensor &intermediate_output,\n const Tensor &grad_scores, int num_tokens,\n int num_experts, int topk, int score_function,\n Tensor &grad_logits, hipStream_t stream) {\n TE_ROUTER_PROBS_TYPE_SWITCH_ALL(\n grad_logits.data.dtype, DataType,\n fused_score_for_moe_aux_loss_backward_kernel_launcher(\n reinterpret_cast(intermediate_output.data.dptr),\n reinterpret_cast(grad_scores.data.dptr), num_tokens, num_experts, topk,\n score_function, reinterpret_cast(grad_logits.data.dptr), stream););\n}\n\n} // namespace fused_router\n} // namespace transformer_engine\n\nvoid nvte_fused_score_for_moe_aux_loss_forward(const NVTETensor logits, int num_tokens,\n int num_experts, int topk, int score_function,\n NVTETensor scores, const NVTETensor routing_map,\n const NVTETensor intermediate_output,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_fused_score_for_moe_aux_loss_forward);\n using namespace transformer_engine;\n fused_router::fused_score_for_moe_aux_loss_forward(\n *convertNVTETensorCheck(logits), num_tokens, num_experts, topk, score_function,\n *convertNVTETensorCheck(scores), *convertNVTETensorCheck(routing_map),\n *convertNVTETensorCheck(intermediate_output), stream);\n}\n\nvoid nvte_fused_score_for_moe_aux_loss_backward(const NVTETensor intermediate_output,\n const NVTETensor grad_scores, int num_tokens,\n int num_experts, int topk, int score_function,\n NVTETensor grad_logits, hipStream_t stream) {\n NVTE_API_CALL(nvte_fused_score_for_moe_aux_loss_backward);\n using namespace transformer_engine;\n fused_router::fused_score_for_moe_aux_loss_backward(\n *convertNVTETensorCheck(intermediate_output), *convertNVTETensorCheck(grad_scores),\n num_tokens, num_experts, topk, score_function, *convertNVTETensorCheck(grad_logits), stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\n\ntemplate \n__device__ void fused_rope_block_forward(const scalar_t *src, const float *freqs, scalar_t *dst,\n const bool interleaved, const int s_id,\n const int offset_block, const int offset_block_dst,\n const int h, const int d, const int d2, const int stride_h,\n const int stride_d, const int o_stride_h,\n const int o_stride_d) {\n extern __shared__ float shared_mem_cos_sin[];\n float *shared_mem_cos = shared_mem_cos_sin;\n float *shared_mem_sin = shared_mem_cos_sin + d2;\n int tid = threadIdx.x * blockDim.y + threadIdx.y;\n for (int i = tid; i < d2; i += blockDim.x * blockDim.y) {\n sincosf(freqs[s_id * d2 + i], &shared_mem_sin[i], &shared_mem_cos[i]);\n }\n __syncthreads();\n\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n#pragma unroll\n for (int d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) {\n float v_cos = shared_mem_cos[d_id];\n float v_sin = shared_mem_sin[d_id];\n int offset_src = offset_block + h_id * stride_h + d_id * stride_d;\n int offset_dst = offset_block_dst + h_id * o_stride_h + d_id * o_stride_d;\n float v_src = src[offset_src];\n float v_src_rotate;\n if (!interleaved) {\n v_src_rotate = (d_id + d2 / 2 < d2)\n ? -static_cast(src[offset_src + (d2 / 2) * stride_d])\n : static_cast(src[offset_src + (d2 / 2 - d2) * stride_d]);\n } else {\n v_src_rotate = (d_id % 2 == 0)\n // d_id + 1\n ? -static_cast(src[offset_src + stride_d])\n // d_id - 1\n : static_cast(src[offset_src - stride_d]);\n }\n dst[offset_dst] = v_src * v_cos + v_src_rotate * v_sin;\n }\n }\n\n // copy the rest\n if (d > d2) {\n#pragma unroll\n for (int d_id = d2 + threadIdx.x; d_id < d; d_id += blockDim.x) {\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n int offset_src = offset_block + h_id * stride_h + d_id * stride_d;\n int offset_dst = offset_block_dst + h_id * o_stride_h + d_id * o_stride_d;\n dst[offset_dst] = src[offset_src];\n }\n }\n }\n}\n\ntemplate \n__device__ void fused_rope_block_backward(const scalar_t *src, const float *freqs, scalar_t *dst,\n const bool interleaved, const int s_id,\n const int offset_block, const int offset_block_dst,\n const int h, const int d, const int d2,\n const int stride_h, const int stride_d,\n const int o_stride_h, const int o_stride_d) {\n extern __shared__ float shared_mem_cos_sin[];\n float *shared_mem_cos = shared_mem_cos_sin;\n float *shared_mem_sin = shared_mem_cos_sin + d2;\n int tid = threadIdx.x * blockDim.y + threadIdx.y;\n for (int i = tid; i < d2; i += blockDim.x * blockDim.y) {\n sincosf(freqs[s_id * d2 + i], &shared_mem_sin[i], &shared_mem_cos[i]);\n }\n __syncthreads();\n\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n#pragma unroll\n for (int d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) {\n int offset_src = offset_block + h_id * stride_h + d_id * stride_d;\n int offset_dst = offset_block_dst + h_id * o_stride_h + d_id * o_stride_d;\n float v_src = src[offset_src];\n float v_cos = shared_mem_cos[d_id];\n float v_src_rotate, v_sin;\n if (!interleaved) {\n if (d_id + d2 / 2 < d2) {\n v_src_rotate = static_cast(src[offset_src + (d2 / 2) * stride_d]);\n v_sin = shared_mem_sin[d_id + d2 / 2];\n } else {\n v_src_rotate = static_cast(src[offset_src + (d2 / 2 - d2) * stride_d]);\n v_sin = -shared_mem_sin[d_id + d2 / 2 - d2];\n }\n } else {\n if (d_id % 2 == 0) {\n v_src_rotate = static_cast(src[offset_src + stride_d]);\n v_sin = shared_mem_sin[d_id + 1];\n } else {\n v_src_rotate = static_cast(src[offset_src - stride_d]);\n v_sin = -shared_mem_sin[d_id - 1];\n }\n }\n dst[offset_dst] = v_src * v_cos + v_src_rotate * v_sin;\n }\n }\n\n // copy the rest\n if (d > d2) {\n#pragma unroll\n for (int d_id = d2 + threadIdx.x; d_id < d; d_id += blockDim.x) {\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n int offset_src = offset_block + h_id * stride_h + d_id * stride_d;\n int offset_dst = offset_block_dst + h_id * o_stride_h + d_id * o_stride_d;\n dst[offset_dst] = src[offset_src];\n }\n }\n }\n}\n\ntemplate \n__global__ void fused_rope_forward_kernel(const scalar_t *src, const int *cu_seqlens,\n const float *freqs, const int *start_positions,\n scalar_t *dst, const bool interleaved, const int cp_size,\n const int cp_rank, const int s, const int h, const int d,\n const int d2, const int stride_s_or_t, const int stride_b,\n const int stride_h, const int stride_d,\n const int o_stride_s_or_t, const int o_stride_b,\n const int o_stride_h, const int o_stride_d) {\n int s_id = blockIdx.x, b_id = blockIdx.y;\n int offset_block, offset_block_dst;\n int cur_seqlens;\n if (cu_seqlens != nullptr) { // THD\n int start = cu_seqlens[b_id] / cp_size;\n int end = cu_seqlens[b_id + 1] / cp_size;\n int t_id = s_id + start;\n if (t_id >= end) return;\n offset_block = t_id * stride_s_or_t;\n offset_block_dst = t_id * o_stride_s_or_t;\n cur_seqlens = end - start;\n } else { // SBHD/BSHD\n offset_block = s_id * stride_s_or_t + b_id * stride_b;\n offset_block_dst = s_id * o_stride_s_or_t + b_id * o_stride_b;\n cur_seqlens = s;\n }\n\n // Offset the RoPE embedding by start_positions if provided.\n int begin_offset = (start_positions == nullptr) ? 0 : start_positions[b_id];\n int s_id_for_freqs = s_id + begin_offset;\n\n // If CP_SIZE > 1, offset the RoPE embedding by cp_rank based on the dual-chunk order.\n if (cp_size > 1) {\n assert(cur_seqlens % 2 == 0);\n if (s_id < cur_seqlens / 2) {\n s_id_for_freqs += cp_rank * cur_seqlens / 2;\n } else {\n s_id_for_freqs += cur_seqlens * cp_size - (cp_rank + 1) * cur_seqlens / 2 - cur_seqlens / 2;\n }\n }\n\n fused_rope_block_forward(src, freqs, dst, interleaved, s_id_for_freqs, offset_block,\n offset_block_dst, h, d, d2, stride_h, stride_d, o_stride_h, o_stride_d);\n}\n\ntemplate \n__global__ void fused_rope_backward_kernel(\n const scalar_t *src, const int *cu_seqlens, const float *freqs, const int *start_positions,\n scalar_t *dst, const bool interleaved, const int cp_size, const int cp_rank, const int s,\n const int h, const int d, const int d2, const int stride_s_or_t, const int stride_b,\n const int stride_h, const int stride_d, const int o_stride_s_or_t, const int o_stride_b,\n const int o_stride_h, const int o_stride_d) {\n int s_id = blockIdx.x, b_id = blockIdx.y;\n int offset_block, offset_block_dst;\n int cur_seqlens;\n if (cu_seqlens != nullptr) { // THD\n int start = cu_seqlens[b_id] / cp_size;\n int end = cu_seqlens[b_id + 1] / cp_size;\n int t_id = s_id + start;\n if (t_id >= end) return;\n offset_block = t_id * stride_s_or_t;\n offset_block_dst = t_id * o_stride_s_or_t;\n cur_seqlens = end - start;\n } else { // SBHD/BSHD\n offset_block = s_id * stride_s_or_t + b_id * stride_b;\n offset_block_dst = s_id * o_stride_s_or_t + b_id * o_stride_b;\n cur_seqlens = s;\n }\n\n // Offset the RoPE embedding by start_positions if provided.\n int begin_offset = (start_positions == nullptr) ? 0 : start_positions[b_id];\n int s_id_for_freqs = s_id + begin_offset;\n\n // If CP_SIZE > 1, offset the RoPE embedding by cp_rank based on the dual-chunk order.\n if (cp_size > 1) {\n assert(cur_seqlens % 2 == 0);\n if (s_id < cur_seqlens / 2) {\n s_id_for_freqs += cp_rank * cur_seqlens / 2;\n } else {\n s_id_for_freqs += cur_seqlens * cp_size - (cp_rank + 1) * cur_seqlens / 2 - cur_seqlens / 2;\n }\n }\n\n fused_rope_block_backward(src, freqs, dst, interleaved, s_id_for_freqs, offset_block,\n offset_block_dst, h, d, d2, stride_h, stride_d, o_stride_h, o_stride_d);\n}\n\ntemplate \n__device__ void fused_qkv_rope_block_forward(const scalar_t *src, const float *freqs, scalar_t *out,\n const bool interleaved, const int s_id,\n const int offset_block, const int offset_block_dst,\n const int h, const int d, const int d2,\n const int row_offset, const int in_row_length,\n const int out_row_length) {\n extern __shared__ float shared_mem_cos_sin_qk[];\n // Split the shared memory into cos and sin parts for q or k\n float *shared_mem_cos = nullptr;\n float *shared_mem_sin = nullptr;\n if (row_offset == 0) { // q\n shared_mem_cos = shared_mem_cos_sin_qk;\n shared_mem_sin = shared_mem_cos_sin_qk + d2;\n } else { // k\n shared_mem_cos = shared_mem_cos_sin_qk + 2 * d2;\n shared_mem_sin = shared_mem_cos_sin_qk + 3 * d2;\n }\n if (freqs != nullptr) {\n int tid = threadIdx.x * blockDim.y + threadIdx.y;\n for (int i = tid; i < d2; i += blockDim.x * blockDim.y) {\n sincosf(freqs[s_id * d2 + i], &shared_mem_sin[i], &shared_mem_cos[i]);\n }\n }\n __syncthreads();\n\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n#pragma unroll\n for (int i = 0; i < out_row_length; i += d) {\n#pragma unroll\n for (int d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) {\n int offset_src = offset_block + h_id * in_row_length + (row_offset + i) + d_id;\n int offset_dst = offset_block_dst + h_id * out_row_length + i + d_id;\n if (freqs != nullptr) {\n float v_cos, v_sin;\n v_cos = shared_mem_cos[d_id];\n v_sin = shared_mem_sin[d_id];\n float v_src = src[offset_src];\n float v_src_rotate;\n if (!interleaved) {\n v_src_rotate = (d_id + d2 / 2 < d2)\n ? -static_cast(src[offset_src + (d2 / 2)])\n : static_cast(src[offset_src + (d2 / 2 - d2)]);\n } else {\n v_src_rotate = (d_id % 2 == 0) ? -static_cast(src[offset_src + 1])\n : static_cast(src[offset_src - 1]);\n }\n out[offset_dst] = v_src * v_cos + v_src_rotate * v_sin;\n } else {\n out[offset_dst] = src[offset_src];\n }\n }\n }\n }\n // copy the rest\n if (d > d2) {\n#pragma unroll\n for (int d_id = d2 + threadIdx.x; d_id < d; d_id += blockDim.x) {\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n#pragma unroll\n for (int i = 0; i < out_row_length; i += d) {\n int offset_src = offset_block + h_id * in_row_length + (row_offset + i) + d_id;\n int offset_dst = offset_block_dst + h_id * out_row_length + i + d_id;\n out[offset_dst] = src[offset_src];\n }\n }\n }\n }\n}\n\ntemplate \n__device__ void fused_qkv_rope_block_backward(const scalar_t *grad_out, const float *freqs,\n scalar_t *out, const bool interleaved, const int s_id,\n const int offset_block, const int offset_block_dst,\n const int h, const int d, const int d2,\n const int row_offset, const int in_row_length,\n const int out_row_length) {\n extern __shared__ float shared_mem_cos_sin_qk[];\n float *shared_mem_cos = nullptr;\n float *shared_mem_sin = nullptr;\n // Split the shared memory into cos and sin parts for q or k\n if (row_offset == 0) { // q\n shared_mem_cos = shared_mem_cos_sin_qk;\n shared_mem_sin = shared_mem_cos_sin_qk + d2;\n } else { // k\n shared_mem_cos = shared_mem_cos_sin_qk + 2 * d2;\n shared_mem_sin = shared_mem_cos_sin_qk + 3 * d2;\n }\n if (freqs != nullptr) {\n int tid = threadIdx.x * blockDim.y + threadIdx.y;\n for (int i = tid; i < d2; i += blockDim.x * blockDim.y) {\n sincosf(freqs[s_id * d2 + i], &shared_mem_sin[i], &shared_mem_cos[i]);\n }\n }\n __syncthreads();\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n#pragma unroll\n for (int i = 0; i < out_row_length; i += d) {\n#pragma unroll\n for (int d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) {\n int offset_dst = offset_block + h_id * in_row_length + (row_offset + i) + d_id;\n int offset_src = offset_block_dst + h_id * out_row_length + i + d_id;\n\n float v_src = grad_out[offset_src];\n if (freqs != nullptr) {\n float v_cos, v_sin;\n v_cos = shared_mem_cos[d_id];\n float v_src_rotate;\n if (!interleaved) {\n if (d_id + d2 / 2 < d2) {\n v_src_rotate = static_cast(grad_out[offset_src + (d2 / 2)]);\n v_sin = shared_mem_sin[d_id + d2 / 2];\n } else {\n v_src_rotate = static_cast(grad_out[offset_src + (d2 / 2 - d2)]);\n v_sin = -shared_mem_sin[d_id + d2 / 2 - d2];\n }\n } else {\n if (d_id % 2 == 0) {\n v_src_rotate = static_cast(grad_out[offset_src + 1]);\n v_sin = shared_mem_sin[d_id + 1];\n } else {\n v_src_rotate = static_cast(grad_out[offset_src - 1]);\n v_sin = -shared_mem_sin[d_id - 1];\n }\n }\n out[offset_dst] = v_src * v_cos + v_src_rotate * v_sin;\n } else {\n out[offset_dst] = grad_out[offset_src];\n }\n }\n }\n }\n // copy the rest\n if (d > d2) {\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n#pragma unroll\n for (int i = 0; i < out_row_length; i += d) {\n#pragma unroll\n for (int d_id = d2 + threadIdx.x; d_id < d; d_id += blockDim.x) {\n int offset_dst = offset_block + h_id * in_row_length + (row_offset + i) + d_id;\n int offset_src = offset_block_dst + h_id * out_row_length + i + d_id;\n out[offset_dst] = grad_out[offset_src];\n }\n }\n }\n }\n}\n\ntemplate \n__global__ void fused_qkv_rope_forward_kernel(\n const scalar_t *qkv_input, const float *q_freqs, const float *k_freqs,\n const int *start_positions, scalar_t *q_out, scalar_t *k_out, scalar_t *v_out,\n const NVTE_QKV_Format qkv_format, const bool interleaved, const int cp_size, const int cp_rank,\n const int s, const int b, const int h, const int d, const int d2, const int q_split_arg,\n const int k_split_arg, const int v_split_arg) {\n int s_id = blockIdx.x, b_id = blockIdx.y;\n int cur_seqlens = s;\n int total_d = q_split_arg + k_split_arg + v_split_arg;\n int offset_block, offset_block_dst_q, offset_block_dst_k, offset_block_dst_v;\n if (qkv_format == NVTE_QKV_Format::NVTE_SBHD) {\n offset_block = s_id * b * h * total_d + b_id * h * total_d;\n offset_block_dst_q = s_id * b * h * q_split_arg + b_id * h * q_split_arg;\n offset_block_dst_k = s_id * b * h * k_split_arg + b_id * h * k_split_arg;\n offset_block_dst_v = s_id * b * h * v_split_arg + b_id * h * v_split_arg;\n } else {\n offset_block = b_id * s * h * total_d + s_id * h * total_d;\n offset_block_dst_q = b_id * s * h * q_split_arg + s_id * h * q_split_arg;\n offset_block_dst_k = b_id * s * h * k_split_arg + s_id * h * k_split_arg;\n offset_block_dst_v = b_id * s * h * v_split_arg + s_id * h * v_split_arg;\n }\n\n int q_limit = q_split_arg;\n int k_limit = q_limit + k_split_arg;\n int s_id_for_freqs;\n if (cp_size > 1) {\n assert(cur_seqlens % 2 == 0);\n if (s_id < cur_seqlens / 2) {\n s_id_for_freqs = s_id + cp_rank * cur_seqlens / 2;\n } else {\n s_id_for_freqs =\n cur_seqlens * cp_size - (cp_rank + 1) * cur_seqlens / 2 + s_id - cur_seqlens / 2;\n }\n } else {\n int begin_offset = (start_positions == nullptr) ? 0 : start_positions[b_id];\n s_id_for_freqs = s_id + begin_offset;\n }\n fused_qkv_rope_block_forward(qkv_input, q_freqs, q_out, interleaved, s_id_for_freqs, offset_block,\n offset_block_dst_q, h, d, d2, 0, total_d, q_split_arg);\n fused_qkv_rope_block_forward(qkv_input, k_freqs, k_out, interleaved, s_id_for_freqs, offset_block,\n offset_block_dst_k, h, d, d2, q_limit, total_d, k_split_arg);\n fused_qkv_rope_block_forward(qkv_input, nullptr, v_out, interleaved, s_id_for_freqs, offset_block,\n offset_block_dst_v, h, d, d2, k_limit, total_d, v_split_arg);\n}\n\ntemplate \n__global__ void fused_qkv_rope_backward_kernel(\n const scalar_t *grad_out_q, const scalar_t *grad_out_k, const scalar_t *grad_out_v,\n const float *q_freqs, const float *k_freqs, scalar_t *qkv_grad,\n const NVTE_QKV_Format qkv_format, const bool interleaved, const int cp_size, const int cp_rank,\n const int s, const int b, const int h, const int d, const int d2, const int q_split_arg,\n const int k_split_arg, const int v_split_arg) {\n int s_id = blockIdx.x, b_id = blockIdx.y;\n int cur_seqlens = s;\n int offset_block, offset_block_dst_q, offset_block_dst_k, offset_block_dst_v;\n int total_d = q_split_arg + k_split_arg + v_split_arg;\n if (qkv_format == NVTE_QKV_Format::NVTE_SBHD) {\n offset_block = s_id * b * h * total_d + b_id * h * total_d;\n offset_block_dst_q = s_id * b * h * q_split_arg + b_id * h * q_split_arg;\n offset_block_dst_k = s_id * b * h * k_split_arg + b_id * h * k_split_arg;\n offset_block_dst_v = s_id * b * h * v_split_arg + b_id * h * v_split_arg;\n } else {\n offset_block = b_id * s * h * total_d + s_id * h * total_d;\n offset_block_dst_q = b_id * s * h * q_split_arg + s_id * h * q_split_arg;\n offset_block_dst_k = b_id * s * h * k_split_arg + s_id * h * k_split_arg;\n offset_block_dst_v = b_id * s * h * v_split_arg + s_id * h * v_split_arg;\n }\n int q_limit = q_split_arg;\n int k_limit = q_limit + k_split_arg;\n int s_id_for_freqs;\n if (cp_size > 1) {\n assert(cur_seqlens % 2 == 0);\n if (s_id < cur_seqlens / 2) {\n s_id_for_freqs = s_id + cp_rank * cur_seqlens / 2;\n } else {\n s_id_for_freqs =\n cur_seqlens * cp_size - (cp_rank + 1) * cur_seqlens / 2 + s_id - cur_seqlens / 2;\n }\n } else {\n s_id_for_freqs = s_id;\n }\n fused_qkv_rope_block_backward(grad_out_q, q_freqs, qkv_grad, interleaved, s_id_for_freqs,\n offset_block, offset_block_dst_q, h, d, d2, 0, total_d,\n q_split_arg);\n fused_qkv_rope_block_backward(grad_out_k, k_freqs, qkv_grad, interleaved, s_id_for_freqs,\n offset_block, offset_block_dst_k, h, d, d2, q_limit, total_d,\n k_split_arg);\n fused_qkv_rope_block_backward(grad_out_v, nullptr, qkv_grad, interleaved, s_id_for_freqs,\n offset_block, offset_block_dst_v, h, d, d2, k_limit, total_d,\n v_split_arg);\n}\n\ntemplate \nvoid fused_rope_forward_launcher(const scalar_t *input, const int *cu_seqlens, const float *freqs,\n const int *start_positions, scalar_t *output,\n const NVTE_QKV_Format qkv_format, const bool interleaved,\n const int cp_size, const int cp_rank, const int s, const int b,\n const int h, const int d, const int d2, const int stride_s_or_t,\n const int stride_b, const int stride_h, const int stride_d,\n cudaStream_t stream) {\n int warps_per_block = h < 16 ? 4 : 8;\n dim3 blocks(s, b);\n dim3 threads(THREADS_PER_WARP, warps_per_block);\n const int shared_mem_size = 2 * d2 * sizeof(float); // cos, sin\n int o_stride_s_or_t, o_stride_b;\n if (qkv_format == NVTE_QKV_Format::NVTE_THD) {\n NVTE_CHECK(cu_seqlens != nullptr, \"cu_seqlens is required for THD format\");\n o_stride_s_or_t = h * d;\n o_stride_b = 0;\n } else if (qkv_format == NVTE_QKV_Format::NVTE_SBHD) {\n o_stride_s_or_t = b * h * d;\n o_stride_b = h * d;\n } else {\n o_stride_s_or_t = h * d;\n o_stride_b = s * h * d;\n }\n const int o_stride_h = d;\n const int o_stride_d = 1;\n\n fused_rope_forward_kernel<<>>(\n input, cu_seqlens, freqs, start_positions, output, interleaved, cp_size, cp_rank, s, h, d, d2,\n stride_s_or_t, stride_b, stride_h, stride_d, o_stride_s_or_t, o_stride_b, o_stride_h,\n o_stride_d);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\ntemplate \nvoid fused_rope_backward_launcher(const scalar_t *output_grads, const int *cu_seqlens,\n const float *freqs, const int *start_positions,\n scalar_t *input_grads, const NVTE_QKV_Format qkv_format,\n const bool interleaved, const int cp_size, const int cp_rank,\n const int s, const int b, const int h, const int d, const int d2,\n const int stride_s_or_t, const int stride_b, const int stride_h,\n const int stride_d, cudaStream_t stream) {\n int warps_per_block = h < 16 ? 4 : 8;\n dim3 blocks(s, b);\n dim3 threads(THREADS_PER_WARP, warps_per_block);\n const int shared_mem_size = 2 * d2 * sizeof(float); // cos, sin\n int o_stride_s_or_t, o_stride_b;\n if (qkv_format == NVTE_QKV_Format::NVTE_THD) {\n NVTE_CHECK(cu_seqlens != nullptr, \"cu_seqlens is required for THD format\");\n o_stride_s_or_t = h * d;\n o_stride_b = 0;\n } else if (qkv_format == NVTE_QKV_Format::NVTE_SBHD) {\n o_stride_s_or_t = b * h * d;\n o_stride_b = h * d;\n } else {\n o_stride_s_or_t = h * d;\n o_stride_b = s * h * d;\n }\n const int o_stride_h = d;\n const int o_stride_d = 1;\n\n fused_rope_backward_kernel<<>>(\n output_grads, cu_seqlens, freqs, start_positions, input_grads, interleaved, cp_size, cp_rank,\n s, h, d, d2, stride_s_or_t, stride_b, stride_h, stride_d, o_stride_s_or_t, o_stride_b,\n o_stride_h, o_stride_d);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\ntemplate \nvoid fused_qkv_rope_forward_launcher(const scalar_t *qkv_input, const float *q_freqs,\n const float *k_freqs, const int *start_positions,\n scalar_t *q_out, scalar_t *k_out, scalar_t *v_out,\n const NVTE_QKV_Format qkv_format, const bool interleaved,\n const int cp_size, const int cp_rank, const int s, const int b,\n const int h, const int d, const int d2,\n const int qkv_split_arg_list_0, const int qkv_split_arg_list_1,\n const int qkv_split_arg_list_2, cudaStream_t stream) {\n const int THREADS_PER_WARP = 32;\n int warps_per_block = (h <= 8) ? h : 8;\n dim3 blocks(s, b);\n dim3 threads(THREADS_PER_WARP, warps_per_block);\n const int shared_mem_size = 4 * d2 * sizeof(float); // cos, sin * q ,k\n\n fused_qkv_rope_forward_kernel<<>>(\n qkv_input, q_freqs, k_freqs, start_positions, q_out, k_out, v_out, qkv_format, interleaved,\n cp_size, cp_rank, s, b, h, d, d2, qkv_split_arg_list_0, qkv_split_arg_list_1,\n qkv_split_arg_list_2);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\ntemplate \nvoid fused_qkv_rope_backward_launcher(const scalar_t *q_grad_out, const scalar_t *k_grad_out,\n const scalar_t *v_grad_out, const float *q_freqs,\n const float *k_freqs, scalar_t *qkv_grad_input,\n const NVTE_QKV_Format qkv_format, const bool interleaved,\n const int cp_size, const int cp_rank, const int s,\n const int b, const int h, const int d, const int d2,\n const int qkv_split_arg_list_0,\n const int qkv_split_arg_list_1,\n const int qkv_split_arg_list_2, cudaStream_t stream) {\n const int THREADS_PER_WARP = 32;\n const int warps_per_block = (h <= 8) ? h : 8;\n dim3 blocks(s, b);\n dim3 threads(THREADS_PER_WARP, warps_per_block);\n const int shared_mem_size = 4 * d2 * sizeof(float); // cos, sin * q ,k\n\n fused_qkv_rope_backward_kernel<<>>(\n q_grad_out, k_grad_out, v_grad_out, q_freqs, k_freqs, qkv_grad_input, qkv_format, interleaved,\n cp_size, cp_rank, s, b, h, d, d2, qkv_split_arg_list_0, qkv_split_arg_list_1,\n qkv_split_arg_list_2);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid fused_rope_forward(const Tensor &input, const Tensor &cu_seqlens, const Tensor &freqs,\n const Tensor &start_positions, Tensor *output,\n const NVTE_QKV_Format qkv_format, const bool interleaved, const int cp_size,\n const int cp_rank, const int s, const int b, const int h, const int d,\n const int d2, const int stride_s_or_t, const int stride_b,\n const int stride_h, const int stride_d, cudaStream_t stream) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.data.dtype, scalar_t,\n fused_rope_forward_launcher(reinterpret_cast(input.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr),\n reinterpret_cast(freqs.data.dptr),\n reinterpret_cast(start_positions.data.dptr),\n reinterpret_cast(output->data.dptr), qkv_format,\n interleaved, cp_size, cp_rank, s, b, h, d, d2, stride_s_or_t,\n stride_b, stride_h, stride_d, stream););\n}\n\nvoid fused_rope_backward(const Tensor &output_grads, const Tensor &cu_seqlens, const Tensor &freqs,\n const Tensor &start_positions, Tensor *input_grads,\n const NVTE_QKV_Format qkv_format, const bool interleaved,\n const int cp_size, const int cp_rank, const int s, const int b,\n const int h, const int d, const int d2, const int stride_s_or_t,\n const int stride_b, const int stride_h, const int stride_d,\n cudaStream_t stream) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n output_grads.data.dtype, scalar_t,\n fused_rope_backward_launcher(reinterpret_cast(output_grads.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr),\n reinterpret_cast(freqs.data.dptr),\n reinterpret_cast(start_positions.data.dptr),\n reinterpret_cast(input_grads->data.dptr), qkv_format,\n interleaved, cp_size, cp_rank, s, b, h, d, d2, stride_s_or_t,\n stride_b, stride_h, stride_d, stream););\n}\n\nvoid fused_qkv_rope_forward(const Tensor &qkv_input, const Tensor &q_freqs, const Tensor &k_freqs,\n const Tensor &start_positions, Tensor *q_out, Tensor *k_out,\n Tensor *v_out, const NVTE_QKV_Format qkv_format, const bool interleaved,\n const int cp_size, const int cp_rank, const int s, const int b,\n const int h, const int d, const int d2, const int qkv_split_arg_list_0,\n const int qkv_split_arg_list_1, const int qkv_split_arg_list_2,\n cudaStream_t stream) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n qkv_input.data.dtype, scalar_t,\n fused_qkv_rope_forward_launcher(reinterpret_cast(qkv_input.data.dptr),\n reinterpret_cast(q_freqs.data.dptr),\n reinterpret_cast(k_freqs.data.dptr),\n reinterpret_cast(start_positions.data.dptr),\n reinterpret_cast(q_out->data.dptr),\n reinterpret_cast(k_out->data.dptr),\n reinterpret_cast(v_out->data.dptr), qkv_format,\n interleaved, cp_size, cp_rank, s, b, h, d, d2,\n qkv_split_arg_list_0, qkv_split_arg_list_1,\n qkv_split_arg_list_2, stream););\n}\n\nvoid fused_qkv_rope_backward(const Tensor &q_grad_out, const Tensor &k_grad_out,\n const Tensor &v_grad_out, const Tensor &q_freqs, const Tensor &k_freqs,\n Tensor *qkv_grad_input, const NVTE_QKV_Format qkv_format,\n const bool interleaved, const int cp_size, const int cp_rank,\n const int s, const int b, const int h, const int d, const int d2,\n const int qkv_split_arg_list_0, const int qkv_split_arg_list_1,\n const int qkv_split_arg_list_2, cudaStream_t stream) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n q_grad_out.data.dtype, scalar_t,\n fused_qkv_rope_backward_launcher(reinterpret_cast(q_grad_out.data.dptr),\n reinterpret_cast(k_grad_out.data.dptr),\n reinterpret_cast(v_grad_out.data.dptr),\n reinterpret_cast(q_freqs.data.dptr),\n reinterpret_cast(k_freqs.data.dptr),\n reinterpret_cast(qkv_grad_input->data.dptr),\n qkv_format, interleaved, cp_size, cp_rank, s, b, h, d, d2,\n qkv_split_arg_list_0, qkv_split_arg_list_1,\n qkv_split_arg_list_2, stream););\n}\n} // end namespace transformer_engine\n\nvoid nvte_fused_rope_forward(const NVTETensor input, const NVTETensor cu_seqlens,\n const NVTETensor freqs, const NVTETensor start_positions,\n NVTETensor output, const NVTE_QKV_Format qkv_format,\n const bool interleaved, const int cp_size, const int cp_rank,\n const int s, const int b, const int h, const int d, const int d2,\n const int stride_s_or_t, const int stride_b, const int stride_h,\n const int stride_d, cudaStream_t stream) {\n NVTE_API_CALL(nvte_fused_rope_forward);\n using namespace transformer_engine;\n fused_rope_forward(*convertNVTETensorCheck(input), *convertNVTETensorCheck(cu_seqlens),\n *convertNVTETensorCheck(freqs), *convertNVTETensorCheck(start_positions),\n convertNVTETensorCheck(output), qkv_format, interleaved, cp_size, cp_rank, s,\n b, h, d, d2, stride_s_or_t, stride_b, stride_h, stride_d, stream);\n}\n\nvoid nvte_fused_rope_backward(const NVTETensor output_grads, const NVTETensor cu_seqlens,\n const NVTETensor freqs, const NVTETensor start_positions,\n NVTETensor input_grads, const NVTE_QKV_Format qkv_format,\n const bool interleaved, const int cp_size, const int cp_rank,\n const int s, const int b, const int h, const int d, const int d2,\n const int stride_s_or_t, const int stride_b, const int stride_h,\n const int stride_d, cudaStream_t stream) {\n NVTE_API_CALL(nvte_fused_rope_backward);\n using namespace transformer_engine;\n fused_rope_backward(*convertNVTETensorCheck(output_grads), *convertNVTETensorCheck(cu_seqlens),\n *convertNVTETensorCheck(freqs), *convertNVTETensorCheck(start_positions),\n convertNVTETensorCheck(input_grads), qkv_format, interleaved, cp_size,\n cp_rank, s, b, h, d, d2, stride_s_or_t, stride_b, stride_h, stride_d, stream);\n}\n\nvoid nvte_fused_qkv_rope_forward(const NVTETensor qkv_input, const NVTETensor q_freqs,\n const NVTETensor k_freqs, const NVTETensor start_positions,\n NVTETensor q_out, NVTETensor k_out, NVTETensor v_out,\n const NVTE_QKV_Format qkv_format, const bool interleaved,\n const int cp_size, const int cp_rank, const int s, const int b,\n const int h, const int d, const int d2,\n const int qkv_split_arg_list_0, const int qkv_split_arg_list_1,\n const int qkv_split_arg_list_2, cudaStream_t stream) {\n NVTE_API_CALL(nvte_fused_qkv_rope_forward);\n using namespace transformer_engine;\n fused_qkv_rope_forward(*convertNVTETensorCheck(qkv_input), *convertNVTETensorCheck(q_freqs),\n *convertNVTETensorCheck(k_freqs), *convertNVTETensorCheck(start_positions),\n convertNVTETensorCheck(q_out), convertNVTETensorCheck(k_out),\n convertNVTETensorCheck(v_out), qkv_format, interleaved, cp_size, cp_rank,\n s, b, h, d, d2, qkv_split_arg_list_0, qkv_split_arg_list_1,\n qkv_split_arg_list_2, stream);\n}\n\nvoid nvte_fused_qkv_rope_backward(const NVTETensor q_grad_out, const NVTETensor k_grad_out,\n const NVTETensor v_grad_out, const NVTETensor q_freqs,\n const NVTETensor k_freqs, NVTETensor qkv_grad_input,\n const NVTE_QKV_Format qkv_format, const bool interleaved,\n const int cp_size, const int cp_rank, const int s, const int b,\n const int h, const int d, const int d2,\n const int qkv_split_arg_list_0, const int qkv_split_arg_list_1,\n const int qkv_split_arg_list_2, cudaStream_t stream) {\n NVTE_API_CALL(nvte_fused_qkv_rope_backward);\n using namespace transformer_engine;\n fused_qkv_rope_backward(*convertNVTETensorCheck(q_grad_out), *convertNVTETensorCheck(k_grad_out),\n *convertNVTETensorCheck(v_grad_out), *convertNVTETensorCheck(q_freqs),\n *convertNVTETensorCheck(k_freqs), convertNVTETensorCheck(qkv_grad_input),\n qkv_format, interleaved, cp_size, cp_rank, s, b, h, d, d2,\n qkv_split_arg_list_0, qkv_split_arg_list_1, qkv_split_arg_list_2, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\n\ntemplate \n__device__ void fused_rope_block_forward(const scalar_t *src, const float *freqs, scalar_t *dst,\n const bool interleaved, const int s_id,\n const int offset_block, const int offset_block_dst,\n const int h, const int d, const int d2, const int stride_h,\n const int stride_d, const int o_stride_h,\n const int o_stride_d) {\n extern __shared__ float shared_mem_cos_sin[];\n float *shared_mem_cos = shared_mem_cos_sin;\n float *shared_mem_sin = shared_mem_cos_sin + d2;\n int tid = threadIdx.x * blockDim.y + threadIdx.y;\n for (int i = tid; i < d2; i += blockDim.x * blockDim.y) {\n sincosf(freqs[s_id * d2 + i], &shared_mem_sin[i], &shared_mem_cos[i]);\n }\n __syncthreads();\n\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n#pragma unroll\n for (int d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) {\n float v_cos = shared_mem_cos[d_id];\n float v_sin = shared_mem_sin[d_id];\n int offset_src = offset_block + h_id * stride_h + d_id * stride_d;\n int offset_dst = offset_block_dst + h_id * o_stride_h + d_id * o_stride_d;\n float v_src = src[offset_src];\n float v_src_rotate;\n if (!interleaved) {\n v_src_rotate = (d_id + d2 / 2 < d2)\n ? -static_cast(src[offset_src + (d2 / 2) * stride_d])\n : static_cast(src[offset_src + (d2 / 2 - d2) * stride_d]);\n } else {\n v_src_rotate = (d_id % 2 == 0)\n // d_id + 1\n ? -static_cast(src[offset_src + stride_d])\n // d_id - 1\n : static_cast(src[offset_src - stride_d]);\n }\n dst[offset_dst] = v_src * v_cos + v_src_rotate * v_sin;\n }\n }\n\n // copy the rest\n if (d > d2) {\n#pragma unroll\n for (int d_id = d2 + threadIdx.x; d_id < d; d_id += blockDim.x) {\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n int offset_src = offset_block + h_id * stride_h + d_id * stride_d;\n int offset_dst = offset_block_dst + h_id * o_stride_h + d_id * o_stride_d;\n dst[offset_dst] = src[offset_src];\n }\n }\n }\n}\n\ntemplate \n__device__ void fused_rope_block_backward(const scalar_t *src, const float *freqs, scalar_t *dst,\n const bool interleaved, const int s_id,\n const int offset_block, const int offset_block_dst,\n const int h, const int d, const int d2,\n const int stride_h, const int stride_d,\n const int o_stride_h, const int o_stride_d) {\n extern __shared__ float shared_mem_cos_sin[];\n float *shared_mem_cos = shared_mem_cos_sin;\n float *shared_mem_sin = shared_mem_cos_sin + d2;\n int tid = threadIdx.x * blockDim.y + threadIdx.y;\n for (int i = tid; i < d2; i += blockDim.x * blockDim.y) {\n sincosf(freqs[s_id * d2 + i], &shared_mem_sin[i], &shared_mem_cos[i]);\n }\n __syncthreads();\n\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n#pragma unroll\n for (int d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) {\n int offset_src = offset_block + h_id * stride_h + d_id * stride_d;\n int offset_dst = offset_block_dst + h_id * o_stride_h + d_id * o_stride_d;\n float v_src = src[offset_src];\n float v_cos = shared_mem_cos[d_id];\n float v_src_rotate, v_sin;\n if (!interleaved) {\n if (d_id + d2 / 2 < d2) {\n v_src_rotate = static_cast(src[offset_src + (d2 / 2) * stride_d]);\n v_sin = shared_mem_sin[d_id + d2 / 2];\n } else {\n v_src_rotate = static_cast(src[offset_src + (d2 / 2 - d2) * stride_d]);\n v_sin = -shared_mem_sin[d_id + d2 / 2 - d2];\n }\n } else {\n if (d_id % 2 == 0) {\n v_src_rotate = static_cast(src[offset_src + stride_d]);\n v_sin = shared_mem_sin[d_id + 1];\n } else {\n v_src_rotate = static_cast(src[offset_src - stride_d]);\n v_sin = -shared_mem_sin[d_id - 1];\n }\n }\n dst[offset_dst] = v_src * v_cos + v_src_rotate * v_sin;\n }\n }\n\n // copy the rest\n if (d > d2) {\n#pragma unroll\n for (int d_id = d2 + threadIdx.x; d_id < d; d_id += blockDim.x) {\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n int offset_src = offset_block + h_id * stride_h + d_id * stride_d;\n int offset_dst = offset_block_dst + h_id * o_stride_h + d_id * o_stride_d;\n dst[offset_dst] = src[offset_src];\n }\n }\n }\n}\n\ntemplate \n__global__ void fused_rope_forward_kernel(const scalar_t *src, const int *cu_seqlens,\n const float *freqs, const int *start_positions,\n scalar_t *dst, const bool interleaved, const int cp_size,\n const int cp_rank, const int s, const int h, const int d,\n const int d2, const int stride_s_or_t, const int stride_b,\n const int stride_h, const int stride_d,\n const int o_stride_s_or_t, const int o_stride_b,\n const int o_stride_h, const int o_stride_d) {\n int s_id = blockIdx.x, b_id = blockIdx.y;\n int offset_block, offset_block_dst;\n int cur_seqlens;\n if (cu_seqlens != nullptr) { // THD\n int start = cu_seqlens[b_id] / cp_size;\n int end = cu_seqlens[b_id + 1] / cp_size;\n int t_id = s_id + start;\n if (t_id >= end) return;\n offset_block = t_id * stride_s_or_t;\n offset_block_dst = t_id * o_stride_s_or_t;\n cur_seqlens = end - start;\n } else { // SBHD/BSHD\n offset_block = s_id * stride_s_or_t + b_id * stride_b;\n offset_block_dst = s_id * o_stride_s_or_t + b_id * o_stride_b;\n cur_seqlens = s;\n }\n\n // Offset the RoPE embedding by start_positions if provided.\n int begin_offset = (start_positions == nullptr) ? 0 : start_positions[b_id];\n int s_id_for_freqs = s_id + begin_offset;\n\n // If CP_SIZE > 1, offset the RoPE embedding by cp_rank based on the dual-chunk order.\n if (cp_size > 1) {\n assert(cur_seqlens % 2 == 0);\n if (s_id < cur_seqlens / 2) {\n s_id_for_freqs += cp_rank * cur_seqlens / 2;\n } else {\n s_id_for_freqs += cur_seqlens * cp_size - (cp_rank + 1) * cur_seqlens / 2 - cur_seqlens / 2;\n }\n }\n\n fused_rope_block_forward(src, freqs, dst, interleaved, s_id_for_freqs, offset_block,\n offset_block_dst, h, d, d2, stride_h, stride_d, o_stride_h, o_stride_d);\n}\n\ntemplate \n__global__ void fused_rope_backward_kernel(\n const scalar_t *src, const int *cu_seqlens, const float *freqs, const int *start_positions,\n scalar_t *dst, const bool interleaved, const int cp_size, const int cp_rank, const int s,\n const int h, const int d, const int d2, const int stride_s_or_t, const int stride_b,\n const int stride_h, const int stride_d, const int o_stride_s_or_t, const int o_stride_b,\n const int o_stride_h, const int o_stride_d) {\n int s_id = blockIdx.x, b_id = blockIdx.y;\n int offset_block, offset_block_dst;\n int cur_seqlens;\n if (cu_seqlens != nullptr) { // THD\n int start = cu_seqlens[b_id] / cp_size;\n int end = cu_seqlens[b_id + 1] / cp_size;\n int t_id = s_id + start;\n if (t_id >= end) return;\n offset_block = t_id * stride_s_or_t;\n offset_block_dst = t_id * o_stride_s_or_t;\n cur_seqlens = end - start;\n } else { // SBHD/BSHD\n offset_block = s_id * stride_s_or_t + b_id * stride_b;\n offset_block_dst = s_id * o_stride_s_or_t + b_id * o_stride_b;\n cur_seqlens = s;\n }\n\n // Offset the RoPE embedding by start_positions if provided.\n int begin_offset = (start_positions == nullptr) ? 0 : start_positions[b_id];\n int s_id_for_freqs = s_id + begin_offset;\n\n // If CP_SIZE > 1, offset the RoPE embedding by cp_rank based on the dual-chunk order.\n if (cp_size > 1) {\n assert(cur_seqlens % 2 == 0);\n if (s_id < cur_seqlens / 2) {\n s_id_for_freqs += cp_rank * cur_seqlens / 2;\n } else {\n s_id_for_freqs += cur_seqlens * cp_size - (cp_rank + 1) * cur_seqlens / 2 - cur_seqlens / 2;\n }\n }\n\n fused_rope_block_backward(src, freqs, dst, interleaved, s_id_for_freqs, offset_block,\n offset_block_dst, h, d, d2, stride_h, stride_d, o_stride_h, o_stride_d);\n}\n\ntemplate \n__device__ void fused_qkv_rope_block_forward(const scalar_t *src, const float *freqs, scalar_t *out,\n const bool interleaved, const int s_id,\n const int offset_block, const int offset_block_dst,\n const int h, const int d, const int d2,\n const int row_offset, const int in_row_length,\n const int out_row_length) {\n extern __shared__ float shared_mem_cos_sin_qk[];\n // Split the shared memory into cos and sin parts for q or k\n float *shared_mem_cos = nullptr;\n float *shared_mem_sin = nullptr;\n if (row_offset == 0) { // q\n shared_mem_cos = shared_mem_cos_sin_qk;\n shared_mem_sin = shared_mem_cos_sin_qk + d2;\n } else { // k\n shared_mem_cos = shared_mem_cos_sin_qk + 2 * d2;\n shared_mem_sin = shared_mem_cos_sin_qk + 3 * d2;\n }\n if (freqs != nullptr) {\n int tid = threadIdx.x * blockDim.y + threadIdx.y;\n for (int i = tid; i < d2; i += blockDim.x * blockDim.y) {\n sincosf(freqs[s_id * d2 + i], &shared_mem_sin[i], &shared_mem_cos[i]);\n }\n }\n __syncthreads();\n\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n#pragma unroll\n for (int i = 0; i < out_row_length; i += d) {\n#pragma unroll\n for (int d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) {\n int offset_src = offset_block + h_id * in_row_length + (row_offset + i) + d_id;\n int offset_dst = offset_block_dst + h_id * out_row_length + i + d_id;\n if (freqs != nullptr) {\n float v_cos, v_sin;\n v_cos = shared_mem_cos[d_id];\n v_sin = shared_mem_sin[d_id];\n float v_src = src[offset_src];\n float v_src_rotate;\n if (!interleaved) {\n v_src_rotate = (d_id + d2 / 2 < d2)\n ? -static_cast(src[offset_src + (d2 / 2)])\n : static_cast(src[offset_src + (d2 / 2 - d2)]);\n } else {\n v_src_rotate = (d_id % 2 == 0) ? -static_cast(src[offset_src + 1])\n : static_cast(src[offset_src - 1]);\n }\n out[offset_dst] = v_src * v_cos + v_src_rotate * v_sin;\n } else {\n out[offset_dst] = src[offset_src];\n }\n }\n }\n }\n // copy the rest\n if (d > d2) {\n#pragma unroll\n for (int d_id = d2 + threadIdx.x; d_id < d; d_id += blockDim.x) {\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n#pragma unroll\n for (int i = 0; i < out_row_length; i += d) {\n int offset_src = offset_block + h_id * in_row_length + (row_offset + i) + d_id;\n int offset_dst = offset_block_dst + h_id * out_row_length + i + d_id;\n out[offset_dst] = src[offset_src];\n }\n }\n }\n }\n}\n\ntemplate \n__device__ void fused_qkv_rope_block_backward(const scalar_t *grad_out, const float *freqs,\n scalar_t *out, const bool interleaved, const int s_id,\n const int offset_block, const int offset_block_dst,\n const int h, const int d, const int d2,\n const int row_offset, const int in_row_length,\n const int out_row_length) {\n extern __shared__ float shared_mem_cos_sin_qk[];\n float *shared_mem_cos = nullptr;\n float *shared_mem_sin = nullptr;\n // Split the shared memory into cos and sin parts for q or k\n if (row_offset == 0) { // q\n shared_mem_cos = shared_mem_cos_sin_qk;\n shared_mem_sin = shared_mem_cos_sin_qk + d2;\n } else { // k\n shared_mem_cos = shared_mem_cos_sin_qk + 2 * d2;\n shared_mem_sin = shared_mem_cos_sin_qk + 3 * d2;\n }\n if (freqs != nullptr) {\n int tid = threadIdx.x * blockDim.y + threadIdx.y;\n for (int i = tid; i < d2; i += blockDim.x * blockDim.y) {\n sincosf(freqs[s_id * d2 + i], &shared_mem_sin[i], &shared_mem_cos[i]);\n }\n }\n __syncthreads();\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n#pragma unroll\n for (int i = 0; i < out_row_length; i += d) {\n#pragma unroll\n for (int d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) {\n int offset_dst = offset_block + h_id * in_row_length + (row_offset + i) + d_id;\n int offset_src = offset_block_dst + h_id * out_row_length + i + d_id;\n\n float v_src = grad_out[offset_src];\n if (freqs != nullptr) {\n float v_cos, v_sin;\n v_cos = shared_mem_cos[d_id];\n float v_src_rotate;\n if (!interleaved) {\n if (d_id + d2 / 2 < d2) {\n v_src_rotate = static_cast(grad_out[offset_src + (d2 / 2)]);\n v_sin = shared_mem_sin[d_id + d2 / 2];\n } else {\n v_src_rotate = static_cast(grad_out[offset_src + (d2 / 2 - d2)]);\n v_sin = -shared_mem_sin[d_id + d2 / 2 - d2];\n }\n } else {\n if (d_id % 2 == 0) {\n v_src_rotate = static_cast(grad_out[offset_src + 1]);\n v_sin = shared_mem_sin[d_id + 1];\n } else {\n v_src_rotate = static_cast(grad_out[offset_src - 1]);\n v_sin = -shared_mem_sin[d_id - 1];\n }\n }\n out[offset_dst] = v_src * v_cos + v_src_rotate * v_sin;\n } else {\n out[offset_dst] = grad_out[offset_src];\n }\n }\n }\n }\n // copy the rest\n if (d > d2) {\n#pragma unroll\n for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {\n#pragma unroll\n for (int i = 0; i < out_row_length; i += d) {\n#pragma unroll\n for (int d_id = d2 + threadIdx.x; d_id < d; d_id += blockDim.x) {\n int offset_dst = offset_block + h_id * in_row_length + (row_offset + i) + d_id;\n int offset_src = offset_block_dst + h_id * out_row_length + i + d_id;\n out[offset_dst] = grad_out[offset_src];\n }\n }\n }\n }\n}\n\ntemplate \n__global__ void fused_qkv_rope_forward_kernel(\n const scalar_t *qkv_input, const float *q_freqs, const float *k_freqs,\n const int *start_positions, scalar_t *q_out, scalar_t *k_out, scalar_t *v_out,\n const NVTE_QKV_Format qkv_format, const bool interleaved, const int cp_size, const int cp_rank,\n const int s, const int b, const int h, const int d, const int d2, const int q_split_arg,\n const int k_split_arg, const int v_split_arg) {\n int s_id = blockIdx.x, b_id = blockIdx.y;\n int cur_seqlens = s;\n int total_d = q_split_arg + k_split_arg + v_split_arg;\n int offset_block, offset_block_dst_q, offset_block_dst_k, offset_block_dst_v;\n if (qkv_format == NVTE_QKV_Format::NVTE_SBHD) {\n offset_block = s_id * b * h * total_d + b_id * h * total_d;\n offset_block_dst_q = s_id * b * h * q_split_arg + b_id * h * q_split_arg;\n offset_block_dst_k = s_id * b * h * k_split_arg + b_id * h * k_split_arg;\n offset_block_dst_v = s_id * b * h * v_split_arg + b_id * h * v_split_arg;\n } else {\n offset_block = b_id * s * h * total_d + s_id * h * total_d;\n offset_block_dst_q = b_id * s * h * q_split_arg + s_id * h * q_split_arg;\n offset_block_dst_k = b_id * s * h * k_split_arg + s_id * h * k_split_arg;\n offset_block_dst_v = b_id * s * h * v_split_arg + s_id * h * v_split_arg;\n }\n\n int q_limit = q_split_arg;\n int k_limit = q_limit + k_split_arg;\n int s_id_for_freqs;\n if (cp_size > 1) {\n assert(cur_seqlens % 2 == 0);\n if (s_id < cur_seqlens / 2) {\n s_id_for_freqs = s_id + cp_rank * cur_seqlens / 2;\n } else {\n s_id_for_freqs =\n cur_seqlens * cp_size - (cp_rank + 1) * cur_seqlens / 2 + s_id - cur_seqlens / 2;\n }\n } else {\n int begin_offset = (start_positions == nullptr) ? 0 : start_positions[b_id];\n s_id_for_freqs = s_id + begin_offset;\n }\n fused_qkv_rope_block_forward(qkv_input, q_freqs, q_out, interleaved, s_id_for_freqs, offset_block,\n offset_block_dst_q, h, d, d2, 0, total_d, q_split_arg);\n fused_qkv_rope_block_forward(qkv_input, k_freqs, k_out, interleaved, s_id_for_freqs, offset_block,\n offset_block_dst_k, h, d, d2, q_limit, total_d, k_split_arg);\n fused_qkv_rope_block_forward(qkv_input, nullptr, v_out, interleaved, s_id_for_freqs, offset_block,\n offset_block_dst_v, h, d, d2, k_limit, total_d, v_split_arg);\n}\n\ntemplate \n__global__ void fused_qkv_rope_backward_kernel(\n const scalar_t *grad_out_q, const scalar_t *grad_out_k, const scalar_t *grad_out_v,\n const float *q_freqs, const float *k_freqs, scalar_t *qkv_grad,\n const NVTE_QKV_Format qkv_format, const bool interleaved, const int cp_size, const int cp_rank,\n const int s, const int b, const int h, const int d, const int d2, const int q_split_arg,\n const int k_split_arg, const int v_split_arg) {\n int s_id = blockIdx.x, b_id = blockIdx.y;\n int cur_seqlens = s;\n int offset_block, offset_block_dst_q, offset_block_dst_k, offset_block_dst_v;\n int total_d = q_split_arg + k_split_arg + v_split_arg;\n if (qkv_format == NVTE_QKV_Format::NVTE_SBHD) {\n offset_block = s_id * b * h * total_d + b_id * h * total_d;\n offset_block_dst_q = s_id * b * h * q_split_arg + b_id * h * q_split_arg;\n offset_block_dst_k = s_id * b * h * k_split_arg + b_id * h * k_split_arg;\n offset_block_dst_v = s_id * b * h * v_split_arg + b_id * h * v_split_arg;\n } else {\n offset_block = b_id * s * h * total_d + s_id * h * total_d;\n offset_block_dst_q = b_id * s * h * q_split_arg + s_id * h * q_split_arg;\n offset_block_dst_k = b_id * s * h * k_split_arg + s_id * h * k_split_arg;\n offset_block_dst_v = b_id * s * h * v_split_arg + s_id * h * v_split_arg;\n }\n int q_limit = q_split_arg;\n int k_limit = q_limit + k_split_arg;\n int s_id_for_freqs;\n if (cp_size > 1) {\n assert(cur_seqlens % 2 == 0);\n if (s_id < cur_seqlens / 2) {\n s_id_for_freqs = s_id + cp_rank * cur_seqlens / 2;\n } else {\n s_id_for_freqs =\n cur_seqlens * cp_size - (cp_rank + 1) * cur_seqlens / 2 + s_id - cur_seqlens / 2;\n }\n } else {\n s_id_for_freqs = s_id;\n }\n fused_qkv_rope_block_backward(grad_out_q, q_freqs, qkv_grad, interleaved, s_id_for_freqs,\n offset_block, offset_block_dst_q, h, d, d2, 0, total_d,\n q_split_arg);\n fused_qkv_rope_block_backward(grad_out_k, k_freqs, qkv_grad, interleaved, s_id_for_freqs,\n offset_block, offset_block_dst_k, h, d, d2, q_limit, total_d,\n k_split_arg);\n fused_qkv_rope_block_backward(grad_out_v, nullptr, qkv_grad, interleaved, s_id_for_freqs,\n offset_block, offset_block_dst_v, h, d, d2, k_limit, total_d,\n v_split_arg);\n}\n\ntemplate \nvoid fused_rope_forward_launcher(const scalar_t *input, const int *cu_seqlens, const float *freqs,\n const int *start_positions, scalar_t *output,\n const NVTE_QKV_Format qkv_format, const bool interleaved,\n const int cp_size, const int cp_rank, const int s, const int b,\n const int h, const int d, const int d2, const int stride_s_or_t,\n const int stride_b, const int stride_h, const int stride_d,\n hipStream_t stream) {\n int warps_per_block = h < 16 ? 4 : 8;\n dim3 blocks(s, b);\n dim3 threads(THREADS_PER_WARP, warps_per_block);\n const int shared_mem_size = 2 * d2 * sizeof(float); // cos, sin\n int o_stride_s_or_t, o_stride_b;\n if (qkv_format == NVTE_QKV_Format::NVTE_THD) {\n NVTE_CHECK(cu_seqlens != nullptr, \"cu_seqlens is required for THD format\");\n o_stride_s_or_t = h * d;\n o_stride_b = 0;\n } else if (qkv_format == NVTE_QKV_Format::NVTE_SBHD) {\n o_stride_s_or_t = b * h * d;\n o_stride_b = h * d;\n } else {\n o_stride_s_or_t = h * d;\n o_stride_b = s * h * d;\n }\n const int o_stride_h = d;\n const int o_stride_d = 1;\n\n fused_rope_forward_kernel<<>>(\n input, cu_seqlens, freqs, start_positions, output, interleaved, cp_size, cp_rank, s, h, d, d2,\n stride_s_or_t, stride_b, stride_h, stride_d, o_stride_s_or_t, o_stride_b, o_stride_h,\n o_stride_d);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\ntemplate \nvoid fused_rope_backward_launcher(const scalar_t *output_grads, const int *cu_seqlens,\n const float *freqs, const int *start_positions,\n scalar_t *input_grads, const NVTE_QKV_Format qkv_format,\n const bool interleaved, const int cp_size, const int cp_rank,\n const int s, const int b, const int h, const int d, const int d2,\n const int stride_s_or_t, const int stride_b, const int stride_h,\n const int stride_d, hipStream_t stream) {\n int warps_per_block = h < 16 ? 4 : 8;\n dim3 blocks(s, b);\n dim3 threads(THREADS_PER_WARP, warps_per_block);\n const int shared_mem_size = 2 * d2 * sizeof(float); // cos, sin\n int o_stride_s_or_t, o_stride_b;\n if (qkv_format == NVTE_QKV_Format::NVTE_THD) {\n NVTE_CHECK(cu_seqlens != nullptr, \"cu_seqlens is required for THD format\");\n o_stride_s_or_t = h * d;\n o_stride_b = 0;\n } else if (qkv_format == NVTE_QKV_Format::NVTE_SBHD) {\n o_stride_s_or_t = b * h * d;\n o_stride_b = h * d;\n } else {\n o_stride_s_or_t = h * d;\n o_stride_b = s * h * d;\n }\n const int o_stride_h = d;\n const int o_stride_d = 1;\n\n fused_rope_backward_kernel<<>>(\n output_grads, cu_seqlens, freqs, start_positions, input_grads, interleaved, cp_size, cp_rank,\n s, h, d, d2, stride_s_or_t, stride_b, stride_h, stride_d, o_stride_s_or_t, o_stride_b,\n o_stride_h, o_stride_d);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\ntemplate \nvoid fused_qkv_rope_forward_launcher(const scalar_t *qkv_input, const float *q_freqs,\n const float *k_freqs, const int *start_positions,\n scalar_t *q_out, scalar_t *k_out, scalar_t *v_out,\n const NVTE_QKV_Format qkv_format, const bool interleaved,\n const int cp_size, const int cp_rank, const int s, const int b,\n const int h, const int d, const int d2,\n const int qkv_split_arg_list_0, const int qkv_split_arg_list_1,\n const int qkv_split_arg_list_2, hipStream_t stream) {\n const int THREADS_PER_WARP = 32;\n int warps_per_block = (h <= 8) ? h : 8;\n dim3 blocks(s, b);\n dim3 threads(THREADS_PER_WARP, warps_per_block);\n const int shared_mem_size = 4 * d2 * sizeof(float); // cos, sin * q ,k\n\n fused_qkv_rope_forward_kernel<<>>(\n qkv_input, q_freqs, k_freqs, start_positions, q_out, k_out, v_out, qkv_format, interleaved,\n cp_size, cp_rank, s, b, h, d, d2, qkv_split_arg_list_0, qkv_split_arg_list_1,\n qkv_split_arg_list_2);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\ntemplate \nvoid fused_qkv_rope_backward_launcher(const scalar_t *q_grad_out, const scalar_t *k_grad_out,\n const scalar_t *v_grad_out, const float *q_freqs,\n const float *k_freqs, scalar_t *qkv_grad_input,\n const NVTE_QKV_Format qkv_format, const bool interleaved,\n const int cp_size, const int cp_rank, const int s,\n const int b, const int h, const int d, const int d2,\n const int qkv_split_arg_list_0,\n const int qkv_split_arg_list_1,\n const int qkv_split_arg_list_2, hipStream_t stream) {\n const int THREADS_PER_WARP = 32;\n const int warps_per_block = (h <= 8) ? h : 8;\n dim3 blocks(s, b);\n dim3 threads(THREADS_PER_WARP, warps_per_block);\n const int shared_mem_size = 4 * d2 * sizeof(float); // cos, sin * q ,k\n\n fused_qkv_rope_backward_kernel<<>>(\n q_grad_out, k_grad_out, v_grad_out, q_freqs, k_freqs, qkv_grad_input, qkv_format, interleaved,\n cp_size, cp_rank, s, b, h, d, d2, qkv_split_arg_list_0, qkv_split_arg_list_1,\n qkv_split_arg_list_2);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid fused_rope_forward(const Tensor &input, const Tensor &cu_seqlens, const Tensor &freqs,\n const Tensor &start_positions, Tensor *output,\n const NVTE_QKV_Format qkv_format, const bool interleaved, const int cp_size,\n const int cp_rank, const int s, const int b, const int h, const int d,\n const int d2, const int stride_s_or_t, const int stride_b,\n const int stride_h, const int stride_d, hipStream_t stream) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.data.dtype, scalar_t,\n fused_rope_forward_launcher(reinterpret_cast(input.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr),\n reinterpret_cast(freqs.data.dptr),\n reinterpret_cast(start_positions.data.dptr),\n reinterpret_cast(output->data.dptr), qkv_format,\n interleaved, cp_size, cp_rank, s, b, h, d, d2, stride_s_or_t,\n stride_b, stride_h, stride_d, stream););\n}\n\nvoid fused_rope_backward(const Tensor &output_grads, const Tensor &cu_seqlens, const Tensor &freqs,\n const Tensor &start_positions, Tensor *input_grads,\n const NVTE_QKV_Format qkv_format, const bool interleaved,\n const int cp_size, const int cp_rank, const int s, const int b,\n const int h, const int d, const int d2, const int stride_s_or_t,\n const int stride_b, const int stride_h, const int stride_d,\n hipStream_t stream) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n output_grads.data.dtype, scalar_t,\n fused_rope_backward_launcher(reinterpret_cast(output_grads.data.dptr),\n reinterpret_cast(cu_seqlens.data.dptr),\n reinterpret_cast(freqs.data.dptr),\n reinterpret_cast(start_positions.data.dptr),\n reinterpret_cast(input_grads->data.dptr), qkv_format,\n interleaved, cp_size, cp_rank, s, b, h, d, d2, stride_s_or_t,\n stride_b, stride_h, stride_d, stream););\n}\n\nvoid fused_qkv_rope_forward(const Tensor &qkv_input, const Tensor &q_freqs, const Tensor &k_freqs,\n const Tensor &start_positions, Tensor *q_out, Tensor *k_out,\n Tensor *v_out, const NVTE_QKV_Format qkv_format, const bool interleaved,\n const int cp_size, const int cp_rank, const int s, const int b,\n const int h, const int d, const int d2, const int qkv_split_arg_list_0,\n const int qkv_split_arg_list_1, const int qkv_split_arg_list_2,\n hipStream_t stream) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n qkv_input.data.dtype, scalar_t,\n fused_qkv_rope_forward_launcher(reinterpret_cast(qkv_input.data.dptr),\n reinterpret_cast(q_freqs.data.dptr),\n reinterpret_cast(k_freqs.data.dptr),\n reinterpret_cast(start_positions.data.dptr),\n reinterpret_cast(q_out->data.dptr),\n reinterpret_cast(k_out->data.dptr),\n reinterpret_cast(v_out->data.dptr), qkv_format,\n interleaved, cp_size, cp_rank, s, b, h, d, d2,\n qkv_split_arg_list_0, qkv_split_arg_list_1,\n qkv_split_arg_list_2, stream););\n}\n\nvoid fused_qkv_rope_backward(const Tensor &q_grad_out, const Tensor &k_grad_out,\n const Tensor &v_grad_out, const Tensor &q_freqs, const Tensor &k_freqs,\n Tensor *qkv_grad_input, const NVTE_QKV_Format qkv_format,\n const bool interleaved, const int cp_size, const int cp_rank,\n const int s, const int b, const int h, const int d, const int d2,\n const int qkv_split_arg_list_0, const int qkv_split_arg_list_1,\n const int qkv_split_arg_list_2, hipStream_t stream) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n q_grad_out.data.dtype, scalar_t,\n fused_qkv_rope_backward_launcher(reinterpret_cast(q_grad_out.data.dptr),\n reinterpret_cast(k_grad_out.data.dptr),\n reinterpret_cast(v_grad_out.data.dptr),\n reinterpret_cast(q_freqs.data.dptr),\n reinterpret_cast(k_freqs.data.dptr),\n reinterpret_cast(qkv_grad_input->data.dptr),\n qkv_format, interleaved, cp_size, cp_rank, s, b, h, d, d2,\n qkv_split_arg_list_0, qkv_split_arg_list_1,\n qkv_split_arg_list_2, stream););\n}\n} // end namespace transformer_engine\n\nvoid nvte_fused_rope_forward(const NVTETensor input, const NVTETensor cu_seqlens,\n const NVTETensor freqs, const NVTETensor start_positions,\n NVTETensor output, const NVTE_QKV_Format qkv_format,\n const bool interleaved, const int cp_size, const int cp_rank,\n const int s, const int b, const int h, const int d, const int d2,\n const int stride_s_or_t, const int stride_b, const int stride_h,\n const int stride_d, hipStream_t stream) {\n NVTE_API_CALL(nvte_fused_rope_forward);\n using namespace transformer_engine;\n fused_rope_forward(*convertNVTETensorCheck(input), *convertNVTETensorCheck(cu_seqlens),\n *convertNVTETensorCheck(freqs), *convertNVTETensorCheck(start_positions),\n convertNVTETensorCheck(output), qkv_format, interleaved, cp_size, cp_rank, s,\n b, h, d, d2, stride_s_or_t, stride_b, stride_h, stride_d, stream);\n}\n\nvoid nvte_fused_rope_backward(const NVTETensor output_grads, const NVTETensor cu_seqlens,\n const NVTETensor freqs, const NVTETensor start_positions,\n NVTETensor input_grads, const NVTE_QKV_Format qkv_format,\n const bool interleaved, const int cp_size, const int cp_rank,\n const int s, const int b, const int h, const int d, const int d2,\n const int stride_s_or_t, const int stride_b, const int stride_h,\n const int stride_d, hipStream_t stream) {\n NVTE_API_CALL(nvte_fused_rope_backward);\n using namespace transformer_engine;\n fused_rope_backward(*convertNVTETensorCheck(output_grads), *convertNVTETensorCheck(cu_seqlens),\n *convertNVTETensorCheck(freqs), *convertNVTETensorCheck(start_positions),\n convertNVTETensorCheck(input_grads), qkv_format, interleaved, cp_size,\n cp_rank, s, b, h, d, d2, stride_s_or_t, stride_b, stride_h, stride_d, stream);\n}\n\nvoid nvte_fused_qkv_rope_forward(const NVTETensor qkv_input, const NVTETensor q_freqs,\n const NVTETensor k_freqs, const NVTETensor start_positions,\n NVTETensor q_out, NVTETensor k_out, NVTETensor v_out,\n const NVTE_QKV_Format qkv_format, const bool interleaved,\n const int cp_size, const int cp_rank, const int s, const int b,\n const int h, const int d, const int d2,\n const int qkv_split_arg_list_0, const int qkv_split_arg_list_1,\n const int qkv_split_arg_list_2, hipStream_t stream) {\n NVTE_API_CALL(nvte_fused_qkv_rope_forward);\n using namespace transformer_engine;\n fused_qkv_rope_forward(*convertNVTETensorCheck(qkv_input), *convertNVTETensorCheck(q_freqs),\n *convertNVTETensorCheck(k_freqs), *convertNVTETensorCheck(start_positions),\n convertNVTETensorCheck(q_out), convertNVTETensorCheck(k_out),\n convertNVTETensorCheck(v_out), qkv_format, interleaved, cp_size, cp_rank,\n s, b, h, d, d2, qkv_split_arg_list_0, qkv_split_arg_list_1,\n qkv_split_arg_list_2, stream);\n}\n\nvoid nvte_fused_qkv_rope_backward(const NVTETensor q_grad_out, const NVTETensor k_grad_out,\n const NVTETensor v_grad_out, const NVTETensor q_freqs,\n const NVTETensor k_freqs, NVTETensor qkv_grad_input,\n const NVTE_QKV_Format qkv_format, const bool interleaved,\n const int cp_size, const int cp_rank, const int s, const int b,\n const int h, const int d, const int d2,\n const int qkv_split_arg_list_0, const int qkv_split_arg_list_1,\n const int qkv_split_arg_list_2, hipStream_t stream) {\n NVTE_API_CALL(nvte_fused_qkv_rope_backward);\n using namespace transformer_engine;\n fused_qkv_rope_backward(*convertNVTETensorCheck(q_grad_out), *convertNVTETensorCheck(k_grad_out),\n *convertNVTETensorCheck(v_grad_out), *convertNVTETensorCheck(q_freqs),\n *convertNVTETensorCheck(k_freqs), convertNVTETensorCheck(qkv_grad_input),\n qkv_format, interleaved, cp_size, cp_rank, s, b, h, d, d2,\n qkv_split_arg_list_0, qkv_split_arg_list_1, qkv_split_arg_list_2, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/fused_rope/fused_rope.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n\n#include \n\n#include \"../common.h\"\n\nstatic __global__ void moe_permute_row_map(const int *sorted_row_id, int *row_id_map,\n const int num_rows, const int topK,\n const int num_out_tokens) {\n // Each block corresponds to one source token\n // row_id_map[topK][num_rows]\n const int bid = blockIdx.x;\n const int tid = threadIdx.x;\n const int idx = bid * blockDim.x + tid;\n\n if (idx >= num_rows * topK) return;\n\n int source_row = sorted_row_id[idx];\n int source_token_id = source_row / topK;\n int source_topK_id = source_row % topK;\n\n if (idx >= num_out_tokens) {\n // Set the indices of dropped tokens to -1\n row_id_map[source_topK_id * num_rows + source_token_id] = -1;\n } else {\n // Create a row id map for subsequent unpermute operation\n row_id_map[source_topK_id * num_rows + source_token_id] = idx;\n }\n}\n\ntemplate \n__global__ void moe_unpermute_kernel(const T *input, T *unpermuted_output, const int *row_id_map,\n const float *prob, const int num_rows, const int topK,\n const int num_cols) {\n extern __shared__ int8_t s_mem[];\n TCompute *s_prob = reinterpret_cast(s_mem);\n\n // Each block corresponds to one dest token\n const int source_token = blockIdx.x;\n const int tid = threadIdx.x;\n\n if (hasProb) {\n for (int i = tid; i < topK; i += blockDim.x * blockDim.y) {\n // Load all the topK probs related to the source row into smem\n s_prob[i] = TCompute(prob[source_token * topK + i]);\n }\n __syncthreads();\n }\n\n // Register buffers for vector type (float4) memory access\n float4 frag_load_store;\n T *frag_load_store_ptr = reinterpret_cast(&frag_load_store);\n\n // Number of elemments in frag_load_store\n static constexpr int kElementsPerAccess = 16 / sizeof(T);\n\n // Traverse along the hidden dimention\n for (int i = tid * kElementsPerAccess; i < num_cols; i += blockDim.x * kElementsPerAccess) {\n TCompute frag_elem[kElementsPerAccess];\n TCompute frag_sum[kElementsPerAccess];\n\n int source_row = row_id_map[source_token];\n\n // source_row == -1 represents a dropped token\n if (source_row != -1) {\n const T *source_row_ptr = input + source_row * num_cols;\n\n frag_load_store = __ldlu(reinterpret_cast(source_row_ptr + i));\n\n for (int e = 0; e < kElementsPerAccess; e++) {\n frag_sum[e] = TCompute(frag_load_store_ptr[e]);\n }\n\n if (hasProb) {\n for (int e = 0; e < kElementsPerAccess; e++) {\n frag_sum[e] = frag_sum[e] * s_prob[0];\n }\n }\n } else {\n for (int e = 0; e < kElementsPerAccess; e++) {\n frag_sum[e] = TCompute(0.0f);\n }\n }\n\n for (int k = 1; k < topK; k++) {\n source_row = row_id_map[k * num_rows + source_token];\n\n if (source_row == -1) continue;\n\n const T *source_row_ptr = input + source_row * num_cols;\n\n frag_load_store = __ldlu(reinterpret_cast(source_row_ptr + i));\n\n for (int e = 0; e < kElementsPerAccess; e++) {\n frag_elem[e] = TCompute(frag_load_store_ptr[e]);\n }\n\n if (hasProb) {\n for (int e = 0; e < kElementsPerAccess; e++) {\n frag_elem[e] = frag_elem[e] * s_prob[k];\n }\n }\n\n for (int e = 0; e < kElementsPerAccess; e++) {\n frag_sum[e] = frag_sum[e] + frag_elem[e];\n }\n }\n\n T *dest_row_ptr = unpermuted_output + source_token * num_cols;\n\n for (int e = 0; e < kElementsPerAccess; e++) {\n if constexpr ((std::is_same_v || std::is_same_v) &&\n (!hasProb)) {\n frag_sum[e] = frag_sum[e] / TCompute(topK);\n }\n frag_load_store_ptr[e] = T(frag_sum[e]);\n }\n\n *reinterpret_cast(dest_row_ptr + i) = frag_load_store;\n }\n}\n\ntemplate \n__global__ void moe_permute_kernel(const T *input_bwd, const T *input_fwd, T *act_grad,\n const float *prob, float *prob_grad, const int *row_id_map,\n const int num_rows, const int topK, const int num_cols) {\n extern __shared__ int8_t s_mem[];\n TCompute *s_prob = reinterpret_cast(s_mem);\n\n // Each block corresponds to one source token\n const int source_token = blockIdx.x;\n const int tid = threadIdx.x;\n\n if (hasProb) {\n for (int i = tid; i < topK; i += blockDim.x) {\n // Load all the topK probs related to the source row into smem\n s_prob[i] = TCompute(prob[source_token * topK + i]);\n }\n __syncthreads();\n }\n\n // Accumulators for the calculation of prob_grad\n float accum[topKTile] = {0.0f};\n\n // Register buffers for vector type (float4) memory access\n float4 frag_load_store;\n T *frag_load_store_ptr = reinterpret_cast(&frag_load_store);\n\n // Number of elemments in frag_load_store\n static constexpr int kElementsPerAccess = 16 / sizeof(T);\n\n // The starting address of each source row\n const T *source_row_ptr = input_bwd + source_token * num_cols;\n\n // Traverse along the hidden dimention\n for (int i = tid * kElementsPerAccess; i < num_cols; i += blockDim.x * kElementsPerAccess) {\n TCompute frag_src[kElementsPerAccess];\n\n frag_load_store = __ldlu(reinterpret_cast(source_row_ptr + i));\n\n for (int e = 0; e < kElementsPerAccess; e++) frag_src[e] = TCompute(frag_load_store_ptr[e]);\n\n int index = source_token;\n\n // Process each row in the corresponding topK rows\n for (int k = 0; k < topKTile; k++) {\n if (k == topK) break;\n\n int dest_row = row_id_map[index];\n index += num_rows;\n\n if (dest_row != -1) {\n if (hasProb) {\n // Calculate act_grad in unpermute bwd\n for (int e = 0; e < kElementsPerAccess; e++)\n frag_load_store_ptr[e] = T(frag_src[e] * s_prob[k]);\n } else {\n // permute fwd\n for (int e = 0; e < kElementsPerAccess; e++) frag_load_store_ptr[e] = T(frag_src[e]);\n }\n\n T *dest_row_ptr = act_grad + dest_row * num_cols;\n *reinterpret_cast(dest_row_ptr + i) = frag_load_store;\n\n if (hasProb) {\n // Inner product calculation for prob_grad in unpermute bwd\n const T *input_fwd_ptr = input_fwd + dest_row * num_cols;\n\n frag_load_store = __ldlu(reinterpret_cast(input_fwd_ptr + i));\n\n TCompute frag_input_fwd[kElementsPerAccess];\n for (int e = 0; e < kElementsPerAccess; e++)\n frag_input_fwd[e] = TCompute(frag_load_store_ptr[e]);\n\n for (int e = 0; e < kElementsPerAccess; e++) {\n accum[k] += static_cast(frag_src[e] * frag_input_fwd[e]);\n }\n }\n }\n }\n }\n\n if (hasProb) {\n for (int k = 0; k < topKTile; k++) {\n if (k == topK) break;\n // Warp-level reduction\n for (int mask = 16; mask > 0; mask /= 2) {\n accum[k] = accum[k] + __shfl_xor_sync(0xffffffff, accum[k], mask, 32);\n }\n }\n\n if (tid == 0) {\n for (int k = 0; k < topKTile; k++) {\n if (k == topK) break;\n prob_grad[source_token * topK + k] = accum[k];\n }\n }\n }\n}\n\ntemplate \nvoid nvte_permute_launcher(const T *input, T *output, const int *sorted_row_id, int *row_id_map,\n const float *prob, float *prob_grad, const T *input_fwd,\n const int num_rows, const int topK, const int num_cols,\n const int num_out_tokens, cudaStream_t stream) {\n using TCompute = typename std::conditional<(std::is_same::value ||\n std::is_same::value),\n half, T>::type;\n\n static constexpr int kElementsPerAccess = 16 / sizeof(T);\n\n if (input_fwd == nullptr) {\n // moe_permute_fwd\n\n int threads = 64;\n int blocks = (num_rows * topK + threads - 1) / threads;\n\n moe_permute_row_map<<>>(sorted_row_id, row_id_map, num_rows, topK,\n num_out_tokens);\n NVTE_CHECK_CUDA(cudaGetLastError());\n\n blocks = num_rows;\n threads = std::min(num_cols / kElementsPerAccess, 1024);\n moe_permute_kernel<<>>(\n input, nullptr, output, nullptr, nullptr, row_id_map, num_rows, topK, num_cols);\n NVTE_CHECK_CUDA(cudaGetLastError());\n } else {\n // moe_unpermute_bwd\n\n int threads = 32;\n int blocks = num_rows;\n\n if (prob == nullptr) {\n // moe_unpermute_bwd without probs\n\n moe_permute_kernel<<>>(\n input, input_fwd, output, nullptr, nullptr, row_id_map, num_rows, topK, num_cols);\n NVTE_CHECK_CUDA(cudaGetLastError());\n } else {\n // moe_unpermute_bwd with probs\n\n size_t smem_bytes = topK * sizeof(TCompute);\n\n if (topK <= 8) {\n moe_permute_kernel<<>>(\n input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, topK, num_cols);\n } else if (topK <= 16) {\n moe_permute_kernel<<>>(\n input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, topK, num_cols);\n } else if (topK <= 32) {\n moe_permute_kernel<<>>(\n input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, topK, num_cols);\n } else if (topK <= 64) {\n moe_permute_kernel<<>>(\n input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, topK, num_cols);\n } else if (topK <= 128) {\n moe_permute_kernel<<>>(\n input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, topK, num_cols);\n } else {\n NVTE_ERROR(\"topK cannot exceed 128.\");\n }\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n }\n}\n\ntemplate \nvoid nvte_unpermute_launcher(const T *input, T *output, int *row_id_map, const float *prob,\n const int num_rows, const int topK, const int num_cols,\n cudaStream_t stream) {\n using TCompute = typename std::conditional<(std::is_same::value ||\n std::is_same::value),\n half, T>::type;\n\n static constexpr int kElementsPerAccess = 16 / sizeof(T);\n\n int blocks = num_rows;\n int threads = std::min(num_cols / kElementsPerAccess, 1024);\n size_t smem_bytes = topK * sizeof(TCompute);\n\n if (prob == nullptr) {\n // moe_permute_bwd\n // moe_unpermute_fwd without probs\n\n moe_unpermute_kernel<<>>(\n input, output, row_id_map, nullptr, num_rows, topK, num_cols);\n NVTE_CHECK_CUDA(cudaGetLastError());\n } else {\n // moe_unpermute_fwd with probs\n\n moe_unpermute_kernel<<>>(\n input, output, row_id_map, prob, num_rows, topK, num_cols);\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\nvoid nvte_permute(const NVTETensor input, NVTETensor output, const NVTETensor sorted_row_id,\n NVTETensor row_id_map, const NVTETensor prob, NVTETensor prob_grad,\n const NVTETensor input_fwd, const int num_rows, const int topK,\n const int num_cols, const int num_out_tokens, cudaStream_t stream) {\n using namespace transformer_engine;\n NVTE_API_CALL(nvte_permute);\n\n const Tensor *input_cu = convertNVTETensorCheck(input);\n const Tensor *output_cu = convertNVTETensorCheck(output);\n const Tensor *sorted_row_id_cu = convertNVTETensorCheck(sorted_row_id);\n const Tensor *row_id_map_cu = convertNVTETensorCheck(row_id_map);\n const Tensor *prob_cu = convertNVTETensorCheck(prob);\n const Tensor *prob_grad_cu = convertNVTETensorCheck(prob_grad);\n const Tensor *input_fwd_cu = convertNVTETensorCheck(input_fwd);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n input_cu->data.dtype, T,\n nvte_permute_launcher(reinterpret_cast(input_cu->data.dptr),\n reinterpret_cast(output_cu->data.dptr),\n reinterpret_cast(sorted_row_id_cu->data.dptr),\n reinterpret_cast(row_id_map_cu->data.dptr),\n reinterpret_cast(prob_cu->data.dptr),\n reinterpret_cast(prob_grad_cu->data.dptr),\n reinterpret_cast(input_fwd_cu->data.dptr), num_rows, topK,\n num_cols, num_out_tokens, stream););\n}\n\nvoid nvte_unpermute(const NVTETensor input, NVTETensor output, NVTETensor row_id_map,\n const NVTETensor prob, const int num_rows, const int topK, const int num_cols,\n cudaStream_t stream) {\n using namespace transformer_engine;\n NVTE_API_CALL(nvte_unpermute);\n\n const Tensor *input_cu = convertNVTETensorCheck(input);\n const Tensor *output_cu = convertNVTETensorCheck(output);\n const Tensor *row_id_map_cu = convertNVTETensorCheck(row_id_map);\n const Tensor *prob_cu = convertNVTETensorCheck(prob);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n input_cu->data.dtype, T,\n nvte_unpermute_launcher(reinterpret_cast(input_cu->data.dptr),\n reinterpret_cast(output_cu->data.dptr),\n reinterpret_cast(row_id_map_cu->data.dptr),\n reinterpret_cast(prob_cu->data.dptr), num_rows, topK,\n num_cols, stream););\n}\n\nvoid nvte_device_radix_sort_pairs(void *temp_storage, size_t *temp_storage_bytes, int *keys_in,\n int *keys_out, int *values_in, int *values_out,\n size_t num_items) {\n NVTE_API_CALL(nvte_device_radix_sort_pairs);\n cub::DeviceRadixSort::SortPairs(temp_storage, *temp_storage_bytes, keys_in, keys_out, values_in,\n values_out, num_items);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \n\n#include \"../common.h\"\n\nstatic __global__ void moe_permute_row_map(const int *sorted_row_id, int *row_id_map,\n const int num_rows, const int topK,\n const int num_out_tokens) {\n // Each block corresponds to one source token\n // row_id_map[topK][num_rows]\n const int bid = blockIdx.x;\n const int tid = threadIdx.x;\n const int idx = bid * blockDim.x + tid;\n\n if (idx >= num_rows * topK) return;\n\n int source_row = sorted_row_id[idx];\n int source_token_id = source_row / topK;\n int source_topK_id = source_row % topK;\n\n if (idx >= num_out_tokens) {\n // Set the indices of dropped tokens to -1\n row_id_map[source_topK_id * num_rows + source_token_id] = -1;\n } else {\n // Create a row id map for subsequent unpermute operation\n row_id_map[source_topK_id * num_rows + source_token_id] = idx;\n }\n}\n\ntemplate \n__global__ void moe_unpermute_kernel(const T *input, T *unpermuted_output, const int *row_id_map,\n const float *prob, const int num_rows, const int topK,\n const int num_cols) {\n extern __shared__ int8_t s_mem[];\n TCompute *s_prob = reinterpret_cast(s_mem);\n\n // Each block corresponds to one dest token\n const int source_token = blockIdx.x;\n const int tid = threadIdx.x;\n\n if (hasProb) {\n for (int i = tid; i < topK; i += blockDim.x * blockDim.y) {\n // Load all the topK probs related to the source row into smem\n s_prob[i] = TCompute(prob[source_token * topK + i]);\n }\n __syncthreads();\n }\n\n // Register buffers for vector type (float4) memory access\n float4 frag_load_store;\n T *frag_load_store_ptr = reinterpret_cast(&frag_load_store);\n\n // Number of elemments in frag_load_store\n static constexpr int kElementsPerAccess = 16 / sizeof(T);\n\n // Traverse along the hidden dimention\n for (int i = tid * kElementsPerAccess; i < num_cols; i += blockDim.x * kElementsPerAccess) {\n TCompute frag_elem[kElementsPerAccess];\n TCompute frag_sum[kElementsPerAccess];\n\n int source_row = row_id_map[source_token];\n\n // source_row == -1 represents a dropped token\n if (source_row != -1) {\n const T *source_row_ptr = input + source_row * num_cols;\n\n frag_load_store = __ldlu(reinterpret_cast(source_row_ptr + i));\n\n for (int e = 0; e < kElementsPerAccess; e++) {\n frag_sum[e] = TCompute(frag_load_store_ptr[e]);\n }\n\n if (hasProb) {\n for (int e = 0; e < kElementsPerAccess; e++) {\n frag_sum[e] = frag_sum[e] * s_prob[0];\n }\n }\n } else {\n for (int e = 0; e < kElementsPerAccess; e++) {\n frag_sum[e] = TCompute(0.0f);\n }\n }\n\n for (int k = 1; k < topK; k++) {\n source_row = row_id_map[k * num_rows + source_token];\n\n if (source_row == -1) continue;\n\n const T *source_row_ptr = input + source_row * num_cols;\n\n frag_load_store = __ldlu(reinterpret_cast(source_row_ptr + i));\n\n for (int e = 0; e < kElementsPerAccess; e++) {\n frag_elem[e] = TCompute(frag_load_store_ptr[e]);\n }\n\n if (hasProb) {\n for (int e = 0; e < kElementsPerAccess; e++) {\n frag_elem[e] = frag_elem[e] * s_prob[k];\n }\n }\n\n for (int e = 0; e < kElementsPerAccess; e++) {\n frag_sum[e] = frag_sum[e] + frag_elem[e];\n }\n }\n\n T *dest_row_ptr = unpermuted_output + source_token * num_cols;\n\n for (int e = 0; e < kElementsPerAccess; e++) {\n if constexpr ((std::is_same_v || std::is_same_v) &&\n (!hasProb)) {\n frag_sum[e] = frag_sum[e] / TCompute(topK);\n }\n frag_load_store_ptr[e] = T(frag_sum[e]);\n }\n\n *reinterpret_cast(dest_row_ptr + i) = frag_load_store;\n }\n}\n\ntemplate \n__global__ void moe_permute_kernel(const T *input_bwd, const T *input_fwd, T *act_grad,\n const float *prob, float *prob_grad, const int *row_id_map,\n const int num_rows, const int topK, const int num_cols) {\n extern __shared__ int8_t s_mem[];\n TCompute *s_prob = reinterpret_cast(s_mem);\n\n // Each block corresponds to one source token\n const int source_token = blockIdx.x;\n const int tid = threadIdx.x;\n\n if (hasProb) {\n for (int i = tid; i < topK; i += blockDim.x) {\n // Load all the topK probs related to the source row into smem\n s_prob[i] = TCompute(prob[source_token * topK + i]);\n }\n __syncthreads();\n }\n\n // Accumulators for the calculation of prob_grad\n float accum[topKTile] = {0.0f};\n\n // Register buffers for vector type (float4) memory access\n float4 frag_load_store;\n T *frag_load_store_ptr = reinterpret_cast(&frag_load_store);\n\n // Number of elemments in frag_load_store\n static constexpr int kElementsPerAccess = 16 / sizeof(T);\n\n // The starting address of each source row\n const T *source_row_ptr = input_bwd + source_token * num_cols;\n\n // Traverse along the hidden dimention\n for (int i = tid * kElementsPerAccess; i < num_cols; i += blockDim.x * kElementsPerAccess) {\n TCompute frag_src[kElementsPerAccess];\n\n frag_load_store = __ldlu(reinterpret_cast(source_row_ptr + i));\n\n for (int e = 0; e < kElementsPerAccess; e++) frag_src[e] = TCompute(frag_load_store_ptr[e]);\n\n int index = source_token;\n\n // Process each row in the corresponding topK rows\n for (int k = 0; k < topKTile; k++) {\n if (k == topK) break;\n\n int dest_row = row_id_map[index];\n index += num_rows;\n\n if (dest_row != -1) {\n if (hasProb) {\n // Calculate act_grad in unpermute bwd\n for (int e = 0; e < kElementsPerAccess; e++)\n frag_load_store_ptr[e] = T(frag_src[e] * s_prob[k]);\n } else {\n // permute fwd\n for (int e = 0; e < kElementsPerAccess; e++) frag_load_store_ptr[e] = T(frag_src[e]);\n }\n\n T *dest_row_ptr = act_grad + dest_row * num_cols;\n *reinterpret_cast(dest_row_ptr + i) = frag_load_store;\n\n if (hasProb) {\n // Inner product calculation for prob_grad in unpermute bwd\n const T *input_fwd_ptr = input_fwd + dest_row * num_cols;\n\n frag_load_store = __ldlu(reinterpret_cast(input_fwd_ptr + i));\n\n TCompute frag_input_fwd[kElementsPerAccess];\n for (int e = 0; e < kElementsPerAccess; e++)\n frag_input_fwd[e] = TCompute(frag_load_store_ptr[e]);\n\n for (int e = 0; e < kElementsPerAccess; e++) {\n accum[k] += static_cast(frag_src[e] * frag_input_fwd[e]);\n }\n }\n }\n }\n }\n\n if (hasProb) {\n for (int k = 0; k < topKTile; k++) {\n if (k == topK) break;\n // Warp-level reduction\n for (int mask = 16; mask > 0; mask /= 2) {\n accum[k] = accum[k] + __shfl_xor_sync(0xffffffff, accum[k], mask, 32);\n }\n }\n\n if (tid == 0) {\n for (int k = 0; k < topKTile; k++) {\n if (k == topK) break;\n prob_grad[source_token * topK + k] = accum[k];\n }\n }\n }\n}\n\ntemplate \nvoid nvte_permute_launcher(const T *input, T *output, const int *sorted_row_id, int *row_id_map,\n const float *prob, float *prob_grad, const T *input_fwd,\n const int num_rows, const int topK, const int num_cols,\n const int num_out_tokens, hipStream_t stream) {\n using TCompute = typename std::conditional<(std::is_same::value ||\n std::is_same::value),\n half, T>::type;\n\n static constexpr int kElementsPerAccess = 16 / sizeof(T);\n\n if (input_fwd == nullptr) {\n // moe_permute_fwd\n\n int threads = 64;\n int blocks = (num_rows * topK + threads - 1) / threads;\n\n moe_permute_row_map<<>>(sorted_row_id, row_id_map, num_rows, topK,\n num_out_tokens);\n NVTE_CHECK_CUDA(hipGetLastError());\n\n blocks = num_rows;\n threads = std::min(num_cols / kElementsPerAccess, 1024);\n moe_permute_kernel<<>>(\n input, nullptr, output, nullptr, nullptr, row_id_map, num_rows, topK, num_cols);\n NVTE_CHECK_CUDA(hipGetLastError());\n } else {\n // moe_unpermute_bwd\n\n int threads = 32;\n int blocks = num_rows;\n\n if (prob == nullptr) {\n // moe_unpermute_bwd without probs\n\n moe_permute_kernel<<>>(\n input, input_fwd, output, nullptr, nullptr, row_id_map, num_rows, topK, num_cols);\n NVTE_CHECK_CUDA(hipGetLastError());\n } else {\n // moe_unpermute_bwd with probs\n\n size_t smem_bytes = topK * sizeof(TCompute);\n\n if (topK <= 8) {\n moe_permute_kernel<<>>(\n input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, topK, num_cols);\n } else if (topK <= 16) {\n moe_permute_kernel<<>>(\n input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, topK, num_cols);\n } else if (topK <= 32) {\n moe_permute_kernel<<>>(\n input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, topK, num_cols);\n } else if (topK <= 64) {\n moe_permute_kernel<<>>(\n input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, topK, num_cols);\n } else if (topK <= 128) {\n moe_permute_kernel<<>>(\n input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, topK, num_cols);\n } else {\n NVTE_ERROR(\"topK cannot exceed 128.\");\n }\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n }\n}\n\ntemplate \nvoid nvte_unpermute_launcher(const T *input, T *output, int *row_id_map, const float *prob,\n const int num_rows, const int topK, const int num_cols,\n hipStream_t stream) {\n using TCompute = typename std::conditional<(std::is_same::value ||\n std::is_same::value),\n half, T>::type;\n\n static constexpr int kElementsPerAccess = 16 / sizeof(T);\n\n int blocks = num_rows;\n int threads = std::min(num_cols / kElementsPerAccess, 1024);\n size_t smem_bytes = topK * sizeof(TCompute);\n\n if (prob == nullptr) {\n // moe_permute_bwd\n // moe_unpermute_fwd without probs\n\n moe_unpermute_kernel<<>>(\n input, output, row_id_map, nullptr, num_rows, topK, num_cols);\n NVTE_CHECK_CUDA(hipGetLastError());\n } else {\n // moe_unpermute_fwd with probs\n\n moe_unpermute_kernel<<>>(\n input, output, row_id_map, prob, num_rows, topK, num_cols);\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\nvoid nvte_permute(const NVTETensor input, NVTETensor output, const NVTETensor sorted_row_id,\n NVTETensor row_id_map, const NVTETensor prob, NVTETensor prob_grad,\n const NVTETensor input_fwd, const int num_rows, const int topK,\n const int num_cols, const int num_out_tokens, hipStream_t stream) {\n using namespace transformer_engine;\n NVTE_API_CALL(nvte_permute);\n\n const Tensor *input_cu = convertNVTETensorCheck(input);\n const Tensor *output_cu = convertNVTETensorCheck(output);\n const Tensor *sorted_row_id_cu = convertNVTETensorCheck(sorted_row_id);\n const Tensor *row_id_map_cu = convertNVTETensorCheck(row_id_map);\n const Tensor *prob_cu = convertNVTETensorCheck(prob);\n const Tensor *prob_grad_cu = convertNVTETensorCheck(prob_grad);\n const Tensor *input_fwd_cu = convertNVTETensorCheck(input_fwd);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n input_cu->data.dtype, T,\n nvte_permute_launcher(reinterpret_cast(input_cu->data.dptr),\n reinterpret_cast(output_cu->data.dptr),\n reinterpret_cast(sorted_row_id_cu->data.dptr),\n reinterpret_cast(row_id_map_cu->data.dptr),\n reinterpret_cast(prob_cu->data.dptr),\n reinterpret_cast(prob_grad_cu->data.dptr),\n reinterpret_cast(input_fwd_cu->data.dptr), num_rows, topK,\n num_cols, num_out_tokens, stream););\n}\n\nvoid nvte_unpermute(const NVTETensor input, NVTETensor output, NVTETensor row_id_map,\n const NVTETensor prob, const int num_rows, const int topK, const int num_cols,\n hipStream_t stream) {\n using namespace transformer_engine;\n NVTE_API_CALL(nvte_unpermute);\n\n const Tensor *input_cu = convertNVTETensorCheck(input);\n const Tensor *output_cu = convertNVTETensorCheck(output);\n const Tensor *row_id_map_cu = convertNVTETensorCheck(row_id_map);\n const Tensor *prob_cu = convertNVTETensorCheck(prob);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n input_cu->data.dtype, T,\n nvte_unpermute_launcher(reinterpret_cast(input_cu->data.dptr),\n reinterpret_cast(output_cu->data.dptr),\n reinterpret_cast(row_id_map_cu->data.dptr),\n reinterpret_cast(prob_cu->data.dptr), num_rows, topK,\n num_cols, stream););\n}\n\nvoid nvte_device_radix_sort_pairs(void *temp_storage, size_t *temp_storage_bytes, int *keys_in,\n int *keys_out, int *values_in, int *values_out,\n size_t num_items) {\n NVTE_API_CALL(nvte_device_radix_sort_pairs);\n hipcub::DeviceRadixSort::SortPairs(temp_storage, *temp_storage_bytes, keys_in, keys_out, values_in,\n values_out, num_items);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/permutation/permutation.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"../util/cuda_driver.h\"\n#include \"../util/logging.h\"\n#include \"nvshmem_waitkernel.h\"\n\n__global__ void __launch_bounds__(1)\n wait_until_on_stream_and_reset(uint64_t* wait_flag, uint64_t wait_value,\n uint64_t signal_reset) {\n nvshmem_uint64_wait_until(wait_flag, NVSHMEM_CMP_EQ, wait_value);\n *wait_flag = signal_reset;\n}\nvoid nvshmem_wait_on_stream(uint64_t* sig_addr, WaitKind wait_kind, cudaStream_t stream) {\n uint64_t wait_value = 1;\n uint64_t signal_reset = 0;\n cudaStream_t cur_stream = stream;\n\n NVTE_CHECK(wait_kind >= WaitKind::KERNEL_WAIT && wait_kind <= WaitKind::STREAM_WAIT,\n \"Invalid wait kind: \", static_cast(wait_kind));\n\n switch (wait_kind) {\n case WaitKind::KERNEL_WAIT:\n wait_until_on_stream_and_reset<<<1, 1, 0, cur_stream>>>(sig_addr, wait_value, signal_reset);\n NVTE_CHECK_CUDA(cudaGetLastError());\n break;\n case WaitKind::NVSHMEM_WAIT:\n nvshmemx_uint64_wait_until_on_stream(sig_addr, NVSHMEM_CMP_EQ, wait_value, cur_stream);\n NVTE_CHECK_CUDA_DRIVER(cuStreamWriteValue64((CUstream)cur_stream, (CUdeviceptr)sig_addr,\n (cuuint64_t)signal_reset,\n CU_STREAM_WRITE_VALUE_DEFAULT));\n break;\n case WaitKind::STREAM_WAIT:\n NVTE_CHECK_CUDA_DRIVER(cuStreamWaitValue64((CUstream)cur_stream, (CUdeviceptr)sig_addr,\n (cuuint64_t)wait_value, CU_STREAM_WAIT_VALUE_GEQ));\n NVTE_CHECK_CUDA_DRIVER(cuStreamWriteValue64((CUstream)cur_stream, (CUdeviceptr)sig_addr,\n (cuuint64_t)signal_reset,\n CU_STREAM_WRITE_VALUE_DEFAULT));\n break;\n }\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"../util/cuda_driver.h\"\n#include \"../util/logging.h\"\n#include \"nvshmem_waitkernel.h\"\n\n__global__ void __launch_bounds__(1)\n wait_until_on_stream_and_reset(uint64_t* wait_flag, uint64_t wait_value,\n uint64_t signal_reset) {\n nvshmem_uint64_wait_until(wait_flag, NVSHMEM_CMP_EQ, wait_value);\n *wait_flag = signal_reset;\n}\nvoid nvshmem_wait_on_stream(uint64_t* sig_addr, WaitKind wait_kind, hipStream_t stream) {\n uint64_t wait_value = 1;\n uint64_t signal_reset = 0;\n hipStream_t cur_stream = stream;\n\n NVTE_CHECK(wait_kind >= WaitKind::KERNEL_WAIT && wait_kind <= WaitKind::STREAM_WAIT,\n \"Invalid wait kind: \", static_cast(wait_kind));\n\n switch (wait_kind) {\n case WaitKind::KERNEL_WAIT:\n wait_until_on_stream_and_reset<<<1, 1, 0, cur_stream>>>(sig_addr, wait_value, signal_reset);\n NVTE_CHECK_CUDA(hipGetLastError());\n break;\n case WaitKind::NVSHMEM_WAIT:\n nvshmemx_uint64_wait_until_on_stream(sig_addr, NVSHMEM_CMP_EQ, wait_value, cur_stream);\n NVTE_CHECK_CUDA_DRIVER(hipStreamWriteValue64((hipStream_t)cur_stream, (hipDeviceptr_t)sig_addr,\n (cuuint64_t)signal_reset,\n CU_STREAM_WRITE_VALUE_DEFAULT));\n break;\n case WaitKind::STREAM_WAIT:\n NVTE_CHECK_CUDA_DRIVER(hipStreamWaitValue64((hipStream_t)cur_stream, (hipDeviceptr_t)sig_addr,\n (cuuint64_t)wait_value, hipStreamWaitValueGte));\n NVTE_CHECK_CUDA_DRIVER(hipStreamWriteValue64((hipStream_t)cur_stream, (hipDeviceptr_t)sig_addr,\n (cuuint64_t)signal_reset,\n CU_STREAM_WRITE_VALUE_DEFAULT));\n break;\n }\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/nvshmem_api/nvshmem_waitkernel.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n\n#if __CUDA_ARCH__ >= 800\n#define half_dtype nv_bfloat16\n#else\n#define half_dtype half\n#endif\n\n#include \n#include \n#include \n\n#include \"common/util/system.h\"\n#include \"common/util/vectorized_pointwise.h\"\n#include \"userbuffers.h\"\n\n#define MAX_THREADS 1024\n\n#define ATOMIC_CONSUMER(chunk) \\\n if (counters) { \\\n if (threadIdx.x == 0 && blockIdx.x == 0) { \\\n while (0 != (atomicCAS(((unsigned int *)counters) + chunk, 0, 0))) { \\\n } \\\n ((unsigned int *)counters)[chunk] = 1; \\\n asm volatile(\"fence.sc.gpu;\\n\"); \\\n } \\\n if (blockIdx.x == 0) __syncthreads(); \\\n }\n\n#define ATOMIC_PRODUCER(chunk) \\\n if (counters) { \\\n ((unsigned int *)counters)[chunk] = 0; \\\n }\n\n// Return true if producer > consumer, otherwise false while preventing integer overflow\n// If we expect that producer will be 2B+ messages behind consumer\n#define CHECK_IDS(producer, consumer) (((unsigned)(producer) - (unsigned)(consumer)) & (~INT_MAX))\n\n// Strip the path from a full filename\n#define FILENAME(file) \\\n ({ \\\n const char *filename = file; \\\n const char *basename = filename; \\\n for (const char *ptr = filename; *ptr != '\\0'; ptr++) { \\\n if (*ptr == '/' || *ptr == '\\\\') { \\\n basename = ptr + 1; \\\n } \\\n } \\\n basename; \\\n })\n\n// Printf to provide enough information so it is easier to attribute failures\n#define UB_PRINT(message, ...) \\\n printf(\"[%s:%s:%d] \" message \"\\n\", FILENAME(__FILE__), __FUNCTION__, __LINE__, __VA_ARGS__)\n\n// Report and error on timeout\n#define CHECK_TIMEOUT(t, timeout) ((clock64() - (t)) > timeout)\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rw(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep, const int lineoffset,\n const int numlines, void **commbuff, const int handleridx,\n const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n int *flagptr, physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n const int blockflagoffset = NVTE_MAX_NVLINK * 2 * blockIdx.x;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset + blockflagoffset;\n myptr += blockflagoffset;\n\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Allreduce reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank,\n blockIdx.x, threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n reduce_id++;\n }\n __syncthreads();\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n __syncthreads();\n for (int line = threadIdx.x + blockDim.x * (myrank + RANKS * blockIdx.x); line < numlines;\n line += blockDim.x * gridDim.x * RANKS) {\n int4 val[RANKS];\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n // int dest = (i+myrank+warp)&(RANKS-1);\n val[i] = userptr[dest[i]][lineoffset + line];\n }\n\n int4 sum = val[0];\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 1; i < RANKS; i++) {\n half_dtype *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < 8; j++) s[j] += x[j];\n }\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n // int dest = (i+myrank+warp)&(RANKS-1);\n userptr[dest[i]][lineoffset + line] = sum;\n }\n }\n\n __syncthreads();\n if (threadIdx.x == 0) __threadfence_system();\n __syncthreads();\n\n if (threadIdx.x < RANKS) {\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&myptr[targetgpu];\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Allreduce Gather: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n if (threadIdx.x == 0 && blockIdx.x == 0) *reduceidptr = reduce_id;\n} // fp16 inplace reduce kernel (Volta,Hopper)\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rr(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep, const int lineoffset,\n const int numlines, void **commbuff, const int handleridx,\n const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n int *flagptr, physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n const int blockflagoffset = NVTE_MAX_NVLINK * 2 * blockIdx.x;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset + blockflagoffset;\n myptr += blockflagoffset;\n\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d ]Allreduce reduce-scatter:SM %d [%d]: expecting %d got %d\", myrank,\n blockIdx.x, threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n reduce_id++;\n }\n __syncthreads();\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n __syncthreads();\n for (int line = threadIdx.x + blockDim.x * (myrank + RANKS * blockIdx.x); line < numlines;\n line += blockDim.x * gridDim.x * RANKS) {\n int4 val[RANKS];\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][lineoffset + line];\n }\n\n int4 sum = val[0];\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 1; i < RANKS; i++) {\n half_dtype *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < 8; j++) s[j] += x[j];\n }\n\n userptr[myrank][lineoffset + line] = sum;\n }\n __syncthreads();\n if (threadIdx.x == 0) __threadfence();\n __syncthreads();\n\n if (threadIdx.x < RANKS) {\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&myptr[targetgpu];\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Allreduce gather: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n\n int skipmy = 0;\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n int dst = (i + warp + myrank) & (RANKS - 1);\n if (dst == myrank) {\n skipmy++;\n continue;\n }\n dest[i - skipmy] = dst;\n }\n __syncthreads();\n\n for (int line = threadIdx.x + blockDim.x * RANKS * blockIdx.x; line < numlines;\n line += blockDim.x * gridDim.x * RANKS) {\n int4 val[RANKS - 1];\n\n#pragma unroll\n for (int i = 0; i < RANKS - 1; i++) {\n val[i] = userptr[dest[i]][lineoffset + line + blockDim.x * dest[i]];\n }\n\n#pragma unroll\n for (int i = 0; i < RANKS - 1; i++) {\n userptr[myrank][lineoffset + line + blockDim.x * dest[i]] = val[i];\n }\n }\n if (threadIdx.x == 0 && blockIdx.x == 0) *reduceidptr = reduce_id;\n} // fp16 inplace reduce kernel (Ampere)\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rr_rs(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep,\n const int mylineoffset, const int totallines,\n void **commbuff, const int handleridx,\n const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n __syncthreads();\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS];\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][mylineoffset + line];\n }\n\n int4 sum = val[0];\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 1; i < RANKS; i++) {\n half_dtype *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < 8; j++) s[j] += x[j];\n }\n\n userptr[myrank][mylineoffset + line] = sum;\n }\n\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 inplace reduce-scatter kernel\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS) userbuffers_fp16_sum_inplace_gpu_rr_rs_oop(\n const int op, const int flagoffset, const int firstrank, const int myrank, const int gpustep,\n const int mylineoffset, const int totallines, const int rowlines, const int skiplines,\n void **commbuff, const int handleridx, void *outbuf, const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n __syncthreads();\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS];\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][mylineoffset + line];\n }\n\n int4 sum = val[0];\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 1; i < RANKS; i++) {\n half_dtype *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < 8; j++) s[j] += x[j];\n }\n\n (reinterpret_cast(outbuf))[(line / rowlines) * skiplines + (line % rowlines)] = sum;\n }\n\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 reduce-scatter kernel (out of place)\n\n#if __CUDA_ARCH__ >= 900 && CUDART_VERSION >= 12010\n// All MC kernels here\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep, const int lineoffset,\n const int numlines, void **commbuff, const int handleridx,\n float4 *mc_ptr, const uint64_t ub_timeout) {\n int *flagptr, physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n const int blockflagoffset = NVTE_MAX_NVLINK * 2 * blockIdx.x;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset + blockflagoffset;\n myptr += blockflagoffset;\n\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (clock64() - s > ub_timeout) {\n UB_PRINT(\"Reduce-scatter: SM %d [%d]: expecting %d got %d\", blockIdx.x, threadIdx.x,\n reduce_id, *flag);\n break;\n }\n }\n reduce_id++;\n }\n __syncthreads();\n#define UNROLL_MC 8\n const int loop_step0 = blockDim.x * gridDim.x * RANKS;\n const int loop_step = loop_step0 * UNROLL_MC;\n const int start_elem = threadIdx.x + blockDim.x * (myrank + RANKS * blockIdx.x);\n const int end_elem = max(start_elem, numlines);\n const int aligned_elem = ((end_elem - start_elem) / loop_step) * loop_step;\n const int end_aligned = start_elem + aligned_elem;\n\n for (int line = start_elem; line < end_aligned; line += loop_step) {\n uint4 val[UNROLL_MC];\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++)\n#if defined(NVTE_UB_FP16)\n asm(\"multimem.ld_reduce.global.add.v4.f16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val[i].x), \"=r\"(val[i].y), \"=r\"(val[i].z), \"=r\"(val[i].w)\n : \"l\"(mc_ptr + (lineoffset + line + i * loop_step0))\n : \"memory\");\n#else\n asm(\"multimem.ld_reduce.global.add.v4.bf16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val[i].x), \"=r\"(val[i].y), \"=r\"(val[i].z), \"=r\"(val[i].w)\n : \"l\"(mc_ptr + (lineoffset + line + i * loop_step0))\n : \"memory\");\n#endif\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++)\n asm volatile(\"multimem.st.global.v4.f32 [%0], {%1,%2,%3,%4};\" ::\"l\"(\n mc_ptr + (lineoffset + line + i * loop_step0)),\n \"r\"(val[i].x), \"r\"(val[i].y), \"r\"(val[i].z), \"r\"(val[i].w)\n : \"memory\");\n }\n for (int line = end_aligned; line < end_elem; line += loop_step0) {\n uint4 val;\n#if defined(NVTE_UB_FP16)\n asm(\"multimem.ld_reduce.global.add.v4.f16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val.x), \"=r\"(val.y), \"=r\"(val.z), \"=r\"(val.w)\n : \"l\"(mc_ptr + (lineoffset + line))\n : \"memory\");\n#else\n asm(\"multimem.ld_reduce.global.add.v4.bf16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val.x), \"=r\"(val.y), \"=r\"(val.z), \"=r\"(val.w)\n : \"l\"(mc_ptr + (lineoffset + line))\n : \"memory\");\n#endif\n asm volatile(\n \"multimem.st.global.v4.f32 [%0], {%1,%2,%3,%4};\" ::\"l\"(mc_ptr + (lineoffset + line)),\n \"r\"(val.x), \"r\"(val.y), \"r\"(val.z), \"r\"(val.w)\n : \"memory\");\n }\n\n __syncthreads();\n if (threadIdx.x == 0) __threadfence_system();\n __syncthreads();\n\n if (threadIdx.x < RANKS) {\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&myptr[targetgpu];\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (clock64() - s > 2ull * ub_timeout) {\n UB_PRINT(\"Allgather: SM %d [%d]: expecting %d got %d\", blockIdx.x, threadIdx.x, reduce_id,\n *flag);\n break;\n }\n }\n }\n if (threadIdx.x == 0 && blockIdx.x == 0) *reduceidptr = reduce_id;\n} // fp16 inplace reduce kernel (Hopper) MC\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc_rs(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep,\n const int mylineoffset, const int totallines,\n void **commbuff, const int handleridx, float4 *mc_ptr,\n const uint64_t ub_timeout) {\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n uint4 *localptr = reinterpret_cast(commbuff[myrank * gpustep + firstrank + handleridx]);\n int lastSM = 0;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n const int loop_step0 = blockDim.x * gridDim.x;\n const int loop_step = loop_step0 * UNROLL_MC;\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = max(start_elem, totallines);\n const int aligned_elem = ((end_elem - start_elem) / loop_step) * loop_step;\n const int end_aligned = start_elem + aligned_elem;\n\n for (int line = start_elem; line < end_aligned; line += loop_step) {\n uint4 val[UNROLL_MC];\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++)\n#if defined(NVTE_UB_FP16)\n asm(\"multimem.ld_reduce.global.add.v4.f16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val[i].x), \"=r\"(val[i].y), \"=r\"(val[i].z), \"=r\"(val[i].w)\n : \"l\"(mc_ptr + (mylineoffset + line + i * loop_step0))\n : \"memory\");\n#else\n asm(\"multimem.ld_reduce.global.add.v4.bf16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val[i].x), \"=r\"(val[i].y), \"=r\"(val[i].z), \"=r\"(val[i].w)\n : \"l\"(mc_ptr + (mylineoffset + line + i * loop_step0))\n : \"memory\");\n#endif\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++) localptr[mylineoffset + line + i * loop_step0] = val[i];\n }\n for (int line = end_aligned; line < end_elem; line += loop_step0) {\n uint4 val;\n#if defined(NVTE_UB_FP16)\n asm(\"multimem.ld_reduce.global.add.v4.f16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val.x), \"=r\"(val.y), \"=r\"(val.z), \"=r\"(val.w)\n : \"l\"(mc_ptr + (mylineoffset + line))\n : \"memory\");\n#else\n asm(\"multimem.ld_reduce.global.add.v4.bf16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val.x), \"=r\"(val.y), \"=r\"(val.z), \"=r\"(val.w)\n : \"l\"(mc_ptr + (mylineoffset + line))\n : \"memory\");\n#endif\n localptr[mylineoffset + line] = val;\n }\n\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 inplace reduce-scatter kernel MC\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc_rs_oop(const int op, const int flagoffset,\n const int firstrank, const int myrank,\n const int gpustep, const int mylineoffset,\n const int totallines, const int rowlines,\n const int skiplines, void **commbuff,\n const int handleridx, void *outbuf, float4 *mc_ptr,\n const uint64_t ub_timeout) {\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n\n const int loop_step0 = blockDim.x * gridDim.x;\n const int loop_step = loop_step0 * UNROLL_MC;\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = max(start_elem, totallines);\n const int aligned_elem = ((end_elem - start_elem) / loop_step) * loop_step;\n const int end_aligned = start_elem + aligned_elem;\n for (int line = start_elem; line < end_aligned; line += loop_step) {\n uint4 val[UNROLL_MC];\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++)\n#if defined(NVTE_UB_FP16)\n asm(\"multimem.ld_reduce.global.add.v4.f16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val[i].x), \"=r\"(val[i].y), \"=r\"(val[i].z), \"=r\"(val[i].w)\n : \"l\"(mc_ptr + (mylineoffset + line + i * loop_step0))\n : \"memory\");\n#else\n asm(\"multimem.ld_reduce.global.add.v4.bf16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val[i].x), \"=r\"(val[i].y), \"=r\"(val[i].z), \"=r\"(val[i].w)\n : \"l\"(mc_ptr + (mylineoffset + line + i * loop_step0))\n : \"memory\");\n#endif\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++)\n (reinterpret_cast(outbuf))[((line + i * loop_step0) / rowlines) * skiplines +\n ((line + i * loop_step0) % rowlines)] = val[i];\n }\n for (int line = end_aligned; line < end_elem; line += loop_step0) {\n uint4 val;\n#if defined(NVTE_UB_FP16)\n asm(\"multimem.ld_reduce.global.add.v4.f16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val.x), \"=r\"(val.y), \"=r\"(val.z), \"=r\"(val.w)\n : \"l\"(mc_ptr + (mylineoffset + line))\n : \"memory\");\n#else\n asm(\"multimem.ld_reduce.global.add.v4.bf16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val.x), \"=r\"(val.y), \"=r\"(val.z), \"=r\"(val.w)\n : \"l\"(mc_ptr + (mylineoffset + line))\n : \"memory\");\n#endif\n reinterpret_cast(outbuf)[(line / rowlines) * skiplines + (line % rowlines)] = val;\n }\n\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 reduce-scatter kernel (out of place) fp16 MC\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc_ag(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep,\n const int mylineoffset, const int totallines,\n void **commbuff, const int handleridx, uint4 *mc_ptr,\n const uint64_t ub_timeout) {\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n uint4 *localptr = reinterpret_cast(commbuff[myrank * gpustep + firstrank + handleridx]);\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n }\n __syncthreads();\n\n const int loop_step0 = blockDim.x * gridDim.x;\n const int loop_step = loop_step0 * UNROLL_MC;\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = max(start_elem, totallines);\n const int aligned_elem = ((end_elem - start_elem) / loop_step) * loop_step;\n const int end_aligned = start_elem + aligned_elem;\n for (int line = start_elem; line < end_aligned; line += loop_step) {\n uint4 val[UNROLL_MC];\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++) val[i] = localptr[mylineoffset + line + i * loop_step0];\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++)\n asm volatile(\"multimem.st.global.v4.f32 [%0], {%1,%2,%3,%4};\" ::\"l\"(\n mc_ptr + (mylineoffset + line + i * loop_step0)),\n \"r\"(val[i].x), \"r\"(val[i].y), \"r\"(val[i].z), \"r\"(val[i].w)\n : \"memory\");\n }\n for (int line = end_aligned; line < end_elem; line += loop_step0) {\n uint4 val = localptr[mylineoffset + line];\n asm volatile(\n \"multimem.st.global.v4.f32 [%0], {%1,%2,%3,%4};\" ::\"l\"(mc_ptr + (mylineoffset + line)),\n \"r\"(val.x), \"r\"(val.y), \"r\"(val.z), \"r\"(val.w)\n : \"memory\");\n }\n\n __syncthreads();\n if (threadIdx.x == 0) __threadfence_system();\n __syncthreads();\n\n __shared__ int lastSM;\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id)\n lastSM = 1;\n else\n lastSM = 0;\n }\n __syncthreads();\n if (lastSM && threadIdx.x < RANKS) {\n if (threadIdx.x == 0) *reduceidptr = reduce_id;\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&myptr[targetgpu];\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Allgather: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x, threadIdx.x,\n reduce_id, *flag);\n break;\n }\n }\n }\n} // fp16 inplace allgather kernel (Hopper) MC\n\n#else\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep, const int lineoffset,\n const int numlines, void **commbuff, const int handleridx,\n float4 *mc_ptr, const uint64_t ub_timeout) {}\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc_rs_oop(const int op, const int flagoffset,\n const int firstrank, const int myrank,\n const int gpustep, const int mylineoffset,\n const int totallines, const int rowlines,\n const int skiplines, void **commbuff,\n const int handleridx, void *outbuf, float4 *mc_ptr,\n const uint64_t ub_timeout) {}\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc_ag(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep,\n const int mylineoffset, const int totallines,\n void **commbuff, const int handleridx, uint4 *mc_ptr,\n const uint64_t ub_timeout) {}\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc_rs(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep,\n const int mylineoffset, const int totallines,\n void **commbuff, const int handleridx, float4 *mc_ptr,\n const uint64_t ub_timeout) {}\n#endif\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS) userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_fp8(\n const int op, const int flagoffset, const int firstrank, const int myrank, const int gpustep,\n const int mylineoffset, const int totallines, const int rowlines, const int skiplines,\n void **commbuff, const int handleridx, void *outbuf, float *scale, const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n half_dtype hscale = (half_dtype)*scale;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n __syncthreads();\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS];\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][mylineoffset + line];\n }\n\n int4 sum[2] = {{0, 0, 0, 0}, {0, 0, 0, 0}};\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n fp8type *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < sizeof(int4) / sizeof(fp8type); j++) s[j] += hscale * (half_dtype)(x[j]);\n }\n int hline = 2 * line;\n (reinterpret_cast(outbuf))[(hline / rowlines) * skiplines + (hline % rowlines)] =\n sum[0];\n hline++;\n (reinterpret_cast(outbuf))[(hline / rowlines) * skiplines + (hline % rowlines)] =\n sum[1];\n }\n\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 reduce-scatter kernel (out of place) (fp8->fp16)\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_atomic_fp8(\n const int op, const int flagoffset, const int firstrank, const int myrank,\n const int gpustep, const int mylineoffset, const int totallines, const int rowlines,\n const int skiplines_out, const int skiplines_in, void **commbuff, const int handleridx,\n void *outbuf, float *scale, void *counters, const int numchunks, const int atomicindex,\n const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n half_dtype hscale = (half_dtype)*scale;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n // const int blockflagoffset = MAX_NVLINK * 2 * blockIdx.x;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr);\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset; // + blockflagoffset;\n }\n\n for (int chunk_i = 0; chunk_i < numchunks; chunk_i++) {\n ATOMIC_CONSUMER(chunk_i);\n\n lastSM = 0;\n if (threadIdx.x < RANKS) {\n reduce_id++;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), /*numchunks * */ adder);\n if (old_val + adder == NVTE_MAX_SMS * (reduce_id /* + numchunks*/)) lastSM = 1;\n }\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n __syncthreads();\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS];\n const int rowlines_in = rowlines / 2;\n const int index_in = skiplines_in == 0\n ? mylineoffset + myrank * totallines + line\n : (numchunks <= 1 ? 1 : chunk_i) * mylineoffset +\n myrank * (totallines * skiplines_in / rowlines_in) +\n (line / rowlines_in) * skiplines_in + (line % rowlines_in);\n const int index1_out = chunk_i * mylineoffset * 2 + ((2 * line) / rowlines) * skiplines_out +\n ((2 * line) % rowlines);\n const int index2_out = chunk_i * mylineoffset * 2 +\n ((2 * line + 1) / rowlines) * skiplines_out +\n ((2 * line + 1) % rowlines);\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][index_in];\n }\n\n int4 sum[2] = {{0, 0, 0, 0}, {0, 0, 0, 0}};\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n fp8type *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < sizeof(int4) / sizeof(fp8type); j++)\n s[j] += hscale * (half_dtype)(x[j]);\n }\n (reinterpret_cast(outbuf))[index1_out] = sum[0];\n (reinterpret_cast(outbuf))[index2_out] = sum[1];\n }\n }\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 reduce-scatter kernel (out of place) (fp8->fp16)\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS) userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_stride(\n const int op, const int flagoffset, const int firstrank, const int myrank, const int gpustep,\n const int mylineoffset, const int totallines, const int rowlines, const int skiplines,\n void **commbuff, const int handleridx, void *outbuf, const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS];\n int index_in = mylineoffset + myrank * (totallines * skiplines / rowlines) +\n (line / rowlines) * skiplines + (line % rowlines);\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][index_in];\n }\n\n int4 sum = val[0];\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 1; i < RANKS; i++) {\n half_dtype *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < 8; j++) s[j] += x[j];\n }\n\n int index_out = (line / rowlines) * skiplines + (line % rowlines);\n (reinterpret_cast(outbuf))[index_out] = sum;\n }\n\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 reduce-scatter kernel (out of place) fp16\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_stride_atomic(\n const int op, const int flagoffset, const int firstrank, const int myrank,\n const int gpustep, const int mylineoffset, const int totallines, const int rowlines,\n const int skiplines, const int numchunks, void **commbuff, const int handleridx,\n void *outbuf, void *counters, const uint64_t ub_timeout) {\n if (counters) {\n if (threadIdx.x == 0) {\n // spin-lock on counter from producer\n while (0 != (atomicCAS(((unsigned int *)counters), 0, 0))) {\n }\n\n // make sure all threadblocks have read/waited on counters.\n atomicInc(((unsigned int *)counters) + numchunks, gridDim.x - 1);\n while (0 != (atomicCAS(((unsigned int *)counters) + numchunks, 0, 0))) {\n }\n\n // reset counter for next producer.\n ((unsigned int *)counters)[0] = 1;\n asm volatile(\"fence.sc.gpu;\\n\");\n }\n }\n __syncthreads();\n\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS];\n int index_in = mylineoffset + myrank * (totallines * skiplines / rowlines) +\n (line / rowlines) * skiplines + (line % rowlines);\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][index_in];\n }\n\n int4 sum = val[0];\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 1; i < RANKS; i++) {\n half_dtype *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < 8; j++) s[j] += x[j];\n }\n\n int index_out = (line / rowlines) * skiplines + (line % rowlines);\n (reinterpret_cast(outbuf))[index_out] = sum;\n }\n\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 reduce-scatter kernel (out of place) fp16\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_stride_multiatomic(\n const int op, const int flagoffset, const int firstrank, const int myrank,\n const int gpustep, const int mylineoffset, const int totallines, const int rowlines,\n const int skiplines, const int numchunks, void **commbuff, const int handleridx,\n void *outbuf, void *counters, const uint64_t ub_timeout) {\n for (int chunk_i = 0; chunk_i < numchunks; chunk_i++) {\n if (counters) {\n if (threadIdx.x == 0) {\n // spin-lock on counter from producer\n while (0 != (atomicCAS(((unsigned int *)counters) + chunk_i, 0, 0))) {\n }\n\n // make sure all threadblocks have read/waited on counters.\n atomicInc(((unsigned int *)counters) + numchunks + chunk_i, gridDim.x - 1);\n while (0 != (atomicCAS(((unsigned int *)counters) + numchunks + chunk_i, 0, 0))) {\n }\n\n // reset counter for next producer.\n ((unsigned int *)counters)[chunk_i] = 1;\n asm volatile(\"fence.sc.gpu;\\n\");\n }\n }\n __syncthreads();\n\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS];\n int index_in = chunk_i * mylineoffset + myrank * (totallines * skiplines / rowlines) +\n (line / rowlines) * skiplines + (line % rowlines);\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][index_in];\n }\n\n int4 sum = val[0];\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 1; i < RANKS; i++) {\n half_dtype *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < 8; j++) s[j] += x[j];\n }\n\n int index_out = chunk_i * mylineoffset + (line / rowlines) * skiplines + (line % rowlines);\n (reinterpret_cast(outbuf))[index_out] = sum;\n }\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n }\n} // fp16 reduce-scatter kernel (out of place) fp16\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rr_ag(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep,\n const int mylineoffset, const int totallines,\n void **commbuff, const int handleridx,\n const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n }\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n\n int skipmy = 0;\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n int dst = (i + warp + myrank) & (RANKS - 1);\n if (dst == myrank) {\n skipmy++;\n continue;\n }\n dest[i - skipmy] = dst;\n }\n __syncthreads();\n\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS - 1];\n\n#pragma unroll\n for (int i = 0; i < RANKS - 1; i++) {\n val[i] = userptr[dest[i]][mylineoffset + line + totallines * dest[i]];\n }\n\n#pragma unroll\n for (int i = 0; i < RANKS - 1; i++) {\n userptr[myrank][mylineoffset + line + totallines * dest[i]] = val[i];\n }\n }\n __shared__ int lastSM;\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id)\n lastSM = 1;\n else\n lastSM = 0;\n }\n __syncthreads();\n if (lastSM && threadIdx.x < RANKS) {\n if (threadIdx.x == 0) *reduceidptr = reduce_id;\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&myptr[targetgpu];\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Allgather: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x, threadIdx.x,\n reduce_id, *flag);\n break;\n }\n }\n }\n} // fp16 inplace reduce kernel (Ampere)\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rw_ag(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep,\n const int mylineoffset, const int totallines,\n void **commbuff, const int handleridx,\n const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int4 *localptr;\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n }\n __syncthreads();\n localptr = userptr[myrank];\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS - 1];\n int skipmy = 0;\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n int dst = (i + warp + myrank) & (RANKS - 1);\n if (dst == myrank) {\n skipmy++;\n continue;\n }\n dest[i - skipmy] = dst;\n }\n#define UNROLLAG 4\n __syncthreads();\n const int loop_step0 = blockDim.x * gridDim.x;\n const int loop_step = loop_step0 * UNROLLAG;\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = max(start_elem, totallines);\n const int aligned_elem = ((end_elem - start_elem) / loop_step) * loop_step;\n const int end_aligned = start_elem + aligned_elem;\n\n for (int line = start_elem; line < end_aligned; line += loop_step) {\n int4 val[UNROLLAG];\n#pragma unroll\n for (int j = 0; j < UNROLLAG; j++) val[j] = localptr[mylineoffset + line + loop_step0 * j];\n\n#pragma unroll\n for (int j = 0; j < UNROLLAG; j++)\n#pragma unroll\n for (int i = 0; i < RANKS - 1; i++) {\n userptr[dest[i]][mylineoffset + line + j * loop_step0] = val[j];\n }\n }\n\n for (int line = end_aligned; line < end_elem; line += loop_step0) {\n int4 sum = localptr[mylineoffset + line];\n#pragma unroll\n for (int i = 0; i < RANKS - 1; i++) {\n userptr[dest[i]][mylineoffset + line] = sum;\n }\n }\n\n __syncthreads();\n if (threadIdx.x == 0) __threadfence_system();\n __syncthreads();\n\n __shared__ int lastSM;\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id)\n lastSM = 1;\n else\n lastSM = 0;\n }\n __syncthreads();\n if (lastSM && threadIdx.x < RANKS) {\n if (threadIdx.x == 0) *reduceidptr = reduce_id;\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&myptr[targetgpu];\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Allgather: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x, threadIdx.x,\n reduce_id, *flag);\n break;\n }\n }\n }\n} // fp16 inplace allgather kernel (Volta,Hopper)\n\n#define SETUP_LAUNCH_CONFIG(sms, threads, stream) \\\n cudaLaunchConfig_t cfg = {sms, threads, 0, stream, NULL, 0}; \\\n cudaLaunchAttribute attribute_ub[2]; \\\n attribute_ub[1].id = cudaLaunchAttributeClusterDimension; \\\n attribute_ub[1].val.clusterDim.x = sms % comm->cga_size == 0 ? comm->cga_size : 1; \\\n attribute_ub[1].val.clusterDim.y = 1; \\\n attribute_ub[1].val.clusterDim.z = 1; \\\n attribute_ub[0].id = cudaLaunchAttributeCooperative; \\\n cfg.attrs = attribute_ub; \\\n cfg.numAttrs = comm->sm_arch >= 9 ? 2 : 1;\n\n#if (CUDART_VERSION >= 12030)\n#define ADD_LAUNCH_COMPLETION_EVENT(attribute_ub, comm_launch_event) \\\n attribute_ub[2].id = cudaLaunchAttributeLaunchCompletionEvent; \\\n attribute_ub[2].val.launchCompletionEvent.event = comm_launch_event;\n#define NUM_LAUNCH_ATTRIBUTE_FOR_FDL_LAUNCH 3\n#else\n#define ADD_LAUNCH_COMPLETION_EVENT(attribute_ub, comm_launch_event)\n#define NUM_LAUNCH_ATTRIBUTE_FOR_FDL_LAUNCH 2\n#endif\n\n#define SETUP_LAUNCH_CONFIG_WITH_COMPLETION_EVENT(sms, threads, stream, comm_launch_event) \\\n cudaLaunchConfig_t cfg = {sms, threads, 0, stream, NULL, 0}; \\\n cudaLaunchAttribute attribute_ub[NUM_LAUNCH_ATTRIBUTE_FOR_FDL_LAUNCH] = {}; \\\n ADD_LAUNCH_COMPLETION_EVENT(attribute_ub, comm_launch_event) \\\n attribute_ub[1].id = cudaLaunchAttributeClusterDimension; \\\n attribute_ub[1].val.clusterDim.x = sms % comm->cga_size == 0 ? comm->cga_size : 1; \\\n attribute_ub[1].val.clusterDim.y = 1; \\\n attribute_ub[1].val.clusterDim.z = 1; \\\n attribute_ub[0].id = cudaLaunchAttributeCooperative; \\\n cfg.attrs = attribute_ub; \\\n cfg.numAttrs = NUM_LAUNCH_ATTRIBUTE_FOR_FDL_LAUNCH;\n\n#define callranks_ag(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8 + (comm->use_rr_kernel ? 0 : arg4 * arg7); \\\n void **arg8 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg9 = handler * comm->nvsize; \\\n uint64_t arg10 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10)}; \\\n NVTE_CHECK_CUDA(cudaLaunchKernelExC( \\\n &cfg, \\\n reinterpret_cast(comm->use_rr_kernel ? userbuffers_fp16_sum_inplace_gpu_rr_ag \\\n : userbuffers_fp16_sum_inplace_gpu_rw_ag), \\\n kernelArgs)); \\\n }\n\n#define callranks_agMC(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8 + arg4 * arg7; \\\n void **arg8 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg9 = handler * comm->nvsize; \\\n uint4 *arg10 = reinterpret_cast(comm->mc_ptr[handler]); \\\n uint64_t arg11 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11)}; \\\n NVTE_CHECK_CUDA(cudaLaunchKernelExC( \\\n &cfg, reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_mc_ag), kernelArgs)); \\\n }\n\n#define callranks_rs(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8 + arg4 * arg7; \\\n void **arg8 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg9 = handler * comm->nvsize; \\\n uint64_t arg10 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10)}; \\\n NVTE_CHECK_CUDA(cudaLaunchKernelExC( \\\n &cfg, reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_rr_rs), kernelArgs)); \\\n }\n\n#define callranks_rsMC(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8 + arg4 * arg7; \\\n void **arg8 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg9 = handler * comm->nvsize; \\\n void *arg10 = comm->mc_ptr[handler]; \\\n uint64_t arg11 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11)}; \\\n NVTE_CHECK_CUDA(cudaLaunchKernelExC( \\\n &cfg, reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_mc_rs), kernelArgs)); \\\n }\n\n#define callranks_rs_oop(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8 + arg4 * arg7, arg8 = rowelements / 8, arg9 = strideelements / 8; \\\n void **arg10 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg11 = handler * comm->nvsize; \\\n void *arg12 = output; \\\n uint64_t arg13 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11), reinterpret_cast(&arg12), \\\n reinterpret_cast(&arg13)}; \\\n NVTE_CHECK_CUDA(cudaLaunchKernelExC( \\\n &cfg, reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_rr_rs_oop), \\\n kernelArgs)); \\\n }\n\n#define callranks_rs_oop_fp8(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 16 / x, \\\n arg6 = offset / 16 + arg4 * arg7, arg8 = rowelements / 8, arg9 = strideelements / 8; \\\n void **arg10 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg11 = handler * comm->nvsize; \\\n void *arg12 = output; \\\n float *arg13 = scale; \\\n uint64_t arg14 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11), reinterpret_cast(&arg12), \\\n reinterpret_cast(&arg13), reinterpret_cast(&arg14)}; \\\n NVTE_CHECK_CUDA(cudaLaunchKernelExC( \\\n &cfg, \\\n reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_fp8), \\\n kernelArgs)); \\\n }\n\n#define callranks_rs_oopMC(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8 + arg4 * arg7, arg8 = rowelements / 8, arg9 = strideelements / 8; \\\n void **arg10 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg11 = handler * comm->nvsize; \\\n void *arg12 = output; \\\n void *arg13 = comm->mc_ptr[handler]; \\\n uint64_t arg14 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11), reinterpret_cast(&arg12), \\\n reinterpret_cast(&arg13), reinterpret_cast(&arg14)}; \\\n NVTE_CHECK_CUDA(cudaLaunchKernelExC( \\\n &cfg, reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_mc_rs_oop), \\\n kernelArgs)); \\\n }\n\n#define callranks_rs_oop_atomic_fp8(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 16 / x, \\\n arg6 = offset / 16, arg8 = rowelements / 8, arg9 = strideelements_out / 8, \\\n arg10 = strideelements_in / 16; \\\n void **arg11 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg12 = handler * comm->nvsize; \\\n void *arg13 = output; \\\n float *arg14 = scale; \\\n void *arg15 = counters; \\\n int arg16 = numchunks, arg17 = atomicindex; \\\n uint64_t arg18 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11), reinterpret_cast(&arg12), \\\n reinterpret_cast(&arg13), reinterpret_cast(&arg14), \\\n reinterpret_cast(&arg15), reinterpret_cast(&arg16), \\\n reinterpret_cast(&arg17), reinterpret_cast(&arg18)}; \\\n NVTE_CHECK_CUDA(cudaLaunchKernelExC( \\\n &cfg, \\\n reinterpret_cast( \\\n userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_atomic_fp8), \\\n kernelArgs)); \\\n }\n\n#define callranks_rs_oop_stride(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8, arg8 = rowelements / 8, arg9 = strideelements / 8; \\\n void **arg10 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg11 = handler * comm->nvsize; \\\n void *arg12 = output; \\\n uint64_t arg13 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11), reinterpret_cast(&arg12), \\\n reinterpret_cast(&arg13)}; \\\n NVTE_CHECK_CUDA(cudaLaunchKernelExC( \\\n &cfg, reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_stride), \\\n kernelArgs)); \\\n }\n\n#define callranks_rs_oop_stride_atomic(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8, arg8 = rowelements / 8, arg9 = strideelements / 8, arg10 = numchunks; \\\n void **arg11 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg12 = handler * comm->nvsize; \\\n void *arg13 = output; \\\n void *arg14 = counters; \\\n uint64_t arg15 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11), reinterpret_cast(&arg12), \\\n reinterpret_cast(&arg13), reinterpret_cast(&arg14), \\\n reinterpret_cast(&arg15)}; \\\n NVTE_CHECK_CUDA(cudaLaunchKernelExC( \\\n &cfg, \\\n reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_stride_atomic), \\\n kernelArgs)); \\\n }\n\n#define callranks_rs_oop_stride_multiatomic(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8, arg8 = rowelements / 8, arg9 = strideelements / 8, arg10 = numchunks; \\\n void **arg11 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg12 = handler * comm->nvsize; \\\n void *arg13 = output; \\\n void *arg14 = counters; \\\n uint64_t arg15 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11), reinterpret_cast(&arg12), \\\n reinterpret_cast(&arg13), reinterpret_cast(&arg14), \\\n reinterpret_cast(&arg15)}; \\\n NVTE_CHECK_CUDA( \\\n cudaLaunchKernelExC(&cfg, \\\n reinterpret_cast( \\\n userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_stride_multiatomic), \\\n kernelArgs)); \\\n }\n\nvoid reducescatter2_userbuff_strided(void *output, const int handler, const int offset,\n const int rowelements, const int colelements,\n const int strideelements, communicator *comm,\n cudaStream_t stream) {\n const int elements = rowelements * colelements;\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n\n if (elements < 64) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n callranks_rs_oop_stride(2) callranks_rs_oop_stride(4) callranks_rs_oop_stride(8)\n callranks_rs_oop_stride(16) callranks_rs_oop_stride(32)\n}\nvoid reducescatter2_userbuff_strided_atomic(void *output, const int handler, const int offset,\n const int rowelements, const int colelements,\n const int strideelements, const int numchunks,\n void *counters, communicator *comm,\n cudaStream_t stream) {\n const int elements = rowelements * colelements;\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n\n if (elements < 64) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n callranks_rs_oop_stride_atomic(2) callranks_rs_oop_stride_atomic(4)\n callranks_rs_oop_stride_atomic(8) callranks_rs_oop_stride_atomic(16)\n callranks_rs_oop_stride_atomic(32)\n}\n\ntemplate \nvoid reducescatter2_userbuff_strided_universal_fp8(void *output, float *scale, const int handler,\n const int offset, const int rowelements,\n const int colelements,\n const int strideelements_out,\n const int strideelements_in, const int numchunks,\n const int atomicindex, void *counters,\n communicator *comm, cudaStream_t stream) {\n const int elements = rowelements * colelements;\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n assert(comm->sm_arch >= 9);\n if (elements < 128) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n callranks_rs_oop_atomic_fp8(2) callranks_rs_oop_atomic_fp8(4) callranks_rs_oop_atomic_fp8(8)\n callranks_rs_oop_atomic_fp8(16) callranks_rs_oop_atomic_fp8(32)\n}\n\ntemplate \nvoid reducescatter2_userbuff_strided_atomic_fp8(void *output, float *scale, const int handler,\n const int offset, const int rowelements,\n const int colelements, const int strideelements_out,\n const int strideelements_in, const int numchunks,\n void *counters, communicator *comm,\n cudaStream_t stream) {\n reducescatter2_userbuff_strided_universal_fp8(\n output, scale, handler, offset, rowelements, colelements, strideelements_out,\n strideelements_in, 1, numchunks, counters /*nullptr*/, comm, stream);\n}\n\ntemplate \nvoid reducescatter2_userbuff_strided_multiatomic_fp8(\n void *output, float *scale, const int handler, const int offset, const int rowelements,\n const int colelements, const int strideelements_out, const int strideelements_in,\n const int numchunks, void *counters, communicator *comm, cudaStream_t stream) {\n reducescatter2_userbuff_strided_universal_fp8(\n output, scale, handler, offset, rowelements, colelements, strideelements_out,\n strideelements_in, numchunks, 0, counters /*nullptr*/, comm, stream);\n}\n\nvoid reducescatter2_userbuff_strided_multiatomic(void *output, const int handler, const int offset,\n const int rowelements, const int colelements,\n const int strideelements, const int numchunks,\n void *counters, communicator *comm,\n cudaStream_t stream) {\n const int elements = rowelements * colelements;\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n\n if (elements < 64) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n callranks_rs_oop_stride_multiatomic(2) callranks_rs_oop_stride_multiatomic(4)\n callranks_rs_oop_stride_multiatomic(8) callranks_rs_oop_stride_multiatomic(16)\n callranks_rs_oop_stride_multiatomic(32)\n}\n\nvoid allgather2_userbuff_inplace(const int handler, const int offset, const int elements,\n communicator *comm, cudaStream_t stream,\n cudaEvent_t comm_launch_event) {\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n\n if (elements < 64) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n if (comm_launch_event) {\n SETUP_LAUNCH_CONFIG_WITH_COMPLETION_EVENT(sms, warps * 32, stream, comm_launch_event);\n if (comm->use_mc && (comm->memflags[handler] & NVTE_UB_MEM_MC_CREATED)) {\n callranks_agMC(2) callranks_agMC(4) callranks_agMC(8) callranks_agMC(16) callranks_agMC(32)\n } else {\n callranks_ag(2) callranks_ag(4) callranks_ag(8) callranks_ag(16) callranks_ag(32)\n }\n } else {\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n if (comm->use_mc && (comm->memflags[handler] & NVTE_UB_MEM_MC_CREATED)) {\n callranks_agMC(2) callranks_agMC(4) callranks_agMC(8) callranks_agMC(16) callranks_agMC(32)\n } else {\n callranks_ag(2) callranks_ag(4) callranks_ag(8) callranks_ag(16) callranks_ag(32)\n }\n }\n}\n\nvoid allgather2_userbuff_inplace_sliced(const int handler, const int offset, const int elements,\n communicator *comm, const int slice_id, const int nslices,\n cudaStream_t stream) {\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n int peerelements = elements / ar_nvsize;\n int saverrkernel = comm->use_rr_kernel;\n comm->use_rr_kernel = 0;\n allgather2_userbuff_inplace(\n handler, offset + ar_nvrank * peerelements * (nslices - 1) + slice_id * peerelements,\n elements, comm, stream);\n comm->use_rr_kernel = saverrkernel;\n}\n\nvoid reducescatter2_userbuff_inplace(const int handler, const int offset, const int elements,\n communicator *comm, cudaStream_t stream,\n cudaEvent_t comm_launch_event) {\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n\n if (elements < 64) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n if (comm_launch_event) {\n SETUP_LAUNCH_CONFIG_WITH_COMPLETION_EVENT(sms, warps * 32, stream, comm_launch_event);\n if (comm->use_mc && (comm->memflags[handler] & NVTE_UB_MEM_MC_CREATED)) {\n callranks_rsMC(2) callranks_rsMC(4) callranks_rsMC(8) callranks_rsMC(16) callranks_rsMC(32)\n } else {\n callranks_rs(2) callranks_rs(4) callranks_rs(8) callranks_rs(16) callranks_rs(32)\n }\n } else {\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n if (comm->use_mc && (comm->memflags[handler] & NVTE_UB_MEM_MC_CREATED)) {\n callranks_rsMC(2) callranks_rsMC(4) callranks_rsMC(8) callranks_rsMC(16) callranks_rsMC(32)\n } else {\n callranks_rs(2) callranks_rs(4) callranks_rs(8) callranks_rs(16) callranks_rs(32)\n }\n }\n}\nvoid reducescatter2_userbuff_stridedoutput(void *output, const int handler, const int offset,\n const int rowelements, const int colelements,\n const int strideelements, communicator *comm,\n cudaStream_t stream, cudaEvent_t comm_launch_event) {\n const int elements = rowelements * colelements;\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n\n if (elements < 64) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n if (comm_launch_event) {\n SETUP_LAUNCH_CONFIG_WITH_COMPLETION_EVENT(sms, warps * 32, stream, comm_launch_event);\n if (comm->use_mc && (comm->memflags[handler] & NVTE_UB_MEM_MC_CREATED)) {\n callranks_rs_oopMC(2) callranks_rs_oopMC(4) callranks_rs_oopMC(8) callranks_rs_oopMC(16)\n callranks_rs_oopMC(32)\n } else {\n callranks_rs_oop(2) callranks_rs_oop(4) callranks_rs_oop(8) callranks_rs_oop(16)\n callranks_rs_oop(32)\n }\n } else {\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n if (comm->use_mc && (comm->memflags[handler] & NVTE_UB_MEM_MC_CREATED)) {\n callranks_rs_oopMC(2) callranks_rs_oopMC(4) callranks_rs_oopMC(8) callranks_rs_oopMC(16)\n callranks_rs_oopMC(32)\n } else {\n callranks_rs_oop(2) callranks_rs_oop(4) callranks_rs_oop(8) callranks_rs_oop(16)\n callranks_rs_oop(32)\n }\n }\n}\nvoid reducescatter2_userbuff(void *output, const int handler, const int offset, const int elements,\n communicator *comm, cudaStream_t stream,\n cudaEvent_t comm_launch_event) {\n reducescatter2_userbuff_stridedoutput(output, handler, offset, elements, 1, 0, comm, stream,\n comm_launch_event);\n}\n\ntemplate \nvoid reducescatter2_userbuff_stridedoutput_fp8(void *output, float *scale, const int handler,\n const int offset, const int rowelements,\n const int colelements, const int strideelements,\n communicator *comm, cudaStream_t stream,\n cudaEvent_t comm_launch_event) {\n const int elements = rowelements * colelements;\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n assert(comm->sm_arch >= 9);\n if (elements < 128) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n if (comm_launch_event) {\n SETUP_LAUNCH_CONFIG_WITH_COMPLETION_EVENT(sms, warps * 32, stream, comm_launch_event);\n callranks_rs_oop_fp8(2) callranks_rs_oop_fp8(4) callranks_rs_oop_fp8(8) callranks_rs_oop_fp8(16)\n callranks_rs_oop_fp8(32)\n } else {\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n callranks_rs_oop_fp8(2) callranks_rs_oop_fp8(4) callranks_rs_oop_fp8(8) callranks_rs_oop_fp8(16)\n callranks_rs_oop_fp8(32)\n }\n}\n\ntemplate void reducescatter2_userbuff_stridedoutput_fp8<__nv_fp8_e5m2>(\n void *output, float *scale, const int handler, const int offset, const int rowelements,\n const int colelements, const int strideelements, communicator *comm, cudaStream_t stream,\n cudaEvent_t comm_launch_event);\n\ntemplate void reducescatter2_userbuff_stridedoutput_fp8<__nv_fp8_e4m3>(\n void *output, float *scale, const int handler, const int offset, const int rowelements,\n const int colelements, const int strideelements, communicator *comm, cudaStream_t stream,\n cudaEvent_t comm_launch_event);\n\ntemplate \nvoid reducescatter2_userbuff_fp8(void *output, float *scale, const int handler, const int offset,\n const int elements, communicator *comm, cudaStream_t stream,\n cudaEvent_t comm_launch_event) {\n reducescatter2_userbuff_stridedoutput_fp8(output, scale, handler, offset, elements, 1, 0,\n comm, stream, comm_launch_event);\n}\n\ntemplate void reducescatter2_userbuff_fp8<__nv_fp8_e5m2>(void *output, float *scale,\n const int handler, const int offset,\n const int elements, communicator *comm,\n cudaStream_t stream,\n cudaEvent_t comm_launch_event);\ntemplate void reducescatter2_userbuff_fp8<__nv_fp8_e4m3>(void *output, float *scale,\n const int handler, const int offset,\n const int elements, communicator *comm,\n cudaStream_t stream,\n cudaEvent_t comm_launch_event);\n\ntemplate void reducescatter2_userbuff_strided_atomic_fp8<__nv_fp8_e4m3>(\n void *output, float *scale, const int handler, const int offset, const int rowelements,\n const int colelements, const int strideelements_out, const int strideelements_in,\n const int numchunks, void *counters, communicator *comm, cudaStream_t stream);\ntemplate void reducescatter2_userbuff_strided_atomic_fp8<__nv_fp8_e5m2>(\n void *output, float *scale, const int handler, const int offset, const int rowelements,\n const int colelements, const int strideelements_out, const int strideelements_in,\n const int numchunks, void *counters, communicator *comm, cudaStream_t stream);\ntemplate void reducescatter2_userbuff_strided_multiatomic_fp8<__nv_fp8_e4m3>(\n void *output, float *scale, const int handler, const int offset, const int rowelements,\n const int colelements, const int strideelements_out, const int strideelements_in,\n const int numchunks, void *counters, communicator *comm, cudaStream_t stream);\ntemplate void reducescatter2_userbuff_strided_multiatomic_fp8<__nv_fp8_e5m2>(\n void *output, float *scale, const int handler, const int offset, const int rowelements,\n const int colelements, const int strideelements_out, const int strideelements_in,\n const int numchunks, void *counters, communicator *comm, cudaStream_t stream);\n\n__global__ void kuserbuffers_pullsend(int myrank, int peer, int *send_id, int *flagptr) {\n atomicAdd_system(flagptr, 1);\n}\n\n__global__ void kuserbuffers_inc(int *id) { atomicAdd(id, 1); }\n\n__global__ void kuserbuffers_dummy(void) {}\n\n__global__ void __launch_bounds__(MAX_THREADS)\n kuserbuffers_pullrecv(int myrank, int peer, int nvrank, int nvpeer, int *recv_id, int *flagptr,\n int4 *srcptr, int4 *dstptr, const int lines, uint64_t ub_timeout) {\n#define UNROLLCOPY 8\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = lines;\n const int aligned_elem = (end_elem - start_elem) & (~(blockDim.x * gridDim.x * UNROLLCOPY - 1));\n const int end_aligned = start_elem + aligned_elem;\n\n if (threadIdx.x == 0) {\n const int signal_id = (*recv_id) + 1;\n volatile int *flag = (volatile int *)flagptr;\n clock_t s = clock64();\n while (CHECK_IDS(*flag, signal_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\n \"pullrecv [grank dst:%d global src:%d][nvrank(GPU) dst: %d src: %d]: expecting %d,\"\n \" observed %d\",\n myrank, peer, nvrank, nvpeer, signal_id, *flag);\n break;\n }\n }\n if (lines == 0) {\n *recv_id = signal_id;\n return;\n } // otherwise need an extra kernel\n }\n __syncthreads();\n\n if (end_elem <= start_elem) return;\n\n for (int line = start_elem; line < end_aligned; line += blockDim.x * gridDim.x * UNROLLCOPY) {\n int4 val[UNROLLCOPY];\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) val[i] = srcptr[line + i * blockDim.x * gridDim.x];\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) dstptr[line + i * blockDim.x * gridDim.x] = val[i];\n }\n for (int line = end_aligned; line < end_elem; line += blockDim.x * gridDim.x)\n dstptr[line] = srcptr[line];\n}\n\n__global__ void __launch_bounds__(MAX_THREADS)\n kuserbuffers_pushsend(int *send_id, int *flagptr, int4 *srcptr, int4 *dstptr, const int lines) {\n if (lines) {\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = lines;\n const int aligned_elem =\n ((end_elem - start_elem) & (~(blockDim.x * gridDim.x * UNROLLCOPY - 1)));\n const int end_aligned = start_elem + aligned_elem;\n if (end_elem > start_elem) {\n for (int line = start_elem; line < end_aligned; line += blockDim.x * gridDim.x * UNROLLCOPY) {\n int4 val[UNROLLCOPY];\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) val[i] = srcptr[line + i * blockDim.x * gridDim.x];\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) dstptr[line + i * blockDim.x * gridDim.x] = val[i];\n }\n for (int line = end_aligned; line < end_elem; line += blockDim.x * gridDim.x)\n dstptr[line] = srcptr[line];\n }\n __syncthreads();\n if (threadIdx.x) return;\n __threadfence_system();\n atomicAdd_system(flagptr,\n 1); // otherwise need local SM sync before sending flag\n } else { // 0 bytes and 1 SM only\n atomicAdd_system(flagptr, 1);\n }\n}\n\n#define CHECK_CE(ce_start, ce_end) \\\n ((ce_start) != nullptr && (ce_end) != nullptr && *(ce_start) != *(ce_end))\n\n__global__ void kuserbuffers_pushrecv(int myrank, int peer, int nvrank, int nvpeer, int *recv_id,\n int *flagptr, int adder, uint64_t ub_timeout,\n int *ce_start_ptr, int *ce_end_ptr) {\n const int signal_id = (*recv_id) + adder;\n *recv_id = signal_id;\n volatile int *flag = (volatile int *)flagptr;\n if (*flag >= signal_id) return;\n clock_t s = clock64();\n while (CHECK_IDS(*flag, signal_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\n \"pushrecv [grank dst:%d global src:%d][nvrank(GPU) dst: %d src: %d]: \"\n \"expecting %d, observed %d\",\n myrank, peer, nvrank, nvpeer, signal_id, *flag);\n if (CHECK_CE(ce_start_ptr, ce_end_ptr))\n UB_PRINT(\"pushrecv: CE deadlock DETECTED: %d (ce_start) != %d (ce_end)\\n\", *ce_start_ptr,\n *ce_end_ptr);\n return;\n }\n }\n}\n\n__global__ void __launch_bounds__(MAX_THREADS)\n kuserbuffers_pushsendrecv(int *send_id, int *send_flagptr, int4 *srcptr, int4 *dstptr,\n const int lines, int send_peer, int recv_peer, int *recv_id,\n int *recv_flagptr, int adder, uint64_t ub_timeout, int nv_send,\n int nv_recv, int *ce_start_ptr, int *ce_end_ptr) {\n if (lines) {\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = lines;\n const int aligned_elem =\n ((end_elem - start_elem) & (~(blockDim.x * gridDim.x * UNROLLCOPY - 1)));\n const int end_aligned = start_elem + aligned_elem;\n if (end_elem > start_elem) {\n for (int line = start_elem; line < end_aligned; line += blockDim.x * gridDim.x * UNROLLCOPY) {\n int4 val[UNROLLCOPY];\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) {\n val[i] = srcptr[line + i * blockDim.x * gridDim.x];\n }\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) {\n dstptr[line + i * blockDim.x * gridDim.x] = val[i];\n }\n }\n for (int line = end_aligned; line < end_elem; line += blockDim.x * gridDim.x) {\n dstptr[line] = srcptr[line];\n }\n }\n __syncthreads();\n if (threadIdx.x) return;\n __threadfence_system();\n atomicAdd_system(send_flagptr,\n 1); // otherwise need local SM sync before sending flag\n } else { // 0 bytes and 1 SM only\n atomicAdd_system(send_flagptr, 1);\n }\n\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n const int signal_id = (*recv_id) + adder;\n *recv_id = signal_id;\n volatile int *flag = (volatile int *)recv_flagptr;\n if (*flag >= signal_id) return;\n clock_t s = clock64();\n while (CHECK_IDS(*flag, signal_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\n \"pushsendrecv [sending peer:%d receiving peer:%d][nvrank(GPU) sending peer: %d\"\n \" receiving peer: %d]: expecting %d, observed %d\",\n send_peer, recv_peer, nv_send, nv_recv, signal_id, *flag);\n if (CHECK_CE(ce_start_ptr, ce_end_ptr))\n UB_PRINT(\"pushrecv: CE deadlock DETECTED: %d (ce_start) != %d (ce_end)\\n\", *ce_start_ptr,\n *ce_end_ptr);\n return;\n }\n }\n }\n}\n\n__global__ void __launch_bounds__(MAX_THREADS)\n kuserbuffers_pushsendrecv_atomic(int *send_id, int *send_flagptr, int4 *srcptr, int4 *dstptr,\n const int lines, int send_peer, int recv_peer, int *recv_id,\n int *recv_flagptr, int adder, void *counters,\n uint64_t ub_timeout, int nv_send, int nv_recv,\n int *ce_start_ptr, int *ce_end_ptr) {\n if (lines) {\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = lines;\n const int aligned_elem =\n ((end_elem - start_elem) & (~(blockDim.x * gridDim.x * UNROLLCOPY - 1)));\n const int end_aligned = start_elem + aligned_elem;\n if (end_elem > start_elem) {\n for (int line = start_elem; line < end_aligned; line += blockDim.x * gridDim.x * UNROLLCOPY) {\n int4 val[UNROLLCOPY];\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) {\n val[i] = srcptr[line + i * blockDim.x * gridDim.x];\n }\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) {\n dstptr[line + i * blockDim.x * gridDim.x] = val[i];\n }\n }\n for (int line = end_aligned; line < end_elem; line += blockDim.x * gridDim.x) {\n dstptr[line] = srcptr[line];\n }\n }\n __syncthreads();\n if (threadIdx.x) return;\n __threadfence_system();\n atomicAdd_system(send_flagptr,\n 1); // otherwise need local SM sync before sending flag\n } else { // 0 bytes and 1 SM only\n atomicAdd_system(send_flagptr, 1);\n }\n\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n const int signal_id = (*recv_id) + adder;\n *recv_id = signal_id;\n volatile int *flag = (volatile int *)recv_flagptr;\n clock_t s = clock64();\n while (CHECK_IDS(*flag, signal_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\n \"pushsendrecv atomic [sending peer:%d receiving peer:%d][nvrank(GPU) sending peer:\"\n \" %d receiving peer: %d]: expecting %d, observed %d\",\n send_peer, recv_peer, nv_send, nv_recv, signal_id, *flag); /*return;*/\n if (CHECK_CE(ce_start_ptr, ce_end_ptr))\n UB_PRINT(\"pushsendrecv atomic: CE deadlock DETECTED: %d (ce_start) != %d (ce_end)\\n\",\n *ce_start_ptr, *ce_end_ptr);\n }\n }\n\n // Decrement atomic val to signal current output tile finish\n if (counters) {\n ((unsigned int *)counters)[0] = 0;\n asm volatile(\"fence.sc.gpu;\\n\");\n }\n }\n}\n\n__global__ void __launch_bounds__(MAX_THREADS) kuserbuffers_pushsendrecv_multiatomic(\n int *send_id, int *send_flagptr, int4 *srcptr, int4 *dstptr, const int lines, int send_peer,\n int recv_peer, int *recv_id, int *recv_flagptr, int adder, void *counters, int nchunks,\n int send_stride, int recv_stride, bool shuffle, uint64_t ub_timeout, int nv_send, int nv_recv) {\n for (int chunk_i = 0; chunk_i < nchunks - 1; chunk_i++) {\n int send_chunk_id = shuffle ? chunk_i : (nchunks + send_peer - chunk_i) % nchunks;\n int recv_chunk_id = shuffle ? chunk_i + 1 : (nchunks + send_peer - chunk_i - 1) % nchunks;\n int send_offset = (send_chunk_id * send_stride) / 16;\n int recv_offset = ((shuffle ? recv_chunk_id : send_chunk_id) * recv_stride) / 16;\n\n if (lines) {\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = lines;\n const int aligned_elem =\n ((end_elem - start_elem) & (~(blockDim.x * gridDim.x * UNROLLCOPY - 1)));\n const int end_aligned = start_elem + aligned_elem;\n if (end_elem > start_elem) {\n for (int line = start_elem; line < end_aligned;\n line += blockDim.x * gridDim.x * UNROLLCOPY) {\n int4 val[UNROLLCOPY];\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) {\n val[i] = srcptr[send_offset + line + i * blockDim.x * gridDim.x];\n }\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) {\n dstptr[recv_offset + line + i * blockDim.x * gridDim.x] = val[i];\n }\n }\n for (int line = end_aligned; line < end_elem; line += blockDim.x * gridDim.x) {\n dstptr[recv_offset + line] = srcptr[send_offset + line];\n }\n }\n __syncthreads();\n if (!threadIdx.x) {\n __threadfence_system();\n atomicAdd_system(send_flagptr,\n 1); // otherwise need local SM sync before sending flag\n }\n } else { // 0 bytes and 1 SM only\n atomicAdd_system(send_flagptr, 1);\n }\n\n // wait for message to arrive.\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n const int signal_id = (*recv_id) + adder;\n *recv_id = signal_id;\n volatile int *flag = (volatile int *)recv_flagptr;\n clock_t s = clock64();\n while (CHECK_IDS(*flag, signal_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\n \"pushsendrecv multiatomic [sending peer:%d receiving peer:%d][nvrank(GPU)\"\n \" sending peer: %d receiving peer: %d]: expecting %d, observed %d\",\n send_peer, recv_peer, nv_send, nv_recv, signal_id, *flag); /*return;*/\n // CE mode is not supported for multi-atomic, so there is no need to check for a deadlock\n return;\n }\n }\n }\n\n // Producer must update counters.\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n // Decrement atomic val to signal current output tile finish\n if (counters) {\n ((unsigned int *)counters)[recv_chunk_id /*chunk_i+1*/] = 0;\n asm volatile(\"fence.sc.gpu;\\n\");\n }\n }\n\n // sync all CTAs before moving to next chunk.\n if (threadIdx.x == 0) {\n atomicInc(((unsigned int *)counters) + nchunks + chunk_i, gridDim.x - 1);\n while (0 != (atomicCAS(((unsigned int *)counters) + nchunks + chunk_i, 0, 0))) {\n }\n }\n __syncthreads();\n }\n}\n\n// Return TRUE if two ranks share the same NV domain\n#define INTRANODE(peer) ((peer / comm->nvsize) == (comm->myrank / comm->nvsize))\n\n// Index corresponds to the type of flag:\n// 0 - Send index counter\n// 1 - CE start index counter\n// 2 - CE end index counter\n#define GET_SEND_PTR_BY_INDEX(peerlocal, comm, dsth, index) \\\n ((reinterpret_cast((comm)->peer_ptr[0][(peerlocal)])) + \\\n ((NVTE_REG0_OFFSET(comm) + NVTE_REG0_RECV + (comm)->myrank * NVTE_MAX_REGIONS + (dsth) + \\\n (index) * NVTE_MAX_NVLINK * NVTE_MAX_REGIONS) * \\\n sizeof(int)))\n\n// Index corresponds to the type of flag:\n// 0 - Receive index counter\n// 1 - CE start index counter\n// 2 - CE end index counter\n#define GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsth, index) \\\n ((reinterpret_cast((comm)->mem_ptr[0])) + \\\n ((NVTE_REG0_OFFSET(comm) + NVTE_REG0_RECV + (recv_peer) * NVTE_MAX_REGIONS + (dsth) + \\\n (index) * NVTE_MAX_NVLINK * NVTE_MAX_REGIONS) * \\\n sizeof(int)))\n\nstatic __global__ void tiny_delay_kern() {\n // Although a loop could be used to add a larger delay here, for\n // the purpose of enforcing proper kernel ordering when using CG,\n // an empty kernel seems to work well enough.\n}\n\nvoid userbuffers_tiny_delay(cudaStream_t stream) { tiny_delay_kern<<<1, 1, 0, stream>>>(); }\n\nvoid userbuffers_send(const int srchandler, const size_t srcoffset, const int dsthandler,\n const size_t dstoffset, const size_t bytes, communicator *comm,\n const int peer, cudaStream_t stream) {\n int peerlocal = peer % comm->nvsize;\n void *flagptr = GET_SEND_PTR_BY_INDEX(peerlocal, comm, dsthandler, 0);\n // void *ce_send_start_ptr = GET_SEND_PTR_BY_INDEX(peerlocal, comm, dsthandler, 1);\n // void *ce_send_end_ptr = GET_SEND_PTR_BY_INDEX(peerlocal, comm, dsthandler, 2);\n bool signalonly = (bytes / 16 == 0) || (comm->use_ce != 0);\n\n assert(INTRANODE(peer));\n\n if (!(comm->launch_mode & NVTE_LAUNCH_GPU)) return;\n if (comm->push == 0) {\n kuserbuffers_pullsend<<<1, 1, 0, stream>>>(comm->myrank, peer, &(comm->send_id[peer]),\n reinterpret_cast(flagptr));\n NVTE_CHECK_CUDA(cudaGetLastError());\n } else {\n void *srcptr = reinterpret_cast(comm->mem_ptr[srchandler]) + srcoffset;\n void *dstptr = reinterpret_cast(comm->peer_ptr[dsthandler][peerlocal]) + dstoffset;\n\n if (comm->use_ce) {\n // kuserbuffers_inc<<<1, 1, 0, stream>>>(reinterpret_cast(ce_send_start_ptr));\n NVTE_CHECK_CUDA(cudaMemcpyAsync(dstptr, srcptr, bytes, cudaMemcpyDeviceToDevice, stream));\n // kuserbuffers_inc<<<1, 1, 0, stream>>>(reinterpret_cast(ce_send_end_ptr));\n }\n SETUP_LAUNCH_CONFIG(signalonly ? 1 : comm->sms, signalonly ? 1 : 1024, stream);\n int *arg1 = &comm->send_id[peer], *arg2 = reinterpret_cast(flagptr);\n int4 *arg3 = reinterpret_cast(srcptr), *arg4 = reinterpret_cast(dstptr);\n int arg5 = signalonly ? 0 : bytes / 16;\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2),\n reinterpret_cast(&arg3), reinterpret_cast(&arg4),\n reinterpret_cast(&arg5)};\n NVTE_CHECK_CUDA(\n cudaLaunchKernelExC(&cfg, reinterpret_cast(kuserbuffers_pushsend), kernelArgs));\n }\n}\n\nvoid userbuffers_sendrecv(const int srchandler, const int dsthandler, const size_t send_offset,\n const size_t recv_offset, const size_t bytes, communicator *comm,\n const int send_peer, const int recv_peer, cudaStream_t stream) {\n bool signalonly = (bytes / 16 == 0) || (comm->use_ce != 0);\n int send_peerlocal = send_peer % comm->nvsize;\n int recv_peerlocal = recv_peer % comm->nvsize;\n void *flagptr_send = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 0);\n // void *ce_send_start_ptr = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 1);\n // void *ce_send_end_ptr = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 2);\n void *flagptr_recv = GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 0);\n\n void *send_srcptr = reinterpret_cast(comm->mem_ptr[srchandler]) + send_offset;\n void *send_dstptr =\n reinterpret_cast(comm->peer_ptr[dsthandler][send_peerlocal]) + send_offset;\n\n if (comm->use_ce) {\n // kuserbuffers_inc<<<1, 1, 0, stream>>>(reinterpret_cast(ce_send_start_ptr));\n NVTE_CHECK_CUDA(\n cudaMemcpyAsync(send_dstptr, send_srcptr, bytes, cudaMemcpyDeviceToDevice, stream));\n // kuserbuffers_inc<<<1, 1, 0, stream>>>(reinterpret_cast(ce_send_end_ptr));\n }\n SETUP_LAUNCH_CONFIG(signalonly ? 1 : comm->sms, signalonly ? 1 : 1024, stream);\n\n int *arg1 = &comm->send_id[send_peer];\n int *arg2 = reinterpret_cast(flagptr_send);\n int4 *arg3 = reinterpret_cast(send_srcptr);\n int4 *arg4 = reinterpret_cast(send_dstptr);\n int arg5 = signalonly ? 0 : bytes / 16;\n int arg6 = send_peer;\n int arg7 = recv_peer;\n int *arg8 = &comm->recv_id[recv_peer * NVTE_MAX_REGIONS + dsthandler];\n int *arg9 = reinterpret_cast(flagptr_recv);\n int arg10 = signalonly ? 1 : comm->sms;\n uint64_t arg11 = comm->ub_timeout;\n int arg12 = send_peerlocal;\n int arg13 = recv_peerlocal;\n int *arg14 = reinterpret_cast(0 ? // temporary disable\n GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 1)\n : nullptr);\n int *arg15 = reinterpret_cast(0 ? // temporary disable\n GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 2)\n : nullptr);\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2),\n reinterpret_cast(&arg3), reinterpret_cast(&arg4),\n reinterpret_cast(&arg5), reinterpret_cast(&arg6),\n reinterpret_cast(&arg7), reinterpret_cast(&arg8),\n reinterpret_cast(&arg9), reinterpret_cast(&arg10),\n reinterpret_cast(&arg11), reinterpret_cast(&arg12),\n reinterpret_cast(&arg13), reinterpret_cast(&arg14),\n reinterpret_cast(&arg15)};\n NVTE_CHECK_CUDA(\n cudaLaunchKernelExC(&cfg, reinterpret_cast(kuserbuffers_pushsendrecv), kernelArgs));\n}\n\nvoid userbuffers_sendrecv_atomic(const int srchandler, const int dsthandler,\n const size_t send_offset, const size_t recv_offset,\n const size_t bytes, communicator *comm, const int send_peer,\n const int recv_peer, void *counters, cudaStream_t stream) {\n assert(comm->push && comm->use_ce == 0);\n bool signalonly = (bytes / 16 == 0) || (comm->use_ce != 0);\n\n int send_peerlocal = send_peer % comm->nvsize;\n int recv_peerlocal = recv_peer % comm->nvsize;\n void *flagptr_send = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 0);\n // void *ce_send_start_ptr = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 1);\n // void *ce_send_end_ptr = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 2);\n void *flagptr_recv = GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 0);\n\n void *send_srcptr = reinterpret_cast(comm->mem_ptr[srchandler]) + send_offset;\n void *send_dstptr =\n reinterpret_cast(comm->peer_ptr[dsthandler][send_peerlocal]) + send_offset;\n if (comm->use_ce) {\n // kuserbuffers_inc<<<1, 1, 0, stream>>>(reinterpret_cast(ce_send_start_ptr));\n NVTE_CHECK_CUDA(\n cudaMemcpyAsync(send_dstptr, send_srcptr, bytes, cudaMemcpyDeviceToDevice, stream));\n // kuserbuffers_inc<<<1, 1, 0, stream>>>(reinterpret_cast(ce_send_end_ptr));\n }\n SETUP_LAUNCH_CONFIG(signalonly ? 1 : comm->sms, signalonly ? 1 : 1024, stream);\n\n int *arg1 = &comm->send_id[send_peer];\n int *arg2 = reinterpret_cast(flagptr_send);\n int4 *arg3 = reinterpret_cast(send_srcptr);\n int4 *arg4 = reinterpret_cast(send_dstptr);\n int arg5 = signalonly ? 0 : bytes / 16;\n int arg6 = send_peer;\n int arg7 = recv_peer;\n int *arg8 = &comm->recv_id[recv_peer * NVTE_MAX_REGIONS + dsthandler];\n int *arg9 = reinterpret_cast(flagptr_recv);\n int arg10 = signalonly ? 1 : comm->sms;\n void *arg11 = counters;\n int arg12 = comm->ub_timeout;\n int arg13 = send_peerlocal;\n int arg14 = recv_peerlocal;\n int *arg15 = reinterpret_cast(0 ? // temporary disable\n GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 1)\n : nullptr);\n int *arg16 = reinterpret_cast(0 ? // temporary disable\n GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 2)\n : nullptr);\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2),\n reinterpret_cast(&arg3), reinterpret_cast(&arg4),\n reinterpret_cast(&arg5), reinterpret_cast(&arg6),\n reinterpret_cast(&arg7), reinterpret_cast(&arg8),\n reinterpret_cast(&arg9), reinterpret_cast(&arg10),\n reinterpret_cast(&arg11), reinterpret_cast(&arg12),\n reinterpret_cast(&arg13), reinterpret_cast(&arg14),\n reinterpret_cast(&arg15), reinterpret_cast(&arg16)};\n NVTE_CHECK_CUDA(cudaLaunchKernelExC(\n &cfg, reinterpret_cast(kuserbuffers_pushsendrecv_atomic), kernelArgs));\n}\n\nvoid userbuffers_sendrecv_multiatomic(const int srchandler, const int dsthandler,\n const size_t send_stride, const size_t recv_stride,\n const size_t bytes, communicator *comm, const int send_peer,\n const int recv_peer, const int nchunks, void *counters,\n bool shuffle, cudaStream_t stream) {\n assert(comm->push && comm->use_ce == 0);\n // CE is not supported\n\n int send_peerlocal = send_peer % comm->nvsize;\n int recv_peerlocal = recv_peer % comm->nvsize;\n void *flagptr_send = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 0);\n void *flagptr_recv = GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 0);\n\n SETUP_LAUNCH_CONFIG(comm->sms, 1024, stream);\n\n int *arg1 = &comm->send_id[send_peer];\n int *arg2 = reinterpret_cast(flagptr_send);\n int4 *arg3 = reinterpret_cast((comm->mem_ptr[srchandler]));\n int4 *arg4 = reinterpret_cast((comm->peer_ptr[dsthandler][send_peerlocal]));\n int arg5 = bytes / 16;\n int arg6 = comm->myrank;\n int arg7 = recv_peer;\n int *arg8 = &comm->recv_id[recv_peer * NVTE_MAX_REGIONS + dsthandler];\n int *arg9 = reinterpret_cast(flagptr_recv);\n int arg10 = comm->sms;\n void *arg11 = counters;\n int arg12 = nchunks;\n int arg13 = send_stride;\n int arg14 = recv_stride;\n bool arg15 = shuffle;\n uint64_t arg16 = comm->ub_timeout;\n int arg17 = send_peerlocal;\n int arg18 = recv_peerlocal;\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2),\n reinterpret_cast(&arg3), reinterpret_cast(&arg4),\n reinterpret_cast(&arg5), reinterpret_cast(&arg6),\n reinterpret_cast(&arg7), reinterpret_cast(&arg8),\n reinterpret_cast(&arg9), reinterpret_cast(&arg10),\n reinterpret_cast(&arg11), reinterpret_cast(&arg12),\n reinterpret_cast(&arg13), reinterpret_cast(&arg14),\n reinterpret_cast(&arg15), reinterpret_cast(&arg16),\n reinterpret_cast(&arg17), reinterpret_cast(&arg18)};\n NVTE_CHECK_CUDA(cudaLaunchKernelExC(\n &cfg, reinterpret_cast(kuserbuffers_pushsendrecv_multiatomic), kernelArgs));\n}\n\nvoid userbuffers_recv(const int srchandler, const size_t srcoffset, const int dsthandler,\n const size_t dstoffset, const size_t bytes, communicator *comm,\n const int peer, cudaStream_t stream) {\n int peerlocal = peer % comm->nvsize;\n void *flagptr = GET_RECV_PTR_BY_INDEX(peer, comm, dsthandler, 0);\n bool signalonly = (bytes / 16 == 0) || (comm->use_ce != 0);\n\n assert(INTRANODE(peer));\n\n if (!(comm->launch_mode & NVTE_LAUNCH_GPU)) return;\n if (comm->push == 0) {\n void *dstptr = reinterpret_cast(comm->mem_ptr[dsthandler]) + dstoffset;\n void *srcptr = reinterpret_cast(comm->peer_ptr[srchandler][peerlocal]) + srcoffset;\n\n kuserbuffers_pullrecv<<sms, signalonly ? 1 : 1024, 0, stream>>>(\n comm->myrank, peer, comm->nvrank, peerlocal,\n &(comm->recv_id[peer * NVTE_MAX_REGIONS + dsthandler]), reinterpret_cast(flagptr),\n reinterpret_cast(srcptr), reinterpret_cast(dstptr),\n signalonly ? 0 : bytes / 16, comm->ub_timeout);\n NVTE_CHECK_CUDA(cudaGetLastError());\n if (!signalonly) {\n kuserbuffers_inc<<<1, 1, 0, stream>>>(&(comm->recv_id[peer * NVTE_MAX_REGIONS + dsthandler]));\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n if (comm->use_ce) {\n NVTE_CHECK_CUDA(cudaMemcpyAsync(dstptr, srcptr, bytes, cudaMemcpyDeviceToDevice, stream));\n }\n } else {\n kuserbuffers_pushrecv<<<1, 1, 0, stream>>>(\n comm->myrank, peer, comm->nvrank, peerlocal,\n &comm->recv_id[peer * NVTE_MAX_REGIONS + dsthandler], reinterpret_cast(flagptr),\n signalonly || comm->sms, comm->ub_timeout,\n reinterpret_cast(0 ? // temporary disable\n GET_RECV_PTR_BY_INDEX(peer, comm, dsthandler, 1)\n : nullptr),\n reinterpret_cast(0 ? // temporary disable\n GET_RECV_PTR_BY_INDEX(peer, comm, dsthandler, 2)\n : nullptr));\n NVTE_CHECK_CUDA(cudaGetLastError());\n }\n}\n\nvoid userbuffers_send_all(const int srchandler, const size_t srcoffset, const int dsthandler,\n const size_t dstoffset, const size_t bytes_per_slice, int tp_rank,\n int tp_size, int world_rank, communicator *comm, cudaStream_t stream) {\n int rank_round_tp = (world_rank / tp_size) * tp_size;\n for (int j = 1; j < tp_size; j++) {\n int i = (tp_rank + j) % tp_size;\n int send_offset = srcoffset + bytes_per_slice * tp_rank;\n int recv_offset = dstoffset + bytes_per_slice * tp_rank;\n userbuffers_send(srchandler, send_offset, dsthandler, recv_offset, bytes_per_slice, comm,\n rank_round_tp + i, stream);\n }\n}\n\nvoid userbuffers_recv_all(const int srchandler, const size_t srcoffset, const int dsthandler,\n const size_t dstoffset, const size_t bytes_per_slice, int tp_rank,\n int tp_size, int world_rank, communicator *comm, cudaStream_t stream) {\n int rank_round_tp = (world_rank / tp_size) * tp_size;\n for (int j = tp_size - 1; j > 0; j--) {\n int i = (tp_rank + j) % tp_size;\n int send_offset = srcoffset + bytes_per_slice * i;\n int recv_offset = dstoffset + bytes_per_slice * i;\n userbuffers_recv(srchandler, send_offset, dsthandler, recv_offset, bytes_per_slice, comm,\n rank_round_tp + i, stream);\n }\n}\n\n// producer\nstatic __global__ void producer_kernel(void *atomic_ptr, int chunk_i) {\n // Decrement atomic val to signal current output tile finish\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n ((unsigned int *)atomic_ptr)[chunk_i] = 0;\n }\n\n // COMM kernel need to explicitely flash gmem.\n // GEMM kernel already executed, and can not see gmem\n // change without COMM kernel explicitely make change\n asm volatile(\"fence.sc.gpu;\\n\");\n}\n\n// consumer\nstatic __global__ void consumer_kernel(void *atomic_ptr, int chunk_i) {\n // Wait for producer to change the val to 0, which signal producer ready\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n while (0 != (atomicCAS((unsigned int *)atomic_ptr + chunk_i, 0, 0))) {\n }\n ((unsigned int *)atomic_ptr)[chunk_i] = 1;\n asm volatile(\"fence.sc.gpu;\\n\");\n }\n}\n\n// consumer_batch\nstatic __global__ void consumer_batch_kernel(void *atomic_ptr, int first_chunk_i, int num_chunks) {\n // Wait for producer to change the val to 0, which signal producer ready\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n for (int i = first_chunk_i; i < num_chunks; i++) {\n while (0 != (atomicCAS((unsigned int *)atomic_ptr + i, 0, 0))) {\n }\n ((unsigned int *)atomic_ptr)[i] = 1;\n asm volatile(\"fence.sc.gpu;\\n\");\n }\n }\n}\n\n// reset counters kernel\nstatic __global__ void reset_counters_kernel(void *atomic_ptr, int num_chunks, bool allgather) {\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n#pragma unroll\n for (int i = 0; i < num_chunks; i++) {\n ((unsigned int *)atomic_ptr)[i] = 1;\n ((unsigned int *)atomic_ptr)[i + num_chunks] = 0;\n }\n if (allgather) ((unsigned int *)atomic_ptr)[0] = 0;\n }\n}\n\nvoid producer(void *atomic_ptr, int chunk_i, cudaStream_t stream) {\n dim3 block(1);\n dim3 grid(1);\n producer_kernel<<>>(atomic_ptr, chunk_i);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid consumer(void *atomic_ptr, int chunk_i, cudaStream_t stream) {\n dim3 block(1);\n dim3 grid(1);\n consumer_kernel<<>>(atomic_ptr, chunk_i);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid consumer_batch(void *atomic_ptr, int first_chunk_i, int num_chunks, cudaStream_t stream) {\n dim3 block(1);\n dim3 grid(1);\n consumer_batch_kernel<<>>(atomic_ptr, first_chunk_i, num_chunks);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid reset_counters(void *atomic_ptr, int num_chunks, bool allgather, cudaStream_t stream) {\n dim3 block(1);\n dim3 grid(1);\n reset_counters_kernel<<>>(atomic_ptr, num_chunks, allgather);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS / 4)\n reduce_fp8_in_bf16_out_cuda(void *inputs, void *output, const float *scale,\n const int num_inputs, const int input_size,\n const int num_aligned_elements_per_input,\n const int tot_input_size) {\n fp8type *inputs_fp8 = reinterpret_cast(inputs);\n half_dtype *output_half = reinterpret_cast(output);\n\n transformer_engine::VectorizedLoader loader(inputs_fp8, tot_input_size);\n transformer_engine::VectorizedStorer storer(output_half, input_size);\n\n const size_t tid = threadIdx.x + blockDim.x * blockIdx.x;\n if (tid >= num_aligned_elements_per_input) {\n return;\n }\n float accum_buf[nvec]; // NOLINT(*)\n\n loader.load(tid, tot_input_size);\n#pragma unroll\n for (int i = 0; i < nvec; ++i) {\n accum_buf[i] = static_cast(loader.separate()[i]) * (*scale);\n }\n for (int input_id = 1; input_id < num_inputs; ++input_id) {\n loader.load(tid + num_aligned_elements_per_input * input_id, tot_input_size);\n#pragma unroll\n for (int i = 0; i < nvec; ++i) {\n accum_buf[i] += static_cast(loader.separate()[i]) * (*scale);\n }\n }\n#pragma unroll\n for (int i = 0; i < nvec; ++i) {\n storer.separate()[i] = static_cast(accum_buf[i]);\n }\n storer.store(tid, input_size);\n}\n\ntemplate \nvoid reduce_fp8_in_bf16_out(void *inputs, void *output, float *scale, int num_inputs,\n int input_size, cudaStream_t stream) {\n constexpr int nvec = 32;\n assert(input_size % nvec == 0);\n const int num_aligned_elements_per_input = input_size / nvec;\n const int tot_input_size = input_size * num_inputs;\n size_t num_threads = MAX_THREADS / 4;\n size_t num_blocks = (num_aligned_elements_per_input + num_threads - 1) / num_threads;\n dim3 block(num_threads);\n dim3 grid(num_blocks);\n reduce_fp8_in_bf16_out_cuda\n <<>>(inputs, output, scale, num_inputs, input_size,\n num_aligned_elements_per_input, tot_input_size);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\ntemplate void reduce_fp8_in_bf16_out<__nv_fp8_e4m3>(void *inputs, void *output, float *scale,\n int num_inputs, int input_size,\n cudaStream_t stream);\ntemplate void reduce_fp8_in_bf16_out<__nv_fp8_e5m2>(void *inputs, void *output, float *scale,\n int num_inputs, int input_size,\n cudaStream_t stream);\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS / 4)\n reduce_bf16_cuda(void *inputs, void *output, const int num_inputs, const int input_size,\n const int num_aligned_elements_per_input, const int tot_input_size) {\n half_dtype *inputs_half = reinterpret_cast(inputs);\n half_dtype *output_half = reinterpret_cast(output);\n\n transformer_engine::VectorizedLoader loader(inputs_half, tot_input_size);\n transformer_engine::VectorizedStorer storer(output_half, input_size);\n\n const size_t tid = threadIdx.x + blockDim.x * blockIdx.x;\n if (tid >= num_aligned_elements_per_input) {\n return;\n }\n float accum_buf[nvec]; // NOLINT(*)\n\n loader.load(tid, tot_input_size);\n#pragma unroll\n for (int i = 0; i < nvec; ++i) {\n accum_buf[i] = static_cast(loader.separate()[i]);\n }\n for (int input_id = 1; input_id < num_inputs; ++input_id) {\n loader.load(tid + num_aligned_elements_per_input * input_id, tot_input_size);\n#pragma unroll\n for (int i = 0; i < nvec; ++i) {\n accum_buf[i] += static_cast(loader.separate()[i]);\n }\n }\n#pragma unroll\n for (int i = 0; i < nvec; ++i) {\n storer.separate()[i] = static_cast(accum_buf[i]);\n }\n storer.store(tid, input_size);\n}\n\nvoid reduce_bf16(void *inputs, void *output, int num_inputs, int input_size, cudaStream_t stream) {\n constexpr int nvec = 32;\n assert(input_size % nvec == 0);\n const int num_aligned_elements_per_input = input_size / nvec;\n const int tot_input_size = input_size * num_inputs;\n size_t num_threads = MAX_THREADS / 4;\n size_t num_blocks = (num_aligned_elements_per_input + num_threads - 1) / num_threads;\n dim3 block(num_threads);\n dim3 grid(num_blocks);\n reduce_bf16_cuda<<>>(\n inputs, output, num_inputs, input_size, num_aligned_elements_per_input, tot_input_size);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#if __CUDA_ARCH__ >= 800\n#define half_dtype hip_bfloat16\n#else\n#define half_dtype half\n#endif\n\n#include \n#include \n#include \n\n#include \"common/util/system.h\"\n#include \"common/util/vectorized_pointwise.h\"\n#include \"userbuffers.h\"\n\n#define MAX_THREADS 1024\n\n#define ATOMIC_CONSUMER(chunk) \\\n if (counters) { \\\n if (threadIdx.x == 0 && blockIdx.x == 0) { \\\n while (0 != (atomicCAS(((unsigned int *)counters) + chunk, 0, 0))) { \\\n } \\\n ((unsigned int *)counters)[chunk] = 1; \\\n asm volatile(\"fence.sc.gpu;\\n\"); \\\n } \\\n if (blockIdx.x == 0) __syncthreads(); \\\n }\n\n#define ATOMIC_PRODUCER(chunk) \\\n if (counters) { \\\n ((unsigned int *)counters)[chunk] = 0; \\\n }\n\n// Return true if producer > consumer, otherwise false while preventing integer overflow\n// If we expect that producer will be 2B+ messages behind consumer\n#define CHECK_IDS(producer, consumer) (((unsigned)(producer) - (unsigned)(consumer)) & (~INT_MAX))\n\n// Strip the path from a full filename\n#define FILENAME(file) \\\n ({ \\\n const char *filename = file; \\\n const char *basename = filename; \\\n for (const char *ptr = filename; *ptr != '\\0'; ptr++) { \\\n if (*ptr == '/' || *ptr == '\\\\') { \\\n basename = ptr + 1; \\\n } \\\n } \\\n basename; \\\n })\n\n// Printf to provide enough information so it is easier to attribute failures\n#define UB_PRINT(message, ...) \\\n printf(\"[%s:%s:%d] \" message \"\\n\", FILENAME(__FILE__), __FUNCTION__, __LINE__, __VA_ARGS__)\n\n// Report and error on timeout\n#define CHECK_TIMEOUT(t, timeout) ((clock64() - (t)) > timeout)\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rw(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep, const int lineoffset,\n const int numlines, void **commbuff, const int handleridx,\n const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n int *flagptr, physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n const int blockflagoffset = NVTE_MAX_NVLINK * 2 * blockIdx.x;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset + blockflagoffset;\n myptr += blockflagoffset;\n\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Allreduce reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank,\n blockIdx.x, threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n reduce_id++;\n }\n __syncthreads();\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n __syncthreads();\n for (int line = threadIdx.x + blockDim.x * (myrank + RANKS * blockIdx.x); line < numlines;\n line += blockDim.x * gridDim.x * RANKS) {\n int4 val[RANKS];\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n // int dest = (i+myrank+warp)&(RANKS-1);\n val[i] = userptr[dest[i]][lineoffset + line];\n }\n\n int4 sum = val[0];\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 1; i < RANKS; i++) {\n half_dtype *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < 8; j++) s[j] += x[j];\n }\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n // int dest = (i+myrank+warp)&(RANKS-1);\n userptr[dest[i]][lineoffset + line] = sum;\n }\n }\n\n __syncthreads();\n if (threadIdx.x == 0) __threadfence_system();\n __syncthreads();\n\n if (threadIdx.x < RANKS) {\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&myptr[targetgpu];\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Allreduce Gather: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n if (threadIdx.x == 0 && blockIdx.x == 0) *reduceidptr = reduce_id;\n} // fp16 inplace reduce kernel (Volta,Hopper)\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rr(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep, const int lineoffset,\n const int numlines, void **commbuff, const int handleridx,\n const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n int *flagptr, physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n const int blockflagoffset = NVTE_MAX_NVLINK * 2 * blockIdx.x;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset + blockflagoffset;\n myptr += blockflagoffset;\n\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d ]Allreduce reduce-scatter:SM %d [%d]: expecting %d got %d\", myrank,\n blockIdx.x, threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n reduce_id++;\n }\n __syncthreads();\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n __syncthreads();\n for (int line = threadIdx.x + blockDim.x * (myrank + RANKS * blockIdx.x); line < numlines;\n line += blockDim.x * gridDim.x * RANKS) {\n int4 val[RANKS];\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][lineoffset + line];\n }\n\n int4 sum = val[0];\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 1; i < RANKS; i++) {\n half_dtype *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < 8; j++) s[j] += x[j];\n }\n\n userptr[myrank][lineoffset + line] = sum;\n }\n __syncthreads();\n if (threadIdx.x == 0) __threadfence();\n __syncthreads();\n\n if (threadIdx.x < RANKS) {\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&myptr[targetgpu];\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Allreduce gather: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n\n int skipmy = 0;\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n int dst = (i + warp + myrank) & (RANKS - 1);\n if (dst == myrank) {\n skipmy++;\n continue;\n }\n dest[i - skipmy] = dst;\n }\n __syncthreads();\n\n for (int line = threadIdx.x + blockDim.x * RANKS * blockIdx.x; line < numlines;\n line += blockDim.x * gridDim.x * RANKS) {\n int4 val[RANKS - 1];\n\n#pragma unroll\n for (int i = 0; i < RANKS - 1; i++) {\n val[i] = userptr[dest[i]][lineoffset + line + blockDim.x * dest[i]];\n }\n\n#pragma unroll\n for (int i = 0; i < RANKS - 1; i++) {\n userptr[myrank][lineoffset + line + blockDim.x * dest[i]] = val[i];\n }\n }\n if (threadIdx.x == 0 && blockIdx.x == 0) *reduceidptr = reduce_id;\n} // fp16 inplace reduce kernel (Ampere)\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rr_rs(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep,\n const int mylineoffset, const int totallines,\n void **commbuff, const int handleridx,\n const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n __syncthreads();\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS];\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][mylineoffset + line];\n }\n\n int4 sum = val[0];\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 1; i < RANKS; i++) {\n half_dtype *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < 8; j++) s[j] += x[j];\n }\n\n userptr[myrank][mylineoffset + line] = sum;\n }\n\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 inplace reduce-scatter kernel\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS) userbuffers_fp16_sum_inplace_gpu_rr_rs_oop(\n const int op, const int flagoffset, const int firstrank, const int myrank, const int gpustep,\n const int mylineoffset, const int totallines, const int rowlines, const int skiplines,\n void **commbuff, const int handleridx, void *outbuf, const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n __syncthreads();\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS];\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][mylineoffset + line];\n }\n\n int4 sum = val[0];\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 1; i < RANKS; i++) {\n half_dtype *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < 8; j++) s[j] += x[j];\n }\n\n (reinterpret_cast(outbuf))[(line / rowlines) * skiplines + (line % rowlines)] = sum;\n }\n\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 reduce-scatter kernel (out of place)\n\n#if __CUDA_ARCH__ >= 900 && CUDART_VERSION >= 12010\n// All MC kernels here\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep, const int lineoffset,\n const int numlines, void **commbuff, const int handleridx,\n float4 *mc_ptr, const uint64_t ub_timeout) {\n int *flagptr, physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n const int blockflagoffset = NVTE_MAX_NVLINK * 2 * blockIdx.x;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset + blockflagoffset;\n myptr += blockflagoffset;\n\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (clock64() - s > ub_timeout) {\n UB_PRINT(\"Reduce-scatter: SM %d [%d]: expecting %d got %d\", blockIdx.x, threadIdx.x,\n reduce_id, *flag);\n break;\n }\n }\n reduce_id++;\n }\n __syncthreads();\n#define UNROLL_MC 8\n const int loop_step0 = blockDim.x * gridDim.x * RANKS;\n const int loop_step = loop_step0 * UNROLL_MC;\n const int start_elem = threadIdx.x + blockDim.x * (myrank + RANKS * blockIdx.x);\n const int end_elem = max(start_elem, numlines);\n const int aligned_elem = ((end_elem - start_elem) / loop_step) * loop_step;\n const int end_aligned = start_elem + aligned_elem;\n\n for (int line = start_elem; line < end_aligned; line += loop_step) {\n uint4 val[UNROLL_MC];\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++)\n#if defined(NVTE_UB_FP16)\n asm(\"multimem.ld_reduce.global.add.v4.f16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val[i].x), \"=r\"(val[i].y), \"=r\"(val[i].z), \"=r\"(val[i].w)\n : \"l\"(mc_ptr + (lineoffset + line + i * loop_step0))\n : \"memory\");\n#else\n asm(\"multimem.ld_reduce.global.add.v4.bf16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val[i].x), \"=r\"(val[i].y), \"=r\"(val[i].z), \"=r\"(val[i].w)\n : \"l\"(mc_ptr + (lineoffset + line + i * loop_step0))\n : \"memory\");\n#endif\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++)\n asm volatile(\"multimem.st.global.v4.f32 [%0], {%1,%2,%3,%4};\" ::\"l\"(\n mc_ptr + (lineoffset + line + i * loop_step0)),\n \"r\"(val[i].x), \"r\"(val[i].y), \"r\"(val[i].z), \"r\"(val[i].w)\n : \"memory\");\n }\n for (int line = end_aligned; line < end_elem; line += loop_step0) {\n uint4 val;\n#if defined(NVTE_UB_FP16)\n asm(\"multimem.ld_reduce.global.add.v4.f16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val.x), \"=r\"(val.y), \"=r\"(val.z), \"=r\"(val.w)\n : \"l\"(mc_ptr + (lineoffset + line))\n : \"memory\");\n#else\n asm(\"multimem.ld_reduce.global.add.v4.bf16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val.x), \"=r\"(val.y), \"=r\"(val.z), \"=r\"(val.w)\n : \"l\"(mc_ptr + (lineoffset + line))\n : \"memory\");\n#endif\n asm volatile(\n \"multimem.st.global.v4.f32 [%0], {%1,%2,%3,%4};\" ::\"l\"(mc_ptr + (lineoffset + line)),\n \"r\"(val.x), \"r\"(val.y), \"r\"(val.z), \"r\"(val.w)\n : \"memory\");\n }\n\n __syncthreads();\n if (threadIdx.x == 0) __threadfence_system();\n __syncthreads();\n\n if (threadIdx.x < RANKS) {\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&myptr[targetgpu];\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (clock64() - s > 2ull * ub_timeout) {\n UB_PRINT(\"Allgather: SM %d [%d]: expecting %d got %d\", blockIdx.x, threadIdx.x, reduce_id,\n *flag);\n break;\n }\n }\n }\n if (threadIdx.x == 0 && blockIdx.x == 0) *reduceidptr = reduce_id;\n} // fp16 inplace reduce kernel (Hopper) MC\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc_rs(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep,\n const int mylineoffset, const int totallines,\n void **commbuff, const int handleridx, float4 *mc_ptr,\n const uint64_t ub_timeout) {\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n uint4 *localptr = reinterpret_cast(commbuff[myrank * gpustep + firstrank + handleridx]);\n int lastSM = 0;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n const int loop_step0 = blockDim.x * gridDim.x;\n const int loop_step = loop_step0 * UNROLL_MC;\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = max(start_elem, totallines);\n const int aligned_elem = ((end_elem - start_elem) / loop_step) * loop_step;\n const int end_aligned = start_elem + aligned_elem;\n\n for (int line = start_elem; line < end_aligned; line += loop_step) {\n uint4 val[UNROLL_MC];\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++)\n#if defined(NVTE_UB_FP16)\n asm(\"multimem.ld_reduce.global.add.v4.f16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val[i].x), \"=r\"(val[i].y), \"=r\"(val[i].z), \"=r\"(val[i].w)\n : \"l\"(mc_ptr + (mylineoffset + line + i * loop_step0))\n : \"memory\");\n#else\n asm(\"multimem.ld_reduce.global.add.v4.bf16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val[i].x), \"=r\"(val[i].y), \"=r\"(val[i].z), \"=r\"(val[i].w)\n : \"l\"(mc_ptr + (mylineoffset + line + i * loop_step0))\n : \"memory\");\n#endif\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++) localptr[mylineoffset + line + i * loop_step0] = val[i];\n }\n for (int line = end_aligned; line < end_elem; line += loop_step0) {\n uint4 val;\n#if defined(NVTE_UB_FP16)\n asm(\"multimem.ld_reduce.global.add.v4.f16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val.x), \"=r\"(val.y), \"=r\"(val.z), \"=r\"(val.w)\n : \"l\"(mc_ptr + (mylineoffset + line))\n : \"memory\");\n#else\n asm(\"multimem.ld_reduce.global.add.v4.bf16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val.x), \"=r\"(val.y), \"=r\"(val.z), \"=r\"(val.w)\n : \"l\"(mc_ptr + (mylineoffset + line))\n : \"memory\");\n#endif\n localptr[mylineoffset + line] = val;\n }\n\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 inplace reduce-scatter kernel MC\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc_rs_oop(const int op, const int flagoffset,\n const int firstrank, const int myrank,\n const int gpustep, const int mylineoffset,\n const int totallines, const int rowlines,\n const int skiplines, void **commbuff,\n const int handleridx, void *outbuf, float4 *mc_ptr,\n const uint64_t ub_timeout) {\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n\n const int loop_step0 = blockDim.x * gridDim.x;\n const int loop_step = loop_step0 * UNROLL_MC;\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = max(start_elem, totallines);\n const int aligned_elem = ((end_elem - start_elem) / loop_step) * loop_step;\n const int end_aligned = start_elem + aligned_elem;\n for (int line = start_elem; line < end_aligned; line += loop_step) {\n uint4 val[UNROLL_MC];\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++)\n#if defined(NVTE_UB_FP16)\n asm(\"multimem.ld_reduce.global.add.v4.f16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val[i].x), \"=r\"(val[i].y), \"=r\"(val[i].z), \"=r\"(val[i].w)\n : \"l\"(mc_ptr + (mylineoffset + line + i * loop_step0))\n : \"memory\");\n#else\n asm(\"multimem.ld_reduce.global.add.v4.bf16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val[i].x), \"=r\"(val[i].y), \"=r\"(val[i].z), \"=r\"(val[i].w)\n : \"l\"(mc_ptr + (mylineoffset + line + i * loop_step0))\n : \"memory\");\n#endif\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++)\n (reinterpret_cast(outbuf))[((line + i * loop_step0) / rowlines) * skiplines +\n ((line + i * loop_step0) % rowlines)] = val[i];\n }\n for (int line = end_aligned; line < end_elem; line += loop_step0) {\n uint4 val;\n#if defined(NVTE_UB_FP16)\n asm(\"multimem.ld_reduce.global.add.v4.f16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val.x), \"=r\"(val.y), \"=r\"(val.z), \"=r\"(val.w)\n : \"l\"(mc_ptr + (mylineoffset + line))\n : \"memory\");\n#else\n asm(\"multimem.ld_reduce.global.add.v4.bf16x2 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(val.x), \"=r\"(val.y), \"=r\"(val.z), \"=r\"(val.w)\n : \"l\"(mc_ptr + (mylineoffset + line))\n : \"memory\");\n#endif\n reinterpret_cast(outbuf)[(line / rowlines) * skiplines + (line % rowlines)] = val;\n }\n\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 reduce-scatter kernel (out of place) fp16 MC\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc_ag(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep,\n const int mylineoffset, const int totallines,\n void **commbuff, const int handleridx, uint4 *mc_ptr,\n const uint64_t ub_timeout) {\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n uint4 *localptr = reinterpret_cast(commbuff[myrank * gpustep + firstrank + handleridx]);\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n }\n __syncthreads();\n\n const int loop_step0 = blockDim.x * gridDim.x;\n const int loop_step = loop_step0 * UNROLL_MC;\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = max(start_elem, totallines);\n const int aligned_elem = ((end_elem - start_elem) / loop_step) * loop_step;\n const int end_aligned = start_elem + aligned_elem;\n for (int line = start_elem; line < end_aligned; line += loop_step) {\n uint4 val[UNROLL_MC];\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++) val[i] = localptr[mylineoffset + line + i * loop_step0];\n#pragma unroll\n for (int i = 0; i < UNROLL_MC; i++)\n asm volatile(\"multimem.st.global.v4.f32 [%0], {%1,%2,%3,%4};\" ::\"l\"(\n mc_ptr + (mylineoffset + line + i * loop_step0)),\n \"r\"(val[i].x), \"r\"(val[i].y), \"r\"(val[i].z), \"r\"(val[i].w)\n : \"memory\");\n }\n for (int line = end_aligned; line < end_elem; line += loop_step0) {\n uint4 val = localptr[mylineoffset + line];\n asm volatile(\n \"multimem.st.global.v4.f32 [%0], {%1,%2,%3,%4};\" ::\"l\"(mc_ptr + (mylineoffset + line)),\n \"r\"(val.x), \"r\"(val.y), \"r\"(val.z), \"r\"(val.w)\n : \"memory\");\n }\n\n __syncthreads();\n if (threadIdx.x == 0) __threadfence_system();\n __syncthreads();\n\n __shared__ int lastSM;\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id)\n lastSM = 1;\n else\n lastSM = 0;\n }\n __syncthreads();\n if (lastSM && threadIdx.x < RANKS) {\n if (threadIdx.x == 0) *reduceidptr = reduce_id;\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&myptr[targetgpu];\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Allgather: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x, threadIdx.x,\n reduce_id, *flag);\n break;\n }\n }\n }\n} // fp16 inplace allgather kernel (Hopper) MC\n\n#else\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep, const int lineoffset,\n const int numlines, void **commbuff, const int handleridx,\n float4 *mc_ptr, const uint64_t ub_timeout) {}\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc_rs_oop(const int op, const int flagoffset,\n const int firstrank, const int myrank,\n const int gpustep, const int mylineoffset,\n const int totallines, const int rowlines,\n const int skiplines, void **commbuff,\n const int handleridx, void *outbuf, float4 *mc_ptr,\n const uint64_t ub_timeout) {}\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc_ag(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep,\n const int mylineoffset, const int totallines,\n void **commbuff, const int handleridx, uint4 *mc_ptr,\n const uint64_t ub_timeout) {}\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_mc_rs(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep,\n const int mylineoffset, const int totallines,\n void **commbuff, const int handleridx, float4 *mc_ptr,\n const uint64_t ub_timeout) {}\n#endif\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS) userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_fp8(\n const int op, const int flagoffset, const int firstrank, const int myrank, const int gpustep,\n const int mylineoffset, const int totallines, const int rowlines, const int skiplines,\n void **commbuff, const int handleridx, void *outbuf, float *scale, const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n half_dtype hscale = (half_dtype)*scale;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n __syncthreads();\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS];\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][mylineoffset + line];\n }\n\n int4 sum[2] = {{0, 0, 0, 0}, {0, 0, 0, 0}};\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n fp8type *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < sizeof(int4) / sizeof(fp8type); j++) s[j] += hscale * (half_dtype)(x[j]);\n }\n int hline = 2 * line;\n (reinterpret_cast(outbuf))[(hline / rowlines) * skiplines + (hline % rowlines)] =\n sum[0];\n hline++;\n (reinterpret_cast(outbuf))[(hline / rowlines) * skiplines + (hline % rowlines)] =\n sum[1];\n }\n\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 reduce-scatter kernel (out of place) (fp8->fp16)\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_atomic_fp8(\n const int op, const int flagoffset, const int firstrank, const int myrank,\n const int gpustep, const int mylineoffset, const int totallines, const int rowlines,\n const int skiplines_out, const int skiplines_in, void **commbuff, const int handleridx,\n void *outbuf, float *scale, void *counters, const int numchunks, const int atomicindex,\n const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n half_dtype hscale = (half_dtype)*scale;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n // const int blockflagoffset = MAX_NVLINK * 2 * blockIdx.x;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr);\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset; // + blockflagoffset;\n }\n\n for (int chunk_i = 0; chunk_i < numchunks; chunk_i++) {\n ATOMIC_CONSUMER(chunk_i);\n\n lastSM = 0;\n if (threadIdx.x < RANKS) {\n reduce_id++;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), /*numchunks * */ adder);\n if (old_val + adder == NVTE_MAX_SMS * (reduce_id /* + numchunks*/)) lastSM = 1;\n }\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n __syncthreads();\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS];\n const int rowlines_in = rowlines / 2;\n const int index_in = skiplines_in == 0\n ? mylineoffset + myrank * totallines + line\n : (numchunks <= 1 ? 1 : chunk_i) * mylineoffset +\n myrank * (totallines * skiplines_in / rowlines_in) +\n (line / rowlines_in) * skiplines_in + (line % rowlines_in);\n const int index1_out = chunk_i * mylineoffset * 2 + ((2 * line) / rowlines) * skiplines_out +\n ((2 * line) % rowlines);\n const int index2_out = chunk_i * mylineoffset * 2 +\n ((2 * line + 1) / rowlines) * skiplines_out +\n ((2 * line + 1) % rowlines);\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][index_in];\n }\n\n int4 sum[2] = {{0, 0, 0, 0}, {0, 0, 0, 0}};\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n fp8type *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < sizeof(int4) / sizeof(fp8type); j++)\n s[j] += hscale * (half_dtype)(x[j]);\n }\n (reinterpret_cast(outbuf))[index1_out] = sum[0];\n (reinterpret_cast(outbuf))[index2_out] = sum[1];\n }\n }\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 reduce-scatter kernel (out of place) (fp8->fp16)\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS) userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_stride(\n const int op, const int flagoffset, const int firstrank, const int myrank, const int gpustep,\n const int mylineoffset, const int totallines, const int rowlines, const int skiplines,\n void **commbuff, const int handleridx, void *outbuf, const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS];\n int index_in = mylineoffset + myrank * (totallines * skiplines / rowlines) +\n (line / rowlines) * skiplines + (line % rowlines);\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][index_in];\n }\n\n int4 sum = val[0];\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 1; i < RANKS; i++) {\n half_dtype *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < 8; j++) s[j] += x[j];\n }\n\n int index_out = (line / rowlines) * skiplines + (line % rowlines);\n (reinterpret_cast(outbuf))[index_out] = sum;\n }\n\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 reduce-scatter kernel (out of place) fp16\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_stride_atomic(\n const int op, const int flagoffset, const int firstrank, const int myrank,\n const int gpustep, const int mylineoffset, const int totallines, const int rowlines,\n const int skiplines, const int numchunks, void **commbuff, const int handleridx,\n void *outbuf, void *counters, const uint64_t ub_timeout) {\n if (counters) {\n if (threadIdx.x == 0) {\n // spin-lock on counter from producer\n while (0 != (atomicCAS(((unsigned int *)counters), 0, 0))) {\n }\n\n // make sure all threadblocks have read/waited on counters.\n atomicInc(((unsigned int *)counters) + numchunks, gridDim.x - 1);\n while (0 != (atomicCAS(((unsigned int *)counters) + numchunks, 0, 0))) {\n }\n\n // reset counter for next producer.\n ((unsigned int *)counters)[0] = 1;\n asm volatile(\"fence.sc.gpu;\\n\");\n }\n }\n __syncthreads();\n\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS];\n int index_in = mylineoffset + myrank * (totallines * skiplines / rowlines) +\n (line / rowlines) * skiplines + (line % rowlines);\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][index_in];\n }\n\n int4 sum = val[0];\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 1; i < RANKS; i++) {\n half_dtype *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < 8; j++) s[j] += x[j];\n }\n\n int index_out = (line / rowlines) * skiplines + (line % rowlines);\n (reinterpret_cast(outbuf))[index_out] = sum;\n }\n\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n} // fp16 reduce-scatter kernel (out of place) fp16\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_stride_multiatomic(\n const int op, const int flagoffset, const int firstrank, const int myrank,\n const int gpustep, const int mylineoffset, const int totallines, const int rowlines,\n const int skiplines, const int numchunks, void **commbuff, const int handleridx,\n void *outbuf, void *counters, const uint64_t ub_timeout) {\n for (int chunk_i = 0; chunk_i < numchunks; chunk_i++) {\n if (counters) {\n if (threadIdx.x == 0) {\n // spin-lock on counter from producer\n while (0 != (atomicCAS(((unsigned int *)counters) + chunk_i, 0, 0))) {\n }\n\n // make sure all threadblocks have read/waited on counters.\n atomicInc(((unsigned int *)counters) + numchunks + chunk_i, gridDim.x - 1);\n while (0 != (atomicCAS(((unsigned int *)counters) + numchunks + chunk_i, 0, 0))) {\n }\n\n // reset counter for next producer.\n ((unsigned int *)counters)[chunk_i] = 1;\n asm volatile(\"fence.sc.gpu;\\n\");\n }\n }\n __syncthreads();\n\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int lastSM = 0;\n\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n if (blockIdx.x == 0) flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&(myptr[targetgpu]);\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Reduce-scatter: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x,\n threadIdx.x, reduce_id, *flag);\n break;\n }\n }\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id) lastSM = 1;\n }\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n#pragma unroll\n for (int i = 0; i < RANKS; i++) dest[i] = (i + myrank + warp) & (RANKS - 1);\n\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS];\n int index_in = chunk_i * mylineoffset + myrank * (totallines * skiplines / rowlines) +\n (line / rowlines) * skiplines + (line % rowlines);\n\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n val[i] = userptr[dest[i]][index_in];\n }\n\n int4 sum = val[0];\n half_dtype *s = reinterpret_cast(&sum);\n\n#pragma unroll\n for (int i = 1; i < RANKS; i++) {\n half_dtype *x = reinterpret_cast(&val[i]);\n#pragma unroll\n for (int j = 0; j < 8; j++) s[j] += x[j];\n }\n\n int index_out = chunk_i * mylineoffset + (line / rowlines) * skiplines + (line % rowlines);\n (reinterpret_cast(outbuf))[index_out] = sum;\n }\n if (threadIdx.x == 0 && lastSM) *reduceidptr = reduce_id;\n }\n} // fp16 reduce-scatter kernel (out of place) fp16\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rr_ag(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep,\n const int mylineoffset, const int totallines,\n void **commbuff, const int handleridx,\n const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n clock_t s = clock64();\n }\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS];\n\n int skipmy = 0;\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n int dst = (i + warp + myrank) & (RANKS - 1);\n if (dst == myrank) {\n skipmy++;\n continue;\n }\n dest[i - skipmy] = dst;\n }\n __syncthreads();\n\n for (int line = threadIdx.x + blockDim.x * blockIdx.x; line < totallines;\n line += blockDim.x * gridDim.x) {\n int4 val[RANKS - 1];\n\n#pragma unroll\n for (int i = 0; i < RANKS - 1; i++) {\n val[i] = userptr[dest[i]][mylineoffset + line + totallines * dest[i]];\n }\n\n#pragma unroll\n for (int i = 0; i < RANKS - 1; i++) {\n userptr[myrank][mylineoffset + line + totallines * dest[i]] = val[i];\n }\n }\n __shared__ int lastSM;\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id)\n lastSM = 1;\n else\n lastSM = 0;\n }\n __syncthreads();\n if (lastSM && threadIdx.x < RANKS) {\n if (threadIdx.x == 0) *reduceidptr = reduce_id;\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&myptr[targetgpu];\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Allgather: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x, threadIdx.x,\n reduce_id, *flag);\n break;\n }\n }\n }\n} // fp16 inplace reduce kernel (Ampere)\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS)\n userbuffers_fp16_sum_inplace_gpu_rw_ag(const int op, const int flagoffset, const int firstrank,\n const int myrank, const int gpustep,\n const int mylineoffset, const int totallines,\n void **commbuff, const int handleridx,\n const uint64_t ub_timeout) {\n __shared__ int4 *userptr[RANKS];\n volatile int *flagptr;\n int physgpu, targetgpu, *myptr;\n int *reduceidptr, reduce_id;\n int4 *localptr;\n if (threadIdx.x < RANKS) {\n physgpu = myrank * gpustep + firstrank;\n targetgpu = threadIdx.x * gpustep + firstrank;\n myptr = (reinterpret_cast(commbuff[physgpu])) + flagoffset;\n reduceidptr = myptr - NVTE_MAX_OPS; // +op;\n reduce_id = (*reduceidptr) + 1;\n flagptr = (reinterpret_cast(commbuff[targetgpu])) + flagoffset;\n userptr[threadIdx.x] = reinterpret_cast(commbuff[targetgpu + handleridx]);\n }\n __syncthreads();\n localptr = userptr[myrank];\n\n int warp = blockIdx.x + (threadIdx.x >> 5);\n int dest[RANKS - 1];\n int skipmy = 0;\n#pragma unroll\n for (int i = 0; i < RANKS; i++) {\n int dst = (i + warp + myrank) & (RANKS - 1);\n if (dst == myrank) {\n skipmy++;\n continue;\n }\n dest[i - skipmy] = dst;\n }\n#define UNROLLAG 4\n __syncthreads();\n const int loop_step0 = blockDim.x * gridDim.x;\n const int loop_step = loop_step0 * UNROLLAG;\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = max(start_elem, totallines);\n const int aligned_elem = ((end_elem - start_elem) / loop_step) * loop_step;\n const int end_aligned = start_elem + aligned_elem;\n\n for (int line = start_elem; line < end_aligned; line += loop_step) {\n int4 val[UNROLLAG];\n#pragma unroll\n for (int j = 0; j < UNROLLAG; j++) val[j] = localptr[mylineoffset + line + loop_step0 * j];\n\n#pragma unroll\n for (int j = 0; j < UNROLLAG; j++)\n#pragma unroll\n for (int i = 0; i < RANKS - 1; i++) {\n userptr[dest[i]][mylineoffset + line + j * loop_step0] = val[j];\n }\n }\n\n for (int line = end_aligned; line < end_elem; line += loop_step0) {\n int4 sum = localptr[mylineoffset + line];\n#pragma unroll\n for (int i = 0; i < RANKS - 1; i++) {\n userptr[dest[i]][mylineoffset + line] = sum;\n }\n }\n\n __syncthreads();\n if (threadIdx.x == 0) __threadfence_system();\n __syncthreads();\n\n __shared__ int lastSM;\n if (threadIdx.x == 0) {\n const int adder = blockIdx.x == 0 ? NVTE_MAX_SMS - gridDim.x + 1 : 1;\n int old_val = atomicAdd(myptr + (NVTE_MAX_NVLINK * 2), adder);\n if (old_val + adder == NVTE_MAX_SMS * reduce_id)\n lastSM = 1;\n else\n lastSM = 0;\n }\n __syncthreads();\n if (lastSM && threadIdx.x < RANKS) {\n if (threadIdx.x == 0) *reduceidptr = reduce_id;\n flagptr[physgpu] = reduce_id;\n volatile int *flag = (volatile int *)&myptr[targetgpu];\n clock_t s = clock64();\n while (CHECK_IDS(*flag, reduce_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\"[%d] Allgather: SM %d [%d]: expecting %d got %d\", myrank, blockIdx.x, threadIdx.x,\n reduce_id, *flag);\n break;\n }\n }\n }\n} // fp16 inplace allgather kernel (Volta,Hopper)\n\n#define SETUP_LAUNCH_CONFIG(sms, threads, stream) \\\n hipLaunchConfig_t cfg = {sms, threads, 0, stream, NULL, 0}; \\\n hipLaunchAttribute attribute_ub[2]; \\\n attribute_ub[1].id = cudaLaunchAttributeClusterDimension; \\\n attribute_ub[1].val.clusterDim.x = sms % comm->cga_size == 0 ? comm->cga_size : 1; \\\n attribute_ub[1].val.clusterDim.y = 1; \\\n attribute_ub[1].val.clusterDim.z = 1; \\\n attribute_ub[0].id = hipLaunchAttributeCooperative; \\\n cfg.attrs = attribute_ub; \\\n cfg.numAttrs = comm->sm_arch >= 9 ? 2 : 1;\n\n#if (CUDART_VERSION >= 12030)\n#define ADD_LAUNCH_COMPLETION_EVENT(attribute_ub, comm_launch_event) \\\n attribute_ub[2].id = cudaLaunchAttributeLaunchCompletionEvent; \\\n attribute_ub[2].val.launchCompletionEvent.event = comm_launch_event;\n#define NUM_LAUNCH_ATTRIBUTE_FOR_FDL_LAUNCH 3\n#else\n#define ADD_LAUNCH_COMPLETION_EVENT(attribute_ub, comm_launch_event)\n#define NUM_LAUNCH_ATTRIBUTE_FOR_FDL_LAUNCH 2\n#endif\n\n#define SETUP_LAUNCH_CONFIG_WITH_COMPLETION_EVENT(sms, threads, stream, comm_launch_event) \\\n hipLaunchConfig_t cfg = {sms, threads, 0, stream, NULL, 0}; \\\n hipLaunchAttribute attribute_ub[NUM_LAUNCH_ATTRIBUTE_FOR_FDL_LAUNCH] = {}; \\\n ADD_LAUNCH_COMPLETION_EVENT(attribute_ub, comm_launch_event) \\\n attribute_ub[1].id = cudaLaunchAttributeClusterDimension; \\\n attribute_ub[1].val.clusterDim.x = sms % comm->cga_size == 0 ? comm->cga_size : 1; \\\n attribute_ub[1].val.clusterDim.y = 1; \\\n attribute_ub[1].val.clusterDim.z = 1; \\\n attribute_ub[0].id = hipLaunchAttributeCooperative; \\\n cfg.attrs = attribute_ub; \\\n cfg.numAttrs = NUM_LAUNCH_ATTRIBUTE_FOR_FDL_LAUNCH;\n\n#define callranks_ag(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8 + (comm->use_rr_kernel ? 0 : arg4 * arg7); \\\n void **arg8 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg9 = handler * comm->nvsize; \\\n uint64_t arg10 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10)}; \\\n NVTE_CHECK_CUDA(hipLaunchKernelExC( \\\n &cfg, \\\n reinterpret_cast(comm->use_rr_kernel ? userbuffers_fp16_sum_inplace_gpu_rr_ag \\\n : userbuffers_fp16_sum_inplace_gpu_rw_ag), \\\n kernelArgs)); \\\n }\n\n#define callranks_agMC(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8 + arg4 * arg7; \\\n void **arg8 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg9 = handler * comm->nvsize; \\\n uint4 *arg10 = reinterpret_cast(comm->mc_ptr[handler]); \\\n uint64_t arg11 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11)}; \\\n NVTE_CHECK_CUDA(hipLaunchKernelExC( \\\n &cfg, reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_mc_ag), kernelArgs)); \\\n }\n\n#define callranks_rs(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8 + arg4 * arg7; \\\n void **arg8 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg9 = handler * comm->nvsize; \\\n uint64_t arg10 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10)}; \\\n NVTE_CHECK_CUDA(hipLaunchKernelExC( \\\n &cfg, reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_rr_rs), kernelArgs)); \\\n }\n\n#define callranks_rsMC(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8 + arg4 * arg7; \\\n void **arg8 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg9 = handler * comm->nvsize; \\\n void *arg10 = comm->mc_ptr[handler]; \\\n uint64_t arg11 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11)}; \\\n NVTE_CHECK_CUDA(hipLaunchKernelExC( \\\n &cfg, reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_mc_rs), kernelArgs)); \\\n }\n\n#define callranks_rs_oop(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8 + arg4 * arg7, arg8 = rowelements / 8, arg9 = strideelements / 8; \\\n void **arg10 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg11 = handler * comm->nvsize; \\\n void *arg12 = output; \\\n uint64_t arg13 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11), reinterpret_cast(&arg12), \\\n reinterpret_cast(&arg13)}; \\\n NVTE_CHECK_CUDA(hipLaunchKernelExC( \\\n &cfg, reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_rr_rs_oop), \\\n kernelArgs)); \\\n }\n\n#define callranks_rs_oop_fp8(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 16 / x, \\\n arg6 = offset / 16 + arg4 * arg7, arg8 = rowelements / 8, arg9 = strideelements / 8; \\\n void **arg10 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg11 = handler * comm->nvsize; \\\n void *arg12 = output; \\\n float *arg13 = scale; \\\n uint64_t arg14 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11), reinterpret_cast(&arg12), \\\n reinterpret_cast(&arg13), reinterpret_cast(&arg14)}; \\\n NVTE_CHECK_CUDA(hipLaunchKernelExC( \\\n &cfg, \\\n reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_fp8), \\\n kernelArgs)); \\\n }\n\n#define callranks_rs_oopMC(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8 + arg4 * arg7, arg8 = rowelements / 8, arg9 = strideelements / 8; \\\n void **arg10 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg11 = handler * comm->nvsize; \\\n void *arg12 = output; \\\n void *arg13 = comm->mc_ptr[handler]; \\\n uint64_t arg14 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11), reinterpret_cast(&arg12), \\\n reinterpret_cast(&arg13), reinterpret_cast(&arg14)}; \\\n NVTE_CHECK_CUDA(hipLaunchKernelExC( \\\n &cfg, reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_mc_rs_oop), \\\n kernelArgs)); \\\n }\n\n#define callranks_rs_oop_atomic_fp8(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 16 / x, \\\n arg6 = offset / 16, arg8 = rowelements / 8, arg9 = strideelements_out / 8, \\\n arg10 = strideelements_in / 16; \\\n void **arg11 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg12 = handler * comm->nvsize; \\\n void *arg13 = output; \\\n float *arg14 = scale; \\\n void *arg15 = counters; \\\n int arg16 = numchunks, arg17 = atomicindex; \\\n uint64_t arg18 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11), reinterpret_cast(&arg12), \\\n reinterpret_cast(&arg13), reinterpret_cast(&arg14), \\\n reinterpret_cast(&arg15), reinterpret_cast(&arg16), \\\n reinterpret_cast(&arg17), reinterpret_cast(&arg18)}; \\\n NVTE_CHECK_CUDA(hipLaunchKernelExC( \\\n &cfg, \\\n reinterpret_cast( \\\n userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_atomic_fp8), \\\n kernelArgs)); \\\n }\n\n#define callranks_rs_oop_stride(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8, arg8 = rowelements / 8, arg9 = strideelements / 8; \\\n void **arg10 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg11 = handler * comm->nvsize; \\\n void *arg12 = output; \\\n uint64_t arg13 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11), reinterpret_cast(&arg12), \\\n reinterpret_cast(&arg13)}; \\\n NVTE_CHECK_CUDA(hipLaunchKernelExC( \\\n &cfg, reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_stride), \\\n kernelArgs)); \\\n }\n\n#define callranks_rs_oop_stride_atomic(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8, arg8 = rowelements / 8, arg9 = strideelements / 8, arg10 = numchunks; \\\n void **arg11 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg12 = handler * comm->nvsize; \\\n void *arg13 = output; \\\n void *arg14 = counters; \\\n uint64_t arg15 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11), reinterpret_cast(&arg12), \\\n reinterpret_cast(&arg13), reinterpret_cast(&arg14), \\\n reinterpret_cast(&arg15)}; \\\n NVTE_CHECK_CUDA(hipLaunchKernelExC( \\\n &cfg, \\\n reinterpret_cast(userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_stride_atomic), \\\n kernelArgs)); \\\n }\n\n#define callranks_rs_oop_stride_multiatomic(x) \\\n if (ar_nvsize == x) { \\\n int arg1 = op - NVTE_MAX_OPS, \\\n arg2 = NVTE_REG0_OFFSET(comm) - \\\n (op == userbuffers_allreduceop_nonsharp ? 2 : 1) * NVTE_REG0_SINGLENODE + \\\n NVTE_MAX_OPS, \\\n arg3 = ar_firstgpu, arg4 = ar_nvrank, arg5 = ar_step, arg7 = elements / 8 / x, \\\n arg6 = offset / 8, arg8 = rowelements / 8, arg9 = strideelements / 8, arg10 = numchunks; \\\n void **arg11 = reinterpret_cast(comm->gpu_ptrs); \\\n int arg12 = handler * comm->nvsize; \\\n void *arg13 = output; \\\n void *arg14 = counters; \\\n uint64_t arg15 = comm->ub_timeout; \\\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2), \\\n reinterpret_cast(&arg3), reinterpret_cast(&arg4), \\\n reinterpret_cast(&arg5), reinterpret_cast(&arg6), \\\n reinterpret_cast(&arg7), reinterpret_cast(&arg8), \\\n reinterpret_cast(&arg9), reinterpret_cast(&arg10), \\\n reinterpret_cast(&arg11), reinterpret_cast(&arg12), \\\n reinterpret_cast(&arg13), reinterpret_cast(&arg14), \\\n reinterpret_cast(&arg15)}; \\\n NVTE_CHECK_CUDA( \\\n hipLaunchKernelExC(&cfg, \\\n reinterpret_cast( \\\n userbuffers_fp16_sum_inplace_gpu_rr_rs_oop_stride_multiatomic), \\\n kernelArgs)); \\\n }\n\nvoid reducescatter2_userbuff_strided(void *output, const int handler, const int offset,\n const int rowelements, const int colelements,\n const int strideelements, communicator *comm,\n hipStream_t stream) {\n const int elements = rowelements * colelements;\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n\n if (elements < 64) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n callranks_rs_oop_stride(2) callranks_rs_oop_stride(4) callranks_rs_oop_stride(8)\n callranks_rs_oop_stride(16) callranks_rs_oop_stride(32)\n}\nvoid reducescatter2_userbuff_strided_atomic(void *output, const int handler, const int offset,\n const int rowelements, const int colelements,\n const int strideelements, const int numchunks,\n void *counters, communicator *comm,\n hipStream_t stream) {\n const int elements = rowelements * colelements;\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n\n if (elements < 64) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n callranks_rs_oop_stride_atomic(2) callranks_rs_oop_stride_atomic(4)\n callranks_rs_oop_stride_atomic(8) callranks_rs_oop_stride_atomic(16)\n callranks_rs_oop_stride_atomic(32)\n}\n\ntemplate \nvoid reducescatter2_userbuff_strided_universal_fp8(void *output, float *scale, const int handler,\n const int offset, const int rowelements,\n const int colelements,\n const int strideelements_out,\n const int strideelements_in, const int numchunks,\n const int atomicindex, void *counters,\n communicator *comm, hipStream_t stream) {\n const int elements = rowelements * colelements;\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n assert(comm->sm_arch >= 9);\n if (elements < 128) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n callranks_rs_oop_atomic_fp8(2) callranks_rs_oop_atomic_fp8(4) callranks_rs_oop_atomic_fp8(8)\n callranks_rs_oop_atomic_fp8(16) callranks_rs_oop_atomic_fp8(32)\n}\n\ntemplate \nvoid reducescatter2_userbuff_strided_atomic_fp8(void *output, float *scale, const int handler,\n const int offset, const int rowelements,\n const int colelements, const int strideelements_out,\n const int strideelements_in, const int numchunks,\n void *counters, communicator *comm,\n hipStream_t stream) {\n reducescatter2_userbuff_strided_universal_fp8(\n output, scale, handler, offset, rowelements, colelements, strideelements_out,\n strideelements_in, 1, numchunks, counters /*nullptr*/, comm, stream);\n}\n\ntemplate \nvoid reducescatter2_userbuff_strided_multiatomic_fp8(\n void *output, float *scale, const int handler, const int offset, const int rowelements,\n const int colelements, const int strideelements_out, const int strideelements_in,\n const int numchunks, void *counters, communicator *comm, hipStream_t stream) {\n reducescatter2_userbuff_strided_universal_fp8(\n output, scale, handler, offset, rowelements, colelements, strideelements_out,\n strideelements_in, numchunks, 0, counters /*nullptr*/, comm, stream);\n}\n\nvoid reducescatter2_userbuff_strided_multiatomic(void *output, const int handler, const int offset,\n const int rowelements, const int colelements,\n const int strideelements, const int numchunks,\n void *counters, communicator *comm,\n hipStream_t stream) {\n const int elements = rowelements * colelements;\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n\n if (elements < 64) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n callranks_rs_oop_stride_multiatomic(2) callranks_rs_oop_stride_multiatomic(4)\n callranks_rs_oop_stride_multiatomic(8) callranks_rs_oop_stride_multiatomic(16)\n callranks_rs_oop_stride_multiatomic(32)\n}\n\nvoid allgather2_userbuff_inplace(const int handler, const int offset, const int elements,\n communicator *comm, hipStream_t stream,\n hipEvent_t comm_launch_event) {\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n\n if (elements < 64) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n if (comm_launch_event) {\n SETUP_LAUNCH_CONFIG_WITH_COMPLETION_EVENT(sms, warps * 32, stream, comm_launch_event);\n if (comm->use_mc && (comm->memflags[handler] & NVTE_UB_MEM_MC_CREATED)) {\n callranks_agMC(2) callranks_agMC(4) callranks_agMC(8) callranks_agMC(16) callranks_agMC(32)\n } else {\n callranks_ag(2) callranks_ag(4) callranks_ag(8) callranks_ag(16) callranks_ag(32)\n }\n } else {\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n if (comm->use_mc && (comm->memflags[handler] & NVTE_UB_MEM_MC_CREATED)) {\n callranks_agMC(2) callranks_agMC(4) callranks_agMC(8) callranks_agMC(16) callranks_agMC(32)\n } else {\n callranks_ag(2) callranks_ag(4) callranks_ag(8) callranks_ag(16) callranks_ag(32)\n }\n }\n}\n\nvoid allgather2_userbuff_inplace_sliced(const int handler, const int offset, const int elements,\n communicator *comm, const int slice_id, const int nslices,\n hipStream_t stream) {\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n int peerelements = elements / ar_nvsize;\n int saverrkernel = comm->use_rr_kernel;\n comm->use_rr_kernel = 0;\n allgather2_userbuff_inplace(\n handler, offset + ar_nvrank * peerelements * (nslices - 1) + slice_id * peerelements,\n elements, comm, stream);\n comm->use_rr_kernel = saverrkernel;\n}\n\nvoid reducescatter2_userbuff_inplace(const int handler, const int offset, const int elements,\n communicator *comm, hipStream_t stream,\n hipEvent_t comm_launch_event) {\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n\n if (elements < 64) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n if (comm_launch_event) {\n SETUP_LAUNCH_CONFIG_WITH_COMPLETION_EVENT(sms, warps * 32, stream, comm_launch_event);\n if (comm->use_mc && (comm->memflags[handler] & NVTE_UB_MEM_MC_CREATED)) {\n callranks_rsMC(2) callranks_rsMC(4) callranks_rsMC(8) callranks_rsMC(16) callranks_rsMC(32)\n } else {\n callranks_rs(2) callranks_rs(4) callranks_rs(8) callranks_rs(16) callranks_rs(32)\n }\n } else {\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n if (comm->use_mc && (comm->memflags[handler] & NVTE_UB_MEM_MC_CREATED)) {\n callranks_rsMC(2) callranks_rsMC(4) callranks_rsMC(8) callranks_rsMC(16) callranks_rsMC(32)\n } else {\n callranks_rs(2) callranks_rs(4) callranks_rs(8) callranks_rs(16) callranks_rs(32)\n }\n }\n}\nvoid reducescatter2_userbuff_stridedoutput(void *output, const int handler, const int offset,\n const int rowelements, const int colelements,\n const int strideelements, communicator *comm,\n hipStream_t stream, hipEvent_t comm_launch_event) {\n const int elements = rowelements * colelements;\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n\n if (elements < 64) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n if (comm_launch_event) {\n SETUP_LAUNCH_CONFIG_WITH_COMPLETION_EVENT(sms, warps * 32, stream, comm_launch_event);\n if (comm->use_mc && (comm->memflags[handler] & NVTE_UB_MEM_MC_CREATED)) {\n callranks_rs_oopMC(2) callranks_rs_oopMC(4) callranks_rs_oopMC(8) callranks_rs_oopMC(16)\n callranks_rs_oopMC(32)\n } else {\n callranks_rs_oop(2) callranks_rs_oop(4) callranks_rs_oop(8) callranks_rs_oop(16)\n callranks_rs_oop(32)\n }\n } else {\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n if (comm->use_mc && (comm->memflags[handler] & NVTE_UB_MEM_MC_CREATED)) {\n callranks_rs_oopMC(2) callranks_rs_oopMC(4) callranks_rs_oopMC(8) callranks_rs_oopMC(16)\n callranks_rs_oopMC(32)\n } else {\n callranks_rs_oop(2) callranks_rs_oop(4) callranks_rs_oop(8) callranks_rs_oop(16)\n callranks_rs_oop(32)\n }\n }\n}\nvoid reducescatter2_userbuff(void *output, const int handler, const int offset, const int elements,\n communicator *comm, hipStream_t stream,\n hipEvent_t comm_launch_event) {\n reducescatter2_userbuff_stridedoutput(output, handler, offset, elements, 1, 0, comm, stream,\n comm_launch_event);\n}\n\ntemplate \nvoid reducescatter2_userbuff_stridedoutput_fp8(void *output, float *scale, const int handler,\n const int offset, const int rowelements,\n const int colelements, const int strideelements,\n communicator *comm, hipStream_t stream,\n hipEvent_t comm_launch_event) {\n const int elements = rowelements * colelements;\n const int op = userbuffers_allreduceop_nonsharp2;\n const int ar_firstgpu =\n op == userbuffers_allreduceop_nonsharp ? comm->ar_firstgpu : comm->ar2_firstgpu;\n const int ar_step = op == userbuffers_allreduceop_nonsharp2 ? 1 : comm->ar2_nvsize;\n const int ar_nvsize = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvsize : comm->ar2_nvsize;\n const int ar_nvrank = op == userbuffers_allreduceop_nonsharp ? comm->ar_nvrank : comm->ar2_nvrank;\n assert(comm->sm_arch >= 9);\n if (elements < 128) return;\n int sms = ar_nvsize == 1 ? 2 : comm->sms;\n int warps = comm->threads / 32;\n if (warps < ar_nvsize) warps = ar_nvsize;\n\n if (comm_launch_event) {\n SETUP_LAUNCH_CONFIG_WITH_COMPLETION_EVENT(sms, warps * 32, stream, comm_launch_event);\n callranks_rs_oop_fp8(2) callranks_rs_oop_fp8(4) callranks_rs_oop_fp8(8) callranks_rs_oop_fp8(16)\n callranks_rs_oop_fp8(32)\n } else {\n SETUP_LAUNCH_CONFIG(sms, warps * 32, stream);\n callranks_rs_oop_fp8(2) callranks_rs_oop_fp8(4) callranks_rs_oop_fp8(8) callranks_rs_oop_fp8(16)\n callranks_rs_oop_fp8(32)\n }\n}\n\ntemplate void reducescatter2_userbuff_stridedoutput_fp8<__hip_fp8_e5m2_fnuz>(\n void *output, float *scale, const int handler, const int offset, const int rowelements,\n const int colelements, const int strideelements, communicator *comm, hipStream_t stream,\n hipEvent_t comm_launch_event);\n\ntemplate void reducescatter2_userbuff_stridedoutput_fp8<__hip_fp8_e4m3_fnuz>(\n void *output, float *scale, const int handler, const int offset, const int rowelements,\n const int colelements, const int strideelements, communicator *comm, hipStream_t stream,\n hipEvent_t comm_launch_event);\n\ntemplate \nvoid reducescatter2_userbuff_fp8(void *output, float *scale, const int handler, const int offset,\n const int elements, communicator *comm, hipStream_t stream,\n hipEvent_t comm_launch_event) {\n reducescatter2_userbuff_stridedoutput_fp8(output, scale, handler, offset, elements, 1, 0,\n comm, stream, comm_launch_event);\n}\n\ntemplate void reducescatter2_userbuff_fp8<__hip_fp8_e5m2_fnuz>(void *output, float *scale,\n const int handler, const int offset,\n const int elements, communicator *comm,\n hipStream_t stream,\n hipEvent_t comm_launch_event);\ntemplate void reducescatter2_userbuff_fp8<__hip_fp8_e4m3_fnuz>(void *output, float *scale,\n const int handler, const int offset,\n const int elements, communicator *comm,\n hipStream_t stream,\n hipEvent_t comm_launch_event);\n\ntemplate void reducescatter2_userbuff_strided_atomic_fp8<__hip_fp8_e4m3_fnuz>(\n void *output, float *scale, const int handler, const int offset, const int rowelements,\n const int colelements, const int strideelements_out, const int strideelements_in,\n const int numchunks, void *counters, communicator *comm, hipStream_t stream);\ntemplate void reducescatter2_userbuff_strided_atomic_fp8<__hip_fp8_e5m2_fnuz>(\n void *output, float *scale, const int handler, const int offset, const int rowelements,\n const int colelements, const int strideelements_out, const int strideelements_in,\n const int numchunks, void *counters, communicator *comm, hipStream_t stream);\ntemplate void reducescatter2_userbuff_strided_multiatomic_fp8<__hip_fp8_e4m3_fnuz>(\n void *output, float *scale, const int handler, const int offset, const int rowelements,\n const int colelements, const int strideelements_out, const int strideelements_in,\n const int numchunks, void *counters, communicator *comm, hipStream_t stream);\ntemplate void reducescatter2_userbuff_strided_multiatomic_fp8<__hip_fp8_e5m2_fnuz>(\n void *output, float *scale, const int handler, const int offset, const int rowelements,\n const int colelements, const int strideelements_out, const int strideelements_in,\n const int numchunks, void *counters, communicator *comm, hipStream_t stream);\n\n__global__ void kuserbuffers_pullsend(int myrank, int peer, int *send_id, int *flagptr) {\n atomicAdd_system(flagptr, 1);\n}\n\n__global__ void kuserbuffers_inc(int *id) { atomicAdd(id, 1); }\n\n__global__ void kuserbuffers_dummy(void) {}\n\n__global__ void __launch_bounds__(MAX_THREADS)\n kuserbuffers_pullrecv(int myrank, int peer, int nvrank, int nvpeer, int *recv_id, int *flagptr,\n int4 *srcptr, int4 *dstptr, const int lines, uint64_t ub_timeout) {\n#define UNROLLCOPY 8\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = lines;\n const int aligned_elem = (end_elem - start_elem) & (~(blockDim.x * gridDim.x * UNROLLCOPY - 1));\n const int end_aligned = start_elem + aligned_elem;\n\n if (threadIdx.x == 0) {\n const int signal_id = (*recv_id) + 1;\n volatile int *flag = (volatile int *)flagptr;\n clock_t s = clock64();\n while (CHECK_IDS(*flag, signal_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\n \"pullrecv [grank dst:%d global src:%d][nvrank(GPU) dst: %d src: %d]: expecting %d,\"\n \" observed %d\",\n myrank, peer, nvrank, nvpeer, signal_id, *flag);\n break;\n }\n }\n if (lines == 0) {\n *recv_id = signal_id;\n return;\n } // otherwise need an extra kernel\n }\n __syncthreads();\n\n if (end_elem <= start_elem) return;\n\n for (int line = start_elem; line < end_aligned; line += blockDim.x * gridDim.x * UNROLLCOPY) {\n int4 val[UNROLLCOPY];\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) val[i] = srcptr[line + i * blockDim.x * gridDim.x];\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) dstptr[line + i * blockDim.x * gridDim.x] = val[i];\n }\n for (int line = end_aligned; line < end_elem; line += blockDim.x * gridDim.x)\n dstptr[line] = srcptr[line];\n}\n\n__global__ void __launch_bounds__(MAX_THREADS)\n kuserbuffers_pushsend(int *send_id, int *flagptr, int4 *srcptr, int4 *dstptr, const int lines) {\n if (lines) {\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = lines;\n const int aligned_elem =\n ((end_elem - start_elem) & (~(blockDim.x * gridDim.x * UNROLLCOPY - 1)));\n const int end_aligned = start_elem + aligned_elem;\n if (end_elem > start_elem) {\n for (int line = start_elem; line < end_aligned; line += blockDim.x * gridDim.x * UNROLLCOPY) {\n int4 val[UNROLLCOPY];\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) val[i] = srcptr[line + i * blockDim.x * gridDim.x];\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) dstptr[line + i * blockDim.x * gridDim.x] = val[i];\n }\n for (int line = end_aligned; line < end_elem; line += blockDim.x * gridDim.x)\n dstptr[line] = srcptr[line];\n }\n __syncthreads();\n if (threadIdx.x) return;\n __threadfence_system();\n atomicAdd_system(flagptr,\n 1); // otherwise need local SM sync before sending flag\n } else { // 0 bytes and 1 SM only\n atomicAdd_system(flagptr, 1);\n }\n}\n\n#define CHECK_CE(ce_start, ce_end) \\\n ((ce_start) != nullptr && (ce_end) != nullptr && *(ce_start) != *(ce_end))\n\n__global__ void kuserbuffers_pushrecv(int myrank, int peer, int nvrank, int nvpeer, int *recv_id,\n int *flagptr, int adder, uint64_t ub_timeout,\n int *ce_start_ptr, int *ce_end_ptr) {\n const int signal_id = (*recv_id) + adder;\n *recv_id = signal_id;\n volatile int *flag = (volatile int *)flagptr;\n if (*flag >= signal_id) return;\n clock_t s = clock64();\n while (CHECK_IDS(*flag, signal_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\n \"pushrecv [grank dst:%d global src:%d][nvrank(GPU) dst: %d src: %d]: \"\n \"expecting %d, observed %d\",\n myrank, peer, nvrank, nvpeer, signal_id, *flag);\n if (CHECK_CE(ce_start_ptr, ce_end_ptr))\n UB_PRINT(\"pushrecv: CE deadlock DETECTED: %d (ce_start) != %d (ce_end)\\n\", *ce_start_ptr,\n *ce_end_ptr);\n return;\n }\n }\n}\n\n__global__ void __launch_bounds__(MAX_THREADS)\n kuserbuffers_pushsendrecv(int *send_id, int *send_flagptr, int4 *srcptr, int4 *dstptr,\n const int lines, int send_peer, int recv_peer, int *recv_id,\n int *recv_flagptr, int adder, uint64_t ub_timeout, int nv_send,\n int nv_recv, int *ce_start_ptr, int *ce_end_ptr) {\n if (lines) {\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = lines;\n const int aligned_elem =\n ((end_elem - start_elem) & (~(blockDim.x * gridDim.x * UNROLLCOPY - 1)));\n const int end_aligned = start_elem + aligned_elem;\n if (end_elem > start_elem) {\n for (int line = start_elem; line < end_aligned; line += blockDim.x * gridDim.x * UNROLLCOPY) {\n int4 val[UNROLLCOPY];\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) {\n val[i] = srcptr[line + i * blockDim.x * gridDim.x];\n }\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) {\n dstptr[line + i * blockDim.x * gridDim.x] = val[i];\n }\n }\n for (int line = end_aligned; line < end_elem; line += blockDim.x * gridDim.x) {\n dstptr[line] = srcptr[line];\n }\n }\n __syncthreads();\n if (threadIdx.x) return;\n __threadfence_system();\n atomicAdd_system(send_flagptr,\n 1); // otherwise need local SM sync before sending flag\n } else { // 0 bytes and 1 SM only\n atomicAdd_system(send_flagptr, 1);\n }\n\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n const int signal_id = (*recv_id) + adder;\n *recv_id = signal_id;\n volatile int *flag = (volatile int *)recv_flagptr;\n if (*flag >= signal_id) return;\n clock_t s = clock64();\n while (CHECK_IDS(*flag, signal_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\n \"pushsendrecv [sending peer:%d receiving peer:%d][nvrank(GPU) sending peer: %d\"\n \" receiving peer: %d]: expecting %d, observed %d\",\n send_peer, recv_peer, nv_send, nv_recv, signal_id, *flag);\n if (CHECK_CE(ce_start_ptr, ce_end_ptr))\n UB_PRINT(\"pushrecv: CE deadlock DETECTED: %d (ce_start) != %d (ce_end)\\n\", *ce_start_ptr,\n *ce_end_ptr);\n return;\n }\n }\n }\n}\n\n__global__ void __launch_bounds__(MAX_THREADS)\n kuserbuffers_pushsendrecv_atomic(int *send_id, int *send_flagptr, int4 *srcptr, int4 *dstptr,\n const int lines, int send_peer, int recv_peer, int *recv_id,\n int *recv_flagptr, int adder, void *counters,\n uint64_t ub_timeout, int nv_send, int nv_recv,\n int *ce_start_ptr, int *ce_end_ptr) {\n if (lines) {\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = lines;\n const int aligned_elem =\n ((end_elem - start_elem) & (~(blockDim.x * gridDim.x * UNROLLCOPY - 1)));\n const int end_aligned = start_elem + aligned_elem;\n if (end_elem > start_elem) {\n for (int line = start_elem; line < end_aligned; line += blockDim.x * gridDim.x * UNROLLCOPY) {\n int4 val[UNROLLCOPY];\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) {\n val[i] = srcptr[line + i * blockDim.x * gridDim.x];\n }\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) {\n dstptr[line + i * blockDim.x * gridDim.x] = val[i];\n }\n }\n for (int line = end_aligned; line < end_elem; line += blockDim.x * gridDim.x) {\n dstptr[line] = srcptr[line];\n }\n }\n __syncthreads();\n if (threadIdx.x) return;\n __threadfence_system();\n atomicAdd_system(send_flagptr,\n 1); // otherwise need local SM sync before sending flag\n } else { // 0 bytes and 1 SM only\n atomicAdd_system(send_flagptr, 1);\n }\n\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n const int signal_id = (*recv_id) + adder;\n *recv_id = signal_id;\n volatile int *flag = (volatile int *)recv_flagptr;\n clock_t s = clock64();\n while (CHECK_IDS(*flag, signal_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\n \"pushsendrecv atomic [sending peer:%d receiving peer:%d][nvrank(GPU) sending peer:\"\n \" %d receiving peer: %d]: expecting %d, observed %d\",\n send_peer, recv_peer, nv_send, nv_recv, signal_id, *flag); /*return;*/\n if (CHECK_CE(ce_start_ptr, ce_end_ptr))\n UB_PRINT(\"pushsendrecv atomic: CE deadlock DETECTED: %d (ce_start) != %d (ce_end)\\n\",\n *ce_start_ptr, *ce_end_ptr);\n }\n }\n\n // Decrement atomic val to signal current output tile finish\n if (counters) {\n ((unsigned int *)counters)[0] = 0;\n asm volatile(\"fence.sc.gpu;\\n\");\n }\n }\n}\n\n__global__ void __launch_bounds__(MAX_THREADS) kuserbuffers_pushsendrecv_multiatomic(\n int *send_id, int *send_flagptr, int4 *srcptr, int4 *dstptr, const int lines, int send_peer,\n int recv_peer, int *recv_id, int *recv_flagptr, int adder, void *counters, int nchunks,\n int send_stride, int recv_stride, bool shuffle, uint64_t ub_timeout, int nv_send, int nv_recv) {\n for (int chunk_i = 0; chunk_i < nchunks - 1; chunk_i++) {\n int send_chunk_id = shuffle ? chunk_i : (nchunks + send_peer - chunk_i) % nchunks;\n int recv_chunk_id = shuffle ? chunk_i + 1 : (nchunks + send_peer - chunk_i - 1) % nchunks;\n int send_offset = (send_chunk_id * send_stride) / 16;\n int recv_offset = ((shuffle ? recv_chunk_id : send_chunk_id) * recv_stride) / 16;\n\n if (lines) {\n const int start_elem = threadIdx.x + blockDim.x * blockIdx.x;\n const int end_elem = lines;\n const int aligned_elem =\n ((end_elem - start_elem) & (~(blockDim.x * gridDim.x * UNROLLCOPY - 1)));\n const int end_aligned = start_elem + aligned_elem;\n if (end_elem > start_elem) {\n for (int line = start_elem; line < end_aligned;\n line += blockDim.x * gridDim.x * UNROLLCOPY) {\n int4 val[UNROLLCOPY];\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) {\n val[i] = srcptr[send_offset + line + i * blockDim.x * gridDim.x];\n }\n#pragma unroll\n for (int i = 0; i < UNROLLCOPY; i++) {\n dstptr[recv_offset + line + i * blockDim.x * gridDim.x] = val[i];\n }\n }\n for (int line = end_aligned; line < end_elem; line += blockDim.x * gridDim.x) {\n dstptr[recv_offset + line] = srcptr[send_offset + line];\n }\n }\n __syncthreads();\n if (!threadIdx.x) {\n __threadfence_system();\n atomicAdd_system(send_flagptr,\n 1); // otherwise need local SM sync before sending flag\n }\n } else { // 0 bytes and 1 SM only\n atomicAdd_system(send_flagptr, 1);\n }\n\n // wait for message to arrive.\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n const int signal_id = (*recv_id) + adder;\n *recv_id = signal_id;\n volatile int *flag = (volatile int *)recv_flagptr;\n clock_t s = clock64();\n while (CHECK_IDS(*flag, signal_id)) {\n if (CHECK_TIMEOUT(s, ub_timeout)) {\n UB_PRINT(\n \"pushsendrecv multiatomic [sending peer:%d receiving peer:%d][nvrank(GPU)\"\n \" sending peer: %d receiving peer: %d]: expecting %d, observed %d\",\n send_peer, recv_peer, nv_send, nv_recv, signal_id, *flag); /*return;*/\n // CE mode is not supported for multi-atomic, so there is no need to check for a deadlock\n return;\n }\n }\n }\n\n // Producer must update counters.\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n // Decrement atomic val to signal current output tile finish\n if (counters) {\n ((unsigned int *)counters)[recv_chunk_id /*chunk_i+1*/] = 0;\n asm volatile(\"fence.sc.gpu;\\n\");\n }\n }\n\n // sync all CTAs before moving to next chunk.\n if (threadIdx.x == 0) {\n atomicInc(((unsigned int *)counters) + nchunks + chunk_i, gridDim.x - 1);\n while (0 != (atomicCAS(((unsigned int *)counters) + nchunks + chunk_i, 0, 0))) {\n }\n }\n __syncthreads();\n }\n}\n\n// Return TRUE if two ranks share the same NV domain\n#define INTRANODE(peer) ((peer / comm->nvsize) == (comm->myrank / comm->nvsize))\n\n// Index corresponds to the type of flag:\n// 0 - Send index counter\n// 1 - CE start index counter\n// 2 - CE end index counter\n#define GET_SEND_PTR_BY_INDEX(peerlocal, comm, dsth, index) \\\n ((reinterpret_cast((comm)->peer_ptr[0][(peerlocal)])) + \\\n ((NVTE_REG0_OFFSET(comm) + NVTE_REG0_RECV + (comm)->myrank * NVTE_MAX_REGIONS + (dsth) + \\\n (index) * NVTE_MAX_NVLINK * NVTE_MAX_REGIONS) * \\\n sizeof(int)))\n\n// Index corresponds to the type of flag:\n// 0 - Receive index counter\n// 1 - CE start index counter\n// 2 - CE end index counter\n#define GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsth, index) \\\n ((reinterpret_cast((comm)->mem_ptr[0])) + \\\n ((NVTE_REG0_OFFSET(comm) + NVTE_REG0_RECV + (recv_peer) * NVTE_MAX_REGIONS + (dsth) + \\\n (index) * NVTE_MAX_NVLINK * NVTE_MAX_REGIONS) * \\\n sizeof(int)))\n\nstatic __global__ void tiny_delay_kern() {\n // Although a loop could be used to add a larger delay here, for\n // the purpose of enforcing proper kernel ordering when using CG,\n // an empty kernel seems to work well enough.\n}\n\nvoid userbuffers_tiny_delay(hipStream_t stream) { tiny_delay_kern<<<1, 1, 0, stream>>>(); }\n\nvoid userbuffers_send(const int srchandler, const size_t srcoffset, const int dsthandler,\n const size_t dstoffset, const size_t bytes, communicator *comm,\n const int peer, hipStream_t stream) {\n int peerlocal = peer % comm->nvsize;\n void *flagptr = GET_SEND_PTR_BY_INDEX(peerlocal, comm, dsthandler, 0);\n // void *ce_send_start_ptr = GET_SEND_PTR_BY_INDEX(peerlocal, comm, dsthandler, 1);\n // void *ce_send_end_ptr = GET_SEND_PTR_BY_INDEX(peerlocal, comm, dsthandler, 2);\n bool signalonly = (bytes / 16 == 0) || (comm->use_ce != 0);\n\n assert(INTRANODE(peer));\n\n if (!(comm->launch_mode & NVTE_LAUNCH_GPU)) return;\n if (comm->push == 0) {\n kuserbuffers_pullsend<<<1, 1, 0, stream>>>(comm->myrank, peer, &(comm->send_id[peer]),\n reinterpret_cast(flagptr));\n NVTE_CHECK_CUDA(hipGetLastError());\n } else {\n void *srcptr = reinterpret_cast(comm->mem_ptr[srchandler]) + srcoffset;\n void *dstptr = reinterpret_cast(comm->peer_ptr[dsthandler][peerlocal]) + dstoffset;\n\n if (comm->use_ce) {\n // kuserbuffers_inc<<<1, 1, 0, stream>>>(reinterpret_cast(ce_send_start_ptr));\n NVTE_CHECK_CUDA(hipMemcpyAsync(dstptr, srcptr, bytes, hipMemcpyDeviceToDevice, stream));\n // kuserbuffers_inc<<<1, 1, 0, stream>>>(reinterpret_cast(ce_send_end_ptr));\n }\n SETUP_LAUNCH_CONFIG(signalonly ? 1 : comm->sms, signalonly ? 1 : 1024, stream);\n int *arg1 = &comm->send_id[peer], *arg2 = reinterpret_cast(flagptr);\n int4 *arg3 = reinterpret_cast(srcptr), *arg4 = reinterpret_cast(dstptr);\n int arg5 = signalonly ? 0 : bytes / 16;\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2),\n reinterpret_cast(&arg3), reinterpret_cast(&arg4),\n reinterpret_cast(&arg5)};\n NVTE_CHECK_CUDA(\n hipLaunchKernelExC(&cfg, reinterpret_cast(kuserbuffers_pushsend), kernelArgs));\n }\n}\n\nvoid userbuffers_sendrecv(const int srchandler, const int dsthandler, const size_t send_offset,\n const size_t recv_offset, const size_t bytes, communicator *comm,\n const int send_peer, const int recv_peer, hipStream_t stream) {\n bool signalonly = (bytes / 16 == 0) || (comm->use_ce != 0);\n int send_peerlocal = send_peer % comm->nvsize;\n int recv_peerlocal = recv_peer % comm->nvsize;\n void *flagptr_send = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 0);\n // void *ce_send_start_ptr = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 1);\n // void *ce_send_end_ptr = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 2);\n void *flagptr_recv = GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 0);\n\n void *send_srcptr = reinterpret_cast(comm->mem_ptr[srchandler]) + send_offset;\n void *send_dstptr =\n reinterpret_cast(comm->peer_ptr[dsthandler][send_peerlocal]) + send_offset;\n\n if (comm->use_ce) {\n // kuserbuffers_inc<<<1, 1, 0, stream>>>(reinterpret_cast(ce_send_start_ptr));\n NVTE_CHECK_CUDA(\n hipMemcpyAsync(send_dstptr, send_srcptr, bytes, hipMemcpyDeviceToDevice, stream));\n // kuserbuffers_inc<<<1, 1, 0, stream>>>(reinterpret_cast(ce_send_end_ptr));\n }\n SETUP_LAUNCH_CONFIG(signalonly ? 1 : comm->sms, signalonly ? 1 : 1024, stream);\n\n int *arg1 = &comm->send_id[send_peer];\n int *arg2 = reinterpret_cast(flagptr_send);\n int4 *arg3 = reinterpret_cast(send_srcptr);\n int4 *arg4 = reinterpret_cast(send_dstptr);\n int arg5 = signalonly ? 0 : bytes / 16;\n int arg6 = send_peer;\n int arg7 = recv_peer;\n int *arg8 = &comm->recv_id[recv_peer * NVTE_MAX_REGIONS + dsthandler];\n int *arg9 = reinterpret_cast(flagptr_recv);\n int arg10 = signalonly ? 1 : comm->sms;\n uint64_t arg11 = comm->ub_timeout;\n int arg12 = send_peerlocal;\n int arg13 = recv_peerlocal;\n int *arg14 = reinterpret_cast(0 ? // temporary disable\n GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 1)\n : nullptr);\n int *arg15 = reinterpret_cast(0 ? // temporary disable\n GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 2)\n : nullptr);\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2),\n reinterpret_cast(&arg3), reinterpret_cast(&arg4),\n reinterpret_cast(&arg5), reinterpret_cast(&arg6),\n reinterpret_cast(&arg7), reinterpret_cast(&arg8),\n reinterpret_cast(&arg9), reinterpret_cast(&arg10),\n reinterpret_cast(&arg11), reinterpret_cast(&arg12),\n reinterpret_cast(&arg13), reinterpret_cast(&arg14),\n reinterpret_cast(&arg15)};\n NVTE_CHECK_CUDA(\n hipLaunchKernelExC(&cfg, reinterpret_cast(kuserbuffers_pushsendrecv), kernelArgs));\n}\n\nvoid userbuffers_sendrecv_atomic(const int srchandler, const int dsthandler,\n const size_t send_offset, const size_t recv_offset,\n const size_t bytes, communicator *comm, const int send_peer,\n const int recv_peer, void *counters, hipStream_t stream) {\n assert(comm->push && comm->use_ce == 0);\n bool signalonly = (bytes / 16 == 0) || (comm->use_ce != 0);\n\n int send_peerlocal = send_peer % comm->nvsize;\n int recv_peerlocal = recv_peer % comm->nvsize;\n void *flagptr_send = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 0);\n // void *ce_send_start_ptr = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 1);\n // void *ce_send_end_ptr = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 2);\n void *flagptr_recv = GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 0);\n\n void *send_srcptr = reinterpret_cast(comm->mem_ptr[srchandler]) + send_offset;\n void *send_dstptr =\n reinterpret_cast(comm->peer_ptr[dsthandler][send_peerlocal]) + send_offset;\n if (comm->use_ce) {\n // kuserbuffers_inc<<<1, 1, 0, stream>>>(reinterpret_cast(ce_send_start_ptr));\n NVTE_CHECK_CUDA(\n hipMemcpyAsync(send_dstptr, send_srcptr, bytes, hipMemcpyDeviceToDevice, stream));\n // kuserbuffers_inc<<<1, 1, 0, stream>>>(reinterpret_cast(ce_send_end_ptr));\n }\n SETUP_LAUNCH_CONFIG(signalonly ? 1 : comm->sms, signalonly ? 1 : 1024, stream);\n\n int *arg1 = &comm->send_id[send_peer];\n int *arg2 = reinterpret_cast(flagptr_send);\n int4 *arg3 = reinterpret_cast(send_srcptr);\n int4 *arg4 = reinterpret_cast(send_dstptr);\n int arg5 = signalonly ? 0 : bytes / 16;\n int arg6 = send_peer;\n int arg7 = recv_peer;\n int *arg8 = &comm->recv_id[recv_peer * NVTE_MAX_REGIONS + dsthandler];\n int *arg9 = reinterpret_cast(flagptr_recv);\n int arg10 = signalonly ? 1 : comm->sms;\n void *arg11 = counters;\n int arg12 = comm->ub_timeout;\n int arg13 = send_peerlocal;\n int arg14 = recv_peerlocal;\n int *arg15 = reinterpret_cast(0 ? // temporary disable\n GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 1)\n : nullptr);\n int *arg16 = reinterpret_cast(0 ? // temporary disable\n GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 2)\n : nullptr);\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2),\n reinterpret_cast(&arg3), reinterpret_cast(&arg4),\n reinterpret_cast(&arg5), reinterpret_cast(&arg6),\n reinterpret_cast(&arg7), reinterpret_cast(&arg8),\n reinterpret_cast(&arg9), reinterpret_cast(&arg10),\n reinterpret_cast(&arg11), reinterpret_cast(&arg12),\n reinterpret_cast(&arg13), reinterpret_cast(&arg14),\n reinterpret_cast(&arg15), reinterpret_cast(&arg16)};\n NVTE_CHECK_CUDA(hipLaunchKernelExC(\n &cfg, reinterpret_cast(kuserbuffers_pushsendrecv_atomic), kernelArgs));\n}\n\nvoid userbuffers_sendrecv_multiatomic(const int srchandler, const int dsthandler,\n const size_t send_stride, const size_t recv_stride,\n const size_t bytes, communicator *comm, const int send_peer,\n const int recv_peer, const int nchunks, void *counters,\n bool shuffle, hipStream_t stream) {\n assert(comm->push && comm->use_ce == 0);\n // CE is not supported\n\n int send_peerlocal = send_peer % comm->nvsize;\n int recv_peerlocal = recv_peer % comm->nvsize;\n void *flagptr_send = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 0);\n void *flagptr_recv = GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 0);\n\n SETUP_LAUNCH_CONFIG(comm->sms, 1024, stream);\n\n int *arg1 = &comm->send_id[send_peer];\n int *arg2 = reinterpret_cast(flagptr_send);\n int4 *arg3 = reinterpret_cast((comm->mem_ptr[srchandler]));\n int4 *arg4 = reinterpret_cast((comm->peer_ptr[dsthandler][send_peerlocal]));\n int arg5 = bytes / 16;\n int arg6 = comm->myrank;\n int arg7 = recv_peer;\n int *arg8 = &comm->recv_id[recv_peer * NVTE_MAX_REGIONS + dsthandler];\n int *arg9 = reinterpret_cast(flagptr_recv);\n int arg10 = comm->sms;\n void *arg11 = counters;\n int arg12 = nchunks;\n int arg13 = send_stride;\n int arg14 = recv_stride;\n bool arg15 = shuffle;\n uint64_t arg16 = comm->ub_timeout;\n int arg17 = send_peerlocal;\n int arg18 = recv_peerlocal;\n void *kernelArgs[] = {reinterpret_cast(&arg1), reinterpret_cast(&arg2),\n reinterpret_cast(&arg3), reinterpret_cast(&arg4),\n reinterpret_cast(&arg5), reinterpret_cast(&arg6),\n reinterpret_cast(&arg7), reinterpret_cast(&arg8),\n reinterpret_cast(&arg9), reinterpret_cast(&arg10),\n reinterpret_cast(&arg11), reinterpret_cast(&arg12),\n reinterpret_cast(&arg13), reinterpret_cast(&arg14),\n reinterpret_cast(&arg15), reinterpret_cast(&arg16),\n reinterpret_cast(&arg17), reinterpret_cast(&arg18)};\n NVTE_CHECK_CUDA(hipLaunchKernelExC(\n &cfg, reinterpret_cast(kuserbuffers_pushsendrecv_multiatomic), kernelArgs));\n}\n\nvoid userbuffers_recv(const int srchandler, const size_t srcoffset, const int dsthandler,\n const size_t dstoffset, const size_t bytes, communicator *comm,\n const int peer, hipStream_t stream) {\n int peerlocal = peer % comm->nvsize;\n void *flagptr = GET_RECV_PTR_BY_INDEX(peer, comm, dsthandler, 0);\n bool signalonly = (bytes / 16 == 0) || (comm->use_ce != 0);\n\n assert(INTRANODE(peer));\n\n if (!(comm->launch_mode & NVTE_LAUNCH_GPU)) return;\n if (comm->push == 0) {\n void *dstptr = reinterpret_cast(comm->mem_ptr[dsthandler]) + dstoffset;\n void *srcptr = reinterpret_cast(comm->peer_ptr[srchandler][peerlocal]) + srcoffset;\n\n kuserbuffers_pullrecv<<sms, signalonly ? 1 : 1024, 0, stream>>>(\n comm->myrank, peer, comm->nvrank, peerlocal,\n &(comm->recv_id[peer * NVTE_MAX_REGIONS + dsthandler]), reinterpret_cast(flagptr),\n reinterpret_cast(srcptr), reinterpret_cast(dstptr),\n signalonly ? 0 : bytes / 16, comm->ub_timeout);\n NVTE_CHECK_CUDA(hipGetLastError());\n if (!signalonly) {\n kuserbuffers_inc<<<1, 1, 0, stream>>>(&(comm->recv_id[peer * NVTE_MAX_REGIONS + dsthandler]));\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n if (comm->use_ce) {\n NVTE_CHECK_CUDA(hipMemcpyAsync(dstptr, srcptr, bytes, hipMemcpyDeviceToDevice, stream));\n }\n } else {\n kuserbuffers_pushrecv<<<1, 1, 0, stream>>>(\n comm->myrank, peer, comm->nvrank, peerlocal,\n &comm->recv_id[peer * NVTE_MAX_REGIONS + dsthandler], reinterpret_cast(flagptr),\n signalonly || comm->sms, comm->ub_timeout,\n reinterpret_cast(0 ? // temporary disable\n GET_RECV_PTR_BY_INDEX(peer, comm, dsthandler, 1)\n : nullptr),\n reinterpret_cast(0 ? // temporary disable\n GET_RECV_PTR_BY_INDEX(peer, comm, dsthandler, 2)\n : nullptr));\n NVTE_CHECK_CUDA(hipGetLastError());\n }\n}\n\nvoid userbuffers_send_all(const int srchandler, const size_t srcoffset, const int dsthandler,\n const size_t dstoffset, const size_t bytes_per_slice, int tp_rank,\n int tp_size, int world_rank, communicator *comm, hipStream_t stream) {\n int rank_round_tp = (world_rank / tp_size) * tp_size;\n for (int j = 1; j < tp_size; j++) {\n int i = (tp_rank + j) % tp_size;\n int send_offset = srcoffset + bytes_per_slice * tp_rank;\n int recv_offset = dstoffset + bytes_per_slice * tp_rank;\n userbuffers_send(srchandler, send_offset, dsthandler, recv_offset, bytes_per_slice, comm,\n rank_round_tp + i, stream);\n }\n}\n\nvoid userbuffers_recv_all(const int srchandler, const size_t srcoffset, const int dsthandler,\n const size_t dstoffset, const size_t bytes_per_slice, int tp_rank,\n int tp_size, int world_rank, communicator *comm, hipStream_t stream) {\n int rank_round_tp = (world_rank / tp_size) * tp_size;\n for (int j = tp_size - 1; j > 0; j--) {\n int i = (tp_rank + j) % tp_size;\n int send_offset = srcoffset + bytes_per_slice * i;\n int recv_offset = dstoffset + bytes_per_slice * i;\n userbuffers_recv(srchandler, send_offset, dsthandler, recv_offset, bytes_per_slice, comm,\n rank_round_tp + i, stream);\n }\n}\n\n// producer\nstatic __global__ void producer_kernel(void *atomic_ptr, int chunk_i) {\n // Decrement atomic val to signal current output tile finish\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n ((unsigned int *)atomic_ptr)[chunk_i] = 0;\n }\n\n // COMM kernel need to explicitely flash gmem.\n // GEMM kernel already executed, and can not see gmem\n // change without COMM kernel explicitely make change\n asm volatile(\"fence.sc.gpu;\\n\");\n}\n\n// consumer\nstatic __global__ void consumer_kernel(void *atomic_ptr, int chunk_i) {\n // Wait for producer to change the val to 0, which signal producer ready\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n while (0 != (atomicCAS((unsigned int *)atomic_ptr + chunk_i, 0, 0))) {\n }\n ((unsigned int *)atomic_ptr)[chunk_i] = 1;\n asm volatile(\"fence.sc.gpu;\\n\");\n }\n}\n\n// consumer_batch\nstatic __global__ void consumer_batch_kernel(void *atomic_ptr, int first_chunk_i, int num_chunks) {\n // Wait for producer to change the val to 0, which signal producer ready\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n for (int i = first_chunk_i; i < num_chunks; i++) {\n while (0 != (atomicCAS((unsigned int *)atomic_ptr + i, 0, 0))) {\n }\n ((unsigned int *)atomic_ptr)[i] = 1;\n asm volatile(\"fence.sc.gpu;\\n\");\n }\n }\n}\n\n// reset counters kernel\nstatic __global__ void reset_counters_kernel(void *atomic_ptr, int num_chunks, bool allgather) {\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n#pragma unroll\n for (int i = 0; i < num_chunks; i++) {\n ((unsigned int *)atomic_ptr)[i] = 1;\n ((unsigned int *)atomic_ptr)[i + num_chunks] = 0;\n }\n if (allgather) ((unsigned int *)atomic_ptr)[0] = 0;\n }\n}\n\nvoid producer(void *atomic_ptr, int chunk_i, hipStream_t stream) {\n dim3 block(1);\n dim3 grid(1);\n producer_kernel<<>>(atomic_ptr, chunk_i);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid consumer(void *atomic_ptr, int chunk_i, hipStream_t stream) {\n dim3 block(1);\n dim3 grid(1);\n consumer_kernel<<>>(atomic_ptr, chunk_i);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid consumer_batch(void *atomic_ptr, int first_chunk_i, int num_chunks, hipStream_t stream) {\n dim3 block(1);\n dim3 grid(1);\n consumer_batch_kernel<<>>(atomic_ptr, first_chunk_i, num_chunks);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid reset_counters(void *atomic_ptr, int num_chunks, bool allgather, hipStream_t stream) {\n dim3 block(1);\n dim3 grid(1);\n reset_counters_kernel<<>>(atomic_ptr, num_chunks, allgather);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS / 4)\n reduce_fp8_in_bf16_out_cuda(void *inputs, void *output, const float *scale,\n const int num_inputs, const int input_size,\n const int num_aligned_elements_per_input,\n const int tot_input_size) {\n fp8type *inputs_fp8 = reinterpret_cast(inputs);\n half_dtype *output_half = reinterpret_cast(output);\n\n transformer_engine::VectorizedLoader loader(inputs_fp8, tot_input_size);\n transformer_engine::VectorizedStorer storer(output_half, input_size);\n\n const size_t tid = threadIdx.x + blockDim.x * blockIdx.x;\n if (tid >= num_aligned_elements_per_input) {\n return;\n }\n float accum_buf[nvec]; // NOLINT(*)\n\n loader.load(tid, tot_input_size);\n#pragma unroll\n for (int i = 0; i < nvec; ++i) {\n accum_buf[i] = static_cast(loader.separate()[i]) * (*scale);\n }\n for (int input_id = 1; input_id < num_inputs; ++input_id) {\n loader.load(tid + num_aligned_elements_per_input * input_id, tot_input_size);\n#pragma unroll\n for (int i = 0; i < nvec; ++i) {\n accum_buf[i] += static_cast(loader.separate()[i]) * (*scale);\n }\n }\n#pragma unroll\n for (int i = 0; i < nvec; ++i) {\n storer.separate()[i] = static_cast(accum_buf[i]);\n }\n storer.store(tid, input_size);\n}\n\ntemplate \nvoid reduce_fp8_in_bf16_out(void *inputs, void *output, float *scale, int num_inputs,\n int input_size, hipStream_t stream) {\n constexpr int nvec = 32;\n assert(input_size % nvec == 0);\n const int num_aligned_elements_per_input = input_size / nvec;\n const int tot_input_size = input_size * num_inputs;\n size_t num_threads = MAX_THREADS / 4;\n size_t num_blocks = (num_aligned_elements_per_input + num_threads - 1) / num_threads;\n dim3 block(num_threads);\n dim3 grid(num_blocks);\n reduce_fp8_in_bf16_out_cuda\n <<>>(inputs, output, scale, num_inputs, input_size,\n num_aligned_elements_per_input, tot_input_size);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\ntemplate void reduce_fp8_in_bf16_out<__hip_fp8_e4m3_fnuz>(void *inputs, void *output, float *scale,\n int num_inputs, int input_size,\n hipStream_t stream);\ntemplate void reduce_fp8_in_bf16_out<__hip_fp8_e5m2_fnuz>(void *inputs, void *output, float *scale,\n int num_inputs, int input_size,\n hipStream_t stream);\n\ntemplate \n__global__ void __launch_bounds__(MAX_THREADS / 4)\n reduce_bf16_cuda(void *inputs, void *output, const int num_inputs, const int input_size,\n const int num_aligned_elements_per_input, const int tot_input_size) {\n half_dtype *inputs_half = reinterpret_cast(inputs);\n half_dtype *output_half = reinterpret_cast(output);\n\n transformer_engine::VectorizedLoader loader(inputs_half, tot_input_size);\n transformer_engine::VectorizedStorer storer(output_half, input_size);\n\n const size_t tid = threadIdx.x + blockDim.x * blockIdx.x;\n if (tid >= num_aligned_elements_per_input) {\n return;\n }\n float accum_buf[nvec]; // NOLINT(*)\n\n loader.load(tid, tot_input_size);\n#pragma unroll\n for (int i = 0; i < nvec; ++i) {\n accum_buf[i] = static_cast(loader.separate()[i]);\n }\n for (int input_id = 1; input_id < num_inputs; ++input_id) {\n loader.load(tid + num_aligned_elements_per_input * input_id, tot_input_size);\n#pragma unroll\n for (int i = 0; i < nvec; ++i) {\n accum_buf[i] += static_cast(loader.separate()[i]);\n }\n }\n#pragma unroll\n for (int i = 0; i < nvec; ++i) {\n storer.separate()[i] = static_cast(accum_buf[i]);\n }\n storer.store(tid, input_size);\n}\n\nvoid reduce_bf16(void *inputs, void *output, int num_inputs, int input_size, hipStream_t stream) {\n constexpr int nvec = 32;\n assert(input_size % nvec == 0);\n const int num_aligned_elements_per_input = input_size / nvec;\n const int tot_input_size = input_size * num_inputs;\n size_t num_threads = MAX_THREADS / 4;\n size_t num_blocks = (num_aligned_elements_per_input + num_threads - 1) / num_threads;\n dim3 block(num_threads);\n dim3 grid(num_blocks);\n reduce_bf16_cuda<<>>(\n inputs, output, num_inputs, input_size, num_aligned_elements_per_input, tot_input_size);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/cuda_runtime.h\"\n#include \"../util/logging.h\"\n\nnamespace transformer_engine {\nnamespace delayed_scaling_recipe {\n\nnamespace {\n\n// amax value to use for updating scaling factor\nenum class AmaxComputeAlgo { INVALID, MOST_RECENT, MAX };\n\nconst char* dtype_name(DType dtype) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(dtype, Type,\n return TypeInfo::name;); // NOLINT(*)\n return \"\";\n}\n\n// Maximum representable value of an FP8 dtype\ninline float fp8_dtype_max(DType dtype) {\n switch (dtype) {\n case DType::kFloat8E4M3:\n return 448;\n case DType::kFloat8E5M2:\n return 57344;\n default:\n NVTE_ERROR(\"Expected FP8 dtype, but got \", dtype_name(dtype));\n }\n return 0;\n}\n\n// struct for amax parameters\nstruct AmaxParam {\n int num_scale = 0;\n float* amax_history = nullptr;\n float* scale = nullptr;\n};\n\n// dummy struct for kernel_bulk's other params\nstruct OtherParams {\n float* a;\n size_t b;\n AmaxComputeAlgo c;\n float d;\n};\n\n#if CUDART_VERSION >= 12010\nconstexpr size_t max_constant_memory_per_kernel = 32768;\nconstexpr size_t AMAX_PARAMS_LIMIT =\n (max_constant_memory_per_kernel - sizeof(OtherParams)) / sizeof(AmaxParam);\n#else\nconstexpr size_t max_constant_memory_per_kernel = 4096;\nconstexpr size_t AMAX_PARAMS_LIMIT =\n (max_constant_memory_per_kernel - sizeof(OtherParams)) / sizeof(AmaxParam);\n#endif\n\nstruct AmaxParams {\n AmaxParam param[AMAX_PARAMS_LIMIT];\n};\n\nnamespace amax_and_scale_update_impl {\n\n// CUDA block size\nconstexpr size_t bsize = 256;\n\n/* CUDA kernel to update amax history and FP8 scaling factors\n *\n * Block dims: bsize x 1 x 1\n *\n * Grid dims: num_scales x 1 x 1\n */\n__global__ void __launch_bounds__(bsize)\n kernel(const float* amax_history_ptr, const float* scale_ptr, float* updated_amax_history_ptr,\n float* updated_scale_ptr, size_t amax_history_length, size_t amax_history_stride,\n AmaxComputeAlgo amax_compute_algo, float scaled_max) {\n const size_t tid = threadIdx.x;\n const size_t bid = blockIdx.x;\n\n // Update amax\n float amax = 0;\n {\n // Roll amax history\n const auto* amax_history = amax_history_ptr + bid;\n auto* updated_amax_history = updated_amax_history_ptr + bid;\n const auto last_amax = amax_history[0];\n const auto& length = amax_history_length;\n const auto& stride = amax_history_stride;\n for (size_t off = 0; off < length; off += bsize) {\n const size_t i = off + tid;\n float a = 0;\n if (i < length) {\n a = (i < length - 1) ? amax_history[(i + 1) * stride] : last_amax;\n amax = fmaxf(amax, a);\n }\n __syncthreads(); // In case roll is in-place\n if (i < length) {\n updated_amax_history[i * stride] = (i > 0) ? a : 0;\n }\n }\n\n // Compute amax to use for scaling factor\n switch (amax_compute_algo) {\n case AmaxComputeAlgo::MOST_RECENT:\n amax = last_amax;\n break;\n case AmaxComputeAlgo::MAX: {\n __shared__ float shared_amax[bsize];\n shared_amax[tid] = amax;\n __syncthreads();\n#pragma unroll\n for (size_t off = bsize / 2; off > 0; off /= 2) {\n if (tid < off) {\n shared_amax[tid] = fmaxf(shared_amax[tid], shared_amax[tid + off]);\n }\n __syncthreads();\n }\n amax = shared_amax[tid];\n } break;\n default:\n amax = 0;\n }\n }\n\n // Update scale\n if (tid == 0) {\n // Update scale\n float scale;\n if (isfinite(amax) && amax > 0) {\n scale = scaled_max / amax;\n } else {\n scale = scale_ptr[bid];\n }\n // When the amax is too tiny that the scale becoming infinite in FP32,\n // we set the scale to the max value of FP32. In this case, the tensor\u2019s\n // amax won't get mapped to the FP8 max representable, but rather\n // something below that, but this is the best thing we can do.\n if (isinf(scale)) {\n scale = std::numeric_limits::max();\n }\n updated_scale_ptr[bid] = scale;\n }\n}\n\n/* CUDA kernel to bulk-update amax history and FP8 scaling factors\n *\n * Block dims: bsize x 1 x 1\n *\n * Grid dims: num_tensors x 1 x 1\n */\n__global__ void __launch_bounds__(bsize)\n kernel_bulk(float* amax_reduction_buffer, AmaxParams p, size_t amax_history_length,\n AmaxComputeAlgo amax_compute_algo, float scaled_max) {\n const size_t bid = blockIdx.x;\n const size_t tid = threadIdx.x;\n const int num_scale = p.param[bid].num_scale;\n\n int offset_in_buffer = 0;\n for (int j = 0; j < bid; j++) {\n offset_in_buffer += p.param[j].num_scale;\n }\n\n for (int count = 0; count < num_scale; count++) {\n // Update amax\n float amax = 0;\n {\n // Roll amax history\n const auto& length = amax_history_length;\n const auto& stride = p.param[bid].num_scale;\n auto* amax_history = p.param[bid].amax_history + count;\n const auto last_amax = ((amax_reduction_buffer != nullptr) &&\n (amax_reduction_buffer[offset_in_buffer + count] != 0.0f))\n ? amax_reduction_buffer[offset_in_buffer + count]\n : amax_history[0];\n if (last_amax != 0.0f) {\n for (size_t off = 0; off < length; off += bsize) {\n const size_t i = off + tid;\n float a = 0;\n if (i < length) {\n a = (i < length - 1) ? amax_history[(i + 1) * stride] : last_amax;\n amax = fmaxf(amax, a);\n }\n __syncthreads(); // Inplace roll\n if (i < length) {\n amax_history[i * stride] = (i > 0) ? a : 0;\n }\n }\n }\n\n // Compute amax to use for scaling factor\n switch (amax_compute_algo) {\n case AmaxComputeAlgo::MOST_RECENT:\n amax = last_amax;\n break;\n case AmaxComputeAlgo::MAX: {\n __shared__ float shared_amax[bsize];\n shared_amax[tid] = amax;\n __syncthreads();\n#pragma unroll\n for (size_t off = bsize / 2; off > 0; off /= 2) {\n if (tid < off) {\n shared_amax[tid] = fmaxf(shared_amax[tid], shared_amax[tid + off]);\n }\n __syncthreads();\n }\n amax = shared_amax[tid];\n } break;\n default:\n amax = 0;\n }\n }\n\n // Update scale\n if (tid == 0) {\n // Computing the scaling factor requires consideration of the following scenarios:\n // 1. amax == 0:\n // No action is possible, set scale to the previous scale (or 1).\n // 2. 0 < amax < tiny_amax\n // The amax is too tiny that the scale becomes infinite in FP32.\n // Set scale = FP32_max\n // 3. tiny_amax <= amax < FP32_max:\n // Set scale = FP8_max (or scaled_max) / amax\n // 4. When amax == inf or amax == nan:\n // No action is possible, set scale to the previous scale (or 1).\n\n float scale;\n if (isfinite(amax) && amax > 0) {\n scale = scaled_max / amax;\n } else {\n scale = p.param[bid].scale[count];\n }\n // When the amax is too tiny that the scale becoming infinite in FP32,\n // we set the scale to the max value of FP32. In this case, the tensor\u2019s\n // amax won't get mapped to the FP8 max representable, but rather\n // something below that, but this is the best thing we can do.\n if (isinf(scale)) {\n scale = std::numeric_limits::max();\n }\n p.param[bid].scale[count] = scale;\n }\n }\n}\n\n} // namespace amax_and_scale_update_impl\n\n} // namespace\n\nvoid amax_and_scale_update(const Tensor& amax_history, const Tensor& scale,\n Tensor* updated_amax_history_, Tensor* updated_scale_,\n const std::string& amax_compute_algo, DType fp8_dtype, float margin,\n cudaStream_t stream) {\n auto& updated_amax_history = *updated_amax_history_;\n auto& updated_scale = *updated_scale_;\n\n // Check tensors\n NVTE_CHECK(amax_history.data.shape.size() == 2, \"Found \", amax_history.data.shape.size(),\n \" dims\");\n const size_t amax_history_length = amax_history.data.shape[0];\n const size_t num_scales = amax_history.data.shape[1];\n NVTE_CHECK(amax_history.data.dtype == DType::kFloat32, \"Found \",\n dtype_name(amax_history.data.dtype), \".\");\n NVTE_CHECK(scale.numel() == num_scales, \"Expected \", num_scales, \" elements, \", \"but found \",\n scale.numel(), \".\");\n NVTE_CHECK(scale.data.dtype == DType::kFloat32, \"Found \", dtype_name(scale.data.dtype), \".\");\n NVTE_CHECK(updated_amax_history.data.shape.size() == 2, \"Found \",\n updated_amax_history.data.shape.size(), \" dims.\");\n NVTE_CHECK(updated_amax_history.data.shape[0] == amax_history_length, \"Expected \",\n amax_history_length, \", \", \"but found \", updated_amax_history.data.shape[0]);\n NVTE_CHECK(updated_amax_history.data.shape[1] == num_scales, \"Expected \", num_scales, \", \",\n \"but found \", updated_amax_history.data.shape[1]);\n NVTE_CHECK(updated_amax_history.data.dtype == DType::kFloat32, \"Got \",\n dtype_name(updated_amax_history.data.dtype), \".\");\n NVTE_CHECK(updated_scale.numel() == num_scales, \"Expected \", num_scales, \" elements, \",\n \"but found \", updated_scale.numel(), \".\");\n NVTE_CHECK(updated_scale.data.dtype == DType::kFloat32, \"Got \",\n dtype_name(updated_scale.data.dtype), \".\");\n\n // amax value to use for updating scaling factor\n AmaxComputeAlgo amax_compute_algo_ = AmaxComputeAlgo::INVALID;\n if (amax_compute_algo == \"max\") {\n amax_compute_algo_ = AmaxComputeAlgo::MAX;\n } else if (amax_compute_algo == \"most_recent\") {\n amax_compute_algo_ = AmaxComputeAlgo::MOST_RECENT;\n } else {\n NVTE_ERROR(\"Unsupported amax compute algorithm (\", amax_compute_algo, \")\");\n }\n\n // Expected maximum value after scale is applied\n const float scaled_max = fp8_dtype_max(fp8_dtype) * std::pow(2.f, -margin);\n\n // Launch CUDA kernel\n constexpr size_t block_size = amax_and_scale_update_impl::bsize;\n const size_t grid_size = num_scales;\n amax_and_scale_update_impl::kernel<<>>(\n static_cast(amax_history.data.dptr), static_cast(scale.data.dptr),\n static_cast(updated_amax_history.data.dptr),\n static_cast(updated_scale.data.dptr), amax_history_length, num_scales,\n amax_compute_algo_, scaled_max);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid amax_and_scale_update_after_reduction(const Tensor& amax_reduction_buffer,\n std::vector amax_histories,\n std::vector scales,\n const std::string& amax_compute_algo, DType fp8_dtype,\n float margin, cudaStream_t stream) {\n using namespace transformer_engine;\n\n // amax value to use for updating scaling factor\n AmaxComputeAlgo amax_compute_algo_ = AmaxComputeAlgo::INVALID;\n if (amax_compute_algo == \"max\") {\n amax_compute_algo_ = AmaxComputeAlgo::MAX;\n } else if (amax_compute_algo == \"most_recent\") {\n amax_compute_algo_ = AmaxComputeAlgo::MOST_RECENT;\n } else {\n NVTE_ERROR(\"Unsupported amax compute algorithm (\", amax_compute_algo, \")\");\n }\n\n // Expected maximum value after scale is applied\n const float scaled_max = fp8_dtype_max(fp8_dtype) * std::pow(2.f, -margin);\n\n // Number of tensors in the bulk\n const size_t num_tensors = amax_histories.size();\n size_t num_remaining_tensors = num_tensors;\n const int num_kernels = (num_tensors + AMAX_PARAMS_LIMIT - 1) / AMAX_PARAMS_LIMIT;\n size_t amax_history_length = 0;\n if (num_tensors > 0) {\n amax_history_length = amax_histories[0]->data.shape[0];\n }\n\n // amax parameters\n float* amax_buffer = static_cast(amax_reduction_buffer.data.dptr);\n AmaxParams p;\n for (int iter = 0; iter < num_kernels; iter++) {\n size_t kernel_num_scales = 0;\n size_t kernel_num_tensors =\n (iter == (num_kernels - 1)) ? num_remaining_tensors : AMAX_PARAMS_LIMIT;\n for (size_t pi = 0; pi < kernel_num_tensors; pi++) {\n size_t i = iter * AMAX_PARAMS_LIMIT + pi;\n\n // Check tensors\n int num_scale = amax_histories[i]->data.shape[1];\n NVTE_CHECK(amax_histories[i]->data.dtype == DType::kFloat32, \"Found \",\n dtype_name(amax_histories[i]->data.dtype), \".\");\n NVTE_CHECK(amax_histories[i]->data.shape.size() == 2, \"Found \",\n amax_histories[i]->data.shape.size(), \" dims\");\n NVTE_CHECK(amax_histories[i]->numel() == amax_history_length * num_scale, \"Expected \",\n amax_history_length * num_scale, \" elements, \", \"but found \",\n amax_histories[i]->numel(), \".\");\n NVTE_CHECK(scales[i]->data.dtype == DType::kFloat32, \"Found \",\n dtype_name(scales[i]->data.dtype), \".\");\n NVTE_CHECK(scales[i]->data.shape.size() == 1, \"Found \", scales[i]->data.shape.size(),\n \" dims\");\n NVTE_CHECK(scales[i]->numel() == num_scale, \"Expected \", num_scale, \" elements, \", \"Found \",\n scales[i]->numel(), \".\");\n\n // amax parameters\n kernel_num_scales += num_scale;\n p.param[pi].num_scale = num_scale;\n p.param[pi].amax_history = static_cast(amax_histories[i]->data.dptr);\n p.param[pi].scale = static_cast(scales[i]->data.dptr);\n }\n\n // Launch CUDA kernel\n size_t grid_size = kernel_num_tensors;\n const size_t block_size = amax_and_scale_update_impl::bsize;\n amax_and_scale_update_impl::kernel_bulk<<>>(\n amax_buffer, p, amax_history_length, amax_compute_algo_, scaled_max);\n NVTE_CHECK_CUDA(cudaGetLastError());\n\n // shift amax buffer pointer\n if (amax_buffer != nullptr) {\n amax_buffer += kernel_num_scales;\n }\n num_remaining_tensors -= AMAX_PARAMS_LIMIT;\n }\n}\n\n} // namespace delayed_scaling_recipe\n} // namespace transformer_engine\n\nvoid nvte_delayed_scaling_recipe_amax_and_scale_update(\n const NVTETensor amax_history, const NVTETensor scale, NVTETensor updated_amax_history,\n NVTETensor updated_scale, const char* amax_compute_algo, NVTEDType fp8_dtype, float margin,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_delayed_scaling_recipe_amax_and_scale_update);\n using namespace transformer_engine;\n delayed_scaling_recipe::amax_and_scale_update(\n *convertNVTETensorCheck(amax_history), *convertNVTETensorCheck(scale),\n convertNVTETensor(updated_amax_history), convertNVTETensor(updated_scale), amax_compute_algo,\n static_cast(fp8_dtype), margin, stream);\n}\n\nvoid nvte_delayed_scaling_recipe_amax_and_scale_update_after_reduction(\n const NVTETensor amax_reduction_buffer, std::vector amax_histories,\n std::vector scales, const char* amax_compute_algo, NVTEDType fp8_dtype,\n float margin, cudaStream_t stream) {\n NVTE_API_CALL(nvte_delayed_scaling_recipe_amax_and_scale_update_after_reduction);\n using namespace transformer_engine;\n size_t num_tensors = amax_histories.size();\n std::vector t_amax_histories, t_scales;\n for (size_t i = 0; i < num_tensors; i++) {\n t_amax_histories.push_back(convertNVTETensor(amax_histories[i]));\n t_scales.push_back(convertNVTETensor(scales[i]));\n }\n delayed_scaling_recipe::amax_and_scale_update_after_reduction(\n *convertNVTETensorCheck(amax_reduction_buffer), t_amax_histories, t_scales, amax_compute_algo,\n static_cast(fp8_dtype), margin, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/hip/hip_runtime.h\"\n#include \"../util/logging.h\"\n\nnamespace transformer_engine {\nnamespace delayed_scaling_recipe {\n\nnamespace {\n\n// amax value to use for updating scaling factor\nenum class AmaxComputeAlgo { INVALID, MOST_RECENT, MAX };\n\nconst char* dtype_name(DType dtype) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(dtype, Type,\n return TypeInfo::name;); // NOLINT(*)\n return \"\";\n}\n\n// Maximum representable value of an FP8 dtype\ninline float fp8_dtype_max(DType dtype) {\n switch (dtype) {\n case DType::kFloat8E4M3:\n return 448;\n case DType::kFloat8E5M2:\n return 57344;\n default:\n NVTE_ERROR(\"Expected FP8 dtype, but got \", dtype_name(dtype));\n }\n return 0;\n}\n\n// struct for amax parameters\nstruct AmaxParam {\n int num_scale = 0;\n float* amax_history = nullptr;\n float* scale = nullptr;\n};\n\n// dummy struct for kernel_bulk's other params\nstruct OtherParams {\n float* a;\n size_t b;\n AmaxComputeAlgo c;\n float d;\n};\n\n#if CUDART_VERSION >= 12010\nconstexpr size_t max_constant_memory_per_kernel = 32768;\nconstexpr size_t AMAX_PARAMS_LIMIT =\n (max_constant_memory_per_kernel - sizeof(OtherParams)) / sizeof(AmaxParam);\n#else\nconstexpr size_t max_constant_memory_per_kernel = 4096;\nconstexpr size_t AMAX_PARAMS_LIMIT =\n (max_constant_memory_per_kernel - sizeof(OtherParams)) / sizeof(AmaxParam);\n#endif\n\nstruct AmaxParams {\n AmaxParam param[AMAX_PARAMS_LIMIT];\n};\n\nnamespace amax_and_scale_update_impl {\n\n// CUDA block size\nconstexpr size_t bsize = 256;\n\n/* CUDA kernel to update amax history and FP8 scaling factors\n *\n * Block dims: bsize x 1 x 1\n *\n * Grid dims: num_scales x 1 x 1\n */\n__global__ void __launch_bounds__(bsize)\n kernel(const float* amax_history_ptr, const float* scale_ptr, float* updated_amax_history_ptr,\n float* updated_scale_ptr, size_t amax_history_length, size_t amax_history_stride,\n AmaxComputeAlgo amax_compute_algo, float scaled_max) {\n const size_t tid = threadIdx.x;\n const size_t bid = blockIdx.x;\n\n // Update amax\n float amax = 0;\n {\n // Roll amax history\n const auto* amax_history = amax_history_ptr + bid;\n auto* updated_amax_history = updated_amax_history_ptr + bid;\n const auto last_amax = amax_history[0];\n const auto& length = amax_history_length;\n const auto& stride = amax_history_stride;\n for (size_t off = 0; off < length; off += bsize) {\n const size_t i = off + tid;\n float a = 0;\n if (i < length) {\n a = (i < length - 1) ? amax_history[(i + 1) * stride] : last_amax;\n amax = fmaxf(amax, a);\n }\n __syncthreads(); // In case roll is in-place\n if (i < length) {\n updated_amax_history[i * stride] = (i > 0) ? a : 0;\n }\n }\n\n // Compute amax to use for scaling factor\n switch (amax_compute_algo) {\n case AmaxComputeAlgo::MOST_RECENT:\n amax = last_amax;\n break;\n case AmaxComputeAlgo::MAX: {\n __shared__ float shared_amax[bsize];\n shared_amax[tid] = amax;\n __syncthreads();\n#pragma unroll\n for (size_t off = bsize / 2; off > 0; off /= 2) {\n if (tid < off) {\n shared_amax[tid] = fmaxf(shared_amax[tid], shared_amax[tid + off]);\n }\n __syncthreads();\n }\n amax = shared_amax[tid];\n } break;\n default:\n amax = 0;\n }\n }\n\n // Update scale\n if (tid == 0) {\n // Update scale\n float scale;\n if (isfinite(amax) && amax > 0) {\n scale = scaled_max / amax;\n } else {\n scale = scale_ptr[bid];\n }\n // When the amax is too tiny that the scale becoming infinite in FP32,\n // we set the scale to the max value of FP32. In this case, the tensor\u2019s\n // amax won't get mapped to the FP8 max representable, but rather\n // something below that, but this is the best thing we can do.\n if (isinf(scale)) {\n scale = std::numeric_limits::max();\n }\n updated_scale_ptr[bid] = scale;\n }\n}\n\n/* CUDA kernel to bulk-update amax history and FP8 scaling factors\n *\n * Block dims: bsize x 1 x 1\n *\n * Grid dims: num_tensors x 1 x 1\n */\n__global__ void __launch_bounds__(bsize)\n kernel_bulk(float* amax_reduction_buffer, AmaxParams p, size_t amax_history_length,\n AmaxComputeAlgo amax_compute_algo, float scaled_max) {\n const size_t bid = blockIdx.x;\n const size_t tid = threadIdx.x;\n const int num_scale = p.param[bid].num_scale;\n\n int offset_in_buffer = 0;\n for (int j = 0; j < bid; j++) {\n offset_in_buffer += p.param[j].num_scale;\n }\n\n for (int count = 0; count < num_scale; count++) {\n // Update amax\n float amax = 0;\n {\n // Roll amax history\n const auto& length = amax_history_length;\n const auto& stride = p.param[bid].num_scale;\n auto* amax_history = p.param[bid].amax_history + count;\n const auto last_amax = ((amax_reduction_buffer != nullptr) &&\n (amax_reduction_buffer[offset_in_buffer + count] != 0.0f))\n ? amax_reduction_buffer[offset_in_buffer + count]\n : amax_history[0];\n if (last_amax != 0.0f) {\n for (size_t off = 0; off < length; off += bsize) {\n const size_t i = off + tid;\n float a = 0;\n if (i < length) {\n a = (i < length - 1) ? amax_history[(i + 1) * stride] : last_amax;\n amax = fmaxf(amax, a);\n }\n __syncthreads(); // Inplace roll\n if (i < length) {\n amax_history[i * stride] = (i > 0) ? a : 0;\n }\n }\n }\n\n // Compute amax to use for scaling factor\n switch (amax_compute_algo) {\n case AmaxComputeAlgo::MOST_RECENT:\n amax = last_amax;\n break;\n case AmaxComputeAlgo::MAX: {\n __shared__ float shared_amax[bsize];\n shared_amax[tid] = amax;\n __syncthreads();\n#pragma unroll\n for (size_t off = bsize / 2; off > 0; off /= 2) {\n if (tid < off) {\n shared_amax[tid] = fmaxf(shared_amax[tid], shared_amax[tid + off]);\n }\n __syncthreads();\n }\n amax = shared_amax[tid];\n } break;\n default:\n amax = 0;\n }\n }\n\n // Update scale\n if (tid == 0) {\n // Computing the scaling factor requires consideration of the following scenarios:\n // 1. amax == 0:\n // No action is possible, set scale to the previous scale (or 1).\n // 2. 0 < amax < tiny_amax\n // The amax is too tiny that the scale becomes infinite in FP32.\n // Set scale = FP32_max\n // 3. tiny_amax <= amax < FP32_max:\n // Set scale = FP8_max (or scaled_max) / amax\n // 4. When amax == inf or amax == nan:\n // No action is possible, set scale to the previous scale (or 1).\n\n float scale;\n if (isfinite(amax) && amax > 0) {\n scale = scaled_max / amax;\n } else {\n scale = p.param[bid].scale[count];\n }\n // When the amax is too tiny that the scale becoming infinite in FP32,\n // we set the scale to the max value of FP32. In this case, the tensor\u2019s\n // amax won't get mapped to the FP8 max representable, but rather\n // something below that, but this is the best thing we can do.\n if (isinf(scale)) {\n scale = std::numeric_limits::max();\n }\n p.param[bid].scale[count] = scale;\n }\n }\n}\n\n} // namespace amax_and_scale_update_impl\n\n} // namespace\n\nvoid amax_and_scale_update(const Tensor& amax_history, const Tensor& scale,\n Tensor* updated_amax_history_, Tensor* updated_scale_,\n const std::string& amax_compute_algo, DType fp8_dtype, float margin,\n hipStream_t stream) {\n auto& updated_amax_history = *updated_amax_history_;\n auto& updated_scale = *updated_scale_;\n\n // Check tensors\n NVTE_CHECK(amax_history.data.shape.size() == 2, \"Found \", amax_history.data.shape.size(),\n \" dims\");\n const size_t amax_history_length = amax_history.data.shape[0];\n const size_t num_scales = amax_history.data.shape[1];\n NVTE_CHECK(amax_history.data.dtype == DType::kFloat32, \"Found \",\n dtype_name(amax_history.data.dtype), \".\");\n NVTE_CHECK(scale.numel() == num_scales, \"Expected \", num_scales, \" elements, \", \"but found \",\n scale.numel(), \".\");\n NVTE_CHECK(scale.data.dtype == DType::kFloat32, \"Found \", dtype_name(scale.data.dtype), \".\");\n NVTE_CHECK(updated_amax_history.data.shape.size() == 2, \"Found \",\n updated_amax_history.data.shape.size(), \" dims.\");\n NVTE_CHECK(updated_amax_history.data.shape[0] == amax_history_length, \"Expected \",\n amax_history_length, \", \", \"but found \", updated_amax_history.data.shape[0]);\n NVTE_CHECK(updated_amax_history.data.shape[1] == num_scales, \"Expected \", num_scales, \", \",\n \"but found \", updated_amax_history.data.shape[1]);\n NVTE_CHECK(updated_amax_history.data.dtype == DType::kFloat32, \"Got \",\n dtype_name(updated_amax_history.data.dtype), \".\");\n NVTE_CHECK(updated_scale.numel() == num_scales, \"Expected \", num_scales, \" elements, \",\n \"but found \", updated_scale.numel(), \".\");\n NVTE_CHECK(updated_scale.data.dtype == DType::kFloat32, \"Got \",\n dtype_name(updated_scale.data.dtype), \".\");\n\n // amax value to use for updating scaling factor\n AmaxComputeAlgo amax_compute_algo_ = AmaxComputeAlgo::INVALID;\n if (amax_compute_algo == \"max\") {\n amax_compute_algo_ = AmaxComputeAlgo::MAX;\n } else if (amax_compute_algo == \"most_recent\") {\n amax_compute_algo_ = AmaxComputeAlgo::MOST_RECENT;\n } else {\n NVTE_ERROR(\"Unsupported amax compute algorithm (\", amax_compute_algo, \")\");\n }\n\n // Expected maximum value after scale is applied\n const float scaled_max = fp8_dtype_max(fp8_dtype) * std::pow(2.f, -margin);\n\n // Launch CUDA kernel\n constexpr size_t block_size = amax_and_scale_update_impl::bsize;\n const size_t grid_size = num_scales;\n amax_and_scale_update_impl::kernel<<>>(\n static_cast(amax_history.data.dptr), static_cast(scale.data.dptr),\n static_cast(updated_amax_history.data.dptr),\n static_cast(updated_scale.data.dptr), amax_history_length, num_scales,\n amax_compute_algo_, scaled_max);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid amax_and_scale_update_after_reduction(const Tensor& amax_reduction_buffer,\n std::vector amax_histories,\n std::vector scales,\n const std::string& amax_compute_algo, DType fp8_dtype,\n float margin, hipStream_t stream) {\n using namespace transformer_engine;\n\n // amax value to use for updating scaling factor\n AmaxComputeAlgo amax_compute_algo_ = AmaxComputeAlgo::INVALID;\n if (amax_compute_algo == \"max\") {\n amax_compute_algo_ = AmaxComputeAlgo::MAX;\n } else if (amax_compute_algo == \"most_recent\") {\n amax_compute_algo_ = AmaxComputeAlgo::MOST_RECENT;\n } else {\n NVTE_ERROR(\"Unsupported amax compute algorithm (\", amax_compute_algo, \")\");\n }\n\n // Expected maximum value after scale is applied\n const float scaled_max = fp8_dtype_max(fp8_dtype) * std::pow(2.f, -margin);\n\n // Number of tensors in the bulk\n const size_t num_tensors = amax_histories.size();\n size_t num_remaining_tensors = num_tensors;\n const int num_kernels = (num_tensors + AMAX_PARAMS_LIMIT - 1) / AMAX_PARAMS_LIMIT;\n size_t amax_history_length = 0;\n if (num_tensors > 0) {\n amax_history_length = amax_histories[0]->data.shape[0];\n }\n\n // amax parameters\n float* amax_buffer = static_cast(amax_reduction_buffer.data.dptr);\n AmaxParams p;\n for (int iter = 0; iter < num_kernels; iter++) {\n size_t kernel_num_scales = 0;\n size_t kernel_num_tensors =\n (iter == (num_kernels - 1)) ? num_remaining_tensors : AMAX_PARAMS_LIMIT;\n for (size_t pi = 0; pi < kernel_num_tensors; pi++) {\n size_t i = iter * AMAX_PARAMS_LIMIT + pi;\n\n // Check tensors\n int num_scale = amax_histories[i]->data.shape[1];\n NVTE_CHECK(amax_histories[i]->data.dtype == DType::kFloat32, \"Found \",\n dtype_name(amax_histories[i]->data.dtype), \".\");\n NVTE_CHECK(amax_histories[i]->data.shape.size() == 2, \"Found \",\n amax_histories[i]->data.shape.size(), \" dims\");\n NVTE_CHECK(amax_histories[i]->numel() == amax_history_length * num_scale, \"Expected \",\n amax_history_length * num_scale, \" elements, \", \"but found \",\n amax_histories[i]->numel(), \".\");\n NVTE_CHECK(scales[i]->data.dtype == DType::kFloat32, \"Found \",\n dtype_name(scales[i]->data.dtype), \".\");\n NVTE_CHECK(scales[i]->data.shape.size() == 1, \"Found \", scales[i]->data.shape.size(),\n \" dims\");\n NVTE_CHECK(scales[i]->numel() == num_scale, \"Expected \", num_scale, \" elements, \", \"Found \",\n scales[i]->numel(), \".\");\n\n // amax parameters\n kernel_num_scales += num_scale;\n p.param[pi].num_scale = num_scale;\n p.param[pi].amax_history = static_cast(amax_histories[i]->data.dptr);\n p.param[pi].scale = static_cast(scales[i]->data.dptr);\n }\n\n // Launch CUDA kernel\n size_t grid_size = kernel_num_tensors;\n const size_t block_size = amax_and_scale_update_impl::bsize;\n amax_and_scale_update_impl::kernel_bulk<<>>(\n amax_buffer, p, amax_history_length, amax_compute_algo_, scaled_max);\n NVTE_CHECK_CUDA(hipGetLastError());\n\n // shift amax buffer pointer\n if (amax_buffer != nullptr) {\n amax_buffer += kernel_num_scales;\n }\n num_remaining_tensors -= AMAX_PARAMS_LIMIT;\n }\n}\n\n} // namespace delayed_scaling_recipe\n} // namespace transformer_engine\n\nvoid nvte_delayed_scaling_recipe_amax_and_scale_update(\n const NVTETensor amax_history, const NVTETensor scale, NVTETensor updated_amax_history,\n NVTETensor updated_scale, const char* amax_compute_algo, NVTEDType fp8_dtype, float margin,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_delayed_scaling_recipe_amax_and_scale_update);\n using namespace transformer_engine;\n delayed_scaling_recipe::amax_and_scale_update(\n *convertNVTETensorCheck(amax_history), *convertNVTETensorCheck(scale),\n convertNVTETensor(updated_amax_history), convertNVTETensor(updated_scale), amax_compute_algo,\n static_cast(fp8_dtype), margin, stream);\n}\n\nvoid nvte_delayed_scaling_recipe_amax_and_scale_update_after_reduction(\n const NVTETensor amax_reduction_buffer, std::vector amax_histories,\n std::vector scales, const char* amax_compute_algo, NVTEDType fp8_dtype,\n float margin, hipStream_t stream) {\n NVTE_API_CALL(nvte_delayed_scaling_recipe_amax_and_scale_update_after_reduction);\n using namespace transformer_engine;\n size_t num_tensors = amax_histories.size();\n std::vector t_amax_histories, t_scales;\n for (size_t i = 0; i < num_tensors; i++) {\n t_amax_histories.push_back(convertNVTETensor(amax_histories[i]));\n t_scales.push_back(convertNVTETensor(scales[i]));\n }\n delayed_scaling_recipe::amax_and_scale_update_after_reduction(\n *convertNVTETensorCheck(amax_reduction_buffer), t_amax_histories, t_scales, amax_compute_algo,\n static_cast(fp8_dtype), margin, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/recipe/delayed_scaling.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n\n#include \"../common.h\"\n#include \"../util/ptx.cuh\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\nnamespace mxfp8_scaling_recipe {\n\nconstexpr int rowwise_row_padding = 128; // Row padding of rowwise_scale and rowwise_amax\nconstexpr int rowwise_col_padding = 4; // Column padding of rowwise_scale and rowwise_amax\nconstexpr int colwise_row_padding = 4; // Row padding of colwise_scale and colwise_amax\nconstexpr int colwise_col_padding = 128; // Column padding of colwise_scale and colwise_amax\n\nconstexpr int kRowsPerTile = 32; // Rows each block processes\nconstexpr int kColsPerTile = 128; // Columns each block processes\n\nconstexpr int kThreadsPerBlock = 128;\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock)\n mxfp8_scaling_compute_partial_amax_kernel(const IType *input, IType *amax_rowwise,\n IType *amax_colwise, int amax_rowwise_stride,\n int amax_colwise_stride, int rows, int cols,\n size_t start_offset, size_t len) {\n __shared__ float smem_amax_rowwise[kRowsPerTile][kColsPerTile / 32];\n\n size_t end_offset = start_offset + len;\n const IType *input_minus_offset = input - start_offset;\n int warp_idx = threadIdx.x / 32;\n int lane_idx = threadIdx.x % 32;\n int c = blockIdx.x * kColsPerTile + threadIdx.x;\n int r = blockIdx.y * kRowsPerTile;\n\n float col_amax = 0.0f;\n#pragma unroll\n for (int i = 0; i < kRowsPerTile; i++) {\n size_t idx = r * cols + c;\n float row_amax = 0.0f;\n\n if (r < rows && c < cols && idx >= start_offset && idx < end_offset) {\n float abs_input = fabs(static_cast(input_minus_offset[idx]));\n row_amax = fmaxf(row_amax, abs_input);\n col_amax = fmaxf(col_amax, abs_input);\n }\n\n#pragma unroll\n for (int delta = 16; delta > 0; delta /= 2) {\n float other_row_amax = __shfl_down_sync(0xFFFFFFFF, row_amax, delta);\n row_amax = fmaxf(row_amax, other_row_amax);\n }\n\n if (lane_idx == 0) {\n smem_amax_rowwise[i][warp_idx] = row_amax;\n }\n\n r++;\n }\n\n amax_colwise[blockIdx.y * amax_colwise_stride + c] = static_cast(col_amax);\n\n __syncthreads();\n\n int r_ = threadIdx.x / (kColsPerTile / 32); // rows in shared memory\n int c_ = threadIdx.x % (kColsPerTile / 32); // cols in shared memory\n r = blockIdx.y * kRowsPerTile + r_;\n c = blockIdx.x * kColsPerTile / 32 + c_;\n amax_rowwise[r * amax_rowwise_stride + c] = static_cast(smem_amax_rowwise[r_][c_]);\n}\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock)\n mxfp8_scaling_partial_cast_kernel(const IType *input, OType *output_rowwise,\n OType *output_colwise, const e8m0_t *scale_inv_rowwise,\n const e8m0_t *scale_inv_colwise, int scale_inv_rowwise_stride,\n int scale_inv_colwise_stride, int rows, int cols,\n size_t start_offset, size_t len) {\n __shared__ float smem_scales_rowwise[kRowsPerTile][kColsPerTile / 32];\n __shared__ float smem_scales_colwise[kColsPerTile];\n\n // Load scales_rowwise\n {\n int r_ = threadIdx.x / (kColsPerTile / 32); // rows in shared memory\n int c_ = threadIdx.x % (kColsPerTile / 32); // cols in shared memory\n int r = blockIdx.y * kRowsPerTile + r_;\n int c = blockIdx.x * kColsPerTile / 32 + c_;\n size_t idx = r * scale_inv_rowwise_stride + c;\n smem_scales_rowwise[r_][c_] = ptx::exp2f_rcp(scale_inv_rowwise[idx]);\n }\n\n // Load scales_colwise\n {\n int c_ = threadIdx.x;\n int r = blockIdx.y * kRowsPerTile / 32;\n int c = blockIdx.x * kColsPerTile + c_;\n size_t idx = r * scale_inv_colwise_stride + c;\n smem_scales_colwise[c_] = ptx::exp2f_rcp(scale_inv_colwise[idx]);\n }\n\n __syncthreads();\n\n size_t end_offset = start_offset + len;\n const IType *input_minus_offset = input - start_offset;\n OType *output_rowwise_minus_offset = output_rowwise - start_offset;\n OType *output_colwise_minus_offset = output_colwise - start_offset;\n int warp_idx = threadIdx.x / 32;\n // int lane_idx = threadIdx.x % 32;\n int c = blockIdx.x * kColsPerTile + threadIdx.x;\n int r = blockIdx.y * kRowsPerTile;\n\n#pragma unroll\n for (int i = 0; i < kRowsPerTile; i++) {\n size_t idx = r * cols + c;\n\n if (r < rows && c < cols && idx >= start_offset && idx < end_offset) {\n float inp = static_cast(input_minus_offset[idx]);\n OType out_rowwise = static_cast(inp * smem_scales_rowwise[i][warp_idx]);\n OType out_colwise = static_cast(inp * smem_scales_colwise[threadIdx.x]);\n output_rowwise_minus_offset[idx] = out_rowwise;\n output_colwise_minus_offset[idx] = out_colwise;\n }\n\n r++;\n }\n}\n\nvoid mxfp8_scaling_compute_partial_amax(const Tensor input, Tensor amax_rowwise,\n Tensor amax_colwise, int rows, int cols,\n size_t start_offset, cudaStream_t stream) {\n NVTE_CHECK(rows % 32 == 0, \"rows must be divisible by 32\");\n NVTE_CHECK(cols % 32 == 0, \"cols must be divisible by 32\");\n\n NVTE_CHECK(input.data.shape.size() == 1, \"input must be a 1D tensor\");\n NVTE_CHECK(start_offset + input.data.shape[0] <= static_cast(rows) * cols,\n \"Invalid start_offset\");\n\n NVTE_CHECK(amax_rowwise.data.shape.size() == 2, \"amax_rowwise must be a 2D tensor\");\n NVTE_CHECK(amax_rowwise.data.shape[0] % rowwise_row_padding == 0,\n \"Wrong padding of amax_rowwise's rows\");\n NVTE_CHECK(amax_rowwise.data.shape[0] >= rows, \"Invalid rows\");\n NVTE_CHECK(amax_rowwise.data.shape[1] % rowwise_col_padding == 0,\n \"Wrong padding of amax_rowwise's cols\");\n NVTE_CHECK(amax_rowwise.data.shape[1] >= cols / 32, \"Invalid cols\");\n NVTE_CHECK(amax_rowwise.dtype() == input.dtype(), \"Wrong dtype of amax_rowwise\");\n\n NVTE_CHECK(amax_colwise.data.shape.size() == 2, \"amax_colwise must be a 2D tensor\");\n NVTE_CHECK(amax_colwise.data.shape[0] % colwise_row_padding == 0,\n \"Wrong padding of amax_colwise's rows\");\n NVTE_CHECK(amax_colwise.data.shape[0] >= rows / 32, \"Invalid rows\");\n NVTE_CHECK(amax_colwise.data.shape[1] % colwise_col_padding == 0,\n \"Wrong padding of amax_colwise's cols\");\n NVTE_CHECK(amax_colwise.data.shape[1] >= cols, \"Invalid cols\");\n NVTE_CHECK(amax_colwise.dtype() == input.dtype(), \"Wrong dtype of amax_colwise\");\n\n int blocks_x = (cols + kColsPerTile - 1) / kColsPerTile;\n int blocks_y = (rows + kRowsPerTile - 1) / kRowsPerTile;\n dim3 grid(blocks_x, blocks_y);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n input.dtype(), IType,\n mxfp8_scaling_compute_partial_amax_kernel<<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(amax_rowwise.data.dptr),\n reinterpret_cast(amax_colwise.data.dptr), amax_rowwise.data.shape[1],\n amax_colwise.data.shape[1], rows, cols, start_offset, input.data.shape[0]);)\n}\n\nvoid mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tensor output_colwise,\n const Tensor scale_inv_rowwise, const Tensor scale_inv_colwise,\n int rows, int cols, size_t start_offset, cudaStream_t stream) {\n NVTE_CHECK(rows % 32 == 0, \"rows must be divisible by 32\");\n NVTE_CHECK(cols % 32 == 0, \"cols must be divisible by 32\");\n\n NVTE_CHECK(input.data.shape.size() == 1, \"input must be a 1D tensor\");\n NVTE_CHECK(start_offset + input.data.shape[0] <= static_cast(rows) * cols,\n \"Invalid start_offset\");\n\n NVTE_CHECK(output_rowwise.data.shape.size() == 1, \"output_rowwise must be a 1D tensor\");\n NVTE_CHECK(output_colwise.data.shape.size() == 1, \"output_colwise must be a 1D tensor\");\n NVTE_CHECK(output_rowwise.data.shape[0] == input.data.shape[0],\n \"Size of input and output_rowwise mismatch\");\n NVTE_CHECK(output_colwise.data.shape[0] == input.data.shape[0],\n \"Size of input and output_colwise mismatch\");\n\n NVTE_CHECK(output_rowwise.dtype() == DType::kFloat8E4M3 || output_rowwise.dtype() == DType::kByte,\n \"output_rowwise should be e4m3 or uint8\");\n NVTE_CHECK(output_colwise.dtype() == DType::kFloat8E4M3 || output_colwise.dtype() == DType::kByte,\n \"output_colwise should be e4m3 or uint8\");\n\n NVTE_CHECK(scale_inv_rowwise.data.shape.size() == 2, \"scale_inv_rowwise must be a 2D tensor\");\n NVTE_CHECK(scale_inv_rowwise.data.shape[0] % rowwise_row_padding == 0,\n \"Wrong padding of scale_inv_rowwise's rows\");\n NVTE_CHECK(scale_inv_rowwise.data.shape[0] >= rows, \"Invalid rows\");\n NVTE_CHECK(scale_inv_rowwise.data.shape[1] % rowwise_col_padding == 0,\n \"Wrong padding of scale_inv_rowwise's cols\");\n NVTE_CHECK(scale_inv_rowwise.data.shape[1] >= cols / 32, \"Invalid cols\");\n NVTE_CHECK(scale_inv_rowwise.dtype() == DType::kByte, \"Wrong dtype of scale_inv_rowwise\");\n\n NVTE_CHECK(scale_inv_colwise.data.shape.size() == 2, \"scale_inv_colwise must be a 2D tensor\");\n NVTE_CHECK(scale_inv_colwise.data.shape[0] % colwise_row_padding == 0,\n \"Wrong padding of scale_inv_colwise's rows\");\n NVTE_CHECK(scale_inv_colwise.data.shape[0] >= rows / 32, \"Invalid rows\");\n NVTE_CHECK(scale_inv_colwise.data.shape[1] % colwise_col_padding == 0,\n \"Wrong padding of scale_inv_colwise's cols\");\n NVTE_CHECK(scale_inv_colwise.data.shape[1] >= cols, \"Invalid cols\");\n NVTE_CHECK(scale_inv_colwise.dtype() == DType::kByte, \"Wrong dtype of scale_inv_colwise\");\n\n int blocks_x = (cols + kColsPerTile - 1) / kColsPerTile;\n int blocks_y = (rows + kRowsPerTile - 1) / kRowsPerTile;\n dim3 grid(blocks_x, blocks_y);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n input.dtype(), IType,\n mxfp8_scaling_partial_cast_kernel<<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(output_rowwise.data.dptr),\n reinterpret_cast(output_colwise.data.dptr),\n reinterpret_cast(scale_inv_rowwise.data.dptr),\n reinterpret_cast(scale_inv_colwise.data.dptr),\n scale_inv_rowwise.data.shape[1], scale_inv_colwise.data.shape[1], rows, cols,\n start_offset, input.data.shape[0]);)\n}\n\n} // namespace mxfp8_scaling_recipe\n} // namespace transformer_engine\n\nvoid nvte_mxfp8_scaling_compute_partial_amax(const NVTETensor input, NVTETensor amax_rowwise,\n NVTETensor amax_colwise, int rows, int cols,\n size_t start_offset, cudaStream_t stream) {\n NVTE_API_CALL(nvte_mxfp8_scaling_compute_partial_amax);\n using namespace transformer_engine;\n mxfp8_scaling_recipe::mxfp8_scaling_compute_partial_amax(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(amax_rowwise),\n *convertNVTETensorCheck(amax_colwise), rows, cols, start_offset, stream);\n}\n\nvoid nvte_mxfp8_scaling_partial_cast(const NVTETensor input, NVTETensor output_rowwise,\n NVTETensor output_colwise, const NVTETensor scale_inv_rowwise,\n const NVTETensor scale_inv_colwise, int rows, int cols,\n size_t start_offset, cudaStream_t stream) {\n NVTE_API_CALL(nvte_mxfp8_scaling_partial_cast);\n using namespace transformer_engine;\n mxfp8_scaling_recipe::mxfp8_scaling_partial_cast(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(output_rowwise),\n *convertNVTETensorCheck(output_colwise), *convertNVTETensorCheck(scale_inv_rowwise),\n *convertNVTETensorCheck(scale_inv_colwise), rows, cols, start_offset, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \"../common.h\"\n#include \"../util/ptx.cuh\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\nnamespace mxfp8_scaling_recipe {\n\nconstexpr int rowwise_row_padding = 128; // Row padding of rowwise_scale and rowwise_amax\nconstexpr int rowwise_col_padding = 4; // Column padding of rowwise_scale and rowwise_amax\nconstexpr int colwise_row_padding = 4; // Row padding of colwise_scale and colwise_amax\nconstexpr int colwise_col_padding = 128; // Column padding of colwise_scale and colwise_amax\n\nconstexpr int kRowsPerTile = 32; // Rows each block processes\nconstexpr int kColsPerTile = 128; // Columns each block processes\n\nconstexpr int kThreadsPerBlock = 128;\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock)\n mxfp8_scaling_compute_partial_amax_kernel(const IType *input, IType *amax_rowwise,\n IType *amax_colwise, int amax_rowwise_stride,\n int amax_colwise_stride, int rows, int cols,\n size_t start_offset, size_t len) {\n __shared__ float smem_amax_rowwise[kRowsPerTile][kColsPerTile / 32];\n\n size_t end_offset = start_offset + len;\n const IType *input_minus_offset = input - start_offset;\n int warp_idx = threadIdx.x / 32;\n int lane_idx = threadIdx.x % 32;\n int c = blockIdx.x * kColsPerTile + threadIdx.x;\n int r = blockIdx.y * kRowsPerTile;\n\n float col_amax = 0.0f;\n#pragma unroll\n for (int i = 0; i < kRowsPerTile; i++) {\n size_t idx = r * cols + c;\n float row_amax = 0.0f;\n\n if (r < rows && c < cols && idx >= start_offset && idx < end_offset) {\n float abs_input = fabs(static_cast(input_minus_offset[idx]));\n row_amax = fmaxf(row_amax, abs_input);\n col_amax = fmaxf(col_amax, abs_input);\n }\n\n#pragma unroll\n for (int delta = 16; delta > 0; delta /= 2) {\n float other_row_amax = __shfl_down_sync(0xFFFFFFFF, row_amax, delta);\n row_amax = fmaxf(row_amax, other_row_amax);\n }\n\n if (lane_idx == 0) {\n smem_amax_rowwise[i][warp_idx] = row_amax;\n }\n\n r++;\n }\n\n amax_colwise[blockIdx.y * amax_colwise_stride + c] = static_cast(col_amax);\n\n __syncthreads();\n\n int r_ = threadIdx.x / (kColsPerTile / 32); // rows in shared memory\n int c_ = threadIdx.x % (kColsPerTile / 32); // cols in shared memory\n r = blockIdx.y * kRowsPerTile + r_;\n c = blockIdx.x * kColsPerTile / 32 + c_;\n amax_rowwise[r * amax_rowwise_stride + c] = static_cast(smem_amax_rowwise[r_][c_]);\n}\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock)\n mxfp8_scaling_partial_cast_kernel(const IType *input, OType *output_rowwise,\n OType *output_colwise, const e8m0_t *scale_inv_rowwise,\n const e8m0_t *scale_inv_colwise, int scale_inv_rowwise_stride,\n int scale_inv_colwise_stride, int rows, int cols,\n size_t start_offset, size_t len) {\n __shared__ float smem_scales_rowwise[kRowsPerTile][kColsPerTile / 32];\n __shared__ float smem_scales_colwise[kColsPerTile];\n\n // Load scales_rowwise\n {\n int r_ = threadIdx.x / (kColsPerTile / 32); // rows in shared memory\n int c_ = threadIdx.x % (kColsPerTile / 32); // cols in shared memory\n int r = blockIdx.y * kRowsPerTile + r_;\n int c = blockIdx.x * kColsPerTile / 32 + c_;\n size_t idx = r * scale_inv_rowwise_stride + c;\n smem_scales_rowwise[r_][c_] = ptx::exp2f_rcp(scale_inv_rowwise[idx]);\n }\n\n // Load scales_colwise\n {\n int c_ = threadIdx.x;\n int r = blockIdx.y * kRowsPerTile / 32;\n int c = blockIdx.x * kColsPerTile + c_;\n size_t idx = r * scale_inv_colwise_stride + c;\n smem_scales_colwise[c_] = ptx::exp2f_rcp(scale_inv_colwise[idx]);\n }\n\n __syncthreads();\n\n size_t end_offset = start_offset + len;\n const IType *input_minus_offset = input - start_offset;\n OType *output_rowwise_minus_offset = output_rowwise - start_offset;\n OType *output_colwise_minus_offset = output_colwise - start_offset;\n int warp_idx = threadIdx.x / 32;\n // int lane_idx = threadIdx.x % 32;\n int c = blockIdx.x * kColsPerTile + threadIdx.x;\n int r = blockIdx.y * kRowsPerTile;\n\n#pragma unroll\n for (int i = 0; i < kRowsPerTile; i++) {\n size_t idx = r * cols + c;\n\n if (r < rows && c < cols && idx >= start_offset && idx < end_offset) {\n float inp = static_cast(input_minus_offset[idx]);\n OType out_rowwise = static_cast(inp * smem_scales_rowwise[i][warp_idx]);\n OType out_colwise = static_cast(inp * smem_scales_colwise[threadIdx.x]);\n output_rowwise_minus_offset[idx] = out_rowwise;\n output_colwise_minus_offset[idx] = out_colwise;\n }\n\n r++;\n }\n}\n\nvoid mxfp8_scaling_compute_partial_amax(const Tensor input, Tensor amax_rowwise,\n Tensor amax_colwise, int rows, int cols,\n size_t start_offset, hipStream_t stream) {\n NVTE_CHECK(rows % 32 == 0, \"rows must be divisible by 32\");\n NVTE_CHECK(cols % 32 == 0, \"cols must be divisible by 32\");\n\n NVTE_CHECK(input.data.shape.size() == 1, \"input must be a 1D tensor\");\n NVTE_CHECK(start_offset + input.data.shape[0] <= static_cast(rows) * cols,\n \"Invalid start_offset\");\n\n NVTE_CHECK(amax_rowwise.data.shape.size() == 2, \"amax_rowwise must be a 2D tensor\");\n NVTE_CHECK(amax_rowwise.data.shape[0] % rowwise_row_padding == 0,\n \"Wrong padding of amax_rowwise's rows\");\n NVTE_CHECK(amax_rowwise.data.shape[0] >= rows, \"Invalid rows\");\n NVTE_CHECK(amax_rowwise.data.shape[1] % rowwise_col_padding == 0,\n \"Wrong padding of amax_rowwise's cols\");\n NVTE_CHECK(amax_rowwise.data.shape[1] >= cols / 32, \"Invalid cols\");\n NVTE_CHECK(amax_rowwise.dtype() == input.dtype(), \"Wrong dtype of amax_rowwise\");\n\n NVTE_CHECK(amax_colwise.data.shape.size() == 2, \"amax_colwise must be a 2D tensor\");\n NVTE_CHECK(amax_colwise.data.shape[0] % colwise_row_padding == 0,\n \"Wrong padding of amax_colwise's rows\");\n NVTE_CHECK(amax_colwise.data.shape[0] >= rows / 32, \"Invalid rows\");\n NVTE_CHECK(amax_colwise.data.shape[1] % colwise_col_padding == 0,\n \"Wrong padding of amax_colwise's cols\");\n NVTE_CHECK(amax_colwise.data.shape[1] >= cols, \"Invalid cols\");\n NVTE_CHECK(amax_colwise.dtype() == input.dtype(), \"Wrong dtype of amax_colwise\");\n\n int blocks_x = (cols + kColsPerTile - 1) / kColsPerTile;\n int blocks_y = (rows + kRowsPerTile - 1) / kRowsPerTile;\n dim3 grid(blocks_x, blocks_y);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n input.dtype(), IType,\n mxfp8_scaling_compute_partial_amax_kernel<<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(amax_rowwise.data.dptr),\n reinterpret_cast(amax_colwise.data.dptr), amax_rowwise.data.shape[1],\n amax_colwise.data.shape[1], rows, cols, start_offset, input.data.shape[0]);)\n}\n\nvoid mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tensor output_colwise,\n const Tensor scale_inv_rowwise, const Tensor scale_inv_colwise,\n int rows, int cols, size_t start_offset, hipStream_t stream) {\n NVTE_CHECK(rows % 32 == 0, \"rows must be divisible by 32\");\n NVTE_CHECK(cols % 32 == 0, \"cols must be divisible by 32\");\n\n NVTE_CHECK(input.data.shape.size() == 1, \"input must be a 1D tensor\");\n NVTE_CHECK(start_offset + input.data.shape[0] <= static_cast(rows) * cols,\n \"Invalid start_offset\");\n\n NVTE_CHECK(output_rowwise.data.shape.size() == 1, \"output_rowwise must be a 1D tensor\");\n NVTE_CHECK(output_colwise.data.shape.size() == 1, \"output_colwise must be a 1D tensor\");\n NVTE_CHECK(output_rowwise.data.shape[0] == input.data.shape[0],\n \"Size of input and output_rowwise mismatch\");\n NVTE_CHECK(output_colwise.data.shape[0] == input.data.shape[0],\n \"Size of input and output_colwise mismatch\");\n\n NVTE_CHECK(output_rowwise.dtype() == DType::kFloat8E4M3 || output_rowwise.dtype() == DType::kByte,\n \"output_rowwise should be e4m3 or uint8\");\n NVTE_CHECK(output_colwise.dtype() == DType::kFloat8E4M3 || output_colwise.dtype() == DType::kByte,\n \"output_colwise should be e4m3 or uint8\");\n\n NVTE_CHECK(scale_inv_rowwise.data.shape.size() == 2, \"scale_inv_rowwise must be a 2D tensor\");\n NVTE_CHECK(scale_inv_rowwise.data.shape[0] % rowwise_row_padding == 0,\n \"Wrong padding of scale_inv_rowwise's rows\");\n NVTE_CHECK(scale_inv_rowwise.data.shape[0] >= rows, \"Invalid rows\");\n NVTE_CHECK(scale_inv_rowwise.data.shape[1] % rowwise_col_padding == 0,\n \"Wrong padding of scale_inv_rowwise's cols\");\n NVTE_CHECK(scale_inv_rowwise.data.shape[1] >= cols / 32, \"Invalid cols\");\n NVTE_CHECK(scale_inv_rowwise.dtype() == DType::kByte, \"Wrong dtype of scale_inv_rowwise\");\n\n NVTE_CHECK(scale_inv_colwise.data.shape.size() == 2, \"scale_inv_colwise must be a 2D tensor\");\n NVTE_CHECK(scale_inv_colwise.data.shape[0] % colwise_row_padding == 0,\n \"Wrong padding of scale_inv_colwise's rows\");\n NVTE_CHECK(scale_inv_colwise.data.shape[0] >= rows / 32, \"Invalid rows\");\n NVTE_CHECK(scale_inv_colwise.data.shape[1] % colwise_col_padding == 0,\n \"Wrong padding of scale_inv_colwise's cols\");\n NVTE_CHECK(scale_inv_colwise.data.shape[1] >= cols, \"Invalid cols\");\n NVTE_CHECK(scale_inv_colwise.dtype() == DType::kByte, \"Wrong dtype of scale_inv_colwise\");\n\n int blocks_x = (cols + kColsPerTile - 1) / kColsPerTile;\n int blocks_y = (rows + kRowsPerTile - 1) / kRowsPerTile;\n dim3 grid(blocks_x, blocks_y);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n input.dtype(), IType,\n mxfp8_scaling_partial_cast_kernel<<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(output_rowwise.data.dptr),\n reinterpret_cast(output_colwise.data.dptr),\n reinterpret_cast(scale_inv_rowwise.data.dptr),\n reinterpret_cast(scale_inv_colwise.data.dptr),\n scale_inv_rowwise.data.shape[1], scale_inv_colwise.data.shape[1], rows, cols,\n start_offset, input.data.shape[0]);)\n}\n\n} // namespace mxfp8_scaling_recipe\n} // namespace transformer_engine\n\nvoid nvte_mxfp8_scaling_compute_partial_amax(const NVTETensor input, NVTETensor amax_rowwise,\n NVTETensor amax_colwise, int rows, int cols,\n size_t start_offset, hipStream_t stream) {\n NVTE_API_CALL(nvte_mxfp8_scaling_compute_partial_amax);\n using namespace transformer_engine;\n mxfp8_scaling_recipe::mxfp8_scaling_compute_partial_amax(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(amax_rowwise),\n *convertNVTETensorCheck(amax_colwise), rows, cols, start_offset, stream);\n}\n\nvoid nvte_mxfp8_scaling_partial_cast(const NVTETensor input, NVTETensor output_rowwise,\n NVTETensor output_colwise, const NVTETensor scale_inv_rowwise,\n const NVTETensor scale_inv_colwise, int rows, int cols,\n size_t start_offset, hipStream_t stream) {\n NVTE_API_CALL(nvte_mxfp8_scaling_partial_cast);\n using namespace transformer_engine;\n mxfp8_scaling_recipe::mxfp8_scaling_partial_cast(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(output_rowwise),\n *convertNVTETensorCheck(output_colwise), *convertNVTETensorCheck(scale_inv_rowwise),\n *convertNVTETensorCheck(scale_inv_colwise), rows, cols, start_offset, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/recipe/mxfp8_scaling.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n\n#include \n\n#include \"../common.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\nnamespace fp8_block_scaling_recipe {\n\nconstexpr int kTileDim = 128;\nconstexpr int kThreadsPerBlock = 256;\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock)\n fp8_block_scaling_compute_partial_amax_kernel(const IType *input, float *amax_ptr,\n const size_t amax_stride_h,\n const size_t amax_stride_w, const size_t h,\n const size_t w, const size_t start_offset,\n const size_t len) {\n constexpr int kThreadsPerWarp = 32;\n constexpr int kLoopsPerRow = kTileDim / kThreadsPerWarp;\n constexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp;\n constexpr int kLoopsPerCol = kTileDim / kNumWarps;\n\n const int tile_col = blockIdx.x;\n const int tile_row = blockIdx.y;\n const size_t end_offset = start_offset + len;\n const IType *input_minus_offset = input - start_offset;\n\n __shared__ float smem[kNumWarps];\n float amax = 0.0f;\n\n for (int loop_col = 0; loop_col < kLoopsPerCol; ++loop_col) {\n size_t r = tile_row * kTileDim + loop_col * kNumWarps + threadIdx.x / kThreadsPerWarp;\n for (int loop_row = 0; loop_row < kLoopsPerRow; ++loop_row) {\n size_t c = tile_col * kTileDim + loop_row * kThreadsPerWarp + (threadIdx.x % kThreadsPerWarp);\n size_t idx = r * w + c;\n if (r < h && c < w && idx >= start_offset && idx < end_offset) {\n float other_amax = fabs(static_cast(input_minus_offset[idx]));\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n }\n }\n\n for (int delta = kThreadsPerWarp / 2; delta > 0; delta /= 2) {\n float other_amax = __shfl_down_sync(0xFFFFFFFF, amax, delta);\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n\n if (threadIdx.x % kThreadsPerWarp == 0) {\n smem[threadIdx.x / kThreadsPerWarp] = amax;\n }\n\n __syncthreads();\n\n if (threadIdx.x == 0) {\n for (int i = 0; i < kNumWarps; ++i) {\n float other_amax = smem[i];\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n amax_ptr[tile_row * amax_stride_h + tile_col * amax_stride_w] = amax;\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock)\n fp8_block_scaling_partial_cast_kernel(const IType *input, OType *output, const float *scale_ptr,\n const size_t scale_stride_h, const size_t scale_stride_w,\n const size_t h, const size_t w, const size_t start_offset,\n const size_t len) {\n using transformer_engine::Vec;\n\n static_assert(sizeof(OType) == 1);\n constexpr int kNumOutputElemsPerBank = 4 / sizeof(OType);\n constexpr int kThreadsPerWarp = 32;\n constexpr int kLoopsPerRow = kTileDim / kThreadsPerWarp;\n constexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp;\n constexpr int kRowsPerWarp = kTileDim / kNumWarps;\n\n __shared__ OType smem[kTileDim][kTileDim + kNumOutputElemsPerBank];\n\n const int tile_w = blockIdx.x;\n const int tile_h = blockIdx.y;\n const size_t end_offset = start_offset + len;\n const IType *input_minus_offset = input - start_offset;\n OType *output_minus_offset = output - start_offset;\n\n const float scale = scale_ptr[tile_h * scale_stride_h + tile_w * scale_stride_w];\n\n // Load input data into shared memory\n bool skip_store = true;\n for (int i = 0; i < kRowsPerWarp; ++i) {\n for (int j = 0; j < kLoopsPerRow; ++j) {\n const int h_in_smem = threadIdx.x / kThreadsPerWarp * kRowsPerWarp + i;\n const int w_in_smem = threadIdx.x % kThreadsPerWarp + kThreadsPerWarp * j;\n const int h_in_input = tile_h * kTileDim + h_in_smem;\n const int w_in_input = tile_w * kTileDim + w_in_smem;\n const size_t idx_in_input = static_cast(h_in_input) * w + w_in_input;\n if (h_in_input < h && w_in_input < w && idx_in_input >= start_offset &&\n idx_in_input < end_offset) {\n float inp = static_cast(input_minus_offset[idx_in_input]) * scale;\n smem[h_in_smem][w_in_smem] = static_cast(inp);\n skip_store = false;\n }\n }\n }\n\n for (int delta = kThreadsPerWarp / 2; delta > 0; delta /= 2) {\n bool other_skip_store = __shfl_down_sync(0xFFFFFFFF, skip_store, delta);\n skip_store = skip_store && other_skip_store;\n }\n skip_store = __shfl_sync(0xFFFFFFFF, skip_store, 0);\n if (skip_store) {\n return;\n }\n\n // Store the casted data into the output.\n // Note that this store operation might write \"out-of-bounds\", but it is intentional:\n // 1. The \"out-of-bounds\" here only crosses the boundary of the \"local shard\" (i.e., the region\n // from start_offset to end_offset), not the boundary of the entire output memory. Therefore,\n // this out-of-bounds write will not cause illegal memory access.\n // 2. We assume that the subsequent all-gather operation happens in-place, so any parts that\n // should not be updated here will be overwritten by the all-gather.\n // This tricky approach allows us to avoid checking whether each output index falls within\n // [start, end), resulting in a significant performance improvement.\n Vec vec_output;\n for (int i = 0; i < kRowsPerWarp; ++i) {\n const int row_in_smem = threadIdx.x / kThreadsPerWarp * kRowsPerWarp + i;\n const int col_in_smem = threadIdx.x % kThreadsPerWarp * kNumOutputElemsPerBank;\n for (int j = 0; j < kNumOutputElemsPerBank; ++j) {\n vec_output.data.elt[j] = smem[row_in_smem][col_in_smem + j];\n }\n const int row_in_output = tile_h * kTileDim + row_in_smem;\n const int col_in_output = tile_w * kTileDim + col_in_smem;\n const size_t idx_in_output = static_cast(row_in_output) * w + col_in_output;\n if (row_in_output < h) {\n if constexpr (kWidthAligned) {\n vec_output.store_to(output_minus_offset + idx_in_output);\n } else {\n int num = min(static_cast(kNumOutputElemsPerBank),\n static_cast(col_in_output < w ? w - col_in_output : 0));\n vec_output.store_to_elts(output_minus_offset, idx_in_output, num);\n }\n }\n }\n}\n\nvoid fp8_block_scaling_compute_partial_amax(const Tensor inp, Tensor amax, size_t h, size_t w,\n size_t amax_stride_h, size_t amax_stride_w,\n size_t start_offset, size_t block_len,\n cudaStream_t stream) {\n NVTE_CHECK(block_len == 128, \"Currently only block_len = 128 is supported\");\n\n size_t len = inp.numel();\n\n assert(h > 0 && w > 0);\n assert(start_offset < h * w);\n assert(start_offset + len <= h * w);\n\n size_t blocks_x = (w + kTileDim - 1) / kTileDim;\n size_t blocks_y = (h + kTileDim - 1) / kTileDim;\n assert(blocks_x <= std::numeric_limits::max());\n assert(blocks_y <= std::numeric_limits::max());\n dim3 grid(blocks_x, blocks_y);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n inp.dtype(), inp_dtype,\n fp8_block_scaling_compute_partial_amax_kernel\n <<>>(reinterpret_cast(inp.data.dptr),\n reinterpret_cast(amax.data.dptr),\n amax_stride_h, amax_stride_w, h, w, start_offset,\n len);)\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid fp8_block_scaling_partial_cast(const Tensor inp, Tensor out, const Tensor scale, size_t h,\n size_t w, size_t scale_stride_h, size_t scale_stride_w,\n size_t start_offset, size_t block_len, const DType out_dtype,\n cudaStream_t stream) {\n NVTE_CHECK(block_len == 128, \"Currently only block_len = 128 is supported\");\n\n size_t len = inp.numel();\n\n assert(h > 0 && w > 0);\n assert(start_offset < h * w);\n assert(start_offset + len <= h * w);\n\n size_t blocks_x = (w + kTileDim - 1) / kTileDim;\n size_t blocks_y = (h + kTileDim - 1) / kTileDim;\n assert(blocks_x <= std::numeric_limits::max());\n assert(blocks_y <= std::numeric_limits::max());\n dim3 grid(blocks_x, blocks_y);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n inp.dtype(), inp_dtype,\n TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(\n out_dtype, fp8_type,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n w % kTileDim == 0, kWidthAligned,\n fp8_block_scaling_partial_cast_kernel\n <<>>(\n reinterpret_cast(inp.data.dptr),\n reinterpret_cast(out.data.dptr),\n reinterpret_cast(scale.data.dptr), scale_stride_h, scale_stride_w,\n h, w, start_offset, len);)))\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n} // namespace fp8_block_scaling_recipe\n} // namespace transformer_engine\n\nvoid nvte_fp8_block_scaling_compute_partial_amax(const NVTETensor inp, NVTETensor amax, size_t h,\n size_t w, size_t amax_stride_h,\n size_t amax_stride_w, size_t start_offset,\n size_t block_len, cudaStream_t stream) {\n NVTE_API_CALL(nvte_fp8_block_scaling_compute_partial_amax);\n using namespace transformer_engine;\n fp8_block_scaling_recipe::fp8_block_scaling_compute_partial_amax(\n *convertNVTETensorCheck(inp), *convertNVTETensorCheck(amax), h, w, amax_stride_h,\n amax_stride_w, start_offset, block_len, stream);\n}\n\nvoid nvte_fp8_block_scaling_partial_cast(const NVTETensor inp, NVTETensor out,\n const NVTETensor scale, size_t h, size_t w,\n size_t scale_stride_h, size_t scale_stride_w,\n size_t start_offset, size_t block_len,\n const NVTEDType out_dtype, cudaStream_t stream) {\n NVTE_API_CALL(nvte_fp8_block_scaling_partial_cast);\n using namespace transformer_engine;\n fp8_block_scaling_recipe::fp8_block_scaling_partial_cast(\n *convertNVTETensorCheck(inp), *convertNVTETensorCheck(out), *convertNVTETensorCheck(scale), h,\n w, scale_stride_h, scale_stride_w, start_offset, block_len, static_cast(out_dtype),\n stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \n\n#include \"../common.h\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\nnamespace fp8_block_scaling_recipe {\n\nconstexpr int kTileDim = 128;\nconstexpr int kThreadsPerBlock = 256;\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock)\n fp8_block_scaling_compute_partial_amax_kernel(const IType *input, float *amax_ptr,\n const size_t amax_stride_h,\n const size_t amax_stride_w, const size_t h,\n const size_t w, const size_t start_offset,\n const size_t len) {\n constexpr int kThreadsPerWarp = 32;\n constexpr int kLoopsPerRow = kTileDim / kThreadsPerWarp;\n constexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp;\n constexpr int kLoopsPerCol = kTileDim / kNumWarps;\n\n const int tile_col = blockIdx.x;\n const int tile_row = blockIdx.y;\n const size_t end_offset = start_offset + len;\n const IType *input_minus_offset = input - start_offset;\n\n __shared__ float smem[kNumWarps];\n float amax = 0.0f;\n\n for (int loop_col = 0; loop_col < kLoopsPerCol; ++loop_col) {\n size_t r = tile_row * kTileDim + loop_col * kNumWarps + threadIdx.x / kThreadsPerWarp;\n for (int loop_row = 0; loop_row < kLoopsPerRow; ++loop_row) {\n size_t c = tile_col * kTileDim + loop_row * kThreadsPerWarp + (threadIdx.x % kThreadsPerWarp);\n size_t idx = r * w + c;\n if (r < h && c < w && idx >= start_offset && idx < end_offset) {\n float other_amax = fabs(static_cast(input_minus_offset[idx]));\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n }\n }\n\n for (int delta = kThreadsPerWarp / 2; delta > 0; delta /= 2) {\n float other_amax = __shfl_down_sync(0xFFFFFFFF, amax, delta);\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n\n if (threadIdx.x % kThreadsPerWarp == 0) {\n smem[threadIdx.x / kThreadsPerWarp] = amax;\n }\n\n __syncthreads();\n\n if (threadIdx.x == 0) {\n for (int i = 0; i < kNumWarps; ++i) {\n float other_amax = smem[i];\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n amax_ptr[tile_row * amax_stride_h + tile_col * amax_stride_w] = amax;\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock)\n fp8_block_scaling_partial_cast_kernel(const IType *input, OType *output, const float *scale_ptr,\n const size_t scale_stride_h, const size_t scale_stride_w,\n const size_t h, const size_t w, const size_t start_offset,\n const size_t len) {\n using transformer_engine::Vec;\n\n static_assert(sizeof(OType) == 1);\n constexpr int kNumOutputElemsPerBank = 4 / sizeof(OType);\n constexpr int kThreadsPerWarp = 32;\n constexpr int kLoopsPerRow = kTileDim / kThreadsPerWarp;\n constexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp;\n constexpr int kRowsPerWarp = kTileDim / kNumWarps;\n\n __shared__ OType smem[kTileDim][kTileDim + kNumOutputElemsPerBank];\n\n const int tile_w = blockIdx.x;\n const int tile_h = blockIdx.y;\n const size_t end_offset = start_offset + len;\n const IType *input_minus_offset = input - start_offset;\n OType *output_minus_offset = output - start_offset;\n\n const float scale = scale_ptr[tile_h * scale_stride_h + tile_w * scale_stride_w];\n\n // Load input data into shared memory\n bool skip_store = true;\n for (int i = 0; i < kRowsPerWarp; ++i) {\n for (int j = 0; j < kLoopsPerRow; ++j) {\n const int h_in_smem = threadIdx.x / kThreadsPerWarp * kRowsPerWarp + i;\n const int w_in_smem = threadIdx.x % kThreadsPerWarp + kThreadsPerWarp * j;\n const int h_in_input = tile_h * kTileDim + h_in_smem;\n const int w_in_input = tile_w * kTileDim + w_in_smem;\n const size_t idx_in_input = static_cast(h_in_input) * w + w_in_input;\n if (h_in_input < h && w_in_input < w && idx_in_input >= start_offset &&\n idx_in_input < end_offset) {\n float inp = static_cast(input_minus_offset[idx_in_input]) * scale;\n smem[h_in_smem][w_in_smem] = static_cast(inp);\n skip_store = false;\n }\n }\n }\n\n for (int delta = kThreadsPerWarp / 2; delta > 0; delta /= 2) {\n bool other_skip_store = __shfl_down_sync(0xFFFFFFFF, skip_store, delta);\n skip_store = skip_store && other_skip_store;\n }\n skip_store = __shfl_sync(0xFFFFFFFF, skip_store, 0);\n if (skip_store) {\n return;\n }\n\n // Store the casted data into the output.\n // Note that this store operation might write \"out-of-bounds\", but it is intentional:\n // 1. The \"out-of-bounds\" here only crosses the boundary of the \"local shard\" (i.e., the region\n // from start_offset to end_offset), not the boundary of the entire output memory. Therefore,\n // this out-of-bounds write will not cause illegal memory access.\n // 2. We assume that the subsequent all-gather operation happens in-place, so any parts that\n // should not be updated here will be overwritten by the all-gather.\n // This tricky approach allows us to avoid checking whether each output index falls within\n // [start, end), resulting in a significant performance improvement.\n Vec vec_output;\n for (int i = 0; i < kRowsPerWarp; ++i) {\n const int row_in_smem = threadIdx.x / kThreadsPerWarp * kRowsPerWarp + i;\n const int col_in_smem = threadIdx.x % kThreadsPerWarp * kNumOutputElemsPerBank;\n for (int j = 0; j < kNumOutputElemsPerBank; ++j) {\n vec_output.data.elt[j] = smem[row_in_smem][col_in_smem + j];\n }\n const int row_in_output = tile_h * kTileDim + row_in_smem;\n const int col_in_output = tile_w * kTileDim + col_in_smem;\n const size_t idx_in_output = static_cast(row_in_output) * w + col_in_output;\n if (row_in_output < h) {\n if constexpr (kWidthAligned) {\n vec_output.store_to(output_minus_offset + idx_in_output);\n } else {\n int num = min(static_cast(kNumOutputElemsPerBank),\n static_cast(col_in_output < w ? w - col_in_output : 0));\n vec_output.store_to_elts(output_minus_offset, idx_in_output, num);\n }\n }\n }\n}\n\nvoid fp8_block_scaling_compute_partial_amax(const Tensor inp, Tensor amax, size_t h, size_t w,\n size_t amax_stride_h, size_t amax_stride_w,\n size_t start_offset, size_t block_len,\n hipStream_t stream) {\n NVTE_CHECK(block_len == 128, \"Currently only block_len = 128 is supported\");\n\n size_t len = inp.numel();\n\n assert(h > 0 && w > 0);\n assert(start_offset < h * w);\n assert(start_offset + len <= h * w);\n\n size_t blocks_x = (w + kTileDim - 1) / kTileDim;\n size_t blocks_y = (h + kTileDim - 1) / kTileDim;\n assert(blocks_x <= std::numeric_limits::max());\n assert(blocks_y <= std::numeric_limits::max());\n dim3 grid(blocks_x, blocks_y);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n inp.dtype(), inp_dtype,\n fp8_block_scaling_compute_partial_amax_kernel\n <<>>(reinterpret_cast(inp.data.dptr),\n reinterpret_cast(amax.data.dptr),\n amax_stride_h, amax_stride_w, h, w, start_offset,\n len);)\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid fp8_block_scaling_partial_cast(const Tensor inp, Tensor out, const Tensor scale, size_t h,\n size_t w, size_t scale_stride_h, size_t scale_stride_w,\n size_t start_offset, size_t block_len, const DType out_dtype,\n hipStream_t stream) {\n NVTE_CHECK(block_len == 128, \"Currently only block_len = 128 is supported\");\n\n size_t len = inp.numel();\n\n assert(h > 0 && w > 0);\n assert(start_offset < h * w);\n assert(start_offset + len <= h * w);\n\n size_t blocks_x = (w + kTileDim - 1) / kTileDim;\n size_t blocks_y = (h + kTileDim - 1) / kTileDim;\n assert(blocks_x <= std::numeric_limits::max());\n assert(blocks_y <= std::numeric_limits::max());\n dim3 grid(blocks_x, blocks_y);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n inp.dtype(), inp_dtype,\n TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(\n out_dtype, fp8_type,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n w % kTileDim == 0, kWidthAligned,\n fp8_block_scaling_partial_cast_kernel\n <<>>(\n reinterpret_cast(inp.data.dptr),\n reinterpret_cast(out.data.dptr),\n reinterpret_cast(scale.data.dptr), scale_stride_h, scale_stride_w,\n h, w, start_offset, len);)))\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n} // namespace fp8_block_scaling_recipe\n} // namespace transformer_engine\n\nvoid nvte_fp8_block_scaling_compute_partial_amax(const NVTETensor inp, NVTETensor amax, size_t h,\n size_t w, size_t amax_stride_h,\n size_t amax_stride_w, size_t start_offset,\n size_t block_len, hipStream_t stream) {\n NVTE_API_CALL(nvte_fp8_block_scaling_compute_partial_amax);\n using namespace transformer_engine;\n fp8_block_scaling_recipe::fp8_block_scaling_compute_partial_amax(\n *convertNVTETensorCheck(inp), *convertNVTETensorCheck(amax), h, w, amax_stride_h,\n amax_stride_w, start_offset, block_len, stream);\n}\n\nvoid nvte_fp8_block_scaling_partial_cast(const NVTETensor inp, NVTETensor out,\n const NVTETensor scale, size_t h, size_t w,\n size_t scale_stride_h, size_t scale_stride_w,\n size_t start_offset, size_t block_len,\n const NVTEDType out_dtype, hipStream_t stream) {\n NVTE_API_CALL(nvte_fp8_block_scaling_partial_cast);\n using namespace transformer_engine;\n fp8_block_scaling_recipe::fp8_block_scaling_partial_cast(\n *convertNVTETensorCheck(inp), *convertNVTETensorCheck(out), *convertNVTETensorCheck(scale), h,\n w, scale_stride_h, scale_stride_w, start_offset, block_len, static_cast(out_dtype),\n stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/recipe/fp8_block_scaling.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"../util/vectorized_pointwise.h\"\n#include \"recipe_common.cuh\"\n\nnamespace transformer_engine {\nnamespace {\n\nconstexpr int amax_kernel_threads = 512;\n\n__launch_bounds__(1) __global__ void zero_amax_kernel(float *amax_ptr, const float *noop_ptr) {\n if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {\n return;\n }\n *amax_ptr = 0;\n}\n\ntemplate \n__launch_bounds__(amax_kernel_threads) __global__\n void amax_kernel(const InputType *input, float *amax, const size_t N,\n const size_t num_aligned_elements, const float *noop_ptr) {\n if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {\n return;\n }\n\n VectorizedLoader loader(input, N);\n InputType max = 0.f;\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const size_t M = num_aligned_elements;\n\n for (size_t tid = blockIdx.x * blockDim.x + threadIdx.x; tid < M; tid += gridDim.x * blockDim.x) {\n loader.load(tid, N);\n#pragma unroll\n for (int i = 0; i < nvec; ++i) {\n const InputType val = static_cast(loader.separate()[i]);\n __builtin_assume(max >= InputType{0.f});\n if constexpr (std::is_same_v) {\n#if __CUDA_ARCH__ >= 800\n max = __hmax(__habs(val), max);\n#else // Turing\n max = static_cast<__nv_bfloat16>(\n fmaxf(fabsf(static_cast(val)), static_cast(max)));\n#endif\n } else if constexpr (std::is_same_v) {\n max = __hmax(__habs(val), max);\n } else {\n max = fmaxf(fabsf(val), max);\n }\n }\n }\n\n // Reduce amax over block\n max = reduce_max(max, warp_id);\n if (threadIdx.x == 0) {\n atomicMaxFloat(amax, max);\n }\n}\n\ntemplate \nvoid launch_amax_kernel(const InputType *input, float *amax, const size_t N, const float *noop_ptr,\n cudaStream_t stream) {\n // Zero out amax so we can update with atomic max\n zero_amax_kernel<<<1, 1, 0, stream>>>(amax, noop_ptr);\n NVTE_CHECK_CUDA(cudaGetLastError());\n\n // Return immediately if tensor is empty\n if (N == 0) {\n return;\n }\n\n // Figure out alignment\n auto align = CheckAlignment(N, nvec, input);\n size_t num_aligned_elements = get_num_aligned_elements(input, N, nvec, sizeof(InputType));\n\n // Figure out CUDA blocks\n constexpr size_t threads = amax_kernel_threads;\n size_t num_blocks = DIVUP(num_aligned_elements, threads);\n constexpr size_t max_blocks = 65535;\n num_blocks = std::min(num_blocks, max_blocks);\n\n // Launch kernel\n switch (align) {\n case Alignment::SAME_ALIGNED:\n amax_kernel\n <<>>(input, amax, N, num_aligned_elements, noop_ptr);\n break;\n case Alignment::SAME_UNALIGNED:\n amax_kernel\n <<>>(input, amax, N, num_aligned_elements, noop_ptr);\n break;\n case Alignment::DIFFERENT: {\n // This case is a logic error, since there is only one pointer (input)\n // in the alignment check. Still safe to process without vectorization.\n amax_kernel<1, true, InputType>\n <<>>(input, amax, N, N, noop_ptr);\n break;\n }\n }\n\n // Check results\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n} // namespace\n} // namespace transformer_engine\n\nnamespace {\n\nvoid compute_amax_impl(const NVTETensor input_, const NVTETensor output_, cudaStream_t stream,\n const NVTEQuantizationConfig config_) {\n using namespace transformer_engine;\n\n // Check input tensor\n NVTE_CHECK(input_ != nullptr, \"Invalid input tensor (got NULL)\");\n const auto &input = *convertNVTETensorCheck(input_);\n NVTE_CHECK(input.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor for amax computation must unquantized, \"\n \"but got scaling_mode=\",\n to_string(input.scaling_mode));\n NVTE_CHECK(!is_fp8_dtype(input.data.dtype),\n \"Input tensor for amax computation must be unquantized, but got dtype=\",\n to_string(input.data.dtype));\n NVTE_CHECK(input.data.dptr != nullptr, \"Input tensor for amax computation has no data\");\n CheckInputTensor(input, \"input_compute_amax\");\n\n // Check output tensor\n NVTE_CHECK(output_ != nullptr, \"Invalid output tensor (got NULL)\");\n auto &output = *convertNVTETensorCheck(output_);\n NVTE_CHECK(output.scaling_mode == NVTE_DELAYED_TENSOR_SCALING ||\n output.scaling_mode == NVTE_NVFP4_1D_SCALING,\n \"Output tensor for amax computation must be FP8 tensor with per-tensor scaling or \"\n \"NVFP4 1D scaling, \"\n \"but got scaling_mode=\",\n to_string(output.scaling_mode));\n NVTE_CHECK(output.amax.numel() == 1,\n \"Output tensor for amax computation has invalid amax tensor \"\n \"(expected 1 entry, got shape=\",\n output.amax.shape, \")\");\n NVTE_CHECK(output.amax.dptr != nullptr || output.columnwise_amax.dptr != nullptr,\n \"Output tensor for amax computation has amax tensor without data\");\n NVTE_CHECK(output.amax.dtype == DType::kFloat32,\n \"Output tensor for amax computation has invalid amax tensor \"\n \"(expected FP32, got dtype=\",\n to_string(output.amax.dtype), \")\");\n CheckOutputTensor(output, \"output_compute_amax\", true);\n\n float *noop_ptr = nullptr;\n if (config_ != nullptr) {\n const QuantizationConfig *config_cpp = reinterpret_cast(config_);\n\n // extract noop tensor from quant_config_cpp if it's not null\n const NVTETensor noop = config_cpp ? config_cpp->noop_tensor : nullptr;\n noop_ptr = reinterpret_cast(\n (noop != nullptr ? convertNVTETensorCheck(noop)->data.dptr : nullptr));\n }\n\n // Compute amax\n float *amax_ptr = reinterpret_cast(\n (output.amax.dptr != nullptr) ? output.amax.dptr : output.columnwise_amax.dptr);\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.data.dtype, IType, constexpr int nvec = 32 / sizeof(IType); launch_amax_kernel(\n reinterpret_cast(input.data.dptr), amax_ptr, input.data.numel(), noop_ptr,\n stream);); // NOLINT(*)\n}\n\n} // anonymous namespace\n\nvoid nvte_compute_amax(const NVTETensor input_, const NVTETensor output_, cudaStream_t stream) {\n NVTE_API_CALL(nvte_compute_amax);\n compute_amax_impl(input_, output_, stream, nullptr);\n}\n\nvoid nvte_compute_amax_with_config(const NVTETensor input_, const NVTETensor output_,\n const NVTEQuantizationConfig config_, cudaStream_t stream) {\n NVTE_API_CALL(nvte_compute_amax_with_config);\n compute_amax_impl(input_, output_, stream, config_);\n}\n\nnamespace transformer_engine {\nnamespace {\n\n__global__ void compute_scale_from_amax_kernel(const float *amax_ptr, float *scale_ptr,\n const float max_fp8, const bool force_pow_2_scales,\n const float epsilon, const float *noop_ptr) {\n if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {\n return;\n }\n\n *scale_ptr = compute_scale_from_amax(*amax_ptr, max_fp8, force_pow_2_scales, epsilon,\n std::numeric_limits::max());\n}\n\n} // namespace\n} // namespace transformer_engine\n\nvoid nvte_compute_scale_from_amax(NVTETensor output_, const NVTEQuantizationConfig config_,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_compute_scale_from_amax);\n using namespace transformer_engine;\n\n // Check output tensor\n NVTE_CHECK(output_ != nullptr, \"Invalid output tensor (got NULL)\");\n auto &output = *convertNVTETensorCheck(output_);\n NVTE_CHECK(output.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Tensor must be FP8 tensor with per-tensor scaling, \"\n \"but got scaling_mode=\",\n to_string(output.scaling_mode));\n NVTE_CHECK(is_fp8_dtype(output.data.dtype),\n \"Tensor must be FP8, but got dtype=\", to_string(output.data.dtype));\n NVTE_CHECK(output.amax.numel() == 1,\n \"Tensor has invalid amax tensor (expected 1 entry, got shape=\", output.amax.shape,\n \")\");\n NVTE_CHECK(output.amax.dptr != nullptr, \"Tensor has amax tensor without data\");\n NVTE_CHECK(output.amax.dtype == DType::kFloat32,\n \"Tensor has invalid amax tensor (expected FP32, got dtype=\",\n to_string(output.amax.dtype), \")\");\n NVTE_CHECK(output.scale.numel() == 1,\n \"Tensor has invalid scale tensor (expected 1 entry, got shape=\", output.scale.shape,\n \")\");\n NVTE_CHECK(output.scale.dptr != nullptr, \"Tensor has scale tensor without data\");\n NVTE_CHECK(output.scale.dtype == DType::kFloat32,\n \"Tensor has invalid scale tensor (expected FP32, got dtype=\",\n to_string(output.scale.dtype), \")\");\n\n // Check config\n NVTE_CHECK(config_ != nullptr, \"Invalid config (got NULL)\");\n const auto &config = *reinterpret_cast(config_);\n\n // Maximum FP8 value\n float max_fp8 = 0.f;\n TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(output.data.dtype, DType,\n max_fp8 = Quantized_Limits::max_norm;);\n\n // noop tensor for cuda graph\n float *noop_ptr = nullptr;\n if (config_ != nullptr) {\n const QuantizationConfig *config_cpp = reinterpret_cast(config_);\n\n // extract noop tensor from quant_config_cpp if it's not null\n const NVTETensor noop = config_cpp ? config_cpp->noop_tensor : nullptr;\n noop_ptr = reinterpret_cast(\n (noop != nullptr ? convertNVTETensorCheck(noop)->data.dptr : nullptr));\n }\n\n // Update scale\n compute_scale_from_amax_kernel<<<1, 1, 0, stream>>>(\n reinterpret_cast(output.amax.dptr),\n reinterpret_cast(output.scale.dptr), max_fp8, config.force_pow_2_scales,\n config.amax_epsilon, noop_ptr);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/logging.h\"\n#include \"../util/vectorized_pointwise.h\"\n#include \"recipe_common.cuh\"\n\nnamespace transformer_engine {\nnamespace {\n\nconstexpr int amax_kernel_threads = 512;\n\n__launch_bounds__(1) __global__ void zero_amax_kernel(float *amax_ptr, const float *noop_ptr) {\n if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {\n return;\n }\n *amax_ptr = 0;\n}\n\ntemplate \n__launch_bounds__(amax_kernel_threads) __global__\n void amax_kernel(const InputType *input, float *amax, const size_t N,\n const size_t num_aligned_elements, const float *noop_ptr) {\n if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {\n return;\n }\n\n VectorizedLoader loader(input, N);\n InputType max = 0.f;\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const size_t M = num_aligned_elements;\n\n for (size_t tid = blockIdx.x * blockDim.x + threadIdx.x; tid < M; tid += gridDim.x * blockDim.x) {\n loader.load(tid, N);\n#pragma unroll\n for (int i = 0; i < nvec; ++i) {\n const InputType val = static_cast(loader.separate()[i]);\n __builtin_assume(max >= InputType{0.f});\n if constexpr (std::is_same_v) {\n#if __CUDA_ARCH__ >= 800\n max = __hmax(__habs(val), max);\n#else // Turing\n max = static_cast<__hip_bfloat16>(\n fmaxf(fabsf(static_cast(val)), static_cast(max)));\n#endif\n } else if constexpr (std::is_same_v) {\n max = __hmax(__habs(val), max);\n } else {\n max = fmaxf(fabsf(val), max);\n }\n }\n }\n\n // Reduce amax over block\n max = reduce_max(max, warp_id);\n if (threadIdx.x == 0) {\n atomicMaxFloat(amax, max);\n }\n}\n\ntemplate \nvoid launch_amax_kernel(const InputType *input, float *amax, const size_t N, const float *noop_ptr,\n hipStream_t stream) {\n // Zero out amax so we can update with atomic max\n zero_amax_kernel<<<1, 1, 0, stream>>>(amax, noop_ptr);\n NVTE_CHECK_CUDA(hipGetLastError());\n\n // Return immediately if tensor is empty\n if (N == 0) {\n return;\n }\n\n // Figure out alignment\n auto align = CheckAlignment(N, nvec, input);\n size_t num_aligned_elements = get_num_aligned_elements(input, N, nvec, sizeof(InputType));\n\n // Figure out CUDA blocks\n constexpr size_t threads = amax_kernel_threads;\n size_t num_blocks = DIVUP(num_aligned_elements, threads);\n constexpr size_t max_blocks = 65535;\n num_blocks = std::min(num_blocks, max_blocks);\n\n // Launch kernel\n switch (align) {\n case Alignment::SAME_ALIGNED:\n amax_kernel\n <<>>(input, amax, N, num_aligned_elements, noop_ptr);\n break;\n case Alignment::SAME_UNALIGNED:\n amax_kernel\n <<>>(input, amax, N, num_aligned_elements, noop_ptr);\n break;\n case Alignment::DIFFERENT: {\n // This case is a logic error, since there is only one pointer (input)\n // in the alignment check. Still safe to process without vectorization.\n amax_kernel<1, true, InputType>\n <<>>(input, amax, N, N, noop_ptr);\n break;\n }\n }\n\n // Check results\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n} // namespace\n} // namespace transformer_engine\n\nnamespace {\n\nvoid compute_amax_impl(const NVTETensor input_, const NVTETensor output_, hipStream_t stream,\n const NVTEQuantizationConfig config_) {\n using namespace transformer_engine;\n\n // Check input tensor\n NVTE_CHECK(input_ != nullptr, \"Invalid input tensor (got NULL)\");\n const auto &input = *convertNVTETensorCheck(input_);\n NVTE_CHECK(input.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor for amax computation must unquantized, \"\n \"but got scaling_mode=\",\n to_string(input.scaling_mode));\n NVTE_CHECK(!is_fp8_dtype(input.data.dtype),\n \"Input tensor for amax computation must be unquantized, but got dtype=\",\n to_string(input.data.dtype));\n NVTE_CHECK(input.data.dptr != nullptr, \"Input tensor for amax computation has no data\");\n CheckInputTensor(input, \"input_compute_amax\");\n\n // Check output tensor\n NVTE_CHECK(output_ != nullptr, \"Invalid output tensor (got NULL)\");\n auto &output = *convertNVTETensorCheck(output_);\n NVTE_CHECK(output.scaling_mode == NVTE_DELAYED_TENSOR_SCALING ||\n output.scaling_mode == NVTE_NVFP4_1D_SCALING,\n \"Output tensor for amax computation must be FP8 tensor with per-tensor scaling or \"\n \"NVFP4 1D scaling, \"\n \"but got scaling_mode=\",\n to_string(output.scaling_mode));\n NVTE_CHECK(output.amax.numel() == 1,\n \"Output tensor for amax computation has invalid amax tensor \"\n \"(expected 1 entry, got shape=\",\n output.amax.shape, \")\");\n NVTE_CHECK(output.amax.dptr != nullptr || output.columnwise_amax.dptr != nullptr,\n \"Output tensor for amax computation has amax tensor without data\");\n NVTE_CHECK(output.amax.dtype == DType::kFloat32,\n \"Output tensor for amax computation has invalid amax tensor \"\n \"(expected FP32, got dtype=\",\n to_string(output.amax.dtype), \")\");\n CheckOutputTensor(output, \"output_compute_amax\", true);\n\n float *noop_ptr = nullptr;\n if (config_ != nullptr) {\n const QuantizationConfig *config_cpp = reinterpret_cast(config_);\n\n // extract noop tensor from quant_config_cpp if it's not null\n const NVTETensor noop = config_cpp ? config_cpp->noop_tensor : nullptr;\n noop_ptr = reinterpret_cast(\n (noop != nullptr ? convertNVTETensorCheck(noop)->data.dptr : nullptr));\n }\n\n // Compute amax\n float *amax_ptr = reinterpret_cast(\n (output.amax.dptr != nullptr) ? output.amax.dptr : output.columnwise_amax.dptr);\n TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(\n input.data.dtype, IType, constexpr int nvec = 32 / sizeof(IType); launch_amax_kernel(\n reinterpret_cast(input.data.dptr), amax_ptr, input.data.numel(), noop_ptr,\n stream);); // NOLINT(*)\n}\n\n} // anonymous namespace\n\nvoid nvte_compute_amax(const NVTETensor input_, const NVTETensor output_, hipStream_t stream) {\n NVTE_API_CALL(nvte_compute_amax);\n compute_amax_impl(input_, output_, stream, nullptr);\n}\n\nvoid nvte_compute_amax_with_config(const NVTETensor input_, const NVTETensor output_,\n const NVTEQuantizationConfig config_, hipStream_t stream) {\n NVTE_API_CALL(nvte_compute_amax_with_config);\n compute_amax_impl(input_, output_, stream, config_);\n}\n\nnamespace transformer_engine {\nnamespace {\n\n__global__ void compute_scale_from_amax_kernel(const float *amax_ptr, float *scale_ptr,\n const float max_fp8, const bool force_pow_2_scales,\n const float epsilon, const float *noop_ptr) {\n if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {\n return;\n }\n\n *scale_ptr = compute_scale_from_amax(*amax_ptr, max_fp8, force_pow_2_scales, epsilon,\n std::numeric_limits::max());\n}\n\n} // namespace\n} // namespace transformer_engine\n\nvoid nvte_compute_scale_from_amax(NVTETensor output_, const NVTEQuantizationConfig config_,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_compute_scale_from_amax);\n using namespace transformer_engine;\n\n // Check output tensor\n NVTE_CHECK(output_ != nullptr, \"Invalid output tensor (got NULL)\");\n auto &output = *convertNVTETensorCheck(output_);\n NVTE_CHECK(output.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Tensor must be FP8 tensor with per-tensor scaling, \"\n \"but got scaling_mode=\",\n to_string(output.scaling_mode));\n NVTE_CHECK(is_fp8_dtype(output.data.dtype),\n \"Tensor must be FP8, but got dtype=\", to_string(output.data.dtype));\n NVTE_CHECK(output.amax.numel() == 1,\n \"Tensor has invalid amax tensor (expected 1 entry, got shape=\", output.amax.shape,\n \")\");\n NVTE_CHECK(output.amax.dptr != nullptr, \"Tensor has amax tensor without data\");\n NVTE_CHECK(output.amax.dtype == DType::kFloat32,\n \"Tensor has invalid amax tensor (expected FP32, got dtype=\",\n to_string(output.amax.dtype), \")\");\n NVTE_CHECK(output.scale.numel() == 1,\n \"Tensor has invalid scale tensor (expected 1 entry, got shape=\", output.scale.shape,\n \")\");\n NVTE_CHECK(output.scale.dptr != nullptr, \"Tensor has scale tensor without data\");\n NVTE_CHECK(output.scale.dtype == DType::kFloat32,\n \"Tensor has invalid scale tensor (expected FP32, got dtype=\",\n to_string(output.scale.dtype), \")\");\n\n // Check config\n NVTE_CHECK(config_ != nullptr, \"Invalid config (got NULL)\");\n const auto &config = *reinterpret_cast(config_);\n\n // Maximum FP8 value\n float max_fp8 = 0.f;\n TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(output.data.dtype, DType,\n max_fp8 = Quantized_Limits::max_norm;);\n\n // noop tensor for cuda graph\n float *noop_ptr = nullptr;\n if (config_ != nullptr) {\n const QuantizationConfig *config_cpp = reinterpret_cast(config_);\n\n // extract noop tensor from quant_config_cpp if it's not null\n const NVTETensor noop = config_cpp ? config_cpp->noop_tensor : nullptr;\n noop_ptr = reinterpret_cast(\n (noop != nullptr ? convertNVTETensorCheck(noop)->data.dptr : nullptr));\n }\n\n // Update scale\n compute_scale_from_amax_kernel<<<1, 1, 0, stream>>>(\n reinterpret_cast(output.amax.dptr),\n reinterpret_cast(output.scale.dptr), max_fp8, config.force_pow_2_scales,\n config.amax_epsilon, noop_ptr);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/recipe/current_scaling.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n\n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/ptx.cuh\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\nnamespace nvfp4_recipe {\n\n#if FP4_TYPE_SUPPORTED\n/*\n * ---------------------------------------------------------------------------\n * NVFP4 2D PARTIAL-SHARD KERNEL DESIGN\n *\n * These kernels mirror the FP8 block-scaling helpers but operate on shard-local\n * slices and nibble-packed FP4 rowwise buffers. One CUDA block covers a logical\n * 16x16 tile (grid = ceil(W/16) x ceil(H/16), blockDim = 256 threads).\n *\n * 1) Partial Amax (`nvfp4_2d_compute_partial_amax_kernel`)\n * - Warps sweep the tile using nested loops, accumulating local maxima only\n * for elements in [start_offset, start_offset + len).\n * - Shared memory reduces the 8 warp maxima; the block writes a float into\n * `amax_ptr[tile_row * stride_h + tile_col * stride_w]`.\n *\n * Tile/warp mapping (each '#' = elements visited by that warp):\n *\n * +------------------+\n * |########..........| Warp 0\n * |########..........| Warp 1\n * | ... |\n * |########..........| Warp 7\n * +------------------+\n *\n * 2) Partial Cast (`nvfp4_2d_partial_cast_kernel`)\n * - Stage the tile into shared memory (same pattern as FP8).\n * - For each 4-value group, build float2 pairs and call\n * `ptx::mul_cvt_fp32_to_fp4_4x`, producing packed FP4 nibbles.\n * - Compute a shard-local byte index and update only the owned nibble(s)\n * using read-modify-write:\n *\n * packed_bits = [mw3 | mw2 | mw1 | mw0]\n * byte_idx = (ref_elem_idx - start_offset) >> 1\n * if elem_idx % 2 == 0: // low nibble\n * byte = (byte & 0xF0) | nibble\n * else: // high nibble\n * byte = (byte & 0x0F) | (nibble << 4)\n *\n * Thread coverage inside a tile:\n *\n * rows: 16 columns: 16\n * Warp 0 -> rows 0-1 lanes sweep cols 0..3, 4..7, ...\n * Warp 1 -> rows 2-3 (groups of 4 elements per thread)\n * ...\n * Warp 7 -> rows 14-15\n * ---------------------------------------------------------------------------\n */\n\n// constexpr float factor = 6.0 * 6.0 * 448.0 * 448.0;\nconstexpr float factor_inv = 1.0 / (6.0 * 6.0 * 448.0 * 448.0);\nconstexpr int kTileDim = 16;\nconstexpr int kThreadsPerBlock = 256;\n\n// Kernel to compute alpha *= amax_A * amax_B / factor\n__global__ void compute_nvfp4_per_tensor_scale_kernel(float alpha_in, const float *amax_A,\n const float *amax_B, float *alpha_out) {\n // factor is defined in the enclosing namespace\n *alpha_out = alpha_in * (*amax_A) * (*amax_B) * factor_inv;\n}\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock)\n nvfp4_2d_compute_partial_amax_kernel(const IType *input, float *amax_ptr,\n const size_t amax_stride_h, const size_t amax_stride_w,\n const size_t h, const size_t w, const size_t start_offset,\n const size_t len) {\n constexpr int kThreadsPerWarp = 32;\n constexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp;\n static_assert(kTileDim * kTileDim == kThreadsPerBlock);\n\n const size_t tile_col = blockIdx.x;\n const size_t tile_row = blockIdx.y;\n const size_t end_offset = start_offset + len;\n const IType *input_minus_offset = input - start_offset;\n\n __shared__ float smem[kNumWarps];\n float amax = 0.0f;\n\n size_t r = tile_row * kTileDim + threadIdx.x / kTileDim;\n size_t c = tile_col * kTileDim + threadIdx.x % kTileDim;\n size_t idx = r * w + c;\n if (r < h && c < w && idx >= start_offset && idx < end_offset) {\n amax = fabs(static_cast(input_minus_offset[idx]));\n }\n\n for (int delta = kThreadsPerWarp / 2; delta > 0; delta /= 2) {\n float other_amax = __shfl_down_sync(0xFFFFFFFF, amax, delta);\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n\n if (threadIdx.x % kThreadsPerWarp == 0) {\n smem[threadIdx.x / kThreadsPerWarp] = amax;\n }\n\n __syncthreads();\n\n if (threadIdx.x == 0) {\n for (int i = 0; i < kNumWarps; ++i) {\n float other_amax = smem[i];\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n amax_ptr[tile_row * amax_stride_h + tile_col * amax_stride_w] = amax;\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock)\n nvfp4_2d_partial_cast_kernel(const IType *input, uint8_t *output, const float *decode_scale_ptr,\n const size_t scale_stride_h, const size_t scale_stride_w,\n const float *global_scale_ptr, const size_t h, const size_t w,\n const size_t start_offset, const size_t len) {\n constexpr int kNumOutputElemsPerBank = 4;\n constexpr int kThreadsPerWarp = 32;\n constexpr int kLoopsPerRow = (kTileDim + kThreadsPerWarp - 1) / kThreadsPerWarp;\n constexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp;\n constexpr int kRowsPerWarp = (kTileDim + kNumWarps - 1) / kNumWarps;\n\n __shared__ float smem[kTileDim][kTileDim + kNumOutputElemsPerBank];\n\n const int tile_w = blockIdx.x;\n const int tile_h = blockIdx.y;\n const size_t shard_end = start_offset + len;\n const IType *input_minus_offset = input - start_offset;\n\n float global_encode_scale = global_scale_ptr[0];\n if (global_encode_scale <= 0.f) {\n global_encode_scale = 1.f;\n }\n const float global_decode_scale = 1.0f / global_encode_scale;\n\n float tile_decode_scale = decode_scale_ptr[tile_h * scale_stride_h + tile_w * scale_stride_w];\n tile_decode_scale = static_cast(static_cast(tile_decode_scale));\n constexpr float kFp32Max = 3.402823466e+38F;\n float tile_encode_val =\n (tile_decode_scale > 0.f) ? 1.0f / (tile_decode_scale * global_decode_scale) : kFp32Max;\n tile_encode_val = fminf(tile_encode_val, kFp32Max);\n const float2 scale_vec = make_float2(tile_encode_val, tile_encode_val);\n\n bool skip_store = true;\n for (int i = 0; i < kRowsPerWarp; ++i) {\n for (int j = 0; j < kLoopsPerRow; ++j) {\n const int h_in_smem = threadIdx.x / kThreadsPerWarp * kRowsPerWarp + i;\n const int w_in_smem = threadIdx.x % kThreadsPerWarp + kThreadsPerWarp * j;\n if (h_in_smem >= kTileDim || w_in_smem >= kTileDim) {\n continue;\n }\n const int h_in_input = tile_h * kTileDim + h_in_smem;\n const int w_in_input = tile_w * kTileDim + w_in_smem;\n const size_t idx_in_input = static_cast(h_in_input) * w + w_in_input;\n if (h_in_input < h && w_in_input < w && idx_in_input >= start_offset &&\n idx_in_input < shard_end) {\n smem[h_in_smem][w_in_smem] = static_cast(input_minus_offset[idx_in_input]);\n skip_store = false;\n }\n }\n }\n\n for (int delta = kThreadsPerWarp / 2; delta > 0; delta /= 2) {\n bool other = __shfl_down_sync(0xFFFFFFFF, skip_store, delta);\n skip_store = skip_store && other;\n }\n skip_store = __shfl_sync(0xFFFFFFFF, skip_store, 0);\n if (skip_store) {\n return;\n }\n\n for (int i = 0; i < kRowsPerWarp; ++i) {\n const int row_in_smem = threadIdx.x / kThreadsPerWarp * kRowsPerWarp + i;\n const int row_in_output = tile_h * kTileDim + row_in_smem;\n if (row_in_output >= h) {\n continue;\n }\n const int col_in_smem = threadIdx.x % kThreadsPerWarp * kNumOutputElemsPerBank;\n if (col_in_smem >= kTileDim) {\n continue;\n }\n const int col_in_output = tile_w * kTileDim + col_in_smem;\n\n float vals[kNumOutputElemsPerBank];\n bool mask[kNumOutputElemsPerBank];\n size_t elem_idx[kNumOutputElemsPerBank];\n bool any_valid = false;\n\n for (int j = 0; j < kNumOutputElemsPerBank; ++j) {\n const int col = col_in_output + j;\n const bool in_width = col < w;\n const size_t idx = static_cast(row_in_output) * w + col;\n elem_idx[j] = idx;\n const bool in_shard = in_width && idx >= start_offset && idx < shard_end;\n mask[j] = in_shard;\n const bool in_tile = (col_in_smem + j) < kTileDim;\n const float tile_val = in_tile ? smem[row_in_smem][col_in_smem + j] : 0.0f;\n vals[j] = in_shard ? tile_val : 0.0f;\n any_valid |= in_shard;\n }\n\n if (!any_valid) {\n continue;\n }\n\n const float2 in01 = make_float2(vals[0], vals[1]);\n const float2 in23 = make_float2(vals[2], vals[3]);\n const auto packed =\n transformer_engine::ptx::mul_cvt_fp32_to_fp4_4x(in01, in23, scale_vec, 0);\n const uint16_t packed_bits = reinterpret_cast(packed);\n\n for (int pair = 0; pair < 2; ++pair) {\n const int first = pair * 2;\n const int second = first + 1;\n if (!mask[first] && !mask[second]) {\n continue;\n }\n const size_t ref_idx = mask[first] ? elem_idx[first] : elem_idx[second];\n const size_t byte_idx = (ref_idx - start_offset) >> 1;\n uint8_t byte = output[byte_idx];\n\n if (mask[first]) {\n const uint8_t nibble = static_cast((packed_bits >> (4 * first)) & 0xF);\n if ((elem_idx[first] & 1u) == 0) {\n byte = static_cast((byte & 0xF0u) | nibble);\n } else {\n byte = static_cast((byte & 0x0Fu) | (nibble << 4));\n }\n }\n\n if (mask[second]) {\n const uint8_t nibble = static_cast((packed_bits >> (4 * second)) & 0xF);\n if ((elem_idx[second] & 1u) == 0) {\n byte = static_cast((byte & 0xF0u) | nibble);\n } else {\n byte = static_cast((byte & 0x0Fu) | (nibble << 4));\n }\n }\n\n output[byte_idx] = byte;\n }\n }\n}\n\nvoid nvfp4_2d_compute_partial_amax(const Tensor inp, Tensor amax, size_t h, size_t w,\n size_t amax_stride_h, size_t amax_stride_w, size_t start_offset,\n size_t block_len, cudaStream_t stream) {\n NVTE_CHECK(block_len == 16, \"NVFP4 2D supports 16x16 tiles only (block_len = 16).\");\n\n size_t len = inp.numel();\n\n assert(h > 0 && w > 0);\n assert(start_offset < h * w);\n assert(start_offset + len <= h * w);\n\n size_t blocks_x = (w + kTileDim - 1) / kTileDim;\n size_t blocks_y = (h + kTileDim - 1) / kTileDim;\n assert(blocks_x <= std::numeric_limits::max());\n assert(blocks_y <= std::numeric_limits::max());\n dim3 grid(blocks_x, blocks_y);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n inp.dtype(), inp_dtype,\n nvfp4_2d_compute_partial_amax_kernel<<>>(\n reinterpret_cast(inp.data.dptr),\n reinterpret_cast(amax.data.dptr), amax_stride_h, amax_stride_w, h, w,\n start_offset, len);)\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid nvfp4_2d_partial_cast(const Tensor inp, Tensor out, const Tensor scale,\n const Tensor global_scale, size_t h, size_t w, size_t scale_stride_h,\n size_t scale_stride_w, size_t start_offset, size_t block_len,\n cudaStream_t stream) {\n NVTE_CHECK(block_len == 16, \"NVFP4 2D supports 16x16 tiles only (block_len = 16).\");\n NVTE_CHECK(out.dtype() == DType::kByte, \"NVFP4 rowwise data must be uint8.\");\n\n size_t len = inp.numel();\n\n assert(h > 0 && w > 0);\n assert(start_offset < h * w);\n assert(start_offset + len <= h * w);\n\n size_t blocks_x = (w + kTileDim - 1) / kTileDim;\n size_t blocks_y = (h + kTileDim - 1) / kTileDim;\n assert(blocks_x <= std::numeric_limits::max());\n assert(blocks_y <= std::numeric_limits::max());\n dim3 grid(blocks_x, blocks_y);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n inp.dtype(), inp_dtype,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n w % kTileDim == 0, kWidthAligned,\n nvfp4_2d_partial_cast_kernel\n <<>>(\n reinterpret_cast(inp.data.dptr),\n reinterpret_cast(out.data.dptr),\n reinterpret_cast(scale.data.dptr), scale_stride_h, scale_stride_w,\n reinterpret_cast(global_scale.data.dptr), h, w, start_offset, len);))\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n/*\n * ---------------------------------------------------------------------------\n * NVFP4 TRANSPOSE KERNEL\n *\n * Unlike FP8, NVFP4 packs two 4-bit values into each byte. A simple byte-wise\n * transpose doesn't work because the packing changes:\n * - Before transpose: elements [m, 2c] and [m, 2c+1] share a byte\n * - After transpose: elements [k, 2*m_packed] and [k, 2*m_packed+1] share a byte\n * which were originally [2*m_packed, k] and [2*m_packed+1, k]\n * ---------------------------------------------------------------------------\n */\n\n// Vectorized transpose kernel parameters\nconstexpr int TRANSPOSE_TILE_DIM = 64; // Logical FP4 elements per tile dimension\n// constexpr int TRANSPOSE_TILE_PACKED = 32; // TILE_DIM / 2 bytes\nconstexpr int TRANSPOSE_BLOCK_SIZE = 256; // threads per block\n\n// Shared memory: store unpacked 4-bit values as bytes for easy transpose\n// Size: TILE_DIM x (TILE_DIM + 4) to avoid bank conflicts\nconstexpr int TRANSPOSE_SHMEM_STRIDE = TRANSPOSE_TILE_DIM + 4;\n\n/*\n * Vectorized transpose kernel with uint2 loads/stores (256 threads)\n * Tile: 64x64 logical FP4 = 64x32 packed bytes\n */\n__global__ void __launch_bounds__(TRANSPOSE_BLOCK_SIZE)\n nvfp4_transpose_kernel(const uint8_t *__restrict__ input, uint8_t *__restrict__ output,\n const size_t M, const size_t K) {\n const size_t K_packed = K / 2;\n const size_t M_packed = M / 2;\n\n const size_t tile_m_start = blockIdx.x * TRANSPOSE_TILE_DIM;\n const size_t tile_k_start = blockIdx.y * TRANSPOSE_TILE_DIM;\n\n __shared__ uint8_t shmem[TRANSPOSE_TILE_DIM][TRANSPOSE_SHMEM_STRIDE];\n\n const int tid = threadIdx.x;\n\n // Phase 1: Load input tile with VECTORIZED uint2 reads\n // 256 threads, each loads 8 bytes (uint2) = 2048 bytes total\n // Input tile: [64 rows, 32 cols] = 2048 bytes\n {\n const int thread_row = tid / 4; // 64 rows, 4 threads per row\n const int thread_col = (tid % 4) * 8; // 4 x 8 = 32 bytes per row\n\n const size_t global_m = tile_m_start + thread_row;\n const size_t global_k_packed_base = tile_k_start / 2 + thread_col;\n\n // Load 8 bytes as uint2\n uint2 loaded = make_uint2(0, 0);\n if (global_m < M && global_k_packed_base + 7 < K_packed) {\n loaded = *reinterpret_cast(&input[global_m * K_packed + global_k_packed_base]);\n } else if (global_m < M) {\n // Boundary: scalar loads\n uint8_t *bytes = reinterpret_cast(&loaded);\n#pragma unroll\n for (int b = 0; b < 8; ++b) {\n size_t col = global_k_packed_base + b;\n bytes[b] = (col < K_packed) ? input[global_m * K_packed + col] : 0;\n }\n }\n\n // Unpack 8 bytes -> 16 nibbles and store to shared memory\n const uint8_t *bytes = reinterpret_cast(&loaded);\n#pragma unroll\n for (int b = 0; b < 8; ++b) {\n const int k0 = thread_col * 2 + b * 2;\n const int k1 = k0 + 1;\n shmem[thread_row][k0] = bytes[b] & 0x0F;\n shmem[thread_row][k1] = (bytes[b] >> 4) & 0x0F;\n }\n }\n\n __syncthreads();\n\n // Phase 2: Write output with VECTORIZED uint2 stores\n // Output tile: [64 rows, 32 cols] = 2048 bytes\n {\n const int thread_row = tid / 4; // output K dimension [0, 64)\n const int thread_col_base = (tid % 4) * 8; // output M_packed [0, 32) in steps of 8\n\n const size_t global_k = tile_k_start + thread_row;\n const size_t global_m_packed_base = tile_m_start / 2 + thread_col_base;\n\n if (global_k >= K) return;\n\n // Build 8 output bytes in registers\n uint8_t out_bytes[8];\n\n#pragma unroll\n for (int b = 0; b < 8; ++b) {\n const int out_m_packed = thread_col_base + b;\n\n if (global_m_packed_base + b >= M_packed) {\n out_bytes[b] = 0;\n continue;\n }\n\n // Two M positions that pack into this output byte\n const int m0 = out_m_packed * 2;\n const int m1 = out_m_packed * 2 + 1;\n const int k = thread_row;\n\n // Read from shared memory (transposed access)\n const uint8_t val0 = shmem[m0][k];\n const uint8_t val1 = shmem[m1][k];\n\n out_bytes[b] = val0 | (val1 << 4);\n }\n\n // Vectorized store as uint2\n if (global_m_packed_base + 7 < M_packed) {\n *reinterpret_cast(&output[global_k * M_packed + global_m_packed_base]) =\n *reinterpret_cast(out_bytes);\n } else {\n // Boundary: scalar stores\n for (int b = 0; b < 8 && global_m_packed_base + b < M_packed; ++b) {\n output[global_k * M_packed + global_m_packed_base + b] = out_bytes[b];\n }\n }\n }\n}\n\nvoid nvfp4_transpose(const Tensor input, Tensor output, cudaStream_t stream) {\n // Input has logical shape [M, K], stored as [M, K/2] bytes\n // Output has logical shape [K, M], stored as [K, M/2] bytes\n\n NVTE_CHECK(input.dtype() == DType::kByte, \"NVFP4 transpose input must be uint8.\");\n NVTE_CHECK(output.dtype() == DType::kByte, \"NVFP4 transpose output must be uint8.\");\n\n // Get dimensions from packed storage\n // input.shape() = [M, K/2], so M = shape[0], K = shape[1] * 2\n const auto in_shape = input.shape();\n NVTE_CHECK(in_shape.size() == 2, \"NVFP4 transpose expects 2D input (packed), got \",\n in_shape.size(), \"D.\");\n const size_t M = in_shape[0];\n const size_t K_packed = in_shape[1];\n const size_t K = K_packed * 2;\n\n // Output should be [K, M/2]\n const size_t M_packed = M / 2;\n NVTE_CHECK(M % 2 == 0, \"NVFP4 transpose requires M (\", M, \") to be even.\");\n\n const auto out_shape = output.shape();\n NVTE_CHECK(out_shape.size() == 2, \"NVFP4 transpose expects 2D output.\");\n NVTE_CHECK(out_shape[0] == K && out_shape[1] == M_packed,\n \"NVFP4 transpose output shape mismatch. Expected [\", K, \", \", M_packed, \"], got [\",\n out_shape[0], \", \", out_shape[1], \"].\");\n\n if (M == 0 || K == 0) return;\n\n // Use vectorized kernel (faster than TMA for pure transpose)\n // 128x128 tiles with 512 threads and uint4 vectorized access\n dim3 block(TRANSPOSE_BLOCK_SIZE);\n dim3 grid((M + TRANSPOSE_TILE_DIM - 1) / TRANSPOSE_TILE_DIM,\n (K + TRANSPOSE_TILE_DIM - 1) / TRANSPOSE_TILE_DIM);\n\n nvfp4_transpose_kernel<<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(output.data.dptr), M, K);\n\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n/*\n * ---------------------------------------------------------------------------\n * NVFP4 SCALE TRANSPOSE KERNEL\n *\n * Transposes tile-level scales from rowwise to columnwise format.\n * Scale values are stored as E4M3 (fp8) in uint8 tensors.\n *\n * Input (rowwise_scale_inv): [M_padded, K_tiles] where scales are stored\n * at every 16th row (i.e., row 0, 16, 32, ... contain the actual scales,\n * and each row i within a tile block has the same scale as row (i // 16) * 16).\n *\n * Output (columnwise_scale_inv): [K_padded, M_tiles] where scales are\n * repeated 16 times per tile row.\n *\n * Mapping:\n * output[k_tile * 16 + i, m_tile] = input[m_tile * 16, k_tile]\n * for i in [0, 16) and valid (k_tile, m_tile) indices.\n * ---------------------------------------------------------------------------\n */\n__global__ void nvfp4_scale_transpose_kernel(\n const uint8_t *__restrict__ input, // [M_padded, K_tiles], E4M3 stored as uint8\n uint8_t *__restrict__ output, // [K_padded, M_tiles], E4M3 stored as uint8\n const size_t M_tiles, // Number of M tiles\n const size_t K_tiles, // Number of K tiles\n const size_t input_stride, // K_tiles (input row stride)\n const size_t output_stride, // M_tiles (output row stride)\n const size_t K_padded // Output height\n) {\n // Each thread handles one output element\n const size_t out_row = blockIdx.y * blockDim.y + threadIdx.y;\n const size_t out_col = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (out_row >= K_padded || out_col >= M_tiles) return;\n\n // Determine which tile row this belongs to\n const size_t k_tile = out_row / kTileDim;\n\n // Read from input: row = m_tile * 16 (first row of the tile), col = k_tile\n // m_tile = out_col\n if (k_tile < K_tiles) {\n const size_t in_row = out_col * kTileDim; // m_tile * 16\n const uint8_t scale = input[in_row * input_stride + k_tile];\n output[out_row * output_stride + out_col] = scale;\n } else {\n output[out_row * output_stride + out_col] = 0;\n }\n}\n\nvoid nvfp4_scale_transpose(const Tensor input, Tensor output, size_t M_tiles, size_t K_tiles,\n cudaStream_t stream) {\n NVTE_CHECK(input.dtype() == DType::kByte, \"NVFP4 scale transpose input must be uint8 (E4M3).\");\n NVTE_CHECK(output.dtype() == DType::kByte, \"NVFP4 scale transpose output must be uint8 (E4M3).\");\n\n const auto in_shape = input.shape();\n const auto out_shape = output.shape();\n NVTE_CHECK(in_shape.size() == 2, \"NVFP4 scale transpose expects 2D input.\");\n NVTE_CHECK(out_shape.size() == 2, \"NVFP4 scale transpose expects 2D output.\");\n\n const size_t input_stride = in_shape[1]; // K_tiles\n const size_t output_stride = out_shape[1]; // M_tiles\n const size_t K_padded = out_shape[0];\n\n if (M_tiles == 0 || K_tiles == 0 || K_padded == 0) return;\n\n constexpr int kBlockDim = 16;\n dim3 block(kBlockDim, kBlockDim);\n dim3 grid((M_tiles + kBlockDim - 1) / kBlockDim, (K_padded + kBlockDim - 1) / kBlockDim);\n\n nvfp4_scale_transpose_kernel<<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(output.data.dptr), M_tiles, K_tiles, input_stride, output_stride,\n K_padded);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n/*\n * ---------------------------------------------------------------------------\n * NVFP4 SCALE EXPANSION KERNEL\n *\n * Expands tile-level scales to row-level scales and converts to FP8 E4M3, used in partial cast.\n *\n * Input (per_block_decode_scale): [tile_rows, tile_cols] in float32\n * Output (target_scale): [rows_padded, tile_cols] in uint8 (E4M3)\n *\n * Each tile row's scale is repeated block_len times in the output.\n * ---------------------------------------------------------------------------\n */\n__global__ void nvfp4_expand_scale_to_fp8_kernel(\n const float *__restrict__ input, // [tile_rows, tile_cols]\n uint8_t *__restrict__ output, // [rows_padded, tile_cols]\n const size_t tile_rows, const size_t tile_cols, const size_t rows_padded,\n const size_t block_len) {\n const size_t out_row = blockIdx.y * blockDim.y + threadIdx.y;\n const size_t out_col = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (out_row >= rows_padded || out_col >= tile_cols) return;\n\n // Determine which tile row this output row belongs to\n const size_t tile_row = out_row / block_len;\n\n float scale_val = 0.0f;\n if (tile_row < tile_rows) {\n scale_val = input[tile_row * tile_cols + out_col];\n }\n\n // Convert float32 to FP8 E4M3\n // Clamp to FP8 E4M3 range and convert\n fp8e4m3 fp8_val = static_cast(scale_val);\n output[out_row * tile_cols + out_col] = reinterpret_cast(fp8_val);\n}\n\nvoid nvfp4_expand_scale_to_fp8(const Tensor input, Tensor output, size_t tile_rows,\n size_t tile_cols, size_t rows_padded, size_t block_len,\n cudaStream_t stream) {\n NVTE_CHECK(input.dtype() == DType::kFloat32, \"Scale input must be float32.\");\n NVTE_CHECK(output.dtype() == DType::kByte, \"Scale output must be uint8 (E4M3).\");\n\n if (tile_rows == 0 || tile_cols == 0 || rows_padded == 0) return;\n\n constexpr int kBlockDim = 16;\n dim3 block(kBlockDim, kBlockDim);\n dim3 grid((tile_cols + kBlockDim - 1) / kBlockDim, (rows_padded + kBlockDim - 1) / kBlockDim);\n\n nvfp4_expand_scale_to_fp8_kernel<<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(output.data.dptr), tile_rows, tile_cols, rows_padded, block_len);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n/*\n * ---------------------------------------------------------------------------\n * NVFP4 COMPUTE PER-BLOCK DECODE SCALE KERNEL\n *\n * Computes per-block decode scale from block amax and global amax:\n * global_scale = (fp8_max * fp4_max) / global_amax = 2688 / global_amax\n * per_block_decode_scale = block_amax * (global_scale * (1 / fp4_max))\n * = block_amax * 448 / global_amax\n *\n * This matches the CUDA device function compute_decoding_scaling_factor() in core_nvfp4.cuh\n *\n * Input (block_amax): [tile_rows, tile_cols] in float32\n * Input (global_amax): scalar float32 (per-tensor amax after all-reduce)\n * Output (scale): [tile_rows, tile_cols] in float32\n * Output (global_scale_out): scalar float32 (the computed global encode scale)\n * ---------------------------------------------------------------------------\n */\n__global__ void nvfp4_compute_per_block_scale_kernel(\n const float *__restrict__ block_amax, // [tile_rows, tile_cols]\n float *__restrict__ scale, // [tile_rows, tile_cols]\n const float *__restrict__ global_amax_ptr, // Pointer to single float value (avoids D2H)\n const size_t numel) {\n const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= numel) return;\n\n constexpr float fp4_max = 6.0f;\n constexpr float fp8_max = 448.0f;\n constexpr float flt_max = 3.402823466e+38f;\n constexpr float tiny = 1.17549435e-38f; // FLT_MIN\n\n // Read global_amax from device memory (avoids D2H transfer)\n float global_amax = *global_amax_ptr;\n\n // Compute global encode scale: S_enc = (fp8_max * fp4_max) / global_amax\n float safe_global_amax = fmaxf(global_amax, tiny);\n float global_scale =\n (global_amax > 0.0f) ? fminf((fp8_max * fp4_max) / safe_global_amax, flt_max) : 1.0f;\n\n // Compute per-block decode scale: S_dec_b = block_amax * (S_enc * (1 / fp4_max))\n float amax_val = block_amax[idx];\n constexpr float fp4_max_inv = 1.0f / fp4_max;\n const float global_scale_multiplier = global_scale * fp4_max_inv;\n float result = fminf(amax_val * global_scale_multiplier, flt_max);\n scale[idx] = result;\n}\n\n// Simple kernel to compute global encode scale from global amax\n__global__ void nvfp4_compute_global_scale_kernel(\n const float *__restrict__ global_amax, // [num_params]\n float *__restrict__ global_scale, // [num_params]\n const size_t num_params) {\n const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= num_params) return;\n\n constexpr float fp4_max = 6.0f;\n constexpr float fp8_max = 448.0f;\n constexpr float flt_max = 3.402823466e+38f;\n constexpr float tiny = 1.17549435e-38f; // FLT_MIN\n\n float amax = global_amax[idx];\n float safe_amax = fmaxf(amax, tiny);\n float scale = (amax > 0.0f) ? fminf((fp8_max * fp4_max) / safe_amax, flt_max) : 1.0f;\n global_scale[idx] = scale;\n}\n\nvoid nvfp4_compute_per_block_scale(const Tensor block_amax, Tensor scale, const Tensor global_amax,\n cudaStream_t stream) {\n NVTE_CHECK(block_amax.dtype() == DType::kFloat32, \"Block amax must be float32.\");\n NVTE_CHECK(scale.dtype() == DType::kFloat32, \"Scale must be float32.\");\n NVTE_CHECK(global_amax.dtype() == DType::kFloat32, \"Global amax must be float32.\");\n NVTE_CHECK(global_amax.numel() == 1, \"Global amax must be a single element tensor.\");\n\n size_t numel = block_amax.numel();\n if (numel == 0) return;\n\n constexpr int kBlockSize = 256;\n int grid_size = (numel + kBlockSize - 1) / kBlockSize;\n\n nvfp4_compute_per_block_scale_kernel<<>>(\n reinterpret_cast(block_amax.data.dptr),\n reinterpret_cast(scale.data.dptr),\n reinterpret_cast(global_amax.data.dptr), numel);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid nvfp4_compute_global_scale(const Tensor global_amax, Tensor global_scale,\n cudaStream_t stream) {\n NVTE_CHECK(global_amax.dtype() == DType::kFloat32, \"Global amax must be float32.\");\n NVTE_CHECK(global_scale.dtype() == DType::kFloat32, \"Global scale must be float32.\");\n\n size_t num_params = global_amax.numel();\n if (num_params == 0) return;\n\n constexpr int kBlockSize = 256;\n int grid_size = (num_params + kBlockSize - 1) / kBlockSize;\n\n nvfp4_compute_global_scale_kernel<<>>(\n reinterpret_cast(global_amax.data.dptr),\n reinterpret_cast(global_scale.data.dptr), num_params);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n/*\n * ---------------------------------------------------------------------------\n * FUSED NVFP4 SCALE COMPUTATION KERNEL\n *\n * Fuses three operations into one kernel:\n * 1. nvfp4_compute_per_block_scale: compute tile-level decode scales from block amax\n * 2. target_amax.copy_: copy global amax to target tensor\n * 3. nvfp4_expand_scale_to_fp8: expand to row-level and convert to FP8 E4M3\n *\n * Input (block_amax): [tile_rows, tile_cols] float32\n * Input (global_amax): [1] float32\n * Output (per_block_scale): [tile_rows, tile_cols] float32 (intermediate, for partial_cast)\n * Output (target_scale): [rows_padded, tile_cols] uint8 (E4M3)\n * Output (target_amax): [1] float32 (copy of global_amax)\n *\n * Saves 2 kernel launches per parameter (eliminates nvfp4_compute_per_block_scale and\n * nvfp4_expand_scale_to_fp8 as separate calls, plus the amax copy).\n * ---------------------------------------------------------------------------\n */\n__global__ void nvfp4_fused_scale_kernel(\n const float *__restrict__ block_amax, // [tile_rows, tile_cols]\n const float *__restrict__ global_amax, // [1]\n float *__restrict__ per_block_scale, // [tile_rows, tile_cols] - for partial_cast\n uint8_t *__restrict__ target_scale, // [rows_padded, tile_cols]\n float *__restrict__ target_amax, // [1]\n const size_t tile_rows, const size_t tile_cols, const size_t rows_padded,\n const size_t block_len) {\n const size_t out_row = blockIdx.y * blockDim.y + threadIdx.y;\n const size_t out_col = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Read global amax once per thread (broadcast)\n const float g_amax = *global_amax;\n\n // Thread (0,0) copies global_amax to target_amax\n if (out_row == 0 && out_col == 0) {\n *target_amax = g_amax;\n }\n\n if (out_row >= rows_padded || out_col >= tile_cols) return;\n\n // Determine which tile row this output row belongs to\n const size_t tile_row = out_row / block_len;\n\n // Compute the scale value\n constexpr float fp4_max = 6.0f;\n constexpr float fp8_max = 448.0f;\n constexpr float flt_max = 3.402823466e+38f;\n constexpr float tiny = 1.17549435e-38f;\n\n float scale_val = 0.0f;\n if (tile_row < tile_rows) {\n float safe_global_amax = fmaxf(g_amax, tiny);\n float global_scale =\n (g_amax > 0.0f) ? fminf((fp8_max * fp4_max) / safe_global_amax, flt_max) : 1.0f;\n constexpr float fp4_max_inv = 1.0f / fp4_max;\n const float global_scale_multiplier = global_scale * fp4_max_inv;\n\n // Read block amax and compute per-block decode scale\n float amax_val = block_amax[tile_row * tile_cols + out_col];\n scale_val = fminf(amax_val * global_scale_multiplier, flt_max);\n\n // Write per-block scale (only once per tile, when out_row % block_len == 0)\n if (out_row % block_len == 0) {\n per_block_scale[tile_row * tile_cols + out_col] = scale_val;\n }\n }\n\n // Convert float32 to FP8 E4M3 and write expanded scale\n fp8e4m3 fp8_val = static_cast(scale_val);\n target_scale[out_row * tile_cols + out_col] = reinterpret_cast(fp8_val);\n}\n\nvoid nvfp4_fused_scale(const Tensor block_amax, const Tensor global_amax, Tensor per_block_scale,\n Tensor target_scale, Tensor target_amax, size_t tile_rows, size_t tile_cols,\n size_t rows_padded, size_t block_len, cudaStream_t stream) {\n NVTE_CHECK(block_amax.dtype() == DType::kFloat32, \"Block amax must be float32.\");\n NVTE_CHECK(global_amax.dtype() == DType::kFloat32, \"Global amax must be float32.\");\n NVTE_CHECK(per_block_scale.dtype() == DType::kFloat32, \"Per-block scale must be float32.\");\n NVTE_CHECK(target_scale.dtype() == DType::kByte, \"Target scale must be uint8 (E4M3).\");\n NVTE_CHECK(target_amax.dtype() == DType::kFloat32, \"Target amax must be float32.\");\n NVTE_CHECK(global_amax.numel() == 1, \"Global amax must be a single element tensor.\");\n NVTE_CHECK(target_amax.numel() == 1, \"Target amax must be a single element tensor.\");\n\n if (tile_rows == 0 || tile_cols == 0 || rows_padded == 0) return;\n\n constexpr int kBlockDim = 16;\n dim3 block(kBlockDim, kBlockDim);\n dim3 grid((tile_cols + kBlockDim - 1) / kBlockDim, (rows_padded + kBlockDim - 1) / kBlockDim);\n\n nvfp4_fused_scale_kernel<<>>(\n reinterpret_cast(block_amax.data.dptr),\n reinterpret_cast(global_amax.data.dptr),\n reinterpret_cast(per_block_scale.data.dptr),\n reinterpret_cast(target_scale.data.dptr),\n reinterpret_cast(target_amax.data.dptr), tile_rows, tile_cols, rows_padded,\n block_len);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n#endif // FP4_TYPE_SUPPORTED\n} // namespace nvfp4_recipe\n} // namespace transformer_engine\n\nvoid nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, size_t tile_rows,\n size_t tile_cols, size_t rows_padded, size_t block_len,\n cudaStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_expand_scale_to_fp8);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_expand_scale_to_fp8(*convertNVTETensorCheck(input),\n *convertNVTETensorCheck(output), tile_rows, tile_cols,\n rows_padded, block_len, stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor scale,\n const NVTETensor global_amax, cudaStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_compute_per_block_scale);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_compute_per_block_scale(*convertNVTETensorCheck(block_amax),\n *convertNVTETensorCheck(scale),\n *convertNVTETensorCheck(global_amax), stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_compute_global_scale(const NVTETensor global_amax, NVTETensor global_scale,\n cudaStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_compute_global_scale);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_compute_global_scale(*convertNVTETensorCheck(global_amax),\n *convertNVTETensorCheck(global_scale), stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_scale_transpose(const NVTETensor input, NVTETensor output, size_t M_tiles,\n size_t K_tiles, cudaStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_scale_transpose);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_scale_transpose(*convertNVTETensorCheck(input),\n *convertNVTETensorCheck(output), M_tiles, K_tiles, stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_data_transpose(const NVTETensor input, NVTETensor output, cudaStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_data_transpose);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_transpose(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output),\n stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, size_t h, size_t w,\n size_t amax_stride_h, size_t amax_stride_w,\n size_t start_offset, size_t block_len,\n cudaStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_2d_compute_partial_amax);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_2d_compute_partial_amax(*convertNVTETensorCheck(inp),\n *convertNVTETensorCheck(amax), h, w, amax_stride_h,\n amax_stride_w, start_offset, block_len, stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale,\n const NVTETensor global_scale, size_t h, size_t w,\n size_t scale_stride_h, size_t scale_stride_w, size_t start_offset,\n size_t block_len, cudaStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_2d_partial_cast);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_2d_partial_cast(*convertNVTETensorCheck(inp), *convertNVTETensorCheck(out),\n *convertNVTETensorCheck(scale),\n *convertNVTETensorCheck(global_scale), h, w, scale_stride_h,\n scale_stride_w, start_offset, block_len, stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_rowwise_amax_A,\n const NVTETensor inpB, const bool use_rowwise_amax_B,\n float alpha_in, NVTETensor alpha_out,\n cudaStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_compute_per_tensor_scale);\n using namespace transformer_engine;\n\n auto *tA = convertNVTETensor(inpA);\n auto *tB = convertNVTETensor(inpB);\n auto *tOut = convertNVTETensor(alpha_out);\n\n void *amax_A_ptr = use_rowwise_amax_A ? tA->amax.dptr : tA->columnwise_amax.dptr;\n void *amax_B_ptr = use_rowwise_amax_B ? tB->amax.dptr : tB->columnwise_amax.dptr;\n void *alpha_ptr = tOut->data.dptr;\n\n // check for not null pointers\n NVTE_CHECK(amax_A_ptr != nullptr, \"amax_A_ptr is null\");\n NVTE_CHECK(amax_B_ptr != nullptr, \"amax_B_ptr is null\");\n NVTE_CHECK(alpha_ptr != nullptr, \"alpha_ptr is null\");\n\n nvfp4_recipe::compute_nvfp4_per_tensor_scale_kernel<<<1, 1, 0, stream>>>(\n alpha_in, reinterpret_cast(amax_A_ptr),\n reinterpret_cast(amax_B_ptr), reinterpret_cast(alpha_ptr));\n NVTE_CHECK_CUDA(cudaGetLastError());\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_fused_scale(const NVTETensor block_amax, const NVTETensor global_amax,\n NVTETensor per_block_scale, NVTETensor target_scale,\n NVTETensor target_amax, size_t tile_rows, size_t tile_cols,\n size_t rows_padded, size_t block_len, cudaStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_fused_scale);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_fused_scale(\n *convertNVTETensorCheck(block_amax), *convertNVTETensorCheck(global_amax),\n *convertNVTETensorCheck(per_block_scale), *convertNVTETensorCheck(target_scale),\n *convertNVTETensorCheck(target_amax), tile_rows, tile_cols, rows_padded, block_len, stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n\n#include \n#include \n\n#include \"../common.h\"\n#include \"../util/ptx.cuh\"\n#include \"../utils.cuh\"\n\nnamespace transformer_engine {\nnamespace nvfp4_recipe {\n\n#if FP4_TYPE_SUPPORTED\n/*\n * ---------------------------------------------------------------------------\n * NVFP4 2D PARTIAL-SHARD KERNEL DESIGN\n *\n * These kernels mirror the FP8 block-scaling helpers but operate on shard-local\n * slices and nibble-packed FP4 rowwise buffers. One CUDA block covers a logical\n * 16x16 tile (grid = ceil(W/16) x ceil(H/16), blockDim = 256 threads).\n *\n * 1) Partial Amax (`nvfp4_2d_compute_partial_amax_kernel`)\n * - Warps sweep the tile using nested loops, accumulating local maxima only\n * for elements in [start_offset, start_offset + len).\n * - Shared memory reduces the 8 warp maxima; the block writes a float into\n * `amax_ptr[tile_row * stride_h + tile_col * stride_w]`.\n *\n * Tile/warp mapping (each '#' = elements visited by that warp):\n *\n * +------------------+\n * |########..........| Warp 0\n * |########..........| Warp 1\n * | ... |\n * |########..........| Warp 7\n * +------------------+\n *\n * 2) Partial Cast (`nvfp4_2d_partial_cast_kernel`)\n * - Stage the tile into shared memory (same pattern as FP8).\n * - For each 4-value group, build float2 pairs and call\n * `ptx::mul_cvt_fp32_to_fp4_4x`, producing packed FP4 nibbles.\n * - Compute a shard-local byte index and update only the owned nibble(s)\n * using read-modify-write:\n *\n * packed_bits = [mw3 | mw2 | mw1 | mw0]\n * byte_idx = (ref_elem_idx - start_offset) >> 1\n * if elem_idx % 2 == 0: // low nibble\n * byte = (byte & 0xF0) | nibble\n * else: // high nibble\n * byte = (byte & 0x0F) | (nibble << 4)\n *\n * Thread coverage inside a tile:\n *\n * rows: 16 columns: 16\n * Warp 0 -> rows 0-1 lanes sweep cols 0..3, 4..7, ...\n * Warp 1 -> rows 2-3 (groups of 4 elements per thread)\n * ...\n * Warp 7 -> rows 14-15\n * ---------------------------------------------------------------------------\n */\n\n// constexpr float factor = 6.0 * 6.0 * 448.0 * 448.0;\nconstexpr float factor_inv = 1.0 / (6.0 * 6.0 * 448.0 * 448.0);\nconstexpr int kTileDim = 16;\nconstexpr int kThreadsPerBlock = 256;\n\n// Kernel to compute alpha *= amax_A * amax_B / factor\n__global__ void compute_nvfp4_per_tensor_scale_kernel(float alpha_in, const float *amax_A,\n const float *amax_B, float *alpha_out) {\n // factor is defined in the enclosing namespace\n *alpha_out = alpha_in * (*amax_A) * (*amax_B) * factor_inv;\n}\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock)\n nvfp4_2d_compute_partial_amax_kernel(const IType *input, float *amax_ptr,\n const size_t amax_stride_h, const size_t amax_stride_w,\n const size_t h, const size_t w, const size_t start_offset,\n const size_t len) {\n constexpr int kThreadsPerWarp = 32;\n constexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp;\n static_assert(kTileDim * kTileDim == kThreadsPerBlock);\n\n const size_t tile_col = blockIdx.x;\n const size_t tile_row = blockIdx.y;\n const size_t end_offset = start_offset + len;\n const IType *input_minus_offset = input - start_offset;\n\n __shared__ float smem[kNumWarps];\n float amax = 0.0f;\n\n size_t r = tile_row * kTileDim + threadIdx.x / kTileDim;\n size_t c = tile_col * kTileDim + threadIdx.x % kTileDim;\n size_t idx = r * w + c;\n if (r < h && c < w && idx >= start_offset && idx < end_offset) {\n amax = fabs(static_cast(input_minus_offset[idx]));\n }\n\n for (int delta = kThreadsPerWarp / 2; delta > 0; delta /= 2) {\n float other_amax = __shfl_down_sync(0xFFFFFFFF, amax, delta);\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n\n if (threadIdx.x % kThreadsPerWarp == 0) {\n smem[threadIdx.x / kThreadsPerWarp] = amax;\n }\n\n __syncthreads();\n\n if (threadIdx.x == 0) {\n for (int i = 0; i < kNumWarps; ++i) {\n float other_amax = smem[i];\n __builtin_assume(amax >= 0);\n __builtin_assume(other_amax >= 0);\n amax = fmaxf(amax, other_amax);\n }\n amax_ptr[tile_row * amax_stride_h + tile_col * amax_stride_w] = amax;\n }\n}\n\ntemplate \n__global__ void __launch_bounds__(kThreadsPerBlock)\n nvfp4_2d_partial_cast_kernel(const IType *input, uint8_t *output, const float *decode_scale_ptr,\n const size_t scale_stride_h, const size_t scale_stride_w,\n const float *global_scale_ptr, const size_t h, const size_t w,\n const size_t start_offset, const size_t len) {\n constexpr int kNumOutputElemsPerBank = 4;\n constexpr int kThreadsPerWarp = 32;\n constexpr int kLoopsPerRow = (kTileDim + kThreadsPerWarp - 1) / kThreadsPerWarp;\n constexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp;\n constexpr int kRowsPerWarp = (kTileDim + kNumWarps - 1) / kNumWarps;\n\n __shared__ float smem[kTileDim][kTileDim + kNumOutputElemsPerBank];\n\n const int tile_w = blockIdx.x;\n const int tile_h = blockIdx.y;\n const size_t shard_end = start_offset + len;\n const IType *input_minus_offset = input - start_offset;\n\n float global_encode_scale = global_scale_ptr[0];\n if (global_encode_scale <= 0.f) {\n global_encode_scale = 1.f;\n }\n const float global_decode_scale = 1.0f / global_encode_scale;\n\n float tile_decode_scale = decode_scale_ptr[tile_h * scale_stride_h + tile_w * scale_stride_w];\n tile_decode_scale = static_cast(static_cast(tile_decode_scale));\n constexpr float kFp32Max = 3.402823466e+38F;\n float tile_encode_val =\n (tile_decode_scale > 0.f) ? 1.0f / (tile_decode_scale * global_decode_scale) : kFp32Max;\n tile_encode_val = fminf(tile_encode_val, kFp32Max);\n const float2 scale_vec = make_float2(tile_encode_val, tile_encode_val);\n\n bool skip_store = true;\n for (int i = 0; i < kRowsPerWarp; ++i) {\n for (int j = 0; j < kLoopsPerRow; ++j) {\n const int h_in_smem = threadIdx.x / kThreadsPerWarp * kRowsPerWarp + i;\n const int w_in_smem = threadIdx.x % kThreadsPerWarp + kThreadsPerWarp * j;\n if (h_in_smem >= kTileDim || w_in_smem >= kTileDim) {\n continue;\n }\n const int h_in_input = tile_h * kTileDim + h_in_smem;\n const int w_in_input = tile_w * kTileDim + w_in_smem;\n const size_t idx_in_input = static_cast(h_in_input) * w + w_in_input;\n if (h_in_input < h && w_in_input < w && idx_in_input >= start_offset &&\n idx_in_input < shard_end) {\n smem[h_in_smem][w_in_smem] = static_cast(input_minus_offset[idx_in_input]);\n skip_store = false;\n }\n }\n }\n\n for (int delta = kThreadsPerWarp / 2; delta > 0; delta /= 2) {\n bool other = __shfl_down_sync(0xFFFFFFFF, skip_store, delta);\n skip_store = skip_store && other;\n }\n skip_store = __shfl_sync(0xFFFFFFFF, skip_store, 0);\n if (skip_store) {\n return;\n }\n\n for (int i = 0; i < kRowsPerWarp; ++i) {\n const int row_in_smem = threadIdx.x / kThreadsPerWarp * kRowsPerWarp + i;\n const int row_in_output = tile_h * kTileDim + row_in_smem;\n if (row_in_output >= h) {\n continue;\n }\n const int col_in_smem = threadIdx.x % kThreadsPerWarp * kNumOutputElemsPerBank;\n if (col_in_smem >= kTileDim) {\n continue;\n }\n const int col_in_output = tile_w * kTileDim + col_in_smem;\n\n float vals[kNumOutputElemsPerBank];\n bool mask[kNumOutputElemsPerBank];\n size_t elem_idx[kNumOutputElemsPerBank];\n bool any_valid = false;\n\n for (int j = 0; j < kNumOutputElemsPerBank; ++j) {\n const int col = col_in_output + j;\n const bool in_width = col < w;\n const size_t idx = static_cast(row_in_output) * w + col;\n elem_idx[j] = idx;\n const bool in_shard = in_width && idx >= start_offset && idx < shard_end;\n mask[j] = in_shard;\n const bool in_tile = (col_in_smem + j) < kTileDim;\n const float tile_val = in_tile ? smem[row_in_smem][col_in_smem + j] : 0.0f;\n vals[j] = in_shard ? tile_val : 0.0f;\n any_valid |= in_shard;\n }\n\n if (!any_valid) {\n continue;\n }\n\n const float2 in01 = make_float2(vals[0], vals[1]);\n const float2 in23 = make_float2(vals[2], vals[3]);\n const auto packed =\n transformer_engine::ptx::mul_cvt_fp32_to_fp4_4x(in01, in23, scale_vec, 0);\n const uint16_t packed_bits = reinterpret_cast(packed);\n\n for (int pair = 0; pair < 2; ++pair) {\n const int first = pair * 2;\n const int second = first + 1;\n if (!mask[first] && !mask[second]) {\n continue;\n }\n const size_t ref_idx = mask[first] ? elem_idx[first] : elem_idx[second];\n const size_t byte_idx = (ref_idx - start_offset) >> 1;\n uint8_t byte = output[byte_idx];\n\n if (mask[first]) {\n const uint8_t nibble = static_cast((packed_bits >> (4 * first)) & 0xF);\n if ((elem_idx[first] & 1u) == 0) {\n byte = static_cast((byte & 0xF0u) | nibble);\n } else {\n byte = static_cast((byte & 0x0Fu) | (nibble << 4));\n }\n }\n\n if (mask[second]) {\n const uint8_t nibble = static_cast((packed_bits >> (4 * second)) & 0xF);\n if ((elem_idx[second] & 1u) == 0) {\n byte = static_cast((byte & 0xF0u) | nibble);\n } else {\n byte = static_cast((byte & 0x0Fu) | (nibble << 4));\n }\n }\n\n output[byte_idx] = byte;\n }\n }\n}\n\nvoid nvfp4_2d_compute_partial_amax(const Tensor inp, Tensor amax, size_t h, size_t w,\n size_t amax_stride_h, size_t amax_stride_w, size_t start_offset,\n size_t block_len, hipStream_t stream) {\n NVTE_CHECK(block_len == 16, \"NVFP4 2D supports 16x16 tiles only (block_len = 16).\");\n\n size_t len = inp.numel();\n\n assert(h > 0 && w > 0);\n assert(start_offset < h * w);\n assert(start_offset + len <= h * w);\n\n size_t blocks_x = (w + kTileDim - 1) / kTileDim;\n size_t blocks_y = (h + kTileDim - 1) / kTileDim;\n assert(blocks_x <= std::numeric_limits::max());\n assert(blocks_y <= std::numeric_limits::max());\n dim3 grid(blocks_x, blocks_y);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n inp.dtype(), inp_dtype,\n nvfp4_2d_compute_partial_amax_kernel<<>>(\n reinterpret_cast(inp.data.dptr),\n reinterpret_cast(amax.data.dptr), amax_stride_h, amax_stride_w, h, w,\n start_offset, len);)\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid nvfp4_2d_partial_cast(const Tensor inp, Tensor out, const Tensor scale,\n const Tensor global_scale, size_t h, size_t w, size_t scale_stride_h,\n size_t scale_stride_w, size_t start_offset, size_t block_len,\n hipStream_t stream) {\n NVTE_CHECK(block_len == 16, \"NVFP4 2D supports 16x16 tiles only (block_len = 16).\");\n NVTE_CHECK(out.dtype() == DType::kByte, \"NVFP4 rowwise data must be uint8.\");\n\n size_t len = inp.numel();\n\n assert(h > 0 && w > 0);\n assert(start_offset < h * w);\n assert(start_offset + len <= h * w);\n\n size_t blocks_x = (w + kTileDim - 1) / kTileDim;\n size_t blocks_y = (h + kTileDim - 1) / kTileDim;\n assert(blocks_x <= std::numeric_limits::max());\n assert(blocks_y <= std::numeric_limits::max());\n dim3 grid(blocks_x, blocks_y);\n\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n inp.dtype(), inp_dtype,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n w % kTileDim == 0, kWidthAligned,\n nvfp4_2d_partial_cast_kernel\n <<>>(\n reinterpret_cast(inp.data.dptr),\n reinterpret_cast(out.data.dptr),\n reinterpret_cast(scale.data.dptr), scale_stride_h, scale_stride_w,\n reinterpret_cast(global_scale.data.dptr), h, w, start_offset, len);))\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n/*\n * ---------------------------------------------------------------------------\n * NVFP4 TRANSPOSE KERNEL\n *\n * Unlike FP8, NVFP4 packs two 4-bit values into each byte. A simple byte-wise\n * transpose doesn't work because the packing changes:\n * - Before transpose: elements [m, 2c] and [m, 2c+1] share a byte\n * - After transpose: elements [k, 2*m_packed] and [k, 2*m_packed+1] share a byte\n * which were originally [2*m_packed, k] and [2*m_packed+1, k]\n * ---------------------------------------------------------------------------\n */\n\n// Vectorized transpose kernel parameters\nconstexpr int TRANSPOSE_TILE_DIM = 64; // Logical FP4 elements per tile dimension\n// constexpr int TRANSPOSE_TILE_PACKED = 32; // TILE_DIM / 2 bytes\nconstexpr int TRANSPOSE_BLOCK_SIZE = 256; // threads per block\n\n// Shared memory: store unpacked 4-bit values as bytes for easy transpose\n// Size: TILE_DIM x (TILE_DIM + 4) to avoid bank conflicts\nconstexpr int TRANSPOSE_SHMEM_STRIDE = TRANSPOSE_TILE_DIM + 4;\n\n/*\n * Vectorized transpose kernel with uint2 loads/stores (256 threads)\n * Tile: 64x64 logical FP4 = 64x32 packed bytes\n */\n__global__ void __launch_bounds__(TRANSPOSE_BLOCK_SIZE)\n nvfp4_transpose_kernel(const uint8_t *__restrict__ input, uint8_t *__restrict__ output,\n const size_t M, const size_t K) {\n const size_t K_packed = K / 2;\n const size_t M_packed = M / 2;\n\n const size_t tile_m_start = blockIdx.x * TRANSPOSE_TILE_DIM;\n const size_t tile_k_start = blockIdx.y * TRANSPOSE_TILE_DIM;\n\n __shared__ uint8_t shmem[TRANSPOSE_TILE_DIM][TRANSPOSE_SHMEM_STRIDE];\n\n const int tid = threadIdx.x;\n\n // Phase 1: Load input tile with VECTORIZED uint2 reads\n // 256 threads, each loads 8 bytes (uint2) = 2048 bytes total\n // Input tile: [64 rows, 32 cols] = 2048 bytes\n {\n const int thread_row = tid / 4; // 64 rows, 4 threads per row\n const int thread_col = (tid % 4) * 8; // 4 x 8 = 32 bytes per row\n\n const size_t global_m = tile_m_start + thread_row;\n const size_t global_k_packed_base = tile_k_start / 2 + thread_col;\n\n // Load 8 bytes as uint2\n uint2 loaded = make_uint2(0, 0);\n if (global_m < M && global_k_packed_base + 7 < K_packed) {\n loaded = *reinterpret_cast(&input[global_m * K_packed + global_k_packed_base]);\n } else if (global_m < M) {\n // Boundary: scalar loads\n uint8_t *bytes = reinterpret_cast(&loaded);\n#pragma unroll\n for (int b = 0; b < 8; ++b) {\n size_t col = global_k_packed_base + b;\n bytes[b] = (col < K_packed) ? input[global_m * K_packed + col] : 0;\n }\n }\n\n // Unpack 8 bytes -> 16 nibbles and store to shared memory\n const uint8_t *bytes = reinterpret_cast(&loaded);\n#pragma unroll\n for (int b = 0; b < 8; ++b) {\n const int k0 = thread_col * 2 + b * 2;\n const int k1 = k0 + 1;\n shmem[thread_row][k0] = bytes[b] & 0x0F;\n shmem[thread_row][k1] = (bytes[b] >> 4) & 0x0F;\n }\n }\n\n __syncthreads();\n\n // Phase 2: Write output with VECTORIZED uint2 stores\n // Output tile: [64 rows, 32 cols] = 2048 bytes\n {\n const int thread_row = tid / 4; // output K dimension [0, 64)\n const int thread_col_base = (tid % 4) * 8; // output M_packed [0, 32) in steps of 8\n\n const size_t global_k = tile_k_start + thread_row;\n const size_t global_m_packed_base = tile_m_start / 2 + thread_col_base;\n\n if (global_k >= K) return;\n\n // Build 8 output bytes in registers\n uint8_t out_bytes[8];\n\n#pragma unroll\n for (int b = 0; b < 8; ++b) {\n const int out_m_packed = thread_col_base + b;\n\n if (global_m_packed_base + b >= M_packed) {\n out_bytes[b] = 0;\n continue;\n }\n\n // Two M positions that pack into this output byte\n const int m0 = out_m_packed * 2;\n const int m1 = out_m_packed * 2 + 1;\n const int k = thread_row;\n\n // Read from shared memory (transposed access)\n const uint8_t val0 = shmem[m0][k];\n const uint8_t val1 = shmem[m1][k];\n\n out_bytes[b] = val0 | (val1 << 4);\n }\n\n // Vectorized store as uint2\n if (global_m_packed_base + 7 < M_packed) {\n *reinterpret_cast(&output[global_k * M_packed + global_m_packed_base]) =\n *reinterpret_cast(out_bytes);\n } else {\n // Boundary: scalar stores\n for (int b = 0; b < 8 && global_m_packed_base + b < M_packed; ++b) {\n output[global_k * M_packed + global_m_packed_base + b] = out_bytes[b];\n }\n }\n }\n}\n\nvoid nvfp4_transpose(const Tensor input, Tensor output, hipStream_t stream) {\n // Input has logical shape [M, K], stored as [M, K/2] bytes\n // Output has logical shape [K, M], stored as [K, M/2] bytes\n\n NVTE_CHECK(input.dtype() == DType::kByte, \"NVFP4 transpose input must be uint8.\");\n NVTE_CHECK(output.dtype() == DType::kByte, \"NVFP4 transpose output must be uint8.\");\n\n // Get dimensions from packed storage\n // input.shape() = [M, K/2], so M = shape[0], K = shape[1] * 2\n const auto in_shape = input.shape();\n NVTE_CHECK(in_shape.size() == 2, \"NVFP4 transpose expects 2D input (packed), got \",\n in_shape.size(), \"D.\");\n const size_t M = in_shape[0];\n const size_t K_packed = in_shape[1];\n const size_t K = K_packed * 2;\n\n // Output should be [K, M/2]\n const size_t M_packed = M / 2;\n NVTE_CHECK(M % 2 == 0, \"NVFP4 transpose requires M (\", M, \") to be even.\");\n\n const auto out_shape = output.shape();\n NVTE_CHECK(out_shape.size() == 2, \"NVFP4 transpose expects 2D output.\");\n NVTE_CHECK(out_shape[0] == K && out_shape[1] == M_packed,\n \"NVFP4 transpose output shape mismatch. Expected [\", K, \", \", M_packed, \"], got [\",\n out_shape[0], \", \", out_shape[1], \"].\");\n\n if (M == 0 || K == 0) return;\n\n // Use vectorized kernel (faster than TMA for pure transpose)\n // 128x128 tiles with 512 threads and uint4 vectorized access\n dim3 block(TRANSPOSE_BLOCK_SIZE);\n dim3 grid((M + TRANSPOSE_TILE_DIM - 1) / TRANSPOSE_TILE_DIM,\n (K + TRANSPOSE_TILE_DIM - 1) / TRANSPOSE_TILE_DIM);\n\n nvfp4_transpose_kernel<<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(output.data.dptr), M, K);\n\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n/*\n * ---------------------------------------------------------------------------\n * NVFP4 SCALE TRANSPOSE KERNEL\n *\n * Transposes tile-level scales from rowwise to columnwise format.\n * Scale values are stored as E4M3 (fp8) in uint8 tensors.\n *\n * Input (rowwise_scale_inv): [M_padded, K_tiles] where scales are stored\n * at every 16th row (i.e., row 0, 16, 32, ... contain the actual scales,\n * and each row i within a tile block has the same scale as row (i // 16) * 16).\n *\n * Output (columnwise_scale_inv): [K_padded, M_tiles] where scales are\n * repeated 16 times per tile row.\n *\n * Mapping:\n * output[k_tile * 16 + i, m_tile] = input[m_tile * 16, k_tile]\n * for i in [0, 16) and valid (k_tile, m_tile) indices.\n * ---------------------------------------------------------------------------\n */\n__global__ void nvfp4_scale_transpose_kernel(\n const uint8_t *__restrict__ input, // [M_padded, K_tiles], E4M3 stored as uint8\n uint8_t *__restrict__ output, // [K_padded, M_tiles], E4M3 stored as uint8\n const size_t M_tiles, // Number of M tiles\n const size_t K_tiles, // Number of K tiles\n const size_t input_stride, // K_tiles (input row stride)\n const size_t output_stride, // M_tiles (output row stride)\n const size_t K_padded // Output height\n) {\n // Each thread handles one output element\n const size_t out_row = blockIdx.y * blockDim.y + threadIdx.y;\n const size_t out_col = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (out_row >= K_padded || out_col >= M_tiles) return;\n\n // Determine which tile row this belongs to\n const size_t k_tile = out_row / kTileDim;\n\n // Read from input: row = m_tile * 16 (first row of the tile), col = k_tile\n // m_tile = out_col\n if (k_tile < K_tiles) {\n const size_t in_row = out_col * kTileDim; // m_tile * 16\n const uint8_t scale = input[in_row * input_stride + k_tile];\n output[out_row * output_stride + out_col] = scale;\n } else {\n output[out_row * output_stride + out_col] = 0;\n }\n}\n\nvoid nvfp4_scale_transpose(const Tensor input, Tensor output, size_t M_tiles, size_t K_tiles,\n hipStream_t stream) {\n NVTE_CHECK(input.dtype() == DType::kByte, \"NVFP4 scale transpose input must be uint8 (E4M3).\");\n NVTE_CHECK(output.dtype() == DType::kByte, \"NVFP4 scale transpose output must be uint8 (E4M3).\");\n\n const auto in_shape = input.shape();\n const auto out_shape = output.shape();\n NVTE_CHECK(in_shape.size() == 2, \"NVFP4 scale transpose expects 2D input.\");\n NVTE_CHECK(out_shape.size() == 2, \"NVFP4 scale transpose expects 2D output.\");\n\n const size_t input_stride = in_shape[1]; // K_tiles\n const size_t output_stride = out_shape[1]; // M_tiles\n const size_t K_padded = out_shape[0];\n\n if (M_tiles == 0 || K_tiles == 0 || K_padded == 0) return;\n\n constexpr int kBlockDim = 16;\n dim3 block(kBlockDim, kBlockDim);\n dim3 grid((M_tiles + kBlockDim - 1) / kBlockDim, (K_padded + kBlockDim - 1) / kBlockDim);\n\n nvfp4_scale_transpose_kernel<<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(output.data.dptr), M_tiles, K_tiles, input_stride, output_stride,\n K_padded);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n/*\n * ---------------------------------------------------------------------------\n * NVFP4 SCALE EXPANSION KERNEL\n *\n * Expands tile-level scales to row-level scales and converts to FP8 E4M3, used in partial cast.\n *\n * Input (per_block_decode_scale): [tile_rows, tile_cols] in float32\n * Output (target_scale): [rows_padded, tile_cols] in uint8 (E4M3)\n *\n * Each tile row's scale is repeated block_len times in the output.\n * ---------------------------------------------------------------------------\n */\n__global__ void nvfp4_expand_scale_to_fp8_kernel(\n const float *__restrict__ input, // [tile_rows, tile_cols]\n uint8_t *__restrict__ output, // [rows_padded, tile_cols]\n const size_t tile_rows, const size_t tile_cols, const size_t rows_padded,\n const size_t block_len) {\n const size_t out_row = blockIdx.y * blockDim.y + threadIdx.y;\n const size_t out_col = blockIdx.x * blockDim.x + threadIdx.x;\n\n if (out_row >= rows_padded || out_col >= tile_cols) return;\n\n // Determine which tile row this output row belongs to\n const size_t tile_row = out_row / block_len;\n\n float scale_val = 0.0f;\n if (tile_row < tile_rows) {\n scale_val = input[tile_row * tile_cols + out_col];\n }\n\n // Convert float32 to FP8 E4M3\n // Clamp to FP8 E4M3 range and convert\n fp8e4m3 fp8_val = static_cast(scale_val);\n output[out_row * tile_cols + out_col] = reinterpret_cast(fp8_val);\n}\n\nvoid nvfp4_expand_scale_to_fp8(const Tensor input, Tensor output, size_t tile_rows,\n size_t tile_cols, size_t rows_padded, size_t block_len,\n hipStream_t stream) {\n NVTE_CHECK(input.dtype() == DType::kFloat32, \"Scale input must be float32.\");\n NVTE_CHECK(output.dtype() == DType::kByte, \"Scale output must be uint8 (E4M3).\");\n\n if (tile_rows == 0 || tile_cols == 0 || rows_padded == 0) return;\n\n constexpr int kBlockDim = 16;\n dim3 block(kBlockDim, kBlockDim);\n dim3 grid((tile_cols + kBlockDim - 1) / kBlockDim, (rows_padded + kBlockDim - 1) / kBlockDim);\n\n nvfp4_expand_scale_to_fp8_kernel<<>>(\n reinterpret_cast(input.data.dptr),\n reinterpret_cast(output.data.dptr), tile_rows, tile_cols, rows_padded, block_len);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n/*\n * ---------------------------------------------------------------------------\n * NVFP4 COMPUTE PER-BLOCK DECODE SCALE KERNEL\n *\n * Computes per-block decode scale from block amax and global amax:\n * global_scale = (fp8_max * fp4_max) / global_amax = 2688 / global_amax\n * per_block_decode_scale = block_amax * (global_scale * (1 / fp4_max))\n * = block_amax * 448 / global_amax\n *\n * This matches the CUDA device function compute_decoding_scaling_factor() in core_nvfp4.cuh\n *\n * Input (block_amax): [tile_rows, tile_cols] in float32\n * Input (global_amax): scalar float32 (per-tensor amax after all-reduce)\n * Output (scale): [tile_rows, tile_cols] in float32\n * Output (global_scale_out): scalar float32 (the computed global encode scale)\n * ---------------------------------------------------------------------------\n */\n__global__ void nvfp4_compute_per_block_scale_kernel(\n const float *__restrict__ block_amax, // [tile_rows, tile_cols]\n float *__restrict__ scale, // [tile_rows, tile_cols]\n const float *__restrict__ global_amax_ptr, // Pointer to single float value (avoids D2H)\n const size_t numel) {\n const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= numel) return;\n\n constexpr float fp4_max = 6.0f;\n constexpr float fp8_max = 448.0f;\n constexpr float flt_max = 3.402823466e+38f;\n constexpr float tiny = 1.17549435e-38f; // FLT_MIN\n\n // Read global_amax from device memory (avoids D2H transfer)\n float global_amax = *global_amax_ptr;\n\n // Compute global encode scale: S_enc = (fp8_max * fp4_max) / global_amax\n float safe_global_amax = fmaxf(global_amax, tiny);\n float global_scale =\n (global_amax > 0.0f) ? fminf((fp8_max * fp4_max) / safe_global_amax, flt_max) : 1.0f;\n\n // Compute per-block decode scale: S_dec_b = block_amax * (S_enc * (1 / fp4_max))\n float amax_val = block_amax[idx];\n constexpr float fp4_max_inv = 1.0f / fp4_max;\n const float global_scale_multiplier = global_scale * fp4_max_inv;\n float result = fminf(amax_val * global_scale_multiplier, flt_max);\n scale[idx] = result;\n}\n\n// Simple kernel to compute global encode scale from global amax\n__global__ void nvfp4_compute_global_scale_kernel(\n const float *__restrict__ global_amax, // [num_params]\n float *__restrict__ global_scale, // [num_params]\n const size_t num_params) {\n const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= num_params) return;\n\n constexpr float fp4_max = 6.0f;\n constexpr float fp8_max = 448.0f;\n constexpr float flt_max = 3.402823466e+38f;\n constexpr float tiny = 1.17549435e-38f; // FLT_MIN\n\n float amax = global_amax[idx];\n float safe_amax = fmaxf(amax, tiny);\n float scale = (amax > 0.0f) ? fminf((fp8_max * fp4_max) / safe_amax, flt_max) : 1.0f;\n global_scale[idx] = scale;\n}\n\nvoid nvfp4_compute_per_block_scale(const Tensor block_amax, Tensor scale, const Tensor global_amax,\n hipStream_t stream) {\n NVTE_CHECK(block_amax.dtype() == DType::kFloat32, \"Block amax must be float32.\");\n NVTE_CHECK(scale.dtype() == DType::kFloat32, \"Scale must be float32.\");\n NVTE_CHECK(global_amax.dtype() == DType::kFloat32, \"Global amax must be float32.\");\n NVTE_CHECK(global_amax.numel() == 1, \"Global amax must be a single element tensor.\");\n\n size_t numel = block_amax.numel();\n if (numel == 0) return;\n\n constexpr int kBlockSize = 256;\n int grid_size = (numel + kBlockSize - 1) / kBlockSize;\n\n nvfp4_compute_per_block_scale_kernel<<>>(\n reinterpret_cast(block_amax.data.dptr),\n reinterpret_cast(scale.data.dptr),\n reinterpret_cast(global_amax.data.dptr), numel);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid nvfp4_compute_global_scale(const Tensor global_amax, Tensor global_scale,\n hipStream_t stream) {\n NVTE_CHECK(global_amax.dtype() == DType::kFloat32, \"Global amax must be float32.\");\n NVTE_CHECK(global_scale.dtype() == DType::kFloat32, \"Global scale must be float32.\");\n\n size_t num_params = global_amax.numel();\n if (num_params == 0) return;\n\n constexpr int kBlockSize = 256;\n int grid_size = (num_params + kBlockSize - 1) / kBlockSize;\n\n nvfp4_compute_global_scale_kernel<<>>(\n reinterpret_cast(global_amax.data.dptr),\n reinterpret_cast(global_scale.data.dptr), num_params);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n/*\n * ---------------------------------------------------------------------------\n * FUSED NVFP4 SCALE COMPUTATION KERNEL\n *\n * Fuses three operations into one kernel:\n * 1. nvfp4_compute_per_block_scale: compute tile-level decode scales from block amax\n * 2. target_amax.copy_: copy global amax to target tensor\n * 3. nvfp4_expand_scale_to_fp8: expand to row-level and convert to FP8 E4M3\n *\n * Input (block_amax): [tile_rows, tile_cols] float32\n * Input (global_amax): [1] float32\n * Output (per_block_scale): [tile_rows, tile_cols] float32 (intermediate, for partial_cast)\n * Output (target_scale): [rows_padded, tile_cols] uint8 (E4M3)\n * Output (target_amax): [1] float32 (copy of global_amax)\n *\n * Saves 2 kernel launches per parameter (eliminates nvfp4_compute_per_block_scale and\n * nvfp4_expand_scale_to_fp8 as separate calls, plus the amax copy).\n * ---------------------------------------------------------------------------\n */\n__global__ void nvfp4_fused_scale_kernel(\n const float *__restrict__ block_amax, // [tile_rows, tile_cols]\n const float *__restrict__ global_amax, // [1]\n float *__restrict__ per_block_scale, // [tile_rows, tile_cols] - for partial_cast\n uint8_t *__restrict__ target_scale, // [rows_padded, tile_cols]\n float *__restrict__ target_amax, // [1]\n const size_t tile_rows, const size_t tile_cols, const size_t rows_padded,\n const size_t block_len) {\n const size_t out_row = blockIdx.y * blockDim.y + threadIdx.y;\n const size_t out_col = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Read global amax once per thread (broadcast)\n const float g_amax = *global_amax;\n\n // Thread (0,0) copies global_amax to target_amax\n if (out_row == 0 && out_col == 0) {\n *target_amax = g_amax;\n }\n\n if (out_row >= rows_padded || out_col >= tile_cols) return;\n\n // Determine which tile row this output row belongs to\n const size_t tile_row = out_row / block_len;\n\n // Compute the scale value\n constexpr float fp4_max = 6.0f;\n constexpr float fp8_max = 448.0f;\n constexpr float flt_max = 3.402823466e+38f;\n constexpr float tiny = 1.17549435e-38f;\n\n float scale_val = 0.0f;\n if (tile_row < tile_rows) {\n float safe_global_amax = fmaxf(g_amax, tiny);\n float global_scale =\n (g_amax > 0.0f) ? fminf((fp8_max * fp4_max) / safe_global_amax, flt_max) : 1.0f;\n constexpr float fp4_max_inv = 1.0f / fp4_max;\n const float global_scale_multiplier = global_scale * fp4_max_inv;\n\n // Read block amax and compute per-block decode scale\n float amax_val = block_amax[tile_row * tile_cols + out_col];\n scale_val = fminf(amax_val * global_scale_multiplier, flt_max);\n\n // Write per-block scale (only once per tile, when out_row % block_len == 0)\n if (out_row % block_len == 0) {\n per_block_scale[tile_row * tile_cols + out_col] = scale_val;\n }\n }\n\n // Convert float32 to FP8 E4M3 and write expanded scale\n fp8e4m3 fp8_val = static_cast(scale_val);\n target_scale[out_row * tile_cols + out_col] = reinterpret_cast(fp8_val);\n}\n\nvoid nvfp4_fused_scale(const Tensor block_amax, const Tensor global_amax, Tensor per_block_scale,\n Tensor target_scale, Tensor target_amax, size_t tile_rows, size_t tile_cols,\n size_t rows_padded, size_t block_len, hipStream_t stream) {\n NVTE_CHECK(block_amax.dtype() == DType::kFloat32, \"Block amax must be float32.\");\n NVTE_CHECK(global_amax.dtype() == DType::kFloat32, \"Global amax must be float32.\");\n NVTE_CHECK(per_block_scale.dtype() == DType::kFloat32, \"Per-block scale must be float32.\");\n NVTE_CHECK(target_scale.dtype() == DType::kByte, \"Target scale must be uint8 (E4M3).\");\n NVTE_CHECK(target_amax.dtype() == DType::kFloat32, \"Target amax must be float32.\");\n NVTE_CHECK(global_amax.numel() == 1, \"Global amax must be a single element tensor.\");\n NVTE_CHECK(target_amax.numel() == 1, \"Target amax must be a single element tensor.\");\n\n if (tile_rows == 0 || tile_cols == 0 || rows_padded == 0) return;\n\n constexpr int kBlockDim = 16;\n dim3 block(kBlockDim, kBlockDim);\n dim3 grid((tile_cols + kBlockDim - 1) / kBlockDim, (rows_padded + kBlockDim - 1) / kBlockDim);\n\n nvfp4_fused_scale_kernel<<>>(\n reinterpret_cast(block_amax.data.dptr),\n reinterpret_cast(global_amax.data.dptr),\n reinterpret_cast(per_block_scale.data.dptr),\n reinterpret_cast(target_scale.data.dptr),\n reinterpret_cast(target_amax.data.dptr), tile_rows, tile_cols, rows_padded,\n block_len);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n#endif // FP4_TYPE_SUPPORTED\n} // namespace nvfp4_recipe\n} // namespace transformer_engine\n\nvoid nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, size_t tile_rows,\n size_t tile_cols, size_t rows_padded, size_t block_len,\n hipStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_expand_scale_to_fp8);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_expand_scale_to_fp8(*convertNVTETensorCheck(input),\n *convertNVTETensorCheck(output), tile_rows, tile_cols,\n rows_padded, block_len, stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor scale,\n const NVTETensor global_amax, hipStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_compute_per_block_scale);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_compute_per_block_scale(*convertNVTETensorCheck(block_amax),\n *convertNVTETensorCheck(scale),\n *convertNVTETensorCheck(global_amax), stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_compute_global_scale(const NVTETensor global_amax, NVTETensor global_scale,\n hipStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_compute_global_scale);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_compute_global_scale(*convertNVTETensorCheck(global_amax),\n *convertNVTETensorCheck(global_scale), stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_scale_transpose(const NVTETensor input, NVTETensor output, size_t M_tiles,\n size_t K_tiles, hipStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_scale_transpose);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_scale_transpose(*convertNVTETensorCheck(input),\n *convertNVTETensorCheck(output), M_tiles, K_tiles, stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_data_transpose(const NVTETensor input, NVTETensor output, hipStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_data_transpose);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_transpose(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output),\n stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, size_t h, size_t w,\n size_t amax_stride_h, size_t amax_stride_w,\n size_t start_offset, size_t block_len,\n hipStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_2d_compute_partial_amax);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_2d_compute_partial_amax(*convertNVTETensorCheck(inp),\n *convertNVTETensorCheck(amax), h, w, amax_stride_h,\n amax_stride_w, start_offset, block_len, stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale,\n const NVTETensor global_scale, size_t h, size_t w,\n size_t scale_stride_h, size_t scale_stride_w, size_t start_offset,\n size_t block_len, hipStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_2d_partial_cast);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_2d_partial_cast(*convertNVTETensorCheck(inp), *convertNVTETensorCheck(out),\n *convertNVTETensorCheck(scale),\n *convertNVTETensorCheck(global_scale), h, w, scale_stride_h,\n scale_stride_w, start_offset, block_len, stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_rowwise_amax_A,\n const NVTETensor inpB, const bool use_rowwise_amax_B,\n float alpha_in, NVTETensor alpha_out,\n hipStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_compute_per_tensor_scale);\n using namespace transformer_engine;\n\n auto *tA = convertNVTETensor(inpA);\n auto *tB = convertNVTETensor(inpB);\n auto *tOut = convertNVTETensor(alpha_out);\n\n void *amax_A_ptr = use_rowwise_amax_A ? tA->amax.dptr : tA->columnwise_amax.dptr;\n void *amax_B_ptr = use_rowwise_amax_B ? tB->amax.dptr : tB->columnwise_amax.dptr;\n void *alpha_ptr = tOut->data.dptr;\n\n // check for not null pointers\n NVTE_CHECK(amax_A_ptr != nullptr, \"amax_A_ptr is null\");\n NVTE_CHECK(amax_B_ptr != nullptr, \"amax_B_ptr is null\");\n NVTE_CHECK(alpha_ptr != nullptr, \"alpha_ptr is null\");\n\n nvfp4_recipe::compute_nvfp4_per_tensor_scale_kernel<<<1, 1, 0, stream>>>(\n alpha_in, reinterpret_cast(amax_A_ptr),\n reinterpret_cast(amax_B_ptr), reinterpret_cast(alpha_ptr));\n NVTE_CHECK_CUDA(hipGetLastError());\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n\nvoid nvte_nvfp4_fused_scale(const NVTETensor block_amax, const NVTETensor global_amax,\n NVTETensor per_block_scale, NVTETensor target_scale,\n NVTETensor target_amax, size_t tile_rows, size_t tile_cols,\n size_t rows_padded, size_t block_len, hipStream_t stream) {\n#if FP4_TYPE_SUPPORTED\n NVTE_API_CALL(nvte_nvfp4_fused_scale);\n using namespace transformer_engine;\n nvfp4_recipe::nvfp4_fused_scale(\n *convertNVTETensorCheck(block_amax), *convertNVTETensorCheck(global_amax),\n *convertNVTETensorCheck(per_block_scale), *convertNVTETensorCheck(target_scale),\n *convertNVTETensorCheck(target_amax), tile_rows, tile_cols, rows_padded, block_len, stream);\n#else\n NVTE_ERROR(\"FP4 support requires CUDA 12.8+, but compile-time CUDA version is \", CUDA_VERSION);\n#endif // FP4_TYPE_SUPPORTED\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/recipe/nvfp4.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n\n#include \"../utils.cuh\"\n#include \"multi_tensor_apply.cuh\"\n\nnamespace transformer_engine {\nnamespace multi_tensor_l2norm {\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntemplate \n__device__ __forceinline__ bool is_aligned(T *p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T *dst, T *src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT *)dst)[dst_offset] = ((LT *)src)[src_offset]; // NOLINT(*)\n}\n\ntemplate \n__device__ __forceinline__ T\nreduce_block_into_lanes(T *x, T val, int lanes = 1,\n bool share_result = false) { // lanes is intended to be <= 32.\n int tid = threadIdx.x + threadIdx.y * blockDim.x;\n int blockSize = blockDim.x * blockDim.y; // blockSize is intended to be a multiple of 32.\n\n if (blockSize >= 64) {\n x[tid] = val;\n __syncthreads();\n }\n\n#pragma unroll\n for (int i = (blockSize >> 1); i >= 64; i >>= 1) {\n if (tid < i) x[tid] = x[tid] + x[tid + i];\n __syncthreads();\n }\n\n T final;\n\n if (tid < 32) {\n if (blockSize >= 64)\n final = x[tid] + x[tid + 32];\n else\n final = val;\n // __SYNCWARP();\n\n#pragma unroll\n for (int i = 16; i >= lanes; i >>= 1) final = final + __shfl_down_sync(0xffffffff, final, i);\n }\n\n if (share_result) {\n if (tid < lanes) x[tid] = final; // EpilogueOp\n // Make sure the smem result is visible to all warps.\n }\n __syncthreads();\n // Avoid potential write before read race when reduce_block_into_lanes is called back to back\n\n return final;\n}\n\ntemplate \n__device__ __forceinline__ T\nreduce_block_into_lanes_max_op(T *x, T val, int lanes = 1,\n bool share_result = false) { // lanes is intended to be <= 32.\n int tid = threadIdx.x + threadIdx.y * blockDim.x;\n int blockSize = blockDim.x * blockDim.y; // blockSize is intended to be a multiple of 32.\n\n if (blockSize >= 64) {\n x[tid] = val;\n __syncthreads();\n }\n\n#pragma unroll\n for (int i = (blockSize >> 1); i >= 64; i >>= 1) {\n if (tid < i) x[tid] = fmaxf(fabsf(x[tid]), fabsf(x[tid + i]));\n __syncthreads();\n }\n\n T final;\n\n if (tid < 32) {\n if (blockSize >= 64)\n final = fmaxf(fabsf(x[tid]), fabsf(x[tid + 32]));\n else\n final = val;\n // __SYNCWARP();\n\n#pragma unroll\n for (int i = 16; i >= lanes; i >>= 1)\n final = fmaxf(fabsf(final), fabsf(__shfl_down_sync(0xffffffff, final, i)));\n }\n\n if (share_result) {\n if (tid < lanes) x[tid] = final; // EpilogueOp\n // Make sure the smem result is visible to all warps.\n __syncthreads();\n }\n\n return final;\n}\n\ntemplate \nstruct L2NormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int *noop_gmem,\n TensorListMetadata<1> &tl, // NOLINT(*),\n float *output, float *output_per_tensor,\n bool per_tensor, int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t *x = reinterpret_cast(tl.addresses[0][tensor_loc]);\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0.f;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size;\n i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]);\n vals[ii] += next * next;\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]);\n vals[ii] += next * next;\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val += vals[i];\n\n float final = reduce_block_into_lanes(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final))\n *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] += final;\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor +\n chunk_idx] = final;\n }\n }\n};\n\ntemplate \nstruct UnscaleL2NormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int *noop_gmem,\n TensorListMetadata<1> &tl, // NOLINT(*),\n const float *inv_scale, float *output,\n float *output_per_tensor, bool per_tensor,\n int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t *x = reinterpret_cast(tl.addresses[0][tensor_loc]);\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0.f;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size;\n i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]) * (*inv_scale);\n vals[ii] += next * next;\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]) * (*inv_scale);\n vals[ii] += next * next;\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val += vals[i];\n\n float final = reduce_block_into_lanes(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final))\n *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] += final;\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor +\n chunk_idx] = final;\n }\n }\n};\n\n// Probably better to template, but since we are not likely to support other norm\ntemplate \nstruct MaxNormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int *noop_gmem,\n TensorListMetadata<1> &tl, // NOLINT(*),\n float *output, float *output_per_tensor,\n bool per_tensor, int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t *x = reinterpret_cast(tl.addresses[0][tensor_loc]);\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size;\n i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]);\n vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]);\n vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val = fmaxf(fabsf(val), fabsf(vals[i]));\n\n float final = reduce_block_into_lanes_max_op(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final))\n *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] = fmaxf(fabsf(output[blockIdx.x]), fabsf(final));\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor +\n chunk_idx] = final;\n }\n }\n};\n\n__global__ void cleanup(float *output, float *output_per_tensor, float *ret, float *ret_per_tensor,\n bool per_tensor, int max_chunks_per_tensor) {\n __shared__ float vals[512];\n\n if (blockIdx.x == 0) {\n float val = 0;\n if (threadIdx.x < 320) val = output[threadIdx.x];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) *ret = sqrt(final);\n }\n\n if (per_tensor) {\n float *output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;\n\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)\n val += output_this_tensor[i];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) ret_per_tensor[blockIdx.x] = sqrt(final);\n }\n}\n\n__global__ void cleanup_v2(float *output, float *output_per_tensor, float *ret,\n float *ret_per_tensor, bool per_tensor, int max_chunks_per_tensor,\n int norm_type, float alpha, float beta) {\n __shared__ float vals[512];\n\n if (blockIdx.x == 0) {\n float val = 0;\n if (threadIdx.x < 320) val = output[threadIdx.x];\n\n if (norm_type == 0) {\n float final = reduce_block_into_lanes_max_op(vals, val);\n if (threadIdx.x == 0) *ret = alpha * (*ret) + beta * final;\n } else {\n float final = reduce_block_into_lanes(vals, val);\n if (threadIdx.x == 0) *ret = sqrt(alpha * (*ret) * (*ret) + beta * final);\n }\n }\n\n if (per_tensor) {\n float *output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;\n\n if (norm_type == 0) {\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)\n val = fmaxf(fabsf(val), fabsf(output_this_tensor[i]));\n\n float final = reduce_block_into_lanes_max_op(vals, val);\n\n if (threadIdx.x == 0)\n ret_per_tensor[blockIdx.x] = alpha * ret_per_tensor[blockIdx.x] + beta * final;\n } else {\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)\n val += output_this_tensor[i];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0)\n ret_per_tensor[blockIdx.x] =\n sqrt(alpha * ret_per_tensor[blockIdx.x] * ret_per_tensor[blockIdx.x] + beta * final);\n }\n }\n}\n\nvoid multi_tensor_l2norm_cuda(int chunk_size, Tensor noop_flag,\n std::vector> tensor_lists, Tensor output,\n Tensor output_per_tensor, Tensor ret, Tensor ret_per_tensor,\n bool per_tensor, int max_chunks_per_tensor, cudaStream_t stream) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n tensor_lists[0][0]->dtype(), dtype,\n multi_tensor_apply<1>(\n BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, L2NormFunctor(), stream,\n reinterpret_cast(output.data.dptr),\n per_tensor ? reinterpret_cast(output_per_tensor.data.dptr) : nullptr, per_tensor,\n max_chunks_per_tensor);)\n\n NVTE_CHECK_CUDA(cudaGetLastError());\n\n // This involves one more small kernel launches, but will be negligible end to end.\n // I could get rid of these by hacking the functor + multi tensor harness with persistence\n // logic, but keeping it simple for now\n cleanup<<>>(\n reinterpret_cast(output.data.dptr),\n per_tensor ? reinterpret_cast(output_per_tensor.data.dptr) : nullptr,\n reinterpret_cast(ret.data.dptr),\n per_tensor ? reinterpret_cast(ret_per_tensor.data.dptr) : nullptr, per_tensor,\n max_chunks_per_tensor);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\nvoid multi_tensor_unscale_l2norm_cuda(int chunk_size, Tensor noop_flag,\n std::vector> tensor_lists,\n Tensor output, Tensor output_per_tensor, Tensor ret,\n Tensor ret_per_tensor, Tensor inv_scale, bool per_tensor,\n int max_chunks_per_tensor, cudaStream_t stream) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n tensor_lists[0][0]->dtype(), dtype,\n multi_tensor_apply<1>(\n BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, UnscaleL2NormFunctor(), stream,\n reinterpret_cast(inv_scale.data.dptr),\n reinterpret_cast(output.data.dptr),\n per_tensor ? reinterpret_cast(output_per_tensor.data.dptr) : nullptr, per_tensor,\n max_chunks_per_tensor);)\n\n NVTE_CHECK_CUDA(cudaGetLastError());\n\n // This involves one more small kernel launches, but will be negligible end to end.\n // I could get rid of these by hacking the functor + multi tensor harness with persistence\n // logic, but keeping it simple for now\n cleanup<<>>(\n reinterpret_cast(output.data.dptr),\n per_tensor ? reinterpret_cast(output_per_tensor.data.dptr) : nullptr,\n reinterpret_cast(ret.data.dptr),\n per_tensor ? reinterpret_cast(ret_per_tensor.data.dptr) : nullptr, per_tensor,\n max_chunks_per_tensor);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n} // namespace multi_tensor_l2norm\n} // namespace transformer_engine\n\nvoid nvte_multi_tensor_l2norm_cuda(int chunk_size, NVTETensor noop_flag, NVTETensor **tensor_lists,\n const size_t num_tensor_lists, const size_t num_tensors_per_list,\n NVTETensor output, NVTETensor output_per_tensor, NVTETensor ret,\n NVTETensor ret_per_tensor, int per_tensor,\n int max_chunks_per_tensor, cudaStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_l2norm_cuda);\n using namespace transformer_engine;\n\n multi_tensor_l2norm::multi_tensor_l2norm_cuda(\n chunk_size, *convertNVTETensorCheck(noop_flag),\n convert_tensor_array(tensor_lists, num_tensor_lists, num_tensors_per_list),\n *convertNVTETensorCheck(output), *convertNVTETensorCheck(output_per_tensor),\n *convertNVTETensorCheck(ret), *convertNVTETensorCheck(ret_per_tensor), per_tensor,\n max_chunks_per_tensor, stream);\n}\n\nvoid nvte_multi_tensor_unscale_l2norm_cuda(int chunk_size, NVTETensor noop_flag,\n NVTETensor **tensor_lists, const size_t num_tensor_lists,\n const size_t num_tensors_per_list, NVTETensor output,\n NVTETensor output_per_tensor, NVTETensor ret,\n NVTETensor ret_per_tensor, NVTETensor inv_scale,\n int per_tensor, int max_chunks_per_tensor,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_unscale_l2norm_cuda);\n using namespace transformer_engine;\n\n multi_tensor_l2norm::multi_tensor_unscale_l2norm_cuda(\n chunk_size, *convertNVTETensorCheck(noop_flag),\n convert_tensor_array(tensor_lists, num_tensor_lists, num_tensors_per_list),\n *convertNVTETensorCheck(output), *convertNVTETensorCheck(output_per_tensor),\n *convertNVTETensorCheck(ret), *convertNVTETensorCheck(ret_per_tensor),\n *convertNVTETensorCheck(inv_scale), per_tensor, max_chunks_per_tensor, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n#include \"../utils.cuh\"\n#include \"multi_tensor_apply.cuh\"\n\nnamespace transformer_engine {\nnamespace multi_tensor_l2norm {\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntemplate \n__device__ __forceinline__ bool is_aligned(T *p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T *dst, T *src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT *)dst)[dst_offset] = ((LT *)src)[src_offset]; // NOLINT(*)\n}\n\ntemplate \n__device__ __forceinline__ T\nreduce_block_into_lanes(T *x, T val, int lanes = 1,\n bool share_result = false) { // lanes is intended to be <= 32.\n int tid = threadIdx.x + threadIdx.y * blockDim.x;\n int blockSize = blockDim.x * blockDim.y; // blockSize is intended to be a multiple of 32.\n\n if (blockSize >= 64) {\n x[tid] = val;\n __syncthreads();\n }\n\n#pragma unroll\n for (int i = (blockSize >> 1); i >= 64; i >>= 1) {\n if (tid < i) x[tid] = x[tid] + x[tid + i];\n __syncthreads();\n }\n\n T final;\n\n if (tid < 32) {\n if (blockSize >= 64)\n final = x[tid] + x[tid + 32];\n else\n final = val;\n // __SYNCWARP();\n\n#pragma unroll\n for (int i = 16; i >= lanes; i >>= 1) final = final + __shfl_down_sync(0xffffffff, final, i);\n }\n\n if (share_result) {\n if (tid < lanes) x[tid] = final; // EpilogueOp\n // Make sure the smem result is visible to all warps.\n }\n __syncthreads();\n // Avoid potential write before read race when reduce_block_into_lanes is called back to back\n\n return final;\n}\n\ntemplate \n__device__ __forceinline__ T\nreduce_block_into_lanes_max_op(T *x, T val, int lanes = 1,\n bool share_result = false) { // lanes is intended to be <= 32.\n int tid = threadIdx.x + threadIdx.y * blockDim.x;\n int blockSize = blockDim.x * blockDim.y; // blockSize is intended to be a multiple of 32.\n\n if (blockSize >= 64) {\n x[tid] = val;\n __syncthreads();\n }\n\n#pragma unroll\n for (int i = (blockSize >> 1); i >= 64; i >>= 1) {\n if (tid < i) x[tid] = fmaxf(fabsf(x[tid]), fabsf(x[tid + i]));\n __syncthreads();\n }\n\n T final;\n\n if (tid < 32) {\n if (blockSize >= 64)\n final = fmaxf(fabsf(x[tid]), fabsf(x[tid + 32]));\n else\n final = val;\n // __SYNCWARP();\n\n#pragma unroll\n for (int i = 16; i >= lanes; i >>= 1)\n final = fmaxf(fabsf(final), fabsf(__shfl_down_sync(0xffffffff, final, i)));\n }\n\n if (share_result) {\n if (tid < lanes) x[tid] = final; // EpilogueOp\n // Make sure the smem result is visible to all warps.\n __syncthreads();\n }\n\n return final;\n}\n\ntemplate \nstruct L2NormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int *noop_gmem,\n TensorListMetadata<1> &tl, // NOLINT(*),\n float *output, float *output_per_tensor,\n bool per_tensor, int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t *x = reinterpret_cast(tl.addresses[0][tensor_loc]);\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0.f;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size;\n i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]);\n vals[ii] += next * next;\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]);\n vals[ii] += next * next;\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val += vals[i];\n\n float final = reduce_block_into_lanes(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final))\n *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] += final;\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor +\n chunk_idx] = final;\n }\n }\n};\n\ntemplate \nstruct UnscaleL2NormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int *noop_gmem,\n TensorListMetadata<1> &tl, // NOLINT(*),\n const float *inv_scale, float *output,\n float *output_per_tensor, bool per_tensor,\n int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t *x = reinterpret_cast(tl.addresses[0][tensor_loc]);\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0.f;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size;\n i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]) * (*inv_scale);\n vals[ii] += next * next;\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]) * (*inv_scale);\n vals[ii] += next * next;\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val += vals[i];\n\n float final = reduce_block_into_lanes(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final))\n *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] += final;\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor +\n chunk_idx] = final;\n }\n }\n};\n\n// Probably better to template, but since we are not likely to support other norm\ntemplate \nstruct MaxNormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int *noop_gmem,\n TensorListMetadata<1> &tl, // NOLINT(*),\n float *output, float *output_per_tensor,\n bool per_tensor, int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t *x = reinterpret_cast(tl.addresses[0][tensor_loc]);\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size;\n i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]);\n vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]);\n vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val = fmaxf(fabsf(val), fabsf(vals[i]));\n\n float final = reduce_block_into_lanes_max_op(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final))\n *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] = fmaxf(fabsf(output[blockIdx.x]), fabsf(final));\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor +\n chunk_idx] = final;\n }\n }\n};\n\n__global__ void cleanup(float *output, float *output_per_tensor, float *ret, float *ret_per_tensor,\n bool per_tensor, int max_chunks_per_tensor) {\n __shared__ float vals[512];\n\n if (blockIdx.x == 0) {\n float val = 0;\n if (threadIdx.x < 320) val = output[threadIdx.x];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) *ret = sqrt(final);\n }\n\n if (per_tensor) {\n float *output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;\n\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)\n val += output_this_tensor[i];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) ret_per_tensor[blockIdx.x] = sqrt(final);\n }\n}\n\n__global__ void cleanup_v2(float *output, float *output_per_tensor, float *ret,\n float *ret_per_tensor, bool per_tensor, int max_chunks_per_tensor,\n int norm_type, float alpha, float beta) {\n __shared__ float vals[512];\n\n if (blockIdx.x == 0) {\n float val = 0;\n if (threadIdx.x < 320) val = output[threadIdx.x];\n\n if (norm_type == 0) {\n float final = reduce_block_into_lanes_max_op(vals, val);\n if (threadIdx.x == 0) *ret = alpha * (*ret) + beta * final;\n } else {\n float final = reduce_block_into_lanes(vals, val);\n if (threadIdx.x == 0) *ret = sqrt(alpha * (*ret) * (*ret) + beta * final);\n }\n }\n\n if (per_tensor) {\n float *output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;\n\n if (norm_type == 0) {\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)\n val = fmaxf(fabsf(val), fabsf(output_this_tensor[i]));\n\n float final = reduce_block_into_lanes_max_op(vals, val);\n\n if (threadIdx.x == 0)\n ret_per_tensor[blockIdx.x] = alpha * ret_per_tensor[blockIdx.x] + beta * final;\n } else {\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)\n val += output_this_tensor[i];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0)\n ret_per_tensor[blockIdx.x] =\n sqrt(alpha * ret_per_tensor[blockIdx.x] * ret_per_tensor[blockIdx.x] + beta * final);\n }\n }\n}\n\nvoid multi_tensor_l2norm_cuda(int chunk_size, Tensor noop_flag,\n std::vector> tensor_lists, Tensor output,\n Tensor output_per_tensor, Tensor ret, Tensor ret_per_tensor,\n bool per_tensor, int max_chunks_per_tensor, hipStream_t stream) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n tensor_lists[0][0]->dtype(), dtype,\n multi_tensor_apply<1>(\n BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, L2NormFunctor(), stream,\n reinterpret_cast(output.data.dptr),\n per_tensor ? reinterpret_cast(output_per_tensor.data.dptr) : nullptr, per_tensor,\n max_chunks_per_tensor);)\n\n NVTE_CHECK_CUDA(hipGetLastError());\n\n // This involves one more small kernel launches, but will be negligible end to end.\n // I could get rid of these by hacking the functor + multi tensor harness with persistence\n // logic, but keeping it simple for now\n cleanup<<>>(\n reinterpret_cast(output.data.dptr),\n per_tensor ? reinterpret_cast(output_per_tensor.data.dptr) : nullptr,\n reinterpret_cast(ret.data.dptr),\n per_tensor ? reinterpret_cast(ret_per_tensor.data.dptr) : nullptr, per_tensor,\n max_chunks_per_tensor);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\nvoid multi_tensor_unscale_l2norm_cuda(int chunk_size, Tensor noop_flag,\n std::vector> tensor_lists,\n Tensor output, Tensor output_per_tensor, Tensor ret,\n Tensor ret_per_tensor, Tensor inv_scale, bool per_tensor,\n int max_chunks_per_tensor, hipStream_t stream) {\n TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(\n tensor_lists[0][0]->dtype(), dtype,\n multi_tensor_apply<1>(\n BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, UnscaleL2NormFunctor(), stream,\n reinterpret_cast(inv_scale.data.dptr),\n reinterpret_cast(output.data.dptr),\n per_tensor ? reinterpret_cast(output_per_tensor.data.dptr) : nullptr, per_tensor,\n max_chunks_per_tensor);)\n\n NVTE_CHECK_CUDA(hipGetLastError());\n\n // This involves one more small kernel launches, but will be negligible end to end.\n // I could get rid of these by hacking the functor + multi tensor harness with persistence\n // logic, but keeping it simple for now\n cleanup<<>>(\n reinterpret_cast(output.data.dptr),\n per_tensor ? reinterpret_cast(output_per_tensor.data.dptr) : nullptr,\n reinterpret_cast(ret.data.dptr),\n per_tensor ? reinterpret_cast(ret_per_tensor.data.dptr) : nullptr, per_tensor,\n max_chunks_per_tensor);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n} // namespace multi_tensor_l2norm\n} // namespace transformer_engine\n\nvoid nvte_multi_tensor_l2norm_cuda(int chunk_size, NVTETensor noop_flag, NVTETensor **tensor_lists,\n const size_t num_tensor_lists, const size_t num_tensors_per_list,\n NVTETensor output, NVTETensor output_per_tensor, NVTETensor ret,\n NVTETensor ret_per_tensor, int per_tensor,\n int max_chunks_per_tensor, hipStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_l2norm_cuda);\n using namespace transformer_engine;\n\n multi_tensor_l2norm::multi_tensor_l2norm_cuda(\n chunk_size, *convertNVTETensorCheck(noop_flag),\n convert_tensor_array(tensor_lists, num_tensor_lists, num_tensors_per_list),\n *convertNVTETensorCheck(output), *convertNVTETensorCheck(output_per_tensor),\n *convertNVTETensorCheck(ret), *convertNVTETensorCheck(ret_per_tensor), per_tensor,\n max_chunks_per_tensor, stream);\n}\n\nvoid nvte_multi_tensor_unscale_l2norm_cuda(int chunk_size, NVTETensor noop_flag,\n NVTETensor **tensor_lists, const size_t num_tensor_lists,\n const size_t num_tensors_per_list, NVTETensor output,\n NVTETensor output_per_tensor, NVTETensor ret,\n NVTETensor ret_per_tensor, NVTETensor inv_scale,\n int per_tensor, int max_chunks_per_tensor,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_multi_tensor_unscale_l2norm_cuda);\n using namespace transformer_engine;\n\n multi_tensor_l2norm::multi_tensor_unscale_l2norm_cuda(\n chunk_size, *convertNVTETensorCheck(noop_flag),\n convert_tensor_array(tensor_lists, num_tensor_lists, num_tensors_per_list),\n *convertNVTETensorCheck(output), *convertNVTETensorCheck(output_per_tensor),\n *convertNVTETensorCheck(ret), *convertNVTETensorCheck(ret_per_tensor),\n *convertNVTETensorCheck(inv_scale), per_tensor, max_chunks_per_tensor, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/multi_tensor/l2norm.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \"common/common.h\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"hadamard_transform_utils.cuh\"\n\nnamespace transformer_engine {\nnamespace {\n\nconstexpr int kMaxTensorsPerKernel = 64; // Args must be <4 KB, expand 64 if needed\nstruct MultiAmaxArgs {\n // (output) Amax buffer for pre-RHT amax buffer\n void* output_pre_rht_amax_list[kMaxTensorsPerKernel];\n // (output) Amax buffer for RHT identity amax buffer\n void* output_identity_amax_list[kMaxTensorsPerKernel];\n // (output) Amax buffer for RHT transpose amax buffer\n void* output_transpose_amax_list[kMaxTensorsPerKernel];\n // Prefix sum (with leading zero) of split_sections of each tensor of input\n int split_sections_range[kMaxTensorsPerKernel + 1];\n // Number of tensors (splits) being processed by kernel\n int num_tensors;\n};\n\nconstexpr int kThreadsPerWarp = 32;\n\ntemplate \n__device__ __forceinline__ void ComputeKernel(uint32_t b_frag_i[4], uint32_t b_frag_t[4],\n IType* in_sh_ptr, int swizzle_idx,\n uint32_t& local_pre_rht_amax_reg,\n uint32_t& local_amax_reg,\n uint32_t& local_amax_t_reg) {\n uint32_t a_frag[4]; // A matrix fragment\n uint32_t c_frag[4]; // Result fragment\n\n uint32_t temp_amax_reg;\n uint32_t temp_amax_t_reg;\n\n if (kReturnIdentityAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0], a_frag[1], a_frag[2], a_frag[3], b_frag_i[0], b_frag_i[1], b_frag_i[2],\n b_frag_i[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_reg);\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_amax_reg)\n : \"r\"(local_amax_reg), \"r\"(temp_amax_reg));\n }\n\n if (kReturnTransposedAmax) {\n if (!kReturnIdentityAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n } else {\n matrix_transpose_m8_n8_b16_inplace(a_frag[0]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[1]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[2]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[3]);\n }\n\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0], a_frag[2], a_frag[1], a_frag[3], b_frag_t[0], b_frag_t[1], b_frag_t[2],\n b_frag_t[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_t_reg);\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_amax_t_reg)\n : \"r\"(local_amax_t_reg), \"r\"(temp_amax_t_reg));\n }\n\n if (kReturnPreRhtAmax) {\n if (!kReturnIdentityAmax && !kReturnTransposedAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n }\n\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[0])\n : \"r\"(a_frag[0]), \"r\"(a_frag[1]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[2])\n : \"r\"(a_frag[2]), \"r\"(a_frag[3]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[0])\n : \"r\"(a_frag[0]), \"r\"(a_frag[2]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_pre_rht_amax_reg)\n : \"r\"(a_frag[0]), \"r\"(local_pre_rht_amax_reg));\n }\n}\n\ntemplate \n__device__ __host__ constexpr int NextPowerOf2() {\n static_assert(kN > 0, \"kN must be > 0\");\n // Round up to the next power of 2 by counting leading zeros.\n return 1 << (32 - __builtin_clz(kN - 1));\n}\n\ntemplate \n__device__ __forceinline__ void ReduceMax(const float pre_rht_amax, const float identity_amax,\n const float transpose_amax, float* staging_for_pre_rht,\n float* staging_for_identity, float* staging_for_transpose,\n float* output_pre_rht_amax_ptr,\n float* output_identity_amax_ptr,\n float* output_transpose_amax_ptr, const int warpid) {\n // intra-warp reduction\n constexpr int kWarpSize = 32;\n int local_rank = threadIdx.x % 32;\n float warp_pre_rht_amax = kReturnPreRhtAmax ? warp_reduce_max(pre_rht_amax) : 0.0f;\n float warp_identity_amax = kReturnIdentityAmax ? warp_reduce_max(identity_amax) : 0.0f;\n float warp_transpose_amax =\n kReturnTransposedAmax ? warp_reduce_max(transpose_amax) : 0.0f;\n\n // inter-warp reduction\n if (threadIdx.x % 32 == 0) {\n if (kReturnPreRhtAmax) {\n staging_for_pre_rht[warpid] = warp_pre_rht_amax;\n }\n if (kReturnIdentityAmax) {\n staging_for_identity[warpid] = warp_identity_amax;\n }\n if (kReturnTransposedAmax) {\n staging_for_transpose[warpid] = warp_transpose_amax;\n }\n }\n __syncthreads();\n constexpr int kNumWarpsPow2 = NextPowerOf2();\n if (warpid == 0) {\n if (kReturnIdentityAmax) {\n float identity_accum = local_rank < kNumWarps ? staging_for_identity[local_rank] : 0.0f;\n identity_accum = warp_reduce_max(identity_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_identity_amax_ptr, identity_accum);\n }\n }\n }\n if (warpid == 1) {\n if (kReturnTransposedAmax) {\n float transpose_accum = local_rank < kNumWarps ? staging_for_transpose[local_rank] : 0.0f;\n transpose_accum = warp_reduce_max(transpose_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_transpose_amax_ptr, transpose_accum);\n }\n }\n }\n if (warpid == 2) {\n if (kReturnPreRhtAmax) {\n float pre_rht_accum = local_rank < kNumWarps ? staging_for_pre_rht[local_rank] : 0.0f;\n pre_rht_accum = warp_reduce_max(pre_rht_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_pre_rht_amax_ptr, pre_rht_accum);\n }\n }\n }\n}\n\n// args: the mult-tensor amax arguments\n__global__ void MultiZeroAmaxKernel(MultiAmaxArgs args) {\n int num_tensors = args.num_tensors;\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n\n for (; tid < num_tensors; tid += stride) {\n float* output_pre_rht_amax_ptr = static_cast(args.output_pre_rht_amax_list[tid]);\n float* output_identity_amax_ptr = static_cast(args.output_identity_amax_list[tid]);\n float* output_transpose_amax_ptr = static_cast(args.output_transpose_amax_list[tid]);\n if (output_pre_rht_amax_ptr != nullptr) {\n *output_pre_rht_amax_ptr = 0;\n }\n if (output_identity_amax_ptr != nullptr) {\n *output_identity_amax_ptr = 0;\n }\n if (output_transpose_amax_ptr != nullptr) {\n *output_transpose_amax_ptr = 0;\n }\n }\n}\n\n// args: the mult-tensor amax arguments\n__global__ void MultiAmaxMemcpyD2DKernelPreRHT(MultiAmaxArgs args) {\n int num_tensors = args.num_tensors;\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n\n for (; tid < num_tensors; tid += stride) {\n float* output_pre_rht_amax_ptr = static_cast(args.output_pre_rht_amax_list[tid]);\n float* output_identity_amax_ptr = static_cast(args.output_identity_amax_list[tid]);\n float* output_transpose_amax_ptr = static_cast(args.output_transpose_amax_list[tid]);\n if (output_pre_rht_amax_ptr != nullptr) {\n float pre_rht_amax = *output_pre_rht_amax_ptr;\n if (output_identity_amax_ptr != nullptr) {\n *output_identity_amax_ptr = pre_rht_amax;\n }\n if (output_transpose_amax_ptr != nullptr) {\n *output_transpose_amax_ptr = pre_rht_amax;\n }\n }\n }\n}\n\ntemplate \n__global__ void GroupHadamardAmaxTmaKernel(const __grid_constant__ CUtensorMap tensor_map_input,\n const MultiAmaxArgs args, uint16_t random_sign_mask,\n uint16_t random_sign_mask_t, uint64_t num_rows,\n uint64_t row_length) {\n#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n\n float* output_pre_rht_amax_ptr;\n float* output_identity_amax_ptr;\n float* output_transpose_amax_ptr;\n\n // calculate the global offset in Y direction to access the correct amax buffer\n int global_offset_y = blockIdx.y * CHUNK_DIM_Y;\n int tensor_id = 0;\n while (args.split_sections_range[tensor_id + 1] <= global_offset_y) {\n ++tensor_id;\n }\n output_pre_rht_amax_ptr = static_cast(args.output_pre_rht_amax_list[tensor_id]);\n output_identity_amax_ptr = static_cast(args.output_identity_amax_list[tensor_id]);\n output_transpose_amax_ptr = static_cast(args.output_transpose_amax_list[tensor_id]);\n\n static_assert(CHUNK_DIM_Y >= BUFF_DIM_Y && CHUNK_DIM_Y % BUFF_DIM_Y == 0);\n static_assert(CHUNK_DIM_X >= BUFF_DIM_X && CHUNK_DIM_X % BUFF_DIM_X == 0);\n\n constexpr size_t STAGES_Y = CHUNK_DIM_Y / BUFF_DIM_Y;\n constexpr size_t STAGES_X = CHUNK_DIM_X / BUFF_DIM_X;\n\n constexpr int kNumWarps = (THREADS_PER_CHUNK * THREADS_PER_Y) / kThreadsPerWarp;\n\n const int input_block_offset_Y = blockIdx.y * CHUNK_DIM_Y;\n const int input_block_offset_X = blockIdx.x * CHUNK_DIM_X;\n\n extern __shared__ __align__(128) char dynamic_shmem[];\n uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem);\n // Manually align dynamic SHMEM per TMA requirements using padding\n // __align__(128) Does not guarantee the pointer to be aligned!\n uint8_t* dshmem = reinterpret_cast((base_shmem_ptr + 127) & ~127ULL);\n\n // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned\n constexpr size_t in_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType);\n IType* in_sh_0 = reinterpret_cast(dshmem);\n dshmem += in_buff_size;\n IType* in_sh_1 = reinterpret_cast(dshmem);\n dshmem += in_buff_size;\n\n IType* in_shs[2] = {in_sh_0, in_sh_1};\n\n constexpr int shmem_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType);\n\n const bool is_master_thread = (threadIdx.x == 0 && threadIdx.y == 0);\n\n // Initialize shared memory barrier with the number of threads participating in the barrier.\n#pragma nv_diag_suppress static_var_with_dynamic_init\n uint64_t* mbar = reinterpret_cast(dshmem);\n dshmem += sizeof(uint64_t) * (STAGES_X * STAGES_Y);\n\n float* max_staging_identity = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n float* max_staging_transpose = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n float* max_staging_pre_rht = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n\n initialize_barriers(mbar,\n is_master_thread);\n\n copy_2d_to_shared(in_shs[0], reinterpret_cast(&tensor_map_input),\n input_block_offset_X, input_block_offset_Y, shmem_buff_size, &mbar[0],\n is_master_thread);\n\n uint32_t had_frag_i[4];\n uint32_t had_frag_t[4];\n get_hadamard_matrix_fragment(\n had_frag_i, random_sign_mask, had_frag_t, random_sign_mask_t);\n\n float local_pre_rht_amax = 0.0;\n float local_amax = 0.0;\n float local_amax_t = 0.0;\n uint32_t local_pre_rht_amax_reg = *reinterpret_cast(&local_pre_rht_amax);\n uint32_t local_amax_reg = *reinterpret_cast(&local_amax);\n uint32_t local_amax_t_reg = *reinterpret_cast(&local_amax_t);\n\n const int warp_id = threadIdx.x / kThreadsPerWarp;\n const int local_rank = threadIdx.x % kThreadsPerWarp;\n const int ld_row_idx = local_rank % kHadamardDimension;\n const int ld_col_idx = local_rank / kHadamardDimension + warp_id * 2;\n const int swizzle_idx = swizzle_128B_atom_32B(ld_row_idx, ld_col_idx);\n\n for (int stage_y = 0; stage_y < STAGES_Y; ++stage_y) {\n for (int stage_x = 0; stage_x < STAGES_X; ++stage_x) {\n int stage = STAGES_X * stage_y + stage_x;\n\n const int next_stage = stage + 1;\n const int next_stage_x = stage_x + 1 == STAGES_X ? 0 : stage_x + 1;\n const int next_stage_y = stage_x + 1 == STAGES_X ? stage_y + 1 : stage_y;\n\n if (next_stage < STAGES_X * STAGES_Y) {\n const int input_global_offset_Y = input_block_offset_Y + next_stage_y * BUFF_DIM_Y;\n const int input_global_offset_X = input_block_offset_X + next_stage_x * BUFF_DIM_X;\n\n copy_2d_to_shared(in_shs[next_stage % 2], // ping-pong\n reinterpret_cast(&tensor_map_input), input_global_offset_X,\n input_global_offset_Y, shmem_buff_size, &mbar[next_stage],\n is_master_thread);\n }\n\n // Wait for the data to have arrived\n ptx::mbarrier_wait_parity(&mbar[stage], 0);\n\n const size_t compute_stage_x_num =\n BUFF_DIM_X / (kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp));\n const size_t compute_stage_y_num = BUFF_DIM_Y / (kHadamardDimension * THREADS_PER_Y);\n\n const size_t in_row_stride = BUFF_DIM_X;\n\n IType* in_sh_ptr = in_shs[stage % 2];\n\n#pragma unroll\n for (size_t compute_stage_y = 0; compute_stage_y < compute_stage_y_num; compute_stage_y++) {\n const int row_idx_offset = (compute_stage_y * kHadamardDimension * THREADS_PER_Y +\n threadIdx.y * kHadamardDimension);\n const int in_row_offset = row_idx_offset * in_row_stride;\n\n#pragma unroll\n for (size_t compute_stage_x = 0; compute_stage_x < compute_stage_x_num; compute_stage_x++) {\n ComputeKernel(\n had_frag_i, had_frag_t,\n in_sh_ptr + in_row_offset +\n (compute_stage_x * kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp)),\n swizzle_idx, local_pre_rht_amax_reg, local_amax_reg, local_amax_t_reg);\n }\n }\n // Ensure all threads have finished their computation before new data over-writes the shared\n // memory.\n __syncthreads();\n\n // Ensure generic shared-memory accesses are visible before the next TMA write.\n ptx::fence_proxy_async_shared_cta();\n }\n }\n\n const int warpid = (threadIdx.x + threadIdx.y * blockDim.x) / kThreadsPerWarp;\n\n if constexpr (kReturnPreRhtAmax) {\n unpack_max_of_packed_bf16(local_pre_rht_amax_reg, local_pre_rht_amax);\n }\n if constexpr (kReturnIdentityAmax) {\n unpack_max_of_packed_bf16(local_amax_reg, local_amax);\n }\n if constexpr (kReturnTransposedAmax) {\n unpack_max_of_packed_bf16(local_amax_t_reg, local_amax_t);\n }\n\n ReduceMax(\n local_pre_rht_amax, local_amax, local_amax_t, max_staging_pre_rht, max_staging_identity,\n max_staging_transpose, output_pre_rht_amax_ptr, output_identity_amax_ptr,\n output_transpose_amax_ptr, warpid);\n\n destroy_barriers(mbar, is_master_thread);\n#else\n NVTE_DEVICE_ERROR(\"Kernel is only supported on SM 10.0+.\");\n#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n}\n\n} // namespace\n\n// broadcast_pre_rht_amax: when it's true, hadamard transform will be disabled\n// if at this time, the amax buffers for output expects both amax_rowwise and amax_colwise\n// then call MultiAmaxMemcpyD2DKernelPreRHT to D2D copy the amax values\nvoid group_hadamard_transform_amax(const Tensor& input_, std::vector& output_list,\n const size_t* split_sections, size_t num_tensors,\n uint16_t random_sign_mask, uint16_t random_sign_mask_t,\n bool broadcast_pre_rht_amax, cudaStream_t stream) {\n NVTE_API_CALL(group_hadamard_transform_amax);\n#if CUDA_VERSION >= 12080\n\n // Check input tensor\n NVTE_CHECK(input_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor must be BF16 tensor, but scaling mode is \",\n to_string(input_.scaling_mode), \".\");\n NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16,\n \"Input tensor must be BF16 tensor, but dtype is \", to_string(input_.dtype()), \".\");\n NVTE_CHECK(input_.dim() >= 2, \"Input must be a 2D tensor.\");\n const SimpleTensor& input = input_.data;\n\n // TODO: validate num_tensors and split_sections\n // assert if num_tensors is greater than kMaxTensorsPerKernel\n // will expand 64 to higher value if needed\n // if input size is going to exceed 4KB kernel launch limit, will then support multi-launch\n NVTE_CHECK(num_tensors <= kMaxTensorsPerKernel,\n \"Number of tensors should be less than or equal to \", kMaxTensorsPerKernel);\n\n // check split_sections\n // TODO: support m_splits_tensor for device initiated API\n NVTE_CHECK(split_sections != nullptr, \"split_sections should not be nullptr\");\n\n MultiAmaxArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.split_sections_range[0] = 0;\n bool all_return_pre_rht_amax = true;\n bool all_return_identity_amax = true;\n bool all_return_transposed_amax = true;\n for (size_t i = 0; i < num_tensors; ++i) {\n void* output_pre_rht_amax_ptr = output_list[i]->amax.dptr;\n // disable RHT(x) for now, only RHT_T(x) should be used\n void* output_identity_amax_ptr = nullptr;\n void* output_transpose_amax_ptr = output_list[i]->columnwise_amax.dptr;\n all_return_pre_rht_amax &= (output_pre_rht_amax_ptr != nullptr);\n all_return_identity_amax &= (output_identity_amax_ptr != nullptr);\n all_return_transposed_amax &= (output_transpose_amax_ptr != nullptr);\n // sanity check split_sections component to see if it's 64 multiple for each element\n NVTE_CHECK(split_sections[i] % 64 == 0, \"component \", i,\n \" of split_sections should be 64 multiple\");\n // also skip adding this tensor to the kernel args there are zero elements in this split\n if (split_sections[i] == 0) {\n continue;\n }\n // fill in kernel arguments\n kernel_args.output_pre_rht_amax_list[kernel_args.num_tensors] = output_pre_rht_amax_ptr;\n kernel_args.output_identity_amax_list[kernel_args.num_tensors] = output_identity_amax_ptr;\n kernel_args.output_transpose_amax_list[kernel_args.num_tensors] = output_transpose_amax_ptr;\n kernel_args.split_sections_range[kernel_args.num_tensors + 1] =\n kernel_args.split_sections_range[kernel_args.num_tensors] + split_sections[i];\n kernel_args.num_tensors++;\n }\n\n NVTE_CHECK(all_return_pre_rht_amax || all_return_identity_amax || all_return_transposed_amax,\n \"At least one of return_pre_rht_amax, return_identity_amax, or return_transposed_amax \"\n \"must be true\");\n // currently we haven't supported all_return_identity_amax, assert error if it's mistakenly enabled\n NVTE_CHECK(!all_return_identity_amax,\n \"Currently RHT transform should only be applied to transposed input\");\n\n if (broadcast_pre_rht_amax) {\n NVTE_CHECK(all_return_pre_rht_amax,\n \"broadcast_pre_rht_amax is only supported when we compute pre-RHT amax\");\n // if all_return_identity_amax and all_return_transposed_amax both are false, there is no need to broadcast anything\n broadcast_pre_rht_amax &= (all_return_identity_amax || all_return_transposed_amax);\n }\n\n // Multi zero out multiple amaxes if needed\n // Currently don't support multi-launch when num_tensors is larger than kMaxTensorsPerKernel\n // let the number of threads equal to number of tensors, use 1 block, kMaxTensorsPerKernel threads per block\n dim3 block_setup_amax(kMaxTensorsPerKernel);\n dim3 grid_setup_amax(1);\n MultiZeroAmaxKernel<<>>(kernel_args);\n NVTE_CHECK_CUDA(cudaGetLastError());\n\n checkCuDriverContext(stream);\n\n using IType = bf16;\n\n const size_t ndim = input.shape.size();\n const size_t row_length = input.shape[ndim - 1];\n size_t num_rows = 1;\n for (size_t i = 0; i < ndim - 1; ++i) {\n num_rows *= input.shape[i];\n }\n\n constexpr int kHadamardDimension = 16;\n NVTE_CHECK(row_length % kHadamardDimension == 0,\n \"row_length must be divisible by hadamard_dimension.\");\n NVTE_CHECK(num_rows % kHadamardDimension == 0,\n \"num_rows must be divisible by hadamard_dimension\");\n\n // four (1x4) 64x64 sub-tiles for ping-pong overlap\n constexpr uint64_t kChunkBlockXSmall = 256;\n constexpr uint64_t kChunkBlockYSmall = 64;\n constexpr uint64_t kBuffDimX = 64;\n constexpr uint64_t kBuffDimY = 64;\n\n alignas(64) CUtensorMap tensor_map_input{};\n\n create_2D_tensor_map(\n /*tensorMap=*/tensor_map_input,\n /*tensor=*/input,\n /*globalY=*/num_rows,\n /*globalX=*/row_length,\n /*shmemY=*/kBuffDimY,\n /*shmemX=*/kBuffDimX,\n /*stride_elems=*/row_length,\n /*offset_elems=*/0,\n /*type_num_bits=*/sizeof(IType) * 8,\n /*swizzle=*/CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B);\n\n constexpr uint64_t kThreadBlockX = 4;\n constexpr uint64_t kThreadBlockY = 1;\n constexpr uint64_t kNumWarps = kThreadBlockX * kThreadBlockY;\n\n dim3 block(kThreadBlockX * kThreadsPerWarp, kThreadBlockY);\n\n dim3 grid(DIVUP(row_length, kChunkBlockXSmall), DIVUP(num_rows, kChunkBlockYSmall));\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n (all_return_transposed_amax && !broadcast_pre_rht_amax), kReturnTransposedAmax,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n (all_return_identity_amax && !broadcast_pre_rht_amax), kReturnIdentityAmax,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n all_return_pre_rht_amax, kReturnPreRhtAmax,\n\n // *2 for ping-pong\n size_t in_sh_size = kBuffDimX * kBuffDimY * 2 * sizeof(IType);\n size_t mbar_size = sizeof(uint64_t) * (kChunkBlockXSmall / kBuffDimX) *\n (kChunkBlockYSmall / kBuffDimY);\n size_t shmem_bytes = in_sh_size + mbar_size + kNumWarps * sizeof(float) * 3;\n // Add padding in case shmem ptr is not aligned to 128 bytes.\n shmem_bytes = (shmem_bytes + 128);\n\n auto kernel = GroupHadamardAmaxTmaKernel<\n IType, kHadamardDimension, kChunkBlockYSmall, kChunkBlockXSmall, kBuffDimY,\n kBuffDimX, kThreadBlockX * kThreadsPerWarp, kThreadBlockY, kReturnPreRhtAmax,\n kReturnIdentityAmax, kReturnTransposedAmax>;\n cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,\n shmem_bytes);\n\n kernel<<>>(tensor_map_input, kernel_args,\n random_sign_mask, random_sign_mask_t,\n num_rows, row_length);\n if (broadcast_pre_rht_amax) {\n MultiAmaxMemcpyD2DKernelPreRHT<<>>(\n kernel_args);\n })));\n\n NVTE_CHECK_CUDA(cudaGetLastError());\n#else\n NVTE_ERROR(\"Hadamard transform requires CUDA 12.8+, but compile-time CUDA version is \",\n CUDA_VERSION);\n#endif // CUDA_VERSION >= 12080\n}\n\n} // namespace transformer_engine\n\n// Naming convention: \"Group\" kernels here means contiguous input concatenated\n// While \"Multi\" kernels are processing a list of pointers, like the zero amax kernel\n\n// Group hadamard transform API is unlike other multi-input & multi-output APIs\n// Group hadamard transform will take in a single input tensor, and directly calculate amax\n// with optional RHT transform. That's because we can assume the input tensor list to be\n// contiguous in memory, so the tensors are only splitted in dimension 0.\n// RHT transform is 16x16, so as long as each split of the input has 16 multiple shape\n// in dimension 0, we can treat the entire input as a single tensor.\n// Although mathmatically 16 multple is enough for this function to be correct,\n// for this kernel, we required 64 multiple of 16 in dimension 0 for better performance.\nvoid nvte_group_hadamard_transform_amax(const NVTETensor input, NVTETensor* outputs,\n const size_t* split_sections, size_t num_tensors,\n int random_sign_mask, int random_sign_mask_t,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_group_hadamard_transform_amax);\n using namespace transformer_engine;\n if (num_tensors == 0) {\n return;\n }\n\n Tensor* input_tensor = convertNVTETensorCheck(input);\n std::vector output_list(num_tensors);\n for (size_t i = 0; i < num_tensors; ++i) {\n output_list[i] = convertNVTETensorCheck(outputs[i]);\n }\n // Call the group tensor Hadamard transform amax implementation.\n group_hadamard_transform_amax(*input_tensor, output_list, split_sections, num_tensors,\n static_cast(random_sign_mask),\n static_cast(random_sign_mask_t), false, stream);\n}\n\n// Grouped-tensor amax without doing hadamard transform\nvoid nvte_group_amax(const NVTETensor input, NVTETensor* outputs, const size_t* split_sections,\n size_t num_tensors, cudaStream_t stream) {\n NVTE_API_CALL(nvte_group_amax);\n using namespace transformer_engine;\n if (num_tensors == 0) {\n return;\n }\n\n Tensor* input_tensor = convertNVTETensorCheck(input);\n std::vector output_list(num_tensors);\n for (size_t i = 0; i < num_tensors; ++i) {\n output_list[i] = convertNVTETensorCheck(outputs[i]);\n }\n\n group_hadamard_transform_amax(*input_tensor, output_list, split_sections, num_tensors, 0, 0, true,\n stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \"common/common.h\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"hadamard_transform_utils.cuh\"\n\nnamespace transformer_engine {\nnamespace {\n\nconstexpr int kMaxTensorsPerKernel = 64; // Args must be <4 KB, expand 64 if needed\nstruct MultiAmaxArgs {\n // (output) Amax buffer for pre-RHT amax buffer\n void* output_pre_rht_amax_list[kMaxTensorsPerKernel];\n // (output) Amax buffer for RHT identity amax buffer\n void* output_identity_amax_list[kMaxTensorsPerKernel];\n // (output) Amax buffer for RHT transpose amax buffer\n void* output_transpose_amax_list[kMaxTensorsPerKernel];\n // Prefix sum (with leading zero) of split_sections of each tensor of input\n int split_sections_range[kMaxTensorsPerKernel + 1];\n // Number of tensors (splits) being processed by kernel\n int num_tensors;\n};\n\nconstexpr int kThreadsPerWarp = 32;\n\ntemplate \n__device__ __forceinline__ void ComputeKernel(uint32_t b_frag_i[4], uint32_t b_frag_t[4],\n IType* in_sh_ptr, int swizzle_idx,\n uint32_t& local_pre_rht_amax_reg,\n uint32_t& local_amax_reg,\n uint32_t& local_amax_t_reg) {\n uint32_t a_frag[4]; // A matrix fragment\n uint32_t c_frag[4]; // Result fragment\n\n uint32_t temp_amax_reg;\n uint32_t temp_amax_t_reg;\n\n if (kReturnIdentityAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0], a_frag[1], a_frag[2], a_frag[3], b_frag_i[0], b_frag_i[1], b_frag_i[2],\n b_frag_i[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_reg);\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_amax_reg)\n : \"r\"(local_amax_reg), \"r\"(temp_amax_reg));\n }\n\n if (kReturnTransposedAmax) {\n if (!kReturnIdentityAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n } else {\n matrix_transpose_m8_n8_b16_inplace(a_frag[0]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[1]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[2]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[3]);\n }\n\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0], a_frag[2], a_frag[1], a_frag[3], b_frag_t[0], b_frag_t[1], b_frag_t[2],\n b_frag_t[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_t_reg);\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_amax_t_reg)\n : \"r\"(local_amax_t_reg), \"r\"(temp_amax_t_reg));\n }\n\n if (kReturnPreRhtAmax) {\n if (!kReturnIdentityAmax && !kReturnTransposedAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n }\n\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[0])\n : \"r\"(a_frag[0]), \"r\"(a_frag[1]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[2])\n : \"r\"(a_frag[2]), \"r\"(a_frag[3]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[0])\n : \"r\"(a_frag[0]), \"r\"(a_frag[2]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_pre_rht_amax_reg)\n : \"r\"(a_frag[0]), \"r\"(local_pre_rht_amax_reg));\n }\n}\n\ntemplate \n__device__ __host__ constexpr int NextPowerOf2() {\n static_assert(kN > 0, \"kN must be > 0\");\n // Round up to the next power of 2 by counting leading zeros.\n return 1 << (32 - __builtin_clz(kN - 1));\n}\n\ntemplate \n__device__ __forceinline__ void ReduceMax(const float pre_rht_amax, const float identity_amax,\n const float transpose_amax, float* staging_for_pre_rht,\n float* staging_for_identity, float* staging_for_transpose,\n float* output_pre_rht_amax_ptr,\n float* output_identity_amax_ptr,\n float* output_transpose_amax_ptr, const int warpid) {\n // intra-warp reduction\n constexpr int kWarpSize = 32;\n int local_rank = threadIdx.x % 32;\n float warp_pre_rht_amax = kReturnPreRhtAmax ? warp_reduce_max(pre_rht_amax) : 0.0f;\n float warp_identity_amax = kReturnIdentityAmax ? warp_reduce_max(identity_amax) : 0.0f;\n float warp_transpose_amax =\n kReturnTransposedAmax ? warp_reduce_max(transpose_amax) : 0.0f;\n\n // inter-warp reduction\n if (threadIdx.x % 32 == 0) {\n if (kReturnPreRhtAmax) {\n staging_for_pre_rht[warpid] = warp_pre_rht_amax;\n }\n if (kReturnIdentityAmax) {\n staging_for_identity[warpid] = warp_identity_amax;\n }\n if (kReturnTransposedAmax) {\n staging_for_transpose[warpid] = warp_transpose_amax;\n }\n }\n __syncthreads();\n constexpr int kNumWarpsPow2 = NextPowerOf2();\n if (warpid == 0) {\n if (kReturnIdentityAmax) {\n float identity_accum = local_rank < kNumWarps ? staging_for_identity[local_rank] : 0.0f;\n identity_accum = warp_reduce_max(identity_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_identity_amax_ptr, identity_accum);\n }\n }\n }\n if (warpid == 1) {\n if (kReturnTransposedAmax) {\n float transpose_accum = local_rank < kNumWarps ? staging_for_transpose[local_rank] : 0.0f;\n transpose_accum = warp_reduce_max(transpose_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_transpose_amax_ptr, transpose_accum);\n }\n }\n }\n if (warpid == 2) {\n if (kReturnPreRhtAmax) {\n float pre_rht_accum = local_rank < kNumWarps ? staging_for_pre_rht[local_rank] : 0.0f;\n pre_rht_accum = warp_reduce_max(pre_rht_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_pre_rht_amax_ptr, pre_rht_accum);\n }\n }\n }\n}\n\n// args: the mult-tensor amax arguments\n__global__ void MultiZeroAmaxKernel(MultiAmaxArgs args) {\n int num_tensors = args.num_tensors;\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n\n for (; tid < num_tensors; tid += stride) {\n float* output_pre_rht_amax_ptr = static_cast(args.output_pre_rht_amax_list[tid]);\n float* output_identity_amax_ptr = static_cast(args.output_identity_amax_list[tid]);\n float* output_transpose_amax_ptr = static_cast(args.output_transpose_amax_list[tid]);\n if (output_pre_rht_amax_ptr != nullptr) {\n *output_pre_rht_amax_ptr = 0;\n }\n if (output_identity_amax_ptr != nullptr) {\n *output_identity_amax_ptr = 0;\n }\n if (output_transpose_amax_ptr != nullptr) {\n *output_transpose_amax_ptr = 0;\n }\n }\n}\n\n// args: the mult-tensor amax arguments\n__global__ void MultiAmaxMemcpyD2DKernelPreRHT(MultiAmaxArgs args) {\n int num_tensors = args.num_tensors;\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n\n for (; tid < num_tensors; tid += stride) {\n float* output_pre_rht_amax_ptr = static_cast(args.output_pre_rht_amax_list[tid]);\n float* output_identity_amax_ptr = static_cast(args.output_identity_amax_list[tid]);\n float* output_transpose_amax_ptr = static_cast(args.output_transpose_amax_list[tid]);\n if (output_pre_rht_amax_ptr != nullptr) {\n float pre_rht_amax = *output_pre_rht_amax_ptr;\n if (output_identity_amax_ptr != nullptr) {\n *output_identity_amax_ptr = pre_rht_amax;\n }\n if (output_transpose_amax_ptr != nullptr) {\n *output_transpose_amax_ptr = pre_rht_amax;\n }\n }\n }\n}\n\ntemplate \n__global__ void GroupHadamardAmaxTmaKernel(const __grid_constant__ CUtensorMap tensor_map_input,\n const MultiAmaxArgs args, uint16_t random_sign_mask,\n uint16_t random_sign_mask_t, uint64_t num_rows,\n uint64_t row_length) {\n#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n\n float* output_pre_rht_amax_ptr;\n float* output_identity_amax_ptr;\n float* output_transpose_amax_ptr;\n\n // calculate the global offset in Y direction to access the correct amax buffer\n int global_offset_y = blockIdx.y * CHUNK_DIM_Y;\n int tensor_id = 0;\n while (args.split_sections_range[tensor_id + 1] <= global_offset_y) {\n ++tensor_id;\n }\n output_pre_rht_amax_ptr = static_cast(args.output_pre_rht_amax_list[tensor_id]);\n output_identity_amax_ptr = static_cast(args.output_identity_amax_list[tensor_id]);\n output_transpose_amax_ptr = static_cast(args.output_transpose_amax_list[tensor_id]);\n\n static_assert(CHUNK_DIM_Y >= BUFF_DIM_Y && CHUNK_DIM_Y % BUFF_DIM_Y == 0);\n static_assert(CHUNK_DIM_X >= BUFF_DIM_X && CHUNK_DIM_X % BUFF_DIM_X == 0);\n\n constexpr size_t STAGES_Y = CHUNK_DIM_Y / BUFF_DIM_Y;\n constexpr size_t STAGES_X = CHUNK_DIM_X / BUFF_DIM_X;\n\n constexpr int kNumWarps = (THREADS_PER_CHUNK * THREADS_PER_Y) / kThreadsPerWarp;\n\n const int input_block_offset_Y = blockIdx.y * CHUNK_DIM_Y;\n const int input_block_offset_X = blockIdx.x * CHUNK_DIM_X;\n\n extern __shared__ __align__(128) char dynamic_shmem[];\n uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem);\n // Manually align dynamic SHMEM per TMA requirements using padding\n // __align__(128) Does not guarantee the pointer to be aligned!\n uint8_t* dshmem = reinterpret_cast((base_shmem_ptr + 127) & ~127ULL);\n\n // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned\n constexpr size_t in_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType);\n IType* in_sh_0 = reinterpret_cast(dshmem);\n dshmem += in_buff_size;\n IType* in_sh_1 = reinterpret_cast(dshmem);\n dshmem += in_buff_size;\n\n IType* in_shs[2] = {in_sh_0, in_sh_1};\n\n constexpr int shmem_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType);\n\n const bool is_master_thread = (threadIdx.x == 0 && threadIdx.y == 0);\n\n // Initialize shared memory barrier with the number of threads participating in the barrier.\n#pragma nv_diag_suppress static_var_with_dynamic_init\n uint64_t* mbar = reinterpret_cast(dshmem);\n dshmem += sizeof(uint64_t) * (STAGES_X * STAGES_Y);\n\n float* max_staging_identity = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n float* max_staging_transpose = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n float* max_staging_pre_rht = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n\n initialize_barriers(mbar,\n is_master_thread);\n\n copy_2d_to_shared(in_shs[0], reinterpret_cast(&tensor_map_input),\n input_block_offset_X, input_block_offset_Y, shmem_buff_size, &mbar[0],\n is_master_thread);\n\n uint32_t had_frag_i[4];\n uint32_t had_frag_t[4];\n get_hadamard_matrix_fragment(\n had_frag_i, random_sign_mask, had_frag_t, random_sign_mask_t);\n\n float local_pre_rht_amax = 0.0;\n float local_amax = 0.0;\n float local_amax_t = 0.0;\n uint32_t local_pre_rht_amax_reg = *reinterpret_cast(&local_pre_rht_amax);\n uint32_t local_amax_reg = *reinterpret_cast(&local_amax);\n uint32_t local_amax_t_reg = *reinterpret_cast(&local_amax_t);\n\n const int warp_id = threadIdx.x / kThreadsPerWarp;\n const int local_rank = threadIdx.x % kThreadsPerWarp;\n const int ld_row_idx = local_rank % kHadamardDimension;\n const int ld_col_idx = local_rank / kHadamardDimension + warp_id * 2;\n const int swizzle_idx = swizzle_128B_atom_32B(ld_row_idx, ld_col_idx);\n\n for (int stage_y = 0; stage_y < STAGES_Y; ++stage_y) {\n for (int stage_x = 0; stage_x < STAGES_X; ++stage_x) {\n int stage = STAGES_X * stage_y + stage_x;\n\n const int next_stage = stage + 1;\n const int next_stage_x = stage_x + 1 == STAGES_X ? 0 : stage_x + 1;\n const int next_stage_y = stage_x + 1 == STAGES_X ? stage_y + 1 : stage_y;\n\n if (next_stage < STAGES_X * STAGES_Y) {\n const int input_global_offset_Y = input_block_offset_Y + next_stage_y * BUFF_DIM_Y;\n const int input_global_offset_X = input_block_offset_X + next_stage_x * BUFF_DIM_X;\n\n copy_2d_to_shared(in_shs[next_stage % 2], // ping-pong\n reinterpret_cast(&tensor_map_input), input_global_offset_X,\n input_global_offset_Y, shmem_buff_size, &mbar[next_stage],\n is_master_thread);\n }\n\n // Wait for the data to have arrived\n ptx::mbarrier_wait_parity(&mbar[stage], 0);\n\n const size_t compute_stage_x_num =\n BUFF_DIM_X / (kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp));\n const size_t compute_stage_y_num = BUFF_DIM_Y / (kHadamardDimension * THREADS_PER_Y);\n\n const size_t in_row_stride = BUFF_DIM_X;\n\n IType* in_sh_ptr = in_shs[stage % 2];\n\n#pragma unroll\n for (size_t compute_stage_y = 0; compute_stage_y < compute_stage_y_num; compute_stage_y++) {\n const int row_idx_offset = (compute_stage_y * kHadamardDimension * THREADS_PER_Y +\n threadIdx.y * kHadamardDimension);\n const int in_row_offset = row_idx_offset * in_row_stride;\n\n#pragma unroll\n for (size_t compute_stage_x = 0; compute_stage_x < compute_stage_x_num; compute_stage_x++) {\n ComputeKernel(\n had_frag_i, had_frag_t,\n in_sh_ptr + in_row_offset +\n (compute_stage_x * kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp)),\n swizzle_idx, local_pre_rht_amax_reg, local_amax_reg, local_amax_t_reg);\n }\n }\n // Ensure all threads have finished their computation before new data over-writes the shared\n // memory.\n __syncthreads();\n\n // Ensure generic shared-memory accesses are visible before the next TMA write.\n ptx::fence_proxy_async_shared_cta();\n }\n }\n\n const int warpid = (threadIdx.x + threadIdx.y * blockDim.x) / kThreadsPerWarp;\n\n if constexpr (kReturnPreRhtAmax) {\n unpack_max_of_packed_bf16(local_pre_rht_amax_reg, local_pre_rht_amax);\n }\n if constexpr (kReturnIdentityAmax) {\n unpack_max_of_packed_bf16(local_amax_reg, local_amax);\n }\n if constexpr (kReturnTransposedAmax) {\n unpack_max_of_packed_bf16(local_amax_t_reg, local_amax_t);\n }\n\n ReduceMax(\n local_pre_rht_amax, local_amax, local_amax_t, max_staging_pre_rht, max_staging_identity,\n max_staging_transpose, output_pre_rht_amax_ptr, output_identity_amax_ptr,\n output_transpose_amax_ptr, warpid);\n\n destroy_barriers(mbar, is_master_thread);\n#else\n NVTE_DEVICE_ERROR(\"Kernel is only supported on SM 10.0+.\");\n#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n}\n\n} // namespace\n\n// broadcast_pre_rht_amax: when it's true, hadamard transform will be disabled\n// if at this time, the amax buffers for output expects both amax_rowwise and amax_colwise\n// then call MultiAmaxMemcpyD2DKernelPreRHT to D2D copy the amax values\nvoid group_hadamard_transform_amax(const Tensor& input_, std::vector& output_list,\n const size_t* split_sections, size_t num_tensors,\n uint16_t random_sign_mask, uint16_t random_sign_mask_t,\n bool broadcast_pre_rht_amax, hipStream_t stream) {\n NVTE_API_CALL(group_hadamard_transform_amax);\n#if CUDA_VERSION >= 12080\n\n // Check input tensor\n NVTE_CHECK(input_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor must be BF16 tensor, but scaling mode is \",\n to_string(input_.scaling_mode), \".\");\n NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16,\n \"Input tensor must be BF16 tensor, but dtype is \", to_string(input_.dtype()), \".\");\n NVTE_CHECK(input_.dim() >= 2, \"Input must be a 2D tensor.\");\n const SimpleTensor& input = input_.data;\n\n // TODO: validate num_tensors and split_sections\n // assert if num_tensors is greater than kMaxTensorsPerKernel\n // will expand 64 to higher value if needed\n // if input size is going to exceed 4KB kernel launch limit, will then support multi-launch\n NVTE_CHECK(num_tensors <= kMaxTensorsPerKernel,\n \"Number of tensors should be less than or equal to \", kMaxTensorsPerKernel);\n\n // check split_sections\n // TODO: support m_splits_tensor for device initiated API\n NVTE_CHECK(split_sections != nullptr, \"split_sections should not be nullptr\");\n\n MultiAmaxArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.split_sections_range[0] = 0;\n bool all_return_pre_rht_amax = true;\n bool all_return_identity_amax = true;\n bool all_return_transposed_amax = true;\n for (size_t i = 0; i < num_tensors; ++i) {\n void* output_pre_rht_amax_ptr = output_list[i]->amax.dptr;\n // disable RHT(x) for now, only RHT_T(x) should be used\n void* output_identity_amax_ptr = nullptr;\n void* output_transpose_amax_ptr = output_list[i]->columnwise_amax.dptr;\n all_return_pre_rht_amax &= (output_pre_rht_amax_ptr != nullptr);\n all_return_identity_amax &= (output_identity_amax_ptr != nullptr);\n all_return_transposed_amax &= (output_transpose_amax_ptr != nullptr);\n // sanity check split_sections component to see if it's 64 multiple for each element\n NVTE_CHECK(split_sections[i] % 64 == 0, \"component \", i,\n \" of split_sections should be 64 multiple\");\n // also skip adding this tensor to the kernel args there are zero elements in this split\n if (split_sections[i] == 0) {\n continue;\n }\n // fill in kernel arguments\n kernel_args.output_pre_rht_amax_list[kernel_args.num_tensors] = output_pre_rht_amax_ptr;\n kernel_args.output_identity_amax_list[kernel_args.num_tensors] = output_identity_amax_ptr;\n kernel_args.output_transpose_amax_list[kernel_args.num_tensors] = output_transpose_amax_ptr;\n kernel_args.split_sections_range[kernel_args.num_tensors + 1] =\n kernel_args.split_sections_range[kernel_args.num_tensors] + split_sections[i];\n kernel_args.num_tensors++;\n }\n\n NVTE_CHECK(all_return_pre_rht_amax || all_return_identity_amax || all_return_transposed_amax,\n \"At least one of return_pre_rht_amax, return_identity_amax, or return_transposed_amax \"\n \"must be true\");\n // currently we haven't supported all_return_identity_amax, assert error if it's mistakenly enabled\n NVTE_CHECK(!all_return_identity_amax,\n \"Currently RHT transform should only be applied to transposed input\");\n\n if (broadcast_pre_rht_amax) {\n NVTE_CHECK(all_return_pre_rht_amax,\n \"broadcast_pre_rht_amax is only supported when we compute pre-RHT amax\");\n // if all_return_identity_amax and all_return_transposed_amax both are false, there is no need to broadcast anything\n broadcast_pre_rht_amax &= (all_return_identity_amax || all_return_transposed_amax);\n }\n\n // Multi zero out multiple amaxes if needed\n // Currently don't support multi-launch when num_tensors is larger than kMaxTensorsPerKernel\n // let the number of threads equal to number of tensors, use 1 block, kMaxTensorsPerKernel threads per block\n dim3 block_setup_amax(kMaxTensorsPerKernel);\n dim3 grid_setup_amax(1);\n MultiZeroAmaxKernel<<>>(kernel_args);\n NVTE_CHECK_CUDA(hipGetLastError());\n\n checkCuDriverContext(stream);\n\n using IType = bf16;\n\n const size_t ndim = input.shape.size();\n const size_t row_length = input.shape[ndim - 1];\n size_t num_rows = 1;\n for (size_t i = 0; i < ndim - 1; ++i) {\n num_rows *= input.shape[i];\n }\n\n constexpr int kHadamardDimension = 16;\n NVTE_CHECK(row_length % kHadamardDimension == 0,\n \"row_length must be divisible by hadamard_dimension.\");\n NVTE_CHECK(num_rows % kHadamardDimension == 0,\n \"num_rows must be divisible by hadamard_dimension\");\n\n // four (1x4) 64x64 sub-tiles for ping-pong overlap\n constexpr uint64_t kChunkBlockXSmall = 256;\n constexpr uint64_t kChunkBlockYSmall = 64;\n constexpr uint64_t kBuffDimX = 64;\n constexpr uint64_t kBuffDimY = 64;\n\n alignas(64) CUtensorMap tensor_map_input{};\n\n create_2D_tensor_map(\n /*tensorMap=*/tensor_map_input,\n /*tensor=*/input,\n /*globalY=*/num_rows,\n /*globalX=*/row_length,\n /*shmemY=*/kBuffDimY,\n /*shmemX=*/kBuffDimX,\n /*stride_elems=*/row_length,\n /*offset_elems=*/0,\n /*type_num_bits=*/sizeof(IType) * 8,\n /*swizzle=*/CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B);\n\n constexpr uint64_t kThreadBlockX = 4;\n constexpr uint64_t kThreadBlockY = 1;\n constexpr uint64_t kNumWarps = kThreadBlockX * kThreadBlockY;\n\n dim3 block(kThreadBlockX * kThreadsPerWarp, kThreadBlockY);\n\n dim3 grid(DIVUP(row_length, kChunkBlockXSmall), DIVUP(num_rows, kChunkBlockYSmall));\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n (all_return_transposed_amax && !broadcast_pre_rht_amax), kReturnTransposedAmax,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n (all_return_identity_amax && !broadcast_pre_rht_amax), kReturnIdentityAmax,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n all_return_pre_rht_amax, kReturnPreRhtAmax,\n\n // *2 for ping-pong\n size_t in_sh_size = kBuffDimX * kBuffDimY * 2 * sizeof(IType);\n size_t mbar_size = sizeof(uint64_t) * (kChunkBlockXSmall / kBuffDimX) *\n (kChunkBlockYSmall / kBuffDimY);\n size_t shmem_bytes = in_sh_size + mbar_size + kNumWarps * sizeof(float) * 3;\n // Add padding in case shmem ptr is not aligned to 128 bytes.\n shmem_bytes = (shmem_bytes + 128);\n\n auto kernel = GroupHadamardAmaxTmaKernel<\n IType, kHadamardDimension, kChunkBlockYSmall, kChunkBlockXSmall, kBuffDimY,\n kBuffDimX, kThreadBlockX * kThreadsPerWarp, kThreadBlockY, kReturnPreRhtAmax,\n kReturnIdentityAmax, kReturnTransposedAmax>;\n hipFuncSetAttribute(reinterpret_cast(kernel), hipFuncAttributeMaxDynamicSharedMemorySize,\n shmem_bytes);\n\n kernel<<>>(tensor_map_input, kernel_args,\n random_sign_mask, random_sign_mask_t,\n num_rows, row_length);\n if (broadcast_pre_rht_amax) {\n MultiAmaxMemcpyD2DKernelPreRHT<<>>(\n kernel_args);\n })));\n\n NVTE_CHECK_CUDA(hipGetLastError());\n#else\n NVTE_ERROR(\"Hadamard transform requires CUDA 12.8+, but compile-time CUDA version is \",\n CUDA_VERSION);\n#endif // CUDA_VERSION >= 12080\n}\n\n} // namespace transformer_engine\n\n// Naming convention: \"Group\" kernels here means contiguous input concatenated\n// While \"Multi\" kernels are processing a list of pointers, like the zero amax kernel\n\n// Group hadamard transform API is unlike other multi-input & multi-output APIs\n// Group hadamard transform will take in a single input tensor, and directly calculate amax\n// with optional RHT transform. That's because we can assume the input tensor list to be\n// contiguous in memory, so the tensors are only splitted in dimension 0.\n// RHT transform is 16x16, so as long as each split of the input has 16 multiple shape\n// in dimension 0, we can treat the entire input as a single tensor.\n// Although mathmatically 16 multple is enough for this function to be correct,\n// for this kernel, we required 64 multiple of 16 in dimension 0 for better performance.\nvoid nvte_group_hadamard_transform_amax(const NVTETensor input, NVTETensor* outputs,\n const size_t* split_sections, size_t num_tensors,\n int random_sign_mask, int random_sign_mask_t,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_group_hadamard_transform_amax);\n using namespace transformer_engine;\n if (num_tensors == 0) {\n return;\n }\n\n Tensor* input_tensor = convertNVTETensorCheck(input);\n std::vector output_list(num_tensors);\n for (size_t i = 0; i < num_tensors; ++i) {\n output_list[i] = convertNVTETensorCheck(outputs[i]);\n }\n // Call the group tensor Hadamard transform amax implementation.\n group_hadamard_transform_amax(*input_tensor, output_list, split_sections, num_tensors,\n static_cast(random_sign_mask),\n static_cast(random_sign_mask_t), false, stream);\n}\n\n// Grouped-tensor amax without doing hadamard transform\nvoid nvte_group_amax(const NVTETensor input, NVTETensor* outputs, const size_t* split_sections,\n size_t num_tensors, hipStream_t stream) {\n NVTE_API_CALL(nvte_group_amax);\n using namespace transformer_engine;\n if (num_tensors == 0) {\n return;\n }\n\n Tensor* input_tensor = convertNVTETensorCheck(input);\n std::vector output_list(num_tensors);\n for (size_t i = 0; i < num_tensors; ++i) {\n output_list[i] = convertNVTETensorCheck(outputs[i]);\n }\n\n group_hadamard_transform_amax(*input_tensor, output_list, split_sections, num_tensors, 0, 0, true,\n stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/util/cuda_runtime.h\"\n#include \"common/util/curanddx.hpp\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"customized_pipeline.cuh\"\n#include \"cutlass/arch/barrier.h\"\n#include \"cutlass/arch/reg_reconfig.h\"\n#include \"cutlass/cluster_launch.hpp\"\n#include \"cutlass/cutlass.h\"\n#include \"cutlass/detail/sm100_blockscaled_layout.hpp\"\n#include \"cutlass/fast_math.h\"\n#include \"cutlass/float8.h\"\n#include \"cutlass/float_subbyte.h\"\n#include \"cutlass/gemm/collective/builders/sm100_common.inl\"\n#include \"cutlass/numeric_conversion.h\"\n#include \"cutlass/numeric_types.h\"\n#include \"cutlass/pipeline/pipeline.hpp\"\n#include \"cutlass/platform/platform.h\"\n#include \"cutlass/util/GPU_Clock.hpp\"\n#include \"cutlass/util/command_line.h\"\n#include \"cutlass/util/print_error.hpp\"\n\n// clang-format off\n\nnamespace transformer_engine {\nnamespace detail {\nnamespace {\n\nusing namespace cute;\n\nstruct CLCResponse { uint32_t data[4] = {0}; };\n\nconstexpr int kFp4ConvertChunkElements = 8;\nconstexpr int kFp4ConvertFullElements = 16;\nconstexpr int kFp4RbitsPerChunk = 2;\nconstexpr int kFp4ChunkCount = kFp4ConvertFullElements / kFp4ConvertChunkElements;\n\n\nCUTLASS_DEVICE\ncutlass::Array StochasticNumericConverterBase(\n cutlass::Array const &input,\n cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n auto output_ptr = reinterpret_cast(&output);\n constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING;\n if constexpr (has_rs) {\n asm volatile(\n \"{\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\\n\"\n \"}\"\n : \"=h\"(output_ptr[0]), \"=h\"(output_ptr[1])\n : \"f\"(input[0]), \"f\"(input[1]), \"f\"(input[2]), \"f\"(input[3]), \"f\"(input[4]), \"f\"(input[5]),\n \"f\"(input[6]), \"f\"(input[7]), \"r\"(rbits[0]), \"r\"(rbits[1]));\n } else {\n NVTE_DEVICE_ERROR(\n \"FP4 cvt PTX instructions are architecture-specific. \"\n \"Try recompiling with sm_XXXa instead of sm_XXX.\");\n }\n return output;\n}\n\nCUTLASS_DEVICE\ncutlass::Array\nStochasticNumericConverter(cutlass::Array const &input,\n cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n cutlass::Array *result_ptr =\n reinterpret_cast *>(&output);\n cutlass::Array const *source_ptr =\n reinterpret_cast const *>(&input);\n cutlass::Array const *rbits_ptr =\n reinterpret_cast const *>(&rbits);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < kFp4ChunkCount; i++) {\n result_ptr[i] = StochasticNumericConverterBase(source_ptr[i], rbits_ptr[i]);\n }\n return output;\n}\n\ntemplate <\n class ElementA,\n class ElementB,\n class ASmemLayout,\n class BSmemLayout,\n class ClusterShape,\n int AccumulatorPipelineStageCount_,\n int EpilogueUnrollFactor_,\n int SchedulerPipelineStageCount_>\nstruct SharedStorage {\n static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;\n static int constexpr EpilogueUnrollFactor = EpilogueUnrollFactor_;\n using AtomThrShapeMNK = cute::Shape<_1, _1, _1>;\n\n using AccumulatorPipeline = cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;\n\n static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline = cutlass::detail::CustomizedPipelineTmaUmmaAsync<\n MainloopPipelineStageCount,\n Shape<_1,_1,_1>,\n AtomThrShapeMNK>;\n using MainloopPipelineStorage = typename MainloopPipeline::SharedStorage;\n using CLCPipeline = cutlass::PipelineCLCFetchAsync;\n using CLCPipelineStorage = typename CLCPipeline::SharedStorage;\n using CLCThrottlePipeline = cutlass::PipelineAsync;\n using CLCThrottlePipelineStorage = typename CLCThrottlePipeline::SharedStorage;\n\n struct TensorStorage : cute::aligned_struct<128, _1> {\n // cute::array_aligned> smem_A;\n cute::array_aligned> smem_A;\n cute::array_aligned> smem_B;\n } tensors;\n\n alignas(16) AccumulatorPipelineStorage accumulator;\n alignas(16) MainloopPipelineStorage mainloop;\n alignas(16) cute::uint64_t tma_barrier[1];\n alignas(16) CLCPipelineStorage clc;\n alignas(16) CLCThrottlePipelineStorage clc_throttle;\n alignas(16) CLCResponse clc_response[SchedulerPipelineStageCount_];\n uint32_t tmem_base_ptr;\n};\n\ntemplate \n__launch_bounds__(512, 1)\n__global__ static void row_col_rht_gemm_device(\n MShape M,\n NShape N,\n KShape K,\n ClusterShape cluster_shape,\n ClusterTileShape cluster_tile,\n TA const* A,\n AStride dA,\n ASmemLayout sAlayout,\n CUTE_GRID_CONSTANT TmaLoadA const tma_load_a,\n TB const* B,\n BStride dB,\n BSmemLayout sBlayout,\n CUTE_GRID_CONSTANT TmaLoadB const tma_load_b,\n TD* D,\n DStride dD,\n DSmemLayout,\n TSFD* SFD,\n TSFDLayout sfd_layout,\n TQA* QA,\n QAStride dQA,\n TSFA* SFA,\n TSFALayout sfa_layout,\n TiledMMA mma,\n float const* a_global_amax,\n float const* c_global_amax,\n const size_t* rng_state) {\n using namespace cute;\n\n // Abort immediately if compilation is not supported\n constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY;\n if constexpr (!is_blackwell_arch) {\n NVTE_DEVICE_ERROR(\"RHT fusion is only supported on Blackwell.\");\n return;\n } else {\n static_assert(kEnableRHTColQuant_ || kEnableRowQuant_,\n \"row_col_rht_gemm_device must generate row-wise \"\n \"and/or column-wise output.\");\n#if !defined(CUTLASS_ARCH_CLC_ENABLED)\n CUTLASS_NOT_IMPLEMENTED();\n return;\n#endif\n\n using X = Underscore;\n // static constexpr bool kApplyStochasticRounding = true;\n using ElementAccumulator = float;\n static int constexpr K_PIPE_MAX = size<3>(ASmemLayout{});\n using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>;\n static uint32_t constexpr kTmaTransactionBytes = cutlass::bits_to_bytes(\n size(AtomThrShapeMNK{}) * cosize(take<0,3>(ASmemLayout{})) * cute::sizeof_bits_v);\n static constexpr bool kEnableStochasticRounding = kEnableStochasticRounding_;\n static constexpr bool kEnableRHTColQuant = kEnableRHTColQuant_;\n static constexpr bool kEnableRowQuant = kEnableRowQuant_;\n static constexpr bool kUseFastMath = kUseFastMath_;\n static int constexpr RhtTensorSize = 16;\n static int constexpr kTmaRhtTensorTransactionBytes = cutlass::bits_to_bytes(\n RhtTensorSize * RhtTensorSize * cute::sizeof_bits_v);\n static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;\n static int constexpr SchedulerPipelineStageCount = SchedulerPipelineStageCount_;\n\n static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline = cutlass::detail::CustomizedPipelineTmaUmmaAsync<\n MainloopPipelineStageCount,\n ClusterShape,\n AtomThrShapeMNK>;\n using MainloopPipelineState = typename MainloopPipeline::PipelineState;\n using CLCPipeline = cutlass::PipelineCLCFetchAsync;\n using CLCPipelineState = typename CLCPipeline::PipelineState;\n using CLCThrottlePipeline = cutlass::PipelineAsync;\n using CLCThrottlePipelineState = typename CLCThrottlePipeline::PipelineState;\n\n static_assert(ClusterShape{} == Shape<_1,_1,_1>{}, \"ClusterShape must be Shape<_1,_1,_1>\");\n\n using TmemAllocator = cute::TMEM::Allocator1Sm;\n static int constexpr VectorSize = RhtTensorSize;\n // Preconditions\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n auto cluster_size = size<0>(cluster_shape);\n auto mainloop_tiler = Shape<_128,_16,_128>{};\n auto epilogue_tiler = Shape<_128,_128,_128>{};\n\n static int constexpr EpilogueUnrollFactor = size<2>(epilogue_tiler) / size<2>(cluster_tile);\n\n // Get the appropriate blocks for this Cluster\n dim3 cluster_coord_in_grid = cluster_id_in_grid();\n\n // Total number of k-tiles\n int const K_TILE_MAX = ceil_div(min(N, K), size<2>(epilogue_tiler));\n\n struct TileScheduler {\n struct WorkTileInfo {\n uint32_t m_idx = 0;\n uint32_t n_idx = 0;\n uint32_t l_idx = 0;\n bool is_valid_tile = false;\n };\n uint32_t tiles_in_m = 0;\n uint32_t tiles_in_n = 0;\n\n int k_tile_max = 0;\n\n int wave_cnt = 0;\n WorkTileInfo work_tile_info;\n WorkTileInfo next_work_tile_info;\n CLCResponse* clc_response_ptr_;\n CUTLASS_DEVICE TileScheduler(uint32_t tiles_m, uint32_t tiles_n, int kmax, CLCResponse* clc_response_ptr)\n : tiles_in_m(tiles_m),\n tiles_in_n(tiles_n),\n\n k_tile_max(kmax),\n work_tile_info({blockIdx.x, blockIdx.y, blockIdx.z, blockIdx.x(\n &clc_response_ptr[state.index()]));\n asm volatile(\n \"{\\n\\t\"\n \"clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes.multicast::cluster::all.b128 [%0], [%1];\\n\\t\"\n \"}\\n\"\n :\n : \"r\"(result_addr), \"r\"(mbarrier_addr));\n #else\n CUTLASS_NOT_IMPLEMENTED();\n #endif\n }\n CUTLASS_DEVICE\n static WorkTileInfo\n work_tile_info_from_clc_response(uint32_t result_addr) {\n WorkTileInfo work_tile_info;\n uint32_t valid = 0;\n #if defined(CUTLASS_ARCH_CLC_ENABLED)\n asm volatile(\n \"{\\n\"\n \".reg .pred p1;\\n\\t\"\n \".reg .b128 clc_result;\\n\\t\"\n \"ld.shared.b128 clc_result, [%4];\\n\\t\"\n \"clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_result;\\n\\t\"\n \"selp.u32 %3, 1, 0, p1;\\n\\t\"\n \"@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_result;\\n\\t\"\n \"}\\n\"\n : \"=r\"(work_tile_info.m_idx), \"=r\"(work_tile_info.n_idx), \"=r\"(work_tile_info.l_idx), \"=r\"(valid)\n : \"r\"(result_addr)\n : \"memory\"\n );\n\n cutlass::arch::fence_view_async_shared();\n #else\n CUTLASS_NOT_IMPLEMENTED();\n #endif\n work_tile_info.is_valid_tile = (valid == 1);\n return work_tile_info;\n }\n };\n\n\n\n // Allocate SMEM\n extern __shared__ char shared_memory[];\n using SharedStorage = SharedStorage;\n SharedStorage& shared_storage = *reinterpret_cast(shared_memory);\n uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile))));\n uint32_t tiles_in_n = uint32_t(size(ceil_div(N, size<2>(epilogue_tiler))));\n TileScheduler scheduler(tiles_in_m, tiles_in_n, K_TILE_MAX, shared_storage.clc_response);\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n auto acc_shape_mma = make_shape(take<0,2>(mainloop_tiler), _1{}, _1{});\n auto acc_shape_epilogue = make_shape(take<0,2>(epilogue_tiler), _1{}, _1{});\n\n auto acc_mainloop_pipelined_shape = append(acc_shape_mma, Int{});\n auto bulk_tmem_mma = TiledMMA::make_fragment_C(acc_mainloop_pipelined_shape);\n\n static int constexpr NumEpilogueColQuantThreadCount = kEnableRHTColQuant ? 128 : 0;\n static int constexpr NumEpilogueRowQuantThreadCount = kEnableRowQuant ? 256 : 0;\n static int constexpr NumMmaThreadCount = kEnableRHTColQuant? 32: 0;\n static int constexpr NumMmaIssueThreadCount = kEnableRHTColQuant? 1: 0;\n static int constexpr NumSchedThreads = 32;\n static int constexpr NumMainloopLoadThreads = 32;\n static int constexpr NumEpilogueThreads = NumEpilogueColQuantThreadCount + NumEpilogueRowQuantThreadCount;\n\n TmemAllocator tmem_allocator{};\n cutlass::arch::NamedBarrier tmem_allocation_result_barrier(\n NumMmaThreadCount + NumEpilogueColQuantThreadCount,\n cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);\n\n int warp_idx = cutlass::canonical_warp_idx_sync();\n\n // warp assignment\n bool is_mma_warp = (warp_idx == 0);\n bool is_dma_warp = (warp_idx == 1);\n bool is_sched_warp = (warp_idx == 2);\n bool is_epilogue_col_quant_warp = (warp_idx >= 4 && warp_idx <= 7);\n bool is_epilogue_row_quant_warp = (warp_idx >= 8 && warp_idx <= 15);\n\n if (is_epilogue_col_quant_warp && elect_one_sync()) {\n cute::prefetch(raw_pointer_cast(c_global_amax));\n }\n if (is_epilogue_row_quant_warp && elect_one_sync()) {\n cute::prefetch(raw_pointer_cast(a_global_amax));\n }\n\n typename MainloopPipeline::Params mainloop_pipeline_params;\n if (is_dma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;\n }\n if (is_mma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;\n }\n mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp;\n mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes;\n mainloop_pipeline_params.initializing_warp = 0;\n mainloop_pipeline_params.num_consumers = NumEpilogueRowQuantThreadCount + NumMmaIssueThreadCount;\n MainloopPipeline mainloop_pipeline(\n shared_storage.mainloop,\n mainloop_pipeline_params,\n cluster_shape,\n cute::true_type{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n\n MainloopPipelineState mainloop_pipe_consumer_state;\n MainloopPipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state();\n\n using AccumulatorPipeline = cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState;\n\n AccumulatorPipelineState accumulator_pipe_consumer_state;\n AccumulatorPipelineState accumulator_pipe_producer_state = cutlass::make_producer_start_state();\n\n typename AccumulatorPipeline::Params accumulator_pipeline_params;\n if (is_mma_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer;\n }\n if (is_epilogue_col_quant_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer;\n }\n // Only one producer thread arrives on this barrier.\n accumulator_pipeline_params.producer_arv_count = 1;\n accumulator_pipeline_params.consumer_arv_count = size(AtomThrShapeMNK{}) * NumEpilogueColQuantThreadCount;\n accumulator_pipeline_params.initializing_warp = 1;\n using IsInitAccumulatorPipeline = cute::conditional_t;\n AccumulatorPipeline accumulator_pipeline(\n shared_storage.accumulator,\n accumulator_pipeline_params,\n cluster_shape,\n IsInitAccumulatorPipeline{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n // CLC pipeline\n typename CLCPipeline::Params clc_pipeline_params;\n if (is_sched_warp) {\n clc_pipeline_params.role = CLCPipeline::ThreadCategory::ProducerConsumer;\n } else {\n clc_pipeline_params.role = CLCPipeline::ThreadCategory::Consumer;\n }\n clc_pipeline_params.producer_blockid = 0;\n clc_pipeline_params.producer_arv_count = 1;\n clc_pipeline_params.consumer_arv_count = NumSchedThreads + cluster_size *\n (NumMainloopLoadThreads + NumEpilogueThreads + NumMmaThreadCount);\n clc_pipeline_params.transaction_bytes = sizeof(CLCResponse);\n clc_pipeline_params.initializing_warp = 3;\n CLCPipeline clc_pipeline(shared_storage.clc, clc_pipeline_params, cluster_shape);\n CLCPipelineState clc_pipeline_consumer_state;\n CLCPipelineState clc_pipeline_producer_state = cutlass::make_producer_start_state();\n\n // CLC throttle pipeline\n typename CLCThrottlePipeline::Params clc_throttle_pipeline_params;\n if (is_dma_warp) {\n clc_throttle_pipeline_params.role = CLCThrottlePipeline::ThreadCategory::Producer;\n }\n if (is_sched_warp) {\n clc_throttle_pipeline_params.role = CLCThrottlePipeline::ThreadCategory::Consumer;\n }\n clc_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads;\n clc_throttle_pipeline_params.consumer_arv_count = NumSchedThreads;\n clc_throttle_pipeline_params.dst_blockid = 0;\n clc_throttle_pipeline_params.initializing_warp = 4;\n\n CLCThrottlePipeline clc_throttle_pipeline(shared_storage.clc_throttle, clc_throttle_pipeline_params);\n CLCThrottlePipelineState clc_pipe_throttle_consumer_state;\n CLCThrottlePipelineState clc_pipe_throttle_producer_state = cutlass::make_producer_start_state();\n\n if (warp_idx == 2 && elect_one_sync()) {\n cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);\n }\n __syncthreads();\n\n if (is_dma_warp) {\n cutlass::arch::warpgroup_reg_dealloc<32>();\n cute::Tensor mA = tma_load_a.get_tma_tensor(make_shape(M,N));\n cute::Tensor mB = tma_load_b.get_tma_tensor(make_shape(RhtTensorSize, RhtTensorSize));\n\n cute::Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_,_, _), Step<_1, X,_1>{});\n cute::Tensor gB_nk = local_tile(mB, cluster_tile, make_coord(_,_, _), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)\n\n cute::Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n cute::Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n cute::Tensor tCgA = thr_mma.partition_A(gA_mk); // (MMA,MMA_M,MMA_K,k)\n cute::Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k)\n\n Layout cta_layout_mnk = make_layout(cluster_shape);\n Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{}));\n auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster);\n\n auto [tAgA, tAsA] = tma_partition(\n tma_load_a,\n get<2>(cta_coord_vmnk),\n make_layout(size<2>(cta_layout_vmnk)),\n group_modes<0,3>(tCsA),\n group_modes<0,3>(tCgA));\n\n auto [tBgB, tBsB] = tma_partition(\n tma_load_b,\n get<1>(cta_coord_vmnk),\n make_layout(size<1>(cta_layout_vmnk)),\n group_modes<0,3>(tCsB),\n group_modes<0,3>(tCgB));\n\n uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);\n uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);\n if constexpr (kEnableRHTColQuant) {\n if (elect_one_sync()) {\n cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0], kTmaRhtTensorTransactionBytes);\n copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_,0,0), tBsB(_,0));\n }\n }\n\n do {\n bool is_first_wave = scheduler.is_first_wave();\n uint32_t skip_wait = is_first_wave;\n auto tAgA_mk = tAgA(_,scheduler.tile_m(),_);\n int k_tile = 0;\n // Throttle CLC producer\n clc_throttle_pipeline.producer_acquire(clc_pipe_throttle_producer_state);\n clc_throttle_pipeline.producer_commit(clc_pipe_throttle_producer_state);\n ++clc_pipe_throttle_producer_state;\n\n CUTLASS_PRAGMA_NO_UNROLL\n while (k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n()) {\n\n int k_tile_idx_n = scheduler.tile_n_base() + k_tile;\n ++k_tile;\n skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount);\n mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state);\n using BarrierType = typename MainloopPipeline::ProducerBarrierType;\n BarrierType* tma_barrier = mainloop_pipeline.producer_get_barrier(\n mainloop_pipe_producer_state);\n int write_stage = mainloop_pipe_producer_state.index();\n ++mainloop_pipe_producer_state;\n if (cute::elect_one_sync()) {\n copy(\n tma_load_a.with(*tma_barrier, tma_mcast_mask_a),\n tAgA_mk(_,k_tile_idx_n),\n tAsA(_,write_stage));\n }\n }\n scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state);\n ++clc_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);\n } else if (is_mma_warp) {\n cutlass::arch::warpgroup_reg_dealloc<32>();\n if constexpr (kEnableRHTColQuant) {\n cute::Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n cute::Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n // Allocate \"fragments\" -- these are actually umma smem descriptors\n cute::Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)\n cute::Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE)\n\n mma.accumulate_ = UMMA::ScaleOut::Zero;\n\n tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);\n __syncwarp();\n tmem_allocation_result_barrier.arrive();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_mma.data() = tmem_base_ptr;\n cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/);\n do {\n uint32_t skip_wait = K_TILE_MAX <= 0;\n\n auto barrier_token = mainloop_pipeline.consumer_try_wait(\n mainloop_pipe_consumer_state,\n skip_wait);\n scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state);\n ++clc_pipeline_consumer_state;\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); ) {\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n int read_stage = mainloop_pipe_consumer_state.index();\n auto tCrA_mk = tCrA(_,_,_,read_stage);\n auto tCrB_nk = tCrB(_,_,0,0);\n CUTLASS_PRAGMA_UNROLL\n for (int k_block = 0; k_block < size<2>(tCrA) / EpilogueUnrollFactor; ++k_block)\n {\n int accumulator_k_block = accumulator_pipe_producer_state.index() * EpilogueUnrollFactor;\n int tCrA_k_block = k_block * EpilogueUnrollFactor;\n accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < EpilogueUnrollFactor; i++) {\n auto accumulators = bulk_tmem_mma(_,_,_,accumulator_k_block + i);\n gemm(mma, tCrA_mk(_,_,tCrA_k_block + i), tCrB_nk, accumulators);\n }\n\n accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);\n ++accumulator_pipe_producer_state;\n }\n auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n skip_wait = k_tile >= K_TILE_MAX;\n mainloop_pipeline.umma_consumer_release(curr_mainloop_pipe_consumer_state);\n barrier_token = mainloop_pipeline.consumer_try_wait(\n mainloop_pipe_consumer_state,\n skip_wait);\n }\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n tmem_allocator.release_allocation_lock();\n accumulator_pipeline.producer_tail(accumulator_pipe_producer_state);\n tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);\n }\n } else if(is_sched_warp) {\n cutlass::arch::warpgroup_reg_dealloc<32>();\n do {\n clc_throttle_pipeline.consumer_wait(clc_pipe_throttle_consumer_state);\n clc_throttle_pipeline.consumer_release(clc_pipe_throttle_consumer_state);\n ++clc_pipe_throttle_consumer_state;\n clc_pipeline_producer_state = scheduler.advance_to_next_work(clc_pipeline, clc_pipeline_producer_state);\n scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state);\n ++clc_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n } else if (is_epilogue_col_quant_warp) {\n cutlass::arch::warpgroup_reg_alloc<192>();\n if constexpr (kEnableRHTColQuant) {\n using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x;\n\n float const c_global_amax_val = *c_global_amax;\n auto acc_epilogue_pipelined_shape = append(acc_shape_epilogue, Int{});\n auto bulk_tmem_epilogue_layout = make_layout(\n acc_epilogue_pipelined_shape,\n make_stride(\n stride<0>(bulk_tmem_mma),\n Int<0>{},\n Int<0>{},\n size<1>(epilogue_tiler)));\n auto bulk_tmem_epilogue = make_tensor(make_tmem_ptr(), bulk_tmem_epilogue_layout);\n\n // leveraging 256-bit writes to global memory\n static int constexpr FragmentSize = 256 / sizeof_bits_v;\n\n tmem_allocation_result_barrier.arrive_and_wait();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_epilogue.data() = tmem_base_ptr;\n int global_thread_idx = threadIdx.x;\n int local_thread_idx = global_thread_idx % cutlass::NumThreadsPerWarpGroup;\n\n size_t rng_seed = 0;\n size_t rng_offset = 0;\n if constexpr (kEnableStochasticRounding) {\n rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0;\n rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0;\n }\n\n cute::Tensor mD = make_tensor(\n cute::subbyte_iterator(D),\n make_shape(M,N),\n dD); // (M,N)\n cute::Tensor gD_mn = local_tile(\n mD,\n epilogue_tiler,\n make_coord(_,_, _),\n Step<_1,_1, X>{}); // (BLK_M,BLK_N)\n cute::Tensor pD = make_identity_tensor(mD.shape());\n cute::Tensor pD_mn = local_tile(\n pD,\n epilogue_tiler,\n make_coord(_,_, _),\n Step<_1,_1, X>{}); // (BLK_M,BLK_N)\n cute::Tensor mSFD = make_tensor(make_gmem_ptr(SFD), sfd_layout);\n cute::Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N)\n cute::Tensor pSFD = make_identity_tensor(mSFD.shape());\n cute::Tensor pSFD_mn = local_tile(pSFD, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N)\n\n cute::Tensor gD_mn_view = tiled_divide(gD_mn, take<0,2>(epilogue_tiler));\n cute::Tensor pD_mn_view = tiled_divide(pD_mn, take<0,2>(epilogue_tiler));\n auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_,_,_,_0{}));\n auto tiled_r2g = make_tiled_copy_D(\n Copy_Atom{},\n tiled_t2r);\n auto thr_t2r = tiled_t2r.get_slice(local_thread_idx);\n auto thr_r2g = tiled_r2g.get_slice(local_thread_idx);\n\n // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release}\n static constexpr float fp4_max = 6.0f;\n static constexpr float fp8_max = 448.0f;\n float const fp4_max_inv = 1.0f / fp4_max;\n float const global_encode_scale = c_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / c_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n\n float const global_decode_scale = 1.0f / global_encode_scale;\n // Scaling factor for fast math path\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n auto sfc_converter = cutlass::NumericConverter{};\n\n do {\n scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state);\n ++clc_pipeline_consumer_state;\n for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); ++k_tile) {\n cute::Tensor tDgD_mn = gD_mn_view(_,_,_,scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n cute::Tensor tDgSFD_mn = gSFD_mn(_,_,scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n cute::Tensor tDpD_mn = pD_mn_view(_,_,_,scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n cute::Tensor tDpSFD_mn = pSFD_mn(_,_,scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n\n accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state);\n\n auto Acc = bulk_tmem_epilogue(_,_,_,accumulator_pipe_consumer_state.index());\n cute::Tensor tDtAcc = thr_t2r.partition_S(Acc); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n cute::Tensor tDgD = thr_t2r.partition_D(tDgD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n cute::Tensor tDpD = thr_t2r.partition_D(tDpD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n cute::Tensor tTR_rAcc = make_tensor(shape(tDgD)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n cute::Tensor tDrD = make_tensor(shape(tDgD));\n cute::Tensor tTR_rAcc_frag = recast>(coalesce(tTR_rAcc));\n cute::Tensor tDrD_frag = recast>(coalesce(tDrD));\n\n cute::Tensor src = thr_r2g.retile_S(tDrD);\n cute::Tensor dst = thr_r2g.retile_D(tDgD);\n cute::Tensor pSrc = thr_r2g.retile_D(tDpD);\n\n cute::Tensor tDgSFD_view = make_tensor(\n tDgSFD_mn.data(),\n make_layout(\n make_shape(shape(tDgSFD_mn), Int<1>{}, Int<1>{}),\n make_stride(stride(tDgSFD_mn), Int<0>{}, Int<0>{})));\n cute::Tensor tDpSFD_view = make_tensor(\n tDpSFD_mn.data(),\n make_layout(\n make_shape(shape(tDpSFD_mn), Int<1>{}, Int<1>{}),\n make_stride(stride(tDpSFD_mn), Int<0>{}, Int<0>{})));\n cute::Tensor tDgSFD = filter(thr_t2r.partition_D(tDgSFD_view));\n cute::Tensor tDrSFD = make_tensor(shape(tDgSFD));\n cute::Tensor tDpSFD = filter(thr_t2r.partition_D(tDpSFD_view));\n static int constexpr NumVecs = size(tDgD) / VectorSize;\n cute::Tensor tD_rRowSFD_frg = recast>(tDrSFD);\n\n cutlass::maximum_absolute_value_reduction, true> amax_reduction;\n cutlass::Array vec_maxs;\n cutlass::Array pvscales;\n // TMEM_LOAD\n copy(tiled_t2r, tDtAcc, tTR_rAcc);\n cutlass::arch::fence_view_async_tmem_load();\n accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state);\n ++accumulator_pipe_consumer_state;\n\n if constexpr (!kUseFastMath) {\n // Downcast to BF16 for bit-wise compatibility with\n // unfused kernels\n auto convert_accum_to_bf16 =\n cutlass::NumericArrayConverter{};\n auto convert_bf16_to_accum =\n cutlass::NumericArrayConverter{};\n tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{})));\n tTR_rAcc_frag(_1{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_1{})));\n }\n\n auto compute_frgs = reinterpret_cast *>(tTR_rAcc_frag.data());\n auto output_frgs = reinterpret_cast *>(tDrD_frag.data());\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]);\n }\n\n pvscales = cutlass::multiplies>{}(\n vec_maxs, global_encode_scale_multiplier);\n auto pvscales_cvted = cutlass::NumericArrayConverter{}(pvscales);\n\n tD_rRowSFD_frg(_0{}) = pvscales_cvted;\n auto qpvscale_ups = cutlass::NumericArrayConverter{}(tD_rRowSFD_frg(_0{}));\n auto qpvscale_scaled = cutlass::multiplies>{}(\n qpvscale_ups,\n global_decode_scale);\n\n cutlass::Array acc_scales;\n if constexpr (kUseFastMath) {\n // fast math: use reciprocal approximate to replace div\n acc_scales = cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled);\n } else {\n // regular path for slower math, use divide to replace div\n acc_scales = cutlass::divides>{}(1.0, qpvscale_scaled);\n }\n\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n transformer_engine::curanddx::detail::philox4x32_native_state rng;\n // \"Prefetch\" a stochastic rounding state for the first tile\n if constexpr (kEnableStochasticRounding) {\n const size_t rng_sequence = global_thread_idx + k_tile * 512 + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512;\n rng.init(rng_seed, rng_sequence, rng_offset);\n }\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(\n acc_scales[v],\n cutlass::platform::numeric_limits::max());\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(cutlass::multiplies>{}(compute_frgs[v], acc_scale), *reinterpret_cast*>(&random_uint4));\n } else {\n output_frgs[v] = cutlass::NumericArrayConverter{}(\n cutlass::multiplies>{}(\n compute_frgs[v],\n acc_scale));\n }\n\n }\n\n cute::Tensor pred_pSrc = cute::lazy::transform(make_tensor(counting_iterator{}, replace<0>(shape(dst), _1{})), [&](auto coord){\n cute::Tensor pSrc_view = group_modes<1,rank(pSrc)>(pSrc);\n return elem_less(pSrc_view(_0{},coord), shape(mD));\n });\n copy_if(tiled_r2g, pred_pSrc, src, dst);\n // 32bit vectorization copy 4 e4m3 SFD for per 64 or(16,4):(0, 1) element\n\n constexpr int vec_len = 32 / sizeof_bits_v;\n cute::Tensor tDrSFD_v = recast>(tDrSFD);\n cute::Tensor tDgSFD_v = recast>(tDgSFD);\n copy_if(\n [&](auto coord){\n cute::Tensor tDpSFD_view = group_modes<1,rank(tDpSFD)>(tDpSFD);\n return elem_less(tDpSFD_view(_0{}, coord * vec_len), shape(mSFD));\n },\n tDrSFD_v, tDgSFD_v);\n }\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n }\n } else if (is_epilogue_row_quant_warp) {\n cutlass::arch::warpgroup_reg_alloc<136>();\n if constexpr (kEnableRowQuant) {\n using S2RVectorType = uint128_t;\n float const a_global_amax_val = *a_global_amax;\n int global_thread_idx = threadIdx.x;\n int local_thread_idx = global_thread_idx % 256;\n size_t rng_seed = 0;\n size_t rng_offset = 0;\n if constexpr (kEnableStochasticRounding) {\n rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0;\n rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0;\n }\n cute::Tensor mQA = make_tensor(cute::subbyte_iterator(QA), make_layout(make_shape(M, N), dQA));\n cute::Tensor gQA_mn = local_tile(mQA, epilogue_tiler, make_coord(_,_, _), Step<_1,X,_1>{});\n cute::Tensor pQA = make_identity_tensor(mQA.shape());\n cute::Tensor pQA_mn = local_tile(pQA, epilogue_tiler, make_coord(_,_, _), Step<_1,X,_1>{});\n\n cute::Tensor mSFA = make_tensor(make_gmem_ptr(SFA), sfa_layout);\n cute::Tensor gSFA_mn = local_tile(mSFA, epilogue_tiler, make_coord(_,_, _), Step<_1,X,_1>{}); // (BLK_M,BLK_N)\n cute::Tensor pSFA = make_identity_tensor(mSFA.shape());\n cute::Tensor pSFA_mn = local_tile(pSFA, epilogue_tiler, make_coord(_,_, _), Step<_1,X,_1>{});\n cute::Tensor sA = as_position_independent_swizzle_tensor(\n group_modes<0,2>(coalesce(make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), sAlayout)))); // (BLOCK_M, BLOCK_M,PIPE)\n using S2RWarpLayout = Layout>;\n using WarpGroupLayout = Layout>;\n using S2RThreadLayout = decltype(blocked_product(S2RWarpLayout{}, WarpGroupLayout{}));\n using S2RValLayout = Layout, _1>>;\n using S2RAtomA = Copy_Atom;\n using R2GAtomQA = Copy_Atom;\n\n auto tiled_s2r = make_tiled_copy(S2RAtomA{}, S2RThreadLayout{}, S2RValLayout{});\n auto tiled_r2g_QA = make_tiled_copy_D(R2GAtomQA{}, tiled_s2r);\n\n auto thr_s2r = tiled_s2r.get_slice(local_thread_idx);\n auto thr_r2g_QA = tiled_r2g_QA.get_slice(local_thread_idx);\n\n cute::Tensor tQAsA = thr_s2r.partition_S(sA); // (Copy, Copy_M, Copy_N, PIPE)\n\n cute::Tensor tQArA = make_tensor_like(make_layout(tQAsA(_, _, _, _0{}).shape())); // (Copy, Copy_M, Copy_N)\n // Tensor tQArA_PI = thr_s2r.partition_S(sA_PI);\n cute::Tensor tQAgQA = thr_r2g_QA.partition_D(gQA_mn);\n cute::Tensor tQArQA = make_tensor_like(tQAgQA(_, _, _, _0{}, _0{}));\n cute::Tensor tQApQA = thr_r2g_QA.partition_D(pQA_mn);\n\n cute::Tensor tQAgSFA = thr_s2r.partition_D(gSFA_mn);\n cute::Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{}));\n cute::Tensor tQApSFA = thr_s2r.partition_D(pSFA_mn);\n\n // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release}\n static constexpr float fp4_max = 6.0f;\n static constexpr float fp8_max = 448.0f;\n float const fp4_max_inv = 1.0f / fp4_max;\n float const global_encode_scale = a_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / a_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n\n float const global_decode_scale = 1.0f / global_encode_scale;\n // Scaling factor for fast math path\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n\n auto sfa_converter = cutlass::NumericConverter{};\n do {\n uint32_t skip_wait = K_TILE_MAX <= 0;\n\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); ) {\n auto tQAgSFA_mn = tQAgSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto tQAgQA_mn = tQAgQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto tQApSFA_mn = tQApSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto tQApQA_mn = tQApQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto barrier_token = mainloop_pipeline.consumer_try_wait(\n mainloop_pipe_consumer_state);\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n copy(tiled_s2r, tQAsA(_, _, _, mainloop_pipe_consumer_state.index()), tQArA);\n cutlass::arch::fence_view_async_shared();\n auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n skip_wait = k_tile >= K_TILE_MAX;\n mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state);\n // static int constexpr NumVecs = size(tQArA) / VectorSize;\n cutlass::maximum_absolute_value_reduction, true> amax_reduction;\n auto compute_frgs = reinterpret_cast *>(tQArA.data());\n auto output_frgs = reinterpret_cast *>(raw_pointer_cast(tQArQA.data()));\n transformer_engine::curanddx::detail::philox4x32_native_state rng;\n if constexpr (kEnableStochasticRounding) {\n const size_t rng_sequence = global_thread_idx + k_tile * 512 + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512;\n rng.init(rng_seed, rng_sequence, rng_offset);\n }\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < size(tQArA)/VectorSize; v++) {\n auto compute_frgs_up = cutlass::NumericArrayConverter{}(compute_frgs[v]);\n auto amax = amax_reduction(ElementAccumulator(0), compute_frgs_up);\n // declare pvscales\n ElementAccumulator pvscales;\n pvscales = cutlass::multiplies{}(amax, global_encode_scale_multiplier);\n filter(tQArSFA)(v) = sfa_converter(pvscales);\n auto qpvscale_ups = cutlass::NumericConverter{}(filter(tQArSFA)(v));\n auto qpvscale_scaled = cutlass::multiplies{}(qpvscale_ups, global_decode_scale);\n ElementAccumulator acc_scales;\n if constexpr (kUseFastMath) {\n // fast math: use reciprocal approximate to replace div\n acc_scales = cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled);\n } else {\n // regular path for slower math, use divide to replace div\n acc_scales = cutlass::divides{}(1.0, qpvscale_scaled);\n }\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(\n acc_scales,\n cutlass::platform::numeric_limits::max());\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(cutlass::multiplies>{}(compute_frgs_up, acc_scale), *reinterpret_cast*>(&random_uint4));\n } else {\n output_frgs[v] = cutlass::NumericArrayConverter{}(\n cutlass::multiplies>{}(\n compute_frgs_up,\n acc_scale));\n }\n }\n\n cute::Tensor pred_tQApQA = cute::lazy::transform(make_tensor(counting_iterator{}, replace<0>(shape(tQAgQA_mn), _1{})), [&](auto coord){\n cute::Tensor tQApQA_view = group_modes<1,rank(tQApQA_mn)>(tQApQA_mn);\n return elem_less(tQApQA_view(_0{}, coord), shape(mQA));\n });\n copy_if(tiled_r2g_QA, pred_tQApQA, tQArQA, tQAgQA_mn);\n // 32bit vectorization copy 4 e4m3 SFA for per 64 or (16,4):(0, 1) element\n constexpr int vec_len = 32 / sizeof_bits_v;\n cute::Tensor tQArSFA_v = recast>(filter(tQArSFA));\n cute::Tensor tQAgSFA_v = recast>(filter(tQAgSFA_mn));\n copy_if(\n [&](auto coord){\n cute::Tensor tQApSFA_view = filter(tQApSFA_mn);\n return elem_less(tQApSFA_view(_0{}, coord * vec_len), shape(mSFA));\n },\n tQArSFA_v, tQAgSFA_v);\n }\n scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state);\n ++clc_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n }while (scheduler.is_valid());\n }\n } else {\n cutlass::arch::warpgroup_reg_dealloc<32>();\n }\n } // sm100 compile guard end\n} // NOLINT(readability/fn_size)\n\n\n// this function computes RHT-GEMM for\n// m = hidden_size, n = sequence_length\n// A: m x n: col-major\n// B: 16 x 16: row-major\n// D: m x n: row-major\n// SFD: m x (n/16): row-major\n// QA: m x n: col-major\n// SFA: m/16 x n: col-major\ntemplate \nvoid row_col_rht_gemm_ntt_w_sfc(\n int sequence_length,\n int hidden_size,\n TA const* A,\n TB const* B,\n TD* D,\n TSFD* SFD,\n TQA* QA,\n TSFA* SFA,\n float const* a_global_amax,\n float const* d_global_amax,\n const size_t* rng_state,\n uint32_t sm_count,\n cudaStream_t stream,\n int k_tile_size = 1024) {\n using namespace cute;\n static int constexpr SFVecSize = 16;\n static int constexpr RhtTensorSize = 16;\n\n static_assert(RhtTensorSize == 16, \"RhtTensorSize must be 16\");\n using LinearSFALayout = decltype(make_layout(make_shape(make_shape(Int{}, 0), 0), make_stride(make_stride(_0{}, _1{}), 0)));\n using LinearSFCLayout = decltype(make_layout(make_shape(0, make_shape(Int{}, 0)), make_stride(0, make_stride(_0{}, _1{}))));\n\n using SwizzledSFALayoutAtom = cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFDLayoutAtom = cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFALayout = decltype(tile_to_shape(SwizzledSFALayoutAtom{}, make_shape(hidden_size,sequence_length), Step<_1,_2>{}));\n using SwizzledSFDLayout = decltype(tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(hidden_size,sequence_length), Step<_2,_1>{}));\n\n using SFALayout = cute::conditional_t;\n using SFCLayout = cute::conditional_t;\n SFALayout sfa_layout;\n SFCLayout sfd_layout;\n\n if constexpr (kEnableSwizzleSFOutput) {\n sfa_layout = tile_to_shape(SwizzledSFALayoutAtom{}, make_shape(hidden_size, sequence_length), Step<_1,_2>{});\n sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(hidden_size, sequence_length), Step<_2,_1>{});\n } else {\n sfa_layout = make_layout(make_shape(make_shape(Int{}, hidden_size/SFVecSize), sequence_length), make_stride(make_stride(_0{}, _1{}), hidden_size/SFVecSize));\n sfd_layout = make_layout(make_shape(hidden_size, make_shape(Int{}, sequence_length/SFVecSize)), make_stride(sequence_length/SFVecSize, make_stride(_0{}, _1{})));\n }\n // Define shapes (dynamic)\n auto M = hidden_size;\n auto N = sequence_length;\n cute::Tensor tensorA = make_tensor(A, make_shape(hidden_size, sequence_length), LayoutLeft{});\n cute::Tensor tensorB = make_tensor(B, make_shape(RhtTensorSize, RhtTensorSize), LayoutLeft{});\n cute::Tensor tensorD = make_tensor(D, make_shape(hidden_size, sequence_length), LayoutRight{});\n cute::Tensor tensorQA = make_tensor(QA, make_shape(hidden_size, sequence_length), LayoutLeft{});\n cute::Tensor tensorSFD = make_tensor(SFD, sfd_layout);\n cute::Tensor tensorSFA = make_tensor(SFA, sfa_layout);\n // Define strides (from tensors)\n auto dA = stride(tensorA); // (dM,dK)\n auto dB = stride(tensorB); // (dN,dK)\n auto dD = stride(tensorD); // (dM,dN)\n auto dQA = stride(tensorQA); // (dM,dK)\n using ClusterShape = Shape< _1, _1, _1>;\n auto cluster_shape = ClusterShape{};\n auto cluster_tile_shape = Shape<_128,Int,Int>{};\n auto cluster_tile_mainloop = Shape<_128,Int,_128>{};\n\n // Each mainloop / epilogue loads 128 x 64 tiles while each MMA proceeds with 128 x 16 tiles\n static int constexpr EpilogueUnrollFactor =\n size<2>(cluster_tile_mainloop) / size<2>(cluster_tile_shape);\n // Construct the MMA\n auto mma = make_tiled_mma(SM100_MMA_F16BF16_SS(cluster_tile_shape), size<1>(cluster_tile_shape),\n UMMA::Major::MN, UMMA::Major::MN>{},\n Layout>{});\n\n // Assert that the TiledMMA uses all CTAs in the CGA.\n CUTE_STATIC_ASSERT_V(size(cluster_shape) == size(mma));\n CUTE_STATIC_ASSERT_V(evenly_divides(cluster_tile_shape, tile_shape(mma)));\n\n // Determine the A and B shapes\n auto mma_shape_B = partition_shape_B(mma, make_shape(size<1>(cluster_tile_shape), size<2>(cluster_tile_shape)));\n\n using TiledMma = decltype(mma);\n using AtomThrID = typename TiledMma::AtomThrID;\n\n using SmemShape_M = decltype(shape_div(shape<0>(cluster_tile_shape), shape_div(shape<0>(cluster_tile_shape), size<0>(cluster_tile_shape) / size(AtomThrID{}))));\n using SmemShape_N = decltype(shape_div(shape<1>(cluster_tile_shape), shape_div(shape<1>(cluster_tile_shape), size<1>(cluster_tile_shape) / size(AtomThrID{}))));\n using SmemShape_K = decltype(cute::get<2>(cluster_tile_shape));\n\n using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<\n cute::UMMA::Major::MN, TB, SmemShape_N, SmemShape_K>());\n\n auto mma_shape_A = partition_shape_A(mma, make_shape(size<0>(cluster_tile_mainloop), size<2>(cluster_tile_mainloop)));\n using SmemShape_M_A = decltype(shape_div(shape<0>(cluster_tile_mainloop), shape_div(shape<0>(cluster_tile_mainloop), size<0>(cluster_tile_mainloop) / size(AtomThrID{}))));\n using SmemShape_K_A = decltype(cute::get<2>(cluster_tile_mainloop));\n using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<\n cute::UMMA::Major::MN, TA, SmemShape_M_A, SmemShape_K_A>());\n\n static uint32_t constexpr TotalTmemRows = 128;\n static uint32_t constexpr Sm100TmemCapacityColumns = 512;\n static uint32_t constexpr TotalTmem = TotalTmemRows * Sm100TmemCapacityColumns;\n static uint32_t constexpr AccumulatorPipelineStageCount =\n TotalTmem /\n (cute::size<0>(cluster_tile_shape) * cute::size<1>(cluster_tile_shape));\n\n // Define the smem layouts (static)\n // Calculate max pipeline stages based on Blackwell SM100's 232KB shared memory\n constexpr int SchedulerPipelineStageCount = 6;\n static int constexpr MainloopPipelineBytes = sizeof(typename cutlass::detail::CustomizedPipelineTmaUmmaAsync<\n 1,\n Shape<_1,_1,_1>,\n Shape<_1, _1, _1>>::SharedStorage);\n\n static int constexpr ClcResponseBytes = sizeof(CLCResponse) * SchedulerPipelineStageCount;\n static int constexpr CLCThrottlePipelineBytes = sizeof(typename cutlass::PipelineAsync::SharedStorage);\n static int constexpr CLCPipelineBytes = sizeof(typename cutlass::PipelineCLCFetchAsync::SharedStorage);\n static int constexpr TmemDeallocBytes = sizeof(cutlass::arch::ClusterBarrier);\n static int constexpr BTensorBytes = cute::size(mma_shape_B) * sizeof(TB);\n static int constexpr AccPipelineBytes = sizeof(typename cutlass::PipelineUmmaAsync>::SharedStorage);\n static int constexpr TmemBasePtrsBytes = sizeof(uint32_t);\n static int constexpr kBlackwellSmemSize = 232448; // 232KB in bytes\n static int constexpr kBytesPerStage =\n cute::size(mma_shape_A) * sizeof(TA) + MainloopPipelineBytes;\n static int constexpr kReservedBytes = ClcResponseBytes + CLCThrottlePipelineBytes + TmemBasePtrsBytes +\n CLCPipelineBytes + TmemDeallocBytes+BTensorBytes + AccPipelineBytes; // Reserve for barriers and other uses\n static int constexpr kMaxStages = (kBlackwellSmemSize - kReservedBytes) / kBytesPerStage;\n auto sP = Int{}; // SMEM pipelines\n auto sA = UMMA::tile_to_mma_shape(\n SmemLayoutAtomA{},\n append(mma_shape_A, sP), Step<_2,_1,_3>{}); // (MMA,MMA_M,MMA_K,PIPE)\n auto sB = UMMA::tile_to_mma_shape(\n SmemLayoutAtomB{},\n append(mma_shape_B, _1{})); // (MMA,MMA_N,MMA_K, _1)\n auto sD = Layout<_1>{}; // XXX Dummy\n\n auto tma_load_a = make_tma_copy_A_sm100(\n SM90_TMA_LOAD{},\n tensorA,\n sA(_,_,_,0),\n cluster_tile_mainloop,\n mma);\n auto tma_load_b = make_tma_copy_B_sm100(\n SM90_TMA_LOAD{},\n tensorB,\n sB(_,_,_,0),\n cluster_tile_shape,\n mma);\n\n // Assert checks problem size should be multiple of 64\n NVTE_CHECK(M % 64 == 0, \"M must be a multiple of 64, but got \", M);\n NVTE_CHECK(N % 64 == 0, \"N must be a multiple of 64, but got \", N);\n\n uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile_shape))));\n uint32_t tiles_in_n = uint32_t(size(ceil_div(N, k_tile_size)));\n uint32_t tiles = tiles_in_m * tiles_in_n;\n\n dim3 dimBlock(512);\n dim3 dimCluster(size<0>(cluster_shape), size<1>(cluster_shape), size<2>(cluster_shape));\n dim3 dimGrid(tiles_in_m, tiles_in_n, 1);\n\n int smem_size = sizeof(\n SharedStorage<\n TA,\n TB,\n decltype(sA),\n decltype(sB),\n ClusterShape,\n AccumulatorPipelineStageCount,\n EpilogueUnrollFactor,\n SchedulerPipelineStageCount>);\n\n auto* kernel_ptr = &row_col_rht_gemm_device<\n decltype(M), decltype(N), decltype(k_tile_size),\n decltype(cluster_shape), decltype(cluster_tile_shape),\n TA, decltype(dA), decltype(sA), decltype(tma_load_a),\n TB, decltype(dB), decltype(sB), decltype(tma_load_b),\n TD, decltype(dD), decltype(sD),\n TSFD, decltype(sfd_layout),\n TQA, decltype(dQA),\n TSFA, decltype(sfa_layout),\n decltype(mma),\n AccumulatorPipelineStageCount,\n SchedulerPipelineStageCount,\n kEnableStochasticRounding,\n kEnableRHTColQuant,\n kEnableRowQuant,\n kUseFastMath>;\n\n NVTE_CHECK_CUDA(cudaFuncSetAttribute(*kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n\n cutlass::ClusterLaunchParams params = {dimGrid, dimBlock, dimCluster, smem_size, stream};\n cutlass::Status status = cutlass::launch_kernel_on_cluster(\n params, (void const *)kernel_ptr, M, N, k_tile_size, cluster_shape, cluster_tile_shape,\n tensorA.data(), dA, sA, tma_load_a,\n tensorB.data(), dB, sB, tma_load_b,\n tensorD.data(), dD, sD,\n tensorSFD.data(), sfd_layout,\n tensorQA.data(), dQA,\n tensorSFA.data(), sfa_layout,\n mma, a_global_amax, d_global_amax, rng_state);\n\n NVTE_CHECK_CUDA(cudaGetLastError());\n NVTE_CHECK(status == cutlass::Status::kSuccess, \"Kernel launch failed.\");\n\n}\n\n} // namespace\n} // namespace detail\n\n// clang-format on\n\nvoid hadamard_transform_cast_fusion(const Tensor &input_, Tensor &output_,\n const Tensor &hadamard_matrix_, QuantizationConfig quant_config,\n cudaStream_t stream) {\n NVTE_API_CALL(hadamard_transform_cast_fusion);\n\n // Check input and output tensors\n NVTE_CHECK(input_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor must be BF16 tensor, but scaling mode is \",\n to_string(input_.scaling_mode), \".\");\n NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16,\n \"Input tensor must be BF16 tensor, but dtype is \", to_string(input_.dtype()), \".\");\n NVTE_CHECK(input_.dim() >= 2, \"Input must be a 2D tensor.\");\n const SimpleTensor &input = input_.data;\n\n // rowwise cast and columnwise cast has different output data pointers\n bool has_rowwise_quant = false;\n bool has_columnwise_quant = false;\n void *rowwise_data_ptr = nullptr;\n void *rowwise_scale_inv_ptr = nullptr;\n void *rowwise_amax_ptr = nullptr;\n void *columnwise_data_ptr = nullptr;\n void *columnwise_scale_inv_ptr = nullptr;\n void *columnwise_amax_ptr = nullptr;\n\n // examine the output tensor (single tensor for dense)\n if (output_.data.dptr != nullptr) {\n has_rowwise_quant = true;\n rowwise_data_ptr = output_.data.dptr;\n rowwise_scale_inv_ptr = output_.scale_inv.dptr;\n rowwise_amax_ptr = output_.amax.dptr;\n }\n\n if (output_.columnwise_data.dptr != nullptr) {\n has_columnwise_quant = true;\n columnwise_data_ptr = output_.columnwise_data.dptr;\n columnwise_scale_inv_ptr = output_.columnwise_scale_inv.dptr;\n columnwise_amax_ptr = output_.columnwise_amax.dptr;\n }\n\n NVTE_CHECK(has_rowwise_quant || has_columnwise_quant,\n \"Output tensor must have rowwise or columnwise quant.\");\n\n // Stochastic rounding config\n const bool use_stochastic_rounding = quant_config.stochastic_rounding;\n const size_t *rng_state = nullptr;\n if (quant_config.rng_state != nullptr) {\n Tensor &rng_state_tensor = *convertNVTETensor(quant_config.rng_state);\n NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64,\n \"RNG state should contain 2 64-bit values.\");\n NVTE_CHECK(rng_state_tensor.data.shape == std::vector{2},\n \"Shape of the RNG state should be [2], but got \", rng_state_tensor.data.shape);\n rng_state = reinterpret_cast(rng_state_tensor.data.dptr);\n }\n\n // Template arguments\n using TA = cute::bfloat16_t;\n using TB = cute::bfloat16_t;\n using TD = cutlass::float_e2m1_t;\n using TSFD = cutlass::float_ue4m3_t;\n using TQA = TD;\n using TSFA = TSFD;\n\n checkCuDriverContext(stream);\n\n // Check Hadamard matrix\n constexpr int kHadamardDimension = 16;\n NVTE_CHECK(hadamard_matrix_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Hadamard matrix must be BF16 tensor, but scaling mode is \",\n to_string(hadamard_matrix_.scaling_mode), \".\");\n NVTE_CHECK(hadamard_matrix_.dtype() == transformer_engine::DType::kBFloat16,\n \"Hadamard matrix must be BF16 tensor, but dtype is \",\n to_string(hadamard_matrix_.dtype()), \".\");\n const SimpleTensor &hadamard_matrix = hadamard_matrix_.data;\n NVTE_CHECK(\n (hadamard_matrix_.shape() == std::vector{kHadamardDimension, kHadamardDimension}),\n \"Hadamard matrix must have shape=\",\n std::vector{kHadamardDimension, kHadamardDimension},\n \", but got shape=\", hadamard_matrix_.shape(), \".\");\n const size_t hadamard_dimension = hadamard_matrix.shape[0];\n\n const size_t ndim = input.shape.size();\n const size_t n = input.shape[ndim - 1];\n size_t m = 1;\n for (size_t i = 0; i < ndim - 1; ++i) {\n m *= input.shape[i];\n }\n\n auto sm_count = transformer_engine::cuda::sm_count();\n\n NVTE_CHECK(n % hadamard_dimension == 0, \"row_length must be divisible by hadamard_dimension.\");\n\n NVTE_CHECK(m % hadamard_dimension == 0, \"num_rows must be divisible by hadamard_dimension\");\n\n int k_tile_size = 1024;\n\n // TODO: add support for swizzle sf output\n const bool use_swizzle_sf_output = false;\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_stochastic_rounding, kEnableStochasticRounding,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n has_columnwise_quant, kEnableRhtColQuant,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n has_rowwise_quant, kEnableRowQuant,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_swizzle_sf_output, kEnableSwizzleSFOutput,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n quant_config.use_fast_math, kUseFastMath,\n\n if constexpr (kEnableRhtColQuant || kEnableRowQuant) {\n detail::row_col_rht_gemm_ntt_w_sfc<\n kEnableStochasticRounding, kEnableRhtColQuant, kEnableRowQuant,\n kEnableSwizzleSFOutput, TA, TB, TD, TSFD, TQA, TSFA, kUseFastMath>(\n /*sequence_length=*/m, /*hidden_size=*/n,\n /*A=*/reinterpret_cast(input.dptr),\n /*B=*/reinterpret_cast(hadamard_matrix.dptr),\n /*D=*/reinterpret_cast(columnwise_data_ptr),\n /*SFD=*/reinterpret_cast(columnwise_scale_inv_ptr),\n /*QA=*/reinterpret_cast(rowwise_data_ptr),\n /*SFA=*/reinterpret_cast(rowwise_scale_inv_ptr),\n /*a_global_amax=*/reinterpret_cast(rowwise_amax_ptr),\n /*d_global_amax=*/reinterpret_cast(columnwise_amax_ptr),\n /*rng_state=*/rng_state, /*sm_count=*/sm_count,\n /*stream=*/stream, /*k_tile_size=*/k_tile_size);\n } else {\n NVTE_ERROR(\"Invalid kernel configuration (kEnableRHTColQuant=\",\n kEnableRhtColQuant, \", kEnableRowQuant=\", kEnableRowQuant, \").\");\n }\n\n );););););\n}\n\n} // namespace transformer_engine\n\nvoid nvte_quantize_with_hadamard_transform(const NVTETensor input, NVTETensor output,\n const NVTETensor hadamard_matrix,\n const NVTEQuantizationConfig quant_config,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_quantize_with_hadamard_transform);\n using namespace transformer_engine;\n QuantizationConfig quant_config_cpp;\n if (quant_config != nullptr) {\n quant_config_cpp = *reinterpret_cast(quant_config);\n }\n hadamard_transform_cast_fusion(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output),\n *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp,\n stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/util/hip/hip_runtime.h\"\n#include \"common/util/curanddx.hpp\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"customized_pipeline.cuh\"\n#include \"cutlass/arch/barrier.h\"\n#include \"cutlass/arch/reg_reconfig.h\"\n#include \"cutlass/cluster_launch.hpp\"\n#include \"cutlass/cutlass.h\"\n#include \"cutlass/detail/sm100_blockscaled_layout.hpp\"\n#include \"cutlass/fast_math.h\"\n#include \"cutlass/float8.h\"\n#include \"cutlass/float_subbyte.h\"\n#include \"cutlass/gemm/collective/builders/sm100_common.inl\"\n#include \"cutlass/numeric_conversion.h\"\n#include \"cutlass/numeric_types.h\"\n#include \"cutlass/pipeline/pipeline.hpp\"\n#include \"cutlass/platform/platform.h\"\n#include \"cutlass/util/GPU_Clock.hpp\"\n#include \"cutlass/util/command_line.h\"\n#include \"cutlass/util/print_error.hpp\"\n\n// clang-format off\n\nnamespace transformer_engine {\nnamespace detail {\nnamespace {\n\nusing namespace cute;\n\nstruct CLCResponse { uint32_t data[4] = {0}; };\n\nconstexpr int kFp4ConvertChunkElements = 8;\nconstexpr int kFp4ConvertFullElements = 16;\nconstexpr int kFp4RbitsPerChunk = 2;\nconstexpr int kFp4ChunkCount = kFp4ConvertFullElements / kFp4ConvertChunkElements;\n\n\nCUTLASS_DEVICE\ncutlass::Array StochasticNumericConverterBase(\n cutlass::Array const &input,\n cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n auto output_ptr = reinterpret_cast(&output);\n constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING;\n if constexpr (has_rs) {\n asm volatile(\n \"{\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\\n\"\n \"}\"\n : \"=h\"(output_ptr[0]), \"=h\"(output_ptr[1])\n : \"f\"(input[0]), \"f\"(input[1]), \"f\"(input[2]), \"f\"(input[3]), \"f\"(input[4]), \"f\"(input[5]),\n \"f\"(input[6]), \"f\"(input[7]), \"r\"(rbits[0]), \"r\"(rbits[1]));\n } else {\n NVTE_DEVICE_ERROR(\n \"FP4 cvt PTX instructions are architecture-specific. \"\n \"Try recompiling with sm_XXXa instead of sm_XXX.\");\n }\n return output;\n}\n\nCUTLASS_DEVICE\ncutlass::Array\nStochasticNumericConverter(cutlass::Array const &input,\n cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n cutlass::Array *result_ptr =\n reinterpret_cast *>(&output);\n cutlass::Array const *source_ptr =\n reinterpret_cast const *>(&input);\n cutlass::Array const *rbits_ptr =\n reinterpret_cast const *>(&rbits);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < kFp4ChunkCount; i++) {\n result_ptr[i] = StochasticNumericConverterBase(source_ptr[i], rbits_ptr[i]);\n }\n return output;\n}\n\ntemplate <\n class ElementA,\n class ElementB,\n class ASmemLayout,\n class BSmemLayout,\n class ClusterShape,\n int AccumulatorPipelineStageCount_,\n int EpilogueUnrollFactor_,\n int SchedulerPipelineStageCount_>\nstruct SharedStorage {\n static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;\n static int constexpr EpilogueUnrollFactor = EpilogueUnrollFactor_;\n using AtomThrShapeMNK = cute::Shape<_1, _1, _1>;\n\n using AccumulatorPipeline = cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;\n\n static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline = cutlass::detail::CustomizedPipelineTmaUmmaAsync<\n MainloopPipelineStageCount,\n Shape<_1,_1,_1>,\n AtomThrShapeMNK>;\n using MainloopPipelineStorage = typename MainloopPipeline::SharedStorage;\n using CLCPipeline = cutlass::PipelineCLCFetchAsync;\n using CLCPipelineStorage = typename CLCPipeline::SharedStorage;\n using CLCThrottlePipeline = cutlass::PipelineAsync;\n using CLCThrottlePipelineStorage = typename CLCThrottlePipeline::SharedStorage;\n\n struct TensorStorage : cute::aligned_struct<128, _1> {\n // cute::array_aligned> smem_A;\n cute::array_aligned> smem_A;\n cute::array_aligned> smem_B;\n } tensors;\n\n alignas(16) AccumulatorPipelineStorage accumulator;\n alignas(16) MainloopPipelineStorage mainloop;\n alignas(16) cute::uint64_t tma_barrier[1];\n alignas(16) CLCPipelineStorage clc;\n alignas(16) CLCThrottlePipelineStorage clc_throttle;\n alignas(16) CLCResponse clc_response[SchedulerPipelineStageCount_];\n uint32_t tmem_base_ptr;\n};\n\ntemplate \n__launch_bounds__(512, 1)\n__global__ static void row_col_rht_gemm_device(\n MShape M,\n NShape N,\n KShape K,\n ClusterShape cluster_shape,\n ClusterTileShape cluster_tile,\n TA const* A,\n AStride dA,\n ASmemLayout sAlayout,\n CUTE_GRID_CONSTANT TmaLoadA const tma_load_a,\n TB const* B,\n BStride dB,\n BSmemLayout sBlayout,\n CUTE_GRID_CONSTANT TmaLoadB const tma_load_b,\n TD* D,\n DStride dD,\n DSmemLayout,\n TSFD* SFD,\n TSFDLayout sfd_layout,\n TQA* QA,\n QAStride dQA,\n TSFA* SFA,\n TSFALayout sfa_layout,\n TiledMMA mma,\n float const* a_global_amax,\n float const* c_global_amax,\n const size_t* rng_state) {\n using namespace cute;\n\n // Abort immediately if compilation is not supported\n constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY;\n if constexpr (!is_blackwell_arch) {\n NVTE_DEVICE_ERROR(\"RHT fusion is only supported on Blackwell.\");\n return;\n } else {\n static_assert(kEnableRHTColQuant_ || kEnableRowQuant_,\n \"row_col_rht_gemm_device must generate row-wise \"\n \"and/or column-wise output.\");\n#if !defined(CUTLASS_ARCH_CLC_ENABLED)\n CUTLASS_NOT_IMPLEMENTED();\n return;\n#endif\n\n using X = Underscore;\n // static constexpr bool kApplyStochasticRounding = true;\n using ElementAccumulator = float;\n static int constexpr K_PIPE_MAX = size<3>(ASmemLayout{});\n using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>;\n static uint32_t constexpr kTmaTransactionBytes = cutlass::bits_to_bytes(\n size(AtomThrShapeMNK{}) * cosize(take<0,3>(ASmemLayout{})) * cute::sizeof_bits_v);\n static constexpr bool kEnableStochasticRounding = kEnableStochasticRounding_;\n static constexpr bool kEnableRHTColQuant = kEnableRHTColQuant_;\n static constexpr bool kEnableRowQuant = kEnableRowQuant_;\n static constexpr bool kUseFastMath = kUseFastMath_;\n static int constexpr RhtTensorSize = 16;\n static int constexpr kTmaRhtTensorTransactionBytes = cutlass::bits_to_bytes(\n RhtTensorSize * RhtTensorSize * cute::sizeof_bits_v);\n static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;\n static int constexpr SchedulerPipelineStageCount = SchedulerPipelineStageCount_;\n\n static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline = cutlass::detail::CustomizedPipelineTmaUmmaAsync<\n MainloopPipelineStageCount,\n ClusterShape,\n AtomThrShapeMNK>;\n using MainloopPipelineState = typename MainloopPipeline::PipelineState;\n using CLCPipeline = cutlass::PipelineCLCFetchAsync;\n using CLCPipelineState = typename CLCPipeline::PipelineState;\n using CLCThrottlePipeline = cutlass::PipelineAsync;\n using CLCThrottlePipelineState = typename CLCThrottlePipeline::PipelineState;\n\n static_assert(ClusterShape{} == Shape<_1,_1,_1>{}, \"ClusterShape must be Shape<_1,_1,_1>\");\n\n using TmemAllocator = cute::TMEM::Allocator1Sm;\n static int constexpr VectorSize = RhtTensorSize;\n // Preconditions\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n auto cluster_size = size<0>(cluster_shape);\n auto mainloop_tiler = Shape<_128,_16,_128>{};\n auto epilogue_tiler = Shape<_128,_128,_128>{};\n\n static int constexpr EpilogueUnrollFactor = size<2>(epilogue_tiler) / size<2>(cluster_tile);\n\n // Get the appropriate blocks for this Cluster\n dim3 cluster_coord_in_grid = cluster_id_in_grid();\n\n // Total number of k-tiles\n int const K_TILE_MAX = ceil_div(min(N, K), size<2>(epilogue_tiler));\n\n struct TileScheduler {\n struct WorkTileInfo {\n uint32_t m_idx = 0;\n uint32_t n_idx = 0;\n uint32_t l_idx = 0;\n bool is_valid_tile = false;\n };\n uint32_t tiles_in_m = 0;\n uint32_t tiles_in_n = 0;\n\n int k_tile_max = 0;\n\n int wave_cnt = 0;\n WorkTileInfo work_tile_info;\n WorkTileInfo next_work_tile_info;\n CLCResponse* clc_response_ptr_;\n CUTLASS_DEVICE TileScheduler(uint32_t tiles_m, uint32_t tiles_n, int kmax, CLCResponse* clc_response_ptr)\n : tiles_in_m(tiles_m),\n tiles_in_n(tiles_n),\n\n k_tile_max(kmax),\n work_tile_info({blockIdx.x, blockIdx.y, blockIdx.z, blockIdx.x(\n &clc_response_ptr[state.index()]));\n asm volatile(\n \"{\\n\\t\"\n \"clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes.multicast::cluster::all.b128 [%0], [%1];\\n\\t\"\n \"}\\n\"\n :\n : \"r\"(result_addr), \"r\"(mbarrier_addr));\n #else\n CUTLASS_NOT_IMPLEMENTED();\n #endif\n }\n CUTLASS_DEVICE\n static WorkTileInfo\n work_tile_info_from_clc_response(uint32_t result_addr) {\n WorkTileInfo work_tile_info;\n uint32_t valid = 0;\n #if defined(CUTLASS_ARCH_CLC_ENABLED)\n asm volatile(\n \"{\\n\"\n \".reg .pred p1;\\n\\t\"\n \".reg .b128 clc_result;\\n\\t\"\n \"ld.shared.b128 clc_result, [%4];\\n\\t\"\n \"clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_result;\\n\\t\"\n \"selp.u32 %3, 1, 0, p1;\\n\\t\"\n \"@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_result;\\n\\t\"\n \"}\\n\"\n : \"=r\"(work_tile_info.m_idx), \"=r\"(work_tile_info.n_idx), \"=r\"(work_tile_info.l_idx), \"=r\"(valid)\n : \"r\"(result_addr)\n : \"memory\"\n );\n\n cutlass::arch::fence_view_async_shared();\n #else\n CUTLASS_NOT_IMPLEMENTED();\n #endif\n work_tile_info.is_valid_tile = (valid == 1);\n return work_tile_info;\n }\n };\n\n\n\n // Allocate SMEM\n extern __shared__ char shared_memory[];\n using SharedStorage = SharedStorage;\n SharedStorage& shared_storage = *reinterpret_cast(shared_memory);\n uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile))));\n uint32_t tiles_in_n = uint32_t(size(ceil_div(N, size<2>(epilogue_tiler))));\n TileScheduler scheduler(tiles_in_m, tiles_in_n, K_TILE_MAX, shared_storage.clc_response);\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n auto acc_shape_mma = make_shape(take<0,2>(mainloop_tiler), _1{}, _1{});\n auto acc_shape_epilogue = make_shape(take<0,2>(epilogue_tiler), _1{}, _1{});\n\n auto acc_mainloop_pipelined_shape = append(acc_shape_mma, Int{});\n auto bulk_tmem_mma = TiledMMA::make_fragment_C(acc_mainloop_pipelined_shape);\n\n static int constexpr NumEpilogueColQuantThreadCount = kEnableRHTColQuant ? 128 : 0;\n static int constexpr NumEpilogueRowQuantThreadCount = kEnableRowQuant ? 256 : 0;\n static int constexpr NumMmaThreadCount = kEnableRHTColQuant? 32: 0;\n static int constexpr NumMmaIssueThreadCount = kEnableRHTColQuant? 1: 0;\n static int constexpr NumSchedThreads = 32;\n static int constexpr NumMainloopLoadThreads = 32;\n static int constexpr NumEpilogueThreads = NumEpilogueColQuantThreadCount + NumEpilogueRowQuantThreadCount;\n\n TmemAllocator tmem_allocator{};\n cutlass::arch::NamedBarrier tmem_allocation_result_barrier(\n NumMmaThreadCount + NumEpilogueColQuantThreadCount,\n cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);\n\n int warp_idx = cutlass::canonical_warp_idx_sync();\n\n // warp assignment\n bool is_mma_warp = (warp_idx == 0);\n bool is_dma_warp = (warp_idx == 1);\n bool is_sched_warp = (warp_idx == 2);\n bool is_epilogue_col_quant_warp = (warp_idx >= 4 && warp_idx <= 7);\n bool is_epilogue_row_quant_warp = (warp_idx >= 8 && warp_idx <= 15);\n\n if (is_epilogue_col_quant_warp && elect_one_sync()) {\n cute::prefetch(raw_pointer_cast(c_global_amax));\n }\n if (is_epilogue_row_quant_warp && elect_one_sync()) {\n cute::prefetch(raw_pointer_cast(a_global_amax));\n }\n\n typename MainloopPipeline::Params mainloop_pipeline_params;\n if (is_dma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;\n }\n if (is_mma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;\n }\n mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp;\n mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes;\n mainloop_pipeline_params.initializing_warp = 0;\n mainloop_pipeline_params.num_consumers = NumEpilogueRowQuantThreadCount + NumMmaIssueThreadCount;\n MainloopPipeline mainloop_pipeline(\n shared_storage.mainloop,\n mainloop_pipeline_params,\n cluster_shape,\n cute::true_type{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n\n MainloopPipelineState mainloop_pipe_consumer_state;\n MainloopPipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state();\n\n using AccumulatorPipeline = cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState;\n\n AccumulatorPipelineState accumulator_pipe_consumer_state;\n AccumulatorPipelineState accumulator_pipe_producer_state = cutlass::make_producer_start_state();\n\n typename AccumulatorPipeline::Params accumulator_pipeline_params;\n if (is_mma_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer;\n }\n if (is_epilogue_col_quant_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer;\n }\n // Only one producer thread arrives on this barrier.\n accumulator_pipeline_params.producer_arv_count = 1;\n accumulator_pipeline_params.consumer_arv_count = size(AtomThrShapeMNK{}) * NumEpilogueColQuantThreadCount;\n accumulator_pipeline_params.initializing_warp = 1;\n using IsInitAccumulatorPipeline = cute::conditional_t;\n AccumulatorPipeline accumulator_pipeline(\n shared_storage.accumulator,\n accumulator_pipeline_params,\n cluster_shape,\n IsInitAccumulatorPipeline{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n // CLC pipeline\n typename CLCPipeline::Params clc_pipeline_params;\n if (is_sched_warp) {\n clc_pipeline_params.role = CLCPipeline::ThreadCategory::ProducerConsumer;\n } else {\n clc_pipeline_params.role = CLCPipeline::ThreadCategory::Consumer;\n }\n clc_pipeline_params.producer_blockid = 0;\n clc_pipeline_params.producer_arv_count = 1;\n clc_pipeline_params.consumer_arv_count = NumSchedThreads + cluster_size *\n (NumMainloopLoadThreads + NumEpilogueThreads + NumMmaThreadCount);\n clc_pipeline_params.transaction_bytes = sizeof(CLCResponse);\n clc_pipeline_params.initializing_warp = 3;\n CLCPipeline clc_pipeline(shared_storage.clc, clc_pipeline_params, cluster_shape);\n CLCPipelineState clc_pipeline_consumer_state;\n CLCPipelineState clc_pipeline_producer_state = cutlass::make_producer_start_state();\n\n // CLC throttle pipeline\n typename CLCThrottlePipeline::Params clc_throttle_pipeline_params;\n if (is_dma_warp) {\n clc_throttle_pipeline_params.role = CLCThrottlePipeline::ThreadCategory::Producer;\n }\n if (is_sched_warp) {\n clc_throttle_pipeline_params.role = CLCThrottlePipeline::ThreadCategory::Consumer;\n }\n clc_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads;\n clc_throttle_pipeline_params.consumer_arv_count = NumSchedThreads;\n clc_throttle_pipeline_params.dst_blockid = 0;\n clc_throttle_pipeline_params.initializing_warp = 4;\n\n CLCThrottlePipeline clc_throttle_pipeline(shared_storage.clc_throttle, clc_throttle_pipeline_params);\n CLCThrottlePipelineState clc_pipe_throttle_consumer_state;\n CLCThrottlePipelineState clc_pipe_throttle_producer_state = cutlass::make_producer_start_state();\n\n if (warp_idx == 2 && elect_one_sync()) {\n cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);\n }\n __syncthreads();\n\n if (is_dma_warp) {\n cutlass::arch::warpgroup_reg_dealloc<32>();\n cute::Tensor mA = tma_load_a.get_tma_tensor(make_shape(M,N));\n cute::Tensor mB = tma_load_b.get_tma_tensor(make_shape(RhtTensorSize, RhtTensorSize));\n\n cute::Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_,_, _), Step<_1, X,_1>{});\n cute::Tensor gB_nk = local_tile(mB, cluster_tile, make_coord(_,_, _), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)\n\n cute::Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n cute::Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n cute::Tensor tCgA = thr_mma.partition_A(gA_mk); // (MMA,MMA_M,MMA_K,k)\n cute::Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k)\n\n Layout cta_layout_mnk = make_layout(cluster_shape);\n Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{}));\n auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster);\n\n auto [tAgA, tAsA] = tma_partition(\n tma_load_a,\n get<2>(cta_coord_vmnk),\n make_layout(size<2>(cta_layout_vmnk)),\n group_modes<0,3>(tCsA),\n group_modes<0,3>(tCgA));\n\n auto [tBgB, tBsB] = tma_partition(\n tma_load_b,\n get<1>(cta_coord_vmnk),\n make_layout(size<1>(cta_layout_vmnk)),\n group_modes<0,3>(tCsB),\n group_modes<0,3>(tCgB));\n\n uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);\n uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);\n if constexpr (kEnableRHTColQuant) {\n if (elect_one_sync()) {\n cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0], kTmaRhtTensorTransactionBytes);\n copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_,0,0), tBsB(_,0));\n }\n }\n\n do {\n bool is_first_wave = scheduler.is_first_wave();\n uint32_t skip_wait = is_first_wave;\n auto tAgA_mk = tAgA(_,scheduler.tile_m(),_);\n int k_tile = 0;\n // Throttle CLC producer\n clc_throttle_pipeline.producer_acquire(clc_pipe_throttle_producer_state);\n clc_throttle_pipeline.producer_commit(clc_pipe_throttle_producer_state);\n ++clc_pipe_throttle_producer_state;\n\n CUTLASS_PRAGMA_NO_UNROLL\n while (k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n()) {\n\n int k_tile_idx_n = scheduler.tile_n_base() + k_tile;\n ++k_tile;\n skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount);\n mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state);\n using BarrierType = typename MainloopPipeline::ProducerBarrierType;\n BarrierType* tma_barrier = mainloop_pipeline.producer_get_barrier(\n mainloop_pipe_producer_state);\n int write_stage = mainloop_pipe_producer_state.index();\n ++mainloop_pipe_producer_state;\n if (cute::elect_one_sync()) {\n copy(\n tma_load_a.with(*tma_barrier, tma_mcast_mask_a),\n tAgA_mk(_,k_tile_idx_n),\n tAsA(_,write_stage));\n }\n }\n scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state);\n ++clc_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);\n } else if (is_mma_warp) {\n cutlass::arch::warpgroup_reg_dealloc<32>();\n if constexpr (kEnableRHTColQuant) {\n cute::Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n cute::Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n // Allocate \"fragments\" -- these are actually umma smem descriptors\n cute::Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)\n cute::Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE)\n\n mma.accumulate_ = UMMA::ScaleOut::Zero;\n\n tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);\n __syncwarp();\n tmem_allocation_result_barrier.arrive();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_mma.data() = tmem_base_ptr;\n cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/);\n do {\n uint32_t skip_wait = K_TILE_MAX <= 0;\n\n auto barrier_token = mainloop_pipeline.consumer_try_wait(\n mainloop_pipe_consumer_state,\n skip_wait);\n scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state);\n ++clc_pipeline_consumer_state;\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); ) {\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n int read_stage = mainloop_pipe_consumer_state.index();\n auto tCrA_mk = tCrA(_,_,_,read_stage);\n auto tCrB_nk = tCrB(_,_,0,0);\n CUTLASS_PRAGMA_UNROLL\n for (int k_block = 0; k_block < size<2>(tCrA) / EpilogueUnrollFactor; ++k_block)\n {\n int accumulator_k_block = accumulator_pipe_producer_state.index() * EpilogueUnrollFactor;\n int tCrA_k_block = k_block * EpilogueUnrollFactor;\n accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < EpilogueUnrollFactor; i++) {\n auto accumulators = bulk_tmem_mma(_,_,_,accumulator_k_block + i);\n gemm(mma, tCrA_mk(_,_,tCrA_k_block + i), tCrB_nk, accumulators);\n }\n\n accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);\n ++accumulator_pipe_producer_state;\n }\n auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n skip_wait = k_tile >= K_TILE_MAX;\n mainloop_pipeline.umma_consumer_release(curr_mainloop_pipe_consumer_state);\n barrier_token = mainloop_pipeline.consumer_try_wait(\n mainloop_pipe_consumer_state,\n skip_wait);\n }\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n tmem_allocator.release_allocation_lock();\n accumulator_pipeline.producer_tail(accumulator_pipe_producer_state);\n tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);\n }\n } else if(is_sched_warp) {\n cutlass::arch::warpgroup_reg_dealloc<32>();\n do {\n clc_throttle_pipeline.consumer_wait(clc_pipe_throttle_consumer_state);\n clc_throttle_pipeline.consumer_release(clc_pipe_throttle_consumer_state);\n ++clc_pipe_throttle_consumer_state;\n clc_pipeline_producer_state = scheduler.advance_to_next_work(clc_pipeline, clc_pipeline_producer_state);\n scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state);\n ++clc_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n } else if (is_epilogue_col_quant_warp) {\n cutlass::arch::warpgroup_reg_alloc<192>();\n if constexpr (kEnableRHTColQuant) {\n using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x;\n\n float const c_global_amax_val = *c_global_amax;\n auto acc_epilogue_pipelined_shape = append(acc_shape_epilogue, Int{});\n auto bulk_tmem_epilogue_layout = make_layout(\n acc_epilogue_pipelined_shape,\n make_stride(\n stride<0>(bulk_tmem_mma),\n Int<0>{},\n Int<0>{},\n size<1>(epilogue_tiler)));\n auto bulk_tmem_epilogue = make_tensor(make_tmem_ptr(), bulk_tmem_epilogue_layout);\n\n // leveraging 256-bit writes to global memory\n static int constexpr FragmentSize = 256 / sizeof_bits_v;\n\n tmem_allocation_result_barrier.arrive_and_wait();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_epilogue.data() = tmem_base_ptr;\n int global_thread_idx = threadIdx.x;\n int local_thread_idx = global_thread_idx % cutlass::NumThreadsPerWarpGroup;\n\n size_t rng_seed = 0;\n size_t rng_offset = 0;\n if constexpr (kEnableStochasticRounding) {\n rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0;\n rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0;\n }\n\n cute::Tensor mD = make_tensor(\n cute::subbyte_iterator(D),\n make_shape(M,N),\n dD); // (M,N)\n cute::Tensor gD_mn = local_tile(\n mD,\n epilogue_tiler,\n make_coord(_,_, _),\n Step<_1,_1, X>{}); // (BLK_M,BLK_N)\n cute::Tensor pD = make_identity_tensor(mD.shape());\n cute::Tensor pD_mn = local_tile(\n pD,\n epilogue_tiler,\n make_coord(_,_, _),\n Step<_1,_1, X>{}); // (BLK_M,BLK_N)\n cute::Tensor mSFD = make_tensor(make_gmem_ptr(SFD), sfd_layout);\n cute::Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N)\n cute::Tensor pSFD = make_identity_tensor(mSFD.shape());\n cute::Tensor pSFD_mn = local_tile(pSFD, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N)\n\n cute::Tensor gD_mn_view = tiled_divide(gD_mn, take<0,2>(epilogue_tiler));\n cute::Tensor pD_mn_view = tiled_divide(pD_mn, take<0,2>(epilogue_tiler));\n auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_,_,_,_0{}));\n auto tiled_r2g = make_tiled_copy_D(\n Copy_Atom{},\n tiled_t2r);\n auto thr_t2r = tiled_t2r.get_slice(local_thread_idx);\n auto thr_r2g = tiled_r2g.get_slice(local_thread_idx);\n\n // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release}\n static constexpr float fp4_max = 6.0f;\n static constexpr float fp8_max = 448.0f;\n float const fp4_max_inv = 1.0f / fp4_max;\n float const global_encode_scale = c_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / c_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n\n float const global_decode_scale = 1.0f / global_encode_scale;\n // Scaling factor for fast math path\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n auto sfc_converter = cutlass::NumericConverter{};\n\n do {\n scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state);\n ++clc_pipeline_consumer_state;\n for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); ++k_tile) {\n cute::Tensor tDgD_mn = gD_mn_view(_,_,_,scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n cute::Tensor tDgSFD_mn = gSFD_mn(_,_,scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n cute::Tensor tDpD_mn = pD_mn_view(_,_,_,scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n cute::Tensor tDpSFD_mn = pSFD_mn(_,_,scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n\n accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state);\n\n auto Acc = bulk_tmem_epilogue(_,_,_,accumulator_pipe_consumer_state.index());\n cute::Tensor tDtAcc = thr_t2r.partition_S(Acc); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n cute::Tensor tDgD = thr_t2r.partition_D(tDgD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n cute::Tensor tDpD = thr_t2r.partition_D(tDpD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n cute::Tensor tTR_rAcc = make_tensor(shape(tDgD)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n cute::Tensor tDrD = make_tensor(shape(tDgD));\n cute::Tensor tTR_rAcc_frag = recast>(coalesce(tTR_rAcc));\n cute::Tensor tDrD_frag = recast>(coalesce(tDrD));\n\n cute::Tensor src = thr_r2g.retile_S(tDrD);\n cute::Tensor dst = thr_r2g.retile_D(tDgD);\n cute::Tensor pSrc = thr_r2g.retile_D(tDpD);\n\n cute::Tensor tDgSFD_view = make_tensor(\n tDgSFD_mn.data(),\n make_layout(\n make_shape(shape(tDgSFD_mn), Int<1>{}, Int<1>{}),\n make_stride(stride(tDgSFD_mn), Int<0>{}, Int<0>{})));\n cute::Tensor tDpSFD_view = make_tensor(\n tDpSFD_mn.data(),\n make_layout(\n make_shape(shape(tDpSFD_mn), Int<1>{}, Int<1>{}),\n make_stride(stride(tDpSFD_mn), Int<0>{}, Int<0>{})));\n cute::Tensor tDgSFD = filter(thr_t2r.partition_D(tDgSFD_view));\n cute::Tensor tDrSFD = make_tensor(shape(tDgSFD));\n cute::Tensor tDpSFD = filter(thr_t2r.partition_D(tDpSFD_view));\n static int constexpr NumVecs = size(tDgD) / VectorSize;\n cute::Tensor tD_rRowSFD_frg = recast>(tDrSFD);\n\n cutlass::maximum_absolute_value_reduction, true> amax_reduction;\n cutlass::Array vec_maxs;\n cutlass::Array pvscales;\n // TMEM_LOAD\n copy(tiled_t2r, tDtAcc, tTR_rAcc);\n cutlass::arch::fence_view_async_tmem_load();\n accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state);\n ++accumulator_pipe_consumer_state;\n\n if constexpr (!kUseFastMath) {\n // Downcast to BF16 for bit-wise compatibility with\n // unfused kernels\n auto convert_accum_to_bf16 =\n cutlass::NumericArrayConverter{};\n auto convert_bf16_to_accum =\n cutlass::NumericArrayConverter{};\n tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{})));\n tTR_rAcc_frag(_1{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_1{})));\n }\n\n auto compute_frgs = reinterpret_cast *>(tTR_rAcc_frag.data());\n auto output_frgs = reinterpret_cast *>(tDrD_frag.data());\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]);\n }\n\n pvscales = cutlass::multiplies>{}(\n vec_maxs, global_encode_scale_multiplier);\n auto pvscales_cvted = cutlass::NumericArrayConverter{}(pvscales);\n\n tD_rRowSFD_frg(_0{}) = pvscales_cvted;\n auto qpvscale_ups = cutlass::NumericArrayConverter{}(tD_rRowSFD_frg(_0{}));\n auto qpvscale_scaled = cutlass::multiplies>{}(\n qpvscale_ups,\n global_decode_scale);\n\n cutlass::Array acc_scales;\n if constexpr (kUseFastMath) {\n // fast math: use reciprocal approximate to replace div\n acc_scales = cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled);\n } else {\n // regular path for slower math, use divide to replace div\n acc_scales = cutlass::divides>{}(1.0, qpvscale_scaled);\n }\n\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n transformer_engine::curanddx::detail::philox4x32_native_state rng;\n // \"Prefetch\" a stochastic rounding state for the first tile\n if constexpr (kEnableStochasticRounding) {\n const size_t rng_sequence = global_thread_idx + k_tile * 512 + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512;\n rng.init(rng_seed, rng_sequence, rng_offset);\n }\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(\n acc_scales[v],\n cutlass::platform::numeric_limits::max());\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(cutlass::multiplies>{}(compute_frgs[v], acc_scale), *reinterpret_cast*>(&random_uint4));\n } else {\n output_frgs[v] = cutlass::NumericArrayConverter{}(\n cutlass::multiplies>{}(\n compute_frgs[v],\n acc_scale));\n }\n\n }\n\n cute::Tensor pred_pSrc = cute::lazy::transform(make_tensor(counting_iterator{}, replace<0>(shape(dst), _1{})), [&](auto coord){\n cute::Tensor pSrc_view = group_modes<1,rank(pSrc)>(pSrc);\n return elem_less(pSrc_view(_0{},coord), shape(mD));\n });\n copy_if(tiled_r2g, pred_pSrc, src, dst);\n // 32bit vectorization copy 4 e4m3 SFD for per 64 or(16,4):(0, 1) element\n\n constexpr int vec_len = 32 / sizeof_bits_v;\n cute::Tensor tDrSFD_v = recast>(tDrSFD);\n cute::Tensor tDgSFD_v = recast>(tDgSFD);\n copy_if(\n [&](auto coord){\n cute::Tensor tDpSFD_view = group_modes<1,rank(tDpSFD)>(tDpSFD);\n return elem_less(tDpSFD_view(_0{}, coord * vec_len), shape(mSFD));\n },\n tDrSFD_v, tDgSFD_v);\n }\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n }\n } else if (is_epilogue_row_quant_warp) {\n cutlass::arch::warpgroup_reg_alloc<136>();\n if constexpr (kEnableRowQuant) {\n using S2RVectorType = uint128_t;\n float const a_global_amax_val = *a_global_amax;\n int global_thread_idx = threadIdx.x;\n int local_thread_idx = global_thread_idx % 256;\n size_t rng_seed = 0;\n size_t rng_offset = 0;\n if constexpr (kEnableStochasticRounding) {\n rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0;\n rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0;\n }\n cute::Tensor mQA = make_tensor(cute::subbyte_iterator(QA), make_layout(make_shape(M, N), dQA));\n cute::Tensor gQA_mn = local_tile(mQA, epilogue_tiler, make_coord(_,_, _), Step<_1,X,_1>{});\n cute::Tensor pQA = make_identity_tensor(mQA.shape());\n cute::Tensor pQA_mn = local_tile(pQA, epilogue_tiler, make_coord(_,_, _), Step<_1,X,_1>{});\n\n cute::Tensor mSFA = make_tensor(make_gmem_ptr(SFA), sfa_layout);\n cute::Tensor gSFA_mn = local_tile(mSFA, epilogue_tiler, make_coord(_,_, _), Step<_1,X,_1>{}); // (BLK_M,BLK_N)\n cute::Tensor pSFA = make_identity_tensor(mSFA.shape());\n cute::Tensor pSFA_mn = local_tile(pSFA, epilogue_tiler, make_coord(_,_, _), Step<_1,X,_1>{});\n cute::Tensor sA = as_position_independent_swizzle_tensor(\n group_modes<0,2>(coalesce(make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), sAlayout)))); // (BLOCK_M, BLOCK_M,PIPE)\n using S2RWarpLayout = Layout>;\n using WarpGroupLayout = Layout>;\n using S2RThreadLayout = decltype(blocked_product(S2RWarpLayout{}, WarpGroupLayout{}));\n using S2RValLayout = Layout, _1>>;\n using S2RAtomA = Copy_Atom;\n using R2GAtomQA = Copy_Atom;\n\n auto tiled_s2r = make_tiled_copy(S2RAtomA{}, S2RThreadLayout{}, S2RValLayout{});\n auto tiled_r2g_QA = make_tiled_copy_D(R2GAtomQA{}, tiled_s2r);\n\n auto thr_s2r = tiled_s2r.get_slice(local_thread_idx);\n auto thr_r2g_QA = tiled_r2g_QA.get_slice(local_thread_idx);\n\n cute::Tensor tQAsA = thr_s2r.partition_S(sA); // (Copy, Copy_M, Copy_N, PIPE)\n\n cute::Tensor tQArA = make_tensor_like(make_layout(tQAsA(_, _, _, _0{}).shape())); // (Copy, Copy_M, Copy_N)\n // Tensor tQArA_PI = thr_s2r.partition_S(sA_PI);\n cute::Tensor tQAgQA = thr_r2g_QA.partition_D(gQA_mn);\n cute::Tensor tQArQA = make_tensor_like(tQAgQA(_, _, _, _0{}, _0{}));\n cute::Tensor tQApQA = thr_r2g_QA.partition_D(pQA_mn);\n\n cute::Tensor tQAgSFA = thr_s2r.partition_D(gSFA_mn);\n cute::Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{}));\n cute::Tensor tQApSFA = thr_s2r.partition_D(pSFA_mn);\n\n // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release}\n static constexpr float fp4_max = 6.0f;\n static constexpr float fp8_max = 448.0f;\n float const fp4_max_inv = 1.0f / fp4_max;\n float const global_encode_scale = a_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / a_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n\n float const global_decode_scale = 1.0f / global_encode_scale;\n // Scaling factor for fast math path\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n\n auto sfa_converter = cutlass::NumericConverter{};\n do {\n uint32_t skip_wait = K_TILE_MAX <= 0;\n\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); ) {\n auto tQAgSFA_mn = tQAgSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto tQAgQA_mn = tQAgQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto tQApSFA_mn = tQApSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto tQApQA_mn = tQApQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto barrier_token = mainloop_pipeline.consumer_try_wait(\n mainloop_pipe_consumer_state);\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n copy(tiled_s2r, tQAsA(_, _, _, mainloop_pipe_consumer_state.index()), tQArA);\n cutlass::arch::fence_view_async_shared();\n auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n skip_wait = k_tile >= K_TILE_MAX;\n mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state);\n // static int constexpr NumVecs = size(tQArA) / VectorSize;\n cutlass::maximum_absolute_value_reduction, true> amax_reduction;\n auto compute_frgs = reinterpret_cast *>(tQArA.data());\n auto output_frgs = reinterpret_cast *>(raw_pointer_cast(tQArQA.data()));\n transformer_engine::curanddx::detail::philox4x32_native_state rng;\n if constexpr (kEnableStochasticRounding) {\n const size_t rng_sequence = global_thread_idx + k_tile * 512 + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512;\n rng.init(rng_seed, rng_sequence, rng_offset);\n }\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < size(tQArA)/VectorSize; v++) {\n auto compute_frgs_up = cutlass::NumericArrayConverter{}(compute_frgs[v]);\n auto amax = amax_reduction(ElementAccumulator(0), compute_frgs_up);\n // declare pvscales\n ElementAccumulator pvscales;\n pvscales = cutlass::multiplies{}(amax, global_encode_scale_multiplier);\n filter(tQArSFA)(v) = sfa_converter(pvscales);\n auto qpvscale_ups = cutlass::NumericConverter{}(filter(tQArSFA)(v));\n auto qpvscale_scaled = cutlass::multiplies{}(qpvscale_ups, global_decode_scale);\n ElementAccumulator acc_scales;\n if constexpr (kUseFastMath) {\n // fast math: use reciprocal approximate to replace div\n acc_scales = cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled);\n } else {\n // regular path for slower math, use divide to replace div\n acc_scales = cutlass::divides{}(1.0, qpvscale_scaled);\n }\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(\n acc_scales,\n cutlass::platform::numeric_limits::max());\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(cutlass::multiplies>{}(compute_frgs_up, acc_scale), *reinterpret_cast*>(&random_uint4));\n } else {\n output_frgs[v] = cutlass::NumericArrayConverter{}(\n cutlass::multiplies>{}(\n compute_frgs_up,\n acc_scale));\n }\n }\n\n cute::Tensor pred_tQApQA = cute::lazy::transform(make_tensor(counting_iterator{}, replace<0>(shape(tQAgQA_mn), _1{})), [&](auto coord){\n cute::Tensor tQApQA_view = group_modes<1,rank(tQApQA_mn)>(tQApQA_mn);\n return elem_less(tQApQA_view(_0{}, coord), shape(mQA));\n });\n copy_if(tiled_r2g_QA, pred_tQApQA, tQArQA, tQAgQA_mn);\n // 32bit vectorization copy 4 e4m3 SFA for per 64 or (16,4):(0, 1) element\n constexpr int vec_len = 32 / sizeof_bits_v;\n cute::Tensor tQArSFA_v = recast>(filter(tQArSFA));\n cute::Tensor tQAgSFA_v = recast>(filter(tQAgSFA_mn));\n copy_if(\n [&](auto coord){\n cute::Tensor tQApSFA_view = filter(tQApSFA_mn);\n return elem_less(tQApSFA_view(_0{}, coord * vec_len), shape(mSFA));\n },\n tQArSFA_v, tQAgSFA_v);\n }\n scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state);\n ++clc_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n }while (scheduler.is_valid());\n }\n } else {\n cutlass::arch::warpgroup_reg_dealloc<32>();\n }\n } // sm100 compile guard end\n} // NOLINT(readability/fn_size)\n\n\n// this function computes RHT-GEMM for\n// m = hidden_size, n = sequence_length\n// A: m x n: col-major\n// B: 16 x 16: row-major\n// D: m x n: row-major\n// SFD: m x (n/16): row-major\n// QA: m x n: col-major\n// SFA: m/16 x n: col-major\ntemplate \nvoid row_col_rht_gemm_ntt_w_sfc(\n int sequence_length,\n int hidden_size,\n TA const* A,\n TB const* B,\n TD* D,\n TSFD* SFD,\n TQA* QA,\n TSFA* SFA,\n float const* a_global_amax,\n float const* d_global_amax,\n const size_t* rng_state,\n uint32_t sm_count,\n hipStream_t stream,\n int k_tile_size = 1024) {\n using namespace cute;\n static int constexpr SFVecSize = 16;\n static int constexpr RhtTensorSize = 16;\n\n static_assert(RhtTensorSize == 16, \"RhtTensorSize must be 16\");\n using LinearSFALayout = decltype(make_layout(make_shape(make_shape(Int{}, 0), 0), make_stride(make_stride(_0{}, _1{}), 0)));\n using LinearSFCLayout = decltype(make_layout(make_shape(0, make_shape(Int{}, 0)), make_stride(0, make_stride(_0{}, _1{}))));\n\n using SwizzledSFALayoutAtom = cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFDLayoutAtom = cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFALayout = decltype(tile_to_shape(SwizzledSFALayoutAtom{}, make_shape(hidden_size,sequence_length), Step<_1,_2>{}));\n using SwizzledSFDLayout = decltype(tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(hidden_size,sequence_length), Step<_2,_1>{}));\n\n using SFALayout = cute::conditional_t;\n using SFCLayout = cute::conditional_t;\n SFALayout sfa_layout;\n SFCLayout sfd_layout;\n\n if constexpr (kEnableSwizzleSFOutput) {\n sfa_layout = tile_to_shape(SwizzledSFALayoutAtom{}, make_shape(hidden_size, sequence_length), Step<_1,_2>{});\n sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(hidden_size, sequence_length), Step<_2,_1>{});\n } else {\n sfa_layout = make_layout(make_shape(make_shape(Int{}, hidden_size/SFVecSize), sequence_length), make_stride(make_stride(_0{}, _1{}), hidden_size/SFVecSize));\n sfd_layout = make_layout(make_shape(hidden_size, make_shape(Int{}, sequence_length/SFVecSize)), make_stride(sequence_length/SFVecSize, make_stride(_0{}, _1{})));\n }\n // Define shapes (dynamic)\n auto M = hidden_size;\n auto N = sequence_length;\n cute::Tensor tensorA = make_tensor(A, make_shape(hidden_size, sequence_length), LayoutLeft{});\n cute::Tensor tensorB = make_tensor(B, make_shape(RhtTensorSize, RhtTensorSize), LayoutLeft{});\n cute::Tensor tensorD = make_tensor(D, make_shape(hidden_size, sequence_length), LayoutRight{});\n cute::Tensor tensorQA = make_tensor(QA, make_shape(hidden_size, sequence_length), LayoutLeft{});\n cute::Tensor tensorSFD = make_tensor(SFD, sfd_layout);\n cute::Tensor tensorSFA = make_tensor(SFA, sfa_layout);\n // Define strides (from tensors)\n auto dA = stride(tensorA); // (dM,dK)\n auto dB = stride(tensorB); // (dN,dK)\n auto dD = stride(tensorD); // (dM,dN)\n auto dQA = stride(tensorQA); // (dM,dK)\n using ClusterShape = Shape< _1, _1, _1>;\n auto cluster_shape = ClusterShape{};\n auto cluster_tile_shape = Shape<_128,Int,Int>{};\n auto cluster_tile_mainloop = Shape<_128,Int,_128>{};\n\n // Each mainloop / epilogue loads 128 x 64 tiles while each MMA proceeds with 128 x 16 tiles\n static int constexpr EpilogueUnrollFactor =\n size<2>(cluster_tile_mainloop) / size<2>(cluster_tile_shape);\n // Construct the MMA\n auto mma = make_tiled_mma(SM100_MMA_F16BF16_SS(cluster_tile_shape), size<1>(cluster_tile_shape),\n UMMA::Major::MN, UMMA::Major::MN>{},\n Layout>{});\n\n // Assert that the TiledMMA uses all CTAs in the CGA.\n CUTE_STATIC_ASSERT_V(size(cluster_shape) == size(mma));\n CUTE_STATIC_ASSERT_V(evenly_divides(cluster_tile_shape, tile_shape(mma)));\n\n // Determine the A and B shapes\n auto mma_shape_B = partition_shape_B(mma, make_shape(size<1>(cluster_tile_shape), size<2>(cluster_tile_shape)));\n\n using TiledMma = decltype(mma);\n using AtomThrID = typename TiledMma::AtomThrID;\n\n using SmemShape_M = decltype(shape_div(shape<0>(cluster_tile_shape), shape_div(shape<0>(cluster_tile_shape), size<0>(cluster_tile_shape) / size(AtomThrID{}))));\n using SmemShape_N = decltype(shape_div(shape<1>(cluster_tile_shape), shape_div(shape<1>(cluster_tile_shape), size<1>(cluster_tile_shape) / size(AtomThrID{}))));\n using SmemShape_K = decltype(cute::get<2>(cluster_tile_shape));\n\n using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<\n cute::UMMA::Major::MN, TB, SmemShape_N, SmemShape_K>());\n\n auto mma_shape_A = partition_shape_A(mma, make_shape(size<0>(cluster_tile_mainloop), size<2>(cluster_tile_mainloop)));\n using SmemShape_M_A = decltype(shape_div(shape<0>(cluster_tile_mainloop), shape_div(shape<0>(cluster_tile_mainloop), size<0>(cluster_tile_mainloop) / size(AtomThrID{}))));\n using SmemShape_K_A = decltype(cute::get<2>(cluster_tile_mainloop));\n using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<\n cute::UMMA::Major::MN, TA, SmemShape_M_A, SmemShape_K_A>());\n\n static uint32_t constexpr TotalTmemRows = 128;\n static uint32_t constexpr Sm100TmemCapacityColumns = 512;\n static uint32_t constexpr TotalTmem = TotalTmemRows * Sm100TmemCapacityColumns;\n static uint32_t constexpr AccumulatorPipelineStageCount =\n TotalTmem /\n (cute::size<0>(cluster_tile_shape) * cute::size<1>(cluster_tile_shape));\n\n // Define the smem layouts (static)\n // Calculate max pipeline stages based on Blackwell SM100's 232KB shared memory\n constexpr int SchedulerPipelineStageCount = 6;\n static int constexpr MainloopPipelineBytes = sizeof(typename cutlass::detail::CustomizedPipelineTmaUmmaAsync<\n 1,\n Shape<_1,_1,_1>,\n Shape<_1, _1, _1>>::SharedStorage);\n\n static int constexpr ClcResponseBytes = sizeof(CLCResponse) * SchedulerPipelineStageCount;\n static int constexpr CLCThrottlePipelineBytes = sizeof(typename cutlass::PipelineAsync::SharedStorage);\n static int constexpr CLCPipelineBytes = sizeof(typename cutlass::PipelineCLCFetchAsync::SharedStorage);\n static int constexpr TmemDeallocBytes = sizeof(cutlass::arch::ClusterBarrier);\n static int constexpr BTensorBytes = cute::size(mma_shape_B) * sizeof(TB);\n static int constexpr AccPipelineBytes = sizeof(typename cutlass::PipelineUmmaAsync>::SharedStorage);\n static int constexpr TmemBasePtrsBytes = sizeof(uint32_t);\n static int constexpr kBlackwellSmemSize = 232448; // 232KB in bytes\n static int constexpr kBytesPerStage =\n cute::size(mma_shape_A) * sizeof(TA) + MainloopPipelineBytes;\n static int constexpr kReservedBytes = ClcResponseBytes + CLCThrottlePipelineBytes + TmemBasePtrsBytes +\n CLCPipelineBytes + TmemDeallocBytes+BTensorBytes + AccPipelineBytes; // Reserve for barriers and other uses\n static int constexpr kMaxStages = (kBlackwellSmemSize - kReservedBytes) / kBytesPerStage;\n auto sP = Int{}; // SMEM pipelines\n auto sA = UMMA::tile_to_mma_shape(\n SmemLayoutAtomA{},\n append(mma_shape_A, sP), Step<_2,_1,_3>{}); // (MMA,MMA_M,MMA_K,PIPE)\n auto sB = UMMA::tile_to_mma_shape(\n SmemLayoutAtomB{},\n append(mma_shape_B, _1{})); // (MMA,MMA_N,MMA_K, _1)\n auto sD = Layout<_1>{}; // XXX Dummy\n\n auto tma_load_a = make_tma_copy_A_sm100(\n SM90_TMA_LOAD{},\n tensorA,\n sA(_,_,_,0),\n cluster_tile_mainloop,\n mma);\n auto tma_load_b = make_tma_copy_B_sm100(\n SM90_TMA_LOAD{},\n tensorB,\n sB(_,_,_,0),\n cluster_tile_shape,\n mma);\n\n // Assert checks problem size should be multiple of 64\n NVTE_CHECK(M % 64 == 0, \"M must be a multiple of 64, but got \", M);\n NVTE_CHECK(N % 64 == 0, \"N must be a multiple of 64, but got \", N);\n\n uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile_shape))));\n uint32_t tiles_in_n = uint32_t(size(ceil_div(N, k_tile_size)));\n uint32_t tiles = tiles_in_m * tiles_in_n;\n\n dim3 dimBlock(512);\n dim3 dimCluster(size<0>(cluster_shape), size<1>(cluster_shape), size<2>(cluster_shape));\n dim3 dimGrid(tiles_in_m, tiles_in_n, 1);\n\n int smem_size = sizeof(\n SharedStorage<\n TA,\n TB,\n decltype(sA),\n decltype(sB),\n ClusterShape,\n AccumulatorPipelineStageCount,\n EpilogueUnrollFactor,\n SchedulerPipelineStageCount>);\n\n auto* kernel_ptr = &row_col_rht_gemm_device<\n decltype(M), decltype(N), decltype(k_tile_size),\n decltype(cluster_shape), decltype(cluster_tile_shape),\n TA, decltype(dA), decltype(sA), decltype(tma_load_a),\n TB, decltype(dB), decltype(sB), decltype(tma_load_b),\n TD, decltype(dD), decltype(sD),\n TSFD, decltype(sfd_layout),\n TQA, decltype(dQA),\n TSFA, decltype(sfa_layout),\n decltype(mma),\n AccumulatorPipelineStageCount,\n SchedulerPipelineStageCount,\n kEnableStochasticRounding,\n kEnableRHTColQuant,\n kEnableRowQuant,\n kUseFastMath>;\n\n NVTE_CHECK_CUDA(hipFuncSetAttribute(reinterpret_cast(*kernel_ptr), hipFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n\n cutlass::ClusterLaunchParams params = {dimGrid, dimBlock, dimCluster, smem_size, stream};\n cutlass::Status status = cutlass::launch_kernel_on_cluster(\n params, (void const *)kernel_ptr, M, N, k_tile_size, cluster_shape, cluster_tile_shape,\n tensorA.data(), dA, sA, tma_load_a,\n tensorB.data(), dB, sB, tma_load_b,\n tensorD.data(), dD, sD,\n tensorSFD.data(), sfd_layout,\n tensorQA.data(), dQA,\n tensorSFA.data(), sfa_layout,\n mma, a_global_amax, d_global_amax, rng_state);\n\n NVTE_CHECK_CUDA(hipGetLastError());\n NVTE_CHECK(status == cutlass::Status::kSuccess, \"Kernel launch failed.\");\n\n}\n\n} // namespace\n} // namespace detail\n\n// clang-format on\n\nvoid hadamard_transform_cast_fusion(const Tensor &input_, Tensor &output_,\n const Tensor &hadamard_matrix_, QuantizationConfig quant_config,\n hipStream_t stream) {\n NVTE_API_CALL(hadamard_transform_cast_fusion);\n\n // Check input and output tensors\n NVTE_CHECK(input_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor must be BF16 tensor, but scaling mode is \",\n to_string(input_.scaling_mode), \".\");\n NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16,\n \"Input tensor must be BF16 tensor, but dtype is \", to_string(input_.dtype()), \".\");\n NVTE_CHECK(input_.dim() >= 2, \"Input must be a 2D tensor.\");\n const SimpleTensor &input = input_.data;\n\n // rowwise cast and columnwise cast has different output data pointers\n bool has_rowwise_quant = false;\n bool has_columnwise_quant = false;\n void *rowwise_data_ptr = nullptr;\n void *rowwise_scale_inv_ptr = nullptr;\n void *rowwise_amax_ptr = nullptr;\n void *columnwise_data_ptr = nullptr;\n void *columnwise_scale_inv_ptr = nullptr;\n void *columnwise_amax_ptr = nullptr;\n\n // examine the output tensor (single tensor for dense)\n if (output_.data.dptr != nullptr) {\n has_rowwise_quant = true;\n rowwise_data_ptr = output_.data.dptr;\n rowwise_scale_inv_ptr = output_.scale_inv.dptr;\n rowwise_amax_ptr = output_.amax.dptr;\n }\n\n if (output_.columnwise_data.dptr != nullptr) {\n has_columnwise_quant = true;\n columnwise_data_ptr = output_.columnwise_data.dptr;\n columnwise_scale_inv_ptr = output_.columnwise_scale_inv.dptr;\n columnwise_amax_ptr = output_.columnwise_amax.dptr;\n }\n\n NVTE_CHECK(has_rowwise_quant || has_columnwise_quant,\n \"Output tensor must have rowwise or columnwise quant.\");\n\n // Stochastic rounding config\n const bool use_stochastic_rounding = quant_config.stochastic_rounding;\n const size_t *rng_state = nullptr;\n if (quant_config.rng_state != nullptr) {\n Tensor &rng_state_tensor = *convertNVTETensor(quant_config.rng_state);\n NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64,\n \"RNG state should contain 2 64-bit values.\");\n NVTE_CHECK(rng_state_tensor.data.shape == std::vector{2},\n \"Shape of the RNG state should be [2], but got \", rng_state_tensor.data.shape);\n rng_state = reinterpret_cast(rng_state_tensor.data.dptr);\n }\n\n // Template arguments\n using TA = cute::bfloat16_t;\n using TB = cute::bfloat16_t;\n using TD = cutlass::float_e2m1_t;\n using TSFD = cutlass::float_ue4m3_t;\n using TQA = TD;\n using TSFA = TSFD;\n\n checkCuDriverContext(stream);\n\n // Check Hadamard matrix\n constexpr int kHadamardDimension = 16;\n NVTE_CHECK(hadamard_matrix_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Hadamard matrix must be BF16 tensor, but scaling mode is \",\n to_string(hadamard_matrix_.scaling_mode), \".\");\n NVTE_CHECK(hadamard_matrix_.dtype() == transformer_engine::DType::kBFloat16,\n \"Hadamard matrix must be BF16 tensor, but dtype is \",\n to_string(hadamard_matrix_.dtype()), \".\");\n const SimpleTensor &hadamard_matrix = hadamard_matrix_.data;\n NVTE_CHECK(\n (hadamard_matrix_.shape() == std::vector{kHadamardDimension, kHadamardDimension}),\n \"Hadamard matrix must have shape=\",\n std::vector{kHadamardDimension, kHadamardDimension},\n \", but got shape=\", hadamard_matrix_.shape(), \".\");\n const size_t hadamard_dimension = hadamard_matrix.shape[0];\n\n const size_t ndim = input.shape.size();\n const size_t n = input.shape[ndim - 1];\n size_t m = 1;\n for (size_t i = 0; i < ndim - 1; ++i) {\n m *= input.shape[i];\n }\n\n auto sm_count = transformer_engine::cuda::sm_count();\n\n NVTE_CHECK(n % hadamard_dimension == 0, \"row_length must be divisible by hadamard_dimension.\");\n\n NVTE_CHECK(m % hadamard_dimension == 0, \"num_rows must be divisible by hadamard_dimension\");\n\n int k_tile_size = 1024;\n\n // TODO: add support for swizzle sf output\n const bool use_swizzle_sf_output = false;\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_stochastic_rounding, kEnableStochasticRounding,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n has_columnwise_quant, kEnableRhtColQuant,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n has_rowwise_quant, kEnableRowQuant,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_swizzle_sf_output, kEnableSwizzleSFOutput,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n quant_config.use_fast_math, kUseFastMath,\n\n if constexpr (kEnableRhtColQuant || kEnableRowQuant) {\n detail::row_col_rht_gemm_ntt_w_sfc<\n kEnableStochasticRounding, kEnableRhtColQuant, kEnableRowQuant,\n kEnableSwizzleSFOutput, TA, TB, TD, TSFD, TQA, TSFA, kUseFastMath>(\n /*sequence_length=*/m, /*hidden_size=*/n,\n /*A=*/reinterpret_cast(input.dptr),\n /*B=*/reinterpret_cast(hadamard_matrix.dptr),\n /*D=*/reinterpret_cast(columnwise_data_ptr),\n /*SFD=*/reinterpret_cast(columnwise_scale_inv_ptr),\n /*QA=*/reinterpret_cast(rowwise_data_ptr),\n /*SFA=*/reinterpret_cast(rowwise_scale_inv_ptr),\n /*a_global_amax=*/reinterpret_cast(rowwise_amax_ptr),\n /*d_global_amax=*/reinterpret_cast(columnwise_amax_ptr),\n /*rng_state=*/rng_state, /*sm_count=*/sm_count,\n /*stream=*/stream, /*k_tile_size=*/k_tile_size);\n } else {\n NVTE_ERROR(\"Invalid kernel configuration (kEnableRHTColQuant=\",\n kEnableRhtColQuant, \", kEnableRowQuant=\", kEnableRowQuant, \").\");\n }\n\n );););););\n}\n\n} // namespace transformer_engine\n\nvoid nvte_quantize_with_hadamard_transform(const NVTETensor input, NVTETensor output,\n const NVTETensor hadamard_matrix,\n const NVTEQuantizationConfig quant_config,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_quantize_with_hadamard_transform);\n using namespace transformer_engine;\n QuantizationConfig quant_config_cpp;\n if (quant_config != nullptr) {\n quant_config_cpp = *reinterpret_cast(quant_config);\n }\n hadamard_transform_cast_fusion(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output),\n *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp,\n stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/util/cuda_runtime.h\"\n#include \"common/util/curanddx.hpp\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"cutlass/arch/barrier.h\"\n#include \"cutlass/cutlass.h\"\n#include \"cutlass/gemm/collective/builders/sm100_common.inl\"\n#include \"cutlass/numeric_conversion.h\"\n#include \"cutlass/pipeline/pipeline.hpp\"\n#include \"cutlass/util/GPU_Clock.hpp\"\n#include \"cutlass/util/command_line.h\"\n#include \"cutlass/util/print_error.hpp\"\n\n// clang-format off\n\nnamespace transformer_engine {\nnamespace detail {\nnamespace {\n\nusing namespace cute;\nusing cute::Tensor; // Ensure unqualified Tensor refers to cute::Tensor, not transformer_engine::Tensor\n\n// calculate the global encode scale factor for a given global amax.\n__device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax) {\n constexpr float kFP8E4M3Max = 448.0f;\n constexpr float kFP4E2M1Max = 6.0f;\n // If scale is infinity, return max value of float32\n float global_encode_scale = cutlass::minimum_with_nan_propagation{}(\n kFP8E4M3Max * kFP4E2M1Max / global_amax, cutlass::platform::numeric_limits::max());\n // If global amax is 0 or infinity, return 1\n return (global_amax == 0.f || global_encode_scale == 0.f) ? 1.f : global_encode_scale;\n}\n\ntemplate \nstruct SharedStorage {\n static constexpr int AccumulatorPipelineStageCount = 16;\n using AtomThrShapeMNK = cute::Shape<_1, _1, _1>;\n\n using AccumulatorPipeline = cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;\n\n static constexpr int MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline = cutlass::PipelineTmaUmmaAsync<\n MainloopPipelineStageCount,\n Shape<_1,_1,_1>,\n AtomThrShapeMNK>;\n using MainloopPipelineStorage = typename MainloopPipeline::SharedStorage;\n\n alignas(16) AccumulatorPipelineStorage accumulator;\n alignas(16) MainloopPipelineStorage mainloop;\n alignas(16) cute::uint64_t tma_barrier[1];\n uint32_t tmem_base_ptr;\n\n struct TensorStorage : cute::aligned_struct<128, _1> {\n // cute::array_aligned> smem_A;\n cute::array_aligned> smem_A;\n cute::array_aligned> smem_B;\n } tensors;\n\n};\n\nCUTLASS_DEVICE\ncutlass::Array\nStochasticNumericConverterBase(cutlass::Array const &input, cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING;\n if constexpr (has_rs) {\n auto output_ptr = reinterpret_cast(&output);\n asm volatile( \\\n \"{\\n\" \\\n \"cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\\n\" \\\n \"cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\\n\" \\\n \"}\" \\\n : \"=h\"(output_ptr[0]),\n \"=h\"(output_ptr[1])\n : \"f\"(input[0]), \"f\"(input[1]), \"f\"(input[2]), \"f\"(input[3]),\n \"f\"(input[4]), \"f\"(input[5]), \"f\"(input[6]), \"f\"(input[7]),\n \"r\"(rbits[0]), \"r\"(rbits[1]));\n } else {\n NVTE_DEVICE_ERROR(\"FP4 cvt PTX instructions are architecture-specific. \"\n \"Try recompiling with sm_XXXa instead of sm_XXX.\");\n }\n return output;\n}\n\nCUTLASS_DEVICE\ncutlass::Array\nStochasticNumericConverter(cutlass::Array const &input, cutlass::Array const *rbits) {\n using result_type = cutlass::Array;\n result_type output;\n cutlass::Array *result_ptr = reinterpret_cast *>(&output);\n cutlass::Array const *source_ptr = reinterpret_cast const *>(&input);\n cutlass::Array const *rbits_ptr = reinterpret_cast const *>(rbits);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 2; i++) {\n result_ptr[i] = StochasticNumericConverterBase(source_ptr[i], rbits_ptr[i]);\n }\n return output;\n}\n\ntemplate \n__global__ static\nvoid\nrht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile,\n TA const* A, AStride dA, ASmemLayout sAlayout, CUTE_GRID_CONSTANT TmaLoadA const tma_load_a,\n TB const* B, BStride dB, BSmemLayout sBlayout, CUTE_GRID_CONSTANT TmaLoadB const tma_load_b,\n TC * C, CStride dC, CSmemLayout ,\n TSFC * SFC,\n TiledMMA mma,\n float const* global_amax,\n const size_t* rng_state)\n{\n using namespace cute;\n constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY;\n if constexpr (!is_blackwell_arch) {\n NVTE_DEVICE_ERROR(\"RHT fusion is only supported on Blackwell.\");\n return;\n } else {\n using X = Underscore;\n // static constexpr bool kApplyStochasticRounding = true;\n using ElementAccumulator = float;\n static constexpr int K_PIPE_MAX = size<3>(ASmemLayout{});\n using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>;\n static constexpr uint32_t kTmaTransactionBytes =\n cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * cosize(take<0,3>(ASmemLayout{})) * cute::sizeof_bits_v);\n\n static constexpr int kTmaRhtTensorTransactionBytes =\n cutlass::bits_to_bytes(16 * 16 * cute::sizeof_bits_v);\n static constexpr int AccumulatorPipelineStageCount = 16;\n\n static constexpr int MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline = cutlass::PipelineTmaUmmaAsync<\n MainloopPipelineStageCount,\n Shape<_1,_1,_1>,\n AtomThrShapeMNK>;\n using MainloopPipelineState = typename MainloopPipeline::PipelineState;\n\n using TmemAllocator = cute::TMEM::Allocator1Sm;\n static constexpr int VectorSize = 16;\n const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0;\n const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0;\n // Preconditions\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n\n // Represent the full tensors\n Tensor mA = tma_load_a.get_tma_tensor(make_shape(M,N));\n Tensor mB = tma_load_b.get_tma_tensor(make_shape(16,16));\n Tensor mC = make_tensor(cute::subbyte_iterator(C), make_shape(M,N), dC); // (M,N)\n\n auto sfc_shape = make_shape(\n M,\n make_shape( make_shape(Int<16>{}, _4{}), N / 64 )\n );\n\n auto sfc_stride = make_stride(\n N / 16,\n make_stride( make_stride(_0{}, _1{}), _4{} )\n );\n\n auto sfc_layout = make_layout(sfc_shape, sfc_stride);\n Tensor mSFC = make_tensor(make_gmem_ptr(SFC), sfc_layout);\n\n auto cluster_shape = Shape< _1, _1, _1>{};\n\n // Get the appropriate blocks for this Cluster\n dim3 cluster_coord_in_grid = cluster_id_in_grid();\n\n // Total number of k-tiles\n const int K_TILE_MAX = min(N, K) / 64;\n uint32_t tiles_in_m = (M + size<0>(cluster_tile) - 1) / size<0>(cluster_tile);\n uint32_t tiles_in_n = (N + 64 - 1) / 64;\n uint32_t linear_tile_idx = blockIdx.x;\n uint32_t tile_idx_m = linear_tile_idx % tiles_in_m;\n uint32_t tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n\n\n auto mainloop_tiler = Shape<_128,_16,_64>{};\n auto epilogue_tiler = Shape<_128,_64,_64>{};\n Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_,_, _), Step<_1, X,_1>{});\n Tensor gB_nk = local_tile(mB, cluster_tile, make_coord(_,_, _), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)\n Tensor gC_mn = local_tile(mC, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N)\n\n Tensor gSFC_mn = local_tile(mSFC, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N)\n // Allocate SMEM\n extern __shared__ char shared_memory[];\n using SharedStorage = SharedStorage;\n SharedStorage& shared_storage = *reinterpret_cast(shared_memory);\n Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n\n //\n // MMA: Define C accumulators and A/B partitioning\n //\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k)\n\n auto mma_epilogue = make_tiled_mma(SM100_MMA_F16BF16_SS{},\n Layout>{});\n ThrMMA thr_mma_epilogue = mma_epilogue.get_slice(block_rank_in_cluster);\n\n\n using TiledMmaEpilogue = decltype(mma_epilogue);\n Tensor tCgA = thr_mma.partition_A(gA_mk);\n // Allocate \"fragments\" -- these are actually umma smem descriptors\n Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)\n Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE)\n\n auto acc_shape_mma = partition_shape_C(TiledMMA{}, take<0,2>(ClusterTileShape{}));\n auto acc_shape_epilogue = partition_shape_C(TiledMmaEpilogue{}, take<0,2>(epilogue_tiler));\n\n auto bulk_tmem_mma = TiledMMA::make_fragment_C(append(acc_shape_mma,\n Int{}));\n\n auto bulk_tmem_epilogue = TiledMmaEpilogue::make_fragment_C(append(acc_shape_epilogue,\n Int{}));\n\n TmemAllocator tmem_allocator{};\n cutlass::arch::NamedBarrier tmem_allocation_result_barrier(32 + 128, cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);\n\n Layout cta_layout_mnk = make_layout(cluster_shape);\n Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{}));\n auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster);\n\n auto [tAgA, tAsA] = tma_partition(tma_load_a,\n get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),\n group_modes<0,3>(tCsA), group_modes<0,3>(tCgA));\n\n auto [tBgB, tBsB] = tma_partition(tma_load_b,\n get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),\n group_modes<0,3>(tCsB), group_modes<0,3>(tCgB));\n\n uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);\n uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);\n\n int warp_idx = cutlass::canonical_warp_idx_sync();\n\n bool is_mma_warp = (warp_idx == 0);\n bool is_dma_warp = (warp_idx == 1);\n bool is_epilogue_warp = (warp_idx >= 4 && warp_idx <= 7);\n\n if (is_epilogue_warp && elect_one_sync()) {\n cute::prefetch(raw_pointer_cast(global_amax));\n }\n\n typename MainloopPipeline::Params mainloop_pipeline_params;\n if (is_dma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;\n }\n if (is_mma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;\n }\n mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp;\n mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes;\n mainloop_pipeline_params.initializing_warp = 0;\n MainloopPipeline mainloop_pipeline(shared_storage.mainloop,\n mainloop_pipeline_params,\n cluster_shape,\n cute::true_type{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n\n MainloopPipelineState mainloop_pipe_consumer_state;\n MainloopPipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state();\n\n\n\n using AccumulatorPipeline = cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState;\n\n AccumulatorPipelineState accumulator_pipe_consumer_state;\n AccumulatorPipelineState accumulator_pipe_producer_state = cutlass::make_producer_start_state();\n\n typename AccumulatorPipeline::Params accumulator_pipeline_params;\n if (is_mma_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer;\n }\n if (is_epilogue_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer;\n }\n // Only one producer thread arrives on this barrier.\n accumulator_pipeline_params.producer_arv_count = 1;\n accumulator_pipeline_params.consumer_arv_count = size(AtomThrShapeMNK{}) * 128;\n accumulator_pipeline_params.initializing_warp = 1;\n AccumulatorPipeline accumulator_pipeline(shared_storage.accumulator,\n accumulator_pipeline_params,\n cluster_shape,\n cute::true_type{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n\n if (warp_idx == 2 && elect_one_sync()) {\n cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);\n }\n __syncthreads();\n using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x;\n\n if (is_dma_warp) {\n if (elect_one_sync()) {\n cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0], kTmaRhtTensorTransactionBytes);\n copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_,0,0), tBsB(_,0));\n }\n cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/);\n do {\n bool is_first_wave = linear_tile_idx == blockIdx.x;\n uint32_t skip_wait = is_first_wave;\n auto tAgA_mk = tAgA(_,tile_idx_m,_);\n int k_tile = 0;\n auto barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state, skip_wait);\n\n\n CUTE_NO_UNROLL\n while (k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n) {\n int k_tile_idx_n = tile_idx_n + k_tile;\n ++k_tile;\n skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount);\n mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state, barrier_token);\n using BarrierType = typename MainloopPipeline::ProducerBarrierType;\n BarrierType* tma_barrier = mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state);\n int write_stage = mainloop_pipe_producer_state.index();\n ++mainloop_pipe_producer_state;\n barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state, skip_wait);\n if (cute::elect_one_sync()) {\n copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_,k_tile_idx_n), tAsA(_,write_stage));\n }\n }\n linear_tile_idx += gridDim.x;\n tile_idx_m = linear_tile_idx % tiles_in_m;\n tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n);\n mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);\n } else if (is_mma_warp) {\n mma.accumulate_ = UMMA::ScaleOut::Zero;\n\n tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);\n __syncwarp();\n tmem_allocation_result_barrier.arrive();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_mma.data() = tmem_base_ptr;\n\n do {\n uint32_t skip_wait = K_TILE_MAX <= 0;\n auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n CUTE_NO_UNROLL\n for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n; )\n {\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n int read_stage = mainloop_pipe_consumer_state.index();\n auto tCrA_mk = tCrA(_,_,_,read_stage);\n auto tCrB_nk = tCrB(_,_,0,0);\n CUTE_UNROLL\n for (int k_block = 0; k_block < size<2>(tCrA) / 4; ++k_block)\n {\n accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);\n CUTE_UNROLL\n for (int i = 0; i < 4; i++) {\n auto accumulators = bulk_tmem_mma(_,_,_,accumulator_pipe_producer_state.index() * 4 + i);\n gemm(mma, tCrA_mk(_,_,k_block * 4 + i), tCrB_nk, accumulators);\n }\n\n accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);\n ++accumulator_pipe_producer_state;\n }\n auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n skip_wait = k_tile >= K_TILE_MAX;\n barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state);\n }\n\n linear_tile_idx += gridDim.x;\n tile_idx_m = linear_tile_idx % tiles_in_m;\n tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n);\n tmem_allocator.release_allocation_lock();\n accumulator_pipeline.producer_tail(accumulator_pipe_producer_state);\n tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);\n } else if (is_epilogue_warp) {\n const float global_amax_val = *global_amax;\n static constexpr int FragmentSize = 256 / sizeof_bits_v;\n\n tmem_allocation_result_barrier.arrive_and_wait();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_epilogue.data() = tmem_base_ptr;\n int thread_idx = threadIdx.x % 128;\n\n Tensor tCgC = thr_mma_epilogue.partition_C(gC_mn); // (MMA,MMA_M,MMA_N) // (MMA,MMA_M,MMA_N)\n auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_,_,_,_0{}));\n auto tiled_r2g = make_tiled_copy_D(Copy_Atom{}, tiled_t2r);\n auto thr_t2r = tiled_t2r.get_slice(thread_idx);\n auto thr_r2g = tiled_r2g.get_slice(thread_idx);\n\n // NVFP4 non-E8 recipe constants and global scales\n static constexpr float fp4_max = 6.0f;\n\n const float global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val);\n const float global_decode_scale = 1.0f / global_encode_scale;\n\n // Scaling factor for fast math path\n static constexpr float fp4_max_inv = 1.0f / fp4_max;\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n\n do {\n for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n; ++k_tile) {\n Tensor tCgC_mn = tCgC(_,_,_,tile_idx_m,tile_idx_n+k_tile);\n\n Tensor tCgSFC_mn = gSFC_mn(_,_,tile_idx_m,tile_idx_n+k_tile);\n accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state);\n\n auto tCtC = bulk_tmem_epilogue(_,_,_,accumulator_pipe_consumer_state.index());\n Tensor tDtC = thr_t2r.partition_S(tCtC); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDgC = thr_t2r.partition_D(tCgC_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n\n Tensor tTR_rAcc = make_tensor(shape(tDgC)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDrC = make_tensor(shape(tDgC));\n Tensor tTR_rAcc_frag = recast>(coalesce(tTR_rAcc));\n Tensor tDrC_frag = recast>(coalesce(tDrC));\n\n Tensor src = thr_r2g.retile_S(tDrC);\n Tensor dst = thr_r2g.retile_D(tDgC);\n\n Tensor tCgSFC = make_tensor(tCgSFC_mn.data(), make_layout(\n make_shape(shape(tCgSFC_mn), Int<1>{}, Int<1>{}),\n make_stride(stride(tCgSFC_mn), Int<0>{}, Int<0>{})\n ));\n\n Tensor tDgSFC = filter(thr_t2r.partition_D(tCgSFC));\n Tensor tDrSFC = make_tensor(shape(tDgSFC));\n\n static constexpr int NumVecs = size(tDgC) / VectorSize;\n Tensor tC_rRowSFD_frg = recast>(tDrSFC);\n\n cutlass::maximum_absolute_value_reduction, true> amax_reduction;\n cutlass::Array vec_maxs;\n cutlass::Array pvscales;\n // TMEM_LOAD\n copy(tiled_t2r, tDtC, tTR_rAcc);\n cutlass::arch::fence_view_async_tmem_load();\n\n accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state);\n\n ++accumulator_pipe_consumer_state;\n\n if constexpr (!kUseFastMath) {\n // Downcast to BF16 for bit-wise compatibility with unfused\n // kernels\n auto convert_accum_to_bf16 = cutlass::NumericArrayConverter{};\n auto convert_bf16_to_accum = cutlass::NumericArrayConverter{};\n tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{})));\n }\n\n auto compute_frgs = reinterpret_cast *>(tTR_rAcc_frag.data());\n auto output_frgs = reinterpret_cast *>(tDrC_frag.data());\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]);\n }\n\n pvscales = cutlass::multiplies>{}(vec_maxs, global_encode_scale_multiplier);\n auto pvscales_cvted = cutlass::NumericArrayConverter{}(pvscales);\n\n tC_rRowSFD_frg(_0{}) = pvscales_cvted;\n auto qpvscale_ups = cutlass::NumericArrayConverter{}(tC_rRowSFD_frg(_0{}));\n auto qpvscale_scaled = cutlass::multiplies>{}(qpvscale_ups, global_decode_scale);\n cutlass::Array acc_scales;\n if constexpr (kUseFastMath) {\n // Fast math: compute approximate reciprocal\n acc_scales = cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled);\n } else {\n // Accurate math: compute reciprocal with division\n acc_scales = cutlass::divides>{}(1.0, qpvscale_scaled);\n }\n\n // Initialize RNG for tile\n const size_t rng_sequence\n = thread_idx + k_tile * 256 + linear_tile_idx * K_TILE_MAX * 256;\n\n transformer_engine::curanddx::detail::philox4x32_native_state\n rng;\n rng.init(rng_seed, rng_sequence, rng_offset);\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(acc_scales[v], cutlass::platform::numeric_limits::max());\n // auto acc_scale = acc_scales[v];\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(\n cutlass::multiplies>{}(\n compute_frgs[v],\n acc_scale\n ),\n reinterpret_cast*>(&random_uint4));\n } else {\n output_frgs[v] = cutlass::NumericArrayConverter{}(cutlass::multiplies>{}(compute_frgs[v], acc_scale));\n }\n }\n\n copy(tiled_r2g, src, dst);\n\n copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFC, tDgSFC);\n\n }\n linear_tile_idx += gridDim.x;\n tile_idx_m = linear_tile_idx % tiles_in_m;\n tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n);\n }\n }\n}\n\n// this function computes RHT-GEMM for\n// A: m x n: col-major\n// B: 16 x 16: row-major\n// C: m x n: row-major\n// SFC: m x (n/16): row-major\ntemplate \nvoid\nrht_gemm_ntt_w_sfc(int m, int n,\n TA const* A,\n TB const* B,\n TC * C,\n TSFC * SFC,\n float const* global_amax,\n const size_t* rng_state,\n uint32_t sm_count,\n cudaStream_t stream,\n int k_tile_size = 2048)\n{\n using namespace cute;\n\n // Define shapes (dynamic)\n auto M = static_cast(m);\n auto N = static_cast(n);\n\n // Define strides (mixed)\n auto dA = make_stride(Int<1>{}, m); // (dM,dK)\n auto dB = make_stride(Int<1>{}, 16); // (dN,dK)\n auto dC = make_stride(n, Int<1>{}); // (dM,dN)\n\n auto cga_shape = Shape< _1, _1, _1>{};\n auto cga_tile_shape = Shape<_128,_16,_16>{};\n auto cluster_tile_mainloop = Shape<_128,_16,_64>{};\n\n // Construct the MMA\n auto mma = make_tiled_mma(SM100_MMA_F16BF16_SS{},\n Layout>{});\n\n // MMA in CGA Layout XXX: Need to generalize synchro? {$nv-release-never}\n\n // Assert that the TiledMMA uses all CTAs in the CGA.\n CUTE_STATIC_ASSERT_V(size(cga_shape) == size(mma));\n CUTE_STATIC_ASSERT_V(evenly_divides(cga_tile_shape, tile_shape(mma)));\n\n // Determine the A and B shapes\n auto mma_shape_B = partition_shape_B(mma, make_shape(size<1>(cga_tile_shape), size<2>(cga_tile_shape)));\n\n using TiledMma = decltype(mma);\n using AtomThrID = typename TiledMma::AtomThrID;\n\n using SmemShape_M = decltype(shape_div(shape<0>(cga_tile_shape), shape_div(shape<0>(cga_tile_shape), size<0>(cga_tile_shape) / size(AtomThrID{}))));\n using SmemShape_N = decltype(shape_div(shape<1>(cga_tile_shape), shape_div(shape<1>(cga_tile_shape), size<1>(cga_tile_shape) / size(AtomThrID{}))));\n using SmemShape_K = decltype(cute::get<2>(cga_tile_shape));\n\n using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<\n cute::UMMA::Major::MN, TB, SmemShape_N, SmemShape_K>());\n\n auto mma_shape_A = partition_shape_A(mma, make_shape(size<0>(cluster_tile_mainloop), size<2>(cluster_tile_mainloop)));\n using SmemShape_M_A = decltype(shape_div(shape<0>(cluster_tile_mainloop), shape_div(shape<0>(cluster_tile_mainloop), size<0>(cluster_tile_mainloop) / size(AtomThrID{}))));\n using SmemShape_K_A = decltype(cute::get<2>(cluster_tile_mainloop));\n using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<\n cute::UMMA::Major::MN, TA, SmemShape_M_A, SmemShape_K_A>());\n\n // Define the smem layouts (static)\n // Calculate max pipeline stages based on Blackwell SM100's 232KB shared memory\n constexpr int kBlackwellSmemSize = 232448; // 232KB in bytes\n constexpr int kBytesPerStage = cute::size(mma_shape_A) * sizeof(TA) + cute::size(mma_shape_B) * sizeof(TB);\n constexpr int kReservedBytes = 256; // Reserve for barriers and other uses\n constexpr int kMaxStages = (kBlackwellSmemSize - kReservedBytes) / kBytesPerStage;\n auto sP = Int{}; // SMEM pipelines\n auto sA = UMMA::tile_to_mma_shape(SmemLayoutAtomA{}, append(mma_shape_A, sP)); // (MMA,MMA_M,MMA_K,PIPE)\n auto sB = UMMA::tile_to_mma_shape(SmemLayoutAtomB{}, append(mma_shape_B, sP)); // (MMA,MMA_N,MMA_K,PIPE)\n auto sC = Layout<_1>{}; // XXX Dummy\n\n // Create GMEM tensors\n Tensor tensorA = make_tensor(A, make_layout(make_shape(M,N), dA)); // (M,N)\n Tensor tensorB = make_tensor(B, make_layout(make_shape(16,16), dB)); // (16,16)\n\n // Create the TiledCopy\n\n auto tma_load_a = make_tma_copy_A_sm100(\n SM90_TMA_LOAD{},\n tensorA,\n sA(_,_,_,0),\n cluster_tile_mainloop,\n mma);\n auto tma_load_b = make_tma_copy_B_sm100(\n SM90_TMA_LOAD{},\n tensorB,\n sB(_,_,_,0),\n cga_tile_shape,\n mma);\n\n // Assert checks on tile sizes -- no predication\n NVTE_CHECK(M % size<0>(cga_tile_shape) == 0,\n \"Inner dimension must be divisible by \", static_cast(size<0>(cga_tile_shape)), \" but got \", M, \".\");\n NVTE_CHECK(N % (4 * size<1>(cga_tile_shape)) == 0,\n \"Outer dimension must be divisible by \", 4 * static_cast(size<1>(cga_tile_shape)),\n \" but got \", N, \".\");\n\n uint32_t tiles = size(ceil_div(M, get<0>(cga_tile_shape))) * size(ceil_div(N, k_tile_size));\n\n tiles = (tiles < sm_count) ? tiles : sm_count;\n\n dim3 dimBlock(256);\n dim3 dimCluster(size<0>(cga_shape), size<1>(cga_shape), size<2>(cga_shape));\n dim3 dimGrid(tiles, 1, 1);\n\n int smem_size = sizeof(SharedStorage);\n auto* kernel_ptr = &rht_gemm_device<\n decltype(M), decltype(N), decltype(k_tile_size), decltype(cga_tile_shape),\n TA, decltype(dA), decltype(sA), decltype(tma_load_a),\n TB, decltype(dB), decltype(sB), decltype(tma_load_b),\n TC, decltype(dC), decltype(sC),\n TSFC,\n decltype(mma),\n kEnableStochasticRounding,\n kUseFastMath>;\n\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(*kernel_ptr,\n cudaFuncAttributeMaxDynamicSharedMemorySize,\n smem_size)\n );\n\n (*kernel_ptr)\n <<< dimGrid, dimBlock, smem_size, stream >>>\n (M, N, k_tile_size, cga_tile_shape,\n A, dA, sA, tma_load_a,\n B, dB, sB, tma_load_b,\n C, dC, sC,\n SFC,\n mma, global_amax,\n rng_state);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n// this function is used to wrap the rht_gemm_ntt_w_sfc function\n//to transpose the input tensor A\ntemplate \nvoid\nrht_gemm_ttt_wrapper(int m, int n,\n TA const* A,\n TB const* B,\n TC * C,\n TSFC * SFC,\n float const* global_amax,\n const size_t* rng_state,\n uint32_t sm_count,\n cudaStream_t stream,\n int k_tile_size = 1024)\n{\n // in addition to transpose the input tensor A\n // we also need to reshape m, n to at best\n // ultilize as many SMs as possible while keeping\n // a relatively large contiguous dimension.\n // for example, after swapping m, n for transpose purposes,\n // the input / output tensor shapes for RHT-GEMM are:\n // A: n x m: col-major\n // B: 16 x 16: row-major\n // C: n x m: row-major\n // SFC: n x (m/16): row-major\n rht_gemm_ntt_w_sfc(\n n, m,\n A, B, C,\n SFC, global_amax,\n rng_state,\n sm_count, stream,\n k_tile_size);\n}\n\n} // namespace\n} // namespace detail\n\n// clang-format on\n\nvoid hadamard_transform_cast_fusion_columnwise(const Tensor &input_, Tensor &output_,\n const Tensor &hadamard_matrix_,\n QuantizationConfig quant_config,\n cudaStream_t stream) {\n NVTE_API_CALL(hadamard_transform_cast_fusion_columnwise);\n\n // Check input and output tensors\n NVTE_CHECK(input_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor must be BF16 tensor, but scaling mode is \",\n to_string(input_.scaling_mode), \".\");\n NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16,\n \"Input tensor must be BF16 tensor, but dtype is \", to_string(input_.dtype()), \".\");\n NVTE_CHECK(input_.dim() >= 2, \"Input must be a 2D tensor.\");\n NVTE_CHECK(!output_.with_gemm_swizzled_scales, \"Output must have scales in compact format.\");\n const SimpleTensor &input = input_.data;\n SimpleTensor &global_amax = output_.amax;\n SimpleTensor &output_t = output_.data;\n SimpleTensor &scale_inv_t = output_.scale_inv;\n\n // Stochastic rounding config\n const bool use_stochastic_rounding = quant_config.stochastic_rounding;\n const size_t *rng_state = nullptr;\n if (quant_config.rng_state != nullptr) {\n Tensor &rng_state_tensor = *convertNVTETensor(quant_config.rng_state);\n NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64,\n \"RNG state should contain 2 64-bit values.\");\n NVTE_CHECK(rng_state_tensor.data.shape == std::vector{2},\n \"Shape of the RNG state should be [2], but got \", rng_state_tensor.data.shape);\n rng_state = reinterpret_cast(rng_state_tensor.data.dptr);\n }\n\n // Template arguments\n using TA = cute::bfloat16_t;\n using TB = cute::bfloat16_t;\n using TC = cutlass::float_e2m1_t;\n using TSFC = cutlass::float_ue4m3_t;\n\n checkCuDriverContext(stream);\n\n // Check Hadamard matrix\n constexpr int kHadamardDimension = 16;\n NVTE_CHECK(hadamard_matrix_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Hadamard matrix must be BF16 tensor, but scaling mode is \",\n to_string(hadamard_matrix_.scaling_mode), \".\");\n NVTE_CHECK(hadamard_matrix_.dtype() == transformer_engine::DType::kBFloat16,\n \"Hadamard matrix must be BF16 tensor, but dtype is \",\n to_string(hadamard_matrix_.dtype()), \".\");\n const SimpleTensor &hadamard_matrix = hadamard_matrix_.data;\n NVTE_CHECK(\n (hadamard_matrix_.shape() == std::vector{kHadamardDimension, kHadamardDimension}),\n \"Hadamard matrix must have shape=\",\n std::vector{kHadamardDimension, kHadamardDimension},\n \", but got shape=\", hadamard_matrix_.shape(), \".\");\n const size_t hadamard_dimension = hadamard_matrix.shape[0];\n\n const size_t ndim = input.shape.size();\n const size_t n = input.shape[ndim - 1];\n size_t m = 1;\n for (size_t i = 0; i < ndim - 1; ++i) {\n m *= input.shape[i];\n }\n\n auto sm_count = transformer_engine::cuda::sm_count();\n\n NVTE_CHECK(n % hadamard_dimension == 0, \"row_length must be divisible by hadamard_dimension.\");\n\n NVTE_CHECK(m % hadamard_dimension == 0, \"num_rows must be divisible by hadamard_dimension\");\n\n int k_tile_size = 1024;\n\n if (m == 8192 && n == 5120) {\n k_tile_size = 512;\n } else if (m == 8192 && n == 10240) {\n k_tile_size = 1024;\n } else if (m == 8192 && n == 2560) {\n k_tile_size = 1280;\n } else if (m == 8192 && n == 11328) {\n k_tile_size = 1024;\n } else if (m == 8192 && n == 512) {\n k_tile_size = 256;\n } else if (m == 8192 && n == 3584) {\n k_tile_size = 512;\n } else if (m == 11328 && n == 8192) {\n k_tile_size = 1024;\n } else if (m == 5120 && n == 8192) {\n k_tile_size = 512;\n } else if (m == 10240 && n == 8192) {\n k_tile_size = 1024;\n } else if (m == 2560 && n == 8192) {\n k_tile_size = 1280;\n } else if (m == 512 && n == 8192) {\n k_tile_size = 256;\n } else if (m == 3584 && n == 8192) {\n k_tile_size = 512;\n } else if (m < 1024 || n < 1024) {\n k_tile_size = 512;\n }\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_stochastic_rounding, kUseStochasticRounding,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n quant_config.use_fast_math, kUseFastMath,\n detail::rht_gemm_ttt_wrapper(\n /*m=*/m,\n /*n=*/n,\n /*A=*/reinterpret_cast(input.dptr),\n /*B=*/reinterpret_cast(hadamard_matrix.dptr),\n /*C=*/reinterpret_cast(output_t.dptr),\n /*SFC=*/reinterpret_cast(scale_inv_t.dptr),\n /*global_amax=*/reinterpret_cast(global_amax.dptr),\n /*rng_state=*/rng_state,\n /*sm_count=*/sm_count,\n /*stream=*/stream,\n /*k_tile_size=*/k_tile_size);););\n}\n\n} // namespace transformer_engine\n\nvoid nvte_hadamard_transform_cast_fusion_columnwise(const NVTETensor input, NVTETensor output,\n const NVTETensor hadamard_matrix,\n const NVTEQuantizationConfig quant_config,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_hadamard_transform_cast_fusion_columnwise);\n using namespace transformer_engine;\n QuantizationConfig quant_config_cpp;\n if (quant_config != nullptr) {\n quant_config_cpp = *reinterpret_cast(quant_config);\n }\n hadamard_transform_cast_fusion_columnwise(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(output),\n *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/util/hip/hip_runtime.h\"\n#include \"common/util/curanddx.hpp\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"cutlass/arch/barrier.h\"\n#include \"cutlass/cutlass.h\"\n#include \"cutlass/gemm/collective/builders/sm100_common.inl\"\n#include \"cutlass/numeric_conversion.h\"\n#include \"cutlass/pipeline/pipeline.hpp\"\n#include \"cutlass/util/GPU_Clock.hpp\"\n#include \"cutlass/util/command_line.h\"\n#include \"cutlass/util/print_error.hpp\"\n\n// clang-format off\n\nnamespace transformer_engine {\nnamespace detail {\nnamespace {\n\nusing namespace cute;\nusing cute::Tensor; // Ensure unqualified Tensor refers to cute::Tensor, not transformer_engine::Tensor\n\n// calculate the global encode scale factor for a given global amax.\n__device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax) {\n constexpr float kFP8E4M3Max = 448.0f;\n constexpr float kFP4E2M1Max = 6.0f;\n // If scale is infinity, return max value of float32\n float global_encode_scale = cutlass::minimum_with_nan_propagation{}(\n kFP8E4M3Max * kFP4E2M1Max / global_amax, cutlass::platform::numeric_limits::max());\n // If global amax is 0 or infinity, return 1\n return (global_amax == 0.f || global_encode_scale == 0.f) ? 1.f : global_encode_scale;\n}\n\ntemplate \nstruct SharedStorage {\n static constexpr int AccumulatorPipelineStageCount = 16;\n using AtomThrShapeMNK = cute::Shape<_1, _1, _1>;\n\n using AccumulatorPipeline = cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;\n\n static constexpr int MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline = cutlass::PipelineTmaUmmaAsync<\n MainloopPipelineStageCount,\n Shape<_1,_1,_1>,\n AtomThrShapeMNK>;\n using MainloopPipelineStorage = typename MainloopPipeline::SharedStorage;\n\n alignas(16) AccumulatorPipelineStorage accumulator;\n alignas(16) MainloopPipelineStorage mainloop;\n alignas(16) cute::uint64_t tma_barrier[1];\n uint32_t tmem_base_ptr;\n\n struct TensorStorage : cute::aligned_struct<128, _1> {\n // cute::array_aligned> smem_A;\n cute::array_aligned> smem_A;\n cute::array_aligned> smem_B;\n } tensors;\n\n};\n\nCUTLASS_DEVICE\ncutlass::Array\nStochasticNumericConverterBase(cutlass::Array const &input, cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING;\n if constexpr (has_rs) {\n auto output_ptr = reinterpret_cast(&output);\n asm volatile( \\\n \"{\\n\" \\\n \"cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\\n\" \\\n \"cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\\n\" \\\n \"}\" \\\n : \"=h\"(output_ptr[0]),\n \"=h\"(output_ptr[1])\n : \"f\"(input[0]), \"f\"(input[1]), \"f\"(input[2]), \"f\"(input[3]),\n \"f\"(input[4]), \"f\"(input[5]), \"f\"(input[6]), \"f\"(input[7]),\n \"r\"(rbits[0]), \"r\"(rbits[1]));\n } else {\n NVTE_DEVICE_ERROR(\"FP4 cvt PTX instructions are architecture-specific. \"\n \"Try recompiling with sm_XXXa instead of sm_XXX.\");\n }\n return output;\n}\n\nCUTLASS_DEVICE\ncutlass::Array\nStochasticNumericConverter(cutlass::Array const &input, cutlass::Array const *rbits) {\n using result_type = cutlass::Array;\n result_type output;\n cutlass::Array *result_ptr = reinterpret_cast *>(&output);\n cutlass::Array const *source_ptr = reinterpret_cast const *>(&input);\n cutlass::Array const *rbits_ptr = reinterpret_cast const *>(rbits);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 2; i++) {\n result_ptr[i] = StochasticNumericConverterBase(source_ptr[i], rbits_ptr[i]);\n }\n return output;\n}\n\ntemplate \n__global__ static\nvoid\nrht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile,\n TA const* A, AStride dA, ASmemLayout sAlayout, CUTE_GRID_CONSTANT TmaLoadA const tma_load_a,\n TB const* B, BStride dB, BSmemLayout sBlayout, CUTE_GRID_CONSTANT TmaLoadB const tma_load_b,\n TC * C, CStride dC, CSmemLayout ,\n TSFC * SFC,\n TiledMMA mma,\n float const* global_amax,\n const size_t* rng_state)\n{\n using namespace cute;\n constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY;\n if constexpr (!is_blackwell_arch) {\n NVTE_DEVICE_ERROR(\"RHT fusion is only supported on Blackwell.\");\n return;\n } else {\n using X = Underscore;\n // static constexpr bool kApplyStochasticRounding = true;\n using ElementAccumulator = float;\n static constexpr int K_PIPE_MAX = size<3>(ASmemLayout{});\n using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>;\n static constexpr uint32_t kTmaTransactionBytes =\n cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * cosize(take<0,3>(ASmemLayout{})) * cute::sizeof_bits_v);\n\n static constexpr int kTmaRhtTensorTransactionBytes =\n cutlass::bits_to_bytes(16 * 16 * cute::sizeof_bits_v);\n static constexpr int AccumulatorPipelineStageCount = 16;\n\n static constexpr int MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline = cutlass::PipelineTmaUmmaAsync<\n MainloopPipelineStageCount,\n Shape<_1,_1,_1>,\n AtomThrShapeMNK>;\n using MainloopPipelineState = typename MainloopPipeline::PipelineState;\n\n using TmemAllocator = cute::TMEM::Allocator1Sm;\n static constexpr int VectorSize = 16;\n const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0;\n const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0;\n // Preconditions\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n\n // Represent the full tensors\n Tensor mA = tma_load_a.get_tma_tensor(make_shape(M,N));\n Tensor mB = tma_load_b.get_tma_tensor(make_shape(16,16));\n Tensor mC = make_tensor(cute::subbyte_iterator(C), make_shape(M,N), dC); // (M,N)\n\n auto sfc_shape = make_shape(\n M,\n make_shape( make_shape(Int<16>{}, _4{}), N / 64 )\n );\n\n auto sfc_stride = make_stride(\n N / 16,\n make_stride( make_stride(_0{}, _1{}), _4{} )\n );\n\n auto sfc_layout = make_layout(sfc_shape, sfc_stride);\n Tensor mSFC = make_tensor(make_gmem_ptr(SFC), sfc_layout);\n\n auto cluster_shape = Shape< _1, _1, _1>{};\n\n // Get the appropriate blocks for this Cluster\n dim3 cluster_coord_in_grid = cluster_id_in_grid();\n\n // Total number of k-tiles\n const int K_TILE_MAX = min(N, K) / 64;\n uint32_t tiles_in_m = (M + size<0>(cluster_tile) - 1) / size<0>(cluster_tile);\n uint32_t tiles_in_n = (N + 64 - 1) / 64;\n uint32_t linear_tile_idx = blockIdx.x;\n uint32_t tile_idx_m = linear_tile_idx % tiles_in_m;\n uint32_t tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n\n\n auto mainloop_tiler = Shape<_128,_16,_64>{};\n auto epilogue_tiler = Shape<_128,_64,_64>{};\n Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_,_, _), Step<_1, X,_1>{});\n Tensor gB_nk = local_tile(mB, cluster_tile, make_coord(_,_, _), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)\n Tensor gC_mn = local_tile(mC, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N)\n\n Tensor gSFC_mn = local_tile(mSFC, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N)\n // Allocate SMEM\n extern __shared__ char shared_memory[];\n using SharedStorage = SharedStorage;\n SharedStorage& shared_storage = *reinterpret_cast(shared_memory);\n Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n\n //\n // MMA: Define C accumulators and A/B partitioning\n //\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k)\n\n auto mma_epilogue = make_tiled_mma(SM100_MMA_F16BF16_SS{},\n Layout>{});\n ThrMMA thr_mma_epilogue = mma_epilogue.get_slice(block_rank_in_cluster);\n\n\n using TiledMmaEpilogue = decltype(mma_epilogue);\n Tensor tCgA = thr_mma.partition_A(gA_mk);\n // Allocate \"fragments\" -- these are actually umma smem descriptors\n Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)\n Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE)\n\n auto acc_shape_mma = partition_shape_C(TiledMMA{}, take<0,2>(ClusterTileShape{}));\n auto acc_shape_epilogue = partition_shape_C(TiledMmaEpilogue{}, take<0,2>(epilogue_tiler));\n\n auto bulk_tmem_mma = TiledMMA::make_fragment_C(append(acc_shape_mma,\n Int{}));\n\n auto bulk_tmem_epilogue = TiledMmaEpilogue::make_fragment_C(append(acc_shape_epilogue,\n Int{}));\n\n TmemAllocator tmem_allocator{};\n cutlass::arch::NamedBarrier tmem_allocation_result_barrier(32 + 128, cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);\n\n Layout cta_layout_mnk = make_layout(cluster_shape);\n Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{}));\n auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster);\n\n auto [tAgA, tAsA] = tma_partition(tma_load_a,\n get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),\n group_modes<0,3>(tCsA), group_modes<0,3>(tCgA));\n\n auto [tBgB, tBsB] = tma_partition(tma_load_b,\n get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),\n group_modes<0,3>(tCsB), group_modes<0,3>(tCgB));\n\n uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);\n uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);\n\n int warp_idx = cutlass::canonical_warp_idx_sync();\n\n bool is_mma_warp = (warp_idx == 0);\n bool is_dma_warp = (warp_idx == 1);\n bool is_epilogue_warp = (warp_idx >= 4 && warp_idx <= 7);\n\n if (is_epilogue_warp && elect_one_sync()) {\n cute::prefetch(raw_pointer_cast(global_amax));\n }\n\n typename MainloopPipeline::Params mainloop_pipeline_params;\n if (is_dma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;\n }\n if (is_mma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;\n }\n mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp;\n mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes;\n mainloop_pipeline_params.initializing_warp = 0;\n MainloopPipeline mainloop_pipeline(shared_storage.mainloop,\n mainloop_pipeline_params,\n cluster_shape,\n cute::true_type{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n\n MainloopPipelineState mainloop_pipe_consumer_state;\n MainloopPipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state();\n\n\n\n using AccumulatorPipeline = cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState;\n\n AccumulatorPipelineState accumulator_pipe_consumer_state;\n AccumulatorPipelineState accumulator_pipe_producer_state = cutlass::make_producer_start_state();\n\n typename AccumulatorPipeline::Params accumulator_pipeline_params;\n if (is_mma_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer;\n }\n if (is_epilogue_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer;\n }\n // Only one producer thread arrives on this barrier.\n accumulator_pipeline_params.producer_arv_count = 1;\n accumulator_pipeline_params.consumer_arv_count = size(AtomThrShapeMNK{}) * 128;\n accumulator_pipeline_params.initializing_warp = 1;\n AccumulatorPipeline accumulator_pipeline(shared_storage.accumulator,\n accumulator_pipeline_params,\n cluster_shape,\n cute::true_type{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n\n if (warp_idx == 2 && elect_one_sync()) {\n cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);\n }\n __syncthreads();\n using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x;\n\n if (is_dma_warp) {\n if (elect_one_sync()) {\n cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0], kTmaRhtTensorTransactionBytes);\n copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_,0,0), tBsB(_,0));\n }\n cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/);\n do {\n bool is_first_wave = linear_tile_idx == blockIdx.x;\n uint32_t skip_wait = is_first_wave;\n auto tAgA_mk = tAgA(_,tile_idx_m,_);\n int k_tile = 0;\n auto barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state, skip_wait);\n\n\n CUTE_NO_UNROLL\n while (k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n) {\n int k_tile_idx_n = tile_idx_n + k_tile;\n ++k_tile;\n skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount);\n mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state, barrier_token);\n using BarrierType = typename MainloopPipeline::ProducerBarrierType;\n BarrierType* tma_barrier = mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state);\n int write_stage = mainloop_pipe_producer_state.index();\n ++mainloop_pipe_producer_state;\n barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state, skip_wait);\n if (cute::elect_one_sync()) {\n copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_,k_tile_idx_n), tAsA(_,write_stage));\n }\n }\n linear_tile_idx += gridDim.x;\n tile_idx_m = linear_tile_idx % tiles_in_m;\n tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n);\n mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);\n } else if (is_mma_warp) {\n mma.accumulate_ = UMMA::ScaleOut::Zero;\n\n tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);\n __syncwarp();\n tmem_allocation_result_barrier.arrive();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_mma.data() = tmem_base_ptr;\n\n do {\n uint32_t skip_wait = K_TILE_MAX <= 0;\n auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n CUTE_NO_UNROLL\n for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n; )\n {\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n int read_stage = mainloop_pipe_consumer_state.index();\n auto tCrA_mk = tCrA(_,_,_,read_stage);\n auto tCrB_nk = tCrB(_,_,0,0);\n CUTE_UNROLL\n for (int k_block = 0; k_block < size<2>(tCrA) / 4; ++k_block)\n {\n accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);\n CUTE_UNROLL\n for (int i = 0; i < 4; i++) {\n auto accumulators = bulk_tmem_mma(_,_,_,accumulator_pipe_producer_state.index() * 4 + i);\n gemm(mma, tCrA_mk(_,_,k_block * 4 + i), tCrB_nk, accumulators);\n }\n\n accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);\n ++accumulator_pipe_producer_state;\n }\n auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n skip_wait = k_tile >= K_TILE_MAX;\n barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state);\n }\n\n linear_tile_idx += gridDim.x;\n tile_idx_m = linear_tile_idx % tiles_in_m;\n tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n);\n tmem_allocator.release_allocation_lock();\n accumulator_pipeline.producer_tail(accumulator_pipe_producer_state);\n tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);\n } else if (is_epilogue_warp) {\n const float global_amax_val = *global_amax;\n static constexpr int FragmentSize = 256 / sizeof_bits_v;\n\n tmem_allocation_result_barrier.arrive_and_wait();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_epilogue.data() = tmem_base_ptr;\n int thread_idx = threadIdx.x % 128;\n\n Tensor tCgC = thr_mma_epilogue.partition_C(gC_mn); // (MMA,MMA_M,MMA_N) // (MMA,MMA_M,MMA_N)\n auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_,_,_,_0{}));\n auto tiled_r2g = make_tiled_copy_D(Copy_Atom{}, tiled_t2r);\n auto thr_t2r = tiled_t2r.get_slice(thread_idx);\n auto thr_r2g = tiled_r2g.get_slice(thread_idx);\n\n // NVFP4 non-E8 recipe constants and global scales\n static constexpr float fp4_max = 6.0f;\n\n const float global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val);\n const float global_decode_scale = 1.0f / global_encode_scale;\n\n // Scaling factor for fast math path\n static constexpr float fp4_max_inv = 1.0f / fp4_max;\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n\n do {\n for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n; ++k_tile) {\n Tensor tCgC_mn = tCgC(_,_,_,tile_idx_m,tile_idx_n+k_tile);\n\n Tensor tCgSFC_mn = gSFC_mn(_,_,tile_idx_m,tile_idx_n+k_tile);\n accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state);\n\n auto tCtC = bulk_tmem_epilogue(_,_,_,accumulator_pipe_consumer_state.index());\n Tensor tDtC = thr_t2r.partition_S(tCtC); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDgC = thr_t2r.partition_D(tCgC_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n\n Tensor tTR_rAcc = make_tensor(shape(tDgC)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDrC = make_tensor(shape(tDgC));\n Tensor tTR_rAcc_frag = recast>(coalesce(tTR_rAcc));\n Tensor tDrC_frag = recast>(coalesce(tDrC));\n\n Tensor src = thr_r2g.retile_S(tDrC);\n Tensor dst = thr_r2g.retile_D(tDgC);\n\n Tensor tCgSFC = make_tensor(tCgSFC_mn.data(), make_layout(\n make_shape(shape(tCgSFC_mn), Int<1>{}, Int<1>{}),\n make_stride(stride(tCgSFC_mn), Int<0>{}, Int<0>{})\n ));\n\n Tensor tDgSFC = filter(thr_t2r.partition_D(tCgSFC));\n Tensor tDrSFC = make_tensor(shape(tDgSFC));\n\n static constexpr int NumVecs = size(tDgC) / VectorSize;\n Tensor tC_rRowSFD_frg = recast>(tDrSFC);\n\n cutlass::maximum_absolute_value_reduction, true> amax_reduction;\n cutlass::Array vec_maxs;\n cutlass::Array pvscales;\n // TMEM_LOAD\n copy(tiled_t2r, tDtC, tTR_rAcc);\n cutlass::arch::fence_view_async_tmem_load();\n\n accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state);\n\n ++accumulator_pipe_consumer_state;\n\n if constexpr (!kUseFastMath) {\n // Downcast to BF16 for bit-wise compatibility with unfused\n // kernels\n auto convert_accum_to_bf16 = cutlass::NumericArrayConverter{};\n auto convert_bf16_to_accum = cutlass::NumericArrayConverter{};\n tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{})));\n }\n\n auto compute_frgs = reinterpret_cast *>(tTR_rAcc_frag.data());\n auto output_frgs = reinterpret_cast *>(tDrC_frag.data());\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]);\n }\n\n pvscales = cutlass::multiplies>{}(vec_maxs, global_encode_scale_multiplier);\n auto pvscales_cvted = cutlass::NumericArrayConverter{}(pvscales);\n\n tC_rRowSFD_frg(_0{}) = pvscales_cvted;\n auto qpvscale_ups = cutlass::NumericArrayConverter{}(tC_rRowSFD_frg(_0{}));\n auto qpvscale_scaled = cutlass::multiplies>{}(qpvscale_ups, global_decode_scale);\n cutlass::Array acc_scales;\n if constexpr (kUseFastMath) {\n // Fast math: compute approximate reciprocal\n acc_scales = cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled);\n } else {\n // Accurate math: compute reciprocal with division\n acc_scales = cutlass::divides>{}(1.0, qpvscale_scaled);\n }\n\n // Initialize RNG for tile\n const size_t rng_sequence\n = thread_idx + k_tile * 256 + linear_tile_idx * K_TILE_MAX * 256;\n\n transformer_engine::curanddx::detail::philox4x32_native_state\n rng;\n rng.init(rng_seed, rng_sequence, rng_offset);\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(acc_scales[v], cutlass::platform::numeric_limits::max());\n // auto acc_scale = acc_scales[v];\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(\n cutlass::multiplies>{}(\n compute_frgs[v],\n acc_scale\n ),\n reinterpret_cast*>(&random_uint4));\n } else {\n output_frgs[v] = cutlass::NumericArrayConverter{}(cutlass::multiplies>{}(compute_frgs[v], acc_scale));\n }\n }\n\n copy(tiled_r2g, src, dst);\n\n copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFC, tDgSFC);\n\n }\n linear_tile_idx += gridDim.x;\n tile_idx_m = linear_tile_idx % tiles_in_m;\n tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n);\n }\n }\n}\n\n// this function computes RHT-GEMM for\n// A: m x n: col-major\n// B: 16 x 16: row-major\n// C: m x n: row-major\n// SFC: m x (n/16): row-major\ntemplate \nvoid\nrht_gemm_ntt_w_sfc(int m, int n,\n TA const* A,\n TB const* B,\n TC * C,\n TSFC * SFC,\n float const* global_amax,\n const size_t* rng_state,\n uint32_t sm_count,\n hipStream_t stream,\n int k_tile_size = 2048)\n{\n using namespace cute;\n\n // Define shapes (dynamic)\n auto M = static_cast(m);\n auto N = static_cast(n);\n\n // Define strides (mixed)\n auto dA = make_stride(Int<1>{}, m); // (dM,dK)\n auto dB = make_stride(Int<1>{}, 16); // (dN,dK)\n auto dC = make_stride(n, Int<1>{}); // (dM,dN)\n\n auto cga_shape = Shape< _1, _1, _1>{};\n auto cga_tile_shape = Shape<_128,_16,_16>{};\n auto cluster_tile_mainloop = Shape<_128,_16,_64>{};\n\n // Construct the MMA\n auto mma = make_tiled_mma(SM100_MMA_F16BF16_SS{},\n Layout>{});\n\n // MMA in CGA Layout XXX: Need to generalize synchro? {$nv-release-never}\n\n // Assert that the TiledMMA uses all CTAs in the CGA.\n CUTE_STATIC_ASSERT_V(size(cga_shape) == size(mma));\n CUTE_STATIC_ASSERT_V(evenly_divides(cga_tile_shape, tile_shape(mma)));\n\n // Determine the A and B shapes\n auto mma_shape_B = partition_shape_B(mma, make_shape(size<1>(cga_tile_shape), size<2>(cga_tile_shape)));\n\n using TiledMma = decltype(mma);\n using AtomThrID = typename TiledMma::AtomThrID;\n\n using SmemShape_M = decltype(shape_div(shape<0>(cga_tile_shape), shape_div(shape<0>(cga_tile_shape), size<0>(cga_tile_shape) / size(AtomThrID{}))));\n using SmemShape_N = decltype(shape_div(shape<1>(cga_tile_shape), shape_div(shape<1>(cga_tile_shape), size<1>(cga_tile_shape) / size(AtomThrID{}))));\n using SmemShape_K = decltype(cute::get<2>(cga_tile_shape));\n\n using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<\n cute::UMMA::Major::MN, TB, SmemShape_N, SmemShape_K>());\n\n auto mma_shape_A = partition_shape_A(mma, make_shape(size<0>(cluster_tile_mainloop), size<2>(cluster_tile_mainloop)));\n using SmemShape_M_A = decltype(shape_div(shape<0>(cluster_tile_mainloop), shape_div(shape<0>(cluster_tile_mainloop), size<0>(cluster_tile_mainloop) / size(AtomThrID{}))));\n using SmemShape_K_A = decltype(cute::get<2>(cluster_tile_mainloop));\n using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<\n cute::UMMA::Major::MN, TA, SmemShape_M_A, SmemShape_K_A>());\n\n // Define the smem layouts (static)\n // Calculate max pipeline stages based on Blackwell SM100's 232KB shared memory\n constexpr int kBlackwellSmemSize = 232448; // 232KB in bytes\n constexpr int kBytesPerStage = cute::size(mma_shape_A) * sizeof(TA) + cute::size(mma_shape_B) * sizeof(TB);\n constexpr int kReservedBytes = 256; // Reserve for barriers and other uses\n constexpr int kMaxStages = (kBlackwellSmemSize - kReservedBytes) / kBytesPerStage;\n auto sP = Int{}; // SMEM pipelines\n auto sA = UMMA::tile_to_mma_shape(SmemLayoutAtomA{}, append(mma_shape_A, sP)); // (MMA,MMA_M,MMA_K,PIPE)\n auto sB = UMMA::tile_to_mma_shape(SmemLayoutAtomB{}, append(mma_shape_B, sP)); // (MMA,MMA_N,MMA_K,PIPE)\n auto sC = Layout<_1>{}; // XXX Dummy\n\n // Create GMEM tensors\n Tensor tensorA = make_tensor(A, make_layout(make_shape(M,N), dA)); // (M,N)\n Tensor tensorB = make_tensor(B, make_layout(make_shape(16,16), dB)); // (16,16)\n\n // Create the TiledCopy\n\n auto tma_load_a = make_tma_copy_A_sm100(\n SM90_TMA_LOAD{},\n tensorA,\n sA(_,_,_,0),\n cluster_tile_mainloop,\n mma);\n auto tma_load_b = make_tma_copy_B_sm100(\n SM90_TMA_LOAD{},\n tensorB,\n sB(_,_,_,0),\n cga_tile_shape,\n mma);\n\n // Assert checks on tile sizes -- no predication\n NVTE_CHECK(M % size<0>(cga_tile_shape) == 0,\n \"Inner dimension must be divisible by \", static_cast(size<0>(cga_tile_shape)), \" but got \", M, \".\");\n NVTE_CHECK(N % (4 * size<1>(cga_tile_shape)) == 0,\n \"Outer dimension must be divisible by \", 4 * static_cast(size<1>(cga_tile_shape)),\n \" but got \", N, \".\");\n\n uint32_t tiles = size(ceil_div(M, get<0>(cga_tile_shape))) * size(ceil_div(N, k_tile_size));\n\n tiles = (tiles < sm_count) ? tiles : sm_count;\n\n dim3 dimBlock(256);\n dim3 dimCluster(size<0>(cga_shape), size<1>(cga_shape), size<2>(cga_shape));\n dim3 dimGrid(tiles, 1, 1);\n\n int smem_size = sizeof(SharedStorage);\n auto* kernel_ptr = &rht_gemm_device<\n decltype(M), decltype(N), decltype(k_tile_size), decltype(cga_tile_shape),\n TA, decltype(dA), decltype(sA), decltype(tma_load_a),\n TB, decltype(dB), decltype(sB), decltype(tma_load_b),\n TC, decltype(dC), decltype(sC),\n TSFC,\n decltype(mma),\n kEnableStochasticRounding,\n kUseFastMath>;\n\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(*kernel_ptr),\n hipFuncAttributeMaxDynamicSharedMemorySize,\n smem_size)\n );\n\n (*kernel_ptr)\n <<< dimGrid, dimBlock, smem_size, stream >>>\n (M, N, k_tile_size, cga_tile_shape,\n A, dA, sA, tma_load_a,\n B, dB, sB, tma_load_b,\n C, dC, sC,\n SFC,\n mma, global_amax,\n rng_state);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n// this function is used to wrap the rht_gemm_ntt_w_sfc function\n//to transpose the input tensor A\ntemplate \nvoid\nrht_gemm_ttt_wrapper(int m, int n,\n TA const* A,\n TB const* B,\n TC * C,\n TSFC * SFC,\n float const* global_amax,\n const size_t* rng_state,\n uint32_t sm_count,\n hipStream_t stream,\n int k_tile_size = 1024)\n{\n // in addition to transpose the input tensor A\n // we also need to reshape m, n to at best\n // ultilize as many SMs as possible while keeping\n // a relatively large contiguous dimension.\n // for example, after swapping m, n for transpose purposes,\n // the input / output tensor shapes for RHT-GEMM are:\n // A: n x m: col-major\n // B: 16 x 16: row-major\n // C: n x m: row-major\n // SFC: n x (m/16): row-major\n rht_gemm_ntt_w_sfc(\n n, m,\n A, B, C,\n SFC, global_amax,\n rng_state,\n sm_count, stream,\n k_tile_size);\n}\n\n} // namespace\n} // namespace detail\n\n// clang-format on\n\nvoid hadamard_transform_cast_fusion_columnwise(const Tensor &input_, Tensor &output_,\n const Tensor &hadamard_matrix_,\n QuantizationConfig quant_config,\n hipStream_t stream) {\n NVTE_API_CALL(hadamard_transform_cast_fusion_columnwise);\n\n // Check input and output tensors\n NVTE_CHECK(input_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor must be BF16 tensor, but scaling mode is \",\n to_string(input_.scaling_mode), \".\");\n NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16,\n \"Input tensor must be BF16 tensor, but dtype is \", to_string(input_.dtype()), \".\");\n NVTE_CHECK(input_.dim() >= 2, \"Input must be a 2D tensor.\");\n NVTE_CHECK(!output_.with_gemm_swizzled_scales, \"Output must have scales in compact format.\");\n const SimpleTensor &input = input_.data;\n SimpleTensor &global_amax = output_.amax;\n SimpleTensor &output_t = output_.data;\n SimpleTensor &scale_inv_t = output_.scale_inv;\n\n // Stochastic rounding config\n const bool use_stochastic_rounding = quant_config.stochastic_rounding;\n const size_t *rng_state = nullptr;\n if (quant_config.rng_state != nullptr) {\n Tensor &rng_state_tensor = *convertNVTETensor(quant_config.rng_state);\n NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64,\n \"RNG state should contain 2 64-bit values.\");\n NVTE_CHECK(rng_state_tensor.data.shape == std::vector{2},\n \"Shape of the RNG state should be [2], but got \", rng_state_tensor.data.shape);\n rng_state = reinterpret_cast(rng_state_tensor.data.dptr);\n }\n\n // Template arguments\n using TA = cute::bfloat16_t;\n using TB = cute::bfloat16_t;\n using TC = cutlass::float_e2m1_t;\n using TSFC = cutlass::float_ue4m3_t;\n\n checkCuDriverContext(stream);\n\n // Check Hadamard matrix\n constexpr int kHadamardDimension = 16;\n NVTE_CHECK(hadamard_matrix_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Hadamard matrix must be BF16 tensor, but scaling mode is \",\n to_string(hadamard_matrix_.scaling_mode), \".\");\n NVTE_CHECK(hadamard_matrix_.dtype() == transformer_engine::DType::kBFloat16,\n \"Hadamard matrix must be BF16 tensor, but dtype is \",\n to_string(hadamard_matrix_.dtype()), \".\");\n const SimpleTensor &hadamard_matrix = hadamard_matrix_.data;\n NVTE_CHECK(\n (hadamard_matrix_.shape() == std::vector{kHadamardDimension, kHadamardDimension}),\n \"Hadamard matrix must have shape=\",\n std::vector{kHadamardDimension, kHadamardDimension},\n \", but got shape=\", hadamard_matrix_.shape(), \".\");\n const size_t hadamard_dimension = hadamard_matrix.shape[0];\n\n const size_t ndim = input.shape.size();\n const size_t n = input.shape[ndim - 1];\n size_t m = 1;\n for (size_t i = 0; i < ndim - 1; ++i) {\n m *= input.shape[i];\n }\n\n auto sm_count = transformer_engine::cuda::sm_count();\n\n NVTE_CHECK(n % hadamard_dimension == 0, \"row_length must be divisible by hadamard_dimension.\");\n\n NVTE_CHECK(m % hadamard_dimension == 0, \"num_rows must be divisible by hadamard_dimension\");\n\n int k_tile_size = 1024;\n\n if (m == 8192 && n == 5120) {\n k_tile_size = 512;\n } else if (m == 8192 && n == 10240) {\n k_tile_size = 1024;\n } else if (m == 8192 && n == 2560) {\n k_tile_size = 1280;\n } else if (m == 8192 && n == 11328) {\n k_tile_size = 1024;\n } else if (m == 8192 && n == 512) {\n k_tile_size = 256;\n } else if (m == 8192 && n == 3584) {\n k_tile_size = 512;\n } else if (m == 11328 && n == 8192) {\n k_tile_size = 1024;\n } else if (m == 5120 && n == 8192) {\n k_tile_size = 512;\n } else if (m == 10240 && n == 8192) {\n k_tile_size = 1024;\n } else if (m == 2560 && n == 8192) {\n k_tile_size = 1280;\n } else if (m == 512 && n == 8192) {\n k_tile_size = 256;\n } else if (m == 3584 && n == 8192) {\n k_tile_size = 512;\n } else if (m < 1024 || n < 1024) {\n k_tile_size = 512;\n }\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_stochastic_rounding, kUseStochasticRounding,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n quant_config.use_fast_math, kUseFastMath,\n detail::rht_gemm_ttt_wrapper(\n /*m=*/m,\n /*n=*/n,\n /*A=*/reinterpret_cast(input.dptr),\n /*B=*/reinterpret_cast(hadamard_matrix.dptr),\n /*C=*/reinterpret_cast(output_t.dptr),\n /*SFC=*/reinterpret_cast(scale_inv_t.dptr),\n /*global_amax=*/reinterpret_cast(global_amax.dptr),\n /*rng_state=*/rng_state,\n /*sm_count=*/sm_count,\n /*stream=*/stream,\n /*k_tile_size=*/k_tile_size);););\n}\n\n} // namespace transformer_engine\n\nvoid nvte_hadamard_transform_cast_fusion_columnwise(const NVTETensor input, NVTETensor output,\n const NVTETensor hadamard_matrix,\n const NVTEQuantizationConfig quant_config,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_hadamard_transform_cast_fusion_columnwise);\n using namespace transformer_engine;\n QuantizationConfig quant_config_cpp;\n if (quant_config != nullptr) {\n quant_config_cpp = *reinterpret_cast(quant_config);\n }\n hadamard_transform_cast_fusion_columnwise(\n *convertNVTETensorCheck(input), *convertNVTETensorCheck(output),\n *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \"common/common.h\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"hadamard_transform_utils.cuh\"\n\nnamespace transformer_engine {\nnamespace {\n\nconstexpr int kThreadsPerWarp = 32;\n\ntemplate \n__device__ __forceinline__ void ComputeKernel(uint32_t b_frag_i[4], uint32_t b_frag_t[4],\n IType* in_sh_ptr, int swizzle_idx,\n uint32_t& local_pre_rht_amax_reg,\n uint32_t& local_amax_reg,\n uint32_t& local_amax_t_reg) {\n uint32_t a_frag[4]; // A matrix fragment\n uint32_t c_frag[4]; // Result fragment\n\n uint32_t temp_amax_reg;\n uint32_t temp_amax_t_reg;\n\n if (kReturnIdentityAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0], a_frag[1], a_frag[2], a_frag[3], b_frag_i[0], b_frag_i[1], b_frag_i[2],\n b_frag_i[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_reg);\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_amax_reg)\n : \"r\"(local_amax_reg), \"r\"(temp_amax_reg));\n }\n\n if (kReturnTransposedAmax) {\n if (!kReturnIdentityAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n } else {\n matrix_transpose_m8_n8_b16_inplace(a_frag[0]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[1]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[2]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[3]);\n }\n\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0], a_frag[2], a_frag[1], a_frag[3], b_frag_t[0], b_frag_t[1], b_frag_t[2],\n b_frag_t[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_t_reg);\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_amax_t_reg)\n : \"r\"(local_amax_t_reg), \"r\"(temp_amax_t_reg));\n }\n\n if (kReturnPreRhtAmax) {\n if (!kReturnIdentityAmax && !kReturnTransposedAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n }\n\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[0])\n : \"r\"(a_frag[0]), \"r\"(a_frag[1]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[2])\n : \"r\"(a_frag[2]), \"r\"(a_frag[3]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[0])\n : \"r\"(a_frag[0]), \"r\"(a_frag[2]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_pre_rht_amax_reg)\n : \"r\"(a_frag[0]), \"r\"(local_pre_rht_amax_reg));\n }\n}\n\ntemplate \n__device__ __host__ constexpr int NextPowerOf2() {\n static_assert(kN > 0, \"kN must be > 0\");\n // Round up to the next power of 2 by counting leading zeros.\n return 1 << (32 - __builtin_clz(kN - 1));\n}\n\ntemplate \n__device__ __forceinline__ void ReduceMax(const float pre_rht_amax, const float identity_amax,\n const float transpose_amax, float* staging_for_pre_rht,\n float* staging_for_identity, float* staging_for_transpose,\n float* output_pre_rht_amax_ptr,\n float* output_identity_amax_ptr,\n float* output_transpose_amax_ptr, const int warpid) {\n // intra-warp reduction\n constexpr int kWarpSize = 32;\n int local_rank = threadIdx.x % 32;\n float warp_pre_rht_amax = kReturnPreRhtAmax ? warp_reduce_max(pre_rht_amax) : 0.0f;\n float warp_identity_amax = kReturnIdentityAmax ? warp_reduce_max(identity_amax) : 0.0f;\n float warp_transpose_amax =\n kReturnTransposedAmax ? warp_reduce_max(transpose_amax) : 0.0f;\n\n // inter-warp reduction\n if (threadIdx.x % 32 == 0) {\n if (kReturnPreRhtAmax) {\n staging_for_pre_rht[warpid] = warp_pre_rht_amax;\n }\n if (kReturnIdentityAmax) {\n staging_for_identity[warpid] = warp_identity_amax;\n }\n if (kReturnTransposedAmax) {\n staging_for_transpose[warpid] = warp_transpose_amax;\n }\n }\n __syncthreads();\n constexpr int kNumWarpsPow2 = NextPowerOf2();\n if (warpid == 0) {\n if (kReturnIdentityAmax) {\n float identity_accum = local_rank < kNumWarps ? staging_for_identity[local_rank] : 0.0f;\n identity_accum = warp_reduce_max(identity_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_identity_amax_ptr, identity_accum);\n }\n }\n }\n if (warpid == 1) {\n if (kReturnTransposedAmax) {\n float transpose_accum = local_rank < kNumWarps ? staging_for_transpose[local_rank] : 0.0f;\n transpose_accum = warp_reduce_max(transpose_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_transpose_amax_ptr, transpose_accum);\n }\n }\n }\n if (warpid == 2) {\n if (kReturnPreRhtAmax) {\n float pre_rht_accum = local_rank < kNumWarps ? staging_for_pre_rht[local_rank] : 0.0f;\n pre_rht_accum = warp_reduce_max(pre_rht_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_pre_rht_amax_ptr, pre_rht_accum);\n }\n }\n }\n}\n\n__launch_bounds__(1) __global__ void ZeroAmaxKernel(float* __restrict__ output_pre_rht_amax_ptr,\n float* __restrict__ output_identity_amax_ptr,\n float* __restrict__ output_transpose_amax_ptr) {\n if (output_pre_rht_amax_ptr != nullptr) {\n *output_pre_rht_amax_ptr = 0;\n }\n if (output_identity_amax_ptr != nullptr) {\n *output_identity_amax_ptr = 0;\n }\n if (output_transpose_amax_ptr != nullptr) {\n *output_transpose_amax_ptr = 0;\n }\n}\n\ntemplate \n__global__ void HadamardAmaxTmaKernel(const __grid_constant__ CUtensorMap tensor_map_input,\n float* __restrict__ output_pre_rht_amax_ptr,\n float* __restrict__ output_identity_amax_ptr,\n float* __restrict__ output_transpose_amax_ptr,\n uint16_t random_sign_mask, uint16_t random_sign_mask_t,\n uint64_t num_rows, uint64_t row_length) {\n#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n\n static_assert(CHUNK_DIM_Y >= BUFF_DIM_Y && CHUNK_DIM_Y % BUFF_DIM_Y == 0);\n static_assert(CHUNK_DIM_X >= BUFF_DIM_X && CHUNK_DIM_X % BUFF_DIM_X == 0);\n\n constexpr size_t STAGES_Y = CHUNK_DIM_Y / BUFF_DIM_Y;\n constexpr size_t STAGES_X = CHUNK_DIM_X / BUFF_DIM_X;\n\n constexpr int kNumWarps = (THREADS_PER_CHUNK * THREADS_PER_Y) / kThreadsPerWarp;\n\n const int input_block_offset_Y = blockIdx.y * CHUNK_DIM_Y;\n const int input_block_offset_X = blockIdx.x * CHUNK_DIM_X;\n\n extern __shared__ __align__(128) char dynamic_shmem[];\n uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem);\n // Manually align dynamic SHMEM per TMA requirements using padding\n // __align__(128) Does not guarantee the pointer to be aligned!\n uint8_t* dshmem = reinterpret_cast((base_shmem_ptr + 127) & ~127ULL);\n\n // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned\n constexpr size_t in_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType);\n IType* in_sh_0 = reinterpret_cast(dshmem);\n dshmem += in_buff_size;\n IType* in_sh_1 = reinterpret_cast(dshmem);\n dshmem += in_buff_size;\n\n IType* in_shs[2] = {in_sh_0, in_sh_1};\n\n constexpr int shmem_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType);\n\n const bool is_master_thread = (threadIdx.x == 0 && threadIdx.y == 0);\n\n // Initialize shared memory barrier with the number of threads participating in the barrier.\n#pragma nv_diag_suppress static_var_with_dynamic_init\n uint64_t* mbar = reinterpret_cast(dshmem);\n dshmem += sizeof(uint64_t) * (STAGES_X * STAGES_Y);\n\n float* max_staging_identity = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n float* max_staging_transpose = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n float* max_staging_pre_rht = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n\n initialize_barriers(mbar,\n is_master_thread);\n\n copy_2d_to_shared(in_shs[0], reinterpret_cast(&tensor_map_input),\n input_block_offset_X, input_block_offset_Y, shmem_buff_size, &mbar[0],\n is_master_thread);\n\n uint32_t had_frag_i[4];\n uint32_t had_frag_t[4];\n get_hadamard_matrix_fragment(\n had_frag_i, random_sign_mask, had_frag_t, random_sign_mask_t);\n\n float local_pre_rht_amax = 0.0;\n float local_amax = 0.0;\n float local_amax_t = 0.0;\n uint32_t local_pre_rht_amax_reg = *reinterpret_cast(&local_pre_rht_amax);\n uint32_t local_amax_reg = *reinterpret_cast(&local_amax);\n uint32_t local_amax_t_reg = *reinterpret_cast(&local_amax_t);\n\n const int warp_id = threadIdx.x / kThreadsPerWarp;\n const int local_rank = threadIdx.x % kThreadsPerWarp;\n const int ld_row_idx = local_rank % kHadamardDimension;\n const int ld_col_idx = local_rank / kHadamardDimension + warp_id * 2;\n const int swizzle_idx = swizzle_128B_atom_32B(ld_row_idx, ld_col_idx);\n\n for (int stage_y = 0; stage_y < STAGES_Y; ++stage_y) {\n for (int stage_x = 0; stage_x < STAGES_X; ++stage_x) {\n int stage = STAGES_X * stage_y + stage_x;\n\n const int next_stage = stage + 1;\n const int next_stage_x = stage_x + 1 == STAGES_X ? 0 : stage_x + 1;\n const int next_stage_y = stage_x + 1 == STAGES_X ? stage_y + 1 : stage_y;\n\n if (next_stage < STAGES_X * STAGES_Y) {\n const int input_global_offset_Y = input_block_offset_Y + next_stage_y * BUFF_DIM_Y;\n const int input_global_offset_X = input_block_offset_X + next_stage_x * BUFF_DIM_X;\n\n copy_2d_to_shared(in_shs[next_stage % 2], // ping-pong\n reinterpret_cast(&tensor_map_input), input_global_offset_X,\n input_global_offset_Y, shmem_buff_size, &mbar[next_stage],\n is_master_thread);\n }\n\n // Wait for the data to have arrived\n ptx::mbarrier_wait_parity(&mbar[stage], 0);\n\n const size_t compute_stage_x_num =\n BUFF_DIM_X / (kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp));\n const size_t compute_stage_y_num = BUFF_DIM_Y / (kHadamardDimension * THREADS_PER_Y);\n\n const size_t in_row_stride = BUFF_DIM_X;\n\n IType* in_sh_ptr = in_shs[stage % 2];\n\n#pragma unroll\n for (size_t compute_stage_y = 0; compute_stage_y < compute_stage_y_num; compute_stage_y++) {\n const int row_idx_offset = (compute_stage_y * kHadamardDimension * THREADS_PER_Y +\n threadIdx.y * kHadamardDimension);\n const int in_row_offset = row_idx_offset * in_row_stride;\n\n#pragma unroll\n for (size_t compute_stage_x = 0; compute_stage_x < compute_stage_x_num; compute_stage_x++) {\n ComputeKernel(\n had_frag_i, had_frag_t,\n in_sh_ptr + in_row_offset +\n (compute_stage_x * kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp)),\n swizzle_idx, local_pre_rht_amax_reg, local_amax_reg, local_amax_t_reg);\n }\n }\n // Ensure all threads have finished their computation before new data over-writes the shared\n // memory.\n __syncthreads();\n // Ensure generic shared-memory accesses are visible before the next TMA write.\n ptx::fence_proxy_async_shared_cta();\n }\n }\n\n const int warpid = (threadIdx.x + threadIdx.y * blockDim.x) / kThreadsPerWarp;\n\n if constexpr (kReturnPreRhtAmax) {\n unpack_max_of_packed_bf16(local_pre_rht_amax_reg, local_pre_rht_amax);\n }\n if constexpr (kReturnIdentityAmax) {\n unpack_max_of_packed_bf16(local_amax_reg, local_amax);\n }\n if constexpr (kReturnTransposedAmax) {\n unpack_max_of_packed_bf16(local_amax_t_reg, local_amax_t);\n }\n\n ReduceMax(\n local_pre_rht_amax, local_amax, local_amax_t, max_staging_pre_rht, max_staging_identity,\n max_staging_transpose, output_pre_rht_amax_ptr, output_identity_amax_ptr,\n output_transpose_amax_ptr, warpid);\n\n destroy_barriers(mbar, is_master_thread);\n#else\n NVTE_DEVICE_ERROR(\"Kernel is only supported on SM 10.0+.\");\n#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n}\n\ntemplate \n__global__ void HadamardTransformKernel(const T* __restrict__ input, T* __restrict__ output,\n T* __restrict__ output_t, uint16_t random_sign_mask,\n uint16_t random_sign_mask_t, uint64_t num_input_rows,\n uint64_t num_input_cols, float* __restrict__ amax,\n float* __restrict__ amax_t, bool inverse_hadamard) {\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900\n static_assert(kHadamardDimension == 16, \"Currently only hadamard dimension 16 is supported.\");\n\n // The whole threadblock will share the same smem.\n extern __shared__ __align__(16) T smem[];\n\n // Each 32 threads process a 16x16 matrix. There is a (y, z) grid of 16x16.\n // If y = 4, z = 4, then each threadblock is processing a 4x4 grid of 16x16 matrices.\n int32_t tid = threadIdx.x;\n int32_t warp_id = threadIdx.y * blockDim.z + threadIdx.z;\n int32_t local_bx = threadIdx.y;\n int32_t local_by = threadIdx.z;\n\n // Define the register fragments\n uint32_t a_frag[4]; // A matrix fragment\n uint32_t b_frag_i[4]; // Transposed Hadamard matrix fragment, used for A @ B(col major)\n uint32_t b_frag_t[4]; // Hadamard matrix fragment, used for A.T @ B.T(col major)\n uint32_t c_frag[4]; // Result fragment\n\n // row and col for each thread. 32 threads will work together in 128 chunk to\n // load the data from global memory to shared memory.\n uint32_t row = tid / (kHadamardDimension * sizeof(T) / sizeof(uint4));\n uint32_t col = tid % (kHadamardDimension * sizeof(T) / sizeof(uint4));\n\n uint32_t smem_index = tid;\n\n uint32_t input_start_col = (blockIdx.x * blockDim.y + local_bx) * kHadamardDimension;\n uint32_t input_start_row = (blockIdx.y * blockDim.z + local_by) * kHadamardDimension;\n\n bool load = (input_start_col < num_input_cols) && (input_start_row < num_input_rows);\n if (!load) {\n // Out of bound, we are returning early. No thread divergence since the whole warp\n // will return early.\n return;\n }\n\n uint64_t global_offset = input_start_col + input_start_row * num_input_cols;\n uint64_t global_offset_t =\n kOutputTrueTransposed ? (input_start_row + input_start_col * num_input_rows) : global_offset;\n\n T* base_smem = smem + kHadamardDimension * kHadamardDimension * warp_id;\n\n uint32_t* smem_b32 = reinterpret_cast(base_smem);\n uint4* smem_b128 = reinterpret_cast(base_smem);\n\n // Asynchronously load the data from global memory to shared memory.\n const uint4* input_b128 = reinterpret_cast(input + global_offset);\n // Each 16x16 chunk is divided into 4 8x8 matrices, we are trying to load each\n // 8x8 chunks consecutively into the smem, so we could leverage ldmatrix m8n8x4\n // to load the data in the tensor core swizzled format.\n __pipeline_memcpy_async(&smem_b128[smem_index],\n &input_b128[row * num_input_cols / (sizeof(uint4) / sizeof(T)) + col],\n sizeof(uint4));\n __pipeline_commit(); // Commit the memcpy. Wait when we are in the computation.\n\n if (inverse_hadamard) {\n get_hadamard_matrix_fragment(b_frag_i, random_sign_mask,\n b_frag_t, random_sign_mask_t);\n } else {\n get_hadamard_matrix_fragment(\n b_frag_i, random_sign_mask, b_frag_t, random_sign_mask_t);\n }\n\n float local_amax = 0.0;\n float local_amax_t = 0.0;\n uint32_t local_amax_reg = *reinterpret_cast(&local_amax);\n uint32_t local_amax_t_reg = *reinterpret_cast(&local_amax_t);\n __pipeline_wait_prior(0);\n\n __syncwarp(); // ensure all lanes finished their cp.async before reading smem\n\n // Load the A to a_frag.\n if constexpr (kComputeIdentity) {\n load_matrix_16x16_from_shared(a_frag[0], a_frag[1], a_frag[2], a_frag[3], smem_b32,\n kHadamardDimension);\n\n // 16x16 @ 16x16 leveraging all threads in the warp.\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0], a_frag[1], a_frag[2], a_frag[3], b_frag_i[0], b_frag_i[1], b_frag_i[2],\n b_frag_i[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], local_amax_reg);\n\n // Store the result to the shared memory in non-transposed order.\n if constexpr (kReturnIdentity) {\n uint4* output_b128 = reinterpret_cast(output + global_offset);\n store_matrix_16x16_to_global(c_frag[0], c_frag[1], c_frag[2], c_frag[3], output_b128,\n num_input_cols);\n }\n }\n\n if constexpr (kComputeTransposed) {\n if (kComputeIdentity) {\n matrix_transpose_m8_n8_b16_inplace(a_frag[0]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[1]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[2]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[3]);\n } else {\n load_matrix_16x16_from_shared(a_frag[0],\n a_frag[2], // NOTE: intentional index swapping\n a_frag[1], // NOTE: intentional index swapping\n a_frag[3], smem_b32, kHadamardDimension);\n }\n\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0],\n // 2,1 is used if we are using movmatrix instruction.\n // Thus loading the matrix in 2,1 order will just be normal.\n // This is to be compatible with the movmatrix instruction.\n a_frag[2], // NOTE: intentional index swapping for transpose purpose.\n a_frag[1], // NOTE: intentional index swapping for transpose purpose.\n a_frag[3], b_frag_t[0], b_frag_t[1], b_frag_t[2], b_frag_t[3], c_frag[0], c_frag[1],\n c_frag[2], c_frag[3], local_amax_t_reg);\n\n // Store the result to the shared memory in non-transposed order.\n if constexpr (kReturnTransposed) {\n uint4* output_t_b128 = reinterpret_cast(output_t + global_offset_t);\n store_matrix_16x16_to_global(\n c_frag[0], c_frag[1], c_frag[2], c_frag[3], output_t_b128,\n kOutputTrueTransposed ? num_input_rows : num_input_cols);\n }\n }\n\n if constexpr (kUpdateIdentityAmax) {\n unpack_max_of_packed_bf16(local_amax_reg, local_amax);\n local_amax = warp_reduce_max(local_amax);\n // broadcast the amax to all threads in a warp from the lane 0\n constexpr int lane_zero = 0;\n local_amax = __shfl_sync(0xFFFFFFFF, local_amax, lane_zero);\n // atomic CAS to output memory.\n if (tid % kThreadsPerWarp == 0) {\n atomicMaxFloat(amax, local_amax);\n }\n }\n if constexpr (kUpdateTransposeAmax) {\n unpack_max_of_packed_bf16(local_amax_t_reg, local_amax_t);\n local_amax_t = warp_reduce_max(local_amax_t);\n // broadcast the amax to all threads in a warp from the lane 0\n constexpr int lane_zero = 0;\n local_amax_t = __shfl_sync(0xFFFFFFFF, local_amax_t, lane_zero);\n // atomic CAS to output memory.\n if (tid % kThreadsPerWarp == 0) {\n atomicMaxFloat(amax_t, local_amax_t);\n }\n }\n#else\n NVTE_DEVICE_ERROR(\"Kernel is only supported on SM 9.0+.\");\n#endif // defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900\n}\n\n} // namespace\n\nvoid hadamard_transform(const Tensor& input_, Tensor& output_, uint16_t random_sign_mask,\n uint16_t random_sign_mask_t, cudaStream_t stream) {\n NVTE_API_CALL(hadamard_transform);\n\n // Check tensors\n // NOTE (frsun): This is non-intuitive, we are writing the result of\n // transposed RHT to the output of rowwise.\n NVTE_CHECK(input_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor must be BF16 tensor, but scaling mode is \",\n to_string(input_.scaling_mode), \".\");\n NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16,\n \"Input tensor must be BF16 tensor, but dtype is \", to_string(input_.dtype()), \".\");\n NVTE_CHECK(input_.dim() >= 2, \"Input must be a 2D tensor.\");\n NVTE_CHECK(output_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Output tensor must be simple tensor, but scaling mode is \",\n to_string(output_.scaling_mode), \".\");\n const SimpleTensor& input = input_.data;\n SimpleTensor output;\n SimpleTensor& output_t = output_.data;\n\n // Check requested outputs\n const bool return_identity = output.dptr != nullptr;\n const bool return_transposed = output_t.dptr != nullptr;\n if (!return_identity && !return_transposed) { // Nothing to do/ill-defined behavior.\n return;\n }\n\n checkCuDriverContext(stream);\n\n const size_t ndim = input.shape.size();\n const size_t row_length = input.shape[ndim - 1];\n size_t num_rows = 1;\n for (size_t i = 0; i < ndim - 1; ++i) {\n num_rows *= input.shape[i];\n }\n\n using IType = bf16;\n\n constexpr int kHadamardDimension = 16;\n NVTE_CHECK(row_length % kHadamardDimension == 0,\n \"row_length must be divisible by hadamard_dimension.\");\n NVTE_CHECK(num_rows % kHadamardDimension == 0,\n \"num_rows must be divisible by hadamard_dimension\");\n\n constexpr uint64_t kThreadBlockX = 4;\n // Configure 4 is used for Hopper, 8 is used for Blackwell for extra memory bandwidth.\n constexpr uint64_t kThreadBlockY = 4;\n\n uint64_t kNumWarpsPerSM = kThreadBlockX * kThreadBlockY;\n\n // The shared memory number of bytes required for **the whole threadblock**.\n size_t shmem_bytes = kHadamardDimension * kHadamardDimension * sizeof(IType) * kNumWarpsPerSM;\n\n dim3 block(kThreadsPerWarp, kThreadBlockX, kThreadBlockY);\n\n dim3 grid(DIVUP(row_length / kHadamardDimension, kThreadBlockX),\n DIVUP(num_rows / kHadamardDimension, kThreadBlockY));\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_transposed, kReturnTransposed,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_identity, kReturnIdentity,\n\n auto kernel =\n HadamardTransformKernel;\n\n cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_bytes);\n\n kernel<<>>(\n reinterpret_cast(input.dptr), reinterpret_cast(output.dptr),\n reinterpret_cast(output_t.dptr), random_sign_mask, random_sign_mask_t,\n num_rows, row_length, nullptr, nullptr, false);););\n\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n// Kernel that will apply the 16x16 hadamard transform the input and input.T, and then\n// get the absolute max value of the result.\nvoid hadamard_transform_amax(const Tensor& input_, Tensor& output_, uint16_t random_sign_mask,\n uint16_t random_sign_mask_t, cudaStream_t stream) {\n NVTE_API_CALL(hadamard_transform_amax);\n#if CUDA_VERSION >= 12080\n\n // Check input tensor\n NVTE_CHECK(input_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor must be BF16 tensor, but scaling mode is \",\n to_string(input_.scaling_mode), \".\");\n NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16,\n \"Input tensor must be BF16 tensor, but dtype is \", to_string(input_.dtype()), \".\");\n NVTE_CHECK(input_.dim() >= 2, \"Input must be a 2D tensor.\");\n const SimpleTensor& input = input_.data;\n\n // Check amax tensors\n SimpleTensor& output_pre_rht_amax = output_.amax;\n SimpleTensor output_identity_amax;\n SimpleTensor& output_transpose_amax = output_.columnwise_amax;\n\n // Check requested outputs\n const bool return_pre_rht_amax = output_pre_rht_amax.dptr != nullptr;\n const bool return_identity_amax = output_identity_amax.dptr != nullptr;\n const bool return_transposed_amax = output_transpose_amax.dptr != nullptr;\n if (!return_identity_amax && !return_transposed_amax &&\n !return_pre_rht_amax) { // Nothing to do/ill-defined behavior.\n return;\n }\n\n // Zero out amaxes if needed\n ZeroAmaxKernel<<<1, 1, 0, stream>>>(reinterpret_cast(output_pre_rht_amax.dptr),\n reinterpret_cast(output_identity_amax.dptr),\n reinterpret_cast(output_transpose_amax.dptr));\n NVTE_CHECK_CUDA(cudaGetLastError());\n\n checkCuDriverContext(stream);\n\n using IType = bf16;\n\n const size_t ndim = input.shape.size();\n const size_t row_length = input.shape[ndim - 1];\n size_t num_rows = 1;\n for (size_t i = 0; i < ndim - 1; ++i) {\n num_rows *= input.shape[i];\n }\n\n constexpr int kHadamardDimension = 16;\n NVTE_CHECK(row_length % kHadamardDimension == 0,\n \"row_length must be divisible by hadamard_dimension.\");\n NVTE_CHECK(num_rows % kHadamardDimension == 0,\n \"num_rows must be divisible by hadamard_dimension\");\n\n constexpr uint64_t kChunkBlockXSmall = 128;\n constexpr uint64_t kChunkBlockYSmall = 128;\n constexpr uint64_t kBuffDimX = 64;\n constexpr uint64_t kBuffDimY = 64;\n\n alignas(64) CUtensorMap tensor_map_input{};\n\n create_2D_tensor_map(\n /*tensorMap=*/tensor_map_input,\n /*tensor=*/input,\n /*globalY=*/num_rows,\n /*globalX=*/row_length,\n /*shmemY=*/kBuffDimY,\n /*shmemX=*/kBuffDimX,\n /*stride_elems=*/row_length,\n /*offset_elems=*/0,\n /*type_num_bits=*/sizeof(IType) * 8,\n /*swizzle=*/CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B);\n\n constexpr uint64_t kThreadBlockX = 4;\n constexpr uint64_t kThreadBlockY = 1;\n constexpr uint64_t kNumWarps = kThreadBlockX * kThreadBlockY;\n\n dim3 block(kThreadBlockX * kThreadsPerWarp, kThreadBlockY);\n\n dim3 grid(DIVUP(row_length, kChunkBlockXSmall), DIVUP(num_rows, kChunkBlockYSmall));\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_transposed_amax, kReturnTransposedAmax,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_identity_amax, kReturnIdentityAmax,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_pre_rht_amax, kReturnPreRhtAmax,\n\n // *2 for ping-pong\n size_t in_sh_size = kBuffDimX * kBuffDimY * 2 * sizeof(IType);\n size_t mbar_size = sizeof(uint64_t) * (kChunkBlockXSmall / kBuffDimX) *\n (kChunkBlockYSmall / kBuffDimY);\n size_t shmem_bytes = in_sh_size + mbar_size + kNumWarps * sizeof(float) * 3;\n // Add padding in case shmem ptr is not aligned to 128 bytes.\n shmem_bytes = (shmem_bytes + 128);\n\n auto kernel = HadamardAmaxTmaKernel<\n IType, kHadamardDimension, kChunkBlockYSmall, kChunkBlockXSmall, kBuffDimY,\n kBuffDimX, kThreadBlockX * kThreadsPerWarp, kThreadBlockY, kReturnPreRhtAmax,\n kReturnIdentityAmax, kReturnTransposedAmax>;\n cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,\n shmem_bytes);\n\n kernel<<>>(\n tensor_map_input, reinterpret_cast(output_pre_rht_amax.dptr),\n reinterpret_cast(output_identity_amax.dptr),\n reinterpret_cast(output_transpose_amax.dptr), random_sign_mask,\n random_sign_mask_t, num_rows, row_length);)));\n\n NVTE_CHECK_CUDA(cudaGetLastError());\n#else\n NVTE_ERROR(\"Hadamard transform requires CUDA 12.8+, but compile-time CUDA version is \",\n CUDA_VERSION);\n#endif // CUDA_VERSION >= 12080\n}\n\n} // namespace transformer_engine\n\nvoid nvte_hadamard_transform(const NVTETensor input, NVTETensor output, int random_sign_mask,\n int random_sign_mask_t, cudaStream_t stream) {\n NVTE_API_CALL(nvte_hadamard_transform);\n using namespace transformer_engine;\n hadamard_transform(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output),\n static_cast(random_sign_mask),\n static_cast(random_sign_mask_t), stream);\n}\n\nvoid nvte_hadamard_transform_amax(const NVTETensor input, NVTETensor output, int random_sign_mask,\n int random_sign_mask_t, cudaStream_t stream) {\n NVTE_API_CALL(nvte_hadamard_transform_amax);\n using namespace transformer_engine;\n hadamard_transform_amax(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output),\n static_cast(random_sign_mask),\n static_cast(random_sign_mask_t), stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \"common/common.h\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"hadamard_transform_utils.cuh\"\n\nnamespace transformer_engine {\nnamespace {\n\nconstexpr int kThreadsPerWarp = 32;\n\ntemplate \n__device__ __forceinline__ void ComputeKernel(uint32_t b_frag_i[4], uint32_t b_frag_t[4],\n IType* in_sh_ptr, int swizzle_idx,\n uint32_t& local_pre_rht_amax_reg,\n uint32_t& local_amax_reg,\n uint32_t& local_amax_t_reg) {\n uint32_t a_frag[4]; // A matrix fragment\n uint32_t c_frag[4]; // Result fragment\n\n uint32_t temp_amax_reg;\n uint32_t temp_amax_t_reg;\n\n if (kReturnIdentityAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0], a_frag[1], a_frag[2], a_frag[3], b_frag_i[0], b_frag_i[1], b_frag_i[2],\n b_frag_i[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_reg);\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_amax_reg)\n : \"r\"(local_amax_reg), \"r\"(temp_amax_reg));\n }\n\n if (kReturnTransposedAmax) {\n if (!kReturnIdentityAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n } else {\n matrix_transpose_m8_n8_b16_inplace(a_frag[0]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[1]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[2]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[3]);\n }\n\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0], a_frag[2], a_frag[1], a_frag[3], b_frag_t[0], b_frag_t[1], b_frag_t[2],\n b_frag_t[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_t_reg);\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_amax_t_reg)\n : \"r\"(local_amax_t_reg), \"r\"(temp_amax_t_reg));\n }\n\n if (kReturnPreRhtAmax) {\n if (!kReturnIdentityAmax && !kReturnTransposedAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n }\n\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[0])\n : \"r\"(a_frag[0]), \"r\"(a_frag[1]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[2])\n : \"r\"(a_frag[2]), \"r\"(a_frag[3]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[0])\n : \"r\"(a_frag[0]), \"r\"(a_frag[2]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_pre_rht_amax_reg)\n : \"r\"(a_frag[0]), \"r\"(local_pre_rht_amax_reg));\n }\n}\n\ntemplate \n__device__ __host__ constexpr int NextPowerOf2() {\n static_assert(kN > 0, \"kN must be > 0\");\n // Round up to the next power of 2 by counting leading zeros.\n return 1 << (32 - __builtin_clz(kN - 1));\n}\n\ntemplate \n__device__ __forceinline__ void ReduceMax(const float pre_rht_amax, const float identity_amax,\n const float transpose_amax, float* staging_for_pre_rht,\n float* staging_for_identity, float* staging_for_transpose,\n float* output_pre_rht_amax_ptr,\n float* output_identity_amax_ptr,\n float* output_transpose_amax_ptr, const int warpid) {\n // intra-warp reduction\n constexpr int kWarpSize = 32;\n int local_rank = threadIdx.x % 32;\n float warp_pre_rht_amax = kReturnPreRhtAmax ? warp_reduce_max(pre_rht_amax) : 0.0f;\n float warp_identity_amax = kReturnIdentityAmax ? warp_reduce_max(identity_amax) : 0.0f;\n float warp_transpose_amax =\n kReturnTransposedAmax ? warp_reduce_max(transpose_amax) : 0.0f;\n\n // inter-warp reduction\n if (threadIdx.x % 32 == 0) {\n if (kReturnPreRhtAmax) {\n staging_for_pre_rht[warpid] = warp_pre_rht_amax;\n }\n if (kReturnIdentityAmax) {\n staging_for_identity[warpid] = warp_identity_amax;\n }\n if (kReturnTransposedAmax) {\n staging_for_transpose[warpid] = warp_transpose_amax;\n }\n }\n __syncthreads();\n constexpr int kNumWarpsPow2 = NextPowerOf2();\n if (warpid == 0) {\n if (kReturnIdentityAmax) {\n float identity_accum = local_rank < kNumWarps ? staging_for_identity[local_rank] : 0.0f;\n identity_accum = warp_reduce_max(identity_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_identity_amax_ptr, identity_accum);\n }\n }\n }\n if (warpid == 1) {\n if (kReturnTransposedAmax) {\n float transpose_accum = local_rank < kNumWarps ? staging_for_transpose[local_rank] : 0.0f;\n transpose_accum = warp_reduce_max(transpose_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_transpose_amax_ptr, transpose_accum);\n }\n }\n }\n if (warpid == 2) {\n if (kReturnPreRhtAmax) {\n float pre_rht_accum = local_rank < kNumWarps ? staging_for_pre_rht[local_rank] : 0.0f;\n pre_rht_accum = warp_reduce_max(pre_rht_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_pre_rht_amax_ptr, pre_rht_accum);\n }\n }\n }\n}\n\n__launch_bounds__(1) __global__ void ZeroAmaxKernel(float* __restrict__ output_pre_rht_amax_ptr,\n float* __restrict__ output_identity_amax_ptr,\n float* __restrict__ output_transpose_amax_ptr) {\n if (output_pre_rht_amax_ptr != nullptr) {\n *output_pre_rht_amax_ptr = 0;\n }\n if (output_identity_amax_ptr != nullptr) {\n *output_identity_amax_ptr = 0;\n }\n if (output_transpose_amax_ptr != nullptr) {\n *output_transpose_amax_ptr = 0;\n }\n}\n\ntemplate \n__global__ void HadamardAmaxTmaKernel(const __grid_constant__ CUtensorMap tensor_map_input,\n float* __restrict__ output_pre_rht_amax_ptr,\n float* __restrict__ output_identity_amax_ptr,\n float* __restrict__ output_transpose_amax_ptr,\n uint16_t random_sign_mask, uint16_t random_sign_mask_t,\n uint64_t num_rows, uint64_t row_length) {\n#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n\n static_assert(CHUNK_DIM_Y >= BUFF_DIM_Y && CHUNK_DIM_Y % BUFF_DIM_Y == 0);\n static_assert(CHUNK_DIM_X >= BUFF_DIM_X && CHUNK_DIM_X % BUFF_DIM_X == 0);\n\n constexpr size_t STAGES_Y = CHUNK_DIM_Y / BUFF_DIM_Y;\n constexpr size_t STAGES_X = CHUNK_DIM_X / BUFF_DIM_X;\n\n constexpr int kNumWarps = (THREADS_PER_CHUNK * THREADS_PER_Y) / kThreadsPerWarp;\n\n const int input_block_offset_Y = blockIdx.y * CHUNK_DIM_Y;\n const int input_block_offset_X = blockIdx.x * CHUNK_DIM_X;\n\n extern __shared__ __align__(128) char dynamic_shmem[];\n uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem);\n // Manually align dynamic SHMEM per TMA requirements using padding\n // __align__(128) Does not guarantee the pointer to be aligned!\n uint8_t* dshmem = reinterpret_cast((base_shmem_ptr + 127) & ~127ULL);\n\n // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned\n constexpr size_t in_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType);\n IType* in_sh_0 = reinterpret_cast(dshmem);\n dshmem += in_buff_size;\n IType* in_sh_1 = reinterpret_cast(dshmem);\n dshmem += in_buff_size;\n\n IType* in_shs[2] = {in_sh_0, in_sh_1};\n\n constexpr int shmem_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType);\n\n const bool is_master_thread = (threadIdx.x == 0 && threadIdx.y == 0);\n\n // Initialize shared memory barrier with the number of threads participating in the barrier.\n#pragma nv_diag_suppress static_var_with_dynamic_init\n uint64_t* mbar = reinterpret_cast(dshmem);\n dshmem += sizeof(uint64_t) * (STAGES_X * STAGES_Y);\n\n float* max_staging_identity = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n float* max_staging_transpose = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n float* max_staging_pre_rht = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n\n initialize_barriers(mbar,\n is_master_thread);\n\n copy_2d_to_shared(in_shs[0], reinterpret_cast(&tensor_map_input),\n input_block_offset_X, input_block_offset_Y, shmem_buff_size, &mbar[0],\n is_master_thread);\n\n uint32_t had_frag_i[4];\n uint32_t had_frag_t[4];\n get_hadamard_matrix_fragment(\n had_frag_i, random_sign_mask, had_frag_t, random_sign_mask_t);\n\n float local_pre_rht_amax = 0.0;\n float local_amax = 0.0;\n float local_amax_t = 0.0;\n uint32_t local_pre_rht_amax_reg = *reinterpret_cast(&local_pre_rht_amax);\n uint32_t local_amax_reg = *reinterpret_cast(&local_amax);\n uint32_t local_amax_t_reg = *reinterpret_cast(&local_amax_t);\n\n const int warp_id = threadIdx.x / kThreadsPerWarp;\n const int local_rank = threadIdx.x % kThreadsPerWarp;\n const int ld_row_idx = local_rank % kHadamardDimension;\n const int ld_col_idx = local_rank / kHadamardDimension + warp_id * 2;\n const int swizzle_idx = swizzle_128B_atom_32B(ld_row_idx, ld_col_idx);\n\n for (int stage_y = 0; stage_y < STAGES_Y; ++stage_y) {\n for (int stage_x = 0; stage_x < STAGES_X; ++stage_x) {\n int stage = STAGES_X * stage_y + stage_x;\n\n const int next_stage = stage + 1;\n const int next_stage_x = stage_x + 1 == STAGES_X ? 0 : stage_x + 1;\n const int next_stage_y = stage_x + 1 == STAGES_X ? stage_y + 1 : stage_y;\n\n if (next_stage < STAGES_X * STAGES_Y) {\n const int input_global_offset_Y = input_block_offset_Y + next_stage_y * BUFF_DIM_Y;\n const int input_global_offset_X = input_block_offset_X + next_stage_x * BUFF_DIM_X;\n\n copy_2d_to_shared(in_shs[next_stage % 2], // ping-pong\n reinterpret_cast(&tensor_map_input), input_global_offset_X,\n input_global_offset_Y, shmem_buff_size, &mbar[next_stage],\n is_master_thread);\n }\n\n // Wait for the data to have arrived\n ptx::mbarrier_wait_parity(&mbar[stage], 0);\n\n const size_t compute_stage_x_num =\n BUFF_DIM_X / (kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp));\n const size_t compute_stage_y_num = BUFF_DIM_Y / (kHadamardDimension * THREADS_PER_Y);\n\n const size_t in_row_stride = BUFF_DIM_X;\n\n IType* in_sh_ptr = in_shs[stage % 2];\n\n#pragma unroll\n for (size_t compute_stage_y = 0; compute_stage_y < compute_stage_y_num; compute_stage_y++) {\n const int row_idx_offset = (compute_stage_y * kHadamardDimension * THREADS_PER_Y +\n threadIdx.y * kHadamardDimension);\n const int in_row_offset = row_idx_offset * in_row_stride;\n\n#pragma unroll\n for (size_t compute_stage_x = 0; compute_stage_x < compute_stage_x_num; compute_stage_x++) {\n ComputeKernel(\n had_frag_i, had_frag_t,\n in_sh_ptr + in_row_offset +\n (compute_stage_x * kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp)),\n swizzle_idx, local_pre_rht_amax_reg, local_amax_reg, local_amax_t_reg);\n }\n }\n // Ensure all threads have finished their computation before new data over-writes the shared\n // memory.\n __syncthreads();\n // Ensure generic shared-memory accesses are visible before the next TMA write.\n ptx::fence_proxy_async_shared_cta();\n }\n }\n\n const int warpid = (threadIdx.x + threadIdx.y * blockDim.x) / kThreadsPerWarp;\n\n if constexpr (kReturnPreRhtAmax) {\n unpack_max_of_packed_bf16(local_pre_rht_amax_reg, local_pre_rht_amax);\n }\n if constexpr (kReturnIdentityAmax) {\n unpack_max_of_packed_bf16(local_amax_reg, local_amax);\n }\n if constexpr (kReturnTransposedAmax) {\n unpack_max_of_packed_bf16(local_amax_t_reg, local_amax_t);\n }\n\n ReduceMax(\n local_pre_rht_amax, local_amax, local_amax_t, max_staging_pre_rht, max_staging_identity,\n max_staging_transpose, output_pre_rht_amax_ptr, output_identity_amax_ptr,\n output_transpose_amax_ptr, warpid);\n\n destroy_barriers(mbar, is_master_thread);\n#else\n NVTE_DEVICE_ERROR(\"Kernel is only supported on SM 10.0+.\");\n#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n}\n\ntemplate \n__global__ void HadamardTransformKernel(const T* __restrict__ input, T* __restrict__ output,\n T* __restrict__ output_t, uint16_t random_sign_mask,\n uint16_t random_sign_mask_t, uint64_t num_input_rows,\n uint64_t num_input_cols, float* __restrict__ amax,\n float* __restrict__ amax_t, bool inverse_hadamard) {\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900\n static_assert(kHadamardDimension == 16, \"Currently only hadamard dimension 16 is supported.\");\n\n // The whole threadblock will share the same smem.\n extern __shared__ __align__(16) T smem[];\n\n // Each 32 threads process a 16x16 matrix. There is a (y, z) grid of 16x16.\n // If y = 4, z = 4, then each threadblock is processing a 4x4 grid of 16x16 matrices.\n int32_t tid = threadIdx.x;\n int32_t warp_id = threadIdx.y * blockDim.z + threadIdx.z;\n int32_t local_bx = threadIdx.y;\n int32_t local_by = threadIdx.z;\n\n // Define the register fragments\n uint32_t a_frag[4]; // A matrix fragment\n uint32_t b_frag_i[4]; // Transposed Hadamard matrix fragment, used for A @ B(col major)\n uint32_t b_frag_t[4]; // Hadamard matrix fragment, used for A.T @ B.T(col major)\n uint32_t c_frag[4]; // Result fragment\n\n // row and col for each thread. 32 threads will work together in 128 chunk to\n // load the data from global memory to shared memory.\n uint32_t row = tid / (kHadamardDimension * sizeof(T) / sizeof(uint4));\n uint32_t col = tid % (kHadamardDimension * sizeof(T) / sizeof(uint4));\n\n uint32_t smem_index = tid;\n\n uint32_t input_start_col = (blockIdx.x * blockDim.y + local_bx) * kHadamardDimension;\n uint32_t input_start_row = (blockIdx.y * blockDim.z + local_by) * kHadamardDimension;\n\n bool load = (input_start_col < num_input_cols) && (input_start_row < num_input_rows);\n if (!load) {\n // Out of bound, we are returning early. No thread divergence since the whole warp\n // will return early.\n return;\n }\n\n uint64_t global_offset = input_start_col + input_start_row * num_input_cols;\n uint64_t global_offset_t =\n kOutputTrueTransposed ? (input_start_row + input_start_col * num_input_rows) : global_offset;\n\n T* base_smem = smem + kHadamardDimension * kHadamardDimension * warp_id;\n\n uint32_t* smem_b32 = reinterpret_cast(base_smem);\n uint4* smem_b128 = reinterpret_cast(base_smem);\n\n // Asynchronously load the data from global memory to shared memory.\n const uint4* input_b128 = reinterpret_cast(input + global_offset);\n // Each 16x16 chunk is divided into 4 8x8 matrices, we are trying to load each\n // 8x8 chunks consecutively into the smem, so we could leverage ldmatrix m8n8x4\n // to load the data in the tensor core swizzled format.\n __pipeline_memcpy_async(&smem_b128[smem_index],\n &input_b128[row * num_input_cols / (sizeof(uint4) / sizeof(T)) + col],\n sizeof(uint4));\n __pipeline_commit(); // Commit the memcpy. Wait when we are in the computation.\n\n if (inverse_hadamard) {\n get_hadamard_matrix_fragment(b_frag_i, random_sign_mask,\n b_frag_t, random_sign_mask_t);\n } else {\n get_hadamard_matrix_fragment(\n b_frag_i, random_sign_mask, b_frag_t, random_sign_mask_t);\n }\n\n float local_amax = 0.0;\n float local_amax_t = 0.0;\n uint32_t local_amax_reg = *reinterpret_cast(&local_amax);\n uint32_t local_amax_t_reg = *reinterpret_cast(&local_amax_t);\n __pipeline_wait_prior(0);\n\n __syncwarp(); // ensure all lanes finished their cp.async before reading smem\n\n // Load the A to a_frag.\n if constexpr (kComputeIdentity) {\n load_matrix_16x16_from_shared(a_frag[0], a_frag[1], a_frag[2], a_frag[3], smem_b32,\n kHadamardDimension);\n\n // 16x16 @ 16x16 leveraging all threads in the warp.\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0], a_frag[1], a_frag[2], a_frag[3], b_frag_i[0], b_frag_i[1], b_frag_i[2],\n b_frag_i[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], local_amax_reg);\n\n // Store the result to the shared memory in non-transposed order.\n if constexpr (kReturnIdentity) {\n uint4* output_b128 = reinterpret_cast(output + global_offset);\n store_matrix_16x16_to_global(c_frag[0], c_frag[1], c_frag[2], c_frag[3], output_b128,\n num_input_cols);\n }\n }\n\n if constexpr (kComputeTransposed) {\n if (kComputeIdentity) {\n matrix_transpose_m8_n8_b16_inplace(a_frag[0]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[1]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[2]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[3]);\n } else {\n load_matrix_16x16_from_shared(a_frag[0],\n a_frag[2], // NOTE: intentional index swapping\n a_frag[1], // NOTE: intentional index swapping\n a_frag[3], smem_b32, kHadamardDimension);\n }\n\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0],\n // 2,1 is used if we are using movmatrix instruction.\n // Thus loading the matrix in 2,1 order will just be normal.\n // This is to be compatible with the movmatrix instruction.\n a_frag[2], // NOTE: intentional index swapping for transpose purpose.\n a_frag[1], // NOTE: intentional index swapping for transpose purpose.\n a_frag[3], b_frag_t[0], b_frag_t[1], b_frag_t[2], b_frag_t[3], c_frag[0], c_frag[1],\n c_frag[2], c_frag[3], local_amax_t_reg);\n\n // Store the result to the shared memory in non-transposed order.\n if constexpr (kReturnTransposed) {\n uint4* output_t_b128 = reinterpret_cast(output_t + global_offset_t);\n store_matrix_16x16_to_global(\n c_frag[0], c_frag[1], c_frag[2], c_frag[3], output_t_b128,\n kOutputTrueTransposed ? num_input_rows : num_input_cols);\n }\n }\n\n if constexpr (kUpdateIdentityAmax) {\n unpack_max_of_packed_bf16(local_amax_reg, local_amax);\n local_amax = warp_reduce_max(local_amax);\n // broadcast the amax to all threads in a warp from the lane 0\n constexpr int lane_zero = 0;\n local_amax = __shfl_sync(0xFFFFFFFF, local_amax, lane_zero);\n // atomic CAS to output memory.\n if (tid % kThreadsPerWarp == 0) {\n atomicMaxFloat(amax, local_amax);\n }\n }\n if constexpr (kUpdateTransposeAmax) {\n unpack_max_of_packed_bf16(local_amax_t_reg, local_amax_t);\n local_amax_t = warp_reduce_max(local_amax_t);\n // broadcast the amax to all threads in a warp from the lane 0\n constexpr int lane_zero = 0;\n local_amax_t = __shfl_sync(0xFFFFFFFF, local_amax_t, lane_zero);\n // atomic CAS to output memory.\n if (tid % kThreadsPerWarp == 0) {\n atomicMaxFloat(amax_t, local_amax_t);\n }\n }\n#else\n NVTE_DEVICE_ERROR(\"Kernel is only supported on SM 9.0+.\");\n#endif // defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900\n}\n\n} // namespace\n\nvoid hadamard_transform(const Tensor& input_, Tensor& output_, uint16_t random_sign_mask,\n uint16_t random_sign_mask_t, hipStream_t stream) {\n NVTE_API_CALL(hadamard_transform);\n\n // Check tensors\n // NOTE (frsun): This is non-intuitive, we are writing the result of\n // transposed RHT to the output of rowwise.\n NVTE_CHECK(input_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor must be BF16 tensor, but scaling mode is \",\n to_string(input_.scaling_mode), \".\");\n NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16,\n \"Input tensor must be BF16 tensor, but dtype is \", to_string(input_.dtype()), \".\");\n NVTE_CHECK(input_.dim() >= 2, \"Input must be a 2D tensor.\");\n NVTE_CHECK(output_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Output tensor must be simple tensor, but scaling mode is \",\n to_string(output_.scaling_mode), \".\");\n const SimpleTensor& input = input_.data;\n SimpleTensor output;\n SimpleTensor& output_t = output_.data;\n\n // Check requested outputs\n const bool return_identity = output.dptr != nullptr;\n const bool return_transposed = output_t.dptr != nullptr;\n if (!return_identity && !return_transposed) { // Nothing to do/ill-defined behavior.\n return;\n }\n\n checkCuDriverContext(stream);\n\n const size_t ndim = input.shape.size();\n const size_t row_length = input.shape[ndim - 1];\n size_t num_rows = 1;\n for (size_t i = 0; i < ndim - 1; ++i) {\n num_rows *= input.shape[i];\n }\n\n using IType = bf16;\n\n constexpr int kHadamardDimension = 16;\n NVTE_CHECK(row_length % kHadamardDimension == 0,\n \"row_length must be divisible by hadamard_dimension.\");\n NVTE_CHECK(num_rows % kHadamardDimension == 0,\n \"num_rows must be divisible by hadamard_dimension\");\n\n constexpr uint64_t kThreadBlockX = 4;\n // Configure 4 is used for Hopper, 8 is used for Blackwell for extra memory bandwidth.\n constexpr uint64_t kThreadBlockY = 4;\n\n uint64_t kNumWarpsPerSM = kThreadBlockX * kThreadBlockY;\n\n // The shared memory number of bytes required for **the whole threadblock**.\n size_t shmem_bytes = kHadamardDimension * kHadamardDimension * sizeof(IType) * kNumWarpsPerSM;\n\n dim3 block(kThreadsPerWarp, kThreadBlockX, kThreadBlockY);\n\n dim3 grid(DIVUP(row_length / kHadamardDimension, kThreadBlockX),\n DIVUP(num_rows / kHadamardDimension, kThreadBlockY));\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_transposed, kReturnTransposed,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_identity, kReturnIdentity,\n\n auto kernel =\n HadamardTransformKernel;\n\n hipFuncSetAttribute(reinterpret_cast(kernel), hipFuncAttributeMaxDynamicSharedMemorySize, shmem_bytes);\n\n kernel<<>>(\n reinterpret_cast(input.dptr), reinterpret_cast(output.dptr),\n reinterpret_cast(output_t.dptr), random_sign_mask, random_sign_mask_t,\n num_rows, row_length, nullptr, nullptr, false);););\n\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n// Kernel that will apply the 16x16 hadamard transform the input and input.T, and then\n// get the absolute max value of the result.\nvoid hadamard_transform_amax(const Tensor& input_, Tensor& output_, uint16_t random_sign_mask,\n uint16_t random_sign_mask_t, hipStream_t stream) {\n NVTE_API_CALL(hadamard_transform_amax);\n#if CUDA_VERSION >= 12080\n\n // Check input tensor\n NVTE_CHECK(input_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING,\n \"Input tensor must be BF16 tensor, but scaling mode is \",\n to_string(input_.scaling_mode), \".\");\n NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16,\n \"Input tensor must be BF16 tensor, but dtype is \", to_string(input_.dtype()), \".\");\n NVTE_CHECK(input_.dim() >= 2, \"Input must be a 2D tensor.\");\n const SimpleTensor& input = input_.data;\n\n // Check amax tensors\n SimpleTensor& output_pre_rht_amax = output_.amax;\n SimpleTensor output_identity_amax;\n SimpleTensor& output_transpose_amax = output_.columnwise_amax;\n\n // Check requested outputs\n const bool return_pre_rht_amax = output_pre_rht_amax.dptr != nullptr;\n const bool return_identity_amax = output_identity_amax.dptr != nullptr;\n const bool return_transposed_amax = output_transpose_amax.dptr != nullptr;\n if (!return_identity_amax && !return_transposed_amax &&\n !return_pre_rht_amax) { // Nothing to do/ill-defined behavior.\n return;\n }\n\n // Zero out amaxes if needed\n ZeroAmaxKernel<<<1, 1, 0, stream>>>(reinterpret_cast(output_pre_rht_amax.dptr),\n reinterpret_cast(output_identity_amax.dptr),\n reinterpret_cast(output_transpose_amax.dptr));\n NVTE_CHECK_CUDA(hipGetLastError());\n\n checkCuDriverContext(stream);\n\n using IType = bf16;\n\n const size_t ndim = input.shape.size();\n const size_t row_length = input.shape[ndim - 1];\n size_t num_rows = 1;\n for (size_t i = 0; i < ndim - 1; ++i) {\n num_rows *= input.shape[i];\n }\n\n constexpr int kHadamardDimension = 16;\n NVTE_CHECK(row_length % kHadamardDimension == 0,\n \"row_length must be divisible by hadamard_dimension.\");\n NVTE_CHECK(num_rows % kHadamardDimension == 0,\n \"num_rows must be divisible by hadamard_dimension\");\n\n constexpr uint64_t kChunkBlockXSmall = 128;\n constexpr uint64_t kChunkBlockYSmall = 128;\n constexpr uint64_t kBuffDimX = 64;\n constexpr uint64_t kBuffDimY = 64;\n\n alignas(64) CUtensorMap tensor_map_input{};\n\n create_2D_tensor_map(\n /*tensorMap=*/tensor_map_input,\n /*tensor=*/input,\n /*globalY=*/num_rows,\n /*globalX=*/row_length,\n /*shmemY=*/kBuffDimY,\n /*shmemX=*/kBuffDimX,\n /*stride_elems=*/row_length,\n /*offset_elems=*/0,\n /*type_num_bits=*/sizeof(IType) * 8,\n /*swizzle=*/CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B);\n\n constexpr uint64_t kThreadBlockX = 4;\n constexpr uint64_t kThreadBlockY = 1;\n constexpr uint64_t kNumWarps = kThreadBlockX * kThreadBlockY;\n\n dim3 block(kThreadBlockX * kThreadsPerWarp, kThreadBlockY);\n\n dim3 grid(DIVUP(row_length, kChunkBlockXSmall), DIVUP(num_rows, kChunkBlockYSmall));\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_transposed_amax, kReturnTransposedAmax,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_identity_amax, kReturnIdentityAmax,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n return_pre_rht_amax, kReturnPreRhtAmax,\n\n // *2 for ping-pong\n size_t in_sh_size = kBuffDimX * kBuffDimY * 2 * sizeof(IType);\n size_t mbar_size = sizeof(uint64_t) * (kChunkBlockXSmall / kBuffDimX) *\n (kChunkBlockYSmall / kBuffDimY);\n size_t shmem_bytes = in_sh_size + mbar_size + kNumWarps * sizeof(float) * 3;\n // Add padding in case shmem ptr is not aligned to 128 bytes.\n shmem_bytes = (shmem_bytes + 128);\n\n auto kernel = HadamardAmaxTmaKernel<\n IType, kHadamardDimension, kChunkBlockYSmall, kChunkBlockXSmall, kBuffDimY,\n kBuffDimX, kThreadBlockX * kThreadsPerWarp, kThreadBlockY, kReturnPreRhtAmax,\n kReturnIdentityAmax, kReturnTransposedAmax>;\n hipFuncSetAttribute(reinterpret_cast(kernel), hipFuncAttributeMaxDynamicSharedMemorySize,\n shmem_bytes);\n\n kernel<<>>(\n tensor_map_input, reinterpret_cast(output_pre_rht_amax.dptr),\n reinterpret_cast(output_identity_amax.dptr),\n reinterpret_cast(output_transpose_amax.dptr), random_sign_mask,\n random_sign_mask_t, num_rows, row_length);)));\n\n NVTE_CHECK_CUDA(hipGetLastError());\n#else\n NVTE_ERROR(\"Hadamard transform requires CUDA 12.8+, but compile-time CUDA version is \",\n CUDA_VERSION);\n#endif // CUDA_VERSION >= 12080\n}\n\n} // namespace transformer_engine\n\nvoid nvte_hadamard_transform(const NVTETensor input, NVTETensor output, int random_sign_mask,\n int random_sign_mask_t, hipStream_t stream) {\n NVTE_API_CALL(nvte_hadamard_transform);\n using namespace transformer_engine;\n hadamard_transform(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output),\n static_cast(random_sign_mask),\n static_cast(random_sign_mask_t), stream);\n}\n\nvoid nvte_hadamard_transform_amax(const NVTETensor input, NVTETensor output, int random_sign_mask,\n int random_sign_mask_t, hipStream_t stream) {\n NVTE_API_CALL(nvte_hadamard_transform_amax);\n using namespace transformer_engine;\n hadamard_transform_amax(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output),\n static_cast(random_sign_mask),\n static_cast(random_sign_mask_t), stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/hadamard_transform/hadamard_transform.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \"common/common.h\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"hadamard_transform_utils.cuh\"\n\nnamespace transformer_engine {\nnamespace {\n\nconstexpr int kMaxTensorsPerKernel = 64;\nconstexpr int kThreadsPerWarp = 32;\n\n__device__ __forceinline__ size_t get_current_tensor_id(\n const ShapeRepresentation shape_rep, const size_t num_tensors, const size_t current_offset,\n const size_t first_logical_dim, const size_t last_logical_dim,\n const int64_t* const __restrict__ offsets_ptr) {\n if (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) {\n const size_t current_row = current_offset / last_logical_dim;\n const size_t rows_per_tensor = first_logical_dim / num_tensors;\n return current_row / rows_per_tensor;\n } else {\n // upper_bound(offsets, current_offset) - 1 in range i in [0..num_tensors)\n size_t low = 0;\n size_t hi = num_tensors; // half-open [low, hi)\n\n while (low < hi) {\n const size_t mid = low + (hi - low) / 2;\n const size_t mid_offset = static_cast(offsets_ptr[mid]);\n\n if (mid_offset <= current_offset) {\n low = mid + 1;\n } else {\n hi = mid;\n }\n }\n\n // low = first index where offsets[low] > current_offset (or low == num_tensors)\n // id = low - 1, but need to evaluate if current_offset < offsets[0]\n return (low == 0) ? 0 : (low - 1);\n }\n}\n\ntemplate \n__device__ __forceinline__ void ComputeKernel(uint32_t b_frag_i[4], uint32_t b_frag_t[4],\n IType* in_sh_ptr, int swizzle_idx,\n uint32_t& local_pre_rht_amax_reg,\n uint32_t& local_amax_reg,\n uint32_t& local_amax_t_reg) {\n uint32_t a_frag[4]; // A matrix fragment\n uint32_t c_frag[4]; // Result fragment\n\n uint32_t temp_amax_reg;\n uint32_t temp_amax_t_reg;\n\n if (kReturnIdentityAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0], a_frag[1], a_frag[2], a_frag[3], b_frag_i[0], b_frag_i[1], b_frag_i[2],\n b_frag_i[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_reg);\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_amax_reg)\n : \"r\"(local_amax_reg), \"r\"(temp_amax_reg));\n }\n\n if (kReturnTransposedAmax) {\n if (!kReturnIdentityAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n } else {\n matrix_transpose_m8_n8_b16_inplace(a_frag[0]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[1]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[2]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[3]);\n }\n\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0], a_frag[2], a_frag[1], a_frag[3], b_frag_t[0], b_frag_t[1], b_frag_t[2],\n b_frag_t[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_t_reg);\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_amax_t_reg)\n : \"r\"(local_amax_t_reg), \"r\"(temp_amax_t_reg));\n }\n\n if (kReturnPreRhtAmax) {\n if (!kReturnIdentityAmax && !kReturnTransposedAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n }\n\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[0])\n : \"r\"(a_frag[0]), \"r\"(a_frag[1]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[2])\n : \"r\"(a_frag[2]), \"r\"(a_frag[3]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[0])\n : \"r\"(a_frag[0]), \"r\"(a_frag[2]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_pre_rht_amax_reg)\n : \"r\"(a_frag[0]), \"r\"(local_pre_rht_amax_reg));\n }\n}\n\ntemplate \n__device__ __host__ constexpr int NextPowerOf2() {\n static_assert(kN > 0, \"kN must be > 0\");\n // Round up to the next power of 2 by counting leading zeros.\n return 1 << (32 - __builtin_clz(kN - 1));\n}\n\ntemplate \n__device__ __forceinline__ void ReduceMax(const float pre_rht_amax, const float identity_amax,\n const float transpose_amax, float* staging_for_pre_rht,\n float* staging_for_identity, float* staging_for_transpose,\n float* output_pre_rht_amax_ptr,\n float* output_identity_amax_ptr,\n float* output_transpose_amax_ptr, const int warpid) {\n // intra-warp reduction\n constexpr int kWarpSize = 32;\n int local_rank = threadIdx.x % 32;\n float warp_pre_rht_amax = kReturnPreRhtAmax ? warp_reduce_max(pre_rht_amax) : 0.0f;\n float warp_identity_amax = kReturnIdentityAmax ? warp_reduce_max(identity_amax) : 0.0f;\n float warp_transpose_amax =\n kReturnTransposedAmax ? warp_reduce_max(transpose_amax) : 0.0f;\n\n // inter-warp reduction\n if (threadIdx.x % 32 == 0) {\n if (kReturnPreRhtAmax) {\n staging_for_pre_rht[warpid] = warp_pre_rht_amax;\n }\n if (kReturnIdentityAmax) {\n staging_for_identity[warpid] = warp_identity_amax;\n }\n if (kReturnTransposedAmax) {\n staging_for_transpose[warpid] = warp_transpose_amax;\n }\n }\n __syncthreads();\n constexpr int kNumWarpsPow2 = NextPowerOf2();\n if (warpid == 0) {\n if (kReturnIdentityAmax) {\n float identity_accum = local_rank < kNumWarps ? staging_for_identity[local_rank] : 0.0f;\n identity_accum = warp_reduce_max(identity_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_identity_amax_ptr, identity_accum);\n }\n }\n }\n if (warpid == 1) {\n if (kReturnTransposedAmax) {\n float transpose_accum = local_rank < kNumWarps ? staging_for_transpose[local_rank] : 0.0f;\n transpose_accum = warp_reduce_max(transpose_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_transpose_amax_ptr, transpose_accum);\n }\n }\n }\n if (warpid == 2) {\n if (kReturnPreRhtAmax) {\n float pre_rht_accum = local_rank < kNumWarps ? staging_for_pre_rht[local_rank] : 0.0f;\n pre_rht_accum = warp_reduce_max(pre_rht_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_pre_rht_amax_ptr, pre_rht_accum);\n }\n }\n }\n}\n\n__global__ void GraphSafeMultiZeroAmaxKernel(const size_t num_tensors, float* amax_rowwise_ptr,\n float* amax_colwise_ptr) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n\n // Assign each thread a range for rowwise and colwise independently\n if (amax_rowwise_ptr != nullptr) {\n for (int i = tid; i < num_tensors; i += stride) {\n amax_rowwise_ptr[i] = 0.f;\n }\n }\n if (amax_colwise_ptr != nullptr) {\n for (int i = tid; i < num_tensors; i += stride) {\n amax_colwise_ptr[i] = 0.f;\n }\n }\n}\n\n__global__ void GraphSafeMultiAmaxMemcpyD2DKernelPreRHT(const size_t num_tensors,\n float* amax_rowwise_ptr,\n float* amax_colwise_ptr) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n\n if (amax_rowwise_ptr != nullptr && amax_colwise_ptr != nullptr) {\n for (; tid < num_tensors; tid += stride) {\n float* output_pre_rht_amax_ptr = amax_rowwise_ptr + tid;\n float* output_transpose_amax_ptr = amax_colwise_ptr + tid;\n *output_transpose_amax_ptr = *output_pre_rht_amax_ptr;\n }\n }\n}\n\ntemplate \n__global__ void GraphSafeGroupHadamardAmaxTmaKernel(\n const __grid_constant__ CUtensorMap tensor_map_input, uint16_t random_sign_mask,\n uint16_t random_sign_mask_t, const ShapeRepresentation shape_rep, const size_t num_tensors,\n const size_t first_logical_dim, const size_t last_logical_dim,\n const int64_t* const __restrict__ offsets_ptr, const int64_t* const __restrict__ first_dims_ptr,\n float* const __restrict__ amax_rowwise_ptr, float* const __restrict__ amax_colwise_ptr) {\n#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n\n float* output_pre_rht_amax_ptr;\n float* output_identity_amax_ptr = nullptr;\n float* output_transpose_amax_ptr;\n\n // calculate the global offset to get tensor id\n size_t global_offset = blockIdx.y * CHUNK_DIM_Y * last_logical_dim;\n // paged stashing: will have input buffer [M, N], where M is larger than sum(first_dims)\n // also need to early return if this CTA is processing a region larger than the last offsets[num_tensors]\n if (global_offset >= offsets_ptr[num_tensors]) {\n return;\n }\n int tensor_id = get_current_tensor_id(shape_rep, num_tensors, global_offset, first_logical_dim,\n last_logical_dim, offsets_ptr);\n output_pre_rht_amax_ptr = static_cast(amax_rowwise_ptr) + tensor_id;\n output_transpose_amax_ptr = static_cast(amax_colwise_ptr) + tensor_id;\n\n static_assert(CHUNK_DIM_Y >= BUFF_DIM_Y && CHUNK_DIM_Y % BUFF_DIM_Y == 0);\n static_assert(CHUNK_DIM_X >= BUFF_DIM_X && CHUNK_DIM_X % BUFF_DIM_X == 0);\n\n constexpr size_t STAGES_Y = CHUNK_DIM_Y / BUFF_DIM_Y;\n constexpr size_t STAGES_X = CHUNK_DIM_X / BUFF_DIM_X;\n\n constexpr int kNumWarps = (THREADS_PER_CHUNK * THREADS_PER_Y) / kThreadsPerWarp;\n\n const int input_block_offset_Y = blockIdx.y * CHUNK_DIM_Y;\n const int input_block_offset_X = blockIdx.x * CHUNK_DIM_X;\n\n extern __shared__ __align__(128) char dynamic_shmem[];\n uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem);\n // Manually align dynamic SHMEM per TMA requirements using padding\n // __align__(128) Does not guarantee the pointer to be aligned!\n uint8_t* dshmem = reinterpret_cast((base_shmem_ptr + 127) & ~127ULL);\n\n // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned\n constexpr size_t in_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType);\n IType* in_sh_0 = reinterpret_cast(dshmem);\n dshmem += in_buff_size;\n IType* in_sh_1 = reinterpret_cast(dshmem);\n dshmem += in_buff_size;\n\n IType* in_shs[2] = {in_sh_0, in_sh_1};\n\n constexpr int shmem_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType);\n\n const bool is_master_thread = (threadIdx.x == 0 && threadIdx.y == 0);\n\n // Initialize shared memory barrier with the number of threads participating in the barrier.\n#pragma nv_diag_suppress static_var_with_dynamic_init\n uint64_t* mbar = reinterpret_cast(dshmem);\n dshmem += sizeof(uint64_t) * (STAGES_X * STAGES_Y);\n\n float* max_staging_identity = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n float* max_staging_transpose = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n float* max_staging_pre_rht = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n\n initialize_barriers(mbar,\n is_master_thread);\n\n copy_2d_to_shared(in_shs[0], reinterpret_cast(&tensor_map_input),\n input_block_offset_X, input_block_offset_Y, shmem_buff_size, &mbar[0],\n is_master_thread);\n\n uint32_t had_frag_i[4];\n uint32_t had_frag_t[4];\n get_hadamard_matrix_fragment(\n had_frag_i, random_sign_mask, had_frag_t, random_sign_mask_t);\n\n float local_pre_rht_amax = 0.0;\n float local_amax = 0.0;\n float local_amax_t = 0.0;\n uint32_t local_pre_rht_amax_reg = *reinterpret_cast(&local_pre_rht_amax);\n uint32_t local_amax_reg = *reinterpret_cast(&local_amax);\n uint32_t local_amax_t_reg = *reinterpret_cast(&local_amax_t);\n\n const int warp_id = threadIdx.x / kThreadsPerWarp;\n const int local_rank = threadIdx.x % kThreadsPerWarp;\n const int ld_row_idx = local_rank % kHadamardDimension;\n const int ld_col_idx = local_rank / kHadamardDimension + warp_id * 2;\n const int swizzle_idx = swizzle_128B_atom_32B(ld_row_idx, ld_col_idx);\n\n for (int stage_y = 0; stage_y < STAGES_Y; ++stage_y) {\n for (int stage_x = 0; stage_x < STAGES_X; ++stage_x) {\n int stage = STAGES_X * stage_y + stage_x;\n\n const int next_stage = stage + 1;\n const int next_stage_x = stage_x + 1 == STAGES_X ? 0 : stage_x + 1;\n const int next_stage_y = stage_x + 1 == STAGES_X ? stage_y + 1 : stage_y;\n\n if (next_stage < STAGES_X * STAGES_Y) {\n const int input_global_offset_Y = input_block_offset_Y + next_stage_y * BUFF_DIM_Y;\n const int input_global_offset_X = input_block_offset_X + next_stage_x * BUFF_DIM_X;\n\n copy_2d_to_shared(in_shs[next_stage % 2], // ping-pong\n reinterpret_cast(&tensor_map_input), input_global_offset_X,\n input_global_offset_Y, shmem_buff_size, &mbar[next_stage],\n is_master_thread);\n }\n\n // Wait for the data to have arrived\n ptx::mbarrier_wait_parity(&mbar[stage], 0);\n\n const size_t compute_stage_x_num =\n BUFF_DIM_X / (kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp));\n const size_t compute_stage_y_num = BUFF_DIM_Y / (kHadamardDimension * THREADS_PER_Y);\n\n const size_t in_row_stride = BUFF_DIM_X;\n\n IType* in_sh_ptr = in_shs[stage % 2];\n\n#pragma unroll\n for (size_t compute_stage_y = 0; compute_stage_y < compute_stage_y_num; compute_stage_y++) {\n const int row_idx_offset = (compute_stage_y * kHadamardDimension * THREADS_PER_Y +\n threadIdx.y * kHadamardDimension);\n const int in_row_offset = row_idx_offset * in_row_stride;\n\n#pragma unroll\n for (size_t compute_stage_x = 0; compute_stage_x < compute_stage_x_num; compute_stage_x++) {\n ComputeKernel(\n had_frag_i, had_frag_t,\n in_sh_ptr + in_row_offset +\n (compute_stage_x * kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp)),\n swizzle_idx, local_pre_rht_amax_reg, local_amax_reg, local_amax_t_reg);\n }\n }\n // Ensure all threads have finished their computation before new data over-writes the shared\n // memory.\n __syncthreads();\n // Ensure generic shared-memory accesses are visible before the next TMA write.\n ptx::fence_proxy_async_shared_cta();\n }\n }\n\n const int warpid = (threadIdx.x + threadIdx.y * blockDim.x) / kThreadsPerWarp;\n\n if constexpr (kReturnPreRhtAmax) {\n unpack_max_of_packed_bf16(local_pre_rht_amax_reg, local_pre_rht_amax);\n }\n if constexpr (kReturnIdentityAmax) {\n unpack_max_of_packed_bf16(local_amax_reg, local_amax);\n }\n if constexpr (kReturnTransposedAmax) {\n unpack_max_of_packed_bf16(local_amax_t_reg, local_amax_t);\n }\n\n ReduceMax(\n local_pre_rht_amax, local_amax, local_amax_t, max_staging_pre_rht, max_staging_identity,\n max_staging_transpose, output_pre_rht_amax_ptr, output_identity_amax_ptr,\n output_transpose_amax_ptr, warpid);\n\n destroy_barriers(mbar, is_master_thread);\n#else\n NVTE_DEVICE_ERROR(\"Kernel is only supported on SM 10.0+.\");\n#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n}\n\n} // namespace\n\n// broadcast_pre_rht_amax: when it's true, hadamard transform will be disabled\n// if at this time, the amax buffers for output expects both amax_rowwise and amax_colwise\n// then call MultiAmaxMemcpyD2DKernelPreRHT to D2D copy the amax values\nvoid group_hadamard_transform_amax_graph_safe(const GroupedTensor* input, GroupedTensor* output,\n uint16_t random_sign_mask,\n uint16_t random_sign_mask_t,\n bool broadcast_pre_rht_amax, cudaStream_t stream) {\n NVTE_API_CALL(group_hadamard_transform_amax_graph_safe);\n#if CUDA_VERSION >= 12080\n\n NVTE_CHECK(input->num_tensors == output->num_tensors,\n \"Number of input and output tensors must be same.\");\n NVTE_CHECK(input->has_data(), \"Cannot quantize tensor without rowwise data.\");\n\n checkCuDriverContext(stream);\n\n bool all_return_pre_rht_amax = output->has_data();\n // there is no rowwise RHT transform in current recipe\n bool all_return_identity_amax = false;\n bool all_return_transposed_amax = output->has_columnwise_data();\n\n NVTE_CHECK(all_return_pre_rht_amax || all_return_identity_amax || all_return_transposed_amax,\n \"At least one of return_pre_rht_amax, return_identity_amax, or return_transposed_amax \"\n \"must be true\");\n\n if (broadcast_pre_rht_amax) {\n NVTE_CHECK(all_return_pre_rht_amax,\n \"broadcast_pre_rht_amax is only supported when we compute pre-RHT amax\");\n // if all_return_identity_amax and all_return_transposed_amax both are false, there is no need to broadcast anything\n broadcast_pre_rht_amax &= (all_return_identity_amax || all_return_transposed_amax);\n }\n\n const size_t num_tensors = input->num_tensors;\n const size_t first_logical_dim = input->logical_shape.data[0];\n const size_t last_logical_dim = input->logical_shape.data[1];\n // const size_t elts_total = first_logical_dim * last_logical_dim;\n NVTE_CHECK(first_logical_dim % 128 == 0,\n \"First dimension of a grouped tensor should be divisible by 128.\");\n NVTE_CHECK(last_logical_dim % 128 == 0,\n \"Last dimension of a grouped tensor should be divisible by 128.\");\n\n float* const amax_rowwise_ptr = reinterpret_cast(output->amax.dptr);\n float* const amax_colwise_ptr = reinterpret_cast(output->columnwise_amax.dptr);\n\n const int64_t* const offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr);\n const int64_t* const first_dims_ptr = reinterpret_cast(output->first_dims.dptr);\n\n // some sanity checks\n if (all_return_pre_rht_amax) {\n NVTE_CHECK(amax_rowwise_ptr != nullptr, \"Amax rowwise pointer should not be nullptr.\");\n }\n if (all_return_transposed_amax) {\n NVTE_CHECK(amax_colwise_ptr != nullptr, \"Amax columnwise pointer should not be nullptr.\");\n }\n\n // Multi zero out multiple amaxes if needed\n dim3 block_setup_amax(kMaxTensorsPerKernel);\n dim3 grid_setup_amax(1);\n GraphSafeMultiZeroAmaxKernel<<>>(\n num_tensors, amax_rowwise_ptr, amax_colwise_ptr);\n NVTE_CHECK_CUDA(cudaGetLastError());\n\n using IType = bf16;\n constexpr int kHadamardDimension = 16;\n\n // four (1x4) 64x64 sub-tiles for ping-pong overlap\n constexpr uint64_t kChunkBlockXSmall = 256;\n constexpr uint64_t kChunkBlockYSmall = 64;\n constexpr uint64_t kBuffDimX = 64;\n constexpr uint64_t kBuffDimY = 64;\n\n alignas(64) CUtensorMap tensor_map_input{};\n\n create_2D_tensor_map(\n /*tensorMap=*/tensor_map_input,\n /*tensor=*/input->data,\n /*globalY=*/first_logical_dim,\n /*globalX=*/last_logical_dim,\n /*shmemY=*/kBuffDimY,\n /*shmemX=*/kBuffDimX,\n /*stride_elems=*/last_logical_dim,\n /*offset_elems=*/0,\n /*type_num_bits=*/sizeof(IType) * 8,\n /*swizzle=*/CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B);\n\n constexpr uint64_t kThreadBlockX = 4;\n constexpr uint64_t kThreadBlockY = 1;\n constexpr uint64_t kNumWarps = kThreadBlockX * kThreadBlockY;\n\n dim3 block(kThreadBlockX * kThreadsPerWarp, kThreadBlockY);\n dim3 grid(DIVUP(last_logical_dim, kChunkBlockXSmall),\n DIVUP(first_logical_dim, kChunkBlockYSmall));\n\n ShapeRepresentation shape_rep = ShapeRepresentation::VARYING_FIRST_DIM;\n if (output->all_same_shape()) {\n shape_rep = ShapeRepresentation::SAME_BOTH_DIMS;\n } else if (output->all_same_first_dim()) {\n shape_rep = ShapeRepresentation::VARYING_LAST_DIM;\n } else if (output->all_same_last_dim()) {\n shape_rep = ShapeRepresentation::VARYING_FIRST_DIM;\n } else if (output->varying_both_dims()) {\n shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS;\n }\n\n const bool is_const_last_dim = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS ||\n shape_rep == ShapeRepresentation::VARYING_FIRST_DIM);\n\n NVTE_CHECK(is_const_last_dim,\n \"Currently we only support const last dimension for graph safe hadamard transform.\");\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n (all_return_transposed_amax && !broadcast_pre_rht_amax), kReturnTransposedAmax,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n (all_return_identity_amax && !broadcast_pre_rht_amax), kReturnIdentityAmax,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n all_return_pre_rht_amax, kReturnPreRhtAmax,\n\n // *2 for ping-pong\n size_t in_sh_size = kBuffDimX * kBuffDimY * 2 * sizeof(IType);\n size_t mbar_size = sizeof(uint64_t) * (kChunkBlockXSmall / kBuffDimX) *\n (kChunkBlockYSmall / kBuffDimY);\n size_t shmem_bytes = in_sh_size + mbar_size + kNumWarps * sizeof(float) * 3;\n // Add padding in case shmem ptr is not aligned to 128 bytes.\n shmem_bytes = (shmem_bytes + 128);\n\n auto kernel = GraphSafeGroupHadamardAmaxTmaKernel<\n IType, kHadamardDimension, kChunkBlockYSmall, kChunkBlockXSmall, kBuffDimY,\n kBuffDimX, kThreadBlockX * kThreadsPerWarp, kThreadBlockY, kReturnPreRhtAmax,\n kReturnIdentityAmax, kReturnTransposedAmax>;\n cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,\n shmem_bytes);\n\n kernel<<>>(\n tensor_map_input, random_sign_mask, random_sign_mask_t, shape_rep, num_tensors,\n first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr,\n amax_rowwise_ptr, amax_colwise_ptr);\n if (broadcast_pre_rht_amax) {\n GraphSafeMultiAmaxMemcpyD2DKernelPreRHT<<>>(num_tensors, amax_rowwise_ptr,\n amax_colwise_ptr);\n })));\n\n NVTE_CHECK_CUDA(cudaGetLastError());\n#else\n NVTE_ERROR(\"Hadamard transform requires CUDA 12.8+, but compile-time CUDA version is \",\n CUDA_VERSION);\n#endif // CUDA_VERSION >= 12080\n}\n\n} // namespace transformer_engine\n\nvoid nvte_group_hadamard_transform_amax_graph_safe(const NVTEGroupedTensor input,\n NVTEGroupedTensor output, int random_sign_mask,\n int random_sign_mask_t, cudaStream_t stream) {\n NVTE_API_CALL(nvte_group_hadamard_transform_amax_graph_safe);\n using namespace transformer_engine;\n\n GroupedTensor* input_tensor = convertNVTEGroupedTensorCheck(input);\n GroupedTensor* output_tensor = convertNVTEGroupedTensorCheck(output);\n\n if (input_tensor->num_tensors == 0) {\n return;\n }\n\n // Call the group tensor Hadamard transform amax implementation.\n group_hadamard_transform_amax_graph_safe(\n input_tensor, output_tensor, static_cast(random_sign_mask),\n static_cast(random_sign_mask_t), false, stream);\n}\n\n// Grouped-tensor amax without doing hadamard transform\nvoid nvte_group_amax_graph_safe(const NVTEGroupedTensor input, NVTEGroupedTensor output,\n cudaStream_t stream) {\n NVTE_API_CALL(nvte_group_amax_graph_safe);\n using namespace transformer_engine;\n\n GroupedTensor* input_tensor = convertNVTEGroupedTensorCheck(input);\n GroupedTensor* output_tensor = convertNVTEGroupedTensorCheck(output);\n\n if (input_tensor->num_tensors == 0) {\n return;\n }\n\n group_hadamard_transform_amax_graph_safe(input_tensor, output_tensor, 0, 0, true, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\n#include \"common/common.h\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"hadamard_transform_utils.cuh\"\n\nnamespace transformer_engine {\nnamespace {\n\nconstexpr int kMaxTensorsPerKernel = 64;\nconstexpr int kThreadsPerWarp = 32;\n\n__device__ __forceinline__ size_t get_current_tensor_id(\n const ShapeRepresentation shape_rep, const size_t num_tensors, const size_t current_offset,\n const size_t first_logical_dim, const size_t last_logical_dim,\n const int64_t* const __restrict__ offsets_ptr) {\n if (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) {\n const size_t current_row = current_offset / last_logical_dim;\n const size_t rows_per_tensor = first_logical_dim / num_tensors;\n return current_row / rows_per_tensor;\n } else {\n // upper_bound(offsets, current_offset) - 1 in range i in [0..num_tensors)\n size_t low = 0;\n size_t hi = num_tensors; // half-open [low, hi)\n\n while (low < hi) {\n const size_t mid = low + (hi - low) / 2;\n const size_t mid_offset = static_cast(offsets_ptr[mid]);\n\n if (mid_offset <= current_offset) {\n low = mid + 1;\n } else {\n hi = mid;\n }\n }\n\n // low = first index where offsets[low] > current_offset (or low == num_tensors)\n // id = low - 1, but need to evaluate if current_offset < offsets[0]\n return (low == 0) ? 0 : (low - 1);\n }\n}\n\ntemplate \n__device__ __forceinline__ void ComputeKernel(uint32_t b_frag_i[4], uint32_t b_frag_t[4],\n IType* in_sh_ptr, int swizzle_idx,\n uint32_t& local_pre_rht_amax_reg,\n uint32_t& local_amax_reg,\n uint32_t& local_amax_t_reg) {\n uint32_t a_frag[4]; // A matrix fragment\n uint32_t c_frag[4]; // Result fragment\n\n uint32_t temp_amax_reg;\n uint32_t temp_amax_t_reg;\n\n if (kReturnIdentityAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0], a_frag[1], a_frag[2], a_frag[3], b_frag_i[0], b_frag_i[1], b_frag_i[2],\n b_frag_i[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_reg);\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_amax_reg)\n : \"r\"(local_amax_reg), \"r\"(temp_amax_reg));\n }\n\n if (kReturnTransposedAmax) {\n if (!kReturnIdentityAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n } else {\n matrix_transpose_m8_n8_b16_inplace(a_frag[0]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[1]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[2]);\n matrix_transpose_m8_n8_b16_inplace(a_frag[3]);\n }\n\n mma_m16_n16_k16_b16_b16_b16_noacc(\n a_frag[0], a_frag[2], a_frag[1], a_frag[3], b_frag_t[0], b_frag_t[1], b_frag_t[2],\n b_frag_t[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_t_reg);\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_amax_t_reg)\n : \"r\"(local_amax_t_reg), \"r\"(temp_amax_t_reg));\n }\n\n if (kReturnPreRhtAmax) {\n if (!kReturnIdentityAmax && !kReturnTransposedAmax) {\n ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3],\n reinterpret_cast(in_sh_ptr) + swizzle_idx);\n }\n\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[0])\n : \"r\"(a_frag[0]), \"r\"(a_frag[1]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[2])\n : \"r\"(a_frag[2]), \"r\"(a_frag[3]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(a_frag[0])\n : \"r\"(a_frag[0]), \"r\"(a_frag[2]));\n asm volatile(\"max.xorsign.abs.bf16x2 %0, %1, %2;\\n\\t\"\n : \"=r\"(local_pre_rht_amax_reg)\n : \"r\"(a_frag[0]), \"r\"(local_pre_rht_amax_reg));\n }\n}\n\ntemplate \n__device__ __host__ constexpr int NextPowerOf2() {\n static_assert(kN > 0, \"kN must be > 0\");\n // Round up to the next power of 2 by counting leading zeros.\n return 1 << (32 - __builtin_clz(kN - 1));\n}\n\ntemplate \n__device__ __forceinline__ void ReduceMax(const float pre_rht_amax, const float identity_amax,\n const float transpose_amax, float* staging_for_pre_rht,\n float* staging_for_identity, float* staging_for_transpose,\n float* output_pre_rht_amax_ptr,\n float* output_identity_amax_ptr,\n float* output_transpose_amax_ptr, const int warpid) {\n // intra-warp reduction\n constexpr int kWarpSize = 32;\n int local_rank = threadIdx.x % 32;\n float warp_pre_rht_amax = kReturnPreRhtAmax ? warp_reduce_max(pre_rht_amax) : 0.0f;\n float warp_identity_amax = kReturnIdentityAmax ? warp_reduce_max(identity_amax) : 0.0f;\n float warp_transpose_amax =\n kReturnTransposedAmax ? warp_reduce_max(transpose_amax) : 0.0f;\n\n // inter-warp reduction\n if (threadIdx.x % 32 == 0) {\n if (kReturnPreRhtAmax) {\n staging_for_pre_rht[warpid] = warp_pre_rht_amax;\n }\n if (kReturnIdentityAmax) {\n staging_for_identity[warpid] = warp_identity_amax;\n }\n if (kReturnTransposedAmax) {\n staging_for_transpose[warpid] = warp_transpose_amax;\n }\n }\n __syncthreads();\n constexpr int kNumWarpsPow2 = NextPowerOf2();\n if (warpid == 0) {\n if (kReturnIdentityAmax) {\n float identity_accum = local_rank < kNumWarps ? staging_for_identity[local_rank] : 0.0f;\n identity_accum = warp_reduce_max(identity_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_identity_amax_ptr, identity_accum);\n }\n }\n }\n if (warpid == 1) {\n if (kReturnTransposedAmax) {\n float transpose_accum = local_rank < kNumWarps ? staging_for_transpose[local_rank] : 0.0f;\n transpose_accum = warp_reduce_max(transpose_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_transpose_amax_ptr, transpose_accum);\n }\n }\n }\n if (warpid == 2) {\n if (kReturnPreRhtAmax) {\n float pre_rht_accum = local_rank < kNumWarps ? staging_for_pre_rht[local_rank] : 0.0f;\n pre_rht_accum = warp_reduce_max(pre_rht_accum);\n if (local_rank == 0) {\n atomicMaxFloat(output_pre_rht_amax_ptr, pre_rht_accum);\n }\n }\n }\n}\n\n__global__ void GraphSafeMultiZeroAmaxKernel(const size_t num_tensors, float* amax_rowwise_ptr,\n float* amax_colwise_ptr) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n\n // Assign each thread a range for rowwise and colwise independently\n if (amax_rowwise_ptr != nullptr) {\n for (int i = tid; i < num_tensors; i += stride) {\n amax_rowwise_ptr[i] = 0.f;\n }\n }\n if (amax_colwise_ptr != nullptr) {\n for (int i = tid; i < num_tensors; i += stride) {\n amax_colwise_ptr[i] = 0.f;\n }\n }\n}\n\n__global__ void GraphSafeMultiAmaxMemcpyD2DKernelPreRHT(const size_t num_tensors,\n float* amax_rowwise_ptr,\n float* amax_colwise_ptr) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n\n if (amax_rowwise_ptr != nullptr && amax_colwise_ptr != nullptr) {\n for (; tid < num_tensors; tid += stride) {\n float* output_pre_rht_amax_ptr = amax_rowwise_ptr + tid;\n float* output_transpose_amax_ptr = amax_colwise_ptr + tid;\n *output_transpose_amax_ptr = *output_pre_rht_amax_ptr;\n }\n }\n}\n\ntemplate \n__global__ void GraphSafeGroupHadamardAmaxTmaKernel(\n const __grid_constant__ CUtensorMap tensor_map_input, uint16_t random_sign_mask,\n uint16_t random_sign_mask_t, const ShapeRepresentation shape_rep, const size_t num_tensors,\n const size_t first_logical_dim, const size_t last_logical_dim,\n const int64_t* const __restrict__ offsets_ptr, const int64_t* const __restrict__ first_dims_ptr,\n float* const __restrict__ amax_rowwise_ptr, float* const __restrict__ amax_colwise_ptr) {\n#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n\n float* output_pre_rht_amax_ptr;\n float* output_identity_amax_ptr = nullptr;\n float* output_transpose_amax_ptr;\n\n // calculate the global offset to get tensor id\n size_t global_offset = blockIdx.y * CHUNK_DIM_Y * last_logical_dim;\n // paged stashing: will have input buffer [M, N], where M is larger than sum(first_dims)\n // also need to early return if this CTA is processing a region larger than the last offsets[num_tensors]\n if (global_offset >= offsets_ptr[num_tensors]) {\n return;\n }\n int tensor_id = get_current_tensor_id(shape_rep, num_tensors, global_offset, first_logical_dim,\n last_logical_dim, offsets_ptr);\n output_pre_rht_amax_ptr = static_cast(amax_rowwise_ptr) + tensor_id;\n output_transpose_amax_ptr = static_cast(amax_colwise_ptr) + tensor_id;\n\n static_assert(CHUNK_DIM_Y >= BUFF_DIM_Y && CHUNK_DIM_Y % BUFF_DIM_Y == 0);\n static_assert(CHUNK_DIM_X >= BUFF_DIM_X && CHUNK_DIM_X % BUFF_DIM_X == 0);\n\n constexpr size_t STAGES_Y = CHUNK_DIM_Y / BUFF_DIM_Y;\n constexpr size_t STAGES_X = CHUNK_DIM_X / BUFF_DIM_X;\n\n constexpr int kNumWarps = (THREADS_PER_CHUNK * THREADS_PER_Y) / kThreadsPerWarp;\n\n const int input_block_offset_Y = blockIdx.y * CHUNK_DIM_Y;\n const int input_block_offset_X = blockIdx.x * CHUNK_DIM_X;\n\n extern __shared__ __align__(128) char dynamic_shmem[];\n uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem);\n // Manually align dynamic SHMEM per TMA requirements using padding\n // __align__(128) Does not guarantee the pointer to be aligned!\n uint8_t* dshmem = reinterpret_cast((base_shmem_ptr + 127) & ~127ULL);\n\n // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned\n constexpr size_t in_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType);\n IType* in_sh_0 = reinterpret_cast(dshmem);\n dshmem += in_buff_size;\n IType* in_sh_1 = reinterpret_cast(dshmem);\n dshmem += in_buff_size;\n\n IType* in_shs[2] = {in_sh_0, in_sh_1};\n\n constexpr int shmem_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType);\n\n const bool is_master_thread = (threadIdx.x == 0 && threadIdx.y == 0);\n\n // Initialize shared memory barrier with the number of threads participating in the barrier.\n#pragma nv_diag_suppress static_var_with_dynamic_init\n uint64_t* mbar = reinterpret_cast(dshmem);\n dshmem += sizeof(uint64_t) * (STAGES_X * STAGES_Y);\n\n float* max_staging_identity = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n float* max_staging_transpose = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n float* max_staging_pre_rht = reinterpret_cast(dshmem);\n dshmem += sizeof(float) * kNumWarps;\n\n initialize_barriers(mbar,\n is_master_thread);\n\n copy_2d_to_shared(in_shs[0], reinterpret_cast(&tensor_map_input),\n input_block_offset_X, input_block_offset_Y, shmem_buff_size, &mbar[0],\n is_master_thread);\n\n uint32_t had_frag_i[4];\n uint32_t had_frag_t[4];\n get_hadamard_matrix_fragment(\n had_frag_i, random_sign_mask, had_frag_t, random_sign_mask_t);\n\n float local_pre_rht_amax = 0.0;\n float local_amax = 0.0;\n float local_amax_t = 0.0;\n uint32_t local_pre_rht_amax_reg = *reinterpret_cast(&local_pre_rht_amax);\n uint32_t local_amax_reg = *reinterpret_cast(&local_amax);\n uint32_t local_amax_t_reg = *reinterpret_cast(&local_amax_t);\n\n const int warp_id = threadIdx.x / kThreadsPerWarp;\n const int local_rank = threadIdx.x % kThreadsPerWarp;\n const int ld_row_idx = local_rank % kHadamardDimension;\n const int ld_col_idx = local_rank / kHadamardDimension + warp_id * 2;\n const int swizzle_idx = swizzle_128B_atom_32B(ld_row_idx, ld_col_idx);\n\n for (int stage_y = 0; stage_y < STAGES_Y; ++stage_y) {\n for (int stage_x = 0; stage_x < STAGES_X; ++stage_x) {\n int stage = STAGES_X * stage_y + stage_x;\n\n const int next_stage = stage + 1;\n const int next_stage_x = stage_x + 1 == STAGES_X ? 0 : stage_x + 1;\n const int next_stage_y = stage_x + 1 == STAGES_X ? stage_y + 1 : stage_y;\n\n if (next_stage < STAGES_X * STAGES_Y) {\n const int input_global_offset_Y = input_block_offset_Y + next_stage_y * BUFF_DIM_Y;\n const int input_global_offset_X = input_block_offset_X + next_stage_x * BUFF_DIM_X;\n\n copy_2d_to_shared(in_shs[next_stage % 2], // ping-pong\n reinterpret_cast(&tensor_map_input), input_global_offset_X,\n input_global_offset_Y, shmem_buff_size, &mbar[next_stage],\n is_master_thread);\n }\n\n // Wait for the data to have arrived\n ptx::mbarrier_wait_parity(&mbar[stage], 0);\n\n const size_t compute_stage_x_num =\n BUFF_DIM_X / (kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp));\n const size_t compute_stage_y_num = BUFF_DIM_Y / (kHadamardDimension * THREADS_PER_Y);\n\n const size_t in_row_stride = BUFF_DIM_X;\n\n IType* in_sh_ptr = in_shs[stage % 2];\n\n#pragma unroll\n for (size_t compute_stage_y = 0; compute_stage_y < compute_stage_y_num; compute_stage_y++) {\n const int row_idx_offset = (compute_stage_y * kHadamardDimension * THREADS_PER_Y +\n threadIdx.y * kHadamardDimension);\n const int in_row_offset = row_idx_offset * in_row_stride;\n\n#pragma unroll\n for (size_t compute_stage_x = 0; compute_stage_x < compute_stage_x_num; compute_stage_x++) {\n ComputeKernel(\n had_frag_i, had_frag_t,\n in_sh_ptr + in_row_offset +\n (compute_stage_x * kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp)),\n swizzle_idx, local_pre_rht_amax_reg, local_amax_reg, local_amax_t_reg);\n }\n }\n // Ensure all threads have finished their computation before new data over-writes the shared\n // memory.\n __syncthreads();\n // Ensure generic shared-memory accesses are visible before the next TMA write.\n ptx::fence_proxy_async_shared_cta();\n }\n }\n\n const int warpid = (threadIdx.x + threadIdx.y * blockDim.x) / kThreadsPerWarp;\n\n if constexpr (kReturnPreRhtAmax) {\n unpack_max_of_packed_bf16(local_pre_rht_amax_reg, local_pre_rht_amax);\n }\n if constexpr (kReturnIdentityAmax) {\n unpack_max_of_packed_bf16(local_amax_reg, local_amax);\n }\n if constexpr (kReturnTransposedAmax) {\n unpack_max_of_packed_bf16(local_amax_t_reg, local_amax_t);\n }\n\n ReduceMax(\n local_pre_rht_amax, local_amax, local_amax_t, max_staging_pre_rht, max_staging_identity,\n max_staging_transpose, output_pre_rht_amax_ptr, output_identity_amax_ptr,\n output_transpose_amax_ptr, warpid);\n\n destroy_barriers(mbar, is_master_thread);\n#else\n NVTE_DEVICE_ERROR(\"Kernel is only supported on SM 10.0+.\");\n#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)\n}\n\n} // namespace\n\n// broadcast_pre_rht_amax: when it's true, hadamard transform will be disabled\n// if at this time, the amax buffers for output expects both amax_rowwise and amax_colwise\n// then call MultiAmaxMemcpyD2DKernelPreRHT to D2D copy the amax values\nvoid group_hadamard_transform_amax_graph_safe(const GroupedTensor* input, GroupedTensor* output,\n uint16_t random_sign_mask,\n uint16_t random_sign_mask_t,\n bool broadcast_pre_rht_amax, hipStream_t stream) {\n NVTE_API_CALL(group_hadamard_transform_amax_graph_safe);\n#if CUDA_VERSION >= 12080\n\n NVTE_CHECK(input->num_tensors == output->num_tensors,\n \"Number of input and output tensors must be same.\");\n NVTE_CHECK(input->has_data(), \"Cannot quantize tensor without rowwise data.\");\n\n checkCuDriverContext(stream);\n\n bool all_return_pre_rht_amax = output->has_data();\n // there is no rowwise RHT transform in current recipe\n bool all_return_identity_amax = false;\n bool all_return_transposed_amax = output->has_columnwise_data();\n\n NVTE_CHECK(all_return_pre_rht_amax || all_return_identity_amax || all_return_transposed_amax,\n \"At least one of return_pre_rht_amax, return_identity_amax, or return_transposed_amax \"\n \"must be true\");\n\n if (broadcast_pre_rht_amax) {\n NVTE_CHECK(all_return_pre_rht_amax,\n \"broadcast_pre_rht_amax is only supported when we compute pre-RHT amax\");\n // if all_return_identity_amax and all_return_transposed_amax both are false, there is no need to broadcast anything\n broadcast_pre_rht_amax &= (all_return_identity_amax || all_return_transposed_amax);\n }\n\n const size_t num_tensors = input->num_tensors;\n const size_t first_logical_dim = input->logical_shape.data[0];\n const size_t last_logical_dim = input->logical_shape.data[1];\n // const size_t elts_total = first_logical_dim * last_logical_dim;\n NVTE_CHECK(first_logical_dim % 128 == 0,\n \"First dimension of a grouped tensor should be divisible by 128.\");\n NVTE_CHECK(last_logical_dim % 128 == 0,\n \"Last dimension of a grouped tensor should be divisible by 128.\");\n\n float* const amax_rowwise_ptr = reinterpret_cast(output->amax.dptr);\n float* const amax_colwise_ptr = reinterpret_cast(output->columnwise_amax.dptr);\n\n const int64_t* const offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr);\n const int64_t* const first_dims_ptr = reinterpret_cast(output->first_dims.dptr);\n\n // some sanity checks\n if (all_return_pre_rht_amax) {\n NVTE_CHECK(amax_rowwise_ptr != nullptr, \"Amax rowwise pointer should not be nullptr.\");\n }\n if (all_return_transposed_amax) {\n NVTE_CHECK(amax_colwise_ptr != nullptr, \"Amax columnwise pointer should not be nullptr.\");\n }\n\n // Multi zero out multiple amaxes if needed\n dim3 block_setup_amax(kMaxTensorsPerKernel);\n dim3 grid_setup_amax(1);\n GraphSafeMultiZeroAmaxKernel<<>>(\n num_tensors, amax_rowwise_ptr, amax_colwise_ptr);\n NVTE_CHECK_CUDA(hipGetLastError());\n\n using IType = bf16;\n constexpr int kHadamardDimension = 16;\n\n // four (1x4) 64x64 sub-tiles for ping-pong overlap\n constexpr uint64_t kChunkBlockXSmall = 256;\n constexpr uint64_t kChunkBlockYSmall = 64;\n constexpr uint64_t kBuffDimX = 64;\n constexpr uint64_t kBuffDimY = 64;\n\n alignas(64) CUtensorMap tensor_map_input{};\n\n create_2D_tensor_map(\n /*tensorMap=*/tensor_map_input,\n /*tensor=*/input->data,\n /*globalY=*/first_logical_dim,\n /*globalX=*/last_logical_dim,\n /*shmemY=*/kBuffDimY,\n /*shmemX=*/kBuffDimX,\n /*stride_elems=*/last_logical_dim,\n /*offset_elems=*/0,\n /*type_num_bits=*/sizeof(IType) * 8,\n /*swizzle=*/CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B);\n\n constexpr uint64_t kThreadBlockX = 4;\n constexpr uint64_t kThreadBlockY = 1;\n constexpr uint64_t kNumWarps = kThreadBlockX * kThreadBlockY;\n\n dim3 block(kThreadBlockX * kThreadsPerWarp, kThreadBlockY);\n dim3 grid(DIVUP(last_logical_dim, kChunkBlockXSmall),\n DIVUP(first_logical_dim, kChunkBlockYSmall));\n\n ShapeRepresentation shape_rep = ShapeRepresentation::VARYING_FIRST_DIM;\n if (output->all_same_shape()) {\n shape_rep = ShapeRepresentation::SAME_BOTH_DIMS;\n } else if (output->all_same_first_dim()) {\n shape_rep = ShapeRepresentation::VARYING_LAST_DIM;\n } else if (output->all_same_last_dim()) {\n shape_rep = ShapeRepresentation::VARYING_FIRST_DIM;\n } else if (output->varying_both_dims()) {\n shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS;\n }\n\n const bool is_const_last_dim = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS ||\n shape_rep == ShapeRepresentation::VARYING_FIRST_DIM);\n\n NVTE_CHECK(is_const_last_dim,\n \"Currently we only support const last dimension for graph safe hadamard transform.\");\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n (all_return_transposed_amax && !broadcast_pre_rht_amax), kReturnTransposedAmax,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n (all_return_identity_amax && !broadcast_pre_rht_amax), kReturnIdentityAmax,\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n all_return_pre_rht_amax, kReturnPreRhtAmax,\n\n // *2 for ping-pong\n size_t in_sh_size = kBuffDimX * kBuffDimY * 2 * sizeof(IType);\n size_t mbar_size = sizeof(uint64_t) * (kChunkBlockXSmall / kBuffDimX) *\n (kChunkBlockYSmall / kBuffDimY);\n size_t shmem_bytes = in_sh_size + mbar_size + kNumWarps * sizeof(float) * 3;\n // Add padding in case shmem ptr is not aligned to 128 bytes.\n shmem_bytes = (shmem_bytes + 128);\n\n auto kernel = GraphSafeGroupHadamardAmaxTmaKernel<\n IType, kHadamardDimension, kChunkBlockYSmall, kChunkBlockXSmall, kBuffDimY,\n kBuffDimX, kThreadBlockX * kThreadsPerWarp, kThreadBlockY, kReturnPreRhtAmax,\n kReturnIdentityAmax, kReturnTransposedAmax>;\n hipFuncSetAttribute(reinterpret_cast(kernel), hipFuncAttributeMaxDynamicSharedMemorySize,\n shmem_bytes);\n\n kernel<<>>(\n tensor_map_input, random_sign_mask, random_sign_mask_t, shape_rep, num_tensors,\n first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr,\n amax_rowwise_ptr, amax_colwise_ptr);\n if (broadcast_pre_rht_amax) {\n GraphSafeMultiAmaxMemcpyD2DKernelPreRHT<<>>(num_tensors, amax_rowwise_ptr,\n amax_colwise_ptr);\n })));\n\n NVTE_CHECK_CUDA(hipGetLastError());\n#else\n NVTE_ERROR(\"Hadamard transform requires CUDA 12.8+, but compile-time CUDA version is \",\n CUDA_VERSION);\n#endif // CUDA_VERSION >= 12080\n}\n\n} // namespace transformer_engine\n\nvoid nvte_group_hadamard_transform_amax_graph_safe(const NVTEGroupedTensor input,\n NVTEGroupedTensor output, int random_sign_mask,\n int random_sign_mask_t, hipStream_t stream) {\n NVTE_API_CALL(nvte_group_hadamard_transform_amax_graph_safe);\n using namespace transformer_engine;\n\n GroupedTensor* input_tensor = convertNVTEGroupedTensorCheck(input);\n GroupedTensor* output_tensor = convertNVTEGroupedTensorCheck(output);\n\n if (input_tensor->num_tensors == 0) {\n return;\n }\n\n // Call the group tensor Hadamard transform amax implementation.\n group_hadamard_transform_amax_graph_safe(\n input_tensor, output_tensor, static_cast(random_sign_mask),\n static_cast(random_sign_mask_t), false, stream);\n}\n\n// Grouped-tensor amax without doing hadamard transform\nvoid nvte_group_amax_graph_safe(const NVTEGroupedTensor input, NVTEGroupedTensor output,\n hipStream_t stream) {\n NVTE_API_CALL(nvte_group_amax_graph_safe);\n using namespace transformer_engine;\n\n GroupedTensor* input_tensor = convertNVTEGroupedTensorCheck(input);\n GroupedTensor* output_tensor = convertNVTEGroupedTensorCheck(output);\n\n if (input_tensor->num_tensors == 0) {\n return;\n }\n\n group_hadamard_transform_amax_graph_safe(input_tensor, output_tensor, 0, 0, true, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/util/cuda_runtime.h\"\n#include \"common/util/curanddx.hpp\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"customized_pipeline.cuh\"\n#include \"cutlass/arch/barrier.h\"\n#include \"cutlass/arch/reg_reconfig.h\"\n#include \"cutlass/cluster_launch.hpp\"\n#include \"cutlass/cutlass.h\"\n#include \"cutlass/detail/sm100_blockscaled_layout.hpp\"\n#include \"cutlass/fast_math.h\"\n#include \"cutlass/float8.h\"\n#include \"cutlass/float_subbyte.h\"\n#include \"cutlass/gemm/collective/builders/sm100_common.inl\"\n#include \"cutlass/numeric_conversion.h\"\n#include \"cutlass/numeric_types.h\"\n#include \"cutlass/pipeline/pipeline.hpp\"\n#include \"cutlass/platform/platform.h\"\n#include \"cutlass/util/GPU_Clock.hpp\"\n#include \"cutlass/util/command_line.h\"\n#include \"cutlass/util/print_error.hpp\"\n\nnamespace transformer_engine {\nnamespace detail {\nnamespace {\n\nusing namespace cute;\n\n// Ensure Tensor refers to cute::Tensor, not transformer_engine::Tensor\nusing cute::Tensor;\n\nconstexpr int kMaxTensorsPerKernel = 64;\nconstexpr int kNVFP4BlockSize = 16;\n\nenum ShapeRepresentation {\n SAME_BOTH_DIMS = 0,\n VARYING_FIRST_DIM = 1,\n VARYING_LAST_DIM = 2,\n VARYING_BOTH_DIMS = 3\n};\n\n__device__ __forceinline__ size_t get_current_tensor_id(\n const ShapeRepresentation shape_rep, const size_t num_tensors, const size_t current_offset,\n const size_t first_logical_dim, const size_t last_logical_dim,\n const int64_t *const __restrict__ offsets_ptr) {\n if (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) {\n const size_t current_row = current_offset / last_logical_dim;\n const size_t rows_per_tensor = first_logical_dim / num_tensors;\n return current_row / rows_per_tensor;\n } else {\n // upper_bound(offsets, current_offset) - 1 in range i in [0..num_tensors)\n size_t low = 0;\n size_t hi = num_tensors; // half-open [low, hi)\n\n while (low < hi) {\n const size_t mid = low + (hi - low) / 2;\n const size_t mid_offset = static_cast(offsets_ptr[mid]);\n\n if (mid_offset <= current_offset) {\n low = mid + 1;\n } else {\n hi = mid;\n }\n }\n\n // low = first index where offsets[low] > current_offset (or low == num_tensors)\n // id = low - 1, but need to evaluate if current_offset < offsets[0]\n return (low == 0) ? 0 : (low - 1);\n }\n}\n\nCUTLASS_DEVICE\ncutlass::Array StochasticNumericConverterBase(\n cutlass::Array const &input, cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n auto output_ptr = reinterpret_cast(&output);\n constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING;\n if constexpr (has_rs) {\n asm volatile(\n \"{\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\\n\"\n \"}\"\n : \"=h\"(output_ptr[0]), \"=h\"(output_ptr[1])\n : \"f\"(input[0]), \"f\"(input[1]), \"f\"(input[2]), \"f\"(input[3]), \"f\"(input[4]), \"f\"(input[5]),\n \"f\"(input[6]), \"f\"(input[7]), \"r\"(rbits[0]), \"r\"(rbits[1]));\n } else {\n NVTE_DEVICE_ERROR(\n \"FP4 cvt PTX instructions are architecture-specific. \"\n \"Try recompiling with sm_XXXa instead of sm_XXX.\");\n }\n return output;\n}\n\nCUTLASS_DEVICE\ncutlass::Array StochasticNumericConverter(\n cutlass::Array const &input, cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n cutlass::Array *result_ptr =\n reinterpret_cast *>(&output);\n cutlass::Array const *source_ptr =\n reinterpret_cast const *>(&input);\n cutlass::Array const *rbits_ptr =\n reinterpret_cast const *>(&rbits);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 2; i++) {\n result_ptr[i] = StochasticNumericConverterBase(source_ptr[i], rbits_ptr[i]);\n }\n return output;\n}\n\ntemplate \nstruct SharedStorage {\n static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;\n static int constexpr EpilogueUnrollFactor = EpilogueUnrollFactor_;\n using AtomThrShapeMNK = cute::Shape<_1, _1, _1>;\n\n using AccumulatorPipeline =\n cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;\n\n static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline =\n cutlass::detail::CustomizedPipelineTmaUmmaAsync,\n AtomThrShapeMNK>;\n using MainloopPipelineStorage = typename MainloopPipeline::SharedStorage;\n using SchedPipeline = cutlass::PipelineCLCFetchAsync;\n using SchedPipelineStorage = typename SchedPipeline::SharedStorage;\n using SchedThrottlePipeline = cutlass::PipelineAsync;\n using SchedThrottlePipelineStorage = typename SchedThrottlePipeline::SharedStorage;\n\n struct TensorStorage : cute::aligned_struct<128, _1> {\n cute::array_aligned> smem_A;\n cute::array_aligned> smem_B;\n } tensors;\n\n alignas(16) AccumulatorPipelineStorage accumulator;\n alignas(16) MainloopPipelineStorage mainloop;\n alignas(16) cute::uint64_t tma_barrier[1];\n alignas(16) SchedPipelineStorage sched;\n alignas(16) SchedThrottlePipelineStorage sched_throttle;\n alignas(16) int32_t atomic_tile_id[SchedulerPipelineStageCount_];\n alignas(16) float global_a_amax[kMaxTensorsPerKernel];\n alignas(16) float global_d_amax[kMaxTensorsPerKernel];\n uint32_t atomic_tile_counter[SchedulerPipelineStageCount_];\n uint32_t tmem_base_ptr;\n};\n\n// Main RHT GEMM kernel entry -- highly templated for flexible architecture/config support\ntemplate \n__launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_graph_safe(\n MShape M, NShape packed_N, KShape K, ClusterShape cluster_shape, ClusterTileShape cluster_tile,\n TA const *A, AStride dA, ASmemLayout sAlayout, CUTE_GRID_CONSTANT TmaLoadA const tma_load_a,\n TB const *B, BStride dB, BSmemLayout sBlayout, CUTE_GRID_CONSTANT TmaLoadB const tma_load_b,\n TQA *QA, QAStride dQA, TSFA *SFA, TSFALayout sfa_layout, TQA *QA_COLWISE, TSFA *SFA_COLWISE,\n float *amax_rowwise, float *amax_colwise, const int64_t *offsets, const int64_t *first_dims,\n size_t num_tensors, ShapeRepresentation shape_rep, uint32_t *tile_scheduler_workspace,\n TiledMMA mma, const size_t *rng_state) {\n using namespace cute;\n\n // Abort immediately if compilation is not supported\n constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY;\n if constexpr (!is_blackwell_arch) {\n NVTE_DEVICE_ERROR(\"RHT fusion is only supported on Blackwell.\");\n return;\n } else {\n static_assert(kEnableRHTColQuant_ || kEnableRowQuant_,\n \"group_row_col_rht_gemm_device_graph_safe must generate row-wise \"\n \"and/or column-wise output.\");\n#if !defined(CUTLASS_ARCH_CLC_ENABLED)\n CUTLASS_NOT_IMPLEMENTED();\n return;\n#endif\n\n using X = Underscore;\n // Accumulator data type for main computation\n using ElementAccumulator = float;\n static int constexpr K_PIPE_MAX = size<3>(ASmemLayout{});\n using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>;\n static uint32_t constexpr kTmaTransactionBytes = cutlass::bits_to_bytes(\n size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v);\n static constexpr bool kEnableStochasticRounding = kEnableStochasticRounding_;\n static constexpr bool kEnableRHTColQuant = kEnableRHTColQuant_;\n static constexpr bool kEnableRowQuant = kEnableRowQuant_;\n static constexpr bool kEnableSwizzleSFOutput = kEnableSwizzleSFOutput_;\n static constexpr bool kUseFastMath = kUseFastMath_;\n\n // Constant for RHT tensor processing (tile size etc)\n static int constexpr RhtTensorSize = 16;\n\n // Get the total number of tokens to process\n // Note that here M is the hidden size, which is the last logical dimension of the input tensor x\n // The kernel is designed in column major, so M is the hidden size\n size_t sum_token_dims = offsets[num_tensors] / M;\n\n // Transaction bytes for TMA transfer on RHT tensor blocks\n static int constexpr kTmaRhtTensorTransactionBytes =\n cutlass::bits_to_bytes(RhtTensorSize * RhtTensorSize * cute::sizeof_bits_v);\n static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;\n static int constexpr SchedulerPipelineStageCount = SchedulerPipelineStageCount_;\n\n // Mainloop pipeline stage calculation, vectorization parameters for scaling factors\n static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{});\n static int constexpr SFVecSize = 16;\n // Swizzle output layout for scaling factor arrays\n using SwizzledSFALayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFDLayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n\n // Mainloop pipeline types for TMA async execution and epilogue cluster scheduling\n using MainloopPipeline =\n cutlass::detail::CustomizedPipelineTmaUmmaAsync;\n using MainloopPipelineState = typename MainloopPipeline::PipelineState;\n using SchedPipeline = cutlass::PipelineCLCFetchAsync;\n using SchedPipelineState = typename SchedPipeline::PipelineState;\n using SchedThrottlePipeline = cutlass::PipelineAsync;\n using SchedThrottlePipelineState = typename SchedThrottlePipeline::PipelineState;\n\n static_assert(ClusterShape{} == Shape<_1, _1, _1>{}, \"ClusterShape must be Shape<_1,_1,_1>\");\n\n using TmemAllocator = cute::TMEM::Allocator1Sm;\n static int constexpr VectorSize = RhtTensorSize;\n\n // Compile-time safety: static shapes required for shared memory layouts\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n // CUTE_STATIC_ASSERT(is_static::value);\n\n auto cluster_size = size<0>(cluster_shape);\n auto mainloop_tiler = Shape<_128, _16, _128>{};\n auto epilogue_tiler = Shape<_128, _128, _128>{};\n\n static int constexpr EpilogueUnrollFactor = size<2>(epilogue_tiler) / size<2>(cluster_tile);\n\n // Get the appropriate blocks for this Cluster\n dim3 cluster_coord_in_grid = cluster_id_in_grid();\n\n // Total number of k-tiles\n int const K_TILE_MAX = min(packed_N, K) / size<2>(epilogue_tiler);\n\n struct TileScheduler {\n uint32_t tiles_in_m = 0;\n uint32_t tiles_in_n = 0;\n uint32_t linear_idx = 0;\n uint32_t next_linear_idx = 0;\n uint32_t start_idx = 0;\n uint32_t tile_m_idx = 0;\n uint32_t tile_n_idx = 0;\n int k_tile_max = 0;\n uint32_t *atomic_tile_index_;\n uint32_t *smem_tile_counter;\n uint32_t atomic_offset;\n cutlass::FastDivmodU64 divmod_tiles_in_m;\n\n CUTLASS_DEVICE TileScheduler(uint32_t tiles_m, uint32_t tiles_n, int kmax,\n uint32_t *atomic_tile_index, uint32_t *smem_tile_counter)\n : tiles_in_m(tiles_m),\n tiles_in_n(tiles_n),\n linear_idx(blockIdx.x),\n next_linear_idx(blockIdx.x),\n start_idx(blockIdx.x),\n k_tile_max(kmax),\n atomic_tile_index_(atomic_tile_index),\n smem_tile_counter(smem_tile_counter),\n atomic_offset(gridDim.x),\n divmod_tiles_in_m(uint64_t(tiles_m)) {\n update_tile_idx();\n }\n CUTLASS_DEVICE void update_tile_idx() {\n uint64_t q, r;\n divmod_tiles_in_m(q, r, uint64_t(linear_idx));\n tile_m_idx = static_cast(r);\n tile_n_idx = static_cast(q) * uint32_t(k_tile_max);\n }\n CUTLASS_DEVICE uint32_t tile_m() const { return tile_m_idx; }\n CUTLASS_DEVICE uint32_t tile_n_base() const { return tile_n_idx; }\n CUTLASS_DEVICE uint32_t tiles_m() const { return tiles_in_m; }\n\n CUTLASS_DEVICE uint32_t tiles_n() const { return tiles_in_n; }\n\n CUTLASS_DEVICE bool is_valid() const {\n return cute::elem_less(cute::make_coord(tile_m(), tile_n_base()),\n cute::make_coord(tiles_in_m, tiles_in_n));\n }\n\n CUTLASS_DEVICE bool is_first_wave() const { return linear_idx == start_idx; }\n\n CUTLASS_DEVICE uint32_t get_linear_tile_idx() const { return linear_idx; }\n\n // Fetch a new tile_id using atomics.\n CUTLASS_DEVICE uint32_t fetch_tile_id_counter(int pred) {\n uint32_t tile_id_counter = 0;\n asm volatile(\n \"{\\n\\t\"\n \".reg .pred p;\\n\\t\"\n \"setp.eq.u32 p, %2, 1;\\n\\t\"\n \"@p atom.global.add.u32 %0, [%1], 1; \\n\\t\"\n \"}\"\n : \"=r\"(tile_id_counter)\n : \"l\"(atomic_tile_index_), \"r\"(pred));\n\n return tile_id_counter;\n }\n\n CUTLASS_DEVICE auto fetch_next_work(SchedPipeline &sched_pipeline,\n SchedPipelineState sched_pipeline_consumer_state) {\n sched_pipeline.consumer_wait(sched_pipeline_consumer_state);\n next_linear_idx = smem_tile_counter[sched_pipeline_consumer_state.index()];\n cutlass::arch::fence_view_async_shared();\n sched_pipeline.consumer_release(sched_pipeline_consumer_state);\n return;\n }\n\n CUTLASS_DEVICE auto advance_to_next_work(SchedPipeline &sched_pipeline,\n SchedPipelineState sched_pipeline_producer_state) {\n uint32_t mbarrier_addr = sched_pipeline.producer_get_barrier(sched_pipeline_producer_state);\n // Wait for clcID buffer to become empty with a flipped phase\n sched_pipeline.producer_acquire(sched_pipeline_producer_state);\n auto is_leading_thread = cute::elect_one_sync();\n uint32_t tile_id_counter = fetch_tile_id_counter(is_leading_thread) + atomic_offset;\n uint32_t smem_addr =\n cute::cast_smem_ptr_to_uint(&smem_tile_counter[sched_pipeline_producer_state.index()]);\n if (is_leading_thread) {\n cute::store_shared_remote(tile_id_counter, smem_addr, mbarrier_addr, 0);\n }\n\n ++sched_pipeline_producer_state;\n return sched_pipeline_producer_state;\n }\n\n CUTLASS_DEVICE auto update_work_tile_info() {\n linear_idx = next_linear_idx;\n update_tile_idx();\n return;\n }\n };\n\n // Allocate and alias shared memory to the kernel's shared storage type\n extern __shared__ char shared_memory[];\n using SharedStorage =\n SharedStorage;\n SharedStorage &shared_storage = *reinterpret_cast(shared_memory);\n\n // Compute the number of tiles in M and N after tiling and assign scheduler\n uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile))));\n uint32_t tiles_in_n = uint32_t(size(ceil_div(sum_token_dims, size<2>(epilogue_tiler))));\n\n TileScheduler scheduler(tiles_in_m, tiles_in_n, K_TILE_MAX, tile_scheduler_workspace,\n shared_storage.atomic_tile_counter);\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n\n // Shapes for accumulated tiles in mainloop and epilogue\n auto acc_shape_mma = make_shape(take<0, 2>(mainloop_tiler), _1{}, _1{});\n auto acc_shape_epilogue = make_shape(take<0, 2>(epilogue_tiler), _1{}, _1{});\n\n // Shape of the accumulator fragment for the main loop pipeline, with pipeline stages appended\n auto acc_mainloop_pipelined_shape = append(acc_shape_mma, Int{});\n auto bulk_tmem_mma = TiledMMA::make_fragment_C(acc_mainloop_pipelined_shape);\n\n // Number of threads assigned for various epilogue roles depending on quantization settings\n static int constexpr NumEpilogueColQuantThreadCount = kEnableRHTColQuant ? 128 : 0;\n static int constexpr NumEpilogueRowQuantThreadCount = kEnableRowQuant ? 256 : 0;\n static int constexpr NumMmaThreadCount = kEnableRHTColQuant ? 32 : 0;\n static int constexpr NumMmaIssueThreadCount = kEnableRHTColQuant ? 1 : 0;\n static int constexpr NumSchedThreads = 32;\n static int constexpr NumMainloopLoadThreads = 32;\n static int constexpr NumEpilogueThreads =\n NumEpilogueColQuantThreadCount + NumEpilogueRowQuantThreadCount;\n\n TmemAllocator tmem_allocator{};\n cutlass::arch::NamedBarrier tmem_allocation_result_barrier(\n NumMmaThreadCount + NumEpilogueColQuantThreadCount,\n cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);\n\n int warp_idx = cutlass::canonical_warp_idx_sync();\n\n // warp assignment\n bool is_mma_warp = (warp_idx == 0);\n bool is_dma_warp = (warp_idx == 1);\n bool is_sched_warp = (warp_idx == 2);\n bool is_epilogue_col_quant_warp = (warp_idx >= 4 && warp_idx <= 7);\n bool is_epilogue_row_quant_warp = (warp_idx >= 8 && warp_idx <= 15);\n\n typename MainloopPipeline::Params mainloop_pipeline_params;\n if (is_dma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;\n }\n if (is_mma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;\n }\n mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp;\n mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes;\n mainloop_pipeline_params.initializing_warp = 0;\n mainloop_pipeline_params.num_consumers =\n NumEpilogueRowQuantThreadCount + NumMmaIssueThreadCount;\n\n MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params,\n cluster_shape, cute::true_type{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n\n MainloopPipelineState mainloop_pipe_consumer_state;\n MainloopPipelineState mainloop_pipe_producer_state =\n cutlass::make_producer_start_state();\n\n using AccumulatorPipeline =\n cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState;\n using AccumulatorPipelineInitBarriers = cute::bool_constant;\n\n AccumulatorPipelineState accumulator_pipe_consumer_state;\n AccumulatorPipelineState accumulator_pipe_producer_state =\n cutlass::make_producer_start_state();\n\n typename AccumulatorPipeline::Params accumulator_pipeline_params;\n if (is_mma_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer;\n }\n if (is_epilogue_col_quant_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer;\n }\n // Only one producer thread arrives on this barrier.\n accumulator_pipeline_params.producer_arv_count = 1;\n accumulator_pipeline_params.consumer_arv_count =\n size(AtomThrShapeMNK{}) * NumEpilogueColQuantThreadCount;\n accumulator_pipeline_params.initializing_warp = 1;\n AccumulatorPipeline accumulator_pipeline(\n shared_storage.accumulator, accumulator_pipeline_params, cluster_shape,\n AccumulatorPipelineInitBarriers{}, cute::true_type{}); // Delay mask calculation\n typename SchedPipeline::Params sched_pipeline_params;\n if (is_sched_warp) {\n sched_pipeline_params.role = SchedPipeline::ThreadCategory::ProducerConsumer;\n } else {\n sched_pipeline_params.role = SchedPipeline::ThreadCategory::Consumer;\n }\n sched_pipeline_params.producer_blockid = 0;\n sched_pipeline_params.producer_arv_count = 1;\n sched_pipeline_params.consumer_arv_count =\n NumSchedThreads +\n cluster_size * (NumMainloopLoadThreads + NumEpilogueThreads + NumMmaThreadCount);\n sched_pipeline_params.transaction_bytes = sizeof(uint32_t);\n sched_pipeline_params.initializing_warp = 3;\n SchedPipeline sched_pipeline(shared_storage.sched, sched_pipeline_params, cluster_shape);\n SchedPipelineState sched_pipeline_consumer_state;\n SchedPipelineState sched_pipeline_producer_state =\n cutlass::make_producer_start_state();\n\n typename SchedThrottlePipeline::Params sched_throttle_pipeline_params;\n if (is_dma_warp) {\n sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Producer;\n }\n if (is_sched_warp) {\n sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Consumer;\n }\n sched_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads;\n sched_throttle_pipeline_params.consumer_arv_count = NumSchedThreads;\n sched_throttle_pipeline_params.dst_blockid = 0;\n sched_throttle_pipeline_params.initializing_warp = 4;\n\n SchedThrottlePipeline sched_throttle_pipeline(shared_storage.sched_throttle,\n sched_throttle_pipeline_params);\n SchedThrottlePipelineState sched_pipeline_throttle_consumer_state;\n SchedThrottlePipelineState sched_pipeline_throttle_producer_state =\n cutlass::make_producer_start_state();\n\n if (warp_idx == 2 && elect_one_sync()) {\n cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);\n }\n __syncthreads();\n\n // Warp group roles: DMA (global->shared copy), MMA (tensor core gemm), scheduler, column quantizer, row quantizer\n if (is_dma_warp) {\n // Warp responsible for loading input from global to shared memory using TMA (Tensor Memory Access).\n cutlass::arch::warpgroup_reg_dealloc<32>();\n // Get TMA tensors for input matrix A and B (Hadamard/transform matrix) from global memory.\n Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, packed_N));\n Tensor mB = tma_load_b.get_tma_tensor(make_shape(RhtTensorSize, RhtTensorSize));\n\n // Partition tensors for tiling according to the mainloop and cluster tilers.\n Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{});\n Tensor gB_nk =\n local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k)\n\n // Shared memory tensors for pipeline\n Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()),\n sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()),\n sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n // Determine warp/tile positioning\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n // Partition global to local fragments for A and B\n Tensor tCgA = thr_mma.partition_A(gA_mk); // (MMA,MMA_M,MMA_K,k)\n Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k)\n\n Layout cta_layout_mnk = make_layout(cluster_shape);\n Layout cta_layout_vmnk =\n tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{}));\n auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster);\n\n auto [tAgA, tAsA] =\n tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),\n group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA));\n\n auto [tBgB, tBsB] =\n tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),\n group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB));\n\n uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);\n uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);\n if constexpr (kEnableRHTColQuant) {\n if (elect_one_sync()) {\n cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0],\n kTmaRhtTensorTransactionBytes);\n copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0),\n tBsB(_, 0));\n }\n }\n\n do {\n // is_first_wave indicates whether this scheduler wave is the first among a group.\n bool is_first_wave = scheduler.is_first_wave();\n uint32_t skip_wait = is_first_wave;\n auto tAgA_mk = tAgA(_, scheduler.tile_m(), _);\n int k_tile = 0;\n\n sched_throttle_pipeline.producer_acquire(sched_pipeline_throttle_producer_state);\n sched_throttle_pipeline.producer_commit(sched_pipeline_throttle_producer_state);\n ++sched_pipeline_throttle_producer_state;\n CUTLASS_PRAGMA_NO_UNROLL\n while (k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n()) {\n int k_tile_idx_n = scheduler.tile_n_base() + k_tile;\n ++k_tile;\n skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount);\n mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state);\n using BarrierType = typename MainloopPipeline::ProducerBarrierType;\n BarrierType *tma_barrier =\n mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state);\n int write_stage = mainloop_pipe_producer_state.index();\n ++mainloop_pipe_producer_state;\n if (cute::elect_one_sync()) {\n copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n),\n tAsA(_, write_stage));\n }\n }\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n // scheduler.advance();\n } while (scheduler.is_valid());\n mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);\n } else if (is_mma_warp) {\n // This warp executes the main tensor core matrix-multiply-accumulate for the Hadamard transform.\n cutlass::arch::warpgroup_reg_dealloc<32>();\n if constexpr (kEnableRHTColQuant) {\n // Setup shared memory fragments for A and B tiles.\n Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()),\n sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()),\n sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n // Allocate \"fragments\" -- these are actually umma smem descriptors\n Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)\n Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE)\n\n mma.accumulate_ = UMMA::ScaleOut::Zero;\n\n tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns,\n &shared_storage.tmem_base_ptr);\n __syncwarp();\n tmem_allocation_result_barrier.arrive();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_mma.data() = tmem_base_ptr;\n // Wait until the B (Hadamard) tensor copy is complete\n cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/);\n do {\n uint32_t skip_wait = K_TILE_MAX <= 0;\n\n auto barrier_token =\n mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0;\n k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) {\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n int read_stage = mainloop_pipe_consumer_state.index();\n auto tCrA_mk = tCrA(_, _, _, read_stage);\n auto tCrB_nk = tCrB(_, _, 0, 0);\n CUTLASS_PRAGMA_UNROLL\n for (int k_block = 0; k_block < size<2>(tCrA) / EpilogueUnrollFactor; ++k_block) {\n int accumulator_k_block =\n accumulator_pipe_producer_state.index() * EpilogueUnrollFactor;\n int tCrA_k_block = k_block * EpilogueUnrollFactor;\n accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < EpilogueUnrollFactor; i++) {\n auto accumulators = bulk_tmem_mma(_, _, _, accumulator_k_block + i);\n gemm(mma, tCrA_mk(_, _, tCrA_k_block + i), tCrB_nk, accumulators);\n }\n\n accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);\n ++accumulator_pipe_producer_state;\n }\n auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n skip_wait = k_tile >= K_TILE_MAX;\n mainloop_pipeline.umma_consumer_release(curr_mainloop_pipe_consumer_state);\n barrier_token =\n mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n }\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n tmem_allocator.release_allocation_lock();\n accumulator_pipeline.producer_tail(accumulator_pipe_producer_state);\n tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);\n }\n } else if (is_sched_warp) {\n // Scheduler warp manages tile assignment and pipeline progress for warps\n cutlass::arch::warpgroup_reg_dealloc<32>();\n do {\n sched_throttle_pipeline.consumer_wait(sched_pipeline_throttle_consumer_state);\n sched_throttle_pipeline.consumer_release(sched_pipeline_throttle_consumer_state);\n ++sched_pipeline_throttle_consumer_state;\n sched_pipeline_producer_state =\n scheduler.advance_to_next_work(sched_pipeline, sched_pipeline_producer_state);\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n } else if (is_epilogue_col_quant_warp) {\n // Warp responsible for quantizing output of Hadamard transform to FP4 for columnwise usage,\n // and writing result tensors/scales to global memory.\n cutlass::arch::warpgroup_reg_alloc<192>();\n if constexpr (kEnableRHTColQuant) {\n using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x;\n\n auto acc_epilogue_pipelined_shape =\n append(acc_shape_epilogue, Int{});\n auto bulk_tmem_epilogue_layout = make_layout(\n acc_epilogue_pipelined_shape,\n make_stride(stride<0>(bulk_tmem_mma), Int<0>{}, Int<0>{}, size<1>(epilogue_tiler)));\n auto bulk_tmem_epilogue = make_tensor(make_tmem_ptr(), bulk_tmem_epilogue_layout);\n\n // Use 256-bit fragments for aligned bulk stores\n static int constexpr FragmentSize = 256 / sizeof_bits_v;\n\n // Wait for TMEM allocation for this pipeline to finish\n tmem_allocation_result_barrier.arrive_and_wait();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_epilogue.data() = tmem_base_ptr;\n int global_thread_idx = threadIdx.x;\n int local_thread_idx = global_thread_idx % cutlass::NumThreadsPerWarpGroup;\n // g2s load all global_d_amax\n CUTLASS_PRAGMA_NO_UNROLL\n for (int g = local_thread_idx; g < num_tensors; g += NumEpilogueColQuantThreadCount) {\n shared_storage.global_d_amax[g] = __ldg(reinterpret_cast(amax_colwise + g));\n }\n\n size_t rng_seed = 0;\n size_t rng_offset = 0;\n // Setup RNG for stochastic rounding\n if constexpr (kEnableStochasticRounding) {\n rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0;\n rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0;\n }\n // TODO(zhongbo): double check the logic here\n int group_idx = get_current_tensor_id(\n shape_rep, num_tensors, (scheduler.tile_n_base() * size<1>(epilogue_tiler)) * M,\n packed_N, M, offsets);\n\n // Determine quantization scale factor layouts/output splits for this group\n TSFDLayout sfd_layout;\n int cur_N = static_cast(first_dims[group_idx]);\n if constexpr (kEnableSwizzleSFOutput) {\n sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{});\n } else {\n sfd_layout = make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)),\n make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{})));\n }\n // Build output tensors for columns and their quant scales\n // TODO(zhongbo): double check the logic here\n Tensor mD = make_tensor(cute::subbyte_iterator(reinterpret_cast(\n reinterpret_cast(QA_COLWISE) + offsets[group_idx] / 2)),\n make_shape(M, cur_N), DStride{}); // (M,packed_N)\n Tensor gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n\n // for every tensor [x, y] row major, x y both a multiple of 128\n // both of its rowwise and colwise scaling factors will have exactly x * y / 16 elements in FP8 E4M3\n Tensor mSFD = make_tensor(\n make_gmem_ptr(reinterpret_cast(reinterpret_cast(SFA_COLWISE) +\n offsets[group_idx] / kNVFP4BlockSize)),\n sfd_layout);\n Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n\n Tensor gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler));\n\n // Setup tile-level TMEM (t2r) and global memory (r2g) copy descriptors\n auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{}));\n auto tiled_r2g =\n make_tiled_copy_D(Copy_Atom{}, tiled_t2r);\n auto thr_t2r = tiled_t2r.get_slice(local_thread_idx);\n auto thr_r2g = tiled_r2g.get_slice(local_thread_idx);\n\n cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount,\n cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);\n // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release}\n static constexpr float fp4_max = 6.0f;\n static constexpr float fp8_max = 448.0f;\n static constexpr float fp4_max_inv = 1.0f / fp4_max;\n float c_global_amax_val = shared_storage.global_d_amax[group_idx];\n float global_encode_scale = c_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / c_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n float global_decode_scale = 1.0f / global_encode_scale;\n\n // Scaling factor for fast math path\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n\n do {\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0;\n k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();\n ++k_tile) {\n int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler);\n\n // TODO(zhongbo): double check the logic here\n int cur_group_idx = get_current_tensor_id(\n shape_rep, num_tensors, global_tile_n_offset * M, packed_N, M, offsets);\n\n if (cur_group_idx != group_idx) {\n group_idx = cur_group_idx;\n c_global_amax_val = shared_storage.global_d_amax[group_idx];\n // update amax\n global_encode_scale = c_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / c_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n global_decode_scale = 1.0f / global_encode_scale;\n global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n // TODO(zhongbo): double check the logic here\n cur_N = first_dims[group_idx];\n if constexpr (kEnableSwizzleSFOutput) {\n sfd_layout =\n tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{});\n } else {\n sfd_layout =\n make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)),\n make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{})));\n }\n // update tensor\n mD = make_tensor(cute::subbyte_iterator(reinterpret_cast(\n reinterpret_cast(QA_COLWISE) + offsets[group_idx] / 2)),\n make_shape(M, cur_N), DStride{});\n gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n mSFD = make_tensor(make_gmem_ptr(reinterpret_cast(\n reinterpret_cast(SFA_COLWISE) +\n offsets[group_idx] / kNVFP4BlockSize)),\n sfd_layout);\n gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n\n gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler));\n }\n int group_start_offset = offsets[group_idx] / M;\n int local_tile_n_idx =\n (global_tile_n_offset - group_start_offset) / size<1>(epilogue_tiler);\n Tensor tDgD_mn = gD_mn_view(_, _, _, scheduler.tile_m(), local_tile_n_idx);\n\n Tensor tDgSFD_mn = gSFD_mn(_, _, scheduler.tile_m(), local_tile_n_idx);\n accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state);\n\n auto Acc = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index());\n Tensor tDtAcc = thr_t2r.partition_S(Acc); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDgD = thr_t2r.partition_D(tDgD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n\n Tensor tTR_rAcc = make_tensor(\n shape(tDgD)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDrD = make_tensor(shape(tDgD));\n Tensor tTR_rAcc_frag =\n recast>(coalesce(tTR_rAcc));\n Tensor tDrD_frag = recast>(coalesce(tDrD));\n\n Tensor src = thr_r2g.retile_S(tDrD);\n Tensor dst = thr_r2g.retile_D(tDgD);\n\n Tensor tDgSFD_view = make_tensor(\n tDgSFD_mn.data(), make_layout(make_shape(shape(tDgSFD_mn), Int<1>{}, Int<1>{}),\n make_stride(stride(tDgSFD_mn), Int<0>{}, Int<0>{})));\n Tensor tDgSFD = filter(thr_t2r.partition_D(tDgSFD_view));\n Tensor tDrSFD = make_tensor(shape(tDgSFD));\n\n static int constexpr NumVecs = size(tDgD) / VectorSize;\n Tensor tD_rRowSFD_frg = recast>(tDrSFD);\n\n // Compute amax and quantization scales for this tile\n cutlass::maximum_absolute_value_reduction<\n cutlass::Array, true>\n amax_reduction;\n cutlass::Array vec_maxs;\n cutlass::Array pvscales;\n // Copy from TMEM to registers\n copy(tiled_t2r, tDtAcc, tTR_rAcc);\n cutlass::arch::fence_view_async_tmem_load();\n accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state);\n ++accumulator_pipe_consumer_state;\n\n if constexpr (!kUseFastMath) {\n // Downcast to BF16 for bit-wise compatibility with\n // unfused kernels\n auto convert_accum_to_bf16 =\n cutlass::NumericArrayConverter{};\n auto convert_bf16_to_accum =\n cutlass::NumericArrayConverter{};\n tTR_rAcc_frag(_0{}) =\n convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{})));\n tTR_rAcc_frag(_1{}) =\n convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_1{})));\n }\n\n auto compute_frgs = reinterpret_cast *>(\n tTR_rAcc_frag.data());\n auto output_frgs = reinterpret_cast *>(tDrD_frag.data());\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]);\n }\n\n pvscales = cutlass::multiplies>{}(\n vec_maxs, global_encode_scale_multiplier);\n auto pvscales_cvted =\n cutlass::NumericArrayConverter{}(pvscales);\n\n tD_rRowSFD_frg(_0{}) = pvscales_cvted;\n auto qpvscale_ups = cutlass::NumericArrayConverter{}(\n tD_rRowSFD_frg(_0{}));\n auto qpvscale_scaled =\n cutlass::multiplies>{}(\n qpvscale_ups, global_decode_scale);\n cutlass::Array acc_scales;\n if constexpr (kUseFastMath) {\n // Fast math: compute approximate reciprocal\n acc_scales =\n cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled);\n } else {\n // Accurate math: compute reciprocal with division\n acc_scales = cutlass::divides>{}(\n 1.0, qpvscale_scaled);\n }\n\n // Prepare stochastic rounding random state if enabled\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n transformer_engine::curanddx::detail::philox4x32_native_state<\n NVTE_BUILD_NUM_PHILOX_ROUNDS>\n rng;\n // \"Prefetch\" a stochastic rounding state for the first tile\n if constexpr (kEnableStochasticRounding) {\n const size_t rng_sequence = global_thread_idx + k_tile * 512 +\n scheduler.get_linear_tile_idx() * K_TILE_MAX * 512;\n rng.init(rng_seed, rng_sequence, rng_offset);\n }\n CUTLASS_PRAGMA_UNROLL\n // Apply round/quantize to each fragment, with or without stochastic rounding\n for (int v = 0; v < NumVecs; v++) {\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(\n acc_scales[v], cutlass::platform::numeric_limits::max());\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(\n cutlass::multiplies>{}(\n compute_frgs[v], acc_scale),\n *reinterpret_cast *>(&random_uint4));\n } else {\n output_frgs[v] =\n cutlass::NumericArrayConverter{}(\n cutlass::multiplies>{}(\n compute_frgs[v], acc_scale));\n }\n }\n\n // Write quantized FP4 tile and dequant scale to gmem\n copy(tiled_r2g, src, dst);\n copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFD, tDgSFD);\n }\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n }\n } else if (is_epilogue_row_quant_warp) {\n // Warp responsible for quantizing the input (before Hadamard transform) to FP4 for row-wise usage.\n cutlass::arch::warpgroup_reg_alloc<136>();\n if constexpr (kEnableRowQuant) {\n using S2RVectorType = uint128_t;\n\n int global_thread_idx = threadIdx.x;\n int local_thread_idx = global_thread_idx % 256;\n size_t rng_seed = 0;\n size_t rng_offset = 0;\n // g2s load all global_a_amax for all groups/tensors\n CUTLASS_PRAGMA_NO_UNROLL\n for (int g = local_thread_idx; g < num_tensors; g += NumEpilogueRowQuantThreadCount) {\n shared_storage.global_a_amax[g] = __ldg(reinterpret_cast(amax_rowwise + g));\n }\n // RNG for stochastic rounding\n if constexpr (kEnableStochasticRounding) {\n rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0;\n rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0;\n }\n // Input/output tensors/partitions for row quant warp\n Tensor mQA =\n make_tensor(cute::subbyte_iterator(QA), make_layout(make_shape(M, packed_N), dQA));\n Tensor gQA_mn = local_tile(mQA, epilogue_tiler, make_coord(_, _, _), Step<_1, X, _1>{});\n Tensor mSFA = make_tensor(make_gmem_ptr(SFA), sfa_layout);\n\n Tensor gSFA_mn = local_tile(mSFA, epilogue_tiler, make_coord(_, _, _),\n Step<_1, X, _1>{}); // (BLK_M,BLK_N)\n // Swizzled shared memory A tile, with layout\n Tensor sA = as_position_independent_swizzle_tensor(group_modes<0, 2>(\n coalesce(make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()),\n sAlayout)))); // (BLOCK_M, BLOCK_M,PIPE)\n\n // Set up layouts for partitioning \u2013 tile-by-warp, with vector granularity\n using S2RWarpLayout = Layout>;\n using WarpGroupLayout = Layout>;\n using S2RThreadLayout = decltype(blocked_product(S2RWarpLayout{}, WarpGroupLayout{}));\n using S2RValLayout = Layout, _1>>;\n using S2RAtomA = Copy_Atom;\n using R2GAtomQA = Copy_Atom;\n using R2GAtomSFA = Copy_Atom;\n auto tiled_s2r = make_tiled_copy(S2RAtomA{}, S2RThreadLayout{}, S2RValLayout{});\n auto tiled_r2g_QA = make_tiled_copy(R2GAtomQA{}, S2RThreadLayout{}, S2RValLayout{});\n auto tiled_r2g_SFA = make_tiled_copy(R2GAtomSFA{}, S2RThreadLayout{}, S2RValLayout{});\n\n auto thr_s2r = tiled_s2r.get_slice(local_thread_idx);\n auto thr_r2g_QA = tiled_r2g_QA.get_slice(local_thread_idx);\n auto thr_r2g_SFA = tiled_r2g_SFA.get_slice(local_thread_idx);\n Tensor tQAsA = thr_s2r.partition_S(sA); // (Copy, Copy_M, Copy_N, PIPE)\n\n // Allocate temporary register tensors for copying quantization => output\n Tensor tQArA = make_tensor_like(\n make_layout(tQAsA(_, _, _, _0{}).shape())); // (Copy, Copy_M, Copy_N)\n Tensor tQAgQA = thr_r2g_QA.partition_S(gQA_mn);\n Tensor tQArQA = make_tensor_like(tQAgQA(_, _, _, _0{}, _0{}));\n\n Tensor tQAgSFA = thr_r2g_SFA.partition_S(gSFA_mn);\n Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{}));\n\n // Will result in barrier_id=10 passed to bar.sync instr as cutlass adds 8\n // in order to go over the reserved named barrier count.\n constexpr int row_quant_barrier_id = 2;\n cutlass::arch::NamedBarrier::sync(NumEpilogueRowQuantThreadCount, row_quant_barrier_id);\n\n int group_idx = get_current_tensor_id(\n shape_rep, num_tensors, (scheduler.tile_n_base() * size<1>(epilogue_tiler)) * M,\n packed_N, M, offsets);\n float a_global_amax_val = shared_storage.global_a_amax[group_idx];\n // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release}\n static constexpr float fp4_max = 6.0f;\n static constexpr float fp8_max = 448.0f;\n static constexpr float fp4_max_inv = 1.0f / fp4_max;\n float global_encode_scale = a_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / a_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n\n float global_decode_scale = 1.0f / global_encode_scale;\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n auto sfa_converter = cutlass::NumericConverter{};\n do {\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0;\n k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) {\n int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler);\n\n int cur_group_idx = get_current_tensor_id(\n shape_rep, num_tensors, global_tile_n_offset * M, packed_N, M, offsets);\n if (cur_group_idx != group_idx) {\n group_idx = cur_group_idx;\n a_global_amax_val = shared_storage.global_a_amax[group_idx];\n // Update group quantization parameters/scaling\n global_encode_scale = a_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / a_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n global_decode_scale = 1.0f / global_encode_scale;\n global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n }\n\n auto tQAgSFA_mn =\n tQAgSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto tQAgQA_mn = tQAgQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state);\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n copy(tiled_s2r, tQAsA(_, _, _, mainloop_pipe_consumer_state.index()), tQArA);\n cutlass::arch::fence_view_async_shared();\n mainloop_pipeline.consumer_release(mainloop_pipe_consumer_state);\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n\n // static int constexpr NumVecs = size(tQArA) / VectorSize;\n cutlass::maximum_absolute_value_reduction<\n cutlass::Array, true>\n amax_reduction;\n auto compute_frgs = reinterpret_cast *>(tQArA.data());\n auto output_frgs = reinterpret_cast *>(\n raw_pointer_cast(tQArQA.data()));\n Tensor amax =\n make_tensor(prepend(take<1, rank(tQArA)>(tQArA.shape()), _1{}));\n Tensor pvscales = make_tensor_like(amax);\n transformer_engine::curanddx::detail::philox4x32_native_state<\n NVTE_BUILD_NUM_PHILOX_ROUNDS>\n rng;\n if constexpr (kEnableStochasticRounding) {\n const size_t rng_sequence = global_thread_idx + k_tile * 512 +\n scheduler.get_linear_tile_idx() * K_TILE_MAX * 512 +\n tiles_in_m * tiles_in_n * K_TILE_MAX * 512;\n rng.init(rng_seed, rng_sequence, rng_offset);\n }\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < size<1>(group_modes<1, rank(tQArA)>(tQArA)); v++) {\n auto amax_view = group_modes<1, rank(amax)>(amax);\n auto pvscales_view = group_modes<1, rank(pvscales)>(pvscales);\n auto compute_frgs_up =\n cutlass::NumericArrayConverter{}(\n compute_frgs[v]);\n amax_view(_0{}, v) = amax_reduction(ElementAccumulator(0), compute_frgs_up);\n pvscales_view(_0{}, v) = cutlass::multiplies{}(\n amax_view(_0{}, v), global_encode_scale_multiplier);\n filter(tQArSFA)(v) = sfa_converter(pvscales_view(_0{}, v));\n auto qpvscale_ups =\n cutlass::NumericConverter{}(filter(tQArSFA)(v));\n auto qpvscale_scaled =\n cutlass::multiplies{}(qpvscale_ups, global_decode_scale);\n ElementAccumulator acc_scales;\n if constexpr (kUseFastMath) {\n // Fast math: compute approximate reciprocal\n acc_scales = cutlass::reciprocal_approximate_ftz{}(\n qpvscale_scaled);\n } else {\n // Accurate math: compute reciprocal with division\n acc_scales = cutlass::divides{}(1.0, qpvscale_scaled);\n }\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(\n acc_scales, cutlass::platform::numeric_limits::max());\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(\n cutlass::multiplies>{}(\n compute_frgs_up, acc_scale),\n *reinterpret_cast *>(&random_uint4));\n } else {\n output_frgs[v] =\n cutlass::NumericArrayConverter{}(\n cutlass::multiplies>{}(\n compute_frgs_up, acc_scale));\n }\n }\n copy(tiled_r2g_QA, tQArQA, tQAgQA_mn);\n copy(tiled_r2g_SFA, filter(tQArSFA), filter(tQAgSFA_mn));\n }\n // scheduler.advance();\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n }\n\n } else {\n cutlass::arch::warpgroup_reg_dealloc<32>();\n }\n }\n} // NOLINT(readability/fn_size)\n\ntemplate \nvoid group_row_col_rht_gemm_ntt_w_sfc_graph_safe(\n int packed_sequence_length, int hidden_size, size_t num_tensors, ShapeRepresentation shape_rep,\n TA const *A, TB const *B, TQA *QA, TSFA *SFA, TQA *QA_COLWISE, TSFA *SFA_COLWISE,\n float *amax_rowwise, float *amax_colwise, const int64_t *offsets, const int64_t *first_dims,\n const size_t *rng_state, uint32_t *tile_scheduler_workspace, uint32_t sm_count,\n cudaStream_t stream, int k_tile_size = 1024) {\n using namespace cute;\n static int constexpr SFVecSize = 16;\n static int constexpr RhtTensorSize = 16;\n\n static_assert(RhtTensorSize == 16, \"RhtTensorSize must be 16\");\n using LinearSFALayout = decltype(make_layout(make_shape(make_shape(Int{}, 0), 0),\n make_stride(make_stride(_0{}, _1{}), 0)));\n using LinearSFDLayout = decltype(make_layout(make_shape(0, make_shape(Int{}, 0)),\n make_stride(0, make_stride(_0{}, _1{}))));\n\n using SwizzledSFALayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFDLayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFALayout = decltype(tile_to_shape(\n SwizzledSFALayoutAtom{}, make_shape(hidden_size, packed_sequence_length), Step<_1, _2>{}));\n using SwizzledSFDLayout = decltype(tile_to_shape(\n SwizzledSFDLayoutAtom{}, make_shape(hidden_size, packed_sequence_length), Step<_2, _1>{}));\n\n using SFALayout = cute::conditional_t;\n using SFDLayout = cute::conditional_t;\n SFALayout sfa_layout;\n SFDLayout sfd_layout;\n\n if constexpr (kEnableSwizzleSFOutput) {\n sfa_layout = tile_to_shape(SwizzledSFALayoutAtom{},\n make_shape(hidden_size, packed_sequence_length), Step<_1, _2>{});\n sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{},\n make_shape(hidden_size, packed_sequence_length), Step<_2, _1>{});\n } else {\n sfa_layout = make_layout(\n make_shape(make_shape(Int{}, hidden_size / SFVecSize), packed_sequence_length),\n make_stride(make_stride(_0{}, _1{}), hidden_size / SFVecSize));\n sfd_layout = make_layout(\n make_shape(hidden_size, make_shape(Int{}, packed_sequence_length / SFVecSize)),\n make_stride(packed_sequence_length / SFVecSize, make_stride(_0{}, _1{})));\n }\n\n // Define shapes (dynamic)\n auto M = hidden_size;\n auto N = packed_sequence_length;\n Tensor tensorA = make_tensor(A, make_shape(hidden_size, packed_sequence_length), LayoutLeft{});\n Tensor tensorB = make_tensor(B, make_shape(RhtTensorSize, RhtTensorSize), LayoutLeft{});\n Tensor tensorQA = make_tensor(QA, make_shape(hidden_size, packed_sequence_length), LayoutLeft{});\n Tensor tensorSFA = make_tensor(SFA, sfa_layout);\n\n // Define strides (from tensors)\n auto dA = stride(tensorA); // (dM,dK)\n auto dB = stride(tensorB); // (dN,dK)\n auto dD = LayoutRight{}; // (dM,dN)\n auto dQA = stride(tensorQA); // (dM,dK)\n using ClusterShape = Shape<_1, _1, _1>;\n auto cluster_shape = ClusterShape{};\n auto cluster_tile_shape = Shape<_128, Int, Int>{};\n auto cluster_tile_mainloop = Shape<_128, Int, _128>{};\n\n // Each mainloop / epilogue loads 128 x 64 tiles while each MMA proceeds with 128 x 16 tiles\n static int constexpr EpilogueUnrollFactor =\n size<2>(cluster_tile_mainloop) / size<2>(cluster_tile_shape);\n // Construct the MMA\n auto mma = make_tiled_mma(\n SM100_MMA_F16BF16_SS(cluster_tile_shape), size<1>(cluster_tile_shape),\n UMMA::Major::MN, UMMA::Major::MN>{},\n Layout>{});\n\n // Assert that the TiledMMA uses all CTAs in the CGA.\n CUTE_STATIC_ASSERT_V(size(cluster_shape) == size(mma));\n CUTE_STATIC_ASSERT_V(evenly_divides(cluster_tile_shape, tile_shape(mma)));\n\n // Determine the A and B shapes\n auto mma_shape_B =\n partition_shape_B(mma, make_shape(size<1>(cluster_tile_shape), size<2>(cluster_tile_shape)));\n\n using TiledMma = decltype(mma);\n using AtomThrID = typename TiledMma::AtomThrID;\n\n using SmemShape_M = decltype(shape_div(\n shape<0>(cluster_tile_shape),\n shape_div(shape<0>(cluster_tile_shape), size<0>(cluster_tile_shape) / size(AtomThrID{}))));\n using SmemShape_N = decltype(shape_div(\n shape<1>(cluster_tile_shape),\n shape_div(shape<1>(cluster_tile_shape), size<1>(cluster_tile_shape) / size(AtomThrID{}))));\n using SmemShape_K = decltype(cute::get<2>(cluster_tile_shape));\n\n using SmemLayoutAtomB =\n decltype(cutlass::gemm::collective::detail::sm100_smem_selector());\n\n auto mma_shape_A = partition_shape_A(\n mma, make_shape(size<0>(cluster_tile_mainloop), size<2>(cluster_tile_mainloop)));\n using SmemShape_M_A =\n decltype(shape_div(shape<0>(cluster_tile_mainloop),\n shape_div(shape<0>(cluster_tile_mainloop),\n size<0>(cluster_tile_mainloop) / size(AtomThrID{}))));\n using SmemShape_K_A = decltype(cute::get<2>(cluster_tile_mainloop));\n using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<\n cute::UMMA::Major::MN, TA, SmemShape_M_A, SmemShape_K_A>());\n\n static uint32_t constexpr TotalTmemRows = 128;\n static uint32_t constexpr Sm100TmemCapacityColumns = 512;\n static uint32_t constexpr TotalTmem = TotalTmemRows * Sm100TmemCapacityColumns;\n static uint32_t constexpr AccumulatorPipelineStageCount =\n TotalTmem / (cute::size<0>(cluster_tile_shape) * cute::size<1>(cluster_tile_shape));\n\n // Define the smem layouts (static)\n // Calculate max pipeline stages based on Blackwell SM100's 232KB shared memory\n constexpr int SchedulerPipelineStageCount = 4;\n static int constexpr MainloopPipelineBytes = sizeof(\n typename cutlass::detail::CustomizedPipelineTmaUmmaAsync<1, Shape<_1, _1, _1>,\n Shape<_1, _1, _1>>::SharedStorage);\n\n static int constexpr SchedulerWorkspaceBytes = sizeof(int) * SchedulerPipelineStageCount;\n static int constexpr SchedulerThrottlePipelineBytes =\n sizeof(typename cutlass::PipelineAsync::SharedStorage);\n static int constexpr SchedulerPipelineBytes =\n sizeof(typename cutlass::PipelineCLCFetchAsync::SharedStorage);\n\n static int constexpr TmemDeallocBytes = sizeof(cutlass::arch::ClusterBarrier);\n static int constexpr BTensorBytes = cute::size(mma_shape_B) * sizeof(TB);\n static int constexpr AccPipelineBytes = sizeof(\n typename cutlass::PipelineUmmaAsync>::SharedStorage);\n static int constexpr TmemBasePtrsBytes = sizeof(uint32_t);\n static int constexpr kBlackwellSmemSize = 232448; // 232KB in bytes\n static int constexpr kBytesPerStage =\n cute::size(mma_shape_A) * sizeof(TA) + MainloopPipelineBytes;\n static int constexpr kReservedBytes = SchedulerWorkspaceBytes + SchedulerThrottlePipelineBytes +\n SchedulerPipelineBytes + TmemBasePtrsBytes +\n TmemDeallocBytes + BTensorBytes +\n AccPipelineBytes; // Reserve for barriers and other uses\n static int constexpr kMaxStages = (kBlackwellSmemSize - kReservedBytes) / kBytesPerStage;\n auto sP = Int{}; // SMEM pipelines\n\n auto sA = UMMA::tile_to_mma_shape(SmemLayoutAtomA{}, append(mma_shape_A, sP),\n Step<_2, _1, _3>{}); // (MMA,MMA_M,MMA_K,PIPE)\n auto sB = UMMA::tile_to_mma_shape(SmemLayoutAtomB{},\n append(mma_shape_B, _1{})); // (MMA,MMA_N,MMA_K, _1)\n auto sD = Layout<_1>{}; // XXX Dummy\n\n auto tma_load_a =\n make_tma_copy_A_sm100(SM90_TMA_LOAD{}, tensorA, sA(_, _, _, 0), cluster_tile_mainloop, mma);\n auto tma_load_b =\n make_tma_copy_B_sm100(SM90_TMA_LOAD{}, tensorB, sB(_, _, _, 0), cluster_tile_shape, mma);\n\n // Assert checks on tile sizes -- no predication\n assert(M % size<0>(cluster_tile_shape) == 0);\n assert(N % size<1>(cluster_tile_shape) == 0);\n\n dim3 dimBlock(512);\n dim3 dimCluster(size<0>(cluster_shape), size<1>(cluster_shape), size<2>(cluster_shape));\n dim3 dimGrid(sm_count, 1, 1);\n\n int smem_size = sizeof(\n SharedStorage);\n\n auto *kernel_ptr = &group_row_col_rht_gemm_device_graph_safe<\n decltype(M), decltype(N), decltype(k_tile_size), decltype(cluster_shape),\n decltype(cluster_tile_shape), TA, decltype(dA), decltype(sA), decltype(tma_load_a), TB,\n decltype(dB), decltype(sB), decltype(tma_load_b), TD, decltype(dD), decltype(sD), TSFD,\n decltype(sfd_layout), TQA, decltype(dQA), TSFA, decltype(sfa_layout), decltype(mma),\n AccumulatorPipelineStageCount, SchedulerPipelineStageCount, kEnableStochasticRounding,\n kEnableRHTColQuant, kEnableRowQuant, kEnableSwizzleSFOutput, kUseFastMath>;\n\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(*kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n\n // Set workspace and set to zero\n NVTE_CHECK_CUDA(cudaMemsetAsync(reinterpret_cast(tile_scheduler_workspace), 0,\n sizeof(uint32_t), stream));\n\n // Launch kernel\n cutlass::ClusterLaunchParams params = {dimGrid, dimBlock, dimCluster, smem_size, stream};\n cutlass::Status status = cutlass::launch_kernel_on_cluster(\n params, (void const *)kernel_ptr, M, N, k_tile_size, cluster_shape, cluster_tile_shape, A, dA,\n sA, tma_load_a, B, dB, sB, tma_load_b, QA, dQA, SFA, sfa_layout, QA_COLWISE, SFA_COLWISE,\n amax_rowwise, amax_colwise, offsets, first_dims, num_tensors, shape_rep,\n tile_scheduler_workspace, mma, rng_state);\n NVTE_CHECK_CUDA(cudaGetLastError());\n NVTE_CHECK(status == cutlass::Status::kSuccess, \"Kernel launch failed.\");\n}\n\n} // namespace\n} // namespace detail\n\nvoid group_hadamard_transform_cast_fusion_graph_safe(const GroupedTensor *input,\n GroupedTensor *output,\n const Tensor &hadamard_matrix_,\n QuantizationConfig &quant_config,\n Tensor &quant_workspace, cudaStream_t stream) {\n NVTE_API_CALL(group_hadamard_transform_cast_fusion_graph_safe);\n\n using transformer_engine::detail::kMaxTensorsPerKernel;\n using transformer_engine::detail::ShapeRepresentation;\n\n void *input_base_ptr = reinterpret_cast(input->data.dptr);\n // TODO(zhongbo): add input sanity checks here\n\n bool all_has_row_quant = output->has_data();\n bool all_has_col_quant = output->has_columnwise_data();\n\n // Stochastic rounding config\n const bool use_stochastic_rounding = quant_config.stochastic_rounding;\n const size_t *rng_state = nullptr;\n if (use_stochastic_rounding) {\n NVTE_CHECK(quant_config.rng_state != nullptr,\n \"Enabled stochastic rounding without providing RNG state\");\n const Tensor &rng_state_tensor = *convertNVTETensorCheck(quant_config.rng_state);\n NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64,\n \"RNG state should contain 2 64-bit values.\");\n NVTE_CHECK(rng_state_tensor.data.shape == std::vector{2},\n \"Shape of the RNG state should be [2], but got \", rng_state_tensor.data.shape);\n rng_state = reinterpret_cast(rng_state_tensor.data.dptr);\n }\n\n uint32_t *tile_scheduler_workspace = nullptr;\n NVTE_CHECK(quant_workspace.data.dptr != nullptr, \"Quantization workspace must be provided.\");\n NVTE_CHECK(quant_workspace.data.buffer_size_bytes() >= sizeof(uint32_t),\n \"Quantization workspace must be at least 4 bytes.\");\n tile_scheduler_workspace = reinterpret_cast(quant_workspace.data.dptr);\n\n // Template arguments\n using TA = cute::bfloat16_t;\n using TB = cute::bfloat16_t;\n using TD = cutlass::float_e2m1_t;\n using TSFD = cutlass::float_ue4m3_t;\n using TQA = TD;\n using TSFA = TSFD;\n\n checkCuDriverContext(stream);\n\n // Check Hadamard matrix\n constexpr int kHadamardDimension = 16;\n\n NVTE_CHECK(hadamard_matrix_.dtype() == transformer_engine::DType::kBFloat16,\n \"Hadamard matrix must be BF16 tensor, but dtype is \",\n to_string(hadamard_matrix_.dtype()), \".\");\n const SimpleTensor &hadamard_matrix = hadamard_matrix_.data;\n NVTE_CHECK(\n (hadamard_matrix_.shape() == std::vector{kHadamardDimension, kHadamardDimension}),\n \"Hadamard matrix must have shape=\",\n std::vector{kHadamardDimension, kHadamardDimension},\n \", but got shape=\", hadamard_matrix_.shape(), \".\");\n const size_t hadamard_dimension = hadamard_matrix.shape[0];\n\n const size_t num_tensors = input->num_tensors;\n const size_t first_logical_dim = input->logical_shape.data[0];\n const size_t last_logical_dim = input->logical_shape.data[1];\n // const size_t elts_total = first_logical_dim * last_logical_dim;\n NVTE_CHECK(first_logical_dim % 128 == 0,\n \"First dimension of a grouped tensor should be divisible by 128.\");\n NVTE_CHECK(last_logical_dim % 128 == 0,\n \"Last dimension of a grouped tensor should be divisible by 128.\");\n NVTE_CHECK(num_tensors <= kMaxTensorsPerKernel,\n \"Number of tensors should be less than or equal to \", kMaxTensorsPerKernel);\n\n ShapeRepresentation shape_rep = ShapeRepresentation::VARYING_FIRST_DIM;\n if (output->all_same_shape()) {\n shape_rep = ShapeRepresentation::SAME_BOTH_DIMS;\n } else if (output->all_same_first_dim()) {\n shape_rep = ShapeRepresentation::VARYING_LAST_DIM;\n } else if (output->all_same_last_dim()) {\n shape_rep = ShapeRepresentation::VARYING_FIRST_DIM;\n } else if (output->varying_both_dims()) {\n shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS;\n }\n\n TQA *const rowwise_data_base_ptr = reinterpret_cast(output->data.dptr);\n TSFA *const rowwise_scale_inv_base_ptr = reinterpret_cast(output->scale_inv.dptr);\n TQA *const colwise_data_base_ptr = reinterpret_cast(output->columnwise_data.dptr);\n TSFA *const colwise_scale_inv_base_ptr =\n reinterpret_cast(output->columnwise_scale_inv.dptr);\n float *const amax_rowwise_base_ptr = reinterpret_cast(output->amax.dptr);\n float *const amax_colwise_base_ptr = reinterpret_cast(output->columnwise_amax.dptr);\n\n const int64_t *const offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr);\n const int64_t *const first_dims_ptr = reinterpret_cast(output->first_dims.dptr);\n\n const bool is_const_last_dim = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS ||\n shape_rep == ShapeRepresentation::VARYING_FIRST_DIM);\n NVTE_CHECK(is_const_last_dim,\n \"Currently we only support const last dimension for graph safe hadamard transform.\");\n\n auto sm_count = transformer_engine::cuda::sm_count();\n\n int k_tile_size = 1024;\n\n const bool use_swizzle_sf_output = output->with_gemm_swizzled_scales;\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_stochastic_rounding, kEnableStochasticRounding,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n all_has_col_quant, kEnableRhtColQuant,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n all_has_row_quant, kEnableRowQuant,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_swizzle_sf_output, kEnableSwizzleSFOutput,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n quant_config.use_fast_math, kUseFastMath,\n\n if constexpr (kEnableRhtColQuant || kEnableRowQuant) {\n detail::group_row_col_rht_gemm_ntt_w_sfc_graph_safe<\n kEnableStochasticRounding, kEnableRhtColQuant, kEnableRowQuant,\n kEnableSwizzleSFOutput, TA, TB, TQA, TSFA, TD, TSFD, kUseFastMath>(\n /*packed_sequence_length=*/first_logical_dim,\n /*hidden_size=*/last_logical_dim,\n /*num_tensors=*/num_tensors,\n /*shape_rep=*/shape_rep,\n /*A=*/reinterpret_cast(input_base_ptr),\n /*B=*/reinterpret_cast(hadamard_matrix.dptr),\n /*QA=*/reinterpret_cast(rowwise_data_base_ptr),\n /*SFA=*/reinterpret_cast(rowwise_scale_inv_base_ptr),\n /*QA_COLWISE=*/reinterpret_cast(colwise_data_base_ptr),\n /*SFA_COLWISE=*/reinterpret_cast(colwise_scale_inv_base_ptr),\n /*amax_rowwise=*/reinterpret_cast(amax_rowwise_base_ptr),\n /*amax_colwise=*/reinterpret_cast(amax_colwise_base_ptr),\n /*offsets=*/offsets_ptr,\n /*first_dims=*/first_dims_ptr,\n /*rng_state=*/rng_state,\n /*tile_scheduler_workspace=*/tile_scheduler_workspace,\n /*sm_count=*/sm_count,\n /*stream=*/stream, /*k_tile_size=*/k_tile_size);\n } else {\n NVTE_ERROR(\"Invalid kernel configuration (kEnableRHTColQuant=\",\n kEnableRhtColQuant, \", kEnableRowQuant=\", kEnableRowQuant, \").\");\n }\n\n );););););\n}\n\n} // namespace transformer_engine\n\nvoid nvte_group_hadamard_transform_cast_fusion_graph_safe(\n const NVTEGroupedTensor input, NVTEGroupedTensor output, const NVTETensor hadamard_matrix,\n const NVTEQuantizationConfig quant_config, NVTETensor quant_workspace, cudaStream_t stream) {\n NVTE_API_CALL(nvte_group_hadamard_transform_cast_fusion_graph_safe);\n using namespace transformer_engine;\n\n GroupedTensor *input_tensor = convertNVTEGroupedTensorCheck(input);\n GroupedTensor *output_tensor = convertNVTEGroupedTensorCheck(output);\n\n Tensor *quant_workspace_tensor = convertNVTETensorCheck(quant_workspace);\n\n QuantizationConfig quant_config_cpp;\n if (quant_config != nullptr) {\n quant_config_cpp = *reinterpret_cast(quant_config);\n }\n\n if (input_tensor->num_tensors == 0) {\n return;\n }\n\n // Call the multi-tensor Hadamard transform amax implementation.\n group_hadamard_transform_cast_fusion_graph_safe(\n input_tensor, output_tensor, *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp,\n *quant_workspace_tensor, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/util/hip/hip_runtime.h\"\n#include \"common/util/curanddx.hpp\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"customized_pipeline.cuh\"\n#include \"cutlass/arch/barrier.h\"\n#include \"cutlass/arch/reg_reconfig.h\"\n#include \"cutlass/cluster_launch.hpp\"\n#include \"cutlass/cutlass.h\"\n#include \"cutlass/detail/sm100_blockscaled_layout.hpp\"\n#include \"cutlass/fast_math.h\"\n#include \"cutlass/float8.h\"\n#include \"cutlass/float_subbyte.h\"\n#include \"cutlass/gemm/collective/builders/sm100_common.inl\"\n#include \"cutlass/numeric_conversion.h\"\n#include \"cutlass/numeric_types.h\"\n#include \"cutlass/pipeline/pipeline.hpp\"\n#include \"cutlass/platform/platform.h\"\n#include \"cutlass/util/GPU_Clock.hpp\"\n#include \"cutlass/util/command_line.h\"\n#include \"cutlass/util/print_error.hpp\"\n\nnamespace transformer_engine {\nnamespace detail {\nnamespace {\n\nusing namespace cute;\n\n// Ensure Tensor refers to cute::Tensor, not transformer_engine::Tensor\nusing cute::Tensor;\n\nconstexpr int kMaxTensorsPerKernel = 64;\nconstexpr int kNVFP4BlockSize = 16;\n\nenum ShapeRepresentation {\n SAME_BOTH_DIMS = 0,\n VARYING_FIRST_DIM = 1,\n VARYING_LAST_DIM = 2,\n VARYING_BOTH_DIMS = 3\n};\n\n__device__ __forceinline__ size_t get_current_tensor_id(\n const ShapeRepresentation shape_rep, const size_t num_tensors, const size_t current_offset,\n const size_t first_logical_dim, const size_t last_logical_dim,\n const int64_t *const __restrict__ offsets_ptr) {\n if (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) {\n const size_t current_row = current_offset / last_logical_dim;\n const size_t rows_per_tensor = first_logical_dim / num_tensors;\n return current_row / rows_per_tensor;\n } else {\n // upper_bound(offsets, current_offset) - 1 in range i in [0..num_tensors)\n size_t low = 0;\n size_t hi = num_tensors; // half-open [low, hi)\n\n while (low < hi) {\n const size_t mid = low + (hi - low) / 2;\n const size_t mid_offset = static_cast(offsets_ptr[mid]);\n\n if (mid_offset <= current_offset) {\n low = mid + 1;\n } else {\n hi = mid;\n }\n }\n\n // low = first index where offsets[low] > current_offset (or low == num_tensors)\n // id = low - 1, but need to evaluate if current_offset < offsets[0]\n return (low == 0) ? 0 : (low - 1);\n }\n}\n\nCUTLASS_DEVICE\ncutlass::Array StochasticNumericConverterBase(\n cutlass::Array const &input, cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n auto output_ptr = reinterpret_cast(&output);\n constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING;\n if constexpr (has_rs) {\n asm volatile(\n \"{\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\\n\"\n \"}\"\n : \"=h\"(output_ptr[0]), \"=h\"(output_ptr[1])\n : \"f\"(input[0]), \"f\"(input[1]), \"f\"(input[2]), \"f\"(input[3]), \"f\"(input[4]), \"f\"(input[5]),\n \"f\"(input[6]), \"f\"(input[7]), \"r\"(rbits[0]), \"r\"(rbits[1]));\n } else {\n NVTE_DEVICE_ERROR(\n \"FP4 cvt PTX instructions are architecture-specific. \"\n \"Try recompiling with sm_XXXa instead of sm_XXX.\");\n }\n return output;\n}\n\nCUTLASS_DEVICE\ncutlass::Array StochasticNumericConverter(\n cutlass::Array const &input, cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n cutlass::Array *result_ptr =\n reinterpret_cast *>(&output);\n cutlass::Array const *source_ptr =\n reinterpret_cast const *>(&input);\n cutlass::Array const *rbits_ptr =\n reinterpret_cast const *>(&rbits);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 2; i++) {\n result_ptr[i] = StochasticNumericConverterBase(source_ptr[i], rbits_ptr[i]);\n }\n return output;\n}\n\ntemplate \nstruct SharedStorage {\n static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;\n static int constexpr EpilogueUnrollFactor = EpilogueUnrollFactor_;\n using AtomThrShapeMNK = cute::Shape<_1, _1, _1>;\n\n using AccumulatorPipeline =\n cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;\n\n static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline =\n cutlass::detail::CustomizedPipelineTmaUmmaAsync,\n AtomThrShapeMNK>;\n using MainloopPipelineStorage = typename MainloopPipeline::SharedStorage;\n using SchedPipeline = cutlass::PipelineCLCFetchAsync;\n using SchedPipelineStorage = typename SchedPipeline::SharedStorage;\n using SchedThrottlePipeline = cutlass::PipelineAsync;\n using SchedThrottlePipelineStorage = typename SchedThrottlePipeline::SharedStorage;\n\n struct TensorStorage : cute::aligned_struct<128, _1> {\n cute::array_aligned> smem_A;\n cute::array_aligned> smem_B;\n } tensors;\n\n alignas(16) AccumulatorPipelineStorage accumulator;\n alignas(16) MainloopPipelineStorage mainloop;\n alignas(16) cute::uint64_t tma_barrier[1];\n alignas(16) SchedPipelineStorage sched;\n alignas(16) SchedThrottlePipelineStorage sched_throttle;\n alignas(16) int32_t atomic_tile_id[SchedulerPipelineStageCount_];\n alignas(16) float global_a_amax[kMaxTensorsPerKernel];\n alignas(16) float global_d_amax[kMaxTensorsPerKernel];\n uint32_t atomic_tile_counter[SchedulerPipelineStageCount_];\n uint32_t tmem_base_ptr;\n};\n\n// Main RHT GEMM kernel entry -- highly templated for flexible architecture/config support\ntemplate \n__launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_graph_safe(\n MShape M, NShape packed_N, KShape K, ClusterShape cluster_shape, ClusterTileShape cluster_tile,\n TA const *A, AStride dA, ASmemLayout sAlayout, CUTE_GRID_CONSTANT TmaLoadA const tma_load_a,\n TB const *B, BStride dB, BSmemLayout sBlayout, CUTE_GRID_CONSTANT TmaLoadB const tma_load_b,\n TQA *QA, QAStride dQA, TSFA *SFA, TSFALayout sfa_layout, TQA *QA_COLWISE, TSFA *SFA_COLWISE,\n float *amax_rowwise, float *amax_colwise, const int64_t *offsets, const int64_t *first_dims,\n size_t num_tensors, ShapeRepresentation shape_rep, uint32_t *tile_scheduler_workspace,\n TiledMMA mma, const size_t *rng_state) {\n using namespace cute;\n\n // Abort immediately if compilation is not supported\n constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY;\n if constexpr (!is_blackwell_arch) {\n NVTE_DEVICE_ERROR(\"RHT fusion is only supported on Blackwell.\");\n return;\n } else {\n static_assert(kEnableRHTColQuant_ || kEnableRowQuant_,\n \"group_row_col_rht_gemm_device_graph_safe must generate row-wise \"\n \"and/or column-wise output.\");\n#if !defined(CUTLASS_ARCH_CLC_ENABLED)\n CUTLASS_NOT_IMPLEMENTED();\n return;\n#endif\n\n using X = Underscore;\n // Accumulator data type for main computation\n using ElementAccumulator = float;\n static int constexpr K_PIPE_MAX = size<3>(ASmemLayout{});\n using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>;\n static uint32_t constexpr kTmaTransactionBytes = cutlass::bits_to_bytes(\n size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v);\n static constexpr bool kEnableStochasticRounding = kEnableStochasticRounding_;\n static constexpr bool kEnableRHTColQuant = kEnableRHTColQuant_;\n static constexpr bool kEnableRowQuant = kEnableRowQuant_;\n static constexpr bool kEnableSwizzleSFOutput = kEnableSwizzleSFOutput_;\n static constexpr bool kUseFastMath = kUseFastMath_;\n\n // Constant for RHT tensor processing (tile size etc)\n static int constexpr RhtTensorSize = 16;\n\n // Get the total number of tokens to process\n // Note that here M is the hidden size, which is the last logical dimension of the input tensor x\n // The kernel is designed in column major, so M is the hidden size\n size_t sum_token_dims = offsets[num_tensors] / M;\n\n // Transaction bytes for TMA transfer on RHT tensor blocks\n static int constexpr kTmaRhtTensorTransactionBytes =\n cutlass::bits_to_bytes(RhtTensorSize * RhtTensorSize * cute::sizeof_bits_v);\n static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;\n static int constexpr SchedulerPipelineStageCount = SchedulerPipelineStageCount_;\n\n // Mainloop pipeline stage calculation, vectorization parameters for scaling factors\n static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{});\n static int constexpr SFVecSize = 16;\n // Swizzle output layout for scaling factor arrays\n using SwizzledSFALayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFDLayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n\n // Mainloop pipeline types for TMA async execution and epilogue cluster scheduling\n using MainloopPipeline =\n cutlass::detail::CustomizedPipelineTmaUmmaAsync;\n using MainloopPipelineState = typename MainloopPipeline::PipelineState;\n using SchedPipeline = cutlass::PipelineCLCFetchAsync;\n using SchedPipelineState = typename SchedPipeline::PipelineState;\n using SchedThrottlePipeline = cutlass::PipelineAsync;\n using SchedThrottlePipelineState = typename SchedThrottlePipeline::PipelineState;\n\n static_assert(ClusterShape{} == Shape<_1, _1, _1>{}, \"ClusterShape must be Shape<_1,_1,_1>\");\n\n using TmemAllocator = cute::TMEM::Allocator1Sm;\n static int constexpr VectorSize = RhtTensorSize;\n\n // Compile-time safety: static shapes required for shared memory layouts\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n // CUTE_STATIC_ASSERT(is_static::value);\n\n auto cluster_size = size<0>(cluster_shape);\n auto mainloop_tiler = Shape<_128, _16, _128>{};\n auto epilogue_tiler = Shape<_128, _128, _128>{};\n\n static int constexpr EpilogueUnrollFactor = size<2>(epilogue_tiler) / size<2>(cluster_tile);\n\n // Get the appropriate blocks for this Cluster\n dim3 cluster_coord_in_grid = cluster_id_in_grid();\n\n // Total number of k-tiles\n int const K_TILE_MAX = min(packed_N, K) / size<2>(epilogue_tiler);\n\n struct TileScheduler {\n uint32_t tiles_in_m = 0;\n uint32_t tiles_in_n = 0;\n uint32_t linear_idx = 0;\n uint32_t next_linear_idx = 0;\n uint32_t start_idx = 0;\n uint32_t tile_m_idx = 0;\n uint32_t tile_n_idx = 0;\n int k_tile_max = 0;\n uint32_t *atomic_tile_index_;\n uint32_t *smem_tile_counter;\n uint32_t atomic_offset;\n cutlass::FastDivmodU64 divmod_tiles_in_m;\n\n CUTLASS_DEVICE TileScheduler(uint32_t tiles_m, uint32_t tiles_n, int kmax,\n uint32_t *atomic_tile_index, uint32_t *smem_tile_counter)\n : tiles_in_m(tiles_m),\n tiles_in_n(tiles_n),\n linear_idx(blockIdx.x),\n next_linear_idx(blockIdx.x),\n start_idx(blockIdx.x),\n k_tile_max(kmax),\n atomic_tile_index_(atomic_tile_index),\n smem_tile_counter(smem_tile_counter),\n atomic_offset(gridDim.x),\n divmod_tiles_in_m(uint64_t(tiles_m)) {\n update_tile_idx();\n }\n CUTLASS_DEVICE void update_tile_idx() {\n uint64_t q, r;\n divmod_tiles_in_m(q, r, uint64_t(linear_idx));\n tile_m_idx = static_cast(r);\n tile_n_idx = static_cast(q) * uint32_t(k_tile_max);\n }\n CUTLASS_DEVICE uint32_t tile_m() const { return tile_m_idx; }\n CUTLASS_DEVICE uint32_t tile_n_base() const { return tile_n_idx; }\n CUTLASS_DEVICE uint32_t tiles_m() const { return tiles_in_m; }\n\n CUTLASS_DEVICE uint32_t tiles_n() const { return tiles_in_n; }\n\n CUTLASS_DEVICE bool is_valid() const {\n return cute::elem_less(cute::make_coord(tile_m(), tile_n_base()),\n cute::make_coord(tiles_in_m, tiles_in_n));\n }\n\n CUTLASS_DEVICE bool is_first_wave() const { return linear_idx == start_idx; }\n\n CUTLASS_DEVICE uint32_t get_linear_tile_idx() const { return linear_idx; }\n\n // Fetch a new tile_id using atomics.\n CUTLASS_DEVICE uint32_t fetch_tile_id_counter(int pred) {\n uint32_t tile_id_counter = 0;\n asm volatile(\n \"{\\n\\t\"\n \".reg .pred p;\\n\\t\"\n \"setp.eq.u32 p, %2, 1;\\n\\t\"\n \"@p atom.global.add.u32 %0, [%1], 1; \\n\\t\"\n \"}\"\n : \"=r\"(tile_id_counter)\n : \"l\"(atomic_tile_index_), \"r\"(pred));\n\n return tile_id_counter;\n }\n\n CUTLASS_DEVICE auto fetch_next_work(SchedPipeline &sched_pipeline,\n SchedPipelineState sched_pipeline_consumer_state) {\n sched_pipeline.consumer_wait(sched_pipeline_consumer_state);\n next_linear_idx = smem_tile_counter[sched_pipeline_consumer_state.index()];\n cutlass::arch::fence_view_async_shared();\n sched_pipeline.consumer_release(sched_pipeline_consumer_state);\n return;\n }\n\n CUTLASS_DEVICE auto advance_to_next_work(SchedPipeline &sched_pipeline,\n SchedPipelineState sched_pipeline_producer_state) {\n uint32_t mbarrier_addr = sched_pipeline.producer_get_barrier(sched_pipeline_producer_state);\n // Wait for clcID buffer to become empty with a flipped phase\n sched_pipeline.producer_acquire(sched_pipeline_producer_state);\n auto is_leading_thread = cute::elect_one_sync();\n uint32_t tile_id_counter = fetch_tile_id_counter(is_leading_thread) + atomic_offset;\n uint32_t smem_addr =\n cute::cast_smem_ptr_to_uint(&smem_tile_counter[sched_pipeline_producer_state.index()]);\n if (is_leading_thread) {\n cute::store_shared_remote(tile_id_counter, smem_addr, mbarrier_addr, 0);\n }\n\n ++sched_pipeline_producer_state;\n return sched_pipeline_producer_state;\n }\n\n CUTLASS_DEVICE auto update_work_tile_info() {\n linear_idx = next_linear_idx;\n update_tile_idx();\n return;\n }\n };\n\n // Allocate and alias shared memory to the kernel's shared storage type\n extern __shared__ char shared_memory[];\n using SharedStorage =\n SharedStorage;\n SharedStorage &shared_storage = *reinterpret_cast(shared_memory);\n\n // Compute the number of tiles in M and N after tiling and assign scheduler\n uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile))));\n uint32_t tiles_in_n = uint32_t(size(ceil_div(sum_token_dims, size<2>(epilogue_tiler))));\n\n TileScheduler scheduler(tiles_in_m, tiles_in_n, K_TILE_MAX, tile_scheduler_workspace,\n shared_storage.atomic_tile_counter);\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n\n // Shapes for accumulated tiles in mainloop and epilogue\n auto acc_shape_mma = make_shape(take<0, 2>(mainloop_tiler), _1{}, _1{});\n auto acc_shape_epilogue = make_shape(take<0, 2>(epilogue_tiler), _1{}, _1{});\n\n // Shape of the accumulator fragment for the main loop pipeline, with pipeline stages appended\n auto acc_mainloop_pipelined_shape = append(acc_shape_mma, Int{});\n auto bulk_tmem_mma = TiledMMA::make_fragment_C(acc_mainloop_pipelined_shape);\n\n // Number of threads assigned for various epilogue roles depending on quantization settings\n static int constexpr NumEpilogueColQuantThreadCount = kEnableRHTColQuant ? 128 : 0;\n static int constexpr NumEpilogueRowQuantThreadCount = kEnableRowQuant ? 256 : 0;\n static int constexpr NumMmaThreadCount = kEnableRHTColQuant ? 32 : 0;\n static int constexpr NumMmaIssueThreadCount = kEnableRHTColQuant ? 1 : 0;\n static int constexpr NumSchedThreads = 32;\n static int constexpr NumMainloopLoadThreads = 32;\n static int constexpr NumEpilogueThreads =\n NumEpilogueColQuantThreadCount + NumEpilogueRowQuantThreadCount;\n\n TmemAllocator tmem_allocator{};\n cutlass::arch::NamedBarrier tmem_allocation_result_barrier(\n NumMmaThreadCount + NumEpilogueColQuantThreadCount,\n cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);\n\n int warp_idx = cutlass::canonical_warp_idx_sync();\n\n // warp assignment\n bool is_mma_warp = (warp_idx == 0);\n bool is_dma_warp = (warp_idx == 1);\n bool is_sched_warp = (warp_idx == 2);\n bool is_epilogue_col_quant_warp = (warp_idx >= 4 && warp_idx <= 7);\n bool is_epilogue_row_quant_warp = (warp_idx >= 8 && warp_idx <= 15);\n\n typename MainloopPipeline::Params mainloop_pipeline_params;\n if (is_dma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;\n }\n if (is_mma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;\n }\n mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp;\n mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes;\n mainloop_pipeline_params.initializing_warp = 0;\n mainloop_pipeline_params.num_consumers =\n NumEpilogueRowQuantThreadCount + NumMmaIssueThreadCount;\n\n MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params,\n cluster_shape, cute::true_type{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n\n MainloopPipelineState mainloop_pipe_consumer_state;\n MainloopPipelineState mainloop_pipe_producer_state =\n cutlass::make_producer_start_state();\n\n using AccumulatorPipeline =\n cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState;\n using AccumulatorPipelineInitBarriers = cute::bool_constant;\n\n AccumulatorPipelineState accumulator_pipe_consumer_state;\n AccumulatorPipelineState accumulator_pipe_producer_state =\n cutlass::make_producer_start_state();\n\n typename AccumulatorPipeline::Params accumulator_pipeline_params;\n if (is_mma_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer;\n }\n if (is_epilogue_col_quant_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer;\n }\n // Only one producer thread arrives on this barrier.\n accumulator_pipeline_params.producer_arv_count = 1;\n accumulator_pipeline_params.consumer_arv_count =\n size(AtomThrShapeMNK{}) * NumEpilogueColQuantThreadCount;\n accumulator_pipeline_params.initializing_warp = 1;\n AccumulatorPipeline accumulator_pipeline(\n shared_storage.accumulator, accumulator_pipeline_params, cluster_shape,\n AccumulatorPipelineInitBarriers{}, cute::true_type{}); // Delay mask calculation\n typename SchedPipeline::Params sched_pipeline_params;\n if (is_sched_warp) {\n sched_pipeline_params.role = SchedPipeline::ThreadCategory::ProducerConsumer;\n } else {\n sched_pipeline_params.role = SchedPipeline::ThreadCategory::Consumer;\n }\n sched_pipeline_params.producer_blockid = 0;\n sched_pipeline_params.producer_arv_count = 1;\n sched_pipeline_params.consumer_arv_count =\n NumSchedThreads +\n cluster_size * (NumMainloopLoadThreads + NumEpilogueThreads + NumMmaThreadCount);\n sched_pipeline_params.transaction_bytes = sizeof(uint32_t);\n sched_pipeline_params.initializing_warp = 3;\n SchedPipeline sched_pipeline(shared_storage.sched, sched_pipeline_params, cluster_shape);\n SchedPipelineState sched_pipeline_consumer_state;\n SchedPipelineState sched_pipeline_producer_state =\n cutlass::make_producer_start_state();\n\n typename SchedThrottlePipeline::Params sched_throttle_pipeline_params;\n if (is_dma_warp) {\n sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Producer;\n }\n if (is_sched_warp) {\n sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Consumer;\n }\n sched_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads;\n sched_throttle_pipeline_params.consumer_arv_count = NumSchedThreads;\n sched_throttle_pipeline_params.dst_blockid = 0;\n sched_throttle_pipeline_params.initializing_warp = 4;\n\n SchedThrottlePipeline sched_throttle_pipeline(shared_storage.sched_throttle,\n sched_throttle_pipeline_params);\n SchedThrottlePipelineState sched_pipeline_throttle_consumer_state;\n SchedThrottlePipelineState sched_pipeline_throttle_producer_state =\n cutlass::make_producer_start_state();\n\n if (warp_idx == 2 && elect_one_sync()) {\n cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);\n }\n __syncthreads();\n\n // Warp group roles: DMA (global->shared copy), MMA (tensor core gemm), scheduler, column quantizer, row quantizer\n if (is_dma_warp) {\n // Warp responsible for loading input from global to shared memory using TMA (Tensor Memory Access).\n cutlass::arch::warpgroup_reg_dealloc<32>();\n // Get TMA tensors for input matrix A and B (Hadamard/transform matrix) from global memory.\n Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, packed_N));\n Tensor mB = tma_load_b.get_tma_tensor(make_shape(RhtTensorSize, RhtTensorSize));\n\n // Partition tensors for tiling according to the mainloop and cluster tilers.\n Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{});\n Tensor gB_nk =\n local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k)\n\n // Shared memory tensors for pipeline\n Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()),\n sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()),\n sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n // Determine warp/tile positioning\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n // Partition global to local fragments for A and B\n Tensor tCgA = thr_mma.partition_A(gA_mk); // (MMA,MMA_M,MMA_K,k)\n Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k)\n\n Layout cta_layout_mnk = make_layout(cluster_shape);\n Layout cta_layout_vmnk =\n tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{}));\n auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster);\n\n auto [tAgA, tAsA] =\n tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),\n group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA));\n\n auto [tBgB, tBsB] =\n tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),\n group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB));\n\n uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);\n uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);\n if constexpr (kEnableRHTColQuant) {\n if (elect_one_sync()) {\n cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0],\n kTmaRhtTensorTransactionBytes);\n copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0),\n tBsB(_, 0));\n }\n }\n\n do {\n // is_first_wave indicates whether this scheduler wave is the first among a group.\n bool is_first_wave = scheduler.is_first_wave();\n uint32_t skip_wait = is_first_wave;\n auto tAgA_mk = tAgA(_, scheduler.tile_m(), _);\n int k_tile = 0;\n\n sched_throttle_pipeline.producer_acquire(sched_pipeline_throttle_producer_state);\n sched_throttle_pipeline.producer_commit(sched_pipeline_throttle_producer_state);\n ++sched_pipeline_throttle_producer_state;\n CUTLASS_PRAGMA_NO_UNROLL\n while (k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n()) {\n int k_tile_idx_n = scheduler.tile_n_base() + k_tile;\n ++k_tile;\n skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount);\n mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state);\n using BarrierType = typename MainloopPipeline::ProducerBarrierType;\n BarrierType *tma_barrier =\n mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state);\n int write_stage = mainloop_pipe_producer_state.index();\n ++mainloop_pipe_producer_state;\n if (cute::elect_one_sync()) {\n copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n),\n tAsA(_, write_stage));\n }\n }\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n // scheduler.advance();\n } while (scheduler.is_valid());\n mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);\n } else if (is_mma_warp) {\n // This warp executes the main tensor core matrix-multiply-accumulate for the Hadamard transform.\n cutlass::arch::warpgroup_reg_dealloc<32>();\n if constexpr (kEnableRHTColQuant) {\n // Setup shared memory fragments for A and B tiles.\n Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()),\n sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()),\n sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n // Allocate \"fragments\" -- these are actually umma smem descriptors\n Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)\n Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE)\n\n mma.accumulate_ = UMMA::ScaleOut::Zero;\n\n tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns,\n &shared_storage.tmem_base_ptr);\n __syncwarp();\n tmem_allocation_result_barrier.arrive();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_mma.data() = tmem_base_ptr;\n // Wait until the B (Hadamard) tensor copy is complete\n cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/);\n do {\n uint32_t skip_wait = K_TILE_MAX <= 0;\n\n auto barrier_token =\n mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0;\n k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) {\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n int read_stage = mainloop_pipe_consumer_state.index();\n auto tCrA_mk = tCrA(_, _, _, read_stage);\n auto tCrB_nk = tCrB(_, _, 0, 0);\n CUTLASS_PRAGMA_UNROLL\n for (int k_block = 0; k_block < size<2>(tCrA) / EpilogueUnrollFactor; ++k_block) {\n int accumulator_k_block =\n accumulator_pipe_producer_state.index() * EpilogueUnrollFactor;\n int tCrA_k_block = k_block * EpilogueUnrollFactor;\n accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < EpilogueUnrollFactor; i++) {\n auto accumulators = bulk_tmem_mma(_, _, _, accumulator_k_block + i);\n gemm(mma, tCrA_mk(_, _, tCrA_k_block + i), tCrB_nk, accumulators);\n }\n\n accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);\n ++accumulator_pipe_producer_state;\n }\n auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n skip_wait = k_tile >= K_TILE_MAX;\n mainloop_pipeline.umma_consumer_release(curr_mainloop_pipe_consumer_state);\n barrier_token =\n mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n }\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n tmem_allocator.release_allocation_lock();\n accumulator_pipeline.producer_tail(accumulator_pipe_producer_state);\n tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);\n }\n } else if (is_sched_warp) {\n // Scheduler warp manages tile assignment and pipeline progress for warps\n cutlass::arch::warpgroup_reg_dealloc<32>();\n do {\n sched_throttle_pipeline.consumer_wait(sched_pipeline_throttle_consumer_state);\n sched_throttle_pipeline.consumer_release(sched_pipeline_throttle_consumer_state);\n ++sched_pipeline_throttle_consumer_state;\n sched_pipeline_producer_state =\n scheduler.advance_to_next_work(sched_pipeline, sched_pipeline_producer_state);\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n } else if (is_epilogue_col_quant_warp) {\n // Warp responsible for quantizing output of Hadamard transform to FP4 for columnwise usage,\n // and writing result tensors/scales to global memory.\n cutlass::arch::warpgroup_reg_alloc<192>();\n if constexpr (kEnableRHTColQuant) {\n using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x;\n\n auto acc_epilogue_pipelined_shape =\n append(acc_shape_epilogue, Int{});\n auto bulk_tmem_epilogue_layout = make_layout(\n acc_epilogue_pipelined_shape,\n make_stride(stride<0>(bulk_tmem_mma), Int<0>{}, Int<0>{}, size<1>(epilogue_tiler)));\n auto bulk_tmem_epilogue = make_tensor(make_tmem_ptr(), bulk_tmem_epilogue_layout);\n\n // Use 256-bit fragments for aligned bulk stores\n static int constexpr FragmentSize = 256 / sizeof_bits_v;\n\n // Wait for TMEM allocation for this pipeline to finish\n tmem_allocation_result_barrier.arrive_and_wait();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_epilogue.data() = tmem_base_ptr;\n int global_thread_idx = threadIdx.x;\n int local_thread_idx = global_thread_idx % cutlass::NumThreadsPerWarpGroup;\n // g2s load all global_d_amax\n CUTLASS_PRAGMA_NO_UNROLL\n for (int g = local_thread_idx; g < num_tensors; g += NumEpilogueColQuantThreadCount) {\n shared_storage.global_d_amax[g] = __ldg(reinterpret_cast(amax_colwise + g));\n }\n\n size_t rng_seed = 0;\n size_t rng_offset = 0;\n // Setup RNG for stochastic rounding\n if constexpr (kEnableStochasticRounding) {\n rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0;\n rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0;\n }\n // TODO(zhongbo): double check the logic here\n int group_idx = get_current_tensor_id(\n shape_rep, num_tensors, (scheduler.tile_n_base() * size<1>(epilogue_tiler)) * M,\n packed_N, M, offsets);\n\n // Determine quantization scale factor layouts/output splits for this group\n TSFDLayout sfd_layout;\n int cur_N = static_cast(first_dims[group_idx]);\n if constexpr (kEnableSwizzleSFOutput) {\n sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{});\n } else {\n sfd_layout = make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)),\n make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{})));\n }\n // Build output tensors for columns and their quant scales\n // TODO(zhongbo): double check the logic here\n Tensor mD = make_tensor(cute::subbyte_iterator(reinterpret_cast(\n reinterpret_cast(QA_COLWISE) + offsets[group_idx] / 2)),\n make_shape(M, cur_N), DStride{}); // (M,packed_N)\n Tensor gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n\n // for every tensor [x, y] row major, x y both a multiple of 128\n // both of its rowwise and colwise scaling factors will have exactly x * y / 16 elements in FP8 E4M3\n Tensor mSFD = make_tensor(\n make_gmem_ptr(reinterpret_cast(reinterpret_cast(SFA_COLWISE) +\n offsets[group_idx] / kNVFP4BlockSize)),\n sfd_layout);\n Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n\n Tensor gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler));\n\n // Setup tile-level TMEM (t2r) and global memory (r2g) copy descriptors\n auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{}));\n auto tiled_r2g =\n make_tiled_copy_D(Copy_Atom{}, tiled_t2r);\n auto thr_t2r = tiled_t2r.get_slice(local_thread_idx);\n auto thr_r2g = tiled_r2g.get_slice(local_thread_idx);\n\n cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount,\n cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);\n // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release}\n static constexpr float fp4_max = 6.0f;\n static constexpr float fp8_max = 448.0f;\n static constexpr float fp4_max_inv = 1.0f / fp4_max;\n float c_global_amax_val = shared_storage.global_d_amax[group_idx];\n float global_encode_scale = c_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / c_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n float global_decode_scale = 1.0f / global_encode_scale;\n\n // Scaling factor for fast math path\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n\n do {\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0;\n k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();\n ++k_tile) {\n int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler);\n\n // TODO(zhongbo): double check the logic here\n int cur_group_idx = get_current_tensor_id(\n shape_rep, num_tensors, global_tile_n_offset * M, packed_N, M, offsets);\n\n if (cur_group_idx != group_idx) {\n group_idx = cur_group_idx;\n c_global_amax_val = shared_storage.global_d_amax[group_idx];\n // update amax\n global_encode_scale = c_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / c_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n global_decode_scale = 1.0f / global_encode_scale;\n global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n // TODO(zhongbo): double check the logic here\n cur_N = first_dims[group_idx];\n if constexpr (kEnableSwizzleSFOutput) {\n sfd_layout =\n tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{});\n } else {\n sfd_layout =\n make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)),\n make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{})));\n }\n // update tensor\n mD = make_tensor(cute::subbyte_iterator(reinterpret_cast(\n reinterpret_cast(QA_COLWISE) + offsets[group_idx] / 2)),\n make_shape(M, cur_N), DStride{});\n gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n mSFD = make_tensor(make_gmem_ptr(reinterpret_cast(\n reinterpret_cast(SFA_COLWISE) +\n offsets[group_idx] / kNVFP4BlockSize)),\n sfd_layout);\n gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n\n gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler));\n }\n int group_start_offset = offsets[group_idx] / M;\n int local_tile_n_idx =\n (global_tile_n_offset - group_start_offset) / size<1>(epilogue_tiler);\n Tensor tDgD_mn = gD_mn_view(_, _, _, scheduler.tile_m(), local_tile_n_idx);\n\n Tensor tDgSFD_mn = gSFD_mn(_, _, scheduler.tile_m(), local_tile_n_idx);\n accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state);\n\n auto Acc = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index());\n Tensor tDtAcc = thr_t2r.partition_S(Acc); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDgD = thr_t2r.partition_D(tDgD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n\n Tensor tTR_rAcc = make_tensor(\n shape(tDgD)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDrD = make_tensor(shape(tDgD));\n Tensor tTR_rAcc_frag =\n recast>(coalesce(tTR_rAcc));\n Tensor tDrD_frag = recast>(coalesce(tDrD));\n\n Tensor src = thr_r2g.retile_S(tDrD);\n Tensor dst = thr_r2g.retile_D(tDgD);\n\n Tensor tDgSFD_view = make_tensor(\n tDgSFD_mn.data(), make_layout(make_shape(shape(tDgSFD_mn), Int<1>{}, Int<1>{}),\n make_stride(stride(tDgSFD_mn), Int<0>{}, Int<0>{})));\n Tensor tDgSFD = filter(thr_t2r.partition_D(tDgSFD_view));\n Tensor tDrSFD = make_tensor(shape(tDgSFD));\n\n static int constexpr NumVecs = size(tDgD) / VectorSize;\n Tensor tD_rRowSFD_frg = recast>(tDrSFD);\n\n // Compute amax and quantization scales for this tile\n cutlass::maximum_absolute_value_reduction<\n cutlass::Array, true>\n amax_reduction;\n cutlass::Array vec_maxs;\n cutlass::Array pvscales;\n // Copy from TMEM to registers\n copy(tiled_t2r, tDtAcc, tTR_rAcc);\n cutlass::arch::fence_view_async_tmem_load();\n accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state);\n ++accumulator_pipe_consumer_state;\n\n if constexpr (!kUseFastMath) {\n // Downcast to BF16 for bit-wise compatibility with\n // unfused kernels\n auto convert_accum_to_bf16 =\n cutlass::NumericArrayConverter{};\n auto convert_bf16_to_accum =\n cutlass::NumericArrayConverter{};\n tTR_rAcc_frag(_0{}) =\n convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{})));\n tTR_rAcc_frag(_1{}) =\n convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_1{})));\n }\n\n auto compute_frgs = reinterpret_cast *>(\n tTR_rAcc_frag.data());\n auto output_frgs = reinterpret_cast *>(tDrD_frag.data());\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]);\n }\n\n pvscales = cutlass::multiplies>{}(\n vec_maxs, global_encode_scale_multiplier);\n auto pvscales_cvted =\n cutlass::NumericArrayConverter{}(pvscales);\n\n tD_rRowSFD_frg(_0{}) = pvscales_cvted;\n auto qpvscale_ups = cutlass::NumericArrayConverter{}(\n tD_rRowSFD_frg(_0{}));\n auto qpvscale_scaled =\n cutlass::multiplies>{}(\n qpvscale_ups, global_decode_scale);\n cutlass::Array acc_scales;\n if constexpr (kUseFastMath) {\n // Fast math: compute approximate reciprocal\n acc_scales =\n cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled);\n } else {\n // Accurate math: compute reciprocal with division\n acc_scales = cutlass::divides>{}(\n 1.0, qpvscale_scaled);\n }\n\n // Prepare stochastic rounding random state if enabled\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n transformer_engine::curanddx::detail::philox4x32_native_state<\n NVTE_BUILD_NUM_PHILOX_ROUNDS>\n rng;\n // \"Prefetch\" a stochastic rounding state for the first tile\n if constexpr (kEnableStochasticRounding) {\n const size_t rng_sequence = global_thread_idx + k_tile * 512 +\n scheduler.get_linear_tile_idx() * K_TILE_MAX * 512;\n rng.init(rng_seed, rng_sequence, rng_offset);\n }\n CUTLASS_PRAGMA_UNROLL\n // Apply round/quantize to each fragment, with or without stochastic rounding\n for (int v = 0; v < NumVecs; v++) {\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(\n acc_scales[v], cutlass::platform::numeric_limits::max());\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(\n cutlass::multiplies>{}(\n compute_frgs[v], acc_scale),\n *reinterpret_cast *>(&random_uint4));\n } else {\n output_frgs[v] =\n cutlass::NumericArrayConverter{}(\n cutlass::multiplies>{}(\n compute_frgs[v], acc_scale));\n }\n }\n\n // Write quantized FP4 tile and dequant scale to gmem\n copy(tiled_r2g, src, dst);\n copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFD, tDgSFD);\n }\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n }\n } else if (is_epilogue_row_quant_warp) {\n // Warp responsible for quantizing the input (before Hadamard transform) to FP4 for row-wise usage.\n cutlass::arch::warpgroup_reg_alloc<136>();\n if constexpr (kEnableRowQuant) {\n using S2RVectorType = uint128_t;\n\n int global_thread_idx = threadIdx.x;\n int local_thread_idx = global_thread_idx % 256;\n size_t rng_seed = 0;\n size_t rng_offset = 0;\n // g2s load all global_a_amax for all groups/tensors\n CUTLASS_PRAGMA_NO_UNROLL\n for (int g = local_thread_idx; g < num_tensors; g += NumEpilogueRowQuantThreadCount) {\n shared_storage.global_a_amax[g] = __ldg(reinterpret_cast(amax_rowwise + g));\n }\n // RNG for stochastic rounding\n if constexpr (kEnableStochasticRounding) {\n rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0;\n rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0;\n }\n // Input/output tensors/partitions for row quant warp\n Tensor mQA =\n make_tensor(cute::subbyte_iterator(QA), make_layout(make_shape(M, packed_N), dQA));\n Tensor gQA_mn = local_tile(mQA, epilogue_tiler, make_coord(_, _, _), Step<_1, X, _1>{});\n Tensor mSFA = make_tensor(make_gmem_ptr(SFA), sfa_layout);\n\n Tensor gSFA_mn = local_tile(mSFA, epilogue_tiler, make_coord(_, _, _),\n Step<_1, X, _1>{}); // (BLK_M,BLK_N)\n // Swizzled shared memory A tile, with layout\n Tensor sA = as_position_independent_swizzle_tensor(group_modes<0, 2>(\n coalesce(make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()),\n sAlayout)))); // (BLOCK_M, BLOCK_M,PIPE)\n\n // Set up layouts for partitioning \u2013 tile-by-warp, with vector granularity\n using S2RWarpLayout = Layout>;\n using WarpGroupLayout = Layout>;\n using S2RThreadLayout = decltype(blocked_product(S2RWarpLayout{}, WarpGroupLayout{}));\n using S2RValLayout = Layout, _1>>;\n using S2RAtomA = Copy_Atom;\n using R2GAtomQA = Copy_Atom;\n using R2GAtomSFA = Copy_Atom;\n auto tiled_s2r = make_tiled_copy(S2RAtomA{}, S2RThreadLayout{}, S2RValLayout{});\n auto tiled_r2g_QA = make_tiled_copy(R2GAtomQA{}, S2RThreadLayout{}, S2RValLayout{});\n auto tiled_r2g_SFA = make_tiled_copy(R2GAtomSFA{}, S2RThreadLayout{}, S2RValLayout{});\n\n auto thr_s2r = tiled_s2r.get_slice(local_thread_idx);\n auto thr_r2g_QA = tiled_r2g_QA.get_slice(local_thread_idx);\n auto thr_r2g_SFA = tiled_r2g_SFA.get_slice(local_thread_idx);\n Tensor tQAsA = thr_s2r.partition_S(sA); // (Copy, Copy_M, Copy_N, PIPE)\n\n // Allocate temporary register tensors for copying quantization => output\n Tensor tQArA = make_tensor_like(\n make_layout(tQAsA(_, _, _, _0{}).shape())); // (Copy, Copy_M, Copy_N)\n Tensor tQAgQA = thr_r2g_QA.partition_S(gQA_mn);\n Tensor tQArQA = make_tensor_like(tQAgQA(_, _, _, _0{}, _0{}));\n\n Tensor tQAgSFA = thr_r2g_SFA.partition_S(gSFA_mn);\n Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{}));\n\n // Will result in barrier_id=10 passed to bar.sync instr as cutlass adds 8\n // in order to go over the reserved named barrier count.\n constexpr int row_quant_barrier_id = 2;\n cutlass::arch::NamedBarrier::sync(NumEpilogueRowQuantThreadCount, row_quant_barrier_id);\n\n int group_idx = get_current_tensor_id(\n shape_rep, num_tensors, (scheduler.tile_n_base() * size<1>(epilogue_tiler)) * M,\n packed_N, M, offsets);\n float a_global_amax_val = shared_storage.global_a_amax[group_idx];\n // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release}\n static constexpr float fp4_max = 6.0f;\n static constexpr float fp8_max = 448.0f;\n static constexpr float fp4_max_inv = 1.0f / fp4_max;\n float global_encode_scale = a_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / a_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n\n float global_decode_scale = 1.0f / global_encode_scale;\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n auto sfa_converter = cutlass::NumericConverter{};\n do {\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0;\n k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) {\n int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler);\n\n int cur_group_idx = get_current_tensor_id(\n shape_rep, num_tensors, global_tile_n_offset * M, packed_N, M, offsets);\n if (cur_group_idx != group_idx) {\n group_idx = cur_group_idx;\n a_global_amax_val = shared_storage.global_a_amax[group_idx];\n // Update group quantization parameters/scaling\n global_encode_scale = a_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / a_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n global_decode_scale = 1.0f / global_encode_scale;\n global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n }\n\n auto tQAgSFA_mn =\n tQAgSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto tQAgQA_mn = tQAgQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state);\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n copy(tiled_s2r, tQAsA(_, _, _, mainloop_pipe_consumer_state.index()), tQArA);\n cutlass::arch::fence_view_async_shared();\n mainloop_pipeline.consumer_release(mainloop_pipe_consumer_state);\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n\n // static int constexpr NumVecs = size(tQArA) / VectorSize;\n cutlass::maximum_absolute_value_reduction<\n cutlass::Array, true>\n amax_reduction;\n auto compute_frgs = reinterpret_cast *>(tQArA.data());\n auto output_frgs = reinterpret_cast *>(\n raw_pointer_cast(tQArQA.data()));\n Tensor amax =\n make_tensor(prepend(take<1, rank(tQArA)>(tQArA.shape()), _1{}));\n Tensor pvscales = make_tensor_like(amax);\n transformer_engine::curanddx::detail::philox4x32_native_state<\n NVTE_BUILD_NUM_PHILOX_ROUNDS>\n rng;\n if constexpr (kEnableStochasticRounding) {\n const size_t rng_sequence = global_thread_idx + k_tile * 512 +\n scheduler.get_linear_tile_idx() * K_TILE_MAX * 512 +\n tiles_in_m * tiles_in_n * K_TILE_MAX * 512;\n rng.init(rng_seed, rng_sequence, rng_offset);\n }\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < size<1>(group_modes<1, rank(tQArA)>(tQArA)); v++) {\n auto amax_view = group_modes<1, rank(amax)>(amax);\n auto pvscales_view = group_modes<1, rank(pvscales)>(pvscales);\n auto compute_frgs_up =\n cutlass::NumericArrayConverter{}(\n compute_frgs[v]);\n amax_view(_0{}, v) = amax_reduction(ElementAccumulator(0), compute_frgs_up);\n pvscales_view(_0{}, v) = cutlass::multiplies{}(\n amax_view(_0{}, v), global_encode_scale_multiplier);\n filter(tQArSFA)(v) = sfa_converter(pvscales_view(_0{}, v));\n auto qpvscale_ups =\n cutlass::NumericConverter{}(filter(tQArSFA)(v));\n auto qpvscale_scaled =\n cutlass::multiplies{}(qpvscale_ups, global_decode_scale);\n ElementAccumulator acc_scales;\n if constexpr (kUseFastMath) {\n // Fast math: compute approximate reciprocal\n acc_scales = cutlass::reciprocal_approximate_ftz{}(\n qpvscale_scaled);\n } else {\n // Accurate math: compute reciprocal with division\n acc_scales = cutlass::divides{}(1.0, qpvscale_scaled);\n }\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(\n acc_scales, cutlass::platform::numeric_limits::max());\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(\n cutlass::multiplies>{}(\n compute_frgs_up, acc_scale),\n *reinterpret_cast *>(&random_uint4));\n } else {\n output_frgs[v] =\n cutlass::NumericArrayConverter{}(\n cutlass::multiplies>{}(\n compute_frgs_up, acc_scale));\n }\n }\n copy(tiled_r2g_QA, tQArQA, tQAgQA_mn);\n copy(tiled_r2g_SFA, filter(tQArSFA), filter(tQAgSFA_mn));\n }\n // scheduler.advance();\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n }\n\n } else {\n cutlass::arch::warpgroup_reg_dealloc<32>();\n }\n }\n} // NOLINT(readability/fn_size)\n\ntemplate \nvoid group_row_col_rht_gemm_ntt_w_sfc_graph_safe(\n int packed_sequence_length, int hidden_size, size_t num_tensors, ShapeRepresentation shape_rep,\n TA const *A, TB const *B, TQA *QA, TSFA *SFA, TQA *QA_COLWISE, TSFA *SFA_COLWISE,\n float *amax_rowwise, float *amax_colwise, const int64_t *offsets, const int64_t *first_dims,\n const size_t *rng_state, uint32_t *tile_scheduler_workspace, uint32_t sm_count,\n hipStream_t stream, int k_tile_size = 1024) {\n using namespace cute;\n static int constexpr SFVecSize = 16;\n static int constexpr RhtTensorSize = 16;\n\n static_assert(RhtTensorSize == 16, \"RhtTensorSize must be 16\");\n using LinearSFALayout = decltype(make_layout(make_shape(make_shape(Int{}, 0), 0),\n make_stride(make_stride(_0{}, _1{}), 0)));\n using LinearSFDLayout = decltype(make_layout(make_shape(0, make_shape(Int{}, 0)),\n make_stride(0, make_stride(_0{}, _1{}))));\n\n using SwizzledSFALayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFDLayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFALayout = decltype(tile_to_shape(\n SwizzledSFALayoutAtom{}, make_shape(hidden_size, packed_sequence_length), Step<_1, _2>{}));\n using SwizzledSFDLayout = decltype(tile_to_shape(\n SwizzledSFDLayoutAtom{}, make_shape(hidden_size, packed_sequence_length), Step<_2, _1>{}));\n\n using SFALayout = cute::conditional_t;\n using SFDLayout = cute::conditional_t;\n SFALayout sfa_layout;\n SFDLayout sfd_layout;\n\n if constexpr (kEnableSwizzleSFOutput) {\n sfa_layout = tile_to_shape(SwizzledSFALayoutAtom{},\n make_shape(hidden_size, packed_sequence_length), Step<_1, _2>{});\n sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{},\n make_shape(hidden_size, packed_sequence_length), Step<_2, _1>{});\n } else {\n sfa_layout = make_layout(\n make_shape(make_shape(Int{}, hidden_size / SFVecSize), packed_sequence_length),\n make_stride(make_stride(_0{}, _1{}), hidden_size / SFVecSize));\n sfd_layout = make_layout(\n make_shape(hidden_size, make_shape(Int{}, packed_sequence_length / SFVecSize)),\n make_stride(packed_sequence_length / SFVecSize, make_stride(_0{}, _1{})));\n }\n\n // Define shapes (dynamic)\n auto M = hidden_size;\n auto N = packed_sequence_length;\n Tensor tensorA = make_tensor(A, make_shape(hidden_size, packed_sequence_length), LayoutLeft{});\n Tensor tensorB = make_tensor(B, make_shape(RhtTensorSize, RhtTensorSize), LayoutLeft{});\n Tensor tensorQA = make_tensor(QA, make_shape(hidden_size, packed_sequence_length), LayoutLeft{});\n Tensor tensorSFA = make_tensor(SFA, sfa_layout);\n\n // Define strides (from tensors)\n auto dA = stride(tensorA); // (dM,dK)\n auto dB = stride(tensorB); // (dN,dK)\n auto dD = LayoutRight{}; // (dM,dN)\n auto dQA = stride(tensorQA); // (dM,dK)\n using ClusterShape = Shape<_1, _1, _1>;\n auto cluster_shape = ClusterShape{};\n auto cluster_tile_shape = Shape<_128, Int, Int>{};\n auto cluster_tile_mainloop = Shape<_128, Int, _128>{};\n\n // Each mainloop / epilogue loads 128 x 64 tiles while each MMA proceeds with 128 x 16 tiles\n static int constexpr EpilogueUnrollFactor =\n size<2>(cluster_tile_mainloop) / size<2>(cluster_tile_shape);\n // Construct the MMA\n auto mma = make_tiled_mma(\n SM100_MMA_F16BF16_SS(cluster_tile_shape), size<1>(cluster_tile_shape),\n UMMA::Major::MN, UMMA::Major::MN>{},\n Layout>{});\n\n // Assert that the TiledMMA uses all CTAs in the CGA.\n CUTE_STATIC_ASSERT_V(size(cluster_shape) == size(mma));\n CUTE_STATIC_ASSERT_V(evenly_divides(cluster_tile_shape, tile_shape(mma)));\n\n // Determine the A and B shapes\n auto mma_shape_B =\n partition_shape_B(mma, make_shape(size<1>(cluster_tile_shape), size<2>(cluster_tile_shape)));\n\n using TiledMma = decltype(mma);\n using AtomThrID = typename TiledMma::AtomThrID;\n\n using SmemShape_M = decltype(shape_div(\n shape<0>(cluster_tile_shape),\n shape_div(shape<0>(cluster_tile_shape), size<0>(cluster_tile_shape) / size(AtomThrID{}))));\n using SmemShape_N = decltype(shape_div(\n shape<1>(cluster_tile_shape),\n shape_div(shape<1>(cluster_tile_shape), size<1>(cluster_tile_shape) / size(AtomThrID{}))));\n using SmemShape_K = decltype(cute::get<2>(cluster_tile_shape));\n\n using SmemLayoutAtomB =\n decltype(cutlass::gemm::collective::detail::sm100_smem_selector());\n\n auto mma_shape_A = partition_shape_A(\n mma, make_shape(size<0>(cluster_tile_mainloop), size<2>(cluster_tile_mainloop)));\n using SmemShape_M_A =\n decltype(shape_div(shape<0>(cluster_tile_mainloop),\n shape_div(shape<0>(cluster_tile_mainloop),\n size<0>(cluster_tile_mainloop) / size(AtomThrID{}))));\n using SmemShape_K_A = decltype(cute::get<2>(cluster_tile_mainloop));\n using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<\n cute::UMMA::Major::MN, TA, SmemShape_M_A, SmemShape_K_A>());\n\n static uint32_t constexpr TotalTmemRows = 128;\n static uint32_t constexpr Sm100TmemCapacityColumns = 512;\n static uint32_t constexpr TotalTmem = TotalTmemRows * Sm100TmemCapacityColumns;\n static uint32_t constexpr AccumulatorPipelineStageCount =\n TotalTmem / (cute::size<0>(cluster_tile_shape) * cute::size<1>(cluster_tile_shape));\n\n // Define the smem layouts (static)\n // Calculate max pipeline stages based on Blackwell SM100's 232KB shared memory\n constexpr int SchedulerPipelineStageCount = 4;\n static int constexpr MainloopPipelineBytes = sizeof(\n typename cutlass::detail::CustomizedPipelineTmaUmmaAsync<1, Shape<_1, _1, _1>,\n Shape<_1, _1, _1>>::SharedStorage);\n\n static int constexpr SchedulerWorkspaceBytes = sizeof(int) * SchedulerPipelineStageCount;\n static int constexpr SchedulerThrottlePipelineBytes =\n sizeof(typename cutlass::PipelineAsync::SharedStorage);\n static int constexpr SchedulerPipelineBytes =\n sizeof(typename cutlass::PipelineCLCFetchAsync::SharedStorage);\n\n static int constexpr TmemDeallocBytes = sizeof(cutlass::arch::ClusterBarrier);\n static int constexpr BTensorBytes = cute::size(mma_shape_B) * sizeof(TB);\n static int constexpr AccPipelineBytes = sizeof(\n typename cutlass::PipelineUmmaAsync>::SharedStorage);\n static int constexpr TmemBasePtrsBytes = sizeof(uint32_t);\n static int constexpr kBlackwellSmemSize = 232448; // 232KB in bytes\n static int constexpr kBytesPerStage =\n cute::size(mma_shape_A) * sizeof(TA) + MainloopPipelineBytes;\n static int constexpr kReservedBytes = SchedulerWorkspaceBytes + SchedulerThrottlePipelineBytes +\n SchedulerPipelineBytes + TmemBasePtrsBytes +\n TmemDeallocBytes + BTensorBytes +\n AccPipelineBytes; // Reserve for barriers and other uses\n static int constexpr kMaxStages = (kBlackwellSmemSize - kReservedBytes) / kBytesPerStage;\n auto sP = Int{}; // SMEM pipelines\n\n auto sA = UMMA::tile_to_mma_shape(SmemLayoutAtomA{}, append(mma_shape_A, sP),\n Step<_2, _1, _3>{}); // (MMA,MMA_M,MMA_K,PIPE)\n auto sB = UMMA::tile_to_mma_shape(SmemLayoutAtomB{},\n append(mma_shape_B, _1{})); // (MMA,MMA_N,MMA_K, _1)\n auto sD = Layout<_1>{}; // XXX Dummy\n\n auto tma_load_a =\n make_tma_copy_A_sm100(SM90_TMA_LOAD{}, tensorA, sA(_, _, _, 0), cluster_tile_mainloop, mma);\n auto tma_load_b =\n make_tma_copy_B_sm100(SM90_TMA_LOAD{}, tensorB, sB(_, _, _, 0), cluster_tile_shape, mma);\n\n // Assert checks on tile sizes -- no predication\n assert(M % size<0>(cluster_tile_shape) == 0);\n assert(N % size<1>(cluster_tile_shape) == 0);\n\n dim3 dimBlock(512);\n dim3 dimCluster(size<0>(cluster_shape), size<1>(cluster_shape), size<2>(cluster_shape));\n dim3 dimGrid(sm_count, 1, 1);\n\n int smem_size = sizeof(\n SharedStorage);\n\n auto *kernel_ptr = &group_row_col_rht_gemm_device_graph_safe<\n decltype(M), decltype(N), decltype(k_tile_size), decltype(cluster_shape),\n decltype(cluster_tile_shape), TA, decltype(dA), decltype(sA), decltype(tma_load_a), TB,\n decltype(dB), decltype(sB), decltype(tma_load_b), TD, decltype(dD), decltype(sD), TSFD,\n decltype(sfd_layout), TQA, decltype(dQA), TSFA, decltype(sfa_layout), decltype(mma),\n AccumulatorPipelineStageCount, SchedulerPipelineStageCount, kEnableStochasticRounding,\n kEnableRHTColQuant, kEnableRowQuant, kEnableSwizzleSFOutput, kUseFastMath>;\n\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(*kernel_ptr), hipFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n\n // Set workspace and set to zero\n NVTE_CHECK_CUDA(hipMemsetAsync(reinterpret_cast(tile_scheduler_workspace), 0,\n sizeof(uint32_t), stream));\n\n // Launch kernel\n cutlass::ClusterLaunchParams params = {dimGrid, dimBlock, dimCluster, smem_size, stream};\n cutlass::Status status = cutlass::launch_kernel_on_cluster(\n params, (void const *)kernel_ptr, M, N, k_tile_size, cluster_shape, cluster_tile_shape, A, dA,\n sA, tma_load_a, B, dB, sB, tma_load_b, QA, dQA, SFA, sfa_layout, QA_COLWISE, SFA_COLWISE,\n amax_rowwise, amax_colwise, offsets, first_dims, num_tensors, shape_rep,\n tile_scheduler_workspace, mma, rng_state);\n NVTE_CHECK_CUDA(hipGetLastError());\n NVTE_CHECK(status == cutlass::Status::kSuccess, \"Kernel launch failed.\");\n}\n\n} // namespace\n} // namespace detail\n\nvoid group_hadamard_transform_cast_fusion_graph_safe(const GroupedTensor *input,\n GroupedTensor *output,\n const Tensor &hadamard_matrix_,\n QuantizationConfig &quant_config,\n Tensor &quant_workspace, hipStream_t stream) {\n NVTE_API_CALL(group_hadamard_transform_cast_fusion_graph_safe);\n\n using transformer_engine::detail::kMaxTensorsPerKernel;\n using transformer_engine::detail::ShapeRepresentation;\n\n void *input_base_ptr = reinterpret_cast(input->data.dptr);\n // TODO(zhongbo): add input sanity checks here\n\n bool all_has_row_quant = output->has_data();\n bool all_has_col_quant = output->has_columnwise_data();\n\n // Stochastic rounding config\n const bool use_stochastic_rounding = quant_config.stochastic_rounding;\n const size_t *rng_state = nullptr;\n if (use_stochastic_rounding) {\n NVTE_CHECK(quant_config.rng_state != nullptr,\n \"Enabled stochastic rounding without providing RNG state\");\n const Tensor &rng_state_tensor = *convertNVTETensorCheck(quant_config.rng_state);\n NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64,\n \"RNG state should contain 2 64-bit values.\");\n NVTE_CHECK(rng_state_tensor.data.shape == std::vector{2},\n \"Shape of the RNG state should be [2], but got \", rng_state_tensor.data.shape);\n rng_state = reinterpret_cast(rng_state_tensor.data.dptr);\n }\n\n uint32_t *tile_scheduler_workspace = nullptr;\n NVTE_CHECK(quant_workspace.data.dptr != nullptr, \"Quantization workspace must be provided.\");\n NVTE_CHECK(quant_workspace.data.buffer_size_bytes() >= sizeof(uint32_t),\n \"Quantization workspace must be at least 4 bytes.\");\n tile_scheduler_workspace = reinterpret_cast(quant_workspace.data.dptr);\n\n // Template arguments\n using TA = cute::bfloat16_t;\n using TB = cute::bfloat16_t;\n using TD = cutlass::float_e2m1_t;\n using TSFD = cutlass::float_ue4m3_t;\n using TQA = TD;\n using TSFA = TSFD;\n\n checkCuDriverContext(stream);\n\n // Check Hadamard matrix\n constexpr int kHadamardDimension = 16;\n\n NVTE_CHECK(hadamard_matrix_.dtype() == transformer_engine::DType::kBFloat16,\n \"Hadamard matrix must be BF16 tensor, but dtype is \",\n to_string(hadamard_matrix_.dtype()), \".\");\n const SimpleTensor &hadamard_matrix = hadamard_matrix_.data;\n NVTE_CHECK(\n (hadamard_matrix_.shape() == std::vector{kHadamardDimension, kHadamardDimension}),\n \"Hadamard matrix must have shape=\",\n std::vector{kHadamardDimension, kHadamardDimension},\n \", but got shape=\", hadamard_matrix_.shape(), \".\");\n const size_t hadamard_dimension = hadamard_matrix.shape[0];\n\n const size_t num_tensors = input->num_tensors;\n const size_t first_logical_dim = input->logical_shape.data[0];\n const size_t last_logical_dim = input->logical_shape.data[1];\n // const size_t elts_total = first_logical_dim * last_logical_dim;\n NVTE_CHECK(first_logical_dim % 128 == 0,\n \"First dimension of a grouped tensor should be divisible by 128.\");\n NVTE_CHECK(last_logical_dim % 128 == 0,\n \"Last dimension of a grouped tensor should be divisible by 128.\");\n NVTE_CHECK(num_tensors <= kMaxTensorsPerKernel,\n \"Number of tensors should be less than or equal to \", kMaxTensorsPerKernel);\n\n ShapeRepresentation shape_rep = ShapeRepresentation::VARYING_FIRST_DIM;\n if (output->all_same_shape()) {\n shape_rep = ShapeRepresentation::SAME_BOTH_DIMS;\n } else if (output->all_same_first_dim()) {\n shape_rep = ShapeRepresentation::VARYING_LAST_DIM;\n } else if (output->all_same_last_dim()) {\n shape_rep = ShapeRepresentation::VARYING_FIRST_DIM;\n } else if (output->varying_both_dims()) {\n shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS;\n }\n\n TQA *const rowwise_data_base_ptr = reinterpret_cast(output->data.dptr);\n TSFA *const rowwise_scale_inv_base_ptr = reinterpret_cast(output->scale_inv.dptr);\n TQA *const colwise_data_base_ptr = reinterpret_cast(output->columnwise_data.dptr);\n TSFA *const colwise_scale_inv_base_ptr =\n reinterpret_cast(output->columnwise_scale_inv.dptr);\n float *const amax_rowwise_base_ptr = reinterpret_cast(output->amax.dptr);\n float *const amax_colwise_base_ptr = reinterpret_cast(output->columnwise_amax.dptr);\n\n const int64_t *const offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr);\n const int64_t *const first_dims_ptr = reinterpret_cast(output->first_dims.dptr);\n\n const bool is_const_last_dim = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS ||\n shape_rep == ShapeRepresentation::VARYING_FIRST_DIM);\n NVTE_CHECK(is_const_last_dim,\n \"Currently we only support const last dimension for graph safe hadamard transform.\");\n\n auto sm_count = transformer_engine::cuda::sm_count();\n\n int k_tile_size = 1024;\n\n const bool use_swizzle_sf_output = output->with_gemm_swizzled_scales;\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_stochastic_rounding, kEnableStochasticRounding,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n all_has_col_quant, kEnableRhtColQuant,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n all_has_row_quant, kEnableRowQuant,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_swizzle_sf_output, kEnableSwizzleSFOutput,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n quant_config.use_fast_math, kUseFastMath,\n\n if constexpr (kEnableRhtColQuant || kEnableRowQuant) {\n detail::group_row_col_rht_gemm_ntt_w_sfc_graph_safe<\n kEnableStochasticRounding, kEnableRhtColQuant, kEnableRowQuant,\n kEnableSwizzleSFOutput, TA, TB, TQA, TSFA, TD, TSFD, kUseFastMath>(\n /*packed_sequence_length=*/first_logical_dim,\n /*hidden_size=*/last_logical_dim,\n /*num_tensors=*/num_tensors,\n /*shape_rep=*/shape_rep,\n /*A=*/reinterpret_cast(input_base_ptr),\n /*B=*/reinterpret_cast(hadamard_matrix.dptr),\n /*QA=*/reinterpret_cast(rowwise_data_base_ptr),\n /*SFA=*/reinterpret_cast(rowwise_scale_inv_base_ptr),\n /*QA_COLWISE=*/reinterpret_cast(colwise_data_base_ptr),\n /*SFA_COLWISE=*/reinterpret_cast(colwise_scale_inv_base_ptr),\n /*amax_rowwise=*/reinterpret_cast(amax_rowwise_base_ptr),\n /*amax_colwise=*/reinterpret_cast(amax_colwise_base_ptr),\n /*offsets=*/offsets_ptr,\n /*first_dims=*/first_dims_ptr,\n /*rng_state=*/rng_state,\n /*tile_scheduler_workspace=*/tile_scheduler_workspace,\n /*sm_count=*/sm_count,\n /*stream=*/stream, /*k_tile_size=*/k_tile_size);\n } else {\n NVTE_ERROR(\"Invalid kernel configuration (kEnableRHTColQuant=\",\n kEnableRhtColQuant, \", kEnableRowQuant=\", kEnableRowQuant, \").\");\n }\n\n );););););\n}\n\n} // namespace transformer_engine\n\nvoid nvte_group_hadamard_transform_cast_fusion_graph_safe(\n const NVTEGroupedTensor input, NVTEGroupedTensor output, const NVTETensor hadamard_matrix,\n const NVTEQuantizationConfig quant_config, NVTETensor quant_workspace, hipStream_t stream) {\n NVTE_API_CALL(nvte_group_hadamard_transform_cast_fusion_graph_safe);\n using namespace transformer_engine;\n\n GroupedTensor *input_tensor = convertNVTEGroupedTensorCheck(input);\n GroupedTensor *output_tensor = convertNVTEGroupedTensorCheck(output);\n\n Tensor *quant_workspace_tensor = convertNVTETensorCheck(quant_workspace);\n\n QuantizationConfig quant_config_cpp;\n if (quant_config != nullptr) {\n quant_config_cpp = *reinterpret_cast(quant_config);\n }\n\n if (input_tensor->num_tensors == 0) {\n return;\n }\n\n // Call the multi-tensor Hadamard transform amax implementation.\n group_hadamard_transform_cast_fusion_graph_safe(\n input_tensor, output_tensor, *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp,\n *quant_workspace_tensor, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/util/cuda_runtime.h\"\n#include \"common/util/curanddx.hpp\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"customized_pipeline.cuh\"\n#include \"cutlass/arch/barrier.h\"\n#include \"cutlass/arch/reg_reconfig.h\"\n#include \"cutlass/cluster_launch.hpp\"\n#include \"cutlass/cutlass.h\"\n#include \"cutlass/detail/sm100_blockscaled_layout.hpp\"\n#include \"cutlass/fast_math.h\"\n#include \"cutlass/float8.h\"\n#include \"cutlass/float_subbyte.h\"\n#include \"cutlass/gemm/collective/builders/sm100_common.inl\"\n#include \"cutlass/numeric_conversion.h\"\n#include \"cutlass/numeric_types.h\"\n#include \"cutlass/pipeline/pipeline.hpp\"\n#include \"cutlass/platform/platform.h\"\n#include \"cutlass/util/GPU_Clock.hpp\"\n#include \"cutlass/util/command_line.h\"\n#include \"cutlass/util/print_error.hpp\"\n\nnamespace transformer_engine {\nnamespace detail {\nnamespace {\n\nusing namespace cute;\n\n// Ensure Tensor refers to cute::Tensor, not transformer_engine::Tensor\nusing cute::Tensor;\n\nconstexpr int kMaxTensorsPerKernel = 64;\n\nstruct MultiAmaxHadamardCastFusionArgs {\n // (output) Amax buffer for input A amax buffer\n void *global_a_amax_list[kMaxTensorsPerKernel];\n // (output) Amax buffer for pre-RHT amax buffer\n void *global_d_amax_list[kMaxTensorsPerKernel];\n // output D pointers for each tensor\n void *output_colwise_list[kMaxTensorsPerKernel];\n // output SFD inverse pointers for each tensor\n void *output_colwise_scale_inv_list[kMaxTensorsPerKernel];\n // split sections of each tensor of input\n int split_sections[kMaxTensorsPerKernel];\n // Prefix sum (with leading zero) of split_sections of each tensor of input\n int split_sections_range[kMaxTensorsPerKernel + 1];\n\n // Number of tensors (splits) being processed by kernel\n int num_tensors;\n};\n\n__device__ __forceinline__ int GetGroupIdx(MultiAmaxHadamardCastFusionArgs *kernel_args_ptr,\n int offset) {\n // Check the kernel args and get the corresponding id\n const int num_tensors = kernel_args_ptr->num_tensors;\n if (offset >= kernel_args_ptr->split_sections_range[num_tensors]) {\n return num_tensors - 1;\n }\n int group_idx = 0;\n while (kernel_args_ptr->split_sections_range[group_idx + 1] <= offset) {\n ++group_idx;\n }\n return group_idx;\n}\n\nCUTLASS_DEVICE\ncutlass::Array StochasticNumericConverterBase(\n cutlass::Array const &input, cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n auto output_ptr = reinterpret_cast(&output);\n constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING;\n if constexpr (has_rs) {\n asm volatile(\n \"{\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\\n\"\n \"}\"\n : \"=h\"(output_ptr[0]), \"=h\"(output_ptr[1])\n : \"f\"(input[0]), \"f\"(input[1]), \"f\"(input[2]), \"f\"(input[3]), \"f\"(input[4]), \"f\"(input[5]),\n \"f\"(input[6]), \"f\"(input[7]), \"r\"(rbits[0]), \"r\"(rbits[1]));\n } else {\n NVTE_DEVICE_ERROR(\n \"FP4 cvt PTX instructions are architecture-specific. \"\n \"Try recompiling with sm_XXXa instead of sm_XXX.\");\n }\n return output;\n}\n\nCUTLASS_DEVICE\ncutlass::Array StochasticNumericConverter(\n cutlass::Array const &input, cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n cutlass::Array *result_ptr =\n reinterpret_cast *>(&output);\n cutlass::Array const *source_ptr =\n reinterpret_cast const *>(&input);\n cutlass::Array const *rbits_ptr =\n reinterpret_cast const *>(&rbits);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 2; i++) {\n result_ptr[i] = StochasticNumericConverterBase(source_ptr[i], rbits_ptr[i]);\n }\n return output;\n}\n\ntemplate \nstruct SharedStorage {\n static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;\n static int constexpr EpilogueUnrollFactor = EpilogueUnrollFactor_;\n using AtomThrShapeMNK = cute::Shape<_1, _1, _1>;\n\n using AccumulatorPipeline =\n cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;\n\n static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline =\n cutlass::detail::CustomizedPipelineTmaUmmaAsync,\n AtomThrShapeMNK>;\n using MainloopPipelineStorage = typename MainloopPipeline::SharedStorage;\n using SchedPipeline = cutlass::PipelineCLCFetchAsync;\n using SchedPipelineStorage = typename SchedPipeline::SharedStorage;\n using SchedThrottlePipeline = cutlass::PipelineAsync;\n using SchedThrottlePipelineStorage = typename SchedThrottlePipeline::SharedStorage;\n\n struct TensorStorage : cute::aligned_struct<128, _1> {\n cute::array_aligned> smem_A;\n cute::array_aligned> smem_B;\n } tensors;\n\n alignas(16) AccumulatorPipelineStorage accumulator;\n alignas(16) MainloopPipelineStorage mainloop;\n alignas(16) cute::uint64_t tma_barrier[1];\n alignas(16) SchedPipelineStorage sched;\n alignas(16) SchedThrottlePipelineStorage sched_throttle;\n alignas(16) int32_t atomic_tile_id[SchedulerPipelineStageCount_];\n alignas(16) float global_a_amax[kMaxTensorsPerKernel];\n alignas(16) float global_d_amax[kMaxTensorsPerKernel];\n uint32_t atomic_tile_counter[SchedulerPipelineStageCount_];\n uint32_t tmem_base_ptr;\n};\n\n// Main RHT GEMM kernel entry -- highly templated for flexible architecture/config support\ntemplate \n__launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device(\n MShape M, NShape packed_N, KShape K, ClusterShape cluster_shape, ClusterTileShape cluster_tile,\n TA const *A, AStride dA, ASmemLayout sAlayout, CUTE_GRID_CONSTANT TmaLoadA const tma_load_a,\n TB const *B, BStride dB, BSmemLayout sBlayout, CUTE_GRID_CONSTANT TmaLoadB const tma_load_b,\n TQA *QA, QAStride dQA, TSFA *SFA, TSFALayout sfa_layout, MultiAmaxHadamardCastFusionArgs args,\n uint32_t *tile_scheduler_workspace, TiledMMA mma, const size_t *rng_state) {\n using namespace cute;\n\n // Abort immediately if compilation is not supported\n constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY;\n if constexpr (!is_blackwell_arch) {\n NVTE_DEVICE_ERROR(\"RHT fusion is only supported on Blackwell.\");\n return;\n } else {\n static_assert(kEnableRHTColQuant_ || kEnableRowQuant_,\n \"group_row_col_rht_gemm_device must generate row-wise \"\n \"and/or column-wise output.\");\n#if !defined(CUTLASS_ARCH_CLC_ENABLED)\n CUTLASS_NOT_IMPLEMENTED();\n return;\n#endif\n\n using X = Underscore;\n // Accumulator data type for main computation\n using ElementAccumulator = float;\n static int constexpr K_PIPE_MAX = size<3>(ASmemLayout{});\n using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>;\n static uint32_t constexpr kTmaTransactionBytes = cutlass::bits_to_bytes(\n size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v);\n static constexpr bool kEnableStochasticRounding = kEnableStochasticRounding_;\n static constexpr bool kEnableRHTColQuant = kEnableRHTColQuant_;\n static constexpr bool kEnableRowQuant = kEnableRowQuant_;\n static constexpr bool kEnableSwizzleSFOutput = kEnableSwizzleSFOutput_;\n static constexpr bool kUseFastMath = kUseFastMath_;\n\n // Constant for RHT tensor processing (tile size etc)\n static int constexpr RhtTensorSize = 16;\n\n // Transaction bytes for TMA transfer on RHT tensor blocks\n static int constexpr kTmaRhtTensorTransactionBytes =\n cutlass::bits_to_bytes(RhtTensorSize * RhtTensorSize * cute::sizeof_bits_v);\n static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;\n static int constexpr SchedulerPipelineStageCount = SchedulerPipelineStageCount_;\n\n // Mainloop pipeline stage calculation, vectorization parameters for scaling factors\n static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{});\n static int constexpr SFVecSize = 16;\n // Swizzle output layout for scaling factor arrays\n using SwizzledSFALayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFDLayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n\n // Mainloop pipeline types for TMA async execution and epilogue cluster scheduling\n using MainloopPipeline =\n cutlass::detail::CustomizedPipelineTmaUmmaAsync;\n using MainloopPipelineState = typename MainloopPipeline::PipelineState;\n using SchedPipeline = cutlass::PipelineCLCFetchAsync;\n using SchedPipelineState = typename SchedPipeline::PipelineState;\n using SchedThrottlePipeline = cutlass::PipelineAsync;\n using SchedThrottlePipelineState = typename SchedThrottlePipeline::PipelineState;\n\n static_assert(ClusterShape{} == Shape<_1, _1, _1>{}, \"ClusterShape must be Shape<_1,_1,_1>\");\n\n using TmemAllocator = cute::TMEM::Allocator1Sm;\n static int constexpr VectorSize = RhtTensorSize;\n\n // Compile-time safety: static shapes required for shared memory layouts\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n // CUTE_STATIC_ASSERT(is_static::value);\n\n auto cluster_size = size<0>(cluster_shape);\n auto mainloop_tiler = Shape<_128, _16, _128>{};\n auto epilogue_tiler = Shape<_128, _128, _128>{};\n\n static int constexpr EpilogueUnrollFactor = size<2>(epilogue_tiler) / size<2>(cluster_tile);\n\n // Get the appropriate blocks for this Cluster\n dim3 cluster_coord_in_grid = cluster_id_in_grid();\n\n // Total number of k-tiles\n int const K_TILE_MAX = min(packed_N, K) / size<2>(epilogue_tiler);\n\n struct TileScheduler {\n uint32_t tiles_in_m = 0;\n uint32_t tiles_in_n = 0;\n uint32_t linear_idx = 0;\n uint32_t next_linear_idx = 0;\n uint32_t start_idx = 0;\n uint32_t tile_m_idx = 0;\n uint32_t tile_n_idx = 0;\n int k_tile_max = 0;\n uint32_t *atomic_tile_index_;\n uint32_t *smem_tile_counter;\n uint32_t atomic_offset;\n cutlass::FastDivmodU64 divmod_tiles_in_m;\n\n CUTLASS_DEVICE TileScheduler(uint32_t tiles_m, uint32_t tiles_n, int kmax,\n uint32_t *atomic_tile_index, uint32_t *smem_tile_counter)\n : tiles_in_m(tiles_m),\n tiles_in_n(tiles_n),\n linear_idx(blockIdx.x),\n next_linear_idx(blockIdx.x),\n start_idx(blockIdx.x),\n k_tile_max(kmax),\n atomic_tile_index_(atomic_tile_index),\n smem_tile_counter(smem_tile_counter),\n atomic_offset(gridDim.x),\n divmod_tiles_in_m(uint64_t(tiles_m)) {\n update_tile_idx();\n }\n CUTLASS_DEVICE void update_tile_idx() {\n uint64_t q, r;\n divmod_tiles_in_m(q, r, uint64_t(linear_idx));\n tile_m_idx = static_cast(r);\n tile_n_idx = static_cast(q) * uint32_t(k_tile_max);\n }\n CUTLASS_DEVICE uint32_t tile_m() const { return tile_m_idx; }\n CUTLASS_DEVICE uint32_t tile_n_base() const { return tile_n_idx; }\n CUTLASS_DEVICE uint32_t tiles_m() const { return tiles_in_m; }\n\n CUTLASS_DEVICE uint32_t tiles_n() const { return tiles_in_n; }\n\n CUTLASS_DEVICE bool is_valid() const {\n return cute::elem_less(cute::make_coord(tile_m(), tile_n_base()),\n cute::make_coord(tiles_in_m, tiles_in_n));\n }\n\n CUTLASS_DEVICE bool is_first_wave() const { return linear_idx == start_idx; }\n\n CUTLASS_DEVICE uint32_t get_linear_tile_idx() const { return linear_idx; }\n\n // Fetch a new tile_id using atomics.\n CUTLASS_DEVICE uint32_t fetch_tile_id_counter(int pred) {\n uint32_t tile_id_counter = 0;\n asm volatile(\n \"{\\n\\t\"\n \".reg .pred p;\\n\\t\"\n \"setp.eq.u32 p, %2, 1;\\n\\t\"\n \"@p atom.global.add.u32 %0, [%1], 1; \\n\\t\"\n \"}\"\n : \"=r\"(tile_id_counter)\n : \"l\"(atomic_tile_index_), \"r\"(pred));\n\n return tile_id_counter;\n }\n\n CUTLASS_DEVICE auto fetch_next_work(SchedPipeline &sched_pipeline,\n SchedPipelineState sched_pipeline_consumer_state) {\n sched_pipeline.consumer_wait(sched_pipeline_consumer_state);\n next_linear_idx = smem_tile_counter[sched_pipeline_consumer_state.index()];\n cutlass::arch::fence_view_async_shared();\n sched_pipeline.consumer_release(sched_pipeline_consumer_state);\n return;\n }\n\n CUTLASS_DEVICE auto advance_to_next_work(SchedPipeline &sched_pipeline,\n SchedPipelineState sched_pipeline_producer_state) {\n uint32_t mbarrier_addr = sched_pipeline.producer_get_barrier(sched_pipeline_producer_state);\n // Wait for clcID buffer to become empty with a flipped phase\n sched_pipeline.producer_acquire(sched_pipeline_producer_state);\n auto is_leading_thread = cute::elect_one_sync();\n uint32_t tile_id_counter = fetch_tile_id_counter(is_leading_thread) + atomic_offset;\n uint32_t smem_addr =\n cute::cast_smem_ptr_to_uint(&smem_tile_counter[sched_pipeline_producer_state.index()]);\n if (is_leading_thread) {\n cute::store_shared_remote(tile_id_counter, smem_addr, mbarrier_addr, 0);\n }\n\n ++sched_pipeline_producer_state;\n return sched_pipeline_producer_state;\n }\n\n CUTLASS_DEVICE auto update_work_tile_info() {\n linear_idx = next_linear_idx;\n update_tile_idx();\n return;\n }\n };\n\n // Allocate and alias shared memory to the kernel's shared storage type\n extern __shared__ char shared_memory[];\n using SharedStorage =\n SharedStorage;\n SharedStorage &shared_storage = *reinterpret_cast(shared_memory);\n\n // Compute the number of tiles in M and N after tiling and assign scheduler\n uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile))));\n uint32_t tiles_in_n = uint32_t(\n size(ceil_div(args.split_sections_range[args.num_tensors], size<2>(epilogue_tiler))));\n\n TileScheduler scheduler(tiles_in_m, tiles_in_n, K_TILE_MAX, tile_scheduler_workspace,\n shared_storage.atomic_tile_counter);\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n\n // Shapes for accumulated tiles in mainloop and epilogue\n auto acc_shape_mma = make_shape(take<0, 2>(mainloop_tiler), _1{}, _1{});\n auto acc_shape_epilogue = make_shape(take<0, 2>(epilogue_tiler), _1{}, _1{});\n\n // Shape of the accumulator fragment for the main loop pipeline, with pipeline stages appended\n auto acc_mainloop_pipelined_shape = append(acc_shape_mma, Int{});\n auto bulk_tmem_mma = TiledMMA::make_fragment_C(acc_mainloop_pipelined_shape);\n\n // Number of threads assigned for various epilogue roles depending on quantization settings\n static int constexpr NumEpilogueColQuantThreadCount = kEnableRHTColQuant ? 128 : 0;\n static int constexpr NumEpilogueRowQuantThreadCount = kEnableRowQuant ? 256 : 0;\n static int constexpr NumMmaThreadCount = kEnableRHTColQuant ? 32 : 0;\n static int constexpr NumMmaIssueThreadCount = kEnableRHTColQuant ? 1 : 0;\n static int constexpr NumSchedThreads = 32;\n static int constexpr NumMainloopLoadThreads = 32;\n static int constexpr NumEpilogueThreads =\n NumEpilogueColQuantThreadCount + NumEpilogueRowQuantThreadCount;\n\n TmemAllocator tmem_allocator{};\n cutlass::arch::NamedBarrier tmem_allocation_result_barrier(\n NumMmaThreadCount + NumEpilogueColQuantThreadCount,\n cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);\n\n int warp_idx = cutlass::canonical_warp_idx_sync();\n\n // warp assignment\n bool is_mma_warp = (warp_idx == 0);\n bool is_dma_warp = (warp_idx == 1);\n bool is_sched_warp = (warp_idx == 2);\n bool is_epilogue_col_quant_warp = (warp_idx >= 4 && warp_idx <= 7);\n bool is_epilogue_row_quant_warp = (warp_idx >= 8 && warp_idx <= 15);\n\n typename MainloopPipeline::Params mainloop_pipeline_params;\n if (is_dma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;\n }\n if (is_mma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;\n }\n mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp;\n mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes;\n mainloop_pipeline_params.initializing_warp = 0;\n mainloop_pipeline_params.num_consumers =\n NumEpilogueRowQuantThreadCount + NumMmaIssueThreadCount;\n\n MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params,\n cluster_shape, cute::true_type{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n\n MainloopPipelineState mainloop_pipe_consumer_state;\n MainloopPipelineState mainloop_pipe_producer_state =\n cutlass::make_producer_start_state();\n\n using AccumulatorPipeline =\n cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState;\n using AccumulatorPipelineInitBarriers = cute::bool_constant;\n\n AccumulatorPipelineState accumulator_pipe_consumer_state;\n AccumulatorPipelineState accumulator_pipe_producer_state =\n cutlass::make_producer_start_state();\n\n typename AccumulatorPipeline::Params accumulator_pipeline_params;\n if (is_mma_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer;\n }\n if (is_epilogue_col_quant_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer;\n }\n // Only one producer thread arrives on this barrier.\n accumulator_pipeline_params.producer_arv_count = 1;\n accumulator_pipeline_params.consumer_arv_count =\n size(AtomThrShapeMNK{}) * NumEpilogueColQuantThreadCount;\n accumulator_pipeline_params.initializing_warp = 1;\n AccumulatorPipeline accumulator_pipeline(\n shared_storage.accumulator, accumulator_pipeline_params, cluster_shape,\n AccumulatorPipelineInitBarriers{}, cute::true_type{}); // Delay mask calculation\n typename SchedPipeline::Params sched_pipeline_params;\n if (is_sched_warp) {\n sched_pipeline_params.role = SchedPipeline::ThreadCategory::ProducerConsumer;\n } else {\n sched_pipeline_params.role = SchedPipeline::ThreadCategory::Consumer;\n }\n sched_pipeline_params.producer_blockid = 0;\n sched_pipeline_params.producer_arv_count = 1;\n sched_pipeline_params.consumer_arv_count =\n NumSchedThreads +\n cluster_size * (NumMainloopLoadThreads + NumEpilogueThreads + NumMmaThreadCount);\n sched_pipeline_params.transaction_bytes = sizeof(uint32_t);\n sched_pipeline_params.initializing_warp = 3;\n SchedPipeline sched_pipeline(shared_storage.sched, sched_pipeline_params, cluster_shape);\n SchedPipelineState sched_pipeline_consumer_state;\n SchedPipelineState sched_pipeline_producer_state =\n cutlass::make_producer_start_state();\n\n typename SchedThrottlePipeline::Params sched_throttle_pipeline_params;\n if (is_dma_warp) {\n sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Producer;\n }\n if (is_sched_warp) {\n sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Consumer;\n }\n sched_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads;\n sched_throttle_pipeline_params.consumer_arv_count = NumSchedThreads;\n sched_throttle_pipeline_params.dst_blockid = 0;\n sched_throttle_pipeline_params.initializing_warp = 4;\n\n SchedThrottlePipeline sched_throttle_pipeline(shared_storage.sched_throttle,\n sched_throttle_pipeline_params);\n SchedThrottlePipelineState sched_pipeline_throttle_consumer_state;\n SchedThrottlePipelineState sched_pipeline_throttle_producer_state =\n cutlass::make_producer_start_state();\n\n if (warp_idx == 2 && elect_one_sync()) {\n cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);\n }\n __syncthreads();\n\n // Warp group roles: DMA (global->shared copy), MMA (tensor core gemm), scheduler, column quantizer, row quantizer\n if (is_dma_warp) {\n // Warp responsible for loading input from global to shared memory using TMA (Tensor Memory Access).\n cutlass::arch::warpgroup_reg_dealloc<32>();\n // Get TMA tensors for input matrix A and B (Hadamard/transform matrix) from global memory.\n Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, packed_N));\n Tensor mB = tma_load_b.get_tma_tensor(make_shape(RhtTensorSize, RhtTensorSize));\n\n // Partition tensors for tiling according to the mainloop and cluster tilers.\n Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{});\n Tensor gB_nk =\n local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k)\n\n // Shared memory tensors for pipeline\n Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()),\n sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()),\n sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n // Determine warp/tile positioning\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n // Partition global to local fragments for A and B\n Tensor tCgA = thr_mma.partition_A(gA_mk); // (MMA,MMA_M,MMA_K,k)\n Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k)\n\n Layout cta_layout_mnk = make_layout(cluster_shape);\n Layout cta_layout_vmnk =\n tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{}));\n auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster);\n\n auto [tAgA, tAsA] =\n tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),\n group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA));\n\n auto [tBgB, tBsB] =\n tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),\n group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB));\n\n uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);\n uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);\n if constexpr (kEnableRHTColQuant) {\n if (elect_one_sync()) {\n cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0],\n kTmaRhtTensorTransactionBytes);\n copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0),\n tBsB(_, 0));\n }\n }\n\n do {\n // is_first_wave indicates whether this scheduler wave is the first among a group.\n bool is_first_wave = scheduler.is_first_wave();\n uint32_t skip_wait = is_first_wave;\n auto tAgA_mk = tAgA(_, scheduler.tile_m(), _);\n int k_tile = 0;\n\n sched_throttle_pipeline.producer_acquire(sched_pipeline_throttle_producer_state);\n sched_throttle_pipeline.producer_commit(sched_pipeline_throttle_producer_state);\n ++sched_pipeline_throttle_producer_state;\n CUTLASS_PRAGMA_NO_UNROLL\n while (k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n()) {\n int k_tile_idx_n = scheduler.tile_n_base() + k_tile;\n ++k_tile;\n skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount);\n mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state);\n using BarrierType = typename MainloopPipeline::ProducerBarrierType;\n BarrierType *tma_barrier =\n mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state);\n int write_stage = mainloop_pipe_producer_state.index();\n ++mainloop_pipe_producer_state;\n if (cute::elect_one_sync()) {\n copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n),\n tAsA(_, write_stage));\n }\n }\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n // scheduler.advance();\n } while (scheduler.is_valid());\n mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);\n } else if (is_mma_warp) {\n // This warp executes the main tensor core matrix-multiply-accumulate for the Hadamard transform.\n cutlass::arch::warpgroup_reg_dealloc<32>();\n if constexpr (kEnableRHTColQuant) {\n // Setup shared memory fragments for A and B tiles.\n Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()),\n sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()),\n sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n // Allocate \"fragments\" -- these are actually umma smem descriptors\n Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)\n Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE)\n\n mma.accumulate_ = UMMA::ScaleOut::Zero;\n\n tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns,\n &shared_storage.tmem_base_ptr);\n __syncwarp();\n tmem_allocation_result_barrier.arrive();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_mma.data() = tmem_base_ptr;\n // Wait until the B (Hadamard) tensor copy is complete\n cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/);\n do {\n uint32_t skip_wait = K_TILE_MAX <= 0;\n\n auto barrier_token =\n mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0;\n k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) {\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n int read_stage = mainloop_pipe_consumer_state.index();\n auto tCrA_mk = tCrA(_, _, _, read_stage);\n auto tCrB_nk = tCrB(_, _, 0, 0);\n CUTLASS_PRAGMA_UNROLL\n for (int k_block = 0; k_block < size<2>(tCrA) / EpilogueUnrollFactor; ++k_block) {\n int accumulator_k_block =\n accumulator_pipe_producer_state.index() * EpilogueUnrollFactor;\n int tCrA_k_block = k_block * EpilogueUnrollFactor;\n accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < EpilogueUnrollFactor; i++) {\n auto accumulators = bulk_tmem_mma(_, _, _, accumulator_k_block + i);\n gemm(mma, tCrA_mk(_, _, tCrA_k_block + i), tCrB_nk, accumulators);\n }\n\n accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);\n ++accumulator_pipe_producer_state;\n }\n auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n skip_wait = k_tile >= K_TILE_MAX;\n mainloop_pipeline.umma_consumer_release(curr_mainloop_pipe_consumer_state);\n barrier_token =\n mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n }\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n tmem_allocator.release_allocation_lock();\n accumulator_pipeline.producer_tail(accumulator_pipe_producer_state);\n tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);\n }\n } else if (is_sched_warp) {\n // Scheduler warp manages tile assignment and pipeline progress for warps\n cutlass::arch::warpgroup_reg_dealloc<32>();\n do {\n sched_throttle_pipeline.consumer_wait(sched_pipeline_throttle_consumer_state);\n sched_throttle_pipeline.consumer_release(sched_pipeline_throttle_consumer_state);\n ++sched_pipeline_throttle_consumer_state;\n sched_pipeline_producer_state =\n scheduler.advance_to_next_work(sched_pipeline, sched_pipeline_producer_state);\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n } else if (is_epilogue_col_quant_warp) {\n // Warp responsible for quantizing output of Hadamard transform to FP4 for columnwise usage,\n // and writing result tensors/scales to global memory.\n cutlass::arch::warpgroup_reg_alloc<192>();\n if constexpr (kEnableRHTColQuant) {\n using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x;\n\n auto acc_epilogue_pipelined_shape =\n append(acc_shape_epilogue, Int{});\n auto bulk_tmem_epilogue_layout = make_layout(\n acc_epilogue_pipelined_shape,\n make_stride(stride<0>(bulk_tmem_mma), Int<0>{}, Int<0>{}, size<1>(epilogue_tiler)));\n auto bulk_tmem_epilogue = make_tensor(make_tmem_ptr(), bulk_tmem_epilogue_layout);\n\n // Use 256-bit fragments for aligned bulk stores\n static int constexpr FragmentSize = 256 / sizeof_bits_v;\n\n // Wait for TMEM allocation for this pipeline to finish\n tmem_allocation_result_barrier.arrive_and_wait();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_epilogue.data() = tmem_base_ptr;\n int global_thread_idx = threadIdx.x;\n int local_thread_idx = global_thread_idx % cutlass::NumThreadsPerWarpGroup;\n // g2s load all global_d_amax\n CUTLASS_PRAGMA_NO_UNROLL\n for (int g = local_thread_idx; g < args.num_tensors; g += NumEpilogueColQuantThreadCount) {\n shared_storage.global_d_amax[g] =\n __ldg(reinterpret_cast(args.global_d_amax_list[g]));\n }\n\n size_t rng_seed = 0;\n size_t rng_offset = 0;\n // Setup RNG for stochastic rounding\n if constexpr (kEnableStochasticRounding) {\n rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0;\n rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0;\n }\n int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler));\n\n // Determine quantization scale factor layouts/output splits for this group\n TSFDLayout sfd_layout;\n int cur_N = args.split_sections[group_idx];\n if constexpr (kEnableSwizzleSFOutput) {\n sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{});\n } else {\n sfd_layout = make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)),\n make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{})));\n }\n // Build output tensors for columns and their quant scales\n Tensor mD = make_tensor(\n cute::subbyte_iterator(reinterpret_cast(args.output_colwise_list[group_idx])),\n make_shape(M, cur_N), DStride{}); // (M,packed_N)\n Tensor gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n\n Tensor mSFD = make_tensor(make_gmem_ptr(reinterpret_cast(\n args.output_colwise_scale_inv_list[group_idx])),\n sfd_layout);\n Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n\n Tensor gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler));\n\n // Setup tile-level TMEM (t2r) and global memory (r2g) copy descriptors\n auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{}));\n auto tiled_r2g =\n make_tiled_copy_D(Copy_Atom{}, tiled_t2r);\n auto thr_t2r = tiled_t2r.get_slice(local_thread_idx);\n auto thr_r2g = tiled_r2g.get_slice(local_thread_idx);\n\n cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount,\n cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);\n // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release}\n static constexpr float fp4_max = 6.0f;\n static constexpr float fp8_max = 448.0f;\n static constexpr float fp4_max_inv = 1.0f / fp4_max;\n float c_global_amax_val = shared_storage.global_d_amax[group_idx];\n float global_encode_scale = c_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / c_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n float global_decode_scale = 1.0f / global_encode_scale;\n\n // Scaling factor for fast math path\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n\n do {\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0;\n k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();\n ++k_tile) {\n int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler);\n\n int cur_group_idx = GetGroupIdx(&args, global_tile_n_offset);\n\n if (cur_group_idx != group_idx) {\n group_idx = cur_group_idx;\n c_global_amax_val = shared_storage.global_d_amax[group_idx];\n // update amax\n global_encode_scale = c_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / c_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n global_decode_scale = 1.0f / global_encode_scale;\n global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n cur_N = args.split_sections[group_idx];\n if constexpr (kEnableSwizzleSFOutput) {\n sfd_layout =\n tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{});\n } else {\n sfd_layout =\n make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)),\n make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{})));\n }\n // update tensor\n mD = make_tensor(cute::subbyte_iterator(\n reinterpret_cast(args.output_colwise_list[group_idx])),\n make_shape(M, cur_N), DStride{});\n gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n mSFD = make_tensor(make_gmem_ptr(reinterpret_cast(\n args.output_colwise_scale_inv_list[group_idx])),\n sfd_layout);\n gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n\n gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler));\n }\n int group_start_offset = args.split_sections_range[group_idx];\n int local_tile_n_idx =\n (global_tile_n_offset - group_start_offset) / size<1>(epilogue_tiler);\n Tensor tDgD_mn = gD_mn_view(_, _, _, scheduler.tile_m(), local_tile_n_idx);\n\n Tensor tDgSFD_mn = gSFD_mn(_, _, scheduler.tile_m(), local_tile_n_idx);\n accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state);\n\n auto Acc = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index());\n Tensor tDtAcc = thr_t2r.partition_S(Acc); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDgD = thr_t2r.partition_D(tDgD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n\n Tensor tTR_rAcc = make_tensor(\n shape(tDgD)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDrD = make_tensor(shape(tDgD));\n Tensor tTR_rAcc_frag =\n recast>(coalesce(tTR_rAcc));\n Tensor tDrD_frag = recast>(coalesce(tDrD));\n\n Tensor src = thr_r2g.retile_S(tDrD);\n Tensor dst = thr_r2g.retile_D(tDgD);\n\n Tensor tDgSFD_view = make_tensor(\n tDgSFD_mn.data(), make_layout(make_shape(shape(tDgSFD_mn), Int<1>{}, Int<1>{}),\n make_stride(stride(tDgSFD_mn), Int<0>{}, Int<0>{})));\n Tensor tDgSFD = filter(thr_t2r.partition_D(tDgSFD_view));\n Tensor tDrSFD = make_tensor(shape(tDgSFD));\n\n static int constexpr NumVecs = size(tDgD) / VectorSize;\n Tensor tD_rRowSFD_frg = recast>(tDrSFD);\n\n // Compute amax and quantization scales for this tile\n cutlass::maximum_absolute_value_reduction<\n cutlass::Array, true>\n amax_reduction;\n cutlass::Array vec_maxs;\n cutlass::Array pvscales;\n // Copy from TMEM to registers\n copy(tiled_t2r, tDtAcc, tTR_rAcc);\n cutlass::arch::fence_view_async_tmem_load();\n accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state);\n ++accumulator_pipe_consumer_state;\n\n if constexpr (!kUseFastMath) {\n // Downcast to BF16 for bit-wise compatibility with\n // unfused kernels\n auto convert_accum_to_bf16 =\n cutlass::NumericArrayConverter{};\n auto convert_bf16_to_accum =\n cutlass::NumericArrayConverter{};\n tTR_rAcc_frag(_0{}) =\n convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{})));\n tTR_rAcc_frag(_1{}) =\n convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_1{})));\n }\n\n auto compute_frgs = reinterpret_cast *>(\n tTR_rAcc_frag.data());\n auto output_frgs = reinterpret_cast *>(tDrD_frag.data());\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]);\n }\n\n pvscales = cutlass::multiplies>{}(\n vec_maxs, global_encode_scale_multiplier);\n auto pvscales_cvted =\n cutlass::NumericArrayConverter{}(pvscales);\n\n tD_rRowSFD_frg(_0{}) = pvscales_cvted;\n auto qpvscale_ups = cutlass::NumericArrayConverter{}(\n tD_rRowSFD_frg(_0{}));\n auto qpvscale_scaled =\n cutlass::multiplies>{}(\n qpvscale_ups, global_decode_scale);\n cutlass::Array acc_scales;\n if constexpr (kUseFastMath) {\n // Fast math: compute approximate reciprocal\n acc_scales =\n cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled);\n } else {\n // Accurate math: compute reciprocal with division\n acc_scales = cutlass::divides>{}(\n 1.0, qpvscale_scaled);\n }\n\n // Prepare stochastic rounding random state if enabled\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n transformer_engine::curanddx::detail::philox4x32_native_state<\n NVTE_BUILD_NUM_PHILOX_ROUNDS>\n rng;\n // \"Prefetch\" a stochastic rounding state for the first tile\n if constexpr (kEnableStochasticRounding) {\n const size_t rng_sequence = global_thread_idx + k_tile * 512 +\n scheduler.get_linear_tile_idx() * K_TILE_MAX * 512;\n rng.init(rng_seed, rng_sequence, rng_offset);\n }\n CUTLASS_PRAGMA_UNROLL\n // Apply round/quantize to each fragment, with or without stochastic rounding\n for (int v = 0; v < NumVecs; v++) {\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(\n acc_scales[v], cutlass::platform::numeric_limits::max());\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(\n cutlass::multiplies>{}(\n compute_frgs[v], acc_scale),\n *reinterpret_cast *>(&random_uint4));\n } else {\n output_frgs[v] =\n cutlass::NumericArrayConverter{}(\n cutlass::multiplies>{}(\n compute_frgs[v], acc_scale));\n }\n }\n\n // Write quantized FP4 tile and dequant scale to gmem\n copy(tiled_r2g, src, dst);\n copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFD, tDgSFD);\n }\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n }\n } else if (is_epilogue_row_quant_warp) {\n // Warp responsible for quantizing the input (before Hadamard transform) to FP4 for row-wise usage.\n cutlass::arch::warpgroup_reg_alloc<136>();\n if constexpr (kEnableRowQuant) {\n using S2RVectorType = uint128_t;\n\n int global_thread_idx = threadIdx.x;\n int local_thread_idx = global_thread_idx % 256;\n size_t rng_seed = 0;\n size_t rng_offset = 0;\n // g2s load all global_a_amax for all groups/tensors\n CUTLASS_PRAGMA_NO_UNROLL\n for (int g = local_thread_idx; g < args.num_tensors; g += NumEpilogueRowQuantThreadCount) {\n shared_storage.global_a_amax[g] =\n __ldg(reinterpret_cast(args.global_a_amax_list[g]));\n }\n // RNG for stochastic rounding\n if constexpr (kEnableStochasticRounding) {\n rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0;\n rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0;\n }\n // Input/output tensors/partitions for row quant warp\n Tensor mQA =\n make_tensor(cute::subbyte_iterator(QA), make_layout(make_shape(M, packed_N), dQA));\n Tensor gQA_mn = local_tile(mQA, epilogue_tiler, make_coord(_, _, _), Step<_1, X, _1>{});\n Tensor mSFA = make_tensor(make_gmem_ptr(SFA), sfa_layout);\n\n Tensor gSFA_mn = local_tile(mSFA, epilogue_tiler, make_coord(_, _, _),\n Step<_1, X, _1>{}); // (BLK_M,BLK_N)\n // Swizzled shared memory A tile, with layout\n Tensor sA = as_position_independent_swizzle_tensor(group_modes<0, 2>(\n coalesce(make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()),\n sAlayout)))); // (BLOCK_M, BLOCK_M,PIPE)\n\n // Set up layouts for partitioning \u2013 tile-by-warp, with vector granularity\n using S2RWarpLayout = Layout>;\n using WarpGroupLayout = Layout>;\n using S2RThreadLayout = decltype(blocked_product(S2RWarpLayout{}, WarpGroupLayout{}));\n using S2RValLayout = Layout, _1>>;\n using S2RAtomA = Copy_Atom;\n using R2GAtomQA = Copy_Atom;\n using R2GAtomSFA = Copy_Atom;\n auto tiled_s2r = make_tiled_copy(S2RAtomA{}, S2RThreadLayout{}, S2RValLayout{});\n auto tiled_r2g_QA = make_tiled_copy(R2GAtomQA{}, S2RThreadLayout{}, S2RValLayout{});\n auto tiled_r2g_SFA = make_tiled_copy(R2GAtomSFA{}, S2RThreadLayout{}, S2RValLayout{});\n\n auto thr_s2r = tiled_s2r.get_slice(local_thread_idx);\n auto thr_r2g_QA = tiled_r2g_QA.get_slice(local_thread_idx);\n auto thr_r2g_SFA = tiled_r2g_SFA.get_slice(local_thread_idx);\n Tensor tQAsA = thr_s2r.partition_S(sA); // (Copy, Copy_M, Copy_N, PIPE)\n\n // Allocate temporary register tensors for copying quantization => output\n Tensor tQArA = make_tensor_like(\n make_layout(tQAsA(_, _, _, _0{}).shape())); // (Copy, Copy_M, Copy_N)\n Tensor tQAgQA = thr_r2g_QA.partition_S(gQA_mn);\n Tensor tQArQA = make_tensor_like(tQAgQA(_, _, _, _0{}, _0{}));\n\n Tensor tQAgSFA = thr_r2g_SFA.partition_S(gSFA_mn);\n Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{}));\n\n // Will result in barrier_id=10 passed to bar.sync instr as cutlass adds 8\n // in order to go over the reserved named barrier count.\n constexpr int row_quant_barrier_id = 2;\n cutlass::arch::NamedBarrier::sync(NumEpilogueRowQuantThreadCount, row_quant_barrier_id);\n\n int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler));\n float a_global_amax_val = shared_storage.global_a_amax[group_idx];\n // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release}\n static constexpr float fp4_max = 6.0f;\n static constexpr float fp8_max = 448.0f;\n static constexpr float fp4_max_inv = 1.0f / fp4_max;\n float global_encode_scale = a_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / a_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n\n float global_decode_scale = 1.0f / global_encode_scale;\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n auto sfa_converter = cutlass::NumericConverter{};\n do {\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0;\n k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) {\n int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler);\n\n int cur_group_idx = GetGroupIdx(&args, global_tile_n_offset);\n if (cur_group_idx != group_idx) {\n group_idx = cur_group_idx;\n a_global_amax_val = shared_storage.global_a_amax[group_idx];\n // Update group quantization parameters/scaling\n global_encode_scale = a_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / a_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n global_decode_scale = 1.0f / global_encode_scale;\n global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n }\n\n auto tQAgSFA_mn =\n tQAgSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto tQAgQA_mn = tQAgQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state);\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n copy(tiled_s2r, tQAsA(_, _, _, mainloop_pipe_consumer_state.index()), tQArA);\n cutlass::arch::fence_view_async_shared();\n mainloop_pipeline.consumer_release(mainloop_pipe_consumer_state);\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n\n // static int constexpr NumVecs = size(tQArA) / VectorSize;\n cutlass::maximum_absolute_value_reduction<\n cutlass::Array, true>\n amax_reduction;\n auto compute_frgs = reinterpret_cast *>(tQArA.data());\n auto output_frgs = reinterpret_cast *>(\n raw_pointer_cast(tQArQA.data()));\n Tensor amax =\n make_tensor(prepend(take<1, rank(tQArA)>(tQArA.shape()), _1{}));\n Tensor pvscales = make_tensor_like(amax);\n transformer_engine::curanddx::detail::philox4x32_native_state<\n NVTE_BUILD_NUM_PHILOX_ROUNDS>\n rng;\n if constexpr (kEnableStochasticRounding) {\n const size_t rng_sequence = global_thread_idx + k_tile * 512 +\n scheduler.get_linear_tile_idx() * K_TILE_MAX * 512 +\n tiles_in_m * tiles_in_n * K_TILE_MAX * 512;\n rng.init(rng_seed, rng_sequence, rng_offset);\n }\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < size<1>(group_modes<1, rank(tQArA)>(tQArA)); v++) {\n auto amax_view = group_modes<1, rank(amax)>(amax);\n auto pvscales_view = group_modes<1, rank(pvscales)>(pvscales);\n auto compute_frgs_up =\n cutlass::NumericArrayConverter{}(\n compute_frgs[v]);\n amax_view(_0{}, v) = amax_reduction(ElementAccumulator(0), compute_frgs_up);\n pvscales_view(_0{}, v) = cutlass::multiplies{}(\n amax_view(_0{}, v), global_encode_scale_multiplier);\n filter(tQArSFA)(v) = sfa_converter(pvscales_view(_0{}, v));\n auto qpvscale_ups =\n cutlass::NumericConverter{}(filter(tQArSFA)(v));\n auto qpvscale_scaled =\n cutlass::multiplies{}(qpvscale_ups, global_decode_scale);\n ElementAccumulator acc_scales;\n if constexpr (kUseFastMath) {\n // Fast math: compute approximate reciprocal\n acc_scales = cutlass::reciprocal_approximate_ftz{}(\n qpvscale_scaled);\n } else {\n // Accurate math: compute reciprocal with division\n acc_scales = cutlass::divides{}(1.0, qpvscale_scaled);\n }\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(\n acc_scales, cutlass::platform::numeric_limits::max());\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(\n cutlass::multiplies>{}(\n compute_frgs_up, acc_scale),\n *reinterpret_cast *>(&random_uint4));\n } else {\n output_frgs[v] =\n cutlass::NumericArrayConverter{}(\n cutlass::multiplies>{}(\n compute_frgs_up, acc_scale));\n }\n }\n copy(tiled_r2g_QA, tQArQA, tQAgQA_mn);\n copy(tiled_r2g_SFA, filter(tQArSFA), filter(tQAgSFA_mn));\n }\n // scheduler.advance();\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n }\n\n } else {\n cutlass::arch::warpgroup_reg_dealloc<32>();\n }\n } // sm100 compile guard end\n} // NOLINT(readability/fn_size)\n\ntemplate \nvoid group_row_col_rht_gemm_ntt_w_sfc(int packed_sequence_length, int hidden_size, TA const *A,\n TB const *B, TQA *QA, TSFA *SFA,\n MultiAmaxHadamardCastFusionArgs &args,\n const size_t *rng_state, uint32_t *tile_scheduler_workspace,\n uint32_t sm_count, cudaStream_t stream,\n int k_tile_size = 1024) {\n using namespace cute;\n static int constexpr SFVecSize = 16;\n static int constexpr RhtTensorSize = 16;\n\n static_assert(RhtTensorSize == 16, \"RhtTensorSize must be 16\");\n using LinearSFALayout = decltype(make_layout(make_shape(make_shape(Int{}, 0), 0),\n make_stride(make_stride(_0{}, _1{}), 0)));\n using LinearSFDLayout = decltype(make_layout(make_shape(0, make_shape(Int{}, 0)),\n make_stride(0, make_stride(_0{}, _1{}))));\n\n using SwizzledSFALayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFDLayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFALayout = decltype(tile_to_shape(\n SwizzledSFALayoutAtom{}, make_shape(hidden_size, packed_sequence_length), Step<_1, _2>{}));\n using SwizzledSFDLayout = decltype(tile_to_shape(\n SwizzledSFDLayoutAtom{}, make_shape(hidden_size, packed_sequence_length), Step<_2, _1>{}));\n\n using SFALayout = cute::conditional_t;\n using SFDLayout = cute::conditional_t;\n SFALayout sfa_layout;\n SFDLayout sfd_layout;\n\n if constexpr (kEnableSwizzleSFOutput) {\n sfa_layout = tile_to_shape(SwizzledSFALayoutAtom{},\n make_shape(hidden_size, packed_sequence_length), Step<_1, _2>{});\n sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{},\n make_shape(hidden_size, packed_sequence_length), Step<_2, _1>{});\n } else {\n sfa_layout = make_layout(\n make_shape(make_shape(Int{}, hidden_size / SFVecSize), packed_sequence_length),\n make_stride(make_stride(_0{}, _1{}), hidden_size / SFVecSize));\n sfd_layout = make_layout(\n make_shape(hidden_size, make_shape(Int{}, packed_sequence_length / SFVecSize)),\n make_stride(packed_sequence_length / SFVecSize, make_stride(_0{}, _1{})));\n }\n\n // Define shapes (dynamic)\n auto M = hidden_size;\n auto N = packed_sequence_length;\n Tensor tensorA = make_tensor(A, make_shape(hidden_size, packed_sequence_length), LayoutLeft{});\n Tensor tensorB = make_tensor(B, make_shape(RhtTensorSize, RhtTensorSize), LayoutLeft{});\n Tensor tensorQA = make_tensor(QA, make_shape(hidden_size, packed_sequence_length), LayoutLeft{});\n Tensor tensorSFA = make_tensor(SFA, sfa_layout);\n\n // Define strides (from tensors)\n auto dA = stride(tensorA); // (dM,dK)\n auto dB = stride(tensorB); // (dN,dK)\n auto dD = LayoutRight{}; // (dM,dN)\n auto dQA = stride(tensorQA); // (dM,dK)\n using ClusterShape = Shape<_1, _1, _1>;\n auto cluster_shape = ClusterShape{};\n auto cluster_tile_shape = Shape<_128, Int, Int>{};\n auto cluster_tile_mainloop = Shape<_128, Int, _128>{};\n\n // Each mainloop / epilogue loads 128 x 64 tiles while each MMA proceeds with 128 x 16 tiles\n static int constexpr EpilogueUnrollFactor =\n size<2>(cluster_tile_mainloop) / size<2>(cluster_tile_shape);\n // Construct the MMA\n auto mma = make_tiled_mma(\n SM100_MMA_F16BF16_SS(cluster_tile_shape), size<1>(cluster_tile_shape),\n UMMA::Major::MN, UMMA::Major::MN>{},\n Layout>{});\n\n // Assert that the TiledMMA uses all CTAs in the CGA.\n CUTE_STATIC_ASSERT_V(size(cluster_shape) == size(mma));\n CUTE_STATIC_ASSERT_V(evenly_divides(cluster_tile_shape, tile_shape(mma)));\n\n // Determine the A and B shapes\n auto mma_shape_B =\n partition_shape_B(mma, make_shape(size<1>(cluster_tile_shape), size<2>(cluster_tile_shape)));\n\n using TiledMma = decltype(mma);\n using AtomThrID = typename TiledMma::AtomThrID;\n\n using SmemShape_M = decltype(shape_div(\n shape<0>(cluster_tile_shape),\n shape_div(shape<0>(cluster_tile_shape), size<0>(cluster_tile_shape) / size(AtomThrID{}))));\n using SmemShape_N = decltype(shape_div(\n shape<1>(cluster_tile_shape),\n shape_div(shape<1>(cluster_tile_shape), size<1>(cluster_tile_shape) / size(AtomThrID{}))));\n using SmemShape_K = decltype(cute::get<2>(cluster_tile_shape));\n\n using SmemLayoutAtomB =\n decltype(cutlass::gemm::collective::detail::sm100_smem_selector());\n\n auto mma_shape_A = partition_shape_A(\n mma, make_shape(size<0>(cluster_tile_mainloop), size<2>(cluster_tile_mainloop)));\n using SmemShape_M_A =\n decltype(shape_div(shape<0>(cluster_tile_mainloop),\n shape_div(shape<0>(cluster_tile_mainloop),\n size<0>(cluster_tile_mainloop) / size(AtomThrID{}))));\n using SmemShape_K_A = decltype(cute::get<2>(cluster_tile_mainloop));\n using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<\n cute::UMMA::Major::MN, TA, SmemShape_M_A, SmemShape_K_A>());\n\n static uint32_t constexpr TotalTmemRows = 128;\n static uint32_t constexpr Sm100TmemCapacityColumns = 512;\n static uint32_t constexpr TotalTmem = TotalTmemRows * Sm100TmemCapacityColumns;\n static uint32_t constexpr AccumulatorPipelineStageCount =\n TotalTmem / (cute::size<0>(cluster_tile_shape) * cute::size<1>(cluster_tile_shape));\n\n // Define the smem layouts (static)\n // Calculate max pipeline stages based on Blackwell SM100's 232KB shared memory\n constexpr int SchedulerPipelineStageCount = 4;\n static int constexpr MainloopPipelineBytes = sizeof(\n typename cutlass::detail::CustomizedPipelineTmaUmmaAsync<1, Shape<_1, _1, _1>,\n Shape<_1, _1, _1>>::SharedStorage);\n\n static int constexpr SchedulerWorkspaceBytes = sizeof(int) * SchedulerPipelineStageCount;\n static int constexpr SchedulerThrottlePipelineBytes =\n sizeof(typename cutlass::PipelineAsync::SharedStorage);\n static int constexpr SchedulerPipelineBytes =\n sizeof(typename cutlass::PipelineCLCFetchAsync::SharedStorage);\n\n static int constexpr TmemDeallocBytes = sizeof(cutlass::arch::ClusterBarrier);\n static int constexpr BTensorBytes = cute::size(mma_shape_B) * sizeof(TB);\n static int constexpr AccPipelineBytes = sizeof(\n typename cutlass::PipelineUmmaAsync>::SharedStorage);\n static int constexpr TmemBasePtrsBytes = sizeof(uint32_t);\n static int constexpr kBlackwellSmemSize = 232448; // 232KB in bytes\n static int constexpr kBytesPerStage =\n cute::size(mma_shape_A) * sizeof(TA) + MainloopPipelineBytes;\n static int constexpr kReservedBytes = SchedulerWorkspaceBytes + SchedulerThrottlePipelineBytes +\n SchedulerPipelineBytes + TmemBasePtrsBytes +\n TmemDeallocBytes + BTensorBytes +\n AccPipelineBytes; // Reserve for barriers and other uses\n static int constexpr kMaxStages = (kBlackwellSmemSize - kReservedBytes) / kBytesPerStage;\n auto sP = Int{}; // SMEM pipelines\n\n auto sA = UMMA::tile_to_mma_shape(SmemLayoutAtomA{}, append(mma_shape_A, sP),\n Step<_2, _1, _3>{}); // (MMA,MMA_M,MMA_K,PIPE)\n auto sB = UMMA::tile_to_mma_shape(SmemLayoutAtomB{},\n append(mma_shape_B, _1{})); // (MMA,MMA_N,MMA_K, _1)\n auto sD = Layout<_1>{}; // XXX Dummy\n\n auto tma_load_a =\n make_tma_copy_A_sm100(SM90_TMA_LOAD{}, tensorA, sA(_, _, _, 0), cluster_tile_mainloop, mma);\n auto tma_load_b =\n make_tma_copy_B_sm100(SM90_TMA_LOAD{}, tensorB, sB(_, _, _, 0), cluster_tile_shape, mma);\n\n // Assert checks on tile sizes -- no predication\n assert(M % size<0>(cluster_tile_shape) == 0);\n assert(N % size<1>(cluster_tile_shape) == 0);\n\n dim3 dimBlock(512);\n dim3 dimCluster(size<0>(cluster_shape), size<1>(cluster_shape), size<2>(cluster_shape));\n dim3 dimGrid(sm_count, 1, 1);\n\n int smem_size = sizeof(\n SharedStorage);\n\n auto *kernel_ptr = &group_row_col_rht_gemm_device<\n decltype(M), decltype(N), decltype(k_tile_size), decltype(cluster_shape),\n decltype(cluster_tile_shape), TA, decltype(dA), decltype(sA), decltype(tma_load_a), TB,\n decltype(dB), decltype(sB), decltype(tma_load_b), TD, decltype(dD), decltype(sD), TSFD,\n decltype(sfd_layout), TQA, decltype(dQA), TSFA, decltype(sfa_layout), decltype(mma),\n AccumulatorPipelineStageCount, SchedulerPipelineStageCount, kEnableStochasticRounding,\n kEnableRHTColQuant, kEnableRowQuant, kEnableSwizzleSFOutput, kUseFastMath>;\n\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(*kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n\n // Set workspace and set to zero\n NVTE_CHECK_CUDA(cudaMemsetAsync(reinterpret_cast(tile_scheduler_workspace), 0,\n sizeof(uint32_t), stream));\n\n // Launch kernel\n cutlass::ClusterLaunchParams params = {dimGrid, dimBlock, dimCluster, smem_size, stream};\n cutlass::Status status = cutlass::launch_kernel_on_cluster(\n params, (void const *)kernel_ptr, M, N, k_tile_size, cluster_shape, cluster_tile_shape, A, dA,\n sA, tma_load_a, B, dB, sB, tma_load_b, QA, dQA, SFA, sfa_layout, args,\n tile_scheduler_workspace, mma, rng_state);\n NVTE_CHECK_CUDA(cudaGetLastError());\n NVTE_CHECK(status == cutlass::Status::kSuccess, \"Kernel launch failed.\");\n}\n\n} // namespace\n} // namespace detail\n\nvoid group_hadamard_transform_cast_fusion(const Tensor &input_, std::vector &output_list,\n const size_t *split_sections, size_t num_tensors,\n const Tensor &hadamard_matrix_,\n QuantizationConfig &quant_config, Tensor &quant_workspace,\n cudaStream_t stream) {\n NVTE_API_CALL(group_hadamard_transform_cast_fusion);\n\n using transformer_engine::detail::kMaxTensorsPerKernel;\n using transformer_engine::detail::MultiAmaxHadamardCastFusionArgs;\n\n NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16,\n \"Input tensor must be BF16 tensor, but dtype is \", to_string(input_.dtype()), \".\");\n NVTE_CHECK(input_.dim() >= 2, \"Input must be a 2D tensor.\");\n const SimpleTensor &input = input_.data;\n\n NVTE_CHECK(output_list.size() == num_tensors,\n \"Number of output tensors should match number of tensors.\");\n\n NVTE_CHECK(num_tensors <= kMaxTensorsPerKernel,\n \"Number of tensors should be less than or equal to \", kMaxTensorsPerKernel);\n\n // construct the multi-tensor args\n MultiAmaxHadamardCastFusionArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.split_sections_range[0] = 0;\n bool all_has_row_quant = true;\n bool all_has_col_quant = true;\n void *rowwise_data_base_ptr = nullptr;\n void *rowwise_scale_inv_base_ptr = nullptr;\n for (size_t i = 0; i < num_tensors; ++i) {\n NVTE_CHECK(split_sections[i] % 128 == 0, \"component \", i,\n \" of split_sections should be 128 multiple\");\n if (split_sections[i] == 0) {\n continue;\n }\n bool has_row_quant = output_list[i]->data.dptr != nullptr;\n bool has_col_quant = output_list[i]->columnwise_data.dptr != nullptr;\n all_has_row_quant = all_has_row_quant && has_row_quant;\n all_has_col_quant = all_has_col_quant && has_col_quant;\n // sanity check, the two bool flags cannot be both false\n NVTE_CHECK(has_row_quant || has_col_quant,\n \"At least one of the output tensors must have row or column quant.\");\n void *amax_rowwise_ptr =\n has_row_quant ? reinterpret_cast(output_list[i]->amax.dptr) : nullptr;\n void *amax_colwise_ptr =\n has_col_quant ? reinterpret_cast(output_list[i]->columnwise_amax.dptr) : nullptr;\n void *rowwise_data_ptr =\n has_row_quant ? reinterpret_cast(output_list[i]->data.dptr) : nullptr;\n void *rowwise_scale_inv_ptr =\n has_row_quant ? reinterpret_cast(output_list[i]->scale_inv.dptr) : nullptr;\n if (all_has_row_quant &&\n (rowwise_data_base_ptr == nullptr || rowwise_scale_inv_base_ptr == nullptr)) {\n rowwise_data_base_ptr = rowwise_data_ptr;\n rowwise_scale_inv_base_ptr = rowwise_scale_inv_ptr;\n }\n void *output_colwise_ptr =\n has_col_quant ? reinterpret_cast(output_list[i]->columnwise_data.dptr) : nullptr;\n void *output_colwise_scale_inv_ptr =\n has_col_quant ? reinterpret_cast(output_list[i]->columnwise_scale_inv.dptr)\n : nullptr;\n kernel_args.global_a_amax_list[kernel_args.num_tensors] = amax_rowwise_ptr;\n kernel_args.global_d_amax_list[kernel_args.num_tensors] = amax_colwise_ptr;\n kernel_args.output_colwise_list[kernel_args.num_tensors] = output_colwise_ptr;\n kernel_args.output_colwise_scale_inv_list[kernel_args.num_tensors] =\n output_colwise_scale_inv_ptr;\n kernel_args.split_sections[kernel_args.num_tensors] = split_sections[i];\n kernel_args.split_sections_range[kernel_args.num_tensors + 1] =\n kernel_args.split_sections_range[kernel_args.num_tensors] + split_sections[i];\n kernel_args.num_tensors++;\n }\n\n // Stochastic rounding config\n const bool use_stochastic_rounding = quant_config.stochastic_rounding;\n const size_t *rng_state = nullptr;\n if (use_stochastic_rounding) {\n NVTE_CHECK(quant_config.rng_state != nullptr,\n \"Enabled stochastic rounding without providing RNG state\");\n const Tensor &rng_state_tensor = *convertNVTETensorCheck(quant_config.rng_state);\n NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64,\n \"RNG state should contain 2 64-bit values.\");\n NVTE_CHECK(rng_state_tensor.data.shape == std::vector{2},\n \"Shape of the RNG state should be [2], but got \", rng_state_tensor.data.shape);\n rng_state = reinterpret_cast(rng_state_tensor.data.dptr);\n }\n\n uint32_t *tile_scheduler_workspace = nullptr;\n NVTE_CHECK(quant_workspace.data.dptr != nullptr, \"Quantization workspace must be provided.\");\n NVTE_CHECK(quant_workspace.data.buffer_size_bytes() >= sizeof(uint32_t),\n \"Quantization workspace must be at least 4 bytes.\");\n tile_scheduler_workspace = reinterpret_cast(quant_workspace.data.dptr);\n\n // Template arguments\n using TA = cute::bfloat16_t;\n using TB = cute::bfloat16_t;\n using TD = cutlass::float_e2m1_t;\n using TSFD = cutlass::float_ue4m3_t;\n using TQA = TD;\n using TSFA = TSFD;\n\n checkCuDriverContext(stream);\n\n // Check Hadamard matrix\n constexpr int kHadamardDimension = 16;\n\n NVTE_CHECK(hadamard_matrix_.dtype() == transformer_engine::DType::kBFloat16,\n \"Hadamard matrix must be BF16 tensor, but dtype is \",\n to_string(hadamard_matrix_.dtype()), \".\");\n const SimpleTensor &hadamard_matrix = hadamard_matrix_.data;\n NVTE_CHECK(\n (hadamard_matrix_.shape() == std::vector{kHadamardDimension, kHadamardDimension}),\n \"Hadamard matrix must have shape=\",\n std::vector{kHadamardDimension, kHadamardDimension},\n \", but got shape=\", hadamard_matrix_.shape(), \".\");\n const size_t hadamard_dimension = hadamard_matrix.shape[0];\n\n const size_t ndim = input.shape.size();\n const size_t n = input.shape[ndim - 1];\n size_t m = 1;\n for (size_t i = 0; i < ndim - 1; ++i) {\n m *= input.shape[i];\n }\n\n auto sm_count = transformer_engine::cuda::sm_count();\n\n NVTE_CHECK(n % hadamard_dimension == 0, \"row_length must be divisible by hadamard_dimension.\");\n\n NVTE_CHECK(m % hadamard_dimension == 0, \"num_rows must be divisible by hadamard_dimension\");\n\n int k_tile_size = 1024;\n\n const bool use_swizzle_sf_output = false;\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_stochastic_rounding, kEnableStochasticRounding,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n all_has_col_quant, kEnableRhtColQuant,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n all_has_row_quant, kEnableRowQuant,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_swizzle_sf_output, kEnableSwizzleSFOutput,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n quant_config.use_fast_math, kUseFastMath,\n\n if constexpr (kEnableRhtColQuant || kEnableRowQuant) {\n detail::group_row_col_rht_gemm_ntt_w_sfc<\n kEnableStochasticRounding, kEnableRhtColQuant, kEnableRowQuant,\n kEnableSwizzleSFOutput, TA, TB, TQA, TSFA, TD, TSFD, kUseFastMath>(\n /*packed_sequence_length=*/m, /*hidden_size=*/n,\n /*A=*/reinterpret_cast(input.dptr),\n /*B=*/reinterpret_cast(hadamard_matrix.dptr),\n /*QA=*/reinterpret_cast(rowwise_data_base_ptr),\n /*SFA=*/reinterpret_cast(rowwise_scale_inv_base_ptr),\n /*args=*/kernel_args,\n /*rng_state=*/rng_state,\n /*tile_scheduler_workspace=*/tile_scheduler_workspace,\n /*sm_count=*/sm_count,\n /*stream=*/stream, /*k_tile_size=*/k_tile_size);\n } else {\n NVTE_ERROR(\"Invalid kernel configuration (kEnableRHTColQuant=\",\n kEnableRhtColQuant, \", kEnableRowQuant=\", kEnableRowQuant, \").\");\n }\n\n );););););\n}\n\n} // namespace transformer_engine\n\nvoid nvte_group_hadamard_transform_cast_fusion(const NVTETensor input, NVTETensor *outputs,\n const NVTETensor hadamard_matrix,\n const size_t *split_sections,\n const size_t num_tensors,\n const NVTEQuantizationConfig quant_config,\n NVTETensor quant_workspace, cudaStream_t stream) {\n NVTE_API_CALL(nvte_group_hadamard_transform_cast_fusion);\n using namespace transformer_engine;\n NVTE_CHECK(num_tensors > 0, \"Number of tensors should be greater than 0.\");\n\n Tensor *input_tensor = convertNVTETensorCheck(input);\n std::vector output_list(num_tensors);\n for (size_t i = 0; i < num_tensors; ++i) {\n output_list[i] = convertNVTETensorCheck(outputs[i]);\n }\n\n Tensor *quant_workspace_tensor = convertNVTETensorCheck(quant_workspace);\n\n QuantizationConfig quant_config_cpp;\n if (quant_config != nullptr) {\n quant_config_cpp = *reinterpret_cast(quant_config);\n }\n\n // Call the multi-tensor Hadamard transform amax implementation.\n group_hadamard_transform_cast_fusion(*input_tensor, output_list, split_sections, num_tensors,\n *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp,\n *quant_workspace_tensor, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/util/hip/hip_runtime.h\"\n#include \"common/util/curanddx.hpp\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"customized_pipeline.cuh\"\n#include \"cutlass/arch/barrier.h\"\n#include \"cutlass/arch/reg_reconfig.h\"\n#include \"cutlass/cluster_launch.hpp\"\n#include \"cutlass/cutlass.h\"\n#include \"cutlass/detail/sm100_blockscaled_layout.hpp\"\n#include \"cutlass/fast_math.h\"\n#include \"cutlass/float8.h\"\n#include \"cutlass/float_subbyte.h\"\n#include \"cutlass/gemm/collective/builders/sm100_common.inl\"\n#include \"cutlass/numeric_conversion.h\"\n#include \"cutlass/numeric_types.h\"\n#include \"cutlass/pipeline/pipeline.hpp\"\n#include \"cutlass/platform/platform.h\"\n#include \"cutlass/util/GPU_Clock.hpp\"\n#include \"cutlass/util/command_line.h\"\n#include \"cutlass/util/print_error.hpp\"\n\nnamespace transformer_engine {\nnamespace detail {\nnamespace {\n\nusing namespace cute;\n\n// Ensure Tensor refers to cute::Tensor, not transformer_engine::Tensor\nusing cute::Tensor;\n\nconstexpr int kMaxTensorsPerKernel = 64;\n\nstruct MultiAmaxHadamardCastFusionArgs {\n // (output) Amax buffer for input A amax buffer\n void *global_a_amax_list[kMaxTensorsPerKernel];\n // (output) Amax buffer for pre-RHT amax buffer\n void *global_d_amax_list[kMaxTensorsPerKernel];\n // output D pointers for each tensor\n void *output_colwise_list[kMaxTensorsPerKernel];\n // output SFD inverse pointers for each tensor\n void *output_colwise_scale_inv_list[kMaxTensorsPerKernel];\n // split sections of each tensor of input\n int split_sections[kMaxTensorsPerKernel];\n // Prefix sum (with leading zero) of split_sections of each tensor of input\n int split_sections_range[kMaxTensorsPerKernel + 1];\n\n // Number of tensors (splits) being processed by kernel\n int num_tensors;\n};\n\n__device__ __forceinline__ int GetGroupIdx(MultiAmaxHadamardCastFusionArgs *kernel_args_ptr,\n int offset) {\n // Check the kernel args and get the corresponding id\n const int num_tensors = kernel_args_ptr->num_tensors;\n if (offset >= kernel_args_ptr->split_sections_range[num_tensors]) {\n return num_tensors - 1;\n }\n int group_idx = 0;\n while (kernel_args_ptr->split_sections_range[group_idx + 1] <= offset) {\n ++group_idx;\n }\n return group_idx;\n}\n\nCUTLASS_DEVICE\ncutlass::Array StochasticNumericConverterBase(\n cutlass::Array const &input, cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n auto output_ptr = reinterpret_cast(&output);\n constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING;\n if constexpr (has_rs) {\n asm volatile(\n \"{\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\\n\"\n \"}\"\n : \"=h\"(output_ptr[0]), \"=h\"(output_ptr[1])\n : \"f\"(input[0]), \"f\"(input[1]), \"f\"(input[2]), \"f\"(input[3]), \"f\"(input[4]), \"f\"(input[5]),\n \"f\"(input[6]), \"f\"(input[7]), \"r\"(rbits[0]), \"r\"(rbits[1]));\n } else {\n NVTE_DEVICE_ERROR(\n \"FP4 cvt PTX instructions are architecture-specific. \"\n \"Try recompiling with sm_XXXa instead of sm_XXX.\");\n }\n return output;\n}\n\nCUTLASS_DEVICE\ncutlass::Array StochasticNumericConverter(\n cutlass::Array const &input, cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n cutlass::Array *result_ptr =\n reinterpret_cast *>(&output);\n cutlass::Array const *source_ptr =\n reinterpret_cast const *>(&input);\n cutlass::Array const *rbits_ptr =\n reinterpret_cast const *>(&rbits);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 2; i++) {\n result_ptr[i] = StochasticNumericConverterBase(source_ptr[i], rbits_ptr[i]);\n }\n return output;\n}\n\ntemplate \nstruct SharedStorage {\n static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;\n static int constexpr EpilogueUnrollFactor = EpilogueUnrollFactor_;\n using AtomThrShapeMNK = cute::Shape<_1, _1, _1>;\n\n using AccumulatorPipeline =\n cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;\n\n static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline =\n cutlass::detail::CustomizedPipelineTmaUmmaAsync,\n AtomThrShapeMNK>;\n using MainloopPipelineStorage = typename MainloopPipeline::SharedStorage;\n using SchedPipeline = cutlass::PipelineCLCFetchAsync;\n using SchedPipelineStorage = typename SchedPipeline::SharedStorage;\n using SchedThrottlePipeline = cutlass::PipelineAsync;\n using SchedThrottlePipelineStorage = typename SchedThrottlePipeline::SharedStorage;\n\n struct TensorStorage : cute::aligned_struct<128, _1> {\n cute::array_aligned> smem_A;\n cute::array_aligned> smem_B;\n } tensors;\n\n alignas(16) AccumulatorPipelineStorage accumulator;\n alignas(16) MainloopPipelineStorage mainloop;\n alignas(16) cute::uint64_t tma_barrier[1];\n alignas(16) SchedPipelineStorage sched;\n alignas(16) SchedThrottlePipelineStorage sched_throttle;\n alignas(16) int32_t atomic_tile_id[SchedulerPipelineStageCount_];\n alignas(16) float global_a_amax[kMaxTensorsPerKernel];\n alignas(16) float global_d_amax[kMaxTensorsPerKernel];\n uint32_t atomic_tile_counter[SchedulerPipelineStageCount_];\n uint32_t tmem_base_ptr;\n};\n\n// Main RHT GEMM kernel entry -- highly templated for flexible architecture/config support\ntemplate \n__launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device(\n MShape M, NShape packed_N, KShape K, ClusterShape cluster_shape, ClusterTileShape cluster_tile,\n TA const *A, AStride dA, ASmemLayout sAlayout, CUTE_GRID_CONSTANT TmaLoadA const tma_load_a,\n TB const *B, BStride dB, BSmemLayout sBlayout, CUTE_GRID_CONSTANT TmaLoadB const tma_load_b,\n TQA *QA, QAStride dQA, TSFA *SFA, TSFALayout sfa_layout, MultiAmaxHadamardCastFusionArgs args,\n uint32_t *tile_scheduler_workspace, TiledMMA mma, const size_t *rng_state) {\n using namespace cute;\n\n // Abort immediately if compilation is not supported\n constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY;\n if constexpr (!is_blackwell_arch) {\n NVTE_DEVICE_ERROR(\"RHT fusion is only supported on Blackwell.\");\n return;\n } else {\n static_assert(kEnableRHTColQuant_ || kEnableRowQuant_,\n \"group_row_col_rht_gemm_device must generate row-wise \"\n \"and/or column-wise output.\");\n#if !defined(CUTLASS_ARCH_CLC_ENABLED)\n CUTLASS_NOT_IMPLEMENTED();\n return;\n#endif\n\n using X = Underscore;\n // Accumulator data type for main computation\n using ElementAccumulator = float;\n static int constexpr K_PIPE_MAX = size<3>(ASmemLayout{});\n using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>;\n static uint32_t constexpr kTmaTransactionBytes = cutlass::bits_to_bytes(\n size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v);\n static constexpr bool kEnableStochasticRounding = kEnableStochasticRounding_;\n static constexpr bool kEnableRHTColQuant = kEnableRHTColQuant_;\n static constexpr bool kEnableRowQuant = kEnableRowQuant_;\n static constexpr bool kEnableSwizzleSFOutput = kEnableSwizzleSFOutput_;\n static constexpr bool kUseFastMath = kUseFastMath_;\n\n // Constant for RHT tensor processing (tile size etc)\n static int constexpr RhtTensorSize = 16;\n\n // Transaction bytes for TMA transfer on RHT tensor blocks\n static int constexpr kTmaRhtTensorTransactionBytes =\n cutlass::bits_to_bytes(RhtTensorSize * RhtTensorSize * cute::sizeof_bits_v);\n static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;\n static int constexpr SchedulerPipelineStageCount = SchedulerPipelineStageCount_;\n\n // Mainloop pipeline stage calculation, vectorization parameters for scaling factors\n static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{});\n static int constexpr SFVecSize = 16;\n // Swizzle output layout for scaling factor arrays\n using SwizzledSFALayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFDLayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n\n // Mainloop pipeline types for TMA async execution and epilogue cluster scheduling\n using MainloopPipeline =\n cutlass::detail::CustomizedPipelineTmaUmmaAsync;\n using MainloopPipelineState = typename MainloopPipeline::PipelineState;\n using SchedPipeline = cutlass::PipelineCLCFetchAsync;\n using SchedPipelineState = typename SchedPipeline::PipelineState;\n using SchedThrottlePipeline = cutlass::PipelineAsync;\n using SchedThrottlePipelineState = typename SchedThrottlePipeline::PipelineState;\n\n static_assert(ClusterShape{} == Shape<_1, _1, _1>{}, \"ClusterShape must be Shape<_1,_1,_1>\");\n\n using TmemAllocator = cute::TMEM::Allocator1Sm;\n static int constexpr VectorSize = RhtTensorSize;\n\n // Compile-time safety: static shapes required for shared memory layouts\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n // CUTE_STATIC_ASSERT(is_static::value);\n\n auto cluster_size = size<0>(cluster_shape);\n auto mainloop_tiler = Shape<_128, _16, _128>{};\n auto epilogue_tiler = Shape<_128, _128, _128>{};\n\n static int constexpr EpilogueUnrollFactor = size<2>(epilogue_tiler) / size<2>(cluster_tile);\n\n // Get the appropriate blocks for this Cluster\n dim3 cluster_coord_in_grid = cluster_id_in_grid();\n\n // Total number of k-tiles\n int const K_TILE_MAX = min(packed_N, K) / size<2>(epilogue_tiler);\n\n struct TileScheduler {\n uint32_t tiles_in_m = 0;\n uint32_t tiles_in_n = 0;\n uint32_t linear_idx = 0;\n uint32_t next_linear_idx = 0;\n uint32_t start_idx = 0;\n uint32_t tile_m_idx = 0;\n uint32_t tile_n_idx = 0;\n int k_tile_max = 0;\n uint32_t *atomic_tile_index_;\n uint32_t *smem_tile_counter;\n uint32_t atomic_offset;\n cutlass::FastDivmodU64 divmod_tiles_in_m;\n\n CUTLASS_DEVICE TileScheduler(uint32_t tiles_m, uint32_t tiles_n, int kmax,\n uint32_t *atomic_tile_index, uint32_t *smem_tile_counter)\n : tiles_in_m(tiles_m),\n tiles_in_n(tiles_n),\n linear_idx(blockIdx.x),\n next_linear_idx(blockIdx.x),\n start_idx(blockIdx.x),\n k_tile_max(kmax),\n atomic_tile_index_(atomic_tile_index),\n smem_tile_counter(smem_tile_counter),\n atomic_offset(gridDim.x),\n divmod_tiles_in_m(uint64_t(tiles_m)) {\n update_tile_idx();\n }\n CUTLASS_DEVICE void update_tile_idx() {\n uint64_t q, r;\n divmod_tiles_in_m(q, r, uint64_t(linear_idx));\n tile_m_idx = static_cast(r);\n tile_n_idx = static_cast(q) * uint32_t(k_tile_max);\n }\n CUTLASS_DEVICE uint32_t tile_m() const { return tile_m_idx; }\n CUTLASS_DEVICE uint32_t tile_n_base() const { return tile_n_idx; }\n CUTLASS_DEVICE uint32_t tiles_m() const { return tiles_in_m; }\n\n CUTLASS_DEVICE uint32_t tiles_n() const { return tiles_in_n; }\n\n CUTLASS_DEVICE bool is_valid() const {\n return cute::elem_less(cute::make_coord(tile_m(), tile_n_base()),\n cute::make_coord(tiles_in_m, tiles_in_n));\n }\n\n CUTLASS_DEVICE bool is_first_wave() const { return linear_idx == start_idx; }\n\n CUTLASS_DEVICE uint32_t get_linear_tile_idx() const { return linear_idx; }\n\n // Fetch a new tile_id using atomics.\n CUTLASS_DEVICE uint32_t fetch_tile_id_counter(int pred) {\n uint32_t tile_id_counter = 0;\n asm volatile(\n \"{\\n\\t\"\n \".reg .pred p;\\n\\t\"\n \"setp.eq.u32 p, %2, 1;\\n\\t\"\n \"@p atom.global.add.u32 %0, [%1], 1; \\n\\t\"\n \"}\"\n : \"=r\"(tile_id_counter)\n : \"l\"(atomic_tile_index_), \"r\"(pred));\n\n return tile_id_counter;\n }\n\n CUTLASS_DEVICE auto fetch_next_work(SchedPipeline &sched_pipeline,\n SchedPipelineState sched_pipeline_consumer_state) {\n sched_pipeline.consumer_wait(sched_pipeline_consumer_state);\n next_linear_idx = smem_tile_counter[sched_pipeline_consumer_state.index()];\n cutlass::arch::fence_view_async_shared();\n sched_pipeline.consumer_release(sched_pipeline_consumer_state);\n return;\n }\n\n CUTLASS_DEVICE auto advance_to_next_work(SchedPipeline &sched_pipeline,\n SchedPipelineState sched_pipeline_producer_state) {\n uint32_t mbarrier_addr = sched_pipeline.producer_get_barrier(sched_pipeline_producer_state);\n // Wait for clcID buffer to become empty with a flipped phase\n sched_pipeline.producer_acquire(sched_pipeline_producer_state);\n auto is_leading_thread = cute::elect_one_sync();\n uint32_t tile_id_counter = fetch_tile_id_counter(is_leading_thread) + atomic_offset;\n uint32_t smem_addr =\n cute::cast_smem_ptr_to_uint(&smem_tile_counter[sched_pipeline_producer_state.index()]);\n if (is_leading_thread) {\n cute::store_shared_remote(tile_id_counter, smem_addr, mbarrier_addr, 0);\n }\n\n ++sched_pipeline_producer_state;\n return sched_pipeline_producer_state;\n }\n\n CUTLASS_DEVICE auto update_work_tile_info() {\n linear_idx = next_linear_idx;\n update_tile_idx();\n return;\n }\n };\n\n // Allocate and alias shared memory to the kernel's shared storage type\n extern __shared__ char shared_memory[];\n using SharedStorage =\n SharedStorage;\n SharedStorage &shared_storage = *reinterpret_cast(shared_memory);\n\n // Compute the number of tiles in M and N after tiling and assign scheduler\n uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile))));\n uint32_t tiles_in_n = uint32_t(\n size(ceil_div(args.split_sections_range[args.num_tensors], size<2>(epilogue_tiler))));\n\n TileScheduler scheduler(tiles_in_m, tiles_in_n, K_TILE_MAX, tile_scheduler_workspace,\n shared_storage.atomic_tile_counter);\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n\n // Shapes for accumulated tiles in mainloop and epilogue\n auto acc_shape_mma = make_shape(take<0, 2>(mainloop_tiler), _1{}, _1{});\n auto acc_shape_epilogue = make_shape(take<0, 2>(epilogue_tiler), _1{}, _1{});\n\n // Shape of the accumulator fragment for the main loop pipeline, with pipeline stages appended\n auto acc_mainloop_pipelined_shape = append(acc_shape_mma, Int{});\n auto bulk_tmem_mma = TiledMMA::make_fragment_C(acc_mainloop_pipelined_shape);\n\n // Number of threads assigned for various epilogue roles depending on quantization settings\n static int constexpr NumEpilogueColQuantThreadCount = kEnableRHTColQuant ? 128 : 0;\n static int constexpr NumEpilogueRowQuantThreadCount = kEnableRowQuant ? 256 : 0;\n static int constexpr NumMmaThreadCount = kEnableRHTColQuant ? 32 : 0;\n static int constexpr NumMmaIssueThreadCount = kEnableRHTColQuant ? 1 : 0;\n static int constexpr NumSchedThreads = 32;\n static int constexpr NumMainloopLoadThreads = 32;\n static int constexpr NumEpilogueThreads =\n NumEpilogueColQuantThreadCount + NumEpilogueRowQuantThreadCount;\n\n TmemAllocator tmem_allocator{};\n cutlass::arch::NamedBarrier tmem_allocation_result_barrier(\n NumMmaThreadCount + NumEpilogueColQuantThreadCount,\n cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);\n\n int warp_idx = cutlass::canonical_warp_idx_sync();\n\n // warp assignment\n bool is_mma_warp = (warp_idx == 0);\n bool is_dma_warp = (warp_idx == 1);\n bool is_sched_warp = (warp_idx == 2);\n bool is_epilogue_col_quant_warp = (warp_idx >= 4 && warp_idx <= 7);\n bool is_epilogue_row_quant_warp = (warp_idx >= 8 && warp_idx <= 15);\n\n typename MainloopPipeline::Params mainloop_pipeline_params;\n if (is_dma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;\n }\n if (is_mma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;\n }\n mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp;\n mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes;\n mainloop_pipeline_params.initializing_warp = 0;\n mainloop_pipeline_params.num_consumers =\n NumEpilogueRowQuantThreadCount + NumMmaIssueThreadCount;\n\n MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params,\n cluster_shape, cute::true_type{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n\n MainloopPipelineState mainloop_pipe_consumer_state;\n MainloopPipelineState mainloop_pipe_producer_state =\n cutlass::make_producer_start_state();\n\n using AccumulatorPipeline =\n cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState;\n using AccumulatorPipelineInitBarriers = cute::bool_constant;\n\n AccumulatorPipelineState accumulator_pipe_consumer_state;\n AccumulatorPipelineState accumulator_pipe_producer_state =\n cutlass::make_producer_start_state();\n\n typename AccumulatorPipeline::Params accumulator_pipeline_params;\n if (is_mma_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer;\n }\n if (is_epilogue_col_quant_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer;\n }\n // Only one producer thread arrives on this barrier.\n accumulator_pipeline_params.producer_arv_count = 1;\n accumulator_pipeline_params.consumer_arv_count =\n size(AtomThrShapeMNK{}) * NumEpilogueColQuantThreadCount;\n accumulator_pipeline_params.initializing_warp = 1;\n AccumulatorPipeline accumulator_pipeline(\n shared_storage.accumulator, accumulator_pipeline_params, cluster_shape,\n AccumulatorPipelineInitBarriers{}, cute::true_type{}); // Delay mask calculation\n typename SchedPipeline::Params sched_pipeline_params;\n if (is_sched_warp) {\n sched_pipeline_params.role = SchedPipeline::ThreadCategory::ProducerConsumer;\n } else {\n sched_pipeline_params.role = SchedPipeline::ThreadCategory::Consumer;\n }\n sched_pipeline_params.producer_blockid = 0;\n sched_pipeline_params.producer_arv_count = 1;\n sched_pipeline_params.consumer_arv_count =\n NumSchedThreads +\n cluster_size * (NumMainloopLoadThreads + NumEpilogueThreads + NumMmaThreadCount);\n sched_pipeline_params.transaction_bytes = sizeof(uint32_t);\n sched_pipeline_params.initializing_warp = 3;\n SchedPipeline sched_pipeline(shared_storage.sched, sched_pipeline_params, cluster_shape);\n SchedPipelineState sched_pipeline_consumer_state;\n SchedPipelineState sched_pipeline_producer_state =\n cutlass::make_producer_start_state();\n\n typename SchedThrottlePipeline::Params sched_throttle_pipeline_params;\n if (is_dma_warp) {\n sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Producer;\n }\n if (is_sched_warp) {\n sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Consumer;\n }\n sched_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads;\n sched_throttle_pipeline_params.consumer_arv_count = NumSchedThreads;\n sched_throttle_pipeline_params.dst_blockid = 0;\n sched_throttle_pipeline_params.initializing_warp = 4;\n\n SchedThrottlePipeline sched_throttle_pipeline(shared_storage.sched_throttle,\n sched_throttle_pipeline_params);\n SchedThrottlePipelineState sched_pipeline_throttle_consumer_state;\n SchedThrottlePipelineState sched_pipeline_throttle_producer_state =\n cutlass::make_producer_start_state();\n\n if (warp_idx == 2 && elect_one_sync()) {\n cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);\n }\n __syncthreads();\n\n // Warp group roles: DMA (global->shared copy), MMA (tensor core gemm), scheduler, column quantizer, row quantizer\n if (is_dma_warp) {\n // Warp responsible for loading input from global to shared memory using TMA (Tensor Memory Access).\n cutlass::arch::warpgroup_reg_dealloc<32>();\n // Get TMA tensors for input matrix A and B (Hadamard/transform matrix) from global memory.\n Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, packed_N));\n Tensor mB = tma_load_b.get_tma_tensor(make_shape(RhtTensorSize, RhtTensorSize));\n\n // Partition tensors for tiling according to the mainloop and cluster tilers.\n Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{});\n Tensor gB_nk =\n local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k)\n\n // Shared memory tensors for pipeline\n Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()),\n sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()),\n sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n // Determine warp/tile positioning\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n // Partition global to local fragments for A and B\n Tensor tCgA = thr_mma.partition_A(gA_mk); // (MMA,MMA_M,MMA_K,k)\n Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k)\n\n Layout cta_layout_mnk = make_layout(cluster_shape);\n Layout cta_layout_vmnk =\n tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{}));\n auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster);\n\n auto [tAgA, tAsA] =\n tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),\n group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA));\n\n auto [tBgB, tBsB] =\n tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),\n group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB));\n\n uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);\n uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);\n if constexpr (kEnableRHTColQuant) {\n if (elect_one_sync()) {\n cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0],\n kTmaRhtTensorTransactionBytes);\n copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0),\n tBsB(_, 0));\n }\n }\n\n do {\n // is_first_wave indicates whether this scheduler wave is the first among a group.\n bool is_first_wave = scheduler.is_first_wave();\n uint32_t skip_wait = is_first_wave;\n auto tAgA_mk = tAgA(_, scheduler.tile_m(), _);\n int k_tile = 0;\n\n sched_throttle_pipeline.producer_acquire(sched_pipeline_throttle_producer_state);\n sched_throttle_pipeline.producer_commit(sched_pipeline_throttle_producer_state);\n ++sched_pipeline_throttle_producer_state;\n CUTLASS_PRAGMA_NO_UNROLL\n while (k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n()) {\n int k_tile_idx_n = scheduler.tile_n_base() + k_tile;\n ++k_tile;\n skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount);\n mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state);\n using BarrierType = typename MainloopPipeline::ProducerBarrierType;\n BarrierType *tma_barrier =\n mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state);\n int write_stage = mainloop_pipe_producer_state.index();\n ++mainloop_pipe_producer_state;\n if (cute::elect_one_sync()) {\n copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n),\n tAsA(_, write_stage));\n }\n }\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n // scheduler.advance();\n } while (scheduler.is_valid());\n mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);\n } else if (is_mma_warp) {\n // This warp executes the main tensor core matrix-multiply-accumulate for the Hadamard transform.\n cutlass::arch::warpgroup_reg_dealloc<32>();\n if constexpr (kEnableRHTColQuant) {\n // Setup shared memory fragments for A and B tiles.\n Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()),\n sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()),\n sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n // Allocate \"fragments\" -- these are actually umma smem descriptors\n Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)\n Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE)\n\n mma.accumulate_ = UMMA::ScaleOut::Zero;\n\n tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns,\n &shared_storage.tmem_base_ptr);\n __syncwarp();\n tmem_allocation_result_barrier.arrive();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_mma.data() = tmem_base_ptr;\n // Wait until the B (Hadamard) tensor copy is complete\n cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/);\n do {\n uint32_t skip_wait = K_TILE_MAX <= 0;\n\n auto barrier_token =\n mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0;\n k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) {\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n int read_stage = mainloop_pipe_consumer_state.index();\n auto tCrA_mk = tCrA(_, _, _, read_stage);\n auto tCrB_nk = tCrB(_, _, 0, 0);\n CUTLASS_PRAGMA_UNROLL\n for (int k_block = 0; k_block < size<2>(tCrA) / EpilogueUnrollFactor; ++k_block) {\n int accumulator_k_block =\n accumulator_pipe_producer_state.index() * EpilogueUnrollFactor;\n int tCrA_k_block = k_block * EpilogueUnrollFactor;\n accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < EpilogueUnrollFactor; i++) {\n auto accumulators = bulk_tmem_mma(_, _, _, accumulator_k_block + i);\n gemm(mma, tCrA_mk(_, _, tCrA_k_block + i), tCrB_nk, accumulators);\n }\n\n accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);\n ++accumulator_pipe_producer_state;\n }\n auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n skip_wait = k_tile >= K_TILE_MAX;\n mainloop_pipeline.umma_consumer_release(curr_mainloop_pipe_consumer_state);\n barrier_token =\n mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n }\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n tmem_allocator.release_allocation_lock();\n accumulator_pipeline.producer_tail(accumulator_pipe_producer_state);\n tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);\n }\n } else if (is_sched_warp) {\n // Scheduler warp manages tile assignment and pipeline progress for warps\n cutlass::arch::warpgroup_reg_dealloc<32>();\n do {\n sched_throttle_pipeline.consumer_wait(sched_pipeline_throttle_consumer_state);\n sched_throttle_pipeline.consumer_release(sched_pipeline_throttle_consumer_state);\n ++sched_pipeline_throttle_consumer_state;\n sched_pipeline_producer_state =\n scheduler.advance_to_next_work(sched_pipeline, sched_pipeline_producer_state);\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n } else if (is_epilogue_col_quant_warp) {\n // Warp responsible for quantizing output of Hadamard transform to FP4 for columnwise usage,\n // and writing result tensors/scales to global memory.\n cutlass::arch::warpgroup_reg_alloc<192>();\n if constexpr (kEnableRHTColQuant) {\n using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x;\n\n auto acc_epilogue_pipelined_shape =\n append(acc_shape_epilogue, Int{});\n auto bulk_tmem_epilogue_layout = make_layout(\n acc_epilogue_pipelined_shape,\n make_stride(stride<0>(bulk_tmem_mma), Int<0>{}, Int<0>{}, size<1>(epilogue_tiler)));\n auto bulk_tmem_epilogue = make_tensor(make_tmem_ptr(), bulk_tmem_epilogue_layout);\n\n // Use 256-bit fragments for aligned bulk stores\n static int constexpr FragmentSize = 256 / sizeof_bits_v;\n\n // Wait for TMEM allocation for this pipeline to finish\n tmem_allocation_result_barrier.arrive_and_wait();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_epilogue.data() = tmem_base_ptr;\n int global_thread_idx = threadIdx.x;\n int local_thread_idx = global_thread_idx % cutlass::NumThreadsPerWarpGroup;\n // g2s load all global_d_amax\n CUTLASS_PRAGMA_NO_UNROLL\n for (int g = local_thread_idx; g < args.num_tensors; g += NumEpilogueColQuantThreadCount) {\n shared_storage.global_d_amax[g] =\n __ldg(reinterpret_cast(args.global_d_amax_list[g]));\n }\n\n size_t rng_seed = 0;\n size_t rng_offset = 0;\n // Setup RNG for stochastic rounding\n if constexpr (kEnableStochasticRounding) {\n rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0;\n rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0;\n }\n int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler));\n\n // Determine quantization scale factor layouts/output splits for this group\n TSFDLayout sfd_layout;\n int cur_N = args.split_sections[group_idx];\n if constexpr (kEnableSwizzleSFOutput) {\n sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{});\n } else {\n sfd_layout = make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)),\n make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{})));\n }\n // Build output tensors for columns and their quant scales\n Tensor mD = make_tensor(\n cute::subbyte_iterator(reinterpret_cast(args.output_colwise_list[group_idx])),\n make_shape(M, cur_N), DStride{}); // (M,packed_N)\n Tensor gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n\n Tensor mSFD = make_tensor(make_gmem_ptr(reinterpret_cast(\n args.output_colwise_scale_inv_list[group_idx])),\n sfd_layout);\n Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n\n Tensor gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler));\n\n // Setup tile-level TMEM (t2r) and global memory (r2g) copy descriptors\n auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{}));\n auto tiled_r2g =\n make_tiled_copy_D(Copy_Atom{}, tiled_t2r);\n auto thr_t2r = tiled_t2r.get_slice(local_thread_idx);\n auto thr_r2g = tiled_r2g.get_slice(local_thread_idx);\n\n cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount,\n cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);\n // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release}\n static constexpr float fp4_max = 6.0f;\n static constexpr float fp8_max = 448.0f;\n static constexpr float fp4_max_inv = 1.0f / fp4_max;\n float c_global_amax_val = shared_storage.global_d_amax[group_idx];\n float global_encode_scale = c_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / c_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n float global_decode_scale = 1.0f / global_encode_scale;\n\n // Scaling factor for fast math path\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n\n do {\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0;\n k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();\n ++k_tile) {\n int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler);\n\n int cur_group_idx = GetGroupIdx(&args, global_tile_n_offset);\n\n if (cur_group_idx != group_idx) {\n group_idx = cur_group_idx;\n c_global_amax_val = shared_storage.global_d_amax[group_idx];\n // update amax\n global_encode_scale = c_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / c_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n global_decode_scale = 1.0f / global_encode_scale;\n global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n cur_N = args.split_sections[group_idx];\n if constexpr (kEnableSwizzleSFOutput) {\n sfd_layout =\n tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{});\n } else {\n sfd_layout =\n make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)),\n make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{})));\n }\n // update tensor\n mD = make_tensor(cute::subbyte_iterator(\n reinterpret_cast(args.output_colwise_list[group_idx])),\n make_shape(M, cur_N), DStride{});\n gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n mSFD = make_tensor(make_gmem_ptr(reinterpret_cast(\n args.output_colwise_scale_inv_list[group_idx])),\n sfd_layout);\n gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{}); // (BLK_M,BLK_N)\n\n gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler));\n }\n int group_start_offset = args.split_sections_range[group_idx];\n int local_tile_n_idx =\n (global_tile_n_offset - group_start_offset) / size<1>(epilogue_tiler);\n Tensor tDgD_mn = gD_mn_view(_, _, _, scheduler.tile_m(), local_tile_n_idx);\n\n Tensor tDgSFD_mn = gSFD_mn(_, _, scheduler.tile_m(), local_tile_n_idx);\n accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state);\n\n auto Acc = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index());\n Tensor tDtAcc = thr_t2r.partition_S(Acc); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDgD = thr_t2r.partition_D(tDgD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n\n Tensor tTR_rAcc = make_tensor(\n shape(tDgD)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDrD = make_tensor(shape(tDgD));\n Tensor tTR_rAcc_frag =\n recast>(coalesce(tTR_rAcc));\n Tensor tDrD_frag = recast>(coalesce(tDrD));\n\n Tensor src = thr_r2g.retile_S(tDrD);\n Tensor dst = thr_r2g.retile_D(tDgD);\n\n Tensor tDgSFD_view = make_tensor(\n tDgSFD_mn.data(), make_layout(make_shape(shape(tDgSFD_mn), Int<1>{}, Int<1>{}),\n make_stride(stride(tDgSFD_mn), Int<0>{}, Int<0>{})));\n Tensor tDgSFD = filter(thr_t2r.partition_D(tDgSFD_view));\n Tensor tDrSFD = make_tensor(shape(tDgSFD));\n\n static int constexpr NumVecs = size(tDgD) / VectorSize;\n Tensor tD_rRowSFD_frg = recast>(tDrSFD);\n\n // Compute amax and quantization scales for this tile\n cutlass::maximum_absolute_value_reduction<\n cutlass::Array, true>\n amax_reduction;\n cutlass::Array vec_maxs;\n cutlass::Array pvscales;\n // Copy from TMEM to registers\n copy(tiled_t2r, tDtAcc, tTR_rAcc);\n cutlass::arch::fence_view_async_tmem_load();\n accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state);\n ++accumulator_pipe_consumer_state;\n\n if constexpr (!kUseFastMath) {\n // Downcast to BF16 for bit-wise compatibility with\n // unfused kernels\n auto convert_accum_to_bf16 =\n cutlass::NumericArrayConverter{};\n auto convert_bf16_to_accum =\n cutlass::NumericArrayConverter{};\n tTR_rAcc_frag(_0{}) =\n convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{})));\n tTR_rAcc_frag(_1{}) =\n convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_1{})));\n }\n\n auto compute_frgs = reinterpret_cast *>(\n tTR_rAcc_frag.data());\n auto output_frgs = reinterpret_cast *>(tDrD_frag.data());\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]);\n }\n\n pvscales = cutlass::multiplies>{}(\n vec_maxs, global_encode_scale_multiplier);\n auto pvscales_cvted =\n cutlass::NumericArrayConverter{}(pvscales);\n\n tD_rRowSFD_frg(_0{}) = pvscales_cvted;\n auto qpvscale_ups = cutlass::NumericArrayConverter{}(\n tD_rRowSFD_frg(_0{}));\n auto qpvscale_scaled =\n cutlass::multiplies>{}(\n qpvscale_ups, global_decode_scale);\n cutlass::Array acc_scales;\n if constexpr (kUseFastMath) {\n // Fast math: compute approximate reciprocal\n acc_scales =\n cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled);\n } else {\n // Accurate math: compute reciprocal with division\n acc_scales = cutlass::divides>{}(\n 1.0, qpvscale_scaled);\n }\n\n // Prepare stochastic rounding random state if enabled\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n transformer_engine::curanddx::detail::philox4x32_native_state<\n NVTE_BUILD_NUM_PHILOX_ROUNDS>\n rng;\n // \"Prefetch\" a stochastic rounding state for the first tile\n if constexpr (kEnableStochasticRounding) {\n const size_t rng_sequence = global_thread_idx + k_tile * 512 +\n scheduler.get_linear_tile_idx() * K_TILE_MAX * 512;\n rng.init(rng_seed, rng_sequence, rng_offset);\n }\n CUTLASS_PRAGMA_UNROLL\n // Apply round/quantize to each fragment, with or without stochastic rounding\n for (int v = 0; v < NumVecs; v++) {\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(\n acc_scales[v], cutlass::platform::numeric_limits::max());\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(\n cutlass::multiplies>{}(\n compute_frgs[v], acc_scale),\n *reinterpret_cast *>(&random_uint4));\n } else {\n output_frgs[v] =\n cutlass::NumericArrayConverter{}(\n cutlass::multiplies>{}(\n compute_frgs[v], acc_scale));\n }\n }\n\n // Write quantized FP4 tile and dequant scale to gmem\n copy(tiled_r2g, src, dst);\n copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFD, tDgSFD);\n }\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n }\n } else if (is_epilogue_row_quant_warp) {\n // Warp responsible for quantizing the input (before Hadamard transform) to FP4 for row-wise usage.\n cutlass::arch::warpgroup_reg_alloc<136>();\n if constexpr (kEnableRowQuant) {\n using S2RVectorType = uint128_t;\n\n int global_thread_idx = threadIdx.x;\n int local_thread_idx = global_thread_idx % 256;\n size_t rng_seed = 0;\n size_t rng_offset = 0;\n // g2s load all global_a_amax for all groups/tensors\n CUTLASS_PRAGMA_NO_UNROLL\n for (int g = local_thread_idx; g < args.num_tensors; g += NumEpilogueRowQuantThreadCount) {\n shared_storage.global_a_amax[g] =\n __ldg(reinterpret_cast(args.global_a_amax_list[g]));\n }\n // RNG for stochastic rounding\n if constexpr (kEnableStochasticRounding) {\n rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0;\n rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0;\n }\n // Input/output tensors/partitions for row quant warp\n Tensor mQA =\n make_tensor(cute::subbyte_iterator(QA), make_layout(make_shape(M, packed_N), dQA));\n Tensor gQA_mn = local_tile(mQA, epilogue_tiler, make_coord(_, _, _), Step<_1, X, _1>{});\n Tensor mSFA = make_tensor(make_gmem_ptr(SFA), sfa_layout);\n\n Tensor gSFA_mn = local_tile(mSFA, epilogue_tiler, make_coord(_, _, _),\n Step<_1, X, _1>{}); // (BLK_M,BLK_N)\n // Swizzled shared memory A tile, with layout\n Tensor sA = as_position_independent_swizzle_tensor(group_modes<0, 2>(\n coalesce(make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()),\n sAlayout)))); // (BLOCK_M, BLOCK_M,PIPE)\n\n // Set up layouts for partitioning \u2013 tile-by-warp, with vector granularity\n using S2RWarpLayout = Layout>;\n using WarpGroupLayout = Layout>;\n using S2RThreadLayout = decltype(blocked_product(S2RWarpLayout{}, WarpGroupLayout{}));\n using S2RValLayout = Layout, _1>>;\n using S2RAtomA = Copy_Atom;\n using R2GAtomQA = Copy_Atom;\n using R2GAtomSFA = Copy_Atom;\n auto tiled_s2r = make_tiled_copy(S2RAtomA{}, S2RThreadLayout{}, S2RValLayout{});\n auto tiled_r2g_QA = make_tiled_copy(R2GAtomQA{}, S2RThreadLayout{}, S2RValLayout{});\n auto tiled_r2g_SFA = make_tiled_copy(R2GAtomSFA{}, S2RThreadLayout{}, S2RValLayout{});\n\n auto thr_s2r = tiled_s2r.get_slice(local_thread_idx);\n auto thr_r2g_QA = tiled_r2g_QA.get_slice(local_thread_idx);\n auto thr_r2g_SFA = tiled_r2g_SFA.get_slice(local_thread_idx);\n Tensor tQAsA = thr_s2r.partition_S(sA); // (Copy, Copy_M, Copy_N, PIPE)\n\n // Allocate temporary register tensors for copying quantization => output\n Tensor tQArA = make_tensor_like(\n make_layout(tQAsA(_, _, _, _0{}).shape())); // (Copy, Copy_M, Copy_N)\n Tensor tQAgQA = thr_r2g_QA.partition_S(gQA_mn);\n Tensor tQArQA = make_tensor_like(tQAgQA(_, _, _, _0{}, _0{}));\n\n Tensor tQAgSFA = thr_r2g_SFA.partition_S(gSFA_mn);\n Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{}));\n\n // Will result in barrier_id=10 passed to bar.sync instr as cutlass adds 8\n // in order to go over the reserved named barrier count.\n constexpr int row_quant_barrier_id = 2;\n cutlass::arch::NamedBarrier::sync(NumEpilogueRowQuantThreadCount, row_quant_barrier_id);\n\n int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler));\n float a_global_amax_val = shared_storage.global_a_amax[group_idx];\n // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release}\n static constexpr float fp4_max = 6.0f;\n static constexpr float fp8_max = 448.0f;\n static constexpr float fp4_max_inv = 1.0f / fp4_max;\n float global_encode_scale = a_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / a_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n\n float global_decode_scale = 1.0f / global_encode_scale;\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n auto sfa_converter = cutlass::NumericConverter{};\n do {\n CUTLASS_PRAGMA_NO_UNROLL\n for (int k_tile = 0;\n k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) {\n int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler);\n\n int cur_group_idx = GetGroupIdx(&args, global_tile_n_offset);\n if (cur_group_idx != group_idx) {\n group_idx = cur_group_idx;\n a_global_amax_val = shared_storage.global_a_amax[group_idx];\n // Update group quantization parameters/scaling\n global_encode_scale = a_global_amax_val > 0.0f\n ? cutlass::minimum_with_nan_propagation{}(\n (fp8_max * fp4_max) / a_global_amax_val,\n cutlass::platform::numeric_limits::max())\n : 1.0f;\n global_decode_scale = 1.0f / global_encode_scale;\n global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n }\n\n auto tQAgSFA_mn =\n tQAgSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto tQAgQA_mn = tQAgQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile);\n auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state);\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n copy(tiled_s2r, tQAsA(_, _, _, mainloop_pipe_consumer_state.index()), tQArA);\n cutlass::arch::fence_view_async_shared();\n mainloop_pipeline.consumer_release(mainloop_pipe_consumer_state);\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n\n // static int constexpr NumVecs = size(tQArA) / VectorSize;\n cutlass::maximum_absolute_value_reduction<\n cutlass::Array, true>\n amax_reduction;\n auto compute_frgs = reinterpret_cast *>(tQArA.data());\n auto output_frgs = reinterpret_cast *>(\n raw_pointer_cast(tQArQA.data()));\n Tensor amax =\n make_tensor(prepend(take<1, rank(tQArA)>(tQArA.shape()), _1{}));\n Tensor pvscales = make_tensor_like(amax);\n transformer_engine::curanddx::detail::philox4x32_native_state<\n NVTE_BUILD_NUM_PHILOX_ROUNDS>\n rng;\n if constexpr (kEnableStochasticRounding) {\n const size_t rng_sequence = global_thread_idx + k_tile * 512 +\n scheduler.get_linear_tile_idx() * K_TILE_MAX * 512 +\n tiles_in_m * tiles_in_n * K_TILE_MAX * 512;\n rng.init(rng_seed, rng_sequence, rng_offset);\n }\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < size<1>(group_modes<1, rank(tQArA)>(tQArA)); v++) {\n auto amax_view = group_modes<1, rank(amax)>(amax);\n auto pvscales_view = group_modes<1, rank(pvscales)>(pvscales);\n auto compute_frgs_up =\n cutlass::NumericArrayConverter{}(\n compute_frgs[v]);\n amax_view(_0{}, v) = amax_reduction(ElementAccumulator(0), compute_frgs_up);\n pvscales_view(_0{}, v) = cutlass::multiplies{}(\n amax_view(_0{}, v), global_encode_scale_multiplier);\n filter(tQArSFA)(v) = sfa_converter(pvscales_view(_0{}, v));\n auto qpvscale_ups =\n cutlass::NumericConverter{}(filter(tQArSFA)(v));\n auto qpvscale_scaled =\n cutlass::multiplies{}(qpvscale_ups, global_decode_scale);\n ElementAccumulator acc_scales;\n if constexpr (kUseFastMath) {\n // Fast math: compute approximate reciprocal\n acc_scales = cutlass::reciprocal_approximate_ftz{}(\n qpvscale_scaled);\n } else {\n // Accurate math: compute reciprocal with division\n acc_scales = cutlass::divides{}(1.0, qpvscale_scaled);\n }\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(\n acc_scales, cutlass::platform::numeric_limits::max());\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(\n cutlass::multiplies>{}(\n compute_frgs_up, acc_scale),\n *reinterpret_cast *>(&random_uint4));\n } else {\n output_frgs[v] =\n cutlass::NumericArrayConverter{}(\n cutlass::multiplies>{}(\n compute_frgs_up, acc_scale));\n }\n }\n copy(tiled_r2g_QA, tQArQA, tQAgQA_mn);\n copy(tiled_r2g_SFA, filter(tQArSFA), filter(tQAgSFA_mn));\n }\n // scheduler.advance();\n scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state);\n ++sched_pipeline_consumer_state;\n scheduler.update_work_tile_info();\n } while (scheduler.is_valid());\n }\n\n } else {\n cutlass::arch::warpgroup_reg_dealloc<32>();\n }\n } // sm100 compile guard end\n} // NOLINT(readability/fn_size)\n\ntemplate \nvoid group_row_col_rht_gemm_ntt_w_sfc(int packed_sequence_length, int hidden_size, TA const *A,\n TB const *B, TQA *QA, TSFA *SFA,\n MultiAmaxHadamardCastFusionArgs &args,\n const size_t *rng_state, uint32_t *tile_scheduler_workspace,\n uint32_t sm_count, hipStream_t stream,\n int k_tile_size = 1024) {\n using namespace cute;\n static int constexpr SFVecSize = 16;\n static int constexpr RhtTensorSize = 16;\n\n static_assert(RhtTensorSize == 16, \"RhtTensorSize must be 16\");\n using LinearSFALayout = decltype(make_layout(make_shape(make_shape(Int{}, 0), 0),\n make_stride(make_stride(_0{}, _1{}), 0)));\n using LinearSFDLayout = decltype(make_layout(make_shape(0, make_shape(Int{}, 0)),\n make_stride(0, make_stride(_0{}, _1{}))));\n\n using SwizzledSFALayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFDLayoutAtom =\n cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom;\n using SwizzledSFALayout = decltype(tile_to_shape(\n SwizzledSFALayoutAtom{}, make_shape(hidden_size, packed_sequence_length), Step<_1, _2>{}));\n using SwizzledSFDLayout = decltype(tile_to_shape(\n SwizzledSFDLayoutAtom{}, make_shape(hidden_size, packed_sequence_length), Step<_2, _1>{}));\n\n using SFALayout = cute::conditional_t;\n using SFDLayout = cute::conditional_t;\n SFALayout sfa_layout;\n SFDLayout sfd_layout;\n\n if constexpr (kEnableSwizzleSFOutput) {\n sfa_layout = tile_to_shape(SwizzledSFALayoutAtom{},\n make_shape(hidden_size, packed_sequence_length), Step<_1, _2>{});\n sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{},\n make_shape(hidden_size, packed_sequence_length), Step<_2, _1>{});\n } else {\n sfa_layout = make_layout(\n make_shape(make_shape(Int{}, hidden_size / SFVecSize), packed_sequence_length),\n make_stride(make_stride(_0{}, _1{}), hidden_size / SFVecSize));\n sfd_layout = make_layout(\n make_shape(hidden_size, make_shape(Int{}, packed_sequence_length / SFVecSize)),\n make_stride(packed_sequence_length / SFVecSize, make_stride(_0{}, _1{})));\n }\n\n // Define shapes (dynamic)\n auto M = hidden_size;\n auto N = packed_sequence_length;\n Tensor tensorA = make_tensor(A, make_shape(hidden_size, packed_sequence_length), LayoutLeft{});\n Tensor tensorB = make_tensor(B, make_shape(RhtTensorSize, RhtTensorSize), LayoutLeft{});\n Tensor tensorQA = make_tensor(QA, make_shape(hidden_size, packed_sequence_length), LayoutLeft{});\n Tensor tensorSFA = make_tensor(SFA, sfa_layout);\n\n // Define strides (from tensors)\n auto dA = stride(tensorA); // (dM,dK)\n auto dB = stride(tensorB); // (dN,dK)\n auto dD = LayoutRight{}; // (dM,dN)\n auto dQA = stride(tensorQA); // (dM,dK)\n using ClusterShape = Shape<_1, _1, _1>;\n auto cluster_shape = ClusterShape{};\n auto cluster_tile_shape = Shape<_128, Int, Int>{};\n auto cluster_tile_mainloop = Shape<_128, Int, _128>{};\n\n // Each mainloop / epilogue loads 128 x 64 tiles while each MMA proceeds with 128 x 16 tiles\n static int constexpr EpilogueUnrollFactor =\n size<2>(cluster_tile_mainloop) / size<2>(cluster_tile_shape);\n // Construct the MMA\n auto mma = make_tiled_mma(\n SM100_MMA_F16BF16_SS(cluster_tile_shape), size<1>(cluster_tile_shape),\n UMMA::Major::MN, UMMA::Major::MN>{},\n Layout>{});\n\n // Assert that the TiledMMA uses all CTAs in the CGA.\n CUTE_STATIC_ASSERT_V(size(cluster_shape) == size(mma));\n CUTE_STATIC_ASSERT_V(evenly_divides(cluster_tile_shape, tile_shape(mma)));\n\n // Determine the A and B shapes\n auto mma_shape_B =\n partition_shape_B(mma, make_shape(size<1>(cluster_tile_shape), size<2>(cluster_tile_shape)));\n\n using TiledMma = decltype(mma);\n using AtomThrID = typename TiledMma::AtomThrID;\n\n using SmemShape_M = decltype(shape_div(\n shape<0>(cluster_tile_shape),\n shape_div(shape<0>(cluster_tile_shape), size<0>(cluster_tile_shape) / size(AtomThrID{}))));\n using SmemShape_N = decltype(shape_div(\n shape<1>(cluster_tile_shape),\n shape_div(shape<1>(cluster_tile_shape), size<1>(cluster_tile_shape) / size(AtomThrID{}))));\n using SmemShape_K = decltype(cute::get<2>(cluster_tile_shape));\n\n using SmemLayoutAtomB =\n decltype(cutlass::gemm::collective::detail::sm100_smem_selector());\n\n auto mma_shape_A = partition_shape_A(\n mma, make_shape(size<0>(cluster_tile_mainloop), size<2>(cluster_tile_mainloop)));\n using SmemShape_M_A =\n decltype(shape_div(shape<0>(cluster_tile_mainloop),\n shape_div(shape<0>(cluster_tile_mainloop),\n size<0>(cluster_tile_mainloop) / size(AtomThrID{}))));\n using SmemShape_K_A = decltype(cute::get<2>(cluster_tile_mainloop));\n using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<\n cute::UMMA::Major::MN, TA, SmemShape_M_A, SmemShape_K_A>());\n\n static uint32_t constexpr TotalTmemRows = 128;\n static uint32_t constexpr Sm100TmemCapacityColumns = 512;\n static uint32_t constexpr TotalTmem = TotalTmemRows * Sm100TmemCapacityColumns;\n static uint32_t constexpr AccumulatorPipelineStageCount =\n TotalTmem / (cute::size<0>(cluster_tile_shape) * cute::size<1>(cluster_tile_shape));\n\n // Define the smem layouts (static)\n // Calculate max pipeline stages based on Blackwell SM100's 232KB shared memory\n constexpr int SchedulerPipelineStageCount = 4;\n static int constexpr MainloopPipelineBytes = sizeof(\n typename cutlass::detail::CustomizedPipelineTmaUmmaAsync<1, Shape<_1, _1, _1>,\n Shape<_1, _1, _1>>::SharedStorage);\n\n static int constexpr SchedulerWorkspaceBytes = sizeof(int) * SchedulerPipelineStageCount;\n static int constexpr SchedulerThrottlePipelineBytes =\n sizeof(typename cutlass::PipelineAsync::SharedStorage);\n static int constexpr SchedulerPipelineBytes =\n sizeof(typename cutlass::PipelineCLCFetchAsync::SharedStorage);\n\n static int constexpr TmemDeallocBytes = sizeof(cutlass::arch::ClusterBarrier);\n static int constexpr BTensorBytes = cute::size(mma_shape_B) * sizeof(TB);\n static int constexpr AccPipelineBytes = sizeof(\n typename cutlass::PipelineUmmaAsync>::SharedStorage);\n static int constexpr TmemBasePtrsBytes = sizeof(uint32_t);\n static int constexpr kBlackwellSmemSize = 232448; // 232KB in bytes\n static int constexpr kBytesPerStage =\n cute::size(mma_shape_A) * sizeof(TA) + MainloopPipelineBytes;\n static int constexpr kReservedBytes = SchedulerWorkspaceBytes + SchedulerThrottlePipelineBytes +\n SchedulerPipelineBytes + TmemBasePtrsBytes +\n TmemDeallocBytes + BTensorBytes +\n AccPipelineBytes; // Reserve for barriers and other uses\n static int constexpr kMaxStages = (kBlackwellSmemSize - kReservedBytes) / kBytesPerStage;\n auto sP = Int{}; // SMEM pipelines\n\n auto sA = UMMA::tile_to_mma_shape(SmemLayoutAtomA{}, append(mma_shape_A, sP),\n Step<_2, _1, _3>{}); // (MMA,MMA_M,MMA_K,PIPE)\n auto sB = UMMA::tile_to_mma_shape(SmemLayoutAtomB{},\n append(mma_shape_B, _1{})); // (MMA,MMA_N,MMA_K, _1)\n auto sD = Layout<_1>{}; // XXX Dummy\n\n auto tma_load_a =\n make_tma_copy_A_sm100(SM90_TMA_LOAD{}, tensorA, sA(_, _, _, 0), cluster_tile_mainloop, mma);\n auto tma_load_b =\n make_tma_copy_B_sm100(SM90_TMA_LOAD{}, tensorB, sB(_, _, _, 0), cluster_tile_shape, mma);\n\n // Assert checks on tile sizes -- no predication\n assert(M % size<0>(cluster_tile_shape) == 0);\n assert(N % size<1>(cluster_tile_shape) == 0);\n\n dim3 dimBlock(512);\n dim3 dimCluster(size<0>(cluster_shape), size<1>(cluster_shape), size<2>(cluster_shape));\n dim3 dimGrid(sm_count, 1, 1);\n\n int smem_size = sizeof(\n SharedStorage);\n\n auto *kernel_ptr = &group_row_col_rht_gemm_device<\n decltype(M), decltype(N), decltype(k_tile_size), decltype(cluster_shape),\n decltype(cluster_tile_shape), TA, decltype(dA), decltype(sA), decltype(tma_load_a), TB,\n decltype(dB), decltype(sB), decltype(tma_load_b), TD, decltype(dD), decltype(sD), TSFD,\n decltype(sfd_layout), TQA, decltype(dQA), TSFA, decltype(sfa_layout), decltype(mma),\n AccumulatorPipelineStageCount, SchedulerPipelineStageCount, kEnableStochasticRounding,\n kEnableRHTColQuant, kEnableRowQuant, kEnableSwizzleSFOutput, kUseFastMath>;\n\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(*kernel_ptr), hipFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n\n // Set workspace and set to zero\n NVTE_CHECK_CUDA(hipMemsetAsync(reinterpret_cast(tile_scheduler_workspace), 0,\n sizeof(uint32_t), stream));\n\n // Launch kernel\n cutlass::ClusterLaunchParams params = {dimGrid, dimBlock, dimCluster, smem_size, stream};\n cutlass::Status status = cutlass::launch_kernel_on_cluster(\n params, (void const *)kernel_ptr, M, N, k_tile_size, cluster_shape, cluster_tile_shape, A, dA,\n sA, tma_load_a, B, dB, sB, tma_load_b, QA, dQA, SFA, sfa_layout, args,\n tile_scheduler_workspace, mma, rng_state);\n NVTE_CHECK_CUDA(hipGetLastError());\n NVTE_CHECK(status == cutlass::Status::kSuccess, \"Kernel launch failed.\");\n}\n\n} // namespace\n} // namespace detail\n\nvoid group_hadamard_transform_cast_fusion(const Tensor &input_, std::vector &output_list,\n const size_t *split_sections, size_t num_tensors,\n const Tensor &hadamard_matrix_,\n QuantizationConfig &quant_config, Tensor &quant_workspace,\n hipStream_t stream) {\n NVTE_API_CALL(group_hadamard_transform_cast_fusion);\n\n using transformer_engine::detail::kMaxTensorsPerKernel;\n using transformer_engine::detail::MultiAmaxHadamardCastFusionArgs;\n\n NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16,\n \"Input tensor must be BF16 tensor, but dtype is \", to_string(input_.dtype()), \".\");\n NVTE_CHECK(input_.dim() >= 2, \"Input must be a 2D tensor.\");\n const SimpleTensor &input = input_.data;\n\n NVTE_CHECK(output_list.size() == num_tensors,\n \"Number of output tensors should match number of tensors.\");\n\n NVTE_CHECK(num_tensors <= kMaxTensorsPerKernel,\n \"Number of tensors should be less than or equal to \", kMaxTensorsPerKernel);\n\n // construct the multi-tensor args\n MultiAmaxHadamardCastFusionArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.split_sections_range[0] = 0;\n bool all_has_row_quant = true;\n bool all_has_col_quant = true;\n void *rowwise_data_base_ptr = nullptr;\n void *rowwise_scale_inv_base_ptr = nullptr;\n for (size_t i = 0; i < num_tensors; ++i) {\n NVTE_CHECK(split_sections[i] % 128 == 0, \"component \", i,\n \" of split_sections should be 128 multiple\");\n if (split_sections[i] == 0) {\n continue;\n }\n bool has_row_quant = output_list[i]->data.dptr != nullptr;\n bool has_col_quant = output_list[i]->columnwise_data.dptr != nullptr;\n all_has_row_quant = all_has_row_quant && has_row_quant;\n all_has_col_quant = all_has_col_quant && has_col_quant;\n // sanity check, the two bool flags cannot be both false\n NVTE_CHECK(has_row_quant || has_col_quant,\n \"At least one of the output tensors must have row or column quant.\");\n void *amax_rowwise_ptr =\n has_row_quant ? reinterpret_cast(output_list[i]->amax.dptr) : nullptr;\n void *amax_colwise_ptr =\n has_col_quant ? reinterpret_cast(output_list[i]->columnwise_amax.dptr) : nullptr;\n void *rowwise_data_ptr =\n has_row_quant ? reinterpret_cast(output_list[i]->data.dptr) : nullptr;\n void *rowwise_scale_inv_ptr =\n has_row_quant ? reinterpret_cast(output_list[i]->scale_inv.dptr) : nullptr;\n if (all_has_row_quant &&\n (rowwise_data_base_ptr == nullptr || rowwise_scale_inv_base_ptr == nullptr)) {\n rowwise_data_base_ptr = rowwise_data_ptr;\n rowwise_scale_inv_base_ptr = rowwise_scale_inv_ptr;\n }\n void *output_colwise_ptr =\n has_col_quant ? reinterpret_cast(output_list[i]->columnwise_data.dptr) : nullptr;\n void *output_colwise_scale_inv_ptr =\n has_col_quant ? reinterpret_cast(output_list[i]->columnwise_scale_inv.dptr)\n : nullptr;\n kernel_args.global_a_amax_list[kernel_args.num_tensors] = amax_rowwise_ptr;\n kernel_args.global_d_amax_list[kernel_args.num_tensors] = amax_colwise_ptr;\n kernel_args.output_colwise_list[kernel_args.num_tensors] = output_colwise_ptr;\n kernel_args.output_colwise_scale_inv_list[kernel_args.num_tensors] =\n output_colwise_scale_inv_ptr;\n kernel_args.split_sections[kernel_args.num_tensors] = split_sections[i];\n kernel_args.split_sections_range[kernel_args.num_tensors + 1] =\n kernel_args.split_sections_range[kernel_args.num_tensors] + split_sections[i];\n kernel_args.num_tensors++;\n }\n\n // Stochastic rounding config\n const bool use_stochastic_rounding = quant_config.stochastic_rounding;\n const size_t *rng_state = nullptr;\n if (use_stochastic_rounding) {\n NVTE_CHECK(quant_config.rng_state != nullptr,\n \"Enabled stochastic rounding without providing RNG state\");\n const Tensor &rng_state_tensor = *convertNVTETensorCheck(quant_config.rng_state);\n NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64,\n \"RNG state should contain 2 64-bit values.\");\n NVTE_CHECK(rng_state_tensor.data.shape == std::vector{2},\n \"Shape of the RNG state should be [2], but got \", rng_state_tensor.data.shape);\n rng_state = reinterpret_cast(rng_state_tensor.data.dptr);\n }\n\n uint32_t *tile_scheduler_workspace = nullptr;\n NVTE_CHECK(quant_workspace.data.dptr != nullptr, \"Quantization workspace must be provided.\");\n NVTE_CHECK(quant_workspace.data.buffer_size_bytes() >= sizeof(uint32_t),\n \"Quantization workspace must be at least 4 bytes.\");\n tile_scheduler_workspace = reinterpret_cast(quant_workspace.data.dptr);\n\n // Template arguments\n using TA = cute::bfloat16_t;\n using TB = cute::bfloat16_t;\n using TD = cutlass::float_e2m1_t;\n using TSFD = cutlass::float_ue4m3_t;\n using TQA = TD;\n using TSFA = TSFD;\n\n checkCuDriverContext(stream);\n\n // Check Hadamard matrix\n constexpr int kHadamardDimension = 16;\n\n NVTE_CHECK(hadamard_matrix_.dtype() == transformer_engine::DType::kBFloat16,\n \"Hadamard matrix must be BF16 tensor, but dtype is \",\n to_string(hadamard_matrix_.dtype()), \".\");\n const SimpleTensor &hadamard_matrix = hadamard_matrix_.data;\n NVTE_CHECK(\n (hadamard_matrix_.shape() == std::vector{kHadamardDimension, kHadamardDimension}),\n \"Hadamard matrix must have shape=\",\n std::vector{kHadamardDimension, kHadamardDimension},\n \", but got shape=\", hadamard_matrix_.shape(), \".\");\n const size_t hadamard_dimension = hadamard_matrix.shape[0];\n\n const size_t ndim = input.shape.size();\n const size_t n = input.shape[ndim - 1];\n size_t m = 1;\n for (size_t i = 0; i < ndim - 1; ++i) {\n m *= input.shape[i];\n }\n\n auto sm_count = transformer_engine::cuda::sm_count();\n\n NVTE_CHECK(n % hadamard_dimension == 0, \"row_length must be divisible by hadamard_dimension.\");\n\n NVTE_CHECK(m % hadamard_dimension == 0, \"num_rows must be divisible by hadamard_dimension\");\n\n int k_tile_size = 1024;\n\n const bool use_swizzle_sf_output = false;\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_stochastic_rounding, kEnableStochasticRounding,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n all_has_col_quant, kEnableRhtColQuant,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n all_has_row_quant, kEnableRowQuant,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_swizzle_sf_output, kEnableSwizzleSFOutput,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n quant_config.use_fast_math, kUseFastMath,\n\n if constexpr (kEnableRhtColQuant || kEnableRowQuant) {\n detail::group_row_col_rht_gemm_ntt_w_sfc<\n kEnableStochasticRounding, kEnableRhtColQuant, kEnableRowQuant,\n kEnableSwizzleSFOutput, TA, TB, TQA, TSFA, TD, TSFD, kUseFastMath>(\n /*packed_sequence_length=*/m, /*hidden_size=*/n,\n /*A=*/reinterpret_cast(input.dptr),\n /*B=*/reinterpret_cast(hadamard_matrix.dptr),\n /*QA=*/reinterpret_cast(rowwise_data_base_ptr),\n /*SFA=*/reinterpret_cast(rowwise_scale_inv_base_ptr),\n /*args=*/kernel_args,\n /*rng_state=*/rng_state,\n /*tile_scheduler_workspace=*/tile_scheduler_workspace,\n /*sm_count=*/sm_count,\n /*stream=*/stream, /*k_tile_size=*/k_tile_size);\n } else {\n NVTE_ERROR(\"Invalid kernel configuration (kEnableRHTColQuant=\",\n kEnableRhtColQuant, \", kEnableRowQuant=\", kEnableRowQuant, \").\");\n }\n\n );););););\n}\n\n} // namespace transformer_engine\n\nvoid nvte_group_hadamard_transform_cast_fusion(const NVTETensor input, NVTETensor *outputs,\n const NVTETensor hadamard_matrix,\n const size_t *split_sections,\n const size_t num_tensors,\n const NVTEQuantizationConfig quant_config,\n NVTETensor quant_workspace, hipStream_t stream) {\n NVTE_API_CALL(nvte_group_hadamard_transform_cast_fusion);\n using namespace transformer_engine;\n NVTE_CHECK(num_tensors > 0, \"Number of tensors should be greater than 0.\");\n\n Tensor *input_tensor = convertNVTETensorCheck(input);\n std::vector output_list(num_tensors);\n for (size_t i = 0; i < num_tensors; ++i) {\n output_list[i] = convertNVTETensorCheck(outputs[i]);\n }\n\n Tensor *quant_workspace_tensor = convertNVTETensorCheck(quant_workspace);\n\n QuantizationConfig quant_config_cpp;\n if (quant_config != nullptr) {\n quant_config_cpp = *reinterpret_cast(quant_config);\n }\n\n // Call the multi-tensor Hadamard transform amax implementation.\n group_hadamard_transform_cast_fusion(*input_tensor, output_list, split_sections, num_tensors,\n *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp,\n *quant_workspace_tensor, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu"} {"input": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/util/cuda_runtime.h\"\n#include \"common/util/curanddx.hpp\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"cutlass/arch/barrier.h\"\n#include \"cutlass/cutlass.h\"\n#include \"cutlass/gemm/collective/builders/sm100_common.inl\"\n#include \"cutlass/numeric_conversion.h\"\n#include \"cutlass/pipeline/pipeline.hpp\"\n#include \"cutlass/util/GPU_Clock.hpp\"\n#include \"cutlass/util/command_line.h\"\n#include \"cutlass/util/print_error.hpp\"\n\nnamespace transformer_engine {\nnamespace detail {\nnamespace {\n\nusing namespace cute;\nusing cute::\n Tensor; // Ensure unqualified Tensor refers to cute::Tensor, not transformer_engine::Tensor\n\nusing Stride2D = cute::Stride>;\n\nconstexpr int kMaxTensorsPerKernel = 64; // Args must be <4 KB, expand 64 if needed\nstruct MultiAmaxHadamardCastFusionArgs {\n // (output) Amax buffer for pre-RHT amax buffer\n void *global_amax_list[kMaxTensorsPerKernel];\n // output C pointers for each tensor\n void *output_colwise_list[kMaxTensorsPerKernel];\n // output scale inverse pointers for each tensor\n void *output_colwise_scale_inv_list[kMaxTensorsPerKernel];\n // split sections of each tensor of input\n int split_sections[kMaxTensorsPerKernel];\n // Prefix sum (with leading zero) of split_sections of each tensor of input\n int split_sections_range[kMaxTensorsPerKernel + 1];\n // stride 2D struct for CUTE\n Stride2D output_stride2d_list[kMaxTensorsPerKernel];\n // Number of tensors (splits) being processed by kernel\n int num_tensors;\n};\n\n__device__ __forceinline__ float *GetGlobalAmaxPtrByTensorId(\n MultiAmaxHadamardCastFusionArgs *kernel_args_ptr, int tensor_id) {\n // directly returns the global amax pointer by tensor id\n if (tensor_id < 0 || tensor_id >= kernel_args_ptr->num_tensors) {\n return nullptr;\n }\n return reinterpret_cast(kernel_args_ptr->global_amax_list[tensor_id]);\n}\n\n__device__ __forceinline__ int GetTensorId(MultiAmaxHadamardCastFusionArgs *kernel_args_ptr,\n int offset) {\n // Check the kernel args and get the corresponding id\n const int num_tensors = kernel_args_ptr->num_tensors;\n if (offset >= kernel_args_ptr->split_sections_range[num_tensors]) {\n return num_tensors - 1;\n }\n int tensor_id = 0;\n while (kernel_args_ptr->split_sections_range[tensor_id + 1] <= offset) {\n ++tensor_id;\n }\n return tensor_id;\n}\n\n// calculate the global encode scale factor for a given global amax.\n__device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax) {\n constexpr float kFP8E4M3Max = 448.0f;\n constexpr float kFP4E2M1Max = 6.0f;\n // If scale is infinity, return max value of float32\n float global_encode_scale = cutlass::minimum_with_nan_propagation{}(\n kFP8E4M3Max * kFP4E2M1Max / global_amax, cutlass::platform::numeric_limits::max());\n // If global amax is 0 or infinity, return 1\n return (global_amax == 0.f || global_encode_scale == 0.f) ? 1.f : global_encode_scale;\n}\n\ntemplate \nstruct SharedStorage {\n static constexpr int AccumulatorPipelineStageCount = 16;\n using AtomThrShapeMNK = cute::Shape<_1, _1, _1>;\n\n using AccumulatorPipeline =\n cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;\n\n static constexpr int MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline =\n cutlass::PipelineTmaUmmaAsync, AtomThrShapeMNK>;\n using MainloopPipelineStorage = typename MainloopPipeline::SharedStorage;\n\n alignas(16) AccumulatorPipelineStorage accumulator;\n alignas(16) MainloopPipelineStorage mainloop;\n alignas(16) cute::uint64_t tma_barrier[1];\n uint32_t tmem_base_ptr;\n\n struct TensorStorage : cute::aligned_struct<128, _1> {\n // cute::array_aligned> smem_A;\n cute::array_aligned> smem_A;\n cute::array_aligned> smem_B;\n } tensors;\n};\n\nCUTLASS_DEVICE\ncutlass::Array StochasticNumericConverterBase(\n cutlass::Array const &input, cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING;\n if constexpr (has_rs) {\n auto output_ptr = reinterpret_cast(&output);\n asm volatile(\n \"{\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\\n\"\n \"}\"\n : \"=h\"(output_ptr[0]), \"=h\"(output_ptr[1])\n : \"f\"(input[0]), \"f\"(input[1]), \"f\"(input[2]), \"f\"(input[3]), \"f\"(input[4]), \"f\"(input[5]),\n \"f\"(input[6]), \"f\"(input[7]), \"r\"(rbits[0]), \"r\"(rbits[1]));\n } else {\n NVTE_DEVICE_ERROR(\n \"FP4 cvt PTX instructions are architecture-specific. \"\n \"Try recompiling with sm_XXXa instead of sm_XXX.\");\n }\n return output;\n}\n\nCUTLASS_DEVICE\ncutlass::Array StochasticNumericConverter(\n cutlass::Array const &input, cutlass::Array const *rbits) {\n using result_type = cutlass::Array;\n result_type output;\n cutlass::Array *result_ptr =\n reinterpret_cast *>(&output);\n cutlass::Array const *source_ptr =\n reinterpret_cast const *>(&input);\n cutlass::Array const *rbits_ptr =\n reinterpret_cast const *>(rbits);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 2; i++) {\n result_ptr[i] = StochasticNumericConverterBase(source_ptr[i], rbits_ptr[i]);\n }\n return output;\n}\n\ntemplate \n__global__ static void group_rht_gemm_device(\n MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, TA const *A, AStride dA,\n ASmemLayout sAlayout, CUTE_GRID_CONSTANT TmaLoadA const tma_load_a, TB const *B, BStride dB,\n BSmemLayout sBlayout, CUTE_GRID_CONSTANT TmaLoadB const tma_load_b, CSmemLayout, TiledMMA mma,\n MultiAmaxHadamardCastFusionArgs kernel_args, const size_t *rng_state) {\n using namespace cute;\n constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY;\n if constexpr (!is_blackwell_arch) {\n NVTE_DEVICE_ERROR(\"RHT fusion is only supported on Blackwell.\");\n return;\n } else {\n using X = Underscore;\n // static constexpr bool kApplyStochasticRounding = true;\n using ElementAccumulator = float;\n static constexpr int K_PIPE_MAX = size<3>(ASmemLayout{});\n using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>;\n static constexpr uint32_t kTmaTransactionBytes = cutlass::bits_to_bytes(\n size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v);\n\n static constexpr int kTmaRhtTensorTransactionBytes =\n cutlass::bits_to_bytes(16 * 16 * cute::sizeof_bits_v);\n static constexpr int AccumulatorPipelineStageCount = 16;\n\n static constexpr int MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline = cutlass::PipelineTmaUmmaAsync, AtomThrShapeMNK>;\n using MainloopPipelineState = typename MainloopPipeline::PipelineState;\n\n using TmemAllocator = cute::TMEM::Allocator1Sm;\n static constexpr int VectorSize = 16;\n const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0;\n const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0;\n // Preconditions\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n\n // Represent the full tensors\n Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, N));\n Tensor mB = tma_load_b.get_tma_tensor(make_shape(16, 16));\n\n using TensorC = decltype(make_tensor(subbyte_iterator(recast_ptr(nullptr)), // engine\n make_shape(int{}, int{}), // (M, N_i)\n Stride2D{} // stride (dM, dN)\n ));\n\n using TensorSFC = decltype(make_tensor(\n make_gmem_ptr(recast_ptr(nullptr)),\n make_layout(make_shape(int{}, // M\n make_shape(make_shape(Int<16>{}, _4{}), // (16, 4)\n int{}) // n_tiles = split / 64\n ),\n make_stride(int{}, // dM = (split / 16)\n make_stride(make_stride(_0{}, _1{}), // inner (16,4) layout\n _4{}) // tiles stride\n ))));\n\n auto cluster_shape = Shape<_1, _1, _1>{};\n\n // Get the appropriate blocks for this Cluster\n dim3 cluster_coord_in_grid = cluster_id_in_grid();\n\n // Total number of k-tiles\n const int K_TILE_MAX = min(N, K) / 64;\n uint32_t tiles_in_m = (M + size<0>(cluster_tile) - 1) / size<0>(cluster_tile);\n uint32_t tiles_in_n = (N + 64 - 1) / 64;\n uint32_t linear_tile_idx = blockIdx.x;\n uint32_t tile_idx_m = linear_tile_idx % tiles_in_m;\n uint32_t tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n\n auto mainloop_tiler = Shape<_128, _16, _64>{};\n auto epilogue_tiler = Shape<_128, _64, _64>{};\n Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{});\n Tensor gB_nk =\n local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k)\n // Tensor gC_mn = local_tile(mC, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N)\n\n using TensorGC = decltype(local_tile(std::declval(), decltype(epilogue_tiler){},\n make_coord(_, _, _), Step<_1, _1, X>{}));\n\n using TensorGSFC = decltype(local_tile(std::declval(), decltype(epilogue_tiler){},\n make_coord(_, _, _), Step<_1, _1, X>{}));\n\n // Allocate SMEM\n extern __shared__ char shared_memory[];\n using SharedStorage = SharedStorage;\n SharedStorage &shared_storage = *reinterpret_cast(shared_memory);\n Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()),\n sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()),\n sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n //\n // MMA: Define C accumulators and A/B partitioning\n //\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k)\n\n auto mma_epilogue = make_tiled_mma(SM100_MMA_F16BF16_SS{},\n Layout>{});\n ThrMMA thr_mma_epilogue = mma_epilogue.get_slice(block_rank_in_cluster);\n\n using TiledMmaEpilogue = decltype(mma_epilogue);\n Tensor tCgA = thr_mma.partition_A(gA_mk);\n // Allocate \"fragments\" -- these are actually umma smem descriptors\n Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)\n Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE)\n\n auto acc_shape_mma = partition_shape_C(TiledMMA{}, take<0, 2>(ClusterTileShape{}));\n auto acc_shape_epilogue = partition_shape_C(TiledMmaEpilogue{}, take<0, 2>(epilogue_tiler));\n\n auto bulk_tmem_mma =\n TiledMMA::make_fragment_C(append(acc_shape_mma, Int{}));\n\n auto bulk_tmem_epilogue = TiledMmaEpilogue::make_fragment_C(\n append(acc_shape_epilogue, Int{}));\n\n TmemAllocator tmem_allocator{};\n cutlass::arch::NamedBarrier tmem_allocation_result_barrier(\n 32 + 128, cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);\n\n Layout cta_layout_mnk = make_layout(cluster_shape);\n Layout cta_layout_vmnk =\n tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{}));\n auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster);\n\n auto [tAgA, tAsA] =\n tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),\n group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA));\n\n auto [tBgB, tBsB] =\n tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),\n group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB));\n\n uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);\n uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);\n\n int warp_idx = cutlass::canonical_warp_idx_sync();\n\n bool is_mma_warp = (warp_idx == 0);\n bool is_dma_warp = (warp_idx == 1);\n bool is_epilogue_warp = (warp_idx >= 4 && warp_idx <= 7);\n\n // if (is_epilogue_warp && elect_one_sync()) {\n // // prefetch to make the global amax in cache\n // for (size_t i = 0; i < kernel_args.num_tensors; ++i) {\n // cute::prefetch(raw_pointer_cast(kernel_args.global_amax_list[i]));\n // }\n // }\n\n typename MainloopPipeline::Params mainloop_pipeline_params;\n if (is_dma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;\n }\n if (is_mma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;\n }\n mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp;\n mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes;\n mainloop_pipeline_params.initializing_warp = 0;\n MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params,\n cluster_shape, cute::true_type{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n\n MainloopPipelineState mainloop_pipe_consumer_state;\n MainloopPipelineState mainloop_pipe_producer_state =\n cutlass::make_producer_start_state();\n\n using AccumulatorPipeline =\n cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState;\n\n AccumulatorPipelineState accumulator_pipe_consumer_state;\n AccumulatorPipelineState accumulator_pipe_producer_state =\n cutlass::make_producer_start_state();\n\n typename AccumulatorPipeline::Params accumulator_pipeline_params;\n if (is_mma_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer;\n }\n if (is_epilogue_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer;\n }\n // Only one producer thread arrives on this barrier.\n accumulator_pipeline_params.producer_arv_count = 1;\n accumulator_pipeline_params.consumer_arv_count = size(AtomThrShapeMNK{}) * 128;\n accumulator_pipeline_params.initializing_warp = 1;\n AccumulatorPipeline accumulator_pipeline(shared_storage.accumulator,\n accumulator_pipeline_params, cluster_shape,\n cute::true_type{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n\n if (warp_idx == 2 && elect_one_sync()) {\n cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);\n }\n __syncthreads();\n using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x;\n\n if (is_dma_warp) {\n if (elect_one_sync()) {\n cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0],\n kTmaRhtTensorTransactionBytes);\n copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0),\n tBsB(_, 0));\n }\n\n do {\n bool is_first_wave = linear_tile_idx == blockIdx.x;\n uint32_t skip_wait = is_first_wave;\n auto tAgA_mk = tAgA(_, tile_idx_m, _);\n int k_tile = 0;\n auto barrier_token =\n mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state, skip_wait);\n\n CUTE_NO_UNROLL\n while (k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n) {\n int k_tile_idx_n = tile_idx_n + k_tile;\n ++k_tile;\n skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount);\n mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state, barrier_token);\n using BarrierType = typename MainloopPipeline::ProducerBarrierType;\n BarrierType *tma_barrier =\n mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state);\n int write_stage = mainloop_pipe_producer_state.index();\n ++mainloop_pipe_producer_state;\n barrier_token =\n mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state, skip_wait);\n if (cute::elect_one_sync()) {\n copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n),\n tAsA(_, write_stage));\n }\n }\n linear_tile_idx += gridDim.x;\n tile_idx_m = linear_tile_idx % tiles_in_m;\n tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n);\n mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);\n } else if (is_mma_warp) {\n mma.accumulate_ = UMMA::ScaleOut::Zero;\n\n tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns,\n &shared_storage.tmem_base_ptr);\n __syncwarp();\n tmem_allocation_result_barrier.arrive();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_mma.data() = tmem_base_ptr;\n\n cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/);\n do {\n uint32_t skip_wait = K_TILE_MAX <= 0;\n auto barrier_token =\n mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n CUTE_NO_UNROLL\n for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n;) {\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n int read_stage = mainloop_pipe_consumer_state.index();\n auto tCrA_mk = tCrA(_, _, _, read_stage);\n auto tCrB_nk = tCrB(_, _, 0, 0);\n CUTE_UNROLL\n for (int k_block = 0; k_block < size<2>(tCrA) / 4; ++k_block) {\n accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);\n CUTE_UNROLL\n for (int i = 0; i < 4; i++) {\n auto accumulators =\n bulk_tmem_mma(_, _, _, accumulator_pipe_producer_state.index() * 4 + i);\n gemm(mma, tCrA_mk(_, _, k_block * 4 + i), tCrB_nk, accumulators);\n }\n\n accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);\n ++accumulator_pipe_producer_state;\n }\n auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n skip_wait = k_tile >= K_TILE_MAX;\n barrier_token =\n mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state);\n }\n\n linear_tile_idx += gridDim.x;\n tile_idx_m = linear_tile_idx % tiles_in_m;\n tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n);\n tmem_allocator.release_allocation_lock();\n accumulator_pipeline.producer_tail(accumulator_pipe_producer_state);\n tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);\n } else if (is_epilogue_warp) {\n static constexpr int FragmentSize = 256 / sizeof_bits_v;\n\n tmem_allocation_result_barrier.arrive_and_wait();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_epilogue.data() = tmem_base_ptr;\n int thread_idx = threadIdx.x % 128;\n\n auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{}));\n auto tiled_r2g =\n make_tiled_copy_D(Copy_Atom{}, tiled_t2r);\n auto thr_t2r = tiled_t2r.get_slice(thread_idx);\n auto thr_r2g = tiled_r2g.get_slice(thread_idx);\n\n // NVFP4 non-E8 recipe constants and global scales\n static constexpr float fp4_max = 6.0f;\n static constexpr float fp4_max_inv = 1.0f / fp4_max;\n\n // get global amax pointer\n int tensor_id = GetTensorId(&kernel_args, tile_idx_n * 64);\n float *global_amax_ptr = GetGlobalAmaxPtrByTensorId(&kernel_args, tensor_id);\n\n TC *cur_output_colwise_ptr =\n reinterpret_cast(kernel_args.output_colwise_list[tensor_id]);\n TSFC *cur_output_colwise_scale_inv_ptr =\n reinterpret_cast(kernel_args.output_colwise_scale_inv_list[tensor_id]);\n int cur_output_colwise_n = kernel_args.split_sections[tensor_id];\n\n TensorC cur_mC = cute::make_tensor(\n cute::subbyte_iterator(cur_output_colwise_ptr),\n cute::make_shape(static_cast(M), cur_output_colwise_n), // (M, N_i)\n kernel_args.output_stride2d_list[tensor_id]);\n\n auto cur_sfc_shape =\n make_shape(M, make_shape(make_shape(Int<16>{}, _4{}), cur_output_colwise_n / 64));\n\n auto cur_sfc_stride =\n make_stride(cur_output_colwise_n / 16, make_stride(make_stride(_0{}, _1{}), _4{}));\n\n TensorSFC cur_mSFC = cute::make_tensor(make_gmem_ptr(cur_output_colwise_scale_inv_ptr),\n make_layout(cur_sfc_shape, cur_sfc_stride));\n\n TensorGC cur_gC_mn = local_tile(\n cur_mC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N)\n );\n\n TensorGSFC cur_gSFC_mn = local_tile(\n cur_mSFC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N-like)\n );\n\n Tensor tCgC = thr_mma_epilogue.partition_C(cur_gC_mn);\n\n float global_amax_val = *global_amax_ptr;\n float global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val);\n\n // Scaling factor for fast math path\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n\n float global_decode_scale = 1.0f / global_encode_scale;\n\n auto sfd_converter = cutlass::NumericConverter{};\n\n do {\n for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n; ++k_tile) {\n // get the starting index of current k-tile in global tensor, to query the correct global amax\n int cur_k_tile_global_elem_idx = (tile_idx_n + k_tile) * 64;\n int new_tensor_id = GetTensorId(&kernel_args, cur_k_tile_global_elem_idx);\n // float* new_global_amax_ptr = GetGlobalAmaxPtr(&kernel_args, cur_k_tile_global_elem_idx);\n global_amax_ptr = GetGlobalAmaxPtrByTensorId(&kernel_args, new_tensor_id);\n // update the scaling factors when it's no longer the same amax pointer\n // TODO(zhongbo): the math operations are very expensive\n // since the kernel is persistent, we can have a cache for all the possible scaling factors\n if (tensor_id != new_tensor_id) {\n global_amax_val = *global_amax_ptr;\n global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val);\n global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n global_decode_scale = 1.0f / global_encode_scale;\n tensor_id = new_tensor_id;\n // went through the cute operations to update the local tensors\n cur_output_colwise_ptr =\n reinterpret_cast(kernel_args.output_colwise_list[tensor_id]);\n cur_output_colwise_scale_inv_ptr =\n reinterpret_cast(kernel_args.output_colwise_scale_inv_list[tensor_id]);\n cur_output_colwise_n = kernel_args.split_sections[tensor_id];\n\n cur_mC = cute::make_tensor(\n cute::subbyte_iterator(cur_output_colwise_ptr),\n cute::make_shape(static_cast(M), cur_output_colwise_n), // (M, N_i)\n kernel_args.output_stride2d_list[tensor_id]);\n\n cur_sfc_shape =\n make_shape(M, make_shape(make_shape(Int<16>{}, _4{}), cur_output_colwise_n / 64));\n\n cur_sfc_stride =\n make_stride(cur_output_colwise_n / 16, make_stride(make_stride(_0{}, _1{}), _4{}));\n\n cur_mSFC = cute::make_tensor(make_gmem_ptr(cur_output_colwise_scale_inv_ptr),\n make_layout(cur_sfc_shape, cur_sfc_stride));\n\n cur_gC_mn = local_tile(\n cur_mC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N)\n );\n\n cur_gSFC_mn = local_tile(cur_mSFC, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{} // (BLK_M, BLK_N-like)\n );\n\n tCgC = thr_mma_epilogue.partition_C(cur_gC_mn);\n }\n // maybe udpated to the new tensor id\n int tensor_start_elem = kernel_args.split_sections_range[tensor_id];\n int local_tile_idx_n = (cur_k_tile_global_elem_idx - tensor_start_elem) / 64;\n\n Tensor tCgC_mn = tCgC(_, _, _, tile_idx_m, local_tile_idx_n);\n Tensor tCgSFC_mn = cur_gSFC_mn(_, _, tile_idx_m, local_tile_idx_n);\n\n accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state);\n\n auto tCtC = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index());\n Tensor tDtC = thr_t2r.partition_S(tCtC); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDgC = thr_t2r.partition_D(tCgC_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n\n Tensor tTR_rAcc =\n make_tensor(shape(tDgC)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDrC = make_tensor(shape(tDgC));\n Tensor tTR_rAcc_frag =\n recast>(coalesce(tTR_rAcc));\n Tensor tDrC_frag = recast>(coalesce(tDrC));\n\n Tensor src = thr_r2g.retile_S(tDrC);\n Tensor dst = thr_r2g.retile_D(tDgC);\n\n Tensor tCgSFC = make_tensor(\n tCgSFC_mn.data(), make_layout(make_shape(shape(tCgSFC_mn), Int<1>{}, Int<1>{}),\n make_stride(stride(tCgSFC_mn), Int<0>{}, Int<0>{})));\n\n Tensor tDgSFC = filter(thr_t2r.partition_D(tCgSFC));\n Tensor tDrSFC = make_tensor(shape(tDgSFC));\n\n static constexpr int NumVecs = size(tDgC) / VectorSize;\n Tensor tC_rRowSFD_frg = recast>(tDrSFC);\n\n cutlass::maximum_absolute_value_reduction,\n true>\n amax_reduction;\n cutlass::Array vec_maxs;\n cutlass::Array pvscales;\n // TMEM_LOAD\n copy(tiled_t2r, tDtC, tTR_rAcc);\n cutlass::arch::fence_view_async_tmem_load();\n\n accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state);\n\n ++accumulator_pipe_consumer_state;\n\n if constexpr (!kUseFastMath) {\n // Downcast to BF16 for bit-wise compatibility with unfused\n // kernels\n auto convert_accum_to_bf16 =\n cutlass::NumericArrayConverter{};\n auto convert_bf16_to_accum =\n cutlass::NumericArrayConverter{};\n tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{})));\n }\n\n auto compute_frgs = reinterpret_cast *>(\n tTR_rAcc_frag.data());\n auto output_frgs = reinterpret_cast *>(tDrC_frag.data());\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]);\n }\n\n pvscales = cutlass::multiplies>{}(\n vec_maxs, global_encode_scale_multiplier);\n auto pvscales_cvted =\n cutlass::NumericArrayConverter{}(pvscales);\n\n tC_rRowSFD_frg(_0{}) = pvscales_cvted;\n auto qpvscale_ups = cutlass::NumericArrayConverter{}(\n tC_rRowSFD_frg(_0{}));\n auto qpvscale_scaled = cutlass::multiplies>{}(\n qpvscale_ups, global_decode_scale);\n cutlass::Array acc_scales;\n if constexpr (kUseFastMath) {\n // Fast math: compute approximate reciprocal\n acc_scales =\n cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled);\n } else {\n // Accurate math: compute reciprocal with division\n acc_scales = cutlass::divides>{}(\n 1.0, qpvscale_scaled);\n }\n\n // Initialize RNG for tile\n const size_t rng_sequence =\n thread_idx + k_tile * 256 + linear_tile_idx * K_TILE_MAX * 256;\n\n transformer_engine::curanddx::detail::philox4x32_native_state<\n NVTE_BUILD_NUM_PHILOX_ROUNDS>\n rng;\n rng.init(rng_seed, rng_sequence, rng_offset);\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(\n acc_scales[v], cutlass::platform::numeric_limits::max());\n // auto acc_scale = acc_scales[v];\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(\n cutlass::multiplies>{}(\n compute_frgs[v], acc_scale),\n reinterpret_cast *>(&random_uint4));\n } else {\n output_frgs[v] = cutlass::NumericArrayConverter{}(\n cutlass::multiplies>{}(\n compute_frgs[v], acc_scale));\n }\n }\n\n copy(tiled_r2g, src, dst);\n\n // copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrC, tDgC);\n\n copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFC, tDgSFC);\n }\n linear_tile_idx += gridDim.x;\n tile_idx_m = linear_tile_idx % tiles_in_m;\n tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n);\n }\n }\n}\n\n// this function computes RHT-GEMM for\n// A: m x n: col-major\n// B: 16 x 16: row-major\n// C: m x n: row-major\n// SFC: m x (n/16): row-major\ntemplate \nvoid group_rht_gemm_ntt_w_sfc(int m, int n, TA const *A, TB const *B,\n MultiAmaxHadamardCastFusionArgs *kernel_args_ptr,\n const size_t *rng_state, uint32_t sm_count, cudaStream_t stream,\n int k_tile_size = 2048) {\n using namespace cute;\n\n // Define shapes (dynamic)\n auto M = static_cast(m);\n auto N = static_cast(n);\n\n // Define strides (mixed)\n auto dA = make_stride(Int<1>{}, m); // (dM,dK)\n auto dB = make_stride(Int<1>{}, 16); // (dN,dK)\n for (size_t i = 0; i < kernel_args_ptr->num_tensors; ++i) {\n kernel_args_ptr->output_stride2d_list[i] =\n make_stride(kernel_args_ptr->split_sections[i], Int<1>{});\n }\n\n auto cga_shape = Shape<_1, _1, _1>{};\n auto cga_tile_shape = Shape<_128, _16, _16>{};\n auto cluster_tile_mainloop = Shape<_128, _16, _64>{};\n\n // Construct the MMA\n auto mma = make_tiled_mma(\n SM100_MMA_F16BF16_SS{},\n Layout>{});\n\n // MMA in CGA Layout XXX: Need to generalize synchro? {$nv-release-never}\n\n // Assert that the TiledMMA uses all CTAs in the CGA.\n CUTE_STATIC_ASSERT_V(size(cga_shape) == size(mma));\n CUTE_STATIC_ASSERT_V(evenly_divides(cga_tile_shape, tile_shape(mma)));\n\n // Determine the A and B shapes\n auto mma_shape_B =\n partition_shape_B(mma, make_shape(size<1>(cga_tile_shape), size<2>(cga_tile_shape)));\n\n using TiledMma = decltype(mma);\n using AtomThrID = typename TiledMma::AtomThrID;\n\n using SmemShape_M = decltype(shape_div(\n shape<0>(cga_tile_shape),\n shape_div(shape<0>(cga_tile_shape), size<0>(cga_tile_shape) / size(AtomThrID{}))));\n using SmemShape_N = decltype(shape_div(\n shape<1>(cga_tile_shape),\n shape_div(shape<1>(cga_tile_shape), size<1>(cga_tile_shape) / size(AtomThrID{}))));\n using SmemShape_K = decltype(cute::get<2>(cga_tile_shape));\n\n using SmemLayoutAtomB =\n decltype(cutlass::gemm::collective::detail::sm100_smem_selector());\n\n auto mma_shape_A = partition_shape_A(\n mma, make_shape(size<0>(cluster_tile_mainloop), size<2>(cluster_tile_mainloop)));\n using SmemShape_M_A =\n decltype(shape_div(shape<0>(cluster_tile_mainloop),\n shape_div(shape<0>(cluster_tile_mainloop),\n size<0>(cluster_tile_mainloop) / size(AtomThrID{}))));\n using SmemShape_K_A = decltype(cute::get<2>(cluster_tile_mainloop));\n using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<\n cute::UMMA::Major::MN, TA, SmemShape_M_A, SmemShape_K_A>());\n\n // Define the smem layouts (static)\n // Calculate max pipeline stages based on Blackwell SM100's 232KB shared memory\n constexpr int kBlackwellSmemSize = 232448; // 232KB in bytes\n constexpr int kBytesPerStage =\n cute::size(mma_shape_A) * sizeof(TA) + cute::size(mma_shape_B) * sizeof(TB);\n constexpr int kReservedBytes = 256; // Reserve for barriers and other uses\n constexpr int kMaxStages = (kBlackwellSmemSize - kReservedBytes) / kBytesPerStage;\n auto sP = Int{}; // SMEM pipelines\n auto sA = UMMA::tile_to_mma_shape(SmemLayoutAtomA{},\n append(mma_shape_A, sP)); // (MMA,MMA_M,MMA_K,PIPE)\n auto sB = UMMA::tile_to_mma_shape(SmemLayoutAtomB{},\n append(mma_shape_B, sP)); // (MMA,MMA_N,MMA_K,PIPE)\n auto sC = Layout<_1>{}; // XXX Dummy\n\n // Create GMEM tensors\n Tensor tensorA = make_tensor(A, make_layout(make_shape(M, N), dA)); // (M,N)\n Tensor tensorB = make_tensor(B, make_layout(make_shape(16, 16), dB)); // (16,16)\n\n // Create the TiledCopy\n\n auto tma_load_a =\n make_tma_copy_A_sm100(SM90_TMA_LOAD{}, tensorA, sA(_, _, _, 0), cluster_tile_mainloop, mma);\n auto tma_load_b =\n make_tma_copy_B_sm100(SM90_TMA_LOAD{}, tensorB, sB(_, _, _, 0), cga_tile_shape, mma);\n\n // Assert checks on tile sizes -- no predication\n NVTE_CHECK(M % size<0>(cga_tile_shape) == 0, \"Inner dimension must be divisible by \",\n static_cast(size<0>(cga_tile_shape)), \" but got \", M, \".\");\n NVTE_CHECK(N % (4 * size<1>(cga_tile_shape)) == 0, \"Outer dimension must be divisible by \",\n 4 * static_cast(size<1>(cga_tile_shape)), \" but got \", N, \".\");\n\n uint32_t tiles = size(ceil_div(M, get<0>(cga_tile_shape))) * size(ceil_div(N, k_tile_size));\n\n tiles = (tiles < sm_count) ? tiles : sm_count;\n\n dim3 dimBlock(256);\n dim3 dimCluster(size<0>(cga_shape), size<1>(cga_shape), size<2>(cga_shape));\n dim3 dimGrid(tiles, 1, 1);\n\n int smem_size = sizeof(SharedStorage);\n auto *kernel_ptr = &group_rht_gemm_device<\n decltype(M), decltype(N), decltype(k_tile_size), decltype(cga_tile_shape), TA, decltype(dA),\n decltype(sA), decltype(tma_load_a), TB, decltype(dB), decltype(sB), decltype(tma_load_b), TC,\n Stride2D, decltype(sC), TSFC, decltype(mma), kEnableStochasticRounding, kUseFastMath>;\n\n NVTE_CHECK_CUDA(\n cudaFuncSetAttribute(*kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n\n (*kernel_ptr)<<>>(M, N, k_tile_size, cga_tile_shape, A, dA,\n sA, tma_load_a, B, dB, sB, tma_load_b, sC,\n mma, *kernel_args_ptr, rng_state);\n NVTE_CHECK_CUDA(cudaGetLastError());\n}\n\n// this function is used to wrap the group_rht_gemm_ntt_w_sfc function\n// to transpose the input tensor A\ntemplate \nvoid group_rht_gemm_ttt_wrapper(int m, int n, TA const *A, TB const *B,\n MultiAmaxHadamardCastFusionArgs *kernel_args_ptr,\n const size_t *rng_state, uint32_t sm_count, cudaStream_t stream,\n int k_tile_size = 1024) {\n // in addition to transpose the input tensor A\n // we also need to reshape m, n to at best\n // ultilize as many SMs as possible while keeping\n // a relatively large contiguous dimension.\n // for example, after swapping m, n for transpose purposes,\n // the input / output tensor shapes for RHT-GEMM are:\n // A: n x m: col-major\n // B: 16 x 16: row-major\n // C: n x m: row-major\n // SFC: n x (m/16): row-major\n group_rht_gemm_ntt_w_sfc(\n n, m, A, B, kernel_args_ptr, rng_state, sm_count, stream, k_tile_size);\n}\n\n} // namespace\n} // namespace detail\n\nvoid group_hadamard_transform_cast_fusion_columnwise(\n const Tensor &input_, std::vector &output_list, const size_t *split_sections,\n size_t num_tensors, const Tensor &hadamard_matrix_, QuantizationConfig &quant_config,\n cudaStream_t stream) {\n NVTE_API_CALL(group_hadamard_transform_cast_fusion_columnwise);\n\n using transformer_engine::detail::kMaxTensorsPerKernel;\n using transformer_engine::detail::MultiAmaxHadamardCastFusionArgs;\n\n NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16,\n \"Input tensor must be BF16 tensor, but dtype is \", to_string(input_.dtype()), \".\");\n NVTE_CHECK(input_.dim() >= 2, \"Input must be a 2D tensor.\");\n const SimpleTensor &input = input_.data;\n\n NVTE_CHECK(output_list.size() == num_tensors,\n \"Number of output tensors should match number of tensors.\");\n\n NVTE_CHECK(num_tensors <= kMaxTensorsPerKernel,\n \"Number of tensors should be less than or equal to \", kMaxTensorsPerKernel);\n\n // construct the multi-tensor args\n MultiAmaxHadamardCastFusionArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.split_sections_range[0] = 0;\n for (size_t i = 0; i < num_tensors; ++i) {\n NVTE_CHECK(split_sections[i] % 64 == 0, \"component \", i,\n \" of split_sections should be 64 multiple\");\n if (split_sections[i] == 0) {\n continue;\n }\n kernel_args.global_amax_list[kernel_args.num_tensors] =\n reinterpret_cast(output_list[i]->amax.dptr);\n // TODO(zhongbo): should we change API assumption to use columnwise_data instead of data?\n kernel_args.output_colwise_list[kernel_args.num_tensors] =\n reinterpret_cast(output_list[i]->data.dptr);\n kernel_args.output_colwise_scale_inv_list[kernel_args.num_tensors] =\n reinterpret_cast(output_list[i]->scale_inv.dptr);\n kernel_args.split_sections[kernel_args.num_tensors] = split_sections[i];\n kernel_args.split_sections_range[kernel_args.num_tensors + 1] =\n kernel_args.split_sections_range[kernel_args.num_tensors] + split_sections[i];\n kernel_args.num_tensors++;\n }\n\n // Stochastic rounding config\n const bool use_stochastic_rounding = quant_config.stochastic_rounding;\n const size_t *rng_state = nullptr;\n if (quant_config.rng_state != nullptr) {\n Tensor &rng_state_tensor = *convertNVTETensor(quant_config.rng_state);\n NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64,\n \"RNG state should contain 2 64-bit values.\");\n NVTE_CHECK(rng_state_tensor.data.shape == std::vector{2},\n \"Shape of the RNG state should be [2], but got \", rng_state_tensor.data.shape);\n rng_state = reinterpret_cast(rng_state_tensor.data.dptr);\n }\n\n // Template arguments\n using TA = cute::bfloat16_t;\n using TB = cute::bfloat16_t;\n using TC = cutlass::float_e2m1_t;\n using TSFC = cutlass::float_ue4m3_t;\n\n checkCuDriverContext(stream);\n\n // Check Hadamard matrix\n constexpr int kHadamardDimension = 16;\n\n NVTE_CHECK(hadamard_matrix_.dtype() == transformer_engine::DType::kBFloat16,\n \"Hadamard matrix must be BF16 tensor, but dtype is \",\n to_string(hadamard_matrix_.dtype()), \".\");\n const SimpleTensor &hadamard_matrix = hadamard_matrix_.data;\n NVTE_CHECK(\n (hadamard_matrix_.shape() == std::vector{kHadamardDimension, kHadamardDimension}),\n \"Hadamard matrix must have shape=\",\n std::vector{kHadamardDimension, kHadamardDimension},\n \", but got shape=\", hadamard_matrix_.shape(), \".\");\n const size_t hadamard_dimension = hadamard_matrix.shape[0];\n\n const size_t ndim = input.shape.size();\n const size_t n = input.shape[ndim - 1];\n size_t m = 1;\n for (size_t i = 0; i < ndim - 1; ++i) {\n m *= input.shape[i];\n }\n\n auto sm_count = transformer_engine::cuda::sm_count();\n\n NVTE_CHECK(n % hadamard_dimension == 0, \"row_length must be divisible by hadamard_dimension.\");\n\n NVTE_CHECK(m % hadamard_dimension == 0, \"num_rows must be divisible by hadamard_dimension\");\n\n int k_tile_size = 1024;\n\n if (m == 8192 && n == 5120) {\n k_tile_size = 512;\n } else if (m == 8192 && n == 10240) {\n k_tile_size = 1024;\n } else if (m == 8192 && n == 2560) {\n k_tile_size = 1280;\n } else if (m == 8192 && n == 11328) {\n k_tile_size = 1024;\n } else if (m == 8192 && n == 512) {\n k_tile_size = 256;\n } else if (m == 8192 && n == 3584) {\n k_tile_size = 512;\n } else if (m == 11328 && n == 8192) {\n k_tile_size = 1024;\n } else if (m == 5120 && n == 8192) {\n k_tile_size = 512;\n } else if (m == 10240 && n == 8192) {\n k_tile_size = 1024;\n } else if (m == 2560 && n == 8192) {\n k_tile_size = 1280;\n } else if (m == 512 && n == 8192) {\n k_tile_size = 256;\n } else if (m == 3584 && n == 8192) {\n k_tile_size = 512;\n } else if (m < 1024 || n < 1024) {\n k_tile_size = 512;\n }\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_stochastic_rounding, kUseStochasticRounding,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n quant_config.use_fast_math, kUseFastMath,\n detail::group_rht_gemm_ttt_wrapper(\n /*m=*/m, /*n=*/n, /*A=*/reinterpret_cast(input.dptr),\n /*B=*/reinterpret_cast(hadamard_matrix.dptr),\n /*kernel_args_ptr=*/&kernel_args, /*rng_state=*/rng_state, /*sm_count=*/sm_count,\n /*stream=*/stream, /*k_tile_size=*/k_tile_size);););\n}\n\n} // namespace transformer_engine\n\nvoid nvte_group_hadamard_transform_cast_fusion_columnwise(\n const NVTETensor input, NVTETensor *outputs, const NVTETensor hadamard_matrix,\n const size_t *split_sections, const size_t num_tensors,\n const NVTEQuantizationConfig quant_config, cudaStream_t stream) {\n NVTE_API_CALL(nvte_multi_hadamard_transform_cast_fusion_columnwise);\n using namespace transformer_engine;\n NVTE_CHECK(num_tensors > 0, \"Number of tensors should be greater than 0.\");\n\n Tensor *input_tensor = convertNVTETensorCheck(input);\n std::vector output_list(num_tensors);\n for (size_t i = 0; i < num_tensors; ++i) {\n output_list[i] = convertNVTETensorCheck(outputs[i]);\n }\n\n QuantizationConfig quant_config_cpp;\n if (quant_config != nullptr) {\n quant_config_cpp = *reinterpret_cast(quant_config);\n }\n\n // Call the multi-tensor Hadamard transform amax implementation.\n group_hadamard_transform_cast_fusion_columnwise(\n *input_tensor, output_list, split_sections, num_tensors,\n *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp, stream);\n}\n", "output": "/*************************************************************************\n * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n *\n * See LICENSE for license information.\n ************************************************************************/\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"common/common.h\"\n#include \"common/util/hip/hip_runtime.h\"\n#include \"common/util/curanddx.hpp\"\n#include \"common/util/ptx.cuh\"\n#include \"common/utils.cuh\"\n#include \"cutlass/arch/barrier.h\"\n#include \"cutlass/cutlass.h\"\n#include \"cutlass/gemm/collective/builders/sm100_common.inl\"\n#include \"cutlass/numeric_conversion.h\"\n#include \"cutlass/pipeline/pipeline.hpp\"\n#include \"cutlass/util/GPU_Clock.hpp\"\n#include \"cutlass/util/command_line.h\"\n#include \"cutlass/util/print_error.hpp\"\n\nnamespace transformer_engine {\nnamespace detail {\nnamespace {\n\nusing namespace cute;\nusing cute::\n Tensor; // Ensure unqualified Tensor refers to cute::Tensor, not transformer_engine::Tensor\n\nusing Stride2D = cute::Stride>;\n\nconstexpr int kMaxTensorsPerKernel = 64; // Args must be <4 KB, expand 64 if needed\nstruct MultiAmaxHadamardCastFusionArgs {\n // (output) Amax buffer for pre-RHT amax buffer\n void *global_amax_list[kMaxTensorsPerKernel];\n // output C pointers for each tensor\n void *output_colwise_list[kMaxTensorsPerKernel];\n // output scale inverse pointers for each tensor\n void *output_colwise_scale_inv_list[kMaxTensorsPerKernel];\n // split sections of each tensor of input\n int split_sections[kMaxTensorsPerKernel];\n // Prefix sum (with leading zero) of split_sections of each tensor of input\n int split_sections_range[kMaxTensorsPerKernel + 1];\n // stride 2D struct for CUTE\n Stride2D output_stride2d_list[kMaxTensorsPerKernel];\n // Number of tensors (splits) being processed by kernel\n int num_tensors;\n};\n\n__device__ __forceinline__ float *GetGlobalAmaxPtrByTensorId(\n MultiAmaxHadamardCastFusionArgs *kernel_args_ptr, int tensor_id) {\n // directly returns the global amax pointer by tensor id\n if (tensor_id < 0 || tensor_id >= kernel_args_ptr->num_tensors) {\n return nullptr;\n }\n return reinterpret_cast(kernel_args_ptr->global_amax_list[tensor_id]);\n}\n\n__device__ __forceinline__ int GetTensorId(MultiAmaxHadamardCastFusionArgs *kernel_args_ptr,\n int offset) {\n // Check the kernel args and get the corresponding id\n const int num_tensors = kernel_args_ptr->num_tensors;\n if (offset >= kernel_args_ptr->split_sections_range[num_tensors]) {\n return num_tensors - 1;\n }\n int tensor_id = 0;\n while (kernel_args_ptr->split_sections_range[tensor_id + 1] <= offset) {\n ++tensor_id;\n }\n return tensor_id;\n}\n\n// calculate the global encode scale factor for a given global amax.\n__device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax) {\n constexpr float kFP8E4M3Max = 448.0f;\n constexpr float kFP4E2M1Max = 6.0f;\n // If scale is infinity, return max value of float32\n float global_encode_scale = cutlass::minimum_with_nan_propagation{}(\n kFP8E4M3Max * kFP4E2M1Max / global_amax, cutlass::platform::numeric_limits::max());\n // If global amax is 0 or infinity, return 1\n return (global_amax == 0.f || global_encode_scale == 0.f) ? 1.f : global_encode_scale;\n}\n\ntemplate \nstruct SharedStorage {\n static constexpr int AccumulatorPipelineStageCount = 16;\n using AtomThrShapeMNK = cute::Shape<_1, _1, _1>;\n\n using AccumulatorPipeline =\n cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;\n\n static constexpr int MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline =\n cutlass::PipelineTmaUmmaAsync, AtomThrShapeMNK>;\n using MainloopPipelineStorage = typename MainloopPipeline::SharedStorage;\n\n alignas(16) AccumulatorPipelineStorage accumulator;\n alignas(16) MainloopPipelineStorage mainloop;\n alignas(16) cute::uint64_t tma_barrier[1];\n uint32_t tmem_base_ptr;\n\n struct TensorStorage : cute::aligned_struct<128, _1> {\n // cute::array_aligned> smem_A;\n cute::array_aligned> smem_A;\n cute::array_aligned> smem_B;\n } tensors;\n};\n\nCUTLASS_DEVICE\ncutlass::Array StochasticNumericConverterBase(\n cutlass::Array const &input, cutlass::Array const &rbits) {\n using result_type = cutlass::Array;\n result_type output;\n constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING;\n if constexpr (has_rs) {\n auto output_ptr = reinterpret_cast(&output);\n asm volatile(\n \"{\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\\n\"\n \"cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\\n\"\n \"}\"\n : \"=h\"(output_ptr[0]), \"=h\"(output_ptr[1])\n : \"f\"(input[0]), \"f\"(input[1]), \"f\"(input[2]), \"f\"(input[3]), \"f\"(input[4]), \"f\"(input[5]),\n \"f\"(input[6]), \"f\"(input[7]), \"r\"(rbits[0]), \"r\"(rbits[1]));\n } else {\n NVTE_DEVICE_ERROR(\n \"FP4 cvt PTX instructions are architecture-specific. \"\n \"Try recompiling with sm_XXXa instead of sm_XXX.\");\n }\n return output;\n}\n\nCUTLASS_DEVICE\ncutlass::Array StochasticNumericConverter(\n cutlass::Array const &input, cutlass::Array const *rbits) {\n using result_type = cutlass::Array;\n result_type output;\n cutlass::Array *result_ptr =\n reinterpret_cast *>(&output);\n cutlass::Array const *source_ptr =\n reinterpret_cast const *>(&input);\n cutlass::Array const *rbits_ptr =\n reinterpret_cast const *>(rbits);\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 2; i++) {\n result_ptr[i] = StochasticNumericConverterBase(source_ptr[i], rbits_ptr[i]);\n }\n return output;\n}\n\ntemplate \n__global__ static void group_rht_gemm_device(\n MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, TA const *A, AStride dA,\n ASmemLayout sAlayout, CUTE_GRID_CONSTANT TmaLoadA const tma_load_a, TB const *B, BStride dB,\n BSmemLayout sBlayout, CUTE_GRID_CONSTANT TmaLoadB const tma_load_b, CSmemLayout, TiledMMA mma,\n MultiAmaxHadamardCastFusionArgs kernel_args, const size_t *rng_state) {\n using namespace cute;\n constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY;\n if constexpr (!is_blackwell_arch) {\n NVTE_DEVICE_ERROR(\"RHT fusion is only supported on Blackwell.\");\n return;\n } else {\n using X = Underscore;\n // static constexpr bool kApplyStochasticRounding = true;\n using ElementAccumulator = float;\n static constexpr int K_PIPE_MAX = size<3>(ASmemLayout{});\n using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>;\n static constexpr uint32_t kTmaTransactionBytes = cutlass::bits_to_bytes(\n size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v);\n\n static constexpr int kTmaRhtTensorTransactionBytes =\n cutlass::bits_to_bytes(16 * 16 * cute::sizeof_bits_v);\n static constexpr int AccumulatorPipelineStageCount = 16;\n\n static constexpr int MainloopPipelineStageCount = size<3>(ASmemLayout{});\n using MainloopPipeline = cutlass::PipelineTmaUmmaAsync, AtomThrShapeMNK>;\n using MainloopPipelineState = typename MainloopPipeline::PipelineState;\n\n using TmemAllocator = cute::TMEM::Allocator1Sm;\n static constexpr int VectorSize = 16;\n const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0;\n const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0;\n // Preconditions\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n CUTE_STATIC_ASSERT(is_static::value);\n\n // Represent the full tensors\n Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, N));\n Tensor mB = tma_load_b.get_tma_tensor(make_shape(16, 16));\n\n using TensorC = decltype(make_tensor(subbyte_iterator(recast_ptr(nullptr)), // engine\n make_shape(int{}, int{}), // (M, N_i)\n Stride2D{} // stride (dM, dN)\n ));\n\n using TensorSFC = decltype(make_tensor(\n make_gmem_ptr(recast_ptr(nullptr)),\n make_layout(make_shape(int{}, // M\n make_shape(make_shape(Int<16>{}, _4{}), // (16, 4)\n int{}) // n_tiles = split / 64\n ),\n make_stride(int{}, // dM = (split / 16)\n make_stride(make_stride(_0{}, _1{}), // inner (16,4) layout\n _4{}) // tiles stride\n ))));\n\n auto cluster_shape = Shape<_1, _1, _1>{};\n\n // Get the appropriate blocks for this Cluster\n dim3 cluster_coord_in_grid = cluster_id_in_grid();\n\n // Total number of k-tiles\n const int K_TILE_MAX = min(N, K) / 64;\n uint32_t tiles_in_m = (M + size<0>(cluster_tile) - 1) / size<0>(cluster_tile);\n uint32_t tiles_in_n = (N + 64 - 1) / 64;\n uint32_t linear_tile_idx = blockIdx.x;\n uint32_t tile_idx_m = linear_tile_idx % tiles_in_m;\n uint32_t tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n\n auto mainloop_tiler = Shape<_128, _16, _64>{};\n auto epilogue_tiler = Shape<_128, _64, _64>{};\n Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{});\n Tensor gB_nk =\n local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k)\n // Tensor gC_mn = local_tile(mC, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N)\n\n using TensorGC = decltype(local_tile(std::declval(), decltype(epilogue_tiler){},\n make_coord(_, _, _), Step<_1, _1, X>{}));\n\n using TensorGSFC = decltype(local_tile(std::declval(), decltype(epilogue_tiler){},\n make_coord(_, _, _), Step<_1, _1, X>{}));\n\n // Allocate SMEM\n extern __shared__ char shared_memory[];\n using SharedStorage = SharedStorage;\n SharedStorage &shared_storage = *reinterpret_cast(shared_memory);\n Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()),\n sAlayout); // (MMA,MMA_M,MMA_N,PIPE)\n Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()),\n sBlayout); // (MMA,MMA_N,MMA_K,PIPE)\n\n //\n // MMA: Define C accumulators and A/B partitioning\n //\n\n int block_rank_in_cluster = cute::block_rank_in_cluster();\n ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx\n Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k)\n\n auto mma_epilogue = make_tiled_mma(SM100_MMA_F16BF16_SS{},\n Layout>{});\n ThrMMA thr_mma_epilogue = mma_epilogue.get_slice(block_rank_in_cluster);\n\n using TiledMmaEpilogue = decltype(mma_epilogue);\n Tensor tCgA = thr_mma.partition_A(gA_mk);\n // Allocate \"fragments\" -- these are actually umma smem descriptors\n Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)\n Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE)\n\n auto acc_shape_mma = partition_shape_C(TiledMMA{}, take<0, 2>(ClusterTileShape{}));\n auto acc_shape_epilogue = partition_shape_C(TiledMmaEpilogue{}, take<0, 2>(epilogue_tiler));\n\n auto bulk_tmem_mma =\n TiledMMA::make_fragment_C(append(acc_shape_mma, Int{}));\n\n auto bulk_tmem_epilogue = TiledMmaEpilogue::make_fragment_C(\n append(acc_shape_epilogue, Int{}));\n\n TmemAllocator tmem_allocator{};\n cutlass::arch::NamedBarrier tmem_allocation_result_barrier(\n 32 + 128, cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);\n\n Layout cta_layout_mnk = make_layout(cluster_shape);\n Layout cta_layout_vmnk =\n tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{}));\n auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster);\n\n auto [tAgA, tAsA] =\n tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),\n group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA));\n\n auto [tBgB, tBsB] =\n tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),\n group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB));\n\n uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);\n uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);\n\n int warp_idx = cutlass::canonical_warp_idx_sync();\n\n bool is_mma_warp = (warp_idx == 0);\n bool is_dma_warp = (warp_idx == 1);\n bool is_epilogue_warp = (warp_idx >= 4 && warp_idx <= 7);\n\n // if (is_epilogue_warp && elect_one_sync()) {\n // // prefetch to make the global amax in cache\n // for (size_t i = 0; i < kernel_args.num_tensors; ++i) {\n // cute::prefetch(raw_pointer_cast(kernel_args.global_amax_list[i]));\n // }\n // }\n\n typename MainloopPipeline::Params mainloop_pipeline_params;\n if (is_dma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;\n }\n if (is_mma_warp) {\n mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;\n }\n mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp;\n mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes;\n mainloop_pipeline_params.initializing_warp = 0;\n MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params,\n cluster_shape, cute::true_type{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n\n MainloopPipelineState mainloop_pipe_consumer_state;\n MainloopPipelineState mainloop_pipe_producer_state =\n cutlass::make_producer_start_state();\n\n using AccumulatorPipeline =\n cutlass::PipelineUmmaAsync;\n using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState;\n\n AccumulatorPipelineState accumulator_pipe_consumer_state;\n AccumulatorPipelineState accumulator_pipe_producer_state =\n cutlass::make_producer_start_state();\n\n typename AccumulatorPipeline::Params accumulator_pipeline_params;\n if (is_mma_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer;\n }\n if (is_epilogue_warp) {\n accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer;\n }\n // Only one producer thread arrives on this barrier.\n accumulator_pipeline_params.producer_arv_count = 1;\n accumulator_pipeline_params.consumer_arv_count = size(AtomThrShapeMNK{}) * 128;\n accumulator_pipeline_params.initializing_warp = 1;\n AccumulatorPipeline accumulator_pipeline(shared_storage.accumulator,\n accumulator_pipeline_params, cluster_shape,\n cute::true_type{}, // Perform barrier init\n cute::true_type{}); // Delay mask calculation\n\n if (warp_idx == 2 && elect_one_sync()) {\n cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);\n }\n __syncthreads();\n using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x;\n\n if (is_dma_warp) {\n if (elect_one_sync()) {\n cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0],\n kTmaRhtTensorTransactionBytes);\n copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0),\n tBsB(_, 0));\n }\n\n do {\n bool is_first_wave = linear_tile_idx == blockIdx.x;\n uint32_t skip_wait = is_first_wave;\n auto tAgA_mk = tAgA(_, tile_idx_m, _);\n int k_tile = 0;\n auto barrier_token =\n mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state, skip_wait);\n\n CUTE_NO_UNROLL\n while (k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n) {\n int k_tile_idx_n = tile_idx_n + k_tile;\n ++k_tile;\n skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount);\n mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state, barrier_token);\n using BarrierType = typename MainloopPipeline::ProducerBarrierType;\n BarrierType *tma_barrier =\n mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state);\n int write_stage = mainloop_pipe_producer_state.index();\n ++mainloop_pipe_producer_state;\n barrier_token =\n mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state, skip_wait);\n if (cute::elect_one_sync()) {\n copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n),\n tAsA(_, write_stage));\n }\n }\n linear_tile_idx += gridDim.x;\n tile_idx_m = linear_tile_idx % tiles_in_m;\n tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n);\n mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);\n } else if (is_mma_warp) {\n mma.accumulate_ = UMMA::ScaleOut::Zero;\n\n tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns,\n &shared_storage.tmem_base_ptr);\n __syncwarp();\n tmem_allocation_result_barrier.arrive();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_mma.data() = tmem_base_ptr;\n\n cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/);\n do {\n uint32_t skip_wait = K_TILE_MAX <= 0;\n auto barrier_token =\n mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n CUTE_NO_UNROLL\n for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n;) {\n mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);\n int read_stage = mainloop_pipe_consumer_state.index();\n auto tCrA_mk = tCrA(_, _, _, read_stage);\n auto tCrB_nk = tCrB(_, _, 0, 0);\n CUTE_UNROLL\n for (int k_block = 0; k_block < size<2>(tCrA) / 4; ++k_block) {\n accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);\n CUTE_UNROLL\n for (int i = 0; i < 4; i++) {\n auto accumulators =\n bulk_tmem_mma(_, _, _, accumulator_pipe_producer_state.index() * 4 + i);\n gemm(mma, tCrA_mk(_, _, k_block * 4 + i), tCrB_nk, accumulators);\n }\n\n accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);\n ++accumulator_pipe_producer_state;\n }\n auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;\n ++mainloop_pipe_consumer_state;\n ++k_tile;\n skip_wait = k_tile >= K_TILE_MAX;\n barrier_token =\n mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);\n mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state);\n }\n\n linear_tile_idx += gridDim.x;\n tile_idx_m = linear_tile_idx % tiles_in_m;\n tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n);\n tmem_allocator.release_allocation_lock();\n accumulator_pipeline.producer_tail(accumulator_pipe_producer_state);\n tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);\n } else if (is_epilogue_warp) {\n static constexpr int FragmentSize = 256 / sizeof_bits_v;\n\n tmem_allocation_result_barrier.arrive_and_wait();\n uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;\n bulk_tmem_epilogue.data() = tmem_base_ptr;\n int thread_idx = threadIdx.x % 128;\n\n auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{}));\n auto tiled_r2g =\n make_tiled_copy_D(Copy_Atom{}, tiled_t2r);\n auto thr_t2r = tiled_t2r.get_slice(thread_idx);\n auto thr_r2g = tiled_r2g.get_slice(thread_idx);\n\n // NVFP4 non-E8 recipe constants and global scales\n static constexpr float fp4_max = 6.0f;\n static constexpr float fp4_max_inv = 1.0f / fp4_max;\n\n // get global amax pointer\n int tensor_id = GetTensorId(&kernel_args, tile_idx_n * 64);\n float *global_amax_ptr = GetGlobalAmaxPtrByTensorId(&kernel_args, tensor_id);\n\n TC *cur_output_colwise_ptr =\n reinterpret_cast(kernel_args.output_colwise_list[tensor_id]);\n TSFC *cur_output_colwise_scale_inv_ptr =\n reinterpret_cast(kernel_args.output_colwise_scale_inv_list[tensor_id]);\n int cur_output_colwise_n = kernel_args.split_sections[tensor_id];\n\n TensorC cur_mC = cute::make_tensor(\n cute::subbyte_iterator(cur_output_colwise_ptr),\n cute::make_shape(static_cast(M), cur_output_colwise_n), // (M, N_i)\n kernel_args.output_stride2d_list[tensor_id]);\n\n auto cur_sfc_shape =\n make_shape(M, make_shape(make_shape(Int<16>{}, _4{}), cur_output_colwise_n / 64));\n\n auto cur_sfc_stride =\n make_stride(cur_output_colwise_n / 16, make_stride(make_stride(_0{}, _1{}), _4{}));\n\n TensorSFC cur_mSFC = cute::make_tensor(make_gmem_ptr(cur_output_colwise_scale_inv_ptr),\n make_layout(cur_sfc_shape, cur_sfc_stride));\n\n TensorGC cur_gC_mn = local_tile(\n cur_mC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N)\n );\n\n TensorGSFC cur_gSFC_mn = local_tile(\n cur_mSFC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N-like)\n );\n\n Tensor tCgC = thr_mma_epilogue.partition_C(cur_gC_mn);\n\n float global_amax_val = *global_amax_ptr;\n float global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val);\n\n // Scaling factor for fast math path\n float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n\n float global_decode_scale = 1.0f / global_encode_scale;\n\n auto sfd_converter = cutlass::NumericConverter{};\n\n do {\n for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n; ++k_tile) {\n // get the starting index of current k-tile in global tensor, to query the correct global amax\n int cur_k_tile_global_elem_idx = (tile_idx_n + k_tile) * 64;\n int new_tensor_id = GetTensorId(&kernel_args, cur_k_tile_global_elem_idx);\n // float* new_global_amax_ptr = GetGlobalAmaxPtr(&kernel_args, cur_k_tile_global_elem_idx);\n global_amax_ptr = GetGlobalAmaxPtrByTensorId(&kernel_args, new_tensor_id);\n // update the scaling factors when it's no longer the same amax pointer\n // TODO(zhongbo): the math operations are very expensive\n // since the kernel is persistent, we can have a cache for all the possible scaling factors\n if (tensor_id != new_tensor_id) {\n global_amax_val = *global_amax_ptr;\n global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val);\n global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;\n global_decode_scale = 1.0f / global_encode_scale;\n tensor_id = new_tensor_id;\n // went through the cute operations to update the local tensors\n cur_output_colwise_ptr =\n reinterpret_cast(kernel_args.output_colwise_list[tensor_id]);\n cur_output_colwise_scale_inv_ptr =\n reinterpret_cast(kernel_args.output_colwise_scale_inv_list[tensor_id]);\n cur_output_colwise_n = kernel_args.split_sections[tensor_id];\n\n cur_mC = cute::make_tensor(\n cute::subbyte_iterator(cur_output_colwise_ptr),\n cute::make_shape(static_cast(M), cur_output_colwise_n), // (M, N_i)\n kernel_args.output_stride2d_list[tensor_id]);\n\n cur_sfc_shape =\n make_shape(M, make_shape(make_shape(Int<16>{}, _4{}), cur_output_colwise_n / 64));\n\n cur_sfc_stride =\n make_stride(cur_output_colwise_n / 16, make_stride(make_stride(_0{}, _1{}), _4{}));\n\n cur_mSFC = cute::make_tensor(make_gmem_ptr(cur_output_colwise_scale_inv_ptr),\n make_layout(cur_sfc_shape, cur_sfc_stride));\n\n cur_gC_mn = local_tile(\n cur_mC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N)\n );\n\n cur_gSFC_mn = local_tile(cur_mSFC, epilogue_tiler, make_coord(_, _, _),\n Step<_1, _1, X>{} // (BLK_M, BLK_N-like)\n );\n\n tCgC = thr_mma_epilogue.partition_C(cur_gC_mn);\n }\n // maybe udpated to the new tensor id\n int tensor_start_elem = kernel_args.split_sections_range[tensor_id];\n int local_tile_idx_n = (cur_k_tile_global_elem_idx - tensor_start_elem) / 64;\n\n Tensor tCgC_mn = tCgC(_, _, _, tile_idx_m, local_tile_idx_n);\n Tensor tCgSFC_mn = cur_gSFC_mn(_, _, tile_idx_m, local_tile_idx_n);\n\n accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state);\n\n auto tCtC = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index());\n Tensor tDtC = thr_t2r.partition_S(tCtC); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDgC = thr_t2r.partition_D(tCgC_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n\n Tensor tTR_rAcc =\n make_tensor(shape(tDgC)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N)\n Tensor tDrC = make_tensor(shape(tDgC));\n Tensor tTR_rAcc_frag =\n recast>(coalesce(tTR_rAcc));\n Tensor tDrC_frag = recast>(coalesce(tDrC));\n\n Tensor src = thr_r2g.retile_S(tDrC);\n Tensor dst = thr_r2g.retile_D(tDgC);\n\n Tensor tCgSFC = make_tensor(\n tCgSFC_mn.data(), make_layout(make_shape(shape(tCgSFC_mn), Int<1>{}, Int<1>{}),\n make_stride(stride(tCgSFC_mn), Int<0>{}, Int<0>{})));\n\n Tensor tDgSFC = filter(thr_t2r.partition_D(tCgSFC));\n Tensor tDrSFC = make_tensor(shape(tDgSFC));\n\n static constexpr int NumVecs = size(tDgC) / VectorSize;\n Tensor tC_rRowSFD_frg = recast>(tDrSFC);\n\n cutlass::maximum_absolute_value_reduction,\n true>\n amax_reduction;\n cutlass::Array vec_maxs;\n cutlass::Array pvscales;\n // TMEM_LOAD\n copy(tiled_t2r, tDtC, tTR_rAcc);\n cutlass::arch::fence_view_async_tmem_load();\n\n accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state);\n\n ++accumulator_pipe_consumer_state;\n\n if constexpr (!kUseFastMath) {\n // Downcast to BF16 for bit-wise compatibility with unfused\n // kernels\n auto convert_accum_to_bf16 =\n cutlass::NumericArrayConverter{};\n auto convert_bf16_to_accum =\n cutlass::NumericArrayConverter{};\n tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{})));\n }\n\n auto compute_frgs = reinterpret_cast *>(\n tTR_rAcc_frag.data());\n auto output_frgs = reinterpret_cast *>(tDrC_frag.data());\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]);\n }\n\n pvscales = cutlass::multiplies>{}(\n vec_maxs, global_encode_scale_multiplier);\n auto pvscales_cvted =\n cutlass::NumericArrayConverter{}(pvscales);\n\n tC_rRowSFD_frg(_0{}) = pvscales_cvted;\n auto qpvscale_ups = cutlass::NumericArrayConverter{}(\n tC_rRowSFD_frg(_0{}));\n auto qpvscale_scaled = cutlass::multiplies>{}(\n qpvscale_ups, global_decode_scale);\n cutlass::Array acc_scales;\n if constexpr (kUseFastMath) {\n // Fast math: compute approximate reciprocal\n acc_scales =\n cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled);\n } else {\n // Accurate math: compute reciprocal with division\n acc_scales = cutlass::divides>{}(\n 1.0, qpvscale_scaled);\n }\n\n // Initialize RNG for tile\n const size_t rng_sequence =\n thread_idx + k_tile * 256 + linear_tile_idx * K_TILE_MAX * 256;\n\n transformer_engine::curanddx::detail::philox4x32_native_state<\n NVTE_BUILD_NUM_PHILOX_ROUNDS>\n rng;\n rng.init(rng_seed, rng_sequence, rng_offset);\n uint4 random_uint4 = uint4{0, 0, 0, 0};\n\n CUTLASS_PRAGMA_UNROLL\n for (int v = 0; v < NumVecs; v++) {\n auto acc_scale = cutlass::minimum_with_nan_propagation{}(\n acc_scales[v], cutlass::platform::numeric_limits::max());\n // auto acc_scale = acc_scales[v];\n if constexpr (kEnableStochasticRounding) {\n random_uint4 = rng.generate4();\n output_frgs[v] = StochasticNumericConverter(\n cutlass::multiplies>{}(\n compute_frgs[v], acc_scale),\n reinterpret_cast *>(&random_uint4));\n } else {\n output_frgs[v] = cutlass::NumericArrayConverter{}(\n cutlass::multiplies>{}(\n compute_frgs[v], acc_scale));\n }\n }\n\n copy(tiled_r2g, src, dst);\n\n // copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrC, tDgC);\n\n copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFC, tDgSFC);\n }\n linear_tile_idx += gridDim.x;\n tile_idx_m = linear_tile_idx % tiles_in_m;\n tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX;\n } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n);\n }\n }\n}\n\n// this function computes RHT-GEMM for\n// A: m x n: col-major\n// B: 16 x 16: row-major\n// C: m x n: row-major\n// SFC: m x (n/16): row-major\ntemplate \nvoid group_rht_gemm_ntt_w_sfc(int m, int n, TA const *A, TB const *B,\n MultiAmaxHadamardCastFusionArgs *kernel_args_ptr,\n const size_t *rng_state, uint32_t sm_count, hipStream_t stream,\n int k_tile_size = 2048) {\n using namespace cute;\n\n // Define shapes (dynamic)\n auto M = static_cast(m);\n auto N = static_cast(n);\n\n // Define strides (mixed)\n auto dA = make_stride(Int<1>{}, m); // (dM,dK)\n auto dB = make_stride(Int<1>{}, 16); // (dN,dK)\n for (size_t i = 0; i < kernel_args_ptr->num_tensors; ++i) {\n kernel_args_ptr->output_stride2d_list[i] =\n make_stride(kernel_args_ptr->split_sections[i], Int<1>{});\n }\n\n auto cga_shape = Shape<_1, _1, _1>{};\n auto cga_tile_shape = Shape<_128, _16, _16>{};\n auto cluster_tile_mainloop = Shape<_128, _16, _64>{};\n\n // Construct the MMA\n auto mma = make_tiled_mma(\n SM100_MMA_F16BF16_SS{},\n Layout>{});\n\n // MMA in CGA Layout XXX: Need to generalize synchro? {$nv-release-never}\n\n // Assert that the TiledMMA uses all CTAs in the CGA.\n CUTE_STATIC_ASSERT_V(size(cga_shape) == size(mma));\n CUTE_STATIC_ASSERT_V(evenly_divides(cga_tile_shape, tile_shape(mma)));\n\n // Determine the A and B shapes\n auto mma_shape_B =\n partition_shape_B(mma, make_shape(size<1>(cga_tile_shape), size<2>(cga_tile_shape)));\n\n using TiledMma = decltype(mma);\n using AtomThrID = typename TiledMma::AtomThrID;\n\n using SmemShape_M = decltype(shape_div(\n shape<0>(cga_tile_shape),\n shape_div(shape<0>(cga_tile_shape), size<0>(cga_tile_shape) / size(AtomThrID{}))));\n using SmemShape_N = decltype(shape_div(\n shape<1>(cga_tile_shape),\n shape_div(shape<1>(cga_tile_shape), size<1>(cga_tile_shape) / size(AtomThrID{}))));\n using SmemShape_K = decltype(cute::get<2>(cga_tile_shape));\n\n using SmemLayoutAtomB =\n decltype(cutlass::gemm::collective::detail::sm100_smem_selector());\n\n auto mma_shape_A = partition_shape_A(\n mma, make_shape(size<0>(cluster_tile_mainloop), size<2>(cluster_tile_mainloop)));\n using SmemShape_M_A =\n decltype(shape_div(shape<0>(cluster_tile_mainloop),\n shape_div(shape<0>(cluster_tile_mainloop),\n size<0>(cluster_tile_mainloop) / size(AtomThrID{}))));\n using SmemShape_K_A = decltype(cute::get<2>(cluster_tile_mainloop));\n using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<\n cute::UMMA::Major::MN, TA, SmemShape_M_A, SmemShape_K_A>());\n\n // Define the smem layouts (static)\n // Calculate max pipeline stages based on Blackwell SM100's 232KB shared memory\n constexpr int kBlackwellSmemSize = 232448; // 232KB in bytes\n constexpr int kBytesPerStage =\n cute::size(mma_shape_A) * sizeof(TA) + cute::size(mma_shape_B) * sizeof(TB);\n constexpr int kReservedBytes = 256; // Reserve for barriers and other uses\n constexpr int kMaxStages = (kBlackwellSmemSize - kReservedBytes) / kBytesPerStage;\n auto sP = Int{}; // SMEM pipelines\n auto sA = UMMA::tile_to_mma_shape(SmemLayoutAtomA{},\n append(mma_shape_A, sP)); // (MMA,MMA_M,MMA_K,PIPE)\n auto sB = UMMA::tile_to_mma_shape(SmemLayoutAtomB{},\n append(mma_shape_B, sP)); // (MMA,MMA_N,MMA_K,PIPE)\n auto sC = Layout<_1>{}; // XXX Dummy\n\n // Create GMEM tensors\n Tensor tensorA = make_tensor(A, make_layout(make_shape(M, N), dA)); // (M,N)\n Tensor tensorB = make_tensor(B, make_layout(make_shape(16, 16), dB)); // (16,16)\n\n // Create the TiledCopy\n\n auto tma_load_a =\n make_tma_copy_A_sm100(SM90_TMA_LOAD{}, tensorA, sA(_, _, _, 0), cluster_tile_mainloop, mma);\n auto tma_load_b =\n make_tma_copy_B_sm100(SM90_TMA_LOAD{}, tensorB, sB(_, _, _, 0), cga_tile_shape, mma);\n\n // Assert checks on tile sizes -- no predication\n NVTE_CHECK(M % size<0>(cga_tile_shape) == 0, \"Inner dimension must be divisible by \",\n static_cast(size<0>(cga_tile_shape)), \" but got \", M, \".\");\n NVTE_CHECK(N % (4 * size<1>(cga_tile_shape)) == 0, \"Outer dimension must be divisible by \",\n 4 * static_cast(size<1>(cga_tile_shape)), \" but got \", N, \".\");\n\n uint32_t tiles = size(ceil_div(M, get<0>(cga_tile_shape))) * size(ceil_div(N, k_tile_size));\n\n tiles = (tiles < sm_count) ? tiles : sm_count;\n\n dim3 dimBlock(256);\n dim3 dimCluster(size<0>(cga_shape), size<1>(cga_shape), size<2>(cga_shape));\n dim3 dimGrid(tiles, 1, 1);\n\n int smem_size = sizeof(SharedStorage);\n auto *kernel_ptr = &group_rht_gemm_device<\n decltype(M), decltype(N), decltype(k_tile_size), decltype(cga_tile_shape), TA, decltype(dA),\n decltype(sA), decltype(tma_load_a), TB, decltype(dB), decltype(sB), decltype(tma_load_b), TC,\n Stride2D, decltype(sC), TSFC, decltype(mma), kEnableStochasticRounding, kUseFastMath>;\n\n NVTE_CHECK_CUDA(\n hipFuncSetAttribute(reinterpret_cast(*kernel_ptr), hipFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n\n (*kernel_ptr)<<>>(M, N, k_tile_size, cga_tile_shape, A, dA,\n sA, tma_load_a, B, dB, sB, tma_load_b, sC,\n mma, *kernel_args_ptr, rng_state);\n NVTE_CHECK_CUDA(hipGetLastError());\n}\n\n// this function is used to wrap the group_rht_gemm_ntt_w_sfc function\n// to transpose the input tensor A\ntemplate \nvoid group_rht_gemm_ttt_wrapper(int m, int n, TA const *A, TB const *B,\n MultiAmaxHadamardCastFusionArgs *kernel_args_ptr,\n const size_t *rng_state, uint32_t sm_count, hipStream_t stream,\n int k_tile_size = 1024) {\n // in addition to transpose the input tensor A\n // we also need to reshape m, n to at best\n // ultilize as many SMs as possible while keeping\n // a relatively large contiguous dimension.\n // for example, after swapping m, n for transpose purposes,\n // the input / output tensor shapes for RHT-GEMM are:\n // A: n x m: col-major\n // B: 16 x 16: row-major\n // C: n x m: row-major\n // SFC: n x (m/16): row-major\n group_rht_gemm_ntt_w_sfc(\n n, m, A, B, kernel_args_ptr, rng_state, sm_count, stream, k_tile_size);\n}\n\n} // namespace\n} // namespace detail\n\nvoid group_hadamard_transform_cast_fusion_columnwise(\n const Tensor &input_, std::vector &output_list, const size_t *split_sections,\n size_t num_tensors, const Tensor &hadamard_matrix_, QuantizationConfig &quant_config,\n hipStream_t stream) {\n NVTE_API_CALL(group_hadamard_transform_cast_fusion_columnwise);\n\n using transformer_engine::detail::kMaxTensorsPerKernel;\n using transformer_engine::detail::MultiAmaxHadamardCastFusionArgs;\n\n NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16,\n \"Input tensor must be BF16 tensor, but dtype is \", to_string(input_.dtype()), \".\");\n NVTE_CHECK(input_.dim() >= 2, \"Input must be a 2D tensor.\");\n const SimpleTensor &input = input_.data;\n\n NVTE_CHECK(output_list.size() == num_tensors,\n \"Number of output tensors should match number of tensors.\");\n\n NVTE_CHECK(num_tensors <= kMaxTensorsPerKernel,\n \"Number of tensors should be less than or equal to \", kMaxTensorsPerKernel);\n\n // construct the multi-tensor args\n MultiAmaxHadamardCastFusionArgs kernel_args;\n kernel_args.num_tensors = 0;\n kernel_args.split_sections_range[0] = 0;\n for (size_t i = 0; i < num_tensors; ++i) {\n NVTE_CHECK(split_sections[i] % 64 == 0, \"component \", i,\n \" of split_sections should be 64 multiple\");\n if (split_sections[i] == 0) {\n continue;\n }\n kernel_args.global_amax_list[kernel_args.num_tensors] =\n reinterpret_cast(output_list[i]->amax.dptr);\n // TODO(zhongbo): should we change API assumption to use columnwise_data instead of data?\n kernel_args.output_colwise_list[kernel_args.num_tensors] =\n reinterpret_cast(output_list[i]->data.dptr);\n kernel_args.output_colwise_scale_inv_list[kernel_args.num_tensors] =\n reinterpret_cast(output_list[i]->scale_inv.dptr);\n kernel_args.split_sections[kernel_args.num_tensors] = split_sections[i];\n kernel_args.split_sections_range[kernel_args.num_tensors + 1] =\n kernel_args.split_sections_range[kernel_args.num_tensors] + split_sections[i];\n kernel_args.num_tensors++;\n }\n\n // Stochastic rounding config\n const bool use_stochastic_rounding = quant_config.stochastic_rounding;\n const size_t *rng_state = nullptr;\n if (quant_config.rng_state != nullptr) {\n Tensor &rng_state_tensor = *convertNVTETensor(quant_config.rng_state);\n NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64,\n \"RNG state should contain 2 64-bit values.\");\n NVTE_CHECK(rng_state_tensor.data.shape == std::vector{2},\n \"Shape of the RNG state should be [2], but got \", rng_state_tensor.data.shape);\n rng_state = reinterpret_cast(rng_state_tensor.data.dptr);\n }\n\n // Template arguments\n using TA = cute::bfloat16_t;\n using TB = cute::bfloat16_t;\n using TC = cutlass::float_e2m1_t;\n using TSFC = cutlass::float_ue4m3_t;\n\n checkCuDriverContext(stream);\n\n // Check Hadamard matrix\n constexpr int kHadamardDimension = 16;\n\n NVTE_CHECK(hadamard_matrix_.dtype() == transformer_engine::DType::kBFloat16,\n \"Hadamard matrix must be BF16 tensor, but dtype is \",\n to_string(hadamard_matrix_.dtype()), \".\");\n const SimpleTensor &hadamard_matrix = hadamard_matrix_.data;\n NVTE_CHECK(\n (hadamard_matrix_.shape() == std::vector{kHadamardDimension, kHadamardDimension}),\n \"Hadamard matrix must have shape=\",\n std::vector{kHadamardDimension, kHadamardDimension},\n \", but got shape=\", hadamard_matrix_.shape(), \".\");\n const size_t hadamard_dimension = hadamard_matrix.shape[0];\n\n const size_t ndim = input.shape.size();\n const size_t n = input.shape[ndim - 1];\n size_t m = 1;\n for (size_t i = 0; i < ndim - 1; ++i) {\n m *= input.shape[i];\n }\n\n auto sm_count = transformer_engine::cuda::sm_count();\n\n NVTE_CHECK(n % hadamard_dimension == 0, \"row_length must be divisible by hadamard_dimension.\");\n\n NVTE_CHECK(m % hadamard_dimension == 0, \"num_rows must be divisible by hadamard_dimension\");\n\n int k_tile_size = 1024;\n\n if (m == 8192 && n == 5120) {\n k_tile_size = 512;\n } else if (m == 8192 && n == 10240) {\n k_tile_size = 1024;\n } else if (m == 8192 && n == 2560) {\n k_tile_size = 1280;\n } else if (m == 8192 && n == 11328) {\n k_tile_size = 1024;\n } else if (m == 8192 && n == 512) {\n k_tile_size = 256;\n } else if (m == 8192 && n == 3584) {\n k_tile_size = 512;\n } else if (m == 11328 && n == 8192) {\n k_tile_size = 1024;\n } else if (m == 5120 && n == 8192) {\n k_tile_size = 512;\n } else if (m == 10240 && n == 8192) {\n k_tile_size = 1024;\n } else if (m == 2560 && n == 8192) {\n k_tile_size = 1280;\n } else if (m == 512 && n == 8192) {\n k_tile_size = 256;\n } else if (m == 3584 && n == 8192) {\n k_tile_size = 512;\n } else if (m < 1024 || n < 1024) {\n k_tile_size = 512;\n }\n\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n use_stochastic_rounding, kUseStochasticRounding,\n TRANSFORMER_ENGINE_SWITCH_CONDITION(\n quant_config.use_fast_math, kUseFastMath,\n detail::group_rht_gemm_ttt_wrapper(\n /*m=*/m, /*n=*/n, /*A=*/reinterpret_cast(input.dptr),\n /*B=*/reinterpret_cast(hadamard_matrix.dptr),\n /*kernel_args_ptr=*/&kernel_args, /*rng_state=*/rng_state, /*sm_count=*/sm_count,\n /*stream=*/stream, /*k_tile_size=*/k_tile_size);););\n}\n\n} // namespace transformer_engine\n\nvoid nvte_group_hadamard_transform_cast_fusion_columnwise(\n const NVTETensor input, NVTETensor *outputs, const NVTETensor hadamard_matrix,\n const size_t *split_sections, const size_t num_tensors,\n const NVTEQuantizationConfig quant_config, hipStream_t stream) {\n NVTE_API_CALL(nvte_multi_hadamard_transform_cast_fusion_columnwise);\n using namespace transformer_engine;\n NVTE_CHECK(num_tensors > 0, \"Number of tensors should be greater than 0.\");\n\n Tensor *input_tensor = convertNVTETensorCheck(input);\n std::vector output_list(num_tensors);\n for (size_t i = 0; i < num_tensors; ++i) {\n output_list[i] = convertNVTETensorCheck(outputs[i]);\n }\n\n QuantizationConfig quant_config_cpp;\n if (quant_config != nullptr) {\n quant_config_cpp = *reinterpret_cast(quant_config);\n }\n\n // Call the multi-tensor Hadamard transform amax implementation.\n group_hadamard_transform_cast_fusion_columnwise(\n *input_tensor, output_list, split_sections, num_tensors,\n *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp, stream);\n}\n", "source": "/kaggle/working/TransformerEngine/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"pt_stable_utils.h\"\n#include \"warp_tensor.h\"\n\nusing namespace xformers::sp24;\n\nnamespace {\ntemplate \nvoid __global__ sparseNM_dense_kernel(\n Element const* in_ptr,\n int64_t in_s0,\n Element* out_ptr,\n int64_t out_s0,\n PreprocFn sort_preproc) {\n constexpr int kBlockSize0 = 64;\n constexpr int kBlockSize1 = 64;\n\n int64_t block0 = blockIdx.x * kBlockSize0;\n int64_t block1 = blockIdx.y * kBlockSize1;\n\n WarpTensor tile;\n static_assert(tile.kElementsPerThread == M);\n CUTLASS_PRAGMA_UNROLL\n for (int tile0 = 0; tile0 < kBlockSize0; tile0 += tile.kRows) {\n CUTLASS_PRAGMA_UNROLL\n for (int tile1 = 0; tile1 < kBlockSize1; tile1 += tile.kCols) {\n // load\n tile.load(in_ptr + (block0 + tile0) * in_s0 + block1 + tile1, in_s0);\n\n // sparsify dense\n tile = tile.template sparsify_dense(sort_preproc);\n\n // store\n tile.store(out_ptr + (block0 + tile0) * out_s0 + block1 + tile1, out_s0);\n }\n }\n}\n\ntemplate \ntorch::stable::Tensor sparseNM_dense(\n torch::stable::Tensor input,\n std::string sort_preproc,\n int64_t kN,\n int64_t kM) {\n std::optional device_guard;\n if (!kIsMeta) {\n STD_TORCH_CHECK(input.is_cuda(), \"All tensors must be on GPU\");\n device_guard.emplace(input.device().index());\n }\n\n STD_TORCH_CHECK(input.dim() == 2, \"Can only sparsify 2d tensors\");\n STD_TORCH_CHECK(\n input.stride(1) == 1,\n \"Can only sparsify contiguous tensors. Sparsify the transpose otherwise.\");\n STD_TORCH_CHECK(input.size(0) % 64 == 0);\n STD_TORCH_CHECK(input.size(1) % 64 == 0);\n\n torch::stable::Tensor out = torch::stable::empty_like(input);\n bool foundKernel = false;\n auto launchKernel = [&](auto dtype, auto N, auto M, auto sort_preproc) {\n if (foundKernel) {\n return;\n }\n foundKernel = true;\n if (kIsMeta) {\n return;\n }\n using Element = decltype(dtype);\n dim3 num_blocks(input.size(0) / 64, input.size(1) / 64, 1);\n sparseNM_dense_kernel\n <<>>(\n (Element const*)input.data_ptr(),\n input.stride(0),\n (Element*)out.data_ptr(),\n out.stride(0),\n sort_preproc);\n };\n STD_TORCH_CHECK(\n input.scalar_type() == torch::headeronly::ScalarType::BFloat16);\n auto dtype = cutlass::bfloat16_t();\n auto Identity = [] __device__(auto x) { return x; };\n auto Abs = [] __device__(auto x) { return cutlass::abs(x); };\n auto testNM = [&](auto N, auto M) {\n if (N.value != kN || M.value != kM) {\n return;\n }\n if (sort_preproc == \"largest\") {\n launchKernel(dtype, N, M, Identity);\n } else if (sort_preproc == \"largest_abs\") {\n launchKernel(dtype, N, M, Abs);\n }\n };\n // 2:8 sparsification\n testNM(std::integral_constant(), std::integral_constant());\n // 2:4 sparsification\n testNM(std::integral_constant(), std::integral_constant());\n STD_TORCH_CHECK(foundKernel, \"Kernel not found\");\n STD_CUDA_KERNEL_LAUNCH_CHECK();\n\n return out;\n}\n} // namespace\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, CUDA, m) {\n m.impl(\"sparseNM_dense\", TORCH_BOX(sparseNM_dense));\n}\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, Meta, m) {\n m.impl(\"sparseNM_dense\", TORCH_BOX(sparseNM_dense));\n}\n\nSTABLE_TORCH_LIBRARY_FRAGMENT(xformers, m) {\n m.def(\n \"sparseNM_dense(Tensor input, str sort_preproc, int N, int M) -> Tensor\");\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"pt_stable_utils.h\"\n#include \"warp_tensor.h\"\n\nusing namespace xformers::sp24;\n\nnamespace {\ntemplate \nvoid __global__ sparseNM_dense_kernel(\n Element const* in_ptr,\n int64_t in_s0,\n Element* out_ptr,\n int64_t out_s0,\n PreprocFn sort_preproc) {\n constexpr int kBlockSize0 = 64;\n constexpr int kBlockSize1 = 64;\n\n int64_t block0 = blockIdx.x * kBlockSize0;\n int64_t block1 = blockIdx.y * kBlockSize1;\n\n WarpTensor tile;\n static_assert(tile.kElementsPerThread == M);\n CUTLASS_PRAGMA_UNROLL\n for (int tile0 = 0; tile0 < kBlockSize0; tile0 += tile.kRows) {\n CUTLASS_PRAGMA_UNROLL\n for (int tile1 = 0; tile1 < kBlockSize1; tile1 += tile.kCols) {\n // load\n tile.load(in_ptr + (block0 + tile0) * in_s0 + block1 + tile1, in_s0);\n\n // sparsify dense\n tile = tile.template sparsify_dense(sort_preproc);\n\n // store\n tile.store(out_ptr + (block0 + tile0) * out_s0 + block1 + tile1, out_s0);\n }\n }\n}\n\ntemplate \ntorch::stable::Tensor sparseNM_dense(\n torch::stable::Tensor input,\n std::string sort_preproc,\n int64_t kN,\n int64_t kM) {\n std::optional device_guard;\n if (!kIsMeta) {\n STD_TORCH_CHECK(input.is_cuda(), \"All tensors must be on GPU\");\n device_guard.emplace(input.device().index());\n }\n\n STD_TORCH_CHECK(input.dim() == 2, \"Can only sparsify 2d tensors\");\n STD_TORCH_CHECK(\n input.stride(1) == 1,\n \"Can only sparsify contiguous tensors. Sparsify the transpose otherwise.\");\n STD_TORCH_CHECK(input.size(0) % 64 == 0);\n STD_TORCH_CHECK(input.size(1) % 64 == 0);\n\n torch::stable::Tensor out = torch::stable::empty_like(input);\n bool foundKernel = false;\n auto launchKernel = [&](auto dtype, auto N, auto M, auto sort_preproc) {\n if (foundKernel) {\n return;\n }\n foundKernel = true;\n if (kIsMeta) {\n return;\n }\n using Element = decltype(dtype);\n dim3 num_blocks(input.size(0) / 64, input.size(1) / 64, 1);\n sparseNM_dense_kernel\n <<>>(\n (Element const*)input.data_ptr(),\n input.stride(0),\n (Element*)out.data_ptr(),\n out.stride(0),\n sort_preproc);\n };\n STD_TORCH_CHECK(\n input.scalar_type() == torch::headeronly::ScalarType::BFloat16);\n auto dtype = cutlass::bfloat16_t();\n auto Identity = [] __device__(auto x) { return x; };\n auto Abs = [] __device__(auto x) { return cutlass::abs(x); };\n auto testNM = [&](auto N, auto M) {\n if (N.value != kN || M.value != kM) {\n return;\n }\n if (sort_preproc == \"largest\") {\n launchKernel(dtype, N, M, Identity);\n } else if (sort_preproc == \"largest_abs\") {\n launchKernel(dtype, N, M, Abs);\n }\n };\n // 2:8 sparsification\n testNM(std::integral_constant(), std::integral_constant());\n // 2:4 sparsification\n testNM(std::integral_constant(), std::integral_constant());\n STD_TORCH_CHECK(foundKernel, \"Kernel not found\");\n STD_CUDA_KERNEL_LAUNCH_CHECK();\n\n return out;\n}\n} // namespace\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, CUDA, m) {\n m.impl(\"sparseNM_dense\", TORCH_BOX(sparseNM_dense));\n}\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, Meta, m) {\n m.impl(\"sparseNM_dense\", TORCH_BOX(sparseNM_dense));\n}\n\nSTABLE_TORCH_LIBRARY_FRAGMENT(xformers, m) {\n m.def(\n \"sparseNM_dense(Tensor input, str sort_preproc, int N, int M) -> Tensor\");\n}\n", "source": "/kaggle/working/xformers/xformers/csrc/sparse24/sparseNM_dense.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"compute_sparse_tile.h\"\n#include \"pt_stable_utils.h\"\n#include \"sparse24_pack.h\"\n\nusing namespace xformers::sp24;\n\nnamespace {\ntemplate \nstruct Params {\n uint64_t const* threads_masks;\n\n T const* input;\n int64_t input_stride;\n int64_t input_dim0;\n int64_t input_dim1;\n\n T* output;\n int64_t output_stride;\n\n T mul0;\n T mul1;\n\n __host__ dim3 getBlocksGrid() const {\n return dim3(\n cutlass::ceil_div(input_dim0, kWarpX),\n cutlass::ceil_div(input_dim1, kWarpY),\n 1);\n }\n\n static CUTLASS_HOST_DEVICE dim3 getThreadsGrid() {\n return dim3(kWarpX / kThreadX, kWarpY / kThreadY, 1);\n }\n\n CUTLASS_DEVICE Tile8x8Masks* getCurrentThreadIndices() const {\n Tile8x8Masks* gmem_threads_masks = (Tile8x8Masks*)threads_masks;\n gmem_threads_masks += blockIdx.y * getThreadsGrid().y + threadIdx.y;\n int64_t strideX = gridDim.y * getThreadsGrid().y;\n gmem_threads_masks +=\n (blockIdx.x * getThreadsGrid().x + threadIdx.x) * strideX;\n return gmem_threads_masks;\n }\n};\n\ntemplate \n__global__ void __launch_bounds__(32 /* num_threads */)\n sparse24_apply_dense_output_k(Params p) {\n using Fragment = cutlass::Array;\n\n // Top-left of the 8x8 tile we own\n int warp_x = blockIdx.x * kWarpX;\n int warp_y = blockIdx.y * kWarpY;\n int x = warp_x + threadIdx.x * kThreadX;\n int y = warp_y + threadIdx.y * kThreadY;\n\n T* output = p.output + x * p.output_stride + y;\n Tile8x8Masks indices = *p.getCurrentThreadIndices();\n\n // Load dense\n Fragment lines[8];\n if (kInputRowMajor) {\n T const* input = p.input + x * p.input_stride + y;\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 8; ++i) {\n cutlass::arch::global_load(\n lines[i], input + i * p.input_stride, true);\n }\n } else {\n T const* input = p.input + x + y * p.input_stride;\n Fragment columns[8];\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 8; ++i) {\n cutlass::arch::global_load(\n columns[i], input + i * p.input_stride, true);\n }\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 8; ++i) {\n CUTLASS_PRAGMA_UNROLL\n for (int j = 0; j < 8; ++j) {\n lines[i][j] = columns[j][i].get();\n }\n }\n }\n\n CUTLASS_PRAGMA_UNROLL\n for (int row = 0; row < 2; ++row) {\n Indices4x4 masks[2];\n if (row == 0) {\n masks[0] = indices.a;\n masks[1] = indices.b;\n } else {\n masks[0] = indices.c;\n masks[1] = indices.d;\n }\n\n // Apply mask\n CUTLASS_PRAGMA_UNROLL\n for (int m = 0; m < 2; ++m) {\n CUTLASS_PRAGMA_UNROLL\n for (int r = 0; r < 4; ++r) {\n CUTLASS_PRAGMA_UNROLL\n for (int c = 0; c < 4; ++c) {\n if ((masks[m] >> (4 * r + c)) & 1) {\n lines[4 * row + r][4 * m + c] =\n p.mul1 * lines[4 * row + r][4 * m + c];\n } else {\n lines[4 * row + r][4 * m + c] =\n p.mul0 * lines[4 * row + r][4 * m + c];\n }\n }\n }\n }\n }\n static_assert(kOutputRowMajor, \"Transpose here for ColMajor output\");\n // Save dense with zeros\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 8; ++i) {\n cutlass::arch::global_store(\n lines[i], output + i * p.output_stride, true);\n }\n}\n\ntemplate \ntorch::stable::Tensor sparse24_apply_dense_output_typed(\n torch::stable::Tensor input,\n torch::stable::Tensor threads_masks,\n float mul0,\n float mul1) {\n STD_TORCH_CHECK(\n input.stride(0) == 1 || input.stride(1) == 1,\n \"`input` should be either RowMajor or ColMajor. Invalid memory layout - try .contiguous()?\");\n\n auto roundedx = cutlass::round_up(input.size(0), kWarpX);\n auto roundedy = cutlass::round_up(input.size(1), kWarpY);\n\n Params p;\n p.input_dim0 = input.size(0);\n p.input_dim1 = input.size(1);\n if (!kIsMeta) {\n p.input = (T const*)input.data_ptr();\n p.threads_masks = (uint64_t const*)threads_masks.data_ptr();\n }\n\n STD_TORCH_CHECK(threads_masks.dim() == 3);\n STD_TORCH_CHECK(\n threads_masks.size(0) == p.getBlocksGrid().x * p.getThreadsGrid().x);\n STD_TORCH_CHECK(\n threads_masks.size(1) == p.getBlocksGrid().y * p.getThreadsGrid().y);\n STD_TORCH_CHECK(threads_masks.stride(1) == sizeof(p.threads_masks[0]));\n STD_TORCH_CHECK(threads_masks.size(2) == sizeof(p.threads_masks[0]));\n STD_TORCH_CHECK(threads_masks.stride(2) == 1);\n STD_TORCH_CHECK(\n threads_masks.scalar_type() == torch::headeronly::ScalarType::Byte);\n\n torch::stable::Tensor output =\n torch::stable::new_empty(input, {input.size(0), input.size(1)});\n STD_TORCH_CHECK(output.stride(-1) == 1, \"expected RowMajor?\");\n if (kIsMeta) {\n return output;\n }\n p.output = (T*)output.data_ptr();\n p.mul0 = T(mul0);\n p.mul1 = T(mul1);\n\n bool inputRowMajor = input.stride(-1) == 1;\n bool outputRowMajor = output.stride(-1) == 1;\n p.input_stride = input.stride(inputRowMajor ? 0 : 1);\n p.output_stride = output.stride(outputRowMajor ? 0 : 1);\n torch::stable::accelerator::DeviceGuard device_guard(input.device().index());\n\n cudaStream_t stream = xf_getCurrentCUDAStream();\n size_t smem_bytes = 0;\n if (inputRowMajor && outputRowMajor) {\n sparse24_apply_dense_output_k\n <<>>(p);\n } else if (!inputRowMajor && outputRowMajor) {\n sparse24_apply_dense_output_k\n <<>>(p);\n } else {\n STD_TORCH_CHECK(\n false,\n \"Unsupported configuration: `input` is \",\n inputRowMajor ? \"RowMajor\" : \"ColMajor\",\n \", and `output` is \",\n outputRowMajor ? \"RowMajor\" : \"ColMajor\");\n }\n STD_CUDA_KERNEL_LAUNCH_CHECK();\n return output;\n}\n\ntemplate \ntorch::stable::Tensor sparse24_apply_dense_output(\n torch::stable::Tensor input,\n torch::stable::Tensor threads_masks,\n double mul0,\n double mul1) {\n STD_TORCH_CHECK(\n input.scalar_type() == torch::headeronly::ScalarType::Half ||\n input.scalar_type() == torch::headeronly::ScalarType::BFloat16,\n \"Unsupported `input` dtype\");\n if (input.scalar_type() == torch::headeronly::ScalarType::Half) {\n return sparse24_apply_dense_output_typed(\n input, threads_masks, mul0, mul1);\n } else {\n return sparse24_apply_dense_output_typed(\n input, threads_masks, mul0, mul1);\n }\n}\n} // namespace\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, CUDA, m) {\n m.impl(\n \"sparse24_apply_dense_output\",\n TORCH_BOX(sparse24_apply_dense_output));\n}\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, Meta, m) {\n m.impl(\n \"sparse24_apply_dense_output\",\n TORCH_BOX(sparse24_apply_dense_output));\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"compute_sparse_tile.h\"\n#include \"pt_stable_utils.h\"\n#include \"sparse24_pack.h\"\n\nusing namespace xformers::sp24;\n\nnamespace {\ntemplate \nstruct Params {\n uint64_t const* threads_masks;\n\n T const* input;\n int64_t input_stride;\n int64_t input_dim0;\n int64_t input_dim1;\n\n T* output;\n int64_t output_stride;\n\n T mul0;\n T mul1;\n\n __host__ dim3 getBlocksGrid() const {\n return dim3(\n cutlass::ceil_div(input_dim0, kWarpX),\n cutlass::ceil_div(input_dim1, kWarpY),\n 1);\n }\n\n static CUTLASS_HOST_DEVICE dim3 getThreadsGrid() {\n return dim3(kWarpX / kThreadX, kWarpY / kThreadY, 1);\n }\n\n CUTLASS_DEVICE Tile8x8Masks* getCurrentThreadIndices() const {\n Tile8x8Masks* gmem_threads_masks = (Tile8x8Masks*)threads_masks;\n gmem_threads_masks += blockIdx.y * getThreadsGrid().y + threadIdx.y;\n int64_t strideX = gridDim.y * getThreadsGrid().y;\n gmem_threads_masks +=\n (blockIdx.x * getThreadsGrid().x + threadIdx.x) * strideX;\n return gmem_threads_masks;\n }\n};\n\ntemplate \n__global__ void __launch_bounds__(32 /* num_threads */)\n sparse24_apply_dense_output_k(Params p) {\n using Fragment = cutlass::Array;\n\n // Top-left of the 8x8 tile we own\n int warp_x = blockIdx.x * kWarpX;\n int warp_y = blockIdx.y * kWarpY;\n int x = warp_x + threadIdx.x * kThreadX;\n int y = warp_y + threadIdx.y * kThreadY;\n\n T* output = p.output + x * p.output_stride + y;\n Tile8x8Masks indices = *p.getCurrentThreadIndices();\n\n // Load dense\n Fragment lines[8];\n if (kInputRowMajor) {\n T const* input = p.input + x * p.input_stride + y;\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 8; ++i) {\n cutlass::arch::global_load(\n lines[i], input + i * p.input_stride, true);\n }\n } else {\n T const* input = p.input + x + y * p.input_stride;\n Fragment columns[8];\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 8; ++i) {\n cutlass::arch::global_load(\n columns[i], input + i * p.input_stride, true);\n }\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 8; ++i) {\n CUTLASS_PRAGMA_UNROLL\n for (int j = 0; j < 8; ++j) {\n lines[i][j] = columns[j][i].get();\n }\n }\n }\n\n CUTLASS_PRAGMA_UNROLL\n for (int row = 0; row < 2; ++row) {\n Indices4x4 masks[2];\n if (row == 0) {\n masks[0] = indices.a;\n masks[1] = indices.b;\n } else {\n masks[0] = indices.c;\n masks[1] = indices.d;\n }\n\n // Apply mask\n CUTLASS_PRAGMA_UNROLL\n for (int m = 0; m < 2; ++m) {\n CUTLASS_PRAGMA_UNROLL\n for (int r = 0; r < 4; ++r) {\n CUTLASS_PRAGMA_UNROLL\n for (int c = 0; c < 4; ++c) {\n if ((masks[m] >> (4 * r + c)) & 1) {\n lines[4 * row + r][4 * m + c] =\n p.mul1 * lines[4 * row + r][4 * m + c];\n } else {\n lines[4 * row + r][4 * m + c] =\n p.mul0 * lines[4 * row + r][4 * m + c];\n }\n }\n }\n }\n }\n static_assert(kOutputRowMajor, \"Transpose here for ColMajor output\");\n // Save dense with zeros\n CUTLASS_PRAGMA_UNROLL\n for (int i = 0; i < 8; ++i) {\n cutlass::arch::global_store(\n lines[i], output + i * p.output_stride, true);\n }\n}\n\ntemplate \ntorch::stable::Tensor sparse24_apply_dense_output_typed(\n torch::stable::Tensor input,\n torch::stable::Tensor threads_masks,\n float mul0,\n float mul1) {\n STD_TORCH_CHECK(\n input.stride(0) == 1 || input.stride(1) == 1,\n \"`input` should be either RowMajor or ColMajor. Invalid memory layout - try .contiguous()?\");\n\n auto roundedx = cutlass::round_up(input.size(0), kWarpX);\n auto roundedy = cutlass::round_up(input.size(1), kWarpY);\n\n Params p;\n p.input_dim0 = input.size(0);\n p.input_dim1 = input.size(1);\n if (!kIsMeta) {\n p.input = (T const*)input.data_ptr();\n p.threads_masks = (uint64_t const*)threads_masks.data_ptr();\n }\n\n STD_TORCH_CHECK(threads_masks.dim() == 3);\n STD_TORCH_CHECK(\n threads_masks.size(0) == p.getBlocksGrid().x * p.getThreadsGrid().x);\n STD_TORCH_CHECK(\n threads_masks.size(1) == p.getBlocksGrid().y * p.getThreadsGrid().y);\n STD_TORCH_CHECK(threads_masks.stride(1) == sizeof(p.threads_masks[0]));\n STD_TORCH_CHECK(threads_masks.size(2) == sizeof(p.threads_masks[0]));\n STD_TORCH_CHECK(threads_masks.stride(2) == 1);\n STD_TORCH_CHECK(\n threads_masks.scalar_type() == torch::headeronly::ScalarType::Byte);\n\n torch::stable::Tensor output =\n torch::stable::new_empty(input, {input.size(0), input.size(1)});\n STD_TORCH_CHECK(output.stride(-1) == 1, \"expected RowMajor?\");\n if (kIsMeta) {\n return output;\n }\n p.output = (T*)output.data_ptr();\n p.mul0 = T(mul0);\n p.mul1 = T(mul1);\n\n bool inputRowMajor = input.stride(-1) == 1;\n bool outputRowMajor = output.stride(-1) == 1;\n p.input_stride = input.stride(inputRowMajor ? 0 : 1);\n p.output_stride = output.stride(outputRowMajor ? 0 : 1);\n torch::stable::accelerator::DeviceGuard device_guard(input.device().index());\n\n hipStream_t stream = xf_getCurrentCUDAStream();\n size_t smem_bytes = 0;\n if (inputRowMajor && outputRowMajor) {\n sparse24_apply_dense_output_k\n <<>>(p);\n } else if (!inputRowMajor && outputRowMajor) {\n sparse24_apply_dense_output_k\n <<>>(p);\n } else {\n STD_TORCH_CHECK(\n false,\n \"Unsupported configuration: `input` is \",\n inputRowMajor ? \"RowMajor\" : \"ColMajor\",\n \", and `output` is \",\n outputRowMajor ? \"RowMajor\" : \"ColMajor\");\n }\n STD_CUDA_KERNEL_LAUNCH_CHECK();\n return output;\n}\n\ntemplate \ntorch::stable::Tensor sparse24_apply_dense_output(\n torch::stable::Tensor input,\n torch::stable::Tensor threads_masks,\n double mul0,\n double mul1) {\n STD_TORCH_CHECK(\n input.scalar_type() == torch::headeronly::ScalarType::Half ||\n input.scalar_type() == torch::headeronly::ScalarType::BFloat16,\n \"Unsupported `input` dtype\");\n if (input.scalar_type() == torch::headeronly::ScalarType::Half) {\n return sparse24_apply_dense_output_typed(\n input, threads_masks, mul0, mul1);\n } else {\n return sparse24_apply_dense_output_typed(\n input, threads_masks, mul0, mul1);\n }\n}\n} // namespace\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, CUDA, m) {\n m.impl(\n \"sparse24_apply_dense_output\",\n TORCH_BOX(sparse24_apply_dense_output));\n}\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, Meta, m) {\n m.impl(\n \"sparse24_apply_dense_output\",\n TORCH_BOX(sparse24_apply_dense_output));\n}\n", "source": "/kaggle/working/xformers/xformers/csrc/sparse24/sparse24_apply_dense_output.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"compute_sparse_tile.h\"\n#include \"pt_stable_utils.h\"\n#include \"sparse24_metadata.h\"\n#include \"sparse24_pack.h\"\n\nusing namespace xformers::sp24;\n\nnamespace {\ntemplate \n__global__ void __launch_bounds__(32 /* num_threads */, 20)\n sparse24_sparsify_both_ways_kernel(\n typename KT::Params p,\n Metadata metadata,\n Algorithm algo) {\n KT::sparse24_sparsify_both_ways_kernel(p, metadata, algo);\n}\n\ntemplate \nstd::\n tuple<\n torch::stable::Tensor, // packed\n torch::stable::Tensor, // packed_meta_reordered\n torch::stable::Tensor, // packed_trans\n torch::stable::Tensor, // packed_trans_meta_reordered\n torch::stable::Tensor // threads_masks\n >\n sparse24_sparsify_both_ways_typed(\n const torch::stable::Tensor input,\n std::string algorithm) {\n using KT = KernelTypes;\n std::optional device_guard;\n if (!kIsMeta) {\n device_guard.emplace(input.device().index());\n }\n\n STD_TORCH_CHECK(input.dim() == 2, \"Can only sparsify 2d tensors\");\n STD_TORCH_CHECK(\n input.stride(1) == 1,\n \"Can only sparsify contiguous tensors. Sparsify the transpose otherwise.\");\n\n auto rows = input.size(0);\n auto cols = input.size(1);\n\n auto [compressed, packed, packed_meta_reordered] =\n MetadataFormat::create_compressed_representation(rows, cols, input, true);\n auto [compressed_trans, packed_trans, packed_trans_meta_reordered] =\n MetadataFormat::create_compressed_representation(cols, rows, input, true);\n STD_TORCH_CHECK(\n input.size(1) % 32 == 0, \"Number of cols should be multiple of 32\");\n\n typename KT::Params p;\n p.input_s0 = input.stride(0);\n p.input_dim0 = input.size(0);\n p.input_dim1 = input.size(1);\n p.packed_stride = packed.stride(0);\n p.packed_trans_stride = packed_trans.stride(0);\n\n MetadataFormat metadata = MetadataFormat(\n packed_meta_reordered, packed_trans_meta_reordered, rows, cols);\n torch::stable::Tensor threads_masks = torch::stable::new_empty(\n input,\n {p.getBlocksGrid().x * p.getThreadsGrid().x,\n p.getBlocksGrid().y * p.getThreadsGrid().y,\n sizeof(p.threads_masks[0])},\n torch::headeronly::ScalarType::Byte);\n if (!kIsMeta) {\n p.input = (Element const*)input.data_ptr();\n p.packed = (Element*)packed.data_ptr();\n p.packed_trans = (Element*)packed_trans.data_ptr();\n p.threads_masks = (uint64_t*)threads_masks.data_ptr();\n }\n\n bool kernel_launched = false;\n auto launchKernel = [&](auto algo, std::string const& algo_name) {\n if (algo_name == algorithm) {\n kernel_launched = true;\n if (kIsMeta) {\n return;\n }\n size_t smem_bytes = 0;\n sparse24_sparsify_both_ways_kernel\n <<>>(p, metadata, algo);\n }\n };\n named_algorithms(launchKernel);\n STD_TORCH_CHECK(kernel_launched, \"Unknown algorithm \\\"\", algorithm, \"\\\"\");\n STD_CUDA_KERNEL_LAUNCH_CHECK();\n return std::make_tuple(\n compressed,\n packed_meta_reordered,\n compressed_trans,\n packed_trans_meta_reordered,\n threads_masks);\n}\n\ntemplate \nstd::\n tuple<\n torch::stable::Tensor, // packed\n torch::stable::Tensor, // packed_meta_reordered\n torch::stable::Tensor, // packed_trans\n torch::stable::Tensor, // packed_trans_meta_reordered\n torch::stable::Tensor // threads_masks\n >\n sparse24_sparsify_both_ways(\n const torch::stable::Tensor input,\n std::string algorithm,\n std::string backend) {\n auto runTyped = [&](auto type) {\n using ElementT = decltype(type);\n if (backend == \"cusparselt\") {\n return sparse24_sparsify_both_ways_typed<\n ElementT,\n MetadataCuSparseLtSm80,\n kIsMeta>(input, algorithm);\n } else {\n STD_TORCH_CHECK(\n backend == \"cutlass\",\n \"backend argument only supports `cutlass` or `cusparselt`\");\n return sparse24_sparsify_both_ways_typed<\n ElementT,\n MetadataCutlassSm80,\n kIsMeta>(input, algorithm);\n }\n };\n\n if (input.scalar_type() == torch::headeronly::ScalarType::Half) {\n return runTyped(cutlass::half_t());\n } else {\n STD_TORCH_CHECK(\n input.scalar_type() == torch::headeronly::ScalarType::Half ||\n input.scalar_type() == torch::headeronly::ScalarType::BFloat16);\n return runTyped(cutlass::bfloat16_t());\n }\n}\n} // namespace\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, CUDA, m) {\n m.impl(\n \"sparse24_sparsify_both_ways\",\n TORCH_BOX(sparse24_sparsify_both_ways));\n}\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, Meta, m) {\n m.impl(\n \"sparse24_sparsify_both_ways\",\n TORCH_BOX(sparse24_sparsify_both_ways));\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"compute_sparse_tile.h\"\n#include \"pt_stable_utils.h\"\n#include \"sparse24_metadata.h\"\n#include \"sparse24_pack.h\"\n\nusing namespace xformers::sp24;\n\nnamespace {\ntemplate \n__global__ void __launch_bounds__(32 /* num_threads */, 20)\n sparse24_sparsify_both_ways_kernel(\n typename KT::Params p,\n Metadata metadata,\n Algorithm algo) {\n KT::sparse24_sparsify_both_ways_kernel(p, metadata, algo);\n}\n\ntemplate \nstd::\n tuple<\n torch::stable::Tensor, // packed\n torch::stable::Tensor, // packed_meta_reordered\n torch::stable::Tensor, // packed_trans\n torch::stable::Tensor, // packed_trans_meta_reordered\n torch::stable::Tensor // threads_masks\n >\n sparse24_sparsify_both_ways_typed(\n const torch::stable::Tensor input,\n std::string algorithm) {\n using KT = KernelTypes;\n std::optional device_guard;\n if (!kIsMeta) {\n device_guard.emplace(input.device().index());\n }\n\n STD_TORCH_CHECK(input.dim() == 2, \"Can only sparsify 2d tensors\");\n STD_TORCH_CHECK(\n input.stride(1) == 1,\n \"Can only sparsify contiguous tensors. Sparsify the transpose otherwise.\");\n\n auto rows = input.size(0);\n auto cols = input.size(1);\n\n auto [compressed, packed, packed_meta_reordered] =\n MetadataFormat::create_compressed_representation(rows, cols, input, true);\n auto [compressed_trans, packed_trans, packed_trans_meta_reordered] =\n MetadataFormat::create_compressed_representation(cols, rows, input, true);\n STD_TORCH_CHECK(\n input.size(1) % 32 == 0, \"Number of cols should be multiple of 32\");\n\n typename KT::Params p;\n p.input_s0 = input.stride(0);\n p.input_dim0 = input.size(0);\n p.input_dim1 = input.size(1);\n p.packed_stride = packed.stride(0);\n p.packed_trans_stride = packed_trans.stride(0);\n\n MetadataFormat metadata = MetadataFormat(\n packed_meta_reordered, packed_trans_meta_reordered, rows, cols);\n torch::stable::Tensor threads_masks = torch::stable::new_empty(\n input,\n {p.getBlocksGrid().x * p.getThreadsGrid().x,\n p.getBlocksGrid().y * p.getThreadsGrid().y,\n sizeof(p.threads_masks[0])},\n torch::headeronly::ScalarType::Byte);\n if (!kIsMeta) {\n p.input = (Element const*)input.data_ptr();\n p.packed = (Element*)packed.data_ptr();\n p.packed_trans = (Element*)packed_trans.data_ptr();\n p.threads_masks = (uint64_t*)threads_masks.data_ptr();\n }\n\n bool kernel_launched = false;\n auto launchKernel = [&](auto algo, std::string const& algo_name) {\n if (algo_name == algorithm) {\n kernel_launched = true;\n if (kIsMeta) {\n return;\n }\n size_t smem_bytes = 0;\n sparse24_sparsify_both_ways_kernel\n <<>>(p, metadata, algo);\n }\n };\n named_algorithms(launchKernel);\n STD_TORCH_CHECK(kernel_launched, \"Unknown algorithm \\\"\", algorithm, \"\\\"\");\n STD_CUDA_KERNEL_LAUNCH_CHECK();\n return std::make_tuple(\n compressed,\n packed_meta_reordered,\n compressed_trans,\n packed_trans_meta_reordered,\n threads_masks);\n}\n\ntemplate \nstd::\n tuple<\n torch::stable::Tensor, // packed\n torch::stable::Tensor, // packed_meta_reordered\n torch::stable::Tensor, // packed_trans\n torch::stable::Tensor, // packed_trans_meta_reordered\n torch::stable::Tensor // threads_masks\n >\n sparse24_sparsify_both_ways(\n const torch::stable::Tensor input,\n std::string algorithm,\n std::string backend) {\n auto runTyped = [&](auto type) {\n using ElementT = decltype(type);\n if (backend == \"cusparselt\") {\n return sparse24_sparsify_both_ways_typed<\n ElementT,\n MetadataCuSparseLtSm80,\n kIsMeta>(input, algorithm);\n } else {\n STD_TORCH_CHECK(\n backend == \"cutlass\",\n \"backend argument only supports `cutlass` or `cusparselt`\");\n return sparse24_sparsify_both_ways_typed<\n ElementT,\n MetadataCutlassSm80,\n kIsMeta>(input, algorithm);\n }\n };\n\n if (input.scalar_type() == torch::headeronly::ScalarType::Half) {\n return runTyped(cutlass::half_t());\n } else {\n STD_TORCH_CHECK(\n input.scalar_type() == torch::headeronly::ScalarType::Half ||\n input.scalar_type() == torch::headeronly::ScalarType::BFloat16);\n return runTyped(cutlass::bfloat16_t());\n }\n}\n} // namespace\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, CUDA, m) {\n m.impl(\n \"sparse24_sparsify_both_ways\",\n TORCH_BOX(sparse24_sparsify_both_ways));\n}\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, Meta, m) {\n m.impl(\n \"sparse24_sparsify_both_ways\",\n TORCH_BOX(sparse24_sparsify_both_ways));\n}\n", "source": "/kaggle/working/xformers/xformers/csrc/sparse24/sparse24_pack.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"pt_stable_utils.h\"\n#include \"sparse24_metadata.h\"\n#include \"sparse24_pack.h\"\n\nusing namespace xformers::sp24;\n\nnamespace {\n\ntemplate \n__global__ void __launch_bounds__(32 /* num_threads */)\n sparse24_apply_kernel(typename KT::Params p) {\n KT::sparse24_apply_kernel(p);\n}\n\n// Apply a 2:4 sparsify pattern computed with\n// `sparse24_sparsify_both_ways_kernel` to another Tensor\ntemplate \nstd::\n tuple<\n torch::stable::Tensor, // packed\n torch::stable::Tensor, // packed_meta_reordered\n torch::stable::Tensor, // packed_trans\n torch::stable::Tensor // packed_trans_meta_reordered\n >\n sparse24_apply_typed(\n torch::stable::Tensor input, // Tensor to sparsify\n torch::stable::Tensor\n threads_masks // Returned by `sparse24_sparsify_both_ways`\n ) {\n using KT = KernelTypes;\n // TODO: Technically we should be able to deal with that\n // by running on the transpose of `input` and swapping\n // `packed` & `packed_t`.\n // This would require to adapt the `threads_masks` a bit tho.\n if (input.stride(1) != 1) {\n input = xf_contiguous(input);\n }\n std::optional device_guard;\n if (!kIsMeta) {\n device_guard.emplace(input.device().index());\n }\n\n STD_TORCH_CHECK(input.dim() == 2);\n STD_TORCH_CHECK(input.stride(1) == 1);\n STD_TORCH_CHECK(input.stride(0) % 8 == 0);\n STD_TORCH_CHECK(input.size(1) % 32 == 0, \"Wrong alignment shape[1]\");\n\n auto rows = input.size(0);\n auto cols = input.size(1);\n\n auto [compressed, packed, packed_meta_reordered] =\n MetadataFormat::create_compressed_representation(\n rows, cols, input, false);\n auto [compressed_trans, packed_trans, packed_trans_meta_reordered] =\n MetadataFormat::create_compressed_representation(\n cols, rows, input, false);\n\n typename KT::Params p;\n p.input_s0 = input.stride(0);\n p.input_dim0 = input.size(0);\n p.input_dim1 = input.size(1);\n\n p.packed_stride = packed.stride(0);\n p.packed_trans_stride = packed_trans.stride(0);\n\n if (!kIsMeta) {\n p.input = (Element const*)input.data_ptr();\n p.packed = (Element*)packed.data_ptr();\n p.packed_trans = (Element*)packed_trans.data_ptr();\n p.threads_masks = (uint64_t*)threads_masks.data_ptr();\n }\n\n STD_TORCH_CHECK(threads_masks.dim() == 3);\n STD_TORCH_CHECK(\n threads_masks.size(0) == p.getBlocksGrid().x * p.getThreadsGrid().x);\n STD_TORCH_CHECK(\n threads_masks.size(1) == p.getBlocksGrid().y * p.getThreadsGrid().y);\n STD_TORCH_CHECK(threads_masks.stride(1) == sizeof(p.threads_masks[0]));\n STD_TORCH_CHECK(threads_masks.size(2) == sizeof(p.threads_masks[0]));\n STD_TORCH_CHECK(threads_masks.stride(2) == 1);\n STD_TORCH_CHECK(\n threads_masks.scalar_type() == torch::headeronly::ScalarType::Byte);\n\n if (!kIsMeta) {\n size_t smem_bytes = 0;\n sparse24_apply_kernel\n <<>>(p);\n STD_CUDA_KERNEL_LAUNCH_CHECK();\n }\n return std::make_tuple(\n compressed,\n packed_meta_reordered,\n compressed_trans,\n packed_trans_meta_reordered);\n}\n\ntemplate \nstd::\n tuple<\n torch::stable::Tensor, // packed\n torch::stable::Tensor, // packed_meta_reordered\n torch::stable::Tensor, // packed_trans\n torch::stable::Tensor // packed_trans_meta_reordered\n >\n sparse24_apply(\n torch::stable::Tensor input, // Tensor to sparsify\n torch::stable::Tensor\n threads_masks, // Returned by `sparse24_sparsify_both_ways`\n std::string backend) {\n auto runTyped = [&](auto type) {\n using ElementT = decltype(type);\n if (backend == \"cusparselt\") {\n return sparse24_apply_typed(\n input, threads_masks);\n } else {\n STD_TORCH_CHECK(\n backend == \"cutlass\",\n \"backend argument only supports `cutlass` or `cusparselt`\");\n return sparse24_apply_typed(\n input, threads_masks);\n }\n };\n\n if (input.scalar_type() == torch::headeronly::ScalarType::Half) {\n return runTyped(cutlass::half_t());\n } else {\n STD_TORCH_CHECK(\n input.scalar_type() == torch::headeronly::ScalarType::Half ||\n input.scalar_type() == torch::headeronly::ScalarType::BFloat16);\n return runTyped(cutlass::bfloat16_t());\n }\n}\n\n} // namespace\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, CUDA, m) {\n m.impl(\"sparse24_apply\", TORCH_BOX(sparse24_apply));\n}\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, Meta, m) {\n m.impl(\"sparse24_apply\", TORCH_BOX(sparse24_apply));\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"pt_stable_utils.h\"\n#include \"sparse24_metadata.h\"\n#include \"sparse24_pack.h\"\n\nusing namespace xformers::sp24;\n\nnamespace {\n\ntemplate \n__global__ void __launch_bounds__(32 /* num_threads */)\n sparse24_apply_kernel(typename KT::Params p) {\n KT::sparse24_apply_kernel(p);\n}\n\n// Apply a 2:4 sparsify pattern computed with\n// `sparse24_sparsify_both_ways_kernel` to another Tensor\ntemplate \nstd::\n tuple<\n torch::stable::Tensor, // packed\n torch::stable::Tensor, // packed_meta_reordered\n torch::stable::Tensor, // packed_trans\n torch::stable::Tensor // packed_trans_meta_reordered\n >\n sparse24_apply_typed(\n torch::stable::Tensor input, // Tensor to sparsify\n torch::stable::Tensor\n threads_masks // Returned by `sparse24_sparsify_both_ways`\n ) {\n using KT = KernelTypes;\n // TODO: Technically we should be able to deal with that\n // by running on the transpose of `input` and swapping\n // `packed` & `packed_t`.\n // This would require to adapt the `threads_masks` a bit tho.\n if (input.stride(1) != 1) {\n input = xf_contiguous(input);\n }\n std::optional device_guard;\n if (!kIsMeta) {\n device_guard.emplace(input.device().index());\n }\n\n STD_TORCH_CHECK(input.dim() == 2);\n STD_TORCH_CHECK(input.stride(1) == 1);\n STD_TORCH_CHECK(input.stride(0) % 8 == 0);\n STD_TORCH_CHECK(input.size(1) % 32 == 0, \"Wrong alignment shape[1]\");\n\n auto rows = input.size(0);\n auto cols = input.size(1);\n\n auto [compressed, packed, packed_meta_reordered] =\n MetadataFormat::create_compressed_representation(\n rows, cols, input, false);\n auto [compressed_trans, packed_trans, packed_trans_meta_reordered] =\n MetadataFormat::create_compressed_representation(\n cols, rows, input, false);\n\n typename KT::Params p;\n p.input_s0 = input.stride(0);\n p.input_dim0 = input.size(0);\n p.input_dim1 = input.size(1);\n\n p.packed_stride = packed.stride(0);\n p.packed_trans_stride = packed_trans.stride(0);\n\n if (!kIsMeta) {\n p.input = (Element const*)input.data_ptr();\n p.packed = (Element*)packed.data_ptr();\n p.packed_trans = (Element*)packed_trans.data_ptr();\n p.threads_masks = (uint64_t*)threads_masks.data_ptr();\n }\n\n STD_TORCH_CHECK(threads_masks.dim() == 3);\n STD_TORCH_CHECK(\n threads_masks.size(0) == p.getBlocksGrid().x * p.getThreadsGrid().x);\n STD_TORCH_CHECK(\n threads_masks.size(1) == p.getBlocksGrid().y * p.getThreadsGrid().y);\n STD_TORCH_CHECK(threads_masks.stride(1) == sizeof(p.threads_masks[0]));\n STD_TORCH_CHECK(threads_masks.size(2) == sizeof(p.threads_masks[0]));\n STD_TORCH_CHECK(threads_masks.stride(2) == 1);\n STD_TORCH_CHECK(\n threads_masks.scalar_type() == torch::headeronly::ScalarType::Byte);\n\n if (!kIsMeta) {\n size_t smem_bytes = 0;\n sparse24_apply_kernel\n <<>>(p);\n STD_CUDA_KERNEL_LAUNCH_CHECK();\n }\n return std::make_tuple(\n compressed,\n packed_meta_reordered,\n compressed_trans,\n packed_trans_meta_reordered);\n}\n\ntemplate \nstd::\n tuple<\n torch::stable::Tensor, // packed\n torch::stable::Tensor, // packed_meta_reordered\n torch::stable::Tensor, // packed_trans\n torch::stable::Tensor // packed_trans_meta_reordered\n >\n sparse24_apply(\n torch::stable::Tensor input, // Tensor to sparsify\n torch::stable::Tensor\n threads_masks, // Returned by `sparse24_sparsify_both_ways`\n std::string backend) {\n auto runTyped = [&](auto type) {\n using ElementT = decltype(type);\n if (backend == \"cusparselt\") {\n return sparse24_apply_typed(\n input, threads_masks);\n } else {\n STD_TORCH_CHECK(\n backend == \"cutlass\",\n \"backend argument only supports `cutlass` or `cusparselt`\");\n return sparse24_apply_typed(\n input, threads_masks);\n }\n };\n\n if (input.scalar_type() == torch::headeronly::ScalarType::Half) {\n return runTyped(cutlass::half_t());\n } else {\n STD_TORCH_CHECK(\n input.scalar_type() == torch::headeronly::ScalarType::Half ||\n input.scalar_type() == torch::headeronly::ScalarType::BFloat16);\n return runTyped(cutlass::bfloat16_t());\n }\n}\n\n} // namespace\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, CUDA, m) {\n m.impl(\"sparse24_apply\", TORCH_BOX(sparse24_apply));\n}\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, Meta, m) {\n m.impl(\"sparse24_apply\", TORCH_BOX(sparse24_apply));\n}\n", "source": "/kaggle/working/xformers/xformers/csrc/sparse24/sparse24_apply.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"pt_stable_utils.h\"\n#include \"static_sort.h\"\n\ntemplate \nstruct Sp24MaskKernel {\n // thread-level tiling\n static constexpr int kNumRows = 4;\n static constexpr int kNumCols = 4;\n static constexpr int kMaxPerRow = 2;\n static constexpr int kMaxPerCol = 2;\n static constexpr int kTakeMaximumCount = kNumRows * kMaxPerRow;\n // block-level tiling\n static constexpr int kRowsPerBlock = 32;\n static constexpr int kColsPerBlock = 16;\n static constexpr int kNumThreadsPerBlock =\n (kRowsPerBlock / kNumRows) * (kRowsPerBlock / kNumRows);\n static_assert(kNumThreadsPerBlock % 32 == 0, \"invalid warp size\");\n\n struct PseudoRandomOrder {\n uint8_t row;\n uint8_t col;\n uint16_t random_score;\n\n CUTLASS_DEVICE bool operator<(PseudoRandomOrder const& other) const {\n return random_score < other.random_score;\n }\n };\n\n struct ElementWithPos {\n Element e;\n uint8_t row;\n uint8_t col;\n\n CUTLASS_DEVICE bool operator<(ElementWithPos const& other) const {\n return e < other.e;\n }\n };\n\n struct Params {\n // Inputs assumed in RowMajor\n Element const* input;\n Element* output;\n int64_t stride0;\n int64_t size0;\n int64_t size1;\n int8_t numRandom;\n\n CUTLASS_HOST_DEVICE int _getNumBlocksRows() const {\n return (size0 + kRowsPerBlock - 1) / kRowsPerBlock;\n }\n CUTLASS_HOST_DEVICE int _getNumBlocksCols() const {\n return (size1 + kColsPerBlock - 1) / kColsPerBlock;\n }\n CUTLASS_HOST_DEVICE int getBlocksGrid() const {\n return _getNumBlocksRows() * _getNumBlocksCols();\n }\n CUTLASS_HOST_DEVICE int getThreadsGrid() const {\n return kNumThreadsPerBlock;\n }\n };\n\n CUTLASS_DEVICE static void run(Params p) {\n // Block-level position\n int block_id = blockIdx.x;\n int thread_row = (block_id % p._getNumBlocksRows()) * kRowsPerBlock;\n int thread_col = (block_id / p._getNumBlocksRows()) * kColsPerBlock;\n\n // Thread-level position\n int thread_id = threadIdx.x;\n constexpr int kTilingRows = kRowsPerBlock / kNumRows;\n thread_row += (thread_id % kTilingRows) * kNumRows;\n thread_col += (thread_id / kTilingRows) * kNumCols;\n\n bool enabled = thread_row < p.size0 && thread_col < p.size1;\n if (!enabled) {\n return;\n }\n\n // We operate on a small 4x4 patch per thread\n using FragmentRow = cutlass::Array;\n cutlass::Array allValues;\n cutlass::Array allOutputs;\n allOutputs.clear();\n uint8_t numOutputsPerRow[kNumRows] = {0};\n uint8_t numOutputsPerCol[kNumCols] = {0};\n uint8_t totalAdded = 0;\n\n CUTLASS_PRAGMA_UNROLL\n for (int row = 0; row < kNumRows; ++row) {\n FragmentRow l = *reinterpret_cast(\n p.input + thread_col + (thread_row + row) * p.stride0);\n CUTLASS_PRAGMA_UNROLL\n for (int col = 0; col < kNumCols; ++col) {\n allValues[col + row * kNumCols].e = l[col];\n allValues[col + row * kNumCols].row = row;\n allValues[col + row * kNumCols].col = col;\n }\n }\n\n // Sort - ascending order\n StaticSort sorter;\n sorter(allValues);\n\n // Take all the values we can starting with the largest\n int i = allValues.size() - 1;\n CUTLASS_PRAGMA_UNROLL\n for (; i >= 0; --i) {\n if (kHasRandom && totalAdded + p.numRandom == kTakeMaximumCount) {\n break;\n }\n\n auto const& v = allValues[i];\n if (numOutputsPerRow[v.row] == kMaxPerRow ||\n numOutputsPerCol[v.col] == kMaxPerCol) {\n continue;\n }\n numOutputsPerRow[v.row] += 1;\n numOutputsPerCol[v.col] += 1;\n totalAdded += 1;\n allOutputs[v.col + v.row * kNumCols] = cutlass::uint1b_t(1);\n }\n\n // Add random elements now\n if (kHasRandom) {\n cutlass::Array randomSorting;\n CUTLASS_PRAGMA_UNROLL\n for (int j = 0; j < randomSorting.size(); ++j) {\n randomSorting[j].row = allValues[j].row;\n randomSorting[j].col = allValues[j].col;\n randomSorting[j].random_score = uint16_t(-1);\n if (j <= i) { // not already considered\n uint16_t seed = allValues[j].e.storage;\n // Assume the lower bits of the significant are more random\n seed = seed & 0x2f;\n randomSorting[j].random_score = (seed ^ (seed << 1)) & 0x7fff;\n }\n }\n sorter(randomSorting);\n CUTLASS_PRAGMA_UNROLL\n for (int j = 0; j < randomSorting.size(); ++j) {\n auto const& v = randomSorting[j];\n if (numOutputsPerRow[v.row] == kMaxPerRow ||\n numOutputsPerCol[v.col] == kMaxPerCol ||\n allOutputs[v.col + v.row * kNumCols].get()) {\n continue;\n }\n numOutputsPerRow[v.row] += 1;\n numOutputsPerCol[v.col] += 1;\n allOutputs[v.col + v.row * kNumCols] = cutlass::uint1b_t(1);\n }\n }\n\n // Write output\n CUTLASS_PRAGMA_UNROLL\n for (int row = 0; row < kNumRows; ++row) {\n FragmentRow fragmentOut;\n fragmentOut.clear();\n CUTLASS_PRAGMA_UNROLL\n for (int col = 0; col < kNumCols; ++col) {\n fragmentOut[col] =\n allOutputs[col + row * kNumCols].get() ? Element(1) : Element(0);\n }\n *reinterpret_cast(\n p.output + thread_col + (thread_row + row) * p.stride0) = fragmentOut;\n }\n }\n};\n\ntemplate \n__global__ void sparse24_largest_mask_2d_cu(typename Kernel::Params p) {\n Kernel::run(p);\n}\n\ntemplate \ntorch::stable::Tensor sparse24_largest_with_random_mask_2d_impl(\n const torch::stable::Tensor input,\n int64_t numRandom) {\n STD_TORCH_CHECK(input.is_cuda(), \"must be a CUDA tensor\");\n STD_TORCH_CHECK(!xf_is_sparse(input), \"must be a dense tensor\");\n STD_TORCH_CHECK(input.is_contiguous(), \"must be contiguous\");\n STD_TORCH_CHECK(input.dim() == 2, \"only works on 2d tensors\");\n STD_TORCH_CHECK(numRandom <= 8, \"There are at most 4x2 elements\")\n\n torch::stable::accelerator::DeviceGuard device_guard(input.device().index());\n cudaStream_t stream = xf_getCurrentCUDAStream();\n torch::stable::Tensor output = torch::stable::empty_like(input);\n\n auto runKernel = [&](auto _) {\n using Element = decltype(_);\n using Kernel = Sp24MaskKernel;\n typename Kernel::Params p;\n p.input = (Element*)input.data_ptr();\n p.output = (Element*)output.data_ptr();\n p.stride0 = input.stride(0);\n p.size0 = input.size(0);\n p.size1 = input.size(1);\n p.numRandom = numRandom;\n STD_TORCH_CHECK((input.size(-1) % Kernel::kNumRows) == 0, \"Wrong shape\");\n STD_TORCH_CHECK((input.size(-2) % Kernel::kNumCols) == 0, \"Wrong shape\");\n\n sparse24_largest_mask_2d_cu\n <<>>(p);\n };\n if (input.scalar_type() == torch::headeronly::ScalarType::Half) {\n runKernel(cutlass::half_t(0));\n } else {\n STD_TORCH_CHECK(\n input.scalar_type() == torch::headeronly::ScalarType::BFloat16,\n \"only f16/bf16 supported\");\n runKernel(cutlass::bfloat16_t(0));\n }\n return output;\n}\n\ntorch::stable::Tensor sparse24_largest_mask_2d(\n const torch::stable::Tensor input) {\n return sparse24_largest_with_random_mask_2d_impl(input, 0);\n}\n\ntorch::stable::Tensor sparse24_largest_with_Krandom_mask_2d(\n const torch::stable::Tensor input,\n int64_t numRandom) {\n return sparse24_largest_with_random_mask_2d_impl(input, numRandom);\n}\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, CUDA, m) {\n m.impl(\"sparse24_largest_mask_2d\", TORCH_BOX(sparse24_largest_mask_2d));\n m.impl(\n \"sparse24_largest_with_Krandom_mask_2d\",\n TORCH_BOX(sparse24_largest_with_Krandom_mask_2d));\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"pt_stable_utils.h\"\n#include \"static_sort.h\"\n\ntemplate \nstruct Sp24MaskKernel {\n // thread-level tiling\n static constexpr int kNumRows = 4;\n static constexpr int kNumCols = 4;\n static constexpr int kMaxPerRow = 2;\n static constexpr int kMaxPerCol = 2;\n static constexpr int kTakeMaximumCount = kNumRows * kMaxPerRow;\n // block-level tiling\n static constexpr int kRowsPerBlock = 32;\n static constexpr int kColsPerBlock = 16;\n static constexpr int kNumThreadsPerBlock =\n (kRowsPerBlock / kNumRows) * (kRowsPerBlock / kNumRows);\n static_assert(kNumThreadsPerBlock % 32 == 0, \"invalid warp size\");\n\n struct PseudoRandomOrder {\n uint8_t row;\n uint8_t col;\n uint16_t random_score;\n\n CUTLASS_DEVICE bool operator<(PseudoRandomOrder const& other) const {\n return random_score < other.random_score;\n }\n };\n\n struct ElementWithPos {\n Element e;\n uint8_t row;\n uint8_t col;\n\n CUTLASS_DEVICE bool operator<(ElementWithPos const& other) const {\n return e < other.e;\n }\n };\n\n struct Params {\n // Inputs assumed in RowMajor\n Element const* input;\n Element* output;\n int64_t stride0;\n int64_t size0;\n int64_t size1;\n int8_t numRandom;\n\n CUTLASS_HOST_DEVICE int _getNumBlocksRows() const {\n return (size0 + kRowsPerBlock - 1) / kRowsPerBlock;\n }\n CUTLASS_HOST_DEVICE int _getNumBlocksCols() const {\n return (size1 + kColsPerBlock - 1) / kColsPerBlock;\n }\n CUTLASS_HOST_DEVICE int getBlocksGrid() const {\n return _getNumBlocksRows() * _getNumBlocksCols();\n }\n CUTLASS_HOST_DEVICE int getThreadsGrid() const {\n return kNumThreadsPerBlock;\n }\n };\n\n CUTLASS_DEVICE static void run(Params p) {\n // Block-level position\n int block_id = blockIdx.x;\n int thread_row = (block_id % p._getNumBlocksRows()) * kRowsPerBlock;\n int thread_col = (block_id / p._getNumBlocksRows()) * kColsPerBlock;\n\n // Thread-level position\n int thread_id = threadIdx.x;\n constexpr int kTilingRows = kRowsPerBlock / kNumRows;\n thread_row += (thread_id % kTilingRows) * kNumRows;\n thread_col += (thread_id / kTilingRows) * kNumCols;\n\n bool enabled = thread_row < p.size0 && thread_col < p.size1;\n if (!enabled) {\n return;\n }\n\n // We operate on a small 4x4 patch per thread\n using FragmentRow = cutlass::Array;\n cutlass::Array allValues;\n cutlass::Array allOutputs;\n allOutputs.clear();\n uint8_t numOutputsPerRow[kNumRows] = {0};\n uint8_t numOutputsPerCol[kNumCols] = {0};\n uint8_t totalAdded = 0;\n\n CUTLASS_PRAGMA_UNROLL\n for (int row = 0; row < kNumRows; ++row) {\n FragmentRow l = *reinterpret_cast(\n p.input + thread_col + (thread_row + row) * p.stride0);\n CUTLASS_PRAGMA_UNROLL\n for (int col = 0; col < kNumCols; ++col) {\n allValues[col + row * kNumCols].e = l[col];\n allValues[col + row * kNumCols].row = row;\n allValues[col + row * kNumCols].col = col;\n }\n }\n\n // Sort - ascending order\n StaticSort sorter;\n sorter(allValues);\n\n // Take all the values we can starting with the largest\n int i = allValues.size() - 1;\n CUTLASS_PRAGMA_UNROLL\n for (; i >= 0; --i) {\n if (kHasRandom && totalAdded + p.numRandom == kTakeMaximumCount) {\n break;\n }\n\n auto const& v = allValues[i];\n if (numOutputsPerRow[v.row] == kMaxPerRow ||\n numOutputsPerCol[v.col] == kMaxPerCol) {\n continue;\n }\n numOutputsPerRow[v.row] += 1;\n numOutputsPerCol[v.col] += 1;\n totalAdded += 1;\n allOutputs[v.col + v.row * kNumCols] = cutlass::uint1b_t(1);\n }\n\n // Add random elements now\n if (kHasRandom) {\n cutlass::Array randomSorting;\n CUTLASS_PRAGMA_UNROLL\n for (int j = 0; j < randomSorting.size(); ++j) {\n randomSorting[j].row = allValues[j].row;\n randomSorting[j].col = allValues[j].col;\n randomSorting[j].random_score = uint16_t(-1);\n if (j <= i) { // not already considered\n uint16_t seed = allValues[j].e.storage;\n // Assume the lower bits of the significant are more random\n seed = seed & 0x2f;\n randomSorting[j].random_score = (seed ^ (seed << 1)) & 0x7fff;\n }\n }\n sorter(randomSorting);\n CUTLASS_PRAGMA_UNROLL\n for (int j = 0; j < randomSorting.size(); ++j) {\n auto const& v = randomSorting[j];\n if (numOutputsPerRow[v.row] == kMaxPerRow ||\n numOutputsPerCol[v.col] == kMaxPerCol ||\n allOutputs[v.col + v.row * kNumCols].get()) {\n continue;\n }\n numOutputsPerRow[v.row] += 1;\n numOutputsPerCol[v.col] += 1;\n allOutputs[v.col + v.row * kNumCols] = cutlass::uint1b_t(1);\n }\n }\n\n // Write output\n CUTLASS_PRAGMA_UNROLL\n for (int row = 0; row < kNumRows; ++row) {\n FragmentRow fragmentOut;\n fragmentOut.clear();\n CUTLASS_PRAGMA_UNROLL\n for (int col = 0; col < kNumCols; ++col) {\n fragmentOut[col] =\n allOutputs[col + row * kNumCols].get() ? Element(1) : Element(0);\n }\n *reinterpret_cast(\n p.output + thread_col + (thread_row + row) * p.stride0) = fragmentOut;\n }\n }\n};\n\ntemplate \n__global__ void sparse24_largest_mask_2d_cu(typename Kernel::Params p) {\n Kernel::run(p);\n}\n\ntemplate \ntorch::stable::Tensor sparse24_largest_with_random_mask_2d_impl(\n const torch::stable::Tensor input,\n int64_t numRandom) {\n STD_TORCH_CHECK(input.is_cuda(), \"must be a CUDA tensor\");\n STD_TORCH_CHECK(!xf_is_sparse(input), \"must be a dense tensor\");\n STD_TORCH_CHECK(input.is_contiguous(), \"must be contiguous\");\n STD_TORCH_CHECK(input.dim() == 2, \"only works on 2d tensors\");\n STD_TORCH_CHECK(numRandom <= 8, \"There are at most 4x2 elements\")\n\n torch::stable::accelerator::DeviceGuard device_guard(input.device().index());\n hipStream_t stream = xf_getCurrentCUDAStream();\n torch::stable::Tensor output = torch::stable::empty_like(input);\n\n auto runKernel = [&](auto _) {\n using Element = decltype(_);\n using Kernel = Sp24MaskKernel;\n typename Kernel::Params p;\n p.input = (Element*)input.data_ptr();\n p.output = (Element*)output.data_ptr();\n p.stride0 = input.stride(0);\n p.size0 = input.size(0);\n p.size1 = input.size(1);\n p.numRandom = numRandom;\n STD_TORCH_CHECK((input.size(-1) % Kernel::kNumRows) == 0, \"Wrong shape\");\n STD_TORCH_CHECK((input.size(-2) % Kernel::kNumCols) == 0, \"Wrong shape\");\n\n sparse24_largest_mask_2d_cu\n <<>>(p);\n };\n if (input.scalar_type() == torch::headeronly::ScalarType::Half) {\n runKernel(cutlass::half_t(0));\n } else {\n STD_TORCH_CHECK(\n input.scalar_type() == torch::headeronly::ScalarType::BFloat16,\n \"only f16/bf16 supported\");\n runKernel(cutlass::bfloat16_t(0));\n }\n return output;\n}\n\ntorch::stable::Tensor sparse24_largest_mask_2d(\n const torch::stable::Tensor input) {\n return sparse24_largest_with_random_mask_2d_impl(input, 0);\n}\n\ntorch::stable::Tensor sparse24_largest_with_Krandom_mask_2d(\n const torch::stable::Tensor input,\n int64_t numRandom) {\n return sparse24_largest_with_random_mask_2d_impl(input, numRandom);\n}\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, CUDA, m) {\n m.impl(\"sparse24_largest_mask_2d\", TORCH_BOX(sparse24_largest_mask_2d));\n m.impl(\n \"sparse24_largest_with_Krandom_mask_2d\",\n TORCH_BOX(sparse24_largest_with_Krandom_mask_2d));\n}\n", "source": "/kaggle/working/xformers/xformers/csrc/sparse24/sparse24_largest_mask_2d.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"pt_stable_utils.h\"\n#include \"sparse24_pack.h\"\n\nusing namespace xformers::sp24;\n\nnamespace {\n__global__ void meta_shuffle_test_kernel(\n torch::headeronly::HeaderOnlyGenericPackedTensorAccessor\n local_meta,\n torch::headeronly::HeaderOnlyGenericPackedTensorAccessor\n final_meta,\n bool transpose) {\n uint32_t meta_ab = 0;\n uint32_t meta_cd = 0;\n for (int i = 0; i < 4; ++i) {\n meta_ab |= uint8b_t(uint32_t(local_meta[threadIdx.x][threadIdx.y][i]))\n << (8 * i);\n meta_cd |= uint8b_t(uint32_t(local_meta[threadIdx.x][threadIdx.y][4 + i]))\n << (8 * i);\n }\n final_meta[threadIdx.x][threadIdx.y][0] =\n warp_shuffle_meta(meta_ab, transpose);\n final_meta[threadIdx.x][threadIdx.y][1] =\n warp_shuffle_meta(meta_cd, transpose);\n}\n\ntorch::stable::Tensor _sparse24_meta_shuffle_test(\n torch::stable::Tensor local_meta,\n bool transpose) {\n auto threads_grid = KernelTypes::Params::getThreadsGrid();\n\n STD_TORCH_CHECK(\n local_meta.scalar_type() == torch::headeronly::ScalarType::Long);\n STD_TORCH_CHECK(local_meta.dim() == 3);\n STD_TORCH_CHECK(local_meta.size(0) == threads_grid.x);\n STD_TORCH_CHECK(local_meta.size(1) == threads_grid.y);\n STD_TORCH_CHECK(local_meta.size(2) == kThreadY);\n torch::stable::accelerator::DeviceGuard device_guard(\n local_meta.device().index());\n cudaStream_t stream = xf_getCurrentCUDAStream();\n\n torch::stable::Tensor final_meta = xf_zeros(\n {threads_grid.x, threads_grid.y, 2},\n local_meta.scalar_type(),\n local_meta.device());\n size_t smem_bytes = 0;\n meta_shuffle_test_kernel<<<1, threads_grid, smem_bytes, stream>>>(\n xf_packed_accessor(local_meta),\n xf_packed_accessor(final_meta),\n transpose);\n return final_meta;\n}\n} // namespace\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, CUDA, m) {\n m.impl(\"_sparse24_meta_shuffle_test\", TORCH_BOX(_sparse24_meta_shuffle_test));\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"pt_stable_utils.h\"\n#include \"sparse24_pack.h\"\n\nusing namespace xformers::sp24;\n\nnamespace {\n__global__ void meta_shuffle_test_kernel(\n torch::headeronly::HeaderOnlyGenericPackedTensorAccessor\n local_meta,\n torch::headeronly::HeaderOnlyGenericPackedTensorAccessor\n final_meta,\n bool transpose) {\n uint32_t meta_ab = 0;\n uint32_t meta_cd = 0;\n for (int i = 0; i < 4; ++i) {\n meta_ab |= uint8b_t(uint32_t(local_meta[threadIdx.x][threadIdx.y][i]))\n << (8 * i);\n meta_cd |= uint8b_t(uint32_t(local_meta[threadIdx.x][threadIdx.y][4 + i]))\n << (8 * i);\n }\n final_meta[threadIdx.x][threadIdx.y][0] =\n warp_shuffle_meta(meta_ab, transpose);\n final_meta[threadIdx.x][threadIdx.y][1] =\n warp_shuffle_meta(meta_cd, transpose);\n}\n\ntorch::stable::Tensor _sparse24_meta_shuffle_test(\n torch::stable::Tensor local_meta,\n bool transpose) {\n auto threads_grid = KernelTypes::Params::getThreadsGrid();\n\n STD_TORCH_CHECK(\n local_meta.scalar_type() == torch::headeronly::ScalarType::Long);\n STD_TORCH_CHECK(local_meta.dim() == 3);\n STD_TORCH_CHECK(local_meta.size(0) == threads_grid.x);\n STD_TORCH_CHECK(local_meta.size(1) == threads_grid.y);\n STD_TORCH_CHECK(local_meta.size(2) == kThreadY);\n torch::stable::accelerator::DeviceGuard device_guard(\n local_meta.device().index());\n hipStream_t stream = xf_getCurrentCUDAStream();\n\n torch::stable::Tensor final_meta = xf_zeros(\n {threads_grid.x, threads_grid.y, 2},\n local_meta.scalar_type(),\n local_meta.device());\n size_t smem_bytes = 0;\n meta_shuffle_test_kernel<<<1, threads_grid, smem_bytes, stream>>>(\n xf_packed_accessor(local_meta),\n xf_packed_accessor(final_meta),\n transpose);\n return final_meta;\n}\n} // namespace\n\nSTABLE_TORCH_LIBRARY_IMPL(xformers, CUDA, m) {\n m.impl(\"_sparse24_meta_shuffle_test\", TORCH_BOX(_sparse24_meta_shuffle_test));\n}\n", "source": "/kaggle/working/xformers/xformers/csrc/sparse24/sparse24_pack_test.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n/* Includes, cuda */\n#include \n#include \n\n#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11000\n// includes cublaslt\n#include \n#endif\n// constants for fused bias+relu kernel\n#define BIAS_RELU_FW_NTHREADS 128 // forward number of thread per block\n#define BIAS_RELU_BW_NTHREADS_X 32 // backward number of thread in feature dim\n#define BIAS_RELU_BW_NTHREADS_Y 16 // backward number of thread in batch dim\n#define BIAS_RELU_RED_PER_THREAD 16 // backward minimal reduction length per thread\n\n// move to a header later on\n#define ILP 4\ntemplate \n__host__ __device__ __forceinline__ bool is_aligned(T* p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\ntemplate \n__device__ __forceinline__ void load_store(T* dst, volatile T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\ntemplate \n__device__ __forceinline__ void load_store(volatile T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\n\n// Keep ReLU in float only. When using half, cast to float before calling.\n__device__ __inline__ float relu(float a) {\n float retf = max(a, 0.f);\n return (retf);\n}\n\n// Keep Sigmoid in float only. When using half, cast to float before calling.\n__device__ __inline__ float sigmoid(float a) {\n float retf = 1.f / (1.f + expf(-a));\n return (retf);\n}\n\n// FP64 Wrapper around cublas GEMMEx\ncublasStatus_t mlp_gemm(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k,\n float* alpha, const double* A, int lda, const double* B, int ldb, const float* beta, double* C,\n int ldc) {\n return cublasGemmEx(handle, transa, transb, m, n, k, alpha, A, CUDA_R_64F, lda, B, CUDA_R_64F, ldb, beta, C,\n CUDA_R_64F, ldc, CUDA_R_64F, CUBLAS_GEMM_DEFAULT);\n}\n\n// FP32 Wrapper around cublas GEMMEx\ncublasStatus_t mlp_gemm(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k,\n float* alpha, const float* A, int lda, const float* B, int ldb, const float* beta, float* C,\n int ldc) {\n return cublasGemmEx(handle, transa, transb, m, n, k, alpha, A, CUDA_R_32F, lda, B, CUDA_R_32F, ldb, beta, C,\n CUDA_R_32F, ldc, CUDA_R_32F, CUBLAS_GEMM_DEFAULT);\n}\n\n// FP16 Tensor core wrapper around cublas GEMMEx\ncublasStatus_t mlp_gemm(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k,\n float* alpha, const at::Half* A, int lda, const at::Half* B, int ldb, float* beta, at::Half* C,\n int ldc) {\n return cublasGemmEx(handle, transa, transb, m, n, k, alpha, A, CUDA_R_16F, lda, B, CUDA_R_16F, ldb, beta, C,\n CUDA_R_16F, ldc, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP);\n}\n#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11000\nint mlp_gemm_lt(cublasLtHandle_t ltHandle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k,\n float* alpha, /* host pointer */\n const at::Half* A, int lda, const at::Half* B, int ldb, float* beta, /* host pointer */\n at::Half* C, int ldc, void* workspace, size_t workspaceSize, cudaStream_t stream, bool use_bias,\n bool use_relu, const void* bias) {\n cublasStatus_t status = CUBLAS_STATUS_SUCCESS;\n\n cublasLtMatmulDescOpaque_t operationDesc = {};\n cublasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};\n cublasLtMatmulPreferenceOpaque_t preference = {};\n\n int returnedResults = 0;\n cublasLtMatmulHeuristicResult_t heuristicResult = {};\n cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DEFAULT;\n\n // Create operation descriptor; see cublasLtMatmulDescAttributes_t\n // for details about defaults; here we just set the transforms for\n // A and B.\n status = cublasLtMatmulDescInit(&operationDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F);\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa));\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n if (use_bias) {\n status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias));\n if (status != CUBLAS_STATUS_SUCCESS) {\n goto CLEANUP;\n }\n if (use_relu) {\n epilogue = CUBLASLT_EPILOGUE_RELU_BIAS;\n } else {\n epilogue = CUBLASLT_EPILOGUE_BIAS;\n }\n } else {\n if (use_relu) {\n epilogue = CUBLASLT_EPILOGUE_RELU;\n }\n }\n\n status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue));\n if (status != CUBLAS_STATUS_SUCCESS) {\n goto CLEANUP;\n }\n\n // Create matrix descriptors. Not setting any extra attributes.\n status =\n cublasLtMatrixLayoutInit(&Adesc, CUDA_R_16F, transa == CUBLAS_OP_N ? m : k, transa == CUBLAS_OP_N ? k : m, lda);\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n status =\n cublasLtMatrixLayoutInit(&Bdesc, CUDA_R_16F, transb == CUBLAS_OP_N ? k : n, transb == CUBLAS_OP_N ? n : k, ldb);\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = cublasLtMatrixLayoutInit(&Cdesc, CUDA_R_16F, m, n, ldc);\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n // Create preference handle; In general, extra attributes can be\n // used here to disable tensor ops or to make sure algo selected\n // will work with badly aligned A, B, C. However, for simplicity\n // here we assume A,B,C are always well aligned (e.g., directly\n // come from cudaMalloc)\n status = cublasLtMatmulPreferenceInit(&preference);\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = cublasLtMatmulPreferenceSetAttribute(&preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize,\n sizeof(workspaceSize));\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n // We just need the best available heuristic to try and run matmul.\n // There is no guarantee that this will work. For example, if A is\n // badly aligned, you can request more (e.g. 32) algos and try to\n // run them one by one until something works.\n status = cublasLtMatmulAlgoGetHeuristic(ltHandle, &operationDesc, &Adesc, &Bdesc, &Cdesc, &Cdesc, &preference, 1,\n &heuristicResult, &returnedResults);\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n if (returnedResults == 0) {\n status = CUBLAS_STATUS_NOT_SUPPORTED;\n goto CLEANUP;\n }\n status = cublasLtMatmul(ltHandle, &operationDesc, alpha, A, &Adesc, B, &Bdesc, beta, C, &Cdesc, C, &Cdesc,\n &heuristicResult.algo, workspace, workspaceSize, stream);\n\nCLEANUP:\n // Descriptors are no longer needed as all GPU work was already\n // enqueued.\n return status == CUBLAS_STATUS_SUCCESS ? 0 : 1;\n}\n\nint mlp_gemm_lt(cublasLtHandle_t ltHandle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k,\n float* alpha, /* host pointer */\n const double* A, int lda, const double* B, int ldb, float* beta, /* host pointer */\n double* C, int ldc, void* workspace, size_t workspaceSize, cudaStream_t stream, bool use_bias,\n bool use_relu, const void* bias) {\n return 1;\n}\n\nint mlp_gemm_lt(cublasLtHandle_t ltHandle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k,\n float* alpha, /* host pointer */\n const float* A, int lda, const float* B, int ldb, float* beta, /* host pointer */\n float* C, int ldc, void* workspace, size_t workspaceSize, cudaStream_t stream, bool use_bias,\n bool use_relu, const void* bias) {\n cublasStatus_t status = CUBLAS_STATUS_SUCCESS;\n\n cublasLtMatmulDescOpaque_t operationDesc = {};\n cublasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};\n cublasLtMatmulPreferenceOpaque_t preference = {};\n\n int returnedResults = 0;\n cublasLtMatmulHeuristicResult_t heuristicResult = {};\n cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DEFAULT;\n\n // Create operation descriptor; see cublasLtMatmulDescAttributes_t\n // for details about defaults; here we just set the transforms for\n // A and B.\n status = cublasLtMatmulDescInit(&operationDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F);\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa));\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n if (use_bias) {\n status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias));\n if (status != CUBLAS_STATUS_SUCCESS) {\n goto CLEANUP;\n }\n if (use_relu) {\n epilogue = CUBLASLT_EPILOGUE_RELU_BIAS;\n } else {\n epilogue = CUBLASLT_EPILOGUE_BIAS;\n }\n } else {\n if (use_relu) {\n epilogue = CUBLASLT_EPILOGUE_RELU;\n }\n }\n\n status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue));\n if (status != CUBLAS_STATUS_SUCCESS) {\n goto CLEANUP;\n }\n\n // Create matrix descriptors. Not setting any extra attributes.\n status =\n cublasLtMatrixLayoutInit(&Adesc, CUDA_R_32F, transa == CUBLAS_OP_N ? m : k, transa == CUBLAS_OP_N ? k : m, lda);\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n status =\n cublasLtMatrixLayoutInit(&Bdesc, CUDA_R_32F, transb == CUBLAS_OP_N ? k : n, transb == CUBLAS_OP_N ? n : k, ldb);\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = cublasLtMatrixLayoutInit(&Cdesc, CUDA_R_32F, m, n, ldc);\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n // Create preference handle; In general, extra attributes can be\n // used here to disable tensor ops or to make sure algo selected\n // will work with badly aligned A, B, C. However, for simplicity\n // here we assume A,B,C are always well aligned (e.g., directly\n // come from cudaMalloc)\n status = cublasLtMatmulPreferenceInit(&preference);\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = cublasLtMatmulPreferenceSetAttribute(&preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize,\n sizeof(workspaceSize));\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n // We just need the best available heuristic to try and run matmul.\n // There is no guarantee that this will work. For example, if A is\n // badly aligned, you can request more (e.g. 32) algos and try to\n // run them one by one until something works.\n status = cublasLtMatmulAlgoGetHeuristic(ltHandle, &operationDesc, &Adesc, &Bdesc, &Cdesc, &Cdesc, &preference, 1,\n &heuristicResult, &returnedResults);\n if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n if (returnedResults == 0) {\n status = CUBLAS_STATUS_NOT_SUPPORTED;\n goto CLEANUP;\n }\n\n status = cublasLtMatmul(ltHandle, &operationDesc, alpha, A, &Adesc, B, &Bdesc, beta, C, &Cdesc, C, &Cdesc,\n &heuristicResult.algo, workspace, workspaceSize, stream);\n\nCLEANUP:\n // Descriptors are no longer needed as all GPU work was already\n // enqueued.\n return status == CUBLAS_STATUS_SUCCESS ? 0 : 1;\n}\n#endif\n\n// Bias ADD. Assume input X is [features x batch size], column major.\n// Bias is one 'features' long vector, with implicit broadcast.\ntemplate \n__global__ void biasAdd_fprop(T* X, T* b, uint batch_size, uint features) {\n T r_x[ILP];\n T r_b[ILP];\n if (is_aligned(X) && is_aligned(b) && features % ILP == 0) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid * ILP < features * batch_size; tid += blockDim.x * gridDim.x) {\n int row = tid % (features / ILP);\n load_store(r_x, X, 0, tid);\n load_store(r_b, b, 0, row);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float bias_sum = static_cast(r_x[ii]) + static_cast(r_b[ii]);\n r_x[ii] = bias_sum;\n }\n load_store(X, r_x, tid, 0);\n }\n } else {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n int row = tid % features;\n r_x[ii] = X[idx];\n r_b[ii] = b[row];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float bias_sum = static_cast(r_x[ii]) + static_cast(r_b[ii]);\n r_x[ii] = bias_sum;\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n X[idx] = r_x[ii];\n }\n }\n }\n }\n}\n\n// Bias ADD + ReLU. Assume input X is [features x batch size], column major.\n// Activation support fuesed ReLU. Safe to call in-place.\ntemplate \n__global__ void biasAddRelu_fprop(T* X, T* b, uint batch_size, uint features) {\n T r_x[ILP];\n T r_b[ILP];\n if (is_aligned(X) && is_aligned(b) && features % ILP == 0) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid * ILP < features * batch_size; tid += blockDim.x * gridDim.x) {\n int row = tid % (features / ILP);\n load_store(r_x, X, 0, tid);\n load_store(r_b, b, 0, row);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float bias_sum = static_cast(r_x[ii]) + static_cast(r_b[ii]);\n r_x[ii] = relu(bias_sum);\n }\n load_store(X, r_x, tid, 0);\n }\n } else {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n int row = tid % features;\n r_x[ii] = X[idx];\n r_b[ii] = b[row];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float bias_sum = static_cast(r_x[ii]) + static_cast(r_b[ii]);\n r_x[ii] = relu(bias_sum);\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n X[idx] = r_x[ii];\n }\n }\n }\n }\n}\n\n// ReLU. Assume input X is [features x batch size], column major.\n// Safe to call in-place.\ntemplate \n__global__ void Relu_fprop(T* X, uint batch_size, uint features) {\n T r_x[ILP];\n if (is_aligned(X) && features % ILP == 0) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid * ILP < features * batch_size; tid += blockDim.x * gridDim.x) {\n load_store(r_x, X, 0, tid);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_x[ii] = relu(static_cast(r_x[ii]));\n }\n load_store(X, r_x, tid, 0);\n }\n } else {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n r_x[ii] = X[idx];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_x[ii] = relu(static_cast(r_x[ii]));\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n X[idx] = r_x[ii];\n }\n }\n }\n }\n}\n\n// Sigmoid. Assume input X is [features x batch size], column major.\n// Safe to call in-place.\ntemplate \n__global__ void Sigmoid_fprop(T* X, uint batch_size, uint features) {\n T r_x[ILP];\n if (is_aligned(X) && features % ILP == 0) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid * ILP < features * batch_size; tid += blockDim.x * gridDim.x) {\n load_store(r_x, X, 0, tid);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_x[ii] = sigmoid(static_cast(r_x[ii]));\n }\n load_store(X, r_x, tid, 0);\n }\n } else {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n r_x[ii] = X[idx];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_x[ii] = sigmoid(static_cast(r_x[ii]));\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n X[idx] = r_x[ii];\n }\n }\n }\n }\n}\n\n// ReLU. Assume input X is [features x batch size], column major.\n// Safe to call in-place.\ntemplate \n__global__ void Relu_bprop(T* dY, T* Y, uint batch_size, uint features, T* dX) {\n T r_dy[ILP];\n T r_y[ILP];\n if (is_aligned(dY) && is_aligned(Y) && is_aligned(dX) && features % ILP == 0) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid * ILP < features * batch_size; tid += blockDim.x * gridDim.x) {\n load_store(r_dy, dY, 0, tid);\n load_store(r_y, Y, 0, tid);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n if ((float)r_y[ii] <= 0.f) r_dy[ii] = 0;\n }\n load_store(dX, r_dy, tid, 0);\n }\n } else {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n r_dy[ii] = dY[idx];\n r_y[ii] = Y[idx];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n if ((float)r_y[ii] <= 0.f) r_dy[ii] = 0;\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n dX[idx] = r_dy[ii];\n }\n }\n }\n }\n}\n\n// Sigmoid. Assume input X is [features x batch size], column major.\n// Safe to call in-place.\ntemplate \n__global__ void Sigmoid_bprop(T* dY, T* Y, uint batch_size, uint features, T* dX) {\n T r_dy[ILP];\n T r_y[ILP];\n if (is_aligned(dY) && is_aligned(Y) && is_aligned(dX) && features % ILP == 0) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid * ILP < features * batch_size; tid += blockDim.x * gridDim.x) {\n load_store(r_dy, dY, 0, tid);\n load_store(r_y, Y, 0, tid);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float grad_out = r_dy[ii];\n float out = r_y[ii];\n float grad_i = out * (1.f - out) * grad_out;\n r_dy[ii] = grad_i;\n }\n load_store(dX, r_dy, tid, 0);\n }\n } else {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n r_dy[ii] = dY[idx];\n r_y[ii] = Y[idx];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float grad_out = r_dy[ii];\n float out = r_y[ii];\n float grad_i = out * (1.f - out) * grad_out;\n r_dy[ii] = grad_i;\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n dX[idx] = r_dy[ii];\n }\n }\n }\n }\n}\n\n// Compute grid size for pointwise backward kernel.\n// block_x/y is total elment being handled per block, not number of threads\nvoid get_biasAddRelu_bprop_grid_size(int yfeat, int batch_size, int block_x, int block_y, int* grid_x, int* grid_y) {\n *grid_x = (yfeat + block_x - 1) / block_x;\n // Get number of SMs for efficient reduction.\n int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;\n // can switch to occupancy calculation. use 4 below now for sm_70\n int max_blocks_y = (num_SMs * 4 + (*grid_x) - 1) / (*grid_x);\n // block_y should be from minimal work per thread\n int nRedSplits = (batch_size + block_y - 1) / block_y;\n // increase number of elem per thread redcution to not launch more than enough\n // kernel adjust work, so here we just launch max block\n *grid_y = std::min(nRedSplits, max_blocks_y);\n return;\n}\n\n// Addition done deterministically via a 2-pass approach. Each CTA writes out partial\n// sum, and the last CTA in grid Y dimension accumulates partials serially and writes to result.\ntemplate \n__global__ void biasAdd_bprop(T* dY, int features, int batch_size, volatile float* intermediate, int* semaphores,\n T* db) {\n // The feature that this thread is responsible for\n int f = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Compute the span this thread is responsible for\n // For this block\n int b_chunkSize = (batch_size + gridDim.y - 1) / gridDim.y;\n int b_nStart = blockIdx.y * b_chunkSize;\n int b_nSpan = min(batch_size, b_nStart + b_chunkSize) - b_nStart;\n // For this thread\n int chunkSize = (b_chunkSize + blockDim.y - 1) / blockDim.y;\n int nStart = threadIdx.y * chunkSize + b_nStart;\n int nSpan = min(b_nStart + b_nSpan, nStart + chunkSize) - nStart;\n\n volatile float* out = intermediate + blockIdx.y * features;\n\n // Flag to trigger last reduction.\n __shared__ bool isLastBlock;\n // we know block size for now\n __shared__ float smem[BIAS_RELU_BW_NTHREADS_X * BIAS_RELU_BW_NTHREADS_Y];\n\n // Accumulate db in FP32 always\n float db_local = 0;\n if (f < features) {\n int nidx = 0;\n // Handle non-multiple of UNROLL_FACTOR residue\n for (; nidx < nSpan % UNROLL_FACTOR; nidx++) {\n int64_t row, col, flat_idx;\n row = f;\n col = nStart + nidx;\n flat_idx = col * features + row;\n db_local += (float)dY[flat_idx];\n }\n\n // Handle meat of work\n for (; (nidx + UNROLL_FACTOR - 1) < nSpan; nidx += UNROLL_FACTOR) {\n int64_t row, col, flat_idx;\n row = f;\n col = nStart + nidx;\n flat_idx = col * features + row;\n#pragma unroll 4\n for (int u = 0; u < UNROLL_FACTOR; u++) {\n db_local += (float)dY[flat_idx];\n flat_idx += features;\n }\n }\n\n // naive block reduction on y-dim\n int linear_idx = threadIdx.y * blockDim.x + threadIdx.x;\n smem[linear_idx] = db_local;\n }\n __syncthreads();\n if (f < features) {\n if (threadIdx.y == 0) {\n for (int yidx = 1; yidx < blockDim.y; yidx++) {\n db_local += smem[yidx * blockDim.x + threadIdx.x];\n }\n\n // block result is in db_local now for all threadIdx.y == 0\n // Write out partial result\n out[f] = db_local;\n }\n }\n __threadfence();\n __syncthreads();\n\n // Increment semaphore and check if this is the last CTA in the grid_y dimension.\n // Only thread (0,0) calls this\n if (threadIdx.x == 0 && threadIdx.y == 0 && f < features) {\n unsigned int sum_idx;\n sum_idx = atomicAdd(&(semaphores[blockIdx.x]), 1);\n isLastBlock = (sum_idx == (gridDim.y - 1));\n }\n __syncthreads();\n\n db_local = 0;\n // No block reduction for now, only thread (*,0) do grid reduction\n if (isLastBlock && f < features) {\n if (threadIdx.y == 0) {\n for (int n = 0; n < gridDim.y; n++) {\n int row, col;\n row = f;\n col = n;\n db_local += (float)(intermediate[col * features + row]);\n }\n db[f] = (T)db_local;\n }\n }\n}\n\n// Addition done deterministically via a 2-pass approach. Each CTA writes out partial\n// sum, and the last CTA in grid Y dimension accumulates partials serially and writes to result.\ntemplate \n__global__ void biasAddRelu_bprop(T* Y, T* dY, int features, int batch_size, T* dX, volatile float* intermediate,\n int* semaphores, T* db) {\n // The feature that this thread is responsible for\n int f = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Compute the span this thread is responsible for\n // For this block\n int b_chunkSize = (batch_size + gridDim.y - 1) / gridDim.y;\n int b_nStart = blockIdx.y * b_chunkSize;\n int b_nSpan = min(batch_size, b_nStart + b_chunkSize) - b_nStart;\n // For this thread\n int chunkSize = (b_chunkSize + blockDim.y - 1) / blockDim.y;\n int nStart = threadIdx.y * chunkSize + b_nStart;\n int nSpan = min(b_nStart + b_nSpan, nStart + chunkSize) - nStart;\n\n volatile float* out = intermediate + blockIdx.y * features;\n\n // Flag to trigger last reduction.\n __shared__ bool isLastBlock;\n // we know block size for now\n __shared__ float smem[BIAS_RELU_BW_NTHREADS_X * BIAS_RELU_BW_NTHREADS_Y];\n\n // Accumulate db in FP32 always\n float db_local = 0;\n if (f < features) {\n int nidx = 0;\n // Handle non-multiple of UNROLL_FACTOR residue\n for (; nidx < nSpan % UNROLL_FACTOR; nidx++) {\n int row, col, flat_idx;\n row = f;\n col = nStart + nidx;\n flat_idx = col * features + row;\n T y_val = Y[flat_idx];\n T dy_val = dY[flat_idx];\n T dx_val;\n if ((float)y_val > 0.f)\n dx_val = dy_val;\n else\n dx_val = 0;\n dX[flat_idx] = dx_val;\n db_local += (float)dx_val;\n }\n\n // Handle meat of work\n for (; (nidx + UNROLL_FACTOR - 1) < nSpan; nidx += UNROLL_FACTOR) {\n int row, col, flat_idx;\n row = f;\n col = nStart + nidx;\n flat_idx = col * features + row;\n#pragma unroll 4\n for (int u = 0; u < UNROLL_FACTOR; u++) {\n T y_val = Y[flat_idx];\n T dy_val = dY[flat_idx];\n T dx_val;\n if ((float)y_val > 0.f)\n dx_val = dy_val;\n else\n dx_val = 0;\n dX[flat_idx] = dx_val;\n db_local += (float)dx_val;\n flat_idx += features;\n }\n }\n\n // naive block reduction on y-dim\n int linear_idx = threadIdx.y * blockDim.x + threadIdx.x;\n smem[linear_idx] = db_local;\n }\n __syncthreads();\n if (f < features) {\n if (threadIdx.y == 0) {\n for (int yidx = 1; yidx < blockDim.y; yidx++) {\n db_local += smem[yidx * blockDim.x + threadIdx.x];\n }\n\n // block result is in db_local now for all threadIdx.y == 0\n // Write out partial result\n out[f] = db_local;\n }\n }\n __threadfence();\n __syncthreads();\n\n // Increment semaphore and check if this is the last CTA in the grid_y dimension.\n // Only thread (0,0) calls this\n if (threadIdx.x == 0 && threadIdx.y == 0 && f < features) {\n unsigned int sum_idx;\n sum_idx = atomicAdd(&(semaphores[blockIdx.x]), 1);\n isLastBlock = (sum_idx == (gridDim.y - 1));\n }\n __syncthreads();\n\n db_local = 0;\n // No block reduction for now, only thread (*,0) do grid reduction\n if (isLastBlock && f < features) {\n if (threadIdx.y == 0) {\n for (int n = 0; n < gridDim.y; n++) {\n int row, col;\n row = f;\n col = n;\n db_local += (float)(intermediate[col * features + row]);\n }\n db[f] = (T)db_local;\n }\n }\n}\n\n// Addition done deterministically via a 2-pass approach. Each CTA writes out partial\n// sum, and the last CTA in grid Y dimension accumulates partials serially and writes to result.\ntemplate \n__global__ void biasAddRelu_bprop_aligned(T* Y, T* dY, int features, int batch_size, T* dX,\n volatile float* intermediate, int* semaphores, T* db) {\n // The feature that this thread is responsible for\n int f = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Compute the span this thread is responsible for\n // For this block\n int b_chunkSize = (batch_size + gridDim.y - 1) / gridDim.y;\n int b_nStart = blockIdx.y * b_chunkSize;\n int b_nSpan = min(batch_size, b_nStart + b_chunkSize) - b_nStart;\n // For this thread\n int chunkSize = (b_chunkSize + blockDim.y - 1) / blockDim.y;\n int nStart = threadIdx.y * chunkSize + b_nStart;\n int nSpan = min(b_nStart + b_nSpan, nStart + chunkSize) - nStart;\n\n volatile float* out = intermediate + blockIdx.y * features;\n\n // Flag to trigger last reduction.\n __shared__ bool isLastBlock;\n\n // Accumulate db in FP32 always\n float db_local[ILP];\n T r_y[ILP];\n T r_dy[ILP];\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n db_local[ii] = 0.f;\n }\n\n // f always <= features in this case\n // if (f < features) {\n int nidx = 0;\n\n // Handle non-multiple of UNROLL_FACTOR residue\n for (; nidx < nSpan % UNROLL_FACTOR; nidx++) {\n int row, col, flat_idx;\n row = f;\n col = nStart + nidx;\n flat_idx = col * features / ILP + row;\n\n load_store(r_y, Y, 0, flat_idx);\n load_store(r_dy, dY, 0, flat_idx);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n if ((float)r_y[ii] <= 0.f) r_dy[ii] = 0;\n db_local[ii] += (float)r_dy[ii];\n }\n load_store(dX, r_dy, flat_idx, 0);\n }\n\n // Handle meat of work\n for (; (nidx + UNROLL_FACTOR - 1) < nSpan; nidx += UNROLL_FACTOR) {\n int row, col, flat_idx;\n row = f;\n col = nStart + nidx;\n flat_idx = col * features / ILP + row; // total threads in x == features/ILP\n#pragma unroll\n for (int u = 0; u < UNROLL_FACTOR; u++) {\n load_store(r_y, Y, 0, flat_idx);\n load_store(r_dy, dY, 0, flat_idx);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n if ((float)r_y[ii] <= 0.f) r_dy[ii] = 0;\n db_local[ii] += (float)r_dy[ii];\n }\n load_store(dX, r_dy, flat_idx, 0);\n flat_idx += features / ILP;\n }\n }\n\n // we know block size for now\n __shared__ float smem[BIAS_RELU_BW_NTHREADS_X * BIAS_RELU_BW_NTHREADS_Y * ILP];\n // naive block reduction on y-dim\n int linear_idx = threadIdx.y * blockDim.x + threadIdx.x;\n float* smem_out = smem + ILP * linear_idx;\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n smem_out[ii] = db_local[ii]; // reuse local dy buffer\n }\n __syncthreads();\n if (threadIdx.y == 0) {\n for (int yidx = 1; yidx < blockDim.y; yidx++) {\n float* smem_in = smem + ILP * (yidx * blockDim.x + threadIdx.x);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n db_local[ii] += smem_in[ii]; // reuse local dy buffer\n }\n }\n\n // block result is in db_local now for all threadIdx.y == 0\n if (gridDim.y == 1) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_dy[ii] = db_local[ii]; // reuse local dy buffer\n }\n load_store(db, r_dy, f, 0);\n return;\n }\n\n // Write out partial result\n load_store(out, db_local, f, 0);\n }\n __threadfence();\n __syncthreads();\n\n // Increment semaphore and check if this is the last CTA in the grid_y dimension.\n // Only thread (0,0) calls this\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n unsigned int sum_idx;\n sum_idx = atomicAdd(&(semaphores[blockIdx.x]), 1);\n isLastBlock = (sum_idx == (gridDim.y - 1));\n }\n __syncthreads();\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n db_local[ii] = 0.f;\n }\n float r_db[ILP];\n\n // No block reduction for now, only thread (*,0) do grid reduction\n if (isLastBlock) {\n if (threadIdx.y == 0) {\n for (int n = 0; n < gridDim.y; n++) {\n int row, col;\n row = f;\n col = n;\n load_store(r_db, intermediate, 0, col * features / ILP + row);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n db_local[ii] += r_db[ii];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_dy[ii] = db_local[ii]; // reuse local dy buffer\n }\n load_store(db, r_dy, f, 0);\n }\n }\n}\n\n// Lists where the num_layers-1 intermediate Y buffers start in reserved space on fprop, starting\n// offset 0. The last Y value is, of course, stored in the user provided output buffer.\nvoid get_y_offsets(int batch_size, int num_layers, const int* output_features, int* y_start_offsets) {\n y_start_offsets[0] = 0;\n for (int i = 1; i < num_layers; i++) {\n y_start_offsets[i] = y_start_offsets[i - 1] + batch_size * output_features[i - 1];\n }\n}\n\n// Returns the reserved space (in elements) needed for the MLP\nsize_t get_mlp_reserved_space(int64_t batch_size, int num_layers, const int* output_features) {\n size_t res_space = 0;\n // Need to store output of every intermediate MLP - size equal to output_features[i] * batch_size\n // for all 'i' in [0, num_layers-1)\n for (int l = 0; l < num_layers; l++) {\n res_space += output_features[l] * batch_size;\n }\n return res_space;\n}\n\n// Returns the size of all fprop activations combined\nsize_t get_all_activations_size(int64_t batch_size, int num_layers, const int* output_features) {\n size_t acts_size = 0;\n for (int l = 0; l < num_layers; l++) {\n acts_size += output_features[l] * batch_size;\n }\n return acts_size;\n}\n\n#if 0\n// Returns the work space (in elements) needed for the MLP bprop.\nsize_t get_mlp_bp_workspace (int batch_size, int num_layers, const int* output_features) {\n /*\n Workspace is partitioned as\n DY_GEMMs : DX_GEMMs\n */\n size_t work_space = 0;\n\n // Store each intermediate dY explicitly. Need 2 dYs per MLP layer (one for o/p\n // of biasReLU_bp and one for o/p of dgrad GEMM).\n work_space += 2*get_all_activations_size(batch_size, num_layers, output_features);\n\n return work_space;\n}\n#endif\n\n// Scratch space needed for reductions in number of elements\nsize_t get_reduction_scratch_space(int batch_size, int num_layers, const int* output_features) {\n size_t max_scratch_space = 0;\n // Loop over all layers to see which one needs the max scratch space\n for (int l = 0; l < num_layers; l++) {\n // need to find max(aligned, not_aligned)\n int tmp, res0, res1;\n\n int block_x = BIAS_RELU_BW_NTHREADS_X;\n int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;\n get_biasAddRelu_bprop_grid_size(output_features[l], batch_size, block_x, block_y, &tmp, &res0);\n\n block_x = ILP * BIAS_RELU_BW_NTHREADS_X;\n get_biasAddRelu_bprop_grid_size(output_features[l], batch_size, block_x, block_y, &tmp, &res1);\n\n max_scratch_space = std::max(max_scratch_space, (size_t)(output_features[l] * res0));\n max_scratch_space = std::max(max_scratch_space, (size_t)(output_features[l] * res1));\n }\n\n return max_scratch_space;\n}\n\n// Buffer for semaphores\nsize_t get_semaphores_size(int num_layers, const int* output_features) {\n // Upper bound on semaphores is one per feature for the layer\n // with the most features.\n int max_features = 0;\n for (int l = 0; l < num_layers; l++) {\n max_features = std::max(max_features, output_features[l]);\n }\n return (size_t)max_features;\n}\n\n// Returns the work space (in elements) needed for the MLP bprop.\ntemplate \nsize_t get_mlp_bp_workspace_in_bytes(int batch_size, int num_layers, const int* output_features) {\n size_t work_space = 0;\n\n // Store each intermediate dY explicitly. Need 2 dYs per MLP layer (one for o/p\n // of biasReLU_bp and one for o/p of dgrad GEMM).\n work_space += 2 * get_all_activations_size(batch_size, num_layers, output_features) * sizeof(T);\n work_space += get_reduction_scratch_space(batch_size, num_layers, output_features) * sizeof(float);\n work_space += get_semaphores_size(num_layers, output_features) * sizeof(int);\n\n return work_space;\n}\n\n// Returns pointers to each segment of the workspace\ntemplate \nvoid partition_mlp_bp_workspace(int batch_size, int num_layers, const int* output_features, void* work_space,\n T** dy_gemms, T** dx_gemms, float** db_scratch, int** semaphores) {\n /*\n Workspace is partitioned as\n DY_GEMMs : DX_GEMMs : DB_SCRATCH : SEMAPHORES\n */\n // Start address where dy_gemm tensors are stored\n *dy_gemms = reinterpret_cast(work_space);\n // Start address where dx_gemm tensors are stored\n *dx_gemms = *dy_gemms + get_all_activations_size(batch_size, num_layers, output_features);\n // Start address where db intermediate tensors are stored\n *db_scratch = reinterpret_cast(*dx_gemms + get_all_activations_size(batch_size, num_layers, output_features));\n // Start address of semaphores\n *semaphores =\n reinterpret_cast(*db_scratch + get_reduction_scratch_space(batch_size, num_layers, output_features));\n\n return;\n}\n\n// Does a simple MLP fprop (GEMM+bias+ReLU).\n// Can handle num_layers number of layers, each with its own shape. Output of layer i is assumed\n// to be input of layer i+1. output_features, WPtr and BPtr are arrays of length num_layers, and\n// must be in the same order i.e. WPtr[i] and BPtr[i] are respectively the weight and bias of layer\n// 'i'.\ntemplate \nint mlp_fp(T* X, int input_features, int batch_size, T** WPtr, int num_layers, int* output_features, T** BPtr, T* Y,\n T* reserved_space, int use_bias, int activation, void* lt_workspace) {\n T *weight, *input, *output, *bias;\n T *reserved_space_x, *reserved_space_y;\n reserved_space_x = NULL;\n reserved_space_y = reserved_space;\n\n // Get cublas handle from Pytorch\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n // Get the stream from cublas handle to reuse for biasReLU kernel.\n cudaStream_t stream;\n cublasGetStream(handle, &stream);\n\n for (int layer = 0; layer < num_layers; layer++) {\n weight = WPtr[layer];\n input = (layer == 0) ? X : reserved_space_x;\n output = (layer == num_layers - 1) ? Y : reserved_space_y;\n if (use_bias) {\n bias = BPtr[layer];\n }\n int ifeat = (layer == 0) ? input_features : output_features[layer - 1];\n int ofeat = output_features[layer];\n\n float one = 1.f;\n float zero = 0.f;\n\n // try with cublaslt first for supported case with valid handle\n int cublaslt_status = 1;\n#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11000\n if (activation < 1) {\n cublaslt_status = mlp_gemm_lt(\n // ltHandle,\n (cublasLtHandle_t)handle, CUBLAS_OP_T, CUBLAS_OP_N, ofeat, batch_size, ifeat, &one, weight, ifeat, input,\n ifeat, &zero, output, ofeat, lt_workspace, 1 << 22, stream, use_bias == 1, activation == 1, bias);\n }\n#endif\n\n // if cublaslt failed or not executed, fallback to cublas\n if (cublaslt_status != 0) {\n cublasStatus_t cublas_status;\n // Call GEMM: fprop is Y = W'X\n cublas_status = mlp_gemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, ofeat, batch_size, ifeat, &one, weight, ifeat, input,\n ifeat, &zero, output, ofeat);\n\n if (cublas_status != CUBLAS_STATUS_SUCCESS) {\n printf(\"GEMM fprop failed with %d\\n\", cublas_status);\n return 1;\n }\n\n const uint& input_size = ofeat;\n int num_blocks = 0;\n int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;\n // Call biasReLU\n if (use_bias == 1) {\n if (activation == 0) { // no activation\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, biasAdd_fprop, BIAS_RELU_FW_NTHREADS, 0);\n biasAdd_fprop<<>>(output, bias, batch_size,\n input_size);\n } else if (activation == 1) { // relu\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, biasAddRelu_fprop, BIAS_RELU_FW_NTHREADS, 0);\n biasAddRelu_fprop<<>>(output, bias, batch_size,\n input_size);\n } else if (activation == 2) { // sigmoid\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, biasAdd_fprop, BIAS_RELU_FW_NTHREADS, 0);\n biasAdd_fprop<<>>(output, bias, batch_size,\n input_size);\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Sigmoid_fprop, BIAS_RELU_FW_NTHREADS, 0);\n Sigmoid_fprop<<>>(output, batch_size, input_size);\n }\n } else {\n // don't need to do anything in case of no activation and no bias\n if (activation == 1) { // relu\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Relu_fprop, BIAS_RELU_FW_NTHREADS, 0);\n Relu_fprop<<>>(output, batch_size, input_size);\n } else if (activation == 2) { // sigmoid\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Sigmoid_fprop, BIAS_RELU_FW_NTHREADS, 0);\n Sigmoid_fprop<<>>(output, batch_size, input_size);\n }\n }\n }\n // Set current output as next layer input\n reserved_space_x = reserved_space_y;\n // Set next layer output\n reserved_space_y += ofeat * batch_size;\n }\n\n return 0;\n}\n\n// Does a simple MLP bprop (GEMM+bias+ReLU).\n// Needs reserved space to come back exactly as it was populated in fprop.\n// Does dgrad and wgrad sequentially.\ntemplate \nint mlp_bp(T* X, T* Y, int input_features, int batch_size, T** WPtr, int num_layers, int* output_features, T* dY,\n T* reserved_space, T* work_space, T* dX, T** dwPtr, T** dbPtr, bool requires_grad, int use_bias,\n int activation) {\n T* weight;\n T *dweight, *dx, *dy, *dbias;\n T *x, *y;\n\n // Where the dx of the biasReLU (== dy of gemm) is stored. Can be thrown away\n // after bp call.\n T* dy_gemm_base;\n // Where the dx after GEMM is stored.\n T* dx_gemm_base;\n // Where partial reduction results are stored.\n float* db_scratch;\n // Semaphores for reduction.\n int* semaphores;\n\n partition_mlp_bp_workspace(batch_size, num_layers, output_features, work_space, &dy_gemm_base, &dx_gemm_base,\n &db_scratch, &semaphores);\n\n size_t semaphore_size = get_semaphores_size(num_layers, output_features) * sizeof(int);\n\n // Get cublas handle from Pytorch\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n // Get the stream from cublas handle to reuse for biasReLU kernel.\n cudaStream_t stream;\n cublasGetStream(handle, &stream);\n\n int* y_offsets = (int*)malloc(num_layers * sizeof(int));\n get_y_offsets(batch_size, num_layers, output_features, y_offsets);\n\n for (int layer = num_layers - 1; layer >= 0; layer--) {\n weight = WPtr[layer];\n dweight = dwPtr[layer];\n\n // x is read from reserved space\n x = (layer == 0) ? X : reserved_space + y_offsets[layer - 1];\n // dx is written in workspace for all but layer==0\n dx = (layer == 0) ? dX : dx_gemm_base + y_offsets[layer - 1];\n\n // y is read from reserved space\n y = (layer == num_layers - 1) ? Y : reserved_space + y_offsets[layer];\n // dx from layer+1\n dy = (layer == num_layers - 1) ? dY : dx_gemm_base + y_offsets[layer];\n // dy_gemm is written to and read immediately\n T* dy_gemm = dy_gemm_base + y_offsets[layer];\n\n dbias = dbPtr[layer];\n int xfeat = (layer == 0) ? input_features : output_features[layer - 1];\n int yfeat = output_features[layer];\n\n float one = 1.f;\n float zero = 0.f;\n\n if (use_bias == 1) {\n if (activation == 0) { // no acitvation\n // bgrad\n dim3 block(BIAS_RELU_BW_NTHREADS_X, BIAS_RELU_BW_NTHREADS_Y);\n int grid_x, grid_y;\n cudaMemsetAsync(semaphores, 0, semaphore_size, stream);\n\n int block_x = BIAS_RELU_BW_NTHREADS_X;\n int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;\n get_biasAddRelu_bprop_grid_size(yfeat, batch_size, block_x, block_y, &grid_x, &grid_y);\n dim3 grid(grid_x, grid_y);\n biasAdd_bprop<<>>(dy, yfeat, batch_size, db_scratch, semaphores, dbias);\n // bypass dgrad through reset pointer\n dy_gemm = dy;\n } else if (activation == 1) { // relu\n dim3 block(BIAS_RELU_BW_NTHREADS_X, BIAS_RELU_BW_NTHREADS_Y);\n int grid_x, grid_y;\n cudaMemsetAsync(semaphores, 0, semaphore_size, stream);\n\n if (yfeat % (ILP * BIAS_RELU_BW_NTHREADS_X) == 0 && is_aligned(y) && is_aligned(dy) && is_aligned(dy_gemm) &&\n is_aligned(dbias)) {\n int block_x = ILP * BIAS_RELU_BW_NTHREADS_X;\n int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;\n get_biasAddRelu_bprop_grid_size(yfeat, batch_size, block_x, block_y, &grid_x, &grid_y);\n dim3 grid(grid_x, grid_y);\n biasAddRelu_bprop_aligned\n <<>>(y, dy, yfeat, batch_size, dy_gemm, db_scratch, semaphores, dbias);\n } else {\n int block_x = BIAS_RELU_BW_NTHREADS_X;\n int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;\n get_biasAddRelu_bprop_grid_size(yfeat, batch_size, block_x, block_y, &grid_x, &grid_y);\n dim3 grid(grid_x, grid_y);\n biasAddRelu_bprop\n <<>>(y, dy, yfeat, batch_size, dy_gemm, db_scratch, semaphores, dbias);\n }\n } else if (activation == 2) { // sigmoid\n // activation backward\n int num_blocks = 0;\n int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Sigmoid_bprop, BIAS_RELU_FW_NTHREADS, 0);\n Sigmoid_bprop<<>>(dy, y, batch_size, yfeat, dy_gemm);\n\n // bgrad, from dy_gemm\n dim3 block(BIAS_RELU_BW_NTHREADS_X, BIAS_RELU_BW_NTHREADS_Y);\n int grid_x, grid_y;\n cudaMemsetAsync(semaphores, 0, semaphore_size, stream);\n\n int block_x = BIAS_RELU_BW_NTHREADS_X;\n int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;\n get_biasAddRelu_bprop_grid_size(yfeat, batch_size, block_x, block_y, &grid_x, &grid_y);\n dim3 grid(grid_x, grid_y);\n biasAdd_bprop<<>>(dy_gemm, yfeat, batch_size, db_scratch, semaphores, dbias);\n }\n } else { // no bias below\n if (activation == 0) {\n // bypass dgrad through reset pointer\n dy_gemm = dy;\n } else if (activation == 1) { // relu\n int num_blocks = 0;\n int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Relu_bprop, BIAS_RELU_FW_NTHREADS, 0);\n Relu_bprop<<>>(dy, y, batch_size, yfeat, dy_gemm);\n } else if (activation == 2) { // sigmoid\n int num_blocks = 0;\n int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Sigmoid_bprop, BIAS_RELU_FW_NTHREADS, 0);\n Sigmoid_bprop<<>>(dy, y, batch_size, yfeat, dy_gemm);\n }\n }\n\n cublasStatus_t cublas_status;\n // Call GEMM dgrad\n if (layer > 0 || requires_grad == 1) {\n cublas_status = mlp_gemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, xfeat, batch_size, yfeat, &one, weight, xfeat, dy_gemm,\n yfeat, &zero, dx, xfeat);\n\n if (cublas_status != CUBLAS_STATUS_SUCCESS) {\n printf(\"GEMM dgrad failed with %d\\n\", cublas_status);\n return 1;\n }\n }\n\n // Call GEMM wgrad\n cublas_status = mlp_gemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, xfeat, yfeat, batch_size, &one, x, xfeat, dy_gemm, yfeat,\n &zero, dweight, xfeat);\n\n if (cublas_status != CUBLAS_STATUS_SUCCESS) {\n printf(\"GEMM wgrad failed with %d\\n\", cublas_status);\n return 1;\n }\n }\n\n return 0;\n}\n\n// Instantiate for floating point types\ntemplate int mlp_fp(float* X, int input_features, int batch_size, float** WPtr, int num_layers,\n int* output_features, float** BPtr, float* Y, float* reserved_space, int use_bias,\n int activation, void* lt_workspace);\n\ntemplate int mlp_bp(float* X, float* Y, int input_features, int batch_size, float** WPtr, int num_layers,\n int* output_features, float* dY, float* reserved_space, float* work_space, float* dX,\n float** dwPtr, float** dbPtr, bool requires_grad, int use_bias, int activation);\n\ntemplate int mlp_fp(at::Half* X, int input_features, int batch_size, at::Half** WPtr, int num_layers,\n int* output_features, at::Half** BPtr, at::Half* Y, at::Half* reserved_space,\n int use_bias, int activation, void* lt_workspace);\n\ntemplate int mlp_bp(at::Half* X, at::Half* Y, int input_features, int batch_size, at::Half** WPtr,\n int num_layers, int* output_features, at::Half* dY, at::Half* reserved_space,\n at::Half* work_space, at::Half* dX, at::Half** dwPtr, at::Half** dbPtr,\n bool requires_grad, int use_bias, int activation);\n\ntemplate int mlp_fp(double* X, int input_features, int batch_size, double** WPtr, int num_layers,\n int* output_features, double** BPtr, double* Y, double* reserved_space, int use_bias,\n int activation, void* lt_workspace);\n\ntemplate int mlp_bp(double* X, double* Y, int input_features, int batch_size, double** WPtr, int num_layers,\n int* output_features, double* dY, double* reserved_space, double* work_space, double* dX,\n double** dwPtr, double** dbPtr, bool requires_grad, int use_bias, int activation);\n\ntemplate size_t get_mlp_bp_workspace_in_bytes(int batch_size, int num_layers, const int* output_features);\ntemplate size_t get_mlp_bp_workspace_in_bytes(int batch_size, int num_layers, const int* output_features);\ntemplate size_t get_mlp_bp_workspace_in_bytes(int batch_size, int num_layers, const int* output_features);\n", "output": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n/* Includes, cuda */\n#include \n#include \n\n#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11000\n// includes cublaslt\n#include \n#endif\n// constants for fused bias+relu kernel\n#define BIAS_RELU_FW_NTHREADS 128 // forward number of thread per block\n#define BIAS_RELU_BW_NTHREADS_X 32 // backward number of thread in feature dim\n#define BIAS_RELU_BW_NTHREADS_Y 16 // backward number of thread in batch dim\n#define BIAS_RELU_RED_PER_THREAD 16 // backward minimal reduction length per thread\n\n// move to a header later on\n#define ILP 4\ntemplate \n__host__ __device__ __forceinline__ bool is_aligned(T* p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\ntemplate \n__device__ __forceinline__ void load_store(T* dst, volatile T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\ntemplate \n__device__ __forceinline__ void load_store(volatile T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\n\n// Keep ReLU in float only. When using half, cast to float before calling.\n__device__ __inline__ float relu(float a) {\n float retf = max(a, 0.f);\n return (retf);\n}\n\n// Keep Sigmoid in float only. When using half, cast to float before calling.\n__device__ __inline__ float sigmoid(float a) {\n float retf = 1.f / (1.f + expf(-a));\n return (retf);\n}\n\n// FP64 Wrapper around cublas GEMMEx\nhipblasStatus_t mlp_gemm(hipblasHandle_t handle, hipblasOperation_t transa, hipblasOperation_t transb, int m, int n, int k,\n float* alpha, const double* A, int lda, const double* B, int ldb, const float* beta, double* C,\n int ldc) {\n return hipblasGemmEx(handle, transa, transb, m, n, k, alpha, A, HIP_R_64F, lda, B, HIP_R_64F, ldb, beta, C,\n HIP_R_64F, ldc, HIP_R_64F, HIPBLAS_GEMM_DEFAULT);\n}\n\n// FP32 Wrapper around cublas GEMMEx\nhipblasStatus_t mlp_gemm(hipblasHandle_t handle, hipblasOperation_t transa, hipblasOperation_t transb, int m, int n, int k,\n float* alpha, const float* A, int lda, const float* B, int ldb, const float* beta, float* C,\n int ldc) {\n return hipblasGemmEx(handle, transa, transb, m, n, k, alpha, A, HIP_R_32F, lda, B, HIP_R_32F, ldb, beta, C,\n HIP_R_32F, ldc, HIP_R_32F, HIPBLAS_GEMM_DEFAULT);\n}\n\n// FP16 Tensor core wrapper around cublas GEMMEx\nhipblasStatus_t mlp_gemm(hipblasHandle_t handle, hipblasOperation_t transa, hipblasOperation_t transb, int m, int n, int k,\n float* alpha, const at::Half* A, int lda, const at::Half* B, int ldb, float* beta, at::Half* C,\n int ldc) {\n return hipblasGemmEx(handle, transa, transb, m, n, k, alpha, A, HIP_R_16F, lda, B, HIP_R_16F, ldb, beta, C,\n HIP_R_16F, ldc, HIP_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP);\n}\n#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11000\nint mlp_gemm_lt(hipblasLtHandle_t ltHandle, hipblasOperation_t transa, hipblasOperation_t transb, int m, int n, int k,\n float* alpha, /* host pointer */\n const at::Half* A, int lda, const at::Half* B, int ldb, float* beta, /* host pointer */\n at::Half* C, int ldc, void* workspace, size_t workspaceSize, hipStream_t stream, bool use_bias,\n bool use_relu, const void* bias) {\n hipblasStatus_t status = HIPBLAS_STATUS_SUCCESS;\n\n hipblasLtMatmulDescOpaque_t operationDesc = {};\n hipblasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};\n hipblasLtMatmulPreferenceOpaque_t preference = {};\n\n int returnedResults = 0;\n hipblasLtMatmulHeuristicResult_t heuristicResult = {};\n hipblasLtEpilogue_t epilogue = HIPBLASLT_EPILOGUE_DEFAULT;\n\n // Create operation descriptor; see cublasLtMatmulDescAttributes_t\n // for details about defaults; here we just set the transforms for\n // A and B.\n status = cublasLtMatmulDescInit(&operationDesc, HIPBLAS_COMPUTE_32F, HIP_R_32F);\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = hipblasLtMatmulDescSetAttribute(&operationDesc, HIPBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = hipblasLtMatmulDescSetAttribute(&operationDesc, HIPBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa));\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n if (use_bias) {\n status = hipblasLtMatmulDescSetAttribute(&operationDesc, HIPBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias));\n if (status != HIPBLAS_STATUS_SUCCESS) {\n goto CLEANUP;\n }\n if (use_relu) {\n epilogue = HIPBLASLT_EPILOGUE_RELU_BIAS;\n } else {\n epilogue = HIPBLASLT_EPILOGUE_BIAS;\n }\n } else {\n if (use_relu) {\n epilogue = HIPBLASLT_EPILOGUE_RELU;\n }\n }\n\n status = hipblasLtMatmulDescSetAttribute(&operationDesc, HIPBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue));\n if (status != HIPBLAS_STATUS_SUCCESS) {\n goto CLEANUP;\n }\n\n // Create matrix descriptors. Not setting any extra attributes.\n status =\n cublasLtMatrixLayoutInit(&Adesc, HIP_R_16F, transa == HIPBLAS_OP_N ? m : k, transa == HIPBLAS_OP_N ? k : m, lda);\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n status =\n cublasLtMatrixLayoutInit(&Bdesc, HIP_R_16F, transb == HIPBLAS_OP_N ? k : n, transb == HIPBLAS_OP_N ? n : k, ldb);\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = cublasLtMatrixLayoutInit(&Cdesc, HIP_R_16F, m, n, ldc);\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n // Create preference handle; In general, extra attributes can be\n // used here to disable tensor ops or to make sure algo selected\n // will work with badly aligned A, B, C. However, for simplicity\n // here we assume A,B,C are always well aligned (e.g., directly\n // come from cudaMalloc)\n status = cublasLtMatmulPreferenceInit(&preference);\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = hipblasLtMatmulPreferenceSetAttribute(&preference, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize,\n sizeof(workspaceSize));\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n // We just need the best available heuristic to try and run matmul.\n // There is no guarantee that this will work. For example, if A is\n // badly aligned, you can request more (e.g. 32) algos and try to\n // run them one by one until something works.\n status = hipblasLtMatmulAlgoGetHeuristic(ltHandle, &operationDesc, &Adesc, &Bdesc, &Cdesc, &Cdesc, &preference, 1,\n &heuristicResult, &returnedResults);\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n if (returnedResults == 0) {\n status = HIPBLAS_STATUS_NOT_SUPPORTED;\n goto CLEANUP;\n }\n status = hipblasLtMatmul(ltHandle, &operationDesc, alpha, A, &Adesc, B, &Bdesc, beta, C, &Cdesc, C, &Cdesc,\n &heuristicResult.algo, workspace, workspaceSize, stream);\n\nCLEANUP:\n // Descriptors are no longer needed as all GPU work was already\n // enqueued.\n return status == HIPBLAS_STATUS_SUCCESS ? 0 : 1;\n}\n\nint mlp_gemm_lt(hipblasLtHandle_t ltHandle, hipblasOperation_t transa, hipblasOperation_t transb, int m, int n, int k,\n float* alpha, /* host pointer */\n const double* A, int lda, const double* B, int ldb, float* beta, /* host pointer */\n double* C, int ldc, void* workspace, size_t workspaceSize, hipStream_t stream, bool use_bias,\n bool use_relu, const void* bias) {\n return 1;\n}\n\nint mlp_gemm_lt(hipblasLtHandle_t ltHandle, hipblasOperation_t transa, hipblasOperation_t transb, int m, int n, int k,\n float* alpha, /* host pointer */\n const float* A, int lda, const float* B, int ldb, float* beta, /* host pointer */\n float* C, int ldc, void* workspace, size_t workspaceSize, hipStream_t stream, bool use_bias,\n bool use_relu, const void* bias) {\n hipblasStatus_t status = HIPBLAS_STATUS_SUCCESS;\n\n hipblasLtMatmulDescOpaque_t operationDesc = {};\n hipblasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};\n hipblasLtMatmulPreferenceOpaque_t preference = {};\n\n int returnedResults = 0;\n hipblasLtMatmulHeuristicResult_t heuristicResult = {};\n hipblasLtEpilogue_t epilogue = HIPBLASLT_EPILOGUE_DEFAULT;\n\n // Create operation descriptor; see cublasLtMatmulDescAttributes_t\n // for details about defaults; here we just set the transforms for\n // A and B.\n status = cublasLtMatmulDescInit(&operationDesc, HIPBLAS_COMPUTE_32F, HIP_R_32F);\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = hipblasLtMatmulDescSetAttribute(&operationDesc, HIPBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = hipblasLtMatmulDescSetAttribute(&operationDesc, HIPBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa));\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n if (use_bias) {\n status = hipblasLtMatmulDescSetAttribute(&operationDesc, HIPBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias));\n if (status != HIPBLAS_STATUS_SUCCESS) {\n goto CLEANUP;\n }\n if (use_relu) {\n epilogue = HIPBLASLT_EPILOGUE_RELU_BIAS;\n } else {\n epilogue = HIPBLASLT_EPILOGUE_BIAS;\n }\n } else {\n if (use_relu) {\n epilogue = HIPBLASLT_EPILOGUE_RELU;\n }\n }\n\n status = hipblasLtMatmulDescSetAttribute(&operationDesc, HIPBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue));\n if (status != HIPBLAS_STATUS_SUCCESS) {\n goto CLEANUP;\n }\n\n // Create matrix descriptors. Not setting any extra attributes.\n status =\n cublasLtMatrixLayoutInit(&Adesc, HIP_R_32F, transa == HIPBLAS_OP_N ? m : k, transa == HIPBLAS_OP_N ? k : m, lda);\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n status =\n cublasLtMatrixLayoutInit(&Bdesc, HIP_R_32F, transb == HIPBLAS_OP_N ? k : n, transb == HIPBLAS_OP_N ? n : k, ldb);\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = cublasLtMatrixLayoutInit(&Cdesc, HIP_R_32F, m, n, ldc);\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n // Create preference handle; In general, extra attributes can be\n // used here to disable tensor ops or to make sure algo selected\n // will work with badly aligned A, B, C. However, for simplicity\n // here we assume A,B,C are always well aligned (e.g., directly\n // come from cudaMalloc)\n status = cublasLtMatmulPreferenceInit(&preference);\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n status = hipblasLtMatmulPreferenceSetAttribute(&preference, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize,\n sizeof(workspaceSize));\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n // We just need the best available heuristic to try and run matmul.\n // There is no guarantee that this will work. For example, if A is\n // badly aligned, you can request more (e.g. 32) algos and try to\n // run them one by one until something works.\n status = hipblasLtMatmulAlgoGetHeuristic(ltHandle, &operationDesc, &Adesc, &Bdesc, &Cdesc, &Cdesc, &preference, 1,\n &heuristicResult, &returnedResults);\n if (status != HIPBLAS_STATUS_SUCCESS) goto CLEANUP;\n\n if (returnedResults == 0) {\n status = HIPBLAS_STATUS_NOT_SUPPORTED;\n goto CLEANUP;\n }\n\n status = hipblasLtMatmul(ltHandle, &operationDesc, alpha, A, &Adesc, B, &Bdesc, beta, C, &Cdesc, C, &Cdesc,\n &heuristicResult.algo, workspace, workspaceSize, stream);\n\nCLEANUP:\n // Descriptors are no longer needed as all GPU work was already\n // enqueued.\n return status == HIPBLAS_STATUS_SUCCESS ? 0 : 1;\n}\n#endif\n\n// Bias ADD. Assume input X is [features x batch size], column major.\n// Bias is one 'features' long vector, with implicit broadcast.\ntemplate \n__global__ void biasAdd_fprop(T* X, T* b, uint batch_size, uint features) {\n T r_x[ILP];\n T r_b[ILP];\n if (is_aligned(X) && is_aligned(b) && features % ILP == 0) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid * ILP < features * batch_size; tid += blockDim.x * gridDim.x) {\n int row = tid % (features / ILP);\n load_store(r_x, X, 0, tid);\n load_store(r_b, b, 0, row);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float bias_sum = static_cast(r_x[ii]) + static_cast(r_b[ii]);\n r_x[ii] = bias_sum;\n }\n load_store(X, r_x, tid, 0);\n }\n } else {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n int row = tid % features;\n r_x[ii] = X[idx];\n r_b[ii] = b[row];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float bias_sum = static_cast(r_x[ii]) + static_cast(r_b[ii]);\n r_x[ii] = bias_sum;\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n X[idx] = r_x[ii];\n }\n }\n }\n }\n}\n\n// Bias ADD + ReLU. Assume input X is [features x batch size], column major.\n// Activation support fuesed ReLU. Safe to call in-place.\ntemplate \n__global__ void biasAddRelu_fprop(T* X, T* b, uint batch_size, uint features) {\n T r_x[ILP];\n T r_b[ILP];\n if (is_aligned(X) && is_aligned(b) && features % ILP == 0) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid * ILP < features * batch_size; tid += blockDim.x * gridDim.x) {\n int row = tid % (features / ILP);\n load_store(r_x, X, 0, tid);\n load_store(r_b, b, 0, row);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float bias_sum = static_cast(r_x[ii]) + static_cast(r_b[ii]);\n r_x[ii] = relu(bias_sum);\n }\n load_store(X, r_x, tid, 0);\n }\n } else {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n int row = tid % features;\n r_x[ii] = X[idx];\n r_b[ii] = b[row];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float bias_sum = static_cast(r_x[ii]) + static_cast(r_b[ii]);\n r_x[ii] = relu(bias_sum);\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n X[idx] = r_x[ii];\n }\n }\n }\n }\n}\n\n// ReLU. Assume input X is [features x batch size], column major.\n// Safe to call in-place.\ntemplate \n__global__ void Relu_fprop(T* X, uint batch_size, uint features) {\n T r_x[ILP];\n if (is_aligned(X) && features % ILP == 0) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid * ILP < features * batch_size; tid += blockDim.x * gridDim.x) {\n load_store(r_x, X, 0, tid);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_x[ii] = relu(static_cast(r_x[ii]));\n }\n load_store(X, r_x, tid, 0);\n }\n } else {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n r_x[ii] = X[idx];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_x[ii] = relu(static_cast(r_x[ii]));\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n X[idx] = r_x[ii];\n }\n }\n }\n }\n}\n\n// Sigmoid. Assume input X is [features x batch size], column major.\n// Safe to call in-place.\ntemplate \n__global__ void Sigmoid_fprop(T* X, uint batch_size, uint features) {\n T r_x[ILP];\n if (is_aligned(X) && features % ILP == 0) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid * ILP < features * batch_size; tid += blockDim.x * gridDim.x) {\n load_store(r_x, X, 0, tid);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_x[ii] = sigmoid(static_cast(r_x[ii]));\n }\n load_store(X, r_x, tid, 0);\n }\n } else {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n r_x[ii] = X[idx];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_x[ii] = sigmoid(static_cast(r_x[ii]));\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n X[idx] = r_x[ii];\n }\n }\n }\n }\n}\n\n// ReLU. Assume input X is [features x batch size], column major.\n// Safe to call in-place.\ntemplate \n__global__ void Relu_bprop(T* dY, T* Y, uint batch_size, uint features, T* dX) {\n T r_dy[ILP];\n T r_y[ILP];\n if (is_aligned(dY) && is_aligned(Y) && is_aligned(dX) && features % ILP == 0) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid * ILP < features * batch_size; tid += blockDim.x * gridDim.x) {\n load_store(r_dy, dY, 0, tid);\n load_store(r_y, Y, 0, tid);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n if ((float)r_y[ii] <= 0.f) r_dy[ii] = 0;\n }\n load_store(dX, r_dy, tid, 0);\n }\n } else {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n r_dy[ii] = dY[idx];\n r_y[ii] = Y[idx];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n if ((float)r_y[ii] <= 0.f) r_dy[ii] = 0;\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n dX[idx] = r_dy[ii];\n }\n }\n }\n }\n}\n\n// Sigmoid. Assume input X is [features x batch size], column major.\n// Safe to call in-place.\ntemplate \n__global__ void Sigmoid_bprop(T* dY, T* Y, uint batch_size, uint features, T* dX) {\n T r_dy[ILP];\n T r_y[ILP];\n if (is_aligned(dY) && is_aligned(Y) && is_aligned(dX) && features % ILP == 0) {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid * ILP < features * batch_size; tid += blockDim.x * gridDim.x) {\n load_store(r_dy, dY, 0, tid);\n load_store(r_y, Y, 0, tid);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float grad_out = r_dy[ii];\n float out = r_y[ii];\n float grad_i = out * (1.f - out) * grad_out;\n r_dy[ii] = grad_i;\n }\n load_store(dX, r_dy, tid, 0);\n }\n } else {\n int tid = blockIdx.x * blockDim.x + threadIdx.x;\n for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n r_dy[ii] = dY[idx];\n r_y[ii] = Y[idx];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float grad_out = r_dy[ii];\n float out = r_y[ii];\n float grad_i = out * (1.f - out) * grad_out;\n r_dy[ii] = grad_i;\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int idx = tid + ii * blockDim.x * gridDim.x;\n if (idx < features * batch_size) {\n dX[idx] = r_dy[ii];\n }\n }\n }\n }\n}\n\n// Compute grid size for pointwise backward kernel.\n// block_x/y is total elment being handled per block, not number of threads\nvoid get_biasAddRelu_bprop_grid_size(int yfeat, int batch_size, int block_x, int block_y, int* grid_x, int* grid_y) {\n *grid_x = (yfeat + block_x - 1) / block_x;\n // Get number of SMs for efficient reduction.\n int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;\n // can switch to occupancy calculation. use 4 below now for sm_70\n int max_blocks_y = (num_SMs * 4 + (*grid_x) - 1) / (*grid_x);\n // block_y should be from minimal work per thread\n int nRedSplits = (batch_size + block_y - 1) / block_y;\n // increase number of elem per thread redcution to not launch more than enough\n // kernel adjust work, so here we just launch max block\n *grid_y = std::min(nRedSplits, max_blocks_y);\n return;\n}\n\n// Addition done deterministically via a 2-pass approach. Each CTA writes out partial\n// sum, and the last CTA in grid Y dimension accumulates partials serially and writes to result.\ntemplate \n__global__ void biasAdd_bprop(T* dY, int features, int batch_size, volatile float* intermediate, int* semaphores,\n T* db) {\n // The feature that this thread is responsible for\n int f = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Compute the span this thread is responsible for\n // For this block\n int b_chunkSize = (batch_size + gridDim.y - 1) / gridDim.y;\n int b_nStart = blockIdx.y * b_chunkSize;\n int b_nSpan = min(batch_size, b_nStart + b_chunkSize) - b_nStart;\n // For this thread\n int chunkSize = (b_chunkSize + blockDim.y - 1) / blockDim.y;\n int nStart = threadIdx.y * chunkSize + b_nStart;\n int nSpan = min(b_nStart + b_nSpan, nStart + chunkSize) - nStart;\n\n volatile float* out = intermediate + blockIdx.y * features;\n\n // Flag to trigger last reduction.\n __shared__ bool isLastBlock;\n // we know block size for now\n __shared__ float smem[BIAS_RELU_BW_NTHREADS_X * BIAS_RELU_BW_NTHREADS_Y];\n\n // Accumulate db in FP32 always\n float db_local = 0;\n if (f < features) {\n int nidx = 0;\n // Handle non-multiple of UNROLL_FACTOR residue\n for (; nidx < nSpan % UNROLL_FACTOR; nidx++) {\n int64_t row, col, flat_idx;\n row = f;\n col = nStart + nidx;\n flat_idx = col * features + row;\n db_local += (float)dY[flat_idx];\n }\n\n // Handle meat of work\n for (; (nidx + UNROLL_FACTOR - 1) < nSpan; nidx += UNROLL_FACTOR) {\n int64_t row, col, flat_idx;\n row = f;\n col = nStart + nidx;\n flat_idx = col * features + row;\n#pragma unroll 4\n for (int u = 0; u < UNROLL_FACTOR; u++) {\n db_local += (float)dY[flat_idx];\n flat_idx += features;\n }\n }\n\n // naive block reduction on y-dim\n int linear_idx = threadIdx.y * blockDim.x + threadIdx.x;\n smem[linear_idx] = db_local;\n }\n __syncthreads();\n if (f < features) {\n if (threadIdx.y == 0) {\n for (int yidx = 1; yidx < blockDim.y; yidx++) {\n db_local += smem[yidx * blockDim.x + threadIdx.x];\n }\n\n // block result is in db_local now for all threadIdx.y == 0\n // Write out partial result\n out[f] = db_local;\n }\n }\n __threadfence();\n __syncthreads();\n\n // Increment semaphore and check if this is the last CTA in the grid_y dimension.\n // Only thread (0,0) calls this\n if (threadIdx.x == 0 && threadIdx.y == 0 && f < features) {\n unsigned int sum_idx;\n sum_idx = atomicAdd(&(semaphores[blockIdx.x]), 1);\n isLastBlock = (sum_idx == (gridDim.y - 1));\n }\n __syncthreads();\n\n db_local = 0;\n // No block reduction for now, only thread (*,0) do grid reduction\n if (isLastBlock && f < features) {\n if (threadIdx.y == 0) {\n for (int n = 0; n < gridDim.y; n++) {\n int row, col;\n row = f;\n col = n;\n db_local += (float)(intermediate[col * features + row]);\n }\n db[f] = (T)db_local;\n }\n }\n}\n\n// Addition done deterministically via a 2-pass approach. Each CTA writes out partial\n// sum, and the last CTA in grid Y dimension accumulates partials serially and writes to result.\ntemplate \n__global__ void biasAddRelu_bprop(T* Y, T* dY, int features, int batch_size, T* dX, volatile float* intermediate,\n int* semaphores, T* db) {\n // The feature that this thread is responsible for\n int f = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Compute the span this thread is responsible for\n // For this block\n int b_chunkSize = (batch_size + gridDim.y - 1) / gridDim.y;\n int b_nStart = blockIdx.y * b_chunkSize;\n int b_nSpan = min(batch_size, b_nStart + b_chunkSize) - b_nStart;\n // For this thread\n int chunkSize = (b_chunkSize + blockDim.y - 1) / blockDim.y;\n int nStart = threadIdx.y * chunkSize + b_nStart;\n int nSpan = min(b_nStart + b_nSpan, nStart + chunkSize) - nStart;\n\n volatile float* out = intermediate + blockIdx.y * features;\n\n // Flag to trigger last reduction.\n __shared__ bool isLastBlock;\n // we know block size for now\n __shared__ float smem[BIAS_RELU_BW_NTHREADS_X * BIAS_RELU_BW_NTHREADS_Y];\n\n // Accumulate db in FP32 always\n float db_local = 0;\n if (f < features) {\n int nidx = 0;\n // Handle non-multiple of UNROLL_FACTOR residue\n for (; nidx < nSpan % UNROLL_FACTOR; nidx++) {\n int row, col, flat_idx;\n row = f;\n col = nStart + nidx;\n flat_idx = col * features + row;\n T y_val = Y[flat_idx];\n T dy_val = dY[flat_idx];\n T dx_val;\n if ((float)y_val > 0.f)\n dx_val = dy_val;\n else\n dx_val = 0;\n dX[flat_idx] = dx_val;\n db_local += (float)dx_val;\n }\n\n // Handle meat of work\n for (; (nidx + UNROLL_FACTOR - 1) < nSpan; nidx += UNROLL_FACTOR) {\n int row, col, flat_idx;\n row = f;\n col = nStart + nidx;\n flat_idx = col * features + row;\n#pragma unroll 4\n for (int u = 0; u < UNROLL_FACTOR; u++) {\n T y_val = Y[flat_idx];\n T dy_val = dY[flat_idx];\n T dx_val;\n if ((float)y_val > 0.f)\n dx_val = dy_val;\n else\n dx_val = 0;\n dX[flat_idx] = dx_val;\n db_local += (float)dx_val;\n flat_idx += features;\n }\n }\n\n // naive block reduction on y-dim\n int linear_idx = threadIdx.y * blockDim.x + threadIdx.x;\n smem[linear_idx] = db_local;\n }\n __syncthreads();\n if (f < features) {\n if (threadIdx.y == 0) {\n for (int yidx = 1; yidx < blockDim.y; yidx++) {\n db_local += smem[yidx * blockDim.x + threadIdx.x];\n }\n\n // block result is in db_local now for all threadIdx.y == 0\n // Write out partial result\n out[f] = db_local;\n }\n }\n __threadfence();\n __syncthreads();\n\n // Increment semaphore and check if this is the last CTA in the grid_y dimension.\n // Only thread (0,0) calls this\n if (threadIdx.x == 0 && threadIdx.y == 0 && f < features) {\n unsigned int sum_idx;\n sum_idx = atomicAdd(&(semaphores[blockIdx.x]), 1);\n isLastBlock = (sum_idx == (gridDim.y - 1));\n }\n __syncthreads();\n\n db_local = 0;\n // No block reduction for now, only thread (*,0) do grid reduction\n if (isLastBlock && f < features) {\n if (threadIdx.y == 0) {\n for (int n = 0; n < gridDim.y; n++) {\n int row, col;\n row = f;\n col = n;\n db_local += (float)(intermediate[col * features + row]);\n }\n db[f] = (T)db_local;\n }\n }\n}\n\n// Addition done deterministically via a 2-pass approach. Each CTA writes out partial\n// sum, and the last CTA in grid Y dimension accumulates partials serially and writes to result.\ntemplate \n__global__ void biasAddRelu_bprop_aligned(T* Y, T* dY, int features, int batch_size, T* dX,\n volatile float* intermediate, int* semaphores, T* db) {\n // The feature that this thread is responsible for\n int f = blockIdx.x * blockDim.x + threadIdx.x;\n\n // Compute the span this thread is responsible for\n // For this block\n int b_chunkSize = (batch_size + gridDim.y - 1) / gridDim.y;\n int b_nStart = blockIdx.y * b_chunkSize;\n int b_nSpan = min(batch_size, b_nStart + b_chunkSize) - b_nStart;\n // For this thread\n int chunkSize = (b_chunkSize + blockDim.y - 1) / blockDim.y;\n int nStart = threadIdx.y * chunkSize + b_nStart;\n int nSpan = min(b_nStart + b_nSpan, nStart + chunkSize) - nStart;\n\n volatile float* out = intermediate + blockIdx.y * features;\n\n // Flag to trigger last reduction.\n __shared__ bool isLastBlock;\n\n // Accumulate db in FP32 always\n float db_local[ILP];\n T r_y[ILP];\n T r_dy[ILP];\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n db_local[ii] = 0.f;\n }\n\n // f always <= features in this case\n // if (f < features) {\n int nidx = 0;\n\n // Handle non-multiple of UNROLL_FACTOR residue\n for (; nidx < nSpan % UNROLL_FACTOR; nidx++) {\n int row, col, flat_idx;\n row = f;\n col = nStart + nidx;\n flat_idx = col * features / ILP + row;\n\n load_store(r_y, Y, 0, flat_idx);\n load_store(r_dy, dY, 0, flat_idx);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n if ((float)r_y[ii] <= 0.f) r_dy[ii] = 0;\n db_local[ii] += (float)r_dy[ii];\n }\n load_store(dX, r_dy, flat_idx, 0);\n }\n\n // Handle meat of work\n for (; (nidx + UNROLL_FACTOR - 1) < nSpan; nidx += UNROLL_FACTOR) {\n int row, col, flat_idx;\n row = f;\n col = nStart + nidx;\n flat_idx = col * features / ILP + row; // total threads in x == features/ILP\n#pragma unroll\n for (int u = 0; u < UNROLL_FACTOR; u++) {\n load_store(r_y, Y, 0, flat_idx);\n load_store(r_dy, dY, 0, flat_idx);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n if ((float)r_y[ii] <= 0.f) r_dy[ii] = 0;\n db_local[ii] += (float)r_dy[ii];\n }\n load_store(dX, r_dy, flat_idx, 0);\n flat_idx += features / ILP;\n }\n }\n\n // we know block size for now\n __shared__ float smem[BIAS_RELU_BW_NTHREADS_X * BIAS_RELU_BW_NTHREADS_Y * ILP];\n // naive block reduction on y-dim\n int linear_idx = threadIdx.y * blockDim.x + threadIdx.x;\n float* smem_out = smem + ILP * linear_idx;\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n smem_out[ii] = db_local[ii]; // reuse local dy buffer\n }\n __syncthreads();\n if (threadIdx.y == 0) {\n for (int yidx = 1; yidx < blockDim.y; yidx++) {\n float* smem_in = smem + ILP * (yidx * blockDim.x + threadIdx.x);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n db_local[ii] += smem_in[ii]; // reuse local dy buffer\n }\n }\n\n // block result is in db_local now for all threadIdx.y == 0\n if (gridDim.y == 1) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_dy[ii] = db_local[ii]; // reuse local dy buffer\n }\n load_store(db, r_dy, f, 0);\n return;\n }\n\n // Write out partial result\n load_store(out, db_local, f, 0);\n }\n __threadfence();\n __syncthreads();\n\n // Increment semaphore and check if this is the last CTA in the grid_y dimension.\n // Only thread (0,0) calls this\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n unsigned int sum_idx;\n sum_idx = atomicAdd(&(semaphores[blockIdx.x]), 1);\n isLastBlock = (sum_idx == (gridDim.y - 1));\n }\n __syncthreads();\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n db_local[ii] = 0.f;\n }\n float r_db[ILP];\n\n // No block reduction for now, only thread (*,0) do grid reduction\n if (isLastBlock) {\n if (threadIdx.y == 0) {\n for (int n = 0; n < gridDim.y; n++) {\n int row, col;\n row = f;\n col = n;\n load_store(r_db, intermediate, 0, col * features / ILP + row);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n db_local[ii] += r_db[ii];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_dy[ii] = db_local[ii]; // reuse local dy buffer\n }\n load_store(db, r_dy, f, 0);\n }\n }\n}\n\n// Lists where the num_layers-1 intermediate Y buffers start in reserved space on fprop, starting\n// offset 0. The last Y value is, of course, stored in the user provided output buffer.\nvoid get_y_offsets(int batch_size, int num_layers, const int* output_features, int* y_start_offsets) {\n y_start_offsets[0] = 0;\n for (int i = 1; i < num_layers; i++) {\n y_start_offsets[i] = y_start_offsets[i - 1] + batch_size * output_features[i - 1];\n }\n}\n\n// Returns the reserved space (in elements) needed for the MLP\nsize_t get_mlp_reserved_space(int64_t batch_size, int num_layers, const int* output_features) {\n size_t res_space = 0;\n // Need to store output of every intermediate MLP - size equal to output_features[i] * batch_size\n // for all 'i' in [0, num_layers-1)\n for (int l = 0; l < num_layers; l++) {\n res_space += output_features[l] * batch_size;\n }\n return res_space;\n}\n\n// Returns the size of all fprop activations combined\nsize_t get_all_activations_size(int64_t batch_size, int num_layers, const int* output_features) {\n size_t acts_size = 0;\n for (int l = 0; l < num_layers; l++) {\n acts_size += output_features[l] * batch_size;\n }\n return acts_size;\n}\n\n#if 0\n// Returns the work space (in elements) needed for the MLP bprop.\nsize_t get_mlp_bp_workspace (int batch_size, int num_layers, const int* output_features) {\n /*\n Workspace is partitioned as\n DY_GEMMs : DX_GEMMs\n */\n size_t work_space = 0;\n\n // Store each intermediate dY explicitly. Need 2 dYs per MLP layer (one for o/p\n // of biasReLU_bp and one for o/p of dgrad GEMM).\n work_space += 2*get_all_activations_size(batch_size, num_layers, output_features);\n\n return work_space;\n}\n#endif\n\n// Scratch space needed for reductions in number of elements\nsize_t get_reduction_scratch_space(int batch_size, int num_layers, const int* output_features) {\n size_t max_scratch_space = 0;\n // Loop over all layers to see which one needs the max scratch space\n for (int l = 0; l < num_layers; l++) {\n // need to find max(aligned, not_aligned)\n int tmp, res0, res1;\n\n int block_x = BIAS_RELU_BW_NTHREADS_X;\n int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;\n get_biasAddRelu_bprop_grid_size(output_features[l], batch_size, block_x, block_y, &tmp, &res0);\n\n block_x = ILP * BIAS_RELU_BW_NTHREADS_X;\n get_biasAddRelu_bprop_grid_size(output_features[l], batch_size, block_x, block_y, &tmp, &res1);\n\n max_scratch_space = std::max(max_scratch_space, (size_t)(output_features[l] * res0));\n max_scratch_space = std::max(max_scratch_space, (size_t)(output_features[l] * res1));\n }\n\n return max_scratch_space;\n}\n\n// Buffer for semaphores\nsize_t get_semaphores_size(int num_layers, const int* output_features) {\n // Upper bound on semaphores is one per feature for the layer\n // with the most features.\n int max_features = 0;\n for (int l = 0; l < num_layers; l++) {\n max_features = std::max(max_features, output_features[l]);\n }\n return (size_t)max_features;\n}\n\n// Returns the work space (in elements) needed for the MLP bprop.\ntemplate \nsize_t get_mlp_bp_workspace_in_bytes(int batch_size, int num_layers, const int* output_features) {\n size_t work_space = 0;\n\n // Store each intermediate dY explicitly. Need 2 dYs per MLP layer (one for o/p\n // of biasReLU_bp and one for o/p of dgrad GEMM).\n work_space += 2 * get_all_activations_size(batch_size, num_layers, output_features) * sizeof(T);\n work_space += get_reduction_scratch_space(batch_size, num_layers, output_features) * sizeof(float);\n work_space += get_semaphores_size(num_layers, output_features) * sizeof(int);\n\n return work_space;\n}\n\n// Returns pointers to each segment of the workspace\ntemplate \nvoid partition_mlp_bp_workspace(int batch_size, int num_layers, const int* output_features, void* work_space,\n T** dy_gemms, T** dx_gemms, float** db_scratch, int** semaphores) {\n /*\n Workspace is partitioned as\n DY_GEMMs : DX_GEMMs : DB_SCRATCH : SEMAPHORES\n */\n // Start address where dy_gemm tensors are stored\n *dy_gemms = reinterpret_cast(work_space);\n // Start address where dx_gemm tensors are stored\n *dx_gemms = *dy_gemms + get_all_activations_size(batch_size, num_layers, output_features);\n // Start address where db intermediate tensors are stored\n *db_scratch = reinterpret_cast(*dx_gemms + get_all_activations_size(batch_size, num_layers, output_features));\n // Start address of semaphores\n *semaphores =\n reinterpret_cast(*db_scratch + get_reduction_scratch_space(batch_size, num_layers, output_features));\n\n return;\n}\n\n// Does a simple MLP fprop (GEMM+bias+ReLU).\n// Can handle num_layers number of layers, each with its own shape. Output of layer i is assumed\n// to be input of layer i+1. output_features, WPtr and BPtr are arrays of length num_layers, and\n// must be in the same order i.e. WPtr[i] and BPtr[i] are respectively the weight and bias of layer\n// 'i'.\ntemplate \nint mlp_fp(T* X, int input_features, int batch_size, T** WPtr, int num_layers, int* output_features, T** BPtr, T* Y,\n T* reserved_space, int use_bias, int activation, void* lt_workspace) {\n T *weight, *input, *output, *bias;\n T *reserved_space_x, *reserved_space_y;\n reserved_space_x = NULL;\n reserved_space_y = reserved_space;\n\n // Get cublas handle from Pytorch\n hipblasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n // Get the stream from cublas handle to reuse for biasReLU kernel.\n hipStream_t stream;\n hipblasGetStream(handle, &stream);\n\n for (int layer = 0; layer < num_layers; layer++) {\n weight = WPtr[layer];\n input = (layer == 0) ? X : reserved_space_x;\n output = (layer == num_layers - 1) ? Y : reserved_space_y;\n if (use_bias) {\n bias = BPtr[layer];\n }\n int ifeat = (layer == 0) ? input_features : output_features[layer - 1];\n int ofeat = output_features[layer];\n\n float one = 1.f;\n float zero = 0.f;\n\n // try with cublaslt first for supported case with valid handle\n int cublaslt_status = 1;\n#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11000\n if (activation < 1) {\n cublaslt_status = mlp_gemm_lt(\n // ltHandle,\n (hipblasLtHandle_t)handle, HIPBLAS_OP_T, HIPBLAS_OP_N, ofeat, batch_size, ifeat, &one, weight, ifeat, input,\n ifeat, &zero, output, ofeat, lt_workspace, 1 << 22, stream, use_bias == 1, activation == 1, bias);\n }\n#endif\n\n // if cublaslt failed or not executed, fallback to cublas\n if (cublaslt_status != 0) {\n hipblasStatus_t cublas_status;\n // Call GEMM: fprop is Y = W'X\n cublas_status = mlp_gemm(handle, HIPBLAS_OP_T, HIPBLAS_OP_N, ofeat, batch_size, ifeat, &one, weight, ifeat, input,\n ifeat, &zero, output, ofeat);\n\n if (cublas_status != HIPBLAS_STATUS_SUCCESS) {\n printf(\"GEMM fprop failed with %d\\n\", cublas_status);\n return 1;\n }\n\n const uint& input_size = ofeat;\n int num_blocks = 0;\n int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;\n // Call biasReLU\n if (use_bias == 1) {\n if (activation == 0) { // no activation\n hipOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, biasAdd_fprop, BIAS_RELU_FW_NTHREADS, 0);\n biasAdd_fprop<<>>(output, bias, batch_size,\n input_size);\n } else if (activation == 1) { // relu\n hipOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, biasAddRelu_fprop, BIAS_RELU_FW_NTHREADS, 0);\n biasAddRelu_fprop<<>>(output, bias, batch_size,\n input_size);\n } else if (activation == 2) { // sigmoid\n hipOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, biasAdd_fprop, BIAS_RELU_FW_NTHREADS, 0);\n biasAdd_fprop<<>>(output, bias, batch_size,\n input_size);\n hipOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Sigmoid_fprop, BIAS_RELU_FW_NTHREADS, 0);\n Sigmoid_fprop<<>>(output, batch_size, input_size);\n }\n } else {\n // don't need to do anything in case of no activation and no bias\n if (activation == 1) { // relu\n hipOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Relu_fprop, BIAS_RELU_FW_NTHREADS, 0);\n Relu_fprop<<>>(output, batch_size, input_size);\n } else if (activation == 2) { // sigmoid\n hipOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Sigmoid_fprop, BIAS_RELU_FW_NTHREADS, 0);\n Sigmoid_fprop<<>>(output, batch_size, input_size);\n }\n }\n }\n // Set current output as next layer input\n reserved_space_x = reserved_space_y;\n // Set next layer output\n reserved_space_y += ofeat * batch_size;\n }\n\n return 0;\n}\n\n// Does a simple MLP bprop (GEMM+bias+ReLU).\n// Needs reserved space to come back exactly as it was populated in fprop.\n// Does dgrad and wgrad sequentially.\ntemplate \nint mlp_bp(T* X, T* Y, int input_features, int batch_size, T** WPtr, int num_layers, int* output_features, T* dY,\n T* reserved_space, T* work_space, T* dX, T** dwPtr, T** dbPtr, bool requires_grad, int use_bias,\n int activation) {\n T* weight;\n T *dweight, *dx, *dy, *dbias;\n T *x, *y;\n\n // Where the dx of the biasReLU (== dy of gemm) is stored. Can be thrown away\n // after bp call.\n T* dy_gemm_base;\n // Where the dx after GEMM is stored.\n T* dx_gemm_base;\n // Where partial reduction results are stored.\n float* db_scratch;\n // Semaphores for reduction.\n int* semaphores;\n\n partition_mlp_bp_workspace(batch_size, num_layers, output_features, work_space, &dy_gemm_base, &dx_gemm_base,\n &db_scratch, &semaphores);\n\n size_t semaphore_size = get_semaphores_size(num_layers, output_features) * sizeof(int);\n\n // Get cublas handle from Pytorch\n hipblasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n // Get the stream from cublas handle to reuse for biasReLU kernel.\n hipStream_t stream;\n hipblasGetStream(handle, &stream);\n\n int* y_offsets = (int*)malloc(num_layers * sizeof(int));\n get_y_offsets(batch_size, num_layers, output_features, y_offsets);\n\n for (int layer = num_layers - 1; layer >= 0; layer--) {\n weight = WPtr[layer];\n dweight = dwPtr[layer];\n\n // x is read from reserved space\n x = (layer == 0) ? X : reserved_space + y_offsets[layer - 1];\n // dx is written in workspace for all but layer==0\n dx = (layer == 0) ? dX : dx_gemm_base + y_offsets[layer - 1];\n\n // y is read from reserved space\n y = (layer == num_layers - 1) ? Y : reserved_space + y_offsets[layer];\n // dx from layer+1\n dy = (layer == num_layers - 1) ? dY : dx_gemm_base + y_offsets[layer];\n // dy_gemm is written to and read immediately\n T* dy_gemm = dy_gemm_base + y_offsets[layer];\n\n dbias = dbPtr[layer];\n int xfeat = (layer == 0) ? input_features : output_features[layer - 1];\n int yfeat = output_features[layer];\n\n float one = 1.f;\n float zero = 0.f;\n\n if (use_bias == 1) {\n if (activation == 0) { // no acitvation\n // bgrad\n dim3 block(BIAS_RELU_BW_NTHREADS_X, BIAS_RELU_BW_NTHREADS_Y);\n int grid_x, grid_y;\n hipMemsetAsync(semaphores, 0, semaphore_size, stream);\n\n int block_x = BIAS_RELU_BW_NTHREADS_X;\n int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;\n get_biasAddRelu_bprop_grid_size(yfeat, batch_size, block_x, block_y, &grid_x, &grid_y);\n dim3 grid(grid_x, grid_y);\n biasAdd_bprop<<>>(dy, yfeat, batch_size, db_scratch, semaphores, dbias);\n // bypass dgrad through reset pointer\n dy_gemm = dy;\n } else if (activation == 1) { // relu\n dim3 block(BIAS_RELU_BW_NTHREADS_X, BIAS_RELU_BW_NTHREADS_Y);\n int grid_x, grid_y;\n hipMemsetAsync(semaphores, 0, semaphore_size, stream);\n\n if (yfeat % (ILP * BIAS_RELU_BW_NTHREADS_X) == 0 && is_aligned(y) && is_aligned(dy) && is_aligned(dy_gemm) &&\n is_aligned(dbias)) {\n int block_x = ILP * BIAS_RELU_BW_NTHREADS_X;\n int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;\n get_biasAddRelu_bprop_grid_size(yfeat, batch_size, block_x, block_y, &grid_x, &grid_y);\n dim3 grid(grid_x, grid_y);\n biasAddRelu_bprop_aligned\n <<>>(y, dy, yfeat, batch_size, dy_gemm, db_scratch, semaphores, dbias);\n } else {\n int block_x = BIAS_RELU_BW_NTHREADS_X;\n int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;\n get_biasAddRelu_bprop_grid_size(yfeat, batch_size, block_x, block_y, &grid_x, &grid_y);\n dim3 grid(grid_x, grid_y);\n biasAddRelu_bprop\n <<>>(y, dy, yfeat, batch_size, dy_gemm, db_scratch, semaphores, dbias);\n }\n } else if (activation == 2) { // sigmoid\n // activation backward\n int num_blocks = 0;\n int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;\n hipOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Sigmoid_bprop, BIAS_RELU_FW_NTHREADS, 0);\n Sigmoid_bprop<<>>(dy, y, batch_size, yfeat, dy_gemm);\n\n // bgrad, from dy_gemm\n dim3 block(BIAS_RELU_BW_NTHREADS_X, BIAS_RELU_BW_NTHREADS_Y);\n int grid_x, grid_y;\n hipMemsetAsync(semaphores, 0, semaphore_size, stream);\n\n int block_x = BIAS_RELU_BW_NTHREADS_X;\n int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;\n get_biasAddRelu_bprop_grid_size(yfeat, batch_size, block_x, block_y, &grid_x, &grid_y);\n dim3 grid(grid_x, grid_y);\n biasAdd_bprop<<>>(dy_gemm, yfeat, batch_size, db_scratch, semaphores, dbias);\n }\n } else { // no bias below\n if (activation == 0) {\n // bypass dgrad through reset pointer\n dy_gemm = dy;\n } else if (activation == 1) { // relu\n int num_blocks = 0;\n int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;\n hipOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Relu_bprop, BIAS_RELU_FW_NTHREADS, 0);\n Relu_bprop<<>>(dy, y, batch_size, yfeat, dy_gemm);\n } else if (activation == 2) { // sigmoid\n int num_blocks = 0;\n int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;\n hipOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Sigmoid_bprop, BIAS_RELU_FW_NTHREADS, 0);\n Sigmoid_bprop<<>>(dy, y, batch_size, yfeat, dy_gemm);\n }\n }\n\n hipblasStatus_t cublas_status;\n // Call GEMM dgrad\n if (layer > 0 || requires_grad == 1) {\n cublas_status = mlp_gemm(handle, HIPBLAS_OP_N, HIPBLAS_OP_N, xfeat, batch_size, yfeat, &one, weight, xfeat, dy_gemm,\n yfeat, &zero, dx, xfeat);\n\n if (cublas_status != HIPBLAS_STATUS_SUCCESS) {\n printf(\"GEMM dgrad failed with %d\\n\", cublas_status);\n return 1;\n }\n }\n\n // Call GEMM wgrad\n cublas_status = mlp_gemm(handle, HIPBLAS_OP_N, HIPBLAS_OP_T, xfeat, yfeat, batch_size, &one, x, xfeat, dy_gemm, yfeat,\n &zero, dweight, xfeat);\n\n if (cublas_status != HIPBLAS_STATUS_SUCCESS) {\n printf(\"GEMM wgrad failed with %d\\n\", cublas_status);\n return 1;\n }\n }\n\n return 0;\n}\n\n// Instantiate for floating point types\ntemplate int mlp_fp(float* X, int input_features, int batch_size, float** WPtr, int num_layers,\n int* output_features, float** BPtr, float* Y, float* reserved_space, int use_bias,\n int activation, void* lt_workspace);\n\ntemplate int mlp_bp(float* X, float* Y, int input_features, int batch_size, float** WPtr, int num_layers,\n int* output_features, float* dY, float* reserved_space, float* work_space, float* dX,\n float** dwPtr, float** dbPtr, bool requires_grad, int use_bias, int activation);\n\ntemplate int mlp_fp(at::Half* X, int input_features, int batch_size, at::Half** WPtr, int num_layers,\n int* output_features, at::Half** BPtr, at::Half* Y, at::Half* reserved_space,\n int use_bias, int activation, void* lt_workspace);\n\ntemplate int mlp_bp(at::Half* X, at::Half* Y, int input_features, int batch_size, at::Half** WPtr,\n int num_layers, int* output_features, at::Half* dY, at::Half* reserved_space,\n at::Half* work_space, at::Half* dX, at::Half** dwPtr, at::Half** dbPtr,\n bool requires_grad, int use_bias, int activation);\n\ntemplate int mlp_fp(double* X, int input_features, int batch_size, double** WPtr, int num_layers,\n int* output_features, double** BPtr, double* Y, double* reserved_space, int use_bias,\n int activation, void* lt_workspace);\n\ntemplate int mlp_bp(double* X, double* Y, int input_features, int batch_size, double** WPtr, int num_layers,\n int* output_features, double* dY, double* reserved_space, double* work_space, double* dX,\n double** dwPtr, double** dbPtr, bool requires_grad, int use_bias, int activation);\n\ntemplate size_t get_mlp_bp_workspace_in_bytes(int batch_size, int num_layers, const int* output_features);\ntemplate size_t get_mlp_bp_workspace_in_bytes(int batch_size, int num_layers, const int* output_features);\ntemplate size_t get_mlp_bp_workspace_in_bytes(int batch_size, int num_layers, const int* output_features);\n", "source": "/kaggle/working/apex/csrc/mlp_cuda.cu"} {"input": "#include \n#include \n#include \n\n__global__ void update_scale_hysteresis_cuda_kernel(float* current_scale, int* growth_tracker, int* hysteresis_tracker,\n const float* found_inf, double growth_factor, double backoff_factor,\n int growth_interval, int hysteresis) {\n if (*found_inf > 0) {\n *hysteresis_tracker -= 1;\n\n // Only reset the growth tracker when hysteresis is larger than zero\n if (*hysteresis_tracker > 0) {\n *growth_tracker = 0;\n return;\n }\n }\n\n if (*found_inf) {\n *current_scale = (*current_scale) * backoff_factor;\n *growth_tracker = 0;\n } else {\n // Entering this branch means we just carried out a successful step,\n // so growth_tracker is incremented before comparing to growth_interval.\n auto successful = (*growth_tracker) + 1;\n if (successful == growth_interval) {\n auto new_scale = static_cast((*current_scale) * growth_factor);\n // Do not grow the scale past fp32 bounds to inf.\n if (isfinite(new_scale)) {\n *current_scale = new_scale;\n }\n *growth_tracker = 0;\n } else {\n *growth_tracker = successful;\n }\n }\n\n // Reset the hysteresis tracker if no infs are found\n if (*found_inf <= 0) {\n *hysteresis_tracker = hysteresis;\n }\n}\n\nat::Tensor update_scale_hysteresis_cuda(at::Tensor current_scale, at::Tensor growth_tracker,\n at::Tensor hysteresis_tracker, at::Tensor found_inf, const double growth_factor,\n const double backoff_factor, const int64_t growth_interval,\n const int hysteresis) {\n update_scale_hysteresis_cuda_kernel<<<1, 1, 0, at::cuda::getCurrentCUDAStream()>>>(\n current_scale.mutable_data_ptr(), growth_tracker.mutable_data_ptr(),\n hysteresis_tracker.mutable_data_ptr(), found_inf.const_data_ptr(), growth_factor, backoff_factor,\n growth_interval, hysteresis);\n\n AT_CUDA_CHECK(cudaGetLastError());\n\n return current_scale;\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n__global__ void update_scale_hysteresis_cuda_kernel(float* current_scale, int* growth_tracker, int* hysteresis_tracker,\n const float* found_inf, double growth_factor, double backoff_factor,\n int growth_interval, int hysteresis) {\n if (*found_inf > 0) {\n *hysteresis_tracker -= 1;\n\n // Only reset the growth tracker when hysteresis is larger than zero\n if (*hysteresis_tracker > 0) {\n *growth_tracker = 0;\n return;\n }\n }\n\n if (*found_inf) {\n *current_scale = (*current_scale) * backoff_factor;\n *growth_tracker = 0;\n } else {\n // Entering this branch means we just carried out a successful step,\n // so growth_tracker is incremented before comparing to growth_interval.\n auto successful = (*growth_tracker) + 1;\n if (successful == growth_interval) {\n auto new_scale = static_cast((*current_scale) * growth_factor);\n // Do not grow the scale past fp32 bounds to inf.\n if (isfinite(new_scale)) {\n *current_scale = new_scale;\n }\n *growth_tracker = 0;\n } else {\n *growth_tracker = successful;\n }\n }\n\n // Reset the hysteresis tracker if no infs are found\n if (*found_inf <= 0) {\n *hysteresis_tracker = hysteresis;\n }\n}\n\nat::Tensor update_scale_hysteresis_cuda(at::Tensor current_scale, at::Tensor growth_tracker,\n at::Tensor hysteresis_tracker, at::Tensor found_inf, const double growth_factor,\n const double backoff_factor, const int64_t growth_interval,\n const int hysteresis) {\n update_scale_hysteresis_cuda_kernel<<<1, 1, 0, at::cuda::getCurrentCUDAStream()>>>(\n current_scale.mutable_data_ptr(), growth_tracker.mutable_data_ptr(),\n hysteresis_tracker.mutable_data_ptr(), found_inf.const_data_ptr(), growth_factor, backoff_factor,\n growth_interval, hysteresis);\n\n AT_CUDA_CHECK(hipGetLastError());\n\n return current_scale;\n}\n", "source": "/kaggle/working/apex/csrc/update_scale_hysteresis.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n// Another possibility:\n// #include \n\n#include \n\n#include \"multi_tensor_apply.cuh\"\n#include \"type_shim.h\"\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntemplate \n__device__ __forceinline__ bool is_aligned(T* p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\n\ntemplate \nstruct L2NormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<1>& tl,\n float* output, float* output_per_tensor, bool per_tensor,\n int max_chunks_per_tensor) {\n if (*noop_gmem) {\n return;\n }\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t* x = (x_t*)tl.addresses[0][tensor_loc];\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]);\n vals[ii] += next * next;\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]);\n vals[ii] += next * next;\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val += vals[i];\n\n float final = reduce_block_into_lanes(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final)) *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] += final;\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;\n }\n }\n};\n\n__global__ void cleanup(float* output, float* output_per_tensor, float* ret, float* ret_per_tensor, bool per_tensor,\n int max_chunks_per_tensor, volatile int* noop_gmem) {\n if (*noop_gmem) {\n return;\n }\n __shared__ float vals[512];\n\n if (blockIdx.x == 0) {\n float val = 0;\n if (threadIdx.x < 320) val = output[threadIdx.x];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) *ret = sqrt(final);\n }\n\n if (per_tensor) {\n float* output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;\n\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x) val += output_this_tensor[i];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) ret_per_tensor[blockIdx.x] = sqrt(final);\n }\n}\n\nstd::tuple multi_tensor_l2norm_mp_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists,\n at::optional per_tensor_python) {\n bool per_tensor = per_tensor_python.has_value() ? per_tensor_python.value() : false;\n\n auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);\n auto output = at::zeros({320}, float_options);\n\n at::Tensor output_per_tensor;\n at::Tensor ret_per_tensor;\n\n int ntensors = tensor_lists[0].size();\n int max_chunks_per_tensor = -1;\n\n if (per_tensor) {\n for (int t = 0; t < ntensors; t++) {\n int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;\n if (max_chunks_this_tensor > max_chunks_per_tensor) max_chunks_per_tensor = max_chunks_this_tensor;\n }\n output_per_tensor = at::zeros({ntensors * max_chunks_per_tensor}, float_options);\n ret_per_tensor = at::empty({ntensors}, float_options);\n } else {\n ret_per_tensor = at::empty({0}, float_options);\n }\n\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n tensor_lists[0][0].scalar_type(), 0, \"multi_tensor_l2norm_mp_cuda\",\n multi_tensor_apply<1>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, L2NormFunctor(),\n output.data_ptr(), per_tensor ? output_per_tensor.data_ptr() : nullptr,\n per_tensor, max_chunks_per_tensor);)\n\n AT_CUDA_CHECK(cudaGetLastError());\n // AT_CUDA_CHECK(cudaDeviceSynchronize());\n\n // This involves one more small kernel launches, but will be negligible end to end.\n // I could get rid of these by hacking the functor + multi tensor harness with persistence\n // logic, but keeping it simple for now\n auto ret = at::empty({1}, output.options());\n const at::cuda::OptionalCUDAGuard device_guard(device_of(output));\n auto stream = at::cuda::getCurrentCUDAStream();\n cleanup<<>>(\n output.data_ptr(), per_tensor ? output_per_tensor.data_ptr() : nullptr, ret.data_ptr(),\n per_tensor ? ret_per_tensor.data_ptr() : nullptr, per_tensor, max_chunks_per_tensor,\n noop_flag.data_ptr());\n\n return std::tuple(ret, ret_per_tensor);\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n// Another possibility:\n// #include \n\n#include \n\n#include \"multi_tensor_apply.cuh\"\n#include \"type_shim.h\"\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntemplate \n__device__ __forceinline__ bool is_aligned(T* p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\n\ntemplate \nstruct L2NormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<1>& tl,\n float* output, float* output_per_tensor, bool per_tensor,\n int max_chunks_per_tensor) {\n if (*noop_gmem) {\n return;\n }\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t* x = (x_t*)tl.addresses[0][tensor_loc];\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]);\n vals[ii] += next * next;\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]);\n vals[ii] += next * next;\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val += vals[i];\n\n float final = reduce_block_into_lanes(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final)) *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] += final;\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;\n }\n }\n};\n\n__global__ void cleanup(float* output, float* output_per_tensor, float* ret, float* ret_per_tensor, bool per_tensor,\n int max_chunks_per_tensor, volatile int* noop_gmem) {\n if (*noop_gmem) {\n return;\n }\n __shared__ float vals[512];\n\n if (blockIdx.x == 0) {\n float val = 0;\n if (threadIdx.x < 320) val = output[threadIdx.x];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) *ret = sqrt(final);\n }\n\n if (per_tensor) {\n float* output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;\n\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x) val += output_this_tensor[i];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) ret_per_tensor[blockIdx.x] = sqrt(final);\n }\n}\n\nstd::tuple multi_tensor_l2norm_mp_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists,\n at::optional per_tensor_python) {\n bool per_tensor = per_tensor_python.has_value() ? per_tensor_python.value() : false;\n\n auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);\n auto output = at::zeros({320}, float_options);\n\n at::Tensor output_per_tensor;\n at::Tensor ret_per_tensor;\n\n int ntensors = tensor_lists[0].size();\n int max_chunks_per_tensor = -1;\n\n if (per_tensor) {\n for (int t = 0; t < ntensors; t++) {\n int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;\n if (max_chunks_this_tensor > max_chunks_per_tensor) max_chunks_per_tensor = max_chunks_this_tensor;\n }\n output_per_tensor = at::zeros({ntensors * max_chunks_per_tensor}, float_options);\n ret_per_tensor = at::empty({ntensors}, float_options);\n } else {\n ret_per_tensor = at::empty({0}, float_options);\n }\n\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n tensor_lists[0][0].scalar_type(), 0, \"multi_tensor_l2norm_mp_cuda\",\n multi_tensor_apply<1>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, L2NormFunctor(),\n output.data_ptr(), per_tensor ? output_per_tensor.data_ptr() : nullptr,\n per_tensor, max_chunks_per_tensor);)\n\n AT_CUDA_CHECK(hipGetLastError());\n // AT_CUDA_CHECK(cudaDeviceSynchronize());\n\n // This involves one more small kernel launches, but will be negligible end to end.\n // I could get rid of these by hacking the functor + multi tensor harness with persistence\n // logic, but keeping it simple for now\n auto ret = at::empty({1}, output.options());\n const at::cuda::OptionalCUDAGuard device_guard(device_of(output));\n auto stream = at::cuda::getCurrentCUDAStream();\n cleanup<<>>(\n output.data_ptr(), per_tensor ? output_per_tensor.data_ptr() : nullptr, ret.data_ptr(),\n per_tensor ? ret_per_tensor.data_ptr() : nullptr, per_tensor, max_chunks_per_tensor,\n noop_flag.data_ptr());\n\n return std::tuple(ret, ret_per_tensor);\n}\n", "source": "/kaggle/working/apex/csrc/multi_tensor_l2norm_kernel_mp.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n// Another possibility:\n// #include \n\n#include \n\n#include \"multi_tensor_apply.cuh\"\n#include \"type_shim.h\"\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntemplate \n__device__ __forceinline__ bool is_aligned(T* p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\n\ntemplate \nstruct L2NormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<1>& tl,\n float* output, float* output_per_tensor, bool per_tensor,\n int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t* x = (x_t*)tl.addresses[0][tensor_loc];\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]);\n vals[ii] += next * next;\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]);\n vals[ii] += next * next;\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val += vals[i];\n\n float final = reduce_block_into_lanes(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final)) *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] += final;\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;\n }\n }\n};\n\ntemplate \nstruct UnscaleL2NormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<1>& tl,\n const float* inv_scale, float* output, float* output_per_tensor,\n bool per_tensor, int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t* x = (x_t*)tl.addresses[0][tensor_loc];\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]) * (*inv_scale);\n vals[ii] += next * next;\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]) * (*inv_scale);\n vals[ii] += next * next;\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val += vals[i];\n\n float final = reduce_block_into_lanes(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final)) *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] += final;\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;\n }\n }\n};\n\n// Probably better to template, but since we are not likely to support other norm\ntemplate \nstruct MaxNormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<1>& tl,\n float* output, float* output_per_tensor, bool per_tensor,\n int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t* x = (x_t*)tl.addresses[0][tensor_loc];\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]);\n vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]);\n vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val = fmaxf(fabsf(val), fabsf(vals[i]));\n\n float final = reduce_block_into_lanes_max_op(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final)) *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] = fmaxf(fabsf(output[blockIdx.x]), fabsf(final));\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;\n }\n }\n};\n\n__global__ void cleanup(float* output, float* output_per_tensor, float* ret, float* ret_per_tensor, bool per_tensor,\n int max_chunks_per_tensor) {\n __shared__ float vals[512];\n\n if (blockIdx.x == 0) {\n float val = 0;\n if (threadIdx.x < 320) val = output[threadIdx.x];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) *ret = sqrt(final);\n }\n\n if (per_tensor) {\n float* output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;\n\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x) val += output_this_tensor[i];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) ret_per_tensor[blockIdx.x] = sqrt(final);\n }\n}\n\n__global__ void cleanup_v2(float* output, float* output_per_tensor, float* ret, float* ret_per_tensor, bool per_tensor,\n int max_chunks_per_tensor, int norm_type, float alpha, float beta) {\n __shared__ float vals[512];\n\n if (blockIdx.x == 0) {\n float val = 0;\n if (threadIdx.x < 320) val = output[threadIdx.x];\n\n if (norm_type == 0) {\n float final = reduce_block_into_lanes_max_op(vals, val);\n if (threadIdx.x == 0) *ret = alpha * (*ret) + beta * final;\n } else {\n float final = reduce_block_into_lanes(vals, val);\n if (threadIdx.x == 0) *ret = sqrt(alpha * (*ret) * (*ret) + beta * final);\n }\n }\n\n if (per_tensor) {\n float* output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;\n\n if (norm_type == 0) {\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)\n val = fmaxf(fabsf(val), fabsf(output_this_tensor[i]));\n\n float final = reduce_block_into_lanes_max_op(vals, val);\n\n if (threadIdx.x == 0) ret_per_tensor[blockIdx.x] = alpha * ret_per_tensor[blockIdx.x] + beta * final;\n } else {\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x) val += output_this_tensor[i];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0)\n ret_per_tensor[blockIdx.x] =\n sqrt(alpha * ret_per_tensor[blockIdx.x] * ret_per_tensor[blockIdx.x] + beta * final);\n }\n }\n}\n\nstd::tuple multi_tensor_l2norm_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists,\n at::optional per_tensor_python) {\n bool per_tensor = per_tensor_python.has_value() ? per_tensor_python.value() : false;\n\n auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);\n auto output = at::zeros({320}, float_options);\n\n at::Tensor output_per_tensor;\n at::Tensor ret_per_tensor;\n\n int ntensors = tensor_lists[0].size();\n int max_chunks_per_tensor = -1;\n\n if (per_tensor) {\n for (int t = 0; t < ntensors; t++) {\n int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;\n if (max_chunks_this_tensor > max_chunks_per_tensor) max_chunks_per_tensor = max_chunks_this_tensor;\n }\n output_per_tensor = at::zeros({ntensors * max_chunks_per_tensor}, float_options);\n ret_per_tensor = at::empty({ntensors}, float_options);\n } else {\n ret_per_tensor = at::empty({0}, float_options);\n }\n\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n tensor_lists[0][0].scalar_type(), 0, \"multi_tensor_l2norm_cuda\",\n multi_tensor_apply<1>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, L2NormFunctor(),\n output.data_ptr(), per_tensor ? output_per_tensor.data_ptr() : nullptr,\n per_tensor, max_chunks_per_tensor);)\n\n AT_CUDA_CHECK(cudaGetLastError());\n // AT_CUDA_CHECK(cudaDeviceSynchronize());\n\n // This involves one more small kernel launches, but will be negligible end to end.\n // I could get rid of these by hacking the functor + multi tensor harness with persistence\n // logic, but keeping it simple for now\n auto ret = at::empty({1}, output.options());\n const at::cuda::OptionalCUDAGuard device_guard(device_of(output));\n auto stream = at::cuda::getCurrentCUDAStream();\n cleanup<<>>(\n output.data_ptr(), per_tensor ? output_per_tensor.data_ptr() : nullptr, ret.data_ptr(),\n per_tensor ? ret_per_tensor.data_ptr() : nullptr, per_tensor, max_chunks_per_tensor);\n\n return std::tuple(ret, ret_per_tensor);\n}\n\nstd::tuple multi_tensor_unscale_l2norm_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists,\n at::Tensor inv_scale,\n at::optional per_tensor_python) {\n bool per_tensor = per_tensor_python.has_value() ? per_tensor_python.value() : false;\n\n auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);\n auto output = at::zeros({320}, float_options);\n\n at::Tensor output_per_tensor;\n at::Tensor ret_per_tensor;\n\n int ntensors = tensor_lists[0].size();\n int max_chunks_per_tensor = -1;\n\n if (per_tensor) {\n for (int t = 0; t < ntensors; t++) {\n int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;\n if (max_chunks_this_tensor > max_chunks_per_tensor) max_chunks_per_tensor = max_chunks_this_tensor;\n }\n output_per_tensor = at::zeros({ntensors * max_chunks_per_tensor}, float_options);\n ret_per_tensor = at::empty({ntensors}, float_options);\n } else {\n ret_per_tensor = at::empty({0}, float_options);\n }\n\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n tensor_lists[0][0].scalar_type(), 0, \"multi_tensor_unscale_l2norm_cuda\",\n multi_tensor_apply<1>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, UnscaleL2NormFunctor(),\n inv_scale.data_ptr(), output.data_ptr(),\n per_tensor ? output_per_tensor.data_ptr() : nullptr, per_tensor,\n max_chunks_per_tensor);)\n\n AT_CUDA_CHECK(cudaGetLastError());\n // AT_CUDA_CHECK(cudaDeviceSynchronize());\n\n // This involves one more small kernel launches, but will be negligible end to end.\n // I could get rid of these by hacking the functor + multi tensor harness with persistence\n // logic, but keeping it simple for now\n auto ret = at::empty({1}, output.options());\n const at::cuda::OptionalCUDAGuard device_guard(device_of(output));\n auto stream = at::cuda::getCurrentCUDAStream();\n cleanup<<>>(\n output.data_ptr(), per_tensor ? output_per_tensor.data_ptr() : nullptr, ret.data_ptr(),\n per_tensor ? ret_per_tensor.data_ptr() : nullptr, per_tensor, max_chunks_per_tensor);\n\n return std::tuple(ret, ret_per_tensor);\n}\n\n// Compute and update grad norm\n// Here use a per tensor norm, and blend new norm(n) and old norm(gn) by\n// L-2: gn = sqrt(a * gn^2 + b * n^2)\n// L-inf: gn = a * gn + b * n\nvoid multi_tensor_norm_out_cuda(int chunk_size, at::Tensor noop_flag, std::vector> tensor_lists,\n at::Tensor out, const float alpha, const float beta, const int norm_type) {\n auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);\n TORCH_CHECK(tensor_lists[0][0].device() == noop_flag.device(), \"noop flag should be on the same device as tensors\");\n // we don't need global thus uses empty here\n auto output = at::empty({320}, float_options);\n\n at::Tensor output_per_tensor;\n at::Tensor ret_per_tensor;\n\n int ntensors = tensor_lists[0].size();\n int max_chunks_per_tensor = -1;\n\n for (int t = 0; t < ntensors; t++) {\n int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;\n if (max_chunks_this_tensor > max_chunks_per_tensor) max_chunks_per_tensor = max_chunks_this_tensor;\n }\n\n // Although it is single write then read, still need to be zero\n // Since tailing element also participate cleanup\n output_per_tensor = at::zeros({ntensors * max_chunks_per_tensor}, float_options);\n\n if (norm_type == 0) {\n DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, \"multi_tensor_maxnorm_cuda\",\n multi_tensor_apply<1>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n MaxNormFunctor(), output.data_ptr(),\n output_per_tensor.data_ptr(), true, max_chunks_per_tensor);)\n } else {\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n tensor_lists[0][0].scalar_type(), 0, \"multi_tensor_l2norm_cuda\",\n multi_tensor_apply<1>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, L2NormFunctor(),\n output.data_ptr(), output_per_tensor.data_ptr(), true,\n max_chunks_per_tensor);)\n }\n AT_CUDA_CHECK(cudaGetLastError());\n\n // AT_CUDA_CHECK(cudaDeviceSynchronize());\n\n // This involves one more small kernel launches, but will be negligible end to end.\n // I could get rid of these by hacking the functor + multi tensor harness with persistence\n // logic, but keeping it simple for now\n auto ret = at::empty({1}, output.options());\n\n // Adding the following device guard since it happens sometimes that the\n // tensors are on one device and the cuda stream is on another device which\n // results in ILLEGAL MEM ACCESS error.\n const at::cuda::OptionalCUDAGuard device_guard(device_of(output));\n auto stream = at::cuda::getCurrentCUDAStream();\n cleanup_v2<<>>(output.data_ptr(), output_per_tensor.data_ptr(),\n ret.data_ptr(), out.data_ptr(), true, max_chunks_per_tensor,\n norm_type, alpha, beta);\n\n return;\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n// Another possibility:\n// #include \n\n#include \n\n#include \"multi_tensor_apply.cuh\"\n#include \"type_shim.h\"\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntemplate \n__device__ __forceinline__ bool is_aligned(T* p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\n\ntemplate \nstruct L2NormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<1>& tl,\n float* output, float* output_per_tensor, bool per_tensor,\n int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t* x = (x_t*)tl.addresses[0][tensor_loc];\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]);\n vals[ii] += next * next;\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]);\n vals[ii] += next * next;\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val += vals[i];\n\n float final = reduce_block_into_lanes(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final)) *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] += final;\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;\n }\n }\n};\n\ntemplate \nstruct UnscaleL2NormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<1>& tl,\n const float* inv_scale, float* output, float* output_per_tensor,\n bool per_tensor, int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t* x = (x_t*)tl.addresses[0][tensor_loc];\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]) * (*inv_scale);\n vals[ii] += next * next;\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]) * (*inv_scale);\n vals[ii] += next * next;\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val += vals[i];\n\n float final = reduce_block_into_lanes(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final)) *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] += final;\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;\n }\n }\n};\n\n// Probably better to template, but since we are not likely to support other norm\ntemplate \nstruct MaxNormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<1>& tl,\n float* output, float* output_per_tensor, bool per_tensor,\n int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t* x = (x_t*)tl.addresses[0][tensor_loc];\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]);\n vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]);\n vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val = fmaxf(fabsf(val), fabsf(vals[i]));\n\n float final = reduce_block_into_lanes_max_op(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final)) *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] = fmaxf(fabsf(output[blockIdx.x]), fabsf(final));\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;\n }\n }\n};\n\n__global__ void cleanup(float* output, float* output_per_tensor, float* ret, float* ret_per_tensor, bool per_tensor,\n int max_chunks_per_tensor) {\n __shared__ float vals[512];\n\n if (blockIdx.x == 0) {\n float val = 0;\n if (threadIdx.x < 320) val = output[threadIdx.x];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) *ret = sqrt(final);\n }\n\n if (per_tensor) {\n float* output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;\n\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x) val += output_this_tensor[i];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) ret_per_tensor[blockIdx.x] = sqrt(final);\n }\n}\n\n__global__ void cleanup_v2(float* output, float* output_per_tensor, float* ret, float* ret_per_tensor, bool per_tensor,\n int max_chunks_per_tensor, int norm_type, float alpha, float beta) {\n __shared__ float vals[512];\n\n if (blockIdx.x == 0) {\n float val = 0;\n if (threadIdx.x < 320) val = output[threadIdx.x];\n\n if (norm_type == 0) {\n float final = reduce_block_into_lanes_max_op(vals, val);\n if (threadIdx.x == 0) *ret = alpha * (*ret) + beta * final;\n } else {\n float final = reduce_block_into_lanes(vals, val);\n if (threadIdx.x == 0) *ret = sqrt(alpha * (*ret) * (*ret) + beta * final);\n }\n }\n\n if (per_tensor) {\n float* output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;\n\n if (norm_type == 0) {\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)\n val = fmaxf(fabsf(val), fabsf(output_this_tensor[i]));\n\n float final = reduce_block_into_lanes_max_op(vals, val);\n\n if (threadIdx.x == 0) ret_per_tensor[blockIdx.x] = alpha * ret_per_tensor[blockIdx.x] + beta * final;\n } else {\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x) val += output_this_tensor[i];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0)\n ret_per_tensor[blockIdx.x] =\n sqrt(alpha * ret_per_tensor[blockIdx.x] * ret_per_tensor[blockIdx.x] + beta * final);\n }\n }\n}\n\nstd::tuple multi_tensor_l2norm_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists,\n at::optional per_tensor_python) {\n bool per_tensor = per_tensor_python.has_value() ? per_tensor_python.value() : false;\n\n auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);\n auto output = at::zeros({320}, float_options);\n\n at::Tensor output_per_tensor;\n at::Tensor ret_per_tensor;\n\n int ntensors = tensor_lists[0].size();\n int max_chunks_per_tensor = -1;\n\n if (per_tensor) {\n for (int t = 0; t < ntensors; t++) {\n int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;\n if (max_chunks_this_tensor > max_chunks_per_tensor) max_chunks_per_tensor = max_chunks_this_tensor;\n }\n output_per_tensor = at::zeros({ntensors * max_chunks_per_tensor}, float_options);\n ret_per_tensor = at::empty({ntensors}, float_options);\n } else {\n ret_per_tensor = at::empty({0}, float_options);\n }\n\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n tensor_lists[0][0].scalar_type(), 0, \"multi_tensor_l2norm_cuda\",\n multi_tensor_apply<1>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, L2NormFunctor(),\n output.data_ptr(), per_tensor ? output_per_tensor.data_ptr() : nullptr,\n per_tensor, max_chunks_per_tensor);)\n\n AT_CUDA_CHECK(hipGetLastError());\n // AT_CUDA_CHECK(cudaDeviceSynchronize());\n\n // This involves one more small kernel launches, but will be negligible end to end.\n // I could get rid of these by hacking the functor + multi tensor harness with persistence\n // logic, but keeping it simple for now\n auto ret = at::empty({1}, output.options());\n const at::cuda::OptionalCUDAGuard device_guard(device_of(output));\n auto stream = at::cuda::getCurrentCUDAStream();\n cleanup<<>>(\n output.data_ptr(), per_tensor ? output_per_tensor.data_ptr() : nullptr, ret.data_ptr(),\n per_tensor ? ret_per_tensor.data_ptr() : nullptr, per_tensor, max_chunks_per_tensor);\n\n return std::tuple(ret, ret_per_tensor);\n}\n\nstd::tuple multi_tensor_unscale_l2norm_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists,\n at::Tensor inv_scale,\n at::optional per_tensor_python) {\n bool per_tensor = per_tensor_python.has_value() ? per_tensor_python.value() : false;\n\n auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);\n auto output = at::zeros({320}, float_options);\n\n at::Tensor output_per_tensor;\n at::Tensor ret_per_tensor;\n\n int ntensors = tensor_lists[0].size();\n int max_chunks_per_tensor = -1;\n\n if (per_tensor) {\n for (int t = 0; t < ntensors; t++) {\n int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;\n if (max_chunks_this_tensor > max_chunks_per_tensor) max_chunks_per_tensor = max_chunks_this_tensor;\n }\n output_per_tensor = at::zeros({ntensors * max_chunks_per_tensor}, float_options);\n ret_per_tensor = at::empty({ntensors}, float_options);\n } else {\n ret_per_tensor = at::empty({0}, float_options);\n }\n\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n tensor_lists[0][0].scalar_type(), 0, \"multi_tensor_unscale_l2norm_cuda\",\n multi_tensor_apply<1>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, UnscaleL2NormFunctor(),\n inv_scale.data_ptr(), output.data_ptr(),\n per_tensor ? output_per_tensor.data_ptr() : nullptr, per_tensor,\n max_chunks_per_tensor);)\n\n AT_CUDA_CHECK(hipGetLastError());\n // AT_CUDA_CHECK(cudaDeviceSynchronize());\n\n // This involves one more small kernel launches, but will be negligible end to end.\n // I could get rid of these by hacking the functor + multi tensor harness with persistence\n // logic, but keeping it simple for now\n auto ret = at::empty({1}, output.options());\n const at::cuda::OptionalCUDAGuard device_guard(device_of(output));\n auto stream = at::cuda::getCurrentCUDAStream();\n cleanup<<>>(\n output.data_ptr(), per_tensor ? output_per_tensor.data_ptr() : nullptr, ret.data_ptr(),\n per_tensor ? ret_per_tensor.data_ptr() : nullptr, per_tensor, max_chunks_per_tensor);\n\n return std::tuple(ret, ret_per_tensor);\n}\n\n// Compute and update grad norm\n// Here use a per tensor norm, and blend new norm(n) and old norm(gn) by\n// L-2: gn = sqrt(a * gn^2 + b * n^2)\n// L-inf: gn = a * gn + b * n\nvoid multi_tensor_norm_out_cuda(int chunk_size, at::Tensor noop_flag, std::vector> tensor_lists,\n at::Tensor out, const float alpha, const float beta, const int norm_type) {\n auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);\n TORCH_CHECK(tensor_lists[0][0].device() == noop_flag.device(), \"noop flag should be on the same device as tensors\");\n // we don't need global thus uses empty here\n auto output = at::empty({320}, float_options);\n\n at::Tensor output_per_tensor;\n at::Tensor ret_per_tensor;\n\n int ntensors = tensor_lists[0].size();\n int max_chunks_per_tensor = -1;\n\n for (int t = 0; t < ntensors; t++) {\n int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;\n if (max_chunks_this_tensor > max_chunks_per_tensor) max_chunks_per_tensor = max_chunks_this_tensor;\n }\n\n // Although it is single write then read, still need to be zero\n // Since tailing element also participate cleanup\n output_per_tensor = at::zeros({ntensors * max_chunks_per_tensor}, float_options);\n\n if (norm_type == 0) {\n DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, \"multi_tensor_maxnorm_cuda\",\n multi_tensor_apply<1>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n MaxNormFunctor(), output.data_ptr(),\n output_per_tensor.data_ptr(), true, max_chunks_per_tensor);)\n } else {\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n tensor_lists[0][0].scalar_type(), 0, \"multi_tensor_l2norm_cuda\",\n multi_tensor_apply<1>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, L2NormFunctor(),\n output.data_ptr(), output_per_tensor.data_ptr(), true,\n max_chunks_per_tensor);)\n }\n AT_CUDA_CHECK(hipGetLastError());\n\n // AT_CUDA_CHECK(cudaDeviceSynchronize());\n\n // This involves one more small kernel launches, but will be negligible end to end.\n // I could get rid of these by hacking the functor + multi tensor harness with persistence\n // logic, but keeping it simple for now\n auto ret = at::empty({1}, output.options());\n\n // Adding the following device guard since it happens sometimes that the\n // tensors are on one device and the cuda stream is on another device which\n // results in ILLEGAL MEM ACCESS error.\n const at::cuda::OptionalCUDAGuard device_guard(device_of(output));\n auto stream = at::cuda::getCurrentCUDAStream();\n cleanup_v2<<>>(output.data_ptr(), output_per_tensor.data_ptr(),\n ret.data_ptr(), out.data_ptr(), true, max_chunks_per_tensor,\n norm_type, alpha, beta);\n\n return;\n}\n", "source": "/kaggle/working/apex/csrc/multi_tensor_l2norm_kernel.cu"} {"input": "#include \n#include \n\n#include \n\n#include \"ATen/ATen.h\"\n#include \"ATen/AccumulateType.h\"\n#include \"ATen/cuda/CUDAContext.h\"\n#include \"ATen/cuda/DeviceUtils.cuh\"\n#include \"static_switch.h\"\n#include \"type_shim.h\"\n\ntemplate \n__device__ void cuWelfordOnlineSum(const U curr, U& mu, U& sigma2, U& count) {\n count = count + U(1);\n U delta = curr - mu;\n U lmean = mu + delta / count;\n mu = lmean;\n U delta2 = curr - lmean;\n sigma2 = sigma2 + delta * delta2;\n}\n\ntemplate \n__device__ void cuChanOnlineSum(const U muB, const U sigma2B, const U countB, U& mu, U& sigma2, U& count) {\n U delta = muB - mu;\n U nA = count;\n U nB = countB;\n count = count + countB;\n U nX = count;\n if (nX > U(0)) {\n nA = nA / nX;\n nB = nB / nX;\n mu = nA * mu + nB * muB;\n sigma2 = sigma2 + sigma2B + delta * delta * nA * nB * nX;\n } else {\n mu = U(0);\n sigma2 = U(0);\n }\n}\n\ntemplate \n__device__ void cuRMSOnlineSum(const U curr, U& sigma2) {\n sigma2 = sigma2 + curr * curr;\n}\n\ntemplate \n__device__ void cuChanRMSOnlineSum(const U sigma2B, U& sigma2) {\n sigma2 = sigma2 + sigma2B;\n}\n\ntemplate \n__device__ void cuWelfordMuSigma2(const T* __restrict__ vals, const int n1, const int n2, const int i1, U& mu,\n U& sigma2, U* buf, bool rms_only) {\n // Assumptions:\n // 1) blockDim.x == warpSize\n // 2) Tensor is contiguous\n // 3) 2*blockDim.y*sizeof(U)+blockDim.y*sizeof(int) shared memory available.\n //\n // compute variance and mean over n2\n U count = U(0);\n mu = U(0);\n sigma2 = U(0);\n if (i1 < n1) {\n // one warp normalizes one n1 index,\n // synchronization is implicit\n // initialize with standard Welford algorithm\n const int numx = blockDim.x * blockDim.y;\n const int thrx = threadIdx.x + threadIdx.y * blockDim.x;\n const T* lvals = vals + i1 * n2;\n int l = 4 * thrx;\n for (; l + 3 < n2; l += 4 * numx) {\n for (int k = 0; k < 4; ++k) {\n U curr = static_cast(lvals[l + k]);\n if (!rms_only) {\n cuWelfordOnlineSum(curr, mu, sigma2, count);\n } else {\n cuRMSOnlineSum(curr, sigma2);\n }\n }\n }\n for (; l < n2; ++l) {\n U curr = static_cast(lvals[l]);\n if (!rms_only) {\n cuWelfordOnlineSum(curr, mu, sigma2, count);\n } else {\n cuRMSOnlineSum(curr, sigma2);\n }\n }\n // intra-warp reductions\n for (int l = 0; l <= 4; ++l) {\n int srcLaneB = (threadIdx.x + (1 << l)) & 31;\n U sigma2B = WARP_SHFL(sigma2, srcLaneB);\n if (!rms_only) {\n U muB = WARP_SHFL(mu, srcLaneB);\n U countB = WARP_SHFL(count, srcLaneB);\n cuChanOnlineSum(muB, sigma2B, countB, mu, sigma2, count);\n } else {\n cuChanRMSOnlineSum(sigma2B, sigma2);\n }\n }\n // threadIdx.x == 0 has correct values for each warp\n // inter-warp reductions\n if (blockDim.y > 1) {\n U* ubuf = (U*)buf;\n U* ibuf = (U*)(ubuf + blockDim.y);\n for (int offset = blockDim.y / 2; offset > 0; offset /= 2) {\n // upper half of warps write to shared\n if (threadIdx.x == 0 && threadIdx.y >= offset && threadIdx.y < 2 * offset) {\n const int wrt_y = threadIdx.y - offset;\n if (!rms_only) {\n ubuf[2 * wrt_y] = mu;\n ibuf[wrt_y] = count;\n }\n ubuf[2 * wrt_y + 1] = sigma2;\n }\n __syncthreads();\n // lower half merges\n if (threadIdx.x == 0 && threadIdx.y < offset) {\n U sigma2B = ubuf[2 * threadIdx.y + 1];\n if (!rms_only) {\n U muB = ubuf[2 * threadIdx.y];\n U countB = ibuf[threadIdx.y];\n cuChanOnlineSum(muB, sigma2B, countB, mu, sigma2, count);\n } else {\n cuChanRMSOnlineSum(sigma2B, sigma2);\n }\n }\n __syncthreads();\n }\n // threadIdx.x = 0 && threadIdx.y == 0 only thread that has correct values\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n if (!rms_only) {\n ubuf[0] = mu;\n }\n ubuf[1] = sigma2;\n }\n __syncthreads();\n if (!rms_only) {\n mu = ubuf[0];\n }\n sigma2 = ubuf[1] / U(n2);\n // don't care about final value of count, we know count == n2\n } else {\n if (!rms_only) {\n mu = WARP_SHFL(mu, 0);\n }\n sigma2 = WARP_SHFL(sigma2 / U(n2), 0);\n }\n }\n}\n\ntemplate <>\n__device__ void cuWelfordMuSigma2(const at::Half* __restrict__ vals, const int n1, const int n2, const int i1,\n float& mu, float& sigma2, float* buf, bool rms_only) {\n // Assumptions:\n // 1) blockDim.x == warpSize\n // 2) Tensor is contiguous\n // 3) 2*blockDim.y*sizeof(U)+blockDim.y*sizeof(int) shared memory available.\n //\n // compute variance and mean over n2\n float count = 0.0f;\n mu = float(0);\n sigma2 = float(0);\n if (i1 < n1) {\n // one warp normalizes one n1 index,\n // synchronization is implicit\n // initialize with standard Welford algorithm\n const int numx = blockDim.x * blockDim.y;\n const int thrx = threadIdx.x + threadIdx.y * blockDim.x;\n const at::Half* lvals = vals + i1 * n2;\n int l = 8 * thrx;\n if ((((size_t)lvals) & 3) != 0) {\n // 16 bit alignment\n // first thread consumes first point\n if (thrx == 0) {\n float curr = static_cast(lvals[0]);\n if (!rms_only) {\n cuWelfordOnlineSum(curr, mu, sigma2, count);\n } else {\n cuRMSOnlineSum(curr, sigma2);\n }\n }\n ++l;\n }\n // at this point, lvals[l] are 32 bit aligned for all threads.\n for (; l + 7 < n2; l += 8 * numx) {\n for (int k = 0; k < 8; k += 2) {\n float2 curr = __half22float2(*((__half2*)(lvals + l + k)));\n if (!rms_only) {\n cuWelfordOnlineSum(curr.x, mu, sigma2, count);\n cuWelfordOnlineSum(curr.y, mu, sigma2, count);\n } else {\n cuRMSOnlineSum(curr.x, sigma2);\n cuRMSOnlineSum(curr.y, sigma2);\n }\n }\n }\n for (; l < n2; ++l) {\n float curr = static_cast(lvals[l]);\n if (!rms_only) {\n cuWelfordOnlineSum(curr, mu, sigma2, count);\n } else {\n cuRMSOnlineSum(curr, sigma2);\n }\n }\n // intra-warp reductions\n for (int l = 0; l <= 4; ++l) {\n int srcLaneB = (threadIdx.x + (1 << l)) & 31;\n float sigma2B = WARP_SHFL(sigma2, srcLaneB);\n if (!rms_only) {\n float muB = WARP_SHFL(mu, srcLaneB);\n float countB = WARP_SHFL(count, srcLaneB);\n cuChanOnlineSum(muB, sigma2B, countB, mu, sigma2, count);\n } else {\n cuChanRMSOnlineSum(sigma2B, sigma2);\n }\n }\n // threadIdx.x == 0 has correct values for each warp\n // inter-warp reductions\n if (blockDim.y > 1) {\n float* ubuf = (float*)buf;\n float* ibuf = (float*)(ubuf + blockDim.y);\n for (int offset = blockDim.y / 2; offset > 0; offset /= 2) {\n // upper half of warps write to shared\n if (threadIdx.x == 0 && threadIdx.y >= offset && threadIdx.y < 2 * offset) {\n const int wrt_y = threadIdx.y - offset;\n ubuf[2 * wrt_y + 1] = sigma2;\n if (!rms_only) {\n ubuf[2 * wrt_y] = mu;\n ibuf[wrt_y] = count;\n }\n }\n __syncthreads();\n // lower half merges\n if (threadIdx.x == 0 && threadIdx.y < offset) {\n float sigma2B = ubuf[2 * threadIdx.y + 1];\n if (!rms_only) {\n float muB = ubuf[2 * threadIdx.y];\n float countB = ibuf[threadIdx.y];\n cuChanOnlineSum(muB, sigma2B, countB, mu, sigma2, count);\n } else {\n cuChanRMSOnlineSum(sigma2B, sigma2);\n }\n }\n __syncthreads();\n }\n // threadIdx.x = 0 && threadIdx.y == 0 only thread that has correct values\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n if (!rms_only) {\n ubuf[0] = mu;\n }\n ubuf[1] = sigma2;\n }\n __syncthreads();\n if (!rms_only) {\n mu = ubuf[0];\n }\n sigma2 = ubuf[1] / float(n2);\n // don't care about final value of count, we know count == n2\n } else {\n if (!rms_only) {\n mu = WARP_SHFL(mu, 0);\n }\n sigma2 = WARP_SHFL(sigma2 / float(n2), 0);\n }\n }\n}\n\ntemplate \nU rsqrt(U v) {\n return U(1) / sqrt(v);\n}\ntemplate <>\nfloat rsqrt(float v) {\n return rsqrtf(v);\n}\ntemplate <>\ndouble rsqrt(double v) {\n return rsqrt(v);\n}\n\nnamespace {\n// This is the un-specialized struct. Note that we prevent instantiation of this\n// struct by putting an undefined symbol in the function body so it won't compile.\n// template \n// struct SharedMemory\n// {\n// // Ensure that we won't compile any un-specialized types\n// __device__ T *getPointer()\n// {\n// extern __device__ void error(void);\n// error();\n// return nullptr;\n// }\n// };\n// https://github.com/NVIDIA/apex/issues/246\ntemplate \nstruct SharedMemory;\n\ntemplate <>\nstruct SharedMemory {\n __device__ float* getPointer() {\n extern __shared__ float s_float[];\n return s_float;\n }\n};\n\ntemplate <>\nstruct SharedMemory {\n __device__ double* getPointer() {\n extern __shared__ double s_double[];\n return s_double;\n }\n};\n} // namespace\n\ntemplate \n__device__ void cuApplyLayerNorm_(V* __restrict__ output_vals, U* __restrict__ mean, U* __restrict__ invvar,\n const T* __restrict__ vals, const int n1, const int n2, const U epsilon,\n const V* __restrict__ gamma, const V* __restrict__ beta, bool rms_only) {\n // Assumptions:\n // 1) blockDim.x == warpSize\n // 2) Tensors are contiguous\n //\n for (auto i1 = blockIdx.y; i1 < n1; i1 += gridDim.y) {\n SharedMemory shared;\n U* buf = shared.getPointer();\n U mu, sigma2;\n cuWelfordMuSigma2(vals, n1, n2, i1, mu, sigma2, buf, rms_only);\n\n const T* lvals = vals + i1 * n2;\n V* ovals = output_vals + i1 * n2;\n U c_invvar = rsqrt(sigma2 + epsilon);\n const int numx = blockDim.x * blockDim.y;\n const int thrx = threadIdx.x + threadIdx.y * blockDim.x;\n if (gamma != nullptr && (beta != nullptr || rms_only)) {\n for (int i = thrx; i < n2; i += numx) {\n U curr = static_cast(lvals[i]);\n if (!rms_only) {\n ovals[i] = gamma[i] * static_cast(c_invvar * (curr - mu)) + beta[i];\n } else {\n ovals[i] = gamma[i] * static_cast(c_invvar * curr);\n }\n }\n } else {\n for (int i = thrx; i < n2; i += numx) {\n U curr = static_cast(lvals[i]);\n if (!rms_only) {\n ovals[i] = static_cast(c_invvar * (curr - mu));\n } else {\n ovals[i] = static_cast(c_invvar * curr);\n }\n }\n }\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n if (!rms_only) {\n mean[i1] = mu;\n }\n invvar[i1] = c_invvar;\n }\n __syncthreads();\n }\n}\n\ntemplate \n__global__ void cuApplyLayerNorm(V* __restrict__ output_vals, U* __restrict__ mean, U* __restrict__ invvar,\n const T* __restrict__ vals, const int n1, const int n2, const U epsilon,\n const V* __restrict__ gamma, const V* __restrict__ beta) {\n cuApplyLayerNorm_(output_vals, mean, invvar, vals, n1, n2, epsilon, gamma, beta, false);\n}\n\ntemplate \n__global__ void cuApplyRMSNorm(V* __restrict__ output_vals, U* __restrict__ invvar, const T* __restrict__ vals,\n const int n1, const int n2, const U epsilon, const V* __restrict__ gamma) {\n cuApplyLayerNorm_(output_vals, nullptr, invvar, vals, n1, n2, epsilon, gamma, nullptr, true);\n}\n\ntemplate \n__device__ V clamp_by_magnitude(V curr_gamma, double eps) {\n const V kMinGamma = V(eps);\n if (curr_gamma >= 0) {\n if (curr_gamma < kMinGamma) {\n return kMinGamma;\n } else {\n return curr_gamma;\n }\n } else {\n if (curr_gamma > -kMinGamma) {\n return -kMinGamma;\n } else {\n return curr_gamma;\n }\n }\n}\n\ntemplate \n__device__ void cuLoadWriteStridedInputs(const int i1_block, const int thr_load_row_off, const int thr_load_col_off,\n const int i2_off, const int row_stride, U* warp_buf1, U* warp_buf2,\n const T* input_or_output, const V* dout, const int i1_end, const int n2,\n const U* __restrict__ mean, const U* __restrict__ invvar,\n const V* __restrict__ gamma, const V* __restrict__ beta, const double eps,\n bool rms_only) {\n int i1 = i1_block + thr_load_row_off;\n if (i1 < i1_end) {\n for (int k = 0; k < blockDim.y; ++k) {\n int i2 = i2_off + k;\n int load_idx = i1 * n2 + i2;\n int write_idx = thr_load_row_off * row_stride + thr_load_col_off + k;\n if (i2 < n2) {\n U c_h = static_cast(input_or_output[load_idx]);\n U curr_dout = static_cast(dout[load_idx]);\n if (!rms_only) {\n warp_buf1[write_idx] = curr_dout;\n if (MemoryEfficient) {\n U curr_beta = static_cast(beta[i2]);\n warp_buf2[write_idx] = curr_dout * (c_h - curr_beta) / static_cast(clamp_by_magnitude(gamma[i2], eps));\n } else {\n warp_buf2[write_idx] = curr_dout * (c_h - mean[i1]) * invvar[i1];\n }\n } else {\n if (MemoryEfficient) {\n warp_buf2[write_idx] = curr_dout * (c_h) / static_cast(clamp_by_magnitude(gamma[i2], eps));\n } else {\n warp_buf2[write_idx] = curr_dout * (c_h)*invvar[i1];\n }\n }\n } else {\n if (!rms_only) {\n warp_buf1[write_idx] = U(0);\n }\n warp_buf2[write_idx] = U(0);\n }\n }\n } else {\n for (int k = 0; k < blockDim.y; ++k) {\n int write_idx = thr_load_row_off * row_stride + thr_load_col_off + k;\n if (!rms_only) {\n warp_buf1[write_idx] = U(0);\n }\n warp_buf2[write_idx] = U(0);\n }\n }\n}\n\ntemplate \n__device__ void cuLoadAddStridedInputs(const int i1_block, const int thr_load_row_off, const int thr_load_col_off,\n const int i2_off, const int row_stride, U* warp_buf1, U* warp_buf2,\n const T* input_or_output, const V* dout, const int i1_end, const int n2,\n const U* __restrict__ mean, const U* __restrict__ invvar,\n const V* __restrict__ gamma, const V* __restrict__ beta, const double eps,\n bool rms_only) {\n int i1 = i1_block + thr_load_row_off;\n if (i1 < i1_end) {\n for (int k = 0; k < blockDim.y; ++k) {\n int i2 = i2_off + k;\n int load_idx = i1 * n2 + i2;\n int write_idx = thr_load_row_off * row_stride + thr_load_col_off + k;\n if (i2 < n2) {\n U c_h = static_cast(input_or_output[load_idx]);\n U curr_dout = static_cast(dout[load_idx]);\n if (!rms_only) {\n U curr_beta = static_cast(beta[i2]);\n warp_buf1[write_idx] += curr_dout;\n if (MemoryEfficient) {\n warp_buf2[write_idx] += curr_dout * (c_h - curr_beta) / static_cast(clamp_by_magnitude(gamma[i2], eps));\n } else {\n warp_buf2[write_idx] += curr_dout * (c_h - mean[i1]) * invvar[i1];\n }\n } else {\n if (MemoryEfficient) {\n warp_buf2[write_idx] += curr_dout * (c_h) / static_cast(clamp_by_magnitude(gamma[i2], eps));\n } else {\n warp_buf2[write_idx] += curr_dout * (c_h)*invvar[i1];\n }\n }\n }\n }\n }\n}\n\ntemplate \n__global__ void cuComputePartGradGammaBeta(const V* __restrict__ dout, const T* __restrict__ input_or_output,\n const int n1, const int n2, const U* __restrict__ mean,\n const U* __restrict__ invvar, U epsilon, const V* __restrict__ gamma,\n const V* __restrict__ beta, U* part_grad_gamma, U* part_grad_beta,\n const double eps, bool rms_only) {\n const int numsegs_n1 = (n1 + blockDim.y * blockDim.y - 1) / (blockDim.y * blockDim.y);\n const int segs_per_block = (numsegs_n1 + gridDim.y - 1) / gridDim.y;\n const int i1_beg = blockIdx.y * segs_per_block * blockDim.y * blockDim.y;\n const int i1_beg_plus_one = (blockIdx.y + 1) * segs_per_block * blockDim.y * blockDim.y;\n const int i1_end = i1_beg_plus_one < n1 ? i1_beg_plus_one : n1;\n const int row_stride = blockDim.x + 1;\n const int thr_load_col_off = (threadIdx.x * blockDim.y) & (blockDim.x - 1);\n const int thr_load_row_off = (threadIdx.x * blockDim.y) / blockDim.x + threadIdx.y * blockDim.y;\n const int i2_off = blockIdx.x * blockDim.x + thr_load_col_off;\n SharedMemory shared;\n U* buf = shared.getPointer(); // buf has at least blockDim.x * blockDim.y * blockDim.y + (blockDim.y -\n // 1)*(blockDim.x/blockDim.y) elements\n U* warp_buf1 = (U*)buf;\n U* warp_buf2 = warp_buf1 + blockDim.y * blockDim.y * row_stride;\n // compute partial sums from strided inputs\n // do this to increase number of loads in flight\n cuLoadWriteStridedInputs(i1_beg, thr_load_row_off, thr_load_col_off, i2_off, row_stride,\n warp_buf1, warp_buf2, input_or_output, dout, i1_end, n2, mean,\n invvar, gamma, beta, eps, rms_only);\n for (int i1_block = i1_beg + blockDim.y * blockDim.y; i1_block < i1_end; i1_block += blockDim.y * blockDim.y) {\n cuLoadAddStridedInputs(i1_block, thr_load_row_off, thr_load_col_off, i2_off, row_stride,\n warp_buf1, warp_buf2, input_or_output, dout, i1_end, n2, mean,\n invvar, gamma, beta, eps, rms_only);\n }\n __syncthreads();\n // inter-warp reductions\n // sum within each warp\n U acc1 = U(0);\n U acc2 = U(0);\n for (int k = 0; k < blockDim.y; ++k) {\n int row1 = threadIdx.y + k * blockDim.y;\n int idx1 = row1 * row_stride + threadIdx.x;\n if (!rms_only) {\n acc1 += warp_buf1[idx1];\n }\n acc2 += warp_buf2[idx1];\n }\n if (!rms_only) {\n warp_buf1[threadIdx.y * row_stride + threadIdx.x] = acc1;\n }\n warp_buf2[threadIdx.y * row_stride + threadIdx.x] = acc2;\n __syncthreads();\n // sum all warps\n for (int offset = blockDim.y / 2; offset > 1; offset /= 2) {\n if (threadIdx.y < offset) {\n int row1 = threadIdx.y;\n int row2 = threadIdx.y + offset;\n int idx1 = row1 * row_stride + threadIdx.x;\n int idx2 = row2 * row_stride + threadIdx.x;\n if (!rms_only) {\n warp_buf1[idx1] += warp_buf1[idx2];\n }\n warp_buf2[idx1] += warp_buf2[idx2];\n }\n __syncthreads();\n }\n int i2 = blockIdx.x * blockDim.x + threadIdx.x;\n if (threadIdx.y == 0 && i2 < n2) {\n int row1 = threadIdx.y;\n int row2 = threadIdx.y + 1;\n int idx1 = row1 * row_stride + threadIdx.x;\n int idx2 = row2 * row_stride + threadIdx.x;\n if (!rms_only) {\n part_grad_beta[blockIdx.y * n2 + i2] = warp_buf1[idx1] + warp_buf1[idx2];\n }\n part_grad_gamma[blockIdx.y * n2 + i2] = warp_buf2[idx1] + warp_buf2[idx2];\n }\n}\n\ntemplate \n__global__ void cuComputeGradGammaBeta(const U* part_grad_gamma, const U* part_grad_beta, const int part_size,\n const int n1, const int n2, V* grad_gamma, V* grad_beta, bool rms_only) {\n // sum partial gradients for gamma and beta\n SharedMemory shared;\n U* buf = shared.getPointer();\n int i2 = blockIdx.x * blockDim.x + threadIdx.x;\n if (i2 < n2) {\n // each warp does sequential reductions until reduced part_size is num_warps\n int num_warp_reductions = part_size / blockDim.y;\n U sum_gamma = U(0);\n U sum_beta = U(0);\n const U* part_grad_gamma_ptr = part_grad_gamma + threadIdx.y * num_warp_reductions * n2 + i2;\n const U* part_grad_beta_ptr = part_grad_beta + threadIdx.y * num_warp_reductions * n2 + i2;\n for (int warp_offset = 0; warp_offset < num_warp_reductions; ++warp_offset) {\n sum_gamma += part_grad_gamma_ptr[warp_offset * n2];\n if (!rms_only) {\n sum_beta += part_grad_beta_ptr[warp_offset * n2];\n }\n }\n // inter-warp reductions\n const int nbsize3 = blockDim.x * blockDim.y / 2;\n for (int offset = blockDim.y / 2; offset >= 1; offset /= 2) {\n // top half write to shared memory\n if (threadIdx.y >= offset && threadIdx.y < 2 * offset) {\n const int write_idx = (threadIdx.y - offset) * blockDim.x + threadIdx.x;\n buf[write_idx] = sum_gamma;\n if (!rms_only) {\n buf[write_idx + nbsize3] = sum_beta;\n }\n }\n __syncthreads();\n // bottom half sums\n if (threadIdx.y < offset) {\n const int read_idx = threadIdx.y * blockDim.x + threadIdx.x;\n sum_gamma += buf[read_idx];\n if (!rms_only) {\n sum_beta += buf[read_idx + nbsize3];\n }\n }\n __syncthreads();\n }\n // write out fully summed gradients\n if (threadIdx.y == 0) {\n grad_gamma[i2] = sum_gamma;\n if (!rms_only) {\n grad_beta[i2] = sum_beta;\n }\n }\n }\n}\n\ntemplate \n__global__ void cuComputeGradInput(const V* __restrict__ dout, const T* __restrict__ input_or_output, const int n1,\n const int n2, const U* __restrict__ mean, const U* __restrict__ invvar, U epsilon,\n const V* gamma, const V* beta, T* grad_input, const double eps, bool rms_only) {\n for (auto i1 = blockIdx.y; i1 < n1; i1 += gridDim.y) {\n U sum_loss1 = U(0);\n U sum_loss2 = U(0);\n const T* k_h = input_or_output + i1 * n2;\n const V* k_dout = dout + i1 * n2;\n const U c_invvar = invvar[i1];\n const U c_mean = !MemoryEfficient ? mean[i1] : 0.;\n const int numx = blockDim.x * blockDim.y;\n const int thrx = threadIdx.x + threadIdx.y * blockDim.x;\n if (gamma != nullptr) {\n int l = 4 * thrx;\n for (; l + 3 < n2; l += 4 * numx) {\n for (int k = 0; k < 4; ++k) {\n const U c_h = static_cast(k_h[l + k]);\n const U c_loss = static_cast(k_dout[l + k]);\n if (!rms_only) {\n sum_loss1 += c_loss * gamma[l + k];\n if (MemoryEfficient) {\n sum_loss2 += c_loss * (c_h - beta[l + k]);\n } else {\n sum_loss2 += c_loss * gamma[l + k] * (c_h - c_mean) * c_invvar;\n }\n } else {\n if (MemoryEfficient) {\n sum_loss2 += c_loss * c_h;\n } else {\n sum_loss2 += c_loss * gamma[l + k] * (c_h)*c_invvar;\n }\n }\n }\n }\n for (; l < n2; ++l) {\n const U c_h = static_cast(k_h[l]);\n const U c_loss = static_cast(k_dout[l]);\n if (!rms_only) {\n sum_loss1 += c_loss * gamma[l];\n if (MemoryEfficient) {\n sum_loss2 += c_loss * (c_h - beta[l]);\n } else {\n sum_loss2 += c_loss * gamma[l] * (c_h - c_mean) * c_invvar;\n }\n } else {\n if (MemoryEfficient) {\n sum_loss2 += c_loss * c_h;\n } else {\n sum_loss2 += c_loss * gamma[l] * (c_h)*c_invvar;\n }\n }\n }\n } else {\n int l = 4 * thrx;\n for (; l + 3 < n2; l += 4 * numx) {\n for (int k = 0; k < 4; ++k) {\n const U c_h = static_cast(k_h[l + k]);\n const U c_loss = static_cast(k_dout[l + k]);\n if (!rms_only) {\n sum_loss1 += c_loss;\n if (MemoryEfficient) {\n sum_loss2 += c_loss * c_h;\n } else {\n sum_loss2 += c_loss * (c_h - c_mean) * c_invvar;\n }\n } else {\n if (MemoryEfficient) {\n sum_loss2 += c_loss * c_h;\n } else {\n sum_loss2 += c_loss * (c_h)*c_invvar;\n }\n }\n }\n }\n for (; l < n2; ++l) {\n const U c_h = static_cast(k_h[l]);\n const U c_loss = static_cast(k_dout[l]);\n if (!rms_only) {\n sum_loss1 += c_loss;\n if (MemoryEfficient) {\n sum_loss2 += c_loss * c_h;\n } else {\n sum_loss2 += c_loss * (c_h - c_mean) * c_invvar;\n }\n } else {\n if (MemoryEfficient) {\n sum_loss2 += c_loss * c_h;\n } else {\n sum_loss2 += c_loss * (c_h)*c_invvar;\n }\n }\n }\n }\n // intra-warp reductions\n for (int mask = blockDim.x / 2; mask > 0; mask /= 2) {\n if (!rms_only) {\n sum_loss1 += WARP_SHFL_XOR(sum_loss1, mask);\n }\n sum_loss2 += WARP_SHFL_XOR(sum_loss2, mask);\n }\n // inter-warp reductions\n if (blockDim.y > 1) {\n SharedMemory shared;\n U* buf = shared.getPointer();\n for (int offset = blockDim.y / 2; offset > 0; offset /= 2) {\n // upper half of warps write to shared\n if (threadIdx.y >= offset && threadIdx.y < 2 * offset) {\n const int wrt_i = (threadIdx.y - offset) * blockDim.x + threadIdx.x;\n if (!rms_only) {\n buf[2 * wrt_i] = sum_loss1;\n }\n buf[2 * wrt_i + 1] = sum_loss2;\n }\n __syncthreads();\n // lower half merges\n if (threadIdx.y < offset) {\n const int read_i = threadIdx.y * blockDim.x + threadIdx.x;\n if (!rms_only) {\n sum_loss1 += buf[2 * read_i];\n }\n sum_loss2 += buf[2 * read_i + 1];\n }\n __syncthreads();\n }\n if (threadIdx.y == 0) {\n if (!rms_only) {\n buf[2 * threadIdx.x] = sum_loss1;\n }\n buf[2 * threadIdx.x + 1] = sum_loss2;\n }\n __syncthreads();\n if (threadIdx.y != 0) {\n if (!rms_only) {\n sum_loss1 = buf[2 * threadIdx.x];\n }\n sum_loss2 = buf[2 * threadIdx.x + 1];\n }\n }\n // all threads now have the two sums over l\n U fH = (U)n2;\n U term1 = (U(1) / fH) * c_invvar;\n T* k_grad_input = grad_input + i1 * n2;\n if (gamma != nullptr) {\n for (int l = thrx; l < n2; l += numx) {\n const U c_h = static_cast(k_h[l]);\n const U c_loss = static_cast(k_dout[l]);\n const U k_gamma = static_cast(clamp_by_magnitude(gamma[l], eps));\n U f_grad_input = fH * c_loss * k_gamma;\n if (!rms_only) {\n const U k_beta = beta[l];\n f_grad_input -= sum_loss1;\n if (MemoryEfficient) {\n f_grad_input -= (c_h - k_beta) / k_gamma * sum_loss2;\n } else {\n f_grad_input -= (c_h - c_mean) * c_invvar * sum_loss2;\n }\n } else {\n if (MemoryEfficient) {\n f_grad_input -= c_h / k_gamma * sum_loss2;\n } else {\n f_grad_input -= c_h * c_invvar * sum_loss2;\n }\n }\n f_grad_input *= term1;\n k_grad_input[l] = static_cast(f_grad_input);\n }\n } else {\n for (int l = thrx; l < n2; l += numx) {\n const U c_h = static_cast(k_h[l]);\n const U c_loss = static_cast(k_dout[l]);\n U f_grad_input = fH * c_loss;\n if (!rms_only) {\n f_grad_input -= sum_loss1;\n if (MemoryEfficient) {\n f_grad_input -= c_h * sum_loss2;\n } else {\n f_grad_input -= (c_h - c_mean) * c_invvar * sum_loss2;\n }\n } else {\n if (MemoryEfficient) {\n f_grad_input -= c_h * sum_loss2;\n } else {\n f_grad_input -= c_h * c_invvar * sum_loss2;\n }\n }\n f_grad_input *= term1;\n k_grad_input[l] = static_cast(f_grad_input);\n }\n }\n // prevent race where buf is written again before reads are done\n __syncthreads();\n }\n}\n\ntemplate \nvoid HostApplyLayerNorm(V* output, U* mean, U* invvar, const T* input, int n1, int n2, double epsilon, const V* gamma,\n const V* beta) {\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n const dim3 threads(32, 4, 1);\n const uint64_t maxGridY = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];\n const dim3 blocks(1, std::min((uint64_t)n1, maxGridY), 1);\n int nshared = threads.y > 1 ? threads.y * sizeof(U) + (threads.y / 2) * sizeof(U) : 0;\n cuApplyLayerNorm<<>>(output, mean, invvar, input, n1, n2, U(epsilon), gamma, beta);\n}\n\ntemplate \nvoid HostApplyRMSNorm(V* output, U* invvar, const T* input, int n1, int n2, double epsilon, const V* gamma) {\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n const dim3 threads(32, 4, 1);\n const uint64_t maxGridY = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];\n const dim3 blocks(1, std::min((uint64_t)n1, maxGridY), 1);\n int nshared = threads.y > 1 ? threads.y * sizeof(U) + (threads.y / 2) * sizeof(U) : 0;\n cuApplyRMSNorm<<>>(output, invvar, input, n1, n2, U(epsilon), gamma);\n}\n\nvoid cuda_layer_norm(at::Tensor& output, at::Tensor& mean, at::Tensor& invvar, const at::Tensor& input, int n1, int n2,\n at::IntArrayRef normalized_shape, const std::optional& gamma,\n const std::optional& beta, double epsilon) {\n using namespace at;\n DISPATCH_DOUBLE_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(\n input.scalar_type(), output.scalar_type(), \"layer_norm_cuda_kernel\",\n using accscalar_t = at::acc_type;\n HostApplyLayerNorm(\n output.data_ptr(), mean.data_ptr(), invvar.data_ptr(),\n input.data_ptr(), n1, n2, epsilon, gamma.has_value() ? gamma->data_ptr() : nullptr,\n beta.has_value() ? beta->data_ptr() : nullptr);)\n}\n\nvoid cuda_rms_norm(at::Tensor& output, at::Tensor& invvar, const at::Tensor& input, int n1, int n2,\n at::IntArrayRef normalized_shape, const std::optional& gamma, double epsilon) {\n using namespace at;\n DISPATCH_DOUBLE_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(\n input.scalar_type(), output.scalar_type(), \"rms_norm_cuda_kernel\",\n using accscalar_t = at::acc_type;\n HostApplyRMSNorm(\n output.data_ptr(), invvar.data_ptr(), input.data_ptr(), n1, n2,\n epsilon, gamma.has_value() ? gamma->data_ptr() : nullptr);)\n}\n\ntemplate \nvoid HostLayerNormGradient(const V* dout, const U* mean, const U* invvar, at::Tensor& input_or_output, int n1, int n2,\n const V* gamma, const V* beta, double epsilon, T* grad_input, V* grad_gamma, V* grad_beta,\n bool memory_efficient) {\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n\n if (gamma != nullptr && beta != nullptr) {\n // compute grad_gamma(j) and grad_beta(j)\n const int part_size = 16;\n const dim3 threads2(32, 4, 1);\n const dim3 blocks2((n2 + threads2.x - 1) / threads2.x, part_size, 1);\n const int nshared2_a = 2 * sizeof(U) * threads2.y * threads2.y * (threads2.x + 1);\n const int nshared2_b = threads2.x * threads2.y * sizeof(U);\n const int nshared2 = nshared2_a > nshared2_b ? nshared2_a : nshared2_b;\n // note (mkozuki): I can hard code part_grad_gamma's dtype as float given that\n // the `cuda_layer_norm_gradient` doesn't support double.\n const auto part_grad_dtype = (input_or_output.scalar_type() == at::ScalarType::Half ||\n input_or_output.scalar_type() == at::ScalarType::BFloat16)\n ? at::ScalarType::Float\n : input_or_output.scalar_type();\n at::Tensor part_grad_gamma = at::empty({part_size, n2}, input_or_output.options().dtype(part_grad_dtype));\n at::Tensor part_grad_beta = at::empty_like(part_grad_gamma);\n BOOL_SWITCH(memory_efficient, MemoryEfficient, [&] {\n auto kernel = &cuComputePartGradGammaBeta;\n kernel<<>>(dout, input_or_output.data_ptr(), n1, n2, mean, invvar,\n U(epsilon), gamma, beta, part_grad_gamma.data_ptr(),\n part_grad_beta.data_ptr(), epsilon, false);\n });\n\n const dim3 threads3(32, 8, 1);\n const dim3 blocks3((n2 + threads2.x - 1) / threads2.x, 1, 1);\n const int nshared3 = threads3.x * threads3.y * sizeof(U);\n cuComputeGradGammaBeta<<>>(\n part_grad_gamma.data_ptr(), part_grad_beta.data_ptr(), part_size, n1, n2, grad_gamma, grad_beta, false);\n }\n\n // compute grad_input\n const uint64_t maxGridY = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];\n const dim3 blocks1(1, std::min((uint64_t)n1, maxGridY), 1);\n const dim3 threads1(32, 4, 1);\n int nshared = threads1.y > 1 ? threads1.y * threads1.x * sizeof(U) : 0;\n BOOL_SWITCH(memory_efficient, MemoryEfficient, [&] {\n auto kernel = cuComputeGradInput;\n kernel<<>>(dout, input_or_output.data_ptr(), n1, n2, mean, invvar,\n U(epsilon), gamma, beta, grad_input, epsilon, false);\n });\n}\n\ntemplate \nvoid HostRMSNormGradient(const V* dout, const U* invvar, at::Tensor& input_or_output, int n1, int n2, const V* gamma,\n double epsilon, T* grad_input, V* grad_gamma, bool memory_efficient) {\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n\n if (gamma != nullptr) {\n const int part_size = 16;\n const dim3 threads2(32, 4, 1);\n const dim3 blocks2((n2 + threads2.x - 1) / threads2.x, part_size, 1);\n const int nshared2_a = 2 * sizeof(U) * threads2.y * threads2.y * (threads2.x + 1);\n const int nshared2_b = threads2.x * threads2.y * sizeof(U);\n const int nshared2 = nshared2_a > nshared2_b ? nshared2_a : nshared2_b;\n // note (mkozuki): I can hard code part_grad_gamma's dtype as float given that\n // the `cuda_layer_norm_gradient` doesn't support double.\n const auto part_grad_dtype = (input_or_output.scalar_type() == at::ScalarType::Half ||\n input_or_output.scalar_type() == at::ScalarType::BFloat16)\n ? at::ScalarType::Float\n : input_or_output.scalar_type();\n at::Tensor part_grad_gamma = at::empty({part_size, n2}, input_or_output.options().dtype(part_grad_dtype));\n BOOL_SWITCH(memory_efficient, MemoryEfficient, [&] {\n auto kernel = &cuComputePartGradGammaBeta;\n kernel<<>>(dout, input_or_output.data_ptr(), n1, n2, invvar, /* unused */\n invvar, U(epsilon), gamma, gamma, /* unused */\n part_grad_gamma.data_ptr(),\n part_grad_gamma.data_ptr(), /* unused */\n epsilon, true);\n });\n\n const dim3 threads3(32, 8, 1);\n const dim3 blocks3((n2 + threads2.x - 1) / threads2.x, 1, 1);\n const int nshared3 = threads3.x * threads3.y * sizeof(U);\n cuComputeGradGammaBeta<<>>(\n part_grad_gamma.data_ptr(), part_grad_gamma.data_ptr(), /* unused */\n part_size, n1, n2, grad_gamma, grad_gamma, /* unused */\n true);\n }\n\n // compute grad_input\n const uint64_t maxGridY = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];\n const dim3 blocks1(1, std::min((uint64_t)n1, maxGridY), 1);\n const dim3 threads1(32, 4, 1);\n int nshared = threads1.y > 1 ? threads1.y * threads1.x * sizeof(U) : 0;\n BOOL_SWITCH(memory_efficient, MemoryEfficient, [&] {\n auto kernel = cuComputeGradInput;\n kernel<<>>(dout, input_or_output.data_ptr(), n1, n2, invvar, /* unused */\n invvar, U(epsilon), gamma, gamma, /* unused */\n grad_input, epsilon, true);\n });\n}\n\nvoid cuda_layer_norm_gradient(at::Tensor& dout, const std::optional& mean, at::Tensor& invvar,\n at::Tensor& input_or_output, int n1, int n2, at::IntArrayRef normalized_shape,\n const std::optional& gamma, const std::optional& beta,\n double epsilon, at::Tensor& grad_input, const std::optional& grad_gamma,\n const std::optional& grad_beta, bool memory_efficient) {\n using namespace at;\n // we can do away with `accscalar_t` as there're only three dtypes: fp32, fp16, bf16\n DISPATCH_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(\n input_or_output.scalar_type(), gamma.has_value() ? gamma->scalar_type() : input_or_output.scalar_type(),\n \"cuComputeGradInput\", using accscalar_t = at::acc_type;\n HostLayerNormGradient(dout.data_ptr(), mean.has_value() ? mean->data_ptr() : nullptr,\n invvar.data_ptr(), input_or_output, n1, n2,\n // TMJ pass nullptr argument for gamma, beta, grad_gamma and grad_beta\n // if gamma Tensor is nullptr on input.\n gamma.has_value() ? gamma->data_ptr() : nullptr,\n gamma.has_value() ? beta->data_ptr() : nullptr, epsilon,\n grad_input.data_ptr(),\n gamma.has_value() ? grad_gamma->data_ptr() : nullptr,\n gamma.has_value() ? grad_beta->data_ptr() : nullptr, memory_efficient);)\n}\n\nvoid cuda_rms_norm_gradient(at::Tensor& dout, at::Tensor& invvar, at::Tensor& input_or_output, int n1, int n2,\n at::IntArrayRef normalized_shape, const std::optional& gamma, double epsilon,\n at::Tensor& grad_input, const std::optional& grad_gamma,\n bool memory_efficient) {\n using namespace at;\n // we can do away with `accscalar_t` as there're only three dtypes: fp32, fp16, bf16\n // DISPATCH_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(\n DISPATCH_DOUBLE_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(\n input_or_output.scalar_type(), gamma.has_value() ? gamma->scalar_type() : input_or_output.scalar_type(),\n \"cuComputeGradInputRMS\", using accscalar_t = at::acc_type;\n HostRMSNormGradient(dout.data_ptr(), invvar.data_ptr(), input_or_output, n1, n2,\n // TMJ pass nullptr argument for gamma, beta, grad_gamma and grad_beta\n // if gamma Tensor is nullptr on input.\n gamma.has_value() ? gamma->data_ptr() : nullptr, epsilon,\n grad_input.data_ptr(),\n gamma.has_value() ? grad_gamma->data_ptr() : nullptr, memory_efficient);)\n}\n", "output": "#include \n\n#include \n\n#include \"ATen/ATen.h\"\n#include \"ATen/AccumulateType.h\"\n#include \"ATen/cuda/CUDAContext.h\"\n#include \"ATen/cuda/DeviceUtils.cuh\"\n#include \"static_switch.h\"\n#include \"type_shim.h\"\n\ntemplate \n__device__ void cuWelfordOnlineSum(const U curr, U& mu, U& sigma2, U& count) {\n count = count + U(1);\n U delta = curr - mu;\n U lmean = mu + delta / count;\n mu = lmean;\n U delta2 = curr - lmean;\n sigma2 = sigma2 + delta * delta2;\n}\n\ntemplate \n__device__ void cuChanOnlineSum(const U muB, const U sigma2B, const U countB, U& mu, U& sigma2, U& count) {\n U delta = muB - mu;\n U nA = count;\n U nB = countB;\n count = count + countB;\n U nX = count;\n if (nX > U(0)) {\n nA = nA / nX;\n nB = nB / nX;\n mu = nA * mu + nB * muB;\n sigma2 = sigma2 + sigma2B + delta * delta * nA * nB * nX;\n } else {\n mu = U(0);\n sigma2 = U(0);\n }\n}\n\ntemplate \n__device__ void cuRMSOnlineSum(const U curr, U& sigma2) {\n sigma2 = sigma2 + curr * curr;\n}\n\ntemplate \n__device__ void cuChanRMSOnlineSum(const U sigma2B, U& sigma2) {\n sigma2 = sigma2 + sigma2B;\n}\n\ntemplate \n__device__ void cuWelfordMuSigma2(const T* __restrict__ vals, const int n1, const int n2, const int i1, U& mu,\n U& sigma2, U* buf, bool rms_only) {\n // Assumptions:\n // 1) blockDim.x == warpSize\n // 2) Tensor is contiguous\n // 3) 2*blockDim.y*sizeof(U)+blockDim.y*sizeof(int) shared memory available.\n //\n // compute variance and mean over n2\n U count = U(0);\n mu = U(0);\n sigma2 = U(0);\n if (i1 < n1) {\n // one warp normalizes one n1 index,\n // synchronization is implicit\n // initialize with standard Welford algorithm\n const int numx = blockDim.x * blockDim.y;\n const int thrx = threadIdx.x + threadIdx.y * blockDim.x;\n const T* lvals = vals + i1 * n2;\n int l = 4 * thrx;\n for (; l + 3 < n2; l += 4 * numx) {\n for (int k = 0; k < 4; ++k) {\n U curr = static_cast(lvals[l + k]);\n if (!rms_only) {\n cuWelfordOnlineSum(curr, mu, sigma2, count);\n } else {\n cuRMSOnlineSum(curr, sigma2);\n }\n }\n }\n for (; l < n2; ++l) {\n U curr = static_cast(lvals[l]);\n if (!rms_only) {\n cuWelfordOnlineSum(curr, mu, sigma2, count);\n } else {\n cuRMSOnlineSum(curr, sigma2);\n }\n }\n // intra-warp reductions\n for (int l = 0; l <= 4; ++l) {\n int srcLaneB = (threadIdx.x + (1 << l)) & 31;\n U sigma2B = WARP_SHFL(sigma2, srcLaneB);\n if (!rms_only) {\n U muB = WARP_SHFL(mu, srcLaneB);\n U countB = WARP_SHFL(count, srcLaneB);\n cuChanOnlineSum(muB, sigma2B, countB, mu, sigma2, count);\n } else {\n cuChanRMSOnlineSum(sigma2B, sigma2);\n }\n }\n // threadIdx.x == 0 has correct values for each warp\n // inter-warp reductions\n if (blockDim.y > 1) {\n U* ubuf = (U*)buf;\n U* ibuf = (U*)(ubuf + blockDim.y);\n for (int offset = blockDim.y / 2; offset > 0; offset /= 2) {\n // upper half of warps write to shared\n if (threadIdx.x == 0 && threadIdx.y >= offset && threadIdx.y < 2 * offset) {\n const int wrt_y = threadIdx.y - offset;\n if (!rms_only) {\n ubuf[2 * wrt_y] = mu;\n ibuf[wrt_y] = count;\n }\n ubuf[2 * wrt_y + 1] = sigma2;\n }\n __syncthreads();\n // lower half merges\n if (threadIdx.x == 0 && threadIdx.y < offset) {\n U sigma2B = ubuf[2 * threadIdx.y + 1];\n if (!rms_only) {\n U muB = ubuf[2 * threadIdx.y];\n U countB = ibuf[threadIdx.y];\n cuChanOnlineSum(muB, sigma2B, countB, mu, sigma2, count);\n } else {\n cuChanRMSOnlineSum(sigma2B, sigma2);\n }\n }\n __syncthreads();\n }\n // threadIdx.x = 0 && threadIdx.y == 0 only thread that has correct values\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n if (!rms_only) {\n ubuf[0] = mu;\n }\n ubuf[1] = sigma2;\n }\n __syncthreads();\n if (!rms_only) {\n mu = ubuf[0];\n }\n sigma2 = ubuf[1] / U(n2);\n // don't care about final value of count, we know count == n2\n } else {\n if (!rms_only) {\n mu = WARP_SHFL(mu, 0);\n }\n sigma2 = WARP_SHFL(sigma2 / U(n2), 0);\n }\n }\n}\n\ntemplate <>\n__device__ void cuWelfordMuSigma2(const at::Half* __restrict__ vals, const int n1, const int n2, const int i1,\n float& mu, float& sigma2, float* buf, bool rms_only) {\n // Assumptions:\n // 1) blockDim.x == warpSize\n // 2) Tensor is contiguous\n // 3) 2*blockDim.y*sizeof(U)+blockDim.y*sizeof(int) shared memory available.\n //\n // compute variance and mean over n2\n float count = 0.0f;\n mu = float(0);\n sigma2 = float(0);\n if (i1 < n1) {\n // one warp normalizes one n1 index,\n // synchronization is implicit\n // initialize with standard Welford algorithm\n const int numx = blockDim.x * blockDim.y;\n const int thrx = threadIdx.x + threadIdx.y * blockDim.x;\n const at::Half* lvals = vals + i1 * n2;\n int l = 8 * thrx;\n if ((((size_t)lvals) & 3) != 0) {\n // 16 bit alignment\n // first thread consumes first point\n if (thrx == 0) {\n float curr = static_cast(lvals[0]);\n if (!rms_only) {\n cuWelfordOnlineSum(curr, mu, sigma2, count);\n } else {\n cuRMSOnlineSum(curr, sigma2);\n }\n }\n ++l;\n }\n // at this point, lvals[l] are 32 bit aligned for all threads.\n for (; l + 7 < n2; l += 8 * numx) {\n for (int k = 0; k < 8; k += 2) {\n float2 curr = __half22float2(*((__half2*)(lvals + l + k)));\n if (!rms_only) {\n cuWelfordOnlineSum(curr.x, mu, sigma2, count);\n cuWelfordOnlineSum(curr.y, mu, sigma2, count);\n } else {\n cuRMSOnlineSum(curr.x, sigma2);\n cuRMSOnlineSum(curr.y, sigma2);\n }\n }\n }\n for (; l < n2; ++l) {\n float curr = static_cast(lvals[l]);\n if (!rms_only) {\n cuWelfordOnlineSum(curr, mu, sigma2, count);\n } else {\n cuRMSOnlineSum(curr, sigma2);\n }\n }\n // intra-warp reductions\n for (int l = 0; l <= 4; ++l) {\n int srcLaneB = (threadIdx.x + (1 << l)) & 31;\n float sigma2B = WARP_SHFL(sigma2, srcLaneB);\n if (!rms_only) {\n float muB = WARP_SHFL(mu, srcLaneB);\n float countB = WARP_SHFL(count, srcLaneB);\n cuChanOnlineSum(muB, sigma2B, countB, mu, sigma2, count);\n } else {\n cuChanRMSOnlineSum(sigma2B, sigma2);\n }\n }\n // threadIdx.x == 0 has correct values for each warp\n // inter-warp reductions\n if (blockDim.y > 1) {\n float* ubuf = (float*)buf;\n float* ibuf = (float*)(ubuf + blockDim.y);\n for (int offset = blockDim.y / 2; offset > 0; offset /= 2) {\n // upper half of warps write to shared\n if (threadIdx.x == 0 && threadIdx.y >= offset && threadIdx.y < 2 * offset) {\n const int wrt_y = threadIdx.y - offset;\n ubuf[2 * wrt_y + 1] = sigma2;\n if (!rms_only) {\n ubuf[2 * wrt_y] = mu;\n ibuf[wrt_y] = count;\n }\n }\n __syncthreads();\n // lower half merges\n if (threadIdx.x == 0 && threadIdx.y < offset) {\n float sigma2B = ubuf[2 * threadIdx.y + 1];\n if (!rms_only) {\n float muB = ubuf[2 * threadIdx.y];\n float countB = ibuf[threadIdx.y];\n cuChanOnlineSum(muB, sigma2B, countB, mu, sigma2, count);\n } else {\n cuChanRMSOnlineSum(sigma2B, sigma2);\n }\n }\n __syncthreads();\n }\n // threadIdx.x = 0 && threadIdx.y == 0 only thread that has correct values\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n if (!rms_only) {\n ubuf[0] = mu;\n }\n ubuf[1] = sigma2;\n }\n __syncthreads();\n if (!rms_only) {\n mu = ubuf[0];\n }\n sigma2 = ubuf[1] / float(n2);\n // don't care about final value of count, we know count == n2\n } else {\n if (!rms_only) {\n mu = WARP_SHFL(mu, 0);\n }\n sigma2 = WARP_SHFL(sigma2 / float(n2), 0);\n }\n }\n}\n\ntemplate \nU rsqrt(U v) {\n return U(1) / sqrt(v);\n}\ntemplate <>\nfloat rsqrt(float v) {\n return rsqrtf(v);\n}\ntemplate <>\ndouble rsqrt(double v) {\n return rsqrt(v);\n}\n\nnamespace {\n// This is the un-specialized struct. Note that we prevent instantiation of this\n// struct by putting an undefined symbol in the function body so it won't compile.\n// template \n// struct SharedMemory\n// {\n// // Ensure that we won't compile any un-specialized types\n// __device__ T *getPointer()\n// {\n// extern __device__ void error(void);\n// error();\n// return nullptr;\n// }\n// };\n// https://github.com/NVIDIA/apex/issues/246\ntemplate \nstruct SharedMemory;\n\ntemplate <>\nstruct SharedMemory {\n __device__ float* getPointer() {\n extern __shared__ float s_float[];\n return s_float;\n }\n};\n\ntemplate <>\nstruct SharedMemory {\n __device__ double* getPointer() {\n extern __shared__ double s_double[];\n return s_double;\n }\n};\n} // namespace\n\ntemplate \n__device__ void cuApplyLayerNorm_(V* __restrict__ output_vals, U* __restrict__ mean, U* __restrict__ invvar,\n const T* __restrict__ vals, const int n1, const int n2, const U epsilon,\n const V* __restrict__ gamma, const V* __restrict__ beta, bool rms_only) {\n // Assumptions:\n // 1) blockDim.x == warpSize\n // 2) Tensors are contiguous\n //\n for (auto i1 = blockIdx.y; i1 < n1; i1 += gridDim.y) {\n SharedMemory shared;\n U* buf = shared.getPointer();\n U mu, sigma2;\n cuWelfordMuSigma2(vals, n1, n2, i1, mu, sigma2, buf, rms_only);\n\n const T* lvals = vals + i1 * n2;\n V* ovals = output_vals + i1 * n2;\n U c_invvar = rsqrt(sigma2 + epsilon);\n const int numx = blockDim.x * blockDim.y;\n const int thrx = threadIdx.x + threadIdx.y * blockDim.x;\n if (gamma != nullptr && (beta != nullptr || rms_only)) {\n for (int i = thrx; i < n2; i += numx) {\n U curr = static_cast(lvals[i]);\n if (!rms_only) {\n ovals[i] = gamma[i] * static_cast(c_invvar * (curr - mu)) + beta[i];\n } else {\n ovals[i] = gamma[i] * static_cast(c_invvar * curr);\n }\n }\n } else {\n for (int i = thrx; i < n2; i += numx) {\n U curr = static_cast(lvals[i]);\n if (!rms_only) {\n ovals[i] = static_cast(c_invvar * (curr - mu));\n } else {\n ovals[i] = static_cast(c_invvar * curr);\n }\n }\n }\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n if (!rms_only) {\n mean[i1] = mu;\n }\n invvar[i1] = c_invvar;\n }\n __syncthreads();\n }\n}\n\ntemplate \n__global__ void cuApplyLayerNorm(V* __restrict__ output_vals, U* __restrict__ mean, U* __restrict__ invvar,\n const T* __restrict__ vals, const int n1, const int n2, const U epsilon,\n const V* __restrict__ gamma, const V* __restrict__ beta) {\n cuApplyLayerNorm_(output_vals, mean, invvar, vals, n1, n2, epsilon, gamma, beta, false);\n}\n\ntemplate \n__global__ void cuApplyRMSNorm(V* __restrict__ output_vals, U* __restrict__ invvar, const T* __restrict__ vals,\n const int n1, const int n2, const U epsilon, const V* __restrict__ gamma) {\n cuApplyLayerNorm_(output_vals, nullptr, invvar, vals, n1, n2, epsilon, gamma, nullptr, true);\n}\n\ntemplate \n__device__ V clamp_by_magnitude(V curr_gamma, double eps) {\n const V kMinGamma = V(eps);\n if (curr_gamma >= 0) {\n if (curr_gamma < kMinGamma) {\n return kMinGamma;\n } else {\n return curr_gamma;\n }\n } else {\n if (curr_gamma > -kMinGamma) {\n return -kMinGamma;\n } else {\n return curr_gamma;\n }\n }\n}\n\ntemplate \n__device__ void cuLoadWriteStridedInputs(const int i1_block, const int thr_load_row_off, const int thr_load_col_off,\n const int i2_off, const int row_stride, U* warp_buf1, U* warp_buf2,\n const T* input_or_output, const V* dout, const int i1_end, const int n2,\n const U* __restrict__ mean, const U* __restrict__ invvar,\n const V* __restrict__ gamma, const V* __restrict__ beta, const double eps,\n bool rms_only) {\n int i1 = i1_block + thr_load_row_off;\n if (i1 < i1_end) {\n for (int k = 0; k < blockDim.y; ++k) {\n int i2 = i2_off + k;\n int load_idx = i1 * n2 + i2;\n int write_idx = thr_load_row_off * row_stride + thr_load_col_off + k;\n if (i2 < n2) {\n U c_h = static_cast(input_or_output[load_idx]);\n U curr_dout = static_cast(dout[load_idx]);\n if (!rms_only) {\n warp_buf1[write_idx] = curr_dout;\n if (MemoryEfficient) {\n U curr_beta = static_cast(beta[i2]);\n warp_buf2[write_idx] = curr_dout * (c_h - curr_beta) / static_cast(clamp_by_magnitude(gamma[i2], eps));\n } else {\n warp_buf2[write_idx] = curr_dout * (c_h - mean[i1]) * invvar[i1];\n }\n } else {\n if (MemoryEfficient) {\n warp_buf2[write_idx] = curr_dout * (c_h) / static_cast(clamp_by_magnitude(gamma[i2], eps));\n } else {\n warp_buf2[write_idx] = curr_dout * (c_h)*invvar[i1];\n }\n }\n } else {\n if (!rms_only) {\n warp_buf1[write_idx] = U(0);\n }\n warp_buf2[write_idx] = U(0);\n }\n }\n } else {\n for (int k = 0; k < blockDim.y; ++k) {\n int write_idx = thr_load_row_off * row_stride + thr_load_col_off + k;\n if (!rms_only) {\n warp_buf1[write_idx] = U(0);\n }\n warp_buf2[write_idx] = U(0);\n }\n }\n}\n\ntemplate \n__device__ void cuLoadAddStridedInputs(const int i1_block, const int thr_load_row_off, const int thr_load_col_off,\n const int i2_off, const int row_stride, U* warp_buf1, U* warp_buf2,\n const T* input_or_output, const V* dout, const int i1_end, const int n2,\n const U* __restrict__ mean, const U* __restrict__ invvar,\n const V* __restrict__ gamma, const V* __restrict__ beta, const double eps,\n bool rms_only) {\n int i1 = i1_block + thr_load_row_off;\n if (i1 < i1_end) {\n for (int k = 0; k < blockDim.y; ++k) {\n int i2 = i2_off + k;\n int load_idx = i1 * n2 + i2;\n int write_idx = thr_load_row_off * row_stride + thr_load_col_off + k;\n if (i2 < n2) {\n U c_h = static_cast(input_or_output[load_idx]);\n U curr_dout = static_cast(dout[load_idx]);\n if (!rms_only) {\n U curr_beta = static_cast(beta[i2]);\n warp_buf1[write_idx] += curr_dout;\n if (MemoryEfficient) {\n warp_buf2[write_idx] += curr_dout * (c_h - curr_beta) / static_cast(clamp_by_magnitude(gamma[i2], eps));\n } else {\n warp_buf2[write_idx] += curr_dout * (c_h - mean[i1]) * invvar[i1];\n }\n } else {\n if (MemoryEfficient) {\n warp_buf2[write_idx] += curr_dout * (c_h) / static_cast(clamp_by_magnitude(gamma[i2], eps));\n } else {\n warp_buf2[write_idx] += curr_dout * (c_h)*invvar[i1];\n }\n }\n }\n }\n }\n}\n\ntemplate \n__global__ void cuComputePartGradGammaBeta(const V* __restrict__ dout, const T* __restrict__ input_or_output,\n const int n1, const int n2, const U* __restrict__ mean,\n const U* __restrict__ invvar, U epsilon, const V* __restrict__ gamma,\n const V* __restrict__ beta, U* part_grad_gamma, U* part_grad_beta,\n const double eps, bool rms_only) {\n const int numsegs_n1 = (n1 + blockDim.y * blockDim.y - 1) / (blockDim.y * blockDim.y);\n const int segs_per_block = (numsegs_n1 + gridDim.y - 1) / gridDim.y;\n const int i1_beg = blockIdx.y * segs_per_block * blockDim.y * blockDim.y;\n const int i1_beg_plus_one = (blockIdx.y + 1) * segs_per_block * blockDim.y * blockDim.y;\n const int i1_end = i1_beg_plus_one < n1 ? i1_beg_plus_one : n1;\n const int row_stride = blockDim.x + 1;\n const int thr_load_col_off = (threadIdx.x * blockDim.y) & (blockDim.x - 1);\n const int thr_load_row_off = (threadIdx.x * blockDim.y) / blockDim.x + threadIdx.y * blockDim.y;\n const int i2_off = blockIdx.x * blockDim.x + thr_load_col_off;\n SharedMemory shared;\n U* buf = shared.getPointer(); // buf has at least blockDim.x * blockDim.y * blockDim.y + (blockDim.y -\n // 1)*(blockDim.x/blockDim.y) elements\n U* warp_buf1 = (U*)buf;\n U* warp_buf2 = warp_buf1 + blockDim.y * blockDim.y * row_stride;\n // compute partial sums from strided inputs\n // do this to increase number of loads in flight\n cuLoadWriteStridedInputs(i1_beg, thr_load_row_off, thr_load_col_off, i2_off, row_stride,\n warp_buf1, warp_buf2, input_or_output, dout, i1_end, n2, mean,\n invvar, gamma, beta, eps, rms_only);\n for (int i1_block = i1_beg + blockDim.y * blockDim.y; i1_block < i1_end; i1_block += blockDim.y * blockDim.y) {\n cuLoadAddStridedInputs(i1_block, thr_load_row_off, thr_load_col_off, i2_off, row_stride,\n warp_buf1, warp_buf2, input_or_output, dout, i1_end, n2, mean,\n invvar, gamma, beta, eps, rms_only);\n }\n __syncthreads();\n // inter-warp reductions\n // sum within each warp\n U acc1 = U(0);\n U acc2 = U(0);\n for (int k = 0; k < blockDim.y; ++k) {\n int row1 = threadIdx.y + k * blockDim.y;\n int idx1 = row1 * row_stride + threadIdx.x;\n if (!rms_only) {\n acc1 += warp_buf1[idx1];\n }\n acc2 += warp_buf2[idx1];\n }\n if (!rms_only) {\n warp_buf1[threadIdx.y * row_stride + threadIdx.x] = acc1;\n }\n warp_buf2[threadIdx.y * row_stride + threadIdx.x] = acc2;\n __syncthreads();\n // sum all warps\n for (int offset = blockDim.y / 2; offset > 1; offset /= 2) {\n if (threadIdx.y < offset) {\n int row1 = threadIdx.y;\n int row2 = threadIdx.y + offset;\n int idx1 = row1 * row_stride + threadIdx.x;\n int idx2 = row2 * row_stride + threadIdx.x;\n if (!rms_only) {\n warp_buf1[idx1] += warp_buf1[idx2];\n }\n warp_buf2[idx1] += warp_buf2[idx2];\n }\n __syncthreads();\n }\n int i2 = blockIdx.x * blockDim.x + threadIdx.x;\n if (threadIdx.y == 0 && i2 < n2) {\n int row1 = threadIdx.y;\n int row2 = threadIdx.y + 1;\n int idx1 = row1 * row_stride + threadIdx.x;\n int idx2 = row2 * row_stride + threadIdx.x;\n if (!rms_only) {\n part_grad_beta[blockIdx.y * n2 + i2] = warp_buf1[idx1] + warp_buf1[idx2];\n }\n part_grad_gamma[blockIdx.y * n2 + i2] = warp_buf2[idx1] + warp_buf2[idx2];\n }\n}\n\ntemplate \n__global__ void cuComputeGradGammaBeta(const U* part_grad_gamma, const U* part_grad_beta, const int part_size,\n const int n1, const int n2, V* grad_gamma, V* grad_beta, bool rms_only) {\n // sum partial gradients for gamma and beta\n SharedMemory shared;\n U* buf = shared.getPointer();\n int i2 = blockIdx.x * blockDim.x + threadIdx.x;\n if (i2 < n2) {\n // each warp does sequential reductions until reduced part_size is num_warps\n int num_warp_reductions = part_size / blockDim.y;\n U sum_gamma = U(0);\n U sum_beta = U(0);\n const U* part_grad_gamma_ptr = part_grad_gamma + threadIdx.y * num_warp_reductions * n2 + i2;\n const U* part_grad_beta_ptr = part_grad_beta + threadIdx.y * num_warp_reductions * n2 + i2;\n for (int warp_offset = 0; warp_offset < num_warp_reductions; ++warp_offset) {\n sum_gamma += part_grad_gamma_ptr[warp_offset * n2];\n if (!rms_only) {\n sum_beta += part_grad_beta_ptr[warp_offset * n2];\n }\n }\n // inter-warp reductions\n const int nbsize3 = blockDim.x * blockDim.y / 2;\n for (int offset = blockDim.y / 2; offset >= 1; offset /= 2) {\n // top half write to shared memory\n if (threadIdx.y >= offset && threadIdx.y < 2 * offset) {\n const int write_idx = (threadIdx.y - offset) * blockDim.x + threadIdx.x;\n buf[write_idx] = sum_gamma;\n if (!rms_only) {\n buf[write_idx + nbsize3] = sum_beta;\n }\n }\n __syncthreads();\n // bottom half sums\n if (threadIdx.y < offset) {\n const int read_idx = threadIdx.y * blockDim.x + threadIdx.x;\n sum_gamma += buf[read_idx];\n if (!rms_only) {\n sum_beta += buf[read_idx + nbsize3];\n }\n }\n __syncthreads();\n }\n // write out fully summed gradients\n if (threadIdx.y == 0) {\n grad_gamma[i2] = sum_gamma;\n if (!rms_only) {\n grad_beta[i2] = sum_beta;\n }\n }\n }\n}\n\ntemplate \n__global__ void cuComputeGradInput(const V* __restrict__ dout, const T* __restrict__ input_or_output, const int n1,\n const int n2, const U* __restrict__ mean, const U* __restrict__ invvar, U epsilon,\n const V* gamma, const V* beta, T* grad_input, const double eps, bool rms_only) {\n for (auto i1 = blockIdx.y; i1 < n1; i1 += gridDim.y) {\n U sum_loss1 = U(0);\n U sum_loss2 = U(0);\n const T* k_h = input_or_output + i1 * n2;\n const V* k_dout = dout + i1 * n2;\n const U c_invvar = invvar[i1];\n const U c_mean = !MemoryEfficient ? mean[i1] : 0.;\n const int numx = blockDim.x * blockDim.y;\n const int thrx = threadIdx.x + threadIdx.y * blockDim.x;\n if (gamma != nullptr) {\n int l = 4 * thrx;\n for (; l + 3 < n2; l += 4 * numx) {\n for (int k = 0; k < 4; ++k) {\n const U c_h = static_cast(k_h[l + k]);\n const U c_loss = static_cast(k_dout[l + k]);\n if (!rms_only) {\n sum_loss1 += c_loss * gamma[l + k];\n if (MemoryEfficient) {\n sum_loss2 += c_loss * (c_h - beta[l + k]);\n } else {\n sum_loss2 += c_loss * gamma[l + k] * (c_h - c_mean) * c_invvar;\n }\n } else {\n if (MemoryEfficient) {\n sum_loss2 += c_loss * c_h;\n } else {\n sum_loss2 += c_loss * gamma[l + k] * (c_h)*c_invvar;\n }\n }\n }\n }\n for (; l < n2; ++l) {\n const U c_h = static_cast(k_h[l]);\n const U c_loss = static_cast(k_dout[l]);\n if (!rms_only) {\n sum_loss1 += c_loss * gamma[l];\n if (MemoryEfficient) {\n sum_loss2 += c_loss * (c_h - beta[l]);\n } else {\n sum_loss2 += c_loss * gamma[l] * (c_h - c_mean) * c_invvar;\n }\n } else {\n if (MemoryEfficient) {\n sum_loss2 += c_loss * c_h;\n } else {\n sum_loss2 += c_loss * gamma[l] * (c_h)*c_invvar;\n }\n }\n }\n } else {\n int l = 4 * thrx;\n for (; l + 3 < n2; l += 4 * numx) {\n for (int k = 0; k < 4; ++k) {\n const U c_h = static_cast(k_h[l + k]);\n const U c_loss = static_cast(k_dout[l + k]);\n if (!rms_only) {\n sum_loss1 += c_loss;\n if (MemoryEfficient) {\n sum_loss2 += c_loss * c_h;\n } else {\n sum_loss2 += c_loss * (c_h - c_mean) * c_invvar;\n }\n } else {\n if (MemoryEfficient) {\n sum_loss2 += c_loss * c_h;\n } else {\n sum_loss2 += c_loss * (c_h)*c_invvar;\n }\n }\n }\n }\n for (; l < n2; ++l) {\n const U c_h = static_cast(k_h[l]);\n const U c_loss = static_cast(k_dout[l]);\n if (!rms_only) {\n sum_loss1 += c_loss;\n if (MemoryEfficient) {\n sum_loss2 += c_loss * c_h;\n } else {\n sum_loss2 += c_loss * (c_h - c_mean) * c_invvar;\n }\n } else {\n if (MemoryEfficient) {\n sum_loss2 += c_loss * c_h;\n } else {\n sum_loss2 += c_loss * (c_h)*c_invvar;\n }\n }\n }\n }\n // intra-warp reductions\n for (int mask = blockDim.x / 2; mask > 0; mask /= 2) {\n if (!rms_only) {\n sum_loss1 += WARP_SHFL_XOR(sum_loss1, mask);\n }\n sum_loss2 += WARP_SHFL_XOR(sum_loss2, mask);\n }\n // inter-warp reductions\n if (blockDim.y > 1) {\n SharedMemory shared;\n U* buf = shared.getPointer();\n for (int offset = blockDim.y / 2; offset > 0; offset /= 2) {\n // upper half of warps write to shared\n if (threadIdx.y >= offset && threadIdx.y < 2 * offset) {\n const int wrt_i = (threadIdx.y - offset) * blockDim.x + threadIdx.x;\n if (!rms_only) {\n buf[2 * wrt_i] = sum_loss1;\n }\n buf[2 * wrt_i + 1] = sum_loss2;\n }\n __syncthreads();\n // lower half merges\n if (threadIdx.y < offset) {\n const int read_i = threadIdx.y * blockDim.x + threadIdx.x;\n if (!rms_only) {\n sum_loss1 += buf[2 * read_i];\n }\n sum_loss2 += buf[2 * read_i + 1];\n }\n __syncthreads();\n }\n if (threadIdx.y == 0) {\n if (!rms_only) {\n buf[2 * threadIdx.x] = sum_loss1;\n }\n buf[2 * threadIdx.x + 1] = sum_loss2;\n }\n __syncthreads();\n if (threadIdx.y != 0) {\n if (!rms_only) {\n sum_loss1 = buf[2 * threadIdx.x];\n }\n sum_loss2 = buf[2 * threadIdx.x + 1];\n }\n }\n // all threads now have the two sums over l\n U fH = (U)n2;\n U term1 = (U(1) / fH) * c_invvar;\n T* k_grad_input = grad_input + i1 * n2;\n if (gamma != nullptr) {\n for (int l = thrx; l < n2; l += numx) {\n const U c_h = static_cast(k_h[l]);\n const U c_loss = static_cast(k_dout[l]);\n const U k_gamma = static_cast(clamp_by_magnitude(gamma[l], eps));\n U f_grad_input = fH * c_loss * k_gamma;\n if (!rms_only) {\n const U k_beta = beta[l];\n f_grad_input -= sum_loss1;\n if (MemoryEfficient) {\n f_grad_input -= (c_h - k_beta) / k_gamma * sum_loss2;\n } else {\n f_grad_input -= (c_h - c_mean) * c_invvar * sum_loss2;\n }\n } else {\n if (MemoryEfficient) {\n f_grad_input -= c_h / k_gamma * sum_loss2;\n } else {\n f_grad_input -= c_h * c_invvar * sum_loss2;\n }\n }\n f_grad_input *= term1;\n k_grad_input[l] = static_cast(f_grad_input);\n }\n } else {\n for (int l = thrx; l < n2; l += numx) {\n const U c_h = static_cast(k_h[l]);\n const U c_loss = static_cast(k_dout[l]);\n U f_grad_input = fH * c_loss;\n if (!rms_only) {\n f_grad_input -= sum_loss1;\n if (MemoryEfficient) {\n f_grad_input -= c_h * sum_loss2;\n } else {\n f_grad_input -= (c_h - c_mean) * c_invvar * sum_loss2;\n }\n } else {\n if (MemoryEfficient) {\n f_grad_input -= c_h * sum_loss2;\n } else {\n f_grad_input -= c_h * c_invvar * sum_loss2;\n }\n }\n f_grad_input *= term1;\n k_grad_input[l] = static_cast(f_grad_input);\n }\n }\n // prevent race where buf is written again before reads are done\n __syncthreads();\n }\n}\n\ntemplate \nvoid HostApplyLayerNorm(V* output, U* mean, U* invvar, const T* input, int n1, int n2, double epsilon, const V* gamma,\n const V* beta) {\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n const dim3 threads(32, 4, 1);\n const uint64_t maxGridY = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];\n const dim3 blocks(1, std::min((uint64_t)n1, maxGridY), 1);\n int nshared = threads.y > 1 ? threads.y * sizeof(U) + (threads.y / 2) * sizeof(U) : 0;\n cuApplyLayerNorm<<>>(output, mean, invvar, input, n1, n2, U(epsilon), gamma, beta);\n}\n\ntemplate \nvoid HostApplyRMSNorm(V* output, U* invvar, const T* input, int n1, int n2, double epsilon, const V* gamma) {\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n const dim3 threads(32, 4, 1);\n const uint64_t maxGridY = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];\n const dim3 blocks(1, std::min((uint64_t)n1, maxGridY), 1);\n int nshared = threads.y > 1 ? threads.y * sizeof(U) + (threads.y / 2) * sizeof(U) : 0;\n cuApplyRMSNorm<<>>(output, invvar, input, n1, n2, U(epsilon), gamma);\n}\n\nvoid cuda_layer_norm(at::Tensor& output, at::Tensor& mean, at::Tensor& invvar, const at::Tensor& input, int n1, int n2,\n at::IntArrayRef normalized_shape, const std::optional& gamma,\n const std::optional& beta, double epsilon) {\n using namespace at;\n DISPATCH_DOUBLE_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(\n input.scalar_type(), output.scalar_type(), \"layer_norm_cuda_kernel\",\n using accscalar_t = at::acc_type;\n HostApplyLayerNorm(\n output.data_ptr(), mean.data_ptr(), invvar.data_ptr(),\n input.data_ptr(), n1, n2, epsilon, gamma.has_value() ? gamma->data_ptr() : nullptr,\n beta.has_value() ? beta->data_ptr() : nullptr);)\n}\n\nvoid cuda_rms_norm(at::Tensor& output, at::Tensor& invvar, const at::Tensor& input, int n1, int n2,\n at::IntArrayRef normalized_shape, const std::optional& gamma, double epsilon) {\n using namespace at;\n DISPATCH_DOUBLE_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(\n input.scalar_type(), output.scalar_type(), \"rms_norm_cuda_kernel\",\n using accscalar_t = at::acc_type;\n HostApplyRMSNorm(\n output.data_ptr(), invvar.data_ptr(), input.data_ptr(), n1, n2,\n epsilon, gamma.has_value() ? gamma->data_ptr() : nullptr);)\n}\n\ntemplate \nvoid HostLayerNormGradient(const V* dout, const U* mean, const U* invvar, at::Tensor& input_or_output, int n1, int n2,\n const V* gamma, const V* beta, double epsilon, T* grad_input, V* grad_gamma, V* grad_beta,\n bool memory_efficient) {\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n\n if (gamma != nullptr && beta != nullptr) {\n // compute grad_gamma(j) and grad_beta(j)\n const int part_size = 16;\n const dim3 threads2(32, 4, 1);\n const dim3 blocks2((n2 + threads2.x - 1) / threads2.x, part_size, 1);\n const int nshared2_a = 2 * sizeof(U) * threads2.y * threads2.y * (threads2.x + 1);\n const int nshared2_b = threads2.x * threads2.y * sizeof(U);\n const int nshared2 = nshared2_a > nshared2_b ? nshared2_a : nshared2_b;\n // note (mkozuki): I can hard code part_grad_gamma's dtype as float given that\n // the `cuda_layer_norm_gradient` doesn't support double.\n const auto part_grad_dtype = (input_or_output.scalar_type() == at::ScalarType::Half ||\n input_or_output.scalar_type() == at::ScalarType::BFloat16)\n ? at::ScalarType::Float\n : input_or_output.scalar_type();\n at::Tensor part_grad_gamma = at::empty({part_size, n2}, input_or_output.options().dtype(part_grad_dtype));\n at::Tensor part_grad_beta = at::empty_like(part_grad_gamma);\n BOOL_SWITCH(memory_efficient, MemoryEfficient, [&] {\n auto kernel = &cuComputePartGradGammaBeta;\n kernel<<>>(dout, input_or_output.data_ptr(), n1, n2, mean, invvar,\n U(epsilon), gamma, beta, part_grad_gamma.data_ptr(),\n part_grad_beta.data_ptr(), epsilon, false);\n });\n\n const dim3 threads3(32, 8, 1);\n const dim3 blocks3((n2 + threads2.x - 1) / threads2.x, 1, 1);\n const int nshared3 = threads3.x * threads3.y * sizeof(U);\n cuComputeGradGammaBeta<<>>(\n part_grad_gamma.data_ptr(), part_grad_beta.data_ptr(), part_size, n1, n2, grad_gamma, grad_beta, false);\n }\n\n // compute grad_input\n const uint64_t maxGridY = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];\n const dim3 blocks1(1, std::min((uint64_t)n1, maxGridY), 1);\n const dim3 threads1(32, 4, 1);\n int nshared = threads1.y > 1 ? threads1.y * threads1.x * sizeof(U) : 0;\n BOOL_SWITCH(memory_efficient, MemoryEfficient, [&] {\n auto kernel = cuComputeGradInput;\n kernel<<>>(dout, input_or_output.data_ptr(), n1, n2, mean, invvar,\n U(epsilon), gamma, beta, grad_input, epsilon, false);\n });\n}\n\ntemplate \nvoid HostRMSNormGradient(const V* dout, const U* invvar, at::Tensor& input_or_output, int n1, int n2, const V* gamma,\n double epsilon, T* grad_input, V* grad_gamma, bool memory_efficient) {\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n\n if (gamma != nullptr) {\n const int part_size = 16;\n const dim3 threads2(32, 4, 1);\n const dim3 blocks2((n2 + threads2.x - 1) / threads2.x, part_size, 1);\n const int nshared2_a = 2 * sizeof(U) * threads2.y * threads2.y * (threads2.x + 1);\n const int nshared2_b = threads2.x * threads2.y * sizeof(U);\n const int nshared2 = nshared2_a > nshared2_b ? nshared2_a : nshared2_b;\n // note (mkozuki): I can hard code part_grad_gamma's dtype as float given that\n // the `cuda_layer_norm_gradient` doesn't support double.\n const auto part_grad_dtype = (input_or_output.scalar_type() == at::ScalarType::Half ||\n input_or_output.scalar_type() == at::ScalarType::BFloat16)\n ? at::ScalarType::Float\n : input_or_output.scalar_type();\n at::Tensor part_grad_gamma = at::empty({part_size, n2}, input_or_output.options().dtype(part_grad_dtype));\n BOOL_SWITCH(memory_efficient, MemoryEfficient, [&] {\n auto kernel = &cuComputePartGradGammaBeta;\n kernel<<>>(dout, input_or_output.data_ptr(), n1, n2, invvar, /* unused */\n invvar, U(epsilon), gamma, gamma, /* unused */\n part_grad_gamma.data_ptr(),\n part_grad_gamma.data_ptr(), /* unused */\n epsilon, true);\n });\n\n const dim3 threads3(32, 8, 1);\n const dim3 blocks3((n2 + threads2.x - 1) / threads2.x, 1, 1);\n const int nshared3 = threads3.x * threads3.y * sizeof(U);\n cuComputeGradGammaBeta<<>>(\n part_grad_gamma.data_ptr(), part_grad_gamma.data_ptr(), /* unused */\n part_size, n1, n2, grad_gamma, grad_gamma, /* unused */\n true);\n }\n\n // compute grad_input\n const uint64_t maxGridY = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];\n const dim3 blocks1(1, std::min((uint64_t)n1, maxGridY), 1);\n const dim3 threads1(32, 4, 1);\n int nshared = threads1.y > 1 ? threads1.y * threads1.x * sizeof(U) : 0;\n BOOL_SWITCH(memory_efficient, MemoryEfficient, [&] {\n auto kernel = cuComputeGradInput;\n kernel<<>>(dout, input_or_output.data_ptr(), n1, n2, invvar, /* unused */\n invvar, U(epsilon), gamma, gamma, /* unused */\n grad_input, epsilon, true);\n });\n}\n\nvoid cuda_layer_norm_gradient(at::Tensor& dout, const std::optional& mean, at::Tensor& invvar,\n at::Tensor& input_or_output, int n1, int n2, at::IntArrayRef normalized_shape,\n const std::optional& gamma, const std::optional& beta,\n double epsilon, at::Tensor& grad_input, const std::optional& grad_gamma,\n const std::optional& grad_beta, bool memory_efficient) {\n using namespace at;\n // we can do away with `accscalar_t` as there're only three dtypes: fp32, fp16, bf16\n DISPATCH_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(\n input_or_output.scalar_type(), gamma.has_value() ? gamma->scalar_type() : input_or_output.scalar_type(),\n \"cuComputeGradInput\", using accscalar_t = at::acc_type;\n HostLayerNormGradient(dout.data_ptr(), mean.has_value() ? mean->data_ptr() : nullptr,\n invvar.data_ptr(), input_or_output, n1, n2,\n // TMJ pass nullptr argument for gamma, beta, grad_gamma and grad_beta\n // if gamma Tensor is nullptr on input.\n gamma.has_value() ? gamma->data_ptr() : nullptr,\n gamma.has_value() ? beta->data_ptr() : nullptr, epsilon,\n grad_input.data_ptr(),\n gamma.has_value() ? grad_gamma->data_ptr() : nullptr,\n gamma.has_value() ? grad_beta->data_ptr() : nullptr, memory_efficient);)\n}\n\nvoid cuda_rms_norm_gradient(at::Tensor& dout, at::Tensor& invvar, at::Tensor& input_or_output, int n1, int n2,\n at::IntArrayRef normalized_shape, const std::optional& gamma, double epsilon,\n at::Tensor& grad_input, const std::optional& grad_gamma,\n bool memory_efficient) {\n using namespace at;\n // we can do away with `accscalar_t` as there're only three dtypes: fp32, fp16, bf16\n // DISPATCH_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(\n DISPATCH_DOUBLE_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(\n input_or_output.scalar_type(), gamma.has_value() ? gamma->scalar_type() : input_or_output.scalar_type(),\n \"cuComputeGradInputRMS\", using accscalar_t = at::acc_type;\n HostRMSNormGradient(dout.data_ptr(), invvar.data_ptr(), input_or_output, n1, n2,\n // TMJ pass nullptr argument for gamma, beta, grad_gamma and grad_beta\n // if gamma Tensor is nullptr on input.\n gamma.has_value() ? gamma->data_ptr() : nullptr, epsilon,\n grad_input.data_ptr(),\n gamma.has_value() ? grad_gamma->data_ptr() : nullptr, memory_efficient);)\n}\n", "source": "/kaggle/working/apex/csrc/layer_norm_cuda_kernel.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n// Another possibility:\n// #include \n\n#include \n\n#include \"multi_tensor_apply.cuh\"\n#include \"type_shim.h\"\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntemplate \n__device__ __forceinline__ bool is_aligned(T* p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\n\ntemplate \nstruct L2NormScaleFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<2>& tl,\n float* output, float* output_per_tensor, float scale, bool per_tensor,\n int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n in_t* in = (in_t*)tl.addresses[0][tensor_loc];\n in += chunk_idx * chunk_size;\n\n out_t* out = (out_t*)tl.addresses[1][tensor_loc];\n out += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n in_t r_in[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_in[i] = 0;\n }\n // bool finite = true;\n out_t r_out[ILP];\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(in) && is_aligned(out)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n load_store(r_in, in, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_in[ii]);\n r_out[ii] = next * scale;\n vals[ii] += next * next;\n // finite = finite && isfinite(r_in[ii]);\n }\n load_store(out, r_out, i_start, 0);\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_in[ii] = 0;\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n r_in[ii] = in[i];\n float next = static_cast(in[i]);\n vals[ii] += next * next;\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_out[ii] = static_cast(r_in[ii]) * scale;\n // finite = finite && isfinite(r_in[ii]);\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) out[i] = r_out[ii];\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val += vals[i];\n\n float final = reduce_block_into_lanes(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final)) *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] += final;\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;\n }\n }\n};\n// Probably better to template, but since we are not likely to support other norm\ntemplate \nstruct MaxNormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<1>& tl,\n float* output, float* output_per_tensor, bool per_tensor,\n int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t* x = (x_t*)tl.addresses[0][tensor_loc];\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]);\n vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]);\n vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val = fmaxf(fabsf(val), fabsf(vals[i]));\n\n float final = reduce_block_into_lanes_max_op(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final)) *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] = fmaxf(fabsf(output[blockIdx.x]), fabsf(final));\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;\n }\n }\n};\n\n__global__ void cleanup_v3(float* output, float* output_per_tensor, float* ret, float* ret_per_tensor, bool per_tensor,\n int max_chunks_per_tensor) {\n __shared__ float vals[512];\n\n if (blockIdx.x == 0) {\n float val = 0;\n if (threadIdx.x < 320) val = output[threadIdx.x];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) *ret = sqrt(final);\n }\n\n if (per_tensor) {\n float* output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;\n\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x) val += output_this_tensor[i];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) ret_per_tensor[blockIdx.x] = sqrt(final);\n }\n}\n\nstd::tuple multi_tensor_l2norm_scale_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists,\n float scale, at::optional per_tensor_python) {\n bool per_tensor = per_tensor_python.has_value() ? per_tensor_python.value() : false;\n\n auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);\n auto output = at::zeros({320}, float_options);\n\n at::Tensor output_per_tensor;\n at::Tensor ret_per_tensor;\n\n int ntensors = tensor_lists[0].size();\n int max_chunks_per_tensor = -1;\n\n if (per_tensor) {\n for (int t = 0; t < ntensors; t++) {\n int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;\n if (max_chunks_this_tensor > max_chunks_per_tensor) max_chunks_per_tensor = max_chunks_this_tensor;\n }\n output_per_tensor = at::zeros({ntensors * max_chunks_per_tensor}, float_options);\n ret_per_tensor = at::empty({ntensors}, float_options);\n } else {\n ret_per_tensor = at::empty({0}, float_options);\n }\n\n DISPATCH_FLOAT_AND_HALF(\n tensor_lists[0][0].scalar_type(), 0, \"multi_tensor_l2norm_scale_cuda\",\n DISPATCH_FLOAT_AND_HALF(\n tensor_lists[1][0].scalar_type(), 1, \"multi_tensor_l2norm_scale_cuda\",\n multi_tensor_apply<2>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n L2NormScaleFunctor(), output.data_ptr(),\n per_tensor ? output_per_tensor.data_ptr() : nullptr, scale, per_tensor,\n max_chunks_per_tensor);))\n\n AT_CUDA_CHECK(cudaGetLastError());\n // AT_CUDA_CHECK(cudaDeviceSynchronize());\n\n // This involves one more small kernel launches, but will be negligible end to end.\n // I could get rid of these by hacking the functor + multi tensor harness with persistence\n // logic, but keeping it simple for now\n auto ret = at::empty({1}, output.options());\n const at::cuda::OptionalCUDAGuard device_guard(device_of(output));\n auto stream = at::cuda::getCurrentCUDAStream();\n cleanup_v3<<>>(\n output.data_ptr(), per_tensor ? output_per_tensor.data_ptr() : nullptr, ret.data_ptr(),\n per_tensor ? ret_per_tensor.data_ptr() : nullptr, per_tensor, max_chunks_per_tensor);\n\n return std::tuple(ret, ret_per_tensor);\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n#include \n// Another possibility:\n// #include \n\n#include \n\n#include \"multi_tensor_apply.cuh\"\n#include \"type_shim.h\"\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntemplate \n__device__ __forceinline__ bool is_aligned(T* p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\n\ntemplate \nstruct L2NormScaleFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<2>& tl,\n float* output, float* output_per_tensor, float scale, bool per_tensor,\n int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n in_t* in = (in_t*)tl.addresses[0][tensor_loc];\n in += chunk_idx * chunk_size;\n\n out_t* out = (out_t*)tl.addresses[1][tensor_loc];\n out += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n in_t r_in[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_in[i] = 0;\n }\n // bool finite = true;\n out_t r_out[ILP];\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(in) && is_aligned(out)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n load_store(r_in, in, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_in[ii]);\n r_out[ii] = next * scale;\n vals[ii] += next * next;\n // finite = finite && isfinite(r_in[ii]);\n }\n load_store(out, r_out, i_start, 0);\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_in[ii] = 0;\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n r_in[ii] = in[i];\n float next = static_cast(in[i]);\n vals[ii] += next * next;\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_out[ii] = static_cast(r_in[ii]) * scale;\n // finite = finite && isfinite(r_in[ii]);\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) out[i] = r_out[ii];\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val += vals[i];\n\n float final = reduce_block_into_lanes(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final)) *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] += final;\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;\n }\n }\n};\n// Probably better to template, but since we are not likely to support other norm\ntemplate \nstruct MaxNormFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<1>& tl,\n float* output, float* output_per_tensor, bool per_tensor,\n int max_chunks_per_tensor) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n x_t* x = (x_t*)tl.addresses[0][tensor_loc];\n x += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n __shared__ float s_vals[512];\n\n float vals[ILP]; // = {0}; // this probably works too but I want to be sure...\n x_t r_x[ILP];\n for (int i = 0; i < ILP; i++) {\n vals[i] = 0.f;\n r_x[i] = 0;\n }\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n load_store(r_x, x, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float next = static_cast(r_x[ii]);\n vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));\n }\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n float next = static_cast(x[i]);\n vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));\n }\n }\n }\n }\n\n float val = 0.f;\n for (int i = 0; i < ILP; i++) val = fmaxf(fabsf(val), fabsf(vals[i]));\n\n float final = reduce_block_into_lanes_max_op(s_vals, val);\n\n if (threadIdx.x == 0) {\n if (!isfinite(final)) *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.\n output[blockIdx.x] = fmaxf(fabsf(output[blockIdx.x]), fabsf(final));\n if (per_tensor)\n output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;\n }\n }\n};\n\n__global__ void cleanup_v3(float* output, float* output_per_tensor, float* ret, float* ret_per_tensor, bool per_tensor,\n int max_chunks_per_tensor) {\n __shared__ float vals[512];\n\n if (blockIdx.x == 0) {\n float val = 0;\n if (threadIdx.x < 320) val = output[threadIdx.x];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) *ret = sqrt(final);\n }\n\n if (per_tensor) {\n float* output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;\n\n float val = 0;\n for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x) val += output_this_tensor[i];\n\n float final = reduce_block_into_lanes(vals, val);\n\n if (threadIdx.x == 0) ret_per_tensor[blockIdx.x] = sqrt(final);\n }\n}\n\nstd::tuple multi_tensor_l2norm_scale_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists,\n float scale, at::optional per_tensor_python) {\n bool per_tensor = per_tensor_python.has_value() ? per_tensor_python.value() : false;\n\n auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);\n auto output = at::zeros({320}, float_options);\n\n at::Tensor output_per_tensor;\n at::Tensor ret_per_tensor;\n\n int ntensors = tensor_lists[0].size();\n int max_chunks_per_tensor = -1;\n\n if (per_tensor) {\n for (int t = 0; t < ntensors; t++) {\n int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;\n if (max_chunks_this_tensor > max_chunks_per_tensor) max_chunks_per_tensor = max_chunks_this_tensor;\n }\n output_per_tensor = at::zeros({ntensors * max_chunks_per_tensor}, float_options);\n ret_per_tensor = at::empty({ntensors}, float_options);\n } else {\n ret_per_tensor = at::empty({0}, float_options);\n }\n\n DISPATCH_FLOAT_AND_HALF(\n tensor_lists[0][0].scalar_type(), 0, \"multi_tensor_l2norm_scale_cuda\",\n DISPATCH_FLOAT_AND_HALF(\n tensor_lists[1][0].scalar_type(), 1, \"multi_tensor_l2norm_scale_cuda\",\n multi_tensor_apply<2>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n L2NormScaleFunctor(), output.data_ptr(),\n per_tensor ? output_per_tensor.data_ptr() : nullptr, scale, per_tensor,\n max_chunks_per_tensor);))\n\n AT_CUDA_CHECK(hipGetLastError());\n // AT_CUDA_CHECK(cudaDeviceSynchronize());\n\n // This involves one more small kernel launches, but will be negligible end to end.\n // I could get rid of these by hacking the functor + multi tensor harness with persistence\n // logic, but keeping it simple for now\n auto ret = at::empty({1}, output.options());\n const at::cuda::OptionalCUDAGuard device_guard(device_of(output));\n auto stream = at::cuda::getCurrentCUDAStream();\n cleanup_v3<<>>(\n output.data_ptr(), per_tensor ? output_per_tensor.data_ptr() : nullptr, ret.data_ptr(),\n per_tensor ? ret_per_tensor.data_ptr() : nullptr, per_tensor, max_chunks_per_tensor);\n\n return std::tuple(ret, ret_per_tensor);\n}\n", "source": "/kaggle/working/apex/csrc/multi_tensor_l2norm_scale_kernel.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"type_shim.h\"\n\n__device__ __forceinline__ int lastpow2(int n) {\n int out = 1 << (31 - __clz(n));\n if (n == out) out >>= 1;\n return out;\n}\n\n__host__ __forceinline__ int h_next_pow2(unsigned int n) {\n n--;\n n |= (n >> 1);\n n |= (n >> 2);\n n |= (n >> 4);\n n |= (n >> 8);\n n |= (n >> 16);\n return ++n;\n}\n\n__host__ __forceinline__ int h_last_pow2(unsigned int n) {\n n |= (n >> 1);\n n |= (n >> 2);\n n |= (n >> 4);\n n |= (n >> 8);\n n |= (n >> 16);\n return n - (n >> 1);\n}\n\n#define WARP_SIZE 32\n\ntemplate \n__device__ __forceinline__ T warp_reduce_sum(T val) {\n#pragma unroll\n for (int i = WARP_SIZE / 2; i > 0; i >>= 1) val = val + __shfl_down_sync(0xffffffff, val, i);\n return val;\n}\n\ntemplate \n__device__ __forceinline__ T reduce_block(T* x, T val) {\n int tid = threadIdx.y * blockDim.x + threadIdx.x;\n int blockSize = blockDim.x * blockDim.y;\n\n if (blockSize > 32) {\n val = warp_reduce_sum(val);\n if (tid % WARP_SIZE == 0) x[tid / WARP_SIZE] = val;\n\n __syncthreads();\n\n val = (tid < blockSize / WARP_SIZE ? x[tid % WARP_SIZE] : T(0));\n }\n\n if (tid / WARP_SIZE == 0) val = warp_reduce_sum(val);\n\n return val;\n}\n\n#define ELEMENTS_PER_ITER 4 // enables concurrency within each thread to hide latency\n#define ELEMENTS_PER_THREAD 16\n#define OPTIMAL_TILE_W 32\n#define MAX_H_BLOCK 128\n#define MAX_BLOCK_SIZE 512\n\n__host__ int div_ru(int x, int y) { return h_last_pow2(1 + (x - 1) / y); }\n\n__host__ void flexible_launch_configs(const int reduction, const int stride, dim3& block, dim3& grid,\n const bool coop_flag = false) {\n int block_x = std::min(h_last_pow2(stride), OPTIMAL_TILE_W);\n int block_y = std::min(h_last_pow2(div_ru(reduction, ELEMENTS_PER_THREAD)), MAX_BLOCK_SIZE / block_x);\n if (block_x * block_y != MAX_BLOCK_SIZE) {\n block_x = std::min(h_last_pow2(stride), MAX_BLOCK_SIZE / block_y);\n }\n\n int grid_x = div_ru(stride, block_x);\n int grid_y = std::min(div_ru(reduction, block_y * ELEMENTS_PER_THREAD), MAX_H_BLOCK);\n if (coop_flag) {\n // it's not worth having a grid reduction if the reduction dimension is not big enough\n grid_y = grid_y < 8 ? 1 : grid_y;\n }\n\n block.x = block_x;\n block.y = block_y;\n block.z = 1;\n grid.x = grid_x;\n grid.y = grid_y;\n grid.z = 1;\n}\n\ntemplate \n__device__ __forceinline__ void welford_merge_element(C& count, T& mean, T& m2n, const C& num_new, const T& mean_new,\n const T& m2n_new) {\n T factor = T(1.0) / max(1, (count + num_new));\n T delta0 = mean - mean_new;\n mean = (mean_new * num_new + mean * count) * factor;\n m2n += m2n_new + delta0 * delta0 * num_new * count * factor;\n count += num_new;\n}\n\ntemplate \n__device__ __forceinline__ void warp_reduce_mean_m2n(T& mean, T& m2n, int& num) {\n#pragma unroll\n for (int i = WARP_SIZE / 2; i > 0; i >>= 1) {\n auto num_new = __shfl_down_sync(0xffffffff, num, i);\n auto mean_new = __shfl_down_sync(0xffffffff, mean, i);\n auto m2n_new = __shfl_down_sync(0xffffffff, m2n, i);\n welford_merge_element(num, mean, m2n, num_new, mean_new, m2n_new);\n }\n}\n\ntemplate \n__device__ void welford_reduce_mean_m2n(T* __restrict__ x, int* __restrict__ count, T& mean, T& m2n, int& num,\n int block_size, int thread_id) {\n int lane = thread_id % WARP_SIZE;\n int wid = thread_id / WARP_SIZE;\n\n if (block_size > 32) {\n warp_reduce_mean_m2n(mean, m2n, num);\n if (lane == 0) {\n x[wid * 2] = mean;\n x[wid * 2 + 1] = m2n;\n count[wid] = num;\n }\n __syncthreads();\n\n if (wid == 0) {\n mean = (thread_id < block_size / WARP_SIZE) ? x[lane * 2] : T(0);\n m2n = (thread_id < block_size / WARP_SIZE) ? x[lane * 2 + 1] : T(0);\n num = (thread_id < block_size / WARP_SIZE) ? count[lane] : int(0);\n }\n }\n\n if (wid == 0) warp_reduce_mean_m2n(mean, m2n, num);\n\n return;\n}\n\n// return spatial size for NC+ Tensors\n__host__ int get_tensor_spatial_size(const at::Tensor& input) {\n auto space_size = input.size(2);\n for (int i = 3; i < input.ndimension(); i++) {\n space_size *= input.size(i);\n }\n return space_size;\n}\n\n// promote accumulation scalar type. promote half to float.\n__host__ at::ScalarType promote_scalartype(const at::Tensor& input) {\n return input.scalar_type() == at::ScalarType::Half ? at::ScalarType::Float : input.scalar_type();\n}\n\n// return single element size, optional accumulation type promotion.\n__host__ size_t get_element_data_size(const at::Tensor& input, bool accumulation = false) {\n auto scalar_type = accumulation ? promote_scalartype(input) : input.scalar_type();\n return at::elementSize(scalar_type);\n}\n\ntemplate \n__device__ __forceinline__ void welford_merge_block_vertical(C& count, T& mean, T& m2n, C* shmem_count, T* shmem_mean,\n T* shmem_m2n) {\n // write to shared memory\n auto address_base = threadIdx.x + threadIdx.y * blockDim.x;\n shmem_mean[address_base] = mean;\n shmem_m2n[address_base] = m2n;\n shmem_count[address_base] = count;\n\n#pragma unroll\n for (int offset = blockDim.y / 2; offset > 0; offset >>= 1) {\n __syncthreads();\n if (threadIdx.y < offset && threadIdx.y + offset < blockDim.y) {\n auto address = address_base + offset * blockDim.x;\n // read shared memory back to register for reduction\n auto num_new = shmem_count[address];\n auto mean_new = shmem_mean[address];\n auto m2n_new = shmem_m2n[address];\n\n welford_merge_element(count, mean, m2n, num_new, mean_new, m2n_new);\n\n // last write is not necessary\n shmem_mean[address_base] = mean;\n shmem_m2n[address_base] = m2n;\n shmem_count[address_base] = count;\n }\n }\n}\n\ntemplate \n__device__ __forceinline__ void merge_block_vertical(T& sum_dy, T& sum_dy_xmu, T* shmem_sum_dy, T* shmem_sum_dy_xmu) {\n // write to shared memory\n auto address_base = threadIdx.x + threadIdx.y * blockDim.x;\n shmem_sum_dy[address_base] = sum_dy;\n shmem_sum_dy_xmu[address_base] = sum_dy_xmu;\n\n#pragma unroll\n for (int offset = blockDim.y / 2; offset > 0; offset >>= 1) {\n __syncthreads();\n if (threadIdx.y < offset && threadIdx.y + offset < blockDim.y) {\n auto address = address_base + offset * blockDim.x;\n\n sum_dy += shmem_sum_dy[address];\n sum_dy_xmu += shmem_sum_dy_xmu[address];\n\n // last write is not necessary\n shmem_sum_dy[address_base] = sum_dy;\n shmem_sum_dy_xmu[address_base] = sum_dy_xmu;\n }\n }\n}\n\n// welford kernel calculating mean/biased_variance/unbiased_variance\ntemplate \n__global__ void welford_kernel(const scalar_t* __restrict__ input, outscalar_t* __restrict__ out_mean,\n outscalar_t* __restrict__ out_var_biased, const int bs, const int fs, const int ss) {\n int block_size = blockDim.x * blockDim.y;\n int count = 0;\n accscalar_t x_mean = accscalar_t(0);\n accscalar_t m_2_n = accscalar_t(0);\n\n int thread_id = threadIdx.y * blockDim.x + threadIdx.x;\n\n for (int batch_id = threadIdx.y; batch_id < bs; batch_id += blockDim.y) {\n int input_base = blockIdx.x * ss + batch_id * ss * fs;\n // sequential welford\n for (int offset = threadIdx.x; offset < ss; offset += blockDim.x) {\n count++;\n auto x_n = static_cast(input[offset + input_base]);\n auto d = x_n - x_mean;\n x_mean += d / count;\n m_2_n += d * (x_n - x_mean);\n }\n }\n\n static __shared__ int s_mem[160];\n accscalar_t* s_mem_ac = (accscalar_t*)&s_mem[32];\n\n welford_reduce_mean_m2n(s_mem_ac, s_mem, x_mean, m_2_n, count, block_size, thread_id);\n\n if (thread_id == 0) {\n out_mean[blockIdx.x] = static_cast(x_mean);\n out_var_biased[blockIdx.x] = static_cast(m_2_n / count);\n }\n}\n\n// elementwise BN kernel\ntemplate \n__global__ void batchnorm_forward_kernel(const scalar_t* __restrict__ input, const accscalar_t* __restrict__ mean,\n const accscalar_t* __restrict__ inv_std,\n const layerscalar_t* __restrict__ weight,\n const layerscalar_t* __restrict__ shift, scalar_t* __restrict__ out,\n const int ss, const int bs) {\n auto m_c = mean[blockIdx.x];\n auto inv_std_c = inv_std[blockIdx.x];\n auto w_c = weight == NULL ? accscalar_t(1.0) : static_cast(weight[blockIdx.x]);\n auto s_c = shift == NULL ? accscalar_t(0.0) : static_cast(shift[blockIdx.x]);\n\n for (int batch_offset = blockIdx.y * blockDim.y + threadIdx.y; batch_offset < bs;\n batch_offset += gridDim.y * blockDim.y) {\n int address_base = blockIdx.x * ss + batch_offset * gridDim.x * ss;\n for (int offset = threadIdx.x + blockIdx.z * blockDim.x; offset < ss; offset += gridDim.z * blockDim.x) {\n out[address_base + offset] =\n static_cast(w_c * (static_cast(input[address_base + offset]) - m_c) * inv_std_c + s_c);\n }\n }\n}\n\n// Backward BN kernel, calculates grad_bias, grad_weight as well as intermediate\n// results to calculating grad_input.\n// Breaking the grad_input to two step to support sync BN, which requires all\n// reduce of the intermediate results across processes.\ntemplate \n__global__ void reduce_bn_kernel(const scalar_t* __restrict__ input, const scalar_t* __restrict__ grad_output,\n const accscalar_t* __restrict__ mean, const accscalar_t* __restrict__ inv_std,\n accscalar_t* __restrict__ sum_dy_o, accscalar_t* __restrict__ sum_dy_xmu_o,\n layerscalar_t* __restrict__ grad_weight, layerscalar_t* __restrict__ grad_bias,\n const int bs, const int fs, const int ss) {\n static __shared__ int s_mem[64];\n // int total_item_num = bs * ss;\n\n int thread_id = threadIdx.y * blockDim.x + threadIdx.x;\n\n auto r_mean = mean[blockIdx.x];\n auto factor = inv_std[blockIdx.x];\n\n // Kahan sum\n accscalar_t sum_dy = 0.0;\n accscalar_t sum_dy_xmu = 0.0;\n accscalar_t sum_dy_c = 0.0;\n accscalar_t sum_dy_xmu_c = 0.0;\n for (int batch_id = threadIdx.y; batch_id < bs; batch_id += blockDim.y) {\n int input_base = blockIdx.x * ss + batch_id * ss * fs;\n for (int offset = threadIdx.x; offset < ss; offset += blockDim.x) {\n auto e_grad = static_cast(grad_output[offset + input_base]);\n auto e_input = static_cast(input[offset + input_base]);\n // calculating sum_dy\n auto sum_dy_y = e_grad - sum_dy_c;\n auto sum_dy_t = sum_dy + sum_dy_y;\n sum_dy_c = (sum_dy_t - sum_dy) - sum_dy_y;\n sum_dy = sum_dy_t;\n\n // calculating sum_dy_xmu\n auto sum_dy_xmu_y = e_grad * (e_input - r_mean) - sum_dy_xmu_c;\n auto sum_dy_xmu_t = sum_dy_xmu + sum_dy_xmu_y;\n sum_dy_xmu_c = (sum_dy_xmu_t - sum_dy_xmu) - sum_dy_xmu_y;\n sum_dy_xmu = sum_dy_xmu_t;\n }\n }\n\n sum_dy = reduce_block((accscalar_t*)s_mem, sum_dy);\n __syncthreads();\n sum_dy_xmu = reduce_block((accscalar_t*)s_mem, sum_dy_xmu);\n\n if (thread_id == 0) {\n if (grad_bias != NULL) {\n grad_bias[blockIdx.x] = static_cast(sum_dy);\n }\n if (grad_weight != NULL) {\n grad_weight[blockIdx.x] = static_cast(sum_dy_xmu * factor);\n }\n // mean_dy[blockIdx.x] = sum_dy / total_item_num;\n // mean_dy_xmu[blockIdx.x] = sum_dy_xmu / total_item_num;\n sum_dy_o[blockIdx.x] = sum_dy;\n sum_dy_xmu_o[blockIdx.x] = sum_dy_xmu;\n }\n}\n\n// elementwise backward BN kernel\ntemplate \n__global__ void batchnorm_backward_kernel(const scalar_t* __restrict__ grad_output, const scalar_t* __restrict__ input,\n const accscalar_t* __restrict__ mean, const accscalar_t* __restrict__ inv_std,\n const layerscalar_t* __restrict__ weight,\n const accscalar_t* __restrict__ sum_dy,\n const accscalar_t* __restrict__ sum_dy_xmu, const int* __restrict__ numel,\n scalar_t* __restrict__ grad_input, const int64_t world_size, const int ss,\n const int bs) {\n int64_t div = 0;\n for (int i = 0; i < world_size; i++) {\n div += numel[i];\n }\n auto m_c = static_cast(mean[blockIdx.x]);\n // auto m_dy_c = static_cast(mean_dy[blockIdx.x]);\n auto m_dy_c = static_cast(sum_dy[blockIdx.x]) / div;\n auto factor_1_c = inv_std[blockIdx.x];\n auto factor_2_c = (weight == NULL ? accscalar_t(1.0) : static_cast(weight[blockIdx.x])) * factor_1_c;\n // factor_1_c = factor_1_c * factor_1_c * mean_dy_xmu[blockIdx.x];\n factor_1_c = factor_1_c * factor_1_c * sum_dy_xmu[blockIdx.x] / div;\n\n for (int batch_offset = blockIdx.y * blockDim.y + threadIdx.y; batch_offset < bs;\n batch_offset += gridDim.y * blockDim.y) {\n int address_base = blockIdx.x * ss + batch_offset * gridDim.x * ss;\n for (int offset = threadIdx.x + blockIdx.z * blockDim.x; offset < ss; offset += gridDim.z * blockDim.x) {\n grad_input[address_base + offset] =\n (static_cast(grad_output[address_base + offset]) - m_dy_c -\n (static_cast(input[address_base + offset]) - m_c) * factor_1_c) *\n factor_2_c;\n }\n }\n}\n\n// welford kernel for c last tensor calculating mean/biased_variance/unbiased_variance\ntemplate \n__global__ void welford_kernel_c_last(const scalar_t* __restrict__ input, outscalar_t* __restrict__ out_mean,\n outscalar_t* __restrict__ out_var_biased, volatile accscalar_t* staging_data,\n int* semaphores, const int reduction_size, const int stride) {\n // hide latency with concurrency\n accscalar_t x_mean[PARALLEL_LOADS];\n accscalar_t m_2_n[PARALLEL_LOADS];\n int count[PARALLEL_LOADS];\n\n#pragma unroll\n for (int i = 0; i < PARALLEL_LOADS; i++) {\n x_mean[i] = accscalar_t(0);\n m_2_n[i] = accscalar_t(0);\n count[i] = accscalar_t(0);\n }\n // tensor dimension (m,c)\n\n // loop along m dimension\n int inner_loop_stride = blockDim.y * gridDim.y;\n\n // offset along m dimension\n int m_offset = blockIdx.y * blockDim.y + threadIdx.y;\n int c_offset = blockIdx.x * blockDim.x + threadIdx.x;\n\n int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);\n int address_base = m_offset * stride + c_offset;\n int address_increment = inner_loop_stride * stride;\n\n for (int i = 0; i < loop_count; i++) {\n accscalar_t x_math[PARALLEL_LOADS];\n accscalar_t x_count_inv[PARALLEL_LOADS];\n accscalar_t is_valid[PARALLEL_LOADS];\n\n // load multiple data in\n#pragma unroll\n for (int j = 0; j < PARALLEL_LOADS; j++) {\n if (c_offset < stride && m_offset < reduction_size) {\n x_math[j] = input[address_base];\n count[j]++;\n x_count_inv[j] = accscalar_t(1) / count[j];\n is_valid[j] = accscalar_t(1);\n } else {\n x_math[j] = accscalar_t(0);\n x_count_inv[j] = accscalar_t(0);\n is_valid[j] = accscalar_t(0);\n }\n m_offset += inner_loop_stride;\n address_base += address_increment;\n }\n\n // calculate mean/m2n with welford\n#pragma unroll\n for (int j = 0; j < PARALLEL_LOADS; j++) {\n accscalar_t delta0 = x_math[j] - x_mean[j];\n x_mean[j] += delta0 * x_count_inv[j];\n accscalar_t delta1 = x_math[j] - x_mean[j];\n m_2_n[j] += delta0 * delta1 * is_valid[j];\n }\n }\n\n // thread reduction to accumulate mean/m_2_n/count between PARALLEL_LOADS\n#pragma unroll\n for (int j = 1; j < PARALLEL_LOADS; j++) {\n welford_merge_element(count[0], x_mean[0], m_2_n[0], count[j], x_mean[j], m_2_n[j]);\n }\n\n // release x_mean / m_2_n\n auto mean_th = x_mean[0];\n auto m2_th = m_2_n[0];\n auto count_th = count[0];\n\n // block-wise reduction with shared memory (since reduction cannot be done within a warp)\n static __shared__ accscalar_t shmem_mean[MAX_BLOCK_SIZE];\n static __shared__ accscalar_t shmem_m2n[MAX_BLOCK_SIZE];\n static __shared__ int shmem_count[MAX_BLOCK_SIZE];\n\n welford_merge_block_vertical(count_th, mean_th, m2_th, shmem_count, shmem_mean, shmem_m2n);\n\n // grid reduction if needed (coop launch used at the first place)\n if (gridDim.y > 1) {\n volatile accscalar_t* staging_mean = staging_data;\n volatile accscalar_t* staging_m2n = &staging_data[stride * gridDim.y];\n volatile int* staging_count = reinterpret_cast(&staging_m2n[stride * gridDim.y]);\n\n address_base = c_offset + blockIdx.y * stride;\n // write data to staging_data;\n if (threadIdx.y == 0 && c_offset < stride) {\n staging_mean[address_base] = mean_th;\n staging_m2n[address_base] = m2_th;\n staging_count[address_base] = count_th;\n }\n\n __threadfence();\n __syncthreads(); // ensuring writes to staging_ is visible to all blocks\n\n __shared__ bool is_last_block_done;\n // mark block done\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n int old = atomicAdd(&semaphores[blockIdx.x], 1);\n is_last_block_done = (old == (gridDim.y - 1));\n }\n\n __syncthreads();\n\n // check that all data is now available in global memory\n if (is_last_block_done) {\n count_th = 0;\n mean_th = accscalar_t(0.0);\n m2_th = accscalar_t(0.0);\n\n for (int y = threadIdx.y; y < gridDim.y; y += blockDim.y) {\n address_base = c_offset + y * stride;\n int num_new = c_offset < stride ? staging_count[address_base] : 0;\n accscalar_t mean_new = c_offset < stride ? staging_mean[address_base] : accscalar_t(0.0);\n accscalar_t m2n_new = c_offset < stride ? staging_m2n[address_base] : accscalar_t(0.0);\n\n welford_merge_element(count_th, mean_th, m2_th, num_new, mean_new, m2n_new);\n }\n\n welford_merge_block_vertical(count_th, mean_th, m2_th, shmem_count, shmem_mean, shmem_m2n);\n if (threadIdx.y == 0 && c_offset < stride) {\n out_mean[c_offset] = static_cast(mean_th);\n out_var_biased[c_offset] = static_cast(m2_th / count_th);\n }\n }\n } else {\n if (blockIdx.y == 0 && threadIdx.y == 0 && c_offset < stride) {\n out_mean[c_offset] = static_cast(mean_th);\n out_var_biased[c_offset] = static_cast(m2_th / count_th);\n }\n }\n}\n\n// parallel welford kernel to further reduce mean / biased_var\n// into mean / unbiased_var / inv_std across multiple processes.\ntemplate \n__global__ void welford_kernel_parallel(const scalar_t* __restrict__ mean, const scalar_t* __restrict__ var_biased,\n const int* __restrict__ numel, scalar_t* __restrict__ out_mean,\n scalar_t* __restrict__ out_var, scalar_t* __restrict__ inv_std,\n const int world_size, const int feature_size, const float eps) {\n for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < feature_size; i += gridDim.x * blockDim.x) {\n // load data;\n int address = i;\n scalar_t x_mean = 0;\n scalar_t m_2_n = 0;\n int count = 0;\n for (int j = 0; j < world_size; j++) {\n welford_merge_element(count, x_mean, m_2_n, numel[j], mean[address], var_biased[address] * numel[j]);\n address += feature_size;\n }\n out_mean[i] = x_mean;\n out_var[i] = m_2_n / (count - 1);\n inv_std[i] = scalar_t(1) / sqrt(m_2_n / count + eps);\n }\n}\n\n// elementwise BN kernel\ntemplate \n__global__ void batchnorm_forward_c_last_kernel(const scalar_t* __restrict__ input, const scalar_t* __restrict__ z,\n const accscalar_t* __restrict__ mean,\n const accscalar_t* __restrict__ inv_std,\n const layerscalar_t* __restrict__ weight,\n const layerscalar_t* __restrict__ shift, scalar_t* __restrict__ out,\n const int reduction_size, const int stride, const bool fuse_relu) {\n // tensor dimension (m,c)\n // loop along m dimension\n int inner_loop_stride = blockDim.y * gridDim.y;\n\n // offset along m dimension\n int m_offset = blockIdx.y * blockDim.y + threadIdx.y;\n int c_offset = blockIdx.x * blockDim.x + threadIdx.x;\n\n auto m_c = mean[c_offset];\n auto inv_std_c = static_cast(inv_std[c_offset]);\n auto w_c = weight == NULL ? accscalar_t(1.0) : static_cast(weight[c_offset]);\n auto s_c = shift == NULL ? accscalar_t(0.0) : static_cast(shift[c_offset]);\n\n int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);\n int address_base = m_offset * stride + c_offset;\n int address_increment = inner_loop_stride * stride;\n\n for (int i = 0; i < loop_count; i++) {\n#pragma unroll\n for (int j = 0; j < PARALLEL_LOADS; j++) {\n if (c_offset < stride && m_offset < reduction_size) {\n auto tmp = w_c * (static_cast(input[address_base]) - m_c) * inv_std_c + s_c;\n if (z != NULL) {\n tmp += z[address_base];\n }\n out[address_base] = (fuse_relu && tmp <= accscalar_t(0.0) ? scalar_t(0.0) : static_cast(tmp));\n }\n m_offset += inner_loop_stride;\n address_base += address_increment;\n }\n }\n}\n\n// elementwise BN kernel\ntemplate \n__global__ void relu_backward_c_last_kernel(const scalar_t* __restrict__ grad_output,\n const scalar_t* __restrict__ input, const scalar_t* __restrict__ z,\n const accscalar_t* __restrict__ mean,\n const accscalar_t* __restrict__ inv_std,\n const layerscalar_t* __restrict__ weight,\n const layerscalar_t* __restrict__ shift, scalar_t* __restrict__ out,\n const int reduction_size, const int stride) {\n // tensor dimension (m,c)\n // loop along m dimension\n int inner_loop_stride = blockDim.y * gridDim.y;\n\n // offset along m dimension\n int m_offset = blockIdx.y * blockDim.y + threadIdx.y;\n int c_offset = blockIdx.x * blockDim.x + threadIdx.x;\n\n auto m_c = mean[c_offset];\n auto inv_std_c = static_cast(inv_std[c_offset]);\n auto w_c = weight == NULL ? accscalar_t(1.0) : static_cast(weight[c_offset]);\n auto s_c = shift == NULL ? accscalar_t(0.0) : static_cast(shift[c_offset]);\n\n int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);\n int address_base = m_offset * stride + c_offset;\n int address_increment = inner_loop_stride * stride;\n\n for (int i = 0; i < loop_count; i++) {\n#pragma unroll\n for (int j = 0; j < PARALLEL_LOADS; j++) {\n if (c_offset < stride && m_offset < reduction_size) {\n auto tmp = w_c * (static_cast(input[address_base]) - m_c) * inv_std_c + s_c;\n if (z != NULL) {\n tmp += z[address_base];\n }\n out[address_base] = (tmp <= accscalar_t(0.0) ? scalar_t(0.0) : grad_output[address_base]);\n }\n m_offset += inner_loop_stride;\n address_base += address_increment;\n }\n }\n}\n\n// batchnorm backward kernel for c last tensor\ntemplate \n__global__ void reduce_bn_c_last_kernel(const scalar_t* __restrict__ input, const scalar_t* __restrict__ grad_output,\n const accscalar_t* __restrict__ mean, const accscalar_t* __restrict__ inv_std,\n accscalar_t* __restrict__ sum_dy_o, accscalar_t* __restrict__ sum_dy_xmu_o,\n layerscalar_t* __restrict__ grad_weight, layerscalar_t* __restrict__ grad_bias,\n volatile accscalar_t* staging_data, int* semaphores, const int reduction_size,\n const int stride) {\n // hide latency with concurrency\n accscalar_t sum_dy[PARALLEL_LOADS];\n accscalar_t sum_dy_xmu[PARALLEL_LOADS];\n\n#pragma unroll\n for (int i = 0; i < PARALLEL_LOADS; i++) {\n sum_dy[i] = accscalar_t(0);\n sum_dy_xmu[i] = accscalar_t(0);\n }\n // tensor dimension (m,c)\n\n // loop along m dimension\n int inner_loop_stride = blockDim.y * gridDim.y;\n\n // offset along m dimension\n int m_offset = blockIdx.y * blockDim.y + threadIdx.y;\n int c_offset = blockIdx.x * blockDim.x + threadIdx.x;\n\n int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);\n int address_base = m_offset * stride + c_offset;\n int address_increment = inner_loop_stride * stride;\n\n auto r_mean = mean[c_offset];\n auto factor = inv_std[c_offset];\n\n for (int i = 0; i < loop_count; i++) {\n accscalar_t x_input[PARALLEL_LOADS];\n accscalar_t x_grad_output[PARALLEL_LOADS];\n\n // load multiple data in\n#pragma unroll\n for (int j = 0; j < PARALLEL_LOADS; j++) {\n if (c_offset < stride && m_offset < reduction_size) {\n x_input[j] = input[address_base];\n x_grad_output[j] = grad_output[address_base];\n } else {\n x_input[j] = accscalar_t(0);\n x_grad_output[j] = accscalar_t(0);\n }\n m_offset += inner_loop_stride;\n address_base += address_increment;\n }\n\n // calculate sum_dy / sum_dy_xmu\n#pragma unroll\n for (int j = 0; j < PARALLEL_LOADS; j++) {\n sum_dy[j] += x_grad_output[j];\n sum_dy_xmu[j] += x_grad_output[j] * (x_input[j] - r_mean);\n }\n }\n\n // thread reduction to accumulate sum_dy / sum_dy_xmu between PARALLEL_LOADS\n#pragma unroll\n for (int j = 1; j < PARALLEL_LOADS; j++) {\n sum_dy[0] += sum_dy[j];\n sum_dy_xmu[0] += sum_dy_xmu[j];\n }\n\n // release array of registers\n auto sum_dy_th = sum_dy[0];\n auto sum_dy_xmu_th = sum_dy_xmu[0];\n\n // block-wise reduction with shared memory (since reduction cannot be done within a warp)\n static __shared__ accscalar_t shmem_sum_dy[MAX_BLOCK_SIZE];\n static __shared__ accscalar_t shmem_sum_dy_xmu[MAX_BLOCK_SIZE];\n\n merge_block_vertical(sum_dy_th, sum_dy_xmu_th, shmem_sum_dy, shmem_sum_dy_xmu);\n\n // grid reduction if needed (coop launch used at the first place)\n if (gridDim.y > 1) {\n volatile accscalar_t* staging_sum_dy = staging_data;\n volatile accscalar_t* staging_sum_dy_xmu = &staging_data[stride * gridDim.y];\n\n address_base = c_offset + blockIdx.y * stride;\n // write data to staging_data;\n if (threadIdx.y == 0 && c_offset < stride) {\n staging_sum_dy[address_base] = sum_dy_th;\n staging_sum_dy_xmu[address_base] = sum_dy_xmu_th;\n }\n\n __threadfence();\n __syncthreads(); // ensuring writes to staging_ is visible to all blocks\n\n __shared__ bool is_last_block_done;\n // mark block done\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n int old = atomicAdd(&semaphores[blockIdx.x], 1);\n is_last_block_done = (old == (gridDim.y - 1));\n }\n\n __syncthreads();\n\n // check that all data is now available in global memory\n if (is_last_block_done) {\n sum_dy_th = accscalar_t(0.0);\n sum_dy_xmu_th = accscalar_t(0.0);\n\n for (int y = threadIdx.y; y < gridDim.y; y += blockDim.y) {\n address_base = c_offset + y * stride;\n sum_dy_th += (c_offset < stride ? staging_sum_dy[address_base] : accscalar_t(0.0));\n sum_dy_xmu_th += (c_offset < stride ? staging_sum_dy_xmu[address_base] : accscalar_t(0.0));\n }\n\n merge_block_vertical(sum_dy_th, sum_dy_xmu_th, shmem_sum_dy, shmem_sum_dy_xmu);\n if (threadIdx.y == 0 && c_offset < stride) {\n if (grad_bias != NULL) {\n grad_bias[c_offset] = static_cast(sum_dy_th);\n }\n if (grad_weight != NULL) {\n grad_weight[c_offset] = static_cast(sum_dy_xmu_th * factor);\n }\n // mean_dy[c_offset] = sum_dy_th / reduction_size;\n // mean_dy_xmu[c_offset] = sum_dy_xmu_th / reduction_size;\n sum_dy_o[c_offset] = sum_dy_th;\n sum_dy_xmu_o[c_offset] = sum_dy_xmu_th;\n }\n }\n } else {\n if (blockIdx.y == 0 && threadIdx.y == 0 && c_offset < stride) {\n if (grad_bias != NULL) {\n grad_bias[c_offset] = static_cast(sum_dy_th);\n }\n if (grad_weight != NULL) {\n grad_weight[c_offset] = static_cast(sum_dy_xmu_th * factor);\n }\n // mean_dy[c_offset] = sum_dy_th / reduction_size;\n // mean_dy_xmu[c_offset] = sum_dy_xmu_th / reduction_size;\n sum_dy_o[c_offset] = sum_dy_th;\n sum_dy_xmu_o[c_offset] = sum_dy_xmu_th;\n }\n }\n}\n\n// elementwise BN kernel\ntemplate \n__global__ void batchnorm_backward_c_last_kernel(\n const scalar_t* __restrict__ grad_output, const scalar_t* __restrict__ input, const accscalar_t* __restrict__ mean,\n const accscalar_t* __restrict__ inv_std, const layerscalar_t* __restrict__ weight,\n const accscalar_t* __restrict__ sum_dy, const accscalar_t* __restrict__ sum_dy_xmu, const int* __restrict__ numel,\n scalar_t* __restrict__ grad_input, const int64_t world_size, const int reduction_size, const int stride) {\n int64_t div = 0;\n for (int i = 0; i < world_size; i++) {\n div += numel[i];\n }\n // tensor dimension (m,c)\n // loop along m dimension\n int inner_loop_stride = blockDim.y * gridDim.y;\n\n // offset along m dimension\n int m_offset = blockIdx.y * blockDim.y + threadIdx.y;\n int c_offset = blockIdx.x * blockDim.x + threadIdx.x;\n\n auto m_c = mean[c_offset];\n auto m_dy_c = sum_dy[c_offset] / div;\n auto factor_1_c = inv_std[c_offset];\n auto factor_2_c = (weight == NULL ? accscalar_t(1.0) : static_cast(weight[c_offset])) * factor_1_c;\n factor_1_c = factor_1_c * factor_1_c * sum_dy_xmu[c_offset] / div;\n\n int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);\n int address_base = m_offset * stride + c_offset;\n int address_increment = inner_loop_stride * stride;\n\n for (int i = 0; i < loop_count; i++) {\n#pragma unroll\n for (int j = 0; j < PARALLEL_LOADS; j++) {\n if (c_offset < stride && m_offset < reduction_size) {\n grad_input[address_base] =\n static_cast((static_cast(grad_output[address_base]) - m_dy_c -\n (static_cast(input[address_base]) - m_c) * factor_1_c) *\n factor_2_c);\n }\n m_offset += inner_loop_stride;\n address_base += address_increment;\n }\n }\n}\n\nstd::vector welford_mean_var_CUDA(const at::Tensor input) {\n const auto batch_size = input.size(0);\n const auto feature_size = input.size(1);\n\n auto space_size = get_tensor_spatial_size(input);\n auto scalar_type = promote_scalartype(input);\n\n at::Tensor out_var_biased = at::empty({feature_size}, input.options().dtype(scalar_type));\n at::Tensor out_mean = at::empty({feature_size}, input.options().dtype(scalar_type));\n\n int block_y = min(h_last_pow2(batch_size), int(MAX_BLOCK_SIZE / 32));\n int block_x = max(1, min(MAX_BLOCK_SIZE / block_y, h_last_pow2(space_size)));\n const dim3 block(block_x, block_y);\n const dim3 grid(feature_size);\n\n auto stream = at::cuda::getCurrentCUDAStream();\n\n {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"welford_mean_var_kernel\", using accscalar_t = at::acc_type;\n welford_kernel\n <<>>(input.data_ptr(), out_mean.data_ptr(),\n out_var_biased.data_ptr(), batch_size, feature_size, space_size););\n }\n\n return {out_mean, out_var_biased};\n}\n\nat::Tensor batchnorm_forward_CUDA(const at::Tensor input, const at::Tensor mean, const at::Tensor inv_std,\n const at::optional weight, const at::optional shift) {\n const auto batch_size = input.size(0);\n const auto feature_size = input.size(1);\n at::Tensor out = at::empty_like(input);\n\n auto space_size = get_tensor_spatial_size(input);\n\n int block_x = max(32, min(MAX_BLOCK_SIZE, h_last_pow2(space_size) / 4));\n int block_y = max(1, min(MAX_BLOCK_SIZE / block_x, h_last_pow2(batch_size) / 4));\n const dim3 block(block_x, block_y);\n int grid_z = max(1, min(65535, h_last_pow2(space_size) / 4 / block_x));\n int batch_group_size = max(1, min(65535, h_last_pow2(batch_size) / block_y));\n const dim3 grid(feature_size, batch_group_size, grid_z);\n auto stream = at::cuda::getCurrentCUDAStream();\n\n if (input.scalar_type() == at::ScalarType::Half && weight.has_value() &&\n weight.value().scalar_type() == at::ScalarType::Float) {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n batchnorm_forward_kernel<<>>(\n input.data_ptr(), mean.data_ptr(), inv_std.data_ptr(),\n weight.has_value() ? weight.value().data_ptr() : NULL,\n shift.has_value() ? shift.value().data_ptr() : NULL, out.data_ptr(), space_size,\n batch_size););\n } else {\n if (weight.has_value()) {\n TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),\n \"input.scalar_type() is not supported with weight.scalar_type()\");\n }\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n batchnorm_forward_kernel<<>>(\n input.data_ptr(), mean.data_ptr(), inv_std.data_ptr(),\n weight.has_value() ? weight.value().data_ptr() : NULL,\n shift.has_value() ? shift.value().data_ptr() : NULL, out.data_ptr(), space_size,\n batch_size););\n }\n return out;\n}\n\nstd::vector reduce_bn_CUDA(const at::Tensor grad_output, const at::Tensor input, const at::Tensor mean,\n const at::Tensor inv_std, const at::optional weight) {\n const auto batch_size = input.size(0);\n const auto feature_size = input.size(1);\n\n auto scalar_type = promote_scalartype(input);\n\n at::Tensor sum_dy = at::empty({feature_size}, mean.options());\n at::Tensor sum_dy_xmu = at::empty({feature_size}, mean.options());\n\n at::Tensor grad_weight;\n at::Tensor grad_bias;\n if (weight.has_value()) {\n grad_weight = at::empty({feature_size}, weight.value().options());\n grad_bias = at::empty({feature_size}, weight.value().options());\n } else {\n grad_weight = at::empty({0}, mean.options());\n grad_bias = at::empty({0}, mean.options());\n }\n\n auto space_size = get_tensor_spatial_size(input);\n\n int block_y = min(h_last_pow2(batch_size), int(MAX_BLOCK_SIZE / 32));\n int block_x = max(1, min(MAX_BLOCK_SIZE / block_y, h_last_pow2(space_size)));\n const dim3 block(block_x, block_y);\n const dim3 grid(feature_size);\n auto stream = at::cuda::getCurrentCUDAStream();\n\n if (input.scalar_type() == at::ScalarType::Half && weight.has_value() &&\n weight.value().scalar_type() == at::ScalarType::Float) {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_backward_reduce\", using accscalar_t = at::acc_type;\n reduce_bn_kernel<<>>(\n input.data_ptr(), grad_output.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), sum_dy.data_ptr(), sum_dy_xmu.data_ptr(),\n weight.has_value() ? grad_weight.data_ptr() : NULL,\n weight.has_value() ? grad_bias.data_ptr() : NULL, batch_size, feature_size, space_size););\n } else {\n if (weight.has_value()) {\n TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),\n \"input.scalar_type() is not supported with weight.scalar_type()\");\n }\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_backward_reduce\", using accscalar_t = at::acc_type;\n reduce_bn_kernel<<>>(\n input.data_ptr(), grad_output.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), sum_dy.data_ptr(), sum_dy_xmu.data_ptr(),\n weight.has_value() ? grad_weight.data_ptr() : NULL,\n weight.has_value() ? grad_bias.data_ptr() : NULL, batch_size, feature_size, space_size););\n }\n\n return {sum_dy, sum_dy_xmu, grad_weight, grad_bias};\n}\n\nat::Tensor batchnorm_backward_CUDA(const at::Tensor grad_output, const at::Tensor input, const at::Tensor mean,\n const at::Tensor inv_std, const at::optional weight,\n const at::Tensor sum_dy, const at::Tensor sum_dy_xmu, const at::Tensor count) {\n const auto batch_size = input.size(0);\n const auto feature_size = input.size(1);\n\n at::Tensor grad_input = at::empty_like(input);\n\n auto space_size = get_tensor_spatial_size(input);\n\n int block_x = max(32, min(MAX_BLOCK_SIZE, h_last_pow2(space_size) / 4));\n int block_y = max(1, min(MAX_BLOCK_SIZE / block_x, h_last_pow2(batch_size) / 4));\n const dim3 block(block_x, block_y);\n int grid_z = max(1, min(65535, h_last_pow2(space_size) / 4 / block_x));\n int batch_group_size = max(1, min(65535, h_last_pow2(batch_size) / block_y));\n const dim3 grid(feature_size, batch_group_size, grid_z);\n\n auto stream = at::cuda::getCurrentCUDAStream();\n\n if (input.scalar_type() == at::ScalarType::Half && weight.has_value() &&\n weight.value().scalar_type() == at::ScalarType::Float) {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_backward\", using accscalar_t = at::acc_type;\n batchnorm_backward_kernel<<>>(\n grad_output.data_ptr(), input.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), weight.has_value() ? weight.value().data_ptr() : NULL,\n sum_dy.data_ptr(), sum_dy_xmu.data_ptr(), count.data_ptr(),\n grad_input.data_ptr(), count.numel(), space_size, batch_size););\n } else {\n if (weight.has_value()) {\n TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),\n \"input.scalar_type() is not supported with weight.scalar_type()\");\n }\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_backward\", using accscalar_t = at::acc_type;\n batchnorm_backward_kernel<<>>(\n grad_output.data_ptr(), input.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), weight.has_value() ? weight.value().data_ptr() : NULL,\n sum_dy.data_ptr(), sum_dy_xmu.data_ptr(), count.data_ptr(),\n grad_input.data_ptr(), count.numel(), space_size, batch_size););\n }\n\n return grad_input;\n}\n\nstd::vector welford_parallel_CUDA(const at::Tensor mean_feature_nodes, const at::Tensor var_biased,\n const at::Tensor numel, const float eps) {\n const auto world_size = mean_feature_nodes.size(0);\n const auto feature_size = mean_feature_nodes.size(1);\n\n at::Tensor out_var = at::empty({feature_size}, var_biased.options());\n at::Tensor inv_std = at::empty_like(out_var);\n at::Tensor out_mean = at::empty_like(out_var);\n\n at::Tensor mean_feature_nodes_ = mean_feature_nodes.contiguous();\n at::Tensor var_biased_ = var_biased.contiguous();\n at::Tensor numel_ = numel.contiguous();\n\n // TODO(jie): tile this for memory coalescing!\n const int block = std::min(h_last_pow2(feature_size), MAX_BLOCK_SIZE);\n const int grid = std::max(1, feature_size / block);\n\n auto stream = at::cuda::getCurrentCUDAStream();\n\n {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(mean_feature_nodes.scalar_type(), 0, \"welford_parallel_kernel\",\n welford_kernel_parallel<<>>(\n mean_feature_nodes_.data_ptr(), var_biased_.data_ptr(),\n numel_.data_ptr(), out_mean.data_ptr(), out_var.data_ptr(),\n inv_std.data_ptr(), world_size, feature_size, eps););\n }\n\n return {out_mean, out_var, inv_std};\n}\n\nstd::vector welford_mean_var_c_last_CUDA(const at::Tensor input) {\n const auto stride = input.size(input.ndimension() - 1);\n const auto reduction_size = input.numel() / stride;\n\n auto scalar_type = promote_scalartype(input);\n auto option = input.options().dtype(scalar_type);\n\n at::Tensor out_var_biased = at::empty({stride}, option);\n at::Tensor out_mean = at::empty({stride}, option);\n\n dim3 block;\n dim3 grid;\n flexible_launch_configs(reduction_size, stride, block, grid, true);\n\n at::Tensor staging_data;\n at::Tensor semaphores;\n if (grid.y > 1) {\n staging_data = at::empty({4 * stride * grid.y}, option);\n semaphores = at::zeros({grid.x}, input.options().dtype(at::kInt));\n }\n\n auto stream = at::cuda::getCurrentCUDAStream();\n\n {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"welford_mean_var_c_last\", using accscalar_t = at::acc_type;\n accscalar_t* staging_data_ptr = grid.y > 1 ? staging_data.data_ptr() : nullptr;\n int* semaphores_ptr = grid.y > 1 ? semaphores.data_ptr() : nullptr;\n welford_kernel_c_last<<>>(\n input.data_ptr(), out_mean.data_ptr(), out_var_biased.data_ptr(),\n staging_data_ptr, semaphores_ptr, reduction_size, stride););\n }\n\n return {out_mean, out_var_biased};\n}\n\nat::Tensor batchnorm_forward_c_last_CUDA(const at::Tensor input, const at::optional z,\n const at::Tensor mean, const at::Tensor inv_std,\n const at::optional weight, const at::optional shift,\n const bool fuse_relu) {\n const auto stride = input.size(input.ndimension() - 1);\n const auto reduction_size = input.numel() / stride;\n\n at::Tensor out = at::empty_like(input);\n\n dim3 block;\n dim3 grid;\n flexible_launch_configs(reduction_size, stride, block, grid);\n\n auto stream = at::cuda::getCurrentCUDAStream();\n\n if (input.scalar_type() == at::ScalarType::Half && weight.has_value() &&\n weight.value().scalar_type() == at::ScalarType::Float) {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n batchnorm_forward_c_last_kernel\n <<>>(input.data_ptr(),\n z.has_value() ? z.value().data_ptr() : NULL,\n mean.data_ptr(), inv_std.data_ptr(),\n weight.has_value() ? weight.value().data_ptr() : NULL,\n shift.has_value() ? shift.value().data_ptr() : NULL,\n out.data_ptr(), reduction_size, stride, fuse_relu););\n } else {\n if (weight.has_value()) {\n TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),\n \"input.scalar_type() is not supported with weight.scalar_type()\");\n }\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n batchnorm_forward_c_last_kernel\n <<>>(input.data_ptr(),\n z.has_value() ? z.value().data_ptr() : NULL,\n mean.data_ptr(), inv_std.data_ptr(),\n weight.has_value() ? weight.value().data_ptr() : NULL,\n shift.has_value() ? shift.value().data_ptr() : NULL,\n out.data_ptr(), reduction_size, stride, fuse_relu););\n }\n return out;\n}\n\nstd::vector reduce_bn_c_last_CUDA(const at::Tensor grad_output, const at::Tensor input,\n const at::Tensor mean, const at::Tensor inv_std,\n const at::optional weight) {\n const auto stride = input.size(input.ndimension() - 1);\n const auto reduction_size = input.numel() / stride;\n\n at::Tensor sumn_dy = at::empty({stride}, mean.options());\n at::Tensor sum_dy_xmu = at::empty({stride}, mean.options());\n\n at::Tensor grad_weight;\n at::Tensor grad_bias;\n if (weight.has_value()) {\n grad_weight = at::empty({stride}, weight.value().options());\n grad_bias = at::empty({stride}, weight.value().options());\n } else {\n // because I cannot return an uninitialized at::Tensor\n grad_weight = at::empty({0}, mean.options());\n grad_bias = at::empty({0}, mean.options());\n }\n\n dim3 block;\n dim3 grid;\n flexible_launch_configs(reduction_size, stride, block, grid, true);\n\n at::Tensor staging_data;\n at::Tensor semaphores;\n if (grid.y > 1) {\n staging_data = at::empty({2 * stride * grid.y}, mean.options());\n semaphores = at::zeros({grid.x}, input.options().dtype(at::kInt));\n }\n auto stream = at::cuda::getCurrentCUDAStream();\n\n if (input.scalar_type() == at::ScalarType::Half && weight.has_value() &&\n weight.value().scalar_type() == at::ScalarType::Float) {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_backward_reduce\", using accscalar_t = at::acc_type;\n accscalar_t* staging_data_ptr = grid.y > 1 ? staging_data.data_ptr() : nullptr;\n int* semaphores_ptr = grid.y > 1 ? semaphores.data_ptr() : nullptr;\n reduce_bn_c_last_kernel<<>>(\n input.data_ptr(), grad_output.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), sumn_dy.data_ptr(), sum_dy_xmu.data_ptr(),\n weight.has_value() ? grad_weight.data_ptr() : NULL,\n weight.has_value() ? grad_bias.data_ptr() : NULL, staging_data_ptr, semaphores_ptr,\n reduction_size, stride););\n } else {\n if (weight.has_value()) {\n TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),\n \"input.scalar_type() is not supported with weight.scalar_type()\");\n }\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_backward_reduce\", using accscalar_t = at::acc_type;\n accscalar_t* staging_data_ptr = grid.y > 1 ? staging_data.data_ptr() : nullptr;\n int* semaphores_ptr = grid.y > 1 ? semaphores.data_ptr() : nullptr;\n reduce_bn_c_last_kernel<<>>(\n input.data_ptr(), grad_output.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), sumn_dy.data_ptr(), sum_dy_xmu.data_ptr(),\n weight.has_value() ? grad_weight.data_ptr() : NULL,\n weight.has_value() ? grad_bias.data_ptr() : NULL, staging_data_ptr, semaphores_ptr,\n reduction_size, stride););\n }\n\n return {sumn_dy, sum_dy_xmu, grad_weight, grad_bias};\n}\n\nat::Tensor batchnorm_backward_c_last_CUDA(const at::Tensor grad_output, const at::Tensor input, const at::Tensor mean,\n const at::Tensor inv_std, const at::optional weight,\n const at::Tensor sum_dy, const at::Tensor sum_dy_xmu,\n const at::Tensor count) {\n const auto stride = input.size(input.ndimension() - 1);\n const auto reduction_size = input.numel() / stride;\n\n at::Tensor grad_input = at::empty_like(input);\n\n dim3 block;\n dim3 grid;\n flexible_launch_configs(reduction_size, stride, block, grid);\n\n auto stream = at::cuda::getCurrentCUDAStream();\n\n if (input.scalar_type() == at::ScalarType::Half && weight.has_value() &&\n weight.value().scalar_type() == at::ScalarType::Float) {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n batchnorm_backward_c_last_kernel\n <<>>(\n grad_output.data_ptr(), input.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), weight.has_value() ? weight.value().data_ptr() : NULL,\n sum_dy.data_ptr(), sum_dy_xmu.data_ptr(), count.data_ptr(),\n grad_input.data_ptr(), count.numel(), reduction_size, stride););\n } else {\n if (weight.has_value()) {\n TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),\n \"input.scalar_type() is not supported with weight.scalar_type()\");\n }\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n batchnorm_backward_c_last_kernel\n <<>>(\n grad_output.data_ptr(), input.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), weight.has_value() ? weight.value().data_ptr() : NULL,\n sum_dy.data_ptr(), sum_dy_xmu.data_ptr(), count.data_ptr(),\n grad_input.data_ptr(), count.numel(), reduction_size, stride););\n }\n\n return grad_input;\n}\n\nat::Tensor relu_backward_c_last_CUDA(const at::Tensor grad_output, const at::Tensor input,\n const at::optional z, const at::Tensor mean, const at::Tensor inv_std,\n const at::optional weight, const at::optional shift) {\n const auto stride = input.size(input.ndimension() - 1);\n const auto reduction_size = input.numel() / stride;\n\n at::Tensor out = at::empty_like(input);\n\n dim3 block;\n dim3 grid;\n flexible_launch_configs(reduction_size, stride, block, grid);\n\n auto stream = at::cuda::getCurrentCUDAStream();\n\n if (input.scalar_type() == at::ScalarType::Half && weight.has_value() &&\n weight.value().scalar_type() == at::ScalarType::Float) {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n relu_backward_c_last_kernel\n <<>>(grad_output.data_ptr(), input.data_ptr(),\n z.has_value() ? z.value().data_ptr() : NULL,\n mean.data_ptr(), inv_std.data_ptr(),\n weight.has_value() ? weight.value().data_ptr() : NULL,\n shift.has_value() ? shift.value().data_ptr() : NULL,\n out.data_ptr(), reduction_size, stride););\n } else {\n if (weight.has_value()) {\n TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),\n \"input.scalar_type() is not supported with weight.scalar_type()\");\n }\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n relu_backward_c_last_kernel<<>>(\n grad_output.data_ptr(), input.data_ptr(),\n z.has_value() ? z.value().data_ptr() : NULL, mean.data_ptr(),\n inv_std.data_ptr(), weight.has_value() ? weight.value().data_ptr() : NULL,\n shift.has_value() ? shift.value().data_ptr() : NULL, out.data_ptr(), reduction_size,\n stride););\n }\n return out;\n}\n", "output": "#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"type_shim.h\"\n\n__device__ __forceinline__ int lastpow2(int n) {\n int out = 1 << (31 - __clz(n));\n if (n == out) out >>= 1;\n return out;\n}\n\n__host__ __forceinline__ int h_next_pow2(unsigned int n) {\n n--;\n n |= (n >> 1);\n n |= (n >> 2);\n n |= (n >> 4);\n n |= (n >> 8);\n n |= (n >> 16);\n return ++n;\n}\n\n__host__ __forceinline__ int h_last_pow2(unsigned int n) {\n n |= (n >> 1);\n n |= (n >> 2);\n n |= (n >> 4);\n n |= (n >> 8);\n n |= (n >> 16);\n return n - (n >> 1);\n}\n\n#define WARP_SIZE 32\n\ntemplate \n__device__ __forceinline__ T warp_reduce_sum(T val) {\n#pragma unroll\n for (int i = WARP_SIZE / 2; i > 0; i >>= 1) val = val + __shfl_down_sync(0xffffffff, val, i);\n return val;\n}\n\ntemplate \n__device__ __forceinline__ T reduce_block(T* x, T val) {\n int tid = threadIdx.y * blockDim.x + threadIdx.x;\n int blockSize = blockDim.x * blockDim.y;\n\n if (blockSize > 32) {\n val = warp_reduce_sum(val);\n if (tid % WARP_SIZE == 0) x[tid / WARP_SIZE] = val;\n\n __syncthreads();\n\n val = (tid < blockSize / WARP_SIZE ? x[tid % WARP_SIZE] : T(0));\n }\n\n if (tid / WARP_SIZE == 0) val = warp_reduce_sum(val);\n\n return val;\n}\n\n#define ELEMENTS_PER_ITER 4 // enables concurrency within each thread to hide latency\n#define ELEMENTS_PER_THREAD 16\n#define OPTIMAL_TILE_W 32\n#define MAX_H_BLOCK 128\n#define MAX_BLOCK_SIZE 512\n\n__host__ int div_ru(int x, int y) { return h_last_pow2(1 + (x - 1) / y); }\n\n__host__ void flexible_launch_configs(const int reduction, const int stride, dim3& block, dim3& grid,\n const bool coop_flag = false) {\n int block_x = std::min(h_last_pow2(stride), OPTIMAL_TILE_W);\n int block_y = std::min(h_last_pow2(div_ru(reduction, ELEMENTS_PER_THREAD)), MAX_BLOCK_SIZE / block_x);\n if (block_x * block_y != MAX_BLOCK_SIZE) {\n block_x = std::min(h_last_pow2(stride), MAX_BLOCK_SIZE / block_y);\n }\n\n int grid_x = div_ru(stride, block_x);\n int grid_y = std::min(div_ru(reduction, block_y * ELEMENTS_PER_THREAD), MAX_H_BLOCK);\n if (coop_flag) {\n // it's not worth having a grid reduction if the reduction dimension is not big enough\n grid_y = grid_y < 8 ? 1 : grid_y;\n }\n\n block.x = block_x;\n block.y = block_y;\n block.z = 1;\n grid.x = grid_x;\n grid.y = grid_y;\n grid.z = 1;\n}\n\ntemplate \n__device__ __forceinline__ void welford_merge_element(C& count, T& mean, T& m2n, const C& num_new, const T& mean_new,\n const T& m2n_new) {\n T factor = T(1.0) / max(1, (count + num_new));\n T delta0 = mean - mean_new;\n mean = (mean_new * num_new + mean * count) * factor;\n m2n += m2n_new + delta0 * delta0 * num_new * count * factor;\n count += num_new;\n}\n\ntemplate \n__device__ __forceinline__ void warp_reduce_mean_m2n(T& mean, T& m2n, int& num) {\n#pragma unroll\n for (int i = WARP_SIZE / 2; i > 0; i >>= 1) {\n auto num_new = __shfl_down_sync(0xffffffff, num, i);\n auto mean_new = __shfl_down_sync(0xffffffff, mean, i);\n auto m2n_new = __shfl_down_sync(0xffffffff, m2n, i);\n welford_merge_element(num, mean, m2n, num_new, mean_new, m2n_new);\n }\n}\n\ntemplate \n__device__ void welford_reduce_mean_m2n(T* __restrict__ x, int* __restrict__ count, T& mean, T& m2n, int& num,\n int block_size, int thread_id) {\n int lane = thread_id % WARP_SIZE;\n int wid = thread_id / WARP_SIZE;\n\n if (block_size > 32) {\n warp_reduce_mean_m2n(mean, m2n, num);\n if (lane == 0) {\n x[wid * 2] = mean;\n x[wid * 2 + 1] = m2n;\n count[wid] = num;\n }\n __syncthreads();\n\n if (wid == 0) {\n mean = (thread_id < block_size / WARP_SIZE) ? x[lane * 2] : T(0);\n m2n = (thread_id < block_size / WARP_SIZE) ? x[lane * 2 + 1] : T(0);\n num = (thread_id < block_size / WARP_SIZE) ? count[lane] : int(0);\n }\n }\n\n if (wid == 0) warp_reduce_mean_m2n(mean, m2n, num);\n\n return;\n}\n\n// return spatial size for NC+ Tensors\n__host__ int get_tensor_spatial_size(const at::Tensor& input) {\n auto space_size = input.size(2);\n for (int i = 3; i < input.ndimension(); i++) {\n space_size *= input.size(i);\n }\n return space_size;\n}\n\n// promote accumulation scalar type. promote half to float.\n__host__ at::ScalarType promote_scalartype(const at::Tensor& input) {\n return input.scalar_type() == at::ScalarType::Half ? at::ScalarType::Float : input.scalar_type();\n}\n\n// return single element size, optional accumulation type promotion.\n__host__ size_t get_element_data_size(const at::Tensor& input, bool accumulation = false) {\n auto scalar_type = accumulation ? promote_scalartype(input) : input.scalar_type();\n return at::elementSize(scalar_type);\n}\n\ntemplate \n__device__ __forceinline__ void welford_merge_block_vertical(C& count, T& mean, T& m2n, C* shmem_count, T* shmem_mean,\n T* shmem_m2n) {\n // write to shared memory\n auto address_base = threadIdx.x + threadIdx.y * blockDim.x;\n shmem_mean[address_base] = mean;\n shmem_m2n[address_base] = m2n;\n shmem_count[address_base] = count;\n\n#pragma unroll\n for (int offset = blockDim.y / 2; offset > 0; offset >>= 1) {\n __syncthreads();\n if (threadIdx.y < offset && threadIdx.y + offset < blockDim.y) {\n auto address = address_base + offset * blockDim.x;\n // read shared memory back to register for reduction\n auto num_new = shmem_count[address];\n auto mean_new = shmem_mean[address];\n auto m2n_new = shmem_m2n[address];\n\n welford_merge_element(count, mean, m2n, num_new, mean_new, m2n_new);\n\n // last write is not necessary\n shmem_mean[address_base] = mean;\n shmem_m2n[address_base] = m2n;\n shmem_count[address_base] = count;\n }\n }\n}\n\ntemplate \n__device__ __forceinline__ void merge_block_vertical(T& sum_dy, T& sum_dy_xmu, T* shmem_sum_dy, T* shmem_sum_dy_xmu) {\n // write to shared memory\n auto address_base = threadIdx.x + threadIdx.y * blockDim.x;\n shmem_sum_dy[address_base] = sum_dy;\n shmem_sum_dy_xmu[address_base] = sum_dy_xmu;\n\n#pragma unroll\n for (int offset = blockDim.y / 2; offset > 0; offset >>= 1) {\n __syncthreads();\n if (threadIdx.y < offset && threadIdx.y + offset < blockDim.y) {\n auto address = address_base + offset * blockDim.x;\n\n sum_dy += shmem_sum_dy[address];\n sum_dy_xmu += shmem_sum_dy_xmu[address];\n\n // last write is not necessary\n shmem_sum_dy[address_base] = sum_dy;\n shmem_sum_dy_xmu[address_base] = sum_dy_xmu;\n }\n }\n}\n\n// welford kernel calculating mean/biased_variance/unbiased_variance\ntemplate \n__global__ void welford_kernel(const scalar_t* __restrict__ input, outscalar_t* __restrict__ out_mean,\n outscalar_t* __restrict__ out_var_biased, const int bs, const int fs, const int ss) {\n int block_size = blockDim.x * blockDim.y;\n int count = 0;\n accscalar_t x_mean = accscalar_t(0);\n accscalar_t m_2_n = accscalar_t(0);\n\n int thread_id = threadIdx.y * blockDim.x + threadIdx.x;\n\n for (int batch_id = threadIdx.y; batch_id < bs; batch_id += blockDim.y) {\n int input_base = blockIdx.x * ss + batch_id * ss * fs;\n // sequential welford\n for (int offset = threadIdx.x; offset < ss; offset += blockDim.x) {\n count++;\n auto x_n = static_cast(input[offset + input_base]);\n auto d = x_n - x_mean;\n x_mean += d / count;\n m_2_n += d * (x_n - x_mean);\n }\n }\n\n static __shared__ int s_mem[160];\n accscalar_t* s_mem_ac = (accscalar_t*)&s_mem[32];\n\n welford_reduce_mean_m2n(s_mem_ac, s_mem, x_mean, m_2_n, count, block_size, thread_id);\n\n if (thread_id == 0) {\n out_mean[blockIdx.x] = static_cast(x_mean);\n out_var_biased[blockIdx.x] = static_cast(m_2_n / count);\n }\n}\n\n// elementwise BN kernel\ntemplate \n__global__ void batchnorm_forward_kernel(const scalar_t* __restrict__ input, const accscalar_t* __restrict__ mean,\n const accscalar_t* __restrict__ inv_std,\n const layerscalar_t* __restrict__ weight,\n const layerscalar_t* __restrict__ shift, scalar_t* __restrict__ out,\n const int ss, const int bs) {\n auto m_c = mean[blockIdx.x];\n auto inv_std_c = inv_std[blockIdx.x];\n auto w_c = weight == NULL ? accscalar_t(1.0) : static_cast(weight[blockIdx.x]);\n auto s_c = shift == NULL ? accscalar_t(0.0) : static_cast(shift[blockIdx.x]);\n\n for (int batch_offset = blockIdx.y * blockDim.y + threadIdx.y; batch_offset < bs;\n batch_offset += gridDim.y * blockDim.y) {\n int address_base = blockIdx.x * ss + batch_offset * gridDim.x * ss;\n for (int offset = threadIdx.x + blockIdx.z * blockDim.x; offset < ss; offset += gridDim.z * blockDim.x) {\n out[address_base + offset] =\n static_cast(w_c * (static_cast(input[address_base + offset]) - m_c) * inv_std_c + s_c);\n }\n }\n}\n\n// Backward BN kernel, calculates grad_bias, grad_weight as well as intermediate\n// results to calculating grad_input.\n// Breaking the grad_input to two step to support sync BN, which requires all\n// reduce of the intermediate results across processes.\ntemplate \n__global__ void reduce_bn_kernel(const scalar_t* __restrict__ input, const scalar_t* __restrict__ grad_output,\n const accscalar_t* __restrict__ mean, const accscalar_t* __restrict__ inv_std,\n accscalar_t* __restrict__ sum_dy_o, accscalar_t* __restrict__ sum_dy_xmu_o,\n layerscalar_t* __restrict__ grad_weight, layerscalar_t* __restrict__ grad_bias,\n const int bs, const int fs, const int ss) {\n static __shared__ int s_mem[64];\n // int total_item_num = bs * ss;\n\n int thread_id = threadIdx.y * blockDim.x + threadIdx.x;\n\n auto r_mean = mean[blockIdx.x];\n auto factor = inv_std[blockIdx.x];\n\n // Kahan sum\n accscalar_t sum_dy = 0.0;\n accscalar_t sum_dy_xmu = 0.0;\n accscalar_t sum_dy_c = 0.0;\n accscalar_t sum_dy_xmu_c = 0.0;\n for (int batch_id = threadIdx.y; batch_id < bs; batch_id += blockDim.y) {\n int input_base = blockIdx.x * ss + batch_id * ss * fs;\n for (int offset = threadIdx.x; offset < ss; offset += blockDim.x) {\n auto e_grad = static_cast(grad_output[offset + input_base]);\n auto e_input = static_cast(input[offset + input_base]);\n // calculating sum_dy\n auto sum_dy_y = e_grad - sum_dy_c;\n auto sum_dy_t = sum_dy + sum_dy_y;\n sum_dy_c = (sum_dy_t - sum_dy) - sum_dy_y;\n sum_dy = sum_dy_t;\n\n // calculating sum_dy_xmu\n auto sum_dy_xmu_y = e_grad * (e_input - r_mean) - sum_dy_xmu_c;\n auto sum_dy_xmu_t = sum_dy_xmu + sum_dy_xmu_y;\n sum_dy_xmu_c = (sum_dy_xmu_t - sum_dy_xmu) - sum_dy_xmu_y;\n sum_dy_xmu = sum_dy_xmu_t;\n }\n }\n\n sum_dy = reduce_block((accscalar_t*)s_mem, sum_dy);\n __syncthreads();\n sum_dy_xmu = reduce_block((accscalar_t*)s_mem, sum_dy_xmu);\n\n if (thread_id == 0) {\n if (grad_bias != NULL) {\n grad_bias[blockIdx.x] = static_cast(sum_dy);\n }\n if (grad_weight != NULL) {\n grad_weight[blockIdx.x] = static_cast(sum_dy_xmu * factor);\n }\n // mean_dy[blockIdx.x] = sum_dy / total_item_num;\n // mean_dy_xmu[blockIdx.x] = sum_dy_xmu / total_item_num;\n sum_dy_o[blockIdx.x] = sum_dy;\n sum_dy_xmu_o[blockIdx.x] = sum_dy_xmu;\n }\n}\n\n// elementwise backward BN kernel\ntemplate \n__global__ void batchnorm_backward_kernel(const scalar_t* __restrict__ grad_output, const scalar_t* __restrict__ input,\n const accscalar_t* __restrict__ mean, const accscalar_t* __restrict__ inv_std,\n const layerscalar_t* __restrict__ weight,\n const accscalar_t* __restrict__ sum_dy,\n const accscalar_t* __restrict__ sum_dy_xmu, const int* __restrict__ numel,\n scalar_t* __restrict__ grad_input, const int64_t world_size, const int ss,\n const int bs) {\n int64_t div = 0;\n for (int i = 0; i < world_size; i++) {\n div += numel[i];\n }\n auto m_c = static_cast(mean[blockIdx.x]);\n // auto m_dy_c = static_cast(mean_dy[blockIdx.x]);\n auto m_dy_c = static_cast(sum_dy[blockIdx.x]) / div;\n auto factor_1_c = inv_std[blockIdx.x];\n auto factor_2_c = (weight == NULL ? accscalar_t(1.0) : static_cast(weight[blockIdx.x])) * factor_1_c;\n // factor_1_c = factor_1_c * factor_1_c * mean_dy_xmu[blockIdx.x];\n factor_1_c = factor_1_c * factor_1_c * sum_dy_xmu[blockIdx.x] / div;\n\n for (int batch_offset = blockIdx.y * blockDim.y + threadIdx.y; batch_offset < bs;\n batch_offset += gridDim.y * blockDim.y) {\n int address_base = blockIdx.x * ss + batch_offset * gridDim.x * ss;\n for (int offset = threadIdx.x + blockIdx.z * blockDim.x; offset < ss; offset += gridDim.z * blockDim.x) {\n grad_input[address_base + offset] =\n (static_cast(grad_output[address_base + offset]) - m_dy_c -\n (static_cast(input[address_base + offset]) - m_c) * factor_1_c) *\n factor_2_c;\n }\n }\n}\n\n// welford kernel for c last tensor calculating mean/biased_variance/unbiased_variance\ntemplate \n__global__ void welford_kernel_c_last(const scalar_t* __restrict__ input, outscalar_t* __restrict__ out_mean,\n outscalar_t* __restrict__ out_var_biased, volatile accscalar_t* staging_data,\n int* semaphores, const int reduction_size, const int stride) {\n // hide latency with concurrency\n accscalar_t x_mean[PARALLEL_LOADS];\n accscalar_t m_2_n[PARALLEL_LOADS];\n int count[PARALLEL_LOADS];\n\n#pragma unroll\n for (int i = 0; i < PARALLEL_LOADS; i++) {\n x_mean[i] = accscalar_t(0);\n m_2_n[i] = accscalar_t(0);\n count[i] = accscalar_t(0);\n }\n // tensor dimension (m,c)\n\n // loop along m dimension\n int inner_loop_stride = blockDim.y * gridDim.y;\n\n // offset along m dimension\n int m_offset = blockIdx.y * blockDim.y + threadIdx.y;\n int c_offset = blockIdx.x * blockDim.x + threadIdx.x;\n\n int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);\n int address_base = m_offset * stride + c_offset;\n int address_increment = inner_loop_stride * stride;\n\n for (int i = 0; i < loop_count; i++) {\n accscalar_t x_math[PARALLEL_LOADS];\n accscalar_t x_count_inv[PARALLEL_LOADS];\n accscalar_t is_valid[PARALLEL_LOADS];\n\n // load multiple data in\n#pragma unroll\n for (int j = 0; j < PARALLEL_LOADS; j++) {\n if (c_offset < stride && m_offset < reduction_size) {\n x_math[j] = input[address_base];\n count[j]++;\n x_count_inv[j] = accscalar_t(1) / count[j];\n is_valid[j] = accscalar_t(1);\n } else {\n x_math[j] = accscalar_t(0);\n x_count_inv[j] = accscalar_t(0);\n is_valid[j] = accscalar_t(0);\n }\n m_offset += inner_loop_stride;\n address_base += address_increment;\n }\n\n // calculate mean/m2n with welford\n#pragma unroll\n for (int j = 0; j < PARALLEL_LOADS; j++) {\n accscalar_t delta0 = x_math[j] - x_mean[j];\n x_mean[j] += delta0 * x_count_inv[j];\n accscalar_t delta1 = x_math[j] - x_mean[j];\n m_2_n[j] += delta0 * delta1 * is_valid[j];\n }\n }\n\n // thread reduction to accumulate mean/m_2_n/count between PARALLEL_LOADS\n#pragma unroll\n for (int j = 1; j < PARALLEL_LOADS; j++) {\n welford_merge_element(count[0], x_mean[0], m_2_n[0], count[j], x_mean[j], m_2_n[j]);\n }\n\n // release x_mean / m_2_n\n auto mean_th = x_mean[0];\n auto m2_th = m_2_n[0];\n auto count_th = count[0];\n\n // block-wise reduction with shared memory (since reduction cannot be done within a warp)\n static __shared__ accscalar_t shmem_mean[MAX_BLOCK_SIZE];\n static __shared__ accscalar_t shmem_m2n[MAX_BLOCK_SIZE];\n static __shared__ int shmem_count[MAX_BLOCK_SIZE];\n\n welford_merge_block_vertical(count_th, mean_th, m2_th, shmem_count, shmem_mean, shmem_m2n);\n\n // grid reduction if needed (coop launch used at the first place)\n if (gridDim.y > 1) {\n volatile accscalar_t* staging_mean = staging_data;\n volatile accscalar_t* staging_m2n = &staging_data[stride * gridDim.y];\n volatile int* staging_count = reinterpret_cast(&staging_m2n[stride * gridDim.y]);\n\n address_base = c_offset + blockIdx.y * stride;\n // write data to staging_data;\n if (threadIdx.y == 0 && c_offset < stride) {\n staging_mean[address_base] = mean_th;\n staging_m2n[address_base] = m2_th;\n staging_count[address_base] = count_th;\n }\n\n __threadfence();\n __syncthreads(); // ensuring writes to staging_ is visible to all blocks\n\n __shared__ bool is_last_block_done;\n // mark block done\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n int old = atomicAdd(&semaphores[blockIdx.x], 1);\n is_last_block_done = (old == (gridDim.y - 1));\n }\n\n __syncthreads();\n\n // check that all data is now available in global memory\n if (is_last_block_done) {\n count_th = 0;\n mean_th = accscalar_t(0.0);\n m2_th = accscalar_t(0.0);\n\n for (int y = threadIdx.y; y < gridDim.y; y += blockDim.y) {\n address_base = c_offset + y * stride;\n int num_new = c_offset < stride ? staging_count[address_base] : 0;\n accscalar_t mean_new = c_offset < stride ? staging_mean[address_base] : accscalar_t(0.0);\n accscalar_t m2n_new = c_offset < stride ? staging_m2n[address_base] : accscalar_t(0.0);\n\n welford_merge_element(count_th, mean_th, m2_th, num_new, mean_new, m2n_new);\n }\n\n welford_merge_block_vertical(count_th, mean_th, m2_th, shmem_count, shmem_mean, shmem_m2n);\n if (threadIdx.y == 0 && c_offset < stride) {\n out_mean[c_offset] = static_cast(mean_th);\n out_var_biased[c_offset] = static_cast(m2_th / count_th);\n }\n }\n } else {\n if (blockIdx.y == 0 && threadIdx.y == 0 && c_offset < stride) {\n out_mean[c_offset] = static_cast(mean_th);\n out_var_biased[c_offset] = static_cast(m2_th / count_th);\n }\n }\n}\n\n// parallel welford kernel to further reduce mean / biased_var\n// into mean / unbiased_var / inv_std across multiple processes.\ntemplate \n__global__ void welford_kernel_parallel(const scalar_t* __restrict__ mean, const scalar_t* __restrict__ var_biased,\n const int* __restrict__ numel, scalar_t* __restrict__ out_mean,\n scalar_t* __restrict__ out_var, scalar_t* __restrict__ inv_std,\n const int world_size, const int feature_size, const float eps) {\n for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < feature_size; i += gridDim.x * blockDim.x) {\n // load data;\n int address = i;\n scalar_t x_mean = 0;\n scalar_t m_2_n = 0;\n int count = 0;\n for (int j = 0; j < world_size; j++) {\n welford_merge_element(count, x_mean, m_2_n, numel[j], mean[address], var_biased[address] * numel[j]);\n address += feature_size;\n }\n out_mean[i] = x_mean;\n out_var[i] = m_2_n / (count - 1);\n inv_std[i] = scalar_t(1) / sqrt(m_2_n / count + eps);\n }\n}\n\n// elementwise BN kernel\ntemplate \n__global__ void batchnorm_forward_c_last_kernel(const scalar_t* __restrict__ input, const scalar_t* __restrict__ z,\n const accscalar_t* __restrict__ mean,\n const accscalar_t* __restrict__ inv_std,\n const layerscalar_t* __restrict__ weight,\n const layerscalar_t* __restrict__ shift, scalar_t* __restrict__ out,\n const int reduction_size, const int stride, const bool fuse_relu) {\n // tensor dimension (m,c)\n // loop along m dimension\n int inner_loop_stride = blockDim.y * gridDim.y;\n\n // offset along m dimension\n int m_offset = blockIdx.y * blockDim.y + threadIdx.y;\n int c_offset = blockIdx.x * blockDim.x + threadIdx.x;\n\n auto m_c = mean[c_offset];\n auto inv_std_c = static_cast(inv_std[c_offset]);\n auto w_c = weight == NULL ? accscalar_t(1.0) : static_cast(weight[c_offset]);\n auto s_c = shift == NULL ? accscalar_t(0.0) : static_cast(shift[c_offset]);\n\n int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);\n int address_base = m_offset * stride + c_offset;\n int address_increment = inner_loop_stride * stride;\n\n for (int i = 0; i < loop_count; i++) {\n#pragma unroll\n for (int j = 0; j < PARALLEL_LOADS; j++) {\n if (c_offset < stride && m_offset < reduction_size) {\n auto tmp = w_c * (static_cast(input[address_base]) - m_c) * inv_std_c + s_c;\n if (z != NULL) {\n tmp += z[address_base];\n }\n out[address_base] = (fuse_relu && tmp <= accscalar_t(0.0) ? scalar_t(0.0) : static_cast(tmp));\n }\n m_offset += inner_loop_stride;\n address_base += address_increment;\n }\n }\n}\n\n// elementwise BN kernel\ntemplate \n__global__ void relu_backward_c_last_kernel(const scalar_t* __restrict__ grad_output,\n const scalar_t* __restrict__ input, const scalar_t* __restrict__ z,\n const accscalar_t* __restrict__ mean,\n const accscalar_t* __restrict__ inv_std,\n const layerscalar_t* __restrict__ weight,\n const layerscalar_t* __restrict__ shift, scalar_t* __restrict__ out,\n const int reduction_size, const int stride) {\n // tensor dimension (m,c)\n // loop along m dimension\n int inner_loop_stride = blockDim.y * gridDim.y;\n\n // offset along m dimension\n int m_offset = blockIdx.y * blockDim.y + threadIdx.y;\n int c_offset = blockIdx.x * blockDim.x + threadIdx.x;\n\n auto m_c = mean[c_offset];\n auto inv_std_c = static_cast(inv_std[c_offset]);\n auto w_c = weight == NULL ? accscalar_t(1.0) : static_cast(weight[c_offset]);\n auto s_c = shift == NULL ? accscalar_t(0.0) : static_cast(shift[c_offset]);\n\n int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);\n int address_base = m_offset * stride + c_offset;\n int address_increment = inner_loop_stride * stride;\n\n for (int i = 0; i < loop_count; i++) {\n#pragma unroll\n for (int j = 0; j < PARALLEL_LOADS; j++) {\n if (c_offset < stride && m_offset < reduction_size) {\n auto tmp = w_c * (static_cast(input[address_base]) - m_c) * inv_std_c + s_c;\n if (z != NULL) {\n tmp += z[address_base];\n }\n out[address_base] = (tmp <= accscalar_t(0.0) ? scalar_t(0.0) : grad_output[address_base]);\n }\n m_offset += inner_loop_stride;\n address_base += address_increment;\n }\n }\n}\n\n// batchnorm backward kernel for c last tensor\ntemplate \n__global__ void reduce_bn_c_last_kernel(const scalar_t* __restrict__ input, const scalar_t* __restrict__ grad_output,\n const accscalar_t* __restrict__ mean, const accscalar_t* __restrict__ inv_std,\n accscalar_t* __restrict__ sum_dy_o, accscalar_t* __restrict__ sum_dy_xmu_o,\n layerscalar_t* __restrict__ grad_weight, layerscalar_t* __restrict__ grad_bias,\n volatile accscalar_t* staging_data, int* semaphores, const int reduction_size,\n const int stride) {\n // hide latency with concurrency\n accscalar_t sum_dy[PARALLEL_LOADS];\n accscalar_t sum_dy_xmu[PARALLEL_LOADS];\n\n#pragma unroll\n for (int i = 0; i < PARALLEL_LOADS; i++) {\n sum_dy[i] = accscalar_t(0);\n sum_dy_xmu[i] = accscalar_t(0);\n }\n // tensor dimension (m,c)\n\n // loop along m dimension\n int inner_loop_stride = blockDim.y * gridDim.y;\n\n // offset along m dimension\n int m_offset = blockIdx.y * blockDim.y + threadIdx.y;\n int c_offset = blockIdx.x * blockDim.x + threadIdx.x;\n\n int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);\n int address_base = m_offset * stride + c_offset;\n int address_increment = inner_loop_stride * stride;\n\n auto r_mean = mean[c_offset];\n auto factor = inv_std[c_offset];\n\n for (int i = 0; i < loop_count; i++) {\n accscalar_t x_input[PARALLEL_LOADS];\n accscalar_t x_grad_output[PARALLEL_LOADS];\n\n // load multiple data in\n#pragma unroll\n for (int j = 0; j < PARALLEL_LOADS; j++) {\n if (c_offset < stride && m_offset < reduction_size) {\n x_input[j] = input[address_base];\n x_grad_output[j] = grad_output[address_base];\n } else {\n x_input[j] = accscalar_t(0);\n x_grad_output[j] = accscalar_t(0);\n }\n m_offset += inner_loop_stride;\n address_base += address_increment;\n }\n\n // calculate sum_dy / sum_dy_xmu\n#pragma unroll\n for (int j = 0; j < PARALLEL_LOADS; j++) {\n sum_dy[j] += x_grad_output[j];\n sum_dy_xmu[j] += x_grad_output[j] * (x_input[j] - r_mean);\n }\n }\n\n // thread reduction to accumulate sum_dy / sum_dy_xmu between PARALLEL_LOADS\n#pragma unroll\n for (int j = 1; j < PARALLEL_LOADS; j++) {\n sum_dy[0] += sum_dy[j];\n sum_dy_xmu[0] += sum_dy_xmu[j];\n }\n\n // release array of registers\n auto sum_dy_th = sum_dy[0];\n auto sum_dy_xmu_th = sum_dy_xmu[0];\n\n // block-wise reduction with shared memory (since reduction cannot be done within a warp)\n static __shared__ accscalar_t shmem_sum_dy[MAX_BLOCK_SIZE];\n static __shared__ accscalar_t shmem_sum_dy_xmu[MAX_BLOCK_SIZE];\n\n merge_block_vertical(sum_dy_th, sum_dy_xmu_th, shmem_sum_dy, shmem_sum_dy_xmu);\n\n // grid reduction if needed (coop launch used at the first place)\n if (gridDim.y > 1) {\n volatile accscalar_t* staging_sum_dy = staging_data;\n volatile accscalar_t* staging_sum_dy_xmu = &staging_data[stride * gridDim.y];\n\n address_base = c_offset + blockIdx.y * stride;\n // write data to staging_data;\n if (threadIdx.y == 0 && c_offset < stride) {\n staging_sum_dy[address_base] = sum_dy_th;\n staging_sum_dy_xmu[address_base] = sum_dy_xmu_th;\n }\n\n __threadfence();\n __syncthreads(); // ensuring writes to staging_ is visible to all blocks\n\n __shared__ bool is_last_block_done;\n // mark block done\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n int old = atomicAdd(&semaphores[blockIdx.x], 1);\n is_last_block_done = (old == (gridDim.y - 1));\n }\n\n __syncthreads();\n\n // check that all data is now available in global memory\n if (is_last_block_done) {\n sum_dy_th = accscalar_t(0.0);\n sum_dy_xmu_th = accscalar_t(0.0);\n\n for (int y = threadIdx.y; y < gridDim.y; y += blockDim.y) {\n address_base = c_offset + y * stride;\n sum_dy_th += (c_offset < stride ? staging_sum_dy[address_base] : accscalar_t(0.0));\n sum_dy_xmu_th += (c_offset < stride ? staging_sum_dy_xmu[address_base] : accscalar_t(0.0));\n }\n\n merge_block_vertical(sum_dy_th, sum_dy_xmu_th, shmem_sum_dy, shmem_sum_dy_xmu);\n if (threadIdx.y == 0 && c_offset < stride) {\n if (grad_bias != NULL) {\n grad_bias[c_offset] = static_cast(sum_dy_th);\n }\n if (grad_weight != NULL) {\n grad_weight[c_offset] = static_cast(sum_dy_xmu_th * factor);\n }\n // mean_dy[c_offset] = sum_dy_th / reduction_size;\n // mean_dy_xmu[c_offset] = sum_dy_xmu_th / reduction_size;\n sum_dy_o[c_offset] = sum_dy_th;\n sum_dy_xmu_o[c_offset] = sum_dy_xmu_th;\n }\n }\n } else {\n if (blockIdx.y == 0 && threadIdx.y == 0 && c_offset < stride) {\n if (grad_bias != NULL) {\n grad_bias[c_offset] = static_cast(sum_dy_th);\n }\n if (grad_weight != NULL) {\n grad_weight[c_offset] = static_cast(sum_dy_xmu_th * factor);\n }\n // mean_dy[c_offset] = sum_dy_th / reduction_size;\n // mean_dy_xmu[c_offset] = sum_dy_xmu_th / reduction_size;\n sum_dy_o[c_offset] = sum_dy_th;\n sum_dy_xmu_o[c_offset] = sum_dy_xmu_th;\n }\n }\n}\n\n// elementwise BN kernel\ntemplate \n__global__ void batchnorm_backward_c_last_kernel(\n const scalar_t* __restrict__ grad_output, const scalar_t* __restrict__ input, const accscalar_t* __restrict__ mean,\n const accscalar_t* __restrict__ inv_std, const layerscalar_t* __restrict__ weight,\n const accscalar_t* __restrict__ sum_dy, const accscalar_t* __restrict__ sum_dy_xmu, const int* __restrict__ numel,\n scalar_t* __restrict__ grad_input, const int64_t world_size, const int reduction_size, const int stride) {\n int64_t div = 0;\n for (int i = 0; i < world_size; i++) {\n div += numel[i];\n }\n // tensor dimension (m,c)\n // loop along m dimension\n int inner_loop_stride = blockDim.y * gridDim.y;\n\n // offset along m dimension\n int m_offset = blockIdx.y * blockDim.y + threadIdx.y;\n int c_offset = blockIdx.x * blockDim.x + threadIdx.x;\n\n auto m_c = mean[c_offset];\n auto m_dy_c = sum_dy[c_offset] / div;\n auto factor_1_c = inv_std[c_offset];\n auto factor_2_c = (weight == NULL ? accscalar_t(1.0) : static_cast(weight[c_offset])) * factor_1_c;\n factor_1_c = factor_1_c * factor_1_c * sum_dy_xmu[c_offset] / div;\n\n int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);\n int address_base = m_offset * stride + c_offset;\n int address_increment = inner_loop_stride * stride;\n\n for (int i = 0; i < loop_count; i++) {\n#pragma unroll\n for (int j = 0; j < PARALLEL_LOADS; j++) {\n if (c_offset < stride && m_offset < reduction_size) {\n grad_input[address_base] =\n static_cast((static_cast(grad_output[address_base]) - m_dy_c -\n (static_cast(input[address_base]) - m_c) * factor_1_c) *\n factor_2_c);\n }\n m_offset += inner_loop_stride;\n address_base += address_increment;\n }\n }\n}\n\nstd::vector welford_mean_var_CUDA(const at::Tensor input) {\n const auto batch_size = input.size(0);\n const auto feature_size = input.size(1);\n\n auto space_size = get_tensor_spatial_size(input);\n auto scalar_type = promote_scalartype(input);\n\n at::Tensor out_var_biased = at::empty({feature_size}, input.options().dtype(scalar_type));\n at::Tensor out_mean = at::empty({feature_size}, input.options().dtype(scalar_type));\n\n int block_y = min(h_last_pow2(batch_size), int(MAX_BLOCK_SIZE / 32));\n int block_x = max(1, min(MAX_BLOCK_SIZE / block_y, h_last_pow2(space_size)));\n const dim3 block(block_x, block_y);\n const dim3 grid(feature_size);\n\n auto stream = at::cuda::getCurrentCUDAStream();\n\n {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"welford_mean_var_kernel\", using accscalar_t = at::acc_type;\n welford_kernel\n <<>>(input.data_ptr(), out_mean.data_ptr(),\n out_var_biased.data_ptr(), batch_size, feature_size, space_size););\n }\n\n return {out_mean, out_var_biased};\n}\n\nat::Tensor batchnorm_forward_CUDA(const at::Tensor input, const at::Tensor mean, const at::Tensor inv_std,\n const at::optional weight, const at::optional shift) {\n const auto batch_size = input.size(0);\n const auto feature_size = input.size(1);\n at::Tensor out = at::empty_like(input);\n\n auto space_size = get_tensor_spatial_size(input);\n\n int block_x = max(32, min(MAX_BLOCK_SIZE, h_last_pow2(space_size) / 4));\n int block_y = max(1, min(MAX_BLOCK_SIZE / block_x, h_last_pow2(batch_size) / 4));\n const dim3 block(block_x, block_y);\n int grid_z = max(1, min(65535, h_last_pow2(space_size) / 4 / block_x));\n int batch_group_size = max(1, min(65535, h_last_pow2(batch_size) / block_y));\n const dim3 grid(feature_size, batch_group_size, grid_z);\n auto stream = at::cuda::getCurrentCUDAStream();\n\n if (input.scalar_type() == at::ScalarType::Half && weight.has_value() &&\n weight.value().scalar_type() == at::ScalarType::Float) {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n batchnorm_forward_kernel<<>>(\n input.data_ptr(), mean.data_ptr(), inv_std.data_ptr(),\n weight.has_value() ? weight.value().data_ptr() : NULL,\n shift.has_value() ? shift.value().data_ptr() : NULL, out.data_ptr(), space_size,\n batch_size););\n } else {\n if (weight.has_value()) {\n TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),\n \"input.scalar_type() is not supported with weight.scalar_type()\");\n }\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n batchnorm_forward_kernel<<>>(\n input.data_ptr(), mean.data_ptr(), inv_std.data_ptr(),\n weight.has_value() ? weight.value().data_ptr() : NULL,\n shift.has_value() ? shift.value().data_ptr() : NULL, out.data_ptr(), space_size,\n batch_size););\n }\n return out;\n}\n\nstd::vector reduce_bn_CUDA(const at::Tensor grad_output, const at::Tensor input, const at::Tensor mean,\n const at::Tensor inv_std, const at::optional weight) {\n const auto batch_size = input.size(0);\n const auto feature_size = input.size(1);\n\n auto scalar_type = promote_scalartype(input);\n\n at::Tensor sum_dy = at::empty({feature_size}, mean.options());\n at::Tensor sum_dy_xmu = at::empty({feature_size}, mean.options());\n\n at::Tensor grad_weight;\n at::Tensor grad_bias;\n if (weight.has_value()) {\n grad_weight = at::empty({feature_size}, weight.value().options());\n grad_bias = at::empty({feature_size}, weight.value().options());\n } else {\n grad_weight = at::empty({0}, mean.options());\n grad_bias = at::empty({0}, mean.options());\n }\n\n auto space_size = get_tensor_spatial_size(input);\n\n int block_y = min(h_last_pow2(batch_size), int(MAX_BLOCK_SIZE / 32));\n int block_x = max(1, min(MAX_BLOCK_SIZE / block_y, h_last_pow2(space_size)));\n const dim3 block(block_x, block_y);\n const dim3 grid(feature_size);\n auto stream = at::cuda::getCurrentCUDAStream();\n\n if (input.scalar_type() == at::ScalarType::Half && weight.has_value() &&\n weight.value().scalar_type() == at::ScalarType::Float) {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_backward_reduce\", using accscalar_t = at::acc_type;\n reduce_bn_kernel<<>>(\n input.data_ptr(), grad_output.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), sum_dy.data_ptr(), sum_dy_xmu.data_ptr(),\n weight.has_value() ? grad_weight.data_ptr() : NULL,\n weight.has_value() ? grad_bias.data_ptr() : NULL, batch_size, feature_size, space_size););\n } else {\n if (weight.has_value()) {\n TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),\n \"input.scalar_type() is not supported with weight.scalar_type()\");\n }\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_backward_reduce\", using accscalar_t = at::acc_type;\n reduce_bn_kernel<<>>(\n input.data_ptr(), grad_output.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), sum_dy.data_ptr(), sum_dy_xmu.data_ptr(),\n weight.has_value() ? grad_weight.data_ptr() : NULL,\n weight.has_value() ? grad_bias.data_ptr() : NULL, batch_size, feature_size, space_size););\n }\n\n return {sum_dy, sum_dy_xmu, grad_weight, grad_bias};\n}\n\nat::Tensor batchnorm_backward_CUDA(const at::Tensor grad_output, const at::Tensor input, const at::Tensor mean,\n const at::Tensor inv_std, const at::optional weight,\n const at::Tensor sum_dy, const at::Tensor sum_dy_xmu, const at::Tensor count) {\n const auto batch_size = input.size(0);\n const auto feature_size = input.size(1);\n\n at::Tensor grad_input = at::empty_like(input);\n\n auto space_size = get_tensor_spatial_size(input);\n\n int block_x = max(32, min(MAX_BLOCK_SIZE, h_last_pow2(space_size) / 4));\n int block_y = max(1, min(MAX_BLOCK_SIZE / block_x, h_last_pow2(batch_size) / 4));\n const dim3 block(block_x, block_y);\n int grid_z = max(1, min(65535, h_last_pow2(space_size) / 4 / block_x));\n int batch_group_size = max(1, min(65535, h_last_pow2(batch_size) / block_y));\n const dim3 grid(feature_size, batch_group_size, grid_z);\n\n auto stream = at::cuda::getCurrentCUDAStream();\n\n if (input.scalar_type() == at::ScalarType::Half && weight.has_value() &&\n weight.value().scalar_type() == at::ScalarType::Float) {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_backward\", using accscalar_t = at::acc_type;\n batchnorm_backward_kernel<<>>(\n grad_output.data_ptr(), input.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), weight.has_value() ? weight.value().data_ptr() : NULL,\n sum_dy.data_ptr(), sum_dy_xmu.data_ptr(), count.data_ptr(),\n grad_input.data_ptr(), count.numel(), space_size, batch_size););\n } else {\n if (weight.has_value()) {\n TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),\n \"input.scalar_type() is not supported with weight.scalar_type()\");\n }\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_backward\", using accscalar_t = at::acc_type;\n batchnorm_backward_kernel<<>>(\n grad_output.data_ptr(), input.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), weight.has_value() ? weight.value().data_ptr() : NULL,\n sum_dy.data_ptr(), sum_dy_xmu.data_ptr(), count.data_ptr(),\n grad_input.data_ptr(), count.numel(), space_size, batch_size););\n }\n\n return grad_input;\n}\n\nstd::vector welford_parallel_CUDA(const at::Tensor mean_feature_nodes, const at::Tensor var_biased,\n const at::Tensor numel, const float eps) {\n const auto world_size = mean_feature_nodes.size(0);\n const auto feature_size = mean_feature_nodes.size(1);\n\n at::Tensor out_var = at::empty({feature_size}, var_biased.options());\n at::Tensor inv_std = at::empty_like(out_var);\n at::Tensor out_mean = at::empty_like(out_var);\n\n at::Tensor mean_feature_nodes_ = mean_feature_nodes.contiguous();\n at::Tensor var_biased_ = var_biased.contiguous();\n at::Tensor numel_ = numel.contiguous();\n\n // TODO(jie): tile this for memory coalescing!\n const int block = std::min(h_last_pow2(feature_size), MAX_BLOCK_SIZE);\n const int grid = std::max(1, feature_size / block);\n\n auto stream = at::cuda::getCurrentCUDAStream();\n\n {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(mean_feature_nodes.scalar_type(), 0, \"welford_parallel_kernel\",\n welford_kernel_parallel<<>>(\n mean_feature_nodes_.data_ptr(), var_biased_.data_ptr(),\n numel_.data_ptr(), out_mean.data_ptr(), out_var.data_ptr(),\n inv_std.data_ptr(), world_size, feature_size, eps););\n }\n\n return {out_mean, out_var, inv_std};\n}\n\nstd::vector welford_mean_var_c_last_CUDA(const at::Tensor input) {\n const auto stride = input.size(input.ndimension() - 1);\n const auto reduction_size = input.numel() / stride;\n\n auto scalar_type = promote_scalartype(input);\n auto option = input.options().dtype(scalar_type);\n\n at::Tensor out_var_biased = at::empty({stride}, option);\n at::Tensor out_mean = at::empty({stride}, option);\n\n dim3 block;\n dim3 grid;\n flexible_launch_configs(reduction_size, stride, block, grid, true);\n\n at::Tensor staging_data;\n at::Tensor semaphores;\n if (grid.y > 1) {\n staging_data = at::empty({4 * stride * grid.y}, option);\n semaphores = at::zeros({grid.x}, input.options().dtype(at::kInt));\n }\n\n auto stream = at::cuda::getCurrentCUDAStream();\n\n {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"welford_mean_var_c_last\", using accscalar_t = at::acc_type;\n accscalar_t* staging_data_ptr = grid.y > 1 ? staging_data.data_ptr() : nullptr;\n int* semaphores_ptr = grid.y > 1 ? semaphores.data_ptr() : nullptr;\n welford_kernel_c_last<<>>(\n input.data_ptr(), out_mean.data_ptr(), out_var_biased.data_ptr(),\n staging_data_ptr, semaphores_ptr, reduction_size, stride););\n }\n\n return {out_mean, out_var_biased};\n}\n\nat::Tensor batchnorm_forward_c_last_CUDA(const at::Tensor input, const at::optional z,\n const at::Tensor mean, const at::Tensor inv_std,\n const at::optional weight, const at::optional shift,\n const bool fuse_relu) {\n const auto stride = input.size(input.ndimension() - 1);\n const auto reduction_size = input.numel() / stride;\n\n at::Tensor out = at::empty_like(input);\n\n dim3 block;\n dim3 grid;\n flexible_launch_configs(reduction_size, stride, block, grid);\n\n auto stream = at::cuda::getCurrentCUDAStream();\n\n if (input.scalar_type() == at::ScalarType::Half && weight.has_value() &&\n weight.value().scalar_type() == at::ScalarType::Float) {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n batchnorm_forward_c_last_kernel\n <<>>(input.data_ptr(),\n z.has_value() ? z.value().data_ptr() : NULL,\n mean.data_ptr(), inv_std.data_ptr(),\n weight.has_value() ? weight.value().data_ptr() : NULL,\n shift.has_value() ? shift.value().data_ptr() : NULL,\n out.data_ptr(), reduction_size, stride, fuse_relu););\n } else {\n if (weight.has_value()) {\n TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),\n \"input.scalar_type() is not supported with weight.scalar_type()\");\n }\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n batchnorm_forward_c_last_kernel\n <<>>(input.data_ptr(),\n z.has_value() ? z.value().data_ptr() : NULL,\n mean.data_ptr(), inv_std.data_ptr(),\n weight.has_value() ? weight.value().data_ptr() : NULL,\n shift.has_value() ? shift.value().data_ptr() : NULL,\n out.data_ptr(), reduction_size, stride, fuse_relu););\n }\n return out;\n}\n\nstd::vector reduce_bn_c_last_CUDA(const at::Tensor grad_output, const at::Tensor input,\n const at::Tensor mean, const at::Tensor inv_std,\n const at::optional weight) {\n const auto stride = input.size(input.ndimension() - 1);\n const auto reduction_size = input.numel() / stride;\n\n at::Tensor sumn_dy = at::empty({stride}, mean.options());\n at::Tensor sum_dy_xmu = at::empty({stride}, mean.options());\n\n at::Tensor grad_weight;\n at::Tensor grad_bias;\n if (weight.has_value()) {\n grad_weight = at::empty({stride}, weight.value().options());\n grad_bias = at::empty({stride}, weight.value().options());\n } else {\n // because I cannot return an uninitialized at::Tensor\n grad_weight = at::empty({0}, mean.options());\n grad_bias = at::empty({0}, mean.options());\n }\n\n dim3 block;\n dim3 grid;\n flexible_launch_configs(reduction_size, stride, block, grid, true);\n\n at::Tensor staging_data;\n at::Tensor semaphores;\n if (grid.y > 1) {\n staging_data = at::empty({2 * stride * grid.y}, mean.options());\n semaphores = at::zeros({grid.x}, input.options().dtype(at::kInt));\n }\n auto stream = at::cuda::getCurrentCUDAStream();\n\n if (input.scalar_type() == at::ScalarType::Half && weight.has_value() &&\n weight.value().scalar_type() == at::ScalarType::Float) {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_backward_reduce\", using accscalar_t = at::acc_type;\n accscalar_t* staging_data_ptr = grid.y > 1 ? staging_data.data_ptr() : nullptr;\n int* semaphores_ptr = grid.y > 1 ? semaphores.data_ptr() : nullptr;\n reduce_bn_c_last_kernel<<>>(\n input.data_ptr(), grad_output.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), sumn_dy.data_ptr(), sum_dy_xmu.data_ptr(),\n weight.has_value() ? grad_weight.data_ptr() : NULL,\n weight.has_value() ? grad_bias.data_ptr() : NULL, staging_data_ptr, semaphores_ptr,\n reduction_size, stride););\n } else {\n if (weight.has_value()) {\n TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),\n \"input.scalar_type() is not supported with weight.scalar_type()\");\n }\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_backward_reduce\", using accscalar_t = at::acc_type;\n accscalar_t* staging_data_ptr = grid.y > 1 ? staging_data.data_ptr() : nullptr;\n int* semaphores_ptr = grid.y > 1 ? semaphores.data_ptr() : nullptr;\n reduce_bn_c_last_kernel<<>>(\n input.data_ptr(), grad_output.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), sumn_dy.data_ptr(), sum_dy_xmu.data_ptr(),\n weight.has_value() ? grad_weight.data_ptr() : NULL,\n weight.has_value() ? grad_bias.data_ptr() : NULL, staging_data_ptr, semaphores_ptr,\n reduction_size, stride););\n }\n\n return {sumn_dy, sum_dy_xmu, grad_weight, grad_bias};\n}\n\nat::Tensor batchnorm_backward_c_last_CUDA(const at::Tensor grad_output, const at::Tensor input, const at::Tensor mean,\n const at::Tensor inv_std, const at::optional weight,\n const at::Tensor sum_dy, const at::Tensor sum_dy_xmu,\n const at::Tensor count) {\n const auto stride = input.size(input.ndimension() - 1);\n const auto reduction_size = input.numel() / stride;\n\n at::Tensor grad_input = at::empty_like(input);\n\n dim3 block;\n dim3 grid;\n flexible_launch_configs(reduction_size, stride, block, grid);\n\n auto stream = at::cuda::getCurrentCUDAStream();\n\n if (input.scalar_type() == at::ScalarType::Half && weight.has_value() &&\n weight.value().scalar_type() == at::ScalarType::Float) {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n batchnorm_backward_c_last_kernel\n <<>>(\n grad_output.data_ptr(), input.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), weight.has_value() ? weight.value().data_ptr() : NULL,\n sum_dy.data_ptr(), sum_dy_xmu.data_ptr(), count.data_ptr(),\n grad_input.data_ptr(), count.numel(), reduction_size, stride););\n } else {\n if (weight.has_value()) {\n TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),\n \"input.scalar_type() is not supported with weight.scalar_type()\");\n }\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n batchnorm_backward_c_last_kernel\n <<>>(\n grad_output.data_ptr(), input.data_ptr(), mean.data_ptr(),\n inv_std.data_ptr(), weight.has_value() ? weight.value().data_ptr() : NULL,\n sum_dy.data_ptr(), sum_dy_xmu.data_ptr(), count.data_ptr(),\n grad_input.data_ptr(), count.numel(), reduction_size, stride););\n }\n\n return grad_input;\n}\n\nat::Tensor relu_backward_c_last_CUDA(const at::Tensor grad_output, const at::Tensor input,\n const at::optional z, const at::Tensor mean, const at::Tensor inv_std,\n const at::optional weight, const at::optional shift) {\n const auto stride = input.size(input.ndimension() - 1);\n const auto reduction_size = input.numel() / stride;\n\n at::Tensor out = at::empty_like(input);\n\n dim3 block;\n dim3 grid;\n flexible_launch_configs(reduction_size, stride, block, grid);\n\n auto stream = at::cuda::getCurrentCUDAStream();\n\n if (input.scalar_type() == at::ScalarType::Half && weight.has_value() &&\n weight.value().scalar_type() == at::ScalarType::Float) {\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n relu_backward_c_last_kernel\n <<>>(grad_output.data_ptr(), input.data_ptr(),\n z.has_value() ? z.value().data_ptr() : NULL,\n mean.data_ptr(), inv_std.data_ptr(),\n weight.has_value() ? weight.value().data_ptr() : NULL,\n shift.has_value() ? shift.value().data_ptr() : NULL,\n out.data_ptr(), reduction_size, stride););\n } else {\n if (weight.has_value()) {\n TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),\n \"input.scalar_type() is not supported with weight.scalar_type()\");\n }\n using namespace at;\n DISPATCH_FLOAT_AND_HALF(\n input.scalar_type(), 0, \"batchnorm_forward\", using accscalar_t = at::acc_type;\n relu_backward_c_last_kernel<<>>(\n grad_output.data_ptr(), input.data_ptr(),\n z.has_value() ? z.value().data_ptr() : NULL, mean.data_ptr(),\n inv_std.data_ptr(), weight.has_value() ? weight.value().data_ptr() : NULL,\n shift.has_value() ? shift.value().data_ptr() : NULL, out.data_ptr(), reduction_size,\n stride););\n }\n return out;\n}\n", "source": "/kaggle/working/apex/csrc/welford.cu"} {"input": "#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"nccl.h\"\n\n/*\n * This file implements a crude but effective mechanism for copying data between tenors owned by different ranks\n * on the same machine using cudaMemcpyAsync peer-to-peer transfers.\n */\n\nnamespace {\n\n__global__ void AddDelay_kernel(const int delay, int* counter) {\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n // waste time while doing something compiler can't predict, thus preventing it from optimizing away this code.\n int new_counter = 0;\n double elapsed = 0;\n clock_t start = clock();\n do {\n clock_t now = clock();\n elapsed = (double)(now - start) * 1e9 / CLOCKS_PER_SEC;\n ++new_counter;\n } while (elapsed < (double)delay);\n *counter = new_counter;\n }\n}\n\nclass NcclCommWrapper {\n private:\n ncclComm_t comm;\n int rank, world_size;\n\n ncclDataType_t get_nccl_type(at::Tensor input) {\n switch (input.scalar_type()) {\n case at::ScalarType::Half:\n return ncclFloat16;\n case at::ScalarType::Float:\n return ncclFloat32;\n case at::ScalarType::Double:\n return ncclFloat64;\n case at::ScalarType::Byte:\n return ncclUint8;\n case at::ScalarType::Char:\n return ncclInt8;\n case at::ScalarType::Int:\n return ncclInt32;\n case at::ScalarType::Long:\n return ncclInt64;\n case at::ScalarType::BFloat16:\n return ncclBfloat16;\n default:\n assert(false);\n }\n }\n\n public:\n NcclCommWrapper() {\n memset(&comm, 0, sizeof(ncclComm_t));\n rank = 0;\n world_size = 0;\n }\n NcclCommWrapper(ncclUniqueId id, int my_rank, int num_ranks) {\n ncclCommInitRank(&comm, num_ranks, id, my_rank);\n rank = my_rank;\n world_size = num_ranks;\n }\n\n ~NcclCommWrapper() {\n printf(\"ncclCommDestroy()\\n\");\n ncclCommDestroy(comm);\n }\n\n void left_right_halo_exchange_inplace(int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo, at::Tensor left_input_halo,\n at::Tensor right_input_halo) {\n auto stream = at::cuda::getCurrentCUDAStream();\n ncclGroupStart();\n ncclDataType_t ncclType = get_nccl_type(left_output_halo);\n bool left_zero = (left_rank < 0);\n bool right_zero = (right_rank < 0);\n size_t left_n = torch::numel(left_output_halo);\n size_t right_n = torch::numel(right_output_halo);\n assert(left_n > 0 && left_n == right_n);\n if (left_zero) {\n left_input_halo.zero_();\n } else {\n AT_DISPATCH_ALL_TYPES_AND3(\n at::ScalarType::Bool, at::ScalarType::BFloat16, at::ScalarType::Half, left_output_halo.scalar_type(),\n \"left_halo_exch\", [&]() {\n // send left (to my_rank - 1)\n ncclSend(left_output_halo.data_ptr(), left_n, ncclType, left_rank, comm, stream);\n // receive left (from my_rank - 1)\n ncclRecv(left_input_halo.data_ptr(), right_n, ncclType, left_rank, comm, stream);\n });\n }\n if (right_zero) {\n right_input_halo.zero_();\n } else {\n AT_DISPATCH_ALL_TYPES_AND3(\n at::ScalarType::Bool, at::ScalarType::BFloat16, at::ScalarType::Half, right_output_halo.scalar_type(),\n \"right_halo_exch\", [&]() {\n // send right (to my_rank + 1 )\n ncclSend(right_output_halo.data_ptr(), right_n, ncclType, right_rank, comm, stream);\n // receive right (from my_rank + 1)\n ncclRecv(right_input_halo.data_ptr(), left_n, ncclType, right_rank, comm, stream);\n });\n }\n ncclGroupEnd();\n }\n\n std::vector left_right_halo_exchange(int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo) {\n // after halo exchange:\n // left_output_halo of rank+1 ends up in right_input_halo of rank\n // right_output_halo of rank-1 ends up in left_input_halo of rank\n auto right_input_halo = torch::empty_like(left_output_halo);\n auto left_input_halo = torch::empty_like(right_output_halo);\n left_right_halo_exchange_inplace(left_rank, right_rank, left_output_halo, right_output_halo, left_input_halo,\n right_input_halo);\n return {left_input_halo, right_input_halo};\n }\n};\n\nclass ManagedObjects {\n public:\n ManagedObjects() {}\n ~ManagedObjects() {\n for (auto it = _nccl_comms.begin(); it != _nccl_comms.end(); ++it) {\n delete *it;\n }\n }\n\n int add_comm(NcclCommWrapper* comm) {\n int handle = _nccl_comms.size();\n _nccl_comms.push_back(comm);\n return handle;\n }\n\n NcclCommWrapper& get_comm(int handle) {\n assert(handle >= 0 && handle < _nccl_comms.size());\n return *_nccl_comms[handle];\n }\n\n private:\n std::vector _nccl_comms;\n};\nclass ManagedObjects mo;\n\n} // end anonymous namespace\n\nnamespace apex {\nnamespace contrib {\nnamespace nccl_p2p {\n\nat::Tensor get_unique_nccl_id(int n) {\n ncclUniqueId id;\n ncclGetUniqueId(&id);\n auto id_tensor = torch::empty({n, (int)sizeof(ncclUniqueId)},\n torch::dtype(torch::kUInt8).device(torch::kCPU).requires_grad(false));\n auto id_ptr = id_tensor.data_ptr();\n size_t offset = 0;\n for (int i = 0; i < n; ++i) {\n ncclUniqueId id;\n ncclGetUniqueId(&id);\n memcpy(id_ptr + offset, &id, sizeof(ncclUniqueId));\n offset += sizeof(ncclUniqueId);\n }\n return id_tensor;\n}\n\nint init_nccl_comm(at::Tensor unique_nccl_id, int my_rank, int num_ranks) {\n ncclUniqueId id;\n auto unique_nccl_id_ptr = unique_nccl_id.data_ptr();\n memcpy(&id, unique_nccl_id_ptr, sizeof(ncclUniqueId));\n NcclCommWrapper* comm = new NcclCommWrapper(id, my_rank, num_ranks);\n int handle = mo.add_comm(comm);\n comm = 0L;\n return handle;\n}\n\nvoid left_right_halo_exchange_inplace(int handle, int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo, at::Tensor left_input_halo,\n at::Tensor right_input_halo) {\n class NcclCommWrapper& communicator = mo.get_comm(handle);\n return communicator.left_right_halo_exchange_inplace(left_rank, right_rank, left_output_halo, right_output_halo,\n left_input_halo, right_input_halo);\n}\n\nstd::vector left_right_halo_exchange(int handle, int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo) {\n class NcclCommWrapper& communicator = mo.get_comm(handle);\n return communicator.left_right_halo_exchange(left_rank, right_rank, left_output_halo, right_output_halo);\n}\n\nvoid add_delay(int delay) {\n auto stream = at::cuda::getCurrentCUDAStream();\n auto t = torch::empty({1}, torch::dtype(torch::kInt32).device(torch::kCUDA));\n AddDelay_kernel<<<1, 1, 0, stream>>>(delay, t.data_ptr());\n}\n\n} // namespace nccl_p2p\n} // namespace contrib\n} // namespace apex\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"nccl.h\"\n\n/*\n * This file implements a crude but effective mechanism for copying data between tenors owned by different ranks\n * on the same machine using cudaMemcpyAsync peer-to-peer transfers.\n */\n\nnamespace {\n\n__global__ void AddDelay_kernel(const int delay, int* counter) {\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n // waste time while doing something compiler can't predict, thus preventing it from optimizing away this code.\n int new_counter = 0;\n double elapsed = 0;\n clock_t start = clock();\n do {\n clock_t now = clock();\n elapsed = (double)(now - start) * 1e9 / CLOCKS_PER_SEC;\n ++new_counter;\n } while (elapsed < (double)delay);\n *counter = new_counter;\n }\n}\n\nclass NcclCommWrapper {\n private:\n ncclComm_t comm;\n int rank, world_size;\n\n ncclDataType_t get_nccl_type(at::Tensor input) {\n switch (input.scalar_type()) {\n case at::ScalarType::Half:\n return ncclFloat16;\n case at::ScalarType::Float:\n return ncclFloat32;\n case at::ScalarType::Double:\n return ncclFloat64;\n case at::ScalarType::Byte:\n return ncclUint8;\n case at::ScalarType::Char:\n return ncclInt8;\n case at::ScalarType::Int:\n return ncclInt32;\n case at::ScalarType::Long:\n return ncclInt64;\n case at::ScalarType::BFloat16:\n return ncclBfloat16;\n default:\n assert(false);\n }\n }\n\n public:\n NcclCommWrapper() {\n memset(&comm, 0, sizeof(ncclComm_t));\n rank = 0;\n world_size = 0;\n }\n NcclCommWrapper(ncclUniqueId id, int my_rank, int num_ranks) {\n ncclCommInitRank(&comm, num_ranks, id, my_rank);\n rank = my_rank;\n world_size = num_ranks;\n }\n\n ~NcclCommWrapper() {\n printf(\"ncclCommDestroy()\\n\");\n ncclCommDestroy(comm);\n }\n\n void left_right_halo_exchange_inplace(int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo, at::Tensor left_input_halo,\n at::Tensor right_input_halo) {\n auto stream = at::cuda::getCurrentCUDAStream();\n ncclGroupStart();\n ncclDataType_t ncclType = get_nccl_type(left_output_halo);\n bool left_zero = (left_rank < 0);\n bool right_zero = (right_rank < 0);\n size_t left_n = torch::numel(left_output_halo);\n size_t right_n = torch::numel(right_output_halo);\n assert(left_n > 0 && left_n == right_n);\n if (left_zero) {\n left_input_halo.zero_();\n } else {\n AT_DISPATCH_ALL_TYPES_AND3(\n at::ScalarType::Bool, at::ScalarType::BFloat16, at::ScalarType::Half, left_output_halo.scalar_type(),\n \"left_halo_exch\", [&]() {\n // send left (to my_rank - 1)\n ncclSend(left_output_halo.data_ptr(), left_n, ncclType, left_rank, comm, stream);\n // receive left (from my_rank - 1)\n ncclRecv(left_input_halo.data_ptr(), right_n, ncclType, left_rank, comm, stream);\n });\n }\n if (right_zero) {\n right_input_halo.zero_();\n } else {\n AT_DISPATCH_ALL_TYPES_AND3(\n at::ScalarType::Bool, at::ScalarType::BFloat16, at::ScalarType::Half, right_output_halo.scalar_type(),\n \"right_halo_exch\", [&]() {\n // send right (to my_rank + 1 )\n ncclSend(right_output_halo.data_ptr(), right_n, ncclType, right_rank, comm, stream);\n // receive right (from my_rank + 1)\n ncclRecv(right_input_halo.data_ptr(), left_n, ncclType, right_rank, comm, stream);\n });\n }\n ncclGroupEnd();\n }\n\n std::vector left_right_halo_exchange(int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo) {\n // after halo exchange:\n // left_output_halo of rank+1 ends up in right_input_halo of rank\n // right_output_halo of rank-1 ends up in left_input_halo of rank\n auto right_input_halo = torch::empty_like(left_output_halo);\n auto left_input_halo = torch::empty_like(right_output_halo);\n left_right_halo_exchange_inplace(left_rank, right_rank, left_output_halo, right_output_halo, left_input_halo,\n right_input_halo);\n return {left_input_halo, right_input_halo};\n }\n};\n\nclass ManagedObjects {\n public:\n ManagedObjects() {}\n ~ManagedObjects() {\n for (auto it = _nccl_comms.begin(); it != _nccl_comms.end(); ++it) {\n delete *it;\n }\n }\n\n int add_comm(NcclCommWrapper* comm) {\n int handle = _nccl_comms.size();\n _nccl_comms.push_back(comm);\n return handle;\n }\n\n NcclCommWrapper& get_comm(int handle) {\n assert(handle >= 0 && handle < _nccl_comms.size());\n return *_nccl_comms[handle];\n }\n\n private:\n std::vector _nccl_comms;\n};\nclass ManagedObjects mo;\n\n} // end anonymous namespace\n\nnamespace apex {\nnamespace contrib {\nnamespace nccl_p2p {\n\nat::Tensor get_unique_nccl_id(int n) {\n ncclUniqueId id;\n ncclGetUniqueId(&id);\n auto id_tensor = torch::empty({n, (int)sizeof(ncclUniqueId)},\n torch::dtype(torch::kUInt8).device(torch::kCPU).requires_grad(false));\n auto id_ptr = id_tensor.data_ptr();\n size_t offset = 0;\n for (int i = 0; i < n; ++i) {\n ncclUniqueId id;\n ncclGetUniqueId(&id);\n memcpy(id_ptr + offset, &id, sizeof(ncclUniqueId));\n offset += sizeof(ncclUniqueId);\n }\n return id_tensor;\n}\n\nint init_nccl_comm(at::Tensor unique_nccl_id, int my_rank, int num_ranks) {\n ncclUniqueId id;\n auto unique_nccl_id_ptr = unique_nccl_id.data_ptr();\n memcpy(&id, unique_nccl_id_ptr, sizeof(ncclUniqueId));\n NcclCommWrapper* comm = new NcclCommWrapper(id, my_rank, num_ranks);\n int handle = mo.add_comm(comm);\n comm = 0L;\n return handle;\n}\n\nvoid left_right_halo_exchange_inplace(int handle, int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo, at::Tensor left_input_halo,\n at::Tensor right_input_halo) {\n class NcclCommWrapper& communicator = mo.get_comm(handle);\n return communicator.left_right_halo_exchange_inplace(left_rank, right_rank, left_output_halo, right_output_halo,\n left_input_halo, right_input_halo);\n}\n\nstd::vector left_right_halo_exchange(int handle, int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo) {\n class NcclCommWrapper& communicator = mo.get_comm(handle);\n return communicator.left_right_halo_exchange(left_rank, right_rank, left_output_halo, right_output_halo);\n}\n\nvoid add_delay(int delay) {\n auto stream = at::cuda::getCurrentCUDAStream();\n auto t = torch::empty({1}, torch::dtype(torch::kInt32).device(torch::kCUDA));\n AddDelay_kernel<<<1, 1, 0, stream>>>(delay, t.data_ptr());\n}\n\n} // namespace nccl_p2p\n} // namespace contrib\n} // namespace apex\n", "source": "/kaggle/working/apex/apex/contrib/csrc/nccl_p2p/nccl_p2p_cuda.cu"} {"input": "#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"nccl.h\"\n\n#define CUDACHECK(cmd) \\\n do { \\\n cudaError_t err = cmd; \\\n if (err != cudaSuccess) { \\\n char hostname[1024]; \\\n gethostname(hostname, 1024); \\\n printf(\"%s: CUDA failure %s:%d '%s'\\n\", hostname, __FILE__, __LINE__, cudaGetErrorString(err)); \\\n } \\\n } while (0)\n\nnamespace {\n\nconstexpr int THREADS_PER_CTA = 128;\n\n/* Basic deleter function for from_blob function.\nvoid deleter(void* ptr)\n{\n printf(\"deleter(ptr=%p)\\n\",ptr);\n cudaFree(ptr);\n}\n*/\n\ntemplate \nat::Tensor blob_view(T* raw_ptr, std::vector shape, const at::TensorOptions& options, bool channels_last) {\n size_t size = 1;\n std::vector strides(shape.size());\n if (channels_last) {\n assert(shape.size() == 4);\n strides[0] = shape[1] * shape[2] * shape[3];\n strides[1] = 1;\n strides[2] = shape[1] * shape[3];\n strides[3] = shape[1];\n } else {\n int idx = strides.size();\n for (auto it = shape.rbegin(); it != shape.rend(); ++it) {\n strides[--idx] = size;\n size *= *it;\n }\n }\n size *= sizeof(T);\n // TODO: Implement dynamic reuse of pooled peer memory.\n // We provide no deleter function because all peer memory allocations are static in this implementation.\n return torch::from_blob((void*)raw_ptr, shape, strides, 0L, options);\n}\n\nvoid tensor_shape(at::Tensor t, bool explicit_nhwc, int& N, int& C, int& H, int& W) {\n if (t.dim() == 3) {\n N = 1;\n if (explicit_nhwc) {\n C = t.size(2);\n H = t.size(0);\n W = t.size(1);\n } else {\n C = t.size(0);\n H = t.size(1);\n W = t.size(2);\n }\n } else if (t.dim() == 4) {\n if (explicit_nhwc) {\n N = t.size(0);\n C = t.size(3);\n H = t.size(1);\n W = t.size(2);\n } else {\n N = t.size(0);\n C = t.size(1);\n H = t.size(2);\n W = t.size(3);\n }\n } else {\n printf(\"%s;%d - t.dim() must be either 3 or 4 (was %d)\\n\", __FILE__, __LINE__, int(t.dim()));\n assert(t.dim() == 3 || t.dim() == 4);\n }\n}\n\nvoid tensor_strides(at::Tensor t, bool explicit_nhwc, int& stride_N, int& stride_C, int& stride_H, int& stride_W) {\n if (t.dim() == 3) {\n if (explicit_nhwc) {\n stride_C = t.stride(2);\n stride_H = t.stride(0);\n stride_W = t.stride(1);\n } else {\n stride_C = t.stride(0);\n stride_H = t.stride(1);\n stride_W = t.stride(2);\n }\n stride_N = t.size(0) * t.size(1) * t.size(2);\n } else if (t.dim() == 4) {\n if (explicit_nhwc) {\n stride_N = t.stride(0);\n stride_C = t.stride(3);\n stride_H = t.stride(1);\n stride_W = t.stride(2);\n } else {\n stride_N = t.stride(0);\n stride_C = t.stride(1);\n stride_H = t.stride(2);\n stride_W = t.stride(3);\n }\n } else {\n printf(\"%s;%d - t.dim() must be either 3 or 4 (was %d)\\n\", __FILE__, __LINE__, t.dim());\n assert(t.dim() == 3 || t.dim() == 4);\n }\n}\n\ntemplate \ninline __device__ void __zero(T* dst) {\n *dst = T(0);\n}\n\ninline __device__ void __zero(int2* dst) { *dst = {0, 0}; }\n\ntemplate \ninline __device__ void zero_tensor(const int dim0, const int dim1, const int dim2, T* __restrict__ data,\n const int data_stride0, const int data_stride1, const int data_stride2,\n const int thread_id, const int block_id, const int num_blocks) {\n const int global_id = thread_id + block_id * THREADS_PER_CTA;\n const int num_threads = num_blocks * THREADS_PER_CTA;\n const int count = dim0 * dim1 * dim2;\n for (int i = global_id; i < count; i += num_threads) {\n int offset;\n if (contiguous) {\n offset = i;\n } else {\n const int j2 = i % dim2;\n const int k = i / dim2;\n const int j1 = k % dim1;\n const int j0 = k / dim1;\n offset = j0 * data_stride0 + j1 * data_stride1 + j2 * data_stride2;\n }\n __zero(data + offset);\n }\n}\n\ntemplate \ninline __device__ void push_pull_tensor(const int dim0, const int dim1, const int dim2, const T* __restrict__ data_in,\n const int data_in_stride0, const int data_in_stride1, const int data_in_stride2,\n T* __restrict__ data_out, const int data_out_stride0,\n const int data_out_stride1, const int data_out_stride2, int4* local_peer,\n int4* remote_peer, const int thread_id, const int block_id,\n const int num_blocks) {\n // 128b=16B NVLink flit\n // Note: Use last 4B as a semaphore\n static_assert(sizeof(T) <= 12);\n union Flit {\n T payload;\n uint uints[4];\n };\n // Communication bit indicates whether flit has been received from\n // a remote GPU\n constexpr uint communication_mask = 1 << 0;\n // Status bit is used to choose the active peer buffer in an\n // alternating double buffer scheme. We use buffer 1 if the bits\n // match, use buffer 2 if the bits differ, and invert the bit\n // after finishing with a buffer.\n constexpr uint status_mask = 1 << 1;\n\n // Split peer memory into two sets of buffers\n // Note: Each block owns a THREADS_PER_CTA*2*16B chunk of peer\n // memory\n const int peer_offset1 = block_id * THREADS_PER_CTA * 2 + thread_id;\n const int peer_offset2 = peer_offset1 + THREADS_PER_CTA;\n volatile int* local_peer1 = reinterpret_cast(local_peer + peer_offset1);\n volatile int* local_peer2 = reinterpret_cast(local_peer + peer_offset2);\n volatile int* remote_peer1 = reinterpret_cast(remote_peer + peer_offset1);\n volatile int* remote_peer2 = reinterpret_cast(remote_peer + peer_offset2);\n\n // Iterate through tensor entries\n const int num_threads = num_blocks * THREADS_PER_CTA;\n const int count = dim0 * dim1 * dim2;\n for (int i0 = block_id * THREADS_PER_CTA; i0 < count; i0 += num_threads) {\n const int i = i0 + thread_id;\n const bool has_data = i < count;\n\n // Calculate buffer positions\n int data_in_offset, data_out_offset;\n if (contiguous) {\n data_in_offset = i;\n data_out_offset = i;\n } else {\n const int j2 = i % dim2;\n const int k = i / dim2;\n const int j1 = k % dim1;\n const int j0 = k / dim1;\n data_in_offset = j0 * data_in_stride0 + j1 * data_in_stride1 + j2 * data_in_stride2;\n data_out_offset = j0 * data_out_stride0 + j1 * data_out_stride1 + j2 * data_out_stride2;\n }\n\n // Determine which peer memory buffer to use\n // Note: The status bit is not affected by asynchronous\n // communication from the remote GPU.\n Flit local_message1, local_message2;\n asm volatile(\"ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(local_message1.uints[0]), \"=r\"(local_message1.uints[1]), \"=r\"(local_message1.uints[2]),\n \"=r\"(local_message1.uints[3])\n : \"l\"(local_peer1)\n : \"memory\");\n asm volatile(\"ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(local_message2.uints[0]), \"=r\"(local_message2.uints[1]), \"=r\"(local_message2.uints[2]),\n \"=r\"(local_message2.uints[3])\n : \"l\"(local_peer2)\n : \"memory\");\n const uint status1 = local_message1.uints[3] & status_mask;\n const uint status2 = local_message2.uints[3] & status_mask;\n const bool peer1_is_active = (status1 ^ status2) == 0;\n volatile int* ox = peer1_is_active ? remote_peer1 : remote_peer2;\n volatile int* ix = peer1_is_active ? local_peer1 : local_peer2;\n const uint status = peer1_is_active ? status1 : status2;\n Flit recv_message = peer1_is_active ? local_message1 : local_message2;\n\n // Send flit to remote GPU\n // Note: Set communication bit and keep status bit\n Flit send_message;\n if (has_data) {\n send_message.payload = data_in[data_in_offset];\n }\n send_message.uints[3] = communication_mask | status;\n asm volatile(\"st.volatile.global.v4.u32 [%0], {%1,%2,%3,%4};\" ::\"l\"(ox), \"r\"(send_message.uints[0]),\n \"r\"(send_message.uints[1]), \"r\"(send_message.uints[2]), \"r\"(send_message.uints[3])\n : \"memory\");\n\n // Recieve flit from peer\n while ((recv_message.uints[3] & communication_mask) == 0) {\n asm volatile(\"ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(recv_message.uints[0]), \"=r\"(recv_message.uints[1]), \"=r\"(recv_message.uints[2]),\n \"=r\"(recv_message.uints[3])\n : \"l\"(ix)\n : \"memory\");\n }\n if (has_data) {\n data_out[data_out_offset] = recv_message.payload;\n }\n\n // Reset semaphore\n // Note: Clear communication bit and invert status bit\n uint flag = ~status & status_mask;\n asm volatile(\"st.volatile.global.v4.u32 [%0], {%1,%2,%3,%4};\" ::\"l\"(ix), \"n\"(0), \"n\"(0), \"n\"(0), \"r\"(flag)\n : \"memory\");\n if (i0 + num_threads < count) {\n __threadfence_system();\n }\n }\n}\n\ntemplate \n#if __CUDA_ARCH__ >= 700\n__launch_bounds__(THREADS_PER_CTA)\n#endif\n __global__ void push_pull_halos_1d_kernel(\n // top halo,\n T* toh, int toh_stride0, int toh_stride1, int toh_stride2, // top output halo (local)\n const T* tih, int tih_stride0, int tih_stride1, int tih_stride2, // top input halo (local)\n int4* tox, // top output transfer buffer (remote peer)\n int4* tix, // top input transfer buffer (local peer)\n // btm halo\n T* boh, int boh_stride0, int boh_stride1, int boh_stride2, // btm output halo (local)\n const T* bih, int bih_stride0, int bih_stride1, int bih_stride2, // btm input halo (local)\n int4* box, // btm output transfer buffer (remote peer)\n int4* bix, // btm input transfer buffer (local peer)\n // dimensions\n int dim0, int dim1, int dim2,\n bool top_first // whether to launch communicate top halo first\n ) {\n const int num_blocks_side = gridDim.x / 2;\n const int block_id_side = (blockIdx.x < num_blocks_side ? blockIdx.x : blockIdx.x - num_blocks_side);\n const bool in_top_block = top_first == (blockIdx.x < num_blocks_side);\n if (in_top_block) {\n if (top_zero) {\n zero_tensor(dim0, dim1, dim2, toh, toh_stride0, toh_stride1, toh_stride2, threadIdx.x,\n block_id_side, num_blocks_side);\n } else {\n push_pull_tensor(dim0, dim1, dim2, tih, tih_stride0, tih_stride1, tih_stride2, toh, toh_stride0,\n toh_stride1, toh_stride2, tix, tox, threadIdx.x, block_id_side, num_blocks_side);\n }\n } else {\n if (btm_zero) {\n zero_tensor(dim0, dim1, dim2, boh, boh_stride0, boh_stride1, boh_stride2, threadIdx.x,\n block_id_side, num_blocks_side);\n } else {\n push_pull_tensor(dim0, dim1, dim2, bih, bih_stride0, bih_stride1, bih_stride2, boh, boh_stride0,\n boh_stride1, boh_stride2, bix, box, threadIdx.x, block_id_side, num_blocks_side);\n }\n }\n}\n\n__global__ void delay_kernel(int delay_nanoseconds, int* counter) {\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n // waste time while doing something compiler can't predict, thus preventing it from optimizing away this code.\n int new_counter = 0;\n double elapsed = 0;\n clock_t start = clock();\n do {\n clock_t now = clock();\n elapsed = (double)(now - start) * 1e9 / CLOCKS_PER_SEC;\n ++new_counter;\n } while (elapsed < (double)delay_nanoseconds);\n *counter = new_counter;\n }\n}\n\n} // namespace\n\nnamespace apex {\nnamespace contrib {\nnamespace peer_memory {\n\nint64_t allocate_raw(int64_t size) {\n float* ptr = 0L;\n cudaMalloc(&ptr, size);\n cudaMemset(ptr, 0, size);\n return (int64_t)ptr;\n}\n\nvoid free_raw(int64_t raw) { cudaFree((void*)raw); }\n\nvoid zero(int64_t raw, int64_t size) { cudaMemset((void*)raw, 0, size); }\n\nat::Tensor get_raw_ipc_address(int64_t raw) {\n cudaIpcMemHandle_t mem_handle;\n CUDACHECK(cudaIpcGetMemHandle(&mem_handle, (void*)raw));\n const int n = sizeof(cudaIpcMemHandle_t);\n auto address_tensor = torch::empty({n}, torch::dtype(torch::kUInt8));\n auto address_tensor_p = address_tensor.data_ptr();\n memcpy(address_tensor_p, (uint8_t*)&mem_handle, n);\n return address_tensor;\n}\n\nstd::vector get_raw_peers(at::Tensor ipc_addresses, int peer_rank, int64_t raw) {\n int peer_group_size = ipc_addresses.size(0);\n std::vector results(peer_group_size);\n for (int i = 0; i < peer_group_size; ++i) {\n if (i != peer_rank) {\n cudaIpcMemHandle_t mem_handle;\n memcpy(&mem_handle, ipc_addresses.index({i}).data_ptr(), sizeof(cudaIpcMemHandle_t));\n void* p = 0L;\n CUDACHECK(cudaIpcOpenMemHandle((void**)&p, mem_handle, cudaIpcMemLazyEnablePeerAccess));\n results[i] = (int64_t)p;\n } else {\n results[i] = (int64_t)raw;\n }\n }\n return results;\n}\n\nat::Tensor blob_view_half(int64_t raw, std::vector shape, bool channels_last) {\n return blob_view((at::Half*)raw, shape, torch::dtype(torch::kFloat16).device(torch::kCUDA), channels_last);\n}\n\nat::Tensor blob_view_float(int64_t raw, std::vector shape, bool channels_last) {\n return blob_view((float*)raw, shape, torch::dtype(torch::kFloat32).device(torch::kCUDA), channels_last);\n}\n\nat::Tensor blob_view_int(int64_t raw, std::vector shape, bool channels_last) {\n return blob_view((int*)raw, shape, torch::dtype(torch::kInt32).device(torch::kCUDA), channels_last);\n}\n\nvoid push_pull_halos_1d(\n bool diagnostics, bool explicit_nhwc,\n int numSM, // number of SMs to use (zero corresponds to all SMs)\n int rank, // rank in spatial parallel group\n bool top_zero, // if top halo should be zeroed\n at::Tensor top_in_halo, // top input halo buffer (in local device memory, sent to top neighbor)\n at::Tensor top_in_transfer, // top input transfer buffer (in local peer memory)\n at::Tensor top_out_transfer, // top output transfer buffer (in top neighbor peer memory)\n at::Tensor top_out_halo, // top output halo buffer (in local device memory, received from top neighbor)\n bool btm_zero, // if btm halo should be zeroed\n at::Tensor btm_in_halo, // btm input halo buffer (in local device memory, sent to btm neighbor)\n at::Tensor btm_in_transfer, // btm input transfer buffer (in local peer memory)\n at::Tensor btm_out_transfer, // btm output transfer buffer (in btm neighbor peer memory)\n at::Tensor btm_out_halo // btm output halo buffer (in local device memory, received from btm neighbor)\n) {\n // basic checks of inputs\n TORCH_CHECK(!(top_zero && btm_zero));\n TORCH_CHECK(top_in_halo.is_cuda());\n TORCH_CHECK(top_out_transfer.is_cuda());\n TORCH_CHECK(top_in_transfer.is_cuda());\n TORCH_CHECK(top_out_halo.is_cuda());\n TORCH_CHECK(btm_in_halo.is_cuda());\n TORCH_CHECK(btm_out_transfer.is_cuda());\n TORCH_CHECK(btm_in_transfer.is_cuda());\n TORCH_CHECK(btm_out_halo.is_cuda());\n\n // tensor shapes\n int tih_N, tih_C, tih_H, tih_W;\n tensor_shape(top_in_halo, explicit_nhwc, tih_N, tih_C, tih_H, tih_W);\n int toh_N, toh_C, toh_H, toh_W;\n tensor_shape(top_out_halo, explicit_nhwc, toh_N, toh_C, toh_H, toh_W);\n int bih_N, bih_C, bih_H, bih_W;\n tensor_shape(btm_in_halo, explicit_nhwc, bih_N, bih_C, bih_H, bih_W);\n int boh_N, boh_C, boh_H, boh_W;\n tensor_shape(btm_out_halo, explicit_nhwc, boh_N, boh_C, boh_H, boh_W);\n TORCH_CHECK(toh_N == tih_N && tih_N == boh_N && boh_N == bih_N && toh_C == tih_C && tih_C == boh_C &&\n boh_C == bih_C && toh_H == tih_H && tih_H == boh_H && boh_H == bih_H && toh_W == tih_W &&\n tih_W == boh_W && boh_W == bih_W);\n int NN = toh_N, NC = toh_C, NH = toh_H, NW = toh_W;\n if (diagnostics) {\n printf(\"rank %d: NN=%d, NC=%d, NH=%d, NW=%d\\n\", rank, NN, NC, NH, NW);\n }\n TORCH_CHECK(NN == 1);\n\n // tensor strides\n int tih_stride_N, tih_stride_C, tih_stride_H, tih_stride_W;\n tensor_strides(top_in_halo, explicit_nhwc, tih_stride_N, tih_stride_C, tih_stride_H, tih_stride_W);\n int toh_stride_N, toh_stride_C, toh_stride_H, toh_stride_W;\n tensor_strides(top_out_halo, explicit_nhwc, toh_stride_N, toh_stride_C, toh_stride_H, toh_stride_W);\n int bih_stride_N, bih_stride_C, bih_stride_H, bih_stride_W;\n tensor_strides(btm_in_halo, explicit_nhwc, bih_stride_N, bih_stride_C, bih_stride_H, bih_stride_W);\n int boh_stride_N, boh_stride_C, boh_stride_H, boh_stride_W;\n tensor_strides(btm_out_halo, explicit_nhwc, boh_stride_N, boh_stride_C, boh_stride_H, boh_stride_W);\n if (diagnostics) {\n printf(\"rank %d: tih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, tih_stride_N, tih_stride_C, tih_stride_H,\n tih_stride_W);\n printf(\"rank %d: toh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, toh_stride_N, toh_stride_C, toh_stride_H,\n toh_stride_W);\n printf(\"rank %d: bih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, bih_stride_N, bih_stride_C, bih_stride_H,\n bih_stride_W);\n printf(\"rank %d: boh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, boh_stride_N, boh_stride_C, boh_stride_H,\n boh_stride_W);\n }\n\n // determine if nhwc\n bool is_nhwc = (toh_stride_C == 1);\n if (diagnostics) {\n printf(\"rank %d: is_nhwc = %s\\n\", rank, is_nhwc ? \"true\" : \"false\");\n }\n\n // determine if contiguous\n bool contiguous = true;\n if ((NN - 1) * toh_stride_N + (NC - 1) * toh_stride_C + (NH - 1) * toh_stride_H + (NW - 1) * toh_stride_W !=\n NN * NC * NH * NW - 1) {\n contiguous = false;\n }\n if ((NN - 1) * boh_stride_N + (NC - 1) * boh_stride_C + (NH - 1) * boh_stride_H + (NW - 1) * boh_stride_W !=\n NN * NC * NH * NW - 1) {\n contiguous = false;\n }\n if (!top_zero) {\n if (toh_stride_N != tih_stride_N || toh_stride_C != tih_stride_C || toh_stride_H != tih_stride_H ||\n toh_stride_W != tih_stride_W) {\n contiguous = false;\n }\n }\n if (!btm_zero) {\n if (boh_stride_N != bih_stride_N || boh_stride_C != bih_stride_C || boh_stride_H != bih_stride_H ||\n boh_stride_W != bih_stride_W) {\n contiguous = false;\n }\n }\n if (diagnostics) {\n printf(\"rank %d: contiguous = %s\\n\", rank, contiguous ? \"true\" : \"false\");\n }\n\n // determine whether to communicate top halo first\n bool top_first = rank % 2 != 0;\n if (diagnostics) {\n printf(\"rank %d: top_first = %s\\n\", rank, top_first ? \"true\" : \"false\");\n }\n\n // peer memory buffers\n int tox_size = top_out_transfer.numel() * top_out_transfer.element_size();\n int tix_size = top_in_transfer.numel() * top_in_transfer.element_size();\n int box_size = btm_out_transfer.numel() * btm_out_transfer.element_size();\n int bix_size = btm_in_transfer.numel() * btm_in_transfer.element_size();\n if (!top_zero) {\n TORCH_CHECK(top_out_transfer.is_contiguous());\n TORCH_CHECK(top_in_transfer.is_contiguous());\n TORCH_CHECK(tox_size == tix_size);\n }\n if (!btm_zero) {\n TORCH_CHECK(btm_out_transfer.is_contiguous());\n TORCH_CHECK(btm_in_transfer.is_contiguous());\n TORCH_CHECK(box_size == bix_size);\n }\n\n // figure out launch parameters\n int device;\n cudaGetDevice(&device);\n cudaDeviceProp prop;\n cudaGetDeviceProperties(&prop, device);\n if (numSM <= 0 || numSM > prop.multiProcessorCount) {\n numSM = prop.multiProcessorCount;\n }\n auto current_stream = at::cuda::getCurrentCUDAStream();\n dim3 block(THREADS_PER_CTA, 1, 1);\n\n // helper macros to launch templated kernel\n#define LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, TOP_ZERO, BTM_ZERO, KERNEL_ARGS, NUM_ELEMENTS) \\\n do { \\\n /* kernel configuration */ \\\n int numBlocksPerSm; \\\n cudaOccupancyMaxActiveBlocksPerMultiprocessor( \\\n &numBlocksPerSm, push_pull_halos_1d_kernel, THREADS_PER_CTA, 0); \\\n dim3 grid(numSM * numBlocksPerSm, 1, 1); \\\n if (grid.x % 2 != 0) { \\\n /* require even number of blocks (half for top, half for bottom) */ \\\n grid.x -= 1; \\\n } \\\n if ((grid.x / 2) * THREADS_PER_CTA > NUM_ELEMENTS) { \\\n /* only need enough blocks to cover top and bottom halo elements */ \\\n grid.x = 2 * ((NUM_ELEMENTS + THREADS_PER_CTA - 1) / THREADS_PER_CTA); \\\n } \\\n if (!TOP_ZERO) { \\\n /* require 2*128b=32B peer memory per thread */ \\\n if ((grid.x / 2) * THREADS_PER_CTA * 32 > tox_size) { \\\n grid.x = 2 * (tox_size / (THREADS_PER_CTA * 32)); \\\n } \\\n } \\\n if (!BTM_ZERO) { \\\n /* require 2*128b=32B peer memory per thread */ \\\n if ((grid.x / 2) * THREADS_PER_CTA * 32 > box_size) { \\\n grid.x = 2 * (box_size / (THREADS_PER_CTA * 32)); \\\n } \\\n } \\\n TORCH_CHECK(grid.x >= 2); \\\n \\\n /* launch kernel */ \\\n cudaLaunchCooperativeKernel((void*)push_pull_halos_1d_kernel, grid, block, \\\n KERNEL_ARGS, 0, current_stream); \\\n } while (false)\n#define LAUNCH_PUSH_PULL_HALO_KERNEL(T, CONTIGUOUS, KERNEL_ARGS, NUM_ELEMENTS) \\\n do { \\\n if (top_zero) { \\\n LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, true, false, KERNEL_ARGS, NUM_ELEMENTS); \\\n } else if (btm_zero) { \\\n LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, false, true, KERNEL_ARGS, NUM_ELEMENTS); \\\n } else { \\\n LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, false, false, KERNEL_ARGS, NUM_ELEMENTS); \\\n } \\\n } while (false)\n\n AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, top_out_halo.scalar_type(), \"push_pull_halos_1d_kernel\", [&] {\n if (diagnostics) {\n printf(\"rank %d: size(scalar_t) = %ld\\n\", rank, sizeof(scalar_t));\n }\n scalar_t* toh_p = top_out_halo.data_ptr();\n scalar_t* tih_p = top_in_halo.data_ptr();\n int4* tox_p = reinterpret_cast(top_out_transfer.data_ptr());\n int4* tix_p = reinterpret_cast(top_in_transfer.data_ptr());\n scalar_t* boh_p = btm_out_halo.data_ptr();\n scalar_t* bih_p = btm_in_halo.data_ptr();\n int4* box_p = reinterpret_cast(btm_out_transfer.data_ptr());\n int4* bix_p = reinterpret_cast(btm_in_transfer.data_ptr());\n if (diagnostics) printf(\"rank %d: choosing halo exchange kernel\\n\", rank);\n\n // do int2 vector loads if channel count permits\n if (contiguous && (NN * NH * NW * NC * sizeof(scalar_t)) % sizeof(int2) == 0) {\n // can do contiguous int2 transfers\n if (diagnostics) {\n }\n toh_stride_N = toh_stride_H = toh_stride_W = toh_stride_C = 1;\n tih_stride_N = tih_stride_H = tih_stride_W = tih_stride_C = 1;\n boh_stride_N = boh_stride_H = boh_stride_W = boh_stride_C = 1;\n bih_stride_N = bih_stride_H = bih_stride_W = bih_stride_C = 1;\n NC = (NN * NH * NW * NC * sizeof(scalar_t)) / sizeof(int2);\n NN = NH = NW = 1;\n if (diagnostics) {\n printf(\"rank %d: launching contiguous int2 halo exchange kernel\\n\", rank);\n printf(\"rank %d: NC=%d, NH=%d, NW=%d\\n\", rank, NC, NH, NW);\n }\n void* kernel_args[] = {(int2**)&toh_p,\n &toh_stride_H,\n &toh_stride_W,\n &toh_stride_C,\n (int2**)&tih_p,\n &tih_stride_H,\n &tih_stride_W,\n &tih_stride_C,\n &tox_p,\n &tix_p,\n (int2**)&boh_p,\n &boh_stride_H,\n &boh_stride_W,\n &boh_stride_C,\n (int2**)&bih_p,\n &bih_stride_H,\n &bih_stride_W,\n &bih_stride_C,\n &box_p,\n &bix_p,\n &NH,\n &NW,\n &NC,\n &top_first};\n int num_elem = NN * NH * NW * NC;\n LAUNCH_PUSH_PULL_HALO_KERNEL(int2, true, kernel_args, num_elem);\n } else if (is_nhwc && (NC * sizeof(scalar_t)) % sizeof(int2) == 0) {\n // can do strided int2 transfers\n int divisor = sizeof(int2) / sizeof(scalar_t);\n if (diagnostics) {\n printf(\"rank %d: launching strided int2 halo exchange kernel\\n\", rank);\n }\n toh_stride_N /= divisor;\n toh_stride_H /= divisor;\n toh_stride_W /= divisor;\n tih_stride_N /= divisor;\n tih_stride_H /= divisor;\n tih_stride_W /= divisor;\n boh_stride_N /= divisor;\n boh_stride_H /= divisor;\n boh_stride_W /= divisor;\n bih_stride_N /= divisor;\n bih_stride_H /= divisor;\n bih_stride_W /= divisor;\n NC /= divisor;\n if (diagnostics) {\n printf(\"rank %d: divisor=%d\\n\", rank, divisor);\n printf(\"rank %d: tih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, tih_stride_N, tih_stride_C, tih_stride_H,\n tih_stride_W);\n printf(\"rank %d: toh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, toh_stride_N, toh_stride_C, toh_stride_H,\n toh_stride_W);\n printf(\"rank %d: bih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, bih_stride_N, bih_stride_C, bih_stride_H,\n bih_stride_W);\n printf(\"rank %d: boh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, boh_stride_N, boh_stride_C, boh_stride_H,\n boh_stride_W);\n printf(\"rank %d: NC=%d, NH=%d, NW=%d\\n\", rank, NC, NH, NW);\n }\n void* kernel_args[] = {(int2**)&toh_p,\n &toh_stride_H,\n &toh_stride_W,\n &toh_stride_C,\n (int2**)&tih_p,\n &tih_stride_H,\n &tih_stride_W,\n &tih_stride_C,\n &tox_p,\n &tix_p,\n (int2**)&boh_p,\n &boh_stride_H,\n &boh_stride_W,\n &boh_stride_C,\n (int2**)&bih_p,\n &bih_stride_H,\n &bih_stride_W,\n &bih_stride_C,\n &box_p,\n &bix_p,\n &NH,\n &NW,\n &NC,\n &top_first};\n int num_elem = NH * NW * NC;\n LAUNCH_PUSH_PULL_HALO_KERNEL(int2, false, kernel_args, num_elem);\n } else {\n // cannot do int2 transfers\n if (diagnostics) {\n printf(\"rank %d: launching non-int2 halo exchange kernel\\n\", rank);\n }\n int num_elem = NC * NH * NW;\n if (is_nhwc) {\n void* kernel_args[] = {&toh_p, &toh_stride_H, &toh_stride_W, &toh_stride_C, &tih_p, &tih_stride_H,\n &tih_stride_W, &tih_stride_C, &tox_p, &tix_p, &boh_p, &boh_stride_H,\n &boh_stride_W, &boh_stride_C, &bih_p, &bih_stride_H, &bih_stride_W, &bih_stride_C,\n &box_p, &bix_p, &NH, &NW, &NC, &top_first};\n LAUNCH_PUSH_PULL_HALO_KERNEL(scalar_t, false, kernel_args, num_elem);\n } else {\n void* kernel_args[] = {&toh_p, &toh_stride_C, &toh_stride_H, &toh_stride_W, &tih_p, &tih_stride_C,\n &tih_stride_H, &tih_stride_W, &tox_p, &tix_p, &boh_p, &boh_stride_C,\n &boh_stride_H, &boh_stride_W, &bih_p, &bih_stride_C, &bih_stride_H, &bih_stride_W,\n &box_p, &bix_p, &NC, &NH, &NW, &top_first};\n LAUNCH_PUSH_PULL_HALO_KERNEL(scalar_t, false, kernel_args, num_elem);\n }\n }\n });\n\n#undef LAUNCH_PUSH_PULL_HALO_KERNEL_BASE\n#undef LAUNCH_PUSH_PULL_HALO_KERNEL\n}\n\n} // namespace peer_memory\n} // namespace contrib\n} // namespace apex\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"nccl.h\"\n\n#define CUDACHECK(cmd) \\\n do { \\\n hipError_t err = cmd; \\\n if (err != hipSuccess) { \\\n char hostname[1024]; \\\n gethostname(hostname, 1024); \\\n printf(\"%s: CUDA failure %s:%d '%s'\\n\", hostname, __FILE__, __LINE__, hipGetErrorString(err)); \\\n } \\\n } while (0)\n\nnamespace {\n\nconstexpr int THREADS_PER_CTA = 128;\n\n/* Basic deleter function for from_blob function.\nvoid deleter(void* ptr)\n{\n printf(\"deleter(ptr=%p)\\n\",ptr);\n cudaFree(ptr);\n}\n*/\n\ntemplate \nat::Tensor blob_view(T* raw_ptr, std::vector shape, const at::TensorOptions& options, bool channels_last) {\n size_t size = 1;\n std::vector strides(shape.size());\n if (channels_last) {\n assert(shape.size() == 4);\n strides[0] = shape[1] * shape[2] * shape[3];\n strides[1] = 1;\n strides[2] = shape[1] * shape[3];\n strides[3] = shape[1];\n } else {\n int idx = strides.size();\n for (auto it = shape.rbegin(); it != shape.rend(); ++it) {\n strides[--idx] = size;\n size *= *it;\n }\n }\n size *= sizeof(T);\n // TODO: Implement dynamic reuse of pooled peer memory.\n // We provide no deleter function because all peer memory allocations are static in this implementation.\n return torch::from_blob((void*)raw_ptr, shape, strides, 0L, options);\n}\n\nvoid tensor_shape(at::Tensor t, bool explicit_nhwc, int& N, int& C, int& H, int& W) {\n if (t.dim() == 3) {\n N = 1;\n if (explicit_nhwc) {\n C = t.size(2);\n H = t.size(0);\n W = t.size(1);\n } else {\n C = t.size(0);\n H = t.size(1);\n W = t.size(2);\n }\n } else if (t.dim() == 4) {\n if (explicit_nhwc) {\n N = t.size(0);\n C = t.size(3);\n H = t.size(1);\n W = t.size(2);\n } else {\n N = t.size(0);\n C = t.size(1);\n H = t.size(2);\n W = t.size(3);\n }\n } else {\n printf(\"%s;%d - t.dim() must be either 3 or 4 (was %d)\\n\", __FILE__, __LINE__, int(t.dim()));\n assert(t.dim() == 3 || t.dim() == 4);\n }\n}\n\nvoid tensor_strides(at::Tensor t, bool explicit_nhwc, int& stride_N, int& stride_C, int& stride_H, int& stride_W) {\n if (t.dim() == 3) {\n if (explicit_nhwc) {\n stride_C = t.stride(2);\n stride_H = t.stride(0);\n stride_W = t.stride(1);\n } else {\n stride_C = t.stride(0);\n stride_H = t.stride(1);\n stride_W = t.stride(2);\n }\n stride_N = t.size(0) * t.size(1) * t.size(2);\n } else if (t.dim() == 4) {\n if (explicit_nhwc) {\n stride_N = t.stride(0);\n stride_C = t.stride(3);\n stride_H = t.stride(1);\n stride_W = t.stride(2);\n } else {\n stride_N = t.stride(0);\n stride_C = t.stride(1);\n stride_H = t.stride(2);\n stride_W = t.stride(3);\n }\n } else {\n printf(\"%s;%d - t.dim() must be either 3 or 4 (was %d)\\n\", __FILE__, __LINE__, t.dim());\n assert(t.dim() == 3 || t.dim() == 4);\n }\n}\n\ntemplate \ninline __device__ void __zero(T* dst) {\n *dst = T(0);\n}\n\ninline __device__ void __zero(int2* dst) { *dst = {0, 0}; }\n\ntemplate \ninline __device__ void zero_tensor(const int dim0, const int dim1, const int dim2, T* __restrict__ data,\n const int data_stride0, const int data_stride1, const int data_stride2,\n const int thread_id, const int block_id, const int num_blocks) {\n const int global_id = thread_id + block_id * THREADS_PER_CTA;\n const int num_threads = num_blocks * THREADS_PER_CTA;\n const int count = dim0 * dim1 * dim2;\n for (int i = global_id; i < count; i += num_threads) {\n int offset;\n if (contiguous) {\n offset = i;\n } else {\n const int j2 = i % dim2;\n const int k = i / dim2;\n const int j1 = k % dim1;\n const int j0 = k / dim1;\n offset = j0 * data_stride0 + j1 * data_stride1 + j2 * data_stride2;\n }\n __zero(data + offset);\n }\n}\n\ntemplate \ninline __device__ void push_pull_tensor(const int dim0, const int dim1, const int dim2, const T* __restrict__ data_in,\n const int data_in_stride0, const int data_in_stride1, const int data_in_stride2,\n T* __restrict__ data_out, const int data_out_stride0,\n const int data_out_stride1, const int data_out_stride2, int4* local_peer,\n int4* remote_peer, const int thread_id, const int block_id,\n const int num_blocks) {\n // 128b=16B NVLink flit\n // Note: Use last 4B as a semaphore\n static_assert(sizeof(T) <= 12);\n union Flit {\n T payload;\n uint uints[4];\n };\n // Communication bit indicates whether flit has been received from\n // a remote GPU\n constexpr uint communication_mask = 1 << 0;\n // Status bit is used to choose the active peer buffer in an\n // alternating double buffer scheme. We use buffer 1 if the bits\n // match, use buffer 2 if the bits differ, and invert the bit\n // after finishing with a buffer.\n constexpr uint status_mask = 1 << 1;\n\n // Split peer memory into two sets of buffers\n // Note: Each block owns a THREADS_PER_CTA*2*16B chunk of peer\n // memory\n const int peer_offset1 = block_id * THREADS_PER_CTA * 2 + thread_id;\n const int peer_offset2 = peer_offset1 + THREADS_PER_CTA;\n volatile int* local_peer1 = reinterpret_cast(local_peer + peer_offset1);\n volatile int* local_peer2 = reinterpret_cast(local_peer + peer_offset2);\n volatile int* remote_peer1 = reinterpret_cast(remote_peer + peer_offset1);\n volatile int* remote_peer2 = reinterpret_cast(remote_peer + peer_offset2);\n\n // Iterate through tensor entries\n const int num_threads = num_blocks * THREADS_PER_CTA;\n const int count = dim0 * dim1 * dim2;\n for (int i0 = block_id * THREADS_PER_CTA; i0 < count; i0 += num_threads) {\n const int i = i0 + thread_id;\n const bool has_data = i < count;\n\n // Calculate buffer positions\n int data_in_offset, data_out_offset;\n if (contiguous) {\n data_in_offset = i;\n data_out_offset = i;\n } else {\n const int j2 = i % dim2;\n const int k = i / dim2;\n const int j1 = k % dim1;\n const int j0 = k / dim1;\n data_in_offset = j0 * data_in_stride0 + j1 * data_in_stride1 + j2 * data_in_stride2;\n data_out_offset = j0 * data_out_stride0 + j1 * data_out_stride1 + j2 * data_out_stride2;\n }\n\n // Determine which peer memory buffer to use\n // Note: The status bit is not affected by asynchronous\n // communication from the remote GPU.\n Flit local_message1, local_message2;\n asm volatile(\"ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(local_message1.uints[0]), \"=r\"(local_message1.uints[1]), \"=r\"(local_message1.uints[2]),\n \"=r\"(local_message1.uints[3])\n : \"l\"(local_peer1)\n : \"memory\");\n asm volatile(\"ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(local_message2.uints[0]), \"=r\"(local_message2.uints[1]), \"=r\"(local_message2.uints[2]),\n \"=r\"(local_message2.uints[3])\n : \"l\"(local_peer2)\n : \"memory\");\n const uint status1 = local_message1.uints[3] & status_mask;\n const uint status2 = local_message2.uints[3] & status_mask;\n const bool peer1_is_active = (status1 ^ status2) == 0;\n volatile int* ox = peer1_is_active ? remote_peer1 : remote_peer2;\n volatile int* ix = peer1_is_active ? local_peer1 : local_peer2;\n const uint status = peer1_is_active ? status1 : status2;\n Flit recv_message = peer1_is_active ? local_message1 : local_message2;\n\n // Send flit to remote GPU\n // Note: Set communication bit and keep status bit\n Flit send_message;\n if (has_data) {\n send_message.payload = data_in[data_in_offset];\n }\n send_message.uints[3] = communication_mask | status;\n asm volatile(\"st.volatile.global.v4.u32 [%0], {%1,%2,%3,%4};\" ::\"l\"(ox), \"r\"(send_message.uints[0]),\n \"r\"(send_message.uints[1]), \"r\"(send_message.uints[2]), \"r\"(send_message.uints[3])\n : \"memory\");\n\n // Recieve flit from peer\n while ((recv_message.uints[3] & communication_mask) == 0) {\n asm volatile(\"ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(recv_message.uints[0]), \"=r\"(recv_message.uints[1]), \"=r\"(recv_message.uints[2]),\n \"=r\"(recv_message.uints[3])\n : \"l\"(ix)\n : \"memory\");\n }\n if (has_data) {\n data_out[data_out_offset] = recv_message.payload;\n }\n\n // Reset semaphore\n // Note: Clear communication bit and invert status bit\n uint flag = ~status & status_mask;\n asm volatile(\"st.volatile.global.v4.u32 [%0], {%1,%2,%3,%4};\" ::\"l\"(ix), \"n\"(0), \"n\"(0), \"n\"(0), \"r\"(flag)\n : \"memory\");\n if (i0 + num_threads < count) {\n __threadfence_system();\n }\n }\n}\n\ntemplate \n#if __CUDA_ARCH__ >= 700\n__launch_bounds__(THREADS_PER_CTA)\n#endif\n __global__ void push_pull_halos_1d_kernel(\n // top halo,\n T* toh, int toh_stride0, int toh_stride1, int toh_stride2, // top output halo (local)\n const T* tih, int tih_stride0, int tih_stride1, int tih_stride2, // top input halo (local)\n int4* tox, // top output transfer buffer (remote peer)\n int4* tix, // top input transfer buffer (local peer)\n // btm halo\n T* boh, int boh_stride0, int boh_stride1, int boh_stride2, // btm output halo (local)\n const T* bih, int bih_stride0, int bih_stride1, int bih_stride2, // btm input halo (local)\n int4* box, // btm output transfer buffer (remote peer)\n int4* bix, // btm input transfer buffer (local peer)\n // dimensions\n int dim0, int dim1, int dim2,\n bool top_first // whether to launch communicate top halo first\n ) {\n const int num_blocks_side = gridDim.x / 2;\n const int block_id_side = (blockIdx.x < num_blocks_side ? blockIdx.x : blockIdx.x - num_blocks_side);\n const bool in_top_block = top_first == (blockIdx.x < num_blocks_side);\n if (in_top_block) {\n if (top_zero) {\n zero_tensor(dim0, dim1, dim2, toh, toh_stride0, toh_stride1, toh_stride2, threadIdx.x,\n block_id_side, num_blocks_side);\n } else {\n push_pull_tensor(dim0, dim1, dim2, tih, tih_stride0, tih_stride1, tih_stride2, toh, toh_stride0,\n toh_stride1, toh_stride2, tix, tox, threadIdx.x, block_id_side, num_blocks_side);\n }\n } else {\n if (btm_zero) {\n zero_tensor(dim0, dim1, dim2, boh, boh_stride0, boh_stride1, boh_stride2, threadIdx.x,\n block_id_side, num_blocks_side);\n } else {\n push_pull_tensor(dim0, dim1, dim2, bih, bih_stride0, bih_stride1, bih_stride2, boh, boh_stride0,\n boh_stride1, boh_stride2, bix, box, threadIdx.x, block_id_side, num_blocks_side);\n }\n }\n}\n\n__global__ void delay_kernel(int delay_nanoseconds, int* counter) {\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n // waste time while doing something compiler can't predict, thus preventing it from optimizing away this code.\n int new_counter = 0;\n double elapsed = 0;\n clock_t start = clock();\n do {\n clock_t now = clock();\n elapsed = (double)(now - start) * 1e9 / CLOCKS_PER_SEC;\n ++new_counter;\n } while (elapsed < (double)delay_nanoseconds);\n *counter = new_counter;\n }\n}\n\n} // namespace\n\nnamespace apex {\nnamespace contrib {\nnamespace peer_memory {\n\nint64_t allocate_raw(int64_t size) {\n float* ptr = 0L;\n hipMalloc(&ptr, size);\n hipMemset(ptr, 0, size);\n return (int64_t)ptr;\n}\n\nvoid free_raw(int64_t raw) { hipFree((void*)raw); }\n\nvoid zero(int64_t raw, int64_t size) { hipMemset((void*)raw, 0, size); }\n\nat::Tensor get_raw_ipc_address(int64_t raw) {\n hipIpcMemHandle_t mem_handle;\n CUDACHECK(hipIpcGetMemHandle(&mem_handle, (void*)raw));\n const int n = sizeof(hipIpcMemHandle_t);\n auto address_tensor = torch::empty({n}, torch::dtype(torch::kUInt8));\n auto address_tensor_p = address_tensor.data_ptr();\n memcpy(address_tensor_p, (uint8_t*)&mem_handle, n);\n return address_tensor;\n}\n\nstd::vector get_raw_peers(at::Tensor ipc_addresses, int peer_rank, int64_t raw) {\n int peer_group_size = ipc_addresses.size(0);\n std::vector results(peer_group_size);\n for (int i = 0; i < peer_group_size; ++i) {\n if (i != peer_rank) {\n hipIpcMemHandle_t mem_handle;\n memcpy(&mem_handle, ipc_addresses.index({i}).data_ptr(), sizeof(hipIpcMemHandle_t));\n void* p = 0L;\n CUDACHECK(hipIpcOpenMemHandle((void**)&p, mem_handle, hipIpcMemLazyEnablePeerAccess));\n results[i] = (int64_t)p;\n } else {\n results[i] = (int64_t)raw;\n }\n }\n return results;\n}\n\nat::Tensor blob_view_half(int64_t raw, std::vector shape, bool channels_last) {\n return blob_view((at::Half*)raw, shape, torch::dtype(torch::kFloat16).device(torch::kCUDA), channels_last);\n}\n\nat::Tensor blob_view_float(int64_t raw, std::vector shape, bool channels_last) {\n return blob_view((float*)raw, shape, torch::dtype(torch::kFloat32).device(torch::kCUDA), channels_last);\n}\n\nat::Tensor blob_view_int(int64_t raw, std::vector shape, bool channels_last) {\n return blob_view((int*)raw, shape, torch::dtype(torch::kInt32).device(torch::kCUDA), channels_last);\n}\n\nvoid push_pull_halos_1d(\n bool diagnostics, bool explicit_nhwc,\n int numSM, // number of SMs to use (zero corresponds to all SMs)\n int rank, // rank in spatial parallel group\n bool top_zero, // if top halo should be zeroed\n at::Tensor top_in_halo, // top input halo buffer (in local device memory, sent to top neighbor)\n at::Tensor top_in_transfer, // top input transfer buffer (in local peer memory)\n at::Tensor top_out_transfer, // top output transfer buffer (in top neighbor peer memory)\n at::Tensor top_out_halo, // top output halo buffer (in local device memory, received from top neighbor)\n bool btm_zero, // if btm halo should be zeroed\n at::Tensor btm_in_halo, // btm input halo buffer (in local device memory, sent to btm neighbor)\n at::Tensor btm_in_transfer, // btm input transfer buffer (in local peer memory)\n at::Tensor btm_out_transfer, // btm output transfer buffer (in btm neighbor peer memory)\n at::Tensor btm_out_halo // btm output halo buffer (in local device memory, received from btm neighbor)\n) {\n // basic checks of inputs\n TORCH_CHECK(!(top_zero && btm_zero));\n TORCH_CHECK(top_in_halo.is_cuda());\n TORCH_CHECK(top_out_transfer.is_cuda());\n TORCH_CHECK(top_in_transfer.is_cuda());\n TORCH_CHECK(top_out_halo.is_cuda());\n TORCH_CHECK(btm_in_halo.is_cuda());\n TORCH_CHECK(btm_out_transfer.is_cuda());\n TORCH_CHECK(btm_in_transfer.is_cuda());\n TORCH_CHECK(btm_out_halo.is_cuda());\n\n // tensor shapes\n int tih_N, tih_C, tih_H, tih_W;\n tensor_shape(top_in_halo, explicit_nhwc, tih_N, tih_C, tih_H, tih_W);\n int toh_N, toh_C, toh_H, toh_W;\n tensor_shape(top_out_halo, explicit_nhwc, toh_N, toh_C, toh_H, toh_W);\n int bih_N, bih_C, bih_H, bih_W;\n tensor_shape(btm_in_halo, explicit_nhwc, bih_N, bih_C, bih_H, bih_W);\n int boh_N, boh_C, boh_H, boh_W;\n tensor_shape(btm_out_halo, explicit_nhwc, boh_N, boh_C, boh_H, boh_W);\n TORCH_CHECK(toh_N == tih_N && tih_N == boh_N && boh_N == bih_N && toh_C == tih_C && tih_C == boh_C &&\n boh_C == bih_C && toh_H == tih_H && tih_H == boh_H && boh_H == bih_H && toh_W == tih_W &&\n tih_W == boh_W && boh_W == bih_W);\n int NN = toh_N, NC = toh_C, NH = toh_H, NW = toh_W;\n if (diagnostics) {\n printf(\"rank %d: NN=%d, NC=%d, NH=%d, NW=%d\\n\", rank, NN, NC, NH, NW);\n }\n TORCH_CHECK(NN == 1);\n\n // tensor strides\n int tih_stride_N, tih_stride_C, tih_stride_H, tih_stride_W;\n tensor_strides(top_in_halo, explicit_nhwc, tih_stride_N, tih_stride_C, tih_stride_H, tih_stride_W);\n int toh_stride_N, toh_stride_C, toh_stride_H, toh_stride_W;\n tensor_strides(top_out_halo, explicit_nhwc, toh_stride_N, toh_stride_C, toh_stride_H, toh_stride_W);\n int bih_stride_N, bih_stride_C, bih_stride_H, bih_stride_W;\n tensor_strides(btm_in_halo, explicit_nhwc, bih_stride_N, bih_stride_C, bih_stride_H, bih_stride_W);\n int boh_stride_N, boh_stride_C, boh_stride_H, boh_stride_W;\n tensor_strides(btm_out_halo, explicit_nhwc, boh_stride_N, boh_stride_C, boh_stride_H, boh_stride_W);\n if (diagnostics) {\n printf(\"rank %d: tih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, tih_stride_N, tih_stride_C, tih_stride_H,\n tih_stride_W);\n printf(\"rank %d: toh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, toh_stride_N, toh_stride_C, toh_stride_H,\n toh_stride_W);\n printf(\"rank %d: bih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, bih_stride_N, bih_stride_C, bih_stride_H,\n bih_stride_W);\n printf(\"rank %d: boh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, boh_stride_N, boh_stride_C, boh_stride_H,\n boh_stride_W);\n }\n\n // determine if nhwc\n bool is_nhwc = (toh_stride_C == 1);\n if (diagnostics) {\n printf(\"rank %d: is_nhwc = %s\\n\", rank, is_nhwc ? \"true\" : \"false\");\n }\n\n // determine if contiguous\n bool contiguous = true;\n if ((NN - 1) * toh_stride_N + (NC - 1) * toh_stride_C + (NH - 1) * toh_stride_H + (NW - 1) * toh_stride_W !=\n NN * NC * NH * NW - 1) {\n contiguous = false;\n }\n if ((NN - 1) * boh_stride_N + (NC - 1) * boh_stride_C + (NH - 1) * boh_stride_H + (NW - 1) * boh_stride_W !=\n NN * NC * NH * NW - 1) {\n contiguous = false;\n }\n if (!top_zero) {\n if (toh_stride_N != tih_stride_N || toh_stride_C != tih_stride_C || toh_stride_H != tih_stride_H ||\n toh_stride_W != tih_stride_W) {\n contiguous = false;\n }\n }\n if (!btm_zero) {\n if (boh_stride_N != bih_stride_N || boh_stride_C != bih_stride_C || boh_stride_H != bih_stride_H ||\n boh_stride_W != bih_stride_W) {\n contiguous = false;\n }\n }\n if (diagnostics) {\n printf(\"rank %d: contiguous = %s\\n\", rank, contiguous ? \"true\" : \"false\");\n }\n\n // determine whether to communicate top halo first\n bool top_first = rank % 2 != 0;\n if (diagnostics) {\n printf(\"rank %d: top_first = %s\\n\", rank, top_first ? \"true\" : \"false\");\n }\n\n // peer memory buffers\n int tox_size = top_out_transfer.numel() * top_out_transfer.element_size();\n int tix_size = top_in_transfer.numel() * top_in_transfer.element_size();\n int box_size = btm_out_transfer.numel() * btm_out_transfer.element_size();\n int bix_size = btm_in_transfer.numel() * btm_in_transfer.element_size();\n if (!top_zero) {\n TORCH_CHECK(top_out_transfer.is_contiguous());\n TORCH_CHECK(top_in_transfer.is_contiguous());\n TORCH_CHECK(tox_size == tix_size);\n }\n if (!btm_zero) {\n TORCH_CHECK(btm_out_transfer.is_contiguous());\n TORCH_CHECK(btm_in_transfer.is_contiguous());\n TORCH_CHECK(box_size == bix_size);\n }\n\n // figure out launch parameters\n int device;\n hipGetDevice(&device);\n hipDeviceProp_t prop;\n hipGetDeviceProperties(&prop, device);\n if (numSM <= 0 || numSM > prop.multiProcessorCount) {\n numSM = prop.multiProcessorCount;\n }\n auto current_stream = at::cuda::getCurrentCUDAStream();\n dim3 block(THREADS_PER_CTA, 1, 1);\n\n // helper macros to launch templated kernel\n#define LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, TOP_ZERO, BTM_ZERO, KERNEL_ARGS, NUM_ELEMENTS) \\\n do { \\\n /* kernel configuration */ \\\n int numBlocksPerSm; \\\n hipOccupancyMaxActiveBlocksPerMultiprocessor( \\\n &numBlocksPerSm, push_pull_halos_1d_kernel, THREADS_PER_CTA, 0); \\\n dim3 grid(numSM * numBlocksPerSm, 1, 1); \\\n if (grid.x % 2 != 0) { \\\n /* require even number of blocks (half for top, half for bottom) */ \\\n grid.x -= 1; \\\n } \\\n if ((grid.x / 2) * THREADS_PER_CTA > NUM_ELEMENTS) { \\\n /* only need enough blocks to cover top and bottom halo elements */ \\\n grid.x = 2 * ((NUM_ELEMENTS + THREADS_PER_CTA - 1) / THREADS_PER_CTA); \\\n } \\\n if (!TOP_ZERO) { \\\n /* require 2*128b=32B peer memory per thread */ \\\n if ((grid.x / 2) * THREADS_PER_CTA * 32 > tox_size) { \\\n grid.x = 2 * (tox_size / (THREADS_PER_CTA * 32)); \\\n } \\\n } \\\n if (!BTM_ZERO) { \\\n /* require 2*128b=32B peer memory per thread */ \\\n if ((grid.x / 2) * THREADS_PER_CTA * 32 > box_size) { \\\n grid.x = 2 * (box_size / (THREADS_PER_CTA * 32)); \\\n } \\\n } \\\n TORCH_CHECK(grid.x >= 2); \\\n \\\n /* launch kernel */ \\\n hipLaunchCooperativeKernel(reinterpret_cast((void*)push_pull_halos_1d_kernel), grid, block, \\\n KERNEL_ARGS, 0, current_stream); \\\n } while (false)\n#define LAUNCH_PUSH_PULL_HALO_KERNEL(T, CONTIGUOUS, KERNEL_ARGS, NUM_ELEMENTS) \\\n do { \\\n if (top_zero) { \\\n LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, true, false, KERNEL_ARGS, NUM_ELEMENTS); \\\n } else if (btm_zero) { \\\n LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, false, true, KERNEL_ARGS, NUM_ELEMENTS); \\\n } else { \\\n LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, false, false, KERNEL_ARGS, NUM_ELEMENTS); \\\n } \\\n } while (false)\n\n AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, top_out_halo.scalar_type(), \"push_pull_halos_1d_kernel\", [&] {\n if (diagnostics) {\n printf(\"rank %d: size(scalar_t) = %ld\\n\", rank, sizeof(scalar_t));\n }\n scalar_t* toh_p = top_out_halo.data_ptr();\n scalar_t* tih_p = top_in_halo.data_ptr();\n int4* tox_p = reinterpret_cast(top_out_transfer.data_ptr());\n int4* tix_p = reinterpret_cast(top_in_transfer.data_ptr());\n scalar_t* boh_p = btm_out_halo.data_ptr();\n scalar_t* bih_p = btm_in_halo.data_ptr();\n int4* box_p = reinterpret_cast(btm_out_transfer.data_ptr());\n int4* bix_p = reinterpret_cast(btm_in_transfer.data_ptr());\n if (diagnostics) printf(\"rank %d: choosing halo exchange kernel\\n\", rank);\n\n // do int2 vector loads if channel count permits\n if (contiguous && (NN * NH * NW * NC * sizeof(scalar_t)) % sizeof(int2) == 0) {\n // can do contiguous int2 transfers\n if (diagnostics) {\n }\n toh_stride_N = toh_stride_H = toh_stride_W = toh_stride_C = 1;\n tih_stride_N = tih_stride_H = tih_stride_W = tih_stride_C = 1;\n boh_stride_N = boh_stride_H = boh_stride_W = boh_stride_C = 1;\n bih_stride_N = bih_stride_H = bih_stride_W = bih_stride_C = 1;\n NC = (NN * NH * NW * NC * sizeof(scalar_t)) / sizeof(int2);\n NN = NH = NW = 1;\n if (diagnostics) {\n printf(\"rank %d: launching contiguous int2 halo exchange kernel\\n\", rank);\n printf(\"rank %d: NC=%d, NH=%d, NW=%d\\n\", rank, NC, NH, NW);\n }\n void* kernel_args[] = {(int2**)&toh_p,\n &toh_stride_H,\n &toh_stride_W,\n &toh_stride_C,\n (int2**)&tih_p,\n &tih_stride_H,\n &tih_stride_W,\n &tih_stride_C,\n &tox_p,\n &tix_p,\n (int2**)&boh_p,\n &boh_stride_H,\n &boh_stride_W,\n &boh_stride_C,\n (int2**)&bih_p,\n &bih_stride_H,\n &bih_stride_W,\n &bih_stride_C,\n &box_p,\n &bix_p,\n &NH,\n &NW,\n &NC,\n &top_first};\n int num_elem = NN * NH * NW * NC;\n LAUNCH_PUSH_PULL_HALO_KERNEL(int2, true, kernel_args, num_elem);\n } else if (is_nhwc && (NC * sizeof(scalar_t)) % sizeof(int2) == 0) {\n // can do strided int2 transfers\n int divisor = sizeof(int2) / sizeof(scalar_t);\n if (diagnostics) {\n printf(\"rank %d: launching strided int2 halo exchange kernel\\n\", rank);\n }\n toh_stride_N /= divisor;\n toh_stride_H /= divisor;\n toh_stride_W /= divisor;\n tih_stride_N /= divisor;\n tih_stride_H /= divisor;\n tih_stride_W /= divisor;\n boh_stride_N /= divisor;\n boh_stride_H /= divisor;\n boh_stride_W /= divisor;\n bih_stride_N /= divisor;\n bih_stride_H /= divisor;\n bih_stride_W /= divisor;\n NC /= divisor;\n if (diagnostics) {\n printf(\"rank %d: divisor=%d\\n\", rank, divisor);\n printf(\"rank %d: tih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, tih_stride_N, tih_stride_C, tih_stride_H,\n tih_stride_W);\n printf(\"rank %d: toh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, toh_stride_N, toh_stride_C, toh_stride_H,\n toh_stride_W);\n printf(\"rank %d: bih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, bih_stride_N, bih_stride_C, bih_stride_H,\n bih_stride_W);\n printf(\"rank %d: boh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, boh_stride_N, boh_stride_C, boh_stride_H,\n boh_stride_W);\n printf(\"rank %d: NC=%d, NH=%d, NW=%d\\n\", rank, NC, NH, NW);\n }\n void* kernel_args[] = {(int2**)&toh_p,\n &toh_stride_H,\n &toh_stride_W,\n &toh_stride_C,\n (int2**)&tih_p,\n &tih_stride_H,\n &tih_stride_W,\n &tih_stride_C,\n &tox_p,\n &tix_p,\n (int2**)&boh_p,\n &boh_stride_H,\n &boh_stride_W,\n &boh_stride_C,\n (int2**)&bih_p,\n &bih_stride_H,\n &bih_stride_W,\n &bih_stride_C,\n &box_p,\n &bix_p,\n &NH,\n &NW,\n &NC,\n &top_first};\n int num_elem = NH * NW * NC;\n LAUNCH_PUSH_PULL_HALO_KERNEL(int2, false, kernel_args, num_elem);\n } else {\n // cannot do int2 transfers\n if (diagnostics) {\n printf(\"rank %d: launching non-int2 halo exchange kernel\\n\", rank);\n }\n int num_elem = NC * NH * NW;\n if (is_nhwc) {\n void* kernel_args[] = {&toh_p, &toh_stride_H, &toh_stride_W, &toh_stride_C, &tih_p, &tih_stride_H,\n &tih_stride_W, &tih_stride_C, &tox_p, &tix_p, &boh_p, &boh_stride_H,\n &boh_stride_W, &boh_stride_C, &bih_p, &bih_stride_H, &bih_stride_W, &bih_stride_C,\n &box_p, &bix_p, &NH, &NW, &NC, &top_first};\n LAUNCH_PUSH_PULL_HALO_KERNEL(scalar_t, false, kernel_args, num_elem);\n } else {\n void* kernel_args[] = {&toh_p, &toh_stride_C, &toh_stride_H, &toh_stride_W, &tih_p, &tih_stride_C,\n &tih_stride_H, &tih_stride_W, &tox_p, &tix_p, &boh_p, &boh_stride_C,\n &boh_stride_H, &boh_stride_W, &bih_p, &bih_stride_C, &bih_stride_H, &bih_stride_W,\n &box_p, &bix_p, &NC, &NH, &NW, &top_first};\n LAUNCH_PUSH_PULL_HALO_KERNEL(scalar_t, false, kernel_args, num_elem);\n }\n }\n });\n\n#undef LAUNCH_PUSH_PULL_HALO_KERNEL_BASE\n#undef LAUNCH_PUSH_PULL_HALO_KERNEL\n}\n\n} // namespace peer_memory\n} // namespace contrib\n} // namespace apex\n", "source": "/kaggle/working/apex/apex/contrib/csrc/peer_memory/peer_memory_cuda.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\ntemplate \n__device__ __forceinline__ scalar_t logSumExp(scalar_t a, scalar_t b) {\n // standard log-sum-exp trick is used here to provide better numerical stability\n return (a >= b) ? a + std::log1p(exp(b - a)) : b + std::log1p(exp(a - b));\n}\n\n// Vanilla transducer loss function (i.e. forward-backward algorithm)\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n\n// Forward (alpha) and backward (beta) path are launched together. Input is assumed to be converted\n// into log scale by the preceding log_softmax layer\n// Diagonal wavefront advancing usually used in dynamic programming is leveraged here.\n// alpha and beta are of acc_t type, as they are essentially accumulators.\n\n// This loss function supports packed input where a tensor of shape [B, T, U, H] is packed into\n// [B_packed, H].\n// Don't-care region (t > audLen) or (u > txtLen) is removed.\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_forward(const scalar_t* x, const int* label, const int* audLen, const int* txtLen,\n const int64_t* batchOffset,\n int64_t dictSize, // 64-bit indexing for data tensor\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n acc_t* alpha, acc_t* beta, scalar_t* loss) {\n const int batch = blockIdx.y;\n const int tid = threadIdx.x;\n const auto myFLen = audLen[batch];\n // Note that start of the sentence is added as 1 here\n const auto myGLen = txtLen[batch] + 1;\n const auto myLabel = label + batch * (maxGLen - 1);\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n const scalar_t* myX = x + myBatchOffset * dictSize;\n int u = tid;\n\n if (blockIdx.x == 0) {\n // alpha path\n acc_t* myAlpha = alpha + batch * maxFLen * maxGLen;\n if (u == 0) myAlpha[0] = 0;\n __syncthreads();\n\n for (int64_t step = 1; step < myFLen + myGLen - 1; ++step) {\n // Move along the diagonal wavefront to leverage available parallelism\n for (u = tid; u < myGLen; u += blockDim.x) {\n int64_t t = step - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(16) in [1]\n if (u == 0) {\n // alpha(t, u) = alpha(t-1, u) * null(t-1, u)\n myAlpha[t * maxGLen + u] = myAlpha[(t - 1) * maxGLen] + myX[((t - 1) * myStrideT) * dictSize + blankIdx];\n } else if (t == 0) {\n // alpha(t, u-1) = alpha(t, u-1) * y(t, u-1)\n myAlpha[u] = myAlpha[u - 1] + myX[(u - 1) * dictSize + myLabel[u - 1]];\n } else {\n // alpha(t, u) = alpha(t-1, u) * null(t-1, u) + alpha(t, u-1) * y(t, u-1)\n acc_t current = myAlpha[(t - 1) * maxGLen + u] + myX[((t - 1) * myStrideT + u) * dictSize + blankIdx];\n acc_t next = myAlpha[t * maxGLen + u - 1] + myX[(t * myStrideT + u - 1) * dictSize + myLabel[u - 1]];\n myAlpha[t * maxGLen + u] = logSumExp(next, current);\n }\n }\n }\n __syncthreads();\n }\n } else if (blockIdx.x == 1) {\n // beta path\n acc_t* myBeta = beta + batch * maxFLen * maxGLen;\n if (u == 0) {\n myBeta[(myFLen - 1) * maxGLen + myGLen - 1] = myX[((myFLen - 1) * myStrideT + myGLen - 1) * dictSize + blankIdx];\n }\n __syncthreads();\n\n for (int64_t step = myFLen + myGLen - 3; step >= 0; --step) {\n for (u = tid; u < myGLen; u += blockDim.x) {\n int64_t t = step - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(18) in [1]\n if (u == myGLen - 1) {\n // beta(t, u) = beta(t+1, u) * null(t, u)\n myBeta[t * maxGLen + u] = myBeta[(t + 1) * maxGLen + u] + myX[(t * myStrideT + u) * dictSize + blankIdx];\n } else if (t == myFLen - 1) {\n // beta(t, u) = beta(t, u+1) * y(t, u)\n myBeta[t * maxGLen + u] = myBeta[t * maxGLen + u + 1] + myX[(t * myStrideT + u) * dictSize + myLabel[u]];\n } else {\n // beta(t, u) = beta(t+1, u)*null(t, u) + beta(t, u+1)*y(t, u)\n acc_t current = myBeta[(t + 1) * maxGLen + u] + myX[(t * myStrideT + u) * dictSize + blankIdx];\n acc_t next = myBeta[t * maxGLen + u + 1] + myX[(t * myStrideT + u) * dictSize + myLabel[u]];\n myBeta[t * maxGLen + u] = logSumExp(next, current);\n }\n }\n }\n __syncthreads();\n }\n if (tid == 0) loss[batch] = -myBeta[0];\n }\n}\n\n// transudcer loss function (i.e. forward-backward algorithm) with batch loading optimization.\n// Compared to the vanilla version, there are two optimizations:\n// 1. load x in batch through loop unrolling to reduce the latency.\n// 2. Use registers and shared memory to hold alpha and beta values passed from one step the next.\n// For simplicity, this kernel currently only supports U <= maxThread, which should be the common\n// case. For cases where U > maxThread, the vanilla kernel is used as a fallback option.\n\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// Forward (alpha) and backward (beta) path are launched together. Input is assumed to be converted\n// into log scale by the preceding log_softmax layer\n// Diagonal wavefront advancing usually used in dynamic programming is leveraged here.\n// alpha and beta are of acc_t type, as they are essentially accumulators.\n\n// This loss function supports packed input where a tensor of shape [B, T, U, H] is packed into\n// [B_packed, H].\n// Don't-care region (t > audLen) or (u > txtLen) is removed.\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_batch_load_forward(const scalar_t* x, const int* label, const int* audLen,\n const int* txtLen, const int64_t* batchOffset, int64_t dictSize,\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n acc_t* alpha, acc_t* beta, scalar_t* loss) {\n const int batch = blockIdx.y;\n int u = threadIdx.x;\n const auto myFLen = audLen[batch];\n const auto myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n const scalar_t* myX = x + myBatchOffset * dictSize;\n scalar_t next[batchLdSize], current[batchLdSize];\n extern __shared__ char smem8[];\n auto smem = reinterpret_cast(smem8);\n\n if (blockIdx.x == 0) {\n // alpha path\n acc_t* myAlpha = alpha + batch * maxFLen * maxGLen;\n // two SMEM regions for double buffering read and write data to avoid data race\n acc_t* const sharedAlpha[2] = {smem, smem + maxGLen};\n\n sharedAlpha[0][u] = 0;\n __syncthreads();\n\n if (u == 0) myAlpha[0] = 0;\n\n auto myAlphaLabel = (u == 0) ? 0 : label[batch * (maxGLen - 1) + u - 1];\n // register used to pass value to the next step for the same thread\n acc_t prvStepAlpha = 0;\n for (int64_t step = 1; step < myFLen + myGLen - 1 + batchLdSize; step += batchLdSize) {\n// Move along the diagonal wavefront to leverage available parallelism\n// Batch loading X through loop unrolling\n#pragma unroll\n for (int i = 0; i < batchLdSize; ++i) {\n if (step + i < myFLen + myGLen - 1) {\n // index computing\n int64_t t = step + i - u;\n int64_t currentId = ((t - 1) * myStrideT + u) * dictSize + blankIdx;\n int64_t nextId = (t * myStrideT + u - 1) * dictSize + myAlphaLabel;\n // main loading loop\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n if (u == 0) {\n current[i] = myX[currentId];\n } else if (t == 0) {\n next[i] = myX[nextId];\n } else {\n current[i] = myX[currentId];\n next[i] = myX[nextId];\n }\n }\n }\n }\n // main computing loop\n for (int i = 0; i < batchLdSize; ++i) {\n // swap the pointer for double buffering\n auto sharedAlphaRd = sharedAlpha[(step + i - 1) % 2];\n auto sharedAlphaWr = sharedAlpha[(step + i) % 2];\n if (step + i < myFLen + myGLen - 1) {\n int64_t t = step + i - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(16) in [1]\n if (u == 0)\n prvStepAlpha = prvStepAlpha + current[i];\n else if (t == 0)\n prvStepAlpha = sharedAlphaRd[u - 1] + next[i];\n else\n prvStepAlpha = logSumExp(prvStepAlpha + current[i], sharedAlphaRd[u - 1] + next[i]);\n sharedAlphaWr[u] = prvStepAlpha;\n myAlpha[t * maxGLen + u] = prvStepAlpha;\n }\n }\n __syncthreads();\n }\n }\n } else if (blockIdx.x == 1) {\n // beta path\n acc_t* myBeta = beta + batch * maxFLen * maxGLen;\n // two SMEM regions for double buffering read and write data to avoid data race\n acc_t* const sharedBeta[2] = {smem, smem + maxGLen};\n sharedBeta[0][u] = myX[((myFLen - 1) * myStrideT + myGLen - 1) * dictSize + blankIdx];\n __syncthreads();\n\n auto myBetaLabel = (u == maxGLen - 1) ? 0 : label[batch * (maxGLen - 1) + u];\n // register used to pass value to the next step for the same thread\n acc_t prvStepBeta = myX[((myFLen - 1) * myStrideT + myGLen - 1) * dictSize + blankIdx];\n if (u == 0) myBeta[(myFLen - 1) * maxGLen + myGLen - 1] = prvStepBeta;\n\n for (int64_t step = 1; step < myFLen + myGLen - 1; step += batchLdSize) {\n// Move along the diagonal wavefront to leverage available parallelism\n// Batch loading X\n#pragma unroll\n for (int i = 0; i < batchLdSize; ++i) {\n if (step + i < myFLen + myGLen - 1) {\n // index computing\n int64_t t = myFLen + myGLen - (step + i) - 2 - u;\n int64_t currentId = (t * myStrideT + u) * dictSize + blankIdx;\n int64_t nextId = (t * myStrideT + u) * dictSize + myBetaLabel;\n // main loading loop\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n if (u == myGLen - 1) {\n current[i] = myX[currentId];\n } else if (t == myFLen - 1) {\n next[i] = myX[nextId];\n } else {\n current[i] = myX[currentId];\n next[i] = myX[nextId];\n }\n }\n }\n }\n // main computing loop\n for (int i = 0; i < batchLdSize; ++i) {\n // swap the pointer for double buffering\n auto sharedBetaRd = sharedBeta[(step + i - 1) % 2];\n auto sharedBetaWr = sharedBeta[(step + i) % 2];\n if (step + i < myFLen + myGLen - 1) {\n int64_t t = myFLen + myGLen - (step + i) - 2 - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(18) in [1]\n if (u == myGLen - 1)\n prvStepBeta = prvStepBeta + current[i];\n else if (t == myFLen - 1)\n prvStepBeta = sharedBetaRd[u + 1] + next[i];\n else\n prvStepBeta = logSumExp(prvStepBeta + current[i], sharedBetaRd[u + 1] + next[i]);\n sharedBetaWr[u] = prvStepBeta;\n myBeta[t * maxGLen + u] = prvStepBeta;\n }\n }\n __syncthreads();\n }\n }\n if (u == 0) loss[batch] = -prvStepBeta;\n }\n}\n\n// Vanilla transudcer loss backward operation.\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// For this backward kernel, bwd op for the preceding softmax is assumed to be handled elsewhere,\n// hence only Eq(20) in [1] is implemented in this kernel.\n\n// Each thread block works on [batch, t, :, :] of data. Each thread works on a specific u at a time\n// Since only gradients for the correct token and null token need to be updated, gradients at other\n// locations are initialized to 0.\n\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_backward(const scalar_t* x, const scalar_t* lossGrad, const int* audLen,\n const int* txtLen, const int* label, const acc_t* alpha, const acc_t* beta,\n const int64_t* batchOffset, int64_t dictSize, int64_t blankIdx,\n int64_t maxFLen, int64_t maxGLen, bool packedInput, scalar_t* xGrad) {\n const int tid = threadIdx.x;\n const int t = blockIdx.x;\n const int batch = blockIdx.y;\n const int64_t myFLen = audLen[batch];\n const int64_t myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n auto myX = x + (myBatchOffset + t * myStrideT) * dictSize;\n auto myAlpha = alpha + batch * maxFLen * maxGLen;\n auto myBeta = beta + batch * maxFLen * maxGLen;\n auto myXGrad = xGrad + (myBatchOffset + t * myStrideT) * dictSize;\n auto myLabel = label + batch * (maxGLen - 1);\n\n int64_t u = tid;\n while (t < myFLen and u < myGLen) {\n // Do the update\n // loss = -ln(Pr(y*|x))\n acc_t grad = std::log(lossGrad[batch]) + myAlpha[t * maxGLen + u] - myBeta[0];\n if (u != myGLen - 1)\n myXGrad[u * dictSize + myLabel[u]] =\n -std::exp(grad + myBeta[t * maxGLen + u + 1] + myX[u * dictSize + myLabel[u]]);\n if (t == myFLen - 1 and u == myGLen - 1)\n myXGrad[u * dictSize + blankIdx] = -std::exp(grad + myX[u * dictSize + blankIdx]);\n else if (t != myFLen - 1)\n myXGrad[u * dictSize + blankIdx] = -std::exp(grad + myBeta[(t + 1) * maxGLen + u] + myX[u * dictSize + blankIdx]);\n\n u += blockDim.x;\n }\n}\n\n// Fused transudcer loss backward operation.\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// The bwd op of the preceding softmax layer is fused in this kernel.\n// Each thread block works on [batch, t, u, :] of data. Each thread works on a specific h at a time\n\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_fused_backward(const scalar_t* x, const scalar_t* lossGrad, const int* audLen,\n const int* txtLen, const int* label, const acc_t* alpha,\n const acc_t* beta, const int64_t* batchOffset, int64_t dictSize,\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n scalar_t* xGrad) {\n const int tid = threadIdx.x;\n const int u = blockIdx.x;\n const int t = blockIdx.y;\n const int batch = blockIdx.z;\n const int64_t myFLen = audLen[batch];\n const int64_t myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n\n __shared__ acc_t commonFactor, myBetaTU, myBetaTUp1, myBetaTp1U, myLabelShared;\n auto myXGrad = xGrad + (myBatchOffset + t * myStrideT + u) * dictSize;\n\n if (t < myFLen and u < myGLen) {\n auto myX = x + (myBatchOffset + t * myStrideT + u) * dictSize;\n auto myAlpha = alpha + batch * maxFLen * maxGLen;\n auto myBeta = beta + batch * maxFLen * maxGLen;\n auto myLabel = label + batch * (maxGLen - 1);\n\n // load and store shared variables in SMEM\n if (tid == 0) {\n commonFactor = std::log(lossGrad[batch]) + myAlpha[t * maxGLen + u] - myBeta[0];\n myBetaTU = myBeta[t * maxGLen + u];\n myBetaTUp1 = myBeta[t * maxGLen + u + 1];\n myBetaTp1U = myBeta[(t + 1) * maxGLen + u];\n myLabelShared = myLabel[u];\n }\n\n __syncthreads();\n\n for (int64_t h = tid; h < dictSize; h += blockDim.x) {\n // Do the update\n acc_t grad = commonFactor + myX[h]; // loss = -ln(Pr(y*|x))\n acc_t myGrad = std::exp(grad + myBetaTU);\n if (u != myGLen - 1 and h == myLabelShared) {\n myGrad -= std::exp(grad + myBetaTUp1);\n } else if (h == blankIdx) {\n if (t == myFLen - 1 and u == myGLen - 1)\n myGrad -= std::exp(grad);\n else if (t != myFLen - 1)\n myGrad -= std::exp(grad + myBetaTp1U);\n }\n myXGrad[h] = myGrad;\n }\n } else if (!packedInput) {\n // In non-pack mode, need to make sure the gradients for don't-care regions are zero.\n for (int64_t h = tid; h < dictSize; h += blockDim.x) {\n myXGrad[h] = 0;\n }\n }\n}\n\n// Vectorized version of fused transudcer loss backward operation.\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// The bwd op of the preceding softmax layer is fused in this kernel.\n// Each thread block works on [batch, t, u, :] of data. Each thread works on a specific h at a time\n\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_fused_vec_backward(const scalar_t* x, const scalar_t* lossGrad, const int* audLen,\n const int* txtLen, const int* label, const acc_t* alpha,\n const acc_t* beta, const int64_t* batchOffset, int64_t dictSize,\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n scalar_t* xGrad) {\n const int tid = threadIdx.x;\n const int u = blockIdx.x;\n const int t = blockIdx.y;\n const int batch = blockIdx.z;\n const int64_t myFLen = audLen[batch];\n const int64_t myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n\n __shared__ acc_t commonFactor, myBetaTU, myBetaTUp1, myBetaTp1U, myLabelShared;\n auto myXGrad = xGrad + (myBatchOffset + t * myStrideT + u) * dictSize;\n auto myX = x + (myBatchOffset + t * myStrideT + u) * dictSize;\n auto myAlpha = alpha + batch * maxFLen * maxGLen;\n auto myBeta = beta + batch * maxFLen * maxGLen;\n auto myLabel = label + batch * (maxGLen - 1);\n\n // Variabels for vectorization\n scalar_t myXBuffer[V], myXGradBuffer[V];\n auto myXVec = reinterpret_cast(myX);\n auto myXGradVec = reinterpret_cast(myXGrad);\n auto myXBufferVec = reinterpret_cast(myXBuffer);\n auto myXGradBufferVec = reinterpret_cast(myXGradBuffer);\n if (t < myFLen and u < myGLen) {\n // load and store shared variables in SMEM\n if (tid == 0) {\n commonFactor = std::log(lossGrad[batch]) + myAlpha[t * maxGLen + u] - myBeta[0];\n myBetaTU = myBeta[t * maxGLen + u];\n if (t != myFLen - 1) myBetaTp1U = myBeta[(t + 1) * maxGLen + u];\n if (u != myGLen - 1) {\n myBetaTUp1 = myBeta[t * maxGLen + u + 1];\n myLabelShared = myLabel[u];\n }\n }\n\n __syncthreads();\n\n#pragma unroll\n for (int64_t h0 = tid * V; h0 < dictSize; h0 += blockDim.x * V) {\n // Load myX in a vector form\n *myXBufferVec = myXVec[h0 / V];\n// Do the update for a vector of input\n#pragma unroll\n for (int i = 0; i < V; ++i) {\n auto h = h0 + i;\n acc_t grad = commonFactor + myXBuffer[i]; // loss = -ln(Pr(y*|x))\n acc_t myGrad = std::exp(grad + myBetaTU);\n if (u != myGLen - 1 and h == myLabelShared) {\n myGrad -= std::exp(grad + myBetaTUp1);\n } else if (h == blankIdx) {\n if (t == myFLen - 1 and u == myGLen - 1)\n myGrad -= std::exp(grad);\n else if (t != myFLen - 1)\n myGrad -= std::exp(grad + myBetaTp1U);\n }\n myXGradBuffer[i] = myGrad;\n }\n\n // Store myXGrad in a vector form\n myXGradVec[h0 / V] = *myXGradBufferVec;\n }\n } else if (!packedInput) {\n // In non-pack mode, need to make sure the gradients for don't-care regions are zero.\n for (int64_t h0 = tid * V; h0 < dictSize; h0 += blockDim.x * V) {\n myXGradVec[h0 / V] = 0;\n }\n }\n}\n\nstd::vector transducer_loss_cuda_forward(torch::Tensor x, torch::Tensor label, torch::Tensor audLen,\n torch::Tensor txtLen, torch::Tensor batchOffset, int maxFLen,\n int blankIdx, int opt, bool packedInput) {\n auto scalarType = x.scalar_type();\n auto tensorOpt = x.options();\n const int batchSize = label.size(0);\n const int maxGLen = label.size(1) + 1;\n const int dictSize = x.size(-1);\n\n TORCH_CHECK(blankIdx >= 0 and blankIdx < dictSize, \"Expected blank index to be in the range of 0 to \", dictSize - 1,\n \", but got \", blankIdx);\n TORCH_CHECK(opt == -1 or opt == 0 or opt == 1, \"Got an invalid optimization level \", opt);\n\n // The data type of alpha and beta will be resolved at dispatch time,\n // hence defined here and assigned later\n torch::Tensor alpha;\n torch::Tensor beta;\n torch::Tensor loss = torch::empty({batchSize}, tensorOpt);\n const auto deviceProperties = at::cuda::getCurrentDeviceProperties();\n const auto maxThreadPerBlock = deviceProperties->maxThreadsPerBlock;\n const auto maxSmemPerBlock = deviceProperties->sharedMemPerBlock;\n const auto batchOffsetPtr = packedInput ? batchOffset.data_ptr() : nullptr;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n scalarType, \"transducer_loss_cuda_forward\", ([&] {\n // resolve accumulation type\n using acc_t = at::acc_type;\n auto accType = c10::CppTypeToScalarType::value;\n auto accTensorOpt = tensorOpt.dtype(accType);\n alpha = torch::empty({batchSize, maxFLen, maxGLen}, accTensorOpt);\n beta = torch::empty({batchSize, maxFLen, maxGLen}, accTensorOpt);\n\n // decide what kernel to launch based on the problem size\n // if the required SMEM size or number threads exceeds the limit, fall back to the vanilla\n // kernel.\n const auto smemSize = 2 * maxGLen * sizeof(acc_t);\n const auto optFallBack = (maxGLen > maxThreadPerBlock or smemSize > maxSmemPerBlock) ? 0\n : (opt == -1) ? 1\n : opt;\n const int threads = std::min(maxThreadPerBlock, maxGLen);\n const dim3 blocks(2, batchSize, 1);\n\n if (optFallBack == 0)\n transducer_loss_forward<<>>(\n x.data_ptr(), label.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n batchOffsetPtr, dictSize, blankIdx, maxFLen, maxGLen, packedInput, alpha.data_ptr(),\n beta.data_ptr(), loss.data_ptr());\n else if (optFallBack == 1)\n transducer_loss_batch_load_forward<<>>(\n x.data_ptr(), label.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n batchOffsetPtr, dictSize, blankIdx, maxFLen, maxGLen, packedInput, alpha.data_ptr(),\n beta.data_ptr(), loss.data_ptr());\n }));\n C10_CUDA_CHECK(cudaGetLastError());\n\n return {alpha, beta, loss};\n}\n\ntorch::Tensor transducer_loss_cuda_backward(torch::Tensor x, torch::Tensor lossGrad, torch::Tensor alpha,\n torch::Tensor beta, torch::Tensor audLen, torch::Tensor txtLen,\n torch::Tensor label, torch::Tensor batchOffset, int maxFLen, int blankIdx,\n int opt, bool fuseSoftmaxBackward, bool packedInput) {\n auto dtype = x.scalar_type();\n torch::Tensor xGrad;\n const int batchSize = label.size(0);\n const int maxGLen = label.size(1) + 1;\n const int dictSize = x.size(-1);\n const auto deviceProperties = at::cuda::getCurrentDeviceProperties();\n const int maxThreadPerBlock = deviceProperties->maxThreadsPerBlock;\n const int warpSize = deviceProperties->warpSize;\n const auto batchOffsetPtr = packedInput ? batchOffset.data_ptr() : nullptr;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fuseSoftmaxBackward) {\n // alloc empty tensors for performance, hence need to ensure zeros are writtern to\n // don't-care region in the kernel.\n xGrad = torch::empty_like(x);\n\n // Would like each thread to work on 4 hidden units\n const int workPerThread = 4;\n // Don't want to have more than 128 threads per thread block\n const int maxThreadPerElmt = std::min(128, maxThreadPerBlock);\n const int threads = std::min(maxThreadPerElmt, std::max(warpSize, (dictSize + workPerThread - 1) / workPerThread));\n const dim3 blocks(maxGLen, maxFLen, batchSize);\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n dtype, \"transducer_loss_cuda_backward\", ([&] {\n using vec_t = uint64_t;\n using acc_t = at::acc_type;\n constexpr int vectFactor = sizeof(vec_t) / sizeof(scalar_t);\n constexpr int vecAlignment = std::alignment_of::value;\n // if all input and output tensors meet the alignment requirement\n bool memAlign = reinterpret_cast(x.data_ptr()) % vecAlignment == 0 and\n reinterpret_cast(xGrad.data_ptr()) % vecAlignment == 0;\n\n if (vectFactor > 1 and dictSize % vectFactor == 0 and memAlign) {\n transducer_loss_fused_vec_backward<<>>(\n x.data_ptr(), lossGrad.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n label.data_ptr(), alpha.data_ptr(), beta.data_ptr(), batchOffsetPtr, dictSize,\n blankIdx, maxFLen, maxGLen, packedInput, xGrad.data_ptr());\n } else {\n transducer_loss_fused_backward<<>>(\n x.data_ptr(), lossGrad.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n label.data_ptr(), alpha.data_ptr(), beta.data_ptr(), batchOffsetPtr, dictSize,\n blankIdx, maxFLen, maxGLen, packedInput, xGrad.data_ptr());\n }\n }));\n } else {\n // for non-fused kernel, the gradients need to be writtern are very sparse, hence initialize\n // the tensor with all zeros.\n xGrad = torch::zeros_like(x);\n // don't launch more threads than needed.\n const int threads = std::min(maxThreadPerBlock, maxGLen);\n const dim3 blocks(maxFLen, batchSize);\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(dtype, \"transducer_loss_cuda_backward\", ([&] {\n using acc_t = at::acc_type;\n transducer_loss_backward<<>>(\n x.data_ptr(), lossGrad.data_ptr(),\n audLen.data_ptr(), txtLen.data_ptr(), label.data_ptr(),\n alpha.data_ptr(), beta.data_ptr(), batchOffsetPtr, dictSize,\n blankIdx, maxFLen, maxGLen, packedInput, xGrad.data_ptr());\n }));\n }\n C10_CUDA_CHECK(cudaGetLastError());\n\n return xGrad;\n}\n", "output": "#include \n#include \n#include \n#include \n#include \n\n#include \n\ntemplate \n__device__ __forceinline__ scalar_t logSumExp(scalar_t a, scalar_t b) {\n // standard log-sum-exp trick is used here to provide better numerical stability\n return (a >= b) ? a + std::log1p(exp(b - a)) : b + std::log1p(exp(a - b));\n}\n\n// Vanilla transducer loss function (i.e. forward-backward algorithm)\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n\n// Forward (alpha) and backward (beta) path are launched together. Input is assumed to be converted\n// into log scale by the preceding log_softmax layer\n// Diagonal wavefront advancing usually used in dynamic programming is leveraged here.\n// alpha and beta are of acc_t type, as they are essentially accumulators.\n\n// This loss function supports packed input where a tensor of shape [B, T, U, H] is packed into\n// [B_packed, H].\n// Don't-care region (t > audLen) or (u > txtLen) is removed.\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_forward(const scalar_t* x, const int* label, const int* audLen, const int* txtLen,\n const int64_t* batchOffset,\n int64_t dictSize, // 64-bit indexing for data tensor\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n acc_t* alpha, acc_t* beta, scalar_t* loss) {\n const int batch = blockIdx.y;\n const int tid = threadIdx.x;\n const auto myFLen = audLen[batch];\n // Note that start of the sentence is added as 1 here\n const auto myGLen = txtLen[batch] + 1;\n const auto myLabel = label + batch * (maxGLen - 1);\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n const scalar_t* myX = x + myBatchOffset * dictSize;\n int u = tid;\n\n if (blockIdx.x == 0) {\n // alpha path\n acc_t* myAlpha = alpha + batch * maxFLen * maxGLen;\n if (u == 0) myAlpha[0] = 0;\n __syncthreads();\n\n for (int64_t step = 1; step < myFLen + myGLen - 1; ++step) {\n // Move along the diagonal wavefront to leverage available parallelism\n for (u = tid; u < myGLen; u += blockDim.x) {\n int64_t t = step - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(16) in [1]\n if (u == 0) {\n // alpha(t, u) = alpha(t-1, u) * null(t-1, u)\n myAlpha[t * maxGLen + u] = myAlpha[(t - 1) * maxGLen] + myX[((t - 1) * myStrideT) * dictSize + blankIdx];\n } else if (t == 0) {\n // alpha(t, u-1) = alpha(t, u-1) * y(t, u-1)\n myAlpha[u] = myAlpha[u - 1] + myX[(u - 1) * dictSize + myLabel[u - 1]];\n } else {\n // alpha(t, u) = alpha(t-1, u) * null(t-1, u) + alpha(t, u-1) * y(t, u-1)\n acc_t current = myAlpha[(t - 1) * maxGLen + u] + myX[((t - 1) * myStrideT + u) * dictSize + blankIdx];\n acc_t next = myAlpha[t * maxGLen + u - 1] + myX[(t * myStrideT + u - 1) * dictSize + myLabel[u - 1]];\n myAlpha[t * maxGLen + u] = logSumExp(next, current);\n }\n }\n }\n __syncthreads();\n }\n } else if (blockIdx.x == 1) {\n // beta path\n acc_t* myBeta = beta + batch * maxFLen * maxGLen;\n if (u == 0) {\n myBeta[(myFLen - 1) * maxGLen + myGLen - 1] = myX[((myFLen - 1) * myStrideT + myGLen - 1) * dictSize + blankIdx];\n }\n __syncthreads();\n\n for (int64_t step = myFLen + myGLen - 3; step >= 0; --step) {\n for (u = tid; u < myGLen; u += blockDim.x) {\n int64_t t = step - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(18) in [1]\n if (u == myGLen - 1) {\n // beta(t, u) = beta(t+1, u) * null(t, u)\n myBeta[t * maxGLen + u] = myBeta[(t + 1) * maxGLen + u] + myX[(t * myStrideT + u) * dictSize + blankIdx];\n } else if (t == myFLen - 1) {\n // beta(t, u) = beta(t, u+1) * y(t, u)\n myBeta[t * maxGLen + u] = myBeta[t * maxGLen + u + 1] + myX[(t * myStrideT + u) * dictSize + myLabel[u]];\n } else {\n // beta(t, u) = beta(t+1, u)*null(t, u) + beta(t, u+1)*y(t, u)\n acc_t current = myBeta[(t + 1) * maxGLen + u] + myX[(t * myStrideT + u) * dictSize + blankIdx];\n acc_t next = myBeta[t * maxGLen + u + 1] + myX[(t * myStrideT + u) * dictSize + myLabel[u]];\n myBeta[t * maxGLen + u] = logSumExp(next, current);\n }\n }\n }\n __syncthreads();\n }\n if (tid == 0) loss[batch] = -myBeta[0];\n }\n}\n\n// transudcer loss function (i.e. forward-backward algorithm) with batch loading optimization.\n// Compared to the vanilla version, there are two optimizations:\n// 1. load x in batch through loop unrolling to reduce the latency.\n// 2. Use registers and shared memory to hold alpha and beta values passed from one step the next.\n// For simplicity, this kernel currently only supports U <= maxThread, which should be the common\n// case. For cases where U > maxThread, the vanilla kernel is used as a fallback option.\n\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// Forward (alpha) and backward (beta) path are launched together. Input is assumed to be converted\n// into log scale by the preceding log_softmax layer\n// Diagonal wavefront advancing usually used in dynamic programming is leveraged here.\n// alpha and beta are of acc_t type, as they are essentially accumulators.\n\n// This loss function supports packed input where a tensor of shape [B, T, U, H] is packed into\n// [B_packed, H].\n// Don't-care region (t > audLen) or (u > txtLen) is removed.\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_batch_load_forward(const scalar_t* x, const int* label, const int* audLen,\n const int* txtLen, const int64_t* batchOffset, int64_t dictSize,\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n acc_t* alpha, acc_t* beta, scalar_t* loss) {\n const int batch = blockIdx.y;\n int u = threadIdx.x;\n const auto myFLen = audLen[batch];\n const auto myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n const scalar_t* myX = x + myBatchOffset * dictSize;\n scalar_t next[batchLdSize], current[batchLdSize];\n extern __shared__ char smem8[];\n auto smem = reinterpret_cast(smem8);\n\n if (blockIdx.x == 0) {\n // alpha path\n acc_t* myAlpha = alpha + batch * maxFLen * maxGLen;\n // two SMEM regions for double buffering read and write data to avoid data race\n acc_t* const sharedAlpha[2] = {smem, smem + maxGLen};\n\n sharedAlpha[0][u] = 0;\n __syncthreads();\n\n if (u == 0) myAlpha[0] = 0;\n\n auto myAlphaLabel = (u == 0) ? 0 : label[batch * (maxGLen - 1) + u - 1];\n // register used to pass value to the next step for the same thread\n acc_t prvStepAlpha = 0;\n for (int64_t step = 1; step < myFLen + myGLen - 1 + batchLdSize; step += batchLdSize) {\n// Move along the diagonal wavefront to leverage available parallelism\n// Batch loading X through loop unrolling\n#pragma unroll\n for (int i = 0; i < batchLdSize; ++i) {\n if (step + i < myFLen + myGLen - 1) {\n // index computing\n int64_t t = step + i - u;\n int64_t currentId = ((t - 1) * myStrideT + u) * dictSize + blankIdx;\n int64_t nextId = (t * myStrideT + u - 1) * dictSize + myAlphaLabel;\n // main loading loop\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n if (u == 0) {\n current[i] = myX[currentId];\n } else if (t == 0) {\n next[i] = myX[nextId];\n } else {\n current[i] = myX[currentId];\n next[i] = myX[nextId];\n }\n }\n }\n }\n // main computing loop\n for (int i = 0; i < batchLdSize; ++i) {\n // swap the pointer for double buffering\n auto sharedAlphaRd = sharedAlpha[(step + i - 1) % 2];\n auto sharedAlphaWr = sharedAlpha[(step + i) % 2];\n if (step + i < myFLen + myGLen - 1) {\n int64_t t = step + i - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(16) in [1]\n if (u == 0)\n prvStepAlpha = prvStepAlpha + current[i];\n else if (t == 0)\n prvStepAlpha = sharedAlphaRd[u - 1] + next[i];\n else\n prvStepAlpha = logSumExp(prvStepAlpha + current[i], sharedAlphaRd[u - 1] + next[i]);\n sharedAlphaWr[u] = prvStepAlpha;\n myAlpha[t * maxGLen + u] = prvStepAlpha;\n }\n }\n __syncthreads();\n }\n }\n } else if (blockIdx.x == 1) {\n // beta path\n acc_t* myBeta = beta + batch * maxFLen * maxGLen;\n // two SMEM regions for double buffering read and write data to avoid data race\n acc_t* const sharedBeta[2] = {smem, smem + maxGLen};\n sharedBeta[0][u] = myX[((myFLen - 1) * myStrideT + myGLen - 1) * dictSize + blankIdx];\n __syncthreads();\n\n auto myBetaLabel = (u == maxGLen - 1) ? 0 : label[batch * (maxGLen - 1) + u];\n // register used to pass value to the next step for the same thread\n acc_t prvStepBeta = myX[((myFLen - 1) * myStrideT + myGLen - 1) * dictSize + blankIdx];\n if (u == 0) myBeta[(myFLen - 1) * maxGLen + myGLen - 1] = prvStepBeta;\n\n for (int64_t step = 1; step < myFLen + myGLen - 1; step += batchLdSize) {\n// Move along the diagonal wavefront to leverage available parallelism\n// Batch loading X\n#pragma unroll\n for (int i = 0; i < batchLdSize; ++i) {\n if (step + i < myFLen + myGLen - 1) {\n // index computing\n int64_t t = myFLen + myGLen - (step + i) - 2 - u;\n int64_t currentId = (t * myStrideT + u) * dictSize + blankIdx;\n int64_t nextId = (t * myStrideT + u) * dictSize + myBetaLabel;\n // main loading loop\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n if (u == myGLen - 1) {\n current[i] = myX[currentId];\n } else if (t == myFLen - 1) {\n next[i] = myX[nextId];\n } else {\n current[i] = myX[currentId];\n next[i] = myX[nextId];\n }\n }\n }\n }\n // main computing loop\n for (int i = 0; i < batchLdSize; ++i) {\n // swap the pointer for double buffering\n auto sharedBetaRd = sharedBeta[(step + i - 1) % 2];\n auto sharedBetaWr = sharedBeta[(step + i) % 2];\n if (step + i < myFLen + myGLen - 1) {\n int64_t t = myFLen + myGLen - (step + i) - 2 - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(18) in [1]\n if (u == myGLen - 1)\n prvStepBeta = prvStepBeta + current[i];\n else if (t == myFLen - 1)\n prvStepBeta = sharedBetaRd[u + 1] + next[i];\n else\n prvStepBeta = logSumExp(prvStepBeta + current[i], sharedBetaRd[u + 1] + next[i]);\n sharedBetaWr[u] = prvStepBeta;\n myBeta[t * maxGLen + u] = prvStepBeta;\n }\n }\n __syncthreads();\n }\n }\n if (u == 0) loss[batch] = -prvStepBeta;\n }\n}\n\n// Vanilla transudcer loss backward operation.\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// For this backward kernel, bwd op for the preceding softmax is assumed to be handled elsewhere,\n// hence only Eq(20) in [1] is implemented in this kernel.\n\n// Each thread block works on [batch, t, :, :] of data. Each thread works on a specific u at a time\n// Since only gradients for the correct token and null token need to be updated, gradients at other\n// locations are initialized to 0.\n\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_backward(const scalar_t* x, const scalar_t* lossGrad, const int* audLen,\n const int* txtLen, const int* label, const acc_t* alpha, const acc_t* beta,\n const int64_t* batchOffset, int64_t dictSize, int64_t blankIdx,\n int64_t maxFLen, int64_t maxGLen, bool packedInput, scalar_t* xGrad) {\n const int tid = threadIdx.x;\n const int t = blockIdx.x;\n const int batch = blockIdx.y;\n const int64_t myFLen = audLen[batch];\n const int64_t myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n auto myX = x + (myBatchOffset + t * myStrideT) * dictSize;\n auto myAlpha = alpha + batch * maxFLen * maxGLen;\n auto myBeta = beta + batch * maxFLen * maxGLen;\n auto myXGrad = xGrad + (myBatchOffset + t * myStrideT) * dictSize;\n auto myLabel = label + batch * (maxGLen - 1);\n\n int64_t u = tid;\n while (t < myFLen and u < myGLen) {\n // Do the update\n // loss = -ln(Pr(y*|x))\n acc_t grad = std::log(lossGrad[batch]) + myAlpha[t * maxGLen + u] - myBeta[0];\n if (u != myGLen - 1)\n myXGrad[u * dictSize + myLabel[u]] =\n -std::exp(grad + myBeta[t * maxGLen + u + 1] + myX[u * dictSize + myLabel[u]]);\n if (t == myFLen - 1 and u == myGLen - 1)\n myXGrad[u * dictSize + blankIdx] = -std::exp(grad + myX[u * dictSize + blankIdx]);\n else if (t != myFLen - 1)\n myXGrad[u * dictSize + blankIdx] = -std::exp(grad + myBeta[(t + 1) * maxGLen + u] + myX[u * dictSize + blankIdx]);\n\n u += blockDim.x;\n }\n}\n\n// Fused transudcer loss backward operation.\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// The bwd op of the preceding softmax layer is fused in this kernel.\n// Each thread block works on [batch, t, u, :] of data. Each thread works on a specific h at a time\n\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_fused_backward(const scalar_t* x, const scalar_t* lossGrad, const int* audLen,\n const int* txtLen, const int* label, const acc_t* alpha,\n const acc_t* beta, const int64_t* batchOffset, int64_t dictSize,\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n scalar_t* xGrad) {\n const int tid = threadIdx.x;\n const int u = blockIdx.x;\n const int t = blockIdx.y;\n const int batch = blockIdx.z;\n const int64_t myFLen = audLen[batch];\n const int64_t myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n\n __shared__ acc_t commonFactor, myBetaTU, myBetaTUp1, myBetaTp1U, myLabelShared;\n auto myXGrad = xGrad + (myBatchOffset + t * myStrideT + u) * dictSize;\n\n if (t < myFLen and u < myGLen) {\n auto myX = x + (myBatchOffset + t * myStrideT + u) * dictSize;\n auto myAlpha = alpha + batch * maxFLen * maxGLen;\n auto myBeta = beta + batch * maxFLen * maxGLen;\n auto myLabel = label + batch * (maxGLen - 1);\n\n // load and store shared variables in SMEM\n if (tid == 0) {\n commonFactor = std::log(lossGrad[batch]) + myAlpha[t * maxGLen + u] - myBeta[0];\n myBetaTU = myBeta[t * maxGLen + u];\n myBetaTUp1 = myBeta[t * maxGLen + u + 1];\n myBetaTp1U = myBeta[(t + 1) * maxGLen + u];\n myLabelShared = myLabel[u];\n }\n\n __syncthreads();\n\n for (int64_t h = tid; h < dictSize; h += blockDim.x) {\n // Do the update\n acc_t grad = commonFactor + myX[h]; // loss = -ln(Pr(y*|x))\n acc_t myGrad = std::exp(grad + myBetaTU);\n if (u != myGLen - 1 and h == myLabelShared) {\n myGrad -= std::exp(grad + myBetaTUp1);\n } else if (h == blankIdx) {\n if (t == myFLen - 1 and u == myGLen - 1)\n myGrad -= std::exp(grad);\n else if (t != myFLen - 1)\n myGrad -= std::exp(grad + myBetaTp1U);\n }\n myXGrad[h] = myGrad;\n }\n } else if (!packedInput) {\n // In non-pack mode, need to make sure the gradients for don't-care regions are zero.\n for (int64_t h = tid; h < dictSize; h += blockDim.x) {\n myXGrad[h] = 0;\n }\n }\n}\n\n// Vectorized version of fused transudcer loss backward operation.\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// The bwd op of the preceding softmax layer is fused in this kernel.\n// Each thread block works on [batch, t, u, :] of data. Each thread works on a specific h at a time\n\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_fused_vec_backward(const scalar_t* x, const scalar_t* lossGrad, const int* audLen,\n const int* txtLen, const int* label, const acc_t* alpha,\n const acc_t* beta, const int64_t* batchOffset, int64_t dictSize,\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n scalar_t* xGrad) {\n const int tid = threadIdx.x;\n const int u = blockIdx.x;\n const int t = blockIdx.y;\n const int batch = blockIdx.z;\n const int64_t myFLen = audLen[batch];\n const int64_t myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n\n __shared__ acc_t commonFactor, myBetaTU, myBetaTUp1, myBetaTp1U, myLabelShared;\n auto myXGrad = xGrad + (myBatchOffset + t * myStrideT + u) * dictSize;\n auto myX = x + (myBatchOffset + t * myStrideT + u) * dictSize;\n auto myAlpha = alpha + batch * maxFLen * maxGLen;\n auto myBeta = beta + batch * maxFLen * maxGLen;\n auto myLabel = label + batch * (maxGLen - 1);\n\n // Variabels for vectorization\n scalar_t myXBuffer[V], myXGradBuffer[V];\n auto myXVec = reinterpret_cast(myX);\n auto myXGradVec = reinterpret_cast(myXGrad);\n auto myXBufferVec = reinterpret_cast(myXBuffer);\n auto myXGradBufferVec = reinterpret_cast(myXGradBuffer);\n if (t < myFLen and u < myGLen) {\n // load and store shared variables in SMEM\n if (tid == 0) {\n commonFactor = std::log(lossGrad[batch]) + myAlpha[t * maxGLen + u] - myBeta[0];\n myBetaTU = myBeta[t * maxGLen + u];\n if (t != myFLen - 1) myBetaTp1U = myBeta[(t + 1) * maxGLen + u];\n if (u != myGLen - 1) {\n myBetaTUp1 = myBeta[t * maxGLen + u + 1];\n myLabelShared = myLabel[u];\n }\n }\n\n __syncthreads();\n\n#pragma unroll\n for (int64_t h0 = tid * V; h0 < dictSize; h0 += blockDim.x * V) {\n // Load myX in a vector form\n *myXBufferVec = myXVec[h0 / V];\n// Do the update for a vector of input\n#pragma unroll\n for (int i = 0; i < V; ++i) {\n auto h = h0 + i;\n acc_t grad = commonFactor + myXBuffer[i]; // loss = -ln(Pr(y*|x))\n acc_t myGrad = std::exp(grad + myBetaTU);\n if (u != myGLen - 1 and h == myLabelShared) {\n myGrad -= std::exp(grad + myBetaTUp1);\n } else if (h == blankIdx) {\n if (t == myFLen - 1 and u == myGLen - 1)\n myGrad -= std::exp(grad);\n else if (t != myFLen - 1)\n myGrad -= std::exp(grad + myBetaTp1U);\n }\n myXGradBuffer[i] = myGrad;\n }\n\n // Store myXGrad in a vector form\n myXGradVec[h0 / V] = *myXGradBufferVec;\n }\n } else if (!packedInput) {\n // In non-pack mode, need to make sure the gradients for don't-care regions are zero.\n for (int64_t h0 = tid * V; h0 < dictSize; h0 += blockDim.x * V) {\n myXGradVec[h0 / V] = 0;\n }\n }\n}\n\nstd::vector transducer_loss_cuda_forward(torch::Tensor x, torch::Tensor label, torch::Tensor audLen,\n torch::Tensor txtLen, torch::Tensor batchOffset, int maxFLen,\n int blankIdx, int opt, bool packedInput) {\n auto scalarType = x.scalar_type();\n auto tensorOpt = x.options();\n const int batchSize = label.size(0);\n const int maxGLen = label.size(1) + 1;\n const int dictSize = x.size(-1);\n\n TORCH_CHECK(blankIdx >= 0 and blankIdx < dictSize, \"Expected blank index to be in the range of 0 to \", dictSize - 1,\n \", but got \", blankIdx);\n TORCH_CHECK(opt == -1 or opt == 0 or opt == 1, \"Got an invalid optimization level \", opt);\n\n // The data type of alpha and beta will be resolved at dispatch time,\n // hence defined here and assigned later\n torch::Tensor alpha;\n torch::Tensor beta;\n torch::Tensor loss = torch::empty({batchSize}, tensorOpt);\n const auto deviceProperties = at::cuda::getCurrentDeviceProperties();\n const auto maxThreadPerBlock = deviceProperties->maxThreadsPerBlock;\n const auto maxSmemPerBlock = deviceProperties->sharedMemPerBlock;\n const auto batchOffsetPtr = packedInput ? batchOffset.data_ptr() : nullptr;\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n scalarType, \"transducer_loss_cuda_forward\", ([&] {\n // resolve accumulation type\n using acc_t = at::acc_type;\n auto accType = c10::CppTypeToScalarType::value;\n auto accTensorOpt = tensorOpt.dtype(accType);\n alpha = torch::empty({batchSize, maxFLen, maxGLen}, accTensorOpt);\n beta = torch::empty({batchSize, maxFLen, maxGLen}, accTensorOpt);\n\n // decide what kernel to launch based on the problem size\n // if the required SMEM size or number threads exceeds the limit, fall back to the vanilla\n // kernel.\n const auto smemSize = 2 * maxGLen * sizeof(acc_t);\n const auto optFallBack = (maxGLen > maxThreadPerBlock or smemSize > maxSmemPerBlock) ? 0\n : (opt == -1) ? 1\n : opt;\n const int threads = std::min(maxThreadPerBlock, maxGLen);\n const dim3 blocks(2, batchSize, 1);\n\n if (optFallBack == 0)\n transducer_loss_forward<<>>(\n x.data_ptr(), label.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n batchOffsetPtr, dictSize, blankIdx, maxFLen, maxGLen, packedInput, alpha.data_ptr(),\n beta.data_ptr(), loss.data_ptr());\n else if (optFallBack == 1)\n transducer_loss_batch_load_forward<<>>(\n x.data_ptr(), label.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n batchOffsetPtr, dictSize, blankIdx, maxFLen, maxGLen, packedInput, alpha.data_ptr(),\n beta.data_ptr(), loss.data_ptr());\n }));\n C10_CUDA_CHECK(hipGetLastError());\n\n return {alpha, beta, loss};\n}\n\ntorch::Tensor transducer_loss_cuda_backward(torch::Tensor x, torch::Tensor lossGrad, torch::Tensor alpha,\n torch::Tensor beta, torch::Tensor audLen, torch::Tensor txtLen,\n torch::Tensor label, torch::Tensor batchOffset, int maxFLen, int blankIdx,\n int opt, bool fuseSoftmaxBackward, bool packedInput) {\n auto dtype = x.scalar_type();\n torch::Tensor xGrad;\n const int batchSize = label.size(0);\n const int maxGLen = label.size(1) + 1;\n const int dictSize = x.size(-1);\n const auto deviceProperties = at::cuda::getCurrentDeviceProperties();\n const int maxThreadPerBlock = deviceProperties->maxThreadsPerBlock;\n const int warpSize = deviceProperties->warpSize;\n const auto batchOffsetPtr = packedInput ? batchOffset.data_ptr() : nullptr;\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fuseSoftmaxBackward) {\n // alloc empty tensors for performance, hence need to ensure zeros are writtern to\n // don't-care region in the kernel.\n xGrad = torch::empty_like(x);\n\n // Would like each thread to work on 4 hidden units\n const int workPerThread = 4;\n // Don't want to have more than 128 threads per thread block\n const int maxThreadPerElmt = std::min(128, maxThreadPerBlock);\n const int threads = std::min(maxThreadPerElmt, std::max(warpSize, (dictSize + workPerThread - 1) / workPerThread));\n const dim3 blocks(maxGLen, maxFLen, batchSize);\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n dtype, \"transducer_loss_cuda_backward\", ([&] {\n using vec_t = uint64_t;\n using acc_t = at::acc_type;\n constexpr int vectFactor = sizeof(vec_t) / sizeof(scalar_t);\n constexpr int vecAlignment = std::alignment_of::value;\n // if all input and output tensors meet the alignment requirement\n bool memAlign = reinterpret_cast(x.data_ptr()) % vecAlignment == 0 and\n reinterpret_cast(xGrad.data_ptr()) % vecAlignment == 0;\n\n if (vectFactor > 1 and dictSize % vectFactor == 0 and memAlign) {\n transducer_loss_fused_vec_backward<<>>(\n x.data_ptr(), lossGrad.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n label.data_ptr(), alpha.data_ptr(), beta.data_ptr(), batchOffsetPtr, dictSize,\n blankIdx, maxFLen, maxGLen, packedInput, xGrad.data_ptr());\n } else {\n transducer_loss_fused_backward<<>>(\n x.data_ptr(), lossGrad.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n label.data_ptr(), alpha.data_ptr(), beta.data_ptr(), batchOffsetPtr, dictSize,\n blankIdx, maxFLen, maxGLen, packedInput, xGrad.data_ptr());\n }\n }));\n } else {\n // for non-fused kernel, the gradients need to be writtern are very sparse, hence initialize\n // the tensor with all zeros.\n xGrad = torch::zeros_like(x);\n // don't launch more threads than needed.\n const int threads = std::min(maxThreadPerBlock, maxGLen);\n const dim3 blocks(maxFLen, batchSize);\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(dtype, \"transducer_loss_cuda_backward\", ([&] {\n using acc_t = at::acc_type;\n transducer_loss_backward<<>>(\n x.data_ptr(), lossGrad.data_ptr(),\n audLen.data_ptr(), txtLen.data_ptr(), label.data_ptr(),\n alpha.data_ptr(), beta.data_ptr(), batchOffsetPtr, dictSize,\n blankIdx, maxFLen, maxGLen, packedInput, xGrad.data_ptr());\n }));\n }\n C10_CUDA_CHECK(hipGetLastError());\n\n return xGrad;\n}\n", "source": "/kaggle/working/apex/apex/contrib/csrc/transducer/transducer_loss_kernel.cu"} {"input": "#include \n#include \n#include \n#include \n#include \n\n#ifdef OLD_GENERATOR_PATH\n#include \n#else\n#include \n#endif\n\n#include \n#include \n\n#include \n\n#include \"philox.cuh\"\n\n// Warp reduce kernels to reduce N groups of data into N numbers, where N = warpSize / width.\n// width should be a power of 2 and should be less than warpSize.\ntemplate \n__device__ __forceinline__ scalar_t warpReduce(scalar_t x, int width = C10_WARP_SIZE) {\n for (unsigned offset = width / 2; offset > 0; offset /= 2) {\n x += __shfl_down_sync(0xffffffff, x, offset, width);\n }\n return x;\n}\n\ninline int largestPowerOfTwo(int x) {\n int y = 1;\n while (y <= x) y <<= 1;\n return y >> 1;\n}\n\n/*\nFigure out vectorization type for masks.\nSimilar to how PyTorch figures out acc_t here:\naten/src/ATen/AccumulateType.h\n*/\ntemplate \nstruct MaskVecType {};\n\ntemplate <>\nstruct MaskVecType<1> {\n using type = uint8_t;\n};\ntemplate <>\nstruct MaskVecType<2> {\n using type = uint16_t;\n};\ntemplate <>\nstruct MaskVecType<4> {\n using type = uint32_t;\n};\n\ntemplate \nusing mvec_type = typename MaskVecType::type;\n\n// Helper class to calculate pointer offset that can be shared by different flavors of kernels.\n// For fwd, batch offset and stride are different for packing and non-packing mode.\nstruct OffsetCalFwd {\n __device__ __forceinline__ OffsetCalFwd(int64_t batch, const int64_t* batchOffset, int64_t maxFLen, int64_t maxGLen,\n int64_t gLen, int64_t hiddenSize, bool packOutput)\n : batch(batch),\n batchOffset(batchOffset),\n maxFLen(maxFLen),\n maxGLen(maxGLen),\n gLen(gLen),\n hiddenSize(hiddenSize),\n packOutput(packOutput) {}\n\n int64_t batch;\n const int64_t* batchOffset;\n int64_t maxFLen;\n int64_t maxGLen;\n int64_t gLen;\n int64_t hiddenSize;\n bool packOutput;\n\n __device__ __forceinline__ int64_t getBatchOffset() {\n return packOutput ? ((batch == 0) ? 0 : batchOffset[batch - 1]) * hiddenSize\n : batch * maxFLen * maxGLen * hiddenSize;\n }\n\n __device__ __forceinline__ int64_t getStrideF() { return packOutput ? gLen * hiddenSize : maxGLen * hiddenSize; }\n};\n\n// Helper class to calculate pointer offset that can be shared by different flavors of kernels\n// For bwd, batch offset and stride are different for packing and non-packing mode.\n// The reducion is done for two input tensors. Therefore, generating two sets of offsets\n// according to bwdFasterDim can lead to a unified implementation in the actual kernel.\nstruct OffsetCalBwd {\n __device__ __forceinline__ OffsetCalBwd(int64_t batch, const int64_t* batchOffset, const int* fLen, const int* gLen,\n int64_t maxFLen, int64_t maxGLen, int64_t hiddenSize, bool packOutput,\n bool bwdFasterDim)\n : batch(batch),\n batchOffset(batchOffset),\n maxFLen(maxFLen),\n maxGLen(maxGLen),\n fLen(fLen),\n gLen(gLen),\n hiddenSize(hiddenSize),\n packOutput(packOutput),\n bwdFasterDim(bwdFasterDim) {}\n\n int64_t batch;\n const int64_t* batchOffset;\n const int* fLen;\n const int* gLen;\n int64_t maxFLen;\n int64_t maxGLen;\n int64_t hiddenSize;\n bool packOutput;\n bool bwdFasterDim; // whether doing bwd on the faster moving dimension\n\n __device__ __forceinline__ int64_t getBatchOffset() {\n return packOutput ? ((batch == 0) ? 0 : batchOffset[batch - 1]) * hiddenSize\n : batch * maxFLen * maxGLen * hiddenSize;\n }\n\n __device__ __forceinline__ int64_t getMaxXLen() { return bwdFasterDim ? maxGLen : maxFLen; }\n\n __device__ __forceinline__ auto getMyXLen() -> decltype(gLen[batch]) {\n return bwdFasterDim ? gLen[batch] : fLen[batch];\n }\n\n __device__ __forceinline__ auto getMyYLen() -> decltype(gLen[batch]) {\n return bwdFasterDim ? fLen[batch] : gLen[batch];\n }\n\n __device__ __forceinline__ int64_t getStrideX() {\n return bwdFasterDim ? hiddenSize : ((packOutput ? gLen[batch] : maxGLen) * hiddenSize);\n }\n\n __device__ __forceinline__ int64_t getStrideY() {\n return bwdFasterDim ? ((packOutput ? gLen[batch] : maxGLen) * hiddenSize) : hiddenSize;\n }\n};\n\n// Vanila transducer joint forward kernel\n// Detail of this joint function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n\n// f is a tensor of shape [batch, T, H]\n// g is a tensor of shape [batch, U, H]\n// the transducer joint does\n// sum = f.unsqueeze(dim=2) + g.unsqueeze(dim=1)\n// The resultant tensor is of shape [batch, T, U, H]\n// Each thread block is working on one \"batch\" of data in the output tensor, [batch, t, u, :]\n\n// This joint function can optionally pack the output where the output tensor with a shape of\n// [B, T, U, H] is packed into [B_packed, H].\n// Don't-care region (t > fLen) or (u > gLen) is removed.\n// To enable packing, the starting offset for each batch need to be specified with batchOffset.\ntemplate \n__global__ void transducer_joint_forward(const scalar_t* f, const scalar_t* g, const int* fLen, const int* gLen,\n const int64_t* batchOffset, int64_t maxFLen, int64_t maxGLen,\n int64_t hiddenSize, bool packOutput, scalar_t* sum) {\n const int batch = blockIdx.z;\n const int t = blockIdx.y;\n const int u = blockIdx.x;\n const auto myFLen = fLen[batch];\n const auto myGLen = gLen[batch];\n\n OffsetCal offsetCal(batch, batchOffset, maxFLen, maxGLen, myGLen, hiddenSize, packOutput);\n const auto myBatchOffset = offsetCal.getBatchOffset();\n const auto strideF = offsetCal.getStrideF();\n scalar_t const* myF = f + batch * maxFLen * hiddenSize + t * hiddenSize;\n scalar_t const* myG = g + batch * maxGLen * hiddenSize + u * hiddenSize;\n scalar_t* mySum = sum + myBatchOffset + t * strideF + u * hiddenSize;\n\n if (t < myFLen and u < myGLen) {\n#pragma unroll\n for (int h = threadIdx.x; h < hiddenSize; h += blockDim.x) {\n if (h < hiddenSize) {\n mySum[h] = myF[h] + myG[h];\n }\n }\n } else if (packOutput == false and t < maxFLen and u < maxGLen) {\n// Need to write finite data to don't-care region because we instantiate the result tensor\n// with torch::empty for performance reasons. Even though it is don't-care region, the\n// contents need to be finite, otherwise could lead to NaN in WGRAD.\n// In packing mode, this write is no longer necessary as we remove the don't-care region\n// from the output.\n// Picking -1 (over 0) here for ease of testing.\n#pragma unroll\n for (int h = threadIdx.x; h < hiddenSize; h += blockDim.x) {\n if (h < hiddenSize) {\n mySum[h] = -1;\n }\n }\n }\n}\n\n/*\nTiled version of the joint forward kernel\nDetail of this joint function can be found in:\n[1] Sequence Transduction with Recurrent Neural Networks.\n\nf is a tensor of shape [batch, T, H]\ng is a tensor of shape [batch, U, H]\nthe transducer joint does\nsum = f.unsqueeze(dim=2) + g.unsqueeze(dim=1)\nThe resultant tensor is of shape [batch, T, U, H]\nEach thread is working on a tile of the shape of tileF x tileG in the result tensor.\nThe input for the tile is first loaded in the register and is reused tileG and tileF times.\n\nThis joint function can optionally pack the output where the output tensor with a shape of\n[B, T, U, H] is packed into [B_packed, H].\nDon't-care region (t > fLen) or (u > gLen) is removed.\nTo enable packing, the starting offset for each batch need to be specified with batchOffset.\n\nOptionally this joint function performs ReLU and/or dropout on the joint output, which is\ncontrolled by arguments relu and dropout, respectively. philoxArgs is argument used for generating\npseudorandom number. When at least one of operations in ReLU and dropout is activated, the joint\nfunction is a masked operation, which is controlled by the template argument masked. In this case,\nmasks are saved to backward.\n*/\ntemplate \n__global__ void transducer_joint_tiled_forward(const scalar_t* f, const scalar_t* g, const int* fLen, const int* gLen,\n const int64_t* batchOffset, int64_t maxFLen, int64_t maxGLen,\n int64_t hiddenSize, int64_t hiddenPerBlock, bool packOutput, bool relu,\n bool dropout, float p, at::PhiloxCudaState philoxArgs, scalar_t* sum,\n uint8_t* mask) {\n static_assert(U == 4, \"U has to be 4, as random numbers are generated in batch of 4\");\n\n const int batch = blockIdx.z;\n const int t = blockIdx.y * tileF;\n const int hiddenBlock = (hiddenSize + hiddenPerBlock - 1) / hiddenPerBlock;\n const int u = blockIdx.x / hiddenBlock * tileG;\n const int hOffset = (blockIdx.x % hiddenBlock) * hiddenPerBlock;\n const int h = threadIdx.x;\n const auto myFLen = fLen[batch];\n const auto myGLen = gLen[batch];\n\n OffsetCal offsetCal(batch, batchOffset, maxFLen, maxGLen, myGLen, hiddenSize, packOutput);\n const auto myBatchOffset = offsetCal.getBatchOffset();\n const auto strideF = offsetCal.getStrideF();\n\n scalar_t const* myF = f + batch * maxFLen * hiddenSize + t * hiddenSize + hOffset;\n scalar_t const* myG = g + batch * maxGLen * hiddenSize + u * hiddenSize + hOffset;\n scalar_t* mySum = sum + myBatchOffset + t * strideF + u * hiddenSize + hOffset;\n uint8_t* myMask = mask + myBatchOffset + t * strideF + u * hiddenSize + hOffset;\n\n // The following code is only needed for dropout. We try to bypass them as much as possible.\n auto seeds = masked ? at::cuda::philox::unpack(philoxArgs)\n : std::make_tuple(static_cast(0), static_cast(0));\n uint64_t tid =\n masked ? (static_cast(blockIdx.z) * gridDim.y * gridDim.x + blockIdx.y * gridDim.x + blockIdx.x) *\n blockDim.x +\n threadIdx.x\n : 0;\n Philox ph(std::get<0>(seeds), tid, std::get<1>(seeds));\n scalar_t scale = masked ? ((p == 0) ? 0 : 1 / p) : 0;\n bool dropoutMask[U];\n\n if (t < myFLen and u < myGLen and hOffset + h < hiddenSize) {\n // register buffers for tiled input reuse\n scalar_t fBuffer[tileF], gBuffer[tileG];\n for (int i = 0; i < tileF; ++i) {\n if (t + i < myFLen) fBuffer[i] = myF[i * hiddenSize + h];\n }\n for (int j = 0; j < tileG; ++j) {\n if (u + j < myGLen) gBuffer[j] = myG[j * hiddenSize + h];\n }\n#pragma unroll\n for (int i = 0; i < tileF; ++i) {\n if (t + i < myFLen) {\n#pragma unroll\n for (int j = 0; j < tileG; ++j) {\n int idx = i * tileG + j;\n if (masked and dropout and idx % U == 0) {\n // For performance, generate 4 random numbers in one shot\n // auto rand4 = curand_uniform4(&state);\n auto rand4 = uniform4(ph());\n dropoutMask[0] = rand4.x < p;\n dropoutMask[1] = rand4.y < p;\n dropoutMask[2] = rand4.z < p;\n dropoutMask[3] = rand4.w < p;\n }\n\n if (u + j < myGLen) {\n scalar_t out = fBuffer[i] + gBuffer[j];\n if (masked) {\n // Apply ReLU here when relu is True\n bool localMask = relu ? (out > 0) : 1;\n localMask = dropout ? localMask & dropoutMask[idx % U] : localMask;\n out = dropout ? out * localMask * scale : out * localMask;\n myMask[i * strideF + j * hiddenSize + h] = static_cast(localMask);\n }\n mySum[i * strideF + j * hiddenSize + h] = out;\n } else if (packOutput == false and u + j < maxGLen)\n mySum[i * strideF + j * hiddenSize + h] = -1;\n }\n } else if (packOutput == false and t + i < maxFLen) {\n// Again need to write finite data to don't-care region\n#pragma unroll\n for (int j = 0; j < tileG; ++j) {\n if (u + j < maxGLen) mySum[i * strideF + j * hiddenSize + h] = -1;\n }\n }\n }\n } else if (packOutput == false and t < maxFLen and u < maxGLen and hOffset + h < hiddenSize) {\n// Only need to ensure the finity in normal mode\n#pragma unroll\n for (int i = 0; i < tileF; ++i) {\n if (t + i < maxFLen) {\n#pragma unroll\n for (int j = 0; j < tileG; ++j) {\n if (u + j < maxGLen) mySum[i * strideF + j * hiddenSize + h] = -1;\n }\n }\n }\n }\n}\n\n/*\nBwd operation (reduction) on one input tensor. Since the operation performed for the two input\ntensors are exactly the same, only one kernel is needed, and the different indexing offsets\nand strides are handled by OffsetCalBwd.\n\nWhen packing is enabled in the fwd op, unpacking is needed to restore the gradients in a\nnon-packed form.\n\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\nand mask contains the mask information.\n*/\ntemplate \n__device__ void transducer_joint_single_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\n int64_t maxGLen, int64_t hiddenSize, bool packOutput,\n bool bwdFasterDim, // whether bwd on the faster moving dimension (u)\n float scale, scalar_t* inGrad, int yBlockOffset = 0) {\n const int batch = blockIdx.z;\n // For the second input tensor, this offset need to be subtracted because the first yBlockOffset\n // sets of thread blocks are for the first input tensor.\n const int x = blockIdx.y - yBlockOffset;\n const int hOffset = blockIdx.x * C10_WARP_SIZE;\n const int wid = threadIdx.y;\n const int lid = threadIdx.x;\n const int numWarp = blockDim.y;\n extern __shared__ char smem8[];\n auto smem = reinterpret_cast(smem8);\n\n OffsetCal offsetCal(batch, batchOffset, fLen, gLen, maxFLen, maxGLen, hiddenSize, packOutput, bwdFasterDim);\n const auto maxXLen = offsetCal.getMaxXLen();\n const auto myXLen = offsetCal.getMyXLen();\n const auto myYLen = offsetCal.getMyYLen();\n scalar_t* myInGrad = inGrad + batch * maxXLen * hiddenSize + x * hiddenSize + hOffset;\n\n if (x < myXLen) {\n const auto myBatchOffset = offsetCal.getBatchOffset();\n const auto strideX = offsetCal.getStrideX();\n const auto strideY = offsetCal.getStrideY();\n const scalar_t* myGrad = grad + myBatchOffset + x * strideX + hOffset;\n const uint8_t* myMask = masked ? mask + myBatchOffset + x * strideX + hOffset : nullptr;\n\n // Each warp reduces numYPerWarp \"y\" first\n acc_t warpSum = 0;\n auto numYPerWarp = (myYLen + numWarp - 1) / numWarp;\n#pragma unroll\n for (int warpY = 0; warpY < numYPerWarp; ++warpY) {\n auto y = wid * numYPerWarp + warpY;\n if (y < myYLen and (hOffset + lid) < hiddenSize)\n if (masked)\n warpSum += static_cast(myGrad[y * strideY + lid]) * myMask[y * strideY + lid] * scale;\n else\n warpSum += myGrad[y * strideY + lid];\n }\n\n // transpose partial sum in SMEM and reduce further using warpReduce\n smem[lid * numWarp + wid] = warpSum;\n __syncthreads();\n auto sum = smem[wid * C10_WARP_SIZE + lid];\n sum = warpReduce(sum, numWarp);\n\n // a a b b c c d d\n // a a b b c c d d\n // a a b b c c d d\n // a a b b c c d d\n // example of 4 warps (a, b, c, d) with 8 threads per warp\n // Each warp need 8 / 4 = 2 threads to write the results.\n if (hOffset + wid * C10_WARP_SIZE / numWarp + lid / numWarp < hiddenSize) {\n if (lid % numWarp == 0) {\n myInGrad[wid * C10_WARP_SIZE / numWarp + lid / numWarp] = sum;\n }\n }\n } else if (wid == 0 and hOffset + lid < hiddenSize) {\n // Need to ensure the grad is zero for don't care region\n myInGrad[lid] = 0;\n }\n}\n\n/*\nActual bwd (reduction) kernel get launched.\nCall transducer_joint_single_backward twice on two input tensors.\nThe two bwd ops are launched together, the first op uses blockIdx.y < maxFLen, and the second op\nuses the rest.\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\nand mask contains the mask information.\n*/\ntemplate \n__global__ void transducer_joint_combined_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\n int64_t maxGLen, int64_t hiddenSize, bool packOutput, float scale,\n scalar_t* fGrad, scalar_t* gGrad) {\n if (blockIdx.y < maxFLen) {\n transducer_joint_single_backward(\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, false, scale, fGrad);\n } else {\n transducer_joint_single_backward(\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, true, scale, gGrad, maxFLen);\n }\n}\n\n/*\nVectorized version of transducer_joint_single_backward\nDoing exact same operation as transducer_joint_single_backward except the load and store are\nvectorized.\nWhen packing is enabled in the fwd op, unpacking is needed to restore the gradients in a\nnon-packed form.\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\nand mask contains the mask information.\n*/\ntemplate \n__device__ void transducer_joint_single_vec_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\n int64_t maxGLen, int64_t hiddenSize, bool packOutput,\n bool bwdFasterDim, float scale, scalar_t* inGrad,\n int yBlockOffset = 0) {\n const int batch = blockIdx.z;\n const int x = blockIdx.y - yBlockOffset;\n const int hOffset = blockIdx.x * C10_WARP_SIZE * V;\n const int wid = threadIdx.y;\n const int lid = threadIdx.x;\n const int numWarp = blockDim.y;\n\n // Figure out the vectorization type for mask\n using mvec_t = mvec_type;\n\n OffsetCal offsetCal(batch, batchOffset, fLen, gLen, maxFLen, maxGLen, hiddenSize, packOutput, bwdFasterDim);\n const auto maxXLen = offsetCal.getMaxXLen();\n const auto myXLen = offsetCal.getMyXLen();\n const auto myYLen = offsetCal.getMyYLen();\n scalar_t* myInGrad = inGrad + batch * maxXLen * hiddenSize + x * hiddenSize + hOffset;\n extern __shared__ char smem8[];\n auto smem = reinterpret_cast(smem8);\n\n acc_t warpSum[V];\n scalar_t inBuffer[V];\n uint8_t maskBuffer[V];\n scalar_t outBuffer[V];\n auto myInGradVec = reinterpret_cast(myInGrad);\n auto outBufferVec = reinterpret_cast(outBuffer);\n\n if (x < myXLen) {\n const auto myBatchOffset = offsetCal.getBatchOffset();\n const auto strideX = offsetCal.getStrideX();\n const auto strideY = offsetCal.getStrideY();\n const scalar_t* myGrad = grad + myBatchOffset + x * strideX + hOffset;\n const uint8_t* myMask = masked ? mask + myBatchOffset + x * strideX + hOffset : nullptr;\n\n for (int i = 0; i < V; ++i) warpSum[i] = 0;\n\n // Each warp reduces numYPerWarp \"y\" first\n auto numYPerWarp = (myYLen + numWarp - 1) / numWarp;\n for (int warpY = 0; warpY < numYPerWarp; ++warpY) {\n auto y = wid * numYPerWarp + warpY;\n auto myGradVec = reinterpret_cast(myGrad + y * strideY);\n auto myMaskVec = masked ? reinterpret_cast(myMask + y * strideY) : nullptr;\n auto inBufferVec = reinterpret_cast(inBuffer);\n auto maskBufferVec = reinterpret_cast(maskBuffer);\n if (hOffset + lid * V < hiddenSize and y < myYLen) {\n *inBufferVec = myGradVec[lid]; // vectorized load\n if (masked) {\n *maskBufferVec = myMaskVec[lid];\n#pragma unroll\n for (int i = 0; i < V; ++i) warpSum[i] += static_cast(inBuffer[i]) * maskBuffer[i] * scale;\n } else {\n#pragma unroll\n for (int i = 0; i < V; ++i) warpSum[i] += inBuffer[i];\n }\n }\n }\n\n // transpose partial sum in SMEM and reduce further using warpReduce\n for (int i = 0; i < V; ++i) {\n smem[lid * numWarp + wid] = warpSum[i];\n __syncthreads();\n auto sum = smem[wid * C10_WARP_SIZE + lid];\n\n if (hOffset + (wid * C10_WARP_SIZE / numWarp) * V < hiddenSize) {\n sum = warpReduce(sum, numWarp);\n if (lid % numWarp == 0) {\n outBuffer[i] = sum;\n }\n }\n __syncthreads();\n }\n\n // a a b b c c d d\n // a a b b c c d d\n // a a b b c c d d\n // a a b b c c d d\n // example of 4 warps (a, b, c, d) with 8 threads per warp\n // Each warp need 8 / 4 = 2 threads to write the results.\n if (lid % numWarp == 0 and hOffset + (wid * C10_WARP_SIZE / numWarp + lid / numWarp) * V < hiddenSize)\n myInGradVec[wid * C10_WARP_SIZE / numWarp + lid / numWarp] = *outBufferVec;\n } else if (wid == 0 and hOffset + lid * V < hiddenSize) {\n // Need to ensure the grad is zero for don't care region\n myInGradVec[lid] = 0;\n }\n}\n\n/*\nVecotrized version of transducer_joint_combined_backward\nCall transducer_joint_single_vec_backward twice on two input tensors.\nThe two bwd ops are launched together, the first op uses blockIdx.y < maxFLen, and the second op\nuses the rest.\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\nand mask contains the mask information.\n*/\ntemplate \n__global__ void transducer_joint_combined_vec_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\n int64_t maxGLen, int64_t hiddenSize, bool packOutput,\n float scale, scalar_t* fGrad, scalar_t* gGrad) {\n if (blockIdx.y < maxFLen) {\n transducer_joint_single_vec_backward(\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, false, scale, fGrad);\n } else {\n transducer_joint_single_vec_backward(\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, true, scale, gGrad, maxFLen);\n }\n}\n\nstd::vector transducer_joint_cuda_forward(torch::Tensor f, torch::Tensor g, torch::Tensor fLen,\n torch::Tensor gLen, torch::Tensor batchOffset,\n int64_t packedBatch, int opt, bool packOutput, bool relu,\n bool dropout, float dropoutProb, int tileSize) {\n auto tensorOpt = f.options();\n auto dtype = f.scalar_type();\n const auto batchSize = f.size(0);\n const auto maxFLen = f.size(1);\n const auto maxGLen = g.size(1);\n const auto hiddenSize = f.size(2);\n bool masked = dropout or relu;\n\n int64_t* batchOffsetPtr = nullptr;\n torch::Tensor sum, mask;\n auto maskOpt = tensorOpt.dtype(torch::kUInt8);\n if (!packOutput) {\n sum = torch::empty({batchSize, maxFLen, maxGLen, hiddenSize}, tensorOpt);\n batchOffsetPtr = nullptr;\n if (masked) mask = torch::empty({batchSize, maxFLen, maxGLen, hiddenSize}, maskOpt);\n } else {\n sum = torch::empty({packedBatch, hiddenSize}, tensorOpt);\n batchOffsetPtr = batchOffset.data_ptr();\n if (masked) mask = torch::empty({packedBatch, hiddenSize}, maskOpt);\n }\n uint8_t* maskPtr = masked ? mask.data_ptr() : nullptr;\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n TORCH_CHECK(opt == 0 or opt == 1, \"Got an invalid optimization level \", opt);\n // Simple heuristics\n const int numThread =\n std::min(128, (static_cast(hiddenSize) + C10_WARP_SIZE - 1) / C10_WARP_SIZE * C10_WARP_SIZE);\n\n if (opt == 0) {\n // vanilla kernel\n const int threads = numThread;\n const dim3 blocks(maxGLen, maxFLen, batchSize);\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n dtype, \"transducer_joint_forward\", ([&] {\n transducer_joint_forward<<>>(\n f.data_ptr(), g.data_ptr(), fLen.data_ptr(), gLen.data_ptr(),\n batchOffsetPtr, maxFLen, maxGLen, hiddenSize, packOutput, sum.data_ptr());\n }));\n }\n if (opt == 1) {\n // tiled version. For simplicity, assume tileF == tileG, even though the kernel can\n // support more general cases.\n const int threads = numThread;\n const int hiddenPerBlock = numThread;\n const int hiddenBlock = (hiddenSize + hiddenPerBlock - 1) / hiddenPerBlock;\n const dim3 blocks((maxGLen + tileSize - 1) / tileSize * hiddenBlock, (maxFLen + tileSize - 1) / tileSize,\n batchSize);\n\n TORCH_CHECK(tileSize == 1 or tileSize == 2 or tileSize == 4, \"Expected tileSize to be in [1, 2, 4], but got \",\n tileSize);\n\n at::PhiloxCudaState rng_engine_inputs;\n if (masked) {\n // set up PRG when the input is masked. rng_engine_inputs will be used as a space filler\n // for non-masked calls.\n // Therefore no need to initialize.\n c10::optional gen_;\n auto gen = at::get_generator_or_default(gen_, at::cuda::detail::getDefaultCUDAGenerator());\n // counterOffset records how many cuRAND calls each thread makes. For a tiled kernel,\n // each thread processes tileF * tileG output elements.\n int64_t counterOffset = tileSize * tileSize;\n {\n std::lock_guard lock(gen->mutex_);\n rng_engine_inputs = gen->philox_cuda_state(counterOffset);\n }\n }\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n dtype, \"transducer_joint_forward\", ([&] {\n void (*kernel)(const scalar_t*, const scalar_t*, const int*, const int*, const int64_t*, int64_t, int64_t,\n int64_t, int64_t, bool, bool, bool, float, at::PhiloxCudaState, scalar_t*, uint8_t*);\n if (masked) {\n switch (tileSize) {\n case 2:\n kernel = &transducer_joint_tiled_forward;\n break;\n case 4:\n kernel = &transducer_joint_tiled_forward;\n break;\n }\n } else {\n switch (tileSize) {\n case 1:\n kernel = &transducer_joint_tiled_forward;\n break;\n case 2:\n kernel = &transducer_joint_tiled_forward;\n break;\n case 4:\n kernel = &transducer_joint_tiled_forward;\n break;\n }\n }\n\n kernel<<>>(f.data_ptr(), g.data_ptr(), fLen.data_ptr(),\n gLen.data_ptr(), batchOffsetPtr, maxFLen, maxGLen, hiddenSize,\n hiddenPerBlock, packOutput, relu, dropout, 1.0f - dropoutProb,\n rng_engine_inputs, sum.data_ptr(), maskPtr);\n }));\n }\n\n C10_CUDA_CHECK(cudaGetLastError());\n if (masked)\n return {sum, mask};\n else\n return {sum};\n}\n\nstd::vector transducer_joint_cuda_backward(std::vector in, torch::Tensor fLen,\n torch::Tensor gLen, torch::Tensor batchOffset, int maxFLen,\n int maxGLen, bool packOutput, float scale) {\n auto grad = in[0];\n bool masked = (in.size() == 2);\n uint8_t* maskPtr = masked ? in[1].data_ptr() : nullptr;\n\n auto tensorOpt = grad.options();\n auto dtype = grad.scalar_type();\n const int batchSize = fLen.size(0);\n const int hiddenSize = grad.size(-1);\n\n const auto deviceProperties = at::cuda::getCurrentDeviceProperties();\n const int maxNumWarp = deviceProperties->maxThreadsPerBlock / C10_WARP_SIZE;\n\n torch::Tensor fGrad = torch::empty({batchSize, maxFLen, hiddenSize}, tensorOpt);\n torch::Tensor gGrad = torch::empty({batchSize, maxGLen, hiddenSize}, tensorOpt);\n\n int64_t* batchOffsetPtr = (!packOutput) ? nullptr : batchOffset.data_ptr();\n\n // The number \"y\" I would like each thread to work on\n const int workPerThread = 32;\n // Since the bwd for f and g have the same thread block size, we need to use the max of the two.\n int numWarp = largestPowerOfTwo((std::max(maxFLen, maxGLen) + workPerThread - 1) / workPerThread);\n // Would like to have at least 2 warps\n numWarp = std::max(2, numWarp);\n // cap on the maximum number of warps allowed\n numWarp = std::min(maxNumWarp, numWarp);\n\n // Need smem for transposing the partial sum. The partial sum is in a matrix of the shape\n // numWarp x warpSize\n const int smemSize = numWarp * C10_WARP_SIZE;\n const dim3 threads(C10_WARP_SIZE, numWarp, 1);\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n dtype, \"transducer_joint_cuda_backward_kernel\", ([&] {\n auto gradPtr = grad.data_ptr();\n auto fLenPtr = fLen.data_ptr();\n auto gLenPtr = gLen.data_ptr();\n auto fGradPtr = fGrad.data_ptr();\n auto gGradPtr = gGrad.data_ptr();\n\n // resolve the acc_t type\n using acc_t = at::acc_type;\n using vec_t = uint64_t;\n\n constexpr int vectFactor = sizeof(vec_t) / sizeof(scalar_t);\n constexpr int vecAlignment = std::alignment_of::value;\n\n // if all input and output tensors meet the alignment requirement\n bool memAlign = (reinterpret_cast(gradPtr) % vecAlignment == 0) and\n (reinterpret_cast(fGradPtr) % vecAlignment == 0) and\n (reinterpret_cast(gGradPtr) % vecAlignment == 0);\n\n if (vectFactor > 1 and hiddenSize % vectFactor == 0 and memAlign) {\n // If vectorization helps and the alignment requirement is met, use the vectorized\n // kernel. For simplicity, hiddenSize needs to be a multiple vecFactor.\n const dim3 blocks((hiddenSize + C10_WARP_SIZE * vectFactor - 1) / (C10_WARP_SIZE * vectFactor),\n maxFLen + maxGLen, batchSize);\n if (masked) {\n transducer_joint_combined_vec_backward\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\n fGradPtr, gGradPtr);\n } else {\n transducer_joint_combined_vec_backward\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\n fGradPtr, gGradPtr);\n }\n } else {\n const dim3 blocks((hiddenSize + C10_WARP_SIZE - 1) / C10_WARP_SIZE, maxFLen + maxGLen, batchSize);\n if (masked) {\n transducer_joint_combined_backward\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\n fGradPtr, gGradPtr);\n } else {\n transducer_joint_combined_backward\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\n fGradPtr, gGradPtr);\n }\n }\n }));\n\n return {fGrad, gGrad};\n}\n", "output": "#include \n#include \n#include \n#include \n\n#ifdef OLD_GENERATOR_PATH\n#include \n#else\n#include \n#endif\n\n#include \n#include \n\n#include \n\n#include \"philox.cuh\"\n\n// Warp reduce kernels to reduce N groups of data into N numbers, where N = warpSize / width.\n// width should be a power of 2 and should be less than warpSize.\ntemplate \n__device__ __forceinline__ scalar_t warpReduce(scalar_t x, int width = C10_WARP_SIZE) {\n for (unsigned offset = width / 2; offset > 0; offset /= 2) {\n x += __shfl_down_sync(0xffffffff, x, offset, width);\n }\n return x;\n}\n\ninline int largestPowerOfTwo(int x) {\n int y = 1;\n while (y <= x) y <<= 1;\n return y >> 1;\n}\n\n/*\nFigure out vectorization type for masks.\nSimilar to how PyTorch figures out acc_t here:\naten/src/ATen/AccumulateType.h\n*/\ntemplate \nstruct MaskVecType {};\n\ntemplate <>\nstruct MaskVecType<1> {\n using type = uint8_t;\n};\ntemplate <>\nstruct MaskVecType<2> {\n using type = uint16_t;\n};\ntemplate <>\nstruct MaskVecType<4> {\n using type = uint32_t;\n};\n\ntemplate \nusing mvec_type = typename MaskVecType::type;\n\n// Helper class to calculate pointer offset that can be shared by different flavors of kernels.\n// For fwd, batch offset and stride are different for packing and non-packing mode.\nstruct OffsetCalFwd {\n __device__ __forceinline__ OffsetCalFwd(int64_t batch, const int64_t* batchOffset, int64_t maxFLen, int64_t maxGLen,\n int64_t gLen, int64_t hiddenSize, bool packOutput)\n : batch(batch),\n batchOffset(batchOffset),\n maxFLen(maxFLen),\n maxGLen(maxGLen),\n gLen(gLen),\n hiddenSize(hiddenSize),\n packOutput(packOutput) {}\n\n int64_t batch;\n const int64_t* batchOffset;\n int64_t maxFLen;\n int64_t maxGLen;\n int64_t gLen;\n int64_t hiddenSize;\n bool packOutput;\n\n __device__ __forceinline__ int64_t getBatchOffset() {\n return packOutput ? ((batch == 0) ? 0 : batchOffset[batch - 1]) * hiddenSize\n : batch * maxFLen * maxGLen * hiddenSize;\n }\n\n __device__ __forceinline__ int64_t getStrideF() { return packOutput ? gLen * hiddenSize : maxGLen * hiddenSize; }\n};\n\n// Helper class to calculate pointer offset that can be shared by different flavors of kernels\n// For bwd, batch offset and stride are different for packing and non-packing mode.\n// The reducion is done for two input tensors. Therefore, generating two sets of offsets\n// according to bwdFasterDim can lead to a unified implementation in the actual kernel.\nstruct OffsetCalBwd {\n __device__ __forceinline__ OffsetCalBwd(int64_t batch, const int64_t* batchOffset, const int* fLen, const int* gLen,\n int64_t maxFLen, int64_t maxGLen, int64_t hiddenSize, bool packOutput,\n bool bwdFasterDim)\n : batch(batch),\n batchOffset(batchOffset),\n maxFLen(maxFLen),\n maxGLen(maxGLen),\n fLen(fLen),\n gLen(gLen),\n hiddenSize(hiddenSize),\n packOutput(packOutput),\n bwdFasterDim(bwdFasterDim) {}\n\n int64_t batch;\n const int64_t* batchOffset;\n const int* fLen;\n const int* gLen;\n int64_t maxFLen;\n int64_t maxGLen;\n int64_t hiddenSize;\n bool packOutput;\n bool bwdFasterDim; // whether doing bwd on the faster moving dimension\n\n __device__ __forceinline__ int64_t getBatchOffset() {\n return packOutput ? ((batch == 0) ? 0 : batchOffset[batch - 1]) * hiddenSize\n : batch * maxFLen * maxGLen * hiddenSize;\n }\n\n __device__ __forceinline__ int64_t getMaxXLen() { return bwdFasterDim ? maxGLen : maxFLen; }\n\n __device__ __forceinline__ auto getMyXLen() -> decltype(gLen[batch]) {\n return bwdFasterDim ? gLen[batch] : fLen[batch];\n }\n\n __device__ __forceinline__ auto getMyYLen() -> decltype(gLen[batch]) {\n return bwdFasterDim ? fLen[batch] : gLen[batch];\n }\n\n __device__ __forceinline__ int64_t getStrideX() {\n return bwdFasterDim ? hiddenSize : ((packOutput ? gLen[batch] : maxGLen) * hiddenSize);\n }\n\n __device__ __forceinline__ int64_t getStrideY() {\n return bwdFasterDim ? ((packOutput ? gLen[batch] : maxGLen) * hiddenSize) : hiddenSize;\n }\n};\n\n// Vanila transducer joint forward kernel\n// Detail of this joint function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n\n// f is a tensor of shape [batch, T, H]\n// g is a tensor of shape [batch, U, H]\n// the transducer joint does\n// sum = f.unsqueeze(dim=2) + g.unsqueeze(dim=1)\n// The resultant tensor is of shape [batch, T, U, H]\n// Each thread block is working on one \"batch\" of data in the output tensor, [batch, t, u, :]\n\n// This joint function can optionally pack the output where the output tensor with a shape of\n// [B, T, U, H] is packed into [B_packed, H].\n// Don't-care region (t > fLen) or (u > gLen) is removed.\n// To enable packing, the starting offset for each batch need to be specified with batchOffset.\ntemplate \n__global__ void transducer_joint_forward(const scalar_t* f, const scalar_t* g, const int* fLen, const int* gLen,\n const int64_t* batchOffset, int64_t maxFLen, int64_t maxGLen,\n int64_t hiddenSize, bool packOutput, scalar_t* sum) {\n const int batch = blockIdx.z;\n const int t = blockIdx.y;\n const int u = blockIdx.x;\n const auto myFLen = fLen[batch];\n const auto myGLen = gLen[batch];\n\n OffsetCal offsetCal(batch, batchOffset, maxFLen, maxGLen, myGLen, hiddenSize, packOutput);\n const auto myBatchOffset = offsetCal.getBatchOffset();\n const auto strideF = offsetCal.getStrideF();\n scalar_t const* myF = f + batch * maxFLen * hiddenSize + t * hiddenSize;\n scalar_t const* myG = g + batch * maxGLen * hiddenSize + u * hiddenSize;\n scalar_t* mySum = sum + myBatchOffset + t * strideF + u * hiddenSize;\n\n if (t < myFLen and u < myGLen) {\n#pragma unroll\n for (int h = threadIdx.x; h < hiddenSize; h += blockDim.x) {\n if (h < hiddenSize) {\n mySum[h] = myF[h] + myG[h];\n }\n }\n } else if (packOutput == false and t < maxFLen and u < maxGLen) {\n// Need to write finite data to don't-care region because we instantiate the result tensor\n// with torch::empty for performance reasons. Even though it is don't-care region, the\n// contents need to be finite, otherwise could lead to NaN in WGRAD.\n// In packing mode, this write is no longer necessary as we remove the don't-care region\n// from the output.\n// Picking -1 (over 0) here for ease of testing.\n#pragma unroll\n for (int h = threadIdx.x; h < hiddenSize; h += blockDim.x) {\n if (h < hiddenSize) {\n mySum[h] = -1;\n }\n }\n }\n}\n\n/*\nTiled version of the joint forward kernel\nDetail of this joint function can be found in:\n[1] Sequence Transduction with Recurrent Neural Networks.\n\nf is a tensor of shape [batch, T, H]\ng is a tensor of shape [batch, U, H]\nthe transducer joint does\nsum = f.unsqueeze(dim=2) + g.unsqueeze(dim=1)\nThe resultant tensor is of shape [batch, T, U, H]\nEach thread is working on a tile of the shape of tileF x tileG in the result tensor.\nThe input for the tile is first loaded in the register and is reused tileG and tileF times.\n\nThis joint function can optionally pack the output where the output tensor with a shape of\n[B, T, U, H] is packed into [B_packed, H].\nDon't-care region (t > fLen) or (u > gLen) is removed.\nTo enable packing, the starting offset for each batch need to be specified with batchOffset.\n\nOptionally this joint function performs ReLU and/or dropout on the joint output, which is\ncontrolled by arguments relu and dropout, respectively. philoxArgs is argument used for generating\npseudorandom number. When at least one of operations in ReLU and dropout is activated, the joint\nfunction is a masked operation, which is controlled by the template argument masked. In this case,\nmasks are saved to backward.\n*/\ntemplate \n__global__ void transducer_joint_tiled_forward(const scalar_t* f, const scalar_t* g, const int* fLen, const int* gLen,\n const int64_t* batchOffset, int64_t maxFLen, int64_t maxGLen,\n int64_t hiddenSize, int64_t hiddenPerBlock, bool packOutput, bool relu,\n bool dropout, float p, at::PhiloxCudaState philoxArgs, scalar_t* sum,\n uint8_t* mask) {\n static_assert(U == 4, \"U has to be 4, as random numbers are generated in batch of 4\");\n\n const int batch = blockIdx.z;\n const int t = blockIdx.y * tileF;\n const int hiddenBlock = (hiddenSize + hiddenPerBlock - 1) / hiddenPerBlock;\n const int u = blockIdx.x / hiddenBlock * tileG;\n const int hOffset = (blockIdx.x % hiddenBlock) * hiddenPerBlock;\n const int h = threadIdx.x;\n const auto myFLen = fLen[batch];\n const auto myGLen = gLen[batch];\n\n OffsetCal offsetCal(batch, batchOffset, maxFLen, maxGLen, myGLen, hiddenSize, packOutput);\n const auto myBatchOffset = offsetCal.getBatchOffset();\n const auto strideF = offsetCal.getStrideF();\n\n scalar_t const* myF = f + batch * maxFLen * hiddenSize + t * hiddenSize + hOffset;\n scalar_t const* myG = g + batch * maxGLen * hiddenSize + u * hiddenSize + hOffset;\n scalar_t* mySum = sum + myBatchOffset + t * strideF + u * hiddenSize + hOffset;\n uint8_t* myMask = mask + myBatchOffset + t * strideF + u * hiddenSize + hOffset;\n\n // The following code is only needed for dropout. We try to bypass them as much as possible.\n auto seeds = masked ? at::cuda::philox::unpack(philoxArgs)\n : std::make_tuple(static_cast(0), static_cast(0));\n uint64_t tid =\n masked ? (static_cast(blockIdx.z) * gridDim.y * gridDim.x + blockIdx.y * gridDim.x + blockIdx.x) *\n blockDim.x +\n threadIdx.x\n : 0;\n Philox ph(std::get<0>(seeds), tid, std::get<1>(seeds));\n scalar_t scale = masked ? ((p == 0) ? 0 : 1 / p) : 0;\n bool dropoutMask[U];\n\n if (t < myFLen and u < myGLen and hOffset + h < hiddenSize) {\n // register buffers for tiled input reuse\n scalar_t fBuffer[tileF], gBuffer[tileG];\n for (int i = 0; i < tileF; ++i) {\n if (t + i < myFLen) fBuffer[i] = myF[i * hiddenSize + h];\n }\n for (int j = 0; j < tileG; ++j) {\n if (u + j < myGLen) gBuffer[j] = myG[j * hiddenSize + h];\n }\n#pragma unroll\n for (int i = 0; i < tileF; ++i) {\n if (t + i < myFLen) {\n#pragma unroll\n for (int j = 0; j < tileG; ++j) {\n int idx = i * tileG + j;\n if (masked and dropout and idx % U == 0) {\n // For performance, generate 4 random numbers in one shot\n // auto rand4 = curand_uniform4(&state);\n auto rand4 = uniform4(ph());\n dropoutMask[0] = rand4.x < p;\n dropoutMask[1] = rand4.y < p;\n dropoutMask[2] = rand4.z < p;\n dropoutMask[3] = rand4.w < p;\n }\n\n if (u + j < myGLen) {\n scalar_t out = fBuffer[i] + gBuffer[j];\n if (masked) {\n // Apply ReLU here when relu is True\n bool localMask = relu ? (out > 0) : 1;\n localMask = dropout ? localMask & dropoutMask[idx % U] : localMask;\n out = dropout ? out * localMask * scale : out * localMask;\n myMask[i * strideF + j * hiddenSize + h] = static_cast(localMask);\n }\n mySum[i * strideF + j * hiddenSize + h] = out;\n } else if (packOutput == false and u + j < maxGLen)\n mySum[i * strideF + j * hiddenSize + h] = -1;\n }\n } else if (packOutput == false and t + i < maxFLen) {\n// Again need to write finite data to don't-care region\n#pragma unroll\n for (int j = 0; j < tileG; ++j) {\n if (u + j < maxGLen) mySum[i * strideF + j * hiddenSize + h] = -1;\n }\n }\n }\n } else if (packOutput == false and t < maxFLen and u < maxGLen and hOffset + h < hiddenSize) {\n// Only need to ensure the finity in normal mode\n#pragma unroll\n for (int i = 0; i < tileF; ++i) {\n if (t + i < maxFLen) {\n#pragma unroll\n for (int j = 0; j < tileG; ++j) {\n if (u + j < maxGLen) mySum[i * strideF + j * hiddenSize + h] = -1;\n }\n }\n }\n }\n}\n\n/*\nBwd operation (reduction) on one input tensor. Since the operation performed for the two input\ntensors are exactly the same, only one kernel is needed, and the different indexing offsets\nand strides are handled by OffsetCalBwd.\n\nWhen packing is enabled in the fwd op, unpacking is needed to restore the gradients in a\nnon-packed form.\n\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\nand mask contains the mask information.\n*/\ntemplate \n__device__ void transducer_joint_single_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\n int64_t maxGLen, int64_t hiddenSize, bool packOutput,\n bool bwdFasterDim, // whether bwd on the faster moving dimension (u)\n float scale, scalar_t* inGrad, int yBlockOffset = 0) {\n const int batch = blockIdx.z;\n // For the second input tensor, this offset need to be subtracted because the first yBlockOffset\n // sets of thread blocks are for the first input tensor.\n const int x = blockIdx.y - yBlockOffset;\n const int hOffset = blockIdx.x * C10_WARP_SIZE;\n const int wid = threadIdx.y;\n const int lid = threadIdx.x;\n const int numWarp = blockDim.y;\n extern __shared__ char smem8[];\n auto smem = reinterpret_cast(smem8);\n\n OffsetCal offsetCal(batch, batchOffset, fLen, gLen, maxFLen, maxGLen, hiddenSize, packOutput, bwdFasterDim);\n const auto maxXLen = offsetCal.getMaxXLen();\n const auto myXLen = offsetCal.getMyXLen();\n const auto myYLen = offsetCal.getMyYLen();\n scalar_t* myInGrad = inGrad + batch * maxXLen * hiddenSize + x * hiddenSize + hOffset;\n\n if (x < myXLen) {\n const auto myBatchOffset = offsetCal.getBatchOffset();\n const auto strideX = offsetCal.getStrideX();\n const auto strideY = offsetCal.getStrideY();\n const scalar_t* myGrad = grad + myBatchOffset + x * strideX + hOffset;\n const uint8_t* myMask = masked ? mask + myBatchOffset + x * strideX + hOffset : nullptr;\n\n // Each warp reduces numYPerWarp \"y\" first\n acc_t warpSum = 0;\n auto numYPerWarp = (myYLen + numWarp - 1) / numWarp;\n#pragma unroll\n for (int warpY = 0; warpY < numYPerWarp; ++warpY) {\n auto y = wid * numYPerWarp + warpY;\n if (y < myYLen and (hOffset + lid) < hiddenSize)\n if (masked)\n warpSum += static_cast(myGrad[y * strideY + lid]) * myMask[y * strideY + lid] * scale;\n else\n warpSum += myGrad[y * strideY + lid];\n }\n\n // transpose partial sum in SMEM and reduce further using warpReduce\n smem[lid * numWarp + wid] = warpSum;\n __syncthreads();\n auto sum = smem[wid * C10_WARP_SIZE + lid];\n sum = warpReduce(sum, numWarp);\n\n // a a b b c c d d\n // a a b b c c d d\n // a a b b c c d d\n // a a b b c c d d\n // example of 4 warps (a, b, c, d) with 8 threads per warp\n // Each warp need 8 / 4 = 2 threads to write the results.\n if (hOffset + wid * C10_WARP_SIZE / numWarp + lid / numWarp < hiddenSize) {\n if (lid % numWarp == 0) {\n myInGrad[wid * C10_WARP_SIZE / numWarp + lid / numWarp] = sum;\n }\n }\n } else if (wid == 0 and hOffset + lid < hiddenSize) {\n // Need to ensure the grad is zero for don't care region\n myInGrad[lid] = 0;\n }\n}\n\n/*\nActual bwd (reduction) kernel get launched.\nCall transducer_joint_single_backward twice on two input tensors.\nThe two bwd ops are launched together, the first op uses blockIdx.y < maxFLen, and the second op\nuses the rest.\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\nand mask contains the mask information.\n*/\ntemplate \n__global__ void transducer_joint_combined_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\n int64_t maxGLen, int64_t hiddenSize, bool packOutput, float scale,\n scalar_t* fGrad, scalar_t* gGrad) {\n if (blockIdx.y < maxFLen) {\n transducer_joint_single_backward(\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, false, scale, fGrad);\n } else {\n transducer_joint_single_backward(\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, true, scale, gGrad, maxFLen);\n }\n}\n\n/*\nVectorized version of transducer_joint_single_backward\nDoing exact same operation as transducer_joint_single_backward except the load and store are\nvectorized.\nWhen packing is enabled in the fwd op, unpacking is needed to restore the gradients in a\nnon-packed form.\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\nand mask contains the mask information.\n*/\ntemplate \n__device__ void transducer_joint_single_vec_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\n int64_t maxGLen, int64_t hiddenSize, bool packOutput,\n bool bwdFasterDim, float scale, scalar_t* inGrad,\n int yBlockOffset = 0) {\n const int batch = blockIdx.z;\n const int x = blockIdx.y - yBlockOffset;\n const int hOffset = blockIdx.x * C10_WARP_SIZE * V;\n const int wid = threadIdx.y;\n const int lid = threadIdx.x;\n const int numWarp = blockDim.y;\n\n // Figure out the vectorization type for mask\n using mvec_t = mvec_type;\n\n OffsetCal offsetCal(batch, batchOffset, fLen, gLen, maxFLen, maxGLen, hiddenSize, packOutput, bwdFasterDim);\n const auto maxXLen = offsetCal.getMaxXLen();\n const auto myXLen = offsetCal.getMyXLen();\n const auto myYLen = offsetCal.getMyYLen();\n scalar_t* myInGrad = inGrad + batch * maxXLen * hiddenSize + x * hiddenSize + hOffset;\n extern __shared__ char smem8[];\n auto smem = reinterpret_cast(smem8);\n\n acc_t warpSum[V];\n scalar_t inBuffer[V];\n uint8_t maskBuffer[V];\n scalar_t outBuffer[V];\n auto myInGradVec = reinterpret_cast(myInGrad);\n auto outBufferVec = reinterpret_cast(outBuffer);\n\n if (x < myXLen) {\n const auto myBatchOffset = offsetCal.getBatchOffset();\n const auto strideX = offsetCal.getStrideX();\n const auto strideY = offsetCal.getStrideY();\n const scalar_t* myGrad = grad + myBatchOffset + x * strideX + hOffset;\n const uint8_t* myMask = masked ? mask + myBatchOffset + x * strideX + hOffset : nullptr;\n\n for (int i = 0; i < V; ++i) warpSum[i] = 0;\n\n // Each warp reduces numYPerWarp \"y\" first\n auto numYPerWarp = (myYLen + numWarp - 1) / numWarp;\n for (int warpY = 0; warpY < numYPerWarp; ++warpY) {\n auto y = wid * numYPerWarp + warpY;\n auto myGradVec = reinterpret_cast(myGrad + y * strideY);\n auto myMaskVec = masked ? reinterpret_cast(myMask + y * strideY) : nullptr;\n auto inBufferVec = reinterpret_cast(inBuffer);\n auto maskBufferVec = reinterpret_cast(maskBuffer);\n if (hOffset + lid * V < hiddenSize and y < myYLen) {\n *inBufferVec = myGradVec[lid]; // vectorized load\n if (masked) {\n *maskBufferVec = myMaskVec[lid];\n#pragma unroll\n for (int i = 0; i < V; ++i) warpSum[i] += static_cast(inBuffer[i]) * maskBuffer[i] * scale;\n } else {\n#pragma unroll\n for (int i = 0; i < V; ++i) warpSum[i] += inBuffer[i];\n }\n }\n }\n\n // transpose partial sum in SMEM and reduce further using warpReduce\n for (int i = 0; i < V; ++i) {\n smem[lid * numWarp + wid] = warpSum[i];\n __syncthreads();\n auto sum = smem[wid * C10_WARP_SIZE + lid];\n\n if (hOffset + (wid * C10_WARP_SIZE / numWarp) * V < hiddenSize) {\n sum = warpReduce(sum, numWarp);\n if (lid % numWarp == 0) {\n outBuffer[i] = sum;\n }\n }\n __syncthreads();\n }\n\n // a a b b c c d d\n // a a b b c c d d\n // a a b b c c d d\n // a a b b c c d d\n // example of 4 warps (a, b, c, d) with 8 threads per warp\n // Each warp need 8 / 4 = 2 threads to write the results.\n if (lid % numWarp == 0 and hOffset + (wid * C10_WARP_SIZE / numWarp + lid / numWarp) * V < hiddenSize)\n myInGradVec[wid * C10_WARP_SIZE / numWarp + lid / numWarp] = *outBufferVec;\n } else if (wid == 0 and hOffset + lid * V < hiddenSize) {\n // Need to ensure the grad is zero for don't care region\n myInGradVec[lid] = 0;\n }\n}\n\n/*\nVecotrized version of transducer_joint_combined_backward\nCall transducer_joint_single_vec_backward twice on two input tensors.\nThe two bwd ops are launched together, the first op uses blockIdx.y < maxFLen, and the second op\nuses the rest.\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\nand mask contains the mask information.\n*/\ntemplate \n__global__ void transducer_joint_combined_vec_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\n int64_t maxGLen, int64_t hiddenSize, bool packOutput,\n float scale, scalar_t* fGrad, scalar_t* gGrad) {\n if (blockIdx.y < maxFLen) {\n transducer_joint_single_vec_backward(\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, false, scale, fGrad);\n } else {\n transducer_joint_single_vec_backward(\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, true, scale, gGrad, maxFLen);\n }\n}\n\nstd::vector transducer_joint_cuda_forward(torch::Tensor f, torch::Tensor g, torch::Tensor fLen,\n torch::Tensor gLen, torch::Tensor batchOffset,\n int64_t packedBatch, int opt, bool packOutput, bool relu,\n bool dropout, float dropoutProb, int tileSize) {\n auto tensorOpt = f.options();\n auto dtype = f.scalar_type();\n const auto batchSize = f.size(0);\n const auto maxFLen = f.size(1);\n const auto maxGLen = g.size(1);\n const auto hiddenSize = f.size(2);\n bool masked = dropout or relu;\n\n int64_t* batchOffsetPtr = nullptr;\n torch::Tensor sum, mask;\n auto maskOpt = tensorOpt.dtype(torch::kUInt8);\n if (!packOutput) {\n sum = torch::empty({batchSize, maxFLen, maxGLen, hiddenSize}, tensorOpt);\n batchOffsetPtr = nullptr;\n if (masked) mask = torch::empty({batchSize, maxFLen, maxGLen, hiddenSize}, maskOpt);\n } else {\n sum = torch::empty({packedBatch, hiddenSize}, tensorOpt);\n batchOffsetPtr = batchOffset.data_ptr();\n if (masked) mask = torch::empty({packedBatch, hiddenSize}, maskOpt);\n }\n uint8_t* maskPtr = masked ? mask.data_ptr() : nullptr;\n\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n\n TORCH_CHECK(opt == 0 or opt == 1, \"Got an invalid optimization level \", opt);\n // Simple heuristics\n const int numThread =\n std::min(128, (static_cast(hiddenSize) + C10_WARP_SIZE - 1) / C10_WARP_SIZE * C10_WARP_SIZE);\n\n if (opt == 0) {\n // vanilla kernel\n const int threads = numThread;\n const dim3 blocks(maxGLen, maxFLen, batchSize);\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n dtype, \"transducer_joint_forward\", ([&] {\n transducer_joint_forward<<>>(\n f.data_ptr(), g.data_ptr(), fLen.data_ptr(), gLen.data_ptr(),\n batchOffsetPtr, maxFLen, maxGLen, hiddenSize, packOutput, sum.data_ptr());\n }));\n }\n if (opt == 1) {\n // tiled version. For simplicity, assume tileF == tileG, even though the kernel can\n // support more general cases.\n const int threads = numThread;\n const int hiddenPerBlock = numThread;\n const int hiddenBlock = (hiddenSize + hiddenPerBlock - 1) / hiddenPerBlock;\n const dim3 blocks((maxGLen + tileSize - 1) / tileSize * hiddenBlock, (maxFLen + tileSize - 1) / tileSize,\n batchSize);\n\n TORCH_CHECK(tileSize == 1 or tileSize == 2 or tileSize == 4, \"Expected tileSize to be in [1, 2, 4], but got \",\n tileSize);\n\n at::PhiloxCudaState rng_engine_inputs;\n if (masked) {\n // set up PRG when the input is masked. rng_engine_inputs will be used as a space filler\n // for non-masked calls.\n // Therefore no need to initialize.\n c10::optional gen_;\n auto gen = at::get_generator_or_default(gen_, at::cuda::detail::getDefaultCUDAGenerator());\n // counterOffset records how many cuRAND calls each thread makes. For a tiled kernel,\n // each thread processes tileF * tileG output elements.\n int64_t counterOffset = tileSize * tileSize;\n {\n std::lock_guard lock(gen->mutex_);\n rng_engine_inputs = gen->philox_cuda_state(counterOffset);\n }\n }\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n dtype, \"transducer_joint_forward\", ([&] {\n void (*kernel)(const scalar_t*, const scalar_t*, const int*, const int*, const int64_t*, int64_t, int64_t,\n int64_t, int64_t, bool, bool, bool, float, at::PhiloxCudaState, scalar_t*, uint8_t*);\n if (masked) {\n switch (tileSize) {\n case 2:\n kernel = &transducer_joint_tiled_forward;\n break;\n case 4:\n kernel = &transducer_joint_tiled_forward;\n break;\n }\n } else {\n switch (tileSize) {\n case 1:\n kernel = &transducer_joint_tiled_forward;\n break;\n case 2:\n kernel = &transducer_joint_tiled_forward;\n break;\n case 4:\n kernel = &transducer_joint_tiled_forward;\n break;\n }\n }\n\n kernel<<>>(f.data_ptr(), g.data_ptr(), fLen.data_ptr(),\n gLen.data_ptr(), batchOffsetPtr, maxFLen, maxGLen, hiddenSize,\n hiddenPerBlock, packOutput, relu, dropout, 1.0f - dropoutProb,\n rng_engine_inputs, sum.data_ptr(), maskPtr);\n }));\n }\n\n C10_CUDA_CHECK(hipGetLastError());\n if (masked)\n return {sum, mask};\n else\n return {sum};\n}\n\nstd::vector transducer_joint_cuda_backward(std::vector in, torch::Tensor fLen,\n torch::Tensor gLen, torch::Tensor batchOffset, int maxFLen,\n int maxGLen, bool packOutput, float scale) {\n auto grad = in[0];\n bool masked = (in.size() == 2);\n uint8_t* maskPtr = masked ? in[1].data_ptr() : nullptr;\n\n auto tensorOpt = grad.options();\n auto dtype = grad.scalar_type();\n const int batchSize = fLen.size(0);\n const int hiddenSize = grad.size(-1);\n\n const auto deviceProperties = at::cuda::getCurrentDeviceProperties();\n const int maxNumWarp = deviceProperties->maxThreadsPerBlock / C10_WARP_SIZE;\n\n torch::Tensor fGrad = torch::empty({batchSize, maxFLen, hiddenSize}, tensorOpt);\n torch::Tensor gGrad = torch::empty({batchSize, maxGLen, hiddenSize}, tensorOpt);\n\n int64_t* batchOffsetPtr = (!packOutput) ? nullptr : batchOffset.data_ptr();\n\n // The number \"y\" I would like each thread to work on\n const int workPerThread = 32;\n // Since the bwd for f and g have the same thread block size, we need to use the max of the two.\n int numWarp = largestPowerOfTwo((std::max(maxFLen, maxGLen) + workPerThread - 1) / workPerThread);\n // Would like to have at least 2 warps\n numWarp = std::max(2, numWarp);\n // cap on the maximum number of warps allowed\n numWarp = std::min(maxNumWarp, numWarp);\n\n // Need smem for transposing the partial sum. The partial sum is in a matrix of the shape\n // numWarp x warpSize\n const int smemSize = numWarp * C10_WARP_SIZE;\n const dim3 threads(C10_WARP_SIZE, numWarp, 1);\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n dtype, \"transducer_joint_cuda_backward_kernel\", ([&] {\n auto gradPtr = grad.data_ptr();\n auto fLenPtr = fLen.data_ptr();\n auto gLenPtr = gLen.data_ptr();\n auto fGradPtr = fGrad.data_ptr();\n auto gGradPtr = gGrad.data_ptr();\n\n // resolve the acc_t type\n using acc_t = at::acc_type;\n using vec_t = uint64_t;\n\n constexpr int vectFactor = sizeof(vec_t) / sizeof(scalar_t);\n constexpr int vecAlignment = std::alignment_of::value;\n\n // if all input and output tensors meet the alignment requirement\n bool memAlign = (reinterpret_cast(gradPtr) % vecAlignment == 0) and\n (reinterpret_cast(fGradPtr) % vecAlignment == 0) and\n (reinterpret_cast(gGradPtr) % vecAlignment == 0);\n\n if (vectFactor > 1 and hiddenSize % vectFactor == 0 and memAlign) {\n // If vectorization helps and the alignment requirement is met, use the vectorized\n // kernel. For simplicity, hiddenSize needs to be a multiple vecFactor.\n const dim3 blocks((hiddenSize + C10_WARP_SIZE * vectFactor - 1) / (C10_WARP_SIZE * vectFactor),\n maxFLen + maxGLen, batchSize);\n if (masked) {\n transducer_joint_combined_vec_backward\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\n fGradPtr, gGradPtr);\n } else {\n transducer_joint_combined_vec_backward\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\n fGradPtr, gGradPtr);\n }\n } else {\n const dim3 blocks((hiddenSize + C10_WARP_SIZE - 1) / C10_WARP_SIZE, maxFLen + maxGLen, batchSize);\n if (masked) {\n transducer_joint_combined_backward\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\n fGradPtr, gGradPtr);\n } else {\n transducer_joint_combined_backward\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\n fGradPtr, gGradPtr);\n }\n }\n }));\n\n return {fGrad, gGrad};\n}\n", "source": "/kaggle/working/apex/apex/contrib/csrc/transducer/transducer_joint_kernel.cu"} {"input": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n#include \n\n#include \n\n#include \"group_norm_nhwc.h\"\n#include \"macros.h\"\n#include \"traits.h\"\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// F O R W A R D\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ void group_norm_nhwc_fwd_sum_kernel(Group_norm_nhwc_fwd_params params) {\n // The traits.\n using Traits = Traits_;\n // The IO traits.\n using IOTraits = typename Traits::IOTraits;\n\n // The IO type\n using IOType = typename IOTraits::Type;\n // The IO doubled type\n using IOType2 = typename IOTraits::Type2;\n\n // The object in charge of doing the sums for the different blocks.\n typedef cub::BlockScan Block_scan;\n\n // Allocate shared memory for Block_scan.\n __shared__ typename Block_scan::TempStorage temp_storage;\n // Allocate shared memory for the groups. We could reduce the amount of shared memory reserved.\n __shared__ float2 smem[THREADS_PER_BLOCK];\n\n // The instance in the batch.\n int ni = blockIdx.z;\n // The channel loaded by that thread (2 channels per thread for F16x2).\n int ci = blockIdx.x * params.channels_per_block + threadIdx.x * 2;\n\n // The first activation loaded by that block.\n int hw_begin = blockIdx.y * params.acts_per_block;\n // The last activation loaded by that block.\n int hw_end = min((int64_t)hw_begin + params.acts_per_block, params.hw);\n\n // The sums.\n float sum = 0.f, sum_sq = 0.f;\n\n // Iterate over the activations to compute the sums.\n for (int hwi = hw_begin; hwi < hw_end; ++hwi) {\n // The offset.\n int64_t offset = (int64_t)ni * params.hwc + hwi * params.c + ci;\n\n // Fetch two channels per thread.\n IOType2 v2 = IOTraits::zero();\n if (ci < params.c) {\n v2 = *reinterpret_cast(&reinterpret_cast(params.x)[offset]);\n }\n\n // Extract the two values.\n float2 f2 = IOTraits::unpack(v2);\n\n // Update the sum.\n sum += f2.x + f2.y;\n // Update the sum of squares.\n sum_sq += f2.x * f2.x + f2.y * f2.y;\n }\n\n // The group that thread works on and the channel in the group (modulus).\n int gj = threadIdx.x * 2 / params.channels_per_group;\n int cj = threadIdx.x * 2 - params.channels_per_group * gj;\n\n // The data for the summations.\n Group_sums inp{cj == 0 ? 1 : 0, sum, sum_sq};\n\n // Do the segmented scan.\n Group_sums out;\n Block_scan(temp_storage).InclusiveScan(inp, out, Group_sums_op());\n\n // Store the results for the groups in shared memory (to produce coalesced stores later).\n if (cj == params.channels_per_group - 2 /* 2 channels per thread */) {\n smem[gj] = make_float2(out.sum, out.sum_sq);\n }\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n // The global group index.\n int gk = blockIdx.x * params.groups_per_block + threadIdx.x;\n\n // Threads that have nothing left to do, exit.\n if (threadIdx.x >= params.groups_per_block || gk >= params.groups) {\n return;\n }\n\n // The first threads (those storing to global memory, load the values).\n float2 sums = smem[threadIdx.x];\n\n // Store to global memory.\n atomicAdd(¶ms.zeroed_red_buffer[(2 * ni + 0) * params.groups + gk], sums.x);\n atomicAdd(¶ms.zeroed_red_buffer[(2 * ni + 1) * params.groups + gk], sums.y);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_fwd_two_passes_setup(Group_norm_nhwc_fwd_params& params, size_t& zeroed_red_buffer_elts) {\n // The pre-computed dimensions.\n params.hw = params.h * params.w;\n params.hwc = params.c * params.hw;\n\n // The number of channels per group.\n params.channels_per_group = params.c / params.groups;\n // The inverse to compute the mean/variance.\n params.inv_hwc_per_group = 1.f / (float)(params.hw * params.channels_per_group);\n\n // Define the number of blocks per activation map. That's a simple heuristic.\n int blocks_per_act_slice = 0;\n if (params.c >= 1280) {\n blocks_per_act_slice = 128 / params.n;\n } else if (params.c >= 640) {\n blocks_per_act_slice = 256 / params.n;\n } else {\n blocks_per_act_slice = 512 / params.n;\n }\n\n // Clamp to at least 1 to avoid divide-by-zero when batch size is large.\n blocks_per_act_slice = max(blocks_per_act_slice, 1);\n\n // Make sure we launch blocks per activation is no less than activations\n blocks_per_act_slice = min(blocks_per_act_slice, div_up(params.hw, params.n));\n\n // Define how many activations are computed per block.\n params.acts_per_block = div_up(params.hw, blocks_per_act_slice);\n // The number of channels per block.\n params.channels_per_block = 320;\n // Special case to deal with 30 channels per group.\n if (params.channels_per_block % params.channels_per_group != 0) {\n params.channels_per_block = 240;\n }\n\n // Special case to deal with 70 channels per group.\n if (params.c == 2240) {\n params.channels_per_block = 280;\n } else if (params.c == 832) {\n params.channels_per_block = 208;\n }\n\n if (params.c % params.channels_per_block != 0) {\n if (params.c % 512 == 0 && params.c != 1536 && params.c != 3072 && params.c % 448 != 0) {\n params.channels_per_block = 512;\n } else if (params.c % 42 == 0) {\n params.channels_per_block = 336;\n } else if (params.c % 384 == 0) {\n params.channels_per_block = 384;\n } else if (params.c % 256 == 0 && params.c % 448 != 0 && params.c % 392 != 0) {\n params.channels_per_block = 256;\n } else if (params.c % 128 == 0 && params.c % 448 != 0 && params.c % 392 != 0) {\n params.channels_per_block = 128;\n } else if (params.c % 448 == 0 && params.c % 392 != 0) {\n params.channels_per_block = 448;\n } else if (params.c % 392 == 0) {\n params.channels_per_block = 392;\n }\n }\n\n // The number of groups per block.\n params.groups_per_block = params.channels_per_block / params.channels_per_group;\n\n // Make sure the number of channels is a multiple of the number of channels per block.\n assert(params.c % params.channels_per_block == 0);\n // Make sure a group does not span multiple blocks.\n assert(params.channels_per_block % params.channels_per_group == 0);\n\n // The number of elements in the reduction buffer (for the sums and sums of squared).\n zeroed_red_buffer_elts = params.n * params.groups * 2;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_fwd_two_passes_sum(const Group_norm_nhwc_fwd_params& params, cudaStream_t stream) {\n // The dimension of the grid.\n dim3 grid;\n\n // The number of blocks to compute all the channels.\n grid.x = params.c / params.channels_per_block;\n // The number of blocks to compute all the activations in a given instance.\n grid.y = div_up(params.hw, params.acts_per_block);\n // The number of instances.\n grid.z = params.n;\n\n // Launch the kernel.\n if (params.precision == PrecisionMode::FP16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp16IOFp16W)\n } else if (params.precision == PrecisionMode::FP16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp16IOBf16W)\n } else if (params.precision == PrecisionMode::FP16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp16IOFp32W)\n } else if (params.precision == PrecisionMode::BF16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Bf16IOFp16W)\n } else if (params.precision == PrecisionMode::BF16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Bf16IOBf16W)\n } else if (params.precision == PrecisionMode::BF16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Bf16IOFp32W)\n } else if (params.precision == PrecisionMode::FP32IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp32IOFp16W)\n } else if (params.precision == PrecisionMode::FP32IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp32IOBf16W)\n } else {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp32IOFp32W)\n }\n\n // Make sure it launched ok.\n CHECK_CUDA(cudaGetLastError());\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ void group_norm_nhwc_fwd_scale_kernel(Group_norm_nhwc_fwd_params params) {\n // The traits.\n using Traits = Traits_;\n // The IO traits.\n using IOTraits = typename Traits::IOTraits;\n // The Weights traits.\n using WTraits = typename Traits::WTraits;\n\n // The IO type\n using IOType = typename IOTraits::Type;\n // The IO doubled type\n using IOType2 = typename IOTraits::Type2;\n\n // Weights type\n using WType = typename WTraits::Type;\n // Weights doubled type\n using WType2 = typename WTraits::Type2;\n\n // The instance in the batch.\n int ni = blockIdx.z;\n // The channel loaded by that thread (2 channels per thread for F16x2).\n int ci = blockIdx.x * params.channels_per_block + threadIdx.x * 2;\n // The group that thread works on and the channel in the group (modulus).\n int gi = ci / params.channels_per_group;\n\n // Load the sum and sum of squares for the group.\n float sum = 0.f, sum_sq = 0.f;\n if (gi < params.groups) {\n sum = params.zeroed_red_buffer[(2 * ni + 0) * params.groups + gi];\n sum_sq = params.zeroed_red_buffer[(2 * ni + 1) * params.groups + gi];\n }\n\n // Load gamma/beta.\n float2 gamma_f2, beta_f2;\n if (ci < params.c) {\n gamma_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.gamma)[ci]));\n beta_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.beta)[ci]));\n }\n\n // Compute the mean.\n float mean = sum * params.inv_hwc_per_group;\n // Compute the variance.\n float var = sum_sq * params.inv_hwc_per_group - (mean * mean);\n // Compute the inverse of the stddev.\n float inv_stddev = var <= 0.f ? 1.f : rsqrtf(var + params.epsilon);\n\n // The first activation loaded by that block.\n int hw_begin = blockIdx.y * params.acts_per_block;\n // The last activation loaded by that block.\n int hw_end = min((int64_t)hw_begin + params.acts_per_block, params.hw);\n\n // Iterate over the activations to compute the sums.\n for (int hwi = hw_begin; hwi < hw_end; ++hwi) {\n // The src/dst offset.\n int64_t offset = (int64_t)ni * params.hwc + hwi * params.c + ci;\n\n // Fetch two channels per thread.\n IOType2 v2 = IOTraits::zero();\n if (ci < params.c) {\n v2 = *reinterpret_cast(&reinterpret_cast(params.x)[offset]);\n }\n\n // Extract the two values.\n float2 f2 = IOTraits::unpack(v2);\n\n // Normalize the channels.\n f2.x = (f2.x - mean) * inv_stddev;\n f2.y = (f2.y - mean) * inv_stddev;\n\n // Scale by gamma and add beta.\n f2.x = gamma_f2.x * f2.x + beta_f2.x;\n f2.y = gamma_f2.y * f2.y + beta_f2.y;\n\n // Apply Swish if needed.\n if (params.with_swish) {\n f2.x = f2.x * sigmoid(f2.x);\n f2.y = f2.y * sigmoid(f2.y);\n }\n\n // Store the scaled values.\n if (ci < params.c) {\n *reinterpret_cast(&reinterpret_cast(params.y)[offset]) = IOTraits::pack(f2);\n }\n }\n\n // Write the sums if needed.\n if (params.sums != nullptr && gi < params.groups) {\n float2 sums;\n sums.x = sum * params.inv_hwc_per_group;\n sums.y = sum_sq * params.inv_hwc_per_group;\n params.sums[ni * params.groups + gi] = sums;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_fwd_two_passes_scale(const Group_norm_nhwc_fwd_params& params, cudaStream_t stream) {\n // The dimension of the grid.\n dim3 grid;\n\n // The number of blocks to compute all the channels.\n grid.x = params.c / params.channels_per_block;\n // The number of blocks to compute all the activations in a given instance.\n grid.y = div_up(params.hw, params.acts_per_block);\n // The number of instances.\n grid.z = params.n;\n\n if (params.precision == PrecisionMode::FP16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp16IOFp16W)\n } else if (params.precision == PrecisionMode::FP16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp16IOBf16W)\n } else if (params.precision == PrecisionMode::FP16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp16IOFp32W)\n } else if (params.precision == PrecisionMode::BF16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Bf16IOFp16W)\n } else if (params.precision == PrecisionMode::BF16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Bf16IOBf16W)\n } else if (params.precision == PrecisionMode::BF16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Bf16IOFp32W)\n } else if (params.precision == PrecisionMode::FP32IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp32IOFp16W)\n } else if (params.precision == PrecisionMode::FP32IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp32IOBf16W)\n } else {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp32IOFp32W)\n }\n\n // Make sure it launched ok.\n CHECK_CUDA(cudaGetLastError());\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n", "output": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \n\n#include \"group_norm_nhwc.h\"\n#include \"macros.h\"\n#include \"traits.h\"\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// F O R W A R D\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ void group_norm_nhwc_fwd_sum_kernel(Group_norm_nhwc_fwd_params params) {\n // The traits.\n using Traits = Traits_;\n // The IO traits.\n using IOTraits = typename Traits::IOTraits;\n\n // The IO type\n using IOType = typename IOTraits::Type;\n // The IO doubled type\n using IOType2 = typename IOTraits::Type2;\n\n // The object in charge of doing the sums for the different blocks.\n typedef hipcub::BlockScan Block_scan;\n\n // Allocate shared memory for Block_scan.\n __shared__ typename Block_scan::TempStorage temp_storage;\n // Allocate shared memory for the groups. We could reduce the amount of shared memory reserved.\n __shared__ float2 smem[THREADS_PER_BLOCK];\n\n // The instance in the batch.\n int ni = blockIdx.z;\n // The channel loaded by that thread (2 channels per thread for F16x2).\n int ci = blockIdx.x * params.channels_per_block + threadIdx.x * 2;\n\n // The first activation loaded by that block.\n int hw_begin = blockIdx.y * params.acts_per_block;\n // The last activation loaded by that block.\n int hw_end = min((int64_t)hw_begin + params.acts_per_block, params.hw);\n\n // The sums.\n float sum = 0.f, sum_sq = 0.f;\n\n // Iterate over the activations to compute the sums.\n for (int hwi = hw_begin; hwi < hw_end; ++hwi) {\n // The offset.\n int64_t offset = (int64_t)ni * params.hwc + hwi * params.c + ci;\n\n // Fetch two channels per thread.\n IOType2 v2 = IOTraits::zero();\n if (ci < params.c) {\n v2 = *reinterpret_cast(&reinterpret_cast(params.x)[offset]);\n }\n\n // Extract the two values.\n float2 f2 = IOTraits::unpack(v2);\n\n // Update the sum.\n sum += f2.x + f2.y;\n // Update the sum of squares.\n sum_sq += f2.x * f2.x + f2.y * f2.y;\n }\n\n // The group that thread works on and the channel in the group (modulus).\n int gj = threadIdx.x * 2 / params.channels_per_group;\n int cj = threadIdx.x * 2 - params.channels_per_group * gj;\n\n // The data for the summations.\n Group_sums inp{cj == 0 ? 1 : 0, sum, sum_sq};\n\n // Do the segmented scan.\n Group_sums out;\n Block_scan(temp_storage).InclusiveScan(inp, out, Group_sums_op());\n\n // Store the results for the groups in shared memory (to produce coalesced stores later).\n if (cj == params.channels_per_group - 2 /* 2 channels per thread */) {\n smem[gj] = make_float2(out.sum, out.sum_sq);\n }\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n // The global group index.\n int gk = blockIdx.x * params.groups_per_block + threadIdx.x;\n\n // Threads that have nothing left to do, exit.\n if (threadIdx.x >= params.groups_per_block || gk >= params.groups) {\n return;\n }\n\n // The first threads (those storing to global memory, load the values).\n float2 sums = smem[threadIdx.x];\n\n // Store to global memory.\n atomicAdd(¶ms.zeroed_red_buffer[(2 * ni + 0) * params.groups + gk], sums.x);\n atomicAdd(¶ms.zeroed_red_buffer[(2 * ni + 1) * params.groups + gk], sums.y);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_fwd_two_passes_setup(Group_norm_nhwc_fwd_params& params, size_t& zeroed_red_buffer_elts) {\n // The pre-computed dimensions.\n params.hw = params.h * params.w;\n params.hwc = params.c * params.hw;\n\n // The number of channels per group.\n params.channels_per_group = params.c / params.groups;\n // The inverse to compute the mean/variance.\n params.inv_hwc_per_group = 1.f / (float)(params.hw * params.channels_per_group);\n\n // Define the number of blocks per activation map. That's a simple heuristic.\n int blocks_per_act_slice = 0;\n if (params.c >= 1280) {\n blocks_per_act_slice = 128 / params.n;\n } else if (params.c >= 640) {\n blocks_per_act_slice = 256 / params.n;\n } else {\n blocks_per_act_slice = 512 / params.n;\n }\n\n // Clamp to at least 1 to avoid divide-by-zero when batch size is large.\n blocks_per_act_slice = max(blocks_per_act_slice, 1);\n\n // Make sure we launch blocks per activation is no less than activations\n blocks_per_act_slice = min(blocks_per_act_slice, div_up(params.hw, params.n));\n\n // Define how many activations are computed per block.\n params.acts_per_block = div_up(params.hw, blocks_per_act_slice);\n // The number of channels per block.\n params.channels_per_block = 320;\n // Special case to deal with 30 channels per group.\n if (params.channels_per_block % params.channels_per_group != 0) {\n params.channels_per_block = 240;\n }\n\n // Special case to deal with 70 channels per group.\n if (params.c == 2240) {\n params.channels_per_block = 280;\n } else if (params.c == 832) {\n params.channels_per_block = 208;\n }\n\n if (params.c % params.channels_per_block != 0) {\n if (params.c % 512 == 0 && params.c != 1536 && params.c != 3072 && params.c % 448 != 0) {\n params.channels_per_block = 512;\n } else if (params.c % 42 == 0) {\n params.channels_per_block = 336;\n } else if (params.c % 384 == 0) {\n params.channels_per_block = 384;\n } else if (params.c % 256 == 0 && params.c % 448 != 0 && params.c % 392 != 0) {\n params.channels_per_block = 256;\n } else if (params.c % 128 == 0 && params.c % 448 != 0 && params.c % 392 != 0) {\n params.channels_per_block = 128;\n } else if (params.c % 448 == 0 && params.c % 392 != 0) {\n params.channels_per_block = 448;\n } else if (params.c % 392 == 0) {\n params.channels_per_block = 392;\n }\n }\n\n // The number of groups per block.\n params.groups_per_block = params.channels_per_block / params.channels_per_group;\n\n // Make sure the number of channels is a multiple of the number of channels per block.\n assert(params.c % params.channels_per_block == 0);\n // Make sure a group does not span multiple blocks.\n assert(params.channels_per_block % params.channels_per_group == 0);\n\n // The number of elements in the reduction buffer (for the sums and sums of squared).\n zeroed_red_buffer_elts = params.n * params.groups * 2;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_fwd_two_passes_sum(const Group_norm_nhwc_fwd_params& params, hipStream_t stream) {\n // The dimension of the grid.\n dim3 grid;\n\n // The number of blocks to compute all the channels.\n grid.x = params.c / params.channels_per_block;\n // The number of blocks to compute all the activations in a given instance.\n grid.y = div_up(params.hw, params.acts_per_block);\n // The number of instances.\n grid.z = params.n;\n\n // Launch the kernel.\n if (params.precision == PrecisionMode::FP16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp16IOFp16W)\n } else if (params.precision == PrecisionMode::FP16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp16IOBf16W)\n } else if (params.precision == PrecisionMode::FP16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp16IOFp32W)\n } else if (params.precision == PrecisionMode::BF16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Bf16IOFp16W)\n } else if (params.precision == PrecisionMode::BF16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Bf16IOBf16W)\n } else if (params.precision == PrecisionMode::BF16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Bf16IOFp32W)\n } else if (params.precision == PrecisionMode::FP32IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp32IOFp16W)\n } else if (params.precision == PrecisionMode::FP32IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp32IOBf16W)\n } else {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp32IOFp32W)\n }\n\n // Make sure it launched ok.\n CHECK_CUDA(hipGetLastError());\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ void group_norm_nhwc_fwd_scale_kernel(Group_norm_nhwc_fwd_params params) {\n // The traits.\n using Traits = Traits_;\n // The IO traits.\n using IOTraits = typename Traits::IOTraits;\n // The Weights traits.\n using WTraits = typename Traits::WTraits;\n\n // The IO type\n using IOType = typename IOTraits::Type;\n // The IO doubled type\n using IOType2 = typename IOTraits::Type2;\n\n // Weights type\n using WType = typename WTraits::Type;\n // Weights doubled type\n using WType2 = typename WTraits::Type2;\n\n // The instance in the batch.\n int ni = blockIdx.z;\n // The channel loaded by that thread (2 channels per thread for F16x2).\n int ci = blockIdx.x * params.channels_per_block + threadIdx.x * 2;\n // The group that thread works on and the channel in the group (modulus).\n int gi = ci / params.channels_per_group;\n\n // Load the sum and sum of squares for the group.\n float sum = 0.f, sum_sq = 0.f;\n if (gi < params.groups) {\n sum = params.zeroed_red_buffer[(2 * ni + 0) * params.groups + gi];\n sum_sq = params.zeroed_red_buffer[(2 * ni + 1) * params.groups + gi];\n }\n\n // Load gamma/beta.\n float2 gamma_f2, beta_f2;\n if (ci < params.c) {\n gamma_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.gamma)[ci]));\n beta_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.beta)[ci]));\n }\n\n // Compute the mean.\n float mean = sum * params.inv_hwc_per_group;\n // Compute the variance.\n float var = sum_sq * params.inv_hwc_per_group - (mean * mean);\n // Compute the inverse of the stddev.\n float inv_stddev = var <= 0.f ? 1.f : rsqrtf(var + params.epsilon);\n\n // The first activation loaded by that block.\n int hw_begin = blockIdx.y * params.acts_per_block;\n // The last activation loaded by that block.\n int hw_end = min((int64_t)hw_begin + params.acts_per_block, params.hw);\n\n // Iterate over the activations to compute the sums.\n for (int hwi = hw_begin; hwi < hw_end; ++hwi) {\n // The src/dst offset.\n int64_t offset = (int64_t)ni * params.hwc + hwi * params.c + ci;\n\n // Fetch two channels per thread.\n IOType2 v2 = IOTraits::zero();\n if (ci < params.c) {\n v2 = *reinterpret_cast(&reinterpret_cast(params.x)[offset]);\n }\n\n // Extract the two values.\n float2 f2 = IOTraits::unpack(v2);\n\n // Normalize the channels.\n f2.x = (f2.x - mean) * inv_stddev;\n f2.y = (f2.y - mean) * inv_stddev;\n\n // Scale by gamma and add beta.\n f2.x = gamma_f2.x * f2.x + beta_f2.x;\n f2.y = gamma_f2.y * f2.y + beta_f2.y;\n\n // Apply Swish if needed.\n if (params.with_swish) {\n f2.x = f2.x * sigmoid(f2.x);\n f2.y = f2.y * sigmoid(f2.y);\n }\n\n // Store the scaled values.\n if (ci < params.c) {\n *reinterpret_cast(&reinterpret_cast(params.y)[offset]) = IOTraits::pack(f2);\n }\n }\n\n // Write the sums if needed.\n if (params.sums != nullptr && gi < params.groups) {\n float2 sums;\n sums.x = sum * params.inv_hwc_per_group;\n sums.y = sum_sq * params.inv_hwc_per_group;\n params.sums[ni * params.groups + gi] = sums;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_fwd_two_passes_scale(const Group_norm_nhwc_fwd_params& params, hipStream_t stream) {\n // The dimension of the grid.\n dim3 grid;\n\n // The number of blocks to compute all the channels.\n grid.x = params.c / params.channels_per_block;\n // The number of blocks to compute all the activations in a given instance.\n grid.y = div_up(params.hw, params.acts_per_block);\n // The number of instances.\n grid.z = params.n;\n\n if (params.precision == PrecisionMode::FP16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp16IOFp16W)\n } else if (params.precision == PrecisionMode::FP16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp16IOBf16W)\n } else if (params.precision == PrecisionMode::FP16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp16IOFp32W)\n } else if (params.precision == PrecisionMode::BF16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Bf16IOFp16W)\n } else if (params.precision == PrecisionMode::BF16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Bf16IOBf16W)\n } else if (params.precision == PrecisionMode::BF16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Bf16IOFp32W)\n } else if (params.precision == PrecisionMode::FP32IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp32IOFp16W)\n } else if (params.precision == PrecisionMode::FP32IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp32IOBf16W)\n } else {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp32IOFp32W)\n }\n\n // Make sure it launched ok.\n CHECK_CUDA(hipGetLastError());\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n", "source": "/kaggle/working/apex/apex/contrib/csrc/group_norm/group_norm_nhwc_fwd_two_pass.cu"} {"input": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n#include \n\n#include \n\n#include \"group_norm_nhwc.h\"\n#include \"macros.h\"\n#include \"traits.h\"\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// B A C K W A R D\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ void group_norm_nhwc_bwd_sum_kernel(Group_norm_nhwc_bwd_params params) {\n // The IO traits.\n using Traits = Traits_;\n // The IO traits.\n using IOTraits = typename Traits::IOTraits;\n // The Weights traits.\n using WTraits = typename Traits::WTraits;\n\n // The IO type\n using IOType = typename IOTraits::Type;\n // The IO doubled type\n using IOType2 = typename IOTraits::Type2;\n\n // Weights type\n using WType = typename WTraits::Type;\n // Weights doubled type\n using WType2 = typename WTraits::Type2;\n\n // The object in charge of doing the sums for the different blocks.\n typedef cub::BlockScan Block_scan;\n\n // Allocate shared memory for Block_scan.\n __shared__ typename Block_scan::TempStorage temp_storage;\n // Allocate shared memory for the groups. We could reduce the amount of shared memory reserved.\n __shared__ float2 smem[THREADS_PER_BLOCK];\n\n // The instance in the batch.\n int ni = blockIdx.z;\n // The channel loaded by that thread (2 channels per thread for F16x2).\n int ci = blockIdx.x * params.channels_per_block + threadIdx.x * 2;\n // The group that thread works on and the channel in the group (modulus).\n int gi = ci / params.channels_per_group;\n\n // The sums from the fwd pass.\n float2 fwd = params.sums[ni * params.groups + gi];\n // The mean of X (computed during the fwd pass -- one value per batch*group).\n float x_mean = fwd.x;\n // The mean of squares of X (computed during the fwd pass -- one value per batch*group).\n float x_sq_mean = fwd.y;\n // The variance.\n float x_var = x_sq_mean - x_mean * x_mean;\n // The reciprocal of the standard deviation (i.e. 1.f / sqrt(var + epsilon)).\n float rcp_x_stddev = x_var <= 0.f ? 1.f : 1.f / sqrtf(x_var + params.epsilon);\n\n // Load gamma.\n float2 gamma_f2 = make_float2(0.f, 0.f);\n float2 beta_f2 = make_float2(0.f, 0.f);\n if (ci < params.c) {\n gamma_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.gamma)[ci]));\n if (params.with_swish) {\n beta_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.beta)[ci]));\n }\n }\n\n // The group that thread works on and the channel in the group (modulus).\n int gj = threadIdx.x * 2 / params.channels_per_group;\n int cj = threadIdx.x * 2 - params.channels_per_group * gj;\n\n // The first activation loaded by that block.\n int hw_begin = blockIdx.y * params.acts_per_block;\n // The last activation loaded by that block.\n int hw_end = min((int64_t)hw_begin + params.acts_per_block, params.hw);\n\n // The gradients for gamma/beta.\n float2 dgamma = make_float2(0.f, 0.f), dbeta = make_float2(0.f, 0.f);\n // Accumulated gradients for dgrad calculation\n float mean_1 = 0.f, mean_2 = 0.f;\n\n // Iterate over the activations to compute the sums.\n for (int hwi = hw_begin; hwi < hw_end; ++hwi) {\n // The offset.\n int64_t offset = (int64_t)ni * params.hwc + hwi * params.c + ci;\n\n // Fetch two channels per thread.\n IOType2 x_v2 = IOTraits::zero();\n IOType2 dy_v2 = IOTraits::zero();\n if (ci < params.c) {\n x_v2 = *reinterpret_cast(&reinterpret_cast(params.x)[offset]);\n dy_v2 = *reinterpret_cast(&reinterpret_cast(params.dy)[offset]);\n }\n\n // Extract the two half values.\n float2 x_f2 = IOTraits::unpack(x_v2);\n float2 dy_f2 = IOTraits::unpack(dy_v2);\n\n // X - X_mean.\n float x_minus_x_mean_x = x_f2.x - x_mean;\n float x_minus_x_mean_y = x_f2.y - x_mean;\n\n // Normalize X.\n float x_norm_x = x_minus_x_mean_x * rcp_x_stddev;\n float x_norm_y = x_minus_x_mean_y * rcp_x_stddev;\n\n if (params.with_swish) {\n float x_gn_x = x_norm_x * gamma_f2.x + beta_f2.x;\n float x_gn_y = x_norm_y * gamma_f2.y + beta_f2.y;\n float s_x = sigmoid(x_gn_x);\n float s_y = sigmoid(x_gn_y);\n dy_f2.x = dy_f2.x * s_x * (1.f + x_gn_x * (1.f - s_x));\n dy_f2.y = dy_f2.y * s_y * (1.f + x_gn_y * (1.f - s_y));\n }\n\n // Update beta.\n dbeta.x += dy_f2.x;\n dbeta.y += dy_f2.y;\n\n // Update dgamma.\n dgamma.x += dy_f2.x * x_norm_x;\n dgamma.y += dy_f2.y * x_norm_y;\n\n // The gradient that enters the x_norm node.\n float dx_norm_x = dy_f2.x * gamma_f2.x;\n float dx_norm_y = dy_f2.y * gamma_f2.y;\n\n // Add to the 1st mean.\n mean_1 += dx_norm_x * x_norm_x;\n mean_1 += dx_norm_y * x_norm_y;\n\n // Add to the 2nd mean.\n mean_2 += dx_norm_x;\n mean_2 += dx_norm_y;\n }\n\n // The data for the summations.\n Group_sums inp{cj == 0 ? 1 : 0, mean_1, mean_2};\n\n // Do the segmented scan.\n Group_sums out;\n Block_scan(temp_storage).InclusiveScan(inp, out, Group_sums_op());\n\n // Store the results for the groups in shared memory (to produce coalesced stores later).\n if (cj == params.channels_per_group - 2 /* 2 channels per thread */) {\n smem[gj] = make_float2(out.sum, out.sum_sq);\n }\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n // The global group index.\n int gk = blockIdx.x * params.groups_per_block + threadIdx.x;\n\n // The first threads (those storing to global memory, load the values).\n float2 sums = smem[threadIdx.x];\n\n // Store to global memory.\n if (threadIdx.x < params.groups_per_block && gk < params.groups) {\n atomicAdd(¶ms.zeroed_red_buffer[(2 * ni + 0) * params.groups + gk], sums.x);\n atomicAdd(¶ms.zeroed_red_buffer[(2 * ni + 1) * params.groups + gk], sums.y);\n }\n\n // The base pointer for the gradients for gamma and beta.\n float* dgamma_beta_ptr = ¶ms.zeroed_red_buffer[params.n * params.groups * 2];\n\n // The 1st channel in the output tensor. NOTE: Two channels per thread store interleaved.\n int ck = blockIdx.x * params.channels_per_block + threadIdx.x;\n\n // Store dgamma and dbeta as well.\n if (ck < params.c) {\n atomicAdd(&dgamma_beta_ptr[0 * params.c + 0 * blockDim.x + ck], dgamma.x);\n atomicAdd(&dgamma_beta_ptr[0 * params.c + 1 * blockDim.x + ck], dgamma.y);\n atomicAdd(&dgamma_beta_ptr[1 * params.c + 0 * blockDim.x + ck], dbeta.x);\n atomicAdd(&dgamma_beta_ptr[1 * params.c + 1 * blockDim.x + ck], dbeta.y);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_bwd_two_passes_setup(Group_norm_nhwc_bwd_params& params, size_t& zeroed_red_buffer_elts) {\n // The pre-computed dimensions.\n params.hw = params.h * params.w;\n params.hwc = params.c * params.hw;\n\n // The number of channels per group.\n params.channels_per_group = params.c / params.groups;\n // The inverse to compute the mean/variance.\n params.inv_hwc_per_group = 1.f / (float)(params.hw * params.channels_per_group);\n\n // Define the number of blocks per activation map. That's a simple heuristic.\n int blocks_per_act_slice = 0;\n if (params.c >= 1280) {\n blocks_per_act_slice = 128 / params.n;\n } else if (params.c >= 640) {\n blocks_per_act_slice = 256 / params.n;\n } else {\n blocks_per_act_slice = 512 / params.n;\n }\n\n // Clamp to at least 1 to avoid divide-by-zero when batch size is large.\n blocks_per_act_slice = max(blocks_per_act_slice, 1);\n\n // Make sure we launch blocks per activation is no less than activations\n blocks_per_act_slice = min(blocks_per_act_slice, div_up(params.hw, params.n));\n\n // Define how many activations are computed per block.\n params.acts_per_block = div_up(params.hw, blocks_per_act_slice);\n\n // The number of channels per block.\n params.channels_per_block = 320;\n // Special case to deal with 30 channels per group.\n if (params.channels_per_block % params.channels_per_group != 0) {\n params.channels_per_block = 240;\n }\n\n // Special case to deal with 70 channels per group.\n if (params.c == 2240) {\n params.channels_per_block = 280;\n } else if (params.c == 832) {\n params.channels_per_block = 208;\n }\n\n if (params.c % params.channels_per_block != 0) {\n if (params.c % 512 == 0 && params.c != 1536 && params.c != 3072 && params.c % 448 != 0) {\n params.channels_per_block = 512;\n } else if (params.c % 42 == 0) {\n params.channels_per_block = 336;\n } else if (params.c % 384 == 0) {\n params.channels_per_block = 384;\n } else if (params.c % 256 == 0 && params.c % 448 != 0 && params.c % 392 != 0) {\n params.channels_per_block = 256;\n } else if (params.c % 128 == 0 && params.c % 448 != 0 && params.c % 392 != 0) {\n params.channels_per_block = 128;\n } else if (params.c % 448 == 0 && params.c % 392 != 0) {\n params.channels_per_block = 448;\n } else if (params.c % 392 == 0) {\n params.channels_per_block = 392;\n }\n }\n\n // The number of groups per block.\n params.groups_per_block = params.channels_per_block / params.channels_per_group;\n\n // Make sure the number of channels is a multiple of the number of channels per block.\n assert(params.c % params.channels_per_block == 0);\n // Make sure a group does not span multiple blocks.\n assert(params.channels_per_block % params.channels_per_group == 0);\n\n // The number of elements in the reduction buffer (for the sums and sums of squared).\n zeroed_red_buffer_elts = params.n * params.groups * 2 + params.c * 2;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_bwd_two_passes_sum(const Group_norm_nhwc_bwd_params& params, cudaStream_t stream) {\n // The dimension of the grid.\n dim3 grid;\n\n // The number of blocks to compute all the channels.\n grid.x = params.c / params.channels_per_block;\n // The number of blocks to compute all the activations in a given instance.\n grid.y = div_up(params.hw, params.acts_per_block);\n // The number of instances.\n grid.z = params.n;\n\n if (params.precision == PrecisionMode::FP16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp16IOFp16W)\n } else if (params.precision == PrecisionMode::FP16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp16IOBf16W)\n } else if (params.precision == PrecisionMode::FP16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp16IOFp32W)\n } else if (params.precision == PrecisionMode::BF16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Bf16IOFp16W)\n } else if (params.precision == PrecisionMode::BF16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Bf16IOBf16W)\n } else if (params.precision == PrecisionMode::BF16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Bf16IOFp32W)\n } else if (params.precision == PrecisionMode::FP32IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp32IOFp16W)\n } else if (params.precision == PrecisionMode::FP32IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp32IOBf16W)\n } else {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp32IOFp32W)\n }\n\n // Make sure it launched ok.\n CHECK_CUDA(cudaGetLastError());\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ void group_norm_nhwc_bwd_scale_kernel(Group_norm_nhwc_bwd_params params) {\n // The IO traits.\n using Traits = Traits_;\n // The IO traits.\n using IOTraits = typename Traits::IOTraits;\n // The Weights traits.\n using WTraits = typename Traits::WTraits;\n\n // The IO type\n using IOType = typename IOTraits::Type;\n // The IO doubled type\n using IOType2 = typename IOTraits::Type2;\n\n // Weights type\n using WType = typename WTraits::Type;\n // Weights doubled type\n using WType2 = typename WTraits::Type2;\n\n // The instance in the batch.\n int ni = blockIdx.z;\n // The channel loaded by that thread (2 channels per thread for F16x2).\n int ci = blockIdx.x * params.channels_per_block + threadIdx.x * 2;\n // The group that thread works on and the channel in the group (modulus).\n int gi = ci / params.channels_per_group;\n\n // Load the gradients for the group.\n float mean_1 = 0.f, mean_2 = 0.f;\n if (gi < params.groups) {\n mean_1 = params.zeroed_red_buffer[(2 * ni + 0) * params.groups + gi];\n mean_2 = params.zeroed_red_buffer[(2 * ni + 1) * params.groups + gi];\n }\n\n // The sums from the fwd pass.\n float2 fwd = params.sums[ni * params.groups + gi];\n // The mean of X (computed during the fwd pass -- one value per batch*group).\n float x_mean = fwd.x;\n // The mean of squares of X (computed during the fwd pass -- one value per batch*group).\n float x_sq_mean = fwd.y;\n // The variance.\n float x_var = x_sq_mean - x_mean * x_mean;\n // The reciprocal of the standard deviation (i.e. 1.f / sqrt(var + epsilon)).\n float rcp_x_stddev = x_var <= 0.f ? 1.f : 1.f / sqrtf(x_var + params.epsilon);\n\n // Mutiply by 1/(HWC) to get real mean\n mean_1 *= params.inv_hwc_per_group;\n mean_2 *= params.inv_hwc_per_group;\n\n // Load gamma.\n float2 gamma_f2 = make_float2(0.f, 0.f);\n float2 beta_f2 = make_float2(0.f, 0.f);\n if (ci < params.c) {\n gamma_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.gamma)[ci]));\n if (params.with_swish) {\n beta_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.beta)[ci]));\n }\n }\n\n // The first activation loaded by that block.\n int hw_begin = blockIdx.y * params.acts_per_block;\n // The last activation loaded by that block.\n int hw_end = min((int64_t)hw_begin + params.acts_per_block, params.hw);\n\n // Iterate over the activations to compute the sums.\n for (int hwi = hw_begin; hwi < hw_end; ++hwi) {\n // The src/dst offset.\n int64_t offset = (int64_t)ni * params.hwc + hwi * params.c + ci;\n\n // Fetch two channels per thread.\n IOType2 x_v2 = IOTraits::zero();\n IOType2 dy_v2 = IOTraits::zero();\n if (ci < params.c) {\n x_v2 = *reinterpret_cast(&reinterpret_cast(params.x)[offset]);\n dy_v2 = *reinterpret_cast(&reinterpret_cast(params.dy)[offset]);\n }\n\n // Extract the two half values.\n float2 x_f2 = IOTraits::unpack(x_v2);\n float2 dy_f2 = IOTraits::unpack(dy_v2);\n\n // X - X_mean.\n float2 x_minus_x_mean_f2;\n x_minus_x_mean_f2.x = x_f2.x - x_mean;\n x_minus_x_mean_f2.y = x_f2.y - x_mean;\n\n // Normalize X.\n float2 x_norm_f2;\n x_norm_f2.x = x_minus_x_mean_f2.x * rcp_x_stddev;\n x_norm_f2.y = x_minus_x_mean_f2.y * rcp_x_stddev;\n\n if (params.with_swish) {\n float x_gn_x = x_norm_f2.x * gamma_f2.x + beta_f2.x;\n float x_gn_y = x_norm_f2.y * gamma_f2.y + beta_f2.y;\n float s_x = sigmoid(x_gn_x);\n float s_y = sigmoid(x_gn_y);\n dy_f2.x = dy_f2.x * s_x * (1.f + x_gn_x * (1.f - s_x));\n dy_f2.y = dy_f2.y * s_y * (1.f + x_gn_y * (1.f - s_y));\n }\n\n // The gradient that enters the x_norm node.\n float2 dx_norm;\n dx_norm.x = dy_f2.x * gamma_f2.x;\n dx_norm.y = dy_f2.y * gamma_f2.y;\n\n // The gradient along the input path.\n float2 dx;\n dx.x = (dx_norm.x - (x_norm_f2.x * mean_1 + mean_2)) * rcp_x_stddev;\n dx.y = (dx_norm.y - (x_norm_f2.y * mean_1 + mean_2)) * rcp_x_stddev;\n\n // Store the scaled values.\n if (ci < params.c) {\n *reinterpret_cast(&reinterpret_cast(params.dx)[offset]) = IOTraits::pack(dx);\n }\n }\n\n // Load gamma/beta and convert to half.\n if (blockIdx.y > 0 || blockIdx.z > 0 || ci >= params.c) {\n return;\n }\n\n // The base pointer for the gradients for gamma and beta.\n float* dgamma_beta_ptr = ¶ms.zeroed_red_buffer[params.n * params.groups * 2];\n\n // The 1st channel in the output tensor. NOTE: Two channels per thread store interleaved.\n int ck = blockIdx.x * params.channels_per_block + threadIdx.x;\n\n // Load the FP32 version of dgamma and dbeta.\n float2 dgamma, dbeta;\n if (ck < params.c) {\n dgamma.x = dgamma_beta_ptr[0 * params.c + 0 * blockDim.x + ck];\n dgamma.y = dgamma_beta_ptr[0 * params.c + 1 * blockDim.x + ck];\n dbeta.x = dgamma_beta_ptr[1 * params.c + 0 * blockDim.x + ck];\n dbeta.y = dgamma_beta_ptr[1 * params.c + 1 * blockDim.x + ck];\n\n // Convert to half2 and store to memory.\n *reinterpret_cast(&reinterpret_cast(params.dgamma)[ci]) = WTraits::pack(dgamma);\n *reinterpret_cast(&reinterpret_cast(params.dbeta)[ci]) = WTraits::pack(dbeta);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_bwd_two_passes_scale(const Group_norm_nhwc_bwd_params& params, cudaStream_t stream) {\n // The dimension of the grid.\n dim3 grid;\n\n // The number of blocks to compute all the channels.\n grid.x = params.c / params.channels_per_block;\n // The number of blocks to compute all the activations in a given instance.\n grid.y = div_up(params.hw, params.acts_per_block);\n // The number of instances.\n grid.z = params.n;\n\n if (params.precision == PrecisionMode::FP16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp16IOFp16W)\n } else if (params.precision == PrecisionMode::FP16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp16IOBf16W)\n } else if (params.precision == PrecisionMode::FP16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp16IOFp32W)\n } else if (params.precision == PrecisionMode::BF16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Bf16IOFp16W)\n } else if (params.precision == PrecisionMode::BF16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Bf16IOBf16W)\n } else if (params.precision == PrecisionMode::BF16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Bf16IOFp32W)\n } else if (params.precision == PrecisionMode::FP32IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp32IOFp16W)\n } else if (params.precision == PrecisionMode::FP32IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp32IOBf16W)\n } else {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp32IOFp32W)\n }\n\n // Make sure it launched ok.\n CHECK_CUDA(cudaGetLastError());\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n", "output": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n#include \"hip/hip_runtime.h\"\n#include \n\n#include \n\n#include \"group_norm_nhwc.h\"\n#include \"macros.h\"\n#include \"traits.h\"\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// B A C K W A R D\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ void group_norm_nhwc_bwd_sum_kernel(Group_norm_nhwc_bwd_params params) {\n // The IO traits.\n using Traits = Traits_;\n // The IO traits.\n using IOTraits = typename Traits::IOTraits;\n // The Weights traits.\n using WTraits = typename Traits::WTraits;\n\n // The IO type\n using IOType = typename IOTraits::Type;\n // The IO doubled type\n using IOType2 = typename IOTraits::Type2;\n\n // Weights type\n using WType = typename WTraits::Type;\n // Weights doubled type\n using WType2 = typename WTraits::Type2;\n\n // The object in charge of doing the sums for the different blocks.\n typedef hipcub::BlockScan Block_scan;\n\n // Allocate shared memory for Block_scan.\n __shared__ typename Block_scan::TempStorage temp_storage;\n // Allocate shared memory for the groups. We could reduce the amount of shared memory reserved.\n __shared__ float2 smem[THREADS_PER_BLOCK];\n\n // The instance in the batch.\n int ni = blockIdx.z;\n // The channel loaded by that thread (2 channels per thread for F16x2).\n int ci = blockIdx.x * params.channels_per_block + threadIdx.x * 2;\n // The group that thread works on and the channel in the group (modulus).\n int gi = ci / params.channels_per_group;\n\n // The sums from the fwd pass.\n float2 fwd = params.sums[ni * params.groups + gi];\n // The mean of X (computed during the fwd pass -- one value per batch*group).\n float x_mean = fwd.x;\n // The mean of squares of X (computed during the fwd pass -- one value per batch*group).\n float x_sq_mean = fwd.y;\n // The variance.\n float x_var = x_sq_mean - x_mean * x_mean;\n // The reciprocal of the standard deviation (i.e. 1.f / sqrt(var + epsilon)).\n float rcp_x_stddev = x_var <= 0.f ? 1.f : 1.f / sqrtf(x_var + params.epsilon);\n\n // Load gamma.\n float2 gamma_f2 = make_float2(0.f, 0.f);\n float2 beta_f2 = make_float2(0.f, 0.f);\n if (ci < params.c) {\n gamma_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.gamma)[ci]));\n if (params.with_swish) {\n beta_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.beta)[ci]));\n }\n }\n\n // The group that thread works on and the channel in the group (modulus).\n int gj = threadIdx.x * 2 / params.channels_per_group;\n int cj = threadIdx.x * 2 - params.channels_per_group * gj;\n\n // The first activation loaded by that block.\n int hw_begin = blockIdx.y * params.acts_per_block;\n // The last activation loaded by that block.\n int hw_end = min((int64_t)hw_begin + params.acts_per_block, params.hw);\n\n // The gradients for gamma/beta.\n float2 dgamma = make_float2(0.f, 0.f), dbeta = make_float2(0.f, 0.f);\n // Accumulated gradients for dgrad calculation\n float mean_1 = 0.f, mean_2 = 0.f;\n\n // Iterate over the activations to compute the sums.\n for (int hwi = hw_begin; hwi < hw_end; ++hwi) {\n // The offset.\n int64_t offset = (int64_t)ni * params.hwc + hwi * params.c + ci;\n\n // Fetch two channels per thread.\n IOType2 x_v2 = IOTraits::zero();\n IOType2 dy_v2 = IOTraits::zero();\n if (ci < params.c) {\n x_v2 = *reinterpret_cast(&reinterpret_cast(params.x)[offset]);\n dy_v2 = *reinterpret_cast(&reinterpret_cast(params.dy)[offset]);\n }\n\n // Extract the two half values.\n float2 x_f2 = IOTraits::unpack(x_v2);\n float2 dy_f2 = IOTraits::unpack(dy_v2);\n\n // X - X_mean.\n float x_minus_x_mean_x = x_f2.x - x_mean;\n float x_minus_x_mean_y = x_f2.y - x_mean;\n\n // Normalize X.\n float x_norm_x = x_minus_x_mean_x * rcp_x_stddev;\n float x_norm_y = x_minus_x_mean_y * rcp_x_stddev;\n\n if (params.with_swish) {\n float x_gn_x = x_norm_x * gamma_f2.x + beta_f2.x;\n float x_gn_y = x_norm_y * gamma_f2.y + beta_f2.y;\n float s_x = sigmoid(x_gn_x);\n float s_y = sigmoid(x_gn_y);\n dy_f2.x = dy_f2.x * s_x * (1.f + x_gn_x * (1.f - s_x));\n dy_f2.y = dy_f2.y * s_y * (1.f + x_gn_y * (1.f - s_y));\n }\n\n // Update beta.\n dbeta.x += dy_f2.x;\n dbeta.y += dy_f2.y;\n\n // Update dgamma.\n dgamma.x += dy_f2.x * x_norm_x;\n dgamma.y += dy_f2.y * x_norm_y;\n\n // The gradient that enters the x_norm node.\n float dx_norm_x = dy_f2.x * gamma_f2.x;\n float dx_norm_y = dy_f2.y * gamma_f2.y;\n\n // Add to the 1st mean.\n mean_1 += dx_norm_x * x_norm_x;\n mean_1 += dx_norm_y * x_norm_y;\n\n // Add to the 2nd mean.\n mean_2 += dx_norm_x;\n mean_2 += dx_norm_y;\n }\n\n // The data for the summations.\n Group_sums inp{cj == 0 ? 1 : 0, mean_1, mean_2};\n\n // Do the segmented scan.\n Group_sums out;\n Block_scan(temp_storage).InclusiveScan(inp, out, Group_sums_op());\n\n // Store the results for the groups in shared memory (to produce coalesced stores later).\n if (cj == params.channels_per_group - 2 /* 2 channels per thread */) {\n smem[gj] = make_float2(out.sum, out.sum_sq);\n }\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n // The global group index.\n int gk = blockIdx.x * params.groups_per_block + threadIdx.x;\n\n // The first threads (those storing to global memory, load the values).\n float2 sums = smem[threadIdx.x];\n\n // Store to global memory.\n if (threadIdx.x < params.groups_per_block && gk < params.groups) {\n atomicAdd(¶ms.zeroed_red_buffer[(2 * ni + 0) * params.groups + gk], sums.x);\n atomicAdd(¶ms.zeroed_red_buffer[(2 * ni + 1) * params.groups + gk], sums.y);\n }\n\n // The base pointer for the gradients for gamma and beta.\n float* dgamma_beta_ptr = ¶ms.zeroed_red_buffer[params.n * params.groups * 2];\n\n // The 1st channel in the output tensor. NOTE: Two channels per thread store interleaved.\n int ck = blockIdx.x * params.channels_per_block + threadIdx.x;\n\n // Store dgamma and dbeta as well.\n if (ck < params.c) {\n atomicAdd(&dgamma_beta_ptr[0 * params.c + 0 * blockDim.x + ck], dgamma.x);\n atomicAdd(&dgamma_beta_ptr[0 * params.c + 1 * blockDim.x + ck], dgamma.y);\n atomicAdd(&dgamma_beta_ptr[1 * params.c + 0 * blockDim.x + ck], dbeta.x);\n atomicAdd(&dgamma_beta_ptr[1 * params.c + 1 * blockDim.x + ck], dbeta.y);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_bwd_two_passes_setup(Group_norm_nhwc_bwd_params& params, size_t& zeroed_red_buffer_elts) {\n // The pre-computed dimensions.\n params.hw = params.h * params.w;\n params.hwc = params.c * params.hw;\n\n // The number of channels per group.\n params.channels_per_group = params.c / params.groups;\n // The inverse to compute the mean/variance.\n params.inv_hwc_per_group = 1.f / (float)(params.hw * params.channels_per_group);\n\n // Define the number of blocks per activation map. That's a simple heuristic.\n int blocks_per_act_slice = 0;\n if (params.c >= 1280) {\n blocks_per_act_slice = 128 / params.n;\n } else if (params.c >= 640) {\n blocks_per_act_slice = 256 / params.n;\n } else {\n blocks_per_act_slice = 512 / params.n;\n }\n\n // Clamp to at least 1 to avoid divide-by-zero when batch size is large.\n blocks_per_act_slice = max(blocks_per_act_slice, 1);\n\n // Make sure we launch blocks per activation is no less than activations\n blocks_per_act_slice = min(blocks_per_act_slice, div_up(params.hw, params.n));\n\n // Define how many activations are computed per block.\n params.acts_per_block = div_up(params.hw, blocks_per_act_slice);\n\n // The number of channels per block.\n params.channels_per_block = 320;\n // Special case to deal with 30 channels per group.\n if (params.channels_per_block % params.channels_per_group != 0) {\n params.channels_per_block = 240;\n }\n\n // Special case to deal with 70 channels per group.\n if (params.c == 2240) {\n params.channels_per_block = 280;\n } else if (params.c == 832) {\n params.channels_per_block = 208;\n }\n\n if (params.c % params.channels_per_block != 0) {\n if (params.c % 512 == 0 && params.c != 1536 && params.c != 3072 && params.c % 448 != 0) {\n params.channels_per_block = 512;\n } else if (params.c % 42 == 0) {\n params.channels_per_block = 336;\n } else if (params.c % 384 == 0) {\n params.channels_per_block = 384;\n } else if (params.c % 256 == 0 && params.c % 448 != 0 && params.c % 392 != 0) {\n params.channels_per_block = 256;\n } else if (params.c % 128 == 0 && params.c % 448 != 0 && params.c % 392 != 0) {\n params.channels_per_block = 128;\n } else if (params.c % 448 == 0 && params.c % 392 != 0) {\n params.channels_per_block = 448;\n } else if (params.c % 392 == 0) {\n params.channels_per_block = 392;\n }\n }\n\n // The number of groups per block.\n params.groups_per_block = params.channels_per_block / params.channels_per_group;\n\n // Make sure the number of channels is a multiple of the number of channels per block.\n assert(params.c % params.channels_per_block == 0);\n // Make sure a group does not span multiple blocks.\n assert(params.channels_per_block % params.channels_per_group == 0);\n\n // The number of elements in the reduction buffer (for the sums and sums of squared).\n zeroed_red_buffer_elts = params.n * params.groups * 2 + params.c * 2;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_bwd_two_passes_sum(const Group_norm_nhwc_bwd_params& params, hipStream_t stream) {\n // The dimension of the grid.\n dim3 grid;\n\n // The number of blocks to compute all the channels.\n grid.x = params.c / params.channels_per_block;\n // The number of blocks to compute all the activations in a given instance.\n grid.y = div_up(params.hw, params.acts_per_block);\n // The number of instances.\n grid.z = params.n;\n\n if (params.precision == PrecisionMode::FP16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp16IOFp16W)\n } else if (params.precision == PrecisionMode::FP16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp16IOBf16W)\n } else if (params.precision == PrecisionMode::FP16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp16IOFp32W)\n } else if (params.precision == PrecisionMode::BF16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Bf16IOFp16W)\n } else if (params.precision == PrecisionMode::BF16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Bf16IOBf16W)\n } else if (params.precision == PrecisionMode::BF16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Bf16IOFp32W)\n } else if (params.precision == PrecisionMode::FP32IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp32IOFp16W)\n } else if (params.precision == PrecisionMode::FP32IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp32IOBf16W)\n } else {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp32IOFp32W)\n }\n\n // Make sure it launched ok.\n CHECK_CUDA(hipGetLastError());\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ void group_norm_nhwc_bwd_scale_kernel(Group_norm_nhwc_bwd_params params) {\n // The IO traits.\n using Traits = Traits_;\n // The IO traits.\n using IOTraits = typename Traits::IOTraits;\n // The Weights traits.\n using WTraits = typename Traits::WTraits;\n\n // The IO type\n using IOType = typename IOTraits::Type;\n // The IO doubled type\n using IOType2 = typename IOTraits::Type2;\n\n // Weights type\n using WType = typename WTraits::Type;\n // Weights doubled type\n using WType2 = typename WTraits::Type2;\n\n // The instance in the batch.\n int ni = blockIdx.z;\n // The channel loaded by that thread (2 channels per thread for F16x2).\n int ci = blockIdx.x * params.channels_per_block + threadIdx.x * 2;\n // The group that thread works on and the channel in the group (modulus).\n int gi = ci / params.channels_per_group;\n\n // Load the gradients for the group.\n float mean_1 = 0.f, mean_2 = 0.f;\n if (gi < params.groups) {\n mean_1 = params.zeroed_red_buffer[(2 * ni + 0) * params.groups + gi];\n mean_2 = params.zeroed_red_buffer[(2 * ni + 1) * params.groups + gi];\n }\n\n // The sums from the fwd pass.\n float2 fwd = params.sums[ni * params.groups + gi];\n // The mean of X (computed during the fwd pass -- one value per batch*group).\n float x_mean = fwd.x;\n // The mean of squares of X (computed during the fwd pass -- one value per batch*group).\n float x_sq_mean = fwd.y;\n // The variance.\n float x_var = x_sq_mean - x_mean * x_mean;\n // The reciprocal of the standard deviation (i.e. 1.f / sqrt(var + epsilon)).\n float rcp_x_stddev = x_var <= 0.f ? 1.f : 1.f / sqrtf(x_var + params.epsilon);\n\n // Mutiply by 1/(HWC) to get real mean\n mean_1 *= params.inv_hwc_per_group;\n mean_2 *= params.inv_hwc_per_group;\n\n // Load gamma.\n float2 gamma_f2 = make_float2(0.f, 0.f);\n float2 beta_f2 = make_float2(0.f, 0.f);\n if (ci < params.c) {\n gamma_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.gamma)[ci]));\n if (params.with_swish) {\n beta_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.beta)[ci]));\n }\n }\n\n // The first activation loaded by that block.\n int hw_begin = blockIdx.y * params.acts_per_block;\n // The last activation loaded by that block.\n int hw_end = min((int64_t)hw_begin + params.acts_per_block, params.hw);\n\n // Iterate over the activations to compute the sums.\n for (int hwi = hw_begin; hwi < hw_end; ++hwi) {\n // The src/dst offset.\n int64_t offset = (int64_t)ni * params.hwc + hwi * params.c + ci;\n\n // Fetch two channels per thread.\n IOType2 x_v2 = IOTraits::zero();\n IOType2 dy_v2 = IOTraits::zero();\n if (ci < params.c) {\n x_v2 = *reinterpret_cast(&reinterpret_cast(params.x)[offset]);\n dy_v2 = *reinterpret_cast(&reinterpret_cast(params.dy)[offset]);\n }\n\n // Extract the two half values.\n float2 x_f2 = IOTraits::unpack(x_v2);\n float2 dy_f2 = IOTraits::unpack(dy_v2);\n\n // X - X_mean.\n float2 x_minus_x_mean_f2;\n x_minus_x_mean_f2.x = x_f2.x - x_mean;\n x_minus_x_mean_f2.y = x_f2.y - x_mean;\n\n // Normalize X.\n float2 x_norm_f2;\n x_norm_f2.x = x_minus_x_mean_f2.x * rcp_x_stddev;\n x_norm_f2.y = x_minus_x_mean_f2.y * rcp_x_stddev;\n\n if (params.with_swish) {\n float x_gn_x = x_norm_f2.x * gamma_f2.x + beta_f2.x;\n float x_gn_y = x_norm_f2.y * gamma_f2.y + beta_f2.y;\n float s_x = sigmoid(x_gn_x);\n float s_y = sigmoid(x_gn_y);\n dy_f2.x = dy_f2.x * s_x * (1.f + x_gn_x * (1.f - s_x));\n dy_f2.y = dy_f2.y * s_y * (1.f + x_gn_y * (1.f - s_y));\n }\n\n // The gradient that enters the x_norm node.\n float2 dx_norm;\n dx_norm.x = dy_f2.x * gamma_f2.x;\n dx_norm.y = dy_f2.y * gamma_f2.y;\n\n // The gradient along the input path.\n float2 dx;\n dx.x = (dx_norm.x - (x_norm_f2.x * mean_1 + mean_2)) * rcp_x_stddev;\n dx.y = (dx_norm.y - (x_norm_f2.y * mean_1 + mean_2)) * rcp_x_stddev;\n\n // Store the scaled values.\n if (ci < params.c) {\n *reinterpret_cast(&reinterpret_cast(params.dx)[offset]) = IOTraits::pack(dx);\n }\n }\n\n // Load gamma/beta and convert to half.\n if (blockIdx.y > 0 || blockIdx.z > 0 || ci >= params.c) {\n return;\n }\n\n // The base pointer for the gradients for gamma and beta.\n float* dgamma_beta_ptr = ¶ms.zeroed_red_buffer[params.n * params.groups * 2];\n\n // The 1st channel in the output tensor. NOTE: Two channels per thread store interleaved.\n int ck = blockIdx.x * params.channels_per_block + threadIdx.x;\n\n // Load the FP32 version of dgamma and dbeta.\n float2 dgamma, dbeta;\n if (ck < params.c) {\n dgamma.x = dgamma_beta_ptr[0 * params.c + 0 * blockDim.x + ck];\n dgamma.y = dgamma_beta_ptr[0 * params.c + 1 * blockDim.x + ck];\n dbeta.x = dgamma_beta_ptr[1 * params.c + 0 * blockDim.x + ck];\n dbeta.y = dgamma_beta_ptr[1 * params.c + 1 * blockDim.x + ck];\n\n // Convert to half2 and store to memory.\n *reinterpret_cast(&reinterpret_cast(params.dgamma)[ci]) = WTraits::pack(dgamma);\n *reinterpret_cast(&reinterpret_cast(params.dbeta)[ci]) = WTraits::pack(dbeta);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_bwd_two_passes_scale(const Group_norm_nhwc_bwd_params& params, hipStream_t stream) {\n // The dimension of the grid.\n dim3 grid;\n\n // The number of blocks to compute all the channels.\n grid.x = params.c / params.channels_per_block;\n // The number of blocks to compute all the activations in a given instance.\n grid.y = div_up(params.hw, params.acts_per_block);\n // The number of instances.\n grid.z = params.n;\n\n if (params.precision == PrecisionMode::FP16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp16IOFp16W)\n } else if (params.precision == PrecisionMode::FP16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp16IOBf16W)\n } else if (params.precision == PrecisionMode::FP16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp16IOFp32W)\n } else if (params.precision == PrecisionMode::BF16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Bf16IOFp16W)\n } else if (params.precision == PrecisionMode::BF16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Bf16IOBf16W)\n } else if (params.precision == PrecisionMode::BF16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Bf16IOFp32W)\n } else if (params.precision == PrecisionMode::FP32IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp32IOFp16W)\n } else if (params.precision == PrecisionMode::FP32IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp32IOBf16W)\n } else {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp32IOFp32W)\n }\n\n // Make sure it launched ok.\n CHECK_CUDA(hipGetLastError());\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n", "source": "/kaggle/working/apex/apex/contrib/csrc/group_norm/group_norm_nhwc_bwd_two_pass.cu"} {"input": "#include \n#include \n#include \n\n#include \n\n#include \"ATen/ATen.h\"\n#include \"ATen/TensorUtils.h\"\n#include \"ATen/cuda/CUDAContext.h\"\n#include \"ATen/cuda/detail/IndexUtils.cuh\"\n// #include \"ATen/Type.h\"\n#include \"ATen/AccumulateType.h\"\n#include \"multi_tensor_apply.cuh\"\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntemplate \n__device__ __forceinline__ bool is_aligned(T* p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\n\n#include \"type_shim.h\"\n\ntypedef enum {\n ADAM_MODE_0 = 0, // eps under square root\n ADAM_MODE_1 = 1 // eps outside square root\n} adamMode_t;\n\ntemplate \n__global__ void adam_cuda_kernel(T* __restrict__ p,\n GRAD_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed\n T* __restrict__ m, T* __restrict__ v, const GRAD_T* __restrict__ g, const float b1,\n const float b2, const float eps, const float grad_scale, const float step_size,\n const size_t tsize, adamMode_t mode, const float decay) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n for (int j = i; j < tsize; j += totThreads) {\n T scaled_grad = g[j] / grad_scale;\n m[j] = b1 * m[j] + (1 - b1) * scaled_grad;\n v[j] = b2 * v[j] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(v[j] + eps);\n else // Mode 1\n denom = sqrtf(v[j]) + eps;\n float update = (m[j] / denom) + (decay * p[j]);\n p[j] = p[j] - (step_size * update);\n if (p_copy != NULL) p_copy[j] = (GRAD_T)p[j];\n }\n}\n\ntemplate \nstruct AdamFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata& tl,\n const float b1, const float b2, const float eps, const float grad_scale,\n const float step_size, adamMode_t mode, const float decay) {\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n T* p = (T*)tl.addresses[0][tensor_loc];\n p += chunk_idx * chunk_size;\n T* m = (T*)tl.addresses[1][tensor_loc];\n m += chunk_idx * chunk_size;\n T* v = (T*)tl.addresses[2][tensor_loc];\n v += chunk_idx * chunk_size;\n GRAD_T* g = (GRAD_T*)tl.addresses[3][tensor_loc];\n g += chunk_idx * chunk_size;\n GRAD_T* p_copy = NULL;\n if (DEPTH == 5) {\n p_copy = (GRAD_T*)tl.addresses[4][tensor_loc];\n p_copy += chunk_idx * chunk_size;\n }\n\n n -= chunk_idx * chunk_size;\n\n T incoming_p[ILP];\n T incoming_m[ILP];\n T incoming_v[ILP];\n T incoming_g[ILP];\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(p) && is_aligned(m) && is_aligned(v) && is_aligned(g) &&\n is_aligned(p_copy)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n GRAD_T tmp_g[ILP];\n load_store(incoming_p, p, 0, i_start);\n load_store(incoming_m, m, 0, i_start);\n load_store(incoming_v, v, 0, i_start);\n load_store(tmp_g, g, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n incoming_g[ii] = static_cast(tmp_g[ii]);\n T scaled_grad = incoming_g[ii] / grad_scale;\n incoming_m[ii] = b1 * incoming_m[ii] + (1 - b1) * scaled_grad;\n incoming_v[ii] = b2 * incoming_v[ii] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(incoming_v[ii] + eps);\n else // Mode 1\n denom = sqrtf(incoming_v[ii]) + eps;\n float update = (incoming_m[ii] / denom) + (decay * incoming_p[ii]);\n incoming_p[ii] = incoming_p[ii] - (step_size * update);\n if (DEPTH == 5) tmp_g[ii] = static_cast(incoming_p[ii]);\n }\n load_store(p, incoming_p, i_start, 0);\n load_store(m, incoming_m, i_start, 0);\n load_store(v, incoming_v, i_start, 0);\n if (DEPTH == 5) load_store(p_copy, tmp_g, i_start, 0);\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n incoming_p[ii] = 0;\n incoming_m[ii] = 0;\n incoming_v[ii] = 0;\n incoming_g[ii] = 0;\n\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n incoming_p[ii] = p[i];\n incoming_m[ii] = m[i];\n incoming_v[ii] = v[i];\n incoming_g[ii] = static_cast(g[i]);\n }\n }\n\n // note for clarification to future michael:\n // From a pure memory dependency perspective, there's likely no point unrolling\n // the write loop, since writes just fire off once their LDGs arrive.\n // Put another way, the STGs are dependent on the LDGs, but not on each other.\n // There is still compute ILP benefit from unrolling the loop though.\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = i_start + threadIdx.x + ii * blockDim.x;\n\n if (j < n && j < chunk_size) {\n T scaled_grad = incoming_g[ii] / grad_scale;\n m[j] = b1 * incoming_m[ii] + (1 - b1) * scaled_grad;\n v[j] = b2 * incoming_v[ii] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(v[j] + eps);\n else // Mode 1\n denom = sqrtf(v[j]) + eps;\n float update = (m[j] / denom) + (decay * incoming_p[ii]);\n p[j] = incoming_p[ii] - (step_size * update);\n if (DEPTH == 5) p_copy[j] = (GRAD_T)p[j];\n }\n }\n }\n }\n }\n};\n\nvoid fused_adam_cuda(at::Tensor& p, at::Tensor& p_copy, at::Tensor& m, at::Tensor& v, at::Tensor& g, float lr,\n float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction,\n float decay) {\n // using namespace at;\n\n // Get tensor size\n int tsize = p.numel();\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (g.scalar_type() == at::ScalarType::Half) {\n // all other values should be fp32 for half gradients\n TORCH_CHECK(p.scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dispatch is done on the gradient type\n using namespace at; // prevents \"toString is undefined\" errors\n DISPATCH_FLOAT_AND_HALF(g.scalar_type(), 0, \"adam_cuda_kernel\", using accscalar_t = at::acc_type;\n adam_cuda_kernel<<>>(\n p.data_ptr(), p_copy.numel() ? p_copy.data_ptr() : NULL,\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n } else {\n using namespace at;\n DISPATCH_DOUBLE_AND_FLOAT(g.scalar_type(), 0, \"adam_cuda_kernel\",\n adam_cuda_kernel<<>>(\n p.data_ptr(),\n NULL, // don't output p_copy for fp32, it's wasted write\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n }\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid fused_adam_cuda_mt(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists, // p, m, v, g, p_copy\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay) {\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n size_t tl_sz = tensor_lists.size();\n TORCH_CHECK(tl_sz == 4 || tl_sz == 5, \"expected tensor lists of size 4 or 5\");\n\n if (tensor_lists[3][0].scalar_type() == at::ScalarType::Half) {\n // alher values should be fp32 for half gradients\n TORCH_CHECK(tensor_lists[0][0].scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dich is done on the gradient type\n if (tl_sz == 5) {\n DISPATCH_FLOAT_AND_HALF(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n using accscalar_t = at::acc_type;\n multi_tensor_apply<5>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<5, accscalar_t, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n } else {\n DISPATCH_FLOAT_AND_HALF(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n using accscalar_t = at::acc_type;\n multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<4, accscalar_t, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n }\n } else {\n if (tl_sz == 5) {\n DISPATCH_DOUBLE_AND_FLOAT(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n multi_tensor_apply<5>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<5, scalar_t_0, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n } else {\n DISPATCH_DOUBLE_AND_FLOAT(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<4, scalar_t_0, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n }\n }\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\ntemplate \n__device__ void convert(const FROM_T vi, TO_T& vo) {\n vo = static_cast(vi);\n}\n\ntemplate <>\n__device__ void convert(const float vi, uint8_t& vo) {\n union S {\n float as_float;\n int as_int;\n };\n S s;\n s.as_float = vi;\n s.as_int = s.as_int & 0xFF800000;\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_half = static_cast(vi + s.as_float / 8.0f);\n vo = t.as_byte[1];\n}\n\ntemplate <>\n__device__ void convert(const uint8_t vi, float& vo) {\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_byte[0] = 0;\n t.as_byte[1] = vi;\n vo = static_cast(t.as_half);\n}\n\ntemplate <>\n__device__ void convert(const at::Half vi, uint8_t& vo) {\n union S {\n float as_float;\n int as_int;\n };\n S s;\n s.as_float = static_cast(vi);\n s.as_int = s.as_int & 0xFF800000;\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_half = static_cast(vi + s.as_float / 8.0f);\n vo = t.as_byte[1];\n}\n\ntemplate <>\n__device__ void convert(const uint8_t vi, at::Half& vo) {\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_byte[0] = 0;\n t.as_byte[1] = vi;\n vo = t.as_half;\n}\n\ntemplate \n__global__ void strided_check_finite_cuda_kernel(volatile int* noop_gmem, GRAD_T* __restrict__ p_copy,\n const size_t tsize, int stride, int clear_overflow_first) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock) * stride;\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock * stride;\n\n if (clear_overflow_first) {\n if (i == 0) {\n *noop_gmem = 0;\n }\n __syncthreads();\n }\n\n for (int j = i; j < tsize; j += totThreads) {\n GRAD_T pi = p_copy[j];\n if (!isfinite(pi)) {\n *noop_gmem = 1;\n }\n }\n}\ntemplate <>\n__global__ void strided_check_finite_cuda_kernel(volatile int* noop_gmem, uint8_t* __restrict__ p_copy,\n const size_t tsize, int stride, int clear_overflow_first) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock) * stride;\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock * stride;\n\n if (clear_overflow_first) {\n if (i == 0) {\n *noop_gmem = 0;\n }\n __syncthreads();\n }\n\n for (int j = i; j < tsize; j += totThreads) {\n at::Half pi;\n convert(p_copy[j], pi);\n if (!isfinite(pi)) {\n *noop_gmem = 1;\n }\n }\n}\n\ntemplate \n__global__ void maybe_cast_kernel(volatile int* overflow_flag, const FROM_T* p_in, TO_T* p_out, const size_t tsize) {\n if (overflow_flag && *overflow_flag != 0) return;\n\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n FROM_T pi[ILP];\n TO_T po[ILP];\n\n for (int j_start = 0; j_start < tsize; j_start += totThreads * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n pi[ii] = 0;\n\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n pi[ii] = p_in[j];\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n convert(pi[ii], po[ii]);\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n p_out[j] = po[ii];\n }\n }\n }\n}\n\ntemplate \n__global__ void reversible_adam_cuda_kernel(\n T* __restrict__ p,\n REDU_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed\n T* __restrict__ m, T* __restrict__ v, const GRAD_T* __restrict__ g, const float b1, const float b2, const float eps,\n const float grad_scale, const float step_size, const size_t tsize, adamMode_t mode, const float decay) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n T mi[ILP];\n T vi[ILP];\n T pi[ILP];\n T gi[ILP];\n\n bool overflow = false;\n for (int j_start = 0; j_start < tsize; j_start += totThreads * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n mi[ii] = T(0);\n vi[ii] = T(0);\n pi[ii] = T(0);\n gi[ii] = GRAD_T(0);\n\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n pi[ii] = p[j];\n mi[ii] = m[j];\n vi[ii] = v[j];\n gi[ii] = static_cast(g[j]);\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n T scaled_grad = gi[ii] / grad_scale;\n if (isfinite(scaled_grad)) {\n mi[ii] = b1 * mi[ii] + (1 - b1) * scaled_grad;\n vi[ii] = b2 * vi[ii] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(vi[ii] + eps);\n else // Mode 1\n denom = sqrtf(vi[ii]) + eps;\n float update = (mi[ii] / denom) + (decay * pi[ii]);\n pi[ii] = pi[ii] - (step_size * update);\n } else {\n overflow = true;\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n m[j] = mi[ii];\n v[j] = vi[ii];\n p[j] = pi[ii];\n if (p_copy != NULL) {\n convert(pi[ii], p_copy[j]);\n }\n }\n }\n }\n\n if (p_copy != NULL) {\n __syncthreads();\n if (overflow) {\n convert(float(INFINITY), p_copy[0]);\n }\n }\n}\n\ntemplate \n__global__ void maybe_adam_undo_cuda_kernel(volatile int* overflow_flag, T* __restrict__ p, T* __restrict__ m,\n T* __restrict__ v, const GRAD_T* __restrict__ g, const float b1,\n const float b2, const float eps, const float grad_scale,\n const float step_size, const size_t tsize, adamMode_t mode,\n const float decay) {\n // NB! Skip undo kernel when overflow flag is NOT set\n if (overflow_flag && *overflow_flag == 0) return;\n\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n T mi[ILP];\n T vi[ILP];\n T pi[ILP];\n T gi[ILP];\n\n for (int j_start = 0; j_start < tsize; j_start += totThreads * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n mi[ii] = T(0);\n vi[ii] = T(0);\n pi[ii] = T(0);\n gi[ii] = GRAD_T(0);\n\n int j = j_start + i * ILP;\n if (j < tsize) {\n pi[ii] = p[j];\n mi[ii] = m[j];\n vi[ii] = v[j];\n gi[ii] = static_cast(g[j]);\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n T scaled_grad = gi[ii] / grad_scale;\n if (isfinite(scaled_grad)) {\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(vi[ii] + eps);\n else // Mode 1\n denom = sqrtf(vi[ii]) + eps;\n pi[ii] = (pi[ii] + step_size * (mi[ii] / denom)) / (1.0f - step_size * decay);\n mi[ii] = (mi[ii] - (1 - b1) * scaled_grad) / b1;\n vi[ii] = (vi[ii] - (1 - b2) * scaled_grad * scaled_grad) / b2;\n // Make sure round off errors don't create (small) negative value.\n // This can happen if we have to revert the very first step.\n vi[ii] = vi[ii] >= 0.0f ? vi[ii] : 0.0f;\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + i * ILP;\n if (j < tsize) {\n m[j] = mi[ii];\n v[j] = vi[ii];\n p[j] = pi[ii];\n }\n }\n }\n}\n\ntemplate \nstruct MaybeCastFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* overflow_flag,\n TensorListMetadata& tl) {\n if (overflow_flag && *overflow_flag != 0) return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n FROM_T* p_in = (FROM_T*)tl.addresses[0][tensor_loc];\n p_in += chunk_idx * chunk_size;\n TO_T* p_out = (TO_T*)tl.addresses[1][tensor_loc];\n p_out += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n int dim = chunk_size < n ? chunk_size : n;\n\n FROM_T pi[ILP];\n TO_T po[ILP];\n\n for (int j_start = 0; j_start < dim; j_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n pi[ii] = FROM_T(0);\n int j = j_start + threadIdx.x + ii * blockDim.x;\n if (j < dim) {\n pi[ii] = p_in[j];\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n convert(pi[ii], po[ii]);\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + threadIdx.x + ii * blockDim.x;\n if (j < dim) {\n p_out[j] = po[ii];\n }\n }\n }\n }\n};\n\nvoid fused_strided_check_finite(at::Tensor& overflow_flag, at::Tensor& p_copy, int stride, int clear_overflow_first) {\n // Get tensor size\n int tsize = p_copy.numel();\n int niter = (tsize + stride - 1) / stride;\n\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n // In order to avoid race condition, blocks must be 1 when clear_overflow_first flag is set.\n const dim3 blocks(clear_overflow_first ? 1 : (niter + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p_copy), \"parameter tensor is too large to be indexed with int32\");\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n using namespace at; // prevents \"toString is undefined\" errors\n DISPATCH_FLOAT_HALF_AND_BYTE(\n p_copy.scalar_type(), 0, \"check_finite_cuda_kernel\",\n strided_check_finite_cuda_kernel<<>>(\n overflow_flag.data_ptr(), p_copy.data_ptr(), tsize, stride, clear_overflow_first););\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid fused_reversible_adam_cuda(at::Tensor& p, at::Tensor& p_copy, at::Tensor& m, at::Tensor& v, at::Tensor& g,\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay) {\n // using namespace at;\n\n // Get tensor size\n int tsize = p.numel();\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (g.scalar_type() == at::ScalarType::Half) {\n // all other values should be fp32 for half gradients\n TORCH_CHECK(p.scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dispatch is done on the gradient type\n using namespace at; // prevents \"toString is undefined\" errors\n if (p_copy.numel() == 0 || p_copy.scalar_type() == g.scalar_type()) {\n DISPATCH_FLOAT_AND_HALF(\n g.scalar_type(), 0, \"adam_cuda_kernel\", using accscalar_t = at::acc_type;\n reversible_adam_cuda_kernel<<>>(\n p.data_ptr(), p_copy.numel() ? p_copy.data_ptr() : NULL,\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1, beta2, eps,\n grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n } else {\n TORCH_CHECK(p_copy.scalar_type() == at::ScalarType::Byte, \"expected parameter to be of byte type\");\n DISPATCH_FLOAT_AND_HALF(\n g.scalar_type(), 0, \"adam_cuda_e5m2_kernel\", using accscalar_t = at::acc_type;\n reversible_adam_cuda_kernel<<>>(\n p.data_ptr(), p_copy.data_ptr(), m.data_ptr(),\n v.data_ptr(), g.data_ptr(), beta1, beta2, eps, grad_scale, step_size, tsize,\n (adamMode_t)mode, decay););\n }\n } else {\n using namespace at;\n DISPATCH_DOUBLE_AND_FLOAT(g.scalar_type(), 0, \"adam_cuda_kernel\",\n reversible_adam_cuda_kernel\n <<>>(\n p.data_ptr(),\n NULL, // don't output p_copy for fp32, it's wasted write\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n }\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid maybe_cast_cuda(at::Tensor& overflow_flag, at::Tensor& p_in, at::Tensor& p_out) {\n // Get tensor size\n int tsize = p_in.numel();\n TORCH_CHECK(tsize == p_out.numel(), \"p_in.numel() must equal p_out.numel()\");\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p_in), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n DISPATCH_FLOAT_HALF_AND_BYTE(p_in.scalar_type(), 0,\n \"maybe_cast_cuda\" DISPATCH_FLOAT_HALF_AND_BYTE(\n p_out.scalar_type(), 1, \"maybe_cast_cuda\",\n maybe_cast_kernel<<>>(\n overflow_flag.numel() ? overflow_flag.data_ptr() : NULL,\n p_in.data_ptr(), p_out.data_ptr(), tsize);))\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid maybe_cast_cuda_mt(int chunk_size, at::Tensor overflow_flag,\n std::vector> tensor_lists) // p_in, p_out\n{\n // Constants\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n size_t tl_sz = tensor_lists.size();\n TORCH_CHECK(tl_sz == 2, \"expected tensor lists of size 2\");\n\n DISPATCH_FLOAT_HALF_AND_BYTE(\n tensor_lists[0][0].scalar_type(), 0, \"maybe_cast_cuda_mt_kernel\",\n DISPATCH_FLOAT_HALF_AND_BYTE(tensor_lists[1][0].scalar_type(), 1, \"maybe_cast_cuda_mt_kernel\",\n multi_tensor_apply<2>(BLOCK_SIZE, chunk_size, overflow_flag, tensor_lists,\n MaybeCastFunctor<2, scalar_t_0, scalar_t_1>());))\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid fused_maybe_adam_undo_cuda(at::Tensor& overflow_flag, at::Tensor& p, at::Tensor& m, at::Tensor& v, at::Tensor& g,\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay) {\n // Get tensor size\n int tsize = p.numel();\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (g.scalar_type() == at::ScalarType::Half) {\n // all other values should be fp32 for half gradients\n TORCH_CHECK(p.scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dispatch is done on the gradient type\n using namespace at; // prevents \"toString is undefined\" errors\n DISPATCH_FLOAT_AND_HALF(g.scalar_type(), 0, \"adam_cuda_kernel\", using accscalar_t = at::acc_type;\n maybe_adam_undo_cuda_kernel\n <<>>(\n overflow_flag.numel() ? overflow_flag.data_ptr() : NULL, p.data_ptr(),\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n } else {\n using namespace at;\n DISPATCH_DOUBLE_AND_FLOAT(\n g.scalar_type(), 0, \"adam_cuda_kernel\",\n maybe_adam_undo_cuda_kernel<<>>(\n overflow_flag.numel() ? overflow_flag.data_ptr() : NULL, p.data_ptr(),\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1, beta2, eps, grad_scale,\n step_size, tsize, (adamMode_t)mode, decay););\n }\n C10_CUDA_CHECK(cudaGetLastError());\n}\n", "output": "#include \n#include \n\n#include \n\n#include \"ATen/ATen.h\"\n#include \"ATen/TensorUtils.h\"\n#include \"ATen/cuda/CUDAContext.h\"\n#include \"ATen/cuda/detail/IndexUtils.cuh\"\n// #include \"ATen/Type.h\"\n#include \"ATen/AccumulateType.h\"\n#include \"multi_tensor_apply.cuh\"\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntemplate \n__device__ __forceinline__ bool is_aligned(T* p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\n\n#include \"type_shim.h\"\n\ntypedef enum {\n ADAM_MODE_0 = 0, // eps under square root\n ADAM_MODE_1 = 1 // eps outside square root\n} adamMode_t;\n\ntemplate \n__global__ void adam_cuda_kernel(T* __restrict__ p,\n GRAD_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed\n T* __restrict__ m, T* __restrict__ v, const GRAD_T* __restrict__ g, const float b1,\n const float b2, const float eps, const float grad_scale, const float step_size,\n const size_t tsize, adamMode_t mode, const float decay) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n for (int j = i; j < tsize; j += totThreads) {\n T scaled_grad = g[j] / grad_scale;\n m[j] = b1 * m[j] + (1 - b1) * scaled_grad;\n v[j] = b2 * v[j] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(v[j] + eps);\n else // Mode 1\n denom = sqrtf(v[j]) + eps;\n float update = (m[j] / denom) + (decay * p[j]);\n p[j] = p[j] - (step_size * update);\n if (p_copy != NULL) p_copy[j] = (GRAD_T)p[j];\n }\n}\n\ntemplate \nstruct AdamFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata& tl,\n const float b1, const float b2, const float eps, const float grad_scale,\n const float step_size, adamMode_t mode, const float decay) {\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n T* p = (T*)tl.addresses[0][tensor_loc];\n p += chunk_idx * chunk_size;\n T* m = (T*)tl.addresses[1][tensor_loc];\n m += chunk_idx * chunk_size;\n T* v = (T*)tl.addresses[2][tensor_loc];\n v += chunk_idx * chunk_size;\n GRAD_T* g = (GRAD_T*)tl.addresses[3][tensor_loc];\n g += chunk_idx * chunk_size;\n GRAD_T* p_copy = NULL;\n if (DEPTH == 5) {\n p_copy = (GRAD_T*)tl.addresses[4][tensor_loc];\n p_copy += chunk_idx * chunk_size;\n }\n\n n -= chunk_idx * chunk_size;\n\n T incoming_p[ILP];\n T incoming_m[ILP];\n T incoming_v[ILP];\n T incoming_g[ILP];\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(p) && is_aligned(m) && is_aligned(v) && is_aligned(g) &&\n is_aligned(p_copy)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n GRAD_T tmp_g[ILP];\n load_store(incoming_p, p, 0, i_start);\n load_store(incoming_m, m, 0, i_start);\n load_store(incoming_v, v, 0, i_start);\n load_store(tmp_g, g, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n incoming_g[ii] = static_cast(tmp_g[ii]);\n T scaled_grad = incoming_g[ii] / grad_scale;\n incoming_m[ii] = b1 * incoming_m[ii] + (1 - b1) * scaled_grad;\n incoming_v[ii] = b2 * incoming_v[ii] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(incoming_v[ii] + eps);\n else // Mode 1\n denom = sqrtf(incoming_v[ii]) + eps;\n float update = (incoming_m[ii] / denom) + (decay * incoming_p[ii]);\n incoming_p[ii] = incoming_p[ii] - (step_size * update);\n if (DEPTH == 5) tmp_g[ii] = static_cast(incoming_p[ii]);\n }\n load_store(p, incoming_p, i_start, 0);\n load_store(m, incoming_m, i_start, 0);\n load_store(v, incoming_v, i_start, 0);\n if (DEPTH == 5) load_store(p_copy, tmp_g, i_start, 0);\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n incoming_p[ii] = 0;\n incoming_m[ii] = 0;\n incoming_v[ii] = 0;\n incoming_g[ii] = 0;\n\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n incoming_p[ii] = p[i];\n incoming_m[ii] = m[i];\n incoming_v[ii] = v[i];\n incoming_g[ii] = static_cast(g[i]);\n }\n }\n\n // note for clarification to future michael:\n // From a pure memory dependency perspective, there's likely no point unrolling\n // the write loop, since writes just fire off once their LDGs arrive.\n // Put another way, the STGs are dependent on the LDGs, but not on each other.\n // There is still compute ILP benefit from unrolling the loop though.\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = i_start + threadIdx.x + ii * blockDim.x;\n\n if (j < n && j < chunk_size) {\n T scaled_grad = incoming_g[ii] / grad_scale;\n m[j] = b1 * incoming_m[ii] + (1 - b1) * scaled_grad;\n v[j] = b2 * incoming_v[ii] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(v[j] + eps);\n else // Mode 1\n denom = sqrtf(v[j]) + eps;\n float update = (m[j] / denom) + (decay * incoming_p[ii]);\n p[j] = incoming_p[ii] - (step_size * update);\n if (DEPTH == 5) p_copy[j] = (GRAD_T)p[j];\n }\n }\n }\n }\n }\n};\n\nvoid fused_adam_cuda(at::Tensor& p, at::Tensor& p_copy, at::Tensor& m, at::Tensor& v, at::Tensor& g, float lr,\n float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction,\n float decay) {\n // using namespace at;\n\n // Get tensor size\n int tsize = p.numel();\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (g.scalar_type() == at::ScalarType::Half) {\n // all other values should be fp32 for half gradients\n TORCH_CHECK(p.scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dispatch is done on the gradient type\n using namespace at; // prevents \"toString is undefined\" errors\n DISPATCH_FLOAT_AND_HALF(g.scalar_type(), 0, \"adam_cuda_kernel\", using accscalar_t = at::acc_type;\n adam_cuda_kernel<<>>(\n p.data_ptr(), p_copy.numel() ? p_copy.data_ptr() : NULL,\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n } else {\n using namespace at;\n DISPATCH_DOUBLE_AND_FLOAT(g.scalar_type(), 0, \"adam_cuda_kernel\",\n adam_cuda_kernel<<>>(\n p.data_ptr(),\n NULL, // don't output p_copy for fp32, it's wasted write\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n }\n C10_CUDA_CHECK(hipGetLastError());\n}\n\nvoid fused_adam_cuda_mt(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists, // p, m, v, g, p_copy\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay) {\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n\n size_t tl_sz = tensor_lists.size();\n TORCH_CHECK(tl_sz == 4 || tl_sz == 5, \"expected tensor lists of size 4 or 5\");\n\n if (tensor_lists[3][0].scalar_type() == at::ScalarType::Half) {\n // alher values should be fp32 for half gradients\n TORCH_CHECK(tensor_lists[0][0].scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dich is done on the gradient type\n if (tl_sz == 5) {\n DISPATCH_FLOAT_AND_HALF(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n using accscalar_t = at::acc_type;\n multi_tensor_apply<5>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<5, accscalar_t, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n } else {\n DISPATCH_FLOAT_AND_HALF(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n using accscalar_t = at::acc_type;\n multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<4, accscalar_t, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n }\n } else {\n if (tl_sz == 5) {\n DISPATCH_DOUBLE_AND_FLOAT(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n multi_tensor_apply<5>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<5, scalar_t_0, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n } else {\n DISPATCH_DOUBLE_AND_FLOAT(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<4, scalar_t_0, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n }\n }\n C10_CUDA_CHECK(hipGetLastError());\n}\n\ntemplate \n__device__ void convert(const FROM_T vi, TO_T& vo) {\n vo = static_cast(vi);\n}\n\ntemplate <>\n__device__ void convert(const float vi, uint8_t& vo) {\n union S {\n float as_float;\n int as_int;\n };\n S s;\n s.as_float = vi;\n s.as_int = s.as_int & 0xFF800000;\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_half = static_cast(vi + s.as_float / 8.0f);\n vo = t.as_byte[1];\n}\n\ntemplate <>\n__device__ void convert(const uint8_t vi, float& vo) {\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_byte[0] = 0;\n t.as_byte[1] = vi;\n vo = static_cast(t.as_half);\n}\n\ntemplate <>\n__device__ void convert(const at::Half vi, uint8_t& vo) {\n union S {\n float as_float;\n int as_int;\n };\n S s;\n s.as_float = static_cast(vi);\n s.as_int = s.as_int & 0xFF800000;\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_half = static_cast(vi + s.as_float / 8.0f);\n vo = t.as_byte[1];\n}\n\ntemplate <>\n__device__ void convert(const uint8_t vi, at::Half& vo) {\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_byte[0] = 0;\n t.as_byte[1] = vi;\n vo = t.as_half;\n}\n\ntemplate \n__global__ void strided_check_finite_cuda_kernel(volatile int* noop_gmem, GRAD_T* __restrict__ p_copy,\n const size_t tsize, int stride, int clear_overflow_first) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock) * stride;\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock * stride;\n\n if (clear_overflow_first) {\n if (i == 0) {\n *noop_gmem = 0;\n }\n __syncthreads();\n }\n\n for (int j = i; j < tsize; j += totThreads) {\n GRAD_T pi = p_copy[j];\n if (!isfinite(pi)) {\n *noop_gmem = 1;\n }\n }\n}\ntemplate <>\n__global__ void strided_check_finite_cuda_kernel(volatile int* noop_gmem, uint8_t* __restrict__ p_copy,\n const size_t tsize, int stride, int clear_overflow_first) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock) * stride;\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock * stride;\n\n if (clear_overflow_first) {\n if (i == 0) {\n *noop_gmem = 0;\n }\n __syncthreads();\n }\n\n for (int j = i; j < tsize; j += totThreads) {\n at::Half pi;\n convert(p_copy[j], pi);\n if (!isfinite(pi)) {\n *noop_gmem = 1;\n }\n }\n}\n\ntemplate \n__global__ void maybe_cast_kernel(volatile int* overflow_flag, const FROM_T* p_in, TO_T* p_out, const size_t tsize) {\n if (overflow_flag && *overflow_flag != 0) return;\n\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n FROM_T pi[ILP];\n TO_T po[ILP];\n\n for (int j_start = 0; j_start < tsize; j_start += totThreads * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n pi[ii] = 0;\n\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n pi[ii] = p_in[j];\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n convert(pi[ii], po[ii]);\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n p_out[j] = po[ii];\n }\n }\n }\n}\n\ntemplate \n__global__ void reversible_adam_cuda_kernel(\n T* __restrict__ p,\n REDU_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed\n T* __restrict__ m, T* __restrict__ v, const GRAD_T* __restrict__ g, const float b1, const float b2, const float eps,\n const float grad_scale, const float step_size, const size_t tsize, adamMode_t mode, const float decay) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n T mi[ILP];\n T vi[ILP];\n T pi[ILP];\n T gi[ILP];\n\n bool overflow = false;\n for (int j_start = 0; j_start < tsize; j_start += totThreads * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n mi[ii] = T(0);\n vi[ii] = T(0);\n pi[ii] = T(0);\n gi[ii] = GRAD_T(0);\n\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n pi[ii] = p[j];\n mi[ii] = m[j];\n vi[ii] = v[j];\n gi[ii] = static_cast(g[j]);\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n T scaled_grad = gi[ii] / grad_scale;\n if (isfinite(scaled_grad)) {\n mi[ii] = b1 * mi[ii] + (1 - b1) * scaled_grad;\n vi[ii] = b2 * vi[ii] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(vi[ii] + eps);\n else // Mode 1\n denom = sqrtf(vi[ii]) + eps;\n float update = (mi[ii] / denom) + (decay * pi[ii]);\n pi[ii] = pi[ii] - (step_size * update);\n } else {\n overflow = true;\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n m[j] = mi[ii];\n v[j] = vi[ii];\n p[j] = pi[ii];\n if (p_copy != NULL) {\n convert(pi[ii], p_copy[j]);\n }\n }\n }\n }\n\n if (p_copy != NULL) {\n __syncthreads();\n if (overflow) {\n convert(float(INFINITY), p_copy[0]);\n }\n }\n}\n\ntemplate \n__global__ void maybe_adam_undo_cuda_kernel(volatile int* overflow_flag, T* __restrict__ p, T* __restrict__ m,\n T* __restrict__ v, const GRAD_T* __restrict__ g, const float b1,\n const float b2, const float eps, const float grad_scale,\n const float step_size, const size_t tsize, adamMode_t mode,\n const float decay) {\n // NB! Skip undo kernel when overflow flag is NOT set\n if (overflow_flag && *overflow_flag == 0) return;\n\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n T mi[ILP];\n T vi[ILP];\n T pi[ILP];\n T gi[ILP];\n\n for (int j_start = 0; j_start < tsize; j_start += totThreads * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n mi[ii] = T(0);\n vi[ii] = T(0);\n pi[ii] = T(0);\n gi[ii] = GRAD_T(0);\n\n int j = j_start + i * ILP;\n if (j < tsize) {\n pi[ii] = p[j];\n mi[ii] = m[j];\n vi[ii] = v[j];\n gi[ii] = static_cast(g[j]);\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n T scaled_grad = gi[ii] / grad_scale;\n if (isfinite(scaled_grad)) {\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(vi[ii] + eps);\n else // Mode 1\n denom = sqrtf(vi[ii]) + eps;\n pi[ii] = (pi[ii] + step_size * (mi[ii] / denom)) / (1.0f - step_size * decay);\n mi[ii] = (mi[ii] - (1 - b1) * scaled_grad) / b1;\n vi[ii] = (vi[ii] - (1 - b2) * scaled_grad * scaled_grad) / b2;\n // Make sure round off errors don't create (small) negative value.\n // This can happen if we have to revert the very first step.\n vi[ii] = vi[ii] >= 0.0f ? vi[ii] : 0.0f;\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + i * ILP;\n if (j < tsize) {\n m[j] = mi[ii];\n v[j] = vi[ii];\n p[j] = pi[ii];\n }\n }\n }\n}\n\ntemplate \nstruct MaybeCastFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* overflow_flag,\n TensorListMetadata& tl) {\n if (overflow_flag && *overflow_flag != 0) return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n FROM_T* p_in = (FROM_T*)tl.addresses[0][tensor_loc];\n p_in += chunk_idx * chunk_size;\n TO_T* p_out = (TO_T*)tl.addresses[1][tensor_loc];\n p_out += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n int dim = chunk_size < n ? chunk_size : n;\n\n FROM_T pi[ILP];\n TO_T po[ILP];\n\n for (int j_start = 0; j_start < dim; j_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n pi[ii] = FROM_T(0);\n int j = j_start + threadIdx.x + ii * blockDim.x;\n if (j < dim) {\n pi[ii] = p_in[j];\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n convert(pi[ii], po[ii]);\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + threadIdx.x + ii * blockDim.x;\n if (j < dim) {\n p_out[j] = po[ii];\n }\n }\n }\n }\n};\n\nvoid fused_strided_check_finite(at::Tensor& overflow_flag, at::Tensor& p_copy, int stride, int clear_overflow_first) {\n // Get tensor size\n int tsize = p_copy.numel();\n int niter = (tsize + stride - 1) / stride;\n\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n // In order to avoid race condition, blocks must be 1 when clear_overflow_first flag is set.\n const dim3 blocks(clear_overflow_first ? 1 : (niter + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p_copy), \"parameter tensor is too large to be indexed with int32\");\n\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n using namespace at; // prevents \"toString is undefined\" errors\n DISPATCH_FLOAT_HALF_AND_BYTE(\n p_copy.scalar_type(), 0, \"check_finite_cuda_kernel\",\n strided_check_finite_cuda_kernel<<>>(\n overflow_flag.data_ptr(), p_copy.data_ptr(), tsize, stride, clear_overflow_first););\n C10_CUDA_CHECK(hipGetLastError());\n}\n\nvoid fused_reversible_adam_cuda(at::Tensor& p, at::Tensor& p_copy, at::Tensor& m, at::Tensor& v, at::Tensor& g,\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay) {\n // using namespace at;\n\n // Get tensor size\n int tsize = p.numel();\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (g.scalar_type() == at::ScalarType::Half) {\n // all other values should be fp32 for half gradients\n TORCH_CHECK(p.scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dispatch is done on the gradient type\n using namespace at; // prevents \"toString is undefined\" errors\n if (p_copy.numel() == 0 || p_copy.scalar_type() == g.scalar_type()) {\n DISPATCH_FLOAT_AND_HALF(\n g.scalar_type(), 0, \"adam_cuda_kernel\", using accscalar_t = at::acc_type;\n reversible_adam_cuda_kernel<<>>(\n p.data_ptr(), p_copy.numel() ? p_copy.data_ptr() : NULL,\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1, beta2, eps,\n grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n } else {\n TORCH_CHECK(p_copy.scalar_type() == at::ScalarType::Byte, \"expected parameter to be of byte type\");\n DISPATCH_FLOAT_AND_HALF(\n g.scalar_type(), 0, \"adam_cuda_e5m2_kernel\", using accscalar_t = at::acc_type;\n reversible_adam_cuda_kernel<<>>(\n p.data_ptr(), p_copy.data_ptr(), m.data_ptr(),\n v.data_ptr(), g.data_ptr(), beta1, beta2, eps, grad_scale, step_size, tsize,\n (adamMode_t)mode, decay););\n }\n } else {\n using namespace at;\n DISPATCH_DOUBLE_AND_FLOAT(g.scalar_type(), 0, \"adam_cuda_kernel\",\n reversible_adam_cuda_kernel\n <<>>(\n p.data_ptr(),\n NULL, // don't output p_copy for fp32, it's wasted write\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n }\n C10_CUDA_CHECK(hipGetLastError());\n}\n\nvoid maybe_cast_cuda(at::Tensor& overflow_flag, at::Tensor& p_in, at::Tensor& p_out) {\n // Get tensor size\n int tsize = p_in.numel();\n TORCH_CHECK(tsize == p_out.numel(), \"p_in.numel() must equal p_out.numel()\");\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p_in), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n DISPATCH_FLOAT_HALF_AND_BYTE(p_in.scalar_type(), 0,\n \"maybe_cast_cuda\" DISPATCH_FLOAT_HALF_AND_BYTE(\n p_out.scalar_type(), 1, \"maybe_cast_cuda\",\n maybe_cast_kernel<<>>(\n overflow_flag.numel() ? overflow_flag.data_ptr() : NULL,\n p_in.data_ptr(), p_out.data_ptr(), tsize);))\n C10_CUDA_CHECK(hipGetLastError());\n}\n\nvoid maybe_cast_cuda_mt(int chunk_size, at::Tensor overflow_flag,\n std::vector> tensor_lists) // p_in, p_out\n{\n // Constants\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n\n size_t tl_sz = tensor_lists.size();\n TORCH_CHECK(tl_sz == 2, \"expected tensor lists of size 2\");\n\n DISPATCH_FLOAT_HALF_AND_BYTE(\n tensor_lists[0][0].scalar_type(), 0, \"maybe_cast_cuda_mt_kernel\",\n DISPATCH_FLOAT_HALF_AND_BYTE(tensor_lists[1][0].scalar_type(), 1, \"maybe_cast_cuda_mt_kernel\",\n multi_tensor_apply<2>(BLOCK_SIZE, chunk_size, overflow_flag, tensor_lists,\n MaybeCastFunctor<2, scalar_t_0, scalar_t_1>());))\n C10_CUDA_CHECK(hipGetLastError());\n}\n\nvoid fused_maybe_adam_undo_cuda(at::Tensor& overflow_flag, at::Tensor& p, at::Tensor& m, at::Tensor& v, at::Tensor& g,\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay) {\n // Get tensor size\n int tsize = p.numel();\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (g.scalar_type() == at::ScalarType::Half) {\n // all other values should be fp32 for half gradients\n TORCH_CHECK(p.scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dispatch is done on the gradient type\n using namespace at; // prevents \"toString is undefined\" errors\n DISPATCH_FLOAT_AND_HALF(g.scalar_type(), 0, \"adam_cuda_kernel\", using accscalar_t = at::acc_type;\n maybe_adam_undo_cuda_kernel\n <<>>(\n overflow_flag.numel() ? overflow_flag.data_ptr() : NULL, p.data_ptr(),\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n } else {\n using namespace at;\n DISPATCH_DOUBLE_AND_FLOAT(\n g.scalar_type(), 0, \"adam_cuda_kernel\",\n maybe_adam_undo_cuda_kernel<<>>(\n overflow_flag.numel() ? overflow_flag.data_ptr() : NULL, p.data_ptr(),\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1, beta2, eps, grad_scale,\n step_size, tsize, (adamMode_t)mode, decay););\n }\n C10_CUDA_CHECK(hipGetLastError());\n}\n", "source": "/kaggle/working/apex/apex/contrib/csrc/optimizers/fused_adam_cuda_kernel.cu"} {"input": "#include \n#include \n#include \n\n#include \n\n__global__ void index_mul_2d_float_dim64(float* out, const float* in1, const float* in2, const int64_t* idx1,\n const int64_t size) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n constexpr int fea_dim = 64;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim) / 4 + tidx;\n int64_t vec_idx2 = (start_idx * fea_dim) / 4 + tidx;\n\n float4 res, src1, src2;\n src1 = reinterpret_cast(in1)[vec_idx1];\n src2 = reinterpret_cast(in2)[vec_idx2];\n res.x = src1.x * src2.x;\n res.y = src1.y * src2.y;\n res.z = src1.z * src2.z;\n res.w = src1.w * src2.w;\n reinterpret_cast(out)[vec_idx2] = res;\n }\n}\n\n__global__ void index_mul_2d_float(float* out, const float* in1, const float* in2, const int64_t* idx1,\n const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim);\n int64_t vec_idx2 = (start_idx * fea_dim);\n\n for (int i = tidx; i < fea_dim; i += stride) {\n out[vec_idx2 + i] = in1[vec_idx1 + i] * in2[vec_idx2 + i];\n }\n }\n}\n\n__global__ void index_mul_2d_half(at::Half* out, const at::Half* in1, const at::Half* in2, const int64_t* idx1,\n const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim);\n int64_t vec_idx2 = (start_idx * fea_dim);\n\n for (int i = tidx; i < fea_dim; i += stride) {\n out[vec_idx2 + i] = at::Half(static_cast(in1[vec_idx1 + i]) * static_cast(in2[vec_idx2 + i]));\n }\n }\n}\n\n__global__ void index_mul_2d_grad_float_dim64(float* grad_in1, float* grad_in2, const float* grad_out, const float* in1,\n const float* in2, const int64_t* idx1, const int64_t size) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n constexpr int fea_dim = 64;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim) / 4 + tidx;\n int64_t vec_idx2 = (start_idx * fea_dim) / 4 + tidx;\n\n float4 src_in1, src_in2, src_grad_out, dst_grad_in2;\n src_grad_out = reinterpret_cast(grad_out)[vec_idx2];\n src_in1 = reinterpret_cast(in1)[vec_idx1];\n src_in2 = reinterpret_cast(in2)[vec_idx2];\n int64_t grad_in1_base_idx = idx1[start_idx] * fea_dim + tidx * 4;\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 0, src_grad_out.x * src_in2.x);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 1, src_grad_out.y * src_in2.y);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 2, src_grad_out.z * src_in2.z);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 3, src_grad_out.w * src_in2.w);\n dst_grad_in2.x = src_grad_out.x * src_in1.x;\n dst_grad_in2.y = src_grad_out.y * src_in1.y;\n dst_grad_in2.z = src_grad_out.z * src_in1.z;\n dst_grad_in2.w = src_grad_out.w * src_in1.w;\n reinterpret_cast(grad_in2)[vec_idx2] = dst_grad_in2;\n }\n}\n\n__global__ void index_mul_2d_grad_float(float* grad_in1, float* grad_in2, const float* grad_out, const float* in1,\n const float* in2, const int64_t* idx1, const int64_t size,\n const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_in1 = in1[vec_idx1 + i];\n float src_in2 = in2[vec_idx2 + i];\n float src_grad_out = grad_out[vec_idx2 + i];\n grad_in2[vec_idx2 + i] = src_grad_out * src_in1;\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, src_grad_out * src_in2);\n }\n }\n}\n\n__global__ void index_mul_2d_grad_half(at::Half* grad_in1, at::Half* grad_in2, const at::Half* grad_out,\n const at::Half* in1, const at::Half* in2, const int64_t* idx1,\n const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_in1 = static_cast(in1[vec_idx1 + i]);\n float src_in2 = static_cast(in2[vec_idx2 + i]);\n float src_grad_out = static_cast(grad_out[vec_idx2 + i]);\n grad_in2[vec_idx2 + i] = at::Half(src_grad_out * src_in1);\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, at::Half(src_grad_out * src_in2));\n }\n }\n}\n\n__global__ void index_mul_2d_grad_grad_float_dim64(float* grad_grad_out, float* grad_in1, float* grad_in2,\n const float* grad_out, const float* grad_grad_in1,\n const float* grad_grad_in2, const float* in1, const float* in2,\n const int64_t* idx1, const int64_t size) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n constexpr int fea_dim = 64;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim) / 4 + tidx;\n int64_t vec_idx2 = (start_idx * fea_dim) / 4 + tidx;\n\n float4 src_grad_grad_in1, src_in1, src_grad_grad_in2, src_in2, src_grad_out;\n float4 dst_grad_grad_out, dst_grad_in2;\n src_grad_grad_in1 = reinterpret_cast(grad_grad_in1)[vec_idx1];\n src_in1 = reinterpret_cast(in1)[vec_idx1];\n src_grad_grad_in2 = reinterpret_cast(grad_grad_in2)[vec_idx2];\n src_in2 = reinterpret_cast(in2)[vec_idx2];\n dst_grad_grad_out.x = src_grad_grad_in1.x * src_in2.x + src_grad_grad_in2.x * src_in1.x;\n dst_grad_grad_out.y = src_grad_grad_in1.y * src_in2.y + src_grad_grad_in2.y * src_in1.y;\n dst_grad_grad_out.z = src_grad_grad_in1.z * src_in2.z + src_grad_grad_in2.z * src_in1.z;\n dst_grad_grad_out.w = src_grad_grad_in1.w * src_in2.w + src_grad_grad_in2.w * src_in1.w;\n reinterpret_cast(grad_grad_out)[vec_idx2] = dst_grad_grad_out;\n src_grad_out = reinterpret_cast(grad_out)[vec_idx2];\n int64_t grad_in1_base_idx = idx1[start_idx] * fea_dim + tidx * 4;\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 0, src_grad_grad_in2.x * src_grad_out.x);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 1, src_grad_grad_in2.y * src_grad_out.y);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 2, src_grad_grad_in2.z * src_grad_out.z);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 3, src_grad_grad_in2.w * src_grad_out.w);\n dst_grad_in2.x = src_grad_grad_in1.x * src_grad_out.x;\n dst_grad_in2.y = src_grad_grad_in1.y * src_grad_out.y;\n dst_grad_in2.z = src_grad_grad_in1.z * src_grad_out.z;\n dst_grad_in2.w = src_grad_grad_in1.w * src_grad_out.w;\n reinterpret_cast(grad_in2)[vec_idx2] = dst_grad_in2;\n }\n}\n\n__global__ void index_mul_2d_grad_grad_float(float* grad_grad_out, float* grad_in1, float* grad_in2,\n const float* grad_out, const float* grad_grad_in1,\n const float* grad_grad_in2, const float* in1, const float* in2,\n const int64_t* idx1, const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_grad_grad_in1 = grad_grad_in1[vec_idx1 + i];\n float src_grad_grad_in2 = grad_grad_in2[vec_idx2 + i];\n float src_in1 = in1[vec_idx1 + i];\n float src_in2 = in2[vec_idx2 + i];\n float src_grad_out = grad_out[vec_idx2 + i];\n grad_grad_out[vec_idx2 + i] = src_grad_grad_in1 * src_in2 + src_grad_grad_in2 * src_in1;\n grad_in2[vec_idx2 + i] = src_grad_grad_in1 * src_grad_out;\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, src_grad_grad_in2 * src_grad_out);\n }\n }\n}\n\n__global__ void index_mul_2d_grad_grad_half(at::Half* grad_grad_out, at::Half* grad_in1, at::Half* grad_in2,\n const at::Half* grad_out, const at::Half* grad_grad_in1,\n const at::Half* grad_grad_in2, const at::Half* in1, const at::Half* in2,\n const int64_t* idx1, const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_grad_grad_in1 = static_cast(grad_grad_in1[vec_idx1 + i]);\n float src_grad_grad_in2 = static_cast(grad_grad_in2[vec_idx2 + i]);\n float src_in1 = static_cast(in1[vec_idx1 + i]);\n float src_in2 = static_cast(in2[vec_idx2 + i]);\n float src_grad_out = static_cast(grad_out[vec_idx2 + i]);\n grad_grad_out[vec_idx2 + i] = at::Half(src_grad_grad_in1 * src_in2 + src_grad_grad_in2 * src_in1);\n grad_in2[vec_idx2 + i] = at::Half(src_grad_grad_in1 * src_grad_out);\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, at::Half(src_grad_grad_in2 * src_grad_out));\n }\n }\n}\n\nvoid index_mul_2d_float_foward_cuda(at::Tensor& out, const at::Tensor& in1, const at::Tensor& in2,\n const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fea_dim == 64) {\n const int BLOCK_THREADS_DIMX = 16;\n const int BLOCK_THREADS_DIMY = 16;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_float_dim64<<>>(\n out.data_ptr(), in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size);\n } else {\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_float<<>>(\n out.data_ptr(), in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n }\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid index_mul_2d_float_backward_cuda(at::Tensor& grad_in1, at::Tensor& grad_in2, const at::Tensor& grad_out,\n const at::Tensor& in1, const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fea_dim == 64) {\n const int BLOCK_THREADS_DIMX = 16;\n const int BLOCK_THREADS_DIMY = 16;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_float_dim64<<>>(\n grad_in1.data_ptr(), grad_in2.data_ptr(), grad_out.data_ptr(), in1.data_ptr(),\n in2.data_ptr(), idx1.data_ptr(), size);\n\n AT_CUDA_CHECK(cudaGetLastError());\n } else {\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_float<<>>(\n grad_in1.data_ptr(), grad_in2.data_ptr(), grad_out.data_ptr(), in1.data_ptr(),\n in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n }\n}\n\nvoid index_mul_2d_float_backward_backward_cuda(at::Tensor& grad_grad_out, at::Tensor& grad_in1, at::Tensor& grad_in2,\n const at::Tensor& grad_out, const at::Tensor& grad_grad_in1,\n const at::Tensor& grad_grad_in2, const at::Tensor& in1,\n const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fea_dim == 64) {\n const int BLOCK_THREADS_DIMX = 16;\n const int BLOCK_THREADS_DIMY = 16;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_grad_float_dim64<<>>(\n grad_grad_out.data_ptr(), grad_in1.data_ptr(), grad_in2.data_ptr(),\n grad_out.data_ptr(), grad_grad_in1.data_ptr(), grad_grad_in2.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size);\n } else {\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_grad_float<<>>(\n grad_grad_out.data_ptr(), grad_in1.data_ptr(), grad_in2.data_ptr(),\n grad_out.data_ptr(), grad_grad_in1.data_ptr(), grad_grad_in2.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n }\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid index_mul_2d_half_foward_cuda(at::Tensor& out, const at::Tensor& in1, const at::Tensor& in2,\n const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_half<<>>(\n out.data_ptr(), in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size,\n fea_dim);\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid index_mul_2d_half_backward_cuda(at::Tensor& grad_in1, at::Tensor& grad_in2, const at::Tensor& grad_out,\n const at::Tensor& in1, const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_half<<>>(\n grad_in1.data_ptr(), grad_in2.data_ptr(), grad_out.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n}\n\nvoid index_mul_2d_half_backward_backward_cuda(at::Tensor& grad_grad_out, at::Tensor& grad_in1, at::Tensor& grad_in2,\n const at::Tensor& grad_out, const at::Tensor& grad_grad_in1,\n const at::Tensor& grad_grad_in2, const at::Tensor& in1,\n const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_grad_half<<>>(\n grad_grad_out.data_ptr(), grad_in1.data_ptr(), grad_in2.data_ptr(),\n grad_out.data_ptr(), grad_grad_in1.data_ptr(), grad_grad_in2.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n", "output": "#include \"hip/hip_runtime.h\"\n#include \n#include \n#include \n\n#include \n\n__global__ void index_mul_2d_float_dim64(float* out, const float* in1, const float* in2, const int64_t* idx1,\n const int64_t size) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n constexpr int fea_dim = 64;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim) / 4 + tidx;\n int64_t vec_idx2 = (start_idx * fea_dim) / 4 + tidx;\n\n float4 res, src1, src2;\n src1 = reinterpret_cast(in1)[vec_idx1];\n src2 = reinterpret_cast(in2)[vec_idx2];\n res.x = src1.x * src2.x;\n res.y = src1.y * src2.y;\n res.z = src1.z * src2.z;\n res.w = src1.w * src2.w;\n reinterpret_cast(out)[vec_idx2] = res;\n }\n}\n\n__global__ void index_mul_2d_float(float* out, const float* in1, const float* in2, const int64_t* idx1,\n const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim);\n int64_t vec_idx2 = (start_idx * fea_dim);\n\n for (int i = tidx; i < fea_dim; i += stride) {\n out[vec_idx2 + i] = in1[vec_idx1 + i] * in2[vec_idx2 + i];\n }\n }\n}\n\n__global__ void index_mul_2d_half(at::Half* out, const at::Half* in1, const at::Half* in2, const int64_t* idx1,\n const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim);\n int64_t vec_idx2 = (start_idx * fea_dim);\n\n for (int i = tidx; i < fea_dim; i += stride) {\n out[vec_idx2 + i] = at::Half(static_cast(in1[vec_idx1 + i]) * static_cast(in2[vec_idx2 + i]));\n }\n }\n}\n\n__global__ void index_mul_2d_grad_float_dim64(float* grad_in1, float* grad_in2, const float* grad_out, const float* in1,\n const float* in2, const int64_t* idx1, const int64_t size) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n constexpr int fea_dim = 64;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim) / 4 + tidx;\n int64_t vec_idx2 = (start_idx * fea_dim) / 4 + tidx;\n\n float4 src_in1, src_in2, src_grad_out, dst_grad_in2;\n src_grad_out = reinterpret_cast(grad_out)[vec_idx2];\n src_in1 = reinterpret_cast(in1)[vec_idx1];\n src_in2 = reinterpret_cast(in2)[vec_idx2];\n int64_t grad_in1_base_idx = idx1[start_idx] * fea_dim + tidx * 4;\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 0, src_grad_out.x * src_in2.x);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 1, src_grad_out.y * src_in2.y);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 2, src_grad_out.z * src_in2.z);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 3, src_grad_out.w * src_in2.w);\n dst_grad_in2.x = src_grad_out.x * src_in1.x;\n dst_grad_in2.y = src_grad_out.y * src_in1.y;\n dst_grad_in2.z = src_grad_out.z * src_in1.z;\n dst_grad_in2.w = src_grad_out.w * src_in1.w;\n reinterpret_cast(grad_in2)[vec_idx2] = dst_grad_in2;\n }\n}\n\n__global__ void index_mul_2d_grad_float(float* grad_in1, float* grad_in2, const float* grad_out, const float* in1,\n const float* in2, const int64_t* idx1, const int64_t size,\n const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_in1 = in1[vec_idx1 + i];\n float src_in2 = in2[vec_idx2 + i];\n float src_grad_out = grad_out[vec_idx2 + i];\n grad_in2[vec_idx2 + i] = src_grad_out * src_in1;\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, src_grad_out * src_in2);\n }\n }\n}\n\n__global__ void index_mul_2d_grad_half(at::Half* grad_in1, at::Half* grad_in2, const at::Half* grad_out,\n const at::Half* in1, const at::Half* in2, const int64_t* idx1,\n const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_in1 = static_cast(in1[vec_idx1 + i]);\n float src_in2 = static_cast(in2[vec_idx2 + i]);\n float src_grad_out = static_cast(grad_out[vec_idx2 + i]);\n grad_in2[vec_idx2 + i] = at::Half(src_grad_out * src_in1);\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, at::Half(src_grad_out * src_in2));\n }\n }\n}\n\n__global__ void index_mul_2d_grad_grad_float_dim64(float* grad_grad_out, float* grad_in1, float* grad_in2,\n const float* grad_out, const float* grad_grad_in1,\n const float* grad_grad_in2, const float* in1, const float* in2,\n const int64_t* idx1, const int64_t size) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n constexpr int fea_dim = 64;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim) / 4 + tidx;\n int64_t vec_idx2 = (start_idx * fea_dim) / 4 + tidx;\n\n float4 src_grad_grad_in1, src_in1, src_grad_grad_in2, src_in2, src_grad_out;\n float4 dst_grad_grad_out, dst_grad_in2;\n src_grad_grad_in1 = reinterpret_cast(grad_grad_in1)[vec_idx1];\n src_in1 = reinterpret_cast(in1)[vec_idx1];\n src_grad_grad_in2 = reinterpret_cast(grad_grad_in2)[vec_idx2];\n src_in2 = reinterpret_cast(in2)[vec_idx2];\n dst_grad_grad_out.x = src_grad_grad_in1.x * src_in2.x + src_grad_grad_in2.x * src_in1.x;\n dst_grad_grad_out.y = src_grad_grad_in1.y * src_in2.y + src_grad_grad_in2.y * src_in1.y;\n dst_grad_grad_out.z = src_grad_grad_in1.z * src_in2.z + src_grad_grad_in2.z * src_in1.z;\n dst_grad_grad_out.w = src_grad_grad_in1.w * src_in2.w + src_grad_grad_in2.w * src_in1.w;\n reinterpret_cast(grad_grad_out)[vec_idx2] = dst_grad_grad_out;\n src_grad_out = reinterpret_cast(grad_out)[vec_idx2];\n int64_t grad_in1_base_idx = idx1[start_idx] * fea_dim + tidx * 4;\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 0, src_grad_grad_in2.x * src_grad_out.x);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 1, src_grad_grad_in2.y * src_grad_out.y);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 2, src_grad_grad_in2.z * src_grad_out.z);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 3, src_grad_grad_in2.w * src_grad_out.w);\n dst_grad_in2.x = src_grad_grad_in1.x * src_grad_out.x;\n dst_grad_in2.y = src_grad_grad_in1.y * src_grad_out.y;\n dst_grad_in2.z = src_grad_grad_in1.z * src_grad_out.z;\n dst_grad_in2.w = src_grad_grad_in1.w * src_grad_out.w;\n reinterpret_cast(grad_in2)[vec_idx2] = dst_grad_in2;\n }\n}\n\n__global__ void index_mul_2d_grad_grad_float(float* grad_grad_out, float* grad_in1, float* grad_in2,\n const float* grad_out, const float* grad_grad_in1,\n const float* grad_grad_in2, const float* in1, const float* in2,\n const int64_t* idx1, const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_grad_grad_in1 = grad_grad_in1[vec_idx1 + i];\n float src_grad_grad_in2 = grad_grad_in2[vec_idx2 + i];\n float src_in1 = in1[vec_idx1 + i];\n float src_in2 = in2[vec_idx2 + i];\n float src_grad_out = grad_out[vec_idx2 + i];\n grad_grad_out[vec_idx2 + i] = src_grad_grad_in1 * src_in2 + src_grad_grad_in2 * src_in1;\n grad_in2[vec_idx2 + i] = src_grad_grad_in1 * src_grad_out;\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, src_grad_grad_in2 * src_grad_out);\n }\n }\n}\n\n__global__ void index_mul_2d_grad_grad_half(at::Half* grad_grad_out, at::Half* grad_in1, at::Half* grad_in2,\n const at::Half* grad_out, const at::Half* grad_grad_in1,\n const at::Half* grad_grad_in2, const at::Half* in1, const at::Half* in2,\n const int64_t* idx1, const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_grad_grad_in1 = static_cast(grad_grad_in1[vec_idx1 + i]);\n float src_grad_grad_in2 = static_cast(grad_grad_in2[vec_idx2 + i]);\n float src_in1 = static_cast(in1[vec_idx1 + i]);\n float src_in2 = static_cast(in2[vec_idx2 + i]);\n float src_grad_out = static_cast(grad_out[vec_idx2 + i]);\n grad_grad_out[vec_idx2 + i] = at::Half(src_grad_grad_in1 * src_in2 + src_grad_grad_in2 * src_in1);\n grad_in2[vec_idx2 + i] = at::Half(src_grad_grad_in1 * src_grad_out);\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, at::Half(src_grad_grad_in2 * src_grad_out));\n }\n }\n}\n\nvoid index_mul_2d_float_foward_cuda(at::Tensor& out, const at::Tensor& in1, const at::Tensor& in2,\n const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fea_dim == 64) {\n const int BLOCK_THREADS_DIMX = 16;\n const int BLOCK_THREADS_DIMY = 16;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_float_dim64<<>>(\n out.data_ptr(), in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size);\n } else {\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_float<<>>(\n out.data_ptr(), in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n }\n\n AT_CUDA_CHECK(hipGetLastError());\n}\n\nvoid index_mul_2d_float_backward_cuda(at::Tensor& grad_in1, at::Tensor& grad_in2, const at::Tensor& grad_out,\n const at::Tensor& in1, const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fea_dim == 64) {\n const int BLOCK_THREADS_DIMX = 16;\n const int BLOCK_THREADS_DIMY = 16;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_float_dim64<<>>(\n grad_in1.data_ptr(), grad_in2.data_ptr(), grad_out.data_ptr(), in1.data_ptr(),\n in2.data_ptr(), idx1.data_ptr(), size);\n\n AT_CUDA_CHECK(hipGetLastError());\n } else {\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_float<<>>(\n grad_in1.data_ptr(), grad_in2.data_ptr(), grad_out.data_ptr(), in1.data_ptr(),\n in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n }\n}\n\nvoid index_mul_2d_float_backward_backward_cuda(at::Tensor& grad_grad_out, at::Tensor& grad_in1, at::Tensor& grad_in2,\n const at::Tensor& grad_out, const at::Tensor& grad_grad_in1,\n const at::Tensor& grad_grad_in2, const at::Tensor& in1,\n const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fea_dim == 64) {\n const int BLOCK_THREADS_DIMX = 16;\n const int BLOCK_THREADS_DIMY = 16;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_grad_float_dim64<<>>(\n grad_grad_out.data_ptr(), grad_in1.data_ptr(), grad_in2.data_ptr(),\n grad_out.data_ptr(), grad_grad_in1.data_ptr(), grad_grad_in2.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size);\n } else {\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_grad_float<<>>(\n grad_grad_out.data_ptr(), grad_in1.data_ptr(), grad_in2.data_ptr(),\n grad_out.data_ptr(), grad_grad_in1.data_ptr(), grad_grad_in2.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n }\n\n AT_CUDA_CHECK(hipGetLastError());\n}\n\nvoid index_mul_2d_half_foward_cuda(at::Tensor& out, const at::Tensor& in1, const at::Tensor& in2,\n const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_half<<>>(\n out.data_ptr(), in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size,\n fea_dim);\n\n AT_CUDA_CHECK(hipGetLastError());\n}\n\nvoid index_mul_2d_half_backward_cuda(at::Tensor& grad_in1, at::Tensor& grad_in2, const at::Tensor& grad_out,\n const at::Tensor& in1, const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_half<<>>(\n grad_in1.data_ptr(), grad_in2.data_ptr(), grad_out.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n}\n\nvoid index_mul_2d_half_backward_backward_cuda(at::Tensor& grad_grad_out, at::Tensor& grad_in1, at::Tensor& grad_in2,\n const at::Tensor& grad_out, const at::Tensor& grad_grad_in1,\n const at::Tensor& grad_grad_in2, const at::Tensor& in1,\n const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n hipStream_t stream = at::cuda::getCurrentCUDAStream();\n\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_grad_half<<>>(\n grad_grad_out.data_ptr(), grad_in1.data_ptr(), grad_in2.data_ptr(),\n grad_out.data_ptr(), grad_grad_in1.data_ptr(), grad_grad_in2.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n\n AT_CUDA_CHECK(hipGetLastError());\n}\n", "source": "/kaggle/working/apex/apex/contrib/csrc/index_mul_2d/index_mul_2d_cuda_kernel.cu"} {"input": "/**\n * From PyTorch:\n *\n * Copyright (c) 2016- Facebook, Inc (Adam Paszke)\n * Copyright (c) 2014- Facebook, Inc (Soumith Chintala)\n * Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)\n * Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu)\n * Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)\n * Copyright (c) 2011-2013 NYU (Clement Farabet)\n * Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)\n * Copyright (c) 2006 Idiap Research Institute (Samy Bengio)\n * Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)\n *\n * From Caffe2:\n *\n * Copyright (c) 2016-present, Facebook Inc. All rights reserved.\n *\n * All contributions by Facebook:\n * Copyright (c) 2016 Facebook Inc.\n *\n * All contributions by Google:\n * Copyright (c) 2015 Google Inc.\n * All rights reserved.\n *\n * All contributions by Yangqing Jia:\n * Copyright (c) 2015 Yangqing Jia\n * All rights reserved.\n *\n * All contributions from Caffe:\n * Copyright(c) 2013, 2014, 2015, the respective contributors\n * All rights reserved.\n *\n * All other contributions:\n * Copyright(c) 2015, 2016 the respective contributors\n * All rights reserved.\n *\n * Caffe2 uses a copyright model similar to Caffe: each contributor holds\n * copyright over their contributions to Caffe2. The project versioning records\n * all such contribution and copyright details. If a contributor wants to further\n * mark their specific copyright on a particular contribution, they should\n * indicate their copyright solely in the commit message of the change when it is\n * committed.\n *\n * All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n *\n * 1. Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n *\n * 2. Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n *\n * 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America\n * and IDIAP Research Institute nor the names of its contributors may be\n * used to endorse or promote products derived from this software without\n * specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE\n * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE\n * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR\n * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF\n * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS\n * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN\n * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)\n * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE\n * POSSIBILITY OF SUCH DAMAGE.\n */\n#include \n#include \n#include \n\n#include \n\n#include \"type_shim.h\"\n\n#define ALIGN_BYTES 16\n\nusing Tensor = at::Tensor;\nusing TensorList = at::TensorList;\nusing ScalarType = at::ScalarType;\nusing at::acc_type;\n\ntemplate \nstruct LogSoftMaxForwardEpilogue {\n __device__ __forceinline__ LogSoftMaxForwardEpilogue(AccumT max_input, AccumT sum)\n : logsum(max_input + std::log(sum)) {}\n\n __device__ __forceinline__ LogSoftMaxForwardEpilogue(AccumT max_log_sum_exp) : logsum(max_log_sum_exp) {}\n\n __device__ __forceinline__ OutT operator()(T input) const { return static_cast(input - logsum); }\n\n const AccumT logsum;\n};\n\ntemplate \nstruct LogSoftMaxBackwardEpilogue {\n __device__ __forceinline__ LogSoftMaxBackwardEpilogue(AccumT sum) : sum(sum) {}\n\n __device__ __forceinline__ T operator()(OutT gradOutput, OutT output) const {\n return static_cast(gradOutput - std::exp(static_cast(output)) * sum);\n }\n\n const AccumT sum;\n};\n\nconst int max_threads = 1024;\n\ninline dim3 SoftMax_getBlockSize(int ILP, uint64_t dim_size) {\n uint64_t block_size = 1;\n uint64_t max_block_size = std::min(dim_size / ILP, static_cast(max_threads));\n while (block_size < (max_block_size / 2)) block_size *= 2;\n // Launch at least a single warp - the kernel assumes that.\n block_size = std::max(block_size, static_cast(32));\n return dim3(block_size);\n}\n\ntemplate \nstruct Add {\n __device__ __forceinline__ T operator()(T a, T b) const { return a + b; }\n};\n\ntemplate \nstruct Max {\n __device__ __forceinline__ T operator()(T a, T b) const { return a < b ? b : a; }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n// Regular kernel (fast when dim_size is large; requires inner_size == 1)\n////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct MaxFloat {\n __device__ __forceinline__ AccumT operator()(AccumT max, T v) const { return ::max(max, (AccumT)v); }\n};\n\ntemplate \nstruct AddFloat {\n __device__ __forceinline__ AccumT operator()(AccumT sum, T v) const { return sum + v; }\n};\n\ntemplate \nstruct SumExpFloat {\n __device__ __forceinline__ SumExpFloat(AccumT v) : max_k(v) {}\n\n __device__ __forceinline__ AccumT operator()(AccumT sum, T v) const { return sum + std::exp(v - max_k); }\n\n const AccumT max_k;\n};\n\ntemplate